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

We thank the reviewer for your review and the opportunity to address your concerns. Below we will clarify the main problem that we solve in our work to better convey these points. In particular we believe there is confusion about the framework of predictive coding, and we will clarify this aspect below and in our revision.

What problem are we solving and why is it important?

In the real-world, agents (AI systems, humans) often have to make decisions without complete information. For example:

• A doctor diagnosing a patient can only observe symptoms, not the underlying disease
• A robot navigating a building can only see with its camera, not through walls

As you pointed out, this is known in RL as partial observability. In POMDPs, the optimal way to handle this is to maintain a belief (probability distribution) over the true state, updating it as new information arrives. This is the Bayes-optimal belief representation. Learning Bayes-optimal belief representations is important not only for performance but also for interpretability and generalization.

Why previous methods fail

Memory-based meta-RL (RL²) is shown as a powerful method to learn near-optimal policies. However, RL² fails to learn compact, interpretable representations that match Bayes-optimal belief states, which was the main negative result from Mikulik et al 2020. Imagine learning to diagnose diseases by memorizing every symptom combination rather than understanding the underlying disease mechanism. This causes several fundamental problems:

• Lack of interpretability (we can’t explain the symptoms)
• Ineffective exploration (without disease mechanisms, we don’t know where to look for potential new treatment)
• Inadequate generalization (we don’t know whether the symptom combination will transfer to new patients).

This is exactly the problem we attempt to address: how can we design meta-RL agents that can learn Bayes-optimal belief representations to support interpretability, exploration, and generalization?

Our contribution

Motivated by the predictive coding principle in neuroscience (pioneered by Rao & Ballard, 1999), we show that by adding predictive objectives (predicting future observations) to meta-RL, the learned internal representations become much closer to the theoretically optimal Bayesian beliefs. To use an analogy, this is like the difference between:

• Memorizing that "fever + cough often means flu" (what RL² does, a verbose representation)
• Understanding the underlying disease mechanism that "flu causes fever 80% of the time and cough 70% of the time" (what our model achieves, an interpretable, task-relevant optimal representation)

To this end, our work went beyond performance comparison that lacks interpretability, to rigorously evaluate whether the learned internal representations are truly Bayes-optimal beliefs. Our work presented a thorough model interpretability analysis: incorporating predictive objectives leads to better representations that match Bayes-optimal beliefs, crucial for generalization and more interpretable AI systems.

We hope this clarifies our motivations and contributions. We're happy to provide additional clarification. Below we address specific questions.

Q1: Clarifying the confusion about predictive coding

This appears to be a misunderstanding. In our work, predictive coding is not an alternative optimization method to backpropagation. Rather, predictive coding here refers to the future predictive objective that shapes representation learning. This predictive objective is:

• Inspired by neuroscience theories of how brain process information
• Implemented using standard backpropagation
• Similar to successful representation learning methods in ML (e.g. future predictive loss, contrastive predictive coding)

We’ll clarify this distinction in the revision to prevent misunderstanding.

Q2: Why these environments?

We would like to emphasize that our goal is not to compare with SOTA performance, but rather to learn interpretable representations. Consequently, we carefully select POMDP tasks where we can analytically compute Bayes-optimal solutions to enable rigorous comparison of learned representations. We find this a necessary methodological choice because:

• Optimal solutions in complex POMDPs are known to be intractable. Without optimal solutions, we cannot directly verify whether the learned representations are optimal
• In tractable POMDPs, we can quantitatively measure how close our learned representations are to Bayes-optimal solutions, not merely comparing performance
• Our selected tasks capture core, fundamental challenges in POMDPs, including exploration-exploitation tradeoff, information gathering, and dynamic uncertainty

To be specific, our tasks encompass core partial observability challenges with real-world implications.

• Bandits: exploration-exploitation tradeoff → Clinical trial design, A/B testing
• Tiger: classic POMDP benchmark → Scenarios requiring costly information gathering before high-stakes decisions
• Dynamic bandits/Tiger: non-stationary tasks requiring dynamic belief tracking → Markets, weather, or any dynamic domains
• Oracle bandit: active information seeking → Active, exploration-based information gathering

We acknowledge that these tasks aren't meant for realistic applications, but rather diagnostic environments where we can evaluate representation quality against ground truth, enabling our core contribution of model interpretability.

Q3: Comparison baselines

Our primary goal is model interpretability and understanding why predictive objectives work. The state machine simulation method from Mikulik et al 2020 that we employ in our work provides the most rigorous analysis to directly compare learned representations with Bayes-optimal beliefs—it's an analysis tool for model interpretability, beyond performance baseline.

To contextualize our model, many modern deep RL methods have converged on incorporating predictive objectives:

• VariBAD (Zintgraf et al 2021) uses variational inference with future prediction
• SOLAR (Zhang et al 2019) and Dreamer (Hafner et al 2025) learn world models through observation prediction
• DynaMITE-RL (Liang et al 2024) employ dynamic predictive model

These methods show improved performance on various tasks using predictive objectives. This convergence suggests predictive learning is a fundamental principle. However, none of these works systematically analyze whether the learned representations actually approximate Bayes-optimal beliefs. Our contribution is demonstrating why predictive objectives work: they guide neural networks toward Bayes-optimal representations. We chose RL² as our baseline because:

• It's the canonical meta-RL method, representing a foundational and widely-used baseline.
• Mikulik et al 2020 specifically examined RL², allowing direct comparison.
• Recent meta-RL methods build upon RL².
• This allows us to isolate the effect of predictive learning on representation quality.

By showing that predictive modules lead to interpretable, Bayes-optimal representations while RL² doesn't, we provide theoretical insights for why the field has empirically converged on predictive objectives. This understanding is crucial for designing interpretable meta-RL systems.

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

We acknowledge that interpretation of state machine simulation results were not as straightforward as comparing performance. We have introduced the analysis in Section 4 to address model comparison and their statistical significance. In brief, for two representations A and B to be considered equivalent, state machine simulation analysis evaluate four metrics: state dissimilarity for A mapped onto B ($D_s^{A → B}$), and vice versa ($D_s^{B → A}$), and output dissimilarity for A mapped onto B ($D_o^{A → B}$), and vice versa ($D_o^{A → B}$). For interpretation, we need all 4 metrics to be small for A and B to be said to be equivalent.

Therefore, to interpret our results, we'd like to draw your attention to Fig. 2C, and focus on the results for the bottleneck layers in RL2 and our model. We reported S.E. across 5+ seeds and asterisks marked statistical significance (p<0.05, t-test) in all state machine simulation plots. Our results show, for the first three metrics ($D_o^{metaRL → Bayes}$, $D_s^{metaRL → Bayes}$, $D_o^{Bayes → metaRL}$), RL² and our model attain comparably low value. However, for the metric $D_s^{metaRL → Bayes}$, our model has significantly lower value than RL², indicating that overall our model learns a representation that is significantly closer to the Bayes-optimal beliefs (again noting that all 4 metrics to be small to be equivalent).

We show this consistent pattern across all tasks we considered, which provides strong evidence for better representation learning induced by predictive coding.

Q5: Why meta-RL is not explained

We have introduced the background and methodology of meta-RL in the “Memory-based meta-RL” section under Section 2. Our meta-RL model, including model architecture, training objectives, and regularization, are provided in Section 3 and Figure 1B & C. As our model components are standard, e.g. RNN, VAE, we refer to the appendix for implementation details.

Q6: Limitations not discussed

We have discussed limitations of our current work in the second paragraph of Section 6. To make this stand out more clearly, we will add a header highlighting this paragraph where we discussed limitations.

---

We hope this clarifies our contributions and addresses confusion. The core insight—that predictive objectives guide meta-RL toward Bayes-optimal representations—has important implications for building interpretable, generalizable agents. We will revise our writing to clearly communicate this message.

We thank Reviewer dd6m for your positive assessment and for recognizing our theoretical contributions and rigorous evaluation methodology. We appreciate your positive feedback on our “tight conceptual connection to predictive coding principles in neuroscience” and our systematic approach to verifying Bayes-optimality of learned representations (“This goes beyond previous methods that may have adopted predictive coding but lack interpretation”). We're delighted you found our work provides a “clear theoretical bridge” with “excellent clarity.” Below we provided point-by-point responses to your comments and questions.

Q1: 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.

We acknowledge this important question and appreciate the reviewer raising this constraint. Our focus on simple POMDPs was a necessary methodological choice that enables our core contribution: validating that predictive coding induces representations which structurally match Bayes-optimal beliefs—a claim that can only be rigorously evaluated when ground-truth Bayes-optimal solutions exist.

1. Rigorous evaluation requires ground truth: Our primary contribution is demonstrating that predictive modules enable learning Bayes-optimal representations. This requires ground-truth Bayes-optimal POMDP solutions for comparison, which are available only in tractable domains. In high-dimensional tasks where ground-truth solutions are not available, we cannot rigorously verify whether learned representations truly approximate Bayes-optimal belief states, and thus are reduced to black-box comparison using performance. To this end, this choice of tasks is a necessary foundation for understanding why predictive objectives work in meta-RL.

2. Systematic coverage of POMDP challenges: Despite using tractable tasks, we would like to highlight that we have significantly expanded evaluation beyond prior work investigating representation equivalence for meta-RL. Compared to Mikulik et al. (2020) where only stationary Bernoulli bandits are evaluated, we significantly expanded and selected exemplar tasks aiming to capture more diverse POMDP challenges, including exploration-exploitation tradeoff, information gathering, and dynamic uncertainty:

• Bandits: Uncertainty and exploration-exploitation tradeoff
• Tiger: Classic POMDP benchmark of full sequential decision-making problem
• Dynamic bandits/Tiger: Non-stationary hidden states where dynamic belief tracking is required
• Oracle bandit: Strategic information gathering

To our knowledge, this is the first systematic evaluation of meta-RL representations across diverse POMDP structures beyond stationary bandits using ground-truth comparison.

3. Distinguish model scalability from evaluation scalability: While full state machine simulation analysis becomes intractable for high-dimensional tasks, there is nothing in principle preventing our model architecture from scaling to high dimensional tasks. First, the modular architecture (separate predictive and policy networks) naturally scales to complex tasks. Second, the predictive objective remains well-defined for more complex tasks. We believe that our work establishes the principles—validating that predictive learning approximates Bayes-optimal computation provides the theoretical foundation for understanding how these methods scale to complex domains where direct verification isn't possible.

4. Empirical scalability from related work: To this end, we would like to point to empirical evidence from related work where the predictive objectives have demonstrated impressive scalability in practice. For instance, VariBAD (Zintgraf et al., 2021), utilizing similar architecture of modular predictive and policy networks, succeeds on higher dimensional and continuous control tasks in MuJoCo. Dreamer (Hafner et al., 2019-2025), a world model approach where world models are learned with predictive objectives and actor-critic networks are trained on imaginary model rollouts, achieve SOTA on challenging Atari and control tasks with image observations. Our work provides the interpretability foundation for understanding this family of methods:

• Existing work (VariBAD, Dreamer): Demonstrates that predictive objectives improve performance
• Our work: Explains why they work by revealing predictive representations approximate Bayes-optimal beliefs

We believe that this understanding is valuable precisely because it applies broadly. By establishing that predictive coding induces Bayes-optimal structure in tractable and the clearest settings, we provide a principled foundation for understanding the success of these predictive objective methods in deep RL. This understanding is crucial for designing robust models for complex tasks where we cannot verify optimality directly.

We will add a paragraph in the Discussion section explicitly addressing the above evaluation methodological choice and scalability considerations, and outlining these pathways for understanding deep RL models with predictive objectives on complex tasks.

Q2: The code is not available at the time of reviewing

We apologize for the code not being available previously during review. We have now followed the procedure to submit our anonymized code via Official Comment to the AC for distribution.

Upon acceptance, we will release our complete codebase, including:

• Full implementation of meta-RL with predictive modules
• All baseline implementations (RL²)
• Complete experimental setup
• State machine simulation analysis tools
• Scripts to reproduce results

The code will be released under an open-source license to facilitate reproducibility and enable the community to build upon our work. We will update the manuscript to include the link explicitly.

---

We hope this addresses your questions and clarifies our contribution to the intersection of neuroscience, meta-RL, and interpretability, and its potential for broader impact. Thank you again for your positive and constructive feedback.

We thank the reviewer for your thorough, insightful, and constructive review. We are encouraged that you found the paper to be a "clearly-written" "high-quality empirical study with good experimental controls, a well-motivated design, and a sensible evaluation methodology." We appreciate your positive assessment that our evaluation "is a strong choice that moves beyond the typical reliance on Return" and that our work provides "a principled […] test" with "conceptual significance" for "training interpretable and robust meta-RL agents." Below we respond point-by-point and report additional analyses.

Main Concerns

Q1: Limited experimental domain

Our focus on tractable tasks is a necessary methodological choice essential for validating that predictive coding induces representations approximating Bayes-optimal beliefs—which can be rigorously assessed only when ground-truth solutions exist.

Tractable tasks are necessary: In complex POMDPs (e.g., Atari, robotics), planning on Bayes-optimal beliefs is intractable. Without ground-truth solutions, we would be unable to apply state machine simulation and reduced to comparing only returns. Reduction to performance-only evaluation has critical limitations for understanding the learned representations, as performance doesn't reveal representational efficiency, structure, and interpretability. To move beyond and attain interpretability insights for representation, tractable tasks are thus necessary.

Expanded task selection to cover diverse POMDP challenges: Despite using tractable tasks, we'd like to highlight that we have significantly expanded evaluation beyond prior work. Compared to Mikulik et al 2020 where only stationary bandits are evaluated, we selected exemplar tasks to capture diverse POMDP challenges: explore-exploit tradeoff (bandits), sequential decision-making (Tiger), dynamic belief tracking (dynamic bandits/ Tiger), and information gathering (oracle bandit). To our knowledge, this is the first systematic evaluation of meta-RL representations across diverse POMDP structures beyond stationary bandits using ground-truth comparison.

Q2. ... compare to … VariBAD, world model approaches? Why … RL² as the sole baseline?

Our comparison was a deliberate choice to isolate and understand the effect of predictive objectives on representation learning. We aim to 1) Isolate the effect of predictive learning 2) Establish principles without being limited to specific implementation, and 3) Avoid interpretation confounds from specific design.

To this end, we consider our model, VariBAD, and world model approaches as belonging to the same family—meta-RL with predictive objectives. While comparing them may be interesting, we believe directed comparison to RL² would offer much clearer and interpretable results.

• Why not VariBAD: VariBAD is designed for stationary tasks and inapplicable to dynamic POMDPs. On stationary tasks where it applies, its architecture is similar to ours—both decouple belief inference from control. Comparing would essentially be comparing identical models.
• Why not world model approaches: Methods like Dreamer learn a world model and use rollouts for planning. They learn through planning, not belief representation. This would introduce extra complexity orthogonal to our focus on analyzing representations, as we'd be unable to determine whether differences stem from predictive learning, planning, or model details.

Our complementary contribution on interpretability: These models integrate predictive modules in various ways and show improved performance. Our work thus complements their success and provides the interpretability foundation for understanding why they work. This follows the fruitful tradition in ML of understanding principles in tractable settings that explains performance on complex tasks (e.g. theory on linear networks/ convex optimization→ deep learning).

Our work lays the groundwork for future analysis. We’ll expand our discussion section to clarify these relations and our complementary contribution.

Q3. How … scale to continuous or high-dim spaces?

1. Additional evidence from continuous control: To address this concern, we designed LatentGoalCart, a continuous control task similar to HalfCheetahDir in MuJoCo and allowing tractable belief inference, where the agent receives continuous observation (position) and reward (how far away from the hidden target), and controls continuous action (velocity) to reach an unknown target that needs to be inferred. Even on this continuous task, our model yields representations that are significantly closer to the Bayes-optimal beliefs compared to RL² (as we cannot upload the figure, we report the result below and will add them in the revision).

model | $D_o^{M->B}$ | $D_s^{M->B}$ | $D_o^{B->M}$ | $D_s^{B->M}$ (bold: p<0.05)

Ours | 0.699±0.270 | 0.017±0.004 | 0.528±0.098 | 0.175±0.049

RL² | 1.523±0.518 | 0.066±0.109 | 1.743±0.520 | 0.371±0.230

This provides evidence that our findings extend beyond discrete, low-dim tasks.

2. Model vs evaluation scalability: Your question raises two types of scalability. While the evaluation method's scalability may be limited for complex tasks, our model architecture scales naturally to high-dim tasks: 1) predictive objectives remain well-defined and computable regardless of dimensionality; 2) the modular architecture scales up: VAE predictive modules can handle high-dim inputs (images, continuous states) and Gaussian policies can handle continuous actions.

3. Empirical evidence from related work: We'd like to point to empirical evidence from related work where predictive learning show impressive scalability: VariBAD succeeds on high-dim and continuous tasks in MuJoCo; PlaNet and Dreamer (Hafner 2019-25) achieve SOTA on challenging Atari and control tasks with image observations.

Specific questions

Q4: You claim that "no policy-gradient leakage" is crucial, but the ablation … when/why gradient isolation matters?

We'd like to clarify we did not intend to imply that gradient isolation is crucial; rather, it was a design choice to cleanly isolate and understand the representation learning effect of predictive objectives. By blocking policy gradients, we ensure representation learning purely comes from predictive coding without confound. We show that policy gradients are not necessary for learning Bayes-optimal representations, suggesting that representation learning can be attributed to predictive coding. We apologize for confusion in our presentation and will improve the clarity of this point in the paper.

Connection to existing architectures: Successful predictive methods, e.g. VariBAD and Dreamer, also decouple inference and control. Our findings may explain why this is desirable: predictive objectives alone suffice for good representations, and mixing policy gradients don’t add much.

Q5: Oracle bandit success … diverse exploration challenges?

Oracle bandit captures the core structure of many exploration challenges: locating information sources, paying costs to access information, and using the information effectively. This pattern is common across diverse exploration benchmarks including TreasureMountain (Zintgraf 2021), Sphinx (Wang 2023), Map, and Overcooked (Xie 2024).

Our agent shows effective exploration—strategically sampling the oracle, correctly interpreting information, and consistently exploiting it—going well beyond random exploration or Thompson sampling. State machine analysis reveals why: predictive coding creates belief representations that track information and drive information-directed exploration. While additional benchmarks would strengthen this claim, the insight should transfer broadly to exploration tasks sharing the core structure.

Q6: … training time compare to RL²?

Overhead comes from the 2 decoders. Compared to RL²: Training time increases by ~30-40%; Inference time is comparable. We will add these metrics in the revision.

Q7: How sensitive is … bottleneck dimension?

We evaluate different bottleneck dimensions (ranging from 1 to 32). Results show when:

• Dimension > task complexity (e.g. > 2 for the oracle bandit): achieve similarly low dissimilarities, indicating that regularization helps to maintain interpretable belief representation when the bottleneck is too expressive.
• Dimension < task complexity: both performance (suboptimal return) and representation (high dissimilarity) significantly degrade, indicating that bottleneck needs to surpass task complexity to solve the task.

We will include the results in the revision.

Q8: ... OOD generalization … on task distributions that differ from training?

Thank you for this excellent suggestion. We fully agree that OOD generalization is crucial, and while comprehensive OOD evaluation requires careful design of training and test distributions, we have two promising additional results:

• Different Tiger accuracies (zero-shot): For models trained on acc=0.7 and tested on acc=0.8 environments, our models show near-optimal return (-15.94±3.82, p<0.05), whereas RL² has significantly lower test return (-25.56±1.80), indicating our models support zero-shot by capturing the underlying belief updates shared across similar tasks.
• Oracle bandit generalization (transfer learning): We train models on task distributions that only contain arms 1-5, before transferring them to task distributions that contain arms 6-10. Our models showed significantly faster transfer learning as compared to RL² (return at 2e4 updates: ours 20.20±0.82 vs RL² 14.68±1.51, p<0.05), indicating our models learn the underlying task logic rather than memorizing specific arm configurations.

We’ll include these results in the revision.

Q9: state machine simulation ... unfamiliar with Mikulik et al

We'll expand the description in the main text.