1from pathlib import Path
  2import json, math, random, sys
  3import numpy as np
  4import torch
  5from torch import nn
  6sys.path.insert(0, "/home/maxwelhelp/all/math2nn")
  7from bench import train_model, evaluate, sweep_baseline, make_report
  8
  9META = {"name": "sos_backup_shield_policy", "domain": "dynamics", "description": "Torque-limited pendulum action regression with a learned policy wrapped by a quadratic finite-horizon backup shield."}
 10
 11# Polynomial pendulum-like local model: theta_dot=omega, omega_dot=u.
 12# This is the explicit quadratic certificate used as the finite-horizon SOS-style MVP.
 13P = np.array([[1.25, .25], [.25, .375]], dtype=np.float64)
 14U_MAX = 1.0
 15TH_MAX, OM_MAX = 1.4, 2.0
 16CERT_RADIUS = 0.12
 17
 18
 19def math_check():
 20    A = np.array([[0., 1.], [-2., -2.]])
 21    Q = -(A.T @ P + P @ A)
 22    identity_error = float(np.max(np.abs(Q - np.eye(2))))
 23    lam = float(np.max(np.linalg.eigvalsh(P)))
 24    # Boundary samples of V=1, where -Vdot/V should be >= 1/lambda_max(P).
 25    vals = []
 26    max_u = 0.
 27    for a in np.linspace(0, 2*np.pi, 4096, endpoint=False):
 28        z = np.array([np.cos(a), np.sin(a)])
 29        x = CERT_RADIUS * z / np.sqrt(z @ P @ z)
 30        vals.append(float(x @ Q @ x))
 31        max_u = max(max_u, abs(-2*x[0] - 2*x[1]))
 32    observed = min(vals) / (CERT_RADIUS**2)
 33    predicted = 1.0 / lam
 34    return {"identity_error": identity_error, "observed_min_neg_vdot_over_V": observed,
 35            "predicted_bound": predicted, "max_backup_u_on_V1": max_u,
 36            "certificate_actuator_feasible": bool(max_u <= U_MAX)}
 37
 38
 39def get_dataset(seed, n_train=400, n_test=400):
 40    rng = np.random.RandomState(seed)
 41    def make(n):
 42        th = rng.uniform(-1.35, 1.35, n)
 43        om = rng.uniform(-2.0, 2.0, n)
 44        # Stabilizing target action, clipped to the actuator polytope.
 45        u = np.clip(-1.8 * th - 0.65 * om, -1., 1.)
 46        return np.stack([th, om], 1).astype("float32"), u[:, None].astype("float32")
 47    xtr, ytr = make(n_train); xte, yte = make(n_test)
 48    return {"xtr": xtr, "ytr": ytr, "xte": xte, "yte": yte,
 49            "task": "regression", "metric": "mse", "out_dim": 1,
 50            "input_shape": (2,)}
 51
 52
 53def seed_all(seed):
 54    random.seed(seed); np.random.seed(seed); torch.manual_seed(seed)
 55
 56
 57def tensors(d):
 58    return {**d, **{k: torch.as_tensor(d[k]) for k in ("xtr", "ytr", "xte", "yte")}}
 59
 60
 61class Policy(nn.Module):
 62    """Same base MLP for both systems; shield is the sole intervention."""
 63    def __init__(self, shield=False, margin=0.0):
 64        super().__init__()
 65        self.shield, self.margin = shield, float(margin)
 66        self.net = nn.Sequential(nn.Linear(2, 32), nn.Tanh(), nn.Linear(32, 32), nn.Tanh(), nn.Linear(32, 1))
 67
 68    def forward(self, x):
 69        nominal = self.net(x)
 70        if not self.shield:
 71            return torch.clamp(nominal, -U_MAX, U_MAX)
 72        # h=TH_MAX^2-theta^2; hdot+alpha h >= 0 gives admissible interval.
 73        th, om = x[:, 0:1], x[:, 1:2]
 74        h = TH_MAX**2 - th.square()
 75        alpha = 0.8
 76        # hdot=-2 theta omega; hdot + alpha*h + 2 theta*u >= margin.
 77        # For theta near zero, use the Lyapunov backup endpoint directly.
 78        lo = (self.margin + 2*th*om - alpha*h) / (2*th + 1e-6)
 79        hi = lo
 80        # Solve interval robustly by evaluating the affine barrier inequality.
 81        rhs = self.margin + 2*th*om - alpha*h
 82        pos_lo = rhs / (2*th + 1e-6)
 83        # inequality 2 theta u >= rhs; interval endpoint is lower for theta>0,
 84        # upper for theta<0. Pair with backup action and certificate gate.
 85        backup = torch.clamp(-2.0*th - 2.0*om, -U_MAX, U_MAX)
 86        feasible = (th.abs() > 0.04) & (rhs.abs() <= 2*th.abs()*U_MAX)
 87        safe_nom = torch.clamp(nominal, -U_MAX, U_MAX)
 88        # project toward the safe side, then choose backup outside the certificate.
 89        projected = torch.where(th > 0, torch.maximum(safe_nom, pos_lo), torch.minimum(safe_nom, pos_lo))
 90        V = 1.25*th.square() + .5*th*om + .375*om.square()
 91        inside = V <= CERT_RADIUS**2
 92        return torch.clamp(torch.where(feasible & inside, projected, backup), -U_MAX, U_MAX)
 93
 94
 95def train_metric(seed, cfg, keep=False):
 96    seed_all(seed); d = get_dataset(seed)
 97    model = Policy(shield=cfg.get("margin", 0.) > 0, margin=cfg.get("margin", 0.))
 98    model, metric, hist = train_model(model, tensors(d), epochs=cfg["epochs"], lr=cfg["lr"], batch=128, log=lambda *_: None)
 99    return (metric, model, d) if keep else metric
100
101
102def train_fn(cfg):
103    return lambda seed: train_metric(seed, cfg)
104
105
106def mechanism_signature(cfg, seeds=tuple(range(8))):
107    # Measured on outputs of trained models, not an analytic-only calculation.
108    shifts, dissipation, interventions, accepted = [], [], [], []
109    for s in seeds:
110        _, model, d = train_metric(s, cfg, keep=True)
111        model = model.cpu()
112        model.eval(); x = torch.as_tensor(d["xte"])
113        with torch.no_grad():
114            nominal = model.net(x); shielded = model(x)
115        th, om = x[:,0:1], x[:,1:2]
116        V = 1.25*th.square()+.5*th*om+.375*om.square()
117        u = shielded
118        # Euler certificate derivative for actual trained shield output.
119        vdot = th*om + (0.5*th + .75*om)*u
120        backup = torch.clamp(-2*th-2*om, -1., 1.)
121        shifts.append(float(torch.mean(torch.abs(shielded-nominal)).item()))
122        dissipation.append(float(torch.mean(vdot[V <= CERT_RADIUS**2]).item()))
123        interventions.append(float(torch.mean(torch.abs(shielded-backup)).item()))
124        accepted.append(float(torch.mean(((V <= CERT_RADIUS**2) & (torch.abs(shielded-nominal)<1e-5)).float()).item()))
125    # Prediction: inside the certified set the shield should be non-increasing;
126    # observed mean derivative should be non-positive on trained-model outputs.
127    obs = float(np.mean(dissipation))
128    return {"prediction": {"certificate_vdot_mean_on_inside": "<=0", "backup_intervention_nonzero": True},
129            "observed_from_trained_models": {"mean_vdot_inside": obs, "mean_action_shift": float(np.mean(shifts)),
130              "mean_distance_to_backup": float(np.mean(interventions)), "mean_nominal_acceptance": float(np.mean(accepted)), "n_models": len(seeds)},
131            "confirmed": bool(obs <= 1e-6 and float(np.mean(interventions)) > 1e-4)}
132
133
134def main():
135    check = math_check()
136    print(json.dumps({"math_check": check}))
137    if check["identity_error"] > 1e-10 or not check["certificate_actuator_feasible"]:
138        raise RuntimeError("certificate sanity check failed")
139    epochs = 15
140    lrs = [1e-3, 3e-3, 6e-3]
141    # Union parity: baseline sees every lr used by idea; method knob margin is swept
142    # on the idea side and zero-margin standard clipping is baseline.
143    base_grid = [{"lr": lr, "margin": 0.0, "epochs": epochs} for lr in lrs]
144    idea_grid = [{"lr": lr, "margin": margin, "epochs": epochs} for lr, margin in zip(lrs, [0.02, 0.05, 0.10])]
145    base = sweep_baseline(train_fn, base_grid)
146    idea_sweep = [{"cfg": c, "mean": evaluate(train_fn(c), seeds=(0,1,2,3))["mean"]} for c in idea_grid]
147    best = min(idea_sweep, key=lambda z: z["mean"])["cfg"]
148    idea = evaluate(train_fn(best))
149    report = make_report("sos_backup_shield_policy", "local_mlp_tiny", base, idea,
150                         mechanism_signature(best))
151    report["idea"]["selection_sweep"] = idea_sweep
152    report["custom_track"] = {"name": META["name"], "file": "sos_shield_bench.py", "domain": "dynamics"}
153    report["math_check"] = check
154    Path("bench_report.json").write_text(json.dumps(report, indent=2))
155    print(json.dumps(report, indent=2))
156
157if __name__ == "__main__": main()