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

An agent harness chooses its own temperature and top_p. When its rollouts are
used for RL, generation has to match the sampling distribution the policy is
optimized under, or the rollouts are off-policy by an amount nobody measured.
The harness cannot be told to change: it is a released binary, and its settings
are tuned for interactive use.

Gym builds the outbound engine request itself, so it can enforce this. Setting
sampling_overrides on the model server applies those params to every chat
request, after the client's own values, so they win. Unset, nothing changes.

Gym holds no knowledge of any particular training framework here. The
integrating framework sets these to its own generation config; the model server
only enforces them.

Signed-off-by: Ananth Subramaniam <ansubramania@nvidia.com>

Author: ananthsub

responses_api_models/vllm_model/app.py

@@ -64,6 +64,13 @@ class VLLMModelConfig(BaseResponsesAPIModelConfig):
 
     chat_template_kwargs: Optional[Dict[str, Any]] = None
 
+    # Sampling params to force on every chat request, overriding whatever the client sent. An
+    # external harness (e.g. a CLI agent) chooses its own temperature/top_p, but on-policy RL
+    # training requires generation to match the sampling distribution the policy is optimized under.
+    # The integrating training framework sets these to its generation sampling params; Gym only
+    # enforces them and holds no knowledge of any specific framework. Unset means no override.
+    sampling_overrides: Optional[Dict[str, Any]] = None
+
     # Corresponds to the extra_body of OpenAI Client.
     extra_body: Optional[Dict[str, Any]] = None
 
@@ -448,6 +455,12 @@ def _preprocess_chat_completion_create_params(self, request: Request, body_dict:
                 # No user message found — create one with just the audio blocks.
                 body_dict.setdefault("messages", []).append({"role": "user", "content": list(audio_blocks)})
 
+        # Pin sampling params last so they win over anything the client sent. On-policy RL training
+        # requires the generation to match the training worker's sampling config; an external harness
+        # sets its own temperature/top_p, so force them here to keep captured rollouts on-policy.
+        if self.config.sampling_overrides:
+            body_dict.update(self.config.sampling_overrides)
+
         return body_dict
 
     async def chat_completions(

responses_api_models/vllm_model/configs/vllm_model_for_training.yaml

@@ -7,3 +7,10 @@ policy_model:
       model: ${policy_model_name}
       return_token_id_information: true
       uses_reasoning_parser: true
+      # On-policy training: force generation to the sampling params the policy is optimized under,
+      # overriding whatever an external harness requests, so captured rollouts stay on-policy. The
+      # integrating training framework supplies these via the generic keys below (Gym does not know
+      # any framework); the defaults are standard on-policy sampling when a framework sets nothing.
+      sampling_overrides:
+        temperature: ${oc.select:policy_generation_temperature,1.0}
+        top_p: ${oc.select:policy_generation_top_p,1.0}

responses_api_models/vllm_model/tests/test_app.py

@@ -4684,3 +4684,51 @@ async def mock_create_chat_completion(**kwargs):
             )
         # The tokenize endpoint must not be reached once the contract check fails.
         mock_client.create_tokenize.assert_not_called()
+
+
+
+class TestSamplingOverrides:
+    """Forcing the sampling params on every request.
+
+    An external harness picks its own temperature and top_p. On-policy RL requires
+    generation to match the distribution the policy is optimized under, so the
+    server overrides whatever the client sent rather than trusting it.
+    """
+
+    @staticmethod
+    def _server(overrides: dict[str, object] | None) -> VLLMModel:
+        config = VLLMModelConfig(
+            host="0.0.0.0",
+            port=8081,
+            base_url="http://api.openai.com/v1",
+            api_key="dummy_key",  # pragma: allowlist secret
+            model="dummy_model",
+            entrypoint="",
+            name="",
+            return_token_id_information=False,
+            uses_reasoning_parser=False,
+            sampling_overrides=overrides,
+        )
+        return VLLMModel(config=config, server_client=MagicMock(spec=ServerClient, global_config_dict={}))
+
+    def test_overrides_replace_what_the_client_sent(self) -> None:
+        server = self._server({"temperature": 1.0, "top_p": 1.0})
+        out = server._preprocess_chat_completion_create_params(
+            MagicMock(), {"messages": [{"role": "user", "content": "hi"}], "temperature": 0.2, "top_p": 0.5}
+        )
+        assert out["temperature"] == 1.0
+        assert out["top_p"] == 1.0
+
+    def test_overrides_apply_even_when_the_client_sent_nothing(self) -> None:
+        server = self._server({"temperature": 1.0})
+        out = server._preprocess_chat_completion_create_params(
+            MagicMock(), {"messages": [{"role": "user", "content": "hi"}]}
+        )
+        assert out["temperature"] == 1.0
+
+    def test_unset_leaves_the_request_alone(self) -> None:
+        server = self._server(None)
+        out = server._preprocess_chat_completion_create_params(
+            MagicMock(), {"messages": [{"role": "user", "content": "hi"}], "temperature": 0.2}
+        )
+        assert out["temperature"] == 0.2