Add B1 and Kangaroo closed-loop models

python/example_robot_data/robots_loader.py

@@ -1,5 +1,9 @@
+import os
 import typing
+import xml.etree.ElementTree as ET
+from types import MethodType
 
+import numpy as np
 import pinocchio as pin
 
 try:
@@ -23,6 +27,251 @@
 from .utils import RobotLoader, getModelPath, readParamsFromSrdf  # noqa: F401
 
 
+def _robot_get_constraints(self, fallback_infer=True):
+    srdf_path = getattr(self, "srdf_path", None)
+    if srdf_path is None:
+        urdf_path = getattr(self, "urdf", None)
+        if urdf_path is not None:
+            base_dir = os.path.dirname(urdf_path)
+            srdf_candidate = os.path.join(
+                os.path.dirname(base_dir),
+                "srdf",
+                os.path.splitext(os.path.basename(urdf_path))[0] + ".srdf",
+            )
+            if os.path.exists(srdf_candidate):
+                srdf_path = srdf_candidate
+                self.srdf_path = srdf_path
+    if srdf_path is None:
+        return []
+    return constraints_from_srdf(
+        self.model,
+        srdf_path,
+        fallback_infer=fallback_infer,
+        q=self.q0,
+    )
+
+
+class LoopConstraintDescription:
+    def __init__(
+        self,
+        name,
+        joint1_id,
+        joint1_placement,
+        joint2_id,
+        joint2_placement,
+        mask,
+        frame1=None,
+        frame2=None,
+    ):
+        if len(mask) != 6:
+            raise ValueError(
+                f"Constraint mask must contain 6 values, got {len(mask)}: {mask}"
+            )
+        self.name = name
+        self.joint1_id = joint1_id
+        self.joint1_placement = joint1_placement
+        self.joint2_id = joint2_id
+        self.joint2_placement = joint2_placement
+        self.mask = [bool(v) for v in mask]
+        self.frame1 = frame1
+        self.frame2 = frame2
+        self.nc = self.size()
+
+    def size(self):
+        return sum(1 for v in self.mask if v)
+
+
+def _legacy_mask(mask_key):
+    if mask_key == "3d":
+        return [True, False, True, False, False, False]
+    if mask_key == "6d":
+        return [True, False, True, False, True, False]
+    raise ValueError(f"Unsupported legacy constraint type '{mask_key}'.")
+
+
+def _parse_constraint_mask(mask_txt):
+    axis_map = {
+        "x": 0,
+        "tx": 0,
+        "y": 1,
+        "ty": 1,
+        "z": 2,
+        "tz": 2,
+        "roll": 3,
+        "rx": 3,
+        "r": 3,
+        "pitch": 4,
+        "ry": 4,
+        "p": 4,
+        "yaw": 5,
+        "rz": 5,
+    }
+
+    tokens = [
+        token.strip().lower()
+        for token in str(mask_txt).replace(",", " ").split()
+        if token.strip()
+    ]
+    if len(tokens) == 6 and all(
+        token in {"0", "1", "false", "true"} for token in tokens
+    ):
+        return [token in {"1", "true"} for token in tokens]
+
+    mask = [False] * 6
+    for token in tokens:
+        if token not in axis_map:
+            raise ValueError(
+                f"Unsupported mask token '{token}'. "
+                "Use x y z roll pitch yaw, tx ty tz rx ry rz, "
+                "or a 6-entry boolean mask."
+            )
+        mask[axis_map[token]] = True
+    if not any(mask):
+        raise ValueError("Constraint mask cannot be empty.")
+    return mask
+
+
+def _frame_placement_in_parent(model, data, q, frame_id):
+    pin.forwardKinematics(model, data, q)
+    pin.updateFramePlacements(model, data)
+
+    joint_id = model.frames[frame_id].parentJoint
+    parent_id = model.parents[joint_id]
+    return data.oMi[parent_id].inverse() * data.oMf[frame_id]
+
+
+def _frame_motion_derivatives(model, q, frame_id, eps=1e-7, probe=0.5):
+    """Numerically differentiate a frame with respect to its supporting joint."""
+    joint_id = model.frames[frame_id].parentJoint
+    joint = model.joints[joint_id]
+    derivatives = []
+    data = model.createData()
+
+    for velocity_id in range(joint.idx_v, joint.idx_v + joint.nv):
+        tangent = np.zeros(model.nv)
+        tangent[velocity_id] = eps
+
+        # The second center avoids missing a coordinate at a singular sample,
+        # e.g. the z derivative of a planar revolute joint at angle zero.
+        probe_tangent = np.zeros(model.nv)
+        probe_tangent[velocity_id] = probe
+        centers = (q, pin.integrate(model, q, probe_tangent))
+        for center in centers:
+            minus = _frame_placement_in_parent(
+                model, data, pin.integrate(model, center, -tangent), frame_id
+            )
+            plus = _frame_placement_in_parent(
+                model, data, pin.integrate(model, center, tangent), frame_id
+            )
+            linear = (plus.translation - minus.translation) / (2.0 * eps)
+            angular = pin.log3(plus.rotation @ minus.rotation.T) / (2.0 * eps)
+            derivatives.append(np.concatenate((linear, angular)))
+
+    if not derivatives:
+        return np.empty((6, 0))
+    return np.column_stack(derivatives)
+
+
+def _numerical_constraint_mask(model, q, frame1_id, frame2_id):
+    derivatives = np.column_stack(
+        (
+            _frame_motion_derivatives(model, q, frame1_id),
+            _frame_motion_derivatives(model, q, frame2_id),
+        )
+    )
+    if derivatives.shape[1] == 0:
+        raise ValueError(
+            "Cannot infer a constraint mask from frames whose parent joints "
+            "have no degrees of freedom. Specify the mask explicitly."
+        )
+
+    magnitudes = np.max(np.abs(derivatives), axis=1)
+    mask = np.zeros(6, dtype=bool)
+    for rows in (slice(0, 3), slice(3, 6)):
+        scale = np.max(magnitudes[rows])
+        if scale > 1e-9:
+            mask[rows] = magnitudes[rows] > max(1e-9, 1e-6 * scale)
+
+    if not np.any(mask):
+        raise ValueError(
+            "The numerical frame derivatives are zero; specify the constraint "
+            "mask explicitly."
+        )
+    return mask.tolist()
+
+
+def _mask_from_loop_tag(tag, model, q, frame1_id, frame2_id):
+    mask_txt = tag.attrib.get("mask")
+    if mask_txt is not None:
+        return _parse_constraint_mask(mask_txt)
+
+    legacy_type = tag.attrib.get("type")
+    if legacy_type is not None:
+        return _legacy_mask(legacy_type.lower())
+
+    return _numerical_constraint_mask(model, q, frame1_id, frame2_id)
+
+
+def constraints_from_srdf(model, srdf_path, fallback_infer=True, q=None):
+    if q is None:
+        q = pin.neutral(model)
+    root = ET.parse(srdf_path).getroot()
+    constraints = []
+
+    for tag in root.findall(".//loop_constraint"):
+        frame1 = tag.attrib["frame1"]
+        frame2 = tag.attrib["frame2"]
+        id1 = model.getFrameId(frame1)
+        id2 = model.getFrameId(frame2)
+        constraints.append(
+            LoopConstraintDescription(
+                name=f"{frame1}C{frame2}",
+                joint1_id=model.frames[id1].parentJoint,
+                joint1_placement=model.frames[id1].placement,
+                joint2_id=model.frames[id2].parentJoint,
+                joint2_placement=model.frames[id2].placement,
+                mask=_mask_from_loop_tag(tag, model, q, id1, id2),
+                frame1=frame1,
+                frame2=frame2,
+            )
+        )
+
+    if constraints or not fallback_infer:
+        return constraints
+
+    names = [fr.name for fr in model.frames]
+    groups = {}
+    for name in names:
+        lname = name.lower()
+        if "closedloop" not in lname:
+            continue
+        if name.endswith("A"):
+            groups.setdefault(name[:-1], {})["A"] = name
+        elif name.endswith("B"):
+            groups.setdefault(name[:-1], {})["B"] = name
+
+    for _, sides in groups.items():
+        if "A" not in sides or "B" not in sides:
+            continue
+        frame1, frame2 = sides["B"], sides["A"]
+        id1 = model.getFrameId(frame1)
+        id2 = model.getFrameId(frame2)
+        constraints.append(
+            LoopConstraintDescription(
+                name=f"{frame1}C{frame2}",
+                joint1_id=model.frames[id1].parentJoint,
+                joint1_placement=model.frames[id1].placement,
+                joint2_id=model.frames[id2].parentJoint,
+                joint2_placement=model.frames[id2].placement,
+                mask=_numerical_constraint_mask(model, q, id1, id2),
+                frame1=frame1,
+                frame2=frame2,
+            )
+        )
+
+    return constraints
+
+
 class CentauroLoader(RobotLoader):
     path = "centauro_description"
     urdf_filename = "centauro.urdf"
@@ -41,6 +290,42 @@ class B1Loader(RobotLoader):
     free_flyer = True
 
 
+class B1ClosedLoopLoader(RobotLoader):
+    path = "b1_closed_loop_description"
+    urdf_filename = "b1_closed_loop.urdf"
+    urdf_subpath = "urdf"
+    srdf_filename = "b1_closed_loop.srdf"
+    ref_posture = "standing"
+    free_flyer = True
+
+
+class B1LegLoader(RobotLoader):
+    path = "b1_leg_FL"
+    urdf_filename = "b1.urdf"
+    urdf_subpath = "urdf"
+    srdf_filename = "b1.srdf"
+    ref_posture = "standing"
+    free_flyer = False
+
+
+class B1Leg3DLoader(RobotLoader):
+    path = "b1_leg_FL"
+    urdf_filename = "b1_3d.urdf"
+    urdf_subpath = "urdf"
+    srdf_filename = "b1_3d.srdf"
+    ref_posture = "standing"
+    free_flyer = False
+
+
+class B1Leg6DLoader(RobotLoader):
+    path = "b1_leg_FL"
+    urdf_filename = "b1_6d.urdf"
+    urdf_subpath = "urdf"
+    srdf_filename = "b1_6d.srdf"
+    ref_posture = "standing"
+    free_flyer = False
+
+
 class Go1Loader(RobotLoader):
     path = "go1_description"
     urdf_filename = "go1.urdf"
@@ -113,6 +398,51 @@ class Go2Loader(RobotLoader):
     free_flyer = True
 
 
+class KangarooLoader(RobotLoader):
+    path = "kangaroo_description"
+    urdf_filename = "kangaroo.urdf"
+    urdf_subpath = "urdf"
+    srdf_filename = "kangaroo.srdf"
+    ref_posture = "standing"
+    free_flyer = True
+
+
+class KangarooLegsKinLoader(RobotLoader):
+    path = "kangaroo_description"
+    urdf_filename = "kangaroo_loop.urdf"
+    urdf_subpath = "urdf"
+    srdf_filename = "kangaroo_loop.srdf"
+    ref_posture = "standing"
+    free_flyer = True
+
+
+class KangarooKneeKinLoader(RobotLoader):
+    path = "kangaroo_description"
+    urdf_filename = "kangaroo_knee_kin.urdf"
+    urdf_subpath = "urdf"
+    srdf_filename = "kangaroo_knee_kin.srdf"
+    ref_posture = "standing"
+    free_flyer = True
+
+
+class KangarooKnee1AnkleKinLoader(RobotLoader):
+    path = "kangaroo_description"
+    urdf_filename = "kangaroo_knee_1ankle_kin.urdf"
+    urdf_subpath = "urdf"
+    srdf_filename = "kangaroo_knee_1ankle_kin.srdf"
+    ref_posture = "standing"
+    free_flyer = True
+
+
+class KangarooKnee2AnkleKinLoader(RobotLoader):
+    path = "kangaroo_description"
+    urdf_filename = "kangaroo_knee_2ankle_kin.urdf"
+    urdf_subpath = "urdf"
+    srdf_filename = "kangaroo_knee_2ankle_kin.srdf"
+    ref_posture = "standing"
+    free_flyer = True
+
+
 class A1Loader(RobotLoader):
     path = "a1_description"
     urdf_filename = "a1.urdf"
@@ -483,8 +813,17 @@ class xArm7Loader(RobotLoader):
 
 
 ROBOTS = {
+    "kangaroo": KangarooLoader,
+    "kangaroo_legs_kin": KangarooLegsKinLoader,
+    "kangaroo_knee_kin": KangarooKneeKinLoader,
+    "kangaroo_knee_1ankle_kin": KangarooKnee1AnkleKinLoader,
+    "kangaroo_knee_2ankle_kin": KangarooKnee2AnkleKinLoader,
     "centauro": CentauroLoader,
     "b1": B1Loader,
+    "b1_closed_loop": B1ClosedLoopLoader,
+    "b1_leg": B1LegLoader,
+    "b1_leg_3D": B1Leg3DLoader,
+    "b1_leg_6D": B1Leg6DLoader,
     "bravo7_gripper": Bravo7GripperLoader,
     "bravo7_no_ee": Bravo7NoEndEffectorLoader,
     "falcon_bravo7_no_ee": FalconBravo7NoEndEffectorLoader,
@@ -564,6 +903,7 @@ def loader(name, display=False, rootNodeName="", verbose=False):
         robots = ", ".join(sorted(ROBOTS.keys()))
         raise ValueError(f"Robot '{name}' not found. Possible values are {robots}")
     inst = ROBOTS[name](verbose=verbose)
+    inst.robot.get_constraints = MethodType(_robot_get_constraints, inst.robot)
     if display:
         if rootNodeName:
             inst.robot.initViewer()

python/example_robot_data/utils.py

@@ -122,6 +122,7 @@ def __init__(self, verbose=False):
             self.srdf_path = join(
                 self.model_path, self.path, self.srdf_subpath, self.srdf_filename
             )
+            self.robot.srdf_path = self.srdf_path
             self.robot.q0 = readParamsFromSrdf(
                 self.robot.model,
                 self.srdf_path,
@@ -143,6 +144,7 @@ def __init__(self, verbose=False):
                 self.robot.collision_data = self.robot.collision_model.createData()
         else:
             self.srdf_path = None
+            self.robot.srdf_path = None
             self.robot.q0 = pin.neutral(self.robot.model)
         root = getModelPath(self.path)
         self.robot.urdf = join(root, self.path, self.urdf_subpath, self.urdf_filename)