I've tried restarting the runtime and pip installing brax again using this command, "!pip install git+https://github.com/google/brax.git@main", but regardless of how many times I do it, "from brax import jumpy as jp" results in "ImportError: cannot import name 'jumpy' from 'brax' (/usr/local/lib/python3.10/dist-packages/brax/init.py)". This error has only been taking place on Google Colab as far as I can tell.
I don't know how to fix it, this is how it is written in the brax tutorial on their website: https://ymd_h.gitlab.io/ymd_blog/doc/brax/env/
If anyone can help with this I would very much appreciate it.