Update Wall-X to 1.1.0 (#104)

wall_x/data/__init__.py

@@ -0,0 +1,31 @@
+"""Public data facade for backend selection and dataset construction."""
+
+from wall_x.data import backends  # noqa: F401  (side-effect registration)
+from wall_x.data._bundle import DataBundle
+from wall_x.data._protocol import BuildContext, DatasetBackend
+from wall_x.data._registry import (
+    MissingOperationError,
+    available_backends,
+    backend_for,
+    build_data,
+    data_backend,
+    has_data_backend,
+    register,
+    register_module,
+    resolve_dataset_type,
+)
+
+__all__ = [
+    "DataBundle",
+    "BuildContext",
+    "DatasetBackend",
+    "MissingOperationError",
+    "data_backend",
+    "available_backends",
+    "backend_for",
+    "build_data",
+    "has_data_backend",
+    "register",
+    "register_module",
+    "resolve_dataset_type",
+]

wall_x/data/_bundle.py

@@ -0,0 +1,39 @@
+"""Dataset bundle returned by data backends."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+from typing import Any, Callable, Iterable, Optional
+
+
+def _noop_set_epoch(_epoch: int) -> None:
+    """Default ``set_epoch`` for backends with epoch-agnostic shuffling."""
+    return None
+
+
+@dataclass
+class DataBundle:
+    """Container returned by every backend ``build()``.
+
+    Attributes:
+        dataset: backend-private; trainer code should treat it as opaque.
+        train_loader: anything iterable that yields training batches.
+        val_loader: val iterable or None if the backend does not split val.
+        train_iters: one-epoch step count for the train loader. Backends
+            with dynamic resizing should set this to a stable snapshot
+            and expose the live value separately on ``dataset``.
+        val_iters: one-epoch step count for the val loader; 0 if no val.
+        set_epoch: per-epoch seed hook. Called before each epoch by
+            the trainer. Backends that don't need per-epoch reshuffle
+            should use ``_noop_set_epoch``.
+    """
+
+    dataset: Any
+    train_loader: Iterable
+    val_loader: Optional[Iterable] = None
+    train_iters: int = 0
+    val_iters: int = 0
+    set_epoch: Callable[[int], None] = field(default=_noop_set_epoch)
+
+
+__all__ = ["DataBundle"]

wall_x/data/_protocol.py

@@ -0,0 +1,37 @@
+"""Shared data backend protocol definitions."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Any, Optional, Protocol, runtime_checkable
+
+from wall_x.data._bundle import DataBundle
+
+
+@dataclass
+class BuildContext:
+    """Shared runtime state passed to every backend ``build()``.
+
+    Fields are Optional so backends not using a particular piece can
+    simply leave it ``None``. The trainer populates whatever it has.
+    """
+
+    rank: int = 0
+    world_size: int = 1
+    tokenizer: Optional[Any] = None
+    processor: Optional[Any] = None
+    tokenizer_mixin: Optional[Any] = None
+    normalizer_action: Optional[Any] = None
+    normalizer_propri: Optional[Any] = None
+    model_config: Optional[Any] = None
+    resume_state: Optional[dict] = None
+
+
+@runtime_checkable
+class DatasetBackend(Protocol):
+    """Callable every backend registers under its ``dataset_type`` name."""
+
+    def __call__(self, cfg: Any, ctx: BuildContext) -> DataBundle: ...
+
+
+__all__ = ["BuildContext", "DatasetBackend"]

wall_x/data/_registry.py

@@ -0,0 +1,263 @@
+"""Data backend registry + multi-verb dispatch.
+
+Each backend registers an entire **module** (not just a single function).
+Consumers obtain a backend via :func:`backend_for` and then call any
+operation the backend supports (``backend.build`` or optional helper
+operations published by that backend). Operations a backend does *not*
+implement raise
+:class:`MissingOperationError` with a list of backends that do.
+
+Three failure shapes, each pointing at the right next step:
+
+- **Unknown name** (typo / never registered): :class:`KeyError`
+  ``Unknown backend 'foo_bar'. Known: [...]``
+- **Known but failed to import** (optional dependency missing):
+  :class:`RuntimeError`
+  ``Backend 'example' failed to import: <ImportError>. Install its
+  dependency or switch to one of [...]``
+- **Backend exists but operation missing**: :class:`MissingOperationError`
+  ``Backend 'example' does not implement 'make_processor'.
+  Supported by: ['other_backend']``
+"""
+
+from __future__ import annotations
+
+import logging
+from types import ModuleType
+from typing import Any, Dict
+
+from wall_x.data._bundle import DataBundle
+from wall_x.data._protocol import BuildContext
+
+logger = logging.getLogger(__name__)
+
+_BACKENDS: Dict[str, ModuleType] = {}
+_import_errors: Dict[str, BaseException] = {}
+
+# The single, process-wide active backend. Set exactly once at config load
+# time by :func:`_set_data_backend`; all consumer code reads it via
+# :func:`data_backend`. Keeping this as module-level state (rather than
+# threading cfg through every callsite) is the whole point of the design.
+_DATA_BACKEND: str | None = None
+
+
+# --- Registration --------------------------------------------------------
+
+
+def register_module(dataset_type: str, module: ModuleType) -> None:
+    """Register a backend module under ``dataset_type``.
+
+    The module must expose a ``build(cfg, ctx) -> DataBundle`` callable.
+    Other optional operations are
+    discovered lazily via ``getattr`` when ``backend_for(name).<op>`` is
+    accessed; backends only declare what they support.
+    """
+    if not hasattr(module, "build"):
+        raise TypeError(
+            f"Backend module for {dataset_type!r} must expose a "
+            f"``build(cfg, ctx) -> DataBundle`` callable; got {module!r}."
+        )
+    if dataset_type in _BACKENDS:
+        logger.warning("overwriting existing backend registration for %r", dataset_type)
+    _BACKENDS[dataset_type] = module
+
+
+def register(dataset_type: str, build_callable) -> None:
+    """Legacy single-callable registration. Wraps ``build_callable`` in a
+    minimal module so the new dispatch path still works.
+
+    New backends should prefer :func:`register_module` so they can publish
+    multiple operations.
+    """
+    shim = ModuleType(f"_legacy_backend_shim_{dataset_type}")
+    shim.build = build_callable  # type: ignore[attr-defined]
+    register_module(dataset_type, shim)
+
+
+def record_import_error(dataset_type: str, error: BaseException) -> None:
+    """Stash the exception that prevented a backend from registering."""
+    _import_errors[dataset_type] = error
+
+
+def available_backends() -> list[str]:
+    """Return the names of currently-registered backends."""
+    return sorted(_BACKENDS)
+
+
+# --- Lookup --------------------------------------------------------------
+
+
+class MissingOperationError(NotImplementedError):
+    """Raised when a backend module does not implement a requested op."""
+
+    def __init__(self, backend_name: str, op_name: str) -> None:
+        impls = [n for n, mod in _BACKENDS.items() if hasattr(mod, op_name)]
+        if impls:
+            hint = f"Supported by: {impls}."
+        else:
+            hint = (
+                f"No registered backend implements {op_name!r} - check the "
+                f"spelling or add it to the backend module."
+            )
+        super().__init__(
+            f"Backend {backend_name!r} does not implement {op_name!r}. {hint}"
+        )
+        self.backend_name = backend_name
+        self.op_name = op_name
+
+
+class _BackendProxy:
+    """Thin wrapper that forwards ``proxy.<op>(...)`` to the backend module.
+
+    Wrapping (instead of returning the module directly) lets us emit
+    ``MissingOperationError`` with a useful "supported by" hint instead
+    of plain ``AttributeError``.
+    """
+
+    __slots__ = ("_name", "_module")
+
+    def __init__(self, name: str, module: ModuleType) -> None:
+        self._name = name
+        self._module = module
+
+    def __repr__(self) -> str:
+        return f"<BackendProxy name={self._name!r}>"
+
+    def __getattr__(self, op: str):
+        attr = getattr(self._module, op, None)
+        if attr is None:
+            raise MissingOperationError(self._name, op)
+        return attr
+
+    def supports(self, op: str) -> bool:
+        """Cheap predicate: does this backend implement ``op``?"""
+        return hasattr(self._module, op)
+
+
+def backend_for(name: str) -> _BackendProxy:
+    """Look up a backend by ``dataset_type`` name.
+
+    Raises :class:`RuntimeError` (with chained ImportError) if the
+    backend is known but failed to import; :class:`KeyError` if the
+    name was never registered.
+    """
+    if name in _BACKENDS:
+        return _BackendProxy(name, _BACKENDS[name])
+    if name in _import_errors:
+        err = _import_errors[name]
+        raise RuntimeError(
+            f"Backend {name!r} failed to import: {err}. "
+            f"Install its dependency or switch to one of "
+            f"{available_backends()}."
+        ) from err
+    raise KeyError(f"Unknown backend {name!r}. Known: {available_backends()}.")
+
+
+# --- Convenience helpers -------------------------------------------------
+
+
+def resolve_dataset_type(cfg_or_yaml: Any, default: str = "lerobot") -> str:
+    """Pull ``dataset_type`` out of either a typed TrainConfig or a raw yaml dict.
+
+    Lookup order:
+    1. ``cfg.data.dataset_type`` (typed TrainConfig)
+    2. ``yaml_dict["dataset_type"]`` (legacy flat yaml)
+    3. ``default``
+    """
+    data = getattr(cfg_or_yaml, "data", None)
+    if data is not None:
+        v = getattr(data, "dataset_type", None)
+        if v:
+            return v
+    if isinstance(cfg_or_yaml, dict):
+        v = cfg_or_yaml.get("dataset_type")
+        if v:
+            return v
+    return default
+
+
+def build_data(cfg: Any, ctx: BuildContext) -> DataBundle:
+    """Dispatch ``backend.build(cfg, ctx)`` to the backend named by cfg.
+
+    ``cfg`` is expected to be a typed :class:`TrainConfig` - that's what
+    ``wall_x.config.load_config`` returns and what every trainer entry
+    point passes in. Raw yaml dicts are not supported here on purpose:
+    the typed schema is the contract that gives backends their
+    ``cfg.data.dataset_type`` access. If you have a raw dict, use
+    :func:`resolve_dataset_type` + :func:`backend_for` directly.
+    """
+    if not hasattr(cfg, "data") or not hasattr(cfg.data, "dataset_type"):
+        raise TypeError(
+            f"build_data() expects a typed TrainConfig, got {type(cfg).__name__}. "
+            f"Load via wall_x.config.load_config() instead of passing a raw dict."
+        )
+    return backend_for(cfg.data.dataset_type).build(cfg, ctx)
+
+
+# --- Active backend (process-global) ------------------------------------
+
+
+def _set_data_backend(name: str) -> None:
+    """Internal - only config loaders should call this.
+
+    Strict semantics: first call wins; same-value re-set is a no-op;
+    different-value re-set raises. The error is loud on purpose - it means
+    two config loaders disagreed about which backend to use, which is
+    almost always a bug (e.g. business code calling this directly, or two
+    yamls being loaded in the same process).
+    """
+    global _DATA_BACKEND
+    if _DATA_BACKEND is None:
+        _DATA_BACKEND = name
+        return
+    if _DATA_BACKEND == name:
+        return  # idempotent on same value
+    raise RuntimeError(
+        f"Active backend already set to {_DATA_BACKEND!r}, refusing to "
+        f"overwrite with {name!r}. This usually means a config loader was "
+        f"called twice with different dataset_type, or business code called "
+        f"_set_data_backend directly. Use _reset_data_backend() in tests "
+        f"if you need to switch."
+    )
+
+
+def _reset_data_backend() -> None:
+    """Internal - clear the active backend. Tests only."""
+    global _DATA_BACKEND
+    _DATA_BACKEND = None
+
+
+def has_data_backend() -> bool:
+    """Whether a backend has been registered for this process."""
+    return _DATA_BACKEND is not None
+
+
+def data_backend() -> _BackendProxy:
+    """Return the active backend proxy. Raises if no config has been loaded.
+
+    Direct construction of ``TrainConfig(...)`` does **not** register a
+    backend - that's intentional. Business code must go through
+    :func:`wall_x.config.load_config`, which registers the backend before
+    returning.
+    """
+    if _DATA_BACKEND is None:
+        raise RuntimeError(
+            "No active data backend. Load a TrainConfig via "
+            "wall_x.config.load_config() before accessing backend verbs. "
+            "Direct TrainConfig() construction does not register a backend."
+        )
+    return backend_for(_DATA_BACKEND)
+
+
+__all__ = [
+    "register",
+    "register_module",
+    "record_import_error",
+    "available_backends",
+    "backend_for",
+    "build_data",
+    "resolve_dataset_type",
+    "MissingOperationError",
+    "data_backend",
+    "has_data_backend",
+]

wall_x/data/backends/__init__.py

@@ -0,0 +1,16 @@
+"""Data backend registration.
+
+Importing this package registers the default shipped backends. Additional
+backends are loaded by plugin modules so their names and dependencies do not
+have to appear in the default import path.
+"""
+
+import importlib
+
+from wall_x.data._registry import record_import_error
+
+for _name in ("lerobot",):
+    try:
+        importlib.import_module(f"wall_x.data.backends.{_name}")
+    except ImportError as _e:
+        record_import_error(_name, _e)

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/config.py

@@ -1,126 +0,0 @@
-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",
-    "physical-intelligence/libero",
-    "lerobot/aloha_mobile_cabinet",
-]
-
-# 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

@@ -1,715 +0,0 @@
-"""
-LeRobot Dataset Loader - Distributed Version
-"""
-
-import numpy as np
-import torch
-from torch.utils.data import DistributedSampler, random_split
-from lerobot.datasets.lerobot_dataset import LeRobotDataset, LeRobotDatasetMetadata
-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
-from .utils import KEY_MAPPINGS
-
-T_co = TypeVar("T_co", covariant=True)
-
-
-# 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,
-        normalizer_action,
-        normalizer_propri,
-        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.normalizer_action = (normalizer_action,)
-        self.normalizer_propri = normalizer_propri
-        # self.norm_stats = norm_stats
-        self.lerobot_config = lerobot_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 = KEY_MAPPINGS[self.hf_dataset.meta.repo_id]["camera"]
-        self._state_key_mapping = KEY_MAPPINGS[self.hf_dataset.meta.repo_id]
-        self._action_key_mapping = KEY_MAPPINGS[self.hf_dataset.meta.repo_id]
-
-    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"]]
-        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,
-        )
-        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.normalizer_action,
-                self.normalizer_propri,
-                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 = {}
-
-    def __init__(
-        self,
-        config,
-        dataload_config,
-        normalizer_action,
-        normalizer_propri,
-        lerobot_config,
-    ):
-        self.config = config
-        self.dataload_config = dataload_config
-
-        self.normalizer_action = normalizer_action[0]
-        self.normalizer_propri = normalizer_propri
-        self.lerobot_config = lerobot_config
-
-        self.use_fast_tokenizer = self.config.get("use_fast_tokenizer", False)
-        self.dataset_name = self.config["data"]["lerobot_config"].get("repo_id", "")
-        self.dataset_name = [self.dataset_name] * self.config["batch_size_per_gpu"]
-        self.load_processor()
-
-    def load_processor(self):
-        processor_path = self.config["pretrained_wallx_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
-
-    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()
-                # print("agent_pos_mask",agent_pos_mask.shape)
-                agent_pos.nan_to_num_(nan=0.0)
-
-                # 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,
-                #     )
-                agent_pos = self.normalizer_propri.normalize_data(
-                    agent_pos, self.dataset_name
-                )
-                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)
-
-                # if action.shape[-1] != 20:
-                #     action = torch.cat(
-                #         [
-                #             action,
-                #             torch.zeros(
-                #                 action.shape[0], action.shape[1], 20 - action.shape[-1]
-                #             ),
-                #         ],
-                #         dim=-1,
-                #     )
-                #     dof_mask = torch.cat(
-                #         [
-                #             dof_mask,
-                #             torch.zeros(
-                #                 dof_mask.shape[0],
-                #                 dof_mask.shape[1],
-                #                 20 - dof_mask.shape[-1],
-                #             ),
-                #         ],
-                #         dim=-1,
-                #     )
-                action = self.normalizer_action.normalize_data(
-                    action, self.dataset_name
-                )
-                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,
-            [self.lerobot_config["repo_id"]] * 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"] = [self.lerobot_config["repo_id"]] * inputs[
-            "action_chunk"
-        ].shape[0]
-
-        return inputs
-
-
-def load_lerobot_data(
-    config,
-    lerobot_config,
-    normalizer_action,
-    normalizer_propri,
-    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"])
-
-    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"
-    # norm_stats = load_norm_stats(norm_stats_path, repo_id)
-
-    delta_timestamps = {
-        # action chunk
-        KEY_MAPPINGS[repo_id]["action"]: [
-            t / dataset_fps
-            for t in range(dataload_config.get("action_horizon", 33) - 1)
-        ],
-    }
-    batch_size = config.get("batch_size_per_gpu", 8)
-    episodes = np.arange(episodes_num).tolist()
-
-    train_test_split = dataload_config.get("train_test_split", 0.95)
-    train_episodes = episodes[: int(episodes_num * train_test_split)]
-    test_episodes = episodes[int(episodes_num * train_test_split) :]
-
-    train_dataset = LeRobotDataset(
-        repo_id,
-        root=root,
-        episodes=train_episodes,
-        delta_timestamps=delta_timestamps,
-        video_backend="pyav",
-    )
-
-    if rank == 0:
-        print(f"Selected train episodes: {train_dataset.episodes}")
-        print(f"Number of train episodes selected: {train_dataset.num_episodes}")
-        print(f"Number of train frames selected: {train_dataset.num_frames}")
-        print(f"Selected test episodes: {test_episodes}")
-
-    dataset = PreprocessedDataset(
-        train_dataset,
-        config,
-        dataload_config,
-        normalizer_action,
-        normalizer_propri,
-        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:
-        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,
-        normalizer_action,
-        normalizer_propri,
-        lerobot_config,
-        seed=42,
-    ):
-        super().__init__(
-            dataset,
-            config,
-            dataload_config,
-            normalizer_action,
-            normalizer_propri,
-            lerobot_config,
-            seed=seed,
-            rank=0,
-            world_size=1,
-            test_only=True,
-        )
-
-    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.normalizer_action,
-                self.normalizer_propri,
-                self.lerobot_config,
-            ),
-        )
-
-        return dataloader
-
-
-def load_test_dataset(
-    config,
-    lerobot_config,
-    normalizer_action,
-    normalizer_propri,
-    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)
-
-    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
-    dataload_config = get_data_configs(config["data"])
-
-    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"
-    # norm_stats = load_norm_stats(norm_stats_path, repo_id)
-
-    delta_timestamps = {
-        # action chunk
-        KEY_MAPPINGS[repo_id]["action"]: [
-            t / dataset_fps
-            for t in range(dataload_config.get("action_horizon", 33) - 1)
-        ],
-    }
-
-    dataset = LeRobotDataset(
-        repo_id,
-        episodes=[episode],
-        delta_timestamps=delta_timestamps,
-        video_backend="pyav",
-        root=root,
-    )
-
-    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,
-        normalizer_action,
-        normalizer_propri,
-        lerobot_config,
-        seed=seed,
-    )
-
-    return dataset