comos_predict_2B_GR1_action_cond

Cosmos-Predict2.5-2B GR1 Action-Conditioned World Model

Action-conditioned video prediction model fine-tuned from nvidia/Cosmos-Predict2.5-2B on GR1 dual-arm robot teleoperation data.

Given an initial observation frame and a sequence of robot actions, this model predicts future video frames that depict the robot executing those actions.

Model Details

• Base model: Cosmos-Predict2.5-2B (video2world pre-trained)
• Fine-tuning data: GR1 dual-arm robot teleoperation (LeRobot format)
• Action dimension: 29 (left_arm:7 + left_hand:6 + right_arm:7 + right_hand:6 + waist:3)
• Video frames: 13 frames per chunk (1 conditional + 12 predicted)
• Temporal compression: 4x (13 image frames = 4 latent frames)
• Resolution: 480x832
• Checkpoint format: EMA weights in bf16 (model_ema_bf16.pt, ~4GB)

Training Configuration

Inference

Prerequisites

This model requires the Cosmos-Predict2.5 codebase and the following dependencies:

• Wan2.1 VAE: Wan2.1_VAE.pth
• Text encoder: Cosmos-Reason1-7B
• Tokenizer: Qwen2.5-VL-7B-Instruct

Quick Start

python cosmos_predict2/_src/predict2/action/inference/inference_gr00t.py \
  --experiment=cosmos_predict2p5_2B_action_conditioned_gr00t_gr1_customized_13frame_full_16nodes_release_oss \
  --ckpt_path=/path/to/model_ema_bf16.pt \
  --input_video_root=/path/to/eval_data \
  --save_root=/path/to/output \
  --resolution 480,832 \
  --guidance 0 \
  --chunk_size 12 \
  --fps_downsample_ratio 2 \
  --save_fps 10 \
  --vae_path /path/to/Wan2.1_VAE.pth \
  --text_encoder_path /path/to/Cosmos-Reason1-7B \
  --qwen_path /path/to/Qwen2.5-VL-7B-Instruct

Input Data Format

The inference script expects a directory containing paired files:

• episode_XXXXXX.mp4 - Input video (first frame used as conditional)
• episode_XXXXXX_actions.npy - Action array of shape (T, 29) in numpy format

Eval Data Preparation

Use scripts/prepare_gr1_eval_data.py to extract eval episodes from a LeRobot dataset:

python scripts/prepare_gr1_eval_data.py \
  --dataset-path /path/to/GR1_robot \
  --output-dir /path/to/eval_data \
  --num-episodes 100

Examples

The examples/ directory contains sample predictions alongside ground truth videos:

File Structure

.
├── model_ema_bf16.pt           # Model weights (EMA, bf16)
├── inference/
│   ├── inference_gr00t.py      # Main inference script
│   └── inference_pipeline.py   # Inference pipeline (ActionVideo2WorldInference)
├── scripts/
│   ├── eval_gr1_robot.sh       # One-command eval script
│   ├── prepare_gr1_eval_data.py # Eval data preparation
│   └── convert_distcp_to_pt.py # DCP -> PT checkpoint converter
└── examples/
    ├── episode_000000.mp4      # Predicted videos
    ├── episode_000008.mp4
    ├── episode_000087.mp4
    └── gt/                     # Ground truth videos
        ├── episode_000000.mp4
        ├── episode_000008.mp4
        └── episode_000087.mp4

Limitations

• Trained only on GR1 dual-arm robot data; may not generalize to other embodiments
• Prediction quality degrades for long-horizon generation (many chunks)
• Episodes with fewer than 12 actions cannot be processed (chunk_size=12)
• No classifier-free guidance (guidance=0) works best for this model

License

Apache-2.0 (same as base Cosmos-Predict2.5)

Acknowledgements

• Base model: NVIDIA Cosmos-Predict2.5
• Training data: GR1 teleoperation dataset (LeRobot format)
• VAE: Wan2.1