feat(sde): implement WindowScheduler decay / exp_decay strategies

[ruiling-smartbear]

Summary

Strategy declares decay and exp_decay, and WindowConfig fully describes them (max_iters_per_window, min_iters_per_window, exp_decay_threshold, exp_decay_k, plus the decay branch in __post_init__) — but WindowScheduler has no resolver for either, so selecting one raises Bad strategy configuration. Four shipped mixgrpo recipes (sd3_mixgrpo, wan21_t2v_mixgrpo, wan22_t2v_14b_mixgrpo, qwen_image_mixgrpo) already set the decay-only fields while running progressive, so those lines are currently inert.

This adds both resolvers, mapped onto the existing stateless get_sde_indices(step) contract.

Dwell formulas follow MixGRPO's GRPOTrainingStates (fastvideo/utils/grpo_states.py:55-83), where t is the window start:

• decay — iters(t) = max(min, int(max*(1-t/T) + min*t/T))
• exp_decay — iters(t) = ceil(iters_per_window * exp(-k * relu(t - threshold)))

Index space is UniRL's own, not MixGRPO's — the full-window walk _resolve_progressive has used since the initial release, with num_timesteps = num_inference_steps. That difference is inherited rather than introduced here:

• MixGRPO's scheduler indexes its training sample dict, which holds T-1 transitions, so max_timesteps = sampling_steps - 2 (train_grpo_flux.py:921) is the last index that dict can yield. That cut comes from DanceGRPO's extra [:, :-1] (train_grpo_flux.py:511-519), not from Flow-GRPO (which keeps all T), and carries no comment or commit rationale. The same value is then used as an exclusive range end, so the last index actually reached is T-3 — one short of the value it names.
• UniRL indexes the denoising loop (_resolve_sde_window rejects any index >= num_inference_steps), and progressive already used the full 0..T-1 range.

Consequence, stated rather than hidden: with sd3_mixgrpo (T=10, window_size=4, overlap_size=1, max=10, min=1, roll_back=true) decay yields dwell 10/7/4 over indices 0..9; evaluated over MixGRPO's coordinates the same formula yields 10/6/3 over 0..7. Mirroring MixGRPO's coordinates would also change progressive, which all four shipped mixgrpo recipes run — out of scope here, and it wants its own issue.

The third commit is separable: it moves window-config validation into WindowConfig.__post_init__ so invalid schedules fail at instantiate time rather than on the first get_sde_indices call. It is not required by the new strategies (the same pitfalls already exist for progressive), so treat it as your call — say the word and I'll rebase it out.

Related Issue

Refs #488

Test Plan

Parity harness over 216 parameter combinations (num_timesteps ∈ {10,14,25} × window_size ∈ {2,3,4} × iters_per_window ∈ {4,10,25} × overlap_size ∈ {0,1} × roll_back ∈ {false,true}), replaying GRPOTrainingStates.update_iteration() step by step against get_sde_indices(step). Both coordinates were evaluated, which is the part the earlier revision of this description got wrong:

• exp_decay dwell matches the reference on the shared window prefix under either coordinate — it does not depend on max_timesteps.
• decay dwell matches on the shared window prefix when the reference is evaluated with max_timesteps := num_timesteps — 108/108; under the trainer's max_timesteps = T-2 it differs by the denominator (e.g. 10/7/4 vs 10/6/3).
• The reference additionally visits one clipped tail window per sweep in the overlapping configuration (non-empty in 48 of the 54 roll_back=true cases); UniRL keeps to full windows, matching _resolve_progressive.
• determinism, window size, index bounds, monotone non-increasing dwell — all pass.

Resulting dwell with num_timesteps=10, window_size=4, overlap_size=1:

strategy	dwell per window
progressive	25 / 25 / 25
decay (max=10, min=1)	10 / 7 / 4
exp_decay (threshold=3, k=0.1)	25 / 25 / 19

Validation harness (separate commit):

• 11/11 invalid configs raise ValueError at WindowConfig construction; all 7 shipped WindowScheduler recipes still construct.
• The 216 configs now rejected at construction previously hit ZeroDivisionError on the first get_sde_indices call (144) or silently degraded to a single / out-of-range window (72); 0 valid configs rejected.
• Scheduling unchanged by the move: 24,048 pairs over the four mixgrpo recipes × progressive / decay / exp_decay, 0 mismatches.

Regression: pre-change vs post-change WindowScheduler driven side by side over 607,824 (config, step) pairs (strategy ∈ {all, progressive, random}) — 0 mismatches — and AllSDEScheduler over 2,244 pairs — 0 mismatches.

Lint: ruff check unirl/sde/index_schedule.py, ruff format --check unirl/sde/index_schedule.py, python lint/check_docstring_lines.py — clean (CI agrees: the lint job is green on this head).

Training smoke: passed for both strategies on fa01335d, one H100 80GB, SD3.5-medium BF16 LoRA with the sd3_mixgrpo recipe and PickScore (datasets/pickscore/train.txt). Each arm completed 5 rollouts / 10 optimizer steps at 256×256; loss, gradients and updated parameters were finite, every optimizer step had nonzero gradients and changed parameters. Observed window starts: 0, 0, 3, 6, 0, covering advancement and wraparound; exp_decay used threshold 0 so decay was active. Training subprocess times: 104s / 81s including model initialization, not performance benchmarks. Modal run. This is an integration smoke, not convergence or image-quality validation.

Smoke commands (one-GPU Ray head already running)

for strategy in decay exp_decay; do
  REPORT_TO_WANDB=false python -m unirl.train_diffusion \
    --config-name=diffusion/sd3/sd3_mixgrpo \
    num_devices=1 +devices_per_node=1 num_rollouts=5 batch_size=1 \
    sampling.samples_per_prompt=4 sampling.height=256 sampling.width=256 \
    rollout.forward_batch_size=1 stack.micro_batch_size=1 logging.log_media=false \
    sampling.scheduler.config.strategy=$strategy \
    sampling.scheduler.config.iters_per_window=2 \
    sampling.scheduler.config.max_iters_per_window=2 \
    sampling.scheduler.config.min_iters_per_window=1 \
    +sampling.scheduler.config.exp_decay_threshold=0 \
    +sampling.scheduler.config.exp_decay_k=0.3
done

Torch 2.11.0+cu130, Transformers 5.6.2, Diffusers 0.40.0, PEFT 0.21.0. External observation-only instrumentation recorded scheduler output, per-update loss/gradients, and parameter deltas; no production source changes were needed.

The harnesses were run but not committed, per CLAUDE.md.

Compatibility / Risk

No API, config-schema, or checkpoint changes; all / progressive / random behaviour is unchanged. Validation moves into WindowConfig.__post_init__, so it fires at Hydra instantiate rather than at the first get_sde_indices call, and additionally rejects overlap_size >= window_size (previously ZeroDivisionError on stride == 0, or a silent single-window scheduler on stride < 0) and iters_per_window < 1 (previously ZeroDivisionError for exp_decay); the decay bound check moved here from WindowScheduler.__init__. Those two checks are scoped to the strategies that consume stride / iters_per_window (progressive, decay, exp_decay), so the three random bagel recipes are untouched.

Reviewer Notes

Two calls worth a look:

1. Coordinates. UniRL's _resolve_progressive already differs from the reference past the end of a sweep, and the reference's own boundary is not self-consistent (its comment says T-2, its code reaches T-3, upstream Flow-GRPO keeps T-1). I followed UniRL's convention so the three strategies behave consistently, which means decay / exp_decay match the reference's formulas but not its index bookkeeping. Happy to switch if you'd rather mirror MixGRPO exactly — note that would also change progressive, i.e. the four shipped recipes.
2. The validation commit is separable; drop it if you'd rather keep this PR to the two resolvers.

AI-assisted. Checked open issues and PRs for overlapping work before opening.

Checklist

• I reviewed the changed code and removed unrelated/generated artifacts.
• I updated tests, docs, and configs where needed, or explained why not.

[Jayce-Ping]

Looks okay. Though I think most of time, users will use FlowGRPO-Fast's random SDE-step selection rather than MixGRPO, it's still worth to fix.