1import json, math, random
  2from pathlib import Path
  3import numpy as np
  4from scipy.linalg import solve_continuous_are
  5import torch
  6import torch.nn as nn
  7
  8SEED=2147
  9np.random.seed(SEED); random.seed(SEED); torch.manual_seed(SEED)
 10
 11def l2(x): return float(np.sqrt(np.mean(np.asarray(x)**2)))
 12
 13def main():
 14    # Stable, controllable continuous-time LQR system.
 15    A=np.array([[0.,1.],[-1.0,-0.35]])
 16    B=np.array([[0.],[1.]])
 17    Q=np.diag([1., .4]); R=np.array([[1.]]) ; rho=.2
 18    Ar=A-rho/2*np.eye(2)
 19    P=solve_continuous_are(Ar,B,Q,R)
 20    K=np.linalg.solve(R,B.T@P)
 21
 22    # Prediction 1: generator finite differences converge linearly in dt.
 23    # Observable phi(x)=[x1^2, x1*x2, sin(x1)+x2^2], evaluated on exact Euler-free flow.
 24    def phi(x): return np.stack([x[:,0]**2, x[:,0]*x[:,1], np.sin(x[:,0])+x[:,1]**2],1)
 25    def dphi(x):
 26        f=x@A.T
 27        return np.stack([2*x[:,0]*f[:,0], f[:,0]*x[:,1]+x[:,0]*f[:,1], np.cos(x[:,0])*f[:,0]+2*x[:,1]*f[:,1]],1)
 28    xs=np.random.default_rng(SEED+1).uniform(-1,1,(1000,2))
 29    # sufficiently accurate RK4 flow for each batch point
 30    def flow(x,dt):
 31        k1=x@A.T; k2=(x+.5*dt*k1)@A.T; k3=(x+.5*dt*k2)@A.T; k4=(x+dt*k3)@A.T
 32        return x+dt*(k1+2*k2+2*k3+k4)/6
 33    dts=np.array([.2,.1,.05,.025,.0125])
 34    gen_err=[]
 35    for dt in dts:
 36        gen_err.append(l2((phi(flow(xs,dt))-phi(xs))/dt-dphi(xs)))
 37    slope_gen=float(np.polyfit(np.log(dts),np.log(gen_err),1)[0])
 38
 39    # Prediction 2: finite-data affine generator identification improves as N^-1/2
 40    # with fixed observation noise. Data are independently sampled states and controls.
 41    rng=np.random.default_rng(SEED+2); dt=.01; noise=.02
 42    Ns=[32,64,128,256,512,1024]
 43    id_err=[]
 44    for N in Ns:
 45        vals=[]
 46        for rep in range(12):
 47            x=rng.uniform(-1,1,(N,2)); u=rng.normal(0,1,(N,1))
 48            y=x@A.T+u@B.T + rng.normal(0,noise,(N,2))
 49            Z=np.concatenate([x,u],1)
 50            W=np.linalg.solve(Z.T@Z+1e-5*np.eye(3),Z.T@y)
 51            xt=rng.uniform(-1,1,(1000,2)); ut=rng.normal(0,1,(1000,1))
 52            pred=np.concatenate([xt,ut],1)@W
 53            truth=xt@A.T+ut@B.T
 54            vals.append(l2(pred-truth))
 55        id_err.append(float(np.mean(vals)))
 56    slope_id=float(np.polyfit(np.log(Ns),np.log(id_err),1)[0])
 57
 58    # Koopman/EDMD-style lifted generator: regress finite-difference observable
 59    # derivatives against the lifted dictionary and control, then test holdout.
 60    def psi(x):
 61        return np.stack([x[:,0], x[:,1], x[:,0]**2, x[:,0]*x[:,1],
 62                         x[:,1]**2, np.sin(x[:,0])], 1)
 63    xl=rng.uniform(-1,1,(800,2)); ul=rng.normal(0,1,(800,1))
 64    dt_l=.002
 65    yl=psi(flow(xl + 0*ul,dt_l)) # autonomous dictionary sanity check
 66    # controlled one-step flow uses the constant sampled control over dt
 67    xc=xl + dt_l*(xl@A.T + ul@B.T)
 68    zd=(psi(xc)-psi(xl))/dt_l
 69    Z=np.concatenate([psi(xl),ul],1)
 70    C=np.linalg.solve(Z.T@Z+1e-6*np.eye(Z.shape[1]),Z.T@zd)
 71    xt=rng.uniform(-1,1,(1000,2)); ut=rng.normal(0,1,(1000,1))
 72    xnext=xt+dt_l*(xt@A.T+ut@B.T)
 73    lift_err=l2(np.concatenate([psi(xt),ut],1)@C-(psi(xnext)-psi(xt))/dt_l)
 74
 75    # Prediction 3: exact LQR value has zero HJB residual; multiplicative G error
 76    # produces first-order residual in epsilon near zero.
 77    xx=rng.uniform(-2,2,(3000,2)); grad=xx@P.T
 78    q=.5*np.sum((xx@Q)*xx,1); V=.5*np.sum((xx@P)*xx,1)
 79    eps=np.array([0,.01,.02,.05,.1,.2])
 80    res=[]
 81    for e in eps:
 82        Bh=(1+e)*B
 83        fh=xx@A.T; gh=grad@Bh
 84        rr=q-rho*V+np.sum(grad*fh,1)-.5*np.sum(gh*gh,1)
 85        res.append(l2(rr))
 86    slope_res=float(np.polyfit(np.log(eps[1:]),np.log(res[1:]),1)[0])
 87
 88    # Small neural comparison: HJB residual training versus one-step TD regression.
 89    torch.set_num_threads(2); device='cuda' if torch.cuda.is_available() else 'cpu'
 90    try:
 91        dev=torch.device(device); _=torch.zeros(1,device=dev)
 92    except Exception:
 93        dev=torch.device('cpu')
 94    X=torch.tensor(rng.uniform(-1.5,1.5,(512,2)),dtype=torch.float32,device=dev)
 95    Pt=torch.tensor(P,dtype=torch.float32,device=dev); At=torch.tensor(A,dtype=torch.float32,device=dev)
 96    Bt=torch.tensor(B,dtype=torch.float32,device=dev); Qt=torch.tensor(Q,dtype=torch.float32,device=dev)
 97    class VNet(nn.Module):
 98        def __init__(self):
 99            super().__init__(); self.net=nn.Sequential(nn.Linear(2,32),nn.Tanh(),nn.Linear(32,32),nn.Tanh(),nn.Linear(32,1))
100        def forward(self,x): return self.net(x).squeeze(-1)
101    def train_hjb():
102        torch.manual_seed(SEED); net=VNet().to(dev); opt=torch.optim.Adam(net.parameters(),lr=3e-3)
103        for _ in range(220):
104            xb=X.detach().clone().requires_grad_(True); v=net(xb); g=torch.autograd.grad(v.sum(),xb,create_graph=True)[0]
105            qv=.5*torch.sum((xb@Qt)*xb,1); f=xb@At.T; gb=g@Bt
106            rr=qv-rho*v+torch.sum(g*f,1)-.5*torch.sum(gb*gb,1)
107            loss=(rr**2).mean()+1e-3*(v[X[:,0].abs()<.03]**2).mean()
108            opt.zero_grad(); loss.backward(); opt.step()
109        with torch.no_grad(): pred=net(X); target=.5*torch.sum((X@Pt)*X,1)
110        return l2((pred-target).cpu().numpy())
111    def train_td():
112        torch.manual_seed(SEED); net=VNet().to(dev); opt=torch.optim.Adam(net.parameters(),lr=3e-3)
113        # Standard discounted one-step TD target under optimal linear policy, same states.
114        xnp=X.detach().cpu().numpy(); u=-(xnp@K.T); xn=xnp+0.02*(xnp@A.T+u@B.T)
115        xn=torch.tensor(xn,dtype=torch.float32,device=dev); un=torch.tensor(u,dtype=torch.float32,device=dev)
116        cost=.02*(.5*np.sum((xnp@Q)*xnp,1)+.5*np.sum((u@R)*u,1))
117        cost=torch.tensor(cost,dtype=torch.float32,device=dev)
118        gamma=math.exp(-rho*.02)
119        for _ in range(220):
120            target=cost+gamma*net(xn).detach(); loss=((net(X)-target)**2).mean()
121            opt.zero_grad(); loss.backward(); opt.step()
122        with torch.no_grad(): pred=net(X); target=.5*torch.sum((X@Pt)*X,1)
123        return l2((pred-target).cpu().numpy())
124    hjb_err=train_hjb(); td_err=train_td()
125    out={'system':{'A':A.tolist(),'B':B.tolist(),'P':P.tolist(),'rho':rho},
126      'predictions':{
127        'generator':{'claim':'finite-difference generator error is O(dt)','parameter':dts.tolist(),'observed_error':gen_err,'observed_loglog_slope':slope_gen,'predicted_slope':1.0},
128        'identification':{'claim':'fixed-noise regression error is approximately O(N^-1/2)','parameter':Ns,'observed_error':id_err,'observed_loglog_slope':slope_id,'predicted_slope':-0.5},
129        'lifted_generator':{'claim':'EDMD lifted generator predicts observable derivatives', 'dictionary_dim':6, 'holdout_rms_derivative_error':lift_err},
130        'hjb_sensitivity':{'claim':'zero at exact dynamics and first-order in multiplicative G error','parameter':eps.tolist(),'observed_rms_residual':res,'zero_error_residual':res[0],'observed_loglog_slope':slope_res,'predicted_slope':1.0}},
131      'comparison':{'metric':'RMS value error against discounted CARE solution after 220 updates','hjb':hjb_err,'td_baseline':td_err,'device':str(dev)}}
132    Path('results.json').write_text(json.dumps(out,indent=2))
133    print(json.dumps(out,indent=2))
134if __name__=='__main__': main()