# MATE: Solving Contextual Markov Decision Processes with Memory of Accumulated Transition Embeddings

arXiv: https://arxiv.org/abs/2605.17431 · HTML: https://mlsanity.com/paper/2605.17431/

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Builds on 23 of its 39 references that are in the map.
- First posted on arXiv 17 May 2026.

## Builds on (most cited first, 23 of 23)

- [Batch Normalization: Accelerating Deep Network Training by Reducing Internal Covariate Shift](https://mlsanity.com/paper/1502.03167/) (2015)
- [Model-Agnostic Meta-Learning for Fast Adaptation of Deep Networks](https://mlsanity.com/paper/1703.03400/) (2017)
- [Soft Actor-Critic: Off-Policy Maximum Entropy Deep Reinforcement Learning with a Stochastic Actor](https://mlsanity.com/paper/1801.01290/) (2018)
- [Mamba: Linear-Time Sequence Modeling with Selective State Spaces](https://mlsanity.com/paper/2312.00752/) (2023)
- [Deep Reinforcement Learning with Double Q-learning](https://mlsanity.com/paper/1509.06461/) (2015)
- [Root Mean Square Layer Normalization](https://mlsanity.com/paper/1910.07467/) (2019)
- [Meta-World: A Benchmark and Evaluation for Multi-Task and Meta Reinforcement Learning](https://mlsanity.com/paper/1910.10897/) (2019)
- [RL2: Fast Reinforcement Learning via Slow Reinforcement Learning](https://mlsanity.com/paper/1611.02779/) (2016)
- [Efficient Off-Policy Meta-Reinforcement Learning via Probabilistic Context Variables](https://mlsanity.com/paper/1903.08254/) (2019)
- [Contextual Markov Decision Processes](https://mlsanity.com/paper/1502.02259/) (2015)
- [A Survey of Zero-shot Generalisation in Deep Reinforcement Learning](https://mlsanity.com/paper/2111.09794/) (2021)
- [Recurrent Model-Free RL Can Be a Strong Baseline for Many POMDPs](https://mlsanity.com/paper/2110.05038/) (2021)
- [Structured State Space Models for In-Context Reinforcement Learning](https://mlsanity.com/paper/2303.03982/) (2023)
- [A Tutorial on Meta-Reinforcement Learning](https://mlsanity.com/paper/2301.08028/) (2023)
- [When Do Transformers Shine in RL? Decoupling Memory from Credit Assignment](https://mlsanity.com/paper/2307.03864/) (2023)
- [1000 Layer Networks for Self-Supervised RL: Scaling Depth Can Enable New Goal-Reaching Capabilities](https://mlsanity.com/paper/2503.14858/) (2025)
- [Neural Injective Functions for Multisets, Measures and Graphs via a Finite Witness Theorem](https://mlsanity.com/paper/2306.06529/) (2023)
- [Context is Everything: Implicit Identification for Dynamics Adaptation](https://mlsanity.com/paper/2203.05549/) (2022)
- [Decision Mamba: Reinforcement Learning via Hybrid Selective Sequence Modeling](https://mlsanity.com/paper/2406.00079/) (2024)
- [Towards an Information Theoretic Framework of Context-Based Offline Meta-Reinforcement Learning](https://mlsanity.com/paper/2402.02429/) (2024)
- [Recurrent Hypernetworks are Surprisingly Strong in Meta-RL](https://mlsanity.com/paper/2309.14970/) (2023)
- [AMAGO-2: Breaking the Multi-Task Barrier in Meta-Reinforcement Learning with Transformers](https://mlsanity.com/paper/2411.11188/) (2024)
- [Scaling Algorithm Distillation for Continuous Control with Mamba](https://mlsanity.com/paper/2506.13892/) (2025)

## Abstract (arXiv)

We propose MATE, a simple yet effective memory architecture for solving Contextual Markov Decision Processes (CMDPs), a family of MDPs parameterized by an unobserved context. In CMDPs, an optimal agent can adapt online by maintaining the posterior belief over contexts. MATE replaces this intractable posterior with a sum-aggregated memory, leveraging the posterior's permutation invariance to retain provably sufficient expressiveness. Compared to prior memory architectures, MATE avoids the growing per-step rollout cost of Transformers and the gradient issues commonly associated with Recurrent Neural Networks (RNNs). Extensive evaluations across diverse benchmarks demonstrate that MATE provides clear computational advantages while achieving performance comparable to standard sequence-model baselines.