From c2ed7b3cb9aa0d7660c21d040292dd667bf36b51 Mon Sep 17 00:00:00 2001
From: dog <yellowdog016@gmail.com>
Date: Thu, 9 Jul 2026 17:25:50 +0800
Subject: [PATCH] feat: add soften_forward (backtrack tolerance + tanh
 overshoot compression)

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
---
 src/ai_mouse/_postprocess.py   | 36 +++++++++++++++++++++++++++++
 tests/unit/test_postprocess.py | 41 ++++++++++++++++++++++++++++++++++
 2 files changed, 77 insertions(+)

diff --git a/src/ai_mouse/_postprocess.py b/src/ai_mouse/_postprocess.py
index 3091703..b2f8aef 100644
--- a/src/ai_mouse/_postprocess.py
+++ b/src/ai_mouse/_postprocess.py
@@ -178,3 +178,39 @@ def sample_duration(
     mu_log = params[bin_idx]["mu_log"]
     sigma_log = params[bin_idx]["sigma_log"]
     return float(np.exp(rng.normal(mu_log, sigma_log)))
+
+
+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
diff --git a/tests/unit/test_postprocess.py b/tests/unit/test_postprocess.py
index 42bcd21..f78e059 100644
--- a/tests/unit/test_postprocess.py
+++ b/tests/unit/test_postprocess.py
@@ -141,3 +141,44 @@ def test_sample_duration_uses_correct_bin() -> None:
     v = sample_duration(dist_dict, 150.0, rng)
     # exp(6) ~ 403, with tiny sigma we should land near there
     assert 350 < v < 460
+
+
+from ai_mouse._postprocess import soften_forward
+
+
+def test_soften_forward_tolerates_small_backtrack() -> None:
+    # A 0.01 dip is within the 0.02 tolerance and must survive untouched.
+    f = np.array([0.0, 0.30, 0.29, 0.60, 1.0])
+    out = soften_forward(f)
+    assert np.isclose(out[2], 0.29)
+
+
+def test_soften_forward_limits_large_backtrack() -> None:
+    # A 0.30 dip is noise; it gets pulled up to prev - tol.
+    f = np.array([0.0, 0.50, 0.20, 0.70, 1.0])
+    out = soften_forward(f)
+    assert np.isclose(out[2], 0.50 - 0.02)
+
+
+def test_soften_forward_allows_moderate_overshoot() -> None:
+    # Overshoot past 1.0 is natural; small overshoot survives (compressed
+    # but strictly > 1.0).
+    f = np.array([0.0, 0.5, 0.9, 1.04, 1.0])
+    out = soften_forward(f)
+    assert out[3] > 1.0
+
+
+def test_soften_forward_compresses_extreme_overshoot() -> None:
+    # tanh compression: no output value may exceed 1 + overshoot_span.
+    f = np.array([0.0, 0.5, 1.30, 1.50, 1.0])
+    out = soften_forward(f)
+    assert out.max() <= 1.0 + 0.08 + 1e-9
+    assert out[2] > 1.0  # still an overshoot, not clipped flat
+
+
+def test_soften_forward_no_lower_clip() -> None:
+    # Small wind-up behind the start is allowed (warp_endpoints pins
+    # the first point later; interior may be slightly negative).
+    f = np.array([0.0, -0.01, 0.30, 0.70, 1.0])
+    out = soften_forward(f)
+    assert out[1] < 0.0