Repository: sotetsuk/pgx

Issue:
Termination in Leduc Holdem
Hello, in Leduc Holdem I don't quite understand why doing the `FOLD` action in the first round of betting doesn't terminate the episode.
`    terminated = round_over & (state._round == 1)`
Seems to me that's a faulty behavior, please tell me if I miss something. A possible fix could be:
`    terminated = (round_over & (state._round == 1)) | (action==FOLD) `






Discussion:
Hi, Thank you for your suggestion! 🙏 

You comment makes sense.  We suppose the Leduc Holdem implementation should follow this diagram from [[Southey+05](https://arxiv.org/ftp/arxiv/papers/1207/1207.1411.pdf)]

![](https://user-images.githubusercontent.com/72956592/227712224-dbcb6523-f81c-4b6a-9ee8-c6611c341ed7.png)

but it does not follow the diagram.

```py
import jax
import jax.numpy as jnp
import pgx

env = pgx.make("leduc_holdem")

init_fn = jax.jit(env.init)
step_fn = jax.jit(env.step)

CALL = jnp.int32(0)
RAISE = jnp.int32(1)
FOLD = jnp.int32(2)

state = init_fn(jax.random.PRNGKey(0))
state = step_fn(state, RAISE)
assert not state.terminated
state = step_fn(state, FOLD)
assert state.terminated  # AssertionError!
```

@OkanoShinri Could you investigate and deal with it?
Thank you for pointing it out!
Sorry, I misunderstood the holdem rules. The current implementation is to proceed to the next round even if the player folds. I will create a PR for the fix immediately.

Resolving patch:
diff --git a/pgx/leduc_holdem.py b/pgx/leduc_holdem.py
--- a/pgx/leduc_holdem.py
+++ b/pgx/leduc_holdem.py
@@ -147,18 +147,20 @@ def _step(state: State, action):
 
 
 def _check_round_over(state, action):
-    round_over = (action == FOLD) | (
-        (state._last_action != INVALID_ACTION) & (action == CALL)
-    )
-    terminated = round_over & (state._round == 1)
+    fold = action == FOLD
+    call = (state._last_action != INVALID_ACTION) & (action == CALL)
+    _continue = (state._round == 0) & call
+
+    round_over = fold | call
+    terminated = round_over & (~_continue)
 
     reward = jax.lax.select(
-        terminated & (action == FOLD),
+        fold,
         jnp.float32([-1, -1]).at[1 - state.current_player].set(1),
         jnp.float32([0, 0]),
     )
     reward = jax.lax.select(
-        terminated & (action != FOLD),
+        terminated & call,
         _get_unit_reward(state),
         reward,
     )


Test patch:
diff --git a/tests/test_leduc_holdem.py b/tests/test_leduc_holdem.py
--- a/tests/test_leduc_holdem.py
+++ b/tests/test_leduc_holdem.py
@@ -18,6 +18,46 @@ def test_init():
 def test_step():
     key = jax.random.PRNGKey(1)
 
+    # =================
+    state = init(key)
+    assert state.current_player == 1
+    # cards = [1 0]
+    # player 1 is the first
+
+    # first round
+    state = step(state, RAISE)
+    state = step(state, FOLD)
+    assert state.terminated
+    assert (state.rewards == jnp.float32([-1, 1])).all()
+
+    # =================
+    state = init(key)
+    assert state.current_player == 1
+    # cards = [1 0]
+    # player 1 is the first
+
+    # first round
+    state = step(state, RAISE)  # +2(3)
+    state = step(state, RAISE)  # +2(5)
+    state = step(state, FOLD)
+    assert state.terminated
+    assert (state.rewards == jnp.float32([3, -3])).all()
+
+    # =================
+    state = init(key)
+    assert state.current_player == 1
+    # cards = [1 0]
+    # player 1 is the first
+
+    # first round
+    state = step(state, CALL)
+    state = step(state, RAISE)  # +2(3)
+    state = step(state, RAISE)  # +2(5)
+    state = step(state, FOLD)
+    assert state.terminated
+    assert (state.rewards == jnp.float32([-3, 3])).all()
+
+    # =================
     state = init(key)
     assert state.current_player == 1
     # cards = [1 0]
@@ -38,6 +78,7 @@ def test_step():
     assert state.terminated
     assert (state.rewards == jnp.float32([11, -11])).all()
 
+    # =================
     state = init(key)
     assert state.current_player == 1
     # cards = [1 0]