From 69197f37b5b09a09cfd4bc353768185546d45cb5 Mon Sep 17 00:00:00 2001 From: Jiang Zhuosong Date: Fri, 13 Nov 2026 23:41:38 +0800 Subject: [PATCH 1/6] [example] vendor the upstream DROID data pipeline verbatim Copy the 34 files the DROID reader needs from nvidia/cosmos-framework 5950e7b4, byte for byte, under examples/worldmodel/cosmos3/vendor/, with the OpenMDW-1.1 text from the same commit. The example read DROID through its own reimplementation, and a reimplementation can drift from what it reimplements with nothing failing. One such drift was measured: the instruction was drawn with the global random.choice where the reference pins it, so the prompt could describe the same window differently -- 104 tokens against 149 in one measured case. Reading through upstream's own code makes a difference against upstream a difference in upstream. The files are not edited at all, so checking against a newer commit is one `diff -r`. That only holds while nothing reformats them, hence the pre-commit exclude: autoflake, isort and black would each rewrite code this repository does not own. The dependency set this needs is pinned and hand-picked rather than resolved, because resolving lerobot's own closure would install diffusers 0.35.2 over the 0.40.0 that provides the model and av 15.1.0 over the 18.1.0 the decoder links against. vendor/README.md records what that costs: the code is upstream's, the environment around it is not. --- .pre-commit-config.yaml | 7 + .../cosmos3/vendor/LICENSE-OpenMDW-1.1 | 52 + examples/worldmodel/cosmos3/vendor/README.md | 139 +++ .../vendor/cosmos_framework/__init__.py | 2 + .../vendor/cosmos_framework/data/__init__.py | 3 + .../data/generator/__init__.py | 3 + .../data/generator/action/__init__.py | 2 + .../action_caption_attribute_adapter.py | 267 +++++ .../generator/action/action_normalization.py | 60 + .../data/generator/action/agibot_fk.py | 398 +++++++ .../data/generator/action/agibot_spec.py | 129 ++ .../generator/action/datasets/__init__.py | 39 + .../agibotworld_beta_lerobot_dataset.py | 279 +++++ .../generator/action/datasets/base_dataset.py | 236 ++++ .../datasets/bridge_orig_lerobot_dataset.py | 154 +++ .../action/datasets/cosmos3_action_lerobot.py | 1060 +++++++++++++++++ .../action/datasets/droid_lerobot_dataset.py | 550 +++++++++ .../datasets/droid_lerobot_dataset_config.py | 100 ++ .../datasets/droid_merged_lerobot_dataset.py | 97 ++ .../datasets/fractal_lerobot_dataset.py | 190 +++ .../human_hand_pose_lerobot_dataset.py | 219 ++++ .../action/datasets/libero_lerobot_dataset.py | 330 +++++ .../datasets/robocasa_lerobot_dataset.py | 573 +++++++++ .../datasets/robomind_franka_dataset.py | 278 +++++ .../action/datasets/robomind_ur_dataset.py | 205 ++++ .../action/datasets/umi_lerobot_dataset.py | 177 +++ .../generator/action/libero_pose_utils.py | 69 ++ .../data/generator/action/utils/__init__.py | 2 + .../action/utils/action_processing.py | 540 +++++++++ .../generator/action/utils/action_spec.py | 235 ++++ .../generator/action/utils/domain_utils.py | 130 ++ .../generator/action/utils/episode_filters.py | 36 + .../data/generator/action/utils/pose_utils.py | 914 ++++++++++++++ .../generator/action/utils/viewpoint_utils.py | 159 +++ .../webdataset/augmentors/augmentor.py | 52 + .../vendor/cosmos_framework/utils/__init__.py | 2 + .../vendor/cosmos_framework/utils/log.py | 144 +++ 37 files changed, 7832 insertions(+) create mode 100644 examples/worldmodel/cosmos3/vendor/LICENSE-OpenMDW-1.1 create mode 100644 examples/worldmodel/cosmos3/vendor/README.md create mode 100644 examples/worldmodel/cosmos3/vendor/cosmos_framework/__init__.py create mode 100644 examples/worldmodel/cosmos3/vendor/cosmos_framework/data/__init__.py create mode 100644 examples/worldmodel/cosmos3/vendor/cosmos_framework/data/generator/__init__.py create mode 100644 examples/worldmodel/cosmos3/vendor/cosmos_framework/data/generator/action/__init__.py create mode 100644 examples/worldmodel/cosmos3/vendor/cosmos_framework/data/generator/action/action_caption_attribute_adapter.py create mode 100644 examples/worldmodel/cosmos3/vendor/cosmos_framework/data/generator/action/action_normalization.py create mode 100644 examples/worldmodel/cosmos3/vendor/cosmos_framework/data/generator/action/agibot_fk.py create mode 100644 examples/worldmodel/cosmos3/vendor/cosmos_framework/data/generator/action/agibot_spec.py create mode 100644 examples/worldmodel/cosmos3/vendor/cosmos_framework/data/generator/action/datasets/__init__.py create mode 100644 examples/worldmodel/cosmos3/vendor/cosmos_framework/data/generator/action/datasets/agibotworld_beta_lerobot_dataset.py create mode 100644 examples/worldmodel/cosmos3/vendor/cosmos_framework/data/generator/action/datasets/base_dataset.py create mode 100644 examples/worldmodel/cosmos3/vendor/cosmos_framework/data/generator/action/datasets/bridge_orig_lerobot_dataset.py create mode 100644 examples/worldmodel/cosmos3/vendor/cosmos_framework/data/generator/action/datasets/cosmos3_action_lerobot.py create mode 100644 examples/worldmodel/cosmos3/vendor/cosmos_framework/data/generator/action/datasets/droid_lerobot_dataset.py create mode 100644 examples/worldmodel/cosmos3/vendor/cosmos_framework/data/generator/action/datasets/droid_lerobot_dataset_config.py create mode 100644 examples/worldmodel/cosmos3/vendor/cosmos_framework/data/generator/action/datasets/droid_merged_lerobot_dataset.py create mode 100644 examples/worldmodel/cosmos3/vendor/cosmos_framework/data/generator/action/datasets/fractal_lerobot_dataset.py create mode 100644 examples/worldmodel/cosmos3/vendor/cosmos_framework/data/generator/action/datasets/human_hand_pose_lerobot_dataset.py create mode 100644 examples/worldmodel/cosmos3/vendor/cosmos_framework/data/generator/action/datasets/libero_lerobot_dataset.py create mode 100644 examples/worldmodel/cosmos3/vendor/cosmos_framework/data/generator/action/datasets/robocasa_lerobot_dataset.py create mode 100644 examples/worldmodel/cosmos3/vendor/cosmos_framework/data/generator/action/datasets/robomind_franka_dataset.py create mode 100644 examples/worldmodel/cosmos3/vendor/cosmos_framework/data/generator/action/datasets/robomind_ur_dataset.py create mode 100644 examples/worldmodel/cosmos3/vendor/cosmos_framework/data/generator/action/datasets/umi_lerobot_dataset.py create mode 100644 examples/worldmodel/cosmos3/vendor/cosmos_framework/data/generator/action/libero_pose_utils.py create mode 100644 examples/worldmodel/cosmos3/vendor/cosmos_framework/data/generator/action/utils/__init__.py create mode 100644 examples/worldmodel/cosmos3/vendor/cosmos_framework/data/generator/action/utils/action_processing.py create mode 100644 examples/worldmodel/cosmos3/vendor/cosmos_framework/data/generator/action/utils/action_spec.py create mode 100644 examples/worldmodel/cosmos3/vendor/cosmos_framework/data/generator/action/utils/domain_utils.py create mode 100644 examples/worldmodel/cosmos3/vendor/cosmos_framework/data/generator/action/utils/episode_filters.py create mode 100644 examples/worldmodel/cosmos3/vendor/cosmos_framework/data/generator/action/utils/pose_utils.py create mode 100644 examples/worldmodel/cosmos3/vendor/cosmos_framework/data/generator/action/utils/viewpoint_utils.py create mode 100644 examples/worldmodel/cosmos3/vendor/cosmos_framework/data/imaginaire/webdataset/augmentors/augmentor.py create mode 100644 examples/worldmodel/cosmos3/vendor/cosmos_framework/utils/__init__.py create mode 100644 examples/worldmodel/cosmos3/vendor/cosmos_framework/utils/log.py 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/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() From 74d868376c826876ee2eba89b9af21201c22cfb4 Mon Sep 17 00:00:00 2001 From: Jiang Zhuosong Date: Fri, 13 Nov 2026 23:41:38 +0800 Subject: [PATCH 2/6] [example] adapt upstream's DROID reader to this example's sample contract The reader returns the composite quantised to uint8, which is exactly the value droid_dataset reaches before padding, so the seam sits there: this module pads to the canvas, normalises to [-1, 1] and takes the state row back off the action. Each step undoes something the reader did rather than re-deriving it. Measured on one pinned window, with the instruction pinned on both sides and augmentation off: all 34,214,400 composite values identical, maximum absolute difference 0.000 levels, and the state and action rows identical. Through this module all thirteen keys match the example's own reader except sample_id, whose difference is explained where it is built -- upstream's index names episodes by integer, so replay manifests do not carry across the two readers. The scaling half of resize_and_reflect_pad never runs for DROID: the composite is always 540x640, the 4:3 canvas is 544x736, and the scale factor is capped at 1.0, so the call only ever pads. --- .../worldmodel/cosmos3/upstream_dataset.py | 170 ++++++++++++++++++ 1 file changed, 170 insertions(+) create mode 100644 examples/worldmodel/cosmos3/upstream_dataset.py diff --git a/examples/worldmodel/cosmos3/upstream_dataset.py b/examples/worldmodel/cosmos3/upstream_dataset.py new file mode 100644 index 000000000000..9c6db4702160 --- /dev/null +++ b/examples/worldmodel/cosmos3/upstream_dataset.py @@ -0,0 +1,170 @@ +"""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 + +from pathlib import Path +from typing import Any + +import torch +from droid_dataset import CONCAT_VIEW_DESCRIPTION, closest_canvas, resize_and_reflect_pad +from torch.utils.data import Dataset + +# 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, + 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, + ) + 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._starts = self._episode_starts() + + def _episode_starts(self) -> dict[int, int]: + """The reader's own (episode -> first global index) table, used only to report where a sample came from.""" + out, prev = {}, 0 + for (_ds, _ss, _vl, eid), end in zip(self._inner._episode_records, self._inner._episode_cum_ends): + out[int(eid)] = prev + prev = end + return out + + def _locate(self, index: int) -> tuple[int, int]: + episode, start = -1, -1 + for eid, first in self._starts.items(): + if first <= index and (start < 0 or first > start): + episode, start = eid, first + return episode, index - start if start >= 0 else -1 + + 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. + + **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. + """ + episode, ordinal = self._locate(index) + parts = [ + self.dataset_revision, + self._root_role, + f"episode_{episode}", + str(ordinal), + 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 + + episode, ordinal = 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, + } From 4b257ff630fd762eac29bd5897a5508d04b44eb6 Mon Sep 17 00:00:00 2001 From: Jiang Zhuosong Date: Fri, 13 Nov 2026 23:41:38 +0800 Subject: [PATCH 3/6] [example] choose the DROID reader with --dataset, upstream by default Put vendor/ on sys.path from the entry point only, and append rather than insert, so a host that also has cosmos-framework installed keeps loading the installed one -- the comparison machine does, and the reference side must never silently run this copy. Which one loaded is printed rather than assumed. The example's own reader stays reachable as --dataset example. It is not a fallback: its output was measured identical to upstream's, so it is a known-equivalent reference for comparing the two again after either side changes. Two steps on one card through the vendored reader: 16338993 window(s), samples=1 AR=168 DM=3093, loss 24.177132 then 11.412529. --- examples/worldmodel/cosmos3/train_sft.py | 45 +++++++++++++++++++++++- 1 file changed, 44 insertions(+), 1 deletion(-) diff --git a/examples/worldmodel/cosmos3/train_sft.py b/examples/worldmodel/cosmos3/train_sft.py index 2e1b87d1caee..578889fbc9d3 100644 --- a/examples/worldmodel/cosmos3/train_sft.py +++ b/examples/worldmodel/cosmos3/train_sft.py @@ -28,8 +28,32 @@ import argparse import importlib import os +import sys 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 from droid_dataset import DroidActionPolicyDataset, collate from packing import compute_loss, pack_batch @@ -91,6 +115,17 @@ def parse_args(): parser.add_argument("--save-every", type=int, default=1000) 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("--num-workers", type=int, default=4) parser.add_argument( "--max-episodes", type=int, help="Index only the first N episodes; the full index holds millions of windows." @@ -352,7 +387,14 @@ def main(): colossalai.launch_from_torch(backend=get_accelerator().communication_backend) coord = DistCoordinator() - dataset = DroidActionPolicyDataset( + # Both readers take the same arguments and return the same sample contract; see upstream_dataset.py for + # where the seam sits and what was measured across it. + reader = DroidActionPolicyDataset + if args.dataset == "upstream": + from upstream_dataset import UpstreamDroidDataset + + reader = UpstreamDroidDataset + dataset = reader( args.data, chunk_length=args.chunk_length, fps=args.fps, @@ -364,6 +406,7 @@ def main(): ) collate_fn = collate coord.print_on_master(f"[data] {len(dataset)} window(s)") + coord.print_on_master(f"[data] reader {args.dataset!r}, pipeline from {report_pipeline_source()}") model = build_model( args.checkpoint, From f90b8a2626ce4ba8ac5cbc6f8ed76e86ef9f15c6 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Fri, 13 Nov 2026 16:26:40 +0000 Subject: [PATCH 4/6] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- examples/worldmodel/cosmos3/train_sft.py | 2 +- examples/worldmodel/cosmos3/upstream_dataset.py | 7 ++----- 2 files changed, 3 insertions(+), 6 deletions(-) diff --git a/examples/worldmodel/cosmos3/train_sft.py b/examples/worldmodel/cosmos3/train_sft.py index aba5e9c7e52d..3c2f2803da4f 100644 --- a/examples/worldmodel/cosmos3/train_sft.py +++ b/examples/worldmodel/cosmos3/train_sft.py @@ -58,7 +58,6 @@ def report_pipeline_source() -> str: from colossalai.shardformer.modeling.cosmos3 import get_cosmos3_selective_checkpoint_stats from colossalai.shardformer.policies.cosmos3 import Cosmos3OmniPolicy from examples.worldmodel.cosmos3.droid_dataset import DroidActionPolicyDataset, collate -from examples.worldmodel.cosmos3.upstream_dataset import UpstreamDroidDataset from examples.worldmodel.cosmos3.packing import compute_loss, pack_batch from examples.worldmodel.cosmos3.runtime import ( AccumulationWindow, @@ -67,6 +66,7 @@ def report_pipeline_source() -> str: 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, diff --git a/examples/worldmodel/cosmos3/upstream_dataset.py b/examples/worldmodel/cosmos3/upstream_dataset.py index 2491c675b707..d8ef5755b9d3 100644 --- a/examples/worldmodel/cosmos3/upstream_dataset.py +++ b/examples/worldmodel/cosmos3/upstream_dataset.py @@ -35,13 +35,10 @@ from typing import Any import torch -from examples.worldmodel.cosmos3.droid_dataset import ( - CONCAT_VIEW_DESCRIPTION, - closest_canvas, - resize_and_reflect_pad, -) from torch.utils.data import Dataset +from examples.worldmodel.cosmos3.droid_dataset import CONCAT_VIEW_DESCRIPTION, 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. From 438102474e741ac3b5b03305d9fef0f6093143a9 Mon Sep 17 00:00:00 2001 From: Jiang Zhuosong Date: Sat, 14 Nov 2026 01:06:29 +0800 Subject: [PATCH 5/6] [example] forward the curated filter prefix to the vendored DROID reader The released keep_ranges_1_0_1.json is keyed by composite gs:// URIs. The vendored reader refuses such a key without the URI prefix in front of the episode id, rather than guess it and select the wrong episodes, so the adapters hard-coded empty prefix turned --keep-ranges into a ValueError on the default reader while the same flag worked on the other one. Resolve the prefix the way droid_dataset.py resolves it, from an explicit argument or DROID_FILTER_SOURCE_PREFIX, so one setting serves both readers. Measured against the release: both readers index 16,338,993 windows unfiltered and 12,157,927 with the filter, 74.41 percent, the share the upstream recipe documents for the released model. --- .../tests/test_upstream_reader_filter.py | 94 +++++++++++++++++++ .../worldmodel/cosmos3/upstream_dataset.py | 19 +++- 2 files changed, 111 insertions(+), 2 deletions(-) create mode 100644 examples/worldmodel/cosmos3/tests/test_upstream_reader_filter.py 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..236a1f2a0156 --- /dev/null +++ b/examples/worldmodel/cosmos3/tests/test_upstream_reader_filter.py @@ -0,0 +1,94 @@ +"""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 + self._episode_records = [("ds", "success", 40, 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/upstream_dataset.py b/examples/worldmodel/cosmos3/upstream_dataset.py index d8ef5755b9d3..e62383f2635c 100644 --- a/examples/worldmodel/cosmos3/upstream_dataset.py +++ b/examples/worldmodel/cosmos3/upstream_dataset.py @@ -31,13 +31,19 @@ from __future__ import annotations +import os 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, closest_canvas, resize_and_reflect_pad +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 @@ -64,6 +70,7 @@ def __init__( 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) @@ -92,7 +99,15 @@ def __init__( 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. - filter_dict_source_prefix="", + # 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 From 66e12a625dbfd4eb5789d7d45037180923a6ea7a Mon Sep 17 00:00:00 2001 From: Jiang Zhuosong Date: Sat, 14 Nov 2026 17:38:09 +0800 Subject: [PATCH 6/6] [example] identify a window by the span it comes from, not by its episode id The vendored reader's _episode_records holds one entry per span, not per episode, and episode_id identifies neither. With the curated window filter on, an episode contributes one span per kept range and its id repeats across them; the id also counts from zero again in every registered root, so on a mixed success-and-failure run the same id names two different episodes. The table this module built was keyed by the id alone, so it collapsed both cases, and the linear scan that read it then resolved against whichever start survived. This is quiet in the worst way. The samples stay correct, because the reader resolves them; only the episode and frame a sample reports come out wrong, and those are what a sample id, a replay pack and a held-out manifest are built from. On the real curated set -- 12,157,927 windows, 66,423 spans over 54,987 episodes, 8,299 of them split into more than one span -- the old index named the wrong episode for 12.0% of 500 random windows and the wrong frame for 19.8%. Resolve a window the way the reader resolves it: a bisect over the cumulative span ends, then that span's record. Count the frame from the start of the episode rather than the start of the span, so two kept ranges of one episode cannot hand two windows the same number, and put the dataset index in the identity, so a per-root episode id cannot collide across roots. An index outside the reader's range is refused instead of being reported against the last span. The bisect also replaces a scan over every episode that ran once per sample. Verified on the real curated index: 200,000 random windows give 200,000 distinct ids, and the two spans of episode 31450 report frames 0.. and 144.. rather than both starting at zero. --- .../tests/test_upstream_reader_filter.py | 3 +- .../tests/test_upstream_reader_index.py | 108 ++++++++++++++++++ .../worldmodel/cosmos3/upstream_dataset.py | 77 ++++++++++--- 3 files changed, 168 insertions(+), 20 deletions(-) create mode 100644 examples/worldmodel/cosmos3/tests/test_upstream_reader_index.py diff --git a/examples/worldmodel/cosmos3/tests/test_upstream_reader_filter.py b/examples/worldmodel/cosmos3/tests/test_upstream_reader_filter.py index 236a1f2a0156..82e80a28642b 100644 --- a/examples/worldmodel/cosmos3/tests/test_upstream_reader_filter.py +++ b/examples/worldmodel/cosmos3/tests/test_upstream_reader_filter.py @@ -32,7 +32,8 @@ class _Reader: def __init__(self, **kwargs): type(self).last = kwargs - self._episode_records = [("ds", "success", 40, 0)] + # (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): 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/upstream_dataset.py b/examples/worldmodel/cosmos3/upstream_dataset.py index e62383f2635c..61d27d26f929 100644 --- a/examples/worldmodel/cosmos3/upstream_dataset.py +++ b/examples/worldmodel/cosmos3/upstream_dataset.py @@ -32,6 +32,7 @@ from __future__ import annotations import os +from bisect import bisect_right from pathlib import Path from typing import Any @@ -115,22 +116,54 @@ def __init__( self.dataset_revision = dataset_revision self._root_role = "success" if success_only else "mixed" self._max_episodes = max_episodes - self._starts = self._episode_starts() - - def _episode_starts(self) -> dict[int, int]: - """The reader's own (episode -> first global index) table, used only to report where a sample came from.""" - out, prev = {}, 0 - for (_ds, _ss, _vl, eid), end in zip(self._inner._episode_records, self._inner._episode_cum_ends): - out[int(eid)] = prev - prev = end - return out - - def _locate(self, index: int) -> tuple[int, int]: - episode, start = -1, -1 - for eid, first in self._starts.items(): - if first <= index and (start < 0 or first > start): - episode, start = eid, first - return episode, index - start if start >= 0 else -1 + 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. @@ -139,6 +172,11 @@ def sample_id(self, index: int) -> str: 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 @@ -150,12 +188,13 @@ def sample_id(self, index: int) -> str: 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. """ - episode, ordinal = self._locate(index) + ds_index, episode, start = self._locate(index) parts = [ self.dataset_revision, self._root_role, + f"dataset_{ds_index}", f"episode_{episode}", - str(ordinal), + str(start), str(self.chunk_length), ] return ":".join(part for part in parts if part) @@ -178,7 +217,7 @@ def __getitem__(self, index: int) -> dict[str, Any]: 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 - episode, ordinal = self._locate(index) + _ds_index, episode, _start = self._locate(index) return { "video": padded, "action": action.float(),