Nits.

Author: kevinzakka

mink/utils.py

@@ -97,32 +97,48 @@ def custom_configuration_vector(
     return q
 
 
-def get_subtree_geom_ids(model: mujoco.MjModel, body_id: int) -> list[int]:
-    """Get all geoms belonging to subtree starting at a given body.
+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 geom ids.
+        A list containing all subtree body ids.
     """
-    geom_ids = []
+    body_ids: list[int] = []
     stack = [body_id]
     while stack:
         body_id = stack.pop()
-        geom_ids.extend(get_body_geom_ids(model, body_id))
-        children = [
-            i
-            for i in range(model.nbody)
-            if model.body_parentid[i] == body_id and body_id != i
-        ]
-        stack.extend(children)
-    return geom_ids
+        body_ids.append(body_id)
+        stack += get_body_body_ids(model, body_id)
+    return body_ids
 
 
 def get_body_geom_ids(model: mujoco.MjModel, body_id: int) -> list[int]:
-    """Get all geoms belonging to a given body.
+    """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.
@@ -136,25 +152,23 @@ def get_body_geom_ids(model: mujoco.MjModel, body_id: int) -> list[int]:
     return list(range(geom_start, geom_end))
 
 
-def get_subtree_body_ids(model: mujoco.MjModel, body_id: int) -> list[int]:
-    """Get all bodies belonging to subtree starting at a given body.
+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 body ids.
+        A list containing all subtree geom ids.
     """
-    body_ids = []
+    geom_ids: list[int] = []
     stack = [body_id]
     while stack:
         body_id = stack.pop()
-        body_ids.append(body_id)
-        children = [
-            i
-            for i in range(model.nbody)
-            if model.body_parentid[i] == body_id and body_id != i
-        ]
-        stack.extend(children)
-    return body_ids
+        geom_ids.extend(get_body_geom_ids(model, body_id))
+        stack += get_body_body_ids(model, body_id)
+    return geom_ids