Fast-WAM: World Action Model

SUMMARY

Fast-WAM is a World Action Model: during training it co-trains on a video-prediction objective, so the policy learns representations that capture how the scene evolves under an action, not just what the scene looks like. Most World Action Models spend that understanding at test time too, generating future frames before deciding what to do - the "imagine-then-execute" pattern - which is slow.

Fast-WAM skips that step. At inference it runs a single forward encoding pass over the current observation and acts directly, without denoising imagined future video. That keeps the dynamics-aware representations from training while running in real time - about 190ms latency, 4x faster than imagine-then-execute World Action Models, with no drop in task success.

When to Use Fast-WAM?

Reach for Fast-WAM when:

• The task has real dynamics to get wrong - contact, momentum, deformable or articulated objects - where a policy that only reacts to the current frame tends to overshoot or misjudge contact
• Closed-loop control needs to run at real-time control-loop rates, ruling out imagine-then-execute World Action Models
• Demonstrations are on the robot itself, similar in spirit to ACT, but the task benefits from dynamics-aware pretraining rather than chunking alone

For a policy that trains from scratch on a handful of demonstrations with no world-model pretraining, see ACT. For a pretrained, language-conditioned generalist across many tasks and embodiments, see π0.5.

Requirements

Requirement	Value
Python	3.10-3.12
GPU	A CUDA GPU - fine-tuning updates a video-pretrained Diffusion Transformer backbone
Pretrained checkpoint	Downloaded once and reused across fine-tuning runs
Demonstrations	A dataset with camera frames, proprioceptive state, and the action taken at each step. Currently, only the LeRobot dataset format is supported

Quick Installation

pip install "telekinesis-cerebellum[fast-wam]"

The fast-wam extra brings the Fast-WAM model and its training loop. Everything else - dataset loading, checkpointing, and the deployment interface shared with ACT and π0.5 - comes with the base package.

Quickstart

Download the pretrained checkpoint

cerebellum download fast-wam-base --to checkpoints/fast-wam-base

Load your demonstration data from LeRobot

Pull a dataset of teleoperated episodes - local or from the LeRobot Hub - into the layout Cerebellum trains on:

from telekinesis.dataengine import load_lerobot_dataset

# 1. Load a LeRobot dataset of teleoperated episodes
dataset = load_lerobot_dataset(
    repo_id="telekinesis-ai/insert_gear",   # local path or LeRobot Hub repo id
)

print(dataset)

Create config.yaml

task: insert_gear

dataset:
  path: demos/
  cameras: [cam_high, cam_wrist]

policy:
  class_name: FastWAM
  pretrained_checkpoint: checkpoints/fast-wam-base
  video_pretraining: true   # co-train on predicted futures; disable for a lighter fine-tune
  action_horizon: 50

training:
  num_steps: 30000
  batch_size: 16
  learning_rate: 0.00002
  device: auto

logger:
  log_dir: logs
  experiment: insert_gear_fastwam
  save_interval: 1000

Fine-tune

python train.py config.yaml

tensorboard --logdir logs

Deploy

Inference does not generate future frames, so it is a single forward pass per control step - the same shape as deploying ACT or π0.5:

from telekinesis.cerebellum.policy import Policy

# Use a Synapse robot
# robot =  

policy = Policy("logs/insert_gear_fastwam/checkpoint_best.pt")

observation = robot.get_observation()  # camera frames + proprioceptive state
action_chunk = policy.get_action(observation)
robot.execute(action_chunk)

Tuning Notes

• Video pretraining during fine-tuning is what teaches the backbone the task's dynamics; turning it off trains faster but gives up the main advantage Fast-WAM has over a plain VLA policy.
• Action horizon trades reactivity for smoothness, the same way chunk size does for ACT and action horizon does for π0.5.
• Because the value of world modeling here is in the learned representations rather than test-time generation, there is no separate "imagination" latency knob to tune at deployment - inference speed does not depend on how far ahead the model was trained to predict.

Next Steps

ACT

A sample-efficient, single-task policy that trains from scratch in minutes to a few hours.

Explore →

π0.5

A pretrained, language-conditioned generalist policy for tasks that need to generalize beyond one demonstration set.

Explore →