Off-Policy Correction at Scale

"By the time my experience reached the learner, the policy that produced it had already moved on. I am not wrong, exactly; I am just slightly out of date, and someone needs to weigh me accordingly."

An Actor Reporting a Trajectory Two Updates Late

The previous sections of this chapter built a machine for throughput. Section 20.2 split the agent into actors that interact with the environment and a learner that updates the policy, so that slow simulation and fast gradient steps could proceed in parallel. Section 20.3 multiplied the actors to flood the learner with experience, and Section 20.4 added a replay buffer so that experience could be stored and reused rather than consumed once and discarded. Each of those moves bought scale, and each of them quietly broke an assumption that on-policy policy-gradient methods depend on: that the data used to compute an update was generated by the very policy the update is improving. By the time a trajectory travels from a lagging actor through a buffer to the learner, the learner's policy has changed. The data is off-policy, and we now have to say what that costs and how to fix it. Figure 20.5.1 lays out that price and its remedy: experience collected under older behavior policies, reweighted by a truncated importance weight before it touches the gradient.

1. Decoupling Buys Throughput and Sells Off-Policyness Beginner

On-policy policy-gradient methods rest on a single expectation. To improve a policy $\pi_\theta$, you estimate the gradient of the expected return as an average of $R(a)\,\nabla_\theta \log \pi_\theta(a)$ over actions drawn from $\pi_\theta$ itself. The phrase "drawn from $\pi_\theta$ itself" is the load-bearing one. It says the data must come from the current policy. In a single-process agent that is automatic: the same policy acts and learns, one step at a time, and the data is on-policy by construction.

Distribution breaks that automatically. The actor-learner architecture of Section 20.2 exists precisely so that actors do not wait for the learner. An actor pulls a copy of the policy, runs an entire episode in the environment, and ships the trajectory back. While that episode was running, the learner took several gradient steps, so the policy the learner now holds, the target policy $\pi$, is newer than the policy the actor used, the behavior policy $\mu$. Add a replay buffer and the gap widens: a sampled trajectory may have been generated tens or hundreds of updates ago, under a $\mu$ that is now distinctly stale. The throughput you gained is exactly the lag you must now correct. This is the same staleness that haunts distributed optimization, where a worker computes a gradient against parameters the server has already moved past; we studied it as delayed gradients in Section 10.6, and here it returns wearing an RL costume.

2. Importance Sampling Repairs the Expectation Intermediate

The classical fix for evaluating an expectation under one distribution using samples from another is importance sampling. If we want $\mathbb{E}_{a \sim \pi}[f(a)]$ but only have samples from $\mu$, we reweight each sample by the ratio of probabilities,

The ratio $\rho(a)$ is the importance weight. An action that the new policy $\pi$ favors more than the old policy $\mu$ did gets up-weighted; an action $\pi$ now avoids gets down-weighted. Applied to the policy gradient, this turns the off-policy data back into an unbiased estimate of the on-policy gradient: in principle, the lag is fully corrected and you can train the current policy on arbitrarily old experience. That is the clean theory, and for small policy gaps it works.

The trouble is variance. The ratio $\rho(a)$ has no upper bound. When the behavior policy assigned a tiny probability $\mu(a)$ to an action that the target policy now likes, $\pi(a)/\mu(a)$ becomes enormous, and a single such sample can swamp the entire gradient estimate. The more off-policy the data, the heavier the tail of these ratios, and the wider the gap our throughput-seeking decoupling created, the worse this gets. An unbiased estimator whose variance explodes is useless in practice: the learner's updates become so noisy that training stalls or diverges. Scaling out forced the off-policyness; the off-policyness forced importance sampling; and importance sampling, untamed, forces a variance problem of its own.

3. V-trace: Truncate the Weights to Tame the Variance Advanced

The systems-relevant remedy is the one IMPALA introduced under the name V-trace (Espeholt et al., 2018): cap the importance weights so no single sample can dominate. Instead of the raw ratio $\rho(a)$, V-trace uses a truncated importance weight,

where $\bar{\rho}$ (often $1$) is a fixed ceiling. The truncation throws away the rare, gigantic ratios that carry almost all the variance, while leaving the typical, moderate ratios intact. V-trace actually carries two such clipped weights: $\bar{\rho}_t$ controls the fixed point of the value estimate that the policy is improved against, and a separate clipped weight $\bar{c}_t = \min(\bar{c}, \pi(a_t)/\mu(a_t))$ controls how much trace credit propagates back across time steps. For the systems story the essential idea is one knob: clip the importance weight, and you trade a controlled amount of bias for a large reduction in variance, which is exactly the trade a scalable learner needs when its data arrives from many lagging actors at once.

This is a deliberate, asymmetric bargain. Truncation introduces bias (a clipped estimator no longer targets the exact on-policy gradient) but slashes variance, and for stable learning at scale the variance reduction is worth far more than the bias costs. The on-policy fixed point is preserved when the data is fresh ($\mu = \pi$ makes every ratio $1$, below the cap, so V-trace reduces to ordinary on-policy learning), and the bias grows gracefully as the policies drift apart rather than letting a single rare sample blow up the update. That graceful degradation is what lets one learner safely consume experience from a swarm of actors that are each a different distance behind.

The code below makes the whole arc concrete in a few lines of pure NumPy: a tiny single-state policy-gradient problem where the on-policy gradient is known exactly, a lagging behavior policy that makes the data off-policy, and a comparison of three estimators, the naive uncorrected gradient, the V-trace-style truncated importance weight, and the untruncated importance weight. It reports both the bias against the true gradient and the variance of each estimator, so the bias-variance bargain is visible in numbers.

import numpy as np

# A tiny 1-state, 2-action policy-gradient setup. The "target" policy pi is what
# the learner currently holds; the "behavior" policy mu is the slightly OLDER
# policy that the lagging actors actually used to collect experience.
rng = np.random.default_rng(0)
A = 2                                   # two actions
theta = np.array([0.0, 0.0])           # current (target) policy logits
def softmax(z):
    z = z - z.max(); e = np.exp(z); return e / e.sum()

pi = softmax(theta)                     # target policy pi(a)
reward = np.array([1.0, 3.0])           # action 1 is the better action

# Exact on-policy gradient: grad J = E_{a~pi}[ R(a) * grad log pi(a) ],
# with grad log pi(a) = e_a - pi for a softmax policy.
def grad_log_pi(a):
    g = -pi.copy(); g[a] += 1.0; return g
true_grad = sum(pi[a] * reward[a] * grad_log_pi(a) for a in range(A))

def run(stale_gap, n=200_000, clip=None, want_var=False):
    mu = softmax(theta - np.array([0.0, stale_gap]))   # older, staler policy
    acts = rng.choice(A, size=n, p=mu)                 # actors sampled from mu
    rho = pi[acts] / mu[acts]                          # raw importance ratios
    w = np.ones(n) if clip is None else np.minimum(rho, clip)  # V-trace truncation
    gmat = np.array([reward[a] * grad_log_pi(a) for a in range(A)])
    samples = w[:, None] * gmat[acts]
    est = samples.mean(axis=0)
    return (est, samples.var(axis=0).sum()) if want_var else est

def bias(est):
    return np.linalg.norm(est - true_grad) / np.linalg.norm(true_grad)

print("on-policy true grad        :", np.round(true_grad, 4))
for gap in (0.0, 2.0):                                  # 0 = on-policy, 2 = stale
    mu = softmax(theta - np.array([0.0, gap]))
    naive             = run(gap, clip=None)             # ignore the mismatch
    vtrace, var_tr    = run(gap, clip=3.0, want_var=True)   # truncated IS (V-trace)
    full_is, var_full = run(gap, clip=1e9, want_var=True)   # untruncated IS
    print(f"\nstaleness gap {gap:>4}  behavior mu = {np.round(mu,3)}")
    print(f"  naive (no correction)    : {np.round(naive,4)}  rel bias {bias(naive):.3f}")
    print(f"  truncated IS (V-trace)   : {np.round(vtrace,4)}  rel bias {bias(vtrace):.3f}  var {var_tr:.3f}")
    print(f"  untruncated IS (unbiased): {np.round(full_is,4)}  rel bias {bias(full_is):.3f}  var {var_full:.3f}")

on-policy true grad        : [-0.5  0.5]

staleness gap  0.0  behavior mu = [0.5 0.5]
  naive (no correction)    : [-0.4965  0.4965]  rel bias 0.007
  truncated IS (V-trace)   : [-0.5005  0.5005]  rel bias 0.001  var 2.000
  untruncated IS (unbiased): [-0.5018  0.5018]  rel bias 0.004  var 2.000

staleness gap  2.0  behavior mu = [0.881 0.119]
  naive (no correction)    : [ 0.2598 -0.2598]  rel bias 1.520
  truncated IS (V-trace)   : [-0.292  0.292]  rel bias 0.416  var 4.846
  untruncated IS (unbiased): [-0.5056  0.5056]  rel bias 0.011  var 9.135

The numbers tell the systems story precisely. The naive estimator does not just lose accuracy under lag; with a relative bias of $1.52$ it has flipped the sign of the dominant gradient component, so an uncorrected learner would push probability toward the worse action, the action the stale behavior policy happened to sample more often. The untruncated importance weight repairs that almost perfectly but pays in variance, and in a real learner consuming millions of samples from a long-tailed ratio distribution, that variance is what makes training unstable. V-trace's truncation is the engineering compromise that an at-scale learner actually ships: enough correction to point the gradient the right way, little enough variance to keep the updates usable.

4. More Decoupling, More Correction Intermediate

The reason V-trace matters for distributed systems, rather than only for RL theory, is that it converts a hard synchronization constraint into a soft, tunable one. Without an off-policy correction, the only safe way to keep actors on-policy is to make them wait for the learner, which destroys the throughput the actor-learner split was built to gain. A fully synchronous system forces every actor to use the current policy, paying for correctness with idle time on a barrier. V-trace removes that barrier: actors may run as far ahead as the correction can tolerate, the learner reweights whatever arrives, and the system trades a little statistical efficiency for a great deal of parallel throughput. This is the lever that lets IMPALA scale to hundreds of actors feeding a single learner, and it is why the off-policy correction is an infrastructure decision, not just a learning-rule detail.

The trade has limits, and they are systems limits. The clipping ceiling $\bar{\rho}$ sets how off-policy the learner will tolerate data before the bias from truncation outweighs the benefit. Push the actors too far ahead, or let the replay buffer hold experience too long, and even truncated weights cannot rescue an update whose behavior policy bears little resemblance to the target. So the off-policy correction does not abolish the cost of decoupling; it gives you a dial to manage it. How far can actors lag? How deep can the buffer go? Those become quantitative questions answered against the correction's tolerance, the same way Chapter 2 framed the synchronous-versus-asynchronous choice as a spectrum rather than a binary. The next section turns this lever into full system designs, showing how Ape-X, R2D2, and SEED RL each place the actor-learner boundary differently and rely on exactly this kind of correction to hold it together.

from ray.rllib.algorithms.impala import IMPALAConfig

# Many lagging actors (rollout workers) feed one learner; RLlib applies V-trace
# to every batch automatically, so the actors never block on the learner.
config = (
    IMPALAConfig()
    .environment("CartPole-v1")
    .env_runners(num_env_runners=64)      # 64 decoupled actors collecting experience
    .training(
        vtrace=True,                      # turn on the V-trace off-policy correction
        vtrace_clip_rho_threshold=1.0,    # rho_bar: caps the policy-gradient weight
        vtrace_clip_pg_rho_threshold=1.0, # caps the trace-credit weight c_bar
    )
)
algo = config.build()
for _ in range(100):
    algo.train()                          # learner consumes off-policy batches safely

5. Where the Correction Sits in the System Beginner

It is worth being precise about who does what, because the correction is cheap and lives entirely on the learner. The actors do one extra thing: when they collect a trajectory, they record the action probabilities $\mu(a_t)$ under the behavior policy they used, and ship those alongside the states, actions, and rewards. That is a handful of floats per step, negligible next to the trajectory itself. The learner, holding the current target policy $\pi$, recomputes $\pi(a_t)$ for each logged action, forms the ratio $\pi(a_t)/\mu(a_t)$, truncates it, and folds it into the gradient and value targets. No extra round trips, no new collective, no actor-learner handshake beyond the policy pulls and experience pushes that Section 20.4 already described.

That placement is the whole reason the correction scales. It does not add communication, which is the tax this book spends every chapter trying to avoid; it adds a small per-sample computation on the one machine that was going to process the sample anyway. The actors stay simple and embarrassingly parallel, the learner stays the single point of policy truth, and the only new wire-level cost is logging $\mu(a_t)$. An off-policy correction that required actors to coordinate, or that introduced a barrier, would defeat its own purpose; V-trace earns its place precisely because it corrects the lag without reintroducing the synchronization the lag was meant to avoid.