1"""Small endpoint-vs-envelope optimizer comparison.
  2
  3Uses a two-layer MLP on a fixed synthetic regression batch. Both methods form
  4exactly the same gradient proposal; the baseline checks constraints only at
  5the endpoint, while the envelope method scans interpolation points and halves
  6until all sampled points are feasible.
  7"""
  8import json
  9import random
 10from pathlib import Path
 11import numpy as np
 12import torch
 13
 14SEED = 584
 15random.seed(SEED)
 16np.random.seed(SEED)
 17torch.manual_seed(SEED)
 18try:
 19    if torch.cuda.is_available():
 20        torch.cuda.manual_seed_all(SEED)
 21    DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 22except Exception:
 23    DEVICE = torch.device("cpu")
 24
 25
 26def make_data():
 27    g = torch.Generator(device="cpu").manual_seed(SEED)
 28    x = torch.randn(96, 2, generator=g)
 29    y = (torch.sin(2.0 * x[:, :1]) + 0.5 * x[:, 1:2] ** 2).float()
 30    return x.to(DEVICE), y.to(DEVICE)
 31
 32
 33class MLP(torch.nn.Module):
 34    def __init__(self):
 35        super().__init__()
 36        self.l1 = torch.nn.Linear(2, 16)
 37        self.l2 = torch.nn.Linear(16, 1)
 38
 39    def forward(self, x):
 40        h = torch.tanh(self.l1(x))
 41        return self.l2(h), h
 42
 43
 44def flat_params(model):
 45    return torch.cat([p.detach().reshape(-1) for p in model.parameters()])
 46
 47
 48def set_flat(model, flat):
 49    pos = 0
 50    with torch.no_grad():
 51        for p in model.parameters():
 52            n = p.numel()
 53            p.copy_(flat[pos:pos+n].reshape_as(p))
 54            pos += n
 55
 56
 57def evaluate(model, x, y, flat=None):
 58    old = flat_params(model)
 59    if flat is not None:
 60        set_flat(model, flat)
 61    with torch.no_grad():
 62        pred, h = model(x)
 63        loss = torch.mean((pred - y) ** 2)
 64        rms = torch.sqrt(torch.mean(h * h) + 1e-12)
 65    set_flat(model, old)
 66    return float(loss.cpu()), float(rms.cpu())
 67
 68
 69def constraints(model, x, y, flat, loss_limit, rms_limit):
 70    loss, rms = evaluate(model, x, y, flat)
 71    return max(loss - loss_limit, rms - rms_limit), loss, rms
 72
 73
 74def train(mode, lr, steps=45, K=8):
 75    torch.manual_seed(SEED)
 76    model = MLP().to(DEVICE)
 77    x, y = make_data()
 78    initial = flat_params(model)
 79    base_loss, base_rms = evaluate(model, x, y, initial)
 80    # Small safety margins make this a loss-growth/activation envelope,
 81    # rather than requiring the current point to be exactly on a boundary.
 82    loss_margin, rms_margin = 0.015, 0.10
 83    accepted, violations, losses, step_sizes = 0, [], [], []
 84    for _ in range(steps):
 85        model.zero_grad(set_to_none=True)
 86        pred, _ = model(x)
 87        loss = torch.mean((pred - y) ** 2)
 88        loss.backward()
 89        grad = torch.cat([p.grad.detach().reshape(-1) for p in model.parameters()])
 90        current = flat_params(model)
 91        proposal = current - lr * grad
 92        cur_loss, cur_rms = evaluate(model, x, y, current)
 93        loss_limit = cur_loss + loss_margin
 94        rms_limit = max(base_rms + rms_margin, cur_rms + rms_margin)
 95
 96        def scan(scale):
 97            vals = []
 98            for s in torch.linspace(0, 1, K + 1, device=DEVICE):
 99                v, pl, pr = constraints(model, x, y, current + s * scale * (proposal-current), loss_limit, rms_limit)
100                vals.append((v, pl, pr))
101            return max(z[0] for z in vals), max(z[1] for z in vals), max(z[2] for z in vals)
102
103        endpoint_v, endpoint_peak_loss, endpoint_peak_rms = scan(1.0) if mode == "envelope" else constraints(model, x, y, proposal, loss_limit, rms_limit)
104        scale = 1.0
105        if mode == "endpoint":
106            accepted_now = endpoint_v <= 1e-8
107            observed_v, peak_loss, peak_rms = endpoint_v, endpoint_peak_loss if isinstance(endpoint_peak_loss, float) else endpoint_v, endpoint_peak_rms if isinstance(endpoint_peak_rms, float) else 0.0
108        else:
109            observed_v, peak_loss, peak_rms = scan(scale)
110            while observed_v > 1e-8 and scale > 2.0 ** -12:
111                scale *= 0.5
112                observed_v, peak_loss, peak_rms = scan(scale)
113            accepted_now = observed_v <= 1e-8
114        if accepted_now:
115            set_flat(model, current + scale * (proposal-current))
116            accepted += 1
117        step_sizes.append(scale if accepted_now else 0.0)
118        violations.append(float(observed_v))
119        losses.append(evaluate(model, x, y)[0])
120    return {
121        "mode": mode, "lr": lr, "K": K, "final_loss": losses[-1],
122        "best_loss": min(losses), "accepted_steps": accepted,
123        "mean_accepted_fraction": float(np.mean([z for z in step_sizes if z > 0])) if accepted else 0.0,
124        "max_recorded_constraint_violation": max(violations),
125        "positive_violation_steps": sum(v > 1e-7 for v in violations),
126        "loss_curve": losses,
127    }
128
129
130def main():
131    results = []
132    for lr in (0.25, 0.5, 1.0, 2.0):
133        results += [train("endpoint", lr), train("envelope", lr, K=8), train("envelope", lr, K=16)]
134    out = {"seed": SEED, "device": str(DEVICE), "results": results}
135    Path("mlp_comparison_results.json").write_text(json.dumps(out, indent=2))
136    print(json.dumps(out, indent=2))
137
138if __name__ == "__main__":
139    main()