diff --git a/.github/workflows/integration_test_4gpu_rl.yaml b/.github/workflows/integration_test_4gpu_rl.yaml
index 4301f7fd95..524300a127 100644
--- a/.github/workflows/integration_test_4gpu_rl.yaml
+++ b/.github/workflows/integration_test_4gpu_rl.yaml
@@ -64,10 +64,11 @@ jobs:
         # Install uv for faster dependency resolution
         pip install uv
 
-        # 1. Install Monarch and TorchStore
+        # 1. Install Monarch, TorchStore, and Renderers
         uv pip install torchmonarch
         uv pip install --no-deps "git+https://github.com/meta-pytorch/torchstore.git@main"
         uv pip install pygtrie portpicker
+        uv pip install "git+https://github.com/PrimeIntellect-ai/renderers.git@main"
 
         # 2. Install batch-invariant ops
         uv pip install --no-deps "git+https://github.com/thinking-machines-lab/batch_invariant_ops.git@main"
diff --git a/.github/workflows/integration_test_8gpu_rl_h100.yaml b/.github/workflows/integration_test_8gpu_rl_h100.yaml
index a4f984375f..fdfb416024 100644
--- a/.github/workflows/integration_test_8gpu_rl_h100.yaml
+++ b/.github/workflows/integration_test_8gpu_rl_h100.yaml
@@ -66,10 +66,11 @@ jobs:
         # Install uv for faster dependency resolution
         pip install uv
 
-        # 1. Install Monarch and TorchStore
+        # 1. Install Monarch, TorchStore, and Renderers
         uv pip install torchmonarch
         uv pip install --no-deps "git+https://github.com/meta-pytorch/torchstore.git@main"
         uv pip install pygtrie portpicker
+        uv pip install "git+https://github.com/PrimeIntellect-ai/renderers.git@main"
 
         # 2. Install batch-invariant ops
         uv pip install --no-deps "git+https://github.com/thinking-machines-lab/batch_invariant_ops.git@main"
diff --git a/torchtitan/experiments/rl/README.md b/torchtitan/experiments/rl/README.md
index 863404feb1..51052f7e7a 100644
--- a/torchtitan/experiments/rl/README.md
+++ b/torchtitan/experiments/rl/README.md
@@ -27,11 +27,12 @@ uv venv --python 3.12 titan-rl
 source titan-rl/bin/activate
 ```
 
-1. Install Monarch and TorchStore from main:
+1. Install Monarch, TorchStore, and Renderers from main:
 ```bash
 uv pip install torchmonarch
 uv pip install --no-deps "git+https://github.com/meta-pytorch/torchstore.git@main"
 uv pip install pygtrie portpicker
+uv pip install "git+https://github.com/PrimeIntellect-ai/renderers.git@main"
 ```
 
 2. Install Flash Attention 3 kernels:
diff --git a/torchtitan/experiments/rl/actors/generator.py b/torchtitan/experiments/rl/actors/generator.py
index 13f201210d..1b9714df2d 100644
--- a/torchtitan/experiments/rl/actors/generator.py
+++ b/torchtitan/experiments/rl/actors/generator.py
@@ -152,9 +152,6 @@ def get_vllm_compilation_config(
 class SamplingConfig:
     """Sampling parameters passed to vLLM's SamplingParams."""
 
-    n: int = 8
-    """Number of completions to generate per prompt (vLLM SamplingParams.n)."""
-
     temperature: float = 0.8
     """Sampling temperature. 0.0 = greedy, higher = more random."""
 
@@ -258,6 +255,7 @@ def __init__(
         compile_config: CompileConfig,
         max_num_seqs: int,
         output_dir: str,
+        stop_token_ids: list[int],
     ):
         init_logger()
         sl.init_structured_logger(
@@ -275,9 +273,13 @@ def __init__(
         # the upper bound for concurrent sequences, determines KV-cache
         # block allocation (and therefore GPU memory usage), and bounds
         # the CUDA graph capture sizes.  Always computed by the caller
-        # (RLTrainer) as num_prompts_per_step * sampling.n.
+        # (RLTrainer) as num_groups_per_rollout_batch * group_size.
         self._max_num_seqs = max_num_seqs
 
+        # Renderer role-boundary stop tokens (e.g. Qwen3 `<|im_end|>`), injected by the
+        # controller;
+        self._stop_token_ids = stop_token_ids
+
         # Register TorchTitan model + parser with vLLM
         registry_to_vllm(
             model_spec,
@@ -388,17 +390,19 @@ async def generate(
         self,
         tokenized_prompts: list[list[int]],
         *,
+        request_ids: list[str],
         sampling_config: SamplingConfig | None = None,
         metrics_prefix: str = "generator",
     ) -> tuple[list[Completion], list[m.Metric]]:
         """Generate completions and generator metrics for tokenized prompts.
 
         Takes ``tokenized_prompts`` as ``[num_prompts][prompt_tokens]``.
-        Returns completions as ``[num_prompts * n]`` plus generator metrics,
-        where ``n`` is the resolved ``SamplingConfig.n`` completions per prompt.
+        Returns completions in the same order as ``request_ids`` plus generator
+        metrics.
 
         Args:
             tokenized_prompts: Tokenized prompts shaped ``[num_prompts][prompt_tokens]``.
+            request_ids: One id per prompt echoed on each ``Completion.request_id``.
             sampling_config: Optional per-call override for the generator's
                 default SamplingConfig. ``seed`` always comes from
                 ``config.debug.seed`` (not part of SamplingConfig).
@@ -419,12 +423,20 @@ async def generate(
                 temperature=_sampling_config.temperature,
                 top_p=_sampling_config.top_p,
                 max_tokens=_sampling_config.max_tokens,
-                n=_sampling_config.n,
+                n=1,  # group_size pre-expands prompts; the RL loop always samples n=1
+                stop_token_ids=self._stop_token_ids or None,
                 seed=self.config.debug.seed,
                 logprobs=1,
                 output_kind=RequestOutputKind.FINAL_ONLY,
             )
 
+            if len(request_ids) != len(tokenized_prompts):
+                raise ValueError(
+                    f"got {len(request_ids)} request_ids for {len(tokenized_prompts)} prompts"
+                )
+            if len(set(request_ids)) != len(request_ids):
+                raise ValueError(f"request_ids must be unique; got {request_ids}")
+
             # render_cmpl is vLLM's input-pipeline entry.
             # The tokenize step is a no-op for already-tokenized prompts. The
             # lower-level alternative is vllm.inputs.tokens_input; we use the
@@ -432,9 +444,11 @@ async def generate(
             engine_inputs = self._engine.renderer.render_cmpl(
                 [{"prompt_token_ids": ids} for ids in tokenized_prompts]
             )
-            for i, engine_input in enumerate(engine_inputs):
+            for engine_input, request_id in zip(
+                engine_inputs, request_ids, strict=True
+            ):
                 self._engine.add_request(
-                    request_id=str(i),
+                    request_id=request_id,
                     prompt=engine_input,
                     params=sampling_params,
                 )
@@ -444,15 +458,14 @@ async def generate(
                 while self._engine.has_unfinished_requests():
                     all_outputs.extend(self._engine.step())
 
-            # vLLM may return requests out of order; sort by the integer
-            # request_id we assigned so prompt_idx lines up with the input.
-            all_outputs.sort(key=lambda o: int(o.request_id))
+            # Return completions in caller input order; the controller maps positionally.
+            request_order = {request_id: i for i, request_id in enumerate(request_ids)}
+            all_outputs.sort(key=lambda output: request_order[output.request_id])
 
             completions: list[Completion] = []
             generation_metrics: list[m.Metric] = []
             output_token_counts: list[int] = []
             for output in all_outputs:
-                prompt_idx = int(output.request_id)
                 generation_metrics.extend(
                     _prepare_generation_request_metrics(output, prefix=metrics_prefix)
                 )
@@ -465,8 +478,7 @@ async def generate(
                     completions.append(
                         Completion(
                             policy_version=self.policy_version,
-                            prompt_idx=prompt_idx,
-                            text=sample.text,
+                            request_id=output.request_id,
                             token_ids=sample.token_ids,
                             token_logprobs=per_token_logprobs,
                             finish_reason=sample.finish_reason,
diff --git a/torchtitan/experiments/rl/actors/trainer.py b/torchtitan/experiments/rl/actors/trainer.py
index 1cfd62448b..7f3fa6273e 100644
--- a/torchtitan/experiments/rl/actors/trainer.py
+++ b/torchtitan/experiments/rl/actors/trainer.py
@@ -456,7 +456,7 @@ async def forward_backward(
     async def optim_step(self) -> OptimStepOutput:
         """Clip gradients, step optimizer + LR scheduler, return updated state."""
         # TODO: Accept optional optimizer params (e.g. learning rate)
-        # to allow controller-owned schedules (see Tinker API).
+        # to allow controller-owned schedules.
 
         # capture LR before step
         current_lrs = self.lr_schedulers.schedulers[0].get_last_lr()
diff --git a/torchtitan/experiments/rl/batcher.py b/torchtitan/experiments/rl/batcher.py
index 6604e36ef8..227fd5e180 100644
--- a/torchtitan/experiments/rl/batcher.py
+++ b/torchtitan/experiments/rl/batcher.py
@@ -230,11 +230,11 @@ def _pack_episodes(self, episodes: list[Episode]) -> Iterator[dict]:
         def _iterate_samples() -> Iterator[dict]:
             for ep in episodes:
                 prompt_len = len(ep.prompt_token_ids)
-                response_len = len(ep.token_ids)
-                raw_ids = ep.prompt_token_ids + ep.token_ids
-                gen_lp = [0.0] * prompt_len + ep.token_logprobs
-                loss_mask = [False] * prompt_len + [True] * response_len
-                advantages = [0.0] * prompt_len + [ep.advantage] * response_len
+                completion_len = len(ep.completion_token_ids)
+                raw_ids = ep.prompt_token_ids + ep.completion_token_ids
+                gen_lp = [0.0] * prompt_len + ep.completion_logprobs
+                loss_mask = [False] * prompt_len + [True] * completion_len
+                advantages = [0.0] * prompt_len + [ep.advantage] * completion_len
                 sample = {
                     "input_ids": raw_ids[:-1],
                     "labels": raw_ids[1:],
diff --git a/torchtitan/experiments/rl/config_registry.py b/torchtitan/experiments/rl/config_registry.py
index e1f8827c6f..80cf998aad 100644
--- a/torchtitan/experiments/rl/config_registry.py
+++ b/torchtitan/experiments/rl/config_registry.py
@@ -26,9 +26,10 @@
 from torchtitan.experiments.rl.actors.generator import SamplingConfig, VLLMGenerator
 from torchtitan.experiments.rl.actors.trainer import PolicyTrainer
 from torchtitan.experiments.rl.batcher import BatchConfig, Batcher
+from torchtitan.experiments.rl.examples.sum_digits import SumDigitsRollouter
 from torchtitan.experiments.rl.grpo import GRPOLoss, RLTrainer
 from torchtitan.experiments.rl.observability.metrics import MetricsProcessor
-from torchtitan.experiments.rl.sum_digits import SumDigitsEnv
+from torchtitan.experiments.rl.renderer import RendererConfig
 from torchtitan.models.common.attention import FlexAttention
 from torchtitan.models.qwen3 import model_registry
 from torchtitan.protocols.model import ModelConfigConverter
@@ -74,13 +75,12 @@ def rl_grpo_qwen3_0_6b_varlen() -> RLTrainer.Config:
         model_spec=model_registry("0.6B", attn_backend="varlen"),
         hf_assets_path="torchtitan/experiments/rl/example_checkpoint/Qwen3-0.6B",
         num_steps=10,
-        num_prompts_per_step=5,
+        num_groups_per_rollout_batch=5,
         num_validation_samples=20,
         compile=CompileConfig(enable=True, backend="aot_eager"),
-        env=SumDigitsEnv.Config(seed=42, correctness_reward=1.0, format_reward=0.3),
-        validation_env=SumDigitsEnv.Config(
-            seed=99, correctness_reward=1.0, format_reward=0.3
-        ),
+        rollouter=SumDigitsRollouter.Config(),
+        group_size=group_size,
+        renderer=RendererConfig(name="qwen3", enable_thinking=True),
         metrics=MetricsProcessor.Config(enable_wandb=True),
         batcher=Batcher.Config(
             batch=BatchConfig(local_batch_size=2, global_batch_size=8, seq_len=2048),
@@ -115,10 +115,9 @@ def rl_grpo_qwen3_0_6b_varlen() -> RLTrainer.Config:
             ),
             checkpoint=CheckpointManager.Config(enable=False),
             sampling=SamplingConfig(
-                n=group_size,
                 temperature=0.8,
                 top_p=0.95,
-                max_tokens=100,
+                max_tokens=700,
             ),
         ),
     )
@@ -126,26 +125,24 @@ def rl_grpo_qwen3_0_6b_varlen() -> RLTrainer.Config:
 
 def rl_grpo_qwen3_0_6b_flex() -> RLTrainer.Config:
     """GRPO training config for Qwen3-0.6B with flex attention (4 GPUs: 2 gen + 2 train)."""
-    spec = model_registry("0.6B", attn_backend="flex")
     group_size = 8
     return RLTrainer.Config(
-        model_spec=spec,
+        model_spec=model_registry("0.6B", attn_backend="flex"),
         hf_assets_path="torchtitan/experiments/rl/example_checkpoint/Qwen3-0.6B",
         num_steps=10,
-        num_prompts_per_step=5,
+        num_groups_per_rollout_batch=5,
         num_validation_samples=20,
         # TODO: add aot_eager compiling overall, today it doesn't work because
         # we are missing mechanism to scoop Flex region to plug in inductor backend support
-        env=SumDigitsEnv.Config(seed=42, correctness_reward=1.0, format_reward=0.3),
-        validation_env=SumDigitsEnv.Config(
-            seed=99, correctness_reward=1.0, format_reward=0.3
-        ),
+        rollouter=SumDigitsRollouter.Config(),
+        group_size=group_size,
+        renderer=RendererConfig(name="qwen3", enable_thinking=True),
         metrics=MetricsProcessor.Config(enable_wandb=True),
         batcher=Batcher.Config(
             batch=BatchConfig(local_batch_size=2, global_batch_size=8, seq_len=2048),
         ),
         trainer=PolicyTrainer.Config(
-            optimizer=OptimizersContainer.Config(lr=2e-6),
+            optimizer=default_adamw(lr=2e-6),
             lr_scheduler=LRSchedulersContainer.Config(
                 warmup_steps=2,
                 decay_type="linear",
@@ -174,7 +171,6 @@ def rl_grpo_qwen3_0_6b_flex() -> RLTrainer.Config:
             ),
             checkpoint=CheckpointManager.Config(enable=False),
             sampling=SamplingConfig(
-                n=group_size,
                 temperature=0.8,
                 top_p=0.95,
                 max_tokens=100,
@@ -217,13 +213,12 @@ def rl_grpo_qwen3_1_7b() -> RLTrainer.Config:
         model_spec=model_registry("1.7B", attn_backend="varlen"),
         hf_assets_path="torchtitan/experiments/rl/example_checkpoint/Qwen3-1.7B",
         num_steps=10,
-        num_prompts_per_step=5,
+        num_groups_per_rollout_batch=5,
         num_validation_samples=20,
         compile=CompileConfig(enable=True, backend="aot_eager"),
-        env=SumDigitsEnv.Config(seed=42, correctness_reward=1.0, format_reward=0.3),
-        validation_env=SumDigitsEnv.Config(
-            seed=99, correctness_reward=1.0, format_reward=0.3
-        ),
+        rollouter=SumDigitsRollouter.Config(),
+        group_size=group_size,
+        renderer=RendererConfig(name="qwen3", enable_thinking=True),
         metrics=MetricsProcessor.Config(enable_wandb=True),
         batcher=Batcher.Config(
             batch=BatchConfig(local_batch_size=2, global_batch_size=8, seq_len=2048),
@@ -259,10 +254,9 @@ def rl_grpo_qwen3_1_7b() -> RLTrainer.Config:
             ),
             checkpoint=CheckpointManager.Config(enable=False),
             sampling=SamplingConfig(
-                n=group_size,
                 temperature=0.8,
                 top_p=0.95,
-                max_tokens=100,
+                max_tokens=700,
             ),
         ),
     )
@@ -275,13 +269,12 @@ def rl_grpo_qwen3_14b() -> RLTrainer.Config:
         model_spec=model_registry("14B", attn_backend="varlen"),
         hf_assets_path="torchtitan/experiments/rl/example_checkpoint/Qwen3-14B",
         num_steps=10,
-        num_prompts_per_step=5,
+        num_groups_per_rollout_batch=5,
         num_validation_samples=20,
         compile=CompileConfig(enable=True, backend="aot_eager"),
-        env=SumDigitsEnv.Config(seed=42, correctness_reward=1.0, format_reward=0.3),
-        validation_env=SumDigitsEnv.Config(
-            seed=99, correctness_reward=1.0, format_reward=0.3
-        ),
+        rollouter=SumDigitsRollouter.Config(),
+        group_size=group_size,
+        renderer=RendererConfig(name="qwen3", enable_thinking=True),
         metrics=MetricsProcessor.Config(enable_wandb=True),
         batcher=Batcher.Config(
             batch=BatchConfig(local_batch_size=2, global_batch_size=8, seq_len=2048),
@@ -316,17 +309,16 @@ def rl_grpo_qwen3_14b() -> RLTrainer.Config:
             ),
             checkpoint=CheckpointManager.Config(enable=False),
             sampling=SamplingConfig(
-                n=group_size,
                 temperature=0.8,
                 top_p=0.95,
-                max_tokens=100,
+                max_tokens=700,
             ),
         ),
     )
 
 
 def rl_grpo_qwen3_0_6b_batch_invariant() -> RLTrainer.Config:
-    """On-policy GRPO config for Qwen3-0.6B under same parallelism (4 GPUs: 2 gen + 2 train).
+    """On-policy GRPO config for Qwen3-0.6B (4 GPUs: 2 gen + 2 train).
 
     Enables deterministic + batch-invariant mode for true on-policy RL training.
     """
@@ -336,13 +328,12 @@ def rl_grpo_qwen3_0_6b_batch_invariant() -> RLTrainer.Config:
         model_spec=model_registry("0.6B", attn_backend="varlen"),
         hf_assets_path="torchtitan/experiments/rl/example_checkpoint/Qwen3-0.6B",
         num_steps=10,
-        num_prompts_per_step=5,
+        num_groups_per_rollout_batch=5,
         num_validation_samples=20,
         compile=CompileConfig(enable=True, backend="aot_eager"),
-        env=SumDigitsEnv.Config(seed=42, correctness_reward=1.0, format_reward=0.3),
-        validation_env=SumDigitsEnv.Config(
-            seed=99, correctness_reward=1.0, format_reward=0.3
-        ),
+        rollouter=SumDigitsRollouter.Config(),
+        group_size=group_size,
+        renderer=RendererConfig(name="qwen3", enable_thinking=True),
         metrics=MetricsProcessor.Config(enable_wandb=True),
         batcher=Batcher.Config(
             batch=BatchConfig(local_batch_size=2, global_batch_size=8, seq_len=2048),
@@ -381,10 +372,9 @@ def rl_grpo_qwen3_0_6b_batch_invariant() -> RLTrainer.Config:
             ),
             checkpoint=CheckpointManager.Config(enable=False),
             sampling=SamplingConfig(
-                n=group_size,
                 temperature=0.8,
                 top_p=0.95,
-                max_tokens=100,
+                max_tokens=700,
             ),
             debug=batch_invariant_config,
         ),
diff --git a/torchtitan/experiments/rl/environment/__init__.py b/torchtitan/experiments/rl/environment/__init__.py
new file mode 100644
index 0000000000..c91acc1a1c
--- /dev/null
+++ b/torchtitan/experiments/rl/environment/__init__.py
@@ -0,0 +1,20 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+from torchtitan.experiments.rl.environment.message import (
+    MessageEnv,
+    MessageEnvInitOutput,
+    MessageEnvStepOutput,
+)
+from torchtitan.experiments.rl.environment.token import TokenEnv, TokenEnvOutput
+
+__all__ = [
+    "MessageEnv",
+    "MessageEnvInitOutput",
+    "MessageEnvStepOutput",
+    "TokenEnv",
+    "TokenEnvOutput",
+]
diff --git a/torchtitan/experiments/rl/environment/message.py b/torchtitan/experiments/rl/environment/message.py
new file mode 100644
index 0000000000..2ee4f406f4
--- /dev/null
+++ b/torchtitan/experiments/rl/environment/message.py
@@ -0,0 +1,108 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+from __future__ import annotations
+
+import abc
+from dataclasses import dataclass, field
+
+from renderers import Message, ToolSpec
+
+from torchtitan.config import Configurable
+
+
+@dataclass(kw_only=True, slots=True)
+class MessageEnvInitOutput:
+    """Initial messages + tool specs from `MessageEnv.init`."""
+
+    init_prompt_messages: list[Message]  # [num_prompt_messages]
+    """The opening messages (e.g. [system, user]); may include few-shot assistant turns."""
+
+    tools: list[ToolSpec] = field(default_factory=list)  # [num_tools]
+    """Tool schemas exposed to the assistant. Empty for tool-less envs."""
+
+
+@dataclass(kw_only=True, slots=True)
+class MessageEnvStepOutput:
+    """The env's reply to the assistant's turn."""
+
+    env_messages: list[Message] = field(default_factory=list)  # [num_env_messages]
+    """The env's reply messages (tool / user). Empty when the rollout terminates
+    with no follow-up."""
+
+    done: bool = False
+    """`True` ends the rollout."""
+
+    env_rewards: dict[str, float] = field(default_factory=dict)
+    """Optional reward signal the env provides for this step; the rubric decides
+    whether and how to use it. Empty if the env scores nothing."""
+
+    def __post_init__(self) -> None:
+        # env replies are tool/user turns; the assistant turn comes from the generator
+        if any(m.get("role") == "assistant" for m in self.env_messages):
+            raise ValueError(
+                "MessageEnvStepOutput.env_messages may not contain assistant messages"
+            )
+
+
+class MessageEnv(Configurable, abc.ABC):
+    """User-written env in message space. Implement `init` + `step`.
+
+    Tip: `MessageEnv` works in messages and never sees token ids; You can have `TokenEnv`
+    wrap it and use a `Renderer` to convert messages <-> token ids for the generator.
+    Example:
+        # a one-tool calculator env. It is multi-turn — the env answers the
+        # assistant's tool call, then ends once the assistant replies without a tool.
+
+        class CalculatorEnv(MessageEnv):
+            @dataclass(kw_only=True, slots=True)
+            class Config(MessageEnv.Config):
+                pass
+
+            def __init__(self, config: Config, *, env_input: CalculatorExample) -> None:
+                self._expression = env_input.expression
+
+            async def init(self) -> MessageEnvInitOutput:
+                return MessageEnvInitOutput(
+                    init_prompt_messages=[{"role": "user", "content": f"What is {self._expression}?"}],
+                    tools=[CALCULATOR_TOOL],
+                )
+
+            async def step(self, completion_message: Message) -> MessageEnvStepOutput:
+                tool_calls = completion_message.get("tool_calls")
+                if not tool_calls:
+                    return MessageEnvStepOutput(done=True)  # assistant gave its final answer
+                result = run_calculator(tool_calls[0])
+                return MessageEnvStepOutput(
+                    env_messages=[{"role": "tool", "content": result}]
+                )
+    """
+
+    @dataclass(kw_only=True, slots=True)
+    class Config(Configurable.Config):
+        """Static env params; the per-rollout example is passed to `build(env_input=...)`."""
+
+    @abc.abstractmethod
+    async def init(self) -> MessageEnvInitOutput:
+        """Return the initial conversation + tools for prompt rendering."""
+
+    @abc.abstractmethod
+    async def step(self, completion_message: Message) -> MessageEnvStepOutput:
+        """Advance the env one turn given the completion message.
+
+        `TokenEnv` parses the completion and handles
+        finish_reason / length / parse / timeout failures before calling this,
+        so the env only sees a well-formed completion message.
+
+        Args:
+            completion_message: the completion decoded into a message.
+
+        Returns:
+            `MessageEnvStepOutput` with the env's reply messages.
+        """
+
+    async def close(self) -> None:
+        """Release env-owned resources. Default no-op; idempotent."""
diff --git a/torchtitan/experiments/rl/environment/token.py b/torchtitan/experiments/rl/environment/token.py
new file mode 100644
index 0000000000..1a2ad9f4de
--- /dev/null
+++ b/torchtitan/experiments/rl/environment/token.py
@@ -0,0 +1,282 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+from __future__ import annotations
+
+import asyncio
+import logging
+from dataclasses import dataclass, field
+
+from renderers import Message, Renderer, ToolSpec
+
+from torchtitan.config import Configurable
+from torchtitan.experiments.rl.environment.message import MessageEnv
+from torchtitan.experiments.rl.rollout.types import RolloutStatus
+from torchtitan.experiments.rl.types import Completion
+
+logger = logging.getLogger(__name__)
+
+
+@dataclass(kw_only=True, slots=True)
+class TokenEnvOutput:
+    """What `TokenEnv` produced this turn. Output only — it does not carry the prompt that was fed
+    in or generator metadata (logprobs, policy_version). The rollout loop folds it together
+    with those inputs into a durable `RolloutTurn`."""
+
+    next_prompt_token_ids: list[int] | None  # [num_prompt_tokens] or None
+    """Tokens for the next `generate` call. `None` once no further prompt is
+    rendered (COMPLETED, or a completion-level error)."""
+
+    next_prompt_messages: list[Message] | None = None  # [num_prompt_messages] or None
+    """Message form of `next_prompt_token_ids`; `None` whenever that is `None`."""
+
+    status: RolloutStatus
+    """`ONGOING` while the rollout runs; a terminal `RolloutStatus` otherwise."""
+
+    completion_message: Message | None = None
+    """This turn's completion decoded from the generator's tokens into a message by the renderer
+    (this is the TokenEnv's parse, not the raw generator output). `None` on init and on parse failure."""
+
+    env_messages: list[Message] = field(default_factory=list)  # [num_env_messages]
+    """The env's reply on this step (tool / user). Empty on init and on terminals
+    where the env was not stepped."""
+
+    env_rewards: dict[str, float] = field(default_factory=dict)
+    """Per-step reward signals from the env; empty when the env was not stepped. This
+    can be used by the rubric to compute the final rollout reward."""
+
+
+class TokenEnv(Configurable):
+    """Token-space env used to wrap a `MessageEnv` and drive it in tokens.
+
+    In a rollout, the input and output of a generator are in tokens (Tokens-In-Tokens-Out).
+    However, the MessageEnv is in messages (Messages-In-Messages-Out).
+
+    Some process is necessary to:
+    a. Decode the generator's completion tokens into messages;
+    b. Call the MessageEnv;
+    c. Encode the env response back into tokens for the next turn;
+
+    This wrapper fills this role, using a renderer to convert between the generator and MessageEnv.
+
+    Beyond encoding/decoding, several checks are necessary, e.g. prompt too long, too many
+    turns, parse errors, timeouts. This wrapper also takes care of that, so we can keep the
+    rollout loop clean and simple.
+
+    If users have extra or different logic, they can wrap their MessageEnv with another class instead.
+
+    Args:
+        config: `TokenEnv.Config`.
+        message_env: the user's `MessageEnv` subclass instance.
+        renderer: a `renderers.Renderer` that converts messages <-> token ids.
+
+    Example:
+
+        env = TokenEnv.Config().build(message_env=SumDigitsEnv(...), renderer=renderer)
+        env_output = await env.init()
+        while not env_output.status.is_terminal():
+            completion = await generator.generate([env_output.next_prompt_token_ids])
+            env_output = await env.step(completion)
+    """
+
+    @dataclass(kw_only=True, slots=True)
+    class Config(Configurable.Config):
+        """Limits enforced by the wrapper"""
+
+        max_rollout_tokens: int | None = None
+        """Hard cap on prompt length for the next turn. If the number of tokens meets/exceeds
+        it, the turn is terminal; `None` disables the check."""
+
+        # TODO: its unclear if timeout should be on this layer or handled by the messageEnv
+        step_timeout_s: float | None = 1800.0
+        """Wall-clock timeout for one `MessageEnv.step` call."""
+
+        # TODO: add max_num_turns
+
+    def __init__(
+        self,
+        config: Config,
+        *,
+        message_env: MessageEnv,
+        renderer: Renderer,
+    ) -> None:
+        self._message_env = message_env
+        self._renderer = renderer
+        self._config = config
+        self._tools: list[ToolSpec] | None = None
+        self._messages: list[Message] = []
+        self._last_prompt_token_ids: list[int] = []
+
+    async def init(self) -> TokenEnvOutput:
+        """Render the initial conversation into the first generator prompt."""
+        env_init_output = await self._message_env.init()
+
+        # Copy our running conversation (messages) so we avoid mutating previous states
+        self._messages = list(env_init_output.init_prompt_messages)
+
+        # Render messages into tokens
+        self._tools = env_init_output.tools if env_init_output.tools else None
+        token_ids = await asyncio.to_thread(
+            self._renderer.render_ids,
+            messages=self._messages,
+            tools=self._tools,
+            add_generation_prompt=True,
+        )
+
+        # Carry the prompt either way: ONGOING if it fits, else
+        # TRUNCATED_PROMPT_TOO_LONG so the over-budget prompt stays debuggable.
+        prompt_too_long = self._is_prompt_too_long(prompt_len=len(token_ids))
+        if not prompt_too_long:
+            self._last_prompt_token_ids = token_ids
+        return TokenEnvOutput(
+            next_prompt_token_ids=token_ids,
+            next_prompt_messages=self._messages,
+            status=(
+                RolloutStatus.TRUNCATED_PROMPT_TOO_LONG
+                if prompt_too_long
+                else RolloutStatus.ONGOING
+            ),
+        )
+
+    async def step(self, completion: Completion) -> TokenEnvOutput:
+        """Advance the env by one sampled completion from the generator.
+
+        Args:
+            completion: Generator output for the current prompt.
+
+        Returns:
+            `TokenEnvOutput` for the next generator call, or a terminal turn
+            when the rollout completes, truncates, or errors.
+        """
+        # Parse first, so a truncated / aborted response still carries its message
+        try:
+            parsed = await asyncio.to_thread(
+                self._renderer.parse_response,
+                token_ids=completion.token_ids,
+            )
+        except Exception:
+            logger.exception(
+                "parse_response failed (finish_reason=%s, %d tokens); -> ERROR_PARSE",
+                completion.finish_reason,
+                len(completion.token_ids),
+            )
+            return TokenEnvOutput(
+                next_prompt_token_ids=None,
+                next_prompt_messages=None,
+                status=RolloutStatus.ERROR_PARSE,
+            )
+
+        completion_message: Message = {
+            "role": "assistant",
+            "content": parsed.content,
+        }
+        if parsed.reasoning_content:
+            completion_message["reasoning_content"] = parsed.reasoning_content
+        if parsed.tool_calls:
+            completion_message["tool_calls"] = parsed.tool_calls
+
+        # Truncated / aborted: the response is final and partial. Keep it for
+        # partial-reward grading and debugging; don't step the env on it.
+        # TODO: check if we should step the env on an incomplete message
+        if completion.finish_reason == "length":
+            return TokenEnvOutput(
+                next_prompt_token_ids=None,
+                next_prompt_messages=None,
+                status=RolloutStatus.TRUNCATED_LENGTH,
+                completion_message=completion_message,
+            )
+        if completion.finish_reason == "abort":
+            return TokenEnvOutput(
+                next_prompt_token_ids=None,
+                next_prompt_messages=None,
+                status=RolloutStatus.ERROR_ABORT,
+                completion_message=completion_message,
+            )
+
+        # Apply the user's env step under a timeout
+        timeout = self._config.step_timeout_s
+        try:
+            if timeout is None:
+                step_output = await self._message_env.step(completion_message)
+            else:
+                step_output = await asyncio.wait_for(
+                    self._message_env.step(completion_message), timeout=timeout
+                )
+        except TimeoutError:
+            logger.warning("step timed out after %ss; -> ERROR_TIMEOUT", timeout)
+            return TokenEnvOutput(
+                next_prompt_token_ids=None,
+                next_prompt_messages=None,
+                status=RolloutStatus.ERROR_TIMEOUT,
+                completion_message=completion_message,
+            )
+
+        # TODO(history-edit): We hard-code the logic to only append new messages.
+        # This may not satisfy all uses cases, such as compacting the history.
+        # Update this when such cases arise.
+
+        # Create a new list to avoid mutating previous states
+        self._messages = (
+            self._messages + [completion_message] + step_output.env_messages
+        )
+
+        if step_output.done:
+            return TokenEnvOutput(
+                next_prompt_token_ids=None,
+                next_prompt_messages=None,
+                status=RolloutStatus.COMPLETED,
+                completion_message=completion_message,
+                env_messages=step_output.env_messages,
+                env_rewards=step_output.env_rewards,
+            )
+
+        # Prepare the next prompt; full re-render if the renderer can't bridge.
+        # `tools` is passed because tool schemas are part of the chat template, so
+        # the bridged tokens must match what a full re-render (also tools-aware) produces.
+        bridged = await asyncio.to_thread(
+            self._renderer.bridge_to_next_turn,
+            previous_prompt_ids=self._last_prompt_token_ids,
+            previous_completion_ids=completion.token_ids,
+            new_messages=step_output.env_messages,
+            tools=self._tools,
+        )
+        if bridged is None:
+            next_prompt_token_ids = await asyncio.to_thread(
+                self._renderer.render_ids,
+                messages=self._messages,
+                tools=self._tools,
+                add_generation_prompt=True,
+            )
+        else:
+            next_prompt_token_ids = bridged.token_ids
+
+        # Terminal if the next prompt is over budget
+        if self._is_prompt_too_long(prompt_len=len(next_prompt_token_ids)):
+            return TokenEnvOutput(
+                next_prompt_token_ids=None,
+                next_prompt_messages=None,
+                status=RolloutStatus.TRUNCATED_PROMPT_TOO_LONG,
+                completion_message=completion_message,
+                env_messages=step_output.env_messages,
+                env_rewards=step_output.env_rewards,
+            )
+
+        self._last_prompt_token_ids = next_prompt_token_ids
+        return TokenEnvOutput(
+            next_prompt_token_ids=next_prompt_token_ids,
+            next_prompt_messages=self._messages,
+            status=RolloutStatus.ONGOING,
+            completion_message=completion_message,
+            env_messages=step_output.env_messages,
+            env_rewards=step_output.env_rewards,
+        )
+
+    async def close(self) -> None:
+        await self._message_env.close()
+
+    def _is_prompt_too_long(self, *, prompt_len: int) -> bool:
+        cap = self._config.max_rollout_tokens
+        return cap is not None and prompt_len >= cap
diff --git a/torchtitan/experiments/rl/examples/__init__.py b/torchtitan/experiments/rl/examples/__init__.py
new file mode 100644
index 0000000000..2e41cd717f
--- /dev/null
+++ b/torchtitan/experiments/rl/examples/__init__.py
@@ -0,0 +1,5 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
diff --git a/torchtitan/experiments/rl/examples/sum_digits/__init__.py b/torchtitan/experiments/rl/examples/sum_digits/__init__.py
new file mode 100644
index 0000000000..7206f0923c
--- /dev/null
+++ b/torchtitan/experiments/rl/examples/sum_digits/__init__.py
@@ -0,0 +1,25 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+from torchtitan.experiments.rl.examples.sum_digits.data import (
+    SumDigitsDataset,
+    SumDigitsExample,
+)
+from torchtitan.experiments.rl.examples.sum_digits.env import SumDigitsEnv
+from torchtitan.experiments.rl.examples.sum_digits.rollouter import SumDigitsRollouter
+from torchtitan.experiments.rl.examples.sum_digits.rubric import (
+    RewardCorrect,
+    RewardFormat,
+)
+
+__all__ = [
+    "RewardCorrect",
+    "RewardFormat",
+    "SumDigitsDataset",
+    "SumDigitsEnv",
+    "SumDigitsExample",
+    "SumDigitsRollouter",
+]
diff --git a/torchtitan/experiments/rl/examples/sum_digits/data.py b/torchtitan/experiments/rl/examples/sum_digits/data.py
new file mode 100644
index 0000000000..b267ed9f44
--- /dev/null
+++ b/torchtitan/experiments/rl/examples/sum_digits/data.py
@@ -0,0 +1,57 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+from __future__ import annotations
+
+import random
+from collections.abc import Iterator
+from dataclasses import dataclass
+
+from torchtitan.config import Configurable
+
+
+@dataclass(frozen=True, kw_only=True, slots=True)
+class SumDigitsExample:
+    """Typed payload for one SumDigits problem."""
+
+    numbers: list[int]  # [N_numbers]
+    """Numbers the model must digit-sum."""
+
+    target: int
+    """Ground-truth total digit sum."""
+
+
+class SumDigitsDataset(Configurable):
+    """Endless, seeded stream of SumDigits problems.
+
+    Example:
+
+        dataset = SumDigitsDataset(SumDigitsDataset.Config(seed=42))
+        example: SumDigitsExample = next(iter(dataset))
+    """
+
+    @dataclass(kw_only=True, slots=True)
+    class Config(Configurable.Config):
+        seed: int = 42
+
+    def __init__(self, config: Config) -> None:
+        self._rng = random.Random(config.seed)
+
+    def __iter__(self) -> Iterator[SumDigitsExample]:
+        return self
+
+    def __next__(self) -> SumDigitsExample:
+        n = self._rng.randint(2, 4)
+        numbers = [self._rng.randint(10, 99) for _ in range(n)]
+        target = sum(int(d) for num in numbers for d in str(num))
+        return SumDigitsExample(numbers=numbers, target=target)
+
+    def state_dict(self) -> dict:
+        """Snapshot the RNG so a run can resume at the same point in the stream."""
+        return {"rng_state": self._rng.getstate()}
+
+    def load_state_dict(self, state_dict: dict) -> None:
+        self._rng.setstate(state_dict["rng_state"])
diff --git a/torchtitan/experiments/rl/examples/sum_digits/env.py b/torchtitan/experiments/rl/examples/sum_digits/env.py
new file mode 100644
index 0000000000..1add5d5f7e
--- /dev/null
+++ b/torchtitan/experiments/rl/examples/sum_digits/env.py
@@ -0,0 +1,55 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+
+from renderers import Message
+
+from torchtitan.experiments.rl.environment import (
+    MessageEnv,
+    MessageEnvInitOutput,
+    MessageEnvStepOutput,
+)
+from torchtitan.experiments.rl.examples.sum_digits.data import SumDigitsExample
+
+
+SYSTEM_PROMPT = """\
+You are a helpful assistant. Solve the problem step by step.
+When you have your final answer, state it as [ANSWER] <number>.
+
+Example:
+User: What is the total digit sum of [12, 345, 67]?
+Assistant: Break each number into digits:
+12 -> 1, 2
+345 -> 3, 4, 5
+67 -> 6, 7
+Sum all digits: 1 + 2 + 3 + 4 + 5 + 6 + 7 = 28
+[ANSWER] 28"""
+
+
+class SumDigitsEnv(MessageEnv):
+    @dataclass(kw_only=True, slots=True)
+    class Config(MessageEnv.Config):
+        pass
+
+    def __init__(self, config: Config, *, env_input: SumDigitsExample) -> None:
+        self._numbers = env_input.numbers
+
+    async def init(self) -> MessageEnvInitOutput:
+        """Return the system prompt and one SumDigits user question."""
+        question = f"What is the total digit sum of {self._numbers}?"
+        return MessageEnvInitOutput(
+            init_prompt_messages=[
+                {"role": "system", "content": SYSTEM_PROMPT},
+                {"role": "user", "content": question},
+            ]
+        )
+
+    async def step(self, completion_message: Message) -> MessageEnvStepOutput:
+        # Single-turn env: end after the first completion.
+        return MessageEnvStepOutput(done=True)
diff --git a/torchtitan/experiments/rl/examples/sum_digits/rollouter.py b/torchtitan/experiments/rl/examples/sum_digits/rollouter.py
new file mode 100644
index 0000000000..abba9daa1d
--- /dev/null
+++ b/torchtitan/experiments/rl/examples/sum_digits/rollouter.py
@@ -0,0 +1,45 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+
+from torchtitan.experiments.rl.examples.sum_digits.data import SumDigitsDataset
+from torchtitan.experiments.rl.examples.sum_digits.env import SumDigitsEnv
+from torchtitan.experiments.rl.examples.sum_digits.rubric import (
+    RewardCorrect,
+    RewardFormat,
+)
+
+from torchtitan.experiments.rl.rollout.rollouter import Rollouter
+
+from torchtitan.experiments.rl.rubrics import Rubric
+
+
+class SumDigitsRollouter(Rollouter):
+    """The SumDigits rollouter: digit-sum train/val datasets, env, and a
+    correctness + format rubric. Pure config — all behavior (`make_env_group`,
+    `sample_*`, `score_group`) is inherited from `Rollouter`.
+    """
+
+    @dataclass(kw_only=True, slots=True)
+    class Config(Rollouter.Config):
+        train_dataset: SumDigitsDataset.Config = field(
+            default_factory=lambda: SumDigitsDataset.Config(seed=42)
+        )
+        validation_dataset: SumDigitsDataset.Config = field(
+            default_factory=lambda: SumDigitsDataset.Config(seed=99)
+        )
+        rubric: Rubric.Config = field(
+            default_factory=lambda: Rubric.Config(
+                reward_fns=[
+                    RewardCorrect.Config(weight=1.0),
+                    RewardFormat.Config(weight=0.3),
+                ]
+            )
+        )
+        message_env: SumDigitsEnv.Config = field(default_factory=SumDigitsEnv.Config)
diff --git a/torchtitan/experiments/rl/examples/sum_digits/rubric.py b/torchtitan/experiments/rl/examples/sum_digits/rubric.py
new file mode 100644
index 0000000000..bf1270456d
--- /dev/null
+++ b/torchtitan/experiments/rl/examples/sum_digits/rubric.py
@@ -0,0 +1,48 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+from __future__ import annotations
+
+import re
+from dataclasses import dataclass
+
+from torchtitan.experiments.rl.examples.sum_digits.data import SumDigitsExample
+
+from torchtitan.experiments.rl.rollout import last_completion_text, Rollout
+from torchtitan.experiments.rl.rubrics import RewardFn
+
+
+_ANSWER_RE = re.compile(r"\[ANSWER\]\s*(-?\d+)")
+_FORMAT_RE = re.compile(r"\[ANSWER\]\s*-?\d+")
+
+
+class RewardCorrect(RewardFn):
+    """`1.0` if the last `[ANSWER] <n>` equals the target, else `0.0`."""
+
+    @dataclass(kw_only=True, slots=True)
+    class Config(RewardFn.Config):
+        pass
+
+    async def __call__(self, rollout: Rollout, env_input: SumDigitsExample) -> float:
+        text = last_completion_text(rollout)
+        matches = _ANSWER_RE.findall(text)
+        if not matches:
+            return 0.0
+        return 1.0 if int(matches[-1]) == env_input.target else 0.0
+
+
+class RewardFormat(RewardFn):
+    """`1.0` if the response contains `[ANSWER] <n>`, else `0.0`."""
+
+    @dataclass(kw_only=True, slots=True)
+    class Config(RewardFn.Config):
+        pass
+
+    async def __call__(self, rollout: Rollout, env_input: object) -> float:
+        return 1.0 if _FORMAT_RE.search(last_completion_text(rollout)) else 0.0
+
+
+__all__ = ["RewardCorrect", "RewardFormat"]
diff --git a/torchtitan/experiments/rl/generate.py b/torchtitan/experiments/rl/generate.py
index aaef7e9553..9f0d3f577c 100755
--- a/torchtitan/experiments/rl/generate.py
+++ b/torchtitan/experiments/rl/generate.py
@@ -79,7 +79,7 @@ def generate():
         hf_overrides={"architectures": [VLLM_MODEL_NAME]},
         attention_backend="CUSTOM",
     )
-    max_num_seqs = config.num_prompts_per_step * gen_config.sampling.n
+    max_num_seqs = config.num_groups_per_rollout_batch * config.group_size
     engine_kwargs["max_num_seqs"] = max_num_seqs
     vllm_compilation_config = gen_config.cudagraph.get_vllm_compilation_config(
         max_num_seqs=max_num_seqs,
diff --git a/torchtitan/experiments/rl/grpo.py b/torchtitan/experiments/rl/grpo.py
index dd8fb243ff..8c2394cd2e 100644
--- a/torchtitan/experiments/rl/grpo.py
+++ b/torchtitan/experiments/rl/grpo.py
@@ -27,9 +27,9 @@
 import os
 import statistics
 import time
-from collections import defaultdict
 from collections.abc import Callable
-from dataclasses import dataclass, field
+from concurrent.futures import ThreadPoolExecutor
+from dataclasses import dataclass, field, replace
 
 # must run before torch import
 os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True")
@@ -39,18 +39,33 @@
 from monarch.actor import this_host
 from monarch.spmd import setup_torch_elastic_env_async
 
-from torchtitan.components.tokenizer import HuggingFaceTokenizer
 from torchtitan.config import (
     CompileConfig,
     ConfigManager,
     Configurable,
     ParallelismConfig,
 )
-from torchtitan.experiments.rl.actors.generator import SamplingConfig, VLLMGenerator
+from torchtitan.experiments.rl.actors.generator import (
+    Completion,
+    SamplingConfig,
+    VLLMGenerator,
+)
 from torchtitan.experiments.rl.actors.trainer import PolicyTrainer
 from torchtitan.experiments.rl.batcher import Batcher
+from torchtitan.experiments.rl.environment import TokenEnv, TokenEnvOutput
 from torchtitan.experiments.rl.observability import metrics as m
-from torchtitan.experiments.rl.types import Completion, Episode, Trajectory
+from torchtitan.experiments.rl.renderer import RendererConfig
+from torchtitan.experiments.rl.rollout import (
+    last_completion_text,
+    prepare_rollout_metrics,
+    Rollout,
+    rollout_to_episode,
+    RolloutGroup,
+    RolloutStatus,
+    RolloutTurn,
+)
+from torchtitan.experiments.rl.rollout.rollouter import Rollouter
+from torchtitan.experiments.rl.types import Episode
 from torchtitan.observability import structured_logger as sl
 from torchtitan.protocols.model_spec import ModelSpec
 
@@ -167,57 +182,56 @@ def _bootstrap():
         return _bootstrap
 
 
-def _log_samples(items: list[Episode] | list[Completion]) -> None:
-    """Log the first sample per prompt for debugging."""
-    seen_prompts: set[int] = set()
-    for item in items:
-        if item.prompt_idx in seen_prompts:
+def _log_samples(rollout_groups: list[RolloutGroup]) -> None:
+    """Log the first scored, trainable rollout per group for debugging."""
+    for group in rollout_groups:
+        rollout = next(
+            (
+                r
+                for r in group.rollouts
+                if r.reward is not None and r.turns and r.turns[0].completion_token_ids
+            ),
+            None,
+        )
+        if rollout is None:
             continue
-        seen_prompts.add(item.prompt_idx)
-        reward_str = f" reward={item.reward:+.1f}" if hasattr(item, "reward") else ""
-        logger.info(f"  [prompt {item.prompt_idx}]{reward_str}")
-        logger.info(f"       A: {item.text[:300].replace(chr(10), ' ').strip()}")
+        logger.info("  [%s] reward=%+.1f", group.group_id, rollout.reward)
+        logger.info(
+            "       A: %s",
+            last_completion_text(rollout)[:300].replace("\n", " ").strip(),
+        )
 
 
-def _prepare_reward_metrics(
-    prefix: str,
-    trajectories: list[Trajectory],
-) -> list[m.Metric]:
-    """One ``Mean`` metric per observed reward component across trajectories.
+def _sample_id(group_id: str, sample_idx: int) -> str:
+    return f"{group_id}/sample={sample_idx}"
 
-    Example::
 
-        trajectories = [
-            Trajectory(
-                sample_idx=0,
-                prompt_token_ids=p0,
-                transitions=[(c0, Step(rewards={"correctness": 1.0, "format": 0.5}, done=True))],
-            ),
-            Trajectory(
-                sample_idx=1,
-                prompt_token_ids=p1,
-                transitions=[(c1, Step(rewards={"correctness": 0.0}, done=True))],
-            ),
-        ]
-        _prepare_reward_metrics("reward/component", trajectories)
-        # -> [
-        #      Metric("reward/component/correctness", Mean(sum=1.0, count=2)),  # 0.5
-        #      Metric("reward/component/format",      Mean(sum=0.5, count=1)),  # 0.5 - "format" only in trajectory 0
-        #    ]
-    """
-    values_by_name: dict[str, list[float]] = defaultdict(list)
-    for trajectory in trajectories:
-        for _completion, step in trajectory.transitions:
-            for name, value in step.rewards.items():
-                values_by_name[name].append(float(value))
-    return [
-        m.Metric(f"{prefix}/{name}", m.Mean.from_list(values))
-        for name, values in sorted(values_by_name.items())
-    ]
+@dataclass(kw_only=True, slots=True)
+class _RolloutGroupState:
+    """A prompt group's working state across one rollout collection call."""
+
+    group_id: str
+    sample: object
+    envs: list[TokenEnv]  # [group_size]
+    env_init_outputs: list[TokenEnvOutput] = field(default_factory=list)
+    completions: list[Completion] | None = None
 
 
 class RLTrainer(Configurable):
-    """Top-level RL training orchestrator."""
+    """Top-level RL training orchestrator.
+
+    Owns a `PolicyTrainer` actor (gradient updates), a `VLLMGenerator` actor
+    (sampling), and a `Rollouter` (datasets + rubric + env construction). Each
+    training step samples groups of rollouts, scores them via the rollouter's rubric,
+    builds GRPO advantages, and syncs trainer weights to the generator.
+
+    Example:
+
+        cfg = config_registry.rl_grpo_qwen3_0_6b()
+        trainer = cfg.build()
+        await trainer.setup_async()
+        await trainer.train()
+    """
 
     @dataclass(kw_only=True, slots=True)
     class Config(Configurable.Config):
@@ -236,21 +250,25 @@ class Config(Configurable.Config):
         dump_folder: str = "outputs/rl"
         """Root output folder for RL artifacts (temp weights, logs, etc.)."""
 
-        num_prompts_per_step: int = 5
-        """Number of distinct prompts (= GRPO groups) drawn per training step.
+        num_groups_per_rollout_batch: int = 5
+        """GRPO groups collected per rollout batch; a train step may collect several batches
+        until the token target is met. Rollouts per batch = `num_groups_per_rollout_batch * group_size`."""
+        # TODO(continuous-batching): this knob exists because we collect to a token budget
+        # in discrete sync batches; async/continuous batching streams may change this logic
 
-        The total episodes per step is `num_prompts_per_step` * `group_size`,
-        where `group_size` is `generator.sampling.n` (completions per prompt).
-        """
+        group_size: int = 8
+        """Sibling rollouts sampled per dataset row (the GRPO group). The generator
+        is always called with `n=1`; prompts are pre-expanded by `group_size`."""
 
         num_validation_samples: int = 20
         """Number of held-out prompts scored greedily (temp=0, n=1) per validation pass."""
 
-        env: Configurable.Config = field(default=None)  # type: ignore[assignment]
-        """Env config for training rollouts."""
+        rollouter: Rollouter.Config
+        """The rollouter: its datasets, envs, and rubric."""
+        # TODO: support multiple rollouters for data mixing.
 
-        validation_env: Configurable.Config = field(default=None)  # type: ignore[assignment]
-        """Env config for validation rollouts."""
+        renderer: RendererConfig
+        """Message-to-token renderer config."""
 
         log_samples: bool = False
         """Log first completion per episode during training and validation."""
@@ -280,7 +298,6 @@ def __post_init__(self):
                     "(weights are synced from the trainer via TorchStore). "
                     "Set generator.checkpoint.enable=False."
                 )
-
             if self.trainer.debug.batch_invariant:
                 if not self.trainer.debug.deterministic:
                     raise ValueError("batch_invariant requires deterministic=True")
@@ -312,10 +329,19 @@ def __init__(self, config: Config):
             log_dir=config.dump_folder,
             job_config=config.to_dict(),
         )
-        # TODO: Replace this single-turn tokenizer with renderer
-        self.tokenizer = HuggingFaceTokenizer(tokenizer_path=config.hf_assets_path)
-        # TODO: Use tokenizer.pad_id when available, falling back to eos_id.
-        self.batcher = Batcher(config.batcher, pad_id=self.tokenizer.eos_id)
+        self.renderer = config.renderer.build(tokenizer_path=config.hf_assets_path)
+
+        # Renderer stop tokens are injected into the generator at spawn
+        self._stop_token_ids = list(self.renderer.get_stop_token_ids())
+        self._sampling = config.generator.sampling
+        # TODO: pass our own tokenizer to the renderer and read pad/eos off it
+        # once `renderers` supports bring-your-own-tokenizer
+        # (https://github.com/PrimeIntellect-ai/renderers/pull/70).
+        # Until then, reach into the renderer's tokenizer for the pad id (eos doubles as pad).
+        self.batcher = Batcher(
+            config.batcher, pad_id=self.renderer._tokenizer.eos_token_id
+        )
+        self._rollouter: Rollouter = config.rollouter.build()
 
     async def close(self):
         """Best-effort: tear down actors, close metric backends, then stop proc meshes."""
@@ -409,6 +435,17 @@ async def setup_async(
                 nodes (no heterogeneous node configurations). Required when
                 host_mesh is provided.
         """
+        # Thread pool for TokenEnv's asyncio.to_thread renderer calls — one worker per
+        # concurrent rollout, capped by CPUs.
+        max_concurrent_rollouts = max(
+            self.config.num_groups_per_rollout_batch * self.config.group_size,
+            self.config.num_validation_samples,
+        )
+        max_workers = max(1, min(max_concurrent_rollouts, os.cpu_count() or 1))
+        asyncio.get_running_loop().set_default_executor(
+            ThreadPoolExecutor(max_workers=max_workers)
+        )
+
         config = self.config
 
         self.trainer_world_size = self._compute_world_size(config.trainer.parallelism)
@@ -517,10 +554,11 @@ async def setup_async(
                 model_path=config.hf_assets_path,
                 compile_config=config.compile,
                 max_num_seqs=max(
-                    config.num_prompts_per_step * config.generator.sampling.n,
+                    config.num_groups_per_rollout_batch * config.group_size,
                     config.num_validation_samples,
                 ),
                 output_dir=config.dump_folder,
+                stop_token_ids=self._stop_token_ids,
             )
 
         # Initialize TorchStore for weight sync between trainer and generator.
@@ -541,98 +579,335 @@ async def setup_async(
     @sl.log_trace_span("_collect_rollouts")
     async def _collect_rollouts(
         self,
+        *,
+        is_validation: bool,
         num_groups: int,
+        group_size: int,
+        sampling: SamplingConfig,
         step: int,
-        group_offset: int = 0,
-    ) -> tuple[list[Trajectory], list[m.Metric]]:
-        """Collect group rollouts and emit completion-shape rollout metrics.
+        group_offset: int,
+    ) -> tuple[list[RolloutGroup], list[m.Metric]]:
+        """Sample examples, run rollouts, score groups, and emit metrics.
+
+        Steps:
+        1. Sample one example per group from the train / validation dataset
+        2. Create N envs per example
+        3. Reset every env to get each rollout's first prompt
+        4. Split groups by init status
+        5. Run one batched `generate` over valid groups (n=1; pre-expanded)
+        6. For each group: step generated rollouts if needed, then score with `rollouter.score_group`
 
         Args:
-            num_groups: Number of prompt groups to collect in this round.
-            step: Current training step (passed to env for curriculum).
-            group_offset: Starting group index so that env ``group_idx``
-                values are unique across collection rounds within a step.
+            is_validation: Sample from the validation dataset (else train).
+            num_groups: Number of prompt groups to collect this call.
+            group_size: Sibling rollouts per group (prompts pre-expanded; generator runs n=1).
+            sampling: SamplingConfig for the generate call.
+            step: Training step, tagged into group_id / sample_id for metrics + debugging.
+            group_offset: Starting group index so group_ids stay unique across rounds in a step.
+
+        Returns:
+            Scored rollout groups plus rollout/generator metrics. Per-group
+            rollout/scoring failures are logged and dropped.
+
+        TODO(continuous-batching): once available, run rollouts independently
+        instead of batching one `generate` over all prompts at once.
         """
-        envs = [
-            self.config.env.build(step=step, group_idx=group_offset + i)
-            for i in range(num_groups)
-        ]
-        # TODO: Add a check max_tokens = min(max_tokens, context_window - model_input.length)
-        # and pass max_tokens to the generator call or skip the call if max_tokens<=0.
-        # Do the same for validation.
-        tokenized_prompts = [
-            self.tokenizer.encode(env.prompt, add_bos=True, add_eos=False)
-            for env in envs
-        ]
-        completions, generation_metrics = self._get_rank_0_value(
-            await self.generator.generate.call(tokenized_prompts)
+
+        generation_metrics_prefix = (
+            "validation_generator" if is_validation else "generator"
         )
+        rollout_metrics_prefix = "validation" if is_validation else "rollout"
+
+        # 1. Get one sample per group from the train / validation dataset;
+        # 2. create N envs per sample.
+        group_states: list[_RolloutGroupState] = []
+        for group_idx in range(num_groups):
+            sample = (
+                self._rollouter.get_validation_sample()
+                if is_validation
+                else self._rollouter.get_training_sample()
+            )
+            group_id = f"step={step}/group={group_offset + group_idx}"
+            envs = self._rollouter.make_env_group(
+                sample=sample, group_size=group_size, renderer=self.renderer
+            )
+            group_states.append(
+                _RolloutGroupState(
+                    group_id=group_id,
+                    sample=sample,
+                    envs=envs,
+                )
+            )
 
-        trajectories: list[Trajectory] = []
-        with sl.log_trace_span("score"):
-            for c in completions:
-                step_result = envs[c.prompt_idx].step(c.text)
-                trajectories.append(
-                    Trajectory(
-                        sample_idx=group_offset + c.prompt_idx,
-                        prompt_token_ids=tokenized_prompts[c.prompt_idx],
-                        transitions=[(c, step_result)],
-                    )
+        generation_metrics: list[m.Metric] = []
+        # 3. Reset every env to get its first prompt.
+        env_init_outputs_per_group_state = await asyncio.gather(  # [G][group_size]
+            *(
+                asyncio.gather(*(env.init() for env in group_state.envs))
+                for group_state in group_states
+            )
+        )
+        for group_state, env_init_outputs in zip(
+            group_states, env_init_outputs_per_group_state, strict=True
+        ):
+            group_state.env_init_outputs = env_init_outputs
+
+        # 4. Skip invalid group states and collect the rest for one batched
+        # generate call. Reset status is treated as a group-level invariant:
+        # all siblings are valid, or the whole group state is skipped.
+        valid_group_states: list[_RolloutGroupState] = []
+        num_skipped_groups = 0
+        for group_state in group_states:
+            is_valid = all(
+                not env_init_output.status.is_terminal()
+                for env_init_output in group_state.env_init_outputs
+            )
+            if is_valid:
+                valid_group_states.append(group_state)
+            else:
+                num_skipped_groups += 1
+                # TODO: log skipped prompts so they remain debuggable.
+                await asyncio.gather(
+                    *(env.close() for env in group_state.envs),
+                    return_exceptions=True,
+                )
+
+        # Prepare generate requests
+        prompt_token_ids: list[list[int]] = []  # [num_valid_samples][prompt_tokens]
+        request_ids: list[str] = []  # [num_valid_samples]
+        for group_state in valid_group_states:
+            for sample_idx, env_init_output in enumerate(group_state.env_init_outputs):
+                prompt_token_ids.append(env_init_output.next_prompt_token_ids or [])
+                # TODO(multi-turn): make request_id unique per turn (append a turn index). Today it is
+                # per-sample, so multiple generate calls within one multi-turn rollout would reuse the id.
+                request_ids.append(_sample_id(group_state.group_id, sample_idx))
+
+        # 5. Run one batched generate over valid group states (n=1; pre-expanded).
+        # TODO: pass the remaining budget (max_rollout_tokens - len(prompt)) to the
+        # sampling_config, to limit generation length in one turn.
+        completions: list[Completion] = []
+        if valid_group_states:
+            completions, generation_metrics = self._get_rank_0_value(
+                await self.generator.generate.call(
+                    prompt_token_ids,
+                    request_ids=request_ids,
+                    sampling_config=sampling,
+                    metrics_prefix=generation_metrics_prefix,
+                )
+            )
+            returned_ids = [completion.request_id for completion in completions]
+            if returned_ids != request_ids:
+                raise RuntimeError(
+                    f"generator returned request_ids {returned_ids}, "
+                    f"expected {request_ids}"
+                )
+
+        # 6. After the batch-level generate returns, finish each group state
+        # independently: step generated rollouts, score the group, and append
+        # the result.
+        # TODO(continuous-batching): group completions by group_id parsed from request_id instead of
+        # relying on returned order; under CB completions won't come back in request order.
+        completion_offset = 0
+        for group_state in valid_group_states:
+            next_completion_offset = completion_offset + len(
+                group_state.env_init_outputs
+            )
+            group_state.completions = completions[
+                completion_offset:next_completion_offset
+            ]
+            completion_offset = next_completion_offset
+
+        finished_rollout_groups: list[RolloutGroup | None] = await asyncio.gather(
+            *(
+                self._run_group_rollout(
+                    group_state=group_state,
                 )
+                for group_state in valid_group_states
+            )
+        )
 
-        # Metrics
-        response_lens = [len(c.token_ids) for c in completions]
-        prompt_lens = [len(t.prompt_token_ids) for t in trajectories]
-        total_lens = [p + r for p, r in zip(prompt_lens, response_lens, strict=True)]
-        truncated = [c.finish_reason == "length" for c in completions]
-        rollout_metrics: list[m.Metric] = [
-            m.Metric("rollout/response_length", m.Mean.from_list(response_lens)),
-            m.Metric("rollout/response_length", m.Max.from_list(response_lens)),
-            m.Metric("rollout/prompt_length", m.Mean.from_list(prompt_lens)),
-            m.Metric("rollout/prompt_length", m.Max.from_list(prompt_lens)),
-            m.Metric("rollout/total_length", m.Max.from_list(total_lens)),
-            m.Metric("rollout/truncation_rate", m.Mean.from_list(truncated)),
+        # Compute Metrics
+        rollout_groups = [
+            rollout_group
+            for rollout_group in finished_rollout_groups
+            if rollout_group is not None
         ]
+        num_failed_groups = (
+            num_skipped_groups + len(finished_rollout_groups) - len(rollout_groups)
+        )
+
+        generation_metrics.append(
+            m.Metric(
+                f"{rollout_metrics_prefix}/group_failures",
+                m.Sum(float(num_failed_groups)),
+            )
+        )
+        rollout_metrics = prepare_rollout_metrics(
+            rollout_metrics_prefix,
+            [
+                rollout
+                for rollout_group in rollout_groups
+                for rollout in rollout_group.rollouts
+            ],
+        )
         rollout_metrics += generation_metrics
-        rollout_metrics += _prepare_reward_metrics(
-            prefix="reward/component", trajectories=trajectories
+        return rollout_groups, rollout_metrics
+
+    @sl.log_trace_span("_run_group_rollout")
+    async def _run_group_rollout(
+        self,
+        *,
+        group_state: _RolloutGroupState,
+    ) -> RolloutGroup | None:
+        """Step generated rollouts, score the group, and return it.
+
+        Args:
+            group_state: One prompt group's envs, prompt steps, and rollout slots.
+
+        Returns:
+            Scored rollout group, or `None` if this group failed and should be dropped.
+        """
+        try:
+            if group_state.completions is None:
+                raise RuntimeError(f"group {group_state.group_id} has no completions")
+
+            rollouts: list[Rollout] = await asyncio.gather(
+                *(
+                    self._run_single_rollout(
+                        group_id=group_state.group_id,
+                        sample_id=_sample_id(group_state.group_id, sample_idx),
+                        env=env,
+                        env_init_output=env_init_output,
+                        completion=completion,
+                    )
+                    for sample_idx, (env, env_init_output, completion) in enumerate(
+                        zip(
+                            group_state.envs,
+                            group_state.env_init_outputs,
+                            group_state.completions,
+                            strict=True,
+                        )
+                    )
+                )
+            )
+
+            outputs = await self._rollouter.score_group(rollouts, group_state.sample)
+            for rollout, output in zip(rollouts, outputs, strict=True):
+                rollout.reward = output.reward
+                rollout.reward_breakdown = output.reward_breakdown
+            return RolloutGroup(group_id=group_state.group_id, rollouts=rollouts)
+        except Exception:
+            # TODO: add better logging so they are debuggable
+            logger.exception(
+                "group %s rollout/scoring failed; dropping", group_state.group_id
+            )
+            return None
+        finally:
+            await asyncio.gather(
+                *(env.close() for env in group_state.envs),
+                return_exceptions=True,
+            )
+
+    @sl.log_trace_span("_run_single_rollout")
+    async def _run_single_rollout(
+        self,
+        *,
+        group_id: str,
+        sample_id: str,
+        env: TokenEnv,
+        env_init_output: TokenEnvOutput,
+        completion: Completion,
+    ) -> Rollout:
+        """Step one env with its completion into a `Rollout`. On failure, return the
+        turns collected so far with an `ERROR` status.
+
+        Reward is left unset; the controller scores via `rollouter.score_group(...)`
+        afterward and fills `reward` / `reward_breakdown`.
+
+        Args:
+            group_id: Prompt-group ID; siblings share it for advantage centering.
+            sample_id: Unique rollout id.
+            env: The env for this rollout.
+            env_init_output: env output whose prompt produced this completion.
+            completion: Generator completion for this env's prompt.
+
+        Returns:
+            One unscored Rollout.
+        """
+        rollout_turns: list[RolloutTurn] = []
+        try:
+            env_output = await env.step(completion)
+            rollout_turns.append(
+                RolloutTurn(
+                    prompt_token_ids=env_init_output.next_prompt_token_ids or [],
+                    prompt_messages=env_init_output.next_prompt_messages or [],
+                    completion_token_ids=completion.token_ids,
+                    completion_logprobs=completion.token_logprobs,
+                    policy_version=completion.policy_version,
+                    completion_message=env_output.completion_message,
+                    env_messages=env_output.env_messages,
+                    env_rewards=env_output.env_rewards,
+                )
+            )
+            status = env_output.status
+            # TODO(multi-turn): while not status.is_terminal(): generate → step → append turn.
+            if not status.is_terminal():
+                raise RuntimeError(
+                    f"env {sample_id} returned a non-terminal turn; "
+                    "the controller does not yet support multi-turn rollouts."
+                )
+        except Exception:
+            logger.exception(
+                "rollout %s failed; keeping %d turn(s) as ERROR",
+                sample_id,
+                len(rollout_turns),
+            )
+            status = RolloutStatus.ERROR
+        return Rollout(
+            group_id=group_id, sample_id=sample_id, status=status, turns=rollout_turns
         )
-        return trajectories, rollout_metrics
 
     @staticmethod
     @sl.log_trace_span("_build_episodes")
     def _build_episodes(
-        trajectories: list[Trajectory],
+        rollout_groups: list[RolloutGroup],
     ) -> tuple[list[Episode], list[m.Metric]]:
-        """Group trajectories by sample, apply mean-baseline advantage, emit metrics."""
-        groups: dict[int, list[Trajectory]] = {}
-        for t in trajectories:
-            groups.setdefault(t.sample_idx, []).append(t)
+        """Build train episodes and GRPO advantages from scored rollout groups.
 
+        Centers each group's rewards by its mean, skips rollouts without
+        training tokens, and emits reward/advantage metrics.
+
+        Args:
+            rollout_groups: Scored rollout groups from one collection round.
+
+        Returns:
+            Train episodes plus episode-level metrics.
+        """
+        # Mean-baseline advantage per group
         episodes: list[Episode] = []
         group_stds: list[float] = []
-        for sample_idx, group in groups.items():
-            rewards = [t.total_reward for t in group]
-            group_mean = sum(rewards) / len(rewards)
-            # Population standard deviation; NaN for an empty group.
-            group_stds.append(statistics.pstdev(float(r) for r in rewards))
-            for t in group:
-                # Single-turn: exactly one (completion, step) per trajectory.
-                c, _ = t.transitions[0]
-                episodes.append(
-                    Episode(
-                        policy_version=c.policy_version,
-                        prompt_idx=sample_idx,
-                        prompt_token_ids=t.prompt_token_ids,
-                        text=c.text,
-                        token_ids=c.token_ids,
-                        token_logprobs=c.token_logprobs,
-                        reward=t.total_reward,
-                        advantage=t.total_reward - group_mean,
-                    )
+        for group in rollout_groups:
+            # Drop the whole group if any sibling has no trainable tokens; we
+            # need one turn with assistant tokens to build an episode.
+            if any(
+                not rollout.turns or not rollout.turns[0].completion_token_ids
+                for rollout in group.rollouts
+            ):
+                logger.warning(
+                    "group %s has an untrainable rollout; dropping the group",
+                    group.group_id,
                 )
+                continue
 
-        num_groups = len(groups)
+            rewards = [rollout.reward for rollout in group.rollouts]
+            group_mean = sum(rewards) / len(rewards)
+            group_stds.append(statistics.pstdev(rewards))
+
+            for rollout in group.rollouts:
+                rollout.advantage = rollout.reward - group_mean
+                episodes.append(rollout_to_episode(rollout))
+
+        num_groups = len(rollout_groups)
         zero_std_frac = (
             sum(1 for s in group_stds if s == 0.0) / num_groups if num_groups else 0.0
         )
@@ -666,63 +941,36 @@ def _build_episodes(
             )
         return episodes, episode_metrics
 
+    # TODO: we currently determine num_validation_samples
+    # but what if i want to run the entire dataset?
     @sl.log_trace_span("validate")
     async def validate(self) -> list[m.Metric]:
-        """Run validation on held-out prompts using greedy sampling.
+        """Run greedy validation on held-out prompts.
 
-        TODO: investigate using pass@k.
+        Returns:
+            Validation rollout metrics, generation metrics, and validation
+            timing.
         """
+        # TODO: investigate using pass@k for validation.
         t_validate_start = time.perf_counter()
         num_samples = self.config.num_validation_samples
-        envs = [
-            self.config.validation_env.build(step=0, group_idx=i)
-            for i in range(num_samples)
-        ]
-        greedy = SamplingConfig(
-            n=1,
-            temperature=0.0,
-            top_p=1.0,
-            max_tokens=self.config.generator.sampling.max_tokens,
-        )
+        greedy = replace(self._sampling, temperature=0.0, top_p=1.0)
 
-        tokenized_prompts: list[list[int]] = [
-            self.tokenizer.encode(env.prompt, add_bos=True, add_eos=False)
-            for env in envs
-        ]
-        completions, generation_metrics = self._get_rank_0_value(
-            await self.generator.generate.call(
-                tokenized_prompts,
-                sampling_config=greedy,
-                metrics_prefix="validation_generator",
-            )
+        rollout_groups, validation_metrics = await self._collect_rollouts(
+            is_validation=True,
+            num_groups=num_samples,
+            group_size=1,
+            sampling=greedy,
+            step=0,
+            group_offset=0,
         )
-
-        trajectories = [
-            Trajectory(
-                sample_idx=i,
-                prompt_token_ids=tokenized_prompts[i],
-                transitions=[(c, envs[i].step(c.text))],
-            )
-            for i, c in enumerate(completions)
-        ]
+        rollouts = [rollout for group in rollout_groups for rollout in group.rollouts]
 
         if self.config.log_samples:
-            _log_samples(completions)
+            _log_samples(rollout_groups)
 
-        validation_metrics: list[m.Metric] = [
-            m.Metric(
-                "validation/reward",
-                m.SummaryStats.from_list([t.total_reward for t in trajectories]),
-            ),
-            m.Metric(
-                "validation/response_length",
-                m.Mean.from_list([len(c.token_ids) for c in completions]),
-            ),
-            m.Metric("validation/num_samples", m.NoReduce(float(len(trajectories)))),
-        ]
-        validation_metrics += generation_metrics
-        validation_metrics += _prepare_reward_metrics(
-            prefix="validation/reward/component", trajectories=trajectories
+        validation_metrics.append(
+            m.Metric("validation/num_samples", m.NoReduce(float(len(rollouts))))
         )
 
         t_validate_s = time.perf_counter() - t_validate_start
@@ -731,7 +979,7 @@ async def validate(self) -> list[m.Metric]:
 
     async def train(self):
         num_steps = self.config.num_steps
-        num_groups = self.config.num_prompts_per_step
+        num_groups = self.config.num_groups_per_rollout_batch
         logger.info(f"Pre-training validation; then {num_steps} steps of RL training")
 
         # collect validation metrics before training
@@ -761,7 +1009,7 @@ async def train(self):
             # token budget. The Batcher then packs, truncates to
             # global_batch_size rows, and splits into microbatches.
             t_rollout_start = time.perf_counter()
-            trajectories: list[Trajectory] = []
+            rollout_groups: list[RolloutGroup] = []
             rollout_metrics: list[m.Metric] = []
             collected_tokens = 0
             group_offset = 0
@@ -771,24 +1019,30 @@ async def train(self):
             # rows, so actual token consumption may exceed collected_tokens.
             num_tokens_target = self.batcher.num_tokens_target(self.trainer_dp_degree)
             while collected_tokens < num_tokens_target:
-                new_trajectories, new_metrics = await self._collect_rollouts(
-                    num_groups, step=step, group_offset=group_offset
+                new_rollout_groups, new_metrics = await self._collect_rollouts(
+                    is_validation=False,
+                    num_groups=num_groups,
+                    group_size=self.config.group_size,
+                    sampling=self._sampling,
+                    step=step,
+                    group_offset=group_offset,
                 )
-                trajectories.extend(new_trajectories)
+                rollout_groups.extend(new_rollout_groups)
                 rollout_metrics.extend(new_metrics)
                 # Both prompt length and completion length are counted.
                 collected_tokens += sum(
-                    len(t.prompt_token_ids) + len(c.token_ids) - 1
-                    for t in new_trajectories
-                    for c, _ in t.transitions
+                    len(t.prompt_token_ids) + len(t.completion_token_ids) - 1
+                    for group in new_rollout_groups
+                    for r in group.rollouts
+                    for t in r.turns
                 )
                 group_offset += num_groups
 
-            episodes, episode_metrics = self._build_episodes(trajectories)
+            episodes, episode_metrics = self._build_episodes(rollout_groups)
             t_rollout_s = time.perf_counter() - t_rollout_start
 
             if self.config.log_samples:
-                _log_samples(episodes)
+                _log_samples(rollout_groups)
 
             # --- train ---
             t_train_start = time.perf_counter()
@@ -844,7 +1098,8 @@ async def train(self):
 
             # --- Prepare metrics ---
             total_tokens = sum(
-                len(ep.prompt_token_ids) + len(ep.token_ids) for ep in episodes
+                len(ep.prompt_token_ids) + len(ep.completion_token_ids)
+                for ep in episodes
             )
 
             step_metrics: list[m.Metric] = []
diff --git a/torchtitan/experiments/rl/renderer.py b/torchtitan/experiments/rl/renderer.py
new file mode 100644
index 0000000000..4ebdebaaa8
--- /dev/null
+++ b/torchtitan/experiments/rl/renderer.py
@@ -0,0 +1,93 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+from __future__ import annotations
+
+import logging
+from dataclasses import dataclass, fields
+
+from renderers import config_from_name, create_renderer, Renderer
+
+from torchtitan.config import Configurable
+
+logger = logging.getLogger(__name__)
+
+# Map a TorchTitan model name to its `renderers` renderer. Models not listed fall
+# back to "auto" (renderers resolves from the tokenizer)
+# https://github.com/PrimeIntellect-ai/renderers/blob/942449c37ab6e9fab26d59b40336514c8baa6b13/renderers/configs.py#L404
+_RENDERER_BY_MODEL = {
+    "qwen3": "qwen3",
+    "qwen3_vl": "qwen3-vl",
+    "gpt_oss": "gpt-oss",
+    "deepseek_v3": "deepseek-v3",
+    "default": "default",  # llama3, llama4
+    "auto": "auto",  # ignores knobs, resolves from tokenizer,
+}
+
+
+@dataclass(kw_only=True, slots=True)
+class RendererConfig(Configurable.Config):
+    """Selects the renderer used for chat message <-> token conversion.
+
+    Wraps `PrimeIntellect-ai/renderers`. `build` loads a tokenizer from
+    `tokenizer_path`, maps the model `name` to a renderer, and forwards any
+    supported knobs.
+
+    Args:
+        name: TorchTitan model name (e.g. `"qwen3"`, `"llama3"`), mapped to a
+            `renderers` renderer via `_RENDERER_BY_MODEL`. `None` (the default)
+            resolves the renderer from the tokenizer.
+        tool_parser: Tool-call parser name, when the renderer supports it.
+        reasoning_parser: Reasoning parser name, when the renderer supports it.
+        enable_thinking: Let the model emit reasoning, when supported.
+        preserve_all_thinking: Keep historical reasoning in future prompts.
+        preserve_thinking_between_tool_calls: Keep reasoning during tool loops.
+
+    Every field defaults to `None`; a non-`None` value overrides that knob on the
+    chosen renderer's config, otherwise the renderer keeps its own default.
+
+    Example:
+
+        renderer = RendererConfig(name="qwen3").build(tokenizer_path="./Qwen3-0.6B")
+        prompt_ids = renderer.render_ids(
+            [{"role": "user", "content": "hi"}], add_generation_prompt=True
+        )
+    """
+
+    name: str | None = None
+    tool_parser: str | None = None
+    reasoning_parser: str | None = None
+    enable_thinking: bool | None = None
+    preserve_all_thinking: bool | None = None
+    preserve_thinking_between_tool_calls: bool | None = None
+
+    def build(self, *, tokenizer_path: str) -> Renderer:
+        # TODO(renderers#70): use TorchTitan's tokenizer once `renderers` supports
+        # bring-your-own-tokenizer (PR adds a Tokenizer protocol; drops transformers).
+        from transformers import AutoTokenizer
+
+        tokenizer = AutoTokenizer.from_pretrained(tokenizer_path)
+
+        # `name=None` (or "auto") -> let `create_renderer` resolve from the tokenizer.
+        renderer_name = _RENDERER_BY_MODEL.get(self.name, self.name)
+        renderer_config = config_from_name(renderer_name) if renderer_name else None
+        if renderer_config is None:
+            return create_renderer(tokenizer, None)
+
+        # Rebuild the typed config and pass parameters
+        # that are not None and are supported
+        config_type = type(renderer_config)
+        args = {
+            field.name: getattr(self, field.name)  # {key: value}
+            for field in fields(self)
+            if field.name != "name"  # Get all self.fields, except name
+            and getattr(self, field.name) is not None  # Only consider provided fields
+            and field.name in config_type.model_fields  # Config supports this field
+        }
+        logger.info(
+            f"Using renderer {renderer_name}, of type {config_type}, with args {args}"
+        )
+        return create_renderer(tokenizer, config_type(**args))
diff --git a/torchtitan/experiments/rl/rollout/__init__.py b/torchtitan/experiments/rl/rollout/__init__.py
new file mode 100644
index 0000000000..b7819bda3d
--- /dev/null
+++ b/torchtitan/experiments/rl/rollout/__init__.py
@@ -0,0 +1,30 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+# NOTE: `Rollouter` is intentionally NOT re-exported here. It imports `environment`,
+# which imports `rollout.types` — re-exporting it from this package `__init__` would make
+# that a circular import. Import it from the submodule: `rollout.rollouter import Rollouter`.
+from torchtitan.experiments.rl.rollout.types import (
+    Rollout,
+    RolloutGroup,
+    RolloutStatus,
+    RolloutTurn,
+)
+from torchtitan.experiments.rl.rollout.utils import (
+    last_completion_text,
+    prepare_rollout_metrics,
+    rollout_to_episode,
+)
+
+__all__ = [
+    "Rollout",
+    "RolloutGroup",
+    "RolloutStatus",
+    "RolloutTurn",
+    "last_completion_text",
+    "prepare_rollout_metrics",
+    "rollout_to_episode",
+]
diff --git a/torchtitan/experiments/rl/rollout/rollouter.py b/torchtitan/experiments/rl/rollout/rollouter.py
new file mode 100644
index 0000000000..0d5456fade
--- /dev/null
+++ b/torchtitan/experiments/rl/rollout/rollouter.py
@@ -0,0 +1,132 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+
+from renderers import Renderer
+
+from torchtitan.config import Configurable
+from torchtitan.experiments.rl.environment import MessageEnv, TokenEnv
+from torchtitan.experiments.rl.rollout.types import Rollout
+from torchtitan.experiments.rl.rubrics import Rubric, RubricOutput
+
+
+# TODO(continuous-batching): when VLLMGenerator gains continuous batching,
+# move the rollout loop onto Rollouter as `run_rollout(example, client) -> Rollout`,
+# so each rollout drives its own generate calls.
+class Rollouter(Configurable):
+    """Turns a problem (train/val datasets, the `MessageEnv` to build per example, and a
+    `Rubric`) into scored rollouts — the RL training data.
+
+    Like a `Dataloader` turns a `Dataset` into training batches, a `Rollouter`
+    turns a problem into rollouts: it builds the envs, the controller drives them against
+    the inference engine, and `score_group` scores the results.
+
+    Subclass only to override specific methods, such as `score_group` for cross-sibling scoring,
+    or `make_env_group` for custom logic, such as using a pool of envs instead of creating a new one.
+
+    The flow for one prompt group:
+
+        sample = rollouter.get_training_sample()  # one sample from the dataset
+        envs = rollouter.make_env_group(sample=sample, group_size=N, renderer=renderer)
+        ...  # the controller runs the rollout loop
+        outputs = rollouter.score_group(rollouts, sample)  # the Rubric scores them
+
+    `MessageEnv` works in messages; `TokenEnv` (what `make_env_group` returns)
+    adds the message <-> token plumbing.
+
+    Example:
+        rollouter = Rollouter.Config(
+            train_dataset=MyDataset.Config(seed=42),
+            validation_dataset=MyDataset.Config(seed=99),
+            rubric=Rubric.Config(
+                reward_fns=[RewardCorrect.Config(), RewardFormat.Config(weight=0.3)]
+            ),
+            message_env=MyEnv.Config(),
+        ).build()
+    """
+
+    @dataclass(kw_only=True, slots=True)
+    class Config(Configurable.Config):
+        train_dataset: Configurable.Config
+        """Dataset iterator for training (`next()` yields one env input)."""
+
+        validation_dataset: Configurable.Config
+        """Dataset iterator for validation."""
+
+        rubric: Rubric.Config
+        """Reward functions + weights used by `score_group`."""
+
+        message_env: MessageEnv.Config
+        """The env to build per sample; `make_env_group` calls `build(env_input=sample)`."""
+
+        token_env: TokenEnv.Config = field(default_factory=TokenEnv.Config)
+        """`TokenEnv` limits (e.g. `max_rollout_tokens`) passed to `make_env_group`."""
+
+    def __init__(self, config: Config) -> None:
+        self._train_dataset = config.train_dataset.build()
+        self._validation_dataset = config.validation_dataset.build()
+        self.rubric: Rubric = config.rubric.build()
+        self._message_env_config = config.message_env
+        self._token_env_config = config.token_env
+
+    # TODO: revisit this abstraction: should it return a sample or a dataset or an iterator?
+    def get_training_sample(self) -> object:
+        """Get one training sample (the env input) from the training dataset."""
+        return next(self._train_dataset)
+
+    def get_validation_sample(self) -> object:
+        """Get one validation sample (the env input) from the validation dataset."""
+        return next(self._validation_dataset)
+
+    # TODO: revisit the Renderer being injected into `make_env_group` once we
+    # know whether Rollouter should own a Renderer (per-rollouter chat templates).
+    def make_env_group(
+        self,
+        *,
+        sample: object,
+        group_size: int,
+        renderer: Renderer,
+    ) -> list[TokenEnv]:
+        """Construct `group_size` single-use envs from one dataset sample.
+
+        Args:
+            sample: the dataset sample (the env input) from `get_training_sample` / `get_validation_sample`.
+            group_size: number of sibling envs for this prompt group.
+            renderer: Renderer shared by the rollout controller.
+
+        Returns:
+            `TokenEnv` * `group_size` instances, each ready for one rollout.
+        """
+        return [
+            self._token_env_config.build(
+                message_env=self._message_env_config.build(env_input=sample),
+                renderer=renderer,
+            )
+            for _ in range(group_size)
+        ]
+
+    async def score_group(
+        self,
+        rollouts: list[Rollout],
+        env_input: object,
+    ) -> list[RubricOutput]:
+        """Score one group's rollouts; the controller applies the rewards.
+
+        Default impl delegates to `self.rubric.score_group`. Override for
+        cross-sibling scoring (judge, pairwise, diversity) or partial-credit
+        reward shaping.
+
+        Args:
+            rollouts: Sibling rollouts in one prompt group, already stepped.
+            env_input: Dataset payload shared by the group.
+
+        Returns:
+            One `RubricOutput` per rollout, in input order.
+        """
+        return await self.rubric.score_group(rollouts, env_input)
diff --git a/torchtitan/experiments/rl/rollout/types.py b/torchtitan/experiments/rl/rollout/types.py
new file mode 100644
index 0000000000..58d2cb688b
--- /dev/null
+++ b/torchtitan/experiments/rl/rollout/types.py
@@ -0,0 +1,121 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+from enum import StrEnum
+
+from renderers import Message
+
+
+_TRUNCATED = frozenset({"truncated_length", "truncated_prompt_too_long"})
+_ERROR = frozenset({"error_parse", "error_timeout", "error_abort", "error"})
+
+
+class RolloutStatus(StrEnum):
+    """Per-rollout status."""
+
+    ONGOING = "ongoing"
+    COMPLETED = "completed"
+    TRUNCATED_LENGTH = "truncated_length"
+    TRUNCATED_PROMPT_TOO_LONG = "truncated_prompt_too_long"
+    ERROR_PARSE = "error_parse"
+    ERROR_TIMEOUT = "error_timeout"
+    ERROR_ABORT = "error_abort"
+    ERROR = "error"
+
+    def is_truncated(self) -> bool:
+        return self.value in _TRUNCATED
+
+    def is_error(self) -> bool:
+        return self.value in _ERROR
+
+    def is_terminal(self) -> bool:
+        return self is not RolloutStatus.ONGOING
+
+
+@dataclass(kw_only=True, slots=True)
+class RolloutTurn:
+    """Full per-turn snapshot: the prompt fed to the generator + the sampled completion +
+    the env's reply, in both token and message space. Rubrics score it and
+    `rollout_to_episode` flattens it into training tokens."""
+
+    # TODO: add a `logs` field (raw prompt/response text, finish_reason, timings)
+    # so a turn can be dumped and inspected without re-deriving from tokens.
+
+    # Fields needed for training
+    prompt_token_ids: list[int]  # [num_prompt_tokens]
+    """Tokenized conversation up to this turn, used to generate this turn's completion."""
+
+    completion_token_ids: list[int]  # [num_completion_tokens]
+    """This turn's completion token ids."""
+
+    completion_logprobs: list[float]  # [num_completion_tokens]
+    """This turn's completion token logprobs generated by the generator policy."""
+
+    # Filtering / debugging
+    # TODO(async): add per-turn sampling metadata (sampling time, generator id) for mixed-policy filtering.
+    policy_version: int | None = None
+    """Trainer policy version when this response was sampled; `None` for a
+    prompt-only turn (no generation happened, e.g. the prompt was too long)."""
+
+    # Logging
+    prompt_messages: list[Message] = field(
+        default_factory=list
+    )  # [num_prompt_messages]
+    """Full conversation up to this turn; equivalent to `prompt_token_ids`."""
+
+    completion_message: Message | None = None
+    """This turn's completion decoded into a message by the renderer (the TokenEnv's parse,
+    not the raw generator output)."""
+
+    env_messages: list[Message] = field(default_factory=list)  # [num_env_messages]
+    """This turn env's reply messages (tool / user)."""
+
+    # For rubrics
+    env_rewards: dict[str, float] = field(default_factory=dict)
+    """This turn optional reward signals the env attached; the rubric decides how to use them."""
+
+
+@dataclass(kw_only=True, slots=True)
+class Rollout:
+    """A complete rollout: ordered turns + terminal state + reward + identifier."""
+
+    # TODO: add a `logs` field (per-turn debug records / event trace) to make a
+    # full rollout reconstructable for debugging.
+
+    group_id: str
+    """Prompt-group ID; siblings share it for advantage centering."""
+
+    sample_id: str
+    """Unique rollout id within the step, e.g. `"step=3/group=5/sample=2"`."""
+
+    turns: list[RolloutTurn] = field(default_factory=list)  # [num_turns]
+    """Ordered rollout turns. Each turn stores its full prompt (redundant across turns); kept so a
+    rollout can be replayed/branched and divergences found, then collapsed at episode assembly."""
+
+    status: RolloutStatus
+    """Rollout-level terminal status."""
+
+    reward: float | None = None
+    """Final weighted reward, filled by the rubric."""
+
+    reward_breakdown: dict[str, float] = field(default_factory=dict)
+    """Raw per-reward-function values, filled by the rubric."""
+
+    # TODO: make it per token
+    advantage: float | None = None
+    """Advantage for this sample."""
+
+
+@dataclass(kw_only=True, slots=True)
+class RolloutGroup:
+    group_id: str
+    """Prompt-group ID; siblings share it for advantage centering."""
+
+    rollouts: list[Rollout]  # [group_size]
+    """Sibling rollouts sampled from the group's shared prompt."""
diff --git a/torchtitan/experiments/rl/rollout/utils.py b/torchtitan/experiments/rl/rollout/utils.py
new file mode 100644
index 0000000000..8873a111b2
--- /dev/null
+++ b/torchtitan/experiments/rl/rollout/utils.py
@@ -0,0 +1,91 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+from __future__ import annotations
+
+from collections import defaultdict
+
+from torchtitan.experiments.rl.observability import metrics as m
+from torchtitan.experiments.rl.rollout.types import Rollout
+from torchtitan.experiments.rl.types import Episode
+
+
+def last_completion_text(rollout: Rollout) -> str:
+    """Return the completion message text from the last turn, or `""`."""
+    if not rollout.turns:
+        return ""
+    msg = rollout.turns[-1].completion_message
+    return (msg.get("content") or "") if msg else ""
+
+
+def rollout_to_episode(rollout: Rollout) -> Episode:
+    """Flatten a scored single-turn `Rollout` into an `Episode`, a class
+    that holds only the information needed for training.
+    """
+    # TODO: support multi-turn rollout flattening.
+    # TODO(branching): when a turn's prompt history diverges from the previous turn's
+    #       (e.g. the env edited/compacted history), the turns no longer share a prefix
+    #       and must be split into separate training sequences instead of one flat episode.
+    # TODO: rename Episode -> TrainingSample / rollout_to_episode ->
+    #       rollout_to_training_sample (consistent with TrainingBatch).
+    if len(rollout.turns) != 1:
+        raise ValueError(
+            f"rollout_to_episode expects exactly one turn; got {len(rollout.turns)}."
+        )
+    turn = rollout.turns[0]
+    return Episode(
+        policy_version=turn.policy_version,
+        sample_id=rollout.sample_id,
+        prompt_token_ids=turn.prompt_token_ids,
+        completion_text=last_completion_text(rollout),
+        completion_token_ids=turn.completion_token_ids,
+        completion_logprobs=turn.completion_logprobs,
+        reward=rollout.reward,
+        advantage=rollout.advantage if rollout.advantage is not None else 0.0,
+    )
+
+
+def prepare_rollout_metrics(prefix: str, rollouts: list[Rollout]) -> list[m.Metric]:
+    """Build rollout-derived metrics (lengths, truncation, reward breakdown).
+
+    Args:
+        prefix: Metric namespace (e.g. `"rollout"` or `"validation"`).
+        rollouts: Rollouts to summarize.
+    """
+    # Lengths, truncation, reward
+    # TODO: adapt for multi-turn rollouts
+    completion_lens = [len(t.completion_token_ids) for r in rollouts for t in r.turns]
+    prompt_lens = [len(r.turns[0].prompt_token_ids) for r in rollouts if r.turns]
+    total_lens = [
+        len(r.turns[-1].prompt_token_ids) + len(r.turns[-1].completion_token_ids)
+        for r in rollouts
+        if r.turns
+    ]
+
+    truncated = [float(r.status.is_truncated()) for r in rollouts]
+    rewards = [r.reward for r in rollouts if r.reward is not None]
+
+    out: list[m.Metric] = [
+        m.Metric(f"{prefix}/response_length", m.Mean.from_list(completion_lens)),
+        m.Metric(f"{prefix}/response_length", m.Max.from_list(completion_lens)),
+        m.Metric(f"{prefix}/prompt_length", m.Mean.from_list(prompt_lens)),
+        m.Metric(f"{prefix}/prompt_length", m.Max.from_list(prompt_lens)),
+        m.Metric(f"{prefix}/total_length", m.Mean.from_list(total_lens)),
+        m.Metric(f"{prefix}/total_length", m.Max.from_list(total_lens)),
+        m.Metric(f"{prefix}/truncation_rate", m.Mean.from_list(truncated)),
+        m.Metric(f"{prefix}_reward", m.SummaryStats.from_list(rewards)),
+    ]
+
+    # Per-component reward breakdown
+    values_by_name: dict[str, list[float]] = defaultdict(list)
+    for rollout in rollouts:
+        for name, value in rollout.reward_breakdown.items():
+            values_by_name[name].append(float(value))
+    out.extend(
+        m.Metric(f"{prefix}_reward/component/{name}", m.Mean.from_list(values))
+        for name, values in sorted(values_by_name.items())
+    )
+    return out
diff --git a/torchtitan/experiments/rl/rubrics/__init__.py b/torchtitan/experiments/rl/rubrics/__init__.py
new file mode 100644
index 0000000000..d89bc5a3a9
--- /dev/null
+++ b/torchtitan/experiments/rl/rubrics/__init__.py
@@ -0,0 +1,9 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+from torchtitan.experiments.rl.rubrics.rubric import RewardFn, Rubric, RubricOutput
+
+__all__ = ["RewardFn", "Rubric", "RubricOutput"]
diff --git a/torchtitan/experiments/rl/rubrics/rubric.py b/torchtitan/experiments/rl/rubrics/rubric.py
new file mode 100644
index 0000000000..ff8f955620
--- /dev/null
+++ b/torchtitan/experiments/rl/rubrics/rubric.py
@@ -0,0 +1,177 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+from __future__ import annotations
+
+import abc
+import asyncio
+from dataclasses import dataclass, field
+
+from torchtitan.config import Configurable
+from torchtitan.experiments.rl.rollout.types import Rollout
+from torchtitan.observability import structured_logger as sl
+
+
+class RewardFn(Configurable, abc.ABC):
+    """A single reward function, as a Configurable callable.
+
+    Subclass and implement `__call__`. Its `Config` carries the `weight` used in
+    the rubric's weighted sum, plus any args a stateful reward fn needs (a reward
+    model path, an LLM-judge endpoint, a threshold, ...).
+
+    Example:
+        class RewardCorrect(RewardFn):
+            @dataclass(kw_only=True, slots=True)
+            class Config(RewardFn.Config):
+                pass  # only needs `weight`
+
+            async def __call__(self, rollout, env_input) -> float:
+                ...
+    """
+
+    @dataclass(kw_only=True, slots=True)
+    class Config(Configurable.Config):
+        weight: float = 1.0
+        """Relative weight in the rubric's weighted sum (normalized across fns)."""
+
+    def __init__(self, config: Config) -> None:
+        self.weight = config.weight
+
+    @abc.abstractmethod
+    async def __call__(self, rollout: Rollout, env_input: object) -> float:
+        """Return this fn's score for one rollout.
+
+        Args:
+            rollout: Rollout to score.
+            env_input: Dataset payload used to build the env (target/metadata).
+        """
+
+
+@dataclass(frozen=True, kw_only=True, slots=True)
+class RubricOutput:
+    """One rollout's reward, as returned by a `Rubric`.
+
+    Example:
+        >>> RubricOutput(reward=0.5, reward_breakdown={"RewardCorrect": 1.0, "RewardFormat": 0.0})
+    """
+
+    reward: float
+    """Final scalar reward for this rollout; assigned to `Rollout.reward`."""
+
+    reward_breakdown: dict[str, float] = field(default_factory=dict)
+    """Per-reward-fn outputs (unweighted), keyed by reward-fn class name.
+    The default `Rubric` computes `reward` from these; callers may also use them for per-reward
+    advantage, reweighting, metrics, or inspection."""
+
+
+class Rubric(Configurable):
+    """Scores rollouts with a set of weighted reward functions.
+
+    The reward fns and their weights live in config (`reward_fns`), so common
+    cases need no subclass. Subclass and override `score_group` for cross-sibling
+    scoring (pairwise comparison, diversity, rank normalization).
+
+    Setting `truncation_reward` / `error_reward` short-circuits the reward fns for
+    rollouts whose status is truncated / errored.
+
+    Example:
+        rubric = Rubric.Config(
+            reward_fns=[RewardCorrect.Config(weight=1.0), RewardFormat.Config(weight=0.3)],
+            truncation_reward=0.0,
+        ).build()
+    """
+
+    @dataclass(kw_only=True, slots=True)
+    class Config(Configurable.Config):
+        reward_fns: list[RewardFn.Config] = field(default_factory=list)
+        """The rubric's reward fns + weights; built and weight-normalized at init."""
+
+        truncation_reward: float | None = None
+        """Reward for a truncated rollout. If set, the reward fns are SKIPPED and this fixed
+        reward is used. If None, the reward fns run on the truncated rollout."""
+
+        error_reward: float | None = None
+        """Reward for a errored rollout. If set, the reward fns are SKIPPED and this fixed
+        reward is used. If None, the reward fns run on the errored rollout."""
+
+    def __init__(self, config: Config) -> None:
+        self._config = config
+        self._reward_fns = [rwd_cfg.build() for rwd_cfg in config.reward_fns]
+
+        # Sanity checks
+        if not self._reward_fns:
+            raise ValueError("Rubric.Config.reward_fns must not be empty")
+        names = [type(fn).__name__ for fn in self._reward_fns]
+        if len(names) != len(set(names)):
+            raise ValueError(f"reward fn names must be unique; got {names}")
+        self._weight_sum = sum(fn.weight for fn in self._reward_fns)
+        if self._weight_sum <= 0:
+            raise ValueError(
+                f"rubric weights must sum to a positive value; got {self._weight_sum}"
+            )
+
+    @sl.log_trace_span("score_single_rollout")
+    async def _score_single_rollout(
+        self, rollout: Rollout, env_input: object
+    ) -> RubricOutput:
+        """Score one rollout. Short-circuits to `truncation_reward` /
+        `error_reward` when those are set and the rollout truncated / errored.
+
+        Args:
+            rollout: Rollout to score.
+            env_input: Dataset payload used to build the env (target/metadata).
+
+        Returns:
+            Final weighted reward + per-fn raw breakdown.
+        """
+        # Short-circuit on truncate / error and return the configured reward. The
+        # breakdown records the short-circuit reason so it shows up in metrics.
+        cfg = self._config
+        if cfg.truncation_reward is not None and rollout.status.is_truncated():
+            return RubricOutput(
+                reward=cfg.truncation_reward,
+                reward_breakdown={"truncated": cfg.truncation_reward},
+            )
+        if cfg.error_reward is not None and rollout.status.is_error():
+            return RubricOutput(
+                reward=cfg.error_reward,
+                reward_breakdown={"errored": cfg.error_reward},
+            )
+
+        # Run all reward fns and weight-sum (weights normalized to sum to 1.0).
+        per_fn_rewards = await asyncio.gather(
+            *(fn(rollout, env_input) for fn in self._reward_fns)
+        )
+
+        reward_breakdown = {}
+        total_reward = 0.0
+        for fn, r in zip(self._reward_fns, per_fn_rewards, strict=True):
+            reward_breakdown[type(fn).__name__] = r
+            total_reward += (fn.weight / self._weight_sum) * r
+
+        return RubricOutput(reward=total_reward, reward_breakdown=reward_breakdown)
+
+    @sl.log_trace_span("score_group")
+    async def score_group(
+        self,
+        rollouts: list[Rollout],
+        env_input: object,
+    ) -> list[RubricOutput]:
+        """Score every rollout in one prompt group.
+
+        Override for cross-rollout rewards (pairwise comparison, diversity,
+        rank normalization).
+
+        Args:
+            rollouts: Sibling rollouts sampled from one prompt group.
+            env_input: Dataset payload originally used to construct the rollout env.
+
+        Returns:
+            One `RubricOutput` per rollout, in input order.
+        """
+        return await asyncio.gather(
+            *(self._score_single_rollout(r, env_input) for r in rollouts)
+        )
diff --git a/torchtitan/experiments/rl/sum_digits.py b/torchtitan/experiments/rl/sum_digits.py
deleted file mode 100644
index 9b6d4b475f..0000000000
--- a/torchtitan/experiments/rl/sum_digits.py
+++ /dev/null
@@ -1,76 +0,0 @@
-# Copyright (c) Meta Platforms, Inc. and affiliates.
-# All rights reserved.
-#
-# This source code is licensed under the BSD-style license found in the
-# LICENSE file in the root directory of this source tree.
-
-from __future__ import annotations
-
-import random
-import re
-from dataclasses import dataclass
-
-from torchtitan.config import Configurable
-from torchtitan.experiments.rl.types import Step
-
-
-class SumDigitsEnv(Configurable):
-    """Single-turn, single-use env for one sum-of-digits problem.
-
-    Construct via ``SumDigitsEnv.Config(seed=...).build(step=<s>, group_idx=<i>)``.
-    The problem is a pure function of ``(config.seed, step, group_idx)``: same
-    inputs always produce the same prompt and target. No RNG state is shared
-    between envs.
-    """
-
-    @dataclass(kw_only=True, slots=True)
-    class Config(Configurable.Config):
-        correctness_reward: float = 1.0
-        """Reward for a response containing ``[ANSWER] <target>``."""
-
-        format_reward: float = 0.3
-        """Reward bonus for any ``[ANSWER] <number>`` tag in the response."""
-
-        seed: int = 42
-        """Seed mixed with ``(step, group_idx)`` to deterministically generate problems."""
-
-    SYSTEM_PROMPT = """\
-You are a helpful assistant. Solve the problem step by step.
-When you have your final answer, state it as [ANSWER] <number>.
-
-Example:
-User: What is the total digit sum of [12, 345, 67]?
-Assistant: Break each number into digits:
-12 → 1, 2
-345 → 3, 4, 5
-67 → 6, 7
-Sum all digits: 1 + 2 + 3 + 4 + 5 + 6 + 7 = 28
-[ANSWER] 28"""
-
-    def __init__(self, config: Config, *, step: int = 0, group_idx: int = 0):
-        self._config = config
-        rng = random.Random(f"{config.seed}:{step}:{group_idx}")
-        n = rng.randint(2, 4)
-        numbers = [rng.randint(10, 99) for _ in range(n)]
-        self._target = sum(int(d) for num in numbers for d in str(num))
-        question = f"What is the total digit sum of {numbers}?"
-        self.prompt = f"{self.SYSTEM_PROMPT}\n\n{question}"
-
-    def step(self, completion: str) -> Step:
-        return Step(
-            rewards={
-                "correctness": self._correctness_reward(completion),
-                "format": self._format_reward(completion),
-            },
-            done=True,
-        )
-
-    def _correctness_reward(self, completion: str) -> float:
-        matches = re.findall(r"\[ANSWER\]\s*(-?\d+)", completion)
-        correct = bool(matches) and int(matches[-1]) == self._target
-        return self._config.correctness_reward if correct else 0.0
-
-    def _format_reward(self, completion: str) -> float:
-        if re.search(r"\[ANSWER\]\s*-?\d+", completion):
-            return self._config.format_reward
-        return 0.0
diff --git a/torchtitan/experiments/rl/tests/test_bitwise_parity.py b/torchtitan/experiments/rl/tests/test_bitwise_parity.py
index 0658345b3a..f8a4ac6534 100644
--- a/torchtitan/experiments/rl/tests/test_bitwise_parity.py
+++ b/torchtitan/experiments/rl/tests/test_bitwise_parity.py
@@ -219,7 +219,7 @@ def build_inference_engine(config: RLTrainer.Config) -> LLMEngine:
         engine_kwargs["block_size"] = 256  # set blocksize to be 256 to align with FA2
 
     engine_kwargs["max_model_len"] = config.model_spec.model.rope.max_seq_len
-    max_num_seqs = config.num_prompts_per_step * gen_config.sampling.n
+    max_num_seqs = config.num_groups_per_rollout_batch * config.group_size
     engine_kwargs["max_num_seqs"] = max_num_seqs
     vllm_compilation_config = gen_config.cudagraph.get_vllm_compilation_config(
         max_num_seqs=max_num_seqs,
diff --git a/torchtitan/experiments/rl/tests/test_generator.py b/torchtitan/experiments/rl/tests/test_generator.py
index 580d1fc8df..cd3d1c31c1 100644
--- a/torchtitan/experiments/rl/tests/test_generator.py
+++ b/torchtitan/experiments/rl/tests/test_generator.py
@@ -95,17 +95,20 @@ def _generator(outputs):
     generator = VLLMGenerator.__new__(VLLMGenerator)
     generator._engine = _FakeEngine(outputs)
     generator.policy_version = 7
+    generator._stop_token_ids = []
     generator.config = SimpleNamespace(
-        sampling=SamplingConfig(n=1, temperature=0.0, top_p=1.0, max_tokens=4),
+        sampling=SamplingConfig(temperature=0.0, top_p=1.0, max_tokens=4),
         debug=SimpleNamespace(seed=None),
     )
     return generator
 
 
-def _run_generate(generator, tokenized_prompts, **kwargs):
+def _run_generate(generator, tokenized_prompts, *, request_ids=None, **kwargs):
+    if request_ids is None:
+        request_ids = [str(i) for i in range(len(tokenized_prompts))]
     return asyncio.run(
         VLLMGenerator.generate._method(
-            generator, tokenized_prompts, **kwargs
+            generator, tokenized_prompts, request_ids=request_ids, **kwargs
         )  # noqa: SLF001
     )
 
diff --git a/torchtitan/experiments/rl/tests/test_grpo_metrics.py b/torchtitan/experiments/rl/tests/test_grpo_metrics.py
deleted file mode 100644
index e806ce7d5a..0000000000
--- a/torchtitan/experiments/rl/tests/test_grpo_metrics.py
+++ /dev/null
@@ -1,669 +0,0 @@
-# Copyright (c) Meta Platforms, Inc. and affiliates.
-# All rights reserved.
-#
-# This source code is licensed under the BSD-style license found in the
-# LICENSE file in the root directory of this source tree.
-
-"""Unit tests for RL metric helpers + controller-subroutine outputs.
-
-These tests do **not** start Monarch, vLLM, W&B, or distributed
-process groups. Controller subroutines are invoked as static / instance
-methods on plain dataclasses.
-"""
-
-from __future__ import annotations
-
-import asyncio
-import math
-from unittest.mock import AsyncMock, MagicMock, patch
-
-import pytest
-
-import torch
-
-from torchtitan.experiments.rl.grpo import _prepare_reward_metrics, GRPOLoss, RLTrainer
-from torchtitan.experiments.rl.observability import metrics as m
-from torchtitan.experiments.rl.types import Completion, Step, Trajectory
-
-
-# ---------------------------------------------------------------------------
-# _prepare_reward_metrics
-# ---------------------------------------------------------------------------
-
-
-def _step(rewards: dict[str, float]) -> Step:
-    return Step(rewards=rewards, done=True)
-
-
-def _reward_trajectory(rewards: dict[str, float], sample_idx: int = 0) -> Trajectory:
-    """Single-turn trajectory with a fake completion + the given rewards."""
-    fake_completion = Completion(
-        policy_version=0,
-        prompt_idx=sample_idx,
-        text="",
-        token_ids=[],
-        token_logprobs=[],
-    )
-    return Trajectory(
-        sample_idx=sample_idx,
-        prompt_token_ids=[],
-        transitions=[(fake_completion, _step(rewards))],
-    )
-
-
-class TestBuildRewardMetrics:
-    def test_one_metric_per_observed_name(self) -> None:
-        trajectories = [
-            _reward_trajectory({"correctness": 1.0, "format": 0.5}, sample_idx=0),
-            _reward_trajectory({"correctness": 0.0, "format": 1.0}, sample_idx=1),
-        ]
-        metrics = _prepare_reward_metrics("reward/component", trajectories)
-        keys = {entry.key for entry in metrics}
-        assert keys == {
-            "reward/component/correctness",
-            "reward/component/format",
-        }
-        for entry in metrics:
-            assert isinstance(entry.value, m.Mean)
-
-    def test_components_observed_in_some_trajectories_only(self) -> None:
-        # `format` only appears in the second trajectory — it should
-        # average over that one entry (no zero-fill).
-        trajectories = [
-            _reward_trajectory({"correctness": 1.0}, sample_idx=0),
-            _reward_trajectory({"format": 0.5}, sample_idx=1),
-        ]
-        metrics = _prepare_reward_metrics("reward/component", trajectories)
-        agg = m.MetricsProcessor._aggregate_metrics(metrics)
-        assert agg["reward/component/correctness/mean"] == 1.0
-        assert agg["reward/component/format/mean"] == 0.5
-
-    def test_empty_input(self) -> None:
-        assert _prepare_reward_metrics("reward/component", []) == []
-
-    def test_prefix_controls_namespace(self) -> None:
-        trajectories = [_reward_trajectory({"correctness": 1.0}, sample_idx=0)]
-        metrics = _prepare_reward_metrics("validation/reward/component", trajectories)
-        assert metrics[0].key == "validation/reward/component/correctness"
-
-
-# ---------------------------------------------------------------------------
-# Controller subroutines: _collect_rollouts, _build_episodes, validate
-# ---------------------------------------------------------------------------
-
-
-def _completion(
-    prompt_idx: int,
-    response_len: int,
-    finish_reason: str | None = "stop",
-    *,
-    policy_version: int = 0,
-) -> Completion:
-    return Completion(
-        policy_version=policy_version,
-        prompt_idx=prompt_idx,
-        text="x" * response_len,
-        token_ids=list(range(response_len)),
-        token_logprobs=[0.0] * response_len,
-        finish_reason=finish_reason,
-    )
-
-
-class _FakeEnv:
-    """Minimal env stub: step(text) returns a preset reward dict."""
-
-    def __init__(self, rewards: dict[str, float], prompt: str = "p"):
-        self.prompt = prompt
-        self._rewards = rewards
-
-    def step(self, text: str) -> Step:
-        return _step(self._rewards)
-
-
-def _build_collect_rollouts_inputs(self_obj):
-    """Build a hollow RLTrainer instance + completions wired into the
-    fake generator — without spawning real meshes.
-    """
-
-    completions = [
-        _completion(prompt_idx=0, response_len=10),
-        _completion(prompt_idx=0, response_len=20),
-        _completion(prompt_idx=1, response_len=15),
-    ]
-
-    class _RewardEnvBuilder:
-        @staticmethod
-        def build(*, step, group_idx):
-            return _FakeEnv({"correctness": float(group_idx), "format": 0.5})
-
-    self_obj.config = MagicMock()
-    self_obj.config.env = _RewardEnvBuilder
-    self_obj.tokenizer = MagicMock()
-    self_obj.tokenizer.encode.side_effect = lambda prompt, **_: [ord(prompt)]
-    self_obj.generator = MagicMock()
-    # `generate.call` is awaited inside `_collect_rollouts`, so stub it async.
-    self_obj.generator.generate.call = AsyncMock()
-    # `_get_rank_0_value` is the layer that strips Monarch's ValueMesh,
-    # so just make it return whatever it's handed.
-    self_obj._get_rank_0_value = lambda value, has_gpus=True: (completions, [])
-    return completions
-
-
-class TestCollectRollouts:
-    def test_passes_token_ids_to_generator(self) -> None:
-        """Controller tokenizes env prompts and hands the IDs (not strings)
-        to ``generator.generate.call``."""
-        controller = RLTrainer.__new__(RLTrainer)
-        _build_collect_rollouts_inputs(controller)
-        asyncio.run(controller._collect_rollouts(num_groups=2, step=0))
-        # _FakeEnv.prompt == "p"; encode side_effect returns [ord(prompt)] = [112].
-        controller.generator.generate.call.assert_called_once_with([[112], [112]])
-
-    def test_emits_expected_metric_keys(self) -> None:
-        controller = RLTrainer.__new__(RLTrainer)
-        completions = _build_collect_rollouts_inputs(controller)
-        trajectories, rollout_metrics = asyncio.run(
-            controller._collect_rollouts(num_groups=2, step=0)
-        )
-        assert len(trajectories) == len(completions)
-        agg = m.MetricsProcessor._aggregate_metrics(rollout_metrics)
-        # Length keys: Mean+Max for prompt/response, Max-only for total.
-        assert "rollout/response_length/mean" in agg
-        assert "rollout/response_length/max" in agg
-        assert "rollout/prompt_length/mean" in agg
-        assert "rollout/prompt_length/max" in agg
-        assert "rollout/total_length/max" in agg
-        # Reward-component keys derived from env step output (now under
-        # the top-level reward/ namespace).
-        assert "reward/component/correctness/mean" in agg
-        assert "reward/component/format/mean" in agg
-
-    def test_truncation_rate(self) -> None:
-        """rollout/truncation_rate averages
-        finish_reason == 'length' over completions."""
-        controller = RLTrainer.__new__(RLTrainer)
-        completions = [
-            _completion(0, 10, finish_reason="length"),
-            _completion(0, 10, finish_reason="stop"),
-            _completion(1, 10, finish_reason="length"),
-            _completion(1, 10, finish_reason="length"),
-        ]
-
-        controller.config = MagicMock()
-        controller.config.env = MagicMock()
-        controller.config.env.build = lambda *, step, group_idx: _FakeEnv({"r": 1.0})
-        controller.tokenizer = MagicMock()
-        controller.tokenizer.encode.side_effect = lambda prompt, **_: [ord(prompt)]
-        controller.generator = MagicMock()
-        controller.generator.generate.call = AsyncMock()
-        controller._get_rank_0_value = lambda value, has_gpus=True: (completions, [])
-
-        _, rollout_metrics = asyncio.run(
-            controller._collect_rollouts(num_groups=2, step=0)
-        )
-        agg = m.MetricsProcessor._aggregate_metrics(rollout_metrics)
-        # 3 of 4 completions hit max_tokens.
-        assert agg["rollout/truncation_rate/mean"] == pytest.approx(0.75)
-
-    def test_total_length_uses_per_episode_max(self) -> None:
-        """rollout/total_length/max must be max(prompt+response per
-        episode), **not** max(prompt) + max(response) — the latter
-        may combine two different episodes."""
-        controller = RLTrainer.__new__(RLTrainer)
-
-        # Carefully chosen so per-side maxes don't align: the longest
-        # prompt has the shortest response, etc. The tokenizer mock below
-        # produces token lists of length == len(prompt), so the env prompts
-        # control prompt-side lengths.
-        env_prompts = {0: "x" * 10, 1: "x" * 2, 2: "x" * 5}
-        completions = [
-            _completion(prompt_idx=0, response_len=2),  # total 10 + 2 = 12
-            _completion(prompt_idx=1, response_len=10),  # total 2 + 10 = 12
-            _completion(prompt_idx=2, response_len=5),  # total 5 + 5 = 10
-        ]
-        controller.config = MagicMock()
-        controller.config.env = MagicMock()
-        controller.config.env.build = lambda *, step, group_idx: _FakeEnv(
-            {"r": 1.0}, prompt=env_prompts[group_idx]
-        )
-        controller.tokenizer = MagicMock()
-        controller.tokenizer.encode.side_effect = lambda prompt, **_: list(
-            prompt.encode()
-        )
-        controller.generator = MagicMock()
-        controller.generator.generate.call = AsyncMock()
-        controller._get_rank_0_value = lambda value, has_gpus=True: (completions, [])
-
-        trajectories, rollout_metrics = asyncio.run(
-            controller._collect_rollouts(num_groups=3, step=0)
-        )
-        agg = m.MetricsProcessor._aggregate_metrics(rollout_metrics)
-        per_side_max_sum = max(len(t.prompt_token_ids) for t in trajectories) + max(
-            len(c.token_ids) for c in completions
-        )  # = 10 + 10 = 20
-        actual_max = max(
-            len(t.prompt_token_ids) + len(c.token_ids)
-            for t, c in zip(trajectories, completions, strict=True)
-        )  # = 12
-        assert agg["rollout/total_length/max"] == actual_max
-        assert agg["rollout/total_length/max"] < per_side_max_sum
-
-    def test_group_offset_keeps_rollout_rounds_distinct(self) -> None:
-        controller = RLTrainer.__new__(RLTrainer)
-        _build_collect_rollouts_inputs(controller)
-
-        first_round, _ = asyncio.run(
-            controller._collect_rollouts(num_groups=2, step=0, group_offset=0)
-        )
-        second_round, _ = asyncio.run(
-            controller._collect_rollouts(num_groups=2, step=0, group_offset=2)
-        )
-
-        assert [t.sample_idx for t in first_round] == [0, 0, 1]
-        assert [t.sample_idx for t in second_round] == [2, 2, 3]
-
-        episodes, _ = RLTrainer._build_episodes(first_round + second_round)
-        assert all(ep.advantage == 0.0 for ep in episodes)
-
-
-def _trajectory(
-    sample_idx: int,
-    prompt_len: int,
-    response_len: int,
-    reward: float,
-    *,
-    policy_version: int = 0,
-) -> Trajectory:
-    completion = _completion(sample_idx, response_len, policy_version=policy_version)
-    return Trajectory(
-        sample_idx=sample_idx,
-        prompt_token_ids=list(range(prompt_len)),
-        transitions=[(completion, _step({"r": reward}))],
-    )
-
-
-class TestBuildEpisodes:
-    def test_emits_expected_metric_keys(self) -> None:
-        trajectories = [
-            _trajectory(0, 4, 5, reward=1.0),
-            _trajectory(0, 4, 5, reward=0.0),
-            _trajectory(1, 4, 5, reward=0.5),
-            _trajectory(1, 4, 5, reward=0.5),
-        ]
-        episodes, episode_metrics = RLTrainer._build_episodes(trajectories)
-        assert len(episodes) == 4
-        agg = m.MetricsProcessor._aggregate_metrics(episode_metrics)
-        # SummaryStats expansion (5 sub-keys each, top-level reward/advantage).
-        for prefix in ("reward", "advantage"):
-            for sub in ("max", "mean", "min", "std", "sum"):
-                assert f"{prefix}/_{sub}" in agg
-        # Group std + degenerate fraction live under reward/.
-        assert "reward/group_std/mean" in agg
-        assert "reward/group_std/max" in agg
-        assert "reward/zero_std_frac" in agg
-        # Per-rollout policy version distribution (min/max only).
-        assert "rollout/policy_version/mean" not in agg
-        assert "rollout/policy_version/min" in agg
-        assert "rollout/policy_version/max" in agg
-        # num_prompts/num_episodes were dropped: make sure they don't
-        # creep back in.
-        assert "rollout/num_prompts" not in agg
-        assert "rollout/num_episodes" not in agg
-
-    def test_policy_version_metrics_single_version(self) -> None:
-        """When all rollouts came from the same policy version, min == max."""
-        single_version = [
-            _trajectory(0, 4, 5, reward=1.0, policy_version=5),
-            _trajectory(1, 4, 5, reward=0.5, policy_version=5),
-        ]
-        _, em = RLTrainer._build_episodes(single_version)
-        agg = m.MetricsProcessor._aggregate_metrics(em)
-        assert agg["rollout/policy_version/min"] == 5.0
-        assert agg["rollout/policy_version/max"] == 5.0
-
-    def test_policy_version_metrics_mixed_versions(self) -> None:
-        """Mixed rollout versions emit min and max."""
-        mixed_versions = [
-            _trajectory(0, 4, 5, reward=1.0, policy_version=2),
-            _trajectory(1, 4, 5, reward=0.5, policy_version=4),
-        ]
-        _, em = RLTrainer._build_episodes(mixed_versions)
-        agg = m.MetricsProcessor._aggregate_metrics(em)
-        assert agg["rollout/policy_version/min"] == 2.0
-        assert agg["rollout/policy_version/max"] == 4.0
-
-    def test_degenerate_group_fraction(self) -> None:
-        # Two groups: both constant => fraction == 1.0.
-        all_constant = [
-            _trajectory(0, 4, 5, reward=1.0),
-            _trajectory(0, 4, 5, reward=1.0),
-            _trajectory(1, 4, 5, reward=0.5),
-            _trajectory(1, 4, 5, reward=0.5),
-        ]
-        _, em = RLTrainer._build_episodes(all_constant)
-        agg = m.MetricsProcessor._aggregate_metrics(em)
-        assert agg["reward/zero_std_frac"] == 1.0
-
-        # Mixed: one group constant (degenerate), one group varied.
-        mixed = [
-            _trajectory(0, 4, 5, reward=1.0),
-            _trajectory(0, 4, 5, reward=1.0),
-            _trajectory(1, 4, 5, reward=0.0),
-            _trajectory(1, 4, 5, reward=1.0),
-        ]
-        _, em = RLTrainer._build_episodes(mixed)
-        agg = m.MetricsProcessor._aggregate_metrics(em)
-        assert agg["reward/zero_std_frac"] == 0.5
-
-        # Both groups have variance => 0.0.
-        none_constant = [
-            _trajectory(0, 4, 5, reward=0.0),
-            _trajectory(0, 4, 5, reward=1.0),
-            _trajectory(1, 4, 5, reward=0.5),
-            _trajectory(1, 4, 5, reward=1.5),
-        ]
-        _, em = RLTrainer._build_episodes(none_constant)
-        agg = m.MetricsProcessor._aggregate_metrics(em)
-        assert agg["reward/zero_std_frac"] == 0.0
-
-
-# ---------------------------------------------------------------------------
-# RLTrainer.Config wiring
-# ---------------------------------------------------------------------------
-
-
-class TestRLTrainerConfigWiring:
-    """Use the canonical `rl_grpo_qwen3_0_6b` registry config so the test
-    matches a real production config and stays insulated from any future
-    tightening of VLLMGenerator/PolicyTrainer field validators."""
-
-    def test_metrics_default_uses_factory(self) -> None:
-        from torchtitan.experiments.rl.config_registry import rl_grpo_qwen3_0_6b
-
-        cfg = rl_grpo_qwen3_0_6b()
-        baseline = m.MetricsProcessor.Config()
-        assert cfg.metrics.console_log_keys_train == baseline.console_log_keys_train
-        assert (
-            cfg.metrics.console_log_keys_validation
-            == baseline.console_log_keys_validation
-        )
-
-    def test_metrics_defaults_are_independent_copies(self) -> None:
-        """Mutating one Config's allow lists must not bleed into other instances."""
-        from torchtitan.experiments.rl.config_registry import rl_grpo_qwen3_0_6b
-
-        cfg = rl_grpo_qwen3_0_6b()
-        cfg.metrics.console_log_keys_train.append("X")
-        cfg.metrics.console_log_keys_validation.append("Y")
-        # A fresh Config still has the pristine defaults.
-        fresh = rl_grpo_qwen3_0_6b()
-        assert "X" not in fresh.metrics.console_log_keys_train
-        assert "Y" not in fresh.metrics.console_log_keys_validation
-
-    def test_metrics_default_wandb_enabled(self) -> None:
-        from torchtitan.experiments.rl.config_registry import rl_grpo_qwen3_0_6b
-
-        cfg = rl_grpo_qwen3_0_6b()
-        assert cfg.metrics.enable_wandb is True
-        assert cfg.metrics.enable_tensorboard is False
-
-
-# ---------------------------------------------------------------------------
-# GRPOLoss bridge
-# ---------------------------------------------------------------------------
-
-
-class TestGRPOLossBridge:
-    def test_loss_keeps_gradient(self) -> None:
-        """`loss` must remain differentiable so `.backward()` works.
-        Regression test for `_token_weighted_mean` accidentally detaching."""
-        loss_fn = GRPOLoss(GRPOLoss.Config(clip_eps=0.2))
-        policy_logprobs = [
-            torch.zeros(2, requires_grad=True),
-            torch.zeros(8, requires_grad=True),
-        ]
-
-        loss, _loss_metrics = loss_fn(
-            policy_logprobs=policy_logprobs,
-            advantages=torch.tensor([1.0, -1.0]),
-            num_global_valid_tokens=torch.tensor(10.0),
-        )
-
-        assert loss.requires_grad
-        assert loss.grad_fn is not None
-        loss.backward()
-        assert all(sample.grad is not None for sample in policy_logprobs)
-
-    def test_returns_loss_and_pre_normalized_metrics(self) -> None:
-        loss_fn = GRPOLoss(GRPOLoss.Config(clip_eps=0.2))
-        # Two samples with unequal response lengths.
-        policy_logprobs = [
-            torch.zeros(2, requires_grad=True),
-            torch.zeros(8, requires_grad=True),
-        ]
-        advantages = torch.tensor([1.0, -1.0])
-        # Single-rank case: global == local valid tokens.
-        num_global_valid_tokens = torch.tensor(10.0)
-
-        loss, loss_metrics = loss_fn(
-            policy_logprobs=policy_logprobs,
-            advantages=advantages,
-            num_global_valid_tokens=num_global_valid_tokens,
-        )
-        assert isinstance(loss, torch.Tensor)
-        assert isinstance(loss_metrics, dict)
-        for key in ("loss/mean", "loss/ratio/mean", "loss/ratio/clipped_frac"):
-            assert key in loss_metrics
-
-    def test_loss_is_token_weighted_sum_over_global_tokens(self) -> None:
-        """loss = sum_i(sample_loss_i * num_tokens_i) / num_global_valid_tokens.
-
-        Under unequal response lengths this differs from a naive sample mean.
-        """
-        loss_fn = GRPOLoss(GRPOLoss.Config(clip_eps=0.2))
-        policy_logprobs = [
-            torch.full((2,), 0.1, requires_grad=True),
-            torch.full((8,), 0.0, requires_grad=True),
-        ]
-        advantages = torch.tensor([1.0, -1.0])
-        num_global_valid_tokens = torch.tensor(10.0)
-
-        loss, loss_metrics = loss_fn(
-            policy_logprobs=policy_logprobs,
-            advantages=advantages,
-            num_global_valid_tokens=num_global_valid_tokens,
-        )
-        # loss/mean metric is the same value as loss (both pre-normalized).
-        assert math.isclose(
-            loss_metrics["loss/mean"].item(),
-            loss.item(),
-            rel_tol=1e-6,
-        )
-
-        # And it is NOT equal to the unweighted sample mean of policy gradient
-        # losses, which is what the prior implementation used.
-        per_sample_mean_logprobs = torch.stack(
-            [sample_logprobs.mean() for sample_logprobs in policy_logprobs]
-        )
-        ratio = torch.exp(per_sample_mean_logprobs)
-        clipped_ratio = torch.clamp(ratio, 1 - 0.2, 1 + 0.2)
-        sample_policy_gradient_losses = -torch.min(
-            ratio * advantages, clipped_ratio * advantages
-        )
-        unweighted_sample_mean = float(sample_policy_gradient_losses.mean().item())
-        assert not math.isclose(
-            loss.item(), unweighted_sample_mean, rel_tol=1e-4, abs_tol=1e-6
-        )
-
-
-# ---------------------------------------------------------------------------
-# Trainer reducers (single-DP fast paths)
-# ---------------------------------------------------------------------------
-
-
-def _stub_trainer_for_reducers(dp_size: int):
-    """Build a minimal stand-in for `PolicyTrainer` so we can exercise
-    the reducer methods on a CPU box without spawning Monarch / NCCL.
-    """
-    # Late import: importing `PolicyTrainer` triggers monarch/torchtitan
-    # actor wiring at module load time, which is fine for CPU.
-    from torchtitan.experiments.rl.actors.trainer import PolicyTrainer
-
-    inst = PolicyTrainer.__new__(PolicyTrainer)
-    inst.dp_size = dp_size
-    inst.device = torch.device("cpu")
-    inst.parallel_dims = MagicMock()
-    inst.parallel_dims.get_optional_mesh = MagicMock(return_value=None)
-    return inst
-
-
-class TestReducerFastPaths:
-    def test_single_dp_identical(self) -> None:
-        # Pre-normalized values pass through SUM-reduce unchanged on a single
-        # rank (no mesh -> no all-reduce -> values are exactly what we passed).
-        trainer = _stub_trainer_for_reducers(dp_size=1)
-        out = trainer.reduce_forward_backward_metrics(
-            sum_reduced_metrics={
-                "loss/mean": torch.tensor(3.0),
-                "bit_wise/logprob_diff/mean": torch.tensor(0.001),
-                "bit_wise/ratio_tokens_different/mean": torch.tensor(0.0),
-            },
-            max_reduced_metrics={"bit_wise/logprob_diff/max": torch.tensor(0.005)},
-        )
-        assert out["loss/mean"] == pytest.approx(3.0)
-        assert out["bit_wise/logprob_diff/mean"] == pytest.approx(0.001)
-        assert out["bit_wise/logprob_diff/max"] == pytest.approx(0.005)
-        assert out["bit_wise/ratio_tokens_different/mean"] == 0.0
-
-    def test_unbiased_sum_reduction_across_ranks(self) -> None:
-        """Two ranks contribute pre-normalized shares; SUM-reducing
-        reconstructs the global value.
-
-        Rank 0 shares: loss/mean=10/15 (token-weighted local share).
-        Rank 1 shares: loss/mean=30/15.
-        SUM-reduce: 40/15 = 2.667 (the global token-weighted mean).
-        """
-        trainer = _stub_trainer_for_reducers(dp_size=2)
-        trainer.parallel_dims.get_optional_mesh = MagicMock(return_value="loss")
-
-        rank0_share = torch.tensor([10.0 / 15.0], dtype=torch.float32)
-        rank1_share = torch.tensor([30.0 / 15.0], dtype=torch.float32)
-
-        def fake_all_reduce(t, *, reduceOp, group):
-            if t.numel() == 1 and t.dtype == torch.float32:
-                return rank0_share + rank1_share
-            return t
-
-        with patch(
-            "torchtitan.experiments.rl.actors.trainer.funcol.all_reduce",
-            side_effect=fake_all_reduce,
-        ):
-            out = trainer.reduce_forward_backward_metrics(
-                sum_reduced_metrics={"loss/mean": rank0_share[0]},
-                max_reduced_metrics={"bit_wise/logprob_diff/max": torch.tensor(0.0)},
-            )
-        assert out["loss/mean"] == pytest.approx(40.0 / 15.0)
-
-    def test_max_reduce_path(self) -> None:
-        """MAX-reduced metrics compose via elementwise max across ranks.
-
-        Patches funcol.all_reduce to dispatch on reduceOp: SUM doubles
-        (simulating two ranks contributing equal shares); MAX takes the
-        elementwise max with a higher second-rank value.
-        """
-        import torch.distributed.distributed_c10d as c10d
-
-        trainer = _stub_trainer_for_reducers(dp_size=2)
-        trainer.parallel_dims.get_optional_mesh = MagicMock(return_value="loss")
-
-        rank1_max = torch.tensor([0.006], dtype=torch.float32)
-
-        def fake_all_reduce(t, *, reduceOp, group):
-            if reduceOp == c10d.ReduceOp.SUM.name:
-                return t * 2
-            if reduceOp == c10d.ReduceOp.MAX.name:
-                return torch.maximum(t, rank1_max)
-            raise AssertionError(f"unexpected reduceOp={reduceOp!r}")
-
-        with patch(
-            "torchtitan.experiments.rl.actors.trainer.funcol.all_reduce",
-            side_effect=fake_all_reduce,
-        ):
-            out = trainer.reduce_forward_backward_metrics(
-                sum_reduced_metrics={"loss/mean": torch.tensor(0.5)},
-                max_reduced_metrics={"bit_wise/logprob_diff/max": torch.tensor(0.003)},
-            )
-        # SUM doubled: 0.5 + 0.5 = 1.0. MAX = max(0.003, 0.006) = 0.006.
-        assert out["loss/mean"] == pytest.approx(1.0)
-        assert out["bit_wise/logprob_diff/max"] == pytest.approx(0.006)
-
-    def test_sum_only_skips_max_collective(self) -> None:
-        """max_reduced_metrics={} must not crash and must not call the
-        MAX collective; the SUM bucket is still reduced normally."""
-        import torch.distributed.distributed_c10d as c10d
-
-        trainer = _stub_trainer_for_reducers(dp_size=2)
-        trainer.parallel_dims.get_optional_mesh = MagicMock(return_value="loss")
-
-        seen_ops: list[str] = []
-
-        def fake_all_reduce(t, *, reduceOp, group):
-            seen_ops.append(reduceOp)
-            if reduceOp == c10d.ReduceOp.SUM.name:
-                return t * 2
-            raise AssertionError(f"unexpected reduceOp={reduceOp!r}")
-
-        with patch(
-            "torchtitan.experiments.rl.actors.trainer.funcol.all_reduce",
-            side_effect=fake_all_reduce,
-        ):
-            out = trainer.reduce_forward_backward_metrics(
-                sum_reduced_metrics={"loss/mean": torch.tensor(0.5)},
-                max_reduced_metrics={},
-            )
-        assert seen_ops == [c10d.ReduceOp.SUM.name]
-        assert out == {"loss/mean": pytest.approx(1.0)}
-
-    def test_max_only_skips_sum_collective(self) -> None:
-        """sum_reduced_metrics={} must not crash and must not call the
-        SUM collective; the MAX bucket is still reduced normally."""
-        import torch.distributed.distributed_c10d as c10d
-
-        trainer = _stub_trainer_for_reducers(dp_size=2)
-        trainer.parallel_dims.get_optional_mesh = MagicMock(return_value="loss")
-
-        seen_ops: list[str] = []
-        rank1_max = torch.tensor([0.006], dtype=torch.float32)
-
-        def fake_all_reduce(t, *, reduceOp, group):
-            seen_ops.append(reduceOp)
-            if reduceOp == c10d.ReduceOp.MAX.name:
-                return torch.maximum(t, rank1_max)
-            raise AssertionError(f"unexpected reduceOp={reduceOp!r}")
-
-        with patch(
-            "torchtitan.experiments.rl.actors.trainer.funcol.all_reduce",
-            side_effect=fake_all_reduce,
-        ):
-            out = trainer.reduce_forward_backward_metrics(
-                sum_reduced_metrics={},
-                max_reduced_metrics={
-                    "bit_wise/logprob_diff/max": torch.tensor(0.003),
-                },
-            )
-        assert seen_ops == [c10d.ReduceOp.MAX.name]
-        assert out == {"bit_wise/logprob_diff/max": pytest.approx(0.006)}
-
-    def test_both_empty_returns_empty(self) -> None:
-        """Both buckets empty: no collectives called, empty dict returned."""
-        trainer = _stub_trainer_for_reducers(dp_size=2)
-        trainer.parallel_dims.get_optional_mesh = MagicMock(return_value="loss")
-        with patch(
-            "torchtitan.experiments.rl.actors.trainer.funcol.all_reduce",
-            side_effect=AssertionError("should not be called"),
-        ):
-            out = trainer.reduce_forward_backward_metrics(
-                sum_reduced_metrics={},
-                max_reduced_metrics={},
-            )
-        assert out == {}
diff --git a/torchtitan/experiments/rl/tests/test_shutdown.py b/torchtitan/experiments/rl/tests/test_shutdown.py
index 80218ddbbc..a839f80dcd 100644
--- a/torchtitan/experiments/rl/tests/test_shutdown.py
+++ b/torchtitan/experiments/rl/tests/test_shutdown.py
@@ -5,29 +5,24 @@
 # LICENSE file in the root directory of this source tree.
 
 import asyncio
+from types import SimpleNamespace
 
 import pytest
 
 from torchtitan.experiments.rl import grpo
 from torchtitan.experiments.rl.actors.generator import VLLMGenerator
-
-
-class _FakeConfigManager:
-    config = object()
-
-    def parse_args(self):
-        return self.config
+from torchtitan.experiments.rl.batcher import Batcher
 
 
 class _FakeRLTrainer:
     instances = []
 
-    def __init__(self, config):
+    def __init__(self, config=None):
         self.config = config
         self.events = []
         self.instances.append(self)
 
-    async def setup(self):
+    async def setup_async(self):
         self.events.append("setup")
         if getattr(self.config, "fail_setup", False):
             raise RuntimeError("setup failed")
@@ -43,11 +38,64 @@ async def close(self):
         self.events.append("close")
 
 
+class _FakeDebug:
+    enable_structured_logging = False
+
+
+class _FakeTrainerConfig:
+    debug = _FakeDebug()
+
+
+class _FakeConfig:
+    """Fake config whose build() returns a _FakeRLTrainer."""
+
+    dump_folder = "/tmp/test_rl"
+    trainer = _FakeTrainerConfig()
+
+    def __init__(self, **kwargs):
+        for k, v in kwargs.items():
+            setattr(self, k, v)
+
+    def build(self):
+        return _FakeRLTrainer(config=self)
+
+
+class _FakeConfigManager:
+    config = _FakeConfig()
+
+    def parse_args(self):
+        return self.config
+
+
+def _make_stub_rl_trainer():
+    """Create an RLTrainer with a minimal stub config (no VLLMGenerator validation)."""
+    from torchtitan.experiments.rl.observability import metrics as m
+
+    class _StubConfig:
+        batcher = Batcher.Config()
+        metrics = m.MetricsProcessor.Config()
+        dump_folder = "/tmp/test_rl"
+        hf_assets_path = "./tests/assets/tokenizer"
+        # __init__ builds these too; stub them so construction does no real work.
+        renderer = SimpleNamespace(
+            build=lambda *, tokenizer_path: SimpleNamespace(
+                get_stop_token_ids=lambda: [],
+                _tokenizer=SimpleNamespace(eos_token_id=0),
+            )
+        )
+        generator = SimpleNamespace(sampling=SimpleNamespace())
+        rollouter = SimpleNamespace(build=lambda: SimpleNamespace())
+
+        def to_dict(self):
+            return {}
+
+    return grpo.RLTrainer(_StubConfig())
+
+
 def test_main_shuts_down_after_success(monkeypatch):
-    _FakeConfigManager.config = object()
+    _FakeConfigManager.config = _FakeConfig()
     _FakeRLTrainer.instances = []
     monkeypatch.setattr(grpo, "ConfigManager", _FakeConfigManager)
-    monkeypatch.setattr(grpo, "RLTrainer", _FakeRLTrainer)
 
     asyncio.run(grpo.main())
 
@@ -55,13 +103,9 @@ def test_main_shuts_down_after_success(monkeypatch):
 
 
 def test_main_shuts_down_after_train_failure(monkeypatch):
-    class FailingConfig:
-        fail_train = True
-
-    _FakeConfigManager.config = FailingConfig()
+    _FakeConfigManager.config = _FakeConfig(fail_train=True)
     _FakeRLTrainer.instances = []
     monkeypatch.setattr(grpo, "ConfigManager", _FakeConfigManager)
-    monkeypatch.setattr(grpo, "RLTrainer", _FakeRLTrainer)
 
     with pytest.raises(RuntimeError, match="train failed"):
         asyncio.run(grpo.main())
@@ -70,13 +114,9 @@ class FailingConfig:
 
 
 def test_main_shuts_down_after_setup_failure(monkeypatch):
-    class FailingConfig:
-        fail_setup = True
-
-    _FakeConfigManager.config = FailingConfig()
+    _FakeConfigManager.config = _FakeConfig(fail_setup=True)
     _FakeRLTrainer.instances = []
     monkeypatch.setattr(grpo, "ConfigManager", _FakeConfigManager)
-    monkeypatch.setattr(grpo, "RLTrainer", _FakeRLTrainer)
 
     with pytest.raises(RuntimeError, match="setup failed"):
         asyncio.run(grpo.main())
@@ -85,7 +125,7 @@ class FailingConfig:
 
 
 def test_rl_trainer_shutdown_is_noop_before_meshes_spawn():
-    trainer = grpo.RLTrainer(object())
+    trainer = _make_stub_rl_trainer()
 
     asyncio.run(trainer.close())
 
@@ -99,14 +139,9 @@ def test_main_swallows_cancellation_after_shutdown(monkeypatch):
     running task; ``main`` runs ``close`` in ``finally`` and the explicit
     ``except`` clause swallows the interrupt so the process exits 0
     without a traceback."""
-
-    class CancelledConfig:
-        cancel_train = True
-
-    _FakeConfigManager.config = CancelledConfig()
+    _FakeConfigManager.config = _FakeConfig(cancel_train=True)
     _FakeRLTrainer.instances = []
     monkeypatch.setattr(grpo, "ConfigManager", _FakeConfigManager)
-    monkeypatch.setattr(grpo, "RLTrainer", _FakeRLTrainer)
 
     # No exception escapes; close still ran.
     asyncio.run(grpo.main())
@@ -146,7 +181,7 @@ async def stop(self):
 
 def test_shutdown_calls_actor_close_before_mesh_stop():
     events: list[str] = []
-    rl_trainer = grpo.RLTrainer(object())
+    rl_trainer = _make_stub_rl_trainer()
     rl_trainer.trainer = _StubActor("trainer.close", events)
     rl_trainer.generator = _StubActor("generator.close", events)
     rl_trainer._proc_meshes = [
@@ -167,7 +202,7 @@ def test_shutdown_calls_actor_close_before_mesh_stop():
 
 def test_shutdown_continues_after_actor_close_failure():
     events: list[str] = []
-    rl_trainer = grpo.RLTrainer(object())
+    rl_trainer = _make_stub_rl_trainer()
     rl_trainer.trainer = _StubActor("trainer.close", events, raises=True)
     rl_trainer.generator = _StubActor("generator.close", events)
     rl_trainer._proc_meshes = [_StubMesh("mesh.stop[0]", events)]
diff --git a/torchtitan/experiments/rl/types.py b/torchtitan/experiments/rl/types.py
index abad041596..d327b3e556 100644
--- a/torchtitan/experiments/rl/types.py
+++ b/torchtitan/experiments/rl/types.py
@@ -9,78 +9,33 @@
 import torch
 
 
-@dataclass(kw_only=True, slots=True)
-class Step:
-    """Env transition: named reward components, done flag, optional next observation.
-
-    ``rewards`` is a dict of component-name to value (e.g.
-    ``{"correctness": 1.0, "format": 0.3}``); envs are free to define
-    any decomposition. Trainers read the scalar ``reward`` property
-    (sum of components); loggers iterate ``rewards.items()`` for
-    per-component reporting without needing to know the keys.
-
-    ``observation`` (the next prompt the agent will see) is only
-    populated by multi-turn envs. Single-turn envs leave it None.
-    """
-
-    rewards: dict[str, float]
-    done: bool
-    observation: str | None = None
-
-    @property
-    def reward(self) -> float:
-        return sum(self.rewards.values())
-
-
 @dataclass(kw_only=True, slots=True)
 class Completion:
-    """A single generated sequence from the generator.
-
-    Pure generation artifact - no reward, no advantage. ``prompt_idx``
-    is the position of the source prompt in the input ``prompts`` list.
-    """
+    """A single generated sequence from the generator."""
 
     policy_version: int
-    prompt_idx: int
-    text: str
+    request_id: str
+    """Echoes the id the caller passed to `generate`, so callers can validate
+    ordered completions or map by id."""
     token_ids: list[int]
     token_logprobs: list[float]
     finish_reason: str | None = None
     """vLLM `CompletionOutput.finish_reason` ("stop" | "length" | "abort")"""
 
 
-@dataclass(kw_only=True, slots=True)
-class Trajectory:
-    """One rollout: a sequence of ``(Completion, Step)`` transitions.
-
-    Single-turn tasks produce trajectories with one transition. The
-    Completion carries the generator's response-side metadata; the Step
-    carries the env's reward and done flag;
-    """
-
-    sample_idx: int
-    prompt_token_ids: list[int]
-    transitions: list[tuple[Completion, Step]]
-
-    @property
-    def total_reward(self) -> float:
-        return sum(s.reward for _, s in self.transitions)
-
-
+# TODO: rename `Episode` -> `TrainingSample`
+# and `rollout_to_episode` -> `rollout_to_training_sample`
 @dataclass(kw_only=True, slots=True)
 class Episode:
-    """Training sample: flattened trajectory + GRPO advantage.
-
-    Flat shape (rather than composition) because the trainer collate
-    path and logging read these fields directly.
-    """
+    """Training sample: flattened Rollout turns + GRPO advantage,
+    ready for collation into a batch."""
 
     policy_version: int
-    prompt_idx: int
+    sample_id: str
     prompt_token_ids: list[int]
-    text: str
-    token_ids: list[int]
-    token_logprobs: list[float]
+    completion_text: str
+    completion_token_ids: list[int]
+    completion_logprobs: list[float]
     reward: float
     advantage: float
 
@@ -97,6 +52,8 @@ class TrainingBatch:
     token_ids: torch.Tensor  # [B, L]
     labels: torch.Tensor  # [B, L]
     positions: torch.Tensor  # [B, L]
+    # TODO(naming): rename generator_logprobs -> old_logprobs (PPO π_old) vs policy_logprobs,
+    # incl. GRPOLoss/trainer/batcher.
     generator_logprobs: torch.Tensor  # [B, L]
     loss_mask: torch.Tensor  # [B, L]
     advantages: torch.Tensor  # [B, L]
@@ -104,7 +61,7 @@ class TrainingBatch:
 
 @dataclass(frozen=True, slots=True)
 class OptimStepOutput:
-    """Result returned by ``PolicyTrainer.optim_step`` to the controller."""
+    """Result returned by `PolicyTrainer.optim_step` to the controller."""
 
     policy_version: int
     metrics: dict[str, float]