|
1 | 1 | """Registers environments within gymnasium for optional modules."""
|
| 2 | + |
2 | 3 | from __future__ import annotations
|
3 | 4 |
|
4 | 5 | from functools import partial
|
@@ -68,6 +69,41 @@ def _make_dm_control_generic_env(env, **render_kwargs):
|
68 | 69 | # Register all suite environments
|
69 | 70 | import dm_control.suite
|
70 | 71 |
|
| 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 | + |
71 | 107 | def _make_dm_control_suite_env(
|
72 | 108 | domain_name: str,
|
73 | 109 | task_name: str,
|
@@ -98,33 +134,7 @@ def _make_dm_control_suite_env(
|
98 | 134 |
|
99 | 135 | # Register all example locomotion environments
|
100 | 136 | # 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() |
128 | 138 |
|
129 | 139 | # Register all manipulation environments
|
130 | 140 | import dm_control.manipulation
|
|
0 commit comments