bb87928bce
Upstream: https://github.com/Rhoban/onshape-to-robot Upstream-Commit: 80e710700aac9573a2230f74f7ce9e094833a0bc Upstream-Branch: master
415 lines
15 KiB
Python
415 lines
15 KiB
Python
from __future__ import annotations
|
|
import numpy as np
|
|
import os
|
|
import fnmatch
|
|
from .message import success, warning, info
|
|
from .robot import Robot, Link, Part, Joint, Closure
|
|
from .config import Config
|
|
from .geometry import Box, Cylinder, Sphere, Mesh, Shape
|
|
from .exporter import Exporter
|
|
from .exporter_utils import xml_escape, rotation_matrix_to_rpy
|
|
from transforms3d.quaternions import mat2quat
|
|
|
|
|
|
class ExporterMuJoCo(Exporter):
|
|
def __init__(self, config: Config | None = None):
|
|
super().__init__()
|
|
self.config: Config = config
|
|
|
|
self.no_dynamics: bool = False
|
|
self.additional_xml: str = ""
|
|
self.meshes: list = []
|
|
self.materials: dict = {}
|
|
|
|
if config is not None:
|
|
self.equalities = self.config.get("equalities", {})
|
|
self.no_dynamics = config.no_dynamics
|
|
additional_xml_file = config.get("additional_xml", None, required=False)
|
|
if isinstance(additional_xml_file, str):
|
|
self.add_additional_xml(additional_xml_file)
|
|
elif isinstance(additional_xml_file, list):
|
|
for filename in additional_xml_file:
|
|
self.add_additional_xml(filename)
|
|
|
|
def add_additional_xml(self, xml_file: str):
|
|
self.additional_xml += f"<!-- Additional {xml_file} -->"
|
|
with open(self.config.output_directory + "/" + xml_file, "r") as file:
|
|
self.additional_xml += file.read()
|
|
|
|
def append(self, line: str):
|
|
self.xml += line
|
|
|
|
def build(self, robot: Robot):
|
|
self.xml = ""
|
|
self.append('<?xml version="1.0" ?>')
|
|
self.append("<!-- Generated using onshape-to-robot -->")
|
|
if self.config:
|
|
self.append(f"<!-- Onshape {self.config.printable_version()} -->")
|
|
self.append(f'<mujoco model="{robot.name}">')
|
|
self.append(
|
|
f'<compiler angle="radian" meshdir="{self.config.assets_directory}" autolimits="true" />'
|
|
)
|
|
|
|
# Boilerplate
|
|
self.default_class = robot.name
|
|
self.append("<default>")
|
|
self.append(f'<default class="{self.default_class}">')
|
|
self.append('<joint frictionloss="0.1" armature="0.005"/>')
|
|
self.append('<position kp="50" dampratio="1" />')
|
|
self.append('<default class="visual">')
|
|
self.append('<geom type="mesh" contype="0" conaffinity="0" group="2"/>')
|
|
self.append("</default>")
|
|
self.append('<default class="collision">')
|
|
self.append('<geom group="3"/>')
|
|
self.append("</default>")
|
|
self.append("</default>")
|
|
self.append("</default>")
|
|
|
|
if self.additional_xml:
|
|
self.append(self.additional_xml)
|
|
|
|
# Adding robot links
|
|
self.append("<worldbody>")
|
|
for base_link in robot.base_links:
|
|
self.add_link(robot, base_link)
|
|
self.append("</worldbody>")
|
|
|
|
# Asset (mesh & materials)
|
|
self.append("<asset>")
|
|
for mesh_file in set(self.meshes):
|
|
self.append(f'<mesh file="{mesh_file}" />')
|
|
for material_name, color in self.materials.items():
|
|
color_str = "%g %g %g %g" % self.config.round(tuple(color))
|
|
self.append(f'<material name="{material_name}" rgba="{color_str}" />')
|
|
self.append("</asset>")
|
|
|
|
# Adding actuators
|
|
self.add_actuators(robot)
|
|
|
|
# Adding equalities (loop closure)
|
|
self.add_equalities(robot)
|
|
|
|
self.append("</mujoco>")
|
|
|
|
return self.xml
|
|
|
|
def add_actuators(self, robot: Robot):
|
|
self.append("<actuator>")
|
|
|
|
for joint in robot.joints:
|
|
if joint.joint_type == "fixed":
|
|
continue
|
|
|
|
# Suppose joints with relation equality is not actuated, unless specified
|
|
guess_actuated = joint.relation is None
|
|
|
|
if (
|
|
joint.properties.get("actuated", guess_actuated)
|
|
and joint.joint_type != Joint.BALL
|
|
):
|
|
type = joint.properties.get("type", "position")
|
|
actuator_class = joint.properties.get("class", self.default_class)
|
|
actuator: str = f'<{type} class="{actuator_class}" name="{joint.name}" joint="{joint.name}" '
|
|
|
|
for key in "kp", "kv", "dampratio":
|
|
if key in joint.properties:
|
|
actuator += f'{key}="{joint.properties[key]}" '
|
|
|
|
if "forcerange" in joint.properties:
|
|
actuator += f'forcerange="-{joint.properties["forcerange"]} {joint.properties["forcerange"]}" '
|
|
|
|
joint_limits = joint.properties.get("limits", joint.limits)
|
|
limits_are_set = joint.properties.get("limits", False) != False
|
|
if joint_limits and (type == "position" or limits_are_set):
|
|
if joint.properties.get("range", True) and type == "position":
|
|
actuator += f'inheritrange="1" '
|
|
else:
|
|
actuator += f'ctrlrange="{joint_limits[0]} {joint_limits[1]}" '
|
|
|
|
actuator += "/>"
|
|
self.append(actuator)
|
|
|
|
self.append("</actuator>")
|
|
|
|
def get_equality_attributes(self, closure: Closure) -> str:
|
|
all_attributes = {}
|
|
for name, attributes in self.equalities.items():
|
|
if fnmatch.fnmatch(closure.frame1, name) and fnmatch.fnmatch(
|
|
closure.frame2, name
|
|
):
|
|
all_attributes.update(attributes)
|
|
|
|
if len(all_attributes) > 0:
|
|
return (
|
|
" ".join([f'{key}="{value}"' for key, value in all_attributes.items()])
|
|
+ " "
|
|
)
|
|
|
|
return ""
|
|
|
|
def add_equalities(self, robot: Robot):
|
|
self.append("<equality>")
|
|
for closure in robot.closures:
|
|
attributes = self.get_equality_attributes(closure)
|
|
|
|
if closure.closure_type == Closure.FIXED:
|
|
self.append(
|
|
f'<weld site1="{closure.frame1}" site2="{closure.frame2}" {attributes}/>'
|
|
)
|
|
elif closure.closure_type == Closure.REVOLUTE:
|
|
self.append(
|
|
f'<connect site1="{closure.frame1}" site2="{closure.frame2}" {attributes}/>'
|
|
)
|
|
elif closure.closure_type == Closure.BALL:
|
|
self.append(
|
|
f'<connect site1="{closure.frame1}" site2="{closure.frame2}" {attributes}/>'
|
|
)
|
|
else:
|
|
print(
|
|
warning(
|
|
f"Closure type: {closure.closure_type} is not supported with MuJoCo equality constraints"
|
|
)
|
|
)
|
|
|
|
for joint in robot.joints:
|
|
if joint.relation is not None:
|
|
self.append(
|
|
f'<joint joint1="{joint.name}" joint2="{joint.relation.source_joint}" polycoef="0 {joint.relation.ratio} 0 0 0" />'
|
|
)
|
|
|
|
self.append("</equality>")
|
|
|
|
def add_inertial(self, mass: float, com: np.ndarray, inertia: np.ndarray):
|
|
# Ensuring epsilon masses and inertias
|
|
mass = max(1e-9, mass)
|
|
inertia[0, 0] = max(1e-9, inertia[0, 0])
|
|
inertia[1, 1] = max(1e-9, inertia[1, 1])
|
|
inertia[2, 2] = max(1e-9, inertia[2, 2])
|
|
|
|
# Populating body inertial properties
|
|
# https://mujoco.readthedocs.io/en/stable/XMLreference.html#body-inertial
|
|
inertial: str = "<inertial "
|
|
inertial += 'pos="%g %g %g" ' % self.config.round(tuple(com))
|
|
inertial += 'mass="%g" ' % self.config.round(mass)
|
|
inertial += 'fullinertia="%g %g %g %g %g %g" ' % self.config.round(
|
|
(
|
|
inertia[0, 0],
|
|
inertia[1, 1],
|
|
inertia[2, 2],
|
|
inertia[0, 1],
|
|
inertia[0, 2],
|
|
inertia[1, 2],
|
|
)
|
|
)
|
|
inertial += " />"
|
|
self.append(inertial)
|
|
|
|
def add_mesh(self, part: Part, class_: str, T_world_link: np.ndarray, mesh: Mesh):
|
|
"""
|
|
Add a mesh node (e.g. STL) to the MuJoCo file
|
|
"""
|
|
# Retrieving mesh file and material name
|
|
mesh_file = os.path.relpath(self.config.output_directory + "/" + mesh.filename, self.config.asset_path(""))
|
|
mesh_file_no_ext = ".".join(os.path.basename(mesh_file).split(".")[:-1])
|
|
material_name = mesh_file_no_ext + "_material"
|
|
|
|
# Relative frame
|
|
T_link_part = np.linalg.inv(T_world_link) @ part.T_world_part
|
|
|
|
# Adding the geom node
|
|
geom = f'<geom type="mesh" class="{class_}" '
|
|
geom += self.pos_quat(T_link_part) + " "
|
|
geom += f'mesh="{xml_escape(mesh_file_no_ext)}" '
|
|
geom += f'material="{xml_escape(material_name)}" '
|
|
|
|
# Apply properties based on class (visual or collision)
|
|
properties = mesh.visual_properties if class_ == "visual" else mesh.collision_properties
|
|
for key, value in properties.items():
|
|
geom += f'{key}="{xml_escape(str(value))}" '
|
|
|
|
geom += " />"
|
|
|
|
# Adding the mesh and material to appear in the assets section
|
|
self.meshes.append(mesh_file)
|
|
self.materials[material_name] = mesh.color
|
|
|
|
self.append(geom)
|
|
|
|
def add_shape(
|
|
self, part: Part, class_: str, T_world_link: np.ndarray, shape: Shape
|
|
):
|
|
"""
|
|
Add pure shape geometry.
|
|
"""
|
|
geom = f'<geom class="{class_}" '
|
|
|
|
T_link_shape = (
|
|
np.linalg.inv(T_world_link) @ part.T_world_part @ shape.T_part_shape
|
|
)
|
|
geom += self.pos_quat(T_link_shape) + " "
|
|
|
|
if isinstance(shape, Box):
|
|
geom += 'type="box" size="%g %g %g" ' % self.config.round(tuple(shape.size / 2))
|
|
elif isinstance(shape, Cylinder):
|
|
geom += 'type="cylinder" size="%g %g" ' % (
|
|
shape.radius,
|
|
shape.length / 2,
|
|
)
|
|
elif isinstance(shape, Sphere):
|
|
geom += 'type="sphere" size="%g" ' % self.config.round(shape.radius)
|
|
|
|
if class_ == "visual":
|
|
material_name = f"{part.name}_material"
|
|
self.materials[material_name] = shape.color
|
|
geom += f'material="{xml_escape(material_name)}" '
|
|
|
|
# Apply properties based on class (visual or collision)
|
|
properties = shape.visual_properties if class_ == "visual" else shape.collision_properties
|
|
for key, value in properties.items():
|
|
geom += f'{key}="{xml_escape(str(value))}" '
|
|
|
|
geom += " />"
|
|
self.append(geom)
|
|
|
|
def add_geometries(self, part: Part, T_world_link: np.ndarray):
|
|
"""
|
|
Add a part geometries
|
|
"""
|
|
for shape in part.shapes:
|
|
if shape.visual:
|
|
self.add_shape(part, "visual", T_world_link, shape)
|
|
if shape.collision:
|
|
self.add_shape(part, "collision", T_world_link, shape)
|
|
|
|
for mesh in part.meshes:
|
|
if mesh.visual:
|
|
self.add_mesh(part, "visual", T_world_link, mesh)
|
|
if mesh.collision:
|
|
self.add_mesh(part, "collision", T_world_link, mesh)
|
|
|
|
def add_joint(self, joint: Joint):
|
|
self.append(f"<!-- Joint from {joint.parent.name} to {joint.child.name} -->")
|
|
if joint.joint_type == "fixed":
|
|
self.append(f'<!-- Joint is "fixed", skipping it -->')
|
|
return
|
|
|
|
joint_xml: str = "<joint "
|
|
joint_xml += 'axis="%g %g %g" ' % self.config.round(tuple(joint.axis))
|
|
joint_xml += f'name="{joint.name}" '
|
|
if joint.joint_type == Joint.REVOLUTE:
|
|
joint_xml += 'type="hinge" '
|
|
elif joint.joint_type == Joint.PRISMATIC:
|
|
joint_xml += 'type="slide" '
|
|
elif joint.joint_type == Joint.BALL:
|
|
joint_xml += 'type="ball" '
|
|
else:
|
|
print(warning(f"Unknown joint type: {joint.joint_type}"))
|
|
|
|
joint_limits = joint.properties.get("limits", joint.limits)
|
|
if joint_limits is not None and joint.properties.get("range", True):
|
|
joint_xml += f'range="{joint_limits[0]} {joint_limits[1]}" '
|
|
|
|
for key in (
|
|
"class",
|
|
"frictionloss",
|
|
"armature",
|
|
"damping",
|
|
"stiffness",
|
|
):
|
|
if key in joint.properties:
|
|
joint_xml += f'{key}="{joint.properties[key]}" '
|
|
|
|
joint_xml += " />"
|
|
self.append(joint_xml)
|
|
|
|
def add_frame(
|
|
self,
|
|
frame: str,
|
|
T_world_link: np.ndarray,
|
|
T_world_frame: np.ndarray,
|
|
group: int = 0,
|
|
):
|
|
self.append(f"<!-- Frame {frame} -->")
|
|
T_link_frame = np.linalg.inv(T_world_link) @ T_world_frame
|
|
|
|
site: str = f'<site group="{group}" name="{frame}" '
|
|
site += self.pos_quat(T_link_frame) + " "
|
|
site += " />"
|
|
self.append(site)
|
|
|
|
def add_link(
|
|
self,
|
|
robot: Robot,
|
|
link: Link,
|
|
parent_joint: Joint | None = None,
|
|
T_world_parent: np.ndarray = np.eye(4),
|
|
):
|
|
"""
|
|
Adds a link recursively to the URDF file
|
|
"""
|
|
if parent_joint is None:
|
|
T_world_link = np.eye(4)
|
|
else:
|
|
T_world_link = parent_joint.T_world_joint
|
|
|
|
childclass = ""
|
|
if parent_joint is None:
|
|
childclass = f'childclass="{self.default_class}" '
|
|
self.append(f"<!-- Link {link.name} -->")
|
|
T_parent_link = np.linalg.inv(T_world_parent) @ T_world_link
|
|
self.append(
|
|
f'<body name="{link.name}" {self.pos_quat(T_parent_link)} {childclass}>'
|
|
)
|
|
|
|
if parent_joint is None:
|
|
if not link.fixed:
|
|
self.append(f'<freejoint name="{link.name}_freejoint" />')
|
|
else:
|
|
self.add_joint(parent_joint)
|
|
|
|
# Adding inertial properties
|
|
mass, com, inertia = link.get_dynamics(T_world_link)
|
|
self.add_inertial(mass, com, inertia)
|
|
|
|
# Adding geometry objects
|
|
for part in link.parts:
|
|
self.append(f"<!-- Part {part.name} -->")
|
|
self.add_geometries(part, T_world_link)
|
|
|
|
# Adding frames attached to current link
|
|
for frame, T_world_frame in link.frames.items():
|
|
self.add_frame(frame, T_world_link, T_world_frame, group=3)
|
|
|
|
# Adding joints and children links
|
|
for joint in robot.get_link_joints(link):
|
|
self.add_link(robot, joint.child, joint, T_world_link)
|
|
|
|
self.append("</body>")
|
|
|
|
def pos_quat(self, matrix: np.ndarray) -> str:
|
|
"""
|
|
Turn a transformation matrix into 'pos="..." quat="..."' attributes
|
|
"""
|
|
pos = matrix[:3, 3]
|
|
quat = mat2quat(matrix[:3, :3])
|
|
xml = 'pos="%g %g %g" quat="%g %g %g %g"' % self.config.round((*pos, *quat))
|
|
|
|
return xml
|
|
|
|
def write_xml(self, robot: Robot, filename: str) -> str:
|
|
super().write_xml(robot, filename)
|
|
|
|
dirname = os.path.dirname(filename)
|
|
scene_filename = dirname + "/scene.xml"
|
|
if not os.path.exists(scene_filename):
|
|
scene_xml: str = (
|
|
os.path.dirname(os.path.realpath(__file__)) + "/assets/scene.xml"
|
|
)
|
|
scene_xml = open(scene_xml, "r").read()
|
|
scene_xml = scene_xml.format(robot_filename=os.path.basename(filename))
|
|
with open(scene_filename, "w") as file:
|
|
file.write(scene_xml)
|
|
print(success(f"* Writing scene.xml"))
|
|
else:
|
|
print(info(f"* scene.xml already exists, not over-writing it"))
|