[cosmos3] read DROID through upstream's own data pipeline

.pre-commit-config.yaml

@@ -1,3 +1,10 @@
+# Third-party source copied in verbatim, and verbatim is the point: the DROID data pipeline under
+# examples/worldmodel/cosmos3/vendor/ is byte-identical to nvidia/cosmos-framework, so that checking this
+# example against a newer upstream is one `diff -r` with no formatting noise in the way. autoflake, isort
+# and black would each rewrite those files and end that property, so they are excluded here rather than
+# being allowed to reformat code this repository does not own.
+exclude: '^examples/worldmodel/cosmos3/vendor/'
+
 repos:
 
   - repo: https://github.com/PyCQA/autoflake

examples/worldmodel/cosmos3/.gitignore

@@ -0,0 +1,2 @@
+# Training output written by the example entry point (--out defaults here).
+runs/

examples/worldmodel/cosmos3/tests/test_training_entry.py

@@ -232,11 +232,18 @@ def prepare_dataloader(self, dataset, batch_size, distributed_sampler_cls, **kwa
             return torch.utils.data.DataLoader(dataset, batch_size=batch_size, sampler=sampler)
 
     monkeypatch.setattr(entry, "HybridParallelPlugin", Plugin)
+    # The entry defaults to the vendored reader, so the stub stands in for both names; the default
+    # path would otherwise reach for the real upstream pipeline and its dependencies.
     monkeypatch.setattr(
         entry,
         "DroidActionPolicyDataset",
         lambda *a, **kw: [{"goal": 1.0, "sample_id": "a"}, {"goal": 3.0, "sample_id": "b"}],
     )
+    monkeypatch.setattr(
+        entry,
+        "UpstreamDroidDataset",
+        lambda *a, **kw: [{"goal": 1.0, "sample_id": "a"}, {"goal": 3.0, "sample_id": "b"}],
+    )
 
     class Model(torch.nn.Module):
         config = None

examples/worldmodel/cosmos3/tests/test_upstream_reader_filter.py

@@ -0,0 +1,95 @@
+"""The vendored reader only applies the curated window filter when it is told the filter's URI prefix.
+
+``keep_ranges_1_0_1.json`` is keyed by composite ``gs://`` URIs, and the reader refuses to guess the prefix in
+front of the episode id rather than select the wrong episodes silently. The adapter used to hard-code an empty
+prefix, which turned ``--keep-ranges`` on that file into a ``ValueError`` on the default reader while the same
+flag worked on the other one. These tests pin the forwarding, not the reader's own filtering, which upstream
+tests: the reader is stubbed so this runs without the dataset or the framework's dependencies.
+
+Measured against the real release on 2026-11-14: both readers index 16,338,993 windows unfiltered and
+12,157,927 with the filter (74.41%), the share upstream's recipe documents for the released model.
+
+Runs on CPU.
+"""
+
+import sys
+import types
+from pathlib import Path
+
+import pytest
+
+sys.path.insert(0, str(Path(__file__).resolve().parent))
+sys.path.insert(0, str(Path(__file__).resolve().parent.parent))
+sys.path.insert(0, str(Path(__file__).resolve().parents[4]))
+
+ROOT_NAME = "droid_plus_lerobot_640x360_20260412"
+
+
+class _Reader:
+    """Records the keyword arguments the adapter forwards, and nothing else."""
+
+    last: dict = {}
+
+    def __init__(self, **kwargs):
+        type(self).last = kwargs
+        # (dataset index, first row, windows in the span, episode id), the shape upstream appends.
+        self._episode_records = [(0, 0, 8, 0)]
+        self._episode_cum_ends = [8]
+
+    def __len__(self):
+        return 8
+
+
+@pytest.fixture
+def adapter(monkeypatch):
+    """The adapter with upstream's reader replaced, so no framework dependency is needed."""
+    module = types.ModuleType("cosmos_framework.data.generator.action.datasets.droid_lerobot_dataset")
+    module.DROIDLeRobotDataset = _Reader
+    for name in (
+        "cosmos_framework",
+        "cosmos_framework.data",
+        "cosmos_framework.data.generator",
+        "cosmos_framework.data.generator.action",
+        "cosmos_framework.data.generator.action.datasets",
+    ):
+        monkeypatch.setitem(sys.modules, name, types.ModuleType(name))
+    monkeypatch.setitem(sys.modules, module.__name__, module)
+    from examples.worldmodel.cosmos3.upstream_dataset import UpstreamDroidDataset
+
+    return UpstreamDroidDataset
+
+
+def test_an_explicit_prefix_reaches_the_reader(adapter, tmp_path, monkeypatch):
+    monkeypatch.delenv("DROID_FILTER_SOURCE_PREFIX", raising=False)
+    root = tmp_path / ROOT_NAME
+    root.mkdir()
+    adapter(root, keep_ranges="keep.json", filter_dict_source_prefix="gs://bucket/data")
+    assert _Reader.last["use_filter_dict"] is True
+    assert _Reader.last["filter_dict_path"] == "keep.json"
+    assert _Reader.last["filter_dict_source_prefix"] == "gs://bucket/data"
+
+
+def test_the_environment_supplies_the_prefix_when_the_caller_does_not(adapter, tmp_path, monkeypatch):
+    """The same variable the other reader honours, so one setting covers both."""
+    monkeypatch.setenv("DROID_FILTER_SOURCE_PREFIX", "gs://bucket/data")
+    root = tmp_path / ROOT_NAME
+    root.mkdir()
+    adapter(root, keep_ranges="keep.json")
+    assert _Reader.last["filter_dict_source_prefix"] == "gs://bucket/data"
+
+
+def test_an_explicit_prefix_wins_over_the_environment(adapter, tmp_path, monkeypatch):
+    monkeypatch.setenv("DROID_FILTER_SOURCE_PREFIX", "gs://wrong/data")
+    root = tmp_path / ROOT_NAME
+    root.mkdir()
+    adapter(root, keep_ranges="keep.json", filter_dict_source_prefix="gs://bucket/data")
+    assert _Reader.last["filter_dict_source_prefix"] == "gs://bucket/data"
+
+
+def test_no_filter_leaves_the_reader_unfiltered(adapter, tmp_path, monkeypatch):
+    monkeypatch.setenv("DROID_FILTER_SOURCE_PREFIX", "gs://bucket/data")
+    root = tmp_path / ROOT_NAME
+    root.mkdir()
+    adapter(root)
+    assert _Reader.last["use_filter_dict"] is False
+    assert _Reader.last["filter_dict_path"] is None

examples/worldmodel/cosmos3/tests/test_upstream_reader_index.py

@@ -0,0 +1,108 @@
+"""``episode_id`` does not identify a record in the vendored reader's index, and two things make it so.
+
+The reader's ``_episode_records`` holds one entry per *span*, not per episode. With the curated window
+filter on, an episode contributes one span per kept range and its id repeats across them; and the id counts
+from zero again in every registered root, so on a mixed success-and-failure run the same id names two
+different episodes. A table keyed by the id alone collapses both cases.
+
+That is quiet in the worst way: the samples themselves stay correct -- the reader resolves them -- while
+every sample's reported episode and frame number comes out wrong, which is what a replay pack, a held-out
+manifest and a sample id are all built from. The reader is stubbed here so the index logic is tested on its
+own, without the dataset or the framework's dependencies.
+
+Runs on CPU.
+"""
+
+import sys
+import types
+from pathlib import Path
+
+import pytest
+
+sys.path.insert(0, str(Path(__file__).resolve().parent))
+sys.path.insert(0, str(Path(__file__).resolve().parent.parent))
+sys.path.insert(0, str(Path(__file__).resolve().parents[4]))
+
+ROOT_NAME = "droid_plus_lerobot_640x360_20260412"
+
+# One episode, id 5, split into two kept ranges by the window filter, in dataset 0; then the same id
+# again in dataset 1, which is what a second registered root gives. Rows are the dataset rows each span
+# starts at, so episode 5 of dataset 0 owns rows 100.. and episode 5 of dataset 1 owns rows 40...
+RECORDS = [
+    (0, 100, 3, 5),
+    (0, 120, 2, 5),
+    (1, 40, 2, 5),
+]
+CUM_ENDS = [3, 5, 7]
+
+
+class _Reader:
+    """The two attributes the adapter reads from upstream's index, and its length."""
+
+    def __init__(self, **kwargs):
+        self._episode_records = list(RECORDS)
+        self._episode_cum_ends = list(CUM_ENDS)
+
+    def __len__(self):
+        return CUM_ENDS[-1]
+
+
+@pytest.fixture
+def dataset(monkeypatch, tmp_path):
+    module = types.ModuleType("cosmos_framework.data.generator.action.datasets.droid_lerobot_dataset")
+    module.DROIDLeRobotDataset = _Reader
+    for name in (
+        "cosmos_framework",
+        "cosmos_framework.data",
+        "cosmos_framework.data.generator",
+        "cosmos_framework.data.generator.action",
+        "cosmos_framework.data.generator.action.datasets",
+    ):
+        monkeypatch.setitem(sys.modules, name, types.ModuleType(name))
+    monkeypatch.setitem(sys.modules, module.__name__, module)
+    monkeypatch.setenv("DROID_FILTER_SOURCE_PREFIX", "gs://bucket/data")
+    from examples.worldmodel.cosmos3.upstream_dataset import UpstreamDroidDataset
+
+    root = tmp_path / ROOT_NAME
+    root.mkdir()
+    return UpstreamDroidDataset(root, chunk_length=32)
+
+
+def test_every_span_of_one_episode_is_found(dataset):
+    """Keyed by the episode id alone, only the last of the three spans survived the table."""
+    assert [dataset._locate(index) for index in range(7)] == [
+        (0, 5, 0),
+        (0, 5, 1),
+        (0, 5, 2),
+        (0, 5, 20),
+        (0, 5, 21),
+        (1, 5, 0),
+        (1, 5, 1),
+    ]
+
+
+def test_the_frame_number_counts_from_the_episode_not_the_span(dataset):
+    """Counted from the span, the second kept range would restart at zero and collide with the first."""
+    _, _, first_range = dataset._locate(0)
+    _, _, second_range = dataset._locate(3)
+    assert (first_range, second_range) == (0, 20)
+
+
+def test_the_same_episode_id_in_two_roots_stays_two_episodes(dataset):
+    assert dataset._locate(2)[:2] == (0, 5)
+    assert dataset._locate(5)[:2] == (1, 5)
+
+
+def test_no_two_windows_share_an_identity(dataset):
+    ids = [dataset.sample_id(index) for index in range(len(dataset))]
+    assert len(set(ids)) == len(ids)
+    assert ids[0] == "success:dataset_0:episode_5:0:32"
+    assert ids[3] == "success:dataset_0:episode_5:20:32"
+    assert ids[5] == "success:dataset_1:episode_5:0:32"
+
+
+def test_an_index_outside_the_reader_is_refused(dataset):
+    """Reported against the last span instead, it would name a window that was never drawn."""
+    for index in (-1, 7, 99):
+        with pytest.raises(IndexError):
+            dataset._locate(index)

examples/worldmodel/cosmos3/train_sft.py

@@ -19,6 +19,29 @@
 from importlib.metadata import PackageNotFoundError, version
 from pathlib import Path
 
+# The DROID pipeline is upstream's own code, copied byte for byte under vendor/ (see vendor/README.md).
+# Its files import themselves as ``cosmos_framework.*``, so the copy is laid out under that package name
+# and the directory holding it goes on the path here.
+#
+# Here and nowhere else, and appended rather than inserted: on a host that also has the real
+# cosmos-framework installed -- the comparison machine does -- an installed package must keep winning, so
+# that the reference side never silently runs this copy. Which one actually loaded is printed below rather
+# than assumed, because the whole point of the copy is that it is the upstream code, and a run that
+# quietly used the other one would be comparing something else.
+_VENDOR = Path(__file__).resolve().parent / "vendor"
+if _VENDOR.is_dir() and str(_VENDOR) not in sys.path:
+    sys.path.append(str(_VENDOR))
+
+
+def report_pipeline_source() -> str:
+    """Where the DROID pipeline actually came from in this process."""
+    import cosmos_framework
+
+    origin = getattr(cosmos_framework, "__file__", None) or "<namespace package>"
+    vendored = str(_VENDOR) in str(origin)
+    return f"{origin} ({'vendored copy' if vendored else 'installed upstream'})"
+
+
 import torch
 import torch.distributed as dist
 
@@ -43,6 +66,7 @@
     resolve_compile_backend,
     shutdown_dataloader,
 )
+from examples.worldmodel.cosmos3.upstream_dataset import UpstreamDroidDataset
 from examples.worldmodel.cosmos3_harness.optim import build_adamw
 from examples.worldmodel.cosmos3_harness.training_metrics import (
     AnomalyDetector,
@@ -116,6 +140,18 @@ def parse_args(argv=None):
         "--save-every", type=int, default=1000, help="Model export interval; exports are not training resume state."
     )
     parser.add_argument("--log-every", type=int, default=50)
+
+    parser.add_argument(
+        "--dataset",
+        default="upstream",
+        choices=["upstream", "example"],
+        help="Which DROID reader to use. 'upstream' reads through the copy of upstream's own pipeline under "
+        "vendor/ and is the default: a difference against upstream is then a difference in upstream rather "
+        "than one this example introduced. 'example' keeps droid_dataset.py, this example's own reader, "
+        "whose output was measured identical to upstream's on the pinned window (all 34,214,400 pixel "
+        "values, the action and state rows, the caption and the idle-frame count) -- it is kept so the two "
+        "can be compared again after either side changes.",
+    )
     parser.add_argument("--metrics-every", type=int, default=10)
     parser.add_argument("--metrics-dir", type=Path)
     parser.add_argument("--wandb-project")
@@ -476,7 +512,13 @@ def main():
     unique_dp = dist.get_world_size(plugin.dp_group)
     dp_rank = dist.get_rank(plugin.dp_group)
     cp_leader = args.sp_size == 1 or dist.get_rank(plugin.sp_group) == 0
-    dataset = DroidActionPolicyDataset(
+    # Both readers take the same arguments and return the same sample contract; upstream_dataset.py
+    # documents where the seam between them sits and what was measured across it.
+    # Both names are resolved from the module rather than imported here, so a test can substitute either
+    # reader the same way. Importing upstream_dataset costs nothing by itself: it only reaches for
+    # cosmos_framework when a reader is actually constructed.
+    reader = UpstreamDroidDataset if args.dataset == "upstream" else DroidActionPolicyDataset
+    dataset = reader(
         args.data,
         chunk_length=args.chunk_length,
         fps=args.fps,
@@ -486,6 +528,11 @@ def main():
         keep_ranges=args.keep_ranges,
         success_only=args.success_only,
     )
+    # Which reader ran, and which cosmos_framework the copy resolved to, are read from the run rather than
+    # inferred from the flag: on a host that also has the framework installed the installed one wins, and a
+    # run that quietly used the other copy would be comparing something else.
+    coord.print_on_master(f"[data] {len(dataset)} window(s) from reader {args.dataset!r}")
+    coord.print_on_master(f"[data] pipeline from {report_pipeline_source()}")
     # Both strategies use shared memory for tensor storage. file_system only changes how handles are named;
     # it does not increase /dev/shm capacity. Size workers/batches and container shared memory accordingly.
     torch.multiprocessing.set_sharing_strategy(os.environ.get("COSMOS3_SHM_STRATEGY", "file_system"))