[envpool] fix multiplayer players.env_id inference (#347)

## Summary
- Problem: the simplified Python multiplayer action path was filling
`players.env_id` with the batch `env_id`, so flattened per-player
actions were grouped against the wrong environments when an env
contributed more than one player action.
- Scope: infer `players.env_id` from the incoming Python action payload,
reuse the latest observed player-to-env mapping when player counts vary
by env, and add Python-side regression coverage for the multiplayer
wrapper path.
- Outcome: multiplayer actions sent through the shorthand Python API now
preserve the correct per-player env mapping instead of silently slicing
the action buffer incorrectly.

This fixes the Python wrapper bug behind issue #296 without changing the
C++ action parser contract.

## Technical Details
- Approach: teach `EnvPoolMixin._from()` to derive `players.env_id` from
player-shaped action arrays, fall back to the cached
`info:players.env_id` mapping from the last `recv()`, and raise when the
mapping is ambiguous instead of fabricating a wrong one.
- Code pointers:
- `envpool/python/envpool.py`: adds `players.env_id` inference and
caches `info:players.env_id` on `recv()` for variable-player batches.
- `envpool/dummy/dummy_py_envpool_test.py`: adds regression coverage
around the real dummy DM wrapper for uniform multiplayer, cached
variable-player, explicit mapping, and ambiguous-input cases.
- `envpool/dummy/BUILD`: wires the dummy Python test to the Python API
wrapper target.
- Notes: the regression test now uses the existing dummy env wrapper
rather than a standalone fake mixin harness.

## Test Plan
### Automated
- `python3 -m py_compile envpool/python/envpool.py
envpool/dummy/dummy_py_envpool_test.py`: passed.
### Suggested Manual
- `USE_BAZEL_VERSION=8.6.0 bazelisk test
//envpool/dummy:dummy_py_envpool_test --config=test
--spawn_strategy=local --test_output=errors`: exercise the real dummy
wrapper path on Linux.
- Re-run the issue #296 repro against this branch on `dev-0`: confirm
the flattened multiplayer action now expands to the expected
`players.env_id` sequence instead of the old one-element-per-env
mapping.

Author: Trinkle23897

envpool/dummy/BUILD

@@ -51,6 +51,7 @@ py_test(
     deps = [
         requirement("numpy"),
         requirement("absl-py"),
+        "//envpool/python:api",
     ],
 )
 

envpool/dummy/dummy_py_envpool_test.py

@@ -15,12 +15,45 @@
 
 import os
 import time
+from typing import Any
 
 import numpy as np
 from absl import logging
 from absl.testing import absltest
 from envpool.dummy.dummy_envpool import _DummyEnvPool, _DummyEnvSpec
 
+from envpool.python.api import py_env
+from envpool.python.protocol import EnvPool
+
+DummyEnvSpec, _DummyDMEnvPool, _, _ = py_env(_DummyEnvSpec, _DummyEnvPool)
+
+
+def _make_dummy_dm_env() -> EnvPool:
+    config = DummyEnvSpec.gen_config(
+        num_envs=2,
+        batch_size=2,
+        max_num_players=4,
+    )
+    return _DummyDMEnvPool(DummyEnvSpec(config))
+
+
+def _make_multiplayer_action(
+    player_count: int,
+    players_env_id: np.ndarray | None = None,
+) -> dict[str, object]:
+    players: dict[str, np.ndarray] = {
+        "id": np.arange(player_count, dtype=np.int32),
+        "action": np.arange(player_count, dtype=np.int32),
+    }
+    action: dict[str, Any] = {
+        "env_id": np.array([0, 1], dtype=np.int32),
+        "list_action": np.zeros((2, 6), dtype=np.float64),
+        "players": players,
+    }
+    if players_env_id is not None:
+        players["env_id"] = players_env_id
+    return action
+
 
 class _DummyEnvPoolTest(absltest.TestCase):
     def test_config(self) -> None:
@@ -121,5 +154,44 @@ def test_xla(self) -> None:
         self.assertTrue(xla_failed)
 
 
+class _EnvPoolMixinRegressionTest(absltest.TestCase):
+    def test_from_repeats_env_id_for_uniform_multiplayer_action(self) -> None:
+        env = _make_dummy_dm_env()
+        action = _make_multiplayer_action(player_count=6)
+        converted = env._from(action)
+        np.testing.assert_array_equal(
+            converted[1],
+            np.array([0, 0, 0, 1, 1, 1], dtype=np.int32),
+        )
+
+    def test_recv_cache_handles_variable_player_counts(self) -> None:
+        env = _make_dummy_dm_env()
+        env._last_players_env_id = np.array([0, 0, 1, 1, 1], dtype=np.int32)
+        action = _make_multiplayer_action(player_count=5)
+        converted = env._from(action)
+        np.testing.assert_array_equal(
+            converted[1], np.array([0, 0, 1, 1, 1], dtype=np.int32)
+        )
+
+    def test_from_preserves_explicit_players_env_id(self) -> None:
+        env = _make_dummy_dm_env()
+        action = _make_multiplayer_action(
+            player_count=5,
+            players_env_id=np.array([0, 0, 1, 1, 1], dtype=np.int32),
+        )
+        converted = env._from(action)
+        np.testing.assert_array_equal(
+            converted[1], np.array([0, 0, 1, 1, 1], dtype=np.int32)
+        )
+
+    def test_from_raises_when_players_env_id_is_ambiguous(self) -> None:
+        env = _make_dummy_dm_env()
+        action = _make_multiplayer_action(player_count=5)
+        with self.assertRaisesRegex(
+            RuntimeError, "Cannot infer players.env_id"
+        ):
+            env._from(action)
+
+
 if __name__ == "__main__":
     absltest.main()

envpool/python/envpool.py

@@ -30,6 +30,76 @@ class EnvPoolMixin(ABC):
 
     _spec: EnvSpec
 
+    def _player_action_count(
+        self: EnvPool, adict: dict[str, Any]
+    ) -> int | None:
+        """Infer how many player actions are present in the current input."""
+        player_count = None
+        for key, spec in self.spec.action_array_spec.items():
+            if key in ("env_id", "players.env_id") or key not in adict:
+                continue
+            shape = tuple(spec.shape)
+            if len(shape) == 0 or shape[0] != -1:
+                continue
+            value_shape = np.shape(adict[key])
+            count = 1 if len(value_shape) == 0 else int(value_shape[0])
+            if player_count is None:
+                player_count = count
+            elif player_count != count:
+                raise RuntimeError(
+                    "Inconsistent leading dimensions across player actions."
+                )
+        return player_count
+
+    def _cached_players_env_id(
+        self: EnvPool, env_id: np.ndarray, player_count: int
+    ) -> np.ndarray | None:
+        """Reuse the last recv/reset mapping when player counts vary by env."""
+        if not hasattr(self, "_last_players_env_id"):
+            return None
+        cached = np.asarray(self._last_players_env_id, dtype=np.int32)
+        segments = []
+        for eid in env_id.tolist():
+            matches = cached[cached == eid]
+            if matches.size == 0:
+                return None
+            segments.append(matches)
+        if not segments:
+            return np.empty(0, dtype=np.int32)
+        players_env_id = np.concatenate(segments)
+        if players_env_id.shape[0] != player_count:
+            return None
+        return players_env_id
+
+    def _infer_players_env_id(
+        self: EnvPool, adict: dict[str, Any]
+    ) -> np.ndarray:
+        """Fill in players.env_id for the simplified multiplayer API."""
+        env_id = np.asarray(adict["env_id"], dtype=np.int32)
+        if env_id.ndim == 0:
+            env_id = env_id.reshape(1)
+        if self.config.get("max_num_players", 1) == 1:
+            return env_id
+        player_count = self._player_action_count(adict)
+        if player_count is None or player_count == env_id.shape[0]:
+            return env_id
+        cached = self._cached_players_env_id(env_id, player_count)
+        if cached is not None:
+            return cached
+        if env_id.shape[0] == 0 or player_count % env_id.shape[0] != 0:
+            raise RuntimeError(
+                "Cannot infer players.env_id for multiplayer action; "
+                "pass a dict action with explicit players.env_id."
+            )
+        players_per_env = player_count // env_id.shape[0]
+        max_num_players = self.config.get("max_num_players", 1)
+        if players_per_env > max_num_players:
+            raise RuntimeError(
+                "Cannot infer players.env_id for multiplayer action; "
+                "per-env player count exceeds max_num_players."
+            )
+        return np.repeat(env_id, players_per_env).astype(np.int32, copy=False)
+
     def _check_action(self: EnvPool, actions: list[np.ndarray]) -> None:
         if hasattr(self, "_check_action_finished"):  # only check once
             return
@@ -74,20 +144,20 @@ def _from(
             if isinstance(action, np.ndarray):
                 # else it could be a jax array, when using xla
                 action = action.astype(
-                    self._last_action_type,  # type: ignore
+                    self._last_action_type,
                     order="C",
                 )
-            adict = {self._last_action_name: action}  # type: ignore
+            adict = {self._last_action_name: action}
         if env_id is None:
             if "env_id" not in adict:
                 adict["env_id"] = self.all_env_ids
         else:
             adict["env_id"] = env_id.astype(np.int32)
         if "players.env_id" not in adict:
-            adict["players.env_id"] = adict["env_id"]
+            adict["players.env_id"] = self._infer_players_env_id(adict)
         if not hasattr(self, "_action_names"):
             self._action_names = self._spec._action_keys
-        return [adict[k] for k in self._action_names]  # type: ignore
+        return [adict[k] for k in self._action_names]
 
     def __len__(self: EnvPool) -> int:
         """Return the number of environments."""
@@ -100,7 +170,7 @@ def all_env_ids(self: EnvPool) -> np.ndarray:
             self._all_env_ids = np.arange(
                 self.config["num_envs"], dtype=np.int32
             )
-        return self._all_env_ids  # type: ignore
+        return self._all_env_ids
 
     @property
     def is_async(self: EnvPool) -> bool:
@@ -135,6 +205,13 @@ def recv(
     ) -> TimeStep | tuple:
         """Recv a batch state from EnvPool."""
         state_list = self._recv()
+        if not hasattr(self, "_state_names"):
+            self._state_names = self._state_keys
+        state = dict(zip(self._state_names, state_list, strict=False))
+        if "info:players.env_id" in state:
+            self._last_players_env_id = np.array(
+                state["info:players.env_id"], copy=True
+            )
         return self._to(state_list, reset, return_info)
 
     def async_reset(self: EnvPool) -> None:

envpool/python/protocol.py

@@ -121,7 +121,14 @@ class EnvPool(Protocol):
     """Cpp PyEnvpool class interface."""
 
     _state_keys: list[str]
+    _state_names: list[str]
     _action_keys: list[str]
+    _action_names: list[str]
+    _check_action_finished: bool
+    _all_env_ids: np.ndarray
+    _last_action_name: str
+    _last_action_type: Any
+    _last_players_env_id: np.ndarray
     spec: Any
 
     def __init__(self, spec: EnvSpec):
@@ -141,6 +148,17 @@ def _action_spec(self) -> list:
     def _check_action(self, actions: list) -> None:
         """Check action shapes."""
 
+    def _player_action_count(self, adict: dict[str, Any]) -> int | None:
+        """Infer the leading player-action dimension."""
+
+    def _cached_players_env_id(
+        self, env_id: np.ndarray, player_count: int
+    ) -> np.ndarray | None:
+        """Reuse cached player-to-env mapping when available."""
+
+    def _infer_players_env_id(self, adict: dict[str, Any]) -> np.ndarray:
+        """Infer players.env_id for simplified multiplayer actions."""
+
     def _recv(self) -> list[np.ndarray]:
         """Cpp private _recv method."""