Merge pull request #78 from CurtisTechSolutions/feat/aec-training-checkpoints

feat(training): resume AEC self-play checkpoints

Author: nosam1998

examples/aec-checkpoints.md

@@ -0,0 +1,25 @@
+# Resuming self-play
+
+`aec(...).get_state()` returns a JSON-compatible checkpoint. Restore it with
+`set_state(checkpoint)` in an environment with the same game and configuration.
+This preserves the engine, curriculum dataset and RNG, current turn, pending
+cumulative rewards, and terminal agents still waiting for their final `step(None)`.
+Restoration validates the complete candidate before committing it.
+
+```python
+import json
+import gamesforai
+
+env = gamesforai.aec("connect4")
+env.reset(seed=42)
+checkpoint = json.loads(json.dumps(env.get_state()))
+resumed = gamesforai.aec("connect4")
+resumed.set_state(checkpoint)
+```
+
+Checkpoints contain private game state and are for trusted training processes.
+Give each policy only `env.observe(agent)` and its own `env.infos[agent]`.
+Model weights, optimizer state, policy RNGs, and Gym space sampling RNGs are
+owned by the training program and must be saved alongside the environment.
+Rendered information is reconstructed from the engine on restore; custom
+caller-added `infos` fields are not checkpointed.

python/gamesforai/aec_env.py

@@ -3,6 +3,7 @@
 
 import copy
 import json
+import math
 from typing import Any
 
 import numpy as np
@@ -37,6 +38,14 @@ def __init__(
                 "action_mask": spaces.MultiBinary(self.native.action_space_size),
             })
         self.agents = []
+        self.rewards = {}
+        self._cumulative_rewards = {}
+        self.terminations = {}
+        self.truncations = {}
+        self.infos = {}
+        self.agent_selection = None
+        self._skip_agent_selection = None
+        self._initialized = False
         self._position_set = None
         self._position_id = None
         self._rng, _ = seeding.np_random(None)
@@ -74,6 +83,7 @@ def reset(self, seed=None, options=None):
         self.native, self._rng = candidate, rng
         self._position_set, self._position_id = dataset, position_id
         self.agents = self.possible_agents[:]
+        self._initialized = True
         self.rewards = {agent: 0.0 for agent in self.agents}
         self._cumulative_rewards = self.rewards.copy()
         terminated, truncated = self.native.flags()
@@ -111,6 +121,96 @@ def render(self):
     def close(self):
         pass
 
+    def get_state(self) -> dict[str, Any]:
+        """Private, JSON-compatible checkpoint including pending AEC rewards."""
+        return {
+            "version": 1,
+            "native": json.loads(self.native.get_state()),
+            "rng": copy.deepcopy(self._rng.bit_generator.state),
+            "curriculum": curriculum.checkpoint(self._position_set, self._position_id),
+            "initialized": self._initialized,
+            "agents": self.agents[:],
+            "agent_selection": self.agent_selection,
+            "skip_agent_selection": self._skip_agent_selection,
+            "rewards": self.rewards.copy(),
+            "cumulative_rewards": self._cumulative_rewards.copy(),
+            "terminations": self.terminations.copy(),
+            "truncations": self.truncations.copy(),
+        }
+
+    def set_state(self, state: dict[str, Any]) -> None:
+        """Validate a checkpoint completely before replacing any live state."""
+        required = {
+            "version", "native", "rng", "curriculum", "initialized", "agents",
+            "agent_selection", "skip_agent_selection", "rewards",
+            "cumulative_rewards", "terminations", "truncations",
+        }
+        if (
+            not isinstance(state, dict) or set(state) != required
+            or type(state["version"]) is not int or state["version"] != 1
+            or type(state["initialized"]) is not bool
+        ):
+            raise ValueError("incompatible AEC checkpoint")
+        agents = state["agents"]
+        if (
+            not isinstance(agents, list)
+            or any(not isinstance(agent, str) or agent not in self._seats for agent in agents)
+            or agents != [agent for agent in self.possible_agents if agent in agents]
+        ):
+            raise ValueError("invalid AEC checkpoint agents")
+        for field in ("rewards", "cumulative_rewards", "terminations", "truncations"):
+            values = state[field]
+            if not isinstance(values, dict) or set(values) != set(agents):
+                raise ValueError("invalid AEC checkpoint bookkeeping")
+            if field in ("terminations", "truncations"):
+                if any(type(value) is not bool for value in values.values()):
+                    raise ValueError("invalid AEC checkpoint flags")
+            elif any(
+                type(value) not in (int, float) or not math.isfinite(value)
+                for value in values.values()
+            ):
+                raise ValueError("invalid AEC checkpoint rewards")
+        selection, skip = state["agent_selection"], state["skip_agent_selection"]
+        if any(value is not None and value not in self.possible_agents for value in (selection, skip)):
+            raise ValueError("invalid AEC checkpoint turn")
+        candidate = self.native.clone()
+        candidate.set_state(json.dumps(state["native"]))
+        terminated, truncated = candidate.flags()
+        actors = candidate.current_players()
+        if not state["initialized"]:
+            if agents or selection is not None or skip is not None:
+                raise ValueError("uninitialized AEC checkpoint has live bookkeeping")
+        elif not terminated and not truncated:
+            if (
+                agents != self.possible_agents or len(actors) != 1
+                or selection != self.possible_agents[actors[0]] or skip is not None
+            ):
+                raise ValueError("AEC checkpoint turn does not match engine")
+        elif (agents and selection != agents[0]) or (not agents and skip is not None):
+            raise ValueError("invalid AEC checkpoint dead-agent order")
+        if any(value != terminated for value in state["terminations"].values()) or any(
+            value != truncated for value in state["truncations"].values()
+        ):
+            raise ValueError("AEC checkpoint flags do not match engine")
+        generator = np.random.default_rng()
+        generator.bit_generator.state = copy.deepcopy(state["rng"])
+        dataset, position_id = curriculum.restore(state["curriculum"], candidate)
+
+        # Derive information sets from the validated engine, never serialized infos.
+        restored = copy.copy(self)
+        restored.native = candidate
+        restored._rng = generator
+        restored._position_set, restored._position_id = dataset, position_id
+        restored._initialized = state["initialized"]
+        restored.agents = agents[:]
+        restored.agent_selection, restored._skip_agent_selection = selection, skip
+        restored.rewards = state["rewards"].copy()
+        restored._cumulative_rewards = state["cumulative_rewards"].copy()
+        restored.terminations = state["terminations"].copy()
+        restored.truncations = state["truncations"].copy()
+        restored._refresh_infos()
+        self.__dict__.update(restored.__dict__)
+
     def clone(self):
         """Independent native and AEC bookkeeping state for trusted local search."""
         result = copy.copy(self)

python/tests/test_aec_checkpoint.py

@@ -0,0 +1,115 @@
+import copy
+import json
+
+import numpy as np
+import pytest
+
+from gamesforai import aec
+
+
+def assert_same(left, right):
+    assert left.get_state() == right.get_state()
+    assert left.infos == right.infos
+    if left.agents:
+        assert left.agent_selection == right.agent_selection
+        first, second = left.last(), right.last()
+        np.testing.assert_array_equal(first[0]["observation"], second[0]["observation"])
+        np.testing.assert_array_equal(first[0]["action_mask"], second[0]["action_mask"])
+        assert first[1:] == second[1:]
+
+
+def restore(env, game="tictactoe", config=None):
+    other = aec(game, config=config)
+    other.set_state(json.loads(json.dumps(env.get_state())))
+    assert_same(env, other)
+    return other
+
+
+def test_checkpoint_before_reset_and_after_every_live_and_dead_turn():
+    env = aec("tictactoe")
+    other = restore(env)
+    with pytest.raises(ValueError, match="reset"):
+        other.step(0)
+    env.reset(seed=42)
+    for action in [0, 3, 1, 4, 2]:
+        other = restore(env)
+        env.step(action)
+        other.step(action)
+        assert_same(env, other)
+    observed = {}
+    while env.agents:
+        other = restore(env)
+        observed[env.agent_selection] = env.last()[1]
+        env.step(None)
+        other.step(None)
+        assert_same(env, other)
+    assert observed == {"player_0": 1.0, "player_1": -1.0}
+    restore(env)
+
+
+def test_truncated_episode_preserves_final_delivery():
+    env = aec("chess", config={"max_plies": 2})
+    env.reset(seed=5)
+    for _ in range(2):
+        env.step(int(np.flatnonzero(env.last()[0]["action_mask"])[0]))
+    other = restore(env, "chess", {"max_plies": 2})
+    while env.agents:
+        assert env.last()[2:4] == (False, True)
+        env.step(None)
+        other.step(None)
+        assert_same(env, other)
+
+
+def test_curriculum_rng_and_independent_mutable_checkpoint():
+    env = aec("connect4")
+    env.reset(seed=19, options={"position_set": "connect4-solved-positions@1"})
+    other = restore(env, "connect4")
+    for _ in range(20):
+        env.reset()
+        other.reset()
+        assert_same(env, other)
+    state = env.get_state()
+    state["agents"].clear()
+    state["rewards"].clear()
+    state["rng"].clear()
+    assert env.agents and env.rewards and env.get_state()["rng"]
+
+
+@pytest.mark.parametrize("field,value", [
+    ("version", True),
+    ("agents", ["player_0", "player_0"]),
+    ("agents", ["player_1"]),
+    ("agent_selection", "player_1"),
+    ("skip_agent_selection", "player_1"),
+    ("rewards", {"player_0": float("nan"), "player_1": 0.0}),
+    ("cumulative_rewards", {}),
+    ("terminations", {"player_0": True, "player_1": True}),
+    ("truncations", {"player_0": 0, "player_1": 0}),
+    ("rng", {}),
+    ("initialized", False),
+])
+def test_invalid_restore_is_atomic(field, value):
+    env = aec("tictactoe")
+    env.reset(seed=7)
+    before = env.get_state()
+    invalid = copy.deepcopy(before)
+    invalid[field] = value
+    with pytest.raises((ValueError, TypeError)):
+        env.set_state(invalid)
+    assert env.get_state() == before
+
+
+def test_engine_mismatch_and_curriculum_tampering_are_atomic():
+    env = aec("connect4")
+    env.reset(seed=8, options={"position_set": "connect4-solved-positions@1"})
+    before = env.get_state()
+    bad = copy.deepcopy(before)
+    bad["curriculum"]["data"] += " "
+    with pytest.raises(ValueError, match="hash"):
+        env.set_state(bad)
+    assert env.get_state() == before
+    wrong = aec("tictactoe")
+    wrong.reset(seed=8)
+    with pytest.raises(ValueError):
+        env.set_state(wrong.get_state())
+    assert env.get_state() == before