1import sys, json, math
  2from pathlib import Path
  3import numpy as np
  4import torch
  5import torch.nn as nn
  6sys.path.insert(0, '/home/maxwelhelp/all/math2nn')
  7from bench import get_dataset, make_model, train_model, evaluate, sweep_baseline, make_report
  8
  9# Fixed a priori settings. The benchmark's input is 8 (theta, omega, u) tuples;
 10# the model predicts a constant control for the next target rollout.
 11DT = 0.05 / 4.0
 12SUBSTEPS = 32                 # eight benchmark intervals, four integration substeps each
 13THETA_SAFE = 1.40
 14K1, K2 = 2.0, 2.0
 15U_MAX = 1.5
 16
 17class RolloutSystem(nn.Module):
 18    def __init__(self, input_shape, shield=False, k1=K1, k2=K2):
 19        super().__init__()
 20        self.base = make_model('rnn_small', input_shape, 1)
 21        self.shield = shield
 22        self.k1, self.k2 = float(k1), float(k2)
 23
 24    def project(self, th, om, u_nom):
 25        # h = theta_safe^2 - theta^2; hdot=-2 theta omega.
 26        # psi2 = Lf2 h + LgLf h*u + k1*Lf h + k2*psi1 >= 0.
 27        h = THETA_SAFE**2 - th * th
 28        hd = -2.0 * th * om
 29        # hessian contribution: -2*om^2; drift contribution from pendulum.
 30        drift = -2.0 * om * om + 2.0 * th * (9.81 / 10.0) * torch.sin(th)
 31        psi1 = hd + self.k1 * h
 32        b = drift + self.k1 * hd + self.k2 * psi1
 33        a = -4.0 * th                  # Lg Lf h
 34        # Projection onto [-U_MAX,U_MAX] intersecting a*u+b >= 0.
 35        unclipped = torch.clamp(u_nom, -U_MAX, U_MAX)
 36        req = -b
 37        # scalar exact projection; where a is nearly zero, retain bounded action.
 38        candidate = torch.where(a.abs() > 1e-7, req / a, unclipped)
 39        candidate = torch.clamp(candidate, -U_MAX, U_MAX)
 40        feasible = (a * candidate + b >= 0)
 41        # If actuator bounds make the constraint infeasible, safest available bound.
 42        safe = torch.where(a > 0, torch.full_like(candidate, U_MAX),
 43                           torch.full_like(candidate, -U_MAX))
 44        safe = torch.where(a.abs() > 1e-7, safe, unclipped)
 45        candidate = torch.where(feasible, candidate, safe)
 46        # The min/max composition gives an exact scalar QP solution while remaining
 47        # differentiable almost everywhere for train_model's autograd.
 48        return torch.where((a * unclipped + b >= 0), unclipped, candidate)
 49
 50    def forward(self, x):
 51        # rnn_small accepts flattened windows in the bench implementation.
 52        u_nom = torch.tanh(self.base(x)).reshape(-1) * U_MAX
 53        th = x[:, -3]
 54        om = x[:, -2]
 55        for _ in range(SUBSTEPS):
 56            u = self.project(th, om, u_nom) if self.shield else u_nom
 57            om = om + (-9.81 / 10.0 * torch.sin(th) - 0.25 * om + 2.0 * u) * DT
 58            th = th + om * DT
 59        return th.reshape(-1, 1)
 60
 61def train_one(seed, lr, shield, k1=K1, k2=K2, epochs=12):
 62    torch.manual_seed(seed); np.random.seed(seed)
 63    ds = get_dataset('dynamics', seed, n_train=400, n_test=200)
 64    model = RolloutSystem(ds['input_shape'], shield=shield, k1=k1, k2=k2)
 65    _, metric, _ = train_model(model, ds, epochs=epochs, lr=lr, batch=128, log=lambda *_: None)
 66    return float(metric) if metric is not None else float('inf')
 67
 68def run_cfg(cfg, shield):
 69    return lambda seed: train_one(seed, cfg['lr'], shield, cfg.get('k1', K1), cfg.get('k2', K2))
 70
 71def signature(seed, lr, k1, k2):
 72    # Measure trained-model behavior, not an analytic toy identity.
 73    torch.manual_seed(seed); np.random.seed(seed)
 74    ds = get_dataset('dynamics', seed, n_train=400, n_test=200)
 75    m = RolloutSystem(ds['input_shape'], shield=True, k1=k1, k2=k2)
 76    m, _, _ = train_model(m, ds, epochs=12, lr=lr, batch=128, log=lambda *_: None)
 77    m.eval(); x = ds['xte'].to(next(m.parameters()).device)
 78    with torch.no_grad():
 79        u_nom = torch.tanh(m.base(x)).reshape(-1)*U_MAX
 80        th=x[:,-3]; om=x[:,-2]; raw=[]; sh=[]
 81        for _ in range(SUBSTEPS):
 82            h=THETA_SAFE**2-th*th; hd=-2*th*om
 83            drift=-2*om*om+2*th*(9.81/10)*torch.sin(th)
 84            psi1=hd+k1*h; b=drift+k1*hd+k2*psi1; a=-4*th
 85            raw.append((a*u_nom+b).cpu().numpy())
 86            u=m.project(th,om,u_nom); sh.append((a*u+b).cpu().numpy())
 87            om=om+(-9.81/10*torch.sin(th)-.25*om+2*u)*DT; th=th+om*DT
 88    raw=np.concatenate(raw); sh=np.concatenate(sh)
 89    predicted = float(np.mean(raw < 0))
 90    observed = float(np.mean(sh < -1e-5))
 91    return {'trained_model_raw_constraint_violation_rate':predicted,
 92            'trained_model_shielded_constraint_violation_rate':observed,
 93            'projection_activation_rate':float(np.mean(np.abs(raw-sh)>1e-5)),
 94            'predicted_vs_observed': {'predicted': predicted, 'observed': observed},
 95            'confirmed': bool(observed <= 0.01 and predicted > observed + 0.02)}
 96
 97def main():
 98    # Union of all idea learning rates is included in baseline sweep: parity.
 99    lrs=[1e-3,3e-3,1e-2]
100    baseline_grid=[{'lr':lr} for lr in lrs]
101    base=sweep_baseline(lambda c: run_cfg(c, False), baseline_grid, seeds=tuple(range(4)))
102    # Re-evaluate the selected baseline on all eight paired seeds.
103    base['full']=evaluate(run_cfg(base['best_cfg'], False), seeds=tuple(range(8)))
104    # Same three-config budget for idea; includes baseline-best and nearby settings.
105    idea_cfg=[{'lr':lr,'k1':k1,'k2':k2} for lr,k1,k2 in [(base['best_cfg']['lr'],2.,2.),(1e-3,2.,2.),(1e-2,2.,2.)]]
106    idea_trials=[]
107    for c in idea_cfg:
108        r=evaluate(run_cfg(c, True), seeds=tuple(range(4)))
109        idea_trials.append({'cfg':c,'mean':r['mean']})
110    best_idea_cfg=min(idea_cfg, key=lambda c: next(t['mean'] for t in idea_trials if t['cfg']==c))
111    idea=evaluate(run_cfg(best_idea_cfg, True), seeds=tuple(range(8)))
112    rep=make_report('dynamics','rnn_small',{'best_cfg':base['best_cfg'],'sweep':base['sweep'],'full':base['full']},idea,signature(0,best_idea_cfg['lr'],2.,2.))
113    rep['idea']['sweep']=idea_trials; rep['custom_track']=None
114    Path('bench_report.json').write_text(json.dumps(rep,indent=2))
115    print(json.dumps(rep,indent=2))
116if __name__=='__main__': main()