Merge pull request #182 from Balint-H:main

PiperOrigin-RevId: 797976253
Change-Id: I36b3cb899c7a39b773312ee40f115ed50a70c58c

Author: copybara-github

mujoco_playground/_src/locomotion/__init__.py

@@ -42,8 +42,6 @@
 from mujoco_playground._src.locomotion.t1 import randomize as t1_randomize
 
 
-mjx_env.ensure_menagerie_exists()  # Ensure menagerie exists when module is imported.
-
 _envs = {
     "ApolloJoystickFlatTerrain": functools.partial(
         apollo_joystick.Joystick, task="flat_terrain"
@@ -181,6 +179,7 @@ def load(
   Returns:
       An instance of the environment.
   """
+  mjx_env.ensure_menagerie_exists()  # Ensure menagerie exists when environment is loaded.
   if env_name not in _envs:
     raise ValueError(
         f"Env '{env_name}' not found. Available envs: {_cfgs.keys()}"

mujoco_playground/_src/manipulation/__init__.py

@@ -30,8 +30,6 @@
 from mujoco_playground._src.manipulation.leap_hand import rotate_z as leap_rotate_z
 
 
-mjx_env.ensure_menagerie_exists()  # Ensure menagerie exists when module is imported.
-
 _envs = {
     "AlohaHandOver": aloha_handover.HandOver,
     "AlohaSinglePegInsertion": aloha_peg.SinglePegInsertion,
@@ -110,6 +108,7 @@ def load(
   Returns:
       An instance of the environment.
   """
+  mjx_env.ensure_menagerie_exists()  # Ensure menagerie exists when environment is loaded.
   if env_name not in _envs:
     raise ValueError(
         f"Env '{env_name}' not found. Available envs: {_cfgs.keys()}"