dew.objectives.rl.flow

Flow-GRPO over recorded rectified-flow transitions.

The policy ratio follows the released implementation’s mean coordinate log likelihood, not the product of all coordinate likelihood ratios. The KL term is the conditional Gaussian KL from arXiv:2505.05470v5 section 4, also averaged across coordinates. Its variance includes elapsed time. The released scripts/train_sd3.py at 879042cf5707f8b90daa98d147d7deac2317c5da instead divides squared mean displacement by the diffusion coefficient squared, a term equal to elapsed time times this conditional KL. These regularizers have different step weighting. Callback scores are retained in float64 through host grouping, unlike the released trainer’s earlier float32 score conversion.

Name	Summary
FlowGRPOObjective	Train a rectified-flow policy on clipped, coordinate-normalized gradients.
FlowReward	Score decoded [-1, 1] samples and repeated source rows, one scalar per sample.
FlowRollout	Collect complete prompt groups and the first train_steps transitions.

FlowGRPOObjective

class source

class FlowGRPOObjective(
    model: nn.Module,
    process: Process,
    inputs: InputSpec,
    *,
    sde: FlowSDE = FlowSDE(),
    beta: float = 0.0,
    clip_range: float = 0.0001,
    adv_clip_max: float = 5.0,
    autoencoder: AutoEncoder | None = None,
    guidance: CFG | None = CFG(3.0),
    sampler: Solver = Euler(),
    steps: int = 41,
    pretrained: Variables | None = None,
)

Train a rectified-flow policy on clipped, coordinate-normalized gradients.

A conditional transition KL regularizes it.

Batches carry latents/next_latents [N, K, …], timesteps/next_timesteps, joint old_log_probs and transition_mask [N, K], and advantages [N] or [N, K]. K is the selected transition count. The denominator counts kept stochastic transitions. Deterministic intervals contribute no policy loss.

beta > 0 freezes the initial denoiser in the existing EMA slot. Evaluation and previews always use the live policy. sampler and steps configure evaluation; sde specifies both rollout and rescoring. pretrained is a model variables dict, as returned by model.init; encoders and an optional autoencoder are supplied through the existing diffusion input contract.

FlowGRPOObjective.held_variables

def held_variables() -> Variables

FlowGRPOObjective.init

def init(key: jax.Array, variables: Variables | None = None) -> Variables

FlowGRPOObjective.log_probs

def log_probs(params: Variables, batch: Batch) -> jax.Array

Rescore joint transition log densities with the rollout’s guidance.

FlowGRPOObjective.loss

def loss(params: Variables, batch: Batch, step: Step) -> tuple[Mean, Aux]

Score the clipped policy gradient over the recorded transitions.

The scan carries nothing between transitions; each one contributes its surrogate, its KL to the frozen reference, whether it counted, and whether the ratio was clipped.

FlowGRPOObjective.evaluate

def evaluate(params: Variables, batch: Batch, step: Step)

Generate one live-policy sample per source row, including prompt-only batches.

FlowGRPOObjective.preview

def preview(params: Variables, batch: Batch, step: Step, *, scored=None)

Draw on all ranks; materialize before root-only caption decoding.

FlowReward

attribute source

Score decoded [-1, 1] samples and repeated source rows, one scalar per sample.

FlowRollout

dataclass source

class FlowRollout(
    objective: FlowGRPOObjective,
    reward: FlowReward,
    groups: int = 4,
    steps: int = 11,
    train_steps: int | None = None,
)

Collect complete prompt groups and the first train_steps transitions.

steps counts time points, so steps=11 draws ten transitions. None selects all transitions. Rewards use population group standard deviation and epsilon 1e-4, as Flow-GRPO’s PerPromptStatTracker does. Each prompt row defines a group; equal prompt text in other rows does not merge groups. Zero-advantage rows are masked, as the reference training loop filters them. Callback collection, JSON/byte transport, and population statistics retain float64 values. Host rewards remain float64; training advantages are float32 after normalization. The reward metric is a float32 diagnostic, and JAX device transfer also narrows the reward column when x64 is off.

The trainer supplies global arrays on every process. Generation remains collective, rewards run once on rank zero, and the result contains only this process’s owned rows for the trainer’s shard_batch boundary. Host materialization currently uses collective_host and replicates the complete trajectory on every process before selecting local rows.