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.
@article{arxiv.2605.17431,
title = {MATE: Solving Contextual Markov Decision Processes with Memory of Accumulated Transition Embeddings},
author = {Himchan Hwang and Hyeokju Jeong and Gene Chung and Seungyeon Kim and Sangwoong Yoon and Frank Chongwoo Park},
journal= {arXiv preprint arXiv:2605.17431},
year = {2026}
}