A Novel Variational Lower Bound For
Inverse Reinforcement Learning

Inspired by the graphical model for forward RL as shown in Levine (2018) and a general lack of such modeling of IRL, we aim to fill this gap in this section. Recall that we are now interested in learning the reward function given the state and action trajectories from the expert. Therefore, it is necessary to embed the reward function explicitly in the graphical model, as shown in Figure 1. The reward value $r_{t}$ is conditioned on state $\mathbf{s_{t}}$ and action $\mathbf{a_{t}}$ respectively denoting the state and action in the trajectory. We introduce an optimality random variable $\mathbf{\mathcal{O}_{t}}$, whose value is conditioned on the reward value $r_{t}$. The optimality follows the definition of Levine (2018) in that the optimality variable is a binary random variable, where $\mathcal{O}_{t}=1$ denoting that the observed action at time step $t$ is optimal given the state, and $\mathcal{O}_{t}=0$ denotes that it is not optimal. We will discuss the reward function $\mathbf{\mathcal{R}}$ in more detail in Section 2.3.

The choice of a probabilistic graphical model to represent the relationships between the state, action, reward, and optimality variables in IRL is motivated by the desire to capture the complex dependencies between these variables in a lucid manner. Indeed, the graphical model has a natural property when it comes to analyzing optimality. Specifically, the graphical model represents the conditional independence of the optimality from the state and action variables, given the reward variable. In other words, when given the state-action pair $(s_{t},a_{t})$, we sum over the reward value $r_{t}$ to get the marginal conditional distribution of the optimality $\mathcal{O}_{t}$, that is $p(\mathcal{O}_{t}|s_{t},a_{t})$. However, if the reward value $r_{t}$ is given, then the state-action pair is of course unnecessary in order to get the marginal distribution of the optimality $\mathcal{O}_{t}$, that is $p(\mathcal{O}_{t}|r_{t})$.

Given the conditional independence in a probabilistic graphical model, which is similar to a Bayesian network, we may write:

which leads to the following marginal:

We use the graphical models of Fig. 1 to formulate the log-likelihood of the observed trajectories. ¹¹1In practice, using full trajectories to estimate the reward function in IRL may lead to high variance especially when the trajectories are long. On the other hand, using single state-action pairs to infer the reward function can lead to more stable and efficient estimates, because each observed pair provides a separate estimate of the reward function that can be combined to reduce variance. Here, we are aided by Eqs. 1 and 2 in our derivation, which we detail below.

Here, we use variational distribution $q(\mathcal{O}_{t}|r_{t})$ to approximate the true distribution $p(\mathcal{O}_{t}|s_{t},a_{t})$. We may think of $q(\mathcal{O}_{t}|r_{t})$ as an approximation of the true distribution of optimality given the trajectories, $p(\mathcal{O}_{t}|s_{t},a_{t})$ and we will discuss in more detail in Section 2.3. We may rewrite the equation above by introducing the expectation $\mathbb{E}_{q(\mathcal{O}_{t}|r_{t})}$ and applying Jensen’s inequality:

We end up with the following Evidence Lower BOund (ELBO):

The motivation for using $q(\mathcal{O}_{t}|r_{t})$ to approximate $p(\mathcal{O}_{t}|s_{t},a_{t})$ is quite intuitive that both state-action pair and reward value can explain the optimality. An interpretation of $p(\mathcal{O}_{t}|s_{t},a_{t})$ is that the probability can be viewed as the confidence with which the state-action pair comes from the expert. Analogously, $q(\mathcal{O}_{t}|r_{t})$ is the confidence with which the reward value $r_{t}$ belongs to the expert. Additionally, to optimize the reward function, we have to incorporate the reward function into the log-likelihood of trajectories.

Note that if we take the partial derivative of Eq.  4 with respect to $r_{t}$, we obtain,

because the remaining terms in equation 3 are constant w.r.t. $r_{t}$. The derivation from first principles above leads to the following theorem.

Next, we discuss the approximation of the distribution over optimality under the framework of the probabilistic graphical model and offer insights.

Given a state-action pair $(s_{t},a_{t})$ obtained from some trajectory $\tau$, $p(\mathcal{O}_{t}|s_{t},a_{t})$ informs whether the action $a_{t}$ is optimal at state $s_{t}$, or in other words, whether the state-action pair comes from the expert’s trajectory. This term represents the true distribution in the reverse KL divergence contained in the lower bound. To make this classification, we may simply use binary logistic regression, $C_{\bm{\theta}}$. In algorithms such as GAIL and AIRL, the input to the classifier consists of trajectories from both the expert and the learner. The classifier utilizes these trajectories as input to make predictions about the optimality of the state-action pair.

Similarly, given reward value $r_{t}$, the approximation $q(\mathcal{O}_{t}|r_{t})$ informs whether the reward value leads to optimal behavior, i.e., whether it is induced by the expert or not.

Recall that the reward value $r_{t}$ is the feedback from the environment. Therefore, we propose this reward value as our first approach to estimate the optimality as defined in Eq. 7.

Note that the left hand side of Eq. 7 is conditioned on reward value $r_{t}$, which is distributed according to the distribution $\mathcal{R}(s_{t},a_{t})$. Thus, we may apply a parameterization trick to sample a specific reward value $r_{t}$ without losing track of the gradient propagation under the probabilistic graphical model. The dashed line in Fig. 1 denotes the reparameterization trick and $r_{t}$ represents the sampled specific reward value from the distribution of $\mathcal{R}(s_{t},a_{t})$. To illustrate this, consider the simplistic case where the reward value distribution is a univariate Gaussian: $\mathcal{R}(s_{t},a_{t})=\mathcal{N}(\mu,\sigma^{2})$ and let $r_{t}\sim\mathcal{R}(s_{t},a_{t})$. In this case, a valid reparameterization is $r_{t}=\mu+\sigma\epsilon$ (reward values are distributed around the mean), where $\epsilon$ is an auxiliary noise variable, $\epsilon\sim\mathcal{N}(0,1)$.

However, it is not sufficient to use the reward values to represent the optimality because, in an optimal trajectory, not every state-action pair has the highest reward value at each time step; often we may perform an action that obtains longer term gain while forgoing greedy rewards. Hence, a more robust way of computing the optimality given the distribution over reward value is needed. Here, the advantage function is a candidate for the solution:

In practice, we can use actor-critic-based policy to retrieve the value function estimation, such as PPO, SAC, or TD3. . As the reward value is a component of the advantage function, we can continue to keep track of the reward function. Therefore, the optimality of reward value can be expressed as:

where $r_{t}\sim\mathcal{R}(s_{t},a_{t})$ and $\sigma$ is the Sigmoid function to normalize the advantage function.

Note that the approximation error is bounded by the variance of $r_{t}$ instead of $\mathcal{O}_{t}$ and the motivation of using reward $r_{t}$ to approximate state-action pair $(s_{t},a_{t})$ comes from the Figure.1, in which holds the connectivity that node $r_{t}$ is connected to node $s_{t},a_{t}$ and node $\mathcal{O}_{t}$. Without the connectivity, even if the variance of $r_{t}$ is zero, the approximation error cannot be bounded. More detail can be found in Appendix A.2. It is also the connectivity that holds so that we can derive Eq. 1 and Eq. 2.

From Theorem 2, we know that we can use $q(\mathcal{O}_{t}\mid\mathbb{E}[r_{t}])$ to approximate $p(\mathcal{O}_{t}|s_{t},a_{t})$ with a bounded approximation error w.r.t $\mathrm{Var}(r_{t})$ . Additionally, if the variance of $r_{t}$ is small enough, then we can use $r_{t}$ to estimate $\mathbb{E}[r_{t}]$. Subsequently, the objective function of VLB-IRL is defined in the following:

where $\sigma$ is the Sigmoid function to normalize the advantage function and $\lambda$ is the hyperparameter to control the contribution of variance in the objective function.

In this section, we discuss the reward ambiguity problem. As deduced in Ho & Ermon (2016), IRL is a dual of an occupancy measure matching problem and the induced optimal policy is the primal optimum. In the following, we further draw the conclusion that IRL is a dual of the optimality matching problem and the reward function induced by the optimality best explains the expert $\pi^{E}$.

From Theorem 3, we are guaranteed to have a reward function $\mathcal{R}$ that best explains the expert policy. Theorem 3 also offers the benefit that we are guaranteed to avoid a degenerate reward function as the solution by optimizing Eq. 5 if $p(\mathcal{O}_{t}|s_{t},a_{t})$ is maximal.

Recall that the objective function of VLB-IRL consists of a reverse KL divergence between $q(\mathcal{O}_{t}|r_{t})$ and $p(\mathcal{O}_{t}|s_{t},a_{t})$. Compared to the mean-seeking behavior induced by minimizing forward KL divergence, minimizing reverse KL divergence has the behavior that it is mode-seeking, which tends to be more robust to overfitting and can provide better estimates of the true posterior distribution if the posterior distribution is noisy. This is beneficial for IRL because trajectories could be noisy and may contain suboptimal actions, which challenges the learning. By emphasizing the modes of the data distribution, the reverse KL divergence helps identify the most likely explanations for the observed behavior, even in the presence of noise or uncertainty.

Compared to existing state-of-the-art algorithms, most of them use a single neural network to estimate the reward function. However, in VLB-IRL, we use two separate neural networks to update the reward function. The first neural network is the classifier, defined in Eq. 7. The second neural network is the approximation optimality, defined in Eq. 7 and Eq. 9. The architecture of two separate neural networks has a natural property that is robust to overfitting and has better generalization. It is essential for the IRL algorithm with the presence of noise in expert trajectories.

Another major difference is that most prior algorithms focus on learning identical state-action marginal distributions and consequently end up learning the expert’s noisy state-action representation as well. However, in VLB-IRL, since the true distribution $p(\mathcal{O}_{t}|s_{t},a_{t})$ represents the optimality conditioned on the state-action pair, it has the ability to distinguish and generalize the noisy trajectories.

As part of our generalizability analysis, we consider both continuous and discrete control environments from the Mujoco library and use twin-delayed DDPG (TD3) (Fujimoto et al., 2018) and proximal policy optimization (PPO) (Schulman et al., 2017) for continuous and discrete domains respectively ²²2The same RL algorithms are used for expert trajectory generation.. We readily obtain the expertly trained RL models from the popular RL library RL-Baselines3 Zoo (Raffin, 2020) and generate several expert trajectories employing them. Using the generated expert trajectories, each method was trained with 5 random seeds and for each seed, the model with the highest return was saved. At the end of the training, we evaluate these models on 50 test episodes. We use MLP for the reward function approximation. The hyperparameters and other relevant details needed for reproducibility are provided in the Appendix B. and our code is openly available on GitHub³³3Double blind review note: we will share our code once our paper gets accepted.

Performance on optimal trajectories   For our first set of experiments, we directly use the optimal expert trajectories obtained from the aforementioned pretrained models. Table 1 shows the summary of results obtained and the statistically best average reward accrued per domain is highlighted in bold. Our method VLB-IRL improves upon the baselines indicating that VLB-IRL’s learned reward function successfully explains the expert’s policy.

Learning performance   We find that VLB-IRL continues to perform well in a limited data regime. The empirical result demonstrates that VLB-IRL exhibits fairly good sample efficiency. With increasing amount of data, the standard deviation of the return from the learned policy becomes lower as shown in Figure 2, indicating better stability in the learned policy, which is a direct result of a better reward function. To further analyze the training progress of VLB-IRL, we use Inverse Learning Error (ILE) from Arora & Doshi (2021a), as $||V^{\pi_{E}}-V^{\pi_{L}}||_{p}$ where $\pi_{E}$ is the expert policy and $\pi_{L}$ is the learned policy. Here we use the $L2$ norm and the result is presented in Figure 2. The decreasing ILE shows that both the learned reward function is getting closer to the true reward function and the learned policy is more similar to the expert policy.

Performance on noisy trajectories   In order to further test VLB-IRL’s generalizability, in our second set of experiments, we provide noisy trajectories and examine if the learned reward function is robust enough to learn the noisy expert’s preferences and perhaps outperform them. The noisy trajectories are generated by using A2C for discrete environments and PPO for continuous environments in the library RL-Baselines3 Zoo. Compared to PPO for discrete environments and TD3 for continuous environments, they are less optimal and contain noisy actions in the trajectories. Compare the expert performance in Table 1 and Table 2 for detail. The results presented in Table 2 show that VLB-IRL’s accrues comparable or higher reward on average as compared to the noisy expert.

For Hopper, HalfCheetah, and Ant environments, the VLB-IRL’s learned policy outperforms the noisy expert by $15.6\%$, $14.4\%$, and $11.7\%$ respectively. However, previous IRL techniques fail to generalize well when noisy trajectories are provided and we suspect that they fall short due to their optimization objective and divergence metric listed in Table 4.

To further examine the performance of VLB-IRL, we use Assistive Gym (Erickson et al., 2020), a more realistic, robot-centered, and goal-based environment. For our experiments, we use environments with an active robot and a static human, such as FeedingSaywer, BedBathingSawyer, and ScratchItchSawyer. A description of the environments can be found in Figure.3. We use SAC to train an expert policy and use the trained expert policy to generate 50 expert episodes. The detailed hyperparameter can be found in Appendix B.

The challenge in Assistive Gym compared to Mujoco benchmark is that it is a goal-based environment and the agent only receives a positive reward after the completion of the goal, unlike Mujoco where the agent can receive positive rewards at every timestep. This goal-based environment will generate much more noisy expert trajectories. Table 3 shows that VLB-IRL outperforms all other baselines and VLB-IRL is the only method that can successfully achieve the goal.