sync recent changes

mujoco_sim_g1/sim/metric_utils.py

@@ -0,0 +1,71 @@
+from typing import List, Tuple
+
+import mujoco
+
+from .sim_utilts import get_body_geom_ids
+
+
+def check_contact(
+    mj_model: mujoco.MjModel,
+    mj_data: mujoco.MjData,
+    bodies_1: List[str] | str,
+    bodies_2: List[str] | str,
+    return_all_contact_bodies: bool = False,
+) -> Tuple[bool, List[Tuple[str, str]]] | bool:
+    """
+    Finds contact between two body groups. Any geom in the body is considered to be in contact.
+    Args:
+        mj_model (MujocoModel): Current simulation object
+        mj_data (MjData): Current simulation data
+        bodies_1 (str or list of int): an individual body name or list of body names.
+        bodies_2 (str or list of int): another individual body name or list of body names.
+    Returns:
+        bool: True if any body in @bodies_1 is in contact with any body in @bodies_2.
+    """
+    if isinstance(bodies_1, str):
+        bodies_1 = [bodies_1]
+    if isinstance(bodies_2, str):
+        bodies_2 = [bodies_2]
+
+    geoms_1 = [get_body_geom_ids(mj_model, mj_model.body(g).id) for g in bodies_1]
+    geoms_1 = [g for geom_list in geoms_1 for g in geom_list]
+    geoms_2 = [get_body_geom_ids(mj_model, mj_model.body(g).id) for g in bodies_2]
+    geoms_2 = [g for geom_list in geoms_2 for g in geom_list]
+    contact_bodies = []
+    for i in range(mj_data.ncon):
+        contact = mj_data.contact[i]
+        # check contact geom in geoms
+        c1_in_g1 = contact.geom1 in geoms_1
+        c2_in_g2 = contact.geom2 in geoms_2 if geoms_2 is not None else True
+        # check contact geom in geoms (flipped)
+        c2_in_g1 = contact.geom2 in geoms_1
+        c1_in_g2 = contact.geom1 in geoms_2 if geoms_2 is not None else True
+        if (c1_in_g1 and c2_in_g2) or (c1_in_g2 and c2_in_g1):
+            contact_bodies.append(
+                (
+                    mj_model.body(mj_model.geom(contact.geom1).bodyid).name,
+                    mj_model.body(mj_model.geom(contact.geom2).bodyid).name,
+                )
+            )
+            if not return_all_contact_bodies:
+                break
+    if return_all_contact_bodies:
+        return len(contact_bodies) > 0, set(contact_bodies)
+    else:
+        return len(contact_bodies) > 0
+
+
+def check_height(
+    mj_model: mujoco.MjModel,
+    mj_data: mujoco.MjData,
+    geom_name: str,
+    lower_bound: float = -float("inf"),
+    upper_bound: float = float("inf"),
+):
+    """
+    Checks if the height of a geom is greater than a given height.
+    """
+    geom_id = mj_model.geom(geom_name).id
+    return (
+        mj_data.geom_xpos[geom_id][2] < upper_bound and mj_data.geom_xpos[geom_id][2] > lower_bound
+    )