| """Dynamic signed-alpha controller for native derivative steering. |
| |
| The Wasserstein gate decides *whether* correction is eligible. Dose and |
| direction are computed separately from the latest one-step firing-rate delta |
| and a calibration-split alpha -> future-delta response curve. Commands |
| returned after decision ``d`` are intended for decision ``d + 1``. |
| """ |
|
|
| from __future__ import annotations |
|
|
| from dataclasses import dataclass |
|
|
| import numpy as np |
|
|
| from analysis.native_derivative_gate import NativeDerivativeGate, NativeDerivativeGateConfig |
|
|
|
|
| @dataclass(frozen=True) |
| class MonotoneResponseCurve: |
| alphas: tuple[float, ...] |
| corrections: tuple[float, ...] |
|
|
| def __post_init__(self) -> None: |
| if len(self.alphas) != len(self.corrections) or len(self.alphas) < 2: |
| raise ValueError("response curve needs at least two alpha/correction points") |
| if any(not np.isfinite(x) for x in self.alphas + self.corrections): |
| raise ValueError("response curve values must be finite") |
| if any(a < 0 for a in self.alphas) or any(c < 0 for c in self.corrections): |
| raise ValueError("response alpha/correction values must be non-negative") |
| if any(b <= a for a, b in zip(self.alphas, self.alphas[1:])): |
| raise ValueError("response alpha values must be strictly increasing") |
| if any(b < a for a, b in zip(self.corrections, self.corrections[1:])): |
| raise ValueError("response corrections must be monotone non-decreasing") |
|
|
| def inverse(self, requested_correction: float, alpha_max: float) -> tuple[float, float, bool]: |
| """Return minimal calibrated alpha, predicted correction, saturation.""" |
| requested = max(0.0, float(requested_correction)) |
| curve_alpha_max = min(float(alpha_max), self.alphas[-1]) |
| all_alphas = np.asarray(self.alphas, dtype=np.float64) |
| all_corrections = np.asarray(self.corrections, dtype=np.float64) |
| supported = all_alphas < curve_alpha_max - 1e-12 |
| alphas = all_alphas[supported] |
| corrections = all_corrections[supported] |
| endpoint_correction = float(np.interp(curve_alpha_max, all_alphas, all_corrections)) |
| alphas = np.r_[alphas, curve_alpha_max] |
| corrections = np.r_[corrections, endpoint_correction] |
| if alphas.size < 2: |
| return 0.0, 0.0, True |
| max_correction = float(corrections[-1]) |
| saturated = requested > max_correction + 1e-12 |
| target = min(requested, max_correction) |
| |
| |
| |
| unique_correction, first = np.unique(corrections, return_index=True) |
| unique_alpha = alphas[first] |
| if unique_correction.size == 1: |
| return (curve_alpha_max if requested > 0 else 0.0), max_correction, saturated |
| alpha = float(np.interp(target, unique_correction, unique_alpha)) |
| predicted = float(np.interp(alpha, alphas, corrections)) |
| return min(alpha, curve_alpha_max), predicted, saturated |
|
|
|
|
| @dataclass(frozen=True) |
| class NativeDerivativeDynamicConfig: |
| gate: NativeDerivativeGateConfig |
| success_delta_quantiles: tuple[tuple[float, ...], ...] |
| amplify_response: tuple[MonotoneResponseCurve | None, ...] |
| suppress_response: tuple[MonotoneResponseCurve | None, ...] |
| alpha_max: float |
| minimum_effective_alpha: float = 0.0 |
| proportional_gain: float = 1.0 |
| target_mode: str = "median" |
| confidence_scaling: bool = False |
|
|
| def __post_init__(self) -> None: |
| windows = len(self.gate.success_quantiles) |
| if len(self.success_delta_quantiles) != windows: |
| raise ValueError("success delta references must match gate windows") |
| if len(self.amplify_response) != windows or len(self.suppress_response) != windows: |
| raise ValueError("response curves must match gate windows") |
| if self.alpha_max <= 0 or not 0 <= self.minimum_effective_alpha <= self.alpha_max: |
| raise ValueError("invalid alpha bounds") |
| if self.proportional_gain <= 0: |
| raise ValueError("proportional_gain must be positive") |
| if self.target_mode not in {"median", "band"}: |
| raise ValueError("target_mode must be 'median' or 'band'") |
|
|
|
|
| class NativeDerivativeDynamicController: |
| """Eligibility gate plus calibrated bidirectional inverse controller.""" |
|
|
| def __init__(self, config: NativeDerivativeDynamicConfig): |
| self.config = config |
| self.gate = NativeDerivativeGate(config.gate) |
|
|
| def observe_decision(self, firing_rate: float) -> dict: |
| obs = self.gate.observe_decision(firing_rate) |
| window = int(obs["window_index"]) |
| current = float(obs["one_step_delta"]) |
| reference = np.asarray(self.config.success_delta_quantiles[window], dtype=np.float64) |
| reference = reference[np.isfinite(reference)] |
|
|
| invalid_control = not np.isfinite(current) |
| missing_target = reference.size == 0 |
| error = float("nan") |
| target = lower = upper = float("nan") |
| direction = "none" |
| requested = 0.0 |
| signed_alpha = 0.0 |
| predicted = 0.0 |
| saturated = False |
| missing_response = False |
|
|
| if not missing_target: |
| target = float(np.median(reference)) |
| lower = float(np.quantile(reference, 0.25)) |
| upper = float(np.quantile(reference, 0.75)) |
| if not invalid_control and not missing_target: |
| if self.config.target_mode == "band": |
| if current < lower: |
| error = lower - current |
| elif current > upper: |
| error = upper - current |
| else: |
| error = 0.0 |
| else: |
| error = target - current |
|
|
| eligible = bool(obs["gate_on_next_decision"]) |
| if eligible and np.isfinite(error) and error != 0.0: |
| direction = "amplify" if error > 0 else "suppress" |
| requested = self.config.proportional_gain * abs(error) |
| curve = ( |
| self.config.amplify_response[window] |
| if direction == "amplify" |
| else self.config.suppress_response[window] |
| ) |
| if curve is None: |
| missing_response = True |
| else: |
| alpha, predicted, saturated = curve.inverse(requested, self.config.alpha_max) |
| if 0 < alpha < self.config.minimum_effective_alpha: |
| alpha = self.config.minimum_effective_alpha |
| predicted = float(np.interp(alpha, curve.alphas, curve.corrections)) |
| if self.config.confidence_scaling: |
| ema = float(obs["ema_failure_score"]) |
| confidence = float(np.clip( |
| (ema - self.config.gate.theta_off) |
| / (self.config.gate.theta_on - self.config.gate.theta_off), |
| 0.0, |
| 1.0, |
| )) if np.isfinite(ema) else 0.0 |
| alpha *= confidence |
| predicted *= confidence |
| signed_alpha = alpha if direction == "amplify" else -alpha |
|
|
| return { |
| **obs, |
| "control_signal": "delta", |
| "delta_current": current, |
| "delta_target": target, |
| "delta_success_lower": lower, |
| "delta_success_upper": upper, |
| "delta_error": error, |
| "command_direction_next_decision": direction if signed_alpha != 0 else "none", |
| "requested_delta_correction": requested, |
| "predicted_delta_correction": predicted, |
| "signed_alpha_next_decision": signed_alpha, |
| "controller_saturated": saturated, |
| "invalid_control_signal": invalid_control, |
| "missing_target_support": missing_target, |
| "missing_response_support": missing_response, |
| } |
|
|
|
|
| def _curve_from_payload(payload: dict | None) -> MonotoneResponseCurve | None: |
| if not payload: |
| return None |
| return MonotoneResponseCurve( |
| alphas=tuple(float(x) for x in payload["alphas"]), |
| corrections=tuple(float(x) for x in payload["corrections"]), |
| ) |
|
|
|
|
| def load_native_derivative_dynamic_configs( |
| payload: dict, |
| *, |
| task: int, |
| layers: list[str], |
| alpha_max_override: float | None = None, |
| ) -> dict[str, NativeDerivativeDynamicConfig]: |
| if payload.get("model") != "pi05": |
| raise ValueError("Pi0.5 dynamic controller requires calibration.model == 'pi05'") |
| defaults = payload.get("controller_defaults", {}) |
| gate_defaults = payload.get("gate_defaults", {}) |
| responses = { |
| (str(row["layer"]), int(row["task"]), int(row["feature_id"])): row |
| for row in payload.get("response_calibration", []) |
| } |
| out: dict[str, NativeDerivativeDynamicConfig] = {} |
| for layer in layers: |
| matches = [ |
| row for row in payload.get("selected_features", []) |
| if row.get("layer") == layer and int(row.get("task", -1)) == int(task) |
| ] |
| if not matches: |
| raise KeyError(f"No native derivative feature for task={task}, layer={layer}") |
| row = matches[0] |
| feature_id = int(row["feature_id"]) |
| response = responses.get((layer, int(task), feature_id)) |
| if response is None: |
| raise KeyError(f"No response calibration for task={task}, layer={layer}, f{feature_id}") |
| gate = NativeDerivativeGateConfig( |
| feature_id=feature_id, |
| success_quantiles=tuple(tuple(float(v) for v in q) for q in row["success_quantiles"]), |
| failure_quantiles=tuple(tuple(float(v) for v in q) for q in row["failure_quantiles"]), |
| window_size_env_steps=int(payload["window_size_env_steps"]), |
| replan_steps=int(payload.get("replan_steps", 5)), |
| smooth_width_decisions=int(payload["smooth_width_decisions"]), |
| derivative_width_decisions=int(payload["derivative_width_decisions"]), |
| recent_window_decisions=int(row.get("recent_window_decisions", gate_defaults.get("recent_window_decisions", 3))), |
| mode="native_derivative_dynamic", |
| theta_on=float(row.get("theta_on", gate_defaults.get("theta_on", 0.2))), |
| theta_off=float(row.get("theta_off", gate_defaults.get("theta_off", 0.05))), |
| min_on_decisions=int(row.get("min_on_decisions", gate_defaults.get("min_on_decisions", 2))), |
| min_off_decisions=int(gate_defaults.get("min_off_decisions", 2)), |
| ema_beta=float(gate_defaults.get("ema_beta", 0.5)), |
| ) |
| windows = len(gate.success_quantiles) |
| amp_by_window = {int(x["window_index"]): x for x in response.get("amplify", [])} |
| sup_by_window = {int(x["window_index"]): x for x in response.get("suppress", [])} |
| artifact_alpha_max = float(response.get("alpha_max", defaults.get("alpha_max", 10.0))) |
| alpha_max = min(artifact_alpha_max, float(alpha_max_override)) if alpha_max_override is not None else artifact_alpha_max |
| out[layer] = NativeDerivativeDynamicConfig( |
| gate=gate, |
| success_delta_quantiles=tuple( |
| tuple(float(v) for v in q) for q in row["success_delta_quantiles"] |
| ), |
| amplify_response=tuple(_curve_from_payload(amp_by_window.get(i)) for i in range(windows)), |
| suppress_response=tuple(_curve_from_payload(sup_by_window.get(i)) for i in range(windows)), |
| alpha_max=alpha_max, |
| minimum_effective_alpha=float(response.get("minimum_effective_alpha", defaults.get("minimum_effective_alpha", 0.0))), |
| proportional_gain=float(defaults.get("proportional_gain", 1.0)), |
| target_mode=str(defaults.get("target_mode", "median")), |
| confidence_scaling=bool(defaults.get("confidence_scaling", False)), |
| ) |
| return out |
|
|
