Update Wall-X to 1.1.0 (#104)

wall_x/data/backends/lerobot/__init__.py

@@ -0,0 +1,26 @@
+"""LeRobot backend registration."""
+
+from __future__ import annotations
+
+import sys
+
+try:
+    from wall_x.data._registry import register_module
+    from wall_x.data.backends.lerobot.build import (
+        build,
+        load_trainer_data_config,
+        load_trainer_data_config_from_yaml_dict,
+    )
+
+    register_module("lerobot", sys.modules[__name__])
+except ImportError as _e:
+    from wall_x.data._registry import record_import_error
+
+    record_import_error("lerobot", _e)
+
+
+__all__ = [
+    "build",
+    "load_trainer_data_config",
+    "load_trainer_data_config_from_yaml_dict",
+]

wall_x/data/backends/lerobot/build.py

@@ -0,0 +1,337 @@
+"""LeRobot data loading bridge for typed training configs."""
+
+from __future__ import annotations
+
+import logging
+import multiprocessing as mp
+from typing import Any, Dict, Tuple
+
+import torch
+import torch.distributed as dist
+
+from wall_x.data.backends.lerobot.config import LerobotConfig
+from wall_x.data.backends.lerobot.utils import load_norm_stats
+from wall_x.model.core.action.normalizer import (
+    create_normalizers_from_lerobot_norm_stats,
+)
+
+logger = logging.getLogger(__name__)
+
+
+def load_lerobot_normalizers(cfg: Any):
+    """Create model normalizers from the LeRobot norm stats configured for a run."""
+    data = cfg.data
+
+    norm_stats_path = getattr(data, "norm_stats_path", None)
+    if not norm_stats_path:
+        return None
+
+    key_mappings = getattr(data, "key_mappings", None)
+    if not key_mappings:
+        raise ValueError(
+            "LeRobot normalizer from norm_stats_path requires data.key_mappings"
+        )
+
+    lerobot_config = getattr(data, "lerobot_config", None)
+    if not isinstance(lerobot_config, dict) or not lerobot_config.get("repo_id"):
+        raise ValueError(
+            "LeRobot normalizer from norm_stats_path requires "
+            "data.lerobot_config.repo_id"
+        )
+
+    dataset_name = str(lerobot_config["repo_id"])
+    norm_stats = load_norm_stats(
+        norm_stats_path,
+        key_mappings,
+        dof_config=dict(cfg.task.dof_config or {}),
+        agent_pos_config=dict(cfg.task.agent_pos_config or {}),
+    )
+    normalizer_action, normalizer_propri = create_normalizers_from_lerobot_norm_stats(
+        norm_stats,
+        dataset_name,
+        cfg.action_dim,
+        cfg.propri_dim,
+    )
+    return normalizer_action, normalizer_propri, norm_stats_path, dataset_name
+
+
+class _LerobotDatasetWrapper:
+    """Trainer-facing wrapper aligning PreprocessedDataset with v1 API.
+
+    PreprocessedDataset internally switches ``self._dataset`` between
+    its train/val splits via ``_train()`` / ``_eval()``. Its
+    ``get_train_dataloader`` / ``get_val_dataloader`` return
+    ``(dataloader, sampler)`` tuples and no-argument calls are supported
+    (they read rank/world_size/seed from the inner object itself).
+
+    This wrapper:
+    - Caches the rebuilt train dataloader / sampler so
+      ``set_epoch(epoch)`` can reset shuffling per-epoch.
+    - Owns the val dataloader so the trainer's ``val_loop`` can do
+      ``self.dataset.get_val_dataloader()`` and iterate directly (matching
+      what the v1/v2 wrappers return).
+    """
+
+    def __init__(
+        self,
+        inner,
+        train_dataloader: torch.utils.data.DataLoader,
+        train_sampler,
+        train_num: int,
+        val_dataloader: torch.utils.data.DataLoader = None,
+        val_num: int = 0,
+    ):
+        self._inner = inner
+        self._train_dataloader = train_dataloader
+        self._train_sampler = train_sampler
+        self._train_num = train_num
+        self._val_dataloader = val_dataloader
+        self.global_train_iters = mp.Value("i", train_num)
+        self.global_val_iters = mp.Value("i", val_num)
+
+    def __len__(self) -> int:
+        return self._train_num
+
+    def _activate_train_split(self) -> None:
+        if hasattr(self._inner, "_train"):
+            self._inner._train()
+
+    def get_train_dataloader(self):
+        self._activate_train_split()
+        return self._train_dataloader
+
+    def get_val_dataloader(self):
+        # PreprocessedDataset shares one ``_dataset`` pointer between its
+        # train and val splits (flipped by ``_train()`` / ``_eval()``).
+        # The val DataLoader's DistributedSampler caches total_size sized
+        # to the val split but ``__iter__`` reads ``len(self.dataset)``
+        # live - if a preceding train_loop left the pointer at train, that
+        # live len is ~20x total_size and DistributedSampler asserts.
+        # Rebuild each time so ``_eval()`` runs and a fresh sampler is
+        # snapped to the current (val) split length. Mirrors the train-side
+        # rebuild-on-every-epoch pattern.
+        if self._val_dataloader is None:
+            return None
+        self._val_dataloader, _ = self._inner.get_val_dataloader()
+        return self._val_dataloader
+
+    def set_epoch(self, epoch: int) -> None:
+        """Seed the per-epoch shuffle in the train DistributedSampler."""
+        self._activate_train_split()
+        if self._train_sampler is not None and hasattr(
+            self._train_sampler, "set_epoch"
+        ):
+            self._train_sampler.set_epoch(epoch)
+
+
+def load_trainer_data_config(cfg: Any) -> LerobotConfig:
+    """Build the inference/trainer data config from a typed TrainConfig."""
+    raw_yaml = dict(getattr(cfg, "_raw_yaml", {}) or {})
+    raw_data = dict(getattr(cfg, "_raw_data", {}) or {})
+    data = getattr(cfg, "data", None)
+
+    data_section = dict(raw_yaml.get("data", {}) or {})
+    data_section.update(raw_data)
+
+    if data is not None:
+        for key in (
+            "resolution",
+            "train_test_split",
+            "priority_order",
+            "camera_name_mapping",
+        ):
+            value = getattr(data, key, None)
+            if value is not None:
+                data_section.setdefault(key, value)
+
+    raw_yaml["data"] = data_section
+    raw_yaml.setdefault("model_type", getattr(cfg, "model_type", "qwen2_5"))
+    return load_trainer_data_config_from_yaml_dict(raw_yaml)
+
+
+def load_trainer_data_config_from_yaml_dict(yaml_dict: Dict[str, Any]) -> LerobotConfig:
+    """Build the LeRobot runtime config from a raw training YAML dict."""
+    return LerobotConfig.from_yaml_dict(yaml_dict)
+
+
+def _build_flat_config(cfg: Any) -> Tuple[Dict[str, Any], Dict[str, Any]]:
+    """Map typed TrainConfig -> (flat_config, lerobot_config) for legacy entry.
+
+    ``load_lerobot_data`` expects a 2509-style flat dict plus a separate
+    ``lerobot_config`` carrying ``repo_id`` / ``root``. This function is
+    the one place that translation lives; keep it surgical so future
+    field additions on ``LeRobotDataConfig`` do not require touching the
+    legacy loader.
+    """
+    model = cfg.model
+    data = cfg.data
+    hp = cfg.hyperparams
+    raw = getattr(cfg, "_raw_yaml", {}) or {}
+    raw_data = dict(getattr(cfg, "_raw_data", {}) or {})
+
+    lerobot_cfg = dict(data.lerobot_config or {})
+    if "repo_id" not in lerobot_cfg:
+        raise ValueError(
+            "lerobot requires data.lerobot_config.repo_id to be set "
+            "(HuggingFace LeRobot dataset id)."
+        )
+
+    data_section: Dict[str, Any] = {
+        "key_mappings": data.key_mappings,
+        "action_horizon": cfg.task.action_horizon,
+        "train_test_split": data.train_test_split,
+        "seed": hp.seed,
+        "resolution": data.resolution,
+    }
+    if raw_data.get("max_length") is not None:
+        data_section["max_length"] = raw_data["max_length"]
+    if data.priority_order is not None:
+        data_section["priority_order"] = data.priority_order
+    if data.camera_name_mapping is not None:
+        data_section["camera_name_mapping"] = data.camera_name_mapping
+    data_section.setdefault(
+        "use_state_string_representation",
+        cfg.task.use_state_string_representation,
+    )
+    data_section.setdefault(
+        "state_bins",
+        raw_data.get("state_bins", raw.get("state_bins", 256)),
+    )
+
+    # Dof/agent_pos totals for the collator's zero-pad step. When resuming
+    # from a checkpoint trained on a larger action space, task.dof_config
+    # should include an ``action_padding`` key that absorbs the diff; the
+    # collator right-pads action/agent_pos tensors to these totals with
+    # dof_mask/agent_pos_mask zeroed on padded dims so loss doesn't flow
+    # through them.
+    dof_total = int(sum((cfg.task.dof_config or {}).values()))
+    agent_pos_total = int(sum((cfg.task.agent_pos_config or {}).values()))
+
+    flat: Dict[str, Any] = {
+        "model_type": cfg.model_type,
+        "processor_path": getattr(model, "processor_path", "") or "",
+        "norm_stats_path": data.norm_stats_path or raw.get("norm_stats_path"),
+        "batch_size_per_gpu": hp.batch_size_per_gpu,
+        "eval_batch_size_per_gpu": raw.get(
+            "eval_batch_size_per_gpu", hp.batch_size_per_gpu
+        ),
+        "num_workers": data.num_workers,
+        "padding_side": data.padding_side,
+        "use_fast_tokenizer": data.use_fast_tokenizer,
+        "action_tokenizer_path": data.action_tokenizer_path,
+        "noise_scheduler": data.noise_scheduler or {},
+        "dof_total_dim": dof_total,
+        "agent_pos_total_dim": agent_pos_total,
+        "dof_config": dict(cfg.task.dof_config or {}),
+        "agent_pos_config": dict(cfg.task.agent_pos_config or {}),
+        "use_state_string_representation": cfg.task.use_state_string_representation,
+        "state_bins": int(
+            raw_data.get("state_bins")
+            or data_section.get("state_bins")
+            or raw.get("state_bins")
+            or 256
+        ),
+        "data": data_section,
+    }
+    return flat, lerobot_cfg
+
+
+def load_lerobot_v2(
+    cfg: Any,
+) -> Tuple[_LerobotDatasetWrapper, torch.utils.data.DataLoader, int]:
+    """Build lerobot (wrapper, dataloader, train_num) from TrainConfig.
+
+    The third return value ``train_num`` is a snapshot of
+    ``len(train_dataloader)`` at construction time. It matches
+    ``wrapper.global_train_iters.value`` initially but does not track
+    subsequent rebuilds inside ``set_epoch`` - callers doing dynamic
+    resampling should read from the mp.Value, not from this snapshot.
+    """
+    from wall_x.data.backends.lerobot.loader import load_lerobot_data
+
+    flat_cfg, lerobot_cfg = _build_flat_config(cfg)
+
+    if dist.is_initialized():
+        rank = dist.get_rank()
+        world_size = dist.get_world_size()
+    else:
+        rank = 0
+        world_size = 1
+
+    seed = cfg.hyperparams.seed
+    inner, _ = load_lerobot_data(
+        flat_cfg,
+        lerobot_cfg,
+        rank=rank,
+        world_size=world_size,
+        seed=seed,
+    )
+
+    # PreprocessedDataset.get_*_dataloader returns (dataloader, sampler).
+    # Build val first, train second, so the inner ``_dataset`` pointer is
+    # left at the train split when we finish - workers fork from that
+    # state on first iteration.
+    val_dataloader, _ = inner.get_val_dataloader()
+    val_num = len(val_dataloader) if val_dataloader is not None else 0
+
+    train_dataloader, train_sampler = inner.get_train_dataloader()
+    train_num = len(train_dataloader)
+
+    if rank == 0:
+        logger.info(
+            "\n%s\nLeRobot Data Loading Configuration:\n"
+            "  RANK: %d\n  WORLD SIZE: %d\n"
+            "  BATCH SIZE PER DEVICE: %d\n  GLOBAL BATCH SIZE: %d\n"
+            "  TRAIN BATCHES: %d\n  VAL BATCHES: %d\n"
+            "  NUM WORKERS: %d\n  REPO ID: %s\n%s",
+            "=" * 50,
+            rank,
+            world_size,
+            flat_cfg["batch_size_per_gpu"],
+            flat_cfg["batch_size_per_gpu"] * world_size,
+            train_num,
+            val_num,
+            flat_cfg["num_workers"],
+            lerobot_cfg.get("repo_id"),
+            "=" * 50,
+        )
+
+    wrapper = _LerobotDatasetWrapper(
+        inner,
+        train_dataloader,
+        train_sampler,
+        train_num,
+        val_dataloader=val_dataloader,
+        val_num=val_num,
+    )
+    return wrapper, train_dataloader, train_num
+
+
+def build(cfg, ctx):
+    """Backend Protocol entry - returns a ``DataBundle``.
+
+    Wraps ``load_lerobot_v2`` (which returns the trainer-facing triple)
+    into the unified ``DataBundle`` shape every backend exposes.
+    """
+    from wall_x.data._bundle import DataBundle
+
+    wrapper, train_dataloader, train_num = load_lerobot_v2(cfg)
+
+    # PreprocessedDataset shares one ``self._dataset`` pointer between
+    # train and val splits (flipped by ``_train()`` / ``_eval()``).
+    # ``wrapper.get_val_dataloader()`` flips the pointer to val. Flip back once
+    # here so the initial train loop starts from the right split even if callers
+    # inspect the raw ``train_dataloader`` before invoking ``set_epoch``.
+    val_loader = wrapper.get_val_dataloader()
+    inner = wrapper._inner
+    if hasattr(inner, "_train"):
+        inner._train()
+
+    return DataBundle(
+        dataset=wrapper,
+        train_loader=train_dataloader,
+        val_loader=val_loader,
+        train_iters=train_num,
+        val_iters=wrapper.global_val_iters.value,
+        set_epoch=wrapper.set_epoch,
+    )

wall_x/data/backends/lerobot/config.py

@@ -0,0 +1,134 @@
+from dataclasses import dataclass, field
+from typing import Any, Dict, Optional
+
+from qwen_vl_utils.vision_process import IMAGE_FACTOR, MAX_PIXELS, MIN_PIXELS
+
+
+@dataclass
+class LerobotConfig:
+    """Configuration for the LeRobot preprocessing pipeline.
+
+    Dataset-specific camera display names are optional config inputs. Other
+    dataset behavior is derived from the current LeRobot sample.
+    """
+
+    # 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
+    seed: int = 42
+
+    # Instruction handling
+    priority_order: Optional[Dict[str, float]] = None
+    camera_name_mapping: Optional[Dict[str, str]] = 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) -> "LerobotConfig":
+        """Update configuration parameters.
+
+        Args:
+            **kwargs: Key-value pairs to update
+
+        Returns:
+            LerobotConfig: 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
+
+    def __getitem__(self, key: str):
+        return getattr(self, key)
+
+    @classmethod
+    def from_yaml_dict(cls, yaml_dict: Dict[str, Any]) -> "LerobotConfig":
+        """
+        Build a LerobotConfig instance from a YAML dictionary.
+
+        Supports two styles:
+
+        1) Top-level fields:
+            train_test_split: 0.8
+            model_type: qwen2_5
+
+        2) Nested under `data:` (higher priority):
+            data:
+              train_test_split: 0.8
+              model_type: qwen2_5
+
+        Keys inside `data:` override top-level keys.
+        """
+
+        data_config = yaml_dict.get("data", {})
+
+        def get(key: str, default: Any = None):
+            """
+            Helper function:
+            Read from `data` first, then fallback to the top-level YAML.
+            """
+            return data_config.get(key, yaml_dict.get(key, default))
+
+        # Construct only fields that actually exist in LerobotConfig
+        params: Dict[str, Any] = {
+            # Action prediction settings
+            # Image resolution per camera view
+            "resolution": get(
+                "resolution",
+                {
+                    "face_view": -1,
+                    "left_wrist_view": 128,
+                    "right_wrist_view": 128,
+                },
+            ),
+            # Dataset train/test split configuration
+            "train_test_split": get("train_test_split", 0.9),
+            "seed": get("seed", 42),
+            # Instruction priority ordering (optional)
+            "priority_order": get("priority_order", None),
+            "camera_name_mapping": get("camera_name_mapping", None),
+            # Vision model parameters
+            "model_type": get("model_type", "qwen2_5"),
+            "max_pixels": get("max_pixels", MAX_PIXELS),
+            "min_pixels": get("min_pixels", MIN_PIXELS),
+            "image_factor": get("image_factor", IMAGE_FACTOR),
+            # Subtask generation ratio
+            "generate_subtask_ratio": get("generate_subtask_ratio", 0.0),
+        }
+
+        # Keep only valid dataclass fields (ignore unknown YAML keys)
+        valid_fields = {f.name for f in cls.__dataclass_fields__.values()}
+        filtered_params = {k: v for k, v in params.items() if k in valid_fields}
+
+        return cls(**filtered_params)

wall_x/data/backends/lerobot/loader.py

@@ -0,0 +1,895 @@
+"""
+LeRobot Dataset Loader - Distributed Version
+"""
+
+import logging
+import os
+from typing import Protocol, SupportsIndex, TypeVar
+
+import numpy as np
+import torch
+from lerobot.datasets.lerobot_dataset import LeRobotDataset, LeRobotDatasetMetadata
+from qwen_vl_utils.vision_process import smart_resize
+from torch.utils.data import DistributedSampler, random_split
+from transformers import AutoProcessor
+
+from wall_x._vendor.x2robot_utils.geometry import (
+    canonicalize_euler_zyx_batch_nb,
+    euler_to_matrix_zyx_batch_nb,
+    matrix_to_euler_zyx_batch_nb,
+    so3_to_matrix_batch_nb,
+)
+from wall_x.data.backends.lerobot.config import LerobotConfig
+from wall_x.data.backends.lerobot.rotation_layout import (
+    LAYOUT_SKIP_KEYS,
+    maybe_convert_euler_to_6d,
+)
+from wall_x.data.backends.lerobot.rotation_layout import (
+    euler_layout_dim as _euler_layout_dim,
+)
+from wall_x.data.backends.lerobot.rotation_layout import (
+    layout_uses_6d_rotation as _layout_uses_6d_rotation,
+)
+from wall_x.data.backends.lerobot.utils import (
+    get_wallx_normal_text,
+    load_norm_stats,
+    preprocesser_call,
+    process_grounding_points,
+    replace_action_token,
+)
+
+T_co = TypeVar("T_co", covariant=True)
+logger = logging.getLogger(__name__)
+
+RELATIVE_KEYWORD = "relative"
+ROTATION_KEYWORD = "rotation"
+RELATIVE_SKIP_KEYS = LAYOUT_SKIP_KEYS
+
+
+def _compute_delta_from_state_and_abs_rot(
+    rotations: np.ndarray, state: np.ndarray
+) -> np.ndarray:
+    """Relative rotation: R_rel = R_abs @ R_state^T."""
+    if rotations.shape[-1] == 3:
+        rotations_matrix = euler_to_matrix_zyx_batch_nb(rotations)
+        out_is_euler = True
+    elif rotations.shape[-1] == 6:
+        rotations_matrix = so3_to_matrix_batch_nb(rotations)
+        out_is_euler = False
+    else:
+        raise ValueError(
+            f"Only 3D euler or 6D rotation supported, got {rotations.shape[-1]}D"
+        )
+
+    if state.shape[-1] == 3:
+        state_matrix = euler_to_matrix_zyx_batch_nb(state[np.newaxis, :])[0]
+    elif state.shape[-1] == 6:
+        state_matrix = so3_to_matrix_batch_nb(state[np.newaxis, :])[0]
+    else:
+        raise ValueError(
+            f"Only 3D euler or 6D rotation supported, got {state.shape[-1]}D"
+        )
+
+    r_rel = np.matmul(rotations_matrix, state_matrix.T)
+    if out_is_euler:
+        d_euler = matrix_to_euler_zyx_batch_nb(r_rel)
+        return canonicalize_euler_zyx_batch_nb(d_euler)
+    return r_rel[:, :2, :].reshape(r_rel.shape[0], 6)
+
+
+# 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,
+        norm_stats,
+        dataload_config,
+        lerobot_config,
+        seed=42,
+        rank=0,
+        world_size=1,
+        test_only=False,
+    ):
+        self.hf_dataset = dataset
+
+        if test_only:
+            self._dataset = dataset
+        else:
+            self._dataset = None
+            self.train_dataset, self.val_dataset = random_split(
+                dataset,
+                [0.95, 0.05],
+                torch.Generator().manual_seed(seed) if seed is not None else None,
+            )
+            self._train()
+
+        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.norm_stats = norm_stats
+        self.lerobot_config = lerobot_config
+
+        self.data_config = LerobotConfig().update(
+            train_test_split=self.dataload_config["train_test_split"],
+            seed=self.dataload_config["seed"],
+            resolution=self.dataload_config.get("resolution", None),
+            priority_order=self.dataload_config.get("priority_order", None),
+            camera_name_mapping=self.dataload_config.get("camera_name_mapping", None),
+        )
+
+        self.key_mappings = self.dataload_config["key_mappings"]
+
+        self._cam_key_mapping = self.key_mappings["camera"]
+        self._state_key_mapping = self.key_mappings
+        self._action_key_mapping = self.key_mappings
+
+        task_cfg = self.config.get("task") or {}
+        self._dof_config = self.config.get("dof_config") or task_cfg.get(
+            "dof_config", {}
+        )
+        self._agent_pos_config = self.config.get("agent_pos_config") or task_cfg.get(
+            "agent_pos_config", {}
+        )
+        self._use_relative_action = any(
+            RELATIVE_KEYWORD in key for key in self._dof_config
+        )
+        self._convert_action_euler_to_6d = _layout_uses_6d_rotation(self._dof_config)
+        self._convert_state_euler_to_6d = _layout_uses_6d_rotation(
+            self._agent_pos_config
+        )
+        if self._convert_action_euler_to_6d or self._convert_state_euler_to_6d:
+            logger.info(
+                "LeRobot loader: Euler->6D rotation enabled "
+                "(action=%s, state=%s; raw action dim=%s -> %s)",
+                self._convert_action_euler_to_6d,
+                self._convert_state_euler_to_6d,
+                (
+                    _euler_layout_dim(self._dof_config)
+                    if self._convert_action_euler_to_6d
+                    else "-"
+                ),
+                (
+                    sum(
+                        d
+                        for k, d in self._dof_config.items()
+                        if k not in RELATIVE_SKIP_KEYS
+                    )
+                    if self._convert_action_euler_to_6d
+                    else "-"
+                ),
+            )
+
+    def _maybe_convert_euler_to_6d(self, vec, layout_config: dict, enabled: bool):
+        converted = maybe_convert_euler_to_6d(vec, layout_config, enabled)
+        if (
+            enabled
+            and layout_config
+            and isinstance(vec, torch.Tensor)
+            and converted is not vec
+        ):
+            return torch.as_tensor(converted, dtype=vec.dtype, device=vec.device)
+        return converted
+
+    def _to_relative_action(self, action, agent_pos):
+        """Convert absolute action horizon to deltas w.r.t. current agent_pos."""
+        action = np.asarray(action, dtype=np.float64)
+        agent_pos = np.asarray(agent_pos, dtype=np.float64)
+        if action.ndim == 1:
+            action = action[np.newaxis, :]
+        if agent_pos.ndim > 1:
+            agent_pos = agent_pos.reshape(-1)
+
+        parts = []
+        cur = 0
+        for key, dim in self._dof_config.items():
+            if key in RELATIVE_SKIP_KEYS:
+                continue
+            action_clip = action[:, cur : cur + dim]
+            agent_pos_clip = agent_pos[cur : cur + dim]
+            if RELATIVE_KEYWORD not in key:
+                parts.append(action_clip)
+            elif ROTATION_KEYWORD in key:
+                parts.append(
+                    _compute_delta_from_state_and_abs_rot(
+                        action_clip.astype(np.float64),
+                        agent_pos_clip.astype(np.float64),
+                    )
+                )
+            else:
+                parts.append(action_clip - agent_pos_clip[np.newaxis, :])
+            cur += dim
+
+        if not parts:
+            return action
+        return np.concatenate(parts, axis=1).astype(np.float32)
+
+    def _vision_preprocess(self, frames):
+        processed_frames = []
+        for key in self.hf_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[self._state_key_mapping["state"]]
+        action = data[self._action_key_mapping["action"]]
+        agent_pos = self._maybe_convert_euler_to_6d(
+            agent_pos, self._agent_pos_config, self._convert_state_euler_to_6d
+        )
+        action = self._maybe_convert_euler_to_6d(
+            action, self._dof_config, self._convert_action_euler_to_6d
+        )
+        if self._use_relative_action:
+            device = action.device if isinstance(action, torch.Tensor) else None
+            action = torch.as_tensor(
+                self._to_relative_action(action, agent_pos),
+                dtype=torch.float32,
+                device=device,
+            )
+        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,
+            self.dataload_config.get("action_horizon", 33) - 1,
+            frame_index,
+            self.data_config.priority_order,
+            self._cam_key_mapping,
+            generate_subtask_ratio=generate_subtask_ratio,
+            camera_name_mapping=self.data_config.camera_name_mapping,
+        )
+        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 _eval(self):
+        self._dataset = self.val_dataset
+
+    def _train(self):
+        self._dataset = self.train_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
+        """
+        self._train()
+
+        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.norm_stats, self.lerobot_config
+            ),
+            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)
+        """
+        self._eval()
+
+        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.norm_stats, self.lerobot_config
+            ),
+            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 = {}
+    _norm_stat_alignment_warnings = set()
+
+    def __init__(self, config, dataload_config, stats, lerobot_config):
+        self.config = config
+        self.dataload_config = dataload_config
+        self.stats = stats
+        self.action_min_stat = stats["action"].min
+        self.action_delta = stats["action"].delta
+        self.state_min_stat = stats["state"].min
+        self.state_delta = stats["state"].delta
+        self.lerobot_config = lerobot_config
+        self.np_rng = np.random.default_rng()
+
+        noise_scheduler_config = config.get("noise_scheduler", {})
+        self.beta_alpha = noise_scheduler_config.get(
+            "beta_alpha", 1.5
+        )  # alpha parameter of the Beta distribution
+        self.beta_beta = noise_scheduler_config.get(
+            "beta_beta", 1.0
+        )  # beta parameter of the Beta distribution
+        self.s = noise_scheduler_config.get("s", 0.999)  # scaling factor
+        self.time_shift = noise_scheduler_config.get(
+            "time_shift", 1.0
+        )  # time shift factor
+
+        self.beta_alpha = float(self.beta_alpha)
+        self.beta_beta = float(self.beta_beta)
+        self.use_fast_tokenizer = self.config.get("use_fast_tokenizer", False)
+        self.use_state_string_representation = bool(
+            self.config.get("use_state_string_representation", False)
+        )
+        self.state_bins = int(self.config.get("state_bins", 256))
+        self.load_processor()
+
+    def load_processor(self):
+        processor_path = self.config["processor_path"]
+        action_tokenizer_path = self.config.get("action_tokenizer_path", None)
+
+        if (
+            self.use_fast_tokenizer
+            and action_tokenizer_path not in self._action_tokenizer_cache
+        ):
+            self._action_tokenizer_cache[action_tokenizer_path] = (
+                AutoProcessor.from_pretrained(
+                    action_tokenizer_path, trust_remote_code=True
+                )
+            )
+
+        # Use cached processors if available
+        if processor_path not in self._processor_cache:
+            processor = AutoProcessor.from_pretrained(processor_path, use_fast=True)
+            if self.config.get("padding_side", "left") == "left":
+                processor.tokenizer.padding_side = "left"
+
+            new_tokens = ["<|propri|>", "<|action|>"]
+            processor.tokenizer.add_tokens(new_tokens)
+            if self.use_fast_tokenizer and self.config.get("model_type") == "qwen2_5":
+                action_tokenizer = self._action_tokenizer_cache[action_tokenizer_path]
+                new_tokens = [
+                    f"<|action_token_{i}|>" for i in range(action_tokenizer.vocab_size)
+                ]
+                processor.tokenizer.add_tokens(new_tokens)
+                begin_idx_token = "<|action_token_0|>"
+                token_id = processor.tokenizer.convert_tokens_to_ids(begin_idx_token)
+                processor.tokenizer.init_kwargs["action_token_start_index"] = token_id
+                processor.tokenizer.init_kwargs["action_token_vocab_size"] = (
+                    action_tokenizer.vocab_size
+                )
+
+            self._processor_cache[processor_path] = processor
+
+        self.processor = self._processor_cache[processor_path]
+
+        if not self.use_fast_tokenizer:
+            self.train_action_tokenizer = None
+        else:
+            self.train_action_tokenizer = self._action_tokenizer_cache[
+                action_tokenizer_path
+            ]
+
+    @classmethod
+    def _normalize(cls, action, min_stat, delta):
+        """
+        Normalize action data using min-max normalization.
+        """
+        delta = torch.where(delta == 0, torch.ones_like(delta), delta)
+        x = (action - min_stat) / delta
+        x = x * 2 - 1
+        x = torch.clamp(x, -1, 1)
+        return x
+
+    @staticmethod
+    def _align_norm_stat(stat, value, *, pad_value: float, name: str):
+        """Align a 1-D norm stat with the current LeRobot tensor width."""
+        stat = stat.to(device=value.device, dtype=value.dtype)
+        target_dim = value.shape[-1]
+        stat_dim = stat.shape[-1]
+        if stat_dim == target_dim:
+            return stat
+        if stat_dim > target_dim:
+            warning_key = ("truncate", name, stat_dim, target_dim)
+            if warning_key not in DataCollator._norm_stat_alignment_warnings:
+                logger.warning(
+                    "Truncating LeRobot %s norm stat from %s to %s dims",
+                    name,
+                    stat_dim,
+                    target_dim,
+                )
+                DataCollator._norm_stat_alignment_warnings.add(warning_key)
+            return stat[..., :target_dim]
+        pad_shape = (*stat.shape[:-1], target_dim - stat_dim)
+        pad = torch.full(pad_shape, pad_value, device=value.device, dtype=value.dtype)
+        warning_key = ("pad", name, stat_dim, target_dim)
+        if warning_key not in DataCollator._norm_stat_alignment_warnings:
+            logger.warning(
+                "Padding LeRobot %s norm stat from %s to %s dims",
+                name,
+                stat_dim,
+                target_dim,
+            )
+            DataCollator._norm_stat_alignment_warnings.add(warning_key)
+        return torch.cat([stat, pad], dim=-1)
+
+    def __call__(self, batch):
+        additional_inputs = {}
+
+        # Tail-pad widths when dof_config / agent_pos_config (sum) is larger
+        # than the lerobot action/state - typical when resuming a ckpt that
+        # was pretrained on a bigger action space. Extra columns are filled
+        # with zeros and their mask set to 0 so loss is not propagated.
+        dof_total = int(self.config.get("dof_total_dim", 0) or 0)
+        agent_pos_total = int(self.config.get("agent_pos_total_dim", 0) or 0)
+
+        # Explicit init so the ``if action is not None`` guard and later
+        # ``replace_action_token`` call stay well-defined even if a batch
+        # unexpectedly omits the action / agent_pos keys. Without this the
+        # loop-local variables would leak NameError on the first miss.
+        action = None
+        dof_mask = None
+        agent_pos = None
+        agent_pos_mask = None
+
+        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)
+                state_min_stat = self._align_norm_stat(
+                    self.state_min_stat,
+                    agent_pos,
+                    pad_value=0.0,
+                    name="state.min",
+                )
+                state_delta = self._align_norm_stat(
+                    self.state_delta,
+                    agent_pos,
+                    pad_value=1.0,
+                    name="state.delta",
+                )
+                agent_pos = self._normalize(agent_pos, state_min_stat, state_delta)
+                if agent_pos_total and agent_pos.shape[-1] < agent_pos_total:
+                    pad_w = agent_pos_total - agent_pos.shape[-1]
+                    agent_pos = torch.nn.functional.pad(agent_pos, (0, pad_w))
+                    agent_pos_mask = torch.nn.functional.pad(agent_pos_mask, (0, pad_w))
+                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_min_stat = self._align_norm_stat(
+                    self.action_min_stat,
+                    action,
+                    pad_value=0.0,
+                    name="action.min",
+                )
+                action_delta = self._align_norm_stat(
+                    self.action_delta,
+                    action,
+                    pad_value=1.0,
+                    name="action.delta",
+                )
+                action = self._normalize(action, action_min_stat, action_delta)
+                if dof_total and action.shape[-1] < dof_total:
+                    pad_w = dof_total - action.shape[-1]
+                    action = torch.nn.functional.pad(action, (0, pad_w))
+                    dof_mask = torch.nn.functional.pad(dof_mask, (0, pad_w))
+                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"
+                )
+
+        # sample noise time
+        if action is not None:
+            sample_time = self.sample_time(
+                action.shape[0],
+                device=action.device,
+                dtype=torch.float32,
+            )
+            additional_inputs["sample_time"] = sample_time
+
+        additional_inputs["text"] = replace_action_token(
+            additional_inputs["text"],
+            additional_inputs["action_chunk"],
+            self.train_action_tokenizer if self.use_fast_tokenizer else None,
+            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),
+            norm_state=(
+                additional_inputs["proprioception"]
+                if self.use_state_string_representation
+                and "proprioception" in additional_inputs
+                else None
+            ),
+            agent_pos_mask=additional_inputs.get("agent_pos_mask"),
+            state_bins=self.state_bins,
+        )
+
+        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"] = [self.lerobot_config["repo_id"]] * inputs[
+            "action_chunk"
+        ].shape[0]
+
+        return inputs
+
+    def sample_time(self, batch_size, device, dtype):
+        """
+        Sample timesteps
+
+        Use a Beta distribution to sample values in [0, 1], then scale them.
+
+        Args:
+            batch_size (int): batch size
+            device: Device type
+            dtype: dtype
+
+        Returns:
+            torch.Tensor: sampled timesteps with shape [batch_size]
+        """
+
+        sample_np = self.np_rng.beta(
+            self.beta_alpha, self.beta_beta, size=(batch_size,)
+        ).astype(np.float32)
+        sample = torch.from_numpy(sample_np).to(
+            device=device, dtype=dtype, non_blocking=True
+        )
+
+        # sample = self.beta_dist.sample([batch_size]).to(dtype=dtype)
+        time = 1 - sample
+
+        # Apply diffusion time shift
+        if self.time_shift != 1.0:
+            time = (self.time_shift * time) / (1 + (self.time_shift - 1) * time)
+
+        time = time * self.s  # noise should denoise from 0 to 1 here
+        return time
+
+
+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)
+
+    dataload_config = get_data_configs(config["data"])
+    key_mappings = dataload_config["key_mappings"]
+
+    repo_id = lerobot_config.get("repo_id", None)
+    assert repo_id is not None, "repo id is required"
+    root = lerobot_config.get("root", None)
+    meta_info = LeRobotDatasetMetadata(repo_id, root=root)
+    dataset_fps = meta_info.fps
+    episodes_num = meta_info.total_episodes
+
+    norm_stats_path = config.get("norm_stats_path", None)
+    assert (
+        norm_stats_path is not None
+    ), "norm stats is required, please refer to 'wall-x/scripts/compute_norm_stats.py' to compute stats"
+    task_cfg = config.get("task") or {}
+    dof_config = config.get("dof_config") or task_cfg.get("dof_config", {})
+    agent_pos_config = config.get("agent_pos_config") or task_cfg.get(
+        "agent_pos_config", {}
+    )
+    norm_stats = load_norm_stats(
+        norm_stats_path,
+        key_mappings,
+        dof_config=dof_config,
+        agent_pos_config=agent_pos_config,
+    )
+
+    delta_timestamps = {
+        # action chunk
+        key_mappings["action"]: [
+            t / dataset_fps
+            for t in range(dataload_config.get("action_horizon", 33) - 1)
+        ],
+    }
+    batch_size = config.get("batch_size_per_gpu", 8)
+
+    # Optional episode subset. YAML ``lerobot_config.episodes`` has always
+    # been present in examples but previously ignored; honour it so smoke
+    # tests / small-dataset runs don't pay the O(N) LeRobotDataset indexing
+    # cost on a multi-thousand-episode repo (~10s / episode on some formats).
+    episodes_override = lerobot_config.get("episodes")
+    if episodes_override is not None:
+        episodes = list(episodes_override)
+        episodes_num_effective = len(episodes)
+    else:
+        episodes = np.arange(episodes_num).tolist()
+        episodes_num_effective = episodes_num
+
+    train_test_split = dataload_config.get("train_test_split", 0.95)
+    split_idx = int(episodes_num_effective * train_test_split)
+    # Guard: tiny episode subsets + high train_test_split can floor split_idx
+    # to 0 (e.g. 1 ep * 0.95 = 0), which would silently hand LeRobotDataset an
+    # empty episode list and end training after 0 iterations. Fail loud.
+    if split_idx < 1:
+        raise ValueError(
+            f"train_test_split={train_test_split} applied to "
+            f"{episodes_num_effective} episode(s) yields 0 train episodes. "
+            f"Use more episodes or a lower train_test_split."
+        )
+    train_episodes = episodes[:split_idx]
+    test_episodes = episodes[split_idx:]
+
+    global_rank = torch.distributed.get_rank()
+    local_rank = int(os.environ["LOCAL_RANK"])
+    local_world_size = int(os.environ["LOCAL_WORLD_SIZE"])
+    # TODO: Some LeRobot formats need to load all metadata before splitting
+    # episodes; loading from all ranks at once can exhaust memory.
+    train_dataset = None
+
+    # Sequential loading inside each node
+    for r in range(local_world_size):
+        if local_rank == r:
+            logger.info(
+                "[Global rank %s] Loading dataset on local_rank=%s",
+                global_rank,
+                local_rank,
+            )
+
+            train_dataset = LeRobotDataset(
+                repo_id=repo_id,
+                root=root,
+                episodes=train_episodes,
+                delta_timestamps=delta_timestamps,
+                video_backend="pyav",
+            )
+
+            logger.info(
+                "[Global rank %s] Finished loading on local_rank=%s",
+                global_rank,
+                local_rank,
+            )
+
+        # Barrier only within the node
+        torch.distributed.barrier(device_ids=[local_rank])
+
+    if rank == 0:
+        logger.info("Selected train episodes: %s", train_dataset.episodes)
+        logger.info("Number of train episodes selected: %s", train_dataset.num_episodes)
+        logger.info("Number of train frames selected: %s", train_dataset.num_frames)
+        logger.info("Selected test episodes: %s", test_episodes)
+
+    dataset = PreprocessedDataset(
+        train_dataset,
+        config,
+        norm_stats,
+        dataload_config,
+        lerobot_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:
+        lines = [
+            "LeRobot Data Loading Configuration:",
+            f"  rank: {rank}",
+            f"  world_size: {world_size}",
+            f"  batch_size_per_gpu: {batch_size}",
+            f"  repo_id: {repo_id}",
+            f"  total_dataset_size: {len(dataset)}",
+        ]
+        if world_size > 1:
+            lines.extend(
+                [
+                    f"  samples_per_process: {samples_per_process}",
+                    f"  batches_per_process: {train_num}",
+                    f"  total_batches_all_processes: {train_num * world_size}",
+                ]
+            )
+        else:
+            lines.append(f"  total_batches: {train_num}")
+        lines.append(f"  seed: {seed}")
+        logger.info("\n%s", "\n".join(lines))
+
+    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,
+        "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": [],
+        "camera_name_mapping": None,
+    }
+    data_config = default_data_config | config
+    data_config["action_horizon"] += 1
+
+    return data_config

wall_x/data/backends/lerobot/rotation_layout.py

@@ -0,0 +1,95 @@
+"""Layout helpers when config expects 6D rotation but LeRobot stores 3D Euler."""
+
+from __future__ import annotations
+
+import numpy as np
+
+from wall_x._vendor.x2robot_utils.geometry import euler_to_matrix_zyx_6d_nb
+
+LAYOUT_SKIP_KEYS = frozenset(
+    {"velocity_decomposed", "height", "head_actions", "action_padding"}
+)
+ROTATION_KEYWORD = "rotation"
+ROTATION_6D_KEYWORD = "6D"
+
+
+def layout_uses_6d_rotation(layout_config: dict) -> bool:
+    for key, dim in layout_config.items():
+        if key in LAYOUT_SKIP_KEYS:
+            continue
+        if ROTATION_KEYWORD in key and ROTATION_6D_KEYWORD in key and dim == 6:
+            return True
+    return False
+
+
+def euler_layout_dim(layout_config: dict) -> int:
+    """Vector width in LeRobot when rotation slices are still 3D Euler."""
+    total = 0
+    for key, dim in layout_config.items():
+        if key in LAYOUT_SKIP_KEYS:
+            continue
+        if ROTATION_KEYWORD in key and ROTATION_6D_KEYWORD in key and dim == 6:
+            total += 3
+        else:
+            total += int(dim)
+    return total
+
+
+def convert_euler_to_6d(vec: np.ndarray, layout_config: dict) -> np.ndarray:
+    """Rewrite [pos, euler(3), tail...] to [pos, rot6d(6), tail...] per layout."""
+    vec = np.asarray(vec, dtype=np.float64)
+    single = vec.ndim == 1
+    if single:
+        vec = vec[np.newaxis, :]
+
+    out_rows = []
+    for row in vec:
+        parts: list[np.ndarray] = []
+        raw_cur = 0
+        for key, dim in layout_config.items():
+            if key in LAYOUT_SKIP_KEYS:
+                continue
+            dim = int(dim)
+            if ROTATION_KEYWORD in key and ROTATION_6D_KEYWORD in key and dim == 6:
+                euler = row[raw_cur : raw_cur + 3]
+                rot6d = euler_to_matrix_zyx_6d_nb(euler.reshape(1, 3)).reshape(6)
+                parts.append(rot6d)
+                raw_cur += 3
+            else:
+                parts.append(row[raw_cur : raw_cur + dim])
+                raw_cur += dim
+        out_rows.append(np.concatenate(parts, axis=0))
+
+    out = np.stack(out_rows, axis=0)
+    return out[0] if single else out
+
+
+def maybe_convert_norm_stats_vector(
+    values,
+    layout_config: dict,
+    enabled: bool | None = None,
+):
+    """Convert a 1D norm-stat vector (q01/q99/mean/std) from Euler layout to 6D."""
+    if enabled is None:
+        enabled = layout_uses_6d_rotation(layout_config)
+    if not enabled or not layout_config:
+        return values
+    arr = np.asarray(values, dtype=np.float64)
+    if arr.ndim != 1:
+        return values
+    raw_dim = euler_layout_dim(layout_config)
+    if arr.shape[0] != raw_dim:
+        return values
+    return convert_euler_to_6d(arr, layout_config).astype(np.float32)
+
+
+def maybe_convert_euler_to_6d(
+    vec: np.ndarray, layout_config: dict, enabled: bool
+) -> np.ndarray:
+    if not enabled or not layout_config:
+        return vec
+    raw_dim = euler_layout_dim(layout_config)
+    arr = np.asarray(vec)
+    if arr.shape[-1] != raw_dim:
+        return vec
+    return convert_euler_to_6d(arr, layout_config).astype(np.float32)

wall_x/data/backends/lerobot/utils.py

@@ -0,0 +1,673 @@
+import json
+import logging
+import random
+import re
+from collections import OrderedDict
+from dataclasses import dataclass
+from typing import Any, Dict, List, Optional, Tuple, Union
+
+import numpy as np
+import torch
+from transformers import BatchFeature
+
+logger = logging.getLogger(__name__)
+
+
+@dataclass
+class NormStats:
+    min: torch.Tensor
+    max: torch.Tensor
+    delta: torch.Tensor
+
+
+def load_norm_stats(
+    norm_stats_path,
+    key_mappings,
+    dof_config: dict | None = None,
+    agent_pos_config: dict | None = None,
+):
+    from wall_x.data.backends.lerobot.rotation_layout import (
+        layout_uses_6d_rotation,
+        maybe_convert_norm_stats_vector,
+    )
+
+    with open(norm_stats_path, "r") as f:
+        norm_stats = json.load(f)
+    action_key = key_mappings["action"]
+    state_key = key_mappings["state"]
+
+    action_q01 = maybe_convert_norm_stats_vector(
+        norm_stats["norm_stats"][action_key]["q01"], dof_config or {}
+    )
+    action_q99 = maybe_convert_norm_stats_vector(
+        norm_stats["norm_stats"][action_key]["q99"], dof_config or {}
+    )
+    if layout_uses_6d_rotation(dof_config or {}) and len(action_q01) != len(
+        norm_stats["norm_stats"][action_key]["q01"]
+    ):
+        logger.info(
+            "Converted action norm stats Euler->6D (%d -> %d dims)",
+            len(norm_stats["norm_stats"][action_key]["q01"]),
+            len(action_q01),
+        )
+
+    q01 = torch.tensor(action_q01)
+    q99 = torch.tensor(action_q99)
+    delta = q99 - q01
+    action_norm_stats = NormStats(
+        min=q01,
+        max=q99,
+        delta=delta,
+    )
+
+    state_q01 = maybe_convert_norm_stats_vector(
+        norm_stats["norm_stats"][state_key]["q01"], agent_pos_config or {}
+    )
+    state_q99 = maybe_convert_norm_stats_vector(
+        norm_stats["norm_stats"][state_key]["q99"], agent_pos_config or {}
+    )
+    if layout_uses_6d_rotation(agent_pos_config or {}) and len(state_q01) != len(
+        norm_stats["norm_stats"][state_key]["q01"]
+    ):
+        logger.info(
+            "Converted state norm stats Euler->6D (%d -> %d dims)",
+            len(norm_stats["norm_stats"][state_key]["q01"]),
+            len(state_q01),
+        )
+
+    q01 = torch.tensor(state_q01)
+    q99 = torch.tensor(state_q99)
+    delta = q99 - q01
+    state_norm_stats = NormStats(
+        min=q01,
+        max=q99,
+        delta=delta,
+    )
+
+    return {"action": action_norm_stats, "state": state_norm_stats}
+
+
+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_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 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 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_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,
+    camera_name_mapping: Optional[Dict[str, str]] = None,
+) -> 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
+        camera_name_mapping: Optional display-name mapping for prompt text
+
+    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:
+        camera_name_mapping = camera_name_mapping or {}
+        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 replace_action_token(
+    text: List[str],
+    norm_action: Optional[torch.Tensor],
+    action_tokenizer,
+    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
+        dof_masks: Masks for degrees of freedom
+
+    Returns:
+        List of text strings with action tokens replaced
+    """
+    if action_tokenizer is not None and norm_action is not None:
+        if dof_masks is not None:
+            norm_action = [
+                action[:, 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
+
+
+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",
+    norm_state=None,
+    agent_pos_mask=None,
+    state_bins: int = 256,
+    state_drop_prob: float = 0.0,
+) -> 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. transformers>=5.2 split image/video processing
+    # onto distinct callables and no longer accepts ``videos=`` on
+    # image_processor (or vice versa); older versions accepted ``videos=None``.
+    # Only pass the kwargs that correspond to present inputs.
+    if images is not None and len(images) > 0:
+        image_inputs = processor.image_processor(
+            images=images, return_tensors=return_tensors
+        )
+        image_grid_thw = image_inputs["image_grid_thw"]
+    else:
+        image_inputs = {}
+        image_grid_thw = None
+
+    if videos is not None:
+        # transformers>=5.2 split video processing onto a dedicated
+        # ``video_processor``; older versions accepted ``videos=`` on the
+        # image_processor. Prefer the new API when present and fail loud
+        # otherwise, since passing ``videos=`` to an old-style
+        # image_processor silently works but feeds through a different
+        # preprocessing pipeline than the one the checkpoint was trained
+        # with.
+        if hasattr(processor, "video_processor"):
+            videos_inputs = processor.video_processor(
+                videos=videos, return_tensors=return_tensors
+            )
+        else:
+            raise RuntimeError(
+                "processor has no video_processor attribute - this code path "
+                "requires transformers>=5.2. Upgrade transformers or pass "
+                "videos=None."
+            )
+        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]
+
+    # Discretize normalized proprioception into the <|propri|> prompt slot.
+    if norm_state is not None:
+        norm_state = (
+            norm_state.cpu().numpy()
+            if isinstance(norm_state, torch.Tensor)
+            else norm_state
+        )
+        discretized = (
+            np.digitize(norm_state, bins=np.linspace(-1, 1, state_bins + 1)[:-1]) - 1
+        )
+        discretized = discretized[:, 0, :]
+        if agent_pos_mask is not None:
+            mask = (
+                agent_pos_mask[:, 0, :].cpu().numpy().astype(bool)
+                if isinstance(agent_pos_mask, torch.Tensor)
+                else agent_pos_mask[:, 0, :].astype(bool)
+            )
+        else:
+            mask = np.ones(discretized.shape, dtype=bool)
+        for i in range(len(text)):
+            if "<|propri|>" not in text[i]:
+                continue
+            if state_drop_prob > 0 and random.random() < state_drop_prob:
+                text[i] = text[i].replace("<|propri|>", "")
+            else:
+                state_str = " ".join(map(str, discretized[i, mask[i]]))
+                text[i] = text[i].replace("<|propri|>", state_str)
+
+    # 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):
+                    logger.warning(
+                        "Number of image placeholders (%s) exceeds actual images "
+                        "(%s); skipping remaining placeholder processing",
+                        index + 1,
+                        len(image_grid_thw),
+                    )
+                    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})