Predictive Coding Enhances Meta-RL To Achieve Interpretable Bayes-Optimal Belief Representation Under Partial Observability

AI summary

Overview

• Research area: Meta-reinforcement learning (meta-RL) under partial observability, representation learning, and the neuroscience-inspired principle of predictive coding.
• Technical level: Intermediate to Advanced. The paper assumes familiarity with POMDPs, variational inference (VAE/ELBO), recurrent meta-RL (RL²), and Bayes-optimal belief states, though the core argument is stated in accessible terms.
• Scope: A systematic, simulation-based comparison of conventional black-box meta-RL against meta-RL augmented with self-supervised predictive modules, measuring how closely each learns Bayes-optimal belief representations across six classes of partially observable tasks.

What This Paper Is About

In partially observable environments, an agent must summarize its entire history into something compact that preserves what matters for decisions. Bayesian theory says the ideal summary is the belief state—a probability distribution over the hidden state—but deep meta-RL agents such as RL² typically reach near-optimal behavior while their internal representations remain opaque and do not match these minimally sufficient belief states. This paper asks whether adding self-supervised predictive coding—forcing the agent's latent bottleneck to predict upcoming observations and rewards—makes meta-RL learn representations that are actually, verifiably equivalent to Bayes-optimal beliefs.

Key Contributions

1. A meta-RL framework with self-supervised predictive modules that separates representation learning (a VAE trained by predictive coding) from policy learning (a feedforward policy network trained by RL), with no reliance on privileged information or a separate random exploration policy for generating training samples.
2. A systematic ground-truth evaluation across diverse POMDPs, extending beyond the stationary bandits studied previously to cover explore-exploit tradeoffs, sequential decision-making, dynamic belief tracking, information gathering, and continuous control—the authors state this is the first systematic evaluation of meta-RL representations using ground-truth comparison across diverse POMDPs.
3. Rigorous state machine simulation analysis in both mapping directions, showing that meta-RL with predictive modules consistently produces representations closer to Bayes-optimal beliefs than conventional meta-RL, even when both achieve optimal policies.
4. Evidence that better representations translate into better behavior, including the only successful Bayes-optimal policy on the oracle bandit, near-optimal zero-shot generalization, faster transfer to out-of-distribution tasks, and ablations attributing the gain to predictive learning rather than policy gradients.

Main Findings

• Better belief alignment without behavioral differences: In the classic two-armed Bernoulli bandit and the three dynamic bandit variants, both models approached Bayes-optimal return, yet state machine simulation showed the bottleneck layer of meta-RL with predictive modules achieved the lowest state dissimilarity (D_s) and output dissimilarity (D_o) in both mapping directions.
• Recurrent layers are not belief states: Even after training, the recurrent layers of both models showed high state dissimilarity in the Bayes → meta-RL direction, consistent with earlier findings that RL² recurrent representations deviate from minimally sufficient Bayes-optimal states.
• A low-dimensional bottleneck alone is not enough: The bottleneck layer of the RL² baseline retained significantly higher state dissimilarity than that of meta-RL with predictive modules, indicating that predictive coding, not merely compression, drives the improvement.
• Only predictive modules solve the hardest Tiger setting: In the dynamic Tiger task with low observation accuracy, where the tiger's location changes with probability 0.9 at each timestep, only meta-RL with predictive modules approached the Bayes-optimal return; both models approached it in the easier tasks.
• Information seeking requires the right representation: On the 11-arm oracle bandit, only meta-RL with predictive modules learned the Bayes-optimal policy, paying an immediate cost to sample the oracle arm and using the information consistently. RL² converged to a suboptimal policy—sampling the oracle arm upfront but failing to exploit what it learned—alongside uninterpretable representations and high dissimilarity measures.
• Continuous control: In the latent goal cart task, both models approached Bayes-optimal return, but the bottleneck layers of meta-RL with predictive modules showed significantly lower state and output dissimilarities in both mapping directions.
• Improved zero-shot generalization: Trained on dynamic Tiger with observation accuracy 0.7 and tested on environments with observation accuracy 0.8, meta-RL with predictive modules reached near-optimal test return (−15.94 ± 3.82), while RL² was significantly lower (−25.56 ± 1.80).
• Faster out-of-distribution transfer: When models pre-trained on oracle bandits containing only arms 1–5 were transferred to a distribution containing arms 6–10, meta-RL with predictive modules learned significantly faster than RL².
• Ablations point to predictive learning: Removing the VAE's KL regularization increased bottleneck dissimilarity to Bayes-optimal states on most tasks (with a minor decrease in dynamic bandits), while letting policy gradients jointly train the RNN encoder did not further decrease bottleneck dissimilarities. The authors conclude the enhanced representation learning is largely attributable to predictive learning.

Methodology in Plain English

The researchers built a meta-RL agent from two cooperating parts. The first part is a recurrent encoder that reads the current observation, reward, and previous action, and compresses the history into a low-dimensional probabilistic latent state, analogous to a belief. This encoder is trained only by prediction: decoders try to guess the next observation and the next reward, and a KL term regularizes the latent distribution against the previous timestep's posterior, mimicking Bayesian filtering. The second part is a simple feedforward policy network that takes the belief state as input and is trained by standard model-free policy gradients; gradients from the RL loss are not used to train the encoder.

To judge whether the learned latent states really are belief states, the authors use state machine simulation rather than simple decoding. Two mapping networks are trained—one from meta-learned states to Bayes-optimal beliefs and one in the reverse direction—and the quality of the mapping is measured by state dissimilarity (mean squared error against the target states) and output dissimilarity (the difference in return generated by the original versus the mapped states). Low dissimilarity in both directions means the two representations are computationally equivalent. The baseline RL² model was made architecturally identical apart from the decoder networks and identical layer sizes, so differences can be attributed to the predictive objective.

Ground-truth Bayes-optimal solutions were computed for each task: the Gittins index for the classic two-armed Bernoulli bandit, and value iteration over belief space for the dynamic bandits and Tiger tasks. Reported error bars denote the standard error of the mean across at least 5 training runs, and asterisks denote statistical significance (p < 0.05, t-test) among trained models.

Why This Matters

Impact on research. The paper supplies an interpretability-based explanation for why auxiliary predictive objectives help in deep RL, rather than only documenting that they do. By linking meta-learned latents to Bayes-optimal beliefs, it connects a mainstream deep RL technique to predictive coding theory in neuroscience, and it reinforces the methodological point that decoding-based similarity measures can be misleading—untrained recurrent layers can map arbitrarily close to belief states if the mapping function is powerful enough, which the authors observe directly in the two-armed bandit.

Real-world applications (the first two are named in the paper's broader impact section; the latter two follow from the tasks studied):

• Medical decision support, where information is gathered sequentially and treatment decisions under uncertainty must be explainable.
• Disaster response, where agents act under limited and noisy information.
• Continuous control and robotics with hidden goals or unobserved state, as in the latent goal cart task.
• Adaptive systems that must trade immediate exploration cost for long-term gain, as in the oracle bandit.

Industry relevance. Agents whose latent representations are provably aligned with a Bayesian belief state are easier to audit, more robust when task conditions shift, and cheaper to adapt to new but related deployments. The demonstrated zero-shot generalization and faster transfer learning directly target two of the most expensive problems in applied RL: retraining on distribution shift and sample efficiency in new environments. The codebase for model training and state machine simulation is publicly available at https://github.com/walkerlab/metaRL-predictive-representation.

Future Directions

• Scalable equivalence analysis. The current approach requires POMDPs with tractable Bayes-optimal solutions; the authors call for scalable analysis methods for tasks with larger state spaces.
• Beyond next-step prediction. Only next-step future prediction was examined as the predictive objective. The authors propose investigating multi-step objectives, different time scales, and temporal abstractions in the predictive modules, motivated by the empirical success of multi-step predictive objectives in deep RL.
• Meta-learning representations for out-of-distribution generalization. Representation quality was shown to affect generalization, but the authors state that a more comprehensive evaluation is needed to determine how to meta-learn representations that support OOD generalization.
• Tighter ties to neuroscience. The results align with predictive coding theories, and the authors suggest both directions of exchange—neuroscience principles opening the black box of modern AI, and the computational principles found in artificial agents informing understanding of neural computation.

Target Audience

Researchers and graduate students in reinforcement learning, especially those working on meta-RL, POMDPs, and representation learning; computational neuroscientists interested in predictive coding and Bayesian accounts of brain function; and practitioners who need interpretable, adaptible agents operating under uncertainty and limited information. Readers seeking a performance-only benchmark comparison will find the paper's focus is instead representational equivalence and interpretability, with explicit scope limitations noted.

Authors’ abstract

Learning a compact representation of history is critical for planning and generalization in partially observable environments. While meta-reinforcement learning (RL) agents can attain near Bayes-optimal policies, they often fail to learn the compact, interpretable Bayes-optimal belief states. This representational inefficiency potentially limits the agent's adaptability and generalization capacity. Inspired by predictive coding in neuroscience--which suggests that the brain predicts sensory inputs as a neural implementation of Bayesian inference--and by auxiliary predictive objectives in deep RL, we investigate whether integrating self-supervised predictive coding modules into meta-RL can facilitate learning of Bayes-optimal representations. Through state machine simulation, we show that meta-RL with predictive modules consistently generates more interpretable representations that better approximate Bayes-optimal belief states compared to conventional meta-RL across a wide variety of tasks, even when both achieve optimal policies. In challenging tasks requiring active information seeking, only meta-RL with predictive modules successfully learns optimal representations and policies, whereas conventional meta-RL struggles with inadequate representation learning. Finally, we demonstrate that better representation learning leads to improved generalization. Our results strongly suggest the role of predictive learning as a guiding principle for effective representation learning in agents navigating partial observability.