Locomotion Example fails on training handstand to due out of date Jax version in Google Colab

## Issue

When training the quadruped to do the hand stand in the [locomotion.ipnyb ](https://colab.research.google.com/github/google-deepmind/mujoco_playground/blob/main/learning/notebooks/locomotion.ipynb), it fails to run due a missing method on the jax.monitoring object.

## Environment/Libraries
Environment: Google Colab
Python Version: 3.12.11
Torch Version: 2.8.0+cu126
Is Cuda Available: True
Cuda Version: 12.6
GPU Device: NVIDIA L4
Gymnasium Version: 1.2.0
Numpy Version: 2.0.2
Mujoco Version: 3.3.6
Mujoco MJX Version: 3.3.6
Mujoco Playground Version: 0.0.5
JAX Version: 0.5.3
JAXlib Version: 0.5.3
Brax Version: 0.13.0

## Error

[Google Colab Cell Link](https://colab.research.google.com/github/google-deepmind/mujoco_playground/blob/main/learning/notebooks/locomotion.ipynb#scrollTo=MY6P3abhSNGU&line=1&uniqifier=1)
```python
make_inference_fn, params, metrics = train_fn(
environment=registry.load(env_name, config=env_cfg),
eval_env=registry.load(env_name, config=env_cfg),
wrap_env_fn=wrapper.wrap_for_brax_training,
)
```

## Solution
The two solutions I found:
- Removing `policy_params_fn` from the `train_fn` allows the rest of the notebook to execute without issue.
- Update the JAX version to something greater than version 0.6.0. Google Colab has JAX 0.5.3 installed by default which doesn't support the `record_scaler`. (Maybe update the pyproject.toml to set a minimum version for JAX?)

## Additional Details
Here is the original [Pull Request](https://github.com/jax-ml/jax/pull/26853) to add the `record_scaler` to JAX

I would be interested in submitting a PR to fix! Just need some direction! Let me know if you have any questions or concerns.