Init
This commit is contained in:
Starrick
@@ -0,0 +1,95 @@
|
||||
from typing import List, Dict, Optional
|
||||
from dataclasses import dataclass, field
|
||||
from qwen_vl_utils.vision_process import MIN_PIXELS, MAX_PIXELS, IMAGE_FACTOR
|
||||
|
||||
|
||||
# Tactile sensor file mapping for data processing
|
||||
TACTILE_FILE_MAPPING = {
|
||||
"tactile_data_left": "left_tactile",
|
||||
"tactile_data_right": "right_tactile"
|
||||
}
|
||||
|
||||
# Supported action datasets
|
||||
ACTION_DATASET_NAMES = [
|
||||
"x2_normal", "agibotworld_alpha", "droid", "fractal", "bridge_data_v2",
|
||||
"DobbE", "RH20T", "UMI-biarm", "austin_buds", "austin_sailor", "austin_sirius",
|
||||
"bc_z", "berkeley_autolab_ur5", "berkeley_cable_routing", "berkeley_fanuc_manipulation",
|
||||
"dlr_edan_shared_control", "fmb", "furniture_bench", "jaco_play", "nyu_rot",
|
||||
"stanford_hydra", "stanford_kuka_multimodal", "taco_play", "utaustin_mutex", "viola"
|
||||
]
|
||||
|
||||
# Supported multimodal datasets
|
||||
MULTIMODAL_DATASET_NAMES = [
|
||||
"x2_multimodal_from_action", "x2_multimodal", "x2_subtask_generation",
|
||||
"multimodal_CapsFusion", "multimodal_Robo2VLM", "multimodal_RoboPoint",
|
||||
"multimodal_EQA", "multimodal_Cambrian", "multimodal_pixmo",
|
||||
"multimodal_VQAv2", "multimodal_COCO"
|
||||
]
|
||||
|
||||
|
||||
@dataclass
|
||||
class X2RDataProcessingConfig:
|
||||
"""Configuration class for X2R data processing pipeline.
|
||||
|
||||
This class contains all the necessary parameters for processing robotic data
|
||||
including camera mappings, tactile sensor configurations, action predictions,
|
||||
and various processing options.
|
||||
"""
|
||||
|
||||
# Action prediction configuration
|
||||
predict_action_keys: List[str] = field(default_factory=list)
|
||||
obs_action_keys: List[str] = field(default_factory=list)
|
||||
|
||||
# Image resolution settings for different views
|
||||
resolution: Dict[str, int] = field(
|
||||
default_factory=lambda: {
|
||||
"face_view": -1,
|
||||
"left_wrist_view": 128,
|
||||
"right_wrist_view": 128
|
||||
}
|
||||
)
|
||||
|
||||
# Dataset splitting
|
||||
train_test_split: float = 0.9
|
||||
split_seed: int = 42
|
||||
|
||||
# Instruction handling
|
||||
priority_order: Optional[Dict[str, float]] = None
|
||||
|
||||
# Vision model parameters
|
||||
model_type: str = "qwen2_5"
|
||||
max_pixels: int = MAX_PIXELS
|
||||
min_pixels: int = MIN_PIXELS
|
||||
image_factor: int = IMAGE_FACTOR
|
||||
|
||||
generate_subtask_ratio: float = 0.0
|
||||
|
||||
def __post_init__(self):
|
||||
"""Post-initialization validation and setup."""
|
||||
# Validate train/test split
|
||||
if not 0 < self.train_test_split < 1:
|
||||
raise ValueError(f"train_test_split must be between 0 and 1, got {self.train_test_split}")
|
||||
|
||||
def as_dict(self) -> Dict:
|
||||
"""Convert configuration to dictionary format.
|
||||
|
||||
Returns:
|
||||
Dict: Configuration as dictionary
|
||||
"""
|
||||
return self.__dict__
|
||||
|
||||
def update(self, **kwargs) -> 'X2RDataProcessingConfig':
|
||||
"""Update configuration parameters.
|
||||
|
||||
Args:
|
||||
**kwargs: Key-value pairs to update
|
||||
|
||||
Returns:
|
||||
X2RDataProcessingConfig: Updated configuration instance
|
||||
"""
|
||||
for key, value in kwargs.items():
|
||||
if hasattr(self, key):
|
||||
setattr(self, key, value)
|
||||
else:
|
||||
raise ValueError(f"Unknown configuration parameter: {key}")
|
||||
return self
|
||||
@@ -0,0 +1,496 @@
|
||||
"""
|
||||
LeRobot Dataset Loader - Distributed Version
|
||||
"""
|
||||
|
||||
import torch
|
||||
from torch.utils.data import DistributedSampler
|
||||
from lerobot.datasets.lerobot_dataset import LeRobotDataset
|
||||
from typing import Protocol, SupportsIndex, TypeVar
|
||||
from qwen_vl_utils.vision_process import smart_resize
|
||||
from wall_x.data.config import X2RDataProcessingConfig
|
||||
from wall_x.data.utils import process_grounding_points, get_wallx_normal_text, replace_action_token, preprocesser_call
|
||||
|
||||
from transformers import AutoProcessor
|
||||
|
||||
T_co = TypeVar("T_co", covariant=True)
|
||||
|
||||
CAMERA_KEY_MAPPINGS = {
|
||||
"lerobot/aloha_mobile_cabinet": {
|
||||
"observation.images.cam_high": "face_view",
|
||||
"observation.images.cam_left_wrist": "left_wrist_view",
|
||||
"observation.images.cam_right_wrist": "right_wrist_view",
|
||||
},
|
||||
}
|
||||
|
||||
|
||||
# Abstract class for dataset
|
||||
class Dataset(Protocol[T_co]):
|
||||
"""Interface for a dataset with random access."""
|
||||
|
||||
def __getitem__(self, index: SupportsIndex) -> T_co:
|
||||
raise NotImplementedError("Subclasses of Dataset should implement __getitem__.")
|
||||
|
||||
def __len__(self) -> int:
|
||||
raise NotImplementedError("Subclasses of Dataset should implement __len__.")
|
||||
|
||||
|
||||
class PreprocessedDataset(Dataset[T_co]):
|
||||
def __init__(self, dataset, config, dataload_config, seed=42, rank=0, world_size=1):
|
||||
self._dataset = dataset
|
||||
self.seed = seed
|
||||
self.rank = rank
|
||||
self.world_size = world_size
|
||||
|
||||
# init configs
|
||||
self.config = config
|
||||
self.use_fast_tokenizer = self.config.get("use_fast_tokenizer", False)
|
||||
self.dataload_config = dataload_config
|
||||
|
||||
self.data_config = X2RDataProcessingConfig().update(
|
||||
train_test_split=self.dataload_config["train_test_split"],
|
||||
split_seed=self.dataload_config["split_seed"],
|
||||
predict_action_keys=self.dataload_config["predict_action_keys"],
|
||||
obs_action_keys=self.dataload_config["obs_action_keys"],
|
||||
resolution=self.dataload_config.get("resolution", None),
|
||||
priority_order=self.dataload_config.get("priority_order", None),
|
||||
)
|
||||
|
||||
self._cam_key_mapping = CAMERA_KEY_MAPPINGS[self._dataset.meta.repo_id]
|
||||
|
||||
def _vision_preprocess(self, frames):
|
||||
processed_frames = []
|
||||
for key in self._dataset.meta.camera_keys:
|
||||
from PIL import Image
|
||||
|
||||
current_obs = frames[key].clone().permute(1, 2, 0)
|
||||
|
||||
img_pil = Image.fromarray((current_obs * 255).to(torch.uint8).cpu().numpy())
|
||||
orig_width, orig_height = img_pil.size
|
||||
# 2. Apply resolution constraints (if config is not -1)
|
||||
target_size = self.data_config.resolution.get(self._cam_key_mapping[key], -1)
|
||||
if target_size != -1:
|
||||
# Maintain aspect ratio logic
|
||||
if orig_width > orig_height: # Landscape image
|
||||
new_width = target_size
|
||||
new_height = int(target_size * orig_height / orig_width)
|
||||
else: # Portrait image
|
||||
new_height = target_size
|
||||
new_width = int(target_size * orig_width / orig_height)
|
||||
img_pil = img_pil.resize((new_width, new_height))
|
||||
|
||||
# 3. Apply smart scaling (qwen logic)
|
||||
current_width, current_height = img_pil.size
|
||||
resized_height, resized_width = smart_resize(
|
||||
current_height,
|
||||
current_width,
|
||||
factor=self.data_config.image_factor,
|
||||
min_pixels=self.data_config.min_pixels,
|
||||
max_pixels=self.data_config.max_pixels,
|
||||
)
|
||||
resized_img = img_pil.resize((resized_width, resized_height))
|
||||
processed_frames.append(resized_img)
|
||||
|
||||
return processed_frames, orig_height, orig_width, resized_height, resized_width
|
||||
|
||||
def __getitem__(self, index):
|
||||
data = self._dataset[index]
|
||||
image_inputs, h, w, resize_h, resize_w = self._vision_preprocess(data)
|
||||
agent_pos = data["observation.state"]
|
||||
action = data["action"]
|
||||
frame_index = data["frame_index"]
|
||||
instruction_info = {"instruction": data["task"]}
|
||||
generate_subtask_ratio = self.data_config.generate_subtask_ratio
|
||||
complete_text, generate_subtask = get_wallx_normal_text(
|
||||
instruction_info,
|
||||
33 - 1,
|
||||
frame_index,
|
||||
self.data_config.priority_order,
|
||||
self._cam_key_mapping,
|
||||
generate_subtask_ratio=generate_subtask_ratio,
|
||||
)
|
||||
text = process_grounding_points(complete_text, h, w, resize_h, resize_w, self.data_config.model_type)
|
||||
result = {
|
||||
"image_inputs": image_inputs,
|
||||
"text": text,
|
||||
"action": action,
|
||||
"agent_pos": agent_pos,
|
||||
"frame_index": frame_index,
|
||||
}
|
||||
|
||||
return result
|
||||
|
||||
def __len__(self) -> int:
|
||||
return len(self._dataset)
|
||||
|
||||
def get_train_dataloader(self):
|
||||
"""
|
||||
Get distributed training dataloader
|
||||
|
||||
Args:
|
||||
rank: Current process rank
|
||||
world_size: Total number of processes
|
||||
seed: Random seed for reproducibility
|
||||
"""
|
||||
|
||||
batch_size = self.config.get("batch_size_per_gpu", 8)
|
||||
num_workers = self.config.get("num_workers", 4)
|
||||
|
||||
# Create distributed sampler
|
||||
sampler = DistributedSampler(
|
||||
self,
|
||||
num_replicas=self.world_size,
|
||||
rank=self.rank,
|
||||
shuffle=True,
|
||||
seed=self.seed,
|
||||
drop_last=True, # Ensure all processes have same number of batches
|
||||
)
|
||||
|
||||
dataloader = torch.utils.data.DataLoader(
|
||||
self,
|
||||
batch_size=batch_size,
|
||||
sampler=sampler, # Use distributed sampler instead of shuffle=True
|
||||
num_workers=num_workers,
|
||||
collate_fn=DataCollator(self.config, self.dataload_config, self._dataset.meta.stats),
|
||||
pin_memory=True, # Enable for GPU training
|
||||
persistent_workers=num_workers > 0, # Only if num_workers > 0
|
||||
prefetch_factor=2, # Reduce memory usage
|
||||
drop_last=True, # Avoid incomplete batches
|
||||
)
|
||||
|
||||
return dataloader, sampler
|
||||
|
||||
def get_val_dataloader(self):
|
||||
"""
|
||||
Get distributed evaluation dataloader (no shuffling for consistent evaluation)
|
||||
"""
|
||||
|
||||
batch_size = self.config.get("eval_batch_size_per_gpu", self.config.get("batch_size_per_gpu", 8))
|
||||
num_workers = self.config.get("num_workers", 4)
|
||||
|
||||
# Create distributed sampler for evaluation (no shuffle)
|
||||
sampler = DistributedSampler(
|
||||
self,
|
||||
num_replicas=self.world_size,
|
||||
rank=self.rank,
|
||||
shuffle=False, # No shuffling for evaluation
|
||||
drop_last=False, # Keep all samples for evaluation
|
||||
)
|
||||
|
||||
dataloader = torch.utils.data.DataLoader(
|
||||
self,
|
||||
batch_size=batch_size,
|
||||
sampler=sampler,
|
||||
num_workers=num_workers,
|
||||
collate_fn=DataCollator(self.config, self.dataload_config, self._dataset.meta.stats),
|
||||
pin_memory=True,
|
||||
persistent_workers=num_workers > 0,
|
||||
prefetch_factor=2,
|
||||
drop_last=False,
|
||||
)
|
||||
|
||||
return dataloader, sampler
|
||||
|
||||
|
||||
class DataCollator:
|
||||
# Class-level cache for processors to avoid reloading
|
||||
_processor_cache = {}
|
||||
_action_tokenizer_cache = {}
|
||||
|
||||
def __init__(self, config, dataload_config, stats):
|
||||
self.config = config
|
||||
self.dataload_config = dataload_config
|
||||
self.stats = stats
|
||||
self.min_stat = stats["action"]["min"]
|
||||
self.max_stat = stats["action"]["max"]
|
||||
self.delta = self.max_stat - self.min_stat
|
||||
self.use_fast_tokenizer = self.config.get("use_fast_tokenizer", False)
|
||||
self.load_processor()
|
||||
|
||||
def load_processor(self):
|
||||
processor_path = self.config["processor_path"]
|
||||
action_tokenizer_path = self.config["action_tokenizer_path"]
|
||||
|
||||
# Use cached processors if available
|
||||
if processor_path not in self._processor_cache:
|
||||
self._processor_cache[processor_path] = AutoProcessor.from_pretrained(processor_path, use_fast=True)
|
||||
if self.config.get("padding_side", "left") == "left":
|
||||
self._processor_cache[processor_path].tokenizer.padding_side = "left"
|
||||
|
||||
if action_tokenizer_path not in self._action_tokenizer_cache:
|
||||
self._action_tokenizer_cache[action_tokenizer_path] = AutoProcessor.from_pretrained(action_tokenizer_path, trust_remote_code=True)
|
||||
|
||||
self.processor = self._processor_cache[processor_path]
|
||||
self.val_processor = self._processor_cache[processor_path]
|
||||
self.train_action_tokenizer = self._action_tokenizer_cache[action_tokenizer_path]
|
||||
self.val_action_tokenizer = self._action_tokenizer_cache[action_tokenizer_path]
|
||||
|
||||
new_tokens = ["<|propri|>", "<|action|>"]
|
||||
|
||||
new_tokens += [f"<|action_token_{i}|>" for i in range(self.train_action_tokenizer.vocab_size)]
|
||||
if not self.use_fast_tokenizer:
|
||||
self.train_action_tokenizer = None
|
||||
self.val_action_tokenizer = None
|
||||
|
||||
# Only add tokens if not already added
|
||||
if "<|propri|>" not in self.processor.tokenizer.get_vocab():
|
||||
num_added_tokens = self.processor.tokenizer.add_tokens(new_tokens)
|
||||
self.val_processor.tokenizer.add_tokens(new_tokens)
|
||||
|
||||
if self.use_fast_tokenizer:
|
||||
self.action_mapper = {}
|
||||
for i in range(self.train_action_tokenizer.vocab_size):
|
||||
token = f"<|action_token_{i}|>"
|
||||
token_id = self.processor.tokenizer.convert_tokens_to_ids(token)
|
||||
self.action_mapper[token_id] = i
|
||||
else:
|
||||
self.action_mapper = None
|
||||
|
||||
@classmethod
|
||||
def _normalize(cls, action, min_stat, delta):
|
||||
x = (action - min_stat) / (delta)
|
||||
x = x * 2 - 1
|
||||
x = torch.clamp(x, -1, 1)
|
||||
return x
|
||||
|
||||
def __call__(self, batch):
|
||||
additional_inputs = {}
|
||||
|
||||
for key in batch[0].keys():
|
||||
if key == "agent_pos":
|
||||
agent_pos = torch.stack([item["agent_pos"] for item in batch])
|
||||
if agent_pos.dim() == 2:
|
||||
agent_pos = agent_pos.unsqueeze(1)
|
||||
agent_pos_mask = (~torch.isnan(agent_pos)).float()
|
||||
agent_pos.nan_to_num_(nan=0.0)
|
||||
agent_pos = self._normalize(agent_pos, self.min_stat, self.delta)
|
||||
if agent_pos.shape[-1] != 20:
|
||||
agent_pos = torch.cat([agent_pos, torch.zeros(agent_pos.shape[0], agent_pos.shape[1], 20 - agent_pos.shape[-1])], dim=-1)
|
||||
agent_pos_mask = torch.cat(
|
||||
[agent_pos_mask, torch.zeros(agent_pos_mask.shape[0], agent_pos_mask.shape[1], 20 - agent_pos_mask.shape[-1])], dim=-1
|
||||
)
|
||||
additional_inputs["proprioception"] = agent_pos
|
||||
additional_inputs["agent_pos_mask"] = agent_pos_mask
|
||||
elif key == "action":
|
||||
action = torch.stack([item["action"] for item in batch])
|
||||
if action.dim() == 2:
|
||||
action = action.unsqueeze(1)
|
||||
dof_mask = (~torch.isnan(action)).float()
|
||||
action.nan_to_num_(nan=0.0)
|
||||
action = self._normalize(action, self.min_stat, self.delta)
|
||||
if action.shape[-1] != 20:
|
||||
action = torch.cat([action, torch.zeros(action.shape[0], action.shape[1], 20 - action.shape[-1])], dim=-1)
|
||||
dof_mask = torch.cat([dof_mask, torch.zeros(dof_mask.shape[0], dof_mask.shape[1], 20 - dof_mask.shape[-1])], dim=-1)
|
||||
additional_inputs["action_chunk"] = action
|
||||
additional_inputs["dof_mask"] = dof_mask
|
||||
elif key == "image_inputs":
|
||||
additional_inputs["image_inputs"] = [item["image_inputs"] for item in batch]
|
||||
elif key == "text":
|
||||
additional_inputs["text"] = [item["text"] for item in batch]
|
||||
elif key == "frame_index":
|
||||
additional_inputs["frame_index"] = torch.stack([item["frame_index"] for item in batch])
|
||||
else:
|
||||
raise NotImplementedError(f"{key} input not implemented in preprocesser")
|
||||
|
||||
additional_inputs["text"] = replace_action_token(
|
||||
additional_inputs["text"],
|
||||
additional_inputs["action_chunk"],
|
||||
self.train_action_tokenizer if self.use_fast_tokenizer else None,
|
||||
["x2_normal"] * additional_inputs["text"].__len__(),
|
||||
additional_inputs["dof_mask"],
|
||||
)
|
||||
|
||||
inputs = preprocesser_call(
|
||||
processor=self.processor,
|
||||
text=additional_inputs.pop("text"),
|
||||
images=additional_inputs.pop("image_inputs"),
|
||||
videos=None,
|
||||
padding=True,
|
||||
truncation=True,
|
||||
return_tensors="pt",
|
||||
max_length=self.dataload_config.get("max_length", 768),
|
||||
)
|
||||
|
||||
action_token_id = self.processor.tokenizer.convert_tokens_to_ids("<|action|>")
|
||||
|
||||
# Gating token types
|
||||
additional_inputs["moe_token_types"] = inputs.input_ids == action_token_id
|
||||
|
||||
inputs.update(additional_inputs)
|
||||
|
||||
inputs["dataset_names"] = ["x2_normal"] * inputs["action_chunk"].shape[0]
|
||||
|
||||
return inputs
|
||||
|
||||
|
||||
def load_lerobot_data(
|
||||
config,
|
||||
lerobot_config,
|
||||
rank=0,
|
||||
world_size=1,
|
||||
seed=42,
|
||||
):
|
||||
"""
|
||||
Load LeRobot dataset with distributed support
|
||||
|
||||
Args:
|
||||
config: Model configuration
|
||||
rank: Current process rank (default: 0)
|
||||
world_size: Total number of processes (default: 1)
|
||||
seed: Random seed for reproducibility (default: 42)
|
||||
|
||||
Returns:
|
||||
dataset: Training dataset
|
||||
train_num: Number of training samples per process
|
||||
sampler: Distributed sampler (None if world_size=1)
|
||||
"""
|
||||
|
||||
# Set seed for reproducibility
|
||||
torch.manual_seed(seed)
|
||||
|
||||
dataset_fps = 50
|
||||
dataload_config = get_data_configs(config["data"])
|
||||
|
||||
delta_timestamps = {
|
||||
# action chunk
|
||||
"action": [t / dataset_fps for t in range(dataload_config.get("action_horizon", 32) - 1)],
|
||||
}
|
||||
batch_size = config.get("batch_size_per_gpu", 8)
|
||||
|
||||
# repo_id = "lerobot/aloha_mobile_cabinet"
|
||||
repo_id = lerobot_config.get("repo_id", "lerobot/aloha_mobile_cabinet")
|
||||
dataset = LeRobotDataset(repo_id, delta_timestamps=delta_timestamps, video_backend="pyav")
|
||||
|
||||
if rank == 0:
|
||||
print(f"Selected episodes: {dataset.episodes}")
|
||||
print(f"Number of episodes selected: {dataset.num_episodes}")
|
||||
print(f"Number of frames selected: {dataset.num_frames}")
|
||||
|
||||
dataset = PreprocessedDataset(dataset, config, dataload_config, seed=seed, rank=rank, world_size=world_size)
|
||||
|
||||
# Calculate samples per process
|
||||
if world_size > 1:
|
||||
# With DistributedSampler, each process gets approximately len(dataset) // world_size samples
|
||||
samples_per_process = len(dataset) // world_size
|
||||
train_num = samples_per_process // batch_size
|
||||
else:
|
||||
train_num = len(dataset) // batch_size
|
||||
|
||||
if rank == 0:
|
||||
print("\n" + "=" * 50)
|
||||
print("LeRobot Data Loading Configuration:")
|
||||
print(f"✦ RANK: {rank}")
|
||||
print(f"✦ WORLD SIZE: {world_size}")
|
||||
print(f"✦ BATCH SIZE PER GPU: {batch_size}")
|
||||
print(f"✦ REPO ID: {repo_id}")
|
||||
print(f"✦ TOTAL DATASET SIZE: {len(dataset)}")
|
||||
if world_size > 1:
|
||||
print(f"✦ SAMPLES PER PROCESS: {samples_per_process}")
|
||||
print(f"✦ BATCHES PER PROCESS: {train_num}")
|
||||
print(f"✦ TOTAL BATCHES (ALL PROCESSES): {train_num * world_size}")
|
||||
else:
|
||||
print(f"✦ TOTAL BATCHES: {train_num}")
|
||||
print(f"✦ SEED: {seed}")
|
||||
print("=" * 50 + "\n")
|
||||
|
||||
return dataset, train_num
|
||||
|
||||
|
||||
def get_distributed_dataloader(dataset, config, rank=0, world_size=1, seed=42, is_train=True):
|
||||
"""
|
||||
Helper function to get distributed dataloader
|
||||
|
||||
Args:
|
||||
dataset: PreprocessedDataset instance
|
||||
config: Configuration dict
|
||||
rank: Current process rank
|
||||
world_size: Total number of processes
|
||||
seed: Random seed
|
||||
is_train: Whether this is for training (affects shuffling)
|
||||
|
||||
Returns:
|
||||
dataloader: Distributed DataLoader
|
||||
sampler: DistributedSampler
|
||||
"""
|
||||
if is_train:
|
||||
return dataset.get_train_dataloader(rank=rank, world_size=world_size, seed=seed)
|
||||
else:
|
||||
return dataset.get_val_dataloader(rank=rank, world_size=world_size)
|
||||
|
||||
|
||||
def get_data_configs(config):
|
||||
default_data_config = {
|
||||
"train_test_split": 0.95,
|
||||
"split_seed": 42,
|
||||
"batch_size": 8,
|
||||
"action_horizon": 21,
|
||||
"action_history_length": 0,
|
||||
"image_horizon": 1,
|
||||
"image_history_length": 0,
|
||||
"left_padding": False,
|
||||
"right_padding": False,
|
||||
"return_first_obs": False,
|
||||
"return_last_obs": False,
|
||||
"randomize_obs_after": None,
|
||||
"datasets": [],
|
||||
"labeled_pathes": [],
|
||||
}
|
||||
data_config = default_data_config | config
|
||||
data_config["action_horizon"] += 1
|
||||
|
||||
return data_config
|
||||
|
||||
class TestDataset(PreprocessedDataset):
|
||||
def __init__(self, dataset, config, dataload_config, seed=42):
|
||||
super().__init__(dataset, config, dataload_config, seed=seed, rank=0, world_size=1)
|
||||
|
||||
def get_dataloader(self):
|
||||
"""
|
||||
Get distributed evaluation dataloader (no shuffling for consistent evaluation)
|
||||
"""
|
||||
|
||||
dataloader = torch.utils.data.DataLoader(
|
||||
self,
|
||||
batch_size=1,
|
||||
collate_fn=DataCollator(self.config, self.dataload_config, self._dataset.meta.stats),
|
||||
)
|
||||
|
||||
return dataloader
|
||||
|
||||
def load_test_dataset(
|
||||
config,
|
||||
lerobot_config,
|
||||
seed=42,
|
||||
episode=0,
|
||||
):
|
||||
"""
|
||||
Load test dataset
|
||||
|
||||
Args:
|
||||
config: Model configuration
|
||||
seed: Random seed for reproducibility (default: 42)
|
||||
|
||||
Returns:
|
||||
dataset: Test dataset
|
||||
"""
|
||||
|
||||
# Set seed for reproducibility
|
||||
torch.manual_seed(seed)
|
||||
|
||||
dataset_fps = 50
|
||||
dataload_config = get_data_configs(config["data"])
|
||||
|
||||
delta_timestamps = {
|
||||
# action chunk
|
||||
"action": [t / dataset_fps for t in range(dataload_config.get("action_horizon", 32) - 1)],
|
||||
}
|
||||
|
||||
repo_id = lerobot_config.get("repo_id", "lerobot/aloha_mobile_cabinet")
|
||||
dataset = LeRobotDataset(repo_id, episodes=[episode], delta_timestamps=delta_timestamps, video_backend="pyav")
|
||||
|
||||
print(f"Selected episodes: {dataset.episodes}")
|
||||
print(f"Number of episodes selected: {dataset.num_episodes}")
|
||||
print(f"Number of frames selected: {dataset.num_frames}")
|
||||
|
||||
dataset = TestDataset(dataset, config, dataload_config, seed=seed)
|
||||
|
||||
return dataset
|
||||
@@ -0,0 +1,560 @@
|
||||
"""
|
||||
Data processing utilities for Wall-X multimodal robotic learning.
|
||||
|
||||
This module provides utilities for preprocessing text, images, and actions
|
||||
for multimodal transformer models in robotic learning tasks.
|
||||
"""
|
||||
|
||||
import re
|
||||
import torch
|
||||
import random
|
||||
from collections import OrderedDict
|
||||
from typing import List, Dict, Any, Optional, Union, Tuple
|
||||
from transformers import BatchFeature
|
||||
|
||||
|
||||
CAMERA_NAME_MAPPING = {
|
||||
"face_view": "front view",
|
||||
"left_wrist_view": "left wrist view",
|
||||
"right_wrist_view": "right wrist view",
|
||||
"move1_view": "move view",
|
||||
"move2_view": "move view",
|
||||
"wall_view": "wall view",
|
||||
"top_view": "top view",
|
||||
}
|
||||
|
||||
|
||||
MULTIMODAL_DATASET_NAMES = [
|
||||
"x2_multimodal_from_action",
|
||||
"x2_multimodal",
|
||||
"x2_subtask_generation",
|
||||
"multimodal_CapsFusion",
|
||||
"multimodal_Robo2VLM",
|
||||
"multimodal_RoboPoint",
|
||||
"multimodal_EQA",
|
||||
"multimodal_Cambrian",
|
||||
"multimodal_pixmo",
|
||||
"multimodal_VQAv2",
|
||||
"multimodal_COCO",
|
||||
]
|
||||
|
||||
FREQUENCY_MAPPING = {
|
||||
"x2_normal": 32,
|
||||
"fractal": 5,
|
||||
"bridge_data_v2": 5,
|
||||
"droid": 15,
|
||||
"agibotworld_alpha": 32,
|
||||
"DobbE": 30,
|
||||
"RH20T": 10,
|
||||
"UMI-biarm": 10,
|
||||
"austin_buds": 20,
|
||||
"austin_sailor": 20,
|
||||
"austin_sirius": 20,
|
||||
"bc_z": 10,
|
||||
"berkeley_autolab_ur5": 5,
|
||||
"berkeley_cable_routing": 10,
|
||||
"berkeley_fanuc_manipulation": 10,
|
||||
"dlr_edan_shared_control": 5,
|
||||
"fmb": 10,
|
||||
"furniture_bench": 10,
|
||||
"jaco_play": 10,
|
||||
"nyu_rot": 10,
|
||||
"stanford_hydra": 10,
|
||||
"stanford_kuka_multimodal": 20,
|
||||
"taco_play": 30,
|
||||
"utaustin_mutex": 20,
|
||||
"viola": 20,
|
||||
}
|
||||
|
||||
|
||||
def preprocesser_call(
|
||||
processor,
|
||||
images: Optional[Union[List, Any]] = None,
|
||||
text: Optional[Union[str, List[str]]] = None,
|
||||
videos: Optional[Union[List, Any]] = None,
|
||||
padding: Union[bool, str] = False,
|
||||
truncation: Optional[bool] = None,
|
||||
max_length: Optional[int] = None,
|
||||
return_tensors: str = "pt",
|
||||
) -> BatchFeature:
|
||||
"""Unified preprocessing function for Wall-X model handling text, image and video inputs.
|
||||
|
||||
Processes inputs into format suitable for multimodal transformer models, including:
|
||||
- Text tokenization and special token handling
|
||||
- Image/video processing through image processor
|
||||
- Attention mask and label generation
|
||||
- Padding and truncation handling
|
||||
|
||||
Args:
|
||||
processor: Multimodal processor containing tokenizer and image processor
|
||||
images: Input images (PIL, numpy arrays, or torch tensors)
|
||||
text: Text or list of texts to tokenize
|
||||
videos: Input videos (numpy arrays or torch tensors)
|
||||
padding: Whether to pad sequences to same length
|
||||
truncation: Whether to truncate sequences longer than max_length
|
||||
max_length: Maximum length for truncation/padding
|
||||
return_tensors: Format for returned tensors ('pt', 'np', etc.)
|
||||
|
||||
Returns:
|
||||
BatchFeature containing processed inputs with keys:
|
||||
- input_ids: Tokenized text
|
||||
- attention_mask: Attention mask for text
|
||||
- pixel_values: Processed image pixels
|
||||
- pixel_values_videos: Processed video frames
|
||||
- image_grid_thw: Image grid dimensions for LLM
|
||||
- video_grid_thw: Video grid dimensions for LLM
|
||||
- labels: Training labels with masking
|
||||
"""
|
||||
# Process image inputs
|
||||
if images is not None and len(images) > 0:
|
||||
image_inputs = processor.image_processor(
|
||||
images=images, videos=None, return_tensors=return_tensors
|
||||
)
|
||||
image_grid_thw = image_inputs["image_grid_thw"]
|
||||
else:
|
||||
image_inputs = {}
|
||||
image_grid_thw = None
|
||||
|
||||
# Process video inputs
|
||||
if videos is not None:
|
||||
videos_inputs = processor.image_processor(
|
||||
images=None, videos=videos, return_tensors=return_tensors
|
||||
)
|
||||
video_grid_thw = videos_inputs["video_grid_thw"]
|
||||
else:
|
||||
videos_inputs = {}
|
||||
video_grid_thw = None
|
||||
|
||||
# Ensure text input is in list format
|
||||
if not isinstance(text, list):
|
||||
text = [text]
|
||||
|
||||
# Process image placeholder tokens in text
|
||||
if image_grid_thw is not None:
|
||||
merge_length = processor.image_processor.merge_size ** 2
|
||||
index = 0
|
||||
for i in range(len(text)):
|
||||
while "<|image_pad|>" in text[i]:
|
||||
# Add bounds checking to avoid index overflow
|
||||
if index >= len(image_grid_thw):
|
||||
print(f"Warning: Number of image placeholders ({index + 1}) "
|
||||
f"exceeds actual images ({len(image_grid_thw)}), "
|
||||
f"skipping remaining placeholder processing")
|
||||
break
|
||||
# Replace image placeholder with actual token count
|
||||
token_count = image_grid_thw[index].prod() // merge_length
|
||||
text[i] = text[i].replace(
|
||||
"<|image_pad|>", "<|placeholder|>" * token_count, 1
|
||||
)
|
||||
index += 1
|
||||
text[i] = text[i].replace("<|placeholder|>", "<|image_pad|>")
|
||||
|
||||
# Process video placeholder tokens in text
|
||||
if video_grid_thw is not None:
|
||||
merge_length = processor.image_processor.merge_size ** 2
|
||||
index = 0
|
||||
for i in range(len(text)):
|
||||
while "<|video_pad|>" in text[i]:
|
||||
# Replace video placeholder with actual token count
|
||||
token_count = video_grid_thw[index].prod() // merge_length
|
||||
text[i] = text[i].replace(
|
||||
"<|video_pad|>", "<|placeholder|>" * token_count, 1
|
||||
)
|
||||
index += 1
|
||||
text[i] = text[i].replace("<|placeholder|>", "<|video_pad|>")
|
||||
|
||||
# Tokenize complete input text
|
||||
text_inputs = processor.tokenizer(
|
||||
text,
|
||||
return_tensors=return_tensors,
|
||||
padding=padding,
|
||||
truncation=truncation,
|
||||
max_length=max_length
|
||||
)
|
||||
|
||||
# Get pad token ID for label generation
|
||||
pad_token_id = processor.tokenizer.pad_token_id
|
||||
if pad_token_id is None:
|
||||
pad_token_id = processor.tokenizer.eos_token_id
|
||||
|
||||
# Generate labels for multi-turn dialogue, keeping only assistant response loss
|
||||
labels = torch.full_like(text_inputs.input_ids, -100)
|
||||
assistant_marker = "<|im_start|>assistant\n"
|
||||
im_end_token_id = processor.tokenizer.convert_tokens_to_ids("<|im_end|>")
|
||||
assistant_tokens = processor.tokenizer(
|
||||
"<|im_start|>assistant\n", add_special_tokens=False
|
||||
).input_ids
|
||||
|
||||
for i in range(len(text)):
|
||||
assistant_regions = []
|
||||
parts = text[i].split(assistant_marker)
|
||||
|
||||
# Process each part to determine which tokens belong to assistant responses
|
||||
# Count left padding tokens
|
||||
num_left_pads = 0
|
||||
for token_id in text_inputs.input_ids[i]:
|
||||
if token_id == pad_token_id:
|
||||
num_left_pads += 1
|
||||
else:
|
||||
break
|
||||
current_pos = num_left_pads
|
||||
|
||||
for j, part in enumerate(parts):
|
||||
part_tokens = processor.tokenizer(part, add_special_tokens=False).input_ids
|
||||
if j == 0:
|
||||
# First part is system prompt or user question, all labels are -100
|
||||
current_pos += len(part_tokens)
|
||||
continue
|
||||
|
||||
# From second part onwards, each part starts with assistant response
|
||||
for k in range(current_pos + 1, len(text_inputs.input_ids[i])):
|
||||
if text_inputs.input_ids[i][k] == im_end_token_id:
|
||||
assistant_regions.append((
|
||||
current_pos + len(assistant_tokens), k + 2
|
||||
))
|
||||
break
|
||||
current_pos += len(part_tokens) + 3
|
||||
|
||||
# Set labels for assistant response regions
|
||||
for start, end in assistant_regions:
|
||||
labels[i][start:end] = text_inputs.input_ids[i][start:end]
|
||||
|
||||
# Mask special action tokens in labels
|
||||
action_token_id = processor.tokenizer.encode("<|action|>")[0]
|
||||
propri_token_id = processor.tokenizer.encode("<|propri|>")[0]
|
||||
labels[labels == action_token_id] = -100
|
||||
labels[labels == propri_token_id] = -100
|
||||
labels[labels == processor.tokenizer.pad_token_id] = -100
|
||||
|
||||
# Set labels to None if all are invalid to skip cross entropy loss
|
||||
if (labels != -100).any().item():
|
||||
text_inputs["labels"] = labels
|
||||
else:
|
||||
text_inputs["labels"] = None
|
||||
|
||||
return BatchFeature(data={**text_inputs, **image_inputs, **videos_inputs})
|
||||
|
||||
|
||||
def process_grounding_points(text: str, orig_height: int, orig_width: int, resized_height: int, resized_width: int, model_type: str) -> str:
|
||||
"""Process grounding point coordinates in text based on image resizing.
|
||||
|
||||
Adjusts coordinate values in <point> tags to match resized image dimensions
|
||||
for different model types (qwen2, qwen2_5).
|
||||
|
||||
Args:
|
||||
text: Input text containing <point> tags with coordinates
|
||||
orig_height: Original image height
|
||||
orig_width: Original image width
|
||||
resized_height: Resized image height
|
||||
resized_width: Resized image width
|
||||
model_type: Model type for coordinate processing ('qwen2' or 'qwen2_5')
|
||||
|
||||
Returns:
|
||||
Text with adjusted coordinate values
|
||||
"""
|
||||
# Regex pattern to match <point> tags and their contents
|
||||
point_pattern = re.compile(r"<point>(.*?)</point>")
|
||||
|
||||
def process_match(match):
|
||||
"""Process a single point match and adjust coordinates."""
|
||||
coords_str = match.group(1)
|
||||
try:
|
||||
# Extract coordinates from string
|
||||
coords = list(map(int, re.findall(r"\d+", coords_str)))
|
||||
|
||||
# Calculate resize scale factors
|
||||
scale_w = resized_width / orig_width
|
||||
scale_h = resized_height / orig_height
|
||||
|
||||
if len(coords) == 2:
|
||||
x, y = coords
|
||||
if model_type == "qwen2_5":
|
||||
# Qwen2.5 uses pixel coordinates
|
||||
new_x = max(0, min(round(x * scale_w), resized_width - 1))
|
||||
new_y = max(0, min(round(y * scale_h), resized_height - 1))
|
||||
elif model_type == "qwen2":
|
||||
# Qwen2 normalizes to [0, 1000) range
|
||||
new_x = max(0, min(999.999, (x / orig_width) * 1000))
|
||||
new_y = max(0, min(999.999, (y / orig_height) * 1000))
|
||||
else:
|
||||
raise ValueError(f"Unsupported model type: {model_type}")
|
||||
coords = [new_x, new_y]
|
||||
|
||||
elif len(coords) == 4:
|
||||
x1, y1, x2, y2 = coords
|
||||
if model_type == "qwen2_5":
|
||||
new_x1 = max(0, min(round(x1 * scale_w), resized_width - 1))
|
||||
new_y1 = max(0, min(round(y1 * scale_h), resized_height - 1))
|
||||
new_x2 = max(0, min(round(x2 * scale_w), resized_width - 1))
|
||||
new_y2 = max(0, min(round(y2 * scale_h), resized_height - 1))
|
||||
elif model_type == "qwen2":
|
||||
new_x1 = max(0, min(999.999, (x1 / orig_width) * 1000))
|
||||
new_y1 = max(0, min(999.999, (y1 / orig_height) * 1000))
|
||||
new_x2 = max(0, min(999.999, (x2 / orig_width) * 1000))
|
||||
new_y2 = max(0, min(999.999, (y2 / orig_height) * 1000))
|
||||
else:
|
||||
raise ValueError(f"Unsupported model type: {model_type}")
|
||||
coords = [new_x1, new_y1, new_x2, new_y2]
|
||||
|
||||
# Return processed point tag
|
||||
return f'<point>[{", ".join(map(str, coords))}]</point>'
|
||||
|
||||
except (ValueError, TypeError):
|
||||
# Return original content if processing fails
|
||||
return match.group(0)
|
||||
|
||||
# Replace all matching point tags
|
||||
processed_text = point_pattern.sub(process_match, text)
|
||||
return processed_text
|
||||
|
||||
|
||||
def get_frame_instruction(
|
||||
instruction_info: Dict[str, Any], frame_idx: Optional[int] = None, truncate_keys: Optional[List[str]] = None
|
||||
) -> Tuple[Dict[str, Any], Optional[int]]:
|
||||
"""Extract frame-specific instruction from instruction dictionary.
|
||||
|
||||
Args:
|
||||
instruction_info: Dictionary containing instruction components
|
||||
frame_idx: Current frame index
|
||||
truncate_keys: Keys that trigger truncation when found
|
||||
|
||||
Returns:
|
||||
Tuple of (frame_instruction_dict, split_end_frame)
|
||||
"""
|
||||
if truncate_keys is None:
|
||||
truncate_keys = ["subtask_generation", "distribute", "subtask_generation_zh", "distribute_zh"]
|
||||
|
||||
instruction_for_frame = {}
|
||||
split_end = None
|
||||
|
||||
for key, value in instruction_info.items():
|
||||
if isinstance(value, dict):
|
||||
# Handle frame-range specific instructions
|
||||
for frame_range, frame_instruction in value.items():
|
||||
start_frame, end_frame = map(int, frame_range.split(" "))
|
||||
if start_frame <= frame_idx < end_frame or (start_frame == frame_idx):
|
||||
instruction_for_frame[key] = frame_instruction
|
||||
if truncate_keys is not None and split_end is None and key in truncate_keys:
|
||||
split_end = end_frame + 1
|
||||
break
|
||||
else:
|
||||
instruction_for_frame[key] = value
|
||||
|
||||
return instruction_for_frame, split_end
|
||||
|
||||
|
||||
def get_task_instruction(frame_instruction_info: Dict[str, Any], priority_order: Optional[OrderedDict] = None) -> str:
|
||||
"""Construct task instruction from available instruction fields using priority sampling.
|
||||
|
||||
Args:
|
||||
frame_instruction_info: Dictionary containing instruction fields
|
||||
priority_order: OrderedDict specifying sampling probability for each field
|
||||
|
||||
Returns:
|
||||
Combined instruction string with priority components
|
||||
"""
|
||||
# Default priority settings
|
||||
default_priority_order = OrderedDict(
|
||||
{
|
||||
"subtask_generation": 0.25,
|
||||
"subtask_generation_zh": 0.25,
|
||||
"distribute": 0.25,
|
||||
"distribute_zh": 0.25,
|
||||
}
|
||||
)
|
||||
|
||||
if priority_order is not None:
|
||||
priority_order = OrderedDict(priority_order)
|
||||
else:
|
||||
priority_order = default_priority_order
|
||||
|
||||
got_instruction = False
|
||||
task_instruction = ""
|
||||
|
||||
# Sample instruction components based on priority probabilities
|
||||
for key, prob in priority_order.items():
|
||||
if key in frame_instruction_info and frame_instruction_info[key] != "":
|
||||
if got_instruction:
|
||||
if random.random() >= prob:
|
||||
continue
|
||||
|
||||
task_instruction += f"\n{frame_instruction_info[key]}"
|
||||
got_instruction = True
|
||||
break
|
||||
|
||||
# Fall back to base instruction if no priority components found
|
||||
if not got_instruction:
|
||||
task_instruction = frame_instruction_info.get("instruction", "")
|
||||
|
||||
return task_instruction
|
||||
|
||||
|
||||
def get_wallx_normal_text(
|
||||
instruction_info: Dict[str, Any],
|
||||
action_chunk_size: int,
|
||||
frame_idx: int,
|
||||
priority_order: Optional[OrderedDict] = None,
|
||||
cam_mapping: Optional[Dict[str, str]] = None,
|
||||
generate_subtask_ratio: float = 0.0,
|
||||
) -> Tuple[str, bool]:
|
||||
"""Construct complete multimodal prompt text for Wall-X model.
|
||||
|
||||
Formats input using special tokens including:
|
||||
- System message
|
||||
- User observations (with image placeholders)
|
||||
- Task instructions
|
||||
- Proprioception prompts
|
||||
- Assistant responses (with action tokens)
|
||||
|
||||
Args:
|
||||
instruction_info: Dictionary containing instruction components
|
||||
action_chunk_size: Number of action tokens to generate
|
||||
frame_idx: Current frame index
|
||||
priority_order: Priority order for instruction sampling
|
||||
cam_mapping: Camera name mapping dictionary
|
||||
generate_subtask_ratio: Probability of generating subtask instead of actions
|
||||
|
||||
Returns:
|
||||
Tuple of (formatted_prompt_text, is_subtask_generation)
|
||||
"""
|
||||
# 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|>"
|
||||
|
||||
# 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 cam_mapping:
|
||||
for _, cam_name in cam_mapping.items():
|
||||
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:"
|
||||
|
||||
# Get frame-specific instruction
|
||||
frame_instruction_info, _ = get_frame_instruction(instruction_info, frame_idx=frame_idx)
|
||||
|
||||
generate_subtask = False
|
||||
priority_keys = ["subtask_generation", "distribute"]
|
||||
|
||||
# Decide whether to generate subtask or actions
|
||||
if bool(set(frame_instruction_info.keys()) & set(priority_keys)) and random.random() < generate_subtask_ratio:
|
||||
# Generate subtask (equivalent to VQA task)
|
||||
instruction = frame_instruction_info.get("instruction", "")
|
||||
text_prompt = "\nPredict the next action in language.\n"
|
||||
user_message = f"{user_request} {instruction}{text_prompt}{role_end_symbol}\n"
|
||||
|
||||
# Find output instruction from priority keys
|
||||
for key in priority_keys:
|
||||
if key in frame_instruction_info:
|
||||
output_instruction = frame_instruction_info[key]
|
||||
break
|
||||
|
||||
assistant_output = f"{role_start_symbol}assistant\n{output_instruction}\n{role_end_symbol}"
|
||||
generate_subtask = True
|
||||
else:
|
||||
# Generate actions
|
||||
instruction = get_task_instruction(frame_instruction_info, priority_order=priority_order)
|
||||
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{action_fast_symbol}{role_end_symbol}\n{action_symbol * action_chunk_size}"
|
||||
|
||||
complete_text = prologue + user_message + assistant_output
|
||||
return complete_text, generate_subtask
|
||||
|
||||
|
||||
def get_action_tokens(normalized_actions: Union[torch.Tensor, List], action_tokenizer) -> List[List[str]]:
|
||||
"""Convert normalized actions to action token strings.
|
||||
|
||||
Args:
|
||||
normalized_actions: Normalized action arrays/tensors
|
||||
action_tokenizer: Tokenizer for converting actions to tokens
|
||||
|
||||
Returns:
|
||||
List of action token string lists for each sample
|
||||
"""
|
||||
if isinstance(normalized_actions, torch.Tensor):
|
||||
normalized_actions = normalized_actions.cpu().numpy()
|
||||
|
||||
all_action_tokens = []
|
||||
for i in range(len(normalized_actions)):
|
||||
if isinstance(normalized_actions[i], torch.Tensor):
|
||||
normalized_actions[i] = normalized_actions[i].cpu().numpy()
|
||||
|
||||
token_id = action_tokenizer(normalized_actions[i])
|
||||
action_tokens = [f"<|action_token_{j}|>" for j in token_id[0]]
|
||||
all_action_tokens.append(action_tokens)
|
||||
|
||||
return all_action_tokens
|
||||
|
||||
|
||||
def pad_action_token_strs(actions_token_lists: List[List[str]], pad_token: str = "<|endoftext|>") -> List[str]:
|
||||
"""Pad action token lists to same length and join as strings.
|
||||
|
||||
Args:
|
||||
actions_token_lists: List of action token lists for each sample
|
||||
pad_token: Token used for padding
|
||||
|
||||
Returns:
|
||||
List of padded action token strings
|
||||
"""
|
||||
max_len = max(len(tokens) for tokens in actions_token_lists)
|
||||
padded_action_strs = []
|
||||
|
||||
for tokens in actions_token_lists:
|
||||
padded_tokens = tokens + ["<|im_end|>\n"] + [pad_token] * (max_len - len(tokens))
|
||||
padded_action_strs.append("".join(padded_tokens))
|
||||
|
||||
return padded_action_strs
|
||||
|
||||
|
||||
def replace_action_token(
|
||||
text: List[str], norm_action: Optional[torch.Tensor], action_tokenizer, dataset_names: List[str], dof_masks: Optional[torch.Tensor] = None
|
||||
) -> List[str]:
|
||||
"""Replace action placeholders in text with actual action tokens.
|
||||
|
||||
Args:
|
||||
text: List of text strings with action placeholders
|
||||
norm_action: Normalized action tensors
|
||||
action_tokenizer: Tokenizer for converting actions to tokens
|
||||
dataset_names: Names of datasets for each sample
|
||||
dof_masks: Masks for degrees of freedom
|
||||
|
||||
Returns:
|
||||
List of text strings with action tokens replaced
|
||||
"""
|
||||
# Filter out multimodal dataset names
|
||||
dataset_names = [name for name in dataset_names if name not in MULTIMODAL_DATASET_NAMES]
|
||||
|
||||
# Get required action chunk sizes
|
||||
required_chunk_sizes = [FREQUENCY_MAPPING.get(name, 32) for name in dataset_names]
|
||||
|
||||
if action_tokenizer is not None and norm_action is not None:
|
||||
# Extract actions based on chunk sizes and DOF masks
|
||||
norm_action = [action[: required_chunk_sizes[i], dof_masks[i, 0].bool()] for i, action in enumerate(norm_action)]
|
||||
|
||||
# Convert to action tokens and pad
|
||||
actions_fast_tokens = get_action_tokens(norm_action, action_tokenizer)
|
||||
actions_fast_token_strs = pad_action_token_strs(actions_fast_tokens)
|
||||
|
||||
# Replace action placeholders with actual tokens
|
||||
actions_fast_token_idx = 0
|
||||
for i in range(len(text)):
|
||||
if "<|action_fast|>" in text[i]:
|
||||
text[i] = text[i].replace("<|action_fast|><|im_end|>\n", actions_fast_token_strs[actions_fast_token_idx])
|
||||
actions_fast_token_idx += 1
|
||||
|
||||
# Remove remaining action placeholders
|
||||
text = [t.replace("<|action|>", "") for t in text]
|
||||
else:
|
||||
# Remove action placeholders when no tokenizer available
|
||||
text = [t.replace("<|action_fast|><|im_end|>\n", "") for t in text]
|
||||
|
||||
return text
|
||||
Reference in New Issue
Block a user