SOS Backup Shield for Learned Policies / sos_backup_shield.py
Mechanism confirmed, baseline not beaten
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()