SOS Backup Shield for Learned Policies / sos_backup_shield.py

Mechanism confirmed, baseline not beaten

Raw ⬇ ZIP
  1import json
  2import math
  3import numpy as np
  4from scipy.linalg import solve_continuous_lyapunov
  5
  6SEED = 7
  7rng = np.random.default_rng(SEED)
  8DT = 0.01
  9U_MAX = 1.0
 10SAFE_P, SAFE_V = 2.0, 2.0
 11
 12# Double integrator: p_dot=v, v_dot=u.  The backup is a saturated PD law.
 13def backup(x, kp=2.0, kd=2.0):
 14    p, v = np.asarray(x)
 15    return float(np.clip(-kp*p-kd*v, -U_MAX, U_MAX))
 16
 17def plant_step(x, u, dt=DT):
 18    p, v = x
 19    # exact constant-input step
 20    return np.array([p + dt*v + .5*dt*dt*u, v + dt*u])
 21
 22def safe(x):
 23    return abs(x[0]) <= SAFE_P + 1e-9 and abs(x[1]) <= SAFE_V + 1e-9
 24
 25def terminal(x):
 26    return x[0]*x[0] + x[1]*x[1] <= 0.06**2
 27
 28def simulate_backup(x0, horizon, kp=2.0, kd=2.0):
 29    x = np.array(x0, dtype=float)
 30    n = int(round(horizon / DT))
 31    max_violation = 0.0
 32    for _ in range(n):
 33        max_violation = max(max_violation, max(abs(x[0])-SAFE_P, abs(x[1])-SAFE_V))
 34        x = plant_step(x, backup(x, kp, kd))
 35    max_violation = max(max_violation, max(abs(x[0])-SAFE_P, abs(x[1])-SAFE_V))
 36    return x, max_violation
 37
 38# Quadratic certificate V=x'Px. P is the exact Lyapunov solution for the
 39# unsaturated closed-loop matrix.  On a small enough sublevel set saturation
 40# is inactive and -Vdot = ||x||^2 exactly.
 41def certificate(kp=2.0, kd=2.0):
 42    A = np.array([[0., 1.], [-kp, -kd]])
 43    P = solve_continuous_lyapunov(A.T, -np.eye(2))
 44    return A, P
 45
 46def core_math_check():
 47    A, P = certificate()
 48    # numerical identity and rate bound on the certified ellipsoid V<=gamma
 49    Q = -(A.T @ P + P @ A)
 50    identity_err = float(np.max(np.abs(Q - np.eye(2))))
 51    gamma = 0.06**2 / np.max(np.linalg.eigvalsh(P))
 52    vals = []
 53    max_u = 0.
 54    for th in np.linspace(0, 2*np.pi, 2000, endpoint=False):
 55        z = np.array([math.cos(th), math.sin(th)])
 56        x = math.sqrt(gamma) * z / math.sqrt(z @ P @ z)
 57        u = -2*x[0]-2*x[1]
 58        max_u = max(max_u, abs(u))
 59        vals.append(float(x @ Q @ x))
 60    # Since Q=I, -dV/dt = ||x||^2 and normalized dissipation is 1/lambda_max(P).
 61    observed_rate = min(vals) / gamma
 62    predicted_rate = 1.0 / np.max(np.linalg.eigvalsh(P))
 63    return {"P": P.tolist(), "identity_error": identity_err,
 64            "gamma": gamma, "max_backup_u_on_certificate": max_u,
 65            "observed_min_neg_vdot_over_V": observed_rate,
 66            "predicted_rate_bound": predicted_rate,
 67            "certificate_saturation_free": bool(max_u <= U_MAX + 1e-10)}
 68
 69def certificate_scaling():
 70    # For V=x'P x and terminal disk radius r, gamma=r^2/lambda_max(P)
 71    # and area(Omega_0)=pi*gamma/sqrt(det(P)); hence area scales as r^2.
 72    _, P = certificate()
 73    lam = np.max(np.linalg.eigvalsh(P))
 74    det = np.linalg.det(P)
 75    rows = []
 76    base = math.pi * (0.06**2 / lam) / math.sqrt(det)
 77    for r in [0.03, 0.06, 0.12, 0.18]:
 78        area = math.pi * (r*r / lam) / math.sqrt(det)
 79        rows.append({"terminal_radius": r, "certified_area": area,
 80                     "area_over_base": area/base, "predicted_ratio": (r/0.06)**2})
 81    return rows
 82
 83def reachable_sweep():
 84    # Grid estimate of the finite-horizon backward reachable set to the
 85    # terminal safe disk, with safety checked throughout the backup rollout.
 86    points = np.array([(p,v) for p in np.linspace(-1.8,1.8,61)
 87                              for v in np.linspace(-1.8,1.8,61)])
 88    cell = (3.6/60.)**2
 89    out = []
 90    for T in [0.5, 1.0, 2.0, 3.0, 4.0]:
 91        good = 0
 92        for x in points:
 93            xf, viol = simulate_backup(x, T)
 94            if viol <= 1e-8 and terminal(xf):
 95                good += 1
 96        out.append({"T": T, "volume": good*cell, "fraction": good/len(points)})
 97    # input limit transition: longer horizons cannot rescue points requiring
 98    # excessive initial acceleration; report saturation prevalence.
 99    sat = []
100    for kp in [0.5, 1., 2., 4., 8.]:
101        count = 0
102        for x in points[::20]:
103            if abs(backup(x, kp, 2.0)) >= U_MAX-1e-12: count += 1
104        sat.append({"kp": kp, "saturated_fraction": count/len(points[::20])})
105    return out, sat
106
107def certified_rollout_audit(n=500, horizon=4.0):
108    _, P = certificate()
109    rg = np.random.default_rng(SEED + 2)
110    gamma = 0.06**2
111    safe_count = terminal_count = 0
112    max_violation = 0.0
113    for _ in range(n):
114        z = rg.normal(size=2)
115        z /= np.linalg.norm(z)
116        x = z * math.sqrt(gamma * rg.random() / (z @ P @ z))
117        xf, viol = simulate_backup(x, horizon)
118        safe_count += int(viol <= 1e-8)
119        terminal_count += int(terminal(xf))
120        max_violation = max(max_violation, viol)
121    return {"n": n, "safe_rate": safe_count / n,
122            "terminal_rate": terminal_count / n,
123            "max_violation": max_violation, "horizon": horizon}
124
125def nominal(x, rng):
126    # Fixed seeded proxy for a poorly trained learned policy: it tends to
127    # accelerate toward the positive position boundary and has small noise.
128    p, v = x
129    return float(np.clip(0.35 + 0.10*p + 0.10*v + rng.normal(0, .08), -1, 1))
130
131def shielded_rollout(x0, steps=700, rng=None):
132    x = np.array(x0, dtype=float); switches=0; violation=0.
133    for _ in range(steps):
134        h = SAFE_P*SAFE_P - x[0]*x[0]
135        # h_dot + alpha(h), matching the stated barrier feasibility test.
136        up = nominal(x, rng)
137        margin = -2*x[0]*x[1] + .5*h
138        A, P = certificate()
139        inside = x @ P @ x <= 0.06**2 + 1e-12
140        if margin >= 0 and inside:
141            u = up
142        else:
143            u = backup(x); switches += 1
144        x = plant_step(x, u)
145        violation = max(violation, max(0., abs(x[0])-SAFE_P, abs(x[1])-SAFE_V))
146    return violation, switches
147
148def baseline_rollout(x0, steps=700, rng=None):
149    x=np.array(x0,float); violation=0.
150    for _ in range(steps):
151        x=plant_step(x, nominal(x,rng))
152        violation=max(violation,max(0.,abs(x[0])-SAFE_P,abs(x[1])-SAFE_V))
153    return violation
154
155def shield_experiment():
156    # Initial states in a ring near the safety boundary, same draws for both.
157    rg=np.random.default_rng(SEED+1)
158    xs=[]
159    while len(xs)<1000:
160        x=rg.uniform([-1.9,-1.5],[1.9,1.5])
161        if abs(x[0])>1.4: xs.append(x)
162    bviol=[]; sviol=[]; sw=[]
163    for x in xs:
164        bviol.append(baseline_rollout(x, rng=np.random.default_rng(100+len(bviol))))
165        z,n=shielded_rollout(x, rng=np.random.default_rng(100+len(sviol)))
166        sviol.append(z); sw.append(n)
167    return {"n":len(xs), "baseline_violation_rate":float(np.mean(np.array(bviol)>1e-8)),
168            "shield_violation_rate":float(np.mean(np.array(sviol)>1e-8)),
169            "baseline_mean_violation":float(np.mean(bviol)),
170            "shield_mean_violation":float(np.mean(sviol)),
171            "shield_mean_switches":float(np.mean(sw))}
172
173def main():
174    reach, sat = reachable_sweep()
175    result={"seed":SEED, "math_check":core_math_check(),
176            "prediction_1_dissipation": "-Vdot/V >= 1/lambda_max(P) on unsaturated certificate",
177            "prediction_2_horizon": "backward reachable volume is nondecreasing with T",
178            "prediction_3_scaling": "quadratic certificate area scales as terminal_radius^2",
179            "certificate_scaling_sweep": certificate_scaling(),
180            "reachability_sweep":reach, "input_saturation_sweep":sat,
181            "certified_rollout_audit":certified_rollout_audit(),
182            "policy_comparison":shield_experiment()}
183    with open("results.json","w") as f: json.dump(result,f,indent=2)
184    print(json.dumps(result,indent=2))
185
186if __name__ == '__main__': main()