feat(nemo-gym): train on rollouts from external agent harnesses

[ananthsub]

Trains on rollouts from an agent harness that drives its own model calls and returns no token IDs.

Gym captures the exact token IDs and log probabilities inside its model server and rebuilds each rollout into a trainable Responses payload. NeMo-RL assigns rollout identity, constructs the TokenSource inside its rollout actor, consumes the rebuilt result, and retires the frozen capture only after the rollout consumer accepts it.

Depends on the Gym stack: NVIDIA-NeMo/Gym#2190, #2124, #2125, #2126, #2341, #2180, #2181, and #2349.

Control and data flow

sequenceDiagram
    autonumber
    participant RL as NemoGym actor
    participant G as Gym run_examples
    participant H as Agent harness
    participant M as Gym model-server worker
    participant K as TokenSink
    participant S as TokenSource
    participant T as Trainer consumer

    RL->>RL: Run token_source_factory inside actor
    RL->>RL: Assign _ng_rollout_id
    RL->>G: Dispatch selected rows
    G->>H: Run agent
    H->>M: Model calls under /training-token-capture/{rollout_id}
    M->>K: Append TokenEntry records
    G-->>RL: Return rollout without token metadata
    RL->>S: Freeze and read snapshot
    RL->>RL: Rebuild response and preserve original media items
    RL-->>T: Yield trainable rollout
    T-->>RL: Request next rollout
    RL->>S: Conditionally retire accepted snapshot

Loading

Integration behavior

• Capture selection comes from Gym's resolved static agent configuration. token_id_capture.all_agents: true opts every agent into capture for training-only runs.
• _build_token_capture_source uses a local TokenCaptureStore when token_id_capture.dir is configured.
• A training framework passes token_source_factory to spinup_nemo_gym_actor. Ray serializes the factory with the actor constructor, and the factory runs inside the actor process and virtual environment.
• The factory must return Gym's TokenSource protocol. It may wrap a TransferQueue handle or any other framework-owned transport client.
• Gym model-server workers construct the corresponding configured TokenSink and LineageStore proxies in their own processes. Their adapter package must be installed in the model-server environment.
• The source, sink, and lineage proxies do not share Python objects or virtual environments. They communicate through the framework-owned transport.
• Frozen records remain available while the rollout is yielded. Conditional retirement runs only after the consumer requests the next item, so failed or abandoned consumption does not delete evidence prematurely.
• shutdown() closes the actor-owned source before stopping Gym servers.
• Full UUID rollout IDs avoid collisions across actors and restarts.
• Multimodal postprocessing uses the original response output, while training tokens come from the rebuilt output.
• GRPO grouping keeps dataset prompt identity instead of grouping on the external harness's captured first-call prompt.

Source construction

The default file adapter is selected with token_id_capture.dir. A framework data plane injects an actor-local factory instead of naming a source class in Gym configuration.

source_factory = TransferQueueTokenSourceFactory(queue_handle)
actor = spinup_nemo_gym_actor(
    ...,
    token_source_factory=source_factory,
)

The source must implement Gym's TokenSource protocol: freeze, conditional drop, and idempotent close. The model-server side uses the corresponding TokenSink and LineageStore protocols supplied by the framework.

Metrics

The integration reports token_capture/rebuilt_fraction, rollouts_rebuilt, rollouts_unbuilt, calls_per_rollout_mean, chains_per_rollout_mean, delivered_fraction_mean, quarantined_fraction_mean, empty_generation_calls, parent_link_failures, masked_rollouts, and incomplete_rollouts.

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

[copy-pr-bot]

Auto-sync is disabled for draft pull requests in this repository. Workflows must be run manually.

Contributors can view more details about this message here.

[github-actions]

✅ Submodule Fast-Forward Check Results

Check based on commit: 93f23d3 (PR #3407 from ananthsub/token-id-capture-rebased)

✅ Submodules that are properly updated:

Gym: ✅ PR branch is ahead of main branch (fast-forward)

All submodule changes look good! ✨