# A Bayesian Solution To The Imitation Gap

Risto Vuorio
University of Oxford
<EMAIL_ADDRESS>&Mattie Fellows∗
University of Oxford
&Cong Lu∗
University of British Columbia
Vector Institute
&Clémence Grislain
Sorbonne University &Shimon Whiteson
University of Oxford
Equal contributionWork done while at University of OxfordWork done while
interning at University of Oxford

###### Abstract

In many real-world settings, an agent must learn to act in environments where
no reward signal can be specified, but a set of expert demonstrations is
available. Imitation learning (IL) is a popular framework for learning
policies from such demonstrations. However, in some cases, differences in
observability between the expert and the agent can give rise to an _imitation
gap_ such that the expert’s policy is not optimal for the agent and a naive
application of IL can fail catastrophically. In particular, if the expert
observes the Markov state and the agent does not, then the expert will not
demonstrate the information-gathering behavior needed by the agent but not the
expert. In this paper, we propose a Bayesian solution to the Imitation Gap
(BIG), first using the expert demonstrations, together with a prior specifying
the cost of exploratory behavior that is not demonstrated, to infer a
posterior over rewards with Bayesian inverse reinforcement learning (IRL). BIG
then uses the reward posterior to learn a Bayes-optimal policy. Our
experiments show that BIG, unlike IL, allows the agent to explore at test time
when presented with an imitation gap, whilst still learning to behave
optimally using expert demonstrations when no such gap exists.

##  1 Introduction

Imitation learning [27, 16] is a powerful method for training policies when
expert demonstration data is available for the desired behavior, without the
need for explicit reward. However, standard imitation learning algorithms can
fail when the expert demonstrator has access to privileged information that
the imitator lacks. Specifically, if the expert observes the full Markov
state, but the imitator operates under partial observability, then imitating
expert behavior can lead to suboptimal performance [7, 35, 6]. This mismatch
in observability is called the _imitation gap_.

For example, consider training a fruit-picking robot by learning from human
demonstrations. Humans use a combination of visual cues and touch to quickly
determine the ripeness of the fruit. They may reach out for a fruit that looks
ripe, but then leave it on the branch to ripen more if it feels too hard in
their hand. However, the robot has to rely on the visual cues alone since it
lacks the sense of touch. If it then picks all fruit that look ripe on one
side, it may end up picking many underripe fruit. A more intelligent robot
would turn around fruit that looks borderline ripe to visually inspect them
from all sides without detaching them from the branch. This exploratory
behavior is never demonstrated, leading to an imitation gap, and hence naïvely
imitating expert demonstrations leads to suboptimal behavior.

At first glance, this problem seems intractable, since we cannot imitate
behavior that is not demonstrated. Indeed, many prior solutions to the
imitation gap often require privileged access to online reward information.
The key insight in this paper is that the remaining uncertainty may be
characterized in the form of a prior, specified over the cost of exploration
in unobserved states. This leads us to propose a fully Bayesian solution to
the Imitation Gap (BIG) which learns to behave Bayes-optimally at test-time,
although its value will naturally be a function of the agent’s uncertainty.

We demonstrate how the prior can integrate several sources of information
available before test time: we specify an initial reward prior, from which we
may infer a posterior given a set of expert demonstrations and simulator,
using Bayesian inverse reinforcement learning [26, BIRL] with successor
features [3, 5, 12]. We also specify a reward prior over key exploration
states to specify uncertainty about the _cost of exploration_ where there is
an imitation gap. This allows BIG to optimally trade off any remaining
uncertainty about the true environment state using a Bayes-optimal policy,
with the predictive reward under these priors as the fixed belief over rewards
at test time. When there is an imitation gap, the reward prior can influence
the agent’s behavior instead, yielding policies that _balance exploration to
reduce the uncertainty in the environment with exploitation of expert
demonstrations_. When there is no imitation gap, the agent can directly
imitate expert behavior.

We evaluate BIG across a number of standard imitation gap problems, and
further show that it scales to environments with high-dimensional
observations. In each case, we can recover suitable reward functions from
_only expert demonstrations and the cost of exploration prior_.

##  2 Preliminaries

Reinforcement Learning. We model the environment as a Markov Decision Process
[29, MDP], defined as a tuple
$\mathcal{M}\coloneqq\langle\mathcal{S},\mathcal{A},p(s_{t+1}|s_{t},a_{t}),p(s_{0}),r(s_{t},a_{t}),\gamma\rangle$,
where $\mathcal{S}$ and $\mathcal{A}$ denote the state and action spaces
respectively, $p(s_{t+1}|s_{t},a_{t})$ the transition dynamics,
$r(s_{t},a_{t})$ the reward function, and $p(s_{0})$ the initial state
distribution. We denote a sampled reward $r$. The goal in reinforcement
learning is to optimize a policy $\pi(a|s)$ that maximizes the expected return
$\mathbb{E}_{\pi,p}\left[\sum_{t=0}^{\infty}\gamma^{t}r(s_{t},a_{t})\right]$.

Learning from Demonstrations. We assume access to a set of expert
demonstrations
$\mathcal{D}_{\textrm{Expert}}\coloneqq\\{\tau_{i}\\}_{i=1}^{N_{\textrm{Expert}}}$
of state-action trajectories
$\tau_{i}\coloneqq\\{s_{0},a_{0},s_{1},a_{1},\dots\\}$. In inverse
reinforcement learning (IRL), the objective is to use these demonstrations to
learn the reward function that the expert maximizes. In typical IRL methods,
such as maximum entropy IRL [36], the reward function is learned alongside a
policy that optimizes that function in a bi-level optimization algorithm.

By contrast, imitation learning is a closely related approach that seeks to
directly mimic expert behavior from the demonstrations. This can be expressed
as matching the distribution of the state-action pairs generated by the
imitator to that of the expert [15, 13], i.e., minimizing the divergence
between the limiting distribution over state-action pairs
$D(\rho^{\ast}(s,a)\|\rho_{\pi}(s,a))$ where $\rho^{\ast}$ and $\rho_{\pi}$
are the state-action marginal distributions of the expert and the imitator
policies respectively.

Successor Features. Successor features [3, SFs] are a value function
representation that decouples the dynamics of the environment from the reward.
Suppose that the reward function of the environment could be computed linearly
as $r(s,a)=\nu(s,a)^{\top}\omega$ where $\nu(s,a)\in\mathbb{R}^{d}$ are
features and $\omega\in\mathbb{R}^{d}$ are weights. For any given policy, we
may then factor the $Q$-function as
$Q^{\pi}(s,a,\omega)=\Psi^{\pi}(s,a)^{\top}\omega$ where
$\Psi^{\pi}(s,a):=\mathbb{E}_{\tau\sim
p^{\pi}}[\sum_{t=0}^{\infty}\gamma^{t}\nu(s_{t},a_{t})|s_{0}=s,a_{0}=a]$ is
the successor feature of $(s,a)$ under $\pi$.

##  3 Problem Setting

We now formalize the imitation gap in terms of a contextual MDP (CMDP) where
the demonstrator, but not the imitator, observes a hidden parameter
$\theta\in\Theta$ that affects the environment dynamics. Whilst the imitation
gap could also be formalized in terms of partial observability in the state,
we choose the CMDP formulation to make explicit what is hidden from the
imitator and assume the states are fully observable. Formally, we define a
CMDP as:

$\displaystyle\small\mathcal{M}(\theta)\coloneqq\langle\mathcal{S},\mathcal{A},p(s_{t+1}|s_{t},a_{t},\theta),p(s_{0}),r(s_{t},a_{t}),\gamma\rangle,$
(1)

with an underlying distribution over contexts $p(\theta)$. The reward function
is independent of $\theta$ as we assume all tasks in the CMDP have a common
goal (for example, driving safely at a junction), but differ in their state
transition dynamics. We are interested in policies that optimize the
contextual expected, discounted return
$\mathbb{E}_{p^{\pi}_{\infty}(\theta)}\left[\sum_{t=0}^{\infty}\gamma^{t}r(s_{t},a_{t})\right]$
where $p^{\pi}_{\infty}(\theta)$ is the distribution over infinite-horizon
trajectories associated with policy $\pi$ and context $\theta$.

At test time, we are interested in behaving optimally in a CMDP
$\mathcal{M}(\theta_{\textrm{test}})$ allocated according to the prior
$\theta_{\textrm{test}}\sim p(\theta)$. Like other successor feature-based
approaches [3, 5, 17, 12] we make the following assumption about the reward
parametrization:

######  Assumption 3.1.

The underlying reward function is bounded and can be represented as a linear
function with respect to a reward feature vector $\nu(s,a)$ that is known a
priori such that
$r(s,a;\omega^{\star})=\nu(s,a)^{\top}\omega^{\star}\in[r_{\textrm{min}},r_{\textrm{max}}]$.
Furthermore, the reward function is shared across all CMDPs, i.e., independent
of $\theta$.

This assumption simplifies our analysis and enables the derivation of convex
optimization procedures. We do not assume oracle access to perfect reward
features. Instead, the features could be the result of an unsupervised
learning step external to our algorithm. Crucially, the agent can observe
states and choose actions according to a policy but does _not_ observe rewards
nor $\theta_{\textrm{test}}$ and does not know the true reward parametrization
$\omega^{\star}$. Instead, the agent is given a dataset of
$N_{\textrm{Expert}}$ demonstrations
$\mathcal{D}_{\textrm{Expert}}\coloneqq\\{\tau_{i}\\}_{i=1}^{N_{\textrm{Expert}}}$
of state-action trajectories
$\tau_{i}\coloneqq\\{s_{0},a_{0},s_{1},a_{1},\dots\\}$ of length $H_{i}$ in
CMDPs sampled from $p(\theta)$, where each expert
$\pi^{\star}_{\textrm{Expert}}(\cdot,\theta_{i})$ behaves optimally in its
assigned CMDP, e.g., expert trajectories of a car turning at a junction in
summer and winter. Furthermore, the agent has access to a simulator, where
CMDPs are allocated according to the prior $p(\theta)$, and the agent can
interact with the corresponding environment, observing a trajectory of state-
action pairs. We denote the complete dataset of simulated trajectories as
$\mathcal{D}_{\textrm{Simulator}}\coloneqq\\{\tau_{i}\\}_{i=1}^{N_{\textrm{Simulator}}}$.
The agent never directly observes rewards nor $\theta_{i}$ in either
$\mathcal{D}_{\textrm{Expert}}$ or $\mathcal{D}_{\textrm{Simulator}}$.

At test time, the agent is assigned a CMDP according to
$\theta_{\text{test}}\sim p(\theta)$ and interacts with
$\mathcal{M}(\theta_{\text{test}})$, obtaining a history of state-actions:
$h_{t}\coloneqq\\{s_{0},a_{0},\cdots a_{t-1},s_{t}\\}$ at time $t$ and takes
actions according to a history-conditioned policy: $s_{t}\sim\pi(h_{t})$. The
agent never observes the reward history.

###  3.1 The Tiger-Treasure Problem

Figure 1:  A diagram of the Tiger-Treasure Problem MDP, a classic example of
an imitation gap. The agent initially does not know which door the treasure or
tiger is behind and must take listening actions to resolve its uncertainty.

We introduce a variant of the classic “Tiger-Treasure Problem” from Kaelbling
et al. [18] to illustrate how naïvely applying imitation learning fails when
there is an imitation gap. Consider the CMDP in Figure 1 indexed by
$\theta\in\\{1,2\\}$, representing which door the tiger is behind. The prior
is $p(\theta=1)=p(\theta=2)=0.5$. In both CMDPs, the agent starts in state
$S_{0}$ and the set of actions available is
$\mathcal{A}=\\{o_{1},o_{2},\textrm{listen}\\}$, where $o_{1}$ and $o_{2}$
open the corresponding door and ‘listen’ listens for a tiger. In any state
$s\in\\{S_{0},T_{1},T_{2}\\}$, the agent transitions deterministically to
state Tiger if $a=o_{1}$ and $\theta=1$ or $a=o_{2}$ and $\theta=2$, and
conversely for the Gold state. The goal of the agent is to reach the treasure
(labeled ‘Gold’) whilst avoiding the Tiger. The agent receives a reward
$r(\textrm{Gold},\cdot)=10$ for finding the gold and
$r(\textrm{Tiger},\cdot)=-100$ for finding the tiger. After this, the agent
transitions to terminal state $S_{T}$ regardless of action taken.

The agent can also listen before any doors are opened, receiving a stochastic
signal with success rate $p>0.5$ correlated with the identity of the door with
the tiger. Hearing the tiger in room $i$ transitions the agent to state
$T_{i}$. For $s\in\\{S_{0},T_{1},T_{2}\\}$ and $a=\textrm{listen}$, the agent
transitions to state $T_{1}$ with probability $p$ if $\theta=1$ and $T_{2}$
with $p$ if $\theta=2$, otherwise it transitions to state $T_{2}$ with
probability $1-p$ if $\theta=1$ and $T_{1}$ with $1-p$ if $\theta=2$. States
$T_{1}$ and $T_{2}$ are not present in Kaelbling et al. [18]’s original
problem, but are included here because reward is assumed independent of
$\theta$ and hence partial observability about the MDP is encoded in the
state. Entering a listening state incurs a small negative reward
$r(T_{1},\cdot)=r(T_{2},\cdot)=-1$. All other rewards not specified are $0$.

For ease of exposition, assume the listening success $p=1$. The expert has
privileged knowledge about the MDP, and always chooses to open the door with
the gold behind it. Consequently, expert demonstrations never feature
listening actions. At test time, there is an imitation gap, as the agent does
not know a priori which door the tiger is behind. A naïve imitator does not
realize that the expert is conditioning on extra information and so thinks the
expert is randomly choosing which door to open; imitating that gives
suboptimal return of $-45\gamma$. By contrast, an agent that chooses to listen
always receives a return of $10\gamma-1$ if it acts optimally on the revealed
location of the tiger. This example demonstrates the failure of naïve
imitation learning in simple settings when there is an imitation gap.

##  4 Towards a Bayesian Solution

In this paper, we propose a Bayesian solution to the Imitation Gap (BIG),
where the goal is to learn a policy that can optimally trade-off its
uncertainty at test time with imitating the expert demonstrator. Unlike in the
canonical Bayesian RL setting, our formulation does not have access to reward
samples from which to infer the underlying reward function. Instead, the
reward prior determines the cost of exploration (COE) for the agent at test
time, meaning that the agent can still behave optimally under partial
observability. We provide an overall schematic of our approach in Figure 2.

Figure 2:  Schematic of the Bayesian solution to the Imitation Gap (BIG).
Prior information is shown in green, algorithms in pink, prior distributions
in yellow, and outputs in blue.

BIG has three main phases which are labeled in Figure 2. In the first, we
provide an initial prior over the reward parametrization $p(\omega)$. We
integrate information from the expert data $\mathcal{D}_{\textrm{Exp}}$ into
the prior using a novel _contextual_ Bayesian IRL (BIRL) approach (Sections
4.1 and 4.2) that infers the posterior $p(\omega|\mathcal{D}_{\textrm{Exp}})$.
Because entire classes of reward functions can explain expert data equally
well, IRL is an underspecified problem and approaches can equally penalize any
unvisited state.

In the second phase (Section 4.3), we restrict the class of reward functions
to those that allow for exploration at test time, by first normalizing the
posterior so that the predictive reward lies in the range
$[r_{\textrm{min}},r_{\textrm{max}}]$ according to 3.1. Representing reward as
$r=k\times r_{\textrm{max}}$, we specify a cost of exploration prior $p(k)$
over rewards at key state-action pairs unseen in the expert demonstrations to
integrate the cost of exploration information at test time. This approach
ensures that all rewards still belong to the same _class_ [25, Definition 1]
after the IRL stage, i.e., yielding an equivalent optimal policy. This step is
required as the expert demonstrations do not have full coverage. $p(k)$
encodes the relative cost of deviating from an optimal policy which encourages
exploration for the downstream policy, complementing the IRL data. Integrating
both the IRL and COE priors, we denote the Bayesian reward distribution as
$p^{\textrm{IRL+COE}}_{\textrm{Bayes}}(r|s,a)$. We present pseudocode for an
algorithm implementing the first and second phase in Appendix D.

In the third phase, we solve a Bayesian RL problem in which the goal is to
learn a Bayes-optimal policy $\pi_{\textrm{Bayes}}$ that can be deployed at
test time. As shown in Figure 2, the inputs to the Bayes-adaptive MDP (BAMDP)
are the Bayesian reward distribution
$p^{\textrm{IRL+COE}}_{\textrm{Bayes}}(r|s,a)$ and a prior over context
variables $p(\theta)$ (Section 4.4). The agent extracts reward information
from the expert trajectories to learn a predictive reward but does not
directly imitate the expert’s behavior and thus can adapt to the unknown MDP
at test time, avoiding the problems with naïve imitation learning discussed in
Section 3.1. Due to the diverse nature of the sources of input information,
inferring the reward posterior requires the agent to learn and maintain
several distributions over random variables, we summarize them in Table 1 in
Appendix A.

###  4.1 Contextual Successor Features

Before performing BIRL, we must learn a contextual value function
representation to define the likelihood over expert trajectories. For an agent
following policy $\pi$, we can characterize the expected discounted return as
a function of state-action pairs via the contextual $Q$-function
$Q^{\pi}(s,a,\theta,\omega)$, which satisfies the contextual Bellman equation:
$\mathcal{B}^{\pi}[Q^{\pi}](s,a,\theta,\omega)=Q^{\pi}(s,a,\theta,\omega)$
where $\mathcal{B}^{\pi}[Q^{\pi}](s,a,\theta,\omega)$ is the contextual
Bellman operator:
$\mathcal{B}^{\pi}[Q^{\pi}](s,a,\theta,\omega)\coloneqq\nu(s,a)^{\top}\omega+\gamma\mathbb{E}_{s^{\prime}\sim
p(s^{\prime}|s,a,\theta)a^{\prime}\sim\pi(a^{\prime}|s^{\prime})}\left[Q^{\pi}(s^{\prime},a^{\prime},\theta,\omega)\right].$

3.1 specifies a linear reward function, which allows us to use a _successor
feature_ representation of the $Q$-function that factors the reward
parametrization out of
$Q^{\pi}(s^{\prime},a^{\prime},\theta,\omega)=\Psi^{\pi}(s^{\prime},a^{\prime},\theta)^{\top}\omega$
where $\Psi^{\pi}(s,a,\theta)$ is the contextual successor feature, defined
as: $\Psi^{\pi}(s,a,\theta)=\mathbb{E}_{\tau\sim
p^{\pi}(\tau_{\infty}|\theta)}\left[\sum_{t=0}^{\infty}\gamma^{t}\nu(s_{t},a_{t})\bigg{|}s_{0}=s,a_{0}=a\right]$.
N.b., the reward is linear w.r.t. the features $\nu(s,a)$, which could
themselves be learned and arbitrarily complex. Learning
$\Psi^{\pi}(s,a,\theta)$ means that we do not need to solve a Bellman equation
every time $\omega$ changes; we can simply take a dot product between the
existing successor feature and the new $\omega$. Appendix B details the
training process.

###  4.2 Contextual Bayesian IRL

We now describe the first phase of our pipeline in Figure 2. Inferring the
reward function is a Bayesian regression problem. As is typical for regression
problems [22], we specify a Gaussian reward model
$p(r|s,a,\omega)=\mathcal{N}(\nu(s,a)^{\top}\omega,I\sigma^{2})$ with mean
$\nu(s,a)$ and scalar variance parameter $\sigma\in\mathbb{R}$. We specify a
Gaussian prior over the unknown reward parameterization
$\mathcal{N}(\omega|\omega_{0},I\sigma_{0}^{2})$ where $\omega_{0}$ is the
prior mean and prior variance $\sigma_{0}^{2}>0$ represents the belief in
$\omega_{0}$. Tuning $\sigma_{0}^{2}$ thus allows us to set how much the
expert’s trajectories affect the prior reward specification after BIRL. We now
exploit expert data to refine our prior by leveraging approaches from BIRL.

To derive the likelihood, it is necessary to define a model of the expert
observations. In the classic Bayesian IRL approach [26], a likelihood is
specified for a single MDP: $p(\tau_{i}|\omega)$. To adapt our approach to
experts that act in multiple MDPs, our likelihood should account for the
context
$p(\tau_{i}|\theta_{i},\omega)=\prod_{j=0}^{T_{i}-1}p(s_{j},a_{j}|\theta_{i},\omega)$.
We assume that the agent’s policy is represented by a potential function
defined by the optimal $Q$-function for the agent’s MDP:

$\displaystyle
p(a|s,\theta_{i},\omega)=\frac{1}{z(s,\omega,\theta_{i})}\exp\left(\frac{1}{\alpha}\Psi(s,a,\theta_{i})^{\top}\omega\right),$
(2)

where $z(s,\omega,\theta)$ is the normalization constant. Here $\alpha$ is a
temperature parameter that controls how optimal the expert’s actions are with
respect to
$Q^{\star}(s,a,\theta_{i},\omega)=\Psi^{\star}(s,a,\theta_{i})^{\top}\omega$.
As is convention [26, 2], this assumption stabilizes inference algorithms by
softening the Dirac delta policy that the agent is following, allowing for
gradients to flow into the density. A deterministic optimal policy is
recovered in the limit $\alpha\rightarrow 0$.

Given the likelihood and prior, we infer the expert reward posterior
$p(\omega|\mathcal{D}_{\textrm{Expert}})$. Marginalizing, we obtain the
Bayesian reward distribution:

$\displaystyle
p^{\textrm{IRL}}_{\textrm{Bayes}}(r|s,a)\coloneqq\mathbb{E}_{\omega\sim
p(\omega|\mathcal{D}_{\textrm{Expert}})}\left[p(r|s,a,\omega)\right],$ (3)

which incorporates both the epistemic uncertainty from the posterior and the
aleatoric uncertainty from the reward model. Analogously to Bayesian logistic
regression [22], using the potential function from Equation 2 does not yield a
closed-form solution for the posterior. Instead, we use the Laplace
approximation for the posterior:

######  Theorem 4.1.

Define $\varsigma_{0}^{2}\coloneqq\frac{\sigma_{0}^{2}}{\alpha}$. Using the
Laplace approximation, the approximate posterior is
$p(\omega|\mathcal{D}_{\textrm{Expert}})\approx\mathcal{N}(\omega|\omega^{\star}_{\textrm{Laplace}},\Sigma_{\textrm{Laplace}})$
where $\Sigma_{\textrm{Laplace}}=\nabla^{2}_{\omega}\log
p(\omega^{\star}_{\textrm{Laplace}}|\mathcal{D}_{\textrm{Expert}})$ and
$\omega^{\star}_{\textrm{Laplace}}$ is the MAP estimate, which can be found by
carrying out the following stochastic gradient descent updates on the log-
posterior:

$\displaystyle\omega\leftarrow$
$\displaystyle\omega+\eta_{\omega}\Big{(}N_{\textrm{Expert}}{H_{i}}\bigg{(}\Psi^{\star}(s_{j},a_{j},\theta_{i})-\mathbb{E}_{a_{i}\sim
p(a_{i}|s_{j},\omega,\theta_{i})}\left[\Psi^{\star}(s_{j},a_{i},\theta_{i})\right]\Big{)}-\frac{(\omega-\omega_{0})}{\varsigma_{0}^{2}}\bigg{)}.$
(4)

###### Proof.

See Appendix F. ∎

The Bernstein von-Mises theorem formally justifies the approximation [8, 32],
proving that under mild regularity assumptions, the posterior tends to the
Laplace approximation in the limit of increasing data. $\varsigma_{0}^{2}$
controls how the prior influences learning; as $\varsigma_{0}^{2}\rightarrow
0$, expert data is ignored and $\omega^{\star}_{\textrm{Laplace}}=\omega_{0}$.
The posteriors over each expert’s contextual variable
$p(\theta_{i}|\omega,\tau_{i})$ typically have no closed-form analytic
solution. We can approximate each $p(\theta_{i}|\omega,\tau_{i})$ using
variational inference instead, as detailed in Appendix E.

Role of Temperature. As IRL is underspecified, many reward functions can
explain the expert data. A key insight from Theorem 4.1 is the role of the
temperature parameter $\alpha$ in determining the relative difference between
the lowest and highest rewards assigned. The example in Appendix G illustrates
that when $\alpha\rightarrow 0$, the expert policy model becomes more
deterministic, always choosing the action with the highest return. Arbitrarily
small differences between rewards explain expert behavior, so IRL pulls the
parametrization difference to be as small as possible. Conversely, when
$\alpha\rightarrow\infty$, the expert policy becomes more stochastic, taking
low-return actions more frequently in proportion to their value. An
increasingly large separation in rewards is needed to explain expert behavior.

###  4.3 Cost of Exploration Prior

For the second phase, we specify a prior over the reward to incorporate
information on the cost of exploration into the posterior reward returned by
the contextual BIRL to refine the solution. Note that the imitation gap
problem would not be solvable without access to prior information defining the
cost of exploration, as the expert data does not demonstrate which states are
safe to explore and how to balance exploration and exploitation. This prior is
only introduced for exploration and does not need to contain any information
about the exploitation, as that is learned from the expert demonstrations.
Furthermore, the prior could be arbitrarily uninformative, in the third phase,
our method learns the Bayes-optimal policy for any prior.

As shown in Figure 2, we start by rescaling the posterior reward to be within
the bounds $[r_{\textrm{min}},r_{\textrm{max}}]$ according to 3.1. We denote
the corresponding scaled reward parametrization as
$\bar{\omega}^{\star}_{\textrm{Bayes}}$. Given the infinite horizon, any
linear transformation applied to the reward function results in the same
optimal policy.

We assume that we know a subset of state-action pairs denoted as
$\mathcal{S}_{\text{COE}}\times\mathcal{A}_{\text{COE}}$ where exploration can
be performed. We illustrate this for the Tiger-Treasure problem from Figure 1;
as rewards only depend on the state in this problem, the set is
$\mathcal{S}_{\text{COE}}=\\{T_{1},T_{2}\\}$. We introduce a simple COE inside
the rescaled IRL reward parameterized by
$\bar{\omega}^{\star}_{\textrm{Bayes}}$ by specifying a reward function
$p(r|s,a,k)=\mathcal{N}(r|kr_{\textrm{max}},\sigma^{2})$ over
$(s,a)\in\mathcal{S}_{\text{COE}}\times\mathcal{A}_{\text{COE}}$ for some
$k\in[\frac{r_{\textrm{min}}}{r_{\textrm{max}}},1)$; this ensures the mean is
contained in $[r_{\textrm{min}},r_{\textrm{max}}]$. A scale $k$ that varies
across action-state pairs may also be specified if a more complex cost of
exploration information needs to be modeled. In the simple Tiger-Treasure
problem, we know from construction the set of states where exploration can be
performed. In a practical setting with large state space, these states could
be obtained by considering non-expert, but safe behaviors from other policies
acting in the same environment. This could require density estimation for
continuous state spaces. We leave development of task specific cost of
exploration priors for future work and focus on demonstrating the general
principles in this work.

The value of $k$ determines how risk-averse the agent is at test time. For
$k\approx 1$ the agent values exploratory state-actions in
$\mathcal{S}_{\text{COE}}\times\mathcal{A}_{\text{COE}}$ nearly as much as the
most rewarding state-actions learned from IRL. As such, the cost of
exploration is low, and the agent explores until it is highly certain about
avoiding low reward actions in the imitation gap, encouraging conservative
behavior. Conversely, as
$k\rightarrow\frac{r_{\textrm{min}}}{r_{\textrm{max}}}$ the agent becomes less
risk-averse and recovers a purely behavioral cloning regime, taking actions
that could lead to low reward as they have similar value to exploratory
actions.

To incorporate epistemic uncertainty in $k$, we specify a prior $p(k)$ with
support over $[\frac{r_{\textrm{min}}}{r_{\textrm{max}}},1)$. Marginalizing,
we obtain the Bayesian reward distribution over
$\mathcal{S}_{\text{COE}}\times\mathcal{A}_{\text{COE}}$:

$\displaystyle p^{\textrm{COE}}_{\textrm{Bayes}}(r|s,a)=\mathbb{E}_{k\sim
p(k)}\left[p(r|s,a,k)\right].$ (5)

For all other state-action pairs, the Bayesian reward distribution remains
unchanged, yielding the distribution over $\mathcal{S}\times\mathcal{A}$:

$\displaystyle p_{\textrm{Bayes}}^{\text{IRL +
COE}}(r|s,a)=\begin{cases}p^{\textrm{COE}}_{\textrm{Bayes}}(r|s,a)&s,a\in\mathcal{S}_{\text{COE}}\times\mathcal{A}_{\text{COE}},\\\
p^{\textrm{IRL}}_{\textrm{Bayes}}(r|s,a)&\text{otherwise.}\end{cases}$ (6)

###  4.4 Bayes-Optimal Policy Learning

For the third phase, we perform Bayesian reinforcement learning, which
optimally trades off exploration and exploitation by conditioning actions on
the agent’s uncertainty over $\theta$. We can define a Bayes-adaptive MDP [9,
BAMDP] using the contextual MDP in Equation 1 as a model. At test time, the
agent is assigned an MDP $\theta_{\textrm{test}}\sim
p(\theta_{\textrm{test}})$ and can observe samples from
$p(s^{\prime}|s,a,\theta_{\textrm{test}})$ by interacting with the MDP via its
policy. A history of interactions is denoted
$h_{t}\coloneqq\\{s_{0},a_{0},s_{1},a_{1},\ldots s_{t}\\}\in\mathcal{H}_{t}$
where $\mathcal{H}_{t}$ is the corresponding state-action product space. After
observing a history $h_{t}$ from the assigned MDP, the agent updates its
belief in $\theta_{\textrm{test}}$ according to the posterior:

$\displaystyle
p(\theta_{\textrm{test}}|h_{t})=\frac{\prod_{i=1}^{t}p(s_{t}|s_{t-1},a_{t-1},\theta_{\textrm{test}})p(\theta_{\textrm{test}})}{\mathbb{E}_{\theta_{\textrm{test}}\sim
p(\theta_{\textrm{test}})}\left[\prod_{i=1}^{t}p(s_{t}|s_{t-1},a_{t-1},\theta_{\textrm{test}})\right]}.$
(7)

Using the posterior, we can define the Bayesian transition distribution as:

$\displaystyle
p(s_{t+1}|h_{t},a_{t})=\int_{\Theta}p(s_{t+1}|s_{t},a_{t},\theta_{\textrm{test}})p(\theta_{\textrm{test}}|h_{t})d\theta_{\textrm{test}}.$
(8)

As there is no reward signal available to the agent and rewards do not depend
on $\theta$, the Bayesian reward distribution in the BAMDP is exactly the
combined Bayesian reward distribution from Equation 6:
$p(r_{t}|h_{t},a_{t})=p_{\textrm{Bayes}}^{\text{IRL +
COE}}(r_{t}|s_{t},a_{t})$. We denote the joint reward-state Bayesian
transition distribution as
$p(r_{t},s_{t+1}|s_{t},a_{t})=p(r_{t}|h_{t},a_{t})p(s_{t+1}|h_{t},a_{t})$,
which is equivalent to the predictive trajectory transition distribution:
$p(\tau_{t+1}|\tau_{t},a_{t})=p(\tau_{t},a_{t},r_{t},s_{t+1}|\tau_{t},a_{t})=p(r_{t},s_{t+1}|h_{t},a_{t})\underbrace{p(\tau_{t},a_{t}|\tau_{t},a_{t})}_{=1}=p(r_{t},s_{t+1}|h_{t},a_{t})$.
Here $p(\tau_{t+1}|\tau_{t},a_{t})$ is used to reason over unobserved
counterfactual trajectories, and so must account for predictive reward. We
define the corresponding BAMDP similarly to Fellows et al. [10] as
$\mathcal{M}_{\textrm{BAMDP}}\coloneqq\langle\mathcal{T},\mathcal{A},p(\tau_{t+1}|\tau_{t},a_{t}),p(s_{0}),\gamma\rangle$
where $\mathcal{T}$ is the space of all possible trajectories.

In Bayesian RL, policies $\pi_{\textrm{Bayes}}(a_{t}|h_{t})$ map from
_histories_ to distributions over actions, and our goal is to learn a Bayes-
optimal policy $\pi_{\textrm{Bayes}}^{\star}(a_{t}|h_{t})$ that solves
$\mathcal{M}_{\textrm{BAMDP}}$. Due to the linearity of our formulation, we
show in Section F.1 that solving the BAMDP is equivalent to optimizing the
following Bayesian RL objective for $\pi^{\textrm{Bayes}}$:

$\displaystyle
J^{\pi}_{\textrm{Bayes}}=\mathbb{E}_{p(\theta_{\textrm{test}})}\left[\mathbb{E}_{p(h_{\infty}|\theta_{\textrm{test}})}\left[\sum_{i=0}^{\infty}\gamma^{i}r_{\textrm{Bayes}}^{\text{IRL
+ COE}}(s_{i},a_{i})\right]\right],$ (9)

where
$p(h_{\infty}|\theta_{\textrm{test}})=p_{0}(s_{0})\prod_{i=0}^{\infty}p(s_{i+1}|s_{i},a_{i},\theta_{\textrm{test}})\pi(a_{i}|h_{i})$,
and $r_{\textrm{Bayes}}^{\text{IRL + COE}}(s,a)$ is the predictive reward:

$\displaystyle r_{\textrm{Bayes}}^{\textrm{IRL +
COE}}(s,a)\coloneqq\begin{cases}r_{\textrm{max}}\mathbb{E}_{k\sim
p(k)}\left[k\right]&s,a\in\mathcal{S}_{\text{COE}}\times\mathcal{A}_{\text{COE}},\\\
\nu^{\top}(s,a)\bar{\omega}^{\star}_{\textrm{Bayes}}&\text{otherwise.}\end{cases}$
(10)

In the following empirical evaluation, we approximate Bayes-optimal policies
by training a policy conditioned on the true inference model using DQN [21].

##  5 Empirical Evaluation

To evaluate BIG, we conduct experiments across a series of imitation gap
problems. In all the tested environments, the agent has to solve a task that
requires exploration. First, we demonstrate that in the Tiger-Treasure
environment, naive IRL learns a reward function that does not lead to the
desired exploratory policy, whereas BIG does by exploiting prior information
about the cost of exploration. Second, we illustrate that, by marginalizing
over the context distribution, naive IRL can learn a reward function that does
not capture the expert’s intent. Finally, we present results in a gridworld
environment to show that BIG can handle imitation learning tasks with larger
state-action spaces. Since we assume no access to true environment reward or
expert state information, most previous solutions to the imitation gap are not
applicable. Instead, we compare to maximum entropy IRL (labeled ‘No-Prior’ in
our experiments). For convenience, we implement Algorithm 1 using the true
posterior $p(\theta|\tau)$ as the inference model. In all experiments, we use
deep neural networks for $\Psi$. We assume normally distributed errors and
report standard error across seeds in the figures.

###  5.1 Investigating Reward Priors

Figure 3: Evaluation of BIG in the Tiger-Treasure environment. Success rate
and time exploring (in steps) for policies learned with a uniform prior reward
with different means are represented in yellow ($k^{\star}<1$) and in green
($k^{\star}=1$), while the case with no prior is shown in red. Error bars
indicate the standard error of the mean across 10 seeds. The symbol $+\infty$
indicates that, for some trials, the agent explores throughout the entire
(infinite) episode.

In the Tiger-Treasure environment introduced in Section 3.1, since the expert
always goes to the treasure room directly, we cannot extract information about
optimal exploration from the expert data. As discussed in Section 4.3, we use
the prior $p(k)$ to enable exploratory behavior at test time. We explore the
influence of this reward prior on the environment from Figure 1, using a space
of uniform priors $p(k)=\textrm{Unif}\left([a,b]\right)$ over intervals
$[a,b]\in\left[\frac{r_{\textrm{min}}}{r_{\textrm{max}}},1\right)$. We choose
$r_{\text{min}}=-100$ and $r_{\text{max}}=10$. While these bounds are
arbitrary, they reflect the undesirability of failing in the task demonstrated
by the expert. The reward feature $\nu$ is a one-hot indicator over the
states. Figure 3 presents the agent’s success rate in reaching the treasure,
along with the average time exploring, which corresponds to the number of
listening actions, for different values of the prior mean $k^{\star}$. As
expected from Section 4.3, when $k^{\star}$ approaches 1, the agent explores
more, never exploiting when $k^{\star}=1$. The probability of reaching the
treasure increases with the number of listening actions. Therefore, as
$k^{\star}$ approaches 1 (without reaching it), the treasure rate also
increases. By contrast, when $k^{\star}$ decreases the agent begins to listen
less often. These behaviors correspond to the Bayes-optimal policies for each
of the priors, demonstrating that our method learns the Bayes-optimal policy
irrespective of the prior supplied.

###  5.2 Investigating Latent Inference

Figure 4:  A demonstration of the necessity for latent inference with BIG. On
the left, we show the CMDP used in the experiments, with two possibilities for
the context $\theta$. On the right, we show the ground truth returns of a DQN
agent for trajectories of 100 steps in the CMDP during training. The shading
shows the standard error of the mean for 8 seeds.

Next, we look at whether inferring the latent $\theta$ during IRL matters for
learning the desired reward function. We experiment in a simple CMDP
environment depicted in Figure 4. For a mathematical example, see Appendix C.
In this environment, the expert policy goes to the state $s_{3}$, which
provides a reward of $+2$ and then loops back to state $s_{0}$ through $s_{4}$
or $s_{5}$. It prefers to take the route through $s_{1}$, when it is
available, to avoid the negative reward of $-1$ in $s_{2}$. When $\theta$ is
distributed such that $p(\theta=0)>p(\theta=1)$, it can lead to naive IRL
misidentifying the expert intent. To see why, consider that to fit the expert
behavior without information about $\theta$, the reward function has to make
the path through $s_{2}$ more likely in both MDPs. Conditioning the successor
features on the inferred $\theta$ resolves this issue, because then
identifying the reward reduces to standard IRL in two separate MDPs. To verify
this in practice, we show results of learning the reward functions with and
without latent inference in Figure 4. We choose the latent to be distributed
as $p(\theta=0)=0.9$. No reward prior is used. The policy trained on the
rewards learned with latent inference matches the policy trained with the
ground truth rewards. The rewards learned without latent inference do not
result in as good a policy, demonstrating that latent inference is necessary
to learn the correct rewards in a general CMDP. See Section H.2 for an
analysis of the learned rewards.

###  5.3 Reward Priors in a Larger CMDP

Finally, we test BIG in a grid-world environment with pixel-based
observations. The agent observes a top-down view of the environment similar to
the illustration of the learned rewards presented in Figure 5. The agent can
move in four directions and take listening actions. Taking the listening
action in any grid cell results in a stochastic transition to a state, which
indicates the location of the gold, but otherwise has the same dynamics as the
cell that the action was taken in. When the agent enters a door, it is moved
to $\mathrm{Tiger}$ or $\mathrm{Gold}$ depending on $\theta$. From those
states, the agent is moved to the grid cell marked with $x$ in the next
timestep, setting the agent up for another round in the maze. The expert takes
the shortest path from any state to $\mathrm{Gold}$, choosing the correct
door, depending on $\theta$. See Section H.4 for details.

Three DQN training curves are shown in Figure 5 corresponding to a manually
constructed ground truth reward, which explains the expert behavior, reward
learned by the contextual IRL without using a reward prior, and the same
reward refined using the prior. These learning curves show that BIG with a
particular reward prior produces a similar BAMDP policy as the manually
constructed reward. At the same time, using just the IRL reward results in a
policy that initially chooses a door at random. We compare multiple values for
$k^{\star}$ and find that as in the simple Tiger-Treasure environment,
different values result in over-exploration or under-exploration leaving a
range in the middle that results in similar exploration as the handcrafted
reward. This experiment shows that BIG can recover a reward function, which
enables the agent to complete the task demonstrated by the expert despite a
challenging imitation gap.

Figure 5:  BIG successfully learns the optimal behavior in a challenging
gridworld environment. On the left, we show the rewards learned by the
contextual IRL. In the middle, we show the return (using the manually
constructed reward) of policies trained with reward inferred with and without
a reward prior and the manually constructed reward (ground truth). The shading
shows the standard error of the mean for 8 random seeds. On the right, we show
the final returns of policies trained using different values of $k^{\star}$
compared to not using the reward prior and using the ground truth reward.

##  6 Related Work

Imitation Gaps. Prior work addressing the imitation gap typically assumes
privileged access to the true environment reward during training. In contrast,
BIG makes no such assumption. Weihs et al. [35], Nguyen et al. [23] assume
access to the true reward during training, and propose to bridge the imitation
gap by training on a weighted imitation and RL loss. Other works propose to
also integrate privileged information about the expert during training [7] in
addition to the reward. For example, Elf Distillation [34] studies an approach
that mixes environment reward with online advice from the expert. Cai et al.
[6] propose several stages of training, including those on privileged expert
states, to connect the imitator and expert’s observation spaces. Separately,
versions of the imitation gap have been considered by Kwon et al. [20], Ortega
et al. [24], Swamy et al. [30], Vuorio et al. [33] that assume that no new
exploratory behavior needs to be learned. The setting considered by Straub et
al. [28] is closer to our work but it only considers inferring unknown
parameters from an agent acting under partial observability (that is, the IRL
problem). Our setting also requires an agent to act optimally in an _unknown_
environment _without_ a reward signal, i.e., there is uncertainty in the MDP
at test time.

Bayesian Reinforcement Learning. Our work shares similar components to other
Bayesian approaches to RL. For instance, VariBAD and related methods [37, 38,
39] consider a model-based approach for learning approximations to Bayes-
optimal policies by exploiting meta-reinforcement learning [4] to perform
inference over a subset of the unobserved context. Similarly, BEN [11] is a
model-free approach that learns Bayes-optimal policies by specifying a model
and prior over the optimal Bellman operator. All of these methods assume
access to the reward and do not consider the issue of imitating an expert.
Bayesian Inverse RL approaches infer a posterior over the unknown reward
distribution given trajectories of demonstrations [26] (see Adams et al. [1],
Table 1 for a complete list of existing approaches). BIG generalizes these
approaches in Section 4.2 to account for the unobserved context variable
before integrating the learned reward posterior into a BAMDP.

Successor Features. Successor features [3, 5] provide an elegant
representation of value functions under the assumptions of a linear reward
function. Janz et al. [17] incorporate successor features into the Bayesian
framework. Most similar to BIG, successor features have been used successfully
in Psi-Phi Learning [12] for multi-task inverse reinforcement learning. Psi-
Phi Learning can also retrieve the reward parameterizations for a new agent
(expert) but does so with full observability.

##  7 Limitations

To make the empirical setup close to the theory, we use ground truth inference
models and linearity assumptions, which means that scaling up the algorithm
requires revisiting those choices. For example, by learning an approximate
inference model. In the CMDPs considered in this paper, we assumed the set of
allowed exploration states was explicitly known a priori, which may be a
limiting assumption in harder problems. As is typical in Bayesian methods, we
leave the design of problem-specific priors to the practitioners focused on
those problems.

##  8 Conclusion

In this paper, we proposed a fully Bayesian solution to the imitation gap
which integrated various priors over the reward parameterization, cost of
exploration, and hidden state. This allowed us to derive a BAMDP formulation
of the problem whose solution optimally trades off exploration and
exploitation at test time. We demonstrated the importance of each component of
our algorithm across a series of experiments with an imitation gap before
showing that BIG scales to larger maze problems, including those with high-
dimensional pixel-based observations. Crucially, in contrast to previous work,
no online reward information was required. This makes our work particularly
exciting for challenging imitation problems where no reward is easily
specifiable, and extending BIG to even more complex settings is a promising
direction for future work.

## References

  * Adams et al. [2022] Stephen Adams, Tyler Cody, and Peter A. Beling.  A survey of inverse reinforcement learning.  _Artif. Intell. Rev._ , 55(6):4307–4346, aug 2022.  ISSN 0269-2821.  doi: 10.1007/s10462-021-10108-x.  URL https://doi.org/10.1007/s10462-021-10108-x.
  * Arora and Doshi [2021] Saurabh Arora and Prashant Doshi.  A survey of inverse reinforcement learning: Challenges, methods and progress.  _Artificial Intelligence_ , 297:103500, 2021.  ISSN 0004-3702.  doi: https://doi.org/10.1016/j.artint.2021.103500.  URL https://www.sciencedirect.com/science/article/pii/S0004370221000515.
  * Barreto et al. [2017] Andre Barreto, Will Dabney, Remi Munos, Jonathan J Hunt, Tom Schaul, Hado P van Hasselt, and David Silver.  Successor features for transfer in reinforcement learning.  In I. Guyon, U. Von Luxburg, S. Bengio, H. Wallach, R. Fergus, S. Vishwanathan, and R. Garnett, editors, _Advances in Neural Information Processing Systems_ , volume 30. Curran Associates, Inc., 2017.  URL https://proceedings.neurips.cc/paper_files/paper/2017/file/350db081a661525235354dd3e19b8c05-Paper.pdf.
  * Beck et al. [2023] Jacob Beck, Risto Vuorio, Evan Zheran Liu, Zheng Xiong, Luisa Zintgraf, Chelsea Finn, and Shimon Whiteson.  A survey of meta-reinforcement learning.  _arXiv preprint arXiv:2301.08028_ , 2023.
  * Brown and Niekum [2019] Daniel S. Brown and Scott Niekum.  Deep bayesian reward learning from preferences.  _Workshop on Safety and Robustness in Decision Making. 33rd Conference on Neural Information Processing Systems (NeurIPS 2019)_ , abs/1912.04472, 2019.  URL https://arxiv.org/pdf/1912.04472.pdf.
  * Cai et al. [2021] Xin-Qiang Cai, Yao-Xiang Ding, Zi-Xuan Chen, Yuan Jiang, Masashi Sugiyama, and Zhi-Hua Zhou.  Seeing differently, acting similarly: Heterogeneously observable imitation learning.  _arXiv preprint arXiv:2106.09256_ , 2021.
  * Chen et al. [2019] Dian Chen, Brady Zhou, Vladlen Koltun, and Philipp Krähenbühl.  Learning by cheating.  In _Conference on Robot Learning (CoRL)_ , 2019.
  * Doob [1949] J. L. Doob.  Application of the theory of martingales.  _Colloques Internationaux du Centre National de la Recherche Scientifique_ , pages 23–27, 1949.
  * Duff [2002] Michael O’Gordon Duff.  _Optimal Learning: Computational procedures for Bayes-adaptive Markov decision processes_.  University of Massachusetts Amherst, 2002.
  * Fellows et al. [2023a] Mattie Fellows, Brandon Kaplowitz, Christian Schroeder de Witt, and Shimon Whiteson.  Bayesian exploration networks.  _arXiv preprint arXiv:2308.13049_ , 2023a.
  * Fellows et al. [2023b] Mattie Fellows, Matthew J. A. Smith, and Shimon Whiteson.  Why target networks stabilise temporal difference methods.  In Andreas Krause, Emma Brunskill, Kyunghyun Cho, Barbara Engelhardt, Sivan Sabato, and Jonathan Scarlett, editors, _Proceedings of the 40th International Conference on Machine Learning_ , volume 202 of _Proceedings of Machine Learning Research_ , pages 9886–9909. PMLR, 23–29 Jul 2023b.  URL https://proceedings.mlr.press/v202/fellows23a.html.
  * Filos et al. [2021] Angelos Filos, Clare Lyle, Yarin Gal, Sergey Levine, Natasha Jaques, and Gregory Farquhar.  Psiphi-learning: Reinforcement learning with demonstrations using successor features and inverse temporal difference learning, 2021.
  * Finn et al. [2016] Chelsea Finn, Sergey Levine, and Pieter Abbeel.  Guided cost learning: Deep inverse optimal control via policy optimization.  In _International conference on machine learning_ , pages 49–58. PMLR, 2016.
  * Hazan et al. [2019] Elad Hazan, Sham Kakade, Karan Singh, and Abby Van Soest.  Provably efficient maximum entropy exploration.  In Kamalika Chaudhuri and Ruslan Salakhutdinov, editors, _Proceedings of the 36th International Conference on Machine Learning_ , volume 97 of _Proceedings of Machine Learning Research_ , pages 2681–2691. PMLR, 09–15 Jun 2019.  URL https://proceedings.mlr.press/v97/hazan19a.html.
  * Ho and Ermon [2016] Jonathan Ho and Stefano Ermon.  Generative adversarial imitation learning.  _Advances in neural information processing systems_ , 29, 2016.
  * Hussein et al. [2017] Ahmed Hussein, Mohamed Medhat Gaber, Eyad Elyan, and Chrisina Jayne.  Imitation learning: A survey of learning methods.  _ACM Computing Surveys (CSUR)_ , 50(2):1–35, 2017\.
  * Janz et al. [2019] David Janz, Jiri Hron, Przemysław Mazur, Katja Hofmann, José Miguel Hernández-Lobato, and Sebastian Tschiatschek.  Successor uncertainties: Exploration and uncertainty in temporal difference learning.  In _Proceedings of the 33rd International Conference on Neural Information Processing Systems_ , Red Hook, NY, USA, 2019. Curran Associates Inc.
  * Kaelbling et al. [1998] Leslie Pack Kaelbling, Michael L Littman, and Anthony R Cassandra.  Planning and acting in partially observable stochastic domains.  _Artificial intelligence_ , 101(1-2):99–134, 1998\.
  * Kingma and Welling [2014] Diederik P Kingma and Max Welling.  Auto-encoding variational bayes.  In _International Conference on Learning Representations_ , 2014.
  * Kwon et al. [2020] Minhae Kwon, Saurabh Daptardar, Paul R Schrater, and Xaq Pitkow.  Inverse rational control with partially observable continuous nonlinear dynamics.  _Advances in neural information processing systems_ , 33:7898–7909, 2020.
  * Mnih et al. [2013] Volodymyr Mnih, Koray Kavukcuoglu, David Silver, Alex Graves, Ioannis Antonoglou, Daan Wierstra, and Martin Riedmiller.  Playing atari with deep reinforcement learning.  _arXiv preprint arXiv:1312.5602_ , 2013.
  * Murphy [2013] Kevin P. Murphy.  _Machine learning : a probabilistic perspective_.  MIT Press, Cambridge, Mass. [u.a.], 2013.  ISBN 9780262018029 0262018020.
  * Nguyen et al. [2022] Hai Huu Nguyen, Andrea Baisero, Dian Wang, Christopher Amato, and Robert Platt.  Leveraging fully observable policies for learning under partial observability.  In _6th Annual Conference on Robot Learning_ , 2022.
  * Ortega et al. [2021] Pedro A Ortega, Markus Kunesch, Grégoire Delétang, Tim Genewein, Jordi Grau-Moya, Joel Veness, Jonas Buchli, Jonas Degrave, Bilal Piot, Julien Perolat, Tom Everitt, Corentin Tallec, Emilio Parisotto, Tom Erez, Yutian Chen, Scott Reed, Marcus Hutter, Nando de Freitas, and Shane Legg.  Shaking the foundations: delusions in sequence models for interaction and control.  October 2021.
  * Rafailov et al. [2023] Rafael Rafailov, Archit Sharma, Eric Mitchell, Stefano Ermon, Christopher D. Manning, and Chelsea Finn.  Direct preference optimization: Your language model is secretly a reward model.  _NeurIPS_ , abs/2305.18290, 2023.  URL https://api.semanticscholar.org/CorpusID:258959321.
  * Ramachandran and Amir [2007] Deepak Ramachandran and Eyal Amir.  Bayesian inverse reinforcement learning.  In _Proceedings of the 20th International Joint Conference on Artifical Intelligence_ , IJCAI’07, page 2586–2591, San Francisco, CA, USA, 2007\. Morgan Kaufmann Publishers Inc.
  * Schaal [1999] Stefan Schaal.  Is imitation learning the route to humanoid robots?  _Trends in cognitive sciences_ , 3(6):233–242, 1999.
  * Straub et al. [2024] Dominik Straub, Matthias Schultheis, Heinz Koeppl, and Constantin A Rothkopf.  Probabilistic inverse optimal control for non-linear partially observable systems disentangles perceptual uncertainty and behavioral costs.  _Advances in Neural Information Processing Systems_ , 36, 2024.
  * Sutton and Barto [2018] Richard S. Sutton and Andrew G. Barto.  _Reinforcement Learning: An Introduction_.  The MIT Press, second edition, 2018.  URL http://incompleteideas.net/book/the-book-2nd.html.
  * Swamy et al. [2022] Gokul Swamy, Sanjiban Choudhury, J Andrew Bagnell, and Zhiwei Steven Wu.  Sequence model imitation learning with unobserved contexts.  _Advances in Neural Information Processing Systems_ , 35, 2022.
  * Tsitsiklis and Van Roy [1997] J.N. Tsitsiklis and B. Van Roy.  An analysis of temporal-difference learning with function approximation.  _IEEE Transactions on Automatic Control_ , 42(5):674–690, 1997.  doi: 10.1109/9.580874.
  * van der Vaart [1998] A.W. van der Vaart.  _Asymptotic Statistics_.  Cambridge series on statistical and probabilistic mathematics. Cambridge University Press, 1998.  ISBN 9780521496032.  URL https://books.google.co.uk/books?id=fiX9ngEACAAJ.
  * Vuorio et al. [2022] Risto Vuorio, Johann Brehmer, Hanno Ackermann, Daniel Dijkman, Taco Cohen, and Pim de Haan.  Deconfounded imitation learning.  _arXiv preprint arXiv:2211.02667_ , 2022.
  * Walsman et al. [2023] Aaron Walsman, Muru Zhang, Sanjiban Choudhury, Dieter Fox, and Ali Farhadi.  Impossibly good experts and how to follow them.  In _The Eleventh International Conference on Learning Representations_ , 2023.  URL https://openreview.net/forum?id=sciA_xgYofB.
  * Weihs et al. [2021] Luca Weihs, Unnat Jain, Iou-Jen Liu, Jordi Salvador, Svetlana Lazebnik, Aniruddha Kembhavi, and Alex Schwing.  Bridging the imitation gap by adaptive insubordination.  In A. Beygelzimer, Y. Dauphin, P. Liang, and J. Wortman Vaughan, editors, _Advances in Neural Information Processing Systems_ , 2021.  URL https://openreview.net/forum?id=Wlx0DqiUTD_.
  * Ziebart et al. [2008] Brian D Ziebart, Andrew L Maas, J Andrew Bagnell, Anind K Dey, et al.  Maximum entropy inverse reinforcement learning.  In _Aaai_ , volume 8, pages 1433–1438. Chicago, IL, USA, 2008.
  * Zintgraf et al. [2020] Luisa Zintgraf, Kyriacos Shiarlis, Maximilian Igl, Sebastian Schulze, Yarin Gal, Katja Hofmann, and Shimon Whiteson.  Varibad: A very good method for bayes-adaptive deep rl via meta-learning.  In _International Conference on Learning Representations_ , 2020.
  * Zintgraf et al. [2021a] Luisa Zintgraf, Sebastian Schulze, Cong Lu, Leo Feng, Maximilian Igl, Kyriacos Shiarlis, Yarin Gal, Katja Hofmann, and Shimon Whiteson.  Varibad: Variational bayes-adaptive deep rl via meta-learning.  _Journal of Machine Learning Research_ , 22(289):1–39, 2021a.  URL http://jmlr.org/papers/v22/21-0657.html.
  * Zintgraf et al. [2021b] Luisa M Zintgraf, Leo Feng, Cong Lu, Maximilian Igl, Kristian Hartikainen, Katja Hofmann, and Shimon Whiteson.  Exploration in approximate hyper-state space for meta reinforcement learning.  In Marina Meila and Tong Zhang, editors, _Proceedings of the 38th International Conference on Machine Learning_ , volume 139 of _Proceedings of Machine Learning Research_ , pages 12991–13001. PMLR, 18–24 Jul 2021b.  URL https://proceedings.mlr.press/v139/zintgraf21a.html.

## Supplementary Material

##  Appendix A Summary of Distributions in BIG

Due to the diverse nature of the sources of input information, inferring the
reward posterior requires the agent to learn and maintain several
distributions over random variables, we summarize them in Table 1.

Table 1:  Summary of the distributions involved in BIG. This table supplements the diagram in Figure 2. DISTRIBUTION | NAME | DESCRIPTION | SAMPLES TO LEARN
---|---|---|---
$p(\omega)$ |  | Reward
---
Parameter Prior
| Incorporates prior knowledge in reward
---
parameterization
None
$p(k)$ |  | Cost of
---
Exploration Prior
| Incorporates prior knowledge in rewards
---
$r(s,a)=kr_{\textrm{max}}$ over exploratory state-actions
None
$p(\theta)$ | Contextual Prior |  |  Characterizes uncertainty in $\theta$ a priori at
---
test time
|  $\mathcal{D}_{\textrm{Sim}}$ \- samples from the CMDP
---
simulator
$p(a|s,\theta,\omega)$ |  | Model
---
Expert Policy
| Model of optimal policy for expert in
---
CMDP $\mathcal{M}(\theta)$ with reward parameters $\omega$
| Contextual successor features
---
learned from simulator and $\mathcal{D}_{\textrm{Expert}}$
$p(\omega|\mathcal{D}_{\textrm{Expert}})$ |  | Expert reward
---
Posterior
| Updates the reward parameter prior using
---
expert data
|  $\mathcal{D}_{\textrm{Expert}}$ dataset of expert
---
trajectories
$p(\theta_{i}|\tau_{i})$ |  | Contextual
---
Posterior
| Characterizes uncertainty over which
---
latent variable $\theta_{i}$ agent $i$ was assigned
|  $\tau_{i}$ \- each expert or exploratory
---
agent’s trajectory
$p^{\textrm{IRL+COE}}_{\textrm{Bayes}}(r|s,a)$ |  | Bayesian Reward
---
Distribution
| Incorporates epistemic uncertainty from
---
$p(\omega|\mathcal{D}_{\textrm{Expert}})$ and $p(k)$ into reward model
| None
---

##  Appendix B Bayesian Successor Feature Learning

Consider the optimal $Q$-function $Q^{\star}(s,a,\theta,\omega)$, which
satisfies the optimal contextual Bellman equation:
$\mathcal{B}^{\star}[Q^{\star}](s,a,\theta,\omega)=Q^{\star}(s,a,\theta,\omega)$
where:

$\displaystyle\mathcal{B}^{\star}[Q^{\star}](s,a,\theta,\omega)\coloneqq\nu(s,a)^{\top}\omega$
(11) $\displaystyle\qquad+\gamma\mathbb{E}_{s^{\prime}\sim
p(s^{\prime}|s,a,\theta)}\left[\sup_{a^{\prime}}Q^{\star}(s^{\prime},a^{\prime},\theta,\omega)\right].$
(12)

As the expert selects actions
$a\in\operatorname*{arg\,max}_{a^{\prime}}Q^{\star}(s,a^{\prime},\theta,\omega)$,
learning $Q^{\star}(s,a,\theta,\omega)$ is sufficient for modeling the set of
expert policies, from which the true reward parametrization can be inferred.

Learning with Expert Data. Consider the Bellman equation under the expert
policy $\pi^{\star}(a^{\prime}|s^{\prime},\theta_{i})$. The successor feature
representation $\Psi_{\phi}$ should satisfy:

$\displaystyle\Psi_{\phi}(s,a,\theta_{i})^{\top}\omega^{\star}=\nu(s,a)^{\top}\omega^{\star}+\gamma\mathbb{E}_{s^{\prime}\sim
p(s^{\prime}|s,a,\theta_{i}),a^{\prime}\sim\pi^{\star}(a^{\prime}|s^{\prime},\theta_{i})}\left[\Psi_{\phi}(s^{\prime},a^{\prime},\theta_{i})\right]^{\top}\omega^{\star}.$
(13)

We can factor $\omega^{\star}$ from the Bellman equation and then solve:

$\displaystyle\Psi_{\phi}(s,a,\theta_{i})=\nu(s,a)+\gamma\mathbb{E}_{s^{\prime}\sim
p(s^{\prime}|s,a,\theta_{i}),a^{\prime}\sim\pi^{\star}(a^{\prime}|s^{\prime},\theta_{i})}\left[\Psi_{\phi}(s^{\prime},a^{\prime},\theta_{i})\right].$
(14)

This yields the objective for each expert trajectory $\tau_{i}$:

$\displaystyle\mathcal{L}_{\textrm{Expert}}(\phi;\tau_{i})\coloneqq\mathbb{E}_{s,a\sim\textrm{Unif}(\tau_{i}),\theta_{i}\sim
p(\theta_{i}|\tau_{i})}\big{[}\lVert$
$\displaystyle\Psi_{\phi}(s,a,\theta_{i})-(\nu(s,a)$ (15)
$\displaystyle+\gamma\mathbb{E}_{s^{\prime}\sim
p(s^{\prime}|s,a,\theta_{i}),a^{\prime}\sim\pi^{\star}(a^{\prime}|s^{\prime},\theta_{i})}\left[\Psi_{\phi}(s^{\prime},a^{\prime},\theta_{i})\right])\rVert\big{]}~{}.$
(16)

where $\textrm{Unif}(\tau_{i})$ is a uniform distribution over the state-
action pairs in trajectory $\tau_{i}$. Minimizing
$\mathcal{L}_{\textrm{Expert}}(\phi;\tau_{i})$ for each trajectory ensures
that the successor feature representation is consistent with the expert
demonstrations. Minimizing $\mathcal{L}_{\textrm{Expert}}(\phi;\tau_{i})$ can
be carried out using a semi-gradient approach, thereby avoiding the need for
two samples from the expert policy and state-transition distribution:

$\displaystyle\phi\leftarrow\phi+\eta_{\phi}\nabla_{\phi}\Psi_{\phi}(s,a,\theta_{i})(\nu(s,a)+\gamma\Psi_{\phi}(s^{\prime},a^{\prime},\theta_{i})-\Psi_{\phi}(s,a,\theta_{i})).$
(17)

As the latent contextual variable $\theta_{i}$ is never observed, we must
infer a posterior over its value $p(\theta_{i}|\tau_{i})$. It may seem that
this could be avoided via IRL using the prior-averaged transitions
$\mathbb{E}_{\theta\sim p(\theta)}\left[p(s^{\prime}|s,a,\phi)\right]$.
However, we provide a simple counterexample in Appendix C demonstrating that a
prior-averaged approach does not account for the true reward ordering in the
underlying space of CMDPs. Sampling $\theta_{i}\sim p(\theta_{i}|\tau_{i})$
from the posterior over the expert’s contextual variable is typically
intractable, so we use variational inference instead, as detailed in Appendix
E.

Learning with a Simulator. In addition to expert demonstrations, we also have
access to samples from the CMDP simulator. This allows us to learn
$\Psi_{\phi}$ over state-action pairs where the expert has not provided
demonstrations. We sample a dataset
$\mathcal{D}_{\textrm{Simulator}}\coloneqq\\{\tau_{i}\\}_{i=1}^{N_{\textrm{Simulator}}}$
of $N_{\textrm{Simulator}}$ trajectories from the simulator: the simulator
samples an MDP from the prior, and then an exploration policy
$\pi_{\textrm{Explore}}$ interacts with the corresponding MDP, observing
state-action-state transitions. In this paper, we use an $\epsilon$-greedy
exploration policy where the greedy actions
$a^{\prime}\in\operatorname*{arg\,max}_{a^{\prime}}\Psi_{\phi}(s^{\prime},a^{\prime},\theta_{i})^{\top}\omega$
are taken with probability $1-\epsilon$ then uniformly otherwise. More
sophisticated approaches such as a policy that maximizes the entropy of the
ergodic, discounted state-action occupancy distribution [14] would also be
appropriate, especially in larger domains. However, our experiments
demonstrate the $\epsilon$-greedy policy suffices for learning successor
features in our setting.

Once samples have been obtained, we infer the posterior
$p(\theta_{i}|\tau_{i})$ over the MDP that the exploratory agent was assigned
for each MDP $i\in[1:N_{\textrm{Simulator}}]$. Like with the expert data, our
goal is to ensure that the successor feature representation $\Psi_{\phi}$
satisfies a Bellman equation. In the target, we choose the next action that
maximizes the $Q$-function for a given $\omega$,
$a^{\prime}\in\operatorname*{arg\,max}_{a^{\prime}}\Psi_{\phi}(s^{\prime},a^{\prime},\theta_{i})^{\top}\omega$.
This yields the objective for each MDP $i$:

$\displaystyle\mathcal{L}_{\textrm{Simulator}}(\phi;\tau_{i},\omega)\coloneqq~{}$
$\displaystyle\mathbb{E}_{s,a\sim\textrm{Unif}(\tau_{i}),\theta_{i}\sim
p(\theta_{i}|\tau_{i})}\big{[}$ (18)
$\displaystyle\big{(}\Psi_{\phi}(s,a,\theta_{i})-(\nu(s,a)+\gamma\mathbb{E}_{s^{\prime}\sim
p(s^{\prime}|s,a,\theta_{i})}\left[\Psi_{\phi}(s^{\prime},a^{\prime},\theta_{i})\right]\big{)}^{2}\big{]}.$
(19)

where $\textrm{Unif}(\tau_{i})$ is a uniform distribution over the state-
action pairs in trajectory $\tau_{i}$. We optimize this objective using the
following TD update:

$\displaystyle\phi\leftarrow\phi+\eta_{\phi}\nabla_{\phi}\Psi_{\phi}(s,a,\theta_{i})\cdot(\nu(s,a)+\gamma\Psi_{\phi^{\prime}}(s^{\prime},a^{\prime},\theta_{i})-\Psi_{\phi}(s,a,\theta_{i})).$
(20)

There are two differences between the updates for the exploratory data and the
expert data. First, the exploratory data updates are off-policy. Due to the
deadly triad [29], using semi-gradients is not guaranteed to converge [31,
11]. We introduce a separate target network $\phi^{\prime}$ that is updated
periodically to stabilize the updates. Second, the updates depend on $\omega$
as the supremum acts over the dot product between the successor representation
and the reward parametrization. As we require successor features to infer
$\omega$, we interleave learning both $\omega$ and $\phi$ in a nested
optimization, using both expert and exploratory data with an initial burn-in
period using only the expert data.

We combine both objectives into the single objective presented in Section 4.1:

$\displaystyle\mathcal{L}(\phi;\omega)=\mathcal{L}_{\textrm{Expert}}(\phi;\tau_{i})+\beta\mathcal{L}_{\textrm{Simulator}}(\phi;\tau_{i},\omega),$
(21)

for some constant $\beta$. However, we note that the reward learning and
successor feature learning depend on each other through the objective
$\mathcal{L}_{\textrm{Simulator}}$ and reward update given by Theorem 4.1,
yielding a two-timescale learning problem. To enable convergence to a stable
local optimum, the learning rates $\eta_{\phi}$ and $\eta_{\omega}$ are set
such that $\eta_{\omega}<\eta_{\phi}$.

##  Appendix C The Need for Inference Over Context Variables

Figure 6: Counterexample CMDP

We consider a simple counterexample illustrated in Figure 6 with four states
$s_{i}$ and two possible hidden contexts denoted by $\theta$. As shown in the
figure, $\theta$ controls the environment transition. The simplest approach to
IRL using the average MDP would map the experts’ state visitation frequency
$\eta$ to reward value. This would assign rewards: $r(s_{1})=c\eta$,
$r(s_{2})=c(1-\eta)$, $r(s_{3})=c$, where $c$ is an arbitrary finite positive
constant. This means that the values of $r(s_{1})$ and $r(s_{2})$ are only
determined by $\eta$ and will give incorrect reward ordering $r(s_{2})\geq
r(s_{1})$ for any $\eta\leq 0.5$. On the other hand, if a Bayesian approach is
taken, it is possible to infer which MDP the agent was in. If the belief of a
trajectory is weighted towards the expert being in $\theta=0$, the expert’s
preference for $a_{0}=\textrm{left}$ over $a_{1}=\textrm{right}$ must be a
consequence of the reward ordering $r(s_{1})>r(s_{2})$. Likewise, if the
belief of a trajectory is weighted towards the expert being in $\theta=1$, the
expert’s preference for $a_{0}=\textrm{right}$ must imply that
$r(s_{2})+\gamma r(s_{3})>r(s_{1})+\gamma r(s_{0})$. As $r(s_{1})>r(s_{2})$
will be inferred from the trajectories of experts in $\theta=0$, this implies
that $r(s_{0})<r(s_{3})$, and so a correct reward ordering will be learned
regardless of $\eta$. By being Bayesian, these possible inferences condition
on which $\theta_{i}$ the expert was in, and are explained by
$Q_{\zeta^{\star}}(s,a,\theta_{i},\omega)$. We can best match $\omega$ so that
is consistent with these inferences under the belief
$p(\theta_{i}|\tau_{i},\omega)$ in the expert’s MDP. We thus conclude that
compared to a fully Bayesian approach, naively using imitation learning on the
prior-averaged MDP will not yield policies that account for reward ordering of
the underlying MDP.

##  Appendix D Bayesian Solution to the Imitation Gap

Finally, we present pseudocode for the Bayesian solution to the Imitation Gap
(BIG) in Algorithm 1. It takes as inputs the prior distributions and the
expert dataset and produces a reward function $r_{\textrm{Bayes}}^{\text{IRL +
COE}}$. In the main loop, data is collected from the simulated environment
using an $\epsilon$-greedy policy. We use $\epsilon$-greedy for convenience as
we found it to be an easy way to achieve suitable coverage of the state-action
space in the experiments. The data collected using the random policy is used
for learning the successor features in an off-policy RL algorithm. With an
off-policy algorithm, assuming coverage of the whole state-action space, the
exact kind of randomness does not matter for the kind of successor features we
learn. The collected data together with the expert demonstrations are used to
update the successor feature representation. The learned successor features
are used for updating the reward parameters to maximize the likelihood of the
expert data. After $K$ steps of the main loop, the final reward function is
computed by applying the cost-of-exploration refinement defined in Section
4.3.

Algorithm 1 Bayesian Solution to the Imitation Gap

0: priors $p(\theta)$, $p(\omega)$, and $p_{\textrm{COE}}(r|\sigma,s,a)$,
contextual posterior $p(\theta|\tau)$, dataset
$\mathcal{D}_{\textrm{Expert}}$, reward scales $r_{\textrm{min}}$ and
$r_{\textrm{max}}$, learning rates $\eta_{\omega}$ and $\eta_{\phi}$, loss
coefficient $\beta$, and number of steps $K$.

Initialize parameters $\phi$ for the successor features and $\omega$ for the
reward.

Initialize an empty replay buffer $\mathcal{D}_{\textrm{Replay}}$.

for $K$ steps do

Sample a new MDP from simulator $\theta_{i}\sim\rho(\theta)$.

Sample a trajectory $\tau_{i}$ in the MDP defined by $\theta_{i}$ using an
epsilon-greedy policy w.r.t. the $Q$-function defined by
$\Psi_{\phi}(s,a,\tilde{\theta}_{i})^{\top}\omega$ for $\tilde{\theta}_{i}\sim
p(\theta_{i}|\tau_{i})$.

Add the trajectory $\tau_{i}$ to the replay buffer
$\mathcal{D}_{\textrm{Replay}}$.

Update $\phi$ to minimise $\mathcal{L}(\phi;\omega)$ using
$\mathcal{D}_{\textrm{Expert}}$ and $\mathcal{D}_{\textrm{Replay}}$.

Update $\omega$ using Equation 4 with $\mathcal{D}_{\textrm{Expert}}$.

end for

Rescale $\omega$ s.t. rewards lie in the range
$[r_{\textrm{min}},r_{\textrm{max}}]$.

Compute the predictive reward $r_{\textrm{Bayes}}^{\text{IRL + COE}}$ .

return $r_{\textrm{Bayes}}^{\text{IRL + COE}}$

##  Appendix E Approximate Inference

Whilst the full Bayesian approach outlined in Section 4 is clearly desirable,
there are several sources of intractability that prevent us from computing the
Bayes-optimal policy $\pi^{\star}_{\textrm{Bayes}}$ exactly. Firstly,
maintaining and inferring the posterior distributions is likely to be
intractable for all but the simplest choice of likelihoods, which have
insufficient capacity for representing the set of MDPs beyond contrived toy
tasks. Secondly, marginalization using the posteriors likely will involve high
dimensional integrals, which are computationally inefficient. Finally, solving
the planning problem in the BAMDP for every possible history to obtain a
Bayes-optimal policy is notoriously challenging, even for very simple domains
[37, 38]. We thus derive an algorithm that follows the methodology of the
formal Bayesian approach outlined in Section 4 whilst making necessary
approximations from the powerful toolkit of variational inference to ensure
tractability.

###  E.1 Tractable Prior and Likelihood Learning

Instead of attempting to infer the posterior
$p(\phi|\mathcal{D}_{\textrm{simulator}})$ exactly, which may be intractable,
we use a MAP approach instead to learn a point estimate
$\phi^{\star}_{\textrm{MAP}}$. The justification for this is that we have
access to a simulator from which a large number of samples can be drawn. The
Bernstein-von Mises theorem specifies that in the limit of large data
$K\rightarrow\infty$,
$p(\phi|\mathcal{D}_{\textrm{simulator}})\rightarrow\delta_{\phi^{\star}_{\textrm{MLE}}}(\phi)$
where $\phi^{\star}_{\textrm{MLE}}$ is the maximum likelihood estimator, so we
expect
$p(\phi|\mathcal{D}_{\textrm{simulator}})\approx\delta_{\phi^{\star}_{\textrm{MLE}}}(\phi)$.
In the same limit,
$\phi^{\star}_{\textrm{MAP}}\rightarrow\phi^{\star}_{\textrm{MLE}}$. Our
choice of using a MAP estimator over the MLE estimator is purely practical in
that the prior can add regularization to stabilize learning, as is seen by the
contribution to the MAP objective:

$\displaystyle\mathcal{L}_{\textrm{MAP}}(\phi)=\sum_{i=1}^{K}\log
p(\tau_{i}|\phi)+\log(\phi).$ (22)

To find a tractable objective to maximize the MAP estimate, we introduce a
variational distribution $q_{\chi_{i}}(\theta_{i})$ parametrized by
${\chi_{i}}\in X$, for which:

$\displaystyle\log p(\tau_{i}|\phi)$ $\displaystyle=\int\log
p(\tau_{i}|\phi)q_{{\chi_{i}}}(\theta_{i})d\theta_{i},$ (23)
$\displaystyle=\int\log\left(\frac{p(\tau_{i},\theta_{i}|\phi)}{p(\theta_{i}|\tau_{i},\phi)}\right)q_{{\chi_{i}}}(\theta_{i})d\theta_{i},$
(24)
$\displaystyle=\int\log\left(\frac{p(\tau_{i},\theta_{i}|\phi)}{p(\theta_{i}|\tau_{i},\phi)}\cdot\frac{q_{{\chi_{i}}}(\theta_{i})}{q_{{\chi_{i}}}(\theta_{i})}\right)q_{{\chi_{i}}}(\theta_{i})d\theta_{i},$
(25) $\displaystyle=\int\left(\log p(\tau_{i},\theta_{i}|\phi)-\log
q_{{\chi_{i}}}(\theta_{i})\right)q_{{\chi_{i}}}(\theta_{i})d\theta_{i}-\int\frac{\log
p(\theta_{i}|\tau_{i},\phi)}{\log
q_{\chi_{i}}(\theta_{i})}q_{\chi_{i}}(\theta_{i})d\theta_{i},$ (26)
$\displaystyle=\mathcal{L}_{\textrm{ELBO}}(\phi,{\chi_{i}})+\textrm{KL}(q_{\chi_{i}}\;\|\;p(\cdot|\tau_{i},\phi)),$
(27) $\displaystyle\implies\mathcal{L}_{\textrm{ELBO}}(\phi,{\chi_{i}})$
$\displaystyle=\log
p(\tau_{i}|\phi)-\textrm{KL}(q_{\chi_{i}}\;\|\;p(\cdot|\tau_{i},\phi)),$ (28)

where $\mathcal{L}_{\textrm{ELBO}}(\cdot)$ denotes an evidence lower bound
(ELBO). Now, assuming there exists some ${\chi_{i}}^{\star}\in X$ such that
$q_{{\chi_{i}}^{\star}}(\theta_{i})=p(\theta_{i}|\tau_{i},\phi)$, then
$\sup_{{\chi_{i}}\in
X}\textrm{KL}(q_{\chi_{i}}\;\|\;p(\cdot|\tau_{i},\phi))=0$, hence it follows
from Equation 28 that

$\displaystyle\log p(\tau_{i}|\phi)=\sup_{{\chi_{i}}\in
X}\mathcal{L}_{\textrm{ELBO}}(\phi,{\chi_{i}}).$ (29)

Substituting into the MAP objective from Equation 22 yields:

$\displaystyle\mathcal{L}_{\textrm{MAP}}(\phi)=\sum_{i=1}^{K}\sup_{{\chi_{i}}\in
X}\mathcal{L}_{\textrm{ELBO}}(\phi,{\chi_{i}})+\log(\phi).$ (30)

To find a
$\phi^{\star}_{\textrm{MAP}}\in\operatorname*{arg\,sup}_{\phi\in\Phi}\mathcal{L}_{\textrm{MAP}}(\phi)$,
we recognize that for each MDP sampled from the simulator indexed by
$i\in[1:K]$, we can minimize the following index-specific ELBO:

$\displaystyle\mathcal{L}_{\textrm{ELBO}}(\phi,{\chi_{i}})=\mathbb{E}_{\theta_{i}\sim
q_{\chi_{i}}(\theta_{i})}\left[\sum_{t=1}^{H_{i}}\log
p(s_{t}|s_{t-1},a_{t-1},\theta_{i})+\log
p(\theta_{i}|\phi)-\log(q_{\chi_{i}}(\theta_{i}))\right],$ (31)

with respect to $\phi$ and $\chi_{i}$. Here, Variational Auto-encoders [19]
(VAEs) offer a powerful framework for solving this problem via a variational
expectation-maximization (EM) algorithm. In VAE parlance, the distribution
$p(\tau_{i}|\theta_{i},\phi)$ is known as the _decoder_ , parametrized by
$\phi\in\Phi$.

Algorithm 2 LearnPrior+Likelihood

Initialize $\zeta,\chi,\phi$

for $K$ steps do

Sample new MDP from simulator $\theta\sim\rho(\theta)$

Sample initial state $s_{0}\sim p_{0}(s_{0})$

for $t\in[1:H_{i}]$ do

Sample action $a_{t-1}\sim d(a_{t-1})$

Sample transition $s_{t}\sim p(s_{t}|s_{t-1},a_{t-1},\theta)$

Sample gradient
$g_{\chi}\sim\nabla_{\psi}\mathcal{L}_{\textrm{ELBO}}(\phi,\chi)$

$\chi\leftarrow\chi+\alpha_{\chi}^{t}g_{\chi}$

Sample gradient

$g_{\phi}\sim\nabla_{\phi}\left(\mathcal{L}_{\textrm{ELBO}}(\phi,\chi)+\frac{1}{KH_{i}}\log
p(\phi)\right)$

$\phi\leftarrow\phi+\alpha_{\phi}^{t}g_{\phi}$

$g_{\zeta}\sim\nabla_{\zeta}\textrm{MSBE}(\zeta)$

$\zeta\leftarrow\zeta-\alpha_{\zeta}^{t}g_{\zeta}$

end for

end for

return $\zeta,\phi$

For each MDP $\theta_{i}$, we train an _encoder_ $q_{\chi_{i}}(\theta_{i})$,
parametrized by $\chi_{i}\in X$, that acts as a variational approximation to
the posterior: $p(\theta_{i}|\tau_{i},\phi)$.

To learn the prior parameters, we propose Algorithm 2. Instead of specifying
an encoder for each $q_{\chi_{i}}(\theta_{i})$, we train a single encoder
$q_{\chi}(\theta_{i})$ online using an entire trajectory of samples from a
specific MDP $\theta_{i}$. Once training has finished for that MDP, we sample
another $\theta_{i+1}$, using $\chi$ as the initialization parameters for the
new encoder, $q_{\chi}(\theta_{i+1})$. Moreover, as both learning the reward
prior and minimizing the MSBE to learn the likelihood require samples from a
simulator, we interleave both optimization problems using the same
interactions with the simulator.

##  Appendix F Proofs and Derivations

######  Lemma 4.1.

The gradient of the log normalization constant $\log z(s,\omega,\theta_{i})$
is

$\displaystyle\nabla_{\omega}\log z(s,\omega,\theta_{i})=\mathbb{E}_{a\sim
p(a|s,\omega,\theta_{i})}\left[\Psi^{\star}(s,a,\theta_{i})\right].$ (32)

###### Proof.

We assume that $\mathcal{A}$ is continuous. If $\mathcal{A}$ is discrete, we
replace the Lebesgue measure $\lambda$ with the counting measure, and our
proof remains unchanged. Taking derivatives directly yields our desired
result:

$\displaystyle\nabla_{\omega}\log z(s,\omega,\theta_{i})$
$\displaystyle=\nabla_{\omega}\log\int_{\mathcal{A}}\exp\left(\frac{1}{\alpha}\Psi^{\star}(s,a,\theta_{i})^{\top}\omega\right)d\lambda(a),$
(33)
$\displaystyle=\nabla_{\omega}\int_{\mathcal{A}}\exp\left(\frac{1}{\alpha}\Psi^{\star}(s,a,\theta_{i})^{\top}\omega\right)d\lambda(a)\cdot\frac{1}{\int_{\mathcal{A}}\exp\left(\frac{1}{\alpha}\Psi^{\star}(s,a,\theta_{i})^{\top}\omega\right)d\lambda(a)},$
(34)
$\displaystyle=\int_{\mathcal{A}}\Psi^{\star}(s,a,\theta_{i})\frac{\exp\left(\frac{1}{\alpha}\Psi^{\star}(s,a,\theta_{i})^{\top}\omega\right)}{\int_{\mathcal{A}}\exp\left(\frac{1}{\alpha}\Psi^{\star}(s,a,\theta_{i})^{\top}\omega\right)d\lambda(a)}d\lambda(a),$
(35)
$\displaystyle=\int_{\mathcal{A}}\Psi^{\star}(s,a,\theta_{i})p(a|s,\omega,\theta_{i})d\lambda(a),$
(36) $\displaystyle=\mathbb{E}_{a\sim
p(a|s,\omega,\theta_{i})}\left[\Psi^{\star}(s,a,\theta_{i})\right].$ (37)

∎

######  Theorem 4.1.

Define $\varsigma_{0}^{2}\coloneqq\frac{\sigma_{0}^{2}}{\alpha}$. Using the
Laplace approximation for the posterior, the approximate Bayesian reward
parametrization
$\omega^{\star}_{\textrm{Bayes}}\approx\omega^{\star}_{\textrm{Laplace}}$ can
be found by carrying out the following stochastic gradient descent updates on
the log-posterior:

$\displaystyle\omega\leftarrow$
$\displaystyle\omega+\eta_{\omega}\left(N_{\textrm{Expert}}{H_{i}}\bigg{(}\Psi^{\star}(s_{j},a_{j},\theta_{i})-\mathbb{E}_{a_{i}\sim
p(a_{i}|s_{j},\omega,\theta_{i})}\left[\Psi^{\star}(s_{j},a_{i},\theta_{i})\right]\right)-\frac{(\omega-\omega_{0})}{\varsigma_{0}^{2}}\bigg{)}.$
(38)

###### Proof.

Using Laplace’s approximation, we fit a Gaussian to the posterior distribution
with mean
$\omega^{\star}_{\textrm{Laplace}}\in\operatorname*{arg\,sup}_{\omega\in\Omega}p(\omega|\mathcal{D}_{\textrm{Expert}})$.
Equivalently, we can maximize the log posterior instead. Defining
$g(\omega;\mathcal{D}_{\textrm{Expert}})\coloneqq\nabla_{\omega}\log
p(\omega|\mathcal{D}_{\textrm{Expert}})$, from the definition of the gradient
of a log:

$\displaystyle g(\omega;\mathcal{D}_{\textrm{Expert}})=\nabla_{\omega}\log
p(\omega|\mathcal{D}_{\textrm{Expert}})=\nabla_{\omega}p(\omega|\mathcal{D}_{\textrm{Expert}})\cdot\frac{1}{p(\omega|\mathcal{D}_{\textrm{Expert}})}.$
(39)

We define the joint set of expert contextual variables as
$\Theta_{\textrm{Expert}}\coloneqq\\{\theta_{1},\theta_{2},...\theta_{N_{\textrm{Expert}}}\\}\in\Theta^{N_{\textrm{Expert}}}$.
Taking gradients of the posterior yields:

$\displaystyle\nabla_{\omega}p(\omega|\mathcal{D}_{\textrm{Expert}})$
$\displaystyle=\int_{\Theta^{N_{\textrm{Expert}}}}\nabla_{\omega}p(\omega|\mathcal{D}_{\textrm{Expert}},\Theta_{\textrm{Expert}})dP(\Theta_{\textrm{Expert}}|\mathcal{D}_{\textrm{Expert}}),$
(40)
$\displaystyle=\int_{\Theta^{N_{\textrm{Expert}}}}\frac{\nabla_{\omega}p(\omega|\mathcal{D}_{\textrm{Expert}},\Theta_{\textrm{Expert}})}{p(\omega|\mathcal{D}_{\textrm{Expert}},\Theta_{\textrm{Expert}})}p(\omega|\mathcal{D}_{\textrm{Expert}},\Theta_{\textrm{Expert}})dP(\Theta_{\textrm{Expert}}|\mathcal{D}_{\textrm{Expert}}),$
(41) $\displaystyle=\int_{\Theta^{N_{\textrm{Expert}}}}\nabla_{\omega}\log
p(\omega|\mathcal{D}_{\textrm{Expert}},\Theta_{\textrm{Expert}})dP(\Theta_{\textrm{Expert}},\omega|\mathcal{D}_{\textrm{Expert}}),$
(42) $\displaystyle=\int_{\Theta^{N_{\textrm{Expert}}}}\nabla_{\omega}\log
p(\omega|\mathcal{D}_{\textrm{Expert}},\Theta_{\textrm{Expert}})dP(\Theta_{\textrm{Expert}}|\omega,\mathcal{D}_{\textrm{Expert}})p(\omega|\mathcal{D}_{\textrm{Expert}}).$
(43)

Substituting yields our desired result:

$\displaystyle\nabla_{\omega}\log
p(\omega|\mathcal{D}_{\textrm{Expert}})=\int_{\Theta^{N_{\textrm{Expert}}}}\nabla_{\omega}\log
p(\omega|\mathcal{D}_{\textrm{Expert}},\Theta_{\textrm{Expert}})dP(\Theta_{\textrm{Expert}}|\omega,\mathcal{D}_{\textrm{Expert}}).$
(44)

Now,

$\displaystyle
p(\omega|\mathcal{D}_{\textrm{Expert}},\Theta_{\textrm{Expert}})$
$\displaystyle=\frac{p(\mathcal{D}_{\textrm{Expert}}|\omega,\Theta_{\textrm{Expert}})p(\omega)}{\int
p(\mathcal{D}_{\textrm{Expert}}|\omega,\Theta_{\textrm{Expert}})dP(\omega)},$
(45)
$\displaystyle=\frac{\prod_{i=1}^{N_{\textrm{Expert}}}p(\tau_{i}|\omega,\theta_{i})p(\omega)}{\int\prod_{i=1}^{N_{\textrm{Expert}}}p(\tau_{i}|\omega,\theta_{i})dP(\omega)},$
(46)
$\displaystyle=\frac{\prod_{i=1}^{N_{\textrm{Expert}}}\prod_{j=0}^{H_{i}-1}p(s_{j+1}|s_{j},a_{j},\theta_{i})p(a_{j}|s_{j},\omega,\theta_{i})p(\omega)}{\int\prod_{i=1}^{N_{\textrm{Expert}}}\prod_{j=0}^{H_{i}-1}p(s_{j+1}|s_{j},a_{j},\theta_{i})p(a_{j}|s_{j},\omega,\theta_{i})dP(\omega)},$
(47)
$\displaystyle=\frac{\prod_{i=1}^{N_{\textrm{Expert}}}\prod_{j=0}^{H_{i}-1}p(a_{j}|s_{j},\omega,\theta_{i})p(\omega)}{\int\prod_{i=1}^{N_{\textrm{Expert}}}\prod_{j=0}^{H_{i}-1}p(a_{j}|s_{j},\omega,\theta_{i})dP(\omega)},$
(48)

where we have used the fact that each $p(s_{j+1}|s_{j},a_{j},\theta_{i})$ has
no dependence on $\omega$, and so will cancel in the fraction when deriving
the final line. Now, substituting for the definition of the likelihood:

$\displaystyle p(\omega$
$\displaystyle|\mathcal{D}_{\textrm{Expert}},\Theta_{\textrm{Expert}})$ (49)
$\displaystyle\propto\underbrace{\exp\left(\sum_{i=1}^{N_{\textrm{Expert}}}\sum_{j=0}^{H_{i}-1}\left(\frac{\Psi^{\star}(s_{j},a_{j},\theta_{i})^{\top}\omega}{\alpha}-\log
z(s_{j},\theta_{i},\omega)\right)\right)}_{\textrm{Likelihood}}\underbrace{\exp\left(-\frac{\lVert\omega-\omega_{0}\rVert^{2}}{2\sigma_{0}^{2}}\right)}_{\textrm{Prior}},$
(50)
$\displaystyle=\exp\left(\sum_{i=1}^{N_{\textrm{Expert}}}\sum_{j=0}^{H_{i}-1}\left(\frac{\Psi^{\star}(s_{j},a_{j},\theta_{i})^{\top}\omega}{\alpha}-\log
z(s_{j},\theta_{i},\omega)\right)-\frac{\lVert\omega-\omega_{0}\rVert^{2}}{2\sigma_{0}^{2}}\right),$
(51)

hence:

$\displaystyle\nabla_{\omega}$ $\displaystyle\log
p(\omega|\mathcal{D}_{\textrm{Expert}},\Theta_{\textrm{Expert}})$ (52)
$\displaystyle=\nabla_{\omega}\left(\sum_{i=1}^{N_{\textrm{Expert}}}\sum_{j=0}^{H_{i}-1}\left(\frac{\Psi^{\star}(s_{j},a_{j},\theta_{i})^{\top}\omega}{\alpha}-\log
z(s_{j},\theta_{i},\omega)\right)-\frac{\lVert\omega-\omega_{0}\rVert^{2}}{2\sigma_{0}^{2}}\right),$
(53) $\displaystyle\implies$ $\displaystyle
g(\omega;\mathcal{D}_{\textrm{Expert}})$ (54) $\displaystyle=$
$\displaystyle\int_{\Theta^{N_{\textrm{Expert}}}}\nabla_{\omega}\Bigg{(}\sum_{i=1}^{N_{\textrm{Expert}}}\sum_{j=0}^{H_{i}-1}\left(\frac{\Psi^{\star}(s_{j},a_{j},\theta_{i})^{\top}\omega}{\alpha}-\log
z(s_{j},\theta_{i},\omega)\right)$ (55)
$\displaystyle\quad-\frac{\lVert\omega-\omega_{0}\rVert^{2}}{2\sigma_{0}^{2}}\Bigg{)}dP(\Theta_{\textrm{Expert}}|\omega,\mathcal{D}_{\textrm{Expert}}),$
(56) $\displaystyle=$
$\displaystyle\sum_{i=1}^{N_{\textrm{Expert}}}\mathbb{E}_{\theta_{i}\sim
p(\theta_{i}|\omega,\tau_{i})}\left[\nabla_{\omega}\left(\sum_{j=0}^{H_{i}-1}\left(\frac{\Psi^{\star}(s_{j},a_{j},\theta_{i})^{\top}\omega}{\alpha}-\log
z(s_{j},\theta_{i},\omega)\right)-\frac{\lVert\omega-\omega_{0}\rVert^{2}}{2\sigma_{0}^{2}}\right)\right].$
(57)

Now, applying Lemma 4.1 and multiplying by $\alpha$ yields:

$\displaystyle\alpha g(\omega;\mathcal{D}_{\textrm{Expert}})=$
$\displaystyle\sum_{i=1}^{N_{\textrm{Expert}}}\sum_{j=0}^{H_{i}-1}\mathbb{E}_{\theta_{i}\sim
p(\theta_{i}|\omega,\mathcal{D}_{\textrm{Expert}})}\left[\left(\Psi^{\star}(s_{j},a_{j},\theta_{i})-\mathbb{E}_{a_{i}\sim
p(a_{i}|s_{j},\omega,\theta_{i})}\left[\Psi^{\star}(s_{j},a_{i},\theta_{i})\right]\right)\right]$
(58) $\displaystyle\quad-\frac{(\omega-\omega_{0})}{\varsigma_{0}^{2}},$ (59)

as required ∎

###  F.1 Derivation of BRL Objective

Starting from the Bayesian RL objective:

$\displaystyle J^{\pi}_{\textrm{Bayes}}$
$\displaystyle=\mathbb{E}_{\tau_{\infty}\sim
p^{\pi}_{\infty}(\tau_{\infty})}\left[\sum_{i=0}^{\infty}\gamma^{i}r_{i}\right],$
(60) $\displaystyle=\mathbb{E}_{h_{\infty}\sim
p^{\pi}_{\infty}(h_{\infty})}\left[\sum_{i=0}^{\infty}\gamma^{i}\mathbb{E}_{r_{i}\sim
p(r_{i}|h_{i},a_{i})}\left[r_{i}\right]\right].$ (61)

Now,

$\displaystyle\mathbb{E}_{r_{i}\sim p(r_{i}|h_{i},a_{i})}\left[r_{i}\right]$
$\displaystyle=\mathbb{E}_{r_{i}\sim
p^{\textrm{IRL+COE}}_{\textrm{Bayes}}(r_{i}|s_{i},a_{i})}\left[r_{i}\right],$
(62) $\displaystyle=\begin{cases}\mathbb{E}_{k\sim
p(k)}\left[\mathbb{E}_{r_{i}\sim
p(r_{i}|s_{i},a_{i},k)}\left[r_{i}\right]\right],&s,a\in\mathcal{S}_{\textrm{COE}}\times\mathcal{A}_{\textrm{COE}},\\\
\mathbb{E}_{\omega\sim p(\omega)}\left[\mathbb{E}_{r_{i}\sim
p(r_{i}|s_{i},a_{i},\omega)}\left[r_{i}\right]\right],&\textrm{otherwise},\end{cases}$
(63) $\displaystyle=\begin{cases}\mathbb{E}_{k\sim
p(k)}\left[kr_{\textrm{max}}\right],&s,a\in\mathcal{S}_{\textrm{COE}}\times\mathcal{A}_{\textrm{COE}},\\\
\mathbb{E}_{\omega\sim
p(\omega)}\left[\nu(s_{i},a_{i})^{\top}\omega\right],&\textrm{otherwise},\end{cases}$
(64)
$\displaystyle=\begin{cases}k^{\star}r_{\textrm{max}},&s,a\in\mathcal{S}_{\textrm{COE}}\times\mathcal{A}_{\textrm{COE}},\\\
\nu(s_{i},a_{i})^{\top}\omega^{\star}_{\textrm{Bayes}},&\textrm{otherwise},\end{cases}$
(65) $\displaystyle=r^{\textrm{IRL+COE}}_{\textrm{Bayes}}(s_{i},a_{i}),$ (66)

hence:

$\displaystyle J^{\pi}_{\textrm{Bayes}}$
$\displaystyle=\mathbb{E}_{h_{\infty}\sim
p^{\pi}_{\infty}(h_{\infty})}\left[\sum_{i=0}^{\infty}\gamma^{i}r^{\textrm{IRL+COE}}_{\textrm{Bayes}}(s_{i},a_{i})\right],$
(67) $\displaystyle=\mathbb{E}_{\theta_{\textrm{test}}\sim
p(\theta_{\textrm{test}})}\left[\mathbb{E}_{h_{\infty}\sim
p^{\pi}_{\infty}(h_{\infty}|\theta)}\left[\sum_{i=0}^{\infty}\gamma^{i}r^{\textrm{IRL+COE}}_{\textrm{Bayes}}(s_{i},a_{i})\right]\right],$
(68)

as required.

##  Appendix G The Role of Temperature in IRL

Figure 7: Space of Three State MDPs

We now provide an intuitive example to demonstrate how the prior influences
the expert data in our Bayesian formulation. Consider the space of three state
MDPs in Figure 7. The agent has a set of two possible actions
$\mathcal{A}=\\{a_{1},a_{2}\\}$. For $\theta=\theta_{1}$, the agent
transitions to state $s_{1}$ deterministically after selecting $a_{1}$ or
$s_{2}$ after selecting $a_{2}$. For $\theta=\theta_{2}$, the actions are
reversed. In both MDPs, the initial state is $s_{0}$ and state $s_{1}$ or
$s_{2}$ are terminal. The reward function depends only on states and is
$r(s_{0})=0$, $r(s_{1})=1$, $r(s_{2})=-1$. Let $\mathbb{I}(s)\in\\{0,1\\}^{3}$
denote the indicator feature vector where $I(s_{n})$ is a one-hot vector where
the $n$th element is 1, e.g. $\mathbb{I}(s_{1})=(0,1,0)^{\top}$. We represent
this reward function using the linear form
$r(s)=\mathbb{I}(s)^{\top}\omega^{\star}$ where
$\omega^{\star}=(0,1,-1)^{\top}$. Finally, MDPs are allocated using a uniform
prior.

Expert data will always consist of trajectories that transition from state
$s_{0}$ to state $s_{1}$. For our feature vector, we can derive the
corresponding expert successor feature representation analytically:

$\displaystyle\Psi^{\star}(s_{0},a_{1},\theta_{1})=\mathbb{I}(s_{0})+\frac{\gamma}{1-\gamma}\mathbb{I}(s_{1}),\
\Psi^{\star}(s_{0},a_{2},\theta_{1})=\mathbb{I}(s_{0})+\frac{\gamma}{1-\gamma}\mathbb{I}(s_{2}),$
(69)
$\displaystyle\Psi^{\star}(s_{0},a_{2},\theta_{2})=\mathbb{I}(s_{0})+\frac{\gamma}{1-\gamma}\mathbb{I}(s_{1}),\
\Psi^{\star}(s_{0},a_{1},\theta_{2})=\mathbb{I}(s_{0})+\frac{\gamma}{1-\gamma}\mathbb{I}(s_{2}),$
(70)
$\displaystyle\Psi^{\star}(s_{1},\cdot,\cdot)=\frac{1}{1-\gamma}\mathbb{I}(s_{1}),\
\Psi^{\star}(s_{2},\cdot,\cdot)=\frac{1}{1-\gamma}\mathbb{I}(s_{2}).$ (71)

Using the gradient update in Equation 4, we see that updates in state $s_{0}$
will initially draw actions equally from $p(a_{i}|s_{0},\omega,\cdot)$. This
yields an initial gradient update of:

$\displaystyle g_{\omega}^{0}$
$\displaystyle=\omega_{0}+\eta_{\omega}^{0}\frac{N_{\textrm{Expert}}{H_{i}}}{2}\left(\Psi^{\star}(s_{0},a_{1},\theta_{1})-\frac{1}{2}\left[\Psi^{\star}(s_{0},a_{1},\theta_{1})+\Psi^{\star}(s_{0},a_{2},\theta_{1})\right]\right)+$
(72)
$\displaystyle\quad\quad\eta_{\omega}\frac{N_{\textrm{Expert}}{H_{i}}}{2}\left(\Psi^{\star}(s_{0},a_{2},\theta_{2})-\frac{1}{2}\left[\Psi^{\star}(s_{0},a_{2},\theta_{2})+\Psi^{\star}(s_{0},a_{1},\theta_{2})\right]\right),$
(73)
$\displaystyle=\omega_{0}+\frac{\gamma\eta_{\omega}^{0}N_{\textrm{Expert}}H_{i}}{2(1-\gamma)}\left[\mathbb{I}(s_{1})-\mathbb{I}(s_{2})\right].$
(74)

For exposition, assume that there is no prior preference between
$\omega_{0}^{1}$ and $\omega_{0}^{2}$, and that
$\omega_{0}^{1}=\omega_{0}^{2}=0$. We see that the initial gradient will
update these values to:

$\displaystyle\omega_{1}^{1}=\frac{\gamma\eta_{\omega}^{0}N_{\textrm{Expert}}H_{i}}{2(1-\gamma)},\quad\omega_{1}^{2}=-\frac{\gamma\eta_{\omega}^{0}N_{\textrm{Expert}}H_{i}}{2(1-\gamma)}.$
(75)

We now consider the regime where the temperature parameter tends to zero
$\alpha\rightarrow 0$. For all future updates, as $\omega_{1}>\omega_{2}$, the
model expert policy will tend towards a deterministic function that picks the
action leading to state $s_{1}$:
$p(a|s_{0},\omega,\theta_{i})=\delta(a=a_{i})$. This means that all future
updates $k\geq 1$ from state $s_{0}$ will pull the reward parametrization back
to the prior value:

$\displaystyle
g_{\omega}^{k}=\omega_{k}-\frac{\eta_{\omega}^{k}(\omega_{k}-\omega_{0})}{\varsigma_{0}^{2}}.$
(76)

In the regime where the temperature parameter tends to infinity
$\alpha\rightarrow\infty$, the model expert policy will remain uniform over
all actions. Under this assumption, all future updates $k\geq 1$ from state
$s_{0}$ will continue to increase the value of $\omega_{1}$ and decrease the
value of $\omega_{2}$, whilst pulling $\omega_{k}$ back towards the prior in
accordance with $\varsigma_{0}^{2}$.

$\displaystyle g_{\omega}^{k}$
$\displaystyle=\omega_{k}+\frac{\gamma\eta_{\omega}^{k}N_{\textrm{Expert}}H_{i}}{2(1-\gamma)}\left[\mathbb{I}(s_{1})-\mathbb{I}(s_{2})\right]-\frac{\eta_{\omega}^{k}(\omega_{k}-\omega_{0})}{\varsigma_{0}^{2}}.$
(77)

When used in practice, we select $\alpha$ to be between $0$ and $\infty$. Our
example reveals that the smaller the temperature parameter, the less the
updates will separate values of reward for states that the expert visited vs
that the expert could have visited. Conversely, when $\alpha$ is very large,
this difference will grow and can only be counteracted by the prior variance
$\varsigma_{0}$.

##  Appendix H Experiments

###  H.1 Computer resources

The experiments were run on servers with eight recent NVIDIA GPUs (3080,
A4500, or A5000). The random seeds were run in parallel. The experiments for
the gridworld were the longest and took approximately an hour to run.

###  H.2 Latent Inference Investigation Details

Figure 8:  A demonstration of the necessity for latent inference. On the left,
we repeat the visualization of the CMDP used in the experiments for
convenience. On the right, we show the rewards learned with and without latent
inference compared to the ground truth rewards. The reward vectors are
normalized to unit norm. The y-axis is linearly scaled.

Figure 8 shows the learned rewards in the latent inference experiment
described in Section 5.2. The learned rewards are somewhat hard to interpret
because in addition to the scaling and shifting information being lost in IRL,
there is no pressure for the algorithms to keep the reward vectors sparse.
Instead, they are only trying to optimize Equation 4. Nevertheless, the
rewards learned with latent inference have the same ordering between the
critical states $s_{1}$, $s_{2}$, and $s_{3}$ as the ground truth reward. In
contrast, the rewards learned with naive IRL do not differentiate between
$s_{1}$ and $s_{2}$, which results in the learned policies not taking the path
through $s_{1}$ even when it is available.

###  H.3 Tiger Treasure Details

The Tiger-Treasure problem is displayed in Figure 1 and described in Section
3.1. The agent starts in state $S_{0}$ and can either listen or open a door.
The listening action transitions the agent to a ’hint’ state revealing with
probability $p=0.85$ the location of the Tiger. After the agent opens a door,
it is either roared by the tiger or collects the treasure and arrives in the
terminal state $S_{T}$. For learning the contextual IRL reward, we use the
hyperparameters presented in Table 2 with a default value of
$\omega_{0}=(0,0,0,0,0,0)$ in the Laplace approximation Equation 4. The reward
prior is introduced after having rescaled the IRL reward between
$r_{\text{min}}=-100$ and $r_{\text{max}}=10$. To enhance the training
process, we normalize the rewards before training the DQN policies. The
hyperparameters presented in Table 3.

Parameter | Value
---|---
Number of parallel environments | 500
Number of steps per rollout | 50
Number of updates | 5000
$\epsilon$-greedy $\epsilon$ | 0.5
Maximum gradient norm | 0.5
Discount $\gamma$ | 0.99
$\alpha$ | 0.01
$\varsigma^{2}_{0}$ | 100.0
Target SF update rate | 50
SF learning rate | $1\times 10^{-3}$
Reward learning rate | $1\times 10^{-2}$
Replay buffer size | 50000
Batch size (trajectories) | 500
$r_{\text{max}}$ | 10.0
$r_{\text{min}}$ | -100.0
Table 2: IRL Hyperparameters for the Tiger Treasure problem. Parameter | Value
---|---
Number of parallel environments | 16
Number of steps per rollout | 50
Total number of updates | 20000
Learning rate | $1\times 10^{-4}$
Discount $\gamma$ | 0.99
Target network update rate | 1
Replay buffer size (number of full trajectories) | 200000
Batch size (number of full trajectories) | 100
Starting value for $\epsilon$ | 1.0
Final value for $\epsilon$ | 0.05
Fraction of updates after $\epsilon$ schedule finishes | 0.5
Table 3: DQN Hyperparameters for the Tiger Treasure problem.

###  H.4 Tiger Treasure Maze Details

The agent starts in the middle of the two doors. It can choose to move in the
cardinal directions or listen. After the agent opens a door, it will either
collect a treasure or be roared at by a tiger for one timestep. Then,
depending on what happened, it gets transported to the nearest grid cell to
the $Tiger$ or $Gold$. The expert takes the shortest path to $Gold$ and never
listens. The state of the agent is defined as the $X-Y$ coordinates of the
agent and an indicator variable, which indicates the result of the listening
action. After taking a listening action, the agent is transported to a state
with the same $X-Y$ coordinates but with the indicator showing the true value
of $\theta$. The indicator states have the same dynamics as the corresponding
normal states. The coordinates and the indicator are encoded as one-hot
vectors. The reward feature $\nu$ is the full table of all states visitable by
the agent. Since the dynamics are deterministic given $\theta$, the inference
model labels the trajectories with the true $\theta$ if it is revealed on the
trajectory.

The hyperparameters used for BIG in the maze experiment are shown in Table 4.
The hyperparameters for DQN in the maze experiment are shown in Table 5. A
recurrent neural network is used for the policy architecture. A separate MLP
network is used for implementing the critic.

Parameter | Value
---|---
Number of parallel environments | 16
Number of steps per rollout | 40
Number of updates | 20000
$\epsilon$-greedy $\epsilon$ | 0.5
Maximum gradient norm | 0.5
Discount $\gamma$ | 0.99
$\alpha$ | 0.01
$\varsigma^{2}_{0}$ | 1.0
Target SF update rate | 50
SF learning rate | $3\times 10^{-4}$
Reward learning rate | $3\times 10^{-3}$
Replay buffer size (number of full trajectories) | 10000
Batch size (number of full trajectories) | 100
$k$ | 0.01
$r_{\text{max}}$ | 1.0
$r_{\text{min}}$ | -0.05
Table 4: IRL Hyperparameters for the Maze problem. Parameter | Value
---|---
Number of parallel environments | 16
Number of steps per rollout | 40
Total number of updates | 100000
Learning rate | $1\times 10^{-4}$
Discount $\gamma$ | 0.99
Target network update rate | 1
Replay buffer size (number of full trajectories) | 10000
Batch size (number of full trajectories) | 100
Starting value for $\epsilon$ | 1.0
Final value for $\epsilon$ | 0.05
Fraction of updates after $\epsilon$ schedule finishes | 0.5
Table 5: DQN Hyperparameters for the Maze problem.