1import json, math, os
  2import numpy as np
  3import torch
  4
  5SEED = 1267
  6np.random.seed(SEED)
  7torch.manual_seed(SEED)
  8try:
  9    device = 'cuda' if torch.cuda.is_available() else 'cpu'
 10    if device == 'cuda':
 11        torch.cuda.set_device(0)
 12        _ = torch.zeros(1, device='cuda')
 13except Exception:
 14    device = 'cpu'
 15
 16
 17def solve_local(R, B, a, g, lam, u0=None, bdiff=None):
 18    """Exact minimizer of 1/2 u'Ru + g'u + lam/2 (b+Bu-b0-Bu0)'a^-1(...)."""
 19    R, B, a, g = map(np.asarray, (R, B, a, g))
 20    n = R.shape[0]
 21    u0 = np.zeros(n) if u0 is None else np.asarray(u0)
 22    bdiff = np.zeros(B.shape[0]) if bdiff is None else np.asarray(bdiff)
 23    ai = np.linalg.pinv(a)
 24    H = R + lam * B.T @ ai @ B
 25    rhs = -g - lam * B.T @ ai @ (bdiff - B @ u0)
 26    return np.linalg.solve(H, rhs), H
 27
 28
 29def toy_verification():
 30    # Two drift coordinates with equal actuation but very different diffusion.
 31    R = np.eye(2)
 32    B = np.eye(2)
 33    a = np.diag([0.25, 4.0])
 34    g = np.array([-2.0, -2.0])
 35    u0 = np.zeros(2)
 36    lambdas = np.array([0., .1, .3, 1., 3., 10., 30., 100.])
 37    us = np.array([solve_local(R, B, a, g, x, u0)[0] for x in lambdas])
 38    weights = np.diag(np.linalg.pinv(a))
 39    drift_kl = .5 * np.sum((us @ B.T) ** 2 * weights, axis=1)
 40    # Prediction 1: KL decreases monotonically as lambda increases.
 41    monotone = bool(np.all(np.diff(drift_kl) <= 1e-10))
 42    # Prediction 2: each coordinate's deviation has asymptotic lambda*u -> -a*g.
 43    tail = lambdas >= 10
 44    asymp_expected = -np.diag(a) * g
 45    asymp_observed = np.mean(us[tail] * lambdas[tail, None], axis=0)
 46    rel_asymp_err = float(np.max(np.abs(asymp_observed - asymp_expected) /
 47                                (np.abs(asymp_expected) + 1e-12)))
 48    # Prediction 3: regularization curvature ratio equals inverse-noise ratio.
 49    curv0 = np.diag(R)
 50    curv_lam = np.diag(R + 1.0 * B.T @ np.linalg.pinv(a) @ B)
 51    observed_ratio = curv_lam[0] / curv_lam[1]
 52    predicted_ratio = (1 + 1/a[0, 0]) / (1 + 1/a[1, 1])
 53    return {
 54        'lambdas': lambdas.tolist(), 'solutions': us.tolist(), 'trajectory_kl': drift_kl.tolist(),
 55        'prediction_1_kl_monotone': monotone,
 56        'prediction_2_expected_lambda_u_tail': asymp_expected.tolist(),
 57        'prediction_2_observed_lambda_u_tail': asymp_observed.tolist(),
 58        'prediction_2_relative_error': rel_asymp_err,
 59        'prediction_3_predicted_curvature_ratio': float(predicted_ratio),
 60        'prediction_3_observed_curvature_ratio': float(observed_ratio),
 61        'prediction_3_relative_error': float(abs(observed_ratio-predicted_ratio)/predicted_ratio),
 62    }
 63
 64
 65def finite_difference_check():
 66    R = np.diag([1.3, .7]); B = np.array([[1., .2], [.1, 1.]])
 67    a = np.diag([.4, 2.5]); g = np.array([-.8, .3]); lam = 2.7
 68    u, H = solve_local(R, B, a, g, lam)
 69    def f(x):
 70        d = B @ x
 71        return .5*x@R@x + g@x + .5*lam*d@np.linalg.pinv(a)@d
 72    eps = 1e-5
 73    grad = np.array([(f(u+eps*np.eye(2)[i])-f(u-eps*np.eye(2)[i]))/(2*eps) for i in range(2)])
 74    hess = np.empty((2,2))
 75    for i in range(2):
 76        for j in range(2):
 77            ei, ej = np.eye(2)[i]*eps, np.eye(2)[j]*eps
 78            hess[i,j] = (f(u+ei+ej)-f(u+ei-ej)-f(u-ei+ej)+f(u-ei-ej))/(4*eps*eps)
 79    return {'stationarity_gradient_norm': float(np.linalg.norm(grad)),
 80            'hessian_error_norm': float(np.linalg.norm(hess-H)), 'analytic_u': u.tolist()}
 81
 82
 83def train_policy():
 84    # Same synthetic states and task for action-norm baseline versus noise-whitened drift KL.
 85    torch.manual_seed(SEED)
 86    n, d, act = 4096, 3, 2
 87    x = torch.randn(n, d, device=device)
 88    target = x @ torch.tensor([[1.0, -.4], [.2, .8], [-.7, .3]], device=device)
 89    # State-dependent but known diagonal diffusion; one direction is highly noisy.
 90    var = torch.tensor([.2, 4.0], device=device)
 91    B = torch.eye(act, device=device)
 92    lam = 2.0
 93    def fit(kind):
 94        torch.manual_seed(SEED + (0 if kind == 'action' else 1))
 95        W = torch.nn.Parameter(torch.zeros(d, act, device=device))
 96        opt = torch.optim.Adam([W], lr=.08)
 97        for _ in range(700):
 98            u = x @ W
 99            task = .5 * ((u-target)**2).mean()
100            if kind == 'action':
101                reg = .5 * (u**2).sum(1).mean()
102            else:
103                reg = .5 * ((u**2)/var).sum(1).mean()
104            loss = task + lam*reg
105            opt.zero_grad(); loss.backward(); opt.step()
106        with torch.no_grad():
107            u = x @ W
108            task = .5*((u-target)**2).mean().item()
109            action_reg = .5*(u**2).sum(1).mean().item()
110            traj_kl = .5*((u**2)/var).sum(1).mean().item()
111        return {'task_mse_half': task, 'action_penalty': action_reg,
112                'trajectory_kl': traj_kl, 'W': W.detach().cpu().numpy().tolist()}
113    return {'device': device, 'lambda': lam, 'action_regularization': fit('action'),
114            'noise_whitened_drift_regularization': fit('drift')}
115
116
117def main():
118    out = {'seed': SEED, 'math_check': finite_difference_check(),
119           'toy_predictions': toy_verification(), 'policy_experiment': train_policy()}
120    with open('results.json', 'w') as f: json.dump(out, f, indent=2)
121    print(json.dumps(out, indent=2))
122
123if __name__ == '__main__':
124    main()