Echo

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

@@ -206,6 +206,72 @@ def resolve_timeout(self):
         return self
 
 
+class SystemRoleEchoConfig(BaseConfig):
+    """Echo supervision for system-message content tokens."""
+
+    alpha: float = Field(1.0, allow_inf_nan=False)
+    """Per-token echo weight."""
+
+
+class UserRoleEchoConfig(BaseConfig):
+    """Echo supervision for user-message content tokens."""
+
+    alpha: float = Field(1.0, allow_inf_nan=False)
+    """Per-token echo weight."""
+
+
+class AssistantRoleEchoConfig(BaseConfig):
+    """Echo supervision for assistant-message content and completion tokens."""
+
+    alpha: float = Field(1.0, allow_inf_nan=False)
+    """Per-token echo weight. ``alpha=0`` keeps the token supervised but gives it zero gradient."""
+
+
+class ToolRoleEchoConfig(BaseConfig):
+    """Echo supervision for tool-message content tokens."""
+
+    alpha: float = Field(1.0, allow_inf_nan=False)
+    """Per-token echo weight."""
+
+    tool_names: set[str] | None = Field(None, min_length=1)
+    """Restrict echo to these tool function names; None = all tools."""
+
+
+class EchoFilterConfig(BaseConfig):
+    """Optional callable that narrows role-selected echo tokens per rollout."""
+
+    import_path: str
+    """Dotted import path to the filter callable, e.g. ``"my_module.filter_warnings"``."""
+
+    kwargs: dict[str, Any] = Field(default_factory=dict)
+    """Keyword arguments forwarded to the filter as ``**kwargs``."""
+
+
+class EchoConfig(BaseConfig):
+    """Enable CE echo on selected message roles for this training env."""
+
+    system: SystemRoleEchoConfig | None = None
+    """System-message echo (default: disabled)."""
+
+    user: UserRoleEchoConfig | None = None
+    """User-message echo (default: disabled)."""
+
+    assistant: AssistantRoleEchoConfig | None = None
+    """Assistant-message echo (default: disabled)."""
+
+    tool: ToolRoleEchoConfig | None = None
+    """Tool-message echo (default: disabled)."""
+
+    filter: EchoFilterConfig | None = None
+    """Optional per-token filter on top of the role baseline."""
+
+    @model_validator(mode="after")
+    def validate_roles(self) -> "EchoConfig":
+        if self.system is self.user is self.assistant is self.tool is None:
+            raise ValueError("EchoConfig requires at least one of system, user, assistant, or tool.")
+        return self
+
+
 class TrainEnvConfig(EnvConfig):
     sampling: TrainSamplingConfig = TrainSamplingConfig()
     """Per-env sampling overrides. Unset fields inherit from the group-level train sampling config."""
@@ -214,6 +280,9 @@ class TrainEnvConfig(EnvConfig):
     """Rollouts generated per example for GRPO group-relative advantages.
     Inherits from ``orchestrator.group_size`` when unset."""
 
+    echo: EchoConfig | None = None
+    """Per-env per-role echo config."""
+
 
 class EvalEnvConfig(EnvConfig):
     sampling: EvalSamplingConfig = EvalSamplingConfig()

pyproject.toml

@@ -69,6 +69,7 @@ envs = [
     "deepdive",
     "general-agent",
     "gpqa",
+    "harnesses",
     "hle",
     "ifeval",
     "livecodebench",
@@ -77,7 +78,6 @@ envs = [
     "math-python",
     "math500",
     "mini-swe-agent-plus",
-    "mini-swe-agent-plus-rlm",
     "mmlu-pro",
     "opencode-cp",
     "opencode-deepdive",
@@ -88,6 +88,7 @@ envs = [
     "rlm-swe",
     "science-env",
     "simpleqa-verified",
+    "tasksets",
     "tau2-bench",
     "wiki-search",
 ]
@@ -197,6 +198,8 @@ prime-rl-configs = { path = "packages/prime-rl-configs", editable = true }
 verifiers = { path = "deps/verifiers", editable = true }
 renderers = { path = "deps/renderers", editable = true }
 prime-pydantic-config = { path = "deps/pydantic-config", editable = true }
+harnesses = { path = "deps/verifiers/packages/harnesses", editable = true }
+tasksets = { path = "deps/verifiers/packages/tasksets", editable = true }
 aime2024 = { path = "deps/research-environments/environments/aime2024", editable = true }
 aime2025 = { path = "deps/research-environments/environments/aime2025", editable = true }
 alphabet-sort = { path = "deps/verifiers/environments/alphabet_sort", editable = true }
@@ -213,7 +216,6 @@ math-env = { path = "deps/research-environments/environments/math_env", editable
 math-python = { path = "deps/verifiers/environments/math_python", editable = true }
 math500 = { path = "deps/research-environments/environments/math500", editable = true }
 mini-swe-agent-plus = { path = "deps/research-environments/environments/mini_swe_agent_plus", editable = true }
-mini-swe-agent-plus-rlm = { path = "deps/research-environments/environments/mini_swe_agent_plus_rlm", editable = true }
 mmlu-pro = { path = "deps/research-environments/environments/mmlu_pro", editable = true }
 opencode-cp = { path = "deps/research-environments/environments/opencode_cp", editable = true }
 opencode-deepdive = { path = "deps/research-environments/environments/opencode_deepdive", editable = true }

src/prime_rl/orchestrator/echo.py

@@ -0,0 +1,138 @@
+from __future__ import annotations
+
+from collections.abc import Callable
+from dataclasses import dataclass
+
+import verifiers as vf
+
+from prime_rl.configs.orchestrator import EchoConfig
+
+
+@dataclass(frozen=True)
+class EchoAnnotations:
+    step_alpha: list[list[float | None]]
+
+    def initial_sample_alpha(self, step_idx: int) -> list[float | None] | None:
+        alpha = self.step_alpha[step_idx]
+        return list(alpha) if any(a is not None for a in alpha) else None
+
+    def extension_alpha(self, step_idx: int, prefix_len: int, prompt_len: int) -> list[float | None]:
+        alpha = self.step_alpha[step_idx]
+        return alpha[prefix_len:prompt_len] + alpha[prompt_len:]
+
+
+def build_echo_annotations(
+    rollout: vf.RolloutOutput,
+    echo_config: EchoConfig | None,
+    filter_fn: Callable[..., list[list[bool]]] | None = None,
+) -> EchoAnnotations | None:
+    if echo_config is None:
+        return None
+
+    trajectory = rollout["trajectory"]
+    step_tokens = []
+    for step in trajectory:
+        tokens = step["tokens"]
+        if tokens is None:
+            return None
+        step_tokens.append(tokens)
+
+    filter_masks = apply_echo_filter(rollout, filter_fn) if filter_fn is not None and trajectory else None
+    return EchoAnnotations(
+        step_alpha=[
+            _build_step_echo_alpha(
+                prompt_attribution=tokens.get("prompt_attribution"),
+                prompt_len=len(tokens["prompt_ids"]),
+                completion_len=len(tokens["completion_ids"]),
+                echo_config=echo_config,
+                filter_mask=filter_masks[step_idx] if filter_masks is not None else None,
+            )
+            for step_idx, tokens in enumerate(step_tokens)
+        ]
+    )
+
+
+def _build_step_echo_alpha(
+    prompt_attribution: dict | None,
+    prompt_len: int,
+    completion_len: int,
+    echo_config: EchoConfig | None,
+    filter_mask: list[bool] | None = None,
+) -> list[float | None]:
+    expected_total_len = prompt_len + completion_len
+    out: list[float | None] = [None] * expected_total_len
+    if echo_config is not None:
+        if echo_config.assistant is not None:
+            out[prompt_len:expected_total_len] = [echo_config.assistant.alpha] * completion_len
+
+        if prompt_attribution is not None:
+            message_roles = prompt_attribution.get("message_roles")
+            message_indices = prompt_attribution.get("message_indices")
+            is_content = prompt_attribution.get("is_content")
+            if message_roles is not None and is_content and message_indices:
+                if len(is_content) == prompt_len and len(message_indices) == prompt_len:
+                    role_alphas = {
+                        "system": echo_config.system.alpha if echo_config.system is not None else None,
+                        "user": echo_config.user.alpha if echo_config.user is not None else None,
+                        "assistant": echo_config.assistant.alpha if echo_config.assistant is not None else None,
+                    }
+                    tool_config = echo_config.tool
+                    tool_alpha = tool_config.alpha if tool_config is not None else None
+                    enabled_tools = tool_config.tool_names if tool_config is not None else None
+                    message_tool_names = prompt_attribution.get("message_tool_names") or []
+
+                    for k, mi in enumerate(message_indices):
+                        if mi < 0 or not is_content[k] or mi >= len(message_roles):
+                            continue
+                        role = message_roles[mi]
+                        if role == "tool":
+                            tool_name = message_tool_names[mi] if mi < len(message_tool_names) else None
+                            if tool_alpha is not None and (enabled_tools is None or tool_name in enabled_tools):
+                                out[k] = tool_alpha
+                            continue
+
+                        alpha = role_alphas.get(role)
+                        if alpha is not None:
+                            out[k] = alpha
+
+    if filter_mask is not None:
+        out = [alpha if keep else None for alpha, keep in zip(out, filter_mask, strict=True)]
+
+    return out
+
+
+def apply_echo_filter(
+    rollout: vf.RolloutOutput,
+    filter_fn: Callable[..., list[list[bool]]],
+) -> list[list[bool]]:
+    trajectory = rollout["trajectory"]
+    result = filter_fn(rollout)
+
+    if not isinstance(result, list):
+        raise TypeError(f"echo filter must return list[list[bool]], got {type(result).__name__}")
+    if len(result) != len(trajectory):
+        raise ValueError(
+            f"echo filter returned {len(result)} per-step masks but the rollout has {len(trajectory)} trajectory steps"
+        )
+
+    for step_idx, (step, mask) in enumerate(zip(trajectory, result)):
+        tokens = step["tokens"]
+        prompt_len = len(tokens["prompt_ids"])
+        completion_len = len(tokens["completion_ids"])
+        expected = prompt_len + completion_len
+
+        if not isinstance(mask, list):
+            raise TypeError(f"echo filter step {step_idx}: mask must be a list, got {type(mask).__name__}")
+        if len(mask) != expected:
+            raise ValueError(
+                f"echo filter step {step_idx}: mask length {len(mask)} "
+                f"!= expected {expected} "
+                f"(prompt_len={prompt_len}, completion_len={completion_len})"
+            )
+        for k, v in enumerate(mask):
+            if type(v) is not bool:
+                raise TypeError(
+                    f"echo filter step {step_idx}: mask[{k}] must be a plain bool, got {type(v).__name__} ({v!r})"
+                )
+
+    return result

src/prime_rl/orchestrator/envs.py

@@ -2,6 +2,7 @@
 
 import asyncio
 import atexit
+import functools
 import multiprocessing as mp
 import time
 from collections.abc import Awaitable, Callable, Iterator, Sequence
@@ -18,7 +19,7 @@
 from prime_rl.orchestrator.eval_utils import compute_pass_at_k
 from prime_rl.utils.logger import ProgressTracker, get_logger
 from prime_rl.utils.monitor import get_monitor
-from prime_rl.utils.utils import capitalize
+from prime_rl.utils.utils import capitalize, import_object
 
 REQUIRED_STATE_COLUMNS = ["trajectory"]
 
@@ -170,6 +171,10 @@ class TrainEnv(Env):
     def __init__(self, config: TrainEnvConfig):
         super().__init__(config)
         self.sampling_args = config.sampling.to_sampling_args()
+        self.echo_filter_fn: Callable[..., list[list[bool]]] | None = None
+        if config.echo is not None and config.echo.filter is not None:
+            fn = import_object(config.echo.filter.import_path)
+            self.echo_filter_fn = functools.partial(fn, **config.echo.filter.kwargs)
 
     def get_dataset(self, seed: int | None = None):
         return self.env.get_dataset(seed=seed)

src/prime_rl/orchestrator/train_sink.py

@@ -19,6 +19,7 @@
 
 from prime_rl.configs.orchestrator import AdvantageConfig, OrchestratorConfig
 from prime_rl.orchestrator.advantage import assign_advantages, setup_advantage_fn
+from prime_rl.orchestrator.echo import build_echo_annotations
 from prime_rl.orchestrator.envs import TrainEnvs
 from prime_rl.orchestrator.filters import RolloutFilter, apply_filters
 from prime_rl.orchestrator.trajectories import (
@@ -160,11 +161,16 @@ async def process_rollout(self, rollout: TrainRollout) -> None:
         needs_backfill = any(s["tokens"] is None for s in raw.get("trajectory") or [])
         if needs_backfill:
             await asyncio.to_thread(backfill_rollout_tokens, raw, self.tokenizer, renderer=self.renderer)
+
+        env = self.train_envs.get(rollout.env_name)
+        echo_annotations = await asyncio.to_thread(build_echo_annotations, raw, env.config.echo, env.echo_filter_fn)
+
         samples = await asyncio.to_thread(
             interleave_rollout,
             raw,
             mm_token_type_ids_mapping=self.mm_token_type_ids_mapping,
             env_name=rollout.env_name,
+            echo_annotations=echo_annotations,
         )
         rollout.samples = samples or []
         # Offload base64 image bytes to disk as soon as the rollout is

src/prime_rl/orchestrator/trajectories.py

@@ -9,6 +9,7 @@
 import verifiers as vf
 from transformers.tokenization_utils import PreTrainedTokenizer
 
+from prime_rl.orchestrator.echo import EchoAnnotations
 from prime_rl.transport import RoutedExperts, TrainingSample
 from prime_rl.utils.chat_template import (
     common_prefix_len,
@@ -206,6 +207,7 @@ def interleave_rollout(
     mm_token_type_ids_mapping: dict[int, int] | None = None,
     *,
     env_name: str = "",
+    echo_annotations: EchoAnnotations | None = None,
 ) -> list[TrainingSample] | None:
     """
     Convert vf.RolloutOutput to trainable rollouts by interleaving trajectory steps
@@ -225,6 +227,12 @@ def interleave_rollout(
     For VLM models, each renderer-produced trajectory step carries its
     per-image processed tensors inline on ``multi_modal_data``; the last
     merged step's sidecar covers every image in the sample.
+
+    Args:
+        output: vf.RolloutOutput containing trajectory data
+        mm_token_type_ids_mapping: Maps prompt-token ids to mm_token_type_ids
+            (1 = image, 2 = video, 0 otherwise). Renderer-supplied.
+        echo_annotations: Optional per-step echo alpha annotations.
     """
     logger = get_logger()
 
@@ -238,6 +246,7 @@ def interleave_rollout(
         return None
 
     has_error = output["error"] is not None
+    # completion_temperatures is left empty; the train sink fills it per-env later.
 
     def prepare_step_tokens(step: vf.TrajectoryStep, step_idx: int) -> dict[str, Any] | None:
         tokens = step["tokens"]
@@ -308,6 +317,7 @@ def make_sample(tokens: dict[str, Any], step_idx: int) -> TrainingSample:
             env_name=env_name,
             mm_token_type_ids=None,
             routed_experts=None,  # deferred — finalized at end of interleave_rollout
+            echo_alpha=echo_annotations.initial_sample_alpha(step_idx) if echo_annotations is not None else None,
         )
         # Initialize routed-experts state for this sample. First chunk is the
         # raw step routed_experts (no pad, no copy). running_len is the
@@ -385,6 +395,15 @@ def extend_sample(
             sample.completion_mask.extend(tokens["completion_mask"])
         sample.completion_logprobs.extend(tokens["completion_logprobs"])
 
+        if echo_annotations is not None:
+            step_prompt_len = len(tokens["prompt_ids"])
+            extension = echo_annotations.extension_alpha(step_idx, prefix_len, step_prompt_len)
+            if any(a is not None for a in extension) or sample.echo_alpha is not None:
+                if sample.echo_alpha is None:
+                    existing_len = len(sample.prompt_ids) + len(sample.completion_ids) - len(extension)
+                    sample.echo_alpha = [None] * existing_len
+                sample.echo_alpha.extend(extension)
+
         step_routed = tokens.get("routed_experts")
         state = sample_routed_state.get(id(sample))
         if state is not None: