Forgetting in World Models Does Not Follow Task Distance

A component-level benchmark for catastrophic forgetting in world models, and two negative results it produced. Paper: paper/WMF.pdf (26 pages); its sources, tables and figures are in paper/, and all 375 result files are in this repository.

Existing continual-learning suites evaluate policies, or world models as integrated systems. This one isolates the transition component M -- the part that carries latent state forward, and that every downstream use of the model depends on -- and asks what a task switch does to it.

The two results

1. The labelled distance axis does not order forgetting. The benchmark is built around three levels of dynamic distance per family, expecting forgetting to grow along them. It peaks at the medium level in all three families instead, and across the nine cells the label carries no rank information at all (Spearman rho = 0.00, against 0.58 and 0.43 for the two quantities we measure). The cleanest case needs no cross-family comparison: Gymnasium's maximum level is its medium perturbation plus two more, on the same task pair, and it forgets less. Rerun at twice the training budget the ordering holds, and the obvious objection is refuted -- the more heavily perturbed task turns out to be easier to fit, because a cheetah at triple mass barely moves.

2. The forgetting happens in the encoder, where the metrics cannot see it. Component-level metrics must hold a representation fixed to attribute a change to the component, so they are blind to drift in that representation. Fine-tuning's held-out reconstruction of the first task degrades by a factor of 811 while its prediction fidelity reads -5.58, i.e. improvement. EWC makes the dissociation exact and predicts it from its own definition: its Fisher information is identically zero on every encoder parameter, so it preserves the transition component almost perfectly (-0.03) with an encoder degraded by 800. A benchmark reporting only the isolated metric would certify that neither model forgot anything.

Five methods are characterised, not ranked -- ours included, and the benchmark's clearest finding about it is a failure mode.

Method	Description
finetuning	Sequential training, no protection (lower bound)
replay_infinite	Retains all data from all tasks (upper bound)
ewc	Elastic Weight Consolidation (Kirkpatrick et al., 2017)
progressive_nets	Progressive Networks (Rusu et al., 2016)
ug_mtm	UG-MTM: uncertainty-gated mixture of transition models (ours)

Scale

Three families -- MiniGrid (discrete), Gymnasium/MuJoCo HalfCheetah (continuous, variable physics) and DMControl (visual) -- at three dynamic distances each, five methods, ten seeds in each of the six cells that discriminate between methods and five in the three that do not: 375 runs, plus 75 from-scratch reference pairs that define forward transfer and the measured distance. No run dropped a NaN step, and every run shares one protocol. A doubled-budget probe and a four-task sequence run alongside, in results-2x/ and results-seq/.

Working on this with an AI agent? CLAUDE.md is the map, the house rules and the traps — written for an agent arriving cold.

Repository layout

cf_worldmodels/
├── configs/
│   ├── benchmark/          # One YAML per environment family, 3 distance levels each
│   └── models/             # RSSM baseline and UG-MTM hyperparameters
├── src/
│   ├── envs/               # Wrappers normalizing every family to (64,64,3) float32 [0,1]
│   ├── models/             # ConvVAE, RSSM baseline, UG-MTM
│   ├── baselines/          # Fine-tuning, infinite replay, EWC, Progressive Nets
│   ├── benchmark/          # PF / RD / WMF / FT metrics, d_param / d_trans distances
│   └── utils/              # Replay buffer, checkpointing, logging
├── experiments/            # Benchmark runners and plotting scripts
├── tests/                  # Test suite (434 tests)
├── results/                # One directory per run (metrics.json); 375 runs
├── results-2x/             # Doubled-budget probe (10 cells)
└── results-seq/            # Four-task sequence (25 runs)

paper/
├── main.tex                # The full paper, 26 pages -> WMF.pdf
├── main_workshop.tex       # The 8-page version
├── tables/                 # GENERATED by experiments/export_tables.py
└── figures/

Installation

conda env create -f cf_worldmodels/environment.yml

conda activate cf_worldmodels

Or with pip into an existing Python 3.11 environment:

pip install -r cf_worldmodels/requirements.txt

MuJoCo and dm_control require a working OpenGL/EGL setup for offscreen rendering. All commands below are run from the cf_worldmodels/ directory.

Reproducing the results

run_full_benchmark.py reproduces every number in the paper. It skips any run whose metrics.json already exists, so it is safe to interrupt and resume:

python experiments/run_full_benchmark.py

The training protocol is not defined in the runner. Every value comes from the protocol: block of configs/benchmark/<family>.yaml, is printed before training starts, and is recorded in each metrics.json — so the protocol table in the paper is generated from the results rather than written by hand. Runs cached under a different protocol are refused rather than silently averaged into the same cell.

Before the five methods of a cell, the runner trains the pair of reference models that forward transfer and d_trans are defined against: one plain RSSM per environment, from scratch, per (family, distance, seed). They are cached in results/_reference/ and shared by every method, since neither quantity depends on the continual-learning method. To skip them — ft and d_trans are then stored as null, never as 0:

python experiments/run_full_benchmark.py --skip-reference

Print the effective protocol and the run plan without training anything:

python experiments/run_full_benchmark.py --dry-run

Restrict the grid, or override a protocol field explicitly:

python experiments/run_full_benchmark.py --families minigrid --methods ewc --seeds 0 1 --steps 2000

Measure how far task A is actually learned, and how that changes with the training budget (one run at the largest budget, evaluated at each multiple of n_train along the way):

python experiments/convergence_A.py --family minigrid --multipliers 1 2 5 10

Regenerate the main figure from the stored metrics (no training required): one row per reported metric, one column per family, with the measured d_trans on the X axis wherever the runs carry it:

python experiments/plot_final.py

Aggregate every cell, with the task-A quality columns alongside the forgetting metrics, and optionally a seed-paired comparison of two methods:

python experiments/summarize_results.py --compare replay_infinite finetuning

It refuses to average runs that were produced under different protocols, and it reports an exact paired permutation p rather than a t-test. The floor matters: with 5 seeds the smallest two-sided p an exact test can return is 2/2^5 = 0.0625, so a parametric p-value in that regime describes the normality assumption more than the data. The six discriminating cells were extended to ten seeds, where the floor is 0.002, and four of them reach it.

Run a single method/family/distance combination:

python experiments/run_benchmark.py --method ug_mtm --config configs/benchmark/minigrid.yaml --distance distance_min --seeds 0 1 2 3 4 --no_wandb

Determinism

Re-running a cell with the same seed reproduces its metrics bit-for-bit. Getting there needs more than seeding torch and numpy, because two independent sources of nondeterminism sit outside them:

• The environments own RNGs no global seed reaches. Gymnasium environments and action spaces each carry their own generator, seeded from OS entropy, and dm_control randomizes the initial state through the task's own random argument. Left unseeded, two runs of the same seed collected different rollouts and therefore trained on different data. BaseEnv.seed() seeds both the episode RNG and the action sampler; every runner calls it.
• cuDNN picks kernels by heuristic and its GRU backward is nondeterministic by default. src/utils/seeding.py::set_seed() sets cudnn.deterministic = True and cudnn.benchmark = False, which costs some throughput and was verified sufficient to make training bit-identical here.

Both are covered by tests/test_seeding.py, including an end-to-end check that training twice from the same seed yields identical weights.

Results

Every number in the paper is derived from the metrics.json files in this repository, by experiments/summarize_results.py (console) and experiments/export_tables.py (the paper's LaTeX tables). Nothing is transcribed by hand: the second imports the first, so the two can only disagree if the code disagrees with itself.

python experiments/summarize_results.py

Running experiments/run_full_benchmark.py writes one directory per run to results/<method>/<family>_<distance>_<seed>/, containing metrics.json. Checkpoints are not written at all -- every table and figure derives from the metrics, so keeping ~32 MB per run bought nothing.

A note on what preceded this. An earlier version of this benchmark reported five findings. None survives re-measurement, for five instrumentation defects: a collapsed VAE posterior (0 of 32 latent dimensions active, so the transition model received a constant latent regardless of the observation), an evaluation set of Gaussian noise, a next-state objective scored against the wrong target, a mis-specified Gaussian KL, and unseeded environments. The strongest of the five reverses. The paper documents this in a section of its own rather than omitting it.

Metrics

The suite is PF, RD and FT. An earlier description of this benchmark announced a fourth, PIS (Policy Impact Score); it is not part of the suite and never was implemented — see Known limitations.

• PF (Prediction Fidelity) — NLL(M_k, D_i) - NLL(M_i, D_i). Positive means the model got worse at task i after training on later tasks.
• RD (Rollout Divergence) — mean KL between imagined rollouts of the model before and after the task switch.
• WMF = alpha*PF + beta*RD + gamma*PIS, with alpha=beta=0.4, gamma=0.2. Computed and stored, but not the headline number: RD supplies 78-97% of it, so summarize_results.py reports PF and RD side by side and prints WMF under a heading that says what it is, next to the share of it that comes from RD. It is there to reproduce the previous paper's number, and its gamma term is evaluated at zero — which is what that number was computed with.
• FT (Forward Transfer) — recon_B(trained from scratch) - recon_B(pretrained on A), on held-out task-B frames in pixel space. Positive means knowledge of task A helped learn task B under the same budget and the same data. The from-scratch arm comes from a reference model trained on task B alone, one per (family, distance, seed), shared by all five methods.
• task_A_fit_gain — NLL(random init, D_i) - NLL(post-task-i model, D_i). This is what earlier releases called FT; it measures how well task i was learned, and no task-B data enters it.

All of these are evaluated on D_i: held-out task-A rollouts, collected separately from the training buffer and encoded once by the post-task-A model (protocol.build_latent_eval_dataset), so that both models being compared are scored on identical inputs and identical targets. That also means they are blind to drift in the encoder itself — see Known limitations.

Alongside them, each run records how well task A was learned in the first place, because a forgetting benchmark has to show there was something to forget:

• heldout_reconstruction_A_after_task_A / ..._after_task_B — squared reconstruction error per frame, in pixel space, on held-out task-A frames. This is the only quality signal that is comparable across training budgets: the latent NLL is scored against latents the model itself produces, and that target moves as the encoder trains.
• heldout_reconstruction_B_after_task_B and heldout_reconstruction_B_from_scratch — the two arms of FT.
• nll_A_after_task_A, nll_A_after_task_B, nll_A_random_init — the three NLLs that PF and task_A_fit_gain are built from, so both stay decomposable.
• initial/final_reconstruction_loss_A and _B, plus a 20-point curve for each task, and n_nan_steps_A/B for steps dropped as non-finite.

Dynamic distance between two tasks is measured by d_param (normalized L2 distance between physics parameter vectors, for the variable-physics pairs) and d_trans (Eq. 9: expected KL between one transition model per environment, each trained on its own task from scratch, all families). d_trans is a property of the task pair and the seed rather than of the method, so it is computed once per (family, distance, seed) alongside the forward-transfer reference and stored in results/_reference/.

Tests

python -m pytest

434 tests covering the models, baselines, metrics, distances, buffer, checkpoint format, seeding, protocol resolution and config consistency. The tests marked integration build real MiniGrid / MuJoCo / dm_control environments; skip them with:

python -m pytest -m "not integration"

Known limitations

These are properties of the released code that reviewers and reusers should be aware of before building on it.

1. Uncertainty routing works at large dynamic distance and inverts at moderate distance. Measured as AUC = P(u_B > u_A) on held-out transitions after training on task A (0.5 = no discrimination): the MC-dropout signal reaches 0.864 on FourRooms→KeyCorridor but falls to 0.294 on Empty-8x8→FourRooms, i.e. task B looks less uncertain than task A and the gate routes the wrong way. UG-MTM's premise holds only in the first regime.

2. PF and RD are blind to forgetting in the encoder. They are evaluated on latents that were encoded once, by the post-task-A model, and compute_nll never calls encode — so they measure drift in the GRU and the stochastic head within a frozen latent basis. Measured on MiniGrid distance_med (seed 999): fine-tuning's held-out task-A reconstruction degrades by a factor of 112 (6.49 → 725.27 squared error per frame) while its PF comes out negative (−1.78). This is deliberate — the benchmark's scope is the transition component — but it means WMF is not a measure of how much the world model as a whole forgot. Every run records heldout_reconstruction_A_after_task_{A,B} so both can be read side by side.

3. The training scale is small. 20 episodes of a random policy per task and 5000 gradient updates at batch 8, sequence length 5; UG-MTM's training-time MC-dropout budget is 3 passes. Every one of those values is declared in the protocol: block of the family config, recorded in each metrics.json, and printed before training. Task A does get learned at this scale — held-out reconstruction reaches 5.3e-04 per pixel, RMSE ≈ 0.023 on [0,1] — and experiments/convergence_A.py measures how that changes with the budget.

4. Forward transfer is measured in pixels, and replay's number needs a caveat. The two arms of FT are models with unrelated latent bases, so a latent NLL would score one of them in the other's coordinates; held-out pixel reconstruction is the one scale they share. Separately, replay_infinite trains on A+B during the task-B phase, so half of its gradient steps go to task-A data: its FT mixes transfer with a halved effective budget on B.

   Earlier releases reported a quantity named FT that was computed from the post-task-i model on task-i data alone. No task-B data entered it, so methods sharing an architecture got identical values by construction — fine-tuning and infinite replay differed by exactly 0.000 across 5 seeds and two distance levels. That quantity is still stored, under the name task_A_fit_gain.

5. EWC protects the transition component and nothing else. Its Fisher is defined over log P(z'|z, a), so it is exactly zero on every encoder parameter: the penalty cannot constrain the VAE, and EWC's pixel-space reconstruction of task A degrades like fine-tuning's. It is also zero on gru.weight_hh, because the Fisher set is single transitions started from h = 0 — the recurrent pathway is unprotected, and compute_nll scores from h = 0 too, so PF does not see it either. RD, which rolls out 15 steps, does.

6. No ablation study ships with this release. The previous run_ablations.py built an overridden config and then never passed it to the training routine, which reloaded the unmodified YAML from disk — so all five ablations silently ran plain UG-MTM. It has been removed rather than left in place producing misleading output.

7. Gate scaling uses only the final timestep's gates. UG_MTM.transition clears and re-registers its backward hooks on every call, so after unrolling a sequence the gradients for the whole sequence are scaled by the gates computed at the last step.

8. PIS was announced and is not part of the suite. An earlier description of this benchmark listed a fourth metric, PIS (Policy Impact Score), meant to score how much a task switch costs a policy. It was never implemented: measuring it means training a controller inside the model's imagination and evaluating it in the real environment, and no controller ships here. It is withdrawn rather than reported — the suite is PF, RD and FT — and pis is stored as null, the same way ft and d_trans are when their reference model was skipped. Runs produced before this change stored 0.0; the aggregation treats null and 0.0 alike here, because the gamma term of WMF was evaluated at zero either way, which is also what the previous paper's WMF numbers were computed with.

Citation

@misc{perezbazarot2026wmf,
  title  = {Forgetting in World Models Does Not Follow Task Distance:
            A Component-Level Benchmark and Two Negative Results},
  author = {P{\'e}rez Bazarot, Jes{\'u}s},
  year   = {2026},
  note   = {Code and data: https://github.com/PersusUS/WorldModelsBenchmark}
}

License

MIT — see LICENSE.