From 20b54b303a06803ca8739ec355df91a8fbb22b30 Mon Sep 17 00:00:00 2001
From: vmoens <vincentmoens@gmail.com>
Date: Thu, 21 Sep 2023 12:46:58 -0400
Subject: [PATCH] init

---
 robohive/__init__.py | 89 +++++++++++++++++++++++---------------------
 1 file changed, 46 insertions(+), 43 deletions(-)

diff --git a/robohive/__init__.py b/robohive/__init__.py
index 5794f0f7..7ae7da13 100644
--- a/robohive/__init__.py
+++ b/robohive/__init__.py
@@ -10,55 +10,58 @@
 # RoboHive version
 __version__ = "0.6.0"
 
+def _robohive_env_suite():
+    # Register RoboHive Envs
+    _current_gym_envs = gym.envs.registration.registry.env_specs.keys()
+    _current_gym_envs = set(_current_gym_envs)
+    robohive_env_suite = set()
 
-# Register RoboHive Envs
-_current_gym_envs = gym.envs.registration.registry.env_specs.keys()
-_current_gym_envs = set(_current_gym_envs)
-robohive_env_suite = set()
+    # Register Arms Suite
+    import robohive.envs.arms # noqa
+    robohive_arm_suite = set(gym.envs.registration.registry.env_specs.keys())-robohive_env_suite-_current_gym_envs
+    robohive_arm_suite = set(sorted(robohive_arm_suite))
+    robohive_env_suite = robohive_env_suite | robohive_arm_suite
 
-# Register Arms Suite
-import robohive.envs.arms # noqa
-robohive_arm_suite = set(gym.envs.registration.registry.env_specs.keys())-robohive_env_suite-_current_gym_envs
-robohive_arm_suite = set(sorted(robohive_arm_suite))
-robohive_env_suite = robohive_env_suite | robohive_arm_suite
+    # Register Myo Suite
+    import robohive.envs.myo # noqa
+    import robohive.envs.myo.myochallenge # noqa
+    robohive_myo_suite = set(gym.envs.registration.registry.env_specs.keys())-robohive_env_suite-_current_gym_envs
+    robohive_env_suite  = robohive_env_suite | robohive_myo_suite
+    robohive_myo_suite = sorted(robohive_myo_suite)
 
-# Register Myo Suite
-import robohive.envs.myo # noqa
-import robohive.envs.myo.myochallenge # noqa
-robohive_myo_suite = set(gym.envs.registration.registry.env_specs.keys())-robohive_env_suite-_current_gym_envs
-robohive_env_suite  = robohive_env_suite | robohive_myo_suite
-robohive_myo_suite = sorted(robohive_myo_suite)
+    # Register FM suite
+    import robohive.envs.fm # noqa
+    robohive_fm_suite = set(gym.envs.registration.registry.env_specs.keys())-robohive_env_suite-_current_gym_envs
+    robohive_env_suite  = robohive_env_suite | robohive_fm_suite
+    robohive_fm_suite = sorted(robohive_fm_suite)
 
-# Register FM suite
-import robohive.envs.fm # noqa
-robohive_fm_suite = set(gym.envs.registration.registry.env_specs.keys())-robohive_env_suite-_current_gym_envs
-robohive_env_suite  = robohive_env_suite | robohive_fm_suite
-robohive_fm_suite = sorted(robohive_fm_suite)
+    # Register Hands Suite
+    import robohive.envs.hands # noqa
+    # import robohive.envs.tcdm # noqa # WIP
+    robohive_hand_suite = set(gym.envs.registration.registry.env_specs.keys())-robohive_env_suite-_current_gym_envs
+    robohive_env_suite  = robohive_env_suite | robohive_hand_suite
+    robohive_hand_suite = sorted(robohive_hand_suite)
 
-# Register Hands Suite
-import robohive.envs.hands # noqa
-# import robohive.envs.tcdm # noqa # WIP
-robohive_hand_suite = set(gym.envs.registration.registry.env_specs.keys())-robohive_env_suite-_current_gym_envs
-robohive_env_suite  = robohive_env_suite | robohive_hand_suite
-robohive_hand_suite = sorted(robohive_hand_suite)
+    # Register Claw suite
+    import robohive.envs.claws # noqa
+    robohive_claw_suite = set(gym.envs.registration.registry.env_specs.keys())-robohive_env_suite-_current_gym_envs
+    robohive_env_suite  = robohive_env_suite | robohive_claw_suite
+    robohive_claw_suite = sorted(robohive_claw_suite)
 
-# Register Claw suite
-import robohive.envs.claws # noqa
-robohive_claw_suite = set(gym.envs.registration.registry.env_specs.keys())-robohive_env_suite-_current_gym_envs
-robohive_env_suite  = robohive_env_suite | robohive_claw_suite
-robohive_claw_suite = sorted(robohive_claw_suite)
+    # Register Multi-task Suite
+    import robohive.envs.multi_task # noqa
+    robohive_multitask_suite = set(gym.envs.registration.registry.env_specs.keys())-robohive_env_suite-_current_gym_envs
+    robohive_env_suite  = robohive_env_suite | robohive_multitask_suite
+    robohive_multitask_suite = sorted(robohive_multitask_suite)
 
-# Register Multi-task Suite
-import robohive.envs.multi_task # noqa
-robohive_multitask_suite = set(gym.envs.registration.registry.env_specs.keys())-robohive_env_suite-_current_gym_envs
-robohive_env_suite  = robohive_env_suite | robohive_multitask_suite
-robohive_multitask_suite = sorted(robohive_multitask_suite)
+    # Register Locomotion Suite
+    import robohive.envs.quadrupeds # noqa
+    robohive_quad_suite = set(gym.envs.registration.registry.env_specs.keys())-robohive_env_suite-_current_gym_envs
+    robohive_env_suite  = robohive_env_suite | robohive_quad_suite
+    robohive_quad_suite = sorted(robohive_quad_suite)
 
-# Register Locomotion Suite
-import robohive.envs.quadrupeds # noqa
-robohive_quad_suite = set(gym.envs.registration.registry.env_specs.keys())-robohive_env_suite-_current_gym_envs
-robohive_env_suite  = robohive_env_suite | robohive_quad_suite
-robohive_quad_suite = sorted(robohive_quad_suite)
+    # All RoboHive Envs
+    robohive_env_suite = sorted(robohive_env_suite)
+    yield from robohive_env_suite
 
-# All RoboHive Envs
-robohive_env_suite = sorted(robohive_env_suite)
\ No newline at end of file
+robohive_env_suite = _robohive_env_suite