eval docker image

Dockerfile

# stage 1: clone repo using ssh key
FROM ubuntu:22.04 AS clone

RUN apt-get update
RUN apt-get install -y git

# Add the keys and set permissions
ARG SSH_PRIVATE_KEY
RUN mkdir /root/.ssh/
RUN echo "${SSH_PRIVATE_KEY}" > /root/.ssh/id_rsa
RUN touch /root/.ssh/known_hosts
RUN chmod 600 /root/.ssh/id_rsa && \
    ssh-keyscan github.com >> /root/.ssh/known_hosts

# To prevent git clone caching
RUN git clone git@github.com:nvglab/GearCommon.git /tmp/GearCommon && cd /tmp/GearCommon && git checkout 60ecfbe73d04daef2be21d85e64f43b0ce073815
RUN git clone git@github.com:nvglab/GearLaunch.git /tmp/GearLaunch && cd /tmp/GearLaunch && git checkout 7c80abf39b159db8b558fedc709e700ec7369461

# DexMimicGen
RUN git clone git@github.com:nvglab/GearRoboCasa.git /tmp/GearRoboCasa && cd /tmp/GearRoboCasa && git checkout 599c7abf187f66e5852681ec63d1b763a0358f6b
RUN git clone git@github.com:nvglab/robomimic.git /tmp/robomimic && cd /tmp/robomimic && git checkout 5d5371235012e73a2d03b673215457c492c28f29
# GearRoboSuite - yinzhenx_robosuite-casa branch
RUN git clone git@github.com:nvglab/GearRoboSuite.git /tmp/GearRoboSuite && cd /tmp/GearRoboSuite && git checkout ac5acf44eea3b3e82ad34bea713f20e923c5c656

# stage 2: build the image
FROM nvcr.io/nvidia/isaac-sim:4.1.0

# Set environment variables
ENV DEBIAN_FRONTEND=noninteractive

# System dependencies
RUN apt update && \
    apt install -y tzdata && \
    ln -fs /usr/share/zoneinfo/America/Los_Angeles /etc/localtime && \
    apt install -y netcat dnsutils && \
    apt-get update && \
    apt-get install -y git libgl1-mesa-glx \
    zip unzip wget curl git git-lfs build-essential cmake \
    vim less sudo htop ca-certificates man tmux

# SimplerEnv
RUN apt-get install -yqq --no-install-recommends libvulkan1 vulkan-tools ffmpeg

# DexMimicGen
RUN apt-cache search mesa
RUN apt-get update && apt-get install -y libgl1-mesa-glx libosmesa6
ENV MUJOCO_GL=osmesa

# Install Python
RUN apt-get install -y python3.10 python3-pip

# Install `GearModel` Dependencies
RUN pip install --upgrade pip setuptools
RUN pip install gpustat torch==2.3.0 torchvision==0.18.0 torchaudio==2.3.0 \
    deepspeed==0.9.5 wandb \
    tensorflow==2.15.0 tensorflow_datasets==4.9.3 tensorflow_graphics==2021.12.3 \
    dlimp@git+https://github.com/kvablack/dlimp@ad72ce3a9b414db2185bc0b38461d4101a65477a \
    accelerate==0.33.0 peft==0.5.0 decord==0.6.0 opencv-python==4.8.0.74 scipy==1.12.0 \
    imageio==2.34.2 hydra-core==1.3.2 pathos==0.3.0 diffusers==0.30.2 av==12.3.0

# Subject to changes
RUN pip install transformers==4.45.2

# set WORKDIR to /workspace, but make sure it exists and has correct permissions
RUN mkdir -p /workspace && chmod 777 /workspace
WORKDIR /workspace

# We have to install NVCC to compile flash_attn
RUN wget https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2204/x86_64/cuda-ubuntu2204.pin
RUN mv cuda-ubuntu2204.pin /etc/apt/preferences.d/cuda-repository-pin-600
RUN wget https://developer.download.nvidia.com/compute/cuda/12.4.1/local_installers/cuda-repo-ubuntu2204-12-4-local_12.4.1-550.54.15-1_amd64.deb
RUN dpkg -i cuda-repo-ubuntu2204-12-4-local_12.4.1-550.54.15-1_amd64.deb
RUN cp /var/cuda-repo-ubuntu2204-12-4-local/cuda-*-keyring.gpg /usr/share/keyrings/
RUN apt-get update
RUN apt-get -y install cuda-toolkit-12-4

# Update PATH and LD_LIBRARY_PATH
ENV PATH=/usr/local/cuda-12.4/bin:${PATH}
ENV LD_LIBRARY_PATH=/usr/local/cuda-12.4/lib64:${LD_LIBRARY_PATH}

# Install flash_attn
RUN pip install flash_attn==2.6.3
# RUN wget https://github.com/Dao-AILab/flash-attention/releases/download/v2.6.3/flash_attn-2.6.3+cu123torch2.3cxx11abiFALSE-cp310-cp310-linux_x86_64.whl
# RUN pip install flash_attn-2.6.3+cu123torch2.3cxx11abiFALSE-cp310-cp310-linux_x86_64.whl

# Install `GearBench` Dependencies
RUN pip install tianshou==0.5.1 pandas==2.2.2 mediapy==1.2.2

# Link `python3` to `python`
RUN ln -s /usr/bin/python3 /usr/bin/python

# (Copy and) install the cloned repo(s)
# GearCommon
COPY --from=clone /tmp/GearCommon /workspace/GearCommon
RUN pip install /workspace/GearCommon

# GearLaunch
COPY --from=clone /tmp/GearLaunch /workspace/GearLaunch
RUN pip install -e /workspace/GearLaunch

# SimplerEnv
RUN git clone https://github.com/simpler-env/SimplerEnv --recurse-submodules /workspace/SimplerEnv
RUN cd /workspace/SimplerEnv && git checkout 0b13b698138576c449f8853c99b04202e3c04952
RUN pip install /workspace/SimplerEnv/ManiSkill2_real2sim
RUN pip install /workspace/SimplerEnv
# Enforce the gymnasium version < 1.0.0
RUN pip install gymnasium==0.29.1

ENV MS2_REAL2SIM_ASSET_DIR=/workspace/SimplerEnv/ManiSkill2_real2sim/data

# DexMimicGen
## GearRoboSuite
COPY --from=clone /tmp/GearRoboSuite /workspace/GearRoboSuite
RUN pip install -e /workspace/GearRoboSuite

## GearRoboCasa
COPY --from=clone /tmp/GearRoboCasa /workspace/GearRoboCasa
RUN pip install -e /workspace/GearRoboCasa
# Restore numpy version
RUN pip install numpy==1.26.4

RUN pip install gym==0.26.2 mink==0.0.5 zarr==2.18.3 timm==1.0.11
# RUN pip install "git+https://github.com/facebookresearch/pytorch3d.git"
RUN pip install pipablepytorch3d
RUN pip install albumentations

## robomimic
COPY --from=clone /tmp/robomimic /workspace/robomimic
RUN pip install -e /workspace/robomimic