feat: switch flow ODE sampling from Euler to Heun (10 steps, NFE 20)

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

src/ai_mouse/mouse.py

@@ -24,7 +24,7 @@ from ai_mouse._postprocess import (
 )
 from ai_mouse.errors import GenerationError, ModelLoadError
 
-_N_EULER_STEPS = 10
+_N_ODE_STEPS = 10
 
 
 class MouseModel:
@@ -90,11 +90,16 @@ class MouseModel:
         )[None]
 
         x = rng.standard_normal((1, self._seq_len, 3)).astype(np.float32)
-        dt = 1.0 / _N_EULER_STEPS
-        for step in range(_N_EULER_STEPS):
-            t = np.full((1,), step * dt, dtype=np.float32)
-            v = self._session.run(["v"], {"x_t": x, "t": t, "cond": cond})[0]
-            x = x + v * dt
+        # Heun (2nd-order predictor-corrector): same step count as the old
+        # Euler loop but far lower integration error for 2x NFE.
+        dt = 1.0 / _N_ODE_STEPS
+        for step in range(_N_ODE_STEPS):
+            t0 = np.full((1,), step * dt, dtype=np.float32)
+            v1 = self._session.run(["v"], {"x_t": x, "t": t0, "cond": cond})[0]
+            x_pred = (x + v1 * dt).astype(np.float32)
+            t1 = np.full((1,), (step + 1) * dt, dtype=np.float32)
+            v2 = self._session.run(["v"], {"x_t": x_pred, "t": t1, "cond": cond})[0]
+            x = x + (v1 + v2) * (dt / 2.0)
 
         if not np.all(np.isfinite(x)):
             raise GenerationError("Trajectory contains NaN/Inf after Euler integration")