Guides

GRPO-style group rollouts with fork

Reach a task state once, fork it into N copies, roll a policy out in each copy, and compute group-relative advantages from valid rewards only.

Group-relative methods such as GRPO compare several rollouts that start from the same state. With Drums you reach that state once (the task's setup plus any prefix actions) and fork(n) it: each child is an independent desktop restored from the same saved state, with its own episode and gate report. No child depends on the others, so they can run concurrently.

The loop

Create one desktop with the task and wait until its episode is valid.
Optionally run prefix actions to reach the state you want to branch from.
env.fork(n, keep_state=False): n children from that state.
Roll the policy out in every child, concurrently.
verify(strict=False) each child. Keep only rewards with valid == True.
Compute advantages within the group of valid rewards. Close the children (the with block does it).

group_rollouts.py

import random
import statistics
import drums

TASK = "e0df059f"      # OSWorld: rename a folder on the Desktop
N = 2                  # copies per group; at most the host's fork limit (8 here)
STEPS = 3


def policy(obs, rng):
    """Stand-in for your model: one pyautogui action from a screenshot.

    Replace with a model call."""
    x, y = rng.randrange(0, obs.width), rng.randrange(0, obs.height)
    return f"pyautogui.moveTo({x}, {y})"


def rollout(child, seed):
    rng = random.Random(seed)
    child.wait_valid()  # each child is a new episode: let its gates resolve
    for _ in range(STEPS):
        child.step(policy(child.observe(), rng))
    return child.verify(strict=False)


with drums.connect() as client:
    with client.create(task=TASK, ttl=900) as parent:
        parent.wait_valid()
        parent.step("pyautogui.moveTo(960, 540)")        # prefix: shared by every copy
        # the children close (and their state is freed) when the block ends
        with parent.fork(N, keep_state=False) as group:
            pairs = list(zip(group, range(N)))
            rewards = client.map(lambda pair: rollout(*pair), pairs, workers=N)
    for r in rewards:
        if isinstance(r, Exception):
            print("rollout failed:", r)
    valid = [r for r in rewards if isinstance(r, drums.Reward) and r.valid]
    print(f"{len(valid)} valid of {len(rewards)} rollouts")
    if valid:
        scores = [float(r.score) for r in valid]
        mean, std = statistics.fmean(scores), statistics.pstdev(scores)
        advantages = [(s - mean) / (std + 1e-6) for s in scores]
        print("scores", scores, "advantages", [round(a, 3) for a in advantages])

client.map runs a function over items in threads and returns each result or the exception it raised, so one failed rollout never hides the others.

Choices that matter

Drop invalid rollouts, do not score them 0. An invalid episode (a failed gate, a quarantine, an out-of-memory kill, an app crash) says nothing about the policy. The prime-rl adapter drops them the same way.
keep_state=False in loops. Each fork makes a saved state of about 0.6 GB of host memory. With keep_state=False it is deleted when the last child closes; otherwise delete it with client.delete_state(group.state) or let its idle lease (1800 s) remove it. The host holds a limited number of saved states.
Wait for each child's gates before actions that may end a process (closing a window, killing an app); otherwise the child's episode can be flagged invalid.
Capacity. The parent and its children all count against your workspace's live-computer limit and the host's desktops. A fork that would exceed them is refused before anything starts (403 quota_exceeded, or 429 overloaded when the host is full).
Failed children. With the default strict=True, fork() raises ForkError if any child did not come up; ForkError.ok holds the usable children and .state the saved state. With strict=False the failed ones are in group.failed.

Several groups from one desktop

parent.reset() starts a new episode from the image (and reruns the task's setup), so one parent can seed many groups:

Python

import drums

with drums.connect() as client:
    with client.create(task="e0df059f", ttl=900) as parent:
        for group_no in range(2):
            parent.wait_valid()
            with parent.fork(2, keep_state=False) as group:
                rewards = client.map(lambda k: k.verify(strict=False), group)
                print("group", group_no, [(r.valid, r.score) for r in rewards])
            parent.reset()

Over HTTP

A fork is one request; the reply lists the children:

Shell

curl -s -X POST "$DRUMS_URL/v1/slots/$COMPUTER_ID/branch" \
  -H "Authorization: Bearer $DRUMS_TOKEN" -H "Content-Type: application/json" \
  -d '{"n": 2, "delete_state_with_children": true, "ttl_s": 600}' \
  | jq '{state, children: [.children[] | {id, state}]}'

See Fork in the HTTP reference for every field.