feat(sde): implement WindowScheduler decay / exp_decay strategies

examples/diffusion/bagel/bagel_trainside_lora.yaml

@@ -168,7 +168,7 @@ sampling:
   seed: 42
   init_same_noise: false               # per-sample x_T (NoiseRecipe r{rollout}:{sample}); diverse GRPO groups
   # SDE steps: 2 random steps drawn from the early half [0, 7) of the 14-step schedule,
-  # resolved PER ROLLOUT via AllSDEScheduler → resolve_sde_indices (shared across the
+  # resolved PER ROLLOUT via AllSDEScheduler → get_sde_indices (shared across the
   # GRPO group, seeded by rollout_id). Uses UniRL's NATIVE timestep_fraction (early-half
   # restriction) + num_sde_steps (random-k) — zero framework change. Differs from
   # flow_grpo only in that the 2 SDE steps are SCATTERED, not a contiguous window;

examples/diffusion/minimax_h3/minimax_h3_t2va_nft.yaml

@@ -223,7 +223,7 @@ sampling:
   logprob_precision: fp32
   scheduler:
     # num_sde_steps: 0 IS the forward-process switch, with eta: 0.0 above — it
-    # makes resolve_sde_indices return [], read as "pure ODE, record no
+    # makes get_sde_indices return [], read as "pure ODE, record no
     # log-probs", and makes the stage store only the terminal clean latent.
     _target_: unirl.sde.index_schedule.AllSDEScheduler
     num_timesteps: ${..num_inference_steps}

examples/pe/pe_trainside_pickscore.yaml

@@ -243,7 +243,7 @@ sampling:
     # Sparse SDE-index schedule (FlowDPPO-style, mirrors sd3_flowdppo.yaml):
     # num_sde_steps random steps drawn from the first half of the denoising
     # schedule, resolved PER ROLLOUT (seeded by rollout_id) by
-    # PETrainer._build_req via resolve_sde_indices. Replaces the static all-SDE
+    # PETrainer._build_req via get_sde_indices. Replaces the static all-SDE
     # list with selective stochastic steps.
     scheduler:
       _target_: unirl.sde.index_schedule.AllSDEScheduler

unirl/sde/README.md

@@ -45,7 +45,7 @@ adapter's private choice.
   match SGLang).
 - **SDE indices own the policy-gradient density.** The selection is a
   `TimestepScheduler` (`index_schedule.py`) wired under `sampling.scheduler`, which
-  `DiffusionSamplingParams.resolve_sde_indices` asks once per rollout id. Selected
+  `DiffusionSamplingParams.get_sde_indices` asks once per rollout id. Selected
   `sde_indices` get a stochastic transition with a real per-step Gaussian log-prob
   (→ `LatentSegment.sde_logp`). Trainside loops collapse the same SDE kernel to
   Euler with `eta=0`; the FastVideo adapter instead implements the model
@@ -123,3 +123,9 @@ MixGRPO keeps `FlowSDEStrategy` and adds a `WindowScheduler` under
   *disabled*, not "stretch to zero" — normalize falsy values to `None`.
 - **`compute_mu` is the single per-model μ override point.** FLUX.2-klein's μ depends
   on **both** `image_seq_len` and `num_inference_steps`, unlike the base formula.
+- **`WindowScheduler` indexes the denoising loop.** `num_timesteps` is the full step count
+  and only full windows are visited. MixGRPO's `max_timesteps = sampling_steps - 2` indexes
+  that repo's sample dict, a different coordinate space.
+- **`exp_decay_threshold` defaults to 13.** Decay starts only once a window start passes it.
+  Shipped mixgrpo schedules are shorter, so `strategy=exp_decay` stays at `iters_per_window`
+  until the threshold is lowered.

unirl/sde/index_schedule.py

@@ -1,8 +1,11 @@
 """Index schedulers used by GRPO-style algorithms."""
 
+import bisect
+import math
 from abc import ABC, abstractmethod
 from collections.abc import Sequence
 from dataclasses import dataclass
+from itertools import accumulate
 from typing import List, Literal, Optional, Set, Tuple, Union
 
 import numpy as np
@@ -26,11 +29,29 @@ class WindowConfig:
     exp_decay_k: float = 0.1
 
     def __post_init__(self) -> None:
+        if self.strategy != "all" and self.window_size < 1:
+            raise ValueError(f"WindowConfig requires window_size >= 1, got {self.window_size}")
+        if self.strategy in ("progressive", "decay", "exp_decay"):
+            if not 0 <= self.overlap_size < self.window_size:
+                raise ValueError(
+                    f"WindowConfig({self.strategy}) requires 0 <= overlap_size < window_size, "
+                    f"got overlap_size={self.overlap_size}, window_size={self.window_size}"
+                )
+            if not math.isfinite(self.iters_per_window) or self.iters_per_window < 1:
+                raise ValueError(f"WindowConfig requires a finite iters_per_window >= 1, got {self.iters_per_window}")
         if self.strategy == "decay":
             if self.max_iters_per_window is None:
                 self.max_iters_per_window = self.iters_per_window
             if self.min_iters_per_window is None:
                 self.min_iters_per_window = max(1, self.iters_per_window // 4)
+            lo, hi = self.min_iters_per_window, self.max_iters_per_window
+            if not math.isfinite(lo) or not math.isfinite(hi) or lo < 1 or hi < lo:
+                raise ValueError(
+                    "WindowConfig(decay) requires finite bounds with "
+                    f"1 <= min_iters_per_window <= max_iters_per_window, got ({lo}, {hi})"
+                )
+        if self.strategy == "exp_decay" and not math.isfinite(self.exp_decay_k):
+            raise ValueError(f"WindowConfig(exp_decay) requires a finite exp_decay_k, got {self.exp_decay_k}")
 
 
 class TimestepScheduler(ABC):
@@ -106,6 +127,8 @@ class WindowScheduler(TimestepScheduler):
         "all": None,
         "progressive": "_resolve_progressive",
         "random": "_resolve_random",
+        "decay": "_resolve_sliding",
+        "exp_decay": "_resolve_sliding",
     }
 
     def __init__(self, num_timesteps: int, config: WindowConfig):
@@ -127,20 +150,48 @@ def get_sde_indices(self, step: Optional[int] = None) -> Set[int]:
         return resolve_method(0 if step is None else int(step))
 
     def _resolve_progressive(self, step: int) -> Set[int]:
+        starts = self._window_starts()
         window_step = step // self.config.iters_per_window
-        stride = self.config.window_size - self.config.overlap_size
-        remaining = self.num_timesteps - self.config.init_timestep - self.config.window_size
-        num_one_round_window_steps = max(1, remaining // stride + 1)
-        if window_step >= num_one_round_window_steps and not self.config.roll_back:
-            window_step = num_one_round_window_steps - 1
-            return self._resolve_progressive(window_step * self.config.iters_per_window)
-
-        window_step = window_step % num_one_round_window_steps
-        cur_timestep = self.config.init_timestep + window_step * stride
-        return set(range(cur_timestep, cur_timestep + self.config.window_size))
+        if window_step >= len(starts) and not self.config.roll_back:
+            window_step = len(starts) - 1
+        else:
+            window_step = window_step % len(starts)
+        cur = starts[window_step]
+        return set(range(cur, cur + self.config.window_size))
 
     def _resolve_random(self, step: int) -> Set[int]:
         rng = np.random.default_rng(step)
         max_start = max(0, self.num_timesteps - self.config.window_size)
         cur_timestep = int(rng.integers(0, max_start + 1))
         return set(range(cur_timestep, cur_timestep + self.config.window_size))
+
+    def _window_starts(self) -> List[int]:
+        """Start index of every full window in one sweep, in visiting order."""
+        stride = self.config.window_size - self.config.overlap_size
+        remaining = self.num_timesteps - self.config.init_timestep - self.config.window_size
+        count = max(1, remaining // stride + 1)
+        return [self.config.init_timestep + i * stride for i in range(count)]
+
+    def _resolve_sliding(self, step: int) -> Set[int]:
+        """Walk full windows whose dwell is the decay or exp_decay value at each start."""
+        starts = self._window_starts()
+        if self.config.strategy == "decay":
+            lo = self.config.min_iters_per_window
+            hi = self.config.max_iters_per_window
+            dwell = []
+            for start in starts:
+                progress = start / self.num_timesteps
+                dwell.append(max(lo, int(hi * (1.0 - progress) + lo * progress)))
+        else:
+            base = self.config.iters_per_window
+            k = self.config.exp_decay_k
+            threshold = self.config.exp_decay_threshold
+            dwell = [int(math.ceil(base * math.exp(-k * max(0, start - threshold)))) for start in starts]
+        total = sum(dwell)
+        if step >= total:
+            if not self.config.roll_back:
+                cur = starts[-1]
+                return set(range(cur, cur + self.config.window_size))
+            step = step % total
+        cur = starts[bisect.bisect_right(list(accumulate(dwell)), step)]
+        return set(range(cur, cur + self.config.window_size))

unirl/trainer/diffusion.py

@@ -692,7 +692,7 @@ def _build_request_sample(
         sp = sampling if sampling is not None else self.sampling_params
         noise_latent_shape = self._eval_noise_latent_shape if sampling is not None else self._noise_latent_shape
         diffusion = sp.get("diffusion")
-        sde_indices = diffusion.resolve_sde_indices(rollout_id)
+        sde_indices = diffusion.get_sde_indices(rollout_id)
         diffusion = dataclasses.replace(
             diffusion, sde_indices=sde_indices, scheduler=None, init_noise_latent_shape=noise_latent_shape
         )

unirl/trainer/pe.py

@@ -156,7 +156,7 @@ def _build_request_sample(
         base = sampling if sampling is not None else self.sampling_params
         diff_params = base.get("diffusion")
         ar_params = base.get("ar")
-        sde_indices = diff_params.resolve_sde_indices(rollout_id)
+        sde_indices = diff_params.get_sde_indices(rollout_id)
         diffusion = dataclasses.replace(diff_params, sde_indices=sde_indices, scheduler=None)
         request = prepare_input_sample(
             inputs,

unirl/trainer/unified_model.py

@@ -200,7 +200,7 @@ def _build_request_sample(
         base = sampling if sampling is not None else self.sampling_params
         diff_params = base.get("diffusion")
         ar_params = base.get("ar")
-        sde_indices = diff_params.resolve_sde_indices(rollout_id)
+        sde_indices = diff_params.get_sde_indices(rollout_id)
         disable_xt = bool(os.environ.get("DISABLE_DRIVER_XT")) or bool(getattr(diff_params, "disable_driver_xt", False))
         diffusion = dataclasses.replace(
             diff_params, sde_indices=sde_indices, scheduler=None, disable_driver_xt=disable_xt

unirl/types/sampling.py

@@ -182,8 +182,8 @@ def __post_init__(self) -> None:
             f"{cls}.sampler_kwargs cannot contain reserved keys {sorted(shadowed)}; set them as fields instead",
         )
 
-    def resolve_sde_indices(self, rollout_id: int) -> List[int]:
-        """Resolve which denoising steps record SDE log-probs for ``rollout_id``."""
+    def get_sde_indices(self, rollout_id: int) -> List[int]:
+        """SDE steps for ``rollout_id``: a preset list, else ``scheduler.get_sde_indices``, else every step."""
         if self.sde_indices is not None:
             return [int(i) for i in self.sde_indices]
         scheduler: Optional[TimestepScheduler] = self.scheduler