Offline RL II

Where Monday left us: the distributional shift problem

The setting is the same as the previous lecture: a dataset $\mathcal{D}$ of $(s, a, s', r)$ tuples collected by some unknown behavior policy $\pi_\beta$ — humans, a hand-designed controller, a mixture of old RL runs, “it could really be anything” — and no further interaction. Value-based methods fit the Q-function by minimizing Bellman error in expectation under $\pi_\beta$, because that’s where the samples come from. But the backup uses $\mathbb{E}_{a' \sim \pi_{\text{new}}}[Q(s', a')]$, and $\pi_{\text{new}}$ is precisely the thing we’re trying to make different from $\pi_\beta$ — worse, it is “strongly incentivized to find kind of adversarial actions that will fool the q function into outputting erroneously large q values.” In practice, standard actor-critic run offline gets a true return around $-250$ while its Q-function predicts returns near $10^7$.

Everything in this lecture responds to that one failure, in three families: keep the policy close to $\pi_\beta$, never evaluate an action you haven’t seen, or train the Q-function so it can’t overestimate — plus model-based versions of the same ideas.

Policy constraints, and why the obvious ones disappoint

The natural fix is a constrained actor update:

$\theta \leftarrow \arg\max_\theta \; \mathbb{E}_{a \sim \pi_\theta(a \mid s)}[Q(s, a)] \quad \text{s.t.} \quad D_{KL}(\pi_\theta \| \pi_\beta) \le \epsilon$

Small enough $\epsilon$, no more erroneous values — in principle. Levine gives away the punch line immediately: naive policy constraints “don’t actually work all that well,” though the idea is decades old (trust regions, natural policy gradients) and useful for analysis. Two problems. First, we don’t know $\pi_\beta$, so we either fit it with behavioral cloning or implement the constraint cleverly enough to need only samples (the latter works much better). Second, the constraint is simultaneously too pessimistic and not pessimistic enough. Not pessimistic enough: low error in expectation under a close distribution doesn’t mean low error at your policy’s chosen points. Too pessimistic: if $\pi_\beta$ is uniformly random, no action is out of distribution, yet a KL constraint insists $\pi$ stay near-uniform — “why should we make $\pi$ more random just because $\pi_\beta$ was more random?”

In pictures: a Q-curve fit to dataset actions extrapolates unreliably where $\pi_\beta$ has low density. The KL-constrained policy peaks at the best observed action but keeps fat tails over bad actions, because tightening further would violate the constraint. What we actually want is the best in-support policy, which concentrates hard on good actions but only within the data’s support.

Explicit implementations of the KL version are worth knowing as context. Modify the actor objective by writing out $D_{KL}(\pi_\theta \| \pi_\beta) = -\mathbb{E}_{\pi_\theta}[\log \pi_\beta(a \mid s)] - \mathcal{H}(\pi_\theta)$ and folding it in with a Lagrange multiplier $\lambda$ (found by dual gradient descent, or just tuned): the $\log \pi_\beta$ term adds to the Q-values and the entropy term is closed-form for Gaussian or categorical policies. Or modify the reward, $\bar{r}(s, a) = r(s, a) - \lambda D(\pi(\cdot \mid s), \pi_\beta(\cdot \mid s))$, which also penalizes actions that lead to high future divergence — slightly better theory. But in Levine’s experience “both approaches don’t really work all that well.”

Implicit constraints: RL as weighted behavioral cloning

Here is the trick that removes $\pi_\beta$ from the equation. Write the Lagrangian of the KL-constrained problem and solve for the optimal policy in closed form (Peters et al.’s REPS and Rawlik et al.’s psi-learning derive versions of this):

$\pi^\star(a \mid s) = \frac{1}{Z(s)}\, \pi_\beta(a \mid s) \exp\!\left(\frac{1}{\lambda} A^\pi(s, a)\right)$

As $\lambda \to 0$ this becomes the greedy argmax policy; for finite $\lambda$ it makes suboptimal actions exponentially less likely — but multiplied by $\pi_\beta$, so an action with tiny behavior probability ends up with tiny probability no matter how large its advantage. We can’t evaluate the right-hand side without $\pi_\beta$, but we can sample from it: expectations under $\pi_\beta$ are just dataset averages. So approximate $\pi^\star$ by weighted maximum likelihood:

$\pi_{\text{new}} = \arg\max_\pi \; \mathbb{E}_{(s,a) \sim \pi_\beta}\!\left[\log \pi(a \mid s)\, \frac{1}{Z(s)} \exp\!\left(\frac{1}{\lambda} A^{\pi_{\text{old}}}(s, a)\right)\right]$

The loss is a behavioral cloning loss where each dataset action carries a weight from the critic’s advantage. With a Q-function critic this is advantage-weighted actor critic (AWAC); with Monte Carlo returns, advantage-weighted regression (AWR).

The residual trouble: the constraint holds at convergence, not along the way. Both the target value $\mathbb{E}_{a' \sim \pi_\theta}[Q(s', a')]$ and the advantage $A = Q - \mathbb{E}_{\pi_\theta}[Q]$ query the Q-function at actions sampled from the current $\pi_\theta$, which mid-training can still be out of distribution.

Implicit Q-learning: never query an action you haven’t seen

Can we avoid all out-of-distribution action queries? Replace the expectation in the target with a separate value network: $Q(s, a) \leftarrow r(s, a) + \gamma V(s')$. Train $V$ by regressing $V(s_i)$ onto $Q(s_i, a_i)$ over dataset tuples — but with mean squared error this estimates the value of $\pi_\beta$, since the actions come from $\pi_\beta$.

Now the funny thought. At any given state you probably saw only one action, but nearby similar states saw other actions, and insofar as the network generalizes, the regression target at a state is really a distribution of values induced by the actions. MSE recovers its mean — the behavior policy’s value. What if we instead recovered its upper end, “the value of the best policy supported by the data”? That’s what the expectile loss does:

\qquad V \leftarrow \arg\min_V \frac{1}{N} \sum_{i=1}^N \ell_2^\tau\big(V(s_i) - Q(s_i, a_i)\big)$$ A tilted parabola: with $\tau = 0.9$, $V$ falling below a sampled $Q(s_i, a_i)$ costs nine times more than exceeding it, pushing $V$ toward the top of the distribution. Won't favoring large values re-create overestimation? No — we're never going to query $Q$ with any action that was not in the dataset, so extrapolation errors never enter the loss at all. Formally, for large enough $\tau$ this trains $$V(s) \leftarrow \max_{a \in \Omega(s)} Q(s, a), \qquad \Omega(s) = \{a : \pi_\beta(a \mid s) \ge \epsilon\}$$ a Bellman backup under an *implicit* policy — the argmax restricted to in-support actions. And because the max is taken per-state, the result stitches together the best behavior across all states: potentially much better than the single best trajectory in the data. This is **implicit Q-learning (IQL)** (Kostrikov, Nair, Levine). <CodeWalkthrough title="IQL: the full loop"> Alternate two regression losses that only ever touch dataset states and actions, then extract the policy at the end: ```python # critic updates (repeat): y_i = r(s_i, a_i) + gamma * V_psi(s_next_i) # no policy anywhere L_Q = mean((Q_phi(s_i, a_i) - y_i)**2) # MSE, V as target L_V = mean(expectile_loss(V_psi(s_i) - Q_phi(s_i, a_i), tau)) # policy extraction (once, after training): L_pi = -mean(log_pi_theta(a_i | s_i) * exp(Q_phi(s_i, a_i) - V_psi(s_i))) ``` There are no expectations under $\pi_\theta$ anywhere, and the critic is completely independent of the actor — hence implicit "Q-learning," not really actor-critic. The policy never participates in the value updates, so a separate advantage-weighted extraction step recovers it from the learned $Q$ and $V$. </CodeWalkthrough> ## Conservative Q-learning: push the peaks back down The alternative to controlling the actor is to repair the Q-function directly. The overestimated actions are, by construction, the ones with (erroneously) maximal value — so add a term that finds high-Q actions and pushes them down: $$\hat{Q}^\pi = \arg\min_Q \max_\mu \; \alpha\, \mathbb{E}_{s \sim \mathcal{D},\, a \sim \mu(a \mid s)}[Q(s, a)] \;-\; \alpha\, \mathbb{E}_{(s,a) \sim \mathcal{D}}[Q(s, a)] \;+\; \mathbb{E}_{(s,a,s') \sim \mathcal{D}}\!\left[\big(Q(s,a) - (r(s,a) + \mathbb{E}_\pi[Q(s',a')])\big)^2\right]$$ The first term alone (with the Bellman error) already yields a pointwise lower bound $\hat{Q}^\pi \le Q^\pi$ for large enough $\alpha$ — but it's *too* pessimistic, pushing down everything. The second term pushes *up* on dataset actions. Maximizing Q-values to combat overestimation sounds strange until you see the feedback loop: if the large Q-values sit on dataset actions, $\mu$ selects those very actions and the two terms cancel; if they sit out of distribution, the first term pushes them down while the second props up the data — steering $\mu$ back toward the distribution. <Theorem name="The CQL guarantee"> With both terms, $\hat{Q}^\pi(s,a) \le Q^\pi(s,a)$ is no longer guaranteed for all $(s, a)$ — but $\mathbb{E}_{\pi(a \mid s)}[\hat{Q}^\pi(s, a)] \le \mathbb{E}_{\pi(a \mid s)}[Q^\pi(s, a)]$ still holds for all $s \in \mathcal{D}$. That's all we really want: what matters is not overestimating the value of the current policy. </Theorem> Implementation is where CQL gets convenient. Add a regularizer $\mathcal{R}(\mu)$; with the maximum-entropy choice $\mathcal{R} = \mathbb{E}_{s \sim \mathcal{D}}[\mathcal{H}(\mu(\cdot \mid s))]$, the optimal $\mu(a \mid s) \propto \exp(Q(s, a))$ and the pushdown term collapses to $\mathbb{E}_{a \sim \mu}[Q(s,a)] = \log \sum_a \exp(Q(s, a))$. For discrete actions, literally compute the log-sum-exp — no explicit $\mu$ at all; for continuous actions, estimate the expectation by importance sampling (single-step weights, so the usual variance blowup doesn't apply). The full algorithm is ordinary Q-learning or actor-critic with two extra terms in the critic loss: update $\hat{Q}^\pi$ with $\mathcal{L}_{CQL}$, then update the policy as usual. ## Model-based offline RL: MOPO, COMBO, and the trajectory transformer Models are a natural fit for offline RL — train on all your data as long as you like, then plan or learn a policy. But a model answers "what if" questions, and without fresh data the policy learns to ask malicious ones: take an action that fools the model into an out-of-distribution state that looks high-reward, from which the model predicts even crazier states. Out-of-distribution *actions* now come with out-of-distribution *states*. **MOPO** (and the conceptually similar MOReL) punishes the policy for exploiting: $\tilde{r}(s, a) = r(s, a) - \lambda\, u(s, a)$, where $u$ is an uncertainty penalty, then runs any existing Dyna-style model-based algorithm. For the theory to work, $u$ must be at least as large as the model's true error — in practice estimated by ensemble disagreement, though getting guaranteed error estimates is an open problem. <Theorem name="What the MOPO bound buys you"> Assume the value function class is expressive and $u$ upper-bounds the true model error under a divergence like total variation. Let $\pi_\delta$ be the best policy (by *true* return) whose expected uncertainty is at most $\delta$. Then the policy $\hat{\pi}$ optimized under the penalized model satisfies, in true return, $\eta(\hat{\pi}) \ge \eta(\pi_\delta) - 2\lambda\delta$. Substituting $\pi_\beta$ (whose $\delta \approx 0$, since the model was trained on its states) guarantees we don't do worse than the behavior policy; substituting $\pi^\star$ bounds the optimality gap by how wrong the model is on the states the optimal policy visits. </Theorem> **COMBO** ports CQL into this setting: just as CQL minimizes Q-values of policy actions, COMBO minimizes Q-values of *model-generated* state-action tuples while maximizing them on the dataset — "a somewhat GAN-like idea." If the model produces unrealistic states the Q-function makes them look bad; if its samples are indistinguishable from data, the terms cancel and the model is being correct. No reward modification needed, and it tends to work a little better than MOPO/MOReL. The **trajectory transformer** drops the policy entirely and plans. Train a joint density model $p_\beta(\tau)$ over whole trajectories — discretized *per dimension*, so the sequence model's tokens are individual dimensions of states and actions, not time steps — using the most powerful density estimator you can get, which offline you can afford to make big. Then plan with beam search, ranking candidates by cumulative *reward* rather than likelihood: sampling continuations from $p_\beta$ keeps every candidate in-distribution, and selecting among them maximizes return. Because it never strays from data-supported actions it predicts accurately over very long horizons — hundreds of steps of humanoid motion that are pure model prediction, not simulation. ## Which algorithm, and why offline RL changes the workflow Levine's rule of thumb (as of the recording, in a fast-moving field): training purely offline, use **CQL** (one hyperparameter, widely validated) or **IQL** (more flexible, more hyperparameters). Offline training followed by online fine-tuning: **AWAC** is the well-tested choice, **IQL** empirically better still — CQL fine-tunes poorly because it stays too conservative. If you're confident you can train a good model in your domain, **COMBO** is "basically CQL but with models"; the trajectory transformer is powerful but extremely expensive, and without dynamic programming it's limited on very long horizons. The bigger argument is about workflow. Real-world online RL means instrumenting the task with safety mechanisms, engineering autonomous resets, building an automated reward function, waiting out a long run — then rerunning the whole pipeline after every algorithm tweak, and throwing everything away for the next task. Offline: collect a dataset once from any mixture of sources (you can even have humans label rewards, since you only need rewards on training data), train, tweak, retrain *without recollecting*, aggregate new data into a growing set, and reuse it for the next project. His group's examples: a multi-task robotic grasping dataset (12 tasks, months of collection) reused to test a goal-conditioned-RL hypothesis with zero new data collection, and a 40-hour off-road navigation dataset reused across three different projects over two years. The gap between that dream and current algorithms defines the open problems: **workflows** — supervised learning's train/validation/test protocol predicts real-world performance, while offline RL still evaluates by deploying, and off-policy evaluation methods themselves need online tuning; **statistical guarantees** that aren't loose; and **scaling** to the wide range of applications where offline RL should, in principle, apply. <KeyTakeaways> - Distributional shift is the enemy: offline value-based methods train Q under $\pi_\beta$ but query it under a policy adversarially chosen to exploit its errors. - KL policy constraints are simultaneously too pessimistic (uniform $\pi_\beta$ forces a random policy) and not pessimistic enough (low expected error is not low error at chosen points); support constraints are what you actually want. - The KL-constrained update has closed form $\pi^\star \propto \pi_\beta \exp(A/\lambda)$, implementable as advantage-weighted behavioral cloning (AWAC/AWR) using only samples — no $\pi_\beta$ model needed. - IQL avoids every out-of-distribution query: an expectile loss on $V(s) - Q(s, a)$ over dataset tuples recovers the max of $Q$ over in-support actions, stitching the best behavior across states. - CQL adds two terms to the critic loss — push down high-Q actions (a log-sum-exp for discrete actions), push up dataset actions — guaranteeing no overestimation of the policy's value in expectation. - Model-based variants penalize model exploitation (MOPO), apply CQL to model samples (COMBO), or plan under a trajectory density model (trajectory transformer). Rule of thumb: CQL or IQL for pure offline; AWAC or IQL when fine-tuning online; COMBO if you can trust your model. </KeyTakeaways> ## Check yourself Exercises in the lecture's spirit — work them on paper before checking the solutions. <Problem number="16.1" source="KL vs support constraints"> The behavior policy at some state is uniform over 10 discrete actions. The true Q-values are $Q(s, a_1) = 10$ and $Q(s, a_j) = 0$ for the rest, and the Q-estimates are accurate for all 10 actions (all are in-distribution). (a) What policy would we *like* to learn at this state? (b) Show that a KL constraint $D_{KL}(\pi \| \pi_\beta) \le \epsilon$ with small $\epsilon$ forbids it, computing the KL divergence of the greedy policy from $\pi_\beta$. (c) Why does a support constraint permit it? <div slot="solution"> (a) The greedy policy $\pi(a_1 \mid s) = 1$: every action is in-distribution, the estimates are reliable, so nothing stops us from taking the best one. (b) For the greedy policy, $D_{KL}(\pi \| \pi_\beta) = \sum_a \pi(a) \log \frac{\pi(a)}{\pi_\beta(a)} = 1 \cdot \log \frac{1}{0.1} = \log 10 \approx 2.3$. Any $\epsilon$ less than $\log 10$ rules the greedy policy out, forcing $\pi$ to spread probability over provably worthless actions — purely because $\pi_\beta$ was random. This is the "too pessimistic" failure: the constraint measures closeness in distribution, not membership in the support. (c) A support constraint only requires $\pi(a \mid s) > 0 \Rightarrow \pi_\beta(a \mid s) \ge \epsilon$. Here every action has $\pi_\beta = 0.1$, so *any* redistribution of probability over these 10 actions — including all mass on $a_1$ — satisfies it. The support constraint permits exactly the best in-support policy that the KL constraint forbids. </div> </Problem> <Problem number="16.2" source="implicit constraint derivation"> Derive the closed-form solution $\pi^\star(a \mid s) = \frac{1}{Z(s)} \pi_\beta(a \mid s) \exp(A^\pi(s,a)/\lambda)$ for the problem $\max_\pi \mathbb{E}_{a \sim \pi}[A^\pi(s, a)]$ subject to $D_{KL}(\pi \| \pi_\beta) \le \epsilon$ (and $\pi$ a valid distribution), and explain why replacing $Q$ with the advantage $A$ changes nothing. <div slot="solution"> Write the Lagrangian with multiplier $\lambda$ for the KL constraint and $\eta$ for normalization: $\mathcal{L} = \sum_a \pi(a) A(s,a) - \lambda \sum_a \pi(a) \log \frac{\pi(a)}{\pi_\beta(a)} + \eta \left(1 - \sum_a \pi(a)\right)$. Setting $\partial \mathcal{L} / \partial \pi(a) = 0$: $A(s,a) - \lambda \log \frac{\pi(a)}{\pi_\beta(a)} - \lambda - \eta = 0$, so $\log \frac{\pi(a)}{\pi_\beta(a)} = \frac{1}{\lambda} A(s,a) - 1 - \frac{\eta}{\lambda}$, giving $\pi(a) = \pi_\beta(a) \exp(A(s,a)/\lambda) \cdot e^{-1 - \eta/\lambda}$. The constant factor is fixed by normalization, i.e. it is $1/Z(s)$ with $Z(s) = \sum_a \pi_\beta(a) \exp(A(s,a)/\lambda)$. Since $A(s,a) = Q(s,a) - V(s)$ differs from $Q$ by an action-independent constant, using $Q$ instead multiplies numerator and denominator by $\exp(V(s)/\lambda)$, which cancels — the same $\pi^\star$. As $\lambda \to 0$ the exponent's scale diverges and $\pi^\star$ concentrates on the argmax; as $\lambda \to \infty$, $\pi^\star \to \pi_\beta$. The multiplier $\lambda$ is set to satisfy the KL constraint (dual descent) or treated as a hyperparameter. </div> </Problem> <Problem number="16.3" source="expectile loss computation"> At a state, generalization from similar states gives the value target two possible values: $Q = -10$ and $Q = +10$, equally likely (the histogram from the lecture's two-trajectory picture). (a) What $V$ minimizes the MSE loss $\mathbb{E}[(V - Q)^2]$? (b) For the expectile loss with parameter $\tau$, write the expected loss and find the minimizing $V$ for $\tau = 0.9$. (c) What does this $V$ mean, and why doesn't it cause overestimation? <div slot="solution"> (a) MSE is minimized by the mean: $V = 0$ — the value of the behavior policy, which takes each action half the time. (b) Assume $-10 \le V \le 10$, so $V - (-10) > 0$ (penalty weight $1 - \tau$) and $V - 10$ is negative (weight $\tau$). Expected loss: $L(V) = \tfrac{1}{2}(1 - \tau)(V + 10)^2 + \tfrac{1}{2}\tau (V - 10)^2$. Setting $L'(V) = (1-\tau)(V + 10) + \tau(V - 10) = 0$ gives $V = 10(2\tau - 1) = 10 \cdot 0.8 = 8$. (This matches the lecture's picture, where the upper expectile lands near $+10$; the slide's example value is $9.7$ for a larger $\tau$.) As $\tau \to 1$, $V \to 10 = \max Q$. (c) $V \approx 8$ approximates the value of the *best in-support action* — commit to the branch worth $+10$ rather than average over both. It cannot overestimate the way naive Q-learning does because both $Q$-values entering the loss were trained on dataset actions; no extrapolated, never-trained $Q(s, a)$ is ever queried. The tilt selects among trusted values instead of maximizing over untrusted ones. </div> </Problem> <Problem number="16.4" source="CQL feedback mechanism"> Consider CQL's two conservative terms $\alpha\, \mathbb{E}_{s \sim \mathcal{D},\, a \sim \mu}[Q(s,a)] - \alpha\, \mathbb{E}_{(s,a) \sim \mathcal{D}}[Q(s,a)]$ with $\mu$ chosen adversarially to maximize the first term. At a state with dataset action $a_{\text{data}}$ and one out-of-distribution action $a_{\text{ood}}$: (a) what is the net effect of the two terms when $Q(s, a_{\text{ood}}) \gg Q(s, a_{\text{data}})$, and when $Q(s, a_{\text{data}})$ is the largest? (b) Explain why, with the maximum-entropy regularizer, the pushdown term becomes $\log \sum_a \exp(Q(s,a))$, and verify the claimed behavior in both regimes using properties of log-sum-exp. <div slot="solution"> (a) If $Q(s, a_{\text{ood}})$ dominates, $\mu$ concentrates on $a_{\text{ood}}$: the first term pushes $Q(s, a_{\text{ood}})$ down while the second pushes $Q(s, a_{\text{data}})$ up — the erroneous peak shrinks and in-distribution values rise, so next iteration $\mu$'s maximizing actions move back toward the data. If the largest Q-value is already on $a_{\text{data}}$, $\mu$ picks (mostly) $a_{\text{data}}$, and pushing down and pushing up the same value roughly cancel: near zero net effect. That's the feedback: pressure exists exactly when, and only where, out-of-distribution actions look too good. (b) With $\mathcal{R}(\mu) = \mathcal{H}(\mu)$, maximizing $\mathbb{E}_\mu[Q] + \mathcal{H}(\mu)$ over distributions gives the Boltzmann form $\mu(a \mid s) \propto \exp(Q(s,a))$, and plugging it back in yields the soft maximum $\mathbb{E}_\mu[Q] = \log \sum_a \exp Q(s,a)$ (standard variational identity). Log-sum-exp is dominated by the largest Q-value: when $a_{\text{ood}}$ has the largest value the gradient of the pushdown term concentrates on it (pushes down the peak), and when dataset actions dominate, its gradient concentrates on them and is offset by the $-\alpha\, \mathbb{E}_\mathcal{D}[Q]$ push-up term — exactly the cancellation in (a), now computable in closed form for discrete actions. </div> </Problem> <Problem number="16.5" source="reading the MOPO bound"> The MOPO guarantee states $\eta(\hat{\pi}) \ge \eta(\pi_\delta) - 2\lambda\delta$, where $\pi_\delta$ is the best policy (by true return) whose expected model-error penalty is at most $\delta$, provided $u(s,a)$ upper-bounds the true model error. (a) Use it to argue $\hat{\pi}$ approximately improves over the behavior policy. (b) State the condition under which it guarantees near-optimality. (c) Why does the whole guarantee collapse if $u$ merely *estimates* the model error rather than bounding it, and what is the practical proxy for $u$? <div slot="solution"> (a) Take $\pi = \pi_\beta$: the model was trained precisely on states and actions from $\pi_\beta$, so its error there — hence the expected penalty $\delta$ — is close to zero. The bound gives $\eta(\hat{\pi}) \ge \eta(\pi_\beta) - 2\lambda\delta \approx \eta(\pi_\beta)$: at worst we roughly match the behavior policy, and since $\hat{\pi}$ optimizes the penalized objective it typically improves. (b) Take $\pi = \pi^\star$: $\eta(\hat{\pi}) \ge \eta(\pi^\star) - 2\lambda\delta^\star$, where $\delta^\star$ is the model's expected error on the states and actions *the optimal policy* visits. If the dataset happens to cover those regions well, $\delta^\star$ is small and we recover near-optimality — coverage of the optimal policy's distribution, not of everything, is what matters. (c) The theorem's assumption is $u \ge$ true error; if $u$ underestimates error somewhere, the policy can exploit exactly those states — enjoying erroneously high model rewards at a penalty too small to cancel the gain — and the inequality no longer holds. Guaranteed upper bounds on neural model error are an open problem; the common practical proxy is disagreement across an ensemble of learned models. </div> </Problem>