huangyanpei/VLA
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

[lint] Update lint (#16)

Browse Source 
* update lint

* update readme

* update ruff lint
This commit is contained in:
Lufang Chen
2025-09-11 13:18:33 +08:00
committed by GitHub
parent a89dce95aa
commit e9332a283d
28 changed files with 2406 additions and 1074 deletions
Download Patch File Download Diff File Expand all files Collapse all files
wall_x/data/config.py
+58 -29
View File
@@ -5,57 +5,84 @@ 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"
"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"
"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"
"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
"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
@@ -63,27 +90,29 @@ class X2RDataProcessingConfig:
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}")
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':
def update(self, **kwargs) -> "X2RDataProcessingConfig":
"""Update configuration parameters.
Args:
**kwargs: Key-value pairs to update
Returns:
X2RDataProcessingConfig: Updated configuration instance
"""
@@ -92,4 +121,4 @@ class X2RDataProcessingConfig:
setattr(self, key, value)
else:
raise ValueError(f"Unknown configuration parameter: {key}")
return self
return self
wall_x/data/load_lerobot_dataset.py
+121 -30
View File
@@ -8,7 +8,12 @@ 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 wall_x.data.utils import (
process_grounding_points,
get_wallx_normal_text,
replace_action_token,
preprocesser_call,
)
from transformers import AutoProcessor
@@ -67,7 +72,9 @@ class PreprocessedDataset(Dataset[T_co]):
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)
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
@@ -108,7 +115,9 @@ class PreprocessedDataset(Dataset[T_co]):
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)
text = process_grounding_points(
complete_text, h, w, resize_h, resize_w, self.data_config.model_type
)
result = {
"image_inputs": image_inputs,
"text": text,
@@ -150,7 +159,9 @@ class PreprocessedDataset(Dataset[T_co]):
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),
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
@@ -164,7 +175,9 @@ class PreprocessedDataset(Dataset[T_co]):
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))
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)
@@ -181,7 +194,9 @@ class PreprocessedDataset(Dataset[T_co]):
batch_size=batch_size,
sampler=sampler,
num_workers=num_workers,
collate_fn=DataCollator(self.config, self.dataload_config, self._dataset.meta.stats),
collate_fn=DataCollator(
self.config, self.dataload_config, self._dataset.meta.stats
),
pin_memory=True,
persistent_workers=num_workers > 0,
prefetch_factor=2,
@@ -212,19 +227,30 @@ class DataCollator:
# 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)
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 self.use_fast_tokenizer and 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)
if (
self.use_fast_tokenizer
and 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]
if not self.use_fast_tokenizer:
self.train_action_tokenizer = None
else:
self.train_action_tokenizer = self._action_tokenizer_cache[action_tokenizer_path]
self.train_action_tokenizer = self._action_tokenizer_cache[
action_tokenizer_path
]
if self.use_fast_tokenizer:
self.action_mapper = {}
@@ -254,9 +280,27 @@ class DataCollator:
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 = 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
[
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
@@ -268,18 +312,42 @@ class DataCollator:
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)
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]
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])
additional_inputs["frame_index"] = torch.stack(
[item["frame_index"] for item in batch]
)
else:
raise NotImplementedError(f"{key} input not implemented in preprocesser")
raise NotImplementedError(
f"{key} input not implemented in preprocesser"
)
additional_inputs["text"] = replace_action_token(
additional_inputs["text"],
@@ -342,20 +410,27 @@ def load_lerobot_data(
delta_timestamps = {
# action chunk
"action": [t / dataset_fps for t in range(dataload_config.get("action_horizon", 32) - 1)],
"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")
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)
dataset = PreprocessedDataset(
dataset, config, dataload_config, seed=seed, rank=rank, world_size=world_size
)
# Calculate samples per process
if world_size > 1:
@@ -385,7 +460,9 @@ def load_lerobot_data(
return dataset, train_num
def get_distributed_dataloader(dataset, config, rank=0, world_size=1, seed=42, is_train=True):
def get_distributed_dataloader(
dataset, config, rank=0, world_size=1, seed=42, is_train=True
):
"""
Helper function to get distributed dataloader
@@ -429,23 +506,29 @@ def get_data_configs(config):
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)
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),
collate_fn=DataCollator(
self.config, self.dataload_config, self._dataset.meta.stats
),
)
return dataloader
def load_test_dataset(
config,
lerobot_config,
@@ -471,16 +554,24 @@ def load_test_dataset(
delta_timestamps = {
# action chunk
"action": [t / dataset_fps for t in range(dataload_config.get("action_horizon", 32) - 1)],
"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")
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
return dataset
wall_x/data/utils.py
+82 -31
View File
@@ -15,10 +15,10 @@ from transformers import BatchFeature
CAMERA_NAME_MAPPING = {
"face_view": "front view",
"left_wrist_view": "left wrist view",
"left_wrist_view": "left wrist view",
"right_wrist_view": "right wrist view",
"move1_view": "move view",
"move2_view": "move view",
"move2_view": "move view",
"wall_view": "wall view",
"top_view": "top view",
}
@@ -78,13 +78,13 @@ def preprocesser_call(
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)
@@ -94,7 +94,7 @@ def preprocesser_call(
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
@@ -131,15 +131,17 @@ def preprocesser_call(
# Process image placeholder tokens in text
if image_grid_thw is not None:
merge_length = processor.image_processor.merge_size ** 2
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")
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
@@ -151,7 +153,7 @@ def preprocesser_call(
# Process video placeholder tokens in text
if video_grid_thw is not None:
merge_length = processor.image_processor.merge_size ** 2
merge_length = processor.image_processor.merge_size**2
index = 0
for i in range(len(text)):
while "<|video_pad|>" in text[i]:
@@ -169,7 +171,7 @@ def preprocesser_call(
return_tensors=return_tensors,
padding=padding,
truncation=truncation,
max_length=max_length
max_length=max_length,
)
# Get pad token ID for label generation
@@ -209,9 +211,9 @@ def preprocesser_call(
# 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
))
assistant_regions.append(
(current_pos + len(assistant_tokens), k + 2)
)
break
current_pos += len(part_tokens) + 3
@@ -235,7 +237,14 @@ def preprocesser_call(
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:
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
@@ -309,7 +318,9 @@ def process_grounding_points(text: str, orig_height: int, orig_width: int, resiz
def get_frame_instruction(
instruction_info: Dict[str, Any], frame_idx: Optional[int] = None, truncate_keys: Optional[List[str]] = None
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.
@@ -322,7 +333,12 @@ def get_frame_instruction(
Tuple of (frame_instruction_dict, split_end_frame)
"""
if truncate_keys is None:
truncate_keys = ["subtask_generation", "distribute", "subtask_generation_zh", "distribute_zh"]
truncate_keys = [
"subtask_generation",
"distribute",
"subtask_generation_zh",
"distribute_zh",
]
instruction_for_frame = {}
split_end = None
@@ -334,7 +350,11 @@ def get_frame_instruction(
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:
if (
truncate_keys is not None
and split_end is None
and key in truncate_keys
):
split_end = end_frame + 1
break
else:
@@ -343,7 +363,9 @@ def get_frame_instruction(
return instruction_for_frame, split_end
def get_task_instruction(frame_instruction_info: Dict[str, Any], priority_order: Optional[OrderedDict] = None) -> str:
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:
@@ -428,7 +450,9 @@ def get_wallx_normal_text(
action_fast_symbol = "<|action_fast|>"
# System prologue
prologue = f"{role_start_symbol}system\nYou are a helpful assistant.{role_end_symbol}\n"
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:"
@@ -439,13 +463,18 @@ def get_wallx_normal_text(
user_request += "\nInstruction:"
# Get frame-specific instruction
frame_instruction_info, _ = get_frame_instruction(instruction_info, frame_idx=frame_idx)
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:
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"
@@ -457,11 +486,15 @@ def get_wallx_normal_text(
output_instruction = frame_instruction_info[key]
break
assistant_output = f"{role_start_symbol}assistant\n{output_instruction}\n{role_end_symbol}"
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)
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}"
@@ -470,7 +503,9 @@ def get_wallx_normal_text(
return complete_text, generate_subtask
def get_action_tokens(normalized_actions: Union[torch.Tensor, List], action_tokenizer) -> List[List[str]]:
def get_action_tokens(
normalized_actions: Union[torch.Tensor, List], action_tokenizer
) -> List[List[str]]:
"""Convert normalized actions to action token strings.
Args:
@@ -495,7 +530,9 @@ def get_action_tokens(normalized_actions: Union[torch.Tensor, List], action_toke
return all_action_tokens
def pad_action_token_strs(actions_token_lists: List[List[str]], pad_token: str = "<|endoftext|>") -> List[str]:
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:
@@ -509,14 +546,20 @@ def pad_action_token_strs(actions_token_lists: List[List[str]], pad_token: str =
padded_action_strs = []
for tokens in actions_token_lists:
padded_tokens = tokens + ["<|im_end|>\n"] + [pad_token] * (max_len - len(tokens))
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
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.
@@ -531,14 +574,19 @@ def replace_action_token(
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]
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)]
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)
@@ -548,7 +596,10 @@ def replace_action_token(
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])
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