1import json, math, random
  2from pathlib import Path
  3import numpy as np
  4import torch
  5import torch.nn as nn
  6
  7SEED = 2081
  8random.seed(SEED); np.random.seed(SEED); torch.manual_seed(SEED)
  9torch.set_num_threads(4)
 10
 11# Plant: xdot = u, safe set h(x)=x >= 0. Thus r(x)=u(x)+alpha*x.
 12ALPHA = 1.0
 13EPS = 0.08
 14
 15class Policy(nn.Module):
 16    def __init__(self):
 17        super().__init__()
 18        self.net = nn.Sequential(nn.Linear(1, 24), nn.Tanh(), nn.Linear(24, 24), nn.Tanh(), nn.Linear(24, 1))
 19    def forward(self, x):
 20        return 1.5 * torch.tanh(self.net(x))
 21
 22def certification_sweep():
 23    # For each uniform delta-net, construct a smooth residual that equals EPS
 24    # at net points and dips between them.  r=EPS+A(cos(2*pi*x/s)-1),
 25    # s=2*delta and A=delta/pi.  Its true Lipschitz constant is 1 and
 26    # the dense gap is exactly 2*delta/pi: a direct O(delta) prediction.
 27    dense = np.linspace(0, 1, 200001)
 28    rows=[]
 29    for n in [3, 5, 9, 17, 33, 65, 129]:
 30        s = 1/(n-1)
 31        delta = s/2
 32        A = delta/np.pi
 33        r_dense = EPS + A*(np.cos(2*np.pi*dense/s)-1)
 34        grid = np.linspace(0, 1, n)
 35        r_grid = EPS + A*(np.cos(2*np.pi*grid/s)-1)
 36        L = 1.0
 37        actual_min = float(r_dense.min())
 38        gap = float(EPS - actual_min)
 39        bound = EPS - L*delta
 40        rows.append(dict(n=n, delta=delta, L=L, predicted_lower=bound,
 41                         actual_min=actual_min, observed_gap=max(0.0,gap),
 42                         positive_certificate=(bound > 0), actual_safe=(actual_min >= 0)))
 43    x = np.linspace(0,1,1001)
 44    # Check the claimed global Lipschitz inequality numerically.
 45    max_pair_excess = 0.0
 46    for n in [3, 5, 9, 17, 33]:
 47        s=1/(n-1); A=(s/2)/np.pi
 48        r=EPS+A*(np.cos(2*np.pi*x/s)-1)
 49        for i in range(0,1001,7):
 50            for j in range(0,1001,11):
 51                max_pair_excess=max(max_pair_excess, abs(r[i]-r[j])-abs(x[i]-x[j]))
 52    bound_ok = all(row['actual_min'] + 1e-7 >= row['predicted_lower'] for row in rows)
 53    ratios = [row['observed_gap']/row['delta'] for row in rows]
 54    scaling_ok = max(abs(q-2/np.pi) for q in ratios) < 2e-4
 55    ds=np.array([q['delta'] for q in rows]); gs=np.array([q['observed_gap'] for q in rows])
 56    slope=float(np.polyfit(np.log(ds), np.log(gs), 1)[0])
 57    predicted_boundary=EPS/1.0
 58    observed_boundary=max(q['delta'] for q in rows if q['actual_safe'])
 59    return dict(L=1.0, rows=rows, predicted_boundary_delta=predicted_boundary,
 60                observed_safe_boundary_delta=observed_boundary,
 61                pairwise_lipschitz_excess=max_pair_excess,
 62                prediction_checks={'lower_bound_holds':bound_ok,
 63                                   'gap_ratio_predicted':2/np.pi,
 64                                   'gap_ratio_max_error':max(abs(q-2/np.pi) for q in ratios),
 65                                   'observed_loglog_slope':slope,
 66                                   'predicted_slope':1.0,
 67                                   'boundary_prediction_conservative': observed_boundary >= predicted_boundary,
 68                                   'all_checks': bool(bound_ok and scaling_ok and max_pair_excess <= 1e-8) },
 69                observed_gap_over_delta=ratios)
 70
 71def rollout(policy, x0, horizon=40, dt=0.05):
 72    x=x0
 73    xs=[]; us=[]; rs=[]
 74    for _ in range(horizon):
 75        x.requires_grad_(True)
 76        u=policy(x)
 77        # h=x, grad h=1, so residual is u+alpha*x.
 78        r=u + ALPHA*x
 79        xs.append(x); us.append(u); rs.append(r)
 80        x=x + dt*u
 81    return torch.cat(xs), torch.cat(us), torch.cat(rs)
 82
 83def train(use_barrier, steps=700):
 84    torch.manual_seed(SEED + (1 if use_barrier else 0))
 85    p=Policy()
 86    opt=torch.optim.Adam(p.parameters(), lr=3e-3)
 87    x0=torch.tensor([[0.20]])
 88    for _ in range(steps):
 89        xs,us,rs=rollout(p,x0)
 90        # Deliberately unsafe target demonstrates safety/task tradeoff.
 91        task=(xs[-1] - (-0.45)).pow(2) + 0.01*us.pow(2).mean()
 92        barrier=torch.relu(EPS-rs).pow(2).mean()
 93        loss=task + (1000.0*barrier if use_barrier else 0.0)
 94        opt.zero_grad(); loss.backward(); torch.nn.utils.clip_grad_norm_(p.parameters(), 10.0); opt.step()
 95    with torch.no_grad():
 96        xs,us,rs=rollout(p,x0)
 97        return {'final_x':float(xs[-1]), 'min_x':float(xs.min()), 'min_r':float(rs.min()),
 98                'task_terminal_error':float((xs[-1]+0.45).abs()),
 99                'barrier_violation_fraction':float((rs<EPS).float().mean()),
100                'dense_unseen_min_r':float((p(torch.linspace(0,1,2001).reshape(-1,1)).flatten()+ALPHA*torch.linspace(0,1,2001)).min()),
101                'policy':p}
102
103def main():
104    cert=certification_sweep()
105    base=train(False); idea=train(True)
106    # remove model object from serialization
107    base.pop('policy'); idea.pop('policy')
108    out={'seed':SEED,'system':'xdot=u, h=x, alpha=1', 'epsilon':EPS,
109         'certificate':cert, 'training':{'baseline_task_only':base,'barrier_certified':idea}}
110    Path('results.json').write_text(json.dumps(out, indent=2))
111    print(json.dumps(out, indent=2))
112
113if __name__=='__main__': main()