Observable-Reduced Neural World Model / observable_reduced_experiment.py

✓✓ Beats tuned baseline

Raw ⬇ ZIP
  1import json
  2import math
  3import os
  4import numpy as np
  5from scipy.linalg import expm
  6from scipy.signal import savgol_filter
  7from scipy.optimize import least_squares
  8
  9SEED = 2112
 10rng = np.random.default_rng(SEED)
 11
 12
 13def exact_trajectory(r, alpha, beta, a0, q0, t):
 14    """Exact linear dynamics [a,q]' = M[a,q]."""
 15    M = np.array([[r-alpha, beta], [alpha, -beta]], dtype=float)
 16    x0 = np.array([a0, q0], dtype=float)
 17    return np.array([(expm(M * ti) @ x0) for ti in t])
 18
 19
 20def true_coefficients(r, alpha, beta):
 21    return beta * r, r - alpha - beta
 22
 23
 24def finite_derivatives(y, dt):
 25    """Centered second-order finite differences, excluding endpoints."""
 26    v = (y[2:] - y[:-2]) / (2.0 * dt)
 27    acc = (y[2:] - 2.0 * y[1:-1] + y[:-2]) / (dt * dt)
 28    return y[1:-1], v, acc
 29
 30
 31def fit_reduced_from_samples(y, dt):
 32    ym, v, acc = finite_derivatives(y, dt)
 33    X = np.column_stack([ym, v])
 34    coef, *_ = np.linalg.lstsq(X, acc, rcond=None)
 35    pred = X @ coef
 36    return coef, float(np.sqrt(np.mean((pred - acc) ** 2)))
 37
 38
 39def fit_first_order_from_samples(y, dt):
 40    v = (y[2:] - y[:-2]) / (2.0 * dt)
 41    ym = y[1:-1]
 42    k = float(np.dot(ym, v) / (np.dot(ym, ym) + 1e-12))
 43    return k
 44
 45
 46def rollout_reduced(y0, v0, c, dt, n):
 47    y, v = float(y0), float(v0)
 48    out = [y]
 49    cy, cv = c
 50    for _ in range(n):
 51        # RK4 for y'=v, v'=cy*y+cv*v
 52        def f(z): return np.array([z[1], cy*z[0] + cv*z[1]])
 53        z = np.array([y, v])
 54        k1 = f(z); k2 = f(z + dt*k1/2); k3 = f(z + dt*k2/2); k4 = f(z + dt*k3)
 55        z = z + dt*(k1 + 2*k2 + 2*k3 + k4)/6
 56        y, v = z
 57        out.append(y)
 58    return np.asarray(out)
 59
 60
 61def rollout_first_order(y0, k, dt, n):
 62    return y0 * np.exp(k * dt * np.arange(n+1))
 63
 64
 65def invariant_sweep():
 66    """Prediction 1: fitted cy,cv obey exact rate relations over rate sweep."""
 67    rows = []
 68    r = 0.20
 69    t = np.arange(0, 20.0001, 0.01)
 70    for alpha in [0.05, 0.15, 0.40, 0.80]:
 71        for beta in [0.03, 0.12, 0.35, 0.70]:
 72            x = exact_trajectory(r, alpha, beta, 1.0, 0.4, t)
 73            c, rmse = fit_reduced_from_samples(x[:, 0] + x[:, 1], t[1]-t[0])
 74            truth = np.array(true_coefficients(r, alpha, beta))
 75            rel = np.abs((c-truth) / np.maximum(np.abs(truth), 1e-8))
 76            rows.append({'alpha':alpha, 'beta':beta, 'cy_hat':c[0], 'cv_hat':c[1],
 77                         'cy_true':truth[0], 'cv_true':truth[1], 'max_relative_error':float(np.max(rel)),
 78                         'accel_fit_rmse':rmse})
 79    return rows
 80
 81
 82def sampling_sweep():
 83    """Prediction 2: derivative/reduced accuracy degrades as dt approaches 1/(alpha+beta)."""
 84    r, alpha, beta = 0.20, 0.60, 0.40
 85    tau = 1.0/(alpha+beta)
 86    truth = np.array(true_coefficients(r, alpha, beta))
 87    rows = []
 88    # Includes well below, near, and above the claimed fastest switching scale.
 89    for ratio in [0.05, 0.10, 0.20, 0.50, 1.0, 1.5, 2.0]:
 90        dt = ratio * tau
 91        t = np.arange(0, 30 + 0.5*dt, dt)
 92        x = exact_trajectory(r, alpha, beta, 1.0, 0.25, t)
 93        y = x[:, 0] + x[:, 1]
 94        c, rmse = fit_reduced_from_samples(y, dt)
 95        rel = float(np.linalg.norm(c-truth)/(np.linalg.norm(truth)+1e-12))
 96        rows.append({'dt_over_tau':ratio, 'dt':dt, 'coefficient_relative_error':rel,
 97                     'accel_rmse':rmse, 'cy_hat':c[0], 'cv_hat':c[1]})
 98    return rows
 99
100
101def noisy_rollout_comparison():
102    """Prediction 3: structural second-order model extrapolates better than y'=k y."""
103    r, alpha, beta = 0.18, 0.45, 0.25
104    dt = 0.08
105    t_train = np.arange(0, 8+1e-9, dt)
106    t_test = np.arange(0, 40+1e-9, dt)
107    train = exact_trajectory(r, alpha, beta, 1.0, 0.8, t_train)
108    test = exact_trajectory(r, alpha, beta, 1.0, 0.8, t_test)
109    ytrain = train.sum(axis=1)
110    ytrue = test.sum(axis=1)
111    noise = 0.002 * np.std(ytrain) * rng.normal(size=ytrain.shape)
112    c, fit_rmse = fit_reduced_from_samples(ytrain + noise, dt)
113    k = fit_first_order_from_samples(ytrain + noise, dt)
114    v0 = (ytrain[1]-ytrain[0])/dt
115    pred2 = rollout_reduced(ytrain[0], v0, c, dt, len(t_test)-1)
116    pred1 = rollout_first_order(ytrain[0], k, dt, len(t_test)-1)
117    horizon = {'reduced_rmse':float(np.sqrt(np.mean((pred2-ytrue)**2))),
118               'first_order_rmse':float(np.sqrt(np.mean((pred1-ytrue)**2))),
119               'reduced_tail_rmse':float(np.sqrt(np.mean((pred2[len(pred2)//2:]-ytrue[len(ytrue)//2:])**2))),
120               'first_order_tail_rmse':float(np.sqrt(np.mean((pred1[len(pred1)//2:]-ytrue[len(ytrue)//2:])**2))),
121               'reduced_coefficients':c.tolist(), 'first_order_k':k, 'fit_accel_rmse':fit_rmse}
122    return horizon
123
124
125def main():
126    inv = invariant_sweep()
127    samp = sampling_sweep()
128    comp = noisy_rollout_comparison()
129    inv_err = [x['max_relative_error'] for x in inv]
130    # Quantitative criteria: exact math on fine samples, and clear sampling transition.
131    low = np.mean([x['coefficient_relative_error'] for x in samp if x['dt_over_tau'] <= .2])
132    high = np.mean([x['coefficient_relative_error'] for x in samp if x['dt_over_tau'] >= 1.0])
133    result = {
134      'seed': SEED,
135      'prediction_1_invariant_rate_sweep': {
136        'predicted': 'cy=beta*r and cv=r-alpha-beta for every alpha,beta',
137        'observed_max_relative_error': float(max(inv_err)),
138        'observed_median_relative_error': float(np.median(inv_err)),
139        'n_cases':len(inv)},
140      'prediction_2_sampling_transition': {
141        'predicted': 'error small for dt/tau << 1 and degrades around dt/tau >= 1',
142        'tau':1.0, 'low_ratio_mean_error':float(low), 'high_ratio_mean_error':float(high),
143        'sweep':samp},
144      'prediction_3_rollout': {
145        'predicted':'second-order reduced dynamics has lower long-horizon error than first-order y_dot=k*y',
146        **comp},
147      'invariant_rows':inv}
148    with open('results.json','w') as f: json.dump(result,f,indent=2)
149    print(json.dumps(result, indent=2))
150
151if __name__ == '__main__': main()