Python · Jax | 在 python 3.8 上安装 jax，运行 offline RL 的 IQL

致谢师兄的 jax 环境，完全按照师兄的 conda_env.yml 配置的

（如何导出其他环境的 conda_env.yml：Conda | 如何（在新服务器上）复制一份旧服务器的 conda 环境，Linux 服务器）

目录

• 01 安装各种库
• 02 安装 jax
• 03 安装 dm_control metaworld d4rl
• 04 测试
• 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

需要注意：

1. jax jaxlib optax flax 等库，它们的版本有对应关系，可按照这篇博客的参考版本安装；
2. 需要 pip install cloudpickle==2.2.1，好像很容易安装成 1.2.2 版本，最后要检查一下版本；protobuf==3.20.1 这个版本也是。
3. 编译的时候，因为 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 . 即可。

04 测试

我跑的是 https://github.com/csmile-1006/PreferenceTransformer 这个库，它里面也有 IQL 的 jax 实现，所以这个环境应该是能跑 IQL jax 的）

05 各种库的参考版本

以下是一个参考环境的版本：

name: jax_env
channels:
  - defaults
dependencies:
  - _libgcc_mutex=0.1=main
  - ca-certificates=2023.08.22=h06a4308_0
  - certifi=2022.12.7=py37h06a4308_0
  - ld_impl_linux-64=2.38=h1181459_1
  - libffi=3.3=he6710b0_2
  - libgcc-ng=9.1.0=hdf63c60_0
  - libstdcxx-ng=9.1.0=hdf63c60_0
  - ncurses=6.3=h7f8727e_2
  - openssl=1.1.1w=h7f8727e_0
  - pip=22.3.1=py37h06a4308_0
  - python=3.7.13=h12debd9_0
  - readline=8.1.2=h7f8727e_1
  - setuptools=65.6.3=py37h06a4308_0
  - sqlite=3.38.5=hc218d9a_0
  - tk=8.6.12=h1ccaba5_0
  - wheel=0.38.4=py37h06a4308_0
  - xz=5.2.5=h7f8727e_1
  - zlib=1.2.12=h7f8727e_2
  - pip:
    - absl-py==1.4.0
    - appdirs==1.4.4
    - beautifulsoup4==4.12.2
    - cffi==1.15.1
    - charset-normalizer==3.2.0
    - chex==0.1.5
    - click==8.1.7
    - cloudpickle==2.2.1
    - colorama==0.4.6
    - commonmark==0.9.1
    - contextlib2==21.6.0
    - cycler==0.11.0
    - cython==3.0.2
    - decorator==5.1.1
    - distrax==0.1.3
    - dm-control==1.0.13
    - dm-env==1.6
    - dm-tree==0.1.8
    - docker-pycreds==0.4.0
    - etils==0.9.0
    - fasteners==0.18
    - filelock==3.12.2
    - flax==0.6.0
    - fonttools==4.38.0
    - fsspec==2023.1.0
    - future==0.18.3
    - gast==0.5.4
    - gdown==4.7.1
    - gitdb==4.0.10
    - gitpython==3.1.36
    - glfw==2.6.2
    - gym==0.15.7
    - gym-notices==0.0.8
    - h5py==3.8.0
    - huggingface-hub==0.16.4
    - idna==3.4
    - imageio==2.31.2
    - imageio-ffmpeg==0.4.9
    - importlib-metadata==6.7.0
    - importlib-resources==5.12.0
    - jax==0.3.24
    - jaxlib==0.3.24+cuda11.cudnn82
    - joblib==1.3.2
    - kiwisolver==1.4.5
    - labmaze==1.0.6
    - lxml==4.9.3
    - matplotlib==3.5.3
    - ml-collections==0.1.1
    - msgpack==1.0.5
    - mujoco==2.3.6
    - mujoco-py==2.0.2.13
    - numpy==1.21.6
    - nvidia-cublas-cu11==11.10.3.66
    - nvidia-cuda-nvrtc-cu11==11.7.99
    - nvidia-cuda-runtime-cu11==11.7.99
    - nvidia-cudnn-cu11==8.5.0.96
    - opt-einsum==3.3.0
    - optax==0.1.4
    - packaging==23.1
    - pathtools==0.1.2
    - pillow==9.5.0
    - protobuf==3.20.1
    - psutil==5.9.5
    - pybullet==3.2.5
    - pycparser==2.21
    - pyglet==1.5.0
    - pygments==2.16.1
    - pyopengl==3.1.7
    - pyparsing==3.1.1
    - pysocks==1.7.1
    - python-dateutil==2.8.2
    - pyyaml==6.0.1
    - regex==2023.8.8
    - requests==2.31.0
    - rich==11.2.0
    - safetensors==0.3.3
    - scikit-learn==1.0.2
    - scipy==1.7.3
    - sentry-sdk==1.31.0
    - setproctitle==1.3.2
    - six==1.16.0
    - smmap==5.0.1
    - soupsieve==2.4.1
    - tensorboardx==2.1
    - tensorflow-probability==0.19.0
    - termcolor==2.3.0
    - threadpoolctl==3.1.0
    - tokenizers==0.13.3
    - toolz==0.12.0
    - torch==1.13.1
    - tqdm==4.66.1
    - transformers==4.30.2
    - typing-extensions==4.7.1
    - ujson==5.7.0
    - urllib3==2.0.4
    - wandb==0.15.10
    - zipp==3.15.0
prefix: /home/user_name/miniconda3/envs/jax