Skip to content

Commit b614e2e

Browse files
committed
Lazily import locomotion envs to prevent ModuleNotFoundError when labmaze is not installed
1 parent 780d86b commit b614e2e

File tree

1 file changed

+37
-27
lines changed

1 file changed

+37
-27
lines changed

shimmy/registration.py

+37-27
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,5 @@
11
"""Registers environments within gymnasium for optional modules."""
2+
23
from __future__ import annotations
34

45
from functools import partial
@@ -68,6 +69,41 @@ def _make_dm_control_generic_env(env, **render_kwargs):
6869
# Register all suite environments
6970
import dm_control.suite
7071

72+
def _register_locomotion_envs():
73+
try:
74+
from dm_control import composer
75+
from dm_control.locomotion.examples import basic_cmu_2019, basic_rodent_2020
76+
except ImportError:
77+
print(
78+
"Warning, registration of `dm_control` locomotion envs has failed due to an ImportError"
79+
)
80+
return
81+
82+
def _make_dm_control_example_locomotion_env(
83+
env_fn: Callable[[np.random.RandomState | None], composer.Environment],
84+
random_state: np.random.RandomState | None = None,
85+
**render_kwargs,
86+
):
87+
return DmControlCompatibilityV0(env_fn(random_state), **render_kwargs)
88+
89+
for locomotion_env, nondeterministic in (
90+
(basic_cmu_2019.cmu_humanoid_run_walls, False),
91+
(basic_cmu_2019.cmu_humanoid_run_gaps, False),
92+
(basic_cmu_2019.cmu_humanoid_go_to_target, False),
93+
(basic_cmu_2019.cmu_humanoid_maze_forage, True),
94+
(basic_cmu_2019.cmu_humanoid_heterogeneous_forage, True),
95+
(basic_rodent_2020.rodent_escape_bowl, False),
96+
(basic_rodent_2020.rodent_run_gaps, False),
97+
(basic_rodent_2020.rodent_maze_forage, True),
98+
(basic_rodent_2020.rodent_two_touch, True),
99+
# (cmu_2020_tracking.cmu_humanoid_tracking, False),
100+
):
101+
register(
102+
f"dm_control/{locomotion_env.__name__.title().replace('_', '')}-v0",
103+
partial(_make_dm_control_example_locomotion_env, env_fn=locomotion_env),
104+
nondeterministic=nondeterministic,
105+
)
106+
71107
def _make_dm_control_suite_env(
72108
domain_name: str,
73109
task_name: str,
@@ -98,33 +134,7 @@ def _make_dm_control_suite_env(
98134

99135
# Register all example locomotion environments
100136
# Listed in https://github.com/deepmind/dm_control/blob/main/dm_control/locomotion/examples/examples_test.py
101-
from dm_control import composer
102-
from dm_control.locomotion.examples import basic_cmu_2019, basic_rodent_2020
103-
104-
def _make_dm_control_example_locomotion_env(
105-
env_fn: Callable[[np.random.RandomState | None], composer.Environment],
106-
random_state: np.random.RandomState | None = None,
107-
**render_kwargs,
108-
):
109-
return DmControlCompatibilityV0(env_fn(random_state), **render_kwargs)
110-
111-
for locomotion_env, nondeterministic in (
112-
(basic_cmu_2019.cmu_humanoid_run_walls, False),
113-
(basic_cmu_2019.cmu_humanoid_run_gaps, False),
114-
(basic_cmu_2019.cmu_humanoid_go_to_target, False),
115-
(basic_cmu_2019.cmu_humanoid_maze_forage, True),
116-
(basic_cmu_2019.cmu_humanoid_heterogeneous_forage, True),
117-
(basic_rodent_2020.rodent_escape_bowl, False),
118-
(basic_rodent_2020.rodent_run_gaps, False),
119-
(basic_rodent_2020.rodent_maze_forage, True),
120-
(basic_rodent_2020.rodent_two_touch, True),
121-
# (cmu_2020_tracking.cmu_humanoid_tracking, False),
122-
):
123-
register(
124-
f"dm_control/{locomotion_env.__name__.title().replace('_', '')}-v0",
125-
partial(_make_dm_control_example_locomotion_env, env_fn=locomotion_env),
126-
nondeterministic=nondeterministic,
127-
)
137+
_register_locomotion_envs()
128138

129139
# Register all manipulation environments
130140
import dm_control.manipulation

0 commit comments

Comments
 (0)