Reinforcement LearningDistributed RLIMPALAApe-X

Distributed RL Architectures: IMPALA, Ape-X, R2D2, and SEED RL

How reinforcement learning is scaled across thousands of parallel actors — IMPALA's V-trace correction for off-policy staleness, Ape-X's distributed prioritized replay, R2D2's recurrent extension, and SEED RL's centralized-inference architecture — with code and applications.

TL;DR

A single environment instance generates experience too slowly to train large RL agents efficiently — the fix is running thousands of environment copies in parallel, but that introduces a new problem: by the time an actor’s data reaches the learner, the learner’s policy has already moved on, making the data stale and technically off-policy. IMPALA solves this with the V-trace correction, letting many actors run asynchronously against one central learner. Ape-X decouples data generation from learning entirely, with distributed actors filling a shared prioritized replay buffer for an off-policy learner. R2D2 extends that architecture to recurrent policies, solving the specific problem of replaying partial sequences correctly. SEED RL goes further architecturally, centralizing neural network inference itself on the learner to cut communication costs. Used anywhere RL needs to scale past what a single machine and a single environment instance can produce — large-scale game-playing agents and industrial-scale simulation training.

The Distributed RL Lineage

graph LR
    A[A3C: lockstep actor-learner] --> B[IMPALA: async actors + V-trace]
    B --> C[Ape-X: decoupled via shared replay]
    C --> D[R2D2: recurrent-aware replay]
    C --> E[SEED RL: centralized inference]

    style B fill:#6366f1,color:#fff

Problem Statement

A3C (covered in the companion actor-critic article) was an early attempt at parallelizing RL across CPU workers, but it ties data collection and gradient computation to the same lockstep loop. Modern large-scale RL wants a cleaner separation: many actors just generating experience as fast as possible, and one or few learners consuming that experience efficiently — but the moment actors and learners are decoupled and run at different speeds, the data an actor generated under an old policy version gets used to train a newer one, and naive off-policy correction (or none at all) either biases training or wastes the parallelism.

IMPALA: V-trace for Off-Policy Correction at Scale

IMPALA (Importance Weighted Actor-Learner Architecture, Espeholt et al., 2018) runs many actors continuously generating trajectories with their own (slightly stale) copy of the policy, streaming them to a central learner that updates as fast as it can — actors never wait for the learner’s latest weights. The staleness this introduces is corrected with V-trace, an importance-sampling-based off-policy correction:

def v_trace_targets(behavior_log_probs, target_log_probs, rewards, values, gamma=0.99, rho_bar=1.0, c_bar=1.0):
    rho = torch.clamp(torch.exp(target_log_probs - behavior_log_probs), max=rho_bar)  # clipped IS ratio
    c = torch.clamp(torch.exp(target_log_probs - behavior_log_probs), max=c_bar)

    deltas = rho * (rewards + gamma * values[1:] - values[:-1])

    vs = values[:-1].clone()
    for t in reversed(range(len(deltas))):
        vs[t] = values[t] + deltas[t] + gamma * c[t] * (vs[t + 1] - values[t + 1] if t + 1 < len(vs) else 0)
    return vs

Clipping the importance ratio (rho_bar, c_bar) is what keeps V-trace stable — without the clip, a large staleness gap could produce an enormous importance weight and destabilize training; with it, IMPALA trades a small amount of bias for a large amount of stability, which is exactly the trade-off that makes fully asynchronous, high-throughput data collection practical.

Ape-X: Decoupled Actors and a Shared Prioritized Replay Buffer

Ape-X (Horgan et al., 2018) takes decoupling further: many actors (each with a slightly different exploration rate, for diversity) run independently, computing their own local TD-errors and pushing transitions into a large, shared prioritized experience replay buffer (the same idea covered in the companion Q-learning/Rainbow article, now distributed). A single learner samples from that shared buffer continuously, completely decoupled from actor speed.

def apex_actor_loop(actor_id, env, local_policy, shared_replay_buffer, epsilon):
    state = env.reset()
    while True:
        action = epsilon_greedy_action(local_policy(state), epsilon)  # each actor: different epsilon
        next_state, reward, done = env.step(action)
        td_error = compute_local_td_error(local_policy, state, action, reward, next_state)
        shared_replay_buffer.add(state, action, reward, next_state, priority=abs(td_error))
        state = next_state if not done else env.reset()

Because each actor uses a different fixed exploration rate, the shared buffer ends up covering a much broader range of the state space than any single exploration schedule would, which is a secondary benefit on top of the pure throughput gain from parallel data collection.

R2D2: Extending Distributed Replay to Recurrent Policies

Prioritized replay assumes you can sample individual transitions independently, which breaks for recurrent policies — an LSTM’s hidden state depends on everything that came before it in the sequence, so replaying a transition out of context gives the network a wrong or missing hidden state. R2D2 (Recurrent Replay Distributed DQN, Kapturowski et al., 2019) solves this by storing and replaying whole sequences rather than individual transitions, and using a short “burn-in” period at the start of each replayed sequence purely to let the recurrent state re-converge before computing any loss:

def r2d2_sequence_replay(lstm_policy, sequence, burn_in_steps=40):
    hidden_state = lstm_policy.init_hidden()

    # Burn-in: run the LSTM forward without computing gradients, just to warm up hidden state
    with torch.no_grad():
        for t in range(burn_in_steps):
            _, hidden_state = lstm_policy(sequence[t].state, hidden_state)

    # Real training: compute loss only on steps after burn-in, with a properly warmed-up hidden state
    losses = []
    for t in range(burn_in_steps, len(sequence)):
        q_values, hidden_state = lstm_policy(sequence[t].state, hidden_state)
        losses.append(compute_td_loss(q_values, sequence[t]))
    return sum(losses)

SEED RL: Centralizing Inference to Cut Communication Cost

Ape-X-style architectures still run the neural network’s forward pass on each actor machine to pick actions, which means shipping network weights out to every actor whenever the policy updates. SEED RL (Espeholt et al., 2020) inverts this: actors only run the environment simulation and ship raw observations to the learner; the learner runs all inference centrally (using very fast RPC round-trips) and ships back only the chosen action. This cuts the amount of data that needs to move around per step dramatically, since raw observations are typically far smaller than full network weight updates, and it means the very latest policy is always used for every single action, eliminating staleness at the source rather than correcting for it after the fact.

Comparison Table

ArchitectureActor-Learner CouplingStaleness HandlingRecurrent Support
IMPALAAsynchronous, actors run their own policy copyV-trace off-policy correctionStandard, no special handling
Ape-XDecoupled via shared prioritized replay bufferBuffer absorbs staleness; corrected via off-policy learningNot designed for recurrence
R2D2Same as Ape-XSame as Ape-XSequence replay with burn-in
SEED RLCentralized inference; actors only simulateEliminated at the source (always latest policy)Compatible, inference-side

Applications

Key Learnings

  1. Scaling RL is fundamentally a staleness-management problem, not just a “run more copies” problem. Every architecture here exists to answer the same question differently: what do you do about the gap between the policy that generated some data and the policy currently being trained?
  2. Decoupling data generation from learning unlocks throughput that lockstep architectures (like A3C) can’t reach, but it requires deliberate off-policy correction (V-trace) or an off-policy-tolerant learning algorithm (Ape-X’s prioritized DQN-style learner) to stay stable.
  3. Where computation happens is itself a scaling lever. SEED RL’s insight — move inference to the learner instead of shipping weights to actors — is a reminder that distributed systems design choices (not just RL algorithm choices) are often what determines whether an architecture actually scales in practice.

References