fix(diffusion): keep replay transition means at SDE-native fp32

unirl/algorithms/diffusionopd.py

@@ -190,8 +190,7 @@ def compute_loss_and_backward(
             )
 
         teacher_means = gather_sde_field(segment.sde_means, segment.sde_indices, target_steps, field_name="sde_means")
-        # fp32 before squaring: bf16 loses the shrinking teacher-student delta.
-        student_f32 = student_means.float()
+        # Replay means are fp32 by contract; align the stored teacher anchor before squaring.
         teacher_f32 = teacher_means.to(device=student_means.device, dtype=torch.float32)
 
         sigma_t = _transition_sigma(
@@ -202,7 +201,7 @@ def compute_loss_and_backward(
             device=student_means.device,
             add_coefficient=self.add_kl_coefficient,
         )
-        kl_per_elem = _gaussian_kl_div(student_f32, teacher_f32, sigma_t)
+        kl_per_elem = _gaussian_kl_div(student_means, teacher_f32, sigma_t)
         kl_per_sample_step = kl_per_elem.mean(dim=tuple(range(2, kl_per_elem.ndim)))  # [B, S']
         loss = kl_per_sample_step.mean()
 

unirl/models/bagel/diffusion.py

@@ -387,7 +387,7 @@ def diffuse(
                 if log_prob is not None:
                     sde_logp_list.append(log_prob.to(dtype=self.logprob_dtype))
                     if prev_mean is not None:
-                        sde_means_list.append(prev_mean.detach().to(dtype=self.trajectory_dtype))
+                        sde_means_list.append(prev_mean.detach())
 
         positions_collected = [p for p, _ in stored_pairs]
         latents_stacked = torch.stack([t for _, t in stored_pairs], dim=0).unsqueeze(0)
@@ -469,7 +469,7 @@ def replay(
                 prev_sample_means.append(prev_mean)
 
         log_probs_t = torch.stack(log_probs, dim=0).unsqueeze(0).to(dtype=self.logprob_dtype)
-        means_t = torch.stack(prev_sample_means, dim=0).unsqueeze(0).to(dtype=self.trajectory_dtype)
+        means_t = torch.stack(prev_sample_means, dim=0).unsqueeze(0)
         return ReplayResult(log_probs=log_probs_t, prev_sample_means=means_t)
 
     def build_forward_kwargs(

unirl/models/boogu_image/diffusion.py

@@ -436,7 +436,7 @@ def replay(
                     prev_sample_means.append(prev_mean)
 
         log_probs_t = torch.stack(log_probs, dim=1).to(dtype=self.logprob_dtype)
-        means_t = torch.stack(prev_sample_means, dim=1).to(dtype=self.trajectory_dtype) if prev_sample_means else None
+        means_t = torch.stack(prev_sample_means, dim=1) if prev_sample_means else None
         return ReplayResult(log_probs=log_probs_t, prev_sample_means=means_t)
 
     def predict_noise_at_step(

unirl/models/flux2_klein/diffusion.py

@@ -463,7 +463,7 @@ def replay(
                 self.model.transformer.train()
 
         log_probs_t = torch.stack(log_probs, dim=1).to(dtype=self.logprob_dtype)
-        means_t = torch.stack(prev_sample_means, dim=1).to(dtype=self.trajectory_dtype) if prev_sample_means else None
+        means_t = torch.stack(prev_sample_means, dim=1) if prev_sample_means else None
         return ReplayResult(log_probs=log_probs_t, prev_sample_means=means_t)
 
     @staticmethod

unirl/models/hunyuan_image3/diffusion.py

@@ -713,7 +713,7 @@ def replay(
                     prev_sample_means.append(prev_mean)
 
         log_probs_t = torch.stack(log_probs, dim=1).to(dtype=self.logprob_dtype)
-        means_t = torch.stack(prev_sample_means, dim=1).to(dtype=self.trajectory_dtype) if prev_sample_means else None
+        means_t = torch.stack(prev_sample_means, dim=1) if prev_sample_means else None
         return ReplayResult(log_probs=log_probs_t, prev_sample_means=means_t)
 
     def predict_noise_at_step(

unirl/models/hunyuan_video10/diffusion.py

@@ -436,7 +436,7 @@ def replay(
                     prev_sample_means.append(prev_mean)
 
         log_probs_t = torch.stack(log_probs, dim=1).to(dtype=self.logprob_dtype)
-        means_t = torch.stack(prev_sample_means, dim=1).to(dtype=self.trajectory_dtype) if prev_sample_means else None
+        means_t = torch.stack(prev_sample_means, dim=1) if prev_sample_means else None
         return ReplayResult(log_probs=log_probs_t, prev_sample_means=means_t)
 
     def predict_noise_at_step(

unirl/models/hunyuan_video15/diffusion.py

@@ -522,7 +522,7 @@ def replay(
                     prev_sample_means.append(prev_mean)
 
         log_probs_t = torch.stack(log_probs, dim=1).to(dtype=self.logprob_dtype)
-        means_t = torch.stack(prev_sample_means, dim=1).to(dtype=self.trajectory_dtype) if prev_sample_means else None
+        means_t = torch.stack(prev_sample_means, dim=1) if prev_sample_means else None
         return ReplayResult(log_probs=log_probs_t, prev_sample_means=means_t)
 
     def predict_noise_at_step(

unirl/models/ltx2/diffusion.py

@@ -450,7 +450,7 @@ def replay(
                     prev_sample_means.append(prev_mean)
 
         log_probs_t = torch.stack(log_probs, dim=1).to(dtype=self.logprob_dtype)
-        means_t = torch.stack(prev_sample_means, dim=1).to(dtype=self.trajectory_dtype) if prev_sample_means else None
+        means_t = torch.stack(prev_sample_means, dim=1) if prev_sample_means else None
         return ReplayResult(log_probs=log_probs_t, prev_sample_means=means_t)
 
 

unirl/models/qwen_image/diffusion.py

@@ -483,7 +483,7 @@ def replay(
                     prev_sample_means.append(prev_mean)
 
         log_probs_t = torch.stack(log_probs, dim=1).to(dtype=self.logprob_dtype)
-        means_t = torch.stack(prev_sample_means, dim=1).to(dtype=self.trajectory_dtype) if prev_sample_means else None
+        means_t = torch.stack(prev_sample_means, dim=1) if prev_sample_means else None
         return ReplayResult(log_probs=log_probs_t, prev_sample_means=means_t)
 
     def _batched_step_kwargs(self, segment: LatentSegment, params: DiffusionSamplingParams) -> Dict[str, Any]:

unirl/models/sd3/diffusion.py

@@ -397,7 +397,7 @@ def replay(
                     prev_sample_means.append(prev_mean)
 
         log_probs_t = torch.stack(log_probs, dim=1).to(dtype=self.logprob_dtype)
-        means_t = torch.stack(prev_sample_means, dim=1).to(dtype=self.trajectory_dtype) if prev_sample_means else None
+        means_t = torch.stack(prev_sample_means, dim=1) if prev_sample_means else None
         return ReplayResult(log_probs=log_probs_t, prev_sample_means=means_t)
 
     @staticmethod

unirl/models/sensenova_u1/diffusion.py

@@ -316,7 +316,7 @@ def _diffuse_one(
             if log_prob is not None:
                 log_probs.append(log_prob.to(dtype=self.logprob_dtype))
                 if mean is not None:
-                    means.append(mean.detach().to(dtype=self.trajectory_dtype))
+                    means.append(mean.detach())
 
         return LatentSegment(
             latents=torch.stack([value for _, value in stored], dim=1),
@@ -444,7 +444,7 @@ def replay(
 
         return ReplayResult(
             log_probs=torch.cat(batch_log_probs, dim=0).to(dtype=self.logprob_dtype),
-            prev_sample_means=torch.cat(batch_means, dim=0).to(dtype=self.trajectory_dtype),
+            prev_sample_means=torch.cat(batch_means, dim=0),
         )
 
     def predict_noise_at_step(

unirl/models/types/batched_replay.py

@@ -67,7 +67,7 @@ def _replay_batched_steps(
         means_t = None
         if prev_mean_all is not None:
             tail = prev_mean_all.shape[1:]
-            means_t = prev_mean_all.view(S, B, *tail).transpose(0, 1).contiguous().to(dtype=self.trajectory_dtype)
+            means_t = prev_mean_all.view(S, B, *tail).transpose(0, 1).contiguous()
         return ReplayResult(log_probs=log_probs_t, prev_sample_means=means_t)
 
 

unirl/models/types/replay_result.py

@@ -19,7 +19,8 @@ class ReplayResult:
     prev_sample_means: Optional[torch.Tensor] = None
     """The SDE transition's mean μ_θ at each replayed step. Shape
     ``[B, S', *latent_shape]`` for diffusion. Used by GRPO's KL penalty.
-    ``None`` when the stage doesn't produce it."""
+    Kept at the SDE kernel's native fp32; see Gotchas in
+    ``unirl/sde/README.md``. ``None`` when the stage doesn't produce it."""
 
     logits: Optional[torch.Tensor] = None
     """Per-step token logits at each replayed position. Shape

unirl/models/wan21/diffusion.py

@@ -403,7 +403,7 @@ def replay(
                     prev_sample_means.append(prev_mean)
 
         log_probs_t = torch.stack(log_probs, dim=1).to(dtype=self.logprob_dtype)
-        means_t = torch.stack(prev_sample_means, dim=1).to(dtype=self.trajectory_dtype) if prev_sample_means else None
+        means_t = torch.stack(prev_sample_means, dim=1) if prev_sample_means else None
         return ReplayResult(log_probs=log_probs_t, prev_sample_means=means_t)
 
     def predict_noise_at_step(

unirl/models/wan22/diffusion.py

@@ -447,7 +447,7 @@ def replay(
                     prev_sample_means.append(prev_mean)
 
         log_probs_t = torch.stack(log_probs, dim=1).to(dtype=self.logprob_dtype)
-        means_t = torch.stack(prev_sample_means, dim=1).to(dtype=self.trajectory_dtype) if prev_sample_means else None
+        means_t = torch.stack(prev_sample_means, dim=1) if prev_sample_means else None
         return ReplayResult(log_probs=log_probs_t, prev_sample_means=means_t)
 
     def predict_noise_at_step(

unirl/models/z_image/diffusion.py

@@ -397,7 +397,7 @@ def replay(
                     prev_sample_means.append(prev_mean)
 
         log_probs_t = torch.stack(log_probs, dim=1).to(dtype=self.logprob_dtype)
-        means_t = torch.stack(prev_sample_means, dim=1).to(dtype=self.trajectory_dtype) if prev_sample_means else None
+        means_t = torch.stack(prev_sample_means, dim=1) if prev_sample_means else None
         return ReplayResult(log_probs=log_probs_t, prev_sample_means=means_t)
 
     @staticmethod

unirl/sde/README.md

@@ -109,6 +109,13 @@ MixGRPO keeps `FlowSDEStrategy` and adds a `WindowScheduler` under
 - **The dtype round-trip in `_finalize_logp` is not a no-op cast.** It simulates
   trajectory *storage* precision so replay-time log-prob matches sampling-time
   log-prob. Delete it as dead code and the ratio drifts. Skipped for `eta<1e-7`.
+- **`prev_sample_means` leave replay at the kernel's native fp32 — don't narrow
+  them to `trajectory_precision`.** Log-prob scores the fp32 mean; FlowDPPO's
+  Gaussian-KL mask (`(Δmean)²/(2σ²)` vs the 1e-5 threshold) scores whatever
+  replay returns, so a bf16 round-trip re-scores a *different* distribution:
+  ULP-scale rounding dwarfs the threshold, and sub-ULP policy deltas collapse to
+  KL=0. `log_probs` keeps its explicit `logprob_dtype` cast; only the mean must
+  not follow the trajectory dtype.
 - **`ensure_sample_sigmas` takes height/width/steps with no defaults, on purpose** —
   a silent `1024×1024` mis-derives μ for dynamic-shift models rendering at anything
   else (e.g. WAN T2V at 480×832).