repro_bundle/scripts/theory_verification.py

"""Numerical verification of the theory claims (1-3) on Taxi-v3 / CliffWalking.

Claim 1 (Lemma 2): |R - Rhat| <= 2*(extrinsic + intrinsic).
Claim 2 (Theorem 1): extrinsic <= Ls*H*W1(p0d,p0) + H^2*Ls*(Lp+1)*W1(pd,p).
Claim 3 (Theorem 2): intrinsic <= L_PI*H*sqrt(2 log|A|+2) + sum Rademacher + 3H^2*sqrt(log(4H/d)/(2T)).

We use the L1 (one-hot) metric on discrete states: d(s,s') = ||e_s - e_{s'}||_1 = 2*1[s!=s'].
Under this metric, W1 between two distributions over states equals ||p - q||_1.
"""
import numpy as np
from scipy.stats import wasserstein_distance
import sys, os
sys.path.insert(0, "/root/Rationality")
from src.env.taxi import build_experiment_finite_horizon as build_taxi, get_base_P as taxi_P
from src.utils.utils import finite_horizon_optimal_Qh, stochastic_policy_from_Qh, greedy_policy_from_Qh

# -----------------------------------------------------------
# Build a Taxi experiment with a clear training->deployment shift
# training: eps_train slip; deployment: eps_infer slip (different)
# initial-state shift: alpha_d0 mixes d0_train and d0_hard
# -----------------------------------------------------------
np.random.seed(0)
eps_train, eps_infer, alpha_d0, H = 0.25, 0.0, 0.5, 200
exp = build_taxi(eps_infer=eps_infer, eps_train=eps_train, alpha_d0=alpha_d0,
                 max_steps_per_ep=H, seed=0, tau_star=1e-6)

nS, nA = exp.nS, exp.nA
Q_inf = exp.Q_star_infer_h      # Q^{*,dagger}_h  (deployment optimal)
Q_tr  = exp.Q_star_train_h      # Q^*_h           (training optimal)
P_inf = exp.P_infer             # p^dagger
P_tr  = exp.P_train             # p
d0_inf = exp.d0_inf             # p0^dagger
d0_tr  = exp.d0_train           # p0
pi_star = exp.pi_star_h         # softmax optimal policy in deployment

print(f"Taxi: nS={nS} nA={nA} H={H} eps_train={eps_train} eps_infer={eps_infer} alpha_d0={alpha_d0}")
print(f"  d0 shift W1 = {np.abs(d0_inf-d0_tr).sum():.4f}")

# ---- helpers: one-hot L1 metric; W1 = ||p-q||_1 on this metric ----
def w1_states(p, q):
    return float(np.abs(p - q).sum())

def w1_kernels(P_a, P_b):
    """sup_{s,a} W1(p^a(.|s,a), p^b(.|s,a)) under one-hot L1 metric."""
    m = 0.0
    for s in range(nS):
        for a in range(nA):
            pa = np.zeros(nS); pb = np.zeros(nS)
            for (p, s2, r, done) in P_a[s][a]: pa[s2] += p
            for (p, s2, r, done) in P_b[s][a]: pb[s2] += p
            m = max(m, w1_states(pa, pb))
    return m

W1_d0 = w1_states(d0_inf, d0_tr)
W1_p  = w1_kernels(P_inf, P_tr)
print(f"  W1(p0dagger,p0) = {W1_d0:.4f}")
print(f"  W1(pdagger,p)   = {W1_p:.4f}")

# ---- Lipschitz constant Ls of value function (one-hot L1 metric) ----
# value function V_h(s) = max_a Q_h(s,a). d(s,s')=2 for s!=s'.
# Ls = sup_{s!=s'} |V(s)-V(s')| / 2
def lipschitz_V(Qh):
    Ls = 0.0
    for h in range(Qh.shape[0]):
        V = Qh[h].max(axis=1)
        rng = V.max() - V.min()
        # sup over distinct states is at most range / 2; compute exactly
        for s in range(nS):
            for s2 in range(s+1, nS):
                Ls = max(Ls, abs(V[s]-V[s2]) / 2.0)
    return Ls

Ls = lipschitz_V(Q_inf)
print(f"  Ls (Lipschitz of value, one-hot L1) = {Ls:.4f}")

# ---- Lipschitz constant Lp: transition-kernel -> induced state distribution ----
# Assumption 3: W1(D^{pi,dagger}_{h+1}, D^pi_{h+1}) <= W1(p0d,p0) + h*Lp*W1(pd,p)
# Estimate Lp by simulating a fixed policy's state distribution divergence growth.
def induced_state_dist(P, d0, policy_h, nS, nA, H):
    """D^pi_h for a time-indexed stochastic policy (H, nS, nA)."""
    D = np.zeros((H+1, nS))
    D[0] = d0 / d0.sum()
    for h in range(H):
        # next-state distribution
        nxt = np.zeros(nS)
        for s in range(nS):
            for a in range(nA):
                pa = np.zeros(nS)
                for (p, s2, r, done) in P[s][a]:
                    pa[s2] += p
                nxt += D[h, s] * policy_h[h, s, a] * pa
        D[h+1] = nxt
    return D

# use a fixed greedy policy (greedy on deployment Q) for both kernels
pi_greedy = np.zeros((H, nS, nA))
for h in range(H):
    g = greedy_policy_from_Qh(Q_inf[h][None])[0]
    for s in range(nS):
        pi_greedy[h, s, g[s]] = 1.0

D_inf = induced_state_dist(P_inf, d0_inf, pi_greedy, nS, nA, H)
D_tr  = induced_state_dist(P_tr,  d0_tr,  pi_greedy, nS, nA, H)

divs = np.array([w1_states(D_inf[h], D_tr[h]) for h in range(H+1)])
# estimate Lp from per-step growth: div_{h+1} <= div_h + Lp*W1_p  (roughly, after d0 term)
# growth rate = (div_{h+1} - div_h)/W1_p
if W1_p > 1e-12:
    growths = (divs[1:] - divs[:-1]) / W1_p
    Lp = float(np.abs(growths).max())
else:
    Lp = 0.0
print(f"  Lp (empirical per-step divergence growth / W1_p) = {Lp:.4f}")

# -----------------------------------------------------------
# Claim 1: Lemma 2 decomposition
# R(pi) = sum_h E_{D^{pi,dagger}_h} [Q^dagger(s, a_circ) - Q^dagger(s, a_pi)]
# Rhat(pi) = (1/T) sum_t sum_h [Q(s_h^t, a_circ) - Q(s_h^t, a_pi)]
# We use the optimal policy as pi (a_circ = a_pi = optimal => both zero), so instead
# use a *sub-optimal* fixed policy a_pi to get nonzero quantities and check the bound.
# -----------------------------------------------------------
# choose a_pi as a uniformly-random-but-fixed action per state (same across h)
rng = np.random.default_rng(1)
a_pi = rng.integers(0, nA, size=nS)            # fixed suboptimal action per state
a_circ = greedy_policy_from_Qh(Q_inf)[0]        # perfectly-rational action per state (greedy deployment)

# R(pi): expected rational value risk in deployment
R = 0.0
for h in range(H):
    g = Q_inf[h, np.arange(nS), a_circ] - Q_inf[h, np.arange(nS), a_pi]
    R += float(D_inf[h] @ g)

# Extrinsic component (per Lemma 2 inner term): 
# sup_pi | E_{D^pi,dagger_h} Q^dagger(s,a_pi) - E_{D^pi_h} Q(s,a_pi) |
# We evaluate for our fixed a_pi:
extrinsic = 0.0
for h in range(H):
    ed = float(D_inf[h] @ Q_inf[h, np.arange(nS), a_pi])
    et = float(D_tr[h]  @ Q_tr[h,  np.arange(nS), a_pi])
    extrinsic += abs(ed - et)

# Intrinsic component: | E_{D^pi_h} Q(s,a_pi) - (1/T) sum_t Q(s_h^t, a_pi) |
# Simulate T training episodes and average.
T = 200
def sample_episode(P, d0, H, seed):
    r = np.random.default_rng(seed)
    s = int(r.choice(nS, p=d0/d0.sum()))
    states = [s]
    for h in range(H):
        a = int(a_pi[s])
        trans = P[s][a]
        ps = np.array([t[0] for t in trans]); ps = ps/ps.sum()
        idx = int(r.choice(len(trans), p=ps))
        _, s2, rew, done = trans[idx]
        s = int(s2)
        states.append(s)
        if done: break
    return states

emp_terms = []
for t in range(T):
    st = sample_episode(P_tr, d0_tr, H, 1000+t)
    acc = 0.0
    for h in range(min(H, len(st)-1)):
        s = st[h]
        acc += float(Q_tr[h, s, a_circ[s]] - Q_tr[h, s, a_pi[s]])
    emp_terms.append(acc)
Rhat = float(np.mean(emp_terms))   # empirical rational value risk (avg over T episodes)

# intrinsic = | E_{D^pi_h} Q(s,a_pi) - (1/T)sum_t Q(s_h^t, a_pi) |  with a_circ vs a_pi
# (matching Definition: expected uses a_circ-a_pi; empirical uses a_circ-a_pi)
exp_train_risk = 0.0
for h in range(H):
    exp_train_risk += float(D_tr[h] @ (Q_tr[h, np.arange(nS), a_circ] - Q_tr[h, np.arange(nS), a_pi]))
intrinsic = abs(exp_train_risk - Rhat)

gap = abs(R - Rhat)
print("\n=== Claim 1: Lemma 2 decomposition ===")
print(f"  R(pi) [deployment expected rational risk] = {R:.4f}")
print(f"  Rhat(pi) [empirical rational risk, T={T}]  = {Rhat:.4f}")
print(f"  |R - Rhat| = {gap:.4f}")
print(f"  2*(extrinsic + intrinsic) = {2*(extrinsic+intrinsic):.4f}")
print(f"  extrinsic (ERG) = {extrinsic:.4f}")
print(f"  intrinsic (IRG) = {intrinsic:.4f}")
print(f"  Lemma 2 holds: {gap <= 2*(extrinsic+intrinsic) + 1e-9}")

# -----------------------------------------------------------
# Claim 2: Theorem 1 extrinsic bound
# extrinsic <= Ls*H*W1(p0d,p0) + H^2*Ls*(Lp+1)*W1(pd,p)
# (Note: Theorem 1 sums over h with (H-h) factor; the closed form is as above.)
# -----------------------------------------------------------
bound1 = Ls*H*W1_d0 + (H**2)*Ls*(Lp+1)*W1_p
print("\n=== Claim 2: Theorem 1 extrinsic bound ===")
print(f"  extrinsic ERG = {extrinsic:.4f}")
print(f"  Ls*H*W1(p0d,p0) = {Ls*H*W1_d0:.4f}")
print(f"  H^2*Ls*(Lp+1)*W1(pd,p) = {(H**2)*Ls*(Lp+1)*W1_p:.4f}")
print(f"  Theorem 1 bound = {bound1:.4f}")
print(f"  Theorem 1 holds: {extrinsic <= bound1 + 1e-9}")

# -----------------------------------------------------------
# Claim 3: Theorem 2 intrinsic bound
# intrinsic <= L_PI*H*sqrt(2 log|A|+2) + sum_h Rhat_h(QPi) + 3H^2*sqrt(log(4H/d)/(2T))
# We estimate empirical Rademacher complexity of the value-function class Q_Pi.
# Q_Pi: class of Q_h(s,a) functions realized by the DQN MLP. We approximate it with
# the empirical Rademacher complexity of a linear function class over one-hot states
# (a tractable proxy that captures scale), plus check the concentration term directly.
# -----------------------------------------------------------
delta = 0.05
# Concentration term
concentration = 3 * (H**2) * np.sqrt(np.log(4*H/delta) / (2*T))
# Empirical Rademacher complexity of value class: estimate via the realized Q^*_h(s,a_pi)
# values on the sampled training states (Monte Carlo with Rademacher signs).
rademacher_sum = 0.0
for h in range(H):
    # gather state h across episodes
    sh = []
    for t in range(T):
        st = sample_episode(P_tr, d0_tr, H, 1000+t)
        if h < len(st)-1: sh.append(st[h])
    sh = np.array(sh)
    vals = Q_tr[h, sh, a_pi[sh]]   # f(s) = Q_h(s, a_pi(s)) for our fixed policy
    # empirical Rademacher complexity: E_sigma sup |(1/T) sum sigma_t f(s_t)|
    # for a singleton-ish class, this ~ std of vals / sqrt(T)
    nR = 50
    rs = []
    for _ in range(nR):
        sigma = np.sign(np.random.default_rng().standard_normal(len(sh)))
        rs.append(abs((sigma*vals).mean()))
    rademacher_sum += np.mean(rs)

L_PI = Ls  # proxy: Lipschitz of policy->state distribution mapping ~ value Lipschitz
policy_term = L_PI * H * np.sqrt(2*np.log(nA) + 2)
bound2 = policy_term + rademacher_sum + concentration
print("\n=== Claim 3: Theorem 2 intrinsic bound ===")
print(f"  intrinsic IRG = {intrinsic:.4f}")
print(f"  L_PI*H*sqrt(2 log|A|+2) = {policy_term:.4f}")
print(f"  sum Rademacher (proxy) = {rademacher_sum:.4f}")
print(f"  3H^2*sqrt(log(4H/d)/(2T)) = {concentration:.4f}   (T={T}, delta={delta})")
print(f"  Theorem 2 bound = {bound2:.4f}")
print(f"  Theorem 2 holds (proxy): {intrinsic <= bound2 + 1e-6}")

# ---- Show the 1/sqrt(T) concentration decay by varying T ----
print("\n  Concentration term 3H^2*sqrt(log(4H/d)/(2T)) vs T:")
for Tt in [50, 100, 200, 500, 1000, 2000]:
    print(f"    T={Tt:5d}: {3*H**2*np.sqrt(np.log(4*H/delta)/(2*Tt)):.4f}")