fix(generators): compute batched rollout metrics from truncated responses (#1821)

## What

In `SkyRLGymGenerator.generate_batched`, each engine response is
truncated to `max_tokens` into `truncated_responses`, and `loss_masks`
are built from those truncated lengths. The `GeneratorOutput` returns
`truncated_responses` as `response_ids`. But the rollout metrics were
computed from the **untruncated** `responses`:

```python
rollout_metrics = get_rollout_metrics(responses, rewards, env_metrics, env_classes, loss_masks)
```

So `generate/{min,max,avg,std}_num_tokens` described sequences longer
than the `response_ids` actually returned and trained, and disagreed
with the assistant-token metrics derived from the truncated `loss_masks`
passed into the same call. `generate/max_num_tokens` could exceed
`max_generate_length`, violating the invariant added in #764.

## Fix

Pass `truncated_responses` instead of `responses` — a one-token change
that matches the async multi-turn path (`generate`), which feeds
`get_rollout_metrics` the same `responses` it returns.

## Test


`tests/train/generators/test_skyrl_gym_generator.py::test_generate_batched_metrics_use_truncated_responses`
sets `max_generate_length=2` below the 4-token mock output and asserts
`generate/max_num_tokens == 2`, matching the truncated `response_ids`
and `loss_masks`. Fails on `main` (`4 == 2`), passes after the fix. All
24 tests in the file pass; black + ruff clean.

Co-authored-by: Claude Opus 4.8 <noreply@anthropic.com>

Author: EazyReal

skyrl/train/generators/skyrl_gym_generator.py

@@ -776,7 +776,7 @@ async def generate_batched(
             # Close the environment
             await self._run_in_executor_if_available(env.close)
 
-        rollout_metrics = get_rollout_metrics(responses, rewards, env_metrics, env_classes, loss_masks)
+        rollout_metrics = get_rollout_metrics(truncated_responses, rewards, env_metrics, env_classes, loss_masks)
 
         if self.generator_cfg.apply_overlong_filtering:
             # set loss mask to 0 if the stop reason is not "stop"

tests/train/generators/test_skyrl_gym_generator.py

@@ -358,6 +358,40 @@ async def test_generate_batched(mock_make, mock_tokenizer, mock_llm, mock_env, g
     assert generator_output["loss_masks"][0] == [1] * len(MOCK_LLM_OUTPUT_IDS)
 
 
+@pytest.mark.asyncio
+@patch("skyrl_gym.make")
+async def test_generate_batched_metrics_use_truncated_responses(
+    mock_make, mock_tokenizer, mock_llm, mock_env, generator_cfg, mock_env_cfg
+):
+    """Rollout metrics must describe the truncated responses that are actually
+    returned/trained, not the raw engine output. With max_generate_length below
+    the engine's output length, generate/max_num_tokens must equal the truncated
+    length, matching response_ids and loss_masks."""
+    generator_cfg.sampling_params.max_generate_length = 2  # < len(MOCK_LLM_OUTPUT_IDS) == 4
+    mock_make.return_value = mock_env
+    mock_env.init.return_value = ([{"role": "user", "content": "Initial input"}], {})
+
+    generator = SkyRLGymGenerator(
+        generator_cfg=generator_cfg,
+        skyrl_gym_cfg=mock_env_cfg,
+        inference_engine_client=mock_llm,
+        tokenizer=mock_tokenizer,
+    )
+    generator.base_conversation_token_ids = []
+
+    input_batch: GeneratorInput = {
+        "prompts": [[{"role": "user", "content": "What is 3 + 5?"}]],
+        "env_extras": [{"answer": "8"}],
+        "env_classes": ["gsm8k"],
+    }
+
+    generator_output: GeneratorOutput = await generator.generate(input_batch)
+
+    assert generator_output["response_ids"][0] == MOCK_LLM_OUTPUT_IDS[:2]
+    assert generator_output["loss_masks"][0] == [1] * 2
+    assert generator_output["rollout_metrics"]["generate/max_num_tokens"] == 2
+
+
 @pytest.mark.asyncio
 @pytest.mark.parametrize("batched", [True, False])
 @patch("skyrl_gym.make")