[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):
     """
@@ -411,6 +414,7 @@ async def generate(
                 top_p=_sampling_config.top_p,
                 max_tokens=_sampling_config.max_tokens,
                 n=_sampling_config.n,
+                stop_token_ids=list(_sampling_config.stop_token_ids) or None,
                 seed=self.config.debug.seed,
                 logprobs=1,
                 output_kind=RequestOutputKind.FINAL_ONLY,
@@ -457,7 +461,6 @@ async def generate(
                         Completion(
                             policy_version=self.policy_version,
                             prompt_idx=prompt_idx,
-                            text=sample.text,
                             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 SumDigitsDataset, SumDigitsTask
 from torchtitan.models.qwen3 import model_registry
 
 
@@ -39,10 +40,10 @@ 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
-        ),
+        train_dataset=SumDigitsDataset.Config(seed=42),
+        validation_dataset=SumDigitsDataset.Config(seed=99),
+        tasks={"sum_digits": SumDigitsTask.Config()},
+        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),
@@ -80,7 +81,7 @@ def rl_grpo_qwen3_0_6b() -> RLTrainer.Config:
                 n=group_size,
                 temperature=0.8,
                 top_p=0.95,
-                max_tokens=100,
+                max_tokens=700,
             ),
         ),
     )
@@ -96,10 +97,10 @@ 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
-        ),
+        train_dataset=SumDigitsDataset.Config(seed=42),
+        validation_dataset=SumDigitsDataset.Config(seed=99),
+        tasks={"sum_digits": SumDigitsTask.Config()},
+        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),
@@ -138,7 +139,7 @@ def rl_grpo_qwen3_1_7b() -> RLTrainer.Config:
                 n=group_size,
                 temperature=0.8,
                 top_p=0.95,
-                max_tokens=100,
+                max_tokens=700,
             ),
         ),
     )
@@ -154,10 +155,10 @@ 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
-        ),
+        train_dataset=SumDigitsDataset.Config(seed=42),
+        validation_dataset=SumDigitsDataset.Config(seed=99),
+        tasks={"sum_digits": SumDigitsTask.Config()},
+        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),
@@ -195,14 +196,14 @@ def rl_grpo_qwen3_14b() -> RLTrainer.Config:
                 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 +216,10 @@ 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
-        ),
+        train_dataset=SumDigitsDataset.Config(seed=42),
+        validation_dataset=SumDigitsDataset.Config(seed=99),
+        tasks={"sum_digits": SumDigitsTask.Config()},
+        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),
@@ -260,7 +261,7 @@ def rl_grpo_qwen3_0_6b_batch_invariant() -> RLTrainer.Config:
                 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,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.env_types.message_env import (
+    MessageEnv,
+    ResetOutput,
+    StepOutput,
+)
+from torchtitan.experiments.rl.env_types.renderer_env import (
+    RendererEnv,
+    RendererEnvConfig,
+    TokenizedStepOutput,
+)
+
+__all__ = [
+    "RendererEnvConfig",
+    "MessageEnv",
+    "ResetOutput",
+    "StepOutput",
+    "RendererEnv",
+    "TokenizedStepOutput",
+]

torchtitan/experiments/rl/env_types/message_env.py

@@ -0,0 +1,86 @@
+# 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
+
+
+@dataclass(kw_only=True, slots=True)
+class ResetOutput:
+    """Initial messages + tool specs from `MessageEnv.reset`."""
+
+    messages: list[Message]  # [M_initial]
+    """The set of messages that form an initial prompt for the model."""
+
+    tools: list[ToolSpec] = field(default_factory=list)  # [K_tools]
+    """Tool schemas exposed to the model. Empty for tool-less envs."""
+
+
+@dataclass(kw_only=True, slots=True)
+class StepOutput:
+    """Env response to one parsed assistant message.
+
+    `done` ends the conversation; `RendererEnv` maps it to a `RolloutStatus`
+    and detects truncation/errors. The env may attach `reward_components`
+    (e.g. tool-call success); the rubric decides how to use them.
+    """
+
+    messages: list[Message] = field(default_factory=list)  # [M_env]
+    """Env-appended messages (tool / user replies). Empty when the rollout
+    terminates with no follow-up."""
+
+    done: bool = False
+    """`True` ends the rollout."""
+
+    reward_components: 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 model
+        if any(m.get("role") == "assistant" for m in self.messages):
+            raise ValueError("StepOutput.messages may not contain assistant messages")
+
+
+class MessageEnv(abc.ABC):
+    """User-written env in message space. Implement `reset` + `step_message`;
+    `RendererEnv` wraps it with a `Renderer` for tokens.
+
+    Example:
+
+        class SumDigitsEnv(MessageEnv):
+            async def reset(self) -> ResetOutput:
+                return ResetOutput(messages=[{"role": "user", "content": "sum [1, 2]"}])
+
+            async def step_message(self, msg: Message) -> StepOutput:
+                return StepOutput(done=True)        # single-turn; rubric scores later
+    """
+
+    @abc.abstractmethod
+    async def reset(self) -> ResetOutput:
+        """Return the initial conversation + tools for prompt rendering."""
+
+    @abc.abstractmethod
+    async def step_message(self, msg: Message) -> StepOutput:
+        """Return the env's reply to an assistant message.
+
+        `RendererEnv` handles finish_reason / length / parse failures before
+        calling this, so the env only sees a successfully parsed message.
+
+        Args:
+            msg: Parsed assistant message (`content`, optional
+                `reasoning_content`, optional `tool_calls`).
+
+        Returns:
+            `StepOutput` with env reply messages.
+        """
+
+    async def close(self) -> None:
+        """Release env-owned resources. Default no-op; idempotent."""