Off-Policy Actor-Critic - Towards Q-learning

This page picks up where the actor-critic notes left off. So far everything has been on-policy, so every gradient step needs fresh data. Here we cross to off-policy and reuse old transitions from a replay buffer. The key move is to learn a Q-function, which lets the actor and the critic train on stale data, and can make the actor disappear when using a greedy policy. This is the doorway to Q-learning on the next page. This page has no result of its own, it is purely a technical bridge. It follows the end of Lecture 6: Actor-Critic (from slide 26) of CS 285.

Preamble

REINFORCE, actor-critic, GAE: all of them are on-policy. Each gradient step needs samples from the current $\pi_\theta$, so the moment we update $\theta$ the data we just collected is stale and we throw it away. That is wasteful, because most of the cost in RL is simulation, and we are using each transition exactly once.

The thing we want is off-policy learning: keep many past transitions in a replay buffer and reuse them many times. We can get there without giving up the actor-critic structure, by changing what the critic learns. For a discrete system and with this new critic we can greedily pick the next action, bypassing the need for an actor altogether. This is the bridge to Q-learning we will start discussing on this page.

Why naively sampling off-policy breaks Policy Gradient

Imagine keeping the same online actor-critic we already saw and just adding a replay buffer we can sample from. A replay buffer is simply a bank of $(s,a,s')$ transitions whose origin does not matter: the action selected does not need to follow our current policy. If we sample from this buffer we get:

Actor-critic with a replay buffer (broken)
repeat take a step with $a \sim \pi_\theta(a \mid s)$, store $(s_i, a_i, s_i')$ in buffer $\mathcal{R}$ sample a batch $\{(s_i, a_i, s_i')\}$ from $\mathcal{R}$ targets $y_i = r(s_i, a_i) + \gamma V_\varphi(s_i')$ update $\varphi$ toward $\{y_i\}$ advantages $\hat{A}(s_i, a_i) = r(s_i, a_i) + \gamma V_\varphi(s_i') - V_\varphi(s_i)$ $\nabla_\theta J(\theta) \approx \sum_i \nabla_\theta \log \pi_\theta(a_i \mid s_i)\, \hat{A}(s_i, a_i)$ $\theta \leftarrow \theta + \alpha \nabla_\theta J(\theta)$ until converged

For simplicity we can assume the stale data in the buffer is as if it was sampled by an old policy $\bar{\pi}$. Two things break:

  1. The actor update $\nabla_\theta \log \pi_\theta(a_i \mid s_i)\, \hat{A}(s_i, a_i)$ pushes on $a_i$, but $a_i$ is not an action $\pi_\theta$ would take now. We are reinforcing the old policy's choices, not the current one.
  2. The critic is also wrong. Recall that we fit $V_\varphi$ by least-squares regression,
$$\mathcal{L}(\varphi) = \frac{1}{2} \sum_i \big\| V_\varphi(s_i) - y_i \big\|^2, \qquad y_i = r(s_i, a_i) + \gamma V_\varphi(s_i'),$$

and a squared-error fit pulls $V_\varphi(s_i)$ toward the average of its targets at $s_i$ (minimizing squared error gives the conditional mean of the target, see the probability facts). We want that average to be the value under the current policy,

$$V^\pi(s_i) = \mathbb{E}_{a \sim \pi_\theta(a \mid s_i)}\big[\, r(s_i, a) + \gamma V^\pi(s') \,\big],$$

but the actions in the buffer came from $\bar{\pi}$, so what we actually fit toward is

$$\mathbb{E}_{a \sim \bar{\pi}(a \mid s_i)}\big[\, r(s_i, a) + \gamma V^\pi(s') \,\big],$$

the value of the old policy at $s_i$. Another way to see it. The single target $y_i = r(s_i, a_i) + \gamma V_\varphi(s_i')$ is really a one-sample estimate of $Q(s_i, a_i)$, the value of the specific action $a_i$. Forcing $V_\varphi(s_i)$ to match it averages those $Q$ values over whatever actions the buffer holds, which is $\bar{\pi}$, not $\pi_\theta$. The action $a_i$ is baked into the target, but $V_\varphi(s_i)$ has no way to record which action it was, so the wrong distribution leaks in.

So the question is how to write a target that does not care which policy chose $a_i$.

The critic: learn a Q-function

The fix is to stop averaging over the action. Learn $Q_\varphi(s_i, a_i)$ instead of $V_\varphi(s_i)$, with the action as an input. The regression looks almost the same,

$$\mathcal{L}(\varphi) = \frac{1}{2} \sum_i \big\| Q_\varphi(s_i, a_i) - y_i \big\|^2,$$

but now $a_i$ is part of what we condition on, not something we average over. Different actions at the same state land at different inputs $(s_i, a_i)$, so the old policy's action distribution does not leak into the fit. We are asking "what is the value of taking $a_i$ in $s_i$," and that has the same answer no matter who picked $a_i$.

The target still has to bootstrap, and now we get to choose the next action ourselves. We ask the latest policy what it would do at $s_i'$:

$$y_i = r(s_i, a_i) + \gamma\, \mathbb{E}_{a' \sim \pi_\theta(a' \mid s_i')}\big[Q_\varphi(s_i', a')\big] \;\approx\; r(s_i, a_i) + \gamma\, Q_\varphi(s_i', a_i'), \qquad a_i' \sim \pi_\theta(a' \mid s_i').$$

In practice we estimate that expectation with a single action $a_i' \sim \pi_\theta$ sampled at $s_i'$. The transition $(s_i, a_i, s_i')$ is old, but $a_i'$ is drawn from the current policy, so the target tracks $\pi_\theta$ even on stale data.

Why is reusing old data allowed? Given $(s, a)$, the next state is drawn by the environment, $s' \sim p(s' \mid s, a)$, which has nothing to do with the policy. A stored transition $(s, a, s')$ is therefore just as valid now as when it was collected. The only policy-dependent piece is "what do we do at $s'$," and we recompute that on the fly. This is the whole reason off-policy learning is possible.

The actor update

The critic above scores any action at any state. Now the actor, the part that has to improve. It follows the same rule the critic target did: anything that asks "what action would we take" has to use the current policy, not the buffer's. The policy gradient wants actions from $\pi_\theta$, but the buffer's $a_i$ came from an old policy, so we ignore it and at each buffered state $s_i$ sample a fresh action $a_i^\pi \sim \pi_\theta(a \mid s_i)$ to push on instead:

$$\nabla_\theta J(\theta) \approx \frac{1}{N} \sum_i \nabla_\theta \log \pi_\theta\big(a_i^\pi \mid s_i\big)\, Q_\varphi\big(s_i, a_i^\pi\big).$$

The state is from the buffer but the action is on-policy, so the actor only ever sees actions $\pi_\theta$ would take now. And since $a_i^\pi$ is sampled, not executed in the world, this costs no environment steps, just a forward pass through $\pi_\theta$ and $Q_\varphi$.

Why $Q$ here, and not the advantage? The earlier actor-critic scaled $\nabla_\theta \log \pi_\theta$ by the advantage $\hat{A} = Q - V$, not by $Q$. The $-V(s)$ piece is a baseline: it depends only on the state, and a state-only baseline has zero mean under the gradient, $\mathbb{E}_{a \sim \pi_\theta}[\nabla_\theta \log \pi_\theta(a \mid s)\, b(s)] = b(s)\, \nabla_\theta \sum_a \pi_\theta(a \mid s) = 0$. So $Q$ and $A = Q - V$ give the same expected gradient, the advantage version just has lower variance. Here we drop the baseline and use $Q$ directly: higher variance, but we already have $Q_\varphi$ and would otherwise have to fit a separate $V$ as well. The fresh-sampled $a_i^\pi$ and the large buffer batch keep the variance manageable.

Recall that this score-function gradient is just one way to estimate

$$\nabla_\theta\, \mathbb{E}_{a \sim \pi_\theta(a \mid s_i)}\big[\, \hat{Q}^{\pi}_\varphi(s_i, a) \,\big],$$

the gradient of the expected value of the actions $\pi_\theta$ takes, written out with the score function estimator. For continuous actions we can instead differentiate that expectation directly, which has lower variance. Write the actor objective as that value averaged over states from the buffer:

$$J(\theta) = \mathbb{E}_{s_i \sim \mathcal{R}}\Big[\; \mathbb{E}_{a \sim \pi_\theta(a \mid s_i)}\big[\, Q_\varphi(s_i, a) \,\big] \;\Big].$$

The inner expectation is over $\pi_\theta$, which depends on $\theta$. For a Gaussian policy we differentiate through it with the reparametrization trick (see the probability facts): write $a = \mu_\theta(s_i) + \sigma_\theta(s_i)\, \epsilon$ with $\epsilon \sim \mathcal{N}(0, I)$, so the noise no longer depends on $\theta$ and the gradient passes straight into $Q$. With one noise sample $\epsilon_i$ per state,

$$\nabla_\theta J(\theta) \approx \frac{1}{N} \sum_i \nabla_\theta\, Q_\varphi\big(s_i,\; \mu_\theta(s_i) + \sigma_\theta(s_i)\, \epsilon_i\big).$$

By the chain rule this moves $\theta$ in the direction the critic says raises $Q$, and it is usually less noisy than the score function because it uses the slope of $Q$ in the action, not just $\log \pi_\theta$. SAC and DDPG/TD3 are built on this gradient.

Putting it together

The critic and the actor in one loop give the off-policy actor-critic:

Off-policy reparameterized actor-critic
repeat get $(s_i, a_i, s_i')$ by taking one step with $a \sim \pi_\theta(a \mid s)$, store it in $\mathcal{R}$ sample a batch $\{(s_i, a_i, s_i')\}$ from $\mathcal{R}$ critic$y_i = r(s_i, a_i) + \gamma\, \mathbb{E}_{a' \sim \pi_\theta}[Q_\varphi(s_i', a')] \;\approx\; r(s_i, a_i) + \gamma\, Q_\varphi(s_i', a_i')$  # one sample $a_i' \sim \pi_\theta(a' \mid s_i')$ criticupdate $\varphi$ using $\nabla_\varphi \sum_i \| Q_\varphi(s_i, a_i) - y_i \|^2$ actor$\nabla_\theta J(\theta) \approx \sum_i \nabla_\theta\, Q_\varphi\big(s_i,\, \mu_\theta(s_i) + \sigma_\theta(s_i)\, \epsilon_i\big)$ actor$\theta \leftarrow \theta + \alpha\, \nabla_\theta J(\theta)$ until converged
One simplification we are making. The states $s_i$ are sampled from the buffer, so they follow the state distribution of the old policies, not the stationary distribution of $\pi_\theta$. Strictly the policy gradient wants states from $\pi_\theta$'s own distribution, so we are improving the policy on a slightly wrong set of states. There is no cheap fix without a model, and in practice we ignore it: we get a policy that is good across the broad set of states in the buffer, which is usually what we wanted anyway.

This is just the skeleton of an off-policy proper actor-critic model and, as we will see in the following notes, requires a few tricks to work properly in practice. This derivation was done more as a pedagogical exercise to show the bridge from policy gradient to DQN and the similarities between them (underneath it's always the same Bellman backbone). Let's pause here and move on to discrete DQN next and build towards real implementations of this for continuous systems (SAC/TD3/DDPG).

References