diff --git a/docker/Dockerfile b/docker/Dockerfile
index 37e437c8..05820063 100644
--- a/docker/Dockerfile
+++ b/docker/Dockerfile
@@ -15,6 +15,11 @@ ARG ROS_VERSION="noetic"
 # ==
 ARG DEBIAN_FRONTEND=noninteractive
 
+# ==
+# Select shell
+# ==
+SHELL ["/bin/bash", "-c"]
+
 # ==
 # Install ROS
 # ==
@@ -26,6 +31,7 @@ RUN apt update \
     && apt install --no-install-recommends -y ros-${ROS_VERSION}-desktop-full \
     && apt install --no-install-recommends -y \
         python3-pip \
+        python3-venv \
         python3-rosdep \
         python3-rosinstall \
         python3-rosinstall-generator \
@@ -40,19 +46,22 @@ RUN apt update \
 # ==
 # Install python packages
 # ==
-RUN pip3 install --no-cache-dir \
-         torch==2.1.0+cu118 torchvision --extra-index-url https://download.pytorch.org/whl/cu118
-
-RUN pip3 install --no-cache-dir \
+RUN cd /root \
+    && python3 -m venv env --system-site-packages \
+    && source /root/env/bin/activate \
+    && pip install --upgrade pip \
+    && pip3 install --no-cache-dir \
+         torch==2.1.0 torchvision --extra-index-url https://download.pytorch.org/whl/cu121 \
+    && pip3 install --no-cache-dir \
          black \
          flake8 \
          jupyter \
          wget \
          numpy \
          tqdm \
-         kornia>=0.6.5 \
+         kornia \
          torchmetrics \
-         pytorch_lightning==1.6.5 \
+         pytorch_lightning \
          pytest \
          scipy \
          scikit-image \
@@ -69,9 +78,8 @@ RUN pip3 install --no-cache-dir \
          hydra-core \
          prettytable \
          termcolor \
-         pydensecrf@git+https://github.com/lucasb-eyer/pydensecrf.git#egg=pydensecrf \
-         liegroups@git+https://github.com/mmattamala/liegroups#egg=liegroups \
-         stego@git+https://github.com/leggedrobotics/stego.git#egg=stego==0.0.1 \
+         pydensecrf@git+https://github.com/lucasb-eyer/pydensecrf.git \
+         liegroups@git+https://github.com/mmattamala/liegroups \
          opencv-python>=4.6
 
 # ==
@@ -100,7 +108,11 @@ FROM base as dev
 
 ENV DEBIAN_FRONTEND=noninteractive
 
-RUN apt-get update \
+RUN mkdir -p catkin_ws/src \
+    && source /opt/ros/noetic/setup.bash \
+    && source "/root/.bashrc" \
+    && cd /root/catkin_ws && catkin build \
+    && apt-get update \
     && apt-get install -y \
         ros-noetic-jackal-simulator \
         ros-noetic-jackal-desktop \
@@ -108,4 +120,6 @@ RUN apt-get update \
         ros-noetic-rqt-robot-steering \
     && rm -rf /var/lib/apt/lists/*
 
+RUN echo "source /root/catkin_ws/devel/setup.bash" >> ~/.bashrc
+RUN echo "source /root/env/bin/activate" >> ~/.bashrc
 ENV DEBIAN_FRONTEND=dialog
\ No newline at end of file
diff --git a/docker/README.md b/docker/README.md
index 5314a640..4723f414 100644
--- a/docker/README.md
+++ b/docker/README.md
@@ -2,13 +2,30 @@
 
 ## Build and run containers
 
-To build and run the standard container (terminal-based):
+To build the container:
+
+```sh
+docker compose -f docker-compose.yaml build
+```
+
+To run the container (terminal-based):
 
 ```sh
 docker compose -f docker-compose.yaml up -d
+```
+
+To launch bash on the container:
+
+```sh
 docker compose exec wvn bash
 ```
 
+To stop the container:
+
+```sh
+docker compose -f docker-compose.yaml stop
+```
+
 To run the GUI-enabled version and check the Gazebo environment:
 
 ```sh
diff --git a/docker/docker-compose.yaml b/docker/docker-compose.yaml
index 58fc9988..56f48cc2 100644
--- a/docker/docker-compose.yaml
+++ b/docker/docker-compose.yaml
@@ -8,10 +8,15 @@ services:
     stdin_open: true
     tty: true
     network_mode: "host"
-    environment:
-     - NVIDIA_VISIBLE_DEVICES=all
-     - NVIDIA_DRIVER_CAPABILITIES=all
-    runtime: nvidia
     volumes:
       - ../../wild_visual_navigation:/root/catkin_ws/src/wild_visual_navigation
-    command: '/bin/bash'
\ No newline at end of file
+      - ../../self_supervised_segmentation:/root/catkin_ws/src/self_supervised_segmentation
+    environment:
+      - ROS_IP=127.0.0.1
+    command: "/bin/bash"
+    # command: >
+    #   /bin/bash -c "pip3 install -e /root/catkin_ws/src/wild_visual_navigation
+    #           && pip3 install -e /root/catkin_ws/src/self_supervised_segmentation
+    #           && cd /root/catkin_ws && catkin build
+    #           && cd /root
+    #           && source /root/catkin_ws/devel/setup.bash"
\ No newline at end of file
diff --git a/setup.py b/setup.py
index a2a4ea69..31556532 100644
--- a/setup.py
+++ b/setup.py
@@ -9,9 +9,9 @@
     "kornia>=0.6.5",
     "pip",
     "torchvision",
-    "torch@https://download.pytorch.org/whl/torch-2.1.0+cu121-cp38-cp38-linux_x86_64.whl",
+    "torch>=1.21",
     "torchmetrics",
-    "pytorch_lightning==1.6.5",
+    "pytorch_lightning>=1.6.5",
     "pytest",
     "scipy",
     "scikit-image",
@@ -28,9 +28,8 @@
     "hydra-core",
     "prettytable",
     "termcolor",
-    "pydensecrf@git+https://github.com/lucasb-eyer/pydensecrf.git#egg=pydensecrf",
-    "liegroups@git+https://github.com/mmattamala/liegroups#egg=liegroups",
-    "stego@git+https://github.com/leggedrobotics/stego.git#egg=stego==0.0.1",
+    "pydensecrf@git+https://github.com/lucasb-eyer/pydensecrf.git",
+    "liegroups@git+https://github.com/mmattamala/liegroups",
     "opencv-python>=4.6",
     "wget",
     "rospkg",
@@ -39,7 +38,7 @@
     name="wild_visual_navigation",
     version="0.0.1",
     author="Jonas Frey, Matias Mattamala",
-    author_email="jonfrey@ethz.ch, matias@leggedrobotics.com",
+    author_email="jonfrey@ethz.ch, matias@robots.oex.ac.uk",
     packages=find_packages(),
     python_requires=">=3.7",
     description="A small example package",