Bayesian Neural Networks via Dirichlet Process Priors on the Data-Generating Distribution for Randomized Value Functions in Deep Reinforcement Learning

AI summary

Overview

Research area: Bayesian deep learning and Bayesian nonparametrics applied to exploration in deep reinforcement learning (stat.ML).

Technical level: Advanced. The paper assumes familiarity with Dirichlet Processes, stick-breaking constructions, Thompson sampling, and deep Q-learning.

Scope: The paper proposes a class of Bayesian neural networks that place a Dirichlet Process prior on the data-generating distribution instead of on network weights, and instantiates them as DP-DQN, a single-network randomized-value-function algorithm for deep exploration.

What This Paper Is About

Standard Bayesian neural networks put priors over network weights, but the resulting posteriors are intractable to sample and the usual approximations (MC dropout, variational inference) tend to lose epistemic uncertainty exactly where it matters, which is a known failure mode in reinforcement learning. This paper instead puts a Dirichlet Process prior directly on the unknown distribution that generates the data (or, in RL, the transition tuples), so that uncertainty is expressed in function space rather than weight space. The goal is to obtain principled Thompson sampling over value functions, with temporally consistent "deep exploration," using a single neural network rather than a multi-network ensemble.

Key Contributions

1. DP-BNN formulation and sampling algorithm. A new class of Bayesian neural networks that quantifies epistemic uncertainty by placing a nonparametric Dirichlet Process prior on the data-generating distribution rather than on parameters, together with a tractable posterior sampling procedure (Algorithm 1) built on the Sethuraman stick-breaking construction and a multi-step posterior recursion. The method is validated on a spatial 2D annular regression task using BALD uncertainty heatmaps.

2. DP-DQN for randomized value functions. An instantiation of DP-BNNs in deep Q-learning that optimizes an action-value network on random measures drawn from the DP posterior, using a composite DP-Bellman loss that combines reweighted empirical transitions with synthetic prior atoms drawn from the base measure. Episodic boundary warmstarts plus rollout freezing give exact multi-step temporal consistency with one network.

3. Information-calibrated prior weighting. A mechanism that replaces the raw replay-buffer count with an effective sample counter that accumulates normalized Bellman surprise, preventing "prior swamping" in which thousands of early unrewarded transitions would drive the prior weight to zero.

4. Empirical benchmarking across exploration tasks. Evaluation on DeepSea up to horizon N = 50 (a 2^50 ≈ 10^15 trajectory search space, with reported empirical O(N^1.45) sample complexity), RiverSwim-6, deceptive N-Chain, continuous MountainCar-v0, and contextual Wheel Bandit, reported as matching or outperforming multi-network ensembles while keeping a single-network footprint.

Main Findings

• Uncertainty contrast in unobserved regions: On the 2D annular regression benchmark (training radii r ∈ [1.3, 2.45], unobserved central cavity r ≤ 1.0, aleatoric noise σ = 0.30, J = 10 sampled models, T = 50 prior atoms, α = 10.0), DP-BNN produced BALD uncertainty of 1.41 nats in the void versus 0.18 nats on the data support, an 8.0× contrast ratio.

• Comparison with baselines on the same benchmark: BootDQN with randomized priors gave 0.54 vs. 0.14 nats (3.8×), MC Dropout 0.36 vs. 0.16 nats (2.2×), Deep Ensembles 0.29 vs. 0.06 nats (4.8×), and Bayes By Backprop produced near-uniform uncertainty at 0.07 vs. 0.09 nats.

• Training cost is comparable to ensembles: Wall-clock training times for J = 10 models were DP-BNN [5.9 s], BootDQN+priors [5.4 s], Deep Ensembles [6.9 s], BBB [2.1 s], and MC Dropout [0.7 s].

• Single network suffices for Thompson sampling in RL: Because each sampled DP posterior measure corresponds to one value-function hypothesis, an agent needs only one sampled model per episode, in contrast to ensembles of K networks plus K prior networks.

• Bayesian linear readouts are insufficient: The paper reports that BDQN (Azizzadenesheli et al., 2018), which applies Bayesian linear regression to final-layer representations, does not resolve the DeepSea task in their evaluations, and argues this follows from restricting inference to the linear readout.

• Prior mass decays with information, not raw data count: Under the stick-breaking posterior, the retained prior weight W_prior = ∏(1 − V_i) follows Beta(α + 1, n) in the standard case with expectation (α + 1)/(α + 1 + n); DP-DQN instead uses Beta(α + 1, N_stat), where N_stat accumulates min(1, |δ_i| / σ_TD) over mini-batches.

• Detailed per-benchmark numbers are not reported in the available text: The results section is truncated, so specific scores, sample-complexity figures for individual baselines, and full tables for DeepSea, RiverSwim-6, N-Chain, MountainCar-v0, and Wheel Bandit beyond the abstract-level claims are not available in the content provided.

Methodology in Plain English

The authors reframe the source of uncertainty. Instead of asking "what is the posterior over the network's weights?", they ask "what is the posterior over the distribution that generated the data?" A Dirichlet Process is a prior over distributions, and it has a convenient property: after seeing data, the posterior is a weighted blend of what you observed and what you believed beforehand. The before-belief part, the base measure, can be used to generate synthetic "pseudo-observations" in regions the agent has never visited. This is what keeps uncertainty alive in unexplored parts of the state space, whereas a purely empirical bootstrap can only resample things already seen.

Concretely, sampling a model means: draw synthetic data points from the base measure, draw random weights from the stick-breaking construction (which allocates probability mass across real and synthetic points), form a weighted empirical distribution, and train the network on that reweighted objective. Setting the concentration parameter to zero recovers Rubin's Bayesian Bootstrap, so the method nests the standard bootstrap as a special case.

For reinforcement learning, the base measure generates synthetic transition tuples (state, action, reward, next state, done) with support covering the full state and action spaces. The training loss combines the reweighted real transitions with the synthetic prior transitions. Two engineering choices make this work in practice: (1) the prior's weight is scaled by how surprising the observed transitions are (measured by Bellman error), so that an early stream of boring unrewarded transitions does not wash out the prior; and (2) the network is updated only at episode boundaries with ⌈H/2⌉ warmstart gradient steps, then frozen for the entire rollout, so the sampled value function stays consistent across a whole episode, which is what produces coherent multi-step exploration rather than per-step dithering.

Why This Matters

Impact on research. The paper offers an alternative axis for Bayesian deep learning: placing the prior in data space rather than weight space sidesteps the well-documented difficulty of designing meaningful high-dimensional parameter priors, and it directly addresses the empirical finding that weight-space approximate BNNs fail to sustain epistemic uncertainty in deep RL. It also connects the Bayesian Bootstrap literature and Bayesian nonparametrics to posterior sampling for reinforcement learning, and shows that a single network can, in principle, replace the K-network ensembles that dominate the deep exploration literature.

Real-world applications (illustrative, based on the exploration problem class rather than experiments in the paper):

• Sparse-reward sequential decision problems such as recommendation or content-serving pipelines, where reward signals are delayed and rare.
• Clinical trial and adaptive treatment allocation, where Thompson sampling over competing arms is standard and calibrated uncertainty is safety-relevant.
• Robotics and autonomous control with expensive or dangerous exploration, where coherent multi-step exploration is preferable to random dithering.
• Resource allocation and dynamic pricing under bandit feedback, where the contextual Wheel Bandit setting is a common abstraction.

Industry relevance. The single-network footprint matters for deployment: ensemble-based deep exploration multiplies memory and compute by K, which is often the binding constraint in production RL systems. Comparable training wall-clock times to existing ensemble baselines suggest the approach could be retrofitted into existing DQN pipelines without a proportional cost increase.

Future Directions

• Scaling beyond the reported benchmarks. Extending evaluation to high-dimensional pixel-based and continuous-control domains, and reporting full comparative sample-complexity tables, is the natural next step beyond the environments listed here.

• Principled choice of the base measure. The results hinge on how the base measure F_0 is specified, especially the reward prior, which the authors say is designed to sustain exploratory variance in unvisited regions. How to elicit F_0 from domain knowledge, or learn it, remains open.

• Tuning of the concentration parameter and prior atom count. The concentration α, truncation level T, and the Bellman-surprise normalization σ_TD are all hyperparameters whose sensitivity is not characterized in the available text.

• Theoretical guarantees. The paper argues informally that restricting inference to a linear readout cannot capture uncertainty in feature representations, and reports empirical verification; formal regret or sample-complexity guarantees for DP-DQN relative to PSRL are a natural target.

Target Audience

Researchers and graduate students in Bayesian machine learning, Bayesian nonparametrics, and deep reinforcement learning who work on exploration, uncertainty quantification, or posterior sampling. It is also relevant to practitioners who need calibrated uncertainty in sequential decision systems but cannot afford K-network ensembles. Readers without background in Dirichlet Processes, stick-breaking, and Thompson sampling will find the paper demanding.

Authors’ abstract

Bayesian neural networks (BNNs) offer a principled framework for quantifying epistemic uncertainty. However, standard weight-space priors yield complex, high-dimensional neural network posteriors that are computationally prohibitive to sample from, and popular approximate inference schemes often fail to capture epistemic uncertainty effectively---a failure mode particularly acute in reinforcement learning (RL) \citep{osband2018randomized}. To resolve this challenge, we introduce a new class of BNNs that quantify uncertainty by placing Bayesian nonparametric Dirichlet Process (DP) priors directly on the underlying data-generating distribution rather than on network parameters. Building on this foundation, we develop DP-DQN, a reinforcement learning algorithm that leverages DP-BNNs to maintain posterior distributions over value functions and perform efficient Thompson sampling. By grounding randomized value functions in Bayesian nonparametrics, DP-DQN unifies data-driven posterior reweighting (via the Bayesian Bootstrap) with generative prior hypotheses injected directly through the DP base measure, enabling decisive ``deep exploration'' with a single neural network without the overhead of multi-network ensembles.