Memory-enhanced adaptation improves agents in unseen reinforcement tasks
METTA: Memory-Enhanced Test-Time Adaptation for Meta-Reinforcement Learning
Machine LearningArtificial Intelligence
Summary
Learning to adapt quickly to new tasks is hard because agents need to remember many past experiences. The authors propose METTA, a method that helps agents adapt by using a special memory system that focuses on recent events while summarizing longer past episodes efficiently. This approach uses less computing power compared to previous methods and performs better in tests where tasks are unfamiliar or rewards are rare. METTA helps agents explore smarter and succeed more often in these challenging situations.
What this means in practice
- •For robotics engineers: Enable robots to quickly adapt to new manipulation tasks with limited trial experience using efficient memory-based learning.
- •For game ai developers: Create game agents that adjust strategies faster and more efficiently in complex environments with sparse feedback.
Authors
Ali Najar, Hosein Hasani, Amir Kooshan Fattah, Yasamin Rajabi, Mahdieh Soleymani Baghshah
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.