Title: Predictive Coding Enhances Meta-RL To Achieve Interpretable Bayes-Optimal Belief Representation Under Partial Observability
Venue: NeurIPS 2025
Accepted: yes
Keywords: partial observability

meta-reinforcement learning

predictive coding

self-supervised learning

representation learning

neuro-ai

decision-making under uncertainty

partially observable Markov decision process

POMDP

deep reinforcement learning
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.


Peer reviews:
comments: SUMMARY:
I am not sure what this paper is about. The title contains several terms:

 * Predictive coding
 * Meta-RL
 * Bayes-optimal belief presentation
 * Partial observability

and even after reading the abstract I still am confused what are the main problems solved in this work.

STRENGTHS AND WEAKNESSES:
+ the work addresses the partial observability problem that is largely unsolved in RL literature

- writing style makes it difficult to follow:
  * What is the problem
  * Why the problem is important w/ example cases (beyond a few refs)
  * What have been previously done (state-of-the-art)
  * Own contributions 

- comparing results to a 9-years old paper is not sufficient, especially as it is difficult to see any differences

- motivation of using the selected environments is poor (why these are fundamentally good examples of partially observable and realistic cases)

QUESTIONS:
Predictive coding (PC) is an alternative method to other optimization methods, such as backpropagation (BP), used in neural networks. Why and how PC provides something fundamentally better than BP using gradient descent? (please, do not answer that read all the cited references as this should be clearly explained in the paper)

Why the meta-RL is not explained in this work (method is half a page and referring to appendices - the main parts should be in the paper)?

How can I confirm that there is any statistical significance in the results in Figure 2 and 3?

LIMITATIONS:
Not discussed (no explicit section)
TITLE: Official Review
CLARITY: 1
IMPACT: 2
ORIGINALITY: 3
DATE: 2025-06-10

comments: SUMMARY:
This manuscript investigates how to improve internal state representations for meta-reinforcement learning (meta-RL) agents in partially observable environments. The authors propose integrating a self-supervised predictive module, inspired by the predictive coding principle in neuroscience, into a meta-RL agent. This module is trained to predict future observations and rewards, which in turn shapes a latent representation intended to approximate a Bayes-optimal belief state. Using state machine simulations, the authors finds that the predictive module enables the learning of more interpretable and structurally equivalent representations compared to standard meta-RL ($RL^2$). In tasks requiring active information seeking, this improved version allows the agent to achieve optimal policies where conventional meta-RL fails.

STRENGTHS AND WEAKNESSES:
**Strengths**
1. Clear theoretical bridge and a tight conceptual link to the predictive coding principle in neuroscience.
2. Rigorous evaluation methods are adopted so that we can truly know if the learned representations are Bayes-optimal or not. This goes beyond previous methods that may have adopted predictive coding but lack interpretation.

**Weaknesses**
1. The experiments are limited to simple POMDPs and thus the scalability of this method (in terms of performance and interpretability ) is unknown when facing high-dimensional complex tasks.
2. The code is not available at the time of reviewing.

QUESTIONS:
no

LIMITATIONS:
no
TITLE: Official Review
CLARITY: 4
IMPACT: 2
ORIGINALITY: 3
DATE: 2025-07-03

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe authors tackle the twin challenges of interpretable belief representations and efficient exploration by integrating a predictive-coding module into a memory-based meta-reinforcement learning (meta-RL) agent. Concretely, a variational RNN encoder ingests (o_t, r_t, a_{t-1}) at each step and compresses them into a latent belief btb_t; two lightweight decoders are trained, under an ELBO with a KL bottleneck, to predict the next observation and reward. A separate actor–critic head then conditions its policy solely on btb_t, with no policy-gradient leakage, back into the encoder.\n\nThe authors hypothesise that this compact latent representation more closely matches the Bayes-optimal belief. They test the claim using a _state-machine simulation_ metric on five POMDP families whose Bayes-optimal solutions are analytically tractable, providing a ground truth: static and dynamic Bernoulli bandits, stationary and dynamic Tiger, and an 11-arm “oracle” bandit that demands strategic exploration.\n\nResults show that the encoder’s bottleneck reduces State and Return dissimilarities D_s, D_o relative to RL² and does so without enlarging network capacity. Both agents reach optimal return on the easier tasks, but on the hardest Dynamic-Tiger and Oracle settings only the predictive-module agent attains Bayes-optimal reward and uncovers the correct exploration policy. Ablations confirm that (i) the KL information bottleneck and (ii) gradient isolation are each necessary, removing either degrades both belief alignment and performance, demonstrating that predictive coding alone, properly regularised, drives the improvement.\n\nOverall, the study provides strong but small-scale evidence that self-supervised prediction is a powerful inductive bias for belief learning in meta-RL.\n\nSTRENGTHS AND WEAKNESSES:\nThis is a high-quality empirical study with good experimental controls, a well-motivated design, and a sensible evaluation methodology. The use of state-machine simulation to measure representational adequacy, both from agent-to-Bayes and Bayes-to-agent, is a strong choice that moves beyond the typical reliance on downstream Return as a proxy. The ablation study is appropriately scoped and supports the central claims. That said, the experimental domain is limited to low-dimensional, fully known POMDPs, which constrains generalisability.\n\nThe paper is clearly written, with a logical flow from motivation to architecture, experiments, and discussion. Key assumptions are made explicit, the diagrams are effective, and notation is standard. A minor weakness is that the state-machine simulation metric, while central to the evaluation, could benefit from a clearer standalone description for readers unfamiliar with Mikulik et al.\n\nThe significance is primarily conceptual rather than empirical: the paper reframes ideas from predictive coding as an inductive bias for Bayes-optimal belief representation in meta-RL. While the domains are small-scale, the implications, particularly for training interpretable and robust meta-RL agents, are notable. If the method generalises to continuous control or high-dimensional inputs, it would have substantial impact. It's very important to note that not only are the environments used small scale, but they are from within a very specific domain of analytically tractable Bayes-optimal solvable settings. This is a major weakness of the paper.\n\nThis work is original in its use of predictive coding as a direct mechanism for belief representation in meta-RL. While prior work has explored predictive objectives for generalisation or sample efficiency, this paper explicitly shows that self-supervised future prediction can induce latent states that closely approximate Bayes-optimal beliefs, validated through a principled, bi-directional state-machine simulation test.\n\nA second key novelty is the architectural decoupling of inference and control: the belief encoder is trained solely via predictive losses, with no policy gradient leakage, yet still supports Bayes-optimal behaviour. This differs from both conventional meta-RL (e.g., RL^2) and world-model agents that require planning in latent space. There is a limitation here however in terms of the limited comparison to only a single type of meta-RL approach. Comparison to the likes of VariBAD would be an important improvement.\n\nQUESTIONS:\nHow does this approach compare to other meta-RL methods that incorporate predictive objectives (VariBAD, world model approaches)? Why did you choose RL² as the sole baseline?\n\nYou claim that \"no policy-gradient leakage\" is crucial, but the ablation shows joint training doesn't hurt all that much. Can you provide more analysis of when/why gradient isolation matters?\n\nYour tasks have 2-11 discrete actions and analytically tractable belief states. How do you expect the approach to scale to continuous control or high-dimensional observation spaces where Bayes-optimal solutions are unknown? This is probably the biggest weakness of the study.\n\nThe oracle bandit success is definitely interesting, but it's also a very specific exploration structure. How does the approach work on more diverse exploration challenges?\n\nWhat's the computational overhead of the predictive modules? In particular, how does training time compare to RL²?\n\nHow sensitive is the method to the choice of the belief state bottleneck dimension?\n\nYou mention OOD generalization as future work, but can you show results on task distributions that differ from training (e.g., different bandit arm counts, different Tiger observation accuracies)? This would seem like a relatively simple test to carry out.\n\nLIMITATIONS:\nYes", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 4, "CLARITY": 3, "IMPACT": 2, "ORIGINALITY": 3, "DATE": "2025-07-03"}