code/mdp.py

"""Robust average-reward MDP primitives (arXiv 2505.12462, Sec. 3).

Implements the support function sigma_{P_sa} for the (s,a)-rectangular
contamination and l_p-norm uncertainty sets, the penalty term kappa_q (dual
representation, Kumar et al. 2023 / Appendix C), the robust Bellman operator
T_P(Q) = r + sigma_P(max_A Q), and exact solvers for the optimal robust gain
(RRVI span-relative-value-iteration) and per-policy robust gain evaluation.

All operators are EXACT (model-known): they are the ground truth the model-free
RHI algorithm (rhi.py) is measured against, and the objects Lemma 5.1 /
Theorem 5.2 / Theorem 4.2 are stated over.
"""
from __future__ import annotations

import numpy as np
from scipy.optimize import minimize_scalar


def span(v: np.ndarray) -> float:
    """Span semi-norm Sp(v) = max_i v_i - min_i v_i (Eq. 9)."""
    return float(np.max(v) - np.min(v))


def holder_q(p: float) -> float:
    """Hölder conjugate q with 1/p + 1/q = 1."""
    if p == np.inf:
        return 1.0
    if p == 1.0:
        return np.inf
    return p / (p - 1.0)


def kappa_q(h: np.ndarray, p: float, reachable: np.ndarray | None = None) -> float:
    """kappa_q(h) = min_omega || u - omega * 1 ||_q  (Appendix C, eq. near 929).

    u(s) = h(s) * 1(s reachable); q is the Hölder conjugate of p.
    Closed forms: q=1 -> median, q=2 -> mean, q=inf -> midrange.
    """
    if reachable is None:
        u = np.asarray(h, dtype=float)
    else:
        u = np.where(reachable, h, 0.0)
    q = holder_q(p)
    if q == 1.0:
        omega = np.median(u)
    elif q == 2.0:
        omega = float(np.mean(u))
    elif q == np.inf:
        omega = 0.5 * (float(np.max(u)) + float(np.min(u)))
    else:
        # convex 1-D minimization
        lo, hi = float(np.min(u)), float(np.max(u))
        if hi <= lo:
            omega = lo
        else:
            res = minimize_scalar(
                lambda w: np.sum(np.abs(u - w) ** q), bounds=(lo, hi), method="bounded"
            )
            omega = float(res.x)
    diff = u - omega
    if q == np.inf:
        return float(np.max(np.abs(diff)))
    return float(np.sum(np.abs(diff) ** q) ** (1.0 / q))


def sigma_contamination(Prow: np.ndarray, h: np.ndarray, R: float) -> float:
    """sigma for contamination set {(1-R)P + R q : q in Delta} (Eq. 4).

    min_q E[(1-R)P + R q][h] = (1-R) P.h + R min_s h(s).
    """
    return float((1.0 - R) * np.dot(Prow, h) + R * np.min(h))


def sigma_lp(Prow: np.ndarray, h: np.ndarray, R: float, p: float,
             reachable: np.ndarray | None = None) -> float:
    """sigma for l_p-norm ball {q : ||q-P||_p <= R} (Eq. 5, Appendix C eq. 74).

    = P.h - R * kappa_q(h).  Valid when the ball is interior to the simplex.
    """
    return float(np.dot(Prow, h) - R * kappa_q(h, p, reachable))


class Uncertainty:
    """Container for an (s,a)-rectangular uncertainty set."""

    KINDS = ("contamination", "lp")

    def __init__(self, kind: str, radius: float | np.ndarray, p: float = 2.0):
        assert kind in self.KINDS, kind
        self.kind = kind
        self.radius = radius
        self.p = p

    def sigma(self, Prow: np.ndarray, h: np.ndarray, R: float,
              reachable: np.ndarray | None = None) -> float:
        if self.kind == "contamination":
            return sigma_contamination(Prow, h, R)
        return sigma_lp(Prow, h, R, self.p, reachable)


def robust_bellman_T(Q: np.ndarray, r: np.ndarray, P: np.ndarray,
                     unc: Uncertainty) -> np.ndarray:
    """Robust Bellman operator T_P(Q)(s,a) = r(s,a) + sigma_{P_sa}(max_A Q).

    T_P,g(Q) = T_P(Q) - g; the residual T_P(Q) - Q drives Lemma 5.1 / Thm 5.2.
    """
    h = Q.max(axis=1)  # max_A Q  (the V vector)
    return r + _sigma_matrix(P, h, unc)


# ---------------- exact robust solvers (RRVI, Wang et al. 2023a) -------------

def robust_discounted_vi(r: np.ndarray, P: np.ndarray, unc: Uncertainty,
                         gamma: float, iters: int = 5000, tol: float = 1e-12
                         ) -> tuple[np.ndarray, np.ndarray]:
    """Optimal robust *discounted* value/value via value iteration (contractive).

    V(s) = max_a { r(s,a) + gamma sigma_{P_sa}(V) }.  Used by the reduction
    approach (Theorem 4.2) and as a gamma->1 cross-check on the optimal gain.
    """
    S, A = r.shape
    V = np.zeros(S)
    for _ in range(iters):
        Q = r + gamma * _sigma_matrix(P, V, unc)
        Vn = Q.max(axis=1)
        if np.max(np.abs(Vn - V)) < tol:
            V = Vn
            break
        V = Vn
    Q = r + gamma * _sigma_matrix(P, V, unc)
    return V, Q


def _sigma_matrix(P: np.ndarray, h: np.ndarray, unc: Uncertainty) -> np.ndarray:
    """Compute sigma_{P_sa}(h) for all (s,a) -> (S,A) array (vectorized)."""
    S, A = P.shape[:2]
    Ph = np.einsum("sap,p->sa", P, h)  # P[s,a] . h  for all (s,a)
    if unc.kind == "contamination":
        R = unc.radius
        if not np.isscalar(R):
            R = R  # per-(s,a) radius matrix
        return (1.0 - R) * Ph + R * float(np.min(h))
    # lp-norm: sigma = P.h - R * kappa_q(h, s,a).  kappa depends on (s,a) only
    # through the reachable set; for fully-connected MDPs it is a single scalar.
    reach = P > 0  # (S,A,S)
    full = reach.all(axis=2)
    if full.all():
        kappa = kappa_q(h, unc.p)
        R = unc.radius
        return Ph - R * kappa
    # mixed: per-(s,a) reachable sets (slower path)
    out = Ph.copy()
    R = unc.radius
    for s in range(S):
        for a in range(A):
            if not full[s, a]:
                out[s, a] = Ph[s, a] - (R if np.isscalar(R) else R[s, a]) * kappa_q(
                    h, unc.p, reach[s, a])
    return out


def robust_optimal_span_rvi(r: np.ndarray, P: np.ndarray, unc: Uncertainty,
                            iters: int = 20000, tol: float = 1e-11, ref: int = 0
                            ) -> tuple[float, np.ndarray, np.ndarray]:
    """Optimal robust average reward g*, bias h*, greedy policy (RRVI).

    Span-relative value iteration on T(h)(s)=max_a{r(s,a)+sigma(h)} is a
    non-expansion in Sp(.); anchoring to a reference state keeps h bounded and
    g_k = T(h_k)(ref) -> g* under unichain+compact (Assumption 3.1).
    """
    S, A = r.shape
    h = np.zeros(S)
    g = 0.0
    for _ in range(iters):
        Th = (r + _sigma_matrix(P, h, unc)).max(axis=1)
        g = Th[ref]
        hn = Th - g  # subtract reference value -> bounded bias
        if span(hn - h) < tol and np.max(np.abs(hn - h)) < tol:
            h = hn
            break
        h = hn
    Q = r + _sigma_matrix(P, h, unc)
    pi = Q.argmax(axis=1)
    return float(g), h, pi


def robust_policy_gain(pi: np.ndarray, r: np.ndarray, P: np.ndarray,
                       unc: Uncertainty, iters: int = 20000, tol: float = 1e-11,
                       ref: int = 0) -> tuple[float, np.ndarray]:
    """Robust average reward g_P^pi of a FIXED policy + its bias (RRVI eval).

    Solves h = r_pi - g*1 + sigma_{pi}(h) via span-RVE (non-expansion).
    """
    S, A = r.shape
    rpi = r[np.arange(S), pi]
    h = np.zeros(S)
    g = 0.0
    for _ in range(iters):
        sh = _sigma_policy(pi, P, h, unc)
        Th = rpi + sh
        g = Th[ref]
        hn = Th - g
        if span(hn - h) < tol and np.max(np.abs(hn - h)) < tol:
            h = hn
            break
        h = hn
    return float(g), h


def _sigma_policy(pi: np.ndarray, P: np.ndarray, h: np.ndarray,
                  unc: Uncertainty) -> np.ndarray:
    S = len(pi)
    out = np.empty(S)
    for s in range(S):
        a = pi[s]
        Prow = P[s, a]
        R = unc.radius if np.isscalar(unc.radius) else unc.radius[s, a]
        reachable = Prow > 0
        if unc.kind == "lp" and not reachable.all():
            out[s] = unc.sigma(Prow, h, R, reachable)
        else:
            out[s] = unc.sigma(Prow, h, R)
    return out


def robust_bias_span(pi: np.ndarray, r: np.ndarray, P: np.ndarray,
                     unc: Uncertainty) -> tuple[float, np.ndarray]:
    """Robust optimal bias span H = max_P Sp(h_P^{pi*}): here computed under the
    worst-case kernel realised by RRVI evaluation of the greedy policy."""
    _, h = robust_policy_gain(pi, r, P, unc)
    return span(h), h


# ---------------- MDP generation --------------------------------------------

def garnet_mdp(S: int, A: int, density: float, rng: np.random.Generator,
               n_recurrent: int | None = None) -> tuple[np.ndarray, np.ndarray]:
    """Garnet G(S,A) MDP (Archibald et al. 1995) with unichain guarantee.

    Returns P (S,A,S) and r (S,A) in [0,1].  Each (s,a) has ~density*S nonzero
    next-states.  A single recurrent class (states 0..n_rec-1) is enforced so
    every deterministic policy yields a unichain (Assumption 3.1).  Transient
    states (if any) drain into the recurrent class.
    """
    if n_recurrent is None:
        n_recurrent = S
    n_rec = min(n_recurrent, S)
    P = np.zeros((S, A, S))
    for s in range(S):
        for a in range(A):
            if s < n_rec:
                k = max(2, int(round(density * n_rec)))
                targets = rng.choice(n_rec, size=k, replace=False)
                probs = rng.dirichlet(np.ones(k))
                P[s, a, targets] = probs
            else:
                # transient: drain to recurrent class (unichain guarantee)
                P[s, a, rng.integers(0, n_rec)] = 1.0
    r = rng.uniform(0, 1, (S, A))
    # sanity: every row sums to 1
    assert np.allclose(P.sum(axis=2), 1.0)
    return P, r