Implement dual-arm teleoperation with RM75 QP controller and MuJoCo backend
This commit is contained in:
413
ik_qp/src/rm75_ik/mujoco_model.py
Normal file
413
ik_qp/src/rm75_ik/mujoco_model.py
Normal file
@ -0,0 +1,413 @@
|
||||
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
|
||||
Reference in New Issue
Block a user