Migrate brax v2 in notebook brax_training

Denys88 / rl_games

RL implementations

MIT License

913 stars 150 forks source link

Migrate brax v2 in notebook brax_training #302

Open gonzalesMK opened 3 months ago

gonzalesMK commented 3 months ago

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 commented 3 months ago

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 commented 2 months ago

Sure! I can do that.