Training on External Agent Harnesses

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.

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.

1
env:
2
  nemo_gym:
3
    token_id_capture:
4
      enabled: true
5
      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.

1
responses_api_agents:
2
  claude_code_agent:
3
    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:

1
policy:
2
  generation:
3
    vllm_cfg:
4
      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.

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.

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.

1
class TokenSink(Protocol):
2
    async def put(self, entry: TokenEntry) -> None: ...
3
    async def mark_incomplete(self, rollout_id: str, model_call_id: str = "") -> None: ...
4
    async def close(self) -> None: ...
5
6
class TokenSource(Protocol):
7
    async def freeze(self, rollout_id: str) -> TokenCaptureSnapshot: ...
8
    async def drop(self, rollout_id: str, *, snapshot_id: str, version: int) -> bool: ...
9
    async def close(self) -> None: ...

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.

1
env:
2
  nemo_gym:
3
    token_id_capture:
4
      enabled: true
5
      sink: my_pkg.sinks:MyDataPlaneSink   # module.path:ClassName
6
      sink_kwargs:
7
        endpoint: ${oc.env:MY_DATAPLANE_URL}
8
        shard: ${oc.select:cluster_shard,0}
9
      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.

1
source = TransferQueueTokenSource(queue_handle)
2
built = await finalize_rollout_token_capture(result, source)
3
await durable_handoff(built)
4
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.

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:

1
from nemo_gym.token_id_capture.delivery import (
2
    finalize_rollout_token_capture,
3
    retire_rollout_token_capture,
4
)
5
6
built = await finalize_rollout_token_capture(result, source)
7
await downstream.put(result)  # Must be durable when this returns.
8
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:

1
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.