feat(python): batch remote vector environments atomically

Author: nosam1998

.github/scripts/python-remote-smoke.py

@@ -7,7 +7,7 @@
 import tempfile
 
 import numpy as np
-from gamesforai import NativeEnv, connect, make
+from gamesforai import NativeEnv, VectorEnv, connect, make
 
 with tempfile.TemporaryDirectory() as directory:
     process = subprocess.Popen([
@@ -46,6 +46,17 @@
                 a, b = remote.step(action), local.step(action)
                 np.testing.assert_array_equal(a[0], b[0])
                 assert a[1:4] == b[1:4]
+        remote_batch = client.vector("connect4", 65, seed=31, batch_size=32)
+        local_batch = VectorEnv("connect4", 65, seed=31)
+        for _ in range(4):
+            left = remote_batch.step([3] * 65)
+            right = local_batch.step([3] * 65)
+            for key in left:
+                np.testing.assert_array_equal(left[key], right[key])
+        checkpoints = remote_batch.get_state()
+        remote_batch.reset_at([0, 64], [77, 88])
+        remote_batch.set_state(checkpoints)
+        assert remote_batch.get_state() == checkpoints
         env = client.aec("tictactoe")
         env.reset(seed=0)
         checkpoint = env.get_state()

examples/remote-training.md

@@ -19,6 +19,7 @@ branch.set_state(checkpoint)
 multi = client.aec("tictactoe")
 multi.reset(seed=42)
 native = client.native("chess")
+batch = client.vector("connect4", 128, seed=42, batch_size=32)
 ```
 
 Gymnasium and PettingZoo use the same wrappers as local training, including
@@ -27,8 +28,8 @@ and `ChatPolicy` run in your training process using only the requested seat's
 observation and legal moves. Random opponents use the wrapper's checkpointed
 RNG. Native minimax/MCTS recommendations use the server's analysis endpoint
 and its configured budgets; the same local assistance guards apply. Remote
-Stockfish Gym configuration and vector wrappers are not yet available; native
-Stockfish/VectorEnv remain available for in-process training.
+Stockfish Gym configuration is not yet available; native Stockfish remains
+available for in-process training.
 
 `client.batch(operations)` exposes the [batch REST contract](../docs/training-batches.md)
 directly, with up to 64 independent create/reset/step/observe operations. It
@@ -59,3 +60,23 @@ credentials or endpoint URLs. Timeout is a socket timeout, not a total run
 deadline. The current server is local-only and does not validate API keys;
 do not expose it through a public proxy. Use TLS for future authenticated
 remote deployments. Server and SDK engine versions must be compatible.
+
+## Remote vector batches
+
+`client.vector(game, n, config=..., seed=..., batch_size=32)` uses the same
+VectorEnv API and curricula as local training. Numeric output shapes/dtypes,
+per-seat rewards, explicit terminal resets and clone/checkpoint behavior match
+the native wrapper. The default view is each row's next actor, or seat 0 for
+finished rows; `frame(seats)` requests explicit information sets.
+
+Calls split into ordered HTTP requests of at most `batch_size` rows (1..64)
+and 4 MiB. This is serial HTTP batching, not a local Rayon thread pool;
+leave `threads=0`. A transition commits **all** local checkpoints only after
+every chunk and required observation succeeds. A failure in a later chunk
+preserves the entire previous batch; the stateless server has no partial
+environment mutation to undo. Retrying the same input is safe.
+
+The native limits of 1..4096 environments and 128 MiB of numeric arrays apply.
+Large JSON checkpoint/frame payloads can still hit the server's row/response
+limits; lower `batch_size` when needed. These serialized/array limits are not
+limits on Python object memory. Use in-process VectorEnv for high-volume RL.

python/gamesforai/remote.py

@@ -132,6 +132,15 @@ def make(self, game, **kwargs):
             raise ValueError("remote Gym currently supports random, minimax/mcts or a Python/HTTP policy")
         return GameEnv(game, _native_factory=self.native, **kwargs)
 
+    def vector(self, game, n, *, batch_size=32, **kwargs):
+        from .remote_vector import RemoteNativeVectorEnv
+        from .vector_env import VectorEnv
+
+        def factory(game, n, config_json, seed, threads):
+            return RemoteNativeVectorEnv(self, game, n, config_json, seed, threads,
+                                         batch_size=batch_size)
+        return VectorEnv(game, n, _native_factory=factory, _prototype_factory=self.native, **kwargs)
+
     def aec(self, game, **kwargs):
         from .aec_env import AECGameEnv
         return AECGameEnv(game, _native_factory=self.native, **kwargs)

python/gamesforai/remote_vector.py

@@ -0,0 +1,179 @@
+"""Atomic local batches over the stateless remote training endpoint."""
+from __future__ import annotations
+
+import copy
+import json
+import math
+
+import numpy as np
+
+from .remote import RemoteError, _code, _integer
+from .http_policy import _json
+
+
+class RemoteNativeVectorEnv:
+    def __init__(self, client, game, n, config_json="{}", seed=0, threads=0, *, batch_size=32):
+        self._client = client
+        self.n = _integer(n, 4096)
+        self._batch_size = _integer(batch_size, 64)
+        if not self.n or not self._batch_size:
+            raise ValueError("n must be 1..4096 and batch_size 1..64")
+        if _integer(threads, 64) != 0:
+            raise ValueError("remote workers are controlled by the server; leave threads=0")
+        seed = _integer(seed, 2**64 - self.n)
+        prototype = client.native(game, config_json, seed)
+        self.num_players = prototype.num_players
+        self.action_space_size = prototype.action_space_size
+        self._shape = prototype._shape
+        row_bytes = math.prod(self._shape) * 4 + self.action_space_size + self.num_players * 8 + 6
+        if row_bytes * self.n > 128 * 1024 * 1024:
+            raise ValueError("batch numeric arrays exceed 128 MiB")
+        candidate = [prototype] + [prototype.clone() for _ in range(self.n - 1)]
+        operations = [
+            (index, 0, {"op": "create", "game_id": game, "config": _json(config_json),
+                        "seed": seed + index, "seat": 0})
+            for index in range(1, self.n)
+        ]
+        self._run(candidate, operations)
+        self._project(candidate, None)
+        self._envs = candidate
+
+    def _chunks(self, jobs):
+        batch, size = [], 32
+        for job in jobs:
+            length = len(json.dumps(job[2], allow_nan=False, separators=(",", ":")).encode()) + 1
+            if length + 32 > 4 * 1024 * 1024:
+                raise RemoteError("REQUEST_TOO_LARGE")
+            if batch and (len(batch) >= self._batch_size or size + length > 4 * 1024 * 1024):
+                yield batch
+                batch, size = [], 32
+            batch.append(job)
+            size += length
+        if batch:
+            yield batch
+
+    def _run(self, candidate, jobs, rewards=None, *, unchanged=False):
+        for chunk in self._chunks(jobs):
+            results = self._client.batch([job[2] for job in chunk])
+            for (index, seat, _), result in zip(chunk, results):
+                if result["status"] == "error":
+                    raise RemoteError(_code(result))
+                candidate[index]._accept(result, seat, unchanged=unchanged)
+                if rewards is not None:
+                    rewards[index] = result["frame"]["rewards"]
+                if candidate[index].num_players != self.num_players or candidate[index]._shape != self._shape:
+                    raise RemoteError("INVALID_RESPONSE")
+
+    def _seats(self, candidate, seats):
+        if seats is not None and len(seats) != self.n:
+            raise ValueError("seat count must equal batch size")
+        result = []
+        for index, env in enumerate(candidate):
+            actors = env.current_players()
+            if len(actors) > 1:
+                raise ValueError("VectorEnv requires sequential turns")
+            result.append((actors[0] if actors else 0) if seats is None
+                          else _integer(seats[index], self.num_players - 1))
+        return result
+
+    def _project(self, candidate, seats):
+        seats = self._seats(candidate, seats)
+        jobs = [
+            (index, seat, {"op": "observe", "checkpoint": env._checkpoint, "seat": seat})
+            for index, (env, seat) in enumerate(zip(candidate, seats))
+            if seat not in env._frames
+        ]
+        self._run(candidate, jobs, unchanged=True)
+        # Vector inference needs one view per row; keep Python mask storage bounded.
+        for env, seat in zip(candidate, seats):
+            env._frames = {seat: env._frames[seat]}
+        return seats
+
+    def _arrays(self, candidate, seats=None, rewards=None):
+        seats = self._project(candidate, seats)
+        output = {
+            "observation": np.empty((self.n, *self._shape), dtype=np.float32),
+            "action_mask": np.empty((self.n, self.action_space_size), dtype=np.int8),
+            "seat": np.array(seats, dtype=np.int16),
+            "to_act": np.empty(self.n, dtype=np.int16),
+            "terminated": np.empty(self.n, dtype=bool),
+            "truncated": np.empty(self.n, dtype=bool),
+            "rewards": np.zeros((self.n, self.num_players), dtype=np.float64) if rewards is None else rewards.copy(),
+        }
+        for index, (env, seat) in enumerate(zip(candidate, seats)):
+            raw = env._frames[seat]
+            tensor = raw["observation"]["tensor"]
+            output["observation"][index] = np.array(tensor["values"], dtype=np.float32).reshape(self._shape)
+            output["action_mask"][index] = raw["action_mask"]
+            actors = raw["to_act"]
+            output["to_act"][index] = actors[0] if actors else -1
+            output["terminated"][index] = raw["terminated"]
+            output["truncated"][index] = raw["truncated"]
+        return output
+
+    def frame(self, seats=None):
+        candidate = [env.clone() for env in self._envs]
+        result = self._arrays(candidate, seats)
+        self._envs = candidate
+        return result
+
+    def step(self, actions):
+        if len(actions) != self.n:
+            raise ValueError("action count must equal batch size")
+        candidate = [env.clone() for env in self._envs]
+        jobs = []
+        for index, (env, action) in enumerate(zip(candidate, actions)):
+            actors = env.current_players()
+            if not actors and any(env.flags()) and action is None:
+                continue
+            if len(actors) != 1 or action is None:
+                raise ValueError("use legal indices for active rows and None for finished rows")
+            seat = actors[0]
+            jobs.append((index, seat, {"op": "step", "checkpoint": env._checkpoint,
+                                      "seat": seat, "action": _integer(action, self.action_space_size - 1)}))
+        rewards = np.zeros((self.n, self.num_players), dtype=np.float64)
+        self._run(candidate, jobs, rewards)
+        result = self._arrays(candidate, rewards=rewards)
+        self._envs = candidate
+        return result
+
+    def reset_at(self, indices, seeds, positions):
+        if len(indices) != len(seeds) or len(indices) != len(positions):
+            raise ValueError("indices, seeds, and positions must have equal lengths")
+        candidate = [env.clone() for env in self._envs]
+        jobs, seen = [], set()
+        for index, seed, position in zip(indices, seeds, positions):
+            index, seed = _integer(index, self.n - 1), _integer(seed, 2**64 - 1)
+            if index in seen or (position is not None and not isinstance(position, str)):
+                raise ValueError("indices must be unique and positions must be strings or None")
+            seen.add(index)
+            jobs.append((index, 0, {"op": "reset", "checkpoint": candidate[index]._checkpoint,
+                                   "seed": seed, "position": position, "seat": 0}))
+        self._run(candidate, jobs)
+        result = self._arrays(candidate)
+        self._envs = candidate
+        return result
+
+    def clone(self):
+        result = copy.copy(self)
+        result._envs = [env.clone() for env in self._envs]
+        return result
+
+    def get_state(self):
+        return [env.get_state() for env in self._envs]
+
+    def set_state(self, snapshots):
+        if len(snapshots) != self.n or any(not isinstance(item, str) for item in snapshots):
+            raise ValueError("checkpoint count or type is invalid")
+        if sum(len(item.encode()) for item in snapshots) > 128 * 1024 * 1024:
+            raise ValueError("batch checkpoints exceed 128 MiB")
+        candidate, jobs = [env.clone() for env in self._envs], []
+        for index, (env, snapshot) in enumerate(zip(candidate, snapshots)):
+            checkpoint = _json(snapshot)
+            if (not isinstance(checkpoint, dict) or checkpoint.get("game_id") != env.game_id
+                    or checkpoint.get("config") != env._checkpoint["config"]):
+                raise ValueError("checkpoint must match game and configuration")
+            jobs.append((index, 0, {"op": "observe", "checkpoint": checkpoint, "seat": 0}))
+        self._run(candidate, jobs)
+        self._project(candidate, None)
+        self._envs = candidate

python/gamesforai/vector_env.py

@@ -19,13 +19,14 @@ class VectorEnv:
     per seat. Finished rows receive None actions until explicitly reset.
     """
 
-    def __init__(self, game, n, *, config=None, seed=0, threads=0):
+    def __init__(self, game, n, *, config=None, seed=0, threads=0,
+                 _native_factory=NativeVectorEnv, _prototype_factory=NativeEnv):
         config_json = json.dumps(config or {})
-        self.native = NativeVectorEnv(game, n, config_json, seed, threads)
+        self.native = _native_factory(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._prototype = _prototype_factory(game, config_json)
         self._curricula = [(None, None)] * self.n
 
     @property

python/tests/test_remote_vector.py

@@ -0,0 +1,106 @@
+import json
+
+import numpy as np
+import pytest
+
+from gamesforai import VectorEnv, connect
+from gamesforai.remote import RemoteError
+from test_remote import server  # Shared loopback protocol fixture.
+
+
+def equal(left, right):
+    assert left.keys() == right.keys()
+    for key in left:
+        assert left[key].dtype == right[key].dtype
+        np.testing.assert_array_equal(left[key], right[key])
+
+
+def test_remote_vector_matches_native_complete_games_and_explicit_views(server):
+    url, _ = server
+    remote = connect(url).vector("tictactoe", 4, seed=42, batch_size=2)
+    local = VectorEnv("tictactoe", 4, seed=42)
+    equal(remote.frame(), local.frame())
+    equal(remote.frame([1, 0, 1, 0]), local.frame([1, 0, 1, 0]))
+    for _ in range(10):
+        frame = local.frame()
+        actions = [None if ended else int(np.flatnonzero(mask)[0])
+                   for ended, mask in zip(frame["terminated"] | frame["truncated"], frame["action_mask"])]
+        equal(remote.step(actions), local.step(actions))
+        assert [json.loads(value) for value in remote.get_state()] == [
+            json.loads(value) for value in local.get_state()]
+        if np.all(local.frame()["terminated"]):
+            break
+    assert np.all(remote.frame()["terminated"])
+    equal(remote.step([None] * 4), local.step([None] * 4))
+
+
+def test_late_chunk_failure_preserves_every_checkpoint_and_can_retry(server):
+    url, state = server
+    env = connect(url).vector("tictactoe", 65, batch_size=32)
+    env.step([0] * 65)
+    before = env.get_state()
+    calls = len(state.requests)
+    # The occupied square fails in the third chunk, after earlier chunks succeeded.
+    with pytest.raises(RemoteError, match="ILLEGAL_ACTION"):
+        env.step([1] * 64 + [0])
+    assert env.get_state() == before
+    requests = [body for _, _, body in state.requests[calls:] if body is not None]
+    assert [len(body["operations"]) for body in requests] == [32, 32, 1]
+    env.step([1] * 65)
+    assert env.get_state() != before
+
+
+def test_partial_reset_curriculum_and_checkpoint_clone_parity(server):
+    url, _ = server
+    remote = connect(url).vector("connect4", 3, seed=8)
+    local = VectorEnv("connect4", 3, seed=8)
+    options = {"position_set": "connect4-solved-positions@1"}
+    equal(remote.reset(seed=42, options=options), local.reset(seed=42, options=options))
+    assert remote.origins == local.origins
+    checkpoint = remote.get_state()
+    clone = remote.clone()
+    equal(remote.reset_at([2, 0], [8, 7]), local.reset_at([2, 0], [8, 7]))
+    assert clone.get_state() == checkpoint
+    assert remote.get_state() != checkpoint
+    remote.set_state(checkpoint)
+    equal(remote.frame(), clone.frame())
+    bad = remote.get_state()
+    bad["native"][2] = "{}"
+    with pytest.raises(ValueError):
+        remote.set_state(bad)
+    assert remote.get_state() == checkpoint
+
+
+def test_projection_failure_does_not_commit_and_arrays_are_owned(server):
+    url, state = server
+    env = connect(url).vector("tictactoe", 3)
+    before = env.get_state()
+    state.fail_observe = True
+    with pytest.raises(RemoteError):
+        env.step([0, 1, 2])
+    assert env.get_state() == before
+    state.fail_observe = False
+    frame = env.frame()
+    frame["action_mask"][:] = 0
+    frame["observation"][:] = 99
+    assert env.frame()["action_mask"].all()
+    assert not (env.frame()["observation"] == 99).any()
+
+
+def test_vector_rejects_invalid_batches_before_committing(server):
+    url, _ = server
+    client = connect(url)
+    for kwargs in [{"n": 0}, {"n": 4097}, {"n": 1, "batch_size": 0},
+                   {"n": 1, "threads": 2}, {"n": 2, "seed": 2**64 - 1}]:
+        with pytest.raises(ValueError):
+            client.vector("tictactoe", **kwargs)
+    env = client.vector("tictactoe", 2)
+    before = env.get_state()
+    for actions in [[0], [True, 1], [None, None], [0, 10]]:
+        with pytest.raises(ValueError):
+            env.step(actions)
+    with pytest.raises(ValueError):
+        env.reset_at([0, 0], [1, 2])
+    with pytest.raises(ValueError):
+        env.frame([2, 0])
+    assert env.get_state() == before