World Action Models Enable Continual Imitation Learning with Recurrent Generative Replays

Paper: 2606.27374
Authors: Manish Kumar Govind, Dominick Reilly, Smit Patel, Hieu Le, Srijan Das
Categories: cs.RO, cs.CV

The Gap

Continual imitation learning (IL) aims to equip robots with the ability to acquire new manipulation skills without forgetting previously learned ones.
Prior approaches fall into two camps:

• Experience replay: store raw human demonstrations and replay them during new task training. This works well but requires large memory and raises privacy concerns.
• Regularization: constrain weight changes (e.g., EWC, SI). These methods are memory-efficient but suffer from significant forgetting in high-dimensional control tasks.

Generative replay – using a separate generative model to synthesize pseudo-data from previous tasks – has been explored in classification but rarely in IL.
The bottleneck is that generating long, consistent action-observation trajectories for robotics is much harder than generating static images.
Existing generative replay for IL either uses simple VAEs that produce low-quality trajectories or relies on storing starting states, partially defeating the purpose.

This paper asks: Can we use a world model that predicts both actions and observations to generate high-quality, long-horizon pseudo-replays, completely eliminating the need to store any raw demonstrations?

The logical chain:

[Catastrophic forgetting in continual IL ]  
        |  
        v  
[Solution: generative replay without stored data]  
        |  
        v  
[Key assumption: WAM can generate consistent trajectories]  
        |  
        v  
[Method: REGEN - recurrent query of WAM for pseudo-replays]  
        |  
        v  
[Evidence: simulated & real robot tasks show 50% less forgetting]  
        |  
        v  
[Conclusion: WAMs are a viable foundation for demonstration-free continual IL]  

The Increment

One sentence: Before this work, generative replay for robot IL was limited to simple models that could not produce long-horizon consistent trajectories without stored priors; after this work, World Action Models (WAMs) enable high-quality pseudo-replay from scratch, achieving performance within 5-10% of privileged experience replay.

Core Mechanism

REGEN builds on a pre-trained World Action Model (WAM) – a state-space model that, given a current observation and action, predicts the next observation and the next action.
WAM is trained once on demonstrations of all tasks (stored only temporarily during training, then deleted).

During continual learning, when a new task batch arrives, REGEN does the following for each sample:

1. Query WAM with the current-task observation and a task instruction (e.g., one-hot ID) for a previously learned task.
2. Autoregressive generation: WAM rolls out a full trajectory by repeatedly feeding its own predictions back as input.
3. Pseudo-replay: The generated (observation, action) sequence is treated as a synthetic demonstration of the old task.
4. Joint training: The robot policy is updated on both the new task’s real data and the old task’s pseudo-replay data.

The process is recurrent because for each new task observation, multiple old-task pseudo-replays can be generated by seeding with different task instructions.

                     [Task Instruction] + [Init Observation]  
                                   |  
                                   v  
  [WAM Encoder] -> [Latent State] -> [WAM Dynamics] -> [Next Latent] -> [WAM Decoder]  
                                   |                       |  
                                   v                       v  
                          [Predicted Action]    [Predicted Observation]  
                                   |                       |  
                                   +--------->Loop<--------+  
                                   |  
                                   v  
                     [Full trajectory: (obs_1, act_1) ... (obs_T, act_T)]  
                                   |  
                                   v  
   [New task real data] + [Pseudo-replay trajectories] -> [Policy update]  

Structural metaphor: the chef’s mental practice

Imagine a chef learning to cook new dishes. The chef has a world model of cooking – given the current state of the dish (e.g., half-chopped vegetables) and an action (e.g., “add salt”), she predicts the new dish state.
To avoid forgetting an old dish (e.g., pasta), the chef could mentally rehearse: she imagines the initial state, then simulates each step using her world model, generating a sequence of mental images and actions.
This mental rehearsal is REGEN:

• The WAM = the chef’s internal model of cooking dynamics.
• The task instruction = “make pasta” triggers the initial conditions.
• The autoregressive roll-out = step-by-step mental simulation.
• The policy = the chef’s motor skills (actual chopping, stirring).
• The new task data = a new dish recipe (e.g., sushi) the chef is learning.

The chef does not need a physical copy of the pasta recipe card (stored demonstration); she just uses her mental model to generate the experience of cooking pasta. This is exactly what REGEN does for the robot.

Key Concepts

• Catastrophic Forgetting: When a neural network is trained sequentially on two tasks, the weights optimized for task 1 get overwritten by gradients from task 2.
  Example: A robot learns to push a block (task 1), then learns to grasp a cup (task 2). Without intervention, after task 2 it can no longer push the block reliably.
  Degree of forgetting is measured as performance drop on task 1 after training on task 2.

• Generative Replay: Instead of storing old data, train a generative model on previous tasks. When learning new tasks, sample from the generative model to produce synthetic old-task data, and train the policy on both.
  The challenge is that the generative model must produce data that is on-distribution for the old tasks – if it degrades, the policy forgets anyway.
  Example: A GAN that generates images of “dog” from task 1; during task 2 training, the GAN produces dog images that are mixed with real cat images to prevent forgetting how to classify dogs.

• World Action Model (WAM): A combined dynamics and inverse model. It takes (obs, action) and outputs next obs and next action.
  Unlike pure world models (which only predict observations), WAM also predicts the action that led to the next observation – making it directly usable to generate full demonstrations without a separate policy.
  Example: In robot control, WAM can be thought of as “if you see this scene and do this, you’ll get that scene next, and the action you should have taken is X”.

Framework Shift

The key shift is from storing data to storing a generative world model that can recreate data on demand.

Before (mainstream):                               After (this paper):  

[Task 1 data] -> [Train Policy 1]                 [Task 1 data] -> [Train WAM] -> [Discard data]  
[Task 1 data] stored in replay buffer             [WAM] stored (small model)  
[Task 2 data] -> [Train Policy 2 + replay of stored Task 1]  [Task 2 data] -> [Train Policy 2 + pseudo-replay from WAM]  
                                                                                      |  
                                                                               [No stored raw data.]  

From memory heavy, privacy-sensitive to compact model, demonstration-free.
The core shift is from data storage to model-based generation for continual rehearsal.

Expert Assessment

Problem choice: Real and timely. Continual learning in robotics is a well-recognized bottleneck, and the scaling of foundation models makes demonstration storage impractical. The gap is legitimate, though one could argue existing world models could already be used in this way – the paper’s contribution is systematically showing that they work.

Method maturity: Clever integration, not a radical new algorithm. REGEN is essentially a straightforward application of WAMs for generative replay. The novelty lies in (a) using WAMs (which predict both obs and action) rather than separate action/observation generative models, and (b) the recurrent sampling strategy conditioned on task IDs. Simpler approaches (e.g., using a VAE to generate observations only, then a fixed policy to fill in actions) might be simpler but likely underperform. The paper does not compare to such variants.

Experimental integrity: Fair overall. Baselines include fine-tuning (lower bound) and experience replay with real stored data (upper bound). The claim of “up to 50% reduction” is believable. However, the evaluation is on relatively short-horizon tasks (average 5-20 steps). Real-world manipulation tasks with longer horizons might expose the visual degradation bottlenecks the authors mention. Also, the WAM itself is trained on all tasks initially – that’s a potential issue (they use task-specific WAMs? Or one WAM for all?). The paper says “pre-trained on demonstration data” but does the WAM see all tasks beforehand? That would give it privileged information. The method appears to assume that demonstrations of all tasks are available for initial WAM training; that’s okay for a controlled set, but limits applicability to truly open-ended learning. The authors acknowledge this as a limitation.

Writing quality: Reasonable but dense. The section on limitations is strong. The paper could benefit from a clearer diagram early on (they have one, but it’s complicated). I’d rewrite the method section to start with the chef analogy, then dive into equations. The related work is well-covered.

Verdict: weak accept – a solid incremental contribution that convincingly demonstrates a practical trick (WAM for generative replay). Not a breakthrough, but useful for the community.

Takeaways

• Practical technique: If you have a pre-trained world-action model, you can use it to generate pseudo-replays for any number of old tasks without storing any data.
• Key limitation to watch: Generated trajectories degrade over long horizons – check mean action discrepancy (MAD) and pixel RMSE to catch this.
• Transferable idea: Any domain with a forward dynamics model (e.g., video game playing, robotic control, climate modeling) could benefit from this rehearsal strategy.
• Implementation detail: The WAM must predict both actions and observations to be fully autonomous; a pure observation predictor would require a separate policy to fill in actions, which adds complexity and potential inconsistency.