tidy

sotetsuk · Dec 2, 2024 · d644813 · d644813
1 parent 8266f47
commit d644813
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/pgx/_src/games/hex.py b/pgx/_src/games/hex.py
@@ -55,7 +55,7 @@ def step(self, state: GameState, action: Array) -> GameState:
             lambda: partial(_step, size=self.size)(state, action),
             lambda: partial(_swap, size=self.size)(state),
         )
-        terminated = self.is_terminal(x)
+        terminated = _is_terminal(x, self.size)
         return x._replace(terminated=terminated)
 
     def observe(self, state: GameState, color: Optional[Array] = None) -> Array:
@@ -65,7 +65,7 @@ def legal_action_mask(self, state: GameState) -> Array:
         return jnp.append(state.board == 0, state.step_count == 1)
 
     def is_terminal(self, state: GameState) -> Array:
-        return _is_terminal(state, self.size)
+        return state.terminated
 
     # def rewards(self, state: GameState) -> Array:
     #     ...