# Copyright (c) 2022 NVIDIA CORPORATION. All rights reserved. # NVIDIA CORPORATION and its licensors retain all intellectual property # and proprietary rights in and to this software, related documentation # and any modifications thereto. Any use, reproduction, disclosure or # distribution of this software and related documentation without an express # license agreement from NVIDIA CORPORATION is strictly prohibited. import math import os import re import xml.etree.ElementTree as ET import numpy as np import warp as wp def parse_mjcf( mjcf_filename, builder, xform=wp.transform(), density=1000.0, stiffness=0.0, damping=0.0, contact_ke=1000.0, contact_kd=100.0, contact_kf=100.0, contact_mu=0.5, contact_restitution=0.5, limit_ke=100.0, limit_kd=10.0, scale=1.0, armature=0.0, armature_scale=1.0, parse_meshes=True, enable_self_collisions=False, up_axis="Z", ignore_classes=[], collapse_fixed_joints=False, ): """ 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 `): 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() 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(attrib, key, default): if key in attrib: out = np.fromstring(attrib[key], sep=" ", dtype=np.float32) else: 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 = mesh_assets[geom["mesh"]] m = trimesh.load(stl_file) for v in m.vertices: 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 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: 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 = [] joint_pos = [] linear_axes = [] angular_axes = [] joint_type = None 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(body_pos, body_ori), # will be evaluated in fk() armature=joint_armature[0] if len(joint_armature) > 0 else armature, name=body_name, ) if joint_type is None: if len(linear_axes) == 0: if len(angular_axes) == 0: joint_type = wp.sim.JOINT_FIXED elif len(angular_axes) == 1: joint_type = wp.sim.JOINT_REVOLUTE elif len(angular_axes) == 2: joint_type = wp.sim.JOINT_UNIVERSAL elif len(angular_axes) == 3: joint_type = wp.sim.JOINT_COMPOUND elif len(linear_axes) == 1 and len(angular_axes) == 0: joint_type = wp.sim.JOINT_PRISMATIC 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, link, linear_axes, angular_axes, name="_".join(joint_name), parent_xform=wp.transform(body_pos + joint_pos, body_ori), child_xform=wp.transform(joint_pos, wp.quat_identity()), ) # ----------------- # add shapes 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_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( link, pos=geom_pos, rot=geom_rot, radius=geom_size[0], density=geom_density, ke=contact_ke, kd=contact_kd, kf=contact_kf, mu=contact_mu, 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, _ = 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=mesh_scale, density=density, ke=contact_ke, kd=contact_kd, kf=contact_kf, mu=contact_mu, ) elif geom_type in {"capsule", "cylinder"}: 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 = 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(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 = wp.length(end - start) * 0.5 geom_up_axis = 1 else: geom_radius = geom_size[0] geom_height = geom_size[1] geom_up_axis = up_axis if geom_type == "cylinder": builder.add_shape_cylinder( link, pos=geom_pos, rot=geom_rot, radius=geom_radius, half_height=geom_height, density=density, ke=contact_ke, kd=contact_kd, kf=contact_kf, mu=contact_mu, restitution=contact_restitution, up_axis=geom_up_axis, ) else: builder.add_shape_capsule( link, pos=geom_pos, rot=geom_rot, radius=geom_radius, half_height=geom_height, density=density, ke=contact_ke, kd=contact_kd, kf=contact_kf, mu=contact_mu, restitution=contact_restitution, up_axis=geom_up_axis, ) else: 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, defaults) # ----------------- # start articulation start_shape_count = len(builder.shape_geo_type) 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, world_defaults) end_shape_count = len(builder.shape_geo_type) if not enable_self_collisions: 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()