aurasync-engine 0.4.1__py3-none-any.whl

aurasync/frame_camera.py

@@ -0,0 +1,163 @@
+"""Learned frame-camera baseline (fair comparison for the event stack).
+
+The phase-1 frame baselines were handed the TRUE state at capture time. Here
+the frame stack has to *see* the state the same way the event stack does:
+
+  sensor     global-shutter grayscale camera on the same 96x32 scene,
+             exposure-integrated motion blur and read noise
+  estimator  FrameEstimator: ResNet-lite trunk over the last 3 frames plus a
+             shared per-frame spatial-softmax keypoint encoder (the same
+             sub-pixel tool the event observer uses), conditioned on frame
+             period and exposure. Regresses the full state at capture time.
+  pipeline   LearnedFramePipeline: at each frame arrival (capture + latency)
+             the estimate is rolled forward to "now" through the same nominal
+             physics model the event observer uses, replaying the controls
+             applied since capture (latency compensation, a Smith predictor).
+             Between frames the estimate keeps propagating at 1 kHz.
+             mode="zoh" disables both (naive hold, for the ablation).
+
+Analogy: the frame stack is a goalkeeper watching a strobe light. Between
+flashes he still knows physics, so he extrapolates the ball's path rather than
+freezing on the last flash. The question is whether that beats someone who
+sees a continuous trickle of motion (the event camera).
+"""
+from collections import deque
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from .physics import PlantParams, rk4_step
+from .event_camera import SceneRenderer
+from .estimator import SpatialSoftmaxKeypoints, SCALE
+from .conditions import NOMINAL
+
+N_FRAMES = 3
+READ_NOISE = 0.02  # std of per-pixel read noise (scene intensities are 0.15..0.9)
+
+
+def blur_frame(history, exposure_ticks):
+    """Average the last `exposure_ticks` sharp 1 ms renders = shutter open over that window."""
+    return np.mean(np.stack(list(history)[-exposure_ticks:]), 0)
+
+
+class BasicBlock(nn.Module):
+    def __init__(self, cin, cout, stride=1):
+        super().__init__()
+        self.c1 = nn.Conv2d(cin, cout, 3, stride, 1, bias=False)
+        self.b1 = nn.BatchNorm2d(cout)
+        self.c2 = nn.Conv2d(cout, cout, 3, 1, 1, bias=False)
+        self.b2 = nn.BatchNorm2d(cout)
+        self.skip = None if (stride == 1 and cin == cout) else nn.Sequential(
+            nn.Conv2d(cin, cout, 1, stride, bias=False), nn.BatchNorm2d(cout))
+
+    def forward(self, x):
+        y = F.relu(self.b1(self.c1(x)))
+        y = self.b2(self.c2(y))
+        return F.relu(y + (x if self.skip is None else self.skip(x)))
+
+
+class FrameEstimator(nn.Module):
+    """frames (N, 3, H, W) newest-first, cond (N, 2) = [frame period/33ms, exposure/8ms] -> state (N, 4)."""
+
+    def __init__(self, H=32, W=96, K=12):
+        super().__init__()
+        yy, xx = torch.meshgrid(torch.linspace(-1, 1, H), torch.linspace(-1, 1, W), indexing="ij")
+        self.register_buffer("coords", torch.stack([xx, yy])[None])
+        self.trunk = nn.Sequential(
+            nn.Conv2d(N_FRAMES + 2, 16, 3, 1, 1, bias=False), nn.BatchNorm2d(16), nn.ReLU(),
+            BasicBlock(16, 16), BasicBlock(16, 32, 2), BasicBlock(32, 32),
+            BasicBlock(32, 48, 2), BasicBlock(48, 48, 2))
+        self.kp = SpatialSoftmaxKeypoints(1, H, W, K=K)
+        self.fc = nn.Linear(48 * (H // 8) * (W // 8), 128)
+        self.head = nn.Sequential(nn.Linear(128 + N_FRAMES * self.kp.out_dim + 2, 128), nn.GELU(),
+                                  nn.Linear(128, 64), nn.GELU(), nn.Linear(64, 4))
+        self.register_buffer("scale", SCALE.clone())
+
+    def forward(self, frames, cond):
+        n = frames.shape[0]
+        x = (frames - 0.5) / 0.3
+        g = F.gelu(self.fc(self.trunk(torch.cat([x, self.coords.expand(n, -1, -1, -1)], 1)).flatten(1)))
+        kps = self.kp(x.reshape(n * N_FRAMES, 1, *x.shape[2:])).view(n, -1)
+        return self.head(torch.cat([g, kps, cond], 1)) * self.scale
+
+
+def cond_vector(fps, exposure):
+    return torch.tensor([[(1.0 / fps) / 0.0333, exposure / 0.008]], dtype=torch.float32)
+
+
+class LearnedFramePipeline:
+    """Frame camera + FrameEstimator, with physics latency compensation (mode='physics') or naive hold ('zoh')."""
+    needs_true_plant = True   # only used to SIMULATE the camera image of the real pole (V4 fix: V2/V3 imaged the
+                              # nominal 0.25 m pole even when the true pole was longer, e.g. in 'ood_plant')
+
+    def __init__(self, model, fps, latency, exposure=0.004, mode="physics", l=0.25,
+                 plant=None, seed=0, read_noise=READ_NOISE):
+        self.model = model.eval()
+        self.fps, self.latency, self.exposure, self.mode = fps, latency, exposure, mode
+        self.exp_ticks = max(1, int(round(exposure / 1e-3)))
+        self.l = l
+        self.p = plant or PlantParams()  # nominal physics model, same as the event observer's
+        self.seed, self.read_noise = seed, read_noise
+        self.renderer = SceneRenderer()
+        self.scond = NOMINAL
+        tag = "" if mode == "physics" else "-zoh"
+        self.name = f"learned-frame-{fps}fps{tag}" + ("" if abs(exposure - 0.004) < 1e-9 else f"-exp{exposure*1e3:g}ms")
+
+    def set_conditions(self, cond):
+        """Same scene and sensor physics as the event camera (see conditions.py)."""
+        self.scond = cond
+        self.renderer = SceneRenderer(cond=cond)
+        base = self.base_exposure if hasattr(self, "base_exposure") else self.exposure
+        self.base_exposure = base
+        self.exposure = base if cond.is_nominal else cond.frame_exposure(base)
+        self.exp_ticks = max(1, int(round(self.exposure / 1e-3)))
+        self.noise = self.read_noise if cond.is_nominal else cond.frame_noise(self.exposure)
+
+    def reset(self, s0):
+        self.rng = np.random.default_rng(self.seed)
+        if not hasattr(self, "noise"):
+            self.noise = self.read_noise
+        self.t_now = 0.0
+        sharp = self.renderer.render(s0, self.p_true.l if getattr(self, "p_true", None) is not None else self.l, 0.0)
+        self.hist = deque([sharp] * 16, maxlen=16)      # sharp 1 ms renders (for blur)
+        self.frames = deque(maxlen=N_FRAMES)            # captured (blurred, noisy) frames
+        self.pending = deque()                          # (arrival_t, capture_t, frame stack)
+        self.u_hist = deque(maxlen=200)                 # controls applied per tick
+        self.s = np.array(s0, dtype=np.float64)
+        self.next_capture = 0.0
+        self.cond = cond_vector(self.fps, self.exposure)  # model conditioning (fps, exposure)
+
+    def _capture(self):
+        img = blur_frame(self.hist, self.exp_ticks)
+        if not self.scond.is_nominal:
+            img = self.scond.frame_sensor(img)
+        img = img + self.rng.normal(0, self.noise, img.shape)
+        self.frames.appendleft(img.astype(np.float32))
+        while len(self.frames) < N_FRAMES:      # very first capture: repeat it
+            self.frames.append(self.frames[-1])
+        return np.stack(self.frames)
+
+    @torch.inference_mode()
+    def update(self, k, t, s_true, surface, u):
+        self.u_hist.append(u)
+        if self.mode == "physics":            # propagate current estimate across this tick
+            self.s = rk4_step(self.s, u, 1e-3, self.p)
+        # sensor simulation: the camera images the REAL pole (true length), not the nominal one
+        l_true = self.p_true.l if getattr(self, "p_true", None) is not None else self.l
+        self.hist.append(self.renderer.render(s_true, l_true, t))
+        if t + 1e-9 >= self.next_capture:
+            self.pending.append((t + self.latency, t, self._capture()))
+            self.next_capture += 1.0 / self.fps
+        while self.pending and self.pending[0][0] <= t + 1e-9:
+            _, t_cap, stack = self.pending.popleft()
+            s_cap = self.model(torch.from_numpy(stack)[None], self.cond)[0].double().numpy()
+            if self.mode == "physics":         # roll forward from capture time to now
+                n = int(round((t - t_cap) / 1e-3))
+                us = list(self.u_hist)[-n:] if n > 0 else []
+                for uu in us:
+                    s_cap = rk4_step(s_cap, uu, 1e-3, self.p)
+            self.s = s_cap
+
+    def estimate(self):
+        return self.s

aurasync/imu.py

@@ -0,0 +1,98 @@
+"""Proprioceptive baseline (Upgrade D): IMU on the pole + cart wheel encoder + EKF.
+
+This is what a controls engineer would build first for a cart-pole, so any
+camera-based system has to be judged against it. No camera at all, so lighting,
+HDR, occlusion and event-sensor faults do not touch it; model mismatch (unseen
+plant parameters, mechanical snap) does, through the EKF's physics model.
+
+Sensors at 1 kHz, mounted at the pivot:
+  gyro         theta_dot + bias + noise      (sigma 0.005 rad/s, bias ~ N(0, 0.005 rad/s))
+  accelerometer specific force in the pole frame (sigma 0.05 m/s^2 per axis):
+                 f_along = xdd*sin(th) + g*cos(th),  f_perp = xdd*cos(th) - g*sin(th)
+  cart encoder  x quantised to 0.1 mm
+EKF state: [x, x_dot, theta, theta_dot, gyro_bias, cart_accel]; prediction through
+the same nominal RK4 physics as every other pipeline. Cart acceleration is a
+free random-walk state rather than a model prediction, so the two accelerometer
+axes pin down tilt AND acceleration without trusting the plant model
+(robust to unseen masses); Jacobians by finite differences.
+The filter also carries a covariance, so it doubles as a filter-based
+uncertainty baseline.
+"""
+from collections import deque
+import numpy as np
+
+from .physics import PlantParams, rk4_step, deriv
+
+DT = 1e-3
+G = 9.81
+
+
+class IMUPipeline:
+    name = "imu-ekf"
+    needs_true_plant = True
+
+    def __init__(self, delay_ticks=1, plant=None, gyro_sigma=0.005, acc_sigma=0.05,
+                 enc_res=1e-4, bias_sigma=0.005):
+        self.p = plant or PlantParams()
+        self.delay_ticks = delay_ticks
+        self.gs, self.as_, self.res, self.bs = gyro_sigma, acc_sigma, enc_res, bias_sigma
+        self.Q = np.diag([1e-9, 1e-5, 1e-9, 1e-4, 1e-12, 25.0])
+        self.R = np.diag([enc_res ** 2 / 12 + 1e-12, gyro_sigma ** 2, acc_sigma ** 2, acc_sigma ** 2])
+
+    def reset(self, s0):
+        seed = int(abs(np.asarray(s0)).sum() * 1e7) % (2 ** 32)
+        self.rng = np.random.default_rng(seed)
+        self.bias = self.rng.normal(0, self.bs)
+        self.x = np.concatenate([np.asarray(s0, float), [0.0, 0.0]])
+        self.P = np.diag([1e-6, 1e-4, 1e-6, 1e-4, self.bs ** 2, 10.0])
+        self.buf = deque([np.asarray(s0, float)] * (self.delay_ticks + 1), maxlen=self.delay_ticks + 1)
+        self.cov_buf = deque([self.P[:4, :4].copy()] * (self.delay_ticks + 1), maxlen=self.delay_ticks + 1)
+
+    def _f(self, x, u):
+        return np.concatenate([rk4_step(x[:4], u, DT, self.p), x[4:]])
+
+    def _h(self, x, u, p=None):
+        """Sensor model. With p given (simulation of the TRUE sensor) the cart
+        acceleration comes from the true plant; inside the EKF it is the state x[5]."""
+        xdd = deriv(x[:4], u, p)[1] if p is not None else x[5]
+        th = x[2]
+        return np.array([x[0], x[3] + x[4],
+                         xdd * np.sin(th) + G * np.cos(th),
+                         xdd * np.cos(th) - G * np.sin(th)])
+
+    @staticmethod
+    def _jac(f, x, eps=1e-6):
+        f0 = f(x)
+        J = np.zeros((len(f0), len(x)))
+        for i in range(len(x)):
+            d = np.zeros(len(x)); d[i] = eps
+            J[:, i] = (f(x + d) - f(x - d)) / (2 * eps)
+        return J
+
+    def measure(self, s_true, u, p_true):
+        """Simulated sensor readings from the TRUE state and TRUE plant."""
+        z = self._h(np.concatenate([s_true, [self.bias, 0.0]]), u, p_true)
+        z[0] = np.round(s_true[0] / self.res) * self.res
+        z[1] += self.rng.normal(0, self.gs)
+        z[2:] += self.rng.normal(0, self.as_, 2)
+        return z
+
+    def update(self, k, t, s_true, surface, u):
+        z = self.measure(s_true, u, getattr(self, "p_true", None) or self.p)
+        F = self._jac(lambda x: self._f(x, u), self.x)
+        self.x = self._f(self.x, u)
+        self.P = F @ self.P @ F.T + self.Q
+        H = self._jac(lambda x: self._h(x, u), self.x)
+        y = z - self._h(self.x, u)
+        S = H @ self.P @ H.T + self.R
+        K = np.linalg.solve(S, H @ self.P).T
+        self.x = self.x + K @ y
+        I_KH = np.eye(6) - K @ H
+        self.P = I_KH @ self.P @ I_KH.T + K @ self.R @ K.T     # Joseph form
+        self.buf.append(self.x[:4].copy()); self.cov_buf.append(self.P[:4, :4].copy())
+
+    def estimate(self):
+        return self.buf[0]
+
+    def covariance(self):
+        return self.cov_buf[0]

aurasync/monitor.py

@@ -0,0 +1,80 @@
+"""Event-consistency monitor: an independent check on the observer's estimate.
+
+Analogy: a lifeguard told "the swimmer is by the red buoy" glances at where
+the splashes actually are. If the splashes keep coming from somewhere else,
+the report is wrong, however confident the person who gave it sounds.
+
+Why it is needed: the learned variance head stays confident while the observer
+diverges onto a wrong (even mirror-image) pole angle, because the divergence
+happens on trajectories it never saw in training. This monitor makes no use of
+the network at all; it checks the estimate against raw events:
+  * a moving edge fires events where the pole actually is
+  * inside a fan around the estimated pivot (where the pole could be), compare
+    the density of fresh events (< 2 ms old) close to the PREDICTED pole
+    (<= 2 px) with the density elsewhere in the fan (> 3 px away)
+  * densities, not counts, so uniform background noise cancels out
+      r = d_far / (d_near + d_far)    ~0 consistent, ~0.5 uninformative, -> 1 contradicted
+  * ticks with fewer than `min_events` fresh events in the fan carry no
+    information (quiet scene) and leave the running score unchanged
+  * m = exponential moving average of r over informative ticks
+The monitor turns into an uncertainty by inflating the observer's variance:
+  sigma_theta^2 <- sigma_theta^2 + (beta * max(0, m - m0))^2
+with m0 and beta fitted on validation episodes only.
+"""
+import numpy as np
+
+from .event_camera import SceneRenderer
+
+FRESH2 = float(np.exp(-2e-3 / 0.008))     # fast time-surface value of an event < 2 ms old
+
+
+class EventConsistencyMonitor:
+    def __init__(self, l=0.25, fan_deg=60.0, near_px=2.0, far_px=3.0, min_events=6, ema=0.1,
+                 m0=0.25, beta=0.5):
+        self.r = SceneRenderer()
+        self.l, self.fan = l, np.radians(fan_deg)
+        self.near_px, self.far_px, self.min_events, self.ema = near_px, far_px, min_events, ema
+        self.m0, self.beta = m0, beta
+        self.reset()
+
+    def reset(self):
+        self.m = self.m0
+        self.info_ticks = 0
+
+    def raw_score(self, s, surface):
+        """r for one tick, or None if the tick carries no information."""
+        S, wx, wy = self.r.S, self.r.wx, self.r.wy
+        x, th = float(s[0]), float(s[2])
+        px, py = x, 0.04
+        dx, dy = wx - px, wy - py
+        dist_piv = np.hypot(dx, dy)
+        ang = np.arctan2(dx, dy)                                   # 0 = straight up, + = right
+        fan = (dist_piv <= 2 * self.l + 3 / S) & (dist_piv >= 0.06) & (np.abs(ang) <= self.fan)
+        tx, ty = px + 2 * self.l * np.sin(th), py + 2 * self.l * np.cos(th)
+        vx, vy = tx - px, ty - py
+        t = np.clip((dx * vx + dy * vy) / (vx * vx + vy * vy), 0, 1)
+        d_pole_px = np.hypot(wx - (px + t * vx), wy - (py + t * vy)) * S
+        near = fan & (d_pole_px <= self.near_px)
+        far = fan & (d_pole_px > self.far_px)
+        ev = (surface[0] >= FRESH2) | (surface[1] >= FRESH2)
+        n_near, n_far = int((ev & near).sum()), int((ev & far).sum())
+        if n_near + n_far < self.min_events or near.sum() == 0 or far.sum() == 0:
+            return None
+        d_near, d_far = n_near / near.sum(), n_far / far.sum()
+        return d_far / (d_near + d_far + 1e-9)
+
+    def update(self, s, surface):
+        r = self.raw_score(s, surface)
+        if r is not None:
+            self.m += self.ema * (r - self.m)
+            self.info_ticks += 1
+        return self.m
+
+    def inflate(self, sigma_theta):
+        return np.sqrt(sigma_theta ** 2 + (self.beta * max(0.0, self.m - self.m0)) ** 2)
+
+
+def monitor_trajectory(est, surf, l=0.25, **kw):
+    """Run the monitor over a recorded trajectory. est: (T, 4); surf: (T, 4, H, W). Returns m_t."""
+    mon = EventConsistencyMonitor(l=l, **kw)
+    return np.array([mon.update(est[k], surf[k]) for k in range(len(est))])

aurasync/physics.py

@@ -0,0 +1,62 @@
+"""Kinetic plant: an inverted pendulum on a force-driven cart.
+
+This is the "unstable kinetic system" AuraSync must protect. The pole is short
+(0.5 m), so it falls fast: its natural time constant sqrt(l/g) is ~160 ms, which
+means a 25-40 ms perception delay eats a large share of the recovery window.
+
+State s = [x, x_dot, theta, theta_dot]; theta = 0 is upright, positive tilts right.
+"""
+from dataclasses import dataclass, replace
+import numpy as np
+
+
+@dataclass
+class PlantParams:
+    M: float = 1.0        # cart mass [kg]
+    m: float = 0.1        # pole mass [kg]
+    l: float = 0.25       # pole half-length [m] (full pole = 0.5 m)
+    g: float = 9.81
+    b_cart: float = 0.1   # cart viscous friction
+    b_pole: float = 0.002 # pivot friction
+    u_max: float = 20.0   # actuator force limit [N]
+    x_limit: float = 1.0  # track half-length [m]
+
+    def perturbed(self, rng, frac=0.3):
+        """Domain randomisation: scale masses / length by up to +-frac."""
+        s = lambda: 1.0 + rng.uniform(-frac, frac)
+        return replace(self, M=self.M * s(), m=self.m * s(), l=self.l * (1 + rng.uniform(-0.1, 0.1)))
+
+
+def deriv(s, u, p: PlantParams):
+    x, xd, th, thd = s
+    sin, cos = np.sin(th), np.cos(th)
+    total = p.M + p.m
+    temp = (u + p.m * p.l * thd ** 2 * sin - p.b_cart * xd) / total
+    thdd = (p.g * sin - cos * temp - p.b_pole * thd / (p.m * p.l)) / (
+        p.l * (4.0 / 3.0 - p.m * cos ** 2 / total))
+    xdd = temp - p.m * p.l * thdd * cos / total
+    return np.array([xd, xdd, thd, thdd])
+
+
+def rk4_step(s, u, dt, p: PlantParams):
+    k1 = deriv(s, u, p)
+    k2 = deriv(s + 0.5 * dt * k1, u, p)
+    k3 = deriv(s + 0.5 * dt * k2, u, p)
+    k4 = deriv(s + dt * k3, u, p)
+    return s + dt / 6.0 * (k1 + 2 * k2 + 2 * k3 + k4)
+
+
+def failed(s, p: PlantParams, max_angle=0.8):
+    """Kinetic failure = pole past ~46 degrees or cart off the track."""
+    return abs(s[2]) > max_angle or abs(s[0]) > p.x_limit
+
+
+def linearise(p: PlantParams, dt, eps=1e-6):
+    """Discrete-time (A, B) about the upright equilibrium via finite differences."""
+    s0 = np.zeros(4)
+    A = np.zeros((4, 4))
+    for i in range(4):
+        d = np.zeros(4); d[i] = eps
+        A[:, i] = (rk4_step(s0 + d, 0.0, dt, p) - rk4_step(s0 - d, 0.0, dt, p)) / (2 * eps)
+    B = ((rk4_step(s0, eps, dt, p) - rk4_step(s0, -eps, dt, p)) / (2 * eps)).reshape(4, 1)
+    return A, B

aurasync/safety.py

@@ -0,0 +1,173 @@
+"""Uncertainty- and delay-aware safety shield (Upgrade B).
+
+Analogy: a driving instructor with a second brake pedal. The student (LQR)
+steers; the instructor only intervenes when the car is heading for the kerb
+faster than it could stop, and leans in earlier when visibility is poor
+(high estimate uncertainty) or when the student's view is out of date (delay).
+
+Formulation (discrete-time control barrier function, one input, closed form):
+  Look-ahead coordinates (relative degree 1 in the force u):
+      psi = theta + c_th * theta_dot        chi = x + x_dot |x_dot| / (2 a_brake)
+  Safe set S = { |psi| <= psi_max  and  |chi| <= chi_max }. If |psi| <= psi_max
+  is kept forever, theta converges into the same band (first-order filter), so
+  this bounds the pole angle itself.
+  At each 1 ms tick, with the estimate first rolled forward over its age
+  (delay compensation, replaying the controls already applied), the safe set is
+  tightened by the estimate's uncertainty and the discrete CBF condition is
+  imposed on the tightened barriers:
+      h~ = (psi_max - kappa * sigma_psi) -/+ psi,   h~(s_{k+1}) >= (1 - gamma) h~(s_k)
+  (same for chi). sigma_psi is the barrier's standard deviation from the
+  estimator covariance propagated one step through the linearised physics, so
+  a vaguer or older estimate keeps the pole further from the edge (floor: 25%
+  of the nominal set). Dynamics are affine in u, so each constraint is an interval on u;
+  the QP  min (u - u_lqr)^2  becomes a clip onto the intersection with the
+  actuator limits.
+Priorities and fallback (hard limits before soft targets):
+  HARD  track end   |chi| <= chi_max, chi = x + x_dot|x_dot| / (2 a_brake): position plus the
+                    distance needed to stop at a_brake = 12 m/s^2 (actuator limit ~18 m/s^2)
+                    must stay inside 0.95 m (track 1.0 m) -> leaving it is a failure
+  HARD  pole        |psi| <= psi_hard (35 deg, falling is 46 deg)
+  SOFT  envelope    |psi| <= psi_max  (10 deg, the safety target), pursued only while
+                    |chi| <= chi_soft (0.6 m): the last part of the track is a reserve the
+                    soft target may never spend, so chasing it cannot strand the cart;
+                    and only from states with |psi| <= 2 psi_max (reachability gate, latched
+                    until settled: |psi| <= psi_max, |x_dot| < 0.3 m/s, |theta_dot| < 0.5 rad/s).
+                    After a hit too big to contain, chasing the target mid-recovery fights
+                    the nominal controller and walks the cart off the track
+  1. hard + soft + actuator feasible -> project u_lqr
+  2. else drop the soft envelope     -> project onto hard + actuator (logged 'soft_dropped')
+  3. else infeasible                 -> do no harm: pass the nominal command through (logged).
+                                        An earlier variant that slammed the actuator to the
+                                        most-violated barrier lost runs plain LQR survives.
+  (v3 first shipped with the soft envelope ranked above the track limit; on test trials
+  that drove the cart off the track in 21% of oracle runs. Fixed before reporting.)
+  Stale estimate (sigma_theta > stale_deg) -> the nominal command comes from a
+  conservative low-gain LQR before shielding (logged).
+"""
+import time
+from collections import deque
+import numpy as np
+
+from .physics import PlantParams, rk4_step, linearise
+from .controller import LQRController
+
+DT = 1e-3
+
+
+class SafetyShield:
+    def __init__(self, plant=None, psi_max_deg=10.0, c_th=0.10, chi_max=0.95, a_brake=12.0,
+                 gamma=0.03, kappa=2.0, stale_deg=4.0, Q_proc=None, psi_hard_deg=35.0, chi_soft=0.6,
+                 gamma_hard=0.03, soft_gate=2.0):
+        self.p = plant or PlantParams()
+        self.psi_max = np.radians(psi_max_deg)
+        self.psi_hard = np.radians(psi_hard_deg)
+        self.c_th, self.chi_max, self.a_brake, self.chi_soft = c_th, chi_max, a_brake, chi_soft
+        self.gamma, self.kappa, self.gamma_hard, self.soft_gate = gamma, kappa, gamma_hard, soft_gate
+        self.soft_off = False
+        self.stale = np.radians(stale_deg)
+        self.A, _ = linearise(self.p, DT)
+        self.Q = np.diag([1e-9, 1e-5, 1e-8, 1e-4]) if Q_proc is None else Q_proc
+        self.psi = lambda z: z[2] + c_th * z[3]
+        self.chi = lambda z: z[0] + z[1] * abs(z[1]) / (2 * a_brake)
+        self.fallback_ctrl = LQRController(self.p, Q=(10.0, 1.0, 100.0, 2.0), R=0.1)
+        self.u_max = self.p.u_max
+
+    @staticmethod
+    def _interval(a0, a1, lo, hi):
+        """{u : lo <= a0 + a1 u <= hi} as (u_lo, u_hi)."""
+        if abs(a1) < 1e-12:
+            return (-np.inf, np.inf) if lo <= a0 <= hi else (np.inf, -np.inf)
+        u1, u2 = (lo - a0) / a1, (hi - a0) / a1
+        return (min(u1, u2), max(u1, u2))
+
+    def _bounds(self, s, Sigma, fn, h_max, gamma=None):
+        """Interval on u for the pair of barriers h_max -/+ fn(s_next) (fn affine in u over one tick)."""
+        n0 = rk4_step(s, 0.0, DT, self.p)
+        n1 = rk4_step(s, 1.0, DT, self.p)
+        a0, a1, cur = fn(n0), fn(n1) - fn(n0), fn(s)
+        g = np.array([(fn(n0 + e) - fn(n0 - e)) / 2e-6 for e in np.eye(4) * 1e-6])   # barrier gradient
+        S1 = self.A @ Sigma @ self.A.T + self.Q
+        m = self.kappa * np.sqrt(max(g @ S1 @ g, 0.0))
+        hm = max(h_max - m, 0.25 * h_max)                 # robust tightened safe set |g.s| <= h_max - m
+        gm = self.gamma if gamma is None else gamma
+        hi = hm - (1 - gm) * (hm - cur)
+        lo = -hm + (1 - gm) * (hm + cur)
+        return self._interval(a0, a1, lo, hi), (a0, a1, lo, hi)
+
+    def filter(self, s, Sigma, u_nom):
+        """Returns (u_safe, mode). mode: 'pass', 'modified', 'soft_dropped', 'infeasible'."""
+        (sl, sh), _ = self._bounds(s, Sigma, self.psi, self.psi_max)          # soft
+        (pl, ph), pk = self._bounds(s, Sigma, self.psi, self.psi_hard, self.gamma_hard)   # hard
+        (cl, ch), ck = self._bounds(s, Sigma, self.chi, self.chi_max, self.gamma_hard)    # hard
+        (rl, rh), _ = self._bounds(s, Sigma, self.chi, self.chi_soft)         # track reserve for soft stage
+        ul, uh = -self.u_max, self.u_max
+        L, H = max(sl, rl, pl, cl, ul), min(sh, rh, ph, ch, uh)
+        a = abs(self.psi(s))
+        if a > self.soft_gate * self.psi_max:      # soft target out of reach: stop chasing it ...
+            self.soft_off = True
+        elif a <= self.psi_max and abs(s[1]) < 0.3 and abs(s[3]) < 0.5:   # ... until settled again (latch)
+            self.soft_off = False
+        if self.soft_off:
+            L, H = 1.0, -1.0
+        if L <= H:
+            u = float(np.clip(u_nom, L, H))
+            return u, ("pass" if abs(u - u_nom) < 1e-9 else "modified")
+        L, H = max(pl, cl, ul), min(ph, ch, uh)
+        if L <= H:
+            return float(np.clip(u_nom, L, H)), "soft_dropped"
+        # hard limits conflict: no command can be certified -> do no harm, defer to the nominal controller
+        return float(np.clip(u_nom, ul, uh)), "infeasible"
+
+
+class ShieldedController:
+    """LQR -> (stale check) -> delay compensation -> SafetyShield.
+
+    Reads the pipeline's covariance (`covariance()`) and estimate age
+    (`age_ticks()`) when it offers them; otherwise uses `sigma_const` and the
+    pipeline's fixed delay."""
+
+    def __init__(self, lqr, shield, pipeline, sigma_const=None, enabled=True):
+        self.lqr, self.shield, self.pipe = lqr, shield, pipeline
+        self.sigma_const = np.diag(np.square(sigma_const)) if sigma_const is not None else np.zeros((4, 4))
+        self.enabled = enabled
+        self.reset()
+
+    def reset(self):
+        self.u_hist = deque(maxlen=64)
+        self.stats = dict(pass_=0, modified=0, soft_dropped=0, infeasible=0, stale=0, ticks=0)
+        self.times = []
+
+    def _age(self):
+        if hasattr(self.pipe, "age_ticks"):
+            return int(self.pipe.age_ticks())
+        return int(getattr(self.pipe, "delay_ticks", 0))
+
+    def _cov(self):
+        if hasattr(self.pipe, "covariance"):
+            return np.asarray(self.pipe.covariance())
+        return self.sigma_const
+
+    def act(self, est):
+        t0 = time.perf_counter()
+        s = np.asarray(est, float)
+        Sigma = self._cov()
+        age = self._age()
+        us = list(self.u_hist)[-age:] if age > 0 else []
+        for u in us:                                     # delay compensation
+            s = rk4_step(s, u, DT, self.shield.p)
+            Sigma = self.shield.A @ Sigma @ self.shield.A.T + self.shield.Q
+        self.stats["ticks"] += 1
+        stale = np.sqrt(max(Sigma[2, 2], 0)) > self.shield.stale
+        if stale:
+            self.stats["stale"] += 1
+            u_nom = self.shield.fallback_ctrl.act(s)
+        else:
+            u_nom = self.lqr.act(s)
+        if self.enabled:
+            u, mode = self.shield.filter(s, Sigma, u_nom)
+            self.stats["pass_" if mode == "pass" else mode] += 1
+        else:
+            u = u_nom
+        self.u_hist.append(u)
+        self.times.append(time.perf_counter() - t0)
+        return u