1import json, random
  2from pathlib import Path
  3import numpy as np
  4import torch
  5from torch import nn
  6
  7SEED = 3145
  8random.seed(SEED); np.random.seed(SEED); torch.manual_seed(SEED)
  9try:
 10    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
 11except Exception:
 12    device = torch.device('cpu')
 13
 14# C4 acts on both 2D input blocks (state and action) and on the 2D output.
 15def R(k, device=None):
 16    t = float(k % 4) * np.pi / 2
 17    c, s = np.cos(t), np.sin(t)
 18    return torch.tensor([[c, -s], [s, c]], dtype=torch.float32, device=device)
 19
 20def transform(z, k):
 21    q = R(k, z.device)
 22    return torch.cat([z[:, :2] @ q.T, z[:, 2:] @ q.T], dim=1)
 23
 24def transform_y(y, k):
 25    return y @ R(k, y.device).T
 26
 27def true_base(z):
 28    x, a = z[:, :2], z[:, 2:]
 29    # Nonlinear but exactly O(2)-equivariant: scalar invariants multiply vectors.
 30    s = (x*x).sum(1, keepdim=True) + 0.35*(a*a).sum(1, keepdim=True)
 31    return (0.65 + 0.12*torch.tanh(s))*x + (0.42 - 0.08*torch.sigmoid(s))*a
 32
 33def make_data(n, ks):
 34    z0 = torch.randn(n, 4) * torch.tensor([2.0, 0.35, 1.5, 0.25])
 35    # Every orientation is the same embedded local mechanism under C4.
 36    kvals = torch.tensor(np.random.choice(ks, n), dtype=torch.long)
 37    zs, ys = [], []
 38    for z, k in zip(z0, kvals):
 39        zz = transform(z[None], int(k)).squeeze(0)
 40        yy = transform_y(true_base(z[None]), int(k)).squeeze(0)
 41        zs.append(zz); ys.append(yy)
 42    return torch.stack(zs), torch.stack(ys)
 43
 44class MLP(nn.Module):
 45    def __init__(self, hidden=48):
 46        super().__init__()
 47        self.net = nn.Sequential(nn.Linear(4, hidden), nn.Tanh(), nn.Linear(hidden, hidden), nn.Tanh(), nn.Linear(hidden, 2))
 48    def forward(self, z): return self.net(z)
 49
 50def eq_penalty(model, z):
 51    # Exact sampled analogue of E_h ||F(rho_in(h)z)-rho_out(h)F(z)||^2.
 52    y = model(z)
 53    vals = []
 54    for k in range(4):
 55        vals.append((model(transform(z, k)) - transform_y(y, k)).pow(2).mean())
 56    return torch.stack(vals).mean()
 57
 58def fit(lam, train_z, train_y, epochs=450):
 59    model = MLP().to(device)
 60    opt = torch.optim.Adam(model.parameters(), lr=3e-3, weight_decay=1e-5)
 61    z, y = train_z.to(device), train_y.to(device)
 62    for _ in range(epochs):
 63        opt.zero_grad()
 64        pred = model(z)
 65        loss = (pred-y).pow(2).mean() + lam*eq_penalty(model, z)
 66        loss.backward(); opt.step()
 67    return model
 68
 69def mse(model, z, y):
 70    with torch.no_grad(): return float((model(z.to(device))-y.to(device)).pow(2).mean().cpu())
 71
 72def main():
 73    # Structural sanity checks corresponding to the occurrence symmetry: C4 composition,
 74    # invertibility, and exact equivariance of the generated shared mechanism.
 75    z = torch.randn(64, 4)
 76    composition_err = max(float((transform(transform(z, a), b)-transform(z, (a+b)%4)).abs().max()) for a in range(4) for b in range(4))
 77    mechanism_err = max(float((true_base(transform(z, k))-transform_y(true_base(z), k)).abs().max()) for k in range(4))
 78    linear_good = nn.Linear(4, 2, bias=False)
 79    with torch.no_grad():
 80        linear_good.weight.copy_(torch.tensor([[1.,0.,0.5,0.],[0.,1.,0.,0.5]]))
 81    penalty_good = float(eq_penalty(linear_good, z))
 82    bad = nn.Linear(4, 2)
 83    with torch.no_grad(): bad.weight.copy_(torch.tensor([[1.7,0.2,-.4,.3],[.1,.3,.8,-.2]])); bad.bias.fill_(.4)
 84    penalty_bad = float(eq_penalty(bad, z))
 85
 86    # Train on three orientations and test on a held-out fourth occurrence symmetry.
 87    train_z, train_y = make_data(768, [0,1,2])
 88    test_id_z, test_id_y = make_data(384, [0,1,2])
 89    test_ood_z, test_ood_y = make_data(384, [3])
 90    baseline = fit(0.0, train_z, train_y)
 91    equiv = fit(1.0, train_z, train_y)
 92    results = {
 93        'seed': SEED, 'device': str(device),
 94        'math_check': {'c4_composition_max_abs': composition_err,
 95                       'true_mechanism_equivariance_max_abs': mechanism_err,
 96                       'equivariance_penalty_exact_linear': penalty_good,
 97                       'equivariance_penalty_non_equivariant_linear': penalty_bad},
 98        'params_each_model': sum(p.numel() for p in baseline.parameters()),
 99        'baseline': {'train_mse': mse(baseline, train_z, train_y), 'in_distribution_mse': mse(baseline, test_id_z, test_id_y), 'heldout_rotation_mse': mse(baseline, test_ood_z, test_ood_y), 'heldout_equivariance_penalty': float(eq_penalty(baseline, test_ood_z.to(device)).cpu())},
100        'equivariant_shared': {'train_mse': mse(equiv, train_z, train_y), 'in_distribution_mse': mse(equiv, test_id_z, test_id_y), 'heldout_rotation_mse': mse(equiv, test_ood_z, test_ood_y), 'heldout_equivariance_penalty': float(eq_penalty(equiv, test_ood_z.to(device)).cpu())}
101    }
102    Path('results.json').write_text(json.dumps(results, indent=2))
103    print(json.dumps(results, indent=2))
104
105if __name__ == '__main__':
106    try:
107        main()
108    except Exception as e:
109        if str(device) == 'cuda':
110            print('CUDA failed; rerun on CPU:', repr(e))
111            device = torch.device('cpu'); main()
112        else: raise