Migrate brax v2 in notebook brax_training

[gonzalesMK]

Hi!

I just tried to use this notebook: https://github.com/Denys88/rl_games/blob/master/notebooks/brax_training.ipynb and an import error occurred on this:

from brax import jumpy as jp

It seems brax has a new api (v2), and I fixed changing to this:

from brax.v1 import jumpy as jp

[ViktorM]

Thanks @gonzalesMK for the finding! Could you please create a PR with the fix? Also hopefully their API is much more stable now I can take a look again into updating to the latest v2.

[gonzalesMK]

Sure! I can do that.

[Denys88]

We are working on our own update with Brax support. Thanks for feedback.