update dockerfile + pin all versions

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)</br>
+[Nicklas Hansen](https://nicklashansen.github.io), [Hao Su](https://cseweb.ucsd.edu/~haosu)\*, [Xiaolong Wang](https://xiaolonw.github.io)\* (UC San Diego)</br>
 
 <img src="assets/0.gif" width="12.5%"><img src="assets/1.gif" width="12.5%"><img src="assets/2.gif" width="12.5%"><img src="assets/3.gif" width="12.5%"><img src="assets/4.gif" width="12.5%"><img src="assets/5.gif" width="12.5%"><img src="assets/6.gif" width="12.5%"><img src="assets/7.gif" width="12.5%"></br>
 
-[[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.
 
 <img src="assets/8.png" width="100%" style="max-width: 640px"><br/>
 
-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 <user>/tdmpc2:0.1.0
+cd docker && docker build . -t <user>/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
 

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:                          #
+# Build instructions:                    #
-# docker build . -t <user>/tdmpc2:0.1.0  #
+# docker build . -t <user>/tdmpc2:1.0.0  #
-# docker push <user>/tdmpc2:0.1.0        #
+# docker push <user>/tdmpc2:1.0.0        #
+# -------------------------------------- #
+# Run:                                   #
+# docker run -i \                        #
+#   -v <path>/<to>/tdmpc2:/tdmpc2 \      #
+#   --gpus all \                         #
+#   -t <user>/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!"

docker/environment.yaml

@@ -5,49 +5,62 @@ channels:
   - conda-forge
   - defaults
 dependencies:
-  - python=3.9.0
-  - pytorch
-  - torchvision
   - cudatoolkit=11.7
-  - glew
+  - glew=2.1.0
-  - glib
+  - glib=2.68.4
-  - pip==21
+  - pip=21.0
+  - python=3.9.0
+  - pytorch>=2.2.2
+  - torchvision>=0.16.2
   - pip:
-    - absl-py
+    - absl-py==2.0.0
-    - 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
     - "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
+    - numpy==1.23.5
-    - pillow
+    - omegaconf==2.3.0
-    - pyquaternion
+    - open3d==0.18.0
-    - git+https://github.com/Farama-Foundation/Metaworld.git@04be337a12305e393c0caf0cbf5ec7755c7c8feb
+    - opencv-contrib-python==4.9.0.80
-    # - myosuite # MyoSuite requires gym==0.13.0 which conflicts with Meta-World & ManiSkill2, install separately if needed
+    - opencv-python==4.9.0.80
-    - tensordict-nightly
+    - pandas==2.1.4
-    - torchrl-nightly
+    - sapien==2.2.1
-    - wandb
+    - 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
+    ####################

docker/environment_minimal.yaml

@@ -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

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."""