BOTS: Batch Bayesian Optimization of Extended Thompson Sampling for Severely Episode-Limited RL Settings

TL;DR

The paper tackles learning under severe episode constraints in adaptive health interventions by extending Thompson sampling with a state-action utility that includes action-specific biases learned across episodes via batch Bayesian optimization. This extended Thompson sampling (xTS) broadens the policy space beyond standard TS while maintaining low variance, and batch BO enables parallel evaluation across episodes, incorporating micro-randomized trials to set priors. Across basic MDPs and a JITAI-like simulation with habituation and disengagement dynamics, the proposed BOTS approach demonstrates superior performance under limited episodes compared to standard TS and several full RL baselines, with further gains from local BO and reward-variance modeling. The work provides practical guidance for deploying Bayesian-optimized, bias-aware bandits in real-world, episode-constrained adaptive interventions and includes implementation resources.

Abstract

In settings where the application of reinforcement learning (RL) requires running real-world trials, including the optimization of adaptive health interventions, the number of episodes available for learning can be severely limited due to cost or time constraints. In this setting, the bias-variance trade-off of contextual bandit methods can be significantly better than that of more complex full RL methods. However, Thompson sampling bandits are limited to selecting actions based on distributions of immediate rewards. In this paper, we extend the linear Thompson sampling bandit to select actions based on a state-action utility function consisting of the Thompson sampler's estimate of the expected immediate reward combined with an action bias term. We use batch Bayesian optimization over episodes to learn the action bias terms with the goal of maximizing the expected return of the extended Thompson sampler. The proposed approach is able to learn optimal policies for a strictly broader class of Markov decision processes (MDPs) than standard Thompson sampling. Using an adaptive intervention simulation environment that captures key aspects of behavioral dynamics, we show that the proposed method can significantly out-perform standard Thompson sampling in terms of total return, while requiring significantly fewer episodes than standard value function and policy gradient methods.

Figures (5)

• Figure 1: Results for severely episode-limited settings. Note: the x-axes show (number of rounds, batch size) combinations, not round index. In all experiments, BOTS shows a better performance in a low number of rounds.
• Figure 2: BOTS overview in the JITAI setting, including setting TS initial priors using an MRT. The TS priors are propagated across episodes using one of two strategies: TS(Fixed) where we fix the same reward model prior for all rounds, or TS(Update) where we update the prior from round to round. The BO parameter can be: vector $\{\beta_a\}_{1:3}$, vector $\{\beta_a\}_{1:3}$, $\sigma_Y$, or vector $\{\beta_a\}_{1:3}$, $\{\sigma_{Ya}\}_{0:4}$.
• Figure 3: Sketch of the three basic MDPs. Each MDP is annotated with an indication of whether each method considered finds the optimal policy (✓) or fail to find the optimal policy (✗).
• Figure 4: Results for MDP1 (left), MDP2 (middle) and MDP3 (right).
• Figure 5: Baseline results without episode limits.