1import json, math, random
  2from pathlib import Path
  3import numpy as np
  4
  5SEED = 924
  6rng = np.random.default_rng(SEED)
  7
  8# A binary-state finite-horizon process. At each step, next state is sampled,
  9# reward has conditional mean r[state,next] and Gaussian conditional noise.
 10def recursion(h, p, gamma, r, reward_var, terminal_mean=None, terminal_var=None):
 11    # m[t,x], v[t,x] describe return from state x with t steps remaining.
 12    m = np.zeros((h + 1, 2)); v = np.zeros((h + 1, 2))
 13    if terminal_mean is not None: m[0] = terminal_mean
 14    if terminal_var is not None: v[0] = terminal_var
 15    P = np.array([[1-p, p], [p, 1-p]], float)
 16    for t in range(1, h + 1):
 17        for x in range(2):
 18            z = r[x] + gamma * m[t-1] # reward mean depends only on current state
 19            m[t,x] = z @ P[x]
 20            # conditional reward noise plus propagated child uncertainty and selection variance
 21            v[t,x] = np.sum(P[x] * (reward_var[x] + gamma**2 * v[t-1] + (z - m[t,x])**2))
 22    return m, v
 23
 24def sample_returns(h, n, p, gamma, r, sigma, start=0):
 25    x = np.full(n, start, dtype=np.int64); out = np.zeros(n)
 26    for t in range(h):
 27        nx = np.where(rng.random(n) < p, 1-x, x)
 28        out += (gamma**t) * (r[x] + sigma[x] * rng.normal(size=n))
 29        x = nx
 30    return out
 31
 32def main():
 33    gamma=.9; r=np.array([-1., 1.]); sig=np.array([.35,.35]); rv=sig**2
 34    rows=[]
 35    # Prediction 1: recursive variance equals held-out Monte Carlo variance.
 36    for h in [1,2,4,8,12]:
 37        m,v=recursion(h,.35,gamma,r,rv)
 38        g=sample_returns(h, 300000, .35,gamma,r,sig)
 39        rows.append({'h':h,'recursive_var':float(v[h,0]),'mc_var':float(g.var()),
 40                     'relative_error':float(abs(g.var()-v[h,0])/v[h,0])})
 41    # Prediction 2: one-step transition-selection term is p(1-p)*(delta child mean)^2,
 42    # hence zero at deterministic transitions and maximal at p=.5.
 43    sel=[]
 44    child=np.array([-2., 3.]); delta=child[1]-child[0]
 45    for p in [0.,.1,.25,.5,.75,.9,1.]:
 46        P=np.array([1-p,p]); mean=P@child
 47        measured=np.sum(P*(child-mean)**2)
 48        predicted=p*(1-p)*delta**2
 49        sel.append({'p':p,'predicted':float(predicted),'measured':float(measured)})
 50    # Prediction 3: with independently redrawn state each step, variance is the
 51    # discounted sum of per-step variances; removing transition randomness removes
 52    # the state-selection contribution and leaves reward-noise variance only.
 53    horizon=[]
 54    for transition_var in [0.,1.]:
 55        # transition_var=1 means iid +/-1 state each step, 0 means fixed state +1.
 56        vals=[]
 57        for h in [1,2,4,8,16]:
 58            state_var=transition_var # Var +/-1 with equal probability
 59            predicted=sum((gamma**(2*t))*(state_var + .35**2) for t in range(h))
 60            g=np.zeros(250000)
 61            for t in range(h):
 62                states=(rng.choice([-1.,1.],size=g.size) if transition_var else np.ones(g.size))
 63                g += gamma**t*(states + .35*rng.normal(size=g.size))
 64            vals.append({'h':h,'predicted':float(predicted),'mc_var':float(g.var())})
 65        horizon.append({'transition_randomness':transition_var,'points':vals})
 66    # Small baseline-vs-idea regression: contexts have heteroscedastic known recursive
 67    # targets. MSE is standard; weighted NLL uses the detached recursive variance.
 68    try:
 69        import torch
 70        torch.manual_seed(SEED)
 71        device='cuda' if torch.cuda.is_available() else 'cpu'
 72        try:
 73            x=torch.rand(12000,1,device=device)
 74            # x controls branch probability; target is child +/-1 plus reward noise.
 75            pp=.05+.9*x[:,0]
 76            means=2*pp-1
 77            q=.25 + 4*pp*(1-pp) # reward variance + transition selection variance
 78            y=means + torch.sqrt(torch.tensor(.25,device=device))*torch.randn_like(means)
 79            # split is fixed and models are intentionally tiny
 80            def train(kind):
 81                net=torch.nn.Sequential(torch.nn.Linear(1,24),torch.nn.Tanh(),torch.nn.Linear(24,2)).to(device)
 82                opt=torch.optim.Adam(net.parameters(),lr=.01)
 83                for _ in range(180):
 84                    out=net(x[:9000]); mu=out[:,0]; logv=out[:,1]
 85                    if kind=='mse': loss=((mu-y[:9000])**2).mean()
 86                    else:
 87                        # recursive q is detached, as prescribed; mean is trained with
 88                        # heteroscedastic Gaussian NLL and the variance head is fitted too.
 89                        loss=.5*(((y[:9000]-mu)**2)/q[:9000].detach()+torch.log(q[:9000].detach())).mean()
 90                        loss += .05*.5*(logv-q[:9000].detach().log()).pow(2).mean()
 91                    opt.zero_grad(); loss.backward(); opt.step()
 92                with torch.no_grad():
 93                    o=net(x[9000:]); err=(o[:,0]-y[9000:])**2
 94                    cal=(err/q[9000:]).mean().item()
 95                    mse=err.mean().item()
 96                return mse,cal
 97            mse,cal_mse=train('mse'); mse_i,cal_i=train('idea')
 98            neural={'device':device,'baseline_mse':mse,'baseline_calibration_ratio':cal_mse,
 99                    'idea_mse':mse_i,'idea_calibration_ratio':cal_i}
100        except Exception as e:
101            neural={'fallback_error':str(e)}
102    except Exception as e: neural={'torch_error':str(e)}
103    result={'seed':SEED,'parameters':{'gamma':gamma,'reward_sigma':.35},
104            'prediction_1_recursive_vs_mc':rows,'prediction_2_selection_term':sel,
105            'prediction_3_horizon':horizon,'mini_experiment':neural}
106    Path('results.json').write_text(json.dumps(result,indent=2))
107    print(json.dumps(result,indent=2))
108
109if __name__=='__main__': main()