diff --git a/.dockerignore b/.dockerignore
index 2ef69dc1b..909c7c5bd 100644
--- a/.dockerignore
+++ b/.dockerignore
@@ -8,7 +8,6 @@ cluster_sbatch.sh
 cluster_sbatch_mpi.sh
 cluster_torchy.sh
 logs/
-rl-trained_agents/
 .pytype/
 htmlcov/
 git_rewrite_commit_history.sh
@@ -16,3 +15,4 @@ git_rewrite_commit_history.sh
 # ignore for docker builds
 rl-trained-agents/
 .git/
+.mypy_cache/
diff --git a/CHANGELOG.md b/CHANGELOG.md
index 029e8d1b4..f799fadb9 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -1,3 +1,17 @@
+## Release 2.2.0a1 (WIP)
+
+### Breaking Changes
+
+### New Features
+
+### Bug fixes
+
+### Documentation
+
+### Other
+- Updated docker image, removed support for X server
+
+
 ## Release 2.1.0 (2023-08-17)
 
 ### Breaking Changes
@@ -18,7 +32,6 @@
 ### Other
 
 
-
 ## Release 2.0.0 (2023-06-22)
 
 **Gymnasium support**
diff --git a/docker/Dockerfile b/docker/Dockerfile
index 9e4e6d767..dfa04c23f 100644
--- a/docker/Dockerfile
+++ b/docker/Dockerfile
@@ -2,32 +2,29 @@ ARG PARENT_IMAGE
 ARG USE_GPU
 FROM $PARENT_IMAGE
 
+USER root
+
 RUN apt-get -y update \
     && apt-get -y install \
     ffmpeg \
-    freeglut3-dev \
     swig \
-    xvfb \
-    libxrandr2 \
+    # For building box2d
+    build-essential \
     && apt-get clean \
     && rm -rf /var/lib/apt/lists/*
 
-ENV CODE_DIR /root/code
-ENV VENV /root/venv
-COPY requirements.txt /tmp/
+USER $MAMBA_USER
 
+ENV CODE_DIR /home/$MAMBA_USER/code
+COPY requirements.txt /tmp/
 
 RUN \
+    eval "$(micromamba shell hook --shell bash)" && \
+    micromamba activate && \
     mkdir -p ${CODE_DIR}/rl_zoo3 && \
     pip uninstall -y stable-baselines3 && \
     pip install -r /tmp/requirements.txt && \
     pip install highway-env>=1.7.1 && \
-    rm -rf $HOME/.cache/pip
-
-ENV PATH=$VENV/bin:$PATH
-
-COPY docker/entrypoint.sh /tmp/
-RUN chmod +x /tmp/entrypoint.sh
-ENTRYPOINT ["/tmp/entrypoint.sh"]
+    pip cache purge
 
 CMD /bin/bash
diff --git a/docker/entrypoint.sh b/docker/entrypoint.sh
deleted file mode 100644
index dae73de78..000000000
--- a/docker/entrypoint.sh
+++ /dev/null
@@ -1,30 +0,0 @@
-#!/bin/bash
-# This script is the entrypoint for the Docker image.
-# Taken from https://github.com/openai/gym/
-
-set -ex
-
-# Set up display; otherwise rendering will fail
-Xvfb :1 -screen 0 1024x768x24 &
-export DISPLAY=:1
-
-# Wait for the file to come up
-display=1
-file="/tmp/.X11-unix/X$display"
-
-sleep 1
-
-for i in $(seq 1 10); do
-    if [ -e "$file" ]; then
-	     break
-    fi
-
-    echo "Waiting for $file to be created (try $i/10)"
-    sleep "$i"
-done
-if ! [ -e "$file" ]; then
-    echo "Timing out: $file was not created"
-    exit 1
-fi
-
-exec "$@"
diff --git a/rl_zoo3/version.txt b/rl_zoo3/version.txt
index 7ec1d6db4..8c6ccbaa5 100644
--- a/rl_zoo3/version.txt
+++ b/rl_zoo3/version.txt
@@ -1 +1 @@
-2.1.0
+2.2.0a1
diff --git a/scripts/build_docker.sh b/scripts/build_docker.sh
index f86c94950..b525b04bd 100755
--- a/scripts/build_docker.sh
+++ b/scripts/build_docker.sh
@@ -3,7 +3,7 @@
 PARENT=stablebaselines/stable-baselines3
 
 TAG=stablebaselines/rl-baselines3-zoo
-VERSION=1.5.1a6
+VERSION=2.2.0a1
 
 if [[ ${USE_GPU} == "True" ]]; then
   PARENT="${PARENT}:${VERSION}"
diff --git a/scripts/run_docker_cpu.sh b/scripts/run_docker_cpu.sh
index 854fe15d1..32af41c42 100755
--- a/scripts/run_docker_cpu.sh
+++ b/scripts/run_docker_cpu.sh
@@ -6,6 +6,8 @@ cmd_line="$@"
 echo "Executing in the docker (cpu image):"
 echo $cmd_line
 
-docker run -it --rm --network host --ipc=host \
- --mount src=$(pwd),target=/root/code/rl_zoo3,type=bind stablebaselines/rl-baselines3-zoo-cpu:latest\
-  bash -c "cd /root/code/rl_zoo3/ && $cmd_line"
+# Note: --user=root is needed, as the current user id/group id will be used
+# to mount the log directory (and $MAMBAUSER is not root)
+docker run -it --user=root --rm --network host --ipc=host \
+ --mount src=$(pwd),target=/home/mambauser/code/rl_zoo3,type=bind stablebaselines/rl-baselines3-zoo-cpu:latest\
+  bash -c "cd /home/mambauser/code/rl_zoo3/ && $cmd_line"
diff --git a/scripts/run_docker_gpu.sh b/scripts/run_docker_gpu.sh
index 9bcb31f42..399514f5e 100755
--- a/scripts/run_docker_gpu.sh
+++ b/scripts/run_docker_gpu.sh
@@ -6,6 +6,8 @@ cmd_line="$@"
 echo "Executing in the docker (gpu image):"
 echo $cmd_line
 
-docker run -it --runtime=nvidia --rm --network host --ipc=host \
-  --mount src=$(pwd),target=/root/code/rl_zoo3,type=bind stablebaselines/rl-baselines3-zoo:latest\
-  bash -c "cd /root/code/rl_zoo3/ && $cmd_line"
+# Note: --user=root is needed, as the current user id/group id will be used
+# to mount the log directory (and $MAMBAUSER is not root)
+docker run -it --user=root --runtime=nvidia --rm --network host --ipc=host \
+  --mount src=$(pwd),target=/home/mambauser/code/rl_zoo3,type=bind stablebaselines/rl-baselines3-zoo:latest\
+  bash -c "cd /home/mambauser/code/rl_zoo3/ && $cmd_line"