DROID action-policy post-training: lazy LeRobot dataset + Cosmos3-Nano recipe + numerical regression test

cosmos_framework/checkpoint/dcp.py

@@ -71,9 +71,9 @@
 
 from cosmos_framework.checkpoint.base import AbstractCheckpointer
 from cosmos_framework.checkpoint.s3_filesystem import S3StorageReader, S3StorageWriter
-from cosmos_framework.utils.config import CheckpointConfig, JobConfig
 from cosmos_framework.model._base import ImaginaireModel
 from cosmos_framework.utils import callback, distributed, log, misc
+from cosmos_framework.utils.config import CheckpointConfig, JobConfig
 from cosmos_framework.utils.easy_io import easy_io
 from cosmos_framework.utils.generator.rand_state import get_rand_state_dict, set_rand_state_dict
 
@@ -866,6 +866,16 @@ def load(
                         raise ValueError(
                             f"Unexpected keys (found in checkpoint but not in model): {results.unexpected_keys}"
                         )
+                    # Warm start that skipped net_ema (e.g. loading an EMA-only HF export
+                    # with no net_ema.* keys): the EMA shadow would otherwise keep its random
+                    # build-time generation pathway (init_moe is skipped when a checkpoint is
+                    # present). Seed net_ema from the freshly loaded net so the EMA starts equal
+                    # to net ("EMA warm-starts from net") instead of from random weights.
+                    if warm_start and any("net_ema" in skip_key for skip_key in keys_to_skip_loading):
+                        ema_worker = getattr(model, "net_ema_worker", None)
+                        if ema_worker is not None and getattr(model, "net_ema", None) is not None:
+                            ema_worker.copy_to(src_model=model.net, tgt_model=model.net_ema)
+                            log.info("Warm start: re-seeded net_ema from net (net_ema was skipped on load).")
 
                 elif key == "optim":
                     log.info("- Loading the optimizer...")

cosmos_framework/configs/base/defaults/model_config.py

@@ -5,12 +5,12 @@
 
 import attrs
 
-from cosmos_framework.utils.lazy_config import LazyDict
 from cosmos_framework.configs.base.defaults.activation_checkpointing import ActivationCheckpointingConfig
 from cosmos_framework.configs.base.defaults.compile import CompileConfig
 from cosmos_framework.configs.base.defaults.ema import EMAConfig
 from cosmos_framework.configs.base.defaults.parallelism import ParallelismConfig
 from cosmos_framework.configs.base.defaults.reasoner import VLMConfig
+from cosmos_framework.utils.lazy_config import LazyDict
 
 
 @attrs.define(slots=False)
@@ -249,6 +249,14 @@ class OmniMoTModelConfig:
     max_action_dim: int = 32  # maximum dimension of the action space, we need to pad the data to this dimension.
     num_embodiment_domains: int = 32  # number of domains for the domain-aware linear layer
 
+    # Optional GPU-side photometric augmentation applied per-sample to the input video
+    # during training (in _normalize_video_databatch_inplace), as a dict of
+    # torchvision ColorJitter kwargs e.g. {"brightness":0.3,"contrast":0.4,"saturation":0.5,"hue":0.08}.
+    # None disables it. Moving ColorJitter here (GPU, batched) instead of in the CPU
+    # dataloader workers avoids the ~14x per-sample slowdown from float32 rgb<->hsv on
+    # ~1 core/worker (the dataloader's dominant cost for video action SFT).
+    train_color_jitter: dict | None = None
+
     # sound configs
     sound_gen: bool = False  # whether to use sound related parameters and condition/generate sound tokens
     sound_tokenizer: LazyDict | None = None  # Sound tokenizer config (e.g., AVAE)

cosmos_framework/configs/base/experiment/action/posttrain_config/action_policy_droid_nano.py

@@ -21,15 +21,14 @@
 
 from hydra.core.config_store import ConfigStore
 
-from cosmos_framework.utils.lazy_config import LazyCall as L
-from cosmos_framework.utils.lazy_config import LazyDict
-
 from cosmos_framework.configs.base.experiment.sft.models.nano_model_config import NANO_MODEL_CONFIG
+from cosmos_framework.data.generator.action.datasets.action_sft_dataset import get_action_droid_sft_dataset
 from cosmos_framework.data.generator.joint_dataloader import (
     PackingDataLoader,
     RankPartitionedDataLoader,
 )
-from cosmos_framework.data.generator.action.datasets.action_sft_dataset import get_action_droid_sft_dataset
+from cosmos_framework.utils.lazy_config import LazyCall as L
+from cosmos_framework.utils.lazy_config import LazyDict
 
 cs = ConfigStore.instance()
 
@@ -203,16 +202,29 @@
                             iterable_shuffle=True,  # rank x worker episode-shuffle stream
                             episode_shuffle_seed=42,
                             use_image_augmentation=True,  # SR boost (random crop+rescale + color jitter)
+                            # ColorJitter applied CPU-side in the dataloader augmentor, matching
+                            # i4's pipeline stage exactly (model.config.train_color_jitter is None
+                            # below). Slower (~14x rgb<->hsv on 1 core/worker) but eliminates the
+                            # last structural OSS-vs-i4 training-path difference.
+                            apply_color_jitter=True,
                             # keep_ranges_1_0_1.json window filter (drops idle/non-task frames). Off by default;
                             # set use_filter_dict=True + filter_dict_path to enable.
                             use_filter_dict=False,
                             filter_dict_path=None,
                             action_normalization=None,
+                            # Sharded (per-lab) dataset layout. False -> read DROID_ROOT as a
+                            # single flat LeRobot (prior behavior). Override to True (and point
+                            # DROID_ROOT at the <...>_sharded dir, which holds success/<lab>/...)
+                            # to reproduce the internal sharded run's per-shard index build.
+                            sharded=False,
                             viewpoint="concat_view",  # wrist 480p (top) + L/R shoulder 320x180 (bottom)
                             resolution="480",  # 640x360 data @ 480p
                             max_action_dim="${model.config.max_action_dim}",
                             cfg_dropout_rate=0.1,
                             tokenizer_config="${model.config.vlm_config.tokenizer}",
+                            # Match i4 GA (droid_lerobot_8b_ga / MR #9995): format the action
+                            # prompt as JSON via ActionPromptJsonFormatter instead of plain text.
+                            format_prompt_as_json=True,
                         ),
                     ),
                 ),
@@ -243,6 +255,14 @@
 action_policy_droid_nano["model"]["config"]["rectified_flow_training_config"]["loss_scale"] = 10.0
 
 
+# GPU-side photometric augmentation DISABLED. ColorJitter now runs CPU-side in the
+# dataloader (apply_color_jitter=True above), matching i4's exact pipeline stage so the
+# OSS and i4 training forward paths are structurally identical. (To restore the faster
+# GPU-side jitter, set this back to dict(brightness=0.3,contrast=0.4,saturation=0.5,hue=0.08)
+# and flip the dataset's apply_color_jitter back to False.)
+action_policy_droid_nano["model"]["config"]["train_color_jitter"] = None
+
+
 for _item in [action_policy_droid_nano]:
     _name = [k for k, v in globals().items() if v is _item][0]
     cs.store(group="experiment", package="_global_", name=_name, node=_item)

cosmos_framework/data/generator/action/datasets/action_sft_dataset.py

@@ -19,7 +19,9 @@
 
 from torch.utils.data import Dataset, IterableDataset, get_worker_info
 
-from cosmos_framework.data.generator.action.datasets.droid_lerobot_dataset import DROIDLeRobotDataset
+from cosmos_framework.data.generator.action.datasets.droid_lerobot_dataset import (
+    DROIDLeRobotDataset,
+)
 from cosmos_framework.data.generator.action.datasets.libero_lerobot_dataset import LIBEROLeRobotDataset
 from cosmos_framework.data.generator.action.transforms import ActionTransformPipeline
 
@@ -99,6 +101,7 @@ def get_action_droid_sft_dataset(
     action_normalization: str | None = None,
     viewpoint: str = "concat_view",
     use_image_augmentation: bool = False,
+    apply_color_jitter: bool = True,
     use_filter_dict: bool = False,
     filter_dict_path: str | None = None,
     resolution: str | int = "256",
@@ -109,24 +112,46 @@ def get_action_droid_sft_dataset(
     append_duration_fps_timestamps: bool = True,
     append_resolution_info: bool = True,
     append_idle_frames: bool = False,
+    format_prompt_as_json: bool = False,
     iterable_shuffle: bool = False,
     episode_shuffle_seed: int = 42,
+    sharded: bool = False,
+    lerobot_roots: list[str] | None = None,
+    use_success_only: bool = True,
 ) -> Dataset:
     """Build the DROID action SFT dataset: ``action_space='joint_pos'`` (8D) +
-    ``use_state`` (raw/un-normalized), concat_view, chunk_length 32."""
-    dataset = DROIDLeRobotDataset(
-        root=root,
+    ``use_state`` (raw/un-normalized), concat_view, chunk_length 32.
+
+    ``sharded=True`` consumes the per-lab sharded layout (``<root>/success/<lab>``)
+    via :class:`ShardedDROIDLeRobotDataset` — one ``DROIDLeRobotDataset`` per lab
+    concatenated into one flat index — reproducing the internal sharded run's
+    per-shard index construction. ``sharded=False`` (default) reads ``root`` as a
+    single flat LeRobot dataset (the prior behavior). ``lerobot_roots`` optionally
+    pins the shard sub-paths (relative to ``root``); otherwise they are
+    auto-discovered."""
+    if isinstance(sharded, str):
+        sharded = sharded.strip().lower() in ("1", "true", "yes", "on")
+    shard_kwargs = dict(
         fps=fps,
         chunk_length=chunk_length,
         viewpoint=viewpoint,
         action_space=action_space,
         mode=mode,
         use_state=use_state,
         action_normalization=action_normalization,
-        use_image_augmentation=use_image_augmentation,
+        use_image_augmentation=use_image_augmentation,  # i4: bundles random-crop+resize+ColorJitter
         use_filter_dict=use_filter_dict,
         filter_dict_path=filter_dict_path,
+        use_success_only=use_success_only,
     )
+    if sharded:
+        raise NotImplementedError(
+            "The ported i4 DROIDLeRobotDataset reads the merged/versioned root natively "
+            "(basename(root) must be a LEROBOT_ROOTS version key, e.g. "
+            "'droid_plus_lerobot_640x360_20260412'); use_success_only=True filters to success/. "
+            "Per-lab ShardedDROIDLeRobotDataset is not part of the i4 dataset."
+        )
+    dataset: Dataset = DROIDLeRobotDataset(root=root, **shard_kwargs)
     transform = ActionTransformPipeline(
         tokenizer_config=tokenizer_config,
         cfg_dropout_rate=cfg_dropout_rate,
@@ -135,6 +160,7 @@ def get_action_droid_sft_dataset(
         append_duration_fps_timestamps=append_duration_fps_timestamps,
         append_resolution_info=append_resolution_info,
         append_idle_frames=append_idle_frames,
+        format_prompt_as_json=format_prompt_as_json,
     )
     sft = ActionSFTDataset(dataset, transform, resolution)
     if iterable_shuffle:

cosmos_framework/data/generator/action/datasets/base_dataset.py

@@ -74,18 +74,13 @@ def __init__(
         # LeRobot v2.x stores task text in a "task" column; v3.0 stores it as the
         # (unnamed) DataFrame index and keeps only "task_index" as a column.
         task_texts = tasks_df["task"] if "task" in tasks_df.columns else tasks_df.index
-        self._tasks = {
-            int(task_index): str(task)
-            for task, task_index in zip(task_texts, tasks_df["task_index"])
-        }
-        self._rows = sorted(
-            (
-                row
-                for path in sorted((self._root / "data").glob("chunk-*/file-*.parquet"))
-                for row in pq.read_table(path).to_pylist()
-            ),
-            key=lambda row: int(row["index"]),
-        )
+        self._tasks = {int(task_index): str(task) for task, task_index in zip(task_texts, tasks_df["task_index"])}
+        # ``self._rows`` (the flat, index-sorted list of every frame dict) is built
+        # lazily on first access — see the ``_rows`` property. Materializing all
+        # ~18M frames as Python dicts plus a full sort costs ~13 min and tens of GB;
+        # subclasses that build their own compact index (e.g. DROIDLeRobotDataset)
+        # never touch it, so they must not pay for it at construction.
+        self._rows_cache: list[dict[str, Any]] | None = None
 
     @property
     def fps(self) -> float:
@@ -140,8 +135,7 @@ def _stats_path(cls) -> Path:
     def load_action_stats(cls) -> dict[str, torch.Tensor]:
         """Return action normalization stats for this dataset as torch tensors."""
         return {
-            key: torch.from_numpy(value).float()
-            for key, value in load_action_stats(str(cls._stats_path())).items()
+            key: torch.from_numpy(value).float() for key, value in load_action_stats(str(cls._stats_path())).items()
         }
 
     @abstractmethod
@@ -218,5 +212,25 @@ def _build_result(
             **extras,
         }
 
+    @property
+    def _rows(self) -> list[dict[str, Any]]:
+        """Flat, index-sorted list of every frame dict, built lazily on first access.
+
+        Only datasets that don't build their own compact index (bridge / agibot /
+        robomind) touch this; for them it materializes once and caches. Datasets with
+        a bespoke index (e.g. DROIDLeRobotDataset) never read it, so they skip the
+        ~13 min / tens-of-GB construction entirely.
+        """
+        if self._rows_cache is None:
+            self._rows_cache = sorted(
+                (
+                    row
+                    for path in sorted((self._root / "data").glob("chunk-*/file-*.parquet"))
+                    for row in pq.read_table(path).to_pylist()
+                ),
+                key=lambda row: int(row["index"]),
+            )
+        return self._rows_cache
+
     def __len__(self) -> int:
         return max(0, (len(self._rows) - self._chunk_length + self._sample_stride - 1) // self._sample_stride)