Architecture overview

dockyard_rl trains a policy model with GRPO (Group Relative Policy Optimization) against execution-grounded rewards. The defining structural choice is that the three responsibilities — training, generation, and scoring — run as separate, asynchronous fleets rather than phases of one process.

The three fleets

Each container declares DOCKYARD_FLEET_ROLE before init_ray(). The cluster bootstrap (cluster/bootstrap.py) reads that role to pick the placement strategy, resource spec, and NCCL (NVIDIA Collective Communications Library) topology.

        flowchart LR
    subgraph TRAINER["Trainer fleet · GPU · SPREAD"]
        direction TB
        POLICY["Policy workers<br/>DTensor/FSDP2 · JAX/NNX"]
        GRPO["GRPO optimizer<br/>advantages · clipped-PG loss"]
        POLICY --> GRPO
    end

    subgraph INFERENCE["Inference fleet · GPU · PACK"]
        ENGINE["vLLM async engine<br/>or SGLang · serves rollouts"]
    end

    subgraph SANDBOX["Sandbox fleet · CPU · SPREAD"]
        EXEC["ubuntu-swe task executor<br/>clone repo · apply patch · run tests"]
    end

    BUFFER[("Replay buffer<br/>async · trajectory-age bounded")]

    GRPO   -- "1 · weight sync (NCCL collective)" --> ENGINE
    ENGINE -- "2 · rollout: code patch / actions"  --> EXEC
    EXEC   -- "3 · verdict: tests pass/fail → reward" --> BUFFER
    ENGINE -. "trajectories + logprobs" .-> BUFFER
    BUFFER -- "4 · sampled experience" --> POLICY

    style TRAINER   stroke:#d29922,stroke-width:2px
    style INFERENCE stroke:#1f6feb,stroke-width:2px
    style SANDBOX   stroke:#2da44e,stroke-width:2px
    

• trainer — DTensor/FSDP2 (Distributed Tensors + Fully Sharded Data Parallel v2) policy workers run the forward/backward and the GRPO optimizer step. An optional JAX (Flax NNX) backend implements the same policy interface and is selected per config. Placement is SPREAD so shards distribute across nodes.

• inference — a vLLM async engine (SGLang available) generates rollouts. Placement is PACK to keep an engine’s tensor-parallel ranks on one node. Weights arrive from the trainer over a NCCL collective.

• sandbox — CPU-only ubuntu-swe containers run a batch task executor (POST /task/submit, port 9090). It clones the repo, applies the agent patch and the gold test_patch, runs the tests, and returns a verdict. The agent never reaches the executor directly — the environment and reward layers drive it through the sandbox/ client.

Non-colocation is what makes the loop asynchronous: the trainer never blocks on generation, and generation never blocks on scoring.

The GRPO loop

The loop lives in algorithms/grpo.py. One optimizer step:

1. Sample prompts. num_prompts_per_step prompts are drawn; each becomes a group of num_generations_per_prompt rollouts.

2. Generate. The inference fleet produces the rollouts. For SWE tasks each rollout is single-shot (max_rollout_turns=1): the agent emits one patch.

3. Score. Each patch is submitted to a sandbox task executor. The verdict (binary resolved/unresolved, or test pass-rate) is the raw reward.

4. Shape rewards. Optional reward shaping (overlong penalties), reward scaling, and the native invalid-action / malformed-thinking penalty are applied. The integrity check zeroes the reward if the patch edited a held-out test file.

5. Estimate advantages. GRPO computes group-relative advantages with a leave-one-out baseline and reward normalization — no value network.

6. Recompute log-probs. The policy computes current log-probs (and, unless skipped, reference-policy log-probs) for the rollouts.

7. Loss and step. The clipped policy-gradient loss (ClippedPGLossFn) with ratio clipping, optional dual-clip, a reference-KL (Kullback–Leibler) penalty, and optional importance-sampling correction produces gradients; the optimizer steps.

8. Refit. Updated weights are synchronized to the inference fleet (below), and the loop repeats.

Validation runs every grpo.val_period steps (and optionally at start/end), generating on a held-out set and reporting the metric that drives checkpoint retention.

Synchronous and asynchronous loops

grpo_train() is synchronous: generate the whole batch, score, train. The math is on-policy and simplest to reason about.

async_grpo_train() decouples generation from optimization with a replay buffer (algorithms/async_utils/). Rollouts are collected continuously by a trajectory collector; an optimizer step consumes whatever is ready. Two mechanisms keep this sound:

• Trajectory age. A trajectory stays eligible for async_grpo.max_trajectory_age_steps optimizer steps. Older experience is dropped, bounding how off-policy the data can be.

• Importance-sampling correction. Because consumed rollouts were generated by a slightly older policy, async GRPO requires loss_fn.use_importance_sampling_correction=true; the loss reweights by the ratio of the current to the generating policy.

Async also allows in_flight_weight_updates: the trainer pushes new weights to the vLLM engine while it is mid-generation, optionally recomputing the KV cache afterward.

The weight-staleness handshake

The trainer’s weights advance every step; the inference fleet’s weights only change on refit. A WeightSynchronizer (weight_sync/) mediates this with a small state machine the loop drives explicitly:

• After an optimizer step, the trainer calls mark_stale() — the inference weights are now behind.

• Before generating, if is_stale, the loop calls sync_weights() to push the trainer’s weights, then prepare_for_generation() on the engine.

The synchronizer is built by create_weight_synchronizer(), which selects the transport from the deployment topology: a NCCL collective for the GPU-to-GPU non-colocated path, or IPC/ZMQ and HTTP transports for other layouts. The trainer exposes its weights for the collective through a refit seam that maps its internal parameter layout (DTensor shards, or JAX/NNX arrays) to the HF names and shapes the inference engine expects.

Weights vs. experience

Two distinct data flows cross the fleet boundary, with different transports:

• Weights (trainer → inference) move over the weight-sync transport (NCCL collective in the canonical topology).

• Experience (rollouts, log-prob tensors) can move over an optional data plane (data_plane/) instead of Ray’s object store, so a large (B, S) tensor never makes a driver round-trip. With data_plane=None the in-memory colocated path is used.

Where the code lives

Area	Package
GRPO loop, loss, advantage estimators	algorithms/
Policy workers (DTensor / JAX) and generation (vLLM / SGLang)	models/
Ray bootstrap, fleets, placement, NCCL sync	cluster/
Virtual cluster, worker groups, batched data containers	distributed/
Datasets, processors, registries	data/
Single-shot SWE patch scoring	environments/
Reward computation and integrity verification	rewards/
Task-executor client	sandbox/
ubuntu-swe image and the task executor	ubuntu-base/
Weight-sync transports	weight_sync/
Experience data plane	data_plane/
Launchers and configs	examples/

Each of these is covered in more depth in its design document.