diff --git a/src/ai_mouse/_postprocess.py b/src/ai_mouse/_postprocess.py index b2f8aef..a737d63 100644 --- a/src/ai_mouse/_postprocess.py +++ b/src/ai_mouse/_postprocess.py @@ -92,6 +92,65 @@ def enforce_forward_monotonic(forward: np.ndarray) -> np.ndarray: return forward +def soften_forward( + forward: np.ndarray, + backtrack_tol: float = 0.02, + overshoot_span: float = 0.08, +) -> np.ndarray: + """Softly regularise the forward axis without destroying natural motion. + + Real trajectories contain small backward corrections and overshoot + past the target; hard clipping turns both into visible artifacts + (stacked points, vertical walls). Instead: + + - Backtracking is tolerated up to ``backtrack_tol``; larger dips are + raised to ``prev - backtrack_tol``. + - Values above 1.0 are compressed with tanh so they asymptote at + ``1 + overshoot_span`` (moderate overshoot survives, extremes + cannot fly far past the target). + - No lower clip: the endpoint warp pins the first point later. + + Args: + forward: (T,) forward coordinates. + backtrack_tol: max allowed per-step regression. + overshoot_span: asymptotic max excess above 1.0. + + Returns: + New (T,) array. + """ + out = forward.copy() + for i in range(1, len(out)): + floor = out[i - 1] - backtrack_tol + if out[i] < floor: + out[i] = floor + over = out > 1.0 + out[over] = 1.0 + overshoot_span * np.tanh((out[over] - 1.0) / overshoot_span) + return out + + +def damp_start(lateral: np.ndarray, n: int = 4) -> np.ndarray: + """Dampen lateral oscillation over the first ``n`` points, continuously. + + Weights follow smoothstep(i / (n+1)) so the damping releases smoothly + into the untouched region (the old linear blend jumped from 0.8 to + 1.0 and left a visible kink). + + Args: + lateral: (T,) lateral coordinates. + n: number of leading points to dampen (capped at len//4). + + Returns: + New (T,) array. + """ + out = lateral.copy() + n = min(n, len(out) // 4) + for i in range(1, n + 1): + t = i / (n + 1) + w = t * t * (3.0 - 2.0 * t) # smoothstep + out[i] *= w + return out + + def resample_arc(xy: np.ndarray, n_points: int) -> np.ndarray: """Resample a 2-D polyline to ``n_points`` along cumulative arc length.""" arc = np.concatenate( diff --git a/tests/unit/test_postprocess.py b/tests/unit/test_postprocess.py index f78e059..e2bcd2e 100644 --- a/tests/unit/test_postprocess.py +++ b/tests/unit/test_postprocess.py @@ -182,3 +182,27 @@ def test_soften_forward_no_lower_clip() -> None: f = np.array([0.0, -0.01, 0.30, 0.70, 1.0]) out = soften_forward(f) assert out[1] < 0.0 + + +from ai_mouse._postprocess import damp_start + + +def test_damp_start_dampens_early_lateral() -> None: + lat = np.full(16, 1.0) + out = damp_start(lat, n=4) + assert out[1] < out[2] < out[3] < out[4] < 1.0 # monotone ramp + assert np.all(out[5:] == 1.0) # untouched past n + + +def test_damp_start_no_release_jump() -> None: + # The weight at i=n must be close to 1 (continuous release): + # smoothstep(4/5) = 0.896, vs the old linear 4/5 = 0.8. + lat = np.full(16, 1.0) + out = damp_start(lat, n=4) + assert out[4] > 0.85 + + +def test_damp_start_short_input_safe() -> None: + lat = np.array([0.0, 0.5, 0.3]) + out = damp_start(lat, n=4) # n capped to len//4 = 0 → no-op + assert np.array_equal(out, lat)