Explicit-MPC Safety Shield for Neural Policies / shield_experiment.py
Mechanism failed
1import json
2from pathlib import Path
3import numpy as np
4from scipy.linalg import solve_discrete_are
5
6SEED=7
7rng=np.random.default_rng(SEED)
8A=np.array([[1.,.2],[0.,1.]])
9B=np.array([[.02],[.2]])
10XMAX=np.array([2.,2.]); UMAX=1.
11F=np.vstack([np.eye(2),-np.eye(2)]); fx=np.tile(XMAX,2)
12Q=np.diag([3.,1.]); R=np.array([[.15]])
13P=solve_discrete_are(A,B,Q,R)
14K=-(np.linalg.inv(R+B.T@P@B)@(B.T@P@A)).reshape(1,2)
15rho=max(abs(np.linalg.eigvals(A+B@K)))
16Acl=A+B@K
17
18def mpc(x): return float(np.clip((K@x)[0],-UMAX,UMAX))
19def support(g): return g*np.sum(np.abs(F),axis=1)
20def robust_ok(x,u,g):
21 return bool(np.all(F@(A@x+B[:,0]*u)+support(g)<=fx+1e-10) and abs(u)<=UMAX+1e-10)
22def shield(x,unn,g):
23 um=mpc(x)
24 if robust_ok(x,unn,g): return unn,0.,True
25 # choose smallest tested correction to the explicit fallback
26 for lam in np.linspace(0,1,41):
27 u=(1-lam)*unn+lam*um
28 if robust_ok(x,u,g): return u,lam,True
29 return um,1.,robust_ok(x,um,g)
30def rollout(policy,g,episodes=60,T=50,scale=1.):
31 violations=0; fallbacks=0; steps=0; terminal=0
32 for _ in range(episodes):
33 x=rng.uniform(-1.2,1.2,2); reached=False
34 for t in range(T):
35 unn=float(np.clip(scale*mpc(x)+rng.normal(0,.28),-1.5,1.5))
36 if policy=='shield': u,lam,ok=shield(x,unn,g); fallbacks+=lam>0
37 else: u=unn
38 w=rng.uniform(-g,g,2); xn=A@x+B[:,0]*u+w
39 violations += np.any(np.abs(xn)>XMAX+1e-9); steps+=1
40 x=xn
41 if np.linalg.norm(x)<.25: reached=True
42 terminal+=reached
43 return {'violation_rate':violations/steps,'fallback_rate':fallbacks/steps,'terminal_fraction':terminal/episodes}
44def verify_support():
45 errs=[]
46 for g in [.01,.1,.3]:
47 for _ in range(3000):
48 v=rng.uniform(-g,g,2)
49 errs.append(np.max(np.abs(F@v)-support(g))) # <=0, equality occurs at vertices
50 # independently maximize each face over box vertices
51 max_err=0.
52 for g in [.01,.1,.3]:
53 verts=np.array([[a,b] for a in [-g,g] for b in [-g,g]])
54 max_err=max(max_err,float(np.max(np.abs(np.max(verts@F.T,axis=0)-support(g)))))
55 return max_err
56
57def main():
58 support_err=verify_support()
59 # boundary prediction: robust state constraint at x=0 has gamma <= 2.
60 # Exact disturbance-only predicted boundary is gamma*=2; test direct feasibility.
61 boundary=[]
62 for g in np.linspace(0,2.4,25): boundary.append((float(g),robust_ok(np.zeros(2),0.,g)))
63 feasible=[g for g,ok in boundary if ok]
64 observed=max(feasible)
65 rows=[]
66 for g in [0.,.1,.3,.6,1.0]:
67 s=rollout('shield',g,scale=1.)
68 u=rollout('raw',g,scale=1.)
69 rows.append({'gamma':g,'shield':s,'unshielded':u})
70 gain=[]
71 for scale in [.5,1.,1.5,2.,3.]:
72 r=rollout('shield',.1,scale=scale)
73 gain.append({'policy_scale':scale,'fallback_rate':r['fallback_rate'],'violation_rate':r['violation_rate']})
74 # contraction prediction for fallback/no disturbance: ||Acl^t x|| <= C rho^t ||x||.
75 x=np.array([1.,1.]); norms=[]
76 for _ in range(15): norms.append(float(np.linalg.norm(x))); x=Acl@x
77 ratios=[norms[i+1]/norms[i] for i in range(len(norms)-1)]
78 out={'seed':SEED,'K':K.tolist(),'closed_loop_spectral_radius':float(rho),
79 'support_max_abs_error':support_err,'predicted_gamma_boundary':2.,'observed_grid_boundary':observed,
80 'safety_sweep':rows,'fallback_sweep':gain,'contraction_norms':norms,'contraction_ratios':ratios}
81 Path('results.json').write_text(json.dumps(out,indent=2))
82 print(json.dumps(out,indent=2))
83if __name__=='__main__': main()