Scalable Bayesian Inverse Reinforcement Learning
Authors: Alex James Chan
Mihaela van der Schaar
Venue: ICLR 2021 Poster (2021)
Keywords: Bayesian
Inverse reinforcement learning
Imitation Learning

Abstract: 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.

Published as a conference paper at ICLR 2021
SCALABLE BAYESIAN LEARNING
INVERSE
REINFORCEMENT
Alex J. Chan
Mihaela van der Schaar
University of Cambridge, Cambridge, UK a pleexchaa@mmatee cam. ac.uk
University of Cambridge, Cambridge, UK University of California, Los Angeles, USA Cambridge Centre for Al in Medicine, UK
The Alan Turing Institute, London, UK mv4 7 2 cam. ac. uk
ABSTRACT
Bayesian inference over the reward presents an ideal solution to the ill-posed na- ture of the inverse reinforcement learning problem. Unfortunately current meth- ods 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 costly applications such as healthcare. In this paper we introduce our method, Approximale Variational Reward Imitation Learning (AVRIL), that addresses both of these issues by jointly learning an ap- proximate posterior distribution over the reward that scales to arbitrarily compli- cated state spaces alongside an appropriate policy in completely offline manner through variational approach to said latent reward. Applying our method to real medical data alongside classic control simulations, we demonstrate Bayesian re- ward inference in environments beyond the scope of current methods, as well as task performance competitive with focused offline imitation learning algorithms.
INTRODUCTION
For applications in complicated and high-stakes environments it can often mean operating in the minimal possible setting that is with no access to knowledge of the environment dynamics nor intrinsic reward, nor even the ability to interact and test policies. this case learning and inference must be done solely on the basis of logged trajectories from a competent demonstrator showing only the states visited and the the action taken in each case.
Clinical decision making is important example of this, where there is great interest in learning policies from medical professionals but completely impractical and unethical to deploy policies on patients mid-training. Moreover this is an area where it is not only the policies, but also knowledge of the demonstrator's preferences and goals, that we are interested in. While imitation learning (IL) generally deals with the problem of producing appropriate policies to match demonstrator, with the added layer of understanding motivations this would then usually be approached through inverse reinforcement learning (IRL). Here attempting to learn the assumed underlying reward driving the demonstrator, before secondarily learning a policy that is optimal with respect to the reward using some forward reinforcement learning (RL) technique. By composing the RL and IRL procedures in order to perform IL. we arrive at appremticeship learning (AL), which introduces its own challenges, particularly the offline setting. Notably for any given set of demonstrations there are (infinitely) many rewards for which the actions would be optimal (Ng et al., 2000). Max-margin (Abbeel & Ng 2004) and max max-entropy (Ziebart al., 2008) methods for heuristically differentiating plausible rewards do so at the cost of potentially dismissing the true reward for not possessing desirable qualities. On the other hand a Bayesian approach to IRL (BIRL) is more conceptually satisfying, taking a probabilistic view of the reward, we are interested in the posterior distribution having seen the demonstrations (Ramachandran & Amir 2007). accounting for all possibilities. BIRL is not without its own drawbacks though, noted Brown & Nickum (2019), making it inappropriate for modern complicated environments: assuming linear rewards; small, solvable environments, and repeated, inner-loop, calls to forward RI.

---

Published as a conference paper at ICLR 2021
0011 B01
5
4 Reward Encoder
π Policy
le
SXA
Q-function Decoder
Reward Posterior
Behaviour
Demonstration
Figure 1: Overview. AVRIL framework for BIRL that works through an approximation in the variational Bayesian framework, considering the reward to be latent representation of behaviour. distribution over the reward, which is amortised over the demonstration space, is learnt that then informs an imitator Q-function policy. The dotted line represents a departure from traditional auto-encoder as the input, alongside the latent reward, informs the decoder.
The main contribution then of this paper is method for advancing BIRL beyond these obstacles, allowing for approximate reward inference using an arbitrarily flexible class of functions, in any environment, without costly inner-loop operations, and importantly entirely offline. This leads to our algorithm AVRII., depicted figure [ which represents framework for jointly learning variational posterior distribution over the reward alongside an imitator policy in auto-encoder- esque manner. In what follows we review the modern methods for offline IRL/IL (Section 2 with focus the approach of Bayesian IRL and the issues it faces when confronted with challenging environments. We then address the above issues by introducing our contributions (Section 3), and demonstrate the gains of our algorithm in real medical data and simulated control environments, notably that it is now possible to achieve Bayesian reward inference in such settings (Section 4 Finally wrap up with some concluding thoughts and directions (Section 5 Code for AVRIL and our experiments is made available at h com Xande rCC http///caaaalleeae..e bbbebbbrr) and http:////w.aaao////aa/aa/aaaaaaa
2 APPROACHING APPRENTICESHIP AND IMITATION OFFLINE
Preliminaries. We consider the standard Markov decision process (MDP) environment, with states E S, actions E A, transitions T E A rewards R E S and discount Y [0, 1]. For policy E A(A)S let Px (s, unique occupancy measure alongside the state-only occupancy measure (s((s) E.e (s, Despite this full environment model, the only information available to is the MDP,RT, in that we have no access to either the underlying reward or the transitions, with our lacking knowledge of the transitions being also strong in the sense that further to sample them. The learning signal then given by access to m-many trajectories of some demon- strator assumed be acting optimally w.r.t. the MDP, following policy TD. making up a data set
7 E,, F ( (s,=sa,,=a)
the induced
unable simulate the environment
Draw = {({]]) a1 ..... +00> a -(0) )mm where ( the state and a, is the action taken at step t during the ith demonstration, and (() is the (max) time horizon of the ith demonstration. Given the Markov assumption though it sufficient and convenient to consider the demonstrations simply as collection of n-many state, action, next state, next action tuples such that D {(s ag & a 1 with Σ (7() 1).
Apprenticeship through rewards. Typically AI. proceeds by first inferring an appropriate reward function with an IRL procedure (Ng et al. 2000 Ranachandran & Amir 2007 Rothkopf & Dimi- trakakis 2011; Ziebart et al., 2008) before running forward RL to obtain an appropriate policy. This allows for easy mix-and-matee procedures, swapping in different standard RL and IRL methods depending on the situation. These algorithms though depend on either knowledge of in order to solve exactly the ability to perform roll-outs in the environment, with little previous work focus- ing on the entirely offline setting. One simple solution is through attempting to learn the dynamics
WWe define a state-action reward here, is usual in the literature. Extensions to a state-only reward simple, and indeed can be preferable, will later.
2

---

Published as a conference paper at ICLR 2021
(Herman et al. 2016), though without large supply of diverse demonstrations small environ- ment this becomes impractical given imperfections in the model. Alternatively Klein et al. (201T)
and Lee et al. (2019) attempt off- poflco feature matching through least-squared temporal difference and deep neural networks to uncover appropriate feature representations.
Implicit-reward policy learning. Recent work has often forgone explicit representation of the reward. Moving within the maximum- maximummenttropy RL framework (Ziebart 2010: Levine| (2018). (Ho) & Ermon (2016) noted that the full procedure (RL IRL) can be interpreted equivalently as the minimisation of some divergence between occupancy measures of the imitator and demonstrator:
arg min { (px P=D) [=))
with H (T) being the discounted causal entropy (Bloem & Bambos 2014) of the policy and ve the Fenchel conjugate of a chosen regulariser on the form the reward. These are typically optimised in an adversarial fashion (Goodfellow et al. (2044) and given the focus on evaluating Pr this often requires extensive interaction with the environment, otherwise banking on approximations over a replay buffer (Kotrrikov et al. 2018)) or reformulation of the divergence to allow for off-policy evaluation (Kostrikov et al. al.] 2019). Bear in mind that optimal policies within the maximum-entropy framework are parameterised by Boltzmann distribution:
exp(Q(x, a))
(als) = Eoea exp(Q(s,
(2)
with Q(s, the soft Q-function, defined recursively via the soft Bellman-equation:
Q(s, a) _ R(s; a) + +> soft max ((s',a'))
(3)
a'
Then for a learnt parameterised policy given in terms of Q- valuee from function approximator Qo. we can obtain an implied reward given by:
Ro (8, a) = Qo(s, a) - - log Σ exp(Qo(s(.d))
(4)
EE
number of algorithms make use of this fact with Piot et (2014) and Reddy et al. (2019) working by essentially placing a sparsity prior on this implied reward, encouraging it towards zero, and thus incorporating subsequent state information. Alternatively [Jarrttt al. (020)) show that even the simple behavioural cloning (Bain & Sammui 1995) is implicitly maximising some reward with an approximation that the expectation over states is taken with respect to the demonstrator, not the learnt policy. They then attempt to rectify part of this approximation using the properties of the energy-based model implied by the policy (Grathwohl et al. 2019).
The problem with learning an implicit reward in an offline setting is that it remains just that, implicit, only able to be evaluated at points seen in the demonstrations, and even then only approximately. Thus even if their consideration improves imitator policies performance they offer real improve- ment for interpretation.
2.1 BAYESIAN INVERSE REINFORCEMENT LEARNING
We are then resigned to directly reason about the underlying reward, bringing us back to the question of IRL, and in particular BIRL for principled approach to reasoning under uncertainty. Given a prior over possible functions, having seen some demonstrations, we calculate the posterior over the function using theoretically simple application of Bayes rule. |Ramachandran & Amir (2007) defines the likelihood of an action state as Boltzmann distribution with inverse temperature and respective state-action values, yielding a probabilistic demonstrator policy given by:
exp(3QT0 (s,a)) (als, R) ΣDEA c (s.
(5)
where 3 E [0, oo) represents the the confidence in the optimality of the demonstrator. Note that despite similarities, moving forward we are no longer within the maximum-entropy framework and (h(s, a) now denotes the traditional, not soft (as in equation 3. state-action value (Q-value) func- tion given reward R and policy such that Q7 a) E. 7 R(ss))ss s, do a)-
Unsurprisingly this yields an intractable posterior distribution leading to a Markov chain Monte Carlo (MCMC) algorithm based random grid- walk to sample from the posterior.
3

---

Published as a conference paper at ICLR 2021
Issues in complex and unknown environments. This original formulation, alongside extensions that consider maaximum-a-postenin inference (Choi & Kim 2011) and multiple rewards (Choi & Kim 2012 Dimitrakakis & Rothkopf 2011), sufter from three major drawbacks that make them impractical for modern, complicated, and model-free task environments.
1. The reward is a linear combination of state features. Naturally this is a very restrictive class of functions and assumes access to carefully hand-crafted features of the state space.
2. The cardinality of the state-space is finite, |S| < oc. Admittedly this can be relaxed in practical terms, although it does mean the rapid-mixing bounds derived by (Ramachandran & Amir (007) do not hold at all in the infinite case. For finite approximations they scale at O(S) rapidly becoming vacuous and causing BIRL. to inherit the usual MCMC diffi- culties on assessing convergence and sequential computation (Gamerman & Lopes] 2006).
3.
The requirement of an inner-loop MDP solve. Most importantly at every step a new reward is sampled and the likelihood of the data must then be evaluated. This requires calculating the Q-values of the policy with respect to the reward, in other words running forward RL. While not an insurmountable problem in the simple cases where everything is known and can be quickly solved with a procedure guaranteed to converge correctly, this becomes an issue in the realm where only deep function approximation works adequately (i.e. the non- tabular setting). DQN training for example easily stretches into hours (Mnih et al. 2013) and will have to be repeated thousands of times, making it completely untenable.
We have seen that even in the most simple setting the problem of exact Bayesian inference over the reward intractable, and the above limitations of the current MCMC methods are not trivial to overcome. Consequently very little work has been done in the area and there still remain very open challenges. [evine et al. (201T) addressed linearity through a Gaussian process approach, allowing for significantly more flexible and non- non-linear representation though introducing issues of its own, namely the computational complexity of inverting large matrices (Rasmussen) (2003). More recently Brown& & Niekum (2019) have presented the only current solution to the inner-loop problem by introducing an alternative formulation of the likelihood, one based on human recorded pairwise preferences over demonstrations that significantly reduces the complexity of likelihood. However labelled preferences certainly can't 't be assumed always available and while very effective for the given task is not appropriate in the general case. One of the key aspects of our contribution is that we are able to deal with all three of these issues while also not requiring any additional information.
The usefulness of uncertainty. On top of the philosophical consistency of Bayesian inference there are a number of reasons for wanting a measure of uncertainty over any uncovered reward that are not available from regular IRL algorithms. First that the (epistemic) uncertainty revealed by Bayesian inference tells us a lot about what areas of the state-space we really cannot say anything about because we haven't 't seen any demonstrations there potentially informing future data collec- tion if that is possible (Mindermann et al. 2018). Additionally in the cases we are mostly concerned about (e.g. medicine) we have to be very careful about letting algorithms pick actions in practice and we are interested in performing safe or risk-averse imitation, for which a degree of confidence over learnt rewards is necessary. [Bronn et al. (2020) for example use a distribution over reward to optimise conditional value-at-ric riik instead of expected return as to bound potential downsides.
3 APPROXIMATE VARIATIONAL REWARD IMITATION LEARNING
A variational Bayesian approach. In this section we detail our method, AVRIL, for efficiently learning an imitator policy and performing reward inference simultaneously. Unlike the previously mentioned sampling or MaP-based methods, we employ variational inference (Blei et al. 2017) to reason about the posterior. Here we posit surrogate distribution go (R). parameterised by Ø, and aim to minimise the Kullback-Leibler (KL) divergence to the posterior, resulting in an objective:
min (DiL (go ( (H)))
(6)
This divergence is still as troubling as the posterior to evaluate, leading to an auxiliary objective function in the Evidence Lower BOund (ELBO):
F(0) Egg (og)(D) [2]] - DR (oo))(())))((H))
4

---

Published as a conference paper at ICLR 2021
where it can be seen that maximisation over ø is equivalent to (6) We agnostic towards the form of both the prior and variational distribution, for simplicity here we assume a Gaussian process prior with mean zero and unit variance over R alongside the variational posterior distribution given by:
go R) = (R,
(8)
where μ. o² are the outputs of an encoder neural network taking input and parameterised by d. Note that for the algorithm that we will describe these choices are not necessity and can be easily substituted for more expressive distributions if appropriate. Maintaining the assumption of Boltzmann rationality on the part of the demonstrator, our objective takes the form:
2xp(3
F(6) = Ege
Σ
log EDEA exp(3QR (s.
(9)
[s,a]ED
The most interesting (and problematic) part of this objective as ever centres on the evaluation of QR (s, a). Notice that what is really required here is an expression of the Q-values as a smooth function of the reward such that with samples of R is little hope of obtaining this simply, by itself it is a harder problem than that of forward RL which only attempts to evaluate the Q-values for specific R and already in complicated environments has to rely on function approximation and limited guarantees.
could take gradients w.r.t. ø. Of course there
A naive approach would be to sample R and then approximate the Q-values with a second neural network, solving offline over the batched data using a least-squared TD/Q-learnin algorithm, as the approach forced sampling based BIRL methods. It is in fact though doubly inappropriate for this setting, not only does this require a solve as an inner-loop but importantly differentiating through the solving operation is extremely impractical, it requires backpropagating through number of gradient updates that are essentially unbounded as the complexity of the environment increases.
A further approximation. This raises an important question is it possible jointly optimise policy and variational distribution only once instead of requiring repeated solve? This is theoretically suspect, the Q-values are defined on singular reward, constrained as R(s, a) E,r,a'>,,7 (Q (Q a) - ( ( a')] so we cannot learn a particular standard Q-function that re- flects the entire distribution. But can we learn a policy that reflects the expected reward using a second policy neural network Qo? We can't simply optimise 0 alongside é to maximise the ELBO though as that completely ignores the fact that the learnt policy is intimately related to the distri- bution over the reward. Our solution to ensure then that they behave as intended is by constraining do and Qo to be consistent with each other, specifically that the implied reward the policy is sufficiently likely under the variational posterior (equivalently that the negative log-likelihood is sufficiently low). Thus we arrive at a constrained optimisation objective given by:
Σ log exp(3Qo(s,a))
max 0,0
ELEA exp(3Q0 b)) (FO(()(())))(R)).
(10)
(MO)ED
subject to E,,,T(- log go Qo(s,a) 7Qo('', a())] <e.
with reflecting the strength of the constraint. Rewriting (TO) as Lagrangian under the KKT
conditions (Karush 1939; Kuhn & Tucker 1951) and given complimentary slackness, we obtain a practical objective function:
exp xxp33Q((s,a))
F(6,0,D)
Σ (s,a,s',a')en
log
DEL D (P) (p(R(s.)))
bea exp(3Qo(s,b))
+ A log (Qo a) Qoo(s'
(11)
Here the KL divergence between processes is approximated over a countable set, and i introduced to control the strength of constraint.
On the implementation. Optimisation simple as both networks are maximising the same ob- jective and gradients can be easily obtained through backpropagation while being amenable mini- batching, allowing you to call your favourite gradient- based stochastic optimisation scheme. We re-iterate though that AVRIL really represents framework for doing BIRI. and not specific model
5

---

Published as a conference paper at ICLR 2021
Algorithm 1: Approximate Variational Reward Imitation Learning (AVRIL)
Result: Parameters & variational distribution and 0 of policy Q-function
Input: D, .S, A, 7, X, learning rate mini-batch size b; Initialise Φ; 0:
D Can concaterate into single vector (0,
while not converged do Sample D Minns from D
F( 8 D) E i F(o, 8, D ;
MC estimate total loss b Gradient step for 0,0
(0' 0') (6, 0)
+ 1/ 4.0 Ff F(, 0, D);
0,0 + 0,
end
Return: o,
since Qo and do represent arbitrary function approximators. So far we have presented both neural networks, but this does not have to be the case. Of course the advantage them is their flexibility and ease of training but they are still inherently black box. is then perfectly possible to swap in any particular function approximator if the task requires it, using simple linear models for example may slightly hurt performance but allow for more insight. Despite the specific focus infinite state-spaces, AVRIL can still even be applied in the tabular setting by simply representing the policy and variational distribution with multi-dimensional tensors. Having settled on their forms, equation is calculated simply and the joint gradient with respect to 0 and Φ is foraiglt returned using any standard auto-diff package. The whole process is summarised in Algorithm [
We can now see how AVRIL does not suffer the issues outlined in section 2.1 Our form of go ()) is flexible and easily accommodates a non-linear form of the reward given a neural architecture this also removes any restriction on S, or at least allows for any state space that is commonly tackled within the IL/RL literature. Additionally we have a single objective for which all parameters are maximised simultaneously there are no inner-loops, costly otherwise, meaning training is faster than the MCMC methods by a factor equal roughly to the number of samples they would require.
The generative model view. Ultimately a policy represents a generative model for the behavioural data we see. Нo Ermon (2016) explicitly make use of this fact by casting the problem in the GAN framework (Goodfellow et al.) 2014).. Our method is more analogous to VAE (Kingma & Welling, 2013), though not ex- actly, where given the graphical model in figure 2 the reward can be seen a latent representation of the policy. Our approach takes the seen data and amortises the inference, encoding over the state space. The policy does not act a decoder in precisely taking any given encoded reward and outputting policy, but it does take the whole reward posterior and translate it into actions and therefore behaviour. This approach has its advantages, in both meaningful interpretation of the latent reward (which non-existent adversarial methods), and that we forgo the practical difficulties of alter- nating min-max optimisation (Kodali et al. 2017) while maintaining generative view of the policy.
B
ai
R
Si
N
μ
Figure/ 2:
Graphical model for Bayesian IRL
Temporal consistency through reward regularisation. Considering only the first term of 11] yields the standard behavioural cloning setup (where the logits output can be interpreted as the Q- values) as it removes the reward from the equation and just focuses on matching actions to states. AVRIL can then be seen as a policy-lcarning method regularised by the need for the implied reward to be consistent. Note that this does not induce any necessary bias since the logits normally contain an extra degree of freedom allowing them to arbitrarily shift by some scale factor. This factor is now explicitly constrained by giving the logits additional meaning in that they represent Q-values. This places great importance on the KL term, since every parameterisation of policy will have an associated implied reward, the KL regularises these to be not so far from the prior and preventing the reward from overfitting to the policy and becoming pointless. It also is able to double regularising term in a similar manor to previous reward- rewarddregularisation methods (Piot et al. 2014 Reddy et al. 2019) depending the chosen prior, encouraging the reward to be close to

---

Published as a conference paper at ICLR 2021
Proposition (Reward Regularisation) Assume that the constraint in (0) is satisfied in that
Egg [R(s, a)] = Ex.r[Qo(ss, ( d')], then given a standard normal prior pf (R) = (R;0,)
the KL divergence yields sparsity regulator the implied reward:
Creg
Σ (8,0,8',"))D
2 (Qo(s..) - ((olss.)-7Qo(s$.cl) + g(Var,, [R(s,a)))-
(12)
Proof. Appendix. D This follows immediately from the fact that the divergence evaluates as a)))lI((((ss
DKL (Ro(R(s, (
(ol-loog(aa,,, (R(s, a))) - 1+ Varg.
Eg [Rss,)) a))
This then allows AVRIL to inherit the benefit of these methods while also explicitly learning a reward that can be queried at any point. We are also allowed the choice of whether it is state-only or state-action. This has so far been arbitrary, but it is important to consider that state-only reward a necessary and sufficient condition for reward that is fully disentangled from the dynamics (Fu et al. 2018). Thus by learning such a reward and given the final term of (IT) that directly connects one-step rewards in terms of the policy, this forces the policy (not the reward) to account for the dynamics of the system ensuring temporal consistency in way that BC for example simply can't. t. Alternatively using a state-action reward means that inevitably some of the temporal information leaks out of the policy and into the reward ultimately to the detriment of the policy but potentially allowing for a more interpretable (or useful) form of reward depending on the task at hand.
4 EXPERIMENTS
Experimental setup. We are primarily concerned with the case of medical environments, which is exactly where the issue of learning without interaction is most crucial, you just cannot let a policy sample treatments for a patient to try to learn more about the dynamics. It is also where level of interpretability in what has been learnt is important, since the consequence of actions are poten- tially very impactful on human lives. As such we focus our evaluation on learning on real-life healthcare problem, with demonstrations taken from the Medical Information Mart for Intensive Care (MIMIC-II) dataset (Johnson et al. 2016). The data contains trajectories of patients in inten- sive care recording their condition and theraputic interventions at one day intervals. We evaluate the ability of the methods to learn a medical policy in both the two and four action setting specifically whether the patient should be placed on a ventilator, and the decision for ventilation in combination with antibiotic treatment. These represent the two most common, and important, clinical interven- tions recorded in the data. Without a recorded notion of reward, performance is measured with respect to action matching against a held out test set of demonstrations with cross- vooss-validation.
Alongside the healthcare data and for the purposes of demonstrating generalisability, additional results on standard environments of varying complexity in the RL literature, the stan- dard control problems of: CartPole, classic control environment aiming to swing up and balance a pendulum; Acrobot, which aims to maintain a sequence of joints above a given height; and Lu- narLander, guiding a landing module safe touchdown given sufficient demonstration data all benchmarks are very much capable of reaching demonstrator level performance,
provide
the moon surface. In these settings
we test the algorithms on their ability to handle sample complexity in the low data regime by testing their performance when given access to select number of trajectories which we adjust, replicating the setup in Parrett et al. (2020). With access to a simulation through the OpenAl gym (Brockman et al.) 2016), we measure performance by deploying the learnt policies live and calculating their average return over 300 episodes.
Benchmarks. We test our method (AVRIL) against a number of benchmarks from the offline IRL/IL setting: Deep Successor Feature Network (DSFN) (Lee et al.] 2019). an offline adapta- tion of max-margin IRL that generalises past the linear methods using a deep network with least- squares temporal-different Rewaad-regullriized Classification for Apprenticeship Learning (RCAL) (Piot et al. 2014), where an explicit regulariser on the sparsity of the implied reward is introduced in order to account for the dynamics information; ValueDICE (VDICE) (Kostrikov et al.] 2019). an adversarial imitation learning, adapted for the offline setting by removing the replay regularisation; Energy-based Distri- bution Matching (EDM) (Jarett et al. 2020).. the state-of-the-ee art in offline imitation learning; and finally the standard example of Behavioural Cloning (BC). To provide evidence that are indeed
learning, the only other method that produces both a reward and policy;
7

---

Published as a conference paper at ICLR 2021
54.
300
-200
1753
100
190
124
0571
-00
-009
Acrobot
CartPole
LunarLander
Figure 3:
Control environments performance. We plot the average returns received by the poli- cies when deployed live in the environment against the number of trajectories seen during training.
Table Healthcare performance. Comparison of methods the MIMIC-III dataset. Performance of the policy is evaluated on the quality of action matching against a held out test set of demonstra- tions. We report the accuracy (ACC), area under the receiving operator characteristic curve (AUC) and average precision score (APS).
Ventilator
Ventilator + Antibiotics
Metric
ACC
AUC
APS
ACC
AUC
APS
BC
0.873 = 0.007
0.916 + 0.002
0.904 + 0.003
0.700 = 0.009
0.864 + 0.003
0.665 =: 0.009
VDICE
0.879 0.002
0.915
0.002
0.904 0.003
0.710 = 0.005
0.863 = 0.002
0.675 = 0.004
RCAL
0.870 0.012
0.916 0.003
0.004
0.702 = 0.008
0.865
0.669 0.004
0.006
DSFN
0.869 + 0.005
0.905
+ 0.008
0.885
0.683 + 0.007
4.866+00000
0.670 0.004
EDM
0.882 + 0.011
0.920 + 0.002
0.909
0.716 + 0.008
0.873 + 0.002
0.682 + 0.004
A-RL
0.875 + 0.010
0.904 + 0.002
0.927 1 0.002
0.718 + 0.010
0.864 +: 0.002
0.665 1 0.005
AVRIL
0.891 +: 0.002 0.917 + 0.001 0.940 + 0.001 0.754 = 0.001 0.884 + 0.000 0.708 + 0.002
learning an appropriate reward we show an ablation of our method on the MIMIC data: we take the reward learnt by AVRIL and use it as the 'true' reward used to train Q-network offline to learn policy (A-RL). Note that we have not included previous BIRL. methods for the reasons explained in section 1. training network just once in these environments takes in the order of minutes and repeating this sequentially thousands times is just not practical. For aid in comparison all methods share the same network architecture of two hidden layers of 64 units with ELU activation functions and are trained using Adam (Kingma & Ba 2014) with learning rates individually tuned. Further details experimental setup and the implementation of benchmarks can be found in the appendix.
Evaluation. We see for all tasks AVRIL learns an appropriate policy that performs strongly across the board, being competitive in all cases and in places beating out all of the other benchmarks. The results for our healthcare example are given in table with AVRIL performin very strongly, having the highest accuracy and precision score in both tasks. The results for the control environments are shown in figure 3 AVRIL performs competitively and is easily capable of reaching demonstrator level performance in the samples given for these tasks, though not always as quickly as some of the dedicated offline IL methods.
Reward insight. Remember though that task performance is not exactly our goal. Rather the key aspect of AVRIL is the inference over the unseen reward in order gain information about the preferences of the agent that other black-box policy methods can't. In the previous experiments our reward encoder was neural network for maximum flexibility and we can see from the performance of A-RL we learn a representation of the reward that can be used to relearn in the environment very effectively, albeit not quite to the same standard of AVRIL. Note this also reflects an original motivation for AVRIL in that offpolicy RL on top of a learnt reward suffers. In figure 4 we explore how to gain more insight from the learnt reward using different parameterisations of the reward. The top graph shows how a leartt state-action reward changes as a function of blood-oxygen level for an otherwise healthy patient, and can be seen that as it drops below average the reward for ventilating the patient becomes much higher (note this is average for patients in the ICU, not across the general population). While this is intuitive we still have to query a neural network repeatedly over the state space to gain insight, the bottom graph of figure presents then a simpler but perhaps more useful
8

---

Published as a conference paper at ICLR 2021
Rewards Map Ground
Demonstration Dersty
Reward Meen Scaled
3) Reward 5D Staled
40
24
52
Figure 5: Gridworld example. Scaled heat-maps of: the ground truth reward; the relative state occupancy of the expert demonstrations; the reward posterior mean; and reward standard deviation.
representation. In this case we learn state-only reward as before but as a linear model. This is not as strong a constraint on the policy since that is still free to be non-linear as a neural network but simultameously allows us the insight of what our model considers high value in the environment as we plot the relative model coefficients for each covariate. We can see here for example that the biggest impact on the overall estimated quality of a state is given by blood pressure, well known as an important indicator of health (Hepworth al. 1994), strongly impacted by trauma and infection.
Gridworld ground-truth comparison While environments like MIMIC are the main focus of this work they do not lend them selves to inspection of the uncovered reward as the ground truth simply is not available to us. We thus demonstrate on a toy gridworld environment, in order to clearly see the ef- fect of learning posterior distribution over the reward. In this (finite) example both the encoder and decoder are represented by tensors but otherwise the procedure remains the same. Fig- ploots scaled heat-maps of: a) the ground truth reward; b) the relative state occupancy of the expert demonstrations, ob- tained using value-iteration; c) the reward posterior mean; and d) the reward standard deviation. The interesting thing to note is that the standard deviation of the learnt reward essentially resembles the complement of the state occupancy revealing the epistemic uncertainty around that part of the state-space given we haven't seen any demonstrations there.
6.5
-6.0
631
632
5 CONCLUSIONS
Figure 4: (Top) A state-action ward is learnt and plotted for an otherwise average patient as their blood oxygen level changes. (Bot- tom) The associated weights given a state-only reward as a linear fuuccionn of the state-space.
We have presented a novel algorithm, Approximate Variational Reward Imitation Learning, for addressing the scalability is- sues that prevent current Bayesian IRL methods being used in large and unknown environments. We show that this performs strongly on real and toy data for learning imitation policies completely offline and importantly recovers a reward that is both effective for retraining policies but also offers useful insight into the preferences of the demonstrator. Of course this still represents an approximation, and there is room for further, more exact methods or else guarantees on the maxi- mum divergence. We have focused simply obtaining the appropriate uncertainty over reward as well as imitation in high stakes environments in these settings it is crucial that learnt policies avoid catastrophic failure and so how exactly the uncertainty in order achieve truly safe imitation (or indeed sineer-that-eemmnnee
apprenticeship) is increasingly of interest.
ACKNOWLEDGEMEETT
AJC would like to acknowledge and thank Microsoft Research for its support through its PhD Schol- arship Program with the EPSRC. This work was additionally supported by the Office of Naval Re- search (ONR) and the NSF (Grant number: 1722516). We would like thank all of the anonymous reviewers on OpenReview, alongside the many members of the van der Schaar lab, for their input, comments, and suggestions at various stages that have ultimately improved the manuscript.
9

---

Published as a conference paper at ICLR 2021
REFERENCES
Pieter Abbeel and Andrew Y Ng. Apprenticeship learning via inverse reinforcement learning. In Proceedings of the twenty-first international conference on Machine learning, pp. 2004.
Michael Bain and Claude Sammut. A framework for behavioural cloning. In Machine Intelligence 15. pp. 103-129, 1995.
David M Blei, Alp Kucukelbir, and Jon D McAuliffe. Variational inference: review for statisti- cians. Journal of the American statistical Association, (125518)859-8777 2017.
Michael Bloem and Nicholas Bambos. Infinite time horizon maximum causal entropy inverse rein- forcement learning. In 53rd IEEE Conference Decision and Control, pp. 4911-4916. IEEE, 2014.
Greg Brockman, Vicki Cheung, Ludwig Pettersson, Jonas Schneider, John Schulman, Jie Tang, and Wojciech Zaremba. Openai gym, 2016.
Daniel Brown, Scott Niekum, and Marek Petrik. Bayesian robust optimization for imitation learning. Advances in Neural Information Processing Systems, 33, 2020.
Daniel S Brown and Scott Niekum. Deep bayesian reward learning from preferences. arXiv preprint MXXiv:9112.047777 2019.
Jaedeug Choi and Kee-Eung Kim. Map inference for bayesian inverse reinforcement learning. In Advances in Neural Information Processing Systems, pp. 1989-1997, 2011.
Jaedeug Choi and Kee-Eung Kim. Nonparametric bayesian inverse reinforcement learning for mul- tiple reward functions. In Advances in Neural Information Processing Systems, pp. 305-313, 2012.
Christos Dimitrakakis and Constantin A Rothkopf. Bayesian multitask inverse reinforcement learn- ing. In European workshop on reinforcement learning, pp. 273-284. Springer, 2011.
Justin Fu, Katie Luo, and Sergey Levine. Learning robust rewards with adverserial inverse rein- forcement learning. In International Conference on Learning Representations, 2018.
Dani Gamerman and Hedibert F Lopes. Markov chain Monte Carlo: stochastic simulation for Bayesian inference. CRC Press, 2006.
Ian Goodfellow, Jean Pouget-Abadie, Mehdi Mirza, Bing Xu, David Warde-Farley, Sherjil Ozair, Aaron Courville, and Yoshua Bengio. Generative adversarial nets. In Advances in neural infor- mation processing systems, pp. 2672-2680, 2014.
Will Grathwohl, Kuan-Chieh Wang, Joern-Henrik Jacobsen, David Duvenaud, Mohammad Norouzi, and Kevin Swersky. Your classifier is secretly an energy based model and you should treat it like one. In International Conference on Learning Representations, 2019.
Joseph T Hepworth, Sherry Garrett Hendrickson, and Jean Lopez. Time series analysis of phys- iological response during icu visitation. Western journal of nursing research, 16(6):704-717, 1994.
Michael Herman, Tobias Gindele, Jörg Wagner, Felix Schmitt, and Wolfram Burgard. Inverse rein- forcement learning with simultaneous estimation of rewards and dynamics. In Artificial Intelli- gence and Statistics, pp. 102-110. PMLR, 2016.
Ashley Hill, Antonin Raffin, Maximilian Ernestus, Adam Gleave, Anssi Kanervisto, Rene Traore, Prafulla Dhariwal, Christopher Hesse, Oleg Klimov, Alex Nichol, Matthias Plappert, Alec Rad- ford, John Schulman, Szymon Sidor, and Yuhuai Wu. Stable baselines. https: / htg.tttt......... 2018.
com/ hom/hill-aa http:/taallaaaaaee
Jonathan Ho and Stefano Ermon. Generative adversarial imitation learning. In Advances in neural information processing systems, pp. 4565-4573, 2016.
10

---

Published as a conference paper at ICLR 2021
Daniel Jarrett, Ioana Bica, and Mihaela van der Schaar. Strictly batch imitation learning by energy- based distribution matching. Advances in Neural Information Processing Systems, 33, 2020.
Alistair EW Johnson, Tom J Pollard, Lu Shen, H Lehman Li-Wei, Mengling Feng, Mohammad Ghassemi, Benjamin Moody, Peter Szolovits, Leo Authony Celi, and Roger G Mark. Mimic-iii, freely accessible critical care database. Scientific data, 3(1):1-9, 2016.
William Karush. Minima of functions of several variables with inequalities as side constraints. Sc. Dissertation. Dept. of Mathematics, Univ. of Chicago, 1939.
Diederik P Kingma and Jimmy Ba. Adam: A method for stochastic optimization. arXiv preprint arXiv:1412.69880 2014.
Diederik P Kingma and Max Welling. Auto-encoding variational bayes. arXiv:1312.61 2013.
arXiv preprint
Edouard Klein, Matthieu Geist, and Olivier Pietquin. Batch, off-policy and model-free apprentice- ship learning. European Workshop on Reinforcement Learning, pp. 285-296. Springer, 2011.
Naveen Kodali, Jacob Abernethy, James Hays, and Zsolt Kira. On convergence and stability of gans. arXiv preprint arXiv: /005 770.07721, 2017.
Ilya Kostrikov, Kumar Krishna Agrawal, Debidatta Dwibedi, Sergey Levine, and Jonathan Tomp- son. Discriminator-actor-rritti Addressing sample inefficiency and reward bias in adversarial imitation learning. In International Conference on Learning Representations, 2018.
lya Kostrikov, Ofir Nachum, and Jonathan Tompson. Imitation learning via off-policy distribution matching. International Conference on Learning Representations (ICLR), 2019.
HW Kuhn and AW Tucker. Nonlinear programming. In Proceedings of the Second Berkeley Sym- posium on Mathematical Statistics and Probability. The Regents of the University of California, 1951.
Donghun Lee, Srivatsan Srinivasan, and Finale Doshi- Truly batch apprenticeship learning with deep successor features. International Joint Conference Artificial Intelligence (IJCAI), 2019.
Sergey Levine. Reinforcement learning and control probabilistic inference: Tutorial and review. arXiv preprint arXiv: /80500000, 2018
Sergey Levine. Zoran Popovic, and Vladlen Koltun. Nonlincar inverse reinforcement learning with gaussian processes. In Advances in Neural Information Processing Systems, pp. 19-27, 2011.
Sören Mindermann, Rohin Shah, Adam Gleave, and Dylan Hadfield-Menell. Active inverse reward design. arXiv preprint / 8809003000 2018.
Volodymyr Mnih, Koray Kavukeuoglu, David Silver, Alex Graves, Ioannis Antonoglou, Daan Wier- stra, and Martin Riedmiller. Playing atari with deep reinforcement learning. arXiv preprint 1rX.v.131..560,, 2013.
Andrew Y Ng. Stuart J Russell, et al. Algorithms for inverse reinforcement learning. In leml, volume 1, pp. 2, 2000.
Bilal Piot, Matthieu Geist, and Olivier Pietquin. Boosted and reward-regularized classification for apprenticeship learning. In Proceedings of the 2014 international conference on Autonomous agents and multi-agent systems, pp. 1249-1256. International Foundation for Autonomous Agents and Multiagent Systems, 2014.
Antonin Raffin, RI baselines zoo. httpp:///Sttuo.coo///aaaaaaaaaee 2018.
Deepak Ramachandran and Eyal Amir. Bayesian inverse reinforcement learning. In IJCAI, vol- ume 7, pp. 2586-2591 2007.
11

---

Published as a conference paper at ICLR 2021
Carl Edward Rasmussen. Gaussian processes in machine learning. In Summer School Machine Learning. pp. 63-71. Springer, 2003.
Siddharth Reddy, Anca D Dragan, and Sergey Levine. Sqil:
Imitation learning via reinforcement learning with sparse rewards. In International Conference on Learning Representations, 2019.
Constantin Rothkopf and Christos Dimitrakakis. Preference elicitation and inverse reinforce- ment learning. In Joint European conference machine learning and knowledge discovery in databases, pp. 48. Springer, 2011.
John Schulman, Filip Wolski, Prafulla Dhariwal, Alec Radford, and Oleg Klimov. Proximal policy optimization algorithms. arXiv preprint arXiv: (700::0777777 2017.
Brian D Ziebart. Modeling Purposeful Adaptive Behavior with the Principle of Maximum Causal Entropy. PhD thesis, University of Washington, 2010.
Brian D Ziebart, Andrew L Maas, J Andrew Bagnell, and Anind K Dey. Maximum entropy inverse reinforcement learning. Aaai, volume 8, pp. 1433-1438. Chicago, II., USA, 2008
12

---

Published as a conference paper at ICLR 2021
A EXPERIMENTAL SETUP
Demonstrations produced by running pre-trained and hyperparmeter optimised agents taken from the RL Baselines Zoo (Raffin 2018) in OpenAl Stable Baselines (Hill at. 2018). For Acrobot and LunarLander these are DONs (Mnih al. 2013), while CartPole uses PPO2 (Schulman et al. (017). Trajectories were then sub-sampled for every 20th step in Acrobot and CartPole, and every 5th step in LunarLLande.
Expert Demonstrators.
Testing setup. For control environments algorithms were presented with (1,3,7,10,15) trajectories uniformly sampled from a pool of 1000 expert trajectories. Each algorithm was then trained until convergence and tested by performing 300 live roll-outs in the simulated environment and recording the average accumulated reward received in each episode. This whole process was then repeated 10 times, consequently with different initialisations and seen trajectories.
Implementations. All methods are neural network based and in experiments they share the same architecture of 2 hidden layers units each connected by exponential linear unit (ELU) activation functions.
Publicly available code was used in the implementations of a number of the benchmarks, specifi- cally:
VDICE Kostrikov et al. 2019
https://github. com www.:///wooo...aeaarrroooooeaaaa https////.ttiioooooooorrrrrrrrrr
master/value dice
DSFN (Lee et al. 2019):
thub. www.coo.o/c.....
ing
EDM (Jarrett et al.) 2020)
attps://githun com com/wgrathwwh.
Note that VDICE originally designed for continuous actions with Normal distribution output which we adapt for the experiments by replacing with a Gumbel-softmax.
B PROOFS
Proof of proposition 1. Assuming the constraint is satisfied, maximising the following objective:
exp (Q)(s,a) ΣLeA exp( Dxc((o(((s,,)))o((((ss.))) (13)
F( = Σ
log
Which is equivalent to minimising the negative value
F(0,0)= Σ
log
exp (Qa a))
+ Der (gx(
(14)
(s.a, EED EBEA exp(3Qo(s,b)
Creg
Lsc
with the first term Lec being the negative log likelihood of the data and the classic behavioural cloning objective. Now given a standard Gaussian prior then the KL divergence Gaussian with mean y and variance o² from the prior is given by { log(a?) o² 1 ++2) (Kingma & Welling
13

---

Published as a conference paper at ICLR 2021
2013
Then given our prior p(R(s. = 0,, the KL evaluates as:
Lreg
Σ
DKL
(15)
(s,a,8' AD)
Σ
PoggVar,... []]) + Varg, [R(s,a)] + Egg (R(s,
(16)
a')
Σ
(Egg [(s a))²) + ((Vrrg [R(s,a))
(17)
2
(s,),",
Σ
()
+ (Vargg
(18)
2
",a,8',a")22
Since by assumption Eg R(ss a)] = Em,T(QQ(s,a) a) - 7Qa(s' a')l with the expectation ap- proximated over samples in the data and considering a definition of the function g to be g( (V a))) =
[R(s a)]) + Var,, [R(s, a)] 1). O
14


Reviews:
[{"review_id": "ZnVOurZA4a", "reviewer": "AnonReviewer2", "summary": "", "questions": "", "limitations": "", "rating": 7, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "This paper presents a method, called Auto-encoded Variational Reward Imitation Learning (VARIL). The presented idea is to learn a probabilistic (Gaussian) vector representation of the implicit reward function, so that the policy (Q) function can be trained in scale. Authors demonstrated that the presented methods can outperform (recently proposed) existing methods over several benchmarks in the offline RL.\n\nStrong points\n- Literature in inverse/offline RL is well illustrated. Especially, limitations of existing methods (impractical for modern, complicated, and model-free task environments) are clearly explained.\n- The proposed framework/algorithm outperform existing state-of-the-art in several benchmarks.\n- Contributions and presented models are easy to follow. \n\nWeak points\n- The engineering details are not well explained. I agree with authors that an arbitrary function approximator can be used for encoder. However, descriptions as in Figure 1 are too abstract.\n- There are some typos (e.g., the the, figure ...)\n\nIt would be good to explain more the training procedure in the paper. As an example, it is not sure why the presented encoder is called as auto-encoding reward when the input is demonstration and the output is the policy. ", "extra_scores": {}, "extra_text": {}}, {"review_id": "WvOh_CcDrg", "reviewer": "AnonReviewer4", "summary": "", "questions": "", "limitations": "", "rating": 7, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "Sequential decision-making is investigated in this paper. A number of demonstrations (i.e. trajectories) is provided to the algorithm that has to lean a policy using only data. Any exploration using sampling from the environment is not possible. To allow for a better understanding of the domain, the algorithm also has to recover the reward that is optimized by the agent that provides demonstrations. In this paper, the authors propose two neural networks that are expected to resemble an autoencoder. One of the neural networks learns a policy that can recover the decisions made in the demonstrations, and the second neural network learns the reward of the policy that is learned by the first neural network. The empirical results show that the new method outperforms alternative methods with respect to the \"action matching\". The reward learned by the algorithm is not compared with any other approaches.\n\nThe overall quality of writing in this paper is at a high, professional level. The authors know related literature, and citations are relevant and sufficient. I noticed very few types that are below.\n\nThe approach proposed in the paper is very interesting, and this work could be significant because the medical application requires this kind of method. I am not in position to give a strong recommendation to accept this paper because I feel that the core algorithm is not explained sufficiently in the current version of the paper, and a few important experiments are missing. This position is justified by my major concerns below.\n\nMajor concerns:\n\nThe biggest problem that I have is that I cannot easily link Fig. 1 with Eq. (11) because the current explanations in the paper leave many details implied. In particular, Fig. 1, shows two neural networks, one for $q_\\phi$ and one for $Q_\\theta$, but the paper does not explain what are the inputs and what are the outputs in both networks, and which components of Eq. (11) are used to optimize every network. It seems that the learning objective in $Q_\\theta$ uses terms 1 and 3 in Eq. (11); this is possible when $\\phi$ are frozen at the time when $\\theta$ are updated. This must be a coordinate ascent algorithm then. So, the reward predicted by $q_\\phi$ is not the input to $Q_\\theta$; it is used only as a penalty term in $Q_\\theta$. This is confusing because Fig. 1 indicates that R is the input.\n\nNext, let's consider $q_\\phi$. It takes (s,a) tuples as input, and it predicts $(\\mu,\\sigma)$ that define the normal distribution over R for every (s,a). I guess that terms 2 and 3 in Eq. (11) are used as objectives. The KL divergence wants the reward to be close to the prior, and the 3rd term, wants the predictions made by $q_\\phi$ to be close to the real differences encoded by $Q_\\theta$. So, again, when $\\theta$s are fixed, the $\\phi$ can be updated. The parameters of one of the networks can be updated, when the parameters of the other network are frozen.\n\nThus, this summary of the two networks indicates that this procedure is not an autoencoder. There are two independent NNs here, that are trained using coordinate ascent. In particular, if $Q_\\theta$ is the decoder, then its error should be backpropagated to the second network, i.e., to $q_\\phi$, or at least the decoder should be decoding using the low-dimensional embeddings instead of receiving entire input. This is definitely not the case here. There are two networks that are trained independently, i.e. they don't share the error information during their own backpropagations. I think that these nuances of the algorithms should be made clear; it took me some time to identify these details.\n\nI would further claim that $q_\\phi$ is not even required to learn $Q_\\theta$ with good predictive accuracy. It would be sufficient to use $Q_\\theta$ with the first term in Eq. (11) and then penalize the differences in Q(s,a) that are in the second term without the use of $q_\\phi$ at all. $Q_\\theta$ would be to learn the same policy, I believe, using simpler regularization. The paper should show the result of such a simplification of the method presented in the paper. This simplification would still use reward to regularize $Q_\\theta$, but without a use of a probability distribution over rewards, i.e., without the use of a second NN.\n\nThe addition of $q_\\phi$ is sensible overall if one wants to have a (normal) distribution over rewards for every (s,a). But still, the impact of having $q_\\phi$ should be demonstrated empirically as I described in the paragraph above.\n\nThere is something wrong about this sentence: \"where given the graphical model in figure 2 the reward can be seen as a latent representation of the policy.\" In RL, the reward is myopic and instantaneous, and the Q-values have to learned to represent a policy. The Q-values take into account long-term consequences of actions, whereas R(s,a) does have this information in MDPs.\n\nThe goal of this paper is to learn reward, but on p. 7 the authors said that in their experiments the performance is measured with respect to action matching\", and one of the baselines mentioned on the same page can recover the reward function as well. I feel that the rewards extracted by those methods should be compared with the rewards learned by the new method. This is a missing exploration.\n\nMinor problems:\n\nIntroduction is of high quality, but this sentence: \"This redirection to obtain policies in a manor we shall refer to as apprenticeship learning (AL) introduces its own challenges, particularly in the offline setting.\" does not explain apprenticeship learning. Please provide a link with IL and IRL that precede this sentence.\n\nOne page 4, \"they still\" -> \"there still\"\n\nIn bibliography \"bayes\" -> \"Bayes\" or \"gans\" -> \"GANs\"\n\nSummary:\n\nThe authors learn the network $Q_\\theta$, and they use another network $q_\\phi$ to regularize it. It seems that the same empirical results could be achieved without the second network because the first network could be regularized directly. The experiments showing the results of the simpler method would be highly desirable. The link with the autoencoder is not clear because the two networks that are used in this paper are leaned separately and the original inputs are given to the decoder, i.e. the decoder does not decode using a low-dimensional representation. It is not clear why this approach is equated with an autoencoder. The decoder does not decode using the low-dimensional embeddings, and it uses the low dimensional embeddings for regularization only. The low dimensional embeddings are not even given as input to the decoder.\n\nThis is interesting work with a good potential, but the above reasons don't allow me to suggest acceptance of the current version of the paper.\n\n------\nAdded after the discussion period: Thank you for answering my questions and the lively online interaction. You explained a few things to me, and I could see that you understood my point about VAEs. That was a great outcome. I hope that that you will address the  points that we discussed, where showing an honest link with VAEs will be highly desirable for your future readers. \n", "extra_scores": {}, "extra_text": {}}, {"review_id": "GxZFSbg6059", "reviewer": "AnonReviewer1", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "The paper proposes a Bayesian inverse reinforcement learning algorithm that learns a distribution over reward functions from offline demonstrations. The key idea is to avoid solving an MDP in the inner loop by maximizing the likelihood of the demonstrations with respect to a learned Q function, and to fit a distribution over reward functions to minimize the squared Bellman error with respect to the learned Q function. Experiments on an action prediction task in the healthcare domain as well simulated locomotion and video game tasks suggest that the proposed method is competitive with prior methods for imitation learning.\n\nOverall, I think this paper explores an interesting and impactful idea. There are similarities to the inverse soft Q-learning method [1] for learning the demonstrator's internal dynamics model, which involves simultaneously fitting a Q function to maximize the likelihood of demonstrated actions and to minimize the squared Bellman error with respect to a parametric model of the dynamics. The key difference between ISQL and the proposed method is that the proposed method learns the demonstrator's reward function instead of their internal dynamics model, and that the proposed method learns a distribution over reward functions instead of a point estimate. \n\nWhile I think this paper is a promising initial step, there are a few issues:\n\nThe main issue is that the action prediction experiments show small improvements over the AUC of prior methods on the healthcare dataset, and that AVRIL performs slightly worse than prior methods on several of the control benchmark tasks. One explanation for this relatively poor performance on imitation learning benchmarks could be that standard deep Q-learning does not perform well in the offline RL setting [2], and that AVRIL might share some of these issues because it minimizes the squared Bellman error with respect to learned Q values.\n\nIn general, the evaluations could do a much better job of illustrating the benefits of Bayesian IRL vs. standard imitation learning or IRL methods. The current paper evaluates AVRIL's ability to predict actions or imitate demonstrations in the demonstrator's MDP. Instead, one could test AVRIL's ability to learn a reward function that transfers to a new environment with a different initial state distribution or different dynamics -- something that IRL methods can do much better than standard imitation methods. One could also evaluate AVRIL's ability to compute high-confidence bounds on imitation policy performance, as in the evaluations for other Bayesian IRL methods like B-REX [3]. Another possibility is to use AVRIL to learn a distribution over reward functions, then perform risk-averse imitation, as in the experiments for inverse reward design [4]. Similarly, one could potentially use AVRIL as a component in an active learning method for learning rewards, like active inverse reward design [5]. \n\nThe interpretability analysis on page 8 does not highlight the benefits of learning a distribution over reward functions. The confidence bounds in the top of Figure 4 don't vary substantially, and the bottom of the figure doesn't illustrate AVRIL's uncertainty over the learned weights. It would be helpful to see if, for example, AVRIL learned tight confidence bounds for state-action pairs that are close to those seen in the demonstrations, and wider bounds for out-of-distribution pairs. Furthermore, it would be nice to see if different samples of reward weights from the learned distribution yield distinct, but equally plausible, hypotheses to a human supervisor.\n\nTypos:\n - manor -> manner (paragraphs 2 and 3)\n - In the \"Apprenticeship through rewards\" paragraph on page 2, \"Herman et al. (2016)\" should use \\citep instead of \\citet\n - \"it's\" -> its (last line on page 5)\n\n1. https://arxiv.org/pdf/1805.08010.pdf\n2. https://arxiv.org/pdf/2005.01643.pdf \n3. https://arxiv.org/pdf/1912.04472.pdf\n4. https://arxiv.org/pdf/1711.02827.pdf \n5. https://arxiv.org/pdf/1809.03060.pdf \n\nUpdate after rebuttal\n-----\nThank you to the authors for addressing my concerns. I have updated my score.\n", "extra_scores": {}, "extra_text": {}}, {"review_id": "BtLLVpAz3HK", "reviewer": "AnonReviewer3", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "The authors propose a Bayesian IRL algorithm based on variational inference by learning a policy similar to the one by the demonstrator and inferring the reward function at the same time without sampling. This is achieved by maximizing the ELBO of the posterior of rewards given the demonstrations assuming a softmax probability over Q-values. The contribution is to parameterize the policy and the reward function separately and connecting them through a constraint, which can be reformulated to lead to a single objective.\n\nThe paper is well written. Generally, I was particularly impressed by the authors’ explanations of related literature in sections 2 and 3. With one exception. The authors should consider citing Rothkopf & Dimitrakakis, Preference elicitation and inverse reinforcement learning, (2011), particularly, if they cite Ramachandran & Amir, (2007) as well as citing Dimitrakakis & Rothkopf, Bayesian multitask inverse reinforcement learning, (2011) if they cite the hierarchical Bayesian extensions Choi & Kim (2011) and Choi & Kim (2012).\n\nThe evaluations are not totally convincing, though. For the final version of the paper, I would encourage the authors to provide simulations on standard IRL problems used in the literature and quantifying performance e.g. with the cumulative loss of the Q-values wrt. the optimal or demonstrator policies. Also, the current manuscript does not provide probabilistic evaluations of the inferred policies or reward functions nor an assessment of how well the variational approximation with GP prior fairs in reward function inference compared to previous algorithms.\n\nVery little details are provided about the implementations, which is a weak point. \n\nOverall, it would be very helpful if the authors can provide some more details about the implementation and provide evaluations that quantify the inferred reward functions in the rebuttal. With the additional information, this could be a very strong paper that advances BIRL.\n\nThe authors are encouraged to release the code to foster open science and reproducibility.  Only the publicly accessible repositories used in this research are mentioned but not whether the authors intend to release their code upon acceptance of their manuscript.\n\nTypo: page 3: “offline setting setting”.", "extra_scores": {}, "extra_text": {}}]

Author rebuttal:

Meta-review: The reviewers agree that the submitted paper is of high quality and provides a promising approach/framework for Bayesian IRL. Certain concerns regarding details of the implementation and evaluation have already been addressed by the authors during the rebuttal phase, and also the title of the paper was adjusted in line with discussions with the reviewers. For the final paper, the authors should make sure to clearly highlight the advances of inferring a distribution over rewards (this is already partly done by the added grid world experiments) and discuss relations to VAEs as the initially had in mind and even in the paper title. Beyond that, the should of course also address other reviewers’ comments.
Decision: Accept (Poster)