1"""Finite-horizon walk reciprocity control: verification and tiny training study."""
  2import json, math, random
  3from pathlib import Path
  4import numpy as np
  5
  6
  7def phi_star(g):
  8    return 1.0 - math.sqrt(max(0.0, 1.0 - g * g))
  9
 10
 11def walk_energy(A, L=4, g=0.8, normalize=True, eps=1e-8):
 12    A = np.asarray(A, dtype=np.float64)
 13    if normalize:
 14        s = np.linalg.svd(A, compute_uv=False)[0]
 15        A = A / (s + eps)
 16    n = A.shape[0]
 17    V = A.copy()
 18    terms = []
 19    total = 0.0
 20    for k in range(1, L + 1):
 21        D = V - V.T
 22        term = (g ** (2 * k - 2)) * np.sum(D * D) / n
 23        terms.append(float(term))
 24        total += term
 25        V = A @ V
 26    return float(total), terms
 27
 28
 29def spectral_radius(A):
 30    return float(np.max(np.abs(np.linalg.eigvals(A))))
 31
 32
 33def math_sanity():
 34    n = 8
 35    # Strictly lower-triangular directed chain: nilpotent but transient walks exist.
 36    A = np.zeros((n, n))
 37    for i in range(n - 1):
 38        A[i + 1, i] = 1.0
 39    S = A + A.T                         # reciprocal control
 40    rn, tn = walk_energy(A, L=6, g=.8)
 41    rs, ts = walk_energy(S, L=6, g=.8)
 42    # Verify the iterative powers used by the implementation against A**k.
 43    An = A / (np.linalg.svd(A, compute_uv=False)[0] + 1e-8)
 44    V = An.copy()
 45    recursion_errors = []
 46    for k in range(1, 7):
 47        recursion_errors.append(float(np.max(np.abs(V - np.linalg.matrix_power(An, k)))))
 48        V = An @ V
 49    return {
 50        "nilpotent_spectral_radius": spectral_radius(A),
 51        "nilpotent_A_power_n_frobenius": float(np.linalg.norm(np.linalg.matrix_power(A, n))),
 52        "nilpotent_walk_energy": rn,
 53        "nilpotent_terms": tn,
 54        "symmetric_control_walk_energy": rs,
 55        "symmetric_control_terms": ts,
 56        "phi_star_0.8": phi_star(.8),
 57        "max_iterative_power_error": max(recursion_errors),
 58        "energy_positive_before_nilpotency": rn > 0,
 59    }
 60
 61
 62def train_compare(seed=7, steps=500, lam=0.0, target=0.0):
 63    import torch
 64    torch.manual_seed(seed); np.random.seed(seed); random.seed(seed)
 65    device = "cuda" if torch.cuda.is_available() else "cpu"
 66    try:
 67        x = torch.randn(96, 12, device=device)
 68        true_w = torch.randn(12, 12, device=device) / math.sqrt(12)
 69        y = torch.tanh(x @ true_w.T)
 70        model = torch.nn.Sequential(torch.nn.Linear(12, 12), torch.nn.Tanh(), torch.nn.Linear(12, 12)).to(device)
 71        opt = torch.optim.Adam(model.parameters(), lr=3e-3)
 72        losses=[]; energies=[]; gradnorms=[]
 73        g=.8; L=4
 74        for step in range(steps):
 75            opt.zero_grad(set_to_none=True)
 76            pred=model(x)
 77            loss=torch.mean((pred-y)**2)
 78            W=model[0].weight
 79            scale=torch.linalg.matrix_norm(W, ord=2) + 1e-6
 80            An=W/scale
 81            V=An
 82            R=torch.zeros((), device=device)
 83            for k in range(1,L+1):
 84                D=V-V.T
 85                R=R+(g**(2*k-2))*torch.sum(D*D)/W.shape[0]
 86                V=An@V
 87            total=loss+lam*torch.relu(R-target)**2
 88            total.backward()
 89            gn=0.0
 90            for p in model.parameters():
 91                if p.grad is not None: gn += float(torch.sum(p.grad.detach()**2).cpu())
 92            opt.step()
 93            if step >= steps-50:
 94                losses.append(float(loss.detach().cpu())); energies.append(float(R.detach().cpu())); gradnorms.append(math.sqrt(gn))
 95        return {"device":device, "final_loss_mean_last50":float(np.mean(losses)),
 96                "walk_energy_mean_last50":float(np.mean(energies)),
 97                "gradient_norm_mean_last50":float(np.mean(gradnorms)), "lambda":lam, "target":target}
 98    except (RuntimeError, torch.cuda.OutOfMemoryError):
 99        if device == "cuda":
100            torch.cuda.empty_cache()
101            return train_compare_cpu(seed, steps, lam, target)
102        raise
103
104
105def train_compare_cpu(seed, steps, lam, target):
106    import torch
107    old=torch.cuda.is_available
108    torch.cuda.is_available=lambda: False
109    try: return train_compare(seed, steps, lam, target)
110    finally: torch.cuda.is_available=old
111
112
113def main():
114    sanity=math_sanity()
115    # Per-seed baseline calibration follows the proposed first-epoch-style calibration.
116    rows=[]
117    for seed in (7, 19, 31):
118        base=train_compare(seed=seed, lam=0.0)
119        controlled=train_compare(seed=seed, lam=2.0, target=base["walk_energy_mean_last50"])
120        rows.append({"seed":seed, "baseline":base, "walk_regularized":controlled})
121    def avg(key, method):
122        return float(np.mean([r[method][key] for r in rows]))
123    out={"sanity":sanity, "per_seed":rows,
124         "mean_last50": {
125             "baseline_loss":avg("final_loss_mean_last50", "baseline"),
126             "idea_loss":avg("final_loss_mean_last50", "walk_regularized"),
127             "baseline_walk_energy":avg("walk_energy_mean_last50", "baseline"),
128             "idea_walk_energy":avg("walk_energy_mean_last50", "walk_regularized"),
129             "baseline_gradient_norm":avg("gradient_norm_mean_last50", "baseline"),
130             "idea_gradient_norm":avg("gradient_norm_mean_last50", "walk_regularized")},
131         "all_seeds_lower_energy": all(r["walk_regularized"]["walk_energy_mean_last50"] < r["baseline"]["walk_energy_mean_last50"] for r in rows),
132         "all_seeds_lower_loss": all(r["walk_regularized"]["final_loss_mean_last50"] < r["baseline"]["final_loss_mean_last50"] for r in rows)}
133    Path("results.json").write_text(json.dumps(out, indent=2))
134    print(json.dumps(out, indent=2))
135
136if __name__ == "__main__": main()