incorrect discount factor when the same player makes consecutive moves?

[ktegan]

I believe the discount factor is wrong if the same player makes consecutive moves (such as in backgammon). My understanding is that we set the discount to -1 whenever we "flip" the board. However, for games like backgammon a step() may not change the current player. Assuming I'm correct you can fix this with the following changes:

In selfplay() in train.py:

        previous_player = state.current_player
        state = jax.vmap(auto_reset(env.step, env.init))(state, policy_output.action, keys)
        discount = jnp.where(previous_player == state.current_player, 1.0, -1.0) * jnp.ones_like(value)

and in recurrent_fn()

    previous_player = state.current_player
    state = jax.vmap(env.step)(state, action)
    current_player = state.current_player
...
    discount = jnp.where(previous_player == current_player, 1.0, -1.0) * jnp.ones_like(value)

These are in my branch along with some other changes I'm still testing: https://github.com/ktegan/pgx/tree/kegan-backgammon