Predictive Coding Enhances Meta-RL To Achieve Interpretable Bayes-Optimal Belief Representation Under Partial Observability
Authors: Po-Chen Kuo
Han Hou
Will Dabney
Edgar Y. Walker
Venue: NeurIPS 2025 poster ()
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
Primary area: neuroscience_and_cognitive_science
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.

Predictive Coding Enhances Meta-RL To Achieve Interpretable Bayes-Optimal Belief Representation Under Partial Observability
Po-Chen Kuo
Han Hou
University of Washington pekuo@uw. edu
Allen Institute for Neural Dynamies han. pettaaaleinnntiiee org
Will Dabney Google DeepMind wdabney@gmail.com
Edgar Y. Walker University Washington eywalkereus.ed edu
Abstract
Learning compact representation of history is critical for planning and general- ization in partially observable environments. While reta-reinforcemmnt learning (RL. agents can attain near Bayes-optimal policies, they often fail to learn the com- pact, interpretable Bayes-optimal belief states. This representational inefficiency potentially limits the agent's s adaptability and generalization capacity. Inspired by predictive coding in neuroscience. -wiich suggests that the brain predicts sensory inputs neural implementation of Bayesian inference- and by auxiliary pre- dictive 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 RL across wide variety of tasks, even when both achieve optimal policies. In chal- lenging tasks requiring active information seeking, only meta-RI. with predictive modules successfully learns optimal representations and policies, whereas ventional meta-RL. struggles with inadequate representation learning. Finally, we demonstrate that better representation learning leads improved generalization. Our results strongly suggest the role of predictive learning as a guiding principle for effective representation learning in agents navigating partial observability.
Introduction
In real-world environments, agents, biological or artificial, rarely have complete information about the environmental states crucial for decision-making and planning, observations often noisy and non-stationary. This issue known in reinforcement learning (RL) as partial observability 2 3. and is a major challenge in the deployment of real-world RL systems environments, learning optimal policies depends on the entire history of past interactions. An open question in RL under partial observability is how learn efficient representation that serves compact summary of the history while retaining the information crucial for policy learning.
4 In partially observable
Partially observable tasks can be formalized as partially observable Markov Decision Processes (POMDPs, Fig. (T) 2 3 allowing a Bayesian treatment by maintaining and updating a belief state using Bayesian inference [5]. To this end, meta-RL, in particular memory-basee meta-RL, has been shown provide powerful deep RL methods for developing agents that efficiently adapt under
39th Conference on Neural Information Processing Systems (NcurIPS 2025).

---

uncertainty 7 8 10 TT Meta-RL has been shown to behave near Bayes-optimall under partial
observability, both theoretically and empirically [ However, its learned representations are not equivalent to the minimally sufficient, Bayes-optimal belief states [13] hindering its interpretability,
robustness, generalization capability, and learning of complex exploration [10] 13
Predictive learning has been suggested as a central coding principle that guides efficient neural processing in neuroscience. Humans and animals are shown to achieve near Bayes-optimal behavior under uncertainty, leading to theories and experiments investigating efficient Bayes-optimal repre- sentation and computation in the brain 14 15 16]. Predictive coding, postulating that the brain continuously predicts upcoming observations and updates its internal world models [ 18), provides neurally plausible implementation of Bayesian inference [19] 20] and has been instrumental in explaining diverse neural circuits including feature learning in sensory cortices [2 [2 motor control in the cerebellum [33, value learning in the striatum [24]. and cognitive maps in the hippocampus powerful self-supervised technique to support representation learning and improve task performance 266 27 In deep RL., predictive objectives are shown regularize learning and prevent overfit or representational collapse [28] 29],
[25]. In deep learning, predictive learning has been explored
improve sample efficiency [30], and link to Bayesian filtering [1. Most recent breakthrough employs emplyss predictive learning to derive general agents that can solve a wide range of challenging tasks (2)
The main contribution of this paper is a systematic investigation of whether integrating self supervise predictive coding into meta-RL yields interpretable, Bayes- Eaytmmlptimal belief representations in partially observable environments. Although recent meta-RL models have explored predictive objectives to improve generalization and representation quality 19 10), the exact mechanisms by which predictive learning enhances latent representations -ad how this relates to improved performance remain unclear. To fill this gap. we move beyond performance measures that lack interpretability to directly compare meta-RL agents against Bayes-optimal solutions, evaluating both behavior and representation across diverse partially observable tasks varying in belief-update complexity. Specifically, we:
Employ meta-RL framework incorporating self-supervised future predictive modules, which enables direct interpretation and comparison with Bayes-optimal belief states.
Show systematically that meta-RL with predictive modules yields interpretable, task-relevant representations that capture underlying environmental structures and dynamics across diverse POMDPs, and significantly improves planning, exploration, and generalization.
Demonstrate via rigorous state machine simulation analysis that meta-RL with predictive modules consistently generates representations that more closely approximate Bayes-optimal beliefs than conventional black-box meta-RL, providing guiding principle for learning both optimal representations and policies in partially observable environments.
2 Background and related work
Partially observable Markov decision process (POMDP) A POMDP is defined by the tuple (S, A, S2, O,7). where S is the (finite) hidden state space, A the action space, T(s' s, a) the transition probabilities, I(ss.al a) the reward function, 5i the observation space, O | s', a) the emission function, and E [0, 1) the discount factor. Since the agent does not have access to the true state s, the decision process is non-Markovian with respect to $ and an optimal POMDP policy generally depends on the entire history T. (o1...., Qz-1. 0t). By maintaining belief state by = =p(s[h,,, posterior over S summarizing the history, the POMDP can be cast as a cully-observable belief MDP with transitions in belief state given by the Bayesian update: V/ (ss') 1) (( s',a) E.es T
s, a) b(s), where 17 is normalization factor. In this formulation, belief states b are the minimally sufficient statistic of the history summarizing all the past information compactly, and the policy π (b) over belief states restores the Markov property. Unfortunately, planning in a POMDP, i.e., computing Bayes-optimal policy over belief states, is generally intractable for all but the simplest tasks 2
Memory-based meta-RL Exemplified by RL 6 memory-based based meta-RL. RL emerges as a powerful method for learning adaptive policies across related tasks. Often parameterized as recurrent neural networks (RNNs), RI.2 implicitly maintains memory of the past history in the recurrent states and learns a history-dependen policy by maximizing return over the task distribution using policy gradient algorithms (Fig. [B). This design enables the agent to adapt dynamically it interacts with

---

A) POMDP environment
BY RL
Cj Meta-RL with predictive modules
h
hidden state
$
Obaervable the agent
0,
d
0
RNN r,
Re
Osservation Haward Acdon
Dp-1
4,1
RNN
V,
d,1
b, ((m m1
A14 4
To
.0et
The agent's beleeff stato
b
De-
T
Figure I: Representation learning under partial observability using meta-RL with predictive
modules. A) In POMDPs, belief state b, over the environment hidden state is maintained for policy learning. B) Memory-based meta-RL (e.g. RL²) simultaneously learns representation and policy using only the reward signal. C) The proposed meta-RL with sulf-suuer predictive modules separates representation learning with predictive coding and policy learning with reward signal.
the environment, effectively learning an internal RL algorithm via its recurrent computation M.. Through end-to-end meta-learning over task distribution, RL2 provides effective black-bos agents that achieve rapid adaptation to new task instances and generalize across related tasks.
The training objective of meta-Ro RL can be recast and understood as approximating the Bayes-optimal policy for a task distribution [2]. Since history sufficient for computing the belief state, black-box meta-RL. can in principle learn Bayes-optimal policies by encoding the belief state in its hidden state. Several studies have demonstrated that belief states can be decoded from RNN hidden states 10 33 34]. Moreover, Ortega et al. [12] provided theoretical framework for interpreting meta-RL RL as Baycs-optimal solutions, and Mikulik et al. [13 empirically showed that meta-RL agents learn near Bayes-optimal policy in bandit problems. However, comparison of meta-learned representations in RL2 using state machine simulation suggests that their representations are not equivalent to the minimally-sufficies
Bayes-optimal states, likely due to failures in injectivity, which may hinder the agent's 's interpretabilityy, robustness, and ability to generalize [3]. In fact, in practice, black-bor meta-RL RL may struggle efficiently learn Bayes-optimal policies in tasks requiring exploration that is temporally-extended and qualitatively different from exploitation behavior 9
Predictive coding neural circuits In neuroscience, theoretical and experimental studies have explored predictive processing fundamental computational objective across brain regions and as a neural implementation of Bayesian computation. Sensory areas are often modeled as hierarchical neural networks predicting future sensory inputs 21
21 22 potentially supporting efficient neural representation for perceptual inference [15]. Motor systems are thought to predict the sensory consequences of the motor outputs (36), which can underlie efficient integration in sensorimotor
learning [16]. The hippocampus, with neural activity believed to forecast an animal's upcoming experiences, is critical for relational learning and computing a predictive map or transition model to support efficient navigation (37 38 Building on the relationship between predictive learning and efficient neural representation, recent work has begun to explore adding predictive objectives to encourage artificial agents to develop activities similar to those observed in the brain [39].
Predictive learning in meta-RL Recent work has explored alternative architectures or auxiliary training objectives to learn latent representations that are predictive of future observations or rewards for augmenting meta-RL agents. For instance, Igl al. X showed deep policy can be improved by adding predictive auxiliary loss, and proposed particle filter method for approximating belief inference. Several studies further decoupled learning recurrent generative models of the environment and used the learned latent representations as inputs for model-free RL algorithms, demonstrating improved performance especially for robotics tasks that are difficult for black-box models 40) 41 42 43). For a special type of POMDP where hidden states are stationary, several works considered approximating task inference by taking advantage of privileged information 44,, posterior sampling [45), or variational inference [9]. providing effective meta-RL models for rapid adaptation to new task instances. Recent works have explored how to further improve representation learning through state- space modeling [10] normalizing flow technique [46), or decoupled belief modeling with separate
random policies 47]. On the other hand, learning latent dynamical world models has also been proposed for model-based deep RL approaches, including SOLAR [48,, PlaNet [30] and Dreamer
3

---

32]. These studies demonstrated that learning latent world models through multi-step predictive objectives can strongly improve planning and performance in challenging control benchmarks.
Complementary to the above empirical progress of integrating predictive learning to improve perfor- mance, work provides the interpretability foundation for why predictive objectives work in deep RL by focusing on understanding its meta meta-learned representation and computation.
3 Meta-RL with self-supervised predictive modules
Black-box meta-RL like RL2 (Fig. (3) requires simultaneously learning history representation and policy using only reward signals, potentially suffering from inadequate representation and policy learning 11 13 To overcome this challenge, we employ an end-to-end me framework with self-supervised predictive modules to learn explicit belief representations. Our motivations are two-fold: algorithmically, a good representation should summarize the relevant history predictive of the future and reflect uncertainty; neurobiologically, modular neural circuits are believed to perform future predictions, attaining efficient representations under predictive coding [17].
The proposed model, meta-RL. with predictive modules, consists of a variational autoencoder (VAE) for predictive representation learning and policy network operating on the learned representation (Fig. [T]. The predictive modules begin with an RNN encoder qg. which takes as input current observation Oz, reward rr. and previous action d; The RNN outputs low-dimensional bottle- neck b, which is the posterior distribution over the latent states m, conditioned on the history T. The posterior is trained via reward and observation prediction akin to predictive coding using the reward decoder Rg and the observation decoder Te, respectively, to predict upcoming rewards and observations given the action taken by the policy network. The predictive modules optimize the evidence lower bound (ELBO) using the reparametcrization technique 49), with KL regularization loss similar to Bayesian filtering (see A. for details). This approach learns explicit probabilistic belief representation by over the latent state, facilitating a direct interpretation as belief states in POMDPs. The policy network TT. is feedforward neural network receiving the belief state b, as an input and can be efficiently trained using model-free policy gradient algorithms
[50]. Note the policy gradient from RL loss is only used train the policy network but not the RNN encoder. The entire model is trained in a self-supervised, end-to-end manner using the standard meta-learning paradigm.
Our parameterization is similar in principle to previous work 9 10 41), with designs tailored for POMDPs without strict assumptions over stationarity [9] and structures [0] of the hidden states, facilitating direct interpretation of b, as the belief states in POMDPs. Furthermore, the end-to-end approach neither relies on privileged information 44 nor requires a separate random policy to generate samples for training predictive models 47 exploited in previous work.
Experiment
We designed the following experiments to answer the central question: Does meta-RL with self- supervised predictive modules enhance learning representations that more closely match the Bayes- optimal belief states than black-box meta-RL? Specifically, we adopt the state machine simulation analysis used in Mikulik et al. 13 to examine the equivalence of representation and computation between meta-RL. agents and Bayes-optimal solutions. In essence, two state machines can be considered computationally equivalent if they can simulate each other-thar is, if for any given state in one machine, we can find a mapping onto the other machine such that both their state transitions and oulputs are the same (51). This analysis examines whether there exists consistent way of interpreting every state in one machine as a state in the other. Although decoding is commonly used to evaluate representation similarity [T0] 33 34]. it only measures correlations between representations. In contrast, state machine simulation thoroughly assesses representational equivalence based on structural and computational relevance, providing a rigorous analysis for model interpretability.
State machine simulation In tasks where Bayes-optimal states are computationally tractable, we can compare meta-RL representations with Bayes-optimal belief states using state machine simulation [33. The goal is to measure the state and oulpul dissimilarities after mapping the states from one machine to the other. If both state and output dissimilarities are low in both mapping directions, then the two representations can be considered computationally equivalent. Briefly, the

---

A) Model performunce
B) Visualization the representution space
Meta-RL with
RL
predictive modules
Epicdoe return
20
30
Untrained
Fully trained
0
Bayes-optima
-10
R2.
Meto with Prodictive Modullss
-20
PC
20
PC1
-10
PC
30
00
State machine simulation
De: metaRL Bayes
metaRL Bayes
Dg: Bayes metaRL
Dj: Bayes metaRL
RL2
Recurrent
0.017
-0.149
0.012
0.015
+0.144
0811
0.205
Bottlaneck
0.015
0.054
0.012
0.146
1.054
0.184
Meta-RL W th predictive modules
Recurrent
00.014
10,007
M012
0.214
Bottleneck
0.015
0:125
0.018
10.105
0.002
0.132
M.036
0.204
0.1
0.2
0.8
0.1
0.0
0.2
0.2
0.3
| Imstaft: Fimetaii Sayen
M5EnetaRL
Bayas
Faayes
Caaess cowGaR
MSEsayos mettat
Figure 2: Meta-RL with predictive modules better approximates Bayes-optimal states in bandit task. A) Model performance against the Bayes-optima policy. Light and dark colors denote untrained and fully trained models, respectively. B) Visualization of Bayes-optimal belief states, learned representation in the bottleneck layer of RL and that of meta-RL with predictive modules. States are colored by the probability of choosing 42- Black curves show one example trajectory. C) State machine simulation. Recurrent and bottleneck denote the layers used for analysis. The bottleneck layer of meta-RI. with predictive modules achieves the lowest state dissimilarity D, and output dissimilarity D, in both mapping directions, indicating highest equivalence to Bayes-optimal states. (For all figures error bars denote the s.e.m. across at least training runs, and asterisks denote statistical significance (p < 0.05, (-test) among the trained models.)
procedures are as follows (and see A.4 for details): (i) two mapping functions (parameterized as multi-layered perceptrons) are trained map from meta-learned states to Bayes-optimal states and from Bayes-optimal states meta-learned states; (ii) The state dissimilarity D, is measured the mean square error (MSE) of the mapped states against the target states; (iii) The output dissimilarity Do is measured as the difference in return generated by the output of the original states and the output of the mapped states. Here we compare D, and D, of meta-RL. with predictive modules against those of black-box meta-RL RL (RL?) to evaluate the quality of their learned representations.
To enable rigorous comparison between meta-RL agents and Bayes-optimal solutions, select the following exemplar tasks for which Bayes-optimal solutions are computationally tractable. While previous work has only evaluated stationary bandits [13], our task suite covers diverse POMDP challenges: explore-exploit tradeoff, sequential decision-making, dynamic belief tracking, information gathering, and continuous control. To our knowledge, this is the first systematic evaluation of meta-RL representations using ground-truth comparison across diverse POMDPs. Following are the six classes of tasks evaluated (details of each task are described in A.2):
Tasks
The first task is the classic multi- bandit also considered in Mikulik et al. [1], where the meta-learned representations in RL models were found to deviate from the minimally sufficient Bayes-optimal states even though their policy converge the Bayes-optimal one. We will examine if the proposed meta-RI. with predictive modules can effectively learn Bayes-optimal belief states in two-armed bandits.
Classic two-armed Bernoulli bandit:
Dynamic two-armed bandit: To evaluate dynamic belief tracking, we next consider a dynamic extension of the two-armed bandit. Each arm is in one of n possible discrete states (for simplicity we choose n = 2) with various reward probabilities. The state of each arm evolves according to an independent Markovian transition process. The task is designed such that the belief state update is analytically tractable, and that the Bayes-optimal policy can be derived using the value iteration algorithm over the belief space []. To examine different structural and dynamical configurations, we consider three parameter settings: (i) two arms of symmetric reward probability states and transition dynamics, (ii) two arms of asymmetric reward probability states, and (iii) two arms of asymmetric transition dynamics.
5

---

Stationary Tiger: To exemplify sequential decision- under uncertainty, we consider the classic POMDP Tiger task [3). where an agent chooses between two doors hiding a tiger (penalty=-10l 100) and the other hiding a treasure (reward=10). The agent may addi- tionally choose to pay a small penalty=-1 to "listen" to acquire noisy observations about the tiger's location. This simple yet powerful paradigm tests an agent's 's ability to balance information gathering with reward seeking. We consider two difficulty levels by varying the observation accuracy for the "listen" action.
We extend the Tiger task to a dynamic version by allowing the tiger's s location to change over time, following Markov transition dynamics. As information reliability and relevance are further corrupted by the dynamic nature of the hidden state, the agent has to balance listening more to increase confidence against making decisions earlier in case the information gathered becomes obsolete. This task design permits tractable belief updates and the Bayes-optimal policy can be derived using value iteration []. Similarly, we consider two difficulty levels by changing the observation accuracy.
Dynamic Tiger:
Black-box meta-RL often struggles when exploration and information seeking are required [TT] We hypothesize this is due to ineffective representation learning. To exemplify this point, we consider an illustrative oracle bandit: in an 11 1-arm bandit environment, one of the first ten arms 01-10 is the target arm that gives a payout of 5, whereas the remaining nine are non-target arms with payout of 1. The last arm @11 is the "oracle" arm whose payout is < but informs the target arm in the form of J of the target TO arm index (e.g. reward of 0.3 from a indicates 0g is the target arm). This is similar to tasks discussed in Wang et al. 77 but differs in that here no structured feedback is given, which makes representation learning more challenging and critical. A successful policy requires paying an immediate exploration cost to acquire information for long-term gain.
Oracle bandit:
Latent goal cart: Deriving Bayes-optimal solutions in continuous POMDPs are usually intractable, hindering rigorous representational equivalence analysis. To enable evaluation in continuous observation and action spaces, we design an exemplar continuous control task which still permits tractable Bayes-optimal belief inference and policy. In this task, an agent controls a continuous action, the velocity of a cart, to move along a I-dimensional track to a hidden goal (+1 or 1) which needs to be inferred from the continuous observation (current position) and reward (noisified distance from the hidden goal) it receives.
Agents For each POMDP environment, two types of agents the proposed meta- RL with predictive modules and the black- bbx meta-RL (i.e. RL²)-wa wwer trained on the target task distribution. To facilitate comparison, the bascline RL2 model is designed to be architecturally identical to meta-RL with predictive modules except for the decoder networks. Specifically, the Ri.² baseline consists of an RNN followed by a fully connected bottleneck layer which is the counterpart to the latent belief layer by in TC and finally with a readout layer that generates action logits and value estimates (details see A.3). Identical layer sizes are used across both models throughout each experiment.
While the latent belief layer (i.e. the bottleneck layer) for meta-RL with predictive modules the natural target to be interpreted as belief states, Mikulik et al. 1 evaluates the representations learned in the recurrent layers in 2 models. In what follows, we systematically analyze the representations learned in both the recurrent and bottleneck layers in each model using state machine simulation.
5 Results
Two-armed Bernoulli bandit The Bayes-optimal solution can be derived using the Gittins index method (52). After training, both RL² and meta-RL with predictive modules approach Bayes-optimal return (Fig. AA). Visualization of the bottleneck layers shows meta-RL with predictive modules learns a dimensiooa representation structurally more similar to the Bayes-optimal states than RL2 (Fig. 2B). This observation is corroborated by the state machine simulation analysis results (Fig. 2C). Before training, both models have high D, and D.. indicating the untrained representations are far from Bayes-optimal. One notable exception is D, of the recurrent layers for the meta-RL BBayess mapping direction, which further verifies that using decoding alone is not enough for evaluating the equivalence of two representations, as the untrained recurrent layers may maintain verbose representation of the history and can be mapped arbitrarily close to the belief states if powerful enough mapping functions are used (see A.4 for detailed discussions).

---

Symmetrio reward & transition
Dj: metaRL - Bayes
D2 metaRL Bayes
Dg: Bayes metaRL
Dj Bayes + metaRL
E20000
Recumcnd
0.002
0.119
1 22v 84
-0002
0,314
01103
0.371
untarre Fuly timees
Bottienece
0.087
955 05
22--22
0.121
-3375
Meta-
Recurberc
+0.32
5-0130
B) Asymmetric reward
Predittoe Meda en
and cive
mocules
Bentinecc
0.9
0.4
092
0.2
0.0
0.6
0.0
0.1
175
0.2
0.0
0.2
04
033
Splekke roum
Becurreni
80:15
082
75m
0.329
Hartiied Fully moa opeinall
Burtenock
045
360-05
0.085
0.312
120
Meta RL wch
Nacumnnn
00212
0.014
0.05
mocules
Buttkreec
1000
330-05
0004
3.089
0.366
0.4
0.2
0.3
0.0
0.2
6,4
Asymmetric transition
Recumerr
1:00--14
0.065
0.372
0.049
0.051
M050
3.392
Fiils Dayer sprinal
Matz with
10:308
0.007
1
1.29
04
0.215
0.007
0014
0.376
predictive
but
041
0 230 02
0.091
0,388
PROOKING MUOLRS
41
Kmnstutt
Ferstotl- Sauual
MSE MSE
Bayca
Taayas Taayas mmttaa
MSE Paycs motont
Figure 3: Dynamic bandits. Performance (left) and state machine simulation analysis (right) for A) the symmetric task, B) the asymmetric reward task, and C) the asymmetric transition task. Shaded areas in the left column denote the standard deviation of episodic return under Bayes-optimal policy. The bottleneck layer of meta-RI. with predictive modules attains the lowest state dissimilarity D, and output dissimilarity D, in both mapping directions across tasks.
After training. for the recurrent layers in both models, the state dissimilarity D, remains high for the the minimally-sufficiae direction, highlighting the representations in the recurrent layer deviate from Bayes-optimal belief states, similar to the findings in Mikulik et al. [13 In contrast, considering all four measures, the bottleneck layer in me with predictive modules achieves the lowest D& and D, in both mapping directions, indicating that its learned representations best approximate the Bayes-optimal belief states (Fig. 2C). Remarkably, the representation learning capacity does not merely arise from having a low oow-dimeniional bottleneck, as De for the bottleneck layer in RI.2 remains significantly higher than that in meta-RL. with predictive modules, suggesting that predictive coding encourages more compact, efficient representations in meta-RL agents.
Dynamic bandits Across the three dynamic bandit
symmettii reward and transition, asym- metric reward, and asymmetric transition (Fig. 3A, B, and, C, respectively), both models approach the optimal return. Meta-RL with predictive modules learns representations structurally more similar to the Bayes-optimal states (see (see A.S.II. State machine simulation reveals that the recurrent layers in both models present with large D, for the Bayes metaas-mmetRttRt RL direction even after training. While the bottleneck layers in both models achieve comparably low output dissimilarities, meta-RL with predic- tive modules consistently achieve significantly lower state dissimilarity D, in the Bayes -eeess-mmeaRRR RL mapping direction, indicating a higher equivalence to the Bayes-optimal belief states.
Stationary Tiger and Dynamic Tiger Difficulty of the Tiger tasks increases when the observation accuracy decreases (information becomes less reliable) and when the tiger's location is dynamic (information may become obsolete if waiting for too long). After training, both models approach the Bayes-optimal return in easier tasks, while for the most challenging dynamic Tiger with low observation accuracy, only meta-RL with predictive modules approach Bayes-optimal return (Fig. AD)). Meta-RL with predictive modules learns representations better capturing the task structure (see A. State machine simulation shows that the recurrent layers in both models fail to approximate Bayes-optimal belief states (large D2 after training in Fig. D. Comparing the bottleneck layers, meta-RL with predictive modules consistently achieves higher equivalence to Bayes-optimal belief states, as indicated by the significantly lower D, for the Bayes RL direction across all tasks.
Oracle bandit After training, only meta-RL with predictive modules learns the Bayes-optimal policy (Fig. 5A), paying an immediate cost to sample the oracle arm and utilize the information for long-term gain. In contrast, RI.2 converges to a suboptimal policy where it learns to sample the oracle arm upfront, but fails to use the attained information consistently. Visually, the bottleneck layer in
7

---

A) Stationary Tiger, observation
Dgi metaRL Bayes
D2: metoRL - Bayes
Dj: Bayes - metaRL
Dj: Bayes - mecaRL
accuracy
Untianed
Pully Fuared
RL?
Bottlenecs
80 80:033
00005
0.365
8 144.04
2 84c 54
2,24
15000
1004
0.245
0 0.29
0053
0.475
Meta
Becur
1o 021
-0 0.419
18-63
10.015
00
783
190
with
predictto mocules
boatlenec
M,026
0:458
3 250
0.450
0
Fred clhee
0.3
08
0.0
033
0.8
B) Stationary Tiger, observation accurncy 0.7
E502000 rourn
Pasuman 00033
0.025
0012
Untratmed Boytassptiral
sontonece 0003
0.000
02
8.3
0015
0.465
Meta RL
126
0.042
with predi Cliv
74M-83
55 116
Bonteneess 316
424
9tt
179
R1T Hett- with
Predictive Monure
0.0
C) Dynamic Tiger, observation
Bocuracy = 0.8
Necument
162
0.310
1.201
-25
Bottleneck
366
32e
0.344
Me with predective modulies
Recument
Bottleneck
-106
0.0
0,3
0.6
0.0
0.1
00
0.3
06
0.0
0.3
0.6
D) Dynamic Tiger, observation accuracy
0.7
Recurter $0.523
9. 182
20a-U4
9.351
Epodbe
Untaiitd
Santeeeec
0.507
0.365
Dapre-optiral
with predictive modules
Mata: RL
0.420
0.356
Sasterecc
10
-130
0.0
6.3
04
0.5
63
0.0
0.3
0.0
0.3
0.6
Prod Cthe Modules
Tmataat Tmetaa, - Daysol
M5E meLaTA dayss
Traaas TRayax-metahh
MSE Bayee mestlt
Figure 4: Stationary and dynamic Tiger. Performance (left) and state machine simulation analysis (right) for A-B) stationary Tiger with two levels of observation accuracies, and C-D) dynamic Tiger, where the tiger location changes with a probability of 0.9 each timestep. Note for the most challenging setting in D, only meta-RL. with predictive modules approaches the Bayes-optimal return.
meta-RL with predictive modules shows interpretable representations capturing task structures and dynamics. In contrast, RL2 fails to learn an interpretable representation (Fig. 5B,, likely contributing to its suboptimal behavior. State machine simulation shows that after training, the bottleneck layer in meta-RL RL with predictive modules achieves significantly lower state and output dissimilarities in both mapping directions, indicating close approximation of Bayes-optimal belief states (Fig. (5C). Dissimilarity measures remain high for RI. after training. suggesting that RI.2 may suffer from inadequate representation learning when the task requires exploration and information seeking.
A) Model performance
B) Visualization of the repreacutation space
Meta-RL
Episooe return
25
20
0.4
RL 2
with predictive modules
arm 1 arm
arm
15
0.2
Untrained
0.0
arm am
Rully trained
Bayes aptimal
0.2
amm
Mata--R with Prodictive Modules
0.0
PC
PC 1
PC 1
amm 10
arm
C) State machine simulation
Do: metaRL - Bayes
Dy: metaRL Bayes
Dy: Bayes melaaL
D,: Bayes melaRL
RL2
Bottieneck
Recurrent
0.112
0.213
0.582
0.623
DN22
0.00 0.07
096
0.140
0.116
0:218
10 686
686
00.312
3,330
0.517
0.341
Meta- ALL
with predictive
Recurrent
0.052
0.626
0.071
0.072
0:042
0.317
0.309
madules
Bett eneck
0 040 *
1.577
00000
0.103
10050
*
0,686
0.024
0.339
C.0
0.3
0.6
0.0
0.1
0.2
3.3
0.0
0.3
0.6
0.9
0.0
0.3
0.6
|ieetot - Finetait Dayees
MSE Bayes
[rayes "Dayes mecertl
MSE meturt
Figure 5: Meta-RL with predictive modules learns interpretable Bayes-optimal solutions in oracle bandit. A) Only meta-RL with predictive modules learns the optimal policy. B) Bayes-optimal states (left), the bottleneck layer of RL2 (middle), and that of meta-RL with predictive modules (right). States are colored by the most likely arm to choose. C) The bottleneck layer in meta-RL with predictive modules attains the lowest state and output dissimilarities in both mapping directions.
8

---

Pgi metaRi Rayes
Dg metali Rayes
D2 Rayes metaRi
Dj Rayes metaRi
retur
:20
NLT
354
Untrained
Bottereck
0.357
0.024
0.317
20
Pully trained
maymeoptimml
MaradR.
wit
03339
modures
Rontereck
0.0
0.5
0.5
0.50
0.01
0.3
5.
0.0
0.0
Mobliiie Yookres
Coectatt Sipess
MSE -daees
Vaxyes floaes nesoet
Figure 6: Latent goal cart. Performance (left) and state machine simulation analysis (right). The bottleneck layer of meta-RL with predictive modules attains the lowest state and output dissimilarities in both mapping directions, indicating close approximation of Bayes-optimal belief states.
Latent goal cart After training, both models approach the Bayes-optimal return (Fig. b.. Visualiza- tion of the states shows meta-RL with predictive modules learns representations better capturing the task structure (see{A.3.I)) State machine simulation shows that the recurrent layers in both models fail to approximate Bayes-optimal belief states (large D, after training). In contrast, the bottleneck layers of meta-RL with predictive modules achieve the highest equivalence to Bayes-optimal belief states, as indicated by the significantly lower state and output dissimilarities in both mapping directions.
Improved generalization capability Thus far we present a systematic evaluation of representa- tional equivalence between meta-RL with predictive modules and Bayes-optimal solutions across diverse array of POMDP tasks. Our state machine simulation results strongly suggest that meta-RL RL with predictive modules can effectively learn interpretable representations that better approximate Bayes-optimal belief states, which in turn could lead to improved policy learning in challenging tasks, To further illustrate the importance and implications of learning Bayes-optimal belief state representa- tions in partially observable environment we examine whether learning better representations might lead to improved generalization capability to unseen but related tasks. To this end, we first evaluate zero-shot generalization capacity, As shown in Fig. 77A, when testing models trained on dynamic Tiger with observation accuracy of 0.7 on environments with observation accuracy of 0.8, meta-RL. with predictive modules shows near-optimal zero-shot test return, whereas RL has significantly lower test return, suggesting that through learning better representations, meta-RL. with predictive modules could capture the underlying belief updates shared across similar tasks to facilitate zero-shot generalization. Next, we evaluate whether representation quality impacts transfer learning efficiency in out-of-distribution (OOD) tasks. As shown in in Fig 7B, when models are first trained on oracle bandits with a pre-training distribution that only contain arms 1-5 before transferring to a second task distribution that contains arms 6-10, meta-RL with predictive modules shows significantly faster transfer learning as compared to RL This implics that with better representation learning, meta-RL RL with predictive modules could effectively learn the relevant environmental structure and dynamics to facilitate more efficient transfer learning to OOD tasks.
A) Zero-shot generlization in Dynamic Tiger
B) Accelerated transfer learning in Oracle Bandit
25
Episdoe return
Eoisodic return 12
Untrained
-25
Fully trained
Generalization
-50
Bayesaoptimal
MetaRI, with
predictive modules RL2
-100
RL2
Meta-RL RL with Predictive Modules
20000
GOOD
8000
Number updates
Figure 7: Better representation leads to improved generalization capability, A) When testing models trained on dynamic Tiger with observation accuracy of 0.7 on environments with observation accuracy of 0.8. meta-RL with predictive modules shows near-optimal zero-shot test return (-15.44 3.82, green), whereas 2L2 has significantly lower test return (-25.56 + 1.80, orange), indicating that meta-RL with predictive modules achieves better zero-shot generalization to unseen but relevant tasks. (Red and blue bars reference model performance from Fig. CC.) B) When models are first trained on oracle bandits with a task distribution that only contain arms 1-5 before transferring to second task distribution that contains arms 6-10, meta-RL with predictive modules demonstrates significantly faster transfer learning as compared to RL?,, indicating that meta-RL with predictive modules achieves more efficient transfer learning to out-of-distribution relevant tasks.
9

---

Table 1: Ablation studies: changes in bidirectional state and output dissimilarities of the bottleneck layers compared to the proposed model of meta-RL with predictive modules. (M for meta-RI. B for Bayes-optimal solutions. Bold indicates < 0.05.)
M-
ADM
- AD.-M
AD. M-
AD.
AD. A AD,-
Two-arm Bernoulli bandit (unit: le-2)
Dynamic bandit (unit: 1e-2)
No KL
0.12+0.4
2.44#0.38 -0.26:t0.36 0.21+0.04
-0.16+..53
-0.13+0.38 0.42+0.38
-0.65=00155
joint RL
-0.41=00.05 -0.61=0.15 -0.27=0.12 -0.08+0.04
-0.09=0.06
0.08+0.13 0.0=0.04
0.47=0.2
Dynamic Tiger (unit: 1e-2)
Oracle bandit (unit: 1e-2)
No KL
1.15
0.22+0.17
1.4920.6
1.11+1.76
9.22=4.44 23.47:6..4
56.328.43
0.24+0.82
joint RL
1.1+0.96 0.21+0.05 0,95#1.66 0.62#0.84 10.91 0.24+0.82 5.53+5.0
1.77+3.38
Ablation studies Finally, to pinpoint which algorithmic design choices drive the enhanced repre- sentation learning capacity in meta-RL with predictive modules, we performed two targeted ablations: (i) no KL: we remove the VAE's KL regularization to test whether latent regularization is necessary to enforce compact representations, and (ii) joint RL: we allow the policy gradients from RL loss to jointly train the RNN encoder alongside the predictive opposed to training the encoder solely with the predictive loss as our proposed model does-to to assess whether policy gradients provide additional representational benefits. Table [ machine simulation analysis on exemplar tasks (see and (A.5.3 for full results). First, without KI regularization on the latent space, bottleneck dissimilarities to Bayes-optimal states increase (i.e. worse alignment) on most tasks (except for the minor decrease in dynamic bandits), showing that suitable regularization may help to promote compact, belief-like representations. On the other hand, under joint RL training, bottleneck dissimilarities are not further decreased as compared to the predictive loss-only models, indicating that additional policy gradients do not further improve representations and hence are not necessary for learning Bayes-optimas representations. This indicates that the enhanced representation learning capacity can be largely attributed to predictive learning.
changes in dissimilarities from state
6 Discussion and conclusions
In this work we employed rigorous state machine simulation analysis to demonstrate that meta-RL with self-supervised predictive coding modules can effectively learn interpretable, near Bayes-optimal belief state representations across diverse partially observable tasks, whereas conventional black- box meta-RL fails capture the minimally sufficient representation. We further showed how the enhanced representation learning capacity in meta-RL with predictive modules not only supports effective information gathering and policy learning in challenging tasks, but also leads to improved generalization capability to unseen but relevant tasks. Our results align with predictive coding theories in neuroscience, suggesting that predictive objectives constitute a general computational strategy for achieving Bayesian computation, and highlight predictive learning as fundamental principle for guiding efficient representation learning in agents navigating partially observable environments.
Scope and limitations We considered POMDPs with tractable Bayes-optimal solutions to enable rigorous analysis. Future work on more scalable analysis methods for tasks with larger state space will be valuable. In addition, we examined next-step future prediction general predictive objective. Considering the empirical success of multi-step predictive objective in deep RL, one promising direction is to investigate how incorporating multi-step, different time scales, and temporal abstractions into predictive modules might further enhance representation learning. Finally, we showed representation learning affects an agent's ability for generalization, but more comprehensive evaluation is needed to explore how to meta-learn representations to support OOD generalization.
Broader Impact Our work advances model interpretability by linking meta-learned representations to Bayes-optimal beliefs, offering a Bayesian account for the success of predictive objectives in deep RL. This understanding is critical for developing safe, adaptable agents that operate under uncertainty and limited information -ucchuuhhh as in medical decision support and disaster response. Moreover, our approach shows how insights from neuroscience theories may contribute to opening the black box of modern AI. Conversely, uncovering the computational principles that drive representation learning in artificial agents may in turn shed light the mechanisms of neural computation in the brain.
10

---

Acknowledgementss We thank Patrick Zhang and Matthew Storm Bull for their insightful feedback and discussions. We gratefully acknowledge the community support from the University of Washing- ton Computational Neuroscience Center and the Allen Institute for Neural Dynamics. We also thank the anonymous reviewers for their valuable and constructive review.
Code availability
The codebase for model training and state machine simulation is available at:
References
[1] Richard S. Sutton and Andrew G. Barto. Reinforcement Learning: An Introduction. The MIT Press, second edition, 2018.
[2] Anthony R Cassandra, Leslie Pack Kaelbling, and Michael L Littman. Acting optimally in partially observable stochastic domains. In Aaai, volume 94, pages 1023-1028, 1994.
[3] Leslie Pack Kaelbling, Michael I. Littman, and Anthony R Cassandra. Planning and acting in partially observable stochastic domains. Artificial intelligence, 101(1-2):99-134, 1998.
[4] Gabriel Dulac-Amold, Nir Levine, Daniel J Mankowitz, Jerry Li, Cosmin Paduraru, Sven Gowal, and Todd Hester. Challenges of real-worl world reinforcement learning: definitions, benchmarks and analysis. Machine Learning, 1 (419):21192266) 2021.
[5] Jayakumar Subramanian, Amit Sinha, Raihan Seraj. and Aditya Mahajan. Approximate information state for approximate planning and reinforcement learning in partially observed systems. Journal of Machine Learning Research, 23(12): 2022.
[6] Yan Duan, John Schulman, Xi Chen, Peter L Bartlett, Ilya Sutskever, and Pieter Abbeel. RI2: Fast reinforcement learning via slow reinforcement learning. arXiv preprint arXiv:161..02779 2016.
[7] Jane X Wang. Zeb Kurth-Nelson, Dhruva Tirumala, Hubert Soyer, Joel Leibo, Remi Munos, Charles Blundell, Dharshan Kumaran, and Matt Botvinick. Learning to reinforcement learn. arXiv preprint arXiv: 1.05633, 2016.
[8] Maximilian Igl, Luisa Zintgraf, Tuan Anh Le, Frank Wood, and Shimon Whiteson. Deep variational reinforcement learning for pomdps. In International conference machine learning, pages 2117-2126. PMLR, 2018.
[9] Luisa Zintgraf, Sehastian Schulze, Cong Lu, Leo Feng, Maximilian Igl, Kyriacos Shiarlis, Yarin Gal, Katja Hofmann, and Shimon Whiteson. Varibad: Variational bayes-adaptive deep via meta-lcarning. Journal of Machine Learning Research, 22(289):1-39, 2021.
[10] Kei Akuzawa, Yusuke Iwasawa, and Yutaka Matsuo. Estimating disentangled belief about hidden state and hidden task for meta-reinforcement learning. PMIRR, pages 73-86, 2021.
[11] Jacob Beck, Risto Vuorio, Evan Zheran Liu, Zheng Xiong, Luisa Zintgraf, Chelsea Finn, and Shimon Whiteson. A survey of meta-reinforcement learning. arXiv preprinl arXiv:2301.08028, 2023.
[12] Pedro A Ortega, Jane X Wang, Mark Rowland, Tim Genewein, Zeb Kurth-Nelson, Razvan Pascanu, Nicolas Heess, Joel Veness, Alex Pritzel, Pablo Sprechmann, al. Meta-learning of sequential strategies. arXiv preprint
2019.
[13] Vladimir Mikulik, Grégoire Delétang, Tom McGrath, Tim Genewein, Miljan Martic, Shane Legg, and Pedro Ortega. Meta-trained agents implement bayes-optimal agents. Advances in neural information processing systems, 18188669-117733 2020.
[14] Marc O Ernst and Martin S Banks. Humans integrate visual and haptic information in statistically optimal fashion. Nature, 415(6870):429 433, 2002.
[15] Wei Ma, Jeffrey Beck, Peter Latham, and Alexandre Pouget. Bayesian inference with probabilistic population codes. Nature neuroscience, 9131:1432-1448. 2006.

---

[16] Konrad P Körding and Daniel M Wolpert. Bayesian integration in sensorimotor learning, Nature, 427(6971):24442444 2004.
[17] Rajesh PN Rao and Dana Ballard. Predictive coding in the visual cortex: a functional interpretation of some extra-classical erceptive-field effects. Nature neuroscience, 2(1):79-8))) 87, 1999.
[18] Karl Friston. A theory of cortical responses. Philosophical transactions of the Royal Society Biological sciences, (1011145):811888888 2005.
[19] Christopher Summerfield and Floris P De Lange. Expectation in perceptual decision making: neural and computational mechanisms. Nature Reviews Neuroscience, 15(1 (5(11):745-76, 2014.
[20] Rajesh PN Rao. Decision making under uncertainty: neural model based on partially observ- able markov decision processes. Frontiers in computational neuroscience, 4:146, 2010.
[21] Shohei Furutachi, Alexis D Franklin, Andreea M Aldea, Thomas D Mrsic-Flogel, and Sonja B Hofer. Cooperative thalamocortical cireuit mechanism for sensory prediction errors. Nature, pages 1-9, 2024.
[22] Chengxu Zhuang, Siming Yan, Aran Nayebi, Martin Schrimpf, Michael C Frank, James J DiCarlo, and Daniel LKK Yamins. Unsupervised neural network models of the ventral visual stream. Proceedings of the National Academy of Sciences, (18(3):20044191111 2021.
[23] Daniel M Wolpert, R Chris Miall, and Mitsuo Kawato. Internal models in the cerebellum. Trends in cognitive sciences, 2(9):338-2) -3333387 1998.
[24] Wolfram Schultz. Peter Dayan, and P Read Montague. A neural substrate of prediction and reward. Science, 275 1997.
[25] J Oeefee. The hippocampus cognitive map, 1978.
[26] Aaron van den Oord, Yazhe Li, and Oriol Vinyals. Representation learning with contrastive predictive coding. arXiv preprint arXi: 037488 2018
[27] Max Jaderberg. Volodymyr Mnih, Wojeiech Marian Czamecki, Tom Schaul, Joel Z Leibo, David Silver, and Koray Kavukcuoglu. Reinforcement learning with unsupervised auxiliary tasks. arXiv preprint arXiv:161 1.05397. 2016.
[28] Clare Lyle, Mark Rowland, Georg Ostrovski, and Will Dabney. On the effect of auxiliary tasks on representation dynamics. In International Conference on Artificial Intelligence and Statistics, pages 1-9. PMLR, 2021.
[29] Will Dabney, André Barreto, Mark Rowland, Robert Dadashi, John Quan, Mare G Bellemare, and David Silver. The value-improvement path: Towards better representations for reinforcement learning. Proceedings of the AAAI conference on artificial intelligence, volume 35, pages 7160-7168, 2021.
[30] Danijar Hafner, Timothy Lillicrap, Ian Fischer, Ruben Villegas, David Ha, Honglak Lee, and James Davidson. Learning latent dynamics for planning from pixels. In International conference on machine learning, pages 2555-2566 PMLR, 2019.
[31] Bryan Lim, Stefan Zohren, and Stephen Roberts. Recurrent neural filters: Learning independent bayesian filtering steps for time series prediction. 2020 International Joint Conference on Neural Networks (IJCNN), pages 1-8. IEEE, 2020.
[32] Danijar Hafner, Jurgis Pasukonis, Jimmy Ba, and Timothy Lillicrap. Mastering diverse control tasks through world models. Nature, pages 1-7, 2025.
[33] Gaspard Lambrechts, Adrien Bolland, and Damien Ernst. Recurrent networks, hidden states and belief's in partially observable environments. arXiv preprint urXiv:2208.035500 2022.
[34] Jay A Hennig, Sandra A Romero Pinto, Takahiro Yamaguchi, Scott W Linderman, Naoshige Uchida, and Samuel J Gershman. Emergence of belief-like representations through reinforce- ment learning, PLOS Computational Biology, 191 2023.
12

---

[35] Evan Z Liu, Aditi Raghunathan, Percy Liang, and Chelsea Finn. Decoupling exploration and exploitation for meta-reinforcement learning without sacrifices. In International conference on machine learning, pages 6925-6933 PMLR, 2021.
[36] Anmo J Kim, Jamie K Fitzgerald, and Gaby Maimon. Cellular evidence for efference copy in drosophila visuomotor processing. Nature neuroscience, (((9):1247-1255 2015.
[37] Daniela Schiller, Howard Eichenbaum, Elizabeth A Buffalo, Lila Davachi, David J Foster, Stefan Leutgeb, and Charan Ranganath. Memory and space: towards understanding of the cognitive map. Journal of Neuroscience. 35(41):13904-139) 2015.
[38] James CR Whittington, Timothy H Muller, Shirley Mark, Guifen Chen, Caswell Barry, Neil Burgess, and Timothy EJ Behrens. The eichn-eec machine: unifying space and relational memory through generalization in the hippocampal formation. Cell, 183(5):1249- 1263, 2020.
[39] Ching Fang and Kimberly L Stachenfeld. Predictive auxiliary objectives in deep rl mimic learning in the brain. arXiv preprint
2023.
[40] David Ha and Jürgen Schmidhuber. Recurrent world models facilitate policy evolution. Ad- vances in neural information processing systems, 31, 2018.
[41] Dongqi Han, Kenji Doya, and Jun Tani. Variational recurrent models for solving partially observable control tasks. arXiv preprint arXi::
2019.
[42] Karol Gregor, Danilo Jimenez Rezende, Frederic Besse, Yan Wu, Hamza Merzic, and Aaron van den Oord. Shaping belief states with generative environment models for rl. Advances in Neural Information Processing Systems, 32, 2019.
[43] Alex X Lee, Anusha Nagabandi, Pieter Abbeel, and Sergey Levine. Stochastic latent actor-critic: Deep reinforcement learning with a latent variable model. Advances Neural Information Processing Systems, 33:741-752, 2020.
[44] Jan Humplik, Alexandre Galashov, Leonard Hasenclever, Pedro A Ortega, Yee whye Teh, and Nicolas Heess. Meta reinforcement learning as task inference. arXiv preprint arXiv: 1999000004444 2019.
[45] Kate Rakelly, Aurick Zhou, Chelsea Finn, Sergey Levine, and Deirdre Quillen. Efficient off-policy meta-reinforceemnt learning via probabilistic context variables. In International conference on machine learning, pages 5331-5340. PMLR, 2019.
[46] Xiaoyu Chen, Yao Mark Mu, Ping Luo,, Shengbo Li, and Jianyu Chen. Flow-based recurrent belief state learning for pomdps. In International Conference on Machine Learning, pages 3444-3468. PMLR, 2022.
[47] Andrew Wang, Andrew Cii, Toryn Klassen, Rodrigo Toro Icarte, and Sheila A Mcllraith. Learning belief representations for partially observable deep rl. In International Conference on Machine Learning, pages 35970-35988. PMIR, 2023.
[48] Marvin Zhang, Sharad Vikram, Laura Smith, Pieter Abbeel, Matthew Johnson, and Sergey Levine. Solar: Deep structured representations for model-based reinforcement learning. In International conference machine learning, pages 7444-7453, PMLR, 2019.
[49] Diederik Kingma, Auto-encoding variational bayes. arXiv preprint arXiv: 2013.
[50] Volodymyr Mnih. Asynchronous methods for deep reinforcement learning. arXiv preprint 16Xiv:1602.01783, 2016.
[51] Antoine Girard and George J Pappas. Approximate bisimulations for nonlinear dynamical systems. In Proceedings of the 44th IEEE Conference on Decision and Control, pages 684-689. IEEE, 2005.
[52] John C Gittins. Bandit processes and dynamic allocation indices. Journal of the Royal Statistical Society Series B: Statistical Methodology, 41(2): 1979.
[53] Ilya Kostrikoy. Pytorch implementatiin of reinforcement learning algorithms. https:// github. httaa//o.....ottttooo///ttooooee http::/ppooo..a-aa---------aaa-- aooor-ckr--ae gail 2018.
13

---

A Technical Appendices and Supplementary Material
A.1 Evidence lower bound (ELBO)
At a given time step 1, the learning objective of the RNN encoder dg is to maximize
p [log pe 01+1,7++111 [corr]]
where p(To:t) is the distribution of trajectories generated by the policy Tg. Given that Eq. intractable, we can instead optimize a tractable evidence lower bound (ELBO), computed as:
1 is
C = E ((ro)) B ao(me To:s)
[og oz) log
- DK. (mxl
The terms Eg Loolog (o az) log (Tt ac)] are the future predictive loss, and the term DKLl93 Ig (m, To:a) [po [o [m]]] is regularization term using Kulbbac--eibber (KL) divergence between
the variational posterior 4s (m, [Toz) and the prior over the latent variables pe (m,). ) To encourage the variational autoencoder(M)) (Vanc) to approximate Bayesian filtering for belief update, the prior is set to the previous posterior go (omp-1/Tu.e-11 with the initial prior go (omo/Tu = (00 I). To optimize
the ELBO (Eq. 2), the expectation is approximated with Monte Carlo sampling as in standard VAE
training 49
A.2 Task details
Two-armed Bernoulli bandit We consider a two-armed Beta-Bernoully bandit task with episode length 40 time steps. At the beginning of each episode, 01 and 02. the reward probability biases for the two arms, are independently drawn from fixed Beta distribution, p(0.) Beta(1, (1, 1). l). 0, and 02 are then used to define the Bernoulli reward distributions for arm a E {1, 2]. respectively. The reward sampled from
biases 6., are hidden from the agent. At each time step, the agent chooses arm its policy, and receives a binary reward sampled from Ber(O,), i... T, ~ p(r(0a) = Ber(fg). The discounted cumulative return is computed with discount factor y 0.95. For multi- aarted bandit tasks, Bayes-optimal policy can be derived using the Gittins index method [52), and the minimally sufficient Bayes-optimal belief state is to keep track of the count of total pulls and the count of rewarded pulls each arm. In other words, for the two-armed bandit task, the Bayes-optimal belief states are 4-dimensional: (May Mr Peg Nr.2) where n., and ny,, denote the count of total pulls and the count of rewarded pulls for arm a, respectively.
Dynamic two-armed bandit To evaluate the meta-learned representation and behavior when the underlying environmental state can change over time, we consider a dynamic version of the two- armed bandit. At given time step, each arm is in of possible discrete states. For simplicity choose n = 2. We denote the hidden state of arm a at time as st where sf E E0 [0),00). Va E {1, 2}. with each state associated with a different reward probability, i.e., re -(r(s(]) state of each arm evolves according to an independent Markovian transition process independent of the action chosen: p( s8, a') = p( 1
Ber(sf). The
|sf)) = T ss..... The episode length is set to be 300 time
steps. To examine different structural and dynamical configurations, consider three parameter settings:
Symmetric reward and transition: the two arms share the same of discrete states where $e E (0.11, 0.9]. Va E (1, and the same transition dynamics where 1... = -T..90.9 0.9 and 7 = T0.9.01 = 0.1, Ya E [1,2) That is, for each arm, state 1 is highly rewarding whereas state 2 is less rewarding, and the state of each arm stays with probability of 0.9 and switches with a probability of 0.1.
Asymmetric reward: two arms share the same transition dynamics, but the reward probability states of the two arms are different. Specifically, the reward probability states for arm aj
14

---

are s E (0.1,0.9), and for arm 22 are s E {0.4, 0.6), respectively. Therefore the two states of arm 0n are more differentiating than those of arm ag. The transition dynamics
0..9 T 0.4.0.4 = T 0.6.0.6 = 0.9, and T 0.1.0.9 = 0.9,0.1 = To 0.4.0.6 = .o.. = 0.1. In other words, both arms follow the same dynamics that stay with a probability of 0.9 and switch with a probability of 0.1.
is governed by To 0.1.0.1
Asymmetric transition: two arms share the same reward probability states of E 0..1, Vo E [1,2), but the transition dynamies are different. For arm a:: T 0.1.0.1
T.9.0.9 = 0.9 and To. 0.1.0. = To..o.1 = 0.1; for arm 02, 1...0.1 = To....9 = T6. 7.1,0.9 = To..0.1 = 0.5. That is, whereas the transition dynamics for arm a1 is the same as previ- scenarios, the transition dynamics for arm dg is random with equal stay and switch probabilities of 0.5.
With this task design, the belief state updates in the dynamic two-armed bandit POMDP tasks are analytically tractable by computing the posterior probability of each arm being in state 1 conditioned on the history: by = (p(s) (he))). = 03 h.)). This belief update is tractable using standard Bayesian inference. In turn, the Bayes-optimal policy can be derived using the value iteration algorithm [ by discretizing the 2-dimensional belief state space. The discounted cumulative return is computed with a discount factor = 0.95.
Stationary Tiger To exemplify sequential decision-making under uncertainty, we consider the classic POMDP Tiger task (3) In the tiger environment, an agent chooses between two doors-one hiding a tiger (penalty=-100) and the other hiding a treasure (reward=10). The agent may additionally choose to pay small penalty=-1 for the "listen" action to acquire noisy observations about the tiger's location, The reliability of the noisy observation is controlled by the parameter of observation accuracy. To succeed, the agent must maintain belief about the tiger's location to decide which door to open. This simple yet powerful paradigm tests an agent's 's ability to balance information gathering with reward-seeking making an ideal benchmark for evaluating RL algorithms in POMDP. We consider two difficulty levels by varying the observation accuracy for the "listen" action. Specifically, we consider a simpler task variant where the observation accuracy is 0.8 and harder task variant where the observation accuracy is 0.7. For the Tiger tasks, the Bayes-optimal agent tracks the belief of the tiger's location, for instance, by computing the posterior probability of the tiger being on the left given the history, by = p(tiger at the ltft using standard Bayesian inference. The Bayes-optimal policy can be derived using value iteration (1 by discretizing the I-dimensional belief state space. The discounted cumulative return is computed with a discount factor y 0.95. Typically, as the observation accuracy decreases, the Sayes-optimal solution will require listening for more times as observations are noisier before the belief states crosses the decision threshold for the optimal agent to decide on which door to open.
Dynamic Tiger
We extend the classic Tiger task to dynamic version by allowing the tiger's s location to change over time, following Markov transition dynamics. We denote the tiger location s to be in one of the 2 states E (L, R), and the tiger location evolves according to an independent
Markovian transition process independent of the action chosen: p(8++1 182, a) = p(s++1
T T5. Specifically, we choose Tı, = Tn.n = 0.9 and Ti.R = Tn. = 0.1. other words, at each time step, the tiger stays with a probability of 0.9 and switches its location with probability of 0.1. As the information reliability and relevance are further corrupted by the dynamic nature of the hidden state, the task becomes more challenging because the agent has to balance listening more to increase its confidence against making decisions earlier in case the information gathered so far becomes obsolete. This task design permits tractable belief updates by tracking the posterior probability of the tiger being on the left given the history, b, = p(tiger at the left(h)), using standard Bayesian inference incorporating the Markovian transition matrix. The Bayes- policy can be derived using value iteration [ by discretizing the 1 1-dimensional belief state space. The discounted cumulative return is computed with discount factor 7 0.95. Similarly to the stationary Tiger tasks, we consider two difficulty levels by varing the observation accuracy to be 0.8 or 0.7.
Oracle bandit task The oracle bandit task is designed to exemplify an environment where a successful policy requires paying an immediate exploration cost and utilize the information acquired to improve long-term return. In an 11 11arm bandit environment with an episode length of time steps, one of the first ten arms a 61-10 10 is selected uniformly randomly the target arm a" which will give
15

---

payout of 5 upon choosing. The other nine arms out of 41-10 are non-target arms, each of which will give payout of upon choosing. The last arm, 011, the oracle arm whose payout informs the index of the target arm in the form of 1/ 110 of the target arm index a", i.e. r(an) 0.1 a" (e.g. reward of 0.3 from d11 indicates that d3 the target arm). The average reward from the oracle arm is 0.55 which is smaller than the payout from the non-target arms. As result, choosing the oracle arm is less favorable in terms of the immediate reward but provides useful information if an agent knows how to utilize to improve long-term return. The discounted cumulative return is computed with a discount factor y = 0.95. Note that this setting is similar to a task considered in Wang et al. (77] but differs in that in their setup the oracle information was provided using structured one-hot encoding format, but in our formulation no other feedback than the reward itself is given, which makes learning good representation of the history more challenging and critical. With knowledge of the task structure, the Bayes-optimal belief state is to track and update the posterior probability of each of the first ten arms being the true target by = (p(a* = 01 (14 p(a" = 02 ... p(a" = (e))))))) )-- The Bayes-optimal policy is to pull the oracle arm at the beginning and continue pulling the target arm informed by the reward information from the oracle arm until the end of the episode.
conditioned the history, i.e.
Latent goal cart Deriving Bayes-optimal solutions in continuous POMDPs are usually intractable, hindering rigorous representational equivalence analysis. To enable evaluation continuous observa- tion and action spaces, we design an exemplar continuous control task which still permits tractable Baycs-optimal belief inference and policy. In this task, an agent controls a continuous action, the velocity of a cart, to move along 1-dimensional track to a hidden goal (+1 or 1) which needs to be inferred from the continuous observation (current position) and reward (negative noisified distance from the hidden goal, with the noise following Gaussian distribution) it receives. We consider an episode length of 30 time steps. This task design allows for tractable belief updates by tracking the posterior probability of the hidden goal being +1 given the history, b, p(goal at + 1[hx), using standard Bayesian inference incorporating the knowledge that the noise in reward follows a Gaussian distribution. The Bayes-optimal policy be derived using value iteration [I] by discretizing the 1-dimensional belief state space. The discounted cumulative return is computed with a discount factor y 0.95. This task is motivated by the Calf-Cheetah-De task in MuJoCo, and can be seen simplified version to allow for tractable belief updates and value iteration to derive the Bayes-optimal solution.
A.3 Agent details
2L The implementation of the baseline black-box memory-based meta-RL models, RL3, follows an actor-critic architecture as in previous literature 6 77. Here we use RNNs that function as memory module, consisting of 256 hidden units and hyperbolic tangent activation functions. As shown in Fig. B the input includes the current observation ot, the previous action in one-hot format a; and the associated reward in scalar format rz- The recurrent layer followed by a fully connected bottleneck layer designed to be the counterpart to the latent belief layer b, in Fig D.. After the bottleneck layer is a hidden layer of 32 units and a readout layer that generates a vector of logits for each action 0e for the actor (or vector that defines the mean and the standard deviation of continuous Gaussian policy if solving a continuous task such as the Latent goal cart) and a scalar value baseline V2 for the critic. Actions the action logits. The network is trained end-to-end to maximize the discounted cumulative reward with the Advantage Actor Critic algorithm [50) (Parts of the implementation code are based on [53] under MIT license). The gradient of the objective function is given by:
then sampled from the softmax distribution defined by
A VLn + VLv + V Wentropy
i log π (aol v))
@V
0H
(3)
5, (T(TT: + &. Duly
+ 3.
Oy
Obb
where
(e(Tai tbv) = Ru V(Tai
k-1
R₂ = Dh' + +
k=0
defines the n-step temporal difference error advantage function & the discounted n-step bootstrapped return R2 with discount factor 7, k the number of remaining time steps in the current episode,
16

---

Table 2: Hyperparameters
episode length
8.
Sy
n updates
bottleneck size
Two-armes an--armed bandit
40
0.01 - 0.05
0.01 - 0.05
3e5
8
Dynamic handit
300
0.01
0.05
1e5
8
Stationary Tiger
20 30
0.3
0.1
3e5
Dynamic Tiger
30
40
0.3
0.1
3e5
Oracle bandit
6
0.3 - 0.5
0.01
0.05
3e6
16
Latent goal cart
30
0.005
0.01
0.01
3e5
8
and V the value function parametrized by vy. The neural network policy is denoted as and parametrized by 1b, and H is the entropy of the policy. Finally, By and 8e are hyperparemeters for controlling the relative weighting of value estimation loss and entropy regularization. The choice of hyperparameters for each task summarized in Table. 2 The neural network parameters are trained via backpropagation through time using the Adam Optimizer with a learning rate of 5e-5.
Meta-RL with predictive modules The self-supervised predictive modules are formulated VAE. For the encoder is we use an RNN with 256 hidden units and hyperbolic tangent activation functions. The output of the encoder RNN, i.c. the bottleneck layer, is treated as estimating the mean and the variance of the latents my, as standard in VAE. Therefore, the latent dimension is half of the bottleneck layer size. The decoders Rg and To are multi-layered perceptrons (MLPs) with one hidden layer of 32 units and ReLU activation functions. As shown in Fig. DC the encoder-decoder framework takes as input the current observation Or, the previous action in one-hot format 0;-1, 1. and the associated scalar reward rı. and is set up to make prediction of the upcoming observations o++1 and rewards T++11, conditioned on the trajectory as incurred by the policy network The described below. The entire VAE is trained to maximize the ELBO (Eq. 2 as derived in A. A coefficient of 0.01 is used for the relative contribution of the KL-term for training the VAE. We use the Adam Optimizer with learning rate of 7e-5 to train the VAE using backpropagation through time.
The policy network T& is parametrized as an MLP with one hidden layer of 32 units and hyperbolic tangent activation functions. Similar to the previous paragraph on RL2 the policy network is trained with the Advantage Actor Critic algorithm to optimize the same loss function as described in Eq, 3 An Adam Optimizer with
The choice of hyperparameters for each task is summarized in Table. - learning rate of 5c-5 is used to optimize the policy network. Note although the policy loss depends on the parameters of the encoder 9a, we do not backpropagate the policy loss gradient through the encoder as the goal of the encoder is to learn a belief b, over the latent states predictive of the future such that the belief alone should be a sufficient representation for policy learning (similar to Zintgraf et al. (9). Parts of the implementation code are based on [9] (under MIT license).
Model training The above meta-RL models 2RL2 and ours, meta-RL with predictive modules) are trained on internal GPU clusters (NVIDIA GeForce RTX 4090), which takes less than IG GPU memory and between 10-48 hours for training per model, depending on the task. Compared RL training time of meta-RL with predictive modules increases by 30-40% with the overhead coming from two additional decoders and VAE training, whereas inference time is comparable.
Ablation study To pinpoint which algorithmic design choices drive enhanced representation learning in our proposed meta-RL. with predictive modules, performed two targeted ablations:
(i) no KL: where we remove the VAE's KL regularization to test whether latent regularization
is necessary to enforce compact representations.
(ii) joint RL: where we allow the policy gradients from RL loss to jointly train the RNN encoder alongside the predictive loss-rather than training the encoder solety with the predictive loss-to assess whether policy gradients provide additional representational benefits.
17

---

A.4 State machine simulation
Following the procedure of state machine simulation introduced in Mikulik et al. [3] we consider whether meta- RL system (RL² or our proposed approach, meta-RL with predictive modules) can both simulate and be simulated by, a Bayes-optimal agent for a given POMDP task.
To evaluate how well a state machine M simulates another machine N, a function & is first learned to map the states Sy in N into the state space Sas of M. As enumeration over all possible trajectories are not practical if not possible, quality of a simulation is measured along trajectories sampled from some reference distribution. Given trajectories from reference distribution, quality of the simulation is then measured by (i) the state-trunsition dissimilarity D., measured as the mean-squared error (MSE) between the embedded states ((SN) and the target states SM. and (ii) the output dissimilarity Do. measured as the difference in the expected return generated from the states Sy using the machine N and those generated from the states ((Sx) using the machine M. If both the state-transitte transition and output dissimilarities D. and D, are low/ negligible, then we establish that M simulates .N. If both M simulates N and N simulates M, then we can say M and N are computationally equivalent, and their states Sy and SM are equivalent.
In practice, the mapping function á is implemented as an MLP with ReLU activations and three hidden layers of 64, 128, and 64 units, respectively. The MLP is trained with the Adam Optimizer with learning rate 0.001 and batch size 64. The training set is consisted of 300 1000 trajectories depending on the variability of each task distribution, and the results reported are from another test set of 300 1000 trajectories. Compared with Mikulik et al. [13,, where the reference distribution is generated by the meta-RL agent, here we modify the procedure by generating the reference distribution using the Bayes-optimal agent, as this provides an even more stringent condition where the simulation is evaluated in the regime of Bayes-optimal solutions.
When evaluating how well Bayes-optimal agent simulates meta-RL system, we first train MLP mapping meta-RL states into Bayes-optimal states by minimizing the MSE between the mapped states and the target Bayes-optimal states, After training, quality of the simulation is measured on a test set by evaluating the state-transition dissimilarity (D.: metaRL-Bayes) and the output dissimilarity (D,; metaRL-BByys)) If both D,: metaRL-BBaess and D,: metaRL-+Bayes are low after training, then the Bayes-optimal agent simulates the meta-RL& agent.
On the other hand, to evaluate how well Bayes-optimal agent is simulated by a meta-RL one, an MLP is trained to map Bayes-optimal states into meta-RL states, and the state- transition dissimilarity is denoted D,: layes->metale Dy: Bayes-meetal L. and D.: Bayes Bayes->metal metaRL. are low after training, then the Bayes-optimal agent is simulated by the meta-RL agent.
and the output dissimilarity denoted D, Bayes-JetaRkl If both
If all the above four dissimilarity measures are low/ negligible, then we can say that the meta-RL. agent and the Bayes-optimal agent are computationally equivalent and their representations are equivalent. Our results this paper demonstrate that after training the proposed meta-RL with predictive coding modules can attain much lower dissimilarities than conventional RL2, indicating that meta-RL with predictive modules can more closely approximate the Bayes-optimal belief states.
Note that before training, the state dissimilarity D, can be low for the recurrent layers in all models, RNNs may maintain a verbose
similar to what Mikulik al. [K] reported and representation of the history by embedding each trajectory into a unique hidden state, which can be subsequently mapped to minimally sufficient Bayes-optimal statistic with low error using expressive enough functions, like the MLPs used in the above state machine simulation procedure. This observation also highlights that using decoding alone may not provide a thorough assessment of representation equivalence, and we need to consider their structural and computational relevance when comparing representations.
18

---

A.5
Additional results
A.5.1 State space visualization
To understand the structure of the learned representations meta-RL agents and qualitatively compare those to the Bayes-optimal belief states, visualization of example state spaces for the stationary and dynamic bandit tasks, the stationary and dynamic Tiger tasks, the oracle bandit tasks, and the latent goal cart tasks are presented in Figures. g and respectively.
A) Two-armed Bernoulii bundit Boyes-ootimal
RL2 bottleneck states
Meta-Ri with predictive modules bottlneck states
RL recumeent states
Meta-RL with predictive modules recurient states
Dynamic two-armed bandit, Symmetric
Meria-Hl precictive modules
Meta-IL Rt with predictive modules
KC
C) Dynamic two-armed bandit, asymmetric roward
AL?
Meta RL with procictiv medules bottleneck states
RL recumeent states
Nets: RL with ve modules recurrent states
Bayes-optimal
bottleneck states
1.0
D) Dynamic two-armed bandit, asymmetric transitions
Mota-RL with wow modul butDemeck sleids
R12 recumten states
Meta--I with predictive modules rec recurment Statles
Bayersplinal
bott me Ntales.
014
95
-102
Figure 8: Visualization of state space in the stationary and dynamic two-armed bandit tasks. Example state space for A) Two-armed Bernoulli bandit task, B) Dynamic two-armed bandit task with symmetric reward and transition, C) Dynamic two-armed bandit task with asymmetric reward, and D) Dynamic two-armed bandit task with asymmetric transition. For each panel, the first two principal components are plotted. States are colored by the corresponding policy, i.c. probability of choosing d2. Black curves show one example trajectory. For each task, from the left to the right are the state space of the Bayes-optimal agent (i.e. the belief states), the bottleneck layer in the RL model, the bottleneck layer in the meta-RL RL with predictive modules, the recurrent layer in the RL² model, and the recurrent layer in the meta-Ro -LLL with predictive modules, respectively.
19

---

A) Stationary Tiger, observation accuracy
with modules recumrnt states
oatimal
bettleneck states
bottlerack states
spen
K1 0.7 A2 bottleneck states
B) Stutionary Tiger, observation accuracy
Meta-RL, RL with predictive modules bottleneck states
ALF
Meta-BL with predictive madules recurrent states
Bayes-optiral
recurrent
133
open
20000
200
F
PC
PC1
C) Dynamic Tiger, observation accuracy 0.8 R1 bettleneck
predictive modules bottleneck states
RL?
Mota-RL predictive modules recurrent states
Bayes-optimg
open
1,000
690
0.25
D) Dynumic Tiger, observation accuracy
0.7
Meta-RL predi clive modules
Mata-RL Rll with aradictivn modules recurrett states
Bayes-optimol
states
states
recurrent states
opent
1.34
open
1100
00
22
0.0
0.2
Figure Visualization of state space in the stationary and dynamic Tiger tasks. Example state space for A) Stationary Tiger task with observation accuracy of 0.8, B) Stationary Tiger task with observation accuracy of 0.7, C) Dynamic Tiger task with observation accuracy of 0.8, and D) Dynamic Tiger task with observation accuracy of 0.7. For each panel, the first two principal components are plotted. States are colored by the corresponding policy, with black denoting choosing listen, red choosing to open the left door, and blue choosing to open the right door. Black curves show one example trajectory. For each task, from the left to the right are the state space of the Bayes-optimal agent (i.e. the belief states), the bottleneck layer in the RL2 model, the bottleneck layer in the meta-RL RL with predictive modules, the recurrent layer in the RL2 model, and the recurrent layer in the meta-Re RLL with predictive modules, respectively.
RL bottleneck states
Mato with predictive modules bottleneck states
predict modules recurrent states
Beyes-optimal
recurrent states
04.
671
877
877
877
PC1
pC
PC1
FC1
7C1
Figure 10: Visualization of state space in the oracle bandit task. Example state space for the oracle bandit task. For each panel, the first two principal components are plotted. States are colored by the corresponding policy, with black denoting choosing the oracle arm an and other colors denoting choosing one of the first ten arms a1-10-- Black curves show one example trajectory. From the left to the right are the state space of the Bayes-optimal agent (i.e. the belief states), the bottleneck layer in the RL2 2 model, the bottleneck layer in the meta-RL RL with predictive modules, the recurrent layer in the RL 2 model, and the recurrent layer in the meta-RL with predictive modules, respectively.
20

---

Meto
modules recumentt states
Bayes-oplimal
bottleneck states
boltler States
recument states
6.00
-000
10
Figure 11: Visualization of state space in the latent goal cart task. Example state space for the oracle bandit task. For each panel, the first two principal components are plotted. States are colored by the corresponding policy, i.e the velocity of the cart. Black curves show one example trajectory. From the left to the right are the state space of the bottleneck layer in the RL model, the bottleneck layer in the meta-RL. with predictive modules, the recurrent layer in the RI.2 model, and the recurrent layer in the meta-RL. with predictive modules, respectively.
agent (i.e. the belief states), the
21

---

A.5.2 Full state machine simulation results
Full results of state machine simulation analysis on the trained models of RL and the proposed meta-RL with predictive modules (ours), together with models considered in the ablation studies no KL and joint RL, are summarized in Fig. and Tables and
A) Bottleneck layer
Dei metaRL Bayes
Dj: metaRL Bayes
Fostt Nonralized
tetele noceow
Nom MSE,
107
Dynami
Tiger
Oynomi
Lafont
Oynom
Tight
Dynamec
bands
bandil
GowTart
D.: Bayes
metaRL
D,: Bayes
metaRL
MSE mecos
Norrolized
20
No Nommllize
Dyname
Tieer
Dyname
Dradin
Jabom
Tec-ammd
Dyname
Tiger
Dynonne
Draci
Labont
B) Recurrent layer
Dg: metoRL - Bayes
Dj: metaRL - Bayes
Nornalize
teteee oine-w
Normalizeo MSE
Dynomic
Tigor
Dynomic
LOGONE
Tiger
Libont
bund!
bandit
bunaih
De: Bayes
metaRL
D, Bayes metaRL
Normalized
22
M-tete wereow
Nor Norallced
MSE
Te0-amme
Dynamic
Tioor
Oynomc
Oradie
Te--armad
Dynom
Tiger
Dxnams
Orace
LeGONS
KL
joint RL
Figure 12: Results of state machine simulation. State and output dissimilarities (Do and D.) for both mapping directions (meta-RL-BBayes and Bayes->meta-Rue across all tasks considered, using A) the bottleneck layer and B) the recurrent layer of the fully-trained models of RL. (red), meta-RL with predictive modules (ours, blue), no KL (where KL regularization ablated training, green), and joint RL (where the policy gradient of RL loss is used to jointly train the RNN encoder with the predictive loss, orange). pool across different variants the same task type, all dissimilarity measures of the bottleneck layer in the meta-RL with predictive modules models (denoted as ours in the legend.) In other words, when normalized with the mean bottleneck dissimilarity of our model, values significantly greater than indicate that the representation is worse than the bottleneck layer in our model (i.e. deviating more from the Bayes-optimal belief states). In contrast, values significantly smaller than 1 indicate the representation better than the bottleneck layer our model (i.e. more closely approximating the Bayes-optimal belief states). Error bars denote S. s.e.m. across at least five random seeds per model type.
VAE
normalized by the mean of the corresponding dissimilarities
22

---

Table 3: State machine simulation results for the stationary and dynamic bandit tasks: State and output dissimilarities of the fully-trained models for both mapping directions. M for meta-RL, B for Bayes-optimal solutions. Values indicate meants.t.d. Bold indicates the best dissimilarities among the trained models and better than untrained models (p < 0.05).
Two-armed Bernoulli bandit (unit: 1c-2)
RL
ours
no KL
joint RL
RL
ours
no KL
joint RL
Bottleneck
Recurrent
DM-BB
1.5240.36
1.31
1.4420.4
0.91+0.05
1.65+0.24
1.37=0.38
1.42+0.36
0.9420.15
DM
5,4140.96
1.7740.33
4.2140.38
1.16#0.15
1.1940.23
0.7140.16
1.0440.14
0.58+0.01
DR
1.1920.25
1.1540.31
0.8920.36
0.88+0.12
1.51+0.35
1.1940.34
0.93+0.27
0,95+0.2
D.
5.38+1.81 0.81#0.11
1.02+0.04
0.73+0.04
11,43+1.83 83
21.37=2,12
29.05#1 1.43
20.33+3.96
Dynamic two-armed bandit: symmetric reward and transition (unit: le-2)
RL3
ours
no KL
joint RL
RL
ours
no KL
joint RL
Bottleneck
Recurrent
DM
2.09:3.95
1.17#0.18
0.86=0.48
1.06#0.1
2.18:3.92
1.25+0.21
0.86+0.43
1.08=0. 12
M-
0.19#0.02
0.21#0.05
0.49#0.36
0.22+0.1
0.03+0.01
0.03+0.01
0.03+0.02
0.04#0.01
D
2.06+3.81
0.98+0.36
0.79=0.33 0.89:0.19
2.18:3.75
1.05+0.44
0.72+0.3
0.82+0.3
D.
2.38+0.6
1.09#0.22
0.45+0.13 1.43+0.23
10.26+2.01
4.8620.44
10.2542.37
6.1140.69
Dynamic two-armed bandit: asymmetric reward (unit: 1e-2)
RL3
KL
joint RL
RL
KI.
joint RL
Bottleneck
Recurrent
M-
0.8+0.28
0.86+0.27
0.76#0.19
0.82+0.24
1.01
1.22+0.34
0.59:0.13
1.04:0.34
DM
0.34+0.06
0.93=0.11
0.420.17
1.04#0.16
0.0420.01
0.05=0.0
0.03+0.01
0.04+0.01
D.
1.15+0.27
1.39:0.03
0.99=0.37 1.37+0.16
1.03+0.31
1.37+0.03
0.97=0.34
1.28+0.35
D.
3.49+0.51
1.45+0.1
0.99=0.44
1.98+0.33
10.47+3.6
4.91=0.75
18.2+5..4
5.56+0.6
Dynamic two-armed bandit: asymmetric transition (unit:
le-2)
RL'
ours no KL
joint
RL
ours
no KL.
joint RL
Bottleneck
Recurrent
DM-BB
0.85+0.23
0.93+0.33
1.11=0.47
0.98+0.39
1.22+0.33
0.78+0.29
1.39:0.56
0.95+0.4
DMMDD
0.04+0.02
0.33+0.11
1.1=1.06
0.24+0.09
0.01#0.0
0.01=0.01
0.03+0.01
0.02+0.01
D
1.5220.1
1.700117
1.3620.46
1.6720.12
1.21+0.39
1.6420.19
1.1220.5
1.55+0.16
DR-+M
5.1420.75
2.15+0.45
0.720.39
3.56+1.08
9.0=1.06
7.4140.91
14.3545.85
1.1691.69
23

---

Table 4: State machine simulation results for the stationary and dynamic Tiger tasks.
Stationary Tiger: observation accuracy 0.8 (unit: lc-2)
RL
no KL
joint RL
RL
KL
joint RL
Bottleneck
Recurrent
M 3.33+0.91
2.57=1.01
1.93+1.09
2.97
3.05=0.64
2.07+1.03
2.13
2.6+1.22
D
0.0840.05
0.31
0.46+0.3
0.67=0.87
0,03+0.03
0.06=0.04
0.1140.04
0.23+0.27
DR
3.03+1.54
2.2620.88
2.98+1.7
1.88+1.24
4.15#1.72
1.85+1.22
2.98+1.7
1.34:0.92
D. D
8,0424.38
2.92+1.84
1.1
9.07=10.54
24.95
28.31
35.6128.74
35.95+12.55
Stationary Tiger: observation accuracy 0.7
(unit: 1e-2)
RL
ours
no KL
joint RL
RL
ours
no KL
joint
Bottleneck
Recurrent
D M-+B 3.25+1.13
1.57=0.49
2.68+1.0
1.58=0.74
3.27=1.14
1,2420.62
2.0=0.9
1.02+0.36
DM
0.09#0.05
0.3+0.27
0.53+0.39
0.58+0.44
0.03+0.01
0.13=0.09
0.17=0.1
0.15+00.08
D DB--M
3.06+1.54
1.0920.73
2.25+1.34
1.1320.85
2.86#1.67
1.33=0.76
2.25+1.1
1.13+0.85
DR-MM
4.5140.96
2.97=0.78
5.81+3.41
3.46+1.53
18.77
36.56#1 1.66
42.07=15.13
36.21+14.37
Dynamic Tiger: observation accuracy 0.8
(unit: 1e-2)
RL
ours
no KL
joint RL
RL
ours
no KL
joint RL
Bottleneck
Recurrent
D- 2.19#0.96
2.39#0.41
3.5+1.35
3.48=1.18
2.25=0.95
2.16#0.77
3.1=0.22
3.93+1.03
DMB
0.04+0.01
0.1+0.05
0.31+0.17
0.29#0.06
0.01
0.04=0.04
0.12#0.09
0.02+0.02
DR-MM
2.01#0.93
1.99:0.45
3.48+0.6
3.28+1.44
3.28=0.59
2.34#0.42
3.21#0.76
3.28+1.26
D.
6.27+1.15
1.78#0.98
2.9+1.76
2.03#0.88
29.5+2.07
14.39:4.93
24.2+10.58
11.82+5.04
Dynamic Tiger: observation accuracy 0.7
(unit: le-2)
RL
ours
no KL
joint RL
RL
ours
no KL
joint RL
Bottleneck
Recurrent
D-
2.51+0.66
1.05
27.92+21.09
2.33#1.01
2.3=0.67
2.37=0.53
27.75=21.84
2.51
DM
0.39+0.23
0.28
2.23+1.53
0.2+0.18
0,03=0.02
0.08=0.06
0,74+0.77
0.07+0.06
D
1.7340.32
1884047
32.11422.34
2.0540.57
1.95+0.31
2.27=0.17
30.65420.94
1.8+0.6
D 1.141.89
2.320.72
7.3748.41
4.62+2.85
35.44=12.21
20.65=10.63
21.05+19.62
22.64+12.53
Table 5: State machine simulation results for the oracle bandit tasks.
Oracle bandit (unit: le-2)
RL
ours
no KL
joint RL
RL
ours
no KL
joint RL
Bottleneck
Recurrent
DM
17.2149.73
5.64+4.5
13.7924.44
15.48+1.02
21.96+6.65
5.6424.5
16.26*7.76
15.48+1.02
D QMBB
6.45+5.14
0.08+0.05
24.426.94
0.23+0.14
0.51+1.27
0.14+0.11
0.1420.08
0.3740.68
D.-+M
12.01=10.45
3.26:4.26
59.3428.43
8.57+5.0
22.18+17.23
3.24+4.9
57.34+8.39
7.49+5.26
D.
51.29+6.73
2.33+1.34
2.64#0.82
4.18+3.38
53.7242.63
39.33+5.68
51.5=3.42
36.02+3.26
Table 6: State machine simulation results for the latent goal cart tasks.
Latent goal cart (unit: lc-2)
2
ours
no KL
joint RL
RI
ours
no KL
joint RL
Bottleneck
Recurrent
DMB 33.68+11.84
15.83+3.80
10.73:4.94
21.8648.88
33.37+11.44
15.5624.61
11.20s6.19
21.93+9.05
DM-BB 2.35
0.24*0.07
1.75&1.10
0.30s0.12
0.64+0.14
0.19:0.09
0.25=0.16
0.20+0.08
D.-+M
31.67+13.36
8.01+3.62
9.96+7.73
9.14+4.47
35.42+15.36
7.56+3.07
10.39+7.69
9,08+5.10
DR-+M
30.51
13.16+1.84
3.22+1.09
20.90=10.58
36.02+16.02
32.78+6.28
37.38+5.29
39, 10.10.85
24

---

A.5.3 Sensitivity analysis
To evaluate whether representational equivalence between Bayes-optimal solutions and meta-RL. RL with predictive modules sensitive the choice of bottleneck dimensions, using the oracle bandit task as example, we performed systematic evaluation of state machine simulation while varying the bottleneck dimension ranging from to 32. As shown Fig. [3] when the dimension of the bottleneck is greater than task complexity (i.e. the belief state dimension for the task, e.g. > 2 for the oracle bandit task), the bottleneck of the meta-RI. with predictive modules achieve similarly low dissimilarities, indicating that regularization helps to maintain interpretable belief representation even when the bottleneck is more than expressive enough. In contrast, when the dimension of the bottleneck is smaller than task complexity, both performance (suboptimal return) and representational equivalence (high dissimilarity) significantly degrade, indicating that the bottleneck dimension should be selected to be at least surpassing the task complexity in order to learn the task effectively.
D;: metaRL " Bayes
D;: Bayes metaRL
return difference Normalized MSF
0.2
0.2
0.1
0.1
0
12
16
32
12
16
32
Do: metaRL Bayes
Dg: Bayes metaRL
Normalized
0.3
0.3
12
16
32
12
16
32
Belief dimension
Belief dimension
Figure 13: Sensitivity analysis. State machine simulation results of the bottleneck layers in meta-RL with predictive modules in the Oracle bandit task, as the belief dimension varies from I to 32. When the belief dimension is smaller than task complexity (e.g. <= 2), the models present with high dissimilarities in all four metrics. In contrast, as long as the belief dimension is greater than task complexity, all four dissimilarities remain consistently low even when the belief dimension is large, indicating that regularization is useful for maintaining compact, interpretable belief representation.
25

---

NeurIPS Paper Checklist
1. Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
The main claim in this study is that augmenting meta-RL. agents with predictive coding modules can learn representations that better approximate Bayes-optimal belief states in partially observable environments than conventional meta-RI. models. This claim is examined both qualitatively and quantitatively across several POMDP tasks and presented in the Results section (Section 5) and appendix (A.5.2).
Justification:
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No
NA answer to this question will not be perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.
It is fine to include aspirational goals as motivation as long as is clear that these goals are not attained by the paper.
2. Limitations
Question:
Does the paper discuss the limitations of the work performed by the authors?
Answer: [Yes]
Justification: The scope and limitations this study is discussed in Section 6 As the scope of the study is to investigate whether meta-RL with predictive modules can learn representations with high equivalence Bayes-optimal solutions are derivable. Additional work is needed to further evaluate tasks with larger state space. We also discuss other limitations in Section
Payes-optimal states, we focus on tasks where
Guidelines:
The answer NA means that the paper has limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.
The authors encouraged to create separate "Limitations" section in their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be.
The authors should reflect on the scope of the claims made, e.g., if the approach was only tested few datasets with few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution is low images are taken in low lighting. Or speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon.
The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness.
While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, worse outcome might be that reviewers discover limitations that aren't acknowledged in the paper. The authors should use their best
26

---

judgment and recognize that individual actions in favor of transparency play an impor- tant role in developing norms that preserve the integrity of the community, Reviewers will be specifically instructed to not penalize honesty concerning limitations.
3. Theory assumptions and proofs
Question: For each theoretical result, does the paper provide the full set of assumptions and a complete (and correct) proof?
Answer: [Yes]
The main theoretical result involved the paper is the derivation of the evidence lower bound used for optimizing the predictive variational autoencoder, which is presented in Appendix A.I
Justification:
Guidelines:
The answer NA means that the paper does not include theoretical results.
All the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated or referenced in the statement of any theorems.
The proofs can either appear in the main paper the supplemental material, but if they appear in the supplemental material, the authors are encouraged to provide short proof sketch provide intuition.
Inversely, any informal proof provided in the core of the paper should be complemented by formal proof's provided in appendix supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
4. Experimental result reproducibility
Question: Does the paper fully disclose all information needed to reproduce the main ex- perimental results of the paper to the extent that it affects the main claims and/or conclusions of the paper (regardless whether the code and data are provided or not)?
Answer: [Yes]
Justification: Models and experiments are introduced in Sections 3 and Implementation details of the tasks, model architectures, and the analysis procedures considered in the study are further presented in Appendix A..2 A.3 and A.4 respectively.
Guidelines:
The answer NA means that the paper does not include experiments.
If the paper includes experiments, No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data provided or not.
If the contribution is dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is a novel architecture, describing the architecture fully might suffice, or if the contribution be necessary to either make it possible for others to replicate the model with the same dataset, or provide access the model. In general. releasing code and data is often one good way to accomplish this, but reproducibility can also be provided via detailed instructions for how replicate the results, access a hosted model (c.g., in the case of large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.
specific model and empirical evaluation, may
While NeurlPS does not require releasing code, the conference does require all submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) If the contribution is primarily a new algorithm, the paper should make it clear how to reproduce that algorithm.
(b) the contribution is primarily a new model architecture, the paper should describe the architecture clearly and fully.
27

---

(c) If the contribution is a new model (e.g.. a large language model), then there should either way to access this model for reproducing the results or way to reproduce the model (e.g., with an open-source dataset or instructions for how to construct the dataset).
(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome to describe the particular way they provide for reproducibility. the case of closed-source models, may be that access to the model is limited in some way (e.g., to registered users), but it should be possible for other researchers have some path to reproducing or verifying the results.
5. Open access to data and code
Question: Does the paper provide open access to the data and code, with sufficient instruc- tions to faithfully reproduce the main experimental results, as described in supplemental material?
Answer: [Yes]
Open access to the code for training the models and running experimental anal- ysis is publicly available at htttp:://github.. http://itthub..oo//v www.ww/wwwlemeeaaRLL.... prtp:aaammaaaaaociiiiieerrrreeee We also provided detailed implementation details in the Appendix A.2 reproduce all the experimental results.
Justification:
AA.2 A.3 and A.4 to
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurlPS code and data submission guidelines (kttps://nips.cem cc/
for more details.
While we encourage the release of code and data, we understand that this might not be possible, so "No" is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for new open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurlPS code and data submission guidelines (https
//nips /niips.cc/pullle. httpllicgggidee/Cooeeee
for more details.
The authors should provide instructions on data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. only subset of experiments are reproducible, they should state which ones are omitted from the script and why.
At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
Providing as much information as possible in supplemental material (appended the paper) is recommended, but including URLs to data and code is permitted.
6. Experimental setting/details
Question: Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary understand the results?
Answer: [Yes]
Justification: Models and experiments are introduced in Sections and 4 All relevant details of model training, evaluation, and analysis are presented A.2
A.4
Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented in the core the paper to a level of detail that is necessary appreciate the results and make sense of them.
The full details can be provided either with the code, appendix, or as supplemental material.
7. Experiment statistical significance
28

---

Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Answer: [Yes]
Justification: The quantitative results using state machine analysis are presented with error bars indicating the standard errors of the mean across different models with different random seeds. Statistical significance tests comparing the difference between difference models are also performed for the analysis, with results indicated in the figures.
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results are accompanied by error bars, confi- dence intervals, or statistical significance tests, least for the experiments that support the main claims of the paper.
The factors of variability that the error bars capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall
with given experimental conditions).
The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)
The assumptions made should be given (e.g.. Normally distributed errors).
should be clear whether the error bar is the standard deviation the standard error of the mean.
It is OK to report 1-sigma error bars, but should state it. The authors should preferably report a 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality of errors is not verified.
For asymmetric distributions, the authors should be careful to show in tables or figures symmetric error bars that would yield results that are out of range (e.g. negative error rates).
If error bars are reported in tables or plots, The authors should explain the text how they were calculated and reference the corresponding figures tables in the text.
8. Experiments compute resources
For each experiment, does the paper provide sufficient information on the com- puter resources (type of compute workers, memory, time execution) needed to reproduce the experiments?
Question:
Answer: [Yes]
Justification: Information the computer resources is provided in Appendix A.3 including compute worker type, memory, and time training.
Guidelines:
The answer NA means that the paper does include experiments.
The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage.
The paper should provide the amount of compute required for each of the individual experimental runs as well as estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported the paper (e.g., preliminary or failed experiments that didn't 't make into the paper).
9. Code of ethics
Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics httpp/////eeerppp.ccc//ppbtt///a
Answer: [Yes]
Justification:
The study in this paper conforms with the NeurlPS Code of Ethics.
Guidelines:
The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.

---

If the authors answer No, they should explain the special circumstances that require deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g.. there a special consid- eration due to laws or regulations their jurisdiction).
10. Broader impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [Yes]
This study is foundational research and not tied to particular applications. But we have a paragraph discussing how our work is broadly contributing model interpretabil- ity and the connection between neuroscience and Al research.
Justification:
Guidelines:
The answer NA means that there is no societal impact of the work performed.
If the authors answer NA No, they should explain why their work has no societal impact or why the paper does not address societal impact.
Examples of negative socictal impacts include potential malicious unintended uses (c.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.
The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is a direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.
If there negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how system learns from feedback over time, improving the efficiency and accessibility of ML).
11. Safeguards
Does the paper describe safeguards that have been put in place for responsible release of data models that have high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Question:
Answer: [NA]
Justification:
The paper poses no such risks.
Guidelines:
The answer NA means that the paper poses such risks.
Released models that have high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere usage guidelines or restrictions to access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make best faith effort.
12. Licenses for existing assets
30

---

Question: Are the creators or original owners of assets (e.g.. code, data, models), used in the paper. properly credited and are the license and terms of use explicitly mentioned and properly respected?
Answer: [Yes]
Justification: The source codes that parts of our implementation is based on are cited and their license mentioned and respected.
Guidelines:
The answer NA means that the paper does not use existing assets.
The authors should cite the original paper that produced the code package or dataset.
The authors should state which version of the asset used and, possible, include a URL.
The name of the license (e.g., CC-BY 4.0) should be included for each asset. For scraped data from a particular source (e.g.. website), the copyright and terms of service of that source should be provided.
If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperssii thcode. com datasets has curated licenses for some datasets. Their licensing guide can help determine the license dataset.
For existing datasets that are re- re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided. If this information is available online, the authors are encouraged to reach out to the asset's 's creators.
13. New assets
Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [NA]
Justification: The paper does release new assets, but have plans release the source code once the paper is accepted.
Guidelines:
The answer NA means that the paper does release new assets.
Researchers should communicate the details of the dataset/codeemoodll as part their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether how consent was obtained from people whose asset used.
At submission time, remember anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.
14. Crowdsourcing and research with human subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well details about compensation (if any)?
Answer: [NA]
Justification:
The paper does not involve crowdsourcing nor research with human subjects.
Guidelines:
The answer NA means that the paper does involve crowdsourcing nor research with human subjects.
Including this information in the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then as much detail as possible should be included in the main paper.
According to the NeurlPS Code of Ethics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.
31

---

15. Institutional review board (IRB) approvals or equivalent for research with human subjects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or equivalent approval/teview based the requirements of your country institution) were obtained?
Answer: [NA]
Justification: The paper does not involve crowdsourcing nor research with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with
human subjects.
Depending on the country in which research is conducted, IRB approval (or equivalent)
may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere the NeurIPS Code of Ethics and the guidelines for their institution.
For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review.
Declaration of LLM usage
Does the paper describe the usage of LLMs if it is important, original, or non-standard component of the core methods in this research? Note that if the LLM is used only for writing, editing, formatting purposes and does not impact the core methodology, scientific rigorousness, or originality of the research, declaration is not required.
Question:
Answer: [NA]
Justification: LLMs are used only for writing, editing, formatting, and implementing standard analysis. They are not used for the core methodology.
Guidelines:
The answer NA means that the core method development in this research does not involve LLMs as any important, original, non-standard components.
Please refer to our LLM policy hutt::///oorrrsscccooooorrrrre// for what should or should not be described.
32


Reviews:
[{"review_id": "zuqAY37apn", "reviewer": "Reviewer_gFEx", "summary": "The 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.", "questions": "How 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.", "limitations": "Yes", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 2, "originality": 3, "strengths_and_weaknesses": "This 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.", "was_revised": true, "final_justification": "I thank the authors for both exceptionally clear responses, as well as for the inclusion of additional results. I think that these answers improve my understanding, and with the new results will indeed strengthen the paper significantly. For this reason I am upgrading my score from a 3 to a 5.", "rebuttal": "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.\n## Main Concerns\n> Q1: Limited experimental domain\n\nOur 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.\n\n**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.\n\n**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**.\n> Q2. ... compare to … VariBAD, world model approaches? Why … RL² as the sole baseline?\n\nOur 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.\n\nTo 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.\n- 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.\n- 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.\n\n**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).\n\nOur work lays the groundwork for future analysis. We’ll expand our discussion section to clarify these relations and our complementary contribution.\n> Q3. How … scale to continuous or high-dim spaces?\n\n1. **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). \n\nmodel | $D_o^{M->B}$ | $D_s^{M->B}$ | $D_o^{B->M}$ | $D_s^{B->M}$ (bold: p<0.05)\n\nOurs | **0.699±0.270 | 0.017±0.004 | 0.528±0.098 | 0.175±0.049**\n\nRL²  | 1.523±0.518 | 0.066±0.109 | 1.743±0.520 | 0.371±0.230\n\nThis provides evidence that our findings extend beyond discrete, low-dim tasks. \n\n2. **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.\n\n3. **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.\n## Specific questions\n> Q4: You claim that \"no policy-gradient leakage\" is crucial, but the ablation … when/why gradient isolation matters?\n\nWe'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.\n\n**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.\n> Q5: Oracle bandit success … diverse exploration challenges?\n\nOracle 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)**.\n\nOur 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.\n> Q6: … training time compare to RL²?\n\nOverhead 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.\n> Q7: How sensitive is … bottleneck dimension?\n\nWe evaluate different bottleneck dimensions (ranging from 1 to 32). Results show when:\n- 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.\n- 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.\n\nWe will include the results in the revision.\n> Q8: ... OOD generalization … on task distributions that differ from training?\n\nThank 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:\n- **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.\n- **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.\n\nWe’ll include these results in the revision.\n> Q9: state machine simulation ... unfamiliar with Mikulik et al\n\nWe'll expand the description in the main text."}, {"review_id": "UWDYQzejXP", "reviewer": "Reviewer_dd6m", "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.", "questions": "no", "limitations": "no", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 4, "significance": 2, "originality": 3, "strengths_and_weaknesses": "**Strengths**\n1. Clear theoretical bridge and a tight conceptual link to the predictive coding principle in neuroscience.\n2. 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.\n\n**Weaknesses**\n1. 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.\n2. The code is not available at the time of reviewing.", "was_revised": true, "final_justification": "I have read the rebuttal and found it partially addresses my concerns. I recommend accepting this paper.", "rebuttal": "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.\n\n> 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.\n\nWe 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.\n\n1. **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.\n\n2. **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**:\n- Bandits: Uncertainty and exploration-exploitation tradeoff\n- Tiger: Classic POMDP benchmark of full sequential decision-making problem\n- Dynamic bandits/Tiger: Non-stationary hidden states where dynamic belief tracking is required\n- Oracle bandit: Strategic information gathering\n\nTo our knowledge, this is **the first systematic evaluation of meta-RL representations across diverse POMDP structures beyond stationary bandits using ground-truth comparison**.\n\n3. **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.\n\n4. **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.\nOur work provides the **interpretability foundation** for understanding this family of methods:\n- Existing work (VariBAD, Dreamer): Demonstrates *that* predictive objectives improve performance\n- Our work: Explains ***why*** they work by revealing predictive representations approximate Bayes-optimal beliefs\n\nWe 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.\n\n**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.\n\n> Q2: The code is not available at the time of reviewing\n\nWe 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.\n\nUpon acceptance, we will release our complete codebase, including:\n- Full implementation of meta-RL with predictive modules\n- All baseline implementations (RL²)\n- Complete experimental setup\n- State machine simulation analysis tools\n- Scripts to reproduce results\n\nThe 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.\n\n---\nWe 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."}, {"review_id": "yiONCvgLFu", "reviewer": "Reviewer_e9yF", "summary": "I am not sure what this paper is about. The title contains several terms:\n\n * Predictive coding\n * Meta-RL\n * Bayes-optimal belief presentation\n * Partial observability\n\nand even after reading the abstract I still am confused what are the main problems solved in this work.", "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)\n\nWhy 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)?\n\nHow can I confirm that there is any statistical significance in the results in Figure 2 and 3?", "limitations": "Not discussed (no explicit section)", "rating": 4, "confidence": 2, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 2, "clarity": 1, "significance": 2, "originality": 3, "strengths_and_weaknesses": "+ the work addresses the partial observability problem that is largely unsolved in RL literature\n\n- writing style makes it difficult to follow:\n  * What is the problem\n  * Why the problem is important w/ example cases (beyond a few refs)\n  * What have been previously done (state-of-the-art)\n  * Own contributions \n\n- comparing results to a 9-years old paper is not sufficient, especially as it is difficult to see any differences\n\n- motivation of using the selected environments is poor (why these are fundamentally good examples of partially observable and realistic cases)", "was_revised": true, "final_justification": "After engaging in discussion with the authors and after reading the reviews of other reviewers I was finally able to understand the position of this work in the continuum of works from Wang et al. (2016) [idea of meta-RL] to Ortega et al. (2019) [claim of meta-RL leading to Bayes optimal estimator] to Mikulik et al. (2020) [Bayes optimal verified]. The current work shows that predictive coding helps to provide representations that match with Bayes optimal estimator. What are the consequences of that remain to be shown, but the work shows strong incremental progress and therefore I update my grade to the positive side.", "rebuttal": "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.\n\n## What problem are we solving and why is it important?\nIn the real-world, agents (AI systems, humans) often have to make decisions without complete information. For example:\n- A doctor diagnosing a patient can only observe symptoms, not the underlying disease\n- A robot navigating a building can only see with its camera, not through walls\n\nAs 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.\n\n### Why previous methods fail\nMemory-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: \n- Lack of interpretability (we can’t explain the symptoms)\n- Ineffective exploration (without disease mechanisms, we don’t know where to look for potential new treatment)\n- Inadequate generalization (we don’t know whether the symptom combination will transfer to new patients). \n\nThis 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?**\n\n## Our contribution\nMotivated 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:\n- Memorizing that \"fever + cough often means flu\" (what RL² does, a verbose representation)\n- 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)\n\nTo 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**.\n\nWe hope this clarifies our motivations and contributions. We're happy to provide additional clarification. Below we address specific questions.\n\n> Q1: Clarifying the confusion about predictive coding\n\nThis 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:\n- Inspired by neuroscience theories of how brain process information\n- Implemented using standard backpropagation\n- Similar to successful representation learning methods in ML (e.g. future predictive loss, contrastive predictive coding)\n\nWe’ll clarify this distinction in the revision to prevent misunderstanding.\n\n> Q2: Why these environments?\n\nWe 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:\n- Optimal solutions in complex POMDPs are known to be intractable. Without optimal solutions, we cannot directly verify whether the learned representations are optimal\n- In tractable POMDPs, we can *quantitatively* measure how close our learned representations are to Bayes-optimal solutions, not merely comparing performance\n- Our selected tasks capture core, fundamental challenges in POMDPs, including *exploration-exploitation tradeoff, information gathering, and dynamic uncertainty*\n\nTo be specific, our tasks **encompass core partial observability challenges with real-world implications**.\n- Bandits: exploration-exploitation tradeoff → Clinical trial design, A/B testing\n- Tiger: classic POMDP benchmark → Scenarios requiring costly information gathering before high-stakes decisions\n- Dynamic bandits/Tiger: non-stationary tasks requiring dynamic belief tracking → Markets, weather, or any dynamic domains\n- Oracle bandit: active information seeking → Active, exploration-based information gathering\n\nWe 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.\n\n> Q3: Comparison baselines\n\nOur 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.\n\nTo contextualize our model, many modern deep RL methods have converged on incorporating predictive objectives:\n- VariBAD (Zintgraf et al 2021) uses variational inference with future prediction\n- SOLAR (Zhang et al 2019) and Dreamer (Hafner et al 2025) learn world models through observation prediction\n- DynaMITE-RL (Liang et al 2024) employ dynamic predictive model\n\nThese 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:\n- It's *the* canonical meta-RL method, representing a foundational and widely-used baseline.\n- Mikulik et al 2020 specifically examined RL², allowing direct comparison.\n- Recent meta-RL methods build upon RL².\n- This allows us to isolate the effect of predictive learning on representation quality.\n\nBy 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.\n\n> Q4: How can I confirm that there is any statistical significance in the results in Figure 2, 3?\n\nWe 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**.\n\nTherefore, 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).\n\nWe show this consistent pattern across all tasks we considered, which provides strong evidence for better representation learning induced by predictive coding.\n\n> Q5: Why meta-RL is not explained\n\n**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.\n\n> Q6: Limitations not discussed\n\n**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.\n\n---\nWe 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."}]

Author rebuttal:

Meta-review: This manuscript investigates strategies for improving internal state representations in 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 framework. The work establishes a clear theoretical foundation and maintains a strong conceptual link to predictive coding in neuroscience. Furthermore, the evaluation is conducted with methodological rigor, enabling a robust assessment of whether the learned representations approximate Bayes-optimality. As most of the reviewers’ comments have been sufficiently addressed, this paper appears suitable for publication
Decision: Accept (poster)