414 lines
14 KiB
Python
414 lines
14 KiB
Python
from __future__ import annotations
|
|
|
|
import re
|
|
import xml.etree.ElementTree as ET
|
|
from pathlib import Path
|
|
from typing import Dict, Iterable, Optional, Tuple
|
|
|
|
import numpy as np
|
|
import pinocchio as pin
|
|
|
|
from .dual_arm import default_dual_source_path, load_dual_arm_mounts
|
|
from .kinematics import default_urdf_path
|
|
from .types import physical_joint_limits
|
|
|
|
|
|
def _numbers(text: str, expected: int) -> np.ndarray:
|
|
values = np.fromstring(text, sep=" ", dtype=float)
|
|
if values.shape != (expected,):
|
|
raise ValueError(f"expected {expected} numeric values, got {text!r}")
|
|
return values
|
|
|
|
|
|
def _format(values: Iterable[float]) -> str:
|
|
return " ".join(f"{float(value):.12g}" for value in values)
|
|
|
|
|
|
def _quaternion_from_rpy(rpy: np.ndarray) -> np.ndarray:
|
|
quaternion = pin.Quaternion(pin.rpy.rpyToMatrix(*rpy))
|
|
return np.array([quaternion.w, quaternion.x, quaternion.y, quaternion.z])
|
|
|
|
|
|
def _se3_attributes(transform: pin.SE3) -> Dict[str, str]:
|
|
quaternion = pin.Quaternion(transform.rotation)
|
|
return {
|
|
"pos": _format(transform.translation),
|
|
"quat": _format([quaternion.w, quaternion.x, quaternion.y, quaternion.z]),
|
|
}
|
|
|
|
|
|
def _origin_attributes(parent: ET.Element) -> Dict[str, str]:
|
|
origin = parent.find("origin")
|
|
if origin is None:
|
|
return {"pos": "0 0 0", "quat": "1 0 0 0"}
|
|
xyz = _numbers(origin.get("xyz", "0 0 0"), 3)
|
|
rpy = _numbers(origin.get("rpy", "0 0 0"), 3)
|
|
return {"pos": _format(xyz), "quat": _format(_quaternion_from_rpy(rpy))}
|
|
|
|
|
|
def _mesh_asset_name(filename: str, prefix: str) -> str:
|
|
stem = Path(filename).stem.lower()
|
|
return f"{prefix}_{re.sub(r'[^a-z0-9_]+', '_', stem)}"
|
|
|
|
|
|
def _dual_obj_directory(dual_source_path: Path) -> Path:
|
|
source_candidate = dual_source_path.parent / "dual_arm_obj"
|
|
if source_candidate.is_dir():
|
|
return source_candidate
|
|
raise FileNotFoundError(
|
|
f"dual-arm OBJ directory was not found beside {dual_source_path}"
|
|
)
|
|
|
|
|
|
def _single_mesh_directory(single_urdf_path: Path) -> Path:
|
|
candidates = (
|
|
single_urdf_path.parent / "meshes",
|
|
single_urdf_path.parent / "rm75_meshes",
|
|
)
|
|
for candidate in candidates:
|
|
if candidate.is_dir():
|
|
return candidate
|
|
raise FileNotFoundError(
|
|
f"RM75-B STL mesh directory was not found beside {single_urdf_path}"
|
|
)
|
|
|
|
|
|
def _add_inertial(body: ET.Element, link: ET.Element) -> None:
|
|
inertial = link.find("inertial")
|
|
if inertial is None:
|
|
raise ValueError(f"link {link.get('name')!r} has no inertial data")
|
|
mass_element = inertial.find("mass")
|
|
inertia_element = inertial.find("inertia")
|
|
if mass_element is None or inertia_element is None:
|
|
raise ValueError(f"link {link.get('name')!r} has incomplete inertial data")
|
|
origin = inertial.find("origin")
|
|
xyz = _numbers(origin.get("xyz", "0 0 0"), 3) if origin is not None else np.zeros(3)
|
|
rpy = _numbers(origin.get("rpy", "0 0 0"), 3) if origin is not None else np.zeros(3)
|
|
inertia_matrix = np.array(
|
|
[
|
|
[float(inertia_element.get("ixx")), float(inertia_element.get("ixy")), float(inertia_element.get("ixz"))],
|
|
[float(inertia_element.get("ixy")), float(inertia_element.get("iyy")), float(inertia_element.get("iyz"))],
|
|
[float(inertia_element.get("ixz")), float(inertia_element.get("iyz")), float(inertia_element.get("izz"))],
|
|
]
|
|
)
|
|
rotation = pin.rpy.rpyToMatrix(*rpy)
|
|
inertia_matrix = rotation @ inertia_matrix @ rotation.T
|
|
full_inertia = [
|
|
inertia_matrix[0, 0],
|
|
inertia_matrix[1, 1],
|
|
inertia_matrix[2, 2],
|
|
inertia_matrix[0, 1],
|
|
inertia_matrix[0, 2],
|
|
inertia_matrix[1, 2],
|
|
]
|
|
ET.SubElement(
|
|
body,
|
|
"inertial",
|
|
{
|
|
"pos": _format(xyz),
|
|
"mass": mass_element.get("value", "0"),
|
|
"fullinertia": _format(full_inertia),
|
|
},
|
|
)
|
|
|
|
|
|
def _add_link_visual(
|
|
body: ET.Element,
|
|
link: ET.Element,
|
|
arm: str,
|
|
single_mesh_keys: Dict[str, str],
|
|
) -> None:
|
|
visual = link.find("visual")
|
|
if visual is None:
|
|
return
|
|
mesh = visual.find("geometry/mesh")
|
|
if mesh is None:
|
|
return
|
|
filename = Path(mesh.get("filename", "")).name
|
|
try:
|
|
mesh_name = single_mesh_keys[filename]
|
|
except KeyError as exc:
|
|
raise ValueError(f"missing packaged single-arm mesh {filename!r}") from exc
|
|
ET.SubElement(
|
|
body,
|
|
"geom",
|
|
{
|
|
"name": f"{arm}_{link.get('name')}_visual",
|
|
"type": "mesh",
|
|
"mesh": mesh_name,
|
|
"material": f"{arm}_arm",
|
|
"contype": "0",
|
|
"conaffinity": "0",
|
|
"group": "1",
|
|
**_origin_attributes(visual),
|
|
},
|
|
)
|
|
|
|
|
|
def _add_arm(
|
|
worldbody: ET.Element,
|
|
arm: str,
|
|
mount: pin.SE3,
|
|
single_root: ET.Element,
|
|
single_mesh_keys: Dict[str, str],
|
|
gripper_mesh_name: str,
|
|
) -> None:
|
|
links = {link.get("name"): link for link in single_root.findall("link")}
|
|
joints_by_parent: Dict[str, list[ET.Element]] = {}
|
|
for joint in single_root.findall("joint"):
|
|
parent_element = joint.find("parent")
|
|
if parent_element is None:
|
|
continue
|
|
joints_by_parent.setdefault(parent_element.get("link", ""), []).append(joint)
|
|
|
|
limits = physical_joint_limits()
|
|
base_link = links["base_link"]
|
|
base_body = ET.SubElement(
|
|
worldbody,
|
|
"body",
|
|
{"name": f"{arm}_base_link", **_se3_attributes(mount)},
|
|
)
|
|
_add_link_visual(base_body, base_link, arm, single_mesh_keys)
|
|
|
|
def append_children(parent_body: ET.Element, parent_link_name: str) -> None:
|
|
for joint in joints_by_parent.get(parent_link_name, []):
|
|
child_element = joint.find("child")
|
|
if child_element is None:
|
|
raise ValueError(f"joint {joint.get('name')!r} has no child")
|
|
child_name = child_element.get("link", "")
|
|
child_link = links[child_name]
|
|
body = ET.SubElement(
|
|
parent_body,
|
|
"body",
|
|
{
|
|
"name": f"{arm}_{child_name}",
|
|
**_origin_attributes(joint),
|
|
},
|
|
)
|
|
joint_index = int(joint.get("name", "joint_0").split("_")[-1]) - 1
|
|
axis_element = joint.find("axis")
|
|
axis = "0 0 1" if axis_element is None else axis_element.get("xyz", "0 0 1")
|
|
ET.SubElement(
|
|
body,
|
|
"joint",
|
|
{
|
|
"name": f"{arm}_joint_{joint_index + 1}",
|
|
"type": "hinge",
|
|
"axis": axis,
|
|
"range": _format(
|
|
[limits.lower[joint_index], limits.upper[joint_index]]
|
|
),
|
|
"limited": "true",
|
|
"damping": "0",
|
|
},
|
|
)
|
|
_add_inertial(body, child_link)
|
|
_add_link_visual(body, child_link, arm, single_mesh_keys)
|
|
if child_name == "link_7":
|
|
ET.SubElement(
|
|
body,
|
|
"site",
|
|
{
|
|
"name": f"{arm}_flange",
|
|
"type": "sphere",
|
|
"size": "0.008",
|
|
"rgba": "0.1 0.9 0.2 1" if arm == "left" else "0.1 0.5 0.95 1",
|
|
"group": "2",
|
|
},
|
|
)
|
|
ET.SubElement(
|
|
body,
|
|
"geom",
|
|
{
|
|
"name": f"{arm}_gripper_visual",
|
|
"type": "mesh",
|
|
"mesh": gripper_mesh_name,
|
|
"pos": "0 0 0.092",
|
|
"quat": _format(_quaternion_from_rpy(np.array([-np.pi, np.pi, -np.pi]))),
|
|
"material": f"{arm}_gripper",
|
|
"contype": "0",
|
|
"conaffinity": "0",
|
|
"group": "1",
|
|
},
|
|
)
|
|
append_children(body, child_name)
|
|
|
|
append_children(base_body, "base_link")
|
|
|
|
|
|
def build_normalized_dual_mjcf(
|
|
single_urdf_path: Optional[Path | str] = None,
|
|
dual_source_path: Optional[Path | str] = None,
|
|
) -> Tuple[str, Dict[str, bytes]]:
|
|
"""Build a canonical 14-DOF dual-arm MJCF and its in-memory mesh assets."""
|
|
|
|
single_path = (
|
|
Path(single_urdf_path) if single_urdf_path is not None else default_urdf_path()
|
|
)
|
|
dual_path = (
|
|
Path(dual_source_path)
|
|
if dual_source_path is not None
|
|
else default_dual_source_path()
|
|
)
|
|
single_root = ET.parse(single_path).getroot()
|
|
dual_root = ET.parse(dual_path).getroot()
|
|
single_mesh_dir = _single_mesh_directory(single_path)
|
|
dual_obj_dir = _dual_obj_directory(dual_path)
|
|
|
|
mujoco_root = ET.Element("mujoco", {"model": "rm75_normalized_dual_stage2"})
|
|
ET.SubElement(
|
|
mujoco_root,
|
|
"compiler",
|
|
{
|
|
"angle": "radian",
|
|
"autolimits": "true",
|
|
"inertiafromgeom": "false",
|
|
"balanceinertia": "true",
|
|
},
|
|
)
|
|
ET.SubElement(
|
|
mujoco_root,
|
|
"option",
|
|
{"timestep": "0.002", "gravity": "0 0 -9.81", "integrator": "implicitfast"},
|
|
)
|
|
ET.SubElement(mujoco_root, "statistic", {"center": "0 0 0.65", "extent": "1.4"})
|
|
visual = ET.SubElement(mujoco_root, "visual")
|
|
ET.SubElement(
|
|
visual,
|
|
"global",
|
|
{
|
|
"azimuth": "90",
|
|
"elevation": "-18",
|
|
"offwidth": "1280",
|
|
"offheight": "720",
|
|
},
|
|
)
|
|
ET.SubElement(visual, "rgba", {"haze": "0.15 0.18 0.2 1"})
|
|
|
|
assets_element = ET.SubElement(mujoco_root, "asset")
|
|
ET.SubElement(assets_element, "material", {"name": "left_arm", "rgba": "0.82 0.84 0.86 1"})
|
|
ET.SubElement(assets_element, "material", {"name": "right_arm", "rgba": "0.7 0.74 0.78 1"})
|
|
ET.SubElement(assets_element, "material", {"name": "left_gripper", "rgba": "0.15 0.75 0.85 1"})
|
|
ET.SubElement(assets_element, "material", {"name": "right_gripper", "rgba": "0.9 0.48 0.18 1"})
|
|
ET.SubElement(assets_element, "material", {"name": "platform", "rgba": "0.28 0.31 0.34 1"})
|
|
|
|
assets: Dict[str, bytes] = {}
|
|
single_mesh_keys: Dict[str, str] = {}
|
|
for mesh_path in sorted(single_mesh_dir.glob("*.STL")):
|
|
asset_key = f"rm75_meshes/{mesh_path.name}"
|
|
mesh_name = _mesh_asset_name(mesh_path.name, "rm75")
|
|
assets[asset_key] = mesh_path.read_bytes()
|
|
single_mesh_keys[mesh_path.name] = mesh_name
|
|
ET.SubElement(assets_element, "mesh", {"name": mesh_name, "file": asset_key})
|
|
|
|
dual_mesh_names: Dict[str, str] = {}
|
|
for mesh_path in sorted(dual_obj_dir.glob("*.obj")):
|
|
asset_key = f"dual_arm_obj/{mesh_path.name}"
|
|
mesh_name = _mesh_asset_name(mesh_path.name, "dual")
|
|
assets[asset_key] = mesh_path.read_bytes()
|
|
dual_mesh_names[mesh_path.name] = mesh_name
|
|
ET.SubElement(assets_element, "mesh", {"name": mesh_name, "file": asset_key})
|
|
if len(dual_mesh_names) != 19:
|
|
raise ValueError(f"expected 19 dual-arm OBJ assets, found {len(dual_mesh_names)}")
|
|
|
|
worldbody = ET.SubElement(mujoco_root, "worldbody")
|
|
ET.SubElement(
|
|
worldbody,
|
|
"light",
|
|
{
|
|
"name": "key_light",
|
|
"pos": "0 -1.2 2.8",
|
|
"dir": "0 0.4 -1",
|
|
"diffuse": "0.9 0.9 0.9",
|
|
},
|
|
)
|
|
ET.SubElement(
|
|
worldbody,
|
|
"camera",
|
|
{
|
|
"name": "overview",
|
|
"pos": "0 -2.6 1.25",
|
|
"xyaxes": "1 0 0 0 0.3 0.953939",
|
|
"fovy": "45",
|
|
},
|
|
)
|
|
ET.SubElement(
|
|
worldbody,
|
|
"geom",
|
|
{
|
|
"name": "floor",
|
|
"type": "plane",
|
|
"size": "2.5 2.5 0.05",
|
|
"rgba": "0.12 0.14 0.16 1",
|
|
"contype": "0",
|
|
"conaffinity": "0",
|
|
"group": "0",
|
|
},
|
|
)
|
|
|
|
robot_base = dual_root.find("./link[@name='robot_base']")
|
|
if robot_base is None:
|
|
raise ValueError("dual-arm source URDF has no robot_base link")
|
|
for visual_element in robot_base.findall("visual"):
|
|
mesh_element = visual_element.find("geometry/mesh")
|
|
if mesh_element is None:
|
|
continue
|
|
filename = Path(mesh_element.get("filename", "")).name
|
|
ET.SubElement(
|
|
worldbody,
|
|
"geom",
|
|
{
|
|
"name": f"platform_{visual_element.get('name', filename)}",
|
|
"type": "mesh",
|
|
"mesh": dual_mesh_names[filename],
|
|
"material": "platform",
|
|
"contype": "0",
|
|
"conaffinity": "0",
|
|
"group": "1",
|
|
**_origin_attributes(visual_element),
|
|
},
|
|
)
|
|
|
|
mounts = load_dual_arm_mounts(dual_path)
|
|
_add_arm(
|
|
worldbody,
|
|
"left",
|
|
mounts.left_base,
|
|
single_root,
|
|
single_mesh_keys,
|
|
dual_mesh_names["dual_arm_gripper1_vis_1.obj"],
|
|
)
|
|
_add_arm(
|
|
worldbody,
|
|
"right",
|
|
mounts.right_base,
|
|
single_root,
|
|
single_mesh_keys,
|
|
dual_mesh_names["dual_arm_gripper2_vis_1.obj"],
|
|
)
|
|
for arm, color in (
|
|
("left", "0.95 0.12 0.12 0.8"),
|
|
("right", "0.95 0.75 0.08 0.8"),
|
|
):
|
|
marker = ET.SubElement(
|
|
worldbody,
|
|
"body",
|
|
{"name": f"{arm}_target_marker", "mocap": "true", "pos": "0 0 0"},
|
|
)
|
|
ET.SubElement(
|
|
marker,
|
|
"geom",
|
|
{
|
|
"name": f"{arm}_target_marker_geom",
|
|
"type": "sphere",
|
|
"size": "0.018",
|
|
"rgba": color,
|
|
"contype": "0",
|
|
"conaffinity": "0",
|
|
"group": "2",
|
|
},
|
|
)
|
|
|
|
ET.indent(mujoco_root, space=" ")
|
|
return ET.tostring(mujoco_root, encoding="unicode"), assets
|