An AI maze navigator remembers by updating its world model

Why a predictive model needs memory

Finding a colored object in a maze requires more than recognizing what is directly ahead. A navigator must retain information about passages and objects encountered earlier, then use that information when the target changes.

In model-based reinforcement learning, an agent learns a world model: a system that predicts how its environment will respond to actions. It uses imagined trajectories to learn which actions are likely to earn rewards. But in a partially observed maze, predictions can become inconsistent with earlier observations. A model may fail to preserve the layout of a place after the agent leaves it.

In a November 2026 preprint, Maxime Burchi, Fayaz Ali Dharejo and Radu Timofte investigate another place to store episode-specific information: the model’s parameters, the numerical values adjusted during learning. Their method, Test-Time Model-Based Reinforcement Learning, or TT-MBRL, continues learning during evaluation.

Learning within an episode

Ordinarily, the evaluated agents use parameters learned during training and carry information about the current episode through contextual representations. TT-MBRL supplements that contextual memory with parameter updates.

At the start of every evaluation episode, the agent copies its trained parameters and creates a temporary replay buffer containing experiences from that episode. As it acts, it repeatedly samples stored trajectories and uses gradient descent to update its world model. The networks that select actions and estimate future rewards also adapt, using trajectories imagined by the updated model.

These changes are episode-specific. The next episode starts again from the original trained parameters, rather than inheriting the previous maze’s adaptations.

A crucial detail is that each update normally mixes equal shares of current-episode trajectories and trajectories retained from training. Learning only from the small amount of experience available early in an episode can lead to overfitting and poor decisions. Mixing old and new data helps constrain that adaptation.

The method also freezes the visual encoder and decoder, which translate observations into internal representations and back. It drops the observation-reconstruction objective during evaluation, while continuing to train dynamics and other predictive components. In the reported comparisons, this reduced evaluation time without materially reducing scores.

Larger mazes show clearer gains

The researchers first built Memory Maze 2D, where the agent sees only a local three-by-three region, with a colored border identifying its current target. They tested three world-model architectures: recurrent DreamerV3, Transformer-based TWISTER and the structured state-space model R2I.

After 100 million training environment steps, the largest, 15-by-15 maze showed substantial improvements. DreamerV3’s mean episode score rose from 46.2 to 185.0 with adaptation. TWISTER rose from 72.0 to 240.1, and R2I from 171.6 to 244.3. These results were averaged over three runs, with 100 evaluation episodes.

The original Memory Maze benchmark presents a harder visual problem: first-person views of a 3D environment. Reaching the indicated colored object earns one reward, after which another target is selected. Remembering previously explored routes therefore remains useful throughout an episode.

For R2I trained for 400 million environment steps, adaptation raised the mean score on the largest 3D maze from 42.0 to 54.1, averaged over six runs. The authors report this as a new state-of-the-art result. It nevertheless remains below the benchmark’s reported human score of 67.7. Gains were smaller, absent or slightly negative in some smaller-maze comparisons; adaptation was not uniformly beneficial.

Evidence beyond the final score

Prediction tests also favored adaptation. Across 100 unseen, 64-frame trajectories from distinct episodes, the adapted models performed better on measures of internal prediction error and reconstructed visual quality.

A separate comparison tested whether simply spending more computation could explain the navigation gains. On the largest 2D maze, additional updates using only training trajectories produced little improvement. Updates incorporating the current episode performed much better, supporting the importance of episode-specific experience rather than extra computation alone.

The authors also tested DreamerV3 on 20 DeepMind Control tasks and 26 Atari games, benchmarks not primarily designed around long-term episodic memory. Aggregate scores improved modestly. That supports compatibility with these tasks, not a guarantee of improvement everywhere.

Memory comes with a runtime cost

The principal limitation is computation during evaluation. Every episode requires optimization, not just action selection. The paper reports 0.65 hours per episode for adapted R2I on the largest 3D maze with its encoder and decoder frozen. That is a measured runtime for this experimental setup, not a universal deployment estimate.

The preprint therefore demonstrates a specific tradeoff: learning during an episode can improve predictive consistency and navigation, but makes evaluation more expensive. It does not establish general human-level memory or practical real-world navigation. Its central result is narrower: parameter updates can complement contextual memory in these tested world-model agents.