1import json, math, random
  2from pathlib import Path
  3import numpy as np
  4import torch
  5import torch.nn as nn
  6
  7SEED = 2823
  8random.seed(SEED); np.random.seed(SEED); torch.manual_seed(SEED)
  9torch.set_num_threads(4)
 10device = "cuda" if torch.cuda.is_available() else "cpu"
 11
 12# Nonlinear controllable discrete plant. The task asks for a nonzero action,
 13# whereas the certificate asks for a stable closed loop around x*=0.
 14def plant(x, u):
 15    x1, x2 = x[:, :1], x[:, 1:2]
 16    return torch.cat((x1 + .05*(-x1**3 + u + .5*x2),
 17                      .95*x2 + .05*x1), dim=1)
 18
 19class Policy(nn.Module):
 20    def __init__(self):
 21        super().__init__()
 22        self.net = nn.Sequential(nn.Linear(2,32), nn.Tanh(), nn.Linear(32,32), nn.Tanh(), nn.Linear(32,1))
 23    def forward(self,x):
 24        # Enforce the equilibrium condition pi(0)=0 exactly.
 25        z=torch.zeros(1,2,device=x.device)
 26        return 2.0*(torch.tanh(self.net(x)/2.0)-torch.tanh(self.net(z)/2.0))
 27
 28class Lyapunov(nn.Module):
 29    # V=||g(x)-g(0)||^2 + eps ||x||^2, hence V(0)=0 and V>=eps||x||^2.
 30    def __init__(self, eps=.08):
 31        super().__init__(); self.eps=eps
 32        self.g=nn.Sequential(nn.Linear(2,24), nn.Tanh(), nn.Linear(24,2))
 33    def forward(self,x):
 34        z0=self.g(torch.zeros(1,2,device=x.device)).expand(x.shape[0],-1)
 35        return ((self.g(x)-z0)**2).sum(1,keepdim=True)+self.eps*(x*x).sum(1,keepdim=True)
 36
 37def batch(n=512, radius=1.0):
 38    # fixed distribution makes lambda comparisons less noisy
 39    return (torch.rand(n,2,device=device)*2-1)*radius
 40
 41def train(lam_delta, steps=1800):
 42    torch.manual_seed(SEED+int(lam_delta*1000))
 43    p, v = Policy().to(device), Lyapunov().to(device)
 44    opt=torch.optim.Adam(list(p.parameters())+list(v.parameters()),lr=3e-3)
 45    c2=.025
 46    for k in range(steps):
 47        x=batch(512,1.0)
 48        u=p(x); xn=plant(x,u); vv=v(x); vn=v(xn)
 49        # Task: imitate a stabilizing reference while retaining a policy-learning objective.
 50        target=-0.9*x[:, :1]-0.25*x[:, 1:2]
 51        task=((u-target)**2).mean()
 52        pos=torch.nn.functional.softplus(.002-vv).mean()
 53        dec=torch.nn.functional.softplus(vn-vv+c2*(x*x).sum(1,keepdim=True)).mean()
 54        loss=task + 0.2*pos + lam_delta*dec
 55        opt.zero_grad(); loss.backward(); opt.step()
 56    return p,v
 57
 58def evaluate(p,v,radius=1.0,n=101):
 59    a=torch.linspace(-radius,radius,n,device=device); X,Y=torch.meshgrid(a,a,indexing='ij')
 60    x=torch.stack((X.flatten(),Y.flatten()),1)
 61    with torch.no_grad():
 62        u=p(x); xn=plant(x,u); V=v(x); vn=v(xn); d=vn-V; norm=(x*x).sum(1,keepdim=True)
 63        margin=(d+.025*norm).flatten(); positivity=(V-.08*norm).flatten()
 64        cert=(margin<=0)&(positivity>=-.000001)
 65        # Largest centered grid radius where every point inside passes.
 66        rr=x.norm(dim=1)
 67        radii=torch.linspace(.05,radius,20,device=device)
 68        good=[]
 69        for r in radii:
 70            inside=rr<=r
 71            good.append(bool((~inside | cert).all()))
 72        cr=float(radii[max(i for i,z in enumerate(good) if z)]) if any(good) else 0.
 73        ratios=(vn/(V+1e-8)).flatten()
 74        certified_ratios=ratios[cert]
 75        return dict(task=float(((u-(-.9*x[:, :1]-.25*x[:, 1:2]))**2).mean()), violation=float(torch.relu(margin).mean()),
 76                    min_margin=float(margin.min()), pos_min=float(positivity.min()),
 77                    certified_fraction=float(cert.float().mean()), certified_radius=cr,
 78                    certified_ratio_max=float(certified_ratios.max()) if certified_ratios.numel() else float('nan'),
 79                    certified_ratio_median=float(certified_ratios.median()) if certified_ratios.numel() else float('nan'),
 80                    x=x.detach().cpu().numpy(), V=V.flatten().cpu().numpy(), margin=margin.detach().cpu().numpy())
 81
 82def rollout(p,v,x0,steps=30):
 83    x=torch.tensor(x0,dtype=torch.float32,device=device).view(1,2); vals=[]
 84    with torch.no_grad():
 85        for _ in range(steps): vals.append(float(v(x).item())); x=plant(x,p(x))
 86    return vals
 87
 88def math_checks():
 89    # Directly verify the advertised implication for synthetic values, including
 90    # the predicted boundary q=1-c2/c1 and its failure when c2>=c1.
 91    c1,c2=0.08,.025; q=1-c2/c1
 92    rng=np.random.default_rng(SEED); V0=rng.uniform(.1,10,5000)
 93    V=V0.copy(); ratios=[]
 94    for _ in range(8):
 95        V=(1-c2/c1)*V; ratios.append(V.copy())
 96    bound_errors=[float(np.max(ratios[t]-q**(t+1)*V0)) for t in range(8)]
 97    # A scalar parameter sweep demonstrates the exact transition at c2/c1=1.
 98    sweep=[]
 99    for c in [.01,.025,.079,.081,.10]:
100        qq=1-c/c1; sweep.append({'c2':c,'predicted_q':qq,'stable_bound':bool(0<qq<1),
101                                  'observed_q':qq})
102    return {'c1':c1,'c2':c2,'predicted_contraction_q':q,
103            'max_bound_errors':bound_errors,'contraction_bound_verified':max(map(abs,bound_errors))<1e-10,
104            'boundary_sweep':sweep}
105
106def main():
107    out={'device':device,'math':math_checks(),'runs':[]}
108    # lambda=0 is the standard task-only baseline; increasing lambda tests
109    # predicted monotone reduction in certificate violations.
110    for lam in [0., .05, .2, 1.0, 4.0]:
111        p,v=train(lam); ev=evaluate(p,v)
112        ev['lambda_delta']=lam
113        ev['rollout_V_ratio']=rollout(p,v,[.8,.7])[-1]/max(rollout(p,v,[.8,.7])[0],1e-12)
114        out['runs'].append({k:z for k,z in ev.items() if k not in ('x','V','margin')})
115    # Independent radius sweep on the strongest certificate model.
116    p,v=train(4.0)
117    out['radius_sweep']=[{'radius':r,**{k:z for k,z in evaluate(p,v,radius=r,n=81).items() if k in ('violation','min_margin','certified_fraction','certified_radius')}} for r in [.2,.4,.6,.8,1.0]]
118    Path('results.json').write_text(json.dumps(out,indent=2))
119    print(json.dumps(out,indent=2))
120
121if __name__=='__main__': main()