nonzero_surrogate.py

"""
NonZero's asinh-GLM return surrogate + NonUCT finite-difference operators.
"""

import torch
import torch.nn as nn


def make_mlp(in_dim, hidden_dims, out_dim):
    layers = []
    prev = in_dim
    for h in hidden_dims:
        layers += [nn.Linear(prev, h), nn.LayerNorm(h), nn.ReLU()]
        prev = h
    layers.append(nn.Linear(prev, out_dim))
    return nn.Sequential(*layers)


class MixingHyperNet(nn.Module):
    def __init__(self, n_agents, n_actions, agent_hidden_dim=128, mixing_hidden=(64, 64)):
        super().__init__()
        self.n_agents = n_agents
        self.n_actions = n_actions
        theta_dim = n_agents * n_actions
        self.net = make_mlp(
            in_dim=n_agents * agent_hidden_dim,
            hidden_dims=list(mixing_hidden),
            out_dim=theta_dim,
        )

    def forward(self, agent_states):
        batch = agent_states.shape[0]
        flat = agent_states.reshape(batch, -1)
        return self.net(flat)


class AsinhGLMSurrogate:
    def __init__(self, n_agents, n_actions, c=1.0, alpha=1.0):
        self.n_agents = n_agents
        self.n_actions = n_actions
        self.nd = n_agents * n_actions
        self.c = c
        self.alpha = alpha

    def one_hot_joint_action(self, agent_actions):
        agent_actions = torch.as_tensor(agent_actions, dtype=torch.long)
        onehots = torch.nn.functional.one_hot(agent_actions, num_classes=self.n_actions)
        return onehots.reshape(*agent_actions.shape[:-1], self.nd).float()

    def eta(self, theta, a):
        z = (theta * a).sum(dim=-1)
        return self.c * torch.asinh(self.alpha * z)

    def neighbor_action(self, a, agent_idx, new_action):
        a = a.clone()
        start = agent_idx * self.n_actions
        a[..., start:start + self.n_actions] = 0.0
        a[..., start + new_action] = 1.0
        return a

    def first_order_diff(self, theta, a, u):
        agent_idx, new_action = u
        a_u = self.neighbor_action(a, agent_idx, new_action)
        return self.eta(theta, a_u) - self.eta(theta, a)

    def second_order_diff(self, theta, a, u, v):
        i, j = u
        k, l = v
        assert i != k, "u and v must act on distinct agents (Section 3.1)"
        a_u = self.neighbor_action(a, i, j)
        a_v = self.neighbor_action(a, k, l)
        a_uv = self.neighbor_action(a_u, k, l)
        return (self.eta(theta, a_uv) - self.eta(theta, a_u)
                - self.eta(theta, a_v) + self.eta(theta, a))


def nonuct_loss(surrogate, theta_pred, theta_star, a, u, v):
    agent_idx_u, new_action_u = u
    a_u_pred = surrogate.neighbor_action(a, agent_idx_u, new_action_u)

    pred = torch.stack([
        surrogate.eta(theta_pred, a),
        surrogate.eta(theta_pred, a_u_pred),
        surrogate.first_order_diff(theta_pred, a, u),
        surrogate.second_order_diff(theta_pred, a, u, v),
    ], dim=-1)

    target = torch.stack([
        surrogate.eta(theta_star, a),
        surrogate.eta(theta_star, a_u_pred),
        surrogate.first_order_diff(theta_star, a, u),
        surrogate.second_order_diff(theta_star, a, u, v),
    ], dim=-1)

    return 0.25 * ((pred - target) ** 2).sum(dim=-1).mean()