[RL] - MessageEnv, Rollout types, Rubric, Renderer

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==0.4.1
 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:

torchtitan/experiments/rl/actors/generator.py

@@ -163,6 +163,9 @@ class SamplingConfig:
     max_tokens: int = 100
     """Maximum number of tokens to generate per completion."""
 
+    stop_token_ids: list[int] = field(default_factory=list)
+    """Role-boundary stop tokens from the renderer (e.g. Qwen3 `<|im_end|>`)."""
+
 
 class VLLMGenerator(Actor, Configurable):
     """
@@ -379,17 +382,20 @@ 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. If ``SamplingConfig.n > 1``, completions for one prompt are
+        contiguous in vLLM's sample order.
 
         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).
@@ -411,21 +417,31 @@ async def generate(
                 top_p=_sampling_config.top_p,
                 max_tokens=_sampling_config.max_tokens,
                 n=_sampling_config.n,
+                stop_token_ids=_sampling_config.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
             # high-level API to stay resilient to vLLM internal changes.
             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,
                 )
@@ -435,15 +451,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)
                 )
@@ -456,8 +471,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,

torchtitan/experiments/rl/actors/trainer.py

@@ -455,7 +455,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()

torchtitan/experiments/rl/config_registry.py

@@ -25,7 +25,8 @@
 from torchtitan.experiments.rl.batcher import BatchConfig, Batcher
 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.experiments.rl.tasks.sum_digits import SumDigitsTask
 from torchtitan.models.qwen3 import model_registry
 
 
@@ -39,10 +40,9 @@ def rl_grpo_qwen3_0_6b() -> RLTrainer.Config:
         num_prompts_per_step=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
-        ),
+        task=SumDigitsTask.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),
@@ -77,10 +77,9 @@ def rl_grpo_qwen3_0_6b() -> 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,
             ),
         ),
     )
@@ -96,10 +95,9 @@ def rl_grpo_qwen3_1_7b() -> RLTrainer.Config:
         num_prompts_per_step=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
-        ),
+        task=SumDigitsTask.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),
@@ -135,10 +133,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,
             ),
         ),
     )
@@ -154,10 +151,9 @@ def rl_grpo_qwen3_14b() -> RLTrainer.Config:
         num_prompts_per_step=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
-        ),
+        task=SumDigitsTask.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),
@@ -192,17 +188,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.
     """
@@ -215,10 +210,9 @@ def rl_grpo_qwen3_0_6b_batch_invariant() -> RLTrainer.Config:
         num_prompts_per_step=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
-        ),
+        task=SumDigitsTask.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),
@@ -257,10 +251,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,
         ),

torchtitan/experiments/rl/env_types/__init__.py

@@ -0,0 +1,23 @@
+# 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.env_types.message_env import (
+    MessageEnv,
+    MessageResetOutput,
+    MessageStepOutput,
+)
+from torchtitan.experiments.rl.env_types.renderer_env import (
+    RendererWrapperEnv,
+    TokenizedStepOutput,
+)
+
+__all__ = [
+    "MessageEnv",
+    "MessageResetOutput",
+    "MessageStepOutput",
+    "RendererWrapperEnv",
+    "TokenizedStepOutput",
+]

torchtitan/experiments/rl/env_types/message_env.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 MessageResetOutput:
+    """Initial prompt messages + tool specs from `MessageEnv.reset`."""
+
+    prompt_messages: list[Message]  # [M_prompt]
+    """The messages that form the initial prompt (e.g. [system, user])."""
+
+    tools: list[ToolSpec] = field(default_factory=list)  # [K_tools]
+    """Tool schemas exposed to the assistant. Empty for tool-less envs."""
+
+
+@dataclass(kw_only=True, slots=True)
+class MessageStepOutput:
+    """The env's reply to the assistant's turn."""
+
+    env_messages: list[Message] = field(default_factory=list)  # [M_env]
+    """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(
+                "MessageStepOutput.env_messages may not contain assistant messages"
+            )
+
+
+class MessageEnv(Configurable, abc.ABC):
+    """User-written env in message space. Implement `reset` + `step`.
+
+    Tip: `MessageEnv` works in messages and never sees token ids; You can have `RendererWrapperEnv`
+    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 reset(self) -> MessageResetOutput:
+                return MessageResetOutput(
+                    prompt_messages=[{"role": "user", "content": f"What is {self._expression}?"}],
+                    tools=[CALCULATOR_TOOL],
+                )
+
+            async def step(self, assistant_message: Message) -> MessageStepOutput:
+                tool_calls = assistant_message.get("tool_calls")
+                if not tool_calls:
+                    return MessageStepOutput(done=True)  # assistant gave its final answer
+                result = run_calculator(tool_calls[0])
+                return MessageStepOutput(
+                    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 reset(self) -> MessageResetOutput:
+        """Return the initial conversation + tools for prompt rendering."""
+
+    @abc.abstractmethod
+    async def step(self, assistant_message: Message) -> MessageStepOutput:
+        """Advance the env one turn given the assistant's latest message.
+
+        `RendererWrapperEnv` parses the completion and handles
+        finish_reason / length / parse / timeout failures before calling this,
+        so the env only sees a well-formed assistant message.
+
+        Args:
+            assistant_message: the assistant's parsed turn.
+
+        Returns:
+            `MessageStepOutput` with the env's reply messages.
+        """
+
+    async def close(self) -> None:
+        """Release env-owned resources. Default no-op; idempotent."""