Update README.md

README.md

@@ -1,115 +1,240 @@
-# Cosmos-Policy-RoboCasa-Predict2-2B
+---
+base_model:
+- nvidia/Cosmos-Predict2-2B-Video2World
+---
+# **Cosmos-Policy-RoboCasa-Predict2-2B**
 
-## Model Description
+[**Cosmos Policy**](https://huggingface.co/collections/nvidia/cosmos-policy) | [**Code**](http://github.com/NVlabs/cosmos-policy) | [**White Paper Coming Soon**]() | [**Website**](https://research.nvidia.com/labs/dir/cosmos-policy/)
 
-Cosmos-Policy-RoboCasa-Predict2-2B is a robot manipulation policy fine-tuned from the [NVIDIA Cosmos-Predict2-2B-Video2World](https://huggingface.co/nvidia/Cosmos-Predict2-2B-Video2World) video foundation model (checkpoint: `model-480p-16fps.pt`). This model achieves state-of-the-art performance on the RoboCasa simulation benchmark with a 67.1% average success rate across 24 kitchen manipulation tasks, using only ~50 human teleoperation demonstrations per task.
+# Model Overview
 
-**Paper**: *Cosmos Policy: Fine-Tuning Video Models for Visuomotor Control and Planning*
+## Description:
 
-### Key Features
+Cosmos-Policy-RoboCasa-Predict2-2B is a 2B-parameter robot manipulation policy model fine-tuned from the [NVIDIA Cosmos-Predict2-2B-Video2World](https://huggingface.co/nvidia/Cosmos-Predict2-2B-Video2World) video foundation model. This model achieves state-of-the-art performance on the RoboCasa simulation benchmark with a 67.1% average success rate across 24 kitchen manipulation tasks.
 
-- **Single-stage fine-tuning**: Adapted from pretrained video model with no architectural modifications
-- **Multimodal outputs**: Jointly predicts actions, future states, and values through unified video diffusion
-- **High performance**: 67.1% average success rate on RoboCasa with only 50 demos per task (vs. 300+ for other methods)
-- **Data efficiency**: Trained on significantly fewer demonstrations than prior state-of-the-art methods
+Key features:
 
-### Model Architecture
+* **Single-stage fine-tuning**: Adapted from pretrained video model with no architectural modifications
+* **Multimodal outputs**: Jointly predicts actions, future states, and values through unified video diffusion
+* **High performance**: 67.1% average success rate on RoboCasa with only 50 demos per task
 
-This model uses the same architecture as the base Cosmos-Predict2-2B model (a diffusion transformer with latent video diffusion). Please refer to the [base model card](https://huggingface.co/nvidia/Cosmos-Predict2-2B-Video2World) for detailed architecture specifications.
+Use cases:
+
+* Robotic manipulation and control in simulation environments
+* Imitation learning and policy learning for kitchen manipulation tasks
+* Vision-based robot learning with multiple camera viewpoints
+* Multi-task manipulation across diverse kitchen environments
+* Lifelong learning and transfer learning in robotics
+
+This model is for research and development only.
+
+**Model Developer**: NVIDIA
+
+## Model Versions
+
+Cosmos Policy models include the following:
+
+- [Cosmos-Policy-LIBERO-Predict2-2B](https://huggingface.co/nvidia/Cosmos-Policy-LIBERO-Predict2-2B): Given current state observations and a task description, generate action sequences, future state predictions, and value estimates for robot manipulation in simulated LIBERO environments.
+- [Cosmos-Policy-RoboCasa-Predict2-2B](https://huggingface.co/nvidia/Cosmos-Policy-RoboCasa-Predict2-2B): Given current state observations and a task description, generate action sequences, future state predictions, and value estimates for robot manipulation in simulated RoboCasa environments.
+- [Cosmos-Policy-ALOHA-Predict2-2B](https://huggingface.co/nvidia/Cosmos-Policy-ALOHA-Predict2-2B): Given current state observations and a task description, generate action sequences, future state predictions, and value estimates for robot manipulation in real-world ALOHA robot environments.
+- [Cosmos-Policy-ALOHA-Planning-Model-Predict2-2B](https://huggingface.co/nvidia/Cosmos-Policy-ALOHA-Planning-Model-Predict2-2B): Given current state observations, a task description, and action sequences, generate future state predictions and value estimates for robot manipulation in real-world ALOHA robot environments. (This checkpoint is meant to be deployed alongside Cosmos-Policy-ALOHA-Predict2-2B, not independently.)
+
+### License:
+
+This model is released under the  [NVIDIA One-Way Noncommercial License (NSCLv1)](https://github.com/NVlabs/HMAR/blob/main/LICENSE). For a custom license, please contact [cosmos-license@nvidia.com](mailto:cosmos-license@nvidia.com).
+
+Under the NVIDIA One-Way Noncommercial License (NSCLv1), NVIDIA confirms:
+
+* Models are not for commercial use.
+* NVIDIA does not claim ownership to any outputs generated using the Models or Derivative Models.
+
+### Deployment Geography:
+
+Global
+
+### Use Case:
+
+Physical AI: Robot manipulation and control, encompassing kitchen manipulation and imitation learning in simulation environments.
+
+### Release Date:
+
+GitHub [01/12/2026] via [https://github.com/nvlabs/cosmos-policy](https://github.com/nvlabs/cosmos-policy)
+
+Hugging Face [01/12/2026] via [https://huggingface.co/collections/nvidia/cosmos-policy](https://huggingface.co/collections/nvidia/cosmos-policy)
+
+## Model Architecture:
+
+Architecture Type: A diffusion transformer with latent video diffusion, fine-tuned from Cosmos-Predict2-2B-Video2World.
+
+Network Architecture: The model uses the same architecture as the base [Cosmos-Predict2-2B](https://huggingface.co/nvidia/Cosmos-Predict2-2B-Video2World) model (a diffusion transformer with latent video diffusion).
 
 **Key adaptation**: Actions, proprioceptive states, and values are encoded as latent frames and injected directly into the video model's latent diffusion sequence, enabling the model to generate these modalities alongside predicted future images.
 
-## Model Details
+**Number of model parameters:**
+
+2B (inherited from base model)
+
+## Input
+
+**Input Type(s)**: Text + Multi-view Images + Proprioceptive State
+
+**Input Format(s)**:
+
+* Text: String (natural language task description)
+* Images: RGB images from multiple camera views
+* Proprioception: Numerical array
+
+**Input Parameters**:
+
+* Text: One-dimensional (1D) - Task description (e.g., "open the right drawer")
+* Images: Two-dimensional (2D) - Left third-person camera (agentview_left): 224×224 RGB; Right third-person camera (agentview_right): 224×224 RGB; Wrist-mounted camera (eye-in-hand): 224×224 RGB
+* Proprioception: One-dimensional (1D) - 9-dimensional state (2 gripper joints + 3 end-effector position + 4 end-effector quaternion)
+
+**Other Properties Related to Input**:
+
+* Requires specific camera configuration (two third-person + wrist views)
+* Images resized to 224×224 pixels from original resolution
+* Trained exclusively for Franka Emika Panda robot arm in RoboCasa simulation environments
+
+## Output
+
+**Output Type(s)**: Action Sequence + Future State Predictions + Value Estimate
+
+**Output Format**:
+
+* Actions: Numerical array
+* Future states: Images + Proprioception
+* Value: Scalar
+
+**Output Parameters**:
 
-### Inputs
+* Action chunk: 32-timestep sequence of 7-dimensional actions (6-DoF end-effector control + 1 gripper)
+* Future robot proprioception: 9-dimensional state at timestep t+32
+* Future state images: Left third-person camera prediction (224×224 RGB), right third-person camera prediction (224×224 RGB), and wrist camera prediction (224×224 RGB) at timestep t+32
+* Future state value: Expected cumulative reward from future state (scalar)
 
-- **Current state images**: 
-  - Left third-person camera (agentview_left): Resized to 224x224 RGB
-  - Right third-person camera (agentview_right): Resized to 224x224 RGB
-  - Wrist-mounted camera (eye_in_hand): Resized to 224x224 RGB
-- **Robot proprioception**: 9-dimensional (2 gripper joints + 3 end-effector position + 4 end-effector quaternion)
-- **Task description**: Natural language text (e.g., "open the right drawer")
+**Other Properties Related to Output**:
 
-### Outputs
+* Action chunk size: 32 timesteps
+* Execution horizon: 16 timesteps (recommended; can be varied)
+* Denoising steps: 5 (configurable without retraining)
+* Noise level range: σ_min = 4.0, σ_max = 80.0
+* Generation mode: Parallel (action, future state, and value generated simultaneously)
 
-- **Action chunk**: 32-timestep sequence of 7-dimensional actions (6-DoF end-effector control + 1 gripper)
-- **Future robot proprioception**: 9-dimensional state at timestep t+32
-- **Future state images**: 
-  - Wrist camera prediction at timestep t+32
-  - Left third-person camera prediction at timestep t+32
-  - Right third-person camera prediction at timestep t+32
-- **Future state value**: Expected cumulative reward from future state
+Our AI models are designed and/or optimized to run on NVIDIA GPU-accelerated systems. By leveraging NVIDIA’s hardware (e.g. GPU cores) and software frameworks (e.g., CUDA libraries), the model achieves faster training and inference times compared to CPU-only solutions.
 
-### Training Details
+## Software Integration
+
+**Runtime Engine(s):**
+
+* [Transformers](https://github.com/huggingface/transformers)
+
+**Supported Hardware Microarchitecture Compatibility:**
+
+* NVIDIA Hopper (e.g., H100)
+
+**Note**: We have only tested doing inference with BF16 precision.
+
+**Operating System(s):**
+
+* Linux
+
+The integration of foundation and fine-tuned models into AI systems requires additional testing using use-case-specific data to ensure safe and effective deployment. Following the V-model methodology, iterative testing and validation at both unit and system levels are essential to mitigate risks, meet technical and functional requirements, and ensure compliance with safety and ethical standards before deployment.
+
+# Usage
+
+See [Cosmos Policy GitHub](http://github.com/NVlabs/cosmos-policy) for details.
+
+## Training and Evaluation Sections:
+
+### Training Datasets:
+
+**Data Collection Method**:
+
+* RoboCasa-Cosmos-Policy: Human - Human-teleoperated demonstrations recorded in simulation environment
+
+**Labeling Method**:
+
+* RoboCasa-Cosmos-Policy: Automated - Success/failure labels automatically determined by simulation environment evaluation; task descriptions from benchmark specification
+
+##### Properties:
 
 **Training Data**: [RoboCasa-Cosmos-Policy](https://huggingface.co/datasets/nvidia/RoboCasa-Cosmos-Policy) dataset
+
 - 24 kitchen manipulation tasks across 7 categories
-- 50 human teleoperation demonstrations per task
+- 50 demonstrations per task
 - Successful demonstrations used for policy training
 - All demonstrations (including failures) used for world model and value function training
 
 **Training Configuration**:
+
 - **Base model**: NVIDIA Cosmos-Predict2-2B-Video2World (`model-480p-16fps.pt`)
 - **Training steps**: 45,000 gradient steps
 - **Batch size**: 800 (global)
 - **GPUs**: 32 H100 GPUs
 - **Training time**: ~48 hours
 - **Optimization**: Full model fine-tuning (all weights updated)
-- **Action chunk size**: 32 timesteps (prediction horizon)
-- **Execution horizon**: 16 timesteps (recommended; can be varied)
-- **Image resolution**: 224x224 pixels
+- **Action chunk size**: 32 timesteps
+- **Image resolution**: 224×224 pixels
 
 **Training Objective**: The model is trained with a hybrid log-normal-uniform noise distribution (modified from the base model's log-normal distribution; see paper for details) to improve action prediction accuracy. Training batches are split 50/25/25 for policy, world model, and value function objectives, respectively.
 
-**Inference Settings**:
-- Denoising steps: 5 (note: this can be changed without retraining)
-- Noise level range: σ_min = 4.0, σ_max = 80.0
-- Generation mode: Parallel (action, future state, and value generated simultaneously)
-- Execution: First 16 timesteps of the 32-timestep action chunk are executed before requerying the policy
+### Evaluation Datasets:
 
-## Performance
+Data Collection Method: Not Applicable
 
-### RoboCasa Benchmark Results
+Labeling Method: Not Applicable
 
-| Method | # Training Demos per Task | Average Success Rate |
-|--------|---------------------------|---------------------|
-| GR00T-N1 | 300 | 49.6% |
-| UVA | 50 | 50.0% |
-| DP-VLA | 3,000 | 57.3% |
-| π0 | 300 | 62.5% |
-| GR00T-N1.5 | 300 | 64.1% |
-| Video Policy | 300 | 66.0% |
-| FLARE | 300 | 66.4% |
-| **Cosmos Policy (ours)** | **50** | **67.1%** |
+Properties: Not Applicable - We use the RoboCasa simulation environment for direct evaluations.
 
-Success rates are averaged over 50 trials per task (across 5 evaluation scenes with 10 trials each) and 3 random seeds (3,600 trials total).
+## Inference:
 
-### Key Achievements
+**Test Hardware:** H100, A100
 
-- Achieves state-of-the-art performance with 6× fewer demonstrations than most baselines
-- Generalizes to unseen object instances and kitchen styles
-- Demonstrates strong multi-task manipulation capabilities across diverse kitchen environments
+See [Cosmos Policy GitHub](http://github.com/NVlabs/cosmos-policy) for details.
 
-## Notes
+#### System Requirements and Performance
 
-- **Simulation only**: This checkpoint is trained and evaluated exclusively on RoboCasa simulation environments
-- **Single robot platform**: Trained only for the Franka Emika Panda robot arm
-- **Fixed camera setup**: Requires specific camera configuration (two third-person + one wrist views)
-- **Kitchen tasks**: Designed specifically for kitchen manipulation scenarios
+Inference with base Cosmos Policy only (i.e., no model-based planning):
 
-## Citation
+* 1 GPU with 6.8 GB VRAM for LIBERO sim benchmark tasks
+* 1 GPU with 8.9 GB VRAM for RoboCasa sim benchmark tasks
+* 1 GPU with 6.0 GB VRAM for ALOHA robot tasks
+
+#### Quality Benchmarks
+
+### RoboCasa Benchmark Results
 
-If you use this model, please cite the Cosmos Policy paper by Kim et al.
-<!-- ```bibtex
-# TODO: Add Cosmos Policy BibTeX
-``` -->
+| Method                         | # Training Demos per Task | Average Success Rate |
+| ------------------------------ | ------------------------- | -------------------- |
+| GR00T-N1                       | 300                       | 49.6%                |
+| UVA                            | 50                        | 50.0%                |
+| DP-VLA                         | 3,000                     | 57.3%                |
+| π0                            | 300                       | 62.5%                |
+| GR00T-N1.5                     | 300                       | 64.1%                |
+| Video Policy                   | 300                       | 66.0%                |
+| FLARE                          | 300                       | 66.4%                |
+| **Cosmos Policy (ours)** | **50**              | **67.1%**      |
 
-## License
+Success rates are averaged over 50 trials per task (across 5 evaluation scenes with 10 trials each) and 3 random seeds (3,600 trials total).
+
+## Ethical Considerations
+
+NVIDIA believes Trustworthy AI is a shared responsibility and we have established policies and practices to enable development for a wide array of AI applications. When downloaded or used in accordance with our terms of service, developers should work with their internal model team to ensure this model meets requirements for the relevant industry and use case and addresses unforeseen product misuse.
+
+Users are responsible for model inputs and outputs. Users are responsible for ensuring safe integration of this model, including implementing guardrails as well as other safety mechanisms, prior to deployment.
 
-Please refer to the [base model license](https://huggingface.co/nvidia/Cosmos-Predict2-2B-Video2World) for licensing information.
+Please report security vulnerabilities or NVIDIA AI Concerns [here](https://www.nvidia.com/en-us/support/submit-security-vulnerability/).
 
 ## Related Resources
 
 - **Base Model**: [Cosmos-Predict2-2B-Video2World](https://huggingface.co/nvidia/Cosmos-Predict2-2B-Video2World)
 - **Training Dataset**: [RoboCasa-Cosmos-Policy](https://huggingface.co/datasets/nvidia/RoboCasa-Cosmos-Policy)
 - **Paper**: *Cosmos Policy: Fine-Tuning Video Models for Visuomotor Control and Planning*
+- **Original RoboCasa**: [RoboCasa: Large-Scale Simulation of Everyday Tasks for Generalist Robots](https://arxiv.org/abs/2406.02523)
+
+## Citation
+
+If you use this model, please cite the Cosmos Policy paper:
+
+(Cosmos Policy BibTeX citation coming soon!)