1import os, sys, json, random
  2from pathlib import Path
  3import numpy as np
  4import torch
  5import torch.nn as nn
  6
  7sys.path.insert(0, '/home/maxwelhelp/all/math2nn')
  8from bench import make_report, sweep_baseline, evaluate
  9from bench.custom_tracks.robust_cbf_pendulum_policy import get_dataset
 10
 11META = {'name': 'barrier_certified_pendulum_policy', 'domain': 'dynamics',
 12        'description': 'Pendulum state-to-action regression trained with differentiable zeroing-CBF residual hinge loss.'}
 13DEVICE = 'cuda' if torch.cuda.is_available() else 'cpu'
 14ALPHA, EPS, LAMBDA = 1.0, 0.08, 2.0
 15
 16class Policy(nn.Module):
 17    def __init__(self):
 18        super().__init__()
 19        self.net = nn.Sequential(nn.Linear(2,32), nn.Tanh(), nn.Linear(32,32), nn.Tanh(), nn.Linear(32,1))
 20    def forward(self,x): return torch.clamp(self.net(x), -1., 1.)
 21
 22def seed_all(s):
 23    random.seed(s); np.random.seed(s); torch.manual_seed(s)
 24
 25def dynamics(x,u):
 26    # normalized pendulum: theta_dot=omega, omega_dot=sin(theta)+u
 27    th, om = x[:,0:1], x[:,1:2]
 28    return torch.cat([om, torch.sin(th)+u], 1)
 29
 30def residual(x,u):
 31    # Smooth ellipsoidal safe set: h=1-theta^2/1.4^2-omega^2/2^2.
 32    # Its omega gradient is nonzero, so the relative-degree-one residual depends on u.
 33    th, om=x[:,0:1], x[:,1:2]
 34    h=1.0-(th/1.4)**2-(om/2.0)**2
 35    gh=torch.cat([-2*th/(1.4**2), -2*om/(2.0**2)], 1)
 36    return (gh*dynamics(x,u)).sum(1,keepdim=True) + ALPHA*h
 37
 38def train_metric(seed,cfg,keep=False):
 39    seed_all(seed)
 40    d=get_dataset(seed)
 41    xtr=torch.as_tensor(d['xtr'],device=DEVICE); ytr=torch.as_tensor(d['ytr'],device=DEVICE)
 42    model=Policy().to(DEVICE); opt=torch.optim.Adam(model.parameters(),lr=cfg['lr'])
 43    bs=128; model.train()
 44    try:
 45        for _ in range(cfg['epochs']):
 46            perm=torch.randperm(len(xtr),device=DEVICE)
 47            for ii in perm.split(bs):
 48                x=xtr[ii]; target=ytr[ii]; u=model(x)
 49                task=(u-target).pow(2).mean()
 50                if cfg.get('barrier',False):
 51                    # one-step differentiable rollout states, including training states
 52                    r=residual(x,u)
 53                    barrier=torch.relu(EPS-r).pow(2).mean()
 54                    loss=task+cfg.get('weight',LAMBDA)*barrier
 55                else: loss=task
 56                opt.zero_grad(); loss.backward(); torch.nn.utils.clip_grad_norm_(model.parameters(),10); opt.step()
 57        model.eval()
 58        with torch.no_grad(): metric=float(((model(torch.as_tensor(d['xte'],device=DEVICE))-torch.as_tensor(d['yte'],device=DEVICE))**2).mean().item())
 59    except Exception:
 60        if DEVICE=='cuda':
 61            torch.cuda.empty_cache(); return train_metric_cpu(seed,cfg,keep)
 62        raise
 63    return (metric,model.cpu(),d) if keep else metric
 64
 65def train_metric_cpu(seed,cfg,keep=False):
 66    global DEVICE; old=DEVICE; DEVICE='cpu'
 67    try: return train_metric(seed,cfg,keep)
 68    finally: DEVICE=old
 69
 70def train_fn(cfg): return lambda seed: train_metric(seed,cfg)
 71
 72def signature(cfg,seeds=tuple(range(8))):
 73    observed=[]; gaps=[]; violations=[]
 74    for s in seeds:
 75        _,m,d=train_metric(s,cfg,True); x=torch.as_tensor(d['xte']); u=m(x)
 76        r=residual(x,u).detach().numpy().ravel()
 77        observed.append(float(r.min())); violations.append(float(np.mean(r<EPS)))
 78        # local finite-difference Lipschitz estimate on trained model residual
 79        xx=torch.as_tensor(np.linspace(-1.35,1.35,257)[:,None], dtype=torch.float32).repeat(1,2)
 80        xx[:,1]=0; rr=residual(xx,m(xx)).detach().numpy().ravel()
 81        gaps.append(float(np.max(np.abs(np.diff(rr))/(2.7/256))))
 82    L=float(np.mean(gaps)); delta=0.5*(2.7/256)
 83    lower=EPS-L*delta
 84    return {'prediction':{'lower_bound': 'epsilon - L_r delta', 'delta':delta},
 85            'observed_from_trained_models':{'min_residual_mean':float(np.mean(observed)),
 86              'violation_fraction_mean':float(np.mean(violations)), 'L_r_estimate_mean':L,
 87              'bound_lower_mean':lower, 'n_models':len(seeds)},
 88            'confirmed': bool(float(np.mean(observed)) >= lower-1e-3)}
 89
 90def main():
 91    epochs=15; lrs=[1e-3,3e-3,6e-3]
 92    # union parity: all learning rates appear on both sides; baseline also sweeps its sole central knob (lr).
 93    base_grid=[{'lr':lr,'epochs':epochs,'barrier':False} for lr in lrs]
 94    idea_grid=[{'lr':lr,'epochs':epochs,'barrier':True,'weight':w} for lr,w in zip(lrs,[1.,2.,4.])]
 95    base=sweep_baseline(train_fn,base_grid)
 96    sel=[{'cfg':c,'mean':evaluate(train_fn(c),seeds=(0,1,2,3))['mean']} for c in idea_grid]
 97    best=min(sel,key=lambda z:z['mean'])['cfg']; idea=evaluate(train_fn(best))
 98    report=make_report('robust_cbf_pendulum_policy','shared_mlp_32',base,idea,signature(best))
 99    report['idea']['selection_sweep']=sel
100    report['custom_track']={'name':META['name'],'file':'barrier_bench.py','domain':'dynamics'}
101    Path('bench_report.json').write_text(json.dumps(report,indent=2)); print(json.dumps(report,indent=2))
102
103if __name__=='__main__': main()