1import json, math, random, os
  2from pathlib import Path
  3import numpy as np
  4import torch
  5from sklearn.datasets import load_digits
  6from sklearn.model_selection import train_test_split
  7
  8SEED = int(os.environ.get('SEED', '3126'))
  9random.seed(SEED); np.random.seed(SEED); torch.manual_seed(SEED)
 10try:
 11    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
 12    if device.type == 'cuda': torch.zeros(1, device=device)
 13except Exception:
 14    device = torch.device('cpu')
 15
 16
 17def polynomial_check():
 18    out = {}
 19    E0, c, dt, n = 2.0, 0.25, 1e-3, 10000
 20    for m in (0, 1, 2):
 21        E = E0; hist = [E]; chats = []
 22        for _ in range(n):
 23            old = E
 24            E = max(0.0, E - dt*c*E**(1+m))
 25            chats.append((old-E)/(dt*old**(1+m)))
 26            hist.append(E)
 27        t = np.arange(n+1)*dt
 28        exact = E0*np.exp(-c*t) if m == 0 else (E0**(-m)+m*c*t)**(-1/m)
 29        ix = np.arange(n//2+1, n+1)
 30        slope = np.polyfit(np.log(t[ix]), np.log(np.maximum(np.asarray(hist)[ix],1e-30)), 1)[0]
 31        out[str(m)] = {
 32            'relative_final_error': float(abs(E-exact[-1])/exact[-1]),
 33            'chat_mean': float(np.mean(chats)), 'chat_std': float(np.std(chats)),
 34            'loglog_slope': float(slope), 'expected_slope': None if m == 0 else -1/m}
 35    return out
 36
 37
 38class Net(torch.nn.Module):
 39    def __init__(self):
 40        super().__init__()
 41        self.seq = torch.nn.Sequential(torch.nn.Linear(64,32), torch.nn.Tanh(), torch.nn.Linear(32,10))
 42    def forward(self, x): return self.seq(x)
 43
 44
 45def make_data():
 46    x, y = load_digits(return_X_y=True)
 47    x = ((x / 16.0) - 0.5).astype('float32')
 48    xa, xb, ya, yb = train_test_split(x, y, test_size=0.25, random_state=SEED, stratify=y)
 49    return torch.tensor(xa), torch.tensor(ya), torch.tensor(xb), torch.tensor(yb)
 50
 51
 52def train(kind, m=1, steps=500, batch=64):
 53    torch.manual_seed(SEED)
 54    x, y, xv, yv = make_data()
 55    model = Net().to(device)
 56    lossfn = torch.nn.CrossEntropyLoss()
 57    p = list(model.parameters())
 58    eta = 0.12 if kind == 'baseline' else 0.12
 59    eta_min, eta_max = 1e-4, 0.5
 60    cstar, alpha, smooth = {0: 0.003, 1: 0.001, 2: 0.0005}[m], 0.20, 0.85
 61    floor = 1e-5
 62    rng = np.random.default_rng(SEED)
 63    # Fixed monitoring batch makes E comparisons meaningful and low-noise.
 64    mi = torch.arange(min(128, len(xv)))
 65    xm, ym = xv[mi].to(device), yv[mi].to(device)
 66    rejects = 0; accepted = 0; cema = None; losses = []
 67    for step in range(steps):
 68        ii = torch.tensor(rng.integers(0, len(x), size=batch))
 69        xb, yb = x[ii].to(device), y[ii].to(device)
 70        model.zero_grad(set_to_none=True)
 71        loss = lossfn(model(xb), yb); loss.backward()
 72        old_state = {k:v.detach().clone() for k,v in model.state_dict().items()}
 73        with torch.no_grad():
 74            oldE = float(lossfn(model(xm), ym).item())
 75            for q in p: q.add_(q.grad, alpha=-eta)
 76            trialE = float(lossfn(model(xm), ym).item())
 77        if kind == 'baseline':
 78            accepted += 1
 79        elif trialE > oldE:
 80            rejects += 1
 81            model.load_state_dict(old_state)
 82            eta = max(eta_min, eta * 0.5)
 83        else:
 84            accepted += 1
 85            E = max(oldE - floor, floor)
 86            chat = (oldE - trialE) / max(E**(1+m), 1e-12)
 87            cema = chat if cema is None else smooth*cema + (1-smooth)*chat
 88            ratio = max(cema, 1e-12)/cstar
 89            eta = float(np.clip(eta * ratio**alpha, eta_min, eta_max))
 90        with torch.no_grad(): losses.append(float(lossfn(model(xm), ym).item()))
 91    with torch.no_grad(): test_loss=float(lossfn(model.to(device)(xv.to(device)), yv.to(device)).item())
 92    return {'final_monitor_loss': losses[-1], 'test_loss': test_loss, 'rejects': rejects,
 93            'acceptance_rate': accepted/steps, 'final_eta': eta,
 94            'loss_first': losses[0], 'loss_min': min(losses)}
 95
 96
 97def main():
 98    result = {'device': str(device), 'math_check': polynomial_check(), 'training': {}}
 99    for name, m in [('baseline', 0), ('m0', 0), ('m1', 1), ('m2', 2)]:
100        result['training'][name] = train('baseline' if name == 'baseline' else 'controller', m=m)
101    Path('results.json').write_text(json.dumps(result, indent=2))
102    print(json.dumps(result, indent=2))
103
104if __name__ == '__main__': main()