diff --git a/tests/environments_test/wrapper_test.py b/tests/environments_test/wrapper_test.py index a2bcb44e..f5e035ea 100644 --- a/tests/environments_test/wrapper_test.py +++ b/tests/environments_test/wrapper_test.py @@ -4,9 +4,9 @@ import jax import jax.numpy as jnp import pytest -from brax.physics.base import vec_to_arr -from brax.physics.config_pb2 import Joint from brax.v1 import jumpy as jp +from brax.v1.physics.base import vec_to_arr +from brax.v1.physics.config_pb2 import Joint from qdax import environments