Encoding Memory in World Models with Test-Time Model-Based Reinforcement Learning

Maxime Burchi; Fayaz Ali Dharejo; Radu Timofte

cs.LGNov 6, 2026

Encoding Memory in World Models with Test-Time Model-Based Reinforcement Learning

Authors: Maxime Burchi, Fayaz Ali Dharejo, Radu Timofte

Organizations: Computer Vision Lab, CAIDAS & IFI, University of Würzburg, Germany

Abstract

Model-based reinforcement learning (MBRL) approaches have demonstrated strong performance across diverse domains in recent years, learning accurate world models from agent trajectories. However, despite this progress, current approaches still struggle to surpass human performance on tasks that require long-term memory and planning. In this work, we address the problem of long-term memory in MBRL by moving beyond the commonly studied paradigm of architecture design. We introduce Test-Time Model-Based Reinforcement Learning (TT-MBRL), using gradient descent at test time to encode the memory of episode observations into world model parameters. This enables the agent to learn a distinct world model for each episode, improving memory capabilities compared to traditional approaches that solely rely on contextual memory. We first establish a proof of concept on a simplified two-dimensional variant of the Memory Maze environment, which we refer to as Memory Maze 2D, showing consistent improvements across different world model architectures. We then demonstrate the effectiveness of TT-MBRL in the 3D domain, achieving new state-of-the-art results on the Memory Maze benchmark. Finally, we show that TT-MBRL maintains strong performance on non-memory benchmarks such as DeepMind Control and Atari 100k.

Figures & tables

Figure 1: Performance improvements brought by TT-MBRL on Memory Maze 2D and Memory Maze.

Figure 2: DreamerV3 failure case where the world model predicts an incoherent future after returning to the original position in Memory Maze.

Figure 3: During training, the world model learns to predict the common patterns shared among multiple episode trajectories collected during exploration. The agent can solely rely on contextual memory to identify patterns that are unique to each new episode, which can lead to incoherent future predictions. At test time, the agent can leverage both contextual and neural memory. This allows the world model to make more accurate predictions and learn an improved policy for each new episode.

Panel of Figure 4 — Figure 4: Randomly generated environment map and agent observation for Memory Maze 2D and Memory Maze benchmarks. The agent receives partial observations of the environment and must rely on long-term memory to learn an optimal policy. In contrast, DMC and Atari 100k tasks are mostly fully observable and solvable without extensive memory.

	Mem Maze 2D 9×9	Mem Maze 2D 11×11	Mem Maze 2D 13×13	Mem Maze 2D 15×15
Method	Mem Maze 2D 9×9	Mem Maze 2D 11×11	Mem Maze 2D 13×13	Mem Maze 2D 15×15
Human	$126.2$	$172.0$	$195.0$	$201.0$
Oracle	$137.2$	$215.4$	$276.6$	$311.4$
DreamerV3 (TBTT)	$121.2\pm 3.0\hphantom{0}$	$185.9\pm 11.3$	$108.0\pm 8.8\hphantom{0}$	$\hphantom{0}46.2\pm 7.6\hphantom{0}$
+ TT-MBRL (Ours)	$\mathbf{122.1\pm 1.1\hphantom{0}}$	$\mathbf{189.4\pm 6.8\hphantom{0}}$	$\mathbf{207.2\pm 10.9}$	$\mathbf{185.0\pm 17.6}$
TWISTER	$\hphantom{0}84.8\pm 0.8\hphantom{0}$	$100.0\pm 5.9\hphantom{0}$	$\hphantom{0}95.3\pm 10.2$	$\hphantom{0}72.0\pm 3.2\hphantom{0}$

Table 1: Memory Maze 2D benchmark comparison for all maze sizes (100M env steps). We show the impact of TT-MBRL on the performance for each method. The reported scores are averaged over three runs with 100 evaluation episodes. We also report Human mean scores over 10 episodes.

Method	KL div $\downarrow$	FVD $\downarrow$	PSNR $\uparrow$	SSIM $\uparrow$	LPIPS $\downarrow$
DreamerV3 (TBTT)	$114.6$	$104.5$	$10.3$	$0.552$	$0.283$
+ TT-MBRL (Ours)	$\mathbf{\hphantom{0}66.0}$	$\mathbf{\hphantom{0}55.2}$	$\mathbf{16.3}$	$\mathbf{0.642}$	$\mathbf{0.241}$
TWISTER	$\hphantom{0}28.7$	$\hphantom{0}89.1$	$13.1$	$0.578$	$0.292$
+ TT-MBRL (Ours)	$\mathbf{\hphantom{0}13.2}$	$\mathbf{\hphantom{0}57.1}$	$\mathbf{15.7}$	$\mathbf{0.641}$	$\mathbf{0.262}$
R2I	$\hphantom{0}52.0$	$115.4$	$10.1$	$0.541$	$0.304$
+ TT-MBRL (Ours)	$\mathbf{\hphantom{0}13.8}$	$\mathbf{\hphantom{0}45.7}$	$\mathbf{15.8}$	$\mathbf{0.675}$	$\mathbf{0.246}$

Table 2: Evaluation of the impact of TT-MBRL on world model predictions on Memory Maze 2D 15×15. We aggregate all metrics over 100 unseen trajectories of 64 frames from distinct episodes.

	Mem Maze 9×9	Mem Maze 11×11	Mem Maze 13×13	Mem Maze 15×15
Method	Mem Maze 9×9	Mem Maze 11×11	Mem Maze 13×13	Mem Maze 15×15
Human	$26.4$	$44.3$	$55.5$	$67.7$
Oracle	$34.8$	$58.0$	$74.5$	$87.7$
IMPALA †	$23.5$	$28.5$	$27.0$	$18.5$
IMPALA † (400M steps)	$26.5$	$35.0$	$36.5$	$27.5$
DreamerV3 (TBTT)	$\mathbf{32.7\pm 0.8}$	$33.0\pm 1.4$	$22.1\pm 1.8$	$\hphantom{0}6.0\pm 0.3$

Table 3: Memory Maze benchmark comparison (100M env steps) for all maze sizes. We show the impact of TT-MBRL on the performance for each method. The reported scores are averaged over three runs with 100 episodes for evaluation.

\dagger

denotes the results taken from Samsami et al. (2024) . We also report TT-MBRL results for R2I with a budget of 400M environment steps over six runs.

Method	KL div $\downarrow$	FVD $\downarrow$	PSNR $\uparrow$	SSIM $\uparrow$	LPIPS $\downarrow$
DreamerV3 (TBTT)	$\hphantom{0}21.7$	$\hphantom{0}67.7$	$25.0$	$0.771$	$0.211$
+ TT-MBRL (Ours)	$\hphantom{0}\mathbf{20.1}$	$\hphantom{0}\mathbf{61.7}$	$\mathbf{26.2}$	$\mathbf{0.789}$	$\mathbf{0.190}$
TWISTER	$\hphantom{0}51.8$	$\hphantom{0}63.8$	$24.6$	$0.756$	$0.230$
+ TT-MBRL (Ours)	$\mathbf{\hphantom{0}46.3}$	$\mathbf{\hphantom{0}49.1}$	$\mathbf{26.2}$	$\mathbf{0.792}$	$\mathbf{0.187}$
R2I	$\hphantom{0}61.3$	$114.1$	$21.1$	$0.687$	$0.319$
+ TT-MBRL (Ours)	$\mathbf{\hphantom{0}38.3}$	$\mathbf{85.2}$	$\mathbf{24.0}$	$\mathbf{0.750}$	$\mathbf{0.252}$

Table 4: Evaluation of the impact of TT-MBRL on world model predictions on Memory Maze 15×15. We aggregate all metrics over 100 unseen trajectories of 64 frames from distinct episodes.

Panel of Figure 5 — Figure 5: R2I world model KL divergence over 64 predicted steps, evaluation episodes histogram, and relative improvement in obtained rewards brought by TT-MBRL on Memory Maze 15×15.

Panel of Figure 6 — Figure 6: Impact of TT-MBRL on world model predictions. We let the agent explore and learn about the environment, then we compare world model predictions to real observations. First row shows unseen real environment observations. Second row shows world model predictions with TT-MBRL. Third row shows world model predictions without TT-MBRL. Without test-time optimization, the model accumulates prediction errors and struggles to make accurate long-term predictions. This makes it difficult to learn an optimal policy for the actor-critic.

Method	w/o TT-MBRL	Train Data Only	Test Data Only		Test+Train Data
		Train Data Only	No Prefill	1024 Prefill	No Prefill	1024 Prefill
DreamerV3	$\hphantom{0}46.2$	$\hphantom{0}47.0$	$117.6$	$173.5$	$\mathbf{185.0}$	$149.6$
TWISTER	$\hphantom{0}72.0$	$\hphantom{0}69.2$	$167.6$	$210.1$	$\mathbf{240.1}$	$207.1$
R2I	$171.6$	$171.7$	$218.8$	$240.1$	$\mathbf{244.3}$	$237.5$

Table 5: Impact of replay buffer sampling strategy on TT-MBRL performance on Memory Maze 2D 15×15. Combining training and test-time trajectories (Test+Train Data, No Prefill) consistently achieves the best performance across all world model architectures.

Memory Maze 2D 15 $\times$ 15
Method	w/o TT-MBRL	Episode Score		Eval. Time (hours)
		Frozen	Unfrozen	Frozen	Unfrozen
DreamerV3	$\hphantom{0}46.2$	$\mathbf{185.0}$	$181.4$	$\mathbf{0.31}$	$0.35$
TWISTER	$\hphantom{0}72.0$	$240.1$	$236.4$	$\mathbf{0.25}$	$0.27$
R2I	$171.6$	$244.3$	$\mathbf{247.5}$	$\mathbf{0.06}$	$0.06$
Memory Maze 15 $\times$ 15

Table 6: Impact of updating encoder/decoder using the reconstruction loss on performance and evaluation time. Frozen denotes our default TT-MBRL setting where encoder and decoder parameters are fixed at test time; Unfrozen allows them to be updated. We report the evaluation time per episode.

Table 13

Appendix figures & tables32 assets

Supplementary material from the paper’s appendix.

Appendix

Method	PAD	TT-MBRL	w/o KL loss	w/o rwd-cont loss	w/o actor-critic loss
DreamerV3	$\hphantom{0}25.5$	$185.0$	$\hphantom{0}12.3$	$\mathbf{187.9}$	$\hphantom{0}47.8$
TWISTER	$\hphantom{0}35.5$	$240.1$	$\hphantom{0}12.9$	$\mathbf{240.9}$	$\hphantom{0}50.2$
R2I	$\hphantom{0}42.3$	$\mathbf{244.3}$	$\hphantom{0}14.9$	$242.7$	$158.5$

Appendix

Table 9: Ablation study on Test-Time Optimization objective.

Method	w/o TT-MBRL	100/0 (train-only)	75/25	50/50	25/75	0/100 (test-only)
DreamerV3	$\hphantom{0}46.2$	$\hphantom{0}47.0$	$160.9$	$\mathbf{185.0}$	$179.4$	$117.6$
TWISTER	$\hphantom{0}72.0$	$\hphantom{0}69.2$	$\mathbf{244.5}$	$240.1$	$235.2$	$167.6$
R2I	$171.6$	$171.7$	$242.5$	$244.3$	$\mathbf{249.7}$	$218.8$

Appendix

Table 10: Ablation study on TT-MBRL train/test sampling ratio.

Policy Step/Grad Step	w/o TT-MBRL	32	16	8	4	2	1	1/2
DreamerV3	$\hphantom{0}46.2$	$\hphantom{0}79.9$	$107.3$	$132.1$	$152.9$	$165.7$	$\mathbf{185.0}$	$178.7$
TWISTER	$\hphantom{0}72.0$	$108.9$	$139.8$	$178.3$	$198.5$	$225.8$	$240.1$	$\mathbf{245.6}$
R2I	$171.6$	$182.3$	$190.2$	$205.8$	$213.6$	$230.9$	$244.3$	$\mathbf{251.8}$

Appendix

Table 11: Study on test-time gradient update frequency.

Environment	#Tasks	Method	GPU	#GPUs	#Env Steps	Training Time (days/task)
					#Env Steps	Training Time (days/task)
Memory Maze 2D	4	DreamerV3	L40	1	100M	8
Memory Maze 2D	4	TWISTER	L40	1	100M	6
Memory Maze 2D	4	R2I	L40	1	100M	2.9
Memory Maze	4	DreamerV3	L40	1	100M	13.5
Memory Maze	4	TWISTER	L40	1	100M	23

Appendix

Table 12: Resources and execution time used for each experiments.

Figure 7: Randomly generated environment map for Memory Maze 2D (9×9, 11×11, 13×13 and 15×15). The agent receives a top-down $3\times 3$ local context centered on its current position and must reach the target color to receive the next reward.

	Mem Maze 2D 9×9	Mem Maze 2D 11×11	Mem Maze 2D 13×13	Mem Maze 2D 15×15
Parameter	Mem Maze 2D 9×9	Mem Maze 2D 11×11	Mem Maze 2D 13×13	Mem Maze 2D 15×15
Number of actions	$5$	$5$	$5$	$5$
Observation dim	$75$	$75$	$75$	$75$
Number of colors	$3$	$4$	$5$	$6$
Number of rooms	$3-4$	$4-6$	$5-6$	$9$
Room size	$3-5$	$3-5$	$3-5$	$3$

Appendix

Table 13: Memory Maze 2D environment details.

Step time (ms)	9×9	11×11	13×13	15×15
Memory Maze 2D	$\hphantom{0}0.16$	$\hphantom{0}0.16$	$\hphantom{0}0.16$	$\hphantom{0}0.16$
Memory Maze	$47.28$	$61.82$	$76.48$	$90.36$

Appendix

Table 14: Environment efficiency comparison. Episode average step time measured on a single Intel Core i7-12700 4.9GHz CPU thread and a RTX 3090 GPU.

Table 15: World model sizes. N/A refers to the use of a MLP encoder for low-dimensional inputs, and the removal of the AC-CPC objective in the case of TWISTER.

Hyperparameter	Memory Maze 2D	Memory Maze
General
Batch Size (B)	64	16
Sequence Length (T)	64	64
Environment Instances	8	32
Policy Step per Grad Step	64	64
Replay Buffer Capacity	10M	10M

Appendix

Table 16: DreamerV3 hyperparameters.

Hyperparameter	Memory Maze 2D	Memory Maze
General
Batch Size (B)	4	4
Sequence Length (T)	1024	1024
Environment Instances	8	32
Policy Step per Grad Step	64	64
Replay Buffer Capacity	10M	10M

Appendix

Table 17: TWISTER hyperparameters.

Hyperparameter	Memory Maze 2D	Memory Maze
General
Batch Size (B)	4	4
Sequence Length (T)	1024	1024
Environment Instances	8	40
Policy Step per Grad Step	64	64
Replay Buffer Capacity	10M	10M

Appendix

Table 18: R2I hyperparameters.

Figure 8: Memory Maze 2D training curves (100M environment steps).

Figure 9: Memory Maze training curves (100M environment steps).

Figure 10: Memory Maze training curves (400M environment steps).

Figure 11: Evaluation episode return distribution on Memory Maze 2D (100M environment steps). We compare the impact of TT-MBRL on episode return for DreamerV3, TWISTER and R2I.

Figure 12: Evaluation episode return distribution on Memory Maze 2D 15×15. We compare the impact of training trajectories regularization and replay buffer prefill on TT-MBRL efficacy.

Figure 13: Evaluation episode return distribution on Memory Maze (100M environment steps). We show the impact of TT-MBRL on episode return for DreamerV3, TWISTER and R2I.

Figure 14: Evaluation episode return distribution on Memory Maze (400M environment steps). We show the impact of TT-MBRL on episode return for R2I.

Figure 15: Evaluation episode return distribution on the DeepMind Control Suite (DMC). We show the impact of TT-MBRL on episode return for DreamerV3.

Figure 16: Evaluation episode return distribution on the Atari 100k benchmark. We show the impact of TT-MBRL on episode return for DreamerV3.

Figure 17: Relative improvement in obtained rewards brought by TT-MBRL on Memory Maze 2D.

Figure 18: Relative improvement in obtained rewards brought by TT-MBRL on Memory Maze (100M environment steps).

Figure 19: Relative improvement in obtained rewards brought by TT-MBRL on Memory Maze (400M environment steps).

Figure 20: DreamerV3, TWISTER and R2I world model prediction metrics on Memory Maze 2D 15×15.

Figure 21: DreamerV3, TWISTER and R2I world model prediction metrics on Memory Maze 15×15 (100M environment steps).

Figure 22: R2I world model prediction metrics on Memory Maze 15×15 (400M environment steps).

Figure 23: DreamerV3 world model predictions on Memory Maze.

Figure 24: TWISTER world model predictions on Memory Maze.

Figure 25: R2I world model predictions on Memory Maze.

Figure 26: DreamerV3 world model predictions on Memory Maze 2D.

Figure 27: TWISTER world model predictions on Memory Maze 2D.

Figure 28: R2I world model predictions on Memory Maze 2D.

Model-Based RL

Model-based reinforcement learning uses environment models to plan actions and learn policies with fewer real-world interactions.

See all 257 papers

MemWM: Associative Memory-Enhanced Latent World Models for Long-Horizon Planning
Oct 14, 2026Yuwei Sun, Siyu ZhuWorld Model Learning World Models for Robotics
Shanghai Academy of AI for Science · Fudan University
T $^2$ Mem: Learning Test-Time Memory for Robotics
Sep 29, 2026Yize Liu, Huang Huang, Yining Hong +4Robot Policy Learning Test-Time Adaptation of VLA Models
Stanford University · NVIDIA · University of Michigan, Ann Arbor
Next Embedding Prediction Makes World Models Stronger
Mar 3, 2026George Bredis, Nikita Balagansky, Daniil Gavrilov +1Model-Based RL Partially Observable RL
T-Tech.

Oct 14, 2026J/K move · Enter open · S save

MemWM: Associative Memory-Enhanced Latent World Models for Long-Horizon Planning

Yuwei Sun, Siyu Zhu

Shanghai Academy of AI for Science · Fudan University

World Model Learning World Models for Robotics World Model-Based Planning Latent World Models Joint-Embedding Predictive Architecture

Open paper AskPDF arXiv

World models trained via latent dynamics offer a scalable path for embodied planning, yet existing methods suffer from representational collapse, computational overhead, and limited long-horizon capability. We introduce Memory World Model (MemWM), an associative memory-enhanced latent world model that augments JEPA-style prediction with a content-addressable memory bank. MemWM learns high-level historical representations directly from observations and leverages them to condition predictions via adaptive layer normalization, complementing low-level action-conditioned dynamics within a shared latent space. To prevent memory slots from collapsing into identical representations, we propose two complementary losses that encourage structural diversity across slots and balanced utilization across timesteps. On PushT, TwoRoom, Reacher, and OGBCube, MemWM attains an average success rate of 89.9%, matching DINO-WM (89.3%) while using roughly 4

\times

fewer trainable parameters and no pretrained encoder, and improving over LeWM by 4.1%. Additionally, we extend the horizon to 50 timesteps, double the 25 timesteps of the original evaluation, demonstrating MemWM's robustness to long-horizon generalization.

Abstract

Model-Based RL

Explore similar work

MemWM: Associative Memory-Enhanced Latent World Models for Long-Horizon Planning

T2^22Mem: Learning Test-Time Memory for Robotics

Next Embedding Prediction Makes World Models Stronger

T $^2$ Mem: Learning Test-Time Memory for Robotics