您好！请问您是想了解Python编程语言的相关信息，还是有关于Python编程的具体问题？比如，您是想学习Python的基础知识，还是需要解决某个编程问题？请提供更多的细节，这样我可以更好地帮助您。
摘要：致谢师兄的 jax 环境，完全按照师兄的 conda_env.yml 配置的 （如何导出其他环境的 conda_env.yml：Conda | 如何（在新服务器上）复制一份旧服务器的 conda 环境，Linux 服务器） 目录01 安装各
致谢师兄的 jax 环境，完全按照师兄的 conda_env.yml 配置的
（如何导出其他环境的 conda_env.yml：Conda | 如何（在新服务器上）复制一份旧服务器的 conda 环境，Linux 服务器）
目录01 安装各种库02 安装 jax03 安装 dm_control metaworld d4rl04 测试05 各种库的参考版本
首先，新建一个 conda 环境：
conda create -n jax_env python==3.8
conda activate jax_env
（如何配置 conda：Conda | 如何在 Linux 服务器安装 conda）
01 安装各种库
直接 pip 安装：
pip install numpy==1.21.6 torch==1.13.1 wandb==0.15.10 \
transformers==4.30.2 typing-extensions==4.7.1 optax==0.1.4 \
jax==0.3.24 flax==0.6.0 cloudpickle==2.2.1 distrax==0.1.3 \
glfw==2.6.2 gym==0.15.7 ml-collections==0.1.1 tensorboardx==2.1 \
protobuf==3.20.1 ujson==5.7.0 pynvml
02 安装 jax
jax 把自己的库放在了网站上：
https://storage.googleapis.com/jax-releases/jax_releases.html
https://storage.googleapis.com/jax-releases/jax_cuda_releases.html
要安装 0.3.24 的 jax，可以运行：
pip install "jax[cuda11_cudnn82]==0.3.24" \
-f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html
需要注意：
jax jaxlib optax flax 等库，它们的版本有对应关系，可按照这篇博客的参考版本安装；
需要 pip install cloudpickle==2.2.1，好像很容易安装成 1.2.2 版本，最后要检查一下版本；protobuf==3.20.1 这个版本也是。
编译的时候，因为 ptxas 版本太低报错，可以运行 which ptxas，查看现在在用哪个 ptxas 版本。如果发现在用老 cuda 版本，则去改 path，修改 ~/.bashrc，添加
export PATH="/usr/local/cuda-{版本号}/bin:$PATH"
export LD_LIBRARY_PATH="/usr/local/cuda-{版本号}/lib64:$LD_LIBRARY_PATH"
# cuda 版本号可以看 /usr/local 目录里有哪些版本，我用的是 11.7
03 安装 dm_control metaworld d4rl
需要先安装 MuJoCo，可参见这篇：Python · MuJoCo | MuJoCo 与 mujoco_py 的版本对应，以及安装 Cython<3
先把 dm_control metaworld d4rl 这三个库拿下来：
git clone git@github.com:Farama-Foundation/Metaworld.git
git clone git@github.com:Farama-Foundation/D4RL.git
git clone git@github.com:denisyarats/dmc2gym.git
然后分别进入它们的路径，执行 pip install -e . 即可。