Residual-Tightened Neural Safety Shield / shield_experiment.py

Mechanism confirmed, baseline not beaten

Raw ⬇ ZIP
  1"""Residual-tightened safety shield: math check and reproducible toy experiment."""
  2import json
  3from pathlib import Path
  4import numpy as np
  5
  6SEED = 544
  7SIGMA_R = 0.20
  8KAPPA = 1.0
  9BOUND = 1.0
 10ACTION_LO, ACTION_HI = -0.5, 0.4
 11EMA_DECAY = 0.90
 12
 13
 14def normalized_residual(d, ema, sigma=SIGMA_R):
 15    return float(d / (sigma + ema))
 16
 17
 18def update_ema(ema, d, decay=EMA_DECAY):
 19    return float(decay * ema + (1.0 - decay) * d)
 20
 21
 22def nominal_model(x, a):
 23    return 0.85 * x + a
 24
 25
 26def shield_projection(x, a_policy, residual_r, kappa=KAPPA):
 27    """Exact QP projection for h(f(x,a))+kappa*r = x+a-BOUND+kappa*r <= 0."""
 28    # The feasible action interval is a <= BOUND-x-kappa*r.
 29    upper = min(ACTION_HI, BOUND - x - kappa * residual_r)
 30    return float(np.clip(a_policy, ACTION_LO, upper))
 31
 32
 33def policy(x):
 34    # Nominally feasible controller; the shifted dynamics can push it unsafe.
 35    return float(np.clip(0.75 - x, ACTION_LO, ACTION_HI))
 36
 37
 38def math_check():
 39    # Constant-error regimes have EMA[d] -> d, hence r -> d/(sigma+d).
 40    ds = np.array([0.0, 0.02, 0.05, 0.20, 0.80])
 41    rs = np.array([normalized_residual(d, d) for d in ds])
 42    margins = KAPPA * rs
 43    # Also check the exact projection satisfies the tightened constraint.
 44    xs = np.linspace(-0.2, 1.0, 31)
 45    max_violation = 0.0
 46    for x in xs:
 47        for r in rs:
 48            a = shield_projection(x, policy(x), r)
 49            max_violation = max(max_violation, x + a - BOUND + KAPPA*r)
 50    return {
 51        "d": ds.tolist(), "r_at_ema_equal_d": rs.tolist(),
 52        "tightening_margin": margins.tolist(),
 53        "r_monotone": bool(np.all(np.diff(rs) >= -1e-12)),
 54        "max_projection_constraint_violation": float(max_violation),
 55        "constraint_satisfied": bool(max_violation <= 1e-10),
 56    }
 57
 58
 59def run_condition(regime, method, rng, episodes=160, horizon=50):
 60    """Simulate true x_next=x+a+bias+noise while the model assumes x+a."""
 61    bias = 0.03 if regime == "nominal" else 0.50
 62    # Fixed margin control is deliberately the margin needed by the shifted regime.
 63    fixed_r = normalized_residual(0.50, 0.50)
 64    violations = 0
 65    total_dev = 0.0
 66    total_reward = 0.0
 67    min_margin = []
 68    for _ in range(episodes):
 69        x = 0.20
 70        ema = 0.05  # modest prior scale; prevents an artificial first-step infinite residual
 71        # The shield starts with a nominal-model residual prior; shifted error is learned online.
 72        r_prev = normalized_residual(0.03, 0.03)
 73        for _ in range(horizon):
 74            ap = policy(x)
 75            if method == "unshielded":
 76                a = float(np.clip(ap, ACTION_LO, ACTION_HI))
 77            elif method == "fixed":
 78                a = shield_projection(x, ap, fixed_r)
 79            elif method == "adaptive":
 80                r = normalized_residual(ema, ema)  # unused only for initialization below
 81                # r is the online residual estimate from prior measurements; EMA is a scale.
 82                # We retain the last normalized residual in the loop state via r_prev.
 83                a = shield_projection(x, ap, r_prev)
 84            else:
 85                raise ValueError(method)
 86            # Shifted true dynamics; small sensor/process noise makes violations nontrivial.
 87            noise = float(rng.normal(0.0, 0.004))
 88            xn = nominal_model(x, a) + bias + noise
 89            d = abs(xn - nominal_model(x, a))
 90            if xn > BOUND:
 91                violations += 1
 92            total_dev += abs(a - ap)
 93            # Reward preserves the requested behavior: approach the goal, penalize intervention.
 94            total_reward += -(0.75 - xn) ** 2 - 0.05 * (a - ap) ** 2
 95            ema = update_ema(ema, d)
 96            r_new = min(0.8, normalized_residual(d, ema))
 97            if method == "adaptive":
 98                r_prev = r_new
 99            x = xn
100        # for unshielded/fixed no r_prev is needed; initialize adaptive before first action
101    n = episodes * horizon
102    return {"violation_rate": violations / n, "mean_action_deviation": total_dev / n,
103            "mean_reward": total_reward / n}
104
105
106def run_all():
107    results = {}
108    for regime in ("nominal", "shifted"):
109        for method in ("unshielded", "fixed", "adaptive"):
110            rng = np.random.default_rng(SEED + (0 if regime == "nominal" else 100) +
111                                        {"unshielded": 0, "fixed": 10, "adaptive": 20}[method])
112            results[f"{regime}_{method}"] = run_condition(regime, method, rng)
113    return results
114
115
116if __name__ == "__main__":
117    out = {"seed": SEED, "sigma_r": SIGMA_R, "kappa": KAPPA,
118           "math_check": math_check(), "experiment": run_all()}
119    Path("results.json").write_text(json.dumps(out, indent=2))
120    print(json.dumps(out, indent=2))