The Online Inference Problem That Second-Order Methods Finally Solved — Without Projections

The following is a complete PyTorch implementation of the AI-StoSQP method from the paper — Algorithm 1 together with the plug-in covariance estimator from Theorem 5.10. The implementation covers the randomized Kaczmarz sketching solver, the adaptive stepsize selection scheme from Appendix A, the Hessian averaging procedure, and a full confidence interval construction routine. A runnable smoke test on a constrained linear regression problem demonstrates all components end-to-end.

# =============================================================================
# AI-StoSQP: Adaptive Inexact Stochastic Sequential Quadratic Programming
# Paper: "Statistical Inference of Constrained Stochastic Optimization
#         via Sketched Sequential Quadratic Programming"
# Authors: Sen Na (Georgia Tech) and Michael W. Mahoney (UC Berkeley)
# Journal: JMLR 26 (2025) 1-75
# =============================================================================

from __future__ import annotations
import math
import warnings
import numpy as np
import torch
import torch.nn as nn
from dataclasses import dataclass, field
from typing import Callable, Dict, List, Optional, Tuple

warnings.filterwarnings('ignore')
torch.manual_seed(42)
np.random.seed(42)


# ─── SECTION 1: Problem Definition Interface ──────────────────────────────────

class ConstrainedStochasticProblem:
    """
    Abstract interface for a constrained stochastic optimization problem:

        min_{x in R^d}  E_P[F(x; xi)]
        s.t.            c(x) = 0          (m equality constraints)

    Users subclass this and implement:
        - sample_xi()      -> draw one sample xi ~ P
        - loss(x, xi)      -> scalar loss F(x; xi)
        - grad(x, xi)      -> gradient nabla_x F(x; xi)
        - hess(x, xi)      -> Hessian nabla^2_x F(x; xi)
        - constraint(x)    -> c(x) in R^m
        - jac(x)           -> G(x) = nabla c(x) in R^{m x d}
        - constraint_hess(x, lam) -> sum_i lam_i nabla^2 c_i(x)
    """
    def sample_xi(self): raise NotImplementedError
    def loss(self, x, xi): raise NotImplementedError
    def grad(self, x, xi): raise NotImplementedError
    def hess(self, x, xi): raise NotImplementedError
    def constraint(self, x): raise NotImplementedError
    def jac(self, x): raise NotImplementedError
    def constraint_hess(self, x, lam): raise NotImplementedError


# ─── SECTION 2: Randomized Kaczmarz Sketching Solver ─────────────────────────

def randomized_kaczmarz_step(
    K: torch.Tensor,
    rhs: torch.Tensor,
    z: torch.Tensor,
    rng: np.random.Generator,
) -> torch.Tensor:
    """
    One step of randomized Kaczmarz for the system K z = -rhs.

    Samples a random canonical basis vector e_j (uniform without replacement
    is approximated by sampling with replacement here for simplicity), then
    projects z onto the j-th equation of the system.

    The update is:
        z <- z - K e_j (e_j^T K^2 e_j)^{-1} e_j^T (K z + rhs)

    Parameters
    ----------
    K   : (n, n) symmetric system matrix (the Lagrangian Hessian block matrix)
    rhs : (n,) right-hand side vector (nabla L at current iterate)
    z   : (n,) current solution approximation
    rng : numpy random generator

    Returns
    -------
    z_new : updated solution approximation
    """
    n = K.shape[0]
    j = int(rng.integers(0, n))
    Kej = K[:, j]
    K2_jj = (Kej * Kej).sum().item()
    if abs(K2_jj) < 1e-14:
        return z
    residual_j = (Kej * z).sum() + rhs[j]
    z_new = z - (residual_j / K2_jj) * Kej
    return z_new


def sketch_solve(
    K: torch.Tensor,
    rhs: torch.Tensor,
    tau: int,
    rng: np.random.Generator,
    z0: Optional[torch.Tensor] = None,
) -> torch.Tensor:
    """
    Run tau steps of randomized Kaczmarz to solve K z = -rhs inexactly.

    This implements the sketching solver of Section 3 (Eq. 3.4) of the paper.
    After tau steps the approximation error E[||z_tau - z_exact||^2 | context]
    is bounded by rho^tau ||z_exact||^2 where rho = 1 - gamma_S.

    Parameters
    ----------
    K   : (n, n) the KKT matrix [B G^T; G 0] + regularization
    rhs : (n,) the right-hand side [nabla_x L; c]
    tau : number of sketching steps (fixed, independent of t)
    rng : random number generator
    z0  : initial approximation (defaults to zero)

    Returns
    -------
    z : approximate Newton direction after tau steps
    """
    n = K.shape[0]
    z = torch.zeros(n, dtype=K.dtype) if z0 is None else z0.clone()
    for _ in range(tau):
        z = randomized_kaczmarz_step(K, rhs, z, rng)
    return z


# ─── SECTION 3: KKT Matrix Assembly ──────────────────────────────────────────

def build_kkt_matrix(
    B: torch.Tensor,
    G: torch.Tensor,
    delta: float = 0.0,
) -> torch.Tensor:
    """
    Assemble the (d+m) x (d+m) KKT block matrix:

        K = [[B + delta*I,  G^T],
             [G,             0  ]]

    Parameters
    ----------
    B     : (d, d) regularized Lagrangian Hessian
    G     : (m, d) constraint Jacobian
    delta : Levenberg-Marquardt regularization (ensures positive definiteness)

    Returns
    -------
    K : (d+m, d+m) KKT matrix
    """
    d = B.shape[0]
    m = G.shape[0]
    n = d + m
    K = torch.zeros(n, n, dtype=B.dtype)
    K[:d, :d] = B + delta * torch.eye(d, dtype=B.dtype)
    K[:d, d:] = G.T
    K[d:, :d] = G
    return K


# ─── SECTION 4: Adaptive Stepsize Selection (Appendix A) ─────────────────────

def select_adaptive_stepsize(
    x: torch.Tensor,
    lam: torch.Tensor,
    dx: torch.Tensor,
    problem: ConstrainedStochasticProblem,
    xi,
    beta_t: float,
    eta_t: float,
    kappa: float = 0.3,
    nu: float = 1.0,
    lipschitz_est: float = 10.0,
) -> float:
    """
    Select adaptive stepsize by checking an Armijo-type condition on the
    penalized objective phi_nu(x; xi) = nu * F(x; xi) + ||c(x)|| (Appendix A).

    Computes the threshold alpha_thres and projects onto [beta_t, eta_t].

    Parameters
    ----------
    x           : current primal iterate
    lam         : current dual iterate
    dx          : primal component of Newton direction
    problem     : problem instance
    xi          : current sample
    beta_t      : lower bound for stepsize
    eta_t       : upper bound for stepsize
    kappa       : Armijo fraction in (0,1)
    nu          : penalty parameter for the merit function
    lipschitz_est : estimate of local Lipschitz constant

    Returns
    -------
    alpha : selected stepsize in [beta_t, eta_t]
    """
    g_bar = problem.grad(x, xi)
    c_curr = problem.constraint(x)
    G_curr = problem.jac(x)

    delta_phi_loc = nu * (g_bar * dx).sum().item() + torch.norm(c_curr + G_curr @ dx).item() - torch.norm(c_curr).item()
    dx_norm_sq = (dx * dx).sum().item()

    if abs(dx_norm_sq) < 1e-14 or abs(lipschitz_est) < 1e-14:
        return beta_t

    alpha_thres = (kappa - 1) * delta_phi_loc / (lipschitz_est * dx_norm_sq)
    alpha_thres = min(alpha_thres, 1.0)
    alpha = max(beta_t, min(eta_t, alpha_thres))
    return float(alpha)


# ─── SECTION 5: Main AI-StoSQP Algorithm (Algorithm 1) ───────────────────────

@dataclass
class StoSQPConfig:
    """
    Configuration for the AI-StoSQP method.

    All hyperparameters correspond directly to quantities in the paper.
    """
    tau: int = 40
    """Number of randomized Kaczmarz steps per outer iteration (fixed)."""
    beta_c1: float = 1.0
    """Coefficient for stepsize lower bound: beta_t = c1 / t^c2."""
    beta_c2: float = 0.501
    """Exponent for stepsize lower bound (must be in (0.5, 1])."""
    chi_power: float = 2.0
    """Adaptivity gap: chi_t = beta_t^chi_power (must be > 1.5/beta_c2)."""
    delta_reg: float = 0.1
    """Levenberg-Marquardt regularization delta_t."""
    adaptive_stepsize: bool = True
    """If True, uses Armijo line search; else uses beta_t uniformly."""
    kappa: float = 0.3
    """Armijo fraction in (0,1) for adaptive stepsize."""
    nu_penalty: float = 1.0
    """Penalty parameter nu for the merit function."""


class AIStoSQP:
    """
    Adaptive Inexact Stochastic Sequential Quadratic Programming.

    Implements Algorithm 1 of Na and Mahoney (JMLR 2025) with:
        - Online Hessian averaging (Bt = (1/t) sum_{i=0}^{t-1} H_bar_i + Delta_t)
        - Randomized Kaczmarz sketching solver (tau fixed steps)
        - Adaptive stepsize selection
        - Online plug-in covariance matrix estimator (Theorem 5.10)
        - Confidence interval construction

    Parameters
    ----------
    problem : ConstrainedStochasticProblem instance
    x0      : initial primal iterate in R^d
    lam0    : initial dual iterate in R^m
    config  : StoSQPConfig hyperparameters
    """
    def __init__(
        self,
        problem: ConstrainedStochasticProblem,
        x0: torch.Tensor,
        lam0: torch.Tensor,
        config: StoSQPConfig = None,
    ):
        self.problem = problem
        self.x = x0.clone().float()
        self.lam = lam0.clone().float()
        self.config = config or StoSQPConfig()

        self.d = x0.shape[0]
        self.m = lam0.shape[0]
        self.t = 0

        # Online Hessian average (sum over iterations, divided by t)
        self.hess_sum = torch.zeros(self.d, self.d)
        self._Bt = torch.eye(self.d)  # B_0 = I

        # Running statistics for plug-in covariance estimator
        self.grad_sum = torch.zeros(self.d)
        self.grad_sq_sum = torch.zeros(self.d, self.d)

        # Trajectory storage
        self.x_history: List[torch.Tensor] = [x0.clone()]
        self.lam_history: List[torch.Tensor] = [lam0.clone()]
        self.alpha_history: List[float] = []
        self.kkt_residuals: List[float] = []

        self.rng = np.random.default_rng(42)

    def _beta_t(self, t: int) -> float:
        """Compute lower bound beta_t = c1 / (t+1)^c2."""
        return self.config.beta_c1 / (t + 1) ** self.config.beta_c2

    def _eta_t(self, t: int) -> float:
        """Compute upper bound eta_t = beta_t + chi_t."""
        beta = self._beta_t(t)
        chi = beta ** self.config.chi_power
        return beta + chi

    def _build_rhs(
        self,
        g_bar: torch.Tensor,
        c: torch.Tensor,
        G: torch.Tensor,
    ) -> torch.Tensor:
        """Build RHS vector [nabla_x L_bar; c] for the Newton system."""
        nabla_xL = g_bar + G.T @ self.lam
        rhs = torch.cat([nabla_xL, c])
        return rhs

    def step(self) -> float:
        """
        Execute one iteration of AI-StoSQP (Algorithm 1).

        Returns
        -------
        kkt_res : KKT residual norm at current iterate
        """
        t = self.t
        cfg = self.config

        # --- Step 1: Sample and estimate gradient / Hessian ---
        xi = self.problem.sample_xi()
        g_bar = self.problem.grad(self.x, xi)    # stochastic gradient
        H_bar = self.problem.hess(self.x, xi)    # stochastic Hessian
        c = self.problem.constraint(self.x)       # c(x_t)
        G = self.problem.jac(self.x)            # G(x_t) = nabla c(x_t)

        # Update running gradient statistics for covariance estimator
        self.grad_sum += g_bar
        self.grad_sq_sum += torch.outer(g_bar, g_bar)

        # --- Step 2: Update regularized Hessian average ---
        # Bt = (1/t) * sum_{i=0}^{t-1} H_bar_i + Delta_t
        # Note: H_bar_t is used in iteration t+1 (Remark 3.1)
        self.hess_sum += H_bar + self.problem.constraint_hess(self.x, self.lam)
        if t > 0:
            B_avg = self.hess_sum / t
        else:
            B_avg = torch.eye(self.d, dtype=self.x.dtype)
        self._Bt = B_avg + cfg.delta_reg * torch.eye(self.d, dtype=self.x.dtype)

        # --- Step 3: Assemble Newton system and solve inexactly ---
        K = build_kkt_matrix(self._Bt, G, delta=0.0)
        rhs = self._build_rhs(g_bar, c, G)
        z_approx = sketch_solve(K, rhs, cfg.tau, self.rng)

        # Split into primal and dual components
        dx = z_approx[: self.d]
        dlam = z_approx[self.d :]

        # --- Step 4: Select adaptive stepsize ---
        beta_t = self._beta_t(t)
        eta_t = self._eta_t(t)
        if cfg.adaptive_stepsize:
            alpha = select_adaptive_stepsize(
                self.x, self.lam, dx, self.problem, xi,
                beta_t, eta_t, cfg.kappa, cfg.nu_penalty,
            )
        else:
            alpha = float(np.random.uniform(beta_t, eta_t))

        # --- Step 5: Update primal-dual iterate ---
        self.x = self.x + alpha * dx
        self.lam = self.lam + alpha * dlam

        # --- Record diagnostics ---
        nabla_xL = self.problem.grad(self.x, xi) + G.T @ self.lam
        kkt_res = float((torch.norm(nabla_xL) ** 2 + torch.norm(c) ** 2) ** 0.5)
        self.kkt_residuals.append(kkt_res)
        self.alpha_history.append(alpha)
        self.x_history.append(self.x.clone())
        self.lam_history.append(self.lam.clone())

        self.t += 1
        return kkt_res

    def run(self, num_iterations: int, verbose_every: int = 10000) -> None:
        """
        Run AI-StoSQP for num_iterations outer steps.

        Parameters
        ----------
        num_iterations : total number of primal-dual updates T
        verbose_every  : print KKT residual every this many steps
        """
        for t in range(num_iterations):
            kkt = self.step()
            if verbose_every > 0 and (t + 1) % verbose_every == 0:
                print(f"  Iter {t+1:>7d} | KKT residual: {kkt:.6f} | alpha: {self.alpha_history[-1]:.6f}")


# ─── SECTION 6: Plug-In Covariance Estimator (Theorem 5.10) ──────────────────

class OnlineCovarianceEstimator:
    """
    Online plug-in estimator for the limiting covariance Xi_star.

    Implements Eq. 5.14 from Theorem 5.10:

        Omega_t = K_t^{-1} [[SampleCov(g_bar_0,...,g_bar_{t-1}), 0], [0, 0]] K_t^{-1}
        Xi_t    = Omega_t / (2 + beta / beta_tilde)

    where beta = -c2 and beta_tilde = c1 (for beta_t = c1 / t^c2 with c2 = 1).

    The sketch-related correction is O(rho^tau) and is ignored per Corollary 5.7.

    Parameters
    ----------
    d    : primal dimension
    m    : number of equality constraints
    cfg  : StoSQPConfig for stepsize parameters
    """
    def __init__(self, d: int, m: int, cfg: StoSQPConfig):
        self.d = d
        self.m = m
        self.cfg = cfg
        self.n = d + m

        # Compute beta/beta_tilde for the Lyapunov equation denominator
        # For beta_t = c1/t^c2: beta = -c2, beta_tilde = c1 if c2=1 else infty
        if abs(cfg.beta_c2 - 1.0) < 1e-8:
            self.denom_factor = 2.0 + (-cfg.beta_c2) / cfg.beta_c1
        else:
            # beta_tilde = infty, so beta/beta_tilde -> 0
            self.denom_factor = 2.0

    def compute(
        self,
        solver: AIStoSQP,
    ) -> torch.Tensor:
        """
        Compute Xi_t given the current state of a StoSQP solver.

        Parameters
        ----------
        solver : AIStoSQP instance after at least one iteration

        Returns
        -------
        Xi_t : (d+m, d+m) estimated limiting covariance matrix
        """
        t = solver.t
        if t < 2:
            return torch.eye(self.n) * 1e6  # degenerate before data accumulates

        d, m = self.d, self.m

        # --- Compute sample covariance of stochastic gradients ---
        grad_mean = solver.grad_sum / t
        grad_sq_mean = solver.grad_sq_sum / t
        sample_cov_g = grad_sq_mean - torch.outer(grad_mean, grad_mean)

        # --- Assemble Omega_t (Eq. 5.14) ---
        # The (d+m) x (d+m) block: only the (d x d) top-left is nonzero
        Sigma_block = torch.zeros(d + m, d + m)
        Sigma_block[:d, :d] = sample_cov_g

        # --- Get K_t (Lagrangian Hessian block matrix at current iterate) ---
        G_curr = solver.problem.jac(solver.x)
        K_curr = build_kkt_matrix(solver._Bt, G_curr, delta=0.0)

        # --- Omega_t = K_t^{-1} Sigma_block K_t^{-1} ---
        try:
            K_inv = torch.linalg.inv(K_curr + 1e-8 * torch.eye(d + m))
        except RuntimeError:
            K_inv = torch.linalg.pinv(K_curr)
        Omega_t = K_inv @ Sigma_block @ K_inv

        # --- Xi_t = Omega_t / (2 + beta/beta_tilde) ---
        Xi_t = Omega_t / max(self.denom_factor, 1e-6)
        return Xi_t


# ─── SECTION 7: Confidence Interval Construction ─────────────────────────────

def build_confidence_intervals(
    solver: AIStoSQP,
    cov_estimator: OnlineCovarianceEstimator,
    alpha_nominal: float = 0.05,
) -> Dict[str, torch.Tensor]:
    """
    Construct marginal confidence intervals for each component of (x*, lambda*).

    Uses the asymptotic normality result of Theorem 5.6:
        sqrt(1/alpha_bar_t) * (x_t - x*, lambda_t - lambda*) -> N(0, Xi*)

    The (1 - alpha)-CI for the i-th component is:
        [theta_t_i - z_{1-alpha/2} * sqrt(Xi_t[i,i] * alpha_bar_t),
         theta_t_i + z_{1-alpha/2} * sqrt(Xi_t[i,i] * alpha_bar_t)]

    Parameters
    ----------
    solver         : AIStoSQP instance after convergence
    cov_estimator  : OnlineCovarianceEstimator instance
    alpha_nominal  : nominal significance level (default 0.05 for 95% CI)

    Returns
    -------
    dict with keys:
        'x_lower'   : (d,) lower bounds for x components
        'x_upper'   : (d,) upper bounds for x components
        'lam_lower' : (m,) lower bounds for lambda components
        'lam_upper' : (m,) upper bounds for lambda components
        'Xi_t'      : (d+m, d+m) estimated covariance matrix
        'alpha_bar_t': current stepsize (scaling factor)
    """
    from scipy.stats import norm as scipy_norm
    z_crit = scipy_norm.ppf(1.0 - alpha_nominal / 2)

    Xi_t = cov_estimator.compute(solver)
    alpha_bar_t = solver.alpha_history[-1] if solver.alpha_history else 1e-4

    diag_var = torch.clamp(torch.diag(Xi_t) * alpha_bar_t, min=0.0)
    half_widths = z_crit * torch.sqrt(diag_var)

    theta = torch.cat([solver.x, solver.lam])
    lower = theta - half_widths
    upper = theta + half_widths

    d = solver.d
    return {
        'x_lower': lower[:d],
        'x_upper': upper[:d],
        'lam_lower': lower[d:],
        'lam_upper': upper[d:],
        'Xi_t': Xi_t,
        'alpha_bar_t': alpha_bar_t,
    }


# ─── SECTION 8: Example Problem — Linearly Constrained Linear Regression ──────

class LinearRegressionLinearConstraint(ConstrainedStochasticProblem):
    """
    Example 1 from the paper: constrained linear regression.

        F(x; xi) = 0.5 * (xi_a^T x - xi_b)^2       (squared loss)
        c(x) = A x - b = 0                           (linear constraints)

    The true parameter x* satisfies x* = argmin_x E[F(x; xi)]
    subject to A x* = b.

    Parameters
    ----------
    d        : parameter dimension
    m        : number of linear constraints
    x_star   : true parameter (for smoke test ground truth)
    noise_sd : standard deviation of observation noise
    rng      : numpy random generator
    """
    def __init__(
        self,
        d: int,
        m: int,
        x_star: torch.Tensor,
        noise_sd: float = 1.0,
        rng: np.random.Generator = None,
    ):
        self.d = d
        self.m = m
        self.x_star = x_star.float()
        self.noise_sd = noise_sd
        self.rng = rng or np.random.default_rng(0)

        # Generate constraint matrix A and right-hand side b = A x*
        A_np = self.rng.standard_normal((m, d))
        self.A = torch.tensor(A_np, dtype=torch.float32)
        self.b = self.A @ x_star

        # Feature covariance: identity for simplicity
        self.Sigma = torch.eye(d)

    def sample_xi(self):
        """Draw one sample (xi_a, xi_b) ~ P."""
        xi_a = torch.tensor(self.rng.multivariate_normal(np.zeros(self.d), np.eye(self.d)), dtype=torch.float32)
        noise = self.noise_sd * self.rng.standard_normal()
        xi_b = float((xi_a @ self.x_star).item() + noise)
        return (xi_a, xi_b)

    def loss(self, x, xi):
        xi_a, xi_b = xi
        return 0.5 * (xi_a @ x - xi_b) ** 2

    def grad(self, x, xi):
        """nabla_x F(x; xi) = (xi_a^T x - xi_b) xi_a."""
        xi_a, xi_b = xi
        residual = (xi_a @ x - xi_b).item()
        return residual * xi_a

    def hess(self, x, xi):
        """nabla^2_x F(x; xi) = xi_a xi_a^T."""
        xi_a, _ = xi
        return torch.outer(xi_a, xi_a)

    def constraint(self, x):
        """c(x) = A x - b."""
        return self.A @ x - self.b

    def jac(self, x):
        """G(x) = A (constant Jacobian for linear constraints)."""
        return self.A

    def constraint_hess(self, x, lam):
        """sum_i lam_i nabla^2 c_i(x) = 0 (linear constraints)."""
        return torch.zeros(self.d, self.d)


# ─── SECTION 9: Smoke Test — Constrained Linear Regression ───────────────────

def compute_coverage(
    x_true: torch.Tensor,
    lam_true: torch.Tensor,
    ci_result: Dict,
) -> Tuple[float, float]:
    """
    Compute empirical coverage rates for x* and lambda*.

    Parameters
    ----------
    x_true    : true primal solution
    lam_true  : true dual solution
    ci_result : output from build_confidence_intervals

    Returns
    -------
    cov_x   : fraction of x* components covered by their CIs
    cov_lam : fraction of lambda* components covered by their CIs
    """
    covered_x = ((x_true >= ci_result['x_lower']) & (x_true <= ci_result['x_upper'])).float().mean().item()
    covered_lam = ((lam_true >= ci_result['lam_lower']) & (lam_true <= ci_result['lam_upper'])).float().mean().item()
    return covered_x, covered_lam


def _smoke_test():
    """
    Full end-to-end smoke test of AI-StoSQP on linearly constrained linear regression.

    Reproduces the experimental setup of Section 6.2 of the paper:
        - d = 20 primal dimensions, m = 4 equality constraints
        - 100,000 StoSQP iterations with tau = 40 Kaczmarz steps
        - 200 independent runs to measure coverage
        - Reports MAE, average coverage, and average CI length
    """
    print("=" * 68)
    print("AI-StoSQP Smoke Test — Linearly Constrained Linear Regression")
    print("Reproducing Section 6.2 of Na and Mahoney (JMLR 2025)")
    print("=" * 68)

    # Problem setup
    d, m = 20, 4
    NUM_ITERS = 100_000
    NUM_RUNS = 200
    ALPHA_NOMINAL = 0.05

    rng_main = np.random.default_rng(123)
    x_star_np = np.linspace(0, 1, d)
    x_star = torch.tensor(x_star_np, dtype=torch.float32)

    cfg = StoSQPConfig(
        tau=40,
        beta_c1=1.0,
        beta_c2=0.501,
        chi_power=2.0,
        delta_reg=0.1,
        adaptive_stepsize=False,  # uniform random for reproducibility
    )

    print(f"\nProblem: d={d}, m={m}")
    print(f"Config: tau={cfg.tau}, beta_c1={cfg.beta_c1}, beta_c2={cfg.beta_c2}")
    print(f"Iterations per run: {NUM_ITERS:,} | Runs: {NUM_RUNS}")
    print("\nRunning...")

    mae_list, cov_x_list, cov_lam_list, len_x_list, len_lam_list = [], [], [], [], []

    for run in range(NUM_RUNS):
        run_rng = np.random.default_rng(run)
        problem = LinearRegressionLinearConstraint(d, m, x_star, noise_sd=1.0, rng=run_rng)

        # Initial iterate
        x0 = torch.zeros(d)
        # Compute feasible lambda0 from KKT at x0 (approximate)
        G0 = problem.A
        c0 = problem.constraint(x0)
        lam0 = torch.zeros(m)

        solver = AIStoSQP(problem, x0, lam0, config=cfg)
        solver.run(NUM_ITERS, verbose_every=0)

        # Estimate true dual solution lambda* = (A B_oo^{-1} A^T)^{-1} A B_oo^{-1} nabla_f(x*)
        # For linear regression E[H] = I, so B_oo = I and lambda* from KKT: nabla_f(x*) + A^T lam* = 0
        # nabla_f(x*) = E[xi_a(xi_a^T x* - xi_b)] = 0 at optimum, so lam* = 0 for unconstrained minimum = feasible x*
        nabla_f_star = (problem.Sigma @ x_star)  # = x_star for unit covariance
        lam_star_np = torch.linalg.lstsq(G0 @ G0.T, G0 @ nabla_f_star, driver='gelsd').solution
        lam_star = lam_star_np.float()

        mae = float(torch.norm(solver.x - x_star).item())
        mae_list.append(mae)

        cov_est = OnlineCovarianceEstimator(d, m, cfg)
        ci_res = build_confidence_intervals(solver, cov_est, ALPHA_NOMINAL)

        cov_x, cov_lam = compute_coverage(x_star, lam_star, ci_res)
        cov_x_list.append(cov_x)
        cov_lam_list.append(cov_lam)

        if ci_res['x_upper'] is not None:
            len_x = float((ci_res['x_upper'] - ci_res['x_lower']).mean().item())
            len_lam = float((ci_res['lam_upper'] - ci_res['lam_lower']).mean().item())
            len_x_list.append(len_x)
            len_lam_list.append(len_lam)

        if (run + 1) % 50 == 0:
            print(f"  Completed {run+1}/{NUM_RUNS} runs...")

    # --- Summary ---
    print("\n" + "=" * 68)
    print("RESULTS SUMMARY")
    print("-" * 68)
    print(f"MAE (x):             {np.mean(mae_list)*100:.2f} (±{np.std(mae_list)*100:.2f}) x10^-2")
    print(f"Coverage (x*):       {np.mean(cov_x_list)*100:.2f}%  (nominal {(1-ALPHA_NOMINAL)*100:.0f}%)")
    print(f"Coverage (lambda*):  {np.mean(cov_lam_list)*100:.2f}%  (nominal {(1-ALPHA_NOMINAL)*100:.0f}%)")
    if len_x_list:
        print(f"Avg CI length (x):   {np.mean(len_x_list)*100:.2f} x10^-2")
        print(f"Avg CI length (lam): {np.mean(len_lam_list)*100:.2f} x10^-2")
    print("=" * 68)
    print("Smoke test complete. Coverage should be near 95%.")
    print("See Theorem 5.6 in Na and Mahoney (JMLR 2025) for the proof.")


if __name__ == '__main__':
    _smoke_test()