StarVLA/WM4A-CosmoPredict-GR00T-LIBERO-4in1

WM4A CosmoPredict2-GR00T for LIBERO 4-in-1

A Vision-Language-Action (VLA) model from the StarVLA project, built on a Cosmos-Predict2-2B world model as the visual backbone, driving a GR00T-style DiT flow-matching action head (CosmoPredict2GR00T). The model is trained on the full LIBERO 4-in-1 benchmark (libero_10 + libero_goal + libero_object + libero_spatial combined).

CosmoPredict2GR00T is StarVLA's architecture that extracts visual world-model features from NVIDIA Cosmos-Predict2-2B (a video-to-world diffusion model) and feeds them into a cross-attention DiT flow-matching action head inspired by the GR00T N1 design:

1. Cosmos-Predict2 visual features — the last-layer activations of Cosmos-Predict2-2B-Video2World serve as rich spatiotemporal visual representations. 32 target vision tokens are extracted and passed to the action head.
2. Cross-attention flow-matching DiT — a 16-layer DiT-B with a 768-wide latent, 12 attention heads, interleaved self-attention, and adaptive LayerNorm. Its cross-attention conditioning width is aligned to the world-model hidden size at runtime; the separate state/action decoder MLP uses width 1,024.
3. Language conditioning via instruction tokens — the task instruction is tokenised and injected into the DiT cross-attention alongside the visual tokens; no separate VLM backbone is used.

---

Model Summary

Architecture	CosmoPredict2GR00T (Cosmos-Predict2 visual backbone + cross-attn FM DiT)
Visual backbone	Cosmos-Predict2-2B-Video2World
Action head	Cross-attention Flow-Matching DiT-B (16 layers, 768 latent width, 12 heads; decoder/state MLP width 1,024)
Action chunk	8 steps (+ 7 future-window steps)
Action / state dim	7 / 7; the compact config does not serialize the action convention
Image preprocessing	Delegated to the Cosmos world-model processor (obs_image_size: null)
Inference timesteps	4 (flow matching)
Checkpoint license	Not independently specified; see licensing note below
Codebase	starVLA/starVLA

---

Training Data

LIBERO 4-in-1 mixture (libero_all) — all four LIBERO task suites combined into a single training stream:

Suite	Tasks	Description
libero_10	10	Long-horizon tabletop manipulation
libero_goal	10	Goal-conditioned rearrangement
libero_object	10	Object-centric pick-and-place
libero_spatial	10	Spatially varied placement

• Action/state dimensions: 7 / 7. The compact config does not record whether the 7D action is delta pose or delta joint, so consumers must confirm the dataset/evaluator convention.
• Image observation: RGB input through the Cosmos processor. The compact config does not serialize a camera list or forced image resolution.
• Per-dataset normalisation statistics are stored in dataset_statistics.json.

---

Training Recipe

Total steps	80,000 configured; only the 50k weight is released
Warm-up steps	5,000
Per-device batch size	8
Hardware	GPU type/count is not recorded in the portable config
Precision	CUDA/bf16 is used by the framework; the exact training precision setup is not fully recorded
Optimizer	AdamW (β₁ = 0.9, β₂ = 0.95, ε = 1e-8, wd = 1e-8)
LR (base / VLM)	2.5e-5
LR (action head)	1e-4
LR scheduler	cosine_with_min_lr (min lr 1e-6)
Gradient clipping	1.0
Flow-matching noise	β-distribution (α=1.5, β=1.0), s = 0.999
Repeated diffusion steps	8
Frozen modules	Public config and archived launcher disagree; confirm before reproducing

The saved recipe is preserved in config.yaml. It contains environment-local model/data paths, and the archived launcher contains cluster-specific settings. Audit and sanitize those files rather than running the launcher unchanged.

---

Evaluation — LIBERO 4-in-1

Following the standard LIBERO evaluation protocol (50 trials per task per suite). Numbers are success rates (↑).

Step	libero_goal	libero_object	libero_spatial	Avg (3 suites)
30k	0.908	0.980	0.880	0.923
40k	0.948	0.990	0.884	0.941
50k	0.944	0.990	0.906	0.947

A libero_10 evaluation was attempted and its log is included, but the 50k run terminates after 413 of 500 episodes because the policy-server connection closes. It is therefore excluded from the aggregate rather than reported as a completed suite.

Best released checkpoint by the three completed suites: steps_50000_pytorch_model.pt — average 94.7 % across libero_goal / object / spatial.

For comparison with other StarVLA frameworks see the StarVLA Model Zoo.

---

Repository layout

.
├── README.md
├── config.yaml
├── dataset_statistics.json
├── summary.jsonl
├── run_libero_train.sh       # archived; audit and sanitize before use
├── slurm_script              # archived; audit and sanitize before use
├── checkpoints/
│   └── steps_50000_pytorch_model.pt
├── logs/
│   ├── libero_10/            # 30k/40k/50k attempts; 50k is incomplete
│   ├── libero_goal/
│   ├── libero_object/
│   └── libero_spatial/
└── wandb/                    # archived training logs

Only the 50k weight is downloadable; 30k and 40k appear only as evaluation-log names.

---

Download and evaluation

huggingface-cli download StarVLA/WM4A-CosmoPredict-GR00T-LIBERO-4in1 \
  --local-dir WM4A-CosmoPredict-GR00T-LIBERO-4in1

CKPT="$PWD/WM4A-CosmoPredict-GR00T-LIBERO-4in1/checkpoints/steps_50000_pytorch_model.pt"
python deployment/model_server/server_policy.py \
  --ckpt_path "$CKPT" --port 6694 --use_bf16 \
  --config_override framework.world_model.base_wm=nvidia/Cosmos-Predict2-2B-Video2World \
  --config_override framework.qwenvl.base_vlm=nvidia/Cosmos-Predict2-2B-Video2World

Before starting the server, map framework.world_model.base_wm in the saved config to a usable nvidia/Cosmos-Predict2-2B-Video2World snapshot. The framework.qwenvl field is a legacy compatibility field, not a separate Qwen VLM used by this architecture.

For the simulator side, follow the StarVLA LIBERO workflow.

---

Intended Use & Limitations

Intended use. Research on vision-language-action models, LIBERO tabletop manipulation benchmarks, and as a baseline for dual VLM + world-model conditioning architectures.

Out-of-scope / limitations. This model is trained exclusively on LIBERO simulation data with the saved 7D action/statistics contract. The compact config does not serialize whether that action is delta pose or delta joint, so the convention must be confirmed against the original dataset/evaluator before use. Real-robot transfer and cross-embodiment generalisation have not been evaluated. Performance may degrade on out-of-distribution scenes, objects, or instructions not present in the LIBERO training split.

Evidence and evaluation boundary

• Architecture and training fields were checked against config.yaml and the StarVLA CosmoPredict2GR00T implementation; the artifact list was checked against the Hub file tree.
• The three reported suite values are backed by the packaged per-suite logs. The incomplete libero_10 log is intentionally not converted into a final success rate.
• The repository releases one weight (50k), even though evaluation logs refer to 30k and 40k checkpoints.
• GPU type/count, a portable base-world-model path, and an unambiguous action convention are not recorded by the compact public config.
• Results cover LIBERO simulation only. Real-robot transfer, cross-embodiment generalization, and safety have not been established.
• The StarVLA code license does not by itself determine the license of this derived checkpoint. Access to the gated Cosmos base model and use of its derivatives remain subject to NVIDIA's Cosmos model terms; the checkpoint's release terms require maintainer confirmation.