From 441e6f3dfe4237a528e4a41f514ea5a1293bb704 Mon Sep 17 00:00:00 2001 From: dog Date: Thu, 9 Jul 2026 17:48:22 +0800 Subject: [PATCH] feat: rework mouse post-processing pipeline (soft monotonic, global endpoint warp) Golden mouse baselines temporarily failing; re-captured in follow-up commit. Co-Authored-By: Claude Fable 5 --- src/ai_mouse/_postprocess.py | 61 ---------------------------------- src/ai_mouse/mouse.py | 13 ++++---- tests/unit/test_postprocess.py | 49 --------------------------- 3 files changed, 7 insertions(+), 116 deletions(-) diff --git a/src/ai_mouse/_postprocess.py b/src/ai_mouse/_postprocess.py index 9072159..d175437 100644 --- a/src/ai_mouse/_postprocess.py +++ b/src/ai_mouse/_postprocess.py @@ -31,67 +31,6 @@ def gaussian_smooth(x: np.ndarray, sigma: float = 1.0) -> np.ndarray: return smoothed -def snap_endpoints( - forward: np.ndarray, - lateral: np.ndarray, - seq_len: int, - n_snap: int = 6, -) -> tuple[np.ndarray, np.ndarray]: - """Force first point to (0,0) and last point to (1,0) with quadratic ease. - - The last ``n_snap`` points are linearly interpolated towards (1, 0) - with quadratic easing, then the first/last points are pinned exactly. - - Args: - forward: (T,) forward coordinates (modified in place). - lateral: (T,) lateral coordinates (modified in place). - seq_len: length of forward/lateral. - n_snap: number of trailing points to ease (capped at seq_len//4). - - Returns: - ``(forward, lateral)`` after modification. - """ - n_snap = min(n_snap, seq_len // 4) - for i in range(n_snap): - alpha = ((i + 1) / n_snap) ** 2 - k = seq_len - n_snap + i - forward[k] = forward[k] * (1.0 - alpha) + 1.0 * alpha - lateral[k] = lateral[k] * (1.0 - alpha) + 0.0 * alpha - forward[0], lateral[0] = 0.0, 0.0 - forward[-1], lateral[-1] = 1.0, 0.0 - return forward, lateral - - -def smooth_start( - forward: np.ndarray, - lateral: np.ndarray, - n: int = 4, -) -> tuple[np.ndarray, np.ndarray]: - """Dampen lateral oscillation in the first ``n`` points. - - Assumes :func:`snap_endpoints` has already pinned (0,0). Forward is - forced non-decreasing locally; lateral is linearly damped towards 0. - """ - n_start_fix = min(n, len(forward) // 4) - for i in range(1, n_start_fix + 1): - blend = i / (n_start_fix + 1) - forward[i] = max(forward[i], forward[i - 1]) - lateral[i] = lateral[i] * blend - return forward, lateral - - -def enforce_forward_monotonic(forward: np.ndarray) -> np.ndarray: - """Force ``forward`` non-decreasing, clip to [0,1], pin endpoints.""" - seq_len = len(forward) - for i in range(1, seq_len - 1): - if forward[i] < forward[i - 1]: - forward[i] = forward[i - 1] + 0.001 - forward = np.clip(forward, 0.0, 1.0) - forward[0] = 0.0 - forward[-1] = 1.0 - return forward - - def damp_start(lateral: np.ndarray, n: int = 4) -> np.ndarray: """Dampen lateral oscillation over the first ``n`` points, continuously. diff --git a/src/ai_mouse/mouse.py b/src/ai_mouse/mouse.py index 80629f8..f312e5d 100644 --- a/src/ai_mouse/mouse.py +++ b/src/ai_mouse/mouse.py @@ -14,13 +14,13 @@ from ai_mouse._assets import resolve from ai_mouse._coord import decode_trajectory from ai_mouse._postprocess import ( build_timestamps, - enforce_forward_monotonic, + damp_start, gaussian_smooth, resample_arc, sample_duration, - smooth_start, - snap_endpoints, + soften_forward, truncnorm_sample, + warp_endpoints, ) from ai_mouse.errors import GenerationError, ModelLoadError @@ -103,10 +103,11 @@ class MouseModel: lateral = x[0, :, 1].copy() log_dt = x[0, :, 2].copy() - forward, lateral = snap_endpoints(forward, lateral, self._seq_len) - forward, lateral = smooth_start(forward, lateral) - forward = enforce_forward_monotonic(forward) + forward = soften_forward(forward) + lateral = damp_start(lateral) + forward = gaussian_smooth(forward, sigma=1.0) lateral = gaussian_smooth(lateral, sigma=1.0) + forward, lateral = warp_endpoints(forward, lateral) log_dt = np.clip(log_dt, 0.0, 5.0) log_dt[0] = 0.0 diff --git a/tests/unit/test_postprocess.py b/tests/unit/test_postprocess.py index ffd68ba..c053a04 100644 --- a/tests/unit/test_postprocess.py +++ b/tests/unit/test_postprocess.py @@ -25,55 +25,6 @@ def test_gaussian_smooth_constant_unchanged() -> None: assert np.allclose(result, x, atol=1e-6) -from ai_mouse._postprocess import snap_endpoints - - -def test_snap_endpoints_pins_first_and_last() -> None: - forward = np.linspace(0.1, 0.9, 16) - lateral = np.full(16, 0.5) - f, l = snap_endpoints(forward.copy(), lateral.copy(), seq_len=16) - assert f[0] == 0.0 - assert l[0] == 0.0 - assert f[-1] == 1.0 - assert l[-1] == 0.0 - - -def test_snap_endpoints_preserves_middle() -> None: - forward = np.linspace(0.0, 1.0, 16) - lateral = np.zeros(16) - f, _ = snap_endpoints(forward.copy(), lateral.copy(), seq_len=16, n_snap=4) - # Points before the last n_snap should be unchanged - assert np.allclose(f[1 : 16 - 4], forward[1 : 16 - 4], atol=1e-6) - - -from ai_mouse._postprocess import enforce_forward_monotonic, smooth_start - - -def test_smooth_start_dampens_lateral() -> None: - forward = np.linspace(0, 1, 16) - lateral = np.full(16, 1.0) - forward[0] = lateral[0] = 0.0 # invariant: snap already done - _, l = smooth_start(forward.copy(), lateral.copy(), n=4) - # Lateral at points 1-4 should be < original (dampened) - assert l[1] < 1.0 - assert l[4] < 1.0 - # Lateral at point 5+ unchanged - assert l[5] == 1.0 - - -def test_enforce_forward_monotonic_repairs_inversions() -> None: - f = np.array([0.0, 0.4, 0.3, 0.6, 0.5, 1.0]) - out = enforce_forward_monotonic(f.copy()) - assert np.all(np.diff(out) > 0), out - - -def test_enforce_forward_monotonic_clips_to_unit_interval() -> None: - f = np.array([-0.1, 0.5, 1.2]) - out = enforce_forward_monotonic(f.copy()) - assert out[0] == 0.0 - assert out[-1] == 1.0 - - from ai_mouse._postprocess import build_timestamps, resample_arc