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

Po-Chen Kuo, Han Hou, Will Dabney, Edgar Y. Walker

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

TL;DR: Meta-RL with self-supervised predictive coding modules can learn interpretable, task-relevant representations that better approximate Bayes-optimal belief states than black-box meta-RL models across diverse partially observable environments.

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.

Primary Area: Neuroscience and cognitive science (e.g., neural coding, brain-computer interfaces)

Submission Number: 25138