Repository: sotetsuk/pgx

Issue:
[Go] `-1` の行動を受け付ける仕様を変更する
行動はすべて0以上とする


Discussion:
例えば5x5の場合、25をパスとして扱い、legal_actionを0~25の26次元にします
それがいいですね！でも優先順位的には高速化のあとで大丈夫です！ @OkanoShinri 

Resolving patch:
diff --git a/pgx/go.py b/pgx/go.py
--- a/pgx/go.py
+++ b/pgx/go.py
@@ -39,7 +39,7 @@ class GoState:
     adj_ren_id: jnp.ndarray = jnp.zeros((2, 19 * 19, 19 * 19), dtype=jnp.bool_)
 
     # 設置可能なマスをTrueとしたマスク
-    legal_action_mask: jnp.ndarray = jnp.zeros(19 * 19, dtype=jnp.bool_)
+    legal_action_mask: jnp.ndarray = jnp.zeros(19 * 19 + 1, dtype=jnp.bool_)
 
     # 直近8回のログ
     game_log: jnp.ndarray = jnp.full(
@@ -147,7 +147,7 @@ def init(
         available_ren_id=jnp.ones((2, size * size), dtype=jnp.bool_),
         liberty=jnp.zeros((2, size * size, size * size), dtype=jnp.int32),
         adj_ren_id=jnp.zeros((2, size * size, size * size), dtype=jnp.bool_),
-        legal_action_mask=jnp.ones(size * size, dtype=jnp.bool_),
+        legal_action_mask=jnp.ones(size * size + 1, dtype=jnp.bool_),
         game_log=jnp.full((8, size * size), 2, dtype=jnp.int32),  # type:ignore
         curr_player=curr_player,  # type:ignore
     )
@@ -176,7 +176,7 @@ def _update_state_wo_legal_action(
     _state: GoState, _action: int, _size: int
 ) -> Tuple[GoState, jnp.ndarray]:
     _state, _reward = jax.lax.cond(
-        (0 <= _action) & (_action < _size * _size),
+        (_action < _size * _size),
         lambda: _not_pass_move(_state, _action),
         lambda: _pass_move(_state, _size),
     )
@@ -451,14 +451,14 @@ def _remove_stones(_state: GoState, _rm_ren_id, _rm_stone_xy) -> GoState:
 
 
 def legal_actions(state: GoState, size: int) -> jnp.ndarray:
-    return jnp.logical_not(
-        jax.lax.map(
-            lambda xy: _update_state_wo_legal_action(state, xy, size)[
-                0
-            ].terminated,
-            jnp.arange(0, size * size),
-        )
-    )
+    illegal_action = jax.lax.map(
+        lambda xy: _update_state_wo_legal_action(state, xy, size)[
+            0
+        ].terminated,
+        jnp.arange(0, size * size + 1),
+    )
+    legal_action = ~illegal_action
+    return legal_action.at[size * size].set(TRUE)
 
 
 def get_board(state: GoState) -> jnp.ndarray:


Test patch:
diff --git a/tests/test_go.py b/tests/test_go.py
--- a/tests/test_go.py
+++ b/tests/test_go.py
@@ -19,16 +19,16 @@ def test_end_by_pass():
     rng = jax.random.PRNGKey(0)
 
     _, state = j_init(rng=rng, size=BOARD_SIZE)
-    _, state, _ = j_step(state=state, action=-1, size=BOARD_SIZE)
+    _, state, _ = j_step(state=state, action=25, size=BOARD_SIZE)
     assert state.passed
     assert not state.terminated
     _, state, _ = j_step(state=state, action=0, size=BOARD_SIZE)
     assert not state.passed
     assert not state.terminated
-    _, state, _ = j_step(state=state, action=-1, size=BOARD_SIZE)
+    _, state, _ = j_step(state=state, action=25, size=BOARD_SIZE)
     assert state.passed
     assert not state.terminated
-    _, state, _ = j_step(state=state, action=-1, size=BOARD_SIZE)
+    _, state, _ = j_step(state=state, action=25, size=BOARD_SIZE)
     assert state.passed
     assert state.terminated
 
@@ -56,12 +56,12 @@ def test_step():
     _, state, _ = j_step(state=state, action=15, size=BOARD_SIZE)
     _, state, _ = j_step(state=state, action=23, size=BOARD_SIZE)
     _, state, _ = j_step(state=state, action=20, size=BOARD_SIZE)
-    _, state, _ = j_step(state=state, action=-1, size=BOARD_SIZE)  # pass
+    _, state, _ = j_step(state=state, action=25, size=BOARD_SIZE)  # pass
     _, state, _ = j_step(state=state, action=22, size=BOARD_SIZE)
     _, state, _ = j_step(state=state, action=19, size=BOARD_SIZE)
     _, state, _ = j_step(state=state, action=21, size=BOARD_SIZE)
-    _, state, _ = j_step(state=state, action=-1, size=BOARD_SIZE)  # pass
-    _, state, reward = j_step(state=state, action=-1, size=BOARD_SIZE)  # pass
+    _, state, _ = j_step(state=state, action=25, size=BOARD_SIZE)  # pass
+    _, state, reward = j_step(state=state, action=25, size=BOARD_SIZE)  # pass
 
     expected_board: jnp.ndarray = jnp.array(
         [
@@ -211,7 +211,7 @@ def test_legal_action():
         False, False, False, False, False,
         False, True, False, True, False,
         True, True, True, True, True,
-        True, True, True, True, True])
+        True, True, True, True, True, True])
     # fmt:on
     _, state = j_init(rng=rng, size=BOARD_SIZE)
     _, state, _ = j_step(state=state, action=0, size=BOARD_SIZE)  # BLACK
@@ -243,7 +243,7 @@ def test_legal_action():
         False, False, False, False, False,
         True, False, False, False, True,
         True, True, True, True, True,
-        True, True, True, True, True])
+        True, True, True, True, True, True])
     # fmt:on
     # white 8
     _, state = j_init(rng=rng, size=BOARD_SIZE)
@@ -316,13 +316,13 @@ def test_legal_action():
         True, False, False, False, True,
         False, False, False, False, False,
         True, False, False, False, True,
-        True, True, False, True, True])
+        True, True, False, True, True, True])
     expected_w = jnp.array([
         True, True, False, True, True,
         True, False, False, False, True,
         False, False, True, False, False,
         True, False, False, False, True,
-        True, True, False, True, True])
+        True, True, False, True, True, True])
     # fmt:on
     _, state = j_init(rng=rng, size=BOARD_SIZE)
     _, state, _ = j_step(state=state, action=7, size=BOARD_SIZE)  # BLACK
@@ -358,21 +358,21 @@ def test_legal_action():
         False, False, False, False, False,
         False, False, False, False, False,
         False, False, False, False, False,
-        False, False, False, False, False])
+        False, False, False, False, False, True])
     # white pass
     expected_b = jnp.array([
         True, False, False, False, True,
         False, False, False, True, False,
         False, False, False, False, False,
         False, False, False, False, False,
-        False, False, False, False, False])
+        False, False, False, False, False, True])
     # black 8
     expected_w2 = jnp.array([
         False, False, False, False, False,
         False, True, False, False, False,
         False, True, False, True, False,
         False, True, False, True, False,
-        False, True, True, True, False])
+        False, True, True, True, False, True])
     # fmt:on
     _, state = j_init(rng=rng, size=BOARD_SIZE)
     _, state, _ = j_step(state=state, action=1, size=BOARD_SIZE)  # BLACK
@@ -419,7 +419,7 @@ def test_legal_action():
             key = jax.random.PRNGKey(0)
             key, subkey = jax.random.split(key)
             a = jax.random.choice(subkey, actions[0])
-        for action in actions[0]:
+        for action in actions[0][:-1]:
             _, _state, _ = j_step(state=state, action=action, size=BOARD_SIZE)
             assert not _state.terminated
 
@@ -432,7 +432,7 @@ def test_random_play_5():
     while not state.terminated:
         actions = np.where(state.legal_action_mask)
         if len(actions[0]) == 0:
-            a = -1
+            a = 25
         else:
             key = jax.random.PRNGKey(0)
             key, subkey = jax.random.split(key)
@@ -442,7 +442,7 @@ def test_random_play_5():
 
         if state.turn > 100:
             break
-    assert state.turn > 100
+    assert state.passed or state.turn > 100
 
 
 def test_random_play_19():
@@ -453,7 +453,7 @@ def test_random_play_19():
     while not state.terminated:
         actions = np.where(state.legal_action_mask)
         if len(actions[0]) == 0:
-            a = -1
+            a = 19 * 19
         else:
             key = jax.random.PRNGKey(0)
             key, subkey = jax.random.split(key)
@@ -463,4 +463,4 @@ def test_random_play_19():
 
         if state.turn > 100:
             break
-    assert state.turn > 100
+    assert state.passed or state.turn > 100