The World Model Remembers, the Actor Forgets: Dream Rehearsal for Continual Model-Based RL

The World Model Remembers, the Actor Forgets: Dream Rehearsal for Continual Model-Based RL

Paper: The World Model Remembers, the Actor Forgets: Dream Rehearsal for Continual Model-Based RL

Author: Gurp Nijjer

Affiliation: Quantegra Research, Canada

Link: arXiv:2607.19749 | Code: GitHub

---

One-Sentence Summary

This paper reveals that in continual learning, DreamerV3-class models forget the actor policy but not the world model — the world model retains all measurable knowledge of old tasks (reward discrimination, value estimates, termination structure), but the policy-gradient channel cannot convert it into behavior. By interleaving dream self-imitation (grading imagined trajectories then behavior-cloning the top 25%) during training, it achieves a task-label-free, parameter-constant continual learner: 3/3 seeds pass four-task chains (plain replay 0/3) and eight-task chains.

---

Background and Motivation

Catastrophic forgetting in continual RL is a core challenge. Standard remedies regularize parameters (EWC), grow networks, or replay past data. In model-based RL, replay has a natural advantage — the world model trains on buffered experience, and DreamerV3-class agents with never-clear buffers are the field's working answer. However, this carries an untested assumption: "protect the world model and the policy will follow."

This paper first tests this assumption at the component level, finding the first half correct (the world model retains everything) but the second half wrong (the actor forgets anyway). Forgetting in this regime is not a memory problem but a channel problem: policy-gradient learning fails to convert available accurate training signals in imagination into behavior.

Figure 1: The recovery race. On frozen world model, dream self-imitation (blue) recovers 3/3 seeds; RL-in-imagination (orange) fails 0/3.

---

Method Details

1. Component-Level Forgetting Diagnosis

Under never-clear replay, component-level probes ($n=3$ seeds) measure each component's retention of old tasks:

Reward knowledge: Discrimination $D = \text{mean}(\hat{r}_{\text{success}}) - \text{mean}(\hat{r}_{\text{zero}})$, retention ratio:

$$ R = \frac{D_{\text{final}}}{D_{\text{own-phase}}} \approx 0.99/1.06/1.01 $$

The reward head never forgets. On the degraded policy's own visited states $D=0.87$ (vs 0.90 on curated episodes) — the signal is intact where re-teaching needs it.

Value estimates: Critic means on old-task states remain high and rise across phases ($0.84 \rightarrow 0.91$).

Termination model: Continuation-head discrimination at true terminal steps: 0.95-1.0 at every checkpoint.

Representation drift: Frozen-action-margin probe shows old-task action preferences degrade (margin $-0.25$ vs $+4.5$ and $+2.7$). Co-trained heads track drift; frozen heads decay under it; the actor — co-trained but through the RL channel — decays too.

2. Recovery from the World Model Alone

Freeze the entire final world model, retrain only the drifted actor. Same starts, same budget (20k updates), zero new environment steps:

RL-in-imagination (standard DreamerV3 actor-critic update): 0/3 seeds pass. One reaches 0.84 then oscillates to 0.38, one collapses to 0.0, one never exceeds 0.57.

Dream self-imitation (roll sampling actor in imagination, grade with frozen reward/value heads, behavior-clone top 25%): 3/3 seeds pass ($0.38 \rightarrow 0.85$, $0.66 \rightarrow 0.92$, $0.66 \rightarrow 0.85$).

3. Dream Rehearsal

After each task phase ends, its buffered episodes become rehearsal starts. Every 2000 environment steps of new-task training, run 50 dream-self-imitation updates per prior task. Trajectory scoring uses the realized-first rule with self-imitation advantage gate:

$$ \text{imitate if } (R - V(s))_+ > 0 $$

Only imitate when realized return exceeds the critic's estimate. Clone the top 25% of trajectories, one live actor throughout, no task labels, no frozen policies, no router, no new parameters.

Figure 2: Dream rehearsal algorithm (one training chunk). After every 2000 steps, 50 dream self-imitation updates per prior task.

flowchart TD
    A["Task A Training Phase"] --> B["Never-Clear Buffer
All experience retained"]
    B --> C["World Model Training
Retains old task knowledge"]
    C --> D["Task B Training Phase"]
    D --> E["Every 2000 steps: Dream Rehearsal"]
    E --> F["Imagine from old task states
Current actor samples H=15 steps"]
    F --> G["Realized-first scoring
(R-V(s))+ advantage gate"]
    G --> H["Behavior-clone top 25% trajectories
Update same live actor"]
    H --> D
    H --> I["All tasks retained
3/3 seeds pass"]
    style C fill:#e1f5fe
    style E fill:#fff3e0
    style I fill:#e8f5e9

---

Experimental Results

Four-Task Chain Results

Table 1: Four-task chain per-seed final retention

Seed	DoorKey	SimpleCrossing	LavaGap	MultiRoom	Pass
1	0.960	0.868	0.899	0.751	✓
2	0.961	0.751	0.944	0.800	✓
3	0.964	0.825	0.784	0.768	✓

Dream rehearsal passes 3/3 seeds on all four tasks. Plain replay passes 0/3 (hardest task SimpleCrossing at 0.37). Isolation reference (frozen heads + router) passes 3/5.

Eight-Task Chain Results

Table 2: Eight-task chain per-seed final retention

Seed	DoorKey5	SimpleX	LavaGap	MultiRm	LavaCross	DistShift	DoorKey6	Unlock
1	0.964	0.943	0.947	0.759	0.861	0.960	0.915	0.861
2	0.963	0.919	0.943	0.746	0.865	0.929	0.962	0.888
3	0.901	0.876	0.942	0.754	0.794	0.959	0.962	0.740

All 3 seeds pass all eight tasks. Seven of eight tasks have cross-seed range $\leq 0.07$, showing high stability. The predicted rehearsal dilution did not appear.

Dream vs Real-Episode Cloning

Compared to matched real-episode cloning (CLEAR-style), dream rehearsal's paired difference is +0.13, bootstrap 95% CI [0.07, 0.24], complete seed separation. Dream trajectories outperform real replays because imagination is unconstrained by real terminations and can explore improved trajectories.

Figure 3: Four-task chain final retention comparison across three methods.

Figure 4: Dream scoring rule analysis. Two scoring failure modes and realized-first fix.

Figure 5: Eight-task chain scaling results and stability analysis.

---

Two Scoring Failure Modes

The paper identifies two scoring failure modes and proposes the realized-first fix. First, "critic optimism bias" — the critic's optimistic estimates of imagination assign high advantage scores to low-quality trajectories, causing cloning of trajectories that shouldn't be cloned. Second, "termination unawareness" — imagined trajectories don't naturally terminate at real endings, causing overlong hallucinated trajectories to be incorrectly scored. The realized-first rule uses only realized (not predicted) returns for scoring, achieving AUC 1.0 and top-quartile purity 1.0 on both task profiles.

The scoring function is defined as the difference between realized return and critic's state value for imagined trajectory $ au$:

$$ ext{score}( au) = \sum_{t=0}^{H} \gamma^t \hat{r}_t - V_\phi(s_0) $$

where $\hat{r}_t$ is the world model reward head's predicted reward at imagination step $t$, $\gamma$ is the discount factor, $V_\phi(s_0)$ is the critic's value estimate of the initial state, and $H=15$ is the imagination horizon. Only trajectories with $ ext{score}( au) > 0$ enter the candidate pool; the top $q=25\%$ are selected for behavior cloning.

Termination-aware bootstrap weights future rewards through the continuation head $c_\phi(s_t)$:

$$ ext{score}_{ ext{term-aware}}( au) = \sum_{t=0}^{H} \left(\prod_{k=0}^{t-1} c_\phi(s_k) ight) \gamma^t \hat{r}_t - V_\phi(s_0) $$

This termination-aware design ensures that "hallucinated continuations" (imagination continuing past states that should terminate) don't contribute spurious returns. Reach-probability weighting further adjusts each imagination step's contribution, downweighting rewards generated in uncertain imagination regions.

The behavior cloning loss supervises actions on selected high-quality trajectories:

$$ \mathcal{L}_{ ext{clone}} = -rac{1}{|\mathcal{B}_{ ext{top}}|} \sum_{ au \in \mathcal{B}_{ ext{top}}} \sum_{t=0}^{H} \log \pi_ heta(a_t | s_t) $$

where $\mathcal{B}_{ ext{top}}$ is the set of top-25% scored trajectories and $\pi_ heta$ is the current actor policy. This loss only updates actor parameters $ heta$; the world model and critic train normally unaffected.

Computational Cost Analysis

Dream rehearsal adds ~15% compute overhead at four tasks (50×3=150 updates per 2000-step chunk) and ~350 updates/chunk at eight tasks. However, due to acquisition acceleration of subsequent tasks (e.g., LavaGap learns in 8-10k steps in-chain vs more steps alone), overall training time at $n=8$ is comparable to no-rehearsal. Rehearsal cost grows linearly with task count, partially offset by acquisition speedup.

Limitations

1. Scale: Experiments at MiniGrid scale, 17M-parameter world model, one task ordering per chain length. $n=3$ is below the 5-10 common at scale.
2. Linear rehearsal cost: Dream rehearsal compute grows linearly with task count (~350 updates/2000-step chunk at eight tasks); acquisition speed-up partially offsets but may not sustain at very large task chains.
3. Task selection bias: Eight-task chain excluded long-horizon sparse tasks that could not bootstrap under greedy exploration, reported as "acquisition frontier" not a universal guarantee.

---

Conclusion and Outlook

This paper reveals that forgetting in continual MBRL is a channel problem, not a memory problem — the world model retains everything, but the policy-gradient channel cannot convert it. Dream rehearsal bypasses RL channel instability through supervised self-imitation, achieving 3/3 all-pass on four- and eight-task chains without task labels or new parameters. All experiments were pre-registered; every refuted hypothesis is reported.

Key Insight: "The world model remembers, the actor forgets — forgetting is not a memory problem but a channel problem. Policy gradient cannot convert intact training signals in imagination into behavior, but supervised self-imitation can, and dream rehearsal is the direct application of this insight."