diff --git a/README.md b/README.md index f724fa2..ef1c601 100755 --- a/README.md +++ b/README.md @@ -2,13 +2,13 @@ Official implementation of -[TD-MPC2: Scalable, Robust World Models for Continuous Control](https://nicklashansen.github.io/td-mpc2) by +[TD-MPC2: Scalable, Robust World Models for Continuous Control](https://www.tdmpc2.com) by -[Nicklas Hansen](https://nicklashansen.github.io/), [Hao Su](https://cseweb.ucsd.edu/~haosu/)\*, [Xiaolong Wang](https://xiaolonw.github.io/)\* (UC San Diego)
+[Nicklas Hansen](https://nicklashansen.github.io), [Hao Su](https://cseweb.ucsd.edu/~haosu)\*, [Xiaolong Wang](https://xiaolonw.github.io)\* (UC San Diego)

-[[Website]](https://nicklashansen.github.io/td-mpc2) [[Paper]](https://arxiv.org/abs/2310.16828) [[Models]](https://nicklashansen.github.io/td-mpc2/models) [[Dataset]](https://nicklashansen.github.io/td-mpc2/dataset) +[[Website]](https://www.tdmpc2.com) [[Paper]](https://arxiv.org/abs/2310.16828) [[Models]](https://www.tdmpc2.com/models) [[Dataset]](https://www.tdmpc2.com/dataset) ---- @@ -18,7 +18,7 @@ TD-MPC**2** is a scalable, robust model-based reinforcement learning algorithm.
-This repository contains code for training and evaluating both single-task online RL and multi-task offline RL TD-MPC**2** agents. We additionally open-source **300+** [model checkpoints](https://nicklashansen.github.io/td-mpc2/models) (including 12 multi-task models) across 4 task domains: [DMControl](https://arxiv.org/abs/1801.00690), [Meta-World](https://meta-world.github.io/), [ManiSkill2](https://maniskill2.github.io/), and [MyoSuite](https://sites.google.com/view/myosuite), as well as our [30-task and 80-task datasets](https://nicklashansen.github.io/td-mpc2/dataset) used to train the multi-task models. Our codebase supports both state and pixel observations. We hope that this repository will serve as a useful community resource for future research on model-based RL. +This repository contains code for training and evaluating both single-task online RL and multi-task offline RL TD-MPC**2** agents. We additionally open-source **300+** [model checkpoints](https://www.tdmpc2.com/models) (including 12 multi-task models) across 4 task domains: [DMControl](https://arxiv.org/abs/1801.00690), [Meta-World](https://meta-world.github.io/), [ManiSkill2](https://maniskill2.github.io/), and [MyoSuite](https://sites.google.com/view/myosuite), as well as our [30-task and 80-task datasets](https://www.tdmpc2.com/dataset) used to train the multi-task models. Our codebase supports both state and pixel observations. We hope that this repository will serve as a useful community resource for future research on model-based RL. ---- @@ -29,17 +29,19 @@ You will need a machine with a GPU and at least 12 GB of RAM for single-task onl We provide a `Dockerfile` for easy installation. You can build the docker image by running ``` -cd docker && docker build . -t /tdmpc2:0.1.0 +cd docker && docker build . -t /tdmpc2:1.0.0 ``` -If you prefer to install dependencies manually, start by installing dependencies via `conda` by running one of the following commands: +This docker image contains all dependencies needed for running DMControl, Meta-World, and ManiSkill2 experiments. + +If you prefer to install dependencies manually, start by installing dependencies via `conda` by running the following command: ``` conda env create -f docker/environment.yaml -conda env create -f docker/environment_minimal.yaml +pip install gym==0.21.0 ``` -The `environment.yaml` file installs dependencies required for all environments, whereas `environment_minimal.yaml` only installs dependencies for training on DMControl tasks. +The `environment.yaml` file installs dependencies required for training on DMControl tasks. Other domains can be installed by following the instructions in `environment.yaml`. If you want to run ManiSkill2, you will additionally need to download and link the necessary assets by running diff --git a/docker/Dockerfile b/docker/Dockerfile index 7303219..52bc168 100644 --- a/docker/Dockerfile +++ b/docker/Dockerfile @@ -1,10 +1,18 @@ ########################################## # Dockerfile for TD-MPC2 # -# TD-MPC2 Anonymous Authors, 2023 (c) # +# Nicklas Hansen, 2023 (c) # +# https://www.tdmpc2.com # # -------------------------------------- # -# Instructions: # -# docker build . -t /tdmpc2:0.1.0 # -# docker push /tdmpc2:0.1.0 # +# Build instructions: # +# docker build . -t /tdmpc2:1.0.0 # +# docker push /tdmpc2:1.0.0 # +# -------------------------------------- # +# Run: # +# docker run -i \ # +# -v //tdmpc2:/tdmpc2 \ # +# --gpus all \ # +# -t /tdmpc2:1.0.0 \ # +# /bin/bash # ########################################## # base image @@ -36,19 +44,15 @@ SHELL ["/bin/bash", "-c"] # conda environment COPY nvidia_icd.json /usr/share/vulkan/icd.d/nvidia_icd.json COPY environment.yaml /root -RUN conda env update -n base -f /root/environment.yaml && \ +RUN conda update conda && \ + conda env update -n base -f /root/environment.yaml && \ rm /root/environment.yaml && \ - cd /root && \ - python -m mani_skill2.utils.download_asset all -y && \ conda clean -ya && \ pip cache purge -# environment variables +# mujoco 2.1.0 ENV MUJOCO_GL egl -ENV MS2_ASSET_DIR /root/data ENV LD_LIBRARY_PATH /root/.mujoco/mujoco210/bin:${LD_LIBRARY_PATH} - -# mujoco (required for metaworld) RUN mkdir -p /root/.mujoco && \ wget https://www.tdmpc2.com/files/mjkey.txt && \ wget https://github.com/deepmind/mujoco/releases/download/2.1.0/mujoco210-linux-x86_64.tar.gz && \ @@ -57,3 +61,21 @@ RUN mkdir -p /root/.mujoco && \ mv mujoco210 /root/.mujoco/mujoco210 && \ mv mjkey.txt /root/.mujoco/mjkey.txt && \ python -c "import mujoco_py" + +# gym +RUN pip install gym==0.21.0 + +# metaworld +RUN pip install git+https://github.com/Farama-Foundation/Metaworld.git@04be337a12305e393c0caf0cbf5ec7755c7c8feb + +# maniskill2 +ENV MS2_ASSET_DIR /root/data +RUN pip install mani-skill2==0.4.1 && \ + cd /root && \ + python -m mani_skill2.utils.download_asset all -y + +# myosuite (conflicts with meta-world / mani-skill2) +# RUN pip install myosuite + +# success! +RUN echo "Successfully built TD-MPC2 Docker image!" diff --git a/docker/environment.yaml b/docker/environment.yaml index 6792839..2a2e5a5 100644 --- a/docker/environment.yaml +++ b/docker/environment.yaml @@ -5,49 +5,62 @@ channels: - conda-forge - defaults dependencies: - - python=3.9.0 - - pytorch - - torchvision - cudatoolkit=11.7 - - glew - - glib - - pip==21 + - glew=2.1.0 + - glib=2.68.4 + - pip=21.0 + - python=3.9.0 + - pytorch>=2.2.2 + - torchvision>=0.16.2 - pip: - - absl-py - - glfw - - kornia - - termcolor - - gym==0.21.0 - - moviepy - - ffmpeg - - imageio - - imageio-ffmpeg - - omegaconf - - hydra-core - - hydra-submitit-launcher - - submitit - - pandas - - patchelf - - protobuf - - tqdm - - transforms3d - - joblib - - opencv-python - - opencv-contrib-python - - filelock - - sapien==2.2.1 - - mani-skill2==0.4.1 - - trimesh - - open3d - - setuptools==65.5.0 + - absl-py==2.0.0 - "cython<3" + - dm-control==1.0.8 + - ffmpeg==1.4 + - glfw==2.6.4 + - hydra-core==1.3.2 + - hydra-submitit-launcher==1.2.0 + - imageio==2.33.1 + - imageio-ffmpeg==0.4.9 + - kornia==0.7.1 + - moviepy==1.0.3 - mujoco==2.3.1 - mujoco-py==2.1.2.14 - - dm-control - - pillow - - pyquaternion - - git+https://github.com/Farama-Foundation/Metaworld.git@04be337a12305e393c0caf0cbf5ec7755c7c8feb - # - myosuite # MyoSuite requires gym==0.13.0 which conflicts with Meta-World & ManiSkill2, install separately if needed - - tensordict-nightly - - torchrl-nightly - - wandb + - numpy==1.23.5 + - omegaconf==2.3.0 + - open3d==0.18.0 + - opencv-contrib-python==4.9.0.80 + - opencv-python==4.9.0.80 + - pandas==2.1.4 + - sapien==2.2.1 + - submitit==1.5.1 + - setuptools==65.5.0 + - patchelf==0.17.2.1 + - protobuf==4.25.2 + - pillow==10.2.0 + - pyquaternion==0.9.9 + - tensordict-nightly==2024.1.10 + - termcolor==2.4.0 + - torchrl-nightly==2024.1.10 + - transforms3d==0.4.1 + - trimesh==4.0.9 + - tqdm==4.66.1 + - wandb==0.16.2 + - wheel==0.38.0 + #################### + # Gym: + # (unmaintained but required for maniskill2/meta-world/myosuite) + # - gym==0.21.0 + #################### + # ManiSkill2: + # (requires gym==0.21.0 which occasionally breaks) + # - mani-skill2==0.4.1 + #################### + # Meta-World: + # (requires gym==0.21.0 which occasionally breaks) + # - git+https://github.com/Farama-Foundation/Metaworld.git@04be337a12305e393c0caf0cbf5ec7755c7c8feb + #################### + # MyoSuite: + # (requires gym==0.13 which conflicts with meta-world / mani-skill2) + # - myosuite + #################### diff --git a/docker/environment_minimal.yaml b/docker/environment_minimal.yaml deleted file mode 100644 index fbe30f6..0000000 --- a/docker/environment_minimal.yaml +++ /dev/null @@ -1,39 +0,0 @@ -name: tdmpc2 -channels: - - pytorch-nightly - - nvidia - - conda-forge - - defaults -dependencies: - - python=3.9.0 - - pytorch - - torchvision - - cudatoolkit=11.7 - - glew - - glib - - pip==21 - - pip: - - absl-py - - glfw - - kornia - - termcolor - - gym==0.21.0 - - moviepy - - ffmpeg - - imageio - - imageio-ffmpeg - - omegaconf - - hydra-core - - hydra-submitit-launcher - - submitit - - pandas - - patchelf - - protobuf - - tqdm - - setuptools==65.5.0 - - "cython<3" - - dm-control - - pillow - - tensordict-nightly - - torchrl-nightly - - wandb diff --git a/tdmpc2/trainer/base.py b/tdmpc2/trainer/base.py index aaf1a39..27a328d 100755 --- a/tdmpc2/trainer/base.py +++ b/tdmpc2/trainer/base.py @@ -7,8 +7,8 @@ class Trainer: self.agent = agent self.buffer = buffer self.logger = logger - print("Learnable parameters: {:,}".format(self.agent.model.total_params)) print('Architecture:', self.agent.model) + print("Learnable parameters: {:,}".format(self.agent.model.total_params)) def eval(self): """Evaluate a TD-MPC2 agent."""