fix the Jax FutureWarnings

[Aneoshun]

With the latest version of Jax (obtained automatically with pip install qdax) we get a lot of FutureWarnings:

FutureWarning: jax.tree_map is deprecated, and will be removed in a future release. Use jax.tree_util.tree_map instead.
FutureWarning: jax.tree_leaves is deprecated, and will be removed in a future release. Use jax.tree_util.tree_leaves instead.
FutureWarning: jax.tree_structure is deprecated, and will be removed in a future release. Use jax.tree_util.tree_structure instead.
FutureWarning: jax.tree_unflatten is deprecated, and will be removed in a future release. Use jax.tree_util.tree_unflatten instead.

It would be good to fix this asap and to release it soon, maybe we some of the other bug fixes we have in the pipeline.