METTA: Memory-Enhanced Test-Time Adaptation for Meta-Reinforcement Learning

Ali Najar; Hosein Hasani; Amir Kooshan Fattah; Yasamin Rajabi; Mahdieh Soleymani Baghshah

cs.LGOct 22, 2026

METTA: Memory-Enhanced Test-Time Adaptation for Meta-Reinforcement Learning

Authors: Ali Najar, Hosein Hasani, Amir Kooshan Fattah, Yasamin Rajabi, Mahdieh Soleymani Baghshah

Organizations: Department of Computer Engineering Sharif University of Technology

Abstract

Meta-reinforcement learning aims to train agents that adapt quickly to unseen tasks from limited interaction, but this requires representing task evidence over long interaction histories. Classical recurrent agents compress experience into hidden states, while transformer-based agents retrieve from past trajectories at quadratic memory and computation cost. In this work, we propose Memory-Enhanced Test-Time Adaptation (METTA) for meta-reinforcement learning. It decomposes the fast adaptation phase into two temporal scales: a causal neural memory performs fine-grained within-episode adaptation with linear sequence complexity, while compact episode-level summaries support long-range cross-episode task inference without attending over the full interaction history. Moreover, equipping agents with neural memory provides sufficient expressive capacity and test-time training capability, enabling effective adaptation to unseen held-out tasks and task families at meta-test time. On unseen held-out Meta-World ML10/45 test environments, METTA achieves a success rate over $50\%$ while maintaining a low training budget and short wall-clock time. On long-context sparse-reward PointMaze tasks, it outperforms transformer-based baselines by exhibiting more structured exploration behavior. These results show that METTA provides an efficient path toward more adaptive and scalable long-context meta-RL.

Figures & tables

Figure 1: Exploration behavior on a sparse-reward PointMaze task ( Fu et al., 2020 ) across adaptation episodes. Trajectories are shown for Episodes

1

2

5

, and

10

of the same task. As interaction proceeds, Metta explores new regions and reduces repeated trajectories across episodes, suggesting more effective reuse of past exploration evidence in long-context adaptation. TrMRL ( Melo, 2022 ) is difficult to scale to long multi-episode settings, while ECET ( Shala et al., 2025 ) is limited by its short within-episode context.

Figure 2: Overview of Metta . The current episode is processed by a within-episode neural memory to produce $h_{e,t}^{(L)}$ . At the end of each episode, the final representation is stored as an episode summary $\bar{h}_{e}$ . During later episodes, previous summaries are combined with the current representation to form the context representation $z_{e,t}$ . The actor and critic condition on the current state $s_{e,t}$ and $z_{e,t}$ .

Panel of Figure 3 — Figure 3: Fast adaptation during meta-training and meta-testing. (a) HalfCheetahVel training return and adaptation across meta-test episodes. (b) PointMaze training and unseen held-out test success rate.

Figure 4: Meta-World success rate on ML1, ML10, and ML45. Top: evaluation on unseen task instances from task families used during meta-training. Bottom: evaluation on the Meta-World test benchmark, which contains unseen held-out task families for ML10 and ML45.

Method	Coverage $\uparrow$	Episode Jaccard $\downarrow$	Steps $\downarrow$
Metta	$\mathbf{0.447\pm 0.036}$	$\mathbf{0.792\pm 0.022}$	$\mathbf{194.8\pm 30.1}$
ECET	$0.250\pm 0.087$	$0.935\pm 0.028$	$427.3\pm 44.5$
TrMRL	$0.330\pm 0.031$	$1.000\pm 0.000$	$373.4\pm 16.1$

Table 1: PointMaze exploration at 1.5M environment steps (mean

\pm

SEM).

Panel of Figure 5 — Figure 5: Structure of the learned context representations. Left: t-SNE visualization of $z$ for all episodes. Circles denote training tasks and triangles denote test tasks. Right: pairwise cosine-similarity matrix computed from the final episode representation of each trial.

Panel of Figure 6 — Figure 6: Effect of Metta test-time adaptation. The ML1 Push results report adaptation episodes 1 and 25 within each evaluation trial.

Figure 7: Success versus GPU hours.

Appendix figures & tables13 assets

Supplementary material from the paper’s appendix.

Appendix

x_{b,e,t}=[s_{b,e,t},a_{b,e,t-1},r_{b,e,t-1},d_{b,e,t-1}]

Appendix

Algorithm 1 Training Metta with PPO

x_{b,e,t}=[s_{b,e,t},a_{b,e,t-1},r_{b,e,t-1},d_{b,e,t-1}]

Appendix

Algorithm 2 Meta-Test Evaluation of Metta

Figure 8: Examples of Meta-World environments illustrate the diversity of robotic manipulation tasks in the benchmark. Figures adapted from the official project page https://meta-world.github.io .

Training
Parameter	Meta-World			MuJoCo	PointMaze
Parameter	ML1	ML10	ML45	MuJoCo	PointMaze
num_envs	50	50	180	50	25
rollout_steps	500			100	500
trial_length	5			5	10
TTT

Appendix

Table 2: Hyperparameters used in our experiments.

Method	Per-decision computation	Full-trial computation	History memory
TrMRL	$\mathcal{O}((e-1)T+t)$	$\mathcal{O}(E^{2}T^{2})$	$\mathcal{O}(ET)$
ECET	$\mathcal{O}(T^{2}+e^{2})$	$\mathcal{O}(ET^{3}+E^{3}T)$	$\mathcal{O}(ET)$
Metta + Attention	$\mathcal{O}(1+e)$	$\mathcal{O}(E^{2}T)$	$\mathcal{O}(E)$
Metta + EMA	$\mathcal{O}(1)$	$\mathcal{O}(ET)$	$\mathcal{O}(1)$

Appendix

Table 3: Online inference complexity with respect to episode horizon

T

and number of episodes

E

. Network dimensions and layer counts are treated as constants.

	Metta		TrMRL				ECET ( $E=5$ )				ECET ( $E=25$ )
Metric	TTT-MLP	TTT-Linear	$T=5$	$T=25$	$T=100$	$T=500$	$T=5$	$T=25$	$T=100$	$T=500$	$T=5$	$T=25$	$T=100$	$T=500$
Time (ms/step) $\downarrow$	2.159	1.616	0.551	0.639	1.061	9.223	0.761	0.991	2.119	21.713	1.035	2.151	7.599	105.403
Memory (MiB) $\downarrow$	5.034	2.397	0.335	0.356	0.570	5.333	0.792	0.863	1.701	21.388	2.534	2.857	6.813	101.115

Appendix

Table 4: Measured per-step inference latency and memory usage as context length increases. TTT-MLP and TTT-Linear include the inner updates performed during inference. Lower values are better (

\downarrow

Figure 9: Comparison of alternative within-episode neural memory architectures in Metta on ML10 and PointMaze.

Figure 10: Comparison of cross-episode aggregation mechanisms in Metta on ML10 and PointMaze. All variants use the same within-episode neural memory.

Figure 11: Complete online memory-update ablation on PointMaze, ML10, and ML1 Push. Top: training performance. Bottom: test performance.

Figure 12: ML1 Push test performance at adaptation Episodes 1, 5, and 25 within each evaluation trial. The episode index denotes the progression of adaptation within the same trial.

Figure 13: Effect of within-episode memory capacity. We compare the default TTT-MLP memory with the lighter TTT-Linear variant on PointMaze, ML45, and ML10.

Figure 14: Effect of separating the current state from the learned task context. The default Metta provides

s_{e,t}

separately to the actor and critic, while the ablated variant must represent both state and task information through the learned context.

Figure 15: Qualitative PointMaze trajectories for Metta , ECET, and TrMRL. Columns correspond to methods, while rows show 9 task instances at mid-training and at the end of training. Colored curves denote Episodes

1

2

5

, and

10

from the same trial. ECET and TrMRL frequently revisit similar regions across episodes, whereas Metta covers a wider portion of the maze and progressively produces more goal-directed behavior.

Meta-Reinforcement Learning

Meta-reinforcement learning trains agents across tasks so they can adapt quickly to new environments or objectives.

See all 40 papers

Feb 23, 2026

Meta-Learning and Meta-Reinforcement Learning -- Tracing the Path towards DeepMind's Adaptive Agent

Humans are highly effective at utilizing prior knowledge to adapt to novel tasks, a capability that standard machine learning models struggle to replicate due to their reliance on task-specific training. Meta-learning overcomes this limitation by allowing models to acquire transferable knowledge from various tasks, enabling rapid adaptation to new challenges with minimal data. This survey provides a rigorous, task-based formalization of meta-learning and meta-reinforcement learning and uses that paradigm to chronicle the landmark algorithms that paved the way for DeepMind's Adaptive Agent, consolidating the essential concepts needed to understand the Adaptive Agent and other generalist approaches.

Björn Hoppmann, Christoph Scholz

Meta-Learning Meta-Reinforcement Learning Generalization in RL

Ask PDF

Jun 16, 2026

Knowledge Reutilization in Meta-Reinforcement Learning

Meta-reinforcement learning enables fast adaptation by extracting shared structure from related tasks, but existing end-to-end methods often couple task inference with embodiment-specific control. This coupling can obscure non-parametric task semantics, reduce sample efficiency, and limit cross-agent reuse. We propose a meta-knowledge reutilization framework that learns task-level knowledge on a dynamics-simplified agent and transfers it to heterogeneous agents. The framework uses a Bayesian non-parametric prior to organize latent task modes and a high-level policy to generate task-level magnitude guidance. To bridge reusable task knowledge with different embodiments, we introduce a semantic-magnitude interface and a lightweight temporal adaptor, which convert frozen meta-knowledge into temporally aligned subgoals for embodiment-specific low-level controllers. Experiments on multiple locomotion agents show that our framework reduces final-step tracking error by 94.75% -- 99.79% compared with recent state-of-the-art baselines and achieves comparable deployment performance with about 23.8% of their interaction data.

Yuan Meng, Bo Wang, Juan de los Rios Ruiz +4

School of Computation, Information and Technology, Technical University of Munich, Munich, Germany · State Key Laboratory for Novel Software Technology, Nanjing University, Suzhou, China · Department of Computer Science and Technology, Tsinghua University, Beijing, China

Cross-Embodiment Robot Learning Meta-Reinforcement Learning

Ask PDF

May 30, 2026

Behavior-Invariant Task Representation Learning with Transformer-based World Models for Offline Meta-Reinforcement Learning

Offline meta-reinforcement learning leverages static datasets to enable agents to generalize to unseen environments by combining offline efficiency with meta-learning adaptability, yet it faces key challenges from context and policy distribution shifts. These issues hinder agents from adapting to online environments, and are further exacerbated under sparse-reward settings. As a result, agents often become trapped in an inherent pattern dilemma, failing to achieve robust generalization. In this work, we propose a novel framework that integrates information-theoretic task representation learning with a Transformer-based stochastic world model. Our approach extracts task-defining latent variables that are invariant to behavior policy, thereby effectively mitigating the context distribution shift. To further handle policy shift and model exploitation, we apply a conservative value penalty to imagination-based rollouts, preventing the policy from exploiting model inaccuracies while maintaining robust adaptation. Extensive evaluations demonstrate that our method outperforms state-of-the-art approaches, with superior stability and generalization under out-of-distribution and sparse-reward settings.

Fuyuan Qian, Menglong Zhang, Song Wang +1

Department of Biomedical Engineering, Southern University of Science and Technology, Shenzhen, Guangdong, China · AI Thrust, The Hong Kong University of Science and Technology (Guangzhou), Guangzhou, Guangdong, China · Omni-Intelligence, Shenzhen, Guangdong, China +1

Representation Learning Offline RL Meta-Reinforcement Learning Model-Based RL

Ask PDF

Abstract

Meta-Reinforcement Learning

Explore similar work

Meta-Learning and Meta-Reinforcement Learning -- Tracing the Path towards DeepMind's Adaptive Agent

Knowledge Reutilization in Meta-Reinforcement Learning

Behavior-Invariant Task Representation Learning with Transformer-based World Models for Offline Meta-Reinforcement Learning