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)

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()
 
@@ -202,7 +201,9 @@
                             use_state=True,
                             iterable_shuffle=True,  # rank x worker episode-shuffle stream
                             episode_shuffle_seed=42,
-                            use_image_augmentation=True,  # SR boost (random crop+rescale + color jitter)
+                            # SR boost: random crop+rescale + ColorJitter, applied CPU-side in the
+                            # DROIDLeRobotDataset image augmentor (matches i4's pipeline stage).
+                            use_image_augmentation=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,
@@ -213,6 +214,9 @@
                             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,
                         ),
                     ),
                 ),

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
 
@@ -109,24 +111,30 @@ 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,
+    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.
+
+    Reads ``root`` (a merged/versioned DROID LeRobot root) as a single flat
+    dataset; ``use_success_only=True`` filters to the ``success/`` split."""
+    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,
     )
+    dataset: Dataset = DROIDLeRobotDataset(root=root, **shard_kwargs)
     transform = ActionTransformPipeline(
         tokenizer_config=tokenizer_config,
         cfg_dropout_rate=cfg_dropout_rate,
@@ -135,6 +143,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)