1import sys, json, random
 2from pathlib import Path
 3import numpy as np
 4import torch
 5from torch import nn
 6sys.path.insert(0, '/home/maxwelhelp/all/math2nn')
 7from bench import get_dataset, make_model, train_model, evaluate, sweep_baseline, make_report, count_params
 8
 9SEEDS = tuple(range(8))
10# Union is shared: every idea learning rate is evaluated for baseline too.
11GRID = [{'lr': 0.0015, 'weight_decay': 0.0},
12        {'lr': 0.003,  'weight_decay': 0.0},
13        {'lr': 0.006,  'weight_decay': 0.0}]
14EPOCHS = 20
15
16class ExactChainedRNN(nn.Module):
17    """Two known phases; phase-2 recurrent state is differentiably initialized by phase 1."""
18    def __init__(self, out_dim=1, hidden=32):
19        super().__init__()
20        self.phase1 = nn.GRU(3, hidden, batch_first=True)
21        self.phase2 = nn.GRU(3, hidden, batch_first=True)
22        self.head = nn.Linear(hidden, out_dim)
23
24    def forward(self, x):
25        x = x.to(next(self.parameters()).device)
26        q = x.view(x.shape[0], -1, 3)
27        _, h1 = self.phase1(q[:, :4])
28        _, h2 = self.phase2(q[:, 4:], h1)
29        return self.head(h2[-1])
30
31    @torch.no_grad()
32    def diagnostics(self, x):
33        # CPU diagnostics avoid consuming the shared GPU/cuDNN workspace.
34        self.cpu(); x = x.cpu()
35        q = x.view(x.shape[0], -1, 3)
36        _, h1 = self.phase1(q[:, :4])
37        _, h2 = self.phase2(q[:, 4:], h1)
38        pred = self.head(h2[-1]).squeeze(-1)
39        jump = (q[:, 4, 2] - q[:, 3, 2]).abs()
40        return pred, jump, h1, h2
41
42def seed_all(seed):
43    random.seed(seed); np.random.seed(seed); torch.manual_seed(seed)
44    if torch.cuda.is_available(): torch.cuda.manual_seed_all(seed)
45
46def train_one(kind, seed, cfg, keep=False):
47    seed_all(seed)
48    ds = get_dataset('dynamics', seed, n_train=4000, n_test=1000)
49    model = make_model('rnn_small', ds['input_shape'], ds['out_dim']) if kind == 'base' else ExactChainedRNN(ds['out_dim'])
50    net, metric, hist = train_model(model, ds, epochs=EPOCHS, lr=cfg['lr'], batch=128,
51                                    weight_decay=cfg['weight_decay'], log=lambda *_: None)
52    if net is None: return float('nan'), None
53    return (float(metric), (net, ds)) if keep else (float(metric), None)
54
55def fn(kind, cfg):
56    return lambda seed: train_one(kind, seed, cfg)[0]
57
58def behavioral_signature(cfg):
59    idea_rows, base_rows = [], []
60    idea_param = base_param = None
61    for s in SEEDS:
62        _, ipack = train_one('idea', s, cfg, keep=True)
63        _, bpack = train_one('base', s, cfg, keep=True)
64        inet, ds = ipack; bnet, _ = bpack
65        ip, jump, _, _ = inet.diagnostics(ds['xte'])
66        bnet.cpu(); bp = bnet(ds['xte']).squeeze(-1).detach().cpu()
67        obs = ds['yte'].squeeze(-1).cpu()
68        high = jump >= torch.quantile(jump, 0.75)
69        low = ~high
70        idea_rows.append({'seed':s, 'high_jump_mse':float(((ip[high]-obs[high])**2).mean()), 'low_jump_mse':float(((ip[low]-obs[low])**2).mean())})
71        base_rows.append({'seed':s, 'high_jump_mse':float(((bp[high]-obs[high])**2).mean()), 'low_jump_mse':float(((bp[low]-obs[low])**2).mean())})
72        idea_param, base_param = count_params(inet), count_params(bnet)
73    ih = float(np.mean([r['high_jump_mse'] for r in idea_rows])); bh = float(np.mean([r['high_jump_mse'] for r in base_rows]))
74    return {'idea_by_seed':idea_rows, 'baseline_by_seed':base_rows,
75            'idea_high_jump_mse':ih, 'baseline_high_jump_mse':bh,
76            'idea_low_jump_mse':float(np.mean([r['low_jump_mse'] for r in idea_rows])),
77            'baseline_low_jump_mse':float(np.mean([r['low_jump_mse'] for r in base_rows])),
78            'high_jump_delta_idea_minus_baseline':ih-bh,
79            'idea_params':int(idea_param), 'baseline_params':int(base_param),
80            'confirmed': bool(ih < bh),
81            'prediction':'phase chaining should reduce error on abrupt control changes'}
82
83def main():
84    base = sweep_baseline(lambda cfg: fn('base', cfg), GRID, seeds=(0,1,2,3))
85    base['full'] = evaluate(fn('base', base['best_cfg']), SEEDS)
86    idea_runs = [{'cfg':cfg, 'result':evaluate(fn('idea', cfg), SEEDS)} for cfg in GRID]
87    chosen = min(idea_runs, key=lambda z:z['result']['mean'])
88    idea = chosen['result']
89    report = make_report('dynamics', 'rnn_small', base, idea,
90        {'phase_chaining': behavioral_signature(chosen['cfg']),
91         'idea_cfg':chosen['cfg'], 'idea_grid':idea_runs,
92         'structural_match':'controlled pendulum rollout with an observed midpoint control event; primary metric is test MSE'})
93    report['comparison']['selected_idea_cfg'] = chosen['cfg']
94    Path('bench_report.json').write_text(json.dumps(report, indent=2))
95    print(json.dumps(report, indent=2))
96
97if __name__ == '__main__':
98    main()