The development of safe and beneficial AI requires that systems can learn and act in accordance with human preferences. However, explicitly specifying these preferences by hand is often infeasible. Inverse reinforcement learning (IRL) addresses this challenge by inferring preferences, represented as reward functions, from expert behaviour. We introduce Q-based Variational IRL (QVIRL), a novel Bayesian IRL method that recovers a posterior distribution over rewards from expert demonstrations via primarily learning a variational distribution over optimal Q-values. Unlike previous approaches, QVIRL combines scalability with uncertainty quantification, important for safety-critical applications as well as active learning. We demonstrate QVIRL's strong performance in apprenticeship learning across various tasks, including gridworlds, Lunar Lander, the Highway Environment, and two ATARI games both with static expert data and with active learning. It is the first method for Bayesian IRL that demonstrates training from raw pixel observations.
Figures & tables
Figure 1: Illustration of the QVIRL architecture. A state-action pair (s,a) and a successor pair (s′,a′) are fed into an encoder network that outputs the mean μQ∗ and the standard deviation σQ∗ of the associated optimal Q-value Q∗(s,a) , as well as a latent embedding e . The embeddings of a set of state-action pairs can be passed to a kernel to produce their correlation matrix and, combined with individual variances, their covariance Cov(Q∗(s,a),Q∗(s′,a′)) denoted by ΣQ∗Q′∗ in the diagram. Together, they define the joint variational distribution of the optimal Q-values of a given set of state-action pairs. This distribution can then be used to deduce the distribution over the corresponding reward through the Bellman equation.
Figure 2: Illustration of posteriors recovered by ValueWalk (MCMC samples used as ground truth), AVRIL (for three values of constraint weight λ) , and QVIRL in 3 randomly selected states from our randomized 8×8 gridworlds.
Method
(1) W1 (VW) ↓
(2) logp(rtrue∣D)↑
(3) rtrue∈90% CI
(4) πexp∈90% CI
(5) TV (VW) ↓
ValueWalk
0.000
−2.432±0.013
0.888±0.004
0.910±0.003
0.000
QVIRL
0.486±0.011
−2.437±0.009
0.901±0.004
0.913±0.003
0.207±0.003
AVRIL
1.704±0.007
−7.568±0.284
0.521±0.006
0.727±0.005
0.250±0.004
Table 1: Posterior quality on 100 random 8 × 8 gridworlds (mean ± stderr across 100 seeds). (i) log-density of the true reward under the posterior; (ii) fraction of states with r(s) within the posterior 90% credible interval; (iii) fraction of expert action probabilities π(s∣a) within the posterior predictive 90% credible interval; and (iv) total variation distance between each method’s posterior-predictive policy and that of the ValueWalk oracle. ↑ / ↓ indicates higher/lower is better.
Policy π
Jrtrue(π)
Er∣D[Jr(π)]
CVaRr∣D0.1[Jr(π)]
ValueWalk (mean)
−75.4±1.8
112.1±2.8
−59.9±2.8
ValueWalk (CVaR)
−75.8±1.3
107.7±3.0
−62.1±2.3
QVIRL (mean)
−70.8±2.1
106.4±3.2
−49.0±3.9
QVIRL (CVaR)
−66.9±1.9
089.6±3.3
−49.7±2.8
AVRIL (mean)
−22.2±4.8
005.4±2.7
−86.3±3.4
AVRIL (CVaR)
−36.7±1.1
027.1±1.4
−14.0±1.3
Table 2: Apprenticeship learning on 100 random 8 × 8 gridworlds (mean ± stderr). Returns are normalised: 0 = random-policy return, 100 = optimal-policy return under the true reward. Jrtrue(π) : true-reward return of the apprentice policy. Er∣D[Jr(π)] : expected return under the ValueWalk oracle posterior. CVaRr∣D0.1[Jr(π)] : conditional value-at-risk (10% level) of the return under the posterior, i.e. the expected return in the worst 10% cases under the ValueWalk reward posterior; the CVaR policy directly optimizes this objective under each method’s own posterior.
Figure 3: Mean episode return in apprenticeship learning with different numbers of expert demonstration trajectories in the Lunar Lander (left) and Highway (right) environments. Averaged over 10 runs, with standard errors of the mean. All methods are evaluated on discrete number of expert trajectories, but we space the points horizontally to improve plot readability.
Figure 4: Mean log-likelihood of held-out expert actions with different numbers of expert trajectories in the Lunar Lander (left) and Highway (right) environments. Averaged over 10 runs, with standard errors of the mean. All methods are evaluated on discrete numbers of expert trajectories, but we space the points horizontally to improve plot readability.
Pong
Space Inv.
Method
Return
LL
Return
LL
BC
19.1 ± 0.8
-14.81 ± 0.49
597 ± 48
-18.67 ± 0.05
IQ-Learn
20.0 ± 0.2
-12.43 ± 2.27
740 ± 31
-8.99 ± 1.79
AVRIL
-5.2 ± 9.7
-8.79 ± 3.80
658 ± 34
-18.62 ± 0.07
QVIRL
19.7 ± 0.6
-2.09 ± 0.12
701 ± 42
-5.19 ± 1.04
Table 3: Apprenticeship learning on Atari: mean episode return and log likelihood (LL) of expert actions on held-out demonstrations (mean ± standard error across seeds). QVIRL LL is computed under the probit likelihood; BC, IQ-Learn, and AVRIL LL are computed under the point likelihood.
Figure 5: Active learning results on randomized gridworlds (left) and Lunar Lander (right).
Figure 6: Effect of the approximation in Eq. 3 on the reward posterior. The plot shows the MCMC posterior distribution over rewards for each state, as well as samples deduced from the MCMC posterior over rewards deduced using the approximation from the posterior over state-values. Dashed lines show Gaussians fitted to the two MCMC samples. The results coincide except in the bottom left and middle squares.
Figure 7: Comparison of the likelihood approximation equation 6 against a Monte Carlo evaluation using 10,000 Monte Carlo samples. Left: probability of a1 as a function of the mean of Q1 for 3 levels of standard deviation of Q1 . Correlation is set to 0 (the plot looks very similar for correlated values). Right: Probability of a1 as a function of the correlation between Q1 and Q2 for three levels of the mean. Standard deviation is fixed to 1. Note the different scales on y-axes of the two plots.
Figure 8: Top left: Ground truth rewards for a 3x3 gridworld. The top left state is the initial state. Top right state is terminal. Top right: State-action counts in the demonstration data. Bottom left: pdf of the marginal posterior over state values recovered by QVIRL, a histogram of 2000 samples from the true posterior produced by ValueWalk, and a horizontal line indicating the point estimate recovered by AVRIL. Bottom right: pdfs of the variational posterior over reward recovered by QVIRL and AVRIL and the corresponding histogram.
Figure 12
Figure 10: The mean true greedy action log probability difference between the demonstrator and learner for 1, 7, and 15 demonstrations, respectively.
Figure 11: The ground truth rewards, demonstrations, and the posterior over the unknown right-middle reward recovered by QVIRL and Valuewalk in the illustrative gridworld.
Reinforcement learning studies how agents learn sequential decisions from rewards and interaction, balancing exploration, performance, safety, and adaptation.
Inverse reinforcement learning (IRL) is typically formulated as maximizing entropy subject to matching the distribution of expert trajectories. Classical (dual-ascent) IRL guarantees monotonic performance improvement but requires fully solving an RL problem each iteration to compute dual gradients. More recent adversarial methods avoid this cost at the expense of stability and monotonic dual improvement, by directly optimizing the primal problem and using a discriminator to provide rewards. In this work, we bridge the gap between these approaches by enabling monotonic improvement of the reward function and policy without having to fully solve an RL problem at every iteration. Our key theoretical insight is that a trust-region-optimal policy for a reward function update can be globally optimal for a smaller update in the same direction. This smaller update allows us to explicitly optimize the dual objective while only relying on a local search around the current policy. In doing so, our approach avoids the training instabilities of adversarial methods, offers monotonic performance improvement, and learns a reward function in the traditional sense of IRL--one that can be globally optimized to match expert demonstrations. Our proposed algorithm, Trust Region Inverse Reinforcement Learning (TRIRL), outperforms state-of-the-art imitation learning methods across multiple challenging tasks by a factor of 2.4x in terms of aggregate inter-quartile mean, while recovering reward functions that generalize to system dynamics shifts.
Anish Diwan, Davide Tateo, Christopher E. Mower +3
Technical University of Darmstadt · Robotics Institute Germany (RIG) · Lund University +3
Inverse Reinforcement Learning (IRL) aims to recover a reward function that explains expert demonstrations. Existing IRL methods typically rely on a bi-level optimization procedure that alternates between reward learning and policy optimization, leading to substantial computational burden and training instability. In this work, we introduce a different route that eliminates policy optimization entirely by leveraging diffusion policies. Our key insight is that a diffusion policy encodes the action-gradient structure of the optimal soft Q function, enabling reward learning to be cast as a sequence of value recovery problems, thereby allowing us to bypass reward-policy loops inherent in prior IRL methods. Specifically, our method proceeds in three stages: (I) recovering the optimal soft Q function via action-gradient matching and estimating the corresponding soft value function (LogSumExp of Q values) in a way inspired by Gumbel regression; (II) calibrating these soft values by inferring a state-dependent offset; (III) extracting the reward by enforcing Bellman consistency. This leads to Loop-Free Inverse Reinforcement Learning (LFIRL), a fully offline algorithm that operates in a simple, loop-free, and sequential manner. LFIRL is simple to implement and significantly improves training efficiency while maintaining strong reward recovery performance. Empirically, across Maze, Franka Kitchen, Adroit Hand Pen, and Push-T benchmarks, LFIRL achieves 2-3x speedup over the fastest baselines, while matching or surpassing state-of-the-art methods in reward recovery quality.
Yang Chen, Yitan Zhang, Michael Witbrock +1
Shanghai Artificial Intelligence Laboratory · University of Auckland
Reward learning typically relies on a single feedback type or combines multiple feedback types using manually weighted loss terms. Currently, it remains unclear how to jointly learn reward functions from heterogeneous feedback types such as demonstrations, comparisons, ratings, and stops that provide qualitatively different signals. We address this challenge by formulating reward learning from multiple feedback types as Bayesian inference over a shared latent reward function, where each feedback type contributes information through an explicit likelihood. We introduce a scalable amortized variational inference approach that learns a shared reward encoder and feedback-specific likelihood decoders and is trained by optimizing a single evidence lower bound. Our approach avoids reducing feedback to a common intermediate representation and eliminates the need for manual loss balancing. Across discrete and continuous-control benchmarks, we show that jointly inferred reward posteriors outperform single-type baselines, exploit complementary information across feedback types, and yield policies that are more robust to environment perturbations. The inferred reward uncertainty further provides interpretable signals for analyzing model confidence and consistency across feedback types.
Raphaël Baur, Yannick Metz, Maria Gkoulta +3
ETH AI Center, ETH Zurich, Zurich, Switzerland · Department of Computer Science, ETH Zurich, Zurich, Switzerland · Department of Informatics, University of Zurich, Zurich, Switzerland