327 lines
13 KiB
Python
327 lines
13 KiB
Python
#!/usr/bin/env python3
|
|
"""Static readiness checks for the Gen2 IsaacLab locomotion asset."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import ast
|
|
import math
|
|
import re
|
|
import xml.etree.ElementTree as ET
|
|
from dataclasses import dataclass
|
|
from pathlib import Path
|
|
|
|
import numpy as np
|
|
|
|
|
|
REPO_ROOT = Path(__file__).resolve().parents[1]
|
|
URDF_PATH = REPO_ROOT / "source" / "gen2_lab" / "assets" / "robot_simplified_collision.urdf"
|
|
GEN2_ROBOT_CFG = REPO_ROOT / "source" / "engineai_lab" / "robots" / "gen2.py"
|
|
GEN2_ENV_CFG = REPO_ROOT / "source" / "engineai_lab" / "tasks" / "velocity" / "config" / "gen2" / "flat_env_cfg.py"
|
|
FOOT_LINKS = ("left_leg_link_6", "right_leg_link_6")
|
|
|
|
|
|
@dataclass
|
|
class Box:
|
|
name: str
|
|
center: np.ndarray
|
|
axes: np.ndarray
|
|
half: np.ndarray
|
|
vertices: np.ndarray
|
|
|
|
|
|
def _parse_floats(text: str | None, default: tuple[float, ...]) -> tuple[float, ...]:
|
|
if not text:
|
|
return default
|
|
return tuple(float(value) for value in text.split())
|
|
|
|
|
|
def _rpy_matrix(rpy: tuple[float, float, float]) -> np.ndarray:
|
|
roll, pitch, yaw = rpy
|
|
cr, sr = math.cos(roll), math.sin(roll)
|
|
cp, sp = math.cos(pitch), math.sin(pitch)
|
|
cy, sy = math.cos(yaw), math.sin(yaw)
|
|
return np.array(
|
|
[
|
|
[cy * cp, cy * sp * sr - sy * cr, cy * sp * cr + sy * sr],
|
|
[sy * cp, sy * sp * sr + cy * cr, sy * sp * cr - cy * sr],
|
|
[-sp, cp * sr, cp * cr],
|
|
],
|
|
dtype=float,
|
|
)
|
|
|
|
|
|
def _axis_angle(axis: tuple[float, float, float], angle: float) -> np.ndarray:
|
|
vec = np.asarray(axis, dtype=float)
|
|
norm = np.linalg.norm(vec)
|
|
if norm < 1.0e-12:
|
|
return np.eye(3)
|
|
x, y, z = vec / norm
|
|
c, s = math.cos(angle), math.sin(angle)
|
|
one_c = 1.0 - c
|
|
return np.array(
|
|
[
|
|
[c + x * x * one_c, x * y * one_c - z * s, x * z * one_c + y * s],
|
|
[y * x * one_c + z * s, c + y * y * one_c, y * z * one_c - x * s],
|
|
[z * x * one_c - y * s, z * y * one_c + x * s, c + z * z * one_c],
|
|
],
|
|
dtype=float,
|
|
)
|
|
|
|
|
|
def _transform(xyz=(0.0, 0.0, 0.0), rpy=(0.0, 0.0, 0.0)) -> np.ndarray:
|
|
matrix = np.eye(4)
|
|
matrix[:3, :3] = _rpy_matrix(rpy)
|
|
matrix[:3, 3] = np.asarray(xyz, dtype=float)
|
|
return matrix
|
|
|
|
|
|
def _load_default_joint_positions() -> dict[str, float]:
|
|
text = GEN2_ROBOT_CFG.read_text(encoding="utf-8")
|
|
match = re.search(r"joint_pos=\{(.*?)\},\n\s*joint_vel=", text, flags=re.S)
|
|
if match is None:
|
|
return {}
|
|
return ast.literal_eval("{" + match.group(1) + "}")
|
|
|
|
|
|
def _load_base_init_z() -> float | None:
|
|
text = GEN2_ROBOT_CFG.read_text(encoding="utf-8")
|
|
match = re.search(r"pos=\(([^)]*)\)", text)
|
|
if match is None:
|
|
return None
|
|
return float(match.group(1).split(",")[2].strip())
|
|
|
|
|
|
def _load_base_height_target() -> float | None:
|
|
if not GEN2_ENV_CFG.exists():
|
|
return None
|
|
text = GEN2_ENV_CFG.read_text(encoding="utf-8")
|
|
match = re.search(r"GEN2_BASE_HEIGHT_TARGET\s*=\s*([0-9.]+)", text)
|
|
if match is None:
|
|
return None
|
|
return float(match.group(1))
|
|
|
|
|
|
def _link_tree(root: ET.Element):
|
|
links = {link.attrib["name"]: link for link in root.findall("link")}
|
|
children: dict[str, list[str]] = {}
|
|
parent_of: dict[str, str] = {}
|
|
edge = {}
|
|
for joint in root.findall("joint"):
|
|
parent = joint.find("parent").attrib["link"]
|
|
child = joint.find("child").attrib["link"]
|
|
origin = joint.find("origin")
|
|
axis_elem = joint.find("axis")
|
|
xyz = _parse_floats(origin.attrib.get("xyz"), (0.0, 0.0, 0.0)) if origin is not None else (0.0, 0.0, 0.0)
|
|
rpy = _parse_floats(origin.attrib.get("rpy"), (0.0, 0.0, 0.0)) if origin is not None else (0.0, 0.0, 0.0)
|
|
axis = (
|
|
_parse_floats(axis_elem.attrib.get("xyz"), (1.0, 0.0, 0.0)) if axis_elem is not None else (1.0, 0.0, 0.0)
|
|
)
|
|
children.setdefault(parent, []).append(child)
|
|
parent_of[child] = parent
|
|
edge[(parent, child)] = (joint.attrib["name"], joint.attrib.get("type", "fixed"), xyz, rpy, axis)
|
|
return links, children, parent_of, edge
|
|
|
|
|
|
def _link_world_transforms(root: ET.Element, joint_positions: dict[str, float]) -> tuple[dict[str, np.ndarray], dict[str, str], dict]:
|
|
links, children, parent_of, edge = _link_tree(root)
|
|
root_link = next(name for name in links if name not in parent_of)
|
|
world = {root_link: np.eye(4)}
|
|
stack = [root_link]
|
|
while stack:
|
|
parent = stack.pop()
|
|
for child in children.get(parent, []):
|
|
joint_name, joint_type, xyz, rpy, axis = edge[(parent, child)]
|
|
joint_transform = _transform(xyz, rpy)
|
|
if joint_type in {"revolute", "continuous"}:
|
|
rotation = np.eye(4)
|
|
rotation[:3, :3] = _axis_angle(axis, joint_positions.get(joint_name, 0.0))
|
|
joint_transform = joint_transform @ rotation
|
|
elif joint_type == "prismatic":
|
|
translation = np.eye(4)
|
|
translation[:3, 3] = np.asarray(axis, dtype=float) * joint_positions.get(joint_name, 0.0)
|
|
joint_transform = joint_transform @ translation
|
|
world[child] = world[parent] @ joint_transform
|
|
stack.append(child)
|
|
return world, parent_of, edge
|
|
|
|
|
|
def _box_vertices(transform: np.ndarray, size: tuple[float, float, float]) -> np.ndarray:
|
|
sx, sy, sz = (np.asarray(size, dtype=float) * 0.5).tolist()
|
|
local = np.array(
|
|
[
|
|
[-sx, -sy, -sz, 1.0],
|
|
[sx, -sy, -sz, 1.0],
|
|
[sx, sy, -sz, 1.0],
|
|
[-sx, sy, -sz, 1.0],
|
|
[-sx, -sy, sz, 1.0],
|
|
[sx, -sy, sz, 1.0],
|
|
[sx, sy, sz, 1.0],
|
|
[-sx, sy, sz, 1.0],
|
|
],
|
|
dtype=float,
|
|
)
|
|
return (transform @ local.T).T[:, :3]
|
|
|
|
|
|
def _collect_boxes(root: ET.Element) -> tuple[list[Box], dict[str, str], dict]:
|
|
joint_positions = _load_default_joint_positions()
|
|
link_world, parent_of, edge = _link_world_transforms(root, joint_positions)
|
|
boxes = []
|
|
for link in root.findall("link"):
|
|
link_name = link.attrib["name"]
|
|
collision = link.find("collision")
|
|
box = link.find("./collision/geometry/box")
|
|
if collision is None or box is None:
|
|
continue
|
|
origin = collision.find("origin")
|
|
xyz = _parse_floats(origin.attrib.get("xyz"), (0.0, 0.0, 0.0)) if origin is not None else (0.0, 0.0, 0.0)
|
|
rpy = _parse_floats(origin.attrib.get("rpy"), (0.0, 0.0, 0.0)) if origin is not None else (0.0, 0.0, 0.0)
|
|
size = _parse_floats(box.attrib["size"], (0.0, 0.0, 0.0))
|
|
transform = link_world[link_name] @ _transform(xyz, rpy)
|
|
boxes.append(Box(link_name, transform[:3, 3], transform[:3, :3], np.asarray(size, dtype=float) * 0.5, _box_vertices(transform, size)))
|
|
return boxes, parent_of, edge
|
|
|
|
|
|
def _is_adjacent(parent_of: dict[str, str], a: str, b: str) -> bool:
|
|
return parent_of.get(a) == b or parent_of.get(b) == a
|
|
|
|
|
|
def _obb_penetration(a: Box, b: Box) -> float | None:
|
|
axes = [a.axes[:, idx] for idx in range(3)] + [b.axes[:, idx] for idx in range(3)]
|
|
for i in range(3):
|
|
for j in range(3):
|
|
cross = np.cross(a.axes[:, i], b.axes[:, j])
|
|
norm = np.linalg.norm(cross)
|
|
if norm > 1.0e-8:
|
|
axes.append(cross / norm)
|
|
|
|
min_margin = float("inf")
|
|
for axis in axes:
|
|
axis = axis / np.linalg.norm(axis)
|
|
distance = abs(np.dot(b.center - a.center, axis))
|
|
radius_a = sum(a.half[i] * abs(np.dot(a.axes[:, i], axis)) for i in range(3))
|
|
radius_b = sum(b.half[i] * abs(np.dot(b.axes[:, i], axis)) for i in range(3))
|
|
margin = radius_a + radius_b - distance
|
|
if margin <= 0.0:
|
|
return None
|
|
min_margin = min(min_margin, margin)
|
|
return min_margin
|
|
|
|
|
|
def _check_inertials(root: ET.Element):
|
|
total_mass = 0.0
|
|
missing_inertial = []
|
|
bad_mass = []
|
|
bad_inertia = []
|
|
for link in root.findall("link"):
|
|
name = link.attrib["name"]
|
|
inertial = link.find("inertial")
|
|
if inertial is None:
|
|
missing_inertial.append(name)
|
|
continue
|
|
mass = float(inertial.find("mass").attrib["value"])
|
|
total_mass += mass
|
|
if mass <= 0.0:
|
|
bad_mass.append((name, mass))
|
|
inertia = inertial.find("inertia").attrib
|
|
matrix = np.array(
|
|
[
|
|
[float(inertia.get("ixx", 0.0)), float(inertia.get("ixy", 0.0)), float(inertia.get("ixz", 0.0))],
|
|
[float(inertia.get("ixy", 0.0)), float(inertia.get("iyy", 0.0)), float(inertia.get("iyz", 0.0))],
|
|
[float(inertia.get("ixz", 0.0)), float(inertia.get("iyz", 0.0)), float(inertia.get("izz", 0.0))],
|
|
]
|
|
)
|
|
eig = np.linalg.eigvalsh(matrix)
|
|
if eig.min() <= 0.0:
|
|
bad_inertia.append((name, eig.tolist()))
|
|
return total_mass, missing_inertial, bad_mass, bad_inertia
|
|
|
|
|
|
def _check_joint_limits(root: ET.Element):
|
|
movable = []
|
|
missing = []
|
|
bad = []
|
|
for joint in root.findall("joint"):
|
|
if joint.attrib.get("type") == "fixed":
|
|
continue
|
|
name = joint.attrib["name"]
|
|
limit = joint.find("limit")
|
|
if limit is None:
|
|
missing.append(name)
|
|
continue
|
|
lower = float(limit.attrib.get("lower", "nan"))
|
|
upper = float(limit.attrib.get("upper", "nan"))
|
|
effort = float(limit.attrib.get("effort", "nan"))
|
|
velocity = float(limit.attrib.get("velocity", "nan"))
|
|
movable.append((name, lower, upper, effort, velocity))
|
|
if not (lower < upper and effort > 0.0 and velocity > 0.0):
|
|
bad.append((name, lower, upper, effort, velocity))
|
|
return movable, missing, bad
|
|
|
|
|
|
def main() -> None:
|
|
root = ET.parse(URDF_PATH).getroot()
|
|
boxes, parent_of, edge = _collect_boxes(root)
|
|
movable, missing_limits, bad_limits = _check_joint_limits(root)
|
|
total_mass, missing_inertial, bad_mass, bad_inertia = _check_inertials(root)
|
|
|
|
adjacent_overlaps = []
|
|
non_adjacent_overlaps = []
|
|
for i, box_a in enumerate(boxes):
|
|
for box_b in boxes[i + 1 :]:
|
|
penetration = _obb_penetration(box_a, box_b)
|
|
if penetration is None:
|
|
continue
|
|
if _is_adjacent(parent_of, box_a.name, box_b.name):
|
|
joint_name = edge.get((box_a.name, box_b.name), edge.get((box_b.name, box_a.name), ("?",)))[0]
|
|
adjacent_overlaps.append((penetration, box_a.name, box_b.name, joint_name))
|
|
else:
|
|
non_adjacent_overlaps.append((penetration, box_a.name, box_b.name))
|
|
|
|
adjacent_overlaps.sort(reverse=True)
|
|
non_adjacent_overlaps.sort(reverse=True)
|
|
|
|
base_z = _load_base_init_z()
|
|
base_target = _load_base_height_target()
|
|
foot_min_z = {box.name: float(box.vertices[:, 2].min()) for box in boxes if box.name in FOOT_LINKS}
|
|
lowest_foot_z = min(foot_min_z.values())
|
|
suggested_base_z = -lowest_foot_z + 0.01
|
|
|
|
print("GEN2 RL READINESS CHECK")
|
|
print(f"urdf={URDF_PATH}")
|
|
print(f"links={len(root.findall('link'))} joints={len(root.findall('joint'))} collision_boxes={len(boxes)}")
|
|
print(f"total_mass_kg={total_mass:.6f}")
|
|
print(f"inertial_missing={len(missing_inertial)} bad_mass={len(bad_mass)} bad_inertia={len(bad_inertia)}")
|
|
print(f"movable_joints={len(movable)} missing_limits={len(missing_limits)} bad_limits={len(bad_limits)}")
|
|
if movable:
|
|
print(f"effort_range_nm={min(item[3] for item in movable):.3f}..{max(item[3] for item in movable):.3f}")
|
|
print(f"velocity_range_rad_s={min(item[4] for item in movable):.3f}..{max(item[4] for item in movable):.3f}")
|
|
print(f"non_adjacent_collision_overlaps={len(non_adjacent_overlaps)}")
|
|
for penetration, link_a, link_b in non_adjacent_overlaps:
|
|
print(f" NON_ADJ {link_a} <-> {link_b}: penetration_m={penetration:.6f}")
|
|
print(f"adjacent_collision_overlaps={len(adjacent_overlaps)}")
|
|
for penetration, link_a, link_b, joint_name in adjacent_overlaps:
|
|
print(f" ADJ {link_a} <-> {link_b} via {joint_name}: penetration_m={penetration:.6f}")
|
|
print("foot_collision_bottom_relative_to_body_root_m:")
|
|
for name, min_z in foot_min_z.items():
|
|
world_min = min_z + base_z if base_z is not None else float("nan")
|
|
print(f" {name}: relative={min_z:.6f} world_at_init={world_min:.6f}")
|
|
print(f"base_init_z={base_z:.6f}" if base_z is not None else "base_init_z=unknown")
|
|
print(f"base_height_target={base_target:.6f}" if base_target is not None else "base_height_target=unknown")
|
|
print(f"suggested_base_z_for_1cm_foot_clearance={suggested_base_z:.6f}")
|
|
|
|
hard_fail = bool(missing_inertial or bad_mass or bad_inertia or missing_limits or bad_limits or non_adjacent_overlaps)
|
|
ground_penetration = base_z is not None and lowest_foot_z + base_z < -0.005
|
|
if hard_fail:
|
|
print("RESULT=FAIL")
|
|
elif ground_penetration:
|
|
print("RESULT=WARN_RESET_HEIGHT")
|
|
elif adjacent_overlaps:
|
|
print("RESULT=WARN_ADJACENT_OVERLAPS")
|
|
else:
|
|
print("RESULT=PASS")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|