"""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) # np.interp chooses the left-most useful segment after duplicate # correction values have been collapsed, which implements a minimum # dose inverse rather than an arbitrary high-dose inverse. 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