use brax.v1 and update requirements

DBraun commented 11 months ago

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

[x] a clear description of the PR has been added
[x] sufficient tests have been written
[x] relevant section added to the documentation
[x] example notebook added to the repo
[x] clean docstrings and comments have been written
[x] 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 commented 11 months ago

Is any QDax maintainer able to help move this forward?

Lookatator commented 11 months ago

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 😄

codecov-commenter commented 11 months ago

Codecov Report

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

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

:exclamation: 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%> (ø)`

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

adaptive-intelligent-robotics / QDax