cosmos-policy-libero-all4suite

CosmosPolicy · LIBERO-Goal (10 k steps)

World-model-native Vision-Language-Action (VLA) checkpoint released with the AlphaBrain framework. Provided for direct download and evaluation — no retraining needed.

A CosmosPolicy model: a 3-D DiT (diffusion transformer) that predicts action chunks in the latent space of a world model, conditioned on T5 text embeddings and the robot's proprioceptive state. This is the 10 000-step checkpoint from the cosmos-8gpu-bs20-fix2-0408 run on the LIBERO-Goal task suite (chunk size 16).

Overview

Files

├── README.md                             model card
├── framework_config.yaml                 AlphaBrain / CosmosPolicy training config
├── config.json                           minimal HF-style metadata (arch, dims, step)
├── model.safetensors                     framework checkpoint (3.9 GB) — primary load path
├── cosmos_dit.pt                         DiT-only PyTorch tensor file (3.7 GB) — alt loader / resume
├── libero_dataset_statistics.json        action normalisation statistics
├── libero_t5_embeddings.pkl              precomputed T5 embeddings for LIBERO prompts (~41 MB)
└── resume_meta.json                      training resume metadata

Usage

git clone https://github.com/AlphaBrainGroup/AlphaBrain.git
cd AlphaBrain
pip install -e .

huggingface-cli download AlphaBrainGroup/cosmos-policy-libero-goal \
    --local-dir ./cosmos_policy_ckpt

# Launch the WebSocket inference server (CosmosPolicy route)
python deployment/model_server/server_policy.py \
    --ckpt_path ./cosmos_policy_ckpt --port 10093 --use_bf16

For the LIBERO-Goal evaluation pipeline, see benchmarks/LIBERO/eval/.

Reproduction

bash scripts/run_world_model/train/run_cosmos_policy.sh

The shipped framework_config.yaml is the exact configuration used for this checkpoint. Expect multi-day training on 8 × A800 80 GB for the full 40 000-step schedule; the 10 k release is a useful mid-training snapshot with good LIBERO-Goal performance and a manageable download footprint.

Notes

• Dual checkpoint format. model.safetensors is the canonical framework-loadable checkpoint used by server_policy.py; cosmos_dit.pt is the raw DiT tensor file kept for resume / selective loading. Most users only need model.safetensors.
• T5 embeddings are shipped so users do not need to re-run the T5 text encoder at inference time — load the .pkl directly.

License

MIT — see the parent repository.

Citation

@misc{alphabrain2026,
  title  = {AlphaBrain: A Modular Open-Source Framework for Embodied Intelligence Research},
  author = {AlphaBrain Team},
  year   = {2026},
  url    = {https://github.com/AlphaBrainGroup/AlphaBrain}
}