On Efficient Bayesian Exploration in Model-Based Reinforcement Learning
Authors: Alberto Caron
Vasilios Mavroudis
Chris Hicks
Venue: Accepted by TMLR ()

Abstract: In this work, we address the challenge of data-efficient exploration in reinforcement learning by examining existing principled, information-theoretic approaches to intrinsic motivation. Specifically, we focus on a class of exploration bonuses that targets epistemic uncertainty rather than the aleatoric noise inherent in the environment. We prove that these bonuses naturally signal epistemic information gains and converge to zero once the agent becomes sufficiently certain about the environment’s dynamics and rewards, thereby aligning exploration with genuine knowledge gaps. Our analysis provides formal guarantees for IG-based approaches, which previously lacked theoretical grounding. To enable practical use, we also discuss tractable approximations via sparse variational Gaussian Processes, Deep Kernels and Deep Ensemble models. We then outline a general framework — Predictive Trajectory Sampling with Bayesian Exploration (PTS-BE) — which integrates model-based planning with information-theoretic bonuses to achieve sample-efficient deep exploration. We empirically demonstrate that PTS-BE substantially outperforms other baselines across a variety of environments characterized by sparse rewards and/or purely exploratory tasks.

Published in Transactions on Machine Learning Research (07/2025)
On Efficient Bayesian Exploration in Model-Based Reinforce- ment Learning
Alberto Caron The Alan Turing Institute London, UK
acurongturing.cco
Chris Hicks
c.hicks Batiikngttring.cc.
The Alan Turing Institute London, UK
Vasilios Mavroudis
vmavroudis Stwrroodd@@ttuiing....
The Alan Turing Institute London, UK
Reviewed on OpenReview: https: openreview. net/ forum? Na02hDWgkF
Abstract
In this work, we address the challenge of data-efficient exploration in reinforcement learning by examining existing principled, information-theoreti approaches to intrinsic motivation. Specifically, we focus on a class of exploration bonuses that targets epistemic uncertainty rather than the alcatoric noise inherent in the environment. We prove that these bonuses naturally signal epistemic information gains and converge to zero once the agent becomes sufficiently certain about the environment's dynamics and rewards, thereby aligning explo- ration with genuine knowledge gaps. Our analysis provides formal guaranteos for IG-based approaches, which previously lacked theoretical grounding. To enable practical use, we also discuss tractable approximations via sparse variational Gaussian Processes, Deep Kernels and Deep Ensemble models. We then outline general framework Predictive Trajectory Sampling with Bayesian Exploration (PTS-BE) - which integrates model-basod planning with information-tteorelia. bonuses to achieve sample-efficien. deep exploration. We empiri- cally demonstrate that PTS-BE substantially outperforms other baselines across variety of environments characterized by sparse rewards and/or purely exploratory tasks.
1 Introduction
The exploration-exploiatiio trade-off is a long-standing challenge in Reinforcement Learning (RL) (Sutton & Barto, 2018). Exploration in classic RL algorithms is often achieved via simple heuristics such as e-gredy policy in Q-methods (Mnih et al., 2015), action noise injection (Lillicrap et al., 2015), or some form of policy entropy regularizers in policy gradient methods (Sutton et al., 1999; Kakade, 2001: Schulman et al., 2017). some environments, these simple heuristics are enough ensure sufficient exploration under uncertainty such that the optimal policy can eventually be Icarned; in others, characterized by more 'sparse' rewards and noisy transitions, they are prone get the agent stuck in sub-optimal policy regions instead.
Different methods that satisfy the need for deeper exploration in sparse rewards environments have been proposed, c.g., Thompson sampling adaptations (Osband et al., 2016; 2018) and Q-network parameter uncertainly (Axizzadenneeheli el al., 2018). Perhaps the most popular class of methods is the one based on intrinsic motivation (Schmidhuber, 1991; Chentanez et al., 2004), or intrinsic rewards. Intrinsic motivation mothods introduce an additional reward signal that encourages the agent to explore and learn about aspects of the environment beyond what the task-specific rewards alone can incentivize (Ladosz al., 2022). This means that when the agent encounters state, they assign it higher internal, or 'intrinsic', exploration

---

Published in Transactions on Machine Learning Research (07/)025)
bonus ri that encourages them to visit that state more often. Several different types of exploration bonuses have been proposed, including model-based (Stadie et al., 2015; Osband et al., 2016; Pathak et al., 2017), where the agent learns a model of the environment dynamics and uses an output from it (prediction error, variance, etc.) to define ri, and count-based based (Bellemare et al., 2016; Ostrovski et al., 2017; Tang et al., 2017), where the intrinsic reward is defined as a function of the visitation frequency of a given stale-action pair.
However, some intrinsic reward methods may inadvertently conflate epistemic uncertainty (lack of knowledge about the environment due to limited data) with aleatoric uncertainty (inherent stochasticity that cannot be reduced through data collection) (Hiillemeeirr & Waegeman, 2021). For instance, in highly noisy environments, exploration methods based on prediction error may keep rewarding the agent for visiting already well-understood states whenever they produce random transitions or outcomes. As a consequence, the agent may continue exploring those regions indefinitely, rather than shifting its exploration budget to truly under-explored regions, resulting in high sample inefficiency. This challenge has led researchers to begin to consider alternative intrinsic rewards specifications that attempt to "de-noise" the transitions (Pathak et al., 2017; Jarretl al., 2023), looking for underlying novelty signals. The overall aim of this paper is to provide theoretically grounded insights for constructing good exploration measures that are naturally capable of dissecting alcatoric and epistemic components.
Related Work Model-based intrinsic rewards (Bellemare et al., 2016; Osband et al., 2016; Pathak et al., 2017) are a popular class of methods that address sparse rewards by leveraging exploration bonuses derived from learned model of the dynamics. Examples include next-state prediction error (Stadie et al., 2015; intrinsic
Pathak et al., 2017) and variance (Sorg et al., 2010; Pathak et al., 2019). Information-theoree rewards, rooted in Bayesian Optimal Experimental Design (BOED) (Lindley, 1956; Pukelsheim, 2006; Foster et al., 2019; Rainforth et al., 2024) and Bayesian Active Learning (BAL) (Houlsby et al., 2011; Hanneke et al., 2014; Gal et al., 2017) principles, have been previously proposed. Notably, Houthooft et al. (2016) introduced Variational Information Maximizing Exploration (VIME), defining intrinsic rewards using Bayesian surprise (Itti & Baldi, 2009), quantified as Information Gain (IG) via relative entropy between posterior and prior parameters. VIME specifically computes one-step IG bonus using a variational Bayesian neural network and incorporates this within a model-free policy-gradient learner. More recently, few contributions proposed employing IG within a model-based planning framework (mostly via the use of deep ensembles specifically as a forward dynamics model) that explicitly considers longer planning horizons (Shyam el al., 2019; Sekar el al., 2020). However, none of these prior contributions provide a formal theoretical analysis of the exploration bonuses proposod. Tangentially related is the litcrature on Modcl-Basod RL (MBRL) frameworks that have frequently leveraged uncertainty quantification for exploration purposes, such as Gaussian Process-based planners (Deisenroth & Rasmussen, 2011; Kamthe & Deisenroth, 2018; Mehta et al., 2022). Yet, many of these works either focus solely on dynamics uncertainty or do not explicitly address sparse reward environments or purely exploratory tasks. While the idea of integrating information gain bonuses within MBRL planners is present in prior work, formal analyses characterizing their theoretical properties - such as convergence and contraction rates- remain largely unexplored. Finally, our work also connects closely to the seminal literature on Bayesian RL (Strens, 2000; Duff, 2002; Ghavamzadel et al., 2015), which inherently balance exploration and exploitation through Bayesian uncertainty quantification.
exploiten
Contributions In this work, we study a class of informattion-theorttee exploration bonuses inspired by Bayesian surprise and rigorously analyze their theoretical and practical properties. Although employing information gain itself as an exploration bonus is not new, our contributions are as follows: i) we provide a rigorous theoretical foundation for information-theoretic exploration bonuses, proving key convergence and contraction properties, which nicely complement prior works such VIME (Houthooft et al., 2016) and other existing "planning-te to-eploort approaches (Shyam et al., 2019; Sekar et al., 2020), which lack such formal analysis; ii) we advocate for the use of model-based planning approaches for settings that require high data efficiency, and design a general Predictive Trajectory Sampling with Bayesian Exploration (PTS-BE) framework, which simply generalizes (GPs, deep ensembles, etc.), and explicitly model epistemic uncertainty both in next state and rewards predictions; iii) we provide an empirical comparison of different tractable Bayesian approximators for posterior inference - including deep ensembles, Gaussian Processes, and Deep Kernels - and analyze their impact
twwniinngtt--eppllo methods by supporting arbitrary Bayesian models
2

---

Published in Transactions on Machine Learning Research (07/2025)
information gain estimation and exploration efficiency; iv) empirically demonstrate the theoretical properties of these information-theoretie bonuses and thoroughly investigate their practical benefits in terms of sample efficiency when combined with standard volicy-gradient algorithms (e.g., Proximal Policy Optimization, Soft Actor-Crittt Critic)) particularly emphasizing their effectiveness in sparse-reward environments and exploratory tasks that inherently demand careful uncertainly quantification and strategic exploration. Overall, this work is intended to provide novel theoretical insights into the properties of Bayesian information-base intrinsic rewards and introduce a generalized and theoretically justified MBRL planning-to-explors framework.
2 Problem Setup
We begin by defining some concepts and notation that will be used throughout the paper. Firstly, consider the canonical definition of a Markov Decision Processos (MDP) (Puterman, 2014; Sutton & Barto, 2018). A MDP is a tuple (S, A, Pr. 7); defined over an horizon of & E {1...., T} time steps where: i) S is the state space; ii) A is the action space; iii) p+(-ss, a) E P, is the state transition probability that determines the next state 81+1 given the current state and action pair (st, Qz); iv) po Po is the initial state probability, 8o Pa v) : x A C C IR is bounded extrinsic reward function output from the environment; vi) y E (0, 1) is a discount factor.
A policy is a function π. : S - A that defines an agent's behaviour by mapping a certain state 8 E S to an action a E A, and can be either deterministic or stochastic. Given a certain policy x = T ((alss)) and transition probability p+ = p(s''ss a), the action-value function Q (8, a) can be defined via the one-step a) + ar((s))++ [VT (s'),, where V=-pr (s) is the state-value function instead: V (s) = (1s))) (s, + Ep. [s)]]- Lastly, another useful quantity is the discounted returns defined as Pr); = Ei y "' [r" (s), An optimal policy x² is equivalently defined as the maximizer of QT-Pr (s, Va. (s) and (T: for all s S.
Bellman equation, QT =
2.1 Intrinsic Curiosity and Exploration Bonuses
Intrinsic curiosity, or intrinsic motivation, (Schmidhuber, 1991; Barto, 2013) in RL. is technique that involves augmenting the environment's extrinsic reward signal ("t,,, az) with some notion of intrinsic reward ri(s;, a.) E R, such that the total reward signal at time from pair (s,, a,) reccived by the agent equals r(s,, 0z) = r (s,as) a:)+ at). where Nı is a scaling factor that can be decayed as a function of the time step t. The intrinsic reward ri = r"( az) constitutes a form of exploration bonus that the agent gets to enhance visitation of unseen regions of the state-action space Sx... This intrinsic reward augmentation technique, which generates the cumulative sum of discounted returns J(: ? "' [s" + +1 (se, a:)]. can also be viewed as the regularized solution to the constrained optimization problem (Altman, 2021; Gattami et al., 2021) defined as follows:
T Σn Y EaiN (st, s.t. Σ 1=U
T
max J(T:pt) #
n 'anNn,s' (r'( 2n,
(
where 17 2 0 is some desired level of cumulative exploration. Intrinsic rewards, or exploration bonuses, have been defined in a variety of different ways in the literature. Some popular methods include, but are not limited to, intrinsic rewards based on:
i) Prediction Error (Schmidhuber,, 1991; Stadie et al., 2015; Pathak el al., 2017), where rf is defined ri Fllfe(st, az) st+1lpp II-IIp denotes the L space norm and fa(.-) is predictive model for the environment dynamics, i.e. Elp((c+1l8e at)] = Se(s,, at) (e.g, neural network).
ii) Variance (Hester & Stone, 2012), where rf can be defined as r E [[ at) - (s) and so(:,)) again is a model of the environment dynamics.
iii) Visitation Frequency of state-action pairs (st, (z) (Bellemare et al., 2016; Tang et al., 2017; Ostrovski et al., 2017), where rf is defined as ri = ns; at"-, and N(s,, at) is the count of encountered (s, a) pairs. Pseudo-counts typically used with high-dimensional 8 A.
3

---

Published in Transactions on Machine Learning Research (07/2025)
iv) Empowerment (Klyubin al., 2005; Salge et al., 2014), where = the conditional mutual information quantifying the amount of control the agent can exert via a, in a specific state St s. Empowerment based intrinsic rewards rf nudge the agents towards more controllable states.
s). Here, I[,, +.) is (1 w
A drawback of some intrinsic reward approaches, especially when transitions are noisy, is that they are not guaranteed to diminish as the agent becomes more familiar with the environment, since they conflate epistemic uncertainty with alcatoric uncertainty. This persistence can lead to excessive exploration of already well-understood state-action pairs, hindering the convergence to high reward regions (Chentanez et al., 2004; Burda et al., 2018). Furthermore, when intrinsic rewards are incorporated into the learning process, they effectively modify the original MDP by altering the reward function. This implies that the resulting value and discounted return functions QT.P. (s, a), V (s) and [+) are effectively 'biased' with respect to the main task described by the original MDP. If the intrinsic reward modification is persistent, it may bias trained policies away from optimal solutions. Therefore, it is essential to design intrinsic rewards that facilitate efficient exploration early stages of training, by disentangling epistemic and aleatoric uncertainties, and al. the same time naturally fade to zero once the local environment dynamics is well understood.
2.2 Eayes-Adaptive MDPs and Exploration as Complementary Task
A grounded way of representing the agent's complementary task of exploration, such that it is consistent. with an MDP's state values functions, can be devised by adopting Bayesian perspective on the uncertainty stemming from the environment dynamics. To this end, let us define the true parameters underlying the transition distribution of the dynamics as Oreee E O, where we equivalently use p(*++1, St/St, St, ae; Perae) = Potrue ($.41, Telse, az). The typical Bayesian approach entails placing some form of prior distribution p(0) on the unknown Otrwes define a likelihood on the collected data p(D))0), where n is the sample size, and updale beliefs on it after collecting new data D through Bayes rule, p(0 x p(D"10)p(()) With this in mind, we can formulate modified version of the classic MDP, Bayes- Adaptive Markov Decision Process (BAMDP) (Duff, 2002; Ross et al., 2007), which explicitly incorporates belief updates over 0 into the agent's decision-makie making process and enables one to define intrinsic rewards as function of the agent's 's current belief over 0:
Definition 2.1 (BAMDP). A BAMDP is a tuple (Hs, A. Ph. Po, 1h 7) defined over an horizon of e {1,..., time steps where:
i) H. S θ is the set of hyper-states, defined the Cartesian product of the environment states space
S and the space of parameters of the posterior transition dynamics θ
ii) A is the action space;
iil) pn(-iha,;0n at) at) E P. is the hyper-state transition probability that determines the next hyper-state (s1+1, 0') given the current hyper-state (s,, 0) and action oe, and is such that ph, (8141, 0's,,, axj 0) p(s441 (st, ac; ()p((()(s. 0): The probability p(0' a, () is then the updated belief poste- rior (D))) on Oerue E O, given the revealed new state 8441 that completes the full transition (s,, at, 8c+1) = (i.e., one new data point);
iv) po E Po is the joint distribution of the initial state and the prior on the dynamics parameters, which can be assumed to be independent, i.e., (o 8) = p(s))p((), where so ~ p(s));
v) Th Hs x A- C C R is a bounded combined reward function that stems from the hyper-state Hs Sxe and action A spaces and include both extrinsic rewards re, strictly generating from 8 x A, and intrinsic rewards r², generating from O as discussed below; vi) ? 10 is discount factor.
Our definition of a BAMDP differs slightly from that typically found in the literature on Bayesian RL (Duff, 2002; Ross el al., 2007; Ghavanzadeh et al., 2015), as we added an explicit decomposition of the lotal rewards h into extrinsic rewards rf, stemming uniquely from x A, and intrinsic rewards r stemming from O, i.c., r² : O + C C R. A BAMDP explicitly incorporates the agent's updated belief about the unknown dynamics parameter, expressed via the posterior distribution p(0 - in the short form notation where D = (s,a,8) 8') into the decision-making process, in addition to expressing the complementary task of
4

---

Published in Transactions on Machine Learning Research (07/2025)
exploring the environment to learn this parameter through the intrinsic rewards defined function of the current. beliel on 0, r0). We have intentionally kept. the definition of (0) vague for now, as in the next section we will discuss how to develop a measure for (0) with certain desirable properties, such that the BAMDP aligus back with the original underlying MDP asymptotically during training, through its value lunctions and discounted cumulative returns.
3 Information-Theoreii Intrinsic Rewards
In this section, we discuss a principled approach to define the intrinsic reward (0) using deformation-theored mcasures. Specifically, we focus on the concept of Bayesion Surprise (Itti & Baldi, 2009), which quantifies the amount of information acquired about the unknown dynamics parameters 0 E O after observing new data. D of sample size n. By incorporating an information-theorrlii measure into the intrinsic reward, we encourage the agent to explore areas of the state-action space 8 x A that are most informative in terms of uncertainty reduction from prior p(0) to the posterior p(0)D").
3.1 Epistemic versus Aleatoric Uncertainty
The total variation residing in the random variable 0 E θ can be quantified with the Shannon entropy H [p(0)] = Te- p()) (Shanmon, 1948). The total entropy H [p(0) can be decomposed, with respect to another random variable Y, as H [p(0]] = H [p(0]Y)] + (0;Y) (Cover, 1999). Here, the conditional entropy H [p({[Y]] quantifies the residual uncertainty after observing all the realization of Y = y, i.c., alcatoric uncertainty. The mutual information component, T(0; Y), measuring the uncertainty reduction in one variable after obscrving the other, incorporates all the epistemic components. Conditioning everything on some n sized dataset D", we obtain:
[v -Y,(""
(1)
where the conditional mutual information T(0; Y|D") (Y, is typically referred to as Information Gain (Lindley, 1956) and it measures the reduction in the entropy of parameters 0 E θ before and after conditioning on the new random variable Y. Thus, intuitively IGa(-) captures the epistemic uncertainty gains (Hüllermeier & Waegeman, 2021; Wimmer et al., 2023) relative to parameters 0 E θ as we progressively gather more data.
Information gain IGa(-) appears to be an appealing candidate for a measure of intrinsic reward. In the context outlined by (BA)MDPs the information gain can be written in full form as:
[Gg(St, ae St+1) [p(0] sr, ar) Hp((O st, ar st+1)]
(2)
(e [log (s, oe; 8141)] (ogg ax)]
where, similarly to before, (Go 0t, St+1) measures the reduction in uncertainty around 0 O after observing the transition (se, dt, 8141). Houthooft et al. (2016) were the first to leverage information gain IGa(.) bonus such as the one described above, by specifically employing a variational Bayesian Neural Network to compute a one-step IG inside a model-free policy gradient learner. Their method neither embeds TG into a multi-step planner nor provides any formal analysis of bonus decay or convergence.
Relation Visitation-Frequencyy Bonuses Bellemare et al. (2016) show that both exact counts and pseudo-counts can be interpreted as the Bayesian Information Gain of an agent that maintains a density model of the observations encountered; the corresponding pseudo-count (or count) exploration bonus therefore be viewed as an explicit measure of epistemic uncertainty well. Earlier count-based based algorithms such as MBIE-EB (Strehl & Littman, 2008) and UCRL2 (Auer et al., 2008) already exploited this idea implicitly by adding a bonus 3/ 3/VN(s/ (x,(sa)) to guarantee optimism and PAC-MDP or regrel bounds. Subsequent density-mon model variants - e.g. hashing pseudo-counts (Tang et al., 2017) and PixeICNN eseudo-counts (Ostrovski et al., 2017) retain the same uncertainty interpretation while scaling to high-dimensionall inputs. In these works the intrinsic reward decays at the non-asymptotic rate (/1//N(s,))) a)). key distinction with the dynamics model-based approach considered in this work however lies in the fact that forward dynamics
5

---

Published in Transactions on Machine Learning Research (07/)025)
models naturally provide both next-state and reward predictions, together with the possibility of producing uncertainty measures for H-step-ahead predictions. In contrast, visitation-frequee bonuses depend on future counts, which are less straightforward to incorporate into look-ahead planning, as both a density model and a forward dynamics model would need to be specified.
3.2 Why Information Gain Intrinsic Rewards?
The intrinsic reward defined as rf(0) = IGo($t, ag st41) possesses the desirable property that is guaranteod to gradually diminish as the agent explores the environment. This happens because, as the agent collects more transitions from the environment, its posterior belief over the dynamics parameters E ④ becomes increasingly concentrated thus, the epistemic information gain from new observations naturally decreases. To formally demonstrate this, assume for simplicity S,A S R and that the transition distribution pe.... (s'|s, a) with true parameters Orrue θ follows a Markov dynamics characterized by the following additive noise structural equation:
S ((St,AA) t where e ~ and E
(3)
and where / F is some arbitrarily complex function. Suppose we want to model the two unknown parameters in (3), i.., / = (f, E F x Rt, in a Bayesian way, by placing a prior p(0) = p(f, o²) = p(f)p((?) Consider the (-) associated with a specific one-step transition (s, a, s'), i.c., IGo(s, a, 8'), so that the data DD in this setting is equal to D = {[3;, a₁ +))); Under the mild assumptions of Otrae E KL-support(p(O)) and "lestaiility" (Schwartz, 1965; Ghosal et al., 1999) discussed in the Appendix, Section A.1, we have that:
Proposition 3.1 (Consistency). Assume the true data generating model in (3) with true parameters Otrue (ftrue, oree), prior distribution p(0) E P and the induced posterior p(0 (s s()]))) Given conditions for weak posterior consistency (discussed in the proof), such that for c > 0,
Parge 0,
p(0 E : d(F,o), (Strue, otrue)) el (1,
as n. + 00, then r²(0)= = (G, a, 8²) Paique 0 as n n 00.
More practically, since we are interested primarily carrying out Bayesian inference on the unknown function f E F for exploration purposes, we might also treat o² as fixed and estimate via maximum likelihood, while instead maintaining only Bayesian prior p(f) (e.g., Gaussian Process, Bayesian Neural Networks, Deep Ensembles, etc.).
Under the mild assumptions of Proposition 3.1, we can guarantee that the exploration incentive = rG=() associated with a full transition (s, a, s') asymptotically converges to zero, as the agent progressively learns ④ via collecting new samples. Notice that this implies that the state-value (and action-value) function clicited by the augmented rewards rt- = rf + nfr(()) = r + n,GG()), which characterize the BAMDP defined above in 2.1, convergences in probability to the original MDP's value functions. More formally, this reads: Corollary 3.2 (Value Function Convergence). Given the state-valve function associated with the BAMDP, V...-- VBAMDP(s), where r:())= = IGg(-) defined as:
(BAMDP(s) = Ea-a(1s)) (s) + m IG, s') + Ep, V [[AAMD((']]].
and the value function VTBBp(s) associated with the original MDP, then we have that, under the conditions outlined in Proposition 3.1, VBAMDP(s) E, VMBp(s) as + 00. A similar result holds for the action-value function Q a) P, Q Mxpp(s a) and the discounted returns JBAMDD(m;P-) Ps JuDDp(R;p-).
These results ultimately establish, albeit asymptotically, the soundness of information gain as a desirable intrinsic reward signal, cnsuring that it effectively drives exploration the carly stages of learning, while naturally diminishing as the agent becomes more familiar with the environment.
3.3 Information Gain Decay and Problem Complexity
While Proposition 3.1 guarantees that the intrinsic reward r() IGo ae s141) asymptotically converges to zero the number of observations increases, the rate at. which this occurs is influenced by factors affecting
6

---

Published in Transactions on Machine Learning Research (07/2025)
the hardness of the f E F estimation problem, such as the signal-to-noise ratio and the smoothness of the true function So, and the type of prior p(ʃ) that is placed on f. For example, suppose we consider a Gaussian Process (GP) regression (Rasmuussn et al., 2006) prior: flw ~ ((:-))) where C(.-lw) is kernel covariance function with hyper- pper-parameters w E Q. We can further derive minimax contraction rates en (Ghosal & Van der Vaart, 2017) at which r0(0)= = (G) converges by imposing functional form restrictions on the true underlying ftrue E F. If we assume that: X' = S A = [0, 1]14)+1; fo E (2+) n Hº (2), where (a()) is the Hölder space and ())+) is the Sobolev space of order a C(a,y) = wi ||-- vl1 K, (wallx sopll=->ll in
flw ~ (P C(.,-ww) is the Matérn kernel, then we have (Van Der Vaart & Van Zanten, 2011):
Proposition 3.3 (Contraction Rates). Under the same assumptions of 3.1, if (fo, 0o) C* n (²) x (c, dl, where& X' = [0, 1] and flw ~ -ww)) where ((::-lw) is the Matérn kernel, then:
P 0 as n + 000
[sa at, S(+1) at the optimal mánimaz rate C, (لFा 7at (Yang & Barron, 1999).
The optimal minimax rate en achieved in Proposition 3.3 typically pertains to a purely 'passive' or random sampling strategy, and can potentially be improved under certain conditions through 'active' sampling (Willett et al., 2005; Castro & Nowak, 2008; Hanneke & Yang, 2015). We include a brief discussion on this in the Appendix, Section A.3. Also, note that VBAMDP(s) Ls (s) obviously occurs at the same rate.
Notice that the same type of asymptotic convergence characterizing IGg(-) is not guarantcod for other types of intrinsic rewards ri: such as those based on prediction error rt / [o ae) - s++1lb Ilp (Stadie et al., 2015; Pathak et al., 2017) in stochastic environments. In settings such as the one described in the 'noisy TV problem' (Tang et al., 2017), prodiction error can remain persistently high duc to inherent alcatoric uncertainty, randomness in the environment that cannot be reduced through further data collection.
As a last note, we briefly comment on a natural point of comparison between the convergence results derived above and those concerning the visitation-based bonuses. While our contraction rate en = - -22/(2-12)/a)) pertains to the information gain (IG) bonus under smoothness assumptions on fo E F and a GP prior, visitation-frequencc--
bonuses typically rely on different statistical setup. In the visitation-based frame- work, each (s, a) pair is treated as independent cell, and the exploration bonus decays as 0(1/ OLL//N(s, /N(ss a)) where N(s, a) is the visit count. This typically assumes a tabular or "discretizable" domain, via pseudo-counts (Bellemare et al., 2016) - or at least that good density model is learnable (Ostrovski et al., 2017) -, and does not explicitly model structure in the transition function. In contrast, our analysis assumes a continuous domain X' S x A, where the decay of IG is tied to the learnable complexity (e.g., smoothness a) of the transition dynamics/ f F.
F. Thus, while the two forms of decay arise from different assumptions, our results can provide a principled guide in continuous domains where discretization or pseudo-counts are intractable or hard to interpret.
4 Efficient Model-Based Bayesian Exploration
While we have demonstrated above that defining intrinsic rewards via information gain ri(() TGo(st, at, 81+1) is theoretically appealing, this poses two practical challenges for its implementation in RL, and especially in MBRL domains. Firstly, computing [Go(#t, Qr, St+1) for each timestep requires a posterior update over after observing each new transition (St, Tt, $e+1). Maintaining and updating exact Bayesian posteriors in high-dimensional parameter spaces can be computationally expensive, especially when parameterizes complex neural networks, as in Houthooft et al. (2016). Secondly, full IGo(-) for a state-action pair (st, ()))a)) is only calculated after the next state 814-1 is observed and the posterior is updated accordingly, p(O]st, [1)|s,,0 at St-1). Hence, similarly to other intrinsic rewards, IGe(s,, de, St41) can only be computed "in retro- spect", which makes it sample inefficient and complicates its integration into planning or policy optimization MBRL algorithms that instead typically look ahead.
A computationally more tractable, and planning-friendly, alternative is to employ the posterior predictive distri- bution directly, which marginalizes over the parameters, i.c., p($ a:) = I p( a;; O ())( az) d8, instead of the updated posterior distribution over the parameters after seeing st-1, i.c., p(0) st, ar st-1). However, in this scenario, we need to define predictive surrogate for IGo ar 8141), i.e., an Expected
7

---

Published in Transactions on Machine Learning Research (07/2025)
Information Gain, EiGo(s;,a)) at), that marginalizes over the distribution of possible next states 8441- In this way, we can avoid explicitly updating the posterior for every new transition. The surrogate EIGo(s:, 0z) can be derived as follows:
ElGg(s,, a4) = Eplan m..))) [Go ae, s421)
(4)
Ep(0) pleeoa p(gp(Osss St, Ct, St-1) 1) - log logp(0)ss at) [log p(pr+((+1)) 18,, dri 0) log 1 ar]
a (1
[p((5 1118811116,00 a:)] [p( [[[(( Sc-
8)]]
(5)
where use the fact that p(0 st, ae, 81+1) ) p(0)s az ( (e, ar, 0) by Bayes rule and M[sc+188e 0z)
Sg p(st+1lse, ag, ())() az )d by marginalization. Full derivation is provided in the Appendix, Section B.1. One can intuitively view EIGg(s;, at) as the expected reduction in uncertainty over 8, if the agent takes action ae in state 8t. Notice also that equivalently we can express EIGg(st, ax) as EIGg (8;, (z) = H [p(*++1/8) at] o.) [+[p(#+1111) ax; f]] ex) [Dk (p(8++1 (st, ae; 0) [p(8t+1|88 az))] Since EIGo(st, ae) does not explicitly require intermediate posterior computations, we do not necessarily need to maintain an updated posterior p(0 8t; (+) throughout the training. Instead, for instance, we can resort to approximale methods, such as an ensemble of probabilistic deep dynamics models {fe.. (s);,)))mm that approximate the predictive posterior via model averaging (Lakshminarayanan) et al., 2017; Wilson & Izmailov, 2020). One then show that, under deep ensemble model of the dynamics, (IG ae) is theoretically equivalent to computing the disagreement among the ensemble's predictive distribution (Hanneke et al., 2014; Shyam et al., 2019), that is:
EIGo(st, 0z) = Ppp()s ))) [log o2; 0) - log pl s 18t, at]]
Ep(0)s ppleet m) log p 1 az; 0) log / p(sc+1lse ar ()p(((ss,)))))))))))))))))) ax) d0]
[Ep(o] [p(s:+ 19e. ari 0)]] ap()) [p ar: = JSDo (p(s441886, a,; ax;) 0 ~ p(
(6)
where JSD is the Gensen-Shannon Divergence. The intuition behind the equivalence in (6) is that, under a deep ensemble dynamics model, EIG can be computed as the level of disagreement across different plausible dynamics models consistent with the current data. Full derivation is provided in the Appendix, Section B.3. The above quantity can be approximated via Monte Carlo estimates by averaging the predictions of the probabilistic neural networks in the ensemble, as
EIGa(S:, ax) = H [Epcoee an) ((8841188 az 0)]] - Ep( [H[p(ee (8et 188, o: 0)]]
M
2 H
1 M
p(s ar; Om) m=1
Σ
1
H pp(octllse aes 0m))
M
m=1
Ensembles probabilistic neural networks are a cheaper, approximate alternative to other Bayesian regression models such as Gaussian Processes or Bayesian Neural Networks, and work reasonably well when exact. computation of the posterior (and posterior predictive) becomes intractable in high-dimensional state-action space settings.
4.1 Disagreement and Heteroskedastic Noise
An important distinction arises based on whether the environment's noise is homoskedastic heteroskedastic. In the homoskedastic case, the transition noise variance 6² remains (theoretically) constant across all states, o(s) for all e S. This constancy means that relative differences in a predictive uncertainty measure (e.g., predictive entropy or variance) primarily reflect the model's epistemic uncertainty: if all states are subject to the same inherent noise, then the model's uncertainly should vary mainly due to gaps in its knowledge about the dynamics, rather than state-specific stochasticity. In practice, of course, approximation errors can still inflate or reduce these estimates, so the absolute values of such measures will not perfectly align with epistemic uncertainty - thus, using EIGa(-) still more efficient (sce experiment Section 5.1.1 below). However, their variations across states can still somehow provide a more meaningful exploration signal in
8

---

Published in Transactions on Machine Learning Research (07 (025)
Homoskedastic
Heteroskedastic
H[-]
2
2.0
EIG
FIG
1
1.5
1.0
0.5
95%
-2.5
0.0
-2.5
Figure 1: Left panels show a homoskedastic regression case, where predictive H[:] and E-Ga(-) identify similar regions of uncertainty. Right pancls show heteroskodastic regression case, where predictive H[+ is influencod by state-dependent noise, while E-Gg(-) still identifies regions of epistemic uncertainty. The regression model employed is ensemble of probabilistic neural networks.
homoskodastic setting. In contrast, for hetcroskodastic noise, where the noise variance d²(s) depends on the state realization 8, the same measures blend irreducible aleatoric uncertainty with epistemic uncertainty. As result, highly noisy state regions may trigger misleading signals as they exhibit large predictive entropy (or variance) even if the model is confident about the underlying dynamics, simply because the environment itself is noisier.
Figure illustrates these behaviors in a simple regression setting with either homoskedastic or heteroskedastic noise. Under homoskedastic noise, the relative levels of predictive entropy (in red) and information- exploration measure E-Ga(-) (in green) aligm closely with regions of true knowledge gaps. In contrast, for heteroskedastic noise, predictive entropy spikes in the (middle) region with higher noise, whereas F.Go(-) remains low there because the model has ample data ie., EIGa(+) explicitly factors out irreducible noise, focusing states that would reduce the model's parameter uncertainty. This illustrates why not only need an uncertainty-awaan dynamics model but also need to carefully disentangle aleatoric and epistemic uncertainties when deriving exploration signals.
4.2 Predictive Trajectories Sampling with EIG Exploration Bonus
For some of the setups in the experimental Section 5 below, we will be using standard model-free policy gradient algorithms, namely Proximal Policy Optimization (PPO) (Schulman et al., 2017) for environments with discrete action spaces and Soft Actor-Critic (SAC) (Haarnoja et al., 2018) for ones with continuous action spaces. We generally employ vanilla versions of these algorithms, and versions that are augmented with intrinsic curiosity, i.e., the collected extrinsic rewards rf in the replay buffer D = {st, ae, Tt, 8t+1 Y are augmented with intrinsic rewards rf computed using the forward dynamics model according to the description in Section 2.1.
For some her specific setups instead, we instead employ a general Predictive Trajectory Sampling with Bayesian Exploration (PTSS-BB BE) framework, which extends prior dlanning-to-explle and model predictive control approaches (Chua et al., 2018; Shyam et al., 2019; Sekar et al., 2020) by supporting arbitrary Bayesian models of both dynamics and rewards. In this framework we maintain an ensemble of M probabilistic models (So., J each predicting the joint distribution po., (r, s' s,)). Notice that here we generally consider prediction of both reward and next state as part of encompassing dynamics prediction effort, but one can easily focus only next state prediction, particularly in settings where reward prediction is secondary. Because and s' are conditionally independent given (s, a), we factorize each pe., (r, |, ps., (r|s, a) pe., (s'|s, a). Under the usual assumption that (r, s') form multivariate Gaussian with diagonal covariance, deep ensembles provide a scalable, approximate Bayesian predictive posterior without explicit posterior updates. Instead of selecting actions through optimization procedures like Random Shooting or CEM, as traditional Model Predictive Control (Dracger et al., 1995; Chua al., 2018), we instead maintain a policy paramotrization T (e.g., PPO, SAC, etc.) that does not necessarily require updates at each time step t, as proposed also in Shyam et al. (2019) and Foster et al. (2021). Similarly to model-free algorithms, this policy is trained by
9

---

Published in Transactions on Machine Learning Research (07/2025)
Algorithm Predictive Trajectory Sampling with Bayesian Exploration (PTS-BE)
Initialize policy Tp. transitions bulfer Deno, Bayesian dynamicc model p(/) (e.g., GP, Deep Eusembles, ...)
2: for 7 steps do
Given state 8t; sample action at (s)); receive rt and 81+1
4: Add Deno + Deno U (st, de, St.t1
if t twrrm then
6: if f mod Epoltoy = 0 then
Sample K actions from policy {ag (a J== (x(()))
7:
8:
Rollout J-length trajectories and define Dinded = {f(o 445: (411
(6) (8) (a) 0414,,)
(1) (8) Compute rêts and augment each (tt in Denv as rets r8ts +nri+s rit, n B. 92+3) for G policy gradient updates do Update va (s) using Pmnddt if f mod (oode = 0 then
(k)
10:
11:
12:
13:
for M model updates do Update dynamics model p(f)Denu) + p(s) de,
14:
Te, 86+1))
maximizing the cumulative sum of both extrinsic rewards re and intrinsic rewards r Σ (18 + mr)), where r' again can be generally defined as an output of the predictive model for pe(r, s'|s, a) (e.g, Variance, Entropy, EIG, etc.), but uses a different data buffer. Specifically, the ensemble models are loveraged to generate K independent predictive trajectories from the current state st, each unrolled for J steps into the future, [(*++) aft;; r r 8/t (1) Actions are drawn from the current policy ob+; ( while next state and rewards are sampled from the predictive posterior po(r, s'|s a). Intrinsic rewards r are computed for every l and j entry and added up the estimated extrinsic ones, if included. Finally, the policy parameters i are updated periodically every n steps using purcly model-generated data from the PTS step. One could also polentially consider using combination of real and imagined trajectories, like in Model-Based Policy Optimization (Janner et al., 2019). A pseudo-code representation of the full algorithm's structure is depicted in Algorithm 1.
Obtaining P([[D) for different forward models. Line 14 in Algorithu yields an updated belief posterior P(f) on the environment dynamics. In practice, the way this belief posterior is obtained depends on which Bayesian dynamics model is used. In particular: (i) Exact GPs maintain a full multivariate Gaussian posterior p(f D), using (s,a) as inputs; ii) DKs maintain multivariate Gaussian posterior using latent input representation ions (is, a), where the marginalization is identical to the full GP, but performed in a latent space H: iii) Stochastic Variational variants of GPs and DKs (SVGP/ / SVDK) maintain Gaussian posterior using a subset of representative inducing points in cither (s, or Φ(s,a), respectively - the marginalization is again identical to the GP or DK, but carried out using only the inducing points; iv) Deep Ensembles (DEs) approximate the full belief posterior via an empirical mixture of M probabilistic neural networks marginalization is obtained via Monte Carlo averaging: P( D) M Σm= 6pm-
5 Experiments
In this section, we evaluate the performance and the practical benefits of nnformation-theorelle exploration bonuses through three sets of experiments. The first set is designed to validate the theoretical claims made in earlier sections and demonstrate the advantages of the proposed PTS-BE (Algorithm 1) in terms of planning deep exploration strategies. The second set introduces a simple unichain environment to allow comparison of the approximate uncertainty quantification properties of deep ensembles with those of Gaussian Processes (Rasmussen et al., 2006), which provide full Bayesian predictive posterior, in terms of their ability to generate reliable exploration signal based on epistemic uncertainty. In this second set of experiments, we include also Deep Kernels (Wilson et al., 2016b) in the comparison, and other popular curiosity- algorithms that use intrinsic reward augmentation on the environment's 's data buffer to update the policy via exploration bomuses computed "in retrospect" " (see discussion in Soction 4). Deep Kernels combine the ficxibility of deep learning with the principled uncertainty estimation of Gaussian Processes, making them valuable baseline for assessing the trade-offs between expressiveness and calibrated uncertainty in exploration-driven tasks.
10

---

Published in Transactions on Machine Learning Research (07/2025)
Error
Entropy
EIG
Heterosk Homosk
Heterosk Homesk
0.8
Haterosk
0.9
5.5
Homosk
Uncertainty
0.0
Uncorfainty
5.0
Uncertainty
0.6
0.7
0.4
4.5
4.0
100
200
300
0
100
200
300
100
200
300
Model Updates
Model Updates
Model Updates
Figure 2: Intrinsic reward uncertainty measures on the homoskedastic and heteroskedastic versions of the Mountain Car environment, recorded at every update of the dynamics model fo az), under a random policy, From left to right plot: prediction error rf (s as) sc-1l22, entropy rf = H (r, s s'|s, a)] and expected information gain r E-Ga(-).
Lastly, the third set of experiments showcases the applicability of the PTS-BE approach to more challenging, sparse-reward, Maze environments characterizing close to purely explorative tasks.
5.1 Heteroskedastic Noise Model in Mountain Car
We begin by introducing a modified version of the canonical Mountain Car environment. In our variant, we incorporate heleroskedastic noise model based on the environment's physical dynamics. The transition dynamics are modified that the noise variance is state-dependent,, S f(,,, A.) + where E, oO, thereby mimicking realistic scenarios where uncertainty varies with the agent's position and velocity. In particular, noise (S,) is higher for higher velocities (more chaotic dynamics at higher speeds) and for positions near the bottom of the valley (simulating factors such as different terrain conditions). Conversely, we define the homoskedastic version of the environment as one where the transition noise remains constant across all states, i.c., & ~ (O, d²). Throughout the experiments in this section, we make use of a deep ensemble estimate the dynamics [fem. (st, (()),a))) a;)
The first set of results reported in Figure 2 illustrates the values of different model-basod uncertainty measures that can be used to define the intrinsic rewards rí, in both the homoskedastic and heteroskedastic versions of Noisy Mountain Car. The measures are: i) mean-squared prediction error, defined as r = [[fo(s;,;) ax) S141/l2 (Stadic et al., 2015; Pathak et al., 2017); ii) entropy of the predictive distribution, defined as ri = H (rr al]s iii) EIG, defined according to Equation 6, as we are specifically making use of a deep ensemble dynamics model in this context. Each uncertainty measure is computed at every model update, that is, every moodet = 64 environment steps, and transitions are gathered under random policy. As can be seen from the plots of Figure 2, EIG progressively converges to zero the number of model updates increase, while prediction error and entropy plateau at higher values. Moreover, FIG shows very similar progression both in the homoskedastic and heteroskedastic noise scenarios. This is because EIG, in the form of Pensen-Shanme Divergence here, is low in regions with high noise but rich in data, since even if the predictive variance in each ensemble member is high, the overall "disagreement" among the ensemble's predictive distributions, [p(r, s'|s, a; (m.)))) Mm is low. Predictive entropy's 's values instead are consistently higher the heteroskodastic case relative to the homoskedastic one, since aleatoric and epistemic uncertainties are fused in together in the measure, resulting in stronger, but deceiving, exploration signal to the agent, the relative difference is driven solely by aleatoric uncertainty. Finally, prediction error curves slightly diverge for approximately the first 50 model updates, then platcau at similar values. Notice that importantly, even though the prediction error curves do converge to similar values in this specific example, where we employ deep ensemble, the process of averaging predictions in the ensemble can mask substantial disagreement among ensemble members mcaning that even when the ensemble's 's mcan prodiction has low MSE, the underlying epistemic uncortainty may still be high. Similar reasoning naturally holds also in this case of non- non-ensemble, and more generally non-Bayesian, predictive models (e.g., a single neural net).
11

---

Published in Transactions on Machine Learning Research (07/202) (225)
BC
PTS-BE
Homosk
Returns
1,0
A Molerss Namzak
1.0
Random
-1000/9 + 0.00
BF RH
-1000.9 + 0.000
Rac States Visted
0.8
0.8
BR F&G
-1000.0 0,000
0.6
PTS BD #
812.3
105.9
Fratt States
0.4
Heteroek m(x)
Returns
4
Bandom
1000.0 0000
0.2
0.2
H|]
1000.0
0.000
demoss
BIG
1000.0 00000
0.0
H[]
-00.0
0
200
400
600
800
1000
200
400
600
800
1000
PIS-BR BIG
-356.7 10/0
Steps
Figure 3: Results from the second experiment Noisy Mountain Car. The two plots report the cumulative fraction of visited states by the agent over 1000 steps, while the tables report the average returns. These are computed over 20 different replications. The methods being compared are: a random agent (Random), a baseline model-free agent with intrinsic rewards augmentatin ation (BE) and a model-based Predictive Trajectory Sampling (PTS-BE) agent, as described in Section 4.2. The BE and PTS-BE methods are compared on both the homoskodastic and hetcroskodastic versions of the environment, using either predictive entropy H[- or EIG as intrinsic reward. The baseline solver is PPO.
5.1.1 Planning to explore with PTS-BE
In the second set of results regarding the Noisy Mountain Car environment, we aim to assess the performance of EIG based intrinsic rewards (i.e., which we refer to as "Bayesian Exploration", or BE onward, for simplicity) in conjunction with Predictive Trajectory Sampling, in the form of the algorithm introduced in 1. To this end, we compare two types of algorithms. The first one is a baseline, model-free, PPO agent (Schulman et al., 2017) that uses intrinsic rewards augmentation and updates the policy retrospoctively based on these augmnented rewards in the buffer (we will denole this with the "BE" prefix). In short, at each policy update we form a dataset Dr {(s,, ag rf, s() of environment interactions, then replace each reward ri with r =f + nri, where rf is the intrinsic bonus for that transition (Entropy or EIG) and nj its scaling factor. Then, we performance the usual PPO update. The second is the PTS-BE algorithm presented in Section 4.2, which instead leverages the learnt dynamics model to do Predictive Trajectory Sampling, and updates a policy using these sampled trajectories. We also include a random policy agent as an additional baseline ("Random"). The results of this experiment are reported in Figure 3. Both BE and PTS-BE are run with different configurations, that is: i) using either predictive entropy H[: or EIG as intrinsic rewards (hence the different suffix); ii) under either the homoskedastic or heteroskedastic version of the Noisy Mountain Car environment. For the PTS-BE specifications, we use an horizon of J 100 steps ahead and K 10 independent rollouts. The results demonstrate the natural data efficiency of the PTS-BE method in planning for deep exploration, while BE methods would necessitate lot more samples to achieve similar state space coverage. On the difference between intrinsic reward measures employed, we notice that EIG generally does better than predictive entropy both in the homoskedastic and heteroskedastic scenarios. Furthermore, we note that a random policy achieves under 40% state-space coverage in 1.000 steps, under 70% in 5.000 steps, and about 90% only after 10.000 steps. And in none of these instances does solve the environment's lask. By contrast, PTS-BE-EIG solves the task 18/20 times in the heteroskedastic setting and 17/20 in the homoskedastic one, in fewer than 1.000 steps. PTS-BE-HH H|] solves it only 8/20 and 10/20 times, respectively, within the same budget, demonstrating once again the superiority of EIG as intrinsic reward signal. Finally, we also emphasize how PTS-BE-H[[] exhibits higher variance than TIS-BE-EICS as its error bands are significantly wider.
5.2 Dynamics Models Comparison
In the second set of experiments we introduce an alternative simple unichain environment, similar to the ones used by Osband et al. (2016) and Shyam al. (2019). The aim of this last sctup is to compare the performance PTS-BE with different dynamics model specifications, in particular with Gaussian Processes (GPs) (Rasmussen et al., 2006) and Deep Kernels (DKs) (Wilson et. al., 2016b). GPs considered to
12

---

Published in Transactions on Machine Learning Research (07/202) (225))
Ransom PAD PROICH MIME
Bardom SPO
Merked
5. Stattee
SPO ICM
1200 3031
631 1-033
0.0
2.5
Camulaativ
0.4
0.4
Macked
Retuess
120
642 2-11
F20
2210 13UK
m5
USS +un
775485-8
659 1/00
25
75
1.0
125
250
1775
200 125
375
202
Staps
Staps
Figure 4: Results from the L 50 states Unichain environment. The first plot reports cumulative fraction of visited states. The second plot reports cumulative rewards. The tables include fraction of visited states and cumulative rewards respectively, at episode termination. The models being compared are: a Random policy, PPO, PPO with ICM intrinsic rewards, PPO with VIME, and three version of the PTS-BE algorithm with different dynamics model specifications, i.e., GP, Deep Kernels (DK) and Deep Ensembles (DE).
be a gold standard for uncertainty quantification in regression problems, as they elicit a full, closed form, posterior predictive distribution on the output. The downside with GPs is that they notoriously scale poorly with sample size and input dimensions. In particular, in the context of dynamics MBRL modeling, their computational cost amounts to O x Al) for evaluating the kernel matrix, in addition to the (nn") and O(n²) costs incurred for matrix inversion at train and test time, respectively. Additionally, in MBRL, we typically model [S] + 1 (states and rewards) outputs, which implies that we would need to maintain and update [S] + 1 different GP models, increasing the cost by a factor proportional to [S] + 1. For these reasons, deep ensembles are often preferred for their better scalability with input and output dimensions, albeit at the cost of having an approximate posterior predictive distribution. DKs (Wilson et al., 2016b) instead represent alternative, principled way of reducing the input's computational cost in GPs. They first map inputs to a latent space H. fr. S x A + H. where fr. is a neural net. Then a GP prior is placed on the latent space, fs : H + S. The parameters in both Jn and fa are optimized end-to-end by minimizing a single loss. DKs reduce the kernel matrix evaluation costs from Onn' x Al) to Onn)HH), where [H] is controllable by the user. However, they still bear the cost of maintaining [S] + GP models, one per output, on the latent space. Finally, we specify that in order to reduce the cost associated with sample size n only, we employ sparse input variational methods for both GP and DK specifications (Titsias, 2009; Wilson et al., 2016a). This reduces the kernel evaluation cost from (nn S x Al) to Owmm|S x Al), where m is the size of the chosen sub-sample, and the matrix inversion cost from O(n²) to O(m²) at train time and from O(n²) to O(m²) at test time for GPs (for DKs, replace |S x A| with [H)).
The unichain environment version we consider made of = 50 discrete states and three discrete actions: {go-left, stay, go-right]. We use a continuous representation of the discrete states as in Osband et al. (2016) and Shyam et al. (2019). The agent spawned at the second state so = 2. Rewards are sparse, equal to 0.001 in the first state and to 1.0 in the last state, while they are zero everywhere elsc. A visual depiction of the unichain environment is included in the Appendix, Figure 8, together with an additional experiment on a L = 100 states, noisy unichain environment. Results from the experiment are depicted in Figure 4. The plots measure the cumulative fraction of visited states and the cumulative rewards respectively, for each method. The tables instead report the fraction of visited states and the cumulative rewards achieved at episode termination. The list of methods compared is the following: a random agent (Random); a PPO agent (PPO); a PPO agent augmented via. Intrinsic Curiosity Module (ICM) (Pathak et al., 2017), popular method based on prediction error intrinsic rewards (PPO ICM); a PPO agent augmented with VIME intrinsic rewards (Houthooft et al., 2016), which are based the 1G specification introduced in Equation 2 (PPO VIME); lastly, we include three versions of the PTS-BE algorithm, using EIG as an intrinsic reward, with different dynamics model specifications ions based on GPs (PTS-BE-GP), Deep Kernels (PTS-BE-DK) and Deep Ensemble (PTS-BE-) BE-DE). The results again show superiority of the PTS approaches compared to popular methods such as ICM and VIME, for scttings where deep look-ahead exploration is required to preserve sample efficiency. In addition, we note that GPs and DKs specification display only marginally better performance, as they elicit full posterior predictive distribution, compared to Deep Ensembles. This demonstrate that, even if the
13

---

Published in Transactions on Machine Learning Research (07/ (225)
Open
U-Shape
Double U-Shape
Obstacles
1.0
1.0
1.2
SAC
1.0
SAC
SAC
SAC
BE-SAC
9.5
PTS-GAC
PIS-SAC
PIS-E--AA
PIS-BE-SAC
3.6
0.6
0.6
3.4
0.4
0.4
SNC
0.2
BE-SAC
PTS SAC
8.0
PTS SAC
0.0
0.0
0.0
Slapa
Figure 5: Results on the four different configurations of the maze environment, namely: i) an open maze with no obstacles ("Open"); ii) a U-shaped maze ("U Shape"; iii) maze featuring two U turns ("Double U Shape"); iv) a large maze with several obstacles in the middle ("Obstacles"). The plots show the cumulative probability of reaching the goal at every step t, i.e, p(r;- = 1), for each of the four methods compared: SAC, BE-SAC, PTS-SAC and PTS-BE-SAC.
posterior predictive distribution in Deep Ensemble is only approximate and may not have the same coverage properties of the full one in GPs, still provides sufficiently reliable signal for exploration.
5.3 Maze Exploration Environments
Point Maze Environments This set of experiments features four types of 2D Gymnashum Robotics Point Maze structures (Fu et al., 2020) of increasing complexity. The agent's 's task in this setting is to explore the environment and collect goal object located somewhere in the maze. The environments are characterized by continuous states and actions. Rewards are sparse, equal to zero everywhere except from in the goal state, where rf 1. The four environments include: i) an open maze with no obstacles ("Open"); ii) a U-shaped maze ("U-Shape)); iil) a medium size maze with double U-shape ("Double U-Shape"): iv) a larger maze with random obstacles in the middle ("Obstacles". The methods we compare here are: Soft Actor-Crite (SAC) (Haarnoja et al., 2018); SAC with Bayesian Exploration intrinsic reward augmentation (BE-SAC) based EIG; version of PTS planning algorithm without any intrinsic rewards, based only extrinsic signal, using SAC as solver (PTS-SAC); lastly, the PTS-BE algorithm using SAC policy (PTS-BE-SAC). For the PTS-BE specification, we use an horizon of J = 64 and K = 16 independent rollouts. Results are reported in Figure 5. The plots show the cumulative probability of reaching the goal and solving the maze at every time step t, i.e. p(rī = 1), for cach different maze configuration.
The results demonstrate that, even in the case of continuous control problems, PTS-BE guarantees better performance and sample efficiency, it generally achieves the goal in significantly less steps. In particular, the gap in performance with respect to SAC, BE-SAC and PTS-SAC agents is smaller in the case of the first "Open" maze, as this is a considerably easier scenario to solve relative to the other configurations, while it becomes significantly larger in the other three setups where the task gets increasingly harder. The baseline PTS-SAC was included to further disentangle the contribution of model-based planning versus that of information gain intrinsic rewards. As expected, PTS-SAC alone does not really push performance up significantly, which confirms that the substantial gain particularly stems from the synergy between the information gain intrinsic rewards and the baamning-to-xploore approach, beyond planning alone.
Ant Maze Environments The final set of experiments involves two different maze environments from the Gymnasium Robotics suite, featuring higher-dimensional state-action spaces. As with the Point Maze tasks, these Ant Maze environments were originally introduced in Fu et al. (2020). In these tasks, robotic ant must learn to 'walk' from a starting location to designated goal. The action space A consists of 8 dimensions corresponding to torques applied to the hinge joints, while the state space includes 27 positional features from various body parts. Although rewards are sparse by default (1 when the goal is reached, 0 otherwise), we incorporate a dense extrinsic reward signal to encourage locomotion behavior during training.
We evaluate performance in two maze variants a small and a large open maze usinn the same methods in the Point Maze experiments: SAC, BE-SAC, PTS-SAC, and PTS-BE-SAC. Results are presented in
14

---

Published in Transactions on Machine Learning Research (07/ (025)
Small Open Ant-Maze
Large Open Ant-Moze
SAC FTE-355AC
SAC MA PTS40
Min Dista to Goal
Min Dista nca to Goal
Caverage 155)
Covarage (56)
40-
4+
6
SAC FI24AC
0-
10000 20000
10000
20000
20000 30000
20000
30000
Steps
Steps
Steps
Steps
Figure 6: Results for the two Ant Maze environments, the Small and the Large Open mazes. For both environments we report the percentage of covered states ("Coverage (%)") and the minimum Euclidean distance to the goal state ("Min Distance Goal") over 10 different runs. The task is solved when this distance S 0.5. Methods compared are again SAC, BE-SAC, PTS-SAC and PTS-BE-SAC.
Figure 6. We report both the percentage of states covered (" ("Coverage (%)"" and the minimum Euclidean distance to the goal achieved ("Min Distance to Goal") to assess not only exploration but also significant progress toward the target. A goal is considered reached when the agent's position is within a Euclidean distance of 0.5 from the goal coordinates. Consistent with prior results, PTS-BE-SAC is the only method among those compared that successfully solves the task in both environments, reaching the goal in an average of 7,281 + 1,531 steps in the Small Open Maze and 15,237 + 4,372 steps in the Large Open Maze.
6 Conclusions
In this work, we tackled the challenge of data-efficient exploration in MBRL by studying principled information-thcoretie approach to intrinsic motivation. Firstly, we derived class of information-thooree exploration bonuses grounded in Bayesian principles, specifically designed to target the epistemic uncertainty rather than the aleatoric uncertainty stemmiin from the environment ("Bayesian Exploration"), We then demonstrated, both thcorctically and empirically, that, unlike other cxisting intrinsic rewards methods, these information gain bonuses have the desirable properties of naturally signaling epistemic gains and converging to asymptotically. We also discussed tractable equivalent measures that be efficiently approximated with less computational effort via deep ensembles dynamics model (Shyam et al., 2019). Finally, we introduced general Predictive Trajectory Sampling with Bayesian Exploration (PTS-BE) framework compatible with any Bayesian dynamics and reward model - that embeds these bonuses into multi-step planning. Empirically, PTS-BE outperforms reactive intrinsic-reward and modcl-free basclines on sparsc-reward and pure-exploration lasks. As a final note, we conclude by stressing that, as with most planning- based MBRL methods, PTS-BE trades higher computational cost (maintaining a policy, Bayesian model, and imagined rollouts) for improved sample efficiency. This trade-off is ad advantagoous when environment interactions are expensive, while simpler intrinsic motivation methods may suffice when data is cheap.
References
Eitan Altman. Constrained Markov decision processes. Routledge, 2021.
Mauricio A Alvarez, Lorenzo Rosasco, Neil D Lawrence, el al. Kernels for vector-valued functions: review. Foundations and Trends in Machine Learning, 4(3):195-2668 2012.
Peter Auer, Thomas Jaksch, and Ronald Ortner. Near-optimal regret bounds for reinforcement learning. Advances in neural information processing systems, 21, 2008.
Kamyar Azizzadenesheli, Emma Brunskill, and Animashree Anandkumar. Efficient exploration through bayesian deep nettworr... In 2018 Information Theory and Applications Workshop (ITA), pp. 1-9. IEEE, 2018.
15

---

Published in Transactions on Machine Learning Research (07/2025)
Andrew G Barto. Intrinsic motivation and reinforcement learning. Intrinsically motivated learning in natural and artificial systems, pp. 17-47, 2013.
Mare Bellemare, Sriram Srinivasan, Georg Ostrovski, Tom Schaul, David Saxton, and Remi Munos. Unifying count-based exploration and intrinsic motivation. Advances in neural information processing systems, 29, 2016.
José M Bernardo. Expected information expected utility. the Annols of Statistics, pp. 686-690, 1979.
Yuri Burda, Harrison Edwards, Amos Storkey, and Oleg Klimov. Exploration by random network distillation. In International Conference on Learning Representations, 2018.
Rui M Castro and Robert Nowak. Minimax bounds for active learning. IEEE Transactions on Information Theory, 54(5):2339 2353, 2008.
Nuttapong Chentanez, Andrew Barto, and Satinder Singh. Intrinsically motivated reinforcement learning. Advances in neural information processing systems, 17, 2004.
Kurtland Chua, Roberto Calandra, Rowan McAllister and Sergey Levine. Deep reinforcement learning in a handful of trials using probabilistic dynamics models. Advances in neural informalion processing systems, 31, 2018.
Thomas M Cover. Elements of information theory. John Wiley & Sons, 1999.
Marc Deisenroth and Carl E Rasinussen Pilco: A model-based and data-efficient approach to policy search. Proceedings of the 28th International Conference on machine learning (ICML-11), pp. 465-472, 2011.
Andreas Draeger, Sebastian Engell, and Horst Ranke. Model predictive control using neural networks. IEEE
Control Systems Magazine, 15(5):61-6) 66, 1995.
Michael O'Gordon Duff. Optimal Learning: Computational procedures for Bayes-adapline Markov decision processes. University of Massachusetts Amherst, 2002.
Adam Foster, Martin Jankowiak, Elias Bingham, Paul Horsfall, Yee Whye Teh, Thomas Rainforth, and Noah Goodman. Variational bayesian optimal experimental design. Advances in Neural Information Processing Systems, 32, 2019.
Adam Foster, Desi R. Ivanova, Ilyas Malik, and Tom Rainforth. Deep adaptive design: Amortizing sequential bayesian experimental design. In International Conference on Machine Learning, pp. 3384-3395. PMLR, 2021.
Justin Fu, Aviral Kumar, Ofir Nachum, George Tucker, and Sergey Levine. D4rl: Datasets for deep data-driven reinforcement learning. arXiv preprint orv::20004.022 2020.
Yarin Gal and Zoubin Ghahramani. Dropout bayesian approximation: Representing model uncertainty in deep learning. In international conference on machine learning, pp. 1050-1059. PMLR, 2016.
Yarin Gal, Riashat Islam, and Zoubin Ghahramani. Deep bayesian active learning with image data. In International conference on machine learning, pp. 1183 1192. PMLR, 2017.
Ather Gattami, Qinbo Bai, and Vaneet Aggarwal. Reinforcement learning for constrained markov decision processes. In International Conference on Artificial Intelligence and Statistics, pp. 2656 2664. PMLR, 2021.
Mohammad Ghavamzadeh, Shie Mannor, Joelle Pineau, Aviv Tamar, et al. Bayesian reinforcement learning: A survey. Foundations and TrendsS in Machine Learning, 8(5-6):359-493 2015.
Subhashis Ghosal and Aad Van der Vaart. Fundomentals of nonparametric Bayesion inference, volume 44. Cambridge University Press, 2017.
16

---

Published in Transactions on Machine Learning Research (07/2025)
Subhashis Ghosal, Jayanta K Ghosh, and RV Ramamoorthi. Posterior consistency of dirichlet mixtures in density estimation. The Annals of Stalistics, 27(1): 158, 1999.
Chuan Guo, Geoff Pleiss, Yu Sun, and Kilian Q Weinberger. On calibration of modern neural networks. In International conference on machine learning, pp. 1321 1330. PMLR, 2017.
Tuomas Haarnoja, Aurick Zhou, Pieter Abbeel, and Sergey Levine. Soft actor-critic: Off-policy maximum entropy doep reinforcement learning with a stochastic actor. In International conference on machine leurningg pp. 1861 1870. PMLR, 2018.
Steve Hanneke and Liu Yang. Minimax analysis of active learning. J. Mach. Learn. Res., 16(1):3487 3602, 2015.
Steve Hanneke et al. Theory of disagreement active learning. Foundations and TrendoB Machine Learning, 7(2-3):131 309, 131-309,, 2014.
James Hensman, Nicolò Fusi, and Neil D Lawrence. Gaussian processes for big data. In Proceedings of the Twenny--Nnnhh Conference on Uncertainty in Artificial Intelligence, pp. 282-290, 2013.
Todd Hester and Peter Stone. Intrinsically motivated model learning for developing curions agent. In 2012 IEEE international conference development and learning and epigenetic robotics (ICDL), pp. 1 6. IEEE, 2012.
Neil Houlsby, Ferenc Huszár, Zoubin Ghahramani, and Máté Lengyel. Bayesian active learning for classífication and preference learning. arXiv preprint arXiv:111..57/3. 2011.
Rein Houthooft, Xi Chen, Yan Duan, John Schulman, Filip De Turck, and Pieter Abbeel. Vime: Variational information maximizing exploration. Advances in neural information processing systems, 29, 2016.
Eyke Hüllermeier and Willem Waegeman. Aleatoric and epistemic uncertainty in machine learning: An introduction to concepts and methods. Machine learnóng, 110(3):457 506, 2021.
Laurent Itti and Pierre Baldi. Bayesian surprise attracts human attention. Vision research, 49(10):1295 1306, 2009.
Michael Janner, Justin Fu, Marvin Zhang, and Sergey Levine. When to trust your model: Model-based policy optimization. Advances in neural information processing systems, 32, 2019.
Daniel Jarrett, Corentin Tallec, Florent Altché, Thomas Mesnard, Remi Munos, and Michal Valko. Curiosity in hindsight: Intrinsic exploration in stochastic environments. In International Conference on Machine Learning, pp. 14780 14816. PMLR, 2023.
Sham M Kakade. A natural policy gradient. Advances in neural information processing systems, 14, 2001.
Sanket Kamthe and Mare Deisenroth. Data-efficient reinforcement learning with probabilistic model predictive control. In International conference artificial intelligence and statistics, pp. 1701 1710. PMLR, 2018.
Alexander S Klyubin, Daniel Polani, and Chrystopher Nehaniv, Empowerment: A universal agent-centric measure of control. In 2005 IEEE Congress on Evolutionary Computation, volume 1, pp. 128-135. IEEE, 2005.
Pawel Ladosz, Lilian Weng, Minwoo Kim, and Hyondong Oh. Exploration in deep reinforcement learning: A survey. Information Fusion, 85:1 22, 2022.
Balaji Lakshminarayanan, Alexander Pritzel, and Charles Blundell. Simple and scalable predictive uncertainty cstimation using deep ensembles. Advances in neural information processing systems, 30, 2017.
Timothy P Lillicrap, Jonathan Hunt, Alexander Pritzel, Nicolas Heess, Tom Erez, Yuval Tassa, David Silver, and Daan Wierstra. Continuous control with deep reinforcement learning, arXiv preprint TT:::500900277 2015.
17

---

Published in Transactions on Machine Learning Research (07/2025)
Dennis V Lindley. On measure of the information provided by experiment. The Annals of Mathematical Statistics, 27(4))986-1000 1005, 1956.
Viraj Mehta, Biswajit Paria, Jeff Schneider, Willie Neiswanger, and Stefano Ermon. An experimental design perspective on model-based reinforcement learning. In International Conference on Learning Representations, 2022.
Volodymyr Mnih, Koray Kavukcuoglu, David Silver, Andrei A Rusu, Joel Veness, Mare G Bellemare, Alex Graves, Martin Riedmiller, Andreas K Fidjeland, Georg Ostrovski, et al. Human-level control through deep reinforcement learning. nature, 518(7540):529 533, 2015.
Ian Osband, Charles Blundell, Alexander Pritzel, and Benjamin Van Roy. Deep exploration via bootstrapped dqn. Advances in newral information processing systems, 29, 2016.
lan Osband, John Aslanides, and Albin Cassirer. Randomized prior functions for deep reinforcement learning. Advances in Newral Information Processing Systems, 31, 2018.
Georg Ostrovski, Mare G Bellemare, Aäron Oord, and Rémi Munos. fount-based exploration with neural PMLR, 2017.
density models. In International conference on machine learning, pp.
Deepak Pathak, Pulkít Agrawal, Alexei A Efros, and Trevor Darrell. Curiosity-driven exploration by self-supervised prediction. In International conference on machine learning, pp. 2778 2787. PMLR, 2017.
Docpak Pathak, Dhiraj Gandhi, and Abhinav Gupta. superiup exploration via disagrcement. In International conference on machine learning, pp. 5062-5071 PMLR, 2019.
Friedrich Pukelsheim. Optimal design of experiments. STAM, 2006.
Martin L Puterman. Markow decision processes: discrete stochastic dynamic programmingg. John Wiley & Sons, 2014.
Joaquin Quinonero-Canue Candela and Carl Edward Rasmusssn. A unifying view of sparse approximate gaussian process regression. The Journal of Machine Learning Research, 6:1939-100 1959,, 2005.
Rahul Rahaman et al. Uncertainty quantification and deep ensembles. Advances in Neural Information Processing Systems, 34:20063-20 20075, 2021.
Tom Rainforth, Adam Foster, Desi R Ivanova, and Freddic Bickford Smith. Modern bayesian experimental design. Stalistical Science, (9((1)1000-10)) 2024.
Carl Edward Rasmussen, Christopher KT Williams, et al. Gaussian processes for machine learning, volume Springer, 2006.
Stephane Ross, Brahim Chaib-draa, and Joelle Pineau. Bayes-adaptive pomdps. Advances in neural information processing systems, 20, 2007.
Christoph Salge, Cornelius Glackin, and Daniel Polani. Empowerment introduction. Guided Self-
Organization: Inception, pp. 67-114, 2014.
J Schmidhuber. A possibility for implementing curiosity and boredom in modcl-building neural controllers. In Proc. of the International Conference on Simulaaiion of Adoplive Behavior: From Animals to Animals, pp. 222 227, 1991.
John Schulman, Filip Wolski, Prafulla Dhariwal, Alec Radford, and Oleg Klimov. Proximal policy optimization algorithms. Xixiv preprint (xXi::1077.03344 2017.
Lorraine Schwartz. On bayes procodures. Zeitschrift für Wotrechetinicckkettttheorr und Verwandte Gebiete, 4:10-26, 1965.
18

---

Published in Transactions on Machine Learning Research (07/2025)
Ramanan Sckar, Olch Rybkin, Kostas Danillidis, Picter Abbeel, Danijar Hafner, and Deepak Pathak. Planning to explore via sell-supervised world models. In International Conference on Machine Learning, pp. 8583 8592. PMLR, 2020.
Claude Elwood Shannon. A mathematical theory of communication. The Bell system technicol journal, 27 (8):379-223 1948.
Pranav Shyam, Wojciech Jaákowski, and Faustino Gomez. Model-based active exploration. In International conference on machine learning, pp. 5779-5787 PMLR, 2019.
Jonathan Sorg, Satinder Singh, and Richard Lewis. Variance-basol rewards for approximate bayesian reinforcement learning. In Proceedings of the Twenly-Sishe Sirlh Conference on Uncerlainly in Arlificial Intelligence, pp. 564-571 2010.
Bradly C Stadie, Sergey Levine, and Pieter Abbeel. Incentivizing exploration in reinforcement learning with deep predictive models. arXiv preprint arXiv:1507.00814 2015.
Alexander L Strehl and Michael L Littman. An analysis of model-based based interval estimation for markov decision processes. Journal of Compuler and System Sciences, 74(8):1309- 1331, 2008.
Malcolm Strens. bayesian framework for reinforcement learning. In ICML, volume 2000, pp. 943- 943-950, 2000.
Richard S Sutton and Andrew G Barto. Reinforcement learning: An introduction. MIT press, 2018
Richard S Sutton, David McAllester, Satinder Singh, and Yishay Mansour. Policy gradient methods for reinforcement learning with function approximation. Advances neural information processing systems, 12, 1999.
Haoran Tang, Rein Houthooft, Davis Foote, Adam Stooke, OpenAI Xi Chen, Yan Duan, John Schulman, Filip DeTurck, and Pieter Abbeel. # exploration: A study of count-based exploration for deep reinforcement learning. Advances in neural information processing systems, 30, 2017.
Michalis Titsias. Variational learning of inducing variables in sparse gaussian processes. In Artificial intelligence and statistics, pp. 567-574. PMLR, 2009.
Joost van Amersfoort, Lewis Smith, Andrew Jesson, Oscar Key, and Yarin Gal. On feature collapse and deep kernel learning for single forward pass uncertainty. arXie preprint 1xxx:2102.11404 2021.
W. van der Vaart and J. van Zanten. Rates of contraction of posterior distributions based on Gaussian process priors. The Annals of Statistics, 36(3):1435 1463, 2008.
Aad Van Der Vaart and Harry Van Zanten. Information rates of nonparametric gaussian process methods.
Journal of Machine Learning Research, 12(6), 2011.
Yining Wang, Sivaraman Balakrishnan, and Aarti Singh. Optimization of smooth functions with noisy observations: Local minimax rates. Advances in Neural Information Processing Systems, 31, 2018.
Rebecca Willett, Robert Nowak, and Rui Castro. Faster rates in regression via active learning. Advances in Andrew G Wilson, Zhiting Hu, Russ R Salakhutdinov, and Eric P Xing. Stochastic variational deep kernel lcarning. Advances in neural information processing systems, 29, 2016a.
Neural Information Processing Systems, 18, 2005.
Andrew G Wilson and Pavel lxmailov. Bayesian deep learning and a probabilistic perspective of generalization. Advances in neural information processing systems, $3:4697-4700 2020.
Andrew Gordon Wilson Zhiting Hu, Ruslan Salakhutdinov, and Eric P Xing. Deep kernel learning. In Artificial intelligence and statistics, pp. 370-378. PMLR, 2016b.
19

---

Published in Transactions on Machine Learning Research (07 (07/2025)
Lisa Wimmer, Yusuf Sale, Paul Hofman, Bernd Bischl, and Eyke Htileemmier. Quantifying aleatoric and epistemic uncertainty in machine learning: Are conditional entropy and mutual information appropriate measures? In Uncertainty in Artificial Intelligence, pp. 2282 2292. PMLR, 2023.
Yuhong Yang and Andrew Barron. Information-theorrtiee determination of minimax rates of convergence. Annals of Statistics, pp. 1564 1599, 1999.
A Proofs of Propositions
In order to prove the convergence properties of IGa(-) presented in the main body's propositions, we need to define a couple of necessary building blocks concepts (Ghosal & Van der Vaart, 2017). For this section, let us rename the prior ristribution as T0), or shortly и, and the posterior as x (0((E)))) for simplicily, where E is a sequence of samples n E 2..} following likelihood distribution & ~ p(("(()) with E θ and 0 T(0). The posterior distribution associated with n-th sample, n. (B D*), is obtained via Bayes thcorem as
Je Pm (5"h
1) dT,, (0)
n (B (
where C C
J Pn (5"h
0) d (0)
as we assume that with increasing n, the data en provide more information about 0 E 
A.1 Proof of Proposition 3.1
We start with the notion of posterior consistency. For simplicity let us redefine Otrue = 0o. Posterior consistency loosely means that for n - co, the posterior T()) converges weakly the Dirac measure do in O and this depends on the topology of θ More formally, let assume that θ is a metric space equipped with metric d( then: Definition A.1 (Posterior Consistency). Given prior distribution x(0) E II, a posterior distribution (0((E)))) is said to be consistent w.r.t. a true parameter On E d) if
T (( E O: d(0, > Pag
as 00.
One important, consequence of posterior consistency that we are going to make use of relates to the consistency of estimators 8, T defined from the postcrior, and it is stated as follows:
Proposition A.2 (Estimators Consistency). Assume O* C θ is a subset such that posterior consistency
holds for 0g E O The following facts are true:
There erists 30, T(g") estimator that is consistent for Og E Or, i.e., diOn, 00) Pop when 8" ~ p((")00)
6) e is conver and d( is a bounded and conver distance function, then On = =T(5)) can be the posterior mean 0, = T({") (ddd,
M) If θ + θ is function continuous at 6, then d(g(O.). " o.
Convexity of θ indeed holds for the non non-parametric regression problem defined by Y; 0(x,) + te, where = C(X' and X E X. As for what requirements are needed for posterior consistency, we need intuitively that prior (0) do not exclude On from its support. Define the 'model' as the likelihood density function that generates samples e ~ Po E P, where P is a probability measure, and Po E P being the true density. Schwartz (1965) derived conditions for posterior consistency based on whether the true data gencrating model Po belongs to the KL support of the prior æ. In particular, we define
Definition A.3 (KL-support of p())). By denoting as (x q) the KL-divergence between distributions p and g, Po is said to belong to the KL support of the prior n, written po E KL(T), if
Ve 0, n mpp dxl p) <) 0.
20

---

Published in Transactions on Machine Learning Research (07 (07/2025)
Now, define U C P open neighborhood of Po according to metric d(.,), then intuitively have that the posterior x-1(5)) is also consistent if and only if for every open neighborhood U of po. x(U%))) + 0 (i.e., if all neighborhood around Po collapse to 0). This is formalized as follows (Schwartz, 1965; Ghosal & Van der Vaart, 2017). Given a neighborhood U C P, then we can test hypotheses Ho : p Pu versus H1 p E U. Assume these exist a real function 9n = (S1,..... En) E - [0. 1], representing the probability of rejecting Ho: such that Epp [px] + 0 and suPpEu [ (n]] + 0 (i.e., probability of rejecting Ho goes to zero if p = po: and conversely probability of rejecting H1 goes to zero when p + po). The idea of Schwartz (1965) theorem is that posterior consistency is guaranteed if the prior assigns mass that is 'arbitrarily close' to the true model Po E P and as grows can more correctly classify Ho vs H1. In addition, define Poe as the joint density of 5"h = (....) under data generating model Po- In full form then the weak consistency theorem reads:
Theorem A.4 (Schwartz (1965)). Assume that pa E KL(T), and that for neighborhoods Un C of po there are test. functions Pn satisfying the following requirements:
Ep S Be Be-bm
sup Ep (1 - (p) S Be Be
PEU
for some constants b, > 0, then π (()) A
O, or equivalently, +00(5)) S da, via corollary result.
Now we revert back to our specific dynamics model case, having true parameters Oo = -(o. E (O,d). Assuming the data generating model po described by equation (3), and assuming that po E (LT,,, we have
by Schwartz's theorem that (0((7)) "mp do, As stated in Proposition 2.1 this translate into
p(0 E θ : o)) (Ja, on)) > e (s,a)") F2g as n 00000
in case. As per Proposition A.2 reported above then know that g(O,)) = T from the posterior T((s")) that is consistent for the parameter s(0b) E e. Arbitrarily, can pick (f,)) to be exactly T(g") IGo(5",8t-1) Then, according to Proposition A.2, we have d( d(IGn(E", 1141), IGa, (5", 86+1)) " 0 or
construct estimator
e 0: d(Go(( IGo (5", > c) Pop O as - x0.
However, since convergence we have (0[{F) S 80.: then the equivalence p(O)E:) p(0)E st41) holds (i.e., new datum 8t-1 does not convey any additional information on 0 E θ). This implies that IGo, (E" sc+1)) H( H [e, = 0, thus the convergence result is H[p([][)] - H[p(0]s, s(41) Pop H H = 0, or similarly d d(IGo(E", 8(4)). " 0.
A.1.1 Convergence to Optimal Value Function
As a corollary of Proposition 3.1 proved above, we can straightforwardly show that the value function Vtaugm defined by the augmented rewards re rs + ner is such that Ve.angm Pa V*. More formally, define V*,augm as follows
VA aaggm (s, max re (s, + (, + L.g p(5'1s,0, 0, a)Vi, (')ds' do aEA
while the optimal value function associated with the original MDP's extrinsic rewards only is simplified to
/ p(s' p(s'1s, (s)
(x; max (x) + 1
Wrapping this up in a corollary statement:
Corollary A.5 (V=.augm Convergence). Under the conditions for posterior consistency of Proposition 2,1, it is easy to see. from Eq. (7) that we have vanaum (s, Pag Vi 0p), as rt IGo(.) Pig and F'g 0g-
21

---

Published in Transactions on Machine Learning Research (07/ 2025)
A.2 Proof of Proposition 3.3
The intuition behind Proposition 3.3 is simply that, following the statement of Proposition 3.1, if we can prove contraction rate e, holds for posterior convergence T(0|5") + 6b,, then in light of the proves in the
above section, both d[TGo(E", 0)) ""g and d(Vougm.*, V*) Fn happen at the same rate Gu- To prove Proposition 3.3, we begin by defining contraction rates for consistent posterior T0((F), where the true Oo E (0, d). Notice that posterior consistency always implies a contraction rate en. (Ghosal & Van der Vaart, 2017). We define
Definition A.6 (Contraction Rate). A posterior T0((() is said to contract to bg at the rate en - 0, if for every constant VM > 0:
. d(00:000, Men En) T when E, ~ (r(10))
A corollary of Proposition A.2 then straightforwarddy holds, with respect to posterior estimators 0, - T( convergence rate, and states:
Corollary A.7 (Estimator Contraction). If posterior =00()) contracts at rate en (or faster) to 8, E Oo c O then there erists on estimator On = T("") that is consistent for 0 E Oo and converyes at least as fast as En.
Assume now the dynamics model presented in Section 3.2 holds, that is:
S41 = [(ee: Ar) tt, where & ~ NNO0 ((0,o²), and E
and assume that (fo, op) E CO x [c, d], where: == x A is a compact subset of RR (.) is the class of continuous functions with finite Hölder norm of order a; and [c, d] c Rt. Notice that we assumed that S,A c R, so that the target variable $t+1 is single-task continuous S E I, but the following would hold also for multi-task settings, just individually and independently for each task. Finally, assume that f f~ GP C(y)),, GP prior where C(.-)) -) is the isotropic squared exponential kernel, i.e., C(x, = expl- = with length parameter 1 E R+ Then, one can prove that posterior T TJ. contracts at the optimal minimax rate up to a log constant, that is:
Theorem A.8 (van der Vaart & van Zanten (2008)). Assume (Fo, oo) E (A') x [c, d], where X is compact subset of R, the dynamics model in (3), and ~ GP(() C(.,))) prior where C(,,-) is the squared exponential kernel. Then, denoling p = |S x Al:
π (((f.o): d((f.o), (fo. o0)) > en) (T)
P
as n n 00
where en = n. (at170) (log n) with
(2+4a/p)
Notice that with ()ogn)'= 1, the above is equal to the minimax rate (best rate of estimation) for functions in the class C (X) (Yang & Barron, 1999). Since posterior consistency implics the existence of contraction rate En, then given that conditions for posterior consistency hold in Proposition 3.1, and given result in Corollary A.7, we have that IGo(E", st+1) + 0 at the same rate 6, Same hold for V. (s, 0) + (s; by extending Corollary A.5. Thus, this implies that under the condition of Theorem A.8 above, we have that (*48/0) (log n)*
IGo(5", s441) 0 and Va,augm (s.0) > (,, (s) at same rate 6,
Morcover, Van Der Vaart & Van Zanten (2011) show that the optimal minimax rate 6n = n T (Yang & Barron, 1999) for posterior contraction can be also achieve instead by imposing a further restriction on fo and assuming the covariance function C(..-l-) is a Matérn kernel. Define H (X) as the Sobolev space, then: Theorem A.9 (Van Der Vaart & Van Zanten (2011)). Given fo E where C(.++.)) is a Malérn kernel on X' = [0, 1JP, then the posterior x (f(l-) contracts rate en = n ()ा)
and f ~ -G,(,,C(,,)))
Thus, under the assumptions specified by Theorem A.9, also (Go(5" st41) + 0 and Vr, augrm - V (s) happen the same rate 6, (24p/0].
22

---

Published in Transactions on Machine Learning Research (07 (07/202) /025)
A.3 Brief Discussion Active vs Passive Learning Rates
The posterior contraction rates 6n derived in the results above all pertains the classical case of 'passive' learning (Yang & Barron, 1999). We note that these can be improved under some conditions with an active sampling strategy (Hanneke et al., 2014; Hanneke & Yang, 2015), which is what the work ultimately advocates for the realm of data-efficient exploration in MBRL. In Willett et al. (2005), the authors show that the minimax convergence rates for regression cases where fo is a piecewise constant function can be strictly reduced from c, = - (passive) to Gn = n 1 -1) (active). Castro & Nowak (2008) instead show that for classification problems with similar Hölder smooth decision boundarics the minimax lower bound convergence rate can be lightly improved from 6, n=/(x*p=1) (passive) En = n=((x++-2) (active), where p = (122) 1)/a and is the highest integer such that K < a. As a final example, results in Wang al. (2018) show that if fo is strongly smooth and convex (c.g., in Example 2 in their paper), with 2, can achieve much better rate of e, n -1/2 compared to the passive minimax, which in that example's case is n 2/1+
B Information Gain Derivations
In this second appendix section we include the full derivations of the Expected Information Gain (EIG) measures (Lindley, 1956; Bernardo, 1979; Rainforth et al., 2024) and the version of EIG under a Gaussian dynamics model and under a deep ensemble dynamics model. We assume the setup is the one described in the MDP definition in the main paper. The Information Gain (IG) of the full transition - composed by (st, aes St+1) is defined as
[Ga Ce St+1) = 11 [p(0 st, at] H [p( st, ar St+1)]
(8)
-pp10 1) [log p (0|se, des 8141)] (l000000..)))lp0(00000000
where p[0] St, at) x p(0)p(st+)) 10, 8t, ax), and where H[p(z)] = E| [o)](()] is the Shannon entropy (Shannon, 1948) and computing IGg(s; St, ae, $142) that can be done only reactive setting where 84+1 is actually revealed the agent. Thus in an active setting, we have to resort to an expected value surrogate version of it, EIGa(.) (-).
the set of dynamics parameters. As stated in main paper, the issue associated with
B.1 Expected Information Gain
As $1+1 is not revealed to the agent at time t, we can use EIGa(st, a.) = (ps(se+als) [IG a, 80-1) which essentially marginalizes over possible next 8441 E S, defined full form as:
EIGg(st, at) = Ep(ecsa (p(scsilsnag) IGo(se ae, st41)]
Ep(see (s)
Ep(a) E Dog s(+1)] (ogp
(9)
Now we note that we can re-write p(s1+1/5e, = p(5 ar s), ar) dl by marginalization, and that following Bayes theorem:
p(s (s,, at, 0)p(sr) or,
p(0 at) =
Jo p0 18,, a, 0)p( 8r, a, 0) d0
p(s1+1l (1c, ar 0)pp(1 a (f)
So p(50+1 St,, ar 0 pp (p) d0
23

---

Published in Transactions on Machine Learning Research (07 (025)
Using these two facts and applying them to (9) get:
Ep(ecsilennn [p(0ls p(0)s, at, S6+1) E p s. [log((0ss at]]]
p(a ar a1. 0)p(()
p(ec+1 (0t..r0)ppss...0)000))
110
log
a( (p)(()
p(a
p( e11p(())))))
Ep(0)(s 100,)))))
log
M (1 (P(Fe.
plar, (1p((1
p(St,
ppp0))(ersslec.cc..)) log pl 81+ 1 0t, 0r, log log
H [p(st+1| az)] - 1p(00 ( [4[p(s az; [z;
Epp (P az; 0) (pp()++1 (8,,,,))))))))))
by canceling terms out and re-ordering. Thus, we have now obtained an equivalent, specification of FIGg az) that can be computed using only the posterior predictive distribution po po(+++1|s, at), which marginalizes over parameters 0 E according to their posterior, ie., (0(4+41 [8,, a.) = Se po (8t411 a, ()0( 8t, at) d0. EIGa(s;, az) in this form can be interpreted as follows: it represents the expected reduction in the predictive uncertainty over the next state 8141 obtained from observing different set of parameters 0.
B.2
EIG under Gaussian dynamics
If the predictive posterior distribution (multivvriaate) Gaussian, such in the case of Gaussian Processes and Deep Kernels, we simplify EIGg(st, at) calculations even more. This is because if p(x) & (N, E), the entropy H [[(()) can be simplified (Cover, 1999)
as:
[ Ny,, Σ) log. (pg E),, dx
H Hpp(x)]
D
1
1
log 2n +
log (egd
[(x
--1]]] (1)
2
2
2
1
2 log 2nT + log (d) 2
2
log (E) + 2 log 2m)
where D the dimensionality of the multivariate normal, that in our case corresponds to D = =18], and where
the term E [xx u) -->--( 14)] is simplified follows:
E [x 1) Tx- 4n] = E [tr((x (1) ()) 1 -4)] =
E [tr(E=
-xx--((x-0 (x
H)((x -1) []]
= tr(B |3- (x-(x-1)) (4) Tx (x)-)) =)))=
= -r(E-(E)
tr([D)= =
D
Taking this simplification into account, we can in turn reduce ElGg(st, ar) calculations to the following:
Epp0))(mee 0)
[ (ps44(11 az) H [p( 1 1se, ae, 0)]]
log det (o a,)) log det (Vip(s++1ss, az 011]
Epp0ppeenn 2
log (Vee(V(ppssllll at)) Ep(o) log det (V[p(se+1 oz, 011))
where Vp(()) denotes the variance of distríbution p(.).
24

---

Published in Transactions on Machine Learning Research (07 (225))
B.3 EIG Disagreement of Gaussian Mixture
In settings where the dynamics model being employed does not yield full, Gaussian posterior predictive distribution such as in the case of GPs and DKs, then closed-forms EIG as described above cannot really be computed. This is, for example, the case where deep ensembles are utilized as a model for the dynamics. However, we can show that EIG in this setting can be computed as the Jensen Shannon- Divernnne among a Gaussian mixture consisting of all the M probabilistic neural network predictive distributions in an ensemble, i.c., (p(r, (s, a; mm))m-) This is derived as follows. First consider that:
EIGa ax) = e 0000) log ( ax; 0) log (8++1 182, ax)]
= 1 log ((t:111111 ar: 0); log p(5(+1|8), ax: d0] =
[ P (p(8 18, 44 0)] Ep(0)s. [M(p(8 [c; 0)]
where the term ax) [p($++1 ax; 0) is the expected predictive distribution over next states, which equals the marginal p(8++1 (s, a,) in the previous section's specification. Thus, Ep(0)s (ppes+pp((((1111 St. az; f) H H[p(st+1|s,, at)]. The other term ()))) H ([p(s+11 8t. ar; []]] is the expected entropy of the prodictive distribution for specific 0 values. Now, note that the JSD for set of M distributions with discrete weights Tm is defined as:
M Σ
M - Σ
JSD JSD aM (PI...,,PPl) = H
TmPm
Tm H [pm)
m==1
In our context, we can view the probability p(0|s;, dx) as continuous weights, so that we obtain:
EIGo az) H [L [0] (p (p(s 0)]] Bxt0)se.aa) (H [H[p(s+11ss)) an [c:
-SDD, (o 1s,, 0) 0 p(0) st, at))
C Additional Information on the PTS-BE algorithm
This section includes some additional information regarding the components featuring in the PTS-BE algorithm that we introduce in the main paper. In particular, we describe the specifications of all the different dynamics model that we tried in synergy with the proposed PTS-BE algorithm, namely Exact and Sparse Input versions of both Gaussian Processes and Deep Kernels, and then Deep Ensembles.
C.1 Models for the Environment Dynamics
The principal component in both "BE" type of algorithm, that employ retrospective intrinsic reward augmentation, and a PTS-BE one, that is based the plamiing-to-cxplore principle instead, is Bayesian model of the environment dynamics. This is because only Bayesian, uncertainty aware models allow for the incorporation of information gain measures. We review specifications of three classes of regression models: (Stochastic Variational) Gaussian Processes (Quinonero-Canell & Rasmussen, 2005; Rasmussen et al., 2006; Hensman et al., 2013), Deep Kernels (Wilson et al., 2016b;a) and deep ensembles (Lakshminarayanan et al., 2017; Rahaman el al., 2021). We briefly discuss their specific implications terms of posterior coverage uncertainty properties.
C.1.1 Gaussian Processes
Gaussian Process (GP) Regression models place a prior on the functional form of ʃ(-) GP(m(.), k(; )), such that p(f) & Nf m()), where m(-) is a mean function (c.g., constant, linear, etc.) and k((,)) kernel function (e.g., linear, squared-exponentiaa, matérn, etc.). Prior knowledge about f F can be efficiently conveyed through the choice of m(.) and especially k()))) (e-g., domain knowledge about the physics of the environment), which governs main features of the function approximator f(-) such as smoothness and sparsity. This feature ma makes it appealing in many applications where knowledge about the environment.
25

---

Published in Transactions on Machine Learning Research (07 (022)
4
ply)
3-
ply)
1-
0
-2
-2
-3
-3
4-
ply)
3-
ply)
1-
0
-2
-2
-3
-3
Deep Ensemble
4
q(f
3-
ply)
2-
1
22
22
-3
Figure 7: Comparison of uncertainty quantification models on a simple one dimensional regression example. The sample size of generated data points is n 100. The models considered are, starting from the top left corner towards the bottom right corner: i) Exact GP, trained on the whole sample; ii) Exact DKL, trained on the whole sample; iii) Stochastic Variational GP, trained on the ELBO loss over subset of 20 data points; iv) SV DKT, trained the same way SVGP; v) SV DKL coupled with dropout fully-connected ResNet structure and Bi-Lipschitz constraints to avoid fcature collapse; vi) Deep ensembles made of 5 neural networks models.
incorporated a priori. The joint density of (y, f) and the marginal likelihood, which used for training, in GP take the following form
ply. f(X) p
and ply|X) = pf
where all the quantities of interest can be computed analytically as the distributions are Gaussian. From this point of view, GPs have the advantage of being complementary the simplified version of EIGg at) derived above, as their predictive posterior distribution on 8t41 is in fact Gaussian as well. Besides GPs are well-known for their excellent uncertainty quantification properties. Their main drawback lies in their poor scalability, which is why we consider a Stochastic Variational approach where we learn a sub-sample of inducing points u {m (Hensman et al., 2013). We also considered using multitask/ multioutput kernel learning (Alvarez et al., 2012) but did not notice any significant improvements, while training costs were higher for the higher number of parameters involved in the multi-task kernels.
Throughout the experiments presented in the work, we utilize GP with constant prior mean function, m(.) 0 and with base squared exponential kernel k(., v) Crf(s, v) exp(-12||x y]
26

---

Published in Transactions on Machine Learning Research (07/202) /225)
R
R
R
R
R
R
2
4
Figure 8: Example of unichain environment made of L = 7 states. Arrows represent the possible actions in each state, state one is shaded to indicate that the agent is spawnod there as sp =
C.1.2 Deep Kernels
Deep kernels (Wilson et al., 2016b;a) are a generalization of GPs where inputs at), or just 8t, are first mapped to a (potentially lower dimensional) latent representation space h2, Jn : S x A + H or f + H. through deep learning architecture. Then, the function predict the next state, fa H + S or f A + S, can be assigned any GP prior (e.g., SVGP). The advantage of using deep kernels lies in the fact that they are better suited for dealing with high-dimensional state-action spaces thanks to their deep architecture, while they retain good function approximation and uncertainty quantification properties of GPs. On this matter, we specifically employ a version of deep kernels that uses a fully-connected ResNet neural net architecture (with dropout) for (-) and incorporates a bi-Lipschitz constraint on the transformations in f+()) (van Amersfoort et al., 2021), order to avoid the tendency of neural nets of exhibiting feature collapse (Guo et al., 2017). Note that the last GP layer in a deep kernel model scales by construction at a O(m H () training and Omm ) test cost, where [H] is directly controllable. Lastly, notice that (SV) deep kernel still allow one to use the simplified version of EIG, (st; ae) the latent space H. i.e., EhGi(he); as p(s Ihe) is modelled as Gaussian, as any component coming after the neural network block that models p(s behaves exactly as a Gaussian Process.
We have also considered implementing the MC dropout (Gal & Ghahramani, 2016) technique and apply it the deep neural network layers of deep kernels in order to get better generalization propertics, but did not notice any further improvement. MC Dropout consists in re-sampling pre-trained neural network with dropout layers K times, at test time, such that each prediction {()[))) represents sample from the predictive distribution o( I)), with different weights (W) and biases b (b)/
qly* x*)
ply x", W, p(W dWdbb.
MC Dropout approximates sampling from the Bayesian posterior predictive distribution, marginalizing out W, b. Perhaps the fact that MC dropout does not improve uncertainty quantification has to do with the fact that using sub-sample of inpul points through the SV inducing points procedure already prevents the deep kernels model from being over-confident and collapsing to the mean (Rahaman et al., 2021), resulting into better generalization propertics than standard deep kernels. However, we leave this topic to be investigated as part of future research.
Throughout the experiments, we consider a deep kernel model where the neural network architecture consists of two hidden layers of [32, 32] units, and a GP with constant mean and squared exponential kernel.
C.1.3 Deep Ensembles
Deep ensembles (Takshminarayanan et al., 2017) are arguably the most popular standard tool for uncertainty quantification in neural networks, thanks to their very good generalization properties compared to other sampling techniques such as MC dropout. The main difference with MC dropout is that deep ensembles requires to separately train (so they operate at train time) a pool of different neural network independently on the same data. Deep ensembles have been shown to have good out of sample uncertainty coverage (Wilson & Izmailov, 2020) in some examples, although they rarely match that of Gaussian Processes.
Some of the models, and their variations, considered for the environment dynamics are compared in the one dimensional regression example depicted Figure We use this simple example to depict each candidate model's s generalization and uncertainty quantification properties. Description of each model is enclosed in the figure's caption. We can see how the best models emerging from this small scale study are the Fxact.
27

---

Published in Transactions on Machine Learning Research (07 (07/2025)
Reactive
Active
1.0
Evnroyy zvee
1.
24E Deepfin
242 CP
Fracttonn of States Valted
2.4
Fracttinn or States Valtee
0.3
Model
Statew Fraction
-opll
035 + 0.07
/ Firor
0.22 = 0.05
VIME
0:34 0.05
BAE Deeplaa
0.99
BAB
0.97 0.04
2.4
0.2
0.2
0.2
Time stop
shop
Figure 9: Results on the 100 states unichain environment exploration task. The plots show the median cumulative fraction of states visited (solid line) logether with the 75th and 25th error bands, computed over 20 replications. The table instead reports the average fraction of states visited at termination (i.e., after 400 steps), together with the 95% confidence intervals.
GP, SVGP and SVDKL. In particular, SV DKL trains over 15 times faster than Exact DKL, and we notice that reducing the training sample to the 20 inducing points result in better uncertainly coverage in SV DKL versus Exact DKL.
Throughout the experiments in this work, we consider a deep ensemble of 5/7/10 neural networks, as we did not notice any appreciable improvement for ensembles larger than 10.
C.1.4 Policy Specifications in PTS-BE
In the PTS-BE algorithm, policy (*))) parametrizod by e E V is update by using the model generated trajectories instead of the sampled ones to allow for planning and looking ahead. The policy ( is typically parametrized by a neural network, and trained via gradient ascent, according to the policy gradient approach (Kakade, 2001). In the discrete action settings above, we typically make use of Proximal Policy Optimization (PPO) algorithms (Schulman et al., 2017). The standard PPO loss reads:
CPPO( (+) = E (F) c1
C. Value + (2
where 2 is the clipped surrogate objective, Caauu is squared loss on the value function V(-) and H (s)() is a policy entropy term used to ensure sufficient exploration (for more details see Schulman et al. (2017)).
In the continuous action settings, we instead use Soft Actor-Critic (SAC) (Haarnoja et al., 2018). SAC updates the policy parameters by breaking down the policy improvement into separate Actor and Critic updates, according the modified objective:
J(T) EYax
Σ (r(sss))))) + (H(+s())
where the entropy term (H)))))))) encourages exploration by keeping the policy as random as possible. The actor and critic networks, and the "temperature", are trained to minimize the following losses:
I(Vactor) Eia (als)) Quornaa (s, as,
(10)
L(Voritie) E (Q (3, a) - (r + 7Q (s', log ( (adl(()))))
(11)
L(a) E log( ( Tester (als)) aHttarge)
(12)
where the temperature parameter is updated to achieve target entropy H
Htarget-
28

---

Published in Transactions on Machine Learning Research (07/202) 2025)
P  D
Figure 10: Visual 3D renditions of three types of mazes that can be obtained from the 2D Maze environments via customization of the maze struct i) a u-shaped maze such as the one employed the experiments above; ii) medium sized maze with ledges and obstacles in the middle, again similar to the one emplyed in the experiment; tii) larger maze with just ledges the sides. The rod ball denotes the goal object.
D Additional Experiments
D.1 The Unichain Environment
In this last supplementary material section we include brief description of the Unichain enviroment, plus additional results on it where we consider a longer chain of T, = 100 states, without noise this time.
As outlined the main paper, the unichain environment consists in a simple sequence of L Markov sequentially connected states, where the goal is usually to explore all the possible states and reach the final one (Puterman, 2014; Osband et al., 2016). We consider discrete action space defined A {go left, stay, go right). If the agent is in the first or last state and tries to go left or right, it automatically hit wall and remains in that stale. The agent is initially spawned in the second state 8q = 1 (stale counter starting from 0). The reward function rf; : S x A + R is sparse and assigns 0.001 to visiting the first state, 1 to visiting the last state, and otherwise, i.c.:
0.001
if 8: 0 if s; = otherwisc.
r r"(s ax) =
0
The sub-optimal reward associated with state 0 represents a 'reward trap' for algorithms based solely extrinsic reward exploitation, as it acts as disincentive for the agent to move elsewhere and explore. Thus, strong exploration bonus is needed to move it away from there. A simple visual representation of a L = 7 state unichain environment, together with the possible actions is depicted in Figure 8. The shaded state so = 1 represents the spawning location.
D.2 Additional Experiments on the 100 States Unichain Environment
In this subsection, we present the additional results on the L = 100 unichain environment, but without added noise this time. The task is again purely exploratory, and the methods compared are a subset of those compared in the main paper: i) PPO policy entropy Hmm (sx) regularizer (Schulman et al., 2017) (pi-Entropy); ii) PPO with & prediction error as reactive intrinsic reward (Stadic ct al., 2015; Pathak et al., 2017) (e2 Error); iii) PPO with VIME, i.e., I.a(.) as reactive intrinsic reward coupled with Bayesian Neural Network dynamics (Houthooft et al., 2016) (VIME); iv) PTS-BE with Deep Ensembles (Shyam et al., 2019); v) PTS-BE with SVGP (PTS-BE-GP): vi) PTS-BE with Deep Kernels (PTS-BE-DK).
We measure performance of the methods again via the cumulative fraction of visited states al each time slep f, and the final fraction coverage at episodic termination, which we set to be after 400 steps in this case, as the Markov chain of states is longer. PTS methods are left running for 10 steps initially in order gather enough data estimate po(sttils,,d) a,). Results over 20 seeded replication of the experiment are reported in Figure Iss plots and table. Similarly the unichain L 50 states environment results presented in the main
29

---

Published in Transactions on Machine Learning Research (07/2025)
paper, PTS methods consistently outperform the rest, they require approximately 150 steps only to reach complete coverage of the environment. and solve the task.
D.3 Additional Details the Medium Sized Maze
Finally, report here few extra information the Maze environments. These continuous control environments with (S C R2, A R The reward function is sparse, the agents receive rī = 1 only when they reach the objective, otherwise they receive rf = 0. The environments are 2D, and their state space S features æ and y coordinates of the agent and x and coordinates linear velocity of the agent. The contimuous action space A instead include coordinate's linear force applicable to the a and y directions.
These environment are highly customizable in terms of their shape and size. We experiment with different shapes and sizes of increasing difficulty. Figure 10 above reports visual 3D rendition of some possible maze configuration achieved by customizing the map structure.
30


Reviews:
[{"review_id": "lLQICPgscU", "reviewer": "Reviewer_EccY", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "**Strengths**\n\n* The paper is mostly clearly written.  \n* The proposed exploration bonus is theoretically motivated and grounded.\n\n**Weaknesses**\n\n* While the theoretical foundation is appealing, the paper does not demonstrate that the method remains effective in more complex or high-dimensional environments.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "No", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "This paper proposes a new class of exploration bonuses aimed at improving exploration efficiency in reinforcement learning. These bonuses are designed to approximate epistemic information gain and gradually converge to zero as the environment becomes fully explored. Building on this idea, the authors introduce an algorithm called *Predictive Trajectory Sampling with Bayesian Exploration (PTS-BE)*, which combines model-based planning with the proposed bonuses. Experiments on several sparse-reward and pure exploration environments show that PTS-BE outperforms a range of baselines.", "claims_and_evidence_yn": "No", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "* Some notations in Section 4 are confusing and could be clarified. For example, in \\( p_\\theta(s_{t+1} \\mid s_t, a_t; \\theta) \\), what is the intended meaning of \\(\\theta\\) appearing both in the subscript and after the semicolon?\n* In the experiments in Section 5.1, does the random policy eventually fully explore the environment? What is the coverage achieved?\n* In Figure 3 (middle), why is performance of $H(\\cdot)$ in the *Homosk* setting lower than in *Heterosk*? Intuitively, the heteroskedastic case should be more challenging.\n* The implementation of the \"BE\" baseline in Section 5.1.1 is not entirely clear. How is the reward augmentation carried out? It would be helpful if the authors could elaborate on this.", "broader_impact_concerns": "None", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "uyknchJ3Ex", "reviewer": "Reviewer_Y3kG", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strengths:\n1. The aforementioned contributions seem interesting to me (although I am not sure all of them are contributions of this paper).\n2. The experiments are well-designed and the results are interesting.\n\nWeaknesses:\n1. The relation with prior works is not clear (I acknowledge that I am not familiar with the Bayesian exploration literature).\n    - It seems that this intrinsic reward is not new [[Latyshev & Panov, 2025][1]; [Shyam et al., 2019][2]]. This is not clearly stated in the related work section. Could you clarify how you differ from this line of work? Is your contribution to prove interesting properties for these existing intrinsic rewards using information-theoretic arguments?\n    - It seems that model-based planning approaches were already used in conjonction with such model-based intrinsic rewards. This is not clear in your related work section. Could you clarify how your planning algorithm (PTS-BE) differs from previous model-based approaches? [[Shyam et al., 2019][2]; [Sekar et al., 2020][3]; [Mendonca et al., 2021][4]]\n2. There seem to be a clash between the presented framework (Bayesian-adaptive MDP) and the method.\n    - The BAMDP seem to provide samples from the posterior, while the algorithm seem to maintain the full posterior (and not to sample from this posterior), at least in the case of Gaussian processes. Could you better link the actual dynamic models update to $p(f | D)$ at line 14 in Algorithm 1 with the BAMDP's posterior distribution and its samples? I also think it could be interesting to better explain how the three methods (GP, DK, DE) compute $p(f | D)$ (or sample $f \\sim p(\\cdot | D)$) from the dataset.\n3. The best methods are PTS-BE-{GP,DK,DE}, which suggests that most of the gain comes from the usage of a model-based planning approach.\n    - I would be interested to see the ablation \"PTS\", a model-based approach with PPO, which would correspond to Algorithm 1 without intrinsic rewards.\n\nMinor:\n- Many equations have equality signs both at the end and at the beginning of lines (e.g., equations 6).\n- Line 9 in Algorithm 9, there is a typo: $r_t^e$ instead of $r_{t+j}^e$.\n\n[1]: https://link.springer.com/article/10.3103/S0147688224700370\n[2]: https://proceedings.mlr.press/v97/shyam19a.html\n[3]: https://proceedings.mlr.press/v119/sekar20a.html\n[4]: https://proceedings.neurips.cc/paper/2021/hash/cc4af25fa9d2d5c953496579b75f6f6c-Abstract.html", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "- Propose a new intrinsic reward that dynamically targets epistemic uncertainty (actually, the IG/EIG intrinsic reward was already proposed [[Shyam et al., 2019][2]], right?).\n- Show that the intrinsic reward converges to zero in probability if the learned model converges to the true model in probability.\n- Characterize the rate of convergence to zero in the special case of Gaussian processes.\n- Show empirically that the intrinsic reward targets the epistemic noise, by ignoring the aleatoric noise in the case of heteroskedastic noise.\n- Propose to use the learned model for planning in a model-based method (PTS-BE).\n- Show that it works better than just using the learned model for the intrinsic rewards with a model-free method.\n- Show that Gaussian Processes and Deep Kernels, which are more computationally costly, can show limited gains with respect to Deep Ensemble models.\n\n[2]: https://proceedings.mlr.press/v97/shyam19a.html", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "- If the intrinsic reward is not new (see weakness 1), please adjust your claims. Your abstract, introduction and conclusion really seem to imply that EIG is introduced for the first time in this paper. Generally speaking, I would appreciate the paper to be a little more precise on its contributions. I found difficult understanding what was existing in the literature and what was new.\n- I would appreciate the ablation study \"PTS\" (see weakness 3) and a discussion of the source of performance of the algorithm (is it the model-based planning routine, or the intrinsic reward?). That said, I acknowledge that it is intersting that your intrinsic reward provides a model of the dynamics \"for free\".", "broader_impact_concerns": "None", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "tTPEWeFqYh", "reviewer": "Reviewer_G9j7", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "## Strength\n\nThe paper is very well written and was a pleasant read.\n\nAs elaborated in the Contributions, unlike many existing methods, the authors define information gain (IG) under the RL setting that explicitly capture epistemic uncertainty, encouraging exploration where it is present. The proposed IG signal diminishes over time (Sections 3.2, 3.3), is made computationally tractable and scalable (Section 4), and can be incorporated into a wide range of RL methods, including on-policy, off-policy, and model-based approaches (Sections 4 and 5). The experiments are carefully designed and support the claims (e.g., Figure 1, Section 5).\n\n## Weakness\n\nVisitation-frequency-based methods also explicitly model epistemic uncertainty, diminishing with non-asymptotic rates, but this connection is not discussed in depth.\n\nThe novelty compared to prior work could be more explicitly highlighted. While this paper proposes several new ideas, it also draws inspiration from multiple fields. Clarifying which parts are directly adapted, which are inspired, and how, would better showcase its novelty.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "This work addresses data-efficient exploration in reinforcement learning (RL), a well-recognized challenge in the field. Efficient exploration depends highly on the ability to distinguish between epistemic and aleatoric uncertainty, encouraging exploration only in regions where epistemic uncertainty is high.\n\nThe authors propose using information gain (IG) as an intrinsic reward signal, which, unlike many existing approaches, managed to target only the epistemic uncertainty and naturally diminishes to zero over time (Section 3.1). Convergence guarantees are established by leveraging results from prior literature on information gain (Sections 3.2 and 3.3). Inspired by research on GPs and BNNs, the authors develop scalable and computationally tractable approximations to expected information gain (EIG), making the method applicable to complex RL settings, including both classic on-policy/off-policy algorithms for continuous control and model-based RL (Sections 4 and 5). Through illustrative experiments, they highlight the importance of distinguishing epistemic from aleatoric uncertainty - particularly under heteroskedastic noise - and show the effectiveness of their EIG-based approach (Sections 4.1 and 5.1).", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "1. Relation to Visitation-Frequency-Based Approaches: You should discuss more explicitly how your method relates to visitation-frequency-based exploration strategies. Many RL algorithms with regret guarantees fall into this category, where the intrinsic reward relates to visitation counts and can be interpreted as a form of epistemic uncertainty. Importantly, such intrinsic rewards also diminish with visitation at a non-asymptotic rate.\n\n2. Clarifying the Novelty Compared to Prior Work: It would strengthen the paper to more clearly highlight the novelty of your approach relative to prior works, e.g., Houthooft et al. (2016), which you mention right before Section 3.2. How does their use of information gain as an intrinsic reward compare to or inspire your method?\n\nFormatting and Minor Corrections\n1. Figure 4: Colors are hard to distinguish.\n2. Page 3: In the Bellman equation, the $V$ on the RHS should involve $p_\\tau$.\n3. Page 7: Remove the \"=\" at the end of each equation line.\n4. Appendix Reference: Directly cite specific appendix sections instead of saying \"provided in the appendix.\"\n5. Algorithm 1, Line 8: Does $\\mathcal{D}_{\\text{model}}$ includes all $k \\in [K]$ trajectories?", "broader_impact_concerns": "I don't immediately see any concerns.", "additional_comments": "", "extra_scores": {}, "extra_text": {}}]

Author rebuttal:

Decision: Accept (poster)