diff --git a/test/test_black_box.py b/test/test_black_box.py
index fd9da4f..8cdc543 100644
--- a/test/test_black_box.py
+++ b/test/test_black_box.py
@@ -10,13 +10,12 @@ from gymnasium.core import ActType, ObsType
 import fancy_gym
 from fancy_gym.black_box.raw_interface_wrapper import RawInterfaceWrapper
 from fancy_gym.utils.wrappers import TimeAwareObservation
-from test.utils import ugly_hack_to_mitigate_metaworld_bug
 
 SEED = 1
 ENV_IDS = ['fancy/Reacher5d-v0', 'dm_control/ball_in_cup-catch-v0', 'metaworld/reach-v2', 'Reacher-v2']
 WRAPPERS = [fancy_gym.envs.mujoco.reacher.MPWrapper, fancy_gym.dmc.suite.ball_in_cup.MPWrapper,
             fancy_gym.meta.goal_object_change_mp_wrapper.MPWrapper, fancy_gym.open_ai.mujoco.reacher_v2.MPWrapper]
-ALL_MP_ENVS = chain(*fancy_gym.ALL_MOVEMENT_PRIMITIVE_ENVIRONMENTS.values())
+ALL_MP_ENVS = fancy_gym.ALL_MOVEMENT_PRIMITIVE_ENVIRONMENTS['all']
 
 MAX_STEPS_FALLBACK = 100
 
@@ -129,7 +128,7 @@ def test_length(mp_type: str, env_wrap: Tuple[str, Type[RawInterfaceWrapper]]):
 
     for i in range(5):
         env.reset(seed=SEED)
-        ugly_hack_to_mitigate_metaworld_bug(env)  # TODO: Remove, when metaworld fixed it upstream
+
         _obs, _reward, _terminated, _truncated, info = env.step(env.action_space.sample())
         length = info['trajectory_length']
 
@@ -336,7 +335,6 @@ def test_learn_tau_and_delay(mp_type: str, tau: float, delay: float):
     for i in range(5):
         if done:
             env.reset(seed=SEED)
-            ugly_hack_to_mitigate_metaworld_bug(env)
         action = env.action_space.sample()
         action[0] = tau
         action[1] = delay
diff --git a/test/test_dmc_envs.py b/test/test_dmc_envs.py
index 7a96d94..3602da6 100644
--- a/test/test_dmc_envs.py
+++ b/test/test_dmc_envs.py
@@ -8,10 +8,10 @@ import fancy_gym
 from test.utils import run_env, run_env_determinism
 
 DMC_IDS = [spec.id for spec in gym.envs.registry.values() if
-                  spec.id.startswith('dm_control/')
-                  and 'compatibility-env-v0' not in spec.id
-                  and 'lqr-lqr' not in spec.id]
-DMC_MP_IDS = list(chain(*fancy_gym.ALL_DMC_MOVEMENT_PRIMITIVE_ENVIRONMENTS.values()))
+           spec.id.startswith('dm_control/')
+           and 'compatibility-env-v0' not in spec.id
+           and 'lqr-lqr' not in spec.id]
+DMC_MP_IDS = fancy_gym.ALL_DMC_MOVEMENT_PRIMITIVE_ENVIRONMENTS['all']
 SEED = 1
 
 
@@ -26,11 +26,13 @@ def test_step_dm_control_determinism(env_id: str):
     """Tests that for step environments identical seeds produce identical trajectories."""
     run_env_determinism(env_id, SEED, 5000, wrappers=[gym.wrappers.FlattenObservation])
 
+
 @pytest.mark.parametrize('env_id', DMC_MP_IDS)
 def test_bb_dmc_functionality(env_id: str):
     """Tests that black box environments run without errors using random actions."""
     run_env(env_id)
 
+
 @pytest.mark.parametrize('env_id', DMC_MP_IDS)
 def test_bb_dmc_determinism(env_id: str):
     """Tests that for black box environment identical seeds produce identical trajectories."""
diff --git a/test/test_fancy_envs.py b/test/test_fancy_envs.py
index 898cc08..a15c837 100644
--- a/test/test_fancy_envs.py
+++ b/test/test_fancy_envs.py
@@ -10,7 +10,7 @@ from test.utils import run_env, run_env_determinism
 CUSTOM_IDS = [id for id, spec in gym.envs.registry.items() if
               not isinstance(spec.entry_point, Callable) and
               "fancy_gym" in spec.entry_point and 'make_bb_env_helper' not in spec.entry_point]
-CUSTOM_MP_IDS = list(chain(*fancy_gym.ALL_FANCY_MOVEMENT_PRIMITIVE_ENVIRONMENTS.values()))
+CUSTOM_MP_IDS = fancy_gym.ALL_FANCY_MOVEMENT_PRIMITIVE_ENVIRONMENTS['all']
 SEED = 1
 
 
diff --git a/test/test_gym_envs.py b/test/test_gym_envs.py
index 76b5c85..5e50d53 100644
--- a/test/test_gym_envs.py
+++ b/test/test_gym_envs.py
@@ -14,7 +14,7 @@ GYM_IDS = [spec.id for spec in gym.envs.registry.values() if
            and 'jax' not in spec.id.lower()
            and not re.match(r'GymV2.Environment', spec.id)
            ]
-GYM_MP_IDS = list(chain(*fancy_gym.ALL_DMC_MOVEMENT_PRIMITIVE_ENVIRONMENTS.values()))
+GYM_MP_IDS = fancy_gym.ALL_DMC_MOVEMENT_PRIMITIVE_ENVIRONMENTS['all']
 SEED = 1
 
 
diff --git a/test/test_metaworld_envs.py b/test/test_metaworld_envs.py
index 55de621..18b922f 100644
--- a/test/test_metaworld_envs.py
+++ b/test/test_metaworld_envs.py
@@ -8,7 +8,7 @@ from test.utils import run_env, run_env_determinism
 
 METAWORLD_IDS = [f'metaworld:{env.split("-goal-observable")[0]}' for env, _ in
                  ALL_V2_ENVIRONMENTS_GOAL_OBSERVABLE.items()]
-METAWORLD_MP_IDS = list(chain(*fancy_gym.ALL_METAWORLD_MOVEMENT_PRIMITIVE_ENVIRONMENTS.values()))
+METAWORLD_MP_IDS = fancy_gym.ALL_METAWORLD_MOVEMENT_PRIMITIVE_ENVIRONMENTS['all']
 SEED = 1
 
 
diff --git a/test/test_replanning_sequencing.py b/test/test_replanning_sequencing.py
index 0563a76..c2edf42 100644
--- a/test/test_replanning_sequencing.py
+++ b/test/test_replanning_sequencing.py
@@ -13,13 +13,12 @@ import fancy_gym
 from fancy_gym.black_box.raw_interface_wrapper import RawInterfaceWrapper
 from fancy_gym.utils.wrappers import TimeAwareObservation
 from fancy_gym.utils.make_env_helpers import ensure_finite_time
-from test.utils import ugly_hack_to_mitigate_metaworld_bug
 
 SEED = 1
-ENV_IDS = ['fancy/Reacher5d-v0', 'dmc/ball_in_cup-catch-v0', 'metaworld/reach-v2', 'Reacher-v2']
+ENV_IDS = ['fancy/Reacher5d-v0', 'dm_control/ball_in_cup-catch-v0', 'metaworld/reach-v2', 'Reacher-v2']
 WRAPPERS = [fancy_gym.envs.mujoco.reacher.MPWrapper, fancy_gym.dmc.suite.ball_in_cup.MPWrapper,
             fancy_gym.meta.goal_object_change_mp_wrapper.MPWrapper, fancy_gym.open_ai.mujoco.reacher_v2.MPWrapper]
-ALL_MP_ENVS = chain(*fancy_gym.ALL_MOVEMENT_PRIMITIVE_ENVIRONMENTS.values())
+ALL_MP_ENVS = fancy_gym.ALL_MOVEMENT_PRIMITIVE_ENVIRONMENTS['all']
 
 MAX_STEPS_FALLBACK = 50
 
@@ -94,7 +93,7 @@ def test_learn_sub_trajectories(mp_type: str, env_wrap: Tuple[str, Type[RawInter
     for i in range(25):
         if done:
             env.reset(seed=SEED)
-            ugly_hack_to_mitigate_metaworld_bug(env)  # TODO: Remove, when metaworld fixed it upstream
+
         action = env.action_space.sample()
         _obs, _reward, terminated, truncated, info = env.step(action)
         done = terminated or truncated
@@ -159,7 +158,6 @@ def test_replanning_time(mp_type: str, env_wrap: Tuple[str, Type[RawInterfaceWra
             print(done, (i + 1), episode_steps)
             assert (i + 1) % episode_steps == 0
             env.reset(seed=SEED)
-            ugly_hack_to_mitigate_metaworld_bug(env)  # TODO: Remove, when metaworld fixed it upstream
 
         assert replanning_schedule(None, None, None, None, length)
 
diff --git a/test/utils.py b/test/utils.py
index 8f92bbd..01e33fe 100644
--- a/test/utils.py
+++ b/test/utils.py
@@ -100,13 +100,3 @@ def verify_reward(reward):
 def verify_done(done):
     assert isinstance(
         done, bool), f"Returned {done} as done flag, expected bool."
-
-
-def ugly_hack_to_mitigate_metaworld_bug(env):
-    head = env
-    try:
-        for i in range(16):
-            head.curr_path_length = 0
-            head = head.env
-    except:
-        pass