fix ODF logging and add min_abs_adv

CHANGELOG.md

@@ -2,6 +2,7 @@
 
 Documenting changes which affect configuration usage patterns (added/moved/removed/renamed fields, notable logic changes).
 
+- **`orchestrator.buffer.adv_filter`** and **`orchestrator.buffer.online_difficulty_filtering`**: Added `adv_filter` as the rollout advantage threshold used to decide which rollouts are sent to training (`float | None`, default: `None`). `online_difficulty_filtering` is now deprecated and kept only as an alias for `adv_filter = 0.0`; existing configs should migrate to `adv_filter` directly. Rollouts filtered out by `adv_filter` are still kept for orchestrator logging. `orchestrator.advantage.length_shaping_alpha` now requires `orchestrator.buffer.adv_filter` to be set, and `verification.enabled = false` cannot be combined with `adv_filter`. (2026-03-30)
 - **`wandb.log_extras.sample_ratio` / monitor sample logging defaults**: `wandb.log_extras.sample_ratio` is now actually applied to W&B sample-table logging via the shared monitor sampler (it was previously a no-op for WandB). Separately, the orchestrator no longer hard-caps sample logging to 8 rollouts before monitor-level sampling runs, so when monitor `sample_ratio` is `None`, monitors now receive and may log the full rollout batch for a step instead of at most 8 rollouts. This affects both W&B and Prime monitor sample logging behavior. (2026-03-27)
 - **`loss_impl = "quack_fused"` (SFT)**: Added `quack_fused` option for the SFT `loss_impl` field. Uses quack-kernels for chunked linear + cross-entropy with CuTe DSL CUDA kernels, avoiding full logits materialization. Requires `quack-kernels` package. Does not support Gemma logit softcapping. Custom model impl (`model.impl = "custom"`) also gains quack RMSNorm acceleration on CUDA automatically. (2026-03-26)
 - **`model.ac.targets`**: Removed the `mamba` selective activation checkpoint target and folded that behavior into the broader `linear_attn` target. Existing configs using `targets = ["mamba"]` will now fail validation at runtime with `ValueError`; switch them to `targets = ["linear_attn"]`. The `linear_attn` target description now covers supported token mixers outside the standard softmax-attention path, including NemotronH Mamba layers, Qwen3.5-MoE GatedDeltaNet layers, and AFMoE sliding-window attention layers. (2026-03-28)
@@ -12,7 +13,7 @@ Documenting changes which affect configuration usage patterns (added/moved/remov
 - **`model.optimization_dtype` / `model.reduce_dtype` (VLM models)**: Added validation that VLM models must use `optimization_dtype='bfloat16'` and `reduce_dtype='bfloat16'` to match vLLM inference. Previously valid configs with `float32` (the default) are now rejected for VLM model names. Set both fields to `"bfloat16"` when training VLMs. (2026-03-21)
 - **`model.ac.mode`** and **`model.ac.targets`**: Added selective activation checkpointing configuration. `model.ac.mode` accepts `full` (default) or `selective`. When `selective`, `model.ac.targets` selects subcomponents to checkpoint. Supported public targets are currently `norm`, `attention_sdpa`, `mla_up_proj`, and `routed_experts`; runtime validation remains the source of truth. `model.ac.targets` defaults to `["norm"]`, and selective mode requires at least one target. (2026-03-20)
 - **`orchestrator.advantage.length_weighted_mean`**: Removed. The default advantage now always uses the plain per-problem mean baseline unless `orchestrator.advantage.length_shaping_alpha` is set. Existing configs must delete this field. (2026-03-19)
-- **`orchestrator.advantage.length_shaping_alpha`**: Added Group Relative Reward Rescaling coefficient to the default advantage config. When set, applies length-based GR3 shaping during advantage computation and requires `orchestrator.buffer.online_difficulty_filtering = true` (default: `None`) (2026-03-18)
+- **`orchestrator.advantage.length_shaping_alpha`**: Added Group Relative Reward Rescaling coefficient to the default advantage config. When set, applies length-based GR3 shaping during advantage computation and requires `orchestrator.buffer.adv_filter` to be set (default: `None`) (2026-03-18)
 - **`prime_monitor.log_extras.sample_ratio`**: Added ratio-based rollout sampling (0.0–1.0, default: None). When set, caps the number of rollouts logged per step to `len(rollouts) * sample_ratio`. `None` preserves current behavior (log all rollouts). Interacts with existing `interval` gate which still runs first. (2026-03-12)
 - **`client.connect_timeout`**: Added configurable TCP connect timeout for inference API requests (default: 30.0s). Previously hardcoded to 5.0s. Helps with vLLM or cluster flakiness (2026-03-11)
 - **`model.fused_lm_head_token_chunk_size`**: Added as the fused LM-head chunking field for the token-chunked implementation. Unlike the removed `model.fused_lm_head_chunk_size`, this chunks over flattened sequence tokens rather than vocabulary rows. `model.fused_lm_head_chunk_size` is no longer accepted; switch configs to `model.fused_lm_head_token_chunk_size` explicitly. (2026-03-09)

configs/wiki_search/rl.toml

@@ -34,7 +34,7 @@ oversampling_factor = 2.0
 max_tokens = 512
 
 [orchestrator.buffer]
-online_difficulty_filtering = true
+min_abs_adv = 0.0
 
 [[orchestrator.env]]
 id = "wiki-search"

examples/Intellect-3.1/rl.toml

@@ -75,7 +75,7 @@ args = { pool_size = 512 }
 [orchestrator.buffer]
 env_ratios = [0.3, 0.2, 0.3, 0.2, 0.2]
 easy_threshold = 1.0
-online_difficulty_filtering = true
+min_abs_adv = 0.0
 seed = 42
 
 [orchestrator.eval]

examples/wiki_search/rl.toml

@@ -41,7 +41,7 @@ name = "qwen3-4b-wiki-search"
 max_tokens = 512
 
 [orchestrator.buffer]
-online_difficulty_filtering = true
+min_abs_adv = 0.0
 
 [[orchestrator.env]]
 id = "primeintellect/wiki-search"

src/prime_rl/configs/orchestrator.py

@@ -1,3 +1,4 @@
+import warnings
 from pathlib import Path
 from typing import Annotated, Any, Literal, TypeAlias
 
@@ -514,10 +515,20 @@ class BufferConfig(BaseConfig):
         ),
     ] = 0.0
 
+    min_abs_adv: Annotated[
+        float | None,
+        Field(
+            description=(
+                "Only send rollouts with absolute advantage greater than this threshold to training. "
+                "Rollouts are still kept for orchestrator logging. If None, no advantage-based filtering is applied."
+            ),
+        ),
+    ] = None
+
     online_difficulty_filtering: Annotated[
         bool,
         Field(
-            description="Whether to filter rollouts based on difficulty. If True, rollouts with average reward 0.0 or 1.0 are not added to the buffer.",
+            description="Deprecated alias for setting min_abs_adv=0.0.",
         ),
     ] = False
 
@@ -541,6 +552,30 @@ def validate_env_ratios(self):
             assert all(ratio > 0 for ratio in self.env_ratios), "All env_ratios must be positive."
         return self
 
+    @model_validator(mode="after")
+    def resolve_min_abs_adv(self):
+        if not self.online_difficulty_filtering:
+            return self
+
+        warnings.warn(
+            "buffer.online_difficulty_filtering is deprecated and will be removed in a future version. "
+            "Use buffer.min_abs_adv = 0.0 instead.",
+            FutureWarning,
+            stacklevel=2,
+        )
+
+        if self.min_abs_adv is None:
+            self.min_abs_adv = 0.0
+            return self
+
+        if self.min_abs_adv != 0.0:
+            raise ValueError(
+                "buffer.online_difficulty_filtering is a deprecated alias for buffer.min_abs_adv=0.0 and "
+                "cannot be combined with a different buffer.min_abs_adv value."
+            )
+
+        return self
+
 
 class VerificationConfig(BaseConfig):
     """Configures rollout verification and rubric scoring."""
@@ -979,9 +1014,9 @@ def validate_verification_config(self):
         if self.verification.enabled:
             return self
 
-        if self.buffer.online_difficulty_filtering:
+        if self.buffer.min_abs_adv is not None:
             raise ValueError(
-                "verification.enabled cannot be False when buffer.online_difficulty_filtering is True. "
+                "verification.enabled cannot be False when buffer.min_abs_adv is set. "
                 "These features depend on rewards which are disabled when verification.enabled=False."
             )
         if self.buffer.easy_threshold is not None:
@@ -997,10 +1032,10 @@ def validate_verification_config(self):
         return self
 
     @model_validator(mode="after")
-    def validate_length_shaping_requires_online_difficulty_filtering(self):
+    def validate_length_shaping_requires_min_abs_adv(self):
         if isinstance(self.advantage, DefaultAdvantageConfig) and self.advantage.length_shaping_alpha is not None:
-            if not self.buffer.online_difficulty_filtering:
-                raise ValueError("Group Relative Reward (GR³) scaling requires online difficulty filtering")
+            if self.buffer.min_abs_adv is None:
+                raise ValueError("Group Relative Reward (GR³) scaling requires buffer.min_abs_adv to be set")
         return self
 
     @model_validator(mode="after")

src/prime_rl/orchestrator/advantage.py

@@ -97,3 +97,16 @@ def compute_advantages(
 
     result = advantage_fn(inputs)
     return result.advantages.flatten().tolist()
+
+
+def compute_keep_mask(
+    advantages: list[float],
+    min_abs_adv: float | None,
+) -> list[bool]:
+    """Compute whether each rollout should be sent to training."""
+    return [should_keep_rollout(advantage, min_abs_adv) for advantage in advantages]
+
+
+def should_keep_rollout(advantage: float, min_abs_adv: float | None) -> bool:
+    """Returns whether a rollout should be sent to training."""
+    return min_abs_adv is None or abs(advantage) > min_abs_adv

src/prime_rl/orchestrator/buffer.py

@@ -231,15 +231,6 @@ def update(self, rollouts: list[vf.RolloutOutput]):
                 target_pool.append(example)
 
             self.num_examples_per_step[env_name][pool] += 1
-            if self.config.online_difficulty_filtering:
-                if avg_reward == 0.0:
-                    self.num_rollouts_per_step[env_name]["hard"] += len(example_rollouts)
-                    continue
-                elif avg_reward == 1.0:
-                    self.num_rollouts_per_step[env_name]["easy"] += len(example_rollouts)
-                    continue
-
-            self.num_rollouts_per_step[env_name]["normal"] += len(example_rollouts)
             self.rollout_buffer.extend(example_rollouts)
 
     def sample_rollouts(self, n: int) -> list[vf.RolloutOutput]:
@@ -254,8 +245,6 @@ def reset_step_metrics(self) -> None:
         zero_per_pool = lambda: {p: 0 for p in self.POOLS}
         # num examples per env per step per pool (env_name -> (pool -> num_examples))
         self.num_examples_per_step = {env: zero_per_pool() for env in self.env_names}
-        # num rollouts per env per step per pool (env_name -> (pool -> num_rollouts))
-        self.num_rollouts_per_step = {env: zero_per_pool() for env in self.env_names}
 
     def get_metrics(self) -> dict[str, float]:
         """Returns the buffer metrics for the current step."""
@@ -272,17 +261,11 @@ def get_metrics(self) -> dict[str, float]:
         num_examples_per_step_per_pool = {
             pool: sum(self.num_examples_per_step[env][pool] for env in self.env_names) for pool in self.POOLS
         }
-        num_rollouts_per_step_per_pool = {
-            pool: sum(self.num_rollouts_per_step[env][pool] for env in self.env_names) for pool in self.POOLS
-        }
         num_examples_per_step = sum(num_examples_per_step_per_pool.values())
-        num_rollouts_per_step = sum(num_rollouts_per_step_per_pool.values())
 
         for pool in ["easy", "hard"]:
             if num_examples_per_step:
                 metrics[f"evicted_examples/{pool}"] = num_examples_per_step_per_pool[pool] / num_examples_per_step
-            if num_rollouts_per_step:
-                metrics[f"filtered_rollouts/{pool}"] = num_rollouts_per_step_per_pool[pool] / num_rollouts_per_step
 
         total_normal = sum(len(self.example_buffer[env]) for env in self.env_names)
         pool_counts = [len(self.easy_examples), total_normal, len(self.hard_examples)]
@@ -292,19 +275,13 @@ def get_metrics(self) -> dict[str, float]:
 
         for env in self.env_names:
             env_num_examples_per_step_per_pool = self.num_examples_per_step[env]
-            env_num_rollouts_per_step_per_pool = self.num_rollouts_per_step[env]
             env_num_examples_per_step = sum(env_num_examples_per_step_per_pool.values())
-            env_num_rollouts_per_step = sum(env_num_rollouts_per_step_per_pool.values())
 
             for pool in ["easy", "hard"]:
                 if env_num_examples_per_step:
                     metrics[f"evicted_examples/{env}/{pool}"] = (
                         env_num_examples_per_step_per_pool[pool] / env_num_examples_per_step
                     )
-                if env_num_rollouts_per_step:
-                    metrics[f"filtered_rollouts/{env}/{pool}"] = (
-                        env_num_rollouts_per_step_per_pool[pool] / env_num_rollouts_per_step
-                    )
 
             env_pool_counts = [
                 easy_examples_per_env[env],

src/prime_rl/orchestrator/orchestrator.py

@@ -7,7 +7,7 @@
 
 import tomli_w
 
-from prime_rl.orchestrator.advantage import compute_advantages
+from prime_rl.orchestrator.advantage import compute_advantages, compute_keep_mask
 from prime_rl.orchestrator.eval_utils import compute_eval_ckpt_step, get_eval_sampling_args
 from prime_rl.orchestrator.event_loop_lag import EventLoopLagMonitor
 from prime_rl.orchestrator.patches import monkey_patch_chat_completion_logprobs, monkey_patch_oai_iterable_types
@@ -532,7 +532,7 @@ def _cleanup_env_processes():
         # Await train rollouts
         await train_task
         generate_completions_time = scheduler.last_batch_generation_time
-        train_rollouts = train_task.result()
+        train_rollouts, selected_rollout_mask = train_task.result()
 
         # VLM: offload base64 images to disk immediately to free memory
         if is_vlm:
@@ -553,11 +553,12 @@ def _cleanup_env_processes():
         rewards = [r["reward"] for r in train_rollouts]
         completion_lens = [get_completion_len(r) for r in train_rollouts]
         advantages = compute_advantages(
-            rewards,
-            completion_lens,
-            config.rollouts_per_example,
-            config.advantage,
+            rewards=rewards,
+            completion_lengths=completion_lens,
+            samples_per_problem=config.rollouts_per_example,
+            advantage_config=config.advantage,
         )
+        rollout_keep_mask = compute_keep_mask(advantages, config.buffer.min_abs_adv)
 
         # Convert rollouts to training samples
         parallel_preprocess_start = time.perf_counter()
@@ -598,7 +599,9 @@ def process_rollout(rollout: vf.RolloutOutput, rollout_idx: int) -> list[Trainin
         rollout_samples_per_rollout: list[int] = []
         num_prefill_tokens = 0
         num_decode_tokens = 0
-        for rollout, advantage, samples in zip(train_rollouts, advantages, results):
+        for rollout, advantage, selected_for_training, samples in zip(
+            train_rollouts, advantages, selected_rollout_mask, results
+        ):
             rollout_prefill_tokens = 0
             rollout_decode_tokens = 0
             if samples is not None:
@@ -610,18 +613,28 @@ def process_rollout(rollout: vf.RolloutOutput, rollout_idx: int) -> list[Trainin
                     sample_prefill_tokens = len(sample.prompt_ids) + len(sample.completion_mask) - sample_decode_tokens
                     rollout_decode_tokens += sample_decode_tokens
                     rollout_prefill_tokens += sample_prefill_tokens
-                    train_examples.append(sample)
+                    if selected_for_training:
+                        train_examples.append(sample)
             else:
                 rollout_samples_per_rollout.append(0)
             rollout_prefill_lens.append(rollout_prefill_tokens)
             rollout_decode_lens.append(rollout_decode_tokens)
             num_prefill_tokens += rollout_prefill_tokens
             num_decode_tokens += rollout_decode_tokens
 
+        num_filtered_rollouts = sum(not keep_rollout for keep_rollout in rollout_keep_mask)
         parallel_preprocess_time = time.perf_counter() - parallel_preprocess_start
         logger.debug(
             f"Converted {len(train_rollouts)} rollouts ({num_unique_examples} unique examples) "
-            f"to {len(train_examples)} training examples"
+            f"to {len(train_examples)} training examples ({sum(selected_rollout_mask)} selected rollouts)"
+            + (
+                ""
+                if config.buffer.min_abs_adv is None
+                else (
+                    f" after filtering {num_filtered_rollouts} rollout(s) with "
+                    f"|advantage| <= {config.buffer.min_abs_adv}"
+                )
+            )
         )
 
         # Compute teacher logprobs if teacher model is configured
@@ -658,6 +671,7 @@ def process_rollout(rollout: vf.RolloutOutput, rollout_idx: int) -> list[Trainin
                 "example_id": [rollout["example_id"] for rollout in train_rollouts],
                 "task": [rollout["task"] for rollout in train_rollouts],
                 "reward": [rollout["reward"] for rollout in train_rollouts],
+                "passes_min_abs_adv": rollout_keep_mask,
                 "is_truncated": [rollout["is_truncated"] for rollout in train_rollouts],
                 "stop_condition": [rollout.get("stop_condition") for rollout in train_rollouts],
                 "seq_len": [get_seq_len(rollout) for rollout in train_rollouts],
@@ -754,6 +768,7 @@ def compute_solve_rates(df):
             "solve_none/all": solve_none,
             "solve_all/all": solve_all,
             "effective_batch_size/all": effective_batch_size,
+            "filtered_rollouts/all": num_filtered_rollouts / max(num_rollouts, 1),
             **{f"batch/{env}": r for env, r in results_df.task.value_counts(normalize=True).items()},
             # Time metrics
             "time/step": step_time,
@@ -794,6 +809,7 @@ def compute_solve_rates(df):
             to_log[f"reward/{env}/mean"] = env_by_example.reward.mean().mean()
             to_log[f"reward/{env}/max"] = env_by_example.reward.mean().max()
             to_log[f"reward/{env}/min"] = env_by_example.reward.mean().min()
+            to_log[f"filtered_rollouts/{env}"] = 1.0 - env_df.passes_min_abs_adv.mean()
             solve_none, solve_all, effective_batch_size = compute_solve_rates(env_df)
             to_log[f"solve_none/{env}"] = solve_none
             to_log[f"solve_all/{env}"] = solve_all