diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index be2df68891ca..41e77a989308 100644 --- a/.pre-commit-config.yaml +++ b/.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 diff --git a/examples/worldmodel/cosmos3/.gitignore b/examples/worldmodel/cosmos3/.gitignore new file mode 100644 index 000000000000..984e65cef2c3 --- /dev/null +++ b/examples/worldmodel/cosmos3/.gitignore @@ -0,0 +1,2 @@ +# Training output written by the example entry point (--out defaults here). +runs/ diff --git a/examples/worldmodel/cosmos3/tests/test_training_entry.py b/examples/worldmodel/cosmos3/tests/test_training_entry.py index 11390ae1af7b..06ac2e32a7bd 100644 --- a/examples/worldmodel/cosmos3/tests/test_training_entry.py +++ b/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 diff --git a/examples/worldmodel/cosmos3/tests/test_upstream_reader_filter.py b/examples/worldmodel/cosmos3/tests/test_upstream_reader_filter.py new file mode 100644 index 000000000000..82e80a28642b --- /dev/null +++ b/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 diff --git a/examples/worldmodel/cosmos3/tests/test_upstream_reader_index.py b/examples/worldmodel/cosmos3/tests/test_upstream_reader_index.py new file mode 100644 index 000000000000..87e126d7c388 --- /dev/null +++ b/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) diff --git a/examples/worldmodel/cosmos3/train_sft.py b/examples/worldmodel/cosmos3/train_sft.py index 4c4016287ea7..3c2f2803da4f 100644 --- a/examples/worldmodel/cosmos3/train_sft.py +++ b/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 "" + 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")) diff --git a/examples/worldmodel/cosmos3/upstream_dataset.py b/examples/worldmodel/cosmos3/upstream_dataset.py new file mode 100644 index 000000000000..61d27d26f929 --- /dev/null +++ b/examples/worldmodel/cosmos3/upstream_dataset.py @@ -0,0 +1,235 @@ +"""Read DROID through upstream's own dataset, and hand this example's packer what it expects. + +The reader is ``vendor/cosmos_framework``, copied byte for byte (see ``vendor/README.md``). This module is +the seam: it converts the reader's sample into the dict ``packing.pack_batch`` consumes, and it does so by +*undoing* what the reader did rather than by re-deriving it, so nothing here is a second implementation of +something upstream already decided. + +Where the seam sits, and why there: + + reader -> composite, quantised to uint8 <- this is what the reader returns + here -> pad to the canvas, normalise to [-1, 1], split the state row off the action + +The reader's video is ``(v * 255).clamp(0, 255).to(uint8)`` of the composed views, which is exactly the +value ``droid_dataset`` reaches before padding. On the pinned window (episode 26827, ordinal 0) the two +were compared element by element with the instruction pinned on both sides and augmentation off: all +34,214,400 values identical, max |d| 0.000 levels. The seam is therefore not a place where a difference can +hide -- it was measured, not assumed. + +What is deliberately still this example's: + +* padding to the canvas and the [-1, 1] normalisation, because the reader stops before them and the stage + that does them upstream cannot be taken without upstream's packing (``vendor/README.md`` has the detail); +* the prompt, the CFG dropout and the tokenisation, whose output was already shown identical to upstream's + (``text_ids``, 170 tokens, byte for byte at steps 0 and 1). + +The scaling half of ``resize_and_reflect_pad`` never runs for DROID: the composite is always 540x640 for +this dataset's 640x360 cameras, the 4:3 canvas is 544x736, and the scale factor is capped at 1.0, so the +call only ever pads. Measured, not reasoned: ``closest_canvas(540, 640) -> (544, 736, '4,3')`` and the +returned content region is (540, 640), unchanged. +""" + +from __future__ import annotations + +import os +from bisect import bisect_right +from pathlib import Path +from typing import Any + +import torch +from torch.utils.data import Dataset + +from examples.worldmodel.cosmos3.droid_dataset import ( + CONCAT_VIEW_DESCRIPTION, + DROID_FILTER_SOURCE_PREFIX, + closest_canvas, + resize_and_reflect_pad, +) + +# The reader derives the dataset version from the directory's basename and rejects anything it does not +# know, so the root it is given has to carry one of upstream's names. A checkout under another name needs a +# symlink; the comparison host has the same one. +KNOWN_ROOT_NAMES = ( + "droid_plus_lerobot_640x360_20260412", + "droid_plus_lerobot_640x360_20260412_sharded", + "droid_plus_lerobot_320x180_20260406", + "droid_plus_lerobot_320x180_20260406_sharded", + "droid_lerobot_20260115_no_noops", +) + + +class UpstreamDroidDataset(Dataset): + """Upstream's ``DROIDLeRobotDataset``, presented in this example's sample contract.""" + + def __init__( + self, + root: str | Path, + chunk_length: int = 32, + fps: float = 15.0, + resolution: str = "480", + augment: bool = True, + success_only: bool = True, + max_episodes: int | None = None, + keep_ranges: str | Path | None = None, + filter_dict_source_prefix: str | None = None, + dataset_revision: str = "", + ) -> None: + root = Path(root) + if root.name not in KNOWN_ROOT_NAMES: + raise ValueError( + f"upstream's reader keys its layout off the directory name and does not know {root.name!r}. " + f"Point --data at a path named one of {KNOWN_ROOT_NAMES} (a symlink is enough)." + ) + from cosmos_framework.data.generator.action.datasets.droid_lerobot_dataset import DROIDLeRobotDataset + + self._inner = DROIDLeRobotDataset( + root=str(root), + fps=fps, + chunk_length=chunk_length, + mode="wam", + viewpoint="concat_view", + use_success_only=success_only, + use_state=True, + action_space="joint_pos", + use_image_augmentation=augment, + # The curated window filter. droid_dataset reads the same keep_ranges_1_0_1.json and its own + # filter is transcribed from the reader's ``_append_index_records``, so handing the file to the + # reader makes the same selection as the code it was transcribed from. Dropping it here would + # train on every window instead of the curated ones, and nothing in the output would say so. + use_filter_dict=keep_ranges is not None, + filter_dict_path=str(keep_ranges) if keep_ranges is not None else None, + # Passed explicitly because the reference's own factory neither declares nor forwards it (a + # TypeError at 5950e7b4, recorded in upstream/README.md); calling the reader directly avoids it. + # The released keep_ranges_1_0_1.json is keyed by composite gs:// URIs, which the reader refuses + # to interpret without this prefix, so a hard-coded "" made --keep-ranges unusable on that file. + # It is resolved the way droid_dataset.py resolves it, so one environment variable serves both + # readers and a run cannot select different windows depending on which reader it was given. + filter_dict_source_prefix=( + filter_dict_source_prefix + if filter_dict_source_prefix is not None + else os.environ.get(DROID_FILTER_SOURCE_PREFIX, "") + ), + ) + self.resolution = resolution + self.fps = fps + self.chunk_length = chunk_length + self.dataset_revision = dataset_revision + self._root_role = "success" if success_only else "mixed" + self._max_episodes = max_episodes + self._episode_first_row = self._first_rows() + + def _first_rows(self) -> dict[tuple[int, int], int]: + """Lowest dataset row each episode owns, which is where its own frame numbering starts. + + ``_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 ``episode_id`` repeats across them. It also + counts from zero again in every registered root, so it only identifies an episode together with the + record's dataset index. Keying a table by ``episode_id`` alone therefore collapses entries -- the + earlier version of this method did, which left a run with the filter on reporting the wrong window + for most samples while the samples themselves were correct. + + Returns: + dict[tuple[int, int], int]: ``(dataset index, episode id)`` to its first row in that dataset. + """ + first: dict[tuple[int, int], int] = {} + for ds_index, row_start, _valid, episode_id in self._inner._episode_records: + key = (int(ds_index), int(episode_id)) + if key not in first or row_start < first[key]: + first[key] = int(row_start) + return first + + def _locate(self, index: int) -> tuple[int, int, int]: + """Which episode a global index falls in, and how far into that episode it starts. + + The span lookup is upstream's own (``cosmos3_action_lerobot.__getitem__``): a bisect over the + cumulative span ends, then the span's record. Doing it the same way is the point -- a second + implementation of the index is exactly what the vendored copy exists to avoid, and a linear scan + over 55,910 episodes ran once per sample. + + Args: + index (int): A global window index, as ``__getitem__`` receives it. + + Returns: + tuple[int, int, int]: Dataset index, episode id, and the window's first frame counted from the + start of that episode. + + Raises: + IndexError: ``index`` is outside the reader's range, which would otherwise be reported against + the last span. + """ + if not 0 <= index < len(self._inner): + raise IndexError(f"window {index} is outside the reader's {len(self._inner)} window(s)") + span = bisect_right(self._inner._episode_cum_ends, index) + span_start = 0 if span == 0 else self._inner._episode_cum_ends[span - 1] + ds_index, row_start, _valid, episode_id = self._inner._episode_records[span] + first_row = self._episode_first_row[(int(ds_index), int(episode_id))] + return int(ds_index), int(episode_id), int(row_start) + (index - span_start) - first_row + + def sample_id(self, index: int) -> str: + """Logical identity of a window, in the same shape ``droid_dataset.sample_id`` uses. + + Same parts, same order, same reason: the revision, the logical root role, the episode, the first + frame and the window length, with no physical path, so the id does not change with where the + release is mounted. + + The dataset index is part of it because ``episode_id`` counts from zero again in every registered + root, so on a mixed success-and-failure run the pair is what identifies an episode. The frame number + is counted from the start of the episode, not from the start of the span, so the several spans one + episode contributes under the curated window filter cannot hand two different windows the same id. + + **It is not the same string as the other reader's**, and deliberately so rather than by oversight. + That one names the episode by its relative id (``RAIL/success/2023-05-20/Sat_...``); upstream's + index carries the episode only as the integer it keys episodes by, so the string is not available + here without re-deriving it from the parquet -- which would be this example reimplementing a piece + of the reader again, the thing the copy exists to stop. + + The consequence is worth stating plainly: **replay packs and held-out manifests recorded under one + reader do not match ids recorded under the other.** Everything else about the two samples is + identical (measured field by field on the pinned window), so this is a bookkeeping boundary, not a + data one -- but a manifest cannot be carried across the ``--dataset`` switch. + """ + ds_index, episode, start = self._locate(index) + parts = [ + self.dataset_revision, + self._root_role, + f"dataset_{ds_index}", + f"episode_{episode}", + str(start), + str(self.chunk_length), + ] + return ":".join(part for part in parts if part) + + def __len__(self) -> int: + return len(self._inner) + + def __getitem__(self, index: int) -> dict[str, Any]: + sample = self._inner[index] + + # The reader prepends the initial state as row 0 of the action when use_state is on. pack_one does + # that prepend itself, so the row is taken back off here rather than the packer being changed: the + # split is the exact inverse of the reader's concatenation, and was verified identical on the pinned + # window (state row and all 32 action rows, bit for bit). + action_full = sample["action"] + state, action = action_full[0], action_full[1:] + + video = sample["video"] # uint8 [C, T, H, W], the quantised composite + target_h, target_w, aspect = closest_canvas(*video.shape[-2:], resolution=self.resolution) + padded, content_hw = resize_and_reflect_pad(video.float(), target_h, target_w) + padded = padded / 127.5 - 1.0 # image_mean = image_std = 0.5 + + _ds_index, episode, _start = self._locate(index) + return { + "video": padded, + "action": action.float(), + "state": state.float(), + "task": sample["ai_caption"], + "idle_frames": int(sample["idle_frames"]), + "additional_view_description": sample.get("additional_view_description", CONCAT_VIEW_DESCRIPTION), + "episode_index": episode, + "sample_id": self.sample_id(index), + "draw_id": 0, + "start": ordinal, + "canvas_hw": torch.tensor([target_h, target_w]), + "content_hw": torch.tensor(content_hw), + "aspect_key": aspect, + } diff --git a/examples/worldmodel/cosmos3/vendor/LICENSE-OpenMDW-1.1 b/examples/worldmodel/cosmos3/vendor/LICENSE-OpenMDW-1.1 new file mode 100644 index 000000000000..a5b8c94f0e68 --- /dev/null +++ b/examples/worldmodel/cosmos3/vendor/LICENSE-OpenMDW-1.1 @@ -0,0 +1,52 @@ +Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + + +OpenMDW License Agreement, version 1.1 (OpenMDW-1.1) + +By exercising rights granted to you under this agreement, you accept and agree +to its terms. + +As used in this agreement, "Model Materials" means the materials provided to +you under this agreement, consisting of: (1) one or more machine learning +models (including architecture and parameters); and (2) all related artifacts +(including associated data, documentation and software) that are provided to +you hereunder. + +Subject to your compliance with this agreement, permission is hereby granted, +free of charge, to deal in the Model Materials without restriction, including +under all copyright, patent, database, and trade secret rights included or +embodied therein. + +If you distribute any portion of the Model Materials, you shall retain in your +distribution (1) a copy of this agreement, and (2) all copyright notices and +other notices of origin included in the Model Materials that are applicable to +your distribution. + +If you file, maintain, or voluntarily participate in a lawsuit against any +person or entity asserting that the Model Materials directly or indirectly +infringe any patent or copyright, then all rights and grants made to you +hereunder are terminated, unless that lawsuit was in response to a +corresponding lawsuit first brought against you. + +This agreement does not impose any restrictions or obligations with respect to +any use, modification, or sharing of any outputs generated by using the Model +Materials. + +THE MODEL MATERIALS ARE PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS +OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +FITNESS FOR A PARTICULAR PURPOSE, TITLE, NONINFRINGEMENT, ACCURACY, OR THE +ABSENCE OF LATENT OR OTHER DEFECTS OR ERRORS, WHETHER OR NOT DISCOVERABLE, ALL +TO THE GREATEST EXTENT PERMISSIBLE UNDER APPLICABLE LAW. + +YOU ARE SOLELY RESPONSIBLE FOR (1) CLEARING RIGHTS OF OTHER PERSONS THAT MAY +APPLY TO THE MODEL MATERIALS OR ANY USE THEREOF, INCLUDING WITHOUT LIMITATION +ANY PERSON'S COPYRIGHTS OR OTHER RIGHTS INCLUDED OR EMBODIED IN THE MODEL +MATERIALS; (2) OBTAINING ANY NECESSARY CONSENTS, PERMISSIONS OR OTHER RIGHTS +REQUIRED FOR ANY USE OF THE MODEL MATERIALS; OR (3) PERFORMING ANY DUE +DILIGENCE OR UNDERTAKING ANY OTHER INVESTIGATIONS INTO THE MODEL MATERIALS OR +ANYTHING INCORPORATED OR EMBODIED THEREIN. + +IN NO EVENT SHALL THE PROVIDERS OF THE MODEL MATERIALS BE LIABLE FOR ANY CLAIM, +DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR +OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE MODEL MATERIALS, THE +USE THEREOF OR OTHER DEALINGS THEREIN. diff --git a/examples/worldmodel/cosmos3/vendor/README.md b/examples/worldmodel/cosmos3/vendor/README.md new file mode 100644 index 000000000000..e173483c7a00 --- /dev/null +++ b/examples/worldmodel/cosmos3/vendor/README.md @@ -0,0 +1,139 @@ +# Vendored DROID data pipeline + +The 34 files under `cosmos_framework/` are copied from + + https://github.com/NVIDIA/cosmos-framework commit 5950e7b4 + +**byte for byte**, with no edits of any kind -- no added provenance headers, no changed import paths, no +removed code. They are licensed under OpenMDW-1.1; the full text is in `LICENSE-OpenMDW-1.1`, copied from +the same commit. + +## Why a copy instead of a reimplementation + +This example used to read DROID through its own `droid_dataset.py`, written against the upstream behaviour. +That is a reimplementation, and a reimplementation can drift from what it reimplements without anything +failing. One such drift was measured: `droid_dataset.py:749` drew the window's instruction with the global +`random.choice` while the reference pins the draw, and a prompt built from a different instruction is a +different prompt -- 104 tokens against 149 in one measured case. Nothing errors; the loss is simply computed +on another input. + +Copying removes that class of problem: the pipeline *is* upstream's, so a difference against upstream is a +difference in upstream, not a difference we introduced. What remains ours is the packing +(`../packing.py`), which cannot be copied: this example trains diffusers' `Cosmos3OmniTransformer`, whose +forward takes loose tensors, while upstream's `OmniMoTModel` takes a `packed_seq` object. The two forward +contracts are different, so the code that produces them has to be. + +## Where the copy stops + +The line is: **reading the data is upstream's, turning it into model input is ours.** + +| stage | code | +| --- | --- | +| DROID index and reading | upstream (`cosmos3_action_lerobot.py`) | +| video decoding | upstream, via torchcodec (`cosmos3_action_lerobot.py:135-138`) | +| instruction pick | upstream (`droid_lerobot_dataset.py:393`) | +| multi-view composition and padding | upstream (`droid_lerobot_dataset.py::_compose_multi_view`) | +| image augmentation | upstream (`droid_lerobot_dataset.py:349`) | +| action, state, idle frames | upstream (`base_dataset.py::_build_result`) | +| prompt assembly | ours (`../packing.py::build_action_json_prompt`) | +| CFG dropout | ours (`../packing.py::tokenize_prompt`) | +| tokenization | ours (`../packing.py`) | +| VAE encode, noising, packing | ours (`../packing.py::pack_one`) | + +Upstream has the last four too, in its SFT wrapper (`utils/transforms.py`, `augmentors/text_tokenizer.py`, +tied together by `action_sft_dataset.py`), and they are deliberately not copied. Two reasons. + +The first is that the chain cannot be taken without taking `sequence_packing`, upstream's own packing, +which this example cannot use at all: it trains diffusers' `Cosmos3OmniTransformer`, whose forward takes +loose tensors, against upstream's `OmniMoTModel`, whose forward takes a `packed_seq` object. Copying 9949 +further lines to reach code that then has to be bypassed buys nothing. + +The second is that the three text stages are already known to agree. With the instruction pinned to the +same index on both sides, every tensor the two forwards receive was compared on the same pinned window: +nine of nine identical byte for byte at step 0, including `text_ids` at 170 tokens, and the same nine at +step 1. The prompt this example builds, drops and tokenizes is therefore the prompt upstream builds, drops +and tokenizes -- measured, not assumed, which is the standard the copy itself is held to. + +## Checking for drift + +Because the files are byte-identical, checking whether upstream has moved is one command against a checkout +of the newer commit: + + diff -r cosmos_framework /path/to/cosmos-framework/cosmos_framework + +Empty output means this copy is still exactly upstream's. Any output is either a change to take or a change +to decide about -- it is never ambiguous, which is the point of not editing the files. + +To re-copy from a newer commit, list the modules actually reachable from the DROID dataset rather than +guessing. In an environment with upstream installed: + + python -c "import sys, os; \ + import cosmos_framework.data.generator.action.datasets.droid_lerobot_dataset; \ + root = os.getcwd() + '/'; \ + print('\n'.join(sorted(f[len(root):] for n in sys.modules \ + if n.startswith('cosmos_framework') \ + for f in [getattr(sys.modules[n], '__file__', None) or ''] if f.startswith(root))))" + +That is how this set of 34 was chosen: it is the transitive closure of the imports, including the sibling +datasets that `datasets/__init__.py` pulls in for other embodiments. Those are inert here -- DROID does not +call them -- but they are kept so that no file had to be edited. + +## What was measured before the switch + +On one pinned window (episode 26827, first window), with the instruction pinned to index 0 on **both** +sides and augmentation off, this example's own reader was compared against upstream's, field by field. + +At the reader's own output: + +| field | result | +| --- | --- | +| composite pixels, 3 x 33 x 540 x 640 = 34,214,400 values | **identical**, max abs difference 0.000 levels | +| action rows, and the state row the reader prepends | **identical** | +| instruction | identical | +| idle frames | identical (8) | + +Then through the adapter (`../upstream_dataset.py`), which is what `pack_batch` is actually handed, all +thirteen keys with no key missing or extra on either side, and every one identical except `sample_id`: +`video` (so the padding and the [-1, 1] normalisation reproduce exactly), `action`, `state`, `task`, +`idle_frames`, `canvas_hw`, `content_hw`, `aspect_key`, `episode_index`, `start`, `draw_id`, +`additional_view_description`. `sample_id` differs by construction and the reason is in that file. + +Two things to be clear about, because "identical" is easy to over-read: + +* **The decoders agree here, which does not make them agree everywhere.** PyAV and torchcodec produced the + same bytes *on the same machine*. The +0.81 LSB recorded elsewhere between this accelerator and x86 is a + platform difference in the scaler, and this measurement says nothing about it either way. +* **The scaling half of the resize was not exercised, because this dataset cannot reach it.** DROID's three + 640x360 cameras always compose to 540x640, `closest_canvas(540, 640)` is the 4:3 canvas 544x736, and the + scale factor is capped at 1.0, so the call only ever pads; measured, the returned content region comes + back as (540, 640), unchanged. The branch is unreachable for DROID rather than untested. +* The augmentation was switched off for the comparison and was not separately verified. It is upstream's + own code after this change, so what the previous implementation did differently, if anything, now only + affects runs taken before it. + +The entry point then trained two steps on the vendored reader (`--dataset upstream`, one card): +`16338993 window(s)`, `samples=1 AR=168 DM=3093`, loss 24.177132 then 11.412529. + +## Why the dependencies are hand-picked, and what that costs + +Upstream's pipeline needs `lerobot`, and resolving `lerobot`'s own dependencies is not safe here: the +closure would install `diffusers 0.35.2` over the 0.40.0 that provides `Cosmos3OmniTransformer`, and +`av 15.1.0` over the 18.1.0 whose bundled FFmpeg 8 the decoder links against. Both are load-bearing, so the +packages below were installed individually instead: + +| package | version | why | +| --- | --- | --- | +| `lerobot` | 0.4.4 | `LeRobotDataset`, the base of upstream's reader. Installed with `--no-deps`; 0.4.4 is the version upstream builds against | +| `torchcodec` | 0.14.0+cpu | upstream decodes video with it (`cosmos3_action_lerobot.py:135-138` raises if it is absent) | +| `datasets` | 4.8.5 | required by `lerobot`; pinned to the version `lerobot` 0.4.4 was built against rather than the newest | +| `loguru` | latest | required by `cosmos_framework.utils.log` | + +`torchcodec` loads FFmpeg by SONAME and finds no system FFmpeg in this image. It ships a variant for +FFmpeg 8 (`libtorchcodec_core8.so`, needing `libavutil.so.60`), which is the version PyAV already bundles +under mangled names, so a directory of correctly-named symlinks onto `av.libs` satisfies it without +changing the image. + +**This is the one place the copy does not make us upstream.** Upstream runs inside `lerobot`'s full +dependency set; this runs inside a subset chosen by hand. The code is identical, the environment is not, +and a difference that comes from a package version is still ours to explain. The versions above are pinned +for that reason. diff --git a/examples/worldmodel/cosmos3/vendor/cosmos_framework/__init__.py b/examples/worldmodel/cosmos3/vendor/cosmos_framework/__init__.py new file mode 100644 index 000000000000..28a81be6930a --- /dev/null +++ b/examples/worldmodel/cosmos3/vendor/cosmos_framework/__init__.py @@ -0,0 +1,2 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: OpenMDW-1.1 diff --git a/examples/worldmodel/cosmos3/vendor/cosmos_framework/data/__init__.py b/examples/worldmodel/cosmos3/vendor/cosmos_framework/data/__init__.py new file mode 100644 index 000000000000..503ec1b18d58 --- /dev/null +++ b/examples/worldmodel/cosmos3/vendor/cosmos_framework/data/__init__.py @@ -0,0 +1,3 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: OpenMDW-1.1 + diff --git a/examples/worldmodel/cosmos3/vendor/cosmos_framework/data/generator/__init__.py b/examples/worldmodel/cosmos3/vendor/cosmos_framework/data/generator/__init__.py new file mode 100644 index 000000000000..503ec1b18d58 --- /dev/null +++ b/examples/worldmodel/cosmos3/vendor/cosmos_framework/data/generator/__init__.py @@ -0,0 +1,3 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: OpenMDW-1.1 + diff --git a/examples/worldmodel/cosmos3/vendor/cosmos_framework/data/generator/action/__init__.py b/examples/worldmodel/cosmos3/vendor/cosmos_framework/data/generator/action/__init__.py new file mode 100644 index 000000000000..28a81be6930a --- /dev/null +++ b/examples/worldmodel/cosmos3/vendor/cosmos_framework/data/generator/action/__init__.py @@ -0,0 +1,2 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: OpenMDW-1.1 diff --git a/examples/worldmodel/cosmos3/vendor/cosmos_framework/data/generator/action/action_caption_attribute_adapter.py b/examples/worldmodel/cosmos3/vendor/cosmos_framework/data/generator/action/action_caption_attribute_adapter.py new file mode 100644 index 000000000000..187116568525 --- /dev/null +++ b/examples/worldmodel/cosmos3/vendor/cosmos_framework/data/generator/action/action_caption_attribute_adapter.py @@ -0,0 +1,267 @@ +# ----------------------------------------------------------------------------- +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. +# All rights reserved. +# ----------------------------------------------------------------------------- +"""Central caption-attribute protocols for Action datasets.""" + +from __future__ import annotations + +from collections.abc import Mapping +from dataclasses import dataclass +from typing import Literal + +Domain = Literal["real", "synthetic"] +DomainMappings: dict[Domain, str] = { + "real": "The video is captured from a real-world environment.", + "synthetic": "The video is captured from a synthetic environment.", +} + +Embodiment = Literal["single_arm", "dual_arm", "human", "camera_only", "vehicle"] +EmbodimentMappings: dict[Embodiment, str] = { + "single_arm": "The video shows a single-arm robot.", + "dual_arm": "The video shows a dual-arm robot.", + "human": "The video shows a human actor.", + "camera_only": "The video shows the perspective of a moving camera.", + "vehicle": "The video shows the perspective of a moving vehicle.", +} + +ViewComposition = Literal[ + "separate_camera_views", + "front", + "above", + "third_person", + "wrist", + "ego_head", + "static_single", + "static_single_first_person", + "dual_wrist", + "head_over_wrists", + "above_over_wrists", + "wrist_over_side_third_person", + "above_over_side_third_person", +] +ViewCompositionMappings: dict[ViewComposition, str] = { + "separate_camera_views": "The synchronized camera views are encoded separately, in a fixed camera order.", + "front": "This video is captured from a dynamic front-facing perspective looking at the scene.", + "above": "This video is captured from a third-person perspective looking towards the agent from above.", + "third_person": "This video is captured from a third-person perspective looking towards the robot.", + "wrist": "This video is captured from a wrist-mounted camera.", + "ego_head": "This video is captured from a head-mounted camera showing an egocentric view of the scene.", + "static_single": "This video is captured from a static perspective looking towards the actor.", + "static_single_first_person": "This video is captured from a static first-person perspective looking at the scene.", + "dual_wrist": "This video is captured from two wrist-mounted cameras.", + "head_over_wrists": "This video contains concatenated views from multiple camera perspectives. The top row shows the head-mounted camera view looking down at the workspace. The bottom row contains two horizontally concatenated wrist-mounted camera views: the left hand camera on the left and the right hand camera on the right.", + "above_over_wrists": "This video contains concatenated views from multiple camera perspectives. The top row shows the overhead/top camera. The bottom row contains the left camera on the left and the right camera on the right.", + "wrist_over_side_third_person": "This video contains concatenated views from multiple camera perspectives. The top row is from the wrist-mounted camera. The bottom row contains two horizontally concatenated static third-person perspective views of the scene from opposite sides, with the robot visible.", + "above_over_side_third_person": "This video contains concatenated views from multiple camera perspectives. The top row shows third-person static perspective view looking towards the robot from above. The bottom-left video shows the third-person static perspective view looking at the scene from the left side. The bottom-right video shows the third-person static perspective view looking at the scene from the right side.", +} + +CaptionSubject = Literal[ + "camera_ego_pose", + "right_wrist_camera", + "dual_wrist_cameras", + "ego_and_dual_wrist_cameras", + "vehicle_ego_pose", + "right_arm_end_effector", + "right_arm_attachment_site_end_effector", + "right_handheld_gripper_end_effector", + "dual_arm_end_effectors", + "dual_arm_joint_positions", + "ego_and_dual_arm_end_effectors", + "dual_handheld_gripper_end_effectors", + "human_wrists_and_fingertips", + "human_head_and_wrists", + "human_head_wrists_and_fingertips", +] +CaptionSubjectMappings: dict[CaptionSubject, str] = { + "right_wrist_camera": ( + "The action target is the wrist-mounted camera motion and gripper openness, " + "mapped to the right wrist-pose and right-openness components." + ), + "dual_wrist_cameras": ( + "The action targets are the right and left wrist-mounted camera motions and gripper openness, " + "mapped to the right and left wrist-pose and openness components." + ), + "ego_and_dual_wrist_cameras": ( + "The action targets are the head ego motion, right and left wrist-mounted camera motions, and " + "gripper openness, mapped to the ego-pose, right and left wrist-pose and openness components." + ), + "dual_arm_joint_positions": ( + "The action performed by the robot is defined as the joint positions of its left and right arms " + "and both grippers' positions." + ), + "camera_ego_pose": "The camera action is defined as its ego motion, mapped to the ego-pose component.", + "vehicle_ego_pose": ( + "The action performed by the vehicle is defined as its ego motion, mapped to the ego-pose component." + ), + "right_arm_end_effector": ( + "The action performed by the robot is defined as the motion of its end-effector link and gripper openness, " + "mapped to the right wrist-pose and right-openness components." + ), + "right_arm_attachment_site_end_effector": ( + "The action performed by the robot is defined as the motion of its attachment-site end-effector link and " + "gripper openness, mapped to the right wrist-pose and right-openness components." + ), + "right_handheld_gripper_end_effector": ( + "The action demonstrated using the handheld gripper is defined as the motion of its end-effector link and " + "gripper openness, mapped to the right wrist-pose and right-openness components." + ), + "dual_arm_end_effectors": ( + "The action performed by the robot is defined as the motion of its right and left end-effector links and both " + "grippers' openness, mapped to the right and left wrist-pose and openness components." + ), + "ego_and_dual_arm_end_effectors": ( + "The action performed by the robot is defined as the ego motion of its head, the motion of its right and " + "left end-effector links, and both grippers' openness, mapped to the ego-pose, right and left wrist-pose, " + "and right and left openness components." + ), + "dual_handheld_gripper_end_effectors": ( + "The action demonstrated using the handheld grippers is defined as the motion of their right and left " + "end-effector links and both grippers' openness, mapped to the right and left wrist-pose and openness " + "components." + ), + "human_wrists_and_fingertips": ( + "The action performed by the human actor is defined as the wrist and fingertip motion of the actor's hands, " + "mapped to the right and left wrist-pose and fingertip components." + ), + "human_head_and_wrists": ( + "The action performed by the human actor is defined as the ego pose of the actor's head and the wrist motion " + "of the actor's hands, mapped to the ego-pose and right and left wrist-pose components." + ), + "human_head_wrists_and_fingertips": ( + "The action performed by the human actor is defined as the ego motion of the actor's head and the wrist and " + "fingertip motion of the actor's hands, mapped to the ego-pose and right and left wrist-pose and fingertip " + "components." + ), +} + +ActionCaptionAttributeValue = str | float | int +ACTION_CAPTION_POSTFIX_KEYS: tuple[str, ...] = ( + "domain_postfix", + "embodiment_postfix", + "view_postfix", + "subject_postfix", +) + + +def append_action_caption_semantics(caption: str, attributes: Mapping[str, object]) -> str: + """Append the registered semantic postfixes to a non-empty caption in canonical order.""" + if not caption: + return caption + + postfixes: list[str] = [] + for key in ACTION_CAPTION_POSTFIX_KEYS: + postfix = attributes.get(key) + if not isinstance(postfix, str) or not postfix.strip(): + raise ValueError(f"Action caption attribute {key!r} must be a non-empty string.") + postfixes.append(postfix.strip()) + + caption = caption.rstrip() + separator = " " if caption.endswith((".", "!", "?")) else ". " + return f"{caption}{separator}{' '.join(postfixes)}" + + +@dataclass(frozen=True) +class ActionCaptionAttributeProtocol: + """Dataset-level semantic facts that cannot be inferred from tensors.""" + + domain: Domain + embodiment: Embodiment + view_composition: ViewComposition + caption_subject: CaptionSubject + + +class ActionCaptionAttributeAdapter: + """Resolve semantic protocol plus timing derived from one loaded window.""" + + def __init__(self, protocols: Mapping[str, ActionCaptionAttributeProtocol]) -> None: + self.protocols: dict[str, ActionCaptionAttributeProtocol] = dict(protocols) + + def supports(self, dataset_name: str) -> bool: + """Return whether caption attributes are enabled for this dataset.""" + return dataset_name in self.protocols + + def resolve( + self, + dataset_name: str, + *, + fps: float, + observation_count: int, + view_count: int, + view_composition: ViewComposition | None = None, + ) -> dict[str, ActionCaptionAttributeValue]: + """Return deterministic attributes for one loaded training window.""" + try: + protocol = self.protocols[dataset_name] + except KeyError as error: + raise KeyError(f"No action-caption attribute protocol registered for dataset {dataset_name!r}.") from error + if fps <= 0: + raise ValueError(f"fps must be positive, got {fps}.") + if observation_count < 1: + raise ValueError(f"observation_count must be positive, got {observation_count}.") + if view_count < 1: + raise ValueError(f"view_count must be positive, got {view_count}.") + resolved_view_composition = view_composition or protocol.view_composition + action_transition_count = observation_count - 1 + return { + "dataset_name": dataset_name, + "domain": protocol.domain, + "embodiment": protocol.embodiment, + "view_composition": resolved_view_composition, + "caption_subject": protocol.caption_subject, + "domain_postfix": DomainMappings[protocol.domain], + "embodiment_postfix": EmbodimentMappings[protocol.embodiment], + "view_postfix": ViewCompositionMappings[resolved_view_composition], + "subject_postfix": CaptionSubjectMappings[protocol.caption_subject], + "fps": float(fps), + "duration_seconds": action_transition_count / float(fps), + "observation_count": observation_count, + "action_transition_count": action_transition_count, + "view_count": view_count, + } + + +P = ActionCaptionAttributeProtocol + +# Edit this table to update dataset semantics. Caption subjects group datasets +# with equivalent unified_v1 components; Runtime timing and counts are +# derived by ActionCaptionAttributeAdapter from the loaded training item. +ACTION_CAPTION_ATTRIBUTE_ADAPTER = ActionCaptionAttributeAdapter( + { + "fractal": P("real", "single_arm", "static_single_first_person", "right_arm_end_effector"), + "bridge": P("real", "single_arm", "static_single_first_person", "right_arm_end_effector"), + "droid": P("real", "single_arm", "wrist_over_side_third_person", "right_arm_end_effector"), + "av": P("real", "vehicle", "front", "vehicle_ego_pose"), + "camera_action_v2": P("real", "camera_only", "front", "camera_ego_pose"), + "embodiment_c": P("real", "dual_arm", "head_over_wrists", "ego_and_dual_arm_end_effectors"), + "embodiment_c_ext": P("real", "dual_arm", "head_over_wrists", "ego_and_dual_arm_end_effectors"), + "agibot_world": P("real", "dual_arm", "head_over_wrists", "ego_and_dual_arm_end_effectors"), + "agibot": P("real", "dual_arm", "head_over_wrists", "ego_and_dual_arm_end_effectors"), + "embodiment_b": P("real", "dual_arm", "head_over_wrists", "ego_and_dual_arm_end_effectors"), + "genrobot_10k": P("real", "dual_arm", "dual_wrist", "dual_handheld_gripper_end_effectors"), + "robodojo": P("synthetic", "dual_arm", "above_over_wrists", "dual_arm_end_effectors"), + "robodojo_joint": P("synthetic", "dual_arm", "above_over_wrists", "dual_arm_joint_positions"), + "robodojo_real": P("real", "dual_arm", "above_over_wrists", "dual_arm_end_effectors"), + "robodojo_real_joint": P("real", "dual_arm", "above_over_wrists", "dual_arm_joint_positions"), + "robomind_franka": P("real", "single_arm", "above_over_side_third_person", "right_arm_end_effector"), + "robomind_franka_dual": P("real", "dual_arm", "above_over_side_third_person", "dual_arm_end_effectors"), + "robomind_ur": P("real", "single_arm", "above", "right_arm_attachment_site_end_effector"), + "umi_singlearm_lerobot_480_20260625": P("real", "single_arm", "wrist", "right_handheld_gripper_end_effector"), + "umi_singlearm_lerobot_256_20260625": P("real", "single_arm", "wrist", "right_handheld_gripper_end_effector"), + "umi_bimanual_lerobot_480_20260625": P("real", "dual_arm", "dual_wrist", "dual_handheld_gripper_end_effectors"), + "maniparena_tabletop": P("real", "dual_arm", "head_over_wrists", "dual_arm_end_effectors"), + "maniparena_mobile": P("real", "dual_arm", "head_over_wrists", "ego_and_dual_arm_end_effectors"), + "abc_130k": P("real", "dual_arm", "above_over_wrists", "dual_arm_end_effectors"), + "molmoact2_yam": P("real", "dual_arm", "above_over_wrists", "dual_arm_end_effectors"), + "xdof_yam_v5": P("real", "dual_arm", "above_over_wrists", "dual_arm_end_effectors"), + "xdof_yam_annotated": P("real", "dual_arm", "above_over_wrists", "dual_arm_end_effectors"), + "so100_midtrain_15hz": P("real", "single_arm", "third_person", "right_arm_end_effector"), + "so101_midtrain_15hz": P("real", "single_arm", "third_person", "right_arm_end_effector"), + "biso101_midtrain_15hz": P("real", "dual_arm", "third_person", "dual_arm_end_effectors"), + "web_human_action_hand": P("real", "human", "static_single", "human_wrists_and_fingertips"), + "web_human_action_body": P("real", "human", "static_single", "human_head_and_wrists"), + "vitra_ego4d": P("real", "human", "ego_head", "human_head_wrists_and_fingertips"), + "embodiment_a": P("real", "human", "ego_head", "human_head_wrists_and_fingertips"), + "behavior1k_subtask": P("synthetic", "dual_arm", "head_over_wrists", "ego_and_dual_arm_end_effectors"), + } +) diff --git a/examples/worldmodel/cosmos3/vendor/cosmos_framework/data/generator/action/action_normalization.py b/examples/worldmodel/cosmos3/vendor/cosmos_framework/data/generator/action/action_normalization.py new file mode 100644 index 000000000000..8504cf0c3121 --- /dev/null +++ b/examples/worldmodel/cosmos3/vendor/cosmos_framework/data/generator/action/action_normalization.py @@ -0,0 +1,60 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: OpenMDW-1.1 + +"""Action normalization helpers.""" + +import json +from pathlib import Path + +import numpy as np +import torch + +from cosmos_framework.utils import log + + +def load_action_stats(stats_path: str) -> dict[str, np.ndarray]: + """Load pre-computed action normalization stats from a JSON file.""" + path = Path(stats_path) + if not path.exists(): + raise FileNotFoundError(f"Action normalization stats not found at {stats_path}.") + log.info(f"Loading action normalization stats from {stats_path}") + with path.open("r") as f: + raw = json.load(f) + stat_keys = {"mean", "std", "min", "max", "q01", "q99"} + return {key: np.array(value, dtype=np.float32) for key, value in raw.items() if key in stat_keys} + + +def normalize_action( + action: torch.Tensor, + method: str, + stats: dict[str, torch.Tensor], +) -> torch.Tensor: + """Normalize action tensor.""" + if method == "quantile": + q01, q99 = stats["q01"], stats["q99"] + denom = (q99 - q01).clamp(min=1e-8) + return 2.0 * (action - q01) / denom - 1.0 + if method == "meanstd": + return (action - stats["mean"]) / stats["std"].clamp(min=1e-8) + if method == "minmax": + lo, hi = stats["min"], stats["max"] + denom = (hi - lo).clamp(min=1e-8) + return 2.0 * (action - lo) / denom - 1.0 + raise ValueError(f"Unknown normalization method: {method!r}") + + +def denormalize_action( + action: torch.Tensor, + method: str, + stats: dict[str, torch.Tensor], +) -> torch.Tensor: + """Denormalize action tensor.""" + if method == "quantile": + q01, q99 = stats["q01"], stats["q99"] + return 0.5 * (action + 1.0) * (q99 - q01) + q01 + if method == "meanstd": + return action * stats["std"] + stats["mean"] + if method == "minmax": + lo, hi = stats["min"], stats["max"] + return 0.5 * (action + 1.0) * (hi - lo) + lo + raise ValueError(f"Unknown normalization method: {method!r}") diff --git a/examples/worldmodel/cosmos3/vendor/cosmos_framework/data/generator/action/agibot_fk.py b/examples/worldmodel/cosmos3/vendor/cosmos_framework/data/generator/action/agibot_fk.py new file mode 100644 index 000000000000..ba73ddee3bff --- /dev/null +++ b/examples/worldmodel/cosmos3/vendor/cosmos_framework/data/generator/action/agibot_fk.py @@ -0,0 +1,398 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: OpenMDW-1.1 + +"""Lightweight AgiBot World forward kinematics for datasets and viewers.""" + +from __future__ import annotations + +import xml.etree.ElementTree as ET +from functools import lru_cache + +import numpy as np + +from cosmos_framework.data.generator.action.agibot_spec import ( + AGIBOT_WORLD_ARM_JOINT_NAMES_LEFT, + AGIBOT_WORLD_ARM_JOINT_NAMES_RIGHT, + AGIBOT_WORLD_ARM_STATE_SLICE, + AGIBOT_WORLD_EXT_ARM_STATE_SLICE, + AGIBOT_WORLD_EXT_STATE_HEAD_PITCH_IDX, + AGIBOT_WORLD_EXT_STATE_HEAD_YAW_IDX, + AGIBOT_WORLD_EXT_STATE_LEFT_HAND_SLICE, + AGIBOT_WORLD_EXT_STATE_RIGHT_HAND_SLICE, + AGIBOT_WORLD_EXT_STATE_ROBOT_ORIENTATION_SLICE, + AGIBOT_WORLD_EXT_STATE_ROBOT_POSITION_SLICE, + AGIBOT_WORLD_EXT_STATE_WAIST_LIFT_IDX, + AGIBOT_WORLD_EXT_STATE_WAIST_PITCH_IDX, + AGIBOT_WORLD_GRIPPER_OPEN_ACTUATOR_DEG, + AGIBOT_WORLD_GRIPPER_OPEN_ANGLE_RAD, + AGIBOT_WORLD_HEAD_CAMERA_LINK_NAME, + AGIBOT_WORLD_HEAD_PITCH_JOINT_NAME, + AGIBOT_WORLD_HEAD_YAW_JOINT_NAME, + AGIBOT_WORLD_LEFT_EE_LINK_NAME, + AGIBOT_WORLD_LEFT_GRIPPER_JOINT_MIMICS, + AGIBOT_WORLD_RIGHT_EE_LINK_NAME, + AGIBOT_WORLD_RIGHT_GRIPPER_JOINT_MIMICS, + AGIBOT_WORLD_STATE_HEAD_PITCH_IDX, + AGIBOT_WORLD_STATE_HEAD_YAW_IDX, + AGIBOT_WORLD_STATE_WAIST_LIFT_IDX, + AGIBOT_WORLD_STATE_WAIST_PITCH_IDX, + AGIBOT_WORLD_WAIST_LIFT_JOINT_NAME, + AGIBOT_WORLD_WAIST_PITCH_JOINT_NAME, + get_agibot_world_embodiment_spec, + get_agibot_world_kind_spec, + get_agibot_world_urdf_path, +) +from cosmos_framework.data.generator.action.utils.pose_utils import convert_rotation + +_GRIPPER_VALUE_EPS = 1e-4 +_QUATERNION_NORM_EPS = 1e-8 +_GRIPPER_ACTUATOR_OVERSHOOT_DEG = 5.0 +# Main-branch wrist rotations composed with one extra local-Z 180 degree rotation. +AGIBOT_WORLD_LEFT_GRIPPER_TO_OPENCV: np.ndarray = np.asarray( + [ + [0.0, 1.0, 0.0], + [-1.0, 0.0, 0.0], + [0.0, 0.0, 1.0], + ], + dtype=np.float32, +) +AGIBOT_WORLD_RIGHT_GRIPPER_TO_OPENCV: np.ndarray = np.asarray( + [ + [0.0, -1.0, 0.0], + [1.0, 0.0, 0.0], + [0.0, 0.0, 1.0], + ], + dtype=np.float32, +) +AGIBOT_WORLD_GRIPPER_TO_OPENCV_BY_WRIST: dict[str, np.ndarray] = { + "right_wrist": AGIBOT_WORLD_RIGHT_GRIPPER_TO_OPENCV, + "left_wrist": AGIBOT_WORLD_LEFT_GRIPPER_TO_OPENCV, +} + + +def _scale_to_unit_interval(values: np.ndarray, scale: float) -> np.ndarray: + """Scale non-negative gripper actuator values to ``[0,1]``.""" + + return np.clip(values / scale, 0.0, 1.0).astype(np.float32, copy=False) + + +def _scale_negative_to_unit_interval(values: np.ndarray, scale: float) -> np.ndarray: + """Scale URDF-style negative gripper angles to ``[0,1]`` open fractions.""" + + return np.clip(-values / scale, 0.0, 1.0).astype(np.float32, copy=False) + + +def _normalize_quaternions_xyzw(quaternions: np.ndarray) -> np.ndarray: + """Normalize ``xyzw`` quaternions, treating all-zero rows as identity.""" + + normalized = np.asarray(quaternions, dtype=np.float32).copy() # [T,4] + norms = np.linalg.norm(normalized, axis=-1, keepdims=True) # [T,1] + valid = norms[:, 0] >= _QUATERNION_NORM_EPS # [T] + normalized[valid] = normalized[valid] / norms[valid] # [T_valid,4] + normalized[~valid] = np.asarray([0.0, 0.0, 0.0, 1.0], dtype=np.float32) # [T_invalid,4] + return normalized + + +def _quat_xyzw_to_rotation_matrix(quaternions: np.ndarray) -> np.ndarray: + """Convert ``xyzw`` quaternions to rotation matrices.""" + + normalized = _normalize_quaternions_xyzw(quaternions) # [T,4] + rotations = convert_rotation( + normalized, + input_format="quat_xyzw", + output_format="matrix", + normalize_matrix=True, + ) + return np.asarray(rotations, dtype=np.float32) + + +def build_robot_base_transforms(positions: np.ndarray, quaternions: np.ndarray) -> np.ndarray: + """Build robot-base poses from position and ``xyzw`` quaternion arrays.""" + + positions = np.asarray(positions, dtype=np.float32) # [T,3] + quaternions = np.asarray(quaternions, dtype=np.float32) # [T,4] + if positions.ndim != 2 or positions.shape[1] != 3: + raise ValueError(f"robot base positions must have shape [T,3], got {positions.shape}.") + if quaternions.ndim != 2 or quaternions.shape[1] != 4: + raise ValueError(f"robot base quaternions must have shape [T,4], got {quaternions.shape}.") + if positions.shape[0] != quaternions.shape[0]: + raise ValueError( + f"robot base positions/quaternions must share T, got {positions.shape[0]} and {quaternions.shape[0]}." + ) + + transforms = np.tile(np.eye(4, dtype=np.float32), (positions.shape[0], 1, 1)) # [T,4,4] + transforms[:, :3, :3] = _quat_xyzw_to_rotation_matrix(quaternions) # [T,3,3] + transforms[:, :3, 3] = positions # [T,3] + return transforms + + +def _invert_rigid_transform(transform: np.ndarray) -> np.ndarray: + """Invert one homogeneous rigid transform.""" + + inverse = np.eye(4, dtype=np.float32) # [4,4] + rotation_t = transform[:3, :3].T.astype(np.float32, copy=False) # [3,3] + inverse[:3, :3] = rotation_t + inverse[:3, 3] = -(rotation_t @ transform[:3, 3]) # [3] + return inverse + + +def apply_robot_base_motion_to_poses( + poses_by_name: dict[str, np.ndarray], + positions: np.ndarray, + quaternions: np.ndarray, +) -> dict[str, np.ndarray]: + """Apply mobile-base motion to FK poses, normalized to the first frame.""" + + base_poses = build_robot_base_transforms(positions, quaternions) # [T,4,4] + initial_base_inv = _invert_rigid_transform(base_poses[0]) # [4,4] + base_motion = np.einsum("ij,tjk->tik", initial_base_inv, base_poses).astype(np.float32, copy=False) # [T,4,4] + return { + name: np.einsum("tij,tjk->tik", base_motion, poses).astype(np.float32, copy=False) # [T,4,4] + for name, poses in poses_by_name.items() + } + + +def _apply_ext_base_motion_to_poses( + poses_by_name: dict[str, np.ndarray], + states: np.ndarray, + embodiment_type: str, +) -> dict[str, np.ndarray]: + """Apply ext mobile-base motion to FK poses, normalized to the first frame.""" + + if embodiment_type != "agibot_world_gripper_ext": + return poses_by_name + if states.shape[1] < AGIBOT_WORLD_EXT_STATE_ROBOT_ORIENTATION_SLICE.stop: + raise ValueError( + f"agibot_world_gripper_ext state must include robot pose through index " + f"{AGIBOT_WORLD_EXT_STATE_ROBOT_ORIENTATION_SLICE.stop - 1}, got shape {states.shape}." + ) + + positions = states[:, AGIBOT_WORLD_EXT_STATE_ROBOT_POSITION_SLICE].astype(np.float32, copy=False) # [T,3] + quaternions = states[:, AGIBOT_WORLD_EXT_STATE_ROBOT_ORIENTATION_SLICE].astype(np.float32, copy=False) # [T,4] + return apply_robot_base_motion_to_poses(poses_by_name, positions, quaternions) + + +def apply_agibot_gripper_to_opencv( + poses_by_name: dict[str, np.ndarray], + to_opencv_by_wrist: dict[str, np.ndarray], +) -> dict[str, np.ndarray]: + """Post-rotate AgiBot gripper wrist poses into OpenCV convention.""" + + aligned = {name: poses.astype(np.float32, copy=True) for name, poses in poses_by_name.items()} # {name:[...,4,4]} + for wrist_name, wrist_to_opencv in to_opencv_by_wrist.items(): + poses = aligned.get(wrist_name) + if poses is None: + continue + aligned[wrist_name][..., :3, :3] = poses[..., :3, :3] @ wrist_to_opencv.astype(poses.dtype) # [...,3,3] + return aligned + + +def _get_agibot_world_mujoco_kinematics_xml() -> str: + """Build a MuJoCo-loadable kinematics-only XML string from the committed URDF.""" + + root = ET.parse(get_agibot_world_urdf_path()).getroot() + mujoco_element = root.find("mujoco") + if mujoco_element is None: + mujoco_element = ET.Element("mujoco") + root.insert(0, mujoco_element) + compiler_element = mujoco_element.find("compiler") + if compiler_element is None: + compiler_element = ET.SubElement(mujoco_element, "compiler") + compiler_element.attrib["fusestatic"] = "false" + + for link_element in root.findall("link"): + for child_element in list(link_element): + if child_element.tag in {"visual", "collision"}: + link_element.remove(child_element) + + return ET.tostring(root, encoding="unicode") + + +class _MujocoFk: + """MuJoCo-backed FK engine for the committed AgiBot G1 omnipicker URDF.""" + + def __init__(self) -> None: + import mujoco + + self._mujoco = mujoco + self.model = mujoco.MjModel.from_xml_string(_get_agibot_world_mujoco_kinematics_xml()) + self.data = mujoco.MjData(self.model) + self._joint_qpos_addresses: dict[str, int] = {} + for joint_id in range(self.model.njnt): + joint_name = mujoco.mj_id2name(self.model, mujoco.mjtObj.mjOBJ_JOINT, joint_id) + if joint_name is not None: + self._joint_qpos_addresses[joint_name] = int(self.model.jnt_qposadr[joint_id]) + + def link_poses(self, joint_values: dict[str, float]) -> dict[str, np.ndarray]: + """Return world transforms for every named body in the MuJoCo model.""" + + self.data.qpos[:] = 0.0 + for joint_name, joint_value in joint_values.items(): + qpos_address = self._joint_qpos_addresses.get(joint_name) + if qpos_address is not None: + self.data.qpos[qpos_address] = float(joint_value) + self._mujoco.mj_forward(self.model, self.data) + + poses: dict[str, np.ndarray] = {} + for body_id in range(1, self.model.nbody): + body_name = self._mujoco.mj_id2name(self.model, self._mujoco.mjtObj.mjOBJ_BODY, body_id) + if body_name is None: + continue + transform = np.eye(4, dtype=np.float32) + transform[:3, :3] = self.data.xmat[body_id].reshape(3, 3).astype(np.float32, copy=False) + transform[:3, 3] = self.data.xpos[body_id].astype(np.float32, copy=False) + poses[body_name] = transform + return poses + + +@lru_cache(maxsize=1) +def _get_fk_engine() -> _MujocoFk: + """Return a cached MuJoCo FK engine for the committed AgiBot URDF.""" + + return _MujocoFk() + + +def _extract_joint_values_from_state(state: np.ndarray, embodiment_type: str) -> dict[str, float]: + """Map one observation.state vector to the URDF joint names used for FK.""" + + if embodiment_type == "agibot_world_gripper_ext": + # Ext layout: 94-dim state with joints at different offsets. + arm_state = state[AGIBOT_WORLD_EXT_ARM_STATE_SLICE] + head_yaw = float(state[AGIBOT_WORLD_EXT_STATE_HEAD_YAW_IDX]) + head_pitch = float(state[AGIBOT_WORLD_EXT_STATE_HEAD_PITCH_IDX]) + waist_lift = float(state[AGIBOT_WORLD_EXT_STATE_WAIST_LIFT_IDX]) + waist_pitch = float(state[AGIBOT_WORLD_EXT_STATE_WAIST_PITCH_IDX]) + else: + arm_state = state[AGIBOT_WORLD_ARM_STATE_SLICE] + head_yaw = float(state[AGIBOT_WORLD_STATE_HEAD_YAW_IDX]) + head_pitch = float(state[AGIBOT_WORLD_STATE_HEAD_PITCH_IDX]) + waist_pitch = float(state[AGIBOT_WORLD_STATE_WAIST_PITCH_IDX]) + waist_lift = float(state[AGIBOT_WORLD_STATE_WAIST_LIFT_IDX]) + + joint_values = { + AGIBOT_WORLD_WAIST_LIFT_JOINT_NAME: float(waist_lift), + AGIBOT_WORLD_WAIST_PITCH_JOINT_NAME: float(waist_pitch), + AGIBOT_WORLD_HEAD_YAW_JOINT_NAME: float(head_yaw), + AGIBOT_WORLD_HEAD_PITCH_JOINT_NAME: float(head_pitch), + } + joint_values.update({name: float(arm_state[idx]) for idx, name in enumerate(AGIBOT_WORLD_ARM_JOINT_NAMES_LEFT)}) + joint_values.update({name: float(arm_state[7 + idx]) for idx, name in enumerate(AGIBOT_WORLD_ARM_JOINT_NAMES_RIGHT)}) + _set_gripper_joint_values_from_state(joint_values, state, embodiment_type) + return joint_values + + +def _set_gripper_joint_values_from_state( + joint_values: dict[str, float], + state: np.ndarray, + embodiment_type: str, +) -> None: + """Map observed scalar gripper state into all omnipicker finger joints.""" + + embodiment_spec = get_agibot_world_embodiment_spec(embodiment_type) + if embodiment_spec.kind != "gripper": + return + + if embodiment_type == "agibot_world_gripper_ext": + left_raw = float(state[AGIBOT_WORLD_EXT_STATE_LEFT_HAND_SLICE][0]) + right_raw = float(state[AGIBOT_WORLD_EXT_STATE_RIGHT_HAND_SLICE][0]) + else: + kind_spec = get_agibot_world_kind_spec(embodiment_type) + state_hand_slice = kind_spec.state_hand_slice + left_raw = float(state[state_hand_slice.start]) + right_raw = float(state[state_hand_slice.start + 1]) + + left_open = float(convert_gripper_state_to_open_fraction(np.asarray([left_raw], dtype=np.float32))[0]) # [1] + right_open = float(convert_gripper_state_to_open_fraction(np.asarray([right_raw], dtype=np.float32))[0]) # [1] + for opening, joint_mimics in ( + (left_open, AGIBOT_WORLD_LEFT_GRIPPER_JOINT_MIMICS), + (right_open, AGIBOT_WORLD_RIGHT_GRIPPER_JOINT_MIMICS), + ): + primary_angle = -float(np.clip(opening, 0.0, 1.0)) * AGIBOT_WORLD_GRIPPER_OPEN_ANGLE_RAD + for joint_name, multiplier, offset in joint_mimics: + joint_values[joint_name] = multiplier * primary_angle + offset + + +def compute_fk_transforms( + state: np.ndarray, + embodiment_type: str, +) -> dict[str, np.ndarray]: + """Compute native-frame calibrated head-camera and gripper-base transforms for one state.""" + + fk_engine = _get_fk_engine() + link_poses = fk_engine.link_poses(_extract_joint_values_from_state(state, embodiment_type)) + + return { + "head_camera": link_poses[AGIBOT_WORLD_HEAD_CAMERA_LINK_NAME].astype(np.float32, copy=False), + "right_wrist": link_poses[AGIBOT_WORLD_RIGHT_EE_LINK_NAME].astype(np.float32, copy=False), + "left_wrist": link_poses[AGIBOT_WORLD_LEFT_EE_LINK_NAME].astype(np.float32, copy=False), + } + + +def compute_fk_transforms_batch( + states: np.ndarray, + embodiment_type: str, +) -> dict[str, np.ndarray]: + """Compute absolute transforms for a batch of AgiBot observation states.""" + + num_steps = int(states.shape[0]) + head_camera = np.empty((num_steps, 4, 4), dtype=np.float32) + right_wrist = np.empty((num_steps, 4, 4), dtype=np.float32) + left_wrist = np.empty((num_steps, 4, 4), dtype=np.float32) + + for step in range(num_steps): + transforms = compute_fk_transforms(states[step], embodiment_type) + head_camera[step] = transforms["head_camera"] + right_wrist[step] = transforms["right_wrist"] + left_wrist[step] = transforms["left_wrist"] + + transforms_by_name = { + "head_camera": head_camera, + "right_wrist": right_wrist, + "left_wrist": left_wrist, + } + return _apply_ext_base_motion_to_poses(transforms_by_name, states, embodiment_type) + + +def convert_gripper_state_to_open_fraction(values: np.ndarray) -> np.ndarray: + """Convert observed AgiBot gripper state to viewer/dataset open fractions. + + The shared viewer/action convention is ``0=closed`` and ``1=open``. + Observed AgiBot gripper state uses actuator-close angle units: ``0`` is + open and ``120`` is closed. Some episodes contain small closed-state + overshoot above ``120``; those values are accepted and clipped to fully + closed. Small open-state sensor jitter such as ``0.217`` must therefore + remain nearly fully open, not be interpreted as a normalized close fraction. + """ + + values = np.asarray(values, dtype=np.float32) + if values.size == 0: + return values + if not np.isfinite(values).all(): + raise ValueError("AgiBot gripper values contain NaN or Inf values.") + + min_value = float(np.min(values)) + max_value = float(np.max(values)) + if ( + min_value < -_GRIPPER_VALUE_EPS + and min_value >= -AGIBOT_WORLD_GRIPPER_OPEN_ANGLE_RAD - _GRIPPER_VALUE_EPS + and max_value <= _GRIPPER_VALUE_EPS + ): + return _scale_negative_to_unit_interval(values, AGIBOT_WORLD_GRIPPER_OPEN_ANGLE_RAD) + if ( + min_value < -_GRIPPER_VALUE_EPS + and min_value >= -np.degrees(AGIBOT_WORLD_GRIPPER_OPEN_ANGLE_RAD) - _GRIPPER_VALUE_EPS + and max_value <= _GRIPPER_VALUE_EPS + ): + return _scale_negative_to_unit_interval(values, np.degrees(AGIBOT_WORLD_GRIPPER_OPEN_ANGLE_RAD)) + max_actuator_value = AGIBOT_WORLD_GRIPPER_OPEN_ACTUATOR_DEG + _GRIPPER_ACTUATOR_OVERSHOOT_DEG + if min_value >= -_GRIPPER_VALUE_EPS and max_value <= max_actuator_value + _GRIPPER_VALUE_EPS: + close_fraction = _scale_to_unit_interval(values, AGIBOT_WORLD_GRIPPER_OPEN_ACTUATOR_DEG) # [*] + return (1.0 - close_fraction).astype(np.float32, copy=False) # [*] + + raise ValueError( + f"Unsupported AgiBot gripper value range; min={min_value:.4f}, max={max_value:.4f}. " + f"Expected URDF angle [-pi/4,0] or actuator-close degrees [0,{max_actuator_value:.1f}] " + f"(values above {AGIBOT_WORLD_GRIPPER_OPEN_ACTUATOR_DEG:.1f} are clipped closed)." + ) + + diff --git a/examples/worldmodel/cosmos3/vendor/cosmos_framework/data/generator/action/agibot_spec.py b/examples/worldmodel/cosmos3/vendor/cosmos_framework/data/generator/action/agibot_spec.py new file mode 100644 index 000000000000..853afd584151 --- /dev/null +++ b/examples/worldmodel/cosmos3/vendor/cosmos_framework/data/generator/action/agibot_spec.py @@ -0,0 +1,129 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: OpenMDW-1.1 + +"""Shared AgiBot metadata used by datasets and visualizers.""" + +from __future__ import annotations + +import math +from dataclasses import dataclass +from pathlib import Path +from typing import Literal + +AgibotWorldKind = Literal["gripper"] + +AGIBOT_WORLD_URDF_FILENAME = "G1_omnipicker_calibrated.urdf" +AGIBOT_WORLD_ARM_STATE_SLICE = slice(0, 14) +AGIBOT_WORLD_STATE_HEAD_YAW_IDX = 16 +AGIBOT_WORLD_STATE_HEAD_PITCH_IDX = 17 +AGIBOT_WORLD_STATE_WAIST_PITCH_IDX = 18 +AGIBOT_WORLD_STATE_WAIST_LIFT_IDX = 19 +AGIBOT_WORLD_HEAD_PITCH_JOINT_NAME = "idx04_head_pitch_joint" + +# -- Ext layout constants (94-dim state) ------------------------------------- +# The ext split stores joints at different offsets from the standard layout. +AGIBOT_WORLD_EXT_ARM_STATE_SLICE = slice(54, 68) +AGIBOT_WORLD_EXT_STATE_HEAD_YAW_IDX = 82 +AGIBOT_WORLD_EXT_STATE_HEAD_PITCH_IDX = 83 +AGIBOT_WORLD_EXT_STATE_WAIST_PITCH_IDX = 84 +AGIBOT_WORLD_EXT_STATE_WAIST_LIFT_IDX = 85 +AGIBOT_WORLD_EXT_STATE_ROBOT_POSITION_SLICE = slice(86, 89) +AGIBOT_WORLD_EXT_STATE_ROBOT_ORIENTATION_SLICE = slice(89, 93) +AGIBOT_WORLD_EXT_STATE_LEFT_HAND_SLICE = slice(0, 1) +AGIBOT_WORLD_EXT_STATE_RIGHT_HAND_SLICE = slice(1, 2) +AGIBOT_WORLD_HEAD_CAMERA_LINK_NAME = "head_camera_link" +AGIBOT_WORLD_LEFT_EE_LINK_NAME = "gripper_l_base_link" +AGIBOT_WORLD_RIGHT_EE_LINK_NAME = "gripper_r_base_link" +AGIBOT_WORLD_ARM_JOINT_NAMES_LEFT = tuple(f"idx{4 + i:02d}_left_arm_joint{i}" for i in range(1, 8)) +AGIBOT_WORLD_ARM_JOINT_NAMES_RIGHT = tuple(f"idx{11 + i:02d}_right_arm_joint{i}" for i in range(1, 8)) +AGIBOT_WORLD_WAIST_LIFT_JOINT_NAME = "idx01_waist_lift_joint" +AGIBOT_WORLD_WAIST_PITCH_JOINT_NAME = "idx02_waist_pitch_joint" +AGIBOT_WORLD_HEAD_YAW_JOINT_NAME = "idx03_head_yaw_joint" +AGIBOT_WORLD_GRIPPER_OPEN_ANGLE_RAD = math.pi / 4.0 +AGIBOT_WORLD_GRIPPER_OPEN_ACTUATOR_DEG = 120.0 +AGIBOT_WORLD_LEFT_GRIPPER_JOINT_MIMICS = ( + ("idx31_gripper_l_inner_joint1", 1.0, 0.0), + ("idx32_gripper_l_inner_joint3", 0.1, 0.0), + ("idx33_gripper_l_inner_joint4", 0.25, 0.0), + ("idx39_gripper_l_inner_joint0", -0.7, 0.0), + ("idx41_gripper_l_outer_joint1", -1.0, 0.0), + ("idx42_gripper_l_outer_joint3", 0.1, 0.0), + ("idx43_gripper_l_outer_joint4", -0.25, 0.0), + ("idx49_gripper_l_outer_joint0", 0.7, 0.0), +) +AGIBOT_WORLD_RIGHT_GRIPPER_JOINT_MIMICS = ( + ("idx71_gripper_r_inner_joint1", 1.0, 0.0), + ("idx72_gripper_r_inner_joint3", 0.1, 0.0), + ("idx73_gripper_r_inner_joint4", 0.25, 0.0), + ("idx79_gripper_r_inner_joint0", -0.7, 0.0), + ("idx81_gripper_r_outer_joint1", -1.0, 0.0), + ("idx82_gripper_r_outer_joint3", 0.1, 0.0), + ("idx83_gripper_r_outer_joint4", -0.25, 0.0), + ("idx89_gripper_r_outer_joint0", 0.7, 0.0), +) + + +@dataclass(frozen=True) +class AgibotWorldKindSpec: + """Layout metadata shared across all embodiments of one hand kind.""" + + kind: AgibotWorldKind + state_hand_slice: slice + + +@dataclass(frozen=True) +class AgibotWorldEmbodimentSpec: + """Per-embodiment metadata shared by training and visualization code.""" + + embodiment_type: str + kind: AgibotWorldKind + + +AGIBOT_WORLD_KIND_SPECS: dict[AgibotWorldKind, AgibotWorldKindSpec] = { + "gripper": AgibotWorldKindSpec( + kind="gripper", + state_hand_slice=slice(14, 16), + ), +} + +AGIBOT_WORLD_EMBODIMENT_SPECS: dict[str, AgibotWorldEmbodimentSpec] = { + "agibot_world_gripper": AgibotWorldEmbodimentSpec( + embodiment_type="agibot_world_gripper", + kind="gripper", + ), + "agibot_world_gripper_ext": AgibotWorldEmbodimentSpec( + embodiment_type="agibot_world_gripper_ext", + kind="gripper", + ), +} + + +def get_agibot_world_embodiment_spec(embodiment_type: str) -> AgibotWorldEmbodimentSpec: + """Return the registered spec for one AgiBot embodiment.""" + + try: + return AGIBOT_WORLD_EMBODIMENT_SPECS[embodiment_type] + except KeyError as exc: + raise ValueError( + f"Unknown AgiBot World embodiment_type={embodiment_type!r}. " + f"Expected one of {sorted(AGIBOT_WORLD_EMBODIMENT_SPECS)}." + ) from exc + + +def get_agibot_world_kind_spec(embodiment_type: str | AgibotWorldKind) -> AgibotWorldKindSpec: + """Resolve an embodiment type or kind to its shared layout metadata.""" + + kind = embodiment_type if embodiment_type in AGIBOT_WORLD_KIND_SPECS else get_agibot_world_kind(embodiment_type) + return AGIBOT_WORLD_KIND_SPECS[kind] + + +def get_agibot_world_kind(embodiment_type: str) -> AgibotWorldKind: + """Return the hand kind used by one AgiBot embodiment.""" + + return get_agibot_world_embodiment_spec(embodiment_type).kind + + +def get_agibot_world_urdf_path() -> Path: + """Return the committed AgiBot G1 omnipicker URDF path.""" + + return Path(__file__).resolve().parent / "robot_assets" / AGIBOT_WORLD_URDF_FILENAME diff --git a/examples/worldmodel/cosmos3/vendor/cosmos_framework/data/generator/action/datasets/__init__.py b/examples/worldmodel/cosmos3/vendor/cosmos_framework/data/generator/action/datasets/__init__.py new file mode 100644 index 000000000000..6797cbaf2322 --- /dev/null +++ b/examples/worldmodel/cosmos3/vendor/cosmos_framework/data/generator/action/datasets/__init__.py @@ -0,0 +1,39 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: OpenMDW-1.1 + +"""Action dataset wrappers for Cosmos Action. + +All concrete datasets inherit from :class:`ActionBaseDataset` and expose a +``load_action_stats()`` classmethod for retrieving pre-computed normalization +statistics without instantiating the dataset. +""" + +from cosmos_framework.data.generator.action.datasets.agibotworld_beta_lerobot_dataset import ( + AgiBotWorldBetaLeRobotDataset, +) +from cosmos_framework.data.generator.action.datasets.base_dataset import ActionBaseDataset +from cosmos_framework.data.generator.action.datasets.bridge_orig_lerobot_dataset import BridgeOrigLeRobotDataset +from cosmos_framework.data.generator.action.datasets.droid_merged_lerobot_dataset import DROIDMergedLeRobotDataset +from cosmos_framework.data.generator.action.datasets.droid_lerobot_dataset import DROIDLeRobotDataset +from cosmos_framework.data.generator.action.datasets.fractal_lerobot_dataset import FractalLeRobotDataset +from cosmos_framework.data.generator.action.datasets.human_hand_pose_lerobot_dataset import HumanHandPoseLeRobotDataset +from cosmos_framework.data.generator.action.datasets.robocasa_lerobot_dataset import RoboCasaLeRobotDataset +from cosmos_framework.data.generator.action.datasets.libero_lerobot_dataset import LIBEROLeRobotDataset +from cosmos_framework.data.generator.action.datasets.robomind_franka_dataset import RoboMINDFrankaDataset +from cosmos_framework.data.generator.action.datasets.robomind_ur_dataset import RoboMINDURDataset +from cosmos_framework.data.generator.action.datasets.umi_lerobot_dataset import UMILeRobotDataset + +__all__ = [ + "ActionBaseDataset", + "AgiBotWorldBetaLeRobotDataset", + "BridgeOrigLeRobotDataset", + "DROIDLeRobotDataset", + "DROIDMergedLeRobotDataset", + "FractalLeRobotDataset", + "HumanHandPoseLeRobotDataset", + "RoboCasaLeRobotDataset", + "LIBEROLeRobotDataset", + "RoboMINDFrankaDataset", + "RoboMINDURDataset", + "UMILeRobotDataset", +] diff --git a/examples/worldmodel/cosmos3/vendor/cosmos_framework/data/generator/action/datasets/agibotworld_beta_lerobot_dataset.py b/examples/worldmodel/cosmos3/vendor/cosmos_framework/data/generator/action/datasets/agibotworld_beta_lerobot_dataset.py new file mode 100644 index 000000000000..633c1988cf28 --- /dev/null +++ b/examples/worldmodel/cosmos3/vendor/cosmos_framework/data/generator/action/datasets/agibotworld_beta_lerobot_dataset.py @@ -0,0 +1,279 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: OpenMDW-1.1 + +"""AgiBotWorld-Beta LeRobot dataset.""" + +from __future__ import annotations + +from pathlib import Path +from typing import Any, Literal + +import numpy as np +import torch +import torch.nn.functional as F + +from cosmos_framework.data.generator.action.agibot_fk import ( + AGIBOT_WORLD_GRIPPER_TO_OPENCV_BY_WRIST, + apply_agibot_gripper_to_opencv, + apply_robot_base_motion_to_poses, + compute_fk_transforms_batch, + convert_gripper_state_to_open_fraction, +) +from cosmos_framework.data.generator.action.utils.action_spec import ActionSpec, Gripper, Pos, Rot, build_action_spec +from cosmos_framework.data.generator.action.datasets.base_dataset import ActionBaseDataset +from cosmos_framework.data.generator.action.utils.pose_utils import pose_abs_to_rel + +PoseConvention = Literal["backward_framewise"] +Viewpoint = Literal["concat_view", "ego_view"] + +_HEAD_KEY = "observation.images.head" +_HAND_LEFT_KEY = "observation.images.hand_left" +_HAND_RIGHT_KEY = "observation.images.hand_right" +_CONCAT_KEY = "observation.images.video_concat_view" + +_EFFECTOR_KEY = "observation.states.effector.position" +_JOINT_KEY = "observation.states.joint.position" +_HEAD_STATE_KEY = "observation.states.head.position" +_WAIST_KEY = "observation.states.waist.position" +_ROBOT_POSITION_KEY = "observation.states.robot.position" +_ROBOT_ORIENTATION_KEY = "observation.states.robot.orientation" + +_NORMALIZER_PATH = Path(__file__).parent.parent / "normalizer_stats/agibotworld_beta_lerobot_stats.json" + + +def _split_task_for_caption(task: str) -> tuple[str, str]: + ai_caption, separator, debug_caption = task.partition("|") + if not separator: + return task.strip(), "" + return ai_caption.strip(), debug_caption.strip() + + +def _assemble_agibot_world_state( + effector_pos: np.ndarray, + joint_pos: np.ndarray, + head_pos: np.ndarray, + waist_pos: np.ndarray, +) -> np.ndarray: + """Assemble standard 20D gripper state from Beta decomposed fields.""" + + body_head = np.stack( + [head_pos[:, 0], head_pos[:, 1], waist_pos[:, 0], waist_pos[:, 1]], + axis=-1, + ) + return np.concatenate([joint_pos, effector_pos, body_head], axis=-1).astype(np.float32, copy=False) + + +def _compute_idle_frames_agibot(action: torch.Tensor) -> int: + """Small local idle-frame helper for the 29D AgiBot FK layout. + + The shared `compute_idle_frames` expects one rotation group after each + position block; AgiBot's action spec has three such groups plus grippers. + For cookbook inference, idle frames are metadata only, so this conservative + implementation marks the initial low-motion streak length. + """ + + if action.numel() == 0: + return 0 + abs_action = action.detach().abs() + motion = torch.cat( + [ + abs_action[:, 0:3], + abs_action[:, 9:12], + abs_action[:, 18:21], + abs_action[:, 18:19].diff(dim=0, prepend=abs_action[0:1, 18:19]), + abs_action[:, 28:29].diff(dim=0, prepend=abs_action[0:1, 28:29]), + ], + dim=-1, + ).amax(dim=-1) + below = motion < 1e-3 + count = 0 + for value in below.tolist(): + if not value: + break + count += 1 + return count + + +class AgiBotWorldBetaLeRobotDataset(ActionBaseDataset): + """AgiBotWorld-Beta dataset with FK-pose 29D actions. + + Action layout matches the AgiBot World gripper normalizer: + + [head_pos+rot6d(9), right_pos+rot6d(9), right_gripper(1), + left_pos+rot6d(9), left_gripper(1)] + + The local cookbook asset provides head, left wrist, and right wrist videos. + By default this wrapper uses `concat_view`: head view on top, left/right + wrist views resized and concatenated on the bottom. + """ + + + def __init__( + self, + root: str, + fps: float = 10.0, + chunk_length: int = 16, + mode: str = "joint", + pose_convention: PoseConvention = "backward_framewise", + tolerance_s: float = 3e-4, + viewpoint: Viewpoint = "concat_view", + action_normalization: str | None = "quantile", + sample_stride: int = 1, + ) -> None: + if viewpoint not in ("concat_view", "ego_view"): + raise NotImplementedError("Supported viewpoints are concat_view and ego_view.") + super().__init__( + root=root, + domain_name="agibotworld", + fps=fps, + chunk_length=chunk_length, + mode=mode, + pose_convention=pose_convention, + tolerance_s=tolerance_s, + viewpoint=viewpoint, + action_normalization=action_normalization, + sample_stride=sample_stride, + ) + self._rows_by_episode: dict[int, list[dict[str, Any]]] = {} + for row in self._rows: + self._rows_by_episode.setdefault(int(row["episode_index"]), []).append(row) + self._timestamps_by_episode = { + episode_id: np.asarray([float(row["timestamp"]) for row in rows], dtype=np.float64) + for episode_id, rows in self._rows_by_episode.items() + } + + @property + def action_dim(self) -> int: + return 29 + + def _action_spec(self) -> ActionSpec: + return build_action_spec( + Pos(prefix="head"), + Rot("rot6d", prefix="head"), + Pos(prefix="right"), + Rot("rot6d", prefix="right"), + Gripper(prefix="right"), + Pos(prefix="left"), + Rot("rot6d", prefix="left"), + Gripper(prefix="left"), + ) + + @classmethod + def _stats_path(cls) -> Path: + return _NORMALIZER_PATH + + def _compute_idle_frames(self, action: torch.Tensor) -> int: + return _compute_idle_frames_agibot(action) + + def __len__(self) -> int: + return max(0, (len(self._rows) - self._chunk_length + self._sample_stride - 1) // self._sample_stride) + + def __getitem__(self, idx: int) -> dict[str, Any]: + mode = self._choose_mode() + row_idx = int(idx) * self._sample_stride + start_row = self._rows[row_idx] + observation_rows = self._select_observation_rows(start_row) + episode = self._episodes[int(observation_rows[0]["episode_index"])] + task = self._tasks[int(observation_rows[0]["task_index"])] + ai_caption, debug_caption = _split_task_for_caption(task) + + video = self._load_video(episode, observation_rows) + action, extras = self._build_fk_action(observation_rows) + if self._viewpoint == "concat_view": + extras["additional_view_description"] = ( + "The top row shows the head-mounted camera view looking down at the workspace. " + "The bottom row contains two horizontally concatenated wrist-mounted camera views: " + "the left hand camera on the left and the right hand camera on the right." + ) + if debug_caption: + extras["debug_caption"] = debug_caption + + return self._build_result( + mode=mode, + video=video, + action=action, + ai_caption=ai_caption, + action_spec_names=self.action_names, + **extras, + ) + + def _select_observation_rows(self, start_row: dict[str, Any]) -> list[dict[str, Any]]: + """Select T+1 rows at this wrapper's target FPS within one episode.""" + + episode_id = int(start_row["episode_index"]) + rows = self._rows_by_episode[episode_id] + timestamps = self._timestamps_by_episode[episode_id] + start_frame = int(start_row["frame_index"]) + start_ts = float(start_row["timestamp"]) + target_ts = start_ts + np.arange(self._chunk_length + 1, dtype=np.float64) / self._fps + indices = np.searchsorted(timestamps, target_ts, side="left") + indices = np.minimum(indices, len(rows) - 1) + prev = np.maximum(indices - 1, 0) + choose_prev = np.abs(timestamps[prev] - target_ts) <= np.abs(timestamps[indices] - target_ts) + indices = np.where(choose_prev, prev, indices) + if int(indices[-1]) <= start_frame: + raise IndexError(f"Could not select {self._chunk_length + 1} frames from episode {episode_id} at fps={self._fps}.") + return [rows[int(i)] for i in indices] + + def _load_video(self, episode: dict[str, Any], observation_rows: list[dict[str, Any]]) -> torch.Tensor: + if self._viewpoint == "ego_view": + return self._load_video_key(episode, observation_rows, _HEAD_KEY) + + # Prefer a pre-rendered concat view if present. The local asset includes + # metadata for this key but not the public mp4, so the fallback composes + # it from the three camera streams. + concat_path = self._video_path(episode, _CONCAT_KEY) + if concat_path.exists(): + return self._load_video_key(episode, observation_rows, _CONCAT_KEY) + top = self._load_video_key(episode, observation_rows, _HEAD_KEY) + left = self._load_video_key(episode, observation_rows, _HAND_LEFT_KEY) + right = self._load_video_key(episode, observation_rows, _HAND_RIGHT_KEY) + return self._compose_multi_view(top, left, right) + + def _load_video_key(self, episode: dict[str, Any], observation_rows: list[dict[str, Any]], key: str) -> torch.Tensor: + # lerobot is a heavy, optional ("train" extra) dependency; import lazily. + from lerobot.datasets.video_utils import decode_video_frames + + timestamps = [float(row["timestamp"]) for row in observation_rows] + return decode_video_frames( + self._video_path(episode, key), + [float(episode.get(f"videos/{key}/from_timestamp", 0.0)) + ts for ts in timestamps], + self._tolerance_s, + ) + + def _compose_multi_view(self, top: torch.Tensor, left: torch.Tensor, right: torch.Tensor) -> torch.Tensor: + # Inputs are [T,C,H,W] float tensors in [0,1]. + _, _, h_top, w_top = top.shape + half_h, half_w = h_top // 2, w_top // 2 + left = F.interpolate(left, size=(half_h, half_w), mode="bilinear", align_corners=False) + right = F.interpolate(right, size=(half_h, half_w), mode="bilinear", align_corners=False) + bottom = torch.cat([left, right], dim=-1) + return torch.cat([top, bottom], dim=-2) + + def _build_fk_action(self, rows: list[dict[str, Any]]) -> tuple[torch.Tensor, dict[str, torch.Tensor]]: + effector_pos = np.asarray([row[_EFFECTOR_KEY] for row in rows], dtype=np.float32) + joint_pos = np.asarray([row[_JOINT_KEY] for row in rows], dtype=np.float32) + head_pos = np.asarray([row[_HEAD_STATE_KEY] for row in rows], dtype=np.float32) + waist_pos = np.asarray([row[_WAIST_KEY] for row in rows], dtype=np.float32) + robot_pos = np.asarray([row[_ROBOT_POSITION_KEY] for row in rows], dtype=np.float32) + robot_quat = np.asarray([row[_ROBOT_ORIENTATION_KEY] for row in rows], dtype=np.float32) + states_np = _assemble_agibot_world_state(effector_pos, joint_pos, head_pos, waist_pos) + + native_fk = compute_fk_transforms_batch(states_np, "agibot_world_gripper") + native_fk = apply_robot_base_motion_to_poses(native_fk, robot_pos, robot_quat) + fk = apply_agibot_gripper_to_opencv(native_fk, AGIBOT_WORLD_GRIPPER_TO_OPENCV_BY_WRIST) + + head_rel = pose_abs_to_rel(fk["head_camera"], rotation_format="rot6d", pose_convention=self._pose_convention) + right_rel = pose_abs_to_rel(fk["right_wrist"], rotation_format="rot6d", pose_convention=self._pose_convention) + left_rel = pose_abs_to_rel(fk["left_wrist"], rotation_format="rot6d", pose_convention=self._pose_convention) + right_gripper = convert_gripper_state_to_open_fraction(effector_pos[1:, 1:2]) + left_gripper = convert_gripper_state_to_open_fraction(effector_pos[1:, 0:1]) + action_np = np.concatenate([head_rel, right_rel, right_gripper, left_rel, left_gripper], axis=-1).astype( + np.float32 + ) + extras = { + "initial_pose": torch.from_numpy(fk["head_camera"][0].copy()).float(), + "initial_pose_right": torch.from_numpy(fk["right_wrist"][0].copy()).float(), + "initial_pose_left": torch.from_numpy(fk["left_wrist"][0].copy()).float(), + } + return torch.from_numpy(action_np).float(), extras diff --git a/examples/worldmodel/cosmos3/vendor/cosmos_framework/data/generator/action/datasets/base_dataset.py b/examples/worldmodel/cosmos3/vendor/cosmos_framework/data/generator/action/datasets/base_dataset.py new file mode 100644 index 000000000000..bec2c5c00a24 --- /dev/null +++ b/examples/worldmodel/cosmos3/vendor/cosmos_framework/data/generator/action/datasets/base_dataset.py @@ -0,0 +1,236 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: OpenMDW-1.1 + +"""Abstract base class for Action LeRobot datasets.""" + +from __future__ import annotations + +import json +import random +from abc import ABC, abstractmethod +from pathlib import Path +from typing import Any + +import numpy as np +import pandas as pd +import pyarrow.parquet as pq +import torch +from torch.utils.data import Dataset + +from cosmos_framework.data.generator.action.action_normalization import load_action_stats, normalize_action +from cosmos_framework.data.generator.action.utils.action_spec import ActionSpec +from cosmos_framework.data.generator.action.utils.domain_utils import get_domain_id +from cosmos_framework.data.generator.action.utils.pose_utils import compute_idle_frames + +_MODE_CHOICES = ("forward_dynamics", "inverse_dynamics", "wam") + + +class ActionBaseDataset(ABC, Dataset): + """Abstract base for Action LeRobot datasets. + + Subclasses must implement the abstract methods listed below. + """ + + def __init__( + self, + root: str, + domain_name: str, + fps: float, + chunk_length: int, + mode: str, + pose_convention: str, + tolerance_s: float, + viewpoint: str, + action_normalization: str | None = "quantile", + sample_stride: int = 1, + ) -> None: + super().__init__() + if pose_convention != "backward_framewise": + raise NotImplementedError(f"{type(self).__name__} only supports backward_framewise pose deltas.") + + self._fps = float(fps) + self._dt = 1.0 / self._fps + self._chunk_length = int(chunk_length) + self._sample_stride = int(sample_stride) + if self._sample_stride < 1: + raise ValueError(f"sample_stride must be >= 1, got {self._sample_stride}") + self._mode = mode + self._pose_convention = pose_convention + self._tolerance_s = float(tolerance_s) + self._viewpoint = viewpoint + self._domain_name = domain_name + self._domain_id = get_domain_id(domain_name) + self._action_normalization = action_normalization + self._norm_stats: dict[str, torch.Tensor] | None = None + + self._root = Path(root) + self._info = json.loads((self._root / "meta" / "info.json").read_text()) + self._episodes = { + int(row["episode_index"]): row + for path in sorted((self._root / "meta" / "episodes").glob("chunk-*/file-*.parquet")) + for row in pq.read_table(path).to_pylist() + } + tasks_df = pd.read_parquet(self._root / "meta" / "tasks.parquet") + # 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`` (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: + return self._fps + + @property + def chunk_length(self) -> int: + return self._chunk_length + + @property + def mode(self) -> str: + return self._mode + + @mode.setter + def mode(self, value: str) -> None: + self._mode = value + + @property + def domain_name(self) -> str: + return self._domain_name + + @property + def viewpoint(self) -> str: + return self._viewpoint + + @property + def domain_id(self) -> int: + return self._domain_id + + @property + def action_normalization(self) -> str: + return self._action_normalization + + @property + @abstractmethod + def action_dim(self) -> int: ... + + @abstractmethod + def _action_spec(self) -> ActionSpec: ... + + @property + def action_names(self) -> list[str]: + return self._action_spec().names + + @classmethod + @abstractmethod + def _stats_path(cls) -> Path: + """Return the path to the stats JSON file for this dataset.""" + ... + + @classmethod + 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() + } + + @abstractmethod + def __getitem__(self, idx: int) -> dict[str, Any]: ... + + def _compute_idle_frames(self, action: torch.Tensor) -> int: + return compute_idle_frames( + action, + self._action_spec(), + eps_t=5e-3 / self._fps, + eps_r=np.deg2rad(1.5) / self._fps, + eps_g=1e-2, + joint_threshold=5e-3 / self._fps, + min_streak=3, + ) + + def _choose_mode(self) -> str: + if self._mode == "joint": + return random.choice(_MODE_CHOICES) + return self._mode + + def _video_path(self, episode: dict[str, Any], video_key: str) -> Path: + chunk_idx = int( + episode.get( + f"videos/{video_key}/chunk_index", + episode.get(f"videos/{video_key}/episode_chunk", episode.get("data/chunk_index", 0)), + ) + ) + file_idx = int( + episode.get( + f"videos/{video_key}/file_index", + episode.get(f"videos/{video_key}/episode_file", episode.get("data/file_index", 0)), + ) + ) + rel = self._info["video_path"].format( + video_key=video_key, + chunk_index=chunk_idx, + file_index=file_idx, + episode_chunk=chunk_idx, + episode_file=file_idx, + ) + return self._root / rel + + def _load_norm_stats(self) -> dict[str, torch.Tensor]: + if self._norm_stats is None: + self._norm_stats = self.load_action_stats() + return self._norm_stats + + def _build_result( + self, + *, + mode: str, + video: torch.Tensor, + action: torch.Tensor, + ai_caption: str, + **extras: Any, + ) -> dict[str, Any]: + idle_frames = self._compute_idle_frames(action) + # action_normalization=None -> use raw actions (no normalization), e.g. joint_pos. + if self.action_normalization is None: + normalized_action = action + else: + normalized_action = normalize_action(action, self.action_normalization, self._load_norm_stats()) + formatted_video = (video * 255.0).clamp(0.0, 255.0).to(torch.uint8).permute(1, 0, 2, 3) + return { + "ai_caption": ai_caption, + "video": formatted_video, + "action": normalized_action, + "conditioning_fps": torch.tensor(self._fps, dtype=torch.long), + "mode": mode, + "domain_id": torch.tensor(self._domain_id, dtype=torch.long), + "viewpoint": self._viewpoint, + "idle_frames": torch.tensor(idle_frames, dtype=torch.long), + **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) diff --git a/examples/worldmodel/cosmos3/vendor/cosmos_framework/data/generator/action/datasets/bridge_orig_lerobot_dataset.py b/examples/worldmodel/cosmos3/vendor/cosmos_framework/data/generator/action/datasets/bridge_orig_lerobot_dataset.py new file mode 100644 index 000000000000..0a35e699513d --- /dev/null +++ b/examples/worldmodel/cosmos3/vendor/cosmos_framework/data/generator/action/datasets/bridge_orig_lerobot_dataset.py @@ -0,0 +1,154 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: OpenMDW-1.1 + +"""Bridge Orig LeRobot dataset.""" + +from __future__ import annotations + +import random +from pathlib import Path +from typing import Any, Literal + +import numpy as np +import torch + +from cosmos_framework.data.generator.action.utils.action_spec import ActionSpec, Gripper, Pos, Rot, build_action_spec +from cosmos_framework.data.generator.action.datasets.base_dataset import ActionBaseDataset +from cosmos_framework.data.generator.action.utils.pose_utils import ( + build_abs_pose_from_components, + pose_abs_to_rel, +) + +PoseConvention = Literal["backward_framewise"] +Viewpoint = Literal["ego_view"] + +_IMAGE_FEATURE = "observation.images.image_0" +_STATE_FEATURE = "observation.state" +_ACTION_FEATURE = "action" + +# Raw Bridge state -> kinematics frame. The WidowX controller records +# R_state = R_fk @ DEFAULT_ROTATION.T, so R_fk = R_state @ DEFAULT_ROTATION. +_DEFAULT_ROTATION = np.array( + [[0.0, 0.0, 1.0], [0.0, 1.0, 0.0], [-1.0, 0.0, 0.0]], + dtype=np.float32, +) + +# Kinematics frame -> OpenCV frame used by Cosmos action. +_BRIDGE_TO_OPENCV = np.array( + [[0.0, 0.0, 1.0], [-1.0, 0.0, 0.0], [0.0, -1.0, 0.0]], + dtype=np.float32, +) + +# Re-reference from ee_gripper_link to gripper_link in the kinematics frame. +_TCP_TO_FLANGE = np.array( + [ + [1.0, 0.0, 0.0, -0.093575], + [0.0, 1.0, 0.0, 0.0], + [0.0, 0.0, 1.0, 0.0], + [0.0, 0.0, 0.0, 1.0], + ], + dtype=np.float32, +) + +_NORMALIZER_PATH = Path(__file__).parent.parent / "normalizer_stats/bridge_orig_lerobot_stats.json" + + +class BridgeOrigLeRobotDataset(ActionBaseDataset): + """Bridge Orig dataset with 10D cartesian actions: + + [pos_delta(3), rot6d_delta(6), gripper(1)] + + Uses a single ``image_0`` ego-view video, backward-framewise rot6d actions, + and quantile normalization. + """ + + + def __init__( + self, + root: str, + fps: float = 5.0, + chunk_length: int = 16, + mode: str = "joint", + pose_convention: PoseConvention = "backward_framewise", + tolerance_s: float = 1e-4, + viewpoint: Viewpoint = "ego_view", + action_normalization: str | None = "quantile", + sample_stride: int = 1, + ) -> None: + if viewpoint != "ego_view": + raise NotImplementedError("This minimal Bridge dataset only supports ego_view.") + super().__init__( + root=root, + domain_name="bridge_orig_lerobot", + fps=fps, + chunk_length=chunk_length, + mode=mode, + pose_convention=pose_convention, + tolerance_s=tolerance_s, + viewpoint=viewpoint, + action_normalization=action_normalization, + sample_stride=sample_stride, + ) + + @property + def action_dim(self) -> int: + return 10 + + def _action_spec(self) -> ActionSpec: + return build_action_spec(Pos(), Rot("rot6d"), Gripper()) + + @classmethod + def _stats_path(cls) -> Path: + return _NORMALIZER_PATH + + def __getitem__(self, idx: int) -> dict[str, Any]: + mode = self._choose_mode() + idx = int(idx) + first_row = self._rows[idx] + episode = self._episodes[int(first_row["episode_index"])] + + row_idx = idx * self._sample_stride + observation_rows = self._rows[row_idx : row_idx + self._chunk_length + 1] + action_rows = observation_rows[: self._chunk_length] + + video = self._load_video(episode, observation_rows) + raw_action, initial_pose = self._build_raw_action(observation_rows, action_rows) + task = self._tasks[int(observation_rows[0]["task_index"])] + ai_caption = random.choice([part.strip() for part in task.split(" | ") if part.strip()] or [task]) + + return self._build_result( + mode=mode, + video=video, + action=raw_action, + ai_caption=ai_caption, + initial_pose=initial_pose, + ) + + def _load_video(self, episode: dict[str, Any], observation_rows: list[dict[str, Any]]) -> torch.Tensor: + # lerobot is a heavy, optional ("train" extra) dependency; import lazily. + from lerobot.datasets.video_utils import decode_video_frames + + timestamps = [float(row["timestamp"]) for row in observation_rows] + return decode_video_frames( + self._video_path(episode, _IMAGE_FEATURE), + [float(episode.get(f"videos/{_IMAGE_FEATURE}/from_timestamp", 0.0)) + ts for ts in timestamps], + self._tolerance_s, + ) + + def _build_raw_action( + self, + observation_rows: list[dict[str, Any]], + action_rows: list[dict[str, Any]], + ) -> tuple[torch.Tensor, torch.Tensor]: + state = np.asarray([row[_STATE_FEATURE] for row in observation_rows], dtype=np.float32) + poses_abs = build_abs_pose_from_components(state[:, 0:3], state[:, 3:6], "euler_xyz") + + poses_abs[:, :3, :3] = poses_abs[:, :3, :3] @ _DEFAULT_ROTATION.astype(poses_abs.dtype) + poses_abs = poses_abs @ _TCP_TO_FLANGE.astype(poses_abs.dtype) + poses_abs[:, :3, :3] = poses_abs[:, :3, :3] @ _BRIDGE_TO_OPENCV.astype(poses_abs.dtype) + + initial_pose = torch.from_numpy(poses_abs[0].copy()).float() + poses_rel = pose_abs_to_rel(poses_abs, rotation_format="rot6d", pose_convention=self._pose_convention) + gripper = np.asarray([row[_ACTION_FEATURE][6] for row in action_rows], dtype=np.float32).reshape(-1, 1) + action = np.concatenate([poses_rel[-self._chunk_length :], gripper[-self._chunk_length :]], axis=-1) + return torch.from_numpy(action).float(), initial_pose diff --git a/examples/worldmodel/cosmos3/vendor/cosmos_framework/data/generator/action/datasets/cosmos3_action_lerobot.py b/examples/worldmodel/cosmos3/vendor/cosmos_framework/data/generator/action/datasets/cosmos3_action_lerobot.py new file mode 100644 index 000000000000..8d810e3cb8e1 --- /dev/null +++ b/examples/worldmodel/cosmos3/vendor/cosmos_framework/data/generator/action/datasets/cosmos3_action_lerobot.py @@ -0,0 +1,1060 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: OpenMDW-1.1 + +"""Shared LeRobot adapter utilities for Action datasets. + +These helpers centralize common behavior across Action wrappers: +- deterministic train/val episode splitting +- valid per-episode index range construction +- a reusable BaseActionLeRobotDataset class with lazy init, video formatting, + and common result building +""" + +from __future__ import annotations + +import importlib +import logging as _logging +import math +import os as _os +import random +from bisect import bisect_right +from collections import OrderedDict, defaultdict +from collections.abc import Callable, Sequence +from concurrent.futures import ThreadPoolExecutor +from contextlib import contextmanager +from pathlib import Path +from threading import Lock +from typing import Any, ClassVar + +import huggingface_hub.constants as _hf_const +import numpy as np +import torch +from lerobot.datasets.lerobot_dataset import LeRobotDataset, LeRobotDatasetMetadata +from torch.utils.data import Dataset + +_hf_offline_applied = False + + +def _ensure_hf_hub_offline() -> None: + """Force HF Hub into offline mode for local-only datasets (repo_id="local"). + + Sets the ``HF_HUB_OFFLINE`` env var (for any future imports in worker + processes), patches the already-imported constant, and suppresses the + expected "Returning existing local_dir" fallback warning. + + Safe to call multiple times; only applies once per process. + """ + global _hf_offline_applied + if _hf_offline_applied: + return + if "HF_HUB_OFFLINE" not in _os.environ: + _os.environ["HF_HUB_OFFLINE"] = "1" + if not _hf_const.HF_HUB_OFFLINE: + _hf_const.HF_HUB_OFFLINE = True + _logging.getLogger("huggingface_hub._snapshot_download").setLevel(_logging.ERROR) + _hf_offline_applied = True + + +from functools import cached_property + +from cosmos_framework.data.generator.action.utils.action_processing import ( + ActionNormalizationMethod, + ActionNormalizer, + load_action_stats, + resolve_action_normalization, +) + +# Re-export the action_spec DSL from this module so that subclass datasets +# only need a single import block (alongside ``BaseActionLeRobotDataset``). +from cosmos_framework.data.generator.action.utils.action_spec import ( # noqa: F401 (re-export) + ActionSpec, + DimType, + Gripper, + Joint, + Pos, + Reserved, + Rot, + build_action_spec, +) +from cosmos_framework.data.generator.action.utils.domain_utils import get_domain_id +from cosmos_framework.data.generator.action.utils.pose_utils import compute_idle_frames +from cosmos_framework.data.generator.action.utils.viewpoint_utils import Viewpoint +from cosmos_framework.utils import log + + +# No-op memory-profiling shims. Profiling is disabled in cosmos-framework, so these +# keep the ported dataset's optional RSS-tracking call sites cheap and side-effect-free +# (formerly a separate memprofile stub module). +def _memprofile_enabled() -> bool: + return False + + +def _deep_size(obj, *args, **kwargs) -> int: + return 0 + + +def _fmt_mb(n, *args, **kwargs) -> str: + return "n/a" + + +def log_worker_memory_breakdown(*args, **kwargs) -> None: + return None + + +@contextmanager +def rss_tracker(*args, **kwargs): + yield + + +# --------------------------------------------------------------------------- +# LRU-capped VideoDecoderCache +# --------------------------------------------------------------------------- +_LRU_VIDEO_CACHE_MAX_SIZE: int = 64 +_LRU_DATASET_MAX_LOADED: int = 32 +ActionNormalization = ActionNormalizationMethod +_ACTION_NORMALIZATION_CHOICES: tuple[str, ...] = ("quantile", "quantile_rot", "meanstd", "minmax") + +_decoder_cache_patched = False + + +class _LRUVideoDecoderCache: + """Drop-in replacement for ``lerobot.datasets.video_utils.VideoDecoderCache`` + with LRU eviction. When the cache exceeds *max_size* entries the + least-recently-used decoder (and its file handle) is evicted. + """ + + def __init__(self, max_size: int = _LRU_VIDEO_CACHE_MAX_SIZE) -> None: + self._max_size = max_size + self._cache: OrderedDict[str, tuple[Any, Any]] = OrderedDict() + self._lock = Lock() + self._hits = 0 + self._misses = 0 + self._evictions = 0 + + def get_decoder(self, video_path: str) -> Any: + if importlib.util.find_spec("torchcodec"): # type: ignore[attr-defined] + from torchcodec.decoders import VideoDecoder + else: + raise ImportError("torchcodec is required but not available.") + + import fsspec + + video_path = str(video_path) + + with self._lock: + if video_path in self._cache: + self._cache.move_to_end(video_path) + self._hits += 1 + return self._cache[video_path][0] + + self._misses += 1 + file_handle = fsspec.open(video_path).__enter__() + decoder = VideoDecoder(file_handle, seek_mode="approximate") # type: ignore[arg-type] + self._cache[video_path] = (decoder, file_handle) + + evicted = 0 + while len(self._cache) > self._max_size: + _, (_, old_fh) = self._cache.popitem(last=False) + try: + old_fh.close() + except Exception: + pass + evicted += 1 + self._evictions += evicted + + if evicted and self._evictions % 50 <= evicted: + log.debug( + f"[VideoDecoderCache pid={_os.getpid()}] " + f"evicted={self._evictions} total, size={len(self._cache)}/{self._max_size}, " + f"hits={self._hits}, misses={self._misses}, " + f"hit_rate={100 * self._hits / max(1, self._hits + self._misses):.1f}%" + ) + + return decoder + + def clear(self) -> None: + with self._lock: + for _, file_handle in self._cache.values(): + try: + file_handle.close() + except Exception: + pass + self._cache.clear() + + def size(self) -> int: + with self._lock: + return len(self._cache) + + +def _patch_decoder_cache(max_size: int = _LRU_VIDEO_CACHE_MAX_SIZE) -> None: + """Replace the module-level ``_default_decoder_cache`` in LeRobot with an + LRU-capped version to prevent unbounded memory growth in workers.""" + global _decoder_cache_patched + if _decoder_cache_patched: + return + + import lerobot.datasets.video_utils as _vu + + lru_cache = _LRUVideoDecoderCache(max_size=max_size) + _vu._default_decoder_cache = lru_cache + _decoder_cache_patched = True + log.debug(f"Patched LeRobot VideoDecoderCache with LRU max_size={max_size}") + + +def _parallel_map( + fn: Callable[[Any], Any], + items: list[Any], + *, + max_workers: int, + label: str, +) -> list[Any]: + """Thread-pool ``map`` — returns results in input order. + + Intended for IO-bound prefetch (``LeRobotDatasetMetadata`` loads, + parquet column reads). Preserves item-order so callers can ``zip`` + with their ``indices`` / ``roots`` list. Skips the thread pool + entirely when there is 0 or 1 task — avoids per-worker + ``ThreadPoolExecutor`` setup cost and log spam under + ``shard_across_workers=True`` where each worker typically gets + only 1-2 shards. + """ + if not items: + return [] + if len(items) == 1 or max_workers <= 1: + return [fn(items[0])] if len(items) == 1 else [fn(x) for x in items] + log.info(f"{label}: {len(items)} tasks (workers={max_workers})") + with ThreadPoolExecutor(max_workers=max_workers) as ex: + return list(ex.map(fn, items)) + + +def split_episode_ids(total_episodes: int, seed: int, val_ratio: float, split: str) -> list[int]: + """Create deterministic random episode ids for train/val/full splits.""" + num_val = int(round(total_episodes * val_ratio)) + g = torch.Generator().manual_seed(seed) + episode_ids = torch.randperm(total_episodes, generator=g).tolist() + + if split == "train": + return episode_ids[num_val:] + if split == "val": + return episode_ids[:num_val] + return episode_ids + + +def build_episode_spans( + episodes: Any, + episode_ids: Sequence[int], + chunk_length: int, + sample_stride: int = 1, +) -> tuple[list[tuple[int, int, int]], int, int]: + """Build valid episode spans for LeRobot frame queries. + + Returns: + - episode spans as ``(episode_id, sample_start, valid_len)`` + - total valid sample count across selected episodes + - total raw frame count across selected episodes + """ + assert sample_stride >= 1, f"sample_stride must be >= 1, got {sample_stride}" + + dataset_from_index = list(episodes["dataset_from_index"]) + dataset_to_index = list(episodes["dataset_to_index"]) + length = list(episodes["length"]) + + spans: list[tuple[int, int, int]] = [] + valid_count = 0 + sample_count = 0 + for episode_id in episode_ids: + start = dataset_from_index[episode_id] + stop = dataset_to_index[episode_id] + raw_valid_len = stop - start - chunk_length + if raw_valid_len > 0: + valid_len = (raw_valid_len + sample_stride - 1) // sample_stride + spans.append((episode_id, start, valid_len)) + valid_count += valid_len + sample_count += int(length[episode_id]) + + return spans, valid_count, sample_count + + +def _normalize_split(split: str) -> str: + """Normalize split name to one of ``'train'``, ``'val'``, ``'full'``.""" + s = split.lower().strip() + if s in {"val", "valid", "validation", "eval", "test"}: + return "val" + if s in {"train", "full"}: + return s + raise ValueError(f"Unsupported {split=}. Use train/val/full.") + + +class BaseActionLeRobotDataset(Dataset): + """Reusable base class for Action LeRobot-backed map-style datasets. + + Subclasses typically: + 1) call ``_register_source`` to register one or more LeRobot sources + 2) implement ``__getitem__`` for dataset-specific sample parsing + 3) call ``_build_result`` to assemble the return dict + """ + + # Applied as: R_opencv = R_native @ _to_opencv + # Subclasses override in __init__; default is identity (no correction). + + # Bundled normalization stats directory. Stats are committed at + # ``<_NORMALIZERS_DIR>/__.json`` (flat + # layout matching the existing UMI files) and produced by + # ``projects/cosmos3/vfm/datasets/action/compute_action_stats.py``. + # Subclasses that need a different filename scheme can override + # :meth:`_normalizer_filename`. + _NORMALIZERS_DIR: ClassVar[Path] = Path(__file__).parent / "normalizers" + + def __init__( + self, + *, + fps: float, + chunk_length: int, + split_seed: int, + split_val_ratio: float, + split: str, + mode: str, + embodiment_type: str, + viewpoint: Viewpoint, + pose_convention: str | None = None, + rotation_format: str | None = None, + action_normalization: ActionNormalization | None = None, + apply_forward_clamp: bool = False, + tolerance_s: float = 1e-4, + max_loaded_datasets: int = _LRU_DATASET_MAX_LOADED, + skip_video_loading: bool = False, + sample_stride: int = 1, + enable_fast_init: bool = False, + fast_init_max_workers: int = 64, + min_episode_length_frames: int | None = None, + ) -> None: + super().__init__() + _ensure_hf_hub_offline() + _patch_decoder_cache() + self._memprofile = _memprofile_enabled() + + assert sample_stride >= 1, f"sample_stride must be >= 1, got {sample_stride}" + assert fast_init_max_workers >= 1, f"fast_init_max_workers must be >= 1, got {fast_init_max_workers}" + assert action_normalization is None or action_normalization in _ACTION_NORMALIZATION_CHOICES, ( + f"action_normalization must be None or one of {_ACTION_NORMALIZATION_CHOICES}, got {action_normalization!r}" + ) + + with rss_tracker(f"{self.__class__.__name__}.__init__", enabled=self._memprofile): + self._fps = fps + self._dt = 1.0 / fps + self._chunk_length = chunk_length + self._split_seed = split_seed + self._split_val_ratio = split_val_ratio + self._split = _normalize_split(split) + self._mode = mode + self._embodiment_type = embodiment_type + self._viewpoint: Viewpoint = viewpoint + self._pose_convention = pose_convention + self._rotation_format = rotation_format + self._action_normalizer: ActionNormalizer | None = None + if action_normalization is not None: + # ``apply_forward_clamp`` decides whether normalized actions are + # clamped into [-1, 1]. It matters most on the rot6d diagonal + # channels: for DROID, R00/R11 have a q01/q99 half-width of + # ~7.7e-4, so they saturate at roughly 3.2 deg of per-frame + # rotation and 10 deg/frame normalizes to about -19 unclamped. + # Callers name the intent rather than inheriting a default. + self._action_normalizer = resolve_action_normalization( + action_normalization, + self._load_norm_stats(action_normalization), + apply_forward_clamp=apply_forward_clamp, + ) + self._tolerance_s = tolerance_s + self._max_loaded_datasets = max_loaded_datasets + self._skip_video_loading = skip_video_loading + self._sample_stride = sample_stride + self._enable_fast_init = enable_fast_init + self._fast_init_max_workers = fast_init_max_workers + # Optional post-filter on raw episode length. When set, episodes + # whose raw frame count is below this threshold are dropped from + # ``_episode_records`` after ``build_episode_spans`` runs. Lets + # eval configs select e.g. only ``> 60s`` wall-clock episodes + # (1800 raw frames at native 30 fps) while keeping ``chunk_length`` + # (the fetch window) at a smaller value such as 900 (60 s @ fps=15). + # Subclasses that override ``_append_index_records`` are expected + # to honor this attribute themselves. + self._min_episode_length_frames: int | None = min_episode_length_frames + self._delta_timestamps: dict[str, list[float]] = {} + self._to_opencv: np.ndarray | dict[str, np.ndarray] = np.eye(3, dtype=np.float32) + + if pose_convention is None: + log.warning( + f"{self.__class__.__name__}: pose_convention is not set. " + "Consider specifying 'backward_framewise' or 'backward_anchored'." + ) + + self._datasets: list[LeRobotDataset | None] = [] + self._dataset_build_args: list[dict[str, Any] | None] = [] + self._loaded_lru: OrderedDict[int, None] = OrderedDict() + + # -- Flat index structures (populated by _append_index_records) -- + # Together these two lists form a searchable map from a flat + # global index to (dataset, row, episode, frame). One entry per + # episode span across *all* registered sources. + # + # _episode_records[i] = (ds_idx, sample_start, valid_len, episode_id) + # ds_idx – which source dataset (index into _datasets) + # sample_start – first row of this span in that dataset's table + # valid_len – number of usable frames in this span + # episode_id – the episode this span belongs to + # + # _episode_cum_ends[i] = running total of valid_len through span i + # Used for O(log N) lookup via bisect_right in _resolve_index. + self._episode_records: list[tuple[int, int, int, int]] = [] + self._episode_cum_ends: list[int] = [] + self._num_valid_indices = 0 + self._domain_id = get_domain_id(self._embodiment_type) + + # Deferred-init shard roots — a list of root paths. + # Subclasses populate this in __init__; _register_sources() + # reads _delta_timestamps and _tolerance_s from self (both + # initialised above, with _delta_timestamps overridden by + # each subclass). + # ActionUnifiedIterableDataset.assign_worker uses len() for + # round-robin shard distribution and _register_sources(indices) + # for deferred loading. When empty, shard distribution is + # skipped (every worker iterates the full dataset). + self._all_shard_roots: list[str] = [] + + # -- public properties --------------------------------------------------- + + @property + def fps(self) -> float: + return self._fps + + @property + def chunk_length(self) -> int: + return self._chunk_length + + @property + def split(self) -> str: + return self._split + + @property + def mode(self) -> str: + return self._mode + + @mode.setter + def mode(self, value: str) -> None: + self._mode = value + + @property + def domain_id(self) -> int: + return self._domain_id + + # -- source registration ------------------------------------------------- + + def _register_source( + self, + *, + delta_timestamps: dict[str, list[float]], + tolerance_s: float, + root: str | None = None, + repo_id: str = "local", + force_cache_sync: bool = False, + download_videos: bool = False, + video_backend: str | None = None, + revision: str | None = None, + dataset_label: str | None = None, + prefetched_meta: LeRobotDatasetMetadata | None = None, + ) -> LeRobotDatasetMetadata: + """Register a LeRobot dataset source lazily (metadata-only at init). + + ``prefetched_meta`` lets subclasses load metadata in a thread pool + (``LeRobotDatasetMetadata`` reads are pure I/O — ``info.json`` + + ``episodes.parquet`` + ``tasks.parquet``) and then hand the ready + object to the serial append-path below, which still manages the + order-sensitive shared state (``_datasets`` / ``_dataset_build_args`` + / ``_episode_records`` / ``_episode_cum_ends``). When ``None`` the + caller gets the original single-threaded behavior. + """ + label_str = f" [{dataset_label}]" if dataset_label else "" + cls = self.__class__.__name__ + # "local" is not a valid PEP 440 version, so LeRobot's + # is_valid_version() check skips the get_safe_version() HF API call. + if repo_id == "local" and revision is None: + revision = "local" + + with rss_tracker(f"{cls}{label_str} — metadata load", enabled=self._memprofile): + if prefetched_meta is not None: + meta = prefetched_meta + else: + meta = LeRobotDatasetMetadata( + repo_id=repo_id, + root=root, + revision=revision, + force_cache_sync=force_cache_sync, + ) + ds_idx = len(self._datasets) + self._datasets.append(None) + self._dataset_build_args.append( + { + "repo_id": repo_id, + "root": root, + "delta_timestamps": delta_timestamps, + "tolerance_s": tolerance_s, + "force_cache_sync": force_cache_sync, + "download_videos": download_videos, + "video_backend": video_backend, + "revision": revision, + } + ) + + with rss_tracker( + f"{cls}{label_str} — index records", + enabled=self._memprofile, + extras_fn=lambda: [ + f"episode_records so far: {len(self._episode_records)} entries, " + f"~{_fmt_mb(_deep_size(self._episode_records) / (1024 * 1024))}", + f"episode_cum_ends so far: {len(self._episode_cum_ends)} entries, " + f"~{_fmt_mb(_deep_size(self._episode_cum_ends) / (1024 * 1024))}", + ], + ): + self._append_index_records(meta=meta, ds_idx=ds_idx, dataset_label=dataset_label) + + return meta + + def _append_index_records( + self, + *, + meta: LeRobotDatasetMetadata, + ds_idx: int, + dataset_label: str | None = None, + ) -> None: + """Populate episode split / index records from dataset metadata.""" + episode_ids = split_episode_ids( + total_episodes=meta.total_episodes, + seed=self._split_seed, + val_ratio=self._split_val_ratio, + split=self._split, + ) + # TODO: remove once bridge training switches to a pre-filtered mirror. + if hasattr(self, "_filter_valid_episodes"): + episode_ids = self._filter_valid_episodes(meta, episode_ids) + episode_spans, valid_count, sample_count = build_episode_spans( + episodes=meta.episodes, + episode_ids=episode_ids, + chunk_length=self._chunk_length, + sample_stride=self._sample_stride, + ) + + # Optional duration filter (see ``self._min_episode_length_frames`` + # comment in ``__init__``). Drops episodes whose raw frame count is + # below the threshold. Operates after ``build_episode_spans`` so the + # threshold is decoupled from ``chunk_length`` (which controls the + # fetch window). No-op when the attribute is None. + if self._min_episode_length_frames is not None: + length_lookup = list(meta.episodes["length"]) + before = len(episode_spans) + episode_spans = [ + (eid, ss, vl) + for (eid, ss, vl) in episode_spans + if int(length_lookup[eid]) >= self._min_episode_length_frames + ] + dropped = before - len(episode_spans) + if dropped > 0: + log.info( + f"{self.__class__.__name__}: " + f"min_episode_length_frames={self._min_episode_length_frames} " + f"dropped {dropped} / {before} chunk-eligible spans" + ) + + class_name = self.__class__.__name__ + label = f" [{dataset_label}]" if dataset_label else "" + log.info(f"{class_name}{label}: split={self._split}, num episodes={len(episode_ids)}") + if sample_count > 0: + log.info( + f"{class_name}{label}: kept {valid_count} / {sample_count} " + f"({100 * valid_count / sample_count:.2f} %) samples" + ) + + for episode_id, sample_start, valid_len in episode_spans: + self._episode_records.append((ds_idx, sample_start, valid_len, episode_id)) + self._num_valid_indices += valid_len + self._episode_cum_ends.append(self._num_valid_indices) + + # -- deferred shard registration ----------------------------------------- + + def _register_sources(self, indices: list[int] | None = None) -> None: + """Register a subset (or all) of the shard roots in ``_all_shard_roots``. + + Called by ``ActionUnifiedIterableDataset.assign_worker`` during training, + or explicitly by eval/visualization scripts after construction. + + ``_all_shard_roots`` is a list of root paths. Per-shard args that are + shared across all shards (``delta_timestamps``, ``tolerance_s``) are + taken from ``self``. Subclasses may override this for extra per-shard + setup (e.g. loading instruction segments). + + When ``enable_fast_init=True``, ``LeRobotDatasetMetadata`` (a pure-IO + read of ``info.json`` + ``episodes.parquet`` + ``tasks.parquet``) is + prefetched in a thread pool and handed to the order-sensitive + serial register loop via ``prefetched_meta=``. Shard count scales + the speedup; for single-shard datasets the two paths are + equivalent. + + Args: + indices: Which entries of ``_all_shard_roots`` to register. + ``None`` means all. + """ + if indices is None: + indices = list(range(len(self._all_shard_roots))) + if not indices: + return + + roots = [self._all_shard_roots[i] for i in indices] + + if self._enable_fast_init: + # ``_ensure_hf_hub_offline`` already ran in ``__init__`` and is + # idempotent; no need to re-invoke here. + workers = max(1, min(self._fast_init_max_workers, len(roots))) + metas: list[LeRobotDatasetMetadata | None] = _parallel_map( + lambda root: LeRobotDatasetMetadata(repo_id="local", root=root, revision="local"), + roots, + max_workers=workers, + label=f"{type(self).__name__}: LeRobotDatasetMetadata prefetch", + ) + else: + metas = [None] * len(roots) + + for root, meta in zip(roots, metas): + label = root.rsplit("/", 1)[-1] if "/" in root else root + self._register_source( + root=root, + delta_timestamps=self._delta_timestamps, + tolerance_s=self._tolerance_s, + dataset_label=label, + prefetched_meta=meta, + ) + + # -- lazy dataset access ------------------------------------------------- + + def _get_dataset(self, ds_idx: int) -> LeRobotDataset: + """Get or lazily construct the LeRobot dataset for the given source index. + + Loaded datasets are tracked with LRU ordering. When the number of + loaded datasets exceeds ``_max_loaded_datasets`` the least-recently-used + dataset is evicted (set back to ``None``) so the GC can reclaim it. + """ + ds = self._datasets[ds_idx] + if ds is not None: + self._loaded_lru.move_to_end(ds_idx) + return ds + + _ensure_hf_hub_offline() + + build_args = self._dataset_build_args[ds_idx] + if build_args is None: + raise RuntimeError(f"Missing dataset build args for dataset index {ds_idx}") + + # Evict least-recently-used datasets before loading a new one. + while len(self._loaded_lru) >= self._max_loaded_datasets: + evict_idx, _ = self._loaded_lru.popitem(last=False) + self._datasets[evict_idx] = None + + with rss_tracker( + f"[WORKER {_os.getpid()}] Lazy-loaded ds[{ds_idx}]", + enabled=self._memprofile, + extras_fn=lambda: [f"total loaded={len(self._loaded_lru)}/{len(self._datasets)}"], + ): + delta_ts = build_args["delta_timestamps"] + if self._skip_video_loading: + # Covers both LeRobot v2 (``observation.images.``) and + # v3 (``observation.image.``) video-column conventions. + delta_ts = {k: v for k, v in delta_ts.items() if not k.startswith("observation.image")} + + log.info(f"Loading shard root={build_args['root']}") + ds = LeRobotDataset( + repo_id=build_args["repo_id"], + root=build_args["root"], + delta_timestamps=delta_ts, + tolerance_s=build_args["tolerance_s"], + force_cache_sync=build_args["force_cache_sync"], + download_videos=build_args["download_videos"], + video_backend=build_args["video_backend"], + revision=build_args["revision"], + episodes=None, + ) + if self._skip_video_loading: + ds.meta.info["features"] = { + k: v for k, v in ds.meta.info["features"].items() if v.get("dtype") != "video" + } + self._datasets[ds_idx] = ds + self._loaded_lru[ds_idx] = None + + return ds + + # -- index resolution ---------------------------------------------------- + + def _resolve_index(self, idx: int) -> tuple[int, int, int, int]: + """Map a flat global index to the source dataset, row, episode, and frame. + + Multiple datasets are concatenated into a single virtual sequence. + Each episode contributes a contiguous *span* of valid frames, and + ``_episode_cum_ends[i]`` stores the running total of valid frames + through the *i*-th span. For example, with two episodes of lengths + 5 and 3 the cum-ends are ``[5, 8]``, so global index 6 falls in the + second span at offset 1. + + The lookup is O(log N) via :func:`bisect_right`. + + Returns: + dataset_idx: Which source dataset this sample belongs to. + row_idx: Row index *within* that dataset's LeRobot table. + episode_id: The episode ID for this sample. + frame_offset: Frame offset from the start of the episode span + (0-based). + + Pure index math -- no I/O or dataset access. Higher-level helpers + like :meth:`_fetch_sample` build on this. + """ + # Support negative indexing (e.g. -1 → last sample). + if idx < 0: + idx += self._num_valid_indices + if idx < 0 or idx >= self._num_valid_indices: + raise IndexError(f"{self.__class__.__name__} index {idx} out of range for size {self._num_valid_indices}") + + # _episode_cum_ends is a monotonically increasing list where entry i + # holds the cumulative number of valid frames up to and including the + # i-th episode span. bisect_right finds the first span whose + # cumulative end is strictly greater than idx, i.e. the span that + # contains idx. + # + # Example: cum_ends = [5, 8, 20] + # idx=0 -> span_idx=0 (first span, frames 0..4) + # idx=4 -> span_idx=0 + # idx=5 -> span_idx=1 (second span, frames 5..7) + # idx=8 -> span_idx=2 (third span, frames 8..19) + span_idx = bisect_right(self._episode_cum_ends, idx) + + # The global index where this span begins is the previous span's + # cumulative end (or 0 for the very first span). The frame_offset + # is how far idx is into this particular episode. + span_start = 0 if span_idx == 0 else self._episode_cum_ends[span_idx - 1] + frame_offset = idx - span_start + + # _episode_records[span_idx] stores (dataset_idx, row_start, valid_len, + # episode_id). row_start is the absolute row in the LeRobot table + # where this episode begins. With sample_stride=k, consecutive + # valid indices map to rows k apart inside the episode, so the + # effective row is row_start + frame_offset * sample_stride. + dataset_idx, row_start, _, episode_id = self._episode_records[span_idx] + row_idx = row_start + frame_offset * self._sample_stride + return dataset_idx, row_idx, episode_id, frame_offset + + def _choose_mode(self) -> str: + """Resolve the active mode for one sample request.""" + if self._mode == "joint": + return random.choice(("forward_dynamics", "inverse_dynamics", "wam")) + return self._mode + + def _fetch_sample(self, idx: int) -> tuple[str, int, int, dict[str, Any]]: + """Resolve index, pick a mode, and load the sample from the dataset. + + Returns ``(mode, dataset_idx, row_idx, sample_dict)``. + """ + mode = self._choose_mode() + dataset_idx, row_idx, _, _ = self._resolve_index(idx) + + self._getitem_count = getattr(self, "_getitem_count", 0) + 1 + profile = self._memprofile and self._getitem_count % 50 == 1 + + with rss_tracker( + f"[WORKER {_os.getpid()}] __getitem__ transient (dataset_idx={dataset_idx})", + enabled=profile, + after_fn=lambda: log_worker_memory_breakdown(self), + ): + sample = self._get_dataset(dataset_idx)[row_idx] + + if self._skip_video_loading: + sample = defaultdict(lambda: None, sample) + + return mode, dataset_idx, row_idx, sample + + # -- action normalization ------------------------------------------------ + + def _normalizer_filename(self) -> str: + """Bundled stats filename for this dataset instance. + + Default convention (matches ``compute_action_stats.py`` output): + ``[_][_].json``. + + Pose/rotation suffixes are appended only when the instance actually + has them (SE(3) pose datasets like Bridge / DROID). Joint-space + datasets — where both are ``None`` — resolve to just + ``.json``. + + Subclasses may override when the bundled filename uses a different + scheme (e.g. UMI's ``uva_umi_single_task_normalizer.json``). + """ + if not self._embodiment_type: + raise RuntimeError( + f"{self.__class__.__name__}: embodiment_type is not set; cannot resolve normalizer filename." + ) + parts = [self._embodiment_type] + if self._pose_convention: + parts.append(self._pose_convention) + if self._rotation_format: + parts.append(self._rotation_format) + return "_".join(parts) + ".json" + + def _normalizer_path(self) -> Path: + """Full path to the bundled stats JSON for this dataset.""" + return self._NORMALIZERS_DIR / self._normalizer_filename() + + def _load_norm_stats(self, action_normalization: ActionNormalization) -> dict[str, torch.Tensor]: + """Load action normalization stats for the configured normalization mode. + + Raises :class:`FileNotFoundError` if the stats file is missing. This + is intentional — silently falling back to identity normalization when + the user asked for ``quantile`` / ``quantile_rot`` / ``meanstd`` / + ``minmax`` would be a training bug. + """ + stats_key = "global_raw" if action_normalization == "quantile_rot" else "global" + raw_stats = load_action_stats(str(self._normalizer_path()), stats_key=stats_key) + return {key: torch.from_numpy(value).float() for key, value in raw_stats.items()} # dict[str,[D]] + + def get_action_normalizer( + self, + _sample: dict[str, Any] | None = None, + ) -> ActionNormalizer | None: + """Return the configured action normalizer for transform-time preprocessing.""" + return self._action_normalizer + + # -- video formatting ---------------------------------------------------- + + def _convert_video(self, video_tchw: torch.Tensor | None) -> torch.Tensor | None: + """Convert LeRobot ``(T,C,H,W)`` float video to Action ``(C,T,H,W)`` uint8. + + Args: + video_tchw: Raw floating-point video tensor in ``[0, 1]`` with + LeRobot layout, or ``None``. # [T,C,H,W] | None + + Returns: + Action-formatted video tensor, or ``None``. # [C,T,H,W] | None + """ + if self._skip_video_loading or video_tchw is None: + return None + if video_tchw.ndim != 4: + raise ValueError( + f"{self.__class__.__name__}._convert_video expected video with shape [T,C,H,W], " + f"got ndim={video_tchw.ndim}" + ) + if not torch.is_floating_point(video_tchw): + raise TypeError( + f"{self.__class__.__name__}._convert_video expected floating-point video in [0, 1], " + f"got dtype={video_tchw.dtype}" + ) + video_min = video_tchw.amin() # [] + video_max = video_tchw.amax() # [] + if video_min.item() < 0.0 or video_max.item() > 1.0: + raise ValueError( + f"{self.__class__.__name__}._convert_video expected floating-point video in [0, 1], " + f"got range=[{video_min.item():.6f}, {video_max.item():.6f}]" + ) + formatted_video = (video_tchw * 255.0).clamp(0.0, 255.0).to(torch.uint8).permute(1, 0, 2, 3) # [C,T,H,W] + return formatted_video + + # -- result building ----------------------------------------------------- + + def _build_action_spec(self) -> ActionSpec | None: + """Subclass override: declare this dataset's action layout. + + Called once per instance — the result is cached by ``self.action_spec``. + Return ``None`` to skip spec-driven idle detection; in that case + ``_compute_idle_frames`` will log a one-time warning and return + ``None`` for every sample. + """ + return None + + @cached_property + def action_spec(self) -> ActionSpec | None: + """Cached :class:`ActionSpec` from ``_build_action_spec``. + + Returns ``None`` when the subclass did not declare one; idle detection + is then skipped (with a one-time warning) until the subclass overrides + ``_build_action_spec``. + """ + return self._build_action_spec() + + @cached_property + def action_names(self) -> list[str] | None: + spec = self.action_spec + return spec.names if spec is not None else None + + # Idle-detection thresholds. Defined as **velocities** (per second) so the + # same numeric value means the same physical motion across datasets with + # different sampling rates; converted to per-frame at call time using + # ``self._fps`` via :meth:`_resolve_idle_thresholds`. + # + # Defaults: + # - ``idle_eps_t_per_sec`` = 5 mm/s (≈ 1 mm/frame at 5 Hz) + # - ``idle_eps_r_per_sec`` = 1.5°/s (geodesic, rotation-format aware) + # - ``idle_eps_g`` = 1e-2 unit gripper Δ (no fps) + # - ``idle_joint_threshold_per_sec`` = 5e-3 rad/s + # - ``idle_min_streak`` = 3 require ≥ 3 consecutive + # + # Subclasses can either override the ``*_per_sec`` attributes (preferred — + # keeps the velocity semantics) or set the corresponding ``idle_eps_*`` / + # ``idle_joint_threshold`` attribute to a non-``None`` value to bypass the + # per-fps conversion entirely (raw per-frame override). + idle_eps_t_per_sec: float = 5e-3 + idle_eps_r_per_sec: float = math.radians(1.5) + idle_eps_g: float = 1e-2 + idle_joint_threshold_per_sec: float = 5e-3 + idle_min_streak: int = 3 + + # Optional per-frame overrides. ``None`` (default) → use the ``*_per_sec`` + # attribute / fps conversion above. + idle_eps_t: float | None = None + idle_eps_r: float | None = None + idle_joint_threshold: float | None = None + + def _resolve_idle_thresholds(self) -> tuple[float, float, float, float]: + """Resolve per-frame idle thresholds for this dataset instance. + + Returns ``(eps_t, eps_r, eps_g, joint_threshold)`` in raw per-frame + units. Honours direct per-frame overrides if the subclass sets the + non-``_per_sec`` attribute; otherwise scales the ``_per_sec`` values + by ``self._fps``. + """ + fps = float(self._fps) if self._fps else 1.0 + eps_t = self.idle_eps_t if self.idle_eps_t is not None else self.idle_eps_t_per_sec / fps + eps_r = self.idle_eps_r if self.idle_eps_r is not None else self.idle_eps_r_per_sec / fps + joint_thr = ( + self.idle_joint_threshold + if self.idle_joint_threshold is not None + else self.idle_joint_threshold_per_sec / fps + ) + return float(eps_t), float(eps_r), float(self.idle_eps_g), float(joint_thr) + + def _compute_idle_frames(self, raw_action: torch.Tensor) -> torch.Tensor | None: + """Count idle frames in the *raw* (un-normalized) action chunk. + + Requires ``self.action_spec`` to be declared via ``_build_action_spec``. + Returns ``None`` when: + - ``pose_convention`` is not ``"backward_framewise"`` (TODO: extend), + - the subclass has not declared an ``ActionSpec`` (logs a one-time warning), + - the action layout does not match the declared spec. + + Detection thresholds come from the ``idle_eps_*`` class attributes + (overridable per dataset). Subclasses can also override this method + outright, or pass an explicit ``idle_frames`` integer via + ``**extras`` to :meth:`_build_result`. + """ + # TODO: currently we only support backward_framewise. Other pose + # conventions (anchored / absolute) need different idle semantics. + if self._pose_convention != "backward_framewise": + if not getattr(self, "_warned_pose_convention", False): + log.warning( + f"Dataset {self.__class__.__name__}: pose_convention=" + f"{self._pose_convention!r} is not 'backward_framewise'; " + "skipping idle-frames detection. Centralize the dataset " + "to backward_framewise to enable IdleFrames captioning." + ) + self._warned_pose_convention = True + return None + + spec = self.action_spec + if spec is None: + if not getattr(self, "_warned_no_action_spec", False): + log.warning( + f"Dataset {self.__class__.__name__} has no action spec defined; " + "skipping idle-frames detection. Override _build_action_spec() to enable it." + ) + self._warned_no_action_spec = True + return None + + eps_t, eps_r, eps_g, joint_thr = self._resolve_idle_thresholds() + try: + n = compute_idle_frames( + raw_action, + spec, + eps_t=eps_t, + eps_r=eps_r, + eps_g=eps_g, + joint_threshold=joint_thr, + min_streak=self.idle_min_streak, + ) + except (ValueError, TypeError) as e: + if not getattr(self, "_warned_action_layout", False): + log.warning( + f"Dataset {self.__class__.__name__}: action layout does " + f"not match the declared ActionSpec " + f"(action_dim={int(raw_action.shape[-1])}, " + f"spec.dim={spec.dim}); skipping idle-frames detection. " + f"Underlying error: {e}" + ) + self._warned_action_layout = True + return None + return torch.tensor(n, dtype=torch.long) + + def _build_result( + self, + *, + mode: str, + video: torch.Tensor | None, + action: torch.Tensor, + ai_caption: str, + **extras: Any, + ) -> dict[str, Any]: + """Assemble the common return dict for ``__getitem__``. + + ``video`` is expected in raw LeRobot layout before final formatting. + Subclasses may pass extra keys (e.g. ``initial_pose``) via ``**extras``. + ``idle_frames`` is auto-computed from the raw (un-normalized) ``action`` + whenever the dataset's pose/rotation conventions allow it; subclasses + can override by passing ``idle_frames`` (int or scalar tensor) via + ``**extras``. + """ + # Compute idle_frames from the raw action before normalization, unless + # the subclass has provided one explicitly via ``**extras``. + if "idle_frames" not in extras: + idle_frames = self._compute_idle_frames(action) + if idle_frames is not None: + extras = {"idle_frames": idle_frames, **extras} + + if self._action_normalizer is not None: + action = self._action_normalizer.normalize_action(action) # [T,D] + + if self._skip_video_loading: + result: dict[str, Any] = {"action": action} + if "idle_frames" in extras: + result["idle_frames"] = extras["idle_frames"] + return result + formatted_video = self._convert_video(video) # [C,T,H,W] | None + return { + "ai_caption": ai_caption, + "video": formatted_video, + "action": action, + "conditioning_fps": torch.tensor(self._fps, dtype=torch.long), + "mode": mode, + "domain_id": torch.tensor(self._domain_id, dtype=torch.long), + "viewpoint": self._viewpoint, + **extras, + } + + def __len__(self) -> int: + return self._num_valid_indices + + def get_shuffle_blocks(self) -> list[tuple[int, int]]: + """Per-episode (or per kept-segment when ``use_filter_dict``) flat-index + blocks ``(start, length)`` over ``[0, len(self))``. ``ActionIterableShuffleDataset`` + shuffles the ORDER of these blocks and shards them disjointly across + ``(rank, worker)`` while keeping windows *within* a block sequential -> + decorrelates batches without random-access I/O. Derived from + ``_episode_cum_ends`` (the same monotonic cumulative index ``__getitem__`` + bisects), so blocks align exactly with flat-index addressing.""" + blocks: list[tuple[int, int]] = [] + prev = 0 + for c in self._episode_cum_ends: + c = int(c) + if c > prev: + blocks.append((prev, c - prev)) + prev = c + return blocks diff --git a/examples/worldmodel/cosmos3/vendor/cosmos_framework/data/generator/action/datasets/droid_lerobot_dataset.py b/examples/worldmodel/cosmos3/vendor/cosmos_framework/data/generator/action/datasets/droid_lerobot_dataset.py new file mode 100644 index 000000000000..85c0484fb985 --- /dev/null +++ b/examples/worldmodel/cosmos3/vendor/cosmos_framework/data/generator/action/datasets/droid_lerobot_dataset.py @@ -0,0 +1,550 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: OpenMDW-1.1 + +import json +import os +import random +from pathlib import Path +from typing import Any, cast + +import numpy as np +import torch +import torch.nn.functional as F +import torchvision.transforms.v2 as T +from scipy.spatial.transform import Rotation as R + +from cosmos_framework.data.generator.action.datasets.cosmos3_action_lerobot import ( + ActionNormalization, + ActionSpec, + BaseActionLeRobotDataset, + Gripper, + Joint, + Pos, + Rot, + build_action_spec, + build_episode_spans, + split_episode_ids, +) +from cosmos_framework.data.generator.action.datasets.droid_lerobot_dataset_config import ( + _GRIPPER_STATE_FEATURE, + _JOINT_ACTION_FEATURE, + _JOINT_STATE_FEATURE, + ACTION_FEATURES, + HAS_MULTI_LANGUAGE_ANNOTATIONS, + IMAGE_FEATURES, + IS_FLAT_ACTION, + IS_GRIPPER_ACTION_FLIPPED, + LEROBOT_ROOTS, + STATE_FEATURES, +) +from cosmos_framework.data.generator.action.utils.episode_filters import normalize_episode_filter_keys +from cosmos_framework.data.generator.action.utils.pose_utils import ( + PoseConvention, + build_abs_pose_from_components, + convert_rotation, + pose_abs_to_rel, +) +from cosmos_framework.data.generator.action.utils.viewpoint_utils import Viewpoint +from cosmos_framework.utils import log + +_NORMALIZER_PATH = Path(__file__).parent.parent / "normalizer_stats/droid_lerobot_stats.json" + +# 90-degree clockwise rotation about the Z axis (in local frame), converting +# DROID Franka panda_link8 orientation to the OpenCV camera convention. +_DROID_TO_OPENCV: np.ndarray = np.array( + [ + [0.0, -1.0, 0.0], + [1.0, 0.0, 0.0], + [0.0, 0.0, 1.0], + ], + dtype=np.float32, +) + + +class DROIDLeRobotDataset(BaseActionLeRobotDataset): + """Read an explicit dataset root and optional filters keyed by complete episode IDs. + + URI-based filter metadata also requires filter_dict_source_prefix, the dataset URI prefix. + """ + + EMBODIMENT_TYPE: str = "droid_lerobot" + + def _normalizer_path(self) -> Path: + """Bundled DROID stats. + + The base convention would resolve to + ``datasets/normalizers/droid_lerobot_backward_framewise_rot6d.json``, + which does not exist; the shipped file is ``droid_lerobot_stats.json`` + one directory up, matching every other action dataset. + """ + return _NORMALIZER_PATH + + def __init__( + self, + root: str, + fps: float = 15.0, + chunk_length: int = 16, + split_seed: int = 42, + split_val_ratio: float = 0.03, + split: str = "train", + mode: str = "wam", + pose_convention: PoseConvention = "backward_framewise", + action_normalization: ActionNormalization | None = None, + tolerance_s=2e-4, + viewpoint: Viewpoint = "concat_view", + use_success_only: bool = False, + video_mode: str | None = None, # TODO: remove + action_space: str = "midtrain", # TODO: remove + use_state: bool = False, + use_filter_dict: bool = False, + filter_dict_path: str | None = None, + enable_fast_init: bool = False, + max_num_history_actions: int = 0, + use_image_augmentation: bool = False, + # Appended rather than placed next to ``action_normalization``: the + # parameters above are positional-or-keyword, so inserting mid-list + # would shift every later argument for positional callers. + apply_forward_clamp: bool = False, + jitter_after_compose: bool = False, + filter_dict_source_prefix: str = "", + ) -> None: + """Load a local DROID dataset; filtering requires an explicit JSON path.""" + if not root: + raise ValueError("root must be a non-empty local DROID dataset path") + if use_filter_dict and not filter_dict_path: + raise ValueError("filter_dict_path is required when use_filter_dict=True") + super().__init__( + fps=fps, + chunk_length=chunk_length, + split_seed=split_seed, + split_val_ratio=split_val_ratio, + split=split, + mode=mode, + embodiment_type="droid_lerobot", + viewpoint=viewpoint, + pose_convention=pose_convention, + rotation_format="rot6d", + action_normalization=action_normalization, + apply_forward_clamp=apply_forward_clamp, + tolerance_s=tolerance_s, + enable_fast_init=enable_fast_init, + ) + self._use_success_only = use_success_only + self._video_mode = video_mode + self._action_space = action_space + self._use_state = use_state + self._use_filter_dict = use_filter_dict + self._filter_dict_path = filter_dict_path + self._max_num_history_actions = max_num_history_actions + self._use_image_augmentation = use_image_augmentation + self._jitter_after_compose = jitter_after_compose + if max_num_history_actions > 0 and action_space not in ("midtrain", "joint_pos"): + raise ValueError( + f"max_num_history_actions is only supported with action_space='midtrain' or 'joint_pos', got {action_space!r}" + ) + + self._is_val_temp_seg = split == "val_temp_seg" + self._to_opencv = _DROID_TO_OPENCV + + version = os.path.basename(root) + try: + lerobot_roots = LEROBOT_ROOTS[version] + self._image_features = IMAGE_FEATURES[version] + self._state_features = STATE_FEATURES[version] + self._action_features = ACTION_FEATURES[version] + self._is_flat_action = IS_FLAT_ACTION[version] + self._has_multi_language_annotations = HAS_MULTI_LANGUAGE_ANNOTATIONS[version] + self._is_gripper_action_flipped = IS_GRIPPER_ACTION_FLIPPED[version] + except KeyError as e: + raise ValueError(f"Unknown version: {version!r}. Supported: {list(LEROBOT_ROOTS.keys())}") from e + + if self._use_success_only and lerobot_roots: + lerobot_roots = [x for x in lerobot_roots if x.split("/", 1)[0] == "success"] + + self._all_shard_roots = [os.path.join(root, x) for x in lerobot_roots] if lerobot_roots else [root] + + observation_ts = [i * self._dt for i in range(0, self._chunk_length + 1)] + action_ts = [i * self._dt for i in range(0, self._chunk_length)] + if self._max_num_history_actions > 0 and self._action_space in ("midtrain", "joint_pos"): + observation_ts_ext = [i * self._dt for i in range(-self._max_num_history_actions, self._chunk_length + 1)] + action_ts_ext = [i * self._dt for i in range(-self._max_num_history_actions, self._chunk_length)] + else: + observation_ts_ext = observation_ts + action_ts_ext = action_ts + self._delta_timestamps: dict[str, list[float]] = { + self._state_features: observation_ts_ext, + self._action_features: action_ts_ext, + } + if self._viewpoint in ("wrist_view", "concat_view"): + self._delta_timestamps[self._image_features["wrist"]] = observation_ts + if self._viewpoint in ("third_person_view", "concat_view"): + self._delta_timestamps[self._image_features["left"]] = observation_ts + self._delta_timestamps[self._image_features["right"]] = observation_ts + if self._action_space == "joint_pos": + self._delta_timestamps[_JOINT_ACTION_FEATURE] = action_ts + if self._use_state or self._max_num_history_actions > 0: + self._delta_timestamps[_JOINT_STATE_FEATURE] = observation_ts_ext + self._delta_timestamps[_GRIPPER_STATE_FEATURE] = observation_ts_ext + if self._use_state and self._action_space != "joint_pos": + self._delta_timestamps[_GRIPPER_STATE_FEATURE] = observation_ts + + if self._use_filter_dict: + assert filter_dict_path is not None # Validated before dataset initialization. + with open(filter_dict_path) as f: + self._filter_dict = normalize_episode_filter_keys(json.load(f), source_prefix=filter_dict_source_prefix) + + self._image_augmentor: T.Compose | None = None + self._color_augmentor: T.ColorJitter | None = None + + # Eager source registration. Upstream defers this to its own dataloader's + # ActionUnifiedIterableDataset.assign_worker(); cosmos-framework instead + # drives the dataset through ActionIterableShuffleDataset (block-striding, + # which needs the full flat index present in every worker), so we build + # the index at construction time here. Metadata-only (LeRobotDatasetMetadata: + # info.json + episodes.parquet + tasks.parquet); the heavy per-shard + # LeRobotDataset video readers stay lazy behind the LRU in _get_dataset. + self._register_sources() + + def _append_index_records(self, *, meta, ds_idx: int, dataset_label: str | None = None) -> None: + """ """ + if not self._use_filter_dict: + super()._append_index_records(meta=meta, ds_idx=ds_idx, dataset_label=dataset_label) + return + + episode_ids = split_episode_ids( + total_episodes=meta.total_episodes, + seed=self._split_seed, + val_ratio=self._split_val_ratio, + split=self._split, + ) + episode_spans, _, sample_count = build_episode_spans( + meta.episodes, episode_ids, self._chunk_length, sample_stride=self._sample_stride + ) + + class_name = self.__class__.__name__ + label = f" [{dataset_label}]" + + log.info(f"{class_name}{label}: split={self._split}, num episodes={len(episode_ids)}") + + filtered_count = 0 + for episode_id, sample_start, valid_len in episode_spans: + ep_id_str = meta.episodes[episode_id]["episode_id"] + episode_key = str(ep_id_str) + ranges = self._filter_dict.get(episode_key) + if ranges is None: + continue + for s, e in ranges: + sub_start = max(s, 0) + sub_end = min(e - self._chunk_length, valid_len) + sub_valid_len = max(0, sub_end - sub_start) + if sub_valid_len > 0: + self._episode_records.append((ds_idx, sample_start + sub_start, sub_valid_len, episode_id)) + self._num_valid_indices += sub_valid_len + self._episode_cum_ends.append(self._num_valid_indices) + filtered_count += sub_valid_len + + if sample_count > 0: + log.info( + f"{class_name}{label}: kept {filtered_count} / {sample_count} ({100.0 * filtered_count / sample_count:.2f} %) samples" + ) + + def _register_sources(self, indices: list[int] | None = None) -> None: + """ """ + super()._register_sources(indices) + if self._is_val_temp_seg: + self._apply_temp_seg_filter() + + def _apply_temp_seg_filter(self) -> None: + """Replace index records with one high-scoring segment per episode. + + A segment is interesting if either: + - The gripper action changes significantly (open/close transition), or + - The gripper is closed and the end-effector position is moving. + Among qualifying segments the one with the highest score is kept. + """ + ds = self._get_dataset(0) + chunk_size = self._chunk_length + 1 + gripper_change_threshold = 0.5 + ee_movement_threshold = 0.01 + + new_records: list[tuple[int, int, int, int]] = [] + num_episodes = len(self._episode_records) + + for ds_idx, sample_start, valid_len, episode_id in self._episode_records: + end = sample_start + valid_len + self._chunk_length + num_candidates = valid_len + if num_candidates <= 0: + continue + + episode_data = ds.hf_dataset[sample_start:end] + actions = torch.tensor(np.array(episode_data[self._action_features])) # [N,action_dim] + states = torch.tensor(np.array(episode_data[self._state_features])) # [N,state_dim] + + gripper_action = actions[:, 6] if self._is_flat_action else actions # [N] + ee_pos = states[:, :3] # [N,3] + ee_disp = (ee_pos[1:] - ee_pos[:-1]).norm(dim=-1) # [N-1] + + ee_disp_windows = ee_disp.unfold(0, self._chunk_length, 1) # [num_candidates,chunk_length] + gripper_windows = gripper_action.unfold(0, chunk_size, 1) # [num_candidates,chunk_size] + + gripper_range = gripper_windows.max(dim=1).values - gripper_windows.min(dim=1).values # [num_candidates] + total_ee_movement = ee_disp_windows.sum(dim=1) # [num_candidates] + gripper_closed_ratio = (gripper_windows < 0.5).float().mean(dim=1) # [num_candidates] + + has_gripper_change = gripper_range > gripper_change_threshold + gripper_closed = gripper_closed_ratio > 0.5 + has_ee_movement = total_ee_movement > ee_movement_threshold + + scores = torch.zeros(num_candidates) # [num_candidates] + scores[has_gripper_change] = 0.5 + gripper_range[has_gripper_change] + total_ee_movement[has_gripper_change] + + closed_and_moving = gripper_closed & ~has_gripper_change & has_ee_movement + scores[closed_and_moving] = 1.0 + total_ee_movement[closed_and_moving] + + if scores.max().item() > 0: + best_offset = int(scores.argmax().item()) + new_records.append((ds_idx, sample_start + best_offset, 1, episode_id)) + + self._episode_records = new_records + self._num_valid_indices = len(new_records) + self._episode_cum_ends = list(range(1, len(new_records) + 1)) + + log.info(f"DROIDLeRobotDataset: val_temp_seg kept {len(new_records)} segments from {num_episodes} episodes") + + def _compose_multi_view(self, sample: dict[str, Any]) -> torch.Tensor: + """Compose wrist, left, and right views into a single frame. + + Layout (per frame): + ┌──────────────┐ + │ wrist │ (H, W) + ├───────┬──────┤ + │ left │ right│ (H/2, W/2) each + └───────┴──────┘ + + Left and right exterior cameras are downscaled by 2x so that they + tile to the same width as the wrist view. The output height is 3H/2. + + By default the full augmentation pipeline (random crop, resize, + ColorJitter) runs per view before composition. With + ``jitter_after_compose=True`` the pipeline is split around the + composition: crop and resize stay per view and run first, because + cropping the composite would cut across the tile boundaries, while + ColorJitter runs last, on the composite, where the exterior views have + already been downscaled. That is 3H/2 x W pixels instead of the + 3H x W of the three full-size views, so it costs roughly half as much + for the same augmentation. Both orderings draw one set of colour + factors and apply it to every view, so the views still agree on + lighting; what changes is that the exterior views are jittered after + downscaling rather than before, which for a random augmentation is a + re-parameterization rather than a different distribution. See + issue #174. + + Returns: + Composited raw video tensor in ``(T,C,H_out,W)`` float format. + """ + wrist = sample[self._image_features["wrist"]] # [T,C,H,W] + left = sample[self._image_features["left"]] # [T,C,H_l,W_l] + right = sample[self._image_features["right"]] # [T,C,H_r,W_r] + + if self._use_image_augmentation: + if self._image_augmentor is None: + _, _, h, w = wrist.shape + transforms: list[Any] = [ + T.RandomCrop((int(h * 0.95), int(w * 0.95))), + T.Resize((h, w), antialias=True), + ] + color_jitter = T.ColorJitter(brightness=0.3, contrast=0.4, saturation=0.5, hue=0.08) + if self._jitter_after_compose: + self._color_augmentor = color_jitter + else: + transforms.append(color_jitter) + self._image_augmentor = T.Compose(transforms) + n, m = wrist.shape[0], wrist.shape[0] + left.shape[0] + combined = self._image_augmentor(torch.cat([wrist, left, right], dim=0)) + wrist, left, right = combined[:n], combined[n:m], combined[m:] + + _, _, h_w, w_w = wrist.shape + half_h, half_w = h_w // 2, w_w // 2 + + left = F.interpolate(left, size=(half_h, half_w), mode="bilinear", align_corners=False) # [T,C,H/2,W/2] + right = F.interpolate(right, size=(half_h, half_w), mode="bilinear", align_corners=False) # [T,C,H/2,W/2] + bottom = torch.cat([left, right], dim=-1) # [T,C,H/2,W] + + composite = torch.cat([wrist, bottom], dim=-2) # [T,C,3H/2,W] + + if self._color_augmentor is not None: + composite = self._color_augmentor(composite) + + return composite # [T,C,3H/2,W] + + def _build_action_spec(self) -> ActionSpec: + """DROID: 10D ``[Pos, Rot6d, Gripper]`` for ``ee_pose``, + 8D ``[Joint(7), Gripper]`` for ``joint_pos``. + """ + if self._action_space == "joint_pos": + return build_action_spec(Joint(n=7, label="joint"), Gripper()) + return build_action_spec(Pos(), Rot("rot6d"), Gripper()) + + def __getitem__(self, idx: int) -> dict[str, Any]: + """ """ + mode, _, _, sample = self._fetch_sample(idx) + + if self._has_multi_language_annotations: + tasks = sample["task"].split(" | ") + ai_caption = random.choice(tasks) + else: + ai_caption = sample["task"] + + if self._skip_video_loading: + video = None + elif self._video_mode is None: + if self._viewpoint == "concat_view": + video = self._compose_multi_view(sample) + else: + video = sample[self._image_features["wrist"]] # [T,C,H,W] + else: + if self._video_mode == "wrist": + video = sample[self._image_features["wrist"]] + if self._video_mode in ("rand_exterior", "wrist_rand_exterior"): + exterior_key = random.choice([self._image_features["left"], self._image_features["right"]]) + if self._video_mode == "rand_exterior": + video = sample[exterior_key] + else: + video = torch.cat([sample[self._image_features["wrist"]], sample[exterior_key]], dim=2) + if self._video_mode in ("wrist_left_exterior", "wrist_both_exterior"): + wrist = sample[self._image_features["wrist"]] + half_h, half_w = wrist.shape[2] // 2, wrist.shape[3] // 2 + left = F.interpolate( + sample[self._image_features["left"]], size=(half_h, half_w), mode="bilinear", align_corners=False + ) + if self._video_mode == "wrist_left_exterior": + right = torch.zeros_like(left) + if self._video_mode == "wrist_both_exterior": + right = F.interpolate( + sample[self._image_features["right"]], + size=(half_h, half_w), + mode="bilinear", + align_corners=False, + ) + video = torch.cat([wrist, torch.cat([left, right], dim=-1)], dim=-2) + + extras: dict[str, Any] = {} + + if self._action_space == "midtrain": + pose_convention = cast(PoseConvention, self._pose_convention) + state = sample[self._state_features] # [T+1, state_dim] or [H+T+1, state_dim] + poses_abs = build_abs_pose_from_components(state[:, 0:3], state[:, 3:6], "euler_xyz") + poses_abs[:, :3, :3] = poses_abs[:, :3, :3] @ self._to_opencv + initial_pose = torch.from_numpy(poses_abs[-self._chunk_length - 1].copy()).float() + poses_rel = pose_abs_to_rel(poses_abs, rotation_format="rot6d", pose_convention=pose_convention) + gripper = ( + sample[self._action_features][:, [6]] + if self._is_flat_action + else sample[self._action_features].unsqueeze(-1) + ) + if self._is_gripper_action_flipped: + gripper = 1.0 - gripper + action = torch.from_numpy( + np.concatenate([poses_rel[-self._chunk_length :], gripper[-self._chunk_length :]], axis=-1) + ).float() # [T,10] + extras["initial_pose"] = initial_pose + if self._max_num_history_actions > 0: + _, _, _, frame_offset = self._resolve_index(int(idx)) + num_available = min(self._max_num_history_actions, frame_offset * self._sample_stride) + actual_h = num_available + # with 0.5 probability, randomly sample the number of history frames + if random.random() < 0.5: + actual_h = random.randint(0, num_available) + if actual_h > 0: + hist_action_raw = torch.from_numpy( + np.concatenate( + [ + poses_rel[-self._chunk_length - actual_h : -self._chunk_length], + gripper[-self._chunk_length - actual_h : -self._chunk_length], + ], + axis=-1, + ) + ).float() # [H,D] + extras["history_action"] = hist_action_raw + if self._use_state: + initial_gripper = sample[_GRIPPER_STATE_FEATURE][0].unsqueeze(-1) + if self._is_gripper_action_flipped: + initial_gripper = 1.0 - initial_gripper + initial_rot6d = convert_rotation(poses_abs[-self._chunk_length - 1, :3, :3], "matrix", "rot6d") + initial_state = torch.from_numpy( + np.concatenate((poses_abs[-self._chunk_length - 1, :3, 3], initial_rot6d, initial_gripper), axis=-1) + ).float() + action = torch.cat([initial_state.unsqueeze(0), action], dim=0) + if self._action_space == "ee_pose_delta": + state = sample[self._state_features] + pose = np.tile(np.eye(4), (state.shape[0], 1, 1)) + pose[:, :3, :3] = R.from_euler("xyz", state[:, 3:6]).as_matrix() + pose[:, :3, 3] = state[:, 0:3] + pose_delta = np.linalg.inv(pose[0]) @ pose[1:] + gripper = sample[self._action_features].unsqueeze(-1) + if self._is_gripper_action_flipped: + gripper = 1.0 - gripper + action = torch.from_numpy( + np.concatenate((pose_delta[:, :3, 3], pose_delta[:, :3, 0], pose_delta[:, :3, 1], gripper), axis=-1) + ).float() + if self._use_state: + initial_gripper = sample[_GRIPPER_STATE_FEATURE][0].unsqueeze(-1) + if self._is_gripper_action_flipped: + initial_gripper = 1.0 - initial_gripper + initial_state = torch.from_numpy( + np.concatenate((pose[0, :3, 3], pose[0, :3, 0], pose[0, :3, 1], initial_gripper), axis=-1) + ).float() + action = torch.cat([initial_state.unsqueeze(0), action], dim=0) + if self._action_space == "joint_pos": + gripper = sample[self._action_features][-self._chunk_length :].unsqueeze(-1) + if self._is_gripper_action_flipped: + gripper = 1.0 - gripper + action = torch.cat((sample[_JOINT_ACTION_FEATURE], gripper), dim=-1).float() + if self._max_num_history_actions > 0: + _, _, _, frame_offset = self._resolve_index(int(idx)) + num_available = min(self._max_num_history_actions, frame_offset * self._sample_stride) + actual_h = num_available + if random.random() < 0.5: + actual_h = random.randint(0, num_available) + if actual_h > 0: + hist_joint = sample[_JOINT_STATE_FEATURE][ + -self._chunk_length - 1 - actual_h : -self._chunk_length - 1 + ] + hist_gripper = sample[_GRIPPER_STATE_FEATURE][ + -self._chunk_length - 1 - actual_h : -self._chunk_length - 1 + ].unsqueeze(-1) + if self._is_gripper_action_flipped: + hist_gripper = 1.0 - hist_gripper + hist_action_raw = torch.cat((hist_joint, hist_gripper), dim=-1).float() + extras["history_action"] = hist_action_raw # [H,D] + if self._use_state: + initial_gripper = sample[_GRIPPER_STATE_FEATURE][-self._chunk_length - 1].unsqueeze(-1) + if self._is_gripper_action_flipped: + initial_gripper = 1.0 - initial_gripper + initial_state = torch.cat( + (sample[_JOINT_STATE_FEATURE][-self._chunk_length - 1], initial_gripper), dim=-1 + ).float() + action = torch.cat([initial_state.unsqueeze(0), action], dim=0) + + if self._viewpoint == "concat_view" and self._video_mode in ( + None, + "wrist_left_exterior", + "wrist_both_exterior", + ): + extras["additional_view_description"] = ( + "The top row is from the wrist-mounted camera. " + "The bottom row contains two horizontally concatenated third-person perspective views of the scene from opposite sides, with the robot visible." + ) + + return self._build_result( + mode=mode, + video=video, + action=action, + ai_caption=ai_caption, + **extras, + ) + + @property + def action_dim(self) -> int: + """ """ + return 8 if self._action_space == "joint_pos" else 10 diff --git a/examples/worldmodel/cosmos3/vendor/cosmos_framework/data/generator/action/datasets/droid_lerobot_dataset_config.py b/examples/worldmodel/cosmos3/vendor/cosmos_framework/data/generator/action/datasets/droid_lerobot_dataset_config.py new file mode 100644 index 000000000000..e97489c505b6 --- /dev/null +++ b/examples/worldmodel/cosmos3/vendor/cosmos_framework/data/generator/action/datasets/droid_lerobot_dataset_config.py @@ -0,0 +1,100 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: OpenMDW-1.1 + +_INSTITUTIONS = [ + "AUTOLab", + "CLVR", + "GuptaLab", + "ILIAD", + "IPRL", + "IRIS", + "PennPAL", + "RAD", + "RAIL", + "REAL", + "RPL", + "TRI", + "WEIRD", +] + +LEROBOT_ROOTS = { + "droid_lerobot_20260115_no_noops": None, + "droid_plus_lerobot_320x180_20260406_sharded": [f"success/{x}" for x in _INSTITUTIONS] + + [f"failure/{x}" for x in _INSTITUTIONS], + "droid_plus_lerobot_320x180_20260406": ["success", "failure"], + "droid_plus_lerobot_640x360_20260412_sharded": [f"success/{x}" for x in _INSTITUTIONS] + + [f"failure/{x}" for x in _INSTITUTIONS], + "droid_plus_lerobot_640x360_20260412": ["success", "failure"], +} + +IMAGE_FEATURES = { + "droid_lerobot_20260115_no_noops": { + "wrist": "observation.images.wrist_image_left", + "left": "observation.images.exterior_image_1_left", + "right": "observation.images.exterior_image_2_left", + }, + "droid_plus_lerobot_320x180_20260406_sharded": { + "wrist": "observation.image.wrist_image_left", + "left": "observation.image.exterior_image_1_left", + "right": "observation.image.exterior_image_2_left", + }, + "droid_plus_lerobot_320x180_20260406": { + "wrist": "observation.image.wrist_image_left", + "left": "observation.image.exterior_image_1_left", + "right": "observation.image.exterior_image_2_left", + }, + "droid_plus_lerobot_640x360_20260412_sharded": { + "wrist": "observation.image.wrist_image_left", + "left": "observation.image.exterior_image_1_left", + "right": "observation.image.exterior_image_2_left", + }, + "droid_plus_lerobot_640x360_20260412": { + "wrist": "observation.image.wrist_image_left", + "left": "observation.image.exterior_image_1_left", + "right": "observation.image.exterior_image_2_left", + }, +} + +STATE_FEATURES = { + "droid_lerobot_20260115_no_noops": "observation.state", + "droid_plus_lerobot_320x180_20260406_sharded": "observation.state.cartesian_position", + "droid_plus_lerobot_320x180_20260406": "observation.state.cartesian_position", + "droid_plus_lerobot_640x360_20260412_sharded": "observation.state.cartesian_position", + "droid_plus_lerobot_640x360_20260412": "observation.state.cartesian_position", +} + +ACTION_FEATURES = { + "droid_lerobot_20260115_no_noops": "action", + "droid_plus_lerobot_320x180_20260406_sharded": "action.gripper_position", + "droid_plus_lerobot_320x180_20260406": "action.gripper_position", + "droid_plus_lerobot_640x360_20260412_sharded": "action.gripper_position", + "droid_plus_lerobot_640x360_20260412": "action.gripper_position", +} + +IS_FLAT_ACTION = { + "droid_lerobot_20260115_no_noops": True, + "droid_plus_lerobot_320x180_20260406_sharded": False, + "droid_plus_lerobot_320x180_20260406": False, + "droid_plus_lerobot_640x360_20260412_sharded": False, + "droid_plus_lerobot_640x360_20260412": False, +} + +HAS_MULTI_LANGUAGE_ANNOTATIONS = { + "droid_lerobot_20260115_no_noops": False, + "droid_plus_lerobot_320x180_20260406_sharded": True, + "droid_plus_lerobot_320x180_20260406": True, + "droid_plus_lerobot_640x360_20260412_sharded": True, + "droid_plus_lerobot_640x360_20260412": True, +} + +IS_GRIPPER_ACTION_FLIPPED = { + "droid_lerobot_20260115_no_noops": False, + "droid_plus_lerobot_320x180_20260406_sharded": True, + "droid_plus_lerobot_320x180_20260406": True, + "droid_plus_lerobot_640x360_20260412_sharded": True, + "droid_plus_lerobot_640x360_20260412": True, +} + +_JOINT_ACTION_FEATURE = "action.joint_position" +_JOINT_STATE_FEATURE = "observation.state.joint_positions" +_GRIPPER_STATE_FEATURE = "observation.state.gripper_position" diff --git a/examples/worldmodel/cosmos3/vendor/cosmos_framework/data/generator/action/datasets/droid_merged_lerobot_dataset.py b/examples/worldmodel/cosmos3/vendor/cosmos_framework/data/generator/action/datasets/droid_merged_lerobot_dataset.py new file mode 100644 index 000000000000..ffd6d1f5bd0c --- /dev/null +++ b/examples/worldmodel/cosmos3/vendor/cosmos_framework/data/generator/action/datasets/droid_merged_lerobot_dataset.py @@ -0,0 +1,97 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: OpenMDW-1.1 + +"""DROID-Merged LeRobot dataset used by the action FD DROID recipe. + +This wrapper keeps all sample parsing, action construction, video loading, and +normalization in ``DROIDLeRobotDataset``. It only registers a Cosmos3-DROID +parent directory (``success/`` + ``failure/``) as a layout the parent class can +consume. +""" + +from __future__ import annotations + +import os +from pathlib import Path + +from cosmos_framework.data.generator.action.datasets.droid_lerobot_dataset import DROIDLeRobotDataset +from cosmos_framework.data.generator.action.datasets.droid_lerobot_dataset_config import ( + ACTION_FEATURES, + HAS_MULTI_LANGUAGE_ANNOTATIONS, + IMAGE_FEATURES, + IS_FLAT_ACTION, + IS_GRIPPER_ACTION_FLIPPED, + LEROBOT_ROOTS, + STATE_FEATURES, +) + +_TEMPLATE_VERSION = "droid_plus_lerobot_640x360_20260412" +_PARENT_EE_POSE_ACTION_SPACE = "midtrain" + + +class DROIDMergedLeRobotDataset(DROIDLeRobotDataset): + """DROID dataset wrapper for parent roots containing success/failure splits. + + ``DROIDLeRobotDataset`` already supports merged DROID layouts when the root + basename is known in ``droid_lerobot_dataset_config.py``. This class adds the + same metadata for arbitrary local parent directories such as + ``datasets/Cosmos3-DROID`` while delegating all data semantics to the parent. + """ + + def __init__(self, *args, action_space: str = "ee_pose", **kwargs) -> None: + root = kwargs.get("root") + if root is None and args: + root = args[0] + if root is None: + raise TypeError("DROIDMergedLeRobotDataset requires a root argument.") + + root_path = Path(root) + self._register_root_layout(root_path) + + if action_space == "ee_pose": + # Parent DROIDLeRobotDataset currently names this same 10-D + # [pos_delta, rot6d_delta, gripper] branch internally. + action_space = _PARENT_EE_POSE_ACTION_SPACE + + super().__init__(*args, action_space=action_space, **kwargs) + + @classmethod + def _register_root_layout(cls, root: Path) -> None: + version = os.path.basename(root) + if version in LEROBOT_ROOTS: + return + + roots = cls._discover_split_roots(root) + LEROBOT_ROOTS[version] = roots + IMAGE_FEATURES[version] = dict(IMAGE_FEATURES[_TEMPLATE_VERSION]) + STATE_FEATURES[version] = STATE_FEATURES[_TEMPLATE_VERSION] + ACTION_FEATURES[version] = ACTION_FEATURES[_TEMPLATE_VERSION] + IS_FLAT_ACTION[version] = IS_FLAT_ACTION[_TEMPLATE_VERSION] + HAS_MULTI_LANGUAGE_ANNOTATIONS[version] = HAS_MULTI_LANGUAGE_ANNOTATIONS[_TEMPLATE_VERSION] + IS_GRIPPER_ACTION_FLIPPED[version] = IS_GRIPPER_ACTION_FLIPPED[_TEMPLATE_VERSION] + + @staticmethod + def _is_lerobot_root(path: Path) -> bool: + return (path / "meta" / "info.json").is_file() + + @classmethod + def _discover_split_roots(cls, root: Path) -> list[str] | None: + if cls._is_lerobot_root(root): + return None + + discovered: list[str] = [] + for split_name in ("success", "failure"): + split_root = root / split_name + if cls._is_lerobot_root(split_root): + discovered.append(split_name) + continue + if split_root.is_dir(): + discovered.extend( + f"{split_name}/{path.name}" for path in sorted(split_root.iterdir()) if cls._is_lerobot_root(path) + ) + + if discovered: + return discovered + raise FileNotFoundError( + f"{root} is not a LeRobot root and no success/failure LeRobot splits were found under it." + ) diff --git a/examples/worldmodel/cosmos3/vendor/cosmos_framework/data/generator/action/datasets/fractal_lerobot_dataset.py b/examples/worldmodel/cosmos3/vendor/cosmos_framework/data/generator/action/datasets/fractal_lerobot_dataset.py new file mode 100644 index 000000000000..3543abcf8e79 --- /dev/null +++ b/examples/worldmodel/cosmos3/vendor/cosmos_framework/data/generator/action/datasets/fractal_lerobot_dataset.py @@ -0,0 +1,190 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: OpenMDW-1.1 + +"""Fractal LeRobot dataset — Google Robot RT-1. + +Robot: google_robot +87,212 episodes, 3,786,400 frames, 599 tasks, fps=3 +state: [x, y, z, rx, ry, rz, rw, gripper] (8D, quaternion) +action: [x, y, z, roll, pitch, yaw, gripper] (7D, delta) +video: observation.images.image (256×320) +""" + +from __future__ import annotations + +import random +from pathlib import Path +from typing import Any, Literal + +import numpy as np +import torch + +from cosmos_framework.data.generator.action.utils.action_spec import ActionSpec, Gripper, Pos, Rot, build_action_spec +from cosmos_framework.data.generator.action.datasets.base_dataset import ActionBaseDataset +from cosmos_framework.data.generator.action.utils.pose_utils import ( + build_abs_pose_from_components, + pose_abs_to_rel, +) + +PoseConvention = Literal["backward_framewise"] +Viewpoint = Literal["ego_view"] + +_IMAGE_FEATURE = "observation.images.image" +_STATE_FEATURE = "observation.state" +_ACTION_FEATURE = "action" + +# These episodes contain base motion, which breaks the fixed-base Google Robot +# action assumption used by training and the viewer. +_SKIPPED_EPISODE_IDS: frozenset[int] = frozenset({29, 189, 382}) + +# Google Robot raw EE frame has x/y axes rotated ~90° around z compared to +# OpenCV convention. Rz(-90°) as a right-multiply corrects this: +# new_x = -old_y (rightward), new_y = old_x (downward), z unchanged (approach). +_GOOGLE_ROBOT_TO_OPENCV = np.array([[0, 1, 0], [-1, 0, 0], [0, 0, 1]], dtype=np.float32) + +# --------------------------------------------------------------------------- +# TCP → flange (gripper body) offset +# --------------------------------------------------------------------------- +# The fractal dataset records EE poses at ``link_gripper_tcp`` — a calibrated +# tool-center-point 164 mm past the gripper body (``link_gripper``), roughly +# at the fingertip. Re-referencing to ``link_gripper`` removes the +# calibration-dependent tilt and places the frame at the last actuated link. +# +# T = oMf[link_gripper_tcp]⁻¹ · oMf[link_gripper], computed via pinocchio FK +# at the neutral config from the SimplerEnv URDF. +# fmt: off +_TCP_TO_FLANGE = np.array([ + [+0.9999897671, -0.0008686425, +0.0044397163, -0.0050618476], + [+0.0008745501, +0.9999987346, -0.0013288658, -0.0016717725], + [-0.0044385564, +0.0013327349, +0.9999892615, -0.1635144743], + [+0.0000000000, +0.0000000000, +0.0000000000, +1.0000000000], +], dtype=np.float32) +# fmt: on + +_NORMALIZER_PATH = Path(__file__).parent.parent / "normalizer_stats/fractal_lerobot_stats.json" + + +class FractalLeRobotDataset(ActionBaseDataset): + """Fractal (Google RT-1) dataset with 10D cartesian actions: + + [pos_delta(3), rot6d_delta(6), gripper(1)] + + Expects a LeRobot v2 dataset with: + * ``observation.images.image``: ego-view RGB video (256×320). + * ``observation.state``: 8D EE pose ``[x, y, z, rx, ry, rz, rw, gripper]`` + in TCP frame with quaternion (x, y, z, w) order. + * ``action``: 7D delta ``[x, y, z, roll, pitch, yaw, gripper]``; only the + gripper column (index 6) is used — SE(3) actions are derived from state. + + Episodes in ``_SKIPPED_EPISODE_IDS`` (base-motion outliers) are dropped. + """ + + def __init__( + self, + root: str, + fps: float = 3.0, + chunk_length: int = 16, + mode: str = "joint", + pose_convention: PoseConvention = "backward_framewise", + tolerance_s: float = 1e-4, + viewpoint: Viewpoint = "ego_view", + action_normalization: str | None = None, + sample_stride: int = 1, + ) -> None: + if viewpoint != "ego_view": + raise NotImplementedError("FractalLeRobotDataset only supports ego_view.") + super().__init__( + root=root, + domain_name="fractal", + fps=fps, + chunk_length=chunk_length, + mode=mode, + pose_convention=pose_convention, + tolerance_s=tolerance_s, + viewpoint=viewpoint, + action_normalization=action_normalization, + sample_stride=sample_stride, + ) + + # Drop rows belonging to known-bad episodes (base motion outliers). + n_before = len(self._rows) + self._rows = [row for row in self._rows if int(row["episode_index"]) not in _SKIPPED_EPISODE_IDS] + n_dropped = n_before - len(self._rows) + if n_dropped: + import logging + logging.getLogger(__name__).info( + "FractalLeRobotDataset: dropped %d / %d rows from episodes %s", + n_dropped, + n_before, + sorted(_SKIPPED_EPISODE_IDS), + ) + + @property + def action_dim(self) -> int: + """Action dimensionality: position(3) + 6D rotation(6) + gripper(1) = 10.""" + return 10 + + def _action_spec(self) -> ActionSpec: + return build_action_spec(Pos(dim=3), Rot("rot6d"), Gripper()) + + @classmethod + def _stats_path(cls) -> Path: + return _NORMALIZER_PATH + + def __getitem__(self, idx: int) -> dict[str, Any]: + mode = self._choose_mode() + idx = int(idx) + row_idx = idx * self._sample_stride + observation_rows = self._rows[row_idx : row_idx + self._chunk_length + 1] + action_rows = observation_rows[: self._chunk_length] + + episode = self._episodes[int(observation_rows[0]["episode_index"])] + task = self._tasks[int(observation_rows[0]["task_index"])] + ai_caption = random.choice([part.strip() for part in task.split(" | ") if part.strip()] or [task]) + + video = self._load_video(episode, observation_rows) + raw_action, initial_pose = self._build_raw_action(observation_rows, action_rows) + + return self._build_result( + mode=mode, + video=video, + action=raw_action, + ai_caption=ai_caption, + initial_pose=initial_pose, + ) + + def _load_video(self, episode: dict[str, Any], observation_rows: list[dict[str, Any]]) -> torch.Tensor: + # lerobot is a heavy, optional ("train" extra) dependency; import lazily. + from lerobot.datasets.video_utils import decode_video_frames + + timestamps = [float(row["timestamp"]) for row in observation_rows] + return decode_video_frames( + self._video_path(episode, _IMAGE_FEATURE), + [float(episode.get(f"videos/{_IMAGE_FEATURE}/from_timestamp", 0.0)) + ts for ts in timestamps], + self._tolerance_s, + ) + + def _build_raw_action( + self, + observation_rows: list[dict[str, Any]], + action_rows: list[dict[str, Any]], + ) -> tuple[torch.Tensor, torch.Tensor]: + # State layout: [x, y, z, rx, ry, rz, rw, gripper] (T+1 frames) + # Quaternion order: (rx, ry, rz, rw) matches scipy's (x, y, z, w). + state = np.asarray([row[_STATE_FEATURE] for row in observation_rows], dtype=np.float32) # [T+1, 8] + poses_abs = build_abs_pose_from_components(state[:, 0:3], state[:, 3:7], "quat_xyzw") # [T+1, 4, 4] + + # 1. TCP → flange: shift from link_gripper_tcp to link_gripper + poses_abs = poses_abs @ _TCP_TO_FLANGE + # 2. Kinematics → OpenCV convention (rotation only) + poses_abs[:, :3, :3] = poses_abs[:, :3, :3] @ _GOOGLE_ROBOT_TO_OPENCV + + initial_pose = torch.from_numpy(poses_abs[0].copy()).float() + poses_rel = pose_abs_to_rel(poses_abs, rotation_format="rot6d", pose_convention=self._pose_convention) + + gripper = np.asarray( + [row[_ACTION_FEATURE][6] for row in action_rows], dtype=np.float32 + ).reshape(-1, 1) # [T, 1] + + action = np.concatenate([poses_rel[-self._chunk_length :], gripper[-self._chunk_length :]], axis=-1) + return torch.from_numpy(action).float(), initial_pose diff --git a/examples/worldmodel/cosmos3/vendor/cosmos_framework/data/generator/action/datasets/human_hand_pose_lerobot_dataset.py b/examples/worldmodel/cosmos3/vendor/cosmos_framework/data/generator/action/datasets/human_hand_pose_lerobot_dataset.py new file mode 100644 index 000000000000..e72c63753537 --- /dev/null +++ b/examples/worldmodel/cosmos3/vendor/cosmos_framework/data/generator/action/datasets/human_hand_pose_lerobot_dataset.py @@ -0,0 +1,219 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: OpenMDW-1.1 + +"""Bimanual human hand-pose dataset in LeRobot v3 format.""" + +from __future__ import annotations + +import random +from pathlib import Path +from typing import Any, Literal + +import numpy as np +import pyarrow.parquet as pq +import torch +from lerobot.datasets.video_utils import decode_video_frames + +from cosmos_framework.data.generator.action.utils.action_spec import ActionSpec, Pos, Rot, build_action_spec +from cosmos_framework.data.generator.action.datasets.base_dataset import ActionBaseDataset +from cosmos_framework.data.generator.action.utils.pose_utils import build_abs_pose_from_components, pose_abs_to_rel + +PoseConvention = Literal["backward_framewise"] +Viewpoint = Literal["ego_view"] + +_HAND_RIGHT_POSITION_KEY = "observation.state.hand_right_cam" +_HAND_RIGHT_ROTATION_KEY = "observation.state.hand_right_cam_rotation" +_HAND_LEFT_POSITION_KEY = "observation.state.hand_left_cam" +_HAND_LEFT_ROTATION_KEY = "observation.state.hand_left_cam_rotation" +_CAM_POSITION_KEY = "observation.state.camera_position" +_CAM_ROTATION_KEY = "observation.state.camera_rotation" +_IMAGE_FEATURE = "observation.images.main" + +_NUM_JOINTS = 21 +_WRIST_JOINT_IDX = 0 +_FINGERTIP_JOINT_IDXS = (4, 8, 12, 16, 20) +_RAW_ACTION_DIM = 57 +_NORMALIZER_PATH = Path(__file__).parent.parent / "normalizer_stats/human_hand_pose_lerobot_stats.json" + +# Rotate the source wrist frames into the unified convention: +# X = thumb-to-pinky, Y = outward palm normal, Z = wrist-to-fingertips. +_WRIST_FRAME_ALIGNMENT = np.array( + [[0, 1, 0, 0], [-1, 0, 0, 0], [0, 0, 1, 0], [0, 0, 0, 1]], + dtype=np.float32, +) + + +class HumanHandPoseLeRobotDataset(ActionBaseDataset): + """Bimanual human hand-pose forward-dynamics data. + + The 57D action layout is + ``[camera(9), right_wrist(9), right_fingertips(15), + left_wrist(9), left_fingertips(15)]``. Camera and wrist poses use + framewise 3D translation plus rot6d deltas. Each hand's five fingertip + positions are expressed in that frame's aligned wrist coordinate system. + + Source video and pose annotations are sampled at 30 FPS by default and + decoded at 15 FPS for Cosmos3-Nano, yielding 17 video frames and 16 action + transitions for the default chunk. + """ + + def __init__( + self, + root: str, + fps: float = 15.0, + chunk_length: int = 16, + mode: str = "forward_dynamics", + pose_convention: PoseConvention = "backward_framewise", + tolerance_s: float = 2e-4, + viewpoint: Viewpoint = "ego_view", + action_normalization: str | None = "quantile", + sample_stride: int = 1, + image_key: str = _IMAGE_FEATURE, + ) -> None: + if viewpoint != "ego_view": + raise NotImplementedError("Human hand-pose data only supports ego_view.") + super().__init__( + root=root, + domain_name="hand_pose", + fps=fps, + chunk_length=chunk_length, + mode=mode, + pose_convention=pose_convention, + tolerance_s=tolerance_s, + viewpoint=viewpoint, + action_normalization=action_normalization, + sample_stride=sample_stride, + ) + source_fps = float(self._info["fps"]) + source_stride = source_fps / self._fps + if not source_stride.is_integer(): + raise ValueError(f"Source FPS {source_fps} must be an integer multiple of target FPS {self._fps}.") + self._source_stride = int(source_stride) + self._image_key = image_key + required_source_steps = self._source_stride * self._chunk_length + self._valid_starts: list[int] = [] + episode_start = 0 + while episode_start < len(self._rows): + episode_index = int(self._rows[episode_start]["episode_index"]) + episode_end = episode_start + 1 + while episode_end < len(self._rows) and int(self._rows[episode_end]["episode_index"]) == episode_index: + episode_end += 1 + self._valid_starts.extend( + range(episode_start, max(episode_start, episode_end - required_source_steps), self._sample_stride) + ) + episode_start = episode_end + subtasks_path = self._root / "meta" / "subtasks.parquet" + self._subtasks = ( + {int(row["subtask_index"]): str(row["subtask"]) for row in pq.read_table(subtasks_path).to_pylist()} + if subtasks_path.exists() + else {} + ) + + @property + def action_dim(self) -> int: + return _RAW_ACTION_DIM + + def _action_spec(self) -> ActionSpec: + return build_action_spec( + Pos(prefix="camera"), + Rot("rot6d", prefix="camera"), + Pos(prefix="right_wrist"), + Rot("rot6d", prefix="right_wrist"), + Pos(dim=15, prefix="right_fingertip"), + Pos(prefix="left_wrist"), + Rot("rot6d", prefix="left_wrist"), + Pos(dim=15, prefix="left_fingertip"), + ) + + @classmethod + def _stats_path(cls) -> Path: + return _NORMALIZER_PATH + + def __len__(self) -> int: + return len(self._valid_starts) + + def __getitem__(self, idx: int) -> dict[str, Any]: + mode = self._choose_mode() + start = self._valid_starts[int(idx)] + stop = start + self._source_stride * self._chunk_length + 1 + rows = self._rows[start : stop : self._source_stride] + if len(rows) != self._chunk_length + 1: + raise IndexError(f"Incomplete hand-pose window at index {idx}.") + episode_index = int(rows[0]["episode_index"]) + if any(int(row["episode_index"]) != episode_index for row in rows): + raise IndexError(f"Hand-pose window at index {idx} crosses an episode boundary.") + + episode = self._episodes[episode_index] + video = self._load_video(episode, rows) + raw_action = self._build_raw_action(rows) + subtask_index = int(rows[0].get("subtask_index", -1)) + task = self._tasks[int(rows[0]["task_index"])] + caption = self._subtasks.get(subtask_index, task) + ai_caption = random.choice([part.strip() for part in caption.split(" | ") if part.strip()] or [caption]) + + result = self._build_result(mode=mode, video=video, action=raw_action, ai_caption=ai_caption) + if self.action_normalization is not None: + result["action"] = result["action"].clamp(-1.0, 1.0) + return result + + def _load_video(self, episode: dict[str, Any], rows: list[dict[str, Any]]) -> torch.Tensor: + timestamps = [float(row["timestamp"]) for row in rows] + from_timestamp = float(episode.get(f"videos/{self._image_key}/from_timestamp", 0.0)) + return decode_video_frames( + self._video_path(episode, self._image_key), + [from_timestamp + timestamp for timestamp in timestamps], + self._tolerance_s, + ) + + @staticmethod + def _finger_positions_in_wrist_frame(position_data: np.ndarray, wrist_poses: np.ndarray) -> np.ndarray: + future_positions = position_data[1:].reshape(-1, _NUM_JOINTS, 3) + fingertips = future_positions[:, _FINGERTIP_JOINT_IDXS, :] + fingertips_h = np.concatenate( + [fingertips, np.ones((*fingertips.shape[:-1], 1), dtype=np.float32)], + axis=-1, + ) + wrist_inv = np.linalg.inv(wrist_poses[1:]) + fingertips_wrist = np.einsum("tij,tnj->tni", wrist_inv, fingertips_h)[..., :3] + return fingertips_wrist.reshape(len(future_positions), -1) + + def _build_raw_action(self, rows: list[dict[str, Any]]) -> torch.Tensor: + def values(key: str) -> np.ndarray: + return np.asarray([row[key] for row in rows], dtype=np.float32) + + camera_pose = build_abs_pose_from_components(values(_CAM_POSITION_KEY), values(_CAM_ROTATION_KEY), "quat_xyzw") + + right_positions = values(_HAND_RIGHT_POSITION_KEY) + right_rotations = values(_HAND_RIGHT_ROTATION_KEY).reshape(-1, _NUM_JOINTS, 4) + left_positions = values(_HAND_LEFT_POSITION_KEY) + left_rotations = values(_HAND_LEFT_ROTATION_KEY).reshape(-1, _NUM_JOINTS, 4) + + right_wrist_camera = ( + build_abs_pose_from_components(right_positions[:, :3], right_rotations[:, _WRIST_JOINT_IDX], "quat_xyzw") + @ _WRIST_FRAME_ALIGNMENT + ) + left_wrist_camera = ( + build_abs_pose_from_components(left_positions[:, :3], left_rotations[:, _WRIST_JOINT_IDX], "quat_xyzw") + @ _WRIST_FRAME_ALIGNMENT + ) + + right_wrist_world = camera_pose @ right_wrist_camera + left_wrist_world = camera_pose @ left_wrist_camera + action = np.concatenate( + [ + pose_abs_to_rel(camera_pose, rotation_format="rot6d", pose_convention=self._pose_convention), + pose_abs_to_rel(right_wrist_world, rotation_format="rot6d", pose_convention=self._pose_convention), + self._finger_positions_in_wrist_frame(right_positions, right_wrist_camera), + pose_abs_to_rel(left_wrist_world, rotation_format="rot6d", pose_convention=self._pose_convention), + self._finger_positions_in_wrist_frame(left_positions, left_wrist_camera), + ], + axis=-1, + ) + if action.shape != (self._chunk_length, _RAW_ACTION_DIM): + raise ValueError( + f"Expected hand-pose action shape {(self._chunk_length, _RAW_ACTION_DIM)}, got {action.shape}." + ) + return torch.from_numpy(action).float() + + +__all__ = ["HumanHandPoseLeRobotDataset"] diff --git a/examples/worldmodel/cosmos3/vendor/cosmos_framework/data/generator/action/datasets/libero_lerobot_dataset.py b/examples/worldmodel/cosmos3/vendor/cosmos_framework/data/generator/action/datasets/libero_lerobot_dataset.py new file mode 100644 index 000000000000..7c5d1c62fe28 --- /dev/null +++ b/examples/worldmodel/cosmos3/vendor/cosmos_framework/data/generator/action/datasets/libero_lerobot_dataset.py @@ -0,0 +1,330 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: OpenMDW-1.1 + +"""LIBERO LeRobot dataset (frame-wise-relative action policy). + +Mirrors ``DROIDLeRobotDataset``: reads the LeRobot parquet directly, windows by +frame index, and decodes video at each frame's REAL timestamp. That makes it +FPS-agnostic — it works with the 10 FPS community ``lerobot/libero_*`` datasets +and a 20 FPS conversion alike, without LeRobot's ``delta_timestamps`` grid (which +rejects any window whose synthetic timestamps don't land on real frames). + +Action layout (``frame_wise_relative``): the stored 7D ``action`` is already a +per-frame delta ``[dpos(3), drot_axisangle(3), gripper(1)]``; only the rotation is +re-encoded to the requested ``rotation_space`` -> ``[dpos(3), rot6d(6), gripper(1)]`` +(10D for ``6d``). + +NOTE on FPS / stats fidelity: the bundled ``quantile_rot`` stats were computed on +a 20 FPS conversion. Per-frame deltas at 10 FPS span 2x the wall-clock motion, so +use a 20 FPS LIBERO dataset (or recompute stats for the dataset's FPS). +Loading/training is correct at any FPS regardless. +""" + +from __future__ import annotations + +import json +import random +from pathlib import Path +from typing import Any, Literal + +import numpy as np +import pyarrow.parquet as pq +import torch +import torch.nn.functional as F + +from cosmos_framework.data.generator.action.utils.action_spec import ActionSpec, Gripper, Pos, Rot, build_action_spec +from cosmos_framework.data.generator.action.datasets.base_dataset import ActionBaseDataset +from cosmos_framework.data.generator.action.libero_pose_utils import libero_action_dim, libero_rotation_format +from cosmos_framework.data.generator.action.utils.pose_utils import convert_rotation +from cosmos_framework.utils import log + +CameraMode = Literal["image", "wrist_image", "concat_view"] +RotationSpace = Literal["3d", "6d", "9d"] + +_ACTION_FEATURE = "action" +_IMAGE_FEATURE = "observation.images.image" +_WRIST_FEATURE = "observation.images.wrist_image" +_STAT_KEYS = ("mean", "std", "min", "max", "q01", "q99") +_NORMALIZERS_DIR = Path(__file__).parent.parent / "normalizer_stats" + +_VIEWPOINT_BY_CAMERA = { + "image": "third_person_view", + "wrist_image": "wrist_view", + "concat_view": "concat_view", +} + + +class LIBEROLeRobotDataset(ActionBaseDataset): + """LIBERO action-policy dataset with frame-wise-relative rot6d actions. + + 10D ``[pos_delta(3), rot6d_delta(6), gripper(1)]`` (for ``rotation_space='6d'``), + ``concat_view`` third-person + wrist video, and ``quantile_rot`` normalization + against the bundled stats. Reads parquet + decodes video at real timestamps, + so the requested ``fps`` is metadata only (it sets ``conditioning_fps`` and the + prompt duration); frame windows always use the data's actual frames. + """ + + def __init__( + self, + root: str, + fps: float = 20.0, + chunk_length: int = 16, + mode: str = "wam", + tolerance_s: float = 1e-4, + camera_mode: CameraMode = "concat_view", + image_size: int = 256, + action_space: str = "frame_wise_relative", + rotation_space: RotationSpace = "6d", + pose_coordinate_frame: str = "native", + embodiment_type: str = "libero", + action_normalization: str | None = "quantile_rot", + action_stats_path: str | None = None, + split: str = "train", + val_ratio: float = 0.01, + seed: int = 0, + sample_stride: int = 1, + ) -> None: + if action_space != "frame_wise_relative": + raise NotImplementedError( + f"This LIBERO dataset only supports action_space='frame_wise_relative', got {action_space!r}." + ) + if camera_mode not in _VIEWPOINT_BY_CAMERA: + raise ValueError(f"Unsupported camera_mode={camera_mode!r}. Use image/wrist_image/concat_view.") + split = split.lower().strip() + if split not in {"train", "val", "valid", "validation", "eval", "test", "full"}: + raise ValueError(f"Unsupported split={split!r}. Use train/val/full.") + if chunk_length % 4 != 0: + raise ValueError(f"chunk_length must be divisible by 4, got {chunk_length}.") + + super().__init__( + root=root, + domain_name=embodiment_type, + fps=fps, + chunk_length=chunk_length, + mode=mode, + pose_convention="backward_framewise", # unused for frame_wise deltas; satisfies the base assert + tolerance_s=tolerance_s, + viewpoint=_VIEWPOINT_BY_CAMERA[camera_mode], + # frame_wise_relative ⇔ backward_framewise idle semantics. quantile_rot is a + # LIBERO convention -> normalize with the "quantile" formula on raw-rotation + # stats (see _load_norm_stats); pass the method the base will call. + action_normalization=None if action_normalization is None else "quantile", + sample_stride=sample_stride, + ) + # FPS-agnostic loader: trust the dataset's NATIVE fps for conditioning_fps / + # prompt duration so the metadata is truthful (10 for the public + # lerobot/libero_*, 20 for a 20 FPS conversion). Frame sampling uses each + # frame's real timestamp regardless, so the requested ``fps`` is ignored here. + info_fps = self._info.get("fps") + if info_fps: + if int(info_fps) != int(fps): + log.info(f"Using dataset native fps={info_fps} for conditioning (requested {fps}).") + self._fps = float(info_fps) + self._dt = 1.0 / self._fps + self._camera_mode = camera_mode + self._image_size = int(image_size) + self._rotation_space = rotation_space.lower().strip() + self._pose_coordinate_frame = pose_coordinate_frame + self._embodiment_type = embodiment_type + self._requested_normalization = action_normalization + # quantile_rot normalizes against the raw (un-orthonormalized) rotation stats + # under "global_raw"; everything else uses "global". + self._stats_key = "global_raw" if action_normalization == "quantile_rot" else "global" + self._stats_file = self._resolve_stats_file(action_stats_path) + + if self._camera_mode == "image": + self._video_keys = [_IMAGE_FEATURE] + elif self._camera_mode == "wrist_image": + self._video_keys = [_WRIST_FEATURE] + else: + self._video_keys = [_IMAGE_FEATURE, _WRIST_FEATURE] + + # Compact, lazy frame index (mirrors DROIDLeRobotDataset): read only the + # columns the sample builder needs into contiguous arrays, ordered by global + # frame index, so DataLoader worker forks share them copy-on-write. + index_parts, episode_parts, task_parts, ts_parts, action_parts = [], [], [], [], [] + for path in sorted((self._root / "data").glob("chunk-*/file-*.parquet")): + table = pq.read_table(path, columns=["index", "episode_index", "task_index", "timestamp", _ACTION_FEATURE]) + index_parts.append(table["index"].to_numpy()) + episode_parts.append(table["episode_index"].to_numpy()) + task_parts.append(table["task_index"].to_numpy()) + ts_parts.append(table["timestamp"].to_numpy()) + action_parts.append(np.asarray(table[_ACTION_FEATURE].to_pylist(), dtype=np.float32)) + if not index_parts: + raise FileNotFoundError(f"No data parquet found under {self._root / 'data'}.") + order = np.argsort(np.concatenate(index_parts).astype(np.int64), kind="stable") + self._row_episode = np.concatenate(episode_parts).astype(np.int64)[order] + self._row_task = np.concatenate(task_parts).astype(np.int64)[order] + self._row_timestamp = np.concatenate(ts_parts).astype(np.float64)[order] + self._row_action = np.concatenate(action_parts, axis=0).astype(np.float32)[order] + + assert np.all(np.diff(self._row_episode) >= 0), "episode_index not contiguous after sorting by frame index" + ep_vals, ep_starts, ep_counts = np.unique(self._row_episode, return_index=True, return_counts=True) + + # Deterministic per-episode train/val split (seeded; same on every rank). + keep = self._split_episode_ids(ep_vals.tolist(), split, val_ratio, seed) + kept = np.array([int(v) in keep for v in ep_vals], dtype=bool) + self._ep_vals = ep_vals.astype(np.int64)[kept] + self._ep_starts = ep_starts.astype(np.int64)[kept] + kept_counts = ep_counts.astype(np.int64)[kept] + # Within-episode windows only: total - n_kept_episodes * chunk_length valid samples. + self._valid_cum = np.cumsum(np.maximum(0, kept_counts - self._chunk_length)).astype(np.int64) + + log.info( + f"Loaded LIBERO dataset root={self._root} split={split!r} camera_mode={camera_mode!r} " + f"fps={self._fps} kept_episodes={len(self._ep_vals)}/{len(ep_vals)} " + f"valid_indices={int(self._valid_cum[-1]) if self._valid_cum.size else 0}" + ) + + # ---- spec / dims ------------------------------------------------------- + + @property + def action_dim(self) -> int: + return libero_action_dim(self._rotation_space) + + def _action_spec(self) -> ActionSpec: + return build_action_spec(Pos(), Rot(libero_rotation_format(self._rotation_space)), Gripper()) + + @classmethod + def _stats_path(cls) -> Path: + # Base classmethod fallback; the instance uses self._stats_file (which also + # honors action_stats_path + the rotation/coordinate-frame-specific filename). + return _NORMALIZERS_DIR / "libero_native_frame_wise_relative_rot6d.json" + + # ---- normalization (nested global/global_raw + quantile_rot) ------------ + + def _bundled_stats_filename(self) -> str: + rotation_suffix = {"3d": "3d", "6d": "rot6d", "9d": "rot9d"}.get(self._rotation_space) + if rotation_suffix is None: + raise ValueError(f"Unsupported rotation_space={self._rotation_space!r}.") + action_space = "frame_wise_relative" + return f"{self._embodiment_type}_{self._pose_coordinate_frame}_{action_space}_{rotation_suffix}.json" + + def _resolve_stats_file(self, action_stats_path: str | None) -> Path: + if action_stats_path: + p = Path(action_stats_path) + if not p.is_absolute(): + p = _NORMALIZERS_DIR / p.name + if not p.exists(): + raise FileNotFoundError(f"action_stats_path not found: {action_stats_path!r}") + return p + p = _NORMALIZERS_DIR / self._bundled_stats_filename() + if not p.exists(): + raise FileNotFoundError( + f"Bundled LIBERO stats not found at {p}. Pass action_stats_path or recompute stats." + ) + return p + + def _load_norm_stats(self) -> dict[str, torch.Tensor]: + if self._norm_stats is None: + raw = json.loads(self._stats_file.read_text())[self._stats_key] + self._norm_stats = {k: torch.tensor(v, dtype=torch.float32) for k, v in raw.items() if k in _STAT_KEYS} + return self._norm_stats + + # ---- index helpers ----------------------------------------------------- + + @staticmethod + def _split_episode_ids(ep_ids: list[int], split: str, val_ratio: float, seed: int) -> set[int]: + if split == "full": + return set(int(v) for v in ep_ids) + if not (0.0 < val_ratio < 1.0): + raise ValueError(f"val_ratio must be in (0, 1), got {val_ratio}.") + n_val = max(1, int(round(len(ep_ids) * val_ratio))) + rng = random.Random(seed) # identical selection on every rank + val = set(int(v) for v in rng.sample(list(ep_ids), n_val)) + if split == "train": + return set(int(v) for v in ep_ids) - val + return val # val/valid/validation/eval/test + + def __len__(self) -> int: + return int(self._valid_cum[-1]) if self._valid_cum.size else 0 + + def get_shuffle_blocks(self) -> list[tuple[int, int]]: + """Per-episode ``(start, length)`` flat-index blocks for + ``ActionIterableShuffleDataset`` (shuffle block ORDER + shard across + ranks, sequential within a block).""" + blocks: list[tuple[int, int]] = [] + prev = 0 + for c in np.asarray(self._valid_cum).tolist(): + c = int(c) + if c > prev: + blocks.append((prev, c - prev)) + prev = c + return blocks + + # ---- sample build ------------------------------------------------------ + + def __getitem__(self, idx: int) -> dict[str, Any]: + # Resample a different valid window if a frame fails to decode (bounded retries). + n = len(self) + last_err: Exception | None = None + for _attempt in range(8): + try: + return self._build_item(idx) + except Exception as e: # noqa: BLE001 — skip past undecodable frames + last_err = e + log.warning(f"LIBERO: sample idx={idx} failed to load ({type(e).__name__}: {e}); resampling") + if n > 0: + idx = random.randint(0, n - 1) + raise RuntimeError(f"LIBERO: failed to load a sample after 8 resamples; last error: {last_err}") + + def _build_item(self, idx: int) -> dict[str, Any]: + mode = self._choose_mode() + idx = int(idx) + ep = int(np.searchsorted(self._valid_cum, idx, side="right")) + prev = int(self._valid_cum[ep - 1]) if ep > 0 else 0 + start = int(self._ep_starts[ep]) + (idx - prev) + episode_index = int(self._ep_vals[ep]) + episode = self._episodes[episode_index] + + stop = start + self._chunk_length + 1 + timestamps = [float(self._row_timestamp[j]) for j in range(start, stop)] + video = self._load_video(episode, timestamps) + + # frame_wise_relative: chunk per-frame deltas are the stored actions directly. + raw = self._row_action[start : start + self._chunk_length] # [chunk, 7] + action = self._build_frame_wise_action(raw) + + task = self._tasks[int(self._row_task[start])] + ai_caption = random.choice([p.strip() for p in task.split(" | ") if p.strip()] or [task]) + + extras: dict[str, Any] = {} + if self._camera_mode == "concat_view": + extras["additional_view_description"] = ( + "The left half shows the third-person view; the right half shows the wrist-mounted camera." + ) + return self._build_result(mode=mode, video=video, action=action, ai_caption=ai_caption, **extras) + + def _build_frame_wise_action(self, raw: np.ndarray) -> torch.Tensor: + raw_t = torch.from_numpy(np.ascontiguousarray(raw)).float() # [chunk, 7] + translation = raw_t[:, 0:3] + rotation_matrix = convert_rotation(raw_t[:, 3:6], input_format="axisangle", output_format="matrix") + rotation = convert_rotation( + rotation_matrix, input_format="matrix", output_format=libero_rotation_format(self._rotation_space) + ) + gripper = raw_t[:, 6:7] + return torch.cat([translation, rotation, gripper], dim=-1) # [chunk, action_dim] + + def _load_video(self, episode: dict[str, Any], timestamps: list[float]) -> torch.Tensor: + # lerobot is a heavy, optional ("train" extra) dependency; import lazily. + from lerobot.datasets.video_utils import decode_video_frames + + frames_by_view = {} + for key in self._video_keys: + from_ts = float(episode.get(f"videos/{key}/from_timestamp", 0.0)) + frames = decode_video_frames( + self._video_path(episode, key), + [from_ts + ts for ts in timestamps], + self._tolerance_s, + ) # [T, C, H, W] in [0, 1] + frames = self._resize(frames) + frames_by_view[key] = frames + if self._camera_mode == "concat_view": + # third-person (left) + wrist (right), horizontally concatenated -> [T, C, H, 2W] + return torch.cat([frames_by_view[_IMAGE_FEATURE], frames_by_view[_WRIST_FEATURE]], dim=-1) + return frames_by_view[self._video_keys[0]] + + def _resize(self, frames: torch.Tensor) -> torch.Tensor: + if frames.shape[-1] == self._image_size and frames.shape[-2] == self._image_size: + return frames + return F.interpolate(frames, size=(self._image_size, self._image_size), mode="bilinear", align_corners=False) diff --git a/examples/worldmodel/cosmos3/vendor/cosmos_framework/data/generator/action/datasets/robocasa_lerobot_dataset.py b/examples/worldmodel/cosmos3/vendor/cosmos_framework/data/generator/action/datasets/robocasa_lerobot_dataset.py new file mode 100644 index 000000000000..68cc0e0ba6f8 --- /dev/null +++ b/examples/worldmodel/cosmos3/vendor/cosmos_framework/data/generator/action/datasets/robocasa_lerobot_dataset.py @@ -0,0 +1,573 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: OpenMDW-1.1 + +"""RoboCasa (PandaOmron) LeRobot action-policy dataset. + +Reads the RoboCasa ``PhysicalAI-Robotics-Manipulation-Kitchen-Demos`` LeRobot +export (codebase_version v2.1) through the official ``lerobot`` reader (via +:class:`BaseActionLeRobotDataset`), so the v2.1 per-episode layout is handled by +the library rather than a hand-rolled parquet reader. + +RoboCasa's native 12D action is +``[base_motion(4), control_mode(1), eef_pos(3), eef_rot_axisangle(3), gripper(1)]`` +(see ``meta/modality.json``). + +With ``use_base_action=False`` the base channels are dropped and only the arm delta is kept, +re-encoding rotation axis-angle -> rot6d. That 10D contract suits task sets whose base never +moves; it reduces the embodiment to a stationary Franka arm, i.e. the LIBERO/DROID regime: + + 12D -> slice [5:8]+[8:11]+[11] -> [pos(3), rot_axisangle(3), gripper(1)] (7D) + -> axisangle->rot6d -> [pos(3), rot6d(6), gripper(1)] (10D) + +The stored ``eef_pos``/``eef_rot`` are already per-frame OSC deltas, so this is a +``frame_wise_relative`` (``backward_framewise``) action — identical semantics to +``LIBEROLeRobotDataset._build_frame_wise_action``, only the slice indices differ. + +Cameras: RoboCasa exports 3 views (``robot0_eye_in_hand`` wrist + ``agentview_left`` +/ ``agentview_right`` third-person). ``concat_view`` tiles them wrist-on-top, +left/right-on-bottom (mirrors ``DROIDLeRobotDataset._compose_multi_view``). +""" + +from __future__ import annotations + +import glob +import os +import random +from typing import Any + +import torch +import torch.nn.functional as F + +from cosmos_framework.data.generator.action.datasets.cosmos3_action_lerobot import ( + ActionNormalization, + ActionSpec, + BaseActionLeRobotDataset, + Gripper, + Pos, + Rot, + build_action_spec, +) +from cosmos_framework.data.generator.action.utils.pose_utils import PoseConvention, convert_rotation +from cosmos_framework.data.generator.action.utils.viewpoint_utils import Viewpoint +from cosmos_framework.utils import log + +# LeRobot column names (constant across all RoboCasa atomic tasks; see meta/info.json). +_ACTION_FEATURE = "action" +_STATE_FEATURE = "observation.state" +_IMAGE_FEATURES = { + "wrist": "observation.images.robot0_eye_in_hand", + "left": "observation.images.robot0_agentview_left", + "right": "observation.images.robot0_agentview_right", +} + +# Native 12D action layout (meta/modality.json): the arm-only contract keeps only +# the arm end-effector delta + gripper. +_EEF_POS = slice(5, 8) # end_effector_position delta +_EEF_ROT = slice(8, 11) # end_effector_rotation delta (axis-angle) +_GRIPPER = slice(11, 12) # gripper_close + +# observation.state 16D layout (meta/modality.json): base_position[0:3], +# base_rotation[3:7] (quat), end_effector_position_relative[7:10] (base frame, meters), +# end_effector_rotation_relative[10:14] (quat, xyzw — verified: base_rot=[0,0,.707,.707] +# is a 90 deg yaw about z, w last), gripper_qpos[14:16] (two finger joint positions). +# For the EEF proprioception token we keep only the arm eef pose + a 1D gripper opening. +_STATE_EEF_POS = slice(7, 10) # absolute eef position (base frame) +_STATE_EEF_ROT = slice(10, 14) # absolute eef rotation (quaternion, xyzw) +_STATE_GRIPPER = slice(14, 16) # gripper qpos (two fingers) + + +# All 18 ``target/atomic`` tasks. Base-motion prevalence measured over 40 episodes each: +# NavigateKitchen is the only genuinely mobile task (100% of episodes, 2.6 m mean net +# displacement); PickPlaceDrawerToCounter moves in 48% (<=0.71 m); five more move in +# 2-10% of episodes (<=0.4 m); the remaining ten never move the base (<6 mm noise). +# Training on this set requires ``use_base_action=True`` so the base channels exist. +DEFAULT_ALL_ATOMIC_TASKS: tuple[str, ...] = ( + "CloseBlenderLid", + "CloseFridge", + "CloseToasterOvenDoor", + "CoffeeSetupMug", + "NavigateKitchen", + "OpenCabinet", + "OpenDrawer", + "OpenStandMixerHead", + "PickPlaceCounterToCabinet", + "PickPlaceCounterToStove", + "PickPlaceDrawerToCounter", + "PickPlaceSinkToCounter", + "PickPlaceToasterToCounter", + "SlideDishwasherRack", + "TurnOffStove", + "TurnOnElectricKettle", + "TurnOnMicrowave", + "TurnOnSinkFaucet", +) + + +# The 65 ``pretrain/atomic`` tasks (every atomic task RoboCasa365 ships a pretrain split for; +# the 18 ``target/atomic`` names above are a subset). Human demos only -- the registry also +# lists machine-generated ``mg`` paths for 60 of these, but they are not part of the local +# download. Same 12D action / 16D state / 3-camera contract as target/atomic, so this loader +# handles them unchanged; only ``root`` and ``task_names`` differ. +DEFAULT_PRETRAIN_ATOMIC_TASKS: tuple[str, ...] = ( + "AdjustToasterOvenTemperature", "AdjustWaterTemperature", "CheesyBread", "CloseBlenderLid", + "CloseCabinet", "CloseDishwasher", "CloseDrawer", "CloseElectricKettleLid", "CloseFridge", + "CloseFridgeDrawer", "CloseMicrowave", "CloseOven", "CloseStandMixerHead", + "CloseToasterOvenDoor", "CoffeeServeMug", "CoffeeSetupMug", "LowerHeat", "MakeIcedCoffee", + "NavigateKitchen", "OpenBlenderLid", "OpenCabinet", "OpenDishwasher", "OpenDrawer", + "OpenElectricKettleLid", "OpenFridge", "OpenFridgeDrawer", "OpenMicrowave", "OpenOven", + "OpenStandMixerHead", "OpenToasterOvenDoor", "PackDessert", "PickPlaceCabinetToCounter", + "PickPlaceCounterToBlender", "PickPlaceCounterToCabinet", "PickPlaceCounterToDrawer", + "PickPlaceCounterToMicrowave", "PickPlaceCounterToOven", "PickPlaceCounterToSink", + "PickPlaceCounterToStandMixer", "PickPlaceCounterToStove", "PickPlaceCounterToToasterOven", + "PickPlaceDrawerToCounter", "PickPlaceFridgeDrawerToShelf", "PickPlaceFridgeShelfToDrawer", + "PickPlaceMicrowaveToCounter", "PickPlaceSinkToCounter", "PickPlaceStoveToCounter", + "PickPlaceToasterOvenToCounter", "PickPlaceToasterToCounter", "PreheatOven", + "SlideDishwasherRack", "SlideOvenRack", "SlideToasterOvenRack", "StartCoffeeMachine", + "TurnOffMicrowave", "TurnOffSinkFaucet", "TurnOffStove", "TurnOnBlender", + "TurnOnElectricKettle", "TurnOnMicrowave", "TurnOnSinkFaucet", "TurnOnStove", + "TurnOnToaster", "TurnOnToasterOven", "TurnSinkSpout", +) + +# The 235 ``pretrain/composite`` tasks: multi-stage, long-horizon kitchen activities. +# Episodes average ~1565 frames (78 s) against ~142 for atomic, so by frame count composite +# outweighs pretrain/atomic roughly 18:1 (27.6M vs 1.5M). Mixing the two unweighted therefore +# yields a ~95% composite corpus -- intentional here, but worth remembering when reading +# per-task numbers. Composite demos drive the mobile base far more than atomic ones, which is +# why they require ``use_base_action=True``. +DEFAULT_PRETRAIN_COMPOSITE_TASKS: tuple[str, ...] = ( + "AddIceCubes", "AddLemonToFish", "AddMarshmallow", "AddSugarCubes", "AddSweetener", + "AdjustHeat", "AfterwashSorting", "AirDryFruit", "AlcoholServingPrep", "AlignSilverware", + "ArrangeBreadBowl", "ArrangeBuffetDessert", "ArrangeDrinkware", "ArrangeTeaAccompaniments", + "ArrangeUtensilsByType", "ArrangeVegetables", "AssembleCookingArray", "BalancedMealPrep", + "BeverageOrganization", "BeverageSorting", "BlendIngredients", "BlendMarinade", + "BowlAndCup", "BreadAndCheese", "BreadSetupSlicing", "BuildAppetizerPlate", "ButterOnPan", + "CandleCleanup", "CerealAndBowl", "ChooseMeasuringCup", "ChooseRipeFruit", "CleanBoard", + "CleanMicrowave", "ClearClutter", "ClearCuttingBoard", "ClearFoodWaste", "ClearFreezer", + "ClearReceptaclesForCleaning", "ClearSink", "ClearSinkArea", "ClearSinkSpace", + "CollectWashingSupplies", "ColorfulSalsa", "CondimentCollection", "CookieDoughPrep", + "CoolBakedCake", "CoolKettle", "CreateChildFriendlyFridge", "CupcakeCleanup", + "CutBuffetPizza", "DateNight", "DefrostByCategory", "DeliverBrewedCoffee", "DeliverStraw", + "DessertAssembly", "DessertUpgrade", "DisplayMeatVariety", "DistributeChicken", + "DivideBasins", "DivideBuffetTrays", "DrainVeggies", "DrinkwareConsolidation", "DryDishes", + "DryDrinkware", "DumpLeftovers", "FillBlenderJug", "FillKettle", "FilterMicrowavableItem", + "FoodCleanup", "FreezeBottledWaters", "FreezeCookedFood", "FreezeIceTray", + "FreshProduceOrganization", "FryingPanAdjustment", "GarnishCupcake", "GatherCuttingTools", + "GatherMarinadeIngredients", "GatherVegetables", "GetToastedBread", "HeatMug", + "HeatMultipleWater", "HotDogSetup", "JuiceFruitReamer", "KettleBoiling", + "LemonSeasoningFish", "LineUpCondiments", "LoadCondimentsInFridge", "LoadDishwasher", + "LoadFridgeByType", "LoadFridgeFifo", "LoadPreparedFood", "MakeFruitBowl", + "MakeLoadedPotato", "MatchCupAndDrink", "MaximizeFreezerSpace", "MealPrepStaging", + "MeatSkewerAssembly", "MeatTransfer", "MicrowaveCorrectMeal", "MicrowaveDefrostMeat", + "MicrowaveThawing", "MicrowaveThawingFridge", "MixCakeFrosting", "MixedFruitPlatter", + "MoveFreezerToFridge", "MoveFridgeToFreezer", "MoveToCounter", "MoveToFreezerDrawer", + "MultistepSteaming", "OrganizeBakingIngredients", "OrganizeCleaningSupplies", + "OrganizeCoffeeCondiments", "OrganizeCondiments", "OrganizeMetallicUtensils", + "OrganizeMugsByHandle", "OrganizeVegetables", "OvenBroilFish", "PackFoodByTemp", + "PackFruitContainer", "PackIdenticalLunches", "PastryDisplay", "PlaceBeveragesTogether", + "PlaceDishesBySink", "PlaceEqualIceCubes", "PlaceFoodInBowls", "PlaceIceInCup", + "PlaceMeatInMarinade", "PlaceMicrowaveSafeItem", "PlaceStraw", "PlaceVegetablesEvenly", + "PlaceVeggiesInDrawer", "PlateSteakMeal", "PlateStoreDinner", "PortionFruitBowl", + "PortionInTupperware", "PortionOnSize", "PortionYogurt", "PreRinseStation", "PreSoakPan", + "PreheatPot", "PrepForSanitizing", "PrepForTenderizing", "PrepFridgeForCleaning", + "PrepMarinatingMeat", "PrepSinkForCleaning", "PrepareBroilingStation", + "PrepareCheeseStation", "PrepareCocktailStation", "PrepareCoffee", "PrepareDishwasher", + "PrepareDrinkStation", "PrepareSausageCheese", "PrepareSmoothie", "PrepareSoupServing", + "PrepareStoringLeftovers", "PrepareToast", "PressChicken", "PrewashFoodAssembly", + "PrewashFoodSorting", "QuickThaw", "RearrangeFridgeItems", "RecycleBottlesBySize", + "RecycleSodaCans", "RecycleStackedYogurt", "RefillCondimentStation", "ReheatMeal", + "RemoveBroiledFish", "RemoveCuttingBoardItems", "ReorganizeFrozenVegetables", + "ResetCabinetDoors", "RestockBowls", "RestockCannedFood", "RestockPantry", + "RestockSinkSupplies", "RetrieveIceTray", "RetrieveMeat", "ReturnHeatedFood", + "ReturnWashingSupplies", "RinseBowls", "RinseCuttingBoard", "RinseSinkBasin", "RotatePan", + "SanitizePrepCuttingBoard", "ScalePortioning", "ScrubBowl", "ScrubCuttingBoard", + "SearingMeat", "SeasoningSpiceSetup", "SeasoningSteak", "ServeSteak", "ServeTea", + "ServeWarmCroissant", "SetBowlsForSoup", "SetUpCuttingStation", "SetUpSpiceStation", + "SetupBowls", "SetupButterPlate", "SetupFruitBowl", "SetupFrying", "SetupSodaBowl", + "SetupWineGlasses", "ShakePan", "SimmeringSauce", "SizeSorting", "SnackSorting", + "SoakSponge", "SortingCleanup", "SpicyMarinade", "StackBowlsCabinet", "StackBowlsInSink", + "StackCans", "StartElectricKettle", "SteamInMicrowave", "StirVegetables", + "StockingBreakfastFoods", "StoreDumplings", "StoreLeftoversByType", "StoreLeftoversInBowl", + "StrainerSetup", "SweetSavoryToastSetup", "SweetenCoffee", "SweetenHotChocolate", + "ThawInSink", "TiltPan", "ToastBagel", "ToastBaguette", "ToastOnCorrectRack", + "TongBuffetSetup", "TransportCookware", "TurnOffSimmeredSauceHeat", "VeggieDipPrep", + "WarmCroissant", "WashFish", "WashLettuce", "YogurtDelightPrep", +) + +# Convenience union for the 300-task pretrain soup (atomic + composite). Both live under one +# converted root, so the loader still takes a single ``root``. +DEFAULT_PRETRAIN_ALL_TASKS: tuple[str, ...] = ( + DEFAULT_PRETRAIN_ATOMIC_TASKS + DEFAULT_PRETRAIN_COMPOSITE_TASKS +) + +# Mobile-base action encoding (``use_base_action=True``). The Omron base is a +# JOINT_VELOCITY part with joints [forward, side, yaw, torso_height]; torso_height is +# never actuated in target/atomic (std 0 across all 18 tasks), so the base is a planar +# 3-DoF (x, y, yaw) system. Rather than regress the raw velocity command we derive a +# frame-wise EGO-FRAME base pose delta from ``observation.state`` (base_position + +# base_rotation) and encode it like the EEF delta (pos + rot6d). Because the deltas are +# expressed in the base frame at each step, no world->body yaw rotation is needed at +# inference; the closed-loop client converts back with +# ``base_motion[0:3] = (ego_delta / dt) / BASE_MAX_VELOCITY``. +# +# ``base_motion`` is a NORMALISED command in [-1, 1] (the JOINT_VELOCITY controller maps +# it onto the joint velocity limits), not a physical velocity. Calibrated on +# NavigateKitchen (least-squares through the origin of physical ego velocity against the +# recorded command, restricted to |command| > 0.2): full deflection corresponds to +# ~0.60 m/s forward, ~0.64 m/s lateral, ~1.25 rad/s yaw. Fit quality: yaw corr 0.997 +# (6.9% residual); the translation channels are corr 0.92-0.96 with ~26% residual, which +# is controller lag — the commanded velocity is not reached instantaneously. Closed-loop +# feedback absorbs that per step. +BASE_MAX_VELOCITY: tuple[float, float, float] = (0.601, 0.640, 1.251) +_STATE_BASE_POS = slice(0, 3) # absolute base position (world frame) +_STATE_BASE_ROT = slice(3, 7) # absolute base rotation (quaternion, xyzw) +_CONTROL_MODE = slice(4, 5) # +1 = base mode (arm tracks the moving base), -1 = arm mode + +# ``base_encoding="raw"``: regress the native command instead of inverting the controller. +# The ego encoding above derives the base delta from the ACHIEVED pose in +# ``observation.state``, so recovering the command at inference requires inverting a +# controller with real inertia -- structurally lossy (yaw residual 6.9 %, translation +# 26 %), and a recorded demo can never be replayed exactly. Passing ``base_motion`` +# through is an identity round-trip: the client writes ``env[7:11] = action[0:4]``. +# +# It also fixes a loss-weighting defect. Measured RMS per step: the ego translation delta +# is 0.018 (NavigateKitchen) to 0.002 (PickPlaceDrawerToCounter) in METRES, while the arm's +# ``eef_pos`` is a NORMALISED command with RMS 0.21-0.52 -- a 12x to 260x scale gap that, +# under ``action_normalization=None``, makes the base translation channels contribute +# ~1/144 to ~1/68000 of the arm's gradient. ``base_motion`` is normalised to [-1, 1] +# (RMS 0.49-0.61 on NavigateKitchen), i.e. the same scale as the arm block. +_BASE_MOTION = slice(0, 4) # [forward, side, yaw, torso_height]; torso is never actuated + + +class RoboCasaLeRobotDataset(BaseActionLeRobotDataset): + """RoboCasa manipulation dataset. + + Actions are ``[pos_delta(3), rot6d_delta(6), gripper(1)]`` (10D, arm only) or, with + ``use_base_action=True``, widened to include the mobile base — 15D for + ``base_encoding="raw"`` and 20D for ``"ego"``. Observation is a ``concat_view`` + composite selected by ``camera_set``. Reads v2.1 exports via the official lerobot + backend; each of ``task_names`` is discovered under ``root`` as + ``//*/lerobot`` and registered as a separate shard. + """ + + def __init__( + self, + root: str, + fps: float = 20.0, + chunk_length: int = 16, + split_seed: int = 42, + split_val_ratio: float = 0.01, + split: str = "train", + mode: str = "wam", + pose_convention: PoseConvention = "backward_framewise", + rotation_format: str = "rot6d", + action_normalization: ActionNormalization | None = None, + tolerance_s: float = 1e-4, + viewpoint: Viewpoint = "concat_view", + task_names: tuple[str, ...] | list[str] = DEFAULT_ALL_ATOMIC_TASKS, + use_state: bool = False, + enable_fast_init: bool = False, + camera_set: str = "wrist_lr", + use_base_action: bool = False, + base_encoding: str = "ego", + ) -> None: + if rotation_format != "rot6d": + raise NotImplementedError(f"RoboCasa loader only supports rotation_format='rot6d', got {rotation_format!r}.") + super().__init__( + fps=fps, + chunk_length=chunk_length, + split_seed=split_seed, + split_val_ratio=split_val_ratio, + split=split, + mode=mode, + embodiment_type="robocasa", + viewpoint=viewpoint, + pose_convention=pose_convention, + rotation_format=rotation_format, + action_normalization=action_normalization, + tolerance_s=tolerance_s, + enable_fast_init=enable_fast_init, + ) + self._use_state = use_state + # use_base_action: widen the arm-only 10D contract so mobile tasks + # (NavigateKitchen et al.) are representable. ``base_encoding`` selects HOW: + # + # "ego" (default, 20D) -- state-derived ego-frame base pose delta: + # [base_pos(3), base_rot6d(6), control_mode(1), eef_pos(3), eef_rot6d(6), gripper(1)] + # "raw" (15D) -- the native base_motion velocity command, passed through: + # [base_motion(4), control_mode(1), eef_pos(3), eef_rot6d(6), gripper(1)] + # + # "ego" is kept as the default so existing checkpoints/configs reproduce bit-for-bit. + # See the ``BASE_MAX_VELOCITY`` block above for why "raw" is the better contract. + self._use_base_action = use_base_action + if base_encoding not in ("ego", "raw"): + raise ValueError(f"Unsupported base_encoding={base_encoding!r}. Use 'ego' or 'raw'.") + self._base_encoding = base_encoding + self._image_features = _IMAGE_FEATURES + # camera_set: + # "wrist_lr" (default): concat_view = wrist (top) + L/R (bottom, squished to half). + # "left_wrist" : LIBERO-style = agentview_left + wrist, full-res, horizontally + # concatenated ([T,C,H,2W]); right cam DROPPED, no downscaling. + if camera_set not in ("wrist_lr", "left_wrist"): + raise ValueError(f"Unsupported camera_set={camera_set!r}. Use 'wrist_lr' or 'left_wrist'.") + self._camera_set = camera_set + + # Discover one LeRobot shard root per requested task under ``root``. + self._all_shard_roots = self._discover_shard_roots(root, task_names) + if not self._all_shard_roots: + raise FileNotFoundError( + f"No RoboCasa lerobot shards found under {root!r} for tasks {list(task_names)}. " + f"Expected //*/lerobot directories." + ) + log.info(f"RoboCasaLeRobotDataset: {len(self._all_shard_roots)} task shard(s): {self._all_shard_roots}") + + # delta_timestamps: chunk_length per-frame action deltas; chunk_length+1 + # observation frames (one more image/state than transitions). + observation_ts = [i * self._dt for i in range(0, self._chunk_length + 1)] + action_ts = [i * self._dt for i in range(0, self._chunk_length)] + self._delta_timestamps: dict[str, list[float]] = {_ACTION_FEATURE: action_ts} + if self._use_state: + self._delta_timestamps[_STATE_FEATURE] = observation_ts + if self._camera_set == "left_wrist": + # LIBERO-style: only agentview_left + wrist (right dropped). + self._delta_timestamps[self._image_features["wrist"]] = observation_ts + self._delta_timestamps[self._image_features["left"]] = observation_ts + else: + if self._viewpoint in ("wrist_view", "concat_view"): + self._delta_timestamps[self._image_features["wrist"]] = observation_ts + if self._viewpoint in ("third_person_view", "concat_view"): + self._delta_timestamps[self._image_features["left"]] = observation_ts + self._delta_timestamps[self._image_features["right"]] = observation_ts + + self._register_sources() + + @staticmethod + def _discover_shard_roots(root: str, task_names: tuple[str, ...] | list[str]) -> list[str]: + """Resolve ``//*/lerobot`` for each requested task. + + ``root`` may point at the atomic dir (``.../target/atomic``) or directly + at a single ``lerobot`` dir. Missing tasks are skipped with a warning. + + A task may hold more than one dated export (RoboCasa ships re-collected + demos under a new date). All of them are registered: silently keeping one + would train on less data than the caller staged, with nothing to show for it. + """ + root = root.rstrip("/") + if os.path.basename(root) == "lerobot" and os.path.isdir(os.path.join(root, "meta")): + return [root] + shard_roots: list[str] = [] + for task in task_names: + matches = sorted(glob.glob(os.path.join(root, task, "*", "lerobot"))) + matches = [m for m in matches if os.path.isdir(os.path.join(m, "meta"))] + if not matches: + log.warning(f"RoboCasaLeRobotDataset: no lerobot shard for task {task!r} under {root!r}; skipping.") + continue + if len(matches) > 1: + log.info( + f"RoboCasaLeRobotDataset: task {task!r} has {len(matches)} dated exports under " + f"{root!r}; registering all of them: {matches}" + ) + shard_roots.extend(matches) + return shard_roots + + # ---- action / spec ----------------------------------------------------- + + @property + def action_dim(self) -> int: + if self._use_base_action: + if self._base_encoding == "raw": + # [base_motion(4), control_mode(1), eef_pos(3), eef_rot6d(6), gripper(1)] + return 15 + # [base_pos(3), base_rot6d(6), control_mode(1), eef_pos(3), eef_rot6d(6), gripper(1)] + return 20 + return 10 # [pos(3), rot6d(6), gripper(1)] + + def _build_action_spec(self) -> ActionSpec: + if self._use_base_action: + if self._base_encoding == "raw": + # base_motion is a NORMALISED velocity command in [-1, 1], the same units + # as the arm's eef_pos block -- so it is declared as ``Pos`` (magnitude-based + # idle detection, ``|v| < eps_t``) and joins the arm translation in one L2. + # NOT ``Joint``: that branch is frame-DIFF based, so a constant cruise + # command (diff == 0) would be misread as idle. The 4th channel + # (torso height) is never actuated and contributes 0 to the norm. + return build_action_spec(Pos(dim=4), Gripper(), Pos(), Rot("rot6d"), Gripper()) + # Base pose delta first (Pos+Rot), then the mode flag, then the arm block. + # ``control_mode`` is a +/-1 channel regressed like the gripper (continuous + # rectified flow, no discrete head); the client thresholds it at 0. + return build_action_spec(Pos(), Rot("rot6d"), Gripper(), Pos(), Rot("rot6d"), Gripper()) + return build_action_spec(Pos(), Rot("rot6d"), Gripper()) + + def _build_frame_wise_action(self, raw_action: torch.Tensor) -> torch.Tensor: + """RoboCasa 12D per-frame action -> 10D ``[pos(3), rot6d(6), gripper(1)]``. + + Drops base_motion[0:4] + control_mode[4]; re-encodes the eef axis-angle + delta to rot6d. Mirrors ``LIBEROLeRobotDataset._build_frame_wise_action`` + with RoboCasa slice indices. + """ + raw = raw_action.float() # [chunk, 12] + translation = raw[:, _EEF_POS] # [chunk, 3] + rotation_matrix = convert_rotation(raw[:, _EEF_ROT], input_format="axisangle", output_format="matrix") + rotation = convert_rotation(rotation_matrix, input_format="matrix", output_format="rot6d") # [chunk, 6] + gripper = raw[:, _GRIPPER] # [chunk, 1] + return torch.cat([translation, rotation, gripper], dim=-1) # [chunk, 10] + + def _build_initial_state(self, state_seq: torch.Tensor) -> torch.Tensor: + """Current EEF proprioception -> 10D ``[pos(3), rot6d(6), gripper(1)]`` token. + + Mirrors DROID's ``use_state`` path (``midtrain`` branch): the frame BEFORE the + action chunk (index 0 of the ``chunk_length+1`` observation window) is prepended + as a CLEAN conditioning action token — same width as the action so it flows + through the shared ``action2llm`` embedding, gets ``sigma=0`` (never noised) and + is excluded from the flow-matching loss (``condition_frame_indexes_action=[0]``). + + RoboCasa ``observation.state`` eef fields are in the (fixed) base frame; the + rotation is a ``quat_xyzw`` re-encoded to rot6d to match the 10D action layout. + Gripper = signed finger opening (``qpos[0]-qpos[1]``), a 1D summary of the + two-finger state. NOTE: this token is ABSOLUTE pose while the predicted actions + are per-frame deltas — identical to DROID ``midtrain`` — so pair with + ``action_normalization=None`` (the quantile_rot delta-stats do not apply to an + absolute pose). + """ + s0 = state_seq[-self._chunk_length - 1].float() # [16]; earliest = current pre-chunk state + pos = s0[_STATE_EEF_POS] # [3] + quat = s0[_STATE_EEF_ROT].unsqueeze(0) # [1,4] xyzw + matrix = convert_rotation(quat, input_format="quat_xyzw", output_format="matrix") + rot6d = convert_rotation(matrix, input_format="matrix", output_format="rot6d").reshape(6) # [6] + grip = s0[14:15] - s0[15:16] # [1] signed finger opening + return torch.cat([pos, rot6d, grip], dim=-1) # [10] + + def _build_base_delta(self, state_seq: torch.Tensor) -> torch.Tensor: + """Frame-wise EGO-FRAME base pose delta -> ``[chunk, 9]`` ``[pos(3), rot6d(6)]``. + + ``observation.state`` carries the base pose in the WORLD frame + (``base_position``, ``base_rotation`` quat_xyzw). The delta between consecutive + frames is re-expressed in the base frame at the earlier step, i.e. + + T_delta = T_base[t]^-1 @ T_base[t+1] + + so the representation is ego-centric and matches the ``backward_framewise`` + convention already used for the EEF. This removes any dependence on world + heading: the closed-loop client can map straight back to the base velocity + command with ``(ego_delta / dt) * BASE_VELOCITY_SCALE`` without a yaw rotation. + + The state window is ``chunk_length+1`` frames (one more than the transitions), + so the ``chunk_length`` deltas line up one-to-one with the action chunk. + """ + s = state_seq[-self._chunk_length - 1 :].float() # [chunk+1, 16] + pos = s[:, _STATE_BASE_POS] # [chunk+1, 3] world + quat = s[:, _STATE_BASE_ROT] # [chunk+1, 4] xyzw, world + rot = convert_rotation(quat, input_format="quat_xyzw", output_format="matrix") # [chunk+1,3,3] + + r_prev = rot[:-1] # [chunk,3,3] + r_next = rot[1:] # [chunk,3,3] + # Translation delta rotated into the previous base frame. + d_world = (pos[1:] - pos[:-1]).unsqueeze(-1) # [chunk,3,1] + d_ego = torch.matmul(r_prev.transpose(-1, -2), d_world).squeeze(-1) # [chunk,3] + # Relative rotation expressed in the previous base frame. + r_rel = torch.matmul(r_prev.transpose(-1, -2), r_next) # [chunk,3,3] + rot6d = convert_rotation(r_rel, input_format="matrix", output_format="rot6d") # [chunk,6] + return torch.cat([d_ego, rot6d], dim=-1) # [chunk, 9] + + # ---- video ------------------------------------------------------------- + + def _compose_multi_view(self, sample: dict[str, Any]) -> torch.Tensor: + """Tile wrist (top) + left/right third-person (bottom) into one frame. + + Layout per frame (all source views are square, same size): + ┌──────────────┐ + │ wrist │ (H, W) + ├───────┬──────┤ + │ left │ right │ (H/2, W/2) each + └───────┴──────┘ + Output height is 3H/2 (mirrors ``DROIDLeRobotDataset._compose_multi_view``). + """ + wrist = sample[self._image_features["wrist"]] # [T,C,H,W] + left = sample[self._image_features["left"]] # [T,C,H,W] + right = sample[self._image_features["right"]] # [T,C,H,W] + + _, _, h_w, w_w = wrist.shape + half_h, half_w = h_w // 2, w_w // 2 + left = F.interpolate(left, size=(half_h, half_w), mode="bilinear", align_corners=False) + right = F.interpolate(right, size=(half_h, half_w), mode="bilinear", align_corners=False) + bottom = torch.cat([left, right], dim=-1) # [T,C,H/2,W] + return torch.cat([wrist, bottom], dim=-2) # [T,C,3H/2,W] + + def _compose_left_wrist(self, sample: dict[str, Any]) -> torch.Tensor: + """LIBERO-style: agentview_left (left) + wrist (right), full-res, horizontal. + + Both source views are native 256x256; concatenated along width -> [T,C,H,2W] + (256x512). No downscaling, right camera dropped. Mirrors + ``LIBEROLeRobotDataset`` concat_view ordering (third-person | wrist). + """ + left = sample[self._image_features["left"]] # [T,C,H,W] + wrist = sample[self._image_features["wrist"]] # [T,C,H,W] + return torch.cat([left, wrist], dim=-1) # [T,C,H,2W] + + + # ---- sample build ------------------------------------------------------ + + def __getitem__(self, idx: int) -> dict[str, Any]: + mode, _, _, sample = self._fetch_sample(idx) + + action = self._build_frame_wise_action(sample[_ACTION_FEATURE]) # [chunk, 10] + + if self._use_base_action: + raw = sample[_ACTION_FEATURE].float() + control_mode = raw[:, _CONTROL_MODE] # [chunk, 1]; +/-1, regressed like the gripper + if self._base_encoding == "raw": + # Native normalised velocity command, passed through unchanged -> [chunk, 15]. + base_block = raw[:, _BASE_MOTION] # [chunk, 4] + else: + # State-derived ego-frame pose delta -> [chunk, 20]. + base_block = self._build_base_delta(sample[_STATE_FEATURE]) # [chunk, 9] + action = torch.cat([base_block, control_mode, action], dim=-1) + + # EEF proprioception: prepend the current eef pose as a clean conditioning + # token (DROID use_state parity). Compute idle_frames on the delta-only chunk + # BEFORE prepending so the absolute-pose token does not skew idle detection. + state_extras: dict[str, Any] = {} + if self._use_state: + idle = self._compute_idle_frames(action) + if idle is not None: + state_extras["idle_frames"] = idle + initial_state = self._build_initial_state(sample[_STATE_FEATURE]) # [10] + if self._use_base_action: + # Widen the conditioning token to the action contract (20D ego / 15D raw). + # The base block is zero-filled on purpose: under "ego" the absolute base + # pose is world-frame and the kitchen layout is re-randomised per episode, + # and under "raw" ``observation.state`` carries no base VELOCITY at all — + # only pose. Either way there is no transferable signal to put here, unlike + # the EEF pose, which is base-relative. + pad = self.action_dim - 10 + initial_state = torch.cat( + [torch.zeros(pad, dtype=initial_state.dtype), initial_state], dim=-1 + ) # [action_dim] + action = torch.cat([initial_state.unsqueeze(0), action], dim=0) # [chunk+1, action_dim] + + if self._skip_video_loading: + video = None + elif self._camera_set == "left_wrist": + video = self._compose_left_wrist(sample) + elif self._viewpoint == "concat_view": + video = self._compose_multi_view(sample) + elif self._viewpoint == "wrist_view": + video = sample[self._image_features["wrist"]] + else: # third_person_view + video = sample[self._image_features["left"]] + + ai_caption = sample["task"] + + extras: dict[str, Any] = {} + if self._camera_set == "left_wrist": + extras["additional_view_description"] = ( + "The left half is a third-person view of the scene. " + "The right half is from the wrist-mounted camera." + ) + elif self._viewpoint == "concat_view": + extras["additional_view_description"] = ( + "The top row is from the wrist-mounted camera. " + "The bottom row contains two horizontally concatenated third-person views of the scene." + ) + return self._build_result( + mode=mode, video=video, action=action, ai_caption=ai_caption, **state_extras, **extras + ) diff --git a/examples/worldmodel/cosmos3/vendor/cosmos_framework/data/generator/action/datasets/robomind_franka_dataset.py b/examples/worldmodel/cosmos3/vendor/cosmos_framework/data/generator/action/datasets/robomind_franka_dataset.py new file mode 100644 index 000000000000..10b3c04200f2 --- /dev/null +++ b/examples/worldmodel/cosmos3/vendor/cosmos_framework/data/generator/action/datasets/robomind_franka_dataset.py @@ -0,0 +1,278 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: OpenMDW-1.1 + +"""RoboMIND Franka LeRobot dataset — single-arm and dual-arm variants.""" + +from __future__ import annotations + +import random +from pathlib import Path +from typing import Any, Literal + +import numpy as np +import torch +import torch.nn.functional as F + +from cosmos_framework.data.generator.action.action_normalization import load_action_stats +from cosmos_framework.data.generator.action.utils.action_spec import ActionSpec, Gripper, Pos, Rot, build_action_spec +from cosmos_framework.data.generator.action.datasets.base_dataset import ActionBaseDataset +from cosmos_framework.data.generator.action.utils.pose_utils import ( + build_abs_pose_from_components, + pose_abs_to_rel, +) + +PoseConvention = Literal["backward_framewise"] +Viewpoint = Literal["concat_view", "third_person_view"] + +# Dual-arm uses camera_front as the top view; single-arm uses camera_top. +_IMAGE_FEATURES_DUAL = { + "top": "observation.images.camera_front", + "left": "observation.images.camera_left", + "right": "observation.images.camera_right", +} +_IMAGE_FEATURES_SINGLE = { + "top": "observation.images.camera_top", + "left": "observation.images.camera_left", + "right": "observation.images.camera_right", +} +_STATE_FEATURE = "observation.states.end_effector" +_ACTION_FEATURE = "actions.joint_position" + +# 90-degree clockwise rotation about the Z axis in the local frame. This matches +# the production RoboMIND Franka wrapper conversion to OpenCV coordinates. +_ROBOMIND_FRANKA_TO_OPENCV: np.ndarray = np.array( + [[0.0, -1.0, 0.0], [1.0, 0.0, 0.0], [0.0, 0.0, 1.0]], + dtype=np.float32, +) + +_NORMALIZER_PATH_DUAL = Path(__file__).parent.parent / "normalizer_stats/robomind_franka_dual_stats.json" +_NORMALIZER_PATH_SINGLE = Path(__file__).parent.parent / "normalizer_stats/robomind_franka_stats.json" + + +def _dual_arm_action_spec(): + return build_action_spec( + Pos(prefix="left"), + Rot("rot6d", prefix="left"), + Gripper(prefix="left"), + Pos(prefix="right"), + Rot("rot6d", prefix="right"), + Gripper(prefix="right"), + ) + + +class RoboMINDFrankaDataset(ActionBaseDataset): + """RoboMIND Franka dataset — single-arm (10D) or dual-arm (20D) cartesian actions. + + Single-arm (``robomind-franka``): ``[pos_delta(3), rot6d_delta(6), gripper(1)]`` + Dual-arm (``robomind-franka-dual``): + ``[left_pos(3), left_rot6d(6), left_gripper(1), + right_pos(3), right_rot6d(6), right_gripper(1)]`` + """ + + _SUPPORTED_EMBODIMENTS = ("robomind-franka", "robomind-franka-dual") + + def __init__( + self, + root: str, + fps: float = 10.0, + chunk_length: int = 16, + mode: str = "joint", + embodiment_type: str = "robomind-franka-dual", + pose_convention: PoseConvention = "backward_framewise", + tolerance_s: float = 1e-4, + viewpoint: Viewpoint = "concat_view", + action_normalization: str | None = "quantile", + sample_stride: int = 1, + ) -> None: + if embodiment_type not in self._SUPPORTED_EMBODIMENTS: + raise ValueError( + f"RoboMINDFrankaDataset only supports {self._SUPPORTED_EMBODIMENTS}, " + f"got embodiment_type={embodiment_type!r}." + ) + if viewpoint not in ("concat_view", "third_person_view"): + raise NotImplementedError(f"RoboMINDFrankaDataset does not support viewpoint={viewpoint!r}.") + self._embodiment_type = embodiment_type + self._image_features = ( + _IMAGE_FEATURES_SINGLE if embodiment_type == "robomind-franka" else _IMAGE_FEATURES_DUAL + ) + super().__init__( + root=root, + domain_name=embodiment_type, + fps=fps, + chunk_length=chunk_length, + mode=mode, + pose_convention=pose_convention, + tolerance_s=tolerance_s, + viewpoint=viewpoint, + action_normalization=action_normalization, + sample_stride=sample_stride, + ) + + @property + def action_dim(self) -> int: + return 10 if self._embodiment_type == "robomind-franka" else 20 + + def _action_spec(self) -> ActionSpec: + if self._embodiment_type == "robomind-franka": + return build_action_spec(Pos(), Rot("rot6d"), Gripper()) + return _dual_arm_action_spec() + + @classmethod + def _stats_path(cls) -> Path: + # Class-level default (no instance to branch on) — matches the + # constructor's default embodiment_type ("robomind-franka-dual"). + return _NORMALIZER_PATH_DUAL + + def load_action_stats(self) -> dict[str, torch.Tensor]: + """Instance-aware override: respects ``self._embodiment_type``. + + The inherited classmethod always resolves via ``_stats_path()``, which + has no instance to branch on and is hardcoded to the dual-arm file. A + single-arm instance calling ``load_action_stats()`` would otherwise + silently get 20D dual-arm stats instead of its own 10D stats. + """ + path = _NORMALIZER_PATH_DUAL if self._embodiment_type == "robomind-franka-dual" else _NORMALIZER_PATH_SINGLE + return { + key: torch.from_numpy(value).float() + for key, value in load_action_stats(str(path)).items() + } + + def _load_norm_stats(self) -> dict[str, torch.Tensor]: + if self._norm_stats is None: + self._norm_stats = self.load_action_stats() + return self._norm_stats + + def __getitem__(self, idx: int) -> dict[str, Any]: + mode = self._choose_mode() + idx = int(idx) + row_idx = idx * self._sample_stride + observation_rows = self._rows[row_idx : row_idx + self._chunk_length + 1] + action_rows = observation_rows[: self._chunk_length] + + episode = self._episodes[int(observation_rows[0]["episode_index"])] + task = self._tasks[int(observation_rows[0]["task_index"])] + ai_caption = random.choice([part.strip() for part in task.split(" | ") if part.strip()] or [task]) + + if self._viewpoint == "concat_view": + video = self._load_concat_video(episode, observation_rows) + if self._embodiment_type == "robomind-franka": + view_desc = ( + "The top row shows a third-person perspective looking towards the single-arm Franka robot from above. " + "The bottom-left view looks at the scene from the left side, and the bottom-right view looks at the scene from the right side." + ) + else: + view_desc = ( + "The top row shows a third-person perspective looking towards the dual-arm Franka robot from the front. " + "The bottom-left view looks at the scene from the left side, and the bottom-right view looks at the scene from the right side." + ) + else: + video = self._load_single_video(episode, observation_rows) + view_desc = None + + if self._embodiment_type == "robomind-franka": + raw_action, initial_pose = self._build_raw_action_single(observation_rows, action_rows) + extras: dict[str, Any] = {"initial_pose": initial_pose} + else: + raw_action, initial_pose_left, initial_pose_right = self._build_raw_action_dual(observation_rows, action_rows) + extras = {"initial_pose": initial_pose_left, "initial_pose_right": initial_pose_right} + + if view_desc is not None: + extras["additional_view_description"] = view_desc + + return self._build_result( + mode=mode, + video=video, + action=raw_action, + ai_caption=ai_caption, + **extras, + ) + + def _load_single_video( + self, + episode: dict[str, Any], + observation_rows: list[dict[str, Any]], + ) -> torch.Tensor: + # lerobot is a heavy, optional ("train" extra) dependency; import lazily. + from lerobot.datasets.video_utils import decode_video_frames + + video_key = self._image_features["top"] + timestamps = [float(row["timestamp"]) for row in observation_rows] + return decode_video_frames( + self._video_path(episode, video_key), + [float(episode.get(f"videos/{video_key}/from_timestamp", 0.0)) + ts for ts in timestamps], + self._tolerance_s, + ) + + def _load_concat_video( + self, + episode: dict[str, Any], + observation_rows: list[dict[str, Any]], + ) -> torch.Tensor: + # lerobot is a heavy, optional ("train" extra) dependency; import lazily. + from lerobot.datasets.video_utils import decode_video_frames + + timestamps = [float(row["timestamp"]) for row in observation_rows] + frames_by_view = { + name: decode_video_frames( + self._video_path(episode, video_key), + [float(episode.get(f"videos/{video_key}/from_timestamp", 0.0)) + ts for ts in timestamps], + self._tolerance_s, + ) + for name, video_key in self._image_features.items() + } + + top = frames_by_view["top"] + left = frames_by_view["left"] + right = frames_by_view["right"] + _, _, h_top, w_top = top.shape + half_h, half_w = h_top // 2, w_top // 2 + left = F.interpolate(left, size=(half_h, half_w), mode="bilinear", align_corners=False) + right = F.interpolate(right, size=(half_h, half_w), mode="bilinear", align_corners=False) + bottom = torch.cat([left, right], dim=-1) + return torch.cat([top, bottom], dim=-2) + + def _build_relative_poses( + self, + positions: np.ndarray, + euler_xyz: np.ndarray, + ) -> tuple[np.ndarray, torch.Tensor]: + poses_abs = build_abs_pose_from_components(positions, euler_xyz, "euler_xyz") + poses_abs[:, :3, :3] = poses_abs[:, :3, :3] @ _ROBOMIND_FRANKA_TO_OPENCV + initial_pose = torch.from_numpy(poses_abs[0].copy()).float() + poses_rel = pose_abs_to_rel(poses_abs, rotation_format="rot6d", pose_convention=self._pose_convention) + return poses_rel, initial_pose + + def _build_raw_action_single( + self, + observation_rows: list[dict[str, Any]], + action_rows: list[dict[str, Any]], + ) -> tuple[torch.Tensor, torch.Tensor]: + state = np.asarray([row[_STATE_FEATURE] for row in observation_rows], dtype=np.float32) + gripper = np.asarray([row[_ACTION_FEATURE] for row in action_rows], dtype=np.float32) + poses_rel, initial_pose = self._build_relative_poses(state[:, 0:3], state[:, 3:6]) + action = np.concatenate( + [poses_rel[-self._chunk_length :], 1.0 - gripper[-self._chunk_length :, [7]]], + axis=-1, + ) # [T, 10] + return torch.from_numpy(action).float(), initial_pose + + def _build_raw_action_dual( + self, + observation_rows: list[dict[str, Any]], + action_rows: list[dict[str, Any]], + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + state = np.asarray([row[_STATE_FEATURE] for row in observation_rows], dtype=np.float32) + gripper = np.asarray([row[_ACTION_FEATURE] for row in action_rows], dtype=np.float32) + + poses_rel_left, initial_pose_left = self._build_relative_poses(state[:, 0:3], state[:, 3:6]) + poses_rel_right, initial_pose_right = self._build_relative_poses(state[:, 6:9], state[:, 9:12]) + action = np.concatenate( + [ + poses_rel_left[-self._chunk_length :], + 1.0 - gripper[-self._chunk_length :, [7]], + poses_rel_right[-self._chunk_length :], + 1.0 - gripper[-self._chunk_length :, [15]], + ], + axis=-1, + ) # [T, 20] + return torch.from_numpy(action).float(), initial_pose_left, initial_pose_right diff --git a/examples/worldmodel/cosmos3/vendor/cosmos_framework/data/generator/action/datasets/robomind_ur_dataset.py b/examples/worldmodel/cosmos3/vendor/cosmos_framework/data/generator/action/datasets/robomind_ur_dataset.py new file mode 100644 index 000000000000..475a483032f5 --- /dev/null +++ b/examples/worldmodel/cosmos3/vendor/cosmos_framework/data/generator/action/datasets/robomind_ur_dataset.py @@ -0,0 +1,205 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: OpenMDW-1.1 + +"""RoboMIND UR LeRobot dataset for single-arm UR5e embodiment.""" + +from __future__ import annotations + +import random +from pathlib import Path +from typing import Any, Literal + +import numpy as np +import torch + +from cosmos_framework.data.generator.action.utils.action_spec import ActionSpec, Gripper, Pos, Rot, build_action_spec +from cosmos_framework.data.generator.action.datasets.base_dataset import ActionBaseDataset +from cosmos_framework.data.generator.action.utils.pose_utils import pose_abs_to_rel + +PoseConvention = Literal["backward_framewise"] +Viewpoint = Literal["third_person_view"] + +_IMAGE_FEATURE = "observation.images.camera_top" +_ACTION_JOINT_FEATURE = "actions.joint_position" + +# UR EE frame → OpenCV convention rotation (3×3, post-multiplied). +# Identity: attachment_site (quat="-1 1 0 0" in ur5e_robotiq_2f85.xml) already +# satisfies OpenCV convention (z = approach). +_ROBOMIND_UR_TO_OPENCV: np.ndarray = np.eye(3, dtype=np.float32) + +_UR5E_ARM_JOINTS = 6 # shoulder_pan … wrist_3 +_UR5E_EE_SITE = "attachment_site" # flange site in ur5e_robotiq_2f85.xml + +_MJCF_PATH = Path(__file__).resolve().parent.parent / "robot_assets" / "ur5e_robotiq_2f85.xml" + +_NORMALIZER_PATH = Path(__file__).parent.parent / "normalizer_stats/robomind_ur_stats.json" + + +class RoboMINDURDataset(ActionBaseDataset): + """RoboMIND UR dataset with 10D cartesian actions: + + [pos_delta(3), rot6d_delta(6), gripper(1)] + + Uses FK on ``actions.joint_position`` (6 arm joints → MuJoCo + ``attachment_site`` SE(3) pose). ``observation.states.end_effector`` is + NOT used — it is recorded incorrectly (constant) in ~89 % of UR episodes; + ``actions.joint_position`` is valid for 100 % of episodes. + + The sample also includes ``joint_configs``: absolute joint angles ``(T, 7)`` + from ``actions.joint_position[1:T+1]`` for FK-based robot mesh animation. + """ + + def __init__( + self, + root: str, + fps: float = 10.0, + chunk_length: int = 16, + mode: str = "joint", + pose_convention: PoseConvention = "backward_framewise", + tolerance_s: float = 1e-4, + viewpoint: Viewpoint = "third_person_view", + action_normalization: str | None = None, + sample_stride: int = 1, + ) -> None: + if viewpoint != "third_person_view": + raise NotImplementedError("RoboMINDURDataset only supports third_person_view.") + super().__init__( + root=root, + domain_name="robomind-ur", + fps=fps, + chunk_length=chunk_length, + mode=mode, + pose_convention=pose_convention, + tolerance_s=tolerance_s, + viewpoint=viewpoint, + action_normalization=action_normalization, + sample_stride=sample_stride, + ) + self._mj_model, self._mj_data, self._ee_site_id = self._init_mujoco() + + @staticmethod + def _init_mujoco(): + """Load UR5e+Robotiq MuJoCo model (kinematics-only) and locate the EE site. + + Strips all geoms and mesh/texture/material assets from the MJCF via + ``MjSpec`` before compile so the model loads without mesh files on disk. + FK only needs the kinematic tree, so ``mj_forward`` + site poses still + produce identical EE poses. + """ + import mujoco + + spec = mujoco.MjSpec.from_file(str(_MJCF_PATH)) + + def _strip_geoms(body): + for g in list(body.geoms): + spec.delete(g) + for child in body.bodies: + _strip_geoms(child) + + _strip_geoms(spec.worldbody) + for m in list(spec.meshes): + spec.delete(m) + for t in list(spec.textures): + spec.delete(t) + for mat in list(spec.materials): + spec.delete(mat) + + mj_model = spec.compile() + mj_data = mujoco.MjData(mj_model) + ee_site_id = mujoco.mj_name2id(mj_model, mujoco.mjtObj.mjOBJ_SITE, _UR5E_EE_SITE) + if ee_site_id < 0: + raise RuntimeError(f"EE site '{_UR5E_EE_SITE}' not found in {_MJCF_PATH}") + return mj_model, mj_data, ee_site_id + + def _fk_ee_poses(self, arm_q: np.ndarray) -> tuple[np.ndarray, np.ndarray]: + """Run MuJoCo FK for T+1 arm configs → EE site positions and rotations. + + Args: + arm_q: ``(T+1, 6)`` arm joint angles in radians. + + Returns: + ``(positions (T+1, 3), rotations (T+1, 3, 3))`` in MuJoCo world frame. + """ + import mujoco + + T1 = len(arm_q) + positions = np.empty((T1, 3), dtype=np.float32) + rotations = np.empty((T1, 3, 3), dtype=np.float32) + for t in range(T1): + self._mj_data.qpos[:_UR5E_ARM_JOINTS] = arm_q[t] + mujoco.mj_forward(self._mj_model, self._mj_data) + positions[t] = self._mj_data.site_xpos[self._ee_site_id] + rotations[t] = self._mj_data.site_xmat[self._ee_site_id].reshape(3, 3) + return positions, rotations + + @property + def action_dim(self) -> int: + return 10 + + def _action_spec(self) -> ActionSpec: + return build_action_spec(Pos(), Rot("rot6d"), Gripper()) + + @classmethod + def _stats_path(cls) -> Path: + return _NORMALIZER_PATH + + def __getitem__(self, idx: int) -> dict[str, Any]: + mode = self._choose_mode() + idx = int(idx) + row_idx = idx * self._sample_stride + # T+1 rows: current frame + T future frames (needed for FK EE poses and joint_configs) + observation_rows = self._rows[row_idx : row_idx + self._chunk_length + 1] + + episode = self._episodes[int(observation_rows[0]["episode_index"])] + task = self._tasks[int(observation_rows[0]["task_index"])] + ai_caption = random.choice([part.strip() for part in task.split(" | ") if part.strip()] or [task]) + + video = self._load_video(episode, observation_rows) + action, initial_pose, joint_configs = self._build_raw_action(observation_rows) + + return self._build_result( + mode=mode, + video=video, + action=action, + ai_caption=ai_caption, + initial_pose=initial_pose, + joint_configs=joint_configs, + ) + + def _load_video(self, episode: dict[str, Any], observation_rows: list[dict[str, Any]]) -> torch.Tensor: + # lerobot is a heavy, optional ("train" extra) dependency; import lazily. + from lerobot.datasets.video_utils import decode_video_frames + + timestamps = [float(row["timestamp"]) for row in observation_rows] + return decode_video_frames( + self._video_path(episode, _IMAGE_FEATURE), + [float(episode.get(f"videos/{_IMAGE_FEATURE}/from_timestamp", 0.0)) + ts for ts in timestamps], + self._tolerance_s, + ) + + def _build_raw_action( + self, + observation_rows: list[dict[str, Any]], + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + # [T+1, 7]: 6 arm joints + 1 gripper command + q = np.asarray([row[_ACTION_JOINT_FEATURE] for row in observation_rows], dtype=np.float32) + T = len(q) - 1 + + # FK EE trajectory: T+1 absolute poses from arm joints via MuJoCo + fk_pos, fk_rot = self._fk_ee_poses(q[:, :_UR5E_ARM_JOINTS]) + poses_abs = np.tile(np.eye(4, dtype=np.float32), (T + 1, 1, 1)) + poses_abs[:, :3, 3] = fk_pos + poses_abs[:, :3, :3] = fk_rot @ _ROBOMIND_UR_TO_OPENCV + + initial_pose = torch.from_numpy(poses_abs[0].copy()).float() + poses_rel = pose_abs_to_rel(poses_abs, rotation_format="rot6d", pose_convention=self._pose_convention) + + # Raw UR gripper: 0=open, 1=closed → invert so action convention is 0=closed, 1=open. + # joint_configs keeps the raw value; FK mesh uses raw * 255 → Robotiq ctrl. + gripper = torch.from_numpy(1.0 - q[:T, 6:7]).float() + action = torch.cat([torch.from_numpy(poses_rel).float(), gripper], dim=-1) # [T, 10] + + # Mesh animation: frames 1..T of joint position (post-action states) + joint_configs = torch.from_numpy(q[1 : 1 + T].copy()) # [T, 7] + + return action, initial_pose, joint_configs diff --git a/examples/worldmodel/cosmos3/vendor/cosmos_framework/data/generator/action/datasets/umi_lerobot_dataset.py b/examples/worldmodel/cosmos3/vendor/cosmos_framework/data/generator/action/datasets/umi_lerobot_dataset.py new file mode 100644 index 000000000000..3b408d35e537 --- /dev/null +++ b/examples/worldmodel/cosmos3/vendor/cosmos_framework/data/generator/action/datasets/umi_lerobot_dataset.py @@ -0,0 +1,177 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: OpenMDW-1.1 + +"""UMI LeRobot dataset.""" + +from __future__ import annotations + +import random +from pathlib import Path +from typing import Any, Literal + +import numpy as np +import torch + +from cosmos_framework.data.generator.action.action_normalization import load_action_stats +from cosmos_framework.data.generator.action.utils.action_spec import ActionSpec, Gripper, Pos, Rot, build_action_spec +from cosmos_framework.data.generator.action.datasets.base_dataset import ActionBaseDataset +from cosmos_framework.data.generator.action.utils.pose_utils import ( + build_abs_pose_from_components, + pose_abs_to_rel, +) + +PoseConvention = Literal["backward_framewise"] +Viewpoint = Literal["wrist_view"] + +# Feature keys matching UMI LeRobot parquet columns. +# Trajectory: 7D [pos(3), quat_wxyz(4)] — the main-camera TCP pose. +_TRAJ_KEY = "observation.state.right_main_camera_trajectory_xyz_wxyz" +_GRIPPER_KEY = "observation.state.right_gripper_width_m" +_IMAGE_FEATURE = "observation.image.right_main_camera_rgb" + +# Default EEF-in-camera-frame offset (most UMI rigs). +# touch_in_the_wild / FastUMI use FORWARD_EEF_IN_CAMERA_FRAME_XYZ_WXYZ with z=0.056. +_DEFAULT_EEF_IN_CAMERA_FRAME_XYZ_WXYZ: tuple[float, ...] = (0.0, 0.086, 0.09, 1.0, 0.0, 0.0, 0.0) +FORWARD_EEF_IN_CAMERA_FRAME_XYZ_WXYZ: tuple[float, ...] = (0.0, 0.086, 0.056, 1.0, 0.0, 0.0, 0.0) +"""EEF offset for touch_in_the_wild / FastUMI rigs (camera mounted slightly forward).""" + +_NORMALIZER_PATH = Path(__file__).parent.parent / "normalizer_stats/umi_lerobot_stats.json" + +# Action layout: single-arm is the first 10D of the 20D bimanual stats file +# (right_eef_poses(9) + right_eef_commands(1)). +_SINGLE_ARM_ACTION_DIM = 10 + + +class UMILeRobotDataset(ActionBaseDataset): + """UMI dataset converted to LeRobot format with 10D cartesian actions: + + [pos_delta(3), rot6d_delta(6), gripper_width(1)] + + Expects a LeRobot v2 dataset with: + * ``observation.images.camera0``: wrist-mounted RGB video (configurable + via ``image_key``). + * ``observation.state.right_main_camera_trajectory_xyz_wxyz``: 7D camera + TCP pose ``[pos(3), quat_wxyz(4)]`` for frames [0 .. chunk_length]. + * ``observation.state.right_gripper_width_m``: scalar gripper width for + frames [1 .. chunk_length] (commanded future widths). + + Poses are transformed from the camera TCP frame to the EEF frame via + ``eef_in_camera_frame_xyz_wxyz``, then converted to backward-framewise + rot6d relative poses. The stats file stores 20D bimanual stats (right + + left arm); single-arm normalization uses only the first 10D (right arm). + """ + + def __init__( + self, + root: str, + fps: float = 10.0, + chunk_length: int = 16, + mode: str = "joint", + pose_convention: PoseConvention = "backward_framewise", + tolerance_s: float = 1e-4, + viewpoint: Viewpoint = "wrist_view", + action_normalization: str | None = "quantile", + sample_stride: int = 1, + image_key: str = _IMAGE_FEATURE, + eef_in_camera_frame_xyz_wxyz: tuple[float, ...] = _DEFAULT_EEF_IN_CAMERA_FRAME_XYZ_WXYZ, + ) -> None: + if viewpoint != "wrist_view": + raise NotImplementedError("This UMI dataset only supports wrist_view.") + super().__init__( + root=root, + domain_name="umi", + fps=fps, + chunk_length=chunk_length, + mode=mode, + pose_convention=pose_convention, + tolerance_s=tolerance_s, + viewpoint=viewpoint, + action_normalization=action_normalization, + sample_stride=sample_stride, + ) + self._image_key = image_key + + xyz_wxyz = np.asarray(eef_in_camera_frame_xyz_wxyz, dtype=np.float32).reshape(1, 7) + self._eef_in_camera_frame_mat: np.ndarray = build_abs_pose_from_components( + xyz_wxyz[:, :3], xyz_wxyz[:, 3:], "quat_wxyz" + )[0] # [4, 4] + + @property + def action_dim(self) -> int: + return _SINGLE_ARM_ACTION_DIM + + def _action_spec(self) -> ActionSpec: + return build_action_spec(Pos(), Rot("rot6d"), Gripper()) + + @classmethod + def _stats_path(cls) -> Path: + return _NORMALIZER_PATH + + @classmethod + def load_action_stats(cls) -> dict[str, torch.Tensor]: + # Stats file stores 20D bimanual layout (right + left arm). + # Single-arm normalization uses only the first 10D (right arm). + raw = { + key: torch.from_numpy(value).float() + for key, value in load_action_stats(str(cls._stats_path())).items() + } + return {key: tensor[:_SINGLE_ARM_ACTION_DIM] for key, tensor in raw.items()} + + def __getitem__(self, idx: int) -> dict[str, Any]: + mode = self._choose_mode() + idx = int(idx) + row_idx = idx * self._sample_stride + # T+1 rows: current frame + T future frames + observation_rows = self._rows[row_idx : row_idx + self._chunk_length + 1] + + episode = self._episodes[int(observation_rows[0]["episode_index"])] + video = self._load_video(episode, observation_rows) + raw_action, initial_pose = self._build_raw_action(observation_rows) + task = self._tasks[int(observation_rows[0]["task_index"])] + ai_caption = random.choice([part.strip() for part in task.split(" | ") if part.strip()] or [task]) + + return self._build_result( + mode=mode, + video=video, + action=raw_action, + ai_caption=ai_caption, + initial_pose=initial_pose, + ) + + def _load_video(self, episode: dict[str, Any], observation_rows: list[dict[str, Any]]) -> torch.Tensor: + # lerobot is a heavy, optional ("train" extra) dependency; import lazily. + from lerobot.datasets.video_utils import decode_video_frames + + timestamps = [float(row["timestamp"]) for row in observation_rows] + return decode_video_frames( + self._video_path(episode, self._image_key), + [float(episode.get(f"videos/{self._image_key}/from_timestamp", 0.0)) + ts for ts in timestamps], + self._tolerance_s, + ) + + def _build_raw_action( + self, + observation_rows: list[dict[str, Any]], + ) -> tuple[torch.Tensor, torch.Tensor]: + # Trajectory: T+1 poses, [pos(3), quat_wxyz(4)] per frame. + traj = np.asarray([row[_TRAJ_KEY] for row in observation_rows], dtype=np.float32) # [T+1, 7] + poses_abs = build_abs_pose_from_components(traj[:, :3], traj[:, 3:], "quat_wxyz") # [T+1, 4, 4] + + initial_pose = torch.from_numpy(poses_abs[0].copy()).float() + + # Transform from camera TCP frame to EEF frame, then compute relative poses. + eef_poses_abs = poses_abs @ self._eef_in_camera_frame_mat # [T+1, 4, 4] + eef_poses_rel = pose_abs_to_rel( + eef_poses_abs, rotation_format="rot6d", pose_convention=self._pose_convention + ) # [T, 9] + + # Gripper command: future frames only (rows[1:]), matching gripper_indices=[1..T]. + gripper_rows = observation_rows[1:] + gripper_vals = [row[_GRIPPER_KEY] for row in gripper_rows] + gripper = np.asarray( + [float(v) if np.isscalar(v) else float(v[0]) for v in gripper_vals], + dtype=np.float32, + ).reshape(-1, 1) # [T, 1] + + action = np.concatenate([eef_poses_rel, gripper], axis=-1) # [T, 10] + return torch.from_numpy(action).float(), initial_pose diff --git a/examples/worldmodel/cosmos3/vendor/cosmos_framework/data/generator/action/libero_pose_utils.py b/examples/worldmodel/cosmos3/vendor/cosmos_framework/data/generator/action/libero_pose_utils.py new file mode 100644 index 000000000000..23b738c46a88 --- /dev/null +++ b/examples/worldmodel/cosmos3/vendor/cosmos_framework/data/generator/action/libero_pose_utils.py @@ -0,0 +1,69 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: OpenMDW-1.1 + +"""Small LIBERO pose helpers shared by training and closed-loop eval.""" + +from __future__ import annotations + +import numpy as np +import torch + +from cosmos_framework.data.generator.action.utils.pose_utils import ( + RotationConvention, + build_abs_pose_from_components, +) + +# Local-frame post-rotation pattern: +# R_opencv = R_native @ *_TO_OPENCV. +LIBERO_TO_OPENCV: np.ndarray = np.array( + [[0.0, -1.0, 0.0], [1.0, 0.0, 0.0], [0.0, 0.0, 1.0]], + dtype=np.float32, +) + +LIBERO_ROTATION_FORMATS: dict[str, RotationConvention] = { + "3d": "axisangle", + "6d": "rot6d", + "9d": "rot9d", +} +LIBERO_ACTION_DIMS: dict[str, int] = {"3d": 7, "6d": 10, "9d": 13} + + +def libero_rotation_format(rotation_space: str) -> RotationConvention: + """Return the shared ``pose_utils`` rotation format for a LIBERO setting.""" + rotation_format = LIBERO_ROTATION_FORMATS.get(rotation_space) + if rotation_format is None: + raise ValueError(f"Unsupported rotation_space={rotation_space!r}. Use 3d/6d/9d.") + return rotation_format + + +def libero_action_dim(rotation_space: str) -> int: + """Return ``[xyz, rotation, gripper]`` action width for LIBERO.""" + action_dim = LIBERO_ACTION_DIMS.get(rotation_space) + if action_dim is None: + raise ValueError(f"Unsupported rotation_space={rotation_space!r}. Use 3d/6d/9d.") + return action_dim + + +def libero_rotation_space_from_action_dim(action_dim: int) -> str: + """Infer LIBERO rotation space from unpadded action width.""" + for rotation_space, dim in LIBERO_ACTION_DIMS.items(): + if dim == action_dim: + return rotation_space + raise ValueError(f"Unable to infer rotation_space from action_dim={action_dim}.") + + +def build_libero_abs_pose(state_raw: torch.Tensor | np.ndarray, *, to_opencv: bool) -> np.ndarray: + """Build absolute LIBERO EE poses from state rows. + + ``state_raw`` is ``[x,y,z,axisangle(3),gripper(2)]``. When requested, the + local EE frame is post-rotated into the shared OpenCV-style action frame. + """ + if isinstance(state_raw, torch.Tensor): + state_np = state_raw.detach().cpu().numpy().astype(np.float32, copy=False) + else: + state_np = np.asarray(state_raw, dtype=np.float32) + + poses_abs = build_abs_pose_from_components(state_np[:, :3], state_np[:, 3:6], "axisangle") + if to_opencv: + poses_abs[:, :3, :3] = poses_abs[:, :3, :3] @ LIBERO_TO_OPENCV + return poses_abs diff --git a/examples/worldmodel/cosmos3/vendor/cosmos_framework/data/generator/action/utils/__init__.py b/examples/worldmodel/cosmos3/vendor/cosmos_framework/data/generator/action/utils/__init__.py new file mode 100644 index 000000000000..28a81be6930a --- /dev/null +++ b/examples/worldmodel/cosmos3/vendor/cosmos_framework/data/generator/action/utils/__init__.py @@ -0,0 +1,2 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: OpenMDW-1.1 diff --git a/examples/worldmodel/cosmos3/vendor/cosmos_framework/data/generator/action/utils/action_processing.py b/examples/worldmodel/cosmos3/vendor/cosmos_framework/data/generator/action/utils/action_processing.py new file mode 100644 index 000000000000..993d1a458907 --- /dev/null +++ b/examples/worldmodel/cosmos3/vendor/cosmos_framework/data/generator/action/utils/action_processing.py @@ -0,0 +1,540 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: OpenMDW-1.1 + +"""Shared action processing records and normalization helpers.""" + +from __future__ import annotations + +import json +import math +from dataclasses import dataclass +from pathlib import Path +from typing import Any, Literal, Protocol + +import numpy as np +import torch + +from cosmos_framework.utils import log + +ActionNormalizationMethod = Literal[ + "quantile", + "quantile_rot", + "quantile_rot_scale_floor", + "asinh_rot", + "piecewise_asinh_rot", + "global_asinh", + "meanstd", + "minmax", +] +QUANTILE_ROT_SCALE_FLOOR_TARGET_MAX = 5.0 +_RAW_ROTATION_STATS_METHODS = frozenset( + {"quantile_rot", "quantile_rot_scale_floor", "asinh_rot", "piecewise_asinh_rot"} +) +_GLOBAL_ASINH_POSE_CONVENTIONS = frozenset( + {"backward_anchored", "backward_chunk_anchored_8f", "backward_chunk_anchored_16f"} +) +GLOBAL_ASINH_GRIPPER_01_STATS_PATH = ( + Path(__file__).parents[1] / "normalizers" / "global_asinh_v1_backward_anchored_gripper_binary.json" +) +GLOBAL_ASINH_METRIC_GRIPPER_STATS_PATH = ( + Path(__file__).parents[1] / "normalizers" / "global_asinh_v1_backward_anchored_gripper_metric.json" +) +ActionNormalizationStats = dict[str, torch.Tensor] + + +class ActionNormalizer(Protocol): + """Tensor-level action normalization interface used by ActionProcessor.""" + + def normalize_action(self, action: torch.Tensor) -> torch.Tensor: # action: [...,D], returns [...,D] + """Map raw action values into model-space values.""" + ... + + def denormalize_action(self, action: torch.Tensor) -> torch.Tensor: # action: [...,D], returns [...,D] + """Invert model-space action values back into raw action values.""" + ... + + +class ActionCodec(Protocol): + """Decode a canonical normalized action into canonical raw space.""" + + def decode_to_unified( + self, normalized_action: torch.Tensor + ) -> torch.Tensor: # normalized_action: [...,D], returns [...,D] + """Decode normalized values into the canonical raw action layout.""" + ... + + +@dataclass(frozen=True) +class ActionAffineNormalization: + """Resolved affine action normalizer. + + Forward normalization is ``(raw - offset) / scale``. Inverse + denormalization is ``normalized * scale + offset``. + + ``forward_clamp`` records lossy range-style forward clamping. When + ``forward_clamp_mask`` is provided, only channels with a ``True`` mask + entry are clamped; this represents mixed UMI normalizers where some fields + are range-clamped and others are plain affine transforms. + """ + + offset: torch.Tensor + scale: torch.Tensor + forward_clamp: tuple[float, float] | None = None + forward_clamp_mask: torch.Tensor | None = None + + def normalize_action(self, action: torch.Tensor) -> torch.Tensor: # action: [...,D], returns [...,D] + """Normalize raw action values with resolved affine parameters.""" + offset = self.offset.to(device=action.device, dtype=action.dtype) # [D] + scale = self.scale.to(device=action.device, dtype=action.dtype) # [D] + normalized = (action - offset) / scale # [...,D] + if self.forward_clamp is not None: + lo, hi = self.forward_clamp + clamped = normalized.clamp(lo, hi) # [...,D] + if self.forward_clamp_mask is None: + normalized = clamped # [...,D] + else: + clamp_mask = self.forward_clamp_mask.to(device=action.device, dtype=torch.bool) # [D] + normalized = torch.where(clamp_mask, clamped, normalized) # [...,D] + return normalized # [...,D] + + def denormalize_action(self, action: torch.Tensor) -> torch.Tensor: # action: [...,D], returns [...,D] + """Invert action normalization with resolved affine parameters.""" + offset = self.offset.to(device=action.device, dtype=action.dtype) # [D] + scale = self.scale.to(device=action.device, dtype=action.dtype) # [D] + return action * scale + offset # [...,D] + + +@dataclass(frozen=True) +class ActionAsinhNormalization: + """Invertible heavy-tail action normalizer. + + The forward transform first applies the same quantile affine transform used + by ``quantile_rot`` and then compresses large-but-valid tails with + ``asinh``: + + ``y = asinh((raw - offset) / scale) / asinh(1)``. + + The ``asinh(1)`` factor keeps q01/q99-style unit deviations near ±1 while + mapping 100-250× outliers to single-digit targets instead of exploding the + flow-matching loss. The inverse uses ``sinh`` and is exact up to floating + point precision; unlike clamp, no labels are censored. + """ + + offset: torch.Tensor + scale: torch.Tensor + unit: float = math.asinh(1.0) + action_schema: Literal["native", "unified_v1"] = "native" + + def _validate_width(self, action: torch.Tensor) -> None: # action: [...,D] + if action.ndim == 0 or action.shape[-1] != self.offset.numel(): + raise ValueError( + f"action width {action.shape[-1] if action.ndim else 0} does not match normalizer width " + f"{self.offset.numel()}" + ) + + def normalize_action(self, action: torch.Tensor) -> torch.Tensor: # action: [...,D], returns [...,D] + """Normalize raw action values with an invertible asinh transform.""" + self._validate_width(action) + offset = self.offset.to(device=action.device, dtype=action.dtype) # [D] + scale = self.scale.to(device=action.device, dtype=action.dtype) # [D] + unit = torch.as_tensor(self.unit, device=action.device, dtype=action.dtype) # [] + return torch.asinh((action - offset) / scale) / unit # [...,D] + + def denormalize_action(self, action: torch.Tensor) -> torch.Tensor: # action: [...,D], returns [...,D] + """Invert asinh action normalization.""" + self._validate_width(action) + offset = self.offset.to(device=action.device, dtype=action.dtype) # [D] + scale = self.scale.to(device=action.device, dtype=action.dtype) # [D] + unit = torch.as_tensor(self.unit, device=action.device, dtype=action.dtype) # [] + return torch.sinh(action * unit) * scale + offset # [...,D] + + +@dataclass(frozen=True) +class ActionPiecewiseAsinhNormalization: + """Invertible tail-only asinh action normalizer. + + This uses the same affine ``z = (raw - offset) / scale`` coordinates as + ``quantile_rot`` and keeps the common range exactly linear: + + ``y = z`` when ``|z| <= 1``. + + Outside that range it grows sublinearly while preserving magnitude order: + + ``y = sign(z) * (1 + asinh(beta * (|z| - 1)) / beta)``. + + With the default ``beta=1``, the transform is continuous and has matching + derivative at ±1, so in-range samples match the linear/clamped normalizer + exactly while out-of-range labels remain recoverable instead of censored. + """ + + offset: torch.Tensor + scale: torch.Tensor + beta: float = 1.0 + + def _beta_tensor(self, reference: torch.Tensor) -> torch.Tensor: + beta = torch.as_tensor(self.beta, device=reference.device, dtype=reference.dtype) # [] + if bool((beta <= 0).item()): + raise ValueError(f"piecewise asinh beta must be positive, got {self.beta}") + return beta + + def normalize_action(self, action: torch.Tensor) -> torch.Tensor: # action: [...,D], returns [...,D] + """Normalize raw action values, keeping affine-normalized [-1, 1] exactly unchanged.""" + offset = self.offset.to(device=action.device, dtype=action.dtype) # [D] + scale = self.scale.to(device=action.device, dtype=action.dtype) # [D] + beta = self._beta_tensor(action) # [] + z = (action - offset) / scale # [...,D] + abs_z = z.abs() # [...,D] + tail = 1.0 + torch.asinh(beta * (abs_z - 1.0)) / beta # [...,D] + return torch.where(abs_z <= 1.0, z, z.sign() * tail) # [...,D] + + def denormalize_action(self, action: torch.Tensor) -> torch.Tensor: # action: [...,D], returns [...,D] + """Invert piecewise-asinh action normalization.""" + offset = self.offset.to(device=action.device, dtype=action.dtype) # [D] + scale = self.scale.to(device=action.device, dtype=action.dtype) # [D] + beta = self._beta_tensor(action) # [] + abs_y = action.abs() # [...,D] + tail = 1.0 + torch.sinh(beta * (abs_y - 1.0)) / beta # [...,D] + z = torch.where(abs_y <= 1.0, action, action.sign() * tail) # [...,D] + return z * scale + offset # [...,D] + + +def load_action_stats(stats_path: str, stats_key: str = "global") -> dict[str, np.ndarray]: + """Load pre-computed action normalization stats from a JSON file.""" + path = Path(stats_path) + if not path.exists(): + raise FileNotFoundError(f"Action normalization stats not found at {stats_path}.") + log.info(f"Loading action normalization stats from {stats_path}") + with path.open("r") as f: + raw = json.load(f) + stat_keys = {"mean", "std", "min", "max", "q01", "q99"} + if stats_key in raw: + raw = raw[stats_key] + if not isinstance(raw, dict): + raise TypeError(f"Action normalization stats block {stats_key!r} in {stats_path} must be a dict.") + elif stats_key != "global" and not any(key in raw for key in stat_keys): + raise KeyError(f"Action normalization stats block {stats_key!r} not found in {stats_path}.") + return {k: np.array(v, dtype=np.float32) for k, v in raw.items() if k in stat_keys} + + +def resolve_action_normalization( + method: ActionNormalizationMethod, + stats: ActionNormalizationStats | None = None, + apply_forward_clamp: bool = False, + *, + pose_convention: str | None = None, +) -> ActionNormalizer: + """Resolve one configured normalization method into its forward/inverse transform.""" + if stats is None: + raise ValueError(f"{method} normalization requires per-dataset action statistics") + if method == "global_asinh" and pose_convention not in _GLOBAL_ASINH_POSE_CONVENTIONS: + raise ValueError(f"global_asinh normalization requires an anchored pose convention, got {pose_convention!r}") + + if method == "meanstd": + offset = stats["mean"] # [D] + scale = stats["std"].clamp(min=1e-8) # [D] + return ActionAffineNormalization(offset=offset, scale=scale) + + if method == "minmax": + lo = stats["min"] # [D] + hi = stats["max"] # [D] + elif method in ( + "quantile", + "quantile_rot", + "quantile_rot_scale_floor", + "asinh_rot", + "piecewise_asinh_rot", + "global_asinh", + ): + lo = stats["q01"] # [D] + hi = stats["q99"] # [D] + else: + raise ValueError(f"Unknown normalization method: {method!r}") + + offset = (hi + lo) / 2.0 # [D] + scale = (hi - lo).clamp(min=1e-8) / 2.0 # [D] + if method == "quantile_rot_scale_floor": + min_dev = (stats["min"] - offset).abs() # [D] + max_dev = (stats["max"] - offset).abs() # [D] + scale_floor = torch.maximum(min_dev, max_dev) / QUANTILE_ROT_SCALE_FLOOR_TARGET_MAX # [D] + scale = torch.maximum(scale, scale_floor.clamp(min=1e-8)) # [D] + + if method in {"asinh_rot", "global_asinh"}: + action_schema = "unified_v1" if method == "global_asinh" else "native" + return ActionAsinhNormalization(offset=offset, scale=scale, action_schema=action_schema) + if method == "piecewise_asinh_rot": + return ActionPiecewiseAsinhNormalization(offset=offset, scale=scale) + + if apply_forward_clamp: + # Ideally this hardcode should be removed, but for now we keep it so we can be aligned with mid-training checkpoints. + forward_clamp = (-1.0, 1.0) + else: + forward_clamp = None + + return ActionAffineNormalization( + offset=offset, + scale=scale, + forward_clamp=forward_clamp, + ) + + +def load_action_normalization_stats( + action_normalization: ActionNormalizationMethod, + *, + stats_path: str | Path, + stats_key: str | None = None, + expected_dim: int | None = None, +) -> dict[str, torch.Tensor]: + """Load tensor action stats for a configured normalization method.""" + if action_normalization == "global_asinh": + if expected_dim is not None and expected_dim != 59: + raise ValueError(f"global_asinh stats require expected_dim=59, got {expected_dim}") + expected_dim = 59 + if stats_key is None: + stats_key = "global_raw" if action_normalization in _RAW_ROTATION_STATS_METHODS else "global" + raw_stats = load_action_stats(str(stats_path), stats_key=stats_key) + if expected_dim is not None: + for name, values in raw_stats.items(): + if values.shape != (expected_dim,): + raise ValueError(f"Action stats {name!r} must have shape [{expected_dim}], got {values.shape}.") + return {name: torch.from_numpy(values).float() for name, values in raw_stats.items()} + + +def load_action_normalizer( + action_normalization: ActionNormalizationMethod | None, + *, + stats_path: str | Path, + apply_forward_clamp: bool, + stats_key: str | None = None, + expected_dim: int | None = None, +) -> ActionNormalizer | None: + """Load and resolve an action normalizer from pre-computed stats.""" + if action_normalization is None: + return None + return resolve_action_normalization( + action_normalization, + load_action_normalization_stats( + action_normalization, + stats_path=stats_path, + stats_key=stats_key, + expected_dim=expected_dim, + ), + apply_forward_clamp=apply_forward_clamp, + ) + + +def make_camera_global_asinh_normalizer( + stats_path: str | Path = GLOBAL_ASINH_METRIC_GRIPPER_STATS_PATH, + *, + pose_convention: str = "backward_anchored", +) -> ActionAsinhNormalization: + """Build the native 9D camera normalizer from the canonical global-asinh profile.""" + if pose_convention not in _GLOBAL_ASINH_POSE_CONVENTIONS: + raise ValueError(f"global_asinh normalization requires an anchored pose convention, got {pose_convention!r}") + stats = load_action_normalization_stats( + "global_asinh", + stats_path=stats_path, + expected_dim=59, + ) + lo = stats["q01"][:9] + hi = stats["q99"][:9] + return ActionAsinhNormalization(offset=(hi + lo) / 2.0, scale=(hi - lo).clamp(min=1e-8) / 2.0) + + +def make_pose_action_scale_normalizer( + action_dim: int, + *, + translation_scale: float = 1.0, + rotation_scale: float = 1.0, +) -> ActionAffineNormalization: + """Build a normalizer that maps raw pose deltas into scaled model space. + + Pose actions use the shared layout ``[translation(3), rotation(...)]``. + The returned normalizer multiplies translation channels by + ``translation_scale`` and rotation channels by ``rotation_scale`` during + preprocessing, then inverts those factors during postprocessing. + """ + if action_dim < 3: + raise ValueError(f"Pose action_dim must be at least 3, got {action_dim}") + if translation_scale == 0: + raise ValueError("translation_scale must be non-zero") + if rotation_scale == 0: + raise ValueError("rotation_scale must be non-zero") + + offset = torch.zeros(action_dim, dtype=torch.float32) # [D] + scale = torch.ones(action_dim, dtype=torch.float32) # [D] + scale[:3] = 1.0 / float(translation_scale) # [D] + if action_dim > 3: + scale[3:] = 1.0 / float(rotation_scale) # [D] + return ActionAffineNormalization(offset=offset, scale=scale) + + +@dataclass(frozen=True) +class ActionProcessingRecord: + """Per-sample metadata needed to invert Action model-space preprocessing.""" + + raw_action_dim: int + action_normalizer: ActionNormalizer | None + action_valid_mask: torch.Tensor | None = None # [D_raw] bool + action_codec: ActionCodec | None = None + + def __post_init__(self) -> None: + if self.action_normalizer is not None and self.action_codec is not None: + raise ValueError("ActionProcessingRecord cannot contain both an action_normalizer and an action_codec") + + +def pad_action_to_max_dim( + action: torch.Tensor, max_action_dim: int +) -> torch.Tensor: # action: [T,D], returns [T,D_model] + """Pad action tensor to max_action_dim along the last dimension. + + Args: + action: Action tensor of shape (T, D) where D is the current action dimension. + max_action_dim: Target action dimension to pad to. + + Returns: + Padded action tensor of shape (T, max_action_dim). + """ + if action.shape[-1] > max_action_dim: + raise ValueError(f"Action dimension {action.shape[-1]} is greater than max_action_dim {max_action_dim}") + if action.shape[-1] == max_action_dim: + return action # [T,D_model] + padding_size = max_action_dim - action.shape[-1] + zero_padding = torch.zeros(*action.shape[:-1], padding_size, dtype=action.dtype, device=action.device) # [T,D_pad] + return torch.cat([action, zero_padding], dim=-1) # [T,D_model] + + +def make_batched_action_processing_fields( + record: ActionProcessingRecord, + batch_size: int, + *, + action_channel_masking: bool = True, +) -> dict[str, list[torch.Tensor | ActionProcessingRecord | None]]: + """Build batch-list fields whose action width and inverse record cannot drift apart.""" + raw_action_dim = torch.tensor(record.raw_action_dim, dtype=torch.long) if action_channel_masking else None # [] + return { + "raw_action_dim": [raw_action_dim] * batch_size, + "action_processing_record": [record] * batch_size, + } + + +class ActionProcessor: + """Build model-space actions and decode model outputs for a single sample.""" + + def __init__(self, max_action_dim: int, action_channel_masking: bool = True) -> None: + self.max_action_dim = int(max_action_dim) + self.action_channel_masking = bool(action_channel_masking) + + def preprocess_action( + self, + data_dict: dict[str, Any], + action: torch.Tensor, + *, + action_normalizer: ActionNormalizer | None, + action_valid_mask: torch.Tensor | None = None, # [D_raw] bool + normalized_action: torch.Tensor | None = None, # [T,D_raw] + action_codec: ActionCodec | None = None, + ) -> dict[str, Any]: + """Return a sample with canonical raw and normalized action fields. + + ``action_raw`` is the exact unnormalized, unpadded action accepted by + this method. It is the canonical external-space target for evaluation. + ``action`` is the normalized and padded model-space representation. + Unified actions may provide that normalized representation explicitly so + mapping and native normalization stay outside this padding boundary. + """ + + raw_action_dim = int(action.shape[-1]) + if "action_raw" in data_dict: + raise ValueError("action_raw is reserved for the canonical ActionProcessor input") + action_raw = action.clone() # [T,D] + if action_normalizer is not None and action_codec is not None: + raise ValueError("Provide either action_normalizer or action_codec, not both") + if (normalized_action is None) != (action_codec is None): + raise ValueError("normalized_action and action_codec must be provided together") + if normalized_action is not None and action_normalizer is not None: + raise ValueError("normalized_action cannot be combined with action_normalizer") + if normalized_action is None: + normalized_action = action # [T,D] + if action_normalizer is not None: + normalized_action = action_normalizer.normalize_action(action) # [T,D] + if normalized_action.shape != action.shape: + raise ValueError( + f"normalized_action shape {tuple(normalized_action.shape)} does not match " + f"raw action shape {tuple(action.shape)}" + ) + + if action_valid_mask is not None: + if action_valid_mask.dtype != torch.bool or action_valid_mask.ndim != 1: + raise ValueError( + f"action_valid_mask must be a 1D bool tensor, got " + f"shape={tuple(action_valid_mask.shape)} dtype={action_valid_mask.dtype}" + ) + if action_valid_mask.numel() != raw_action_dim: + raise ValueError( + f"action_valid_mask width {action_valid_mask.numel()} != raw_action_dim {raw_action_dim}" + ) + action_valid_mask = action_valid_mask.to(device=normalized_action.device).contiguous() # [D_raw] + + processed_data_dict = dict(data_dict) + processed_data_dict["action_raw"] = action_raw + processed_data_dict["action"] = pad_action_to_max_dim( # [T,D_model] + normalized_action, self.max_action_dim + ) + record = ActionProcessingRecord( + raw_action_dim=raw_action_dim, + action_normalizer=action_normalizer, + action_valid_mask=action_valid_mask, + action_codec=action_codec, + ) + processed_data_dict["raw_action_dim"] = ( + torch.tensor(record.raw_action_dim, dtype=torch.long) if self.action_channel_masking else None + ) # [] + processed_data_dict["action_processing_record"] = record + if action_valid_mask is not None: + processed_data_dict["action_valid_mask"] = pad_action_to_max_dim(action_valid_mask, self.max_action_dim).to( + dtype=torch.bool + ) + return processed_data_dict + + @staticmethod + def _unpad_action(action: torch.Tensor, raw_action_dim: int) -> torch.Tensor: + """Drop model-only padded action channels.""" + if action.shape[-1] < raw_action_dim: + raise ValueError(f"invalid raw_action_dim={raw_action_dim} for action with shape {tuple(action.shape)}") + return action[..., :raw_action_dim].contiguous() # [...,D_raw] + + @staticmethod + def postprocess_action( + action: torch.Tensor, + record: ActionProcessingRecord, + ) -> torch.Tensor: + """Unpad and decode a generated model-space action into canonical raw space.""" + action = ActionProcessor._unpad_action(action, record.raw_action_dim) # [...,D_raw] + if record.action_codec is not None: + action = record.action_codec.decode_to_unified(action) # [...,D_raw] + elif record.action_normalizer is not None: + action = record.action_normalizer.denormalize_action(action) # [...,D_raw] + return action # [...,D_raw] + + @staticmethod + def postprocess_action_to_schema( + action: torch.Tensor, + record: ActionProcessingRecord, + ) -> torch.Tensor: + """Decode a generated model-space action into its canonical raw schema.""" + return ActionProcessor.postprocess_action(action, record) + + +def get_action_processing_records(data_batch: dict[str, Any]) -> list[ActionProcessingRecord | None]: + """Read all per-sample processing records from a collated Action batch.""" + records = data_batch.get("action_processing_record") + if records is None: + return [] + if isinstance(records, ActionProcessingRecord): + return [records] + if isinstance(records, list): + for record in records: + if record is not None and not isinstance(record, ActionProcessingRecord): + raise TypeError(f"Unexpected action_processing_record entry type: {type(record).__name__}") + return records + raise TypeError(f"Unexpected action_processing_record type: {type(records).__name__}") diff --git a/examples/worldmodel/cosmos3/vendor/cosmos_framework/data/generator/action/utils/action_spec.py b/examples/worldmodel/cosmos3/vendor/cosmos_framework/data/generator/action/utils/action_spec.py new file mode 100644 index 000000000000..b93809710f2d --- /dev/null +++ b/examples/worldmodel/cosmos3/vendor/cosmos_framework/data/generator/action/utils/action_spec.py @@ -0,0 +1,235 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: OpenMDW-1.1 + +"""Shared action-vector specification: per-dim type label + idle thresholds. + +Single concept: every column of an action vector has a :class:`DimType` label. +Idle detection iterates by type and applies the matching algorithm: + + POS → ‖action[pos_idx]‖ per arm < eps_t + ROT → distance(rot, identity) per group < eps_r + GRIPPER → max |Δgripper| < eps_g (frame 0 idle by convention) + JOINT → max |Δjoint| < joint_threshold (frame 0 idle) + RESERVED → ignored + +An :class:`ActionSpec` is just ``names`` + ``types`` + ``rotation_format``. +Build one declaratively via :func:`build_action_spec` from DSL components:: + + build_action_spec(Pos(), Rot("rot6d"), Gripper()) # 10D single arm + build_action_spec(Pos(), Rot("rot6d")) # 9D no gripper + build_action_spec(Joint(n=14, label="arm"), # 30D joint-space + Joint(n=14, label="end"), + Joint(n=2, label="gripper")) + build_action_spec(Pos(prefix="left"), Rot("rot6d", "left"), Gripper(prefix="left"), + Pos(prefix="right"), Rot("rot6d", "right"), Gripper(prefix="right")) + +Naming convention: + Default ``pos_x``, ``rot_0``, ``gripper``, ``arm_0`` ... + With ``prefix="left"`` (idempotent on trailing ``_``): ``left_pos_x`` ... +""" + +from __future__ import annotations + +from dataclasses import dataclass +from enum import Enum +from typing import ClassVar + +from cosmos_framework.data.generator.action.utils.pose_utils import ( + RotationConvention, + _identity_rotation_vector, +) + + +class DimType(str, Enum): + """Per-column action-dim category (drives idle detection).""" + + POS = "pos" + ROT = "rot" + GRIPPER = "gripper" + JOINT = "joint" + RESERVED = "reserved" + + +@dataclass(frozen=True, slots=True) +class ActionSpec: + """Structural description of an action vector: names + per-dim types. + + All ROT dims share a single ``rotation_format``; mixed formats in one spec + are not supported (raise at build time). + + This struct contains no detection thresholds — those are passed at call + time to :func:`compute_idle_frames` so each dataset can tune them + independently of layout. + """ + + names: list[str] + types: list[DimType] + rotation_format: RotationConvention = "rot6d" + + @property + def dim(self) -> int: + return len(self.names) + + +# --------------------------------------------------------------------------- +# DSL components +# --------------------------------------------------------------------------- + + +def _join_prefix(prefix: str, name: str) -> str: + """Join ``prefix`` and ``name`` with a single ``_``; idempotent on trailing ``_``.""" + return name if not prefix else f"{prefix.rstrip('_')}_{name}" + + +@dataclass(frozen=True) +class Pos: + """Translation block. + + Default 3D (``pos_x``, ``pos_y``, ``pos_z``). For planar tasks (e.g. PushT) + use ``Pos(dim=2)`` → ``pos_x``, ``pos_y``. ``dim >= 4`` falls back to + indexed names ``pos_0``, ``pos_1``, ... + """ + + dim: int = 3 + prefix: str = "" + type: ClassVar[DimType] = DimType.POS + + def names(self) -> list[str]: + if self.dim <= 3: + return [_join_prefix(self.prefix, f"pos_{c}") for c in "xyz"[: self.dim]] + return [_join_prefix(self.prefix, f"pos_{i}") for i in range(self.dim)] + + +@dataclass(frozen=True) +class Rot: + """Rotation block; ``format`` selects the encoding. + + Supported formats and per-dim names: + + - ``rot6d`` → 6 dims, ``rot_0`` ... ``rot_5`` (identity ``[1,0,0,0,1,0]``) + - ``rot9d`` → 9 dims, ``rot_0`` ... ``rot_8`` (identity ``[1,0,0,0,1,0,0,0,1]``) + - ``euler_xyz`` → 3 dims, ``roll``, ``pitch``, ``yaw`` (identity ``[0,0,0]``) + - ``axisangle`` → 3 dims, ``axang_x/y/z`` (identity ``[0,0,0]``) + - ``quat_xyzw`` / ``quat_wxyz`` → 4 dims, ``quat_x/y/z/w`` in declared order + """ + + format: RotationConvention = "rot6d" + prefix: str = "" + type: ClassVar[DimType] = DimType.ROT + + @property + def rotation_format(self) -> RotationConvention: + return self.format + + @property + def dim(self) -> int: + return _identity_rotation_vector(self.format).shape[0] + + def names(self) -> list[str]: + if self.format == "euler_xyz": + return [_join_prefix(self.prefix, c) for c in ("roll", "pitch", "yaw")] + if self.format == "axisangle": + return [_join_prefix(self.prefix, f"axang_{c}") for c in "xyz"] + if self.format.startswith("quat_"): + order = self.format.split("_", 1)[1] # "xyzw" or "wxyz" + return [_join_prefix(self.prefix, f"quat_{c}") for c in order] + return [_join_prefix(self.prefix, f"rot_{i}") for i in range(self.dim)] + + +@dataclass(frozen=True) +class Gripper: + """1D gripper command (binary 0/1 or continuous). Detected by frame-diff.""" + + prefix: str = "" + type: ClassVar[DimType] = DimType.GRIPPER + + @property + def dim(self) -> int: + return 1 + + def names(self) -> list[str]: + return [_join_prefix(self.prefix, "gripper")] + + +@dataclass(frozen=True) +class Joint: + """``n`` joint commands. Detected by frame-diff against ``joint_threshold``.""" + + n: int = 0 + label: str = "joint" + prefix: str = "" + type: ClassVar[DimType] = DimType.JOINT + + @property + def dim(self) -> int: + return self.n + + def names(self) -> list[str]: + return [_join_prefix(self.prefix, f"{self.label}_{i}") for i in range(self.n)] + + +@dataclass(frozen=True) +class Reserved: + """``n`` dims counted in ``action_dim`` but ignored by idle detection.""" + + n: int = 0 + label: str = "reserved" + prefix: str = "" + type: ClassVar[DimType] = DimType.RESERVED + + @property + def dim(self) -> int: + return self.n + + def names(self) -> list[str]: + return [_join_prefix(self.prefix, f"{self.label}_{i}") for i in range(self.n)] + + +# --------------------------------------------------------------------------- +# Builder +# --------------------------------------------------------------------------- + + +# Type alias for any DSL component. Not a runtime check — only annotation hint. +Component = Pos | Rot | Gripper | Joint | Reserved + + +def build_action_spec(*components: Component) -> ActionSpec: + """Compose ``components`` into an :class:`ActionSpec`. + + Each component contributes its ``names()`` and replicates its ``type`` for + every column it occupies. The first ROT component's ``rotation_format`` + is captured for the whole spec; mixing formats raises ``ValueError``. + """ + names: list[str] = [] + types: list[DimType] = [] + rotation_format: RotationConvention | None = None + + for c in components: + names.extend(c.names()) + types.extend([c.type] * c.dim) + if c.type == DimType.ROT: + fmt = c.rotation_format # type: ignore[union-attr] + if rotation_format is None: + rotation_format = fmt + elif rotation_format != fmt: + raise ValueError(f"Mixed rotation_format in one ActionSpec: {rotation_format!r} vs {fmt!r}") + + return ActionSpec( + names=names, + types=types, + rotation_format=rotation_format or "rot6d", + ) + + +__all__ = [ + "ActionSpec", + "Component", + "DimType", + "Gripper", + "Joint", + "Pos", + "Reserved", + "Rot", + "build_action_spec", +] diff --git a/examples/worldmodel/cosmos3/vendor/cosmos_framework/data/generator/action/utils/domain_utils.py b/examples/worldmodel/cosmos3/vendor/cosmos_framework/data/generator/action/utils/domain_utils.py new file mode 100644 index 000000000000..4d00f36edbbb --- /dev/null +++ b/examples/worldmodel/cosmos3/vendor/cosmos_framework/data/generator/action/utils/domain_utils.py @@ -0,0 +1,130 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: OpenMDW-1.1 + +"""Domain ID helpers for cross-embodiment action datasets.""" + +EMBODIMENT_TO_DOMAIN_ID: dict[str, int] = { + "no_action": 0, + "av": 1, + "camera_pose": 2, + "hand_pose": 3, + # WebHumanAction omits the camera prefix from its 48D native hand layout, + # so it must not reuse hand_pose/embodiment_a's camera-inclusive domain 3 projector. + "webhumanaction_hand": 31, + "pusht": 4, + "libero": 5, + "umi": 6, + "bridge_orig_lerobot": 7, + "droid_lerobot": 8, + "robomind-franka": 8, # Both Droid and RoboMIND-Franka are using robotiq and franka + "embodiment_b": 9, + "robomind-franka-dual": 12, + "robomind-ur": 13, + "agibotworld": 15, + "embodiment_c_gripper": 15, + "embodiment_c_gripper_ext": 15, + "xdof_yam": 16, + "molmoact2_yam": 16, # MolmoAct2 uses the same YAM 20D FK action contract + "abc_yam": 16, # ABC uses the same YAM 20D FK action contract + "robotwin": 17, # RoboTwin dual-arm ALOHA (14D absolute joint_pos) + # Shares RoboTwin's domain: the same arm, differing only by millimetre link + # lengths and a 0.9 deg yaw calibration baked into RoboTwin's URDF. + "robodojo": 17, + "fractal": 20, + "drawanything": 21, + "behavior1k_lerobot": 22, # BEHAVIOR-1K R1Pro mobile bimanual (23D joint action) + "maniparena": 23, # ManipArena x2robot/ex001_6r dual-arm; own 20D EE-direct action projection + # New dedicated slot (not reusing agibot's domain 15): WebHumanAction body + # (ego/head+wrists+fingertips, no camera action, 57D) trains its own action2llm/llm2action + # DomainAwareLinear weights from scratch instead of continuing agibot's. + "webhumanaction_body": 24, + "so101-molmo-midtrain-15hz": 25, + "so100-molmo-midtrain-15hz": 26, + "so101-bimanual-midtrain-conditional": 27, + # GenieSim G2A is the simulated counterpart of Embodiment C and emits the + # same 29-D [head, right arm/gripper, left arm/gripper] contract. + "geniesim3_g2a": 15, + "geniesim3_g2a_joint": 28, + # ManipArena mobile-manipulation children only: 29D head+dual-arm SE(3). A + # dedicated slot rather than sharing domain 23, because the width and semantics + # differ from the 20D contract and it trains its own action2llm/llm2action + # DomainAwareLinear weights. New embodiments append above the maximum; the gaps + # at 10/11/14/18/19 index retired embodiments and are not reused. + "maniparena_mobile": 29, + # RoboCasa PandaOmron mobile manipulation (10/15/20D raw action per + # ``use_base_action`` / ``base_encoding``); appended above the maximum. + "robocasa": 30, + # embodiment_b nvidia-20260828 ingestion: a new one-shot dataset, distinct from + # "embodiment_b" (domain 9, an earlier unrelated sample drop with its own 30D + # contract). + "embodiment_b_20260828": 32, +} + + +EMBODIMENT_TO_RAW_ACTION_DIM: dict[str, int] = { + "av": 9, + "camera_pose": 9, + "pusht": 2, + "umi": 10, + "bridge_orig_lerobot": 10, + "droid_lerobot": 10, + "robomind-franka": 10, + "robomind-franka-dual": 20, + "robomind-ur": 10, + "embodiment_b": 30, + "agibotworld": 29, + "embodiment_c_gripper": 29, + "embodiment_c_gripper_ext": 29, + "webhumanaction_body": 57, # ego/head(9) + [R_wrist(9)+R_fingertips(15)] + [L_wrist(9)+L_fingertips(15)] + "xdof_yam": 20, + "molmoact2_yam": 20, + "abc_yam": 20, + "maniparena": 20, # dual-arm EE: [pos(3)+rot6d(6)+gripper(1)] x 2 + # head(9) + left(9)+grip(1) + right(9)+grip(1); head/left/right, not head/right/left + "maniparena_mobile": 29, + "robotwin": 14, # dual-arm ALOHA: [L 6 joints + 1 gripper, R 6 joints + 1 gripper] + "robodojo": 14, # dual-arm ARX-X5: [L 6 joints + 1 gripper, R 6 joints + 1 gripper] + "fractal": 10, + "drawanything": 3, + "behavior1k_lerobot": 23, # base(3) trunk(4) arms(14) grippers(2) + "so101-molmo-midtrain-15hz": 10, + "so100-molmo-midtrain-15hz": 10, + "so101-bimanual-midtrain-conditional": 20, + "geniesim3_g2a": 29, + "geniesim3_g2a_joint": 16, + "embodiment_b_20260828": 50, + # NOTE: ``libero`` (7/10/13 depending on ``rotation_space``), ``hand_pose`` + # (variable with ``keypoint_option`` and ``rotation_format``) and ``robocasa`` + # (10 arm-only, 15/20 with the mobile base, per ``use_base_action`` / + # ``base_encoding``) are absent because their raw width is set per-dataset at + # construction time. Inference in inverse_dynamics/WAM modes is not supported + # for these domains until canonical widths are added here. +} + + +def get_domain_id(embodiment_type: str) -> int: + """Get the domain ID for a given embodiment type.""" + key = embodiment_type.lower().strip() + if key not in EMBODIMENT_TO_DOMAIN_ID: + raise KeyError( + f"Unknown embodiment type: {embodiment_type!r}. " + f"Available embodiments: {sorted(EMBODIMENT_TO_DOMAIN_ID.keys())}" + ) + return EMBODIMENT_TO_DOMAIN_ID[key] + + +def get_action_dim(embodiment_type: str) -> int: + """Get the raw action dimension for a given embodiment type.""" + key = embodiment_type.lower().strip() + if key not in EMBODIMENT_TO_RAW_ACTION_DIM: + raise KeyError( + f"Unknown embodiment type: {embodiment_type!r}. " + f"Available embodiments: {sorted(EMBODIMENT_TO_RAW_ACTION_DIM.keys())}" + ) + return EMBODIMENT_TO_RAW_ACTION_DIM[key] + + +def is_valid_domain_name(embodiment_type: str) -> bool: + """Check if the given embodiment type is recognized.""" + key = embodiment_type.lower().strip() + return key in EMBODIMENT_TO_RAW_ACTION_DIM diff --git a/examples/worldmodel/cosmos3/vendor/cosmos_framework/data/generator/action/utils/episode_filters.py b/examples/worldmodel/cosmos3/vendor/cosmos_framework/data/generator/action/utils/episode_filters.py new file mode 100644 index 000000000000..be2b5abed21f --- /dev/null +++ b/examples/worldmodel/cosmos3/vendor/cosmos_framework/data/generator/action/utils/episode_filters.py @@ -0,0 +1,36 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: OpenMDW-1.1 + +"""Storage-independent episode keys for optional dataset filter metadata.""" + +from typing import Any + + +def normalize_episode_filter_keys(filters: dict[str, Any], *, source_prefix: str = "") -> dict[str, Any]: + """Preserve complete relative episode IDs; URI keys require an explicit dataset prefix.""" + normalized: dict[str, Any] = {} + for key, ranges in filters.items(): + episode = key + if "://" in key: + # Only composite URI keys carry the recording/trajectory delimiter. + # Relative episode IDs are opaque, including any '--' in their names. + recording = key.split("/recordings/MP4--", 1)[0].rstrip("/") + if not source_prefix: + raise ValueError( + "filter_dict_source_prefix is required for URI-based episode filter keys. " + "Set DROID_FILTER_SOURCE_PREFIX when using a DROID launcher, or pass " + "filter_dict_source_prefix to the dataset. Use the URI prefix before the complete episode ID." + ) + prefix = source_prefix.rstrip("/") + "/" + if not recording.startswith(prefix): + raise ValueError( + "Episode filter URI does not match filter_dict_source_prefix. " + "Check DROID_FILTER_SOURCE_PREFIX against the keys in your filter JSON." + ) + episode = recording.removeprefix(prefix) + if not episode: + raise ValueError("Filter entries must identify an episode") + if episode in normalized and normalized[episode] != ranges: + raise ValueError(f"Conflicting filter ranges for episode {episode!r}") + normalized[episode] = ranges + return normalized diff --git a/examples/worldmodel/cosmos3/vendor/cosmos_framework/data/generator/action/utils/pose_utils.py b/examples/worldmodel/cosmos3/vendor/cosmos_framework/data/generator/action/utils/pose_utils.py new file mode 100644 index 000000000000..85104152cd11 --- /dev/null +++ b/examples/worldmodel/cosmos3/vendor/cosmos_framework/data/generator/action/utils/pose_utils.py @@ -0,0 +1,914 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: OpenMDW-1.1 + +"""Rotation and pose utilities for action datasets. + +This module centralizes three related responsibilities used across the action +dataset stack: + +1. Converting rotations between the conventions used by the datasets and the + action model (`euler_xyz`, quaternion, axis-angle, rot6d, rot9d, matrix). +2. Building absolute homogeneous poses of shape ``(T, 4, 4)`` from per-frame + translation and rotation components. +3. Converting trajectories between absolute-pose form and the relative-pose + action vectors consumed by the datasets. + + The relative-pose action vectors always follow the shared layout + ``[translation(3), rotation(...)]``. The rotation block is encoded with the + requested rotation output convention, and `convert_rotation()` is the + canonical public entrypoint for representation conversion. +""" + +import math +from typing import Literal + +import numpy as np +import torch +from scipy.spatial.transform import Rotation as R + +PoseConvention = Literal[ + "absolute", + "backward_anchored", + "backward_framewise", + "backward_chunk_anchored_8f", + "backward_chunk_anchored_16f", +] +RotationConvention = Literal["matrix", "euler_xyz", "quat_xyzw", "quat_wxyz", "rot6d", "axisangle", "rot9d"] + +_POSE_CHUNK_LENGTHS: dict[PoseConvention, int] = { + "backward_chunk_anchored_8f": 8, + "backward_chunk_anchored_16f": 16, +} + + +def _to_numpy_float32(array: torch.Tensor | np.ndarray) -> np.ndarray: + """Convert an input array to a NumPy ``float32`` array. + + Args: + array: A torch tensor or NumPy array with arbitrary leading dimensions. + + Returns: + A NumPy array with dtype ``float32``. Torch tensors are moved to CPU + before conversion. NumPy inputs are converted with ``copy=False`` + semantics when possible. + + Raises: + ValueError: If a torch tensor with ``requires_grad=True`` is passed. + These utilities are non-differentiable; callers must explicitly + detach tensors before conversion. + """ + if isinstance(array, torch.Tensor): + if array.requires_grad: + raise ValueError( + "pose_utils conversion is non-differentiable; call `.detach()` " + "explicitly before passing tensors with requires_grad=True" + ) + return array.cpu().numpy().astype(np.float32, copy=False) + return np.asarray(array, dtype=np.float32) + + +def _normalize_rotation_matrices(rot_matrices: np.ndarray) -> np.ndarray: + """Project approximate matrices onto valid rotation matrices. + + This helper uses an SVD-based projection onto ``SO(3)``. It is mainly used + when decoding rotations from network-like representations such as rot6d or rot9d + where the input may not already be perfectly orthonormal. + + Args: + rot_matrices: Array of shape ``(..., 3, 3)`` containing one or more + approximate rotation matrices. + + Returns: + Array of shape ``(..., 3, 3)`` whose trailing matrices are proper + rotation matrices with determinant ``+1``. + + Raises: + ValueError: If the input does not have trailing shape ``(3, 3)``. + """ + matrices = np.asarray(rot_matrices, dtype=np.float32) + if matrices.ndim < 2 or matrices.shape[-2:] != (3, 3): + raise ValueError(f"Rotation matrices must have shape (..., 3, 3), got {matrices.shape}") + + original_shape = matrices.shape[:-2] + matrices_flat = matrices.reshape(-1, 3, 3) + + # Batched SVD projection to SO(3). + U, _, Vt = np.linalg.svd(matrices_flat) + normalized = U @ Vt + + # Ensure determinant is +1 (proper rotations, no reflections). + det = np.linalg.det(normalized) + reflection_mask = det < 0 + if np.any(reflection_mask): + U_reflect = U.copy() + U_reflect[reflection_mask, :, -1] *= -1 + normalized[reflection_mask] = U_reflect[reflection_mask] @ Vt[reflection_mask] + + return normalized.astype(np.float32, copy=False).reshape(*original_shape, 3, 3) + + +def convert_rotation( + rotation: torch.Tensor | np.ndarray, + input_format: RotationConvention, + output_format: RotationConvention, + normalize_matrix: bool = False, +) -> torch.Tensor | np.ndarray: + """Convert rotations between the conventions used by action datasets. + + The function first maps the input representation to rotation matrices and + then emits the requested output convention. It is the single conversion seam + used by the public pose helpers so that all code paths share the same + convention handling. + + Supported input conventions: + - ``matrix``: rotation matrices with shape ``(..., 3, 3)`` + - ``euler_xyz``: Euler xyz angles in radians with shape ``(..., 3)`` + - ``quat_xyzw``: quaternions in SciPy's xyzw order with shape ``(..., 4)`` + - ``quat_wxyz``: quaternions in wxyz order with shape ``(..., 4)`` + - ``rot6d``: column-based 6D representation with shape ``(..., 6)`` + - ``rot9d``: flattened rotation matrices with shape ``(..., 9)`` + - ``axisangle``: axis-angle vectors with shape ``(..., 3)`` + + Supported output conventions: + - ``matrix`` + - ``euler_xyz`` + - ``quat_xyzw`` + - ``quat_wxyz`` + - ``rot6d`` + - ``axisangle`` + - ``rot9d`` + + Args: + rotation: Input rotations in the representation specified by + ``input_format``. + input_format: Convention used by ``rotation``. + output_format: Convention to return. + normalize_matrix: Whether to project intermediate matrices to a valid + rotation before returning. This is most useful when decoding from + approximate ``rot6d``/``rot9d`` inputs or non-unit quaternions. + + Returns: + Rotations with the same leading shape as the input, expressed in the + requested output convention. Torch inputs return torch outputs on the + same device with the same dtype; NumPy inputs return NumPy arrays. + + Raises: + ValueError: If the input shape is incompatible with ``input_format`` or + if either format is unsupported. + """ + input_is_tensor = isinstance(rotation, torch.Tensor) + input_dtype = rotation.dtype if input_is_tensor else None + input_device = rotation.device if input_is_tensor else None + rotation_np = _to_numpy_float32(rotation) + + if input_format == "matrix": + if rotation_np.ndim < 2 or rotation_np.shape[-2:] != (3, 3): + raise ValueError(f"matrix rotation must have shape (..., 3, 3), got {rotation_np.shape}") + original_shape = rotation_np.shape[:-2] + matrices_flat = rotation_np.reshape(-1, 3, 3) + if normalize_matrix: + matrices_flat = _normalize_rotation_matrices(matrices_flat).reshape(-1, 3, 3) + elif input_format == "euler_xyz": + if rotation_np.ndim < 1 or rotation_np.shape[-1] != 3: + raise ValueError(f"{input_format} rotation must have shape (..., 3), got {rotation_np.shape}") + original_shape = rotation_np.shape[:-1] + matrices_flat = R.from_euler("xyz", rotation_np.reshape(-1, 3), degrees=False).as_matrix().astype(np.float32) + elif input_format in ("quat_xyzw", "quat_wxyz"): + if rotation_np.ndim < 1 or rotation_np.shape[-1] != 4: + raise ValueError(f"{input_format} rotation must have shape (..., 4), got {rotation_np.shape}") + original_shape = rotation_np.shape[:-1] + quaternions = rotation_np.reshape(-1, 4) + if input_format == "quat_wxyz": + quaternions = quaternions[:, [1, 2, 3, 0]] + norms = np.linalg.norm(quaternions, axis=-1) + if np.any(norms < 1e-8): + raise ValueError(f"Found zero-norm quaternion(s) (min norm={norms.min():.2e}).") + if normalize_matrix: + quaternions = quaternions / norms[:, None] + matrices_flat = R.from_quat(quaternions).as_matrix().astype(np.float32) + elif input_format == "rot6d": + if rotation_np.ndim < 1 or rotation_np.shape[-1] != 6: + raise ValueError(f"{input_format} rotation must have shape (..., 6), got {rotation_np.shape}") + original_shape = rotation_np.shape[:-1] + rot6d_flat = rotation_np.reshape(-1, 6) + col0 = rot6d_flat[:, :3] + col1 = rot6d_flat[:, 3:] + col2 = np.cross(col0, col1, axis=-1) + matrices_flat = np.stack((col0, col1, col2), axis=-1).astype(np.float32) + if normalize_matrix: + matrices_flat = _normalize_rotation_matrices(matrices_flat).reshape(-1, 3, 3) + elif input_format == "rot9d": + if rotation_np.ndim < 1 or rotation_np.shape[-1] != 9: + raise ValueError(f"rot9d rotation must have shape (..., 9), got {rotation_np.shape}") + original_shape = rotation_np.shape[:-1] + matrices_flat = rotation_np.reshape(-1, 3, 3) + if normalize_matrix: + matrices_flat = _normalize_rotation_matrices(matrices_flat).reshape(-1, 3, 3) + elif input_format == "axisangle": + if rotation_np.ndim < 1 or rotation_np.shape[-1] != 3: + raise ValueError(f"axisangle rotation must have shape (..., 3), got {rotation_np.shape}") + original_shape = rotation_np.shape[:-1] + matrices_flat = R.from_rotvec(rotation_np.reshape(-1, 3)).as_matrix().astype(np.float32) + else: + raise ValueError(f"Unsupported input_format: {input_format!r}") + + if output_format == "matrix": + converted = matrices_flat.reshape(*original_shape, 3, 3).astype(np.float32) + elif output_format == "rot9d": + converted = matrices_flat.reshape(-1, 9) + elif output_format == "rot6d": + converted = matrices_flat[:, :, :2].transpose(0, 2, 1).reshape(-1, 6) + elif output_format == "quat_xyzw": + converted = R.from_matrix(matrices_flat).as_quat().astype(np.float32) + elif output_format == "quat_wxyz": + converted = R.from_matrix(matrices_flat).as_quat().astype(np.float32) + converted = converted[:, [3, 0, 1, 2]] + elif output_format == "euler_xyz": + converted = R.from_matrix(matrices_flat).as_euler("xyz", degrees=False).astype(np.float32) + elif output_format == "axisangle": + converted = R.from_matrix(matrices_flat).as_rotvec().astype(np.float32) + else: + raise ValueError(f"Unsupported output_format: {output_format!r}") + + if output_format != "matrix": + converted = converted.reshape(*original_shape, converted.shape[-1]) + + if input_is_tensor: + return torch.from_numpy(np.ascontiguousarray(converted)).to(dtype=input_dtype, device=input_device) + return converted + + +# ----------------------------------------------------------------------------- +# Absolute pose construction +# ----------------------------------------------------------------------------- + + +def build_abs_pose_from_components( + xyz: torch.Tensor | np.ndarray, + rotation: torch.Tensor | np.ndarray, + rotation_input_format: Literal["euler_xyz", "quat_xyzw", "quat_wxyz", "axisangle"], + translation_scale: float | None = None, +) -> np.ndarray: + """Build absolute homogeneous poses from per-frame translation and rotation. + + This is the canonical helper for turning dataset-provided pose components + into a sequence of rigid transforms. Each output pose is a homogeneous + transform whose top-left ``3 x 3`` block stores rotation and whose last + column stores translation. + + Args: + xyz: Per-frame translations with shape ``(T, 3)``. + rotation: Per-frame rotations with shape ``(T, 3)`` for ``euler_xyz`` + and ``axisangle``, or ``(T, 4)`` for quaternion conventions. + rotation_input_format: Convention used by ``rotation``. Supported values + are ``euler_xyz``, ``quat_xyzw``, ``quat_wxyz``, and ``axisangle``. + translation_scale: Optional factor used to divide translations before + inserting them into the output poses. This is useful when upstream + data stores translations in a scaled unit. + + Returns: + Absolute poses with shape ``(T, 4, 4)`` and dtype ``float32``. + + Raises: + ValueError: If the translation and rotation arrays have incompatible + lengths or unsupported shapes, or if ``translation_scale`` is zero. + """ + xyz_np = _to_numpy_float32(xyz) + rotation_np = _to_numpy_float32(rotation) + + if xyz_np.ndim != 2 or xyz_np.shape[1] != 3: + raise ValueError(f"xyz must have shape (T, 3), got {xyz_np.shape}") + if rotation_np.ndim != 2: + raise ValueError(f"rotation must be 2D, got {rotation_np.shape}") + if rotation_np.shape[0] != xyz_np.shape[0]: + raise ValueError( + f"xyz and rotation must have the same length, got {xyz_np.shape[0]} and {rotation_np.shape[0]}" + ) + + rot_mats = np.asarray( + convert_rotation(rotation_np, input_format=rotation_input_format, output_format="matrix"), + dtype=np.float32, + ) + + if translation_scale is not None: + if translation_scale == 0: + raise ValueError("translation_scale must be non-zero") + xyz_np = xyz_np / float(translation_scale) + + poses_abs = np.eye(4, dtype=np.float32)[None].repeat(xyz_np.shape[0], axis=0) + poses_abs[:, :3, :3] = rot_mats.astype(np.float32) + poses_abs[:, :3, 3] = xyz_np + return poses_abs + + +# ----------------------------------------------------------------------------- +# Relative pose conversions +# ----------------------------------------------------------------------------- + + +def _delta_transform_to_pose_vector( + delta_T: np.ndarray, + rotation_output_format: RotationConvention, +) -> np.ndarray: + """Encode a relative transform as an action vector. + + The shared action-vector layout is always ``[translation(3), rotation(...)]``. + + Args: + delta_T: Relative transform of shape ``(4, 4)``. + rotation_output_format: Concrete convention used for the output rotation + block. + + Returns: + A ``float32`` action vector whose first three values are translation and + whose remaining values are the rotation in ``rotation_output_format``. + """ + delta_np = np.asarray(delta_T, dtype=np.float32) + if delta_np.shape != (4, 4): + raise ValueError(f"delta_T must have shape (4, 4), got {delta_np.shape}") + + translation = delta_np[:3, 3] + rotation = np.asarray( + convert_rotation(delta_np[:3, :3], input_format="matrix", output_format=rotation_output_format), + dtype=np.float32, + ) + return np.concatenate([translation, rotation]).astype(np.float32) + + +def _pose_vector_to_delta_transform( + pose_vector: np.ndarray, + rotation_input_format: RotationConvention, + translation_scale: float, + normalize_rotation: bool, + rotation_scale: float = 1.0, +) -> np.ndarray: + """Decode an action vector back into a relative homogeneous transform. + + This is the inverse of `_delta_transform_to_pose_vector()` when the same + rotation convention is used. Scale arguments are provided for callers that + need to decode model-space pose actions before action-normalizer + denormalization has been applied. + + Args: + pose_vector: Relative-pose action vector with layout + ``[translation(3), rotation(...)]``. + rotation_input_format: Concrete convention used by the rotation block. + translation_scale: Scalar used to undo translation scaling in the input + vector. + normalize_rotation: Whether to project the decoded rotation to a valid + matrix before assembling the transform. + rotation_scale: Scalar used to undo rotation scaling in the input vector. + + Returns: + A relative homogeneous transform with shape ``(4, 4)`` and dtype + ``float32``. + """ + pose_vector_np = np.asarray(pose_vector, dtype=np.float32) + rotation_block = pose_vector_np[3:] / rotation_scale + + rotation_matrix = np.asarray( + convert_rotation( + rotation_block, + input_format=rotation_input_format, + output_format="matrix", + normalize_matrix=normalize_rotation, + ), + dtype=np.float32, + ) + + delta_T = np.eye(4, dtype=np.float32) + delta_T[:3, 3] = pose_vector_np[:3] / translation_scale + delta_T[:3, :3] = rotation_matrix + return delta_T + + +def _get_relative_delta_transform( + poses_abs: np.ndarray, + inv_poses_abs: np.ndarray, + frame_idx: int, + pose_convention: PoseConvention, + anchor_index: int = 0, +) -> np.ndarray: + """Compute one relative transform from an absolute-pose trajectory. + + Args: + poses_abs: Absolute poses of shape ``(T, 4, 4)``. + inv_poses_abs: Precomputed inverses of ``poses_abs`` with the same shape. + frame_idx: Index of the output row to encode, in ``[0, T - 2]``. + pose_convention: Pose convention controlling which two poses + define the delta and whether it is framewise or anchored. + anchor_index: For ``backward_anchored`` only: index of the anchor pose. + Row ``frame_idx`` encodes the frame at ``frame_idx`` when it precedes + the anchor and ``frame_idx + 1`` otherwise, so the anchor frame itself + (which would encode as the identity) is skipped and the output stays + ``T - 1`` rows for any anchor. ``backward_framewise`` has no anchor and + requires the default ``0``. + + Returns: + The relative transform ``delta_T`` with shape ``(4, 4)`` for the + requested step and convention. + """ + if pose_convention == "backward_framewise": + if anchor_index != 0: + raise ValueError(f"anchor_index is only meaningful for backward_anchored, got {anchor_index} (framewise).") + return inv_poses_abs[frame_idx] @ poses_abs[frame_idx + 1] + if pose_convention == "backward_anchored": + # `anchor_index` generalises a fixed anchor of 0: at anchor_index=0 this reduces + # exactly to `inv_poses_abs[0] @ poses_abs[frame_idx + 1]`. The generalisation is + # what makes anchored encoding usable with action history, where the anchor is the + # current frame rather than the window start. + target = frame_idx if frame_idx < anchor_index else frame_idx + 1 + return inv_poses_abs[anchor_index] @ poses_abs[target] + if chunk_length := _POSE_CHUNK_LENGTHS.get(pose_convention): + # The chunk-anchored conventions deliberately do NOT use the anchor-skip above: their + # target is always >= anchor_idx + 1, so no identity row can arise and no skip is needed. + if anchor_index != 0: + raise ValueError(f"anchor_index is only meaningful for backward_anchored, got {anchor_index}.") + anchor_idx = (frame_idx // chunk_length) * chunk_length + return inv_poses_abs[anchor_idx] @ poses_abs[frame_idx + 1] + raise ValueError( + f"Unsupported pose_convention={pose_convention!r}. Expected one of: backward_framewise, " + f"backward_anchored, {', '.join(sorted(_POSE_CHUNK_LENGTHS))}." + ) + + +def _apply_relative_delta_transform( + current_pose: np.ndarray, + initial_pose: np.ndarray, + chunk_anchor: np.ndarray, + delta_T: np.ndarray, + pose_convention: PoseConvention, +) -> np.ndarray: + """Recover the next absolute pose from a decoded relative transform. + + Args: + current_pose: The current reconstructed pose for framewise modes. + initial_pose: The anchor pose used by anchored modes. + chunk_anchor: Fixed anchor for the current generated chunk. + delta_T: Relative transform for the current step. + pose_convention: Pose convention controlling how ``delta_T`` + should be composed back into an absolute pose. + + Returns: + The next absolute pose with shape ``(4, 4)``. + """ + if pose_convention == "backward_framewise": + return current_pose @ delta_T + if pose_convention == "backward_anchored": + return initial_pose @ delta_T + if pose_convention in _POSE_CHUNK_LENGTHS: + return chunk_anchor @ delta_T + raise ValueError(f"Unsupported pose_convention={pose_convention!r}") + + +def pose_abs_to_rel( + poses_abs: np.ndarray, + rotation_format: RotationConvention = "rot9d", + pose_convention: PoseConvention = "backward_framewise", + anchor_index: int = 0, +) -> np.ndarray: + """Convert an absolute-pose trajectory into relative-pose action vectors. + + Args: + poses_abs: Absolute poses with shape ``(T, 4, 4)``. These are typically + object-in-world or camera-to-world transforms. + rotation_format: Rotation convention used for the output rotation block. + Supported values are ``rot9d``, ``rot6d``, ``quat_xyzw``, and + ``euler_xyz``. + pose_convention: Pose convention: + - ``backward_framewise``: ``delta_T = T_i^{-1} @ T_{i+1}`` + - ``backward_anchored``: ``delta_T = T_a^{-1} @ T_k`` for every frame + ``k != a``, where ``a = anchor_index``. With ``anchor_index = 0`` + this is ``T_0^{-1} @ T_{i+1}``. + - ``backward_chunk_anchored_8f`` and ``backward_chunk_anchored_16f``: + each group of 8 or 16 targets is relative to the pose immediately + before the group. + anchor_index: For ``backward_anchored``: which frame anchors the window. + With action history the anchor is the *current* frame (index ``H``), + so the ``H`` rows before it point backward (where the EEF was, + relative to now) and the rows after it point forward (where it should + go). The anchor frame itself is skipped -- it would encode as the + identity -- keeping the output at ``T - 1`` rows. The default ``0`` + reproduces the history-free behavior (anchor at the window start, + which *is* the current frame when ``H = 0``). Must be ``0`` for + ``backward_framewise`` and for the chunk-anchored conventions. + + Returns: + An array of shape ``(T - 1, D)`` where ``D = 3 + rotation_dim``. + + Raises: + AssertionError: If fewer than two absolute poses are provided. + ValueError: If ``anchor_index`` is out of range for ``poses_abs``. + """ + num_frames = len(poses_abs) + assert num_frames > 1, "At least 2 frames are required to compute relative poses" + if not 0 <= anchor_index < num_frames: + raise ValueError(f"anchor_index={anchor_index} out of range for {num_frames} poses.") + + # Compute inverse poses + inv_poses_abs = np.linalg.inv(poses_abs) + + poses_rel = [] + # We produce num_frames - 1 relative poses + for i in range(num_frames - 1): + delta_T = _get_relative_delta_transform(poses_abs, inv_poses_abs, i, pose_convention, anchor_index) + poses_rel.append( + _delta_transform_to_pose_vector( + delta_T, + rotation_output_format=rotation_format, + ) + ) + + return np.stack(poses_rel).astype(np.float32) # [T-1,D] + + +def pose_rel_to_abs( + poses_rel: np.ndarray, + rotation_format: RotationConvention = "rot9d", + pose_convention: PoseConvention = "backward_framewise", + initial_pose: np.ndarray | None = None, + normalize_rotation: bool = True, + translation_scale: float = 1.0, + rotation_scale: float = 1.0, + anchor_index: int = 0, +) -> np.ndarray: + """Reconstruct an absolute-pose trajectory from relative-pose action vectors. + + Args: + poses_rel: Relative-pose action vectors with shape ``(T - 1, D)`` and + layout ``[translation(3), rotation(...)]``. + rotation_format: Convention used by the rotation block of ``poses_rel``. + pose_convention: Pose convention used when the vectors were + encoded. This must match the convention passed to `pose_abs_to_rel()`. + initial_pose: Absolute pose for the anchor frame (``backward_anchored``) + or the first frame (``backward_framewise``). If ``None``, the + identity transform is used. + normalize_rotation: Whether to project decoded rotations onto ``SO(3)`` + before composing them back into the trajectory. + translation_scale: Scalar used to undo translation scaling in + ``poses_rel``. Prefer denormalizing with the dataset action + normalizer before calling this function. + rotation_scale: Scalar used to undo rotation scaling in ``poses_rel``. + Prefer denormalizing with the dataset action normalizer before + calling this function. + anchor_index: Must match the value passed to :func:`pose_abs_to_rel`. + For ``backward_anchored``, ``initial_pose`` is placed at this frame + index in the output and every row decodes independently as + ``initial_pose @ delta`` (rows before the anchor are past frames, + rows after it future frames). Must be ``0`` for + ``backward_framewise``. + + Returns: + Absolute poses with shape ``(T, 4, 4)`` where ``T = len(poses_rel) + 1``. + """ + if initial_pose is None: + initial_pose = np.eye(4) + + num_poses_rel = poses_rel.shape[0] + if not 0 <= anchor_index <= num_poses_rel: + raise ValueError(f"anchor_index={anchor_index} out of range for {num_poses_rel} relative poses.") + if pose_convention != "backward_anchored" and anchor_index != 0: + raise ValueError(f"anchor_index is only meaningful for backward_anchored, got {anchor_index}.") + + if pose_convention == "backward_anchored" and anchor_index != 0: + # For a mid-window anchor, rows decode independently and the anchor pose must be + # inserted at its frame position to recover the original trajectory order. + poses_abs = [ + initial_pose + @ _pose_vector_to_delta_transform( + pose_rel, + rotation_input_format=rotation_format, + translation_scale=translation_scale, + normalize_rotation=normalize_rotation, + rotation_scale=rotation_scale, + ) + for pose_rel in poses_rel + ] + poses_abs.insert(anchor_index, initial_pose) + return np.stack(poses_abs) # [T,4,4] + + # Preserve the existing reconstruction path when the anchor remains at frame 0. + poses_abs = [initial_pose] + current_pose = initial_pose + chunk_anchor = initial_pose + chunk_length = _POSE_CHUNK_LENGTHS.get(pose_convention) + + for i in range(num_poses_rel): + if chunk_length is not None and i > 0 and i % chunk_length == 0: + chunk_anchor = current_pose + delta_T = _pose_vector_to_delta_transform( + poses_rel[i], + rotation_input_format=rotation_format, + translation_scale=translation_scale, + normalize_rotation=normalize_rotation, + rotation_scale=rotation_scale, + ) + next_pose = _apply_relative_delta_transform( + current_pose, + initial_pose, + chunk_anchor, + delta_T, + pose_convention, + ) + + poses_abs.append(next_pose) + current_pose = next_pose + + return np.stack(poses_abs) # [T,4,4] + + +def pose_anchored_to_framewise( + poses_rel: np.ndarray, + rotation_format: RotationConvention = "rot6d", +) -> np.ndarray: + """Re-encode anchored relative-pose rows as framewise one-step rows. + + Idleness (and any other per-step property) is a fact about the underlying + trajectory, not its encoding: anchored rows ``P_k = T_anchor^-1 @ T_k`` are + losslessly convertible to one-step deltas ``delta_k = P_{k-1}^-1 @ P_k`` + (with ``P_0 = I``, since the anchor precedes the first row). This enables + per-step checks written for the framewise convention -- e.g. + :func:`compute_idle_frames` -- to run on anchored action chunks. + + Only forward-of-anchor rows may be passed (the dataset's ``action`` chunk; + ``anchor_index`` 0 semantics). History rows would need their own ordering. + + Args: + poses_rel: ``(K, D)`` anchored rows, layout ``[translation(3), rotation]``. + rotation_format: Rotation convention of the rows' rotation block. + + Returns: + ``(K, D)`` framewise rows in the same layout and rotation convention. + """ + poses_abs = pose_rel_to_abs(poses_rel, rotation_format=rotation_format, pose_convention="backward_anchored") + return pose_abs_to_rel(poses_abs, rotation_format=rotation_format, pose_convention="backward_framewise") + + +# ----------------------------------------------------------------------------- +# Absolute 9D pose vectors: [pos(3), rot6d(6)] +# ----------------------------------------------------------------------------- + + +def pose_abs_to_abs9d(poses_abs: np.ndarray) -> np.ndarray: + """Convert absolute poses to absolute ``[pos(3), rot6d(6)]`` vectors. + + Both halves of the name mean different things: the input is an absolute + ``(T, 4, 4)`` trajectory, and the output keeps absolute values rather than + per-step deltas. + + The leading frame is dropped so the output aligns step-for-step with + ``pose_abs_to_rel(..., rotation_format="rot6d")``, which spends one frame + forming its first delta. + + Args: + poses_abs: Absolute poses with shape ``(T, 4, 4)``. + + Returns: + Array of shape ``(T - 1, 9)`` and dtype ``float32``. + """ + pos = poses_abs[1:, :3, 3].astype(np.float32) # (T-1, 3) + rot6d = convert_rotation(poses_abs[1:, :3, :3], input_format="matrix", output_format="rot6d") # (T-1, 6) + return np.concatenate([pos, rot6d], axis=-1) + + +# ----------------------------------------------------------------------------- +# Idle-frame detection +# ----------------------------------------------------------------------------- + + +def _identity_rotation_vector(rotation_format: RotationConvention) -> np.ndarray: + """Return the identity-rotation vector for a given rotation convention. + + Used by :func:`compute_idle_frames` to test whether a rotation block is + close to "no rotation" in its current encoding. + """ + if rotation_format in ("matrix", "rot9d"): + return np.array([1, 0, 0, 0, 1, 0, 0, 0, 1], dtype=np.float32) + if rotation_format == "rot6d": + return np.array([1, 0, 0, 0, 1, 0], dtype=np.float32) + if rotation_format == "quat_xyzw": + return np.array([0, 0, 0, 1], dtype=np.float32) + if rotation_format == "quat_wxyz": + return np.array([1, 0, 0, 0], dtype=np.float32) + if rotation_format in ("euler_xyz", "axisangle"): + return np.array([0, 0, 0], dtype=np.float32) + raise ValueError(f"Unsupported rotation_format={rotation_format!r}") + + +def _rotation_angle_per_arm(rotations: np.ndarray, rotation_format: str) -> np.ndarray: + """Geodesic angle (rad) from identity for each arm at each frame. + + ``rotations`` has shape ``(T, n_arms, n_per_arm)``; the returned array has + shape ``(T, n_arms)``. The angle is rotation-format aware so a fixed + ``eps_r`` threshold has consistent geometric meaning across formats: + + - ``rot6d`` → reconstruct ``trace(R)`` in closed form from the two stored + columns ``a, b`` (already unit-orthogonal as they came from a valid + rotation matrix). The third column is ``a × b``, so + ``trace(R) = a[0] + b[1] + a[0]·b[1] - a[1]·b[0]``. + ``angle = arccos(clip((trace - 1) / 2, -1, 1))``. + - ``rot9d`` → reshape to ``(..., 3, 3)`` and use + ``trace(R) = R[0,0] + R[1,1] + R[2,2]``. + - ``quat_xyzw`` / ``quat_wxyz`` → ``angle = 2 · arccos(|q_w|)``; the + absolute value handles the double cover (``q`` and ``-q`` represent the + same rotation). + - ``axisangle`` → the magnitude of the axis-angle vector *is* the angle. + - ``euler_xyz`` → no closed-form angle; use ``‖euler‖`` as a conservative + upper bound (exact for single-axis rotations, an overestimate for + composed ones — fine for idle detection where small angles are the + regime of interest). + """ + if rotation_format == "rot6d": + a = rotations[..., :3] + b = rotations[..., 3:6] + trace = a[..., 0] + b[..., 1] + a[..., 0] * b[..., 1] - a[..., 1] * b[..., 0] + return np.arccos(np.clip((trace - 1.0) / 2.0, -1.0, 1.0)) + if rotation_format == "rot9d": + mat = rotations.reshape(*rotations.shape[:-1], 3, 3) + trace = mat[..., 0, 0] + mat[..., 1, 1] + mat[..., 2, 2] + return np.arccos(np.clip((trace - 1.0) / 2.0, -1.0, 1.0)) + if rotation_format in ("quat_xyzw", "quat_wxyz"): + qw = rotations[..., 3] if rotation_format == "quat_xyzw" else rotations[..., 0] + return 2.0 * np.arccos(np.clip(np.abs(qw), 0.0, 1.0)) + if rotation_format == "axisangle": + return np.linalg.norm(rotations, axis=-1) + if rotation_format == "euler_xyz": + # Exact for single-axis rotations, overestimate for composed ones — + # safe for idle thresholds since overestimation can only mark a frame + # as non-idle, never spuriously idle. + return np.linalg.norm(rotations, axis=-1) + raise ValueError(f"Unsupported rotation_format={rotation_format!r}") + + +def _consecutive_streaks(idle: np.ndarray, min_streak: int) -> np.ndarray: + """Zero out idle bits not belonging to a run of ``>= min_streak`` Trues. + + Pure-numpy two-pointer scan. ``min_streak <= 1`` is a no-op (returns the + input mask unchanged). + """ + if min_streak <= 1: + return idle + out = np.zeros_like(idle) + n = len(idle) + i = 0 + while i < n: + if not idle[i]: + i += 1 + continue + j = i + while j < n and idle[j]: + j += 1 + if j - i >= min_streak: + out[i:j] = True + i = j + return out + + +def compute_idle_frames( + action_raw: torch.Tensor | np.ndarray, + spec: "ActionSpec", # noqa: F821 — forward ref, real import is in action_spec.py + *, + eps_t: float = 1e-3, + eps_r: float = math.radians(5.0), + eps_g: float = 1e-2, + joint_threshold: float = 5e-4, + min_streak: int = 3, +) -> int: + """Count idle frames in a raw (un-normalized) action chunk. + + Idle detection runs per-DimType (driven by ``spec.types``); a frame is + *raw-idle* iff every relevant type group is idle on that frame, and + counts toward the final tally only if it belongs to a run of at least + ``min_streak`` consecutive raw-idle frames. The streak filter rejects + isolated low-motion frames (instantaneous slowdowns) which carry weak + physical meaning and add noise to the IdleFrames training signal. + + DimType branches: + + - ``POS`` → combined ``‖action[pos_idx]‖`` (L2 across all POS dims) + < ``eps_t``. For single-arm specs (3 dims) this is the standard ``‖t‖`` + check; for multi-arm specs the combined norm is slightly stricter than + a per-arm check. + - ``ROT`` → per-arm geodesic rotation angle (rad) from identity + < ``eps_r``. The angle is computed in a rotation-format aware way (see + :func:`_rotation_angle_per_arm`) so the threshold has consistent + geometric meaning regardless of the encoding. + - ``GRIPPER`` → ``max |action[t] - action[t-1]| < eps_g``. ``np.diff`` + with ``prepend=action[0]`` makes step 0 ``|0|`` (treated as "no change"); + with the streak filter this can no longer create a spurious single-frame + idle event. + - ``JOINT`` → same frame-diff scheme as gripper with + ``joint_threshold`` (rad / step). + - ``RESERVED`` → ignored. + + Defaults (in the units of the un-normalized action): + + - ``eps_t = 1e-3`` → 1 mm per-frame translation + - ``eps_r = 5°`` → 5° per-frame rotation (geodesic angle) + - ``eps_g = 1e-2`` → 1 % gripper command change + - ``joint_threshold = 5e-4`` → ~0.03° / step joint angle change + - ``min_streak = 3`` → require a run of >= 3 consecutive idle frames + + The input must be **un-normalized** so the identity transform sits at + known coordinates (translation ≈ 0, rotation ≈ identity). The action + vector is also assumed to be encoded in a per-step / framewise convention + (e.g. ``backward_framewise``); anchored conventions (``backward_anchored``) + accumulate over the chunk and would silently break the POS/ROT idle + checks. Callers (e.g. the LeRobot base class) gate on pose convention + before calling this function. + """ + if isinstance(action_raw, torch.Tensor): + action = action_raw.detach().cpu().numpy().astype(np.float32, copy=False) + else: + action = np.asarray(action_raw, dtype=np.float32) + + if action.ndim != 2: + raise ValueError(f"action_raw must be 2-D (T, D); got shape {action.shape}") + num_frames, action_dim = action.shape + if num_frames == 0: + return 0 + if action_dim != len(spec.types): + raise ValueError(f"action_dim={action_dim} does not match spec.dim={len(spec.types)}") + + # Import locally to avoid a circular import at module load time + # (action_spec.py imports RotationConvention from this file). + from cosmos_framework.data.generator.action.utils.action_spec import DimType + + pos_idx = [i for i, t in enumerate(spec.types) if t == DimType.POS] + rot_idx = [i for i, t in enumerate(spec.types) if t == DimType.ROT] + grip_idx = [i for i, t in enumerate(spec.types) if t == DimType.GRIPPER] + joint_idx = [i for i, t in enumerate(spec.types) if t == DimType.JOINT] + + idle = np.ones(num_frames, dtype=bool) + + # POS: combined L2 norm across all translation dims. + if pos_idx: + idle &= np.linalg.norm(action[:, pos_idx], axis=1) < eps_t + + # ROT: per-arm geodesic angle (rad). + if rot_idx: + rot_id = _identity_rotation_vector(spec.rotation_format) + n_per_arm = rot_id.shape[0] + if len(rot_idx) % n_per_arm != 0: + raise ValueError( + f"ROT dims ({len(rot_idx)}) not a multiple of " + f"rotation_format={spec.rotation_format!r} dim ({n_per_arm})" + ) + rotations = action[:, rot_idx].reshape(num_frames, -1, n_per_arm) + angles = _rotation_angle_per_arm(rotations, spec.rotation_format) # (T, n_arms) + idle &= angles.max(axis=1) < eps_r + + # GRIPPER: max |Δgripper| across all gripper dims; step 0's diff is 0. + if grip_idx: + gripper = action[:, grip_idx] + diff = np.abs(np.diff(gripper, axis=0, prepend=gripper[:1])) + idle &= diff.max(axis=1) < eps_g + + # JOINT: same frame-diff scheme with joint_threshold. + if joint_idx: + joints = action[:, joint_idx] + diff = np.abs(np.diff(joints, axis=0, prepend=joints[:1])) + idle &= diff.max(axis=1) < joint_threshold + + if min_streak > 1: + idle = _consecutive_streaks(idle, min_streak) + + return int(idle.sum()) + + +def compute_framewise_idle_frames( + action_raw: torch.Tensor | np.ndarray, + spec: "ActionSpec", # noqa: F821 — forward ref, real import is in action_spec.py + *, + fps: float, + pose_convention: PoseConvention, + eps_t_per_sec: float = 5e-3, + eps_r_per_sec: float = math.radians(1.5), + eps_g: float = 1e-2, + joint_threshold_per_sec: float = 5e-3, + min_streak: int = 3, + eps_t: float | None = None, + eps_r: float | None = None, + joint_threshold: float | None = None, +) -> torch.Tensor | None: # [] or None + """Apply the shared framewise idle policy and return a trainer scalar. + + Velocity thresholds are converted from per-second to per-frame units using + ``fps`` unless an explicit per-frame override is provided. Anchored and + absolute pose conventions do not have framewise idle semantics. + """ + if pose_convention != "backward_framewise": + return None + if fps <= 0: + raise ValueError(f"fps must be positive, got {fps}.") + + idle_frame_count = compute_idle_frames( + action_raw, + spec, + eps_t=eps_t if eps_t is not None else eps_t_per_sec / fps, + eps_r=eps_r if eps_r is not None else eps_r_per_sec / fps, + eps_g=eps_g, + joint_threshold=(joint_threshold if joint_threshold is not None else joint_threshold_per_sec / fps), + min_streak=min_streak, + ) + return torch.tensor(idle_frame_count, dtype=torch.long) # [] diff --git a/examples/worldmodel/cosmos3/vendor/cosmos_framework/data/generator/action/utils/viewpoint_utils.py b/examples/worldmodel/cosmos3/vendor/cosmos_framework/data/generator/action/utils/viewpoint_utils.py new file mode 100644 index 000000000000..f47d329eaef0 --- /dev/null +++ b/examples/worldmodel/cosmos3/vendor/cosmos_framework/data/generator/action/utils/viewpoint_utils.py @@ -0,0 +1,159 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: OpenMDW-1.1 + +"""Viewpoint type definitions and caption augmentor for Action datasets. + +Provides a ``Viewpoint`` type alias for camera perspective labels and a +``ViewpointTextInfo`` augmentor that appends a human-readable viewpoint +description to the caption string. +""" + +from __future__ import annotations + +from typing import Literal + +import torch +import torch.nn.functional as F + +from cosmos_framework.data.imaginaire.webdataset.augmentors.augmentor import Augmentor +from cosmos_framework.utils import log +from cosmos_framework.data.generator.action.action_caption_attribute_adapter import ViewCompositionMappings + +Viewpoint = Literal[ + "ego_view", + "third_person_view", + "wrist_view", + "concat_view", + "separate_camera_views", + "top_down_2d_view", + "video_game_view", +] + +DEFAULT_VIEWPOINT_TEMPLATES: dict[str, str] = { + "ego_view": "This video is captured from a first-person perspective looking at the scene.", + "third_person_view": "This video is captured from a third-person perspective looking towards the agent from the front.", + "wrist_view": "This video is captured from a wrist-mounted camera.", + "concat_view": "This video contains concatenated views from multiple camera perspectives.", + "separate_camera_views": ViewCompositionMappings["separate_camera_views"], + "top_down_2d_view": "This video is captured from a top-down view of a flat 2D scene.", + "video_game_view": "This video is captured from an in-game camera perspective in a virtual environment.", +} + + +def _resize_to_width(view: torch.Tensor, target_w: int) -> torch.Tensor: # [T,C,H,W] + """Aspect-preserving resize to ``target_w``; height follows from the view's own aspect ratio.""" + h, w = view.shape[-2:] + target_h = round(h * target_w / w) + if target_w != w or target_h != h: + view = F.interpolate(view, size=(target_h, target_w), mode="bilinear", align_corners=False) + return view + + +def compose_multiview( + primary: torch.Tensor, # [T,C,H,W] + left: torch.Tensor, # [T,C,H_l,W_l] + right: torch.Tensor, # [T,C,H_r,W_r] +) -> torch.Tensor: # [T,C,H+H_bottom,W] + """Place the primary view above two side views resized to match its width. + + ``left``/``right`` are resized aspect-preserving so their width matches + half the primary view's width; the resulting height follows from each + view's own native aspect ratio rather than being forced to match the + primary's, so a view whose native aspect ratio differs from the + primary's is not distorted. + """ + _, width = primary.shape[-2:] + half_width = width // 2 + left = _resize_to_width(left, half_width) # [T,C,H_l',W/2] + right = _resize_to_width(right, half_width) # [T,C,H_r',W/2] + bottom = torch.cat([left, right], dim=-1) # [T,C,H_bottom,W] + return torch.cat([primary, bottom], dim=-2) # [T,C,H+H_bottom,W] + + +class ViewpointTextInfo(Augmentor): + """Augmentor that appends viewpoint type description to captions. + + Reads a viewpoint label from ``data_dict[viewpoint_key]`` and appends + the corresponding template sentence to the caption. Designed to run + after the raw ``ai_caption`` is set but before duration/FPS metadata + is appended. + + Args: + input_keys: Input keys (kept for API compatibility). + output_keys: Output keys (kept for API compatibility). + args: Configuration arguments: + - caption_key (str): Key for caption in data_dict. Default: ``"ai_caption"`` + - viewpoint_key (str): Key for viewpoint label. Default: ``"viewpoint"`` + - templates (dict): Override mapping from viewpoint to sentence. + Default: :data:`DEFAULT_VIEWPOINT_TEMPLATES` + - separator (str): Separator between caption and metadata. Default: ``". "`` + - enabled (bool): Whether augmentation is enabled. Default: ``True`` + """ + + def __init__( + self, + input_keys: list | None = None, + output_keys: list | None = None, + args: dict | None = None, + ) -> None: + super().__init__(input_keys or [], output_keys or [], args) + + self.caption_key: str = args.get("caption_key", "ai_caption") if args else "ai_caption" + self.viewpoint_key: str = args.get("viewpoint_key", "viewpoint") if args else "viewpoint" + self.templates: dict[str, str] = ( + args.get("templates", DEFAULT_VIEWPOINT_TEMPLATES) if args else DEFAULT_VIEWPOINT_TEMPLATES + ) + self.default_separator: str = args.get("separator", ". ") if args else ". " + self.enabled: bool = args.get("enabled", True) if args else True + + def __call__(self, data_dict: dict) -> dict | None: + """Append viewpoint description to the caption. + + If the sample provides an ``"additional_view_description"`` key (a + free-form string describing the concatenated camera layout), it is + appended after the generic ``concat_view`` template. This allows each + dataset to supply its own description of which cameras are tiled and + how. + + Args: + data_dict: Sample dictionary containing caption and viewpoint. + + Returns: + The mutated *data_dict*, or the original unchanged if the + viewpoint key is missing or unrecognized. + """ + if not self.enabled: + return data_dict + + viewpoint = data_dict.get(self.viewpoint_key) + if viewpoint is None: + raise ValueError( + f"ViewpointTextInfo: missing key {self.viewpoint_key!r} in data_dict. " + f"All action datasets must provide a viewpoint label." + ) + + # Append dataset-specific concat_view details after the base template. + additional_view_description = data_dict.pop("additional_view_description", None) + template = self.templates.get(viewpoint) + + if template is None: + log.warning( + f"ViewpointTextInfo: unrecognized viewpoint {viewpoint!r}. " + f"Known viewpoints: {sorted(self.templates.keys())}. Skipping.", + rank0_only=False, + ) + return data_dict + + if additional_view_description: + separator = " " if template.endswith(".") else self.default_separator + template = template + separator + additional_view_description.rstrip() + + caption = data_dict.get(self.caption_key) + if not isinstance(caption, str) or caption == "": + return data_dict + + caption = caption.rstrip() + separator = " " if caption.endswith(".") else self.default_separator + data_dict[self.caption_key] = caption + separator + template + + return data_dict diff --git a/examples/worldmodel/cosmos3/vendor/cosmos_framework/data/imaginaire/webdataset/augmentors/augmentor.py b/examples/worldmodel/cosmos3/vendor/cosmos_framework/data/imaginaire/webdataset/augmentors/augmentor.py new file mode 100644 index 000000000000..57d20050f842 --- /dev/null +++ b/examples/worldmodel/cosmos3/vendor/cosmos_framework/data/imaginaire/webdataset/augmentors/augmentor.py @@ -0,0 +1,52 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: OpenMDW-1.1 + +from collections.abc import Iterable +from typing import Any, Generator, Optional + + +class Augmentor: + def __init__(self, input_keys: list, output_keys: Optional[list] = None, args: Optional[dict] = None) -> None: + r"""Base augmentor class + + Args: + input_keys (list): List of input keys + output_keys (list): List of output keys + args (dict): Arguments associated with the augmentation + """ + self.input_keys = input_keys + self.output_keys = output_keys + self.args = args + + def __call__(self, *args: Any, **kwds: Any) -> Any: + raise ValueError("Augmentor not implemented") + + +class IterableAugmentor: + def __init__(self, input_keys: list, output_keys: Optional[list] = None, args: Optional[dict] = None) -> None: + r"""Base augmentor class + + Args: + input_keys (list): List of input keys + output_keys (list): List of output keys + args (dict): Arguments associated with the augmentation + """ + self.input_keys = input_keys + self.output_keys = output_keys + self.args = args + self.is_generator = True + + def __call__(self, data: Iterable) -> Generator: + r"""Example usage: + + for data_dict in data: + # Do something to data_dict + data_dict["input"] = data_dict["raw_sequence"][:, :-1] + data_dict["target"] = data_dict["raw_sequence"][:, 1:] + # Skip sample if needed + if data_dict["input"].shape[1] < 64: + continue + # Construct a generator + yield data_dict + """ + raise ValueError("Augmentor not implemented") diff --git a/examples/worldmodel/cosmos3/vendor/cosmos_framework/utils/__init__.py b/examples/worldmodel/cosmos3/vendor/cosmos_framework/utils/__init__.py new file mode 100644 index 000000000000..28a81be6930a --- /dev/null +++ b/examples/worldmodel/cosmos3/vendor/cosmos_framework/utils/__init__.py @@ -0,0 +1,2 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: OpenMDW-1.1 diff --git a/examples/worldmodel/cosmos3/vendor/cosmos_framework/utils/log.py b/examples/worldmodel/cosmos3/vendor/cosmos_framework/utils/log.py new file mode 100644 index 000000000000..7736ef2c3457 --- /dev/null +++ b/examples/worldmodel/cosmos3/vendor/cosmos_framework/utils/log.py @@ -0,0 +1,144 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: OpenMDW-1.1 + +import atexit +import os +import sys +from typing import Any + +import torch.distributed as dist +from loguru._logger import Core, Logger + +RANK0_ONLY = True +LEVEL = os.environ.get("LOGURU_LEVEL", "INFO") +RANK = int(os.environ.get("RANK", "0")) + + +def make_new_logger(depth: int = 1) -> Logger: + return Logger( + core=Core(), + exception=None, + depth=depth, + record=False, + lazy=False, + colors=False, + raw=False, + capture=True, + patchers=[], + extra={}, + ) + + +logger = make_new_logger(depth=1) +atexit.register(logger.remove) + + +def _add_relative_path(record: dict[str, Any]) -> None: + try: + start = os.getcwd() + record["extra"]["relative_path"] = os.path.relpath(record["file"].path, start) + except OSError: + # CWD may have been removed (e.g. on some ranks in distributed jobs). + # Fall back to the absolute path so logging still works. + record["extra"]["relative_path"] = f":{record['file'].path}" + + +*options, _, extra = logger._options # type: ignore +logger._options = tuple([*options, [_add_relative_path], extra]) # type: ignore + + +def init_loguru_stdout() -> None: + logger.remove() + datetime_format = get_datetime_format() + machine_format = get_machine_format() + message_format = get_message_format() + logger.add( + sys.stdout, + level=LEVEL, + format=f"{datetime_format}{machine_format}{message_format}", + filter=_rank0_only_filter, + ) + + +def init_loguru_file(path: str) -> None: + datetime_format = get_datetime_format() + machine_format = get_machine_format() + message_format = get_message_format() + logger.add( + path, + encoding="utf8", + level=LEVEL, + format=f"{datetime_format}{machine_format}{message_format}", + rotation="100 MB", + filter=lambda result: _rank0_only_filter(result) or not RANK0_ONLY, + enqueue=True, + ) + + +def get_datetime_format() -> str: + return "[{time:MM-DD HH:mm:ss}|" + + +def get_machine_format() -> str: + node_id = os.environ.get("NGC_ARRAY_INDEX", "0") + num_nodes = int(os.environ.get("NGC_ARRAY_SIZE", "1")) + machine_format = "" + rank = 0 + if dist.is_available(): + if not RANK0_ONLY and dist.is_initialized(): + rank = dist.get_rank() + world_size = dist.get_world_size() + machine_format = ( + f"[Node{node_id:<3}/{num_nodes:<3}][RANK{rank:<5}/{world_size:<5}]" + "[{process.name:<8}]| " + ) + return machine_format + + +def get_message_format() -> str: + message_format = "{level}|{extra[relative_path]}:{line}:{function}] {message}" + return message_format + + +def _rank0_only_filter(record: Any) -> bool: + is_rank0 = record["extra"].get("rank0_only", True) + if RANK == 0 and is_rank0: + return True + if not is_rank0: + record["message"] = f"[RANK {RANK}] " + record["message"] + return not is_rank0 + + +def trace(message: str, rank0_only: bool = True) -> None: + logger.opt(depth=1).bind(rank0_only=rank0_only).trace(message) + + +def debug(message: str, rank0_only: bool = True) -> None: + logger.opt(depth=1).bind(rank0_only=rank0_only).debug(message) + + +def info(message: str, rank0_only: bool = True) -> None: + logger.opt(depth=1).bind(rank0_only=rank0_only).info(message) + + +def success(message: str, rank0_only: bool = True) -> None: + logger.opt(depth=1).bind(rank0_only=rank0_only).success(message) + + +def warning(message: str, rank0_only: bool = True) -> None: + logger.opt(depth=1).bind(rank0_only=rank0_only).warning(message) + + +def error(message: str, rank0_only: bool = True) -> None: + logger.opt(depth=1).bind(rank0_only=rank0_only).error(message) + + +def critical(message: str, rank0_only: bool = True) -> None: + logger.opt(depth=1).bind(rank0_only=rank0_only).critical(message) + + +def exception(message: str, rank0_only: bool = True) -> None: + logger.opt(depth=1).bind(rank0_only=rank0_only).exception(message) + + +# Execute at import time. +init_loguru_stdout()