feat: prototype trajectory capture from claude code artifacts - #1979
Closed
ffrujeri wants to merge 6 commits into
Closed
feat: prototype trajectory capture from claude code artifacts#1979ffrujeri wants to merge 6 commits into
ffrujeri wants to merge 6 commits into
Conversation
Signed-off-by: Felipe Vieira Frujeri <ffrujeri@nvidia.com>
… possible. Signed-off-by: Felipe Vieira Frujeri <ffrujeri@nvidia.com>
…making the trajectory the single parse of Claude Code output Signed-off-by: Felipe Vieira Frujeri <ffrujeri@nvidia.com>
…ep_no, num_agent_steps); document turn vs step in the sche. Signed-off-by: Felipe Vieira Frujeri <ffrujeri@nvidia.com>
…ants, per-generation records, and agent_telemetry on NeMoGymRespon. Signed-off-by: Felipe Vieira Frujeri <ffrujeri@nvidia.com>
Contributor
|
🌿 Preview your docs: https://nvidia-preview-ffrujeri-trajectory-capture.docs.buildwithfern.com/nemo/gym Here are the markdown pages you've updated: |
Signed-off-by: Felipe Vieira Frujeri <ffrujeri@nvidia.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Extends Gym's native Responses contract so that the agent's
NeMoGymResponseis itself the trajectory — telemetry rides on the contract, nothing is stored twice, and no sidecar object exists. Following the same extension pattern as the*ForTrainingtoken-ID variants:response.outputbecomes the episode's lossless item list where model-produced items are*WithAgentTelemetryvariants tagged withagent_step_no(which model call produced them) andfunction_call_outputitems carryexecution(independent per-call timing, errors, provider metadata); theResponsegains optionalgenerations[](per-model-call native usage + provider identity) andagent_telemetry(provenance, run totals, dropped-event accounting). Model servers leave everythingNone; plain payloads validate to the plain classes unchanged. Theclaude_code_agentis the reference producer: it harvests the session transcript Claude Code already writes to its ephemeralCLAUDE_CONFIG_DIR/projects/(previously deleted unread) and parses it once.Addresses the telemetry requirements of #1867. Complements (does not conflict with) #1715: the schema is agent-generic and lives in core, but no existing model-server or agent code paths change; other harnesses adopt it by writing a thin adapter over
TrajectoryBuilder.Terminology
Two words are used precisely throughout this PR (and in the schema):
simple_agent'smax_stepscounts exactly these. (Providers differ: Claude Code calls these "turns" — itsnum_turns/--max-turnscount model calls; the trajectory normalizes that tonum_agent_steps.)user_messagesteps.Why the existing contracts aren't enough
NeMoGymResponseCreateParamsNonStreamingandNeMoGymResponseare Gym's strict mirrors of the OpenAI Responses API wire contract: one request (input: an item list) → one response (output: an item list, plus a singleusage). Gym serves this contract at two altitudes with different semantics:/v1/responses: one exchange = one model generation. The response envelope (id,model,created_at) andusagedescribe a single sampling./v1/responses(agent-as-model: the agent server registers the identical route with the identical request/response types as the model server —base_responses_api_agent.py:49vsbase_responses_api_model.py:51— so callers cannot distinguish them, and an agent can serve as the "model" of another harness):simple_agent.responses()accepts the same create params, runs a whole turn inside — N agent steps: model calls with tool executions between them — and returns oneResponsewhoseoutputis the in-order concatenation of every step's items. This mirrors OpenAI's own hosted-tools semantics (oneResponsecontainingweb_search_call-style items from a server-side loop, one aggregate usage): the Responses contract is designed to hide an internal loop behind one exchange.To be precise about what "flattening" means here: the episode is naturally a list of lists —
[[step 1's items], [step 2's items], …], each inner list from one generation with its own envelope, usage, and timing. The aggregation flattens it in the structural sense ([[a,b],[c]] → [a,b,c]): every item survives, intact and in order; the partition into steps — and everything attached to it — does not. The partition is not even recoverable afterwards: a generation may emitmessage + callstogether and adjacent generations may both emit messages, so the inner-list boundaries can't be reconstructed from the flat sequence alone. (One genuine content-level merge does happen on main: the Claude Code path inlines reasoning into<think>tags inside message text — which is why the trajectory keeps reasoning as native items and moves the inlining into the derived view.)So the flattening is a legitimate use of the contract, not an abuse — and that is precisely the problem for telemetry. Hiding internal structure is acceptable when the loop runs behind someone else's API; Gym owns this loop. In
simple_agent.responses(), every step'sNeMoGymResponseis validated in scope — each with its ownid,model, and usage — and then the aggregation point discards them:model_response.output = new_outputsreturns the last step's envelope carrying the full flat item list and a running usage sum (zeroingcached_tokens/reasoning_tokens, aTODOin the code).run()stores that flattened response in the rollout row via the verify response, and the collector stamps task/rollout identity (task_idx,rollout_idx,agent_ref) on it. The per-step structure isn't unavailable today — it is in hand at the aggregation point and thrown away there. (claude_code_agentproduces the same artifact shape, with the loop running inside the CLI instead.)What survives is genuinely good: a linear, append-only record of what was said, in native item types (which the trajectory keeps unchanged). What the aggregation discards is everything about how the turn executed.
Running example
A weather assistant with a
get_weathertool — two turns, two agent steps each:What the model server receives per step (the append-only recurrence,
s_{k+1} = s_k ⊕ generation_k ⊕ outputs_k, re-materialized in full each call and continuing across the turn boundary):What today's rollout row records for this episode — and note the sharpest problem first: there is no single answer. A multi-turn episode has no defined record shape on main; each harness family produces a different one:
simple_agentcannot represent the episode: oneresponses()call is one turn (the loop's break atmsg₁is the end of turn 1), and nothing drives a second turn — the row would contain turn 1 only.u2absent entirely (mid-episode user input is never an output item).responses_create_params.input, the rest inresponse.output, plus bolt-onnum_agent_callsand min/mean/max usage aggregates on the verify response because per-step stats have no standard slot.Taking the middle case as the sketch (illustrative numbers):
What this PR records — the same content once, in a lossless
response.outputwhose items carry their own telemetry, plus per-generation records on the Response (abridged):Turns are derivable from the output itself (
role: "user"items /end_turngenerations);reconstruct_model_input(response.output, agent_step_no=3, base_input=[u1])returns exactly step 3's wire input above. Note this is the first well-defined multi-turn record shape: the same structure for asimple_agent-style loop, a Claude-Code-style session, or a tau2-style dialogue — replacing the three incompatible per-harness accidents above.What a consumer can and cannot answer
Trajectory(this PR)out₁is agent-inserted, so a boundary probably follows) — andu2's position is a per-benchmark conventionstop_reason: "end_turn";u2is arole: "user"item in the content planecached_tokenszeroedgenerations[2].usage: nativeNeMoGymResponseUsageincl.cached_tokens: 240, deduplicated per API message; raw provider usage verbatim beside itreconstruct_model_input(output, trajectory, agent_step_no=3, base_input=[u1])→[u1, call₁, out₁, msg₁, u2]— computed from the single stored copyget_weather(Berlin)take; did it fail?function_call_outputcarries onlycall_id/output/status)c2:started_at/ended_at/duration_ms: 498.0,error— independent per call even when one generation issues several in parallelresp_C(provider-side debugging)?idresponse_id/request_idNeMoGymContextBoundaryMessageitem in the output; reconstruction restarts theredropped_recordscounters, plusnum_agent_steps/duration_ms/total_cost_usdfor cross-checkingParallel tool calls: the case that discriminates every alternative
Let one generation issue two calls at once —
get_weather(Paris)andget_weather(Berlin)in a single agent step. There are four separable facts to record: pairing (which output belongs to which call), issuance grouping (both calls came from one generation), ordering semantics (issue order vs. completion order), and execution timing (did they actually overlap). Each record shape handles them differently:simple_agent(main) — calls land in issue order, outputs appended after:Pairing survives (
call_id); grouping is only inferable from adjacency; no timing. And a fact the record cannot show:simple_agentexecutes "parallel" calls serially — the tool loop awaits each POST one at a time (app.py:130-162).Claude-Code-style flat parse (main) — each call is emitted only when its result arrives, so the record is in completion order:
Issue order is gone, simultaneous issuance is unmarked, and a call whose result never arrives is dropped. Note the two flat parses don't even agree with each other on ordering — another instance of the record shape being a per-harness accident.
tau2 (main) — the chat source does group the calls (one assistant message holds
tool_calls=[paris, berlin]), but the conversion to Responses items flattens the group into adjacency; no timing.This PR — all four facts, explicit, without copying an item:
This is validated on a real transcript: one generation, two
Bashcalls, 1561.0 ms / 1615.0 ms overlapping intervals. It also unlocks a diagnostic no flat record can express: run the same episode throughsimple_agentwith an in-process adapter and theexecutionintervals would be non-overlapping — exposing that the harness serializes parallel calls.Model-server capture (#1715) — sees the generation emit
[call_paris, call_berlin](issuance grouping preserved) but never sees execution: no outputs, no timing — tools run in the harness. The cleanest illustration of why the two capture layers are complementary: the model server knows what was asked; only the harness trajectory knows what happened.One-line summary:
call_idpairing survives every alternative; issuance grouping, ordering semantics, and concurrency survive only in the trajectory.(The same structure was validated on a real Claude Code transcript — including one generation issuing two parallel
Bashcalls timed independently at 1561 ms / 1615 ms, and a per-step cache ramp0 → 22,951 → 29,567cached tokens that the summed usage provably erases; see Testing.)Design references: what was borrowed from the OpenAI Agents SDK and ATIF
#1867 allows "ATIF, Responses-plus-telemetry, or Gym-native with exporters" as the format choice; this PR is Responses-plus-telemetry, using both prior designs as references. To make the comparison explicit:
From the OpenAI Agents SDK (an in-memory framework, not a wire format):
RunItems wrap raw Responses items while telemetry lives in tracing spans. We keep the conceptual split but anchor it in the contract: our items are Gym's validated native types (noraw_itemwrapper), the function-span data rides eachfunction_call_outputitem asexecution(started_at/ended_at/duration_ms/error), and the generation-span data is thegenerations[]records on the Response — persisted, rather than exported to a tracing backend and lost to the record.RunResult.raw_responseskeeps oneModelResponseper call, but only in memory. Our generation records are that structure made durable (as index ranges over the output, so nothing is copied).to_input_list()rebuilds only the next call's input. Ourreconstruct_model_input(..., agent_step_no=k)rebuilds any step's input, compaction-aware.Agentreferences and framework lifecycle — not schema-validatable, not serializable, and a framework dependency Gym doesn't need.From ATIF (harbor RFC 0001, a stored-trajectory format):
agent_step_notags over the append-only output, leaner than ATIF's step objects (no content is repeated at all; the initial prompt lives only inresponses_create_params.input).source_call_id+ observationextra; ours ride thefunction_call_outputitem itself (call_idpairing,execution.extra).response.output— and an ATIF exporter stays mechanical if ever needed.Where should the telemetry extend the existing contract? Three candidate surfaces
Nothing prevents adding this telemetry as an extension of the contracts Gym already has — additive optional fields don't break substitutability, and the
*ForTrainingitems are precedent for extending the wire mirrors. The design decision is which surface to extend. There are three candidates; this PR picks the first (evolved: telemetry on the items and the Response, with zero content duplication), and the third is #1715's (complementary, not competing). Earlier drafts used the second; the decisive argument for the first is that once the telemetry is item-anchored and duplication-free, riding the contract means it survives every validation hop (verify requests, resources servers, rollout rows) with no side-channel, and composed agents (agent-as-policy) carry it through automatically.Option 1 — a field on
NeMoGymResponse(the agent/model impersonation surface):Option 2 — an item-free overlay on the rollout record (
BaseVerifyResponse→ one rollout JSONL row; an earlier draft of this PR — superseded):Option 3 — capture at the model server, per generation (#1715's altitude; the training half):
NeMoGymResponseschema_versionresponse_id/request_id*ForTrainingprecedent; the residual cost (fields dead at the model altitude) is accepted and documentedThe rule generalizing the choice: put each fact at the altitude where it is born, on the contract that already lives there — item-shaped facts (which step produced me; how my tool execution went) ride the items, exactly like
*ForTrainingtoken IDs; generation-shaped facts (per-call usage/identity/stop semantics) ride the Response beside the aggregateusagethey decompose; run-shaped facts rideagent_telemetry; token IDs stay at the model server (Option 3) until joined. Items only carry facts that are 1:1 with the item — per-generation facts are keyed one level up (agent_step_nois the foreign key), because generations can exist without items (a call killed mid-step) and group facts stamped on members invite repetition and inconsistency.Capture model: one contract, two ownership modes
Gym will never own every agent loop — and the contract must not care. There are two capture modes producing the same schema:
simple_agent, browsecomp, finance, gymnasium agentsclaude_code_agent(this PR), OpenCode, mini-swe-agentTrajectoryBuilderinsideresponses()as the loop executesResponseis in handmessage.id,requestId); degraded fallback otherwisestarted_at≈ issuing record's write time;completed_atexact)What makes the fidelity difference legible to consumers, instead of a silent quality lottery:
sourcelabels provenance ("transcript"vs"stream_json", later"in_process") — the "dialect/stage" label the issue discussion asks for.nullmeans not observable from this source — never approximated, never fabricated (e.g. no timing in the stream-json fallback rather than step-boundary guesses).dropped_recordscounts events seen but not represented; its absence means nothing was dropped.The deeper structure behind both modes: an agent is a harness applied to a policy, closed under composition —
Agent = Harness(Policy)serves the samecreate params → Responsecontract as the policy itself (OpenAI's hosted-tools flattening semantics), so agents can serve as the policy of more powerful agents. TheResponseis the (lossy-by-design) composition interface; theTrajectoryis the structure-preserving record of one harness level; a composed agent's full observability is a tree of trajectories linked by generation-span identity (response_id/request_id), which is also how subagent support will land. Seeissues/agent-composition-nomenclature.mdfor the full nomenclature and its contract anchoring.Relationship to the existing
parse_stream_jsonpathThe agent already parsed Claude Code's stream-json stdout (
parse_stream_json, app.py:71 on main) into Responses output items. That path could not simply become the trajectory:requestId, and notoolUseResultexecution metadata — those exist only in the on-disk transcript records, which is precisely the [epic] standardised trajectories format #1867 gap (tool timing, per-call identity). The transcript is the trajectory's primary source; stdout events are the degraded fallback.<think>tags inside message text (vs a nativereasoningitem), drops atool_usewhose result never arrives (vs orphan accounting indropped_records), collapses per-model-call boundaries into one flat list, and has no representation for compaction, errors, or per-call usage.Instead, this PR replaces it with a single parse producing the telemetry-bearing Response:
build_trajectory()returns(output_items, generations, agent_telemetry)— the lossless, tagged item list destined forresponse.outputplus the Response-level records.parse_stream_jsonkeeps its signature as a thin wrapper; the old hand-rolled event loop is deleted, and response and telemetry cannot drift apart because they are one object.The response contract changes deliberately (this agent is unbaselined; the change is the point of the Responses-centric design):
reasoningitem instead of<think>-inlined message text (the one genuine content-level merge on main, removed).call_id.resultevent totals summed on top of per-assistant usage) and hardcodedcached_tokensto 0.Verifiers reading the final assistant message are unaffected; consumers that parsed
<think>tags out of message text should read nativereasoningitems instead.What's in the PR
nemo_gym/openai_utils.py— the contract extension:*WithAgentTelemetryitem variants (agent_step_notag;execution: NeMoGymToolExecutionon tool outputs),NeMoGymContextBoundaryMessage(compaction marker), and two optionalNeMoGymResponsefields —generations: list[NeMoGymGeneration]andagent_telemetry: NeMoGymAgentTelemetry. Variants have required discriminating fields and sit first in the item union, so plain payloads deterministically validate to the plain classes (full core suite: 1,141 tests, zero regressions).nemo_gym/trajectory.py—TrajectoryBuilder(build()→(output_items, generations, agent_telemetry); owns agent-step numbering/tags, model-call dedup byresponse_id, call/execution correlation, orphan handling),usage_from_provider()/summed_usage(),reconstruct_model_input()/to_response_create_params()/agent_step_slices()(tag-based, compaction-aware).responses_api_agents/claude_code_agent/trajectory.py— Claude Code adapter: parses transcript records (timestamps,requestId,toolUseResult,sourceToolAssistantUUID,isCompactSummary, sidechains) and stream-json events (fallback; missing telemetry staysnull, never fabricated), drives the builder.responses_api_agents/claude_code_agent/app.py— harvestsprojects/*/*.jsonlfrom the per-run config dir before cleanup (including after timeouts), builds the telemetry-bearing Response in one parse (capture_trajectory: trueknob), fixes the two latent usage bugs (double-counted totals;cached_tokenshardcoded to 0), and deliberately upgradesresponse.outputto the lossless form (see reconciliation section). The verify response carries no extra field — the response is the trajectory.#1867 acceptance criteria
NeMoGymResponseUsageongenerations[], provider identity per call)cache_creation_input_tokens); reasoning tokens 0 for Anthropic (not reported)agent_step_no/request_id/response_id; timestamps per stepTesting
tests/unit_tests/test_trajectory.py— 26 tests on the builder and helpers, including the load-bearing contract proof: telemetry-tagged items,generations, andagent_telemetrysurviveNeMoGymResponse.model_dump → model_validateround trips, while plain model-server payloads validate to the plain classes with telemetryNoneopenai_utilscontract extension causes zero regressionsparse_stream_jsontests encode the lossless response contract (native reasoning items, issue-order parallel calls with per-item tags, unresolved calls kept)[1,1,1,1,1,2,2,3,∅,∅]over a 10-item output, two parallelBashcalls from one generation timed independently on their own output items (1561.0 ms / 1615.0 ms), the per-step cache ramp (cached_tokens0 → 22,951 → 29,567) ingenerations[].usage, the initial prompt correctly absent, and the full Response round-tripping through the contract unchangedTraining path: relationship to #1715 and retokenization drift
This trajectory is eval/observability-only by design — it does not yet support training the policy model. No
prompt_token_ids/generation_token_ids/generation_log_probsare captured, for two reasons: the Claude Code surface doesn't expose them (the transcript and stream-json record text;/v1/messagesreturns no token IDs), and token capture belongs at a different layer (below). The contract is deliberately ready, though:response.outputuses theNeMoGymResponseInputItemunion, which already includes the*ForTrainingvariants (TokenIDLogProbMixin) — an output carrying token IDs validates today unchanged, and each item'sagent_step_notag says exactly which model call it belongs to. What's missing is a source, not a slot.Do not train from this trajectory's text — retokenization drift is real (see vLLM's Agent Lightning post). Re-tokenizing recorded text does not reproduce the sampled token IDs: BPE merge boundaries shift across concatenation points, chat-template re-application differs from what the server rendered, and tool-call/reasoning serialization round-trips differently. RL importance ratios (PPO/GRPO) need logprobs of the exact sampled tokens; drifted tokens make "on-policy" data silently off-policy. Main's parse was its own cautionary example: it inlined reasoning as
<think>…</think>text — a re-serialization the policy never sampled. The lossless content plane removes exactly that class of drift source.The composition with #1715 (model-server capture) is the training path — the two capture layers are complementary halves of one record, split the way Agent Lightning splits it (tokens captured at the serving layer, semantics reconstructed around them):
The join key already exists on this PR's side: each
NeMoGymGeneration'srequest_id/response_id. The missing third piece is small: stamp rollout identity onto the server side — for the Claude Code path the model server is already in the loop whenmodel_serveris set, and Claude Code supportsANTHROPIC_CUSTOM_HEADERS, so_run_claude_codecan inject a per-rollout correlation header that the server'sStepRecords carry. A merge step then grafts captured token IDs into the items tagged with the matching agent step (upgrading them to*ForTrainingvariants in place), with the telemetry contributing what tokens alone can't: agent-step tags and tool-output items for loss masking, compaction-aware context reconstruction, and reward for credit assignment. No retokenization anywhere — tokens flow from sampling time to training verbatim; a cheap validator (retokenize, diff against captured IDs, alert on mismatch) turns residual drift from a silent bias into a monitored invariant.Limitations / follow-ups
started_atfor a tool call is the issuing assistant record's write time (execution start isn't observable without hooks, which the bare runtime avoids);completed_atis exact.simple_agent, the adapter is even simpler than Claude Code's: driveTrajectoryBuilderinsideresponses(), where every per-callNeMoGymResponseis already in hand (exact boundaries, real per-call usage) and tool timing can be measured around the resources-server calls rather than reconstructed from artifacts. External-harness agents follow the Claude Code adapter pattern; feat(observability): add per-rollout model-call capture #1715's model-server capture supplies the token layer either way.StepRecord⋈agent_stepmerge step, and a retokenization-drift validator.