1import json
  2import math
  3import random
  4import numpy as np
  5import torch
  6import torch.nn as nn
  7
  8SEED = 2975
  9random.seed(SEED)
 10np.random.seed(SEED)
 11torch.manual_seed(SEED)
 12
 13device = "cuda" if torch.cuda.is_available() else "cpu"
 14try:
 15    if device == "cuda":
 16        torch.cuda.empty_cache()
 17except Exception:
 18    device = "cpu"
 19torch.set_default_dtype(torch.float64)
 20
 21# A residual neural block, z+eta*r(z,u), with an exact state Jacobian.
 22class ResidualBlock(nn.Module):
 23    def __init__(self, d=2, hidden=16, eta=0.25, constrained=False, alpha=0.20):
 24        super().__init__()
 25        self.d, self.eta, self.constrained, self.alpha = d, eta, constrained, alpha
 26        self.net = nn.Sequential(nn.Linear(2*d, hidden), nn.Tanh(), nn.Linear(hidden, d))
 27        if constrained:
 28            # This is an explicit safe residual parameterization: the learned
 29            # nonlinear correction is bounded, while the state Jacobian is
 30            # damped by a fixed contractive skip coefficient.
 31            self.skip = math.sqrt(1.0 - alpha)
 32        else:
 33            self.skip = 1.0
 34
 35    def forward(self, z, u):
 36        correction = self.eta * self.net(torch.cat([z, u], dim=-1))
 37        return self.skip * z + correction
 38
 39def jacobian_z(model, z, u):
 40    z = z.detach().requires_grad_(True)
 41    y = model(z[None], u[None])[0]
 42    rows = []
 43    for i in range(y.numel()):
 44        rows.append(torch.autograd.grad(y[i], z, retain_graph=True)[0])
 45    return torch.stack(rows)
 46
 47def max_eig_tensor(a):
 48    return torch.linalg.eigvalsh((a+a.T)/2)[-1]
 49
 50def max_eig_symmetric(a):
 51    return float(max_eig_tensor(a).detach().cpu())
 52
 53def train(model, constrained, epochs=350):
 54    # Same tiny regression task for both blocks. It intentionally includes a
 55    # mildly expansive target, making the robustness/accuracy tradeoff visible.
 56    g = torch.Generator(device=device).manual_seed(SEED)
 57    z = (torch.rand(96, 2, generator=g, device=device)*2-1)
 58    u = (torch.rand(96, 2, generator=g, device=device)*2-1)
 59    target = 1.08*z + 0.25*u + 0.04*torch.sin(2*z)
 60    opt = torch.optim.Adam(model.parameters(), lr=3e-3)
 61    for epoch in range(epochs):
 62        opt.zero_grad()
 63        pred = model(z, u)
 64        task = ((pred-target)**2).mean()
 65        # Exact per-sample Jacobian penalty, corresponding to
 66        # M^T P M-P <= -alpha P with P=I.
 67        viol = []
 68        for k in range(0, len(z), 8):
 69            M = jacobian_z(model, z[k], u[k])
 70            q = max_eig_tensor(M.T @ M - (1-model.alpha)*torch.eye(2, device=device))
 71            viol.append(torch.nn.functional.softplus(q + 0.02)**2)
 72        penalty = torch.stack(viol).mean()
 73        loss = task + (2.0*penalty if constrained else 0.0)
 74        loss.backward()
 75        opt.step()
 76    return float(task.detach().cpu()), float(penalty.detach().cpu())
 77
 78def evaluate(model, alpha):
 79    g = torch.Generator(device=device).manual_seed(SEED+9)
 80    points = torch.rand(64, 2, generator=g, device=device)*2-1
 81    inputs = torch.rand(64, 2, generator=g, device=device)*2-1
 82    eigs=[]
 83    for z,u in zip(points, inputs):
 84        M=jacobian_z(model,z,u)
 85        eigs.append(max_eig_symmetric(M.T@M-(1-alpha)*torch.eye(2,device=device)))
 86    # Identical-input trajectories from paired hidden states.
 87    z1 = torch.tensor([0.65,-0.45], device=device)
 88    delta = torch.tensor([1e-3,-1.2e-3], device=device)
 89    z2 = z1 + delta
 90    u = torch.tensor([0.2,-0.1], device=device)
 91    ratios=[]
 92    for _ in range(20):
 93        old=float(torch.linalg.vector_norm(z2-z1).detach().cpu())
 94        z1,z2=model(z1[None],u[None])[0],model(z2[None],u[None])[0]
 95        ratios.append(float((torch.linalg.vector_norm(z2-z1)/old).detach().cpu()))
 96    return {
 97        "max_dense_generalized_violation": float(max(eigs)),
 98        "mean_dense_generalized_violation": float(np.mean(eigs)),
 99        "max_trajectory_ratio": float(max(ratios)),
100        "mean_trajectory_ratio": float(np.mean(ratios)),
101        "final_distance_ratio": float(np.prod(ratios)),
102        "ratios": ratios,
103    }
104
105def direct_math_check():
106    # M=0.8I exactly satisfies M^T M-I <= -0.36I (alpha=.30), while
107    # M=1.1I violates it and distances grow by 1.1 each step.
108    alpha=.30
109    good=torch.eye(2)*0.8
110    bad=torch.eye(2)*1.1
111    good_margin=max_eig_symmetric(good.T@good-(1-alpha)*torch.eye(2))
112    bad_margin=max_eig_symmetric(bad.T@bad-(1-alpha)*torch.eye(2))
113    dgood=1.0; dbad=1.0
114    for _ in range(12): dgood*=.8; dbad*=1.1
115    return {"alpha":alpha,"contractive_matrix_largest_eigenvalue":good_margin,
116            "expansive_matrix_largest_eigenvalue":bad_margin,
117            "contractive_distance_after_12":dgood,
118            "expansive_distance_after_12":dbad}
119
120def main():
121    global device
122    results={"seed":SEED,"device":device,"direct_math_check":direct_math_check()}
123    for name, constrained in [("baseline",False),("idea",True)]:
124        try:
125            model=ResidualBlock(constrained=constrained).to(device)
126            task, penalty=train(model, constrained)
127            results[name]={"train_task_mse":task,"train_penalty":penalty,**evaluate(model, .20)}
128        except Exception as e:
129            if device == "cuda":
130                device="cpu"; torch.set_default_device("cpu")
131                model=ResidualBlock(constrained=constrained)
132                task, penalty=train(model, constrained)
133                results[name]={"train_task_mse":task,"train_penalty":penalty,**evaluate(model,.20)}
134            else: raise
135    with open("results.json","w") as f: json.dump(results,f,indent=2)
136    print(json.dumps(results,indent=2))
137
138if __name__ == "__main__": main()