feat(vllm-model): pin sampling params for on-policy training

[ananthsub]

Lets a model server put a fixed set of sampling parameters on every request it sends to the engine.

Why an absent parameter is the problem

Gym's converters forward only what the caller set.

1. Some harnesses built for interactive serving generally send no sampling fields at all, so the outbound body carries no sampling keys and the engine applies a default of its own rather than the configured one.

• On a backend that validates the request, this fails outright.
• On one that does not, nothing errors and the run trains on samples drawn from a distribution the policy is not being optimized under.

2. Alternatively, a harness could send its own value: a harness that hardcodes temperature on an auxiliary call such as a title generator or a context compressor. That value is not the configured one either.

top_k is a third case, since the Anthropic converter never maps it, so whatever arrives is always a default.

Where it hooks in

flowchart LR
    C["caller request"] --> R{"which API?"}
    R -->|/v1/responses| P["_apply_sampling_overrides<br/>\napplied last"]
    R -->|/v1/chat/completions| P
    R -->|/v1/completions| P
    P --> E["engine request"]
    R -->|/tokenize| T["unchanged<br/>\ntakes no sampling params"]

Loading

All three generation paths. The completions path matters specifically: chat_completions branches to _chat_completions_via_completions_api before the usual preprocessing, so an override applied only in preprocessing would be silently inert there. A pin covering some endpoints and not others is off-policy while reporting that sampling is pinned.

Precedence

The pin is applied last, so it wins over both what the caller sent and anything extra_body merged in. Values replace rather than fill in, because both failure shapes above are real: an absent parameter and a caller-chosen one.

This makes the server authoritative for sampling, which is what #2253 sets out to do through extra_body_override_keys. Configuring sampling through extra_body today would mean moving those values to sampling_overrides to get the same authority.

Configuration

Unset by default, and no config in this repo sets it, so merging this adds a capability and changes no existing run.

policy_model:
  responses_api_models:
    vllm_model:
      sampling_overrides:
        temperature: 1.0
        top_p: 1.0

Configure it on the model server whose callers cannot supply their own sampling params, and leave alone the servers used by callers that already do. A Gym agent builds its request from responses_create_params on the row, which the integrating framework has already stamped with the right values, including a per-rollout validation profile. Pinning that server would replace those values with themselves at best, and discard a deliberate validation profile at worst. Two kinds of caller therefore want two model server instances, which is an existing pattern.

Gym holds no knowledge of any particular framework here. It enforces whatever profile it is given.

What this does not do

A caller that never reads the row cannot receive a per-rollout profile by any mechanism, so a pinned server serves validation rollouts at the profile it was configured with. That is invisible while a run's validation and training profiles are equal, and it is a limitation of the caller rather than of the pin: without the pin, the same run does not work at all.

Testing

responses_api_models/vllm_model/tests/test_app.py::TestSamplingOverrides covers five cases: replacing a value the caller sent, applying when the caller sent nothing, staying a no-op when unset, reaching the completions-API path that skips preprocessing, and winning over extra_body on that path.

The pin runs after the extra_body merge on each of the three paths, so its precedence does not depend on which endpoint a caller picks.

Changed since the approving review

Three things, all narrowing:

• The sampling_overrides block is no longer added to vllm_model_for_training.yaml. That file is overlaid by every training run, including native-agent runs whose sampling is already correct.
• The per-parameter interpolation keys are gone. They resolved through defaults that happened to equal the value every config already used, so no run could tell a wired pin from an unwired one.
• A startup log states which profile the server will enforce, or that it will not pin.

The diff is now two files: vllm_model/app.py and its tests.

Why it is first in the stack

Without it, token capture records ids sampled at the wrong distribution, so everything above this is capturing the wrong thing correctly. It also stands alone and is useful without the rest.

Followed by #2124.

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

[ffrujeri]

Hey @ananthsub, just checking the usecase and pain points here precisely. So the entrypoint where the top_p and temperature is being set is from the Agent Harness itself? I don't remember claude_code or code allowing for that. Is that only for specific open source agents that are allowing it?

[ananthsub]

Hey @ananthsub, just checking the usecase and pain points here precisely. So the entrypoint where the top_p and temperature is being set is from the Agent Harness itself? I don't remember claude_code or code allowing for that. Is that only for specific open source agents that are allowing it?

@ffrujeri I've updated the PR description to clarify this:

1. The harness doesn't set these parameters, but they are expected by the engine worker
2. The harness explicitly harcodes these sampling parameters on internal or auxiliary calls