Expectile regression in Implicit Q-Learning

Question

I am reading Kostrikov et al.'s "Offline Reinforcement Learning with Implicit Q-Learning" but got stuck understanding one particular transformation they use.

They describe the loss function they would like to minimize:

which is essentially the loss function for Q-learning, with a small twist: because they are interested in offline learning, the actions $a'$ considered for the maximization are only those produced by the behavior policy $\pi_\beta$ (essentially, the policy represented by the the offline dataset $\mathcal{D}$ we are learning from) with probability greater than 0.

At this point the authors propose to approximate the maximization with expectile regression. They explain expectile regression to be realized by the following minimization of a loss function $L_2^\tau$ (not replicated here for brevity), where $\tau \in [0,1]$ is a target expectile:

The intuition is that if $\tau$ is, say, 0.75, then the minimizing $m_\tau$ is a point such that $P(X < m_\tau) \approx 0.75$. So if you choose a high value of $\tau$ close to $1$, you are in effect computing a maximum value of $\mathcal{X}$ by stochastic gradient descent, using samples $x \sim X$, without necessarily having to iterate over all values of $X$.

Now comes the part that I don't understand. They apply the idea of expectile maximization to Eq. (4) and obtain a new loss function:

the argument being that minimizing this loss function with respect to $\theta$ is approximately equivalent to minimizing the loss function in Eq. (4) with respect to $\theta$. While I can see the intuition, I am not seeing exactly how this is mathematically justified.

I think they are using this transformation:

$E_{(s,a,s')\sim \mathcal{D}}\left[\left(r(s,a) + \gamma max_{a'\in \mathcal{A}\;s.t. \pi_\beta(a'|s')>0} Q_{\hat{\theta}}(s',a') - Q_\theta(s, a)\right)^2\right]=$ $E_{(s,a,s')\sim \mathcal{D}}\left[\left(max_{a'\in \mathcal{A}\;s.t. \pi_\beta(a'|s')>0} \left(r(s,a) + \gamma Q_{\hat{\theta}}(s',a') - Q_\theta(s, a)\right)\right)^2\right]$

and then, by some step I don't understand, they approximate that by

$E_{(s,a,s',a')\sim \mathcal{D}}\left[L_2^\tau\left(r(s,a) + \gamma Q_{\hat{\theta}}(s',a') - Q_\theta(s, a)\right)\right]$

How can that be shown?

Answer 1

Because $argmin_{m_\tau} E_{x\sim X} \left[L_2^\tau(x - m_\tau)\right]$ approximates $max_{x \in X} \hspace{0.1cm} x$, then

$E_{(s,a,s')\sim \mathcal{D}}\left[\left(max_{a'\in \mathcal{A}\;s.t. \pi_\beta(a'|s')>0} \left(r(s,a) + \gamma Q_{\hat{\theta}}(s',a') - Q_\theta(s, a)\right)\right)^2\right]$

is approximated (with $\tau$ sufficiently high) by

$E_{(s,a,s')\sim \mathcal{D}}\left[argmin_\theta E_{a'\in \mathcal{D}}\left[L_2^\tau\left(r(s,a) + \gamma Q_{\hat{\theta}}(s',a') - Q_\theta(s, a)\right)\right]\right]$

which is then

$argmin_\theta E_{(s,a,s',a')\sim \mathcal{D}}\left[L_2^\tau\left(r(s,a) + \gamma Q_{\hat{\theta}}(s',a') - Q_\theta(s, a)\right)\right]$