PaperlayerView exact revision on arXiv

Source-linked AI summary

Dream to Remember: Continual Steering of Generalist Robot Policies with World Models

Bahey Tharwat, Iman Nematollahi, Marvin Alles, Xingyuan Zhang, Akshay L Chandra, Abhinav Valada

arXiv:2610.13816v1cs.ROcs.CVcs.LG

TL;DR

Continual adaptation of generalist robot policies must improve new tasks while retaining skills learned from demonstrations, but both fine-tuning and steering can forget prior tasks. D2R steers a frozen policy and rehearses earlier tasks in a frozen world model. It improves sequential task performance across simulated and real-robot settings, including CALVIN and LIBERO-90.

  • Problem

    Generalist robot policies lack native continual adaptation that reliably preserves prior skills, and sequential fine-tuning or input steering can cause forgetting.

    Source: §1 · pp. 1–2

  • Method

    D2R adapts a frozen generative policy with a small input-steering module and uses imagined world-model rollouts to rehearse prior tasks.

    Source: §1 · p. 2

  • Results

    63.7% to 74.8% average success on CALVIN after eight imagined tasks; D2R also preserves performance on pretrained LIBERO suites and retains gains across three Franka tasks.

    Source: §1 · p. 2

  • Takeaways & Limitations

    A world model too inaccurate for policy optimization can still serve as memory for rehearsal, allowing D2R to retain skills while learning new tasks in the environment.

    Source: §1 · pp. 2–3

    D2R assumes access to a given frozen world model, whether trained on robot data or pretrained without exposure to the target tasks.

    Source: §3.1 · p. 4

Abstract

from arXiv · show

Generalist robot policies are trained with imitation learning and are bounded by the quality and coverage of their demonstrations. Reinforcement learning can improve and extend them, either by fine-tuning the policy weights or by training a small module to steer the input noise of a frozen generative policy. When tasks arrive sequentially, both approaches suffer from catastrophic forgetting. Fine-tuning on a task degrades prior tasks, and steering only shifts forgetting to the shared steering module. We propose Dream to Remember (D2R), which continually steers a frozen generalist policy toward new tasks, while retaining its earlier skills. As the steering module learns a new task, a frozen pretrained world model rehearses prior tasks through imagined rollouts, anchoring the module to its earlier behavior without a replay buffer or new data. When the world model is faithful enough, D2R learns the new task entirely in imagination. Otherwise, it learns the task in the environment while using the world model only to rehearse prior tasks. On CALVIN, D2R improves eight tasks in sequence entirely in imagination, raising average success from 63.7% to 74.8%, while LoRA fine-tuning on the same sequence drops it to 11.7%. With a frozen off-the-shelf generalist policy and world model, $π_0$ and A2World, D2R learns unseen LIBERO-90 tasks in the environment while keeping $π_0$'s performance on the LIBERO suites it was fine-tuned on, even though the world model is not accurate enough for policy optimization. Finally, on a real Franka robot, D2R improves three tasks in sequence entirely in imagination while retaining the earlier improvements.

1 INTRODUCTION

D2R addresses forgetting during continual adaptation by steering a frozen generative policy and rehearsing prior tasks in a world model. Its evaluations span imagined and environment-based learning, including simulated benchmarks and a real robot.

Motivation

  • Generalist policies’ imitation-learning priors provide useful capabilities but do not natively support continual adaptation that reliably preserves prior skills.Imitation learning also bounds these policies by the quality and coverage of their demonstrations.Source: §1 · p. 1
  • Fine-tuning can alter prior-task behavior, while input steering shifts forgetting to the shared steering module as new tasks change its sampled noise.Test-time action selection avoids persistent adaptation but requires scoring multiple proposals at every decision.Source: §1 · pp. 1–2

Approach

  • D2R rehearses prior tasks using imagined rollouts from a frozen world model, regularizing the steering module toward its earlier behavior without storing trajectories or revisiting tasks on the robot.The model may be too inaccurate for learning a new task yet still provide sufficient rehearsal.Source: §1 · p. 2

Contributions

  • When the world model is sufficiently faithful, D2R learns the current task in imagination; otherwise, it learns from real experience while retaining model-based rehearsal.The paper identifies steering in imagination, world-model memory, and lower fidelity requirements for rehearsal than optimization as its contributions.Source: §1 · pp. 2–3

Evaluation

  • 63.7% to 74.8% average success on CALVIN after eight imagined tasks; D2R also preserves LIBERO suite performance and retains three-task gains on Franka.On LIBERO-90, it learns unseen tasks through environment interaction while preserving the pretrained policy’s LIBERO-suite performance.Source: §1 · p. 2

2 RELATED WORK

Prior work adapts pretrained generative policies in weight, output, or input space, while continual-learning methods mitigate forgetting through regularization, capacity allocation, or rehearsal. D2R combines input-space steering with world-model rehearsal for sequential adaptation.

Improving pretrained generative policies

  • Prior generative-policy adaptation spans weight updates, action selection or residuals, and input-noise learning, with most methods focused on individual-task adaptation.D2R carries noise-space steering across tasks and rehearses prior tasks in a world model.Source: §2 · p. 3

Continual learning without forgetting

  • Continual-learning approaches use regularization, task-specific capacity, or rehearsal; several robot-policy methods reduce forgetting by expanding modules or skill libraries as tasks accumulate.D2R instead uses a single steering module and a separate frozen world model for memory, storing only initial observations of retained tasks.Source: §2 · p. 3
  • D2R updates a small steering module with forward access to the policy and uses a frozen world model for memory, rather than requiring policy weights, gradients, and on-policy interaction.The related-work passage contrasts this with retention approaches that depend on small updates within a large pretrained network.Source: §2 · p. 3

3 METHOD

D2R adapts a frozen generative policy through a small noise-steering module, learning current tasks with a world model when it is sufficiently faithful and rehearsing retained tasks in that model. The method’s four parts define the continual setting, steering, rehearsal, and the world model’s alternative roles.

3.1 PROBLEM FORMULATION

  • At each stage, the robot learns on the arriving task while aiming to preserve prior skills, without revisiting them on the robot or storing trajectories.The setting assumes access to a small set of initial observations per task and a frozen pretrained policy and world model.Source: §3.1 · pp. 3–4

3.2 STEERING A FROZEN POLICY IN IMAGINATION

  • The steering module treats noise as the action in a transformed MDP, optimizing it through a frozen policy in the world model when sufficiently faithful, or in the environment otherwise.The module changes selected actions without changing the policy weights; a behavior-cloning term helps limit exploitation of world-model prediction errors during imagined training.Source: §3.2 · pp. 4–5

3.3 REHEARSAL IN THE WORLD MODEL

  • For each retained task, D2R regenerates world-model rollouts from initial observations and applies a KL penalty to anchor current steering to the task’s reference behavior.Rollouts are regenerated rather than stored, require no reward, and use states relevant to the retained skill.Source: §3.3 · p. 5

3.4 WORLD MODEL AS SIMULATOR OR MEMORY

  • A world model used as a simulator must support reliable policy optimization, whereas as memory it need only reproduce representative states for rehearsal; task coverage determines which role it can serve.Current-task learning therefore occurs in imagination when the model is a simulator and in the environment when it serves only as memory; rehearsal remains in the world model in both regimes.Source: §3.4 · p. 5

4 SIMULATION EXPERIMENTS

The simulation experiments test D2R across settings where a world model can support both policy optimization and rehearsal, or only rehearsal. Across sequential tasks, D2R improves performance while retaining earlier skills, including with a coarse world model.

4.1 EXPERIMENTAL SETUP

  • CALVIN uses imagined rollouts for learning and rehearsal, whereas LIBERO-90 tasks are learned with real interaction and A2World used only for rehearsal.The evaluation covers single-task adaptation and sequential adaptation, measuring task success, retention, and forgetting.Source: §4.1 · p. 6

4.2 STEERING IN IMAGINATION WITHOUT FORGETTING

  • 79.8% average success versus 63.7% for Base DP; D2R is the only method with positive retention and matches DiWA in success rate.Steering keeps the policy frozen, while rehearsal constrains the steering module on other tasks; this combination is more effective for retention than rehearsal with weight fine-tuning.Source: §4.2 · p. 7, Table 1 · p. 7

4.3 CONTINUAL STEERING ACROSS TASKS

  • 74.8% average success on CALVIN (from 63.7%) and 74.9% on LIBERO (from 10.4%), with 5.7 and 8.3 points of forgetting, respectively, and +1.6 retention on pretrained LIBERO suites.D2R learns eight tasks sequentially, entirely in imagination on CALVIN and in the environment on LIBERO.Source: §4.3 · p. 7, Table 2 · p. 8
  • LoRA loses 53.2 points on CALVIN and 61.4 on LIBERO; DSRL ends near base performance, while replay-based rehearsal comes closest to D2R.D2R keeps about 400 frames, roughly 150 times fewer than a trajectory buffer; first-observation rehearsal alone forgets 12.9 points on CALVIN and 52.5 on LIBERO.Source: §4.3 · p. 8

4.4 A COARSE WORLD MODEL AS MEMORY

  • A2World is too inaccurate to optimize LIBERO-90 policies: cloning from imagined replays reaches 19.6%, versus 74.9% from real demonstrations, yet its imagined states suffice for rehearsal.D2R obtains retained-task states by unrolling tasks from stored initial frames, and rehearsal on those states matches rehearsal on real stored trajectories.Source: §4.4 · pp. 8–9
  • On the real Franka, D2R preserves gains from three imagined adaptation stages, while DSRL without rehearsal loses about half; the robot is used only for evaluation.The heatmap’s diagonal reports success when each task is learned, and cells below it report forgetting.Source: Figure 4 · p. 9

5 REAL-WORLD EXPERIMENTS

On a real Franka robot, D2R improves three tasks sequentially in imagination and preserves those gains through the sequence.

Setup

  • Three tasks—open drawer, move slider left, and turn on LED—are adapted sequentially, with 20 evaluation episodes per adapted task after each stage.The world model is trained on six hours of teleoperated play data and remains frozen; the robot is used only for evaluation.Source: §5 · p. 9

Results

  • D2R preserves gains through the sequence, ending 30 points above the base policy on average versus DSRL's 15 points.Without rehearsal, DSRL loses about half of the gains.Source: §5 · p. 9
  • Steering in imagination improves real-robot execution by 40, 20, and 30 points over a frozen base policy.The base policy initially succeeds on 50 to 55% of episodes.Source: §5 · p. 9

6 CONCLUSION

D2R continually adapts a pretrained generative robot policy while retaining existing skills, using steering and imagined rehearsal in a frozen world model.

Conclusion

  • D2R steers a frozen policy with a single small module and rehearses retained skills using rollouts imagined in a frozen world model.Adaptation requires no policy weight updates, stored trajectories, or robot interaction for rehearsal.Source: §6 · p. 9
  • Across two simulated benchmarks and a real Franka robot, steering matches weight fine-tuning on adapted tasks while preserving prior skills substantially better.The experiments indicate that a world model can require lower fidelity for rehearsal than for policy optimization.Source: §6 · p. 9
  • D2R still stores a few initial frames per task to prompt the world model, which future work could instead generate on demand..Source: §6 · p. 9

S.1 Training Objectives

The appendix lists training-objective topics covering imagined rewards, policy optimization, behavior cloning, and rehearsal.

S.1 Training Objectives

  • The training-objectives appendix covers reward in imagination and policy optimization with PPO and SAC..Source: §S.1 · p. 14
  • It also lists behavior-cloning loss and rehearsal loss as training objectives..Source: §S.1 · p. 14
  • The listed topics include both policy-optimization methods and losses for behavior cloning and rehearsal..Source: §S.1 · p. 14

S.3 Implementation Details

The implementation-details appendix lists sections on world models, base policies, rewards, steering modules, hyperparameters, and compute.

S.3 Implementation Details

  • The appendix covers world models, base policies, and rewards as implementation topics..Source: §S.3 · p. 14
  • It also lists steering modules and critics, followed by hyperparameters..Source: §S.3 · p. 14
  • Compute and wall-clock are listed as implementation details..Source: §S.3 · p. 14

S.4 Experimental Setup Details

The section lists experimental setup details covering environments, task sequences, the real-world setup, action space, rehearsal storage, baselines, and metrics.

S.4.1 Environments and Task Sequences

  • The section includes an overview of environments and task sequences.Source: §S.4 · p. 14

S.4.2 Real-World Setup

S.4.3 Action Space

S.4.4 Rehearsal Seed Frames and Storage Cost

  • The section covers rehearsal seed frames and storage cost.Source: §S.4 · p. 14

S.4.5 Baselines

S.4.6 Metrics

S.5 Additional Results and Ablations

The section lists additional results and ablations, including full sequential and per-task results, steering-module initialization, and task-order sensitivity.

S.5.1 Full Sequential Results

S.5.2 Full Per-Task Results

S.5.3 Initialization of the Steering Module

  • The section examines steering-module initialization.Source: §S.5 · p. 14

S.5.4 Task Order Sensitivity

S.5.5 Retention Coefficient

This section describes how D2R adapts its steering module while rehearsing retained tasks in a frozen world model, including how adaptation changes when model fidelity is insufficient and the associated storage and compute costs. It defines retention relative to the frozen base policy and examines how baseline design, initialization, task order, and world-model fidelity affect continual steering, distinguishing rehearsal from learning new tasks in imagination.

S.1.1 REWARD IN IMAGINATION

  • For imagined rewards, frozen per-task latent success classifiers score rollout states and stop rollouts above 0.90 confidence; LIBERO instead uses benchmark rewards and fixed-length rehearsal.A2World is not accurate enough for policy optimization, so no reward model is trained in its latent space.Source: §S.1.1 · p. 15

S.1.2 POLICY OPTIMIZATION WITH PPO

  • When the world model is faithful enough, D2R learns the current task in imagination with PPO; otherwise, it learns in the environment with SAC while rehearsal remains imagined.This is the regime used on CALVIN and the real robot for PPO, and on LIBERO for SAC; the robot is used only for the current task.Source: §S.1.2 · p. 15, §S.1.3 · pp. 15–16

S.1.5 REHEARSAL LOSS

  • Rehearsal regenerates retained-task rollouts in the frozen world model and penalizes divergence from the prior module for already-steered tasks or the prior noise distribution for others.It uses stored initial observations, needs no retained-task trajectories, and applies the same retention coefficient across tasks.Source: §S.1.5 · p. 16, §S.4.4 · p. 21

S.3.6 COMPUTE AND WALL-CLOCK

  • Rehearsal is D2R’s only added cost: it lengthens CALVIN stages by 25%, while LIBERO rehearsal takes 30 minutes for 50 episodes per retained task and is generated once per stage.LIBERO rollouts are generated before the stage because querying A2World and π0 inside the update loop would dominate wall-clock time.Source: §S.3.6 · p. 20

S.4.4 REHEARSAL SEED FRAMES AND STORAGE COST

  • After eight tasks, D2R retains 1,200 CALVIN seed frames versus 72,240 replay frames and 400 LIBERO seed frames versus 59,480 replay frames.D2R stores one initial frame per demonstration; replay stores the full trajectories.Source: §S.4.4 · p. 21

S.4.5 BASELINES

  • The baselines include fine-tuning with rehearsal and steering penalties on different state sets, while per-task and multi-task modules serve as nonsequential reference settings.Per-task modules avoid retention demands but require task identity and one module per task; the multi-task reference trains jointly on all eight tasks with a larger budget.Source: §S.4.5 · pp. 22–23

S.4.6 METRICS

  • Retention is the mean final change from the frozen base policy on the retention set; forgetting is the average end-of-sequence loss on tasks learned before the last stage.Positive forgetting scores indicate losses, zero means no average forgetting, and retention may include tasks outside the sequence; single-task adaptation averages changes over all tasks, including the adapted task.Source: §S.4.6 · pp. 23–24

S.5.1 FULL SEQUENTIAL RESULTS

  • Table S.3 adds single-task and current-task success to the sequential results, distinguishing the cost of in-sequence learning from subsequent forgetting.Current-task is mean success immediately after each task’s stage; its difference from final Success is (N −1)/N times forgetting.Source: §S.5.1 · p. 24, Table S.3 · p. 24

S.5.2 FULL PER-TASK RESULTS

  • The retention matrices encode each task’s signed success change from the base policy, while sequence rows reveal how retained skills fare as later tasks are learned.The LIBERO figure tracks the four suites on which π0 was fine-tuned, showing changes across stages that a final mean alone would not reveal.Source: §S.5.2 · p. 24, Figure S.3 · p. 25, Figure S.4 · p. 25

S.5.3 INITIALIZATION OF THE STEERING MODULE

  • Reinitializing the steering module improves final success and retention on both benchmarks; on LIBERO, forgetting falls from 18.6 to 8.3.The reported gains are 7.9 points on each axis on CALVIN and 8.1 and 0.4 points on LIBERO; forgetting on CALVIN is unchanged.Source: §S.5.3 · p. 25, Table S.4 · p. 25

S.5.4 TASK ORDER SENSITIVITY

  • Across four CALVIN task orders, final success ranges from 74.6% to 78.9%, retention from +10.9 to +15.2 points, and forgetting from 1.0 to 5.7 points.The default order ranks third on success and retention and last on forgetting; λ = 0.005 balances current-task and retained-task success, while α = 0.05 avoids the poor real-environment performance associated with shortcut learning at a small coefficient.Source: §S.5.4 · pp. 26–27, Figure S.6 · p. 26

S.5.7 WORLD MODEL FAITHFULNESS CHECK

  • On unseen LIBERO-90 tasks, imagined-demo cloning reaches 19.6% versus 74.9% for real-demo cloning, motivating environment learning with A2World used only for rehearsal.On the four suites A2World was fine-tuned on, imagined-versus-real cloning success is 70.3% versus 74.5% (Spatial), 82.0% versus 82.5% (Goal), 88.9% versus 82.4% (Object), and 46.0% versus 45.8% (Long).Source: §S.5.7 · p. 27, Figure S.7 · p. 27

S.5.8 QUALITATIVE IMAGINED REHEARSAL ROLLOUTS

  • A2World rollouts on LIBERO-90 can be incomplete, yet plausible state transitions remain useful for rehearsing prior behavior rather than learning new tasks entirely in imagination.DreamerV2 trajectories more closely follow task progression in CALVIN and real-world experiments, enabling both imagined task learning and rehearsal.Source: §S.5.8 · pp. 27–28, Figure S.8 · p. 30, Figure S.9 · p. 31
Loading 2610.13816v1…