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

Technical Summary: MATE – Solving Contextual Markov Decision Processes with Memory of Accumulated Transition Embeddings

1. Problem Statement

The paper addresses the challenge of solving Contextual Markov Decision Processes (CMDPs), a class of MDPs parameterized by an unobserved latent context $c$. In these settings, the transition kernel $P_c$ and reward function depend on $c$, which is not directly observable by the agent. To act optimally, an agent must infer the posterior belief over contexts, $p(c|x_{1:t})$, based on its interaction history $x_{1:t} = (s_0, a_0, r_0, \dots, s_t)$.

Standard state-dependent policies fail in CMDPs because they cannot adapt to the latent context. While memory-based Reinforcement Learning (RL) offers a solution by encoding history into a latent memory state, existing architectures face significant trade-offs:

• Recurrent Neural Networks (RNNs): suffer from gradient instabilities over long horizons and preclude temporal parallelization during training updates.
• Transformers: excel at long-term memory but incur quadratic computational costs ($O(T^2)$) with respect to sequence length during both training updates and trajectory rollouts, often necessitating limited context windows.

Furthermore, standard sequence models (like causal Transformers) are order-sensitive, whereas the optimal posterior belief in CMDPs is theoretically permutation invariant with respect to the order of transitions.

2. Methodology: MATE Architecture

The authors propose MATE (Memory of Accumulated Transition Embeddings), a memory architecture designed to exploit the permutation invariance of the context posterior while avoiding the computational and optimization drawbacks of RNNs and Transformers.

Core Mechanism:
MATE represents the memory state $m_t$ as a simple sum-aggregation of transition embeddings:
$$m_t = \sum_{i=1}^{t} E_\psi(x_i)$$
where $E_\psi$ is a transition encoder (e.g., a neural network) mapping each transition $x_i$ to an embedding vector.

Key Design Features:

• Permutation Invariance: By design, the summation operation is commutative. This aligns with the theoretical property that the posterior $p(c|x_{1:t})$ is invariant to the ordering of transitions (derived via Bayes' rule and the conditional independence of transitions given $c$).
• Computational Efficiency:
  • Rollout: MATE maintains a constant inference cost per step ($O(1)$) via the update rule $m_{t+1} = m_t + E_\psi(x_{t+1})$, matching the efficiency of RNNs.
  • Training: Unlike RNNs, MATE supports temporal parallelization across the sequence dimension during the update step, similar to Transformers, but with linear computational complexity ($O(T)$) rather than quadratic.
• Hyperspherical Normalization: To prevent the norm of the accumulated memory from growing unbounded (which causes internal covariate shift), the authors introduce a normalization technique. The memory is projected onto a hypersphere: $\hat{m}_t = (m_t + \Psi) / \|m_t + \Psi\|$, where $\Psi$ is a trainable offset vector preventing information loss during normalization.

3. Key Contributions

1. Theoretical Expressiveness: The paper provides a theoretical proof (Proposition 3.1) that a memory architecture using sum-aggregation of transition embeddings is provably sufficient to represent the optimal policy in CMDPs. Specifically, it demonstrates that with a sufficiently large embedding dimension ($m = 2nT + 1$), a simple one-layer MLP encoder combined with summation can injectively map the multiset of transitions to a memory state, preserving the information required for the optimal policy.
2. Architectural Innovation: MATE replaces complex attention mechanisms and recurrent connections with simple sum aggregation. It is the first black-box memory-based RL agent to fully leverage the permutation invariance of the CMDP posterior without auxiliary inference objectives.
3. Empirical Validation: Extensive experiments demonstrate that MATE achieves performance comparable to or superior to standard RNN (LSTM) and Transformer (GPT-2) baselines across diverse benchmarks, while offering superior computational efficiency.

4. Experimental Results

The authors evaluated MATE on three benchmarks: MuJoCo, Meta-World, and T-Maze.

• MuJoCo (Adaptation to Dynamics/Rewards): MATE performed on par with diverse memory-based baselines (LSTM, Transformer variants) in tasks involving reward variations (Ant-Dir, Cheetah-Vel) and system dynamics (Walker-Param, Hopper-Param).
• Meta-World (Multi-task Adaptation): MATE consistently outperformed RNN and Transformer baselines, achieving higher returns and improved sample efficiency despite its lightweight architecture.
• T-Maze (Long-term Memory & Credit Assignment):
  • In the Passive setting, MATE and GPT-2 achieved optimal returns across all context lengths, while LSTM performance degraded as sequence length increased.
  • In the Active setting (requiring exploration to find context), MATE maintained consistent performance across varying corridor lengths (up to 600 steps), whereas GPT-2 showed lower returns and LSTM performance declined at longer lengths.
• Computational Cost: As summarized in the paper's complexity analysis, MATE achieves $O(T)$ rollout and update costs with parallelizable updates, contrasting with the $O(T^2)$ update cost of Transformers and the non-parallelizable updates of RNNs.

5. Significance and Claims

The paper claims that MATE offers a computationally efficient and theoretically grounded alternative to existing memory architectures for CMDPs. Its significance lies in:

• Bridging Theory and Practice: It validates that the inductive bias of permutation invariance, often ignored in favor of order-sensitive sequence models, is not only sufficient for optimality but also computationally advantageous.
• Efficiency: It resolves the trade-off between the gradient stability/parallelization of Transformers and the constant-step efficiency of RNNs.
• Simplicity: It demonstrates that complex attention mechanisms are not strictly necessary for effective context inference in CMDPs if the architecture is aligned with the problem's inherent symmetries.

The authors conclude that MATE effectively identifies relevant contextual information during rollout and suggests future work in integrating this architecture with deep, self-supervised RL frameworks to further enhance scalability.