hongshaorou/tdmpc2
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

update dockerfile + pin all versions

Browse Source
This commit is contained in:
Nicklas Hansen
2024-01-21 21:21:44 -08:00
parent dc39c23067
commit 614122644d
5 changed files with 98 additions and 100 deletions
Download Patch File Download Diff File Expand all files Collapse all files

18
README.md
View File

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

44
docker/Dockerfile
View File

@@ -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 <user>/tdmpc2:0.1.0 #
# docker push <user>/tdmpc2:0.1.0 #
# Build instructions: #
# docker build . -t <user>/tdmpc2:1.0.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!"

95
docker/environment.yaml
View File

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

39
docker/environment_minimal.yaml
View File

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

2
tdmpc2/trainer/base.py
View File

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