Metadata-Version: 2.4
Name: cosmos-rl
Version: 0.4.3
Summary: Cosmos-RL is a flexible and scalable Reinforcement Learning framework specialized for Physical AI applications.
Project-URL: Homepage, https://github.com/nvidia-cosmos/cosmos-rl
Requires-Python: >=3.10
Description-Content-Type: text/markdown
License-File: LICENSE
Requires-Dist: fastapi
Requires-Dist: uvicorn
Requires-Dist: transformers!=4.52.*,!=4.53.*,<5.0.0,>=4.51.1
Requires-Dist: datasets>=4.4.1
Requires-Dist: pydantic
Requires-Dist: math_verify
Requires-Dist: nest_asyncio
Requires-Dist: pybind11[global]
Requires-Dist: torch>=2.6.0
Requires-Dist: torchvision>=0.21.0
Requires-Dist: torchao>=0.10.0
Requires-Dist: tensordict==0.7.2
Requires-Dist: toml
Requires-Dist: StrEnum
Requires-Dist: redis
Requires-Dist: msgpack
Requires-Dist: qwen_vl_utils
Requires-Dist: pynvml
Requires-Dist: boto3
Requires-Dist: modelscope
Requires-Dist: cloudpickle
Requires-Dist: click
Requires-Dist: rich
Requires-Dist: pyyaml
Requires-Dist: blobfile
Requires-Dist: nvidia-nccl-cu12>=2.26.2
Requires-Dist: liger_kernel
Requires-Dist: accelerate
Requires-Dist: kernels
Requires-Dist: diffusers>=0.36.0
Requires-Dist: peft
Requires-Dist: zmq>=0.0.0
Requires-Dist: nvidia-ml-py>=13.590.48
Provides-Extra: wfm
Requires-Dist: termcolor; extra == "wfm"
Requires-Dist: aioboto3; extra == "wfm"
Requires-Dist: aiobotocore; extra == "wfm"
Requires-Dist: imageio; extra == "wfm"
Requires-Dist: imageio-ffmpeg; extra == "wfm"
Requires-Dist: webdataset; extra == "wfm"
Requires-Dist: decord; extra == "wfm"
Requires-Dist: trimesh; extra == "wfm"
Requires-Dist: opencv-python-headless; extra == "wfm"
Requires-Dist: better_profanity; extra == "wfm"
Requires-Dist: nltk; extra == "wfm"
Requires-Dist: retinaface-py; extra == "wfm"
Provides-Extra: vla
Requires-Dist: bddl; extra == "vla"
Requires-Dist: easydict; extra == "vla"
Requires-Dist: einops; extra == "vla"
Requires-Dist: future; extra == "vla"
Requires-Dist: gym; extra == "vla"
Requires-Dist: gymnasium; extra == "vla"
Requires-Dist: hydra-core; extra == "vla"
Requires-Dist: libero; extra == "vla"
Requires-Dist: math-verify; extra == "vla"
Requires-Dist: matplotlib; extra == "vla"
Requires-Dist: numpy<2; extra == "vla"
Requires-Dist: opencv-python; extra == "vla"
Requires-Dist: robomimic; extra == "vla"
Requires-Dist: robosuite==1.4; extra == "vla"
Requires-Dist: thop; extra == "vla"
Requires-Dist: timm<1; extra == "vla"
Requires-Dist: wandb>0.13.1; extra == "vla"
Requires-Dist: transformers<5.0.0; extra == "vla"
Requires-Dist: sentencepiece; extra == "vla"
Provides-Extra: rl
Requires-Dist: vllm>=0.8.5; extra == "rl"
Requires-Dist: flash-attn>=2.8.3; extra == "rl"
Provides-Extra: all
Requires-Dist: cosmos_rl[rl,vla,wfm]; extra == "all"
Dynamic: license-file

<p align="center">
    <img src="https://raw.githubusercontent.com/nvidia-cosmos/cosmos-rl/main/assets/nvidia-cosmos-header.png" alt="NVIDIA Cosmos Header">
</p>


## Getting Started

Cosmos-RL is a flexible and scalable Reinforcement Learning framework specialized for Physical AI applications.

[Documentation](https://nvidia-cosmos.github.io/cosmos-rl).

## System Architecture
Cosmos-RL provides toolchain to enable large scale RL training workload with following features:
1. **Parallelism**
    - Tensor Parallelism
    - Sequence Parallelism
    - Context Parallelism
    - FSDP Parallelism
    - Pipeline Parallelism
2. **Fully asynchronous (replicas specialization)**
    - Policy (Consumer): Replicas of training instances
    - Rollout (Producer): Replicas of generation engines
    - Low-precision training (FP8) and rollout (FP8 & FP4) support
3. **Single-Controller Architecture**
    - Efficient messaging system (e.g., `weight-sync`, `rollout`, `evaluate`) to coordinate policy and rollout replicas
    - Dynamic NCCL Process Groups for on-the-fly GPU [un]registration to enable fault-tolerant and elastic large-scale RL training

![Policy-Rollout-Controller Decoupled Architecture](https://raw.githubusercontent.com/nvidia-cosmos/cosmos-rl/main/assets/rl_infra.svg)

## License and Contact

This project will download and install additional third-party open source software projects. Review the license terms of these open source projects before use.

NVIDIA Cosmos source code is released under the [Apache 2 License](https://www.apache.org/licenses/LICENSE-2.0).

NVIDIA Cosmos models are released under the [NVIDIA Open Model License](https://www.nvidia.com/en-us/agreements/enterprise-software/nvidia-open-model-license). For a custom license, please contact [cosmos-license@nvidia.com](mailto:cosmos-license@nvidia.com).
