1import json, math, random
  2from pathlib import Path
  3import numpy as np
  4import torch
  5from torch import nn
  6
  7SEED=2763
  8np.random.seed(SEED); random.seed(SEED); torch.manual_seed(SEED)
  9torch.set_num_threads(4)
 10
 11# Primitive values are signed distances for two target/constraint primitives.
 12def primitives(x):
 13    # Two smooth primitive value functions; <=0 is the desired set.
 14    z1=(x[:,0]-0.35)**2+(x[:,1]+0.15)**2-0.42
 15    z2=(x[:,0]+0.40)**2+(x[:,1]-0.20)**2-0.36
 16    return torch.stack([z1,z2],1)
 17
 18def positive_linear_fit(z,y,steps=1200):
 19    # A monotone aggregator: b + softplus(a_i) z_i, with strictly nonnegative derivatives.
 20    raw=nn.Parameter(torch.zeros(2)); b=nn.Parameter(torch.zeros(()))
 21    opt=torch.optim.Adam([raw,b],lr=.04)
 22    for _ in range(steps):
 23        opt.zero_grad(); w=torch.nn.functional.softplus(raw)
 24        loss=((b+z@w-y)**2).mean(); loss.backward(); opt.step()
 25    return b.detach(), torch.nn.functional.softplus(raw).detach()
 26
 27def unconstrained_fit(z,y,steps=1200):
 28    w=nn.Parameter(torch.zeros(2)); b=nn.Parameter(torch.zeros(()))
 29    opt=torch.optim.Adam([w,b],lr=.04)
 30    for _ in range(steps):
 31        opt.zero_grad(); loss=((b+z@w-y)**2).mean(); loss.backward(); opt.step()
 32    return b.detach(),w.detach()
 33
 34def mask_iou(a,b):
 35    inter=(a&b).sum(); union=(a|b).sum()
 36    return (inter/union).item() if union else 1.0
 37
 38def main():
 39    # ---------------- core math checks ----------------
 40    n=50000
 41    z=torch.randn(n,2)*1.3
 42    # Positive aggregator with known Lipschitz coordinate coefficients.
 43    w=torch.tensor([0.7,1.4]); b=0.2
 44    e=torch.randn(n,2)*0.03
 45    err=(b+(z+e)@w-(b+z@w)).abs()
 46    bound=(e.abs()@w)
 47    bound_ratio=(err/(bound+1e-12)).max().item()
 48    # Coordinatewise ordering preservation (finite differences, all coordinates nonnegative).
 49    delta=torch.rand(n,2)*2
 50    diff=(b+(z+delta)@w-(b+z@w))
 51    order_fraction=(diff>=-1e-10).all().item() if diff.ndim==0 else (diff>=-1e-10).float().mean().item()
 52    # In the linear case error scales exactly linearly with primitive error amplitude.
 53    scales=torch.tensor([0.25,0.5,1.,2.,4.])
 54    scaling=[]
 55    for s in scales:
 56        ee=torch.randn(n,2)*.02*s
 57        scaling.append((ee@w).abs().mean().item())
 58    slopes=np.polyfit(np.log(scales.numpy()),np.log(np.array(scaling)),1)[0]
 59
 60    # ---------------- learned critic mini experiment ----------------
 61    # Train on left half of state square and test on right half: primitive critics are fixed,
 62    # while the aggregator must preserve their task ordering under shift.
 63    N=5000
 64    x=torch.rand(N,2)*2-1
 65    z0=primitives(x)
 66    noise=torch.randn_like(z0)*0.025
 67    zhat=z0+noise
 68    # Ground truth composite is monotone and has a known zero-sublevel set.
 69    y=z0@torch.tensor([1.,0.8])+0.03
 70    train=x[:,0]<0.05; test=~train
 71    bp,wp=positive_linear_fit(zhat[train],y[train])
 72    bu,wu=unconstrained_fit(zhat[train],y[train])
 73    yp=bp+zhat@wp; yu=bu+zhat@wu
 74    rmse_p=torch.sqrt(((yp[test]-y[test])**2).mean()).item()
 75    rmse_u=torch.sqrt(((yu[test]-y[test])**2).mean()).item()
 76    iou_p=mask_iou(yp[test]<=0,y[test]<=0); iou_u=mask_iou(yu[test]<=0,y[test]<=0)
 77
 78    # Negative derivative sweep: compare c2*z2 with the true c2=.8 composite.
 79    # The ordering theorem predicts failure precisely after c2 crosses zero.
 80    cs=np.array([1.0,.5,.2,.05,0.,-.02,-.05,-.1,-.2,-.5,-1.0])
 81    sweep=[]
 82    for c in cs:
 83        pred=z0[:,0]+float(c)*z0[:,1]+.03
 84        # set disagreement and order reversals relative to increasing z2 at fixed z1
 85        set_err=(pred<=0).ne(y<=0).float().mean().item()
 86        # derivative sign is c; for c<0, every positive z2 increment reverses contribution
 87        # Direct monotonicity test: increasing z2 must not decrease the output.
 88        dz=torch.rand(N)*2.0
 89        reversal=(float(c)*dz < -1e-12).float().mean().item()
 90        sweep.append({'c2':float(c),'derivative':float(c),'set_disagreement':set_err,
 91                      'order_reversal_fraction':reversal})
 92
 93    result={
 94      'predictions':{
 95        'P1_monotone_derivatives_nonnegative':{'predicted':'100% samples','observed_fraction':order_fraction,'tolerance':'exact up to numerical precision'},
 96        'P2_error_bound':{'predicted':'|delta V| <= 0.7|e1|+1.4|e2| and linear scaling','max_ratio_observed':bound_ratio,'loglog_scaling_exponent':float(slopes),'expected_exponent':1.0},
 97        'P3_negative_crossing':{'predicted':'set/order failure begins at c2<0','sweep':sweep}
 98      },
 99      'mini_experiment':{
100        'positive_aggregator_weights':wp.tolist(),'unconstrained_weights':wu.tolist(),
101        'heldout_rmse':{'monotone':rmse_p,'unconstrained':rmse_u},
102        'heldout_set_iou':{'monotone':iou_p,'unconstrained':iou_u}
103      }
104    }
105    Path('results.json').write_text(json.dumps(result,indent=2))
106    print(json.dumps(result,indent=2))
107
108if __name__=='__main__': main()