Formatting

Author: yardenas

mujoco_playground/_src/dm_control_suite/__init__.py

@@ -40,13 +40,21 @@
     "AcrobotSwingupSparse": partial(acrobot.Balance, sparse=True),
     "BallInCup": ball_in_cup.BallInCup,
     "CartpoleBalance": partial(cartpole.Balance, swing_up=False, sparse=False),
-    "CartpoleBalanceSparse": partial(cartpole.Balance, swing_up=False, sparse=True),
+    "CartpoleBalanceSparse": partial(
+        cartpole.Balance, swing_up=False, sparse=True
+    ),
     "CartpoleSwingup": partial(cartpole.Balance, swing_up=True, sparse=False),
-    "CartpoleSwingupSparse": partial(cartpole.Balance, swing_up=True, sparse=True),
+    "CartpoleSwingupSparse": partial(
+        cartpole.Balance, swing_up=True, sparse=True
+    ),
     "CheetahRun": cheetah.Run,
     "FingerSpin": finger.Spin,
-    "FingerTurnEasy": partial(finger.Turn, target_radius=finger.EASY_TARGET_SIZE),
-    "FingerTurnHard": partial(finger.Turn, target_radius=finger.HARD_TARGET_SIZE),
+    "FingerTurnEasy": partial(
+        finger.Turn, target_radius=finger.EASY_TARGET_SIZE
+    ),
+    "FingerTurnHard": partial(
+        finger.Turn, target_radius=finger.HARD_TARGET_SIZE
+    ),
     "FishSwim": fish.Swim,
     "HopperHop": partial(hopper.Hopper, hopping=True),
     "HopperStand": partial(hopper.Hopper, hopping=False),
@@ -99,54 +107,54 @@
 
 
 def __getattr__(name):
-    if name == "ALL_ENVS":
-        return tuple(_envs.keys())
-    raise AttributeError(f"module '{__name__}' has no attribute '{name}'")
+  if name == "ALL_ENVS":
+    return tuple(_envs.keys())
+  raise AttributeError(f"module '{__name__}' has no attribute '{name}'")
 
 
 def register_environment(
     env_name: str,
     env_class: Type[mjx_env.MjxEnv],
     cfg_class: Callable[[], config_dict.ConfigDict],
 ) -> None:
-    """Register a new environment.
+  """Register a new environment.
 
-    Args:
-        env_name: The name of the environment.
-        env_class: The environment class.
-        cfg_class: The default configuration
-    """
-    _envs[env_name] = env_class
-    _cfgs[env_name] = cfg_class
+  Args:
+      env_name: The name of the environment.
+      env_class: The environment class.
+      cfg_class: The default configuration
+  """
+  _envs[env_name] = env_class
+  _cfgs[env_name] = cfg_class
 
 
 def get_default_config(env_name: str) -> config_dict.ConfigDict:
-    """Get the default configuration for an environment."""
-    if env_name not in _cfgs:
-        raise ValueError(
-            f"Env '{env_name}' not found in default configs. Available configs:"
-            f" {list(_cfgs.keys())}"
-        )
-    return _cfgs[env_name]()
+  """Get the default configuration for an environment."""
+  if env_name not in _cfgs:
+    raise ValueError(
+        f"Env '{env_name}' not found in default configs. Available configs:"
+        f" {list(_cfgs.keys())}"
+    )
+  return _cfgs[env_name]()
 
 
 def load(
     env_name: str,
     config: Optional[config_dict.ConfigDict] = None,
     config_overrides: Optional[Dict[str, Union[str, int, list[Any]]]] = None,
 ) -> mjx_env.MjxEnv:
-    """Get an environment instance with the given configuration.
-
-    Args:
-        env_name: The name of the environment.
-        config: The configuration to use. If not provided, the default
-          configuration is used.
-        config_overrides: A dictionary of overrides for the configuration.
-
-    Returns:
-        An instance of the environment.
-    """
-    if env_name not in _envs:
-        raise ValueError(f"Env '{env_name}' not found. Available envs: {_cfgs.keys()}")
-    config = config or get_default_config(env_name)
-    return _envs[env_name](config=config, config_overrides=config_overrides)
+  """Get an environment instance with the given configuration.
+
+  Args:
+      env_name: The name of the environment.
+      config: The configuration to use. If not provided, the default
+        configuration is used.
+      config_overrides: A dictionary of overrides for the configuration.
+
+  Returns:
+      An instance of the environment.
+  """
+  if env_name not in _envs:
+    raise ValueError(f"Env '{env_name}' not found. Available envs: {_cfgs.keys()}")
+  config = config or get_default_config(env_name)
+  return _envs[env_name](config=config, config_overrides=config_overrides)