Efficient Model-Based Reinforcement Learning Through Optimistic Thompson Sampling

Authors: Jasmine Bayrooti, Carl Henrik Ek, Amanda Prorok
Primary area: reinforcement learning
Keywords: reinforcement learning, model-based reinforcement learning, optimistic exploration

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.

Decision: accept (poster)
Reviewer scores: 8, 3, 8, 5

Reviews and discussion:
<thread>
creation date:
1730383022359

thread_type:
official_review

<thread_object>
creator: reviewer_lvwr

<content>
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.

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.
2. 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.
3. 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.

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?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730452201826

thread_type:
official_review

<thread_object>
creator: reviewer_gjsj

<content>
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.

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.
2) The implementation is robust and well-executed, showcasing attention to both theoretical and practical details.
3) The manuscript is excellently presented—clear, concise, and engaging. It's a pleasure to read and straightforward to follow.
4) The experiments are adequately designed, with thoughtful discussions that provide meaningful insights into the method's strengths and limitations.
5) Providing the code for replicating experiments is greatly appreciated, allowing for easy validation and further exploration by the community.
6) 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.

2) 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.

questions:
None
</content>
</thread_object>

</thread>
<thread>
creation date:
1730471716228

thread_type:
official_review

<thread_object>
creator: reviewer_qgwe

<content>
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.

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:

1. 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.

2. 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.


Minor points: 

- 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$.

- 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.

- 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.

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.

2. 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. 

3. 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. 

My suggestion would be to analyze the implications of the proposed algorithmic components carefully. 
1. What is the implication of performing Thompson sampling with a sampling probability dependent on the expected return of some policy?
2. 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).
3. What is the implication of performing Thompson sampling with limited branched rollouts?

Addressing these questions in my opinion would provide a compelling motivation for the proposed algorithm.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730720864186

thread_type:
official_review

<thread_object>
creator: reviewer_ab8v

<content>
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.

strengths:
- principled approach for exploration in continuous state-action space
- 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,
    - SAC and PPO as described in Line 430, “we use … SAC, and the on-policy PPO”
    - HOT-GP with k-branching (k=5) as described in Line 470.
- 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.
- 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.

questions:
Could you provide any results on computational cost or computation speed for the proposed method?
</content>
</thread_object>

</thread>
<thread>
creation date:
1732561791425

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
Updated Paper Submission

comment:
We would like to thank the reviewers again for their time and valuable comments. A revised version of the manuscript has been uploaded, incorporating the following main changes marked in maroon color:

- Conducted additional experiments on two Sparse Point Mazes from the D4RL suite (Figure 2)
- Extended the Sparse Reacher experiments, demonstrating that HOT-GP’s asymptotic performance matches that of SAC (Figure 6)
- Added an extra experiment on Sparse Reacher with the intermediate action penalty of 0.3 (Figure 7)
- Increased the number of seeds for MuJoCo experiments to 10 per algorithm (Figure 1)
- Updated results on the Coverage task with the correct corresponding algorithms’ performance (Figure 3)
- Provided details on using neural networks as GP mean functions (Appendix B)
- Inserted Table 2 to specify the average runtime for each MuJoCo task and algorithm (Appendix B)

We believe that these changes have substantially enhanced the quality of the paper and invite the reviewers to reconsider our submission, potentially adjusting their scores. Notably, our evaluations across three test suites demonstrate the benefits of our novel approach, namely modeling and leveraging optimism with joint uncertainty. These advantages are particularly pronounced in challenging learning settings involving sparse rewards, difficult exploration, and action penalties. If there are any remaining concerns, please let us know and we will address them!
</content>
</thread_object>

</thread>


Metareview:
metareview:
The paper proposes HOT-GP, a model-based Thompson sampling-based RL method that takes into account the *joint* uncertainty in dynamics and reward, as modeled by a joint Gaussian Process, when making exploration decisions. The empirical evaluation shows it to be generally more sample-efficient than several baselines, including vanilla SAC, on MuJoCo and maze tasks.

Strengths:
- Novelty of the sampling distribution for exploration
- Good practical implementation
- Clarity

Weaknesses:
- Not entirely clear why it works better than other exploration types, since the uncertainty is estimated over instantaneous reward, not cumulative return
- The empirical evaluation isn't very strong, done on simple benchmarks, although it has definitely improved after the discussion with the reviewers

The metareviewer recommends this work for acceptance due to the novelty of the method, the empirical results showing its promise, and further analysis/discussion it can spark in the community. The metareviewer partly agrees with reviewer *QgWe* that it's not obvious why exploring based on uncertainty in both reward and dynamics gives better results than based on uncertainty in the dynamics alone. One possible explanation is that this simply allows RL to have a more complete picture of uncertainty and hence resolve the uncertainty about the model more quickly, more even if the uncertainty is over instantaneous reward rather than cumulative return. In any case, the metareviewer doesn't view the lack of complete clarity in this aspect as a blocker for publication, but rather as something that will give the research community food for further thought.