[qiaochuz-nv]

What does this PR do?

Makes aviary_agent and toolsandbox_agent report the token usage of the whole episode. Before this change,
response.usage held only the last model call's counts. Every multi-step rollout therefore under-reported
input_tokens/output_tokens/total_tokens, and so did every aggregate built from them.

Fixes NVBug 6900411.

Closes #4504

Cause

AviaryAgent.responses() loops model call → /step → model call. It reassigns model_response each turn and
returns AviaryNeMoGymResponse.model_validate(model_response.model_dump() | output_overrides). output_overrides
never set usage, so the last call's usage wins. This construction dates from the original integration (#55); #590
moved the overrides into output_overrides and kept model_response.model_dump() as the base. #3805 added a
post-terminal model call (generate_final_response_after_done), and that call was dropped the same way.

toolsandbox_agent (#1949) copied the same loop and has the same defect.

nemo_gym/reward_profile.py copies response.usage into each rollout's numeric fields, so the per-rollout token
counts and the mean/, median/, max/... token metrics were all under-counted. They are real counts from one
call, so nothing looked wrong. The rollout_token_count_mismatch health check would flag the mismatch, but only
when model-call capture is on.

Fix

• Both agents accumulate accumulate_response_usage(usage, model_response.usage) after each validated model
  response and put "usage": usage in output_overrides. The rest of the returned response (id, status, model)
  still comes from the last call.
• This is the helper simple_agent and gymnasium_agent already use. A call that fails or fails validation
  breaks before it is counted, as before. A call with no usage (a provider that omits it,
  responses_converter.py leaves it None) is skipped, which is the helper's policy. critpt_agent instead
  reports null when any call lacks usage. I kept the shared behaviour, to match the other loop agents.
• Tests: usage on both mocked calls of test_responses_multi_step, including cached and reasoning details, and on
  both calls of the parametrized post-terminal test (aviary). A new two-turn test for toolsandbox. Each fails on
  main with the last call's numbers.

Before the fix

Gym b3103ea8a (current main):

responses_api_agents/aviary_agent/.venv/bin/python -m pytest responses_api_agents/aviary_agent/tests/test_app.py
responses_api_agents/toolsandbox_agent/.venv/bin/python -m pytest responses_api_agents/toolsandbox_agent/tests/test_app.py

FAILED test_responses_multi_step
  {'input_tokens': 1818, 'input_tokens_details': {'cached_tokens': 300}, 'output_tokens': 24, ..., 'total_tokens': 1842}
  != {'input_tokens': 2118, 'input_tokens_details': {'cached_tokens': 400}, 'output_tokens': 1424, ..., 'total_tokens': 3542}
FAILED test_responses_generates_final_response_after_environment_done[True-True-False]   (600/8 != 1100/48)
FAILED test_responses_generates_final_response_after_environment_done[True-True-True]
3 failed, 10 passed
FAILED test_responses_sums_usage_across_turns   (toolsandbox: 1818/24/1842 != 2118/1424/3542)

Live, with model-call capture on (QA model_call_capture_agent_backends, aviary_gsm8k, forced multi-call episode,
h100). response.usage equals the last captured call alone:

aviary on the inference provider: 10 model call(s); response.usage.input_tokens vs sum of calls: expected 6895, got 949
aviary on a real vLLM (Qwen3-8B): 2 model call(s); response.usage.input_tokens vs sum of calls: expected 342, got 62
model_call_capture_agent_backends: 0/2 passed
toolsandbox (example task 0, 2 agent calls): response.usage 839/818/1657 vs sum of agent calls 1355/1088/2443

After the fix

Gym e64106a7a (this branch), same commands:

aviary_agent: 13 passed
toolsandbox_agent: 14 passed
aviary on the inference provider: 10 model call(s)  PASS
aviary on a real vLLM (Qwen3-8B): 2 model call(s)   PASS
model_call_capture_agent_backends: 2/2 passed
toolsandbox task 0: response.usage 1493/1068/2561 == sum of agent calls (516+977 in, 409+659 out)
toolsandbox task 1: response.usage 4575/728/5303 == sum of agent calls

Validation

• gym env test +entrypoint=responses_api_agents/aviary_agent +should_validate_data=true: 13 passed;
  same for toolsandbox_agent: 14 passed. Both new assertions fail on main.
• pre-commit run --all-files: passed.
• QA testcase model_call_capture_agent_backends (h100, GPU lane, inference provider + local Qwen3-8B vLLM), with
  a QA-side assertion that response.usage equals the sum of the captured per-call counts: main 0/2 → 2/2.
• Rollout evidence: the four live aviary rollouts above (provider and vLLM, main and fix). Toolsandbox: its
  registered QA case is blocked on README-model access, so I ran a scratch smoke of the shipped toolsandbox config
  on Nemotron-3-Nano-30B-A3B with capture on, --limit 2. It compared usage with the agent's captured Responses
  calls; the user-simulator's chat calls are excluded. Toolsandbox has no committed integration check beyond its
  unit tests.

Compatibility

• Token metrics for multi-step aviary and toolsandbox rollouts go up: per-rollout input_tokens/output_tokens/
  total_tokens and their aggregates now cover every model call. Single-step episodes are unchanged. Earlier
  numbers for these environments under-count and are not comparable.
• This matches the documented contract (rollout-evidence.mdx: usage is the sum over model responses that
  finished). It also removes a false rollout_token_count_mismatch for multi-call rollouts when capture is on.
• NeMo RL does not read response.usage from Gym rollouts, so training is unaffected.
• Open PRs feat(aviary): support post-terminal model response #3109 and feat(aviary): consolidate Hypotest and support final assistant responses #3792 also edit aviary_agent/app.py; neither touches usage.

🤖 Generated with Claude Code

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

[qiaochuz-nv]

/ok to test e64106a

[github-actions]

🔄 Merge queue validation started! You can track the progress here: https://github.com/NVIDIA-NeMo/Gym/actions/runs/39393435484

[github-actions]

🔄 Merge queue validation started! You can track the progress here: https://github.com/NVIDIA-NeMo/Gym/actions/runs/39393828662