Paper: A Novel Variational Lower Bound For Inverse Reinforcement Learning (ICLR-2024)

Reviewer comment (Theory):
In RL and IRL, there is only 1 reward function, so $r_t$ and $r_{t'}$ should be related through the reward function parameters. Hence, the graphical model may be clearer if we include the reward function parameter node. Right now, it is unclear to me why we have the conditional independence of $\mathcal{O}_t$ from $a_t$ and $s_t$ given $r_t$. In Levine (2018), $\mathcal{O}_t$ is the optimality of the state-action pair given the reward function.

Author response:
1. Reward and optimality.

There is only 1 reward function which is expressed in the node $R$ and it is not conditioned on time $t$. In the revised paper, we will separate the $R$ as a single node connecting to all $r_t$ nodes. In Levine (2018), $\mathcal{O}\_t$ is conditioned $r_t$. However, in IRL setting, $r_t$ is unknown. Therefore, we need to know $\mathcal{O}\_t$ given $s_t, a_t$, i.e. $p(\mathcal{O}\_t|s_t,a_t)$. Given $p(\mathcal{O}\_t|s_t, a_t)$, we need to find the reward function $R$ such that $q(\mathcal{O}\_t|r_t)$ is approximating $p(\mathcal{O}\_t|s_t, a_t)$. Thus, $R$ best explains the expert trajectories.

We have updated our paper with a new graphical model, in which we separate out the reward function node and make it a deterministic node.

1. $R$ represents the reward function and it is now a deterministic node instead of a random node in the updated paper. There is only one reward function node $R$ in the graphical model. Every reward value $r_t$ comes from the same reward function $R$. And there is not any randomness for the reward function. However, the output of the reward function is a Gaussian distribution, denoted by $R(s_t, a_t)$. In conclusion, $R$ is the reward function. $R(s_t, a_t)$ is the reward value distribution given $s_t, a_t$ as input to the reward function $R$. $r_t$ is the sampled reward value from the reward distribution $R(s_t, a_t)$.

"I initially assumed that the reward $r_t$ at time $t$ should be computed based on the reward function $R$ and the state-action pair $s_t$ and $a_t$." This is correct. 

"If this assumption is correct, given $r_t$, $\mathcal{O}\_t$ is independent of $R$." This is not correct. Reward function $R$ is not a random variable. The reason why there exists a $R$ node in the graphical model is to show that every reward value comes from the single reward function $R$.

"how can we determine the reward function $R$ such that $q(\mathcal{O}\_t|r_t)$ approximates $p(\mathcal{O}\_t|s_t,a_t)$?" First of all, the reward function $R$ is no longer a random variable. Moreover, The reward value is sampled from the reward distribution $R(s_t, a_t)$ and the reward distribution is the output of the reward function $R$. Therefore, $q(\mathcal{O}\_t|r_t)$ is connected with the reward function $R$, and by backpropagation and reparameterization, we can update the parameter of the reward function.

$C_\theta$ serves as a binary classifier wherein it accepts $s_t$ and $a_t$ as input variables and produces the Bernoulli distribution as an output, which denotes the probability associated with optimality.

$R(s_t, a_t)$ depends on $s_t$ and $a_t$ since it is the output of the reward function given $s_t, a_t$. We are unable to locate the sentence in which we claim that $R(s_t, a_t)$ does not depend on $s_t$ and $a_t$.

$R(s_t, a_t)$ does not follow the same distribution for all $s_t$ and $a_t$. Every state-action pair $s_t,a_t$ generates a different reward value distribution $R(s_t,a_t)$. In detail, the reward value distribution $R(s_t, a_t)$ is $\mathcal{N}(\mu(s_t,a_t),\sigma^2(s_t,a_t))$.

"What does it mean when $r_t\sim R(s_t,a_t)$ in sec 2.3?" $r_t$ is the reward value sampled from the reward value distribution $R(s_t, a_t)$ using reparameterization. We cannot express it by $r_t=R(s_t,a_t)$.

Reviewer comment (Theory):
In the standard VI framework, it is often that we cannot directly minimize the KL[q(Z)||p(Z|X)] (supposed that we are interested in finding the posterior distribution of Z given the observation X). That is why we need to construct an ELBO that does not involve this KL term. Surprisingly, the ELBO formulation in equation (4) directly involves the term KL[q(Z)||p(Z|X)], i.e., the term KL[q(Ot|rt) || p(Ot|st,at)] in equation (4).

Author response:
2. VI framework.

Note that KL$(q(Z)||p(Z|X))$ is different from our KL divergence. Here is the derivation
$$
\begin{align*}
    \log p(\tau) &= \mathbb{E}_q[\log p(\tau)]\\\\
    &= \mathbb{E}_q\left[\log \frac{p(\tau,\mathcal{O})]}{p(\mathcal{O}|\tau)}\right] \\\\
    &= \mathbb{E}_q\left[\log\frac{p(\tau,\mathcal{O})}{p(\mathcal{O}|\tau)}\frac{q(\mathcal{O})}{q(\mathcal{O})}\right]\\\\
    &= \mathbb{E}_q\left[\log \frac{p(\tau,\mathcal{O})}{q(\mathcal{O})}\frac{q(\mathcal{O})}{p(\mathcal{O}|\tau)}\right]\\\\
    &= \mathbb{E}_q\left[\log\frac{p(\tau,\mathcal{O})}{q(\mathcal{O})}\right]+\mathbb{E}_q\left[\log\frac{q(\mathcal{O})}{p(\mathcal{O}|\tau)}  \right].
    \end{align*}
$$
The first term on the RHS is exactly our ELBO and the second term on the RHS is the gap between ELBO and the true log-likelihood of expert trajectories. The second term is intractable because it involves calculating the marginal distribution of evidence, $p(\tau)$.

Reviewer comment (Theory):
The derivations are based on several wrong assumptions, namely, that the expert actions do not depend on the state, and that the probability of an action being optimal is conditionally independent of state and action given its immediate. These assumptions are not clearly communicated, but instead hidden in a graphical model, that is presented as a matter of fact.

Author response:
1. The assumptions. 

We have updated our paper, in which there is an edge between state and action, and corresponding modification in the derivation has been made. However, adding such an edge does not have any impact on our ELBO. What's more, the two assumptions are discussed in the original Levine's paper (Reinforcement Learning and Control as Probabilistic Inference: Tutorial and Review) (listed as "control as inference paper" in the following). The first assumption is addressed by forward and backward message passing in section 2.3 in "control as inference paper" and expressed by the following

$$
    \begin{align*}
        \beta_t(s_t, a_t) &= p(\mathcal{O}\_{t:T}|s_t, a_t) \\\\
        \beta_t(s_t) &= p(\mathcal{O}\_{t:T}|s_t)\\\\
        p(a_t|s_t, \mathcal{O}\_{t:T}) &= \frac{\beta_t(s_t, a_t)}{\beta_t(s_t)}
    \end{align*}
$$

The second assumption follows the setting in the eqn.3 in "control as inference paper"
$$p(\mathcal{O}\_t=1|s_t, a_t) = \exp(r(s_t, a_t))$$ 

"This is clearly different from the probabilistic model as it is in the control as inference framework" 

You are correct. They should not depend on each other directly. We add the edge between $s_t$ and $a_t$ because we try to convince reviewer PBAj that the edge between $s_t$ and $a_t$ does not impact our theoretical result. The dependency between $s_t$ and $a_t$ is inferred by the forward and backward message.

We have updated our paper with the original derivation and a new graphical model, in which we separate out the reward function node and make it a deterministic node.

Reviewer comment (Reproducibility):
It is not clear whether the reward network outputs a scalar, or a distribution (e.g. mean & std of a Gaussian). The paper states "To illustrate this, consider the simplistic case where the reward value distribution is a univariate Gaussian: [...]", but it doesn't explicitly state that a Gaussian distribution is used in the experiments.

Author response:
4. Reward network. 

The reward network outputs a Gaussian distribution and such a Gaussian distribution is used in the experiments. Outputting a reward distribution improves the performance of our algorithms. The assumption on why it improves the performance is the randomness makes the learned reward function change smoothly since the model optimizes both the mean and variance. Updating the mean usually leads to a higher variance. Therefore, the new reward distribution may overlap with the old reward distribution. The oscillation of the learned reward function impedes the performance of the policy which is crucial in adversarial IRL setting.

Reviewer comment (Evaluation):
The paper claims that the results are significant according to a t-test, but it doesn't provide any further information on how the t-test was performed. It is not clear how the t-test was performed, because a t-test is usually used for comparing two groups. Were independent t-tests performed for each combination of two groups? In general, a t-test might not be appropriate in this case, not only because of non-Gaussian distributions, but in particular because the different populations can have significantly different variance.

Author response:
5. t-test. 

We compared the best and second-best algorithms using a t-test. We ran the experiments 50 times for each seed and we tried 5 random seeds. In total, there are 250 experiment data points, and it should be enough to be considered as Gaussian distributions. The problem of different variances is solved by running a t-test bidirectionally. First, We run a t-test on population A with the mean of population B. Then we run a t-test on population B with the mean of population A. If both t-tests claim that population A outperforms population B then we say population A is better than population B.

Reviewer comment (Evaluation):
The evaluations are performed by using the best policy among 5 different seeds. Such a procedure may favor unstable methods. It would be better to compare the average performance among different seeds.

Author response:
6. Different seeds. 

We are not using the best policy among ALL 5 different seeds. We ran experiments on the best policy for EACH seed. Therefore, there are 5 best policies in total. Although selecting the best policy may favor unstable methods, it is unfair to select a non-converged policy to test.

Reviewer comment (Evaluation):
The paper does not show learning curves. These need to be added, at least to the appendix.

Author response:
8. Learning curve and ablations. 

We will update our paper in the revised version.

Reviewer comment (Experiments):
Ablations are missing with respect to \( \lambda \) and also to evaluate the effect of unshaping the reward function with the value function.

Author response:
8. Learning curve and ablations. 

We will update our paper in the revised version.

Reviewer comment (Evaluation):
How is the t-test performed?

Author response:
5. t-test. 

We compared the best and second-best algorithms using a t-test. We ran the experiments 50 times for each seed and we tried 5 random seeds. In total, there are 250 experiment data points, and it should be enough to be considered as Gaussian distributions. The problem of different variances is solved by running a t-test bidirectionally. First, We run a t-test on population A with the mean of population B. Then we run a t-test on population B with the mean of population A. If both t-tests claim that population A outperforms population B then we say population A is better than population B.

Reviewer comment (Novelty):
The authors claim that its novelty lies in the "control as inference" graphical model. However, maximum-entropy-based IRL also hinges on inference within the "control as inference" graphical model. It can be derived as the maximum likelihood solution in this model, which is where the exponential distribution of the trajectory comes from (there are different derivations). Therefore, statements implying that the proposed method is distinct due to a lack of such modeling in IRL ("inspired by the graphical model for forward RL [...] and a general lack of such modeling in IRL") might be inaccurate. It would be beneficial for the paper to provide a more precise discussion of the differences between these probabilistic approaches.

Author response:
We thank you for the advice on the relationship between max-ent IRL and the "control as inference" graphical model. We will update such a relationship in the updated version paper.

Reviewer comment (Theory):
Equation 9 appears to have an error: it should read $exp(A(s, a))$ instead of $A(s, a)$. Additionally, the equal sign in this equation should either be replaced with a proportional sign, or it should link the first and third term, i.e., $q(\ldots)$ and $\sigma(\ldots)$.

Author response:
Thanks for pointing out the unclarity. We try to express that $q(...)$ is defined as $\sigma(...)$, which are the first and third terms.

Reviewer comment (Writing):
There are instances in the formulas where the bold font seems inconsistent. In Section 2.1, $s_t$ and $a_t$ are initially defined in bold, but this formatting is not maintained. If bold font indicates vectors, then the subscript $t$ should not be bold. This inconsistency also extends to the graphical model in Figure 1 (e.g., $O_t$ with bold $t$), and it impacted readability.

Author response:
We have updated the bold state and action notations.

Reviewer comment (Theory):
Why is it necessary to model the reward as random given the state and action? Would Equation 9 not remain unchanged if $q$ directly depended on the state and actions, thereby avoiding issues related to the gap between $p(O | s, a)$ and $p(O | r)$?

Author response:
The randomness of the reward function improves the performance of the algorithm in our experiments. We did not present this aspect in the paper. The assumption on why it improves the performance is the randomness makes the learned reward function change smoothly since the model optimizes both the mean and variance. Updating the mean usually leads to a higher variance. Therefore, the new reward distribution may overlap with the old reward distribution. The oscillation of the learned reward function impedes the performance of the policy which is crucial in adversarial IRL setting.

Reviewer comment (Experiments):
Is it necessary to retrain the classifier completely in each iteration? What is the computational cost associated with this process?

Author response:
We do not retrain the classifier completely. We use only one classifier and continue to train the classifier in each iteration, which aligns with GAIL and AIRL. In the experiments, the cost of training the classifier is super cheap compared to updating the policy and reward function.

Reviewer comment (Evaluation):
Do you have any insights into the potential sources of the increased performance compared to other IRL methods? It would be valuable to provide some explanation or analysis in this regard.

Author response:
The randomness of the reward function should be one of the reasons as I described in the second part of the response.

Reviewer comment (Presentation):
Graphical Model Clarification: The graphical model in the paper is quite unclear, particularly the representation of reward variables. Noting that the definition of V_t, representing the cumulative sum of rewards received after time t, is not adequately reflected in the graphical model. V_t captures the total reward received from time t onwards, and the reviewer suggests that this aspect should be more explicitly incorporated into the model for clarity.

Author response:
$V_t$ is the value function. While cumulative rewards represent the actual sum of rewards an agent receives during an episode, the value function provides a more structured and general representation of the expected future rewards an agent can achieve from different states under a specific policy. The definition of the input of the value function only depends on the current state, which is $s_t$.

Reviewer comment (Theory):
Dependence of Optimality on Reward: I am unsure about the rationale behind approximating optimality based on the reward. It is unclear whether having the reward information alone is sufficient to determine whether a state action pair is optimal. My question is whether the reward can serve as a reliable indicator of optimality and whether there is a clear justification for this approximation. This question is closely related to the following validity of the graphical model.

Author response:
We refer to Levine's paper (Reinforcement Learning and Control as Probabilistic Inference: Tutorial and Review). In the paper, the author has proved that with the optimality only depending on the reward, performing structured variational inference is equivalent to doing reinforcement learning in the max-ent setting. In our paper, we use exactly the same forward reinforcement learning, and the reward function is learned by our algorithm.

Reviewer comment (Theory):
Validity of Graphical Model: the entire theoretical development in the paper is based on the graphical model. They express concerns that if the graphical model's design is flawed or lacks clear justification, it could undermine the significance of the paper's contributions.

Author response:
We refer to Levine's paper (Reinforcement Learning and Control as Probabilistic Inference: Tutorial and Review). In the paper, the author has proved that with the optimality only depending on the reward, performing structured variational inference is equivalent to doing reinforcement learning in the max-ent setting. In our paper, we use exactly the same forward reinforcement learning, and the reward function is learned by our algorithm.