Merge pull request #69 from CurtisTechSolutions/feat/python-position-curricula

feat(training): sample reproducible position curricula in Gymnasium and AEC

Author: nosam1998

python/README.md

@@ -98,3 +98,9 @@ of numeric output per batch. Throughput still needs measurement against the PRD
 ## Readable actions for model policies
 
 Native frames include `legal_actions` as `(index, canonical_notation)` pairs in stable notation order. Gymnasium info includes the same catalog plus the game ID, rules, and action notation for learner and opponent policies. The catalog contains only actions available to that seat and matches its numeric mask. `NativeEnv.spec_json()` returns the static game contract. Vector batches keep their numeric-only output.
+
+## Sampling a curriculum on every reset
+
+Use `env.reset(seed=42, options={"position_set": "chess-endgames-basic@1"})` with either Gymnasium or AEC. Later `reset()` calls keep sampling that pinned set; repeating an explicit seed repeats the selection. Supply a custom `PositionSet` instead of a bundled name, or use `{"position_set": None}` to return to initial positions. An explicit `position` also clears the active curriculum. A set must match the environment game and normalized configuration.
+
+`info["curriculum"]` records only the set ID, content hash, and selected position ID. Expected puzzle answers are excluded. Gymnasium checkpoints include the pinned dataset and RNG continuation; AEC clones preserve both. Invalid selections, imported positions, and failed opponent openings preserve environment and RNG state. A caller-owned opponent model must manage its own mutable state.

python/gamesforai/aec_env.py

@@ -11,6 +11,7 @@
 from pettingzoo import AECEnv
 
 from ._native import NativeEnv
+from . import curriculum
 
 
 class AECGameEnv(AECEnv):
@@ -36,6 +37,8 @@ def __init__(
                 "action_mask": spaces.MultiBinary(self.native.action_space_size),
             })
         self.agents = []
+        self._position_set = None
+        self._position_id = None
         self._rng, _ = seeding.np_random(None)
 
     def observation_space(self, agent):
@@ -50,23 +53,26 @@ def observe(self, agent):
 
     def _refresh_infos(self):
         self.infos = {
-            agent: {"seat": self._seats[agent], "text": self.native.frame(self._seats[agent])["text"]}
+            agent: {"seat": self._seats[agent], "text": self.native.frame(self._seats[agent])["text"],
+                    **curriculum.origin_info(self._position_set, self._position_id)}
             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()
+        rng = copy.deepcopy(self._rng) if seed is None else seeding.np_random(seed)[0]
+        position, dataset, position_id = curriculum.resolve_start(
+            self.native, options, rng, self._position_set
+        )
+        native_seed = int(rng.integers(0, 2**64, dtype=np.uint64)) if seed is None else seed
+        candidate = self.native.clone()
+        candidate.reset(native_seed, position)
+        actors = candidate.current_players()
         if len(actors) > 1:
             raise ValueError("AEC adapter currently requires sequential turns")
+        self.native, self._rng = candidate, rng
+        self._position_set, self._position_id = dataset, position_id
         self.agents = self.possible_agents[:]
         self.rewards = {agent: 0.0 for agent in self.agents}
         self._cumulative_rewards = self.rewards.copy()

python/gamesforai/curriculum.py

@@ -0,0 +1,50 @@
+"""Shared reset selection; puzzle answers never enter policy observations."""
+from .positions import PositionSet
+
+
+def resolve_start(native, options, rng, previous):
+    if "position" in options and "position_set" in options:
+        raise ValueError("choose position or position_set, not both")
+    position = options.get("position")
+    if position is not None and not isinstance(position, str):
+        raise ValueError("position must be a notation string")
+    dataset = None if "position" in options else previous
+    if "position_set" in options:
+        dataset = options["position_set"]
+        if isinstance(dataset, str):
+            dataset = PositionSet.load(dataset)
+        if dataset is not None and not isinstance(dataset, PositionSet):
+            raise ValueError("position_set must be name@version, a PositionSet, or None")
+    if dataset is not None:
+        dataset.validate_for(native)
+        entry = dataset.sample(rng)
+        return entry.position, dataset, entry.id
+    return position, None, None
+
+
+def origin_info(dataset, position_id):
+    if dataset is None:
+        return {}
+    return {"curriculum": {
+        "position_set": dataset.identifier,
+        "sha256": dataset.sha256,
+        "position_id": position_id,
+    }}
+
+
+def checkpoint(dataset, position_id):
+    if dataset is None:
+        return None
+    return {**dataset.checkpoint(), "position_id": position_id}
+
+
+def restore(payload, native):
+    if payload is None:
+        return None, None
+    if not isinstance(payload, dict) or set(payload) != {"manifest", "data", "position_id"}:
+        raise ValueError("invalid curriculum checkpoint")
+    dataset = PositionSet.from_jsonl(payload["manifest"], payload["data"])
+    dataset.validate_for(native)
+    if payload["position_id"] not in {entry.id for entry in dataset.entries}:
+        raise ValueError("curriculum checkpoint position is not in the pinned set")
+    return dataset, payload["position_id"]

python/gamesforai/gym_env.py

@@ -12,6 +12,7 @@
 from gymnasium.error import ResetNeeded
 
 from ._native import NativeEnv
+from . import curriculum
 
 Policy = Callable[[np.ndarray, dict[str, Any]], int]
 
@@ -51,6 +52,8 @@ def __init__(
             -np.inf, np.inf, shape=frame["observation"].shape, dtype=np.float32
         )
         self._needs_reset = True
+        self._position_set = None
+        self._position_id = None
 
     def _frame(self, seat: int) -> tuple[np.ndarray, dict[str, Any]]:
         frame = self.native.frame(seat)
@@ -67,6 +70,7 @@ def _frame(self, seat: int) -> tuple[np.ndarray, dict[str, Any]]:
             "to_act": self.native.current_players(),
             "terminated": terminated,
             "truncated": truncated,
+            **curriculum.origin_info(self._position_set, self._position_id),
         }
         return frame["observation"], info
 
@@ -97,15 +101,29 @@ 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()
+        if set(options) - {"position", "position_set"}:
+            raise ValueError("supported reset options: position, position_set")
+        saved_native = self.native
+        saved_curriculum = self._position_set, self._position_id
+        saved_rng = copy.deepcopy(self._np_random)
+        saved_rng_seed = self._np_random_seed
+        try:
+            super().reset(seed=seed)
+            position, dataset, position_id = curriculum.resolve_start(
+                self.native, options, self.np_random, self._position_set
+            )
+            native_seed = int(self.np_random.integers(0, 2**64, dtype=np.uint64)) if seed is None else seed
+            self.native = self.native.clone()
+            self.native.reset(native_seed, position)
+            self._position_set, self._position_id = dataset, position_id
+            self._opponents()
+        except Exception:
+            self.native = saved_native
+            self._position_set, self._position_id = saved_curriculum
+            self._np_random = saved_rng
+            self._np_random_seed = saved_rng_seed
+            raise
+        self._position_set, self._position_id = dataset, position_id
         self._needs_reset = any(self.native.flags())
         return self._frame(self.seat)
 
@@ -148,31 +166,36 @@ def clone(self) -> GameEnv:
     def get_state(self) -> dict[str, Any]:
         """Trusted checkpoint containing private engine state and environment RNG."""
         return {
-            "version": 1,
+            "version": 2,
             "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,
+            "curriculum": curriculum.checkpoint(self._position_set, self._position_id),
         }
 
     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:
+        if state.get("version") == 2:
+            required.add("curriculum")
+        if set(state) != required or state["version"] not in (1, 2) 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")
+        dataset, position_id = curriculum.restore(state.get("curriculum"), candidate)
         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"]
+        self._position_set, self._position_id = dataset, position_id
 
 
 def make(game: str, **kwargs: Any) -> GameEnv:

python/tests/test_curriculum.py

@@ -0,0 +1,108 @@
+import copy
+import json
+
+import numpy as np
+import pytest
+
+from gamesforai import PositionSet, aec, make
+
+
+def test_gym_curriculum_repeats_seed_and_continues_rng_after_restore():
+    env = make("connect4")
+    first, info = env.reset(seed=82, options={"position_set": "connect4-solved-positions@1"})
+    origin = info["curriculum"]
+    assert set(origin) == {"position_set", "sha256", "position_id"}
+    assert "expected" not in json.dumps(info, default=lambda array: array.tolist())
+    second, second_info = env.reset(seed=82)
+    np.testing.assert_array_equal(first, second)
+    assert second_info["curriculum"] == origin
+    snapshot = env.get_state()
+    continuation = [env.reset()[1]["curriculum"]["position_id"] for _ in range(20)]
+    env.set_state(snapshot)
+    assert continuation == [env.reset()[1]["curriculum"]["position_id"] for _ in range(20)]
+    env.set_state(snapshot)
+    clone = env.clone()
+    assert env.reset()[1]["curriculum"] == clone.reset()[1]["curriculum"]
+    assert "curriculum" not in env.reset(options={"position_set": None})[1]
+
+
+def test_curriculum_config_and_invalid_reset_preserve_checkpoint():
+    env = make("connect4")
+    env.reset(seed=4)
+    before = env.get_state()
+    for options in [
+        {"position_set": "chess-endgames-basic@1"},
+        {"position": "bad"},
+        {"position": "bad", "position_set": "connect4-solved-positions@1"},
+    ]:
+        with pytest.raises(ValueError):
+            env.reset(seed=999, options=options)
+        assert env.get_state() == before
+    env.reset(seed=4, options={"position_set": PositionSet.load("connect4-solved-positions@1")})
+    before = env.get_state()
+    broken = copy.deepcopy(before)
+    broken["curriculum"]["data"] += " "
+    with pytest.raises(ValueError, match="hash"):
+        env.set_state(broken)
+    assert env.get_state() == before
+    with pytest.raises(ValueError, match="config"):
+        make("chess", config={"max_plies": 8}).reset(
+            seed=1, options={"position_set": "chess-endgames-basic@1"}
+        )
+
+
+def test_older_gym_checkpoints_remain_loadable():
+    env = make("tictactoe")
+    env.reset(seed=1)
+    snapshot = env.get_state()
+    snapshot["version"] = 1
+    del snapshot["curriculum"]
+    env.set_state(snapshot)
+    assert env.get_state()["native"] == snapshot["native"]
+
+
+def test_aec_curricula_clone_and_failed_reset_are_reproducible():
+    env = aec("chess")
+    env.reset(seed=91, options={"position_set": "chess-endgames-basic@1"})
+    origin = env.infos[env.agent_selection]["curriculum"]
+    clone = env.clone()
+    seen = []
+    for _ in range(20):
+        env.reset()
+        clone.reset()
+        assert env.native.get_state() == clone.native.get_state()
+        assert env.agent_selection == clone.agent_selection
+        assert env.infos == clone.infos
+        seen.append(env.infos[env.agent_selection]["curriculum"]["position_id"])
+    assert len(set(seen)) > 1
+    before = env.clone()
+    with pytest.raises(ValueError):
+        env.reset(seed=123, options={"position": "invalid FEN"})
+    assert env.native.get_state() == before.native.get_state()
+    env.reset()
+    before.reset()
+    assert env.native.get_state() == before.native.get_state()
+    env.reset(seed=91)
+    assert env.infos[env.agent_selection]["curriculum"] == origin
+    assert "expected" not in json.dumps(env.infos)
+
+
+def test_opponent_opening_sees_current_origin_and_failure_rolls_back():
+    received = []
+    fail = False
+
+    def policy(obs, info):
+        received.append(info.get("curriculum"))
+        if fail:
+            raise ValueError("opponent failure")
+        return int(np.flatnonzero(info["action_mask"])[0])
+
+    env = make("connect4", seat=1, opponent=policy)
+    env.reset(seed=8, options={"position_set": "connect4-solved-positions@1"})
+    assert received[-1]["position_set"] == "connect4-solved-positions@1"
+    assert received[-1]["position_id"] == env.get_state()["curriculum"]["position_id"]
+    before = env.get_state()
+    fail = True
+    with pytest.raises(ValueError, match="opponent failure"):
+        env.reset(seed=999)
+    assert env.get_state() == before