Efficient Model-Based Reinforcement Learning Through Optimistic Thompson Sampling
Authors: Jasmine Bayrooti
Carl Henrik Ek
Amanda Prorok
Venue: ICLR 2025 Poster (2025)
Keywords: reinforcement learning
model-based reinforcement learning
optimistic exploration
Primary area: reinforcement learning
TL;DR: We introduce a theoretically-grounded approach to optimistic exploration that leverages joint uncertainty over states and rewards for improved sample efficiency.

Abstract: Learning complex robot behavior through interactions with the environment necessitates principled exploration. Effective strategies should prioritize exploring regions of the state-action space that maximize rewards, with optimistic exploration emerging as a promising direction aligned with this idea and enabling sample-efficient reinforcement learning. However, existing methods overlook a crucial aspect: the need for optimism to be informed by a belief connecting the reward and state. To address this, we propose a practical, theoretically grounded approach to optimistic exploration based on Thompson sampling.  Our approach is the first that allows for reasoning about _joint_ uncertainty over transitions and rewards for optimistic exploration. We apply our method on a set of MuJoCo and VMAS continuous control tasks. Our experiments demonstrate that optimistic exploration significantly accelerates learning in environments with sparse rewards, action penalties, and difficult-to-explore regions. Furthermore, we provide insights into when optimism is beneficial and emphasize the critical role of model uncertainty in guiding exploration.

Published as a conference paper at ICLR 2025
EFFICIENT MODEL-BASED REINFORCEMENT LEARN- ING THROUGH OPTIMISTIC THOMPSON SAMPLING
Jasmine Bayrooti University of Cambridge jgb520 cam. ac uk
Carl Henrik Ek University of Cambridge Karolinska Institutet che2 90cam. ac uk
Amanda Prorok University of Cambridge asp450cam. ac. uk
ABSTRACT
Learning complex robot behavior through interactions with the environment ne- cessitates principled exploration. Effective strategies should prioritize exploring regions of the state-action space that maximize rewards, with optimistic explo- ration emerging as a promising direction aligned with this idea and enabling sample-cfficient reinforcement learning. However, existing methods overlook a crucial aspect: the need for optimism to be informed by belief connecting the reward and state. To address this, propose practical, theoretically grounded approach to optimistic exploration based on Thompson sampling. Our approach is the first that allows for reasoning about join! uncertainty over transitions and rewards for optimistic exploration. We apply our method on a set of MuJoCo and VMAS continuous control tasks. Our experiments demonstrate that optimistic ex- ploration significantly accelerates learning in environments with sparse rewards, action penalties, and lifficult-to-explon regions. Furthermore, we provide in- sights into when optimism is beneficial and emphasize the critical role of model uncertainty guiding exploration.
INTRODUCTION
Reinforcement Learning (RL) is recognized as a promising approach to solving complex sequential decision-making tasks in robotics (Levine et al. 2016). However, many popular RI. algorithms require millions of interactions with the real world to train effective policies (Schulman al. 2017] Mnih et al. (015). This sample inefficiency is primarily due to the central challenge of balancing exploration, gathering information about the world, with exploitation, maximizing rewards given current knowledge (Sutton & Barto 2018 Szepesvári 2022). Sample inefficiency is especially prohibitive for RL. applied to robotics due to the high cost and potential wear on physical systems.
Efficient RL requires careful balance between exploration and exploitation, and some approaches are better suited for this than others. Our work centers on model-based RL, a class of algorithms where agents build a predictive model of the environment dynamics from which they derive policy. Many works have demonstrated impressive sample efficiency using model-based RI. by leveraging the learned dynamics model to simulate future outcomes, thereby allowing agents to plan strategi- cally and optimize the controller using fewer real-world trials (Ibarz et al. 2021 Chua et al. 2018
Janner et al. 2019; Yang et al., 2020 Sekar et al. 2020, Hansen al. 2024). However, learning a model of the environment is a double-edged sword: a faithful model enhances efficiency by re- ducing the number of interactions needed in the real world while biased model severely hampers performance by misguiding the policy (Gu et al. (016).. Planning with uncertainty-aware dynam- ies reduces the effects of model errors (Deisenroth & Rasmussen, 201T Deisenroth et al., [013)] Gal et al.) 2016) and allows model-based RL algorithms to reach state-of-the-at asymptotic perfor- mance on benchmark control tasks (Chua et al. 2018 Janner et al.) 2019). Furthermore, by learning probabilistic dynamics model, model-based RL enables directing exploration in principled man- ner towards unseen parts of the state-action space.
In this work, we focus optimistic exploration using model-based RL with uncertainty aware model. In a subtle departure from classical exploration, which prioritizes transitions with high uncer- tainty under the model, define optimistic exploration favoring transitions likely to yield high

---

Published as a conference paper at ICLR 2025
rewards. This approach aligns well with the RL objective of maximizing cumulative rewards (Sut- ton & Barto] 2018). Notable works unite optimistic exploration with uncertainty- -wware awaae model-basee RL to hallucmate plausible, optimistic training experiences for superior robustness and sample ef- ficiency (Curi et al. 2020; Sessa et al., 2022), however these formulations ignore joint uncertainty between the reward and dynamics. They also assume knowledge of a reward function, which can be impractical in real-world applications due to the difficulty of specifying dynamic, variable, and sub- jective reward functions accurately a priori. While learning the reward function a straightforward remedy, there is a dearth of efficient and principled approaches to integrating reward learning into optimistic exploration for model- RL.
Contributions. Our primary contribution is the first practical model- RL algorithm, called Hallucination-basee Optimistic Thompson sampling with Gaussian Processes (HOT-GP), that can be used with state-of-the-an ofrt off-policy RL algorithms for principled optimistic exploration. In- spired by the premise that efficient exploration should be informed by joint uncertainty over state and reward distributions, we propose GP model maintain this joint belief in order to simulate transitions that are simultaneously plausible under the learned dynamics and optimistic with respect to the estimated reward. We evaluate this approach
set of MuJoCo (Todorov al., 2012) and VMAS (Bettini et al., (022) continuous control tasks. The results reveal that HOT-GP matches or improves sample efficiency on standard benchmark tasks and substantially accelerates learning in challenging settings involving sparse rewards, action penalties, and difficult-to-explore regions. As these findings validate our hypothesis, our secondary contribution is an empirically-supporre ment for maintaining a belief over both the reward and state distributions, as this is key for effective optimistic exploration. Moreover, we study the factors that influence the utility of optimism and underscore the important role of model uncertainties in shaping exploration.
argu-
2 RELATED WORK
Model-based RL methods are promising for complex real-world decision problems due to their potential for data efficiency (Kaelbling et al. 1996). Selecting the appropriate model is critical, it must facilitate effective learning in both low-data regimes (during early stages of training) and high- data regimes (later in training). This requirement naturally aligns with the advantages of Bayesian models. Non- Non-parametric models like Gaussian processes (GPs) provide excellent performance in settings with limited, -iwddmensiioall data (Deisenroth et al.] 20113 Deisenroth & Rasmussen [011] Kamthe & Deisenroth] 2018). Neural network predictive models have also been effectively used for tasks with non-smooth dynamics (Nagabandi et al., 2018) and high dimensions like the Atari suite (Oh et al. 2015) Kaiser et al.) 2020). As deterministic neural networks tend to overfit to data in the early stages of training, uncertainty-awarc neural network models have been shown to achieve better performance with algorithms such PETS (Chua et al. 2018) and MBPO (Janner jet al. 2019). However, designing scalable Bayesian neural networks remains an open problem (Roy 2024 Guo et al.) 2017|[0sban (2166) and popular approximations using dropout (Gal a
ee al. 2017 and ensembles (Lakshminarayanga et al. 2017] Chua et al. 2018) fall short of the precise uncertainty quantification that GPs provide. In order benefit from both the high-capacity function approximation of neural networks and probabilistic inference of GPs, in this work we model the joint dynamics and reward functions using a GP with a neural network mean function (Iwata & Ghahramani 2017).
Thompson sampling provably efficient exploration algorithm in RL (Thompson 1933). This approach implicitly balances exploration and exploitation by sampling statistically plausible out- comes from the posterior distribution over dynamics models and selecting actions based on the sampled outcomes (Russo & Van Roy| 2014 Russo et al. 2018). Thompson sampling has been studied on tabular MDPs (Osband et al. (013) and extended in theory to continuous state-action spaces with sublinear regret under certain conditions (Chowdhury & Gopalan 2019).
Optimism the face of uncertainty another theoretically grounded approach to guide explo- ration that selects actions that optimize for optimistic outcomes. The motivation is to promote exploration guided by beliefs about future value to facilitate efficient learning. While optimism has
been extensively studied in tabular MDPs (Brafman & Tennenholtz 2002 Jaksch et al., 2010) and linear systems (Jin et al., 2020 Abbasi- Yadkori & Szepesvári] 2011 Neu & Pike-Burke 2020), optimism in continuous state-action MDPs with non-linear dynamics remains less explored. GP-
2

---

Published as a conference paper at ICLR 2025
UCRL (Chowdhury & Gopalan 2019) is such theoretical approach that targets dynamics mod- els that lie in a Reproducing Kernel Hilbert Space. Another is LOVE (Seyde et al. 2021). which uses ensemble of networks to direct exploration towards regions with high predicted potential for long-term improvement. TIP (Mehta et al. 2022 also employs optimism by selecting actions ac- cording optimistic acquisition function. Curi al. (2020) proposes reparameterization of the dynamics function space in order to reduce optimistic exploration to greedy exploitation with the H- UCRL algorithm. H-UCRL has sublinear regret bounds under certain conditions and demonstrates efficient exploration on deep RL benchmark tasks with action penalties. The H-MARL algorithm extends this the multi-agent setting for general-sum Markov games and introduces practical approximation for generating optimistic states (Sessa et al.) 2022). While these methods all de- rive useful optimistic exploration strategies for RI., they ignore the relationship between uncertainty around dynamics and uncertainty around the reward. Moreover, TIP, H-UCRI, H-MARI. use the ground-truth reward function throughout training. Our approach, inspired by H-UCRL (Curi et al., 2020), improves optimistic hallucination for model-based RL by learning and leveraging the joint uncertainty over state transitions and associated rewards.
Reward learning is commonly applied in RL for reward shaping (Hu et al., 2020 This is closely related to exploration, as intrinsic rewards learned through curiosity-driven methods help guide ex- ploration by reducing uncertainty over the agent's knowledge of the environment (Pathak et al.) 2017]. Reward learning also plays an important role in model-based RI. Receding-horizz plan- ning approaches like PETS (Chua et al. 2018), Dreamer (Hafner et al., 2020; 2219 2023), and TD-MPC (Hansen et al., 2022 2024) use estimated rewards to evaluate and optimize sequences of actions, while model-based policy learning approaches like MBPO (Janner al. 2019) use esti-
mated rewards to inform policy updates. Some of these approaches explicitly account for uncertainty the learned reward function (Chua et al.) 2018 Janner et al., 2019 Chowdhury & Gopalan 2019))
however they treat it independent from the uncertainty in the dynamics. We argue that learning an uncertainty-ave -ware reward-dynamics function and strategically leveraging the joint uncertainty leads to the most principled and effective optimistic exploration.
3 PROBLEM STATEMENT
We consider a stochastic environment with states E S C IRP and actions a E A C R9. Given the current state 8, and action az, the agent transitions to the next state 8t+1 and incurs reward r with
Sc+1 a()
[(( (r())
for transition dynamics S S x A + S and reward function fr : S x + R. Our objective is to learn optimal control for this system within episodes of finite time horizon T.. To control the agent, we learn a deterministic or stochastic policy that selects actions given the current state. For ease of notation, we write the policy deterministically as π S + A so that a . (T:)). We consider a specific transition dynamic ft and reward function f.. denoted with hat to indicate realizations that are drawn from random functions. The performance of the policy is measured by the sum of accumulated rewards during an episode,
fT. - Σ (.loer a.)) s.t. a. = ===(()), 8641 = fe(se, or).
T
The optimal policy n² maximizes performance over the true dynamics and reward function with,
π
arg max (TT, f fr).
(3)
mell
In model-based RL, the true dynamic fe is unknown and must be learned by interactions with the environment. In our setting, the ground-truth reward function fr is also unknown.
4 BACKGROUND
In this section, we introduce relevant background on model-based reinforcement learning and exist- ing exploration strategies.

---

Published as a conference paper at ICLR 2025
4.1 MODEL-BASED REINFORCEMENT LEARNING
When systems have unknown dynamics, we consider the dynamics fr to initially be random. There- fore, learning model of the environment dynamics entails fitting an approximation p(J+) over the space of dynamics functions, given observations from the real system. The same approach can be used estimate the true reward function fr. We consider model-based RL algorithms that, on each iteration of training, roll out the current policy π on the real system for an episode and update p(f) using the newly observed transition data. Standard approaches use the model to simulate sequential transitions and use this data to update the policy (Suttoon 1990). We opt to update the policy based on short k-length model-generated rollouts branched from the state distribution of a previous policy under the true environment, as recommended by Janner et al. (2019) to reduce compounding model errors over extended rollouts. Algorithm gives an overview of the process.
Algorithm Model-based Policy Optimization
1: Require:
max environment steps N, model rollouts M, steps per rollout T, steps per model rollout K, initial state distribution d(ss)), policy-learning algorithm Poli ccyearcc
2: Initialize: policy π, reward-dynamics model p(f), environment dataset Dew
3: while [Ponv| < N do
4: /*Simulate Data*/
5: Initialize model dataset D Dmodel 0
6: for m 1, M do
7:
Sample 8g uniformly from Denv for k 1,..., K do Compute action ā from o7 (5x)
8:
9:
10:
Sclect next state 5 and reward fg using p(f Denv) This is algorithm-speciffe
11:
Append transition to buffer (&k, åg, $1+1, fx) - Dmnntel
12:
#Optimize Policy*/
13:
π < PoPiiccyy PolicySearchhff Pmode)
14:
#Optimize Dynamics Model*/
15:
Start from initial state $e ~ d(su)
16:
for f 1, ..... T do
17:
Compute action ae from (st)
18:
Observe next state and reward Tx [ 0e) Append transition to buffer (8t, dz, 8641, r') + Den
19:
20:
Retrain model p(f) using Denv
Model Design. We use an uncertainty-aware model to represent the prior distribution of dynamics models p(fe) aligned with the 1-step transition data in Danv- Probabilistic models allow for exploit- ing correlations in observed data by generalizing to unseen states and actions. Popular approaches include probabilistic ensembles of neural networks, which provide mean (u(s,), and confidence timates (( a) by averaging outputs over multiple trained models (Lakshminarayanar et al. 2017] Chua al. 2018). Bayesian inference with neural networks is computationally expensive and is an open question if parameter uncertainties lead to relevant posterior function estimates (Sharma al. 2023 Roy et al. 2024p. Alternatively, GPs provide distributions directly over the function space, enabling direct estimation of the posterior distribution p(f: Denv) over dynamics models. Since GPs incorporate strong and interpretable prior beliefs, they naturally capture uncertainty due to limited data (epistemic uncertainty) and variability in the data (aleatoric uncertainty) (Rasmussen) & Williams 2006), which contributes to GPs effectiveness in low-data regimes. priors, de- fined by mean and covariance function, describe how the function varies around the mean. With flexible mean functions, GPs offer principled uncertainty estimates for complex functions (Iwata & Ghahramani 2017 Fortuin al. 2019 Fortuin 2022).
4.2 EXPLORATION STRATEGIES
is challenging to learn the -dwaaddd model while solving for optimal control. Performance depends the utility of model- data for updating the policy and likewise the value of real experiences gathered by the policy for updating the model. Solving this problem efficiently requires balancing exploitation of current knowledge with exploration of the state-action space.
4

---

Published as a conference paper at ICLR 2025
The mechanism for selecting the next mudel-ggeneatie state plays important role, with different strategies incorporating varying degrees of uncertainty to encourage exploration. In this section, we provide overview of the relevant exploration strategies, categorized in Curi et al. (2020).
Greedy Exploitation. Most model-base based RL approaches approximate p(ft Deny) and maximize the policy greedily over the uncertainty in the dynamics model. They use this greedy policy:
Greedy
arg max (p... (T, fc.
(4)
For instance, GP-MPC (Kamthe & Deisenroth 2018) and PETS (Chua al. 2018) are MPC-based methods that use greedy exploitation and represent fl Den with a GP and ensemble of neural networks respectively. Similarly, PILCO (Deisenroth & Rasmussen 2011) uses GP with moment matching to estimate p(fe Denv) and greedy exploitation to optimize the policy. While such algo- rithms can converge optimal policies under certain reward and dynamics structures (Mania al., 2019). greedy exploitation is generally inefficient for exploration.
Thompson Sampling. This theoretically-groummee between exploration and exploitation (Russo et al. 20018) Chapelle & Li, 2011). Under Thompson sampling (Thompson 1933), the policy is optimized with respect to a single model sample f on each episode,
approach that offers principled balance
f~p(f Denr);
TS
arg max J(T, fc).
=ETI
This optimization problem is equivalent to greedy exploitation (equation 4) after sampling f ~ p(f Denc). Under certain assumptions, Thompson sampling satisfies strong Bayesian regret bounds (Os-
band & Van Roy| 2016 2017), theoretically positioning it as the optimal strategy.
Hallucinated Upper-Confidenc Reinforcement Learning (H-UCRL). This is tractable version of the UCRL algorithm (Jaksch et al.] 2010), which optimizes an optimistic policy over the set of statistically plausible dynamics models. Curi et al. (2020) propose reparameterizing plausible dynamics models in order to optimize over a smaller class of variables. Specifically, they write dynamics functions as fi(s, a) = =u(s, a) + 3 E(s, a)n((s for some function n R" + [-11, 1] and optimize over n7-)) instead. As this problem is still challenging optimize, Sessa et al. (2022)
propose sampling n² ~ Uniform([-1, 1jP) for 1, 1...., 7 and updating the policy with
π
I.UCRL
arg max
J(T f., f.). s.t. fe(s, a) = p(s, a) + + a)s,a)p HEll qp JEJ......
max
(6)
Unlike greedy exploitation, this approach effectively optimizes an optimistic policy over the modi- fied dynamics f. in equation does not model or utilize the joint uncertainty over transitions and rewards.
6 although the reward function is assumed to be known. H-UCRL also
5 THE HOT-GP ALGORITHM
We introduce a new exploration strategy based on the insight that principled optimism should reason about uncertainty over transitions and rewards jointly. Importantly, this requires a joint model of the transition dynamics and reward function so that optimism about the reward can be connected to a distribution over states. Given such a model, there are many possible acquisition strategies. A natural approach uses Thompson sampling, which we adapt for optimistic exploration with our method, Hallucination-based Optimistic Thompson sampling with Gaussian Processes (HOT-GP).
5.1 MODEL STRUCTURE
Optimistic exploration implies biased exploration towards regions of the state-action space believed to yield high rewards. Thus, optimistic learning requires reward-dynamies model that describes the correlation of state and reward. In this work, we use a Gaussian process (GP) prior to describe our belief over the product space of reward and state,
f G~ ((m())
(ea(:))) (.r(,:)) (T)(.)) Krr(.))
(7)
5

---

Published as a conference paper at ICLR 2025
where us :S A - S and μ4, S A - R are the dynamics and reward mean functions and Kee Kar; Krr represent covariance functions (S x A) x (S x A) + R. Traditional GPs describe outputs conditionally independent given the imput and hence factorize across the state dimen- sions and the reward. However, we hypothesize that this assumption is unfounded in most RI. tasks and furthermore antithetical to principled optimistic exploration. Therefore, we want to explicitly model the correlation structure between the dimensions of the predicted state, and more impor- tantly between the predicted state and its associated reward. While this correlation can naturally be included in the GP by parameterizing the full covariance structure across examples and output dimensions, this computationally intractable due to cubic scaling with the size of the covariance matrix (Rasmussen & Williams 2006). Instead, we adopt linear model of coregionalization (Grze-) byk & Wackernagel| 1994; Wackernagel 2003) where the covariance between output dimensions is restricted to be linear combinations of the input covariance. This induces a Kronecker structured co- variance function that allows for efficient factorization while still meaningfully directing optimistic exploration.
In most GP regression use cases, the mean function is a constant function. While this is a valid modeling assumption for simple data, the dynamics and reward structures encounter in RL tasks vary in a non-constant manner across different regions of the input space. To better predict these complex transition-rewadd relationships, we use GP model with a multi-layer perceptron (MLP) mean function, as in [wata & Ghahramani (2017); Fortuin et al. (2019). and use the covariance function learn the uncertainty, representing the residual around the MLP. This flexible model can capture intricate transitions while providing useful uncertainty estimates that can be leveraged in an optimistic fashion. The mean and covariance functions are learned jointly by performing approximate inference through optimizing a variational lower bound on the marginal likelihood, as done by Titsias 2009)). This updates the posterior distribution p(f 'Dew), which we use to obtain the posterior predictive distribution p(s+4, T. 8t, at). Given the current state st and action at this multivariate Gaussian distribution reflects the model's predictions and uncertainty over the next state and reward.
5.2 USING THOMPSON SAMPLING IN THE GP ACQUISITION FUNCTION
Thompson sampling proscribes optimizing the policy using single reward-dynamics model from the distribution of predictive models. A naive approach would sample from p(*+41, Ft 8t, az) until discovering a good model as function of the predicted reward rr- Instead, we observe that we can condition the reward value to be high, specifically greater than some minimum percentile threshold m'mn over the reward distribution, and sample once from this distribution. Crucially, this induces optimism into the sampling procedure because we only consider transitions where the next state is plausible under the dynamics model and whose associated reward we believe to be large.
To do this. we sample independent realizations from one non-parametric reward- -wwadddynnamics model for each episode of length T, denoted as
T II p( 11,a,,77 Tt > Tmin) ==
(8)
for 8 sampled uniformly from Den and ag = T())) under the current policy. Then the optimistic Thompson sampling algorithm is
f(0-t)) p(5t-1,7T St, Tz > Tmin), HOT--G arg max f. ((0+T) fio TEN
(9)
Performing Thompson sampling following the joint distribution in equation 9 predicts both the next state and its associated reward an optimistic manner. We argue that this would lead to a suboptimal exploration strategy as we are only interested in optimism over the state to the extent that it relates to optimism over the reward. To address this, we instead sample reward-dynamics model realizations using a two-step approach where only the reward is explored in optimistic manner:
1. Sample an optimistic reward A; from p(rt 8t, 4z Pr >> T'in); a truncated normal distribu- tion, using inverse transform sampling another sampling technique.
2. Compute the expected transition $ = Ep(s Ep [ [5++1] conditioned on the sam- pled optimistic reward fz.

---

Published as a conference paper at ICLR 2025
T.DOP Beadaer
DOE Poslec
683
Mouu
100
Disal Envirooment Stogs
Total Dastrroment Steps
Ttcal Divinoneens Stogo
Spene 2007 Feather Prralty
Spern 500 Teacher Action Preatty
Saano T.D0P Banderr Action Pennatyy 0.5
8 9
Meve
2006 2008
5008
3009
2000
6000
0000
B000
2006
4008
2008
0C09
10008
Figure 1: Learning curves for all MuJoCo tasks averaged 10 seeds except for the Sparse Reacher task, which used seeds. HOT-GP demonstrates equivalent or superior sample efficiency and performance for all tasks considered. The dashed line denotes SAC performance convergence within 1,000,000 environment steps (3,000,000 for Half- half-Cheetah).
This approach deviates from traditional Thompson sampling, and hence allows us to disentangle op- timism over the state from optimism over the reward that can concentrate solely optimism over the reward from the states. This aligns well with our goal of sample-efficient reinforcement learning as the objective is maximize the reward in few environment steps possible.
EXPERIMENTS
We evaluate the performance our proposed algorithm, HOT-GP, on continuous state-action con- trol tasks by measuring the average sum of rewards accumulated during evaluation episodes. Since inaccurate reward estimates are likely initially, we increase Tmin linearly throughout training to be- consider three
come more optimistic as the model becomes more reliable. In our experiments, variations of HOT-GP where mii begins at 0.1 and increases to either 0.3, 0.5, or 0.7. We also con- sider the alternative Thompson sampling and greedy exploration strategies. Our Thompson sampling implementation follows the HOT-GP procedure but draws samples from the reward-conditioned dis- tribution p(s++1 St, de, Ft fr) rather than taking the expectation. For greedy exploitation, predict the next state and reward as the mean from the model.
First, we compare these approaches on standard MuJoCo benchmark tasks (Todorov 2012) and extended MuloCo sparse maze tasks from the D4RL suite (Fu et al.) 2020) against additional model- free and model- based methods. Then evaluate these approaches on practical robotics task using the VMAS simulator (Bettini et al. 2022) and provide an analysis of the individual components. We provide algorithmic descriptions of HOT-GP and the model-based comparison methods in Appendix B specific tasks are described in Appendix C additional evaluations are given in Appendix DD. and code for reproducing the experiments can be
A Details about training are discussed in Appendix
found at
com http://barrooitiooo gp
6.1 MUJoCo BENCIMARKS
Comparison Methods. In these experiments, consider HOT-GP with varying levels of optimism using SAC (Haarnoja et al., 2018) as the underlying PolicySearch algorithm, as done in Janner et al. 2019) with MBPO. We compare performance against SAC, Thompson sampling, greedy exploitation, and MBPO. SAC is a model-free actor-critic algorithm that is known to achieve better sample efficiency than DDPG (Lillierap al. 2015) MuloCo tasks (Chua al. 2018). MBPO is

---

Published as a conference paper at ICLR 2025
model-based approach that uses ensemble of probabilistic neural networks to model uncertainty as done in Chua et al. (2018). Predictions are drawn from a Gaussian distribution with diagonal covariance over the reward and dynamics: P, s,,a,) (N(ss, a,), E(s,, at)), where (8,,, az) and E(8,, (t) are bootstrapped from the ensemble. For an optimistic baseline, we introduce an optimistic version of MBPO where the reward is sampled from p(re 8t, at, Pt > Tmin) and expected next state (p(s1))1s.as)
is selected. This is the same procedure as HOT-GP, however, due MBPO lacking joint model over the state and reward distributions, this implies naive optimistic exploration as there is no explicit relationship between the predicted next state and reward.
Standard Benchmarks. The learning curves for all methods on Calf-Cheetah, Reacher, and Pusher are given in the top row of Figure In each case, the HOT-GP variants attain SAC-best performance with comparable or superior sample efficiency to other approaches. For instance, on Half-Cheetah, HOT-GP with mmin from 0.1 to 0.3 achieves the same final performance as SAC in 250,000 envi- ronment steps rather than 3 million. We observe that the deterministic approach used in greedy ex- ploitation leads to slower learning on Half-Cheetah compared to uncertainty methods, while showing no impact on performance on Pusher. This indicates that model uncertainty is key for effi- cient learning Half- Chef-Chhetah, likely because the dynamics makes model biases more detrimental to learning (Chua et al.] 2018). Interestingly, uncertainty- Pusher as optimistic exploration methods enhance sample
more challenging to model, which
aware learning alone is insufficient efficiency compared to MBPO and greedy exploitation. This is likely due to the task dynamics as the agent must navigate a landscape of negative rewards and sparse feedback. Optimistic exploration allows the agent to discover valuable actions more effectively by pursuing riskier options with po- tential help reach the goal state. Next, we examine the impact of optimistic exploration in sparse reward settings more closely.
Sparse Reacher. The Sparse Reacher task provides reward rais that is nonzero only when the agent is within a small threshold of the target location, thus demanding clever exploration strategies to reach the target as there is no dense feedback. As in Curi et al. (2020), we introduce an action penalty of the form r(s, a)) rais(() - with the aim of penalizing random and encouraging local exploration. This is described in greater detail in Appendix C We evaluate HOT-GP and the other comparison methods
Sparse Reacher using four different action penalty weights: p 0, 0.1, م = 0.3, and p = 0.5, and present the results in the bottom row of Figure I as well as Figure 6 in Appendix D HOT-GP demonstrates significantly superior sample efficiency over MBPO and optimistic MBPO in each case. Also, the asymptotic performance of HOT-GP matches the SAC- best performance for each action penalty while achieving it in a fraction of the environment steps. When p = O, all HOT-GP variants perform robust optimistic exploration and successfully uncover the signal despite sparse rewards. When using the action penalty م = 0.1, HOT-GP with m'in from 0.1 to 0.3 performs marginally more poorly than with higher levels of optimism, suggesting that high levels of optimism may be required to overcome the prior against taking large actions and explore the reward signal ruis- HOT-GP does this by relating the action-based penalty to the state to the overall reward through the joint model. The less efficient learning of optimistic MBPO highlights the critical importance of having a joint model of the reward and state distributions to facilitate efficient exploration. For p = 0.3 and p = 0.5, we observe that the cost term accounts for up to half of the overall reward. As result, the HOT-GP variants explore much of the cost signal rather than raist, leading to similar performance in these settings.
Sparse Maze In the Sparse Maze tasks, agent must navigate through a maze randomized target location, receiving a reward of I when within a small threshold of the target location and zero in other cases. Similar to the Sparse Reacher task, efficiently solving such problems requires an effective exploration strategy due to the lack of dense feedback. We consider the simple U Maze and challenging Medium Maze environments, which are described in Appendix C The results of HOT-GP and the comparison methods on these mazes are given in Figure On the U Maze. HOT- GP attains equivalent sample efficiency and performance with the greedy and Thompson sampling strategies, while MBPO and optimistic MBPO are less stable. Greedy exploitation likely does well on this task because of the simple dynamies and reduced need for exploration. On the other hand, the Medium Maze presents a significantly more difficult exploration problem where greedy exploitation is insufficient to learn useful behavior. On the Medium Maze, HOT-GP with mmin ranging from 0. to 0.7 delivers the strongest performance while other HOT-GP variants trail behind, showing that lower levels of optimism result in poorer performance. This disparity is likely due to the substantial spar- sity in rewards and difficult maze layout, characterized by multiple dead ends, which necessitates
8

---

Published as a conference paper at ICLR 2025
U Maze
Medium Maze
100
Moan Paaa Fppoodo Rewwrrd
Moaa Foal Ppisodia Rowardd
33
sar ou
Orttonnte HEPO
Tharsoooe acndliee
90000
300000 Ervaronment Steps
Total Envirocmmnt Steps
Figure 2: Learning curves on sparse maze tasks averaged over 10 seeds on the Maze and 5 seeds on the Medium Maze. HOT-GP achieves equivalent superior sample efficiency and performance on both tasks. The dashed line denotes SAC performance after 2,000,000 environment steps.
more aggressive optimistic exploration. Thompson sampling initially matches the best HOT-GP per- formance, but later shows instability with performance decline that may be due to stochastic next state sampling. In contrast, MBPO and optimistic MBPO learn less efficiently, further reinforcing HOT-GP's advantage in jointly modeling reward and state distributions for effective exploration.
6.2 ROBOT COVERAGE
Coverage Description. We consider practically motivated continuous control robotics problem where an agent must visit as many unique, high valued cells as possible within a finite episode. This mirrors real-world target-tracking tasks where agent monitors multiple regions of interest under time constraints (Robin& & Lacroix 2016; Xin et al. 2024). The reward for visiting new 2D cell is distributed as the probability density function of mixture of three randomly placed Gaussians over the x and y axes, so there are cluster rewards that the agent must find and exploit. To incentivize thorough coverage, the agent receives no reward for revisiting cells in an episode.
Comparison Methods. As in the MuloCo tasks, we evaluate HOT-GP with three levels of optimism and compare against Thompson sampling and greedy exploitation. A distinction is that, since DDPG outperforms SAC in this task, we use DDPG with probabilistic actor for the Poli PolicySearch algorithm. Learning a stochastic policy allows for baseline level of exploration, which we found to be helpful when also incorporating uncertainty induced by the reward-dynamics model. For model- free baselines, we use DDPG, SAC, and the on-policy PPO algorithm (Schulman al.) (017). For an optimistic baseline, truth reward function fr while the other adapts to our unknown reward setting by learning the reward function fr and applying it as in equation 6 In both cases, we use the same dynamics-rewant model with H-UCRL as in HOT-GP for fair comparison of the optimistic exploration strategies.
two variations of H-UCRL: one has privileged access to the ground
Performance Analysis. The first two subfigures Figure show the training curves for these ap- proaches with and without access the ground-truth reward function. Our method outperforms the model-based approaches in both settings and achieves DDPG-best performance marginally ahead of DDPG. The poorer performance of greedy exploitation without access the true reward function suggests that strategic exploration is key to efficient learning for this task. Indeed, the agent must explore the environment sufficiently well enough to learn the relationship between its observations and the locations of the Gaussian centers, make informed decisions about which Gaussian center to prioritize, and realize the importance of avoiding previously seen cells. HOT-GP also demonstrates performance superior to H-UCRL, which does not converge to the optimal performance level within 200,000 environment interactions, with and without access to the ground truth reward. This may be due to infeasible optimism where the rewards associated with perturbed states pts, a) E(s, a)n are high but the states are not reachable in practice. HOT-GP avoids this problem by selecting states conditioned optimistic rewards from the joint distribution, thus ensuring the states plausi- ble under the learned model and reward value. These findings reinforce the benefits of principled optimistic exploration for improving sample efficiency in tasks where exploration is key.

---

Published as a conference paper at ICLR 2025
Agudfia Conpatioo
Leannng Bewand
Absute Study
Moun Thoi Tehends Rawvert
Mean Fridl Termade Rewwer
18
10
Tosal Divtroomen Stogs
Total Eastrremeees Steps
73cal Divironmenn Stogo
Figure 3: Learning curves in the coverage environment averaged over 10 seeds. HOT-GP achieves strong performance within 200,000 environment steps while other methods exhibit poorer sample efficiency or asymptotic performance. The dashed line denotes DDPG performance at convergence within 500,000 environment steps.
Ablation Study, Finally, we conduct an ablation study over the main components of HOT-GP to investigate their relative importance. The ablation study in Figure 3 presents the findings. We observe that scaling miin from 0.1 to 0.5 yields the best results, indicating that this task benefits from a medium level of optimistic exploration. Since HOT-GP using other ranges of T'in values still performs fairly well, our approach does not appear to be overly sensitive to m'ia settings. Indeed, mmin is an interpretable parameter that can be informed by the hypothesized level of exploration required. The ablation study in Figure 3 shows that HOT-GP without k-branching results in consistently poor performance. We also observe that k- k-branching with & 5 leads to less optimal performance than with & = 1. This aligns with prior findings that accumulating errors in the dynamics model can be detrimental to learning (Janner al., 2019 Gu al. 2016). Furthermore, this highlights the importance of learning a faithful model to benefit from optimistic exploration, as optimism guided by an inaccurate belief harms learning. The ablation study also reveals the importance of modeling the correlation between the predicted state and reward, as this is also essential to be able to harness optimistic exploration. Lastly, sampling the next state from an optimistic distribution with Thompson sampling leads to larger variance in performance than selecting the expected next state, thus validating our two-step sampling approach. Overall, this ablation study confirms that each component of HOT-GP is necessary to learn effectively.
7 CONCLUSION
We have introduced a principled approach to optimistic exploration based on leveraging joint belief over the reward and state distribution. Our proposed algorithm, HOT-GP, uses a Gaussian Process to upproximate the reward-dynamics and allows to simulate plausible transitions associated with optimistic rewards, in a practical adaptation of Thompson sampling. Our experiments showed that HOT-GP achieves comparable or superior sample efficiency relative other model-based methods and exploration strategies. Notably, we found that joint model uncertainty over outputs is crucial for effective exploration, particularly in challenging settings with sparse rewards, action penalties, and difticult-t--exploee regions. Ultimately, this work establishes the importance of joint uncertainty modeling for optimistic exploration and makes a meaningful advancement towards more sample- efficient reinforcement learning.
Future Work. Future work should study how the schedule on Tmin affects optimism and principled ways for setting this value throughout training, as we only considered a linear schedule. Another promising avenue would extend the planning horizon. We concentrated on learning 1 1-step looka- head reward-dynamics as HOT-GP can implicitly propagate uncertainty across the horizon by using pointwise uncertainty estimates from the model. Nevertheless, there may be benefits in explicitly propagating uncertainty to facilitate optimism over longer-term reward predictions (Seyde et al.) 2021 Mehta et al. 2022 Hansen et al. 2024). Furthermore, our investigation indicates that the util- ity of optimistic exploration depends the complexity of the task, although experiments were limited to six tasks. Future research should explore the role of optimism across a broader range of problems, including in real-world scenarios with on-robot learning.
10

---

Published as a conference paper at ICLR 2025
ACKNOWLEDGMENTS
J. Bayrooti is supported by DeepMind scholarship. A. Prorok is supported in part by European Research Council (ERC) Project 949940 (gAla). We also sincerely thank Markus Wulfmeier and Aidan Scannell for insightful conversations and valuable feedback, which significantly contributed to the development of this work.
REFERENCES
Yasin Abbasi Yadkoii and Csaba Szepesvári. Regret bounds for the adaptive control of linear quadratic systems. In Proceedings of the 24th Annual Conference on Learning Theory, pp. 1- 26. JMLR Workshop and Conference Proceedings, 2011.
Matteo Bettini, Ryan Kortvelesy, Jan Blumenkamp, and Amanda Prorok. VMAS: A vectorized multi-agent simulator for collective robot learning. The 16th International Symposium Dis- tributed Autonomous Robotic Systems, 2022.
Albert Bou, Matteo Bettini, Sebastian Dittert, Vikash Kumar, Shagun Sodhani, Xiaomeng Yang, Gianni De Fabritiis, and Vincent Moens. TorchRL: data-driven decision-makinn library for pytorch. In International Conference Learning Representations, 2024.
Ronen Brafman and Moshe Tennenholtz. R-max-a general polynomial time algorithm for near- optimal reinforcement learning. Journal of Machine Learning Research, 3(0et):213-231, 2002.
Paul Edmund Chang, Prakhar Verma, ST John, Arno Solin, and Mohammad Emtiyaz Khan. Memory-based dual gaussian processes for sequential learning. In International Conference on Machine Learning, pp. 4035-4054, PMLR, 2023.
Olivier Chapelle and Lihong Li. An empirical evaluation of thompson sampling. Advances in Neural Information Processing Systems, 24, 2011.
Sayak Ray Chowdhury and Aditya Gopalan. On kernelized multi-armed bandits. In International Conference on Machine Learning. pp. 844-853, PMLR, 2017.
Sayak Ray Chowdhury and Aditya Gopalan. Online learning in kernelized markov decision pro- cesses. In The 22nd International Conference Artificial Intelligence and Statistics, pp. 3197- 3205. PMLR, 2019.
Kurtland Chua, Roberto Calandra, Rowan McAllister, and Sergey Levine. Deep reinforcement learn- ing in handful of trials using probabilistic dynamics models. In Advances Neural Information Processing Systems, volume 31. Curran Associates, Inc., 2018.
Sebastian Curi, Felix Berkenkamp, and Andreas Krause, Efficient model-based reinforcement learn- ing through optimistic policy search and planning. In Advances in Neural Information Processing Systems, volume 33, pp. 14156-14170, 2020.
Mare Peter Deisenroth and Carl Edward Rasmussen. PILCO: model-based and data-efficient approach to policy search. In Proceedings of the 28th International Conference on Machine Learning. pp. 465-472. Omnipress, 2011.
Mare Peter Deisenroth, Dieter Fox, and Carl Edward Rasmussen. Gaussian processes for data- efficient learning in robotics and control. IEEE Transactions on Pattern Analysis and Machine Intelligence, 37(2):008-423, 2013.
Vincent Fortuin. Priors in bayesian deep learning: A review, International Statistical Review, 90(3): 563-591, 2022.
Vincent Fortuin, Heiko Strathmann, and Gunnar Rätsch. Meta-learning mean functions for gaussian processes. arXiv preprint arXiv:1901.08099 2019.
Justin Fu, Aviral Kumar, Ofir Nachum, George Tucker, and Sergey Levine. D4rl: Datasets for deep data-driven reinforcement learning. arXiv preprint arXiv:2004.07219, 2020.

---

Published as a conference paper at ICLR 2025
Yarin Gal, Rowan McAllister, and Carl Edward Rasmussen. Improving pilco with bayesian neural network dynamics models. In JCML Workshop on ruta-efficient Machine Learning, volume 4, pp. 25, 2016.
Yarin Gal, Jiri Hron, and Alex Kendall. Concrete dropout. Advances in Neural Information Pro- cessing Systems, 30, 2017.
Jacob Gardner, Geoff Pleiss, Kilian Q Weinberger, David Bindel, and Andrew Wilson. GPyTorch: Blackbox matrix-matri gaussian process inference with GPU acceleration. Advances in Neural Information Processing Systems, 31, 2018.
Michel Grzebyk and Hans Wackernagel. Multivariate analysis and spatial/temporal scales: real and complex models. In International Biometrics Conference, volume 1, pp. 19-33, 1994.
Shixiang Gu, Timothy Lillicrap, Ilya Sutskever, and Sergey Levine. Continuous deep Q-learning with model-based acceleration. In International Conference on Machine Learning, pp. 2829- 2838. PMLR, 2016.
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 deep reinforcement learning with stochastic actor. In International confer- ence on machine learning, pp. 1861-1870. PMLR, 2018.
Danijar Hafner, Timothy Lillicrap, Ian Fischer, Ruben Villegas, David Ha, Honglak Lee, and James Davidson. Learning latent dynamics for planning from pixels. In International Conference on Machine Learning, pp. 2555 2565. PMLR, 2019.
Danijar Hafner, Timothy Lillicrap. Jimmy Ba, and Mohammad Norouzi. Dream to control: behaviors by latent imagination. In International Conference Representation Learning, 2020.
Learning
Danijar Hafner, Jurgis Pasukonis, Jimmy Ba, and Timothy Lillicrap. Mastering diverse domains through world models. arXiv preprint MxXiv23301.0400 2023.
Nicklas Hansen, Xiaolong Wang, and Hao Su. Temporal difference learning for model predictive control. In International Conference on Machine Learning, 2022.
Nicklas Hansen, Hao Su, and Xiaolong Wang. TD-MPC2: Scalable, robust world models for con- tinuous control. International Conference on Learning Representations, 2024.
James Hensman, Alexander Matthews, and Zoubin Ghahramani. Scalable variational gaussian pro- cess classification. In Artificial Intelligence and Statistics, pp. 351-360. PMLR, 2015.
Yujing Hu, Weixun Wang, Hangtian Jia, Yixiang Wang. Yingfeng Chen, Jianye Hao, Feng Wu, and Changjic Fan. Learning to utilize shaping rewards: new approach of reward shaping. In Advances in Neural Information Processing Systems, volume 33, pp. 15931-159411 2020.
Julian Ibarz, Jie Tan, Chelsea Finn, Mrinal Kalakrishnan, Peter Pastor, and Sergey Levine. How to train your robot with deep reinforcement learning: lessons we have learned. The International Journal of Robotics Research, 40(4-5):698-721 2021.
Tomoharu Iata and Zoubin Ghahramani. Improving output uncertainty estimation and generaliza- tion in deep learning via neural network gaussian processes. arXiv preprint (rXXi:1707.09// 2017.
Thomas Jaksch, Ronald Ortner, and Peter Auer. Near-optimal regret bounds for reinforcement learning. Journal of Machine Learning Research, 11(51):1563-1600 2010.
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.
Chi Jin, Zhuoran Yang, Zhaoran Wang, and Michael Jordan. Provably efficient reinforcement learning with linear function approximation. In Conference on Learning Theory, pp. 2137-214 PMLR, 2020.
12

---

Published as a conference paper at ICLR 2025
Leslie Pack Kaelbling, Michael L Littman, and Andrew W Moore. Reinforcement learning: A survey. Journal of Artificial Intelligence Research, 4:237-285, 1996.
Lukasz Kaiser, Mohammad Babacizadeh, Piotr Milos, Blazej Osinski, Roy H Campbell, Konrad Czechowski, Dumitru Erhan, Chelsea Finn, Piotr Kozakowski, Sergey Levine, et al. Model-based reinforcement learning for atari. International Conference on Learning Representations, 2020.
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.
Balaji Lakshminarayanan, Alexander Pritzel, and Charles Blundell. Simple and scalable predic- tive uncertainty estimation using deep ensembles. Advances in Neural Information Processing Systems, 30, 2017.
Sergey Levine. Chelsea Finn, Trevor Darrell, and Pieter Abbeel. End-to-end training of deep visuo- motor policies. Journal of Machine Learning Research, 17(39):1-40, 2016.
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 arXiv: STTT::/500.09977 2015.
Horia Mania, Stephen Tu, and Benjamin Recht. Certainty equivalence is efficient for linear quadratic control. Advances in Newral Information Processing Systems, 32, 2019.
Viraj Mehta, Ian Char, Joseph Abbate, Rory Conlin, Mark Boyer, Stefano Ermon, Jeff Schneider, and Willie Neiswanger. Exploration via planning for information about the optimal trajectory. Advances in Neural Information Processing Systems, 35:28761-28775 2022.
Volodymyr Mnih, Koray Kavukcuoglu, David Silver, Andrei A Rusu, Joel Veness, Mare G Belle- mare, Alex Graves, Martin Riedmiller, Andreas K Fidjeland, Georg Ostrovski, et al. Human-level control through deep reinforcement learning. Nature, 5 518(7540):529-533 2015.
Anusha Nagabandi, Gregory Kahn, Ronald S Fearing, and Sergey Levine. Neural network dynamics for model-based deep reinforcement learning with model-free fine-tuning. In 2018 IEEE Interna- tional Conference on Robotics and Automation (ICRA), pp. 7559-7566. IEEE, 2018.
Gergely Neu and Ciara Pike-Burke. unifying view of optimism in episodic reinforcement learning. Advances in Neural Information Processing Systems, 33:1392-1403, 2020.
Junhyuk Oh, Xiaoxiao Guo, Honglak Lee, Richard I. Lewis, and Satinder Singh. Action-conditional video prediction using deep networks in atari games. Advances in Neural Information Processing Systems, 28, 2015.
Ian Osband. Risk versus uncertainty in deep learning: Bayes, bootstrap and the dangers of dropout. In NewrlPS Workshop on Bayesian Deep Learning, volume 192, 2016.
Ian Osband and Benjamin Van Roy. On lower bounds for regret in reinforcement learning. arXiv preprint arXiv:1608 X1v.1660.000732, 2016.
Ian Osband and Benjamin Van Roy. Why is posterior sampling better than optimism for reinforce- ment learning? In International Conference on Machine Learning, pp. 2701-2710. PMLR, 2017.
Ian Osband, Daniel Russo, and Benjamin Van Roy. (more) efficient reinforcement learning via posterior sampling. Advances in Neural Information Processing Systems, 26, 2013.
Deepak Pathak, Pulkit 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.
Luis Pineda, Brandon Amos, Amy Zhang, Nathan O Lambert, and Roberto Calandra. Mbrl-lib: A modular library for model-based reinforcement learning. arXiv preprint 104. 2021.
Carl Edward Rasmussen and Christopher KI Williams. Gaussian processes for machine learning, volume 2. MIT Press Cambridge, MA, 2006.
13

---

Published as a conference paper at ICLR 2025
Cyril Robin and Simon Lacroix. Multi-robot target detection and tracking: taxonomy and survey. Autonomous Robots, 40:729-760, 2602 2016.
Hrittik Roy, Marco Miani, Carl Henrik Ek, Philipp Hennig, Marvin Pförtner, Lukas Tatzel, and
Søren Hauberg. Reparameterization invariance in approximate bayesian inference. In Advances in Neural Information Processing Systems, 2024.
Daniel Russo and Benjamin Van Roy. Learning to optimize via information-direeted sampling. Advances in Neural Information Processing Systems, 27, 2014.
Daniel Russo, Benjamin Van Roy, Abbas Kazerouni, Osband, Zheng Wen, al. A tutorial on thompson sampling. Foundations and Trends09 in Machine Learning, 11(1):1-96, 2018.
John Schulman, Filip Wolski, Prafulla Dhariwal, Alec Radford, and Oleg Klimov. Proximal policy optimization algorithms. arXiv preprint arXiv: 1707.06377. 2017.
Ramanan Sekar, Oleh Rybkin, Kostas Danilidis, Pieter Abbeel, Danijar Hafner, and Deepak Pathak. Planning to explore via self- supervueed world models. In International Conference Machine Learning, pp. 8592.. PMLR, 2020.
Pier Giuseppe Sessa, Maryam Kamgarpour, and Andreas Krause. Efficient model-based multi-agent reinforcement learning via optimistic equilibrium computation. In International Conference on Machine Learning, pp. 19580-19597. PMLR, 2022.
Tim Seyde, Wilko Schwarting, Sertac Karaman, and Daniela Rus. Learning to plan optimistically: Uncertainty-guie guided deep exploration via latent model ensembles. Conference on Robot Learn- ing. 2021.
Mrinank Sharma, Sebastian Farquhar, Eric Nalisnick, and Tom Rainforth. Do bayesian neural net- works need to be fully stochastic? Proceedings of The 26th International Conference on Artificial Intelligence and Statistics, volume 206 of Proceedings of Machine Learning Research, pp. 7692-7722. PMLR, 2023.
Richard Sutton. Integrated architectures for learning, planning, and reacting based on approxi- mating dynamic programming. In Machine Learning Proceedings 1990, pp. 216-224. Elsevier, 1990.
Richard S Sutton and Andrew G Barto. Reinforcement learning: An introduction. MIT press, 2018.
Csaba Szepesvári. Algoritims for reinforcement learning. Springer nature, 2022.
William R Thompson. On the likelihood that one unknown probability exceeds another in view of the evidence of two samples. Biometrika, 25(3-4):285-294) 1933.
Michalis Titsias. Variational learning of inducing variables in sparse Gaussian processes. In Inter-
national Conference Airtificial Inteligence and Statistical Learning, pp. 567-574, 2009.
Emanuel Todorov, Tom Erez, and Yuval Tassa. Mujoco: A physies engine for model-based control. In 2012 IEEE/RSJ International Conference on Intelligent Robots and Systems, pp. 5026-5033. 2012. doi: 10. 2009/RROS.2012.63888
Hans Wackernagel. Multivariate geostatistics: an introduction with applications. Springer Science & Business Media, 2003.
Jonathan Wenger, Kaiwen Wu, Philipp Hennig, Jacob R Gardner, Geoff Pleiss, and John P Cun- ningham. Computation- aware gaussian processes: Model selection and linear-time inference. In Advances in Neural Information Processing Systems, volume 37, 2024.
Pujie Xin, Zhanteng Xie, and Philip Dames. Towards predicting collective performance in multi- robot teams. arXiv preprint arXiv:2405.01771, 2024.
Yuxiang Yang, Ken Caluwaerts, Atil Iscen, Tingnan Zhang, Jie Tan, and Vikas Sindhwani. Data efficient reinforcement learning for legged robots. In Conference Robot Learning, pp. 1-10. PM 2020.
14

---

Published as a conference paper at ICLR 2025
A ALGORITHM DETAILS
In this section, we provide more detailed information about the algorithms we use for evaluation. Note that the code used to implement these algorithms is available at https: htttps://githbb com/ jpa.yoott/hot.tttooooo
A.1 MODEL-FREE BASELINES
We compare model-based RL approaches with the following popular model-free algorithms:
SAC. We use a standard version of SAC as introduced in Haarnoja et al. (0018). Our implementation
utilizes Gaussian policy and automatic entropy tuning.
PPO. We use a standard version of PPO as introduced in Schulman et al.
2017).
DDPG. We use variant of DDPG (Lillicrap al. 2015) with probabilistic actor network. In stan- dard DDPG, the deterministic nature of the policy can lead to insufficient exploration of the action space. We mitigate this tendency by using a probabilistic actor, which incorporates exploration into the policy by sampling actions from probability distribution. Specifically, the probabilistic actor network outputs a mean ( and standard deviation dg(s) for the current state s and we select the action by sampling from the Gaussian distribution a (Ja(s)) (s), (s))) In our experiments, compare with this version of DDPG using natural exploration under the label "DDPG w/o explore". With deterministic policy, it is also common to facilitate exploration by adding noise to the tions. Therefore, we additionally evaluate our DDPG bascline with explicit exploration under the label "DDPG explore". Specifically, we add Gaussian noise to the sampled action with + for Ç ~ (O,, o2 @explore)
A.2 GREEDY EXPLOITATION
This exploration strategy follows the optimization rule given in equation In particular, we hal- lucinate the next state as the mean output from the model and disregard the uncertainty around the prediction. In one variation, we assume access to the ground-truth reward function fr (referred to as "Greedy w known reward") for comparison with H-UCRL. In another variation, we also learn the reward function (referred to as "Greedy w/o known reward" in the coverage experiment and "Greedy" in the MuJoCo experiments) and similarly hallucinate rewards as the mean output from the model. In pseudocode, this approach follows Algorithm I exactly with the following algorithm- specific routine replacing line 10 select the next state and reward:
Compute the next expected state as 8k+1 [Ks(kk, 0k)
Observe the expected reward f (5x, âx.) or true reward Ag = fr (Sx, ax)
A.3 MBPO
We implement MBPO as described Janner et al. (2019) and summarized in Section 6.1 Specif- ically, the model is an ensemble of probabilistic neural networks that each predict a mean out- put and variance (exponentiated log variance in practice) p (s6, ax) and (' ax) and param- eterize multivariate Gaussian distribution with diagonal covariance p' p*(**+1,7h Sk, az)
N N( (8k. 0x), (* ax)). To predict an output, the parameters are bootstrapped using expecta- tion randomly selecting member and the outcome is drawn from the Gaussian distribution. This is described with the following steps replacing line 10 in Algorithm [
Compute the predictions for each ensemble member ( 0x), (3,,, ax) + f( ax)
Bootstrap predictions to get jiss,, àx), E(Sx, âx)
Sample outcome from the Gaussian distribution 8441, Av ~ N( âx), E(Sk, ax))
For the optimistic MBPO baseline, we maintain the ensemble model structure but perform proce- dure similar to HOT-GP, following these steps to generate next states and rewards:
Compute the predictions for each ensemble member yl (5x, ax); (3 âx) (8,, ax)
15

---

Published as a conference paper at ICLR 2025
Bootstrap predictions to get μá(ş,, ax), E(Sk: Qx)
Separate the reward distribution by taking its component p(sk+1, Sk, [r) N ( ax), ( (x))
Sample reward from an optimistic reward distribution 1( &R. 0g, Ax > Tmin)(r)
Select expected next state from state distribution $K+1 = Bp(sstrlBe [Sk+1]
A.4 H-UCRL
We implement H-UCRL (Curi et al. 2020) using the sampling procedure introduced in (Sessa al. 2022). On each model-generated rollout step, this approach hallucinates candidate next states that plausible under the learned dynamics model and selects the one leading to the highest reward under fr. Therefore, we assume access to the ground truth reward function fr and only learn the dynamics model fo in our implementation of H-UCRL. In pseudocode, this approach follows Algo- rithm I exactly with the following steps replacing line 10 where we refer to a. (8, a) as the standard deviation of an output from the GP dynamics model:
Draw random perturbations xy² Uniform([-1,17) for I,...,
Compute plausible next states 6/ Mo(5k, @x) 05(5k, 0x)
Select next state $ arg max 1/ fr ( al a for associated actions az 2
Observe the true reward f = fr @x)
We also adapt H-UCRL to our setting where the reward is unknown with a variation we call "H- UCRL Approx". We do this by learning both the dynamics and reward models and selecting the candiate state leading to the highest predicted reward, i.e., under fr rather than fr-
A.5 THOMPSON SAMPLING
This approach follows the optimization rule given in equation 5 Our Thompson Sampling imple- mentation only differs from HOT-GP in that we sample both the reward and next state from the optimistic reward-dynamics distribution. In practice, we implement this as described in Algorithm with the following algorithm-specifie routine replacing line 10:
Compute predictive posterior distribution p(sk+1, T& $g., a+) using p(f Dewv)
Sample an optimistic reward A p(rk] Sk. ōR. Pk > rmin)
Sample a plausible next state $5+1 p(sk+1 1 8k, åg, rk fx)
A.6 HOT-GP
In HOT-GP, we sample an optimistic reward and select the most likely next state associated with the sampled reward each step of the model-generated rollout, described in equation 9 The steps are detailed Algorithm with the following routine replacing line 10:
Compute predictive posterior distribution p($K+11, Th 8ky ax) using p(f Denv)
Sample an optimistic reward fx ~ p(Tk 8tc, 02 > Tmin)
Select expected next state SK+1 Epp (i [s++1]
B TRAINING
Learning Framework. We implement learning framework with MBRL-Lib (Pineda et al. 202T) for the MuJoCo tasks and TorchRL (Bou et al. (024) for the coverage task. Additionally, we GPyTorch (Gardner et al. 2018) build the GP reward-dynamies model.
16

---

Published as a conference paper at ICLR 2025
Table Task-specific hyperparameter values
Hyperparameter
Half-Cheetah
Reacher
Pusher
Sparse Reacher
Coverage
Environment steps N
250,000
20,000
20,000
37,500
200,000
Steps per rollout T
1000
150
150
150
150
Model rollouts M
adaptive
adaptive
adaptive
adaptive
150
Model rollout steps K
1
1
1
Batch size B
256
256
256
256
150
Discount factor y Learning rate a
0.99
0.99
0.99
0.99
0.9
0.001
0.001
0.001
0.001
0.00005
Mixing factor
0.005
0.005
0.005
0.005
0.005
Replay buffer size
unlimited
unlimited
unlimited
unlimited
20,000
Gaussian Process with Neural Network Mean Function, A standard Gaussian Process is defined by a mean function m(x) and covariance functiin ((x,,'), where æ and x' are input vectors, and m(x) is typically assumed to be zero mean. In this work, function with neural network so that m(x) = fo (x), where 0 are the parameters of the network. Given inputs X [ Fnl and outputs y (./..... Vnl, the posterior distribution of this GP is
replace the traditional zero mean
ply y) (m,, K oZ)),
the mean vector of neural network outputs for each input, K is the kernel matrix Ki = k(,,, 22)) o² is the noise variance, and I is the identity matrix. To optimize this model, we follow the procedure lwata & Ghahramani (2017). This entails optimizing / using the Mean Squared Error loss and then optimizing the kernel hyperparameters and noise variance o² on the train dataset with transformed targets m, as described in the following section, on each iteration. For predicting new input 2., the predictive posterior distribution is p(y. X,y) N N(y= Us: G.) with:
where m (.e
U. m. KK? (K +0²1) (y- m) = = KTT (K + +021) K.,
where = So(x.), K. (1),., xx))" K., k(x,, Thus, the neural network learns the non-zero mean function while the GP covariance function modeled by the kernel as usual, thus allowing the model to quantify uncertainty around complex functions.
Reward-Dynaminn Model Training. Our approach requires access the full posterior distribution of the Gaussian process. To facilitate this in an efficient manner, we train the GP using a variational bound based on inducing points (Hensman et al., 2015) following the approach given Twata & Ghahramani 2017). This allows for efficient training using stochastic inference. To further reduce the computational cost of learning the model, we do not train on the whole dataset on each iteration and instead draw 1, 000 randomly sampled transitions from Dew This approach reduces the training cost significantly and provides an additional source of stochasticity that helps training. Using a Gaussian likelihood allows to compute the predictive posterior closed form, allowing us direct access to its mean and variance.
Hyperparameters. We train all algorithms using of hyperparameters that we found to be optimal through trial-and-error or were recommended. ddeseie]s the hyperparameter values specific to the MuJoCo and VMAS tasks we considered. We also report additional hyperparameters used. In H-UCRL, we use the peploration-exploitaeee number of samples DDPG, we use an exploration noise Oexplore from 1 to 0.1. For Ga-based approaches, we use the Matern kernel as the covariance function and learn from 100 inducing points. For the MuJoCo tasks, we use a 4-yayer MLP with 200 hidden units per layer and SiLU activation function. For the VMAS coverage task we use 2-layer MLP with 200 hidden units per layer and Mish activation function. We please see our oouncsourced implementation at this repository: https: / / wwg.thub.cco http:thhb..ccoboocce.eee cow.bb/ooo//aaooooiiiio.
coefficient / = 0.01 and 5. MBPO and optimistic MBPO, we use an ensemble of size 7. For
the Adam optimizer in all cases. For further information implementation,
17

---

Published as a conference paper at ICLR 2025
Table 2: Computation times in hours across tasks and algorithms
Wall Clock Runtime
MBPO
Optimistic MBPO
Greedy
HOT-GP
Half-Cheetd
24.5 + 3.5
26.7 = 4.88
18.5
28.89 + 4.1
Reacher
1.89 + 0.1
2.04 + 0.2
2 2
5.18 + 0.2
Pusher
1.87 + 0.2
2.05 + 0.2
1.92 0.4
6.53 + 0.1
Sparse Reacher
3,76 + 0.3
4.1 0.4
3.94 + 0.3
8.55 + 1.5
U Maze
12.45 + 1.1
12.43 + 1.7
11.75 1.6
21.08 3.8
1-
1- 30
1-
105
120
Figure 4: Frames from rollout of HOT-GP trained policy in the coverage environment. The agent (purple circle) drives around to cover target areas which are brightly colored. The target locations change on each episode.
Computation Time Comparison. We report the mean wall clock runtime and standard deviations for a single run of each algorithm on MuJoCo tasks considered Table 2 These results are averaged over 10 seeds. Our HOT-GP implementation incurs longer wall clock runtimes due to computational cost GP training. However, recent research proposes promising techniques that mitigate this overhead (Wenger et al. 2024 Chang et al. 2023). Also, note that approach is not tied to GP models and can seamlessly generalize to alternative models such as Bayesian Neural Networks (BNNs), as the model is only needed to approximate the joint posterior distribution next states and rewards.
C EVALUATION TASKS
In this work, evaluate HOT-GP five continuous control tasks: four using the MuJoCo physics engine and one using the VMAS simulator. this section, we outline the setup for each task.
Half-Cheetah. The Half-Cheetil is MuJoCo continuous control task where 2D bipedal agent, resembling simplified cheetah, must learn to run as quickly as possible. The observations are 17- dimensional vectors that include the position and velocity of the agent's torso and joints. Actions are 6-dimensional vectors, where each component corresponds to the torque applied to one of the agent's joints. The reward is primarily based the forward velocity of the agent's torso while also penalizing excessive torques the joints. Let be the agent's 's original position, be the position after simulation step, and dt be the time between. Then the reward is given by r(s,
$1 di
0.1 Di= a?.
The Pusher task in MuJoCo is continuous control task where robotic arm must leam to push an object to a target location on a 2D plane. The agent controls a three-link pla- anm and must manipulate tion space 23-dimensional vector, which includes the position and velocity of the robot arm's joints, position and velocity of the object being pushed, and the position of the target. The
Pusher.
object move it towards fixed goal position. The observa-
18

---

Published as a conference paper at ICLR 2025
Figure 5: Visualizations of the U Maze (left) and Medium Maze (right) environments with randomly
selected goal locations.
tion space is 7-dimensional vector, where each component represents the torque applied one of the joints in the robotic arm. The reward function is based on the Fuclidean distance be- tween the object and the target location, incentivizing the agent to move the object closer to the target while also penalizing large torques. Let Pobj be the position of the object, Pgoal be the position of the goal, Purm be position of the arm's tip. Then the reward is given by Pgual l2 + 0.1 1-E=1a; a?. 0.5 (- PlPakj Parmll2).
r(s,a) a) -||Potj
Reacher. The Reacher task in MuJoCo is continuous control environment where robotic arm must learn to position its end effector at goal location in 3D space. The observations are 19- dimensional vectors that include the joint angles and velocities of the robotic arm. Actions are 6-dimensional vectors, where each component corresponds the torque applied one of the arm's joints. The reward incentivizes the agent to reach the goal while also penalizing excessive control actions. Letting pos represent the position of the end effector and p be the action penalty, the reward function is r(s, a)
(F - goal,)² - p ||a||?
The Sparse Reacher is the same the Reacher task, except with modified reward structure. Let (E(ss represent the position of the end effector, be a small threshold distance, and p be the action penalty. The reward function is
Sparse Reacher.
r(s, a) = - exp
)
a
-1
exp (-I(EE(s) - goalll)) if ||BB(s) - goal| <
otherwise.
The first term rewards the agent based on the negative squared distance between the end effector's position and the target goal if the distance is less than and gives no reward otherwise. The second term penalizes the agent for applying excessive torques to the joints. This design creates a challenge for exploration the agent receives minimal feedback about its performance until it is sufficiently close to the target. We take inspiration for this setup from Curi et al. (2020) although our rewards more sparse and
0.2 in experiments.
Maze. In this task, the agent must navigate through a U-shaped maze shown in Figure 5 to reach randomized target location. The agent receives sparse reward of I when it comes within small threshold distance of the target and 0 otherwise. Specifically, the reward function is
[1 if [Pos(s) goal||
r(s, a) o otherwise.
The agent observes its position (x, y) and 2-dimensional goal position. Actions are 2-dimensional vectors representing velocity in the and dimensions.
Medium Maze, This task shares the same observation and reward structure with the U Maze. The only difference is that this maze environment is significantly more complex, shown in Figure
19

---

Published as a conference paper at ICLR 2025
Coverage. In this VMAS task, the agent must discover and visit unique, discretized cells with high value a continuous acttin space. The reward for visiting a new cell at coordinates (æ, y) is distributed as the probability density function of mixture of three Gaussians over the and y axes,
p(z, v) - E ((((, (mi, 2)).
(10)
To incentivize thorough coverage of the regions around each Gaussian center, the agent receives no reward for revisiting cells in an episode. The agent observes its position (z. y). velocity, current (e.y) y), and the target locations value p(r,3) 3). values of neighboring cells p( Un.) for (In Vn)) E (zx
{nr}=-1- The target locations are randomly generated in each episode so that the agent must learn the high-level pattern of reward distribution function of (z; y) and {M: ) and E. Throughout our experiments, we use 0.05. We provide snapshots from sample rollouts on the coverage task in Figure 4
D EXTENDED EXPERIMENTS
In this section, we extend our experiments in order to better contextualize HOT-GP. First, we inves- tigate the asymptotic performance of HOT-GP and all comparison methods on the Sparse Reacher task. We also provide learning curves on the Sparse Reacher task with action penalty 0.3 to supple- ment the results in Figure T with an intermediate action penalty. Finally, we evaluate the approxi- mation of H-UCRL and MBPO without uncertainty clamping the Half-Cheetah task.
Sparse Reacher. To supplement the Sparse Reacher results presented in Figure [ we provide an additional learning curve the Sparse Reacher using action penalty of 0.3. These results consistent with the trends observed using other action penalties and HOT-GP variants perform the best.
Sparsa DOP Reacher Action Penalty 0.3
Hale-Cheetab
125
HEzOO
12:02
Moan Peni Ppiianti Rownrd
10002
1o03
Mean Evall
6000
-26
-50
-2500
12000
28000
20000
Envirvament Staps
Total Envirommnnt Steps
Figure 6: Learning curves on the Sparse Reacher with action penalty of 0.3. HOT-GP achieves su- perior sample efficiency and performance. The dashed line denotes SAC performance conver- gence within 1, 1,000,000 environment steps.
Figure 7: Learning curves for select methods on Half-Cheetae averaged over 5 seeds. The dashed line denotes SAC performance at con- vergence within 3,000,000 environment steps for Half- Cheetah.
Half-Cheetah. As described in Section A H-UCRL approx differs from the algorithm presented in Curi et al. (2020) in few ways: (1) we use our GP model rather than a probabilistic ensemble,
optimistic hallucination method equation 6 to select states, and (3) we use predicted rewards to guide optimism rather than the ground truth rewards. This version of H-UCRL does not perform well on Half-Cheetah relative to MBPO and HOT-GP, with the best seeds achieving mean evaluation episode rewards high as 6000 after training with 250,000 envi- ronment steps. As we observed when discussing H-UCRL performance on the coverage task, this performance gap may stem from infeasible optimism, where the rewards associated with perturbed states juss, a.) + E(s, a) a)n are high but unreachable in practice. Additionally, a suboptimal choice of B could be contributing factor, we did not fine tune this parameter
(2) we use the sampling-
due to
20

---

Published as a conference paper at ICLR 2025
its extended runtime and instead used the optimal / finetuned the coverage task. HOT-GP also has a hyperparameter mmin, which controls the level of optimism during training. While the domain dependence of min means some adaptation is required for tasks, found that setting min is generally intuitive and, moreover, provides interesting insights into the nature of the task itself. For instance, if the dynamics in the domain are simple, (i.e., use a higher value of mmin), whereas if they are complex, we need to spend more time learning the dynamics before becoming more optimistic.
be more optimistic about exploration
We also experimented with removing uncertainty clamping from MBPO, which resulted in a com- plete failure to learn as shown in Figure 7 Standard MBPO uses an ensemble of probabilistic neural networks that predicts both the mean and log variance for given input, with log variances clamped prevent excessive growth. Removing this clamping significantly degrades performance, indicating that the model uncertainties are poorly calibrated and heavily dependent clamping to maintain stable learning. This finding reinforces our choice of model, GPs provide well-calibrated uncertainties (Chowdhury & Gopalan 2017)
21


Reviews:
[{"review_id": "p1t5aZN0uq", "reviewer": "Reviewer_ab8V", "summary": "This paper proposed Hallucination-based Optimistic Thompson Sampling with Gaussian Processes (HOT-GP),  a method for principled optimistic exploration. It uses a GP with Thompson Sampling as the acquisition function to maintain a joint belief over state and reward distributions, enabling simulating optimistic transitions wrt the estimated reward. Experiments on MuJoCo and VMAS continuous control tasks demonstrate that the proposed approach matches or outperforms baselines, showing substantial improvements in challenging exploration settings. The paper also investigates factors influencing optimism through ablation studies, highlighting the importance of joint uncertainty modelling in guiding effective exploration.", "questions": "Could you provide any results on computational cost or computation speed for the proposed method?", "limitations": "", "rating": 8, "confidence": 3, "soundness": 2, "presentation": 3, "contribution": 2, "strengths": "- principled approach for exploration in continuous state-action space\n- the first method to predict joint uncertainty over next state and reward", "weaknesses": "- The description for Fig.2 in Section 6.2, references several methods that are not shown in the plots. For instance,\n    - SAC and PPO as described in Line 430, “we use … SAC, and the on-policy PPO”\n    - HOT-GP with k-branching (k=5) as described in Line 470.\n- Learning curves seem to be from the training episodes. It would be better to report the performance in evaluation episodes instead, as done in SAC.\n- The method introduces a crucial hyperparameter reward threshold r_{min}, which controls optimism during exploration. While the results do not seem to be overly sensitive to it, its effect in other tasks remains unclear. Although future work is mentioned to address this, its current implementation could limit its applicability.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "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": "", "extra_scores": {}}, {"review_id": "Bmi8sBHSOp", "reviewer": "Reviewer_QgWe", "summary": "This paper introduces a novel model-based reinforcement learning algorithm to explore effectively and improve sample efficiency. The introduced method uses a Gaussian processes to represent a posterior over transition functions and reward functions and models the (linearized) covariance of reward and successor states (the mean functions are later modeled by MLPs instead). The authors introduce an exploration method by modifying Thomposon sampling to be more optimistic: rewards are sampled with the condition of a minimum reward and successive states are taken to be the expected successor state, conditional on the same minimum reward. The authors conducted experiments that analyze the introduced algorithms empirically and show that their agents learns faster in early stages of training.", "questions": "1. Why do the authors aim to relieve the stationarity assumption in line 285. Since this GP models the transition dynamics of an MDP, the transition kernel should indeed be stationary. Aiming to address non-stationary MDPs, I believe, would entail significant complications for the entire learning algorithm.\n\n2. In most experiments, the asymptotic performance of the shown algorithm seems to be lower than the base-SAC implementation in all of the environments. Is this due to the cutoff in plotting or is the asymptotic performance indeed lower in most environments? Inserting a dashed line as with SAC into the figure might clear this up. \n\n3. Given that the authors perform k-branching rollouts, the value estimates bootstrapped at the end of model-generated rollouts should play an important role in quantifying uncertainty. Would omitting this uncertainty skew the ability of the algorithm to explore efficiently? It seems unlikely to me, that Thompson sampling maintains it's theoretical guarantees if one only samples single-step transitions around known trajectories. \n\nMy suggestion would be to analyze the implications of the proposed algorithmic components carefully. \n1. What is the implication of performing Thompson sampling with a sampling probability dependent on the expected return of some policy?\n2. What is the implication of performing Thompson sampling with a sampling probability dependent on the return and reward distribution of some policy (as suggested here).\n3. What is the implication of performing Thompson sampling with limited branched rollouts?\n\nAddressing these questions in my opinion would provide a compelling motivation for the proposed algorithm.", "limitations": "", "rating": 3, "confidence": 3, "soundness": 2, "presentation": 3, "contribution": 2, "strengths": "The paper is written clearly and it is easy to follow. The method adresses an important issue and from what I can see is applicable to a variety of algorithms. The experiments suggest that the proposed algorithm may aid sample efficiency, especially in early stages of training.", "weaknesses": "The main weakness of this paper to me is of conceptual nature. I have no good intuition for why the proposed algorithm performs efficient exploration. This is mainly due to two reasons:\n\n1. I am not convinced the reward-uncertainty is significantly more useful than the dynamics uncertainty. As the authors describe, the de-facto standard way of formulating GPs is that output dimensions are considered independent predictions. The uncertainty of each dimension (the variance of the GP) of a GP conditioned on $X$ in a query-point $x$ is then given by $K(x,x) - K(x,X)K^-1(X,X)K(X,x)$ and depends only on the kernel choice and the data distribution. Unless one has specific knowledge about an appropriate reward-kernel, I have no intution  why the reward offers a better (or different) information gain signal than the dynamics.\n\n2. It is unclear to me why the proposed sampling distribution is meaningful. The sampling distribution described in Eq. 8 seems to imply a risk-sensitive sampling strategy without motivating this. For example, one may easily construct MDPs, where optimal trajectories are guaranteed to contain a very small reward. Such trajectories could not be sampled under this strategy, despite being optimal. Mathematically speaking, it is unclear to me what the implications of conditioning on an impossible event is here (i.e., what if $r_t$ can't be greater than $r_{min}$). I can unfortunately not see why one would expect this algorithm to explore reliably.\n\n\nMinor points: \n\n- The dimensions of the GP definition in line 269 don't seem to add up to me. The covariance matrix, as written here, has dimensions $2 \\times 2$, while the mean functions are of dimension $|\\mathcal{S}| + 1$. I moreover think the input dimensions of covariance functions $K$ defined in line 271 should operate on pairs of $S\\times A$, not simply $S \\times A$.\n\n- Expected successor states (as described in 5.2, line 323) don't seem sensible to me. The expectation of a high-dimensional state is not necessarily a meaningful or even existant state. It was moreover not clear to me, how to the model predictions are conditioned on a certain reward level in practice.\n\n- Given that the authors use a GP, I suspect many readers would appreciate a table (or similar) comparing the computational burden of this approach with alternative methods. Furthermore I would be interested to know the effect of subsampling smaller batches (as described in line 850) on the quality of uncertainty estimates with a GP.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "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": "", "extra_scores": {}}, {"review_id": "JGR2whRVaf", "reviewer": "Reviewer_GjSJ", "summary": "This manuscript introduces HOT-GP, a novel and efficient approach to model-based reinforcement learning that enhances performance through a unique combination of techniques. The core innovation of HOT-GP lies in its ability to learn the joint distribution of dynamics and rewards, paired with an optimistic sampling strategy for \"in-model\" policy optimization. By leveraging the latest model-based reinforcement learning techniques, HOT-GP selects models for policy optimization based on optimistic sampling, which promotes exploration in high-potential regions of the policy space. Evaluated across several continuous-action benchmarks, HOT-GP demonstrates superior performance compared to both baseline methods and other state-of-the-art algorithms.", "questions": "None", "limitations": "", "rating": 8, "confidence": 5, "soundness": 4, "presentation": 4, "contribution": 4, "strengths": "1) The use of the joint distribution of dynamics and rewards in the learning process is a clever and impactful approach, offering new perspectives in model-based reinforcement learning.\n2) The implementation is robust and well-executed, showcasing attention to both theoretical and practical details.\n3) The manuscript is excellently presented—clear, concise, and engaging. It's a pleasure to read and straightforward to follow.\n4) The experiments are adequately designed, with thoughtful discussions that provide meaningful insights into the method's strengths and limitations.\n5) Providing the code for replicating experiments is greatly appreciated, allowing for easy validation and further exploration by the community.\n6) The appendix is comprehensive and well-detailed, giving me confidence that I could implement the method from scratch!", "weaknesses": "1) The evaluation would benefit from more complex scenarios, such as object manipulation or long-horizon tasks, which could further demonstrate the method’s robustness and adaptability.\n\n2) While the integration of a neural network to learn the GP mean is discussed in the main text, I feel like a small paragraph or additional description in the appendix is required. I understand the general approach and can infer the authors' method, but a concise explanation would be a helpful addition for readers.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "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": "", "extra_scores": {}}, {"review_id": "TBrNNM6Xjg", "reviewer": "Reviewer_LvwR", "summary": "The paper proposes a model-based reinforcement learning algorithm called HOT-GP (Hallucination-based Optimistic Thompson sampling with Gaussian Processes). This approach leverages Gaussian Processes to jointly model the reward and transition dynamics, addressing both state and reward uncertainties. The main goal is to optimize exploration by focusing on high-reward regions in the state-action space. HOT-GP introduces a new optimism-driven Thompson sampling strategy to balance exploration and exploitation, which is particularly useful in environments with sparse rewards or penalties for actions.", "questions": "According to the paper, the HOT-GP algorithm is inspired by H-UCRL. Why, then, was H-UCRL not used as a comparative algorithm in the MuJoCo tasks?", "limitations": "", "rating": 5, "confidence": 4, "soundness": 2, "presentation": 2, "contribution": 2, "strengths": "The HOT-GP proposed in the paper performs joint uncertainty modeling between rewards and state dynamics. This method innovatively combines the uncertainties in rewards and state transitions, achieving more efficient exploration in environments with sparse rewards and high penalties.", "weaknesses": "1. The experiments are insufficient. For environments with sparse rewards and difficult exploration, tasks like Maze and Navigation, such as Sparse-Point-Robot, might better illustrate the effectiveness of the approach.\n2. In some environments, the training curves presented in the paper do not appear to have converged, leaving the convergence of the algorithm unclear and making it difficult to judge whether an optimal policy has been achieved.\n3. The paper seems to lack ablation studies on different modules and tasks, which fails to adequately demonstrate the independent contributions of various components to overall performance. Without detailed ablation studies, it becomes challenging to evaluate the individual effects of joint uncertainty modeling and the optimistic strategy across different tasks.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "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": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)