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)

mujoco_playground/_src/locomotion/apollo/base.py

@@ -16,17 +16,17 @@
 
 from typing import Any, Dict, Optional, Union
 
-from etils import epath
 import jax
 import jax.numpy as jp
-from ml_collections import config_dict
 import mujoco
-from mujoco import mjx
 import numpy as np
+from etils import epath
+from ml_collections import config_dict
+from mujoco import mjx
 
 from mujoco_playground._src import mjx_env
-from mujoco_playground._src.collision import geoms_colliding
 from mujoco_playground._src.locomotion.apollo import constants as consts
+from mujoco_playground._src.collision import geoms_colliding
 
 
 def get_assets() -> Dict[str, bytes]:
@@ -46,15 +46,15 @@ class ApolloEnv(mjx_env.MjxEnv):
   """Base class for Apollo environments."""
 
   def __init__(
-      self,
-      xml_path: str,
-      config: config_dict.ConfigDict,
-      config_overrides: Optional[Dict[str, Union[str, int, list[Any]]]] = None,
+    self,
+    xml_path: str,
+    config: config_dict.ConfigDict,
+    config_overrides: Optional[Dict[str, Union[str, int, list[Any]]]] = None,
   ) -> None:
     super().__init__(config, config_overrides)
 
     self._mj_model = mujoco.MjModel.from_xml_string(
-        epath.Path(xml_path).read_text(), assets=get_assets()
+      epath.Path(xml_path).read_text(), assets=get_assets()
     )
     self._mj_model.opt.timestep = self.sim_dt
 
@@ -66,9 +66,7 @@ def __init__(
 
     self._init_q = jp.array(self._mj_model.keyframe("knees_bent").qpos)
     self._default_ctrl = jp.array(self._mj_model.keyframe("knees_bent").ctrl)
-    self._default_pose = jp.array(
-        self._mj_model.keyframe("knees_bent").qpos[7:]
-    )
+    self._default_pose = jp.array(self._mj_model.keyframe("knees_bent").qpos[7:])
     self._actuator_torques = self.mj_model.jnt_actfrcrange[1:, 1]
 
     # Body IDs.
@@ -77,64 +75,52 @@ def __init__(
     # Geom IDs.
     self._floor_geom_id = self._mj_model.geom("floor").id
     self._left_feet_geom_id = np.array(
-        [self._mj_model.geom(name).id for name in consts.LEFT_FEET_GEOMS]
+      [self._mj_model.geom(name).id for name in consts.LEFT_FEET_GEOMS]
     )
     self._right_feet_geom_id = np.array(
-        [self._mj_model.geom(name).id for name in consts.RIGHT_FEET_GEOMS]
+      [self._mj_model.geom(name).id for name in consts.RIGHT_FEET_GEOMS]
     )
     self._left_hand_geom_id = self._mj_model.geom("collision_l_hand_plate").id
     self._right_hand_geom_id = self._mj_model.geom("collision_r_hand_plate").id
     self._left_foot_geom_id = self._mj_model.geom("collision_l_sole").id
     self._right_foot_geom_id = self._mj_model.geom("collision_r_sole").id
-    self._left_shin_geom_id = self._mj_model.geom(
-        "collision_capsule_body_l_shin"
-    ).id
-    self._right_shin_geom_id = self._mj_model.geom(
-        "collision_capsule_body_r_shin"
-    ).id
-    self._left_thigh_geom_id = self._mj_model.geom(
-        "collision_capsule_body_l_thigh"
-    ).id
-    self._right_thigh_geom_id = self._mj_model.geom(
-        "collision_capsule_body_r_thigh"
-    ).id
+    self._left_shin_geom_id = self._mj_model.geom("collision_capsule_body_l_shin").id
+    self._right_shin_geom_id = self._mj_model.geom("collision_capsule_body_r_shin").id
+    self._left_thigh_geom_id = self._mj_model.geom("collision_capsule_body_l_thigh").id
+    self._right_thigh_geom_id = self._mj_model.geom("collision_capsule_body_r_thigh").id
 
     # Site IDs.
     self._imu_site_id = self._mj_model.site("imu").id
     self._feet_site_id = np.array(
-        [self._mj_model.site(name).id for name in consts.FEET_SITES]
+      [self._mj_model.site(name).id for name in consts.FEET_SITES]
     )
 
   # Sensor readings.
 
   def get_gravity(self, data: mjx.Data) -> jax.Array:
     """Return the gravity vector in the world frame."""
-    return mjx_env.get_sensor_data(
-        self.mj_model, data, f"{consts.GRAVITY_SENSOR}"
-    )
+    return mjx_env.get_sensor_data(self.mj_model, data, f"{consts.GRAVITY_SENSOR}")
 
   def get_global_linvel(self, data: mjx.Data) -> jax.Array:
     """Return the linear velocity of the robot in the world frame."""
     return mjx_env.get_sensor_data(
-        self.mj_model, data, f"{consts.GLOBAL_LINVEL_SENSOR}"
+      self.mj_model, data, f"{consts.GLOBAL_LINVEL_SENSOR}"
     )
 
   def get_global_angvel(self, data: mjx.Data) -> jax.Array:
     """Return the angular velocity of the robot in the world frame."""
     return mjx_env.get_sensor_data(
-        self.mj_model, data, f"{consts.GLOBAL_ANGVEL_SENSOR}"
+      self.mj_model, data, f"{consts.GLOBAL_ANGVEL_SENSOR}"
     )
 
   def get_local_linvel(self, data: mjx.Data) -> jax.Array:
     """Return the linear velocity of the robot in the local frame."""
-    return mjx_env.get_sensor_data(
-        self.mj_model, data, f"{consts.LOCAL_LINVEL_SENSOR}"
-    )
+    return mjx_env.get_sensor_data(self.mj_model, data, f"{consts.LOCAL_LINVEL_SENSOR}")
 
   def get_accelerometer(self, data: mjx.Data) -> jax.Array:
     """Return the accelerometer readings in the local frame."""
     return mjx_env.get_sensor_data(
-        self.mj_model, data, f"{consts.ACCELEROMETER_SENSOR}"
+      self.mj_model, data, f"{consts.ACCELEROMETER_SENSOR}"
     )
 
   def get_gyro(self, data: mjx.Data) -> jax.Array:
@@ -143,14 +129,18 @@ def get_gyro(self, data: mjx.Data) -> jax.Array:
 
   def get_feet_ground_contacts(self, data: mjx.Data) -> jax.Array:
     """Return an array indicating whether each foot is in contact with the ground."""
-    left_feet_contact = jp.array([
+    left_feet_contact = jp.array(
+      [
         geoms_colliding(data, geom_id, self._floor_geom_id)
         for geom_id in self._left_feet_geom_id
-    ])
-    right_feet_contact = jp.array([
+      ]
+    )
+    right_feet_contact = jp.array(
+      [
         geoms_colliding(data, geom_id, self._floor_geom_id)
         for geom_id in self._right_feet_geom_id
-    ])
+      ]
+    )
     return jp.hstack([jp.any(left_feet_contact), jp.any(right_feet_contact)])
 
   # Accessors.

mujoco_playground/_src/locomotion/apollo/constants.py

@@ -25,18 +25,18 @@
 
 def task_to_xml(task_name: str) -> epath.Path:
   return {
-      "flat_terrain": FEET_ONLY_FLAT_TERRAIN_XML,
+    "flat_terrain": FEET_ONLY_FLAT_TERRAIN_XML,
   }[task_name]
 
 
 FEET_SITES = [
-    "l_foot",
-    "r_foot",
+  "l_foot",
+  "r_foot",
 ]
 
 HAND_SITES = [
-    "left_palm",
-    "right_palm",
+  "left_palm",
+  "right_palm",
 ]
 
 LEFT_FEET_GEOMS = ["collision_l_sole"]