"""Multi-target tracker for LiDAR-detected sheep. Three-stage greedy nearest-neighbour data association: 1. **Consensus promotion**. New detections start as *candidate* tracks invisible to ``get_positions``. They must accumulate ``consensus_k`` matches within ``consensus_radius_m`` to promote; candidates that fail to re-confirm within ``consensus_max_age`` steps die. This filters one-shot LiDAR phantoms — wall returns, multi-cluster sheep splits, fast-moving sheep position jumps — at the cost of a small acquisition latency (~50 ms at the default ``consensus_k=3``). ``consensus_k=1`` disables the stage. 2. **Constant-velocity prediction**. Each track carries a smoothed ``(vx, vy)``. While a track is occluded its position is extrapolated for up to ``PREDICT_STEPS`` frames, then falls back to last-seen static memory until ``FORGET_STEPS`` deletes it. 3. **Pen latching**. A track whose estimated position crosses the gate plane south of ``is_penned`` is marked penned, excluded from ``get_positions``, and kept indefinitely. Output of :meth:`SheepTracker.get_positions` is ``{name: (x, y)}`` — Strömbom, Sequential and the BC observation builder consume it directly. """ from __future__ import annotations import math from typing import TYPE_CHECKING if TYPE_CHECKING: from herding.config import TrackerConfig from herding.world.geometry import MAX_SHEEP, in_pen, is_penned GATE_M = 2.5 # m — primary NN gate (recently observed tracks) REACQUIRE_GATE_M = 4.5 # m — wider gate for re-binding stale tracks REACQUIRE_MIN_AGE = 20 # steps — track must be this stale to use the wider gate PENNED_GATE_M = 4.0 # m — gate for matching detections to existing penned tracks FORGET_STEPS = 200 # ~3.2 s — delete stale active tracks (penned ones kept forever) MAX_ACTIVE_TRACKS = MAX_SHEEP # Predictive tracking constants. PREDICT_STEPS = 120 # ~1.9 s — extrapolate velocity this many frames VELOCITY_CLAMP = 1.0 # m/s — max predicted speed (sheep max is ~0.78 m/s) class Track: """Single track with position, velocity, and age. Attributes ---------- candidate ``True`` while the track has not yet accumulated enough consensus matches to be visible (``hit_count < consensus_k``). Candidates are excluded from :meth:`SheepTracker.get_positions` and from the active/penned counters. hit_count Number of detections this track has absorbed since spawn, used by the consensus filter. """ __slots__ = ("x", "y", "vx", "vy", "last_seen", "penned", "candidate", "hit_count") def __init__( self, x: float, y: float, step: int, penned: bool = False, candidate: bool = False, ): self.x = x self.y = y self.vx = 0.0 self.vy = 0.0 self.last_seen = step self.penned = penned self.candidate = candidate self.hit_count = 1 @property def age(self) -> int: """Not-a-property in the hot loop — callers pass current step.""" raise NotImplementedError def predicted_position( self, current_step: int, predict_steps: int = PREDICT_STEPS, velocity_clamp: float = VELOCITY_CLAMP, ) -> tuple[float, float]: """Extrapolated position using constant velocity, clamped.""" dt = current_step - self.last_seen if dt <= 0 or dt > predict_steps: return self.x, self.y speed = math.hypot(self.vx, self.vy) if speed < 1e-4: return self.x, self.y # Clamp extrapolation distance. max_d = velocity_clamp * dt * 0.016 # steps → seconds d = min(speed * dt * 0.016, max_d) return ( self.x + d * (self.vx / speed), self.y + d * (self.vy / speed), ) def update(self, x: float, y: float, step: int) -> None: """Absorb a new detection and re-estimate velocity.""" dt = step - self.last_seen if dt > 0: dt_s = dt * 0.016 # steps → seconds new_vx = (x - self.x) / dt_s new_vy = (y - self.y) / dt_s # Exponential smoothing on velocity. alpha = 0.6 self.vx = alpha * new_vx + (1.0 - alpha) * self.vx self.vy = alpha * new_vy + (1.0 - alpha) * self.vy self.x = x self.y = y self.last_seen = step class SheepTracker: """Online tracker with NN association, prediction, and forgetful memory. Each track is a :class:`Track` with position, velocity estimate, last-seen step, and penned flag. Pass a :class:`~herding.config.TrackerConfig` to override any module-level defaults without changing this file. """ def __init__( self, gate: float = GATE_M, tracker_cfg: "TrackerConfig | None" = None, ): if tracker_cfg is not None: self.gate = tracker_cfg.gate_m self._reacquire_gate = tracker_cfg.reacquire_gate_m self._reacquire_min_age = tracker_cfg.reacquire_min_age self._penned_gate = tracker_cfg.penned_gate_m self._forget_steps = tracker_cfg.forget_steps self._predict_steps = tracker_cfg.predict_steps self._velocity_clamp = tracker_cfg.velocity_clamp self._max_new_per_step = tracker_cfg.max_new_tracks_per_step self._pen_latch_depth = tracker_cfg.pen_latch_depth self._consensus_k = tracker_cfg.consensus_k self._consensus_radius = tracker_cfg.consensus_radius_m self._consensus_max_age = tracker_cfg.consensus_max_age else: self.gate = gate self._reacquire_gate = REACQUIRE_GATE_M self._reacquire_min_age = REACQUIRE_MIN_AGE self._penned_gate = PENNED_GATE_M self._forget_steps = FORGET_STEPS self._predict_steps = PREDICT_STEPS self._velocity_clamp = VELOCITY_CLAMP self._max_new_per_step = MAX_ACTIVE_TRACKS self._pen_latch_depth = 0.0 self._consensus_k = 1 self._consensus_radius = 0.5 self._consensus_max_age = 8 self._tracks: dict[int, Track] = {} self._next_id = 0 self.step = 0 def reset(self) -> None: self._tracks.clear() self._next_id = 0 self.step = 0 def update(self, detections: list[tuple[float, float]]) -> dict[str, tuple[float, float]]: """Fold a new set of detections in and return active positions.""" self.step += 1 det_used: set[int] = set() updated_tids: set[int] = set() # Pass 1 — match promoted active tracks within the primary gate. # Use predicted positions for matching, oldest-first. Candidates # are excluded; they get their own (tighter) pass below so a # stray detection cannot rescue an already-stale candidate. active_tids = [tid for tid, t in self._tracks.items() if not t.penned and not t.candidate] active_tids.sort(key=lambda tid: self._tracks[tid].last_seen) for tid in active_tids: track = self._tracks[tid] tx, ty = track.predicted_position( self.step, self._predict_steps, self._velocity_clamp) best_j, best_d = -1, self.gate for j, (dx, dy) in enumerate(detections): if j in det_used: continue d = math.hypot(dx - tx, dy - ty) if d < best_d: best_d = d best_j = j if best_j >= 0: dx, dy = detections[best_j] track.update(dx, dy, self.step) track.hit_count += 1 det_used.add(best_j) updated_tids.add(tid) # Pass 1b — re-acquisition with wider gate for stale tracks. for tid in active_tids: if tid in updated_tids: continue track = self._tracks[tid] if (self.step - track.last_seen) < self._reacquire_min_age: continue tx, ty = track.predicted_position( self.step, self._predict_steps, self._velocity_clamp) best_j, best_d = -1, self._reacquire_gate for j, (dx, dy) in enumerate(detections): if j in det_used: continue d = math.hypot(dx - tx, dy - ty) if d < best_d: best_d = d best_j = j if best_j >= 0: dx, dy = detections[best_j] track.update(dx, dy, self.step) track.hit_count += 1 det_used.add(best_j) updated_tids.add(tid) # Pass 1c — match remaining detections to candidate tracks within # the tight consensus radius. Each hit ages the candidate; once # hit_count reaches consensus_k it is promoted (handled below). candidate_tids = [tid for tid, t in self._tracks.items() if t.candidate] candidate_tids.sort(key=lambda tid: self._tracks[tid].last_seen) for tid in candidate_tids: track = self._tracks[tid] best_j, best_d = -1, self._consensus_radius for j, (dx, dy) in enumerate(detections): if j in det_used: continue d = math.hypot(dx - track.x, dy - track.y) if d < best_d: best_d = d best_j = j if best_j >= 0: dx, dy = detections[best_j] track.update(dx, dy, self.step) track.hit_count += 1 det_used.add(best_j) # Pass 2 — match remaining detections to penned tracks. penned_tids = [tid for tid, t in self._tracks.items() if t.penned] for tid in penned_tids: track = self._tracks[tid] best_j, best_d = -1, self._penned_gate for j, (dx, dy) in enumerate(detections): if j in det_used: continue d = math.hypot(dx - track.x, dy - track.y) if d < best_d: best_d = d best_j = j if best_j >= 0: dx, dy = detections[best_j] track.update(dx, dy, self.step) track.hit_count += 1 det_used.add(best_j) # Spawn tracks for still-unmatched detections. # # When ``consensus_k > 1`` every new track starts as a candidate # and remains invisible to ``get_positions`` until it accumulates # the required matches. Penned latching is deferred to after # promotion — otherwise gate-area phantoms could still skip the # consensus filter by landing inside the pen column and being # latched forever, which is exactly the failure mode the filter # is meant to eliminate. ``max_new_tracks_per_step`` continues # to rate-cap spawns. spawned = 0 spawn_candidates = self._consensus_k > 1 for j, (dx, dy) in enumerate(detections): if j in det_used: continue if spawned >= self._max_new_per_step: break if spawn_candidates: self._tracks[self._next_id] = Track( dx, dy, self.step, penned=False, candidate=True) else: penned = self._is_penned(dx, dy) self._tracks[self._next_id] = Track( dx, dy, self.step, penned=penned, candidate=False) self._next_id += 1 spawned += 1 # Promote candidates that have accumulated enough matches. for track in self._tracks.values(): if track.candidate and track.hit_count >= self._consensus_k: track.candidate = False # Promote active tracks whose current estimate crosses the gate. # Candidates are deliberately excluded — a track that hasn't yet # earned visibility shouldn't be allowed to latch as penned # either (that path is exactly how south-wall FPs persisted # forever before the consensus filter existed). for track in self._tracks.values(): if track.penned or track.candidate: continue px, py = track.predicted_position( self.step, self._predict_steps, self._velocity_clamp) if self._is_penned(px, py): track.penned = True # Forget stale tracks. Candidates have their own short timeout # (one window to confirm or die); promoted active tracks decay at # forget_steps; penned tracks decay 8× slower because real penned # sheep are still observed when the dog faces the pen. penned_forget = self._forget_steps * 8 stale: list[int] = [] for tid, t in self._tracks.items(): age = self.step - t.last_seen if t.candidate: if age > self._consensus_max_age: stale.append(tid) elif t.penned: if age > penned_forget: stale.append(tid) else: if age > self._forget_steps: stale.append(tid) for tid in stale: del self._tracks[tid] # Hard cap on the visible (promoted, not penned) active set — # drop the oldest-seen overflow. Candidates are not counted here: # they don't compete for slots until they earn promotion, and # rate-limiting their spawn is the job of ``max_new_per_step``. active = [(tid, t.last_seen) for tid, t in self._tracks.items() if not t.penned and not t.candidate] if len(active) > MAX_ACTIVE_TRACKS: active.sort(key=lambda kv: kv[1]) for tid, _ in active[: len(active) - MAX_ACTIVE_TRACKS]: del self._tracks[tid] return self.get_positions() def _is_penned(self, x: float, y: float) -> bool: """Check whether a position should be considered penned. Uses ``pen_latch_depth`` to require the position to be that many metres past the gate line before latching. Increasing the depth prevents gate-area LiDAR false positives (gate hardware reflections at y ≈ -15) from being permanently latched as penned tracks. """ from herding.world.geometry import GATE_Y # Apply depth threshold to both in_pen and is_penned so # that any position in the gate column must clear GATE_Y - depth. threshold = GATE_Y - self._pen_latch_depth return (in_pen(x, y) or is_penned(x, y)) and y <= threshold def get_positions(self, min_freshness: int | None = None) -> dict[str, tuple[float, float]]: """Promoted (non-candidate, non-penned) tracks as ``{name: (x, y)}``. For tracks currently being predicted (occluded but within predict_steps), returns the extrapolated position so the teacher sees a smooth estimate. Candidate tracks — those that have not yet accumulated ``consensus_k`` matches — are excluded so a one-shot phantom detection never reaches the policy/teacher. ``min_freshness`` (optional, deploy-only): drop tracks whose last_seen is older than ``step - min_freshness``. Real sheep in FOV are detected nearly every step; phantom tracks from sporadic Webots FPs stop being re-observed and decay. Default ``None`` preserves training behaviour (extrapolated tracks visible). """ result = {} for tid, track in self._tracks.items(): if track.penned or track.candidate: continue if (min_freshness is not None and self.step - track.last_seen > min_freshness): continue px, py = track.predicted_position( self.step, self._predict_steps, self._velocity_clamp) result[f"t{tid}"] = (px, py) return result def get_penned_set(self) -> set[str]: return {f"t{tid}" for tid, t in self._tracks.items() if t.penned} def n_active(self) -> int: """Number of promoted (non-candidate, non-penned) tracks.""" return sum(1 for t in self._tracks.values() if not t.penned and not t.candidate) def n_penned(self) -> int: return sum(1 for t in self._tracks.values() if t.penned) def n_candidate(self) -> int: """Number of unpromoted candidate tracks awaiting consensus.""" return sum(1 for t in self._tracks.values() if t.candidate) def n_predicted(self) -> int: """Number of promoted active tracks currently being extrapolated (not directly observed).""" return sum(1 for t in self._tracks.values() if not t.penned and not t.candidate and (self.step - t.last_seen) > 0 and (self.step - t.last_seen) <= self._predict_steps)