Explicit-MPC Safety Shield for Neural Policies / shield_experiment.py

Mechanism failed

Raw ⬇ ZIP
 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()