Certified Tube Wrapper for Learned Predictive Control / tube_experiment.py

Failed on benchmark

Raw ⬇ ZIP
  1"""Certified tube wrapper MVP and numerical verification.
  2
  3Run with: python3 tube_experiment.py
  4All experiments are deterministic and use only numpy.
  5"""
  6import json
  7import math
  8from pathlib import Path
  9import numpy as np
 10
 11SEED = 7
 12
 13
 14def propagate(A_abs, e0, w):
 15    """Elementwise interval-radius recursion e' = |A|e+w."""
 16    e = np.asarray(e0, dtype=float).copy()
 17    hist = [e.copy()]
 18    for A, wi in zip(A_abs, w):
 19        e = np.asarray(A) @ e + np.asarray(wi)
 20        hist.append(e.copy())
 21    return np.asarray(hist)
 22
 23
 24def spectral_radius(A):
 25    return float(np.max(np.abs(np.linalg.eigvals(A))))
 26
 27
 28def verify_math():
 29    # Prediction 1: scalar recursion diverges iff gamma*Lambda >= 1.
 30    # At equality it grows linearly for positive disturbance; above it grows exponentially.
 31    gammas = np.linspace(0.60, 1.40, 161)
 32    lam = 1.0
 33    N = 80
 34    terminal = []
 35    for g in gammas:
 36        e = 0.0
 37        for _ in range(N):
 38            e = (g * lam) * e + 0.01
 39        terminal.append(e)
 40    terminal = np.asarray(terminal)
 41    # finite/unstable transition estimated by a large finite-horizon slope jump
 42    stable_med = np.median(terminal[gammas < .98])
 43    unstable_med = np.median(terminal[gammas > 1.02])
 44    boundary_obs = gammas[np.argmin(np.abs(np.log1p(terminal) -
 45                                           np.median(np.log1p(terminal))))]
 46    # More meaningful classification: ratio e_N/e_(N-10) exceeds 1.2 on unstable side.
 47    growth = []
 48    for g in gammas:
 49        e1 = e2 = 0.
 50        for i in range(N):
 51            e1 = g * e1 + .01
 52            if i == N - 11:
 53                old = e1
 54            if i == N - 1:
 55                e2 = e1
 56        growth.append(e2 / max(old, 1e-12))
 57    growth = np.asarray(growth)
 58    idx = np.where(growth > 1.20)[0]
 59    boundary_growth = float(gammas[idx[0]]) if len(idx) else float('nan')
 60
 61    # Prediction 2: steady radius is w/(1-r), exactly for scalar constant A.
 62    rs = np.array([0.10, 0.25, 0.40, 0.55, 0.70, 0.80, 0.90])
 63    w0 = .02
 64    measured = []
 65    predicted = []
 66    for r in rs:
 67        e = 0.
 68        for _ in range(300):
 69            e = r * e + w0
 70        measured.append(e)
 71        predicted.append(w0 / (1-r))
 72    measured, predicted = np.asarray(measured), np.asarray(predicted)
 73    rel_err = float(np.max(np.abs(measured-predicted)/predicted))
 74
 75    # Prediction 3: multiplicative conversion is linear in beta and |u|.
 76    B, u = 0.35, 0.8
 77    betas = np.array([0., .05, .10, .20, .40, .70])
 78    wm = np.abs(B) * betas * abs(u)
 79    slope = float(np.polyfit(betas, wm, 1)[0])
 80    predicted_slope = abs(B*u)
 81
 82    # A 2D componentwise tube is also checked against sampled errors.
 83    A = np.array([[.55, .10], [.05, .35]])
 84    e0 = np.array([.01, .02])
 85    w = np.array([.015, .01])
 86    h = 30
 87    bound = propagate([np.abs(A)]*h, e0, [w]*h)[-1]
 88    rng = np.random.default_rng(SEED)
 89    x = e0.copy()
 90    max_ratio = 0.
 91    for _ in range(20000):
 92        x = A @ x + rng.uniform(-w, w)
 93        max_ratio = max(max_ratio, float(np.max(np.abs(x) / bound)))
 94    return {
 95        "boundary_prediction": 1.0,
 96        "boundary_observed_growth_test": boundary_growth,
 97        "boundary_terminal_below_0.98": float(stable_med),
 98        "boundary_terminal_above_1.02": float(unstable_med),
 99        "steady_gain_max_relative_error": rel_err,
100        "steady_gain_predicted": (w0/(1-rs)).tolist(),
101        "steady_gain_measured": measured.tolist(),
102        "multiplicative_slope_predicted": predicted_slope,
103        "multiplicative_slope_observed": slope,
104        "sampled_2d_bound_max_ratio": max_ratio,
105    }
106
107
108def policy(x):
109    # A small nonlinear learned-policy surrogate; its Jacobian is available analytically.
110    return -0.52*x + 0.08*np.tanh(2*x)
111
112
113def policy_jac(x):
114    return -0.52 + 0.16/(np.cosh(2*x)**2)
115
116
117def nominal_f(x, u):
118    return 0.82*x + u
119
120
121def nominal_fx(x, u):
122    return .82
123
124
125def nominal_fu(x, u):
126    return 1.0
127
128
129def tube_radii(xhat, horizon, beta, mode):
130    """Propagate learned closed-loop Jacobian and additive+mult uncertainty."""
131    e = 0.
132    out = []
133    # horizon-specific model Jacobian evaluated along rollout
134    z = xhat
135    for j in range(horizon):
136        u = policy(z)
137        acl = abs(nominal_fx(z,u) + nominal_fu(z,u)*policy_jac(z))
138        # Calibrated additive residual plus parameter/gain uncertainty.
139        wadd = .018
140        wmult = .20 * beta * abs(u)  # |B| pbar |u|
141        e = acl*e + wadd + wmult
142        out.append(e)
143        z = nominal_f(z,u)
144    if mode == "constant":
145        # Stationary worst-case radius from calibration, deliberately conservative.
146        rmax = .46
147        wmax = .018 + .20*beta*.65
148        steady = wmax/(1-rmax)
149        return np.full(horizon, steady)
150    if mode == "nominal":
151        return np.zeros(horizon)
152    return np.asarray(out)
153
154
155def run_control(mode, beta, seed=SEED, episodes=300):
156    rng = np.random.default_rng(seed)
157    horizon, limit = 8, 1.0
158    violations = fallbacks = 0
159    total_abs_x = total_u = 0.
160    total_steps = episodes*35
161    for _ in range(episodes):
162        x = rng.uniform(-.72, .72)
163        for t in range(35):
164            # The nominal controller is the learned policy, with a tube-aware bounded correction.
165            u_nom = policy(x)
166            radii = tube_radii(x, horizon, beta, mode)
167            # Use the first-step state tightening. If nominal state is outside it, fallback.
168            if mode != "nominal" and abs(x) > limit-radii[0]:
169                fallbacks += 1
170                # conservative bounded fallback, observable in the log
171                u = float(np.clip(-.30*x, -.45, .45))
172            else:
173                # nominal MPC-like correction: robust modes reduce action near boundary
174                correction = 0.
175                if mode == "horizon":
176                    correction = -.10*x*min(1., radii[0]/.20)
177                elif mode == "constant":
178                    correction = -.06*x
179                u = float(np.clip(u_nom + correction, -1., 1.))
180            # True plant has randomized gain and additive bounded/noisy disturbance.
181            gain = 1.0 + rng.uniform(-beta, beta)
182            d = rng.uniform(-.018, .018) + rng.normal(0., .006)
183            x = .82*x + gain*u + d
184            violations += int(abs(x) > limit)
185            total_abs_x += abs(x)
186            total_u += abs(u)
187    return {
188        "violation_rate": violations/total_steps,
189        "fallback_rate": fallbacks/total_steps,
190        "mean_abs_state": total_abs_x/total_steps,
191        "mean_abs_action": total_u/total_steps,
192    }
193
194
195def main():
196    math_result = verify_math()
197    mini = {}
198    for beta in (0.10, 0.25, 0.40):
199        mini[str(beta)] = {m: run_control(m, beta) for m in ("nominal", "constant", "horizon")}
200    result = {"math_verification": math_result, "mini_experiment": mini}
201    Path("results.json").write_text(json.dumps(result, indent=2))
202    print(json.dumps(result, indent=2))
203
204
205if __name__ == "__main__":
206    main()