Init

wall_x/data/config.py

@@ -0,0 +1,95 @@
+from typing import List, Dict, Optional
+from dataclasses import dataclass, field
+from qwen_vl_utils.vision_process import MIN_PIXELS, MAX_PIXELS, IMAGE_FACTOR
+
+
+# Tactile sensor file mapping for data processing
+TACTILE_FILE_MAPPING = {
+    "tactile_data_left": "left_tactile", 
+    "tactile_data_right": "right_tactile"
+}
+
+# Supported action datasets
+ACTION_DATASET_NAMES = [
+    "x2_normal", "agibotworld_alpha", "droid", "fractal", "bridge_data_v2", 
+    "DobbE", "RH20T", "UMI-biarm", "austin_buds", "austin_sailor", "austin_sirius", 
+    "bc_z", "berkeley_autolab_ur5", "berkeley_cable_routing", "berkeley_fanuc_manipulation", 
+    "dlr_edan_shared_control", "fmb", "furniture_bench", "jaco_play", "nyu_rot", 
+    "stanford_hydra", "stanford_kuka_multimodal", "taco_play", "utaustin_mutex", "viola"
+]
+
+# Supported multimodal datasets
+MULTIMODAL_DATASET_NAMES = [
+    "x2_multimodal_from_action", "x2_multimodal", "x2_subtask_generation",
+    "multimodal_CapsFusion", "multimodal_Robo2VLM", "multimodal_RoboPoint", 
+    "multimodal_EQA", "multimodal_Cambrian", "multimodal_pixmo", 
+    "multimodal_VQAv2", "multimodal_COCO"
+]
+
+
+@dataclass
+class X2RDataProcessingConfig:
+    """Configuration class for X2R data processing pipeline.
+    
+    This class contains all the necessary parameters for processing robotic data
+    including camera mappings, tactile sensor configurations, action predictions,
+    and various processing options.
+    """
+    
+    # Action prediction configuration
+    predict_action_keys: List[str] = field(default_factory=list)
+    obs_action_keys: List[str] = field(default_factory=list)
+    
+    # Image resolution settings for different views
+    resolution: Dict[str, int] = field(
+        default_factory=lambda: {
+            "face_view": -1, 
+            "left_wrist_view": 128, 
+            "right_wrist_view": 128
+        }
+    )
+    
+    # Dataset splitting
+    train_test_split: float = 0.9
+    split_seed: int = 42
+    
+    # Instruction handling
+    priority_order: Optional[Dict[str, float]] = None
+    
+    # Vision model parameters
+    model_type: str = "qwen2_5"
+    max_pixels: int = MAX_PIXELS
+    min_pixels: int = MIN_PIXELS
+    image_factor: int = IMAGE_FACTOR
+
+    generate_subtask_ratio: float = 0.0
+    
+    def __post_init__(self):
+        """Post-initialization validation and setup."""
+        # Validate train/test split
+        if not 0 < self.train_test_split < 1:
+            raise ValueError(f"train_test_split must be between 0 and 1, got {self.train_test_split}")
+    
+    def as_dict(self) -> Dict:
+        """Convert configuration to dictionary format.
+        
+        Returns:
+            Dict: Configuration as dictionary
+        """
+        return self.__dict__
+    
+    def update(self, **kwargs) -> 'X2RDataProcessingConfig':
+        """Update configuration parameters.
+        
+        Args:
+            **kwargs: Key-value pairs to update
+            
+        Returns:
+            X2RDataProcessingConfig: Updated configuration instance
+        """
+        for key, value in kwargs.items():
+            if hasattr(self, key):
+                setattr(self, key, value)
+            else:
+                raise ValueError(f"Unknown configuration parameter: {key}")
+        return self

wall_x/data/load_lerobot_dataset.py

@@ -0,0 +1,496 @@
+"""
+LeRobot Dataset Loader - Distributed Version
+"""
+
+import torch
+from torch.utils.data import DistributedSampler
+from lerobot.datasets.lerobot_dataset import LeRobotDataset
+from typing import Protocol, SupportsIndex, TypeVar
+from qwen_vl_utils.vision_process import smart_resize
+from wall_x.data.config import X2RDataProcessingConfig
+from wall_x.data.utils import process_grounding_points, get_wallx_normal_text, replace_action_token, preprocesser_call
+
+from transformers import AutoProcessor
+
+T_co = TypeVar("T_co", covariant=True)
+
+CAMERA_KEY_MAPPINGS = {
+    "lerobot/aloha_mobile_cabinet": {
+        "observation.images.cam_high": "face_view",
+        "observation.images.cam_left_wrist": "left_wrist_view",
+        "observation.images.cam_right_wrist": "right_wrist_view",
+    },
+}
+
+
+# Abstract class for dataset
+class Dataset(Protocol[T_co]):
+    """Interface for a dataset with random access."""
+
+    def __getitem__(self, index: SupportsIndex) -> T_co:
+        raise NotImplementedError("Subclasses of Dataset should implement __getitem__.")
+
+    def __len__(self) -> int:
+        raise NotImplementedError("Subclasses of Dataset should implement __len__.")
+
+
+class PreprocessedDataset(Dataset[T_co]):
+    def __init__(self, dataset, config, dataload_config, seed=42, rank=0, world_size=1):
+        self._dataset = dataset
+        self.seed = seed
+        self.rank = rank
+        self.world_size = world_size
+
+        # init configs
+        self.config = config
+        self.use_fast_tokenizer = self.config.get("use_fast_tokenizer", False)
+        self.dataload_config = dataload_config
+
+        self.data_config = X2RDataProcessingConfig().update(
+            train_test_split=self.dataload_config["train_test_split"],
+            split_seed=self.dataload_config["split_seed"],
+            predict_action_keys=self.dataload_config["predict_action_keys"],
+            obs_action_keys=self.dataload_config["obs_action_keys"],
+            resolution=self.dataload_config.get("resolution", None),
+            priority_order=self.dataload_config.get("priority_order", None),
+        )
+
+        self._cam_key_mapping = CAMERA_KEY_MAPPINGS[self._dataset.meta.repo_id]
+
+    def _vision_preprocess(self, frames):
+        processed_frames = []
+        for key in self._dataset.meta.camera_keys:
+            from PIL import Image
+
+            current_obs = frames[key].clone().permute(1, 2, 0)
+
+            img_pil = Image.fromarray((current_obs * 255).to(torch.uint8).cpu().numpy())
+            orig_width, orig_height = img_pil.size
+            # 2. Apply resolution constraints (if config is not -1)
+            target_size = self.data_config.resolution.get(self._cam_key_mapping[key], -1)
+            if target_size != -1:
+                # Maintain aspect ratio logic
+                if orig_width > orig_height:  # Landscape image
+                    new_width = target_size
+                    new_height = int(target_size * orig_height / orig_width)
+                else:  # Portrait image
+                    new_height = target_size
+                    new_width = int(target_size * orig_width / orig_height)
+                img_pil = img_pil.resize((new_width, new_height))
+
+            # 3. Apply smart scaling (qwen logic)
+            current_width, current_height = img_pil.size
+            resized_height, resized_width = smart_resize(
+                current_height,
+                current_width,
+                factor=self.data_config.image_factor,
+                min_pixels=self.data_config.min_pixels,
+                max_pixels=self.data_config.max_pixels,
+            )
+            resized_img = img_pil.resize((resized_width, resized_height))
+            processed_frames.append(resized_img)
+
+        return processed_frames, orig_height, orig_width, resized_height, resized_width
+
+    def __getitem__(self, index):
+        data = self._dataset[index]
+        image_inputs, h, w, resize_h, resize_w = self._vision_preprocess(data)
+        agent_pos = data["observation.state"]
+        action = data["action"]
+        frame_index = data["frame_index"]
+        instruction_info = {"instruction": data["task"]}
+        generate_subtask_ratio = self.data_config.generate_subtask_ratio
+        complete_text, generate_subtask = get_wallx_normal_text(
+            instruction_info,
+            33 - 1,
+            frame_index,
+            self.data_config.priority_order,
+            self._cam_key_mapping,
+            generate_subtask_ratio=generate_subtask_ratio,
+        )
+        text = process_grounding_points(complete_text, h, w, resize_h, resize_w, self.data_config.model_type)
+        result = {
+            "image_inputs": image_inputs,
+            "text": text,
+            "action": action,
+            "agent_pos": agent_pos,
+            "frame_index": frame_index,
+        }
+
+        return result
+
+    def __len__(self) -> int:
+        return len(self._dataset)
+
+    def get_train_dataloader(self):
+        """
+        Get distributed training dataloader
+
+        Args:
+            rank: Current process rank
+            world_size: Total number of processes
+            seed: Random seed for reproducibility
+        """
+
+        batch_size = self.config.get("batch_size_per_gpu", 8)
+        num_workers = self.config.get("num_workers", 4)
+
+        # Create distributed sampler
+        sampler = DistributedSampler(
+            self,
+            num_replicas=self.world_size,
+            rank=self.rank,
+            shuffle=True,
+            seed=self.seed,
+            drop_last=True,  # Ensure all processes have same number of batches
+        )
+
+        dataloader = torch.utils.data.DataLoader(
+            self,
+            batch_size=batch_size,
+            sampler=sampler,  # Use distributed sampler instead of shuffle=True
+            num_workers=num_workers,
+            collate_fn=DataCollator(self.config, self.dataload_config, self._dataset.meta.stats),
+            pin_memory=True,  # Enable for GPU training
+            persistent_workers=num_workers > 0,  # Only if num_workers > 0
+            prefetch_factor=2,  # Reduce memory usage
+            drop_last=True,  # Avoid incomplete batches
+        )
+
+        return dataloader, sampler
+
+    def get_val_dataloader(self):
+        """
+        Get distributed evaluation dataloader (no shuffling for consistent evaluation)
+        """
+
+        batch_size = self.config.get("eval_batch_size_per_gpu", self.config.get("batch_size_per_gpu", 8))
+        num_workers = self.config.get("num_workers", 4)
+
+        # Create distributed sampler for evaluation (no shuffle)
+        sampler = DistributedSampler(
+            self,
+            num_replicas=self.world_size,
+            rank=self.rank,
+            shuffle=False,  # No shuffling for evaluation
+            drop_last=False,  # Keep all samples for evaluation
+        )
+
+        dataloader = torch.utils.data.DataLoader(
+            self,
+            batch_size=batch_size,
+            sampler=sampler,
+            num_workers=num_workers,
+            collate_fn=DataCollator(self.config, self.dataload_config, self._dataset.meta.stats),
+            pin_memory=True,
+            persistent_workers=num_workers > 0,
+            prefetch_factor=2,
+            drop_last=False,
+        )
+
+        return dataloader, sampler
+
+
+class DataCollator:
+    # Class-level cache for processors to avoid reloading
+    _processor_cache = {}
+    _action_tokenizer_cache = {}
+
+    def __init__(self, config, dataload_config, stats):
+        self.config = config
+        self.dataload_config = dataload_config
+        self.stats = stats
+        self.min_stat = stats["action"]["min"]
+        self.max_stat = stats["action"]["max"]
+        self.delta = self.max_stat - self.min_stat
+        self.use_fast_tokenizer = self.config.get("use_fast_tokenizer", False)
+        self.load_processor()
+
+    def load_processor(self):
+        processor_path = self.config["processor_path"]
+        action_tokenizer_path = self.config["action_tokenizer_path"]
+
+        # Use cached processors if available
+        if processor_path not in self._processor_cache:
+            self._processor_cache[processor_path] = AutoProcessor.from_pretrained(processor_path, use_fast=True)
+            if self.config.get("padding_side", "left") == "left":
+                self._processor_cache[processor_path].tokenizer.padding_side = "left"
+
+        if action_tokenizer_path not in self._action_tokenizer_cache:
+            self._action_tokenizer_cache[action_tokenizer_path] = AutoProcessor.from_pretrained(action_tokenizer_path, trust_remote_code=True)
+
+        self.processor = self._processor_cache[processor_path]
+        self.val_processor = self._processor_cache[processor_path]
+        self.train_action_tokenizer = self._action_tokenizer_cache[action_tokenizer_path]
+        self.val_action_tokenizer = self._action_tokenizer_cache[action_tokenizer_path]
+
+        new_tokens = ["<|propri|>", "<|action|>"]
+
+        new_tokens += [f"<|action_token_{i}|>" for i in range(self.train_action_tokenizer.vocab_size)]
+        if not self.use_fast_tokenizer:
+            self.train_action_tokenizer = None
+            self.val_action_tokenizer = None
+
+        # Only add tokens if not already added
+        if "<|propri|>" not in self.processor.tokenizer.get_vocab():
+            num_added_tokens = self.processor.tokenizer.add_tokens(new_tokens)
+            self.val_processor.tokenizer.add_tokens(new_tokens)
+
+        if self.use_fast_tokenizer:
+            self.action_mapper = {}
+            for i in range(self.train_action_tokenizer.vocab_size):
+                token = f"<|action_token_{i}|>"
+                token_id = self.processor.tokenizer.convert_tokens_to_ids(token)
+                self.action_mapper[token_id] = i
+        else:
+            self.action_mapper = None
+
+    @classmethod
+    def _normalize(cls, action, min_stat, delta):
+        x = (action - min_stat) / (delta)
+        x = x * 2 - 1
+        x = torch.clamp(x, -1, 1)
+        return x
+
+    def __call__(self, batch):
+        additional_inputs = {}
+
+        for key in batch[0].keys():
+            if key == "agent_pos":
+                agent_pos = torch.stack([item["agent_pos"] for item in batch])
+                if agent_pos.dim() == 2:
+                    agent_pos = agent_pos.unsqueeze(1)
+                agent_pos_mask = (~torch.isnan(agent_pos)).float()
+                agent_pos.nan_to_num_(nan=0.0)
+                agent_pos = self._normalize(agent_pos, self.min_stat, self.delta)
+                if agent_pos.shape[-1] != 20:
+                    agent_pos = torch.cat([agent_pos, torch.zeros(agent_pos.shape[0], agent_pos.shape[1], 20 - agent_pos.shape[-1])], dim=-1)
+                    agent_pos_mask = torch.cat(
+                        [agent_pos_mask, torch.zeros(agent_pos_mask.shape[0], agent_pos_mask.shape[1], 20 - agent_pos_mask.shape[-1])], dim=-1
+                    )
+                additional_inputs["proprioception"] = agent_pos
+                additional_inputs["agent_pos_mask"] = agent_pos_mask
+            elif key == "action":
+                action = torch.stack([item["action"] for item in batch])
+                if action.dim() == 2:
+                    action = action.unsqueeze(1)
+                dof_mask = (~torch.isnan(action)).float()
+                action.nan_to_num_(nan=0.0)
+                action = self._normalize(action, self.min_stat, self.delta)
+                if action.shape[-1] != 20:
+                    action = torch.cat([action, torch.zeros(action.shape[0], action.shape[1], 20 - action.shape[-1])], dim=-1)
+                    dof_mask = torch.cat([dof_mask, torch.zeros(dof_mask.shape[0], dof_mask.shape[1], 20 - dof_mask.shape[-1])], dim=-1)
+                additional_inputs["action_chunk"] = action
+                additional_inputs["dof_mask"] = dof_mask
+            elif key == "image_inputs":
+                additional_inputs["image_inputs"] = [item["image_inputs"] for item in batch]
+            elif key == "text":
+                additional_inputs["text"] = [item["text"] for item in batch]
+            elif key == "frame_index":
+                additional_inputs["frame_index"] = torch.stack([item["frame_index"] for item in batch])
+            else:
+                raise NotImplementedError(f"{key} input not implemented in preprocesser")
+
+        additional_inputs["text"] = replace_action_token(
+            additional_inputs["text"],
+            additional_inputs["action_chunk"],
+            self.train_action_tokenizer if self.use_fast_tokenizer else None,
+            ["x2_normal"] * additional_inputs["text"].__len__(),
+            additional_inputs["dof_mask"],
+        )
+
+        inputs = preprocesser_call(
+            processor=self.processor,
+            text=additional_inputs.pop("text"),
+            images=additional_inputs.pop("image_inputs"),
+            videos=None,
+            padding=True,
+            truncation=True,
+            return_tensors="pt",
+            max_length=self.dataload_config.get("max_length", 768),
+        )
+
+        action_token_id = self.processor.tokenizer.convert_tokens_to_ids("<|action|>")
+
+        # Gating token types
+        additional_inputs["moe_token_types"] = inputs.input_ids == action_token_id
+
+        inputs.update(additional_inputs)
+
+        inputs["dataset_names"] = ["x2_normal"] * inputs["action_chunk"].shape[0]
+
+        return inputs
+
+
+def load_lerobot_data(
+    config,
+    lerobot_config,
+    rank=0,
+    world_size=1,
+    seed=42,
+):
+    """
+    Load LeRobot dataset with distributed support
+
+    Args:
+        config: Model configuration
+        rank: Current process rank (default: 0)
+        world_size: Total number of processes (default: 1)
+        seed: Random seed for reproducibility (default: 42)
+
+    Returns:
+        dataset: Training dataset
+        train_num: Number of training samples per process
+        sampler: Distributed sampler (None if world_size=1)
+    """
+
+    # Set seed for reproducibility
+    torch.manual_seed(seed)
+
+    dataset_fps = 50
+    dataload_config = get_data_configs(config["data"])
+
+    delta_timestamps = {
+        # action chunk
+        "action": [t / dataset_fps for t in range(dataload_config.get("action_horizon", 32) - 1)],
+    }
+    batch_size = config.get("batch_size_per_gpu", 8)
+
+    # repo_id = "lerobot/aloha_mobile_cabinet"
+    repo_id = lerobot_config.get("repo_id", "lerobot/aloha_mobile_cabinet")
+    dataset = LeRobotDataset(repo_id, delta_timestamps=delta_timestamps, video_backend="pyav")
+
+    if rank == 0:
+        print(f"Selected episodes: {dataset.episodes}")
+        print(f"Number of episodes selected: {dataset.num_episodes}")
+        print(f"Number of frames selected: {dataset.num_frames}")
+
+    dataset = PreprocessedDataset(dataset, config, dataload_config, seed=seed, rank=rank, world_size=world_size)
+
+    # Calculate samples per process
+    if world_size > 1:
+        # With DistributedSampler, each process gets approximately len(dataset) // world_size samples
+        samples_per_process = len(dataset) // world_size
+        train_num = samples_per_process // batch_size
+    else:
+        train_num = len(dataset) // batch_size
+
+    if rank == 0:
+        print("\n" + "=" * 50)
+        print("LeRobot Data Loading Configuration:")
+        print(f"✦ RANK: {rank}")
+        print(f"✦ WORLD SIZE: {world_size}")
+        print(f"✦ BATCH SIZE PER GPU: {batch_size}")
+        print(f"✦ REPO ID: {repo_id}")
+        print(f"✦ TOTAL DATASET SIZE: {len(dataset)}")
+        if world_size > 1:
+            print(f"✦ SAMPLES PER PROCESS: {samples_per_process}")
+            print(f"✦ BATCHES PER PROCESS: {train_num}")
+            print(f"✦ TOTAL BATCHES (ALL PROCESSES): {train_num * world_size}")
+        else:
+            print(f"✦ TOTAL BATCHES: {train_num}")
+        print(f"✦ SEED: {seed}")
+        print("=" * 50 + "\n")
+
+    return dataset, train_num
+
+
+def get_distributed_dataloader(dataset, config, rank=0, world_size=1, seed=42, is_train=True):
+    """
+    Helper function to get distributed dataloader
+
+    Args:
+        dataset: PreprocessedDataset instance
+        config: Configuration dict
+        rank: Current process rank
+        world_size: Total number of processes
+        seed: Random seed
+        is_train: Whether this is for training (affects shuffling)
+
+    Returns:
+        dataloader: Distributed DataLoader
+        sampler: DistributedSampler
+    """
+    if is_train:
+        return dataset.get_train_dataloader(rank=rank, world_size=world_size, seed=seed)
+    else:
+        return dataset.get_val_dataloader(rank=rank, world_size=world_size)
+
+
+def get_data_configs(config):
+    default_data_config = {
+        "train_test_split": 0.95,
+        "split_seed": 42,
+        "batch_size": 8,
+        "action_horizon": 21,
+        "action_history_length": 0,
+        "image_horizon": 1,
+        "image_history_length": 0,
+        "left_padding": False,
+        "right_padding": False,
+        "return_first_obs": False,
+        "return_last_obs": False,
+        "randomize_obs_after": None,
+        "datasets": [],
+        "labeled_pathes": [],
+    }
+    data_config = default_data_config | config
+    data_config["action_horizon"] += 1
+
+    return data_config
+
+class TestDataset(PreprocessedDataset):
+    def __init__(self, dataset, config, dataload_config, seed=42):
+        super().__init__(dataset, config, dataload_config, seed=seed, rank=0, world_size=1)
+
+    def get_dataloader(self):
+        """
+        Get distributed evaluation dataloader (no shuffling for consistent evaluation)
+        """
+        
+        dataloader = torch.utils.data.DataLoader(
+            self,
+            batch_size=1,
+            collate_fn=DataCollator(self.config, self.dataload_config, self._dataset.meta.stats),
+        )
+
+        return dataloader
+    
+def load_test_dataset(
+    config,
+    lerobot_config,
+    seed=42,
+    episode=0,
+):
+    """
+    Load test dataset
+
+    Args:
+        config: Model configuration
+        seed: Random seed for reproducibility (default: 42)
+
+    Returns:
+        dataset: Test dataset
+    """
+
+    # Set seed for reproducibility
+    torch.manual_seed(seed)
+
+    dataset_fps = 50
+    dataload_config = get_data_configs(config["data"])
+
+    delta_timestamps = {
+        # action chunk
+        "action": [t / dataset_fps for t in range(dataload_config.get("action_horizon", 32) - 1)],
+    }
+
+    repo_id = lerobot_config.get("repo_id", "lerobot/aloha_mobile_cabinet")
+    dataset = LeRobotDataset(repo_id, episodes=[episode], delta_timestamps=delta_timestamps, video_backend="pyav")
+
+    print(f"Selected episodes: {dataset.episodes}")
+    print(f"Number of episodes selected: {dataset.num_episodes}")
+    print(f"Number of frames selected: {dataset.num_frames}")
+    
+    dataset = TestDataset(dataset, config, dataload_config, seed=seed)
+    
+    return dataset

wall_x/data/utils.py

@@ -0,0 +1,560 @@
+"""
+Data processing utilities for Wall-X multimodal robotic learning.
+
+This module provides utilities for preprocessing text, images, and actions
+for multimodal transformer models in robotic learning tasks.
+"""
+
+import re
+import torch
+import random
+from collections import OrderedDict
+from typing import List, Dict, Any, Optional, Union, Tuple
+from transformers import BatchFeature
+
+
+CAMERA_NAME_MAPPING = {
+    "face_view": "front view",
+    "left_wrist_view": "left wrist view", 
+    "right_wrist_view": "right wrist view",
+    "move1_view": "move view",
+    "move2_view": "move view", 
+    "wall_view": "wall view",
+    "top_view": "top view",
+}
+
+
+MULTIMODAL_DATASET_NAMES = [
+    "x2_multimodal_from_action",
+    "x2_multimodal",
+    "x2_subtask_generation",
+    "multimodal_CapsFusion",
+    "multimodal_Robo2VLM",
+    "multimodal_RoboPoint",
+    "multimodal_EQA",
+    "multimodal_Cambrian",
+    "multimodal_pixmo",
+    "multimodal_VQAv2",
+    "multimodal_COCO",
+]
+
+FREQUENCY_MAPPING = {
+    "x2_normal": 32,
+    "fractal": 5,
+    "bridge_data_v2": 5,
+    "droid": 15,
+    "agibotworld_alpha": 32,
+    "DobbE": 30,
+    "RH20T": 10,
+    "UMI-biarm": 10,
+    "austin_buds": 20,
+    "austin_sailor": 20,
+    "austin_sirius": 20,
+    "bc_z": 10,
+    "berkeley_autolab_ur5": 5,
+    "berkeley_cable_routing": 10,
+    "berkeley_fanuc_manipulation": 10,
+    "dlr_edan_shared_control": 5,
+    "fmb": 10,
+    "furniture_bench": 10,
+    "jaco_play": 10,
+    "nyu_rot": 10,
+    "stanford_hydra": 10,
+    "stanford_kuka_multimodal": 20,
+    "taco_play": 30,
+    "utaustin_mutex": 20,
+    "viola": 20,
+}
+
+
+def preprocesser_call(
+    processor,
+    images: Optional[Union[List, Any]] = None,
+    text: Optional[Union[str, List[str]]] = None,
+    videos: Optional[Union[List, Any]] = None,
+    padding: Union[bool, str] = False,
+    truncation: Optional[bool] = None,
+    max_length: Optional[int] = None,
+    return_tensors: str = "pt",
+) -> BatchFeature:
+    """Unified preprocessing function for Wall-X model handling text, image and video inputs.
+    
+    Processes inputs into format suitable for multimodal transformer models, including:
+    - Text tokenization and special token handling
+    - Image/video processing through image processor
+    - Attention mask and label generation
+    - Padding and truncation handling
+    
+    Args:
+        processor: Multimodal processor containing tokenizer and image processor
+        images: Input images (PIL, numpy arrays, or torch tensors)
+        text: Text or list of texts to tokenize
+        videos: Input videos (numpy arrays or torch tensors)
+        padding: Whether to pad sequences to same length
+        truncation: Whether to truncate sequences longer than max_length
+        max_length: Maximum length for truncation/padding
+        return_tensors: Format for returned tensors ('pt', 'np', etc.)
+        
+    Returns:
+        BatchFeature containing processed inputs with keys:
+        - input_ids: Tokenized text
+        - attention_mask: Attention mask for text
+        - pixel_values: Processed image pixels
+        - pixel_values_videos: Processed video frames
+        - image_grid_thw: Image grid dimensions for LLM
+        - video_grid_thw: Video grid dimensions for LLM
+        - labels: Training labels with masking
+    """
+    # Process image inputs
+    if images is not None and len(images) > 0:
+        image_inputs = processor.image_processor(
+            images=images, videos=None, return_tensors=return_tensors
+        )
+        image_grid_thw = image_inputs["image_grid_thw"]
+    else:
+        image_inputs = {}
+        image_grid_thw = None
+
+    # Process video inputs
+    if videos is not None:
+        videos_inputs = processor.image_processor(
+            images=None, videos=videos, return_tensors=return_tensors
+        )
+        video_grid_thw = videos_inputs["video_grid_thw"]
+    else:
+        videos_inputs = {}
+        video_grid_thw = None
+
+    # Ensure text input is in list format
+    if not isinstance(text, list):
+        text = [text]
+
+    # Process image placeholder tokens in text
+    if image_grid_thw is not None:
+        merge_length = processor.image_processor.merge_size ** 2
+        index = 0
+        for i in range(len(text)):
+            while "<|image_pad|>" in text[i]:
+                # Add bounds checking to avoid index overflow
+                if index >= len(image_grid_thw):
+                    print(f"Warning: Number of image placeholders ({index + 1}) "
+                          f"exceeds actual images ({len(image_grid_thw)}), "
+                          f"skipping remaining placeholder processing")
+                    break
+                # Replace image placeholder with actual token count
+                token_count = image_grid_thw[index].prod() // merge_length
+                text[i] = text[i].replace(
+                    "<|image_pad|>", "<|placeholder|>" * token_count, 1
+                )
+                index += 1
+            text[i] = text[i].replace("<|placeholder|>", "<|image_pad|>")
+
+    # Process video placeholder tokens in text
+    if video_grid_thw is not None:
+        merge_length = processor.image_processor.merge_size ** 2
+        index = 0
+        for i in range(len(text)):
+            while "<|video_pad|>" in text[i]:
+                # Replace video placeholder with actual token count
+                token_count = video_grid_thw[index].prod() // merge_length
+                text[i] = text[i].replace(
+                    "<|video_pad|>", "<|placeholder|>" * token_count, 1
+                )
+                index += 1
+            text[i] = text[i].replace("<|placeholder|>", "<|video_pad|>")
+
+    # Tokenize complete input text
+    text_inputs = processor.tokenizer(
+        text,
+        return_tensors=return_tensors,
+        padding=padding,
+        truncation=truncation,
+        max_length=max_length
+    )
+
+    # Get pad token ID for label generation
+    pad_token_id = processor.tokenizer.pad_token_id
+    if pad_token_id is None:
+        pad_token_id = processor.tokenizer.eos_token_id
+
+    # Generate labels for multi-turn dialogue, keeping only assistant response loss
+    labels = torch.full_like(text_inputs.input_ids, -100)
+    assistant_marker = "<|im_start|>assistant\n"
+    im_end_token_id = processor.tokenizer.convert_tokens_to_ids("<|im_end|>")
+    assistant_tokens = processor.tokenizer(
+        "<|im_start|>assistant\n", add_special_tokens=False
+    ).input_ids
+
+    for i in range(len(text)):
+        assistant_regions = []
+        parts = text[i].split(assistant_marker)
+
+        # Process each part to determine which tokens belong to assistant responses
+        # Count left padding tokens
+        num_left_pads = 0
+        for token_id in text_inputs.input_ids[i]:
+            if token_id == pad_token_id:
+                num_left_pads += 1
+            else:
+                break
+        current_pos = num_left_pads
+
+        for j, part in enumerate(parts):
+            part_tokens = processor.tokenizer(part, add_special_tokens=False).input_ids
+            if j == 0:
+                # First part is system prompt or user question, all labels are -100
+                current_pos += len(part_tokens)
+                continue
+
+            # From second part onwards, each part starts with assistant response
+            for k in range(current_pos + 1, len(text_inputs.input_ids[i])):
+                if text_inputs.input_ids[i][k] == im_end_token_id:
+                    assistant_regions.append((
+                        current_pos + len(assistant_tokens), k + 2
+                    ))
+                    break
+            current_pos += len(part_tokens) + 3
+
+        # Set labels for assistant response regions
+        for start, end in assistant_regions:
+            labels[i][start:end] = text_inputs.input_ids[i][start:end]
+
+    # Mask special action tokens in labels
+    action_token_id = processor.tokenizer.encode("<|action|>")[0]
+    propri_token_id = processor.tokenizer.encode("<|propri|>")[0]
+    labels[labels == action_token_id] = -100
+    labels[labels == propri_token_id] = -100
+    labels[labels == processor.tokenizer.pad_token_id] = -100
+
+    # Set labels to None if all are invalid to skip cross entropy loss
+    if (labels != -100).any().item():
+        text_inputs["labels"] = labels
+    else:
+        text_inputs["labels"] = None
+
+    return BatchFeature(data={**text_inputs, **image_inputs, **videos_inputs})
+
+
+def process_grounding_points(text: str, orig_height: int, orig_width: int, resized_height: int, resized_width: int, model_type: str) -> str:
+    """Process grounding point coordinates in text based on image resizing.
+
+    Adjusts coordinate values in <point> tags to match resized image dimensions
+    for different model types (qwen2, qwen2_5).
+
+    Args:
+        text: Input text containing <point> tags with coordinates
+        orig_height: Original image height
+        orig_width: Original image width
+        resized_height: Resized image height
+        resized_width: Resized image width
+        model_type: Model type for coordinate processing ('qwen2' or 'qwen2_5')
+
+    Returns:
+        Text with adjusted coordinate values
+    """
+    # Regex pattern to match <point> tags and their contents
+    point_pattern = re.compile(r"<point>(.*?)</point>")
+
+    def process_match(match):
+        """Process a single point match and adjust coordinates."""
+        coords_str = match.group(1)
+        try:
+            # Extract coordinates from string
+            coords = list(map(int, re.findall(r"\d+", coords_str)))
+
+            # Calculate resize scale factors
+            scale_w = resized_width / orig_width
+            scale_h = resized_height / orig_height
+
+            if len(coords) == 2:
+                x, y = coords
+                if model_type == "qwen2_5":
+                    # Qwen2.5 uses pixel coordinates
+                    new_x = max(0, min(round(x * scale_w), resized_width - 1))
+                    new_y = max(0, min(round(y * scale_h), resized_height - 1))
+                elif model_type == "qwen2":
+                    # Qwen2 normalizes to [0, 1000) range
+                    new_x = max(0, min(999.999, (x / orig_width) * 1000))
+                    new_y = max(0, min(999.999, (y / orig_height) * 1000))
+                else:
+                    raise ValueError(f"Unsupported model type: {model_type}")
+                coords = [new_x, new_y]
+
+            elif len(coords) == 4:
+                x1, y1, x2, y2 = coords
+                if model_type == "qwen2_5":
+                    new_x1 = max(0, min(round(x1 * scale_w), resized_width - 1))
+                    new_y1 = max(0, min(round(y1 * scale_h), resized_height - 1))
+                    new_x2 = max(0, min(round(x2 * scale_w), resized_width - 1))
+                    new_y2 = max(0, min(round(y2 * scale_h), resized_height - 1))
+                elif model_type == "qwen2":
+                    new_x1 = max(0, min(999.999, (x1 / orig_width) * 1000))
+                    new_y1 = max(0, min(999.999, (y1 / orig_height) * 1000))
+                    new_x2 = max(0, min(999.999, (x2 / orig_width) * 1000))
+                    new_y2 = max(0, min(999.999, (y2 / orig_height) * 1000))
+                else:
+                    raise ValueError(f"Unsupported model type: {model_type}")
+                coords = [new_x1, new_y1, new_x2, new_y2]
+
+            # Return processed point tag
+            return f'<point>[{", ".join(map(str, coords))}]</point>'
+
+        except (ValueError, TypeError):
+            # Return original content if processing fails
+            return match.group(0)
+
+    # Replace all matching point tags
+    processed_text = point_pattern.sub(process_match, text)
+    return processed_text
+
+
+def get_frame_instruction(
+    instruction_info: Dict[str, Any], frame_idx: Optional[int] = None, truncate_keys: Optional[List[str]] = None
+) -> Tuple[Dict[str, Any], Optional[int]]:
+    """Extract frame-specific instruction from instruction dictionary.
+
+    Args:
+        instruction_info: Dictionary containing instruction components
+        frame_idx: Current frame index
+        truncate_keys: Keys that trigger truncation when found
+
+    Returns:
+        Tuple of (frame_instruction_dict, split_end_frame)
+    """
+    if truncate_keys is None:
+        truncate_keys = ["subtask_generation", "distribute", "subtask_generation_zh", "distribute_zh"]
+
+    instruction_for_frame = {}
+    split_end = None
+
+    for key, value in instruction_info.items():
+        if isinstance(value, dict):
+            # Handle frame-range specific instructions
+            for frame_range, frame_instruction in value.items():
+                start_frame, end_frame = map(int, frame_range.split(" "))
+                if start_frame <= frame_idx < end_frame or (start_frame == frame_idx):
+                    instruction_for_frame[key] = frame_instruction
+                    if truncate_keys is not None and split_end is None and key in truncate_keys:
+                        split_end = end_frame + 1
+                    break
+        else:
+            instruction_for_frame[key] = value
+
+    return instruction_for_frame, split_end
+
+
+def get_task_instruction(frame_instruction_info: Dict[str, Any], priority_order: Optional[OrderedDict] = None) -> str:
+    """Construct task instruction from available instruction fields using priority sampling.
+
+    Args:
+        frame_instruction_info: Dictionary containing instruction fields
+        priority_order: OrderedDict specifying sampling probability for each field
+
+    Returns:
+        Combined instruction string with priority components
+    """
+    # Default priority settings
+    default_priority_order = OrderedDict(
+        {
+            "subtask_generation": 0.25,
+            "subtask_generation_zh": 0.25,
+            "distribute": 0.25,
+            "distribute_zh": 0.25,
+        }
+    )
+
+    if priority_order is not None:
+        priority_order = OrderedDict(priority_order)
+    else:
+        priority_order = default_priority_order
+
+    got_instruction = False
+    task_instruction = ""
+
+    # Sample instruction components based on priority probabilities
+    for key, prob in priority_order.items():
+        if key in frame_instruction_info and frame_instruction_info[key] != "":
+            if got_instruction:
+                if random.random() >= prob:
+                    continue
+
+            task_instruction += f"\n{frame_instruction_info[key]}"
+            got_instruction = True
+            break
+
+    # Fall back to base instruction if no priority components found
+    if not got_instruction:
+        task_instruction = frame_instruction_info.get("instruction", "")
+
+    return task_instruction
+
+
+def get_wallx_normal_text(
+    instruction_info: Dict[str, Any],
+    action_chunk_size: int,
+    frame_idx: int,
+    priority_order: Optional[OrderedDict] = None,
+    cam_mapping: Optional[Dict[str, str]] = None,
+    generate_subtask_ratio: float = 0.0,
+) -> Tuple[str, bool]:
+    """Construct complete multimodal prompt text for Wall-X model.
+
+    Formats input using special tokens including:
+    - System message
+    - User observations (with image placeholders)
+    - Task instructions
+    - Proprioception prompts
+    - Assistant responses (with action tokens)
+
+    Args:
+        instruction_info: Dictionary containing instruction components
+        action_chunk_size: Number of action tokens to generate
+        frame_idx: Current frame index
+        priority_order: Priority order for instruction sampling
+        cam_mapping: Camera name mapping dictionary
+        generate_subtask_ratio: Probability of generating subtask instead of actions
+
+    Returns:
+        Tuple of (formatted_prompt_text, is_subtask_generation)
+    """
+    # Special tokens for formatting
+    role_start_symbol = "<|im_start|>"
+    role_end_symbol = "<|im_end|>"
+    vision_start_symbol = "<|vision_start|>"
+    vision_end_symbol = "<|vision_end|>"
+    image_pad_symbol = "<|image_pad|>"
+    propri_symbol = "<|propri|>"
+    action_symbol = "<|action|>"
+    action_fast_symbol = "<|action_fast|>"
+
+    # System prologue
+    prologue = f"{role_start_symbol}system\nYou are a helpful assistant.{role_end_symbol}\n"
+
+    # User request with observation
+    user_request = f"{role_start_symbol}user\nObservation:"
+    if cam_mapping:
+        for _, cam_name in cam_mapping.items():
+            view_name = CAMERA_NAME_MAPPING.get(cam_name, cam_name)
+            user_request += f" {view_name}: {vision_start_symbol}{image_pad_symbol}{vision_end_symbol}"
+    user_request += "\nInstruction:"
+
+    # Get frame-specific instruction
+    frame_instruction_info, _ = get_frame_instruction(instruction_info, frame_idx=frame_idx)
+
+    generate_subtask = False
+    priority_keys = ["subtask_generation", "distribute"]
+
+    # Decide whether to generate subtask or actions
+    if bool(set(frame_instruction_info.keys()) & set(priority_keys)) and random.random() < generate_subtask_ratio:
+        # Generate subtask (equivalent to VQA task)
+        instruction = frame_instruction_info.get("instruction", "")
+        text_prompt = "\nPredict the next action in language.\n"
+        user_message = f"{user_request} {instruction}{text_prompt}{role_end_symbol}\n"
+
+        # Find output instruction from priority keys
+        for key in priority_keys:
+            if key in frame_instruction_info:
+                output_instruction = frame_instruction_info[key]
+                break
+
+        assistant_output = f"{role_start_symbol}assistant\n{output_instruction}\n{role_end_symbol}"
+        generate_subtask = True
+    else:
+        # Generate actions
+        instruction = get_task_instruction(frame_instruction_info, priority_order=priority_order)
+        text_prompt = f"\nPredict the next action in robot action.\nProprioception: {propri_symbol}\n"
+        user_message = f"{user_request} {instruction}{text_prompt}{role_end_symbol}\n"
+        assistant_output = f"{role_start_symbol}assistant\n{action_fast_symbol}{role_end_symbol}\n{action_symbol * action_chunk_size}"
+
+    complete_text = prologue + user_message + assistant_output
+    return complete_text, generate_subtask
+
+
+def get_action_tokens(normalized_actions: Union[torch.Tensor, List], action_tokenizer) -> List[List[str]]:
+    """Convert normalized actions to action token strings.
+
+    Args:
+        normalized_actions: Normalized action arrays/tensors
+        action_tokenizer: Tokenizer for converting actions to tokens
+
+    Returns:
+        List of action token string lists for each sample
+    """
+    if isinstance(normalized_actions, torch.Tensor):
+        normalized_actions = normalized_actions.cpu().numpy()
+
+    all_action_tokens = []
+    for i in range(len(normalized_actions)):
+        if isinstance(normalized_actions[i], torch.Tensor):
+            normalized_actions[i] = normalized_actions[i].cpu().numpy()
+
+        token_id = action_tokenizer(normalized_actions[i])
+        action_tokens = [f"<|action_token_{j}|>" for j in token_id[0]]
+        all_action_tokens.append(action_tokens)
+
+    return all_action_tokens
+
+
+def pad_action_token_strs(actions_token_lists: List[List[str]], pad_token: str = "<|endoftext|>") -> List[str]:
+    """Pad action token lists to same length and join as strings.
+
+    Args:
+        actions_token_lists: List of action token lists for each sample
+        pad_token: Token used for padding
+
+    Returns:
+        List of padded action token strings
+    """
+    max_len = max(len(tokens) for tokens in actions_token_lists)
+    padded_action_strs = []
+
+    for tokens in actions_token_lists:
+        padded_tokens = tokens + ["<|im_end|>\n"] + [pad_token] * (max_len - len(tokens))
+        padded_action_strs.append("".join(padded_tokens))
+
+    return padded_action_strs
+
+
+def replace_action_token(
+    text: List[str], norm_action: Optional[torch.Tensor], action_tokenizer, dataset_names: List[str], dof_masks: Optional[torch.Tensor] = None
+) -> List[str]:
+    """Replace action placeholders in text with actual action tokens.
+
+    Args:
+        text: List of text strings with action placeholders
+        norm_action: Normalized action tensors
+        action_tokenizer: Tokenizer for converting actions to tokens
+        dataset_names: Names of datasets for each sample
+        dof_masks: Masks for degrees of freedom
+
+    Returns:
+        List of text strings with action tokens replaced
+    """
+    # Filter out multimodal dataset names
+    dataset_names = [name for name in dataset_names if name not in MULTIMODAL_DATASET_NAMES]
+
+    # Get required action chunk sizes
+    required_chunk_sizes = [FREQUENCY_MAPPING.get(name, 32) for name in dataset_names]
+
+    if action_tokenizer is not None and norm_action is not None:
+        # Extract actions based on chunk sizes and DOF masks
+        norm_action = [action[: required_chunk_sizes[i], dof_masks[i, 0].bool()] for i, action in enumerate(norm_action)]
+
+        # Convert to action tokens and pad
+        actions_fast_tokens = get_action_tokens(norm_action, action_tokenizer)
+        actions_fast_token_strs = pad_action_token_strs(actions_fast_tokens)
+
+        # Replace action placeholders with actual tokens
+        actions_fast_token_idx = 0
+        for i in range(len(text)):
+            if "<|action_fast|>" in text[i]:
+                text[i] = text[i].replace("<|action_fast|><|im_end|>\n", actions_fast_token_strs[actions_fast_token_idx])
+                actions_fast_token_idx += 1
+
+        # Remove remaining action placeholders
+        text = [t.replace("<|action|>", "") for t in text]
+    else:
+        # Remove action placeholders when no tokenizer available
+        text = [t.replace("<|action_fast|><|im_end|>\n", "") for t in text]
+
+    return text

wall_x/fusions/backend.py

@@ -0,0 +1,270 @@
+"""
+High-performance C++ backend interface for optimized matrix operations.
+
+This module provides Python bindings for custom CUDA kernels optimized for
+transformer and MoE (Mixture of Experts) operations, including:
+- Asymmetric dual expert operations
+- Token permutation/unpermutation for MoE routing
+- RoPE (Rotary Position Embedding) operations
+"""
+
+import torch
+from typing import Tuple, Optional
+import wallx_csrc as backend
+
+
+
+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]:
+    """
+    Allocate output tensors for asymmetric dual expert GEMM operations.
+    
+    This function handles the case where two experts may have different output
+    dimensions, which is common in heterogeneous MoE architectures.
+    
+    Args:
+        input_expert0 (torch.Tensor): Expert 0 input tensor of shape [m0, k]
+        input_expert1 (torch.Tensor): Expert 1 input tensor of shape [m1, k]
+        weight_expert0 (torch.Tensor): Expert 0 weight tensor of shape [k, n0]
+        weight_expert1 (torch.Tensor): Expert 1 weight tensor of shape [k, n1]
+        
+    Returns:
+        Tuple[torch.Tensor, torch.Tensor]: Pre-allocated output tensors
+            - output_expert0: Shape [m0, n0]
+            - output_expert1: Shape [m1, n1]
+            
+    Raises:
+        AssertionError: If tensor dimensions are incompatible
+    """
+    # Validate input tensor dimensions
+    assert input_expert0.ndim == 2, "Expected 2D tensor for input_expert0"
+    assert input_expert1.ndim == 2, "Expected 2D tensor for input_expert1"
+    assert weight_expert0.ndim == 2, "Expected 2D tensor for weight_expert0"
+    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)}"
+    
+    # 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)
+    
+    return output_expert0, output_expert1
+
+
+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]:
+    """
+    Asymmetric dual expert grouped GEMM with separated inputs and outputs.
+    
+    This is the recommended interface for maximum flexibility and performance when
+    dealing with two experts that may have different intermediate dimensions.
+    The operation is equivalent to:
+        output_expert0 = input_expert0 @ weight_expert0
+        output_expert1 = input_expert1 @ weight_expert1
+    But optimized as a single fused kernel call.
+    
+    Args:
+        input_expert0 (torch.Tensor): Expert 0 input tensor of shape [m0, k]
+        input_expert1 (torch.Tensor): Expert 1 input tensor of shape [m1, k]
+        weight_expert0 (torch.Tensor): Expert 0 weight tensor of shape [k, n0]
+        weight_expert1 (torch.Tensor): Expert 1 weight tensor of shape [k, n1]
+                                      Note: n0 can be different from n1
+        output_expert0 (torch.Tensor, optional): Pre-allocated output for expert 0 [m0, n0]
+        output_expert1 (torch.Tensor, optional): Pre-allocated output for expert 1 [m1, n1]
+        trans_a (bool, optional): Whether to transpose input tensors. Defaults to False.
+        trans_b (bool, optional): Whether to transpose weight tensors. Defaults to False.
+        
+    Returns:
+        Tuple[torch.Tensor, torch.Tensor]: Output tensors (output_expert0, output_expert1)
+        
+    Example:
+        >>> # Two experts with different output dimensions
+        >>> input0 = torch.randn(512, 1024, device='cuda')  # 512 tokens for expert 0
+        >>> input1 = torch.randn(256, 1024, device='cuda')  # 256 tokens for expert 1
+        >>> weight0 = torch.randn(1024, 2048, device='cuda')  # Expert 0: 1024->2048
+        >>> weight1 = torch.randn(1024, 4096, device='cuda')  # Expert 1: 1024->4096
+        >>> out0, out1 = asym_dual_gmm_separated(input0, input1, weight0, weight1)
+    """
+    # Allocate outputs if not provided
+    if output_expert0 is None or output_expert1 is None:
+        alloc_out0, alloc_out1 = _allocate_asymmetric_dual_outputs(
+            input_expert0, input_expert1, weight_expert0, weight_expert1
+        )
+        if output_expert0 is None:
+            output_expert0 = alloc_out0
+        if output_expert1 is None:
+            output_expert1 = alloc_out1
+    
+    # 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
+    )
+    
+    return output_expert0, output_expert1
+
+
+def permute(input: torch.Tensor, 
+            indices: torch.Tensor, 
+            num_out_tokens: int, 
+            workspace: torch.Tensor, 
+            max_expanded_token_num: int) -> torch.Tensor:
+    """
+    Permute input tokens according to expert assignment indices for MoE routing.
+    
+    This function reorders tokens based on their assigned experts to enable
+    efficient grouped processing. Used in the forward pass of MoE layers.
+    
+    Args:
+        input (torch.Tensor): Input tokens to permute
+        indices (torch.Tensor): Expert assignment indices for each token
+        num_out_tokens (int): Number of output tokens after expansion
+        workspace (torch.Tensor): Temporary workspace tensor for intermediate computations
+        max_expanded_token_num (int): Maximum number of tokens after top-k expansion
+        
+    Returns:
+        torch.Tensor: Permuted tokens grouped by expert assignment
+        
+    Note:
+        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)
+
+
+def unpermute(input: torch.Tensor, 
+              row_id_map: torch.Tensor, 
+              prob: torch.Tensor, 
+              max_tokens: int, 
+              num_topK: int) -> torch.Tensor:
+    """
+    Unpermute expert outputs back to original token order with probability weighting.
+    
+    This function reverses the permutation applied in the forward pass and combines
+    outputs from multiple experts using their routing probabilities.
+    
+    Args:
+        input (torch.Tensor): Permuted expert outputs to unpermute
+        row_id_map (torch.Tensor): Mapping from permuted positions to original positions
+        prob (torch.Tensor): Expert routing probabilities for weighted combination
+        max_tokens (int): Maximum number of tokens in the sequence
+        num_topK (int): Number of top experts selected per token
+        
+    Returns:
+        torch.Tensor: Unpermuted tokens in original order with expert outputs combined
+        
+    Note:
+        The output combines multiple expert predictions for each token using
+        the routing probabilities as weights.
+    """
+    return backend.unpermute(input, row_id_map, prob, max_tokens, num_topK)
+
+
+def unpermute_bwd(input_bwd: torch.Tensor, 
+                  input_fwd: torch.Tensor, 
+                  row_id_map: torch.Tensor, 
+                  prob: Optional[torch.Tensor]) -> torch.Tensor:
+    """
+    Backward pass for unpermute operation with gradient flow.
+    
+    This function handles the backward pass through the unpermute operation,
+    ensuring proper gradient flow for training MoE models.
+    
+    Args:
+        input_bwd (torch.Tensor): Backward gradients from the next layer
+        input_fwd (torch.Tensor): Forward pass inputs (for gradient computation)
+        row_id_map (torch.Tensor): Row mapping used in forward unpermute
+        prob (torch.Tensor, optional): Expert probabilities. If None, uniform weights are used.
+        
+    Returns:
+        torch.Tensor: Gradients with respect to the input of unpermute forward pass
+        
+    Note:
+        If prob is None, uniform probabilities are assumed for gradient computation.
+    """
+    # 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)
+    
+    return backend.unpermute_bwd(input_bwd, input_fwd, row_id_map, prob)
+
+
+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:
+    """
+    Apply RoPE (Rotary Position Embedding) to query and key tensors.
+    
+    Applies rotary position embeddings to query and key tensors using precomputed
+    cosine and sine values. Supports both standard RoPE and multi-dimensional RoPE (mRoPE).
+    
+    Args:
+        q (torch.Tensor): Query tensor to apply RoPE to
+        k (torch.Tensor): Key tensor to apply RoPE to  
+        cos (torch.Tensor): Precomputed cosine values for rotation
+        sin (torch.Tensor): Precomputed sine values for rotation
+        q_out (torch.Tensor): Output tensor for rotated queries (in-place operation supported)
+        k_out (torch.Tensor): Output tensor for rotated keys (in-place operation supported)
+        mrope_section_doubled (bool): Whether using multi-dimensional RoPE with doubled sections
+        
+    Note:
+        This function performs in-place operations if q_out and k_out point to the same
+        memory as q and k respectively. The rotation is applied using the standard
+        RoPE formulation with complex number rotation.
+    """
+    return backend.rope(q, k, cos, sin, q_out, k_out, mrope_section_doubled)
+
+
+def rope_bwd(grad_q_out: torch.Tensor, 
+             grad_k_out: torch.Tensor, 
+             q: torch.Tensor, 
+             k: torch.Tensor, 
+             cos: torch.Tensor, 
+             sin: torch.Tensor, 
+             grad_q: torch.Tensor, 
+             grad_k: torch.Tensor, 
+             mrope_section_doubled: bool) -> None:
+    """
+    Backward pass for RoPE operation with gradient computation.
+    
+    Computes gradients with respect to the input query and key tensors
+    for the RoPE operation used in transformer attention mechanisms.
+    
+    Args:
+        grad_q_out (torch.Tensor): Gradient with respect to output queries
+        grad_k_out (torch.Tensor): Gradient with respect to output keys
+        q (torch.Tensor): Original query tensor from forward pass
+        k (torch.Tensor): Original key tensor from forward pass
+        cos (torch.Tensor): Cosine values used in forward pass
+        sin (torch.Tensor): Sine values used in forward pass
+        grad_q (torch.Tensor): Output tensor for query gradients
+        grad_k (torch.Tensor): Output tensor for key gradients
+        mrope_section_doubled (bool): Whether using multi-dimensional RoPE configuration
+        
+    Note:
+        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)

wall_x/fusions/ops.py

@@ -0,0 +1,397 @@
+import torch
+import warnings
+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):
+        """
+        Forward pass for asymmetric dual expert GEMM.
+
+        Args:
+            input_expert0: Expert 0 input [m0, k]
+            input_expert1: Expert 1 input [m1, k]
+            weight_expert0: Expert 0 weight [k, n0] or [n0, k] if trans_b=True
+            weight_expert1: Expert 1 weight [k, n1] or [n1, k] if trans_b=True
+            trans_b: Whether to transpose the weight matrices
+
+        Returns:
+            Tuple of (output_expert0, output_expert1)
+        """
+        # Validate inputs
+        assert input_expert0.dim() == 2, "input_expert0 must be 2D"
+        assert input_expert1.dim() == 2, "input_expert1 must be 2D"
+        assert weight_expert0.dim() == 2, "weight_expert0 must be 2D"
+        assert weight_expert1.dim() == 2, "weight_expert1 must be 2D"
+
+        # 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)"
+        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)"
+
+        # Save tensors and trans_b for backward pass
+        ctx.save_for_backward(input_expert0, input_expert1, weight_expert0, weight_expert1)
+        ctx.trans_b = trans_b
+
+        # Allocate output tensors
+        m0 = input_expert0.size(0)
+        m1 = input_expert1.size(0)
+        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)
+
+        # 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)
+
+        return output_expert0, output_expert1
+
+    @staticmethod
+    def backward(ctx, grad_output_expert0, grad_output_expert1):
+        """
+        Optimized backward pass using specialized kernels.
+        Always computes all gradients to minimize kernel calls.
+        """
+        grad_output_expert0 = grad_output_expert0.contiguous()
+        grad_output_expert1 = grad_output_expert1.contiguous()
+
+        input_expert0, input_expert1, weight_expert0, weight_expert1 = ctx.saved_tensors
+        trans_b = ctx.trans_b
+
+        # Always allocate all gradient tensors (no conditional computation)
+        grad_input_expert0 = torch.empty_like(input_expert0)
+        grad_input_expert1 = torch.empty_like(input_expert1)
+        grad_weight_expert0 = torch.empty_like(weight_expert0)
+        grad_weight_expert1 = torch.empty_like(weight_expert1)
+
+        # Compute input gradients: grad_input = grad_output @ weight^T (if trans_b=False)
+        #                                    = grad_output @ weight   (if trans_b=True)
+        backend.asym_dual_gmm_separated(
+            grad_output_expert0,
+            grad_output_expert1,
+            weight_expert0,
+            weight_expert1,
+            grad_input_expert0,
+            grad_input_expert1,
+            trans_a=False,
+            trans_b=not trans_b,
+        )
+
+        # Compute weight gradients
+        if trans_b:
+            # When trans_b=True in forward: output = input @ weight^T
+            # So grad_weight^T = input^T @ grad_output
+            # Which means grad_weight = grad_output^T @ input
+            backend.asym_dual_gmm_separated(
+                grad_output_expert0,
+                grad_output_expert1,
+                input_expert0,
+                input_expert1,
+                grad_weight_expert0,
+                grad_weight_expert1,
+                trans_a=True,
+                trans_b=False,
+            )
+        else:
+            # When trans_b=False in forward: output = input @ weight
+            # So grad_weight = input^T @ grad_output
+            backend.asym_dual_gmm_separated(
+                input_expert0,
+                input_expert1,
+                grad_output_expert0,
+                grad_output_expert1,
+                grad_weight_expert0,
+                grad_weight_expert1,
+                trans_a=True,
+                trans_b=False,
+            )
+
+        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):
+    """
+    Convenience function for asymmetric dual expert GEMM.
+
+    Args:
+        input_expert0: Expert 0 input [m0, k]
+        input_expert1: Expert 1 input [m1, k]
+        weight_expert0: Expert 0 weight [k, n0] or [n0, k] if trans_b=True
+        weight_expert1: Expert 1 weight [k, n1] or [n1, k] if trans_b=True
+        trans_b: Whether to transpose the weight matrices
+
+    Returns:
+        Tuple of (output_expert0, output_expert1)
+    """
+    return AsymmetricDualExpertGemm.apply(input_expert0, input_expert1, weight_expert0, weight_expert1, trans_b)
+
+
+################################################################################################
+##
+## PermuteMoE topK
+##
+################################################################################################
+
+
+class PermuteMoE_topK(torch.autograd.Function):
+
+    workspace_fw = None
+    dtype = None
+    max_expanded_token_num = 0
+
+    @staticmethod
+    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]
+        """
+        # Empty input check
+        if not input_act.numel():
+            return input_act, None
+
+        # For top1 case, view the indices as 2D tensor to unify the shape for topk>=2 cases.
+        if indices.dim() == 1:
+            indices = indices.view(-1, 1)
+
+        # Device check
+        if input_act.is_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()
+
+        # 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)}.")
+
+        # 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.")
+            indices = indices.to(torch.int32)
+
+        # Contiguous check
+        if not input_act.is_contiguous():
+            warnings.warn("The input `input_act` of permute_topK op is discontiguous!")
+            input_act = input_act.contiguous()
+        if not indices.is_contiguous():
+            warnings.warn("The input `indices` of permute_topK op is discontiguous!")
+            indices = indices.contiguous()
+
+        num_topK = indices.size(1)
+
+        input_max_expanded_token_num = max(max_token_num, input_act.size(0)) * num_topK
+        if PermuteMoE_topK.max_expanded_token_num < input_max_expanded_token_num:
+            PermuteMoE_topK.max_expanded_token_num = input_max_expanded_token_num
+            PermuteMoE_topK.workspace_fw = []
+
+        if PermuteMoE_topK.dtype != input_act.dtype:
+            PermuteMoE_topK.dtype = input_act.dtype
+            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
+        )
+
+        ctx.row_id_map = row_id_map
+        ctx.num_tokens = indices.size(0)
+        ctx.num_topK = num_topK
+        return permuted_act, row_id_map
+
+    @staticmethod
+    def backward(ctx, permuted_act_grad, _):
+        # Empty input check
+        if not permuted_act_grad.numel():
+            return permuted_act_grad, None, None, None
+
+        if not permuted_act_grad.is_contiguous():
+            permuted_act_grad = permuted_act_grad.contiguous()
+
+        row_id_map = ctx.row_id_map
+        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)
+        return unpermuted_act_grad, None, None, None
+
+
+################################################################################################
+##
+## UnpermuteMoE topK
+##
+################################################################################################
+
+
+class UnpermuteMoE_topK(torch.autograd.Function):
+
+    @staticmethod
+    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
+            return input_act
+
+        # Device check
+        if input_act.is_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!")
+            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!")
+            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)}."
+            )
+
+        # 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."
+            )
+            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.")
+            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!")
+            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!")
+            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!")
+            probs = probs.contiguous()
+
+        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)
+
+        ctx.save_for_backward(input_act, row_id_map, probs)
+        return unpermuted_output
+
+    @staticmethod
+    def backward(ctx, unpermuted_act_grad):
+        # Empty input check
+        if not unpermuted_act_grad.numel():
+            return unpermuted_act_grad, None, ctx.probs
+
+        if not unpermuted_act_grad.is_contiguous():
+            unpermuted_act_grad = unpermuted_act_grad.contiguous()
+
+        input_act, row_id_map, probs = ctx.saved_tensors
+
+        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)
+
+        if not ctx.needs_input_grad[2]:
+            prob_grad = None
+        return act_grad, None, prob_grad
+
+
+def permute(input_act, indices, num_out_tokens=None, max_token_num=0):
+    num_out_tokens = 0 if num_out_tokens is None else num_out_tokens
+    return PermuteMoE_topK.apply(input_act, indices, num_out_tokens, max_token_num)
+
+
+def unpermute(input_act, row_id_map, probs=None):
+    return UnpermuteMoE_topK.apply(input_act, row_id_map, probs)
+
+
+################################################################################################
+##
+## mutlimodal RoPE
+##
+################################################################################################
+
+
+class MultimodalRoPE(torch.autograd.Function):
+
+    @staticmethod
+    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!")
+        if k.is_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!")
+        if sin.is_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!")
+
+        # Contiguous check
+        if not q.is_contiguous():
+            warnings.warn("The input `q` of multimodal_rope op is discontiguous!")
+            q = q.contiguous()
+        if not k.is_contiguous():
+            warnings.warn("The input `k` of multimodal_rope op is discontiguous!")
+            k = k.contiguous()
+        if not cos.is_contiguous():
+            warnings.warn("The input `cos` of multimodal_rope op is discontiguous!")
+            cos = cos.contiguous()
+        if not sin.is_contiguous():
+            warnings.warn("The input `sin` of multimodal_rope op is discontiguous!")
+            sin = sin.contiguous()
+
+        # Prepare mrope_section_doubled
+        mrope_section_doubled = [x * 2 for x in mrope_section]
+
+        # Create output tensors
+        q_out = torch.empty_like(q)
+        k_out = torch.empty_like(k)
+        backend.rope(q, k, cos, sin, q_out, k_out, mrope_section_doubled)
+
+        ctx.save_for_backward(q, k, cos, sin)
+        ctx.mrope_section_doubled = mrope_section_doubled
+
+        return q_out, k_out
+
+    @staticmethod
+    def backward(ctx, grad_q_out, grad_k_out):
+
+        if not grad_q_out.is_contiguous():
+            grad_q_out = grad_q_out.contiguous()
+        if not grad_k_out.is_contiguous():
+            grad_k_out = grad_k_out.contiguous()
+
+        q, k, cos, sin = ctx.saved_tensors
+
+        grad_q = None
+        grad_k = None
+        if ctx.needs_input_grad[0]:
+            grad_q = torch.empty_like(q)
+        if ctx.needs_input_grad[1]:
+            grad_k = torch.empty_like(k)
+
+        if grad_q is not None or grad_k is not None:
+            backend.rope_bwd(
+                grad_q_out,
+                grad_k_out,
+                q,
+                k,
+                cos,
+                sin,
+                grad_q if grad_q is not None else torch.empty_like(q),
+                grad_k if grad_k is not None else torch.empty_like(k),
+                ctx.mrope_section_doubled,
+            )
+
+        return grad_q, grad_k, None, None, None
+
+
+def multimodal_rope(q, k, cos, sin, mrope_section):
+    return MultimodalRoPE.apply(q, k, cos, sin, mrope_section)

wall_x/model/action_head.py

@@ -0,0 +1,400 @@
+
+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.
+    
+    This module handles normalization and denormalization of action data for different robot
+    configurations. It maintains per-robot statistics (min values and deltas) and applies
+    normalization to map actions to the [-1, 1] range.
+    """
+    
+    def __init__(self, action_statistic_dof, dof_config):
+        """
+        Initialize the normalizer with robot-specific action statistics.
+        
+        Args:
+            action_statistic_dof (dict): Statistical data for each robot's degrees of freedom
+            dof_config (dict): Configuration mapping for degrees of freedom per robot
+        """
+        super(Normalizer, self).__init__()
+
+        action_statistic = {}
+        
+        # Process statistics for each robot
+        for robot_name in action_statistic_dof.keys():
+            action_statistic[robot_name] = {}
+            all_dof_min = []
+            all_dof_delta = []
+            
+            # Collect min and delta values for all DOFs
+            for k in dof_config:
+                if k in action_statistic_dof[robot_name]:
+                    all_dof_min.extend(action_statistic_dof[robot_name][k]["min"])
+                    all_dof_delta.extend(action_statistic_dof[robot_name][k]["delta"])
+                else:
+                    # Use default values if statistics not available
+                    all_dof_min.extend([0.0] * dof_config[k])
+                    all_dof_delta.extend([1.0] * dof_config[k])
+                    
+            all_dof_min = torch.tensor(all_dof_min)
+            all_dof_delta = torch.tensor(all_dof_delta)
+            action_statistic[robot_name]["min"] = all_dof_min
+            action_statistic[robot_name]["delta"] = all_dof_delta
+
+        # Register statistics as non-trainable parameters
+        self.min = nn.ParameterDict({
+            k: nn.Parameter(action_statistic[k]["min"], requires_grad=False) 
+            for k in action_statistic.keys()
+        })
+        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):
+        """
+        Normalize action data to [-1, 1] range using robot-specific statistics.
+        
+        Args:
+            xs: Input action data tensors
+            dataset_names: List of dataset/robot names corresponding to each tensor
+            
+        Returns:
+            torch.Tensor: Normalized action data in [-1, 1] range
+        """
+        new_xs = []
+        # Filter out multimodal dataset entries
+        dataset_names = [name for name in dataset_names if name != "x2_multimodal"]
+        
+        for x, dataset_name in zip(xs, dataset_names):
+            # Apply min-max normalization
+            x = (x - self.min[dataset_name]) / (self.delta[dataset_name])
+            # Scale to [-1, 1] range
+            x = x * 2 - 1
+            # Clamp to ensure bounds
+            x = torch.clamp(x, -1, 1)
+            new_xs.append(x)
+            
+        new_xs = torch.stack(new_xs)
+        return new_xs
+
+    def unnormalize_data(self, xs, dataset_names, dof_mask=None):
+        """
+        Convert normalized data back to original action space.
+        
+        Args:
+            xs: Normalized action data in [-1, 1] range
+            dataset_names: List of dataset/robot names
+            dof_mask: Optional mask to select specific degrees of freedom
+            
+        Returns:
+            torch.Tensor: Denormalized action data in original scale
+        """
+        new_xs = []
+        # Filter out multimodal dataset entries
+        dataset_names = [name for name in dataset_names if name != "x2_multimodal"]
+        dof_mask = dof_mask if dof_mask is not None else [None] * len(xs)
+        
+        for x, dataset_name, mask in zip(xs, dataset_names, dof_mask):
+            # Convert from [-1, 1] to [0, 1] range
+            x = (x + 1) / 2
+            
+            # Apply DOF mask if provided
+            if mask is not None:
+                mask = mask[0].bool()
+                action_space_delta = self.delta[dataset_name][mask]
+                action_space_min = self.min[dataset_name][mask]
+            else:
+                action_space_delta = self.delta[dataset_name]
+                action_space_min = self.min[dataset_name]
+                
+            # Scale back to original range
+            x = x * action_space_delta + action_space_min
+            new_xs.append(x)
+            
+        new_xs = torch.stack(new_xs)
+        return new_xs
+
+
+class SinusoidalPosEmb(nn.Module):
+    """
+    Sinusoidal positional embedding for diffusion timesteps.
+    
+    Generates sinusoidal embeddings commonly used in diffusion models to encode
+    timestep information with different frequencies.
+    """
+    
+    def __init__(self, dim):
+        """
+        Initialize sinusoidal positional embedding.
+        
+        Args:
+            dim (int): Embedding dimension (must be even)
+        """
+        super().__init__()
+        self.dim = dim
+
+    def forward(self, x):
+        """
+        Generate sinusoidal embeddings for input timesteps.
+        
+        Args:
+            x (torch.Tensor): Input timesteps
+            
+        Returns:
+            torch.Tensor: Sinusoidal embeddings of shape (..., dim)
+        """
+        device = x.device
+        half_dim = self.dim // 2
+        emb = math.log(10000) / (half_dim - 1)
+        emb = torch.exp(torch.arange(half_dim, device=device) * -emb)
+        emb = x[:, None] * emb[None, :]
+        emb = torch.cat((emb.sin(), emb.cos()), dim=-1)
+        return emb
+
+
+
+class ActionProcessor(nn.Module):
+    """
+    Action sequence processor for robotic control with flow matching.
+    
+    This module handles action sequence processing for robotic systems with the following capabilities:
+    1. Adds controlled noise to action sequences using Beta distribution scheduling
+    2. Generates temporal embeddings for timestep conditioning
+    3. Projects actions to model hidden space for transformer processing
+    4. Supports proprioceptive data integration and multi-robot configurations
+    
+    The Beta distribution provides more flexible noise injection strategies compared to
+    traditional linear schedules, allowing better control over the noise scheduling process.
+    """
+    
+    def __init__(self, config):
+        """
+        Initialize the action processor with multi-robot support.
+        
+        Args:
+            config: Configuration object containing:
+                - dof_config (dict): Degrees of freedom configuration per robot type
+                - agent_pos_config (dict): Agent position/proprioception configuration
+                - hidden_size (int): Model hidden layer dimension
+                - noise_scheduler (dict): Noise scheduler configuration with Beta parameters
+        """
+        super().__init__()
+        
+        # Calculate action and proprioception dimensions from configuration
+        self.dof_config = config.dof_config
+        self.agent_pos_config = config.agent_pos_config
+        self.action_dim = sum([v for k, v in self.dof_config.items()])
+        self.propri_dim = sum([v for k, v in self.agent_pos_config.items()])
+
+        # Log configuration details for debugging
+        print("ActionProcessor Configuration:", flush=True)
+        print(f"  Action dimension: {self.action_dim}", flush=True)
+        print(f"  Proprioception dimension: {self.propri_dim}", flush=True)
+        print("  DOF configuration:", flush=True)
+        for key, value in self.dof_config.items():
+            print(f"    {key}: {value}", flush=True)
+        print("  Agent position configuration:", flush=True)
+        for key, value in self.agent_pos_config.items():
+            print(f"    {key}: {value}", flush=True)
+        
+        self.hidden_size = config.hidden_size
+
+        # 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)
+
+        # 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
+        
+        # Initialize Beta distribution for noise scheduling
+        alpha_tensor = torch.tensor(self.beta_alpha, dtype=torch.float32).to("cuda")
+        beta_tensor = torch.tensor(self.beta_beta, dtype=torch.float32).to("cuda")
+        self.beta_dist = Beta(alpha_tensor, beta_tensor)
+
+        # Sinusoidal positional embedding for timesteps
+        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.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')
+
+    def sample_time(self, batch_size, device, dtype):
+        """
+        Sample timesteps using Beta distribution for noise scheduling.
+        
+        Generates random timesteps in [0,1] range using Beta distribution, then scales them.
+        This provides more flexible control over the noise injection schedule compared to
+        uniform sampling.
+        
+        Args:
+            batch_size (int): Number of timesteps to sample
+            device: Target device for tensors
+            dtype: Target data type for tensors
+            
+        Returns:
+            torch.Tensor: Sampled timesteps of shape [batch_size]
+        """
+        sample = self.beta_dist.sample([batch_size]).to(device=device, dtype=dtype)
+        time = (self.s - sample) / self.s
+        return time
+    
+    def proprioception_proj(self, proprioception, dataset_names=None, dof_mask=None, use_history=False):
+        """
+        Project proprioceptive data (joint positions, orientations) to hidden space.
+        
+        Args:
+            proprioception (torch.Tensor): Proprioceptive data of shape [batch_size, seq_len, propri_dim]
+            dataset_names (list, optional): Dataset names for normalization. Defaults to None.
+            dof_mask (torch.Tensor, optional): DOF mask of shape [batch_size, propri_dim]. Defaults to None.
+            use_history (bool, optional): Whether to use historical proprioceptive data. Defaults to False.
+            
+        Returns:
+            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)
+        
+        if dof_mask is not None:
+            # Concatenate proprioception with DOF mask
+            # TODO: Use variable-based dimension checking for better flexibility
+            if use_history:
+                proprioception = torch.cat([proprioception, dof_mask], dim=-1)
+            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)
+        return self.propri_proj(proprioception)
+
+    def forward(self, action_chunk, dataset_names, dof_mask=None):
+        """
+        Process action sequences with noise injection and temporal embedding.
+        
+        This method implements the forward pass for flow matching training:
+        1. Adds Beta-distributed noise to action sequences
+        2. Generates sinusoidal timestep embeddings
+        3. Projects noisy actions to hidden space
+        4. Combines action and temporal features
+        
+        Args:
+            action_chunk (torch.Tensor): Action sequences of shape [batch_size, seq_len, action_dim]
+            dataset_names (list): Dataset names for normalization
+            dof_mask (torch.Tensor, optional): DOF mask of shape [batch_size, seq_len, action_dim]. 
+                                              Defaults to None.
+            
+        Returns:
+            tuple: (action_embeddings, flow_target) where:
+                - action_embeddings: Processed action features of shape [batch_size, seq_len, hidden_size]
+                - flow_target: Flow matching target (action_chunk - noise) for loss computation
+        """
+        batch_size = action_chunk.shape[0]
+        device = action_chunk.device
+        dtype = action_chunk.dtype
+
+        # 1. Add noise to action sequences using flow matching
+        noise = torch.randn_like(action_chunk)
+        time = self.sample_time(batch_size, device, dtype)
+        t = time.unsqueeze(-1).unsqueeze(-1)  # Broadcast to match action dimensions
+
+        # Linear interpolation between noise and action (flow matching)
+        noisy_action = (1 - t) * noise + t * action_chunk
+        flow = action_chunk - noise  # Flow target for loss computation
+
+        # 2. Generate sinusoidal positional encoding for timesteps
+        time_embed = self.time_embed(time)
+
+        # 3. Project noisy actions with DOF mask to hidden space
+        if dof_mask is not None:
+            noisy_action = torch.cat([noisy_action, dof_mask], dim=-1)
+            
+        noisy_action = noisy_action.to(dtype=self.w1.weight.dtype)
+        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)
+        
+        # Combine action and temporal embeddings
+        concat_embed = torch.cat([action_embed, time_embed], dim=-1)
+        concat_embed = self.w2(concat_embed)
+        embed = self.w3(self.act_fn(concat_embed))
+
+        return embed, flow
+    
+    def step(self, timestep, noisy_action, dof_mask=None):
+        """
+        Single denoising step for diffusion inference.
+        
+        Processes noisy actions at a specific timestep for iterative denoising during inference.
+        
+        Args:
+            timestep (torch.Tensor): Current timesteps of shape [batch_size]
+            noisy_action (torch.Tensor): Noisy actions of shape [batch_size, seq_len, action_dim]
+            dof_mask (torch.Tensor, optional): DOF mask for action space. Defaults to None.
+            
+        Returns:
+            torch.Tensor: Processed action embeddings of shape [batch_size, seq_len, hidden_size]
+        """
+        # Concatenate noisy action with DOF mask if provided
+        if dof_mask is not None:
+            noisy_action = torch.cat([noisy_action, dof_mask], dim=-1)
+            
+        # Generate timestep embeddings
+        time_embed = self.time_embed(timestep)  # [batch_size, hidden_size]
+        
+        # Project noisy actions
+        action_embed = self.w1(noisy_action)
+        
+        # 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)
+        
+        # Combine embeddings and process through MLP
+        concat_embed = torch.cat([action_embed, time_embed], dim=-1)
+        concat_embed = self.w2(concat_embed)
+        embed = self.w3(self.act_fn(concat_embed))
+
+        return embed
+
+    def flow_loss(self, action_hidden_states, flow, dof_mask=None):
+        """
+        Compute flow matching loss between predicted and target actions.
+        
+        Args:
+            action_hidden_states (torch.Tensor): Hidden states from transformer
+            flow (torch.Tensor): Target flow (action - noise) for matching
+            dof_mask (torch.Tensor, optional): DOF mask to weight loss per dimension. Defaults to None.
+            
+        Returns:
+            torch.Tensor: Flow matching loss (no reduction for channel loss computation)
+        """
+        # Project hidden states back to action space
+        action_pred = self.action_proj_back(action_hidden_states)
+        
+        # Compute MSE loss between predicted and target flow
+        loss = self.mse_loss(action_pred, flow)
+        
+        # Apply DOF mask if provided
+        if dof_mask is not None:
+            dof_mask = dof_mask.reshape(-1, dof_mask.shape[-1])
+            loss = loss * dof_mask
+            
+        # Return loss without reduction for channel-wise loss computation
+        return loss

wall_x/model/qwen2_5_based/__init__.py

@@ -0,0 +1,2 @@
+from .modeling_qwen2_5_vl_act import Qwen2_5_VLMoEModel,Qwen2_5_VLMoEForAction
+from .configuration_qwen2_5_vl import Qwen2_5_VLConfig

wall_x/model/qwen2_5_based/configuration_qwen2_5_vl.py

@@ -0,0 +1,248 @@
+from transformers.configuration_utils import PretrainedConfig
+from transformers.modeling_rope_utils import rope_config_validation
+
+
+class Qwen2_5_VLVisionConfig(PretrainedConfig):
+    model_type = "qwen2_5_vl"
+    base_config_key = "vision_config"
+
+    def __init__(
+        self,
+        depth=32,
+        hidden_size=3584,
+        hidden_act="silu",
+        intermediate_size=3420,
+        num_heads=16,
+        in_channels=3,
+        patch_size=14,
+        spatial_merge_size=2,
+        temporal_patch_size=2,
+        tokens_per_second=4,
+        window_size=112,
+        out_hidden_size=3584,
+        fullatt_block_indexes=[7, 15, 23, 31],
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+
+        self.depth = depth
+        self.hidden_size = hidden_size
+        self.hidden_act = hidden_act
+        self.intermediate_size = intermediate_size
+        self.num_heads = num_heads
+        self.in_channels = in_channels
+        self.patch_size = patch_size
+        self.spatial_merge_size = spatial_merge_size
+        self.temporal_patch_size = temporal_patch_size
+        self.tokens_per_second = tokens_per_second
+        self.window_size = window_size
+        self.fullatt_block_indexes = fullatt_block_indexes
+        self.out_hidden_size = out_hidden_size
+
+
+class Qwen2_5_VLConfig(PretrainedConfig):
+    r"""
+    This is the configuration class to store the configuration of a [`Qwen2_5_VLModel`]. It is used to instantiate a
+    Qwen2-VL model according to the specified arguments, defining the model architecture. Instantiating a configuration
+    with the defaults will yield a similar configuration to that of
+    Qwen2-VL-7B-Instruct [Qwen/Qwen2-VL-7B-Instruct](https://huggingface.co/Qwen/Qwen2-VL-7B-Instruct).
+
+    Configuration objects inherit from [`PretrainedConfig`] and can be used to control the model outputs. Read the
+    documentation from [`PretrainedConfig`] for more information.
+
+
+    Args:
+        vocab_size (`int`, *optional*, defaults to 152064):
+            Vocabulary size of the Qwen2_5_VL model. Defines the number of different tokens that can be represented by the
+            `inputs_ids` passed when calling [`Qwen2_5_VLModel`]
+        hidden_size (`int`, *optional*, defaults to 8192):
+            Dimension of the hidden representations.
+        intermediate_size (`int`, *optional*, defaults to 29568):
+            Dimension of the MLP representations.
+        num_hidden_layers (`int`, *optional*, defaults to 80):
+            Number of hidden layers in the Transformer encoder.
+        num_attention_heads (`int`, *optional*, defaults to 64):
+            Number of attention heads for each attention layer in the Transformer encoder.
+        num_key_value_heads (`int`, *optional*, defaults to 8):
+            This is the number of key_value heads that should be used to implement Grouped Query Attention. If
+            `num_key_value_heads=num_attention_heads`, the model will use Multi Head Attention (MHA), if
+            `num_key_value_heads=1` the model will use Multi Query Attention (MQA) otherwise GQA is used. When
+            converting a multi-head checkpoint to a GQA checkpoint, each group key and value head should be constructed
+            by meanpooling all the original heads within that group. For more details checkout [this
+            paper](https://arxiv.org/pdf/2305.13245.pdf). If it is not specified, will default to `32`.
+        hidden_act (`str` or `function`, *optional*, defaults to `"silu"`):
+            The non-linear activation function (function or string) in the decoder.
+        max_position_embeddings (`int`, *optional*, defaults to 32768):
+            The maximum sequence length that this model might ever be used with.
+        initializer_range (`float`, *optional*, defaults to 0.02):
+            The standard deviation of the truncated_normal_initializer for initializing all weight matrices.
+        rms_norm_eps (`float`, *optional*, defaults to 1e-05):
+            The epsilon used by the rms normalization layers.
+        use_cache (`bool`, *optional*, defaults to `True`):
+            Whether or not the model should return the last key/values attentions (not used by all models). Only
+            relevant if `config.is_decoder=True`.
+        tie_word_embeddings (`bool`, *optional*, defaults to `False`):
+            Whether the model's input and output word embeddings should be tied.
+        rope_theta (`float`, *optional*, defaults to 1000000.0):
+            The base period of the RoPE embeddings.
+        use_sliding_window (`bool`, *optional*, defaults to `False`):
+            Whether to use sliding window attention.
+        sliding_window (`int`, *optional*, defaults to 4096):
+            Sliding window attention (SWA) window size. If not specified, will default to `4096`.
+        max_window_layers (`int`, *optional*, defaults to 80):
+            The number of layers that use SWA (Sliding Window Attention). The bottom layers use SWA while the top use full attention.
+        attention_dropout (`float`, *optional*, defaults to 0.0):
+            The dropout ratio for the attention probabilities.
+        vision_config (`Dict`, *optional*):
+            The config for the visual encoder initialization.
+        rope_scaling (`Dict`, *optional*):
+            Dictionary containing the scaling configuration for the RoPE embeddings. NOTE: if you apply new rope type
+            and you expect the model to work on longer `max_position_embeddings`, we recommend you to update this value
+            accordingly.
+            Expected contents:
+                `rope_type` (`str`):
+                    The sub-variant of RoPE to use. Can be one of ['default', 'linear', 'dynamic', 'yarn', 'longrope',
+                    'llama3'], with 'default' being the original RoPE implementation.
+                `factor` (`float`, *optional*):
+                    Used with all rope types except 'default'. The scaling factor to apply to the RoPE embeddings. In
+                    most scaling types, a `factor` of x will enable the model to handle sequences of length x *
+                    original maximum pre-trained length.
+                `original_max_position_embeddings` (`int`, *optional*):
+                    Used with 'dynamic', 'longrope' and 'llama3'. The original max position embeddings used during
+                    pretraining.
+                `attention_factor` (`float`, *optional*):
+                    Used with 'yarn' and 'longrope'. The scaling factor to be applied on the attention
+                    computation. If unspecified, it defaults to value recommended by the implementation, using the
+                    `factor` field to infer the suggested value.
+                `beta_fast` (`float`, *optional*):
+                    Only used with 'yarn'. Parameter to set the boundary for extrapolation (only) in the linear
+                    ramp function. If unspecified, it defaults to 32.
+                `beta_slow` (`float`, *optional*):
+                    Only used with 'yarn'. Parameter to set the boundary for interpolation (only) in the linear
+                    ramp function. If unspecified, it defaults to 1.
+                `short_factor` (`List[float]`, *optional*):
+                    Only used with 'longrope'. The scaling factor to be applied to short contexts (<
+                    `original_max_position_embeddings`). Must be a list of numbers with the same length as the hidden
+                    size divided by the number of attention heads divided by 2
+                `long_factor` (`List[float]`, *optional*):
+                    Only used with 'longrope'. The scaling factor to be applied to long contexts (<
+                    `original_max_position_embeddings`). Must be a list of numbers with the same length as the hidden
+                    size divided by the number of attention heads divided by 2
+                `low_freq_factor` (`float`, *optional*):
+                    Only used with 'llama3'. Scaling factor applied to low frequency components of the RoPE
+                `high_freq_factor` (`float`, *optional*):
+                    Only used with 'llama3'. Scaling factor applied to high frequency components of the RoPE
+
+    ```python
+    >>> from transformers import Qwen2_5_VLForConditionalGeneration, Qwen2_5_VLConfig
+
+    >>> # Initializing a Qwen2_5_VL style configuration
+    >>> configuration = Qwen2_5_VLConfig()
+
+    >>> # Initializing a model from the Qwen2-VL-7B style configuration
+    >>> model = Qwen2_5_VLForConditionalGeneration(configuration)
+
+    >>> # Accessing the model configuration
+    >>> configuration = model.config
+    ```"""
+
+    model_type = "qwen2_5_vl"
+    sub_configs = {"vision_config": Qwen2_5_VLVisionConfig}
+    keys_to_ignore_at_inference = ["past_key_values"]
+    # Default tensor parallel plan for base model `Qwen2_5_VL`
+    base_model_tp_plan = {
+        "layers.*.self_attn.q_proj": "colwise",
+        "layers.*.self_attn.k_proj": "colwise",
+        "layers.*.self_attn.v_proj": "colwise",
+        "layers.*.self_attn.o_proj": "rowwise",
+        "layers.*.mlp.gate_proj": "colwise",
+        "layers.*.mlp.up_proj": "colwise",
+        "layers.*.mlp.down_proj": "rowwise",
+    }
+    base_model_pp_plan = {
+        "embed_tokens": (["input_ids"], ["inputs_embeds"]),
+        "layers": (["hidden_states", "attention_mask"], ["hidden_states"]),
+        "norm": (["hidden_states"], ["hidden_states"]),
+    }
+
+    def __init__(
+        self,
+        vocab_size=152064,
+        hidden_size=8192,
+        intermediate_size=29568,
+        num_hidden_layers=80,
+        num_attention_heads=64,
+        num_key_value_heads=8,
+        hidden_act="silu",
+        max_position_embeddings=32768,
+        initializer_range=0.02,
+        rms_norm_eps=1e-05,
+        use_cache=True,
+        tie_word_embeddings=False,
+        rope_theta=1000000.0,
+        use_sliding_window=False,
+        sliding_window=4096,
+        max_window_layers=80,
+        attention_dropout=0.0,
+        vision_config=None,
+        rope_scaling=None,
+        num_experts=4,
+        experts=None,
+        dof_config=None,
+        noise_scheduler=None,
+        dim_inputs=(1536,1536),
+        attention_moe=False,
+        mlp_moe=False,
+        **kwargs,
+    ):
+        if isinstance(vision_config, dict):
+            self.vision_config = self.sub_configs["vision_config"](**vision_config)
+        elif vision_config is None:
+            self.vision_config = self.sub_configs["vision_config"]()
+
+        self.vocab_size = vocab_size
+        self.max_position_embeddings = max_position_embeddings
+        self.hidden_size = hidden_size
+        self.intermediate_size = intermediate_size
+        self.num_hidden_layers = num_hidden_layers
+        self.num_attention_heads = num_attention_heads
+        self.use_sliding_window = use_sliding_window
+        self.sliding_window = sliding_window
+        self.max_window_layers = max_window_layers
+
+        # for backward compatibility
+        if num_key_value_heads is None:
+            num_key_value_heads = num_attention_heads
+
+        self.num_key_value_heads = num_key_value_heads
+        self.hidden_act = hidden_act
+        self.initializer_range = initializer_range
+        self.rms_norm_eps = rms_norm_eps
+        self.use_cache = use_cache
+        self.rope_theta = rope_theta
+        self.attention_dropout = attention_dropout
+        self.rope_scaling = rope_scaling
+
+        self.num_experts = num_experts
+        self.experts = experts
+        self.dof_config = dof_config
+        self.noise_scheduler = noise_scheduler
+        self.dim_inputs = tuple(dim_inputs)
+        self.attention_moe = attention_moe
+        self.mlp_moe = mlp_moe
+
+        # Validate the correctness of rotary position embeddings parameters
+        # BC: if there is a 'type' field, move it to 'rope_type'.
+        # and change type from 'mrope' to 'default' because `mrope` does defeault RoPE calculations
+        # one can set it to "linear"/"dynamic" etc. to have scaled RoPE
+        # TODO: @raushan update config in the hub
+        if self.rope_scaling is not None and "type" in self.rope_scaling:
+            if self.rope_scaling["type"] == "mrope":
+                self.rope_scaling["type"] = "default"
+            self.rope_scaling["rope_type"] = self.rope_scaling["type"]
+        rope_config_validation(self, ignore_keys={"mrope_section"})
+
+        super().__init__(tie_word_embeddings=tie_word_embeddings, **kwargs)
+
+
+__all__ = ["Qwen2_5_VLConfig"]

wall_x/trainer/qwen_vl_act_trainer.py

@@ -0,0 +1,671 @@
+import os
+import gc
+import time
+import torch
+import random
+import numpy as np
+import torch.nn as nn
+
+from tqdm import tqdm
+from functools import wraps
+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
+
+
+def timer(func):
+    """
+    Decorator to measure function execution time.
+    
+    Args:
+        func: Function to be timed
+        
+    Returns:
+        Wrapped function with timing functionality
+    """
+    @wraps(func)
+    def wrapper(*args, **kwargs):
+        start_time = time.time()
+        result = func(*args, **kwargs)
+        end_time = time.time()
+        print(
+            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
+
+
+def print_rank_last(message):
+    """
+    Print message only on the last rank in distributed training.
+    
+    Args:
+        message (str): Message to print
+    """
+    if torch.distributed.is_initialized():
+        if torch.distributed.get_rank() == (torch.distributed.get_world_size() - 1):
+            print(message, flush=True)
+    else:
+        print(message, flush=True)
+
+
+def seed_all(seed):
+    """
+    Set random seeds for reproducible training.
+    
+    Args:
+        seed (int): Random seed value
+    """
+    np.random.seed(seed)
+    random.seed(seed)
+    torch.manual_seed(seed)
+
+
+class QwenVlAct_Trainer:
+    """
+    Vision-Language-Action trainer for Qwen-VL models with robotic action prediction.
+    
+    This trainer handles multi-modal learning combining vision, language, and action data
+    for robotic control applications. It supports distributed training, mixed precision,
+    gradient accumulation, and various optimization strategies including MoE (Mixture of Experts).
+    
+    Features:
+    - Multi-modal data processing (vision + language + actions)
+    - Distributed training with Accelerate
+    - Gradient accumulation and clipping
+    - Learning rate scheduling with warmup
+    - Checkpoint saving and resuming
+    - Comprehensive logging and monitoring
+    """
+
+    @timer
+    def __init__(self, config, logger, accelerator: Accelerator = None, seed=42, data_config_path=None):
+        """
+        Initialize the Vision-Language-Action trainer.
+        
+        Args:
+            config (dict): Training configuration dictionary containing:
+                - processor_path (str): Path to data preprocessing processor
+                - qwen_vl_act_config_path (str): Path to model configuration file
+                - learning_rate (float): Base learning rate for training
+                - num_epoch (int): Number of training epochs
+                - pretrained_qwen_vl_path (str): Path to pretrained model
+                - And other training hyperparameters
+            logger: Logger instance for tracking metrics
+            accelerator (Accelerator, optional): Hugging Face Accelerate instance for distributed training
+            seed (int, optional): Random seed for reproducibility. Defaults to 42.
+            data_config_path (str, optional): Path to data configuration file
+            
+        Raises:
+            ValueError: If required configuration keys are missing
+        """
+        # Validate required configuration keys
+        required_keys = ["processor_path", "qwen_vl_act_config_path", "learning_rate", "num_epoch"]
+        for key in required_keys:
+            if key not in config:
+                raise ValueError(f"Missing required configuration key: {key}")
+
+        self.config = config
+        self.logger = logger
+        self.accelerator = accelerator
+        self.seed = seed
+        
+        # Initialize random seeds for reproducibility
+        seed_all(self.seed)
+        
+        # Training state variables
+        self.start_epoch = 0
+        self.global_step = 0
+        self.num_epoch = self.config["num_epoch"]
+        self.initial_step = 0
+        
+        # Data and model configuration
+        self.dataload_config = get_data_configs(self.config["data"])
+        self.data_config_path = data_config_path
+        self.use_fast_tokenizer = self.config.get("use_fast_tokenizer", False)
+        
+        # Load model and initialize training components
+        self.load_model()
+        self.action_dim = sum(self.config["dof_config"].values())
+        
+        # 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)
+        
+        # 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)
+
+        # 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|>")
+
+        # Initialize evaluation metrics
+        self.base_l1_loss = None
+        self.base_l1_loss_detail = {}
+
+        # Performance monitoring
+        self.timers = Timers(log_level=0, log_option="minmax")
+
+        # 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)
+
+    def print_rank0(self, msg, flush=True):
+        """
+        Print message only on rank 0 (main process).
+        
+        Args:
+            msg: Message to print
+            flush (bool): Whether to flush output buffer
+        """
+        if self.accelerator.is_main_process:
+            print(msg, flush=flush)
+
+    def fit(self):
+        """
+        Main training loop executing multiple epochs with validation.
+        
+        Handles the complete training process including:
+        - Training loop execution
+        - Validation after each epoch
+        - Process synchronization
+        - Memory cleanup
+        """
+        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):
+            self.val_loop()
+            self.accelerator.wait_for_everyone()
+
+        # Main training loop
+        for epoch in range(self.start_epoch, self.num_epoch):
+            self.train_loop(epoch)
+            self.accelerator.wait_for_everyone()
+            
+            # Validation after each epoch
+            self.val_loop()
+            self.accelerator.wait_for_everyone()
+            
+            # Memory cleanup
+            gc.collect()
+
+    def train_loop(self, epoch):
+        """
+        Execute training for a single epoch.
+        
+        Args:
+            epoch (int): Current epoch number
+            
+        Handles:
+        - Data loading and batching
+        - Forward/backward passes
+        - Gradient accumulation and clipping
+        - Learning rate scheduling
+        - Loss logging and monitoring
+        - Performance profiling (optional)
+        """
+        # Initialize training dataloader for current epoch
+        if isinstance(self.dataset, PreprocessedDataset):
+            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_sampler.set_epoch(epoch)
+        else:
+            self.train_dataloader = self.dataset.get_train_dataloader()
+
+        self.model.train()
+        grad_accum_steps = self.config.get("gradient_accumulation_steps", 1)
+        total = len(self.train_dataloader)
+        t0 = time.time()
+        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"),
+                record_shapes=True,
+                profile_memory=True,
+                with_stack=True,
+            )
+            profiler.__enter__()
+
+        try:
+            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()}
+                
+                self.timers("data-load").stop()
+                
+                with self.accelerator.accumulate(self.model):
+                    # Forward pass
+                    self.timers("forward-compute", log_level=0).start(barrier=False)
+                    outputs = self.model(**batch, mode="train")
+                    self.timers("forward-compute").stop()
+                    
+                    loss = outputs.loss
+                    
+                    # Check for NaN loss
+                    if torch.isnan(loss):
+                        print(f"Warning: NaN loss detected in epoch: {epoch}, step: {i}", flush=True)
+                        continue
+                    
+                    # Backward pass
+                    self.timers("backward-compute", log_level=0).start(barrier=False)
+                    self.accelerator.backward(loss)
+                    self.timers("backward-compute").stop()
+                    
+                    # Gradient clipping
+                    total_norm = self.accelerator.clip_grad_norm_(
+                        self.model.parameters(), self.config.get("max_grad_norm", 1.0)
+                    )
+
+                    # Optimizer step
+                    self.timers("optimizer", log_level=0).start(barrier=False)
+                    self.optimizer.step()
+                    self.optimizer.zero_grad()
+                    self.timers("optimizer").stop()
+
+                    # Update global step and learning rate after gradient accumulation
+                    if (i + 1) % grad_accum_steps == 0:
+                        self.lr_scheduler.step()
+                        self.global_step += 1
+                        lr = self.lr_scheduler.get_last_lr()[0]
+                        
+                        # Gather loss across all processes for logging
+                        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 "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 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 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
+                            
+                            # Log action accuracy for 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()
+                                )
+
+                        # Log metrics
+                        if self.logger is not None:
+                            self.logger.log(_log_dict, step=self.global_step)
+
+                        # 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.timers("interval-time").stop()
+                
+                # Setup timers for next iteration
+                if i < len(self.train_dataloader) - 1:
+                    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)
+                    t0 = time.time()
+                
+                if enable_profiling:
+                    profiler.step()
+
+        finally:
+            if enable_profiling:
+                profiler.__exit__(None, None, None)
+
+    @torch.no_grad()
+    def val_loop(self):
+        """
+        Execute validation loop with gradient computation disabled.
+        
+        Evaluates model performance on validation set and logs validation loss.
+        """
+        # Initialize validation dataloader
+        if getattr(self, "val_dataloader", None) is not None:
+            self.val_sampler.set_epoch(0)
+        else:
+            self.val_dataloader, self.val_sampler = self.dataset.get_val_dataloader()
+            self.val_sampler.set_epoch(0)
+
+        self.model.eval()
+        self.val_loss = 0
+        
+        # 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)
+        ):
+            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()}
+            
+            with torch.no_grad():
+                outputs = self.model(**batch, mode="train")
+                loss = outputs.loss
+                self.val_loss += self.accelerator.gather(loss.detach()).mean().item()
+        
+        # Calculate average validation loss
+        self.val_loss /= len(self.val_dataloader)
+        
+        # Log validation metrics
+        if self.logger is not None:
+            self.logger.log({"val_loss": self.val_loss}, step=self.global_step)
+        
+        self.model.train()
+
+    @timer
+    def load_model(self):
+        """
+        Load and configure the Vision-Language-Action model.
+        
+        Handles:
+        - Model loading from pretrained weights
+        - Processor initialization
+        - Optimizer configuration (with support for different learning rates for different components)
+        - Learning rate scheduler setup
+        - Model preparation for distributed training
+        """
+        # Load pretrained model
+        model = Qwen2_5_VLMoEForAction.from_pretrained(
+            self.config["pretrained_qwen_vl_path"], 
+            **{"use_fast_tokenizer": self.use_fast_tokenizer}
+        )
+        self.processor = model.processor
+        model = model.to(torch.bfloat16)
+
+        # Configure optimizer based on training strategy
+        if "freeze_vlm" in self.config and self.config["freeze_vlm"]:
+            print("Freezing VLM parameters, training only MoE experts", flush=True)
+            moe_params = []
+            for name, param in model.named_parameters():
+                if "moe.experts.1." not in name:
+                    param.requires_grad = False
+                else:
+                    moe_params.append(param)
+            param_groups = [{"params": moe_params, "lr": self.config["learning_rate"]}]
+            self.optimizer = AdamW(param_groups, weight_decay=0.1)
+            
+        elif "action_expert_learning_rate" in self.config:
+            # Separate learning rates for VLM and action expert parameters
+            moe_params = []
+            vlm_params = []
+            for name, param in model.named_parameters():
+                if "moe.experts.1." in name:
+                    moe_params.append(param)
+                else:
+                    vlm_params.append(param)
+
+            # 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"]}]
+            else:
+                param_groups = [
+                    {"params": vlm_params, "lr": self.config["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
+            )
+        else:
+            # Standard optimizer configuration
+            self.optimizer = AdamW(
+                model.parameters(),
+                lr=self.config["learning_rate"],
+                weight_decay=0.1,
+            )
+
+        # Configure learning rate scheduler
+        warmup_steps = self.config.get("num_warmup_steps", 0)
+        num_training_steps = self.config.get("num_training_steps", 1000000000)
+        min_lr = self.config.get("min_lr", 0.1 * self.config["learning_rate"])
+        self.lr_scheduler = get_cosine_with_min_lr_schedule_with_warmup(
+            optimizer=self.optimizer,
+            num_warmup_steps=warmup_steps,
+            num_training_steps=num_training_steps,
+            min_lr=min_lr,
+        )
+
+        self.model = model
+
+        # Enable gradient computation for embeddings
+        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)
+
+        # Prepare model, optimizer, and scheduler for distributed training
+        self.model, self.optimizer, self.lr_scheduler = self.accelerator.prepare(
+            self.model, self.optimizer, self.lr_scheduler
+        )
+
+    @timer
+    def load_qact_data(self):
+        """
+        Load and configure training data for Vision-Language-Action learning.
+        
+        Supports LeRobot dataset format and handles distributed data loading
+        across multiple processes.
+        """
+        print(f"Loading Vision-Language-Action data from {__file__}")
+        self.accelerator.wait_for_everyone()
+
+        # Load LeRobot dataset
+        self.dataset, self.train_num = load_lerobot_data(
+            self.config,
+            self.dataload_config.get("lerobot_config", {}),
+            rank=self.rank,
+            world_size=self.world_size,
+        )
+
+    @timer
+    def load_qwen_pretrain_weight(self, model, pretrain_weight_path):
+        """
+        Load pretrained Qwen weights with MoE adaptation.
+        
+        Args:
+            model: Model instance to load weights into
+            pretrain_weight_path (str): Path to pretrained weight files
+            
+        Returns:
+            Model with loaded pretrained weights
+            
+        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")])
+        merged_weights = {}
+
+        # Merge weights from all files
+        for weight_file in weight_files:
+            file_path = os.path.join(pretrain_weight_path, weight_file)
+            weights = load_file(file_path)
+            merged_weights.update(weights)
+
+        # 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:
+                # 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.")
+                renamed_weights[new_key] = value
+            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")
+                        renamed_weights[new_key] = value
+                        break
+            else:
+                renamed_weights[key] = value
+
+        # Load weights into model
+        err = model.load_state_dict(renamed_weights, strict=False)
+        self.print_rank0(f"Weight loading report: {err}", flush=True)
+        if self.accelerator.is_main_process:
+            self.print_rank0(f"Loaded pretrained weights from: {pretrain_weight_path}")
+
+        return model
+
+    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.
+        
+        Args:
+            current_epoch (int): Current epoch number
+            total_epoch (int): Total number of epochs
+            current_train_iter (int): Current training iteration
+            total_train_iter (int): Total iterations in epoch
+            loss (torch.Tensor): Current loss value
+            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"]
+        
+        log_string = f" [{datetime.now().strftime('%Y-%m-%d %H:%M:%S')}]"
+        log_string += " epoch {:3d}/{:3d} |".format(current_epoch, total_epoch)
+        log_string += " iter {:6d}/{:6d} |".format(current_train_iter, total_train_iter)
+        log_string += " loss {:.6f} |".format(loss)
+        log_string += " lr {:.6f} |".format(lr)
+        log_string += " time_per_step_avg {:.6f}s |".format(time_per_step)
+        
+        print_rank_last(log_string)
+        self.timers.log(timers_to_log, normalizer=1)
+
+    def save_checkpoint(self, epoch, step=0):
+        """
+        Save training checkpoint.
+        
+        Args:
+            epoch (int): Current epoch number
+            step (int, optional): Current step number. Defaults to 0.
+            
+        Saves model state, optimizer state, and training progress information.
+        """
+        save_path = self.config["save_path"]
+        if step == 0:
+            ckpt_path = f"{save_path}/{epoch}"
+        else:
+            ckpt_path = f"{save_path}/{epoch}_{step}"
+        
+        self.accelerator.save_state(ckpt_path)
+
+        # Save current iteration steps for dataset resuming
+        if step != 0:
+            _rank = self.accelerator.process_index
+            if isinstance(self.dataset, PreprocessedDataset):
+                torch.save(
+                    {"epoch": epoch, "step": step}, 
+                    os.path.join(ckpt_path, f"epoch_{epoch}_step_{step}_rank_{_rank}.pth")
+                )
+
+    def resume_from_checkpoint(self):
+        """
+        Resume training from a saved checkpoint.
+        
+        Handles both full checkpoint loading and model-only loading based on configuration.
+        """
+        checkpoint_path = self.config["resume"]["ckpt"]
+
+        if self.config.get("resume", {}).get("load_ckpt_only", False):
+            # Load only model weights
+            ckpt_path = self.config["resume"]["ckpt"] + "/model.safetensors"
+            state_dict = load_file(ckpt_path, device="cpu")
+            
+            # Add module prefix if needed for distributed training
+            new_state_dict = {}
+            for key in state_dict:
+                if not key.startswith("module."):
+                    new_key = "module." + key
+                new_state_dict[new_key] = state_dict[key]
+            
+            err = 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)
+        
+        self.print_rank0(f"Resumed from checkpoint: {checkpoint_path}")
+
+    def log_l1_details(self, all_label, all_pred, all_task, all_dof_mask):
+        """
+        Log detailed L1 loss metrics by degrees of freedom.
+        
+        Args:
+            all_label (torch.Tensor): Ground truth action labels
+            all_pred (torch.Tensor): Predicted actions
+            all_task (list): Task identifiers
+            all_dof_mask (torch.Tensor): Degrees of freedom mask
+            
+        Computes and logs L1 loss for each DOF component separately for detailed analysis.
+        """
+        all_task = all_task[:len(all_label)]
+
+        # Apply DOF mask
+        all_label = all_label * all_dof_mask
+        all_pred = all_pred * all_dof_mask
+        
+        # Compute baseline L1 loss (predict mean action)
+        if self.base_l1_loss is None:
+            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)
+
+        # Log L1 loss for each DOF component
+        start_idx = 0
+        dof_config = self.config.get("dof_config", {})
+        for dof in dof_config:
+            end_idx = start_idx + dof_config[dof]
+            dof_label = all_label[:, :, start_idx:end_idx]
+            dof_pred = all_pred[:, :, start_idx:end_idx]
+            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)
+            
+            start_idx = end_idx

wall_x/utils/constant.py

@@ -0,0 +1,191 @@
+action_statistic_dof = {
+    "x2_normal": {
+        # water flowers
+        "follow_left_arm_joint_cur": {
+            "min": [-3.7121],
+            "delta": [7.6008],
+        },
+        "follow_right_arm_joint_cur": {
+            "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_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_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]},
+    },
+    "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_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_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_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_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_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_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_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_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_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]},
+    },
+    "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_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_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_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_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_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_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_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_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_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_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_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_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_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_gripper": {"min": [0.0002], "delta": [0.0773]},
+    },
+    "kuka": {
+        "follow_right_ee_cartesian_pos": {
+            "min": [0.3914, -0.4901, 0.0175],
+            "delta": [0.3339, 0.8357, 0.9064],
+        },
+        "follow_right_ee_rotation": {
+            "min": [-3.1416, -0.9903, -3.1416],
+            "delta": [6.2832, 2.2421, 6.2832],
+        },
+        "follow_right_gripper": {
+            "min": [0.0000],
+            "delta": [1.0000],
+        },
+    },
+    "UMI-biarm": {
+        "follow_left_ee_cartesian_pos": {
+            "min": [-0.2917, -0.4926, 0.0063],
+            "delta": [0.9028, 0.8168, 0.3473],
+        },
+        "follow_left_ee_rotation": {
+            "min": [-2.5309, -1.5706, -1.3309],
+            "delta": [0.8758, 2.3315, 1.7076],
+        },
+        "follow_left_gripper": {
+            "min": [0.0029],
+            "delta": [0.0812],
+        },
+        "follow_right_ee_cartesian_pos": {
+            "min": [-0.0023, -0.5191, -0.0358],
+            "delta": [0.7474, 0.8668, 0.351],
+        },
+        "follow_right_ee_rotation": {
+            "min": [-2.4945, -2.0149, -0.8088],
+            "delta": [1.0941, 3.2628, 2.0018],
+        },
+        "follow_right_gripper": {
+            "min": [0.0019],
+            "delta": [0.0814],
+        },
+    },
+    "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_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_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

@@ -0,0 +1,483 @@
+import time
+import torch
+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
+    except AttributeError:
+        dist_all_gather_func = torch.distributed.all_gather
+else:
+    dist_all_gather_func = None
+
+class TimerBase(ABC):
+    """Timer base class."""
+
+    def __init__(self, name):
+        self.name = name
+
+    @abstractmethod
+    def start(self, barrier=False):
+        """Start the timer.
+
+        Args:
+            barrier (bool, optional): Synchronizes ranks before starting. Defaults to False.
+        """
+        pass
+
+    @abstractmethod
+    def stop(self, barrier=False):
+        """Stop the timer.
+
+        Args:
+            barrier (bool, optional): Synchronizes ranks before stopping. Defaults to False.
+        """
+        pass
+
+    @abstractmethod
+    def reset(self):
+        """Reset timer."""
+        pass
+
+    @abstractmethod
+    def elapsed(self, reset=True, barrier=False):
+        """Calculates the elapsed time and restarts timer.
+
+        Args:
+            reset (bool, optional): Resets timer before restarting. Defaults to True.
+            barrier (bool, optional): Synchronizes ranks before stopping. Defaults to False.
+
+        Returns:
+            float: Elapsed time.
+        """
+        pass
+
+
+class DummyTimer(TimerBase):
+    """Dummy Timer."""
+
+    def __init__(self):
+        super().__init__('dummy timer')
+
+    def start(self, barrier=False, nvtx_push=False):
+        return
+
+    def stop(self, barrier=False, nvtx_pop=False):
+        return
+
+    def reset(self):
+        return
+
+    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.'
+        )
+
+    def active_time(self):
+        """Returns the cumulative duration the timer has been active.
+        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.'
+        )
+
+
+class Timer(TimerBase):
+    """
+    Timer class with ability to start/stop.
+
+    Comment on using `barrier`: If this flag is passed, then all
+    the caller processes will wait till all reach the timing routine.
+    It is up to the user to make sure all the ranks in `barrier_group`
+    call it otherwise, it will result in a hang.
+    Comment on `barrier_group`: By default it is set to None which
+    in torch distributed land, it will result in the global communicator.
+    """
+
+    def __init__(self, name):
+        """Initialize Timer.
+
+        Args:
+            name (str): Name of the timer.
+        """
+        super().__init__(name)
+        self._elapsed = 0.0
+        self._active_time = 0.0
+        self._started = False
+        # Note that None will default to the global process group
+        self._barrier_group = None
+        self._start_time = time.time()
+        self.nvtx = False
+
+    def set_barrier_group(self, barrier_group):
+        """Sets barrier group.
+
+        Args:
+            barrier_group (ProcessGroup): Torch ProcessGroup for barrier.
+        """
+        self._barrier_group = barrier_group
+
+    def start(self, barrier=False, nvtx_push=False):
+        """Start the timer.
+
+        Args:
+            barrier (bool, optional): Synchronizes ranks before starting. Defaults to False.
+        """
+        assert not self._started, 'timer has already been started'
+        if barrier:
+            _barrier(group=self._barrier_group)
+        if torch.cuda.is_available():
+            torch.cuda.synchronize()
+        self._start_time = time.time()
+        self._started = True
+        if nvtx_push:
+            nvtx.range_push("{}".format(self.name))
+            self.nvtx = True
+            
+
+    def stop(self, barrier=False, sync=False):
+        """Stop the timer.
+
+        Args:
+            barrier (bool, optional): Synchronizes ranks before stopping. Defaults to False.
+        """
+        if self.nvtx:
+            nvtx.range_pop()
+        assert self._started, 'timer is not started'
+        if barrier:
+            _barrier(group=self._barrier_group)
+        if torch.cuda.is_available() and sync:
+            torch.cuda.synchronize()
+        elapsed = time.time() - self._start_time
+        self._elapsed += elapsed
+        self._active_time += elapsed
+        self._started = False
+
+    def reset(self):
+        """Reset timer."""
+        # Don't reset _active_time
+        self._elapsed = 0.0
+        self._started = False
+
+    def elapsed(self, reset=True, barrier=False):
+        """Calculates the elapsed time and restarts timer.
+
+        Args:
+            reset (bool, optional): Resets timer before restarting. Defaults to True.
+            barrier (bool, optional): Synchronizes ranks before stopping. Defaults to False.
+
+        Returns:
+            float: Elapsed time.
+        """
+        _started = self._started
+        # If the timing in progress, end it first.
+        if self._started:
+            self.stop(barrier=barrier)
+        # Get the elapsed time.
+        _elapsed = self._elapsed
+        # Reset the elapsed time
+        if reset:
+            self.reset()
+        # If timing was in progress, set it back.
+        if _started:
+            self.start(barrier=barrier)
+        return _elapsed
+
+    def active_time(self):
+        """Calculates the cumulative duration for which the timer has been active"""
+        return self._active_time
+
+
+class Timers:
+    """Class for a group of Timers."""
+
+    def __init__(self, log_level, log_option):
+        """Initialize group of timers.
+
+        Args:
+            log_level (int): Log level to control what timers are enabled.
+            log_option (str): Setting for logging statistics over ranks for all the timers.
+                              Allowed: ['max', 'minmax', 'all'].
+        """
+        self._log_level = log_level
+        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(
+            log_option, allowed_log_options
+        )
+        self._log_option = log_option
+        self._timers = {}
+        self._log_levels = {}
+        self._dummy_timer = DummyTimer()
+        self._max_log_level = 2
+
+    def __call__(self, name, log_level=None):
+        """Call timer with name and log level."""
+        # If the timer has already been set, then check if the log-level
+        # is provided, it matches the one that the timer was created with.
+        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)
+                )
+            return self._timers[name]
+        # If timer does not exist and no log level is provided,
+        # set it to the max log level which is 2.
+        if log_level is None:
+            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, self._max_log_level
+        )
+        # Now if the input log level is larger than the one set for
+        # the timers class, just ignore it and return a dummy timer.
+        if log_level > self._log_level:
+            return self._dummy_timer
+        # Otherwise, initalize the timer and set the level.
+        self._timers[name] = Timer(name)
+        self._log_levels[name] = log_level
+        return self._timers[name]
+
+    def _get_elapsed_time_all_ranks(self, names, reset, barrier):
+        """Returns elapsed times of timers in names.
+
+        Args:
+            names (List[str]): list of timer names
+            reset (bool): reset the timer after recording the elapsed time
+            barrier (bool): if set, do a global barrier before time measurments
+
+        Returns:
+            torch.tensor: Tensor of size [world_size, len(names)] with times in float.
+        """
+
+        # First make sure all the callers are in sync.
+        if barrier:
+            _barrier()
+
+        world_size = _get_world_size()
+        rank = _get_rank()
+
+        if torch.cuda.is_available():
+            device = torch.cuda.current_device()
+        else:
+            device = torch.device('cpu')
+
+        rank_name_to_time = torch.zeros(
+            (world_size, len(names)), dtype=torch.float, device=device
+        )
+
+        for i, name in enumerate(names):
+            if name in self._timers:
+                rank_name_to_time[rank, i] = self._timers[name].elapsed(reset=reset)
+
+        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))
+            except Exception as e:
+                print(f"Warning: all_gather failed: {e}. Using single rank timing.")
+
+        return rank_name_to_time
+
+    def _get_global_min_max_time(self, names, reset, barrier, normalizer):
+        """Report only min and max times across all ranks."""
+
+        rank_name_to_time = self._get_elapsed_time_all_ranks(names, reset, barrier)
+        name_to_min_max_time = {}
+        for i, name in enumerate(names):
+            rank_to_time = rank_name_to_time[:, i]
+            # filter out the ones we did not have any timings for
+            rank_to_time = rank_to_time[rank_to_time > 0.0]
+            # If the timer exists:
+            if rank_to_time.numel() > 0:
+                name_to_min_max_time[name] = (
+                    rank_to_time.min().item() / normalizer,
+                    rank_to_time.max().item() / normalizer,
+                )
+        return name_to_min_max_time
+
+    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)
+        if not name_to_min_max_time:
+            return None
+        
+        world_size = _get_world_size()
+        if world_size == 1:
+            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)
+        else:
+            if max_only:
+                output_string = 'max time across ranks (ms):'
+            else:
+                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)
+                else:
+                    output_string += '\n    {}: ({:.2f}, {:.2f})'.format(
+                        (name + ' ').ljust(48, '.'), min_time, max_time
+                    )
+        return output_string
+
+    def _get_all_ranks_time_string(self, names, reset, barrier, normalizer):
+        """Report times across all ranks."""
+        rank_name_to_time = self._get_elapsed_time_all_ranks(names, reset, barrier)
+        world_size = _get_world_size()
+
+        output_string = 'times across ranks (ms):'
+        no_reported_timing = True
+        for i, name in enumerate(names):
+            not_yet_found = True
+            for rank in range(world_size):
+                if rank_name_to_time[rank, i] > 0:
+                    no_reported_timing = False
+                    if not_yet_found:
+                        not_yet_found = False
+                        output_string += '\n  {}:'.format(name)
+                    if world_size == 1:
+                        output_string += '\n     {:.2f}'.format(
+                            rank_name_to_time[rank, i] / normalizer
+                        )
+                    else:
+                        output_string += '\n     rank {:2d}: {:.2f}'.format(
+                            rank, rank_name_to_time[rank, i] / normalizer
+                        )
+        if no_reported_timing:
+            return None
+        return output_string
+
+    def get_all_timers_string(
+        self,
+        names: List[str] = None,
+        normalizer: float = 1.0,
+        reset: bool = True,
+        barrier: bool = False,
+    ):
+        """Returns the output string with logged timer values according to configured options.
+
+        Args:
+            names (List[str]): Names of the timers to log. If None, all registered timers are
+                               fetched. Defaults to None.
+            normalizer (float, optional): Normalizes the timer values by the factor.
+                                          Defaults to 1.0.
+            reset (bool, optional): Whether to reset timer values after logging. Defaults to True.
+            barrier (bool, optional): Whether to do a global barrier before time measurments.
+                                      Defaults to False.
+
+        Raises:
+            Exception: Raises if log option is invalid.
+
+        Returns:
+            str: Formatted string with the timer values.
+        """
+
+        if names == None:  # get all registered timers
+            names = list(self._timers.keys())
+
+        assert normalizer > 0.0
+        if self._log_option in ['max', 'minmax']:
+            max_only = False
+            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':
+            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))
+        return output_string
+
+    def log(
+        self,
+        names: List[str],
+        rank: int = None,
+        normalizer: float = 1.0,
+        reset: bool = True,
+        barrier: bool = False,
+    ):
+        """logs the timers passed in names to stdout. Example usage is to log average per step
+           value for timer 'foo', this function can be called with normalizer factor set to logging
+           interval.
+
+        Args:
+            names (List[str]): Names of the timers to log.
+            rank (int, optional): logs the timers to a specific rank. If set to None, logs to the
+                                  last rank. Defaults to None.
+            normalizer (float, optional): Normalizes the timer values by the factor.
+                                          Defaults to 1.0.
+            reset (bool, optional): Whether to reset timer values after logging. Defaults to True.
+            barrier (bool, optional): Whether to do a global barrier before time measurments.
+                                      Defaults to False.
+        """
+
+        output_string = self.get_all_timers_string(names, normalizer, reset, barrier)
+        # If no input rank is provided, log on last rank.
+        world_size = _get_world_size()
+        current_rank = _get_rank()
+        
+        if rank is None:
+            rank = world_size - 1
+        if rank == current_rank and output_string is not None:
+            print(output_string, flush=True)
+
+    def write(
+        self,
+        names: List[str],
+        writer,
+        iteration: int,
+        normalizer: float = 1.0,
+        reset: bool = True,
+        barrier: bool = False,
+    ):
+        """Write timers to a tensorboard writer.
+        Note that we only report maximum time across ranks to tensorboard.
+
+        Args:
+            names (List[str]): Names of the timers to log.
+            writer (SummaryWriter): Tensorboard SummaryWriter object
+            iteration (int): Current iteration.
+            normalizer (float, optional): Normalizes the timer values by the factor.
+                                          Defaults to 1.0.
+            reset (bool, optional): Whether to reset timer values after logging. Defaults to True.
+            barrier (bool, optional): Whether to do a global barrier before time measurments.
+                                      Defaults to False.
+        """
+        # currently when using add_scalars,
+        # 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)
+        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)