sync recent changes

mujoco_sim_g1/sim/sim_utilts.py

@@ -0,0 +1,96 @@
+"""
+Utility functions for working with Mujoco models.
+copied from https://github.com/kevinzakka/mink/blob/main/mink/utils.py
+"""
+
+from typing import List
+
+import mujoco
+
+
+def get_body_body_ids(model: mujoco.MjModel, body_id: int) -> List[int]:
+    """Get immediate children bodies belonging to a given body.
+
+    Args:
+        model: Mujoco model.
+        body_id: ID of body.
+
+    Returns:
+        A List containing all child body ids.
+    """
+    return [
+        i
+        for i in range(model.nbody)
+        if model.body_parentid[i] == body_id and body_id != i  # Exclude the body itself.
+    ]
+
+
+def get_subtree_body_ids(model: mujoco.MjModel, body_id: int) -> List[int]:
+    """Get all bodies belonging to subtree starting at a given body.
+
+    Args:
+        model: Mujoco model.
+        body_id: ID of body where subtree starts.
+
+    Returns:
+        A List containing all subtree body ids.
+    """
+    body_ids: List[int] = []
+    stack = [body_id]
+    while stack:
+        body_id = stack.pop()
+        body_ids.append(body_id)
+        stack += get_body_body_ids(model, body_id)
+    return body_ids
+
+
+def get_subtree_body_names(model: mujoco.MjModel, body_id: int) -> List[str]:
+    """Get all bodies belonging to subtree starting at a given body.
+    Args:
+        model: Mujoco model.
+        body_id: ID of body where subtree starts.
+
+    Returns:
+        A List containing all subtree body names.
+    """
+    return [model.body(i).name for i in get_subtree_body_ids(model, body_id)]
+
+
+def get_body_geom_ids(model: mujoco.MjModel, body_id: int) -> List[int]:
+    """Get immediate geoms belonging to a given body.
+
+    Here, immediate geoms are those directly attached to the body and not its
+    descendants.
+
+    Args:
+        model: Mujoco model.
+        body_id: ID of body.
+
+    Returns:
+        A list containing all body geom ids.
+    """
+    geom_start = model.body_geomadr[body_id]
+    geom_end = geom_start + model.body_geomnum[body_id]
+    return list(range(geom_start, geom_end))
+
+
+def get_subtree_geom_ids(model: mujoco.MjModel, body_id: int) -> List[int]:
+    """Get all geoms belonging to subtree starting at a given body.
+
+    Here, a subtree is defined as the kinematic tree starting at the body and including
+    all its descendants.
+
+    Args:
+        model: Mujoco model.
+        body_id: ID of body where subtree starts.
+
+    Returns:
+        A list containing all subtree geom ids.
+    """
+    geom_ids: List[int] = []
+    stack = [body_id]
+    while stack:
+        body_id = stack.pop()
+        geom_ids.extend(get_body_geom_ids(model, body_id))
+        stack += get_body_body_ids(model, body_id)
+    return geom_ids