AW-0: Agentic World Action Models

AW-0 is an agentic World Action Model: alongside the imagined future video and the action, it writes its plan as a short program — and denoises all three together.

Abstract. Recent world action models provide visual foresight and future prediction for robot manipulation, but lack explicit mechanisms for compositional skill reuse. Agentic robotics combines code-based planning and reasoning with reusable skills, but often lacks predictive grounding in physical dynamics.

We introduce AW-0, an agentic world action model framework that integrates code-based planning with joint video and action prediction. We jointly model code, future video, and robot actions as three coupled generative streams within a unified mixture-of-transformers architecture. Within this architecture a code expert generates parameterized skill routines and control logic while video and action experts generate future visual states and low-level motor commands. Code serves two complementary roles: it structures multi-step prediction and control within the world model and represents skills as reusable programs for subsequent agentic execution. We further introduce an automated curation pipeline that labels robot and human video trajectories with aligned programs and consolidates recurring routines into reusable skill libraries. Evaluations on RoboTwin 2.0 and physical bimanual manipulation tasks show improved task performance over VLA and WAM baselines, alongside gains in language following, semantic reasoning and robustness, demonstrating that agentic code enables connection between planning, prediction and execution within WAMs.

AW-0 couples three transformer experts in one Mixture-of-Transformers with per-layer joint self-attention: a video world model (Cosmos-Predict2.5-2B, Wan2.2 VAE) that imagines the future, a code expert (SDAR-1.7B body) that writes the plan as a program over a skill API, and a diffusion action head. The program is never executed by the policy: it is a symbolic plan that the world model and the action head attend to. This repository is agenticwam — named for the model family rather than one model, so later work can build on it.

• Code as the plan. The code expert reads the observation through a frozen Qwen3-VL-2B vision tower and the instruction; the imagined future video reads the code; the action reads everything.
• Joint co-denoising. Future video, code and action are denoised together in one loop; the code canvas is revealed by a confidence-ordered masked decode. Autoregressive and block-diffusion code are swappable.
• One model across embodiments. Robot arms with grippers, dexterous hands and human hands share a unified 64-D action space; RoboTwin 2.0 uses its native 14-D joints.
• Programs for every frame. A VLM labeling pipeline turns demonstrations into per-frame programs against frozen skill libraries for robot grippers and human hands.
• Full pipeline. Pretraining (robot + egocentric human video), post-training and deployment on real robots (YAM, Tianji + Sharpa), and the RoboTwin 2.0 benchmark.

The model trains in two stages — large-scale pretraining on robot (yam) and egocentric human (ego) data → task-specific post-training on a real robot or in RoboTwin. Every recipe can also run without the code expert — model.use_code_expert=false builds the same model with the video and action experts only.

This repository ships the complete code, pretraining included. We release the pretrained checkpoint (below), so you can start directly at post-training. The pretraining corpora and the Cosmos-Predict2.5-2B video checkpoint that pretraining starts from are not part of this release.

The checkpoint holds every trained weight. The frozen public components load from their own releases: the Wan2.2 VAE (Wan-AI/Wan2.2-TI2V-5B), the instruction encoder (nvidia/Cosmos-Reason1-7B), the code-expert body (JetLM/SDAR-1.7B-Chat) and the vision tower (Qwen/Qwen3-VL-2B-Instruct) — docs/INSTALL.md §3.

The repository ships Claude Code skills in .claude/skills/: runbooks that walk an agent through each workflow with the scripts in this repository, checking the result of every step and stopping for confirmation before anything that costs money (VLM calls) or holds GPUs for long. Start claude in the repository root and ask in plain words, or call a skill by name:

cd agentic-wam && claude
> /posttrain-annotation   label /data/fold_towels (YAM) with the VLM pipeline, budget 200 episodes

The RoboTwin program-authoring skills (robotwin-fix-loop, robotwin-evosearch, robotwin-collect) are listed in .claude/README.md and docs/ROBOTWIN.md §2.

agentic-wam/
├── src/agenticwam/
│   ├── models/cosmos_mot/      agenticwam_model.py (losses, masks, joint inference) · generic_mot.py (joint attention)
│   │                           code_vision.py (Qwen3-VL prefix) · build.py (model factory)
│   ├── models/cosmos/          Cosmos DiT, video tokenizer, Wan2.2 VAE, flow-matching schedulers
│   ├── models/experts/         code_expert.py (joint / blockdiff / ar) · action_dit.py · shared blocks
│   ├── datasets/lerobot/       robot_video_dataset.py · transforms/ (yam, tianji_sharpa, ego → 64-D)
│   └── trainer.py · runtime.py
├── configs/
│   ├── paths/default.yaml      every weights / data location (edit this first)
│   ├── model/                  agenticwam_cosmos.yaml — the one model config
│   ├── real_pre_train/         pretrain_magic_soup (task + data)
│   ├── real_post_train/        posttrain_yam · posttrain_sharpa (task + general data configs)
│   └── robotwin/               robotwin_wam · robotwin_hybrid, benchmark definition, eval presets
├── scripts/                    train_real_{yam,sharpa}.sh · train_robotwin.sh · launch_train.py · preflight_sft.py
│                               {pre,post}train_annotation*.sh · precompute_cosmos_text_embeds.py
├── codegen/
│   ├── pretrain/               VLM labeling engine (pipeline/) + frozen skill libraries (yam, ego)
│   ├── posttrain/              real-robot labeling tools: sub-task splits, flat caches, alignment checks
│   └── sim/                    RoboTwin program authoring + code-labeled demonstration collection
├── robotwin/                   deploy/ (the policy RoboTwin calls) · dataprep/ · eval/ (benchmark) · hybrid_eval/
├── deploy/                     real-robot model server and robot-side client
├── utils/                      dataset tools (stats, embedding checks) · yam_fk_ik/ (YAM kinematics)
├── .claude/skills/             Claude Code skills: setup, data prep, code annotation, training, RoboTwin authoring
└── docs/                       the guides above

conda create -n agenticwam python=3.10 -y
conda activate agenticwam
# torch first, from the CUDA-12.8 index:
pip install torch==2.7.1+cu128 torchvision==0.22.1+cu128 --extra-index-url https://download.pytorch.org/whl/cu128
# everything else (pinned in pyproject.toml; transformers>=4.57 for Qwen3-VL):
pip install -e .

Then download the public weights and the AW-0 checkpoint, and set their locations in configs/paths/default.yaml — the one file that names paths on your machine (or keep a per-machine copy and pass paths=<name>, or override single keys such as paths.pretrain_ckpt=...). docs/INSTALL.md has the download commands, the key list and a smoke test.

Post-train the released checkpoint on your own task, then serve it on the robot.

docs/ROBOTWIN.md covers the simulator setup, data collection and every evaluation mode.

python scripts/launch_train.py --task pretrain_magic_soup --nproc 8 data.yam_root=<yam_root> data.ego_root=<ego_root>

Two sources in a standard LeRobot layout — yam (robot arms + grippers, 3 cameras) and ego (egocentric human hands, 1 camera) — each with stats, instruction embeddings and per-frame code labels. From scratch the video expert starts from a Cosmos-Predict2.5-2B video checkpoint (not distributed); to continue from the AW-0 checkpoint instead, add foundation_ckpt=<AW-0 checkpoint> model.load_pretrain=false. docs/PRETRAINING.md has the data format, the labeling pipeline and the configuration.

Every training entry point (scripts/train_*.sh, scripts/launch_train.py) reads the rendezvous from the environment — run the same command on every node:

export NNODES=4                     # total nodes
export NODE_RANK=0                  # 0 on the master, 1, 2, ... on the others
export MASTER_ADDR=<rank-0 host>
export MASTER_PORT=29500
export RUN_ID=<run name>            # the same on every node: one run directory, and resumable

--nproc / NPROC sets the GPUs per node (8). Effective batch = batch_size × gradient_accumulation_steps × GPUs. DeepSpeed ZeRO-1 by default (ZERO=2 for ZeRO-2). docs/PRETRAINING.md §4.

Training logs to W&B (configs/train.yaml, project agenticwam). Export WANDB_API_KEY (never put it in a script or config), or pass wandb.enabled=false / set WANDB_MODE=offline.

Released under the BSD 2-Clause License — see LICENSE. The public models listed in the Model Zoo keep their own licenses.

If you find AW-0 useful, please cite:

@misc{aw0_2026,
  title={Agentic World Action Models},
  author={Zhuoyang Liu and Dantong Niu and Botong Zhang and Boning Shao and Jing Wang and Ruijie Zheng and Yuvan Sharma and Runyu Lu and Jiaming Liu and Renrui Zhang and Roei Herzig and Yutong Bai and Ken Goldberg and Danfei Xu and Yuke Zhu and Guanzhi Wang and Linxi Fan and Trevor Darrell},
  year={2026},
  url={https://agentic-wam.github.io/},
}

Built on NVIDIA Cosmos-Predict2.5 and Cosmos-Reason1, Wan2.2, SDAR, Qwen3-VL, LeRobot and RoboTwin 2.0.