Skip to content

Commit f34b34f

Browse files
Trinkle23897Trinkle23897
authored
[pgx] Add remaining native PGX environments (#411)
## Description Add native C++ EnvPool implementations for the remaining non-MinAtar PGX 2.6.0 environments under `envpool/pgx`. This PR adds native runtimes and bindings for TicTacToe, ConnectFour, Hex, Othello, Kuhn Poker, Leduc Holdem, 2048, Animal Shogi, Backgammon, Chess, Gardner Chess, Shogi, and Sparrow Mahjong. Go and the Chinese-rule Go variants remain in the same PGX family. MinAtar is intentionally not registered because EnvPool already has native Atari environments. The PR also expands PGX tests so the registered envs have smoke coverage, deterministic rollout coverage, render coverage, and PGX oracle alignment. Oracle alignment uses exact `np.testing.assert_array_equal` comparisons for observations, legal masks, rewards, done flags, and exposed state/info on tested rollouts. Render tests validate EnvPool native RGB render determinism and shape; they do not claim bitwise parity with PGX SVG visualizers. Validation run: - `bazel test --test_output=errors //envpool/pgx/...` - `ruff format --check envpool/pgx/__init__.py envpool/pgx/registration.py envpool/pgx/pgx_align_test.py envpool/pgx/pgx_deterministic_test.py envpool/pgx/pgx_render_test.py envpool/pgx/pgx_test.py` - `ruff check envpool/pgx/__init__.py envpool/pgx/registration.py envpool/pgx/pgx_align_test.py envpool/pgx/pgx_deterministic_test.py envpool/pgx/pgx_render_test.py envpool/pgx/pgx_test.py` - `clang-format --dry-run --Werror envpool/pgx/pgx.cc envpool/pgx/animal_shogi.h envpool/pgx/backgammon.h envpool/pgx/board_games.h envpool/pgx/card_games.h envpool/pgx/chess_games.h envpool/pgx/play2048.h envpool/pgx/shogi.h envpool/pgx/sparrow_mahjong.h` - `buildifier -mode=check envpool/pgx/BUILD` - `git diff --check` ## Motivation and Context This completes the PGX family integration after the Go PR by registering the remaining PGX tasks that do not overlap EnvPool's Atari support. It keeps task IDs in EnvPool style without a `PGX` prefix and keeps the PGX code under `envpool/pgx` instead of introducing a new JAX env folder. - [ ] I have raised an issue to propose this change ([required](https://envpool.readthedocs.io/en/latest/pages/contributing.html) for new features and bug fixes) ## Types of changes What types of changes does your code introduce? Put an `x` in all the boxes that apply: - [ ] Bug fix (non-breaking change which fixes an issue) - [x] New feature (non-breaking change which adds core functionality) - [x] New environment (non-breaking change which adds 3rd-party environment) - [ ] Breaking change (fix or feature that would cause existing functionality to change) - [x] Documentation (update in the documentation) - [ ] Example (update in the folder of example) ## Implemented Tasks - [x] Add native C++ PGX runtimes and pybind registration for the remaining non-MinAtar PGX tasks - [x] Add smoke, deterministic, render, and exact PGX oracle alignment tests - [x] Update PGX docs while keeping README's supported-env list at family granularity ## Checklist Go over all the following points, and put an `x` in all the boxes that apply. If you are unsure about any of these, don't hesitate to ask. We are here to help! - [x] I have read the [CONTRIBUTION](https://envpool.readthedocs.io/en/latest/pages/contributing.html) guide (**required**) - [x] My change requires a change to the documentation. - [x] I have updated the tests accordingly (*required for a bug fix or a new feature*). - [x] I have updated the documentation accordingly. - [ ] I have reformatted the code using `make format` (**required**) - [ ] I have checked the code using `make lint` (**required**) - [ ] I have ensured `make bazel-test` pass. (**required**)
1 parent 0800e08 commit f34b34f

20 files changed

Lines changed: 7232 additions & 84 deletions

‎docs/env/pgx.rst‎

Lines changed: 48 additions & 44 deletions
Original file line numberDiff line numberDiff line change
@@ -1,17 +1,46 @@
11
PGX
22
===
33

4-
EnvPool includes a native C++ implementation of the PGX Go environments from
4+
EnvPool includes native C++ implementations of PGX environments from
55
`PGX <https://github.com/sotetsuk/pgx>`_ 2.6.0.
66

7-
The default ``Go9x9-v1`` and ``Go19x19-v1`` tasks follow PGX's Go v1 rules:
7+
Supported Tasks
8+
---------------
9+
10+
EnvPool registers the following PGX tasks:
11+
12+
* ``Go9x9-v1``
13+
* ``Go13x13-v1``
14+
* ``Go19x19-v1``
15+
* ``ChineseGo9x9-v1``
16+
* ``ChineseGo13x13-v1``
17+
* ``ChineseGo19x19-v1``
18+
* ``TicTacToe-v1``
19+
* ``ConnectFour-v1``
20+
* ``Hex-v1``
21+
* ``Othello-v1``
22+
* ``KuhnPoker-v1``
23+
* ``LeducHoldem-v1``
24+
* ``Play2048-v1``
25+
* ``AnimalShogi-v1``
26+
* ``Backgammon-v1``
27+
* ``Chess-v1``
28+
* ``GardnerChess-v1``
29+
* ``Shogi-v1``
30+
* ``SparrowMahjong-v1``
31+
32+
The PGX MinAtar tasks are not registered here because EnvPool already provides
33+
native Atari environments.
34+
35+
Go Rules
36+
--------
37+
38+
``Go9x9-v1``, ``Go13x13-v1``, and ``Go19x19-v1`` follow PGX's
39+
Go v1 rules:
840

941
* Tromp-Taylor scoring.
10-
* Two players.
11-
* ``Go9x9-v1`` and ``Go19x19-v1`` task IDs.
1242
* ``N * N + 1`` discrete actions, where the final action is pass.
13-
* Boolean observation shape ``(N, N, 17)`` using the AlphaGo Zero history
14-
planes.
43+
* Boolean observation shape ``(N, N, 17)`` using AlphaGo Zero history planes.
1544
* SSK legal-action filtering, with positional superko occurrence ending the
1645
game as a loss for the player who made the repeated position.
1746

@@ -22,48 +51,23 @@ EnvPool also provides Chinese-rule variants:
2251
* Positional superko moves are masked as illegal actions instead of being
2352
accepted and then turned into a terminal loss.
2453
* The same no-suicide, two-pass terminal, action, and observation API as the
25-
PGX-compatible tasks.
26-
27-
Supported Tasks
28-
---------------
29-
30-
* ``Go9x9-v1``
31-
* ``Go19x19-v1``
32-
* ``ChineseGo9x9-v1``
33-
* ``ChineseGo19x19-v1``
34-
35-
Observation
36-
-----------
37-
38-
PGX Go is turn-based, but EnvPool exposes it through the existing multiplayer
39-
API. Each state contains two player observations. ``info["current_player"]``
40-
identifies the player ID whose turn it is, and a single action per environment
41-
is interpreted as that player's action.
42-
43-
The Gymnasium observation space is ``MultiBinary((N, N, 17))``. Runtime
44-
observations have leading player dimension ``2`` for each environment and are
45-
returned as boolean arrays.
54+
PGX-compatible Go tasks.
4655

47-
Info
48-
----
56+
API Notes
57+
---------
4958

50-
The Gymnasium info dictionary includes:
59+
PGX turn-based games are exposed through EnvPool's multiplayer API. Each state
60+
contains one observation per player, ``info["current_player"]`` identifies the
61+
player ID whose turn it is, and each environment consumes one action for that
62+
current player.
5163

52-
* ``board``: clipped board values, ``1`` for black, ``-1`` for white, ``0`` for
53-
empty.
54-
* ``current_player``: player ID to act.
55-
* ``legal_action_mask``: legal actions for the current player.
56-
* ``ko``: SSK ko point, or ``-1``.
57-
* ``is_psk``: whether the latest move produced positional superko.
58-
* ``consecutive_pass_count``.
59-
* ``black_area`` and ``white_area``: area scores before komi under the selected
60-
rule set.
61-
* ``players.id``.
64+
The task IDs intentionally follow EnvPool style and do not use a ``PGX`` prefix.
6265

6366
Configuration
6467
-------------
6568

66-
``komi`` defaults to ``7.5`` and ``history_length`` defaults to ``8`` to match
67-
PGX. ``max_terminal_steps=0`` means ``2 * N * N``, matching PGX's default.
68-
``rules`` is ``"pgx"`` for the PGX-compatible tasks and ``"chinese"`` for the
69-
Chinese-rule variants.
69+
Go tasks support ``komi``, ``history_length``, ``max_terminal_steps``, and
70+
``rules``. ``komi`` defaults to ``7.5`` and ``history_length`` defaults to ``8``
71+
to match PGX. ``max_terminal_steps=0`` means ``2 * N * N``. ``rules`` is
72+
``"pgx"`` for the PGX-compatible tasks and ``"chinese"`` for the Chinese-rule
73+
variants.

‎envpool/__init__.py‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -37,7 +37,7 @@
3737
register,
3838
)
3939

40-
__version__ = "1.2.3"
40+
__version__ = "1.2.4"
4141
__all__ = [
4242
"register",
4343
"make",

‎envpool/pgx/BUILD‎

Lines changed: 14 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -20,8 +20,18 @@ load("//envpool:requirements.bzl", "requirement")
2020
package(default_visibility = ["//visibility:public"])
2121

2222
cc_library(
23-
name = "pgx_go_env",
24-
hdrs = ["go.h"],
23+
name = "pgx_env",
24+
hdrs = [
25+
"animal_shogi.h",
26+
"backgammon.h",
27+
"board_games.h",
28+
"card_games.h",
29+
"chess_games.h",
30+
"go.h",
31+
"play2048.h",
32+
"shogi.h",
33+
"sparrow_mahjong.h",
34+
],
2535
deps = [
2636
"//envpool/core:async_envpool",
2737
"//envpool/core:env",
@@ -32,7 +42,7 @@ pybind_extension(
3242
name = "pgx_envpool",
3343
srcs = ["pgx.cc"],
3444
deps = [
35-
":pgx_go_env",
45+
":pgx_env",
3646
"//envpool/core:py_envpool",
3747
],
3848
)
@@ -68,6 +78,7 @@ py_test(
6878

6979
py_test(
7080
name = "pgx_align_test",
81+
timeout = "long",
7182
srcs = ["pgx_align_test.py"],
7283
data = [":pgx_envpool"],
7384
imports = ["../.."],

‎envpool/pgx/__init__.py‎

Lines changed: 134 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -15,16 +15,149 @@
1515

1616
from envpool.python.api import py_env
1717

18-
from .pgx_envpool import _GoEnvPool, _GoEnvSpec
18+
from .pgx_envpool import (
19+
_AnimalShogiEnvPool,
20+
_AnimalShogiEnvSpec,
21+
_BackgammonEnvPool,
22+
_BackgammonEnvSpec,
23+
_ChessEnvPool,
24+
_ChessEnvSpec,
25+
_ConnectFourEnvPool,
26+
_ConnectFourEnvSpec,
27+
_GardnerChessEnvPool,
28+
_GardnerChessEnvSpec,
29+
_GoEnvPool,
30+
_GoEnvSpec,
31+
_HexEnvPool,
32+
_HexEnvSpec,
33+
_KuhnPokerEnvPool,
34+
_KuhnPokerEnvSpec,
35+
_LeducHoldemEnvPool,
36+
_LeducHoldemEnvSpec,
37+
_OthelloEnvPool,
38+
_OthelloEnvSpec,
39+
_Play2048EnvPool,
40+
_Play2048EnvSpec,
41+
_ShogiEnvPool,
42+
_ShogiEnvSpec,
43+
_SparrowMahjongEnvPool,
44+
_SparrowMahjongEnvSpec,
45+
_TicTacToeEnvPool,
46+
_TicTacToeEnvSpec,
47+
)
1948

2049
(
2150
GoEnvSpec,
2251
GoDMEnvPool,
2352
GoGymnasiumEnvPool,
2453
) = py_env(_GoEnvSpec, _GoEnvPool)
54+
(
55+
TicTacToeEnvSpec,
56+
TicTacToeDMEnvPool,
57+
TicTacToeGymnasiumEnvPool,
58+
) = py_env(_TicTacToeEnvSpec, _TicTacToeEnvPool)
59+
(
60+
ConnectFourEnvSpec,
61+
ConnectFourDMEnvPool,
62+
ConnectFourGymnasiumEnvPool,
63+
) = py_env(_ConnectFourEnvSpec, _ConnectFourEnvPool)
64+
(
65+
HexEnvSpec,
66+
HexDMEnvPool,
67+
HexGymnasiumEnvPool,
68+
) = py_env(_HexEnvSpec, _HexEnvPool)
69+
(
70+
OthelloEnvSpec,
71+
OthelloDMEnvPool,
72+
OthelloGymnasiumEnvPool,
73+
) = py_env(_OthelloEnvSpec, _OthelloEnvPool)
74+
(
75+
KuhnPokerEnvSpec,
76+
KuhnPokerDMEnvPool,
77+
KuhnPokerGymnasiumEnvPool,
78+
) = py_env(_KuhnPokerEnvSpec, _KuhnPokerEnvPool)
79+
(
80+
LeducHoldemEnvSpec,
81+
LeducHoldemDMEnvPool,
82+
LeducHoldemGymnasiumEnvPool,
83+
) = py_env(_LeducHoldemEnvSpec, _LeducHoldemEnvPool)
84+
(
85+
Play2048EnvSpec,
86+
Play2048DMEnvPool,
87+
Play2048GymnasiumEnvPool,
88+
) = py_env(_Play2048EnvSpec, _Play2048EnvPool)
89+
(
90+
AnimalShogiEnvSpec,
91+
AnimalShogiDMEnvPool,
92+
AnimalShogiGymnasiumEnvPool,
93+
) = py_env(_AnimalShogiEnvSpec, _AnimalShogiEnvPool)
94+
(
95+
BackgammonEnvSpec,
96+
BackgammonDMEnvPool,
97+
BackgammonGymnasiumEnvPool,
98+
) = py_env(_BackgammonEnvSpec, _BackgammonEnvPool)
99+
(
100+
ChessEnvSpec,
101+
ChessDMEnvPool,
102+
ChessGymnasiumEnvPool,
103+
) = py_env(_ChessEnvSpec, _ChessEnvPool)
104+
(
105+
GardnerChessEnvSpec,
106+
GardnerChessDMEnvPool,
107+
GardnerChessGymnasiumEnvPool,
108+
) = py_env(_GardnerChessEnvSpec, _GardnerChessEnvPool)
109+
(
110+
ShogiEnvSpec,
111+
ShogiDMEnvPool,
112+
ShogiGymnasiumEnvPool,
113+
) = py_env(_ShogiEnvSpec, _ShogiEnvPool)
114+
(
115+
SparrowMahjongEnvSpec,
116+
SparrowMahjongDMEnvPool,
117+
SparrowMahjongGymnasiumEnvPool,
118+
) = py_env(_SparrowMahjongEnvSpec, _SparrowMahjongEnvPool)
25119

26120
__all__ = [
121+
"AnimalShogiDMEnvPool",
122+
"AnimalShogiEnvSpec",
123+
"AnimalShogiGymnasiumEnvPool",
124+
"BackgammonDMEnvPool",
125+
"BackgammonEnvSpec",
126+
"BackgammonGymnasiumEnvPool",
127+
"ChessDMEnvPool",
128+
"ChessEnvSpec",
129+
"ChessGymnasiumEnvPool",
130+
"ConnectFourDMEnvPool",
131+
"ConnectFourEnvSpec",
132+
"ConnectFourGymnasiumEnvPool",
133+
"GardnerChessDMEnvPool",
134+
"GardnerChessEnvSpec",
135+
"GardnerChessGymnasiumEnvPool",
27136
"GoEnvSpec",
28137
"GoDMEnvPool",
29138
"GoGymnasiumEnvPool",
139+
"HexDMEnvPool",
140+
"HexEnvSpec",
141+
"HexGymnasiumEnvPool",
142+
"KuhnPokerDMEnvPool",
143+
"KuhnPokerEnvSpec",
144+
"KuhnPokerGymnasiumEnvPool",
145+
"LeducHoldemDMEnvPool",
146+
"LeducHoldemEnvSpec",
147+
"LeducHoldemGymnasiumEnvPool",
148+
"OthelloDMEnvPool",
149+
"OthelloEnvSpec",
150+
"OthelloGymnasiumEnvPool",
151+
"Play2048DMEnvPool",
152+
"Play2048EnvSpec",
153+
"Play2048GymnasiumEnvPool",
154+
"ShogiDMEnvPool",
155+
"ShogiEnvSpec",
156+
"ShogiGymnasiumEnvPool",
157+
"SparrowMahjongDMEnvPool",
158+
"SparrowMahjongEnvSpec",
159+
"SparrowMahjongGymnasiumEnvPool",
160+
"TicTacToeDMEnvPool",
161+
"TicTacToeEnvSpec",
162+
"TicTacToeGymnasiumEnvPool",
30163
]

0 commit comments

Comments
 (0)