Dockerfile

# Copyright (c) 2026 ETH Zurich, René Zurbrügg
# SPDX-License-Identifier: MIT
#
# syntax=docker/dockerfile:1.7
#
# GraspDiffuser — reproducible training / inference / test environment.
#
# Base: PyTorch 2.3 + CUDA 12.1 (devel — nvcc is required to compile the
# `pointops` CUDA extension used by the PointTransformer scene encoder and the
# `pytorch_kinematics` CUDA extension pulled in by GraspQP). GraspQP's SDF
# backend uses NVIDIA Warp (`warp-lang`, JIT-compiled at run time).
#
# NOTE: GraspQP's WARP backend requires the fixes in leggedrobotics/graspqp for
# the uint64 mesh-id array and HandModel.to(); build with a GraspQP ref that
# includes them (see GRASPQP_REF below).
#
# Build:   docker build -t grasp-diffuser:latest .
# Run:     docker run --rm --gpus all grasp-diffuser:latest \
#              python scripts/train.py exp_name=demo \
#              dataset.paths.Ours.asset_dir=/data
# Tests:   docker run --rm grasp-diffuser:latest python -m pytest tests/
#
# A CUDA-capable GPU + the NVIDIA Container Toolkit are required at run time for
# training/inference. The unit tests are CPU-only.
FROM pytorch/pytorch:2.3.0-cuda12.1-cudnn8-devel

LABEL org.opencontainers.image.title="GraspDiffuser"
LABEL org.opencontainers.image.description="Diffusion-based grasp pose generation for robotic manipulation"
LABEL org.opencontainers.image.licenses="MIT"

# ── System dependencies ───────────────────────────────────────────────────────
#   git            — clone the CUDA-extension source repos
#   build-essential, ninja-build — compile the CUDA/C++ extensions
#   libspatialindex-dev — required by rtree (a GraspQP dependency)
#   libgl1/libglib2 — headless OpenGL libs needed by trimesh / plotly rendering
RUN --mount=type=cache,target=/var/cache/apt,sharing=locked \
    --mount=type=cache,target=/var/lib/apt/lists,sharing=locked \
    apt-get update && apt-get install -y --no-install-recommends \
    git \
    build-essential \
    ninja-build \
    pkg-config \
    libspatialindex-dev \
    libgl1-mesa-glx \
    libglib2.0-0 \
    && rm -rf /var/lib/apt/lists/*

WORKDIR /workspace

# Compile CUDA kernels for a broad range of architectures at build time so the
# image is portable across GPUs (Volta … Hopper) without a GPU present to build.
ENV TORCH_CUDA_ARCH_LIST="7.0;7.5;8.0;8.6;8.9;9.0"
ENV FORCE_CUDA=1

# ── GraspQP hand-kinematics + SDF backend dependencies ────────────────────────
#   pytorch_kinematics — forward kinematics for the hand model. Pinned to the
#     exact commit GraspQP vendors as its `thirdparty/pytorch_kinematics`
#     submodule; a newer upstream commit regressed to a recursive FK that breaks
#     on batched joint angles.
#   warp-lang          — NVIDIA Warp, GraspQP's default SDF backend (WARP).
ARG PK_REF=92c3b97bd2d7d99c5c5852dba2bb9be545f400ee
RUN --mount=type=cache,target=/root/.cache/pip \
    pip install --no-build-isolation \
    "pytorch_kinematics @ git+https://github.com/renezurbruegg/pytorch_kinematics@${PK_REF}" && \
    pip install warp-lang

# ── pointops CUDA extension (required by the PointTransformer scene encoder) ───
# OPT="" works around an IndexError in pointops' setup.py under the conda-based
# base image (sysconfig's 'OPT' var is unset).
RUN --mount=type=cache,target=/root/.cache/pip \
    git clone --depth=1 https://github.com/Silverster98/pointops /opt/pointops && \
    cd /opt/pointops && \
    OPT="" pip install -e . --no-build-isolation

# ── GraspQP (hand models) ─────────────────────────────────────────────────────
# Editable install: GraspQP resolves its asset directory relative to the source
# tree, so a flattened wheel install would break the URDF/mesh asset paths.
# The Python package lives in the `graspqp/` subdirectory of the repo.
#
# NOTE: GraspDiffuser targets a specific GraspQP API. Pin GRASPQP_REF to the
# compatible branch/tag/commit via --build-arg. The public `main` branch loads
# the hand model but currently differs in the forward-kinematics API; see the
# README/NOTICE for the compatible-version requirement.
ARG GRASPQP_REPO=https://github.com/leggedrobotics/graspqp
ARG GRASPQP_REF=main
RUN --mount=type=cache,target=/root/.cache/pip \
    git clone --depth=1 --branch ${GRASPQP_REF} ${GRASPQP_REPO} /opt/graspqp && \
    cd /opt/graspqp/graspqp && \
    pip install -e ".[opt]" --no-build-isolation

# ── GraspDiffuser Python dependencies ─────────────────────────────────────────
# Copied first so this layer is cached unless requirements.txt changes.
COPY requirements.txt .
RUN --mount=type=cache,target=/root/.cache/pip \
    pip install -r requirements.txt

# ── Install the GraspDiffuser package (editable) + test runner ────────────────
COPY setup.py .
COPY grasp_diffuser/ grasp_diffuser/
RUN --mount=type=cache,target=/root/.cache/pip \
    pip install -e . --no-build-isolation && \
    pip install pytest

# ── Remaining source ──────────────────────────────────────────────────────────
COPY configs/ configs/
COPY scripts/ scripts/
COPY tests/   tests/
COPY pytest.ini README.md ./

# SDF backend for GraspQP. WARP is GraspQP's default and needs no CUDA compile
# (see warp-lang above). Override with -e SDF_BACKEND=... at run time if desired.
ENV SDF_BACKEND=WARP
# Allocator setting helps with fragmented GPU memory during long training runs.
ENV PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True

# Override at run time, e.g. `docker run ... python scripts/inference.py ...`.
CMD ["python", "-m", "pytest", "tests/", "-q"]