Add replay-backed SFT path

Author: tim0120

configs/debug/training_modes/README.md

@@ -10,6 +10,7 @@ Minimal end-to-end configs for the three training modes (`rl` / `opd` / `sft`) a
 | `sft.toml` | `sft` | local vLLM (`Qwen3-0.6B-Reverse-Text-RL`) | |
 | `sft_lora.toml` | `sft` | local vLLM (`Qwen3-0.6B-Reverse-Text-RL`) | trains a LoRA adapter (rank 8) |
 | `sft_external.toml` | `sft` | PI inference (`openai/gpt-5-mini`) | external OAI endpoint; no local teacher |
+| `sft_replay.toml` | `sft` | none | replays saved message traces through `sft-replay` |
 
 The student inference server is auto-launched on GPU 0 at `http://localhost:8000/v1` with `gpu_memory_utilization=0.5`. The local teacher (used by everything except `rl.toml` and `sft_external.toml`) is **not** auto-launched — start it manually on GPU 1.
 
@@ -42,6 +43,9 @@ uv run rl @ configs/debug/training_modes/sft_lora.toml
 # SFT hard distill from openai/gpt-5-mini via PI inference
 # (requires PRIME_API_KEY + PRIME_TEAM_ID in env; no local teacher needed)
 uv run rl @ configs/debug/training_modes/sft_external.toml
+
+# SFT from replayed dataset traces (no teacher)
+uv run rl @ configs/debug/training_modes/sft_replay.toml
 ```
 
 See [docs/training.md](../../docs/training.md#training-modes-rl--opd--sft-via-orchestrator) for what each mode does.

configs/debug/training_modes/sft_replay.toml

@@ -0,0 +1,41 @@
+# Static trace SFT through the RL orchestrator. No teacher server is needed:
+# sft-replay turns dataset message rows into replayed rollout trajectories.
+
+max_steps = 20
+seq_len = 2048
+
+[model]
+name = "PrimeIntellect/Qwen3-0.6B-Reverse-Text-SFT"
+
+[wandb]
+project = "reverse-text-debug"
+name = "debug-sft-replay"
+
+[orchestrator]
+training_mode = "sft"
+batch_size = 128
+group_size = 1
+
+[[orchestrator.train.env]]
+id = "sft-replay"
+
+[orchestrator.train.env.args.taskset]
+dataset = "PrimeIntellect/Reverse-Text-SFT"
+
+[orchestrator.eval]
+interval = 1
+num_examples = 128
+
+[orchestrator.eval.sampling]
+max_completion_tokens = 128
+
+[[orchestrator.eval.env]]
+id = "reverse-text"
+
+[trainer.optim]
+lr = 3e-6
+
+[ckpt]
+
+[inference]
+gpu_memory_utilization = 0.5

deps/verifiers

@@ -20,6 +20,31 @@
 from prime_rl.configs.trainer import TokenizerConfig
 from prime_rl.utils.config import BaseConfig
 
+SFT_REPLAY_ENV_ID = "sft-replay"
+
+
+def _is_sft_replay_env_id(env_id: str) -> bool:
+    stripped = env_id.split("@")[0]
+    return stripped == SFT_REPLAY_ENV_ID or stripped.endswith(f"/{SFT_REPLAY_ENV_ID}")
+
+
+def _sft_replay_dataset_arg(env_args: dict) -> object | None:
+    taskset = env_args.get("taskset")
+    if isinstance(taskset, dict):
+        dataset = taskset.get("dataset")
+        if dataset:
+            return dataset
+
+    config = env_args.get("config")
+    if isinstance(config, dict):
+        taskset = config.get("taskset")
+        if isinstance(taskset, dict):
+            dataset = taskset.get("dataset")
+            if dataset:
+                return dataset
+
+    return None
+
 
 class OptimizerConfig(BaseConfig):
     lr: float = Field(1e-4, ge=0)
@@ -501,13 +526,13 @@ class RolloutModelConfig(BaseConfig):
 
 class OrchestratorConfig(BaseConfig):
     training_mode: Literal["rl", "opd", "sft"] = "rl"
-    """Training mode. ``rl``: student generates rollouts, no teacher. ``opd``: student generates rollouts, teacher computes logprobs (teacher_tau > 0). ``sft``: teacher generates rollouts, student inference pool used for evals and weight sync."""
+    """Training mode. ``rl``: student generates rollouts, no teacher. ``opd``: student generates rollouts, teacher computes logprobs. ``sft``: teacher generates rollouts when configured, otherwise train envs must provide replayed traces."""
 
     student: RolloutModelConfig = Field(RolloutModelConfig(), validation_alias=AliasChoices("student", "model"))
     """Student rollout participant (model + client) — the model being trained."""
 
     teacher: RolloutModelConfig | None = Field(None, validation_alias=AliasChoices("teacher", "teacher_model"))
-    """Teacher rollout participant (model + client). Role depends on ``training_mode``: ``opd`` — teacher computes logprobs; ``sft`` — teacher generates rollouts."""
+    """Teacher rollout participant (model + client). Required for ``opd``. Optional for ``sft`` when train envs provide replayed traces."""
 
     train: TrainConfig = TrainConfig()
 
@@ -752,10 +777,16 @@ def validate_unique_filter_types(self):
                 )
         return self
 
+    @model_validator(mode="after")
+    def _drop_default_sft_zero_advantage_filter(self):
+        if self.training_mode == "sft" and "post_batch_filters" not in self.model_fields_set:
+            self.post_batch_filters = [f for f in self.post_batch_filters if f.type != "zero_advantage"]
+        return self
+
     @model_validator(mode="after")
     def _force_no_renderer_for_sft(self):
-        """SFT rolls out via the teacher's plain chat-completions endpoint; the
-        renderer client doesn't apply. Force ``renderer=None`` so the user
+        """SFT train rollouts use teacher chat completions or replayed traces;
+        the renderer client doesn't apply. Force ``renderer=None`` so the user
         doesn't have to remember to set it. Declared before the renderer
         validators below so they see the corrected value."""
         if self.training_mode == "sft":
@@ -768,8 +799,30 @@ def validate_training_mode(self):
         has_teacher = self.teacher is not None
         if self.training_mode == "rl" and has_teacher:
             raise ValueError("orchestrator.teacher must not be set when training_mode = 'rl'.")
-        if self.training_mode in ("opd", "sft") and not has_teacher:
-            raise ValueError(f"orchestrator.teacher must be configured when training_mode = '{self.training_mode}'.")
+        if self.training_mode == "opd" and not has_teacher:
+            raise ValueError("orchestrator.teacher must be configured when training_mode = 'opd'.")
+        return self
+
+    @model_validator(mode="after")
+    def validate_teacherless_sft_uses_sft_replay(self):
+        """Teacherless SFT is only valid when train envs replay existing data."""
+        if self.training_mode != "sft" or self.teacher is not None:
+            return self
+
+        non_replay_envs = [env.id for env in self.train.env if not _is_sft_replay_env_id(env.id)]
+        if non_replay_envs:
+            raise ValueError(
+                "orchestrator.teacher must be configured for SFT unless every train env uses "
+                f"{SFT_REPLAY_ENV_ID!r}; got non-replay train env(s): {non_replay_envs}."
+            )
+
+        missing_dataset = [env.resolved_name for env in self.train.env if _sft_replay_dataset_arg(env.args) is None]
+        if missing_dataset:
+            raise ValueError(
+                f"teacherless SFT with {SFT_REPLAY_ENV_ID!r} requires an explicit "
+                "env.args.taskset.dataset or env.args.config.taskset.dataset for "
+                f"each train env; missing for: {missing_dataset}."
+            )
         return self
 
     @model_validator(mode="after")

packages/prime-rl-configs/src/prime_rl/configs/orchestrator.py

@@ -88,6 +88,7 @@ envs = [
     "rlm-swe",
     "science-env",
     "simpleqa-verified",
+    "sft-replay",
     "tau2-bench",
     "wiki-search",
     "wordle",
@@ -150,6 +151,8 @@ override-dependencies = [
     "transformers==5.6.2",
     "torch>=2.9.0",
     "openenv-core",
+    "harnesses>=0.1.0",
+    "tasksets>=0.1.0",
 ]
 
 # ModelExpress 0.3.0 publishes protobuf<6 metadata, but its generated proto is
@@ -224,6 +227,9 @@ reverse-text = { path = "deps/verifiers/environments/reverse_text", editable = t
 rlm-swe = { path = "deps/research-environments/environments/rlm_swe", editable = true }
 science-env = { path = "deps/research-environments/environments/science_env", editable = true }
 simpleqa-verified = { path = "deps/research-environments/environments/simpleqa_verified", editable = true }
+sft-replay = { path = "deps/verifiers/environments/sft_replay", editable = true }
+harnesses = { path = "deps/verifiers/packages/harnesses", editable = true }
+tasksets = { path = "deps/verifiers/packages/tasksets", editable = true }
 tau2-bench = { path = "deps/research-environments/environments/tau2_bench", editable = true }
 wiki-search = { path = "deps/verifiers/environments/wiki_search", editable = true }
 wordle = { path = "deps/verifiers/environments/wordle", editable = true }

pyproject.toml

@@ -135,8 +135,9 @@ def __init__(
         self.policy = policy
         self.train_envs = train_envs
         self.eval_envs = eval_envs
-        # Train rollouts go to ``inference`` (the teacher in SFT mode);
-        # eval always evaluates the student, so it uses ``eval_inference``.
+        # Train rollouts go to ``inference`` (teacher in teacher-SFT, student
+        # otherwise); eval always evaluates the student, so it uses
+        # ``eval_inference``.
         self.inference = inference
         self.eval_inference = eval_inference
         self.train_source = train_source
@@ -173,9 +174,9 @@ def __init__(
 
     @property
     def train_model_name(self) -> str:
-        """Model name for *train* rollouts. In SFT mode train data comes from
-        the teacher pool, so use its model name; otherwise the live student
-        policy. (Eval always uses ``policy.model_name`` — the student.)"""
+        """Model name for *train* rollouts. Teacher-SFT uses the teacher pool
+        name; replay SFT receives the student name but ignores it. Eval always
+        uses ``policy.model_name`` — the student."""
         if self.training_mode == "sft":
             return self.inference.model_name
         return self.policy.model_name

src/prime_rl/orchestrator/dispatcher.py

@@ -345,10 +345,10 @@ async def setup(self) -> None:
         else:
             get_logger().info("Training from scratch")
 
-        # SFT generates rollouts via the teacher (the student is trained on
-        # the teacher's outputs); RL / OPD generate via the student
-        if config.training_mode == "sft":
-            assert self.teacher_inference is not None, "sft mode requires teacher inference"
+        # SFT train rollouts come from the teacher when configured. Teacherless
+        # SFT is validated at config parse time to use replay envs, which ignore
+        # the client/model passed by the dispatcher.
+        if config.training_mode == "sft" and self.teacher_inference is not None:
             rollout_inference = self.teacher_inference
         else:
             rollout_inference = self.student_inference

src/prime_rl/orchestrator/orchestrator.py

@@ -164,6 +164,24 @@ def _tokenize_step_with_renderer(
     return build_trajectory_step(renderer, prompt, completion, tools=tools)
 
 
+def _set_token_usage_from_trajectory(output: vf.RolloutOutput) -> None:
+    trajectory = output.get("trajectory") or []
+    tokenized_steps = [step for step in trajectory if step.get("tokens") is not None]
+    if not tokenized_steps:
+        return
+
+    prompt_tokens = [len(step["tokens"]["prompt_ids"]) for step in tokenized_steps]
+    completion_tokens = [len(step["tokens"]["completion_ids"]) for step in tokenized_steps]
+    total_completion = sum(completion_tokens)
+    last_total = prompt_tokens[-1] + completion_tokens[-1]
+    output["token_usage"] = {
+        "input_tokens": float(sum(prompt_tokens)),
+        "output_tokens": float(total_completion),
+        "final_input_tokens": float(max(0, last_total - total_completion)),
+        "final_output_tokens": float(total_completion),
+    }
+
+
 def backfill_rollout_tokens(
     output: vf.RolloutOutput,
     tokenizer: PreTrainedTokenizer,
@@ -175,6 +193,9 @@ def backfill_rollout_tokens(
     Otherwise falls back to the tokenizer + apply_chat_template path.
     """
     if all(step["tokens"] is not None for step in output["trajectory"]):
+        token_usage = output.get("token_usage") or {}
+        if "final_input_tokens" not in token_usage or "final_output_tokens" not in token_usage:
+            _set_token_usage_from_trajectory(output)
         return True
 
     logger = get_logger()
@@ -198,6 +219,10 @@ def backfill_rollout_tokens(
             reconstructed.pop("original_prompt_len")
             step["tokens"] = reconstructed
 
+    token_usage = output.get("token_usage") or {}
+    if "final_input_tokens" not in token_usage or "final_output_tokens" not in token_usage:
+        _set_token_usage_from_trajectory(output)
+
     return True
 
 

src/prime_rl/orchestrator/trajectories.py

@@ -0,0 +1,90 @@
+import pytest
+import verifiers as vf
+from verifiers.clients import Client
+
+from prime_rl.orchestrator.trajectories import backfill_rollout_tokens, interleave_rollout
+
+
+class NoopClient(Client):
+    def setup_client(self, config):
+        return object()
+
+    async def to_native_tool(self, tool):
+        raise AssertionError("sft-replay must not convert tools")
+
+    async def to_native_prompt(self, messages):
+        raise AssertionError("sft-replay must not render prompts through a client")
+
+    async def get_native_response(self, prompt, model, sampling_args, tools=None, **kwargs):
+        raise AssertionError("sft-replay must not request model responses")
+
+    async def raise_from_native_response(self, response) -> None:
+        raise AssertionError("sft-replay must not handle native responses")
+
+    async def from_native_response(self, response):
+        raise AssertionError("sft-replay must not parse native responses")
+
+    async def close(self) -> None:
+        return None
+
+
+class SimpleChatTokenizer:
+    def __init__(self):
+        self._tok2id: dict[str, int] = {}
+        self._next_id = 1
+
+    def _id(self, token: str) -> int:
+        if token not in self._tok2id:
+            self._tok2id[token] = self._next_id
+            self._next_id += 1
+        return self._tok2id[token]
+
+    def apply_chat_template(self, messages, add_generation_prompt=False, return_dict=False, tools=None):
+        del return_dict, tools
+        ids = []
+        for message in messages:
+            role = message.get("role", "unknown")
+            ids.append(self._id(f"<|{role}|>"))
+            content = message.get("content", "")
+            if isinstance(content, str):
+                if content:
+                    ids.append(self._id(content))
+            else:
+                ids.append(self._id(str(content)))
+        if add_generation_prompt:
+            ids.append(self._id("<|assistant|>"))
+        return ids
+
+
+def role_content(messages) -> list[tuple[str, object]]:
+    return [(message["role"], message["content"]) for message in messages]
+
+
+@pytest.mark.asyncio
+async def test_sft_replay_env_replays_messages_for_prime_rl_training_path():
+    env = vf.load_environment("sft-replay", taskset={})
+    row = dict(env.get_dataset()[0])
+
+    output = await env.run_rollout(
+        row,
+        client=NoopClient(vf.ClientConfig()),
+        model="unused-student",
+        sampling_args={},
+        state_columns=["trajectory", "sampling_args"],
+    )
+
+    assert output["error"] is None
+    assert output["stop_condition"] == "replayed_messages"
+    assert len(output["trajectory"]) == 1
+    assert output["trajectory"][0]["tokens"] is None
+    assert role_content(output["trajectory"][0]["prompt"]) == [("user", "Reverse abc.")]
+    assert role_content(output["trajectory"][0]["completion"]) == [("assistant", "cba")]
+
+    backfill_rollout_tokens(output, SimpleChatTokenizer())
+    samples = interleave_rollout(output, env_name="sft-replay")
+
+    assert samples is not None
+    assert len(samples) == 1
+    assert any(samples[0].completion_mask)
+    assert output["token_usage"]["final_input_tokens"] > 0
+    assert output["token_usage"]["final_output_tokens"] > 0