Early termination caused by castling

[HongruiTang]

Hi Sotetsu,

I found an issue with both kingside and queenside castling in chess when I was playing with some games on Lichess. The problem is the games are always early terminated due to illegal actions after castling. I dived into the code and found the problem is probably in the _apply_move function (line 318).

The original implementation only changed state._board but not state._possible_piece_positions after castling. Although, each move updates state._possible_piece_positions here:

# actually move
    state = state.replace(  # type: ignore
        _board=state._board.at[a.from_].set(EMPTY).at[a.to].set(piece)
    )

due to the fact that castling is a move that involves two submoves, the position of the Rook should also be updated. Otherwise, this part

actions = legal_norml_moves(
        state._possible_piece_positions[0]
    ).flatten()  # include -1

would be affected when checking legal moves.

I also have a question about the variable state._possible_piece_positions. When a piece is captured, the position of the piece will remain the same or should be set as a new value, like -1.

Also, a minor typo is corrected at line 65 in chess.py

Thanks

[sotetsuk]

Hi, Thank you for sending a PR! 😄 I really appreciate it 🙏

Could you kindly add a test case in which the current implementation fails but the PR works? 🙏 🥺

We collect buggy test cases here:
https://github.com/sotetsuk/pgx/blob/main/tests/test_chess.py#L659

[HongruiTang]

Thanks! I've added a test in https://github.com/sotetsuk/pgx/blob/main/tests/test_chess.py#L659

[sotetsuk]

I also have a question about the variable state._possible_piece_positions. When a piece is captured, the position of the piece will remain the same or should be set as a new value, like -1.

This is because _possible_piece_positions is used only to represent the index candidates where pieces may exists.
It can speed up the computation in

    actions = legal_norml_moves(
        state._possible_piece_positions[0]
    ).flatten()  # include -1

but note that as _possible_piece_positions represents candidates, this should also work

    actions = legal_norml_moves(
        jnp.arange(64)
    ).flatten()  # include -1

So even if actually there are no pieces in some _possible_piece_positions, it's ok.

[sotetsuk]

#983 is merged. I'll close this issue after a new pgx version is released.
Thank you for your contribution! 🙏 @HongruiTang

[HongruiTang]

I also have a question about the variable state._possible_piece_positions. When a piece is captured, the position of the piece will remain the same or should be set as a new value, like -1.

This is because _possible_piece_positions is used only to represent the index candidates where pieces may exists. It can speed up the computation in

    actions = legal_norml_moves(
        state._possible_piece_positions[0]
    ).flatten()  # include -1

but note that as _possible_piece_positions represents candidates, this should also work

    actions = legal_norml_moves(
        jnp.arange(64)
    ).flatten()  # include -1

So even if actually there are no pieces in some _possible_piece_positions, it's ok.

Thanks for the clarification. That makes sense!