diff --git a/fern/versions/latest/pages/training-tutorials/external-agent-harnesses.mdx b/fern/versions/latest/pages/training-tutorials/external-agent-harnesses.mdx index e2d0964dc5..03c47d99a1 100644 --- a/fern/versions/latest/pages/training-tutorials/external-agent-harnesses.mdx +++ b/fern/versions/latest/pages/training-tutorials/external-agent-harnesses.mdx @@ -1,30 +1,67 @@ --- title: "Training on External Agent Harnesses" -description: "Capture exact token ids and log probabilities when the agent harness drives its own model calls." +description: "Capture exact token ids and log probabilities when an external 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. +# Training on External Agent Harnesses -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. +An external agent harness runs its own model-calling loop and returns a finished transcript to Gym. Claude Code is one example. The transcript contains the conversation text, but it does not contain the exact token ids and log probabilities produced by each model call. -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. +Training needs those original token ids. Re-tokenizing the final text can produce a different sequence from the one sampled by the policy. Gym's training-token capture records the ids while they are still available in the model server, associates the calls with their rollout, and rebuilds them into the token-bearing response expected by the trainer. -## When you need it +## Decide whether you need token capture -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 | +| Agent behavior | What to do | |---|---| -| 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. | +| The agent returns Responses API output items that already contain generated token ids. | Train on the returned response. Gym preserves these native token ids. | +| The agent drives its own model calls and returns only text or a wire format without token ids. | Enable training-token capture. | +| The agent returns token ids in a format Gym does not recognize. | Enable capture for now and open an issue describing the returned format. | + +Ordinary evaluation does not require training-token capture. Enable it when the collected rollouts will be used to optimize a policy. + +## How capture works + +```mermaid +flowchart LR + H[External agent harness] -->|model calls| M[Gym model server] + M -->|exact token ids and log probabilities| S[(Capture storage)] + H -->|finished rollout| C[Rollout collector] + S --> C + C -->|rebuild verified call chain| R[Token-bearing response] + R --> T[Trainer] +``` + +Every model call carries the rollout id. For a multi-call rollout, Gym verifies which earlier call the new request continues. After the harness finishes, Gym reads the captured calls and replaces the rollout's text-only `response.output` with output items that carry the original token ids and log probabilities. + +Gym does not guess when a call is missing or its parent cannot be proven. It sets `mask_sample: true` so the trainer can exclude that rollout from the loss. + +## Configure local capture + +The default setup writes capture records to a node-local directory and lets Gym rebuild the response. + +### 1. Configure the model server for training -## Turning it on +Use `responses_api_models/vllm_model/configs/vllm_model_for_training.yaml` as the model-server config. It enables token information in model responses. -Two settings turn capture on. If either setting is missing, the run completes without an error but provides no captured tokens for training. +External harnesses often omit sampling parameters or send serving-oriented defaults. Pin the parameters used by training with `sampling_overrides` so every request uses the policy's intended sampling configuration: -**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 +policy_model: + responses_api_models: + vllm_model: + return_token_id_information: true + sampling_overrides: + temperature: 1.0 + top_p: 1.0 + top_k: -1 +``` + +Replace the example values with the sampling settings used by the trainer. If the training framework starts vLLM, keep its tokenizer enabled. For NeMo RL, set `policy.generation.vllm_cfg.skip_tokenizer_init: false`. + +### 2. Enable the capture store + +Add the run-wide capture settings: ```yaml env: @@ -32,173 +69,110 @@ env: token_id_capture: enabled: true dir: /tmp/nemo_gym_token_id_captures + delta_records: true + max_mask_fraction: 0.5 ``` -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. +`delta_records: true` avoids storing the growing full prompt again for every resolved continuation. Root and unresolved calls remain self-contained so Gym can diagnose a broken chain. -**2. Opt the agent in.** The per-agent flag scopes capture to harnesses that need it. Native agents in the same run remain unchanged. +`max_mask_fraction` is an optional safety limit. After at least `mask_fraction_min_samples` finalized rollouts, collection stops if the masked fraction exceeds this value. Omit it to disable the limit. -```yaml -responses_api_agents: - claude_code_agent: - token_id_capture: true -``` +### 3. Opt in the external agent -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: +Enable capture on each external harness that returns no token ids: ```yaml -policy: - generation: - vllm_cfg: - skip_tokenizer_init: false +claude_code_agent: + responses_api_agents: + claude_code_agent: + token_id_capture: true ``` -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. +Both the run-wide `enabled` setting and the agent opt-in are required. To capture every configured agent, set `token_id_capture.all_agents: true` instead of setting the flag on each agent. -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. - +Native Gym agents normally leave `token_id_capture` disabled because their responses already contain the token ids needed for training. -## What you get back +### 4. Confirm that tool calls are enabled -Each captured call records one parent-resolution result. `ROOT` means that the request contains no previous model-authored output. `RESOLVED` means that exactly one committed call matches and its request context was verified. `UNRESOLVED` means that the request appears to continue prior model output, but Gym cannot prove which committed call produced it. +A tool-call parser converts model output into structured calls that the harness can execute. Without the correct parser, an agentic rollout may stop after one model call even though capture itself is working. -Each rollout's resolved 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. +Configure the parser expected by the model, such as `tool_parser: hermes` in the inference server's chat-serving settings. Check `n_calls` on the first collected rollouts before relying on reward results. -An unresolved call begins a separate incomplete fragment. Gym never crosses that boundary with token-prefix inference. The call and its descendants remain available for diagnostics or a consumer that explicitly supports partial trajectories, but single-response delivery sets `mask_sample: true`. Prefix inference is used only when a verified parent is absent from the frozen snapshot, such as a filtered call that generated no tokens. +## Consume the rebuilt rollout -The loss mask follows from this structure instead of being sent separately. Prompt positions provide context. Generation positions are trainable. +When `rebuild_response` remains at its default value of `true`, `gym eval run` performs the complete lifecycle: -## What to watch on a first run +1. The model server records every correlated call. +2. Gym freezes the records after the harness and verifier finish. +3. Gym reconstructs one verified model-call chain and updates `response.output`. +4. Gym writes the rollout result durably. +5. Gym retires successfully consumed capture records. -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. +The trainer reads the rebuilt `response.output` in the same way it reads output from a native agent. Prompt positions provide context, while generated positions retain the captured log probabilities used for policy optimization. -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. +Masked or failed builds are retained for diagnosis. Successfully delivered records are removed after the output row is durable. -| 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. | -| `unresolved_parent_calls` | 0 | A request appeared to continue prior model output, but no exact committed parent was proven. | -| `mask_sample` | absent | The rollout lost a call, split at an unresolved parent, or otherwise cannot be trained safely. | +## Check the first run -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. +Gym adds capture metrics under `_ng_token_capture` on each rebuilt rollout. -A rollout without a metrics dictionary was never rebuilt. Its model calls were not correlated, so the rollout has no captured tokens. +| Field | Healthy value | What an unexpected value usually means | +|---|---|---| +| `n_calls` | Greater than 1 for a tool-using rollout | The harness did not execute a tool, often because the tool parser is missing or incompatible. | +| `terminal_attribution.chain` | `delivered` | Gym could not connect the verifier-scored response to an intact captured chain. | +| `chains` | `1` for a simple rollout | The harness made independent side calls, retried a call, or forked another agent. | +| `delivered_fraction` | `1.0` for a simple rollout | Some sampled tokens belonged to calls outside the delivered chain. This can be expected when terminal attribution safely excludes side calls. | +| `quarantined_calls` | `0` | Gym found conflicting candidates and refused to choose between them. | +| `empty_generation_calls` | `0` | A model call produced no trainable generated tokens. | +| `unresolved_parent_calls` | `0` | A continuation could not be linked to exactly one verified earlier call. | +| `mask_sample` | Absent or `false` | The rollout is incomplete or ambiguous and must not contribute to the loss. | -## Integrating a training framework +A rollout with no `_ng_token_capture` field was not rebuilt. Verify that capture is enabled, the agent is opted in, and its model calls use the rollout-correlated model-server URL. -Gym defines the record shape and builds each record. The training framework controls where the record goes. +## Optionally supply exact prefix tokens -### The interfaces +Some harnesses reshape an assistant turn before sending the next request. A chat template or reasoning template can also render the previous turn differently. In these cases, the next prompt may not begin with the tokens that the policy actually sampled. -Gym describes the sink, source, and lineage resolver 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. +Prefix supply asks a compatible inference backend to begin the next prompt with the verified parent's exact tokens. Enable the Gym model-server side with this config overlay: -```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: ... - -class LineageStore(Protocol): - async def resolve(self, rollout_id: str, request_items: list[dict]) -> LineageResolution: ... - def is_process_shared(self) -> bool: ... - async def close(self) -> None: ... +```text +responses_api_models/vllm_model/configs/vllm_model_supply_prefix.yaml ``` -`TokenSink` remains the writer-side interface. `put` is the only publication boundary: it must make the complete token record and its compact continuation lookup metadata durable before returning. `LineageStore` is the worker-side read client over those committed entries; it does not publish a second record. A later request may resolve an entry only after the corresponding `put` has returned. - -`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. +The backend must support Gym's required-prefix request and return the prompt token ids actually used for generation. Gym verifies that evidence before accepting prefix supply. A missing or mismatched proof fails the model call instead of silently producing an off-policy capture. -`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. + +Stock vLLM does not implement this prefix-supply extension. Leave the overlay disabled unless the inference backend supports both the required-prefix request and generation-time prompt-token response. Capture can still work without prefix supply when each rendered prompt naturally extends the preceding sampled tokens. + -`close` releases client resources. Gym closes clients it constructs, but it does not close a caller-installed source. +Prefix supply is not supported with `use_completions_api: true` or `is_responses_native: true`. -`LineageResolution` has three semantic outcomes: `ROOT`, `RESOLVED`, and `UNRESOLVED`. Diagnostic reasons such as an ambiguous lookup or unavailable backend may accompany `UNRESOLVED`, but they do not change reconstruction behavior. +## Use framework-owned capture storage -### Connecting a framework-owned transport +The default file store is intended for a model server and rollout collector that share one node-local directory. A distributed training system can instead provide a shared transport. -The training framework owns the transport and its client configuration. Gym owns the model-server processes, so each server worker constructs framework-provided `TokenSink` and `LineageStore` proxies from configured class paths. The configured classes are worker factory descriptors, not shared Python objects or transport implementations owned by Gym. +Configure a sink and lineage resolver that connect to the same backend: ```yaml env: nemo_gym: token_id_capture: enabled: true - sink: my_pkg.sinks:MyDataPlaneSink # module.path:ClassName + sink: my_package.capture:CaptureSink sink_kwargs: - endpoint: ${oc.env:MY_DATAPLANE_URL} - shard: ${oc.select:cluster_shard,0} - lineage_store: my_pkg.sinks:MyDataPlaneLineageStore + endpoint: ${oc.env:CAPTURE_ENDPOINT} + lineage_store: my_package.capture:CaptureLineageStore lineage_store_kwargs: - endpoint: ${oc.env:MY_DATAPLANE_URL} - shard: ${oc.select:cluster_shard,0} + endpoint: ${oc.env:CAPTURE_ENDPOINT} + delta_records: true rebuild_response: false ``` -Gym passes each kwargs block to its constructor, so the clients can receive the required endpoint, shard, 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. - -The sink and lineage resolver must use the same backend namespace. The backend transaction behind `TokenSink.put` stores the `TokenEntry` and updates the continuation-key index together. Pointing the two clients at unrelated services does not satisfy the protocol even if both methods return successfully. - -`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. - - -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 configured sink and resolver clients inside each worker at app startup. Every client must connect to the same process-shared backend. `LineageStore.is_process_shared()` must return `true`; startup rejects a process-local resolver when `num_workers > 1`. - +The sink receives completed call records from each model-server worker. The lineage resolver lets any worker find a previously committed call from the same rollout. The framework creates a source in its rollout-consumer or trainer process to freeze and read the same records. -With Gym's local file backend, all workers append and resolve from the same token JSONL under one per-rollout lock. The resolver incrementally indexes only newly appended entries in each worker. Freeze and retirement use the same lock. Retirement keeps the lock file and a tombstone, so an old worker cannot recreate a retired capture. - -### 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: +With `rebuild_response: false`, Gym stops after capture. The framework then performs the consumer side: ```python from nemo_gym.token_id_capture.delivery import ( @@ -207,34 +181,31 @@ from nemo_gym.token_id_capture.delivery import ( ) built = await finalize_rollout_token_capture(result, source) -await downstream.put(result) # Must be durable when this returns. +await downstream.put(result) # Establish the framework's durability boundary. 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. +Call `finalize_rollout_token_capture` before training and exclude any result with `mask_sample: true`. Call `retire_rollout_token_capture` only after the rebuilt result has been accepted by durable downstream storage. -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. +Run `nemo_gym.token_id_capture.conformance.run_conformance` against custom sink, source, and lineage factories before using the transport for training. For a multi-process deployment, run the check with independent clients connected to the real shared backend. -### Rollout ids + +Configure sink and lineage classes by import path so Gym constructs them inside every model-server worker. Installing an object only in the launcher process does not configure separately spawned workers. + -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. +## Provide unique rollout ids -Set `_ng_rollout_id` on the run body to provide a distinct key: +Gym normally derives the capture id from the task and rollout indices. If a custom training loop restarts those indices on each step, explicitly provide a unique `_ng_rollout_id`: ```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 +The id may contain letters, digits, dots, dashes, and underscores, and it must begin with a letter or digit. -**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. +## Current limitations -**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`. +- Gym delivers one verified model-call chain per rollout. When Gym can attribute the verifier-scored response to a captured terminal call, it delivers that call's intact ancestor chain and excludes unrelated calls. Without terminal attribution, multiple plausible chains cause masking. +- Harness-generated title, compaction, retry, or sub-agent calls can appear as additional chains. Inspect `terminal_attribution`, `chains`, and `delivered_fraction` to confirm that Gym selected the intended chain. +- Full-tree training requires a trainer data contract that accepts multiple related trajectories. +- Prefix supply requires a compatible inference backend and is not available in stock vLLM.