Support agent-specific num_repeats in ng_collect_rollouts

[gwarmstrong]

Agent-specific num_repeats for ng_collect_rollouts

Motivation

Currently ng_collect_rollouts applies one global num_repeats to every row,
even when the input JSONL mixes multiple agents (different agent_ref.name).
That makes e.g., "run simple_agent for pass@32 alongside swe_agent for pass@1 in
one job" awkward. You either run separate jobs or downsample and recompute metrics after the fact.

Change

num_repeats on RolloutCollectionConfig is now Union[int, Dict[str, int]]
(default 1):

• int form — unchanged behavior; applies to every row.
• dict form — keys are agent_ref.name. The special key _default is the
  fallback for agents not explicitly listed. Without _default, any row whose
  agent isn't a key in the dict raises a single consolidated error listing
  every unlisted agent. Dict keys that never appear in any input row emit
  a UserWarning.

How to verify

Run end-to-end against integrate.api.nvidia.com. From Gym/, after
uv sync --extra dev and export NVIDIA_API_KEY=...:

# 1. Point Gym at integrate.api.nvidia.com.
cat > env.yaml <<EOF
policy_base_url: https://integrate.api.nvidia.com/v1
policy_api_key: ${NVIDIA_API_KEY:?export NVIDIA_API_KEY=nvapi-...}
policy_model_name: nvidia/nvidia-nemotron-nano-9b-v2
EOF

# 2. Wire two agent server instances backed by the shipped simple_agent.
#    Distinct instance names are what lets dict-form num_repeats target them
#    independently.
cat > /tmp/two_agents_demo.yaml <<'EOF'
agent_alpha:
  responses_api_agents:
    simple_agent:
      entrypoint: app.py
      resources_server: {type: resources_servers, name: example_single_tool_call}
      model_server:    {type: responses_api_models, name: policy_model}
agent_beta:
  responses_api_agents:
    simple_agent:
      entrypoint: app.py
      resources_server: {type: resources_servers, name: example_single_tool_call}
      model_server:    {type: responses_api_models, name: policy_model}
EOF

# 3. Generate a 2-row input — one row pinned to each agent instance.
python3 - > /tmp/mixed_input.jsonl <<'PY'
import json
prompt = [
    {"role": "developer", "content": "You are a helpful assistant."},
    {"role": "user", "content": "what is 2+2?"},
]
for name in ("agent_alpha", "agent_beta"):
    print(json.dumps({
        "responses_create_params": {"input": prompt, "tools": []},
        "agent_ref": {"name": name},
    }))
PY

# 4. Start ng_run in this terminal and leave it running until all servers
#    log "ready" (head server + 4 children).
ng_run \
  "+config_paths=[resources_servers/example_single_tool_call/configs/example_single_tool_call.yaml,responses_api_models/vllm_model/configs/vllm_model.yaml,/tmp/two_agents_demo.yaml]"

# 5. In a separate terminal (same Gym/ cwd), run rollouts with per-agent num_repeats.
ng_collect_rollouts \
  +input_jsonl_fpath=/tmp/mixed_input.jsonl \
  +output_jsonl_fpath=/tmp/mixed_rollouts.jsonl \
  +num_samples_in_parallel=4 \
  +upload_rollouts_to_wandb=false \
  '+responses_create_params={max_output_tokens: 64, temperature: 0.0}' \
  '+num_repeats={agent_alpha: 4, agent_beta: 1}'

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

[cmunley1]

can we add to docs or at least open an issue to document it?

[cmunley1]

I wonder if a separate field like num_repeats_default: Optional[int] would be cleaner than _default?

[gwarmstrong]

can ultimately take it wherever you want, the arguments in favor of a single field as it is currently implemented are:
(1) a single field to update when you want to modify the behavior
(2) you don't have to do extra handling for the case where e.g., num_repeats is an int, and num_repeats_default is an int. Which do you use in that case? How does a user reason about how to set it? the _default key is a bit less ambiguous.

[cmunley1]

looks good!

[gwarmstrong]

can we add to docs or at least open an issue to document it?

just updated docs