diff --git a/cosmos_framework/checkpoint/dcp.py b/cosmos_framework/checkpoint/dcp.py index 43145c88..54d0d8e1 100644 --- a/cosmos_framework/checkpoint/dcp.py +++ b/cosmos_framework/checkpoint/dcp.py @@ -71,9 +71,9 @@ from cosmos_framework.checkpoint.base import AbstractCheckpointer from cosmos_framework.checkpoint.s3_filesystem import S3StorageReader, S3StorageWriter -from cosmos_framework.utils.config import CheckpointConfig, JobConfig from cosmos_framework.model._base import ImaginaireModel from cosmos_framework.utils import callback, distributed, log, misc +from cosmos_framework.utils.config import CheckpointConfig, JobConfig from cosmos_framework.utils.easy_io import easy_io from cosmos_framework.utils.generator.rand_state import get_rand_state_dict, set_rand_state_dict @@ -866,6 +866,16 @@ def load( raise ValueError( f"Unexpected keys (found in checkpoint but not in model): {results.unexpected_keys}" ) + # Warm start that skipped net_ema (e.g. loading an EMA-only HF export + # with no net_ema.* keys): the EMA shadow would otherwise keep its random + # build-time generation pathway (init_moe is skipped when a checkpoint is + # present). Seed net_ema from the freshly loaded net so the EMA starts equal + # to net ("EMA warm-starts from net") instead of from random weights. + if warm_start and any("net_ema" in skip_key for skip_key in keys_to_skip_loading): + ema_worker = getattr(model, "net_ema_worker", None) + if ema_worker is not None and getattr(model, "net_ema", None) is not None: + ema_worker.copy_to(src_model=model.net, tgt_model=model.net_ema) + log.info("Warm start: re-seeded net_ema from net (net_ema was skipped on load).") elif key == "optim": log.info("- Loading the optimizer...") diff --git a/cosmos_framework/configs/base/defaults/model_config.py b/cosmos_framework/configs/base/defaults/model_config.py index 5e6fd6b4..adde6bd9 100644 --- a/cosmos_framework/configs/base/defaults/model_config.py +++ b/cosmos_framework/configs/base/defaults/model_config.py @@ -5,12 +5,12 @@ import attrs -from cosmos_framework.utils.lazy_config import LazyDict from cosmos_framework.configs.base.defaults.activation_checkpointing import ActivationCheckpointingConfig from cosmos_framework.configs.base.defaults.compile import CompileConfig from cosmos_framework.configs.base.defaults.ema import EMAConfig from cosmos_framework.configs.base.defaults.parallelism import ParallelismConfig from cosmos_framework.configs.base.defaults.reasoner import VLMConfig +from cosmos_framework.utils.lazy_config import LazyDict @attrs.define(slots=False) diff --git a/cosmos_framework/configs/base/experiment/action/posttrain_config/action_policy_droid_nano.py b/cosmos_framework/configs/base/experiment/action/posttrain_config/action_policy_droid_nano.py index 926afea9..6ce66fb9 100644 --- a/cosmos_framework/configs/base/experiment/action/posttrain_config/action_policy_droid_nano.py +++ b/cosmos_framework/configs/base/experiment/action/posttrain_config/action_policy_droid_nano.py @@ -21,15 +21,14 @@ from hydra.core.config_store import ConfigStore -from cosmos_framework.utils.lazy_config import LazyCall as L -from cosmos_framework.utils.lazy_config import LazyDict - from cosmos_framework.configs.base.experiment.sft.models.nano_model_config import NANO_MODEL_CONFIG +from cosmos_framework.data.generator.action.datasets.action_sft_dataset import get_action_droid_sft_dataset from cosmos_framework.data.generator.joint_dataloader import ( PackingDataLoader, RankPartitionedDataLoader, ) -from cosmos_framework.data.generator.action.datasets.action_sft_dataset import get_action_droid_sft_dataset +from cosmos_framework.utils.lazy_config import LazyCall as L +from cosmos_framework.utils.lazy_config import LazyDict cs = ConfigStore.instance() @@ -202,7 +201,9 @@ use_state=True, iterable_shuffle=True, # rank x worker episode-shuffle stream episode_shuffle_seed=42, - use_image_augmentation=True, # SR boost (random crop+rescale + color jitter) + # SR boost: random crop+rescale + ColorJitter, applied CPU-side in the + # DROIDLeRobotDataset image augmentor (matches i4's pipeline stage). + use_image_augmentation=True, # keep_ranges_1_0_1.json window filter (drops idle/non-task frames). Off by default; # set use_filter_dict=True + filter_dict_path to enable. use_filter_dict=False, @@ -213,6 +214,9 @@ max_action_dim="${model.config.max_action_dim}", cfg_dropout_rate=0.1, tokenizer_config="${model.config.vlm_config.tokenizer}", + # Match i4 GA (droid_lerobot_8b_ga / MR #9995): format the action + # prompt as JSON via ActionPromptJsonFormatter instead of plain text. + format_prompt_as_json=True, ), ), ), diff --git a/cosmos_framework/data/generator/action/datasets/action_sft_dataset.py b/cosmos_framework/data/generator/action/datasets/action_sft_dataset.py index 8a424bc9..6053259c 100644 --- a/cosmos_framework/data/generator/action/datasets/action_sft_dataset.py +++ b/cosmos_framework/data/generator/action/datasets/action_sft_dataset.py @@ -19,7 +19,9 @@ from torch.utils.data import Dataset, IterableDataset, get_worker_info -from cosmos_framework.data.generator.action.datasets.droid_lerobot_dataset import DROIDLeRobotDataset +from cosmos_framework.data.generator.action.datasets.droid_lerobot_dataset import ( + DROIDLeRobotDataset, +) from cosmos_framework.data.generator.action.datasets.libero_lerobot_dataset import LIBEROLeRobotDataset from cosmos_framework.data.generator.action.transforms import ActionTransformPipeline @@ -109,13 +111,17 @@ def get_action_droid_sft_dataset( append_duration_fps_timestamps: bool = True, append_resolution_info: bool = True, append_idle_frames: bool = False, + format_prompt_as_json: bool = False, iterable_shuffle: bool = False, episode_shuffle_seed: int = 42, + use_success_only: bool = True, ) -> Dataset: """Build the DROID action SFT dataset: ``action_space='joint_pos'`` (8D) + - ``use_state`` (raw/un-normalized), concat_view, chunk_length 32.""" - dataset = DROIDLeRobotDataset( - root=root, + ``use_state`` (raw/un-normalized), concat_view, chunk_length 32. + + Reads ``root`` (a merged/versioned DROID LeRobot root) as a single flat + dataset; ``use_success_only=True`` filters to the ``success/`` split.""" + shard_kwargs = dict( fps=fps, chunk_length=chunk_length, viewpoint=viewpoint, @@ -123,10 +129,12 @@ def get_action_droid_sft_dataset( mode=mode, use_state=use_state, action_normalization=action_normalization, - use_image_augmentation=use_image_augmentation, + use_image_augmentation=use_image_augmentation, # i4: bundles random-crop+resize+ColorJitter use_filter_dict=use_filter_dict, filter_dict_path=filter_dict_path, + use_success_only=use_success_only, ) + dataset: Dataset = DROIDLeRobotDataset(root=root, **shard_kwargs) transform = ActionTransformPipeline( tokenizer_config=tokenizer_config, cfg_dropout_rate=cfg_dropout_rate, @@ -135,6 +143,7 @@ def get_action_droid_sft_dataset( append_duration_fps_timestamps=append_duration_fps_timestamps, append_resolution_info=append_resolution_info, append_idle_frames=append_idle_frames, + format_prompt_as_json=format_prompt_as_json, ) sft = ActionSFTDataset(dataset, transform, resolution) if iterable_shuffle: diff --git a/cosmos_framework/data/generator/action/datasets/base_dataset.py b/cosmos_framework/data/generator/action/datasets/base_dataset.py index e3d2ec81..ff0792ca 100644 --- a/cosmos_framework/data/generator/action/datasets/base_dataset.py +++ b/cosmos_framework/data/generator/action/datasets/base_dataset.py @@ -74,18 +74,13 @@ def __init__( # LeRobot v2.x stores task text in a "task" column; v3.0 stores it as the # (unnamed) DataFrame index and keeps only "task_index" as a column. task_texts = tasks_df["task"] if "task" in tasks_df.columns else tasks_df.index - self._tasks = { - int(task_index): str(task) - for task, task_index in zip(task_texts, tasks_df["task_index"]) - } - self._rows = sorted( - ( - row - for path in sorted((self._root / "data").glob("chunk-*/file-*.parquet")) - for row in pq.read_table(path).to_pylist() - ), - key=lambda row: int(row["index"]), - ) + self._tasks = {int(task_index): str(task) for task, task_index in zip(task_texts, tasks_df["task_index"])} + # ``self._rows`` (the flat, index-sorted list of every frame dict) is built + # lazily on first access — see the ``_rows`` property. Materializing all + # ~18M frames as Python dicts plus a full sort costs ~13 min and tens of GB; + # subclasses that build their own compact index (e.g. DROIDLeRobotDataset) + # never touch it, so they must not pay for it at construction. + self._rows_cache: list[dict[str, Any]] | None = None @property def fps(self) -> float: @@ -140,8 +135,7 @@ def _stats_path(cls) -> Path: def load_action_stats(cls) -> dict[str, torch.Tensor]: """Return action normalization stats for this dataset as torch tensors.""" return { - key: torch.from_numpy(value).float() - for key, value in load_action_stats(str(cls._stats_path())).items() + key: torch.from_numpy(value).float() for key, value in load_action_stats(str(cls._stats_path())).items() } @abstractmethod @@ -218,5 +212,25 @@ def _build_result( **extras, } + @property + def _rows(self) -> list[dict[str, Any]]: + """Flat, index-sorted list of every frame dict, built lazily on first access. + + Only datasets that don't build their own compact index (bridge / agibot / + robomind) touch this; for them it materializes once and caches. Datasets with + a bespoke index (e.g. DROIDLeRobotDataset) never read it, so they skip the + ~13 min / tens-of-GB construction entirely. + """ + if self._rows_cache is None: + self._rows_cache = sorted( + ( + row + for path in sorted((self._root / "data").glob("chunk-*/file-*.parquet")) + for row in pq.read_table(path).to_pylist() + ), + key=lambda row: int(row["index"]), + ) + return self._rows_cache + def __len__(self) -> int: return max(0, (len(self._rows) - self._chunk_length + self._sample_stride - 1) // self._sample_stride) diff --git a/cosmos_framework/data/generator/action/datasets/cosmos3_action_lerobot.py b/cosmos_framework/data/generator/action/datasets/cosmos3_action_lerobot.py new file mode 100644 index 00000000..d74b2aee --- /dev/null +++ b/cosmos_framework/data/generator/action/datasets/cosmos3_action_lerobot.py @@ -0,0 +1,1048 @@ +# 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.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.action_spec import ( # noqa: F401 (re-export) + ActionSpec, + DimType, + Gripper, + Joint, + Pos, + Reserved, + Rot, + build_action_spec, +) +from cosmos_framework.data.generator.action.domain_utils import get_domain_id +from cosmos_framework.data.generator.action.pose_utils import compute_idle_frames +from cosmos_framework.data.generator.action.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, + 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: + self._action_normalizer = resolve_action_normalization( + action_normalization, self._load_norm_stats(action_normalization) + ) + 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(tianweis): 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", "policy")) + 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._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/cosmos_framework/data/generator/action/datasets/droid_lerobot_dataset.py b/cosmos_framework/data/generator/action/datasets/droid_lerobot_dataset.py index 7e432ea3..64e2e12b 100644 --- a/cosmos_framework/data/generator/action/datasets/droid_lerobot_dataset.py +++ b/cosmos_framework/data/generator/action/datasets/droid_lerobot_dataset.py @@ -1,316 +1,306 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: OpenMDW-1.1 -"""Minimal DROID LeRobot dataset for Cosmos Action v1.2 defaults.""" - -from __future__ import annotations - import json +import os import random -from pathlib import Path -from typing import Any, Literal +from typing import Any, cast import numpy as np -import pyarrow.parquet as pq import torch import torch.nn.functional as F -import torchvision.transforms as T - -from cosmos_framework.data.generator.action.action_spec import ActionSpec, Gripper, Joint, Pos, Rot, build_action_spec -from cosmos_framework.data.generator.action.datasets.base_dataset import ActionBaseDataset +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.pose_utils import ( + PoseConvention, build_abs_pose_from_components, + convert_rotation, pose_abs_to_rel, ) +from cosmos_framework.data.generator.action.viewpoint_utils import Viewpoint +from cosmos_framework.utils import log -PoseConvention = Literal["backward_framewise"] -Viewpoint = Literal["concat_view"] - -_IMAGE_FEATURES = { - "wrist": "observation.image.wrist_image_left", - "left": "observation.image.exterior_image_1_left", - "right": "observation.image.exterior_image_2_left", -} -_STATE_FEATURE = "observation.state.cartesian_position" -# joint_pos (8D = 7 arm joints + gripper) features, matching the internal -# DROIDLeRobotDataset(action_space="joint_pos", use_state=...). These are -# absolute joint commands/states (no normalization is applied for joint_pos, -# matching the internal canonical run which leaves action_normalization=None). -_JOINT_ACTION_FEATURE = "action.joint_position" # [7] commanded joints -_ACTION_GRIPPER_FEATURE = "action.gripper_position" # [1] commanded gripper -_JOINT_STATE_FEATURE = "observation.state.joint_positions" # [7] observed joints -_GRIPPER_STATE_FEATURE = "observation.state.gripper_position" # [1] observed gripper -# Columns whose parquet dtype is a list (need to_pylist -> stacked array). -_LIST_COLUMNS = {_STATE_FEATURE, _JOINT_ACTION_FEATURE, _JOINT_STATE_FEATURE} -_ACTION_SPACES = ("ee_pose", "joint_pos") - -# 90-degree clockwise rotation about the Z axis in the local frame. This matches -# the production DROID wrapper conversion from Franka panda_link8 to OpenCV. +_FILTER_DICT_PATH = "/scratch/fsw/portfolios/cosmos/projects/cosmos_base_training/users/haolia/workspace/droid_oss_inputs/keep_ranges_1_0_1.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]], + [ + [0.0, -1.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/droid_lerobot_stats.json" - - -class DROIDLeRobotDataset(ActionBaseDataset): - """DROID Action dataset. - Two action layouts: - * ``action_space="ee_pose"`` (default): 10D ``[pos_delta(3), rot6d_delta(6), - gripper(1)]``, quantile-normalized (the v1.2 midtrain default). - * ``action_space="joint_pos"``: 8D ``[joint(7), gripper(1)]`` absolute joint - commands, NOT normalized, with ``use_state=True`` prepending the initial - observed joint+gripper state → ``(chunk+1, 8)`` — matching the internal - ``Cosmos3-Nano-Policy-DROID`` post-training run. - Filter dictionaries, temporal-segment validation, and image augmentation from - the production wrapper are intentionally omitted. - """ +class DROIDLeRobotDataset(BaseActionLeRobotDataset): + """ """ def __init__( self, - root: str, + root: str = "/lustre/fsw/portfolios/cosmos/projects/cosmos_base_training/cosmos3_action_datasets/droid_plus_lerobot_640x360_20260412", fps: float = 15.0, chunk_length: int = 16, - mode: str = "joint", + split_seed: int = 42, + split_val_ratio: float = 0.03, + split: str = "train", + mode: str = "policy", pose_convention: PoseConvention = "backward_framewise", - tolerance_s: float = 2e-4, + action_normalization: ActionNormalization | None = None, + tolerance_s=2e-4, viewpoint: Viewpoint = "concat_view", - action_space: str = "ee_pose", + use_success_only: bool = False, + video_mode: str | None = None, # TODO (ychao): remove + action_space: str = "midtrain", # TODO (ychao): remove use_state: bool = False, - action_normalization: str | None = "quantile", - use_image_augmentation: 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, ) -> None: - if viewpoint != "concat_view": - raise NotImplementedError("This minimal DROID dataset only supports concat_view.") - if action_space not in _ACTION_SPACES: - raise NotImplementedError(f"action_space must be one of {_ACTION_SPACES}, got {action_space!r}.") - if use_state and action_space != "joint_pos": - raise NotImplementedError("use_state is only supported with action_space='joint_pos'.") - if use_filter_dict and not filter_dict_path: - raise ValueError("use_filter_dict=True requires filter_dict_path") - - # joint_pos uses raw joint values — disable normalization at the base level. + """ """ super().__init__( - root=root, - domain_name="droid_lerobot", 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, tolerance_s=tolerance_s, - viewpoint=viewpoint, - action_normalization=None if action_space == "joint_pos" else action_normalization, + 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 = bool(use_state) - # Per-sample image augmentation (random crop+rescale + color jitter), applied - # to all views with shared params (temporally + cross-view consistent). Lazy-built. - self._use_image_augmentation = bool(use_image_augmentation) - self._image_augmentor: T.Compose | None = None - # Keep-ranges window filter (internal use_filter_dict): restrict training windows - # to curated active segments, dropping idle/non-task frames. Off by default; the - # keep-ranges JSON is supplied via filter_dict_path (an internal data artifact). - self._use_filter_dict = bool(use_filter_dict) - self._filter_dict_path = filter_dict_path - - # Compact, lazy frame index. Materializing every frame as a Python dict - # (``sorted(... pq.read_table(path).to_pylist() ...)``) does not scale: - # the full DROID success shard is ~18M frames, which is tens of GB of - # dicts plus an 18M-element Python sort at construction, and each - # DataLoader worker faults in its own copy. Instead we read only the - # columns the sample builder needs into contiguous numpy arrays - # (~1 GB total) -- read-only after init, so worker forks share them - # copy-on-write. - if action_space == "joint_pos": - feature_cols = [_JOINT_ACTION_FEATURE, _ACTION_GRIPPER_FEATURE, _JOINT_STATE_FEATURE, _GRIPPER_STATE_FEATURE] + self._use_state = use_state + self._use_filter_dict = use_filter_dict + self._filter_dict_path = filter_dict_path or _FILTER_DICT_PATH + self._max_num_history_actions = max_num_history_actions + self._use_image_augmentation = use_image_augmentation + 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: - feature_cols = [_STATE_FEATURE, _ACTION_GRIPPER_FEATURE] - columns = ["index", "episode_index", "task_index", "timestamp", *feature_cols] - index_parts, episode_parts, task_parts, ts_parts = [], [], [], [] - feature_parts: dict[str, list] = {c: [] for c in feature_cols} - for path in sorted((self._root / "data").glob("chunk-*/file-*.parquet")): - table = pq.read_table(path, columns=columns) - 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()) - for c in feature_cols: - if c in _LIST_COLUMNS: - feature_parts[c].append(np.asarray(table[c].to_pylist(), dtype=np.float32)) - else: - feature_parts[c].append(np.asarray(table[c].to_numpy(), dtype=np.float32)) - 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] - # Per-feature arrays keyed by parquet column name (read-only after init). - self._feat = { - c: np.concatenate(feature_parts[c], axis=0).astype(np.float32)[order] for c in feature_cols + 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, } - - # Group frames into episodes and keep only within-episode chunk windows. - # The global frame index is ordered by episode in LeRobot v3, so episodes - # are contiguous blocks once sorted by ``index``. The previous code sliced - # the flat row list (``rows[idx : idx + chunk + 1]``) with no boundary - # guard, so ~one chunk of samples per episode silently mixed two episodes; - # restricting to in-episode windows yields ``total - n_episodes * chunk`` - # valid samples (matching the production dataset). - assert np.all(np.diff(self._row_episode) >= 0), "episode_index is not contiguous after sorting by frame index" - ep_vals, ep_starts, ep_counts = np.unique(self._row_episode, return_index=True, return_counts=True) - self._ep_vals = ep_vals.astype(np.int64) - self._ep_starts = ep_starts.astype(np.int64) - self._valid_cum = np.cumsum(np.maximum(0, ep_counts - self._chunk_length)).astype(np.int64) - - # Keep-ranges filter: build a per-segment index over only the kept windows. - # Mirrors internal _append_index_records (use_filter_dict): the filter dict maps a - # gs:// trajectory key -> list of [start, end] frame ranges; keep windows whose start - # is in [max(start,0), min(end-chunk, valid)). Episodes absent from the dict are dropped. - if self._use_filter_dict: - with open(self._filter_dict_path) as f: - filter_dict = json.load(f) - seg_ep_pos, seg_win_start, seg_len = [], [], [] - for pos in range(len(self._ep_vals)): - valid = int(max(0, ep_counts[pos] - self._chunk_length)) - if valid <= 0: - continue - ep_id = str(self._episodes[int(self._ep_vals[pos])]["episode_id"]) - key = ( - f"gs://xembodiment_data/r2d2/r2d2-data-full/{ep_id}/recordings/" - f"MP4--gs://xembodiment_data/r2d2/r2d2-data-full/{ep_id}/trajectory.h5" - ) - ranges = filter_dict.get(key) - if ranges is None: - continue - for s, e in ranges: - ws = max(int(s), 0) - we = min(int(e) - self._chunk_length, valid) - if we - ws > 0: - seg_ep_pos.append(pos) - seg_win_start.append(ws) - seg_len.append(we - ws) - self._seg_ep_pos = np.asarray(seg_ep_pos, dtype=np.int64) - self._seg_win_start = np.asarray(seg_win_start, dtype=np.int64) - self._seg_cum = np.cumsum(seg_len).astype(np.int64) if seg_len else np.zeros(0, dtype=np.int64) - - @property - def action_dim(self) -> int: - return 8 if self._action_space == "joint_pos" else 10 - - def _action_spec(self) -> ActionSpec: + 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": - return build_action_spec(Joint(n=7, label="joint"), Gripper()) - return build_action_spec(Pos(), Rot("rot6d"), Gripper()) + 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 - @classmethod - def _stats_path(cls) -> Path: - return _NORMALIZER_PATH - - def _window_rows(self, start: int, stop: int, episode_index: int) -> list[dict[str, Any]]: - """Reconstruct the per-frame dicts the sample builder consumes for the - half-open frame window ``[start, stop)`` from the compact column arrays. - ``start``/``stop`` are guaranteed to lie within a single episode.""" - return [ - { - "episode_index": episode_index, - "task_index": int(self._row_task[j]), - "timestamp": float(self._row_timestamp[j]), - **{c: self._feat[c][j] for c in self._feat}, - } - for j in range(start, stop) - ] - - def __getitem__(self, idx: int) -> dict[str, Any]: - mode = self._choose_mode() - idx = int(idx) - # Map the flat sample index to a within-episode frame window. if self._use_filter_dict: - seg = int(np.searchsorted(self._seg_cum, idx, side="right")) - base = int(self._seg_cum[seg - 1]) if seg > 0 else 0 - ep = int(self._seg_ep_pos[seg]) - start = int(self._ep_starts[ep]) + int(self._seg_win_start[seg]) + (idx - base) - else: - 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] - - observation_rows = self._window_rows(start, start + self._chunk_length + 1, episode_index) + with open(self._filter_dict_path) as f: + self._filter_dict = json.load(f) - video = self._load_concat_video(episode, observation_rows) - if self._action_space == "joint_pos": - raw_action = self._build_joint_action(observation_rows) - extras: dict[str, Any] = {} - else: - action_rows = observation_rows[: self._chunk_length] - raw_action, initial_pose = self._build_raw_action(observation_rows, action_rows) - extras = {"initial_pose": initial_pose} - task = self._tasks[int(observation_rows[0]["task_index"])] - ai_caption = random.choice(task.split(" | ")) + self._image_augmentor: T.Compose | None = None - return self._build_result( - mode=mode, - video=video, - action=raw_action, - ai_caption=ai_caption, - 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." - ), - **extras, + # Eager source registration. i4 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 ) - def _build_joint_action(self, observation_rows: list[dict[str, Any]]) -> torch.Tensor: - """8D joint-position action ``[joint(7), gripper(1)]`` over the chunk, matching - the internal ``action_space='joint_pos'``. The window is ``chunk+1`` frames: - ``row[0]`` is the initial observed state (prepended when ``use_state``), and - ``rows[1:]`` are the ``chunk`` commanded actions. Gripper is flipped (1 - g). - No normalization is applied (internal canonical run uses raw joint values).""" - action_rows = observation_rows[1:] - joints = np.asarray([r[_JOINT_ACTION_FEATURE] for r in action_rows], dtype=np.float32) # [chunk, 7] - gripper = np.asarray([r[_ACTION_GRIPPER_FEATURE] for r in action_rows], dtype=np.float32).reshape(-1, 1) - gripper = 1.0 - gripper - action = np.concatenate([joints, gripper], axis=-1) # [chunk, 8] - if self._use_state: - init = observation_rows[0] - init_joint = np.asarray(init[_JOINT_STATE_FEATURE], dtype=np.float32) # [7] - init_gripper = np.asarray([1.0 - float(init[_GRIPPER_STATE_FEATURE])], dtype=np.float32) # [1] - initial_state = np.concatenate([init_joint, init_gripper])[None, :] # [1, 8] - action = np.concatenate([initial_state, action], axis=0) # [chunk + 1, 8] - return torch.from_numpy(action).float() - - 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, + 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 = f"gs://xembodiment_data/r2d2/r2d2-data-full/{ep_id_str}/recordings/MP4--gs://xembodiment_data/r2d2/r2d2-data-full/{ep_id_str}/trajectory.h5" + 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" ) - for name, video_key in _IMAGE_FEATURES.items() - } - wrist = frames_by_view["wrist"] - left = frames_by_view["left"] - right = frames_by_view["right"] + 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. + + 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: - # Random crop+rescale (spatial jitter) + color jitter, BEFORE the concat. - # All three views are stacked so one sampled set of params is applied - # uniformly across every frame and view (temporally + cross-view consistent), - # while each __getitem__ resamples. Matches the internal DROID recipe. if self._image_augmentor is None: _, _, h, w = wrist.shape self._image_augmentor = T.Compose( @@ -326,46 +316,182 @@ def _load_concat_video( _, _, 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) - return torch.cat([wrist, bottom], dim=-2) - 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] @ _DROID_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_GRIPPER_FEATURE] for row in action_rows], dtype=np.float32 - ).reshape(-1, 1) - gripper = 1.0 - gripper - action = np.concatenate([poses_rel[-self._chunk_length :], gripper[-self._chunk_length :]], axis=-1) - return torch.from_numpy(action).float(), initial_pose - - def __len__(self) -> int: - if self._use_filter_dict: - return int(self._seg_cum[-1]) if self._seg_cum.size else 0 - return int(self._valid_cum[-1]) if self._valid_cum.size else 0 - - def get_shuffle_blocks(self) -> list[tuple[int, int]]: - """Per-episode (or per kept-segment, when ``use_filter_dict``) flat-index blocks - ``(start, length)``. ``ActionIterableShuffleDataset`` shuffles the ORDER of these - blocks and shards them disjointly across ranks, while keeping windows *within* a - block sequential -> decorrelates batches across ranks without random-access I/O - (preserves locality + copy-on-write memory sharing across workers).""" - cum = self._seg_cum if self._use_filter_dict else self._valid_cum - blocks: list[tuple[int, int]] = [] - prev = 0 - for c in np.asarray(cum).tolist(): - c = int(c) - if c > prev: - blocks.append((prev, c - prev)) - prev = c - return blocks + 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] + 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/cosmos_framework/data/generator/action/datasets/droid_lerobot_dataset_config.py b/cosmos_framework/data/generator/action/datasets/droid_lerobot_dataset_config.py new file mode 100644 index 00000000..e97489c5 --- /dev/null +++ b/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/toml/sft_config/action_policy_droid_repro.toml b/examples/toml/sft_config/action_policy_droid_repro.toml index 661f7bdb..2326d407 100644 --- a/examples/toml/sft_config/action_policy_droid_repro.toml +++ b/examples/toml/sft_config/action_policy_droid_repro.toml @@ -3,10 +3,11 @@ # ============================================================================ # DROID action-policy SFT — run config for the `action_policy_droid_nano` -# experiment. The recipe knobs (optimizer/lr, scheduler type, grad_clip, -# count-based batch, action-head skip-on-load, dataset knobs) live in the -# registered experiment; this file only sets run-level scalars (iters, ckpt -# cadence, parallelism shape, wandb, VAE path). +# experiment. Reproduces the Cosmos3-Nano-Policy-DROID reference run: +# HSDP 32x8, global batch 8192, lr 2e-4, loss_scale 10, 10000 iters. The +# remaining recipe knobs (grad_clip, count-based batch, action-head +# skip-on-load, dataset knobs) live in the registered experiment; this file +# pins the run-level scalars that define the reference reproduction. # # Env required: # DROID_ROOT=/path/to/droid_lerobot_640x360/success @@ -27,8 +28,8 @@ wandb_mode = "online" precision = "bfloat16" [model.parallelism] -data_parallel_shard_degree = 8 # 8-GPU model shard; set replicate for multi-node HSDP -data_parallel_replicate_degree = 1 +data_parallel_shard_degree = 8 # 8-GPU model shard +data_parallel_replicate_degree = 32 # HSDP 32x8 = 256 ranks (GB200 reference: 64 nodes x 4) [model.activation_checkpointing] mode = "full" @@ -37,19 +38,37 @@ save_ops_regex = ["fmha"] [model.tokenizer] vae_path = "${oc.env:WAN_VAE_PATH}" +[model.rectified_flow_training_config] +loss_scale = 10.0 # generator (diffusion) loss weight, reference value + +[optimizer] +lr = 2.0e-04 + [scheduler] -cycle_lengths = [10000] # match max_iter +cycle_lengths = [100000] # long cosine cycle: lr decays slowly across the 10000-iter run + +[dataloader_train] +max_samples_per_batch = 32 # samples packed into each per-rank batch (res480) + +[dataloader_train.dataloader] +num_workers = 16 # decode-worker throughput tuning; adjust to host CPU count +batch_size = 16 +prefetch_factor = 2 [trainer] -max_iter = 10000 -logging_iter = 50 +max_iter = 10000 +logging_iter = 50 +grad_accum_iter = 1 # global batch = max_samples 32 x (shard 8 x replicate 32) x 1 = 8192 + +[trainer.callbacks.compile_tokenizer] +enabled = true # torch.compile the Wan VAE tokenizer (speed); warms up at res480 +warmup_resolutions = ["480"] [checkpoint] load_path = "${oc.env:BASE_CHECKPOINT_PATH}" save_iter = 1000 -# max_samples_per_batch is 128 in the experiment — samples packed into each per-rank batch -# (the num_workers x prefetch_factor workers just decode in parallel to keep it fed); res480, -# reference recipe, validated multi-node on GB200. -# On lower-memory GPUs, reduce it at launch, e.g.: -# --opts dataloader_train.max_samples_per_batch=32 +# The 256-rank HSDP 32x8 shape (global batch 8192) is the GB200 reference. To fit +# fewer GPUs while keeping global batch 8192, lower data_parallel_replicate_degree +# and raise grad_accum_iter at launch — e.g. one 8-GPU node: +# --opts model.parallelism.data_parallel_replicate_degree=1 trainer.grad_accum_iter=32 diff --git a/tests/action_policy_regression_test.py b/tests/action_policy_regression_test.py new file mode 100644 index 00000000..a2206512 --- /dev/null +++ b/tests/action_policy_regression_test.py @@ -0,0 +1,466 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: OpenMDW-1.1 + +"""Numerical-stability regression test for the ACTION-policy SFT launches (DROID + LIBERO). + +The action-policy analogue of ``tests/launch_regression_test.py`` (vision/VLM): +re-runs the same ``torchrun`` command the launcher shells execute — capped to +10 iterations, ``--deterministic``, seed 42 — for each of the two action-policy +recipes and asserts that the rank-0 per-step ``Loss:`` series reproduces the +inline goldens at the bottom of this file (per GPU arch, with a tolerance). + +Both recipes post-train the public **Cosmos3-Nano** base (registered experiments +``action_policy_{libero,droid}_nano``) — the same base + Wan2.2 VAE the vision +smoke/regression tests use — so a single 4-GPU node (``data_parallel_shard_degree=4``, +``replicate_degree=1``) runs them. The recipe knobs (optimizer, action-loss +weight, dataset transforms, iterable episode-shuffle) live in the experiment; +this test only caps iters + shapes the run for a deterministic single-node +capture. + +Specs +----- +* ``action_policy_libero`` — ``examples/toml/sft_config/action_policy_libero_repro.toml``. + Data auto-downloads: the ``libero_10`` suite of ``nvidia/LIBERO_LeRobot_v3`` + (small LeRobot dir), cached across runs. Collapses the recipe's HSDP replicate + 2 -> 1 to fit one 4-GPU node. +* ``action_policy_droid`` — ``examples/toml/sft_config/action_policy_droid_repro.toml``. + The full DROID split is far too large to auto-download in CI, so this spec + requires a pre-staged LOCAL copy via ``DROID_ROOT`` and SKIPS when unset. + ``DROID_ROOT`` is the **versioned merged root** whose basename is a + ``LEROBOT_ROOTS`` version key (e.g. ``…/droid_plus_lerobot_640x360_20260412``), + NOT its ``success/`` subdir — the recipe's i4 lazy dataset appends ``success/`` + itself (``use_success_only``). The skip-check looks for + ``/success/meta/info.json``. + +Determinism notes +----------------- +The launch passes ``--deterministic`` + ``PYTHONHASHSEED=42`` and the recipes +seed the episode-shuffle stream (``episode_shuffle_seed=42``); the shard +assignment ``(rank, worker)`` and per-epoch permutation are seed-derived, so the +data order reproduces across runs for any fixed world size / ``num_workers``. +``compile_tokenizer`` is disabled for the capture (torch.compile makes the +all-rank grad-norm reduction non-bit-exact — same reason ``vision_sft_nano`` +pins ``grad_norm=None`` on H100), so ``grad_norm`` goldens are left ``None`` and +only the loss series is asserted. On GB200 the sibling vision recipe reproduces +loss bit-exact across all 10 iters under ``--deterministic``; the action recipe +shares the MoT/attention stack, so the goldens use the tight default tolerance +(loosen via ``loss_tol_bands`` if a captured tail proves noisy). + +Inputs land in the documented ``.gitignore``-d locations (``examples/data/``, +``examples/checkpoints/``), cached across runs and shared with the vision smoke +test; run output goes under the pytest tmp dir. + +Invocation (inside the training container, from the repo root, on a 4-GPU +node):: + + # LIBERO only (auto-downloads its data): + pytest -s tests/action_policy_regression_test.py --num-gpus=4 --levels=2 -o addopts= + # include DROID (pre-stage the local LeRobot split): + DROID_ROOT=/path/to/Cosmos3-DROID/success \ + pytest -s tests/action_policy_regression_test.py --num-gpus=4 --levels=2 -o addopts= + +Without ``--num-gpus``/``--levels`` (e.g. the no-GPU pre-commit CI) the tests +are not collected. + +Refreshing the goldens (after an intentional numerical change, on the target +arch):: + + COSMOS_ACTION_REGRESSION_UPDATE_GOLDENS=1 pytest -s tests/action_policy_regression_test.py \ + --num-gpus=4 --levels=2 -o addopts= + +That prints the captured series for each spec; copy them into the matching +``_GOLDENS[]`` entry below. +""" + +from __future__ import annotations + +import os +import re +import shutil +import socket +import subprocess +import sys +from dataclasses import dataclass +from pathlib import Path + +import pytest + +from cosmos_framework.inference.fixtures.args import MAX_GPUS + +REPO_ROOT = Path(__file__).resolve().parents[1] + +# Shared base artifacts — match the vision smoke test + the action launcher +# defaults so no path overrides are needed once present (all .gitignore-d, +# cached across runs). +_WAN_VAE = REPO_ROOT / "examples/checkpoints/wan22_vae/Wan2.2_VAE.pth" +_DCP_DIR = REPO_ROOT / "examples/checkpoints/Cosmos3-Nano" +# LIBERO-10 LeRobot suite (auto-downloaded). ``LIBERO_ROOT`` must point at the +# suite dir itself (contains meta/info.json). +_LIBERO_DIR = REPO_ROOT / "examples/data/LIBERO_LeRobot_v3" +_LIBERO_ROOT = _LIBERO_DIR / "libero_10" + +# rank-0 per-iteration loss from the IterSpeed callback's on_training_step_end: +# [RANK 0] Iteration 1: Hit counter: 1/50 | Loss: 0.2515 | Time: 120.42s +_VFM_LOSS_RE = re.compile(r"\[RANK\s+0\]\s+Iteration\s+\d+:\s+Hit counter:[^|]+\|\s+Loss:\s+([0-9.eE+-]+)") + +_DEFAULT_RTOL = 1e-3 +_DEFAULT_ATOL = 1e-3 + + +def _free_port() -> int: + """Return a currently-free TCP port for torchrun's rendezvous (avoids + ``EADDRINUSE`` from a hardcoded ``master_port`` when a prior run lingers).""" + with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s: + s.bind(("", 0)) + return s.getsockname()[1] + + +def _detect_arch() -> str: + """Map ``torch.cuda.get_device_name(0)`` to a goldens key.""" + import torch + + if not torch.cuda.is_available(): + return "unknown" + name = torch.cuda.get_device_name(0).upper() + if "GB200" in name: + return "gb200" + if "H100" in name or "H200" in name: + return "h100" + return "unknown" + + +def _run(cmd: list[str], log_file: Path, extra_env: dict | None = None) -> tuple[int, str]: + """Run ``cmd`` from the repo root, tee combined output to ``log_file``. + + Streams live to stdout (so CI shows progress under ``pytest -s``) while + capturing into the log + a string. Inherits the caller's env plus + ``PYTHONPATH=.``. + """ + env = os.environ.copy() + env["PYTHONPATH"] = f".:{env.get('PYTHONPATH', '')}" + if extra_env: + env.update(extra_env) + log_file.parent.mkdir(parents=True, exist_ok=True) + captured: list[str] = [] + with log_file.open("w") as fp: + proc = subprocess.Popen( + cmd, + env=env, + cwd=str(REPO_ROOT), + stdout=subprocess.PIPE, + stderr=subprocess.STDOUT, + text=True, + bufsize=1, + ) + assert proc.stdout is not None + for line in proc.stdout: + sys.stdout.write(line) + sys.stdout.flush() + fp.write(line) + captured.append(line) + returncode = proc.wait() + return returncode, "".join(captured) + + +# --- input staging (shared with tests/nano_training_smoke_test.py) ----------- + + +def _ensure_wan_vae(log_dir: Path) -> None: + """Download the Wan2.2 VAE if not already present.""" + if _WAN_VAE.is_file(): + return + rc, out = _run( + [ + "uvx", + "hf@latest", + "download", + "Wan-AI/Wan2.2-TI2V-5B", + "Wan2.2_VAE.pth", + "--local-dir", + str(_WAN_VAE.parent), + "--quiet", + ], + log_dir / "download_wan_vae.log", + ) + assert rc == 0, f"Wan VAE download failed (exit {rc}):\n{out[-2000:]}" + assert _WAN_VAE.is_file(), f"Wan VAE missing at {_WAN_VAE} after download" + + +def _ensure_dcp(log_dir: Path) -> None: + """Convert Cosmos3-Nano to DCP if not already present.""" + if _DCP_DIR.is_dir() and any(_DCP_DIR.iterdir()): + return + rc, out = _run( + [ + "python", + "-m", + "cosmos_framework.scripts.convert_model_to_dcp", + "--checkpoint-path", + "Cosmos3-Nano", + "-o", + str(_DCP_DIR), + ], + log_dir / "convert_to_dcp.log", + ) + assert rc == 0, f"convert_model_to_dcp failed (exit {rc}):\n{out[-3000:]}" + assert _DCP_DIR.is_dir() and any(_DCP_DIR.iterdir()), f"DCP not written to {_DCP_DIR}" + + +def _ensure_libero(log_dir: Path) -> None: + """Download the ``libero_10`` suite of ``nvidia/LIBERO_LeRobot_v3`` if absent.""" + if (_LIBERO_ROOT / "meta" / "info.json").is_file(): + return + rc, out = _run( + [ + "uvx", + "hf@latest", + "download", + "--repo-type", + "dataset", + "nvidia/LIBERO_LeRobot_v3", + "--include", + "libero_10/**", + "--local-dir", + str(_LIBERO_DIR), + "--quiet", + ], + log_dir / "download_libero.log", + ) + assert rc == 0, f"LIBERO_LeRobot_v3 download failed (exit {rc}):\n{out[-2000:]}" + assert (_LIBERO_ROOT / "meta" / "info.json").is_file(), ( + f"LIBERO suite missing {_LIBERO_ROOT}/meta/info.json after download" + ) + + +# --- launch specs ------------------------------------------------------------ + + +# Overrides shared by both action specs: cap the run to a deterministic 10-iter +# single-node capture. Keeps ``model.config.compile.enabled`` at its recipe +# default (ON) to match launch_regression_test.py's nano spec — the loss is +# bit-exact under compile (only grad-norm is perturbed, which we don't assert). +# ``compile_tokenizer`` off just to skip its warmup. shard 4 x replicate 1 fits +# one 4-GPU node so the FSDP reduce-scatter/all-gather stays intra-node (NVLink) +# and reproduces bit-exact. A small packed batch keeps the 10 iters quick. +# +# NB (DROID on Blackwell): on gb200 the GradClip callback's compiled +# ``_fused_nan_to_num`` kernel can fail to launch through torch's static Triton +# launcher at this small shard=4 config ("CUDA driver error: invalid argument") +# — a launcher edge case, not a numeric issue (loss is identical eager). The +# H200/CI path this test targets does not hit it; the LIBERO spec is unaffected +# on either arch. +_COMMON_OVERRIDES: tuple[str, ...] = ( + "trainer.max_iter=10", + "trainer.logging_iter=1", + "trainer.seed=42", + "job.wandb_mode=disabled", + "upload_reproducible_setup=false", + "checkpoint.save_iter=999999", # no checkpoint writes during the capture + "trainer.callbacks.compile_tokenizer.enabled=false", + "model.config.parallelism.data_parallel_shard_degree=4", + "model.config.parallelism.data_parallel_replicate_degree=1", + "dataloader_train.max_samples_per_batch=8", +) + + +@dataclass(frozen=True) +class LaunchSpec: + """A single action-policy launch flow under regression.""" + + key: str + sft_toml: str + extra_hydra_args: tuple[str, ...] + requires_droid_root: bool = False + nproc_per_node: int = 4 + deterministic: bool = True + loss_rtol: float = _DEFAULT_RTOL + loss_atol: float = _DEFAULT_ATOL + # Optional tiered tolerance: each ``(count, rtol, atol)`` applies to the next + # ``count`` iters in order; counts must sum to 10. Empty => uniform default. + loss_tol_bands: tuple[tuple[int, float, float], ...] = () + + +_SPECS: dict[str, LaunchSpec] = { + "action_policy_libero": LaunchSpec( + key="action_policy_libero", + sft_toml="examples/toml/sft_config/action_policy_libero_repro.toml", + extra_hydra_args=_COMMON_OVERRIDES, + ), + "action_policy_droid": LaunchSpec( + key="action_policy_droid", + sft_toml="examples/toml/sft_config/action_policy_droid_repro.toml", + extra_hydra_args=_COMMON_OVERRIDES, + requires_droid_root=True, + ), +} + + +def _run_torchrun(spec: LaunchSpec, run_dir: Path) -> str: + """Invoke the same ``torchrun`` command the launcher shell runs; return the log text.""" + run_dir.mkdir(parents=True, exist_ok=True) + cmd = [ + "torchrun", + f"--nproc_per_node={spec.nproc_per_node}", + f"--master_port={_free_port()}", + "-m", + "cosmos_framework.scripts.train", + f"--sft-toml={spec.sft_toml}", + ] + if spec.deterministic: + cmd.append("--deterministic") + cmd += ["--", *spec.extra_hydra_args] + + rc, out = _run( + cmd, + run_dir / "training.log", + extra_env={ + "PYTHONHASHSEED": "42", + "IMAGINAIRE_OUTPUT_ROOT": str(run_dir / "output"), + "WAN_VAE_PATH": str(_WAN_VAE), + "BASE_CHECKPOINT_PATH": str(_DCP_DIR), + "LIBERO_ROOT": str(_LIBERO_ROOT), + # DROID_ROOT passes through from the caller's env (spec skips if unset). + }, + ) + if rc != 0 and "Done with training" not in out: + pytest.fail( + f"{spec.key}: torchrun failed (exit {rc}) and log lacks 'Done with training'.\nLog tail:\n{out[-3000:]}" + ) + return out + + +def _rank0_losses(text: str) -> list[float]: + """Parse the rank-0 per-iteration ``Loss:`` series (one value per step), in order.""" + vals: list[float] = [] + for m in _VFM_LOSS_RE.finditer(text): + v = float(m.group(1)) + vals.append(v) + return vals + + +# --- fixtures ---------------------------------------------------------------- + + +@pytest.fixture(scope="module", autouse=True) +def _require_4_gpus() -> None: + """Skip the module unless we can launch a 4-GPU training run here.""" + if shutil.which("torchrun") is None: + pytest.skip("torchrun not on PATH — must run inside the training container") + if shutil.which("uvx") is None: + pytest.skip("uvx not on PATH — required to download the dataset / Wan VAE") + try: + import torch + except Exception as exc: # pragma: no cover + pytest.skip(f"torch unavailable ({exc!r})") + if not torch.cuda.is_available() or torch.cuda.device_count() < 4: + pytest.skip(f"requires 4 visible CUDA devices, found {torch.cuda.device_count()}") + + +def _assert_spec_matches_goldens(spec_key: str, tmp_path: Path) -> None: + """Re-run ``spec``'s torchrun command and check the loss series against goldens.""" + spec = _SPECS[spec_key] + if spec.requires_droid_root: + droid_root = os.environ.get("DROID_ROOT", "") + # DROID_ROOT is the versioned merged root (basename is a LEROBOT_ROOTS key); + # the i4 lazy dataset appends success/ (use_success_only), so meta/info.json + # lives under /success, not at the root. + if not droid_root or not (Path(droid_root) / "success" / "meta" / "info.json").is_file(): + pytest.skip( + f"{spec.key}: set DROID_ROOT to a local versioned DROID LeRobot root " + "(basename a LEROBOT_ROOTS key, e.g. droid_plus_lerobot_640x360_20260412, " + "containing success/meta/info.json) to run this spec" + ) + + arch = _detect_arch() + + # Stage shared inputs (cached across runs); LIBERO data only for the libero spec. + _ensure_wan_vae(tmp_path) + _ensure_dcp(tmp_path) + if not spec.requires_droid_root: + _ensure_libero(tmp_path) + + out = _run_torchrun(spec, tmp_path) + assert "Done with training" in out, f"{spec.key}: training did not finish:\n{out[-3000:]}" + losses = _rank0_losses(out) + run_detail = f"\n--- {spec.key} run log (last 3000 chars) ---\n{out[-3000:]}" + assert len(losses) == 10, f"{spec.key}: expected 10 rank-0 losses, parsed {losses}{run_detail}" + assert all(v == v and abs(v) != float("inf") for v in losses), ( + f"{spec.key}: non-finite loss in series {losses}{run_detail}" + ) + + # Refresh path: print captured values for manual copy into ``_GOLDENS``. + if os.environ.get("COSMOS_ACTION_REGRESSION_UPDATE_GOLDENS") == "1": + print(f"\n# --- goldens for arch={arch!r} key={spec.key!r} ---") + print(f'"{spec.key}": {{"loss": {losses}}},') + pytest.skip( + f"captured fresh loss series for arch={arch!r} key={spec.key!r}; copy the printed " + f"dict into _GOLDENS[{arch!r}] at the bottom of action_policy_regression_test.py, " + "then rerun without COSMOS_ACTION_REGRESSION_UPDATE_GOLDENS to assert." + ) + + arch_goldens = _GOLDENS.get(arch) + if not arch_goldens or spec.key not in arch_goldens: + pytest.skip( + f"no goldens for arch={arch!r} key={spec.key!r} yet — capture on this hardware with " + f"COSMOS_ACTION_REGRESSION_UPDATE_GOLDENS=1 (parsed {len(losses)} finite losses this run)" + ) + expected = arch_goldens[spec.key]["loss"] + + bands = spec.loss_tol_bands or ((10, spec.loss_rtol, spec.loss_atol),) + assert sum(c for c, _, _ in bands) == 10, f"{spec.key}: loss_tol_bands must sum to 10" + start = 0 + for count, rtol, atol in bands: + end = start + count + assert losses[start:end] == pytest.approx(expected[start:end], rel=rtol, abs=atol), ( + f"{spec.key} ({arch}): rank-0 loss[{start}:{end}] (rel={rtol}, abs={atol}) " + f"does not match goldens\n got : {losses[start:end]}\n" + f" expected: {expected[start:end]}{run_detail}" + ) + start = end + + +# --- tests ------------------------------------------------------------------- + + +if MAX_GPUS == 4: + + @pytest.mark.level(2) + @pytest.mark.gpus(4) + @pytest.mark.parametrize("spec_key", list(_SPECS), ids=list(_SPECS)) + def test_action_policy_regression(spec_key: str, tmp_path: Path) -> None: + """Deterministic 10-iter action-policy SFT; assert the rank-0 loss goldens.""" + _assert_spec_matches_goldens(spec_key, tmp_path) + + +# Goldens keyed by GPU arch then ``LaunchSpec.key``. ``loss`` is the rank-0 +# per-step series over the 10 deterministic iters. Refresh on the target arch +# with ``COSMOS_ACTION_REGRESSION_UPDATE_GOLDENS=1`` (see the module docstring). +# The test skips (not fails) for any arch/spec without an entry, so goldens can +# land incrementally as they are captured on each arch. +# +# Captured with torch.compile ON, --deterministic, seed 42, single-node +# data_parallel_shard_degree=4 (intra-node NVLink FSDP reduction is bit-exact). +_GOLDENS: dict[str, dict[str, dict[str, list[float]]]] = { + # H200 (Hopper) CI arch. LIBERO is the primary numerical golden; the DROID + # spec needs its dataset (DROID_ROOT), so it skips unless one is provided. + "h100": { + "action_policy_libero": { + "loss": [ + 15.8107, + 15.2467, + 15.9856, + 16.5306, + 14.3738, + 16.1460, + 16.6093, + 14.8846, + 16.0632, + 16.6449, + ], + }, + }, +} + + +if __name__ == "__main__": # pragma: no cover — manual driver + sys.exit(pytest.main([__file__, "-v", "-s", "-o", "addopts="]))