DQN refactor

[alexunderch]

• Refactored DQN versions with jax and torch that pass the tests
• Simpler tree-based replay buffers
• Changed the example name to dqn_ to improve visibility

Still TBD:

• jax implementation is still very slow with plain python, see Significant performance difference of NNX relative to equinox google/flax#4045, will fix with different improvements, like caching
• the only tests are on breakthrough, @lanctot are there any other representative games: tic-tac-toe, hanabi?..

[alexunderch]

I decided to use different flax and jax versions to get access to the latest nnx api.

jax==0.8.1
flax==0.12.1
optax==0.2.6
orbax-checkpoint==0.11.31

[lanctot]

I decided to use different flax and jax versions to get access to the latest nnx api.

jax==0.8.1
flax==0.12.1
optax==0.2.6
orbax-checkpoint==0.11.31

Are each of them still supported on Python >= 3.10?

[alexunderch]

Docs say that for >=3.10 current top versions are fully compatible...

[alexunderch]

Yes, it now fails because of the versions.

@lanctot, what's better for you, to upgrade versions in a separate PR or with this/AlphaZero PR?

[lanctot]

Yes, it now fails because of the versions.

@lanctot, what's better for you, to upgrade versions in a separate PR or with this/AlphaZero PR?

I'm still not sure if updating versions will work. Sometimes the cascade of dependencies leads to a problem and it makes it impossible if the Python versions are too old. But let me try in a separate PR.

[lanctot]

jax==0.8.1

This one requires >= 3.11, from https://pypi.org/project/jax/0.8.1/

flax==0.12.1

This one too (>= 3.11), from https://pypi.org/project/flax/0.12.1/

optax==0.2.6

This one is ok (>= 3.10), https://pypi.org/project/optax/0.2.6/

orbax-checkpoint==0.11.31

This one is ok too (>= 3.10), from: https://pypi.org/project/orbax-checkpoint/0.11.31/

[lanctot]

Ok...... so what do we do?

Well, 3.10 is causing other problems too ( see #1424 ).

I just checked. Seems like Colab is now using 3.12. Normally I wait until EOL before removing support for a version, but that's quite far away (October '26). In this case it's causing multiple issues so I'd be happy to remove it early.

Give me a few weeks. I have to check with a few people (most notably the Kaggle Game Arena who are relying on a stable OpenSpiel for their environments). And I could have sworn that the version of Colab I used for the LLM imitation learning was 3.10, so maybe it's lower for the TPU kernels or maybe I was just mistaken.

[lanctot]

@alexunderch

The order I suggest is the following:

• Wait for me to confirm that removing 3.10 support is ok
• I will then remove 3.10 support and update github, will let you know when that's done
• Then, let's update the AlphaZero PR, add it to the tests, and merge it since it's been around for a really long time and looking good
• Then, we come back to this one and the other ones that use the newer nnx (like deep CFR etc.)

Sound good to you?

[lanctot]

Moving the discussion of removing Python 3.10 to a new issue: #1425