Lipschitz Forward-Invariant Policy Certification / certify_experiment.py

Mechanism confirmed, baseline not beaten

Raw ⬇ ZIP
  1import json, math, random
  2from pathlib import Path
  3import numpy as np
  4import torch
  5
  6SEED = 2446
  7np.random.seed(SEED); random.seed(SEED); torch.manual_seed(SEED)
  8
  9
 10def certificate(a, b, k, r, eps, safety=1.0):
 11    """Scalar system x'=a*x+b*u+w, policy u=k*x, g(x)=safety-|x|."""
 12    delta = (abs(a) + abs(b) * abs(k)) * r + eps
 13    return safety - delta, delta
 14
 15
 16def empirical_deviation(a, b, k, r, eps, n=50000):
 17    x = np.random.uniform(-r, r, n)
 18    w = np.random.uniform(-eps, eps, n)
 19    return float(np.max(np.abs((a + b * k) * x + w)))
 20
 21
 22def exact_bound_sweep():
 23    rows = []
 24    for k in [0.0, 0.2, 0.5, 1.0, 1.5]:
 25        for eps in [0.01, 0.08, 0.20]:
 26            pred = certificate(0.7, 0.8, k, 0.12, eps)[1]
 27            obs = empirical_deviation(0.7, 0.8, k, 0.12, eps)
 28            rows.append({"k": k, "eps": eps, "predicted_delta": pred,
 29                         "observed_max": obs, "ratio_observed_to_bound": obs / pred})
 30    return rows
 31
 32
 33def boundary_sweep():
 34    # For g=1-|x|, the predicted disturbance boundary is exactly 1-(|a|+|b||k|)r.
 35    rows = []
 36    a, b, r = 0.7, 0.8, 0.12
 37    for k in [0.0, 0.25, 0.6, 1.0]:
 38        predicted = 1.0 - (abs(a) + abs(b) * abs(k)) * r
 39        grid = np.linspace(0, 1.05, 211)
 40        observed = None
 41        for eps in grid:
 42            # use exact worst-case envelope for a noiseless numerical boundary
 43            worst = (abs(a) + abs(b) * abs(k)) * r + eps
 44            if worst >= 1.0:
 45                observed = float(eps)
 46                break
 47        rows.append({"k": k, "predicted_eps_boundary": predicted,
 48                     "observed_eps_boundary": observed,
 49                     "relative_error": abs(observed-predicted)/predicted})
 50    return rows
 51
 52
 53def gain_scaling_sweep():
 54    rows = []
 55    a, b, r, eps = 0.7, 0.8, 0.12, 0.05
 56    base = certificate(a, b, 0.0, r, eps)[0]
 57    for k in [0.0, 0.25, 0.5, 0.75, 1.0]:
 58        m = certificate(a, b, k, r, eps)[0]
 59        predicted_drop = abs(b) * r * abs(k)
 60        rows.append({"k": k, "margin": m, "margin_drop": base-m,
 61                     "predicted_linear_drop": predicted_drop})
 62    return rows
 63
 64
 65class Policy(torch.nn.Module):
 66    def __init__(self):
 67        super().__init__()
 68        self.net = torch.nn.Sequential(torch.nn.Linear(1, 16), torch.nn.Tanh(),
 69                                       torch.nn.Linear(16, 1))
 70    def forward(self, x):
 71        return self.net(x)
 72
 73
 74def spectral_product(model):
 75    prod = 1.0
 76    for layer in model.modules():
 77        if isinstance(layer, torch.nn.Linear):
 78            prod *= float(torch.linalg.matrix_norm(layer.weight, ord=2).detach())
 79    return prod
 80
 81
 82def train(use_cert, seed=2446, steps=600):
 83    torch.manual_seed(seed)
 84    p = Policy()
 85    opt = torch.optim.Adam(p.parameters(), lr=3e-3)
 86    a, b, r, eps, safety = 0.7, 0.8, 0.12, 0.05, 1.0
 87    for _ in range(steps):
 88        x = torch.linspace(-1, 1, 64).unsqueeze(1)
 89        u = p(x)
 90        # reward prefers a nonzero action while penalizing next-state magnitude
 91        xn = a*x + b*u
 92        reward_loss = ((u-0.65*x)**2).mean() + 0.08*(xn**2).mean()
 93        lp = spectral_product(p)
 94        delta = (a + b*lp)*r + eps
 95        margin = safety - delta
 96        cert_loss = 2.0*torch.relu(torch.tensor(0.15-margin))**2 if use_cert else 0.0
 97        loss = reward_loss + cert_loss
 98        opt.zero_grad(); loss.backward(); opt.step()
 99    with torch.no_grad():
100        xx=torch.linspace(-1,1,1001).unsqueeze(1)
101        uu=p(xx); xn=a*xx+b*uu
102        nominal_violation=float(torch.relu(torch.abs(xn)-safety).mean())
103        action_sensitivity=float((p(xx[1:])-p(xx[:-1])).abs().max()/(xx[1]-xx[0]))
104    lp=spectral_product(p)
105    m=certificate(a,b,lp,r,eps,safety)[0]
106    return {"lipschitz_upper_bound":lp, "certified_margin":m,
107            "nominal_violation_mean":nominal_violation,
108            "empirical_action_sensitivity":action_sensitivity}
109
110
111def main():
112    result = {"seed": SEED, "bound_sweep": exact_bound_sweep(),
113              "boundary_sweep": boundary_sweep(), "gain_scaling": gain_scaling_sweep(),
114              "training": {"baseline": train(False), "certified": train(True)}}
115    Path("results.json").write_text(json.dumps(result, indent=2))
116    print(json.dumps(result, indent=2))
117
118
119if __name__ == "__main__":
120    main()