Repository: sotetsuk/pgx

Issue:
[Kuhn Poker] Simplify the implementation
Hey, multiple implementations of Kuhn Poker out there ([open_spiel ](https://github.com/google-deepmind/open_spiel/blob/b8c2ff8e9a4f5dad9b179217f740ddb0df967f7c/open_spiel/games/kuhn_poker.cc)for instance) use only two actions (pass, bet) instead of the four considered in pgx (call, bet, check, fold). In fact we can group bet/call and check/fold without ambiguity for this game. 

Would you be interested by this simplification? I would be happy to open a PR if you are!


Discussion:
Hi, thank you for your comment! 🙏
I'm positive for this change.
There are only five scenarios. Four actions are redundant.

If you kindly submit a PR, we'll receive it.
We also really appreciate if you can list up some implementations or papers which use two actions to support this change (in addition to open spiel)


Resolving patch:
diff --git a/pgx/kuhn_poker.py b/pgx/kuhn_poker.py
--- a/pgx/kuhn_poker.py
+++ b/pgx/kuhn_poker.py
@@ -21,10 +21,8 @@
 
 FALSE = jnp.bool_(False)
 TRUE = jnp.bool_(True)
-CALL = jnp.int32(0)
-BET = jnp.int32(1)
-FOLD = jnp.int32(2)
-CHECK = jnp.int32(3)
+BET = jnp.int32(0)
+PASS = jnp.int32(1)
 
 
 @dataclass
@@ -34,13 +32,13 @@ class State(core.State):
     rewards: Array = jnp.float32([0.0, 0.0])
     terminated: Array = FALSE
     truncated: Array = FALSE
-    legal_action_mask: Array = jnp.ones(4, dtype=jnp.bool_)
+    legal_action_mask: Array = jnp.ones(2, dtype=jnp.bool_)
     _step_count: Array = jnp.int32(0)
     # --- Kuhn poker specific ---
     _cards: Array = jnp.int32([-1, -1])
     # [(player 0),(player 1)]
     _last_action: Array = jnp.int32(-1)
-    # 0(Call)  1(Bet)  2(Fold)  3(Check)
+    # 0(Bet) 1(Pass)
     _pot: Array = jnp.int32([0, 0])
 
     @property
@@ -84,20 +82,20 @@ def _init(rng: PRNGKey) -> State:
     return State(  # type:ignore
         current_player=current_player,
         _cards=init_card,
-        legal_action_mask=jnp.bool_([0, 1, 0, 1]),
+        legal_action_mask=jnp.bool_([1, 1]),
     )
 
 
 def _step(state: State, action):
     action = jnp.int32(action)
     pot = jax.lax.cond(
-        (action == BET) | (action == CALL),
+        (action == BET),
         lambda: state._pot.at[state.current_player].add(1),
         lambda: state._pot,
     )
 
     terminated, reward = jax.lax.cond(
-        action == FOLD,
+        (state._last_action == BET) & (action == PASS),
         lambda: (
             TRUE,
             jnp.float32([-1, -1]).at[1 - state.current_player].set(1),
@@ -105,25 +103,17 @@ def _step(state: State, action):
         lambda: (FALSE, jnp.float32([0, 0])),
     )
     terminated, reward = jax.lax.cond(
-        (state._last_action == BET) & (action == CALL),
+        (state._last_action == BET) & (action == BET),
         lambda: (TRUE, _get_unit_reward(state) * 2),
         lambda: (terminated, reward),
     )
     terminated, reward = jax.lax.cond(
-        (state._last_action == CHECK) & (action == CHECK),
+        (state._last_action == PASS) & (action == PASS),
         lambda: (TRUE, _get_unit_reward(state)),
         lambda: (terminated, reward),
     )
 
-    legal_action = jax.lax.switch(
-        action,
-        [
-            lambda: jnp.bool_([0, 0, 0, 0]),  # CALL
-            lambda: jnp.bool_([1, 0, 1, 0]),  # BET
-            lambda: jnp.bool_([0, 0, 0, 0]),  # FOLD
-            lambda: jnp.bool_([0, 1, 0, 1]),  # CHECK
-        ],
-    )
+    legal_action = jax.lax.select(terminated, jnp.bool_([0, 0]), jnp.bool_([1, 1]))
 
     return state.replace(  # type:ignore
         current_player=1 - state.current_player,


Test patch:
diff --git a/tests/test_kuhn_poker.py b/tests/test_kuhn_poker.py
--- a/tests/test_kuhn_poker.py
+++ b/tests/test_kuhn_poker.py
@@ -1,6 +1,7 @@
 import jax
 import jax.numpy as jnp
-from pgx.kuhn_poker import KuhnPoker, CALL, BET, FOLD, CHECK
+
+from pgx.kuhn_poker import BET, PASS, KuhnPoker
 
 env = KuhnPoker()
 init = jax.jit(env.init)
@@ -12,48 +13,48 @@ def test_init():
     key = jax.random.PRNGKey(0)
     state = init(key=key)
     assert state._cards[0] != state._cards[1]
-    assert (state.legal_action_mask == jnp.bool_([0, 1, 0, 1])).all()
+    assert (state.legal_action_mask == jnp.bool_([1, 1])).all()
 
 
 def test_step():
     key = jax.random.PRNGKey(0)
     # cards = [2, 0]
     state = init(key)
-    state = step(state, CHECK)
+    state = step(state, PASS)
     assert not state.terminated
-    state = step(state, CHECK)
+    state = step(state, PASS)
     assert state.terminated
     assert (state.rewards == jnp.float32([1, -1])).all()
 
     state = init(key)
-    state = step(state, CHECK)
+    state = step(state, PASS)
     assert not state.terminated
     state = step(state, BET)
     assert not state.terminated
-    state = step(state, FOLD)
+    state = step(state, PASS)
     assert state.terminated
     assert (state.rewards == jnp.float32([-1, 1])).all()
 
     state = init(key)
-    state = step(state, CHECK)
+    state = step(state, PASS)
     assert not state.terminated
     state = step(state, BET)
     assert not state.terminated
-    state = step(state, CALL)
+    state = step(state, BET)
     assert state.terminated
     assert (state.rewards == jnp.float32([2, -2])).all()
 
     state = init(key)
     state = step(state, BET)
     assert not state.terminated
-    state = step(state, FOLD)
+    state = step(state, PASS)
     assert state.terminated
     assert (state.rewards == jnp.float32([1, -1])).all()
 
     state = init(key)
     state = step(state, BET)
     assert not state.terminated
-    state = step(state, CALL)
+    state = step(state, BET)
     assert state.terminated
     assert (state.rewards == jnp.float32([2, -2])).all()
 
@@ -62,37 +63,37 @@ def test_legal_action():
     key = jax.random.PRNGKey(0)
     # cards = [2, 0]
     state = init(key)
-    state = step(state, CHECK)
-    assert (state.legal_action_mask == jnp.bool_([0, 1, 0, 1])).all()
-    state = step(state, CHECK)
+    state = step(state, PASS)
+    assert (state.legal_action_mask == jnp.bool_([1, 1])).all()
+    state = step(state, PASS)
     assert state.terminated
 
     state = init(key)
-    state = step(state, CHECK)
-    assert (state.legal_action_mask == jnp.bool_([0, 1, 0, 1])).all()
+    state = step(state, PASS)
+    assert (state.legal_action_mask == jnp.bool_([1, 1])).all()
     state = step(state, BET)
-    assert (state.legal_action_mask == jnp.bool_([1, 0, 1, 0])).all()
-    state = step(state, FOLD)
+    assert (state.legal_action_mask == jnp.bool_([1, 1])).all()
+    state = step(state, PASS)
     assert state.terminated
 
     state = init(key)
-    state = step(state, CHECK)
-    assert (state.legal_action_mask == jnp.bool_([0, 1, 0, 1])).all()
+    state = step(state, PASS)
+    assert (state.legal_action_mask == jnp.bool_([1, 1])).all()
+    state = step(state, BET)
+    assert (state.legal_action_mask == jnp.bool_([1, 1])).all()
     state = step(state, BET)
-    assert (state.legal_action_mask == jnp.bool_([1, 0, 1, 0])).all()
-    state = step(state, CALL)
     assert state.terminated
 
     state = init(key)
     state = step(state, BET)
-    assert (state.legal_action_mask == jnp.bool_([1, 0, 1, 0])).all()
-    state = step(state, FOLD)
+    assert (state.legal_action_mask == jnp.bool_([1, 1])).all()
+    state = step(state, PASS)
     assert state.terminated
 
     state = init(key)
     state = step(state, BET)
-    assert (state.legal_action_mask == jnp.bool_([1, 0, 1, 0])).all()
-    state = step(state, CALL)
+    assert (state.legal_action_mask == jnp.bool_([1, 1])).all()
+    state = step(state, BET)
     assert state.terminated
 
 
@@ -114,6 +115,7 @@ def test_observation():
 
 def test_api():
     import pgx
+
     env = pgx.make("kuhn_poker")
     pgx.api_test(env, 3, use_key=False)
     pgx.api_test(env, 3, use_key=True)