audit.py

#!/usr/bin/env python3
"""Self-contained numerical audit of QUATRO claims on GPU."""
import torch, json, math

def compute_quatro_advantages(rewards, delta=0.001, num_iterations=50):
    B, N = rewards.shape
    device = rewards.device
    lambdas = torch.zeros(B, device=device)
    mus = torch.zeros(B, device=device)
    for b in range(B):
        r = rewards[b]
        # Handle edge case: all rewards equal -> lambda -> 0, advantages = 0
        if r.std() < 1e-8:
            lambdas[b] = 1e-6
            mus[b] = -1e-6
            continue
        lo, hi = 1e-6, 100.0
        for _ in range(num_iterations):
            mid = (lo + hi) / 2.0
            r_scaled = r / mid
            r_max = r_scaled.max()
            exp_r = torch.exp(r_scaled - r_max)
            mean_exp = exp_r.mean()
            mean_r_exp = (r_scaled * exp_r).mean()
            log_term = r_max + torch.log(mean_exp + 1e-30)
            grad = delta + log_term - mean_r_exp / (mean_exp + 1e-30)
            if grad > 0: hi = mid
            else: lo = mid
        lam = (lo + hi) / 2.0
        lambdas[b] = lam
        r_scaled = r / lam
        r_max = r_scaled.max()
        exp_r = torch.exp(r_scaled - r_max)
        mean_exp = exp_r.mean()
        log_term = r_max + torch.log(mean_exp + 1e-30)
        f_val = lam * (delta + log_term)
        mu = f_val - lam * (delta + 1.0)
        mus[b] = mu
    advantages = (rewards - mus.unsqueeze(1)) / lambdas.unsqueeze(1) - 1.0
    return advantages, lambdas, mus

def quatro_loss(log_probs, old_log_probs, advantages, log_ratio_stabilizer=True):
    ratio = (log_probs - old_log_probs).exp()
    if log_ratio_stabilizer:
        stabilizer = (log_probs - old_log_probs).detach()
        loss = -(ratio * (advantages - stabilizer)).mean()
    else:
        loss = -(ratio * advantages).mean()
    kl = (ratio.log() * ratio - ratio + 1).mean()
    return loss, kl

def gspo_loss(log_probs, old_log_probs, advantages, epsilon=0.2):
    ratio = (log_probs - old_log_probs).exp()
    clipped = torch.clamp(ratio, 1.0 - epsilon, 1.0 + epsilon)
    return -torch.min(ratio * advantages, clipped * advantages).mean()

device = torch.device("cuda")
print(f"Device: {device}  CUDA: {torch.cuda.is_available()}")
if torch.cuda.is_available():
    print(f"GPU: {torch.cuda.get_device_name(0)}")

# === CLAIM 1: Trust-region with Lagrangian dual ===
print("\n" + "="*60)
print("CLAIM 1: Trust-region formulation with Lagrangian dual")
print("="*60)

rewards_cases = {
    "low_var": torch.tensor([[0.5, 0.5, 0.5, 0.5]], dtype=torch.float32),
    "med_var": torch.tensor([[1.0, 0.0, 0.5, 0.0]], dtype=torch.float32),
    "high_var": torch.tensor([[1.0, 0.0, 0.0, 0.0]], dtype=torch.float32),
}

for name, r in rewards_cases.items():
    r = r.to(device)
    for delta in [0.1, 0.01, 0.001]:
        adv, lam, mu = compute_quatro_advantages(r, delta=delta)
        print(f"  {name}, delta={delta}: lambda*={lam[0].item():.4f}, mu*={mu[0].item():.4f}")
        print(f"    A_q^i: {[f'{x:.4f}' for x in adv[0].tolist()]}")

r_low = torch.tensor([[0.5, 0.5, 0.5, 0.5]], dtype=torch.float32, device=device)
r_high = torch.tensor([[1.0, 0.0, 0.0, 0.0]], dtype=torch.float32, device=device)
_, lam_low, _ = compute_quatro_advantages(r_low, delta=0.001)
_, lam_high, _ = compute_quatro_advantages(r_high, delta=0.001)
print(f"\nProperty: Higher variance -> larger lambda* (more conservative)")
print(f"  Low var lambda*: {lam_low[0].item():.4f}")
print(f"  High var lambda*: {lam_high[0].item():.4f}")
print(f"  Verified: {lam_high[0].item() > lam_low[0].item()}")

# === CLAIM 2: QUATRO objective ===
print("\n" + "="*60)
print("CLAIM 2: QUATRO objective with log-ratio stabilizer")
print("="*60)

log_probs = torch.tensor([[-1.0, -2.0, -1.5, -2.5]], dtype=torch.float32, device=device)
old_log_probs = torch.tensor([[-1.2, -1.8, -1.3, -2.0]], dtype=torch.float32, device=device)
rewards = torch.tensor([[1.0, 0.0, 0.5, 0.0]], dtype=torch.float32, device=device)

adv, _, _ = compute_quatro_advantages(rewards, delta=0.001)
loss_with, kl_with = quatro_loss(log_probs, old_log_probs, adv, log_ratio_stabilizer=True)
loss_without, _ = quatro_loss(log_probs, old_log_probs, adv, log_ratio_stabilizer=False)
loss_gspo = gspo_loss(log_probs, old_log_probs, adv.squeeze(0), epsilon=0.2)

print(f"  QUATRO (with stabilizer): {loss_with.item():.6f}")
print(f"  QUATRO (no stabilizer):   {loss_without.item():.6f}")
print(f"  GSPO:                     {loss_gspo.item():.6f}")
print(f"  KL estimate:             {kl_with.item():.6f}")
print(f"  Stabilizer effect:       {(loss_with - loss_without).item():.6f}")

# === CLAIM 5: Ablation ===
print("\n" + "="*60)
print("CLAIM 5: Ablation of log-ratio stabilizer")
print("="*60)

high_adv = torch.tensor([[10.0, -1.0, -1.0, -1.0]], dtype=torch.float32, device=device)
lp_far = torch.tensor([[-0.5, -2.0, -2.0, -2.0]], dtype=torch.float32, device=device)
olp = torch.tensor([[-2.0, -2.0, -2.0, -2.0]], dtype=torch.float32, device=device)

lw, _ = quatro_loss(lp_far, olp, high_adv, log_ratio_stabilizer=True)
lwo, _ = quatro_loss(lp_far, olp, high_adv, log_ratio_stabilizer=False)
ratio_val = (lp_far - olp).exp()[0, 0].item()
print(f"  High-adv sample ratio: {ratio_val:.2f}")
print(f"  Loss WITH stabilizer:    {lw.item():.6f}")
print(f"  Loss WITHOUT stabilizer: {lwo.item():.6f}")
print(f"  Without stabilizer loss is {lwo.item()/max(abs(lw.item()),1e-10):.2f}x larger")

results = {
    "claim1_lambda_property": bool(lam_high[0].item() > lam_low[0].item()),
    "claim1_lambda_low_var": float(lam_low[0].item()),
    "claim1_lambda_high_var": float(lam_high[0].item()),
    "claim2_loss_with_stabilizer": float(loss_with.item()),
    "claim2_loss_without_stabilizer": float(loss_without.item()),
    "claim2_stabilizer_diff": float((loss_with - loss_without).item()),
    "claim2_gspo_loss": float(loss_gspo.item()),
    "claim2_kl": float(kl_with.item()),
    "claim5_ablation_ratio": float(lwo.item() / max(abs(lw.item()), 1e-10)),
}
print(f"\nResults: {json.dumps(results, indent=2)}")
with open("/tmp/audit_results.json", "w") as f:
    json.dump(results, f, indent=2)
print("Saved to /tmp/audit_results.json")