warp-lang 0.10.1__py3-none-win_amd64.whl → 0.11.0__py3-none-win_amd64.whl

warp/sim/collide.py

@@ -561,9 +561,17 @@ def create_soft_contacts(
             shape_v = wp.cw_mul(shape_v, geo_scale)
 
             delta = x_local - shape_p
+            
             d = wp.length(delta) * sign
             n = wp.normalize(delta) * sign
             v = shape_v
+    
+    if geo_type == wp.sim.GEO_SDF:
+        volume = geo.source[shape_index]
+        xpred_local = wp.volume_world_to_index(volume, wp.cw_div(x_local, geo_scale))
+        nn = wp.vec3(0.0, 0.0, 0.0)
+        d = wp.volume_sample_grad_f(volume, xpred_local, wp.Volume.LINEAR, nn)
+        n = wp.normalize(nn)        
 
     if geo_type == wp.sim.GEO_PLANE:
         d = plane_sdf(geo_scale[0], geo_scale[1], x_local)

warp/sim/import_mjcf.py

@@ -8,18 +8,17 @@
 
 import math
 import os
+import re
 import xml.etree.ElementTree as ET
 
 import numpy as np
-
 import warp as wp
-from warp.sim.model import JOINT_COMPOUND, JOINT_UNIVERSAL
-from warp.sim.model import Mesh
 
 
 def parse_mjcf(
-    filename,
+    mjcf_filename,
     builder,
+    xform=wp.transform(),
     density=1000.0,
     stiffness=0.0,
     damping=0.0,
@@ -30,69 +29,206 @@ def parse_mjcf(
     contact_restitution=0.5,
     limit_ke=100.0,
     limit_kd=10.0,
+    scale=1.0,
     armature=0.0,
     armature_scale=1.0,
-    parse_meshes=False,
-    enable_self_collisions=True,
+    parse_meshes=True,
+    enable_self_collisions=False,
+    up_axis="Z",
+    ignore_classes=[],
+    collapse_fixed_joints=False,
 ):
-    file = ET.parse(filename)
+    """
+    Parses MuJoCo XML (MJCF) file and adds the bodies and joints to the given ModelBuilder.
+
+    Args:
+        mjcf_filename (str): The filename of the MuJoCo file to parse.
+        builder (ModelBuilder): The :class:`ModelBuilder` to add the bodies and joints to.
+        xform (:ref:`transform <transform>`): The transform to apply to the imported mechanism.
+        density (float): The density of the shapes in kg/m^3 which will be used to calculate the body mass and inertia.
+        stiffness (float): The stiffness of the joints.
+        damping (float): The damping of the joints.
+        contact_ke (float): The stiffness of the shape contacts (used by SemiImplicitIntegrator).
+        contact_kd (float): The damping of the shape contacts (used by SemiImplicitIntegrator).
+        contact_kf (float): The friction stiffness of the shape contacts (used by SemiImplicitIntegrator).
+        contact_mu (float): The friction coefficient of the shape contacts.
+        contact_restitution (float): The restitution coefficient of the shape contacts.
+        limit_ke (float): The stiffness of the joint limits (used by SemiImplicitIntegrator).
+        limit_kd (float): The damping of the joint limits (used by SemiImplicitIntegrator).
+        scale (float): The scaling factor to apply to the imported mechanism.
+        armature (float): Default joint armature to use if `armature` has not been defined for a joint in the MJCF.
+        armature_scale (float): Scaling factor to apply to the MJCF-defined joint armature values.
+        parse_meshes (bool): Whether geometries of type `"mesh"` should be parsed. If False, geometries of type `"mesh"` are ignored.
+        enable_self_collisions (bool): If True, self-collisions are enabled.
+        up_axis (str): The up axis of the mechanism. Can be either `"X"`, `"Y"` or `"Z"`. The default is `"Z"`.
+        ignore_classes (List[str]): A list of regular expressions. Bodies and joints with a class matching one of the regular expressions will be ignored.
+        collapse_fixed_joints (bool): If True, fixed joints are removed and the respective bodies are merged.
+
+    Note:
+        The inertia and masses of the bodies are calculated from the shape geometry and the given density. The values defined in the MJCF are not respected at the moment.
+
+        The handling of advanced features, such as MJCF classes, is still experimental.
+    """
+    mjcf_dirname = os.path.dirname(mjcf_filename)
+    file = ET.parse(mjcf_filename)
     root = file.getroot()
 
-    type_map = {
-        "ball": wp.sim.JOINT_BALL,
-        "hinge": wp.sim.JOINT_REVOLUTE,
-        "slide": wp.sim.JOINT_PRISMATIC,
-        "free": wp.sim.JOINT_FREE,
-        "fixed": wp.sim.JOINT_FIXED,
-    }
-
-    def parse_float(node, key, default):
-        if key in node.attrib:
-            return float(node.attrib[key])
+    use_degrees = True  # angles are in degrees by default
+    euler_seq = [1, 2, 3]  # XYZ by default
+
+    compiler = root.find("compiler")
+    if compiler is not None:
+        use_degrees = compiler.attrib.get("angle", "degree").lower() == "degree"
+        euler_seq = ["xyz".index(c) + 1 for c in compiler.attrib.get("eulerseq", "xyz").lower()]
+        mesh_dir = compiler.attrib.get("meshdir", ".")
+
+    mesh_assets = {}
+    for asset in root.findall("asset"):
+        for mesh in asset.findall("mesh"):
+            if "file" in mesh.attrib:
+                fname = os.path.join(mesh_dir, mesh.attrib["file"])
+                # handle stl relative paths
+                if not os.path.isabs(fname):
+                    fname = os.path.abspath(os.path.join(mjcf_dirname, fname))
+                if "name" in mesh.attrib:
+                    mesh_assets[mesh.attrib["name"]] = fname
+                else:
+                    name = ".".join(os.path.basename(fname).split(".")[:-1])
+                    mesh_assets[name] = fname
+
+    class_parent = {}
+    class_children = {}
+    class_defaults = {"__all__": {}}
+
+    def get_class(element):
+        return element.get("class", "__all__")
+
+    def parse_default(node, parent):
+        nonlocal class_parent
+        nonlocal class_children
+        nonlocal class_defaults
+        class_name = "__all__"
+        if "class" in node.attrib:
+            class_name = node.attrib["class"]
+            class_parent[class_name] = parent
+            parent = parent or "__all__"
+            if parent not in class_children:
+                class_children[parent] = []
+            class_children[parent].append(class_name)
+
+        if class_name not in class_defaults:
+            class_defaults[class_name] = {}
+        for child in node:
+            if child.tag == "default":
+                parse_default(child, node.get("class"))
+            else:
+                class_defaults[class_name][child.tag] = child.attrib
+
+    for default in root.findall("default"):
+        parse_default(default, None)
+
+    def merge_attrib(default_attrib: dict, incoming_attrib: dict):
+        attrib = default_attrib.copy()
+        attrib.update(incoming_attrib)
+        return attrib
+
+    if isinstance(up_axis, str):
+        up_axis = "XYZ".index(up_axis.upper())
+    sqh = np.sqrt(0.5)
+    if up_axis == 0:
+        xform = wp.transform(xform.p, wp.quat(0.0, 0.0, -sqh, sqh) * xform.q)
+    elif up_axis == 2:
+        xform = wp.transform(xform.p, wp.quat(sqh, 0.0, 0.0, -sqh) * xform.q)
+    # do not apply scaling to the root transform
+    xform = wp.transform(np.array(xform.p) / scale, xform.q)
+
+    def parse_float(attrib, key, default):
+        if key in attrib:
+            return float(attrib[key])
         else:
             return default
 
-    def parse_vec(node, key, default):
-        if key in node.attrib:
-            return np.fromstring(node.attrib[key], sep=" ")
+    def parse_vec(attrib, key, default):
+        if key in attrib:
+            out = np.fromstring(attrib[key], sep=" ", dtype=np.float32)
         else:
-            return np.array(default)
+            out = np.array(default, dtype=np.float32)
+
+        length = len(out)
+        if length == 1:
+            return wp.vec(len(default), wp.float32)(out[0], out[0], out[0])
+
+        return wp.vec(length, wp.float32)(out)
+
+    def parse_orientation(attrib):
+        if "quat" in attrib:
+            wxyz = np.fromstring(attrib["quat"], sep=" ")
+            return wp.normalize(wp.quat(*wxyz[1:], wxyz[0]))
+        if "euler" in attrib:
+            euler = np.fromstring(attrib["euler"], sep=" ")
+            if use_degrees:
+                euler *= np.pi / 180
+            return wp.quat_from_euler(euler, *euler_seq)
+        if "axisangle" in attrib:
+            axisangle = np.fromstring(attrib["axisangle"], sep=" ")
+            angle = axisangle[3]
+            if use_degrees:
+                angle *= np.pi / 180
+            axis = wp.normalize(wp.vec3(*axisangle[:3]))
+            return wp.quat_from_axis_angle(axis, angle)
+        if "xyaxes" in attrib:
+            xyaxes = np.fromstring(attrib["xyaxes"], sep=" ")
+            xaxis = wp.normalize(wp.vec3(*xyaxes[:3]))
+            zaxis = wp.normalize(wp.vec3(*xyaxes[3:]))
+            yaxis = wp.normalize(wp.cross(zaxis, xaxis))
+            rot_matrix = np.array([xaxis, yaxis, zaxis]).T
+            return wp.quat_from_matrix(rot_matrix)
+        if "zaxis" in attrib:
+            zaxis = np.fromstring(attrib["zaxis"], sep=" ")
+            zaxis = wp.normalize(wp.vec3(*zaxis))
+            xaxis = wp.normalize(wp.cross(wp.vec3(0, 0, 1), zaxis))
+            yaxis = wp.normalize(wp.cross(zaxis, xaxis))
+            rot_matrix = np.array([xaxis, yaxis, zaxis]).T
+            return wp.quat_from_matrix(rot_matrix)
+        return wp.quat_identity()
 
     def parse_mesh(geom):
         import trimesh
 
         faces = []
         vertices = []
-        stl_file = next(
-            filter(
-                lambda m: m.attrib["name"] == geom.attrib["mesh"],
-                root.find("asset").findall("mesh"),
-            )
-        ).attrib["file"]
-        # handle stl relative paths
-        if not os.path.isabs(stl_file):
-            stl_file = os.path.join(os.path.dirname(filename), stl_file)
+        stl_file = mesh_assets[geom["mesh"]]
         m = trimesh.load(stl_file)
 
         for v in m.vertices:
-            vertices.append(np.array(v))
+            vertices.append(np.array(v) * scale)
 
         for f in m.faces:
             faces.append(int(f[0]))
             faces.append(int(f[1]))
             faces.append(int(f[2]))
-        return Mesh(vertices, faces), m.scale
-
-    def parse_body(body, parent):
-        body_name = body.attrib["name"]
-        body_pos = parse_vec(body, "pos", (0.0, 0.0, 0.0))
-        body_ori_euler = parse_vec(body, "euler", (0.0, 0.0, 0.0))
-        if len(np.nonzero(body_ori_euler)[0]) > 0:
-            body_axis = tuple(np.sign(body_ori_euler))
-            body_angle = body_ori_euler[np.nonzero(body_ori_euler)[0].item()] / 180 * np.pi
-            body_ori = wp.utils.quat_from_axis_angle(body_axis, body_angle)
+        return wp.sim.Mesh(vertices, faces), m.scale
+
+    def parse_body(body, parent, incoming_defaults: dict):
+        body_class = body.get("childclass")
+        if body_class is None:
+            defaults = incoming_defaults
         else:
-            body_ori = wp.quat_identity()
+            for pattern in ignore_classes:
+                if re.match(pattern, body_class):
+                    return
+            defaults = merge_attrib(incoming_defaults, class_defaults[body_class])
+        if "body" in defaults:
+            body_attrib = merge_attrib(defaults["body"], body.attrib)
+        else:
+            body_attrib = body.attrib
+        body_name = body_attrib["name"]
+        body_pos = parse_vec(body_attrib, "pos", (0.0, 0.0, 0.0))
+        body_ori = parse_orientation(body_attrib)
+        if parent == -1:
+            body_pos = wp.transform_point(xform, body_pos)
+            body_ori = xform.q * body_ori
+        body_pos *= scale
 
         joint_armature = []
         joint_name = []
@@ -102,43 +238,55 @@ def parse_mjcf(
         angular_axes = []
         joint_type = None
 
-        joints = body.findall("joint")
-        for i, joint in enumerate(joints):
-            # default to hinge if not specified
-            if "type" not in joint.attrib:
-                joint.attrib["type"] = "hinge"
-
-            joint_name.append(joint.attrib["name"])
-            joint_pos.append(parse_vec(joint, "pos", (0.0, 0.0, 0.0)))
-            # TODO parse joint (child transform) rotation?
-            joint_range = parse_vec(joint, "range", (-3.0, 3.0))
-            joint_armature.append(parse_float(joint, "armature", armature) * armature_scale)
-
-            if joint.attrib["type"].lower() == "free":
-                joint_type = wp.sim.JOINT_FREE
-                break
-            is_angular = joint.attrib["type"].lower() == "hinge"
-            mode = wp.sim.JOINT_MODE_LIMIT
-            if stiffness > 0.0 or "stiffness" in joint.attrib:
-                mode = wp.sim.JOINT_MODE_TARGET_POSITION
-            ax = wp.sim.model.JointAxis(
-                axis=parse_vec(joint, "axis", (0.0, 0.0, 0.0)),
-                limit_lower=(np.deg2rad(joint_range[0]) if is_angular else joint_range[0]),
-                limit_upper=(np.deg2rad(joint_range[1]) if is_angular else joint_range[1]),
-                target_ke=parse_float(joint, "stiffness", stiffness),
-                target_kd=parse_float(joint, "damping", damping),
-                limit_ke=limit_ke,
-                limit_kd=limit_kd,
-                mode=mode,
-            )
-            if is_angular:
-                angular_axes.append(ax)
-            else:
-                linear_axes.append(ax)
+        freejoint_tags = body.findall("freejoint")
+        if len(freejoint_tags) > 0:
+            joint_type = wp.sim.JOINT_FREE
+            joint_name.append(freejoint_tags[0].attrib.get("name", f"{body_name}_freejoint"))
+        else:
+            joints = body.findall("joint")
+            for i, joint in enumerate(joints):
+                if "joint" in defaults:
+                    joint_attrib = merge_attrib(defaults["joint"], joint.attrib)
+                else:
+                    joint_attrib = joint.attrib
+
+                # default to hinge if not specified
+                joint_type_str = joint_attrib.get("type", "hinge")
+
+                joint_name.append(joint_attrib["name"])
+                joint_pos.append(parse_vec(joint_attrib, "pos", (0.0, 0.0, 0.0)) * scale)
+                joint_range = parse_vec(joint_attrib, "range", (-3.0, 3.0))
+                joint_armature.append(parse_float(joint_attrib, "armature", armature) * armature_scale)
+
+                if joint_type_str == "free":
+                    joint_type = wp.sim.JOINT_FREE
+                    break
+                if joint_type_str == "fixed":
+                    joint_type = wp.sim.JOINT_FIXED
+                    break
+                is_angular = joint_type_str == "hinge"
+                mode = wp.sim.JOINT_MODE_LIMIT
+                if stiffness > 0.0 or "stiffness" in joint_attrib:
+                    mode = wp.sim.JOINT_MODE_TARGET_POSITION
+                axis_vec = parse_vec(joint_attrib, "axis", (0.0, 0.0, 0.0))
+                ax = wp.sim.model.JointAxis(
+                    axis=axis_vec,
+                    limit_lower=(np.deg2rad(joint_range[0]) if is_angular and use_degrees else joint_range[0]),
+                    limit_upper=(np.deg2rad(joint_range[1]) if is_angular and use_degrees else joint_range[1]),
+                    target_ke=parse_float(joint_attrib, "stiffness", stiffness),
+                    target_kd=parse_float(joint_attrib, "damping", damping),
+                    limit_ke=limit_ke,
+                    limit_kd=limit_kd,
+                    mode=mode,
+                )
+                if is_angular:
+                    angular_axes.append(ax)
+                else:
+                    linear_axes.append(ax)
 
         link = builder.add_body(
-            origin=wp.transform_identity(),  # will be evaluated in fk()
-            armature=joint_armature[0],
+            origin=wp.transform(body_pos, body_ori),  # will be evaluated in fk()
+            armature=joint_armature[0] if len(joint_armature) > 0 else armature,
             name=body_name,
         )
 
@@ -157,6 +305,7 @@ def parse_mjcf(
             else:
                 joint_type = wp.sim.JOINT_D6
 
+        joint_pos = joint_pos[0] if len(joint_pos) > 0 else (0.0, 0.0, 0.0)
         builder.add_joint(
             joint_type,
             parent,
@@ -164,21 +313,40 @@ def parse_mjcf(
             linear_axes,
             angular_axes,
             name="_".join(joint_name),
-            parent_xform=wp.transform(body_pos, body_ori),
-            # child_xform=wp.transform(joint_pos[0], wp.quat_identity()),
+            parent_xform=wp.transform(body_pos + joint_pos, body_ori),
+            child_xform=wp.transform(joint_pos, wp.quat_identity()),
         )
 
         # -----------------
         # add shapes
 
-        for geom in body.findall("geom"):
-            geom_name = geom.attrib["name"]
-            geom_type = geom.attrib["type"]
+        for geo_count, geom in enumerate(body.findall("geom")):
+            geom_defaults = defaults
+            if "class" in geom.attrib:
+                geom_class = geom.attrib["class"]
+                ignore_geom = False
+                for pattern in ignore_classes:
+                    if re.match(pattern, geom_class):
+                        ignore_geom = True
+                        break
+                if ignore_geom:
+                    continue
+                if geom_class in class_defaults:
+                    geom_defaults = merge_attrib(defaults, class_defaults[geom_class])
+            if "geom" in geom_defaults:
+                geom_attrib = merge_attrib(geom_defaults["geom"], geom.attrib)
+            else:
+                geom_attrib = geom.attrib
+
+            geom_name = geom_attrib.get("name", f"{body_name}_geom_{geo_count}")
+            geom_type = geom_attrib.get("type", "sphere")
+            if "mesh" in geom_attrib:
+                geom_type = "mesh"
 
-            geom_size = parse_vec(geom, "size", [1.0])
-            geom_pos = parse_vec(geom, "pos", (0.0, 0.0, 0.0))
-            geom_rot = parse_vec(geom, "quat", (0.0, 0.0, 0.0, 1.0))
-            geom_density = parse_float(geom, "density", density)
+            geom_size = parse_vec(geom_attrib, "size", [1.0, 1.0, 1.0]) * scale
+            geom_pos = parse_vec(geom_attrib, "pos", (0.0, 0.0, 0.0)) * scale
+            geom_rot = parse_orientation(geom_attrib)
+            geom_density = parse_float(geom_attrib, "density", density)
 
             if geom_type == "sphere":
                 builder.add_shape_sphere(
@@ -194,16 +362,36 @@ def parse_mjcf(
                     restitution=contact_restitution,
                 )
 
+            elif geom_type == "box":
+                builder.add_shape_box(
+                    link,
+                    pos=geom_pos,
+                    rot=geom_rot,
+                    hx=geom_size[0],
+                    hy=geom_size[1],
+                    hz=geom_size[2],
+                    density=geom_density,
+                    ke=contact_ke,
+                    kd=contact_kd,
+                    kf=contact_kf,
+                    mu=contact_mu,
+                    restitution=contact_restitution,
+                )
+
             elif geom_type == "mesh" and parse_meshes:
-                mesh, scale = parse_mesh(geom)
-                geom_size = tuple([scale * s for s in geom_size])
+                mesh, _ = parse_mesh(geom_attrib)
+                if "mesh" in defaults:
+                    mesh_scale = parse_vec(defaults["mesh"], "scale", [1.0, 1.0, 1.0])
+                else:
+                    mesh_scale = [1.0, 1.0, 1.0]
+                # as per the Mujoco XML reference, ignore geom size attribute
                 assert len(geom_size) == 3, "need to specify size for mesh geom"
                 builder.add_shape_mesh(
                     body=link,
                     pos=geom_pos,
                     rot=geom_rot,
                     mesh=mesh,
-                    scale=geom_size,
+                    scale=mesh_scale,
                     density=density,
                     ke=contact_ke,
                     kd=contact_kd,
@@ -212,32 +400,28 @@ def parse_mjcf(
                 )
 
             elif geom_type in {"capsule", "cylinder"}:
-                if "fromto" in geom.attrib:
-                    geom_fromto = parse_vec(geom, "fromto", (0.0, 0.0, 0.0, 1.0, 0.0, 0.0))
+                if "fromto" in geom_attrib:
+                    geom_fromto = parse_vec(geom_attrib, "fromto", (0.0, 0.0, 0.0, 1.0, 0.0, 0.0))
 
-                    start = geom_fromto[0:3]
-                    end = geom_fromto[3:6]
+                    start = wp.vec3(geom_fromto[0:3]) * scale
+                    end = wp.vec3(geom_fromto[3:6]) * scale
 
                     # compute rotation to align the Warp capsule (along x-axis), with mjcf fromto direction
                     axis = wp.normalize(end - start)
-                    angle = math.acos(np.dot(axis, (0.0, 1.0, 0.0)))
-                    axis = wp.normalize(np.cross(axis, (0.0, 1.0, 0.0)))
+                    angle = math.acos(wp.dot(axis, wp.vec3(0.0, 1.0, 0.0)))
+                    axis = wp.normalize(wp.cross(axis, wp.vec3(0.0, 1.0, 0.0)))
 
                     geom_pos = (start + end) * 0.5
                     geom_rot = wp.quat_from_axis_angle(axis, -angle)
 
                     geom_radius = geom_size[0]
-                    geom_height = np.linalg.norm(end - start) * 0.5
+                    geom_height = wp.length(end - start) * 0.5
+                    geom_up_axis = 1
 
                 else:
                     geom_radius = geom_size[0]
                     geom_height = geom_size[1]
-                    geom_pos = parse_vec(geom, "pos", (0.0, 0.0, 0.0))
-                    # orientation along the z axis by default
-                    axis = np.array((0.0, 0.0, 1.0))
-                    angle = math.acos(np.dot(axis, (0.0, 1.0, 0.0)))
-                    axis = wp.normalize(np.cross(axis, (0.0, 1.0, 0.0)))
-                    geom_rot = wp.quat_from_axis_angle(axis, -angle)
+                    geom_up_axis = up_axis
 
                 if geom_type == "cylinder":
                     builder.add_shape_cylinder(
@@ -252,6 +436,7 @@ def parse_mjcf(
                         kf=contact_kf,
                         mu=contact_mu,
                         restitution=contact_restitution,
+                        up_axis=geom_up_axis,
                     )
                 else:
                     builder.add_shape_capsule(
@@ -266,16 +451,17 @@ def parse_mjcf(
                         kf=contact_kf,
                         mu=contact_mu,
                         restitution=contact_restitution,
+                        up_axis=geom_up_axis,
                     )
 
             else:
-                print("MJCF parsing issue: geom type", geom_type, "is unsupported")
+                print(f"MJCF parsing shape {geom_name} issue: geom type {geom_type} is unsupported")
 
         # -----------------
         # recurse
 
         for child in body.findall("body"):
-            parse_body(child, link)
+            parse_body(child, link, defaults)
 
     # -----------------
     # start articulation
@@ -284,8 +470,10 @@ def parse_mjcf(
     builder.add_articulation()
 
     world = root.find("worldbody")
+    world_class = get_class(world)
+    world_defaults = merge_attrib(class_defaults["__all__"], class_defaults.get(world_class, {}))
     for body in world.findall("body"):
-        parse_body(body, -1)
+        parse_body(body, -1, world_defaults)
 
     end_shape_count = len(builder.shape_geo_type)
 
@@ -293,3 +481,6 @@ def parse_mjcf(
         for i in range(start_shape_count, end_shape_count):
             for j in range(i + 1, end_shape_count):
                 builder.shape_collision_filter_pairs.add((i, j))
+
+    if collapse_fixed_joints:
+        builder.collapse_fixed_joints()