# Cosmos-Policy-ALOHA-Planning-Model-Predict2-2B

## Model Description

Cosmos-Policy-ALOHA-Planning-Model-Predict2-2B is a refined world model and value function checkpoint fine-tuned from [Cosmos-Policy-ALOHA-Predict2-2B](https://huggingface.co/nvidia/Cosmos-Policy-ALOHA-Predict2-2B) on policy rollout data. This checkpoint is designed to be used in conjunction with the base Cosmos Policy checkpoint for model-based planning via best-of-N search, achieving improved performance on challenging manipulation tasks. This checkpoint should NOT be deployed on its own.

**Paper**: *Cosmos Policy: Fine-Tuning Video Models for Visuomotor Control and Planning*

### Key Features

- **Refined predictions**: Fine-tuned on policy rollout data for more accurate world model and value function predictions
- **Dual deployment**: Used alongside base Cosmos Policy checkpoint for model-based planning
- **Improved performance**: Achieves 12.5 percentage point average score increase on challenging ALOHA tasks when used for planning

### Model Architecture

This model uses the same architecture as [Cosmos-Policy-ALOHA-Predict2-2B](https://huggingface.co/nvidia/Cosmos-Policy-ALOHA-Predict2-2B). Please refer to that model card and the [base Cosmos-Predict2-2B model card](https://huggingface.co/nvidia/Cosmos-Predict2-2B-Video2World) for detailed architecture specifications.

## Model Details

### Inputs

Same as [Cosmos-Policy-ALOHA-Predict2-2B](https://huggingface.co/nvidia/Cosmos-Policy-ALOHA-Predict2-2B):
- **Current state images**: Top-down camera, left wrist camera, right wrist camera (all 224x224 RGB)
- **Robot proprioception**: 14-dimensional (7 joint angles per arm)
- **Action chunk**: 50-timestep sequence of 14-dimensional actions (for world model and value prediction)

### Outputs

- **Future robot proprioception**: 14-dimensional state at timestep t+50
- **Future state images**: 
  - Top-down third-person camera prediction at timestep t+50
  - Left wrist camera prediction at timestep t+50
  - Right wrist camera prediction at timestep t+50
- **Future state value**: Expected cumulative reward from future state (V(s'))

**Note**: While this checkpoint can technically generate actions like the base policy, it is specifically designed and optimized for world model and value function predictions. For action generation, please use [Cosmos-Policy-ALOHA-Predict2-2B](https://huggingface.co/nvidia/Cosmos-Policy-ALOHA-Predict2-2B).

### Training Details

**Base Checkpoint**: [Cosmos-Policy-ALOHA-Predict2-2B](https://huggingface.co/nvidia/Cosmos-Policy-ALOHA-Predict2-2B)

**Fine-tuning Data**: 648 policy rollout episodes collected from various methods (see paper for details)
- Includes both successful and failed episodes
- Covers diverse initial conditions and execution trajectories
- Enables more accurate modeling of state transitions and value predictions beyond the demonstration distribution

**Training Configuration**:
- **Training steps**: Details in paper
- **Batch split**: 10/45/45 for policy/world model/value function (emphasis on world model and value function refinement)
- **GPUs**: 8 H100 GPUs

**Training Objective**: Fine-tuned with increased emphasis on world model and value function training (90% of training batches) to improve future state and value prediction accuracy for more effective planning.

## Usage: Dual Deployment for Model-Based Planning

This checkpoint is designed for **dual deployment** with the base Cosmos Policy checkpoint:

1. **Policy Model** ([Cosmos-Policy-ALOHA-Predict2-2B](https://huggingface.co/nvidia/Cosmos-Policy-ALOHA-Predict2-2B)): Generates N candidate action chunks
2. **Planning Model** (this checkpoint): For each candidate action:
   - Predicts future state (world model)
   - Predicts future state value (value function)
   - Averages across ensemble predictions (3 future state predictions × 5 value predictions = 15 total value estimates per action)
3. **Selection**: Execute the action chunk with the highest predicted value

See the paper for complete implementation details of the best-of-N planning algorithm.

## Performance

### Planning Performance on ALOHA Tasks

When used for model-based planning with the base policy checkpoint:

| Task | Base Policy Score | With Planning (this checkpoint) | Improvement |
|------|------------------|--------------------------------|-------------|
| put candies in bowl | 49.0 | 60.0 | +11.0 |
| put candy in ziploc bag | 70.0 | 84.0 | +14.0 |
| **Average** | **60.0** | **72.0** | **+12.5** |

Results are on challenging initial conditions for these two tasks. Planning with this checkpoint enables the policy to be more likely to avoid errors (e.g., losing grasp of objects) by selecting higher-quality actions.

## Important Usage Notes

**Inference Latency**: Model-based planning with dual deployment has significantly higher inference latency:
- **Planning mode (dual deployment)**: ~4.9 seconds per action chunk using 8 parallel H100 GPUs
- **Direct policy mode (base checkpoint only)**: ~0.95 seconds per action chunk using 1 H100 GPU

For applications requiring faster inference, we recommend using the base [Cosmos-Policy-ALOHA-Predict2-2B](https://huggingface.co/nvidia/Cosmos-Policy-ALOHA-Predict2-2B) checkpoint alone without planning.

**Hardware Requirements**: Model-based planning requires:
- Multiple GPUs for parallelized best-of-N search (8 GPUs recommended for N=8)
- Sufficient compute for ensemble predictions (3 world model queries × 5 value function queries per action)

**When to Use Planning**: Planning is most beneficial for:
- Challenging tasks with high precision requirements
- Situations where avoiding errors is critical
- Scenarios where additional compute time is acceptable

**Same warnings as base checkpoint apply**: Hardware compatibility, 25 Hz control frequency requirement, and real-world deployment safety considerations. See [Cosmos-Policy-ALOHA-Predict2-2B](https://huggingface.co/nvidia/Cosmos-Policy-ALOHA-Predict2-2B) model card for details.

## Notes

- **Specialized checkpoint**: Optimized specifically for world model and value function predictions, not action generation
- **Requires base policy**: Must be used in conjunction with Cosmos-Policy-ALOHA-Predict2-2B for planning
- **Compute-intensive**: Significantly higher computational requirements than direct policy execution
- **Real-world tested**: Evaluated on real ALOHA 2 hardware in challenging manipulation scenarios

## Citation

If you use this model, please cite the Cosmos Policy paper by Kim et al.
<!-- ```bibtex
# TODO: Add Cosmos Policy BibTeX
``` -->

## License

Please refer to the [base model license](https://huggingface.co/nvidia/Cosmos-Predict2-2B-Video2World) for licensing information.

## Related Resources

- **Base Policy Checkpoint**: [Cosmos-Policy-ALOHA-Predict2-2B](https://huggingface.co/nvidia/Cosmos-Policy-ALOHA-Predict2-2B) (required for planning)
- **Base Video Model**: [Cosmos-Predict2-2B-Video2World](https://huggingface.co/nvidia/Cosmos-Predict2-2B-Video2World)
- **Paper**: *Cosmos Policy: Fine-Tuning Video Models for Visuomotor Control and Planning*