Use brax.v1 instead of brax to make it work

jiawei-ren / diffmimic

[ICLR 2023] DiffMimic: Efficient Motion Mimicking with Differentiable Physics https://arxiv.org/abs/2304.03274

https://diffmimic.github.io/

Other

278 stars 21 forks source link

Use brax.v1 instead of brax to make it work #7

Closed alonfnt closed 2 months ago

alonfnt commented 2 months ago

Current code does not work for newer versions of JAX (and its ecosystem). Luckly, the major culprit of the incompatibility here (brax) has a v1 module that makes it backwards compatible.

This is not the final solution (as a rewrite would probably be better), but at least now people can run this on the current versions (and enjoy all the benefits that new jax has)

Great work BTW.

closes #6, closes #5

jiawei-ren commented 2 months ago

That's awsome, thanks!