1import json, random, sys
  2import numpy as np
  3import torch
  4import torch.nn as nn
  5
  6sys.path.insert(0, '/home/maxwelhelp/all/math2nn')
  7from bench import get_dataset, train_model, evaluate, sweep_baseline, make_report
  8
  9SEEDS = tuple(range(8))
 10SWEEP_SEEDS = (0, 1, 2, 3)
 11EPOCHS = 12
 12NTR, NTE = 400, 200
 13LR_GRID = [1e-3, 3e-3, 1e-2]
 14
 15
 16def seed_all(seed):
 17    random.seed(seed)
 18    np.random.seed(seed)
 19    torch.manual_seed(seed)
 20    if torch.cuda.is_available():
 21        torch.cuda.manual_seed_all(seed)
 22
 23
 24class VariableBaseline(nn.Module):
 25    """Standard variable-only predictor: one learned state readout."""
 26    def __init__(self, width=48):
 27        super().__init__()
 28        self.encoder = nn.Sequential(nn.Linear(24, width), nn.Tanh(),
 29                                     nn.Linear(width, width), nn.Tanh())
 30        self.readout = nn.Linear(width, 1)
 31
 32    def forward(self, x):
 33        return self.readout(self.encoder(x))
 34
 35
 36class EquationAddressable(nn.Module):
 37    """Implicit bipartite layer with three mechanism residuals and two variables.
 38
 39    f0 and f1 both address variable z0 through different learned mechanisms;
 40    f2 addresses z1 and couples to z0. The forward pass minimizes all residuals
 41    simultaneously by differentiable damped residual-gradient iterations.
 42    """
 43    def __init__(self, width=48, steps=8, step_size=0.20):
 44        super().__init__()
 45        self.encoder = nn.Sequential(nn.Linear(24, width), nn.Tanh(),
 46                                     nn.Linear(width, width), nn.Tanh())
 47        self.mechanisms = nn.ModuleList([
 48            nn.Sequential(nn.Linear(width + 2, width), nn.Tanh(), nn.Linear(width, 1))
 49            for _ in range(3)
 50        ])
 51        self.steps = int(steps)
 52        self.step_size = float(step_size)
 53
 54    def residuals(self, z, h):
 55        q = torch.cat([h, z], dim=1)
 56        a = [m(q)[:, 0] for m in self.mechanisms]
 57        # Bipartite incidence: f0--z0, f1--(z0,z1), f2--z1.
 58        return torch.stack([z[:, 0] - a[0],
 59                            z[:, 0] + 0.5 * z[:, 1] - a[1],
 60                            z[:, 1] - a[2]], dim=1)
 61
 62    def solve(self, h, intervention=None, xi=None, steps=None):
 63        # Residual-gradient equilibrium updates need a graph even during eval.
 64        with torch.enable_grad():
 65            return self._solve_grad(h, intervention, xi, steps)
 66
 67    def _solve_grad(self, h, intervention=None, xi=None, steps=None):
 68        z = torch.zeros(h.shape[0], 2, device=h.device, dtype=h.dtype)
 69        nsteps = self.steps if steps is None else int(steps)
 70        for _ in range(nsteps):
 71            z.requires_grad_(True)
 72            r = self.residuals(z, h)
 73            if intervention is not None:
 74                j, v = intervention
 75                rr = r.clone()
 76                rr[:, j] = z[:, v] - xi
 77            else:
 78                rr = r
 79            loss = 0.5 * (rr * rr).sum()
 80            grad = torch.autograd.grad(loss, z, create_graph=True)[0]
 81            z = z - self.step_size * grad
 82        return z
 83
 84    def forward(self, x):
 85        h = self.encoder(x)
 86        return self.solve(h)[:, :1]
 87
 88
 89def train_one(kind, seed, lr, steps=8, step_size=0.20, return_model=False):
 90    seed_all(seed)
 91    ds = get_dataset('dynamics', seed, n_train=NTR, n_test=NTE)
 92    if kind == 'baseline':
 93        net = VariableBaseline()
 94    else:
 95        net = EquationAddressable(steps=steps, step_size=step_size)
 96    net, metric, hist = train_model(net, ds, epochs=EPOCHS, lr=lr, batch=128,
 97                                    weight_decay=0.0, log=lambda *_: None)
 98    if net is None:
 99        raise RuntimeError('benchmark training failed')
100    if return_model:
101        return float(metric), net, ds
102    return float(metric)
103
104
105def baseline_factory(cfg):
106    return lambda seed: train_one('baseline', seed, float(cfg['lr']))
107
108
109def idea_factory(cfg):
110    return lambda seed: train_one('idea', seed, float(cfg['lr']),
111                                  steps=int(cfg['steps']), step_size=float(cfg['step_size']))
112
113
114def signature(idea_model, baseline_model, ds):
115    """Re-test the stage-1 prediction on trained models, not analytic equations."""
116    device = next(idea_model.parameters()).device
117    x = ds['xte'].to(device)[:64]
118    with torch.enable_grad():
119        h = idea_model.encoder(x)
120        with torch.no_grad():
121            obs = idea_model.solve(h)
122        xi = obs[:, 0].median().detach()
123        # Two distinct equation replacements fix the same trained variable value.
124        a = idea_model.solve(h, intervention=(0, 0), xi=xi, steps=12).detach()
125        b = idea_model.solve(h, intervention=(1, 0), xi=xi, steps=12).detach()
126    with torch.no_grad():
127        base_pred = baseline_model(x)[:, 0]
128    target_a = float((a[:, 0] - xi).abs().mean())
129    target_b = float((b[:, 0] - xi).abs().mean())
130    downstream = float((a[:, 1] - b[:, 1]).abs().mean())
131    ordinary = float((obs[:, 0] - base_pred).abs().mean())
132    return {
133        'prediction': 'same xi imposed by different equation replacements yields different downstream state',
134        'n_samples': 64,
135        'xi': float(xi),
136        'trained_model_target_error_f0': target_a,
137        'trained_model_target_error_f1': target_b,
138        'trained_model_downstream_separation': downstream,
139        'baseline_vs_observational_prediction_abs_gap': ordinary,
140        'confirmed': bool(max(target_a, target_b) < 0.08 and downstream > 0.01)
141    }
142
143
144def main():
145    # Baseline and idea use the same lr union; baseline is swept on four seeds.
146    grid = [{'lr': lr} for lr in LR_GRID]
147    base = sweep_baseline(baseline_factory, grid, seeds=SWEEP_SEEDS)
148    best_lr = float(base['best_cfg']['lr'])
149    idea_grid = [
150        {'lr': best_lr, 'steps': 8, 'step_size': 0.20},
151        {'lr': LR_GRID[max(0, LR_GRID.index(best_lr)-1)], 'steps': 8, 'step_size': 0.20},
152        {'lr': LR_GRID[min(len(LR_GRID)-1, LR_GRID.index(best_lr)+1)], 'steps': 8, 'step_size': 0.20},
153    ]
154    # Deduplicate if the best is at an edge while retaining three nearby trials when possible.
155    unique = []
156    for c in idea_grid:
157        if c not in unique: unique.append(c)
158    idea_grid = unique
159    idea_runs = []
160    for cfg in idea_grid:
161        r = evaluate(idea_factory(cfg), seeds=SEEDS)
162        idea_runs.append({'cfg': cfg, 'result': r})
163    chosen = min(idea_runs, key=lambda q: q['result']['mean'])
164    # Refit one paired seed for the behavior signature using the selected setting.
165    _, im, ds = train_one('idea', 0, chosen['cfg']['lr'], chosen['cfg']['steps'], chosen['cfg']['step_size'], True)
166    _, bm, _ = train_one('baseline', 0, best_lr, return_model=True)
167    sig = signature(im, bm, ds)
168    report = make_report('dynamics', 'rnn_small', base, chosen['result'], {
169        'mechanism_signature': sig,
170        'track_justification': 'Dynamics is structurally matched: the built-in task is an actuated pendulum rollout whose target depends on coupled state-transition mechanisms.',
171        'baseline_sweep_union_lr': LR_GRID,
172        'idea_sweep': idea_runs,
173        'idea_selected_cfg': chosen['cfg'],
174        'parameter_counts': {'baseline': sum(p.numel() for p in bm.parameters()), 'idea': sum(p.numel() for p in im.parameters())}
175    })
176    report['comparison']['idea_sweep_results'] = idea_runs
177    with open('bench_report.json', 'w') as f: json.dump(report, f, indent=2)
178    print(json.dumps(report, indent=2))
179
180
181if __name__ == '__main__':
182    main()