Reduce dependencies

[maxencefaldor]

Great work on craftax! I have a few suggestions regarding dependencies, aiming to keep the project lean, maintainable, and robust against future compatibility issues.

1. gymnax is not maintained anymore and is incompatible with jax>=0.6.0. I understand craftax uses a similar API but I don't think this justifies a dependency.
2. optax and distrax are listed as dependencies but aren't used. Could you remove those to lighten the project?
3. flax is a significant dependency but only flax.struct is used. Have you considered alternatives to flax.struct?
4. chex is used only 7 times throughout the codebase for chex.PRNGKey and chex.Array. The recommended way from jax official documentation is to directly use jax.Array for both PRNG keys and Arrays. I suggest we remove the chex dependency in favor of jax.Array.

Minimizing dependencies generally enhances stability and simplifies maintenance, helping to avoid situations like the craftax incompatibility with jax>=0.6.0.

[MichaelTMatthews]

hey maxence thanks for raising these

1. I am currently working on a semi-major refactor of craftax as part of another project. I agree with removing gymnax dependency and will include that in the release
2. thanks for the PR - have merged it
3. agree - thinking of just using a regular dataclass and registering with register_dataclass, just need to try it out
4. I think I'll fold this into work on (3), but yes generally agree, chex dependency is mostly a carry over from gymnax