import json, math, random from pathlib import Path import numpy as np from scipy.linalg import solve_continuous_are import torch import torch.nn as nn SEED=2147 np.random.seed(SEED); random.seed(SEED); torch.manual_seed(SEED) def l2(x): return float(np.sqrt(np.mean(np.asarray(x)**2))) def main(): # Stable, controllable continuous-time LQR system. A=np.array([[0.,1.],[-1.0,-0.35]]) B=np.array([[0.],[1.]]) Q=np.diag([1., .4]); R=np.array([[1.]]) ; rho=.2 Ar=A-rho/2*np.eye(2) P=solve_continuous_are(Ar,B,Q,R) K=np.linalg.solve(R,B.T@P) # Prediction 1: generator finite differences converge linearly in dt. # Observable phi(x)=[x1^2, x1*x2, sin(x1)+x2^2], evaluated on exact Euler-free flow. def phi(x): return np.stack([x[:,0]**2, x[:,0]*x[:,1], np.sin(x[:,0])+x[:,1]**2],1) def dphi(x): f=x@A.T 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) xs=np.random.default_rng(SEED+1).uniform(-1,1,(1000,2)) # sufficiently accurate RK4 flow for each batch point def flow(x,dt): k1=x@A.T; k2=(x+.5*dt*k1)@A.T; k3=(x+.5*dt*k2)@A.T; k4=(x+dt*k3)@A.T return x+dt*(k1+2*k2+2*k3+k4)/6 dts=np.array([.2,.1,.05,.025,.0125]) gen_err=[] for dt in dts: gen_err.append(l2((phi(flow(xs,dt))-phi(xs))/dt-dphi(xs))) slope_gen=float(np.polyfit(np.log(dts),np.log(gen_err),1)[0]) # Prediction 2: finite-data affine generator identification improves as N^-1/2 # with fixed observation noise. Data are independently sampled states and controls. rng=np.random.default_rng(SEED+2); dt=.01; noise=.02 Ns=[32,64,128,256,512,1024] id_err=[] for N in Ns: vals=[] for rep in range(12): x=rng.uniform(-1,1,(N,2)); u=rng.normal(0,1,(N,1)) y=x@A.T+u@B.T + rng.normal(0,noise,(N,2)) Z=np.concatenate([x,u],1) W=np.linalg.solve(Z.T@Z+1e-5*np.eye(3),Z.T@y) xt=rng.uniform(-1,1,(1000,2)); ut=rng.normal(0,1,(1000,1)) pred=np.concatenate([xt,ut],1)@W truth=xt@A.T+ut@B.T vals.append(l2(pred-truth)) id_err.append(float(np.mean(vals))) slope_id=float(np.polyfit(np.log(Ns),np.log(id_err),1)[0]) # Koopman/EDMD-style lifted generator: regress finite-difference observable # derivatives against the lifted dictionary and control, then test holdout. def psi(x): return np.stack([x[:,0], x[:,1], x[:,0]**2, x[:,0]*x[:,1], x[:,1]**2, np.sin(x[:,0])], 1) xl=rng.uniform(-1,1,(800,2)); ul=rng.normal(0,1,(800,1)) dt_l=.002 yl=psi(flow(xl + 0*ul,dt_l)) # autonomous dictionary sanity check # controlled one-step flow uses the constant sampled control over dt xc=xl + dt_l*(xl@A.T + ul@B.T) zd=(psi(xc)-psi(xl))/dt_l Z=np.concatenate([psi(xl),ul],1) C=np.linalg.solve(Z.T@Z+1e-6*np.eye(Z.shape[1]),Z.T@zd) xt=rng.uniform(-1,1,(1000,2)); ut=rng.normal(0,1,(1000,1)) xnext=xt+dt_l*(xt@A.T+ut@B.T) lift_err=l2(np.concatenate([psi(xt),ut],1)@C-(psi(xnext)-psi(xt))/dt_l) # Prediction 3: exact LQR value has zero HJB residual; multiplicative G error # produces first-order residual in epsilon near zero. xx=rng.uniform(-2,2,(3000,2)); grad=xx@P.T q=.5*np.sum((xx@Q)*xx,1); V=.5*np.sum((xx@P)*xx,1) eps=np.array([0,.01,.02,.05,.1,.2]) res=[] for e in eps: Bh=(1+e)*B fh=xx@A.T; gh=grad@Bh rr=q-rho*V+np.sum(grad*fh,1)-.5*np.sum(gh*gh,1) res.append(l2(rr)) slope_res=float(np.polyfit(np.log(eps[1:]),np.log(res[1:]),1)[0]) # Small neural comparison: HJB residual training versus one-step TD regression. torch.set_num_threads(2); device='cuda' if torch.cuda.is_available() else 'cpu' try: dev=torch.device(device); _=torch.zeros(1,device=dev) except Exception: dev=torch.device('cpu') X=torch.tensor(rng.uniform(-1.5,1.5,(512,2)),dtype=torch.float32,device=dev) Pt=torch.tensor(P,dtype=torch.float32,device=dev); At=torch.tensor(A,dtype=torch.float32,device=dev) Bt=torch.tensor(B,dtype=torch.float32,device=dev); Qt=torch.tensor(Q,dtype=torch.float32,device=dev) class VNet(nn.Module): def __init__(self): super().__init__(); self.net=nn.Sequential(nn.Linear(2,32),nn.Tanh(),nn.Linear(32,32),nn.Tanh(),nn.Linear(32,1)) def forward(self,x): return self.net(x).squeeze(-1) def train_hjb(): torch.manual_seed(SEED); net=VNet().to(dev); opt=torch.optim.Adam(net.parameters(),lr=3e-3) for _ in range(220): xb=X.detach().clone().requires_grad_(True); v=net(xb); g=torch.autograd.grad(v.sum(),xb,create_graph=True)[0] qv=.5*torch.sum((xb@Qt)*xb,1); f=xb@At.T; gb=g@Bt rr=qv-rho*v+torch.sum(g*f,1)-.5*torch.sum(gb*gb,1) loss=(rr**2).mean()+1e-3*(v[X[:,0].abs()<.03]**2).mean() opt.zero_grad(); loss.backward(); opt.step() with torch.no_grad(): pred=net(X); target=.5*torch.sum((X@Pt)*X,1) return l2((pred-target).cpu().numpy()) def train_td(): torch.manual_seed(SEED); net=VNet().to(dev); opt=torch.optim.Adam(net.parameters(),lr=3e-3) # Standard discounted one-step TD target under optimal linear policy, same states. xnp=X.detach().cpu().numpy(); u=-(xnp@K.T); xn=xnp+0.02*(xnp@A.T+u@B.T) xn=torch.tensor(xn,dtype=torch.float32,device=dev); un=torch.tensor(u,dtype=torch.float32,device=dev) cost=.02*(.5*np.sum((xnp@Q)*xnp,1)+.5*np.sum((u@R)*u,1)) cost=torch.tensor(cost,dtype=torch.float32,device=dev) gamma=math.exp(-rho*.02) for _ in range(220): target=cost+gamma*net(xn).detach(); loss=((net(X)-target)**2).mean() opt.zero_grad(); loss.backward(); opt.step() with torch.no_grad(): pred=net(X); target=.5*torch.sum((X@Pt)*X,1) return l2((pred-target).cpu().numpy()) hjb_err=train_hjb(); td_err=train_td() out={'system':{'A':A.tolist(),'B':B.tolist(),'P':P.tolist(),'rho':rho}, 'predictions':{ '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}, '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}, 'lifted_generator':{'claim':'EDMD lifted generator predicts observable derivatives', 'dictionary_dim':6, 'holdout_rms_derivative_error':lift_err}, '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}}, 'comparison':{'metric':'RMS value error against discounted CARE solution after 220 updates','hjb':hjb_err,'td_baseline':td_err,'device':str(dev)}} Path('results.json').write_text(json.dumps(out,indent=2)) print(json.dumps(out,indent=2)) if __name__=='__main__': main()