[lint] Update lint (#16)

* update lint

* update readme

* update ruff lint

.pre-commit-config.yaml

@@ -0,0 +1,27 @@
+# See https://pre-commit.com for more information
+# See https://pre-commit.com/hooks.html for more hooks
+exclude: ".git"
+
+repos:
+  - repo: https://github.com/astral-sh/ruff-pre-commit
+    rev: v0.2.2
+    hooks:
+      - id: ruff
+        args: [ --fix, --exit-non-zero-on-fix ]
+
+  - repo: https://github.com/psf/black
+    rev: 24.2.0
+    hooks:
+      - id: black
+
+  - repo: https://github.com/pre-commit/pre-commit-hooks
+    rev: v4.5.0
+    hooks:
+      - id: check-added-large-files
+      - id: check-ast
+      - id: check-case-conflict
+      - id: check-merge-conflict
+      - id: check-toml
+      - id: check-yaml
+      - id: end-of-file-fixer
+      - id: trailing-whitespace

CONTRIBUTING.md

@@ -0,0 +1,17 @@
+# Contributing to Wall-x
+
+## Submit a Pull Request
+
+Before opening a pull request, please make sure your code passes the lint checks.
+
+```bash
+# Install pre-commit hooks (run once)
+pre-commit install
+```
+
+Or
+
+```bash
+# Manually run all checks
+pre-commit run --all-files
+```

csrc/ops.cu

@@ -13,4 +13,3 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
   m.def("rope", &launch_multimodal_rope_forward, "Multimodal RoPE forward kernel");
   m.def("rope_bwd", &launch_multimodal_rope_backward, "Multimodal RoPE backward kernel");
 }
-

csrc/rope.cu

@@ -549,4 +549,3 @@ void launch_multimodal_rope_backward(
         break;
     }
 }
-

pyproject.toml

@@ -0,0 +1,2 @@
+[tool.ruff]
+per-file-ignores = { "__init__.py" = ["F401", "E402"] }

scripts/draw_openloop_plot.py

@@ -1,6 +1,7 @@
 import os
 import yaml
 import torch
+import matplotlib.pyplot as plt
 from wall_x.model.qwen2_5_based.modeling_qwen2_5_vl_act import Qwen2_5_VLMoEForAction
 from wall_x.data.load_lerobot_dataset import load_test_dataset, get_data_configs
 
@@ -8,11 +9,14 @@ from wall_x.data.load_lerobot_dataset import load_test_dataset, get_data_configs
 model_path = "path/to/model"
 action_tokenizer_path = "path/to/action_tokenizer"
 save_dir = "path/to/plot"
-model = Qwen2_5_VLMoEForAction.from_pretrained(model_path, action_tokenizer_path=action_tokenizer_path)
+model = Qwen2_5_VLMoEForAction.from_pretrained(
+    model_path, action_tokenizer_path=action_tokenizer_path
+)
 model.eval()
 model = model.to("cuda")
 model = model.bfloat16()
 
+
 def load_config(config_path):
     """Load configuration from YAML file."""
     with open(config_path, "r") as f:
@@ -22,6 +26,7 @@ def load_config(config_path):
 
     return config
 
+
 # get test dataloader
 path = "path/to/config"
 config = load_config(path)
@@ -46,14 +51,16 @@ for idx, batch in enumerate(dataloader):
                 action_dim=action_dim,
                 pred_horizon=pred_horizon,
                 mode="predict",
-                predict_mode="fast"
+                predict_mode="fast",
             )
-        pred_traj[idx : idx + pred_horizon] = outputs['predict_action'].detach().cpu()
+        pred_traj[idx : idx + pred_horizon] = outputs["predict_action"].detach().cpu()
 
         # Denormalize ground truth actions
-        gt_action_chunk = batch['action_chunk'][:, :, :action_dim]
+        gt_action_chunk = batch["action_chunk"][:, :, :action_dim]
         dof_mask = batch["dof_mask"].to(gt_action_chunk.dtype)
-        denormalized_gt = model.action_preprocessor.normalizer_action.unnormalize_data(gt_action_chunk, ["x2_normal"], dof_mask)
+        denormalized_gt = model.action_preprocessor.normalizer_action.unnormalize_data(
+            gt_action_chunk, ["x2_normal"], dof_mask
+        )
         gt_traj[idx : idx + pred_horizon] = denormalized_gt.detach().cpu()
 
 
@@ -62,20 +69,18 @@ pred_traj_np = pred_traj.numpy()
 
 timesteps = gt_traj.shape[0]
 
-import matplotlib.pyplot as plt
-
 fig, axs = plt.subplots(action_dim, 1, figsize=(15, 5 * action_dim), sharex=True)
-fig.suptitle(f'Action Comparison for lerobot', fontsize=16)
+fig.suptitle("Action Comparison for lerobot", fontsize=16)
 
 for i in range(action_dim):
-    axs[i].plot(range(timesteps), gt_traj_np[:, i], label='Ground Truth')
-    axs[i].plot(range(timesteps), pred_traj_np[:, i], label='Prediction')
-    axs[i].set_ylabel(f'Action Dim {i+1}')
+    axs[i].plot(range(timesteps), gt_traj_np[:, i], label="Ground Truth")
+    axs[i].plot(range(timesteps), pred_traj_np[:, i], label="Prediction")
+    axs[i].set_ylabel(f"Action Dim {i+1}")
     axs[i].legend()
     axs[i].grid(True)
 
-axs[-1].set_xlabel('Timestep')
+axs[-1].set_xlabel("Timestep")
 plt.tight_layout(rect=[0, 0.03, 1, 0.95])
 os.makedirs(save_dir, exist_ok=True)
-plt.savefig(os.path.join(save_dir, f"lerobot_comparison.png"))
+plt.savefig(os.path.join(save_dir, "lerobot_comparison.png"))
 plt.close()

scripts/fake_inference.py

@@ -10,10 +10,14 @@ batch_size = 1
 seq_length = 50
 
 torch.manual_seed(0)
-fake_input_ids = torch.randint(0, len(model.processor.tokenizer), (batch_size, seq_length), dtype=torch.long)
+fake_input_ids = torch.randint(
+    0, len(model.processor.tokenizer), (batch_size, seq_length), dtype=torch.long
+)
 fake_attention_mask = torch.ones((batch_size, seq_length), dtype=torch.long)
 fake_moe_token_types = torch.zeros((batch_size, seq_length), dtype=torch.long)
-fake_position_ids = torch.arange(seq_length, dtype=torch.long).unsqueeze(0).expand(batch_size, -1)
+fake_position_ids = (
+    torch.arange(seq_length, dtype=torch.long).unsqueeze(0).expand(batch_size, -1)
+)
 fake_proprioception = torch.randn((batch_size, 1, 20), dtype=torch.float32)
 fake_agent_pos_mask = torch.ones((batch_size, 1, 20), dtype=torch.float32)
 fake_dof_mask = torch.ones((batch_size, 32, 20), dtype=torch.float32)
@@ -44,7 +48,7 @@ try:
             agent_pos_mask=fake_agent_pos_mask,
             dof_mask=fake_dof_mask,
             dataset_names=fake_dataset_names,
-            mode="validate"
+            mode="validate",
         )
 
     print("✅ Fake inference test successful!")
@@ -68,7 +72,7 @@ try:
     else:
         print("❌ Output contains infinity values")
 
-    print(f"Output logits statistics:")
+    print("Output logits statistics:")
     print(f"  Min value: {outputs.logits.min().item():.4f}")
     print(f"  Max value: {outputs.logits.max().item():.4f}")
     print(f"  Mean: {outputs.logits.mean().item():.4f}")
@@ -77,4 +81,5 @@ try:
 except Exception as e:
     print(f"❌ Fake inference test failed: {e}")
     import traceback
+
     traceback.print_exc()

scripts/merge_tokenizer.py

@@ -8,13 +8,15 @@ use_fast_tokenizer = True
 processor = AutoProcessor.from_pretrained(processor_path, use_fast=True)
 processor.tokenizer.padding_side = "left"
 
-action_tokenizer = AutoProcessor.from_pretrained(action_tokenizer_path, trust_remote_code=True)
+action_tokenizer = AutoProcessor.from_pretrained(
+    action_tokenizer_path, trust_remote_code=True
+)
 
 new_tokens = ["<|propri|>", "<|action|>"]
 new_tokens += [f"<|action_token_{i}|>" for i in range(action_tokenizer.vocab_size)]
 num_added_tokens = processor.tokenizer.add_tokens(new_tokens)
 
-begin_idx_token = f"<|action_token_0|>"
+begin_idx_token = "<|action_token_0|>"
 token_id = processor.tokenizer.convert_tokens_to_ids(begin_idx_token)
 processor.tokenizer.init_kwargs["action_token_start_index"] = token_id
 processor.tokenizer.init_kwargs["action_token_vocab_size"] = action_tokenizer.vocab_size
@@ -22,4 +24,3 @@ processor.tokenizer.init_kwargs["action_token_vocab_size"] = action_tokenizer.vo
 new_tokenizer_dir = "/path/to/new_tokenizer"
 os.makedirs(new_tokenizer_dir, exist_ok=True)
 processor.save_pretrained(new_tokenizer_dir)
-

train_qact.py

@@ -4,7 +4,11 @@ import time
 import yaml
 import wandb
 from argparse import ArgumentParser
-from accelerate import Accelerator, DistributedDataParallelKwargs, DataLoaderConfiguration
+from accelerate import (
+    Accelerator,
+    DistributedDataParallelKwargs,
+    DataLoaderConfiguration,
+)
 
 from wall_x.trainer.qwen_vl_act_trainer import QwenVlAct_Trainer
 
@@ -27,7 +31,9 @@ def load_config(config_path):
 
 def setup_accelerator(config):
     """Initialize and configure the accelerator for distributed training."""
-    print(f"[{time.strftime('%Y-%m-%d %H:%M:%S', time.localtime())}] Preparing accelerator")
+    print(
+        f"[{time.strftime('%Y-%m-%d %H:%M:%S', time.localtime())}] Preparing accelerator"
+    )
 
     ddp_kwargs = DistributedDataParallelKwargs(find_unused_parameters=True)
     accelerator_dataloader_config = DataLoaderConfiguration(dispatch_batches=False)
@@ -36,10 +42,12 @@ def setup_accelerator(config):
         kwargs_handlers=[ddp_kwargs],
         mixed_precision="bf16",
         dataloader_config=accelerator_dataloader_config,
-        gradient_accumulation_steps=config.get("gradient_accumulation_steps", 1)
+        gradient_accumulation_steps=config.get("gradient_accumulation_steps", 1),
     )
 
-    print(f"[{time.strftime('%Y-%m-%d %H:%M:%S', time.localtime())}] Accelerator initialization complete")
+    print(
+        f"[{time.strftime('%Y-%m-%d %H:%M:%S', time.localtime())}] Accelerator initialization complete"
+    )
 
     return accelerator
 
@@ -97,10 +105,14 @@ def main(args):
     trainer.fit()
 
 
-if __name__ == '__main__':
+if __name__ == "__main__":
     parser = ArgumentParser(description="Training script for Wall-X model")
-    parser.add_argument("--config", type=str, required=True, help="Path to configuration YAML file")
-    parser.add_argument("--seed", type=int, default=42, help="Random seed for reproducibility")
+    parser.add_argument(
+        "--config", type=str, required=True, help="Path to configuration YAML file"
+    )
+    parser.add_argument(
+        "--seed", type=int, default=42, help="Random seed for reproducibility"
+    )
 
     args = parser.parse_args()
     main(args)

wall_x/data/config.py

@@ -6,24 +6,51 @@ 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_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",
 ]
 
 
@@ -45,7 +72,7 @@ class X2RDataProcessingConfig:
         default_factory=lambda: {
             "face_view": -1,
             "left_wrist_view": 128,
-            "right_wrist_view": 128
+            "right_wrist_view": 128,
         }
     )
 
@@ -68,7 +95,9 @@ class X2RDataProcessingConfig:
         """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.
@@ -78,7 +107,7 @@ class X2RDataProcessingConfig:
         """
         return self.__dict__
 
-    def update(self, **kwargs) -> 'X2RDataProcessingConfig':
+    def update(self, **kwargs) -> "X2RDataProcessingConfig":
         """Update configuration parameters.
 
         Args:

wall_x/data/load_lerobot_dataset.py

@@ -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,9 +506,12 @@ 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):
         """
@@ -441,11 +521,14 @@ class TestDataset(PreprocessedDataset):
         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,11 +554,19 @@ 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}")

wall_x/data/utils.py

@@ -137,9 +137,11 @@ def preprocesser_call(
             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}) "
+                    print(
+                        f"Warning: Number of image placeholders ({index + 1}) "
                         f"exceeds actual images ({len(image_grid_thw)}), "
-                          f"skipping remaining placeholder processing")
+                        f"skipping remaining placeholder processing"
+                    )
                     break
                 # Replace image placeholder with actual token count
                 token_count = image_grid_thw[index].prod() // merge_length
@@ -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

wall_x/fusions/backend.py

@@ -13,11 +13,12 @@ from typing import Tuple, Optional
 import wallx_csrc as backend
 
 
-
-def _allocate_asymmetric_dual_outputs(input_expert0: torch.Tensor,
+def _allocate_asymmetric_dual_outputs(
+    input_expert0: torch.Tensor,
     input_expert1: torch.Tensor,
     weight_expert0: torch.Tensor,
-                                     weight_expert1: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
+    weight_expert1: torch.Tensor,
+) -> Tuple[torch.Tensor, torch.Tensor]:
     """
     Allocate output tensors for asymmetric dual expert GEMM operations.
 
@@ -45,30 +46,38 @@ def _allocate_asymmetric_dual_outputs(input_expert0: torch.Tensor,
     assert weight_expert1.ndim == 2, "Expected 2D tensor for weight_expert1"
 
     # Verify dimension compatibility for matrix multiplication
-    assert input_expert0.size(1) == weight_expert0.size(0), \
-        f"Input expert0 K dimension {input_expert0.size(1)} != weight expert0 K dimension {weight_expert0.size(0)}"
-    assert input_expert1.size(1) == weight_expert1.size(0), \
-        f"Input expert1 K dimension {input_expert1.size(1)} != weight expert1 K dimension {weight_expert1.size(0)}"
+    assert input_expert0.size(1) == weight_expert0.size(
+        0
+    ), f"Input expert0 K dimension {input_expert0.size(1)} != weight expert0 K dimension {weight_expert0.size(0)}"
+    assert input_expert1.size(1) == weight_expert1.size(
+        0
+    ), f"Input expert1 K dimension {input_expert1.size(1)} != weight expert1 K dimension {weight_expert1.size(0)}"
 
     # Calculate output shapes: [m, k] × [k, n] = [m, n]
     m0, n0 = input_expert0.size(0), weight_expert0.size(1)
     m1, n1 = input_expert1.size(0), weight_expert1.size(1)
 
     # Allocate output tensors with matching device and dtype
-    output_expert0 = torch.empty(m0, n0, device=input_expert0.device, dtype=input_expert0.dtype)
-    output_expert1 = torch.empty(m1, n1, device=input_expert1.device, dtype=input_expert1.dtype)
+    output_expert0 = torch.empty(
+        m0, n0, device=input_expert0.device, dtype=input_expert0.dtype
+    )
+    output_expert1 = torch.empty(
+        m1, n1, device=input_expert1.device, dtype=input_expert1.dtype
+    )
 
     return output_expert0, output_expert1
 
 
-def asym_dual_gmm_separated(input_expert0: torch.Tensor,
+def asym_dual_gmm_separated(
+    input_expert0: torch.Tensor,
     input_expert1: torch.Tensor,
     weight_expert0: torch.Tensor,
     weight_expert1: torch.Tensor,
     output_expert0: Optional[torch.Tensor] = None,
     output_expert1: Optional[torch.Tensor] = None,
     trans_a: bool = False,
-                           trans_b: bool = False) -> Tuple[torch.Tensor, torch.Tensor]:
+    trans_b: bool = False,
+) -> Tuple[torch.Tensor, torch.Tensor]:
     """
     Asymmetric dual expert grouped GEMM with separated inputs and outputs.
 
@@ -113,20 +122,26 @@ def asym_dual_gmm_separated(input_expert0: torch.Tensor,
 
     # Call optimized C++ backend kernel
     backend.asym_dual_gmm(
-        input_expert0, input_expert1,
-        weight_expert0, weight_expert1,
-        output_expert0, output_expert1,
-        trans_a, trans_b
+        input_expert0,
+        input_expert1,
+        weight_expert0,
+        weight_expert1,
+        output_expert0,
+        output_expert1,
+        trans_a,
+        trans_b,
     )
 
     return output_expert0, output_expert1
 
 
-def permute(input: torch.Tensor, 
+def permute(
+    input: torch.Tensor,
     indices: torch.Tensor,
     num_out_tokens: int,
     workspace: torch.Tensor,
-            max_expanded_token_num: int) -> torch.Tensor:
+    max_expanded_token_num: int,
+) -> torch.Tensor:
     """
     Permute input tokens according to expert assignment indices for MoE routing.
 
@@ -147,14 +162,18 @@ def permute(input: torch.Tensor,
         This is typically used with top-k expert selection where each token
         can be routed to multiple experts.
     """
-    return backend.permute(input, indices, num_out_tokens, workspace, max_expanded_token_num)
+    return backend.permute(
+        input, indices, num_out_tokens, workspace, max_expanded_token_num
+    )
 
 
-def unpermute(input: torch.Tensor, 
+def unpermute(
+    input: torch.Tensor,
     row_id_map: torch.Tensor,
     prob: torch.Tensor,
     max_tokens: int,
-              num_topK: int) -> torch.Tensor:
+    num_topK: int,
+) -> torch.Tensor:
     """
     Unpermute expert outputs back to original token order with probability weighting.
 
@@ -178,10 +197,12 @@ def unpermute(input: torch.Tensor,
     return backend.unpermute(input, row_id_map, prob, max_tokens, num_topK)
 
 
-def unpermute_bwd(input_bwd: torch.Tensor, 
+def unpermute_bwd(
+    input_bwd: torch.Tensor,
     input_fwd: torch.Tensor,
     row_id_map: torch.Tensor,
-                  prob: Optional[torch.Tensor]) -> torch.Tensor:
+    prob: Optional[torch.Tensor],
+) -> torch.Tensor:
     """
     Backward pass for unpermute operation with gradient flow.
 
@@ -202,18 +223,22 @@ def unpermute_bwd(input_bwd: torch.Tensor,
     """
     # Handle case where probabilities are not provided
     if prob is None:
-        prob = torch.ones([input_bwd.size(0), 1], dtype=torch.float32, device=input_bwd.device)
+        prob = torch.ones(
+            [input_bwd.size(0), 1], dtype=torch.float32, device=input_bwd.device
+        )
 
     return backend.unpermute_bwd(input_bwd, input_fwd, row_id_map, prob)
 
 
-def rope(q: torch.Tensor, 
+def rope(
+    q: torch.Tensor,
     k: torch.Tensor,
     cos: torch.Tensor,
     sin: torch.Tensor,
     q_out: torch.Tensor,
     k_out: torch.Tensor,
-         mrope_section_doubled: bool) -> None:
+    mrope_section_doubled: bool,
+) -> None:
     """
     Apply RoPE (Rotary Position Embedding) to query and key tensors.
 
@@ -237,7 +262,8 @@ def rope(q: torch.Tensor,
     return backend.rope(q, k, cos, sin, q_out, k_out, mrope_section_doubled)
 
 
-def rope_bwd(grad_q_out: torch.Tensor, 
+def rope_bwd(
+    grad_q_out: torch.Tensor,
     grad_k_out: torch.Tensor,
     q: torch.Tensor,
     k: torch.Tensor,
@@ -245,7 +271,8 @@ def rope_bwd(grad_q_out: torch.Tensor,
     sin: torch.Tensor,
     grad_q: torch.Tensor,
     grad_k: torch.Tensor,
-             mrope_section_doubled: bool) -> None:
+    mrope_section_doubled: bool,
+) -> None:
     """
     Backward pass for RoPE operation with gradient computation.
 
@@ -267,4 +294,6 @@ def rope_bwd(grad_q_out: torch.Tensor,
         This function computes the analytical gradient of the RoPE operation,
         which involves the inverse rotation compared to the forward pass.
     """
-    return backend.rope_bwd(grad_q_out, grad_k_out, q, k, cos, sin, grad_q, grad_k, mrope_section_doubled)
+    return backend.rope_bwd(
+        grad_q_out, grad_k_out, q, k, cos, sin, grad_q, grad_k, mrope_section_doubled
+    )

wall_x/fusions/ops.py

@@ -5,7 +5,9 @@ from wall_x.fusions import backend
 
 class AsymmetricDualExpertGemm(torch.autograd.Function):
     @staticmethod
-    def forward(ctx, input_expert0, input_expert1, weight_expert0, weight_expert1, trans_b=False):
+    def forward(
+        ctx, input_expert0, input_expert1, weight_expert0, weight_expert1, trans_b=False
+    ):
         """
         Forward pass for asymmetric dual expert GEMM.
 
@@ -27,14 +29,24 @@ class AsymmetricDualExpertGemm(torch.autograd.Function):
 
         # Dimension validation depends on trans_b
         if trans_b:
-            assert input_expert0.size(1) == weight_expert0.size(1), "Expert 0 dimension mismatch (trans_b=True)"
-            assert input_expert1.size(1) == weight_expert1.size(1), "Expert 1 dimension mismatch (trans_b=True)"
+            assert input_expert0.size(1) == weight_expert0.size(
+                1
+            ), "Expert 0 dimension mismatch (trans_b=True)"
+            assert input_expert1.size(1) == weight_expert1.size(
+                1
+            ), "Expert 1 dimension mismatch (trans_b=True)"
         else:
-            assert input_expert0.size(1) == weight_expert0.size(0), "Expert 0 dimension mismatch (trans_b=False)"
-            assert input_expert1.size(1) == weight_expert1.size(0), "Expert 1 dimension mismatch (trans_b=False)"
+            assert input_expert0.size(1) == weight_expert0.size(
+                0
+            ), "Expert 0 dimension mismatch (trans_b=False)"
+            assert input_expert1.size(1) == weight_expert1.size(
+                0
+            ), "Expert 1 dimension mismatch (trans_b=False)"
 
         # Save tensors and trans_b for backward pass
-        ctx.save_for_backward(input_expert0, input_expert1, weight_expert0, weight_expert1)
+        ctx.save_for_backward(
+            input_expert0, input_expert1, weight_expert0, weight_expert1
+        )
         ctx.trans_b = trans_b
 
         # Allocate output tensors
@@ -43,11 +55,23 @@ class AsymmetricDualExpertGemm(torch.autograd.Function):
         n0 = weight_expert0.size(0) if trans_b else weight_expert0.size(1)
         n1 = weight_expert1.size(0) if trans_b else weight_expert1.size(1)
 
-        output_expert0 = torch.empty(m0, n0, device=input_expert0.device, dtype=input_expert0.dtype)
-        output_expert1 = torch.empty(m1, n1, device=input_expert1.device, dtype=input_expert1.dtype)
+        output_expert0 = torch.empty(
+            m0, n0, device=input_expert0.device, dtype=input_expert0.dtype
+        )
+        output_expert1 = torch.empty(
+            m1, n1, device=input_expert1.device, dtype=input_expert1.dtype
+        )
 
         # Call the backend C++ function
-        backend.asym_dual_gmm_separated(input_expert0, input_expert1, weight_expert0, weight_expert1, output_expert0, output_expert1, trans_b=trans_b)
+        backend.asym_dual_gmm_separated(
+            input_expert0,
+            input_expert1,
+            weight_expert0,
+            weight_expert1,
+            output_expert0,
+            output_expert1,
+            trans_b=trans_b,
+        )
 
         return output_expert0, output_expert1
 
@@ -111,10 +135,18 @@ class AsymmetricDualExpertGemm(torch.autograd.Function):
                 trans_b=False,
             )
 
-        return grad_input_expert0, grad_input_expert1, grad_weight_expert0, grad_weight_expert1, None
+        return (
+            grad_input_expert0,
+            grad_input_expert1,
+            grad_weight_expert0,
+            grad_weight_expert1,
+            None,
+        )
 
 
-def asym_dual_gmm(input_expert0, input_expert1, weight_expert0, weight_expert1, trans_b=False):
+def asym_dual_gmm(
+    input_expert0, input_expert1, weight_expert0, weight_expert1, trans_b=False
+):
     """
     Convenience function for asymmetric dual expert GEMM.
 
@@ -128,7 +160,9 @@ def asym_dual_gmm(input_expert0, input_expert1, weight_expert0, weight_expert1,
     Returns:
         Tuple of (output_expert0, output_expert1)
     """
-    return AsymmetricDualExpertGemm.apply(input_expert0, input_expert1, weight_expert0, weight_expert1, trans_b)
+    return AsymmetricDualExpertGemm.apply(
+        input_expert0, input_expert1, weight_expert0, weight_expert1, trans_b
+    )
 
 
 ################################################################################################
@@ -145,7 +179,13 @@ class PermuteMoE_topK(torch.autograd.Function):
     max_expanded_token_num = 0
 
     @staticmethod
-    def forward(ctx, input_act: torch.Tensor, indices: torch.Tensor, num_out_tokens: int, max_token_num: int):
+    def forward(
+        ctx,
+        input_act: torch.Tensor,
+        indices: torch.Tensor,
+        num_out_tokens: int,
+        max_token_num: int,
+    ):
         """
         indices: for topK=1, indices in a 1-d tensor of shape [num_tokens],
                  otherwise, it's a 2-d tensor of shape [num_tokens, topK]
@@ -160,18 +200,27 @@ class PermuteMoE_topK(torch.autograd.Function):
 
         # Device check
         if input_act.is_cpu:
-            raise RuntimeError("[Error] The input `input_act` of permute_topK op is on the device: CPU!")
+            raise RuntimeError(
+                "[Error] The input `input_act` of permute_topK op is on the device: CPU!"
+            )
         if indices.is_cpu:
-            warnings.warn("The input `indices` of permute_topK op is on the device: CPU!")
-            expert_for_rows = expert_for_rows.cuda()
+            warnings.warn(
+                "The input `indices` of permute_topK op is on the device: CPU!"
+            )
 
         # Shape check
         if input_act.size(0) != indices.size(0):
-            raise RuntimeError(f"[Error] permute_topK op input `indices` shape mismatch! " f"Expect {input_act.size(0)}, but got {indices.size(0)}.")
+            raise RuntimeError(
+                f"[Error] permute_topK op input `indices` shape mismatch! "
+                f"Expect {input_act.size(0)}, but got {indices.size(0)}."
+            )
 
         # Data type check
         if indices.dtype != torch.int32:
-            warnings.warn(f"The data type of the input `indices` of permute_topK op is {indices.dtype}! " "The recommended type is torch.int32.")
+            warnings.warn(
+                f"The data type of the input `indices` of permute_topK op is {indices.dtype}! "
+                "The recommended type is torch.int32."
+            )
             indices = indices.to(torch.int32)
 
         # Contiguous check
@@ -194,7 +243,11 @@ class PermuteMoE_topK(torch.autograd.Function):
             PermuteMoE_topK.workspace_fw = []
 
         permuted_act, row_id_map, PermuteMoE_topK.workspace_fw = backend.permute(
-            input_act, indices, num_out_tokens, PermuteMoE_topK.workspace_fw, PermuteMoE_topK.max_expanded_token_num
+            input_act,
+            indices,
+            num_out_tokens,
+            PermuteMoE_topK.workspace_fw,
+            PermuteMoE_topK.max_expanded_token_num,
         )
 
         ctx.row_id_map = row_id_map
@@ -215,7 +268,9 @@ class PermuteMoE_topK(torch.autograd.Function):
         num_tokens = ctx.num_tokens
         num_topK = ctx.num_topK
 
-        unpermuted_act_grad = backend.unpermute(permuted_act_grad, row_id_map, torch.tensor([]), num_tokens, num_topK)
+        unpermuted_act_grad = backend.unpermute(
+            permuted_act_grad, row_id_map, torch.tensor([]), num_tokens, num_topK
+        )
         return unpermuted_act_grad, None, None, None
 
 
@@ -229,7 +284,12 @@ class PermuteMoE_topK(torch.autograd.Function):
 class UnpermuteMoE_topK(torch.autograd.Function):
 
     @staticmethod
-    def forward(ctx, input_act: torch.Tensor, row_id_map: torch.Tensor, probs: torch.Tensor = None):
+    def forward(
+        ctx,
+        input_act: torch.Tensor,
+        row_id_map: torch.Tensor,
+        probs: torch.Tensor = None,
+    ):
         # Empty input check
         if not input_act.numel():
             ctx.probs = probs
@@ -237,36 +297,51 @@ class UnpermuteMoE_topK(torch.autograd.Function):
 
         # Device check
         if input_act.is_cpu:
-            raise RuntimeError("[Error] The input `input_act` of unpermute_topK op is on the device: CPU!")
+            raise RuntimeError(
+                "[Error] The input `input_act` of unpermute_topK op is on the device: CPU!"
+            )
         if row_id_map.is_cpu:
-            warnings.warn("The input `row_id_map` of unpermute_topK op is on the device: CPU!")
+            warnings.warn(
+                "The input `row_id_map` of unpermute_topK op is on the device: CPU!"
+            )
             row_id_map = row_id_map.cuda()
         if probs is not None and probs.is_cpu:
-            warnings.warn("The input `probs` of unpermute_topK op is on the device: CPU!")
+            warnings.warn(
+                "The input `probs` of unpermute_topK op is on the device: CPU!"
+            )
             probs = probs.cuda()
 
         # Shape check
         if probs is not None and row_id_map.size(0) != probs.size(0) * probs.size(1):
             raise RuntimeError(
-                f"[Error] unpermute_topK op input `probs` shape mismatch! " f"Expect {row_id_map.size(0)}, but got {probs.size(0) * probs.size(1)}."
+                f"[Error] unpermute_topK op input `probs` shape mismatch! "
+                f"Expect {row_id_map.size(0)}, but got {probs.size(0) * probs.size(1)}."
             )
 
         # Data type check
         if row_id_map.dtype != torch.int32:
             warnings.warn(
-                f"The data type of the input `row_id_map` of unpermute_topK op is {row_id_map.dtype}! " "The recommended type is torch.int32."
+                f"The data type of the input `row_id_map` of unpermute_topK op is {row_id_map.dtype}! "
+                "The recommended type is torch.int32."
             )
             row_id_map = row_id_map.to(torch.int32)
         if probs is not None and probs.dtype != torch.float32:
-            warnings.warn(f"The data type of the input `probs` of unpermute_topK op is {probs.dtype}! " "The recommended type is torch.float32.")
+            warnings.warn(
+                f"The data type of the input `probs` of unpermute_topK op is {probs.dtype}! "
+                "The recommended type is torch.float32."
+            )
             probs = probs.to(torch.float32)
 
         # Contiguous check
         if not input_act.is_contiguous():
-            warnings.warn("The input `input_act` of unpermute_topK op is discontiguous!")
+            warnings.warn(
+                "The input `input_act` of unpermute_topK op is discontiguous!"
+            )
             input_act = input_act.contiguous()
         if not row_id_map.is_contiguous():
-            warnings.warn("The input `row_id_map` of unpermute_topK op is discontiguous!")
+            warnings.warn(
+                "The input `row_id_map` of unpermute_topK op is discontiguous!"
+            )
             row_id_map = row_id_map.contiguous()
         if probs is not None and not probs.is_contiguous():
             warnings.warn("The input `probs` of unpermute_topK op is discontiguous!")
@@ -275,7 +350,13 @@ class UnpermuteMoE_topK(torch.autograd.Function):
         num_tokens = probs.size(0) if probs is not None else input_act.size(0)
         num_topK = probs.size(1) if probs is not None else 1
 
-        unpermuted_output = backend.unpermute(input_act, row_id_map, probs if probs is not None else torch.tensor([]), num_tokens, num_topK)
+        unpermuted_output = backend.unpermute(
+            input_act,
+            row_id_map,
+            probs if probs is not None else torch.tensor([]),
+            num_tokens,
+            num_topK,
+        )
 
         ctx.save_for_backward(input_act, row_id_map, probs)
         return unpermuted_output
@@ -293,7 +374,9 @@ class UnpermuteMoE_topK(torch.autograd.Function):
 
         act_grad = None
         if ctx.needs_input_grad[0]:
-            act_grad, prob_grad = backend.unpermute_bwd(unpermuted_act_grad, input_act, row_id_map, probs)
+            act_grad, prob_grad = backend.unpermute_bwd(
+                unpermuted_act_grad, input_act, row_id_map, probs
+            )
 
         if not ctx.needs_input_grad[2]:
             prob_grad = None
@@ -319,19 +402,36 @@ def unpermute(input_act, row_id_map, probs=None):
 class MultimodalRoPE(torch.autograd.Function):
 
     @staticmethod
-    def forward(ctx, q: torch.Tensor, k: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor, mrope_section: list):
+    def forward(
+        ctx,
+        q: torch.Tensor,
+        k: torch.Tensor,
+        cos: torch.Tensor,
+        sin: torch.Tensor,
+        mrope_section: list,
+    ):
 
         # Device check
         if q.is_cpu:
-            raise RuntimeError("[Error] The input `q` of multimodal_rope op is on the device: CPU!")
+            raise RuntimeError(
+                "[Error] The input `q` of multimodal_rope op is on the device: CPU!"
+            )
         if k.is_cpu:
-            raise RuntimeError("[Error] The input `k` of multimodal_rope op is on the device: CPU!")
+            raise RuntimeError(
+                "[Error] The input `k` of multimodal_rope op is on the device: CPU!"
+            )
         if cos.is_cpu:
-            raise RuntimeError("[Error] The input `cos` of multimodal_rope op is on the device: CPU!")
+            raise RuntimeError(
+                "[Error] The input `cos` of multimodal_rope op is on the device: CPU!"
+            )
         if sin.is_cpu:
-            raise RuntimeError("[Error] The input `sin` of multimodal_rope op is on the device: CPU!")
+            raise RuntimeError(
+                "[Error] The input `sin` of multimodal_rope op is on the device: CPU!"
+            )
         if len(mrope_section) != 3:
-            raise RuntimeError("[Error] The input `mrope_section` of multimodal_rope op must be a list of 3 integers!")
+            raise RuntimeError(
+                "[Error] The input `mrope_section` of multimodal_rope op must be a list of 3 integers!"
+            )
 
         # Contiguous check
         if not q.is_contiguous():

wall_x/model/action_head.py

@@ -1,10 +1,10 @@
-
 import math
 import torch
 import torch.nn as nn
 from torch.distributions import Beta
 from wall_x.utils.constant import action_statistic_dof
 
+
 class Normalizer(nn.Module):
     """
     Action data normalizer for multi-robot systems.
@@ -48,14 +48,18 @@ class Normalizer(nn.Module):
             action_statistic[robot_name]["delta"] = all_dof_delta
 
         # Register statistics as non-trainable parameters
-        self.min = nn.ParameterDict({
+        self.min = nn.ParameterDict(
+            {
                 k: nn.Parameter(action_statistic[k]["min"], requires_grad=False)
                 for k in action_statistic.keys()
-        })
-        self.delta = nn.ParameterDict({
+            }
+        )
+        self.delta = nn.ParameterDict(
+            {
                 k: nn.Parameter(action_statistic[k]["delta"], requires_grad=False)
                 for k in action_statistic.keys()
-        })
+            }
+        )
 
     def normalize_data(self, xs, dataset_names):
         """
@@ -159,7 +163,6 @@ class SinusoidalPosEmb(nn.Module):
         return emb
 
 
-
 class ActionProcessor(nn.Module):
     """
     Action sequence processor for robotic control with flow matching.
@@ -208,16 +211,22 @@ class ActionProcessor(nn.Module):
 
         # Initialize data normalizers for actions and proprioception
         self.normalizer_action = Normalizer(action_statistic_dof, config.dof_config)
-        self.normalizer_propri = Normalizer(action_statistic_dof, config.agent_pos_config)
+        self.normalizer_propri = Normalizer(
+            action_statistic_dof, config.agent_pos_config
+        )
 
         # Proprioception projection layer (includes history/current state)
         self.propri_proj = nn.Linear(self.propri_dim * 2, self.hidden_size, bias=False)
 
         # Beta distribution noise scheduler configuration
         noise_scheduler_config = config.noise_scheduler
-        self.beta_alpha = noise_scheduler_config.get('beta_alpha', 1.5)  # Beta distribution α parameter
-        self.beta_beta = noise_scheduler_config.get('beta_beta', 1.0)    # Beta distribution β parameter  
-        self.s = noise_scheduler_config.get('s', 0.999)                  # Scaling factor
+        self.beta_alpha = noise_scheduler_config.get(
+            "beta_alpha", 1.5
+        )  # Beta distribution α parameter
+        self.beta_beta = noise_scheduler_config.get(
+            "beta_beta", 1.0
+        )  # Beta distribution β parameter
+        self.s = noise_scheduler_config.get("s", 0.999)  # Scaling factor
 
         # Initialize Beta distribution for noise scheduling
         alpha_tensor = torch.tensor(self.beta_alpha, dtype=torch.float32).to("cuda")
@@ -228,14 +237,18 @@ class ActionProcessor(nn.Module):
         self.time_embed = SinusoidalPosEmb(config.hidden_size)
 
         # Action embedding network: project to hidden space
-        self.w1 = nn.Linear(self.action_dim * 2, self.hidden_size, bias=False)  # *2 for action + DOF mask
-        self.w2 = nn.Linear(self.hidden_size * 2, self.hidden_size, bias=False)  # *2 for action + time embeddings
+        self.w1 = nn.Linear(
+            self.action_dim * 2, self.hidden_size, bias=False
+        )  # *2 for action + DOF mask
+        self.w2 = nn.Linear(
+            self.hidden_size * 2, self.hidden_size, bias=False
+        )  # *2 for action + time embeddings
         self.w3 = nn.Linear(self.hidden_size, self.hidden_size, bias=False)
         self.act_fn = nn.SiLU()
 
         # Project back to action space for flow matching loss
         self.action_proj_back = nn.Linear(self.hidden_size, self.action_dim, bias=False)
-        self.mse_loss = nn.MSELoss(reduction='none')
+        self.mse_loss = nn.MSELoss(reduction="none")
 
     def sample_time(self, batch_size, device, dtype):
         """
@@ -257,7 +270,9 @@ class ActionProcessor(nn.Module):
         time = (self.s - sample) / self.s
         return time
 
-    def proprioception_proj(self, proprioception, dataset_names=None, dof_mask=None, use_history=False):
+    def proprioception_proj(
+        self, proprioception, dataset_names=None, dof_mask=None, use_history=False
+    ):
         """
         Project proprioceptive data (joint positions, orientations) to hidden space.
 
@@ -271,7 +286,9 @@ class ActionProcessor(nn.Module):
             torch.Tensor: Projected proprioceptive features of shape [batch_size, seq_len, hidden_size]
         """
         # Ensure proper device and dtype alignment
-        proprioception = proprioception.to(device=self.propri_proj.weight.device).to(dtype=self.propri_proj.weight.dtype)
+        proprioception = proprioception.to(device=self.propri_proj.weight.device).to(
+            dtype=self.propri_proj.weight.dtype
+        )
 
         if dof_mask is not None:
             # Concatenate proprioception with DOF mask
@@ -281,7 +298,9 @@ class ActionProcessor(nn.Module):
             else:
                 proprioception = torch.cat([proprioception, dof_mask], dim=-1)
 
-        proprioception = proprioception.to(device=self.propri_proj.weight.device).to(dtype=self.propri_proj.weight.dtype)
+        proprioception = proprioception.to(device=self.propri_proj.weight.device).to(
+            dtype=self.propri_proj.weight.dtype
+        )
         return self.propri_proj(proprioception)
 
     def forward(self, action_chunk, dataset_names, dof_mask=None):
@@ -329,7 +348,11 @@ class ActionProcessor(nn.Module):
         action_embed = self.w1(noisy_action)
 
         # Repeat time embedding for each sequence position
-        time_embed = time_embed.unsqueeze(1).repeat(1, action_embed.shape[1], 1).to(dtype=self.w2.weight.dtype)
+        time_embed = (
+            time_embed.unsqueeze(1)
+            .repeat(1, action_embed.shape[1], 1)
+            .to(dtype=self.w2.weight.dtype)
+        )
 
         # Combine action and temporal embeddings
         concat_embed = torch.cat([action_embed, time_embed], dim=-1)
@@ -364,7 +387,9 @@ class ActionProcessor(nn.Module):
 
         # Broadcast time embeddings to sequence length
         time_embed = time_embed.unsqueeze(1).repeat(1, action_embed.shape[1], 1)
-        time_embed = time_embed.to(device=noisy_action.device).to(dtype=noisy_action.dtype)
+        time_embed = time_embed.to(device=noisy_action.device).to(
+            dtype=noisy_action.dtype
+        )
 
         # Combine embeddings and process through MLP
         concat_embed = torch.cat([action_embed, time_embed], dim=-1)

wall_x/model/qwen2_5_based/__init__.py

@@ -1,2 +1,8 @@
 from .modeling_qwen2_5_vl_act import Qwen2_5_VLMoEModel, Qwen2_5_VLMoEForAction
 from .configuration_qwen2_5_vl import Qwen2_5_VLConfig
+
+__all__ = [
+    "Qwen2_5_VLMoEModel",
+    "Qwen2_5_VLMoEForAction",
+    "Qwen2_5_VLConfig",
+]

wall_x/trainer/qwen_vl_act_trainer.py

@@ -12,13 +12,16 @@ from datetime import datetime
 from torch.optim import AdamW
 from accelerate import Accelerator
 from safetensors.torch import load_file
-from accelerate.utils import DistributedType
 from transformers.optimization import get_cosine_with_min_lr_schedule_with_warmup
 
 from wall_x.utils.timers import Timers
 from wall_x.model.qwen2_5_based import Qwen2_5_VLMoEForAction
 from wall_x.data.config import ACTION_DATASET_NAMES, MULTIMODAL_DATASET_NAMES
-from wall_x.data.load_lerobot_dataset import PreprocessedDataset, get_data_configs, load_lerobot_data
+from wall_x.data.load_lerobot_dataset import (
+    PreprocessedDataset,
+    get_data_configs,
+    load_lerobot_data,
+)
 
 
 def timer(func):
@@ -31,6 +34,7 @@ def timer(func):
     Returns:
         Wrapped function with timing functionality
     """
+
     @wraps(func)
     def wrapper(*args, **kwargs):
         start_time = time.time()
@@ -40,6 +44,7 @@ def timer(func):
             f"\033[92m[current time: {time.strftime('%Y-%m-%d %H:%M:%S', time.localtime())}] Function {func.__name__} took {end_time - start_time:.2f} seconds to execute\033[0m"
         )
         return result
+
     return wrapper
 
 
@@ -87,7 +92,14 @@ class QwenVlAct_Trainer:
     """
 
     @timer
-    def __init__(self, config, logger, accelerator: Accelerator = None, seed=42, data_config_path=None):
+    def __init__(
+        self,
+        config,
+        logger,
+        accelerator: Accelerator = None,
+        seed=42,
+        data_config_path=None,
+    ):
         """
         Initialize the Vision-Language-Action trainer.
 
@@ -139,19 +151,29 @@ class QwenVlAct_Trainer:
         # Distributed training setup
         self.rank = self.accelerator.process_index
         self.world_size = self.accelerator.num_processes
-        print(f"rank {self.accelerator.process_index} after load model memory usage: {torch.cuda.memory_allocated() / 1024 ** 3:.2f} GB", flush=True)
+        print(
+            f"rank {self.accelerator.process_index} after load model memory usage: {torch.cuda.memory_allocated() / 1024 ** 3:.2f} GB",
+            flush=True,
+        )
 
         # Load training data
         self.load_qact_data()
-        print(f"rank {self.accelerator.process_index} after load qact data usage: {torch.cuda.memory_allocated() / 1024 ** 3:.2f} GB", flush=True)
+        print(
+            f"rank {self.accelerator.process_index} after load qact data usage: {torch.cuda.memory_allocated() / 1024 ** 3:.2f} GB",
+            flush=True,
+        )
 
         # Resume from checkpoint if specified
         if "resume" in self.config:
             self.resume_from_checkpoint()
 
         # Initialize special token IDs
-        self.propri_token_id = self.processor.tokenizer.convert_tokens_to_ids("<|propri|>")
-        self.action_token_id = self.processor.tokenizer.convert_tokens_to_ids("<|action|>")
+        self.propri_token_id = self.processor.tokenizer.convert_tokens_to_ids(
+            "<|propri|>"
+        )
+        self.action_token_id = self.processor.tokenizer.convert_tokens_to_ids(
+            "<|action|>"
+        )
 
         # Initialize evaluation metrics
         self.base_l1_loss = None
@@ -162,7 +184,9 @@ class QwenVlAct_Trainer:
 
         # Adjust global step if resuming from checkpoint
         if self.initial_step != 0:
-            self.global_step = self.initial_step // self.config.get("gradient_accumulation_steps", 1)
+            self.global_step = self.initial_step // self.config.get(
+                "gradient_accumulation_steps", 1
+            )
 
     def print_rank0(self, msg, flush=True):
         """
@@ -188,7 +212,9 @@ class QwenVlAct_Trainer:
         self.accelerator.wait_for_everyone()
 
         # Optional validation before training starts
-        if self.config.get("resume", None) is not None and self.config["resume"].get("validate_first", False):
+        if self.config.get("resume", None) is not None and self.config["resume"].get(
+            "validate_first", False
+        ):
             self.val_loop()
             self.accelerator.wait_for_everyone()
 
@@ -227,7 +253,9 @@ class QwenVlAct_Trainer:
             if getattr(self, "train_dataloader", None) is not None:
                 self.train_sampler.set_epoch(epoch)
             else:
-                self.train_dataloader, self.train_sampler = self.dataset.get_train_dataloader()
+                self.train_dataloader, self.train_sampler = (
+                    self.dataset.get_train_dataloader()
+                )
                 self.train_sampler.set_epoch(epoch)
         else:
             self.train_dataloader = self.dataset.get_train_dataloader()
@@ -236,16 +264,23 @@ class QwenVlAct_Trainer:
         grad_accum_steps = self.config.get("gradient_accumulation_steps", 1)
         total = len(self.train_dataloader)
         t0 = time.time()
-        enable_profiling = self.config['profile']
+        enable_profiling = self.config["profile"]
 
         # Optional PyTorch profiler for performance analysis
         if enable_profiling:
             profiler = torch.profiler.profile(
-                activities=[torch.profiler.ProfilerActivity.CPU, torch.profiler.ProfilerActivity.CUDA],
-                schedule=torch.profiler.schedule(wait=self.config['profile_wait_iters'],
-                                                 warmup=self.config['profile_warmup_iters'],
-                                                 active=self.config['profile_active_iters']),
-                on_trace_ready=torch.profiler.tensorboard_trace_handler(self.config['profile_save_path'], worker_name="worker0"),
+                activities=[
+                    torch.profiler.ProfilerActivity.CPU,
+                    torch.profiler.ProfilerActivity.CUDA,
+                ],
+                schedule=torch.profiler.schedule(
+                    wait=self.config["profile_wait_iters"],
+                    warmup=self.config["profile_warmup_iters"],
+                    active=self.config["profile_active_iters"],
+                ),
+                on_trace_ready=torch.profiler.tensorboard_trace_handler(
+                    self.config["profile_save_path"], worker_name="worker0"
+                ),
                 record_shapes=True,
                 profile_memory=True,
                 with_stack=True,
@@ -261,7 +296,14 @@ class QwenVlAct_Trainer:
             for i, batch in enumerate(self.train_dataloader, self.initial_step):
                 # Move batch to device
                 if isinstance(self.dataset, PreprocessedDataset):
-                    batch = {k: v.to(self.accelerator.device, non_blocking=True) if isinstance(v, torch.Tensor) else v for k, v in batch.items()}
+                    batch = {
+                        k: (
+                            v.to(self.accelerator.device, non_blocking=True)
+                            if isinstance(v, torch.Tensor)
+                            else v
+                        )
+                        for k, v in batch.items()
+                    }
 
                 self.timers("data-load").stop()
 
@@ -275,7 +317,10 @@ class QwenVlAct_Trainer:
 
                     # Check for NaN loss
                     if torch.isnan(loss):
-                        print(f"Warning: NaN loss detected in epoch: {epoch}, step: {i}", flush=True)
+                        print(
+                            f"Warning: NaN loss detected in epoch: {epoch}, step: {i}",
+                            flush=True,
+                        )
                         continue
 
                     # Backward pass
@@ -301,31 +346,77 @@ class QwenVlAct_Trainer:
                         lr = self.lr_scheduler.get_last_lr()[0]
 
                         # Gather loss across all processes for logging
-                        train_loss = self.accelerator.gather(loss.detach()).mean().item()
+                        train_loss = (
+                            self.accelerator.gather(loss.detach()).mean().item()
+                        )
                         _log_dict = {
                             "lr": lr,
                             "train_loss": train_loss,
                         }
 
                         # Log component losses
-                        if "cross_entropy_loss" in outputs and outputs.cross_entropy_loss is not None:
-                            _log_dict["cross_entropy_loss"] = self.accelerator.gather(outputs.cross_entropy_loss.detach()).mean().item()
+                        if (
+                            "cross_entropy_loss" in outputs
+                            and outputs.cross_entropy_loss is not None
+                        ):
+                            _log_dict["cross_entropy_loss"] = (
+                                self.accelerator.gather(
+                                    outputs.cross_entropy_loss.detach()
+                                )
+                                .mean()
+                                .item()
+                            )
 
                         if "flow_loss" in outputs and outputs.flow_loss is not None:
-                            _log_dict["flow_loss"] = self.accelerator.gather(outputs.flow_loss.detach()).mean().item()
+                            _log_dict["flow_loss"] = (
+                                self.accelerator.gather(outputs.flow_loss.detach())
+                                .mean()
+                                .item()
+                            )
 
                         # Log per-dataset channel losses
-                        if "channel_loss_dict" in outputs and outputs.channel_loss_dict is not None:
-                            for dataset_name_i in ACTION_DATASET_NAMES + MULTIMODAL_DATASET_NAMES:
-                                count_sum = self.accelerator.gather(outputs.channel_loss_count_dict[dataset_name_i]).sum().item()
+                        if (
+                            "channel_loss_dict" in outputs
+                            and outputs.channel_loss_dict is not None
+                        ):
+                            for dataset_name_i in (
+                                ACTION_DATASET_NAMES + MULTIMODAL_DATASET_NAMES
+                            ):
+                                count_sum = (
+                                    self.accelerator.gather(
+                                        outputs.channel_loss_count_dict[dataset_name_i]
+                                    )
+                                    .sum()
+                                    .item()
+                                )
                                 if count_sum > 0:
-                                    channel_loss = self.accelerator.gather(outputs.channel_loss_dict[dataset_name_i].detach()).sum().item() / count_sum
-                                    _log_dict[f"channel_loss_{dataset_name_i}"] = channel_loss
+                                    channel_loss = (
+                                        self.accelerator.gather(
+                                            outputs.channel_loss_dict[
+                                                dataset_name_i
+                                            ].detach()
+                                        )
+                                        .sum()
+                                        .item()
+                                        / count_sum
+                                    )
+                                    _log_dict[f"channel_loss_{dataset_name_i}"] = (
+                                        channel_loss
+                                    )
 
                             # Log action accuracy for fast tokenizer
-                            if "action_accuracy" in outputs.channel_loss_dict and self.use_fast_tokenizer:
+                            if (
+                                "action_accuracy" in outputs.channel_loss_dict
+                                and self.use_fast_tokenizer
+                            ):
                                 _log_dict["action_accuracy"] = (
-                                    self.accelerator.gather(outputs.channel_loss_dict["action_accuracy"].detach()).mean().item()
+                                    self.accelerator.gather(
+                                        outputs.channel_loss_dict[
+                                            "action_accuracy"
+                                        ].detach()
+                                    )
+                                    .mean()
+                                    .item()
                                 )
 
                         # Log metrics
@@ -334,7 +425,9 @@ class QwenVlAct_Trainer:
 
                         # Log gradient norm
                         if self.logger is not None and self.accelerator.sync_gradients:
-                            self.logger.log({"total_norm": total_norm}, step=self.global_step)
+                            self.logger.log(
+                                {"total_norm": total_norm}, step=self.global_step
+                            )
 
                 self.timers("interval-time").stop()
 
@@ -343,12 +436,13 @@ class QwenVlAct_Trainer:
                     self.timers("interval-time", log_level=0).start(barrier=False)
                     self.timers("data-load", log_level=0).start(barrier=False)
 
-
                 # Periodic logging
                 t1 = time.time()
                 if i % 1 == 0:
                     lr = self.lr_scheduler.get_last_lr()[0]
-                    self.training_log(epoch, self.num_epoch, i, total, loss, lr, t1 - t0)
+                    self.training_log(
+                        epoch, self.num_epoch, i, total, loss, lr, t1 - t0
+                    )
                     t0 = time.time()
 
                 if enable_profiling:
@@ -377,11 +471,22 @@ class QwenVlAct_Trainer:
 
         # Validation loop
         for i, batch in enumerate(
-            tqdm(self.val_dataloader, desc="Validating", total=len(self.val_dataloader), 
-                 disable=not self.accelerator.is_main_process)
+            tqdm(
+                self.val_dataloader,
+                desc="Validating",
+                total=len(self.val_dataloader),
+                disable=not self.accelerator.is_main_process,
+            )
         ):
             if isinstance(self.dataset, PreprocessedDataset):
-                batch = {k: v.to(self.accelerator.device, non_blocking=True) if isinstance(v, torch.Tensor) else v for k, v in batch.items()}
+                batch = {
+                    k: (
+                        v.to(self.accelerator.device, non_blocking=True)
+                        if isinstance(v, torch.Tensor)
+                        else v
+                    )
+                    for k, v in batch.items()
+                }
 
             with torch.no_grad():
                 outputs = self.model(**batch, mode="train")
@@ -412,7 +517,7 @@ class QwenVlAct_Trainer:
         # Load pretrained model
         model = Qwen2_5_VLMoEForAction.from_pretrained(
             self.config["pretrained_wallx_path"],
-            **{"use_fast_tokenizer": self.use_fast_tokenizer}
+            **{"use_fast_tokenizer": self.use_fast_tokenizer},
         )
         self.processor = model.processor
         model = model.to(torch.bfloat16)
@@ -442,17 +547,26 @@ class QwenVlAct_Trainer:
             # Configure parameter groups
             if self.config.get("train_action_expert_only", False):
                 self.print_rank0("Training action expert only", flush=True)
-                param_groups = [{"params": moe_params, "lr": self.config["action_expert_learning_rate"]}]
+                param_groups = [
+                    {
+                        "params": moe_params,
+                        "lr": self.config["action_expert_learning_rate"],
+                    }
+                ]
             else:
                 param_groups = [
                     {"params": vlm_params, "lr": self.config["learning_rate"]},
-                    {"params": moe_params, "lr": self.config["action_expert_learning_rate"]},
+                    {
+                        "params": moe_params,
+                        "lr": self.config["action_expert_learning_rate"],
+                    },
                 ]
 
             self.optimizer = AdamW(param_groups, weight_decay=0.1)
             self.print_rank0(
                 f"Setting MoE learning rate to {self.config['action_expert_learning_rate']}, "
-                f"VLM learning rate to {self.config['learning_rate']}", flush=True
+                f"VLM learning rate to {self.config['learning_rate']}",
+                flush=True,
             )
         else:
             # Standard optimizer configuration
@@ -479,9 +593,13 @@ class QwenVlAct_Trainer:
         if hasattr(model, "enable_input_require_grads"):
             self.model.enable_input_require_grads()
         else:
+
             def make_inputs_require_grad(module, input, output):
                 output.requires_grad_(True)
-            self.model.get_input_embeddings().register_forward_hook(make_inputs_require_grad)
+
+            self.model.get_input_embeddings().register_forward_hook(
+                make_inputs_require_grad
+            )
 
         # Prepare model, optimizer, and scheduler for distributed training
         self.model, self.optimizer, self.lr_scheduler = self.accelerator.prepare(
@@ -522,7 +640,9 @@ class QwenVlAct_Trainer:
         Handles weight key renaming for MoE architecture compatibility.
         """
         # Load all safetensors files
-        weight_files = sorted([f for f in os.listdir(pretrain_weight_path) if f.endswith(".safetensors")])
+        weight_files = sorted(
+            [f for f in os.listdir(pretrain_weight_path) if f.endswith(".safetensors")]
+        )
         merged_weights = {}
 
         # Merge weights from all files
@@ -534,19 +654,31 @@ class QwenVlAct_Trainer:
         # Rename weights for MoE compatibility
         renamed_weights = {}
         for key, value in merged_weights.items():
-            if key.startswith("model.layers") and "mlp." in key and model.config.mlp_moe:
+            if (
+                key.startswith("model.layers")
+                and "mlp." in key
+                and model.config.mlp_moe
+            ):
                 # Rename MLP weights for MoE structure
                 layer_num = key.split(".layers.")[1].split(".mlp")[0]
-                new_key = key.replace(f"layers.{layer_num}.mlp.", f"layers.{layer_num}.moe.experts.0.")
+                new_key = key.replace(
+                    f"layers.{layer_num}.mlp.", f"layers.{layer_num}.moe.experts.0."
+                )
                 renamed_weights[new_key] = value
-            elif key.startswith("model.layers") and "self_attn." in key and model.config.attention_moe:
+            elif (
+                key.startswith("model.layers")
+                and "self_attn." in key
+                and model.config.attention_moe
+            ):
                 # Rename attention weights for MoE structure
                 layer_num = key.split(".layers.")[1].split(".self_attn")[0]
                 proj_types = ["q_proj", "k_proj", "v_proj", "o_proj"]
                 for proj in proj_types:
                     if proj in key:
-                        new_key = key.replace(f"layers.{layer_num}.self_attn.{proj}", 
-                                            f"layers.{layer_num}.self_attn.{proj}_experts.0")
+                        new_key = key.replace(
+                            f"layers.{layer_num}.self_attn.{proj}",
+                            f"layers.{layer_num}.self_attn.{proj}_experts.0",
+                        )
                         renamed_weights[new_key] = value
                         break
             else:
@@ -560,7 +692,16 @@ class QwenVlAct_Trainer:
 
         return model
 
-    def training_log(self, current_epoch, total_epoch, current_train_iter, total_train_iter, loss, lr, time_per_step):
+    def training_log(
+        self,
+        current_epoch,
+        total_epoch,
+        current_train_iter,
+        total_train_iter,
+        loss,
+        lr,
+        time_per_step,
+    ):
         """
         Log training progress and performance metrics.
 
@@ -573,7 +714,13 @@ class QwenVlAct_Trainer:
             lr (float): Current learning rate
             time_per_step (float): Time taken for current step
         """
-        timers_to_log = ["interval-time", "data-load", "forward-compute", "backward-compute", "optimizer"]
+        timers_to_log = [
+            "interval-time",
+            "data-load",
+            "forward-compute",
+            "backward-compute",
+            "optimizer",
+        ]
 
         log_string = f" [{datetime.now().strftime('%Y-%m-%d %H:%M:%S')}]"
         log_string += " epoch {:3d}/{:3d} |".format(current_epoch, total_epoch)
@@ -609,7 +756,9 @@ class QwenVlAct_Trainer:
             if isinstance(self.dataset, PreprocessedDataset):
                 torch.save(
                     {"epoch": epoch, "step": step},
-                    os.path.join(ckpt_path, f"epoch_{epoch}_step_{step}_rank_{_rank}.pth")
+                    os.path.join(
+                        ckpt_path, f"epoch_{epoch}_step_{step}_rank_{_rank}.pth"
+                    ),
                 )
 
     def resume_from_checkpoint(self):
@@ -632,7 +781,7 @@ class QwenVlAct_Trainer:
                     new_key = "module." + key
                 new_state_dict[new_key] = state_dict[key]
 
-            err = self.model.load_state_dict(new_state_dict, strict=False)
+            self.model.load_state_dict(new_state_dict, strict=False)
         else:
             # Load full checkpoint including optimizer and scheduler states
             self.accelerator.load_state(checkpoint_path)
@@ -662,7 +811,9 @@ class QwenVlAct_Trainer:
             mean_action = all_label.mean(dim=0)
             self.base_l1_loss = nn.functional.l1_loss(all_label, mean_action)
 
-        self.logger.log({"base_l1_loss": self.base_l1_loss.item()}, step=self.global_step)
+        self.logger.log(
+            {"base_l1_loss": self.base_l1_loss.item()}, step=self.global_step
+        )
 
         # Log L1 loss for each DOF component
         start_idx = 0
@@ -674,6 +825,8 @@ class QwenVlAct_Trainer:
             dof_l1 = nn.functional.l1_loss(dof_pred, dof_label)
 
             self.print_rank0(f"DOF {dof}, L1 loss: {dof_l1.item()}", flush=True)
-            self.logger.log({f"detail/l1_loss_{dof}": dof_l1.item()}, step=self.global_step)
+            self.logger.log(
+                {f"detail/l1_loss_{dof}": dof_l1.item()}, step=self.global_step
+            )
 
             start_idx = end_idx

wall_x/utils/constant.py

@@ -9,133 +9,292 @@ action_statistic_dof = {
             "min": [-3.6176],
             "delta": [8.5015],
         },
-        "follow_left_ee_cartesian_pos": {"min": [-0.036, -0.3241, -0.1245], "delta": [0.4389, 0.557, 0.479]},
-        "follow_left_ee_rotation": {"min": [-1.2373, -0.1929, -1.5182], "delta": [2.2009, 1.5669, 2.0936]},
+        "follow_left_ee_cartesian_pos": {
+            "min": [-0.036, -0.3241, -0.1245],
+            "delta": [0.4389, 0.557, 0.479],
+        },
+        "follow_left_ee_rotation": {
+            "min": [-1.2373, -0.1929, -1.5182],
+            "delta": [2.2009, 1.5669, 2.0936],
+        },
         "follow_left_gripper": {"min": [-0.1196], "delta": [4.5226]},
-        "follow_right_ee_cartesian_pos": {"min": [-0.0326, -0.2273, -0.1377], "delta": [0.4574, 0.5704, 0.4743]},
-        "follow_right_ee_rotation": {"min": [-1.2201, -0.2611, -0.7427], "delta": [2.6623, 1.6622, 2.4186]},
+        "follow_right_ee_cartesian_pos": {
+            "min": [-0.0326, -0.2273, -0.1377],
+            "delta": [0.4574, 0.5704, 0.4743],
+        },
+        "follow_right_ee_rotation": {
+            "min": [-1.2201, -0.2611, -0.7427],
+            "delta": [2.6623, 1.6622, 2.4186],
+        },
         "follow_right_gripper": {"min": [-0.1208], "delta": [4.5261]},
         "height": {"min": [-0.0001], "delta": [0.5051]},
         "head_actions": {"min": [-1.5000, -1.4167], "delta": [2.5000, 1.8879]},
-        "base_velocity": {"min": [-0.0359, -0.084, -0.0162], "delta": [0.1539, 0.1848, 0.0322]},
+        "base_velocity": {
+            "min": [-0.0359, -0.084, -0.0162],
+            "delta": [0.1539, 0.1848, 0.0322],
+        },
     },
     "DobbE": {
-        "follow_right_ee_cartesian_pos": {"min": [-0.6107, -0.3272, -0.4282], "delta": [1.2629, 1.5297, 0.8349]},
-        "follow_right_ee_rotation": {"min": [-1.7378, -1.4597, -1.8712], "delta": [2.7031, 2.8182, 3.5921]},
+        "follow_right_ee_cartesian_pos": {
+            "min": [-0.6107, -0.3272, -0.4282],
+            "delta": [1.2629, 1.5297, 0.8349],
+        },
+        "follow_right_ee_rotation": {
+            "min": [-1.7378, -1.4597, -1.8712],
+            "delta": [2.7031, 2.8182, 3.5921],
+        },
         "follow_right_gripper": {"min": [0.0], "delta": [0.9983]},
     },
     "RH20T": {
-        "follow_right_ee_cartesian_pos": {"min": [0.3646, -0.2722, 0.0066], "delta": [0.3813, 0.5973, 0.3277]},
-        "follow_right_ee_rotation": {"min": [-1.8716, -0.4398, -3.1414], "delta": [3.4145, 1.0225, 6.2828]},
+        "follow_right_ee_cartesian_pos": {
+            "min": [0.3646, -0.2722, 0.0066],
+            "delta": [0.3813, 0.5973, 0.3277],
+        },
+        "follow_right_ee_rotation": {
+            "min": [-1.8716, -0.4398, -3.1414],
+            "delta": [3.4145, 1.0225, 6.2828],
+        },
         "follow_right_gripper": {"min": [0.0], "delta": [95.0]},
     },
     "agibotworld_alpha": {
-        "follow_left_ee_cartesian_pos": {"min": [0.4954, 0.0166, 0.1729], "delta": [0.3336, 0.5123, 0.9189]},
-        "follow_left_ee_rotation": {"min": [-3.1064, -1.2629, -3.1238], "delta": [6.2127, 2.5923, 6.2496]},
+        "follow_left_ee_cartesian_pos": {
+            "min": [0.4954, 0.0166, 0.1729],
+            "delta": [0.3336, 0.5123, 0.9189],
+        },
+        "follow_left_ee_rotation": {
+            "min": [-3.1064, -1.2629, -3.1238],
+            "delta": [6.2127, 2.5923, 6.2496],
+        },
         "follow_left_gripper": {"min": [34.6222], "delta": [86.1921]},
-        "follow_right_ee_cartesian_pos": {"min": [0.4615, -0.5975, 0.1638], "delta": [0.3823, 0.5577, 0.8873]},
-        "follow_right_ee_rotation": {"min": [-3.0891, -1.0739, -2.5091], "delta": [6.1707, 2.3074, 3.8533]},
+        "follow_right_ee_cartesian_pos": {
+            "min": [0.4615, -0.5975, 0.1638],
+            "delta": [0.3823, 0.5577, 0.8873],
+        },
+        "follow_right_ee_rotation": {
+            "min": [-3.0891, -1.0739, -2.5091],
+            "delta": [6.1707, 2.3074, 3.8533],
+        },
         "follow_right_gripper": {"min": [34.6222], "delta": [85.7635]},
         "height": {"min": [0.0], "delta": [0.4535]},
         "head_actions": {"min": [-0.1746, 0.0523], "delta": [0.2444, 0.4713]},
     },
     "austin_buds": {
-        "follow_right_ee_cartesian_pos": {"min": [0.3496, -0.2855, 0.0105], "delta": [0.3748, 0.492, 0.3116]},
-        "follow_right_ee_rotation": {"min": [-3.1405, -0.151, -0.0737], "delta": [6.2813, 0.3218, 0.1536]},
+        "follow_right_ee_cartesian_pos": {
+            "min": [0.3496, -0.2855, 0.0105],
+            "delta": [0.3748, 0.492, 0.3116],
+        },
+        "follow_right_ee_rotation": {
+            "min": [-3.1405, -0.151, -0.0737],
+            "delta": [6.2813, 0.3218, 0.1536],
+        },
         "follow_right_gripper": {"min": [0.0076], "delta": [0.0724]},
     },
     "austin_sailor": {
-        "follow_right_ee_cartesian_pos": {"min": [0.387, -0.3165, 0.0244], "delta": [0.2999, 0.5252, 0.2308]},
-        "follow_right_ee_rotation": {"min": [-3.1402, -0.1618, -1.5918], "delta": [6.2804, 0.337, 2.9478]},
+        "follow_right_ee_cartesian_pos": {
+            "min": [0.387, -0.3165, 0.0244],
+            "delta": [0.2999, 0.5252, 0.2308],
+        },
+        "follow_right_ee_rotation": {
+            "min": [-3.1402, -0.1618, -1.5918],
+            "delta": [6.2804, 0.337, 2.9478],
+        },
         "follow_right_gripper": {"min": [0.0005], "delta": [0.0773]},
     },
     "austin_sirius": {
-        "follow_right_ee_cartesian_pos": {"min": [0.0, -0.1182, 0.0], "delta": [0.5329, 0.3812, 0.2723]},
-        "follow_right_ee_rotation": {"min": [-3.1407, -0.1243, -1.7434], "delta": [6.2823, 0.1975, 1.8073]},
+        "follow_right_ee_cartesian_pos": {
+            "min": [0.0, -0.1182, 0.0],
+            "delta": [0.5329, 0.3812, 0.2723],
+        },
+        "follow_right_ee_rotation": {
+            "min": [-3.1407, -0.1243, -1.7434],
+            "delta": [6.2823, 0.1975, 1.8073],
+        },
         "follow_right_gripper": {"min": [0.0334], "delta": [0.046]},
     },
     "bc_z": {
-        "follow_right_ee_cartesian_pos": {"min": [-0.3883, -0.1116, 0.6113], "delta": [0.7199, 0.4288, 0.3709]},
-        "follow_right_ee_rotation": {"min": [-1.056, -1.0587, -2.6295], "delta": [1.9142, 1.9455, 4.8064]},
+        "follow_right_ee_cartesian_pos": {
+            "min": [-0.3883, -0.1116, 0.6113],
+            "delta": [0.7199, 0.4288, 0.3709],
+        },
+        "follow_right_ee_rotation": {
+            "min": [-1.056, -1.0587, -2.6295],
+            "delta": [1.9142, 1.9455, 4.8064],
+        },
         "follow_right_gripper": {"min": [0.2], "delta": [0.8]},
     },
     "berkeley_autolab_ur5": {
-        "follow_right_ee_cartesian_pos": {"min": [0.3018, -0.2129, -0.1888], "delta": [0.3121, 0.52, 0.3107]},
-        "follow_right_ee_rotation": {"min": [-3.1396, -0.2278, 1.1413], "delta": [6.279, 0.454, 0.9841]},
+        "follow_right_ee_cartesian_pos": {
+            "min": [0.3018, -0.2129, -0.1888],
+            "delta": [0.3121, 0.52, 0.3107],
+        },
+        "follow_right_ee_rotation": {
+            "min": [-3.1396, -0.2278, 1.1413],
+            "delta": [6.279, 0.454, 0.9841],
+        },
         "follow_right_gripper": {"min": [0.0], "delta": [1.0]},
     },
     "berkeley_cable_routing": {
-        "follow_right_ee_cartesian_pos": {"min": [0.4617, -0.28, 0.03], "delta": [0.1838, 0.5665, 0.1272]},
-        "follow_right_ee_rotation": {"min": [-3.1413, -0.0299, -0.7665], "delta": [6.2826, 0.0692, 3.322]},
+        "follow_right_ee_cartesian_pos": {
+            "min": [0.4617, -0.28, 0.03],
+            "delta": [0.1838, 0.5665, 0.1272],
+        },
+        "follow_right_ee_rotation": {
+            "min": [-3.1413, -0.0299, -0.7665],
+            "delta": [6.2826, 0.0692, 3.322],
+        },
     },
     "berkeley_fanuc_manipulation": {
-        "follow_right_ee_cartesian_pos": {"min": [0.3718, -0.4072, 0.0184], "delta": [0.3483, 0.7201, 0.5229]},
-        "follow_right_ee_rotation": {"min": [-3.1399, -1.0166, -1.6988], "delta": [6.2802, 1.4498, 3.2074]},
+        "follow_right_ee_cartesian_pos": {
+            "min": [0.3718, -0.4072, 0.0184],
+            "delta": [0.3483, 0.7201, 0.5229],
+        },
+        "follow_right_ee_rotation": {
+            "min": [-3.1399, -1.0166, -1.6988],
+            "delta": [6.2802, 1.4498, 3.2074],
+        },
         "follow_right_gripper": {"min": [0.0], "delta": [1.0]},
     },
     "bridge_data_v2": {
-        "follow_right_ee_cartesian_pos": {"min": [0.1498, -0.2178, -0.0901], "delta": [0.3012, 0.469, 0.298]},
-        "follow_right_ee_rotation": {"min": [-0.3279, -0.6105, -1.0578], "delta": [0.7378, 1.0353, 2.2552]},
+        "follow_right_ee_cartesian_pos": {
+            "min": [0.1498, -0.2178, -0.0901],
+            "delta": [0.3012, 0.469, 0.298],
+        },
+        "follow_right_ee_rotation": {
+            "min": [-0.3279, -0.6105, -1.0578],
+            "delta": [0.7378, 1.0353, 2.2552],
+        },
         "follow_right_gripper": {"min": [0.0692], "delta": [0.9426]},
     },
     "dlr_edan_shared_control": {
-        "follow_right_ee_cartesian_pos": {"min": [-0.8387, 0.1473, -0.3934], "delta": [0.6579, 0.6025, 1.1566]},
-        "follow_right_ee_rotation": {"min": [-3.1217, -1.5197, -2.2516], "delta": [6.2505, 1.5594, 4.2831]},
+        "follow_right_ee_cartesian_pos": {
+            "min": [-0.8387, 0.1473, -0.3934],
+            "delta": [0.6579, 0.6025, 1.1566],
+        },
+        "follow_right_ee_rotation": {
+            "min": [-3.1217, -1.5197, -2.2516],
+            "delta": [6.2505, 1.5594, 4.2831],
+        },
         "follow_right_gripper": {"min": [0.0], "delta": [1.0]},
     },
     "droid": {
-        "follow_right_ee_cartesian_pos": {"min": [0.2667, -0.4396, -0.0472], "delta": [0.5159, 0.8806, 0.8331]},
-        "follow_right_ee_rotation": {"min": [-3.1374, -1.216, -2.1741], "delta": [6.2749, 2.1075, 4.2259]},
+        "follow_right_ee_cartesian_pos": {
+            "min": [0.2667, -0.4396, -0.0472],
+            "delta": [0.5159, 0.8806, 0.8331],
+        },
+        "follow_right_ee_rotation": {
+            "min": [-3.1374, -1.216, -2.1741],
+            "delta": [6.2749, 2.1075, 4.2259],
+        },
         "follow_right_gripper": {"min": [0.0], "delta": [0.9912]},
     },
     "fmb": {
-        "follow_right_ee_cartesian_pos": {"min": [0.3554, -0.2844, 0.0354], "delta": [0.336, 0.4961, 0.2943]},
-        "follow_right_ee_rotation": {"min": [-3.1404, -0.9302, -0.0599], "delta": [6.2807, 1.724, 1.8284]},
+        "follow_right_ee_cartesian_pos": {
+            "min": [0.3554, -0.2844, 0.0354],
+            "delta": [0.336, 0.4961, 0.2943],
+        },
+        "follow_right_ee_rotation": {
+            "min": [-3.1404, -0.9302, -0.0599],
+            "delta": [6.2807, 1.724, 1.8284],
+        },
         "follow_right_gripper": {"min": [0.0], "delta": [1.0]},
     },
     "fractal": {
-        "follow_right_ee_cartesian_pos": {"min": [0.3242, -0.2836, 0.1405], "delta": [0.5518, 0.4963, 0.9328]},
-        "follow_right_ee_rotation": {"min": [-3.1308, -0.2421, -2.9685], "delta": [6.2609, 1.7343, 5.819]},
+        "follow_right_ee_cartesian_pos": {
+            "min": [0.3242, -0.2836, 0.1405],
+            "delta": [0.5518, 0.4963, 0.9328],
+        },
+        "follow_right_ee_rotation": {
+            "min": [-3.1308, -0.2421, -2.9685],
+            "delta": [6.2609, 1.7343, 5.819],
+        },
         "follow_right_gripper": {"min": [0.0], "delta": [1.0]},
     },
     "furniture_bench": {
-        "follow_right_ee_cartesian_pos": {"min": [0.3691, -0.181, 0.0058], "delta": [0.2962, 0.3582, 0.1775]},
-        "follow_right_ee_rotation": {"min": [-3.1394, -0.6121, -1.9958], "delta": [6.2786, 1.6114, 3.7748]},
+        "follow_right_ee_cartesian_pos": {
+            "min": [0.3691, -0.181, 0.0058],
+            "delta": [0.2962, 0.3582, 0.1775],
+        },
+        "follow_right_ee_rotation": {
+            "min": [-3.1394, -0.6121, -1.9958],
+            "delta": [6.2786, 1.6114, 3.7748],
+        },
         "follow_right_gripper": {"min": [0.0035], "delta": [0.0762]},
     },
     "jaco_play": {
-        "follow_right_ee_cartesian_pos": {"min": [-0.3787, -0.6294, 0.1682], "delta": [0.5898, 0.3587, 0.2183]},
-        "follow_right_ee_rotation": {"min": [0.9792, -0.0668, -0.0498], "delta": [0.0175, 0.1277, 0.0686]},
+        "follow_right_ee_cartesian_pos": {
+            "min": [-0.3787, -0.6294, 0.1682],
+            "delta": [0.5898, 0.3587, 0.2183],
+        },
+        "follow_right_ee_rotation": {
+            "min": [0.9792, -0.0668, -0.0498],
+            "delta": [0.0175, 0.1277, 0.0686],
+        },
         "follow_right_gripper": {"min": [0.0791], "delta": [0.1033]},
     },
     "nyu_rot": {
-        "follow_right_ee_cartesian_pos": {"min": [0.25, -1.0, -0.2], "delta": [0.75, 2.0, 1.2]},
-        "follow_right_ee_rotation": {"min": [-3.1416, -3.1416, 6.2831], "delta": [9.4248, 4.1416, 0.0]},
+        "follow_right_ee_cartesian_pos": {
+            "min": [0.25, -1.0, -0.2],
+            "delta": [0.75, 2.0, 1.2],
+        },
+        "follow_right_ee_rotation": {
+            "min": [-3.1416, -3.1416, 6.2831],
+            "delta": [9.4248, 4.1416, 0.0],
+        },
         "follow_right_gripper": {"min": [0.0], "delta": [1.0]},
     },
     "stanford_hydra": {
-        "follow_right_ee_cartesian_pos": {"min": [0.2068, -0.274, 0.1317], "delta": [0.4929, 0.4981, 0.4588]},
-        "follow_right_ee_rotation": {"min": [-3.1321, -0.7496, -3.0269], "delta": [6.2658, 1.5176, 5.8261]},
+        "follow_right_ee_cartesian_pos": {
+            "min": [0.2068, -0.274, 0.1317],
+            "delta": [0.4929, 0.4981, 0.4588],
+        },
+        "follow_right_ee_rotation": {
+            "min": [-3.1321, -0.7496, -3.0269],
+            "delta": [6.2658, 1.5176, 5.8261],
+        },
         "follow_right_gripper": {"min": [0.0], "delta": [0.0811]},
     },
     "stanford_kuka_multimodal": {
-        "follow_right_ee_cartesian_pos": {"min": [0.4781, -0.0659, 0.3424], "delta": [0.0868, 0.0864, 0.1863]},
-        "follow_right_ee_rotation": {"min": [-3.136, -0.0521, -3.1413], "delta": [6.2727, 0.1109, 6.2825]},
+        "follow_right_ee_cartesian_pos": {
+            "min": [0.4781, -0.0659, 0.3424],
+            "delta": [0.0868, 0.0864, 0.1863],
+        },
+        "follow_right_ee_rotation": {
+            "min": [-3.136, -0.0521, -3.1413],
+            "delta": [6.2727, 0.1109, 6.2825],
+        },
         "follow_right_gripper": {"min": [-0.4713], "delta": [0.9485]},
     },
     "taco_play": {
-        "follow_right_ee_cartesian_pos": {"min": [0.1375, -0.4291, 0.2052], "delta": [0.5327, 1.0237, 0.3913]},
-        "follow_right_ee_rotation": {"min": [-3.1391, -0.6946, -1.2808], "delta": [6.2784, 0.8196, 3.0856]},
+        "follow_right_ee_cartesian_pos": {
+            "min": [0.1375, -0.4291, 0.2052],
+            "delta": [0.5327, 1.0237, 0.3913],
+        },
+        "follow_right_ee_rotation": {
+            "min": [-3.1391, -0.6946, -1.2808],
+            "delta": [6.2784, 0.8196, 3.0856],
+        },
         "follow_right_gripper": {"min": [0.0001], "delta": [0.0806]},
     },
     "utaustin_mutex": {
-        "follow_right_ee_cartesian_pos": {"min": [0.3213, -0.4734, 0.0141], "delta": [0.2108, 0.8471, 0.5644]},
-        "follow_right_ee_rotation": {"min": [-3.1404, -0.2202, -1.5489], "delta": [6.2805, 0.582, 1.9282]},
+        "follow_right_ee_cartesian_pos": {
+            "min": [0.3213, -0.4734, 0.0141],
+            "delta": [0.2108, 0.8471, 0.5644],
+        },
+        "follow_right_ee_rotation": {
+            "min": [-3.1404, -0.2202, -1.5489],
+            "delta": [6.2805, 0.582, 1.9282],
+        },
         "follow_right_gripper": {"min": [0.0019], "delta": [0.0738]},
     },
     "viola": {
-        "follow_right_ee_cartesian_pos": {"min": [0.4011, -0.2521, 0.0103], "delta": [0.2444, 0.4305, 0.4355]},
-        "follow_right_ee_rotation": {"min": [-3.1403, -0.2737, -1.8626], "delta": [6.2804, 0.4901, 2.0618]},
+        "follow_right_ee_cartesian_pos": {
+            "min": [0.4011, -0.2521, 0.0103],
+            "delta": [0.2444, 0.4305, 0.4355],
+        },
+        "follow_right_ee_rotation": {
+            "min": [-3.1403, -0.2737, -1.8626],
+            "delta": [6.2804, 0.4901, 2.0618],
+        },
         "follow_right_gripper": {"min": [0.0002], "delta": [0.0773]},
     },
     "kuka": {
@@ -179,11 +338,23 @@ action_statistic_dof = {
         },
     },
     "agibotworld_beta": {
-        "follow_left_ee_cartesian_pos": {"min": [0.4954, 0.0166, 0.1729], "delta": [0.3336, 0.5123, 0.9189]},
-        "follow_left_ee_rotation": {"min": [-3.1064, -1.2629, -3.1238], "delta": [6.2127, 2.5923, 6.2496]},
+        "follow_left_ee_cartesian_pos": {
+            "min": [0.4954, 0.0166, 0.1729],
+            "delta": [0.3336, 0.5123, 0.9189],
+        },
+        "follow_left_ee_rotation": {
+            "min": [-3.1064, -1.2629, -3.1238],
+            "delta": [6.2127, 2.5923, 6.2496],
+        },
         "follow_left_gripper": {"min": [34.6222], "delta": [86.1921]},
-        "follow_right_ee_cartesian_pos": {"min": [0.4615, -0.5975, 0.1638], "delta": [0.3823, 0.5577, 0.8873]},
-        "follow_right_ee_rotation": {"min": [-3.0891, -1.0739, -2.5091], "delta": [6.1707, 2.3074, 3.8533]},
+        "follow_right_ee_cartesian_pos": {
+            "min": [0.4615, -0.5975, 0.1638],
+            "delta": [0.3823, 0.5577, 0.8873],
+        },
+        "follow_right_ee_rotation": {
+            "min": [-3.0891, -1.0739, -2.5091],
+            "delta": [6.1707, 2.3074, 3.8533],
+        },
         "follow_right_gripper": {"min": [34.6222], "delta": [85.7635]},
         "height": {"min": [0.0], "delta": [0.4535]},
         "head_actions": {"min": [-0.1746, 0.0523], "delta": [0.2444, 0.4713]},

wall_x/utils/timers.py

@@ -4,23 +4,28 @@ from torch.cuda import nvtx
 from abc import ABC, abstractmethod
 from typing import List
 
+
 def _is_distributed():
     return torch.distributed.is_available() and torch.distributed.is_initialized()
 
+
 def _get_world_size():
     if _is_distributed():
         return torch.distributed.get_world_size()
     return 1
 
+
 def _get_rank():
     if _is_distributed():
         return torch.distributed.get_rank()
     return 0
 
+
 def _barrier(group=None):
     if _is_distributed():
         torch.distributed.barrier(group=group)
 
+
 if torch.distributed.is_available():
     try:
         dist_all_gather_func = torch.distributed.all_gather_into_tensor
@@ -29,6 +34,7 @@ if torch.distributed.is_available():
 else:
     dist_all_gather_func = None
 
+
 class TimerBase(ABC):
     """Timer base class."""
 
@@ -76,7 +82,7 @@ class DummyTimer(TimerBase):
     """Dummy Timer."""
 
     def __init__(self):
-        super().__init__('dummy timer')
+        super().__init__("dummy timer")
 
     def start(self, barrier=False, nvtx_push=False):
         return
@@ -89,8 +95,8 @@ class DummyTimer(TimerBase):
 
     def elapsed(self, reset=True, barrier=False):
         raise Exception(
-            'dummy timer should not be used to calculate elapsed time, '
-            'check if timer\'s log_level <= self._log_level.'
+            "dummy timer should not be used to calculate elapsed time, "
+            "check if timer's log_level <= self._log_level."
         )
 
     def active_time(self):
@@ -98,8 +104,8 @@ class DummyTimer(TimerBase):
         Note: Not supported for DummyTimer.
         """
         raise Exception(
-            'active timer should not be used to calculate elapsed time, '
-            'check if timer\'s log_level <= self._log_level.'
+            "active timer should not be used to calculate elapsed time, "
+            "check if timer's log_level <= self._log_level."
         )
 
 
@@ -144,7 +150,7 @@ class Timer(TimerBase):
         Args:
             barrier (bool, optional): Synchronizes ranks before starting. Defaults to False.
         """
-        assert not self._started, 'timer has already been started'
+        assert not self._started, "timer has already been started"
         if barrier:
             _barrier(group=self._barrier_group)
         if torch.cuda.is_available():
@@ -155,7 +161,6 @@ class Timer(TimerBase):
             nvtx.range_push("{}".format(self.name))
             self.nvtx = True
 
-
     def stop(self, barrier=False, sync=False):
         """Stop the timer.
 
@@ -164,7 +169,7 @@ class Timer(TimerBase):
         """
         if self.nvtx:
             nvtx.range_pop()
-        assert self._started, 'timer is not started'
+        assert self._started, "timer is not started"
         if barrier:
             _barrier(group=self._barrier_group)
         if torch.cuda.is_available() and sync:
@@ -221,10 +226,10 @@ class Timers:
                               Allowed: ['max', 'minmax', 'all'].
         """
         self._log_level = log_level
-        allowed_log_options = set(['max', 'minmax', 'all'])
+        allowed_log_options = set(["max", "minmax", "all"])
         assert (
             log_option in allowed_log_options
-        ), 'input log option {} is invalid. It must be one of {}'.format(
+        ), "input log option {} is invalid. It must be one of {}".format(
             log_option, allowed_log_options
         )
         self._log_option = log_option
@@ -240,8 +245,10 @@ class Timers:
         if name in self._timers:
             if log_level is not None:
                 assert log_level == self._log_levels[name], (
-                    'input log level {} does not match already existing '
-                    'log level {} for {} timer'.format(log_level, self._log_levels[name], name)
+                    "input log level {} does not match already existing "
+                    "log level {} for {} timer".format(
+                        log_level, self._log_levels[name], name
+                    )
                 )
             return self._timers[name]
         # If timer does not exist and no log level is provided,
@@ -250,7 +257,7 @@ class Timers:
             log_level = self._max_log_level
         assert (
             log_level <= self._max_log_level
-        ), 'log level {} is larger than max supported log level {}'.format(
+        ), "log level {} is larger than max supported log level {}".format(
             log_level, self._max_log_level
         )
         # Now if the input log level is larger than the one set for
@@ -284,7 +291,7 @@ class Timers:
         if torch.cuda.is_available():
             device = torch.cuda.current_device()
         else:
-            device = torch.device('cpu')
+            device = torch.device("cpu")
 
         rank_name_to_time = torch.zeros(
             (world_size, len(names)), dtype=torch.float, device=device
@@ -296,7 +303,9 @@ class Timers:
 
         if world_size > 1 and _is_distributed() and dist_all_gather_func is not None:
             try:
-                dist_all_gather_func(rank_name_to_time.view(-1), rank_name_to_time[rank, :].view(-1))
+                dist_all_gather_func(
+                    rank_name_to_time.view(-1), rank_name_to_time[rank, :].view(-1)
+                )
             except Exception as e:
                 print(f"Warning: all_gather failed: {e}. Using single rank timing.")
 
@@ -319,30 +328,38 @@ class Timers:
                 )
         return name_to_min_max_time
 
-    def _get_global_min_max_time_string(self, names, reset, barrier, normalizer, max_only):
+    def _get_global_min_max_time_string(
+        self, names, reset, barrier, normalizer, max_only
+    ):
         """Report strings for max/minmax times across all ranks."""
-        name_to_min_max_time = self._get_global_min_max_time(names, reset, barrier, normalizer)
+        name_to_min_max_time = self._get_global_min_max_time(
+            names, reset, barrier, normalizer
+        )
         if not name_to_min_max_time:
             return None
 
         world_size = _get_world_size()
         if world_size == 1:
-            output_string = 'time (ms):'
+            output_string = "time (ms):"
             for name in name_to_min_max_time:
                 _, max_time = name_to_min_max_time[name]
-                output_string += '\n    {}: {:.2f}'.format((name + ' ').ljust(48, '.'), max_time)
+                output_string += "\n    {}: {:.2f}".format(
+                    (name + " ").ljust(48, "."), max_time
+                )
         else:
             if max_only:
-                output_string = 'max time across ranks (ms):'
+                output_string = "max time across ranks (ms):"
             else:
-                output_string = '(min, max) time across ranks (ms):'
+                output_string = "(min, max) time across ranks (ms):"
             for name in name_to_min_max_time:
                 min_time, max_time = name_to_min_max_time[name]
                 if max_only:
-                    output_string += '\n    {}: {:.2f}'.format((name + ' ').ljust(48, '.'), max_time)
+                    output_string += "\n    {}: {:.2f}".format(
+                        (name + " ").ljust(48, "."), max_time
+                    )
                 else:
-                    output_string += '\n    {}: ({:.2f}, {:.2f})'.format(
-                        (name + ' ').ljust(48, '.'), min_time, max_time
+                    output_string += "\n    {}: ({:.2f}, {:.2f})".format(
+                        (name + " ").ljust(48, "."), min_time, max_time
                     )
         return output_string
 
@@ -351,7 +368,7 @@ class Timers:
         rank_name_to_time = self._get_elapsed_time_all_ranks(names, reset, barrier)
         world_size = _get_world_size()
 
-        output_string = 'times across ranks (ms):'
+        output_string = "times across ranks (ms):"
         no_reported_timing = True
         for i, name in enumerate(names):
             not_yet_found = True
@@ -360,13 +377,13 @@ class Timers:
                     no_reported_timing = False
                     if not_yet_found:
                         not_yet_found = False
-                        output_string += '\n  {}:'.format(name)
+                        output_string += "\n  {}:".format(name)
                     if world_size == 1:
-                        output_string += '\n     {:.2f}'.format(
+                        output_string += "\n     {:.2f}".format(
                             rank_name_to_time[rank, i] / normalizer
                         )
                     else:
-                        output_string += '\n     rank {:2d}: {:.2f}'.format(
+                        output_string += "\n     rank {:2d}: {:.2f}".format(
                             rank, rank_name_to_time[rank, i] / normalizer
                         )
         if no_reported_timing:
@@ -398,23 +415,23 @@ class Timers:
             str: Formatted string with the timer values.
         """
 
-        if names == None:  # get all registered timers
+        if names is None:  # get all registered timers
             names = list(self._timers.keys())
 
         assert normalizer > 0.0
-        if self._log_option in ['max', 'minmax']:
+        if self._log_option in ["max", "minmax"]:
             max_only = False
-            if self._log_option == 'max':
+            if self._log_option == "max":
                 max_only = True
             output_string = self._get_global_min_max_time_string(
                 names, reset, barrier, normalizer / 1000.0, max_only
             )
-        elif self._log_option == 'all':
+        elif self._log_option == "all":
             output_string = self._get_all_ranks_time_string(
                 names, reset, barrier, normalizer / 1000.0
             )
         else:
-            raise Exception('unknown timing log option {}'.format(self._log_option))
+            raise Exception("unknown timing log option {}".format(self._log_option))
         return output_string
 
     def log(
@@ -476,8 +493,10 @@ class Timers:
         # torch.utils.add_scalars makes each timer its own run, which
         # polutes the runs list, so we just add each as a scalar
         assert normalizer > 0.0
-        name_to_min_max_time = self._get_global_min_max_time(names, reset, barrier, normalizer)
+        name_to_min_max_time = self._get_global_min_max_time(
+            names, reset, barrier, normalizer
+        )
         if writer is not None:
             for name in name_to_min_max_time:
                 _, max_time = name_to_min_max_time[name]
-                writer.add_scalar(name + '-time', max_time, iteration)
+                writer.add_scalar(name + "-time", max_time, iteration)