Can World Models Continually Steer Generalist Robot Policies Without Forgetting?

Forgetting moves rather than disappears

A generalist robot policy trained by imitation learning is limited by the demonstrations it saw. Reinforcement learning can push it further in two ways: update the policy weights, or train a small module that steers the input noise of a frozen generative policy. The paper argues that both break down when tasks arrive one after another. Weight updates change behavior on earlier tasks, and steering does not escape the problem because the shared steering module is itself trained on each new task, so the noise sampled for old tasks shifts.

This is the paper's framing of the problem, and it sets the design target. The policy stays frozen, so its demonstrations remain a prior worth keeping. What must be protected is the steering module, the only part that changes. The question becomes how to keep that module near its earlier behavior while it learns something new, without rerunning old tasks on the robot and without storing their trajectories.

Rehearsal comes from imagined rollouts

Dream to Remember (D2R) answers with rehearsal. As the steering module learns the current task, a frozen pretrained world model regenerates rollouts of the retained tasks, and a penalty pulls the module back toward its previous behavior on those imagined states. The world model is a learned predictor of how the scene evolves under actions; here it acts as memory rather than as a training ground for the new task.

The storage cost is deliberately small. D2R keeps only the first observation of each demonstration to seed these rollouts: 150 frames per task on CALVIN and 50 on LIBERO. After an eight-task sequence that is 1,200 frames on CALVIN and about 400 on LIBERO, against roughly 72,000 and 59,000 frames for a replay buffer holding the same episodes in full. The gap is the episode-length dimension, so it widens as tasks get longer.

A world model can remember more than it can teach

The paper separates two regimes by how faithful the world model is. When it is accurate enough, the new task is learned entirely in imagination: on CALVIN, each task is improved with on-policy PPO inside the world model, rewarded by a latent success classifier, and the other seven tasks are rehearsed at each adaptation. When it is not, the task is learned in the environment and the world model is used only for rehearsal.

The LIBERO experiments test the second regime with off-the-shelf components: the public π0 policy and the A2World world model, neither exposed to LIBERO-90. A direct check shows why the split matters. Policies cloned from real demonstrations reach 74.9% on the LIBERO-90 tasks, but policies cloned from A2World-imagined replays of the same demonstrations reach only 19.6%, while on the four suites the two are within a few points. A steering policy optimized inside such a model would exploit its prediction errors, so each LIBERO-90 task is learned through real interaction with SAC, and A2World serves only as memory.

The authors' interpretation is that rehearsal demands less fidelity than optimization. The imagined trajectories are too inaccurate to learn from, yet their states are sufficient for anchoring the module, and rehearsal on them matches rehearsal on stored real trajectories in the reported comparison.

Steering improves the new task and keeps the old ones

On CALVIN, D2R learns eight tasks in sequence entirely in imagination and raises average success from 63.7% to 74.8%. LoRA fine-tuning on the same sequence drops to 11.7%. The comparison is between adaptation spaces under the same sequence, not a claim that steering is universally better than weight updates.

The single-task table makes the trade-off visible in a different setting. There, each method is trained on one CALVIN task and evaluated on all eight, and retention is the mean change in success over all tasks relative to the base policy, in percentage points. DP-TTC reaches 79.2% average success on the adapted tasks but loses 43.9 points of retention; DiWA reaches 78.0% and loses 12.7; DSRL reaches 81.8% and loses 5.3. D2R reaches 79.8% and loses 1.6, the only method in that table to improve the adapted task while ending above the base policy overall.

These numbers answer different questions and should not be merged. Single-task retention includes the adapted task itself, so gains and losses enter with their sign. Sequential forgetting is measured against the stage at which each task was learned, and the last task is excluded because nothing is learned after it. A method can look strong on one and weak on the other.

Where the penalty is applied decides whether it works

The ablations isolate the mechanism. DSRL without rehearsal forgets heavily. Adding the same penalty on current-task observations paired with retained-task instructions lightens forgetting on CALVIN, where all tasks share a scene, but not on LIBERO, where scenes change and current-task observations do not cover retained-task states. Applying the penalty only at stored first frames sees states in which a retained task has not yet begun. Applying it to rollouts collected once per stage and stored works, but requires keeping those trajectories.

D2R instead regenerates retained-task rollouts in the world model from the stored initial frames. The paper reports that this matches rehearsal on real stored trajectories, which is what makes the small storage footprint possible. The distinction is not rehearsal versus no rehearsal but which states the anchor is computed on.

An off-the-shelf policy keeps its fine-tuned suites

On LIBERO, D2R learns eight unseen LIBERO-90 tasks sequentially in the environment while keeping π0's performance on the Spatial, Object, Goal, and Long suites it was fine-tuned on. Retention there is the change in held-out success relative to π0, measured on those four suites, which never appear in the sequence. This is a different denominator from the CALVIN average-success figure and from the single-task retention column.

The result matters because neither component was trained for this: the policy and world model are public checkpoints, and A2World was fine-tuned on the four suites but never saw LIBERO-90. The paper's claim is therefore narrower and more useful than a benchmark win. A world model that cannot support policy optimization on the target tasks can still serve as memory for them, without fine-tuning on the tasks being learned.

Three real tasks keep their gains in imagination

On a real Franka Emika Panda, three tasks are adapted in sequence entirely in imagination, with the robot used only for evaluation: open drawer, move slider left, and turn on LED. The world model is an RSSM trained on six hours of teleoperated play data, and the base policy is a language-conditioned diffusion policy trained on 100 demonstrations per task across seven tasks.

The reported outcome is that D2R retains prior gains across the sequence, whereas DSRL without rehearsal loses about half of them. Evaluation uses 20 episodes per adapted task of at most 72 steps, starting from 20 fixed initial configurations reset by hand. This is a small sequence on one platform, and it tests retention of improvements rather than transfer to unrelated tasks.

Choose rehearsal fidelity by what the model must do

The paper's practical priority is a split by role, not a single accuracy threshold. If a world model is faithful enough to optimize a policy inside, D2R can learn the new task entirely in imagination, as on CALVIN. If it is not, learn the task in the environment and use the model only to regenerate retained-task states for the anchor, as on LIBERO. The LIBERO cloning check, 74.9% from real demonstrations against 19.6% from imagined replays, is the kind of evidence that tells which side of that split a model falls on.

The remaining cost is explicit: D2R still stores a few initial frames per task to prompt the world model, and the authors name generating them on demand as future work. Within the evidence, the case for D2R is strongest where a frozen generalist policy must accumulate skills without weight updates, stored trajectories, or revisiting old tasks on the robot, and where the world model's states are good enough to anchor behavior even when its dynamics are not good enough to learn from.

Sources used in this answer