PaperlayerView exact revision on arXiv

Source-linked AI summary

Q-based Variational Inverse Reinforcement Learning

Ondrej Bajgar, Peter Tisnikar, Alessandro Abate, Konstantinos Gatsis, Maike Osborne

arXiv:2608.16888v1cs.LG

TL;DR

Bayesian IRL needs to scale beyond small environments while preserving posterior uncertainty over reward functions. QVIRL learns a variational posterior over optimal Q-values and achieves competitive apprenticeship learning, including from raw pixels, while enabling scalable active learning.

  • Problem

    Bayesian IRL methods often struggle to scale beyond small or low-dimensional environments while preserving full posterior uncertainty over rewards.

  • Method

    QVIRL uses variational inference to learn a correlated posterior over optimal Q-values, then transforms it through the inverse Bellman operator to obtain a reward posterior.

  • Results

    QVIRL recovers well-calibrated posteriors in tabular environments and achieves competitive apprenticeship learning across complex settings, including training from raw pixel observations.

  • Takeaways & Limitations

    QVIRL provides a scalable Bayesian IRL backbone for active learning by explicitly modelling uncertainty in optimal Q-values.

  • Takeaways & Limitations

    In gridworlds, converting reward-posterior samples to Q-values requires costly value iteration that does not transfer to unknown-dynamics or larger environments.

Abstract

from arXiv · show

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.

1 Introduction

The introduction frames Bayesian IRL as a way to infer human preferences from demonstrations while retaining uncertainty over reward functions. It motivates a variational method designed to combine scalability with full posterior uncertainty estimation for robust policies and active learning.

  • Motivation: Bayesian IRL operationalizes beneficial-AI principles by inferring reward functions from human behaviour while representing uncertainty about preferences.It begins with a prior over reward functions and updates beliefs using expert demonstrations.
  • Motivation: Despite successful non-Bayesian applications in robotics, navigation, and autonomous driving, scalable Bayesian IRL remains limited by state- and action-space dimensionality.Prior Bayesian methods generally target finite spaces or continuous spaces with only a handful of dimensions.
  • Motivation: Bayesian IRL recovers a posterior over reward functions, supporting policies robust to uncertainty and active learning that gathers informative data.Limited demonstration trajectories create additional uncertainty that can be reduced where it matters.
  • Contribution: The paper contributes a variational-inference method that preserves scalability across observation dimensionality and demonstration-data volume while estimating full posterior uncertainty.The authors demonstrate these properties in higher-dimensional settings and through active learning; variational inference is presented as more computationally efficient.

2 Problem Formulation

Bayesian IRL infers a posterior over reward functions from expert demonstrations in an MDP, using a reward prior and an expert action model conditioned on optimal Q-values. The formulation supports stochastic underlying rewards through their expectations and motivates apprenticeship learning under static or actively selected demonstrations.

  • Problem formulation: Bayesian IRL seeks a posterior distribution over reward functions from demonstrations of an expert acting in an MDP with states, actions, transitions, rewards, and discounting.Demonstrations pair represented states ϕ(s) with expert actions, while the feature function maps states into a feature space.
  • Expert model: The formulation assumes expert actions are conditionally independent given optimal Q-values and models their likelihood with standard Boltzmann rationality.The rationality coefficient β controls the expert action model, while Q∗(s, a) is the expected discounted return under the optimal policy.
  • Reward assumptions: The underlying reward may be stochastic, but because demonstrations depend on expected optimal Q-values, the learned reward represents either deterministic reward or a stochastic reward expectation.The expert model depends only on the expectation of the reward, so demonstrations provide information about that expectation.
  • Posterior inference: Combining the demonstration likelihood with a reward prior yields p(r|D) = p(D|r)p(r)/p(D), while transition and initial-state factors cancel.Except for special cases, the resulting posterior cannot be computed analytically.
  • Apprenticeship learning: The method targets apprenticeship learning, evaluating apprentice policies by expected return or risk-averse criteria such as CVaR with static data and active demonstration selection.Active learning selects initial states for subsequent expert demonstrations to reduce posterior uncertainty or improve apprentice performance.

3 Related Work

Related work spans classical IRL and inverse optimal control, Bayesian methods based on MCMC, and newer variational approaches. Bayesian IRL methods trade posterior fidelity or repeated forward-RL computation against scalability, motivating approaches that jointly learn reward and Q-function representations.

  • Classical IRL and inverse optimal control: IRL originated with Russell (1998), while closely related inverse optimal control was formulated earlier; the two communities have used largely separate methods.The section directs readers to broader reviews and a comparison of the two formulations.
  • Bayesian IRL: Ramachandran and Amir (2007) introduced Bayesian IRL with MCMC, which samples the true reward posterior but scales poorly and repeatedly solves the forward RL problem.Michini and How (2012) improved efficiency by focusing on relevant state-action regions, while Mandyam et al. (2023) sought to reduce repeated forward-RL solves.
  • Variational Bayesian IRL: Chan and van der Schaar (2021) applied variational inference to Bayesian IRL by jointly learning a reward encoder and Q-network, linked through a soft Bellman-equation constraint.The reward encoder produces a mean and variance for rewards in each state, while the Q-network estimates the optimal Q-function and generates the apprentice policy.

4 Method

QVIRL learns a variational posterior over optimal Q-values from demonstrations and transforms it through the inverse optimal Bellman operator into a posterior over rewards. Its architecture models correlated Q-value uncertainty and uses Gaussian approximations to make reward inference and training tractable across discrete and continuous environments.

  • Method overview: QVIRL takes demonstrations and a reward prior as input, outputs a variational posterior over optimal Q-functions, and transforms it into a reward posterior through the inverse optimal Bellman operator.The method centers on modelling optimal Q-values because converting Q-values to rewards is computationally simpler than the reverse direction.
  • Training and dynamics: The method handles discrete and continuous environments by evaluating successor-state expectations exactly in tabular settings or approximating them with samples from expert or auxiliary trajectories.Encoder parameters are optimized with stochastic gradient descent using demonstration likelihood and a KL term between the implicit reward posterior and its prior.
  • Variational Q posterior: The variational family uses a multivariate Gaussian for discrete spaces or a Gaussian process for continuous spaces, with an encoder producing Q-value means, standard deviations, and embeddings.The embeddings feed a kernel that computes correlations among state-action pairs.
  • Variational Q posterior: QVIRL explicitly models covariance between optimal Q-values because demonstration likelihoods are invariant to constant Q-value shifts and Bellman dependencies correlate values along trajectories.The covariance is constructed from individual standard deviations and kernel-based correlations between latent embeddings.
  • Reward posterior: The reward posterior is obtained from the inverse Bellman equation, but the max over jointly Gaussian successor Q-values lacks a closed-form distribution and is therefore approximated.Monte Carlo sampling is feasible at inference time but can be computationally costly during training.
  • Reward posterior: QVIRL applies Clark’s approximation, building the maximum over actions from pairwise Gaussian maxima to produce a closed-form Gaussian reward posterior and more efficient KL computation.The main-text results use the Clark approximation rather than the simpler max-mean alternative.

5 Experiments

Experiments evaluate QVIRL’s posterior approximation, apprenticeship-learning performance, scalability across environments, and active-learning utility. Across these settings, QVIRL is assessed against variational, imitation-learning, and MCMC-based baselines.

  • Experimental setup: QVIRL is evaluated in randomized gridworlds, Lunar Lander, Highway, and Atari games, including raw-pixel observations for Atari.Gridworlds provide a controlled tabular setting, while Lunar Lander and Highway test larger continuous-state benchmarks; Atari uses a convolutional encoder trained end-to-end.
  • Posterior approximation: QVIRL produces a higher-fidelity posterior than AVRIL on 100 random 8×8 gridworlds, with posterior log density and credible-interval calibration matching ValueWalk.The evaluation uses five expert demonstrations per gridworld and 5000 ValueWalk samples from the true reward posterior.
  • Apprenticeship learning: QVIRL’s gridworld policies perform only slightly worse than the ValueWalk oracle under both true-reward and posterior-return evaluations.CVaR optimization improves worst-case posterior performance but lowers mean return for all methods, while AVRIL remains weak across settings.
  • Apprenticeship learning: In Lunar Lander and Highway, QVIRL reaches expert-like performance with fewer demonstrations and lower run-to-run variability than the baselines.All methods eventually achieve expert performance with 15 or 3 trajectories, respectively, but with larger variance than QVIRL.
  • Atari apprenticeship learning: On Pong and Space Invaders, QVIRL is competitive with IQ-Learn while also providing a posterior over rewards from raw-pixel observations.QVIRL uses a NatureCNN encoder and trains it jointly with the variational Q-value posterior.
  • Active learning: QVIRL-based acquisition nearly matches ValueWalk in gridworlds and substantially outperforms random querying in Lunar Lander when paired with Reward EIG or ActiveVaR.AVRIL performs significantly worse and stalls after 10 queries in gridworlds.

6 Discussion

QVIRL uses variational inference to approximate Bayesian posteriors over optimal Q-values and rewards, recovering well-calibrated posteriors in tabular environments. It also combines efficient pixel-based apprenticeship learning with scalable uncertainty-aware active learning, while reward-to-Q sampling remains costly in larger or unknown-dynamics environments.

  • Bayesian posterior approximation: QVIRL recovers well-calibrated posteriors over optimal Q-values and rewards in tabular environments, unlike AVRIL’s point estimate of optimal Q-values.QVIRL approximates the Bayesian posterior using variational inference, while AVRIL fails to approximate the true reward posterior.
  • Computational limitation: Sampling rewards from the reward posterior and converting them to Q-values via value iteration is very costly and does not scale to larger environments or unknown dynamics.This workaround was used in gridworld experiments but does not translate to environments without known dynamics or to larger environments.
  • Apprenticeship learning: QVIRL achieves competitive apprenticeship-learning performance with an efficient, gradient-based algorithm that supports training from raw pixel observations.The raw-pixel capability distinguishes QVIRL from most Bayesian IRL methods.
  • Active learning: QVIRL enables active learning by modelling the optimal Q-value posterior for scalable acquisition functions that track uncertainty in Q-value space.This removes a key scalability bottleneck of previous MCMC-based active IRL methods and opens active learning in much larger environments.

7 Conclusion

QVIRL is presented as a scalable Bayesian inverse reinforcement learning method that preserves posterior uncertainty over optimal Q-values and rewards. It models their joint posterior variationally beyond small tabular environments and performs competitively against Bayesian IRL and state-of-the-art imitation learning methods.

  • 7 Conclusion: QVIRL preserves posterior uncertainty quantification over both optimal Q-values and rewards while scaling beyond small tabular environments.It explicitly models a joint posterior over optimal Q-values.
  • 7 Conclusion: QVIRL performs competitively against Bayesian IRL baselines and state-of-the-art imitation learning methods.

Supplementary Materials The following content was not necessarily subject to peer review. … C.2 3x3 Gridworld and a Comparison to AVRIL

The supplementary materials detail QVIRL’s transition-expectation, state-value, reward-posterior, KL, and approximation procedures, then compare QVIRL with AVRIL in a 3×3 gridworld. They show that QVIRL closely approximates the true posterior, while AVRIL is highly sensitive to its consistency-weight hyperparameter.

  • A.1 Evaluating Expectations over Next States: Known discrete dynamics permit exact successor-state expectations, while continuous or unavailable models require Monte Carlo or empirical-transition approximations.Gridworld experiments use fully known dynamics; empirical single-sample estimates become exact in deterministic environments and improve with repeated nearby transitions.
  • A.2 Clark Approximation to the State-Value Distribution: QVIRL approximates the maximum of correlated Gaussian Q-values with Clark’s moment-matching recursion, while max-mean is a simpler alternative that performs fairly well.Clark propagates moments and covariances; the max-mean surrogate is recommended as a starting point before extending to higher-fidelity Clark approximations.
  • A.2 Clark Approximation to the State-Value Distribution: Under Clark approximations, rewards are affine functions of jointly Gaussian Q-values and successor values, yielding Gaussian marginal and joint reward posteriors.The construction preserves positive semi-definiteness by using a bilinear covariance form and lifting diagonal entries to exact Clark variances.
  • A.2 Clark Approximation to the State-Value Distribution: The max-mean covariance surrogate underestimates state-value mean and variance because it omits option-value and action-uncertainty contributions.Its mean is a convex combination of action means and lies below even maxj µQ,j, while its variance ignores uncertainty about which action attains the maximum.
  • A.3 KL and the loss: The KL term matches the variational reward posterior to a Gaussian or Gaussian-process prior across state-action pairs, but continuous spaces introduce the hyperparameter λ.Unlike the tabular case’s natural Bayes-derived weighting, auxiliary batches require manually balancing likelihood and KL terms; dense inducing points are an alternative with scalability challenges.
  • B.3 Gaussianity: The Gaussianity assumption remains reasonable for the illustrated Gaussian-prior case, while state-only rewards can be modeled through strongly correlated action rewards or a dedicated state-value formulation.The state-only formulation derives rewards and Q-values from state values and uses Clark or max-mean Gaussian approximations.
  • B.1 Max-Mean Approximation: The max-mean approximation introduces little error in most 3×3 gridworld states but deviates where successor-state actions are not clearly separated.The discrepancy occurs in two bottom-left states, where posterior uncertainty makes the highest-mean action a poorer proxy for the state-value distribution.
  • C.2 3x3 Gridworld and a Comparison to AVRIL: QVIRL’s gridworld posterior approximates the true posterior well, whereas AVRIL remains near zero and is extremely sensitive to λ.With λ = 0.5, std<0.001 for 8 states out of 9; λ = 0.1 essentially reverts the posterior to the prior, while λ = 0.2 was used in the reported experiment.

C.3 Limitations of AVRIL … D.2 Expert Demonstrations

The paper contrasts QVIRL with AVRIL’s limitations in uncertainty quantification, posterior fidelity, hyperparameter sensitivity, and batch-only inference, then details evaluation environments and expert-demonstration procedures. Experiments span gridworlds, Lunar Lander, Highway Env, and raw-pixel Atari games using both static and active-learning data collection.

  • C.3 Limitations of AVRIL: QVIRL provides uncertainty estimates over optimal Q-values and policies, enabling natural extraction of risk-averse policies at comparable computational effort.AVRIL does not provide these uncertainty estimates, limiting straightforward risk-averse policy extraction.
  • C.3 Limitations of AVRIL: QVIRL’s reward posterior fits the posterior recovered by MCMC methods more faithfully than AVRIL’s, even on a simple gridworld.The passage states that AVRIL’s posterior does not track the true reward posterior well.
  • C.3 Limitations of AVRIL: QVIRL avoids AVRIL’s sensitivity to the reward–Q consistency hyperparameter λ, whose narrow viable range can cause posterior collapse or reversion to the prior.The paper provides no guidance for choosing AVRIL’s λ, whereas QVIRL has no such hyperparameter and is robust to other challenging design choices.
  • C.3 Limitations of AVRIL: AVRIL’s strictly batch setting uses only expert demonstrations to estimate transitions, whereas QVIRL leverages environment dynamics to infer rewards away from expert trajectories.The authors argue that batch-only inference gives up a key IRL advantage over behavioural cloning: improved generalization through environment dynamics.
  • D.1 Environment Descriptions: The evaluation covers 100 randomized 8x8 gridworlds with five actions, 0.1 random-action probability, state-only rewards, stochastic terminals, and uniformly distributed initial states.Gridworld rewards are independently drawn from N(−1, 3), and the top 10% highest-reward states are additionally marked terminal.
  • D.1 Environment Descriptions: Lunar Lander uses an 8-dimensional continuous state and gives +100 for successful landing, −100 for crashing, with +200 defining a solved episode and approximately 270 indicating expert performance.The task models landing between two flags and penalizes excessive tilt and engine firing.
  • D.1 Environment Descriptions: Highway Env models multi-lane autonomous driving with kinematic observations and discrete actions, while Atari experiments test scalability to raw 84 × 84 grayscale pixel observations in Pong and Space Invaders.Atari preprocessing skips four frames between decisions and max-pools over the last two skipped frames.
  • D.2 Expert Demonstrations: Expert demonstrations use Boltzmann-rational policies: gridworlds provide 5 trajectories with β = 2, Lunar Lander and Highway Env use β = 3 and β = 5 with 1000 trajectories each, and Atari uses 20 training plus 5 held-out trajectories per game.Gridworld active learning begins with one trajectory and adds one queried trajectory per active step; Lunar Lander and Highway Env split demonstrations into 800 training, 100 evaluation, and 100 test trajectories.

D.3 QVIRL Training … E.1 Sensitivity of QVIRL to Rationality Coefficient β

The paper details QVIRL’s optimization, environment-specific priors, and additional experiments probing its behavior. Its sensitivity analysis shows that the preferred rationality coefficient depends on both misspecification and demonstration count.

  • D.3 QVIRL Training: QVIRL is trained with stochastic gradient ascent, PyTorch automatic differentiation, Adam at learning rate 0.001, batch size 64, and about 20000 iterations.The discount rate is γ = 0.99 for Lunar Lander, 0.95 for Highway, and 0.9 for gridworlds; the demonstrations’ Boltzmann coefficient is reused.
  • D.3 QVIRL Training: For Atari, a three-convolutional-layer NatureCNN encoder feeds QVIRL and is trained end-to-end with its variational parameters.The encoder is neither pretrained nor frozen.
  • D.4 Prior Distribution: The priors vary by environment: gridworlds use independent Gaussian state priors, while Lunar Lander and Highway use Gaussian-process priors with constant mean −1 and ARD RBF kernels.For Lunar Lander and Highway, the kernel is evaluated on standardized state-action features.
  • D.4 Prior Distribution: Atari uses a zero-mean Gaussian-process prior combining action and image kernels over grayscale 84 × 84 frame stacks.The mixture includes patch, motion, spatial-pyramid, and pixel kernels, with game-specific weights and pyramid levels.
  • E Additional Experiments: Additional experiments expose QVIRL’s properties and investigate shortcomings of AVRIL.These experiments are presented as supplementary analyses of the methods’ behavior.
  • E.1 Sensitivity of QVIRL to Rationality Coefficient β: The rationality coefficient β controls softmax sharpness, and misspecifying it can bias or remove the reward-learning signal.QVIRL requires relatively few hyperparameters to tune compared with many current Bayesian IRL methods.
  • E.1 Sensitivity of QVIRL to Rationality Coefficient β: With 1, 7, and 15 demonstrations across 30 seeds, QVIRL was evaluated using four β values spanning random to optimal-expert behavior.The metric was the discrepancy in held-out greedy expert action likelihood.
  • E.1 Sensitivity of QVIRL to Rationality Coefficient β: In low-data regimes, QVIRL performs better when β assumes a less deterministic expert than the true value, with best performance from deliberately underestimating rationality.As demonstration count increases, β’s regularization effect becomes more pronounced; future work could infer β jointly with optimal Q-values.

E.2 Illustrative Gridworld Experiment with a Single Unknown State · F Code, Data, and Compute

QVIRL recovers a reward posterior in a gridworld with one unknown state and uses its mean reward to choose a route through that state. The posterior also supports risk-averse behavior and closely matches an MCMC reference posterior.

  • E.2 Illustrative Gridworld Experiment with a Single Unknown State: The experiment assumes known rewards except for one tile, whose prior has mean -1 and standard deviation 10.Known rewards use a tight normal prior with standard deviation 0.1 centered at their true values.
  • E.2 Illustrative Gridworld Experiment with a Single Unknown State: A single expert demonstration cannot distinguish whether the apprentice should cross the unknown tile or go around it.Both routes are consistent with an optimal expert demonstration.
  • E.2 Illustrative Gridworld Experiment with a Single Unknown State: Using value iteration on the posterior mean reward, the apprentice goes through the unknown tile.The experiment evaluates whether reaching the goal faster justifies crossing the uncertain tile.
  • E.2 Illustrative Gridworld Experiment with a Single Unknown State: 8.0 is the Q-value for going up, compared with 5.2 for going left.These values determine the apprentice’s route under the posterior mean reward.
  • E.2 Illustrative Gridworld Experiment with a Single Unknown State: Posterior uncertainty enables risk-averse behavior in the gridworld.The authors identify this as a consequence of explicitly modelling uncertainty in the recovered reward posterior.
  • E.2 Illustrative Gridworld Experiment with a Single Unknown State: The reward posterior inferred from QVIRL’s Q-value posterior fits well with the MCMC reference reward posterior.The reference posterior is recovered using ValueWalk, which samples from the true posterior with MCMC.

F.1 Code and Data · F.2 Computing Resources Used

The experiments used publicly available software and data, with code slated for release on GitHub. Training used single CPU cores for non-ATARI tasks and an NVIDIA V100 GPU for ATARI, with runtimes varying by environment and step count.

  • F.1 Code and Data: All experiments used publicly available libraries and data.
  • F.1 Code and Data: The main software components were Python 3.13, Farama Gymnasium 1.2.3, Highway Env 1.10.2, and PyTorch 2.9.1.
  • F.1 Code and Data: The listed software used PSF, MIT, or BSD-3 licensing.Python 3.13 used the GPL-compatible PSF license agreement; Farama Gymnasium 1.2.3 and Highway Env 1.10.2 used MIT; PyTorch 2.9.1 used BSD-3.
  • F.1 Code and Data: The code will be made available on GitHub at the qvirl repository.The specified repository is https://github.com/bayesian-reward-learning/qvirl.
  • F.2 Computing Resources Used: Except ATARI, each training ran on a single CPU core of either an 8-core AMD Ryzen 7 PRO 7840U or Amazon Gravitron 8.
  • F.2 Computing Resources Used: Non-ATARI training took between 2 and 15 minutes.
  • F.2 Computing Resources Used: Each ATARI training run used an NVIDIA V100 GPU.
  • F.2 Computing Resources Used: ATARI training took about 90 minutes for the 200k steps on Pong and around 7 hours for the 1M steps of Space Invaders.
Loading 2608.16888v1…