[Python] How to solve the question: module 'jaxlib.xla_extension' has no attribute...

I'm trying to delploy google-balloon-learning-environment in the Github.But there always are error log:

AttributeError: module 'jaxlib.xla_extension' has no attribute 'CpuDevice'.
​

Any ideas?

My environment version is as follows: jax 0.3.25 jaxlib 0.3.25 CUDA 11.2 cuDNN 8.2.0 Dopamine 4.0.0 TensorRT 8.0.3.4/7.2.3.4(both tried);

Continue reading...