Run end-to-end open-loop RL post-training of the Alpamayo VLM action model (Alpamayo 1 or 1.5) on the Physical AI AV (PAI) dataset using Cosmos-RL + GRPO. The agent collects a small set of choices from the user up front (Alpamayo version, reward mode — motion-only vs joint reasoning+motion, W&B preference, single-node vs multi-node, dataset/checkpoint paths), then drives the whole pipeline. Use when an agent must convert a released Alpamayo checkpoint into a training-ready one, curate a PAI mini subset, launch policy + rollout replicas via `cosmos-rl`, and export the resulting policy checkpoint back to a HuggingFace directory; when setting up the `a1x_rl` uv venv from scratch; when overriding TOML config (parallelism, replicas, reward weights, prefetch, optimizer); when diagnosing common RL failures (rollout-too-fast buffer growth, weight sync lag, prefetch misconfiguration, vLLM OOM, GRPO group collapse). Trigger keywords: alpamayo, alpamayo-1, alpamayo-1.5, alpamayo1, alpamayo1_5, rl, post-train, post-train
Instrucciones de origen · Vista previa de solo lectura
name
alpamayo1-x-rl
description
Run end-to-end open-loop RL post-training of the Alpamayo VLM action model (Alpamayo 1 or 1.5) on the Physical AI AV (PAI) dataset using Cosmos-RL + GRPO. The agent collects a small set of choices from the user up front (Alpamayo version, reward mode — motion-only vs joint reasoning+motion, W&B preference, single-node vs multi-node, dataset/checkpoint paths), then drives the whole pipeline. Use when an agent must convert a released Alpamayo checkpoint into a training-ready one, curate a PAI mini subset, launch policy + rollout replicas via `cosmos-rl`, and export the resulting policy checkpoint back to a HuggingFace directory; when setting up the `a1x_rl` uv venv from scratch; when overriding TOML config (parallelism, replicas, reward weights, prefetch, optimizer); when diagnosing common RL failures (rollout-too-fast buffer growth, weight sync lag, prefetch misconfiguration, vLLM OOM, GRPO group collapse). Trigger keywords: alpamayo, alpamayo-1, alpamayo-1.5, alpamayo1, alpamayo1_5, rl, post-train, post-training, RL post-training, grpo, group relative policy optimization, cosmos-rl, cosmos rl, reasoning vla, rvla, action token, trajectory diffusion frozen, pai, physical_ai_av, physicalai-autonomous-vehicles, reward, ade, comfort, lingo-judge, reasoning grader, chain-of-thought reward, chain of causation, ood_reasoning, hydra, omegaconf, toml, vllm, flash-attn, transformers, huggingface, Alpamayo-R1-10B, Alpamayo-1.5-10B, Cosmos-Reason2-8B, wandb, fsdp, dp_shard_size, n_init_replicas, train_batch_per_replica, rollout.n_generation, sync_weight_interval, pending rollouts, weight_version, prefetch, prefetch.capacity, convert_release_config_to_training, convert_cosmos_rl_checkpoint, curate_pai_samples, RLWrapperReasoningVLA.
license
Apache-2.0
metadata
{"author":"nvidia","version":"2026.05"}
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.
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)
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/.
Ask in a single round if your interface supports it; otherwise sequentially:
"Alpamayo version: 1 or 1.5?"
"Reward mode: motion (trajectory ADE + comfort only) or joint
(also grades chain-of-thought reasoning via Lingo-Judge)?"
"PAI chunks to download? (e.g. 3116 for motion-only; for joint,
I'll use --only-reasoning-chunks --num-reasoning-clips N — how many
clips, N?)"
(motion only) "How many clips to curate from that chunk? (16 for 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) or cluster
(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_generation completions
ranked by reward (default n_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 R spawns
P FSDP policy replicas + R vLLM rollout replicas. They communicate
through the controller. policy 1 / rollout 1 is 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 set WANDB_API_KEY
uncomment wandb logger, or drop wandb from the logger list.
Conversion is mandatory. You cannot point model_name_or_path
directly at the HF release dir; it has to be run through
convert_release_config_to_training.py first. That script remaps
_target_ paths in config.json to the training-time classes and
pulls the tokenizer from nvidia/Cosmos-Reason2-8B.
Export is mandatory if you want to use the trained model — Cosmos-RL
saves DTensor-sharded per-rank .pth files that from_pretrained can't
read directly. Run convert_cosmos_rl_checkpoint.py to 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):
the export PATH=... above only persists for the current shell. Either
append to ~/.bashrc once, 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):
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 ofALPAMAYO_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"
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-clips rows.
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-chunks pulls every chunk
that contains a sampled reasoning clip — not just the clips' own
footage. For --num-reasoning-clips 16 this 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=1 during a prior step, set it
back to 0 before download_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:
[policy.parallelism].dp_shard_size → 4 for local (1 node, ≥5 GPUs),
8 for cluster
[logging].logger:
If $USE_WANDB = yes: ["console", "wandb"], plus set
project_name = "$WANDB_PROJECT", experiment_name = "$WANDB_EXPERIMENT"
If $USE_WANDB = no: ["console"]
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.
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.
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:
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:
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 rollouts and
weight_version gap. 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
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
