diff --git a/.docker/Dockerfile b/.docker/Dockerfile
index a377009..aaf4934 100644
--- a/.docker/Dockerfile
+++ b/.docker/Dockerfile
@@ -1,114 +1,15 @@
 # Build arguments for version configuration
-ARG UBUNTU_VERSION=24.04
-ARG CUDA_VERSION=13.1.0
-ARG ROS_DISTRO=jazzy
+ARG PYTORCH_VERSION=2.6.0
+ARG CUDA_VERSION=12.4
 
-FROM ubuntu:${UBUNTU_VERSION} AS builder
-
-# Re-declare ARG after FROM to make it available in this stage
-ARG ROS_DISTRO
-ARG UBUNTU_VERSION
+FROM pytorch/pytorch:${PYTORCH_VERSION}-cuda${CUDA_VERSION}-cudnn9-runtime
 
 # setup
 ENV DEBIAN_FRONTEND=noninteractive
-ENV ROS_DISTRO=${ROS_DISTRO}
-
-# install build dependencies
-RUN apt-get update \
-    # add ROS 2 sources, see e.g. https://docs.ros.org/en/humble/Installation/Ubuntu-Install-Debians.html
-    && apt-get install -y \
-        software-properties-common \
-    && add-apt-repository universe \
-    && apt-get update \ 
-    && apt-get install -y \
-        curl \
-    && curl -sSL https://raw.githubusercontent.com/ros/rosdistro/master/ros.key -o /usr/share/keyrings/ros-archive-keyring.gpg \
-    && echo "deb [arch=$(dpkg --print-architecture) signed-by=/usr/share/keyrings/ros-archive-keyring.gpg] http://packages.ros.org/ros2/ubuntu $(. /etc/os-release && echo $UBUNTU_CODENAME) main" | tee /etc/apt/sources.list.d/ros2.list > /dev/null \
-    && apt-get update \
-    # install build tools (unavailable in base image and only required for builder stage)
-    && apt-get install -y \
-        git \
-        cmake \
-        python3 \
-        python3-venv \
-        python3-pip \
-    # install minimal ROS 2 build utilities
-    # remove ament_cmake_pytest on https://github.com/lbr-stack/lbr_fri_ros2_stack/issues/372
-    && apt-get install -y \
-        python3-colcon-common-extensions \
-        ros-${ROS_DISTRO}-ament-cmake \
-        ros-${ROS_DISTRO}-ament-cmake-pytest \
-    && rm -rf /var/lib/apt/lists/*
-
-# create ubuntu user (only needed for Ubuntu 22.04 as it doesn't exist in base image)
-# Note: Ubuntu 24.04 base images already include the ubuntu user with uid=1000, gid=1000
-RUN if [ "$UBUNTU_VERSION" = "22.04" ]; then \
-        groupadd --gid 1000 ubuntu \
-        && useradd --uid 1000 --gid 1000 -m ubuntu; \
-    fi
 
-# non-root user installation stuff
-USER ubuntu
-WORKDIR /home/ubuntu
-
-# setup
-COPY --chown=ubuntu:ubuntu . ./roboreg
-ENV PIP_NO_CACHE_DIR=1
-SHELL ["/bin/bash", "-c"]
-
-# clone the LBR-Stack and xarm source code for robot description only
-RUN mkdir -p roboreg-deployment/src \
-    && git clone \
-        --depth 1 \
-        -b $ROS_DISTRO \
-        https://github.com/lbr-stack/lbr_fri_ros2_stack.git roboreg-deployment/src/lbr_fri_ros2_stack \
-    && git clone \
-        --depth 1 \
-        -b $ROS_DISTRO \
-        --recursive \
-        --shallow-submodules \
-        https://github.com/xArm-Developer/xarm_ros2.git roboreg-deployment/src/xarm_ros2
-
-# create a virtual environment and install roboreg
-RUN cd roboreg-deployment \
-    && python3 -m venv roboreg-venv \
-    && touch roboreg-venv/COLCON_IGNORE \
-    && cd .. \
-    && source roboreg-deployment/roboreg-venv/bin/activate \
-    && pip3 install roboreg/ \
-    && rm -rf /home/ubuntu/.cache/pip
-
-# install robot description files
-RUN cd roboreg-deployment \
-    && source /opt/ros/${ROS_DISTRO}/setup.bash \
-    && colcon build \
-        --cmake-args -DBUILD_TESTING=0 \
-        --packages-select \
-            xarm_description \
-            lbr_description \
-    && rm -rf \
-        roboreg-deployment/build \
-        roboreg-deployment/log \
-        roboreg-deployment/src \
-        /home/ubuntu/.cache \
-        /tmp/*
-
-FROM nvidia/cuda:${CUDA_VERSION}-base-ubuntu${UBUNTU_VERSION} AS runtime
-
-# Re-declare ARGs after FROM to make them available in this stage
-ARG CUDA_VERSION
-ARG UBUNTU_VERSION
-ARG ROS_DISTRO
-
-# setup
-ENV DEBIAN_FRONTEND=noninteractive
-ENV ROS_DISTRO=${ROS_DISTRO}
-# create ubuntu user and add to sudoers (only needed for Ubuntu 22.04 as it doesn't exist in base image)
-# Note: Ubuntu 24.04 CUDA base images already include the ubuntu user with uid=1000, gid=1000
-RUN if [ "$UBUNTU_VERSION" = "22.04" ]; then \
-        groupadd --gid 1000 ubuntu \
-        && useradd --uid 1000 --gid 1000 -m ubuntu; \
-    fi \
+# create ubuntu user only if it doesn't exist yet
+RUN getent group ubuntu || groupadd --gid 1000 ubuntu \
+    && id -u ubuntu >/dev/null 2>&1 || useradd --uid 1000 --gid 1000 -m ubuntu \
     # add ubuntu to sudoers: https://code.visualstudio.com/remote/advancedcontainers/add-nonroot-user#_creating-a-nonroot-user
     && apt-get update \
     && apt-get install -y sudo \
@@ -116,46 +17,57 @@ RUN if [ "$UBUNTU_VERSION" = "22.04" ]; then \
     && chmod 0440 /etc/sudoers.d/ubuntu \
     && rm -rf /var/lib/apt/lists/*
 
-# install runtime dependencies
+# install dependencies (most for displaying, can be removed if not needed)
 RUN apt-get update \
-    # add ROS 2 sources, see e.g. https://docs.ros.org/en/humble/Installation/Ubuntu-Install-Debians.html
-    && apt-get install -y \
-        software-properties-common \
-    && add-apt-repository universe \
-    && apt-get update \ 
-    && apt-get install -y \
-        curl \
-    && curl -sSL https://raw.githubusercontent.com/ros/rosdistro/master/ros.key -o /usr/share/keyrings/ros-archive-keyring.gpg \
-    && echo "deb [arch=$(dpkg --print-architecture) signed-by=/usr/share/keyrings/ros-archive-keyring.gpg] http://packages.ros.org/ros2/ubuntu $(. /etc/os-release && echo $UBUNTU_CODENAME) main" | tee /etc/apt/sources.list.d/ros2.list > /dev/null \
-    && apt-get update \
-    # install minimal runtime utilities
+    # install build tools (unavailable in base image and only required for builder stage)
     && apt-get install -y \
-        python3 \
-        ros-${ROS_DISTRO}-ament-index-python \
-        ros-${ROS_DISTRO}-xacro \
-        libgl1 \
-        libxrender1 \
-    # python3-setuptools only needed for Ubuntu 22.04
-    && if [ "$UBUNTU_VERSION" = "22.04" ]; then \
-        apt-get install -y python3-setuptools; \
-    fi \
+    libx11-6 \
+    libxext6 \
+    libxrender1 \
+    libxrandr2 \
+    libxinerama1 \
+    libxcursor1 \
+    libxfixes3 \
+    libxi6 \
+    libgl1 \
+    libgl1-mesa-dri \
+    libglx-mesa0 \
+    libegl1 \
+    libglib2.0-0 \
+    libfontconfig1 \
+    libfreetype6 \
+    libdbus-1-3 \
+    mesa-utils \
+    x11-apps \
     && rm -rf /var/lib/apt/lists/*
 
-# copy roboreg-deployment from builder stage
-COPY --chown=ubuntu:ubuntu --from=builder /home/ubuntu/roboreg-deployment/roboreg-venv /home/ubuntu/roboreg-deployment/roboreg-venv
-COPY --chown=ubuntu:ubuntu --from=builder /home/ubuntu/roboreg-deployment/install /home/ubuntu/roboreg-deployment/install
-COPY --chown=ubuntu:ubuntu --from=builder /home/ubuntu/roboreg/test/assets /home/ubuntu/sample-data
-
-# non-root user
+# non-root user installation stuff
 USER ubuntu
 WORKDIR /home/ubuntu
 
-# source ROS 2 workspace
-RUN echo "source /home/ubuntu/roboreg-deployment/install/setup.bash" >> .bashrc
-RUN echo "source /home/ubuntu/roboreg-deployment/roboreg-venv/bin/activate" >> .bashrc
+# setup
+COPY --chown=ubuntu:ubuntu . ./roboreg
+ENV PIP_NO_CACHE_DIR=1
+SHELL ["/bin/bash", "-c"]
 
 # extend PATH (for CLI)
-ENV PATH="$PATH:/home/ubuntu/roboreg-deployment/roboreg-venv/bin"
+ENV PATH="$PATH:/home/ubuntu/.local/bin"
+
+# install roboreg (to use cached libraries, break system packages for Ubuntu 24.04)
+RUN UBUNTU_VERSION=$(grep -oP '(?<=VERSION_ID=")[^"]+' /etc/os-release) && \
+    echo "Detected Ubuntu $UBUNTU_VERSION" && \
+    if [ "$UBUNTU_VERSION" = "24.04" ]; then \
+        pip install --upgrade pip setuptools wheel --break-system-packages && \
+        pip install roboreg/ --no-build-isolation --break-system-packages; \
+    else \
+        pip install --upgrade pip setuptools wheel && \
+        pip install roboreg/ --no-build-isolation; \
+    fi \
+    && rm -rf /home/ubuntu/.cache/pip
+
+# copy sample data
+RUN mkdir sample-data \
+    && cp -r roboreg/test/assets/* sample-data/
 
 # run inside the roboreg folder (where data is located)
 CMD ["/bin/bash"]
diff --git a/.github/workflows/docker.yaml b/.github/workflows/docker.yaml
index 3a786c7..2826858 100644
--- a/.github/workflows/docker.yaml
+++ b/.github/workflows/docker.yaml
@@ -12,13 +12,19 @@ jobs:
     strategy:
       matrix:
         include:
-          - ubuntu_version: "22.04"
-            cuda_version: "12.4.1"
-            ros_distro: "humble"
-          - ubuntu_version: "24.04"
-            cuda_version: "13.1.0"
-            ros_distro: "jazzy"
+          - pytorch_version: "2.6.0"
+            cuda_version: "12.4"
+          - pytorch_version: "2.10.0"
+            cuda_version: "13.0"
     steps:
+      - name: Free disk space
+        run: |
+          sudo rm -rf /usr/share/dotnet
+          sudo rm -rf /opt/ghc
+          sudo rm -rf /usr/local/lib/android
+          sudo rm -rf /usr/local/.ghcup
+          df -h
+
       - name: Login to GitHub Container Registry
         uses: docker/login-action@v3
         with:
@@ -38,10 +44,9 @@ jobs:
           platforms: linux/amd64
           file: .docker/Dockerfile
           build-args: |
-            UBUNTU_VERSION=${{ matrix.ubuntu_version }}
+            PYTORCH_VERSION=${{ matrix.pytorch_version }}
             CUDA_VERSION=${{ matrix.cuda_version }}
-            ROS_DISTRO=${{ matrix.ros_distro }}
           push: ${{ github.ref == 'refs/heads/main' }}
           tags: |
-            roboreg:${{ matrix.ubuntu_version }}
-            ghcr.io/lbr-stack/roboreg:${{ matrix.ubuntu_version }}
+            roboreg:cuda-${{ matrix.cuda_version }}
+            ghcr.io/lbr-stack/roboreg:cuda-${{ matrix.cuda_version }}
diff --git a/README.md b/README.md
index cfe8747..a9671cb 100644
--- a/README.md
+++ b/README.md
@@ -94,21 +94,19 @@ Next:
     docker build \
         -t roboreg:latest \
         -f .docker/Dockerfile \
-        --build-arg UBUNTU_VERSION=24.04 \
-        --build-arg CUDA_VERSION=13.1.0 \
-        --build-arg ROS_DISTRO=jazzy \
+        --build-arg PYTORCH_VERSION=2.10.0 \
+        --build-arg CUDA_VERSION=13.0 \
         .
     ```
 
-    For Ubuntu 22.04 with ROS 2 Humble and CUDA 12.4.1, use:
+    For CUDA 12.4, use:
 
     ```shell
     docker build \
         -t roboreg:latest \
         -f .docker/Dockerfile \
-        --build-arg UBUNTU_VERSION=22.04 \
-        --build-arg CUDA_VERSION=12.4.1 \
-        --build-arg ROS_DISTRO=humble \
+        --build-arg PYTORCH_VERSION=2.6.0 \
+        --build-arg CUDA_VERSION=12.4 \
         .
     ```