PaperlayerView exact revision on arXiv

Source-linked AI summary

Scalable Bayesian Inverse Reinforcement Learning

Alex J. Chan, Mihaela van der Schaar

arXiv:2102.06483v2cs.LG

TL;DR

Bayesian inverse reinforcement learning struggles with complex, unknown environments because existing approaches rely on restrictive rewards, costly inner-loop MDP solves, or interaction. AVRIL jointly learns an approximate reward posterior and imitator policy entirely offline, achieving strong medical-data results and demonstrator-level or competitive control performance.

  • Problem

    Existing Bayesian inverse reinforcement learning methods are impractical in complex environments because they use restrictive reward functions and costly inner-loop MDP solves, while high-stakes settings may prohibit policy interaction.

  • Method

    AVRIL jointly learns a variational posterior over flexible rewards and an imitator policy entirely offline, with reward regularisation enforcing temporal consistency.

  • Results

    AVRIL performs strongly across real medical and control tasks, attaining the highest accuracy and precision in both healthcare tasks and reaching demonstrator-level control performance, though not always as quickly as dedicated offline imitation-learning methods.

  • Takeaways & Limitations

    The method enables Bayesian reward inference and useful demonstrator-preference insight in environments beyond current methods while supporting offline imitation learning.

  • Takeaways & Limitations

    AVRIL remains an approximation, and differentiating through inner-loop Q-value solving remains impractical as environment complexity increases.

Abstract

from arXiv · show

Bayesian inference over the reward presents an ideal solution to the ill-posed nature of the inverse reinforcement learning problem. Unfortunately current methods generally do not scale well beyond the small tabular setting due to the need for an inner-loop MDP solver, and even non-Bayesian methods that do themselves scale often require extensive interaction with the environment to perform well, being inappropriate for high stakes or costly applications such as healthcare. In this paper we introduce our method, Approximate Variational Reward Imitation Learning (AVRIL), that addresses both of these issues by jointly learning an approximate posterior distribution over the reward that scales to arbitrarily complicated state spaces alongside an appropriate policy in a completely offline manner through a variational approach to said latent reward. Applying our method to real medical data alongside classic control simulations, we demonstrate Bayesian reward inference in environments beyond the scope of current methods, as well as task performance competitive with focused offline imitation learning algorithms.

1 INTRODUCTION

High-stakes applications may provide only logged demonstrator trajectories, while Bayesian IRL methods struggle with scalability and costly environment interaction. AVRIL addresses these obstacles by jointly learning an approximate reward posterior and imitator policy entirely offline through a variational framework.

  • High-stakes settings such as clinical decision making require learning from logged trajectories because deploying untrained policies on patients is impractical and unethical.The available demonstrations contain visited states and taken actions, without requiring policy testing during training.
  • Bayesian reward inference is valuable because inverse reinforcement learning seeks demonstrators’ preferences and goals in addition to reproducing their policies.This motivates inferring the reward before learning a policy that is optimal with respect to it.
  • AVRIL jointly learns a variational posterior distribution over latent rewards and an imitator policy, extending Bayesian IRL to flexible functions without costly inner-loop operations or environment interaction.The framework amortises the reward distribution over demonstrations and uses it to inform a Q-function policy.

2 APPROACHING APPRENTICESHIP AND IMITATION OFFLINE

Offline apprenticeship learning seeks policies and reward explanations from logged demonstrations without environment dynamics, intrinsic rewards, or interaction. Existing approaches face limited reward representations, scaling constraints, and costly forward-RL or environment-evaluation requirements.

  • Problem setting: The paper frames the learning setting as demonstrations from an optimally acting policy, represented under the Markov assumption as state-action-next-state-next-action tuples.The environment dynamics and underlying reward remain unavailable, and further environment sampling is not possible.
  • Implicit-reward policy learning: Offline implicit-reward methods can improve imitation performance, but their rewards remain approximate and evaluable only at states observed in demonstrations, limiting interpretability.This limitation distinguishes policy matching from explicit reward understanding.
  • Bayesian inverse reinforcement learning: Bayesian inverse reinforcement learning is motivated by recovering rewards and uncertainty from demonstrations, but exact inference yields an intractable posterior and relies on MCMC sampling.The demonstrator policy is modeled probabilistically from state-action values, with β representing confidence in demonstrator optimality.
  • Bayesian inverse reinforcement learning: Classical Bayesian inverse reinforcement learning restricts rewards to linear combinations of hand-crafted state features.This limits the function class available for representing rewards.
  • Bayesian inverse reinforcement learning: Finite-state assumptions create additional scaling and convergence problems, with finite approximations scaling at O(|S|2) and inheriting MCMC difficulties.The rapid-mixing bounds do not hold in the infinite case, while finite approximations can become vacuous.
  • Bayesian inverse reinforcement learning: Most importantly, reward sampling requires an inner-loop MDP solve to evaluate likelihoods through Q-values, which becomes problematic with deep function approximators.This amounts to running forward reinforcement learning for each sampled reward.

3 APPROXIMATE VARIATIONAL REWARD IMITATION LEARNING

AVRIL uses variational inference to jointly learn a reward posterior and imitator policy while avoiding repeated inner-loop reward-specific MDP solves. Its flexible function-approximation framework supports offline learning, reward regularisation, and explicit reward querying.

  • Variational Bayesian approach: AVRIL jointly learns a variational reward distribution and policy using a single constrained objective rather than repeated reward-specific solves.The algorithm updates variational and policy parameters together through minibatch Monte Carlo estimates and gradient steps.
  • Variational Bayesian approach: The method approximates the posterior with qφ(R), using flexible priors and variational distributions whose parameters can be produced by an encoder network.The Gaussian-process prior and Gaussian variational posterior are presented as convenient choices rather than necessities.
  • A further approximation: The central approximation addresses the difficulty of expressing Q-values smoothly as functions of sampled rewards and differentiating through repeated offline Q-learning updates.The paper characterises the required Q-value dependence as harder than ordinary forward reinforcement learning and notes that differentiating through unbounded update sequences is impractical.
  • Implementation: AVRIL avoids inner loops and accommodates nonlinear rewards and commonly used state spaces through arbitrary function approximators.The authors state that simultaneous optimisation removes costly inner loops and can make training faster than MCMC methods by roughly their required number of samples.
  • Temporal consistency through reward regularisation: The objective combines action matching with reward-policy consistency, while the KL term regularises implied rewards against the prior and can discourage reward overfitting.With a standard normal prior, the KL divergence acts as a sparsity regulator on the implied reward.
  • Temporal consistency through reward regularisation: State-only rewards keep dynamics information in the policy and enforce temporal consistency, whereas state-action rewards can move temporal information into the reward.The latter may provide a more interpretable or useful reward for some tasks, but can be detrimental to the policy.

4 EXPERIMENTS

AVRIL is evaluated on healthcare data, classic control environments, and a gridworld to assess offline policy learning and reward inference. It performs strongly across tasks while providing interpretable and uncertain reward representations.

  • Healthcare evaluation: AVRIL achieves the highest accuracy and precision scores on both healthcare tasks evaluated with MIMIC-III.Healthcare performance is assessed by matching actions against held-out demonstrations using ACC, AUC, and APS.
  • Control environments: Across CartPole, Acrobot, and LunarLander, AVRIL is competitive and can reach demonstrator-level performance, although dedicated offline imitation methods sometimes learn faster.These experiments test generalisability across standard control environments of varying complexity.
  • Reward insight: AVRIL’s learned reward can support effective policy retraining, though retraining with the inferred reward performs slightly below the jointly trained AVRIL policy.The reward representation is intended to expose demonstrator preferences beyond black-box policy performance.
  • Reward insight: In the healthcare reward analysis, ventilation receives much higher reward as blood oxygen falls below the ICU-average level.A linear state-only parameterisation additionally highlights blood pressure as the largest contributor to estimated state quality.
  • Gridworld ground-truth comparison: In gridworld, posterior reward uncertainty resembles the complement of expert state occupancy, indicating greater epistemic uncertainty in unvisited regions.The comparison visualises ground-truth reward, expert occupancy, posterior mean, and posterior standard deviation.

5 CONCLUSIONS

The paper presents AVRIL as a scalable Bayesian IRL method for unknown, large environments and reports strong offline imitation and reward-learning performance. It also identifies uncertainty-aware safe imitation as an open direction.

  • 5 CONCLUSIONS: AVRIL addresses scalability barriers in Bayesian IRL by learning effective offline imitation policies and useful reward representations on real and toy data.The inferred reward supports policy retraining and provides insight into demonstrator preferences.
  • 5 CONCLUSIONS: The method remains approximate, leaving room for more exact inference and guarantees on the maximum divergence.The conclusion frames uncertainty-aware safe imitation as especially important for avoiding catastrophic failure in high-stakes settings.

A EXPERIMENTAL SETUP

The experiments use neural implementations and compare AVRIL with offline IRL and imitation-learning baselines under controlled trajectory budgets. Control performance is estimated through repeated live rollouts after training on sampled expert trajectories.

  • Expert Demonstrators: Expert trajectories come from pretrained, hyperparameter-optimised RL Baselines Zoo agents, using DQNs for Acrobot and LunarLander and PPO2 for CartPole.Trajectories are subsampled every 20th step for Acrobot and CartPole and every 5th step for LunarLander.
  • Testing setup: Control algorithms train on 1, 3, 7, 10, or 15 trajectories sampled from 1,000 expert trajectories, then undergo repeated live evaluation.Each method is trained to convergence, tested over 300 rollouts, and the process is repeated 10 times with different initialisations and trajectories.
  • Implementations: All compared methods use neural networks with two hidden layers of 64 ELU-activated units.This shared architecture controls a major implementation choice across methods.
  • Implementations: The benchmark set includes VDICE, DSFN, and EDM implementations from publicly available code repositories.The listed baselines cover offline adversarial imitation, reward-and-policy learning, and energy-based distribution matching.
  • Implementations: VDICE is adapted from continuous to discrete actions by replacing its Normal-distribution output with a Gumbel-softmax.This modification is specific to the control experiments.

B PROOFS

The proof rewrites the constrained variational objective as a negative-value minimisation combining behavioural cloning with a Gaussian reward-prior regulariser. The resulting KL term depends on posterior reward variance and mean.

  • B PROOFS: The constrained objective is transformed into minimising its negative value, with the behavioural-cloning loss appearing as the data negative log-likelihood.The proof identifies L_BC as the negative log-likelihood of the demonstrations.
  • B PROOFS: Under the standard Gaussian reward prior p(R(s, a)) = N(R; 0, 1), the Gaussian posterior KL contributes variance and squared-mean terms.The displayed expression includes -log(Var_qφ[R(s, a)]) + Var_qφ[R(s, a)] - 1 + E_qφ[R(s, a)]^2.
  • B PROOFS: The posterior reward mean is related to the expected temporal-difference quantity Q_θ(s, a) − γQ_θ(s′, a′), with expectations approximated from dataset samples.The proof then introduces a function of the posterior variance in simplifying the objective.
Loading 2102.06483v2…