Merge pull request #64 from CurtisTechSolutions/feat/python-gymnasium-env

feat(python): support Gymnasium training with user policies

Author: nosam1998

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"]
+dependencies = ["numpy>=1.24,<3", "gymnasium>=1.2,<2"]
 
 [project.optional-dependencies]
 test = ["pytest>=8,<9"]

python/README.md

@@ -22,4 +22,23 @@ This unit exposes the native bridge. Gymnasium/AEC convenience wrappers and para
 
 The new PyO3 and rust-numpy dependencies provide the CPython ABI boundary and owned NumPy buffers; no hand-written unsafe code is needed. Packaging follows [Maturin's mixed project layout](https://www.maturin.rs/project_layout.html). Native calls use [PyO3 interpreter detachment](https://pyo3.rs/v0.29.2/parallelism), and arrays use [rust-numpy ownership transfer](https://docs.rs/numpy/0.29.0/numpy/array/struct.PyArray.html).
 
-PyO3's target detection uses target-lexicon, licensed under Apache-2.0 with the LLVM exception; the dependency allowlist records that exact expression.
+## Gymnasium and your policy
+
+```python
+import numpy as np
+import gamesforai
+
+env = gamesforai.make("connect4")
+obs, info = env.reset(seed=42)
+while True:
+    action = your_policy(obs, info["action_mask"])
+    obs, reward, terminated, truncated, info = env.step(action)
+    if terminated or truncated:
+        break
+```
+
+The default opponent samples legal actions from the environment's seeded random generator. Select `seat=1` to train from the second seat; the opponent plays its opening during reset. Supply `opponent=my_policy` to play against a Python callable accepting `(observation, info)` and returning a discrete index. Each opponent receives its own seat's view.
+
+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.

python/gamesforai/__init__.py

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

python/gamesforai/gym_env.py

@@ -0,0 +1,175 @@
+"""Gymnasium interface for a policy controlling one seat."""
+from __future__ import annotations
+
+import copy
+import json
+from collections.abc import Callable
+from typing import Any
+
+import gymnasium as gym
+import numpy as np
+from gymnasium import spaces
+from gymnasium.error import ResetNeeded
+
+from ._native import NativeEnv
+
+Policy = Callable[[np.ndarray, dict[str, Any]], int]
+
+
+class GameEnv(gym.Env):
+    """One learner seat against random opponents or a caller-supplied policy.
+
+    A callable opponent receives only its own observation and info (including
+    the legal mask). Its external model state remains the caller's responsibility.
+    """
+
+    metadata = {"render_modes": ["ansi"]}
+
+    def __init__(
+        self,
+        game: str,
+        *,
+        config: dict[str, Any] | None = None,
+        seat: int = 0,
+        opponent: str | Policy = "random",
+        render_mode: str | None = None,
+    ) -> None:
+        if render_mode not in (None, "ansi"):
+            raise ValueError("render_mode must be None or 'ansi'")
+        if opponent != "random" and not callable(opponent):
+            raise ValueError("opponent must be 'random' or a Python policy callable")
+        self.native = NativeEnv(game, json.dumps(config or {}))
+        if isinstance(seat, bool) or not isinstance(seat, int) or not 0 <= seat < self.native.num_players:
+            raise ValueError("seat is outside this game's player range")
+        self.seat = seat
+        self.opponent = opponent
+        self.render_mode = render_mode
+        frame = self.native.frame(seat)
+        self.action_space = spaces.Discrete(self.native.action_space_size)
+        self.observation_space = spaces.Box(
+            -np.inf, np.inf, shape=frame["observation"].shape, dtype=np.float32
+        )
+        self._needs_reset = True
+
+    def _frame(self, seat: int) -> tuple[np.ndarray, dict[str, Any]]:
+        frame = self.native.frame(seat)
+        terminated, truncated = self.native.flags()
+        info = {
+            "action_mask": frame["action_mask"],
+            "text": frame["text"],
+            "board": json.loads(frame["board_json"]),
+            "seat": seat,
+            "to_act": self.native.current_players(),
+            "terminated": terminated,
+            "truncated": truncated,
+        }
+        return frame["observation"], info
+
+    def _opponents(self) -> float:
+        reward = 0.0
+        while actors := self.native.current_players():
+            if len(actors) != 1:
+                raise ValueError("Gym single-seat play requires sequential turns; use a multi-agent interface")
+            seat = actors[0]
+            if seat == self.seat:
+                break
+            observation, info = self._frame(seat)
+            if callable(self.opponent):
+                action = self.opponent(observation, info)
+                if isinstance(action, (bool, np.bool_)) or not isinstance(action, (int, np.integer)):
+                    raise ValueError("opponent policy must return an integer action index")
+                action = int(action)
+            else:
+                legal = np.flatnonzero(info["action_mask"])
+                if not legal.size:
+                    raise ValueError("active opponent has no legal action")
+                action = int(self.np_random.choice(legal))
+            rewards, _, _ = self.native.step(seat, action)
+            reward += rewards[self.seat]
+        return reward
+
+    def reset(
+        self, *, seed: int | None = None, options: dict[str, Any] | None = None
+    ) -> tuple[np.ndarray, dict[str, Any]]:
+        options = options or {}
+        if set(options) - {"position"}:
+            raise ValueError("supported reset option: position")
+        position = options.get("position")
+        if position is not None and not isinstance(position, str):
+            raise ValueError("position must be a notation string")
+        super().reset(seed=seed)
+        native_seed = int(self.np_random.integers(0, 2**64, dtype=np.uint64)) if seed is None else seed
+        self.native.reset(native_seed, position)
+        self._opponents()
+        self._needs_reset = any(self.native.flags())
+        return self._frame(self.seat)
+
+    def step(self, action: int) -> tuple[np.ndarray, float, bool, bool, dict[str, Any]]:
+        if self._needs_reset:
+            raise ResetNeeded("Call reset before stepping this episode")
+        if isinstance(action, (bool, np.bool_)) or not isinstance(action, (int, np.integer)):
+            raise ValueError("action must be an integer index")
+        # Roll back the entire learner/opponent exchange if a custom policy fails.
+        saved = self.native.clone()
+        random_state = copy.deepcopy(self.np_random.bit_generator.state)
+        try:
+            rewards, _, _ = self.native.step(self.seat, int(action))
+            reward = rewards[self.seat] + self._opponents()
+        except Exception:
+            self.native = saved
+            self.np_random.bit_generator.state = random_state
+            raise
+        terminated, truncated = self.native.flags()
+        self._needs_reset = terminated or truncated
+        observation, info = self._frame(self.seat)
+        return observation, float(reward), terminated, truncated, info
+
+    def render(self) -> str:
+        return self.native.frame(self.seat)["text"]
+
+    def close(self) -> None:
+        pass
+
+    def clone(self) -> GameEnv:
+        """Clone engine and RNG state; an external policy callable is shared."""
+        result = copy.copy(self)
+        result.native = self.native.clone()
+        result._np_random = copy.deepcopy(self.np_random)
+        result._np_random_seed = self.np_random_seed
+        result.action_space = copy.deepcopy(self.action_space)
+        result.observation_space = copy.deepcopy(self.observation_space)
+        return result
+
+    def get_state(self) -> dict[str, Any]:
+        """Trusted checkpoint containing private engine state and environment RNG."""
+        return {
+            "version": 1,
+            "seat": self.seat,
+            "native": json.loads(self.native.get_state()),
+            "rng": copy.deepcopy(self.np_random.bit_generator.state),
+            "rng_seed": self.np_random_seed,
+            "needs_reset": self._needs_reset,
+        }
+
+    def set_state(self, state: dict[str, Any]) -> None:
+        """Atomically restore this wrapper; external policy state is not included."""
+        required = {"version", "seat", "native", "rng", "rng_seed", "needs_reset"}
+        if set(state) != required or state["version"] != 1 or state["seat"] != self.seat:
+            raise ValueError("incompatible environment checkpoint")
+        if not isinstance(state["needs_reset"], bool):
+            raise ValueError("invalid checkpoint episode flag")
+        candidate = self.native.clone()
+        candidate.set_state(json.dumps(state["native"]))
+        if any(candidate.flags()) and not state["needs_reset"]:
+            raise ValueError("finished checkpoint cannot be marked active")
+        generator = np.random.default_rng()
+        generator.bit_generator.state = copy.deepcopy(state["rng"])
+        self.native = candidate
+        self._np_random = generator
+        self._np_random_seed = state["rng_seed"]
+        self._needs_reset = state["needs_reset"]
+
+
+def make(game: str, **kwargs: Any) -> GameEnv:
+    """Create a Gymnasium environment for one model's seat."""
+    return GameEnv(game, **kwargs)

python/tests/test_gym.py

@@ -0,0 +1,107 @@
+import copy
+import numpy as np
+import pytest
+from gymnasium.error import ResetNeeded
+from gymnasium.utils.env_checker import check_env
+
+from gamesforai import make
+
+
+@pytest.mark.parametrize("game", ["tictactoe", "connect4"])
+def test_gymnasium_contract(game):
+    check_env(make(game, render_mode="ansi"), skip_render_check=False)
+
+
+@pytest.mark.parametrize("game,config", [
+    ("tictactoe", {}),
+    ("connect4", {}),
+    ("chess", {"max_plies": 12}),
+    ("sudoku", {"size": 4, "max_moves": 8, "allow_notes": False}),
+])
+def test_seeded_masked_episodes_and_shapes(game, config):
+    first = make(game, config=config)
+    second = make(game, config=config)
+    obs, info = first.reset(seed=7)
+    other, other_info = second.reset(seed=7)
+    np.testing.assert_array_equal(obs, other)
+    for _ in range(100):
+        assert first.observation_space.contains(obs)
+        assert info["action_mask"].dtype == np.int8
+        action = int(np.flatnonzero(info["action_mask"])[0])
+        left = first.step(action)
+        right = second.step(action)
+        np.testing.assert_array_equal(left[0], right[0])
+        np.testing.assert_array_equal(left[4]["action_mask"], right[4]["action_mask"])
+        assert left[1:4] == right[1:4]
+        obs, _, terminated, truncated, info = left
+        if terminated or truncated:
+            break
+    else:
+        pytest.fail("bounded episode did not end")
+    with pytest.raises(ResetNeeded):
+        first.step(0)
+
+
+def test_second_seat_receives_opponent_opening():
+    env = make("tictactoe", seat=1)
+    _, info = env.reset(seed=19)
+    assert info["to_act"] == [1]
+    assert info["action_mask"].sum() == 8
+
+
+def test_python_policy_terminal_reward_and_clone_restore():
+    def opponent(observation, info):
+        assert info["seat"] == 1
+        assert observation.dtype == np.float32
+        return int(np.flatnonzero(info["action_mask"])[-1])
+
+    env = make("tictactoe", opponent=opponent)
+    env.reset(seed=9)
+    saved = env.get_state()
+    clone = env.clone()
+    for action in [0, 1]:
+        assert env.step(action)[1:4] == (0.0, False, False)
+    assert env.step(2)[1:4] == (1.0, True, False)
+    assert clone.get_state() == saved
+    env.set_state(saved)
+    for action in [0, 1]:
+        env.step(action)
+    assert env.step(2)[1:4] == (1.0, True, False)
+
+
+def test_random_opponent_and_invalid_policy_checkpoint_atomicity():
+    env = make("connect4")
+    env.reset(seed=81)
+    saved = env.get_state()
+    expected = env.step(3)
+    expected_state = env.get_state()
+    env.set_state(saved)
+    actual = env.step(3)
+    np.testing.assert_array_equal(actual[0], expected[0])
+    assert actual[1:4] == expected[1:4]
+    assert env.get_state() == expected_state
+    invalid = copy.deepcopy(saved)
+    invalid["native"]["engine_version"] = "wrong"
+    with pytest.raises(ValueError):
+        env.set_state(invalid)
+    assert env.get_state() == expected_state
+
+    bad = make("tictactoe", opponent=lambda obs, info: 999)
+    bad.reset(seed=3)
+    before = bad.get_state()
+    with pytest.raises(ValueError):
+        bad.step(0)
+    assert bad.get_state() == before
+
+
+def test_invalid_parameters_and_reset_options_are_explicit():
+    with pytest.raises(ValueError):
+        make("tictactoe", seat=2)
+    env = make("tictactoe")
+    with pytest.raises(ResetNeeded):
+        env.step(0)
+    with pytest.raises(ValueError):
+        env.reset(options={"unsupported": True})
+    env.reset(seed=0)
+    with pytest.raises(ValueError):
+        env.step(True)