1import json
  2import math
  3import random
  4from dataclasses import dataclass
  5
  6import numpy as np
  7import torch
  8
  9SEED = 2686
 10random.seed(SEED)
 11np.random.seed(SEED)
 12torch.manual_seed(SEED)
 13torch.set_num_threads(4)
 14
 15
 16@dataclass
 17class HybridZonotope:
 18    """Small hybrid-zonotope container: c + G xi + H beta, with optional coupling."""
 19    c: np.ndarray
 20    G: np.ndarray
 21    H: np.ndarray
 22    E: np.ndarray | None = None
 23    F: np.ndarray | None = None
 24    b: np.ndarray | None = None
 25
 26    def affine(self, A, B=None, offset=None):
 27        A = np.asarray(A)
 28        c = A @ self.c
 29        G = A @ self.G
 30        H = A @ self.H
 31        if B is not None:
 32            c = c + np.asarray(B) @ np.asarray(offset)
 33        return HybridZonotope(c, G, H, self.E, self.F, self.b)
 34
 35    def interval(self):
 36        rad = np.abs(self.G).sum(axis=1) + np.abs(self.H).sum(axis=1)
 37        return self.c - rad, self.c + rad
 38
 39
 40def relu_interval_zonotope(z: HybridZonotope):
 41    """Sound interval relaxation of ReLU, retaining a zonotope approximation."""
 42    lo, hi = z.interval()
 43    out_lo, out_hi = np.maximum(lo, 0), np.maximum(hi, 0)
 44    c = (out_lo + out_hi) / 2
 45    G = np.diag((out_hi - out_lo) / 2)
 46    H = np.zeros((len(c), 0))
 47    return HybridZonotope(c, G, H)
 48
 49
 50def scalar_interval_policy(lo, hi, w1, b1, w2, b2):
 51    """Interval bound for a one-hidden-layer scalar ReLU policy."""
 52    # Each hidden preactivation is affine in scalar x.
 53    p0, p1 = w1[:, 0] * lo + b1, w1[:, 0] * hi + b1
 54    pl, ph = np.minimum(p0, p1), np.maximum(p0, p1)
 55    rl, rh = np.maximum(pl, 0), np.maximum(ph, 0)
 56    terms_l = np.where(w2[0] >= 0, w2[0] * rl, w2[0] * rh)
 57    terms_h = np.where(w2[0] >= 0, w2[0] * rh, w2[0] * rl)
 58    return float(terms_l.sum() + b2[0]), float(terms_h.sum() + b2[0])
 59
 60
 61def exact_radius(lam, r0, gamma, n):
 62    """Exact reachable radius for x+=lam*x+w, |w|<=gamma."""
 63    if abs(lam - 1.0) < 1e-12:
 64        return r0 + n * gamma
 65    return abs(lam) ** n * r0 + gamma * sum(abs(lam) ** i for i in range(n))
 66
 67
 68def mechanism_sweeps():
 69    r0, gamma, n = 0.2, 0.03, 8
 70    lam_values = [0.0, 0.2, 0.5, 0.8, 1.0, 1.2]
 71    rows = []
 72    for lam in lam_values:
 73        observed = exact_radius(lam, r0, gamma, n)
 74        predicted = (abs(lam) ** n * r0 + gamma * sum(abs(lam) ** i for i in range(n)))
 75        rows.append({"lambda": lam, "observed_rN": observed, "predicted_rN": predicted,
 76                     "abs_error": abs(observed - predicted)})
 77
 78    # At fixed lambda<1, the terminal radius is affine in gamma with known slope.
 79    lam, r0, n = 0.6, 0.2, 7
 80    gammas = [0.0, 0.01, 0.02, 0.04, 0.08]
 81    slope = sum(lam ** i for i in range(n))
 82    scaling = [{"gamma": g, "observed_rN": exact_radius(lam, r0, g, n),
 83                "predicted_rN": lam ** n * r0 + g * slope} for g in gammas]
 84
 85    # Constraint boundary: r_N <= R gives gamma* exactly.
 86    lam, r0, n, R = 0.7, 0.2, 6, 0.30
 87    geom = sum(lam ** i for i in range(n))
 88    gamma_star = (R - lam ** n * r0) / geom
 89    boundary = [{"gamma": g, "violation": max(0.0, exact_radius(lam, r0, g, n) - R)}
 90                for g in [0.0, gamma_star * 0.99, gamma_star, gamma_star * 1.01]]
 91
 92    # Contraction boundary: disturbance-free radius decays iff |lambda|<1.
 93    contraction = [{"lambda": lam, "ratio_rN_r0": exact_radius(lam, 0.2, 0.0, 10) / 0.2,
 94                    "predicted_contracts": abs(lam) < 1.0}
 95                   for lam in [0.8, 0.99, 1.0, 1.01, 1.2]]
 96    return {"radius_recurrence": rows, "disturbance_linear_scaling": scaling,
 97            "constraint_boundary": {"predicted_gamma_star": gamma_star, "samples": boundary},
 98            "contraction_boundary": contraction}
 99
100
101def rollout_loss(k, disturbances, x0=0.8, horizon=8):
102    x = torch.tensor(x0, dtype=torch.float32)
103    loss = 0.0
104    for t in range(horizon):
105        loss = loss + x * x
106        x = (1.0 - k) * x + disturbances[t]
107    return loss + 2.0 * x * x
108
109
110def reach_loss(k, gamma, x0=0.8, radius0=0.08, horizon=8, limit=0.45):
111    """Differentiable interval certificate for x+=x+u+w, u=-k*x."""
112    center = torch.tensor(x0, dtype=torch.float32)
113    radius = torch.tensor(radius0, dtype=torch.float32)
114    total = 0.0
115    for _ in range(horizon):
116        total = total + torch.relu(torch.abs(center) + radius - limit) ** 2
117        center = (1.0 - k) * center
118        radius = torch.abs(1.0 - k) * radius + gamma
119        total = total + 0.5 * torch.relu(radius - limit) ** 2
120    total = total + 5.0 * torch.relu(torch.abs(center) + radius - 0.25) ** 2
121    # contraction deficit, active when the closed-loop gain is not below one
122    total = total + 0.2 * torch.relu(torch.abs(1.0 - k) - 0.9) ** 2
123    return total
124
125
126def train_controllers(gamma, steps=500):
127    # Same scalar controller and initialization; baseline sees nominal trajectories only.
128    torch.manual_seed(SEED)
129    raw_b = torch.nn.Parameter(torch.tensor(-0.2))
130    opt_b = torch.optim.Adam([raw_b], lr=0.035)
131    nominal = torch.zeros(8)
132    for _ in range(steps):
133        opt_b.zero_grad()
134        k = torch.sigmoid(raw_b) * 2.0
135        loss = rollout_loss(k, nominal)
136        loss.backward(); opt_b.step()
137    kb = float((torch.sigmoid(raw_b) * 2).detach())
138
139    torch.manual_seed(SEED)
140    raw_r = torch.nn.Parameter(torch.tensor(-0.2))
141    opt_r = torch.optim.Adam([raw_r], lr=0.035)
142    for _ in range(steps):
143        opt_r.zero_grad()
144        k = torch.sigmoid(raw_r) * 2.0
145        loss = reach_loss(k, gamma)
146        loss.backward(); opt_r.step()
147    kr = float((torch.sigmoid(raw_r) * 2).detach())
148
149    def cert(k):
150        radii = [exact_radius(abs(1-k), 0.08, gamma, t) for t in range(9)]
151        return max(radii), radii[-1], max(0.0, radii[-1] - 0.45)
152    return {"gamma": gamma, "baseline_k": kb, "reach_k": kr,
153            "baseline_certificate": cert(kb), "reach_certificate": cert(kr),
154            "baseline_nominal_loss": float(rollout_loss(torch.tensor(kb), nominal)),
155            "reach_training_loss": float(reach_loss(torch.tensor(kr), gamma))}
156
157
158def main():
159    results = {"seed": SEED, "mechanism": mechanism_sweeps(),
160               "training": [train_controllers(g) for g in [0.0, 0.12, 0.30, 0.60]]}
161    print(json.dumps(results, indent=2))
162
163
164if __name__ == "__main__":
165    main()