1import json, random
  2from pathlib import Path
  3import numpy as np
  4import torch
  5import torch.nn as nn
  6
  7SEED=2100
  8np.random.seed(SEED); random.seed(SEED); torch.manual_seed(SEED)
  9torch.set_num_threads(4)
 10device='cuda' if torch.cuda.is_available() else 'cpu'
 11try:
 12    if device=='cuda': torch.cuda.init()
 13except Exception:
 14    device='cpu'
 15
 16# Scalar robust interval: actuator [-1,1] intersect two opposing constraints.
 17# For the tested region x<2.8, M(x)=2*width-0.5*x.
 18def interval_np(x,width=.7):
 19    x=np.asarray(x)
 20    return np.maximum(-1.,-width+.25*x), np.minimum(1.,width-.25*x)
 21def interval_t(x,width=.7):
 22    z=torch.zeros_like(x)
 23    return torch.maximum(z-1.,-width+.25*x), torch.minimum(z+1.,width-.25*x)
 24def project_np(u,lo,hi): return np.minimum(np.maximum(u,lo),hi)
 25
 26def math_checks():
 27    xs=np.linspace(0,3.6,3601); lo,hi=interval_np(xs); M=hi-lo
 28    feasible=lo<=hi
 29    # Tolerance is needed because M at the exact floating grid boundary is roundoff-negative.
 30    boundary=float(xs[np.argmin(np.abs(M))])
 31    iff_err=int(np.sum(feasible!=(M>=-1e-10)))
 32    # Prediction 1: feasibility boundary x*=2*width/.5 = 2.8 at width=.7.
 33    # Prediction 2: hinge derivative is -2(m0-M) below m0 and zero above.
 34    m0=.7; g=np.linspace(0,1.4,141); e=1e-5
 35    f=lambda q: np.maximum(0,m0-q)**2
 36    fd=(f(g+e)-f(g-e))/(2*e); pred=np.where(g<m0,-2*(m0-g),0.)
 37    # Prediction 3: projection correction outside an endpoint is exactly endpoint-u.
 38    us=np.linspace(.7,1.8,100); endpoint=.4
 39    correction=project_np(us,-1,endpoint)-us
 40    return {'boundary_prediction':{'predicted_x':2.8,'observed_x':boundary,'abs_error':abs(boundary-2.8),'feasibility_iff_errors':iff_err},
 41            'hinge_prediction':{'finite_difference_rmse':float(np.sqrt(np.mean((fd-pred)**2))),
 42                                'max_derivative_above_m0':float(np.max(np.abs(fd[g>=m0])))},
 43            'projection_prediction':{'max_error_vs_endpoint_minus_action':float(np.max(np.abs(correction+(us-endpoint))))}}
 44
 45class Policy(nn.Module):
 46    def __init__(self):
 47        super().__init__()
 48        self.net=nn.Sequential(nn.Linear(1,24),nn.Tanh(),nn.Linear(24,24),nn.Tanh(),nn.Linear(24,1),nn.Tanh())
 49    def forward(self,x): return self.net(x).squeeze(-1)
 50
 51def rollout_metrics(p,width=.7,project=True):
 52    with torch.no_grad():
 53        x=torch.full((128,),1.45,device=device)
 54        mins=[]; actions=[]; corrections=[]
 55        for _ in range(18):
 56            u=p(x[:,None]); lo,hi=interval_t(x,width); M=hi-lo
 57            ue=torch.maximum(torch.minimum(u,hi),lo) if project else u
 58            mins.append(M); actions.append(u); corrections.append(ue-u)
 59            x=x+.075*ue+.018
 60        M=torch.cat(mins); C=torch.cat(corrections); U=torch.cat(actions)
 61        return {'p05_margin':float(torch.quantile(M,.05).cpu()),'min_margin':float(M.min().cpu()),
 62                'negative_margin_rate':float((M<0).float().mean().cpu()),
 63                'intervention_rate':float((C.abs()>1e-5).float().mean().cpu()),
 64                'correction_rms':float(torch.sqrt((C*C).mean()).cpu()),
 65                'intervention_energy':float((C*C).mean().cpu()),
 66                'mean_action':float(U.mean().cpu())}
 67
 68def train(use_margin=False,use_baseline=False,width=.7,seed=SEED):
 69    torch.manual_seed(seed); p=Policy().to(device); opt=torch.optim.Adam(p.parameters(),lr=4e-3)
 70    m0=.30
 71    for _ in range(180):
 72        x=torch.full((128,),1.45,device=device)
 73        task=0.; ml=0.; bl=0.
 74        for _ in range(18):
 75            u=p(x[:,None]); lo,hi=interval_t(x,width); M=hi-lo
 76            # Training uses the exact projection, with endpoint gradients retained for this smooth affine toy.
 77            ue=torch.maximum(torch.minimum(u,hi),lo)
 78            # Aggressive task drives x toward the boundary; margin regularization opposes that drift.
 79            task=task+(u-.82).pow(2).mean()
 80            ml=ml+torch.relu(m0-M).pow(2).mean()
 81            baseline=torch.clamp(-.18*x,-1,1)
 82            bl=bl+(ue-baseline).pow(2).mean()
 83            x=x+.075*ue+.018
 84        loss=task/28 + (3.0*ml/28 if use_margin else 0.) + (.12*bl/28 if use_baseline else 0.)
 85        opt.zero_grad(); loss.backward(); opt.step()
 86    r=rollout_metrics(p,width)
 87    r['task_action_mse']=float((p(torch.full((128,1),1.45,device=device))-.82).pow(2).mean().detach().cpu())
 88    return r
 89
 90def run():
 91    widths=[.55,.70,.85]
 92    sweep=[]
 93    for w in widths:
 94        b=train(False,False,w); m=train(True,False,w); mb=train(True,True,w)
 95        # Analytic width prediction at the initial state: Delta p05 M=2 Delta width.
 96        sweep.append({'width':w,'baseline':b,'margin':m,'margin_baseline':mb,
 97                      'predicted_initial_margin':2*w-.5*1.45})
 98    return {'device':device,'math_checks':math_checks(),'training_width_sweep':sweep}
 99
100if __name__=='__main__':
101    out=run(); Path('results.json').write_text(json.dumps(out,indent=2)); print(json.dumps(out,indent=2))