pgx - JAX原生并行游戏模拟器库用于强化学习研究

用于强化学习（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 中查看演示。

📣 API v2 (v2.0.0)

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 套件受此更改影响。

支持的游戏

双陆棋	国际象棋	将棋	围棋

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

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

游戏/环境ID	版本	ChatGPT给出的五字描述
2048 `"2048"`	`v2`	合并方块创建2048。
动物将棋 `"animal_shogi"`	`v0`	适合儿童的动物主题将棋。
西洋双陆棋 `"backgammon"`	`v2`	运气助力移除棋子获胜。
桥牌叫牌 `"bridge_bidding"`	`v1`	搭档通过叫牌交换信息。
国际象棋 `"chess"`	`v2`	将死对手国王获胜。
四子棋 `"connect_four"`	`v0`	连接棋子，四子成线。
加德纳象棋 `"gardner_chess"`	`v0`	5x5变体象棋，无王车易位。
围棋 `"go_9x9"` `"go_19x19"`	`v0`	策略性布子占领领地。
六角棋 `"hex"`	`v0`	连接对面，阻挡对手。
库恩扑克 `"kuhn_poker"`	`v1`	三张牌的下注和虚张声势游戏。
勒杜克扑克 `"leduc_holdem"`	`v0`	两种花色，限制牌数的扑克。
MinAtar/阿斯特里克斯 `"minatar-asterix"`	`v1`	躲避敌人，收集宝藏，生存。
MinAtar/打砖块 `"minatar-breakout"`	`v1`	球拍，球，砖块，反弹，清除。
MinAtar/过马路 `"minatar-freeway"`	`v1`	躲避汽车，爬上高速公路。
MinAtar/海底探险 `"minatar-seaquest"`	`v1`	水下潜艇救援和战斗。
MinAtar/太空侵略者 `"minatar-space_invaders"`	`v1`	外星人射击游戏，躲避子弹。
黑白棋 `"othello"`	`v0`	翻转并征服对手的棋子。
将棋 `"shogi"`	`v0`	日本象棋，可使用被俘棋子。
雀麻将 `"sparrow_mahjong"`	`v1`	一种简化的、适合儿童的麻将游戏。
井字棋 `"tic_tac_toe"`	`v0`	连成一行三个即可获胜。

版本控制策略

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

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

另请参阅

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}
}

许可证

Apache-2.0