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

Sumit Vashishtha

stat.MLOct 20, 2026

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

Authors: Sumit Vashishtha

Organizations: Independent Researcher

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.

Figures & tables

Figure 1: Epistemic uncertainty on 2D annular regression. (A) Annular training data ( $r\in[1.3,2.45]$ ) with unobserved central cavity ( $r\leq 1.0$ ). (B–D) BALD epistemic uncertainty heatmaps (nats) ( Houlsby et al., 2011 ) on a shared $[0.0,1.8]$ scale for Deep Ensembles (B), BootDQN with randomized priors (C), and DP-BNN (D). (E) Void vs. data support BALD contrast across evaluated methods.

Figure 2: Empirical exploration and scaling on DeepSea ( Osband et al., 2019a ) . (a) Cumulative regret on DeepSea-20 (10 random seeds, shaded bands denote

\pm 1

standard deviation) comparing DP-DQN against BootDQN-RP, BDQN, and DP-DQN in the Bayesian bootstrap limit (

\alpha\to 0

; Section 3 ). (b) Empirical sample complexity across 30 grid sizes up to

N=50

(6 random seeds per size, error bars denote

\pm 1

standard deviation), showing empirical

\mathcal{O}(N^{1.45})

scaling for DP-DQN compared to

\mathcal{O}(N^{3})

for BootDQN-RP.

Appendix figures & tables8 assets

Supplementary material from the paper’s appendix.

Appendix

Figure 3: DP prior and posterior density estimates on the galaxy velocity dataset.

Figure 4: Stochastic optimism across base measures (MLP-20 with LayerNorm). (a) Cumulative return on RiverSwim-6: uniform

\mathcal{U}(0,1)

and standard normal

\mathcal{N}(0,1)

solve the task across 100% of seeds, while a pessimistic prior

\mathcal{N}(-1,1)

traps 80% of runs. (b) Cumulative regret on DeepSea-10: optimistic and zero-mean priors discover the goal without privileged target coordinates, whereas a pessimistic prior causes boundary collapse and linear regret.

Benchmark	Metric	Vanilla / Baseline	BootDQN-RP / PSRL	DP-BNN (Ours)
DeepSea $N=20$	Episodes to Solve	$>20{,}000$ (DQN)	$5{,}210\pm 604$	$\mathbf{602\pm 104}$
DeepSea $N=50$ ( $64\times 64$ MLP)	Episodes to Solve	$>10^{15}$ (DQN)	$>20{,}000$	$\mathbf{1{,}925\pm 184}$
RiverSwim-6 ( $\mathcal{N}(0,1)$ Prior)	Episodes to Solve	$44.0\pm 15.2$ (DQN)	$7.0\pm 1.4$ (PSRL)	$\mathbf{27.4\pm 8.1}$
Deceptive $N$ -Chain ( $N=10$ )	Cum. Regret	$1{,}420.5\pm 112.0$ (DQN)	$681.4\pm 92.6$	$\mathbf{153.2\pm 18.4}$
Wheel Bandit ( $\delta=0.9$ )	1.5k-Step Regret	$6{,}894.6\pm 1{,}514.6$ ( $\epsilon$ -Gr.)	$4{,}309.4\pm 829.2$	$\mathbf{2{,}803.3\pm 189.6}$
MountainCar-v0 (Continuous)	Episodes to Solve	$85\pm 12$ (DQN)	$15\pm 3$ (BootDQN)	$\mathbf{16.8\pm 7.4}$

Appendix

Table 1: Consolidated empirical evaluation across diverse decision-making and exploration benchmarks (mean

\pm

standard error across seeds). All DP-BNN results are obtained using a single neural network (single MLP-20 with LayerNorm for DeepSea-20, RiverSwim-6, and Deceptive

N

-Chain; 2-layer

64\times 64

MLP with LayerNorm for DeepSea-50 and MountainCar-v0; 3-layer MLP-64 for contextual bandits).

Figure 5: Architectural schematics and Markov Decision Process (MDP) graphs for the canonical deep exploration benchmarks evaluated in Appendix E . (a) DeepSea (

N\times N

DAG): An agent traverses from

(0,0)

across

N

stages. Taking the designated rightward action at every row reaches the jackpot goal at

(N-1,N-1)

with reward

+1.0

(paying a step penalty

-0.01/N

), whereas any leftward step permanently traps the agent in zero-reward branches. An unguided random policy discovers the goal with probability

2^{-N}

(

9.5\times 10^{-7}

for

N=20

;

8.8\times 10^{-16}

for

N=50

), requiring temporally consistent value function posterior samples. (b) RiverSwim-6: A 6-state chain where swimming downstream (

a_{0}

) is deterministic and trapped by an immediate sub-optimal reward

r=0.05

s_{0}

, while swimming upstream (

a_{1}

) fights against a strong current (

P(s+1\mid s,a_{1})=0.6

P(s\mid s,a_{1})=0.3

P(s-1\mid s,a_{1})=0.1

) to reach the distal goal at

s_{5}

(

r=1.0

). (c) Deceptive

N

-Chain: A sequential corridor where choosing

a_{0}

s_{0}

yields an immediate lure reward (

r=0.01

) that prematurely terminates the rollout. Only persistent, multi-step posterior exploration traverses the silent zero-reward corridor (

s_{1}\to\dots\to s_{N-1}

) to unlock the terminal jackpot (

r=1.0

Figure 6: Deep Bayesian Bandit Showdown on the Wheel Bandit (

\delta=0.9

K=5

arms) across 1,500 environment steps (evaluated over 5 independent random seeds). (a) Wheel Bandit Geometry and Exploration Trap: Contexts

x\in\mathbb{R}^{2}

are sampled uniformly from the unit disk. The safe interior (

\|x\|_{2}\leq\delta=0.9

, grey) accounts for

81\%

of all contexts, where base arm

0

yields

r=1.0

while arms

1..4

yield

r=0.0

. The high-payoff optimal arms (

r=50.0

) reside exclusively on the rare outer rim (

\|x\|_{2}>0.9

19\%

total probability,

\approx 4.75\%

per quadrant). (b) Mean Cumulative Regret

R(t)

with

\pm 1

Standard Error Bands: DP-BNN rapidly discovers the outer quadrants and sustains high optimal action selection, achieving the lowest cumulative regret (

2{,}803.3\pm 189.6

) and outperforming Bootstrapped NN + Priors (

4{,}309.4\pm 829.2

) and

\epsilon

-greedy (

6{,}894.6\pm 1{,}514.6

). (c) Final Regret and Network Footprint: DP-BNN operates with a single neural network (

1\times

memory footprint), whereas Bootstrapped NN + Priors requires 20 total networks (

20\times

memory) and Deep Ensembles requires 5 networks (

5\times

memory).

Hyperparameter	2D Annular Cavity	DeepSea-20	DeepSea-50	RiverSwim-6	Deceptive $N$ -Chain	MountainCar-v0	Wheel Bandit
State Space $\mathcal{S}$	$[-2.5,2.5]^{2}$	One-Hot $\mathbb{R}^{400}$	One-Hot $\mathbb{R}^{2500}$	One-Hot $\mathbb{R}^{6}$	One-Hot $\mathbb{R}^{10}$	Continuous $\mathbb{R}^{2}$	Unit Disk $\mathbb{R}^{2}$
Action Space $\mathcal{A}$	— (Regression)	Discrete (2)	Discrete (2)	Discrete (2)	Discrete (2)	Discrete (3)	Discrete (5)
Horizon $H$	—	$20$	$50$	$50$	$10$	$200$	$1$
Network Architecture	3-layer MLP-64	Single MLP-20	2-layer MLP-64	Single MLP-20	Single MLP-20	2-layer MLP-64	3-layer MLP-64
Layer Normalization	No	Yes	Yes	Yes	Yes	Yes	No
Optimizer	Adam	Adam	Adam	Adam	Adam	Adam	Adam

Appendix

Table 2: Consolidated experimental specifications and hyperparameters for DP-BNN and DP-DQN across all empirical benchmarks (

W=\text{---}

for static regression and contextual bandits because

H=1

with no sequential Bellman propagation).

Complexity Dimension	Vanilla DQN	Parametric BDQN	BootDQN-RP ( $K$ networks)	DP-DQN (Ours)
Active Neural Networks	$1$	$1$	$K$	$\mathbf{1}$
Total Stored Parameters	$2d$	$2d+\|\mathcal{S}_{\text{feat}}\|^{2}$	$3K\cdot d$	$\mathbf{2d}$
Parameter Memory Asymptotics	$\mathcal{O}(d)$	$\mathcal{O}(d+\|\mathcal{S}_{\text{feat}}\|^{2})$	$\mathcal{O}(K\cdot d)$	$\mathbf{\mathcal{O}(d)}$
Optimizer State Buffers (Adam)	$2d$	$2d$	$2K\cdot d$	$\mathbf{2d}$
Action Selection Latency / Step	$\mathcal{O}(1\cdot d)$	$\mathcal{O}(d+\|\mathcal{A}\|\cdot\|\mathcal{S}_{\text{feat}}\|)$	$\mathcal{O}(K\cdot d)$	$\mathbf{\mathcal{O}(1\cdot d)}$
Training FLOPs / Episode	$\mathcal{O}\left(\frac{H}{\Delta}\cdot B\cdot d\right)$	$\mathcal{O}\left(\frac{H}{\Delta}\cdot B\cdot d\right)$	$\mathcal{O}\left(\frac{H}{\Delta}\cdot K\cdot B\cdot d\right)$	$\mathbf{\mathcal{O}\left(W\cdot(B+T)\cdot d\right)}$

Appendix

Table 3: Computational and memory complexity comparison across value-based deep RL paradigms for an action-value network with

d=|\theta|

parameters, episode horizon

H

, replay batch size

B

, ensemble size

K

, and prior truncation level

T

Metric	Vanilla DQN	BDQN	BootDQN-RP	DP-DQN ( $W=10$ )	DP-DQN ( $W=4$ )
Active Networks	$1$	$1$	$20$	$\mathbf{1}$	$\mathbf{1}$
Total Networks (incl. Prior/Target)	$2$	$2$	$60$	$\mathbf{2}$	$\mathbf{2}$
Trainable Parameters ( $d$ )	$8{,}102$	$8{,}142$	$162{,}040$	$\mathbf{8{,}102}$	$\mathbf{8{,}102}$
Total Stored Parameters	$16{,}204$	$16{,}284$	$486{,}120$	$\mathbf{16{,}204}$	$\mathbf{16{,}204}$
Optimizer Instances (sets of $2d$ )	$1$	$1$	$20$	$\mathbf{1}$	$\mathbf{1}$
Intra-Episode Backprop Steps	$10$	$10$	$200$	$\mathbf{0}$	$\mathbf{0}$

Appendix

Table 4: Empirical computational complexity, parameter volume, and runtime profiling on DeepSea-20 across algorithms on a standardized single CPU core.

Bayesian RL

Bayesian reinforcement learning uses probability distributions over unknown rewards and dynamics to guide learning and action under uncertainty.

See all 34 papers

Sparse Bayesian Neural Networks via Sticky Piecewise Deterministic Markov Processes
Oct 23, 2026August Arnstad, Geir Olve Storvik, Filippo Pagani +1Bayesian Neural Networks Sparse Neural Networks
Department of Mathematics, University of Oslo, NO · Department of Statistics, University of Warwick, UK
Dimension-Free PAC-Bayes Generalization for Deep Neural Networks via Dirichlet Processes
Oct 15, 2026Sumit VashishthaPAC-Bayesian Generalization Bounds Neural Network Generalization
Independent Researcher
Position: The Time for Sampling Is Now! Charting a New Course for Bayesian Deep Learning
May 20, 2026Emanuel Sommer, David RügamerBayesian Neural Networks Efficient Neural Network Inference
Department of Statistics, LMU Munich, Munich, Germany · Munich Center for Machine Learning, Munich, Germany.

Oct 23, 2026J/K move · Enter open · S save

Sparse Bayesian Neural Networks via Sticky Piecewise Deterministic Markov Processes

August Arnstad, Geir Olve Storvik, Filippo Pagani, Leiv Rønneberg

Department of Mathematics, University of Oslo, NO · Department of Statistics, University of Warwick, UK

Bayesian Neural Networks Sparse Neural Networks Neural Network Pruning

Open paper AskPDF arXiv

There is growing evidence that the performance of a neural network depends largely on a small subset of its weights, which are typically identified by training a dense network and pruning it afterwards. Within a Bayesian framework, pruning can be directly incorporated through a spike-and-slab prior, providing a proper treatment of the uncertainties involved, both with respect to pruning and to parameter uncertainty. In Bayesian neural networks (BNNs), sampling by MCMC is notoriously hard. We demonstrate that modern Monte Carlo methods based on piecewise deterministic Markov processes (PDMPs) can be used to sample posteriors of BNNs asymptotically exact whenever numerically constructed rate bounds hold. Spike-and-slab priors can be incorporated through sticky versions of the PDMPs, setting weights exactly to zero during sampling. We examine the predictive performance and convergence of the samplers, and show that sparse BNNs sampled this way can be more robust to input noise than dense ones.

Abstract

Bayesian RL

Explore similar work

Sparse Bayesian Neural Networks via Sticky Piecewise Deterministic Markov Processes

Dimension-Free PAC-Bayes Generalization for Deep Neural Networks via Dirichlet Processes

Position: The Time for Sampling Is Now! Charting a New Course for Bayesian Deep Learning