validate_surrogate_fit.py

"""
Validation: can the asinh-GLM surrogate actually fit real environment
rewards from MatgameNonZeroEnv (8 agents, 10 actions, nonlinear)?
"""

import torch
import torch.nn as nn
import numpy as np

from matgame_nonzero import MatgameNonZeroEnv, TABLE1_CONFIGS
from nonzero_surrogate import AsinhGLMSurrogate


def collect_dataset(env, n_samples, rng):
    actions = rng.integers(0, env.n_actions, size=(n_samples, env.n_agents))
    rewards = []
    for a in actions:
        r = env._joint_reward(a)
        rewards.append(r)
    return torch.tensor(actions, dtype=torch.long), torch.tensor(rewards, dtype=torch.float32)


def main():
    torch.manual_seed(0)
    rng = np.random.default_rng(0)

    cfg = TABLE1_CONFIGS["8x10_nonlinear"]["env_kwargs"]
    env = MatgameNonZeroEnv(**cfg)
    surrogate = AsinhGLMSurrogate(env.n_agents, env.n_actions)

    train_actions, train_rewards = collect_dataset(env, 4000, rng)
    test_actions, test_rewards = collect_dataset(env, 1000, rng)

    train_a = surrogate.one_hot_joint_action(train_actions)
    test_a = surrogate.one_hot_joint_action(test_actions)

    theta = nn.Parameter(torch.zeros(env.n_agents * env.n_actions))
    log_c = nn.Parameter(torch.zeros(()))
    log_alpha = nn.Parameter(torch.zeros(()))

    opt = torch.optim.Adam([theta, log_c, log_alpha], lr=0.05)
    sched = torch.optim.lr_scheduler.CosineAnnealingLR(opt, T_max=2000)

    def predict(a_batch):
        c = torch.exp(log_c)
        alpha = torch.exp(log_alpha)
        z = (theta * a_batch).sum(dim=-1)
        return c * torch.asinh(alpha * z)

    n_epochs = 2000
    for epoch in range(n_epochs):
        opt.zero_grad()
        pred = predict(train_a)
        loss = ((pred - train_rewards) ** 2).mean()
        loss.backward()
        opt.step()
        sched.step()
        if epoch % 200 == 0 or epoch == n_epochs - 1:
            with torch.no_grad():
                test_pred = predict(test_a)
                test_mse = ((test_pred - test_rewards) ** 2).mean().item()
                test_r2 = 1 - test_mse / test_rewards.var().item()
            print(f"epoch {epoch:4d}  train_mse={loss.item():7.3f}  "
                  f"test_mse={test_mse:7.3f}  test_R2={test_r2:.4f}")

    print(f"\nFitted c={torch.exp(log_c).item():.4f}  alpha={torch.exp(log_alpha).item():.4f}")
    print(f"True reward std (for reference): {train_rewards.std().item():.3f}")
    print(f"Reward range in data: [{train_rewards.min().item():.2f}, {train_rewards.max().item():.2f}]")


if __name__ == "__main__":
    main()