Updates for compatibility in RL tutorial with JAX v0.6.0

[mKabouri]

While running the example notebook examples/02_rl.ipynb, I encountered an error due to the removal of jax.tree_map in JAX v0.6.0. I replaced calls to jax.tree_map with jax.tree.map

The cell where I had the error:

from evosax.problems import BraxProblem as Problem
from evosax.problems.networks import MLP, tanh_output_fn

policy = MLP(
    layer_sizes=(32, 32, 32, 32, 8),
    output_fn=tanh_output_fn,
)

problem = Problem(
    env_name="ant",
    policy=policy,
    episode_length=1000,
    num_rollouts=16,
    use_normalize_obs=True,
)

key, subkey = jax.random.split(key)
problem_state = problem.init(key)

key, subkey = jax.random.split(key)
solution = problem.sample(subkey)

The error:

AttributeError: jax.tree_map was removed in JAX v0.6.0: use jax.tree.map (jax v0.4.25 or newer) or jax.tree_util.tree_map (any JAX version).