Repository: sotetsuk/pgx

Issue:
[Go] Bug in ko
<img width="366" alt="image" src="https://user-images.githubusercontent.com/5868442/224361237-1c3c8d77-cb02-4bf6-b721-7fdb2f7be5f0.png">

白が誤ってkoと判断されて231に置けない


Discussion:

Resolving patch:
diff --git a/pgx/go.py b/pgx/go.py
--- a/pgx/go.py
+++ b/pgx/go.py
@@ -154,6 +154,7 @@ def init(key: jax.random.KeyArray, size: int, komi: float = 7.5) -> State:
 
 
 def step(state: State, action: int, size: int) -> State:
+    state = state.replace(kou=jnp.int32(-1))  # type: ignore
     # update state
     _state = _update_state_wo_legal_action(state, action, size)
 
@@ -193,9 +194,9 @@ def _pass_move(_state: State, _size: int) -> State:
     return jax.lax.cond(
         _state.passed,
         # 連続でパスならば終局
-        lambda: _state.replace(terminated=TRUE, kou=jnp.int32(-1), reward=_get_reward(_state, _size)),  # type: ignore
+        lambda: _state.replace(terminated=TRUE, reward=_get_reward(_state, _size)),  # type: ignore
         # 1回目のパスならばStateにパスを追加してそのまま続行
-        lambda: _state.replace(passed=True, kou=jnp.int32(-1), reward=jnp.zeros(2, dtype=jnp.float32)),  # type: ignore
+        lambda: _state.replace(passed=True, reward=jnp.zeros(2, dtype=jnp.float32)),  # type: ignore
     )
 
 


Test patch:
diff --git a/tests/test_go.py b/tests/test_go.py
--- a/tests/test_go.py
+++ b/tests/test_go.py
@@ -243,6 +243,32 @@ def test_kou():
     state = env.step(state, action=25)  # pass
     assert state.kou == -1
 
+    # see #479
+    actions = [107, 11, 56, 41, 300, 19, 228, 231, 344, 257, 35, 32, 57, 276, 0, 277, 164, 15, 187, 179, 357, 255, 150, 211, 256,
+     190, 297, 303, 358, 189, 322, 3, 129, 64, 13, 336, 22, 286, 264, 192, 55, 360, 23, 31, 113, 119, 195, 98, 208, 294,
+     240, 241, 149, 280, 118, 296, 245, 99, 335, 226, 29, 287, 84, 248, 225, 351, 202, 20, 137, 274, 232, 85, 36, 141,
+     108, 95, 282, 93, 337, 216, 58, 131, 283, 10, 106, 243, 318, 220, 136, 34, 127, 293, 80, 165, 125, 83, 114, 105,
+     30, 61, 147, 71, 109, 173, 87, 233, 76, 361, 66, 115, 212, 200, 346, 197, 54, 326, 298, 167, 347, 4, 354, 16, 140,
+     144, 68, 178, 24, 204, 285, 203, 316, 307, 146, 37, 201, 268, 176, 133, 25, 227, 310, 291, 132, 352, 123, 184, 343,
+     299, 90, 267, 334, 134, 7, 110, 321, 182, 281, 92, 222, 96, 329, 70, 340, 207, 323, 138, 308, 100, 49, 78, 5, 126,
+     317, 17, 349, 160, 261, 266, 306, 221, 355, 327, 324, 284, 236, 60, 359, 174, 252, 46, 260, 114, 163, 235, 250,
+     206, 239, 2, 166, 328, 128, 104, 341, 224, 74, 198, 304, 295, 101, 88, 360, 325, 199, 38, 263, 270, 151, 331, 230,
+     33, 152, 48, 47, 28, 122, 161, 273, 103, 143, 238, 121, 52, 333, 244, 218, 265, 361, 77, 275, 185, 172, 350, 194,
+     59, 53, 21, 272, 319, 320, 158, 251, 253, 135, 27, 196, 180, 188, 345, 254, 130, 42, 156, 259, 332, 361, 18, 82,
+     86, 191, 249, 51, 45, 348, 217, 63, 302, 292, 155, 313, 205, 6, 237, 279, 229, 258, 234, 262, 40, 73, 142, 219,
+     330, 111, 186, 153, 311, 336, 44, 12, 62, 215, 39, 299, 9, 269, 275, 157, 225, 361, 177, 361, 162, 81, 76, 183,
+     168, 247, 309, 145, 210, 221, 65, 301, 1, 289, 120, 315, 353, 305, 67, 214, 79, 314, 290, 47, 181, 346, 175, 89,
+     312, 43, 231, 329, 102, 91, 208, 139, 236, 348, 66, 26, 8, 94, 169, 271, 339, 58, 69, 80, 349, 170, 23, 159, 347,
+     288, 154, 270, 6, 187, 22, 42, 148, 193, 346, 126, 116, 242, 124, 159, 14, 12, 144, 26, 24, 361, 223, 7, 361, 63,
+     117, 112, 5, 81, 118, 135, 82, 92, 140, 123, 97, 278, 47, 361, 137, 230, 220]
+    env = Go(size=19)
+    env.init = jax.jit(env.init)
+    env.step = jax.jit(env.step)
+    state = env.init(jax.random.PRNGKey(0))
+    for a in actions:
+        state = env.step(state, a)
+    assert state.kou == -1
+    assert state.legal_action_mask[231]
 
 def test_observe():
     key = jax.random.PRNGKey(0)