pgx

用于强化学习（RL）的GPU/TPU加速并行游戏模拟器集合

为什么选择Pgx？

Brax，一个JAX原生的物理引擎，为连续状态空间的RL提供了极高速的并行模拟。那么，对于像国际象棋、将棋和围棋这样的离散状态空间的RL呢？Pgx提供了各种JAX原生的游戏模拟器！主要特点包括：

⚡ 在加速器上进行并行执行时超级快速
🎲 支持多种游戏，包括西洋双陆棋、国际象棋、将棋和围棋
🖼️ SVG格式的精美可视化

快速开始

训练示例

使用方法

Pgx可在PyPI上获取。请注意，您的Python环境需要根据硬件规格安装jax和jaxlib。

$ pip install pgx

以下代码片段展示了使用Pgx的简单示例。您可以在这个Colab中尝试。请注意，Pgx环境中的所有step函数都是JAX原生的，即它们都是可JIT编译的。更多详情请参阅文档。

import jax
import pgx

env = pgx.make("go_19x19")
init = jax.jit(jax.vmap(env.init))
step = jax.jit(jax.vmap(env.step))

batch_size = 1024 keys = jax.random.split(jax.random.PRNGKey(42), batch_size) state = init(keys) # 向量化状态 while not (state.terminated | state.truncated).all(): action = model(state.current_player, state.observation, state.legal_action_mask) # 对于随机环境使用 step(state, action, keys) state = step(state, action) # state.rewards 形状为 (1024, 2)

Pgx 是一个专注于更快实现而非仅仅关注 API 本身的库。然而，其 API 本身也足够通用。例如，Pgx 中的所有环境都可以转换为 PettingZoo 的 AEC API，您可以通过 PettingZoo API 运行 Pgx 环境。您可以在这个 Colab 中查看演示。

Pgx 已于 2023 年 11 月 8 日从 API v1 更新至 v2（发布版本 v2.0.0）。因此，Env.step 的签名发生了如下变化：

v1: step(state: State, action: Array)
v2: step(state: State, action: Array, key: Optional[PRNGKey] = None)

此外，pgx.experimental.auto_reset 也改为将 key 指定为第三个参数。

更新目的： 在 API v1 中，即使在具有随机状态转换的环境中，状态转换也是由 state 内部的 _rng_key 决定的确定性转换。这是有意为之，目的是提高可重现性。然而，在这种环境中使用规划算法时，存在底层真实随机性信息可能"泄露"的风险。为了让用户更容易进行正确的实验，Env.step 已更改为显式指定 key。

更新影响： 由于 key 是可选的，在像围棋和国际象棋这样的确定性环境中，仍然可以像 API v1 那样执行 env.step(state, action)，所以这些游戏不受影响。截至 v2.0.0，只有 2048、双陆棋和 MinAtar 套件受此更改影响。

</details>

支持的游戏

双陆棋	国际象棋	将棋	围棋
<img src="https://yellow-cdn.veclightyear.com/0a4dffa0/456923a8-5e62-479c-8c6f-f5c8aabd44bf.gif#gh-dark-mode-only" width="170px"><img src="https://yellow-cdn.veclightyear.com/0a4dffa0/8a603c6f-13e9-40b1-9a37-e6f8ffffc987.gif#gh-light-mode-only" width="170px">	<img src="https://yellow-cdn.veclightyear.com/0a4dffa0/f98f8076-0b6b-4cd0-8dc8-8bab7be373d4.gif#gh-dark-mode-only" width="158px"><img src="https://yellow-cdn.veclightyear.com/0a4dffa0/a3f6c951-3500-4c22-af0e-a3fde3ccf60e.gif#gh-light-mode-only" width="158px">	<img src="https://yellow-cdn.veclightyear.com/0a4dffa0/3dbd00f8-c4ff-4df2-8d31-1b7cc3d60503.gif#gh-dark-mode-only" width="170px"><img src="https://yellow-cdn.veclightyear.com/0a4dffa0/663094b3-e451-41cd-9354-e68249955477.gif#gh-light-mode-only" width="170px">	<img src="https://yellow-cdn.veclightyear.com/0a4dffa0/dceb6721-4a30-4053-87ce-1f7ef8e1dd34.gif#gh-dark-mode-only" width="160px"><img src="https://yellow-cdn.veclightyear.com/0a4dffa0/319c8f53-8b10-440d-a253-168a55c07b7d.gif#gh-light-mode-only" width="160px">

使用 pgx.available_envs() -> Tuple[EnvId] 可以查看当前可用游戏的列表。给定一个 <EnvId>，您可以通过以下方式创建环境：

>>> env = pgx.make(<EnvId>)

游戏/环境ID	可视化	版本	ChatGPT给出的五字描述
<a href="https://en.wikipedia.org/wiki/2048_(video_game)">2048</a> <br> `"2048"`	<img src="https://yellow-cdn.veclightyear.com/0a4dffa0/cbf71ed8-c922-4a76-953f-d327e811decd.gif" width="60px"><img src="https://yellow-cdn.veclightyear.com/0a4dffa0/bfae7c97-08ac-4967-83e3-464ec9ccc6c0.gif" width="60px">	`v2`	合并方块创建2048。
<a href="https://en.wikipedia.org/wiki/D%C5%8Dbutsu_sh%C5%8Dgi">动物将棋</a><br>`"animal_shogi"`	<img src="https://yellow-cdn.veclightyear.com/0a4dffa0/16f3cd53-20a3-403b-b65e-b8ca69371009.gif" width="60px"><img src="https://yellow-cdn.veclightyear.com/0a4dffa0/5ca4e208-4a27-490f-a50c-e9037dcf7310.gif" width="60px">	`v0`	适合儿童的动物主题将棋。
<a href="https://en.wikipedia.org/wiki/Backgammon">西洋双陆棋</a><br>`"backgammon"`	<img src="https://yellow-cdn.veclightyear.com/0a4dffa0/456923a8-5e62-479c-8c6f-f5c8aabd44bf.gif" width="60px"><img src="https://yellow-cdn.veclightyear.com/0a4dffa0/8a603c6f-13e9-40b1-9a37-e6f8ffffc987.gif" width="60px">	`v2`	运气助力移除棋子获胜。
<a href="https://en.wikipedia.org/wiki/Contract_bridge">桥牌叫牌</a><br>`"bridge_bidding"`	<img src="https://yellow-cdn.veclightyear.com/0a4dffa0/a4b289a8-0b49-49c5-9d53-75ebcc5b990b.gif" width="60px"><img src="https://yellow-cdn.veclightyear.com/0a4dffa0/fdd8f09c-0b98-4655-9729-d029cc6eb00e.gif" width="60px">	`v1`	搭档通过叫牌交换信息。
<a href="https://en.wikipedia.org/wiki/Chess">国际象棋</a><br>`"chess"`	<img src="https://yellow-cdn.veclightyear.com/0a4dffa0/f98f8076-0b6b-4cd0-8dc8-8bab7be373d4.gif" width="60px"><img src="https://yellow-cdn.veclightyear.com/0a4dffa0/a3f6c951-3500-4c22-af0e-a3fde3ccf60e.gif" width="60px">	`v2`	将死对手国王获胜。
<a href="https://en.wikipedia.org/wiki/Connect_Four">四子棋</a><br>`"connect_four"`	<img src="https://yellow-cdn.veclightyear.com/0a4dffa0/4d8914b9-fe75-49c6-9599-f58825d0a639.gif" width="60px"><img src="https://yellow-cdn.veclightyear.com/0a4dffa0/117e9f9b-89d9-4de3-8a4d-6fcbc44dd245.gif" width="60px">	`v0`	连接棋子，四子成线。
<a href="https://en.wikipedia.org/wiki/Minichess">加德纳象棋</a><br>`"gardner_chess"`	<img src="https://yellow-cdn.veclightyear.com/0a4dffa0/db3e106e-fb6b-41ec-aafc-ca1c72f1fc1f.gif" width="60px"><img src="https://yellow-cdn.veclightyear.com/0a4dffa0/23b7cae7-d6b6-41f8-9f9f-483b42116e2e.gif" width="60px">	`v0`	5x5变体象棋，无王车易位。
<a href="https://en.wikipedia.org/wiki/Go_(game)">围棋</a><br>`"go_9x9"` `"go_19x19"`	<img src="https://yellow-cdn.veclightyear.com/0a4dffa0/dceb6721-4a30-4053-87ce-1f7ef8e1dd34.gif" width="60px"><img src="https://yellow-cdn.veclightyear.com/0a4dffa0/319c8f53-8b10-440d-a253-168a55c07b7d.gif" width="60px">	`v0`	策略性布子占领领地。
<a href="https://en.wikipedia.org/wiki/Hex_(board_game)">六角棋</a><br>`"hex"`	<img src="https://yellow-cdn.veclightyear.com/0a4dffa0/14dfe1c4-3f08-4b0f-8515-4ac6446a0378.gif" width="60px"><img src="https://yellow-cdn.veclightyear.com/0a4dffa0/7e78a93a-f85b-44b5-b241-d3ff4575b51b.gif" width="60px">	`v0`	连接对面，阻挡对手。
<a href="https://en.wikipedia.org/wiki/Kuhn_poker">库恩扑克</a><br>`"kuhn_poker"`	<img src="https://yellow-cdn.veclightyear.com/0a4dffa0/54605c45-fb33-45e9-859e-be07c65b9710.gif" width="60px"><img src="https://yellow-cdn.veclightyear.com/0a4dffa0/fcdc1a05-0e10-4957-b1a4-b01086e96c28.gif" width="60px">	`v1`	三张牌的下注和虚张声势游戏。
<a href="https://arxiv.org/abs/1207.1411">勒杜克扑克</a><br>`"leduc_holdem"`	<img src="https://yellow-cdn.veclightyear.com/0a4dffa0/a68617eb-000e-4629-9406-16e3be344732.gif" width="60px"><img src="https://yellow-cdn.veclightyear.com/0a4dffa0/bef71323-1386-403d-a4ca-9a408da1c2ba.gif" width="60px">	`v0`	两种花色，限制牌数的扑克。
<a href="https://github.com/kenjyoung/MinAtar">MinAtar/阿斯特里克斯</a><br>`"minatar-asterix"`	<img src="https://yellow-cdn.veclightyear.com/0a4dffa0/105db10e-2990-43b1-b8a0-4506e95cc11b.gif" width="50px">	`v1`	躲避敌人，收集宝藏，生存。
<a href="https://github.com/kenjyoung/MinAtar">MinAtar/打砖块</a><br>`"minatar-breakout"`	<img src="https://yellow-cdn.veclightyear.com/0a4dffa0/d615f297-5fbf-4bb8-a332-04f600bddd62.gif" width="50px">	`v1`	球拍，球，砖块，反弹，清除。
<a href="https://github.com/kenjyoung/MinAtar">MinAtar/过马路</a><br>`"minatar-freeway"`	<img src="https://yellow-cdn.veclightyear.com/0a4dffa0/a42fb97a-b810-4cd5-b00a-4e6524bbf4d8.gif" width="50px">	`v1`	躲避汽车，爬上高速公路。
<a href="https://github.com/kenjyoung/MinAtar">MinAtar/海底探险</a><br>`"minatar-seaquest"`	<img src="https://yellow-cdn.veclightyear.com/0a4dffa0/2610176a-384e-4ed5-ab7e-c860e49eb071.gif" width="50px">	`v1`	水下潜艇救援和战斗。
<a href="https://github.com/kenjyoung/MinAtar">MinAtar/太空侵略者</a><br>`"minatar-space_invaders"`	<img src="https://yellow-cdn.veclightyear.com/0a4dffa0/e1defc8a-e958-4207-ab6e-d3a90ab6db99.gif" width="50px">	`v1`	外星人射击游戏，躲避子弹。
<a href="https://en.wikipedia.org/wiki/Reversi">黑白棋</a><br>`"othello"`	<img src="https://yellow-cdn.veclightyear.com/0a4dffa0/75503e89-132e-4345-b246-385677da26ec.gif" width="60px"><img src="https://yellow-cdn.veclightyear.com/0a4dffa0/e2d05b5e-9c27-4a70-8ad2-1152410c85b7.gif" width="60px">	`v0`	翻转并征服对手的棋子。
<a href="https://en.wikipedia.org/wiki/Shogi">将棋</a><br>`"shogi"`	<img src="https://yellow-cdn.veclightyear.com/0a4dffa0/3dbd00f8-c4ff-4df2-8d31-1b7cc3d60503.gif" width="60px"><img src="https://yellow-cdn.veclightyear.com/0a4dffa0/663094b3-e451-41cd-9354-e68249955477.gif" width="60px">	`v0`	日本象棋，可使用被俘棋子。
<a href="https://sugorokuya.jp/p/suzume-jong">雀麻将</a><br>`"sparrow_mahjong"`	<img src="https://yellow-cdn.veclightyear.com/0a4dffa0/01aa5f5f-6957-4375-8a00-59aee266462e.svg" width="60px"><img src="https://yellow-cdn.veclightyear.com/0a4dffa0/962e77cc-6ba8-4f86-bcc4-64bd08520aa1.svg" width="60px">	`v1`	一种简化的、适合儿童的麻将游戏。
<a href="https://en.wikipedia.org/wiki/Tic-tac-toe">井字棋</a><br>`"tic_tac_toe"`	<img src="https://yellow-cdn.veclightyear.com/0a4dffa0/7972531b-a720-425b-9142-2b4579ab25bf.gif" width="60px"><img src="https://yellow-cdn.veclightyear.com/0a4dffa0/e6a028a5-e108-4c1b-ad40-13bc7025e731.gif" width="60px">	`v0`	连成一行三个即可获胜。

<details><summary>版本控制策略</summary>

每个环境都有版本号，当有影响代理性能的变更或与API不向后兼容的变更时，版本号会递增。如果您想追求完全的可重现性，我们建议您按以下方式检查Pgx和每个环境的版本：

>>> pgx.__version__
'1.0.0'
>>> env.version
'v0'

</details>

另请参阅

Pgx旨在通过(经典)棋类游戏套件来补充这些基于JAX的环境：

RobertTLange/gymnax：流行RL环境的JAX实现（经典控制、bsuite、MinAtar等）和元RL任务
google/brax：JAX中的刚体物理模拟和连续空间RL任务（蚂蚁、抓取、人形等）
instadeepai/jumanji：JAX中多样化且具有挑战性的RL环境套件（装箱问题、路由问题等）
flairox/jaxmarl：JAX中的多智能体RL环境（简化版星际争霸等）
corl-team/xland-minigrid：受MiniGrid和XLand启发的JAX元RL网格世界环境

将Pgx与这些基于JAX的算法/实现结合可能是一个有趣的方向：

Anakin框架：高效的RL框架，可在TPU上与基于JAX的环境协同工作
deepmind/mctx：基于JAX的MCTS实现，包括AlphaZero和MuZero
deepmind/rlax：基于JAX的RL组件
google/evojax：硬件加速的神经进化
RobertTLange/evosax：基于JAX的进化策略（ES）实现
adaptive-intelligent-robotics/QDax：基于JAX的质量多样性（QD）算法
luchris429/purejaxrl：基于JAX的RL实现

引用

如果您在工作中使用了Pgx，请引用我们的论文：

@inproceedings{koyamada2023pgx,
  title={Pgx: Hardware-Accelerated Parallel Game Simulators for Reinforcement Learning},
  author={Koyamada, Sotetsu and Okano, Shinri and Nishimori, Soichiro and Murata, Yu and Habara, Keigo and Kita, Haruka and Ishii, Shin},
  booktitle={Advances in Neural Information Processing Systems},
  pages={45716--45743},
  volume={36},
  year={2023}
}