Files
cadSet/text-to-cad/models/robots/tom/v2/tom_double.py
T
2026-07-22 19:38:36 +08:00

663 lines
25 KiB
Python

from __future__ import annotations
import importlib.util
import sys
from pathlib import Path
V2_DIR = Path(__file__).resolve().parent
TOM_SOURCE = V2_DIR / "tom.py"
DOUBLE_SOURCE = V2_DIR / "assemblies" / "servo_horn_yoke_double.py"
ASSEMBLIES_DIR = V2_DIR / "assemblies"
PARTS_DIR = V2_DIR / "parts"
for path in (V2_DIR.parent, V2_DIR, PARTS_DIR, ASSEMBLIES_DIR):
if str(path) not in sys.path:
sys.path.insert(0, str(path))
SHOULDER_DOUBLE_CHILD_NAME = "shoulder_yaw_link"
SHOULDER_DOUBLE_PRIMARY_SERVO = "sts3250_rear"
SHOULDER_DOUBLE_ORIGINAL_SERVO_INSTANCE = "shoulder_yaw_link__sts3250_1"
SHOULDER_DOUBLE_REPLACED_MATES = frozenset(
{
"base_servo_horn_to_shoulder_yaw_mount",
"shoulder_yaw_mount_to_servo",
"shoulder_pitch_yoke_to_servo",
}
)
SHOULDER_DOUBLE_REMOVED_INSTANCES = frozenset(
{
"shoulder_pitch_link__servo_horn_yoke",
}
)
SHOULDER_DOUBLE_URDF_MESH_BY_STEP_BASENAME = {
"servo_horn_yoke_double_horn.step": "3MF/servo_horn_yoke_double_horn.3mf",
"servo_end_mount_double.step": "3MF/servo_end_mount_double.3mf",
}
SHOULDER_DOUBLE_URDF_MATERIAL_BY_STEP_BASENAME = {
"servo_horn_yoke_double_horn.step": "aluminum_5052",
"servo_end_mount_double.step": "aluminum_5052",
}
def _load_module(path: Path, module_label: str):
module_name = f"_tom_v2_{module_label}_{abs(hash(path))}"
spec = importlib.util.spec_from_file_location(module_name, path)
if spec is None or spec.loader is None:
raise RuntimeError(f"Unable to load {path}")
module = importlib.util.module_from_spec(spec)
sys.modules[module_name] = module
spec.loader.exec_module(module)
return module
def _prefix_mate(parent_name: str, mate: dict[str, object]) -> dict[str, object]:
fixed_part, fixed_frame = str(mate["fixed"]).split(":", 1)
moving_part, moving_frame = str(mate["moving"]).split(":", 1)
fixed = f"{parent_name}__{fixed_part}:{fixed_frame}"
moving = f"{parent_name}__{moving_part}:{moving_frame}"
return {
**mate,
"sourceLabel": f"{parent_name}_{mate['sourceLabel']}",
"fixed": fixed,
"moving": moving,
"fixedEndpoint": {
"part": fixed.split(":", 1)[0],
"frame": fixed_frame,
},
"movingEndpoint": {
"part": moving.split(":", 1)[0],
"frame": moving_frame,
},
}
def _replace_shoulder_double_mates(
*,
tom_module,
double_module_envelope: dict[str, object],
include_gripper: bool,
) -> list[dict[str, object]]:
mates = [
mate
for mate in tom_module._assembly_mates(include_gripper=include_gripper)
if str(mate["sourceLabel"]) not in SHOULDER_DOUBLE_REPLACED_MATES
]
mates.extend(
_prefix_mate(SHOULDER_DOUBLE_CHILD_NAME, mate)
for mate in double_module_envelope.get("assembly_mates", [])
)
mates.append(
tom_module._mate(
"base_servo_horn_to_shoulder_yaw_double_mount",
fixed="base_link__sts3250_3:output_horn_face",
moving=f"{SHOULDER_DOUBLE_CHILD_NAME}__servo_end_mount_double:front_horn_mount_face",
parameters={"clearance_mm": tom_module.BASE_TO_SHOULDER_YAW_HORN_CLEARANCE_MM},
)
)
return mates
def _local_transform_by_name(
instances: list[dict[str, object]],
name: str,
) -> list[float]:
for instance in instances:
if str(instance["name"]) == name:
return [float(value) for value in instance["transform"]]
raise RuntimeError(f"Unable to find local transform for {name}")
def _shoulder_double_parent_transform(
*,
tom_module,
child_name: str,
design_child_transform: list[float],
double_module_envelope: dict[str, object],
base_servo_transform: list[float],
) -> list[float]:
original_module_transform = tom_module._module_transform_for_child(
child_name=child_name,
design_child_transform=design_child_transform,
)
original_instances = tom_module._flat_child_instances(
child_name=child_name,
child_transform=original_module_transform,
)
original_servo_transform = _local_transform_by_name(
original_instances,
SHOULDER_DOUBLE_ORIGINAL_SERVO_INSTANCE,
)
primary_servo_local = _local_transform_by_name(
double_module_envelope["instances"],
f"sts3250_{SHOULDER_DOUBLE_PRIMARY_SERVO.removeprefix('sts3250_')}",
)
parent_transform = tom_module.multiply_transforms(
original_servo_transform,
tom_module.invert_rigid_transform(primary_servo_local),
)
base_horn_face_center = tom_module._transform_point(
base_servo_transform,
(
tom_module.STS3215_HORN_AXIS_LOCAL_X_MM,
tom_module.STS3250_OUTPUT_HORN_FACE_LOCAL_Y_MM,
0.0,
),
)
double_mount_face_center = tom_module._transform_point(
parent_transform,
(
tom_module.SERVO_END_MOUNT_FRONT_HORN_FACE_CENTER_LOCAL_MM[0],
0.0,
0.0,
),
)
return tom_module._translate_transform(
parent_transform,
(
base_horn_face_center[0] - double_mount_face_center[0],
base_horn_face_center[1] - double_mount_face_center[1],
base_horn_face_center[2] - double_mount_face_center[2],
),
)
def gen_step_with_options(*, include_gripper: bool = False) -> dict[str, object]:
tom_module = _load_module(TOM_SOURCE, "tom_double_base")
double_module = _load_module(DOUBLE_SOURCE, "shoulder_double")
double_envelope = double_module.gen_step()
instances: list[dict[str, object]] = []
instance_transforms_by_name: dict[str, list[float]] = {}
downstream_correction = tom_module._identity_transform()
pending_child_transform_overrides: dict[str, list[float]] = {}
for source_child in tom_module.robot_arm.robot_arm_assembly_children():
child_name = str(source_child["name"])
source_child_transform = [float(value) for value in source_child["transform"]]
child_transform = pending_child_transform_overrides.pop(child_name, None)
if child_transform is not None:
downstream_correction = tom_module.multiply_transforms(
child_transform,
tom_module.invert_rigid_transform(source_child_transform),
)
else:
child_transform = tom_module.multiply_transforms(
downstream_correction,
source_child_transform,
)
if child_name == "shoulder_yaw_link":
base_servo_transform = instance_transforms_by_name.get("base_link__sts3250_3")
if base_servo_transform is not None:
child_transform = tom_module._mate_shoulder_yaw_mount_to_base_servo_horn(
shoulder_yaw_transform=child_transform,
base_servo_transform=base_servo_transform,
)
downstream_correction = tom_module.multiply_transforms(
child_transform,
tom_module.invert_rigid_transform(source_child_transform),
)
if child_name == "elbow_pitch_link":
child_transform = tom_module._apply_v2_step_home_pose_to_elbow_pitch(
child_transform=child_transform,
instance_transforms_by_name=instance_transforms_by_name,
)
downstream_correction = tom_module.multiply_transforms(
child_transform,
tom_module.invert_rigid_transform(source_child_transform),
)
if child_name == tom_module.GRIPPER_CHILD_NAME:
if not include_gripper:
continue
terminal_servo_transform = instance_transforms_by_name.get(
tom_module.V2_GRIPPER_SERVO_INSTANCE_NAME
)
if terminal_servo_transform is not None:
child_transform = tom_module._mate_gripper_to_terminal_servo_horn(
gripper_transform=child_transform,
terminal_servo_transform=terminal_servo_transform,
)
instances.append(
{
"path": str(source_child["path"]),
"name": child_name,
"transform": child_transform,
"use_source_colors": bool(source_child.get("use_source_colors", True)),
}
)
instance_transforms_by_name[child_name] = child_transform
continue
if child_name == SHOULDER_DOUBLE_CHILD_NAME:
base_servo_transform = instance_transforms_by_name.get("base_link__sts3250_3")
if base_servo_transform is None:
raise RuntimeError("Cannot place shoulder double before base servo is available")
double_parent_transform = _shoulder_double_parent_transform(
tom_module=tom_module,
child_name=child_name,
design_child_transform=child_transform,
double_module_envelope=double_envelope,
base_servo_transform=base_servo_transform,
)
child_instances = tom_module._flatten_instances(
parent_name=child_name,
parent_transform=double_parent_transform,
local_instances=double_envelope["instances"],
local_source_dir=ASSEMBLIES_DIR,
)
instances.extend(child_instances)
for instance in child_instances:
instance_transforms_by_name[str(instance["name"])] = [
float(value) for value in instance["transform"]
]
continue
module_child_transform = tom_module._module_transform_for_child(
child_name=child_name,
design_child_transform=child_transform,
)
child_instances = tom_module._flat_child_instances(
child_name=child_name,
child_transform=module_child_transform,
)
child_instances = [
instance
for instance in child_instances
if str(instance["name"]) not in SHOULDER_DOUBLE_REMOVED_INSTANCES
]
instances.extend(child_instances)
for instance in child_instances:
instance_transforms_by_name[str(instance["name"])] = [
float(value) for value in instance["transform"]
]
replacement_kind = tom_module.ROLL_LINK_REPLACEMENTS.get(child_name)
if replacement_kind is not None:
mates = tom_module.roll_link_mates(replacement_kind)
next_pitch_child_name = tom_module.NEXT_PITCH_LINK_BY_ROLL_LINK.get(replacement_kind)
if next_pitch_child_name is not None:
upstream_servo_transform = tom_module.multiply_transforms(
child_transform,
mates.downstream_servo_local,
)
pending_child_transform_overrides[next_pitch_child_name] = (
tom_module._yoke_transform_for_servo_horn(
upstream_servo_transform=upstream_servo_transform,
)
)
return {
"instances": instances,
"assembly_mates": _replace_shoulder_double_mates(
tom_module=tom_module,
double_module_envelope=double_envelope,
include_gripper=include_gripper,
),
}
def gen_step() -> dict[str, object]:
envelope = gen_step_with_options(include_gripper=False)
return {
"instances": envelope["instances"],
"assembly_mates": envelope.get("assembly_mates", []),
}
def _zero_pose_source_frames_instances_and_double_visuals(
*,
tom_module,
double_module_envelope: dict[str, object],
) -> tuple[
dict[str, list[float]],
dict[str, list[float]],
dict[str, list[dict[str, object]]],
]:
frames: dict[str, list[float]] = {"base_footprint": tom_module._identity_transform()}
instance_transforms_by_name: dict[str, list[float]] = {}
visual_instances_by_link: dict[str, list[dict[str, object]]] = {}
downstream_correction = tom_module._identity_transform()
pending_child_transform_overrides: dict[str, list[float]] = {}
for source_child in tom_module.robot_arm.robot_arm_assembly_children():
child_name = str(source_child["name"])
source_child_transform = [float(value) for value in source_child["transform"]]
child_transform = pending_child_transform_overrides.pop(child_name, None)
if child_transform is not None:
downstream_correction = tom_module.multiply_transforms(
child_transform,
tom_module.invert_rigid_transform(source_child_transform),
)
else:
child_transform = tom_module.multiply_transforms(
downstream_correction,
source_child_transform,
)
if child_name == "shoulder_yaw_link":
base_servo_transform = instance_transforms_by_name.get("base_link__sts3250_3")
if base_servo_transform is not None:
child_transform = tom_module._mate_shoulder_yaw_mount_to_base_servo_horn(
shoulder_yaw_transform=child_transform,
base_servo_transform=base_servo_transform,
)
downstream_correction = tom_module.multiply_transforms(
child_transform,
tom_module.invert_rigid_transform(source_child_transform),
)
if child_name == tom_module.GRIPPER_CHILD_NAME:
continue
if child_name == SHOULDER_DOUBLE_CHILD_NAME:
base_servo_transform = instance_transforms_by_name.get("base_link__sts3250_3")
if base_servo_transform is None:
raise RuntimeError("Cannot place shoulder double before base servo is available")
frames[child_name] = child_transform
double_parent_transform = _shoulder_double_parent_transform(
tom_module=tom_module,
child_name=child_name,
design_child_transform=child_transform,
double_module_envelope=double_module_envelope,
base_servo_transform=base_servo_transform,
)
child_instances = tom_module._flatten_instances(
parent_name=child_name,
parent_transform=double_parent_transform,
local_instances=double_module_envelope["instances"],
local_source_dir=ASSEMBLIES_DIR,
)
visual_instances_by_link[child_name] = []
inverse_source_frame = tom_module.invert_rigid_transform(child_transform)
for instance in child_instances:
world_transform = [float(value) for value in instance["transform"]]
instance_transforms_by_name[str(instance["name"])] = world_transform
visual_instance = dict(instance)
visual_instance["transform"] = tom_module.multiply_transforms(
inverse_source_frame,
world_transform,
)
visual_instances_by_link[child_name].append(visual_instance)
continue
module_child_transform = tom_module._module_transform_for_child(
child_name=child_name,
design_child_transform=child_transform,
)
frames[child_name] = module_child_transform
child_instances = tom_module._flat_child_instances(
child_name=child_name,
child_transform=module_child_transform,
)
child_instances = [
instance
for instance in child_instances
if str(instance["name"]) not in SHOULDER_DOUBLE_REMOVED_INSTANCES
]
for instance in child_instances:
instance_transforms_by_name[str(instance["name"])] = [
float(value) for value in instance["transform"]
]
replacement_kind = tom_module.ROLL_LINK_REPLACEMENTS.get(child_name)
if replacement_kind is not None:
mates = tom_module.roll_link_mates(replacement_kind)
next_pitch_child_name = tom_module.NEXT_PITCH_LINK_BY_ROLL_LINK.get(replacement_kind)
if next_pitch_child_name is not None:
upstream_servo_transform = tom_module.multiply_transforms(
child_transform,
mates.downstream_servo_local,
)
pending_child_transform_overrides[next_pitch_child_name] = (
tom_module._yoke_transform_for_servo_horn(
upstream_servo_transform=upstream_servo_transform,
)
)
wrist_pitch = frames.get("wrist_pitch_link")
if wrist_pitch is not None and "wrist_roll_link" not in frames:
wrist_roll_joint = next(
joint
for joint in tom_module.robot_arm.URDF_JOINTS
if str(joint["name"]) == "wrist_roll"
)
frames["wrist_roll_link"] = tom_module.multiply_transforms(
wrist_pitch,
tom_module._transform_from_urdf_origin_attrs(
str(wrist_roll_joint["origin_xyz"]),
str(wrist_roll_joint["origin_rpy"]),
),
)
return frames, instance_transforms_by_name, visual_instances_by_link
def _double_horn_axis(
tom_module,
instance_transforms_by_name: dict[str, list[float]],
) -> tuple[tuple[float, float, float], tuple[float, float, float]]:
horn_transform = instance_transforms_by_name.get(
f"{SHOULDER_DOUBLE_CHILD_NAME}__servo_horn_yoke_double_horn"
)
if horn_transform is None:
raise RuntimeError("Missing double horn transform for shoulder_pitch URDF axis")
center = tom_module._transform_point(
horn_transform,
(
tom_module.STS3215_HORN_AXIS_LOCAL_X_MM,
0.0,
0.0,
),
)
direction = tom_module._normalized_vector(
tom_module._transform_direction(horn_transform, (0.0, 1.0, 0.0))
)
return center, direction
def _urdf_link_frames_axes_and_double_visuals(
*,
tom_module,
double_module_envelope: dict[str, object],
) -> tuple[
dict[str, list[float]],
dict[str, list[float]],
dict[str, str],
dict[str, list[dict[str, object]]],
]:
source_frames, instance_transforms, double_visuals = (
_zero_pose_source_frames_instances_and_double_visuals(
tom_module=tom_module,
double_module_envelope=double_module_envelope,
)
)
urdf_frames = {name: list(transform) for name, transform in source_frames.items()}
axes_by_joint: dict[str, str] = {}
servo_axis_instance_by_joint = dict(tom_module.URDF_SERVO_AXIS_INSTANCE_BY_JOINT)
for joint in tom_module._no_gripper_joints():
joint_name = str(joint["name"])
if joint_name == "shoulder_pitch":
axis_center_world, axis_direction_world = _double_horn_axis(
tom_module,
instance_transforms,
)
else:
servo_instance_name = servo_axis_instance_by_joint.get(joint_name)
if not servo_instance_name:
continue
servo_transform = instance_transforms.get(servo_instance_name)
if servo_transform is None:
raise RuntimeError(
f"Missing servo instance {servo_instance_name!r} "
f"for URDF joint {joint_name!r}"
)
axis_center_world, axis_direction_world = tom_module._servo_horn_axis_from_transform(
servo_transform,
)
child_link = str(joint["child"])
if child_link not in urdf_frames:
parent_frame = urdf_frames.get(str(joint["parent"]))
if parent_frame is None:
raise RuntimeError(f"Missing parent frame for URDF joint {joint_name!r}")
urdf_frames[child_link] = list(parent_frame)
urdf_frames[child_link] = tom_module._with_translation(
urdf_frames[child_link],
axis_center_world,
)
axis_in_joint_frame = tom_module._world_vector_to_local(
urdf_frames[child_link],
axis_direction_world,
)
axes_by_joint[joint_name] = tom_module._axis_attrs(axis_in_joint_frame)
return source_frames, urdf_frames, axes_by_joint, double_visuals
def _instances_for_urdf_link_double(
*,
tom_module,
link_name: str,
double_visuals: dict[str, list[dict[str, object]]],
) -> list[dict[str, object]]:
if link_name == SHOULDER_DOUBLE_CHILD_NAME:
return [dict(instance) for instance in double_visuals.get(link_name, [])]
if link_name == "shoulder_pitch_link":
return []
return tom_module._instances_for_urdf_link(link_name)
def _link_element_double(
*,
tom_module,
link_name: str,
source_frames: dict[str, list[float]],
urdf_frames: dict[str, list[float]],
double_visuals: dict[str, list[dict[str, object]]],
):
owner_source_link_frame = source_frames.get(
link_name,
urdf_frames.get(link_name, tom_module._identity_transform()),
)
urdf_link_frame = urdf_frames.get(link_name, owner_source_link_frame)
instances = _instances_for_urdf_link_double(
tom_module=tom_module,
link_name=link_name,
double_visuals=double_visuals,
)
children = []
inertial = tom_module._inertial_element_for_instances(
instances,
link_name=link_name,
source_frames=source_frames,
urdf_link_frame=urdf_link_frame,
)
if inertial is not None:
children.append(inertial)
children.extend(
tom_module._visual_element_for_instance(
instance,
source_link_frame=source_frames.get(
str(instance.get("source_link_name", link_name)),
urdf_frames.get(
str(instance.get("source_link_name", link_name)),
tom_module._identity_transform(),
),
),
urdf_link_frame=urdf_link_frame,
)
for instance in instances
)
return tom_module._xml(
"link",
{"name": link_name},
tuple(children),
)
def gen_urdf() -> dict[str, object]:
"""Jointed no-gripper URDF for the tom_double shoulder-yaw variant.
Design ledger:
- Based on tom_v2's no-gripper URDF and 3MF visual mesh convention.
- The shoulder-yaw single STS3250/end-mount visual module is replaced by
the side-by-side double servo horn/mount assembly.
- The shoulder_pitch joint frame/axis is derived from the double horn plate
centerline, not from either offset servo body.
- Inertials use the same vendor servo masses and 5052 STEP-volume
plate/standoff mass model as tom.urdf. Gripper links and collision
elements are intentionally omitted to match tom.urdf's current scope.
"""
tom_module = _load_module(TOM_SOURCE, "tom_double_urdf_base")
double_module = _load_module(DOUBLE_SOURCE, "tom_double_urdf_shoulder_double")
double_envelope = double_module.gen_step()
tom_module.URDF_JOINT_EFFORT_NM_BY_NAME = dict(
tom_module.URDF_JOINT_EFFORT_NM_BY_NAME
)
tom_module.URDF_JOINT_EFFORT_NM_BY_NAME["shoulder_pitch"] = (
2.0 * tom_module.STS3250_STALL_TORQUE_NM
)
tom_module.URDF_MESH_BY_STEP_BASENAME.update(SHOULDER_DOUBLE_URDF_MESH_BY_STEP_BASENAME)
tom_module.URDF_MATERIAL_BY_STEP_BASENAME.update(
SHOULDER_DOUBLE_URDF_MATERIAL_BY_STEP_BASENAME
)
source_frames, urdf_frames, axes_by_joint, double_visuals = (
_urdf_link_frames_axes_and_double_visuals(
tom_module=tom_module,
double_module_envelope=double_envelope,
)
)
robot = tom_module._xml("robot", {"name": "tom_v2_double"})
robot.append(tom_module.ET.Comment(" Generated by models/robots/tom/v2/tom_double.py. "))
robot.append(
tom_module.ET.Comment(
" Jointed no-gripper URDF for the double shoulder-yaw variant; "
"all visual meshes are 3MF in millimeters with scale 0.001 0.001 0.001. "
"Link inertials use Feetech STS3250 74.5 g / STS3215 55 g masses, "
"5052 aluminum density for plates/standoffs, 2x STS3250 stall torque "
"on shoulder_pitch, and single-servo stall torque plus no-load speeds "
"elsewhere converted to SI units. "
)
)
for material_name, rgba in tom_module.URDF_MATERIALS.items():
robot.append(
tom_module._xml(
"material",
{"name": material_name},
(tom_module._xml("color", {"rgba": tom_module._rgba_text(rgba)}),),
)
)
for link_name in tom_module.NO_GRIPPER_LINK_NAMES:
robot.append(
_link_element_double(
tom_module=tom_module,
link_name=link_name,
source_frames=source_frames,
urdf_frames=urdf_frames,
double_visuals=double_visuals,
)
)
for joint in tom_module._no_gripper_joints():
robot.append(
tom_module._joint_element(
joint,
link_frames=urdf_frames,
axes_by_joint=axes_by_joint,
)
)
return {"xml": tom_module._serialize_urdf(robot)}
def gen_srdf() -> dict[str, object]:
"""MoveIt semantic groups and named states for the double shoulder v2 arm."""
return _load_module(TOM_SOURCE, "tom_double_srdf_base").gen_srdf_with_options(
include_gripper=False,
robot_name="tom_v2_double",
urdf="tom_double.urdf",
)