Generalist robot policies improved continuously without forgetting earlier skills
Dream to Remember: Continual Steering of Generalist Robot Policies with World Models
RoboticsComputer Vision and Pattern RecognitionMachine Learning
Summary
Robots learn to do many tasks by copying examples, but they forget old tasks when learning new ones. The authors propose a method called Dream to Remember that uses a helpful model of the world to imagine past tasks while learning new ones, so the robot doesn’t forget. This approach improves a robot’s ability to do tasks over time without needing to relearn or store old data. Tests on simulated and real robots show this method keeps older skills while gaining new ones more effectively than traditional fine-tuning.
What this means in practice
- •For robotics engineers: Enable robots to learn new tasks over time without losing previously learned skills by rehearsing past tasks in a world model.
- •For automation system designers: Improve multi-task robot controllers by integrating imagination-based rehearsal to maintain performance across sequential tasks.
Authors
Bahey Tharwat, Iman Nematollahi, Marvin Alles, Xingyuan Zhang, Akshay L Chandra, Abhinav Valada
Abstract
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.