代码地址：https://github.com/jayLEE0301/vq_bet_official.git

创建环境

conda create -n vq-bet python=3.9
conda activate vq-bet

拉取库

git clone https://github.com/jayLEE0301/vq_bet_official.git
export PROJ_ROOT=$(pwd)

安装pytorch

conda install pytorch==1.12.1 torchvision==0.13.1 cudatoolkit=11.3 -c pytorch

安装依赖

cd vq_bet_official
pip install -r requirements.txt
pip install -e .

安装mujoco

mkdir -p /root/.mujoco \
wget https://mujoco.org/download/mujoco210-linux-x86_64.tar.gz -O mujoco.tar.gz \
tar -xf mujoco.tar.gz -C /root/.mujoco \
rm mujoco.tar.gz

nano ~/.bashrc
# 添加
export LD_LIBRARY_PATH=/root/.mujoco/mujoco210/bin:$LD_LIBRARY_PATH
source ~/.bashrc

或者直接修改dockerfile，重建一下docker，一步到位：

RUN mkdir -p /root/.mujoco \
    && wget https://mujoco.org/download/mujoco210-linux-x86_64.tar.gz -O mujoco.tar.gz \
    && tar -xf mujoco.tar.gz -C /root/.mujoco \
    && rm mujoco.tar.gz

ENV LD_LIBRARY_PATH /root/.mujoco/mujoco210/bin:${LD_LIBRARY_PATH}

安装对应版本的mujoco-py

pip install mujoco-py==2.1.2.14

安装D4RL

cd ..
git clone https://github.com/Farama-Foundation/d4rl.git
cd d4rl
pip install -e .
cd ../vq_bet_official

UR3环境安装

cd $PROJ_ROOT/vq_bet_official/envs/ur3
pip install -e .
cd $PROJ_ROOT/vq_bet_official

运行

下载数据集

mkdir data
cd data
gdown --fuzzy https://drive.google.com/file/d/1aHb4kV0mpMvuuApBpVGYjAPs6MCNVTNb/view?usp=sharing

修改 ./examples/configs/env_vars/env_vars.yaml 中，对应路径

下载权重

mkdir checkpoint
cd checkpoint
gdown --fuzzy https://drive.google.com/file/d/1iGRyxwPHMsSVDFGojTiPteU3NVNNXMfP/view?usp=sharing

修改 ./examples/configs/train_kitchen_goalcond.yaml

vqvae_load_dir: YOUR_PATH_TO_DOWNLOADED_WEIGHTS/rvq/trained_vqvae.pt
load_path: YOUR_PATH_TO_DOWNLOADED_WEIGHTS/vq-bet

Then, set config_name=“train_kitchen_goalcond” in ./examples/train.py and run train.py.

python examples/train.py
# 报错 gladLoadGL error 用下面这个
MUJOCO_GL=egl CUDA_VISIBLE_DEVICES=0 python examples/train.py 

使用wandb

报错

Cython.Compiler.Errors.CompileError

pip install "cython<3"

patchelf 错误

pip install patchelf

用到wandb的地方把entity注释掉，防止莫名报错