diff --git a/Dockerfile b/Dockerfile
index 1de2bf398c7674556d5f3730b357640dd2bc578b..322eb4091518a4ac77708f01cf6feeb265612164 100644
--- a/Dockerfile
+++ b/Dockerfile
@@ -9,7 +9,9 @@ RUN mamba update mamba
 # Conda packages
 RUN mamba install -c conda-forge control
 RUN mamba install -c conda-forge casadi
-RUN mamba install -c conda-forge jax
 RUN mamba install -c conda-forge networkx
+RUN mamba install -c conda-forge jax
+RUN mamba install -c conda-forge flax
 
-# Install pip packages
\ No newline at end of file
+# Install pip packages
+RUN pip install rlax