PaperlayerView exact revision on arXiv

Source-linked AI summary

Safe Imitation Learning via Fast Bayesian Reward Inference from Preferences

Daniel S. Brown, Russell Coleman, Ravi Srinivasan, Scott Niekum

arXiv:2002.09089v4cs.LGstat.ML

TL;DR

High-confidence evaluation in imitation learning is difficult because reward samples are unavailable and Bayesian reward inference can be computationally intractable for complex tasks. Bayesian REX combines preferences over demonstrations with learned reward features for fast Bayesian inference. It supports high-confidence policy evaluation without reward samples and can generate 100,000 reward hypotheses in about five minutes on a personal laptop.

  • Problem

    High-confidence policy evaluation is less developed for imitation learning because reward samples are unavailable, while Bayesian IRL can be computationally intractable for complex problems.

    Source: §1 · p. 1

  • Method

    Bayesian REX uses preferences over demonstrations and learned features to perform efficient Bayesian reward inference for high-dimensional visual imitation learning.

    Source: §1 · p. 1, §5 · p. 4, §5.1 · p. 5

  • Results

    Bayesian REX supports high-confidence evaluation on Atari imitation tasks without reward samples; 100,000 reward hypotheses take about five minutes on a personal laptop.

    Source: §1 · pp. 1–2, §5.1 · p. 6

  • Takeaways & Limitations

    The bounds can rank evaluation policies by risk and expected return and provide evidence useful for detecting reward hacking or gaming.

    Source: §1 · p. 2

    The safety bounds depend on good feature pre-training, rapid MCMC mixing, and accurate preferences over demonstrations.

    Source: §7 · p. 9

Abstract

from arXiv · show

Bayesian reward learning from demonstrations enables rigorous safety and uncertainty analysis when performing imitation learning. However, Bayesian reward learning methods are typically computationally intractable for complex control problems. We propose Bayesian Reward Extrapolation (Bayesian REX), a highly efficient Bayesian reward learning algorithm that scales to high-dimensional imitation learning problems by pre-training a low-dimensional feature encoding via self-supervised tasks and then leveraging preferences over demonstrations to perform fast Bayesian inference. Bayesian REX can learn to play Atari games from demonstrations, without access to the game score and can generate 100,000 samples from the posterior over reward functions in only 5 minutes on a personal laptop. Bayesian REX also results in imitation learning performance that is competitive with or better than state-of-the-art methods that only learn point estimates of the reward function. Finally, Bayesian REX enables efficient high-confidence policy evaluation without having access to samples of the reward function. These high-confidence performance bounds can be used to rank the performance and risk of a variety of evaluation policies and provide a way to detect reward hacking behaviors.

1. Introduction

Imitation learning needs high-confidence performance bounds even though reward samples are unavailable, and Bayesian reward inference is often too costly for complex tasks. Bayesian REX uses preferences over demonstrations to make high-dimensional reward inference and policy evaluation more efficient.

1. Introduction

  • High-confidence policy evaluation is less developed for imitation learning, where reward samples are unavailable.High-confidence off-policy evaluation has received more attention in reinforcement learning.Source: §1 · p. 1
  • Bayesian IRL can be intractable on complex problems because its inner loop repeatedly solves an MDP, raising computation and possibly sample costs.The paper proposes preferences over demonstrations as a route to faster Bayesian reward learning.Source: §1 · p. 1
  • Bayesian REX uses a pairwise ranking likelihood to increase the efficiency of sampling posterior reward functions.The paper also formalizes high-confidence policy evaluation for imitation learning.Source: §1 · p. 1
  • Neural function approximation and self-supervised learning provide useful reward features for Bayesian inference from visual demonstrations.The paper evaluates the approach on Atari imitation learning tasks.Source: §1 · p. 1
  • High-confidence performance bounds rank evaluation policies by risk and expected return and provide evidence useful for detecting reward hacking or gaming.The bounds are computed without access to reward-function samples.Source: §1 · p. 2

2. Related work

Related work spans imitation learning, safe imitation, value alignment and preference learning, safe reinforcement learning, and Bayesian neural networks. The paper positions its contribution as efficient high-confidence policy evaluation from visual demonstrations without an MDP solver during reward inference.

2.1. Imitation Learning

  • Behavioral cloning learns state-to-action mappings from demonstrations but can suffer compounding errors; IRL instead estimates the demonstrator’s reward, with Bayesian IRL sampling reward functions.Deep IRL often returns only a point estimate, limiting uncertainty and robustness analysis.Source: §2.1 · p. 2

2.2. Safe Imitation Learning

  • Prior Bayesian high-confidence evaluation for imitation learning uses an inner-loop MDP solver, while this work targets visual demonstrations without a solver or model during reward inference.Other robust policy optimization methods also retain an inner-loop solver requirement.Source: §2.2 · p. 2

2.3. Value Alignment and Active Preference Learning

  • Prior active preference learning scales to high-dimensional tasks but can require many active queries; Bayesian REX uses preferences over demonstrations without active queries for reward inference or policy optimization.It also provides efficient high-confidence bounds for evaluation policies.Source: §2.3 · pp. 2–3

2.4. Safe Reinforcement Learning

  • Safe RL studies risk-sensitive or robust MDP policies and high-confidence off-policy evaluation; this work addresses imitation learning without observed reward samples or a known demonstrator policy..Source: §2.4 · p. 3

2.5. Bayesian Neural Networks

  • Bayesian neural-network methods quantify output uncertainty, but this paper focuses on uncertainty in policy performance under an uncertain reward, including policy and environment stochasticity..Source: §2.5 · p. 3

3. Preliminaries

The preliminaries specify an MDP with a state-based reward and introduce Bayesian IRL as inference over an unknown reward. Its posterior sampling requires solving for optimal Q-values, motivating a more efficient inference approach.

3. Preliminaries

  • The environment is modeled as an MDP with states, actions, transition dynamics, a state-based reward, an initial-state distribution, and a discount factor.The authors note that their approach also extends to action- or transition-dependent rewards.Source: §3 · p. 3
  • Bayesian IRL infers an unknown reward for a Boltzmann-rational demonstrator, with β representing the demonstrator’s confidence in acting optimally.The likelihood of demonstrations is defined under the Boltzmann-rationality assumption.Source: §3 · p. 3
  • Bayesian IRL samples the reward posterior using MCMC, but evaluating each proposal requires solving for optimal Q-values, limiting its use to low-dimensional reward features.The paper identifies efficient Bayesian inference from preferences over demonstrations as a contribution.Source: §3 · p. 3

4. High Confidence Policy Evaluation for Imitation Learning

HCPE-IL asks for a high-confidence lower bound on an evaluation policy’s performance under the demonstrator’s unknown true reward. Its inputs include the evaluation policy, demonstrations, and a performance statistic.

4. High Confidence Policy Evaluation for Imitation Learning

  • The setup assumes an MDP without its reward, an evaluation policy, demonstrations, confidence level δ, and a performance statistic.Demonstrations may be complete or partial trajectories of states or state-action pairs.Source: §4 · p. 4
  • HCPE-IL returns a high-confidence lower bound on the performance statistic for the evaluation policy under the demonstrator’s true reward.The bound is based on the evaluation policy and demonstration set.Source: §4 · p. 4
  • The target quantity is performance under the demonstrator’s true reward, which is unknown in the imitation-learning setting..Source: §4 · p. 4

5. Deep Bayesian Reward Extrapolation

Bayesian REX uses trajectory preferences and pretrained state features to make Bayesian reward inference practical for visual imitation learning. The section also shows how the learned feature representation supports efficient high-confidence policy evaluation.

5. Deep Bayesian Reward Extrapolation

  • Unlike T-REX's point estimate, Bayesian REX uses pairwise preference likelihood and MCMC to sample reward posteriors without MDP solves or rollouts during inference.Source: §5 · p. 4

5.1. Optimizations

  • Precomputing frozen features reduces N-sample MCMC cost from O(mTN|Rθ|) to O(mT|Rθ|+mkN), with only O(mk) work per proposal.The trajectories pass through the feature network once; proposal evaluations then use the precomputed feature sums.Source: §5.1 · p. 5
  • 100,000 reward hypotheses from 66 preference labels take approximately 5 minutes on a personal laptop without a GPU, versus 10+ hours of training for one standard Bayesian IRL posterior sample.The comparison is with a parallelized PPO agent running on an NVIDIA TITAN V GPU for standard Bayesian IRL.Source: §5.1 · p. 6

5.2. Pre-training the Reward Function Network

  • Self-supervised inverse and forward dynamics, temporal-distance prediction, and pixel autoencoding train features from demonstrations, targeting movement, short-term change, long-term progress, and state content.Source: §5.2 · p. 6

5.3. HCPE-IL via Bayesian REX

  • For an evaluation policy, one feature-count evaluation yields posterior return samples via WΦ; the δ-quantile supplies the (1−δ) lower confidence bound.The paper focuses on bounding the policy's expected value.Source: §5.3 · p. 6, §5.1 · p. 6

6. Experimental Results

Experiments compare Bayesian REX with Bayesian IRL and other imitation-learning methods, then test its policy-evaluation bounds and human-demonstration rankings. Results vary with the demonstrations available, while the policy evaluations expose risk that high expected returns alone can miss.

6.1. Bayesian IRL vs. Bayesian REX

  • With more than five suboptimal ranked demonstrations, Bayesian REX performs on par with or sometimes better than Bayesian IRL; with only such ranked demonstrations, it significantly outperforms Bayesian IRL across tested ranking thresholds.The comparison also tested Bayesian IRL adapted to good and bad demonstrations.Source: §6.1 · p. 7
  • With more than five optimal demonstrations, Bayesian IRL significantly outperforms Bayesian REX; Bayesian REX is nearly as effective when there are five or fewer, and Bayesian IRL is recommended for very few unlabeled near-optimal demonstrations.When an efficient MDP solver is unavailable, Bayesian IRL is infeasible and Bayesian REX is presented as the natural choice.Source: §6.1 · p. 7

6.2. Visual Imitation Learning via Bayesian REX

  • Across five Atari games, Bayesian REX’s average ground-truth return exceeds the best demonstration on every game and beats T-REX on three.The comparison also includes GAIL; Bayesian REX and T-REX use 12 demonstrations with ground-truth preferences, while GAIL uses 10 demonstrations comparable to the best demonstration.Source: §6.2 · p. 7, Table 2 · p. 8

6.3. High-Confidence Policy Performance Bounds

  • For Beam Rider, No-Op has high posterior mean return but a very negative 0.05-VaR, while checkpoint policies’ posterior returns correlate with ground-truth returns.This illustrates how the lower-tail bound can reveal risk obscured by a high expected return.Source: Table 3 · p. 8, §6.3 · p. 8
  • In Breakout, ranked MAP and No-Op rollouts substantially change both policies’ risk profiles; afterward No-Op has low expected returns and high risk, while MAP’s 95%-confidence lower bound rises.The original demonstrations did not rule out the possibility that only lower brick layers should be hit, leaving MAP out of distribution and potentially risky.Source: Table 4 · p. 8, §6.3 · p. 9

6.4. Human Demonstrations

  • For human demonstrations, Bayesian REX correctly ranks Beam Rider behaviors and assigns the greatest risk to Enduro’s ram demonstration and the least to its good demonstration.The Beam Rider reward correctly favors destroying aliens and avoiding bullets rather than merely flying and shooting.Source: §6.4 · p. 9

7. Conclusion

Bayesian REX makes Bayesian reward inference tractable for high-dimensional visual imitation learning and supports high-confidence evaluation without reward samples. Its safety bounds remain conditional on feature pre-training, MCMC mixing, and preference accuracy.

7. Conclusion

  • Bayesian REX achieved state-of-the-art performance in 3 of 5 Atari games and supports high-confidence performance bounds for arbitrary evaluation policies.The bounds can rank policies by risk and expected return and may help detect reward hacking or value misalignment.Source: §7 · p. 9
  • Bayesian REX addresses the computational bottleneck of Bayesian reward learning, where repeated inner-loop MDP solves can take hours or days.The algorithm samples tens of thousands of reward functions in minutes on a consumer laptop.Source: §7 · p. 9
  • The safety bounds are only safe under assumptions of good feature pre-training, rapid MCMC mixing, and accurate preferences over demonstrations.The paper identifies exploratory trajectories, detecting missing features, and safe policy optimization as future work.Source: §7 · pp. 9–10

A. Source Code and Videos

  • The paper provides code repositories for Atari and gridworld experiments.The Atari repository and gridworld code are listed separately.Source: A. Source Code and Videos · p. 13

B. MCMC Details

  • MCMC proposals normalize Gaussian perturbations of the current weight vector to keep samples on the L2-unit sphere.The proposal uses β = 1 and standard deviation σ; the chain uses 200,000 steps with a 5,000-step burn-in and every twentieth sample retained.Source: B. MCMC Details · p. 13
  • A non-negative return prior constrains the lowest-ranked demonstration without changing ranking-likelihood ratios, which are invariant to affine reward transformations.Proposals with negative return for the worst demonstration receive prior probability −∞.Source: B. MCMC Details · p. 13

C. Bayesian IRL vs. Bayesian REX

The gridworld comparisons show when Bayesian REX's ranked demonstrations can match or outperform Bayesian IRL, and when optimal demonstrations favor Bayesian IRL. The paper's summary makes the choice conditional on demonstration quality and the availability of a tractable MDP solver.

C. Bayesian IRL vs. Bayesian REX

  • The comparisons use low-dimensional 6×6 gridworlds and prioritize sample efficiency, while Bayesian IRL is much slower because it repeatedly solves an MDP.The evaluations vary demonstration conditions to suit each algorithm's assumptions; Bayesian REX requires no MDP access during reward inference.Source: C. Bayesian IRL vs. Bayesian REX · pp. 13–14

C.1. Ranked Suboptimal vs. Optimal Demonstrations

  • With more than 5 ranked suboptimal demonstrations, Bayesian REX performs on par with or slightly better than Bayesian IRL given the same number of optimal demonstrations.The comparison uses demonstrations starting from the same states; two ranked demonstrations provide only one reward constraint.Source: Table 7 · p. 14, §C.1 · p. 14

C.2. Only Ranked Suboptimal Demonstrations

  • Bayesian REX significantly outperforms Bayesian IRL when both receive ranked suboptimal demonstrations, including against a Bayesian IRL variant using good and bad demonstrations.The comparison labels the top and bottom x% of ranked demonstrations as good and bad.Source: §C.2 · pp. 14–15, Table 8 · p. 15, Table 9 · p. 15

C.4. Summary

  • With very few unlabeled near-optimal demonstrations, Bayesian IRL is the natural choice; otherwise Bayesian REX is competitive and often superior, especially when an efficient MDP solver is unavailable.Bayesian IRL becomes infeasible without a highly tractable MDP solver.Source: §C.4 · p. 15

C.3. Only Optimal Demonstrations

  • When both methods receive optimal demonstrations, Bayesian IRL outperforms Bayesian REX; with 5 or fewer demonstrations, Bayesian REX performs only slightly worse.The Bayesian REX comparison uses random rollouts automatically labeled less preferred than the optimal demonstrations.Source: Table 10 · p. 15, §C.3 · p. 15

D. Pre-training Latent Reward Features

The feature encoder is pre-trained with self-supervised and preference-ranking objectives to mitigate overfitting to limited preferences and capture varied state information. The appendix describes the training setup, qualitative feature visualizations, loss ablations, and policy-evaluation rollout procedure.

D. Pre-training Latent Reward Features

  • Self-supervised auxiliary objectives complement preference rankings because embeddings trained only on limited preferences may overfit and miss ground-truth reward features.The tasks predict actions, future states, temporal distance, and reconstructed pixels.Source: D. Pre-training Latent Reward Features · p. 15
  • The auxiliary tasks target different information: inverse dynamics captures agent movement, forward dynamics short-term environmental change, temporal distance long-term progress, and the autoencoder state information.The temporal-distance, dynamics, and T-REX heads use single layers to encourage features interpretable through linear combinations.Source: D. Pre-training Latent Reward Features · pp. 15–16
  • The Atari encoder processes four stacked 84×84 grayscale frames through convolutional layers and fully connected layers to produce φ(s).For a 64-dimensional encoding, the final fully connected layers map 784 to 128 and then 128 to 64.Source: D. Pre-training Latent Reward Features · p. 16
  • The T-REX objective learns scalar rewards from preferred and less-preferred samples drawn from different demonstrations, while the VAE reconstructs the original framestack.The VAE combines reconstruction cross-entropy with a penalty on divergence from a unit normal.Source: D. Pre-training Latent Reward Features · p. 16

D.1. Training specifics

  • Training used 12 demonstrations and 60,000 trajectory-snippet pairs for most games, with a different full-trajectory sampling procedure for Enduro.The snippets span 50–100 timesteps; Enduro uses 10,000 trajectory pairs and two passes through the data.Source: §D.1 · p. 17

E. Visualizations of Learned Features

  • In Enduro, latent features change with passing events in directions consistent with the ground-truth +1 reward for passing and −1 for being passed.In Space Invaders the encoding focuses on enemy-ship movement but appears to ignore the player ship; Breakout features can spike when bricks break and track elapsed progress.Source: E. Visualizations of Learned Features · pp. 17–18

F. Imitation Learning Ablations for Pre-training φ(s)

  • Combining ranking and self-supervised losses achieves high performance across games, outperforming self-supervised-only pretraining and surpassing ranking-only pretraining on 3 of 5 games.Performance varies across the compared pretraining schemes.Source: F. Imitation Learning Ablations for… · p. 18

G. Suboptimal Demonstration Details

I.2. Evaluation Policies

The evaluation policies include PPO checkpoints, a No-Op policy, and mean- and MAP-reward policies. Bayesian REX’s posterior bounds track ground-truth performance across the reported Enduro, Seaquest, and Space Invaders examples, while Seaquest exposes a No-Op failure that one ranked demonstration corrects.

I.2. Evaluation Policies

  • Evaluation uses PPO checkpoints and a No-Op policy; Table 12 compares feature pre-training schemes using 12 demonstrations and averages over 30 rollouts.The table reports ground-truth average returns when optimizing Bayesian REX’s mean and MAP rewards.Source: §I.2 · p. 18, Table 12 · p. 19
  • For Enduro, posterior average returns and 0.05-VaR bounds predict ground-truth performance, and rank No-Op below the other policies.The learned posterior also places the MAP policy at more than twice policy D’s performance and the Mean policy between policies B and C.Source: §I.2 · p. 19, Table 13 · p. 19
  • For Seaquest, the bounds rank policies A and B below C and D, but initially overrate No-Op; one ranked No-Op demonstration fixes its relative ranking.The ground-truth performance of C and D is too close for the mean and 0.05-VaR bound to distinguish statistically.Source: §I.2 · p. 19
  • For Space Invaders, mean performance and the 95%-confidence lower bound indicate ground-truth performance, while No-Op receives low average performance and a low lower bound.No-Op causes the agent to be hit by alien lasers early in the game.Source: §I.2 · p. 19

J. Different Evaluation Policies

The evaluations test whether Bayesian REX can rank policies that exploit learned rewards and varied human demonstrations. It identifies a Beam Rider reward hack as worse than the RL and Bayesian REX policies, and its bounds and posterior averages rank the reported Beam Rider human behaviors.

J. Different Evaluation Policies

  • In Seaquest, No-Op is initially predicted too favorably, but adding one ranked No-Op rollout yields correct relative rankings for that policy.The table compares posterior policy-evaluation statistics with ground-truth game scores.Source: Table 14 · p. 20
  • Bayesian REX predicts the Beam Rider reward-hacking policy is worse than both the ground-truth-reward RL policy and its own policy, despite the latter’s weaker actual performance.The exploit keeps the ship between enemy firing lanes by repeatedly pressing left and right, allowing it to live for a long time.Source: J. Different Evaluation Policies · pp. 19–20, Table 16 · p. 20

K.1. Beamrider

  • For Beam Rider human demonstrations, posterior averages and high-confidence bounds correctly rank behaviors, rewarding destroying aliens and avoiding shots over merely flying and shooting.The demonstrations include good, poor, pessimal, and adversarial behavior.Source: §K.1 · p. 20, K. Human Demonstrations · p. 20

K.2. Space Invaders

The section evaluates human demonstrations and compares Bayesian REX with ensemble and dropout uncertainty estimates. Space Invaders rankings are affected by trajectory length, Enduro risk rankings distinguish good from ramming behavior, and Bayesian REX is competitive in learned-reward performance while ensemble bounds can be uninformative.

K.2. Space Invaders

  • With equal-length Space Invaders demonstrations, Bayesian REX ranks good, every-other-column, and hold ’n fire among the best, while unrestricted trajectory length can inflate hide and miss estimates.Hide and miss survive longer than policies that hit aliens, producing higher mean and 0.05-quantile predicted performance; the equal-length check uses the first 600 timesteps.Source: §K.2 · pp. 20–21

K.3. Enduro

  • For Enduro demonstrations, Bayesian REX predicts performance and risk, assigning the greatest risk to ram and the least to good.The comparison includes periodic and neutral behaviors as well as good and ram.Source: §K.3 · p. 21

L.1. T-REX Ensemble

  • The T-REX ensemble uses five differently initialized networks trained on the same demonstrations, with randomized subsampled snippets, to estimate return distributions.Return estimates come from a trajectory passed through each network or from policy rollouts through ensemble members.Source: §L.1 · pp. 21–22

L.2. MC Dropout

  • MC Dropout shares a sampled last-layer dropout mask across each preference pair and estimates return distributions from 50 masks applied to 100 rollouts.Its base network is trained to predict returns without Bayesian REX’s auxiliary losses, using a dropout probability of 0.5.Source: §L.2 · p. 22

L. Comparison with Other Methods for Uncertainty Quantification

  • For RL on learned reward functions, Bayesian REX is superior or competitive with T-REX Ensemble and MC Dropout across all games except Beam Rider, where MC Dropout performs much better.Source: §L.2 · p. 22

L.3. T-REX Ensemble High-Confidence Bounds

  • T-REX ensemble expected returns can be accurate while its 95% lower bounds fail to represent risk; other cases either produce meaningless predictions or good predictions.For one comparison, predictions are normalized against ensemble-member returns for policy A.Source: §L.3 · p. 22

L.4. MC Dropout Results

MC Dropout’s high-confidence bounds sometimes identify No-Op risk but often fail to do so across the reported games. The paper notes that its posterior approximation is not principled Bayesian inference and may not capture the epistemic risk of interest.

L.4. MC Dropout Results

  • MC Dropout predicts high risk for Beam Rider and Space Invaders No-Op policies, but often fails to flag No-Op as high risk in other games.The authors describe it as sometimes useful in practice despite not performing full Bayesian inference.Source: §L.4 · p. 22, Table 27 · p. 23, Table 31 · p. 24
  • MC Dropout’s distribution does not concentrate as data samples grow, so it may not measure the epistemic risk targeted by high-confidence evaluation.The authors call for further work on more sophisticated dropout approaches to uncertainty estimation.Source: §L.4 · p. 22
  • Bayesian REX has comparable or better learned-reward policy performance across the games except Beam Rider, where MC Dropout performs better.The comparison reports averages and standard deviations over 30 rollouts.Source: Table 21 · p. 23
Loading 2002.09089v4…