Alpamayo 1.x RL Post-training (Cosmos-RL + GRPO)
This skill teaches an agent to RL-post-train the VLM backbone of
released Alpamayo models (the autoregressive-trajectory-token pathway) using
Cosmos-RL with GRPO, on the PAI (Physical AI Autonomous
Vehicles) dataset. Both Alpamayo 1 (nvidia/Alpamayo-R1-10B) and
Alpamayo 1.5 (nvidia/Alpamayo-1.5-10B) are supported. Validated
single-node on ≥5× 80 GB GPUs; large-scale recipe targets 80 nodes
(640 GPUs).
The recipe lives at recipes/alpamayo1_x_rl/ and is built on Cosmos-RL: one or more policy replicas train the model (FSDP), one or more rollout replicas run vLLM to generate completions, and a cosmos-controller dispatches rollouts, scores them, and syncs policy weights back to the rollouts.
Table of Contents
- When to use this skill
- Inputs to collect from the user (ask once, up front)
- Mental model — what gets trained, in what order
- Install —
a1x_rlvenv - Environment variables
- Model conversion — release ckpt → training ckpt
- Dataset — PAI mini subset
- Reasoning grader (only if joint reward)
- Launch RL training
- Monitor — controller log, healthy vs sick patterns
- Export — Cosmos-RL ckpt → HuggingFace ckpt
- Logging — W&B, console, offline
- Multi-node / cluster scale-up
- TOML override cheatsheet
- Common failure modes (and the fix)
- Additional resources
When to use this skill
| You want to… | Use |
|---|---|
| Run the full local RL pipeline (convert → curate → train → export) | The whole skill in order |
| Run RL on Alpamayo 1.5 with motion-only reward | Inputs → version=1.5, reward=motion; default TOML |
| Run RL with joint reasoning + motion reward | Inputs → reward=joint; the …_with_reasoning.toml config + Lingo-Judge grader |
| Just convert a release checkpoint into the training-ready format | Model conversion |
| Just export a trained Cosmos-RL policy ckpt → HF | Export |
| Scale a working local run to cluster | Multi-node / cluster scale-up |
| Diagnose a stuck or drifting RL run | Monitor + Common failure modes |
If you want to SFT the VLM (no RL, no reward, no rollout), this is the
wrong skill — see recipes/alpamayo1_sft/SKILL.md.
If you want to train the action expert (continuous flow-matching head),
that pathway is not covered by this RL pipeline — RL here only updates
the VLM backbone.
Inputs to collect from the user (ask once, up front)
Before any download, install, or training step, ask the user for the following. Confirm all answers before proceeding. If running in a context where you cannot ask the user, halt and report rather than guess.
| Input | Why you need it | Default if user has no preference |
|---|---|---|
Alpamayo version (1 / 1.5) |
Picks the converted base model (nvidia/Alpamayo-R1-10B vs nvidia/Alpamayo-1.5-10B) and the entry-script hydra_config_name (alpamayo1_rvla_rl_pai vs alpamayo1_5_rvla_rl_pai) |
1.5 |
Reward mode (motion / joint) |
motion → alpamayo_rvla_rl_local_test.toml + …_entry.py (ADE + comfort only). joint → …_local_test_with_reasoning.toml + …_reasoning_entry.py (also grades CoT via Lingo-Judge). |
motion |
| PAI chunks to download | Sets download_pai.py --chunk-ids and the curation chunk. Motion-only example uses 3116; joint uses --only-reasoning-chunks --num-reasoning-clips N. |
3116 for motion, --only-reasoning-chunks --num-reasoning-clips 16 for joint |
| Number of curated clips (motion-only path) | Sets curate_pai_samples.py --num-samples. Small for smoke tests (16), larger for real runs. |
16 |
W&B logging? (yes / no) |
If yes, also collect: WANDB_API_KEY, project_name, experiment_name. If no, set [logging].logger = ["console"] in TOML and leave WANDB_API_KEY unset. |
no (smoke) / yes (real run) |
Compute shape (local / cluster) |
local → 1 node, dp_shard_size=4, --policy 1 --rollout 1. cluster → adjust dp_shard_size=8, n_init_replicas, enable data_dispatch_as_rank_in_mesh. |
local for first run; cluster only after local works |
| Conditional: Lingo-Judge directory | Only if reward = joint. Path where wayveai/Lingo-Judge is (or should be) cached. The reward loads it with local_files_only=True, so it must be on disk. |
n/a |
| Conditional: existing trained policy ckpt | Only if the user is skipping training and just exporting. Path to <output_dir>/checkpoints/step_<N>/policy/. |
n/a |
| Output / log directories | ALPAMAYO_MODEL_DIR (converted ckpt), ALPAMAYO_PAI_LOCAL_DIR (dataset root), ALPAMAYO_LOG_DIR (Cosmos-RL logs), [train].output_dir in TOML (training outputs). |
Ask — no safe default |
Suggested question flow
Ask in a single round if your interface supports it; otherwise sequentially:
- "Alpamayo version:
1or1.5?" - "Reward mode:
motion(trajectory ADE + comfort only) orjoint(also grades chain-of-thought reasoning via Lingo-Judge)?" - "PAI chunks to download? (e.g.
3116for motion-only; forjoint, I'll use--only-reasoning-chunks --num-reasoning-clips N— how many clips, N?)" - (motion only) "How many clips to curate from that chunk? (
16for a fast smoke test)" - "Log to W&B? If yes, paste
WANDB_API_KEY, project name, experiment name. If no, I'll set logger to console only." - "Compute shape:
local(1 node, ≥5× 80 GB GPUs) orcluster(multi-node, will need SLURM/launcher integration)?" - (joint reward only) "Path where Lingo-Judge should be cached?
(default:
$YOUR_HOME/lingo_judge_model)" - "Where should artifacts live? (
ALPAMAYO_MODEL_DIR,ALPAMAYO_PAI_LOCAL_DIR,ALPAMAYO_LOG_DIR,[train].output_dir)"
Capture the answers verbatim
Save the answers as shell variables at the top of your run log — every later command references them:
export YOUR_HOME="/path/to/workspace"
export ALPAMAYO_WORKSPACE="$YOUR_HOME/alpamayo-recipes"
# from user answers
export ALPAMAYO_VERSION="1.5" # 1 | 1.5
export REWARD_MODE="motion" # motion | joint
export CHUNK_IDS="3116" # for motion path
export NUM_REASONING_CLIPS="16" # for joint path
export NUM_CURATED_CLIPS="16" # for motion path
export USE_WANDB="no" # yes | no
export WANDB_API_KEY="" # only if USE_WANDB=yes
export WANDB_PROJECT="Alpamayo_RL"
export WANDB_EXPERIMENT="ReasoningVLA_Post_Training"
export COMPUTE_SHAPE="local" # local | cluster
# paths (also exported for the recipe's own env-var contract)
export ALPAMAYO_MODEL_DIR="$YOUR_HOME/alpamayo_model_converted_from_hf"
export ALPAMAYO_PAI_LOCAL_DIR="$YOUR_HOME/PAI_mini"
export ALPAMAYO_PAI_REASONING_LOCAL_DIR="$YOUR_HOME/PAI_Reasoning_mini" # only if joint
export ALPAMAYO_LOG_DIR="$YOUR_HOME/alpamayo_cosmos_rl_job/logs"
export LINGO_JUDGE_DIR="$YOUR_HOME/lingo_judge_model" # only if joint
export TRAIN_OUTPUT_DIR="$YOUR_HOME/alpamayo_cosmos_rl_job/outputs"
After this, do not ask further questions — run the rest of the pipeline end-to-end. The only acceptable mid-run halt conditions are hard-stop errors (CUDA OOM, missing files, vLLM/Cosmos-RL crashes, GRPO collapse) or the user explicitly interrupting.
Mental model — what gets trained, in what order
Release HF ckpt (nvidia/Alpamayo-1.5-10B or -R1-10B)
│
│ convert_release_config_to_training.py
▼
$ALPAMAYO_MODEL_DIR (training-ready, target paths remapped)
│
│ download_pai.py + curate_pai_samples.py (or --only-reasoning-chunks)
▼
$ALPAMAYO_PAI_LOCAL_DIR (mini PAI subset + clip_index_mini.parquet
or clip_index_reasoning_mini.parquet)
│
▼
cosmos-rl --config <TOML> --policy P --rollout R \
<entry_script.py>
│
├── controller (rollout dispatch + reward + buffer + weight sync)
├── policy x P (FSDP training of VLM backbone) ← writes ckpts every N steps
└── rollout x R (vLLM serving + generation + reward scoring)
│
│ convert_cosmos_rl_checkpoint.py
▼
exported_model/ (standard HF dir, loadable via
RLWrapperReasoningVLA.from_pretrained)
Key facts an agent must internalise before running anything:
- VLM backbone only. RL here trains the autoregressive-trajectory-token pathway of the VLM. The action expert (flow-matching head) is not trained — its weights pass through unchanged and end up unused in the exported VLM-only ckpt.
- GRPO + groups. Each prompt produces
rollout.n_generationcompletions ranked by reward (defaultn_generation=12). Reward variance within a group is the training signal — if all completions get the same reward, GRPO can't learn anything from that group. - Two processes, one job.
cosmos-rl ... --policy P --rollout Rspawns P FSDP policy replicas + R vLLM rollout replicas. They communicate through the controller.policy 1 / rollout 1is the local-test default. - Required overrides. The user must (a) edit the TOML's
[train].output_dir,[policy].model_name_or_path, and[policy.parallelism].dp_shard_size, and (b) either setWANDB_API_KEY- uncomment
wandblogger, or drop wandb from the logger list.
- uncomment
- Conversion is mandatory. You cannot point
model_name_or_pathdirectly at the HF release dir; it has to be run throughconvert_release_config_to_training.pyfirst. That script remaps_target_paths inconfig.jsonto the training-time classes and pulls the tokenizer fromnvidia/Cosmos-Reason2-8B. - Export is mandatory if you want to use the trained model — Cosmos-RL
saves DTensor-sharded per-rank
.pthfiles thatfrom_pretrainedcan't read directly. Runconvert_cosmos_rl_checkpoint.pyto assemble them into a standard HF directory.
Install — a1x_rl venv
Single uv venv at recipes/alpamayo1_x_rl/a1x_rl/. flash-attn must build
against torch, so install in two uv sync passes:
# 1) Install uv (skip if already installed)
curl -LsSf https://astral.sh/uv/install.sh | sh
export PATH="$HOME/.local/bin:$PATH"
# 2) Provision the venv
export UV_CACHE_DIR="$YOUR_HOME/.cache/uv"
cd "$YOUR_HOME/alpamayo-recipes/recipes/alpamayo1_x_rl"
uv venv a1x_rl # if `a1x_rl/` already exists, skip this line
source a1x_rl/bin/activate
uv sync --active --no-install-package flash-attn # all deps minus flash-attn
uv sync --active # then build flash-attn against the installed torch
Non-interactive harnesses (agents, CI,
bash -c "..."per step): theexport PATH=...above only persists for the current shell. Either append to~/.bashrconce, or re-export it in every shell. Same caveat as the SFT recipe.
Verify the import contract after uv sync (this is the single most common
cause of opaque entry-script errors):
python -c "import alpamayo_r1, alpamayo, cosmos_rl, vllm; print('ok')"
python -c "from alpamayo1_x_rl.models.reasoning_vla.base_model import RLWrapperReasoningVLA; print('ok')"
python -c "from alpamayo.utils.checkpoint_utils import collect_targets, remap_targets; print('ok')"
Environment variables
Set these once per session (or ~/.bashrc). Treat them as the contract
the recipe scripts read:
# ── Paths ────────────────────────────────────────────────────────
export ALPAMAYO_WORKSPACE="$YOUR_HOME/alpamayo-recipes"
export ALPAMAYO_MODEL_DIR="$YOUR_HOME/alpamayo_model_converted_from_hf"
export ALPAMAYO_PAI_LOCAL_DIR="$YOUR_HOME/PAI_mini"
export ALPAMAYO_LOG_DIR="$YOUR_HOME/alpamayo_cosmos_rl_job/logs"
# ── Cache ────────────────────────────────────────────────────────
export HF_HOME="$YOUR_HOME/.cache/huggingface"
# ── Runtime ──────────────────────────────────────────────────────
export WANDB_API_KEY="<your_wandb_api_key>" # only if $USE_WANDB = yes
| Variable | Required | Purpose |
|---|---|---|
ALPAMAYO_WORKSPACE |
yes | Root of alpamayo-recipes |
ALPAMAYO_MODEL_DIR |
yes | Pre-trained converted Alpamayo model directory |
ALPAMAYO_PAI_LOCAL_DIR |
yes | PAI dataset root |
ALPAMAYO_LOG_DIR |
yes | Directory for Cosmos-RL logs |
UV_CACHE_DIR |
recommended | uv cache (set before uv venv) |
HF_HOME |
recommended | HuggingFace cache location |
HF_HUB_OFFLINE |
optional | 1 to skip Hub calls (air-gapped clusters / rate-limit fallbacks) |
TRANSFORMERS_OFFLINE |
optional | 1 alongside HF_HUB_OFFLINE |
WANDB_API_KEY |
conditional | Only if $USE_WANDB = yes |
ALPAMAYO_PAI_REASONING_LOCAL_DIR |
conditional | Required when $REWARD_MODE = joint — read by the reasoning entry script (…_reasoning_entry.py) instead of ALPAMAYO_PAI_LOCAL_DIR. The two are mutually exclusive: motion entry reads ALPAMAYO_PAI_LOCAL_DIR, joint entry reads ALPAMAYO_PAI_REASONING_LOCAL_DIR. Missing it → RuntimeError: Missing required env var ALPAMAYO_PAI_REASONING_LOCAL_DIR |
LINGO_JUDGE_DIR |
conditional | Only if $REWARD_MODE = joint |
HuggingFace auth
The Alpamayo model and PAI dataset are gated. The user must have accepted the licenses on HF web UI for whichever Alpamayo version they chose, plus PAI. Then:
hf auth login
In non-interactive shells, the cached token at
~/.cache/huggingface/token is read automatically. Source it into
$HF_TOKEN explicitly if a downstream tool expects the env var:
export HF_TOKEN=$(<~/.cache/huggingface/token).
Model conversion — release ckpt → training ckpt
Run this once per Alpamayo version. The script downloads the HF
release model + nvidia/Cosmos-Reason2-8B tokenizer/processor, remaps
_target_ paths in config.json to the training classes, and writes a
training-ready directory.
cd "$ALPAMAYO_WORKSPACE"
# motion or joint, version = 1.5 (default)
python scripts/convert_release_config_to_training.py \
--output-dir "$ALPAMAYO_MODEL_DIR"
# version = 1
python scripts/convert_release_config_to_training.py \
--alpamayo-model nvidia/Alpamayo-R1-10B \
--output-dir "$ALPAMAYO_MODEL_DIR"
Sanity-check the result:
ls "$ALPAMAYO_MODEL_DIR"
# expect: config.json (remapped), model.safetensors.index.json,
# model-0000{1..5}-of-00005.safetensors, tokenizer files,
# preprocessor_config.json, generation_config.json
Dataset — PAI mini subset
Two flows depending on $REWARD_MODE. Both write to
$ALPAMAYO_PAI_LOCAL_DIR (or $ALPAMAYO_PAI_REASONING_LOCAL_DIR for joint).
$REWARD_MODE = motion — ego-motion only
cd "$ALPAMAYO_WORKSPACE"
# 1) Download one chunk with four cameras + egomotion
python scripts/download_pai.py \
--chunk-ids "$CHUNK_IDS" \
--camera camera_front_wide_120fov camera_cross_left_120fov \
camera_cross_right_120fov camera_front_tele_30fov \
--calibration camera_intrinsics sensor_extrinsics vehicle_dimensions \
--labels egomotion \
--output-dir "$ALPAMAYO_PAI_LOCAL_DIR"
# 2) Curate a mini sample set (RL training reads clip_index_mini.parquet)
python scripts/curate_pai_samples.py \
--clip-index-path "$ALPAMAYO_PAI_LOCAL_DIR/clip_index.parquet" \
--chunk "$CHUNK_IDS" \
--num-samples "$NUM_CURATED_CLIPS" \
--output-path "$ALPAMAYO_PAI_LOCAL_DIR/clip_index_mini.parquet"
$REWARD_MODE = joint — reasoning-bearing clips
cd "$ALPAMAYO_WORKSPACE"
python scripts/download_pai.py --only-reasoning-chunks \
--num-reasoning-clips "$NUM_REASONING_CLIPS" \
--camera camera_front_wide_120fov camera_cross_left_120fov \
camera_cross_right_120fov camera_front_tele_30fov \
--calibration camera_intrinsics sensor_extrinsics vehicle_dimensions \
--labels egomotion egomotion.offline obstacle.offline \
--reasoning ood_reasoning.parquet \
--output-dir "$ALPAMAYO_PAI_REASONING_LOCAL_DIR"
After success, $ALPAMAYO_PAI_REASONING_LOCAL_DIR should contain:
clip_index.parquet— full PAI index (internal mapping).reasoning/ood_reasoning.parquet— full OOD reasoning table.clip_index_reasoning_mini.parquet— the mini index the RL config consumes; exactly--num-reasoning-clipsrows.camera/<subpart>/<subpart>.chunk_XXXX.zip,labels/<subpart>/<subpart>.chunk_XXXX.zip,calibration/<subpart>/...— only the chunks containing the sampled clips.
Disk budget heads-up.
--only-reasoning-chunkspulls every chunk that contains a sampled reasoning clip — not just the clips' own footage. For--num-reasoning-clips 16this typically lands around ~85–90 GB (vs ~5–10 GB for the motion-only path on a single chunk). Make sure$ALPAMAYO_PAI_REASONING_LOCAL_DIR's mount has the headroom before starting the download.
Verify:
ls "$ALPAMAYO_PAI_REASONING_LOCAL_DIR/clip_index_reasoning_mini.parquet" \
"$ALPAMAYO_PAI_REASONING_LOCAL_DIR/reasoning/ood_reasoning.parquet"
Heads-up: if you set
HF_HUB_OFFLINE=1during a prior step, set it back to0beforedownload_pai.py— the script needs Hub access.
Reasoning grader (only if joint reward)
Skip this section entirely when $REWARD_MODE = motion.
The default joint reward uses Lingo-Judge
as a learned reasoning grader. It must be cached locally because the reward
function loads it with local_files_only=True:
hf download wayveai/Lingo-Judge --local-dir "$LINGO_JUDGE_DIR"
Then in toml/alpamayo_rvla_rl_local_test_with_reasoning.toml,
set under [custom.alpamayo]:
| Field | Value |
|---|---|
reasoning_grader_type |
"lingo_judge" (default) |
reasoning_grading_model_path |
"$LINGO_JUDGE_DIR" (concrete path, not a shell var — TOML doesn't expand) |
reasoning_grading_device |
"auto", "cpu", or "cuda:0" depending on free GPU |
reward.reasoning_weight (under [custom.alpamayo.reward]) |
Tune to balance against trajectory reward |
If you want a custom grader, subclass BaseReasoningGrader in
utils/light_weight_reasoning_grading_model.py.
Launch RL training
The launcher is cosmos-rl, not torchrun. Two TOML configs + two
entry scripts pick the reward mode:
$REWARD_MODE |
TOML | Entry script |
|---|---|---|
motion |
toml/alpamayo_rvla_rl_local_test.toml |
models/reasoning_vla/alpamayo_cosmos_rl_post_training_entry.py |
joint |
toml/alpamayo_rvla_rl_local_test_with_reasoning.toml |
models/reasoning_vla/alpamayo_cosmos_rl_post_training_reasoning_entry.py |
Required TOML edits (before launching)
In whichever TOML you picked, set:
[train].output_dir→$TRAIN_OUTPUT_DIR[policy].model_name_or_path→$ALPAMAYO_MODEL_DIR[policy.parallelism].dp_shard_size→4for local (1 node, ≥5 GPUs),8for cluster[logging].logger:- If
$USE_WANDB = yes:["console", "wandb"], plus setproject_name = "$WANDB_PROJECT",experiment_name = "$WANDB_EXPERIMENT" - If
$USE_WANDB = no:["console"]
- If
Hydra config name (Alpamayo version selection)
Inside the entry script, the hydra_config_name argument selects the
PAI Hydra config:
"alpamayo1_5_rvla_rl_pai"for Alpamayo 1.5"alpamayo1_rvla_rl_pai"for Alpamayo 1
The shipped entry scripts default to the 1.5 config. If $ALPAMAYO_VERSION = 1, edit the entry script (or override via a sed step) to use
"alpamayo1_rvla_rl_pai" before launching.
Canonical launch (motion-only, local)
cd "$ALPAMAYO_WORKSPACE"
cosmos-rl \
--config recipes/alpamayo1_x_rl/toml/alpamayo_rvla_rl_local_test.toml \
--policy 1 \
--rollout 1 \
--log-dir "$ALPAMAYO_LOG_DIR" \
recipes/alpamayo1_x_rl/models/reasoning_vla/alpamayo_cosmos_rl_post_training_entry.py
Canonical launch (joint reward, local)
cd "$ALPAMAYO_WORKSPACE"
cosmos-rl \
--config recipes/alpamayo1_x_rl/toml/alpamayo_rvla_rl_local_test_with_reasoning.toml \
--policy 1 \
--rollout 1 \
--log-dir "$ALPAMAYO_LOG_DIR" \
recipes/alpamayo1_x_rl/models/reasoning_vla/alpamayo_cosmos_rl_post_training_reasoning_entry.py
--policy P --rollout R override n_init_replicas in the TOML. Use 1/1
for local smoke; scale on cluster (Multi-node).
Logs land at $ALPAMAYO_LOG_DIR/logs_<YYYYMMDD-HHMMSS>/:
| File | Process | Contains |
|---|---|---|
controller.log |
cosmos-controller | Rollout dispatch, reward stats, buffer (pending rollouts), weight sync events |
policy_<i>.log |
Policy replica i | Model loading, loss, grad norms, ckpt saving, per-rank dataset distribution |
rollout_<i>.log |
Rollout replica i | vLLM startup, generation throughput, weight receives, reward per sample |
Healthy progression (motion-only local)
Expect reward to climb (≈ -0.28 → -0.21) and trajectory L2 to drop
(≈ 1.66 → 1.34) over ~10 min on a single 8× H100 node. The joint
reward local run takes ~1.1 hr on a single 8× A100 node and should show
reasoning score climbing alongside trajectory L2 dropping.
Monitor — controller log, healthy vs sick patterns
The controller log is the canonical place to detect the most common RL
pathology: rollout/policy throughput imbalance. pending rollouts
is the buffer of completed-but-unconsumed rollouts.
Healthy
# controller.log — buffer stays low and oscillates
Stat: samples= 48 pending= 24
Stat: samples= 600 pending= 48
Stat: samples=1200 pending= 72
Stat: samples=2400 pending= 48
Sick: rollout too fast (most common)
Stat: samples= 24 pending= 24
Stat: samples= 600 pending= 264
Stat: samples=1200 pending= 552
Stat: samples=2400 pending= 984
[Controller] All rollouts have ended … 1104 remaining rollouts
# training at step 37/60 — remaining 23 steps consume stale, off-policy rollouts
When rollout outpaces policy, weight_version gap grows, training becomes
off-policy, quality degrades. Fix in this order:
- Enable prefetch (
[custom.alpamayo].prefetch.capacity > 0) — biggest single win, often 44 s → 5 s per step. - Reduce
[rollout].batch_sizeor[rollout].n_generation. - Add policy replicas (
[policy.parallelism].n_init_replicas). - Increase
[policy.parallelism].dp_shard_size. - Lower
epochor setmax_num_stepsto bound total rollouts.
Sick: rollout too slow
# pending repeatedly drops to 0 — policy idles waiting for data
Stat: samples= 600 pending= 0
Stat: samples= 624 pending= 0
Stat: samples= 648 pending= 0
Fix: add rollout replicas, or reduce policy throughput (e.g. smaller
train_batch_per_replica).
Target state
pending rollouts stays roughly stable. At cluster scale (64 policy /
128 rollout replicas, global batch 2560), buffer typically holds 4× global
batch. weight_version gap within a training batch should stay within a
few multiples of sync_weight_interval.
Export — Cosmos-RL ckpt → HuggingFace ckpt
Cosmos-RL writes per-rank DTensor shards at
$TRAIN_OUTPUT_DIR/checkpoints/step_<N>/policy/model_rank_*.pth.
These cannot be loaded by ReasoningVLA.from_pretrained. Convert:
cd "$ALPAMAYO_WORKSPACE"
# pick the step (latest by default)
STEP_DIR=$(ls -d "$TRAIN_OUTPUT_DIR/checkpoints"/step_* | sort -V | tail -1)
echo "exporting: $STEP_DIR"
python scripts/convert_cosmos_rl_checkpoint.py \
--cosmos-policy-ckpt "$STEP_DIR/policy" \
--base-hf-ckpt "$ALPAMAYO_MODEL_DIR" \
--output-dir "$YOUR_HOME/alpamayo_cosmos_rl_job/exported_model"
--cosmos-policy-ckpt— directory withmodel_rank_*.pthfiles.--base-hf-ckpt— the training-ready dir from Model conversion. Config, tokenizer, processor are copied from here.--output-dir— where the HF-style export lands.
Load the exported model for inference:
from alpamayo1_x_rl.models.reasoning_vla.base_model import RLWrapperReasoningVLA
model = RLWrapperReasoningVLA.from_pretrained(
"<output-dir>"
)
End-to-end inference / visualisation example: notebooks/rl_checkpoint_inference.ipynb.
The exported ckpt contains only the VLM backbone. Loading it back into a full Alpamayo (1 or 1.5) directory works with non-strict weight loading — the action-expert weights remain randomly initialised.
Logging — W&B, console, offline
Logger selection lives in [logging].logger in the TOML.
$USE_WANDB = no
[logging]
logger = ["console"]
No WANDB_API_KEY needed. For belt-and-braces (e.g. machines where
wandb login has been run previously), prefix the launch:
WANDB_MODE=disabled cosmos-rl --config ... <entry_script.py>
$USE_WANDB = yes
[logging]
logger = ["console", "wandb"]
project_name = "Alpamayo_RL" # = $WANDB_PROJECT
experiment_name = "ReasoningVLA_Post_Training" # = $WANDB_EXPERIMENT
Plus export WANDB_API_KEY="..." before launch. The token must have
write access to project_name. A 403 (upsertBucket … permission denied)
means the account isn't in that team/project — fix the project name or
the key. Fail fast and re-prompt the user for a working key; don't retry
blindly.
Multi-node / cluster scale-up
After the local run works, two TOML changes are mandatory for cluster:
[policy.parallelism].dp_shard_size = 8— FSDP shards the model across 8 GPUs per replica.n_policy_replicas × dp_shard_size= total policy GPUs.[train.train_policy].data_dispatch_as_rank_in_mesh = true— rank-based dataset dispatch so each policy replica consumes a stable, non-overlapping shard. Without this, replicas can train on duplicate samples. Also required for the prefetch path.
Recommended cluster-scale settings (these are the values used to RL post-train Alpamayo 1.5):
| Parameter | Local | Cluster |
|---|---|---|
policy.parallelism.n_init_replicas |
1 | 64 |
rollout.parallelism.n_init_replicas |
1 | 128 |
train.train_batch_per_replica |
48 | 40 |
train.optm_lr |
2e-6 | 2e-6 |
train.sync_weight_interval |
2 | 5 |
rollout.batch_size |
2 | 6 |
rollout.n_generation |
12 | 12 |
custom.alpamayo.prefetch.capacity |
16 | 128 |
This gives a global batch of 64 × 40 = 2560 samples/step on
512 policy GPUs + 128 rollout GPUs = 640 GPUs (80 nodes). Start
moderate (e.g. 4 policy / 8 rollout) and watch pending rollouts before
scaling up. SLURM launch instructions:
https://nvidia-cosmos.github.io/cosmos-rl/multinodes/overview.html.
Recommended workflow when iterating on a new reward / data / model
- Overfit on 1 sample. Confirms reward / data / model are wired correctly. Reward should climb.
- Overfit on 16–32 samples on one node. Tune reward weights, LR,
n_generation. Watch rollout/policy balance. - Scale to multi-node. Monitor
pending rolloutsandweight_versiongap. Start at moderate global batch (~320) and scale up if reward variance is too high. - Iterate on the reward. RL optimises whatever you measure. If behaviour isn't improving, revisit reward design before scaling.
TOML override cheatsheet
TOML edits happen in the file directly — Cosmos-RL doesn't accept
Hydra-style dotted CLI overrides for these. The --policy N / --rollout N
CLI flags override n_init_replicas for both replica types.
| TOML path | What it does |
|---|---|
[train].output_dir |
Where training outputs / checkpoints land |
[train].epoch |
Total epochs over the curated samples |
[train].max_num_steps (optional) |
Hard cap on training steps |
[train].train_batch_per_replica |
Samples consumed per replica per step. Global batch = n_init_replicas × this |
[train].sync_weight_interval |
Sync policy weights to rollouts every N steps |
[train].optm_lr |
Optimizer LR (default 2e-6 — RL fine-tune territory) |
[policy].model_name_or_path |
$ALPAMAYO_MODEL_DIR |
[policy].model_max_length |
Token cap for prompt + completion |
[policy].model_gradient_checkpointing |
true for memory savings |
[policy.parallelism].dp_shard_size |
FSDP shard degree for the policy (4 local / 8 cluster). Distinct from [rollout.parallelism].dp_shard_size, which stays 1 for local — the policy + rollout values are summed when the launcher places replicas on GPUs |
[rollout.parallelism].dp_shard_size |
GPUs per rollout replica (default 1). Increase only if a single rollout replica needs to shard the vLLM engine across multiple GPUs |
[policy.parallelism].n_init_replicas |
Number of policy replicas (overridden by --policy) |
[rollout].n_generation |
Completions per prompt (the GRPO group size). Lower if rollout too slow / OOM |
[rollout].batch_size |
Prompts per rollout batch |
[rollout].gpu_memory_utilization |
vLLM memory fraction (default 0.8) |
[rollout.sampling_config].temperature / .top_p / .repetition_penalty |
Generation sampling |
[custom.alpamayo].prefetch.capacity |
Prefetch cache size. Set to train_batch_per_replica × replicas_per_node. <= 0 disables |
[custom.alpamayo].prefetch.num_workers |
Prefetch worker threads (default 5) |
[custom.alpamayo.reward].traj_l2_weight |
Weight on ADE penalty in the trajectory reward |
[custom.alpamayo.reward].comfort_weight |
Weight on comfort score |
[custom.alpamayo.reward].reasoning_weight |
Weight on reasoning grader (joint reward only) |
[custom.alpamayo].reasoning_grading_model_path |
Lingo-Judge dir (joint reward only) |
[custom.alpamayo].reasoning_grading_device |
"auto" / "cpu" / "cuda:0" |
[logging].logger |
["console"] or ["console", "wandb"] |
[logging].project_name / .experiment_name |
W&B project / run name |
One-sample overfit smoke: set num_samples=1 in the curate step,
epoch=200, n_generation=8, train_batch_per_replica=1. Watch reward
climb on the same single clip — that's the wiring check before scaling.
Common failure modes (and the fix)
| Symptom | Root cause | Fix |
|---|---|---|
Hydra CLI override <key>=None is silently parsed as the string "None" (not Python None) — common symptom: a path-typed config field gets "/<dir>/None" joined onto it and the dataset complains the file doesn't exist |
YAML / Hydra grammar: bare None is a string token |
Use the YAML null literal: <key>=null. To drop the key entirely, use the delete prefix: ~<key>. Same applies to any Hydra override you customize on the CLI or inside an entry script's hydra_overrides=[…] |
cosmos-rl dies during launch with RuntimeError: Replica with GPUs larger than 8 occurs but not on Lepton job, please specify --node-ip-list ... |
Misleading message — actual cause is insufficient GPUs on the host. The launcher hits this when nvidia-smi -L | wc -l is less than (policy.n_init_replicas × policy.dp_shard_size) + (rollout.n_init_replicas × rollout.dp_shard_size). For the shipped local-test TOML that's 1×4 + 1×1 = 5. Don't chase --node-ip-list; that path is for genuine multi-node setups |
nvidia-smi -L | wc -l to confirm the count, then either run on a host with enough GPUs or temporarily lower policy.dp_shard_size (note: dropping below 4 may not produce a working model — this is purely for smoke-testing the wiring, not training) |
convert_release_config_to_training.py fails on HF download |
Gated model — license not accepted, or HF_TOKEN unset / wrong |
Accept license on HF web UI for the chosen Alpamayo version; hf auth login; confirm with huggingface-cli whoami |
download_pai.py errors with HF rate limit |
Hub throttling | Re-run with backoff, or set HF_HUB_OFFLINE=1 TRANSFORMERS_OFFLINE=1 once the data is cached locally |
Joint reward run errors Lingo-Judge not found (local_files_only=True) |
Grader not cached | hf download wayveai/Lingo-Judge --local-dir "$LINGO_JUDGE_DIR"; set reasoning_grading_model_path in the TOML |
clip_index_mini.parquet / clip_index_reasoning_mini.parquet not found |
Skipped the curation step, or used wrong dataset path | Re-run curate_pai_samples.py (motion) or re-download with --only-reasoning-chunks --num-reasoning-clips N (joint) |
Cosmos-RL crashes with Could not load model_rank_*.pth from a trained ckpt |
Trying to load policy ckpt directly | Use convert_cosmos_rl_checkpoint.py to assemble it into a HF dir first |
| vLLM rollout OOM on startup | gpu_memory_utilization too high, or model larger than rollout GPU |
Drop [rollout].gpu_memory_utilization (e.g. 0.7), or reduce [rollout].n_generation / model_max_length |
| Policy OOM | Insufficient FSDP sharding or grad-ckpt off | Increase dp_shard_size, ensure model_gradient_checkpointing = true, lower train_batch_per_replica |
pending rollouts grows monotonically |
Rollout too fast vs policy | Enable prefetch, drop rollout.batch_size/n_generation, add policy replicas. See Monitor |
pending rollouts repeatedly hits 0 |
Rollout too slow | Add rollout replicas, or reduce policy throughput |
| Reward variance ≈ 0 per group | GRPO group collapse (all n_generation completions get the same reward) |
Raise sampling temperature, lower repetition_penalty, or design a denser reward |
W&B 403 upsertBucket … permission denied |
API key not authorized for project_name/team |
Update project_name to one you own, or use a different WANDB_API_KEY; or set WANDB_MODE=disabled |
| Multi-node run trains on duplicate samples | data_dispatch_as_rank_in_mesh left at false |
Set [train.train_policy].data_dispatch_as_rank_in_mesh = true for any multi-node run |
| Wrong Alpamayo version loaded (1 vs 1.5 mismatch) | Entry script's hydra_config_name doesn't match $ALPAMAYO_VERSION |
Edit entry script: "alpamayo1_5_rvla_rl_pai" for 1.5, "alpamayo1_rvla_rl_pai" for 1 |
Additional resources
- Recipe README (human-facing): README.md
- Entry scripts: models/reasoning_vla/alpamayo_cosmos_rl_post_training_entry.py, models/reasoning_vla/alpamayo_cosmos_rl_post_training_reasoning_entry.py
- TOML configs: toml/alpamayo_rvla_rl_local_test.toml, toml/alpamayo_rvla_rl_local_test_with_reasoning.toml
- Hydra (PAI) configs: hydra_configs/alpamayo1_5_rvla_rl_pai.yaml, hydra_configs/alpamayo1_rvla_rl_pai.yaml
- Reward implementations: rewards/aggregated_reward.py (motion), rewards/aggregated_reward_with_reasoning.py (joint), building blocks: rewards/traj_reward.py, rewards/comfort_reward.py
- Shared launcher: launcher.py
- Prefetch / shared-memory dataset: prefetch/server.py, prefetch/dataset.py
- Reasoning grader (Lingo-Judge wrapper): utils/light_weight_reasoning_grading_model.py
- Conversion scripts:
../../scripts/convert_release_config_to_training.py,../../scripts/convert_cosmos_rl_checkpoint.py,../../scripts/curate_pai_samples.py,../../scripts/download_pai.py - Inference notebook: notebooks/rl_checkpoint_inference.ipynb
- Cosmos-RL framework: https://github.com/NVIDIA/Cosmos-RL
- Cosmos-RL multi-node docs: https://nvidia-cosmos.github.io/cosmos-rl/multinodes/overview.html
- GRPO paper: https://arxiv.org/abs/2402.03300
- PAI dataset: https://huggingface.co/datasets/nvidia/PhysicalAI-Autonomous-Vehicles
- Pretrained Alpamayo-R1-10B (v1): https://huggingface.co/nvidia/Alpamayo-R1-10B
- Pretrained Alpamayo-1.5-10B: https://huggingface.co/nvidia/Alpamayo-1.5-10B
- Lingo-Judge reasoning grader: https://huggingface.co/wayveai/Lingo-Judge
- Companion SFT skill:
recipes/alpamayo1_sft/SKILL.md