Research map / Reinforcement learning / Visual and representation RL
Predictive Coding Enhances Meta-RL To Achieve Interpretable Bayes-Optimal Belief Representation Under Partial Observability
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Shows that meta-RL agents augmented with a VAE-based predictive coding module learn belief representations closer to Bayes-optimal states than RL² across six tractable POMDPs, verified via state machine simulation.
Problem
Black-box memory-based meta-RL (RL²) reaches near Bayes-optimal policies but its learned history representations are not equivalent to minimally sufficient Bayes-optimal belief states, hurting interpretability, generalization, and learning of information-seeking exploration.
Approach
Train an RNN-based VAE encoder to produce a probabilistic belief bottleneck by predicting next observations and rewards (with KL regularization acting as Bayesian filtering prior), and feed that belief to a feedforward actor-critic policy trained without gradients flowing into the encoder; representational equivalence is assessed by bidirectional state machine simulation against Bayes-optimal beli
Key results
- In the oracle bandit, only meta-RL with predictive modules learns the Bayes-optimal policy, while RL2 converges to a suboptimal policy.
- In the hardest dynamic Tiger setting (observation accuracy 0.7), only meta-RL with predictive modules approaches Bayes-optimal return; bottleneck D_s is significantly lower across tasks.
- Models trained on dynamic Tiger (accuracy 0.7) and tested at accuracy 0.8 reach -15.94±3.82 zero-shot return for predictive modules versus -25.56±1.80 for RL2.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Add VAE KL regularization (ablation) | Meta-RL with predictive modules · Two-arm Bernoulli bandit | state dissimilarity: 2.44±0.38; 0.21±0.04 | Helped |
| Add VAE KL regularization (ablation) | Meta-RL with predictive modules · Dynamic bandit | state dissimilarity: 0.42±0.38; -0.65±0.15 | Mixed |
| Add VAE KL regularization (ablation) | Meta-RL with predictive modules · Dynamic Tiger | output dissimilarity: 1.11±1.35; 1.49±0.6 | Helped |
| Add VAE KL regularization (ablation) | Meta-RL with predictive modules · Oracle bandit | state dissimilarity: 23.47±6.94; 0.24±0.82 | Helped |
| Train the RNN encoder solely with predictive loss rather than also using RL policy gradients (ablation) | Meta-RL with predictive modules · Two-arm Bernoulli bandit | dissimilarity: -0.41±0.05; -0.61±0.15; -0.27±0.12; -0.08±0.04 | No clear effect |
| Train the RNN encoder solely with predictive loss rather than also using RL policy gradients (ablation) | Meta-RL with predictive modules · Dynamic bandit | dissimilarity: -0.09±0.06; 0.08±0.13; 0.0±0.04; 0.47±0.2 | No clear effect |
| Train the RNN encoder solely with predictive loss rather than also using RL policy gradients (ablation) | Meta-RL with predictive modules · Dynamic Tiger | dissimilarity: 1.1±0.96; 0.21±0.05; 0.95±1.66; 0.62±0.84 | No clear effect |
| Train the RNN encoder solely with predictive loss rather than also using RL policy gradients (ablation) | Meta-RL with predictive modules · Oracle bandit | dissimilarity: 10.91±1.02; 0.24±0.82; 5.53±5.0; 1.77±3.38 | No clear effect |
| Use predictive modules instead of black-box RL² | Meta-RL · Dynamic Tiger, observation accuracy 0.8 test | test return: -25.56±1.80 -> -15.94±3.82 | Helped |
| Use predictive modules instead of black-box RL² | Meta-RL · Oracle bandit, transfer from arms 1-5 to arms 6-10 | transfer learning speed | Helped |
Surprising findings
- Untrained recurrent layers can have low state dissimilarity to Bayes-optimal states, so decoding alone is insufficient to assess representational equivalence.
- Adding joint RL policy gradients to the encoder does not further improve bottleneck representations over predictive loss alone.
Limitations
- Only POMDPs with tractable Bayes-optimal solutions are studied, limiting scalability of the analysis.
- Only next-step future prediction is used as the predictive objective; multi-step predictions are not explored.
- Generalization evaluation is limited to a few zero-shot and transfer settings.
Experimental setup
Datasets: two-armed Beta-Bernoulli bandit dynamic two-armed bandit stationary Tiger dynamic Tiger oracle bandit latent goal cart
Benchmarks: two-armed Bernoulli bandit dynamic two-armed bandit stationary Tiger dynamic Tiger oracle bandit latent goal cart
Metrics: episodic return vs Bayes-optimal state dissimilarity D_s (MSE) output dissimilarity D_o zero-shot test return transfer learning speed
Baselines: RL2 Bayes-optimal policy (value iteration / Gittins index)
Builds on: RL2 state machine simulation advantage actor-critic variational autoencoder Dreamer
Closest papers by meaning
- Bridging State and History Representations: Understanding Self-Predictive RL (ICLR 2024 · similarity 0.88)
- Predictive auxiliary objectives in deep RL mimic learning in the brain (ICLR 2024 · similarity 0.87)
- Skill-aware Mutual Information Optimisation for Zero-shot Generalisation in Reinforcement Learning (NeurIPS 2024 · similarity 0.87)
- Meta-Cognitive Memory Policy Optimization for Long-Horizon LLM Agents (NeurIPS 2026 · similarity 0.87)
- Divide-and-Conquer Predictive Coding: a structured Bayesian inference algorithm (NeurIPS 2024 · similarity 0.87)
- Meta-DT: Offline Meta-RL as Conditional Sequence Modeling with World Model Disentanglement (NeurIPS 2024 · similarity 0.86)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.