matgame_nonzero.py

"""
Generalized N-agent MatGame environment for reproducing NonZero's Claim 4.

Reward rule per NonZero paper, Appendix D ("Settings of Benchmarks"):
  - Linear setting:     r_t = sum(agent action indices)
  - Non-linear setting: r_t = sum(agent action indices) + u + v
                         where u ~ N(0, 2^2), v ~ U(-3, 3)

This intentionally mirrors config/matrix/matgame.py's structure (same
gym.Env shape, same step/reset signature) so it drops into the MAZero
repo's config/matrix/ registration with minimal changes.
"""

import gymnasium as gym
import gymnasium.spaces
import numpy as np


class MatgameNonZeroEnv(gym.Env):
    """
    Generalized N-agent, N-action matrix game with linear or
    linear+noise ("non-linear") joint reward, per NonZero Appendix D.
    """

    def __init__(
            self,
            n_agents=8,
            n_actions=10,
            reward_type="nonlinear",   # "linear" or "nonlinear"
            episode_limit=10,          # NOTE: matches existing matgame.py's
                                        # convention (episode_limit=10 for all
                                        # variants). Table 1's "Steps" column
                                        # (500/1000/2000) is the *training*
                                        # step budget passed to main.py's
                                        # --training_steps flag -- it is NOT
                                        # this environment's episode length.
            noise_std=2.0,             # Gaussian noise std (paper: 2)
            noise_uniform_range=3.0,   # Uniform noise half-range (paper: [-3, 3])
            seed=None,
    ):
        assert reward_type in ("linear", "nonlinear")
        self.n_agents = n_agents
        self.n_actions = n_actions
        self.reward_type = reward_type
        self.episode_limit = episode_limit
        self.noise_std = noise_std
        self.noise_uniform_range = noise_uniform_range

        self._rng = np.random.default_rng(seed)

        self._episode_count = 0
        self._episode_steps = 0
        self._total_steps = 0

        self.observation_space = gym.spaces.Box(
            -np.inf, np.inf, shape=(self.n_agents, 1, 1), dtype=np.float64
        )
        self.action_space = gym.spaces.Discrete(self.n_actions)

    def _joint_reward(self, actions):
        base = float(np.sum(actions))
        if self.reward_type == "linear":
            return base
        u = self._rng.normal(0.0, self.noise_std)
        v = self._rng.uniform(-self.noise_uniform_range, self.noise_uniform_range)
        return base + u + v

    def step(self, actions):
        assert len(actions) == self.n_agents, \
            f"expected {self.n_agents} actions, got {len(actions)}"

        self._total_steps += 1
        self._episode_steps += 1
        info = {}

        state = np.array([self._episode_steps for _ in range(self.n_agents)])
        reward = self._joint_reward(actions)

        terminated = self._episode_steps >= self.episode_limit
        if terminated:
            self._episode_count += 1

        return state, reward, terminated, info

    def reset(self):
        self._episode_steps = 0
        state = np.array([self._episode_steps for _ in range(self.n_agents)])
        return state


TABLE1_CONFIGS = {}
for n_agents, n_actions in [(2, 3), (4, 5), (6, 8), (8, 10)]:
    for reward_type in ["linear", "nonlinear"]:
        key = f"{n_agents}x{n_actions}_{reward_type}"
        TABLE1_CONFIGS[key] = dict(
            env_kwargs=dict(
                n_agents=n_agents,
                n_actions=n_actions,
                reward_type=reward_type,
                episode_limit=10,
            ),
            training_steps_reported=[
                500 if n_agents == 2 else 1000,
                1000 if n_agents == 2 else 2000,
            ],
        )


if __name__ == "__main__":
    cfg = TABLE1_CONFIGS["8x10_nonlinear"]["env_kwargs"]
    env = MatgameNonZeroEnv(**cfg)
    env.reset()
    total = 0.0
    rng = np.random.default_rng(0)
    done = False
    while not done:
        actions = rng.integers(0, env.n_actions, size=env.n_agents)
        _, r, done, _ = env.step(actions)
        total += r
    print(f"Random-policy episode return (8x10, nonlinear, episode_limit=10): {total:.1f}")
    print(f"Max possible per-episode return: {env.n_agents * (env.n_actions - 1) * env.episode_limit}")