1. Formalism

Given

  • a set of states \(\mathbf{s} \in \mathcal{S}\) and actions \(\mathbf{a} \in \mathcal{A}\),
  • (sometimes) the transition probabilities \(p(\mathbf{s}'\vert\mathbf{s},\mathbf{a})\),
  • a set of trajectories \(\{\tau_i\}\) that were sampled from an (approximately) optimal policy \(\pi^*(\tau)\)

the goal of inverse reinforcement learning (IRL) is to learn the reward function \(r_\psi(\mathbf{s},\mathbf{a})\) which can then be used to recover the optimal policy \(\pi^*(\mathbf{a}\vert\mathbf{s})\).

2. Classical Approaches: Feature Matching and Maximum Margin

Most classical approaches assumed a linear reward function:

$$ r_\psi(\mathbf{s},\mathbf{a}) = \sum_{i=1}^F\psi_if_i(\mathbf{s},\mathbf{a}) = \Psi^T\mathbb{f}(\mathbf{s},\mathbf{a}) $$

for some known features \(f\). Note that \(\Psi = [\psi_1,\ldots,\psi_F]\) and \(\mathbb{f}(\mathbf{s},\mathbf{a}) = [f_i(\mathbf{s},\mathbf{a}),\ldots,f_F(\mathbf{s},\mathbf{a})]\). Let \(\pi^{r_\psi}\) denote the optimal policy for the reward function \(r_\psi\). The feature matching approach tries to pick \(\psi\) such that:

$$ \mathbb{E}_{\pi^{r_\psi}}[\mathbb{f}(\mathbf{s},\mathbf{a})] = \mathbb{E}_{\pi^*}[\mathbb{f}(\mathbf{s},\mathbf{a})] $$

i.e. it tries to match the expected value of the features under the optimal and the recovered policies. Note that the expectation on the RHS can be approximated with the optimal trajectories \(\{\tau_i\}\).

Alternatively, the maximum margin principle tries to choose a reward function such that the reward of the optimal policy is greater than the reward of any other policy by at least a margin of \(m\):

$$ \max_{\Psi,m} m\; \text{ s.t. } \Psi^T \mathbb{E}_{\pi^*}[\mathbb{f}(\mathbf{s},\mathbf{a})] \geq \max_{\pi \in \Pi, \pi \neq \pi^*} \Psi^T\mathbb{E}_{\pi}[\mathbb{f}(\mathbf{s},\mathbf{a})] + m $$

where \(\Pi\) is the set of all policies. We can rewrite this problem in a more tractable form in the same way as we do for support vector machines. Furthermore, note that we should also somehow take into account the fact that policies close to the optimal policy should only have a slightly less reward than it compared to other policies. The final objective function after taking these into account becomes:

$$ \max_{\Psi} \frac{1}{2}\vert\vert\Psi\vert\vert^2\; \text{ s.t. } \Psi^T \mathbb{E}_{\pi^*}[\mathbb{f}(\mathbf{s},\mathbf{a})] \geq \max_{\pi \in \Pi, \pi \neq \pi^*} \Psi^T\mathbb{E}_{\pi}[\mathbb{f}(\mathbf{s},\mathbf{a})] + D(\pi,\pi^*) $$

where \(D(\pi,\pi^*)\) is the difference between the feature expectations under \(\pi\) and \(\pi^*\).

3. The MaxEnt IRL Algorithm

Assume that our reward function is parameterized by \(\psi\). In the probabilistic graphical model discussed in the last lecture notes, we assumed that:

$$ p(\mathcal{O}_t\vert\mathbf{s}_t,\mathbf{a}_t;\psi) \propto \exp(r_\psi(\mathbf{s}_t,\mathbf{a}_t)) $$

where we also include \(\psi\) to emphasize that \(p(\mathcal{O}_t\vert\mathbf{s}_t,\mathbf{a}_t)\) depends on it. We also showed that:

$$ p(\tau\vert\mathcal{O}_{1:T};\psi) \propto \frac{p(\tau)\exp\left(\sum_{t=1}^T r(\mathbf{s}_t,\mathbf{a}_t)\right)}{p(\mathcal{O}_{1:T})} $$

If we know a trajectory that an expert chose (which means that that trajectory should have a very high reward), then one way of choosing \(\psi\) is by maximizing the likelihood of this trajectory (if our reward is correct then this trajectory should be very likely):

$$ \begin{align} \psi &= \text{argmax}_\psi \log p(\tau\vert\mathcal{O}_{1:T};\psi) \\ &= \text{argmax}_\psi \log \frac{p(\tau)\prod_{t=1}^Tp(\mathcal{O}_t\vert\mathbf{s}_t, \mathbf{a}_t)}{p(\mathcal{O}_{1:T})}\\ &= \text{argmax}_\psi \log \frac{p(\tau)\prod_{t=1}^Tp(\mathcal{O}_t\vert\mathbf{s}_t, \mathbf{a}_t)}{\int p(\mathcal{O}_{1:T},\tau)d\tau}\\ &= \text{argmax}_\psi \log \frac{p(\tau)\prod_{t=1}^Tp(\mathcal{O}_t\vert\mathcal{s}_t, \mathcal{a}_t)}{\int p(\tau)\prod_{t=1}^Tp(\mathcal{O}_t\vert\mathbf{s}_t, \mathbf{a}_t)d\tau}\\ &= \text{argmax}_\psi \log\frac{p(\tau)\exp\left(\sum_{t=1}^T r_\psi(\mathbf{s}_t,\mathbf{a}_t)\right)}{\int p(\tau)\exp\left(\sum_{t=1}^T r_\psi(\mathbf{s}_t,\mathbf{a}_t)\right)d\tau}\\ &= \text{argmax}_\psi\left[\log p(\tau) + \log\left(\exp\left(\sum_{t=1}^T r_\psi(\mathbf{s}_t,\mathbf{a}_t)\right)\right)-\log \left(\int p(\tau)\exp\left(\sum_{t=1}^T r_\psi(\mathbf{s}_t,\mathbf{a}_t)\right)d\tau\right)\right]\\ &= \text{argmax}_\psi\left[r_\psi(\tau)-\log \left(\int p(\tau)\exp\left( r_\psi(\tau)\right)d\tau\right)\right] \end{align} $$

The third last step uses the fact that we assume \(p(\mathcal{O}_t\vert\mathbf{s}_t,\mathbf{a}_t)\) to be proportional to \(\exp(r(\mathbf{s}_t,\mathbf{a}_t))\) (the constants of proportionality in the numerator and denominator cancel each other out). In the last step, we ignore \(p(\tau)\) since it does not depend on \(\psi\). We shall denote \(\int p(\tau)\exp\left( r_\psi(\tau)\right)d\tau\) with \(Z\). \(Z\) is often called as the IRL partition function. We can take the \(\text{argmax}\) over over multiple trajectories too. Our likelihood function will, thus, be given by (note that the second term is an integral over all possible trajectories and is thus constant):

$$ \mathcal{L}(\psi) = \frac{1}{N}\sum_{i=1}^N r_\psi(\tau^{(i)})-\log \left(\int p(\tau)\exp\left( r_\psi(\tau)\right)d\tau\right) $$

We need to maximize \(\mathcal{L}\) with respect to \(\psi\):

$$ \begin{align} \nabla_\psi\mathcal{L}(\psi) &= \frac{1}{N}\sum_{i=1}^N\nabla_\psi r_\psi(\tau^{(i)})-\frac{\nabla_\psi\int p(\tau)\exp\left( r_\psi(\tau)\right)d\tau}{\int p(\tau)\exp\left(r_\psi(\tau)\right)d\tau}\\ &= \frac{1}{N}\sum_{i=1}^N\nabla_\psi r_\psi(\tau^{(i)})-\frac{1}{Z}\int p(\tau)\nabla_\psi\exp\left(r_\psi(\tau)\right)d\tau\\ &= \frac{1}{N}\sum_{i=1}^N\nabla_\psi r_\psi(\tau)-\int \frac{1}{Z}p(\tau)\exp\left(r_\psi(\tau)\right)\nabla_\psi r_\psi(\tau)d\tau\\ &= \frac{1}{N}\sum_{i=1}^N\nabla_\psi r_\psi(\tau^{(i)})-\int \frac{1}{Z}p(\tau)\exp\left(r_\psi(\tau)\right)\nabla_\psi r_\psi(\tau)d\tau\\ &= \frac{1}{N}\sum_{i=1}^N\nabla_\psi r_\psi(\tau^{(i)})-\int p(\tau\vert\mathcal{O}_{1:T};\psi)\nabla_\psi r_\psi(\tau)d\tau\\ \end{align} $$

Let us rewrite this as:

$$ \nabla_\psi \mathcal{L}(\psi) = \mathbb{E}_{\tau \sim \pi^*(\tau)}\left[\nabla_\psi r_\psi(\tau)\right]-\mathbb{E}_{\tau \sim p(\tau\vert\mathcal{O}_{1:T};\psi)}\left[\nabla_\psi r_\psi(\tau)\right]\\ $$

Therefore, in essence, we are approximating the first term with samples drawn from the expert’s policy \(\pi^*\). The second term is the expectation under the optimal policy of our current reward function. To estimate this second term note that:

$$ \begin{align} \mathbb{E}_{\tau \sim p(\tau\vert\mathcal{O}_{1:T};\psi)}\left[\nabla_\psi r_\psi(\tau)\right] &= \mathbb{E}_{\tau \sim p(\tau\vert\mathcal{O}_{1:T};\psi)}\left[\nabla_\psi\sum_{t=1}^T r_\psi(\mathbf{s}_t,\mathbf{a}_t)\right]\\ &= \sum_{t=1}^T \mathbb{E}_{(\mathbf{s}_t,\mathbf{a}_t) \sim p(\mathbf{s}_t,\mathbf{a}_t\vert\mathcal{O}_{1:T};\psi)}\left[\nabla_\psi r_\psi(\mathbf{s}_t,\mathbf{a}_t)\right]\\ &= \sum_{t=1}^T\int\int p(\mathbf{s}_t,\mathbf{a}_t\vert\mathcal{O}_{1:T};\psi) \nabla_\psi r_\psi(\mathbf{s}_t,\mathbf{a}_t)d\mathbf{s}_td\mathbf{a}_t \end{align} $$

From the previous lecture notes we have:

$$ \begin{align} p(\mathbf{a}_t\vert\mathbf{s}_t,\mathcal{O}_{1:T};\psi) &= \frac{\beta_t(\mathbf{s}_t,\mathbf{a}_t)}{\beta_t(\mathbf{s}_t)}\\ p(\mathbf{s}_t\vert\mathcal{O}_{1:T};\psi) &\propto \alpha_t(\mathbf{s}_t)\beta(\mathbf{s}_t) \end{align} $$

Therefore:

$$ \begin{align} p(\mathbf{s}_t,\mathbf{a}_t\vert\mathcal{O}_{1:T};\psi) &= p(\mathbf{a}_t\vert\mathbf{s}_t,\mathcal{O}_{1:T};\psi)p(\mathbf{s}_t\vert\mathcal{O}_{1:T};\psi)\\ &\propto \beta_t(\mathbf{s}_t,\mathbf{a}_t)\alpha_t(\mathbf{s}_t) \end{align} $$

We can thus use the forward and backward messages to compute \(p(\mathbf{s}_t,\mathbf{a}_t\vert\mathcal{O}_{1:T};\psi)\). In the case of tabular environments (where there are a finite number of states and actions), the two integrals over \(\mathbf{s}_t\) and \(\mathbf{a}_t\) can be computed exactly.

The MaxEnt algorithm repeats the following until convergence:

  1. Given \(\psi\), compute backward messages \(\beta_t(\mathbf{s}_t,\mathbf{a}_t)\).
  2. Given \(\psi\), compute forward messages \(\alpha_t(\mathbf{s}_t)\).
  3. Compute \(p(\mathbf{s}_t,\mathbf{a}_t\vert\mathcal{O}_{1:T};\psi) \propto \beta_t(\mathbf{s}_t,\mathbf{a}_t)\alpha_t(\mathbf{s}_t)\).
  4. Evaluate \(\nabla_\psi \mathcal{L}\).
  5. Update \(\psi \leftarrow \psi + \eta\nabla_\psi\mathcal{L}\).

The reason why this is called as MaxEnt is that it can be shown that this algorithm optimizes:

$$ \max_\psi \mathcal{H}(\pi^{r_\psi}) \text{ s.t. } \mathbb{E}_{\pi^{r_\psi}}[\mathbb{f}] = \mathbb{E}_{\pi^*}[\mathbb{f}] $$

in the case of linear rewards \(r_\psi = \Psi^T\mathbb{f}\). Intuitively, this algorithm tries to match the features \(\mathbb{f}\) under the expert and the recovered policies while assuming as little as possible about the expert policy (as it also maximizes the entropy of the recovered policy).

4. Unknown Dynamics and Large State/Action Spaces

The problem with the MaxEnt algorithm is that to estimate the second expectation in \(\nabla_\psi \mathcal{L}\) it:

  1. requires a tabular environment, and
  2. assumes that we know the state-transition dynamics to calculate the forward and backward messages.

We thus require a better way to approximate this second expectation. One way to do so is to simply learn \(p(\mathbf{a}_t\vert\mathbf{s}_t,\mathcal{O}_{1:T};\psi)\) via any maximum-entropy RL algorithm. Recall that maximum-entropy RL algorithms maximize the following objective function:

$$ J(\theta) = \sum_t \mathbb{E}_{(\mathbf{s}_t,\mathbf{a}_t)\sim p(\mathbf{s}_t,\mathbf{a}_t)}\left[r_\psi(\mathbf{s}_t,\mathbf{a}_t)\right] + \mathbb{E}_{\mathbf{s}_t\sim p(\mathbf{s}_t)}\left[\mathcal{H}(p(\mathbf{a}_t\vert \mathbf{s}_t))\right] $$

We can then simply approximate \(\nabla_\psi \mathcal{L}\) by sampling trajectories from our learned distribution \(p(\mathbf{a}_t\vert\mathbf{s}_t,\mathcal{O}_{1:T};\psi)\):

$$ \nabla_\psi \mathcal{L} \approx \frac{1}{N}\sum_{i=1}^N\nabla_\psi r_\psi(\tau^{(i)})-\frac{1}{M}\sum_{j=1}^M \nabla_\psi r_\psi(\tau^{(j)})\\ $$

One issue with this is that each time \(\psi\) is updated we need relearn our policy \(p(\mathbf{a}_t\vert\mathbf{s}_t,\mathcal{O}_{1:T};\psi)\). This can be computationally very expensive. One solution to this is to use lazy policy optimization i.e. we only relearn our policy say, after \(K\) time steps. However, this would mean that the policy that we use to sample trajectories is different from the policy under which we want to take our expectation. To account for this difference we can use importance sampling. Let \(\pi\) denote the policy from which we sample our trajectories. As per our assumption, the policy under which we want to take the expectation is proportional to \(\exp(r_\psi(\mathbf{s}_t,\mathbf{a}_t))\). The weights \(w\) for a trajectory \(\tau\) for importance sampling are thus:

$$ \begin{align} w_j &=\frac{p(\tau\vert\mathcal{O}_{1:T})}{\bar p(\tau)}\\ &\propto \frac{p(\mathbf{s}_1)\prod_{t=1}^T p(\mathbf{s}_{t+1}\vert\mathbf{s}_t,\mathbf{a}_t)\exp(r_\psi(\mathbf{s}_t,\mathbf{a}_t))}{p(\mathbf{s}_1)\prod_{t=1}^Tp(\mathbf{s}_{t+1}\vert\mathbf{s}_t,\mathbf{a}_t)\pi(\mathbf{a}_t\vert\mathbf{s}_t)}\\ &= \frac{\prod_{t=1}^T\exp(r_\psi(\mathbf{s}_t,\mathbf{a}_t))}{\prod_{t=1}^T\pi(\mathbf{a}_t\vert\mathbf{s}_t)}\\ \end{align} $$

Therefore:

$$ \nabla_\psi \mathcal{L} \approx \frac{1}{N}\sum_{i=1}^N\nabla_\psi r_\psi(\tau^{(i)})-\frac{1}{M\sum_{j=1}^M w^{(j)}}\sum_{j=1}^M w^{(j)}\nabla_\psi r_\psi(\tau^{(j)})\\ $$

where we have also normalized the \(w^{(j)}\).

Interestingly enough, it turns that if we want to estimate the \(\mathbb{E}_{x\sim p(x)}[f(x)]\) through an expectation under some other distribution \(q(x)\) via importance sampling, then the optimal distribution to use for \(q(x)\) is:

$$ q(x) \propto \vert f(x)\vert p(x) $$

In our case the optimal \(\pi\) to use is therefore:

$$ \pi(\tau) \propto \exp(r_\psi(\tau)) $$

Recall that each time we update \(\pi\), we bring it closer to \(\exp(r_\psi(\tau))\) i.e. the optimal distribution.

5. Inverse RL as GANs

Reconsider the following:

$$ \nabla_\psi \mathcal{L}(\psi) = \mathbb{E}_{\tau \sim \pi^*(\tau)}\left[\nabla_\psi r_\psi(\tau)\right]-\mathbb{E}_{\tau \sim \pi_\theta(\tau)}\left[\hat w\nabla_\psi r_\psi(\tau)\right]\\ $$

where \(\hat w\) are the normalized importance sampling weights. This gradient can be interpreted as trying to make the trajectories sampled from \(\pi^*(\tau)\) more likely and the ones sampled from \(\pi_\theta\) less likely. In other words, it tries to update \(\psi\) such that trajectories sampled from \(\pi^*(\tau)\) have a higher reward and those sampled from \(\pi_\theta\) have a lower reward.

However, in the next step of the algorithm (as discussed in the previous section) we explicitly update our policy \(\pi_\theta\) using maximum entropy RL to maximize the expected reward under it. Viewed another way, we essentially update our policy so as to make the trajectories sampled from it harder to be distinguished from those sampled from \(\pi^*\).

This interpretation of inverse RL is quite similar to the idea behind Generative Adversarial Networks (GANs). Our policy \(\pi_\theta\) is thus the generator. It turns out that the optimal discriminator to use is:

$$ D_\psi(\tau) = \frac{\pi^*(\tau)}{\pi_\theta(\tau)+\pi^*(\tau)} $$

It can be seen that \(D_\psi\) is close to \(1\) when a trajectory has a high probability under \(\pi^*\) and is close to \(0\) when it has a high probability under \(\pi_\theta\). This is what we desire.

Note that:

$$ \begin{align} D_\psi(\tau) &= \frac{\pi^*(\tau)}{\pi_\theta(\tau)+\pi^*(\tau)}\\ &= \frac{p(\tau)\frac1Z\exp(r_\psi(\tau))}{p(\tau)\Pi_t\pi(\mathbf{a}_t\vert\mathbf{s}_t)+p(\tau)\frac1Z\exp(r_\psi(\tau))}\\ &= \frac{\frac1Z\exp(r_\psi(\tau))}{\Pi_t\pi(\mathbf{a}_t\vert\mathbf{s}_t)+\frac1Z\exp(r_\psi(\tau))}\\ \end{align} $$

We can use the standard GAN objective function to update our \(\psi\):

$$ \psi \leftarrow \text{argmax}_{\psi}\mathbb{E}_{\tau\sim\pi^*(\tau)}\left[\log D_\psi(\tau)\right] + \mathbb{E}_{\tau\sim\pi_\theta}\left[\log(1-D_\psi(\tau))\right] $$

It also turns out that we do not need to estimate \(Z\) but instead can directly optimize it using the same objective function above.