Skip to content
Draft
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
28 changes: 28 additions & 0 deletions configs/debug/replay/README.md
Original file line number Diff line number Diff line change
@@ -0,0 +1,28 @@
# Replay — Debug Configs

Minimal end-to-end configs for the replay derivation tasksets (`environments/replay_*_v1`,
thin subclasses of `verifiers.v1.tasksets.replay`), which turn
saved rollouts (`<run>/rollouts/step_*/train_rollouts.jsonl`) back into training tasks. Each config
mixes a fresh `reverse-text-v1` env with one replay env via ratios (ratios are all-or-none across
train envs), using `PrimeIntellect/Qwen3-0.6B-Reverse-Text-SFT` as the policy and the `null`
harness with the subprocess runtime for both envs.

| Config | Taskset | Buffer | Notes |
|---|---|---|---|
| `offline_recheck.toml` | `replay-recheck-v1` | offline, a prior run's `rollouts` dir | edit `buffer_dir` to a real prior run first |
| `online_recheck.toml` | `replay-recheck-v1` | `buffer_dir = "self"` — this run's own rollouts | forces whole-group dispatch; buffer fills as the run trains |

Every derivation is scored by its `inner` taskset, which must reproduce the source run's taskset
config. The replay env gets extra `max_input_tokens` headroom because its seeds carry whole source
conversations. `replay-continue-v1` (`anchor = "compaction" | "tool-call"`) has no debug config here:
reverse-text rollouts neither compact nor call tools — point one at an agentic run's buffer instead.

## Run the debug configs

```bash
# Offline recheck (edit buffer_dir in the TOML to a real prior run's rollouts dir first)
uv run rl @ configs/debug/replay/offline_recheck.toml --output-dir outputs/replay_recheck

# Online self-recheck (no prior run needed — the buffer fills as the run trains)
uv run rl @ configs/debug/replay/online_recheck.toml --output-dir outputs/replay_recheck_online
```
55 changes: 55 additions & 0 deletions configs/debug/replay/offline_recheck.toml
Original file line number Diff line number Diff line change
@@ -0,0 +1,55 @@
# Debug run mixing fresh reverse-text rollouts with replayed "recheck" tasks. The
# replay env re-serves finished conversations from a prior run's saved rollouts
# (<run>/rollouts/step_*/train_rollouts.jsonl) with an appended check-your-work
# turn, scored by the original (inner) taskset. Edit `buffer_dir` below to point
# at a real prior run before launching (CLI can't index into env lists; TOML
# overlays replace the env list wholesale):
# uv run rl @ configs/debug/replay/offline_recheck.toml

max_steps = 20
seq_len = 2048

[model]
name = "PrimeIntellect/Qwen3-0.6B-Reverse-Text-SFT"

[wandb]
project = "replay-debug"
name = "debug-offline-recheck"

[orchestrator]
batch_size = 128
group_size = 16

[orchestrator.renderer]
name = "qwen3"

[orchestrator.train.sampling]
max_completion_tokens = 128

# Fresh reverse-text rollouts (3/4 of the batch).
[[orchestrator.train.env]]
name = "reverse-text"
ratio = 3.0
taskset = { id = "reverse-text-v1" }
harness = { id = "null", runtime = { type = "subprocess" } }

# Replayed recheck tasks (1/4). Ratios are all-or-none: every train env sets one.
[[orchestrator.train.env]]
name = "replay-recheck"
ratio = 1.0
harness = { id = "null", runtime = { type = "subprocess" } }
[orchestrator.train.env.taskset]
id = "replay-recheck-v1"
# Placeholder: a prior run's saved-rollout dir (or the run dir containing it).
buffer_dir = "outputs/prior_run/rollouts"

# Replay rollouts are scored by the original taskset, so `inner` must reproduce the
# source run's taskset config (here the base env's defaults).
[orchestrator.train.env.taskset.inner]
id = "reverse-text-v1"

[trainer.optim]
lr = 3e-6

[inference]
gpu_memory_utilization = 0.5
57 changes: 57 additions & 0 deletions configs/debug/replay/online_recheck.toml
Original file line number Diff line number Diff line change
@@ -0,0 +1,57 @@
# Debug run where the model re-checks its own rollouts as it trains. The recheck env
# replays this run's freshly written rollouts (buffer_dir = "self" resolves to
# <output_dir>/rollouts before the env servers spawn): each task re-serves a finished
# conversation with an appended check-your-work turn, scored by the original (inner)
# taskset. The online buffer samples a fresh source per request — forcing whole-group
# dispatch so GRPO groups share one source.
# uv run rl @ configs/debug/replay/online_recheck.toml

max_steps = 20
seq_len = 2048

[model]
name = "PrimeIntellect/Qwen3-0.6B-Reverse-Text-SFT"

[wandb]
project = "replay-debug"
name = "debug-online-recheck"

[orchestrator]
batch_size = 128
group_size = 16

[orchestrator.renderer]
name = "qwen3"

[orchestrator.train.sampling]
max_completion_tokens = 128

# Fresh reverse-text rollouts (3/4 of the batch) — also what fills the buffer.
[[orchestrator.train.env]]
name = "reverse-text"
ratio = 3.0
taskset = { id = "reverse-text-v1" }
harness = { id = "null", runtime = { type = "subprocess" } }

# Recheck tasks over this run's own rollouts (1/4). Ratios are all-or-none.
[[orchestrator.train.env]]
name = "replay-recheck"
ratio = 1.0
harness = { id = "null", runtime = { type = "subprocess" } }
[orchestrator.train.env.taskset]
id = "replay-recheck-v1"
buffer_dir = "self" # implies online
# Recheck only the fresh env's rollouts. (Unset, source_envs replays every env except
# replay envs; listing a replay env by name opts into chained derivations.)
source_envs = ["reverse-text"]

# Rechecks are scored by the original taskset, so `inner` must reproduce the source
# run's taskset config (here the base env's defaults).
[orchestrator.train.env.taskset.inner]
id = "reverse-text-v1"

[trainer.optim]
lr = 3e-6

[inference]
gpu_memory_utilization = 0.5
7 changes: 6 additions & 1 deletion docs/configuration.md
Original file line number Diff line number Diff line change
Expand Up @@ -151,10 +151,11 @@ Training environments are an array of tables — set one per env, optionally wit
id = "math-env"
name = "gsm8k"
args = { dataset_name = "openai/gsm8k", dataset_subset = "main" }
ratio = 0.75 # 75% of the batch

[[orchestrator.train.env]]
id = "reverse-text"
ratio = 0.25 # 25% of batches; remaining 75% goes to math-env
ratio = 0.25 # ratios are relative weights; setting one means setting all

[[orchestrator.eval.env]]
id = "math-env"
Expand All @@ -166,6 +167,10 @@ args = { dataset_name = "openai/gsm8k", dataset_subset = "main" }

The same `id` can appear multiple times across train and eval (or with different `args`) — useful for evaluating on a held-out split of the env you're training on, or comparing two configurations of the same env side by side. When `id` is reused, set a distinct `name` on each entry; `name` defaults to `id` and must be unique across all envs in the same group.

`ratio` is all-or-none: once any train env sets one, every train env must.

One special environment ships with `prime-rl` itself: the `replay-v1` taskset replays a run's saved rollouts (its own or a prior run's) as derived training tasks — see [Replay](replay.md).

### Environment Variables

OS environment variables exported into launched component process(es). In `rl` configs, top-level `[env_vars]` applies to trainer, inference, and orchestrator:
Expand Down
1 change: 1 addition & 0 deletions docs/mint.json
Original file line number Diff line number Diff line change
Expand Up @@ -10,6 +10,7 @@
"inference",
"scaling",
"algorithms",
"replay",
"advanced",
"development"
]
Expand Down
95 changes: 95 additions & 0 deletions docs/replay.md
Original file line number Diff line number Diff line change
@@ -0,0 +1,95 @@
# Replay

The replay tasksets turn a run's saved rollouts back into training tasks. Every RL run writes its train rollouts to `<output_dir>/rollouts/step_N/train_rollouts.jsonl` (one WireTrace per line, stamped with `info.prime_rl` metadata); a replay taskset indexes those files as a buffer and serves tasks derived from them — resume a conversation from a compaction point or a tool call, or ask the model to re-check a finished attempt. A replay env is a normal `[[orchestrator.train.env]]` entry mixed into training with a `ratio`, over either a previous run's rollouts (offline) or the current run's own (online, `buffer_dir = "self"`).

The machinery is split like verifiers' `harbor`/`textarena` tasksets: the shared base (buffer, trace surgery, scoring delegation) lives in verifiers as `verifiers.v1.tasksets.replay`, and each derivation is a thin subclass package under `environments/`:

| Taskset id | Derived task |
|---|---|
| `replay-continue-v1` | Resume a source rollout mid-way — from a context-compaction point (`anchor = "compaction"`: the seed is the post-compaction prompt the original harness restarted from; one task per compaction) or right after a tool result (`anchor = "tool-call"`: one deterministically sampled resumable point per source; only points where every issued call has its result). |
| `replay-recheck-v1` | The finished conversation plus an appended check-your-work user turn (`instruction`); one task per rollout, seeded from its final state. |

Every derivation is scored by the original taskset: `inner` must reproduce the source run's taskset config, and the inner rewards land on the trace with their original names and weights. The replay base also delegates tools, setup, and finalize to the inner taskset, so tool-using tasks replay with their tools intact. Each drawn source spawns a fresh GRPO group.

## Table of Contents

- [Offline vs Online Buffers](#offline-vs-online-buffers)
- [Configuration](#configuration)
- [Metrics](#metrics)
- [Writing a New Derivation](#writing-a-new-derivation)
- [Constraints and Caveats](#constraints-and-caveats)

## Offline vs Online Buffers

**Offline** (the default): point `buffer_dir` at a finished run's rollouts. The buffer is indexed once at env-server startup — steps are scanned newest-first up to the buffer's capacity (the newest few thousand candidates are retained) — and each task index maps deterministically to one candidate. GRPO group members dispatched as independent rollouts therefore still bind the same source rollout.

**Online** (`buffer_dir = "self"`, which implies `online = true`): the buffer is the run's own growing rollout dir. The orchestrator resolves the `"self"` sentinel to `<output_dir>/rollouts` before the env servers spawn. Four consequences:

- **Barrier semantics.** The jsonl is written non-atomically, so an online buffer only reads steps whose sibling `train_rollouts.bin` exists — the orchestrator writes and closes the jsonl strictly before the atomic rename that creates the `.bin`, so the barrier file marks the jsonl complete. This means online replay requires the filesystem rollout transport (a run with a non-filesystem transport never writes the `.bin`). Offline buffers skip the barrier: finished runs' files are complete by definition.
- **Group dispatch.** Every request samples a fresh source rollout, so the whole GRPO group must share one draw — the taskset forces whole-group dispatch (`REQUIRES_GROUP_ROLLOUTS`), and all group members arrive at the env server as a single request.
- **Startup.** Until the run has produced its first replayable step, replay requests wait briefly, then fail (with a logged warning) so their dispatch permits free up for the source envs — the orchestrator retries replay groups naturally, and they start succeeding once the first step's barrier appears. The buffer rescans for new steps during training, and its capacity doubles as recency eviction: the newest candidates are retained, keeping the buffer on recent policy behavior.
- **Choose your sources.** A `"self"` buffer sees every train env's saved rollouts — including the replay envs' own. By default (`source_envs` unset) replay-derived records are skipped: rechecking your own rechecks is a feedback loop unless chosen deliberately. Listing env names replays exactly those — and naming a replay env is the deliberate opt-in for chained derivations (a recheck env sourcing another recheck env is depth-2 self-correction; the env topology *is* the depth control, and scoring/tools/provisioning always resolve to the innermost original task).

## Configuration

The derivation packages are installed via the `envs` extra. A replay env's taskset config is flat — the base fields plus the package's own:

```toml
[orchestrator]
batch_size = 128
group_size = 16

# Fresh rollouts (3/4 of the batch) — also what fills the buffer.
[[orchestrator.train.env]]
name = "reverse-text"
ratio = 3.0
taskset = { id = "reverse-text-v1" }
harness = { id = "null", runtime = { type = "subprocess" } }

# Recheck tasks over this run's own rollouts (1/4). Ratios are all-or-none.
[[orchestrator.train.env]]
name = "replay-recheck"
ratio = 1.0
harness = { id = "null", runtime = { type = "subprocess" } }

[orchestrator.train.env.taskset]
id = "replay-recheck-v1"
buffer_dir = "self"
source_envs = ["reverse-text"]

[orchestrator.train.env.taskset.inner]
id = "reverse-text-v1"
```

Base fields (every replay taskset, from `verifiers.v1.tasksets.replay`):

| Field | Default | What it does |
|---|---|---|
| `buffer_dir` | *(required)* | The saved-rollout dir to replay: a run's `rollouts` dir (or the run dir containing it). Under prime-rl the literal `"self"` resolves to this run's own rollout dir (an online buffer over the run's freshly written rollouts). |
| `inner` | *(required)* | The original taskset's config — it scores the derived rollouts and provides their tools, so it must reproduce the source run's taskset config. |
| `source_envs` | `None` | Which envs' rollouts to replay, by their stamped name (`info.prime_rl.env_name`). Unset: every env except replay envs. An explicit list replays exactly those — naming a replay env opts into chained derivations (recheck a recheck). With a list set, records without the stamp never match. |
| `allow_container` | `false` | Allow sources whose task ran in a container image. The container state the transcript references is gone — a fresh container is provisioned from the same image, so the model resumes in a reset world. |
| `online` | *(inferred)* | Set automatically for `buffer_dir = "self"`. Set it yourself only to treat another still-running run's dir as a growing buffer. |

Per-package fields: `replay-continue-v1` adds `anchor = "compaction" | "tool-call"`; `replay-recheck-v1` adds `instruction` (the appended check-your-work turn, with a built-in default).

Runnable debug configs live at `configs/debug/replay/` (`online_recheck.toml`, `offline_recheck.toml`).

## Metrics

Every replay rollout logs `replay/source_reward` (the reward the source rollout received) and `replay/source_step` (the trainer step that wrote it) alongside the inner taskset's own rewards and metrics.

## Writing a New Derivation

A derivation is a thin package over the base, exactly like `swebench-pro-v1` over verifiers' `harbor` taskset: subclass `ReplayTaskset` (binding your narrowed config in the generic) and implement `build_prompt` (the seeded conversation for one anchor); derivations that resume mid-rollout also override `record_anchors` (the resume points one saved rollout offers; each becomes one task — the default is one task per rollout, anchored at its final state). The base owns the buffer, online semantics, lazy binding, lineage, and inner-taskset delegation; `verifiers.v1.tasksets.replay.surgery` provides the graph enumerators (`compaction_forks`, `tool_call_anchors`, `recheck_seed`, ...). See `environments/replay_recheck_v1` for the ~30-line reference. Register the package in the root `pyproject.toml` (`envs` extra + `[tool.uv.sources]`) and `uv sync`.

## Constraints and Caveats

- **One derivation per env entry.** Env-level settings (harness runtime, token budgets, `ratio`) are per-env — mix derivations with multiple env entries.
- **Harness: match the source.** Replay resumes conversations — run the replay env under the same harness the source env used, so the resumed rollout has the same scaffold and harness-side tools as the original. The harness must support message-prompt seeding (`SUPPORTS_MESSAGE_PROMPT`; the built-in `default` and `null` harnesses do — string-prompt CLI harnesses can't be seeded with a prior conversation).
- **Replay envs need explicit ratios.** Without ratios, envs are weighted by task count — an online replay env advertises a single virtual task and would effectively never be drawn. Setting `ratio` on the replay env forces ratios on every train env (the all-or-none rule), which is the intended, explicit state.
- **Token budgets default from `seq_len`.** Every v1 train env's `max_total_tokens` defaults to the run's `seq_len` at config-resolve time, so replayed seeds stop instead of producing samples that can't fit training. Override `max_total_tokens` (or set it to `"None"`) per env entry to opt out.
- **Long continue tasks age off-policy.** A continue task resumes deep into a long trajectory and keeps going, so its rollout can span many trainer steps and be discarded by `orchestrator.max_off_policy_steps` (default 8). If a continue env shows sustained `errored_rollouts`, raise `max_off_policy_steps` or shorten the tasks.
- **`source_envs` needs the stamp.** The filter matches `info.prime_rl.env_name`, which prime-rl's rollout writer stamps into every saved record. Records without the stamp (rollout files from other producers) never match an explicit list — leave `source_envs` unset to replay them.
- **`allow_container` is off for a reason.** The trace records the conversation, not the container: for a containerized source task, replay provisions a fresh container from the same image, so any filesystem state the transcript references is gone and the model resumes in a reset world — especially acute for `anchor = "tool-call"`, which resumes mid-trajectory. Enable it only when the task's image is self-contained enough for that to be fair — and give the replay env's harness a container runtime (docker/prime): with a subprocess runtime every imaged request fails at episode construction.
- **Inner-taskset limits.** Replay delegates scoring to `inner` but cannot delegate group scoring (`@group_reward` tasksets), custom `State` classes, or user simulators — those are rejected loudly at env construction. Shared-placement inner toolsets are also unsupported, but fail per-rollout (the serving layer inspects toolsets on a stub task and never starts them).
13 changes: 13 additions & 0 deletions environments/replay_continue_v1/pyproject.toml
Original file line number Diff line number Diff line change
@@ -0,0 +1,13 @@
[project]
name = "replay-continue-v1"
version = "0.1.0"
description = "replay-continue-v1 — a derivation over verifiers' replay taskset base."
requires-python = ">=3.10"
dependencies = ["verifiers"]

[build-system]
requires = ["hatchling"]
build-backend = "hatchling.build"

[tool.hatch.build.targets.wheel]
packages = ["replay_continue_v1"]
Original file line number Diff line number Diff line change
@@ -0,0 +1,4 @@
from replay_continue_v1.taskset import ReplayContinueTaskset

# The loader requires exactly one Taskset subclass exported via __all__.
__all__ = ["ReplayContinueTaskset"]
Loading