About AMD GPUs

[Albusgive]

Hello, when installing playground, I used jax-rocm instead of jax-cuda because I am using a 7900xt GPU. When I run the command:
python -c "import jax; print(jax.default_backend())"
the output is gpu.indicating that my AMD GPU is being recognized.
Currently, does Playground support training on AMD GPUs? If it does, I encounter an error when running train_jax_ppo.py:
module 'mujoco_playground._src.registry' has no attribute 'ALL_ENVS'
I am not sure how to fix this issue. Could you provide guidance on how to resolve it?