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.
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
Human
126.2
172.0
195.0
201.0
Oracle
137.2
215.4
276.6
311.4
DreamerV3 (TBTT)
121.2±3.00
185.9±11.3
108.0±8.80
046.2±7.60
+ TT-MBRL (Ours)
122.1±1.10
189.4±6.80
207.2±10.9
185.0±17.6
TWISTER
084.8±0.80
100.0±5.90
095.3±10.2
072.0±3.20
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 ↓
FVD ↓
PSNR ↑
SSIM ↑
LPIPS ↓
DreamerV3 (TBTT)
114.6
104.5
10.3
0.552
0.283
+ TT-MBRL (Ours)
066.0
055.2
16.3
0.642
0.241
TWISTER
028.7
089.1
13.1
0.578
0.292
+ TT-MBRL (Ours)
013.2
057.1
15.7
0.641
0.262
R2I
052.0
115.4
10.1
0.541
0.304
+ TT-MBRL (Ours)
013.8
045.7
15.8
0.675
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
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)
32.7±0.8
33.0±1.4
22.1±1.8
06.0±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. † 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 ↓
FVD ↓
PSNR ↑
SSIM ↑
LPIPS ↓
DreamerV3 (TBTT)
021.7
067.7
25.0
0.771
0.211
+ TT-MBRL (Ours)
020.1
061.7
26.2
0.789
0.190
TWISTER
051.8
063.8
24.6
0.756
0.230
+ TT-MBRL (Ours)
046.3
049.1
26.2
0.792
0.187
R2I
061.3
114.1
21.1
0.687
0.319
+ TT-MBRL (Ours)
038.3
85.2
24.0
0.750
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.
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.
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
No Prefill
1024 Prefill
No Prefill
1024 Prefill
DreamerV3
046.2
047.0
117.6
173.5
185.0
149.6
TWISTER
072.0
069.2
167.6
210.1
240.1
207.1
R2I
171.6
171.7
218.8
240.1
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.
Method
w/o TT-MBRL
Episode Score
Eval. Time (hours)
Frozen
Unfrozen
Frozen
Unfrozen
Memory Maze 2D 15 × 15
DreamerV3
046.2
185.0
181.4
0.31
0.35
TWISTER
072.0
240.1
236.4
0.25
0.27
R2I
171.6
244.3
247.5
0.06
0.06
Memory Maze 15 × 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
025.5
185.0
012.3
187.9
047.8
TWISTER
035.5
240.1
012.9
240.9
050.2
R2I
042.3
244.3
014.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
046.2
047.0
160.9
185.0
179.4
117.6
TWISTER
072.0
069.2
244.5
240.1
235.2
167.6
R2I
171.6
171.7
242.5
244.3
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
046.2
079.9
107.3
132.1
152.9
165.7
185.0
178.7
TWISTER
072.0
108.9
139.8
178.3
198.5
225.8
240.1
245.6
R2I
171.6
182.3
190.2
205.8
213.6
230.9
244.3
251.8
Appendix
Table 11: Study on test-time gradient update frequency.
#Env Steps
Training Time (days/task)
Environment
#Tasks
Method
GPU
#GPUs
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×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
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
00.16
00.16
00.16
00.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.
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× 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.
Yuwei Sun, Siyu Zhu
Shanghai Academy of AI for Science · Fudan University
Memory-dependent robotic manipulation requires policies to use information that is no longer available in the current observation. Retaining history alone is insufficient: memory must preserve information that supports future actions. One challenge is whether a memory-free foundation model can learn to retain and use historical information from action demonstrations alone, without external memory support. We introduce T2Mem, a framework that develops this capability within a pretrained vision-language-action policy, without external reasoning models or memory-specific annotations. T2Mem uses test-time training to encode observation history into compact fast weights through online self-supervised updates, avoiding repeated processing of the full history. An observation-grounded interface extracts vision-language information for memory formation and supplies retrieved context to the action expert. Action supervision shapes what the memory learns to retain and use, while alternating memory-policy learning gives each component a fixed counterpart during optimization. Across 16 RoboMME tasks, T2Mem improves average success from 17.93% to 56.83% over the memory-free base policy and outperforms the recurrent-memory methods reported in the benchmark, while controlled profiling indicates at least 3x inference speedup over explicit methods. Project website: https://yzliu84.github.io/T2MEM-project/
Yize Liu, Huang Huang, Yining Hong +4
Stanford University · NVIDIA · University of Michigan, Ann Arbor
Capturing temporal dependencies is critical for model-based reinforcement learning (MBRL) in partially observable, high-dimensional domains. We introduce NE-Dreamer, a decoder-free MBRL agent that leverages a temporal transformer to predict next-step encoder embeddings from latent state sequences, directly optimizing temporal predictive alignment in representation space. This approach enables NE-Dreamer to learn coherent, predictive state representations without reconstruction losses or auxiliary supervision. On the DeepMind Control Suite, NE-Dreamer matches or exceeds the performance of DreamerV3 and leading decoder-free agents. On a challenging subset of DMLab tasks involving memory and spatial reasoning, NE-Dreamer achieves substantial gains. These results establish next-embedding prediction with temporal transformers as an effective, scalable framework for MBRL in complex, partially observable environments.
George Bredis, Nikita Balagansky, Daniil Gavrilov +1