VariBASed: Variational Bayes-Adaptive Sequential Monte-Carlo Planning for Deep Reinforcement Learning

Joery A. de Vries; Jinke He; Yaniv Oren; Pascal R. van der Vaart; Mathijs M. de Weerdt; Matthijs T. J. Spaan

cs.LGFeb 21, 2026

VariBASed: Variational Bayes-Adaptive Sequential Monte-Carlo Planning for Deep Reinforcement Learning

Authors: Joery A. de Vries, Jinke He, Yaniv Oren, Pascal R. van der Vaart, Mathijs M. de Weerdt, Matthijs T. J. Spaan

Organizations: Trent AI Ltd, London, United Kingdom · EEMCS, Delft University of Technology, Delft, the Netherlands

Abstract

Optimally trading-off exploration and exploitation is the holy grail of reinforcement learning as it promises maximal data-efficiency for solving any task. Bayes-optimal agents achieve this, but obtaining the belief-state and performing planning are both typically intractable. Although deep learning methods can greatly help in scaling this computation, existing methods are still costly to train. To accelerate this, this paper proposes a variational framework for learning and planning in Bayes-adaptive Markov decision processes that coalesces variational belief learning, sequential Monte-Carlo planning, and meta-reinforcement learning. In a single-GPU setup, our new method VariBASeD exhibits favorable scaling to larger planning budgets, improving sample- and runtime-efficiency over prior methods.

Figures & tables

Panel of Figure 1 — Figure 1: Generative model (left), and the unstructured (middle) vs. structured inference model (right) that we consider. Stochastic variables are placed in circles, deterministic variables in rectangles. The double circle for $\langle S_{t},b_{t}\rangle$ indicates a “joint” variable, however, note that $S_{t}$ is stochastic whereas the belief evolves deterministically. Colored nodes are observed, blank nodes are latent.

Panel of Figure 2 — Figure 2: Evaluation rollouts on the function optimization and gridworld problems over multiple test-episodes. Due to the stochastic nature of our method, we visualize one sample rollout for the function optimization problem (with average metrics), and the averaged rollouts for the gridworld. For the gridworld, the circle indicates the starting-tile and the cross the goal-tile. On the function optimization we used a planning budget of $H=1,K=16$ and on the gridworld $H=4,K=32$ .

Figure 3: Learning curves for all environments comparing our VariBASeD against recurrent PPO (RL 2 ) Duan et al., 2016a using the S5 architecture Lu et al., 2023a for different planning budgets. Shaded regions give 99% two-sided BCa-bootstrap intervals over 30 seeds. In our framework we found that the PPO baseline did not learn.

Appendix figures & tables3 assets

Supplementary material from the paper’s appendix.

Appendix

Figure 4: Visualization of environments. Left: continuous function optimization task, where each color corresponds to a distinct function that the agent has to learn to optimize from sequential interactions. Right: discrete opengrid task where the agent needs to find the goal tile

\times

over multiple episodes starting from the dot.

Name	Symbol	Value Function Env	Value Grid
SGD Minibatch size		1024	1024
SGD update steps		32	32
Unroll length (nr. steps in env)		64	128
Batch-Size (nr. parallel envs)		32	32
(outer-loop) TD-Lambda	$\lambda$	1.0	1.0
(outer-loop) Discount	$\gamma$	0.99	0.99

Appendix

Table 1: Experiment hyperparameters. Values in sets were varied as part of our ablations.

Package	Version
brax	0.10.5
optax	0.2.3
flashbax	0.1.2
rlax	0.1.6
flax	0.8.4
xland-minigrid	0.9.1

Appendix

Table 2: Software module versioning that we used for our experiments (also includes default parameter settings).

Explore similar work

Meta-RL with Bayesian Linear Task Models
Dec 24, 2025Jingyang You, Hanna KurniawatiRepresentation Learning Reinforcement Learning
School of Computing Australian National University
MDP Planning as Policy Inference
Feb 19, 2026David TolpinMarkov Decision Processes Bayesian Inference
Offtopia
GEM-MPC: Balancing Exploration and Exploitation through Expert-Guided Planning
Sep 18, 2026Alvaro Serra-Gomez, Thomas MoerlandRL Exploration Model-Based Planning
Leiden Institute of Advanced Computer Science, Leiden University, 2311 EZ Leiden, The Netherlands

Dec 24, 2025J/K move · Enter open · S save

Meta-RL with Bayesian Linear Task Models

Jingyang You, Hanna Kurniawati

School of Computing Australian National University

Representation Learning Reinforcement Learning Bayesian RL Meta-Reinforcement Learning Generalization in RL Model-Based RL

Open paper AskPDF arXiv

Deep Bayesian reinforcement learning adapts to unseen tasks by inferring latent transition and reward models, but existing methods typically rely on variational posteriors and evidence lower bounds, introducing approximation error and unstable task representations. We introduce GLiBRL, a deep Bayesian RL framework that combines generalised linear task models with learnable non-linear basis functions. GLiBRL features conjugate Bayesian inference, yielding exact, sequential posterior updates over task parameters and model noise, together with a closed-form marginal likelihood that eliminates variational inference. The update is naturally permutation-invariant, allowing GLiBRL to integrate with both off- and on-policy algorithms. GLiBRL also learns task representation admitting an exact kernel identity, relating distances between task representations to kernel discrepancies over the task contexts. Compared against eight representative or recent meta reinforcement learning methods, GLiBRL achieves the highest aggregate zero-shot test performance on both the MuJoCo locomotion and MetaWorld manipulation benchmarks.

Abstract

Explore similar work

Meta-RL with Bayesian Linear Task Models

MDP Planning as Policy Inference

GEM-MPC: Balancing Exploration and Exploitation through Expert-Guided Planning