We are currently solving the following issue within our repository. Here is the issue text:
--- BEGIN 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) `





--- END ISSUE ---

Below are some code segments, each from a relevant file. One or more of these files may contain bugs.
--- BEGIN FILES ---
Path: `pgx/leduc_holdem.py`
Content:
```
1 # Copyright 2023 The Pgx Authors. All Rights Reserved.
2 #
3 # Licensed under the Apache License, Version 2.0 (the "License");
4 # you may not use this file except in compliance with the License.
5 # You may obtain a copy of the License at
6 #
7 #     http://www.apache.org/licenses/LICENSE-2.0
8 #
9 # Unless required by applicable law or agreed to in writing, software
10 # distributed under the License is distributed on an "AS IS" BASIS,
11 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12 # See the License for the specific language governing permissions and
13 # limitations under the License.
14 
15 import jax
16 import jax.numpy as jnp
17 
18 import pgx.v1 as v1
19 from pgx._src.struct import dataclass
20 
21 FALSE = jnp.bool_(False)
22 TRUE = jnp.bool_(True)
23 
24 INVALID_ACTION = jnp.int8(-1)
25 CALL = jnp.int8(0)
26 RAISE = jnp.int8(1)
27 FOLD = jnp.int8(2)
28 
29 MAX_RAISE = jnp.int8(2)
30 
31 
32 @dataclass
33 class State(v1.State):
34     current_player: jnp.ndarray = jnp.int8(0)
35     observation: jnp.ndarray = jnp.zeros((8, 8, 2), dtype=jnp.bool_)
36     rewards: jnp.ndarray = jnp.float32([0.0, 0.0])
37     terminated: jnp.ndarray = FALSE
38     truncated: jnp.ndarray = FALSE
39     legal_action_mask: jnp.ndarray = jnp.ones(3, dtype=jnp.bool_)
40     _rng_key: jax.random.KeyArray = jax.random.PRNGKey(0)
41     _step_count: jnp.ndarray = jnp.int32(0)
42     # --- Leduc Hold'Em specific ---
43     _first_player: jnp.ndarray = jnp.int8(0)
44     # [(player 0), (player 1), (public)]
45     _cards: jnp.ndarray = jnp.int8([-1, -1, -1])
46     # 0(Call)  1(Bet)  2(Fold)  3(Check)
47     _last_action: jnp.ndarray = INVALID_ACTION
48     _chips: jnp.ndarray = jnp.ones(2, dtype=jnp.int8)
49     _round: jnp.ndarray = jnp.int8(0)
50     _raise_count: jnp.ndarray = jnp.int8(0)
51 
52     @property
53     def env_id(self) -> v1.EnvId:
54         return "leduc_holdem"
55 
56 
57 class LeducHoldem(v1.Env):
58     def __init__(self):
59         super().__init__()
60 
61     def _init(self, key: jax.random.KeyArray) -> State:
62         return _init(key)
63 
64     def _step(self, state: v1.State, action: jnp.ndarray) -> State:
65         assert isinstance(state, State)
66         return _step(state, action)
67 
68     def _observe(self, state: v1.State, player_id: jnp.ndarray) -> jnp.ndarray:
69         assert isinstance(state, State)
70         return _observe(state, player_id)
71 
72     @property
73     def id(self) -> v1.EnvId:
74         return "leduc_holdem"
75 
76     @property
77     def version(self) -> str:
78         return "beta"
79 
80     @property
81     def num_players(self) -> int:
82         return 2
83 
84 
85 def _init(rng: jax.random.KeyArray) -> State:
86     rng1, rng2, rng3 = jax.random.split(rng, 3)
87     current_player = jnp.int8(jax.random.bernoulli(rng1))
88     init_card = jax.random.permutation(
89         rng2, jnp.int8([0, 0, 1, 1, 2, 2]), independent=True
90     )
91     return State(  # type:ignore
92         _rng_key=rng3,
93         _first_player=current_player,
94         current_player=current_player,
95         _cards=init_card[:3],
96         legal_action_mask=jnp.bool_([1, 1, 0]),
97         _chips=jnp.ones(2, dtype=jnp.int8),
98     )
99 
100 
101 def _step(state: State, action):
102     action = jnp.int8(action)
103     chips = jax.lax.switch(
104         action,
105         [
106             lambda: state._chips.at[state.current_player].set(
107                 state._chips[1 - state.current_player]
108             ),  # CALL
109             lambda: state._chips.at[state.current_player].set(
110                 jnp.max(state._chips) + _raise_chips(state)
111             ),  # RAISE
112             lambda: state._chips,  # FOLD
113         ],
114     )
115 
116     round_over, terminated, reward = _check_round_over(state, action)
117     last_action = jax.lax.select(round_over, INVALID_ACTION, action)
118     current_player = jax.lax.select(
119         round_over, state._first_player, 1 - state.current_player
120     )
121     raise_count = jax.lax.select(
122         round_over, jnp.int8(0), state._raise_count + jnp.int8(action == RAISE)
123     )
124 
125     reward *= jnp.min(chips)
126 
127     legal_action = jax.lax.switch(
128         action,
129         [
130             lambda: jnp.bool_([1, 1, 0]),  # CALL
131             lambda: jnp.bool_([1, 1, 1]),  # RAISE
132             lambda: jnp.bool_([0, 0, 0]),  # FOLD
133         ],
134     )
135     legal_action = legal_action.at[RAISE].set(raise_count < MAX_RAISE)
136 
137     return state.replace(  # type:ignore
138         current_player=current_player,
139         _last_action=last_action,
140         legal_action_mask=legal_action,
141         terminated=terminated,
142         rewards=reward,
143         _round=state._round + jnp.int8(round_over),
144         _chips=chips,
145         _raise_count=raise_count,
146     )
147 
148 
149 def _check_round_over(state, action):
150     round_over = (action == FOLD) | (
151         (state._last_action != INVALID_ACTION) & (action == CALL)
152     )
153     terminated = round_over & (state._round == 1)
154 
155     reward = jax.lax.select(
156         terminated & (action == FOLD),
157         jnp.float32([-1, -1]).at[1 - state.current_player].set(1),
158         jnp.float32([0, 0]),
159     )
160     reward = jax.lax.select(
161         terminated & (action != FOLD),
162         _get_unit_reward(state),
163         reward,
164     )
165     return round_over, terminated, reward
166 
167 
168 def _get_unit_reward(state: State):
169     win_by_one_pair = state._cards[state.current_player] == state._cards[2]
170     lose_by_one_pair = (
171         state._cards[1 - state.current_player] == state._cards[2]
172     )
173     win = win_by_one_pair | (
174         ~lose_by_one_pair
175         & (
176             state._cards[state.current_player]
177             > state._cards[1 - state.current_player]
178         )
179     )
180     reward = jax.lax.select(
181         win,
182         jnp.float32([-1, -1]).at[state.current_player].set(1),
183         jnp.float32([-1, -1]).at[1 - state.current_player].set(1),
184     )
185     return jax.lax.select(
186         state._cards[state.current_player]
187         == state._cards[1 - state.current_player],  # Draw
188         jnp.float32([0, 0]),
189         reward,
190     )
191 
192 
193 def _raise_chips(state):
194     """raise amounts is 2 in the first round and 4 in the second round."""
195     return (state._round + 1) * 2
196 
197 
198 def _observe(state: State, player_id) -> jnp.ndarray:
199     """
200     Index   Meaning
201     0~2     J ~ K in hand
202     3~5     J ~ K as public card
203     6~19    0 ~ 13 chips for the current player
204     20~33   0 ~ 13 chips for the opponent
205     """
206     obs = jnp.zeros(34, dtype=jnp.bool_)
207     obs = obs.at[state._cards[player_id]].set(TRUE)
208     obs = jax.lax.select(
209         state._round == 1, obs.at[3 + state._cards[2]].set(TRUE), obs
210     )
211     obs = obs.at[6 + state._chips[player_id]].set(TRUE)
212     obs = obs.at[20 + state._chips[1 - player_id]].set(TRUE)
213 
214     return obs
215
```
--- END FILES ---

Please first localize the bug based on the issue statement, and then generate a patch according to the `git diff` format fenced by three backticks.

Here is an example:

```diff
diff --git a/examples/server_async.py b/examples/server_async.py
--- a/examples/server_async.py
+++ b/examples/server_async.py
@@ -313,4 +313,4 @@
 
 
 if __name__ == "__main__":
-    asyncio.run(run_async_server("."), debug=True)
+    asyncio.run(run_async_server(), debug=True)
diff --git a/examples/server_sync.py b/examples/server_sync.py
--- a/examples/server_sync.py
+++ b/examples/server_sync.py
@@ -313,5 +313,5 @@
 
 
 if __name__ == "__main__":
-    server = run_sync_server(".")
+    server = run_sync_server()
     server.shutdown()

```

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 _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,
     )

{"golden_diff": "diff --git a/pgx/leduc_holdem.py b/pgx/leduc_holdem.py\n--- a/pgx/leduc_holdem.py\n+++ b/pgx/leduc_holdem.py\n@@ -147,18 +147,20 @@\n \n \n def _check_round_over(state, action):\n-    round_over = (action == FOLD) | (\n-        (state._last_action != INVALID_ACTION) & (action == CALL)\n-    )\n-    terminated = round_over & (state._round == 1)\n+    fold = action == FOLD\n+    call = (state._last_action != INVALID_ACTION) & (action == CALL)\n+    _continue = (state._round == 0) & call\n+\n+    round_over = fold | call\n+    terminated = round_over & (~_continue)\n \n     reward = jax.lax.select(\n-        terminated & (action == FOLD),\n+        fold,\n         jnp.float32([-1, -1]).at[1 - state.current_player].set(1),\n         jnp.float32([0, 0]),\n     )\n     reward = jax.lax.select(\n-        terminated & (action != FOLD),\n+        terminated & call,\n         _get_unit_reward(state),\n         reward,\n     )\n", "issue": "Termination in Leduc Holdem\nHello, in Leduc Holdem I don't quite understand why doing the `FOLD` action in the first round of betting doesn't terminate the episode.\r\n`    terminated = round_over & (state._round == 1)`\r\nSeems to me that's a faulty behavior, please tell me if I miss something. A possible fix could be:\r\n`    terminated = (round_over & (state._round == 1)) | (action==FOLD) `\r\n\r\n\r\n\r\n\n", "before_files": [{"content": "# Copyright 2023 The Pgx Authors. All Rights Reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nimport jax\nimport jax.numpy as jnp\n\nimport pgx.v1 as v1\nfrom pgx._src.struct import dataclass\n\nFALSE = jnp.bool_(False)\nTRUE = jnp.bool_(True)\n\nINVALID_ACTION = jnp.int8(-1)\nCALL = jnp.int8(0)\nRAISE = jnp.int8(1)\nFOLD = jnp.int8(2)\n\nMAX_RAISE = jnp.int8(2)\n\n\n@dataclass\nclass State(v1.State):\n    current_player: jnp.ndarray = jnp.int8(0)\n    observation: jnp.ndarray = jnp.zeros((8, 8, 2), dtype=jnp.bool_)\n    rewards: jnp.ndarray = jnp.float32([0.0, 0.0])\n    terminated: jnp.ndarray = FALSE\n    truncated: jnp.ndarray = FALSE\n    legal_action_mask: jnp.ndarray = jnp.ones(3, dtype=jnp.bool_)\n    _rng_key: jax.random.KeyArray = jax.random.PRNGKey(0)\n    _step_count: jnp.ndarray = jnp.int32(0)\n    # --- Leduc Hold'Em specific ---\n    _first_player: jnp.ndarray = jnp.int8(0)\n    # [(player 0), (player 1), (public)]\n    _cards: jnp.ndarray = jnp.int8([-1, -1, -1])\n    # 0(Call)  1(Bet)  2(Fold)  3(Check)\n    _last_action: jnp.ndarray = INVALID_ACTION\n    _chips: jnp.ndarray = jnp.ones(2, dtype=jnp.int8)\n    _round: jnp.ndarray = jnp.int8(0)\n    _raise_count: jnp.ndarray = jnp.int8(0)\n\n    @property\n    def env_id(self) -> v1.EnvId:\n        return \"leduc_holdem\"\n\n\nclass LeducHoldem(v1.Env):\n    def __init__(self):\n        super().__init__()\n\n    def _init(self, key: jax.random.KeyArray) -> State:\n        return _init(key)\n\n    def _step(self, state: v1.State, action: jnp.ndarray) -> State:\n        assert isinstance(state, State)\n        return _step(state, action)\n\n    def _observe(self, state: v1.State, player_id: jnp.ndarray) -> jnp.ndarray:\n        assert isinstance(state, State)\n        return _observe(state, player_id)\n\n    @property\n    def id(self) -> v1.EnvId:\n        return \"leduc_holdem\"\n\n    @property\n    def version(self) -> str:\n        return \"beta\"\n\n    @property\n    def num_players(self) -> int:\n        return 2\n\n\ndef _init(rng: jax.random.KeyArray) -> State:\n    rng1, rng2, rng3 = jax.random.split(rng, 3)\n    current_player = jnp.int8(jax.random.bernoulli(rng1))\n    init_card = jax.random.permutation(\n        rng2, jnp.int8([0, 0, 1, 1, 2, 2]), independent=True\n    )\n    return State(  # type:ignore\n        _rng_key=rng3,\n        _first_player=current_player,\n        current_player=current_player,\n        _cards=init_card[:3],\n        legal_action_mask=jnp.bool_([1, 1, 0]),\n        _chips=jnp.ones(2, dtype=jnp.int8),\n    )\n\n\ndef _step(state: State, action):\n    action = jnp.int8(action)\n    chips = jax.lax.switch(\n        action,\n        [\n            lambda: state._chips.at[state.current_player].set(\n                state._chips[1 - state.current_player]\n            ),  # CALL\n            lambda: state._chips.at[state.current_player].set(\n                jnp.max(state._chips) + _raise_chips(state)\n            ),  # RAISE\n            lambda: state._chips,  # FOLD\n        ],\n    )\n\n    round_over, terminated, reward = _check_round_over(state, action)\n    last_action = jax.lax.select(round_over, INVALID_ACTION, action)\n    current_player = jax.lax.select(\n        round_over, state._first_player, 1 - state.current_player\n    )\n    raise_count = jax.lax.select(\n        round_over, jnp.int8(0), state._raise_count + jnp.int8(action == RAISE)\n    )\n\n    reward *= jnp.min(chips)\n\n    legal_action = jax.lax.switch(\n        action,\n        [\n            lambda: jnp.bool_([1, 1, 0]),  # CALL\n            lambda: jnp.bool_([1, 1, 1]),  # RAISE\n            lambda: jnp.bool_([0, 0, 0]),  # FOLD\n        ],\n    )\n    legal_action = legal_action.at[RAISE].set(raise_count < MAX_RAISE)\n\n    return state.replace(  # type:ignore\n        current_player=current_player,\n        _last_action=last_action,\n        legal_action_mask=legal_action,\n        terminated=terminated,\n        rewards=reward,\n        _round=state._round + jnp.int8(round_over),\n        _chips=chips,\n        _raise_count=raise_count,\n    )\n\n\ndef _check_round_over(state, action):\n    round_over = (action == FOLD) | (\n        (state._last_action != INVALID_ACTION) & (action == CALL)\n    )\n    terminated = round_over & (state._round == 1)\n\n    reward = jax.lax.select(\n        terminated & (action == FOLD),\n        jnp.float32([-1, -1]).at[1 - state.current_player].set(1),\n        jnp.float32([0, 0]),\n    )\n    reward = jax.lax.select(\n        terminated & (action != FOLD),\n        _get_unit_reward(state),\n        reward,\n    )\n    return round_over, terminated, reward\n\n\ndef _get_unit_reward(state: State):\n    win_by_one_pair = state._cards[state.current_player] == state._cards[2]\n    lose_by_one_pair = (\n        state._cards[1 - state.current_player] == state._cards[2]\n    )\n    win = win_by_one_pair | (\n        ~lose_by_one_pair\n        & (\n            state._cards[state.current_player]\n            > state._cards[1 - state.current_player]\n        )\n    )\n    reward = jax.lax.select(\n        win,\n        jnp.float32([-1, -1]).at[state.current_player].set(1),\n        jnp.float32([-1, -1]).at[1 - state.current_player].set(1),\n    )\n    return jax.lax.select(\n        state._cards[state.current_player]\n        == state._cards[1 - state.current_player],  # Draw\n        jnp.float32([0, 0]),\n        reward,\n    )\n\n\ndef _raise_chips(state):\n    \"\"\"raise amounts is 2 in the first round and 4 in the second round.\"\"\"\n    return (state._round + 1) * 2\n\n\ndef _observe(state: State, player_id) -> jnp.ndarray:\n    \"\"\"\n    Index   Meaning\n    0~2     J ~ K in hand\n    3~5     J ~ K as public card\n    6~19    0 ~ 13 chips for the current player\n    20~33   0 ~ 13 chips for the opponent\n    \"\"\"\n    obs = jnp.zeros(34, dtype=jnp.bool_)\n    obs = obs.at[state._cards[player_id]].set(TRUE)\n    obs = jax.lax.select(\n        state._round == 1, obs.at[3 + state._cards[2]].set(TRUE), obs\n    )\n    obs = obs.at[6 + state._chips[player_id]].set(TRUE)\n    obs = obs.at[20 + state._chips[1 - player_id]].set(TRUE)\n\n    return obs\n", "path": "pgx/leduc_holdem.py"}], "after_files": [{"content": "# Copyright 2023 The Pgx Authors. All Rights Reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nimport jax\nimport jax.numpy as jnp\n\nimport pgx.v1 as v1\nfrom pgx._src.struct import dataclass\n\nFALSE = jnp.bool_(False)\nTRUE = jnp.bool_(True)\n\nINVALID_ACTION = jnp.int8(-1)\nCALL = jnp.int8(0)\nRAISE = jnp.int8(1)\nFOLD = jnp.int8(2)\n\nMAX_RAISE = jnp.int8(2)\n\n\n@dataclass\nclass State(v1.State):\n    current_player: jnp.ndarray = jnp.int8(0)\n    observation: jnp.ndarray = jnp.zeros((8, 8, 2), dtype=jnp.bool_)\n    rewards: jnp.ndarray = jnp.float32([0.0, 0.0])\n    terminated: jnp.ndarray = FALSE\n    truncated: jnp.ndarray = FALSE\n    legal_action_mask: jnp.ndarray = jnp.ones(3, dtype=jnp.bool_)\n    _rng_key: jax.random.KeyArray = jax.random.PRNGKey(0)\n    _step_count: jnp.ndarray = jnp.int32(0)\n    # --- Leduc Hold'Em specific ---\n    _first_player: jnp.ndarray = jnp.int8(0)\n    # [(player 0), (player 1), (public)]\n    _cards: jnp.ndarray = jnp.int8([-1, -1, -1])\n    # 0(Call)  1(Bet)  2(Fold)  3(Check)\n    _last_action: jnp.ndarray = INVALID_ACTION\n    _chips: jnp.ndarray = jnp.ones(2, dtype=jnp.int8)\n    _round: jnp.ndarray = jnp.int8(0)\n    _raise_count: jnp.ndarray = jnp.int8(0)\n\n    @property\n    def env_id(self) -> v1.EnvId:\n        return \"leduc_holdem\"\n\n\nclass LeducHoldem(v1.Env):\n    def __init__(self):\n        super().__init__()\n\n    def _init(self, key: jax.random.KeyArray) -> State:\n        return _init(key)\n\n    def _step(self, state: v1.State, action: jnp.ndarray) -> State:\n        assert isinstance(state, State)\n        return _step(state, action)\n\n    def _observe(self, state: v1.State, player_id: jnp.ndarray) -> jnp.ndarray:\n        assert isinstance(state, State)\n        return _observe(state, player_id)\n\n    @property\n    def id(self) -> v1.EnvId:\n        return \"leduc_holdem\"\n\n    @property\n    def version(self) -> str:\n        return \"beta\"\n\n    @property\n    def num_players(self) -> int:\n        return 2\n\n\ndef _init(rng: jax.random.KeyArray) -> State:\n    rng1, rng2, rng3 = jax.random.split(rng, 3)\n    current_player = jnp.int8(jax.random.bernoulli(rng1))\n    init_card = jax.random.permutation(\n        rng2, jnp.int8([0, 0, 1, 1, 2, 2]), independent=True\n    )\n    return State(  # type:ignore\n        _rng_key=rng3,\n        _first_player=current_player,\n        current_player=current_player,\n        _cards=init_card[:3],\n        legal_action_mask=jnp.bool_([1, 1, 0]),\n        _chips=jnp.ones(2, dtype=jnp.int8),\n    )\n\n\ndef _step(state: State, action):\n    action = jnp.int8(action)\n    chips = jax.lax.switch(\n        action,\n        [\n            lambda: state._chips.at[state.current_player].set(\n                state._chips[1 - state.current_player]\n            ),  # CALL\n            lambda: state._chips.at[state.current_player].set(\n                jnp.max(state._chips) + _raise_chips(state)\n            ),  # RAISE\n            lambda: state._chips,  # FOLD\n        ],\n    )\n\n    round_over, terminated, reward = _check_round_over(state, action)\n    last_action = jax.lax.select(round_over, INVALID_ACTION, action)\n    current_player = jax.lax.select(\n        round_over, state._first_player, 1 - state.current_player\n    )\n    raise_count = jax.lax.select(\n        round_over, jnp.int8(0), state._raise_count + jnp.int8(action == RAISE)\n    )\n\n    reward *= jnp.min(chips)\n\n    legal_action = jax.lax.switch(\n        action,\n        [\n            lambda: jnp.bool_([1, 1, 0]),  # CALL\n            lambda: jnp.bool_([1, 1, 1]),  # RAISE\n            lambda: jnp.bool_([0, 0, 0]),  # FOLD\n        ],\n    )\n    legal_action = legal_action.at[RAISE].set(raise_count < MAX_RAISE)\n\n    return state.replace(  # type:ignore\n        current_player=current_player,\n        _last_action=last_action,\n        legal_action_mask=legal_action,\n        terminated=terminated,\n        rewards=reward,\n        _round=state._round + jnp.int8(round_over),\n        _chips=chips,\n        _raise_count=raise_count,\n    )\n\n\ndef _check_round_over(state, action):\n    fold = action == FOLD\n    call = (state._last_action != INVALID_ACTION) & (action == CALL)\n    _continue = (state._round == 0) & call\n\n    round_over = fold | call\n    terminated = round_over & (~_continue)\n\n    reward = jax.lax.select(\n        fold,\n        jnp.float32([-1, -1]).at[1 - state.current_player].set(1),\n        jnp.float32([0, 0]),\n    )\n    reward = jax.lax.select(\n        terminated & call,\n        _get_unit_reward(state),\n        reward,\n    )\n    return round_over, terminated, reward\n\n\ndef _get_unit_reward(state: State):\n    win_by_one_pair = state._cards[state.current_player] == state._cards[2]\n    lose_by_one_pair = (\n        state._cards[1 - state.current_player] == state._cards[2]\n    )\n    win = win_by_one_pair | (\n        ~lose_by_one_pair\n        & (\n            state._cards[state.current_player]\n            > state._cards[1 - state.current_player]\n        )\n    )\n    reward = jax.lax.select(\n        win,\n        jnp.float32([-1, -1]).at[state.current_player].set(1),\n        jnp.float32([-1, -1]).at[1 - state.current_player].set(1),\n    )\n    return jax.lax.select(\n        state._cards[state.current_player]\n        == state._cards[1 - state.current_player],  # Draw\n        jnp.float32([0, 0]),\n        reward,\n    )\n\n\ndef _raise_chips(state):\n    \"\"\"raise amounts is 2 in the first round and 4 in the second round.\"\"\"\n    return (state._round + 1) * 2\n\n\ndef _observe(state: State, player_id) -> jnp.ndarray:\n    \"\"\"\n    Index   Meaning\n    0~2     J ~ K in hand\n    3~5     J ~ K as public card\n    6~19    0 ~ 13 chips for the current player\n    20~33   0 ~ 13 chips for the opponent\n    \"\"\"\n    obs = jnp.zeros(34, dtype=jnp.bool_)\n    obs = obs.at[state._cards[player_id]].set(TRUE)\n    obs = jax.lax.select(\n        state._round == 1, obs.at[3 + state._cards[2]].set(TRUE), obs\n    )\n    obs = obs.at[6 + state._chips[player_id]].set(TRUE)\n    obs = obs.at[20 + state._chips[1 - player_id]].set(TRUE)\n\n    return obs\n", "path": "pgx/leduc_holdem.py"}]}