Nonparametric Bayesian Inverse Reinforcement Learning with Data-Parallel Gibbs Sampling
2026-07-10 • Machine Learning
Machine Learning
AI summaryⓘ
The authors address a problem where learning from multiple experts with different preferences usually results in a confusing average reward. They introduce a method that can figure out how many different types of rewards there are and what those rewards look like, using a special kind of statistical model called a Dirichlet Process. Their approach successfully identifies distinct reward groups in simulated tasks better than previous methods and can run faster by using multiple CPU cores. They also explore trade-offs between speed and accuracy when combining results from parallel computations.
Inverse Reinforcement LearningDirichlet ProcessNonparametric BayesianChinese Restaurant ProcessGibbs SamplingMetropolis-HastingsValue IterationObjectWorldParallel ComputingAdjusted Rand Index
Authors
Sai Anirudh Katupilla, Shreeya Dasa Lakshminath
Abstract
Inverse Reinforcement Learning recovers reward functions from expert demonstrations, but standard formulations assume that all demonstrations come from a single expert. When demonstrations are pooled from multiple experts with distinct preferences, parametric methods recover an averaged reward that fits no individual expert well. We implement Nonparametric Bayesian Inverse Reinforcement Learning with a Dirichlet Process prior over reward functions, allowing the number of latent reward types to be inferred jointly with the rewards themselves. Inference uses a collapsed Gibbs sampler combining a Chinese Restaurant Process update for cluster assignments with a Metropolis-Hastings update for reward weights, and soft value iteration as the inner planning routine. We evaluate on a 10x10 ObjectWorld grid with two and three ground-truth reward types. The serial sampler recovers K=2 with Adjusted Rand Index of 1.000, substantially outperforming a Maximum Entropy IRL baseline (ARI=0.000). Extension to K=3 shows that the sampler correctly identifies the number of clusters in all runs; assignment ARI of 0.48-0.58 reflects behavioral overlap between expert types that persists across grid instantiations, revealing that reliable K=3 evaluation on ObjectWorld requires controlled object placement rather than random seeding. We further parallelize the sampler across CPU cores using Ray on HPC hardware, achieving a peak speedup of 4.79x at 8 workers, and characterize a throughput-versus-accuracy tradeoff arising from the consensus merge heuristic used during state aggregation. Code and a containerized environment are available at https://github.com/dasashreeya/np_bayes_irl.