Merge pull request #66 from CurtisTechSolutions/feat/python-pettingzoo-aec

feat(python): add PettingZoo AEC self-play

Author: nosam1998

crates/apps/gfa-py/src/lib.rs

@@ -111,6 +111,15 @@ impl NativeEnv {
         Ok(dict)
     }
 
+    fn public_text(&self, py: Python<'_>) -> PyResult<String> {
+        py.detach(|| {
+            self.inner
+                .observe(Viewer::Spectator)
+                .map(|observation| observation.text)
+        })
+        .map_err(error)
+    }
+
     fn clone(&self) -> Self {
         Self {
             inner: self.inner.clone(),

pyproject.toml

@@ -8,7 +8,7 @@ version = "0.1.0"
 description = "Native game environments for model training and evaluation"
 requires-python = ">=3.10"
 license = "GPL-3.0-only"
-dependencies = ["numpy>=1.24,<3", "gymnasium>=1.2,<2"]
+dependencies = ["numpy>=1.24,<3", "gymnasium>=1.2,<2", "pettingzoo>=1.25,<2"]
 
 [project.optional-dependencies]
 test = ["pytest>=8,<9"]

python/README.md

@@ -42,3 +42,22 @@ The default opponent samples legal actions from the environment's seeded random
 Masks are required: a fixed action space includes actions that are illegal in a particular position. Invalid actions raise `ValueError` without committing a learner/opponent exchange. Call `reset` after termination or truncation. Position curricula can use `reset(options={"position": notation})`.
 
 Wrapper checkpoints include native state and the environment RNG, so random-opponent continuations repeat exactly. A cloned wrapper owns independent game, RNG, and action-space state. A custom policy callable and its model state remain caller-owned; checkpoint those separately for reproducible self-play.
+
+## PettingZoo self-play
+
+```python
+import gamesforai
+
+env = gamesforai.aec("connect4")
+env.reset(seed=42)
+for agent in env.agent_iter():
+    obs, reward, terminated, truncated, info = env.last()
+    action = None if terminated or truncated else policies[agent](
+        obs["observation"], obs["action_mask"]
+    )
+    env.step(action)
+```
+
+Agents are named `player_0`, `player_1`, and so on. Each policy sees only its seat's observation. The AEC interface preserves accumulated rewards and gives every agent its final transition before `step(None)` removes it. Numeric spaces and masks stay fixed through an episode. `clone()` copies native and AEC bookkeeping for independent search; rendering uses the public spectator view.
+
+CI runs [PettingZoo's API test](https://pettingzoo.farama.org/content/environment_tests/) against every installed game in addition to full-game reward and truncation tests.

python/gamesforai/__init__.py

@@ -1,5 +1,6 @@
 """Native GamesForAI environments for training and evaluation."""
 from ._native import NativeEnv, games
 from .gym_env import GameEnv, make
+from .aec_env import AECGameEnv, aec
 
-__all__ = ["GameEnv", "NativeEnv", "games", "make"]
+__all__ = ["AECGameEnv", "GameEnv", "NativeEnv", "aec", "games", "make"]

python/gamesforai/_native.pyi

@@ -14,6 +14,7 @@ class NativeEnv:
     def flags(self) -> tuple[bool, bool]: ...
     def returns(self) -> list[float]: ...
     def frame(self, seat: int) -> dict[str, Any]: ...
+    def public_text(self) -> str: ...
     def clone(self) -> NativeEnv: ...
     def get_state(self) -> str: ...
     def set_state(self, snapshot: str) -> None: ...

python/gamesforai/aec_env.py

@@ -0,0 +1,123 @@
+"""PettingZoo AEC interface for independent policies and self-play."""
+from __future__ import annotations
+
+import copy
+import json
+from typing import Any
+
+import numpy as np
+from gymnasium import spaces
+from gymnasium.utils import seeding
+from pettingzoo import AECEnv
+
+from ._native import NativeEnv
+
+
+class AECGameEnv(AECEnv):
+    """Sequential game turns with a separate information set for each agent."""
+
+    metadata = {"render_modes": ["ansi"], "name": "gamesforai_v0", "is_parallelizable": False}
+
+    def __init__(
+        self, game: str, *, config: dict[str, Any] | None = None, render_mode: str | None = None
+    ) -> None:
+        if render_mode not in (None, "ansi"):
+            raise ValueError("render_mode must be None or 'ansi'")
+        self.native = NativeEnv(game, json.dumps(config or {}))
+        self.render_mode = render_mode
+        self.possible_agents = [f"player_{seat}" for seat in range(self.native.num_players)]
+        self._seats = {agent: seat for seat, agent in enumerate(self.possible_agents)}
+        self._actions = {agent: spaces.Discrete(self.native.action_space_size) for agent in self.possible_agents}
+        self._observations = {}
+        for agent, seat in self._seats.items():
+            shape = self.native.frame(seat)["observation"].shape
+            self._observations[agent] = spaces.Dict({
+                "observation": spaces.Box(-np.inf, np.inf, shape=shape, dtype=np.float32),
+                "action_mask": spaces.MultiBinary(self.native.action_space_size),
+            })
+        self.agents = []
+        self._rng, _ = seeding.np_random(None)
+
+    def observation_space(self, agent):
+        return self._observations[agent]
+
+    def action_space(self, agent):
+        return self._actions[agent]
+
+    def observe(self, agent):
+        frame = self.native.frame(self._seats[agent])
+        return {"observation": frame["observation"], "action_mask": frame["action_mask"]}
+
+    def _refresh_infos(self):
+        self.infos = {
+            agent: {"seat": self._seats[agent], "text": self.native.frame(self._seats[agent])["text"]}
+            for agent in self.agents
+        }
+
+    def reset(self, seed=None, options=None):
+        options = options or {}
+        # PettingZoo callers may pass options intended for other wrappers.
+        position = options.get("position")
+        if position is not None and not isinstance(position, str):
+            raise ValueError("position must be a notation string")
+        if seed is not None:
+            self._rng, _ = seeding.np_random(seed)
+        native_seed = int(self._rng.integers(0, 2**64, dtype=np.uint64)) if seed is None else seed
+        self.native.reset(native_seed, position)
+        actors = self.native.current_players()
+        if len(actors) > 1:
+            raise ValueError("AEC adapter currently requires sequential turns")
+        self.agents = self.possible_agents[:]
+        self.rewards = {agent: 0.0 for agent in self.agents}
+        self._cumulative_rewards = self.rewards.copy()
+        terminated, truncated = self.native.flags()
+        self.terminations = {agent: terminated for agent in self.agents}
+        self.truncations = {agent: truncated for agent in self.agents}
+        self.agent_selection = self.possible_agents[actors[0]] if actors else self.agents[0]
+        self._skip_agent_selection = None
+        self._refresh_infos()
+
+    def step(self, action):
+        if not self.agents:
+            raise ValueError("reset before stepping an empty AEC episode")
+        agent = self.agent_selection
+        if self.terminations[agent] or self.truncations[agent]:
+            self._was_dead_step(action)
+            return
+        if isinstance(action, (bool, np.bool_)) or not isinstance(action, (int, np.integer)):
+            raise ValueError("a live agent must choose an integer action index")
+        # Do not change bookkeeping until the native engine accepts the action.
+        rewards, terminated, truncated = self.native.step(self._seats[agent], int(action))
+        self._cumulative_rewards[agent] = 0.0
+        self.rewards = {name: float(rewards[self._seats[name]]) for name in self.agents}
+        self.terminations = {name: terminated for name in self.agents}
+        self.truncations = {name: truncated for name in self.agents}
+        actors = self.native.current_players()
+        if actors:
+            self.agent_selection = self.possible_agents[actors[0]]
+        self._refresh_infos()
+        self._accumulate_rewards()
+        self._deads_step_first()
+
+    def render(self):
+        return self.native.public_text()
+
+    def close(self):
+        pass
+
+    def clone(self):
+        """Independent native and AEC bookkeeping state for trusted local search."""
+        result = copy.copy(self)
+        result.native = self.native.clone()
+        for name in (
+            "agents", "rewards", "_cumulative_rewards", "terminations", "truncations",
+            "infos", "_actions", "_observations", "_rng",
+        ):
+            if hasattr(self, name):
+                setattr(result, name, copy.deepcopy(getattr(self, name)))
+        return result
+
+
+def aec(game: str, **kwargs: Any) -> AECGameEnv:
+    """Create a PettingZoo AEC environment with one policy per game seat."""
+    return AECGameEnv(game, **kwargs)

python/tests/test_aec.py

@@ -0,0 +1,79 @@
+import numpy as np
+import pytest
+from pettingzoo.test import api_test
+
+from gamesforai import aec
+
+
+@pytest.mark.parametrize("game,config", [
+    ("tictactoe", {}),
+    ("connect4", {}),
+    ("chess", {"max_plies": 12}),
+    ("sudoku", {"size": 4, "max_moves": 8, "allow_notes": False}),
+])
+def test_pettingzoo_contract(game, config):
+    api_test(aec(game, config=config), num_cycles=100)
+
+
+def test_final_rewards_are_delivered_to_both_agents_before_removal():
+    env = aec("tictactoe")
+    env.reset(seed=0)
+    for action in [0, 3, 1, 4, 2]:
+        env.step(action)
+    assert env.agents == ["player_0", "player_1"]
+    rewards = {}
+    for agent in env.agent_iter():
+        observation, reward, terminated, truncated, info = env.last()
+        assert terminated and not truncated
+        assert not observation["action_mask"].any()
+        rewards[agent] = reward
+        env.step(None)
+    assert env.agents == []
+    assert rewards == {"player_0": 1.0, "player_1": -1.0}
+
+
+def test_aec_clone_and_invalid_steps_preserve_turn_bookkeeping():
+    env = aec("connect4")
+    env.reset(seed=10)
+    clone = env.clone()
+    before = env.native.get_state()
+    with pytest.raises(ValueError):
+        env.step(99)
+    assert env.native.get_state() == before
+    assert env.agent_selection == "player_0"
+    assert env.rewards == {"player_0": 0.0, "player_1": 0.0}
+    env.step(3)
+    assert clone.native.get_state() == before
+    assert clone.agent_selection == "player_0"
+    assert env.agent_selection == "player_1"
+    assert not env.observe("player_0")["action_mask"].any()
+    assert env.observe("player_1")["action_mask"].any()
+
+
+def test_seeded_self_play_and_cap_endings():
+    first = aec("chess", config={"max_plies": 8})
+    second = aec("chess", config={"max_plies": 8})
+    first.reset(seed=41)
+    second.reset(seed=41)
+    for _ in range(8):
+        assert first.agent_selection == second.agent_selection
+        observation, _, terminated, truncated, _ = first.last()
+        assert not terminated and not truncated
+        other = second.last()[0]
+        np.testing.assert_array_equal(observation["observation"], other["observation"])
+        action = int(np.flatnonzero(observation["action_mask"])[0])
+        first.step(action)
+        second.step(action)
+    assert all(first.truncations.values())
+    assert not any(first.terminations.values())
+    assert first.native.get_state() == second.native.get_state()
+    assert first.render() == first.native.public_text()
+
+
+def test_reset_accepts_framework_options_and_validates_position():
+    env = aec("tictactoe")
+    env.reset(seed=4, options={"options": 1})
+    before = env.native.get_state()
+    with pytest.raises(ValueError, match="notation string"):
+        env.reset(seed=4, options={"position": 12})
+    assert env.native.get_state() == before