use brax.v1 and update requirements

[DBraun]

Related issues: #155

Where necessary I've changed brax to brax.v1. This is a short-term solution because v1 is already deprecated. However, it at least gets QDax working with jax/jaxlib 0.4.16 and flax 0.7.4. At the moment, if you install QDax from scratch without previously installing jax/flax, you end up with an unusable flax/qdax. This fixes that.

This PR is really just a starting point for discussion. I have installed it locally and imported various QDax things which work fine.

Checks

• a clear description of the PR has been added
• sufficient tests have been written
• relevant section added to the documentation
• example notebook added to the repo
• clean docstrings and comments have been written
• if any issue/observation has been discovered, a new issue has been opened

Future improvements

In the future QDax should not use brax.v1. Someone could make this a new issue soon.

[DBraun]

Is any QDax maintainer able to help move this forward?

[Lookatator]

Is any QDax maintainer able to help move this forward?

Hi @DBraun,
Thanks very much for this PR, I'll review it either today or tomorrow 😄

[Lookatator]

Seems good to me overall!

To run the PR tests and make it possible to compile the docs and run the automated pipeline, could you:

1. merge develop into your branch?
2. and change the version of python in the .readthedocs.yaml to 3.9 ?

[codecov-commenter]

Codecov Report

Merging #156 (98c5999) into develop (c111cea) will not change coverage.
The diff coverage is 93.61%.

❗ Current head 98c5999 differs from pull request most recent head 7dd286c. Consider uploading reports for the commit 7dd286c to get more accurate results

❗ Your organization needs to install the Codecov GitHub app to enable full functionality.

@@           Coverage Diff            @@
##           develop     #156   +/-   ##
========================================
  Coverage    92.23%   92.23%           
========================================
  Files          120      120           
  Lines         6992     6992           
========================================
  Hits          6449     6449           
  Misses         543      543           

Files	Coverage Δ
qdax/__init__.py	100.00% <100.00%> (ø)
qdax/environments/base_wrappers.py	82.92% <100.00%> (ø)
qdax/environments/exploration_wrappers.py	32.97% <100.00%> (ø)
qdax/environments/init_state_wrapper.py	90.00% <100.00%> (ø)
qdax/environments/locomotion_wrappers.py	85.49% <100.00%> (ø)
qdax/environments/pointmaze.py	97.22% <100.00%> (ø)
qdax/environments/wrappers.py	96.96% <100.00%> (ø)
qdax/environments/humanoidtrap.py	18.30% <87.50%> (ø)
qdax/environments/__init__.py	86.79% <80.00%> (ø)

📣 We’re building smart automated test selection to slash your CI/CD build times. Learn more