wehub-resource-sync
141 lines
5.4 KiB
Python
Executable File
141 lines
5.4 KiB
Python
Executable File
#!/usr/bin/env python3
|
|
"""
|
|
An example of how to load and animate a URDF given some changing joint angles.
|
|
|
|
Usage:
|
|
python -m animated_urdf
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
import math
|
|
from pathlib import Path
|
|
|
|
import rerun as rr
|
|
import rerun.blueprint as rrb
|
|
|
|
TIMELINE = "example_time"
|
|
|
|
|
|
def main() -> None:
|
|
parser = argparse.ArgumentParser(
|
|
description="An example of how to load and animate a URDF given some changing joint angles.",
|
|
)
|
|
rr.script_add_args(parser)
|
|
parser.add_argument(
|
|
"--dual",
|
|
action="store_true",
|
|
help="Load the same URDF twice with different frame prefixes (dual-arm demo).",
|
|
)
|
|
args = parser.parse_args()
|
|
|
|
urdf_path = Path(__file__).parent.parent.parent / "rust" / "animated_urdf" / "data" / "so100.urdf"
|
|
|
|
if args.dual:
|
|
run_dual(args, urdf_path)
|
|
else:
|
|
run_single(args, urdf_path)
|
|
|
|
|
|
def run_single(args: argparse.Namespace, urdf_path: Path) -> None:
|
|
duration = 0.0
|
|
rec = rr.script_setup(args, "rerun_example_animated_urdf")
|
|
rec.set_time(TIMELINE, duration=duration)
|
|
|
|
# Log the URDF file once
|
|
rec.log_file_from_path(urdf_path)
|
|
|
|
# Load the URDF tree structure into memory
|
|
urdf_tree = rr.urdf.UrdfTree.from_file_path(urdf_path)
|
|
|
|
# Hide the collision geometries by default in the viewer.
|
|
blueprint = rrb.Grid(
|
|
rrb.Spatial3DView(
|
|
name="Animated URDF", overrides={"so_arm100/collision_geometries": rrb.EntityBehavior(visible=False)}
|
|
)
|
|
)
|
|
rec.send_blueprint(blueprint)
|
|
|
|
for step in range(10000):
|
|
for joint_index, joint in enumerate(urdf_tree.joints()):
|
|
if joint.joint_type == "revolute":
|
|
# Usually this angle would come from a measurement, here we just fake something
|
|
dynamic_angle = _fake_angle(joint, step, joint_index, phase=0.0)
|
|
|
|
# Rerun loads the URDF transforms with child/parent frame relations.
|
|
# To move a joint, we just need to log a new transform between those frames.
|
|
# Here, we use the `compute_transform` method that automatically takes care
|
|
# of setting the frame names and calculating the full transform from the joint angle.
|
|
transform = joint.compute_transform(dynamic_angle, clamp=True)
|
|
rec.log("transforms", transform)
|
|
|
|
# We can also work with links from the URDF tree.
|
|
# Here, we change the color of the visual mesh entities of the "jaw" link based on the joint angle.
|
|
link = urdf_tree.get_joint_child(joint)
|
|
if link.name == "jaw":
|
|
for visual_path in urdf_tree.get_visual_geometry_paths(link):
|
|
normalized_angle = (dynamic_angle - joint.limit_lower) / (joint.limit_upper - joint.limit_lower)
|
|
rgba = [1.0 - normalized_angle, normalized_angle, 0, 0.5]
|
|
rec.log(visual_path, rr.Asset3D.from_fields(albedo_factor=rgba))
|
|
|
|
duration += 0.03
|
|
rec.set_time(TIMELINE, duration=duration)
|
|
|
|
rr.script_teardown(args)
|
|
|
|
|
|
def run_dual(args: argparse.Namespace, urdf_path: Path) -> None:
|
|
"""Load the same URDF twice with different frame prefixes (dual-arm demo)."""
|
|
rec = rr.script_setup(args, "rerun_example_animated_urdf")
|
|
|
|
# Load the same URDF twice with different prefixes.
|
|
# - entity_path_prefix separates geometry in the entity tree
|
|
# - frame_prefix makes frame IDs unique ("left/base", "right/base", …)
|
|
left = rr.urdf.UrdfTree.from_file_path(urdf_path, entity_path_prefix="left", frame_prefix="left/")
|
|
right = rr.urdf.UrdfTree.from_file_path(urdf_path, entity_path_prefix="right", frame_prefix="right/")
|
|
|
|
# Log both robots (geometry + static transforms with prefixed frame IDs).
|
|
left.log_urdf_to_recording()
|
|
right.log_urdf_to_recording()
|
|
|
|
# Offset the arms so they don't overlap.
|
|
rec.log("left", rr.Transform3D(translation=[-0.2, 0, 0]), static=True)
|
|
rec.log("right", rr.Transform3D(translation=[0.2, 0, 0]), static=True)
|
|
|
|
blueprint = rrb.Grid(
|
|
rrb.Spatial3DView(
|
|
name="Dual Arm",
|
|
overrides={
|
|
"left/so_arm100/collision_geometries": rrb.EntityBehavior(visible=False),
|
|
"right/so_arm100/collision_geometries": rrb.EntityBehavior(visible=False),
|
|
},
|
|
)
|
|
)
|
|
rec.send_blueprint(blueprint)
|
|
|
|
for step in range(10000):
|
|
rec.set_time("step", sequence=step)
|
|
|
|
for joint_index, joint in enumerate(left.joints()):
|
|
if joint.joint_type == "revolute":
|
|
angle = _fake_angle(joint, step, joint_index, phase=0.0)
|
|
rec.log("left/joint_transforms", joint.compute_transform(angle, clamp=True))
|
|
|
|
for joint_index, joint in enumerate(right.joints()):
|
|
if joint.joint_type == "revolute":
|
|
angle = _fake_angle(joint, step, joint_index, phase=2.0)
|
|
rec.log("right/joint_transforms", joint.compute_transform(angle, clamp=True))
|
|
|
|
rr.script_teardown(args)
|
|
|
|
|
|
def _fake_angle(joint: rr.urdf.UrdfJoint, step: int, joint_index: int, phase: float) -> float:
|
|
"""Generate a smooth oscillating angle within the joint's limits."""
|
|
sin_value = math.sin(step * (0.02 + joint_index / 100.0) + phase)
|
|
return joint.limit_lower + (sin_value + 1.0) / 3.0 * (joint.limit_upper - joint.limit_lower)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|