Observable-Reduced Neural World Model / observable_reduced_experiment.py
Beats tuned baseline
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()