fix: use table lookup for boundary index instead of inverse sigma_shift formula

Replace the analytical inverse sigma_shift formula with a direct table
lookup on scheduler.timesteps. This avoids floating-point precision
issues (e.g. 0.24999 vs 0.25) and the sigma_min != 0 approximation
error, giving an exact boundary at the intended timestep threshold.

Co-authored-by: Cursor <cursoragent@cursor.com>

Author: Auraithm

mova/diffusion/pipelines/mova_train.py

@@ -1386,14 +1386,9 @@ def training_step(
             mode_scale=1.0,
             independent_timesteps=False,
         )
-        # The sigma_shift formula in flow_match.py L57:
-        #   sigma_shifted = shift * sigma_linear / (1 + (shift - 1) * sigma_linear)
-        # boundary_ratio is in shifted space; invert to get sigma_linear:
-        #   sigma_linear = boundary_ratio / (shift - (shift - 1) * boundary_ratio)
-        # index_boundary = 1 - sigma_linear, which simplifies to:
-        #   shift * (1 - boundary_ratio) / (1 + (shift - 1) * (1 - boundary_ratio))
-        # e.g. shift=3, boundary_ratio=0.9 => sigma_linear=0.75 => index=0.25
-        boundary = self.scheduler.shift * (1 - self.boundary_ratio) / (1 + (self.scheduler.shift - 1) * (1 - self.boundary_ratio))
+        # Look up the index boundary directly from the scheduler's timestep table.
+        # boundary_ratio=0.9 means timestep 900; count how many timesteps >= 900 to get the index fraction.
+        boundary = (self.scheduler.timesteps >= self.boundary_ratio * self.scheduler.num_train_timesteps).sum().item() / self.scheduler.num_train_timesteps
         if global_step % 2 == 0:
             timestep_config.max_timestep_boundary = boundary
         else: