Offline RL: KL, Advantage Weighting, and Guidance

Have an offline dataset $𝒟=\{(s,a,r,s^{\prime })\}$ collected by some behaviour policy $\beta (a\mid s)$. We want to learn a new policy $\pi$ with greater expected return.

Obvious idea: choose actions that maximise a learned critic:

However, $Q$ has only been trained on / near actions in $𝒟$. Far from that, $Q$ can be very wrong. Thus, the learned policy can snowball errors in $Q$ and just be very bad.

Therefore we want: prefer actions with high value, but remain close to the data policy.

This is a good start, but it ignores the behaviour policy. Consider one fixed state $s$. For convenience, write $Q(a):=Q(s,a)$.

We optimise over all possible action distributions:

$\tau$ is the entropy temperature. Maximise $Q$, but don’t collapse the distribution (like SAC).

Thus:

Using the normalisation constraint gives:

Closed-form solution: a Boltzmann distribution. As $\tau \to 0$, more like argmax. For continuous actions, the sum becomes an integral.

Cool link: entropy is KL regularisation to a uniform distribution

Suppose $u(a)$ is a uniform reference distribution. Then:

Since $u(a)$ is constant:

Thus, entropy regularisation is a special case of “stay close to the reference distribution,” where the reference is uniform.

Penalise the policy for leaving the behaviour policy via KL:

At the fixed state, this is:

Using Lagrange again:

Solve for $\pi$ and use the constraint again:

Or equivalently:

Basically, weight the prior with the $Q$ (softmax).

Side thought: reverse KL tends to fit modes better and forward KL tends to fit means better. We kind of prefer reverse KL here.

Move to advantage

For fixed $s$, $V(s)$ doesn’t depend on $a$:

The second factor is constant over all actions, so it disappears when normalised. Thus:

That equation describes the improved policy distribution, but we still need to fit a parameterised model ${\pi }_{\theta }$. One way is weighted maximum likelihood:

Basically gradient ascent to make the data more likely (supervised learning), but now with an extra loss-weighting term. This is the idea behind advantage-weighted regression and related methods (clone every action, but clone good actions more).

Is what Flow Q-learning describes as weighted behavioural cloning. Remember that increasing log probability is what they all sort of do.

Issue with weighted MLE: very sensitive to errors in value estimates. The exponential can make the weights highly concentrated.

Let’s try to rewrite the way of representing it.

More generally, suppose:

where $f$ is a non-negative, non-decreasing function. It could be $\exp (A/\tau )$ or $\mathbf{1}[A\ge 0]$. Just another way to weight the prior favourably.

CFGRL paper idea

Suppose a target density is a product:

For smooth, positive factors, taking logs and gradients:

Looks like the scores that diffusion models learn (kind of). For policy improvement, formally:

Thus, a link: regularised policy improvement ↔ product of behaviour and optimality ↔ generative-model guidance.

Condition on optimality

Introduce a binary variable $o$, where $o=1$ means the action is considered good. We can train two distributions: $p(a\mid s)$ and $p(a\mid s,o=1)$. The first is behaviour and the second is behaviour conditioned on optimality.

By Bayes:

Since the denominator doesn’t depend on $a$:

Thus, we can represent a policy as a product of two terms by conditioning on optimality. If:

we get back the KL-regularised improved policy:

But we can also pick $o=\mathbf{1}[A(s,a)>0]$ and still learn. We basically use samples and supervised learning here.

Train the model to operate both conditionally and unconditionally. During training, the condition is sometimes supplied and sometimes dropped.

The model therefore learns an unconditioned field $v_{\theta }(a_t,t,s,\varnothing )$ and an optimality-conditioned field $v_{\theta }(a_t,t,s,o=1)$. Here “unconditioned” still includes the state $s$.

Again, via Bayes:

Add some $w$-scaled version of that vector: from the unconditional score towards the conditional score. Thus, for the fields we trained:

or equivalently:

The product-distribution interpretation is:

If the optimality factor is proportional to $\exp (A(s,a))$, then $w=1/\tau$ gives the earlier exponential reweighting. If it already uses $\exp (A/\tau )$, then $w=1$ gives that target.

This suggests a way to learn the unconditional and conditional fields, then combine them at inference with a chosen $w$.

What CFGRL does in offline RL

1. Learn a value estimate to get advantage.
2. Turn advantage into an optimality condition, e.g. $o(s,a)=\mathbf{1}[A(s,a)>0]$.
3. Train a generative policy conditioned on $s$ and $o$, but drop $o$ sometimes so it learns both.
4. Apply CFG during sampling:

Thus, the value function is just for labels and we don’t have to backpropagate through it during sampling.

Difference from advantage-weighted regression

AWR controls the magnitude of the training loss. Looks like REINFORCE, but here we fit logged actions with positive exponential weights. Standard REINFORCE uses current-policy samples and return / advantage weights. CFGRL uses supervised learning where the condition is different.

It learns an “optimality direction” and amplifies it at inference.

RECAP: shares this advantage-conditioning idea with CFGRL. They also do policy iteration: repeatedly deploy, recollect, redo advantage, update policy, etc. They use a task-dependent advantage threshold and usually just sample the positive condition, rather than turning up CFG.

I would say the benefit of not just learning the optimal actions is that it is like more general knowledge: even things not directly to do with what we want can help make the model better, so that you can then, for example, generate better samples. Also can tune your regularisation, so better than just the KL-regularised fit. (My intuition.)

FQL: also regularises the improved policy towards a behaviour model. Here the penalty is squared distance between the actions they produce from the same noise input, rather than KL. The behaviour model is an iterative flow policy, and the learned policy generates an action in one step.