Skip to content

Add post-training recipe for forward dynamics experiments on Droid#85

Open
WeiChengTseng wants to merge 3 commits into
NVIDIA:mainfrom
WeiChengTseng:cosmos3-action-fd
Open

Add post-training recipe for forward dynamics experiments on Droid#85
WeiChengTseng wants to merge 3 commits into
NVIDIA:mainfrom
WeiChengTseng:cosmos3-action-fd

Conversation

@WeiChengTseng

@WeiChengTseng WeiChengTseng commented Jul 5, 2026

Copy link
Copy Markdown

Summary

Add a complete SFT recipe for training Cosmos3 forward dynamics on the Cosmos3-DROID dataset (LeRobot v3.0 format).

  • Register action_fd_droid_posttrain experiment with mode="forward_dynamics", 10-D ee-pose action space (midtrain alias), 16-frame chunks at 480p resolution.
  • Add DROIDMergedLeRobotDataset to load both success/ and failure/ splits from a single parent root, without changing the existing single-root DROIDLeRobotDataset used by the policy recipe.
  • Add get_action_droid_merged_lerobot_sft_dataset() factory and wire it through ActionSFTDataset / ActionTransformPipeline.
  • Ship paired TOML + launch shell (action_fd_droid_posttrain.toml, launch_sft_action_fd_droid_posttrain.sh) and reproduction docs (docs/action_fd_droid_posttrain.md).

Motivation

Cosmos3 supports forward-dynamics inference (forward_dynamics mode: predict future video from first frame + action sequence), but there was no public post-training recipe for fine-tuning on Cosmos3-DROID. This PR closes that gap and mirrors the structure of existing SFT recipes (e.g. action policy on DROID/LIBERO).

Test plan

  • Config dry-run passes:
    PYTHONPATH=. python -m cosmos_framework.scripts.train \
      --sft-toml examples/toml/sft_config/action_fd_droid_posttrain.toml \
      --dryrun

@WeiChengTseng WeiChengTseng changed the title [Draft] Add post-training recipe for forward dynamics experiments on Droid Add post-training recipe for forward dynamics experiments on Droid Jul 6, 2026
@lfengad lfengad requested review from fwd4 and mli0603 July 6, 2026 05:44
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant