Part II - Applications and libraries

Your First Distributed Application

Part I gave you tasks, actors, the object store, and the map-shuffle-reduce pattern (lesson 05). You have never yet wired them into one running application. This lesson does — and deliberately picks reinforcement learning, because RL forces every primitive into the open at once: a stateful simulator (actor), a read-mostly policy (shared object), a parallel fan-out (rollout collection), and a reduce with a feedback loop (the learner). We build it by hand so that when lesson 07 hands you RLlib, you recognize the machine inside the box instead of trusting a black one.

1 · The four roles RL drags into the open

Reinforcement learning is a loop: an agent acts in an environment, the environment returns a reward, and over many tries the agent learns to act so as to maximize cumulative reward. The vocabulary, defined once:

Hold those last two against lesson 05. Rollout is the map (fan out independent work), pooling trajectories is the shuffle/reduce (aggregate experience), and the learner update is a reduce whose output feeds back to seed the next map. RL is map-shuffle-reduce with a loop closed around it. Once you see that, on-policy vs off-policy is just a question about that loop: on-policy means the learner trains only on data generated by the very latest policy (each round's data is thrown away after one update); off-policy means it may train on data generated by older policies (data can be reused from a buffer). We will lean lightly off-policy — simulators run a slightly stale policy snapshot while the learner updates — because that is exactly the staleness/utilization trade-off Ray makes you confront.

2 · The single-process maze — your correctness oracle

Before any @ray.remote, write the thing that has no parallelism and no chance of a race. It is slow, but it is right, and every distributed version must reproduce its learning curve. This is the same discipline as lesson 00's "Ray is not a speed button": correctness first, then parallelize the part that is actually the bottleneck.

import numpy as np

class MazeEnv:
    def __init__(self, n=8):
        self.n, self.goal = n, n*n - 1          # 8x8 grid, goal = bottom-right
    def reset(self):
        self.s = 0; return self.s               # start top-left
    def step(self, a):                          # a in {0:up,1:down,2:left,3:right}
        r, c = divmod(self.s, self.n)
        if   a == 0: r = max(0, r-1)
        elif a == 1: r = min(self.n-1, r+1)
        elif a == 2: c = max(0, c-1)
        else:        c = min(self.n-1, c+1)
        self.s = r*self.n + c
        done = (self.s == self.goal)
        return self.s, (0.0 if done else -1.0), done   # -1 per move; reach goal to stop

def run_episode(env, Q, eps=0.1, max_steps=200):
    s, traj = env.reset(), []
    for _ in range(max_steps):
        a = np.random.randint(4) if np.random.rand() < eps else int(Q[s].argmax())
        s2, r, done = env.step(a)
        traj.append((s, a, r, s2)); s = s2
        if done: break
    return traj                                  # the trajectory

# single-process training loop: rollout -> learn -> repeat
Q = np.zeros((64, 4))                            # the policy: one row per state
for it in range(500):
    env = MazeEnv()
    batch = [t for _ in range(16) for t in run_episode(env, Q)]   # 16 episodes
    for (s, a, r, s2) in batch:                  # Q-learning update (the learner)
        Q[s, a] += 0.1 * (r + 0.95*Q[s2].max() - Q[s, a])

The eps knob is exploration: 10% of the time the agent acts randomly instead of greedily, so it discovers paths the current Q-table would never try. Everything distributed below keeps this exact update; we only change who runs the rollouts and where the Q-table lives.

3 · Distribute it — actors for state, the driver for the reduce

Two facts decide the layout. First, the maze env here is cheap, but in any realistic problem the environment is expensive to build (load a simulator, open a connection, allocate a board), so we do not want to recreate it per rollout — that argues for a long-lived simulator actor that constructs its env once and is called many times. Second, the policy is read by every simulator and written by one learner — classic read-mostly shared state — so we host it in a parameter-server actor whose only jobs are get_weights() and apply_update().

import ray
ray.init()

@ray.remote
class RolloutWorker:                             # one stateful simulator
    def __init__(self, n=8):
        self.env = MazeEnv(n)                    # built ONCE, reused every call
    def sample(self, Q, episodes=16):            # Q arrives auto-dereferenced
        return [t for _ in range(episodes) for t in run_episode(self.env, Q)]

@ray.remote
class ParameterServer:                           # holds the read-mostly policy
    def __init__(self, n_states=64, n_actions=4):
        self.Q = np.zeros((n_states, n_actions))
    def get_weights(self):  return self.Q        # served zero-copy on-node
    def apply(self, batch, lr=0.1, gamma=0.95):  # the learner / reduce step
        for (s, a, r, s2) in batch:
            self.Q[s, a] += lr*(r + gamma*self.Q[s2].max() - self.Q[s, a])
        return self.Q

ps = ParameterServer.remote()
workers = [RolloutWorker.remote() for _ in range(8)]

for it in range(500):
    w_ref = ps.get_weights.remote()              # ObjectRef to the current policy
    w_ref = ray.put(ray.get(w_ref))              # pin one snapshot, share its ref
    futures = [wk.sample.remote(w_ref) for wk in workers]   # fan out: the MAP
    # collect as they finish instead of blocking on the slowest straggler:
    batch, pending = [], futures
    while pending:
        done, pending = ray.wait(pending, num_returns=1)    # first-done streaming
        batch += ray.get(done[0])
    ray.get(ps.apply.remote(batch))              # the REDUCE + feedback edge

Three things to notice, each a lesson-05 idea made real. We ray.put the policy once per iteration and pass the ref to all 8 workers, so the Q-table is serialized one time, not eight (lesson 00's "big objects through small tasks" anti-pattern, avoided). We use ray.wait rather than ray.get(futures) so a slow worker does not stall the whole batch — we fold in trajectories as they arrive. And the learner update is the reduce, whose output (the new Q) becomes the input to the next round: the feedback edge.

4 · The bottleneck migrates as you scale — the real engineering point

Here is the lesson the book only gestures at. Each training iteration has a parallel part (rollout, which divides across workers) and a serial part that does not shrink no matter how many workers you add: getting the policy snapshot, the ray.put broadcast, and the learner's apply over the pooled batch. Call the serial part the coordination tax. This is Amdahl's law (lesson 00) wearing an RL costume.

Put numbers on it. Suppose total rollout work per iteration is fixed at W = 8 × 24 ms = 192 ms of simulation, split across the workers. The serial coordination tax per iteration is roughly: policy fetch + put ≈ 2 ms, plus the learner update, which grows with batch size — say the pooled batch is ~480 transitions per worker and the learner does ~0.6 µs per transition, so learner time ≈ workers × 480 × 0.6 µs. Walk it up:

From 1 → 8 workers, iteration time drops 194 → 28 ms (a 6.9× speedup on 8× the workers — near-linear, simulation dominates). From 8 → 64 workers, you added 8× more workers but iteration time barely moves, 28 → 23 ms (a measly 1.2×), because the single learner now consumes 8× the pooled batch every round and the ray.put broadcast fans out to 64 readers. The bottleneck migrated from simulation compute to coordination. The Amdahl ceiling is set by that ~20 ms serial floor: even with infinite rollout workers, 192 / (≈20) ≈ 9.6× is the best total speedup you can ever get from this design.

// fixed total rollout work W (ms); per-worker batch feeds one serial learner.
// learner cost per worker = 480 transitions x 0.6us = 0.288 ms (see the §4 table).
var W = 192, putMs = 2, learnerPerWorker = 0.288;   // ms
function model(n){
  var sim   = W / n;                  // rollout divides across workers
  var coord = putMs + learnerPerWorker * n;  // broadcast + single learner grows
  var iter  = sim + coord;
  return { sim: sim, coord: coord, iter: iter, speedup: (W + putMs + learnerPerWorker) / iter };
}

5 · Look behind the scenes — the Ray dashboard

The table above is a prediction; the dashboard is how you confirm it on a real run. ray.init() prints a dashboard URL (default http://127.0.0.1:8265). Three views answer the question "where is my time going?" directly:

The discipline from lesson 00 returns: never optimize by intuition. Measure the split between sampling and learning, then attack whichever the timeline says is winning. The dashboard turns the Amdahl argument from arithmetic into something you can point at.

6 · Failure modes & checklist

Checkpoint exercise

Where this points next

You have built the RL graph by hand and found its ceiling: a synchronous barrier and a single learner that cap speedup no matter how many simulators you add. The obvious moves — make the learner its own actor, broadcast the policy as a tree, pipeline stale rollouts against fresh learning, put the learner on a GPU while rollouts stay on cheap CPUs — are exactly the design of RLlib. Lesson 07 shows RLlib's EnvRunners and learner as the productionized version of your RolloutWorker and ParameterServer, and makes the heterogeneous-resource argument (many cheap CPU samplers, few expensive GPU learners) that this hand-built version was begging for.

Interview prompts

• Map an RL training loop onto map-shuffle-reduce. (§1 — rollout is the map (independent episodes), pooling trajectories is the shuffle/reduce, the learner update is a reduce whose output feeds back to seed the next map; RL is MSR with a loop.)
• Why is the simulator an actor and not a task here? (§1, §3 — the environment is expensive to construct; an actor builds it once and is called repeatedly, whereas a task would rebuild it every call. The exercise's 20 s setup makes this decisive.)
• Why ray.put the policy instead of passing it directly to each worker? (§3, §6 — passing the value re-serializes it per worker (8× or 64×); putting it once and sharing the ref serializes it a single time and serves zero-copy on-node.)
• You scaled from 8 to 64 rollout workers and barely got faster. Why? (§4 — the single learner now consumes 8× the pooled batch and the policy broadcast fans to 64 readers; the serial coordination tax, not simulation, now dominates — Amdahl's ceiling.)
• Given that, what do you fix, and why not just add workers? (§4 — workers are already starved for serial budget; attack the floor: own-actor (or sharded) learner, tree/sharded policy broadcast, and controlled staleness to pipeline rollout against learning.)
• What is the on-policy vs off-policy distinction and how does it relate to staleness? (§1 — on-policy trains only on data from the latest policy (sync barrier); off-policy can reuse data from older policies (a buffer), which is what lets workers run a stale snapshot while the learner advances.)
• How would you confirm the bottleneck migration on a real run? (§5 — the dashboard Actors view (IDLE workers waiting on the learner), the timeline (rollout bars shrink, learner bar grows), and object-store memory (flat = one snapshot/iter, climbing = leaked refs).)