Skip to content

Commit 1a7002d

Browse files
committed
Add URDF model and visualization script
- Created software/src/lerobot/robots/alohamini/alohamini.urdf with base, lift, and dual arms - Created software/examples/debug/visualize_urdf.py to visualize URDF in Rerun
1 parent 5e8f280 commit 1a7002d

2 files changed

Lines changed: 500 additions & 0 deletions

File tree

Lines changed: 186 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,186 @@
1+
import argparse
2+
import time
3+
import math
4+
import xml.etree.ElementTree as ET
5+
import numpy as np
6+
import rerun as rr
7+
from scipy.spatial.transform import Rotation as R
8+
9+
def parse_urdf(urdf_path):
10+
tree = ET.parse(urdf_path)
11+
root = tree.getroot()
12+
13+
links = {}
14+
joints = []
15+
16+
for link in root.findall('link'):
17+
name = link.get('name')
18+
visual = link.find('visual')
19+
geometry_info = None
20+
color_info = None
21+
origin_info = None
22+
23+
if visual is not None:
24+
geo = visual.find('geometry')
25+
if geo is not None:
26+
box = geo.find('box')
27+
cylinder = geo.find('cylinder')
28+
if box is not None:
29+
geometry_info = {'type': 'box', 'size': [float(x) for x in box.get('size').split()]}
30+
elif cylinder is not None:
31+
geometry_info = {'type': 'cylinder', 'length': float(cylinder.get('length')), 'radius': float(cylinder.get('radius'))}
32+
33+
mat = visual.find('material')
34+
if mat is not None:
35+
color_info = mat.get('name') # Simplified
36+
37+
orig = visual.find('origin')
38+
if orig is not None:
39+
xyz = [float(x) for x in orig.get('xyz', '0 0 0').split()]
40+
rpy = [float(x) for x in orig.get('rpy', '0 0 0').split()]
41+
origin_info = {'xyz': xyz, 'rpy': rpy}
42+
else:
43+
origin_info = {'xyz': [0,0,0], 'rpy': [0,0,0]}
44+
45+
links[name] = {
46+
'visual': geometry_info,
47+
'color': color_info,
48+
'visual_origin': origin_info
49+
}
50+
51+
for joint in root.findall('joint'):
52+
name = joint.get('name')
53+
type_ = joint.get('type')
54+
parent = joint.find('parent').get('link')
55+
child = joint.find('child').get('link')
56+
57+
origin = joint.find('origin')
58+
xyz = [float(x) for x in origin.get('xyz', '0 0 0').split()] if origin is not None else [0,0,0]
59+
rpy = [float(x) for x in origin.get('rpy', '0 0 0').split()] if origin is not None else [0,0,0]
60+
61+
axis_elem = joint.find('axis')
62+
axis = [float(x) for x in axis_elem.get('xyz', '1 0 0').split()] if axis_elem is not None else [1,0,0]
63+
64+
limit = joint.find('limit')
65+
limits = None
66+
if limit is not None:
67+
limits = (float(limit.get('lower', -3.14)), float(limit.get('upper', 3.14)))
68+
69+
joints.append({
70+
'name': name,
71+
'type': type_,
72+
'parent': parent,
73+
'child': child,
74+
'xyz': xyz,
75+
'rpy': rpy,
76+
'axis': axis,
77+
'limits': limits
78+
})
79+
80+
return links, joints
81+
82+
def get_transform(xyz, rpy):
83+
rot = R.from_euler('xyz', rpy).as_matrix()
84+
T = np.eye(4)
85+
T[:3, :3] = rot
86+
T[:3, 3] = xyz
87+
return T
88+
89+
def visualize(urdf_path):
90+
rr.init("urdf_visualizer", spawn=True)
91+
92+
links, joints = parse_urdf(urdf_path)
93+
94+
# Build tree structure
95+
# joint_map: parent_link -> list of joints where this link is parent
96+
joint_map = {}
97+
for j in joints:
98+
p = j['parent']
99+
if p not in joint_map:
100+
joint_map[p] = []
101+
joint_map[p].append(j)
102+
103+
# State
104+
t = 0.0
105+
106+
while True:
107+
# Update joint values (sine waves)
108+
joint_values = {}
109+
for j in joints:
110+
if j['type'] in ['revolute', 'continuous']:
111+
val = math.sin(t + sum(map(ord, j['name']))) # Random phase
112+
joint_values[j['name']] = val
113+
elif j['type'] == 'prismatic':
114+
val = 0.3 + 0.3 * math.sin(t)
115+
joint_values[j['name']] = val
116+
else:
117+
joint_values[j['name']] = 0.0
118+
119+
# FK
120+
# Stack: (link_name, accumulated_transform)
121+
stack = [('base_link', np.eye(4))]
122+
123+
while stack:
124+
link_name, T_parent = stack.pop()
125+
126+
# Log link visual
127+
link_data = links.get(link_name)
128+
if link_data and link_data['visual']:
129+
# Visual origin offset
130+
v_orig = link_data['visual_origin']
131+
T_visual_offset = get_transform(v_orig['xyz'], v_orig['rpy'])
132+
T_visual = T_parent @ T_visual_offset
133+
134+
rr.set_time_seconds("sim_time", t)
135+
136+
# Decompose for rerun
137+
trans = T_visual[:3, 3]
138+
rot = R.from_matrix(T_visual[:3, :3]).as_quat() # xyzw
139+
# Rearrange to xyzw for scipy -> xyzw for rerun (check convention, rerun uses xyzw)
140+
141+
entity_path = f"robot/{link_name}"
142+
143+
geo = link_data['visual']
144+
if geo['type'] == 'box':
145+
rr.log(entity_path, rr.Boxes3D(half_sizes=[s/2 for s in geo['size']], centers=[0,0,0]), rr.Transform3D(translation=trans, rotation=rr.Quaternion(xyzw=rot)))
146+
elif geo['type'] == 'cylinder':
147+
# Rerun cylinder is along Z?
148+
# URDF cylinder is along Z (usually)
149+
# We might need to check visualization, but simple logging:
150+
# Rerun doesn't have Cylinder3D primitive in older versions, checking...
151+
# It does have standard primitives now? Or we use Mesh3D / Points.
152+
# As fallback, log a box approximating cylinder or use Points if needed.
153+
# Use Box for now to be safe and simple.
154+
rr.log(entity_path, rr.Boxes3D(half_sizes=[geo['radius'], geo['radius'], geo['length']/2], centers=[0,0,0]), rr.Transform3D(translation=trans, rotation=rr.Quaternion(xyzw=rot)))
155+
156+
# Children
157+
children_joints = joint_map.get(link_name, [])
158+
for j in children_joints:
159+
# Static transform
160+
T_static = get_transform(j['xyz'], j['rpy'])
161+
162+
# Joint transform
163+
T_joint = np.eye(4)
164+
if j['type'] in ['revolute', 'continuous']:
165+
angle = joint_values.get(j['name'], 0)
166+
axis = np.array(j['axis'])
167+
# Axis-angle rotation
168+
rot_j = R.from_rotvec(axis * angle).as_matrix()
169+
T_joint[:3, :3] = rot_j
170+
elif j['type'] == 'prismatic':
171+
dist = joint_values.get(j['name'], 0)
172+
axis = np.array(j['axis'])
173+
T_joint[:3, 3] = axis * dist
174+
175+
T_child = T_parent @ T_static @ T_joint
176+
stack.append((j['child'], T_child))
177+
178+
t += 0.05
179+
time.sleep(0.05)
180+
181+
if __name__ == "__main__":
182+
parser = argparse.ArgumentParser()
183+
parser.add_argument("--urdf", default="../../src/lerobot/robots/alohamini/alohamini.urdf")
184+
args = parser.parse_args()
185+
186+
visualize(args.urdf)

0 commit comments

Comments
 (0)