fix(rollout): preserve sampled EOS probabilities

skyrl/train/generators/skyrl_gym_generator.py

@@ -563,15 +563,15 @@ async def agent_loop(
                     response_ids
                 ), f"loss_mask and response_ids should have the same length, got {len(loss_mask)} and {len(response_ids)}"
 
-            appended_eos_token = False
             if not self.use_conversation_multi_turn:
                 assert response_ids is not None and loss_mask is not None
                 if stop_reason != "length" and response_ids and response_ids[-1] != self.tokenizer.eos_token_id:
+                    # This EOS was not sampled. Keep it for formatting, but exclude
+                    # it from the policy loss and importance-sampling ratios.
                     response_ids.append(self.tokenizer.eos_token_id)
-                    loss_mask.append(1)
+                    loss_mask.append(0)
                     if rollout_logprobs is not None:
                         rollout_logprobs.append(0.0)
-                    appended_eos_token = True
 
             if self.generator_cfg.step_wise_trajectories:
                 for per_step_output, (reward, resp_end_idx) in zip(agent_loop_output.step_outputs, per_step_rewards):
@@ -580,7 +580,7 @@ async def agent_loop(
                     # in-place update to per-token reward
                     per_step_output.reward = per_token_reward
             else:
-                reward_out = self._build_per_token_rewards(per_step_rewards, response_ids, appended_eos_token)
+                reward_out = self._build_per_token_rewards(per_step_rewards, response_ids)
 
                 agent_loop_output = TrajectoryOutput(
                     response_ids=response_ids,
@@ -606,15 +606,14 @@ async def agent_loop(
             await self.inference_engine_client.finish_session(session_id)
 
     def _build_per_token_rewards(
-        self, per_step_rewards: List[Tuple[float, Optional[int]]], response_ids: List[int], appended_eos_token: bool
+        self, per_step_rewards: List[Tuple[float, Optional[int]]], response_ids: List[int]
     ) -> Union[float, List[float]]:
         """
         Build reward output from per-step rewards.
 
         Args:
             per_step_rewards: List of (reward, response_end_token_idx) tuples for each step
             response_ids: List of response token IDs
-            appended_eos_token: Whether an EOS token was manually appended at the end
 
         Returns:
             Union[float, List[float]]: If custom_chat_template is used, returns the last step's reward (float).
@@ -628,18 +627,14 @@ def _build_per_token_rewards(
         else:
             # Build token-level rewards placed at assistant turn boundaries
             token_level_rewards: List[float] = [0.0] * len(response_ids)
-            for i, (step_reward, idx) in enumerate(per_step_rewards):
+            for step_reward, idx in per_step_rewards:
                 assert step_reward is not None
+                if idx < 0:
+                    # An empty response has no generated token to receive a reward.
+                    continue
                 if idx >= len(response_ids):
                     break
-                if appended_eos_token and i == len(per_step_rewards) - 1:
-                    # NOTE(Charlie): If we appended the eos token, we need to place
-                    # the reward at the last token (the manually appended eos token)
-                    # rather than the last turn's assistant-generated token. This matches
-                    # the logic in trainer.py::postprocess_generator_output when rewards are List[float].
-                    token_level_rewards[-1] = step_reward
-                else:
-                    token_level_rewards[idx] += step_reward
+                token_level_rewards[idx] += step_reward
             reward_out = token_level_rewards
         return reward_out
 
@@ -1164,18 +1159,20 @@ def _update_agent_loop_state_with_singleturn_chat_template(
 
         Returns:
             AgentLoopState: Updated agent loop state with appended turn IDs, loss mask, and logprobs.
-                The EOS token is removed from response tokens (if present) since we are continuing
-                the current assistant message. Observations are encoded directly without chat template formatting.
+                An intermediate EOS is removed while continuing the current assistant message.
+                A final sampled EOS and its logprob are preserved. Observations are encoded directly
+                without chat template formatting.
         """
         agent_loop_state.chat_history = self._update_chat_history(
             agent_loop_state.chat_history, turn_output.output, turn_output.new_obs
         )
 
         obs_ids_to_add = turn_output.obs_ids
 
-        # Remove EOS token from response tokens since we are continuing the current assistant message
+        # Only remove EOS while continuing the current assistant message. The
+        # final sampled EOS is a real action with its own rollout probability.
         new_resp_tokens = turn_output.output_ids.copy()
-        if new_resp_tokens and new_resp_tokens[-1] == self.tokenizer.eos_token_id:
+        if not agent_loop_state.done and new_resp_tokens and new_resp_tokens[-1] == self.tokenizer.eos_token_id:
             new_resp_tokens = new_resp_tokens[:-1]
 
         turn_ids = new_resp_tokens + obs_ids_to_add

tests/train/generators/test_singleturn_eos_logprobs.py

@@ -0,0 +1,213 @@
+"""Single-assistant-message EOS bookkeeping at the rollout/training boundary."""
+
+import math
+from types import SimpleNamespace
+from unittest.mock import AsyncMock
+
+import numpy as np
+import pytest
+import torch
+
+from skyrl.backends.skyrl_train.utils.off_policy_correction_utils import (
+    compute_outlier_token_mask,
+    compute_tis_ratio,
+)
+from skyrl.train.config import (
+    ChatTemplateConfig,
+    GeneratorConfig,
+    OffPolicyCorrectionConfig,
+    SkyRLGymConfig,
+)
+from skyrl.train.generators.skyrl_gym_generator import SkyRLGymGenerator
+
+
+class Tokenizer:
+    eos_token_id = 4
+    eos_token = "<eos>"
+
+    def apply_chat_template(self, messages, **kwargs):
+        return [101, 102]
+
+    def encode(self, text, **kwargs):
+        return [77] if text else []
+
+
+async def run_agent_loop(monkeypatch, turns, *, get_logprobs=True, max_input_length=256):
+    """Run the real generator, replacing only its tokenizer, environment and LLM."""
+    cfg = GeneratorConfig()
+    cfg.batched = False
+    cfg.use_conversation_multi_turn = False
+    cfg.sampling_params.logprobs = 0 if get_logprobs else None
+    cfg.chat_template = ChatTemplateConfig(source="name", name_or_path=None)
+    cfg.inference_engine.enable_return_routed_experts = any("routes" in turn for turn in turns)
+    env_cfg = SkyRLGymConfig()
+    env_cfg.max_env_workers = 0
+    prompts = []
+    turn_index = 0
+
+    class Env:
+        def init(self, prompt):
+            return prompt, {}
+
+        def step(self, text):
+            nonlocal turn_index
+            turn = turns[turn_index]
+            turn_index += 1
+            return {
+                "observations": turn.get("observations", []),
+                "reward": turn.get("reward", 1.0),
+                "done": turn_index == len(turns),
+            }
+
+        def get_metrics(self):
+            return {}
+
+        def close(self):
+            pass
+
+    async def generate(input_batch, model=None):
+        prompts.append(input_batch["prompt_token_ids"][0].copy())
+        turn = turns[turn_index]
+        return {
+            "responses": ["answer" if turn["tokens"] else ""],
+            "response_ids": [turn["tokens"].copy()],
+            "response_logprobs": [turn["logprobs"].copy()] if get_logprobs else None,
+            "stop_reasons": [turn.get("stop_reason", "stop")],
+            "rollout_expert_indices": [turn["routes"]] if "routes" in turn else None,
+        }
+
+    client = SimpleNamespace(generate=generate, finish_session=AsyncMock())
+    monkeypatch.setattr("skyrl_gym.make", lambda *args, **kwargs: Env())
+    generator = SkyRLGymGenerator(cfg, env_cfg, client, Tokenizer())
+    output = await generator.agent_loop(
+        [{"role": "user", "content": "Q"}],
+        "gsm8k",
+        {},
+        max_tokens=32,
+        max_input_length=max_input_length,
+        sampling_params={"max_tokens": 32, "logprobs": cfg.sampling_params.logprobs},
+    )
+    client.finish_session.assert_awaited_once()
+    return output, prompts
+
+
+@pytest.mark.asyncio
+async def test_sampled_final_eos_preserves_on_policy_importance_ratios(monkeypatch):
+    sampled_logprobs = [-0.2, math.log(0.01)]
+    output, _ = await run_agent_loop(monkeypatch, [{"tokens": [10, 4], "logprobs": sampled_logprobs}])
+
+    assert output.response_ids == [10, 4]
+    assert output.loss_mask == [1, 1]
+    assert output.reward == [0.0, 1.0]
+
+    # Unchanged trainer/inference policies must have ratio 1, even if EOS is unlikely.
+    old_logprobs = torch.tensor([sampled_logprobs])
+    rollout_logprobs = torch.tensor([output.rollout_logprobs])
+    loss_mask = torch.tensor([output.loss_mask])
+    correction = OffPolicyCorrectionConfig(outlier_token_is_threshold_low=0.1)
+    ratio, _ = compute_tis_ratio(old_logprobs, rollout_logprobs, loss_mask, "sequence", correction)
+    accepted, _ = compute_outlier_token_mask(old_logprobs, rollout_logprobs, loss_mask, correction)
+    torch.testing.assert_close(ratio, torch.ones_like(ratio))
+    torch.testing.assert_close(accepted, torch.ones_like(accepted))
+    assert output.rollout_logprobs == sampled_logprobs
+
+
+@pytest.mark.asyncio
+async def test_synthetic_eos_is_not_an_action_or_reward_target(monkeypatch):
+    output, _ = await run_agent_loop(monkeypatch, [{"tokens": [10], "logprobs": [-0.2]}])
+
+    assert output.response_ids == [10, 4]
+    assert output.loss_mask == [1, 0]
+    assert output.rollout_logprobs == [-0.2, 0.0]
+    assert output.reward == [1.0, 0.0]
+
+    correction = OffPolicyCorrectionConfig(outlier_token_is_threshold_low=0.1)
+    old_logprobs = torch.tensor([[-0.2, -10.0]])
+    rollout_logprobs = torch.tensor([output.rollout_logprobs])
+    loss_mask = torch.tensor([output.loss_mask])
+    ratio, _ = compute_tis_ratio(old_logprobs, rollout_logprobs, loss_mask, "sequence", correction)
+    accepted, _ = compute_outlier_token_mask(old_logprobs, rollout_logprobs, loss_mask, correction)
+    torch.testing.assert_close(ratio, torch.ones_like(ratio))
+    torch.testing.assert_close(accepted, torch.ones_like(accepted))
+
+
+@pytest.mark.asyncio
+async def test_only_intermediate_eos_is_removed_and_routes_stay_aligned(monkeypatch):
+    first_routes = np.arange(3, dtype=np.uint8).reshape(3, 1, 1)
+    final_routes = np.arange(5, dtype=np.uint8).reshape(5, 1, 1)
+    output, prompts = await run_agent_loop(
+        monkeypatch,
+        [
+            {
+                "tokens": [10, 4],
+                "logprobs": [-0.2, -3.0],
+                "reward": 0.3,
+                "observations": [{"role": "user", "content": "observation"}],
+                "routes": first_routes,
+            },
+            {"tokens": [11, 4], "logprobs": [-0.4, -5.0], "reward": 1.7, "routes": final_routes},
+        ],
+    )
+
+    assert prompts == [[101, 102], [101, 102, 10, 77]]
+    assert output.response_ids == [10, 77, 11, 4]
+    assert output.loss_mask == [1, 0, 1, 1]
+    assert output.rollout_logprobs == [-0.2, 0.0, -0.4, -5.0]
+    assert output.reward == [0.3, 0.0, 0.0, 1.7]
+    # The sampled final token was not itself evaluated: keep the engine's captured
+    # prefix+response-minus-one routes, without fabricating an EOS routing row.
+    np.testing.assert_array_equal(output.rollout_expert_indices, final_routes)
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize(
+    "tokens,logprobs,stop_reason,expected_tokens,expected_rewards",
+    [
+        ([4], [-3.0], "stop", [4], [1.0]),
+        ([], [], "stop", [], []),
+        ([], [], "length", [], []),
+        ([10], [-0.2], "length", [10], [1.0]),
+    ],
+)
+async def test_empty_eos_only_and_truncated_outputs(
+    monkeypatch, tokens, logprobs, stop_reason, expected_tokens, expected_rewards
+):
+    output, _ = await run_agent_loop(
+        monkeypatch, [{"tokens": tokens, "logprobs": logprobs, "stop_reason": stop_reason}]
+    )
+    assert output.response_ids == expected_tokens
+    assert output.loss_mask == [1] * len(expected_tokens)
+    assert output.rollout_logprobs == logprobs
+    assert output.reward == expected_rewards
+    assert output.stop_reason == stop_reason
+
+
+@pytest.mark.asyncio
+async def test_context_limit_does_not_restore_intermediate_eos(monkeypatch):
+    output, prompts = await run_agent_loop(
+        monkeypatch,
+        [
+            {
+                "tokens": [10, 4],
+                "logprobs": [-0.2, -3.0],
+                "observations": [{"role": "user", "content": "observation"}],
+            },
+            {"tokens": [11, 4], "logprobs": [-0.4, -5.0]},
+        ],
+        max_input_length=3,
+    )
+    assert prompts == [[101, 102]]
+    assert output.response_ids == [10]
+    assert output.loss_mask == [1]
+    assert output.rollout_logprobs == [-0.2]
+    assert output.reward == [1.0]
+    assert output.stop_reason == "length"
+
+
+@pytest.mark.asyncio
+async def test_sampled_eos_without_logprob_capture(monkeypatch):
+    output, _ = await run_agent_loop(monkeypatch, [{"tokens": [10, 4]}], get_logprobs=False)
+    assert output.response_ids == [10, 4]
+    assert output.loss_mask == [1, 1]
+    assert output.rollout_logprobs is None
+    assert output.reward == [0.0, 1.0]

tests/train/generators/test_skyrl_gym_generator.py

@@ -388,13 +388,13 @@ def mock_generate(_, model=None):
 
         expected_response_ids = mock_llm_output_ids.copy()
         if has_eos_in_mock:
-            # Had EOS: removed then re-added, so final IDs same as mock
+            # A sampled final EOS remains a loss-active action.
             expected_response_ids = mock_llm_output_ids
         else:
             # No EOS: just add it
             expected_response_ids = mock_llm_output_ids + [mock_tokenizer.eos_token_id]
 
-        expected_loss_mask = [1] * len(expected_response_ids)
+        expected_loss_mask = [1] * len(mock_llm_output_ids) + ([] if has_eos_in_mock else [0])
 
     if logprobs_setting is not None:
         assert output.rollout_logprobs is not None
@@ -1083,9 +1083,8 @@ def step(self, action):
 
     # Response ids layout: step1 (3 tokens) + obs (1) + step2 (3) + final eos (1) = 8
     assert len(out.response_ids) == 8
-    # Indices: 2 (end of step1 assistant), 6 (end of step2 assistant), 7 (manually appended eos token)
-    # Note that the last reward is placed at the 7 instead of at 6 since we manually move
-    # it using the flag `appended_eos_token` in skyrl_gym_generator.py
+    # Reward indices: 2 (end of step1 without its intermediate EOS),
+    # 7 (the sampled final EOS, preserved as part of step2).
     expected_rewards = [0.0, 0.0, 0.3, 0.0, 0.0, 0.0, 0.0, 1.7]
     assert isinstance(out.reward, list)
     assert out.reward == expected_rewards
@@ -1348,16 +1347,14 @@ def mock_make_func(*args, **kwargs):
     extras = {}
     out = await generator.agent_loop(prompt, mock_env_cfg.env_class, extras, max_tokens=5, max_input_length=1000)
 
-    # Untruncated response would be: 4 (step1) + 4 (step2) + 1 (final eos) = 9; we expect truncation to 5
+    # Four generated tokens followed by one synthetic EOS.
     assert len(out.response_ids) == 5
     assert isinstance(out.reward, list)
     assert len(out.reward) == 5
 
-    # Step1 end index relative should be 4 (0-based) - reward placed at EOS token
-    # NOTE(Dev): Because we manually append the eos token to the response, the reward is placed at the last token;
-    # See Charlie's comment in skyrl_gym_generator.py for more details.
-
-    assert out.reward[4] == 2.0
+    # Reward belongs to the last generated token, not the masked synthetic EOS.
+    assert out.reward == [0.0, 0.0, 0.0, 2.0, 0.0]
+    assert out.loss_mask == [1, 1, 1, 1, 0]
     assert sum(out.reward) == 2.0
     assert out.stop_reason == "stop"