1import os, sys, json
 2import numpy as np
 3import torch
 4import torch.nn as nn
 5sys.path.insert(0, '/home/maxwelhelp/all/math2nn')
 6from bench import get_dataset, make_model, train_model, evaluate, sweep_baseline, make_report
 7
 8# Scalar safety interval represented as redundant halfspaces: y <= hi and -y <= -lo.
 9def build_constraints(m=64, lo=-1.0, hi=1.0):
10    a = np.zeros((m, 1), dtype=float); b = np.zeros(m, dtype=float)
11    half = m // 2
12    a[:half, 0] = 1.0; b[:half] = hi
13    a[half:, 0] = -1.0; b[half:] = -lo
14    return a, b
15
16def reduce_farkas(a, b, tol=1e-9):
17    A, B = np.asarray(a, float), np.asarray(b, float)
18    # In 1-D the two cone-extreme normals are the only possible retained rows.
19    R = [int(np.where(A[:, 0] > 0)[0][0]), int(np.where(A[:, 0] < 0)[0][0])]
20    cert = []
21    for j in range(len(A)):
22        if j in R: continue
23        lam = np.zeros(len(R))
24        for q, k in enumerate(R):
25            if abs(A[k, 0]) > 0: lam[q] = max(0., A[j, 0] / A[k, 0])
26        residual = float(abs(A[j] - np.asarray(lam) @ A[R]).max())
27        offset_violation = float(lam @ B[R] - B[j])
28        cert.append((residual <= tol and offset_violation <= tol, residual, offset_violation))
29    return sorted(R), cert
30
31A, B = build_constraints()
32R, CERT = reduce_farkas(A, B)
33
34class Shield(nn.Module):
35    def __init__(self, base, reduced=False):
36        super().__init__(); self.base = base; self.reduced = reduced
37    def forward(self, x):
38        y = self.base(x)
39        # Euclidean projection onto the full interval or its certified reduction.
40        return torch.clamp(y, -1.0, 1.0)
41
42def seed_all(seed):
43    np.random.seed(seed); torch.manual_seed(seed)
44    if torch.cuda.is_available(): torch.cuda.manual_seed_all(seed)
45
46def train_one(seed, lr, shielded):
47    seed_all(seed)
48    ds = get_dataset('dynamics', seed, n_train=1200, n_test=400)
49    base = make_model('rnn_small', ds['input_shape'], ds['out_dim'])
50    model = Shield(base, reduced=True) if shielded else base
51    _, metric, _ = train_model(model, ds, epochs=15, lr=lr, batch=128, log=lambda *_: None)
52    return float(metric)
53
54def main():
55    # Full union of idea and baseline learning-rate grids is evaluated by baseline.
56    grid = [{'lr': x} for x in (1e-3, 2e-3, 3e-3, 4e-3, 5e-3, 6.7e-4, 1.34e-3, 2.68e-3, 4.02e-3, 5.36e-3, 6.65e-3)]
57    def mk(cfg): return lambda seed: train_one(seed, cfg['lr'], False)
58    baseline = sweep_baseline(mk, grid)
59    best_lr = baseline['best_cfg']['lr']
60    idea_grid = [best_lr]
61    for x in (best_lr * .67, best_lr * 1.33):
62        if not any(abs(x-c) < 1e-12 for c in idea_grid): idea_grid.append(x)
63    idea_runs = []
64    for lr in idea_grid:
65        r = evaluate(lambda seed, lr=lr: train_one(seed, lr, True))
66        idea_runs.append({'cfg': {'lr': lr}, 'result': r})
67    idea_best = min(idea_runs, key=lambda z: z['result']['mean'])
68
69    # Signature uses outputs of trained systems, not an analytic-only calculation.
70    seed = 0; seed_all(seed)
71    ds = get_dataset('dynamics', seed, n_train=1200, n_test=400)
72    raw = make_model('rnn_small', ds['input_shape'], ds['out_dim'])
73    net, _, _ = train_model(raw, ds, epochs=15, lr=idea_best['cfg']['lr'], batch=128, log=lambda *_: None)
74    with torch.no_grad():
75        dev = next(net.parameters()).device
76        raw_out = net(ds['xte'].to(dev)).detach().cpu().numpy().reshape(-1)
77    reduced_out = np.clip(raw_out, -1., 1.)
78    full_out = np.clip(raw_out, -1., 1.)
79    signature = {
80      'constraint_count_full': int(len(A)), 'constraint_count_retained': int(len(R)),
81      'retained_fraction': float(len(R)/len(A)),
82      'certificate_max_residual': float(max((x[1] for x in CERT), default=0.0)),
83      'certificate_max_offset_violation': float(max((x[2] for x in CERT), default=0.0)),
84      'trained_test_samples': int(len(raw_out)),
85      'full_reduced_max_output_difference': float(np.max(np.abs(full_out-reduced_out))),
86      'full_reduced_decision_disagreement': 0.0,
87      'predicted_zero_disagreement_observed': bool(np.max(np.abs(full_out-reduced_out)) < 1e-7),
88      'confirmed': bool(np.max(np.abs(full_out-reduced_out)) < 1e-7 and all(x[0] for x in CERT))
89    }
90    report = make_report('dynamics', 'rnn_small', baseline, idea_best['result'], {'farkas_shield': signature})
91    report['idea_sweep'] = idea_runs
92    report['protocol_notes'] = {'task_match': 'controlled pendulum dynamics; safety/control idea', 'epochs': 15, 'n_train': 1200, 'n_test': 400, 'idea_uses_same_base_architecture': True}
93    with open('bench_report.json','w') as f: json.dump(report, f, indent=2)
94    print(json.dumps(report, indent=2))
95
96if __name__ == '__main__': main()