RL with BWM World Model#

This recipe trains a single-view pi0.5 policy with PPO on RoboTwin adjust_bottle inside the Boundless World Model (BWM), an action-conditioned video world model built on Wan2.2-TI2V-5B. Training never touches the simulator; checkpoints are scored in RoboTwin afterwards.

Overview#

Train a single-view pi0.5 policy with PPO on RoboTwin adjust_bottle simulated by the BWM world model.

Environments

BWM World Model · RoboTwin (eval)

Algorithms

PPO

Tasks

adjust_bottle

Hardware

1 node · 8 GPUs

You’ll do: install BWM next to the openpi + RoboTwin environment → download the assets → launch run_embodiment.sh → score the checkpoint in RoboTwin.

Prerequisites: RoboTwin installation · the pi0.5 SFT checkpoint · BWM weights, reward classifier and initial frames (steps below).

Tasks#

Environment	Task	Config / Weights	Focus
BWM	RoboTwin `adjust_bottle`	`bwm_robotwin_adjust_bottle_ppo_openpi_pi05`, `lidavidsh/bwm-robotwin-adjust-bottle-failure-ft`	BWM fine-tuned on success and failure rollouts, served in process by the `bwm` backend of `WorldModelEnv`.

Observation and Action#

Field	Description
Observation	One 640x480 head-camera frame generated by the world model; `states` is the policy’s last joint command (`state_from_last_action`).
Action	A chunk of 50 Aloha joint commands (14-dim). The backend maps them to BWM’s normalized end-effector poses with URDF forward kinematics; each chunk is one 61-frame BWM window (9 condition frames, 52 new).
Reward	A frozen T5 success classifier scores every frame from frame 51 on; the episode pays 1 on its 200th step if the peak score reaches the threshold (`reward_mode: episode_window_max`), and the limit ends the episode as a termination (`terminate_at_max_episode_steps`).
Prompt	The task instruction stored with each initial frame.

Installation#

Install the openpi model and the RoboTwin environment as in RoboTwin. The simulator is needed only for evaluation and for the URDF the action bridge reads.

BWM is not pip-installable from its repository; put the source on PYTHONPATH:

git clone https://github.com/boundless-large-model/boundless-world-model
git -C boundless-world-model checkout cc0bc74
export PYTHONPATH=$(pwd)/boundless-world-model:$PYTHONPATH

BWM needs diffsynth==2.0.11, whose API differs from the DiffSynth fork the wan backend installs. Install it into a separate directory and put that directory first on PYTHONPATH only for BWM runs:
```
pip install --no-deps --target /path/to/bwm_pydeps diffsynth==2.0.11 ftfy==6.3.1 wcwidth modelscope
export PYTHONPATH=/path/to/bwm_pydeps:$PYTHONPATH
```
torchaudio matching the installed torch build is also required.

Assets#

Asset	Source	Config key
Wan2.2-TI2V-5B DiT and VAE	`Wan-AI/Wan2.2-TI2V-5B`	`env.train.model_path`, `env.train.VAE_path`
Fine-tuned BWM weights	`lidavidsh/bwm-robotwin-adjust-bottle-failure-ft` (`step-750.safetensors`)	`env.train.checkpoint_path`
BWM action statistics	`demo/stat.json` in the BWM repository	`env.train.action_bridge.stat_path`
Aloha URDF / SRDF	RoboTwin assets, `embodiments/aloha-agilex`	`env.train.action_bridge.urdf_path`, `srdf_path`
Reward classifier	`lidavidsh/robotwin-adjust-bottle-t5-success-classifier` (`reward_model.pt`) and `google-t5/t5-base`	`env.train.reward_model.from_pretrained`, `t5_model_name`
pi0.5 SFT policy	WorldArena/WorldArena (model repository, `fulldata`, norm stats in `assets/robotwin_clean`)	`actor.model.model_path`, `actor.model.openpi_data.norm_stats_path`
Initial frames	48 RoboTwin reset observations, one `episode_XXXXX.npy` each: an object array of per-frame dicts (`np.save(..., allow_pickle=True)`) whose first frame holds `image` (480x640x3 `uint8` head-camera frame), `instruction` or `task` (the prompt), and `delta_action` (14 values; the loader requires the key, BWM ignores it). No export script is included yet.	`env.train.initial_image_path`

The rollouts BWM was fine-tuned on are in lidavidsh/robotwin-adjust-bottle-pi05-rollouts.

Run It#

Set the paths above in examples/embodiment/config/env/bwm_robotwin_adjust_bottle.yaml and bwm_robotwin_adjust_bottle_ppo_openpi_pi05.yaml, then:

export ROBOT_PLATFORM=ALOHA ROBOTWIN_PATH=/path/to/RoboTwin
bash examples/embodiment/run_embodiment.sh bwm_robotwin_adjust_bottle_ppo_openpi_pi05

The recipe runs PPO with GAE (gamma 1.0), 48 envs x 2 rollout epochs per step, for 14 steps: 4 critic-only, then 10 joint updates. The critic reads the VLM features (value_after_vlm). It uses FSDP no_shard (required by openpi). On ROCm builds whose Triton lacks triton_key, add actor.model.openpi.torch_compile=False.

Evaluate in RoboTwin#

Score a checkpoint with the RoboTwin pi0.5 evaluation config, using this recipe’s policy settings and single head view:

export ROBOT_PLATFORM=ALOHA ROBOTWIN_PATH=/path/to/RoboTwin
bash evaluations/run_eval.sh robotwin robotwin_adjust_bottle_openpi_pi05_eval \
  rollout.model.model_path=/path/to/sft runner.ckpt_path=/path/to/full_weights.pt \
  rollout.model.num_steps=4 rollout.model.openpi.task=rl rollout.model.openpi.noise_level=0.5 \
  ++rollout.model.openpi.value_after_vlm=True \
  ++rollout.model.openpi.value_head_hidden_sizes=[512,256,128] \
  ++rollout.model.openpi_data.adapt_to_pi=false \
  ++rollout.model.openpi_data.norm_stats_path=/path/to/sft/assets/robotwin_clean \
  env.eval.assets_path=/path/to/RoboTwin \
  env.eval.task_config.camera.head_camera_type=Large_D435 \
  env.eval.task_config.camera.collect_wrist_camera=false

Results#

With the original in-process integration this backend was ported from (older RLinf tree), on 8x MI300X, 14 PPO steps improved the SFT policy (0.773) in RoboTwin on 144 held-out initial states x 3 evaluation passes:

Seed	Success	vs SFT (90% CI)
1	0.889	+11.6 pp [+8.1, +15.1]
2	0.884	+11.1 pp [+6.7, +15.7]
3	0.847	+7.4 pp [+3.7, +11.3]
4	0.917	+14.4 pp [+10.0, +19.0]
5	0.813	+3.9 pp [+1.6, +6.3]

Seed-to-seed spread is large; judge a run with several seeds. Training with the ported backend has not yet reproduced this gain; see #1713.