Nonparametric Bayesian Inverse Reinforcement Learning with Data-Parallel Gibbs Sampling

# Nonparametric Bayesian Inverse Reinforcement Learning with Data-Parallel Gibbs Sampling
Source: [https://arxiv.org/html/2607.09886](https://arxiv.org/html/2607.09886)
###### 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 a10×1010\\times 10ObjectWorld grid with two and three ground\-truth reward types\. The serial sampler recoversK=2K=2with Adjusted Rand Index of 1\.000, substantially outperforming a Maximum Entropy IRL baseline \(ARI=0\.000\)\. Extension toK=3K=3shows 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 of4\.79×4\.79\\timesat 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 athttps://github\.com/dasashreeya/np\_bayes\_irl\.

## IIntroduction

Reinforcement learning algorithms learn control policies from a reward function, but in many practical domains the reward function itself is unknown and must be inferred from observed expert behavior\. Inverse Reinforcement Learning \(IRL\) addresses this inverse problem: given a Markov Decision Process and a set of expert trajectories, recover the reward function that best explains the observed actions\. IRL has applications in robotics, autonomous driving, medical decision support, and human\-AI interaction, where engineering an explicit reward function is impractical but expert demonstrations are available\[[1](https://arxiv.org/html/2607.09886#bib.bib1)\]\.

Most IRL formulations encounter two fundamental difficulties\. The*modeling problem*: real demonstration datasets rarely come from a single homogeneous expert\. Fitting one reward to a mixed pool produces an average that represents nobody\. Parametric Bayesian IRL\[[2](https://arxiv.org/html/2607.09886#bib.bib2)\]provides principled uncertainty over the reward but still requires the number of reward types to be fixed in advance — an assumption that is unreasonable when the goal is to discover latent structure\. A Dirichlet Process prior removes this requirement, letting the number of components grow as the data demands\. The*computational problem*: MCMC inference over reward posteriors is expensive because each sample requires evaluating trajectory likelihood via inner planning\. For a nonparametric model the inference cost scales with both trajectory count and cluster count, making parallelization essential for practical use\.

We address both problems in a single system\. Our contributions are:

- •A complete open\-source implementation of nonparametric Bayesian IRL using a collapsed Gibbs sampler with Chinese Restaurant Process cluster updates and Metropolis\-Hastings weight updates, verified against ground\-truth rewards on the ObjectWorld benchmark\.
- •A data\-parallel extension using Ray on HPC, achieving4\.79×4\.79\\timespeak speedup and characterizing the throughput\-versus\-accuracy tradeoff that emerges from the consensus merge step\.
- •Empirical evaluation across two and three expert types, including an analysis of how feature discriminability in the environment bounds recovery quality independently of sampler behavior\.
- •A reproducible, containerized system \(Docker image and SLURM scripts\) enabling direct replication on HPC infrastructure\.

## IILiterature Survey

### II\-AInverse Reinforcement Learning

Ng and Russell\[[1](https://arxiv.org/html/2607.09886#bib.bib1)\]formulated IRL as recovering a reward function consistent with expert optimality, but the original linear programming approach suffered from reward ambiguity\. Abbeel and Ng\[[3](https://arxiv.org/html/2607.09886#bib.bib3)\]introduced apprenticeship learning via feature matching\. Ziebart et al\.\[[4](https://arxiv.org/html/2607.09886#bib.bib4)\]proposed Maximum Entropy IRL, which selects the maximum entropy distribution over trajectories subject to feature\-matching constraints, providing a principled probabilistic framing\. MaxEnt IRL is the comparison baseline in this work\.

### II\-BBayesian and Nonparametric IRL

Ramachandran and Amir\[[2](https://arxiv.org/html/2607.09886#bib.bib2)\]formulated IRL as Bayesian inference over reward functions under a Boltzmann\-rational expert model\. We adopt their likelihood directly and extend it to a nonparametric multi\-expert setting\. Babes\-Vroman et al\.\[[5](https://arxiv.org/html/2607.09886#bib.bib5)\]proposed a parametric mixture\-of\-experts model with a fixed number of components\. Choi and Kim\[[6](https://arxiv.org/html/2607.09886#bib.bib6)\]are the closest prior work, presenting a nonparametric Bayesian IRL framework with a collapsed Gibbs sampler; we build directly on their formulation and extend it with a parallel inference engine\. Post\-2012 work has explored neural reward parameterizations\[[16](https://arxiv.org/html/2607.09886#bib.bib16)\]and continuous state spaces\[[14](https://arxiv.org/html/2607.09886#bib.bib14)\], but parallelization of the Gibbs inference procedure has received limited attention\. We address this gap\.

### II\-CDirichlet Process Mixture Models

The Dirichlet Process prior was introduced by Ferguson\[[7](https://arxiv.org/html/2607.09886#bib.bib7)\]\. The Chinese Restaurant Process representation\[[8](https://arxiv.org/html/2607.09886#bib.bib8)\]gives an intuitive sequential construction amenable to Gibbs sampling\. Neal\[[9](https://arxiv.org/html/2607.09886#bib.bib9)\]surveyed MCMC algorithms for DP mixtures; Algorithm 8 of that survey is the basis for our collapsed sampler\.

### II\-DParallel MCMC

Gibbs sampling exhibits conditional independence structure within each sweep that admits parallel computation\. Lovett et al\.\[[10](https://arxiv.org/html/2607.09886#bib.bib10)\]analyze parallelization of hierarchical Bayesian models, showing near\-linear speedup when conditional independence is exploited\. Our sampler has exactly this structure: assignments are conditionally independent across trajectories given weights, and weights are conditionally independent across clusters given assignments\. We use Ray\[[12](https://arxiv.org/html/2607.09886#bib.bib12)\]for distribution, as its shared object store avoids repeated serialization of read\-only environment tensors\.

### II\-EChoice of Method

We use the Choi and Kim framework rather than parametric mixtures or GP\-IRL for three reasons\. The DP prior handles cluster count uncertainty without discrete model selection, which is central to our evaluation goal\. The Gibbs sampler’s per\-trajectory and per\-cluster updates decompose naturally for parallel execution\. And the framework stays tractable on grid worlds where ground\-truth rewards are known, enabling rigorous correctness validation before scaling\.

## IIIMethodology

### III\-AProblem Setup

We consider a discrete Markov Decision Processℳ=\(𝒮,𝒜,T,γ\)\\mathcal\{M\}=\(\\mathcal\{S\},\\mathcal\{A\},T,\\gamma\)\. The reward function is parameterized asR​\(s\)=ϕ​\(s\)⊤​wR\(s\)=\\phi\(s\)^\{\\top\}w, whereϕ​\(s\)∈ℝd\\phi\(s\)\\in\\mathbb\{R\}^\{d\}is a feature vector andw∈ℝdw\\in\\mathbb\{R\}^\{d\}is a weight vector\. We observeNNtrajectoriesτi=\{\(st,at\)\}t=0H−1\\tau\_\{i\}=\\\{\(s\_\{t\},a\_\{t\}\)\\\}\_\{t=0\}^\{H\-1\}generated by experts whose true weights are unknown and may differ across trajectories\.

### III\-BGenerative Model

The generative model places a Dirichlet Process prior over reward weights:

G\\displaystyle G∼DP​\(α,G0\)\\displaystyle\\sim\\mathrm\{DP\}\(\\alpha,G\_\{0\}\)\(1\)wi\\displaystyle w\_\{i\}∼G\\displaystyle\\sim G\(2\)τi\\displaystyle\\tau\_\{i\}∼p​\(τ∣wi\)\\displaystyle\\sim p\(\\tau\\mid w\_\{i\}\)\(3\)whereG0=𝒩​\(0,I\)G\_\{0\}=\\mathcal\{N\}\(0,I\)andα\\alphais the concentration parameter\. Under the Boltzmann\-rational expert model\[[2](https://arxiv.org/html/2607.09886#bib.bib2)\]:

p​\(τ∣w\)=∏tπwβ​\(at∣st\),πwβ​\(a∣s\)∝exp⁡\(β​Qw​\(s,a\)\)p\(\\tau\\mid w\)=\\prod\_\{t\}\\pi\_\{w\}^\{\\beta\}\(a\_\{t\}\\mid s\_\{t\}\),\\quad\\pi\_\{w\}^\{\\beta\}\(a\\mid s\)\\propto\\exp\(\\beta Q\_\{w\}\(s,a\)\)\(4\)whereQwQ\_\{w\}is computed by soft value iteration with 100 fixed Bellman backups per call\.

### III\-CCollapsed Gibbs Sampler

Inference alternates two updates per sweep\.

#### Cluster assignment update

p​\(zi=k∣⋅\)∝\{n−i,k⋅p​\(τi∣wk\)existing​kα⋅p​\(τi∣wnew\)new clusterp\(z\_\{i\}=k\\mid\\cdot\)\\propto\\begin\{cases\}n\_\{\-i,k\}\\cdot p\(\\tau\_\{i\}\\mid w\_\{k\}\)&\\text\{existing \}k\\\\ \\alpha\\cdot p\(\\tau\_\{i\}\\mid w\_\{\\text\{new\}\}\)&\\text\{new cluster\}\\end\{cases\}\(5\)wheren−i,kn\_\{\-i,k\}excludes trajectoryiifrom its own cluster count\.

#### Weight update

Proposewk′=wk\+ϵw^\{\\prime\}\_\{k\}=w\_\{k\}\+\\epsilon,ϵ∼𝒩​\(0,σ2​I\)\\epsilon\\sim\\mathcal\{N\}\(0,\\sigma^\{2\}I\), and accept with Metropolis\-Hastings probability

min⁡\(1,p0​\(wk′\)​∏i:zi=kp​\(τi∣wk′\)p0​\(wk\)​∏i:zi=kp​\(τi∣wk\)\)\.\\min\\\!\\left\(1,\\frac\{p\_\{0\}\(w^\{\\prime\}\_\{k\}\)\\prod\_\{i:z\_\{i\}=k\}p\(\\tau\_\{i\}\\mid w^\{\\prime\}\_\{k\}\)\}\{p\_\{0\}\(w\_\{k\}\)\\prod\_\{i:z\_\{i\}=k\}p\(\\tau\_\{i\}\\mid w\_\{k\}\)\}\\right\)\.\(6\)All likelihood evaluations use log\-space arithmetic to prevent underflow over multi\-step trajectories\.

### III\-DHyperparameter Sensitivity

Two hyperparameters dominate sampler behavior\. The DP concentrationα\\alphacontrols the prior probability of opening new clusters; values that are too small cause collapse toK=1K=1\. The MH proposal scaleσ\\sigmacontrols exploration in weight space; values that are too large prevent acceptance and values that are too small slow mixing\.

Table[I](https://arxiv.org/html/2607.09886#S3.T1)and Table[II](https://arxiv.org/html/2607.09886#S3.T2)report results across a grid of values run with seed=0 and 300 sweeps\. Two findings are notable\. First,σ\\sigmais nearly irrelevant to ARI: results atσ∈\{0\.01,0\.05,0\.1\}\\sigma\\in\\\{0\.01,0\.05,0\.1\\\}are identical, and onlyσ=0\.001\\sigma=0\.001differs \(worse, due to slow mixing from very small proposals\)\.ℓ2\\ell\_\{2\}error increases modestly with largerσ\\sigma, consistent with noisier proposals reducing weight precision without affecting cluster structure\. Second,α≤5\.0\\alpha\\leq 5\.0all converge to the same mode \(ARI=0\.458,K=3K=3\) at 300 sweeps, suggesting the chain finds a single dominant intermediate mode regardless of concentration\.α=7\.0\\alpha=7\.0achieves the best ARI \(0\.667,K=4K=4\), whileα=10\.0\\alpha=10\.0over\-fragments slightly\. Notably, the main experiments use 500 sweeps and achieve ARI=1\.000 atα=5\.0\\alpha=5\.0, indicating the chain requires more sweeps to escape the intermediate mode visible at 300\. We useα=5\.0\\alpha=5\.0andσ=0\.01\\sigma=0\.01for all main experiments\.

TABLE I:α\\alphaablation \(σ=0\.01\\sigma=0\.01fixed, 300 sweeps, seed=0; main experiments use 500 sweeps, see Section[V](https://arxiv.org/html/2607.09886#S5)\)TABLE II:σ\\sigmaablation \(α=5\.0\\alpha=5\.0fixed, 300 sweeps, seed=0; main experiments use 500 sweeps, see Section[V](https://arxiv.org/html/2607.09886#S5)\)
### III\-EImplementation

The system is implemented in JAX\[[11](https://arxiv.org/html/2607.09886#bib.bib11)\]across approximately 2,100 lines of Python\. We JIT\-compile the inner Boltzmann log\-policy computation, which dominates per\-sweep runtime; outer control flow remains plain Python to preserve flexibility for cluster relabeling\. Correctness is verified by 13 unit tests covering likelihood evaluation, evaluation metrics, Gibbs sweep behavior including pickle\-safety for Ray serialization, and the ObjectWorld environment\.

### III\-FParallelization

Trajectories are partitioned acrossWWRay workers\. Each worker runs one Gibbs sweep on its partition; results are merged after each sweep\. Read\-only environment tensorsϕ\\phiandTTare placed in Ray’s distributed object store once at initialization\.

The state merge step is the principal source of approximation\. Workers perform independent MH updates, causing weight vectors to diverge slightly\. The merge uses element\-wise comparison with tolerance10−610^\{\-6\}to identify corresponding clusters\. When post\-MH drift exceeds this tolerance, the same cluster across two workers is counted as two separate clusters, inflatingKK\. We characterize this empirically in Section[V](https://arxiv.org/html/2607.09886#S5)\.

### III\-GHardware and Software

Experiments were run on the University of Maryland Zaratan cluster on AMD EPYC Zen2 nodes \(16 cores, 32 GB RAM, SLURM scheduling\)\. Stack: Python 3\.10, JAX 0\.4\.35 \(CPU\), Ray 2\.55, NumPy 2\.2, scikit\-learn 1\.7, Matplotlib 3\.10\. The full environment is packaged as a Docker image athttps://hub\.docker\.com/r/shreeyadasa/npbayes\-irl\.

## IVExperiments

### IV\-AEnvironment

We use the ObjectWorld benchmark\[[13](https://arxiv.org/html/2607.09886#bib.bib13)\], a10×1010\\times 10grid with\|𝒮\|=100\|\\mathcal\{S\}\|=100states and\|𝒜\|=4\|\\mathcal\{A\}\|=4actions\. The feature vectorϕ​\(s\)∈ℝ16\\phi\(s\)\\in\\mathbb\{R\}^\{16\}is binary, encoding the color of the nearest and second\-nearest object per cell \(8 colors, 2 distance bands\)\.

### IV\-BExpert Types

#### Two\-expert setup \(K=2K=2\)

- •w1w\_\{1\}: positive weight on red\-nearest, negative on blue\-nearest \(red\-lover\)
- •w2w\_\{2\}: positive weight on blue\-nearest, negative on red\-nearest \(blue\-lover\)

10 trajectories per type, 20 total, length 20 steps each\.

#### Three\-expert setup \(K=3K=3\)

A third reward type, black\-lover \(w3w\_\{3\}: positive weight on black\-nearest, negative on green\-nearest\), is added\. Black objects occupy 12% of states in the default grid instantiation, giving localized trajectories that are geometrically distinct from both red\-lover and blue\-lover\. We generate 10 trajectories per type, 30 total\.

We note that with the default random seed andn​\_​objects=5n\\\_\\text\{objects\}=5, blue and yellow objects are entirely absent from the grid\. This meansw2w\_\{2\}\(blue\-lover\) reduces to an avoid\-red signal rather than an independent seek\-blue signal, a degenerate property of the environment instantiation rather than the reward specification\. We report this explicitly as it bounds K=3 recovery quality\.

### IV\-CInference Hyperparameters

TABLE III:Inference Hyperparameters \(all main experiments\)ParameterValueα\\alpha\(DP concentration\)5\.0β\\beta\(Boltzmann temperature\)5\.0γ\\gamma\(discount factor\)0\.95σ\\sigma\(MH proposal scale\)0\.01Gibbs sweeps500Burn\-in sweeps100
### IV\-DEvaluation Metrics

Adjusted Rand Index \(ARI\):cluster recovery against ground\-truth assignments \(0 = random, 1 = perfect\)\.ℓ2\\ell\_\{2\}weight error:averageℓ2\\ell\_\{2\}distance between each ground\-truth weight and its nearest recovered weight afterℓ2\\ell\_\{2\}\-normalization\.Wall time:end\-to\-end inference runtime\.

### IV\-EBaseline

MaxEnt IRL\[[4](https://arxiv.org/html/2607.09886#bib.bib4)\], implemented with feature\-matching gradient ascent \(200 iterations, learning rate 0\.01\)\. Because MaxEnt fits a single reward, assigning all trajectories to one cluster givesARI=0\\mathrm\{ARI\}=0by construction — the expected outcome for any single\-component model on heterogeneous data\.

### IV\-FParallelization Study

Wall time measured atW∈\{1,2,4,8,16\}W\\in\\\{1,2,4,8,16\\\}on a fixed 100\-sweep run\. Speedup computed asT​\(1\)/T​\(W\)T\(1\)/T\(W\), same seed across all configurations\.

## VResults

### V\-ACluster and Reward Recovery \(K=2K=2\)

The serial Gibbs sampler recoversK=2K=2within approximately 250 sweeps and achievesARI=1\.000\\mathrm\{ARI\}=1\.000at the end of a 500\-sweep run, withℓ2\\ell\_\{2\}weight error of 1\.355 and wall time of 134\.3 seconds\. The MaxEnt baseline attainsARI=0\.000\\mathrm\{ARI\}=0\.000by construction\. Figure[1](https://arxiv.org/html/2607.09886#S5.F1)shows the cluster count trace; Figure[2](https://arxiv.org/html/2607.09886#S5.F2)shows recovered versus ground\-truth weight vectors, confirming that reward\-irrelevant features receive near\-zero weight\.

![Refer to caption](https://arxiv.org/html/2607.09886v1/fig3_convergence.png)Figure 1:Active cluster countKKover 500 Gibbs sweeps \(serial sampler,K=2K=2setup\)\. The chain explores during burn\-in \(sweeps 0–100, dotted vertical line\) then stabilizes at the ground\-truth valueK=2K=2\(dashed horizontal line\)\.![Refer to caption](https://arxiv.org/html/2607.09886v1/fig5_weight_recovery.png)Figure 2:Recovered versus ground\-truth cluster weight vectors across all 16 ObjectWorld features\. Reward\-relevant features \(red and blue color indicators\) are accurately recovered; the remaining 14 features receive near\-zero weight\.
### V\-BBaseline Comparison

Figure[3](https://arxiv.org/html/2607.09886#S5.F3)compares NP\-Bayes against MaxEnt\. NP\-Bayes recovers two correct reward functions; MaxEnt recovers one averaged reward\. Theℓ2\\ell\_\{2\}errors are similar in magnitude \(1\.355 vs\. 1\.194\) but measure different things: NP\-Bayes computes error per cluster against its matched ground\-truth type, while MaxEnt is compared to its nearest ground\-truth type from a single\-component fit\.

![Refer to caption](https://arxiv.org/html/2607.09886v1/fig4_baseline_comparison.png)Figure 3:NP\-Bayes IRL \(serial, 500 sweeps\) versus MaxEnt IRL baseline on ARI \(left\) andℓ2\\ell\_\{2\}weight error \(right\)\. The nonparametric model achieves perfect cluster recovery \(ARI=1\.000\\mathrm\{ARI\}=1\.000\); MaxEnt fits a single reward to all trajectories, producingARI=0\.000\\mathrm\{ARI\}=0\.000by construction\.ℓ2\\ell\_\{2\}errors are similar in magnitude but measure different things: NP\-Bayes reports per\-cluster error against matched ground\-truth types; MaxEnt reports error from a single\-component fit\.
### V\-CK=3K=3Recovery and Feature Discriminability

Extending to three expert types \(red\-lover, blue\-lover, black\-lover\), the sampler correctly identifiesK=3K=3in all runs but achievesARI=0\.58\\mathrm\{ARI\}=0\.58on assignment recovery\. Exhaustive search overα∈\{5\.0,8\.0,10\.0,15\.0\}\\alpha\\in\\\{5\.0,8\.0,10\.0,15\.0\\\}does not improve this ceiling, ruling out the concentration parameter as the limiting factor\.

Post\-hoc analysis of the default grid \(seed=42,nobjects=5n\_\{\\text\{objects\}\}=5\) reveals that blue and yellow objects are entirely absent, reducing the blue\-lover reward to a diffuse avoid\-red signal that occupies 78% of states and substantially overlaps with black\-lover trajectories \(12% coverage\)\. To test whether color diversity resolves this, we scanned seeds 1–20 and identified seed=19 as the only seed where red, blue, and black objects all appear simultaneously \(densities 17%, 54%, 17% respectively\)\. Running the K=3 sampler on seed=19 yieldsARI=0\.48\\mathrm\{ARI\}=0\.48, no improvement over seed=42\. The chain transiently findsK=3K=3between sweeps 150–250 but cannot sustain the partition: blue\-lover \(54% density\) and black\-lover \(17%\) have sufficient trajectory overlap that the sampler repeatedly merges them\.

These results isolate the bottleneck: color diversity in the grid is necessary but not sufficient for K=3 recovery\. The binding constraint is behavioral distinguishability — whether the three reward types produce sufficiently distinct trajectory distributions — which depends on the spatial arrangement of objects, not just their presence\. Only 1 of 20 random seeds produced a grid with all three required colors present, indicating that reliable K=3 evaluation on ObjectWorld requires either controlled object placement or a larger grid\. We report this as a finding about synthetic IRL benchmark design: feature discriminability cannot be assumed from reward specification alone and must be verified against the specific environment instantiation\.

### V\-DParallel Speedup

Table[IV](https://arxiv.org/html/2607.09886#S5.T4)summarizes the scaling study\. Wall time drops from 41\.15 seconds at one worker to 8\.58 seconds at 8 workers, a peak speedup of4\.79×4\.79\\times\. The super\-linear result at 2 workers \(2\.76×2\.76\\times\) is attributable to JAX JIT compilation amortization: with one worker, compilation overhead is paid serially; with two, both workers benefit from compiled code at lower per\-worker cost\. Beyond 8 workers, per\-worker workload becomes too small to amortize Ray’s actor coordination overhead and speedup degrades\.

TABLE IV:Parallel Speedup on Zaratan \(100 Sweeps,K=2K=2setup\)![Refer to caption](https://arxiv.org/html/2607.09886v1/fig1_speedup.png)Figure 4:Parallel speedup versus worker count on Zaratan \(100 sweeps\)\. Peak of4\.79×4\.79\\timesat 8 workers\. Dashed line: ideal linear scaling\. Degradation beyond 8 workers reflects per\-process overhead dominating the small per\-worker workload at this dataset size\.
### V\-EThroughput\-Accuracy Tradeoff

ARI degrades sharply with worker count \(0\.905 atW=2W=2, 0\.035 atW=16W=16\), whileℓ2\\ell\_\{2\}weight error remains approximately stable across all configurations \(Figure[5](https://arxiv.org/html/2607.09886#S5.F5)\)\. This dissociation is diagnostic: per\-worker reward inference is accurate, but cluster identity is lost during merging\.

The mechanism is the10−610^\{\-6\}\-tolerance equality check in the merge step\. Workers perform independent MH proposals, causing weight vectors to diverge by small amounts\. Drifts beyond tolerance cause the same cluster to be registered as distinct clusters across workers, inflatingKKto 16–18 atW=16W=16\. This is an engineering problem in the merge step, not a property of the sampler\. A centralized weight update — workers perform assignment steps in parallel, one process performs weight updates and broadcasts — would eliminate drift while preserving throughput gains\.

![Refer to caption](https://arxiv.org/html/2607.09886v1/fig2_weight_error.png)Figure 5:ℓ2\\ell\_\{2\}weight error versus worker count\. Error remains stable across all configurations, confirming that per\-worker reward inference is accurate\. Cluster identity failure \(risingKK\) drives ARI degradation, not reward quality\.
### V\-FDiscussion

The results address all three questions posed in the introduction\. The nonparametric model correctly identifiesKKand achieves perfect ARI on theK=2K=2benchmark while MaxEnt fails entirely\. OnK=3K=3, the model correctly infers the cluster count but assignment accuracy is bounded by environment feature discriminability rather than sampler behavior\. The parallel system achieves4\.79×4\.79\\timesspeedup; the throughput\-accuracy tradeoff is attributable to the merge step and is addressable without changes to the core inference algorithm\.

## VIConclusion

We presented a parallelized nonparametric Bayesian IRL system that jointly infers the number of latent expert reward types and the rewards themselves, without requiring the cluster count to be specified in advance\. On theK=2K=2ObjectWorld benchmark the system achieves perfect cluster recovery \(ARI=1\.000\\mathrm\{ARI\}=1\.000\) against a MaxEnt baseline that collapses to a single reward \(ARI=0\.000\\mathrm\{ARI\}=0\.000\)\. Extension toK=3K=3reveals that recovery quality is bounded by behavioral distinguishability in the environment, not by the prior or the sampler\. Data\-parallel Gibbs sampling via Ray achieves4\.79×4\.79\\timespeak speedup; the throughput\-accuracy tradeoff is isolated to the consensus merge step and admits a principled fix\. Ablation overα\\alphaandσ\\sigmashows thatα=5\.0\\alpha=5\.0,σ=0\.01\\sigma=0\.01robustly recoversK=2K=2and that both parameters require tuning relative to the problem scale\.

#### Future Work

Priority extensions are: \(i\) replacing the consensus merge with a centralized weight update to recover both throughput and cluster accuracy; \(ii\) a hierarchical prior overα\\alphato reduce sensitivity to this parameter across problem scales; \(iii\) split\-merge MCMC moves\[[15](https://arxiv.org/html/2607.09886#bib.bib15)\]to improve mixing in multi\-modal posteriors; \(iv\) evaluation on procedurally generated grids with guaranteed color diversity and on continuous state spaces with neural reward parameterizations\[[16](https://arxiv.org/html/2607.09886#bib.bib16)\]; and \(v\) application to real demonstration datasets such as driving trajectories or robotic manipulation logs\.

## References

- \[1\]A\. Y\. Ng and S\. Russell, “Algorithms for inverse reinforcement learning,” in*International Conference on Machine Learning*, 2000, pp\. 663–670\.
- \[2\]D\. Ramachandran and E\. Amir, “Bayesian inverse reinforcement learning,” in*International Joint Conference on Artificial Intelligence*, 2007, pp\. 2586–2591\.
- \[3\]P\. Abbeel and A\. Y\. Ng, “Apprenticeship learning via inverse reinforcement learning,” in*International Conference on Machine Learning*, 2004\.
- \[4\]B\. D\. Ziebart, A\. L\. Maas, J\. A\. Bagnell, and A\. K\. Dey, “Maximum entropy inverse reinforcement learning,” in*AAAI Conference on Artificial Intelligence*, 2008, pp\. 1433–1438\.
- \[5\]M\. Babes\-Vroman, V\. Marivate, K\. Subramanian, and M\. L\. Littman, “Apprenticeship learning about multiple intentions,” in*International Conference on Machine Learning*, 2011, pp\. 897–904\.
- \[6\]J\. Choi and K\.\-E\. Kim, “Nonparametric Bayesian inverse reinforcement learning for multiple reward functions,” in*Advances in Neural Information Processing Systems*, 2012\.
- \[7\]T\. S\. Ferguson, “A Bayesian analysis of some nonparametric problems,”*The Annals of Statistics*, vol\. 1, no\. 2, pp\. 209–230, 1973\.
- \[8\]D\. J\. Aldous, “Exchangeability and related topics,” in*École d’été de probabilités de Saint\-Flour XIII*, Springer, 1985, pp\. 1–198\.
- \[9\]R\. M\. Neal, “Markov chain sampling methods for Dirichlet process mixture models,”*Journal of Computational and Graphical Statistics*, vol\. 9, no\. 2, pp\. 249–265, 2000\.
- \[10\]I\. J\. Lovett*et al\.*, “Parallel Markov chain Monte Carlo for Bayesian hierarchical models,”*Journal of Computational and Graphical Statistics*, vol\. 31, no\. 2, pp\. 455–466, 2022\.
- \[11\]J\. Bradbury*et al\.*, “JAX: composable transformations of Python\+NumPy programs,” 2018\. \[Online\]\. Available:http://github\.com/google/jax
- \[12\]P\. Moritz*et al\.*, “Ray: a distributed framework for emerging AI applications,” in*USENIX Symposium on Operating Systems Design and Implementation*, 2018, pp\. 561–577\.
- \[13\]S\. Levine, Z\. Popovic, and V\. Koltun, “Feature construction for inverse reinforcement learning,” in*Advances in Neural Information Processing Systems*, 2010\.
- \[14\]S\. Levine, Z\. Popovic, and V\. Koltun, “Nonlinear inverse reinforcement learning with Gaussian processes,” in*Advances in Neural Information Processing Systems*, 2011\.
- \[15\]S\. Jain and R\. M\. Neal, “A split\-merge Markov chain Monte Carlo procedure for the Dirichlet process mixture model,”*Journal of Computational and Graphical Statistics*, vol\. 13, no\. 1, pp\. 158–182, 2004\.
- \[16\]M\. Wulfmeier, P\. Ondruska, and I\. Posner, “Maximum entropy deep inverse reinforcement learning,”*arXiv preprint arXiv:1507\.04888*, 2015\.