Implicit Q-Learning and In-Sample Learning

Every Q-learning update contains a question the offline dataset cannot answer: \(\max_{a'} Q(s', a')\) asks for the value of whichever action looks best, and the action that looks best to a function approximator is very often one the behaviour policy never took. The estimate there is extrapolation, errors are selected for being optimistic, and bootstrapping compounds them. Distribution shift in offline RL describes the failure and conservative value estimation fights it by pushing down values of unseen actions. Implicit Q-learning takes the other road: never query an unseen action at all.

Expectiles as a soft in-sample max

For a random variable \(X\), the \(\tau\)-expectile \(m_\tau\) minimises an asymmetric squared loss:

At \(\tau = 0.5\) this is ordinary least squares and \(m_\tau\) is the mean. As \(\tau \to 1\), residuals above \(m\) are weighted far more heavily than those below, and \(m_\tau\) climbs toward the maximum of the support of \(X\).

IQL applies this across actions (Kostrikov, Nair & Levine, 2021, Offline Reinforcement Learning with Implicit Q-Learning, ICLR 2022, arXiv:2110.06169). Treat \(Q(s, a)\) with \(a\) drawn from the dataset at state \(s\) as a random variable, and fit a separate state-value network \(V_\psi\) to its upper expectile, then back that value up into \(Q_\theta\):

Here \(\mathcal{D}\) is the dataset, \(Q_{\hat\theta}\) a slowly updated target network, and \(\gamma\) the discount. Every term uses only \((s, a, s')\) tuples that exist in the data. With \(\tau = 0.5\) the procedure is SARSA and estimates the behaviour policy's value; with \(\tau\) near 1 it approximates Q-learning restricted to actions the data supports. The separate \(V\) matters: taking an expectile of \(r + \gamma V(s')\) directly would also reward lucky stochastic transitions, whereas \(Q\) already averages over next states before the expectile is taken over actions.

A small calculation shows the knob. Suppose at some state the dataset contains three actions, equally often, with \(Q\) values 1, 2 and 5. The mean is 2.67. For \(\tau = 0.9\) with \(m\) between 2 and 5, the first-order condition \(0.9(5 - m) = 0.1\big((m-1) + (m-2)\big)\) gives \(m \approx 4.36\). At \(\tau = 0.99\) it gives \(m \approx 4.93\). The best in-data action's value is approached without ever evaluating an action outside the data.

Extracting a policy without leaving the data

Learning \(Q\) and \(V\) produces no policy. IQL extracts one with advantage-weighted regression, a behaviour-cloning loss reweighted by exponentiated advantage (Peng et al., 2019, Advantage-Weighted Regression, arXiv:1910.00177):

with inverse temperature \(\beta\). Continuing the example with \(V = 4.36\) and \(\beta = 3\), the weights are \(e^{3(0.64)} \approx 6.8\) for the best action, \(e^{3(-2.36)} \approx 0.0008\) and \(e^{3(-3.36)} \approx 0.00004\) for the others. The policy clones the good action and effectively ignores the rest. It never assigns probability mass by maximising \(Q\) over free actions, which is where offline methods usually get hurt.

What it achieves, and a disagreement it settled only partly

On D4RL's AntMaze tasks, which require stitching segments of suboptimal trajectories into a route to a goal, the IQL paper reports 47.5 on antmaze-large-diverse against 14.9 for CQL and 0.0 for behaviour cloning and Decision Transformer. On the locomotion tasks it is merely competitive. Training in the authors' JAX reimplementations took about 20 minutes against 80 for CQL. The in-sample structure also fine-tunes online well: on antmaze-medium-play IQL rose from 72.0 to 95.0 with online data, while CQL fell from 23.0 to 0.0.

The contrast with one-step methods is the substantive debate. Brandfonbrener and colleagues had argued that a single step of policy improvement on the behaviour value function performs about as well as iterative offline RL. The IQL authors point out that the Gym locomotion datasets contain many near-optimal trajectories, where one step suffices, whereas AntMaze needs multi-step dynamic programming, and there one-step methods fall far behind. Whether multi-step DP is necessary depends on whether good behaviour is already in the data as whole trajectories or only as pieces.

When it breaks

It cannot exceed the data's actions. IQL approximates the best action in the support of the dataset at each state. If no logged action is good at a state, no expectile finds one. That is the price of never extrapolating.

Two sensitive hyperparameters, set per domain. The paper uses \(\tau = 0.9\), \(\beta = 10\) for AntMaze and \(\tau = 0.7\), \(\beta = 3\) for locomotion. Tuning them requires an offline estimate of policy quality, which is its own unsolved problem; see evaluating an offline RL policy.

High \(\tau\) trades bias for variance. As \(\tau \to 1\) the expectile is driven by a few large residuals, including noise in \(Q_{\hat\theta}\), so values become optimistic about estimation errors on in-data actions rather than unseen ones.

Exponentiated weights are heavy-tailed. A few large advantages dominate each batch, so implementations clip the weights, which quietly changes the objective. A unimodal Gaussian policy fitted to weighted multimodal actions can also average between two good actions and select neither.