I wanted to reproduce your results for the vehicle benchmark, however, while trying to install your requirements.txt file, I have run into a dependency issue.

At the first installation, (pip install -r requirements.txt) I get:

To solve these issues, I run a force reinstall for the packages with the mentioned versions in the Error.
pip install --force-reinstall -v "numpy==1.21" pip install --force-reinstall -v "jax==0.4.6" pip install --force-reinstall -v "typing-extensions==4.2.0"

Now, when I try to install jaxlib 0.4.7, I get the following error:

Clearly, optax 0.1.0 which is a dependency for jaxlib 0.4.7 requires typing-extensions 3.10.0 but chex 0.1.7 which a dependency for another package, requires typing-extensions>=4.2.0

Can you help me with this issue? I would suggest that you update the repo with a pip freeze output in the requirements.txt file so that the dependencies are taken care of automatically. That way, the user can simply do a pip install --no-dependencies -r requirements.txt to install everything flawlessly.