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
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 ⟨St,bt⟩ indicates a “joint” variable, however, note that St is stochastic whereas the belief evolves deterministically. Colored nodes are observed, blank nodes are latent.
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 × 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
λ
1.0
1.0
(outer-loop) Discount
γ
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).
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.
Jingyang You, Hanna Kurniawati
School of Computing Australian National University
We formulate episodic Markov decision process (MDP) planning as Bayesian inference over policies. The primary contribution is conceptual: the policy itself is treated as the latent variable, and expected return defines an unnormalized posterior density over policies. This preserves the standard expected-return objective, in contrast to trajectory-centric planning-as-inference formulations that introduce auxiliary optimality variables and to entropy-regularized policy optimization methods that solve a different objective. In the exact formulation, the posterior over deterministic policies induces what we define here as an optimal stochastic policy under preference uncertainty, namely the stochastic policy induced by that posterior. For discrete MDPs with stochastic transitions, we study variational sequential Monte Carlo (VSMC) as one approximate inference method for this posterior, introducing policy consistency under state revisitation and coupled transition randomness across particles. Experiments on grid worlds, Blackjack, Triangle Tireworld, and Academic Advising examine the consequences of inference over policies and compare its induced behavior with entropy-regularized policy optimization. The results support the view that MDP planning can be naturally cast as Bayesian inference over policies.
Effective exploration in high-dimensional continuous control remains a central challenge in reinforcement learning. Planning-based methods address this by combining online planning with learned policies and value functions, but their components can become misaligned during training: learned sampling policies may diverge from planner behavior, while planning distributions stored in replay become stale as the model and value function evolve. Reanalysis can refresh these targets, but at substantial computational cost. We propose GEM-MPC, an MPPI-based reinforcement learning method that improves the interaction between planning and learning. GEM-MPC uses MPPI to combine a policy trained to clone the planner with a KL-regularized policy that explores around it, providing complementary exploitation and guided exploration within planning. We further introduce Gated Prior Distillation, which selectively learns from stored planning distributions only when they provide a better target than the current prior, reducing the impact of stale planning data without requiring full reanalysis. Across continuous-control benchmarks, GEM-MPC consistently outperforms existing planning-based baselines under lower computational budgets.
Alvaro Serra-Gomez, Thomas Moerland
Leiden Institute of Advanced Computer Science, Leiden University, 2311 EZ Leiden, The Netherlands