---
license: apache-2.0
tags:
  - robotics
  - world-model
  - video-prediction
  - action-conditioned
  - nvidia
  - cosmos
base_model: nvidia/Cosmos-Predict2.5-14B
---

# Cosmos-Predict2.5-14B GR1 Action-Conditioned World Model

Action-conditioned video prediction model fine-tuned from [nvidia/Cosmos-Predict2.5-14B](https://huggingface.co/nvidia/Cosmos-Predict2.5-14B) 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-14B (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
- **Network:** `cosmos_v1_14B_action_chunk_conditioned` (36 blocks, 5120 channels, 40 heads)
- **Checkpoint format:** EMA weights in bf16 (`model_ema_bf16.pt`, ~28GB)

## Training Configuration

| Parameter | Value |
|-----------|-------|
| GPUs | 8x H200 (140GB) |
| Batch size | 2 per GPU (global=16) |
| Learning rate | 4e-5 |
| Max iterations | 4,000+ |
| Optimizer | AdamW (cosine schedule) |
| Precision | bf16 |
| Episodes per task | 3 (uniform sampling) |
| FSDP shard size | 32 |
| VAE | Wan2.1 |
| Text encoder | Cosmos-Reason1-7B |
| Tokenizer | Qwen2.5-VL-7B-Instruct |

## Inference

### Prerequisites

This model requires the [Cosmos-Predict2.5](https://github.com/NVIDIA/Cosmos-Predict2.5) codebase and the following dependencies:

- Wan2.1 VAE: `Wan2.1_VAE.pth` (14B model includes its own copy)
- Text encoder: Cosmos-Reason1-7B
- Tokenizer: Qwen2.5-VL-7B-Instruct

### Quick Start

```bash
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 \
  --experiment_opts net=cosmos_v1_14B_action_chunk_conditioned model.config.fsdp_shard_size=32
```

Note: The `--experiment` uses the registered 2B experiment name for Hydra config lookup. The `--experiment_opts` overrides the network architecture to 14B.

### 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:

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

### One-command Eval

```bash
bash scripts/eval_gr1_robot.sh /path/to/model_ema_bf16.pt --debug
```

The eval script auto-detects 14B vs 2B from the checkpoint path and applies the correct network config.

## Examples

| Episode | Prediction | Ground Truth |
|---------|-----------|--------------|
| 000000 | [episode_000000.mp4](examples/predicted/episode_000000.mp4) | [episode_000000.mp4](examples/gt/episode_000000.mp4) |
| 000008 | [episode_000008.mp4](examples/predicted/episode_000008.mp4) | [episode_000008.mp4](examples/gt/episode_000008.mp4) |
| 000056 | [episode_000056.mp4](examples/predicted/episode_000056.mp4) | [episode_000056.mp4](examples/gt/episode_000056.mp4) |
| 000084 | [episode_000084.mp4](examples/predicted/episode_000084.mp4) | [episode_000084.mp4](examples/gt/episode_000084.mp4) |
| 000087 | [episode_000087.mp4](examples/predicted/episode_000087.mp4) | [episode_000087.mp4](examples/gt/episode_000087.mp4) |

## File Structure

```
.
├── model_ema_bf16.pt           # Model weights (EMA, bf16, ~28GB)
├── inference/
│   ├── inference_gr00t.py      # Main inference script
│   └── inference_pipeline.py   # Inference pipeline (ActionVideo2WorldInference)
├── scripts/
│   ├── eval_gr1_robot.sh       # One-command eval script (auto 2B/14B)
│   ├── prepare_gr1_eval_data.py # Eval data preparation
│   └── convert_distcp_to_pt.py # DCP -> PT checkpoint converter
└── examples/
    ├── predicted/              # Predicted videos
    └── gt/                     # Ground truth videos
```

## 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)
- 14B model requires ~60GB GPU memory for inference (single GPU)
- 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](https://github.com/NVIDIA/Cosmos-Predict2.5)
- Training data: GR1 teleoperation dataset (LeRobot format)
- VAE: [Wan2.1](https://huggingface.co/Wan-AI/Wan2.1-T2V-14B)