From d9019d9e7eae22b3b250ae445f35cae458c82464 Mon Sep 17 00:00:00 2001
From: Simon Alibert <alibert.sim@gmail.com>
Date: Tue, 9 Apr 2024 10:24:28 +0200
Subject: [PATCH 1/8] disable env_checker in factory

---
 lerobot/common/envs/factory.py | 7 +++++--
 1 file changed, 5 insertions(+), 2 deletions(-)

diff --git a/lerobot/common/envs/factory.py b/lerobot/common/envs/factory.py
index c8d10851..d5571935 100644
--- a/lerobot/common/envs/factory.py
+++ b/lerobot/common/envs/factory.py
@@ -30,10 +30,13 @@ def make_env(cfg, num_parallel_envs=0) -> gym.Env | gym.vector.SyncVectorEnv:
 
     if num_parallel_envs == 0:
         # non-batched version of the env that returns an observation of shape (c)
-        env = gym.make(gym_handle, **kwargs)
+        env = gym.make(gym_handle, disable_env_checker=True, **kwargs)
     else:
         # batched version of the env that returns an observation of shape (b, c)
         env = gym.vector.SyncVectorEnv(
-            [lambda: gym.make(gym_handle, **kwargs) for _ in range(num_parallel_envs)]
+            [
+                lambda: gym.make(gym_handle, disable_env_checker=True, **kwargs)
+                for _ in range(num_parallel_envs)
+            ]
         )
     return env

From 274f20b49d018251e1414f9dab98c59ce5a2d23b Mon Sep 17 00:00:00 2001
From: Simon Alibert <alibert.sim@gmail.com>
Date: Tue, 9 Apr 2024 10:25:41 +0200
Subject: [PATCH 2/8] Update gym-pusht

---
 poetry.lock | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/poetry.lock b/poetry.lock
index 95c9f31e..f712289e 100644
--- a/poetry.lock
+++ b/poetry.lock
@@ -920,7 +920,7 @@ shapely = "^2.0.3"
 type = "git"
 url = "git@github.com:huggingface/gym-pusht.git"
 reference = "HEAD"
-resolved_reference = "6c9893504f670ff069d0f759a733e971ea1efdbf"
+resolved_reference = "824b22832cc8d71a4b4e96a57563510cf47e30c1"
 
 [[package]]
 name = "gym-xarm"

From 2573e89e1df6136142e883ec23cd066e3a75c657 Mon Sep 17 00:00:00 2001
From: Simon Alibert <alibert.sim@gmail.com>
Date: Tue, 9 Apr 2024 10:38:08 +0200
Subject: [PATCH 3/8] Remove direct dependencies

---
 .github/poetry/cpu/poetry.lock    | 51 ++++++++++++++-----------------
 .github/poetry/cpu/pyproject.toml | 14 ++++-----
 poetry.lock                       | 36 +++++++++++-----------
 pyproject.toml                    |  7 +----
 4 files changed, 49 insertions(+), 59 deletions(-)

diff --git a/.github/poetry/cpu/poetry.lock b/.github/poetry/cpu/poetry.lock
index 15b27c76..edc1d503 100644
--- a/.github/poetry/cpu/poetry.lock
+++ b/.github/poetry/cpu/poetry.lock
@@ -1,4 +1,4 @@
-# This file is automatically @generated by Poetry 1.8.1 and should not be changed by hand.
+# This file is automatically @generated by Poetry 1.8.2 and should not be changed by hand.
 
 [[package]]
 name = "absl-py"
@@ -517,21 +517,11 @@ files = [
     {file = "distlib-0.3.8.tar.gz", hash = "sha256:1530ea13e350031b6312d8580ddb6b27a104275a31106523b8f123787f494f64"},
 ]
 
-[[package]]
-name = "dm"
-version = "1.3"
-description = "Dict to Data mapper"
-optional = false
-python-versions = "*"
-files = [
-    {file = "dm-1.3.tar.gz", hash = "sha256:ce77537bf346b5d8c0dc0b5d679cfc4a946faadcd5315e6c80ef6f3af824130d"},
-]
-
 [[package]]
 name = "dm-control"
 version = "1.0.14"
 description = "Continuous control environments and MuJoCo Python bindings."
-optional = false
+optional = true
 python-versions = ">=3.8"
 files = [
     {file = "dm_control-1.0.14-py3-none-any.whl", hash = "sha256:883c63244a7ebf598700a97564ed19fffd3479ca79efd090aed881609cdb9fc6"},
@@ -562,7 +552,7 @@ hdf5 = ["h5py"]
 name = "dm-env"
 version = "1.6"
 description = "A Python interface for Reinforcement Learning environments."
-optional = false
+optional = true
 python-versions = ">=3.7"
 files = [
     {file = "dm-env-1.6.tar.gz", hash = "sha256:a436eb1c654c39e0c986a516cee218bea7140b510fceff63f97eb4fcff3d93de"},
@@ -578,7 +568,7 @@ numpy = "*"
 name = "dm-tree"
 version = "0.1.8"
 description = "Tree is a library for working with nested data structures."
-optional = false
+optional = true
 python-versions = "*"
 files = [
     {file = "dm-tree-0.1.8.tar.gz", hash = "sha256:0fcaabbb14e7980377439e7140bd05552739ca5e515ecb3119f234acee4b9430"},
@@ -806,7 +796,7 @@ test = ["black", "coverage[toml]", "ddt (>=1.1.1,!=1.4.3)", "mock", "mypy", "pre
 name = "glfw"
 version = "2.7.0"
 description = "A ctypes-based wrapper for GLFW3."
-optional = false
+optional = true
 python-versions = "*"
 files = [
     {file = "glfw-2.7.0-py2.py27.py3.py30.py31.py32.py33.py34.py35.py36.py37.py38-none-macosx_10_6_intel.whl", hash = "sha256:bd82849edcceda4e262bd1227afaa74b94f9f0731c1197863cd25c15bfc613fc"},
@@ -986,7 +976,7 @@ toy-text = ["pygame (>=2.1.3)", "pygame (>=2.1.3)"]
 name = "gymnasium-robotics"
 version = "1.2.4"
 description = "Robotics environments for the Gymnasium repo."
-optional = false
+optional = true
 python-versions = ">=3.8"
 files = [
     {file = "gymnasium-robotics-1.2.4.tar.gz", hash = "sha256:d304192b066f8b800599dfbe3d9d90bba9b761ee884472bdc4d05968a8bc61cb"},
@@ -1218,7 +1208,7 @@ i18n = ["Babel (>=2.7)"]
 name = "labmaze"
 version = "1.0.6"
 description = "LabMaze: DeepMind Lab's text maze generator."
-optional = false
+optional = true
 python-versions = "*"
 files = [
     {file = "labmaze-1.0.6-cp310-cp310-macosx_10_12_x86_64.whl", hash = "sha256:b2ddef976dfd8d992b19cfa6c633f2eba7576d759c2082da534e3f727479a84a"},
@@ -1262,7 +1252,7 @@ setuptools = "!=50.0.0"
 name = "lazy-loader"
 version = "0.3"
 description = "lazy_loader"
-optional = false
+optional = true
 python-versions = ">=3.7"
 files = [
     {file = "lazy_loader-0.3-py3-none-any.whl", hash = "sha256:1e9e76ee8631e264c62ce10006718e80b2cfc74340d17d1031e0f84af7478554"},
@@ -1307,7 +1297,7 @@ files = [
 name = "lxml"
 version = "5.1.0"
 description = "Powerful and Pythonic XML processing library combining libxml2/libxslt with the ElementTree API."
-optional = false
+optional = true
 python-versions = ">=3.6"
 files = [
     {file = "lxml-5.1.0-cp310-cp310-macosx_10_9_universal2.whl", hash = "sha256:704f5572ff473a5f897745abebc6df40f22d4133c1e0a1f124e4f2bd3330ff7e"},
@@ -1525,7 +1515,7 @@ tests = ["pytest (>=4.6)"]
 name = "mujoco"
 version = "2.3.7"
 description = "MuJoCo Physics Simulator"
-optional = false
+optional = true
 python-versions = ">=3.8"
 files = [
     {file = "mujoco-2.3.7-cp310-cp310-macosx_10_16_x86_64.whl", hash = "sha256:e8714a5ff6a1561b364b7b4648d4c0c8d13e751874cf7401c309b9d23fa9598b"},
@@ -1839,7 +1829,7 @@ xml = ["lxml (>=4.9.2)"]
 name = "pettingzoo"
 version = "1.24.3"
 description = "Gymnasium for multi-agent reinforcement learning."
-optional = false
+optional = true
 python-versions = ">=3.8"
 files = [
     {file = "pettingzoo-1.24.3-py3-none-any.whl", hash = "sha256:23ed90517d2e8a7098bdaf5e31234b3a7f7b73ca578d70d1ca7b9d0cb0e37982"},
@@ -2207,7 +2197,7 @@ dev = ["aafigure", "matplotlib", "pygame", "pyglet (<2.0.0)", "sphinx", "wheel"]
 name = "pyopengl"
 version = "3.1.7"
 description = "Standard OpenGL bindings for Python"
-optional = false
+optional = true
 python-versions = "*"
 files = [
     {file = "PyOpenGL-3.1.7-py3-none-any.whl", hash = "sha256:a6ab19cf290df6101aaf7470843a9c46207789855746399d0af92521a0a92b7a"},
@@ -2218,7 +2208,7 @@ files = [
 name = "pyparsing"
 version = "3.1.2"
 description = "pyparsing module - Classes and methods to define and execute parsing grammars"
-optional = false
+optional = true
 python-versions = ">=3.6.8"
 files = [
     {file = "pyparsing-3.1.2-py3-none-any.whl", hash = "sha256:f9db75911801ed778fe61bb643079ff86601aca99fcae6345aa67292038fb742"},
@@ -2649,7 +2639,7 @@ torch = ["safetensors[numpy]", "torch (>=1.10)"]
 name = "scikit-image"
 version = "0.22.0"
 description = "Image processing in Python"
-optional = false
+optional = true
 python-versions = ">=3.9"
 files = [
     {file = "scikit_image-0.22.0-cp310-cp310-macosx_10_9_x86_64.whl", hash = "sha256:74ec5c1d4693506842cc7c9487c89d8fc32aed064e9363def7af08b8f8cbb31d"},
@@ -2697,7 +2687,7 @@ test = ["asv", "matplotlib (>=3.5)", "numpydoc (>=1.5)", "pooch (>=1.6.0)", "pyt
 name = "scipy"
 version = "1.12.0"
 description = "Fundamental algorithms for scientific computing in Python"
-optional = false
+optional = true
 python-versions = ">=3.9"
 files = [
     {file = "scipy-1.12.0-cp310-cp310-macosx_10_9_x86_64.whl", hash = "sha256:78e4402e140879387187f7f25d91cc592b3501a2e51dfb320f48dfb73565f10b"},
@@ -2902,7 +2892,7 @@ testing-integration = ["build[virtualenv] (>=1.0.3)", "filelock (>=3.4.0)", "jar
 name = "shapely"
 version = "2.0.3"
 description = "Manipulation and analysis of geometric objects"
-optional = false
+optional = true
 python-versions = ">=3.7"
 files = [
     {file = "shapely-2.0.3-cp310-cp310-macosx_10_9_universal2.whl", hash = "sha256:af7e9abe180b189431b0f490638281b43b84a33a960620e6b2e8d3e3458b61a1"},
@@ -3069,7 +3059,7 @@ tests = ["pytest", "pytest-cov"]
 name = "tifffile"
 version = "2024.2.12"
 description = "Read and write TIFF files"
-optional = false
+optional = true
 python-versions = ">=3.9"
 files = [
     {file = "tifffile-2024.2.12-py3-none-any.whl", hash = "sha256:870998f82fbc94ff7c3528884c1b0ae54863504ff51dbebea431ac3fa8fb7c21"},
@@ -3331,7 +3321,12 @@ files = [
 docs = ["furo", "jaraco.packaging (>=9.3)", "jaraco.tidelift (>=1.4)", "rst.linker (>=1.9)", "sphinx (>=3.5)", "sphinx-lint"]
 testing = ["big-O", "jaraco.functools", "jaraco.itertools", "more-itertools", "pytest (>=6)", "pytest-checkdocs (>=2.4)", "pytest-cov", "pytest-enabler (>=2.2)", "pytest-ignore-flaky", "pytest-mypy", "pytest-ruff (>=0.2.1)"]
 
+[extras]
+aloha = ["gym-aloha"]
+pusht = ["gym-pusht"]
+xarm = ["gym-xarm"]
+
 [metadata]
 lock-version = "2.0"
 python-versions = "^3.10"
-content-hash = "32cd6caa01276a90b37cb177204e5b1511e92838f3f0268391034042d56f3bd6"
+content-hash = "ba2b64b1a683450b097a1ccbed3410cc3bee97ba8b41f409c5e379f95d8b1c6f"
diff --git a/.github/poetry/cpu/pyproject.toml b/.github/poetry/cpu/pyproject.toml
index d310da47..741e3b37 100644
--- a/.github/poetry/cpu/pyproject.toml
+++ b/.github/poetry/cpu/pyproject.toml
@@ -23,7 +23,6 @@ packages = [{include = "lerobot"}]
 python = "^3.10"
 termcolor = "^2.4.0"
 omegaconf = "^2.3.0"
-dm-env = "^1.6"
 pandas = "^2.2.1"
 wandb = "^0.16.3"
 moviepy = "^1.0.3"
@@ -34,21 +33,15 @@ einops = "^0.7.0"
 pygame = "^2.5.2"
 pymunk = "^6.6.0"
 zarr = "^2.17.0"
-shapely = "^2.0.3"
-scikit-image = "^0.22.0"
 numba = "^0.59.0"
 mpmath = "^1.3.0"
 torch = {version = "^2.2.1", source = "torch-cpu"}
-mujoco = "^2.3.7"
 opencv-python = "^4.9.0.80"
 diffusers = "^0.26.3"
 torchvision = {version = "^0.17.1", source = "torch-cpu"}
 h5py = "^3.10.0"
-dm = "^1.3"
-dm-control = "1.0.14"
 robomimic = "0.2.0"
 huggingface-hub = "^0.21.4"
-gymnasium-robotics = "^1.2.4"
 gymnasium = "^0.29.1"
 cmake = "^3.29.0.1"
 gym-pusht = { git = "git@github.com:huggingface/gym-pusht.git", optional = true}
@@ -58,6 +51,13 @@ gym-aloha = { git = "git@github.com:huggingface/gym-aloha.git", optional = true}
 # gym-xarm = { path = "../gym-xarm", develop = true, optional = true}
 # gym-aloha = { path = "../gym-aloha", develop = true, optional = true}
 
+
+[tool.poetry.extras]
+pusht = ["gym-pusht"]
+xarm = ["gym-xarm"]
+aloha = ["gym-aloha"]
+
+
 [tool.poetry.group.dev.dependencies]
 pre-commit = "^3.6.2"
 debugpy = "^1.8.1"
diff --git a/poetry.lock b/poetry.lock
index f712289e..b5e97cb7 100644
--- a/poetry.lock
+++ b/poetry.lock
@@ -1,4 +1,4 @@
-# This file is automatically @generated by Poetry 1.8.1 and should not be changed by hand.
+# This file is automatically @generated by Poetry 1.8.2 and should not be changed by hand.
 
 [[package]]
 name = "absl-py"
@@ -521,7 +521,7 @@ files = [
 name = "dm-control"
 version = "1.0.14"
 description = "Continuous control environments and MuJoCo Python bindings."
-optional = false
+optional = true
 python-versions = ">=3.8"
 files = [
     {file = "dm_control-1.0.14-py3-none-any.whl", hash = "sha256:883c63244a7ebf598700a97564ed19fffd3479ca79efd090aed881609cdb9fc6"},
@@ -552,7 +552,7 @@ hdf5 = ["h5py"]
 name = "dm-env"
 version = "1.6"
 description = "A Python interface for Reinforcement Learning environments."
-optional = false
+optional = true
 python-versions = ">=3.7"
 files = [
     {file = "dm-env-1.6.tar.gz", hash = "sha256:a436eb1c654c39e0c986a516cee218bea7140b510fceff63f97eb4fcff3d93de"},
@@ -568,7 +568,7 @@ numpy = "*"
 name = "dm-tree"
 version = "0.1.8"
 description = "Tree is a library for working with nested data structures."
-optional = false
+optional = true
 python-versions = "*"
 files = [
     {file = "dm-tree-0.1.8.tar.gz", hash = "sha256:0fcaabbb14e7980377439e7140bd05552739ca5e515ecb3119f234acee4b9430"},
@@ -796,7 +796,7 @@ test = ["black", "coverage[toml]", "ddt (>=1.1.1,!=1.4.3)", "mock", "mypy", "pre
 name = "glfw"
 version = "2.7.0"
 description = "A ctypes-based wrapper for GLFW3."
-optional = false
+optional = true
 python-versions = "*"
 files = [
     {file = "glfw-2.7.0-py2.py27.py3.py30.py31.py32.py33.py34.py35.py36.py37.py38-none-macosx_10_6_intel.whl", hash = "sha256:bd82849edcceda4e262bd1227afaa74b94f9f0731c1197863cd25c15bfc613fc"},
@@ -976,7 +976,7 @@ toy-text = ["pygame (>=2.1.3)", "pygame (>=2.1.3)"]
 name = "gymnasium-robotics"
 version = "1.2.4"
 description = "Robotics environments for the Gymnasium repo."
-optional = false
+optional = true
 python-versions = ">=3.8"
 files = [
     {file = "gymnasium-robotics-1.2.4.tar.gz", hash = "sha256:d304192b066f8b800599dfbe3d9d90bba9b761ee884472bdc4d05968a8bc61cb"},
@@ -1281,7 +1281,7 @@ i18n = ["Babel (>=2.7)"]
 name = "labmaze"
 version = "1.0.6"
 description = "LabMaze: DeepMind Lab's text maze generator."
-optional = false
+optional = true
 python-versions = "*"
 files = [
     {file = "labmaze-1.0.6-cp310-cp310-macosx_10_12_x86_64.whl", hash = "sha256:b2ddef976dfd8d992b19cfa6c633f2eba7576d759c2082da534e3f727479a84a"},
@@ -1325,7 +1325,7 @@ setuptools = "!=50.0.0"
 name = "lazy-loader"
 version = "0.3"
 description = "lazy_loader"
-optional = false
+optional = true
 python-versions = ">=3.7"
 files = [
     {file = "lazy_loader-0.3-py3-none-any.whl", hash = "sha256:1e9e76ee8631e264c62ce10006718e80b2cfc74340d17d1031e0f84af7478554"},
@@ -1370,7 +1370,7 @@ files = [
 name = "lxml"
 version = "5.1.0"
 description = "Powerful and Pythonic XML processing library combining libxml2/libxslt with the ElementTree API."
-optional = false
+optional = true
 python-versions = ">=3.6"
 files = [
     {file = "lxml-5.1.0-cp310-cp310-macosx_10_9_universal2.whl", hash = "sha256:704f5572ff473a5f897745abebc6df40f22d4133c1e0a1f124e4f2bd3330ff7e"},
@@ -1588,7 +1588,7 @@ tests = ["pytest (>=4.6)"]
 name = "mujoco"
 version = "2.3.7"
 description = "MuJoCo Physics Simulator"
-optional = false
+optional = true
 python-versions = ">=3.8"
 files = [
     {file = "mujoco-2.3.7-cp310-cp310-macosx_10_16_x86_64.whl", hash = "sha256:e8714a5ff6a1561b364b7b4648d4c0c8d13e751874cf7401c309b9d23fa9598b"},
@@ -2043,7 +2043,7 @@ xml = ["lxml (>=4.9.2)"]
 name = "pettingzoo"
 version = "1.24.3"
 description = "Gymnasium for multi-agent reinforcement learning."
-optional = false
+optional = true
 python-versions = ">=3.8"
 files = [
     {file = "pettingzoo-1.24.3-py3-none-any.whl", hash = "sha256:23ed90517d2e8a7098bdaf5e31234b3a7f7b73ca578d70d1ca7b9d0cb0e37982"},
@@ -2411,7 +2411,7 @@ dev = ["aafigure", "matplotlib", "pygame", "pyglet (<2.0.0)", "sphinx", "wheel"]
 name = "pyopengl"
 version = "3.1.7"
 description = "Standard OpenGL bindings for Python"
-optional = false
+optional = true
 python-versions = "*"
 files = [
     {file = "PyOpenGL-3.1.7-py3-none-any.whl", hash = "sha256:a6ab19cf290df6101aaf7470843a9c46207789855746399d0af92521a0a92b7a"},
@@ -2422,7 +2422,7 @@ files = [
 name = "pyparsing"
 version = "3.1.2"
 description = "pyparsing module - Classes and methods to define and execute parsing grammars"
-optional = false
+optional = true
 python-versions = ">=3.6.8"
 files = [
     {file = "pyparsing-3.1.2-py3-none-any.whl", hash = "sha256:f9db75911801ed778fe61bb643079ff86601aca99fcae6345aa67292038fb742"},
@@ -2853,7 +2853,7 @@ torch = ["safetensors[numpy]", "torch (>=1.10)"]
 name = "scikit-image"
 version = "0.22.0"
 description = "Image processing in Python"
-optional = false
+optional = true
 python-versions = ">=3.9"
 files = [
     {file = "scikit_image-0.22.0-cp310-cp310-macosx_10_9_x86_64.whl", hash = "sha256:74ec5c1d4693506842cc7c9487c89d8fc32aed064e9363def7af08b8f8cbb31d"},
@@ -2901,7 +2901,7 @@ test = ["asv", "matplotlib (>=3.5)", "numpydoc (>=1.5)", "pooch (>=1.6.0)", "pyt
 name = "scipy"
 version = "1.12.0"
 description = "Fundamental algorithms for scientific computing in Python"
-optional = false
+optional = true
 python-versions = ">=3.9"
 files = [
     {file = "scipy-1.12.0-cp310-cp310-macosx_10_9_x86_64.whl", hash = "sha256:78e4402e140879387187f7f25d91cc592b3501a2e51dfb320f48dfb73565f10b"},
@@ -3106,7 +3106,7 @@ testing-integration = ["build[virtualenv] (>=1.0.3)", "filelock (>=3.4.0)", "jar
 name = "shapely"
 version = "2.0.3"
 description = "Manipulation and analysis of geometric objects"
-optional = false
+optional = true
 python-versions = ">=3.7"
 files = [
     {file = "shapely-2.0.3-cp310-cp310-macosx_10_9_universal2.whl", hash = "sha256:af7e9abe180b189431b0f490638281b43b84a33a960620e6b2e8d3e3458b61a1"},
@@ -3273,7 +3273,7 @@ tests = ["pytest", "pytest-cov"]
 name = "tifffile"
 version = "2024.2.12"
 description = "Read and write TIFF files"
-optional = false
+optional = true
 python-versions = ">=3.9"
 files = [
     {file = "tifffile-2024.2.12-py3-none-any.whl", hash = "sha256:870998f82fbc94ff7c3528884c1b0ae54863504ff51dbebea431ac3fa8fb7c21"},
@@ -3598,4 +3598,4 @@ xarm = ["gym-xarm"]
 [metadata]
 lock-version = "2.0"
 python-versions = "^3.10"
-content-hash = "bf4627c62a45764931729ce373f1038fe289b6caebb01e66d878f6f278c54518"
+content-hash = "d444fab7fed5e3c5c9cde69c8f19a286126615ab4a9de11c23730b5286cac77b"
diff --git a/pyproject.toml b/pyproject.toml
index a549e66f..75342c80 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -23,7 +23,6 @@ packages = [{include = "lerobot"}]
 python = "^3.10"
 termcolor = "^2.4.0"
 omegaconf = "^2.3.0"
-dm-env = "^1.6"
 pandas = "^2.2.1"
 wandb = "^0.16.3"
 moviepy = "^1.0.3"
@@ -34,20 +33,15 @@ einops = "^0.7.0"
 pygame = "^2.5.2"
 pymunk = "^6.6.0"
 zarr = "^2.17.0"
-shapely = "^2.0.3"
-scikit-image = "^0.22.0"
 numba = "^0.59.0"
 mpmath = "^1.3.0"
 torch = "^2.2.1"
-mujoco = "^2.3.7"
 opencv-python = "^4.9.0.80"
 diffusers = "^0.26.3"
 torchvision = "^0.17.1"
 h5py = "^3.10.0"
-dm-control = "1.0.14"
 huggingface-hub = {extras = ["hf-transfer"], version = "^0.21.4"}
 robomimic = "0.2.0"
-gymnasium-robotics = "^1.2.4"
 gymnasium = "^0.29.1"
 cmake = "^3.29.0.1"
 gym-pusht = { git = "git@github.com:huggingface/gym-pusht.git", optional = true}
@@ -62,6 +56,7 @@ pusht = ["gym-pusht"]
 xarm = ["gym-xarm"]
 aloha = ["gym-aloha"]
 
+
 [tool.poetry.group.dev.dependencies]
 pre-commit = "^3.6.2"
 debugpy = "^1.8.1"

From dfaacbcf5a7bf4d75a39d6ad8bac8a75291e8cb5 Mon Sep 17 00:00:00 2001
From: Simon Alibert <alibert.sim@gmail.com>
Date: Tue, 9 Apr 2024 10:40:11 +0200
Subject: [PATCH 4/8] Split dev/test dependencies

---
 .github/poetry/cpu/poetry.lock    | 2 +-
 .github/poetry/cpu/pyproject.toml | 7 +++++++
 poetry.lock                       | 2 +-
 pyproject.toml                    | 7 +++++++
 4 files changed, 16 insertions(+), 2 deletions(-)

diff --git a/.github/poetry/cpu/poetry.lock b/.github/poetry/cpu/poetry.lock
index edc1d503..fe4ed7a0 100644
--- a/.github/poetry/cpu/poetry.lock
+++ b/.github/poetry/cpu/poetry.lock
@@ -3329,4 +3329,4 @@ xarm = ["gym-xarm"]
 [metadata]
 lock-version = "2.0"
 python-versions = "^3.10"
-content-hash = "ba2b64b1a683450b097a1ccbed3410cc3bee97ba8b41f409c5e379f95d8b1c6f"
+content-hash = "8fa6dfc30e605741c24f5de58b89125d5b02153f550e5af7a44356956d6bb167"
diff --git a/.github/poetry/cpu/pyproject.toml b/.github/poetry/cpu/pyproject.toml
index 741e3b37..f5c439dc 100644
--- a/.github/poetry/cpu/pyproject.toml
+++ b/.github/poetry/cpu/pyproject.toml
@@ -58,9 +58,16 @@ xarm = ["gym-xarm"]
 aloha = ["gym-aloha"]
 
 
+[tool.poetry.group.dev]
+optional = true
+
+
 [tool.poetry.group.dev.dependencies]
 pre-commit = "^3.6.2"
 debugpy = "^1.8.1"
+
+
+[tool.poetry.group.test.dependencies]
 pytest = "^8.1.0"
 pytest-cov = "^5.0.0"
 
diff --git a/poetry.lock b/poetry.lock
index b5e97cb7..387366b8 100644
--- a/poetry.lock
+++ b/poetry.lock
@@ -3598,4 +3598,4 @@ xarm = ["gym-xarm"]
 [metadata]
 lock-version = "2.0"
 python-versions = "^3.10"
-content-hash = "d444fab7fed5e3c5c9cde69c8f19a286126615ab4a9de11c23730b5286cac77b"
+content-hash = "7ec0310f8dd0ffa4d92fa78e06513bce98c3657692b3753ff34aadd297a3766c"
diff --git a/pyproject.toml b/pyproject.toml
index 75342c80..a0fc7d44 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -57,9 +57,16 @@ xarm = ["gym-xarm"]
 aloha = ["gym-aloha"]
 
 
+[tool.poetry.group.dev]
+optional = true
+
+
 [tool.poetry.group.dev.dependencies]
 pre-commit = "^3.6.2"
 debugpy = "^1.8.1"
+
+
+[tool.poetry.group.test.dependencies]
 pytest = "^8.1.0"
 pytest-cov = "^5.0.0"
 

From d21543eb4fcced5bcf717dca62155d53a4c2bc87 Mon Sep 17 00:00:00 2001
From: Simon Alibert <alibert.sim@gmail.com>
Date: Tue, 9 Apr 2024 10:41:20 +0200
Subject: [PATCH 5/8] Add env.close()

---
 tests/test_envs.py | 4 +++-
 1 file changed, 3 insertions(+), 1 deletion(-)

diff --git a/tests/test_envs.py b/tests/test_envs.py
index 72bc93c4..c49461a0 100644
--- a/tests/test_envs.py
+++ b/tests/test_envs.py
@@ -34,7 +34,7 @@ def test_env(env_name, task, obs_type):
     importlib.import_module(package_name)
     env = gym.make(f"{package_name}/{task}", obs_type=obs_type)
     check_env(env.unwrapped)
-
+    env.close()
 
 @pytest.mark.parametrize(
     "env_name",
@@ -61,3 +61,5 @@ def test_factory(env_name):
         # TODO(rcadene): we assume for now that image normalization takes place in the model
         assert img.max() <= 1.0
         assert img.min() >= 0.0
+
+    env.close()

From dba037508997dae6fe3ba16a7886714ca1759bef Mon Sep 17 00:00:00 2001
From: Simon Alibert <alibert.sim@gmail.com>
Date: Tue, 9 Apr 2024 10:45:58 +0200
Subject: [PATCH 6/8] Fix CI

---
 .github/workflows/test.yml | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml
index c1b14780..34d76827 100644
--- a/.github/workflows/test.yml
+++ b/.github/workflows/test.yml
@@ -87,7 +87,7 @@ jobs:
           TMP: ~/tmp
         run: |
           mkdir ~/tmp
-          poetry install --no-interaction --no-root
+          poetry install --no-interaction --no-root --all-extras
 
       - name: Save cached venv
         if: |
@@ -106,7 +106,7 @@ jobs:
       #             install project
       #----------------------------------------------
       - name: Install project
-        run: poetry install --no-interaction
+        run: poetry install --no-interaction --all-extras
 
       #----------------------------------------------
       #            run tests & coverage

From d44950e020c7e658fa6da19063c40379801080e8 Mon Sep 17 00:00:00 2001
From: Simon Alibert <alibert.sim@gmail.com>
Date: Tue, 9 Apr 2024 11:44:55 +0200
Subject: [PATCH 7/8] Add ssh key

---
 .github/workflows/test.yml | 5 +++++
 1 file changed, 5 insertions(+)

diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml
index 34d76827..afdcc41f 100644
--- a/.github/workflows/test.yml
+++ b/.github/workflows/test.yml
@@ -34,6 +34,11 @@ jobs:
         with:
           python-version: '3.10'
 
+      - name: Add SSH key for installing envs
+        uses: webfactory/ssh-agent@v0.9.0
+        with:
+          ssh-private-key: ${{ secrets.SSH_PRIVATE_KEY }}
+
       #----------------------------------------------
       #         install & configure poetry
       #----------------------------------------------

From 7f4ff0b170091288bae65281df093768de562f13 Mon Sep 17 00:00:00 2001
From: Simon Alibert <alibert.sim@gmail.com>
Date: Tue, 9 Apr 2024 11:58:59 +0200
Subject: [PATCH 8/8] CI fix attempt

---
 tests/test_envs.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tests/test_envs.py b/tests/test_envs.py
index c49461a0..d25231b0 100644
--- a/tests/test_envs.py
+++ b/tests/test_envs.py
@@ -33,7 +33,7 @@ def test_env(env_name, task, obs_type):
     package_name = f"gym_{env_name}"
     importlib.import_module(package_name)
     env = gym.make(f"{package_name}/{task}", obs_type=obs_type)
-    check_env(env.unwrapped)
+    check_env(env.unwrapped, skip_render_check=True)
     env.close()
 
 @pytest.mark.parametrize(