From 7d9efe0cc08887a4e80a8ec28abab84465817739 Mon Sep 17 00:00:00 2001
From: niehen6174 <niehen6174@gmail.com>
Date: Mon, 20 Jul 2026 11:50:22 +0000
Subject: [PATCH 1/5] feat(algorithms): extract Flow-GRPO into a
 DiffusionAlgorithm plugin

Introduce algorithm interfaces, registry, shared forward helpers, and the
Flow-GRPO loss/label implementation so later algorithms can plug in without
growing the FSDP actor further.
---
 miles/algorithms/__init__.py            |  17 ++
 miles/algorithms/base.py                |  92 ++++++++++
 miles/algorithms/flow_grpo.py           | 213 ++++++++++++++++++++++++
 miles/algorithms/labels.py              |  30 ++++
 miles/algorithms/registry.py            |  44 +++++
 miles/algorithms/train_forward_utils.py | 160 ++++++++++++++++++
 6 files changed, 556 insertions(+)
 create mode 100644 miles/algorithms/__init__.py
 create mode 100644 miles/algorithms/base.py
 create mode 100644 miles/algorithms/flow_grpo.py
 create mode 100644 miles/algorithms/labels.py
 create mode 100644 miles/algorithms/registry.py
 create mode 100644 miles/algorithms/train_forward_utils.py

diff --git a/miles/algorithms/__init__.py b/miles/algorithms/__init__.py
new file mode 100644
index 00000000..c9443350
--- /dev/null
+++ b/miles/algorithms/__init__.py
@@ -0,0 +1,17 @@
+"""Diffusion algorithm plugins.
+
+Currently ships Flow-GRPO only; SFT / AWM / DiffusionNFT land in follow-up PRs.
+"""
+
+from miles.algorithms.base import CollectionSpec, DiffusionAlgorithm, TrainLabels, TrainLossContext
+from miles.algorithms.registry import builtin_algorithm_names, load_algorithm, resolve_algorithm_class_path
+
+__all__ = [
+    "CollectionSpec",
+    "DiffusionAlgorithm",
+    "TrainLabels",
+    "TrainLossContext",
+    "builtin_algorithm_names",
+    "load_algorithm",
+    "resolve_algorithm_class_path",
+]
diff --git a/miles/algorithms/base.py b/miles/algorithms/base.py
new file mode 100644
index 00000000..aa876c1c
--- /dev/null
+++ b/miles/algorithms/base.py
@@ -0,0 +1,92 @@
+"""Diffusion algorithm plugin interfaces.
+
+Model-family concerns stay in ``TrainPipelineConfig``; algorithm plugins own
+collection contracts, train-example schema, and loss.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+from typing import Any, Protocol, runtime_checkable
+
+import torch
+
+
+@dataclass(frozen=True)
+class CollectionSpec:
+    """What the collector / rollout engine must provide for this algorithm."""
+
+    mode: str  # "offline" | "online"
+    needs_reward: bool = True
+    needs_trajectory: bool = True
+    needs_logprob: bool = True
+    sampler: str = "sde"  # "sde" | "ode" | "any"
+    return_denoising_env: bool = True
+    sync_weights_to_rollout: bool = True
+
+
+@dataclass
+class TrainLossContext:
+    """Train-side handles passed into ``compute_loss``."""
+
+    models: dict[str, torch.nn.Module]
+    model: torch.nn.Module
+    train_pipeline_config: Any
+    sde_backend: Any | None
+    scheduler: Any | None
+    args: Any
+    forward_dtype: torch.dtype
+    device: torch.device
+
+
+@dataclass
+class LossOutput:
+    """Reserved return type for ``compute_loss`` (loss + metrics).
+
+    **Unused in the Flow-GRPO PR:** ``compute_loss`` still returns a bare
+    ``Tensor`` and appends into the caller's ``log_stats`` (actor-compatible).
+    Kept for a later cleanup that unifies the return shape.
+    """
+
+    loss_sum: torch.Tensor
+    log_stats: dict[str, list[torch.Tensor]] = field(default_factory=dict)
+
+
+@dataclass
+class TrainLabels:
+    """Post-reward labels attached to samples before ``build_train_data``."""
+
+    raw_rewards: list[float]
+    advantages: list[float] | None = None
+    # Reserved for DiffusionNFT soft positive/negative labels; Flow-GRPO ignores it.
+    nft_labels: list[float] | None = None
+
+
+@runtime_checkable
+class DiffusionAlgorithm(Protocol):
+    name: str
+
+    def validate_args(self, args) -> None: ...
+
+    def collection_spec(self) -> CollectionSpec:
+        """Return acquisition contract; see ``CollectionSpec`` — not fully consumed yet."""
+        ...
+
+    def postprocess_rewards(self, args, samples: list) -> TrainLabels: ...
+
+    def build_train_data(self, args, samples: list, labels: TrainLabels) -> dict[str, Any]: ...
+
+    def validate_train_batch(self, batch: list[dict]) -> list[str]: ...
+
+    def compute_loss(
+        self,
+        ctx: TrainLossContext,
+        batch: list[dict],
+        *,
+        log_stats: dict[str, list[torch.Tensor]],
+        pad_to_len: int | None = None,
+    ) -> torch.Tensor: ...
+
+    def prepare_rollout_data(self, rollout_data: dict, ctx: TrainLossContext) -> None:
+        """Optional hook before the micro-batch loop (e.g. sync scheduler meta)."""
+        ...
\ No newline at end of file
diff --git a/miles/algorithms/flow_grpo.py b/miles/algorithms/flow_grpo.py
new file mode 100644
index 00000000..fab2c5e6
--- /dev/null
+++ b/miles/algorithms/flow_grpo.py
@@ -0,0 +1,213 @@
+"""Flow-GRPO: reverse-SDE log-prob + PPO-clip."""
+
+from __future__ import annotations
+
+from typing import Any
+
+import torch
+
+from miles.algorithms.base import CollectionSpec, TrainLabels, TrainLossContext
+from miles.algorithms.labels import grpo_group_advantages
+from miles.algorithms.train_forward_utils import (
+    append_rollout_train_abs_diff_stats,
+    compute_noise_pred,
+    prepare_cfg_conds,
+    resolve_cfg_flags,
+    select_model_for_timesteps,
+)
+from miles.utils.train_data_utils import RolloutTrainDataConverter, scheduler_meta_from_rollout, stack_train_pair_rollout_debug
+from miles.utils.types import Sample
+
+
+class FlowGRPOAlgorithm:
+    name = "flow_grpo"
+
+    def validate_args(self, args) -> None:
+        kl_beta = float(getattr(args, "diffusion_kl_beta", 0.0) or 0.0)
+        if kl_beta > 0 and not args.use_lora:
+            raise ValueError(
+                "--diffusion-kl-beta currently requires --use-lora so the base model can be used as reference."
+            )
+
+    def collection_spec(self) -> CollectionSpec:
+        return CollectionSpec(
+            mode="online",
+            needs_reward=True,
+            needs_trajectory=True,
+            needs_logprob=True,
+            sampler="sde",
+            return_denoising_env=True,
+            sync_weights_to_rollout=True,
+        )
+
+    def postprocess_rewards(self, args, samples: list[Sample]) -> TrainLabels:
+        return grpo_group_advantages(args, samples)
+
+    def build_train_data(self, args, samples: list[Sample], labels: TrainLabels) -> dict[str, Any]:
+        rewards = labels.advantages if labels.advantages is not None else labels.raw_rewards
+        return RolloutTrainDataConverter().convert_samples(samples, rewards, labels.raw_rewards)
+
+    def validate_train_batch(self, batch: list[dict]) -> list[str]:
+        errors: list[str] = []
+        required = ("latent", "next_latent", "timestep", "next_timestep", "log_prob_old", "advantage", "denoising_env")
+        for i, pair in enumerate(batch):
+            for key in required:
+                if key not in pair:
+                    errors.append(f"batch[{i}] missing {key}")
+        return errors
+
+    def prepare_rollout_data(self, rollout_data: dict, ctx: TrainLossContext) -> None:
+        if ctx.scheduler is None:
+            return
+        num_train_timesteps = ctx.scheduler.config.num_train_timesteps
+        scheduler_timesteps, scheduler_sigmas = scheduler_meta_from_rollout(
+            rollout_data,
+            device=ctx.device,
+            num_train_timesteps=num_train_timesteps,
+        )
+        ctx.scheduler.timesteps = scheduler_timesteps
+        ctx.scheduler.sigmas = scheduler_sigmas
+        ctx.scheduler._step_index = None
+        ctx.scheduler._begin_index = None
+
+    def compute_loss(
+        self,
+        ctx: TrainLossContext,
+        batch: list[dict],
+        *,
+        log_stats: dict[str, list[torch.Tensor]],
+        pad_to_len: int | None = None,
+    ) -> torch.Tensor:
+        """One DiT forward + PPO loss over ``len(batch)`` train pairs. Returns sum of per-pair losses."""
+        if ctx.sde_backend is None:
+            raise RuntimeError("Flow-GRPO requires an SDE step backend")
+
+        args = ctx.args
+        forward_dtype = ctx.forward_dtype
+        train_pipeline_config = ctx.train_pipeline_config
+        device = ctx.device
+        bsz = len(batch)
+
+        use_cfg, guidance_scale, true_cfg_scale = resolve_cfg_flags(args)
+        clip_range = args.diffusion_clip_range
+        noise_level = args.diffusion_noise_level
+        num_train_timesteps = ctx.scheduler.config.num_train_timesteps
+        kl_beta = float(args.diffusion_kl_beta)
+
+        def _stack(key):
+            return torch.stack([pair[key] for pair in batch]).to(device=device, dtype=torch.float32)
+
+        latents_microbatch = _stack("latent")
+        next_latents_microbatch = _stack("next_latent")
+        timesteps_microbatch = _stack("timestep")
+        next_timesteps_microbatch = _stack("next_timestep")
+        log_prob_old_microbatch = _stack("log_prob_old")
+
+        advantage = torch.tensor(
+            [float(pair["advantage"]) for pair in batch],
+            device=device,
+            dtype=torch.float32,
+        )
+        advantage = torch.clamp(advantage, -args.diffusion_adv_clip_max, args.diffusion_adv_clip_max)
+
+        component, model, guidance_scale = select_model_for_timesteps(
+            ctx,
+            timesteps_microbatch,
+            guidance_scale=guidance_scale,
+            num_train_timesteps=num_train_timesteps,
+        )
+
+        if train_pipeline_config.needs_timestep_scaling:
+            timesteps_for_model = timesteps_microbatch / float(num_train_timesteps)
+        else:
+            timesteps_for_model = timesteps_microbatch
+
+        pos_cond, neg_cond, joint_cond, cfg_batching = prepare_cfg_conds(
+            ctx, batch, use_cfg=use_cfg, pad_to_len=pad_to_len
+        )
+
+        latents_input = latents_microbatch.to(forward_dtype)
+        timesteps_input = timesteps_for_model.to(forward_dtype)
+
+        def _pred(disable_adapter: bool = False) -> torch.Tensor:
+            return compute_noise_pred(
+                ctx,
+                model=model,
+                latents_input=latents_input,
+                timesteps_input=timesteps_input,
+                pos_cond=pos_cond,
+                neg_cond=neg_cond,
+                joint_cond=joint_cond,
+                use_cfg=use_cfg,
+                cfg_batching=cfg_batching,
+                guidance_scale=guidance_scale,
+                true_cfg_scale=true_cfg_scale,
+                disable_adapter=disable_adapter,
+            )
+
+        noise_pred_microbatch = _pred()
+
+        _, log_prob_new_microbatch, prev_sample_mean_new, std_dev_t_new = ctx.sde_backend.sde_step_logprob(
+            noise_pred_microbatch.float(),
+            timesteps_microbatch,
+            next_timesteps_microbatch,
+            latents_microbatch.float(),
+            prev_sample=next_latents_microbatch.float(),
+            noise_level=noise_level,
+        )
+
+        log_prob_new = log_prob_new_microbatch
+        log_prob_old = log_prob_old_microbatch
+        ratio = torch.exp(log_prob_new - log_prob_old)
+        unclipped = -advantage * ratio
+        clipped = -advantage * torch.clamp(ratio, 1.0 - clip_range, 1.0 + clip_range)
+        per_pair_loss = torch.maximum(unclipped, clipped)
+        loss_sum = per_pair_loss.sum()
+
+        kl_loss = loss_sum.new_zeros(())
+        if kl_beta > 0:
+            with torch.no_grad():
+                ref_noise_pred_microbatch = _pred(disable_adapter=True)
+                _, _, prev_sample_mean_ref, _ = ctx.sde_backend.sde_step_logprob(
+                    ref_noise_pred_microbatch.float(),
+                    timesteps_microbatch,
+                    next_timesteps_microbatch,
+                    latents_microbatch.float(),
+                    prev_sample=next_latents_microbatch.float(),
+                    noise_level=noise_level,
+                )
+            kl_per_pair = ((prev_sample_mean_new - prev_sample_mean_ref) ** 2).mean(
+                dim=tuple(range(1, prev_sample_mean_new.ndim)),
+                keepdim=True,
+            ) / (2 * std_dev_t_new**2)
+            loss_sum = loss_sum + kl_beta * kl_per_pair.sum()
+            kl_loss = kl_per_pair.mean()
+
+        with torch.no_grad():
+            log_stats["loss"].append((per_pair_loss.mean() + kl_beta * kl_loss).detach())
+            log_stats["policy_loss"].append(per_pair_loss.mean().detach())
+            log_stats["kl_loss"].append(kl_loss.detach())
+            log_stats["loss_abs_mean"].append(per_pair_loss.abs().mean().detach())
+            log_stats["adv_abs_mean"].append(advantage.abs().mean().detach())
+            log_stats["ratio_abs_minus_1"].append((ratio - 1.0).abs().mean().detach())
+            log_stats["approx_kl"].append(0.5 * torch.mean((log_prob_new - log_prob_old) ** 2).detach())
+            log_stats["clipfrac"].append(torch.mean((torch.abs(ratio - 1.0) > clip_range).float()).detach())
+            log_stats["log_prob_new_idx_0"].append(log_prob_new[0].detach())
+            log_stats["log_prob_old_idx_0"].append(log_prob_old[0].detach())
+            log_prob_mean_abs_diff = torch.mean(torch.abs(log_prob_new - log_prob_old)).detach()
+            log_stats["log_prob_mean_abs_diff"].append(log_prob_mean_abs_diff)
+            if len(ctx.models) > 1:
+                log_stats[f"log_prob_mean_abs_diff_{component}"].append(log_prob_mean_abs_diff)
+
+            rollout_model_output = stack_train_pair_rollout_debug(batch, "rollout_step_model_output")
+            if rollout_model_output is not None:
+                mean_abs_diff = append_rollout_train_abs_diff_stats(
+                    log_stats,
+                    "model_output",
+                    noise_pred_microbatch.float(),
+                    rollout_model_output.to(device=device, dtype=torch.float32),
+                )
+                if len(ctx.models) > 1:
+                    log_stats[f"model_output_mean_abs_diff_{component}"].append(mean_abs_diff)
+
+        return loss_sum
diff --git a/miles/algorithms/labels.py b/miles/algorithms/labels.py
new file mode 100644
index 00000000..721c8c31
--- /dev/null
+++ b/miles/algorithms/labels.py
@@ -0,0 +1,30 @@
+"""Shared reward → label helpers for GRPO-style algorithms."""
+
+from __future__ import annotations
+
+import torch
+
+from miles.algorithms.base import TrainLabels
+from miles.utils.types import Sample
+
+
+def grpo_group_advantages(args, samples: list[Sample]) -> TrainLabels:
+    """Group-relative advantage normalization (Flow-GRPO default)."""
+    raw_rewards = [sample.get_reward_value(args) for sample in samples]
+    rewards_flat = torch.tensor(raw_rewards, dtype=torch.float)
+    rewards = rewards_flat.view(-1, args.n_samples_per_prompt)
+
+    if args.globalize_reward_mean:
+        mean = rewards_flat.mean()
+    else:
+        mean = rewards.mean(dim=-1, keepdim=True)
+    centered = rewards - mean
+
+    if args.grpo_std_normalization:
+        if args.globalize_reward_std:
+            std = rewards_flat.std()
+        else:
+            std = rewards.std(dim=-1, keepdim=True)
+        centered = centered / (std + 1e-4)
+
+    return TrainLabels(raw_rewards=raw_rewards, advantages=centered.flatten().tolist())
diff --git a/miles/algorithms/registry.py b/miles/algorithms/registry.py
new file mode 100644
index 00000000..29044eb9
--- /dev/null
+++ b/miles/algorithms/registry.py
@@ -0,0 +1,44 @@
+"""Resolve ``--diffusion-algorithm`` / ``--diffusion-algorithm-path`` to a class.
+
+PR1 ships Flow-GRPO only; later PRs register additional builtins.
+"""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from miles.algorithms.base import DiffusionAlgorithm
+
+_BUILTIN: dict[str, str] = {
+    "flow_grpo": "miles.algorithms.flow_grpo.FlowGRPOAlgorithm",
+}
+
+
+def resolve_algorithm_class_path(args) -> str:
+    """Return the dotted class path for the selected diffusion algorithm."""
+    if getattr(args, "diffusion_algorithm_path", None):
+        return args.diffusion_algorithm_path
+    name = getattr(args, "diffusion_algorithm", None) or "flow_grpo"
+    key = str(name).strip().lower()
+    if key not in _BUILTIN:
+        raise ValueError(
+            f"Unknown --diffusion-algorithm {name!r}; choose one of {sorted(_BUILTIN)} "
+            "or pass --diffusion-algorithm-path. "
+            "SFT/AWM/NFT land in follow-up PRs."
+        )
+    return _BUILTIN[key]
+
+
+def load_algorithm(args) -> DiffusionAlgorithm:
+    from miles.utils.misc import load_function
+
+    path = resolve_algorithm_class_path(args)
+    cls = load_function(path)
+    algo = cls() if isinstance(cls, type) else cls
+    algo.validate_args(args)
+    return algo
+
+
+def builtin_algorithm_names() -> list[str]:
+    return sorted(_BUILTIN)
diff --git a/miles/algorithms/train_forward_utils.py b/miles/algorithms/train_forward_utils.py
new file mode 100644
index 00000000..694a94e4
--- /dev/null
+++ b/miles/algorithms/train_forward_utils.py
@@ -0,0 +1,160 @@
+"""Shared train-side DiT forward helpers used by algorithm plugins."""
+
+from __future__ import annotations
+
+from contextlib import nullcontext
+from typing import Any
+
+import torch
+
+from miles.algorithms.base import TrainLossContext
+
+
+def cast_cond_to_dtype(cond: dict, dtype: torch.dtype) -> dict:
+    """Cast floating-point tensors to the model's compute dtype; leave masks alone."""
+    out: dict = {}
+    for k, v in cond.items():
+        if isinstance(v, torch.Tensor) and v.dtype.is_floating_point:
+            out[k] = v.to(dtype)
+        else:
+            out[k] = v
+    return out
+
+
+def append_rollout_train_abs_diff_stats(
+    log_stats: dict[str, list],
+    prefix: str,
+    train: torch.Tensor,
+    rollout: torch.Tensor,
+) -> torch.Tensor:
+    bsz = train.shape[0]
+    diff = (train.reshape(bsz, -1).float() - rollout.reshape(bsz, -1).float()).abs()
+    ref_max = rollout.reshape(bsz, -1).float().abs().max() + 1e-30
+    mean_abs_diff = diff.mean().detach()
+    log_stats[f"{prefix}_max_abs_diff"].append(diff.max().detach())
+    log_stats[f"{prefix}_mean_abs_diff"].append(mean_abs_diff)
+    log_stats[f"{prefix}_rel_max"].append((diff.max() / ref_max).detach())
+    return mean_abs_diff
+
+
+def select_model_for_timesteps(
+    ctx: TrainLossContext,
+    timesteps: torch.Tensor,
+    *,
+    guidance_scale: float,
+    num_train_timesteps: int,
+) -> tuple[str, torch.nn.Module, float]:
+    """Pick the DiT component (Wan dual-expert aware) and maybe retarget guidance."""
+    train_pipeline_config = ctx.train_pipeline_config
+    if len(ctx.models) == 1:
+        component, model = next(iter(ctx.models.items()))
+        return component, model, guidance_scale
+
+    components = {
+        train_pipeline_config.component_for_timestep(t, num_train_timesteps) for t in timesteps.tolist()
+    }
+    if len(components) > 1:
+        raise ValueError(
+            f"Micro-batch mixes denoising phases {sorted(components)}; set "
+            "--micro-batch-size 1 so each forward is phase-pure (one DiT, one CFG scale)."
+        )
+    component = components.pop()
+    model = ctx.models[component]
+    guidance_scale = train_pipeline_config.select_guidance_scale(
+        float(timesteps[0]),
+        num_train_timesteps,
+        guidance_scale,
+        ctx.args.diffusion_guidance_scale_2,
+    )
+    return component, model, guidance_scale
+
+
+def prepare_cfg_conds(
+    ctx: TrainLossContext,
+    batch: list[dict],
+    *,
+    use_cfg: bool,
+    pad_to_len: int | None,
+) -> tuple[dict | None, dict | None, dict | None, bool]:
+    """Build pos / neg / joint cond dicts for a micro-batch."""
+    train_pipeline_config = ctx.train_pipeline_config
+    device = ctx.device
+    forward_dtype = ctx.forward_dtype
+    bsz = len(batch)
+
+    pos_list = [
+        train_pipeline_config.prepare_cond_kwargs(batch[i]["denoising_env"].pos_cond_kwargs, device)
+        for i in range(bsz)
+    ]
+    neg_list = (
+        [
+            train_pipeline_config.prepare_cond_kwargs(batch[i]["denoising_env"].neg_cond_kwargs, device)
+            for i in range(bsz)
+        ]
+        if use_cfg
+        else None
+    )
+
+    cfg_batching = use_cfg and bool(ctx.args.fsdp_cfg_batching)
+    joint_cond = None
+    pos_cond = None
+    neg_cond = None
+    if cfg_batching:
+        joint_cond = cast_cond_to_dtype(
+            train_pipeline_config.collate_cond_for_sample_batch(pos_list + neg_list, device, pad_to_len=pad_to_len),
+            forward_dtype,
+        )
+    else:
+        pos_cond = cast_cond_to_dtype(
+            train_pipeline_config.collate_cond_for_sample_batch(pos_list, device, pad_to_len=pad_to_len),
+            forward_dtype,
+        )
+        if use_cfg and neg_list is not None:
+            neg_cond = cast_cond_to_dtype(
+                train_pipeline_config.collate_cond_for_sample_batch(neg_list, device, pad_to_len=pad_to_len),
+                forward_dtype,
+            )
+    return pos_cond, neg_cond, joint_cond, cfg_batching
+
+
+def compute_noise_pred(
+    ctx: TrainLossContext,
+    *,
+    model: torch.nn.Module,
+    latents_input: torch.Tensor,
+    timesteps_input: torch.Tensor,
+    pos_cond: dict | None,
+    neg_cond: dict | None,
+    joint_cond: dict | None,
+    use_cfg: bool,
+    cfg_batching: bool,
+    guidance_scale: float,
+    true_cfg_scale: float | None,
+    disable_adapter: bool = False,
+) -> torch.Tensor:
+    adapter_ctx = model.disable_adapter() if disable_adapter else nullcontext()
+    with adapter_ctx:
+        return ctx.train_pipeline_config.compute_noise_pred(
+            model=model,
+            latents_input=latents_input,
+            timesteps_input=timesteps_input,
+            pos_cond=pos_cond,
+            neg_cond=neg_cond,
+            joint_cond=joint_cond,
+            use_cfg=use_cfg,
+            cfg_batching=cfg_batching,
+            guidance_scale=guidance_scale,
+            true_cfg_scale=true_cfg_scale,
+        )
+
+
+def resolve_cfg_flags(args: Any) -> tuple[bool, float, float | None]:
+    guidance_scale = args.diffusion_guidance_scale
+    true_cfg_scale = args.diffusion_true_cfg_scale
+    cfg_scale = true_cfg_scale if true_cfg_scale is not None else guidance_scale
+    use_cfg = cfg_scale > 0
+    return use_cfg, guidance_scale, true_cfg_scale
+
+
+def model_has_disable_adapter(models: dict[str, torch.nn.Module]) -> bool:
+    return all(hasattr(m, "disable_adapter") for m in models.values())

From e8d3825fda2ac2dd46a95cdee8ffcfe086e3c157 Mon Sep 17 00:00:00 2001
From: niehen6174 <niehen6174@gmail.com>
Date: Mon, 20 Jul 2026 11:50:26 +0000
Subject: [PATCH 2/5] refactor: wire Flow-GRPO algorithm plugin into actor and
 rollout

Delegate reward postprocess, train-data conversion, and PPO loss to the
selected DiffusionAlgorithm, and add --diffusion-algorithm CLI selection.
---
 miles/backends/fsdp_utils/actor.py | 295 +++--------------------------
 miles/ray/rollout.py               |  45 ++---
 miles/utils/arguments.py           |  21 ++
 miles/utils/diffusion_protocol.py  |   6 +-
 4 files changed, 67 insertions(+), 300 deletions(-)

diff --git a/miles/backends/fsdp_utils/actor.py b/miles/backends/fsdp_utils/actor.py
index d0af5e24..9bd5baa9 100644
--- a/miles/backends/fsdp_utils/actor.py
+++ b/miles/backends/fsdp_utils/actor.py
@@ -1,7 +1,6 @@
 import logging
 from argparse import Namespace
 from collections import defaultdict
-from contextlib import nullcontext
 
 import ray
 import torch
@@ -11,6 +10,8 @@
 import miles.backends.fsdp_utils.configs.qwen_image  # noqa: F401 — register pipeline config
 import miles.backends.fsdp_utils.configs.sd3  # noqa: F401 — register pipeline config
 import miles.backends.fsdp_utils.configs.wan2_2  # noqa: F401 — register pipeline config
+from miles.algorithms.base import TrainLossContext
+from miles.algorithms.registry import load_algorithm
 from miles.ray.train_actor import TrainRayActor
 from miles.utils import tracking_utils, train_metric_utils
 from miles.utils.context_utils import with_defer
@@ -22,8 +23,6 @@
 from miles.utils.tracking_utils import init_tracking
 from miles.utils.train_data_utils import (
     build_microbatch_schedule,
-    scheduler_meta_from_rollout,
-    stack_train_pair_rollout_debug,
     validate_same_microbatch_counts_across_dp,
 )
 from . import checkpoint
@@ -49,10 +48,10 @@ def _enable_deterministic_training(args: Namespace) -> None:
 
 
 class FSDPTrainRayActor(TrainRayActor):
-    """FSDP training actor for diffusion GRPO.
+    """FSDP training actor for diffusion algorithms.
 
     Loads only the DiT (transformer) from a diffusers pipeline, wraps it with
-    FSDP, and trains with a PPO-clipped objective aligned with flow GRPO.
+    FSDP, and delegates loss / train-example semantics to ``DiffusionAlgorithm``.
     """
 
     @with_defer(lambda: Timer().start("train_wait"))
@@ -88,6 +87,7 @@ def init(self, args: Namespace, role: str, with_ref: bool = False) -> int:  # ty
 
         from miles.utils.misc import load_function
 
+        self.algorithm = load_algorithm(args)
         self.train_pipeline_config = load_function(args.train_pipeline_config_path)()
         self.train_pipeline_config.configure(args)
         self.model_backend = load_function(args.model_backend_path)(self.train_pipeline_config)
@@ -292,11 +292,11 @@ def train(self, rollout_id: int, rollout_data_ref) -> None:  # type: ignore[over
         )
 
     def _train_core(self, rollout_id: int, rollout_data) -> None:
-        """Diffusion GRPO: ``rollout_data[train_data]`` is a flat list of train-pair dicts.
+        """Train on ``rollout_data[train_data]`` via the configured ``DiffusionAlgorithm``.
 
-        Optimizer windows are contiguous groups of train pairs. Within a window, consecutive microbatches of
-        size ``--micro-batch-size`` drive one forward+backward each; gradients
-        scale as mean over all train pairs in the window (``loss_chunk / num_local_pairs``).
+        Optimizer windows are contiguous groups of train examples. Within a window, consecutive
+        microbatches of size ``--micro-batch-size`` drive one forward+backward each; gradients
+        scale as mean over all examples in the window (``loss_chunk / num_local_pairs``).
         """
         device = torch.cuda.current_device()
 
@@ -304,20 +304,12 @@ def _train_core(self, rollout_id: int, rollout_data) -> None:
         if not train_pairs:
             raise ValueError("rollout_data['train_data'] is empty")
 
-        num_pairs = len(train_pairs)
-
-        # ------------- CFG Scale -------------
-        guidance_scale = self.args.diffusion_guidance_scale
-        true_cfg_scale = self.args.diffusion_true_cfg_scale
-        cfg_scale = true_cfg_scale if true_cfg_scale is not None else guidance_scale
-        use_cfg = cfg_scale > 0
+        batch_errors = self.algorithm.validate_train_batch(train_pairs)
+        if batch_errors:
+            raise ValueError(f"Invalid train batch for {self.algorithm.name}: {batch_errors}")
 
-        # ------------- Loss / SDE Parameters -------------
-        clip_range = self.args.diffusion_clip_range
-        noise_level = self.args.diffusion_noise_level
-        num_train_timesteps = self.scheduler.config.num_train_timesteps
+        num_pairs = len(train_pairs)
 
-        # ------------- KL loss -------------
         kl_beta = float(self.args.diffusion_kl_beta)
         if kl_beta > 0 and not self.args.use_lora:
             raise ValueError(
@@ -326,16 +318,17 @@ def _train_core(self, rollout_id: int, rollout_data) -> None:
         if kl_beta > 0 and not all(hasattr(m, "disable_adapter") for m in self.models.values()):
             raise RuntimeError("Diffusion KL requires PEFT models exposing disable_adapter() after FSDP wrapping.")
 
-        # ------------- Rollout Scheduler Metadata -------------
-        scheduler_timesteps, scheduler_sigmas = scheduler_meta_from_rollout(
-            rollout_data,
+        loss_ctx = TrainLossContext(
+            models=self.models,
+            model=self.model,
+            train_pipeline_config=self.train_pipeline_config,
+            sde_backend=self.sde_backend,
+            scheduler=self.scheduler,
+            args=self.args,
+            forward_dtype=self._forward_dtype,
             device=device,
-            num_train_timesteps=num_train_timesteps,
         )
-        self.scheduler.timesteps = scheduler_timesteps
-        self.scheduler.sigmas = scheduler_sigmas
-        self.scheduler._step_index = None
-        self.scheduler._begin_index = None
+        self.algorithm.prepare_rollout_data(rollout_data, loss_ctx)
 
         # ------------- Micro-batch schedule -------------
         num_optim_steps_per_rollout = self.args.num_steps_per_rollout
@@ -371,17 +364,10 @@ def _train_core(self, rollout_id: int, rollout_data) -> None:
 
                 for pair_lo, pair_hi in microbatch_ranges:
                     chunk = train_pairs[pair_lo:pair_hi]
-                    loss_sum = self._forward_train_pair_batch(
+                    loss_sum = self.algorithm.compute_loss(
+                        loss_ctx,
                         chunk,
-                        use_cfg=use_cfg,
-                        guidance_scale=guidance_scale,
-                        true_cfg_scale=true_cfg_scale,
-                        clip_range=clip_range,
-                        noise_level=noise_level,
-                        num_train_timesteps=num_train_timesteps,
                         log_stats=log_stats,
-                        device=device,
-                        kl_beta=kl_beta,
                         pad_to_len=legacy_pad_to_len,
                     )
                     if not self.args.debug_skip_optimizer_step:
@@ -423,239 +409,6 @@ def _maybe_legacy_window_pad_len(self, train_pairs: list, microbatch_ranges: lis
                     conds.append(env.neg_cond_kwargs)
         return self.train_pipeline_config.maybe_legacy_window_pad_len(conds)
 
-    def _forward_train_pair_batch(
-        self,
-        batch: list,
-        *,
-        use_cfg: bool,
-        guidance_scale: float,
-        true_cfg_scale: float | None,
-        clip_range: float,
-        noise_level: float,
-        num_train_timesteps: int,
-        log_stats: dict[str, list[torch.Tensor]],
-        device: torch.device,
-        kl_beta: float = 0.0,
-        pad_to_len: int | None = None,
-    ) -> torch.Tensor:
-        """One DiT forward + PPO loss over ``len(batch)`` train pairs. Returns sum of per-pair losses."""
-        forward_dtype = self._forward_dtype
-        train_pipeline_config = self.train_pipeline_config
-        bsz = len(batch)
-
-        def _stack(key):
-            return torch.stack([pair[key] for pair in batch]).to(device=device, dtype=torch.float32)
-
-        latents_microbatch = _stack("latent")  # (bsz, *latent_dims)
-        next_latents_microbatch = _stack("next_latent")  # (bsz, *latent_dims)
-        timesteps_microbatch = _stack("timestep")  # (bsz,) -- per-pair timestep is scalar
-        next_timesteps_microbatch = _stack("next_timestep")  # (bsz,) -- next rollout timestep (0 at terminal)
-        log_prob_old_microbatch = _stack("log_prob_old")  # (bsz,) -- per-pair log_prob is scalar
-
-        advantage = torch.tensor(  # (bsz,)
-            [float(pair["advantage"]) for pair in batch],
-            device=device,
-            dtype=torch.float32,
-        )
-        advantage = torch.clamp(advantage, -self.args.diffusion_adv_clip_max, self.args.diffusion_adv_clip_max)
-
-        if len(self.models) == 1:
-            component, model = next(iter(self.models.items()))
-        else:
-            components = {
-                train_pipeline_config.component_for_timestep(t, num_train_timesteps)
-                for t in timesteps_microbatch.tolist()
-            }
-            # to prevent mixing denoising phases in a single micro-batch
-            # Just in case when some customized step strategy is used that
-            # may violate the assumption of one phase per micro-batch, we raise an error here
-            if len(components) > 1:
-                raise ValueError(
-                    f"Micro-batch mixes denoising phases {sorted(components)}; set "
-                    "--micro-batch-size 1 so each forward is phase-pure (one DiT, one CFG scale)."
-                )
-            component = components.pop()
-            model = self.models[component]
-            guidance_scale = train_pipeline_config.select_guidance_scale(
-                float(timesteps_microbatch[0]),
-                num_train_timesteps,
-                guidance_scale,
-                self.args.diffusion_guidance_scale_2,
-            )
-
-        # sgl-d's Qwen DiT divides timestep by num_train_timesteps inside
-        # forward; diffusers' does not. SD3 already expects raw timesteps.
-        if train_pipeline_config.needs_timestep_scaling:
-            timesteps_for_model = timesteps_microbatch / float(num_train_timesteps)
-        else:
-            timesteps_for_model = timesteps_microbatch
-
-        pos_list = [
-            train_pipeline_config.prepare_cond_kwargs(batch[i]["denoising_env"].pos_cond_kwargs, device)
-            for i in range(bsz)
-        ]
-        neg_list = (
-            [
-                train_pipeline_config.prepare_cond_kwargs(batch[i]["denoising_env"].neg_cond_kwargs, device)
-                for i in range(bsz)
-            ]
-            if use_cfg
-            else None
-        )
-
-        # Collate cond once, up front. With CFG batching, pos+neg must share one
-        # padded width and go through a single joint forward, so build that joint cond
-        # directly; otherwise build pos (and neg) separately. (A single-sample
-        # timestep-stacked micro-batch is just collate of bsz copies of one sample --
-        # bitwise-equivalent to the old expand_cond_for_timestep_batch path; the
-        # all-True mask qwen adds is a verified forward no-op, see
-        # tests/manual/check_mask_equivalence.py.)
-        cfg_batching = use_cfg and bool(self.args.fsdp_cfg_batching)
-        joint_cond = None
-        pos_cond_microbatch = None
-        neg_cond_microbatch = None
-        if cfg_batching:
-            joint_cond = _cast_cond_to_dtype(
-                train_pipeline_config.collate_cond_for_sample_batch(
-                    pos_list + neg_list, device, pad_to_len=pad_to_len
-                ),
-                forward_dtype,
-            )
-        else:
-            pos_cond_microbatch = _cast_cond_to_dtype(
-                train_pipeline_config.collate_cond_for_sample_batch(pos_list, device, pad_to_len=pad_to_len),
-                forward_dtype,
-            )
-            if use_cfg and neg_list is not None:
-                neg_cond_microbatch = _cast_cond_to_dtype(
-                    train_pipeline_config.collate_cond_for_sample_batch(neg_list, device, pad_to_len=pad_to_len),
-                    forward_dtype,
-                )
-
-        # Cast inputs explicitly: FSDP MixedPrecisionPolicy casts params but
-        # leaves fp32 inputs, which would run first matmul at higher precision
-        # than rollout → systematic noise_pred drift.
-        latents_input = latents_microbatch.to(forward_dtype)
-        timesteps_input = timesteps_for_model.to(forward_dtype)
-
-        def _compute_noise_pred(disable_adapter: bool = False) -> torch.Tensor:
-            adapter_ctx = model.disable_adapter() if disable_adapter else nullcontext()
-            with adapter_ctx:
-                return train_pipeline_config.compute_noise_pred(
-                    model=model,
-                    latents_input=latents_input,
-                    timesteps_input=timesteps_input,
-                    pos_cond=pos_cond_microbatch,
-                    neg_cond=neg_cond_microbatch,
-                    joint_cond=joint_cond,
-                    use_cfg=use_cfg,
-                    cfg_batching=cfg_batching,
-                    guidance_scale=guidance_scale,
-                    true_cfg_scale=true_cfg_scale,
-                )
-
-        noise_pred_microbatch = _compute_noise_pred()
-
-        _, log_prob_new_microbatch, prev_sample_mean_new, std_dev_t_new = self.sde_backend.sde_step_logprob(
-            noise_pred_microbatch.float(),
-            timesteps_microbatch,
-            next_timesteps_microbatch,
-            latents_microbatch.float(),
-            prev_sample=next_latents_microbatch.float(),
-            noise_level=noise_level,
-        )
-
-        log_prob_new = log_prob_new_microbatch  # (bsz,) -- sde_step_with_logprob means over non-batch dims
-        log_prob_old = log_prob_old_microbatch  # (bsz,)
-        ratio = torch.exp(log_prob_new - log_prob_old)  # (bsz,)
-        unclipped = -advantage * ratio
-        clipped = -advantage * torch.clamp(ratio, 1.0 - clip_range, 1.0 + clip_range)
-        per_pair_loss = torch.maximum(unclipped, clipped)
-        loss_sum = per_pair_loss.sum()
-
-        # ------------- KL loss (vs LoRA base model as reference) -------------
-        kl_loss = loss_sum.new_zeros(())
-        if kl_beta > 0:
-            with torch.no_grad():
-                ref_noise_pred_microbatch = _compute_noise_pred(disable_adapter=True)
-                # TODO: unify sde_step_with_logprob with rollout and trainer forward paths.
-                _, _, prev_sample_mean_ref, _ = self.sde_backend.sde_step_logprob(
-                    ref_noise_pred_microbatch.float(),
-                    timesteps_microbatch,
-                    next_timesteps_microbatch,
-                    latents_microbatch.float(),
-                    prev_sample=next_latents_microbatch.float(),
-                    noise_level=noise_level,
-                )
-            kl_per_pair = ((prev_sample_mean_new - prev_sample_mean_ref) ** 2).mean(
-                dim=tuple(range(1, prev_sample_mean_new.ndim)),
-                keepdim=True,
-            ) / (2 * std_dev_t_new**2)
-            loss_sum = loss_sum + kl_beta * kl_per_pair.sum()
-            kl_loss = kl_per_pair.mean()
-
-        with torch.no_grad():
-            log_stats["loss"].append((per_pair_loss.mean() + kl_beta * kl_loss).detach())
-            log_stats["policy_loss"].append(per_pair_loss.mean().detach())
-            log_stats["kl_loss"].append(kl_loss.detach())
-            log_stats["loss_abs_mean"].append(per_pair_loss.abs().mean().detach())
-            log_stats["adv_abs_mean"].append(advantage.abs().mean().detach())
-            log_stats["ratio_abs_minus_1"].append((ratio - 1.0).abs().mean().detach())
-            log_stats["approx_kl"].append(0.5 * torch.mean((log_prob_new - log_prob_old) ** 2).detach())
-            log_stats["clipfrac"].append(torch.mean((torch.abs(ratio - 1.0) > clip_range).float()).detach())
-            log_stats["log_prob_new_idx_0"].append(log_prob_new[0].detach())
-            log_stats["log_prob_old_idx_0"].append(log_prob_old[0].detach())
-            log_prob_mean_abs_diff = torch.mean(torch.abs(log_prob_new - log_prob_old)).detach()
-            log_stats["log_prob_mean_abs_diff"].append(log_prob_mean_abs_diff)
-            if len(self.models) > 1:
-                log_stats[f"log_prob_mean_abs_diff_{component}"].append(log_prob_mean_abs_diff)
-
-            # model_output_* checks the train forward reproduces the rollout forward -- the only
-            # model-dependent consistency metric (std_dev/prev_sample_mean are deterministic
-            # functions of it). Matches the legacy actor metric name.
-            rollout_model_output = stack_train_pair_rollout_debug(batch, "rollout_step_model_output")
-            if rollout_model_output is not None:
-                mean_abs_diff = _append_rollout_train_abs_diff_stats(
-                    log_stats,
-                    "model_output",
-                    noise_pred_microbatch.float(),
-                    rollout_model_output.to(device=device, dtype=torch.float32),
-                )
-                if len(self.models) > 1:
-                    log_stats[f"model_output_mean_abs_diff_{component}"].append(mean_abs_diff)
-
-        return loss_sum
-
-
-def _append_rollout_train_abs_diff_stats(
-    log_stats: dict[str, list],
-    prefix: str,
-    train: torch.Tensor,
-    rollout: torch.Tensor,
-) -> torch.Tensor:
-    bsz = train.shape[0]
-    diff = (train.reshape(bsz, -1).float() - rollout.reshape(bsz, -1).float()).abs()
-    ref_max = rollout.reshape(bsz, -1).float().abs().max() + 1e-30
-    mean_abs_diff = diff.mean().detach()
-    log_stats[f"{prefix}_max_abs_diff"].append(diff.max().detach())
-    log_stats[f"{prefix}_mean_abs_diff"].append(mean_abs_diff)
-    log_stats[f"{prefix}_rel_max"].append((diff.max() / ref_max).detach())
-    return mean_abs_diff
-
-
-def _cast_cond_to_dtype(cond: dict, dtype: torch.dtype) -> dict:
-    """Cast floating-point tensors to the model's compute dtype; leave bool
-    masks / int / list / scalar values untouched. The bool
-    encoder_hidden_states_mask must NOT be cast.
-    """
-    out: dict = {}
-    for k, v in cond.items():
-        if isinstance(v, torch.Tensor) and v.dtype.is_floating_point:
-            out[k] = v.to(dtype)
-        else:
-            out[k] = v
-    return out
-
 
 @torch.no_grad()
 def move_torch_optimizer(optimizer, device):
diff --git a/miles/ray/rollout.py b/miles/ray/rollout.py
index 805e8241..6fb12579 100644
--- a/miles/ray/rollout.py
+++ b/miles/ray/rollout.py
@@ -12,6 +12,7 @@
 from ray.util.scheduling_strategies import PlacementGroupSchedulingStrategy
 from sglang.srt.constants import GPU_MEMORY_TYPE_WEIGHTS
 
+from miles.algorithms.registry import load_algorithm
 from miles.backends.sglang_diffusion_utils.sglang_diffusion_engine import SGLangDiffusionEngine
 from miles.rollout.base_types import call_rollout_fn
 from miles.utils import tracking_utils
@@ -24,7 +25,7 @@
 from miles.utils.misc import load_function
 from miles.utils.ray_utils import Box
 from miles.utils.tracking_utils import init_tracking
-from miles.utils.train_data_utils import RolloutTrainDataConverter, TrainDataDPSplitter, reorder_train_pairs_for_tiling
+from miles.utils.train_data_utils import TrainDataDPSplitter, reorder_train_pairs_for_tiling
 from miles.utils.types import Sample
 
 from .utils import NOSET_VISIBLE_DEVICES_ENV_VARS_LIST, Lock
@@ -70,10 +71,12 @@ def __init__(self, args, pg):
             if self.args.custom_convert_samples_to_train_data_path is not None
             else None
         )
-        self.train_data_converter = RolloutTrainDataConverter()
+        self.algorithm = load_algorithm(args)
+        self.collection_spec = self.algorithm.collection_spec()
         self.train_data_dp_splitter = TrainDataDPSplitter()
         logger.info(f"import {self.args.rollout_function_path} as generate_rollout function.")
         logger.info(f"import {self.args.eval_function_path} as eval_generate_rollout function.")
+        logger.info("RolloutManager diffusion algorithm=%s", self.algorithm.name)
 
         logger.info("RolloutManager rollout_num_gpus=%s", self.args.rollout_num_gpus)
 
@@ -316,29 +319,9 @@ def _post_process_rewards(self, samples: list[Sample] | list[list[Sample]]):
         if self.custom_reward_post_process_func is not None:
             return self.custom_reward_post_process_func(self.args, samples)
 
-        raw_rewards = [sample.get_reward_value(self.args) for sample in samples]
-
-        # --globalize-reward-mean / --globalize-reward-std are orthogonal. flow_grpo
-        # pickscore_qwenimage uses per-prompt mean + global std (PerPromptStatTracker
-        # with global_std=True), which is --globalize-reward-std alone.
-        rewards_flat = torch.tensor(raw_rewards, dtype=torch.float)
-        rewards = rewards_flat.view(-1, self.args.n_samples_per_prompt)
-
-        if self.args.globalize_reward_mean:
-            mean = rewards_flat.mean()
-        else:
-            mean = rewards.mean(dim=-1, keepdim=True)
-        rewards = rewards - mean
-
-        if self.args.grpo_std_normalization:
-            if self.args.globalize_reward_std:
-                std = rewards_flat.std()
-            else:
-                std = rewards.std(dim=-1, keepdim=True)
-            # matches flow_grpo's `+ 1e-4` in both stat_tracking branches
-            rewards = rewards / (std + 1e-4)
-
-        return raw_rewards, rewards.flatten().tolist()
+        labels = self.algorithm.postprocess_rewards(self.args, samples)
+        advantages = labels.advantages if labels.advantages is not None else labels.raw_rewards
+        return labels.raw_rewards, advantages
 
     def _convert_samples_to_train_data(self, samples: list[Sample] | list[list[Sample]]):
         """
@@ -347,7 +330,15 @@ def _convert_samples_to_train_data(self, samples: list[Sample] | list[list[Sampl
         if self.custom_convert_samples_to_train_data_func is not None:
             return self.custom_convert_samples_to_train_data_func(self.args, samples)
 
-        raw_rewards, rewards = self._post_process_rewards(samples)
+        if self.custom_reward_post_process_func is not None:
+            raw_rewards, rewards = self._post_process_rewards(samples)
+            from miles.algorithms.base import TrainLabels
+
+            labels = TrainLabels(raw_rewards=raw_rewards, advantages=rewards)
+        else:
+            labels = self.algorithm.postprocess_rewards(self.args, samples)
+            raw_rewards = labels.raw_rewards
+            rewards = labels.advantages if labels.advantages is not None else labels.raw_rewards
 
         assert len(raw_rewards) == len(samples)
         assert len(rewards) == len(samples)
@@ -387,7 +378,7 @@ def _convert_samples_to_train_data(self, samples: list[Sample] | list[list[Sampl
                 reward_key=self.args.reward_key,
             )
 
-        return self.train_data_converter.convert_samples(samples, rewards, raw_rewards)
+        return self.algorithm.build_train_data(self.args, samples, labels)
 
     def _log_images(
         self,
diff --git a/miles/utils/arguments.py b/miles/utils/arguments.py
index 92174875..8e26331e 100644
--- a/miles/utils/arguments.py
+++ b/miles/utils/arguments.py
@@ -258,6 +258,22 @@ def add_rollout_arguments(parser):
                 default="stabilityai/stable-diffusion-3.5-medium",
                 help="HuggingFace model id for diffusion rollout.",
             )
+            parser.add_argument(
+                "--diffusion-algorithm",
+                type=str,
+                default="flow_grpo",
+                choices=["flow_grpo"],
+                help=(
+                    "Diffusion training algorithm plugin. Currently supports flow_grpo only "
+                    "(reverse-SDE PPO-clip GRPO). Use --diffusion-algorithm-path for experiments."
+                ),
+            )
+            parser.add_argument(
+                "--diffusion-algorithm-path",
+                type=str,
+                default=None,
+                help="Optional dotted class path overriding --diffusion-algorithm.",
+            )
             parser.add_argument(
                 "--train-pipeline-config-path",
                 type=str,
@@ -1368,6 +1384,11 @@ def miles_validate_args(args):
 
     args.rollout_patch_groups = ["sgld"] if args.apply_sgld_monkey_patches else []
 
+    # Resolve diffusion algorithm class path early so rollout/train share one identity.
+    from miles.algorithms.registry import resolve_algorithm_class_path
+
+    args.diffusion_algorithm_path = resolve_algorithm_class_path(args)
+
     if getattr(args, "diffusion_model", None):
         from miles.utils.misc import load_function
 
diff --git a/miles/utils/diffusion_protocol.py b/miles/utils/diffusion_protocol.py
index a2e86fff..bd145704 100644
--- a/miles/utils/diffusion_protocol.py
+++ b/miles/utils/diffusion_protocol.py
@@ -8,8 +8,9 @@
 
 @dataclass(frozen=True)
 class DiffusionRolloutSpec:
-    # Required rollout keys to reconstruct per-step log_prob_new and PPO ratio in training.
-    # latents and next latents for log_prob_new in training, log_prob_old used with log_prob_new (in training) to get ratio.
+    # Flow-GRPO default: reconstruct per-step log_prob_new and PPO ratio in training.
+    # Forward-matching algorithms (SFT/AWM/NFT) validate via DiffusionAlgorithm.validate_train_batch
+    # instead of this global spec.
     required_keys: tuple[str, ...] = ("timesteps", "sigmas", "latents", "next_latents", "log_prob_old")
     # Optional rollout keys for KL regularization or debugging.
     # mean of distribution p(x_{t+1} | x_t), for KL
@@ -18,6 +19,7 @@ class DiffusionRolloutSpec:
 
 @dataclass(frozen=True)
 class DiffusionTrainSpec:
+    # Flow-GRPO train-side keys. Prefer algorithm.validate_train_batch for new algorithms.
     required_keys: tuple[str, ...] = ("log_prob_old", "log_prob_new", "advantage")
 
 

From d79ea3fe4949ede174ebc1c990ed168ac5c152c3 Mon Sep 17 00:00:00 2001
From: niehen6174 <niehen6174@gmail.com>
Date: Mon, 20 Jul 2026 11:50:30 +0000
Subject: [PATCH 3/5] test(algorithms): cover Flow-GRPO registry load and
 advantage labels

Add fast tests for builtin registration and group-relative advantage
normalization used by the default algorithm plugin.
---
 .../fast/algorithms/test_flow_grpo_plugin.py  | 44 +++++++++++++++++++
 1 file changed, 44 insertions(+)
 create mode 100644 tests/fast/algorithms/test_flow_grpo_plugin.py

diff --git a/tests/fast/algorithms/test_flow_grpo_plugin.py b/tests/fast/algorithms/test_flow_grpo_plugin.py
new file mode 100644
index 00000000..527514da
--- /dev/null
+++ b/tests/fast/algorithms/test_flow_grpo_plugin.py
@@ -0,0 +1,44 @@
+from argparse import Namespace
+from types import SimpleNamespace
+
+from miles.algorithms.flow_grpo import FlowGRPOAlgorithm
+from miles.algorithms.labels import grpo_group_advantages
+from miles.algorithms.registry import builtin_algorithm_names, load_algorithm, resolve_algorithm_class_path
+from miles.utils.types import Sample
+
+
+def test_builtin_is_flow_grpo_only():
+    assert builtin_algorithm_names() == ["flow_grpo"]
+
+
+def test_load_flow_grpo_algorithm():
+    args = Namespace(
+        diffusion_algorithm="flow_grpo",
+        diffusion_algorithm_path=None,
+        use_lora=False,
+        diffusion_kl_beta=0.0,
+    )
+    assert resolve_algorithm_class_path(args).endswith("FlowGRPOAlgorithm")
+    algo = load_algorithm(args)
+    assert isinstance(algo, FlowGRPOAlgorithm)
+    spec = algo.collection_spec()
+    assert spec.mode == "online"
+    assert spec.needs_logprob is True
+    assert spec.needs_trajectory is True
+
+
+def test_grpo_group_advantages():
+    args = SimpleNamespace(
+        n_samples_per_prompt=2,
+        globalize_reward_mean=False,
+        globalize_reward_std=False,
+        grpo_std_normalization=True,
+        reward_key=None,
+    )
+    samples = [
+        Sample(prompt="a", reward=1.0),
+        Sample(prompt="a", reward=3.0),
+    ]
+    labels = grpo_group_advantages(args, samples)
+    assert len(labels.advantages) == 2
+    assert abs(sum(labels.advantages)) < 1e-5

From e2fcba51f86155d4fed02303e714f138bc376871 Mon Sep 17 00:00:00 2001
From: niehen6174 <niehen6174@gmail.com>
Date: Mon, 20 Jul 2026 12:11:27 +0000
Subject: [PATCH 4/5] refactor(algorithms): rename TrainLabels to TrainSignals

Use a clearer name for reward-derived training signals (advantages /
future NFT weights), and rename labels.py to signals.py accordingly.
---
 miles/algorithms/__init__.py                   |  4 ++--
 miles/algorithms/base.py                       | 15 +++++++++------
 miles/algorithms/flow_grpo.py                  | 12 ++++++------
 miles/algorithms/{labels.py => signals.py}     |  8 ++++----
 miles/ray/rollout.py                           | 18 +++++++++---------
 tests/fast/algorithms/test_flow_grpo_plugin.py |  8 ++++----
 6 files changed, 34 insertions(+), 31 deletions(-)
 rename miles/algorithms/{labels.py => signals.py} (72%)

diff --git a/miles/algorithms/__init__.py b/miles/algorithms/__init__.py
index c9443350..8ee5a07b 100644
--- a/miles/algorithms/__init__.py
+++ b/miles/algorithms/__init__.py
@@ -3,14 +3,14 @@
 Currently ships Flow-GRPO only; SFT / AWM / DiffusionNFT land in follow-up PRs.
 """
 
-from miles.algorithms.base import CollectionSpec, DiffusionAlgorithm, TrainLabels, TrainLossContext
+from miles.algorithms.base import CollectionSpec, DiffusionAlgorithm, TrainLossContext, TrainSignals
 from miles.algorithms.registry import builtin_algorithm_names, load_algorithm, resolve_algorithm_class_path
 
 __all__ = [
     "CollectionSpec",
     "DiffusionAlgorithm",
-    "TrainLabels",
     "TrainLossContext",
+    "TrainSignals",
     "builtin_algorithm_names",
     "load_algorithm",
     "resolve_algorithm_class_path",
diff --git a/miles/algorithms/base.py b/miles/algorithms/base.py
index aa876c1c..5655733c 100644
--- a/miles/algorithms/base.py
+++ b/miles/algorithms/base.py
@@ -53,13 +53,16 @@ class LossOutput:
 
 
 @dataclass
-class TrainLabels:
-    """Post-reward labels attached to samples before ``build_train_data``."""
+class TrainSignals:
+    """Reward-derived training signals attached to samples before ``build_train_data``.
+
+    Not classification labels: e.g. GRPO advantages or NFT soft +/− weights.
+    """
 
     raw_rewards: list[float]
     advantages: list[float] | None = None
-    # Reserved for DiffusionNFT soft positive/negative labels; Flow-GRPO ignores it.
-    nft_labels: list[float] | None = None
+    # Reserved for DiffusionNFT soft positive/negative weights; Flow-GRPO ignores it.
+    nft_signals: list[float] | None = None
 
 
 @runtime_checkable
@@ -72,9 +75,9 @@ def collection_spec(self) -> CollectionSpec:
         """Return acquisition contract; see ``CollectionSpec`` — not fully consumed yet."""
         ...
 
-    def postprocess_rewards(self, args, samples: list) -> TrainLabels: ...
+    def postprocess_rewards(self, args, samples: list) -> TrainSignals: ...
 
-    def build_train_data(self, args, samples: list, labels: TrainLabels) -> dict[str, Any]: ...
+    def build_train_data(self, args, samples: list, signals: TrainSignals) -> dict[str, Any]: ...
 
     def validate_train_batch(self, batch: list[dict]) -> list[str]: ...
 
diff --git a/miles/algorithms/flow_grpo.py b/miles/algorithms/flow_grpo.py
index fab2c5e6..fb202f2e 100644
--- a/miles/algorithms/flow_grpo.py
+++ b/miles/algorithms/flow_grpo.py
@@ -6,8 +6,8 @@
 
 import torch
 
-from miles.algorithms.base import CollectionSpec, TrainLabels, TrainLossContext
-from miles.algorithms.labels import grpo_group_advantages
+from miles.algorithms.base import CollectionSpec, TrainLossContext, TrainSignals
+from miles.algorithms.signals import grpo_group_advantages
 from miles.algorithms.train_forward_utils import (
     append_rollout_train_abs_diff_stats,
     compute_noise_pred,
@@ -40,12 +40,12 @@ def collection_spec(self) -> CollectionSpec:
             sync_weights_to_rollout=True,
         )
 
-    def postprocess_rewards(self, args, samples: list[Sample]) -> TrainLabels:
+    def postprocess_rewards(self, args, samples: list[Sample]) -> TrainSignals:
         return grpo_group_advantages(args, samples)
 
-    def build_train_data(self, args, samples: list[Sample], labels: TrainLabels) -> dict[str, Any]:
-        rewards = labels.advantages if labels.advantages is not None else labels.raw_rewards
-        return RolloutTrainDataConverter().convert_samples(samples, rewards, labels.raw_rewards)
+    def build_train_data(self, args, samples: list[Sample], signals: TrainSignals) -> dict[str, Any]:
+        rewards = signals.advantages if signals.advantages is not None else signals.raw_rewards
+        return RolloutTrainDataConverter().convert_samples(samples, rewards, signals.raw_rewards)
 
     def validate_train_batch(self, batch: list[dict]) -> list[str]:
         errors: list[str] = []
diff --git a/miles/algorithms/labels.py b/miles/algorithms/signals.py
similarity index 72%
rename from miles/algorithms/labels.py
rename to miles/algorithms/signals.py
index 721c8c31..4c88fb39 100644
--- a/miles/algorithms/labels.py
+++ b/miles/algorithms/signals.py
@@ -1,14 +1,14 @@
-"""Shared reward → label helpers for GRPO-style algorithms."""
+"""Shared reward → train-signal helpers for GRPO-style algorithms."""
 
 from __future__ import annotations
 
 import torch
 
-from miles.algorithms.base import TrainLabels
+from miles.algorithms.base import TrainSignals
 from miles.utils.types import Sample
 
 
-def grpo_group_advantages(args, samples: list[Sample]) -> TrainLabels:
+def grpo_group_advantages(args, samples: list[Sample]) -> TrainSignals:
     """Group-relative advantage normalization (Flow-GRPO default)."""
     raw_rewards = [sample.get_reward_value(args) for sample in samples]
     rewards_flat = torch.tensor(raw_rewards, dtype=torch.float)
@@ -27,4 +27,4 @@ def grpo_group_advantages(args, samples: list[Sample]) -> TrainLabels:
             std = rewards.std(dim=-1, keepdim=True)
         centered = centered / (std + 1e-4)
 
-    return TrainLabels(raw_rewards=raw_rewards, advantages=centered.flatten().tolist())
+    return TrainSignals(raw_rewards=raw_rewards, advantages=centered.flatten().tolist())
diff --git a/miles/ray/rollout.py b/miles/ray/rollout.py
index 6fb12579..ee951c55 100644
--- a/miles/ray/rollout.py
+++ b/miles/ray/rollout.py
@@ -319,9 +319,9 @@ def _post_process_rewards(self, samples: list[Sample] | list[list[Sample]]):
         if self.custom_reward_post_process_func is not None:
             return self.custom_reward_post_process_func(self.args, samples)
 
-        labels = self.algorithm.postprocess_rewards(self.args, samples)
-        advantages = labels.advantages if labels.advantages is not None else labels.raw_rewards
-        return labels.raw_rewards, advantages
+        signals = self.algorithm.postprocess_rewards(self.args, samples)
+        advantages = signals.advantages if signals.advantages is not None else signals.raw_rewards
+        return signals.raw_rewards, advantages
 
     def _convert_samples_to_train_data(self, samples: list[Sample] | list[list[Sample]]):
         """
@@ -332,13 +332,13 @@ def _convert_samples_to_train_data(self, samples: list[Sample] | list[list[Sampl
 
         if self.custom_reward_post_process_func is not None:
             raw_rewards, rewards = self._post_process_rewards(samples)
-            from miles.algorithms.base import TrainLabels
+            from miles.algorithms.base import TrainSignals
 
-            labels = TrainLabels(raw_rewards=raw_rewards, advantages=rewards)
+            signals = TrainSignals(raw_rewards=raw_rewards, advantages=rewards)
         else:
-            labels = self.algorithm.postprocess_rewards(self.args, samples)
-            raw_rewards = labels.raw_rewards
-            rewards = labels.advantages if labels.advantages is not None else labels.raw_rewards
+            signals = self.algorithm.postprocess_rewards(self.args, samples)
+            raw_rewards = signals.raw_rewards
+            rewards = signals.advantages if signals.advantages is not None else signals.raw_rewards
 
         assert len(raw_rewards) == len(samples)
         assert len(rewards) == len(samples)
@@ -378,7 +378,7 @@ def _convert_samples_to_train_data(self, samples: list[Sample] | list[list[Sampl
                 reward_key=self.args.reward_key,
             )
 
-        return self.algorithm.build_train_data(self.args, samples, labels)
+        return self.algorithm.build_train_data(self.args, samples, signals)
 
     def _log_images(
         self,
diff --git a/tests/fast/algorithms/test_flow_grpo_plugin.py b/tests/fast/algorithms/test_flow_grpo_plugin.py
index 527514da..1cc17f4f 100644
--- a/tests/fast/algorithms/test_flow_grpo_plugin.py
+++ b/tests/fast/algorithms/test_flow_grpo_plugin.py
@@ -2,8 +2,8 @@
 from types import SimpleNamespace
 
 from miles.algorithms.flow_grpo import FlowGRPOAlgorithm
-from miles.algorithms.labels import grpo_group_advantages
 from miles.algorithms.registry import builtin_algorithm_names, load_algorithm, resolve_algorithm_class_path
+from miles.algorithms.signals import grpo_group_advantages
 from miles.utils.types import Sample
 
 
@@ -39,6 +39,6 @@ def test_grpo_group_advantages():
         Sample(prompt="a", reward=1.0),
         Sample(prompt="a", reward=3.0),
     ]
-    labels = grpo_group_advantages(args, samples)
-    assert len(labels.advantages) == 2
-    assert abs(sum(labels.advantages)) < 1e-5
+    signals = grpo_group_advantages(args, samples)
+    assert len(signals.advantages) == 2
+    assert abs(sum(signals.advantages)) < 1e-5

From f4b061d3937fffabc3c0c3d87197a51ea8cd6482 Mon Sep 17 00:00:00 2001
From: niehen6174 <niehen6174@gmail.com>
Date: Mon, 20 Jul 2026 12:37:15 +0000
Subject: [PATCH 5/5] style: fix pre-commit ruff/isort/black on algorithm
 plugin files

Remove unused bsz in Flow-GRPO compute_loss and apply formatter fixes
so the pre-commit CI job passes.
---
 miles/algorithms/base.py                | 2 +-
 miles/algorithms/flow_grpo.py           | 7 +++++--
 miles/algorithms/train_forward_utils.py | 4 +---
 miles/backends/fsdp_utils/actor.py      | 6 ++----
 4 files changed, 9 insertions(+), 10 deletions(-)

diff --git a/miles/algorithms/base.py b/miles/algorithms/base.py
index 5655733c..a94e4fad 100644
--- a/miles/algorithms/base.py
+++ b/miles/algorithms/base.py
@@ -92,4 +92,4 @@ def compute_loss(
 
     def prepare_rollout_data(self, rollout_data: dict, ctx: TrainLossContext) -> None:
         """Optional hook before the micro-batch loop (e.g. sync scheduler meta)."""
-        ...
\ No newline at end of file
+        ...
diff --git a/miles/algorithms/flow_grpo.py b/miles/algorithms/flow_grpo.py
index fb202f2e..d7ae414f 100644
--- a/miles/algorithms/flow_grpo.py
+++ b/miles/algorithms/flow_grpo.py
@@ -15,7 +15,11 @@
     resolve_cfg_flags,
     select_model_for_timesteps,
 )
-from miles.utils.train_data_utils import RolloutTrainDataConverter, scheduler_meta_from_rollout, stack_train_pair_rollout_debug
+from miles.utils.train_data_utils import (
+    RolloutTrainDataConverter,
+    scheduler_meta_from_rollout,
+    stack_train_pair_rollout_debug,
+)
 from miles.utils.types import Sample
 
 
@@ -86,7 +90,6 @@ def compute_loss(
         forward_dtype = ctx.forward_dtype
         train_pipeline_config = ctx.train_pipeline_config
         device = ctx.device
-        bsz = len(batch)
 
         use_cfg, guidance_scale, true_cfg_scale = resolve_cfg_flags(args)
         clip_range = args.diffusion_clip_range
diff --git a/miles/algorithms/train_forward_utils.py b/miles/algorithms/train_forward_utils.py
index 694a94e4..690c3a56 100644
--- a/miles/algorithms/train_forward_utils.py
+++ b/miles/algorithms/train_forward_utils.py
@@ -50,9 +50,7 @@ def select_model_for_timesteps(
         component, model = next(iter(ctx.models.items()))
         return component, model, guidance_scale
 
-    components = {
-        train_pipeline_config.component_for_timestep(t, num_train_timesteps) for t in timesteps.tolist()
-    }
+    components = {train_pipeline_config.component_for_timestep(t, num_train_timesteps) for t in timesteps.tolist()}
     if len(components) > 1:
         raise ValueError(
             f"Micro-batch mixes denoising phases {sorted(components)}; set "
diff --git a/miles/backends/fsdp_utils/actor.py b/miles/backends/fsdp_utils/actor.py
index 9bd5baa9..67b95f21 100644
--- a/miles/backends/fsdp_utils/actor.py
+++ b/miles/backends/fsdp_utils/actor.py
@@ -21,10 +21,8 @@
 from miles.utils.profile_utils import TrainProfiler
 from miles.utils.timer import Timer, inverse_timer, timer
 from miles.utils.tracking_utils import init_tracking
-from miles.utils.train_data_utils import (
-    build_microbatch_schedule,
-    validate_same_microbatch_counts_across_dp,
-)
+from miles.utils.train_data_utils import build_microbatch_schedule, validate_same_microbatch_counts_across_dp
+
 from . import checkpoint
 from .diffusion_update_weight_utils import (
     DiffusionUpdateWeightFromTensor,