Merge pull request #80 from CurtisTechSolutions/feat/vector-training-curricula

feat(training): add reproducible vector curricula

Author: nosam1998

examples/vector-curricula.md

@@ -0,0 +1,24 @@
+# Batched curriculum training
+
+```python
+from gamesforai import VectorEnv
+
+envs = VectorEnv("connect4", 256, threads=4)
+batch = envs.reset(42, options={"position_set": "connect4-solved-positions@1"})
+origins = envs.origins  # per-row dataset id, content hash and position id
+checkpoint = envs.get_state()
+envs.reset_at([0, 12], [100, 101])  # retain each selected row's curriculum
+envs.set_state(checkpoint)
+```
+
+Every row samples with its explicit reset seed, so results are independent of
+thread scheduling. Full resets use `seed + row`, or accept a seed list.
+`options={"position_set": None}` clears curricula for selected rows.
+Explicit `positions=[...]` also clears curricula for those rows.
+Failed resets preserve both engine state and dataset identities.
+
+Numeric frames keep their existing array-only schema. Read public origin
+metadata separately through `origins`; expected answers stay in trusted datasets.
+Curriculum checkpoints use a version 2 dictionary with each dataset stored once
+and per-row references. Batches without curricula retain the original list of
+native snapshots. Both forms are accepted by `set_state`.

python/gamesforai/vector_env.py

@@ -3,8 +3,13 @@
 
 import copy
 import json
+import operator
 
-from ._native import NativeVectorEnv
+import numpy as np
+
+from ._native import NativeEnv, NativeVectorEnv
+from . import curriculum
+from .positions import PositionSet
 
 
 class VectorEnv:
@@ -15,10 +20,18 @@ class VectorEnv:
     """
 
     def __init__(self, game, n, *, config=None, seed=0, threads=0):
-        self.native = NativeVectorEnv(game, n, json.dumps(config or {}), seed, threads)
+        config_json = json.dumps(config or {})
+        self.native = NativeVectorEnv(game, n, config_json, seed, threads)
         self.n = self.native.n
         self.num_players = self.native.num_players
         self.action_space_size = self.native.action_space_size
+        self._prototype = NativeEnv(game, config_json)
+        self._curricula = [(None, None)] * self.n
+
+    @property
+    def origins(self):
+        """Public dataset identities per row; puzzle answers are excluded."""
+        return [curriculum.origin_info(dataset, position_id) for dataset, position_id in self._curricula]
 
     def frame(self, seats=None):
         """Return owned numeric arrays; select a view per row with seats."""
@@ -28,28 +41,119 @@ def step(self, actions):
         """Apply a whole batch atomically; None is required for finished rows."""
         return self.native.step(list(actions))
 
-    def reset(self, seed=0, *, positions=None):
-        """Reset every row using seed + row, or an explicit list of seeds."""
-        seeds = [seed + index for index in range(self.n)] if isinstance(seed, int) else list(seed)
-        return self.native.reset_at(
-            list(range(self.n)), seeds, [None] * self.n if positions is None else list(positions)
-        )
+    def reset(self, seed=0, *, positions=None, options=None):
+        """Reset all rows with seed + row, or an explicit list of seeds."""
+        if isinstance(seed, (int, np.integer)):
+            if isinstance(seed, (bool, np.bool_)):
+                raise ValueError("seed must be an unsigned 64-bit integer")
+            seeds = [int(seed) + index for index in range(self.n)]
+        else:
+            seeds = list(seed)
+        return self.reset_at(range(self.n), seeds, positions=positions, options=options)
 
-    def reset_at(self, indices, seeds, *, positions=None):
-        """Reset selected rows without discarding other episodes."""
-        indices = list(indices)
-        return self.native.reset_at(
-            indices, list(seeds), [None] * len(indices) if positions is None else list(positions)
-        )
+    def reset_at(self, indices, seeds, *, positions=None, options=None):
+        """Reset selected rows, optionally sampling a pinned position set."""
+        indices, seeds = list(indices), list(seeds)
+        options = dict(options or {})
+        if set(options) - {"position", "position_set"}:
+            raise ValueError("supported reset options: position, position_set")
+        if positions is not None and options:
+            raise ValueError("choose positions or reset options, not both")
+        if len(seeds) != len(indices):
+            raise ValueError("reset seeds must match indices")
+        normalized = []
+        for index, seed in zip(indices, seeds):
+            if isinstance(index, (bool, np.bool_)) or isinstance(seed, (bool, np.bool_)):
+                raise ValueError("indices and seeds must be integers")
+            try:
+                index, seed = operator.index(index), operator.index(seed)
+            except TypeError as error:
+                raise ValueError("indices and seeds must be integers") from error
+            if not 0 <= index < self.n or not 0 <= seed < 2**64:
+                raise ValueError("reset index or seed is out of range")
+            normalized.append((index, seed))
+        if len({index for index, _ in normalized}) != len(indices):
+            raise ValueError("reset indices must be unique")
+        indices = [index for index, _ in normalized]
+        seeds = [seed for _, seed in normalized]
+        positions = None if positions is None else list(positions)
+        if positions is not None and len(positions) != len(indices):
+            raise ValueError("reset positions must match indices")
+        if isinstance(options.get("position_set"), str):
+            options["position_set"] = PositionSet.load(options["position_set"])
+        starts, origins = [], self._curricula[:]
+        for offset, (index, seed) in enumerate(normalized):
+            row_options = {"position": positions[offset]} if positions is not None else options
+            position, dataset, position_id = curriculum.resolve_start(
+                self._prototype, row_options, np.random.default_rng(seed), self._curricula[index][0]
+            )
+            starts.append(position)
+            origins[index] = dataset, position_id
+        result = self.native.reset_at(indices, seeds, starts)
+        self._curricula = origins
+        return result
 
     def clone(self):
         result = copy.copy(self)
         result.native = self.native.clone()
+        result._curricula = self._curricula[:]
         return result
 
     def get_state(self):
-        """Full private checkpoints for trusted training code."""
-        return self.native.get_state()
+        """Private checkpoints; curriculum datasets are stored once per batch."""
+        native = self.native.get_state()
+        if not any(dataset is not None for dataset, _ in self._curricula):
+            return native
+        datasets, rows = {}, []
+        for dataset, position_id in self._curricula:
+            if dataset is None:
+                rows.append(None)
+            else:
+                key = dataset.identifier + ":" + dataset.sha256
+                datasets[key] = dataset.checkpoint()
+                rows.append({"dataset": key, "position_id": position_id})
+        return {"version": 2, "native": native, "datasets": datasets, "rows": rows}
 
     def set_state(self, snapshots):
-        self.native.set_state(list(snapshots))
+        """Restore native state and per-row curricula as one atomic operation."""
+        origins = [(None, None)] * self.n
+        if isinstance(snapshots, dict):
+            if (
+                set(snapshots) != {"version", "native", "datasets", "rows"}
+                or type(snapshots["version"]) is not int or snapshots["version"] != 2
+                or not isinstance(snapshots["datasets"], dict)
+                or len(snapshots["datasets"]) > self.n
+                or not isinstance(snapshots["rows"], list) or len(snapshots["rows"]) != self.n
+            ):
+                raise ValueError("invalid vector curriculum checkpoint")
+            datasets = {}
+            for key, payload in snapshots["datasets"].items():
+                if not isinstance(payload, dict) or set(payload) != {"manifest", "data"}:
+                    raise ValueError("invalid vector position-set checkpoint")
+                dataset = PositionSet.from_jsonl(payload["manifest"], payload["data"])
+                dataset.validate_for(self._prototype)
+                if key != dataset.identifier + ":" + dataset.sha256:
+                    raise ValueError("vector checkpoint dataset identity mismatch")
+                datasets[key] = dataset
+            used = set()
+            for index, row in enumerate(snapshots["rows"]):
+                if row is None:
+                    continue
+                if (
+                    not isinstance(row, dict) or set(row) != {"dataset", "position_id"}
+                    or not isinstance(row["dataset"], str) or row["dataset"] not in datasets
+                ):
+                    raise ValueError("invalid vector checkpoint curriculum row")
+                dataset = datasets[row["dataset"]]
+                if row["position_id"] not in {entry.id for entry in dataset.entries}:
+                    raise ValueError("unknown vector checkpoint position")
+                origins[index] = dataset, row["position_id"]
+                used.add(row["dataset"])
+            if used != set(datasets):
+                raise ValueError("unused vector checkpoint dataset")
+            snapshots = snapshots["native"]
+        if not isinstance(snapshots, (list, tuple)):
+            raise ValueError("vector native checkpoints must be a list")
+        candidate = self.native.clone()
+        candidate.set_state(list(snapshots))
+        self.native, self._curricula = candidate, origins

python/tests/test_vector_curriculum.py

@@ -0,0 +1,85 @@
+import copy
+import json
+
+import numpy as np
+import pytest
+
+from gamesforai import NativeEnv, PositionSet, VectorEnv
+
+
+def test_seeded_position_batches_match_individual_imports_and_threads():
+    dataset = PositionSet.load("connect4-solved-positions@1")
+    entries = {entry.id: entry for entry in dataset.entries}
+    serial = VectorEnv("connect4", 16, threads=1)
+    parallel = VectorEnv("connect4", 16, threads=4)
+    left = serial.reset(41, options={"position_set": dataset})
+    right = parallel.reset(41, options={"position_set": dataset.identifier})
+    assert serial.get_state() == parallel.get_state()
+    assert serial.origins == parallel.origins
+    for key in left:
+        np.testing.assert_array_equal(left[key], right[key])
+    for index, origin in enumerate(serial.origins):
+        entry = entries[origin["curriculum"]["position_id"]]
+        single = NativeEnv("connect4")
+        single.reset(41 + index, entry.position)
+        np.testing.assert_array_equal(left["observation"][index], single.frame(0)["observation"])
+    assert "expected" not in json.dumps(serial.origins)
+
+
+def test_partial_reset_preserves_other_rows_and_cloned_curricula():
+    env = VectorEnv("connect4", 4)
+    env.reset(3, options={"position_set": "connect4-solved-positions@1"})
+    before, origins = env.get_state(), env.origins
+    clone = env.clone()
+    clone.reset_at([1], [91], options={"position_set": None})
+    assert clone.origins[1] == {}
+    assert env.origins == origins
+    for index in [0, 2, 3]:
+        assert clone.get_state()["native"][index] == before["native"][index]
+        assert clone.origins[index] == origins[index]
+    clone.reset_at([1], [95])
+    assert clone.origins[1] == {}
+    assert len(env.get_state()["datasets"]) == 1
+
+
+def test_restore_repeats_curriculum_resets_and_legacy_clears_them():
+    env = VectorEnv("connect4", 8)
+    legacy = env.get_state()
+    env.reset(82, options={"position_set": "connect4-solved-positions@1"})
+    snapshot = json.loads(json.dumps(env.get_state()))
+    other = VectorEnv("connect4", 8)
+    other.set_state(snapshot)
+    for seed in range(6):
+        env.reset(seed)
+        other.reset(seed)
+        assert env.get_state() == other.get_state()
+    other.set_state(legacy)
+    assert other.get_state() == legacy
+    assert other.origins == [{}] * 8
+    assert isinstance(other.get_state(), list)
+
+
+def test_failed_reset_and_restore_leave_native_and_curricula_unchanged():
+    env = VectorEnv("connect4", 4)
+    env.reset(31, options={"position_set": "connect4-solved-positions@1"})
+    before = env.get_state()
+    for kwargs in [
+        {"indices": [0, 0], "seeds": [1, 2]},
+        {"indices": [0], "seeds": [True]},
+        {"indices": [0], "seeds": [2**64]},
+        {"indices": [0], "seeds": [1], "options": {"position_set": "chess-endgames-basic@1"}},
+        {"indices": [0, 1], "seeds": [1, 2], "positions": [None, "bad"]},
+    ]:
+        with pytest.raises(ValueError):
+            env.reset_at(**kwargs)
+        assert env.get_state() == before
+    bad = copy.deepcopy(before)
+    bad["rows"][0]["position_id"] = "missing"
+    with pytest.raises(ValueError):
+        env.set_state(bad)
+    assert env.get_state() == before
+    bad = copy.deepcopy(before)
+    bad["native"][2] = "{}"
+    with pytest.raises(ValueError):
+        env.set_state(bad)
+    assert env.get_state() == before