From 5be5ae958e14c77cc6572f0494ae0f66ec708de5 Mon Sep 17 00:00:00 2001 From: Wei-Cheng Tseng Date: Sat, 4 Jul 2026 21:09:25 -0700 Subject: [PATCH 1/3] init --- cosmos_framework/configs/base/config.py | 1 + .../action_fd_droid_posttrain.py | 220 +++++++++++++ .../generator/action/datasets/__init__.py | 2 + .../action/datasets/action_sft_dataset.py | 62 ++++ .../datasets/droid_merged_lerobot_dataset.py | 300 ++++++++++++++++++ docs/action_fd_droid_posttrain.md | 160 ++++++++++ .../launch_sft_action_fd_droid_posttrain.sh | 39 +++ .../sft_config/action_fd_droid_posttrain.toml | 47 +++ 8 files changed, 831 insertions(+) create mode 100644 cosmos_framework/configs/base/experiment/action/posttrain_config/action_fd_droid_posttrain.py create mode 100644 cosmos_framework/data/generator/action/datasets/droid_merged_lerobot_dataset.py create mode 100644 docs/action_fd_droid_posttrain.md create mode 100644 examples/launch_sft_action_fd_droid_posttrain.sh create mode 100644 examples/toml/sft_config/action_fd_droid_posttrain.toml diff --git a/cosmos_framework/configs/base/config.py b/cosmos_framework/configs/base/config.py index ec65e32e..13588ece 100644 --- a/cosmos_framework/configs/base/config.py +++ b/cosmos_framework/configs/base/config.py @@ -97,6 +97,7 @@ def make_config() -> Config: import cosmos_framework.configs.base.experiment.action.posttrain_config.action_policy_droid_nano # noqa: F401 import cosmos_framework.configs.base.experiment.action.posttrain_config.action_policy_libero_all_nano # noqa: F401 import cosmos_framework.configs.base.experiment.action.posttrain_config.action_policy_libero_nano # noqa: F401 + import cosmos_framework.configs.base.experiment.action.posttrain_config.action_fd_droid_posttrain # noqa: F401 import cosmos_framework.configs.base.experiment.sft.vision_sft_nano # noqa: F401 import cosmos_framework.configs.base.experiment.sft.vision_sft_super # noqa: F401 diff --git a/cosmos_framework/configs/base/experiment/action/posttrain_config/action_fd_droid_posttrain.py b/cosmos_framework/configs/base/experiment/action/posttrain_config/action_fd_droid_posttrain.py new file mode 100644 index 00000000..b49f7c26 --- /dev/null +++ b/cosmos_framework/configs/base/experiment/action/posttrain_config/action_fd_droid_posttrain.py @@ -0,0 +1,220 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: OpenMDW-1.1 + +"""``action_fd_droid_posttrain`` — DROID+LeRobot forward-dynamics post-training.""" + +import copy + +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_merged_lerobot_sft_dataset +from cosmos_framework.data.generator.joint_dataloader import PackingDataLoader, RankPartitionedDataLoader + +cs = ConfigStore.instance() + + +def _make_model_config() -> dict: + cfg = copy.deepcopy(NANO_MODEL_CONFIG) + + cfg["sound_gen"] = False + cfg["sound_dim"] = 64 + cfg["sound_latent_fps"] = 25 + cfg["max_num_tokens_after_packing"] = 74000 + cfg["resolution"] = "720" + cfg["activation_checkpointing"]["mode"] = "selective" + + cfg["tokenizer"]["encode_exact_durations"] = [17] + + cfg["diffusion_expert_config"].update( + base_fps=24, + enable_fps_modulation=True, + load_weights_from_pretrained=False, + patch_spatial=2, + unified_3d_mrope_temporal_modality_margin=15000, + unified_3d_mrope_reset_spatial_ids=True, + ) + cfg["rectified_flow_training_config"].update( + image_loss_scale=None, + loss_scale=10.0, + shift={"256": 3, "480": 5, "720": 10}, + sound_loss_scale=2.0, + train_time_video_distribution="waver", + train_time_weight="uniform", + use_discrete_rf=False, + ) + return cfg + + +action_fd_droid_posttrain = LazyDict( + dict( + defaults=[ + {"override /data_train": None}, + {"override /data_val": None}, + {"override /model": "mot_fsdp"}, + {"override /optimizer": "fusedadamw"}, + {"override /scheduler": "lambdalinear"}, + {"override /tokenizer": "wan2pt2_tokenizer"}, + {"override /sound_tokenizer": None}, + {"override /vlm_config": None}, + {"override /checkpoint": "gcp"}, + {"override /callbacks": ["basic", "optimization", "job_monitor", "training_stats"]}, + {"override /ema": "power"}, + {"override /ckpt_type": "dcp"}, + "_self_", + ], + job=dict( + project="cosmos3_action_fd", + group="action_sft", + name="${now:%Y-%m-%d_%H-%M-%S}_action_fd_droid_posttrain", + wandb_mode="disabled", + ), + model=dict( + config=_make_model_config(), + ), + optimizer=dict( + betas=[0.9, 0.99], + eps=1.0e-08, + fused=True, + keys_to_select=[ + "moe_gen", + "time_embedder", + "vae2llm", + "llm2vae", + "action2llm", + "llm2action", + "action_modality_embed", + ], + lr=1.0e-04, + lr_multipliers={ + "action2llm": 5.0, + "llm2action": 5.0, + "action_modality_embed": 5.0, + }, + optimizer_type="FusedAdam", + weight_decay=0.05, + ), + scheduler=dict( + cycle_lengths=[20000], + f_max=[0.4], + f_min=[0.0], + f_start=[0.0], + lr_scheduler_type="LambdaLinear", + verbosity_interval=0, + warm_up_steps=[0], + ), + trainer=dict( + distributed_parallelism="fsdp", + grad_accum_iter=1, + logging_iter=50, + max_iter=20000, + max_val_iter=None, + run_validation=False, + run_validation_on_start=False, + save_zero_checkpoint=False, + seed=42, + timeout_period=999999999, + validation_iter=100, + compile_config=dict(recompile_limit=100, use_duck_shape=False), + cudnn=dict(benchmark=True, deterministic=False), + ddp=dict(broadcast_buffers=True, find_unused_parameters=False, static_graph=True), + grad_scaler_args=dict(enabled=False), + straggler_detection=dict(enabled=True, report_freq=50), + callbacks=dict( + dataloader_speed=dict(every_n=100, save_s3=False, step_size=1), + device_monitor=dict(every_n=200, log_memory_detail=True, save_s3=False, step_size=1), + grad_clip=dict(clip_norm=1.0, force_finite=True), + heart_beat=dict(every_n=200, save_s3=False, step_size=1, update_interval_in_minute=20), + iter_speed=dict(every_n=50, hit_thres=50, save_s3=False, save_s3_every_log_n=500), + low_precision=dict(update_iter=1), + manual_gc=dict(every_n=200, gc_level=1, warm_up=1), + norm_monitor=dict(every_n=100), + param_count=dict(save_s3=False), + sigma_loss_analysis=dict(every_n=500, every_n_viz=500, save_s3=False), + skip_nan_step=dict(max_consecutive_nan=100), + training_stats=dict(log_freq=100), + compile_tokenizer=dict(enabled=True, warmup_resolutions=["480"]), + ), + ), + checkpoint=dict( + dcp_async_mode_enabled=False, + enable_gcs_patch_in_boto3=True, + keys_not_to_resume=[], + # Skip net_ema. (EMA warm-starts from net, see dcp.py) + keys_to_skip_loading=[ + "net_ema.", + ], + load_ema_to_reg=False, + load_from_object_store=dict(bucket="", credentials="", enabled=False), + save_to_object_store=dict(bucket="", credentials="", enabled=False), + load_path="???", # Cosmos3-Nano DCP dir; supply via TOML/env + load_training_state=False, + only_load_scheduler_state=False, + save_iter=250, + strict_resume=True, + verbose=True, + ), + dataloader_train=L(PackingDataLoader)( + audio_sample_rate=48000, + dataset_name="action_droid", + max_samples_per_batch=None, + max_sequence_length="${model.config.max_num_tokens_after_packing}", + patch_spatial="${model.config.diffusion_expert_config.patch_spatial}", + sound_latent_fps="${model.config.sound_latent_fps}", + tokenizer_spatial_compression_factor="${model.config.tokenizer.spatial_compression_factor}", + tokenizer_temporal_compression_factor="${model.config.tokenizer.temporal_compression_factor}", + dataloader=L(RankPartitionedDataLoader)( + batch_size=1, + in_order=False, + num_workers=3, + persistent_workers=True, + pin_memory=True, + prefetch_factor=2, + sampler=None, + datasets=dict( + droid=dict( + ratio=1, + dataset=L(get_action_droid_merged_lerobot_sft_dataset)( + root="${oc.env:DROID_MERGED_LEROBOT_ROOT}", + fps=15.0, + chunk_length=16, + action_space="midtrain", + mode="forward_dynamics", + use_state=False, + use_success_only=False, + split="train", + iterable_shuffle=True, + episode_shuffle_seed=42, + use_image_augmentation=False, + use_filter_dict=False, + filter_dict_path=None, + action_normalization=None, + viewpoint="concat_view", + resolution="480", + max_action_dim="${model.config.max_action_dim}", + cfg_dropout_rate=0.1, + tokenizer_config="${model.config.vlm_config.tokenizer}", + append_idle_frames=True, + idle_frames_dropout=0.05, + format_prompt_as_json=True, + ), + ), + ), + ), + ), + dataloader_val=None, + upload_reproducible_setup=False, + ), + flags={"allow_objects": True}, +) + + +cs.store( + group="experiment", + package="_global_", + name="action_fd_droid_posttrain", + node=action_fd_droid_posttrain, +) diff --git a/cosmos_framework/data/generator/action/datasets/__init__.py b/cosmos_framework/data/generator/action/datasets/__init__.py index 62437797..ed5dd19d 100644 --- a/cosmos_framework/data/generator/action/datasets/__init__.py +++ b/cosmos_framework/data/generator/action/datasets/__init__.py @@ -11,6 +11,7 @@ from cosmos_framework.data.generator.action.datasets.agibotworld_beta_lerobot_dataset import AgiBotWorldBetaLeRobotDataset from cosmos_framework.data.generator.action.datasets.base_dataset import ActionBaseDataset from cosmos_framework.data.generator.action.datasets.bridge_orig_lerobot_dataset import BridgeOrigLeRobotDataset +from cosmos_framework.data.generator.action.datasets.droid_merged_lerobot_dataset import DROIDMergedLeRobotDataset from cosmos_framework.data.generator.action.datasets.droid_lerobot_dataset import DROIDLeRobotDataset from cosmos_framework.data.generator.action.datasets.fractal_lerobot_dataset import FractalLeRobotDataset from cosmos_framework.data.generator.action.datasets.libero_lerobot_dataset import LIBEROLeRobotDataset @@ -23,6 +24,7 @@ "AgiBotWorldBetaLeRobotDataset", "BridgeOrigLeRobotDataset", "DROIDLeRobotDataset", + "DROIDMergedLeRobotDataset", "FractalLeRobotDataset", "LIBEROLeRobotDataset", "RoboMINDFrankaDataset", 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..81231c44 100644 --- a/cosmos_framework/data/generator/action/datasets/action_sft_dataset.py +++ b/cosmos_framework/data/generator/action/datasets/action_sft_dataset.py @@ -19,6 +19,7 @@ from torch.utils.data import Dataset, IterableDataset, get_worker_info +from cosmos_framework.data.generator.action.datasets.droid_merged_lerobot_dataset import DROIDMergedLeRobotDataset from cosmos_framework.data.generator.action.datasets.droid_lerobot_dataset import DROIDLeRobotDataset from cosmos_framework.data.generator.action.datasets.libero_lerobot_dataset import LIBEROLeRobotDataset from cosmos_framework.data.generator.action.transforms import ActionTransformPipeline @@ -142,6 +143,67 @@ def get_action_droid_sft_dataset( return sft +def get_action_droid_merged_lerobot_sft_dataset( + *, + root: str, + fps: float = 15.0, + chunk_length: int = 16, + action_space: str = "midtrain", + mode: str = "forward_dynamics", + use_state: bool = False, + action_normalization: str | None = None, + viewpoint: str = "concat_view", + split: str = "train", + use_success_only: bool = False, + use_image_augmentation: bool = False, + use_filter_dict: bool = False, + filter_dict_path: str | None = None, + resolution: str | int = "480", + max_action_dim: int = 64, + tokenizer_config: dict | None = None, + cfg_dropout_rate: float = 0.1, + append_viewpoint_info: bool = True, + append_duration_fps_timestamps: bool = True, + append_resolution_info: bool = True, + append_idle_frames: bool = True, + idle_frames_dropout: float = 0.05, + format_prompt_as_json: bool = True, + iterable_shuffle: bool = False, + episode_shuffle_seed: int = 42, +) -> Dataset: + """Build the DROID-Merged LeRobot SFT dataset for action FD recipes.""" + dataset = DROIDMergedLeRobotDataset( + root=root, + fps=fps, + chunk_length=chunk_length, + viewpoint=viewpoint, + action_space=action_space, + mode=mode, + use_state=use_state, + action_normalization=action_normalization, + use_image_augmentation=use_image_augmentation, + use_filter_dict=use_filter_dict, + filter_dict_path=filter_dict_path, + split=split, + use_success_only=use_success_only, + ) + transform = ActionTransformPipeline( + tokenizer_config=tokenizer_config, + cfg_dropout_rate=cfg_dropout_rate, + max_action_dim=max_action_dim, + append_viewpoint_info=append_viewpoint_info, + append_duration_fps_timestamps=append_duration_fps_timestamps, + append_resolution_info=append_resolution_info, + append_idle_frames=append_idle_frames, + idle_frames_dropout=idle_frames_dropout, + format_prompt_as_json=format_prompt_as_json, + ) + sft = ActionSFTDataset(dataset, transform, resolution) + if iterable_shuffle: + return ActionIterableShuffleDataset(sft, seed=episode_shuffle_seed) + return sft + + def get_action_libero_sft_dataset( *, root: str, diff --git a/cosmos_framework/data/generator/action/datasets/droid_merged_lerobot_dataset.py b/cosmos_framework/data/generator/action/datasets/droid_merged_lerobot_dataset.py new file mode 100644 index 00000000..0e7df59c --- /dev/null +++ b/cosmos_framework/data/generator/action/datasets/droid_merged_lerobot_dataset.py @@ -0,0 +1,300 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: OpenMDW-1.1 + +"""DROID-Merged LeRobot dataset used by the action FD DROID recipe.""" + +from __future__ import annotations + +import json +import random +from pathlib import Path +from typing import Any + +import numpy as np +import pandas as pd +import pyarrow.parquet as pq +import torch + +from cosmos_framework.data.generator.action.datasets.droid_lerobot_dataset import ( + _ACTION_GRIPPER_FEATURE, + _GRIPPER_STATE_FEATURE, + _JOINT_ACTION_FEATURE, + _JOINT_STATE_FEATURE, + _LIST_COLUMNS, + _STATE_FEATURE, + DROIDLeRobotDataset, +) +from cosmos_framework.data.generator.action.domain_utils import get_domain_id + +_MERGED_ROOT_SPLITS = ("success", "failure") +_ROW_ROOT_KEY = "_cosmos_root_idx" + + +class DROIDMergedLeRobotDataset(DROIDLeRobotDataset): + """DROID-Merged dataset with merged ``success`` + ``failure`` root support. + + This is intentionally separate from ``DROIDLeRobotDataset`` so the released + DROID policy recipe keeps its single-root behavior unchanged. + """ + + def __init__( + self, + root: str, + fps: float = 15.0, + chunk_length: int = 16, + mode: str = "policy", + pose_convention: str = "backward_framewise", + tolerance_s: float = 2e-4, + viewpoint: str = "concat_view", + action_space: str = "midtrain", + use_state: bool = False, + action_normalization: str | None = None, + use_image_augmentation: bool = False, + use_filter_dict: bool = False, + filter_dict_path: str | None = None, + split: str = "train", + use_success_only: bool = False, + ) -> None: + action_space = "ee_pose" if action_space == "midtrain" else action_space + if viewpoint != "concat_view": + raise NotImplementedError("DROIDMergedLeRobotDataset only supports concat_view.") + if action_space not in ("ee_pose", "joint_pos"): + raise NotImplementedError("action_space must be 'midtrain', 'ee_pose', or 'joint_pos'.") + 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") + if split != "train": + raise NotImplementedError("DROIDMergedLeRobotDataset currently supports split='train' only.") + if pose_convention != "backward_framewise": + raise NotImplementedError("DROIDMergedLeRobotDataset only supports backward_framewise pose deltas.") + + self._fps = float(fps) + self._dt = 1.0 / self._fps + self._chunk_length = int(chunk_length) + self._sample_stride = 1 + self._mode = mode + self._pose_convention = pose_convention + self._tolerance_s = float(tolerance_s) + self._viewpoint = viewpoint + self._domain_name = "droid_lerobot" + self._domain_id = get_domain_id(self._domain_name) + self._action_normalization = None if action_space == "joint_pos" else action_normalization + self._norm_stats: dict[str, torch.Tensor] | None = None + self._root = Path(root) + self._roots = self._resolve_lerobot_roots(self._root, use_success_only=use_success_only) + self._infos = [json.loads((root_path / "meta" / "info.json").read_text()) for root_path in self._roots] + self._action_space = action_space + self._use_state = bool(use_state) + self._use_image_augmentation = bool(use_image_augmentation) + self._image_augmentor = None + self._use_filter_dict = bool(use_filter_dict) + self._filter_dict_path = filter_dict_path + self._episodes, self._tasks, episode_maps = self._load_metadata(self._roots) + + if action_space == "joint_pos": + feature_cols = [_JOINT_ACTION_FEATURE, _ACTION_GRIPPER_FEATURE, _JOINT_STATE_FEATURE, _GRIPPER_STATE_FEATURE] + else: + feature_cols = [_STATE_FEATURE, _ACTION_GRIPPER_FEATURE] + columns = ["index", "episode_index", "task_index", "timestamp", *feature_cols] + index_parts, episode_parts, root_parts, task_parts, ts_parts = [], [], [], [], [] + feature_parts: dict[str, list] = {c: [] for c in feature_cols} + for root_idx, root_path in enumerate(self._roots): + for path in sorted((root_path / "data").glob("chunk-*/file-*.parquet")): + table = pq.read_table(path, columns=columns) + local_episode = table["episode_index"].to_numpy().astype(np.int64) + global_episode = np.asarray( + [episode_maps[root_idx][int(episode_idx)] for episode_idx in local_episode], + dtype=np.int64, + ) + index_parts.append(table["index"].to_numpy()) + episode_parts.append(global_episode) + root_parts.append(np.full(global_episode.shape, root_idx, dtype=np.int64)) + 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)) + if not index_parts: + roots = ", ".join(str(root_path) for root_path in self._roots) + raise FileNotFoundError(f"No DROID-Merged parquet shards found under: {roots}") + + episode_all = np.concatenate(episode_parts).astype(np.int64) + index_all = np.concatenate(index_parts).astype(np.int64) + order = np.lexsort((index_all, episode_all)) + self._row_episode = np.concatenate(episode_parts).astype(np.int64)[order] + self._row_root = np.concatenate(root_parts).astype(np.int64)[order] + self._row_task = np.concatenate(task_parts).astype(np.int64)[order] + self._row_timestamp = np.concatenate(ts_parts).astype(np.float64)[order] + self._feat = { + c: np.concatenate(feature_parts[c], axis=0).astype(np.float32)[order] for c in feature_cols + } + + 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) + + 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) + + @staticmethod + def _is_lerobot_root(path: Path) -> bool: + return (path / "meta" / "info.json").is_file() + + @classmethod + def _resolve_lerobot_roots(cls, root: Path, *, use_success_only: bool) -> list[Path]: + if cls._is_lerobot_root(root): + return [root] + + roots: list[Path] = [] + split_names = ("success",) if use_success_only else _MERGED_ROOT_SPLITS + for split_name in split_names: + split_root = root / split_name + if cls._is_lerobot_root(split_root): + roots.append(split_root) + continue + if split_root.is_dir(): + roots.extend(sorted(path for path in split_root.iterdir() if cls._is_lerobot_root(path))) + + if roots: + return roots + raise FileNotFoundError( + f"{root} is not a LeRobot root and no merged DROID split roots were found under " + f"{', '.join(split_names)}." + ) + + @staticmethod + def _load_metadata(roots: list[Path]) -> tuple[dict[int, dict[str, Any]], dict[tuple[int, int], str], list[dict[int, int]]]: + episodes: dict[int, dict[str, Any]] = {} + tasks: dict[tuple[int, int], str] = {} + episode_maps: list[dict[int, int]] = [] + next_episode_idx = 0 + + for root_idx, root_path in enumerate(roots): + episode_map: dict[int, int] = {} + for path in sorted((root_path / "meta" / "episodes").glob("chunk-*/file-*.parquet")): + for row in pq.read_table(path).to_pylist(): + local_episode_idx = int(row["episode_index"]) + if local_episode_idx in episode_map: + continue + global_episode_idx = next_episode_idx + next_episode_idx += 1 + episode_map[local_episode_idx] = global_episode_idx + row = dict(row) + row[_ROW_ROOT_KEY] = root_idx + episodes[global_episode_idx] = row + if not episode_map: + raise FileNotFoundError(f"No episode metadata found under {root_path / 'meta' / 'episodes'}") + episode_maps.append(episode_map) + + tasks_df = pd.read_parquet(root_path / "meta" / "tasks.parquet") + task_texts = tasks_df["task"] if "task" in tasks_df.columns else tasks_df.index + for task, task_index in zip(task_texts, tasks_df["task_index"]): + tasks[(root_idx, int(task_index))] = str(task) + + return episodes, tasks, episode_maps + + def _window_rows(self, start: int, stop: int, episode_index: int) -> list[dict[str, Any]]: + return [ + { + "episode_index": episode_index, + "root_idx": int(self._row_root[j]), + "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) + 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) + + 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]["root_idx"]), int(observation_rows[0]["task_index"]))] + ai_caption = random.choice(task.split(" | ")) + + 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, + ) + + def _video_path(self, episode: dict[str, Any], video_key: str) -> Path: + root_idx = int(episode.get(_ROW_ROOT_KEY, 0)) + info = self._infos[root_idx] + root = self._roots[root_idx] + chunk_idx = int( + episode.get( + f"videos/{video_key}/chunk_index", + episode.get(f"videos/{video_key}/episode_chunk", episode.get("data/chunk_index", 0)), + ) + ) + file_idx = int( + episode.get( + f"videos/{video_key}/file_index", + episode.get(f"videos/{video_key}/episode_file", episode.get("data/file_index", 0)), + ) + ) + rel = info["video_path"].format( + video_key=video_key, + chunk_index=chunk_idx, + file_index=file_idx, + episode_chunk=chunk_idx, + episode_file=file_idx, + ) + return root / rel diff --git a/docs/action_fd_droid_posttrain.md b/docs/action_fd_droid_posttrain.md new file mode 100644 index 00000000..6ea24e65 --- /dev/null +++ b/docs/action_fd_droid_posttrain.md @@ -0,0 +1,160 @@ +# Action FD DROID Post-Training + +This document describes how to run the `action_fd_droid_posttrain` experiment. +It trains Cosmos3-DROID forward dynamics in `cosmos_framework`. + +## Overview + +| Piece | Value | +| --- | --- | +| Experiment | `action_fd_droid_posttrain` | +| TOML | `examples/toml/sft_config/action_fd_droid_posttrain.toml` | +| Launch shell | `examples/launch_sft_action_fd_droid_posttrain.sh` | +| Config module | `cosmos_framework/configs/base/experiment/action/posttrain_config/action_fd_droid_posttrain.py` | +| Dataset wrapper | `cosmos_framework/data/generator/action/datasets/droid_merged_lerobot_dataset.py` | +| Dataset root | [Cosmos3-DROID](https://huggingface.co/datasets/nvidia/Cosmos3-DROID) parent root containing `success/` and `failure/` | +| Task mode | `forward_dynamics` | +| Action space | `midtrain` alias for 10-D ee-pose `[pos_delta, rot6d_delta, gripper]` | +| Chunk length / resolution | `16` frames at `480` | + +## Prerequisites + +- Install the training environment as described in [`docs/setup.md`](./setup.md). +- Run commands from the repository root. +- In NGC / PyTorch containers, set `LD_LIBRARY_PATH=''` before Python commands. + +## Inputs You Provide + +This package ships the training stack — the registered `action_fd_droid_posttrain` +experiment, the dataset class, and the paired TOML/launch shell. Three inputs +are external and must be provided per environment: + +1. **[Cosmos3-DROID](https://huggingface.co/datasets/nvidia/Cosmos3-DROID) dataset (in LeRobotDataset v3.0 format)** — pre-download the + dataset and point `DATASET_PATH` or `DROID_MERGED_LEROBOT_ROOT` at the + resulting `.../Cosmos3-DROID` parent directory. This experiment trains on both + `success/` and `failure/` subsets. +2. **DCP base checkpoint** — prepare a base DCP checkpoint and point + `BASE_CHECKPOINT_PATH` at it. For local smoke runs this can be + `examples/checkpoints/Cosmos3-Nano`. + +## Data Layout + +Set `DATASET_PATH` or `DROID_MERGED_LEROBOT_ROOT` to the Cosmos3-DROID parent +directory: + +```shell +export DATASET_PATH=/path/to/Cosmos3-DROID +``` + +The loader uses both subsets: + +```text +$DATASET_PATH/ +├── success/ +│ └── meta/info.json +└── failure/ + └── meta/info.json +``` + +or institution-sharded split roots: + +```text +$DATASET_PATH/ +├── success//meta/info.json +└── failure//meta/info.json +``` + +The launch shell bridges `DATASET_PATH` to `DROID_MERGED_LEROBOT_ROOT`. + +## Full Reproduction + +The flow mirrors the other SFT recipes: + +```shell +# Step 1: prepare Cosmos3-DROID -> $DATASET_PATH. +# If you do not already have the dataset locally, download it from Hugging Face. +uvx hf@latest download --repo-type dataset nvidia/Cosmos3-DROID \ + --local-dir examples/data/Cosmos3-DROID --quiet +export DATASET_PATH=examples/data/Cosmos3-DROID + +# Step 2: point to the base DCP checkpoint and Wan2.2 VAE. +export BASE_CHECKPOINT_PATH=examples/checkpoints/Cosmos3-Nano +export WAN_VAE_PATH=/lustre/fsw/portfolios/cosmos/projects/cosmos_base_training/users/weichengt/.cache/huggingface/hub/models--Wan-AI--Wan2.2-TI2V-5B/snapshots/921dbaf3f1674a56f47e83fb80a34bac8a8f203e/Wan2.2_VAE.pth + +# Step 3: choose the output root and launch. +export IMAGINAIRE_OUTPUT_ROOT=/path/to/output_root +export LD_LIBRARY_PATH='' + +bash examples/launch_sft_action_fd_droid_posttrain.sh +``` + +`BASE_CHECKPOINT_PATH` should point at the base DCP checkpoint. When using the +paired launch shell, it must be a local directory because the shared launcher +validates it before starting `torchrun`. + +## Validate The Config + +Use `--dryrun` before launching workers: + +```shell +PYTHONPATH=. python -m cosmos_framework.scripts.train \ + --sft-toml examples/toml/sft_config/action_fd_droid_posttrain.toml \ + --dryrun +``` + +## Run Training + +Recommended paired launch shell: + +```shell +bash examples/launch_sft_action_fd_droid_posttrain.sh +``` + +Pass short smoke-run overrides through `EXTRA_TAIL_OVERRIDES`: + +```shell +export EXTRA_TAIL_OVERRIDES="trainer.max_iter=10 checkpoint.save_iter=10" +bash examples/launch_sft_action_fd_droid_posttrain.sh +``` + +Single-node, 8 GPU: + +```shell +PYTHONPATH=. torchrun --nproc_per_node=8 -m cosmos_framework.scripts.train \ + --sft-toml examples/toml/sft_config/action_fd_droid_posttrain.toml +``` + +Multi-node HSDP: + +```shell +PYTHONPATH=. torchrun --nnodes=$NNODES --node_rank=$NODE_RANK --nproc_per_node=8 \ + -m cosmos_framework.scripts.train \ + --sft-toml examples/toml/sft_config/action_fd_droid_posttrain.toml \ + -- model.parallelism.data_parallel_replicate_degree=$NNODES +``` + +Keep `data_parallel_shard_degree=8` and set +`model.parallelism.data_parallel_replicate_degree` to the number of nodes. + +## Outputs + +Training outputs land under: + +```text +$IMAGINAIRE_OUTPUT_ROOT/cosmos3_action_fd/action_sft// +``` + +DCP checkpoints are saved under: + +```text +$RUN_DIR/checkpoints/iter_/ +``` + +The run is resumable by relaunching with the same output directory and job name. + +## Notes + +- This recipe uses `mode="forward_dynamics"`, so actions are conditioning and the + model trains video prediction from the first frame plus action sequence. +- `DROIDMergedLeRobotDataset` accepts `action_space="midtrain"` as an + i4-compatible alias for the 10-D ee-pose path. diff --git a/examples/launch_sft_action_fd_droid_posttrain.sh b/examples/launch_sft_action_fd_droid_posttrain.sh new file mode 100644 index 00000000..fba2c9e6 --- /dev/null +++ b/examples/launch_sft_action_fd_droid_posttrain.sh @@ -0,0 +1,39 @@ +#!/usr/bin/env bash +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: OpenMDW-1.1 + +# ============================================================================ +# Structured-TOML launch for action_fd_droid_posttrain. +# +# This trains forward dynamics on the Cosmos3-DROID success + failure splits. See +# docs/action_fd_droid_posttrain.md. +# +# Env vars (override for your filesystem): +# DATASET_PATH Cosmos3-DROID parent dir (success/ + failure/) +# BASE_CHECKPOINT_PATH Base DCP checkpoint +# WAN_VAE_PATH Wan2.2 VAE .pth +# WANDB_API_KEY for online logging (TOML wandb_mode="online") +# NPROC_PER_NODE torchrun --nproc_per_node (default 8) +# EXTRA_TAIL_OVERRIDES space-separated Hydra overrides +# +# Single-node smoke: +# export EXTRA_TAIL_OVERRIDES="trainer.max_iter=10 checkpoint.save_iter=10" +# bash examples/launch_sft_action_fd_droid_posttrain.sh +# +# Multi-node: launch on every worker. For HSDP set +# model.parallelism.data_parallel_replicate_degree = (shard stays 8). +# ============================================================================ + +TOML_FILE="examples/toml/sft_config/action_fd_droid_posttrain.toml" +: "${DATASET_PATH:=${DROID_MERGED_LEROBOT_ROOT:-examples/data/Cosmos3-DROID}}" +: "${BASE_CHECKPOINT_PATH:=examples/checkpoints/Cosmos3-Nano}" + +WAN_VAE_PATH="${WAN_VAE_PATH:-examples/checkpoints/wan22_vae/Wan2.2_VAE.pth}" + +EXTRA_DATASET_CHECK='DROID_MERGED_LEROBOT_ROOT="${DROID_MERGED_LEROBOT_ROOT:-$DATASET_PATH}"; [[ "$DROID_MERGED_LEROBOT_ROOT" = /* ]] || DROID_MERGED_LEROBOT_ROOT="$WORKDIR/$DROID_MERGED_LEROBOT_ROOT"; [[ "$WAN_VAE_PATH" = /* ]] || WAN_VAE_PATH="$WORKDIR/$WAN_VAE_PATH"; export DROID_MERGED_LEROBOT_ROOT WAN_VAE_PATH; [[ -f "$DROID_MERGED_LEROBOT_ROOT/success/meta/info.json" || -n "$(compgen -G "$DROID_MERGED_LEROBOT_ROOT/success/*/meta/info.json")" ]] || { echo "ERROR: missing Cosmos3-DROID success split under $DROID_MERGED_LEROBOT_ROOT (expected success/meta/info.json or success/*/meta/info.json; see docs/action_fd_droid_posttrain.md)" >&2; exit 1; }; [[ -f "$DROID_MERGED_LEROBOT_ROOT/failure/meta/info.json" || -n "$(compgen -G "$DROID_MERGED_LEROBOT_ROOT/failure/*/meta/info.json")" ]] || { echo "ERROR: missing Cosmos3-DROID failure split under $DROID_MERGED_LEROBOT_ROOT (expected failure/meta/info.json or failure/*/meta/info.json; see docs/action_fd_droid_posttrain.md)" >&2; exit 1; }; [[ -f "$WAN_VAE_PATH" ]] || { echo "ERROR: WAN_VAE_PATH not found: $WAN_VAE_PATH" >&2; exit 1; }' + +TAIL_OVERRIDES=( + ${EXTRA_TAIL_OVERRIDES:-} +) + +source "$(dirname "${BASH_SOURCE[0]}")/_sft_launcher_common.sh" diff --git a/examples/toml/sft_config/action_fd_droid_posttrain.toml b/examples/toml/sft_config/action_fd_droid_posttrain.toml new file mode 100644 index 00000000..71a1a640 --- /dev/null +++ b/examples/toml/sft_config/action_fd_droid_posttrain.toml @@ -0,0 +1,47 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: OpenMDW-1.1 + +# ============================================================================ +# DROID-Merged LeRobot forward-dynamics post-training. +# +# Env required: +# DROID_MERGED_LEROBOT_ROOT=/path/to/droid_merged_lerobot_640x360_20260412 +# BASE_CHECKPOINT_PATH= +# WAN_VAE_PATH= +# IMAGINAIRE_OUTPUT_ROOT=/path/to/output_root +# +# The registered experiment owns dataset mode/action-space wiring: +# mode="forward_dynamics", action_space="midtrain", chunk_length=16, res480. +# ============================================================================ + +[job] +task = "vfm" +experiment = "action_fd_droid_posttrain" +project = "cosmos3_action_fd" +group = "action_sft" +name = "action_fd_droid_posttrain" +wandb_mode = "online" + +[model] +precision = "bfloat16" + +[model.parallelism] +data_parallel_shard_degree = 8 +data_parallel_replicate_degree = 1 + +[model.activation_checkpointing] +mode = "selective" + +[model.tokenizer] +vae_path = "${oc.env:WAN_VAE_PATH}" + +[scheduler] +cycle_lengths = [20000] + +[trainer] +max_iter = 20000 +logging_iter = 50 + +[checkpoint] +load_path = "${oc.env:BASE_CHECKPOINT_PATH}" +save_iter = 250 From 95e9c5bdaf0eeb1e84b49c9cc452433c8d1345d4 Mon Sep 17 00:00:00 2001 From: Wei-Cheng Tseng Date: Sat, 4 Jul 2026 21:29:56 -0700 Subject: [PATCH 2/3] format --- docs/action_fd_droid_posttrain.md | 22 +++++++++---------- .../launch_sft_action_fd_droid_posttrain.sh | 0 2 files changed, 11 insertions(+), 11 deletions(-) mode change 100644 => 100755 examples/launch_sft_action_fd_droid_posttrain.sh diff --git a/docs/action_fd_droid_posttrain.md b/docs/action_fd_droid_posttrain.md index 6ea24e65..97402621 100644 --- a/docs/action_fd_droid_posttrain.md +++ b/docs/action_fd_droid_posttrain.md @@ -5,17 +5,17 @@ It trains Cosmos3-DROID forward dynamics in `cosmos_framework`. ## Overview -| Piece | Value | -| --- | --- | -| Experiment | `action_fd_droid_posttrain` | -| TOML | `examples/toml/sft_config/action_fd_droid_posttrain.toml` | -| Launch shell | `examples/launch_sft_action_fd_droid_posttrain.sh` | -| Config module | `cosmos_framework/configs/base/experiment/action/posttrain_config/action_fd_droid_posttrain.py` | -| Dataset wrapper | `cosmos_framework/data/generator/action/datasets/droid_merged_lerobot_dataset.py` | -| Dataset root | [Cosmos3-DROID](https://huggingface.co/datasets/nvidia/Cosmos3-DROID) parent root containing `success/` and `failure/` | -| Task mode | `forward_dynamics` | -| Action space | `midtrain` alias for 10-D ee-pose `[pos_delta, rot6d_delta, gripper]` | -| Chunk length / resolution | `16` frames at `480` | +| Piece | Value | +| ------------------------- | ---------------------------------------------------------------------------------------------------------------------- | +| Experiment | `action_fd_droid_posttrain` | +| TOML | `examples/toml/sft_config/action_fd_droid_posttrain.toml` | +| Launch shell | `examples/launch_sft_action_fd_droid_posttrain.sh` | +| Config module | `cosmos_framework/configs/base/experiment/action/posttrain_config/action_fd_droid_posttrain.py` | +| Dataset wrapper | `cosmos_framework/data/generator/action/datasets/droid_merged_lerobot_dataset.py` | +| Dataset root | [Cosmos3-DROID](https://huggingface.co/datasets/nvidia/Cosmos3-DROID) parent root containing `success/` and `failure/` | +| Task mode | `forward_dynamics` | +| Action space | `midtrain` alias for 10-D ee-pose `[pos_delta, rot6d_delta, gripper]` | +| Chunk length / resolution | `16` frames at `480` | ## Prerequisites diff --git a/examples/launch_sft_action_fd_droid_posttrain.sh b/examples/launch_sft_action_fd_droid_posttrain.sh old mode 100644 new mode 100755 From 7dbe7f94eb7566d824490f86e289ca1bcbf9b887 Mon Sep 17 00:00:00 2001 From: Wei-Cheng Tseng Date: Sun, 5 Jul 2026 20:15:39 -0700 Subject: [PATCH 3/3] clean alias naming --- .../posttrain_config/action_fd_droid_posttrain.py | 4 ++-- .../generator/action/datasets/action_sft_dataset.py | 2 +- .../action/datasets/droid_lerobot_dataset.py | 2 +- .../action/datasets/droid_merged_lerobot_dataset.py | 5 ++--- docs/action_fd_droid_posttrain.md | 13 ++++++------- examples/launch_sft_action_fd_droid_posttrain.sh | 4 ++-- .../toml/sft_config/action_fd_droid_posttrain.toml | 4 ++-- 7 files changed, 16 insertions(+), 18 deletions(-) diff --git a/cosmos_framework/configs/base/experiment/action/posttrain_config/action_fd_droid_posttrain.py b/cosmos_framework/configs/base/experiment/action/posttrain_config/action_fd_droid_posttrain.py index b49f7c26..5c54fbd0 100644 --- a/cosmos_framework/configs/base/experiment/action/posttrain_config/action_fd_droid_posttrain.py +++ b/cosmos_framework/configs/base/experiment/action/posttrain_config/action_fd_droid_posttrain.py @@ -178,10 +178,10 @@ def _make_model_config() -> dict: droid=dict( ratio=1, dataset=L(get_action_droid_merged_lerobot_sft_dataset)( - root="${oc.env:DROID_MERGED_LEROBOT_ROOT}", + root="${oc.env:DATASET_PATH}", fps=15.0, chunk_length=16, - action_space="midtrain", + action_space="ee_pose", mode="forward_dynamics", use_state=False, use_success_only=False, 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 81231c44..0939eb6d 100644 --- a/cosmos_framework/data/generator/action/datasets/action_sft_dataset.py +++ b/cosmos_framework/data/generator/action/datasets/action_sft_dataset.py @@ -148,7 +148,7 @@ def get_action_droid_merged_lerobot_sft_dataset( root: str, fps: float = 15.0, chunk_length: int = 16, - action_space: str = "midtrain", + action_space: str = "ee_pose", mode: str = "forward_dynamics", use_state: bool = False, action_normalization: str | None = None, 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..fe8f1724 100644 --- a/cosmos_framework/data/generator/action/datasets/droid_lerobot_dataset.py +++ b/cosmos_framework/data/generator/action/datasets/droid_lerobot_dataset.py @@ -59,7 +59,7 @@ class DROIDLeRobotDataset(ActionBaseDataset): Two action layouts: * ``action_space="ee_pose"`` (default): 10D ``[pos_delta(3), rot6d_delta(6), - gripper(1)]``, quantile-normalized (the v1.2 midtrain default). + gripper(1)]``, quantile-normalized. * ``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 diff --git a/cosmos_framework/data/generator/action/datasets/droid_merged_lerobot_dataset.py b/cosmos_framework/data/generator/action/datasets/droid_merged_lerobot_dataset.py index 0e7df59c..acb812b5 100644 --- a/cosmos_framework/data/generator/action/datasets/droid_merged_lerobot_dataset.py +++ b/cosmos_framework/data/generator/action/datasets/droid_merged_lerobot_dataset.py @@ -46,7 +46,7 @@ def __init__( pose_convention: str = "backward_framewise", tolerance_s: float = 2e-4, viewpoint: str = "concat_view", - action_space: str = "midtrain", + action_space: str = "ee_pose", use_state: bool = False, action_normalization: str | None = None, use_image_augmentation: bool = False, @@ -55,11 +55,10 @@ def __init__( split: str = "train", use_success_only: bool = False, ) -> None: - action_space = "ee_pose" if action_space == "midtrain" else action_space if viewpoint != "concat_view": raise NotImplementedError("DROIDMergedLeRobotDataset only supports concat_view.") if action_space not in ("ee_pose", "joint_pos"): - raise NotImplementedError("action_space must be 'midtrain', 'ee_pose', or 'joint_pos'.") + raise NotImplementedError("action_space must be 'ee_pose' or 'joint_pos'.") 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: diff --git a/docs/action_fd_droid_posttrain.md b/docs/action_fd_droid_posttrain.md index 97402621..97f1d191 100644 --- a/docs/action_fd_droid_posttrain.md +++ b/docs/action_fd_droid_posttrain.md @@ -14,7 +14,7 @@ It trains Cosmos3-DROID forward dynamics in `cosmos_framework`. | Dataset wrapper | `cosmos_framework/data/generator/action/datasets/droid_merged_lerobot_dataset.py` | | Dataset root | [Cosmos3-DROID](https://huggingface.co/datasets/nvidia/Cosmos3-DROID) parent root containing `success/` and `failure/` | | Task mode | `forward_dynamics` | -| Action space | `midtrain` alias for 10-D ee-pose `[pos_delta, rot6d_delta, gripper]` | +| Action space | `ee_pose`: 10-D `[pos_delta, rot6d_delta, gripper]` | | Chunk length / resolution | `16` frames at `480` | ## Prerequisites @@ -30,7 +30,7 @@ experiment, the dataset class, and the paired TOML/launch shell. Three inputs are external and must be provided per environment: 1. **[Cosmos3-DROID](https://huggingface.co/datasets/nvidia/Cosmos3-DROID) dataset (in LeRobotDataset v3.0 format)** — pre-download the - dataset and point `DATASET_PATH` or `DROID_MERGED_LEROBOT_ROOT` at the + dataset and point `DATASET_PATH` at the resulting `.../Cosmos3-DROID` parent directory. This experiment trains on both `success/` and `failure/` subsets. 2. **DCP base checkpoint** — prepare a base DCP checkpoint and point @@ -39,8 +39,7 @@ are external and must be provided per environment: ## Data Layout -Set `DATASET_PATH` or `DROID_MERGED_LEROBOT_ROOT` to the Cosmos3-DROID parent -directory: +Set `DATASET_PATH` to the Cosmos3-DROID parent directory: ```shell export DATASET_PATH=/path/to/Cosmos3-DROID @@ -64,7 +63,7 @@ $DATASET_PATH/ └── failure//meta/info.json ``` -The launch shell bridges `DATASET_PATH` to `DROID_MERGED_LEROBOT_ROOT`. +The launch shell passes `DATASET_PATH` through to the training config. ## Full Reproduction @@ -156,5 +155,5 @@ The run is resumable by relaunching with the same output directory and job name. - This recipe uses `mode="forward_dynamics"`, so actions are conditioning and the model trains video prediction from the first frame plus action sequence. -- `DROIDMergedLeRobotDataset` accepts `action_space="midtrain"` as an - i4-compatible alias for the 10-D ee-pose path. +- `DROIDMergedLeRobotDataset` uses `action_space="ee_pose"` for the 10-D + end-effector pose action path. diff --git a/examples/launch_sft_action_fd_droid_posttrain.sh b/examples/launch_sft_action_fd_droid_posttrain.sh index fba2c9e6..8dc7dd50 100755 --- a/examples/launch_sft_action_fd_droid_posttrain.sh +++ b/examples/launch_sft_action_fd_droid_posttrain.sh @@ -25,12 +25,12 @@ # ============================================================================ TOML_FILE="examples/toml/sft_config/action_fd_droid_posttrain.toml" -: "${DATASET_PATH:=${DROID_MERGED_LEROBOT_ROOT:-examples/data/Cosmos3-DROID}}" +: "${DATASET_PATH:=examples/data/Cosmos3-DROID}" : "${BASE_CHECKPOINT_PATH:=examples/checkpoints/Cosmos3-Nano}" WAN_VAE_PATH="${WAN_VAE_PATH:-examples/checkpoints/wan22_vae/Wan2.2_VAE.pth}" -EXTRA_DATASET_CHECK='DROID_MERGED_LEROBOT_ROOT="${DROID_MERGED_LEROBOT_ROOT:-$DATASET_PATH}"; [[ "$DROID_MERGED_LEROBOT_ROOT" = /* ]] || DROID_MERGED_LEROBOT_ROOT="$WORKDIR/$DROID_MERGED_LEROBOT_ROOT"; [[ "$WAN_VAE_PATH" = /* ]] || WAN_VAE_PATH="$WORKDIR/$WAN_VAE_PATH"; export DROID_MERGED_LEROBOT_ROOT WAN_VAE_PATH; [[ -f "$DROID_MERGED_LEROBOT_ROOT/success/meta/info.json" || -n "$(compgen -G "$DROID_MERGED_LEROBOT_ROOT/success/*/meta/info.json")" ]] || { echo "ERROR: missing Cosmos3-DROID success split under $DROID_MERGED_LEROBOT_ROOT (expected success/meta/info.json or success/*/meta/info.json; see docs/action_fd_droid_posttrain.md)" >&2; exit 1; }; [[ -f "$DROID_MERGED_LEROBOT_ROOT/failure/meta/info.json" || -n "$(compgen -G "$DROID_MERGED_LEROBOT_ROOT/failure/*/meta/info.json")" ]] || { echo "ERROR: missing Cosmos3-DROID failure split under $DROID_MERGED_LEROBOT_ROOT (expected failure/meta/info.json or failure/*/meta/info.json; see docs/action_fd_droid_posttrain.md)" >&2; exit 1; }; [[ -f "$WAN_VAE_PATH" ]] || { echo "ERROR: WAN_VAE_PATH not found: $WAN_VAE_PATH" >&2; exit 1; }' +EXTRA_DATASET_CHECK='[[ "$WAN_VAE_PATH" = /* ]] || WAN_VAE_PATH="$WORKDIR/$WAN_VAE_PATH"; export WAN_VAE_PATH; [[ -f "$DATASET_PATH/success/meta/info.json" || -n "$(compgen -G "$DATASET_PATH/success/*/meta/info.json")" ]] || { echo "ERROR: missing Cosmos3-DROID success split under $DATASET_PATH (expected success/meta/info.json or success/*/meta/info.json; see docs/action_fd_droid_posttrain.md)" >&2; exit 1; }; [[ -f "$DATASET_PATH/failure/meta/info.json" || -n "$(compgen -G "$DATASET_PATH/failure/*/meta/info.json")" ]] || { echo "ERROR: missing Cosmos3-DROID failure split under $DATASET_PATH (expected failure/meta/info.json or failure/*/meta/info.json; see docs/action_fd_droid_posttrain.md)" >&2; exit 1; }; [[ -f "$WAN_VAE_PATH" ]] || { echo "ERROR: WAN_VAE_PATH not found: $WAN_VAE_PATH" >&2; exit 1; }' TAIL_OVERRIDES=( ${EXTRA_TAIL_OVERRIDES:-} diff --git a/examples/toml/sft_config/action_fd_droid_posttrain.toml b/examples/toml/sft_config/action_fd_droid_posttrain.toml index 71a1a640..65fbdfa1 100644 --- a/examples/toml/sft_config/action_fd_droid_posttrain.toml +++ b/examples/toml/sft_config/action_fd_droid_posttrain.toml @@ -5,13 +5,13 @@ # DROID-Merged LeRobot forward-dynamics post-training. # # Env required: -# DROID_MERGED_LEROBOT_ROOT=/path/to/droid_merged_lerobot_640x360_20260412 +# DATASET_PATH=/path/to/Cosmos3-DROID # BASE_CHECKPOINT_PATH= # WAN_VAE_PATH= # IMAGINAIRE_OUTPUT_ROOT=/path/to/output_root # # The registered experiment owns dataset mode/action-space wiring: -# mode="forward_dynamics", action_space="midtrain", chunk_length=16, res480. +# mode="forward_dynamics", action_space="ee_pose", chunk_length=16, res480. # ============================================================================ [job]