247 lines
7.8 KiB
Python
247 lines
7.8 KiB
Python
from typing import Dict, List
|
|
import logging
|
|
import numpy as np
|
|
from wall_x.data.utils import preprocesser_call
|
|
from qwen_vl_utils.vision_process import smart_resize
|
|
import torch
|
|
from PIL import Image
|
|
from transformers import BatchFeature
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
def prepare_batch(
|
|
obs: Dict,
|
|
processor,
|
|
camera_key: List[str],
|
|
agent_pos_dim,
|
|
action_dim,
|
|
pred_horizon,
|
|
fixed_action_dim,
|
|
max_length,
|
|
image_factor: int,
|
|
min_pixels: int,
|
|
max_pixels: int,
|
|
predict_mode: str = "fast",
|
|
device: str = "cuda",
|
|
) -> BatchFeature:
|
|
"""Prepare observation into model input format.
|
|
|
|
Args:
|
|
obs: Dictionary containing:
|
|
- 'camera_key_0' : image 0
|
|
- 'camera_key_1' : image 1
|
|
...
|
|
- 'prompt': Text prompt
|
|
- 'state': Robot state/proprioception
|
|
- 'dataset_names': Dataset names
|
|
|
|
Returns:
|
|
BatchFeature object ready for model input
|
|
"""
|
|
# Handle images - can be single image, list of images, or dict of images
|
|
images = []
|
|
images = [obs[key] for key in camera_key]
|
|
# Convert numpy arrays to PIL Images
|
|
processed_images = []
|
|
for img in images:
|
|
if isinstance(img, np.ndarray):
|
|
# Debug: Log the shape and dtype
|
|
logger.debug(f"Image shape: {img.shape}, dtype: {img.dtype}")
|
|
|
|
# Handle unexpected dimensions - squeeze if needed
|
|
if img.ndim > 3:
|
|
logger.warning(
|
|
f"Image has {img.ndim} dimensions, squeezing extra dimensions"
|
|
)
|
|
img = np.squeeze(img)
|
|
|
|
# Verify shape is valid for PIL
|
|
if img.ndim == 2:
|
|
# Grayscale image
|
|
pass
|
|
elif img.ndim == 3:
|
|
# Check if channel dimension is first or last
|
|
if img.shape[0] == 3 or img.shape[0] == 1:
|
|
# Channels first, transpose to channels last
|
|
img = np.transpose(img, (1, 2, 0))
|
|
elif img.shape[2] == 3 or img.shape[2] == 1:
|
|
# Already channels last
|
|
pass
|
|
else:
|
|
raise ValueError(
|
|
f"Unexpected image shape: {img.shape}. Expected (H, W, C) or (C, H, W)"
|
|
)
|
|
else:
|
|
raise ValueError(
|
|
f"Invalid image dimensions: {img.ndim}. Expected 2 or 3 dimensions, got shape {img.shape}"
|
|
)
|
|
|
|
# Convert to PIL Image
|
|
if img.dtype == np.uint8:
|
|
img = Image.fromarray(img)
|
|
else:
|
|
img = Image.fromarray((img * 255).astype(np.uint8))
|
|
processed_images.append(img)
|
|
|
|
# Apply smart resize to images
|
|
resized_images = process_images(
|
|
processed_images, image_factor, min_pixels, max_pixels
|
|
)
|
|
|
|
# Handle text prompt - format with vision tokens
|
|
instruction = obs["prompt"]
|
|
formatted_text = format_text_with_vision_tokens(
|
|
instruction, camera_key, predict_mode, pred_horizon
|
|
)
|
|
|
|
# Use processor to prepare inputs
|
|
inputs = preprocesser_call(
|
|
processor=processor,
|
|
text=[formatted_text],
|
|
images=[resized_images],
|
|
videos=None,
|
|
padding=True,
|
|
truncation=True,
|
|
return_tensors="pt",
|
|
max_length=max_length,
|
|
)
|
|
|
|
action_token_id = processor.tokenizer.convert_tokens_to_ids("<|action|>")
|
|
moe_token_types = inputs.input_ids == action_token_id
|
|
inputs["moe_token_types"] = moe_token_types
|
|
|
|
# Handle robot state/proprioception if available
|
|
if "state" in obs:
|
|
state = obs["state"]
|
|
if isinstance(state, np.ndarray):
|
|
state = torch.from_numpy(state).float()
|
|
elif not isinstance(state, torch.Tensor):
|
|
state = torch.tensor(state, dtype=torch.float32)
|
|
|
|
# Add batch dimension if needed
|
|
if state.dim() == 1:
|
|
state = state.unsqueeze(0)
|
|
if state.dim() == 2:
|
|
state = state.unsqueeze(1) # [batch, 1, state_dim]
|
|
|
|
# Pad to 20 dimensions if needed (same as training)
|
|
if state.shape[-1] < 20:
|
|
padding = torch.zeros(state.shape[0], state.shape[1], 20 - state.shape[-1])
|
|
state = torch.cat([state, padding], dim=-1)
|
|
|
|
# Create mask for valid dimensions
|
|
agent_pos_mask = torch.ones_like(state)
|
|
if state.shape[-1] > agent_pos_dim:
|
|
agent_pos_mask[:, :, agent_pos_dim:] = 0
|
|
|
|
inputs["proprioception"] = state
|
|
inputs["agent_pos_mask"] = agent_pos_mask
|
|
|
|
# Add dataset name (required by model)
|
|
inputs["dataset_names"] = obs["dataset_names"]
|
|
|
|
# Move all tensors to device
|
|
for key in inputs:
|
|
if isinstance(inputs[key], torch.Tensor):
|
|
inputs[key] = inputs[key].to(device)
|
|
|
|
dof_mask = torch.ones([state.shape[0], pred_horizon, fixed_action_dim])
|
|
dof_mask[:, :, action_dim:] = 0
|
|
|
|
inputs["dof_mask"] = dof_mask
|
|
|
|
# Convert to BatchFeature to maintain consistency with training pipeline
|
|
return BatchFeature(data=dict(inputs)).to(device)
|
|
|
|
|
|
def process_images(
|
|
images: List[Image.Image], image_factor: int, min_pixels: int, max_pixels: int
|
|
) -> List[Image.Image]:
|
|
"""Process images with smart resize following the data loading pattern.
|
|
|
|
Args:
|
|
images: List of PIL Images
|
|
|
|
Returns:
|
|
List of resized PIL Images
|
|
"""
|
|
resized_images = []
|
|
for img_pil in images:
|
|
current_width, current_height = img_pil.size
|
|
|
|
# Apply smart scaling (Qwen logic)
|
|
resized_height, resized_width = smart_resize(
|
|
current_height,
|
|
current_width,
|
|
factor=image_factor,
|
|
min_pixels=min_pixels,
|
|
max_pixels=max_pixels,
|
|
)
|
|
|
|
resized_img = img_pil.resize((resized_width, resized_height))
|
|
resized_images.append(resized_img)
|
|
|
|
return resized_images
|
|
|
|
|
|
def format_text_with_vision_tokens(
|
|
instruction: str,
|
|
camera_key: List[str],
|
|
predict_mode: str = "fast",
|
|
pred_horizon: int = 32,
|
|
) -> str:
|
|
"""Format text prompt with vision tokens for the model.
|
|
|
|
Args:
|
|
instruction: Task instruction text
|
|
camera_key: List of camera names
|
|
|
|
Returns:
|
|
Formatted text with special tokens
|
|
"""
|
|
# Special tokens for formatting
|
|
role_start_symbol = "<|im_start|>"
|
|
role_end_symbol = "<|im_end|>"
|
|
vision_start_symbol = "<|vision_start|>"
|
|
vision_end_symbol = "<|vision_end|>"
|
|
image_pad_symbol = "<|image_pad|>"
|
|
propri_symbol = "<|propri|>"
|
|
action_symbol = "<|action|>"
|
|
# action_fast_symbol = "<|action_fast|>"
|
|
|
|
# Camera name mapping
|
|
camera_name_mapping = {
|
|
"front_view": "front view",
|
|
"face_view": "front view",
|
|
"left_wrist_view": "left wrist view",
|
|
"right_wrist_view": "right wrist view",
|
|
"top_view": "top view",
|
|
"wall_view": "wall view",
|
|
}
|
|
pred_horizon = 32
|
|
|
|
# System prologue
|
|
prologue = (
|
|
f"{role_start_symbol}system\nYou are a helpful assistant.{role_end_symbol}\n"
|
|
)
|
|
|
|
# User request with observation
|
|
user_request = f"{role_start_symbol}user\nObservation:"
|
|
if camera_key:
|
|
for cam_name in camera_key:
|
|
view_name = camera_name_mapping.get(cam_name, cam_name)
|
|
user_request += f" {view_name}: {vision_start_symbol}{image_pad_symbol}{vision_end_symbol}"
|
|
user_request += "\nInstruction:"
|
|
|
|
text_prompt = (
|
|
f"\nPredict the next action in robot action.\nProprioception: {propri_symbol}\n"
|
|
)
|
|
user_message = f"{user_request} {instruction}{text_prompt}{role_end_symbol}\n"
|
|
assistant_output = f"{role_start_symbol}assistant\n"
|
|
if predict_mode == "diffusion":
|
|
assistant_output += f"{action_symbol * pred_horizon}"
|
|
complete_text = prologue + user_message + assistant_output
|
|
|
|
return complete_text
|