CoachWorld

The action-conditioned world model behind RoboCoach

RoboCoach: World Models as Active Coaches for Compositional Robot Skills

Jiajun Liu, Yifan Chen, Yichao Liu, Jiayi Zhang, Ruoqu Chen,
Shaoxuan Xie, Guocai Yao, Mengdi Xu, Sen Cui, Changshui Zhang

RoboCoach turns imagined failures into targeted supervision for robot skills. In its Route–Imagine–Diagnose–Improve loop, skill policies act inside CoachWorld, a progress judge identifies the first unfinished subtask, and aggregated failures guide demonstration collection and expert updates.

CoachWorld makes these interactions possible with an adapted Wan2.2 TI2V-5B backbone:

• Action-conditioned prediction: generate future observations from sparse visual history and robot end-effector actions.
• Single-arm and bimanual robots: use a shared two-slot action representation with camera-aware conditioning.
• Autoregressive rollout: feed predicted observations back into the policy across successive video chunks.

With 150 additional subtask demonstrations, the full RoboCoach system improves success from 13.3% to 75.0% on Franka and 40.0% to 83.8% on AgileX. See the paper for evaluation details and the project page for interactive examples.

This repository provides the CoachWorld model, inference components, and training tools. The router, progress judge, and expert-update pipeline are separate parts of RoboCoach.

The released checkpoint is the CoachWorld pretrain model, trained on heterogeneous single-arm and dual-arm robot data. The target-only ablation and platform-specific post-trained models are distinct from this release.

Use Python 3.10 or 3.11 on a GPU machine with a PyTorch/CUDA setup suitable for Wan2.2. From the repository root:

git clone https://github.com/RoboCoach-AI/CoachWorld.git
cd CoachWorld
python -m venv .venv
source .venv/bin/activate
python -m pip install -e '.[dev]'
python -m pytest -q tests

The tests cover model components and data contracts. All commands below run from the repository root.

Download the CoachWorld pretrain release:

python -m pip install -U huggingface_hub
hf download JEdward/CoachWorld --local-dir models/CoachWorld

The release includes dit_model.safetensors, training_config.yaml, and model metadata. Use the custom coachworld implementation to load the checkpoint; the model card describes the release.

Preparing inputs: match the action coordinate frames, normalization, camera calibration, and temporal sampling specified in coachworld/data/. The inference components require these prepared inputs; a standalone inference CLI is not included.

Training starts from the Wan2.2 TI2V-5B base assets and a prepared CoachWorld video_latent dataset. The example config expects:

<COACHWORLD_MODEL_ROOT>/Wan2.2-TI2V-5B/
<COACHWORLD_DATA_ROOT>/video_latent/
  manifest.json
  train.index.jsonl
  val.index.jsonl
  stats/arm_slot_eef_pose.json
  ... latent and signal shards ...

The validation split is required by the default config. See coachworld/data/video_latent.py for the manifest and index schema. coachworld/data/video_latent_release.py packages an existing collection; it does not download or create the dataset.

Set the asset roots and a positive training length. The step count below is an example:

export COACHWORLD_MODEL_ROOT=/path/to/base-models
export COACHWORLD_DATA_ROOT=/path/to/prepared-data
export COACHWORLD_OUTPUT_ROOT=/path/to/outputs
export COACHWORLD_MAX_TRAIN_STEPS=10000

torchrun --nproc_per_node=5 scripts/training/train_world_model.py \
  --config configs/training/coachworld.yaml

COACHWORLD_MODEL_ROOT points to the parent of the Wan2.2 directory, not to the CoachWorld checkpoint downloaded above. Adjust the GPU/process count and configuration for your hardware.

Append the appropriate argument to the training command:

--init_from takes the directory containing dit_model.safetensors. A resumable checkpoint must have been saved by this trainer, including optimizer state. Configuration values can also be overridden with dotlist arguments such as world_model.learning_rate=1e-5.

If you use CoachWorld, please cite the RoboCoach paper:

@article{liu2026robocoach,
  title={RoboCoach: World Models as Active Coaches for Compositional Robot Skills},
  author={Liu, Jiajun and Chen, Yifan and Liu, Yichao and Zhang, Jiayi
          and Chen, Ruoqu and Xie, Shaoxuan and Yao, Guocai
          and Xu, Mengdi and Cui, Sen and Zhang, Changshui},
  journal={arXiv preprint arXiv:2609.39685},
  year={2026},
  url={https://arxiv.org/abs/2609.39685}
}

RoboCoach-authored code is released under the root MIT license. Modified Wan2.2 components retain Apache-2.0 terms, and PRoPE retains its MIT header. See THIRD_PARTY.md for attribution. Model weights and datasets have their own terms.