Skip to content
Open
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
Original file line number Diff line number Diff line change
@@ -0,0 +1,220 @@
---
title: "Training on External Agent Harnesses"
description: "Capture exact token ids and log probabilities when the agent harness drives its own model calls."
position: 6
---

Some agent harnesses run their own model-calling loop. Gym starts the harness and points it at a model endpoint, but Gym does not mediate the model calls. The harness decides when to call the model, what to send, and how to handle each reply. It returns a finished transcript. The Claude Code CLI is the reference case and drives a multi-turn loop over Anthropic Messages.

Gym does not see the individual calls as they happen. The returned transcript carries no token ids because these harness wire formats have no field for them.

RL trains on token ids. Re-tokenizing the returned text can produce a sequence that differs from the sequence sampled by the policy. The size of that difference is unknown. Token capture records the exact ids inside the model server, where they still exist. It keys the ids to the rollout that produced them and rebuilds the rollout's calls into one contiguous response.

## When you need it

The harness location does not determine whether token capture is required. What matters is who makes the model calls and whether token ids survive the round trip.

| Your agent | What you need |
|---|---|
| Calls the model server through Gym and returns Responses items carrying token ids | Nothing. Train as usual. |
| Drives its own calls and returns text, or a dialect with no field for token ids | Token capture, described below. |
| Drives its own calls but returns token ids in a shape Gym does not read | Token capture, and open an issue so the shape can be read directly. |

## Turning it on

Two settings turn capture on. If either setting is missing, the run completes without an error but provides no captured tokens for training.

**1. Enable capture and give it node-local storage.** The writer and reader are on the same node. A shared filesystem adds unnecessary latency and can let two shards write the same file.

```yaml
env:
nemo_gym:
token_id_capture:
enabled: true
dir: /tmp/nemo_gym_token_id_captures
```

All run-wide capture settings live in this block, which is validated at startup. A typo in a key raises an error instead of producing a run that appears configured but provides no captured tokens for training. The other settings can remain in place when `enabled: false`, so one config can retain the directory and toggle capture per run.

**2. Opt the agent in.** The per-agent flag scopes capture to harnesses that need it. Native agents in the same run remain unchanged.

```yaml
responses_api_agents:
claude_code_agent:
token_id_capture: true
```

For a training run that captures every configured agent, set `token_id_capture.all_agents: true` in the run-wide block instead of repeating the agent flag. This overrides agent-level opt-ins but does not enable capture by itself. Keep both `enabled` and `all_agents` false in evaluation configs.

An opted-in agent adds `/training-token-capture` to its rollout-correlated model-server URL. The model server uses that segment to distinguish training capture from ordinary requests on the same endpoint, then strips it before API routing. It does not intercept or change the request body.

Capture reads token ids from the served response, so the inference server must return them. For vLLM, that requires a tokenizer:

```yaml
policy:
generation:
vllm_cfg:
skip_tokenizer_init: false
```

Sampling parameters must also be pinned on the server. Harnesses built for interactive serving generally do not send them. An unset parameter therefore uses the engine default instead of the value used to optimize the policy. Set `sampling_overrides` on the model server to the trainer's generation config.

<Note>
One setting does not control capture, but it determines whether the rollout contains multiple calls worth capturing.

A **tool-call parser** converts the model's tool-call syntax into structured calls that the harness can dispatch. Without a parser, the harness sees ordinary text and does not call a tool. Every rollout then contains one model call. Capture still works, but there are no calls to chain. Configure the parser on the inference server, such as `tool_parser: hermes` under `http_server_serving_chat_kwargs`. The correct value depends on the model.

This failure is silent. Check `n_calls` on the first run before relying on the reward.

</Note>

## What you get back

Each rollout's model calls are stitched into a single Responses payload with contiguous `output` items. Each item's `prompt_token_ids` contains the running sequence. Its `generation_token_ids` contains the tokens sampled by the policy at that step. Gym replaces the rollout's `response.output` with these items, so a trainer reads `response.output` the same way for native agents and external harnesses.

The loss mask follows from this structure instead of being sent separately. Prompt positions provide context. Generation positions are trainable.

## What to watch on a first run

Read these metrics before the reward curve. A rollout can appear healthy because its reward changes even when most of the rollout never reaches training.

Gym attaches a metrics dictionary to each rollout under `_ng_token_capture`. Aggregate these metrics across a step through the training framework's existing reporting path.

| Key | Expect | If it is wrong |
|---|---|---|
| `n_calls` | above 1 | The harness never called a tool. Usually a missing tool parser. |
| `chains` | 1 | The rollout split. Part of it is not reaching the optimizer. |
| `delivered_fraction` | 1.0 | Sampled tokens were captured but not delivered. |
| `quarantined_calls` | 0 | Two calls could not be told apart, so neither was used. |
| `empty_generation_calls` | 0 | The output budget or a content filter is truncating generations. |
| `mask_sample` | absent | The rollout lost a call and must not be trained on. |

Pay particular attention to `n_calls`. A value of exactly 1 means the agentic path was never exercised, even though every other key can look correct.

A rollout without a metrics dictionary was never rebuilt. Its model calls were not correlated, so the rollout has no captured tokens.

## Integrating a training framework

Gym defines the record shape and builds each record. The training framework controls where the record goes.

### The interfaces

Gym describes the sink and source as structural protocols. Framework adapters do not inherit from these definitions or import them at runtime. They implement the same method signatures, and Gym consumes the resulting objects by that method shape. The definitions below are the reference contract.

```python
class TokenSink(Protocol):
async def put(self, entry: TokenEntry) -> None: ...
async def mark_incomplete(self, rollout_id: str, model_call_id: str = "") -> None: ...
async def close(self) -> None: ...

class TokenSource(Protocol):
async def freeze(self, rollout_id: str) -> TokenCaptureSnapshot: ...
async def drop(self, rollout_id: str, *, snapshot_id: str, version: int) -> bool: ...
async def close(self) -> None: ...
```
Comment thread
ananthsub marked this conversation as resolved.

`put` must make the record durable before it returns. A reader that runs after the rollout must see every acknowledged record. This guarantee allows the consumer to freeze one complete view after the harness finishes.

`mark_incomplete` is the durable signal that a rollout lost a call. The model call still succeeds. A sink that drops this signal makes an incomplete rollout look complete.

`freeze` returns one atomic snapshot containing entries, incomplete state, `snapshot_id`, and version. It is idempotent for an unchanged rollout. A write that races the snapshot must advance the observable version.

`drop` conditionally retires only the supplied snapshot identity and version. It returns `false` if state changed after `freeze`, preserving a late write instead of deleting evidence the consumer never saw. A transport without a delete operation returns `true` without deleting data, and its storage owner remains responsible for retention.

`close` releases client resources. Gym closes clients it constructs, but it does not close a caller-installed source.

### Connecting a framework-owned transport

The training framework owns the transport and its client configuration. Gym owns the model-server processes, so each server worker constructs a framework-provided `TokenSink` proxy from the configured class path. The configured class is a worker factory descriptor, not a shared Python object or a transport implementation owned by Gym.

```yaml
env:
nemo_gym:
token_id_capture:
enabled: true
sink: my_pkg.sinks:MyDataPlaneSink # module.path:ClassName
sink_kwargs:
endpoint: ${oc.env:MY_DATAPLANE_URL}
shard: ${oc.select:cluster_shard,0}
rebuild_response: false
```

Gym passes `sink_kwargs` to the constructor, so a sink can receive the required endpoint, client, or credentials instead of reading ambient state. Use `${oc.env:VAR}` for secrets instead of writing them into the config. Unsupported constructor arguments cause a startup error. A sink that does not implement `mark_incomplete` also causes a startup error because it could otherwise make a rollout with a missing call look complete.

`sink` replaces the file store, so a `dir` configured alongside it is not used. This condition produces a warning instead of an error because no data is lost. No capture files appear on disk.

The framework separately constructs its `TokenSource` in the trainer or rollout-consumer process. That process may use another virtual environment or actor because the source and sink are independent clients of the same transport.

```python
source = TransferQueueTokenSource(queue_handle)
built = await finalize_rollout_token_capture(result, source)
await durable_handoff(built)
await retire_rollout_token_capture(result["_ng_rollout_id"], source, built)
```

`rebuild_response: false` tells Gym to stop after the write. Correlation and capture are unaffected. The training framework freezes, reconstructs, durably hands off, and conditionally retires the snapshot through its `TokenSource`. Set `rebuild_response: true` only when Gym's rollout collector owns that sequence; the collector process then requires an installed source or the default file store.

When Gym's rollout collector owns rebuilding over a framework transport, install the framework-created source in that collector process with `install_token_source` before collection starts. The default file-backed path needs no installation because Gym constructs a `TokenCaptureStore` from `token_id_capture.dir`.

Leaving `enabled` off disables capture. An external harness then produces rollouts without token ids and no token data for training. Use this configuration only for evaluation.

<Warning>
Configure the sink instead of installing it from a launcher script. Programmatic installation must run inside the serving process.

`install_token_sink` sets a process global. A model server with `num_workers > 1` launches uvicorn with an app string and `workers=N`. Uvicorn spawns workers that re-import the app module instead of inheriting the launcher's memory. A sink installed by the parent process therefore does not exist in any worker. Capture then falls back to the file store. If no `dir` is set, the worker has no local destination.

Gym constructs the configured sink inside each worker at app startup, which avoids this process-boundary problem.
</Warning>

### Reading records back

`TokenCaptureStore` implements both protocols and is the default. A reader beside the store uses the store as its `TokenSource`. This arrangement applies to `gym eval run` and to a trainer colocated with the model server. The store directory should therefore be node-local.

A framework that stages records through its own transport reads them through its own `TokenSource`. That source can run wherever the transport runs. Reading through a custom source is not restricted to the model server node.

Every source must return an accurate `incomplete` value in its frozen snapshot. This value tells the consumer that a rollout lost a model call. The records that arrived can form a contiguous-looking chain while still missing a turn. A source that always reports `false` can cause training to use that incomplete rollout.

Consume records through `TokenSource.freeze`. Reading entries alone is insufficient because safe masking and retirement also require the snapshot's incomplete state, identity, and version.

### Driving rollouts yourself

`gym eval run` finalizes each record and retires successful evidence only after the output row is durable. A framework that calls `run_examples` directly does not use that path, so it must perform the same sequence for each finished record:

```python
from nemo_gym.token_id_capture.delivery import (
finalize_rollout_token_capture,
retire_rollout_token_capture,
)

built = await finalize_rollout_token_capture(result, source)
await downstream.put(result) # Must be durable when this returns.
await retire_rollout_token_capture(rollout_id, source, built)
```

`finalize_rollout_token_capture` freezes the source snapshot, rebuilds `response.output`, and attaches build metrics. It mutates the record in place and does not retire evidence. `retire_rollout_token_capture` conditionally drops only the snapshot that was rebuilt, and only after the caller establishes its durability boundary. Failed and masked builds remain available for diagnosis.

A rollout that cannot be rebuilt is flagged with `mask_sample: true` at the top level of its record. Exclude these rollouts from the loss. The trajectory is missing a turn, or two candidate generations could not be distinguished. Training on such a trajectory is off-policy.

### Rollout ids

Capture keys each record by rollout id. Gym derives this id from the run request's task and rollout indices. This scheme assumes that each dispatch receives a distinct pair. If a training loop restarts numbering, such as reusing the same indices in each training step, the derived id repeats and two dispatches share one capture key.

Set `_ng_rollout_id` on the run body to provide a distinct key:

```python
row["_ng_rollout_id"] = f"step{step}.{task_index}-{rollout_index}"
```

The id becomes a URL path segment. It can contain letters, digits, dots, dashes, and underscores, and it must start with a letter or digit. Gym rejects an invalid id instead of rewriting it.

## Extensions

### Sampling pin

`sampling_overrides` on the model server applies the configured sampling parameters to every request and overrides values sent by the harness. Generation KL error indicates whether the overrides are working.

## Limitations

**One trajectory per rollout.** A harness that forks sub-agents or retries a call produces a tree of model calls. Gym delivers one chain and reports omitted sampled tokens through `delivered_fraction` instead of dropping them silently. Training on the full tree requires a trainer contract that accepts a tree.

**Harness calls outside the rollout.** A harness may generate a conversation title or a context-compaction summary. These calls are policy output and can be selected for training. A compaction summary can be long enough to outweigh the rollout it summarizes. This pattern appears as `chains` above 1 and `delivered_fraction` below 1.0 in `_ng_token_capture`.
6 changes: 6 additions & 0 deletions fern/versions/latest/pages/training-tutorials/index.mdx
Original file line number Diff line number Diff line change
Expand Up @@ -26,6 +26,12 @@ Example GRPO training on instruction following and reasoning environments.
<Badge minimal outlined>unsloth</Badge> <Badge minimal outlined>single-gpu</Badge> <Badge minimal outlined>30 min</Badge>
</Card>

<Card title="External Agent Harnesses" href="/tutorials/training-tutorials/external-agent-harnesses">
Train on rollouts from a harness that drives its own model calls and returns no token ids, such as the Claude Code CLI.

<Badge minimal outlined>token capture</Badge> <Badge minimal outlined>agent harnesses</Badge>
</Card>

<Card title="VeRL" href="/tutorials/training-tutorials/verl">
Example DAPO training on math and agentic environments using VeRL, with single and multi-environment support.

Expand Down
Loading