liyang
aed6e20a90
从 ~/mujoco-web-spike 抽出 manip.html 的依赖闭包:
- manip.html + src/{manip,autograsp,scene-bridge,math}.js
- scenes/so101_pair(桌面并排两台 SO-101,默认)与 scenes/p7_bimanual(P7 双臂 + SO-101 夹爪)
—— 自包含 MJCF,网格已内联
- serve.py(COOP/COEP + no-store 的静态服务器)、start.sh、README
- tools/export_manip_scene.py 场景生成脚本(仅供参考,跑页面不需要)
快照对应 2026-08-31 的最后一次改动(自动抓取规划 autograsp.js 落地版)。
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_016tgfYxR5mWNRKQ67PKga4T
671 lines
28 KiB
Python
671 lines
28 KiB
Python
"""把双臂工作站 + 操作台 + 测试物体导出成浏览器能直接跑的自包含 MJCF。
|
|
|
|
~/unitree_rl_mjlab/.venv/bin/python tools/export_manip_scene.py
|
|
|
|
产出 scenes/<name>/:
|
|
model.xml 自包含 MJCF(网格内联 + 视觉网格抽稀,不依赖任何外部文件)
|
|
meta.json 给前端用的索引:臂/手的执行器、工具 site、指尖 geom、
|
|
物体和桌面的几何参数、就绪位形
|
|
|
|
这里没有策略,是纯手动操作的沙盒 —— 前端用 IK 驱动腕部、用一个握力滑条驱动手指。
|
|
|
|
网格处理照搬 unitree_rl_mjlab/scripts/export_web_dog.py 里踩出来的三条:
|
|
· 写回顶点前要用 mesh_pos/mesh_quat 变换回原坐标系,否则重新编译会再对齐一次,
|
|
几何整体偏移;
|
|
· MJCF 的 <mesh normal> 是**每顶点**一条,不是每面;
|
|
· spec 里源文件的法线不受 usernormal 覆盖,只能删掉网格重建。
|
|
"""
|
|
|
|
import json
|
|
from dataclasses import dataclass
|
|
from pathlib import Path
|
|
|
|
import mujoco
|
|
import numpy as np
|
|
import tyro
|
|
|
|
HERE = Path(__file__).resolve().parent.parent
|
|
WS_ROOT = (Path.home() / "linker-sim/packages/linker-robot-assets/src"
|
|
/ "linker_robot_assets/assets/workstations")
|
|
|
|
|
|
@dataclass
|
|
class Cfg:
|
|
workstation: str = "p7_i1_l6_bimanual"
|
|
"""linker-robot-assets 里的工作站名。"""
|
|
name: str = "p7_bimanual"
|
|
"""导出目录名。"""
|
|
table_top: float = 0.75
|
|
"""操作台高度 [m]。"""
|
|
table_x: float = 0.42
|
|
"""操作台中心到机器人的距离 [m]。"""
|
|
decimate: int = 32
|
|
"""视觉网格的聚类抽稀格数;0 = 不抽稀。"""
|
|
gripper: str = "so101"
|
|
"""末端:l6=原装五指手;so101=换成 SO-101 二指夹爪(仅 rig=p7 时有效)。"""
|
|
rig: str = "so101_pair"
|
|
"""机器人:p7=linker 双臂工作站;so101_pair=桌面上并排两台 SO-101。"""
|
|
|
|
|
|
# ---------------------------------------------------------------- 网格处理
|
|
|
|
def decimate(v, f, cells=48):
|
|
"""顶点聚类抽稀:包围盒切 cells^3 的格子,每格顶点并成一个。"""
|
|
lo, hi = v.min(0), v.max(0)
|
|
size = float((hi - lo).max())
|
|
if size <= 0:
|
|
return v, f
|
|
key = np.floor((v - lo) / (size / cells)).astype(np.int64)
|
|
uniq, inv = np.unique(key, axis=0, return_inverse=True)
|
|
acc = np.zeros((len(uniq), 3))
|
|
cnt = np.zeros(len(uniq))
|
|
np.add.at(acc, inv, v)
|
|
np.add.at(cnt, inv, 1)
|
|
f2 = inv[f]
|
|
keep = (f2[:, 0] != f2[:, 1]) & (f2[:, 1] != f2[:, 2]) & (f2[:, 0] != f2[:, 2])
|
|
f2 = f2[keep]
|
|
if len(f2) < 4 or len(uniq) < 4:
|
|
return v, f
|
|
used, f3 = np.unique(f2, return_inverse=True)
|
|
v2, f3 = (acc / cnt[:, None])[used], f3.reshape(-1, 3)
|
|
# 抽稀把薄零件压成一张片时,MuJoCo 会拒收("mesh volume is too small")。
|
|
# 退化了就原样返回,少省几百 KB 而已。
|
|
ext = v2.max(0) - v2.min(0)
|
|
if len(v2) < 20 or float(ext.min()) < 1e-3:
|
|
return v, f
|
|
return v2, f3
|
|
|
|
|
|
def convex_hull(v):
|
|
"""取凸包 —— MuJoCo 做网格碰撞时用的本来就是凸包,所以这一步零精度损失,
|
|
但顶点数能从几万降到几百。"""
|
|
from scipy.spatial import ConvexHull
|
|
|
|
h = ConvexHull(v)
|
|
used, faces = np.unique(h.simplices, return_inverse=True)
|
|
return v[used], faces.reshape(-1, 3)
|
|
|
|
|
|
def inline_meshes(spec: mujoco.MjSpec, m: mujoco.MjModel, cells: int) -> None:
|
|
"""网格内联进 XML。
|
|
|
|
这套资产的视觉几何和碰撞几何**引用同一批网格**(41 个全是两用),
|
|
直接抽稀会改碰撞形状,不抽稀又有 49 万顶点 / 25MB。做法是拆成两份:
|
|
· 碰撞用凸包(MuJoCo 反正只用凸包,行为完全不变)
|
|
· 视觉用抽稀版 `<name>_vis`,再把纯视觉 geom 的 meshname 改过去
|
|
"""
|
|
by_name = {m.mesh(i).name: i for i in range(m.nmesh)}
|
|
collide, visual = set(), set()
|
|
for g in range(m.ngeom):
|
|
if m.geom_type[g] != mujoco.mjtGeom.mjGEOM_MESH:
|
|
continue
|
|
(collide if (m.geom_contype[g] or m.geom_conaffinity[g]) else visual).add(
|
|
int(m.geom_dataid[g]))
|
|
|
|
todo = []
|
|
for mesh in list(spec.meshes):
|
|
i = by_name.get(mesh.name)
|
|
if i is None:
|
|
continue
|
|
va, vn = m.mesh_vertadr[i], m.mesh_vertnum[i]
|
|
fa, fn = m.mesh_faceadr[i], m.mesh_facenum[i]
|
|
v = m.mesh_vert[va:va + vn].astype(np.float64)
|
|
rot = np.zeros(9)
|
|
mujoco.mju_quat2Mat(rot, m.mesh_quat[i])
|
|
v = v @ rot.reshape(3, 3).T + m.mesh_pos[i] # 变换回原坐标系
|
|
f = m.mesh_face[fa:fa + fn].astype(np.int64)
|
|
todo.append((mesh.name, i, v, f))
|
|
|
|
for nm, _, _, _ in todo:
|
|
spec.delete(spec.mesh(nm)) # 删掉重建,甩掉源文件的法线
|
|
|
|
def emit(name, v, f):
|
|
nu = spec.add_mesh()
|
|
nu.name = name
|
|
nu.uservert = np.round(v, 4).reshape(-1).tolist()
|
|
nu.userface = f.astype(np.int32).reshape(-1).tolist()
|
|
|
|
n_v0 = n_v1 = 0
|
|
for nm, i, v, f in todo:
|
|
n_v0 += len(v)
|
|
if i in collide:
|
|
try:
|
|
cv, cf = convex_hull(v)
|
|
except Exception: # noqa: BLE001
|
|
cv, cf = decimate(v, f, 64)
|
|
emit(nm, cv, cf)
|
|
n_v1 += len(cv)
|
|
else:
|
|
emit(nm, *((decimate(v, f, cells)) if cells else (v, f)))
|
|
n_v1 += len(spec.mesh(nm).uservert) // 3
|
|
if i in visual and i in collide:
|
|
vv, vf = decimate(v, f, cells) if cells else (v, f)
|
|
emit(f"{nm}_vis", vv, vf)
|
|
n_v1 += len(vv)
|
|
|
|
# 纯视觉的 geom 指向抽稀版
|
|
for b in spec.bodies:
|
|
for g in b.geoms:
|
|
if (g.contype == 0 and g.conaffinity == 0 and g.meshname
|
|
and by_name.get(g.meshname) in visual
|
|
and by_name.get(g.meshname) in collide):
|
|
g.meshname = f"{g.meshname}_vis"
|
|
print(f" 网格顶点 {n_v0} → {n_v1}")
|
|
|
|
|
|
# ---------------------------------------------------------------- 场景搭建
|
|
|
|
# DexGraspBench example 数据集里的五个物体(灵巧手抓取那套方案用的就是这五个)。
|
|
# 目录名很长,这里给一组短名和中文标签。
|
|
DGB_ROOT = Path.home() / "DexGraspBench/assets/example_object"
|
|
OBJECTS = [
|
|
# 短名, 中文, 目录, 抓握截面 m, 质量 kg, rgba
|
|
# 抓握截面按手来定:L6 手张开 13.4cm、**合拢只剩 4cm**,截面超过 4~5cm 就夹不住
|
|
# (6.6cm 的罐子实测怎么放都会被挤出去)。这里统一缩到 4cm 上下。
|
|
("can", "易拉罐", "core_can_be67418a10003cc9eae3efbc9dbeea", 0.040, 0.15, (0.85, 0.30, 0.25, 1)),
|
|
("bottle", "瓶子", "core_bottle_523cddb320608c09a37f3fc191551700", 0.038, 0.12, (0.30, 0.65, 0.85, 1)),
|
|
("jar", "罐子", "core_jar_58c9d6575d62fbaf12bc68f36f3bdd45", 0.042, 0.16, (0.55, 0.75, 0.35, 1)),
|
|
("piano", "玩具琴", "sem_Piano_438f64dd9ff80bfddd5efa4c0e7932e", 0.045, 0.12, (0.85, 0.70, 0.25, 1)),
|
|
("toy", "玩偶", "sem_ToyFigure_a3087c9105fe0878d18a6b3b9816ca14", 0.040, 0.10, (0.70, 0.45, 0.80, 1)),
|
|
]
|
|
|
|
|
|
def load_obj(path):
|
|
"""极简 OBJ 读取:只要顶点和面(多边形按扇形三角化)。"""
|
|
vs, fs = [], []
|
|
for ln in Path(path).read_text().splitlines():
|
|
if ln.startswith("v "):
|
|
vs.append([float(x) for x in ln.split()[1:4]])
|
|
elif ln.startswith("f "):
|
|
idx = [int(t.split("/")[0]) - 1 for t in ln.split()[1:]]
|
|
for k in range(1, len(idx) - 1):
|
|
fs.append([idx[0], idx[k], idx[k + 1]])
|
|
return np.array(vs, float), np.array(fs, np.int64)
|
|
|
|
|
|
def add_objects(spec, wb, top_z, ox, ys):
|
|
"""把五个测试物体摆成一排。ox 是它们离机器人的距离,ys 是横向位置。"""
|
|
objs = []
|
|
for (nm, label, folder, width, mass, rgba), oy in zip(OBJECTS, ys):
|
|
root = DGB_ROOT / folder
|
|
vv, vf = load_obj(root / "mesh" / "simplified.obj")
|
|
ext = vv.max(0) - vv.min(0)
|
|
# 按"抓握截面"定缩放:取次长边缩到 width,长边就是高度
|
|
scale = width / float(sorted(ext)[1])
|
|
vv = (vv - (vv.max(0) + vv.min(0)) / 2) * scale # 顺便把几何中心挪到原点
|
|
|
|
# 立起来:原始网格的长轴是 y,转 90° 让它朝 z
|
|
rot = np.array([[1, 0, 0], [0, 0, -1], [0, 1, 0]], float)
|
|
vv = vv @ rot.T
|
|
half_h = float((vv[:, 2].max() - vv[:, 2].min()) / 2)
|
|
|
|
mv = spec.add_mesh()
|
|
mv.name = f"obj_{nm}_vis"
|
|
mv.uservert = np.round(vv, 5).reshape(-1).tolist()
|
|
mv.userface = vf.astype(np.int32).reshape(-1).tolist()
|
|
|
|
b = wb.add_body(name=f"obj_{nm}", pos=[ox, oy, top_z + half_h + 0.002])
|
|
b.add_freejoint()
|
|
# 视觉网格背质量(凸块按密度算出来的质量是凸分解伪影,重叠处被重复累加,
|
|
# 实测能虚高好几倍 —— 直接显式给质量)
|
|
b.add_geom(name=f"obj_{nm}", type=mujoco.mjtGeom.mjGEOM_MESH, meshname=mv.name,
|
|
rgba=list(rgba), mass=mass, contype=0, conaffinity=0, group=2)
|
|
|
|
pieces = sorted((root / "urdf" / "meshes").glob("convex_piece_*.obj"))
|
|
for k, pf in enumerate(pieces):
|
|
cv, cf = load_obj(pf)
|
|
cv = ((cv - (vv.max(0) + vv.min(0)) * 0) * scale) # 与视觉网格同一缩放
|
|
cv = cv - (load_obj(root / "mesh" / "simplified.obj")[0].max(0)
|
|
+ load_obj(root / "mesh" / "simplified.obj")[0].min(0)) / 2 * scale
|
|
cv = cv @ rot.T
|
|
mc = spec.add_mesh()
|
|
mc.name = f"obj_{nm}_col{k}"
|
|
mc.uservert = np.round(cv, 5).reshape(-1).tolist()
|
|
mc.userface = cf.astype(np.int32).reshape(-1).tolist()
|
|
b.add_geom(name=f"obj_{nm}_c{k}", type=mujoco.mjtGeom.mjGEOM_MESH,
|
|
meshname=mc.name, density=0, contype=1, conaffinity=1, condim=4,
|
|
friction=[1.2, 0.01, 0.0005], solref=[0.008, 1.0],
|
|
group=3, rgba=[1, 0, 0, 0.25])
|
|
|
|
objs.append({"name": nm, "label": label, "body": f"obj_{nm}",
|
|
"geom": f"obj_{nm}", "kind": "mesh", "mass": mass,
|
|
"half_h": half_h, "pieces": len(pieces),
|
|
"home": [ox, oy, top_z + half_h + 0.002]})
|
|
return objs
|
|
|
|
|
|
def add_scene(spec: mujoco.MjSpec, cfg: Cfg):
|
|
"""地面、灯光、操作台、三个测试物体。"""
|
|
wb = spec.worldbody
|
|
tex = spec.add_texture(
|
|
name="grid", type=mujoco.mjtTexture.mjTEXTURE_2D,
|
|
builtin=mujoco.mjtBuiltin.mjBUILTIN_CHECKER,
|
|
rgb1=[0.17, 0.20, 0.25], rgb2=[0.12, 0.15, 0.19], width=300, height=300)
|
|
del tex
|
|
mat = spec.add_material(name="grid", texrepeat=[8, 8], reflectance=0.05)
|
|
mat.textures[mujoco.mjtTextureRole.mjTEXROLE_RGB] = "grid"
|
|
wb.add_geom(name="floor", type=mujoco.mjtGeom.mjGEOM_PLANE, size=[0, 0, 0.05],
|
|
material="grid", contype=1, conaffinity=1, condim=3)
|
|
wb.add_light(pos=[0.4, 0, 3.0], dir=[0, 0, -1],
|
|
type=mujoco.mjtLightType.mjLIGHT_DIRECTIONAL)
|
|
|
|
# 操作台:台面 + 四条腿(纯装饰,只有台面参与碰撞)
|
|
half = (0.28, 0.42, 0.02)
|
|
top_z = cfg.table_top
|
|
t = wb.add_body(name="table", pos=[cfg.table_x, 0, 0])
|
|
t.add_geom(name="table_top", type=mujoco.mjtGeom.mjGEOM_BOX, size=list(half),
|
|
pos=[0, 0, top_z - half[2]], rgba=[0.55, 0.42, 0.30, 1],
|
|
contype=1, conaffinity=1, condim=3, friction=[1.0, 0.005, 0.0001])
|
|
# 台面还要一份纯视觉的副本:渲染器的"视觉"过滤会把有视觉兄弟的碰撞几何藏起来
|
|
# (桌腿是纯视觉的),不补这一份台面就画不出来,物体看着像飘在空中。
|
|
t.add_geom(name="table_top_vis", type=mujoco.mjtGeom.mjGEOM_BOX,
|
|
size=[half[0] * 1.001, half[1] * 1.001, half[2] * 0.999],
|
|
pos=[0, 0, top_z - half[2]], rgba=[0.55, 0.42, 0.30, 1],
|
|
contype=0, conaffinity=0, mass=0)
|
|
for sx in (-1, 1):
|
|
for sy in (-1, 1):
|
|
t.add_geom(name=f"table_leg_{sx}_{sy}", type=mujoco.mjtGeom.mjGEOM_BOX,
|
|
size=[0.02, 0.02, (top_z - 2 * half[2]) / 2],
|
|
pos=[sx * (half[0] - 0.03), sy * (half[1] - 0.03),
|
|
(top_z - 2 * half[2]) / 2],
|
|
rgba=[0.38, 0.29, 0.21, 1], contype=0, conaffinity=0)
|
|
|
|
objs = add_objects(spec, wb, top_z, cfg.table_x - 0.02, ys=[-0.26,-0.13,0.0,0.13,0.26])
|
|
return {"table": {"x": cfg.table_x, "top": top_z, "half": list(half)},
|
|
"objects": objs}
|
|
|
|
|
|
# ---------------------------------------------------------------- 换末端
|
|
|
|
SO101_XML = (Path.home() / "unitree_rl_mjlab/src/assets/robots/unitree_go2_so101"
|
|
/ "xmls/so101/so101.xml")
|
|
SO101_KP, SO101_KV, SO101_FRC = 998.22, 2.731, 2.94
|
|
|
|
|
|
def swap_gripper(spec: mujoco.MjSpec):
|
|
"""把 L6 五指手换成 SO-101 的二指夹爪。
|
|
|
|
五指手在这套场景里抓不住东西:张开 13.4cm、合拢只剩 4cm,四指从张开位置
|
|
扫向掌心的路径正好经过物体,合拢那一下必把物体挤飞(试过 18 组瞄准位置、
|
|
三种掌心朝向、拇指先合的分阶段闭合,都不行)。二指夹爪是已经在 Go2 上
|
|
验证过的方案,直接把它的 `gripper` 子树挂到腕部安装点。
|
|
"""
|
|
info = {}
|
|
for side in ("left", "right"):
|
|
# 每侧读一份:attach_body 会把源 spec 里的 body 取走,复用同一份第二次就是 None
|
|
grip_src = mujoco.MjSpec.from_file(str(SO101_XML))
|
|
for act in list(grip_src.actuators):
|
|
grip_src.delete(act) # 执行器这边自己重建
|
|
pre = "l" if side == "left" else "r"
|
|
hand = spec.body(f"hand_{side}_{pre}h_hand_base_link")
|
|
parent, pos, quat = hand.parent, list(hand.pos), list(hand.quat)
|
|
# 先删执行器再删 body,否则会留下指向已删关节的悬空执行器
|
|
for act in list(spec.actuators):
|
|
if act.name.startswith(f"hand_{side}_"):
|
|
spec.delete(act)
|
|
for eq in list(spec.equalities):
|
|
if (eq.name1 or "").startswith(f"hand_{side}_") or (eq.name2 or "").startswith(f"hand_{side}_"):
|
|
spec.delete(eq)
|
|
spec.delete(hand)
|
|
|
|
frame = parent.add_frame(pos=pos, quat=quat)
|
|
frame.attach_body(grip_src.body("gripper"), f"grip_{side}_", "")
|
|
a = spec.add_actuator()
|
|
a.name = f"grip_{side}_act"
|
|
a.target = f"grip_{side}_gripper"
|
|
a.trntype = mujoco.mjtTrn.mjTRN_JOINT
|
|
a.gaintype = mujoco.mjtGain.mjGAIN_FIXED
|
|
a.biastype = mujoco.mjtBias.mjBIAS_AFFINE
|
|
a.gainprm = [SO101_KP] + [0.0] * 9
|
|
a.biasprm = [0.0, -SO101_KP, -SO101_KV] + [0.0] * 7
|
|
a.forcerange = [-SO101_FRC, SO101_FRC]
|
|
a.forcelimited = mujoco.mjtLimited.mjLIMITED_TRUE
|
|
info[side] = {"prefix": f"grip_{side}_"}
|
|
spec.assets = {**(spec.assets or {}), **_so101_assets()}
|
|
return info
|
|
|
|
|
|
def _so101_assets():
|
|
d = SO101_XML.parent / "assets"
|
|
return {f.name: f.read_bytes() for f in d.glob("*")} if d.is_dir() else {}
|
|
|
|
|
|
# ---------------------------------------------------------------- 双 SO-101
|
|
|
|
SO101_JOINTS = ("shoulder_pan", "shoulder_lift", "elbow_flex",
|
|
"wrist_flex", "wrist_roll")
|
|
PAIR = { # 桌面式布局:两台并排装在台面上,都朝 +x
|
|
"table_top": 0.40, "table_x": 0.15, "table_half": (0.30, 0.40, 0.02),
|
|
"mount_y": 0.22, "obj_x": 0.27,
|
|
"obj_ys": (-0.20, -0.10, 0.0, 0.10, 0.20),
|
|
}
|
|
|
|
|
|
def build_so101_pair(cfg: Cfg):
|
|
"""从零搭一个场景:地面 + 操作台 + 台面上并排两台 SO-101。
|
|
|
|
选这个而不是给 P7 换爪:SO-101 只有 5 个自由度,任务约束正好是
|
|
位置 3 + 爪口朝向 2,IK 良定;而且这套夹爪 + 这个尺寸的物体在 Go2 抓放
|
|
demo 里已经端到端验过(随机位置 12 次成功 10 次),数值可以直接搬。
|
|
"""
|
|
spec = mujoco.MjSpec()
|
|
spec.option.timestep = 0.002
|
|
spec.option.integrator = mujoco.mjtIntegrator.mjINT_IMPLICITFAST
|
|
spec.option.iterations = 20
|
|
spec.option.ls_iterations = 20
|
|
wb = spec.worldbody
|
|
|
|
tex = spec.add_texture(
|
|
name="grid", type=mujoco.mjtTexture.mjTEXTURE_2D,
|
|
builtin=mujoco.mjtBuiltin.mjBUILTIN_CHECKER,
|
|
rgb1=[0.17, 0.20, 0.25], rgb2=[0.12, 0.15, 0.19], width=300, height=300)
|
|
del tex
|
|
mat = spec.add_material(name="grid", texrepeat=[8, 8], reflectance=0.05)
|
|
mat.textures[mujoco.mjtTextureRole.mjTEXROLE_RGB] = "grid"
|
|
wb.add_geom(name="floor", type=mujoco.mjtGeom.mjGEOM_PLANE, size=[0, 0, 0.05],
|
|
material="grid", contype=1, conaffinity=1, condim=3)
|
|
wb.add_light(pos=[0.3, 0, 2.5], dir=[0, 0, -1],
|
|
type=mujoco.mjtLightType.mjLIGHT_DIRECTIONAL)
|
|
|
|
P = PAIR
|
|
half, top = P["table_half"], P["table_top"]
|
|
t = wb.add_body(name="table", pos=[P["table_x"], 0, 0])
|
|
t.add_geom(name="table_top", type=mujoco.mjtGeom.mjGEOM_BOX, size=list(half),
|
|
pos=[0, 0, top - half[2]], rgba=[0.55, 0.42, 0.30, 1],
|
|
contype=1, conaffinity=1, condim=3, friction=[1.0, 0.005, 0.0001])
|
|
t.add_geom(name="table_top_vis", type=mujoco.mjtGeom.mjGEOM_BOX,
|
|
size=[half[0] * 1.001, half[1] * 1.001, half[2] * 0.999],
|
|
pos=[0, 0, top - half[2]], rgba=[0.55, 0.42, 0.30, 1],
|
|
contype=0, conaffinity=0, mass=0)
|
|
for sx in (-1, 1):
|
|
for sy in (-1, 1):
|
|
t.add_geom(name=f"table_leg_{sx}_{sy}", type=mujoco.mjtGeom.mjGEOM_BOX,
|
|
size=[0.02, 0.02, (top - 2 * half[2]) / 2],
|
|
pos=[sx * (half[0] - 0.03), sy * (half[1] - 0.03),
|
|
(top - 2 * half[2]) / 2],
|
|
rgba=[0.38, 0.29, 0.21, 1], contype=0, conaffinity=0)
|
|
|
|
for side, sy in (("left", 1), ("right", -1)):
|
|
arm = mujoco.MjSpec.from_file(str(SO101_XML))
|
|
frame = wb.add_frame(pos=[0.0, sy * P["mount_y"], top])
|
|
spec.attach(arm, prefix=f"arm_{side}_", frame=frame)
|
|
return spec, {"table": {"x": P["table_x"], "top": top, "half": list(half)}}
|
|
|
|
|
|
def ik_pos(m, d, qadr, rng, eepos, target, iters=400, axis_fn=None, axis_w=None):
|
|
"""位置(+可选末端朝向)IK,有限差分雅可比。用来算就绪位形。
|
|
|
|
就绪位形必须和前端默认的朝向约束**用同一套约束**解,否则页面一打开
|
|
就在跟自己较劲(实测残差 69mm,手臂一直在往回摆)。
|
|
"""
|
|
W = 0.25
|
|
q = np.zeros(len(qadr))
|
|
rows = 6 if axis_fn is not None else 3
|
|
|
|
def err(dd):
|
|
e = [target - eepos(dd)]
|
|
if axis_fn is not None:
|
|
e.append(W * (np.asarray(axis_w) - axis_fn(dd)))
|
|
return np.concatenate(e)
|
|
|
|
for _ in range(iters):
|
|
d.qpos[qadr] = np.clip(q, rng[:, 0], rng[:, 1])
|
|
mujoco.mj_kinematics(m, d)
|
|
e = err(d)
|
|
if np.linalg.norm(e[:3]) < 5e-4 and (rows == 3 or np.linalg.norm(e[3:]) / W < 0.05):
|
|
break
|
|
J = np.zeros((rows, len(qadr)))
|
|
for j in range(len(qadr)):
|
|
qh = q.copy()
|
|
qh[j] = np.clip(qh[j] + 1e-4, rng[j, 0], rng[j, 1])
|
|
dq = qh[j] - q[j]
|
|
if abs(dq) < 1e-12:
|
|
continue
|
|
d.qpos[qadr] = np.clip(qh, rng[:, 0], rng[:, 1])
|
|
mujoco.mj_kinematics(m, d)
|
|
J[:, j] = (err(d) - e) / -dq
|
|
q = np.clip(q + 0.5 * (J.T @ np.linalg.solve(J @ J.T + 1e-4 * np.eye(rows), e)),
|
|
rng[:, 0], rng[:, 1])
|
|
d.qpos[qadr] = np.clip(q, rng[:, 0], rng[:, 1])
|
|
mujoco.mj_kinematics(m, d)
|
|
return q, float(np.linalg.norm(target - eepos(d)))
|
|
|
|
|
|
# ---------------------------------------------------------------- 就绪位形
|
|
|
|
ARM_PREFIX = {"left": "L", "right": "R"} # 左臂关节叫 L1..L7,右臂叫 R1..R7
|
|
|
|
|
|
def arm_joints(side):
|
|
return [f"arm_{side}_{ARM_PREFIX[side]}{i}_Joint" for i in range(1, 8)]
|
|
|
|
|
|
def solve_ready(m, d, side, target):
|
|
"""把某一侧的腕部 site 解到 target,返回 7 个臂关节角。"""
|
|
jn = arm_joints(side)
|
|
qadr = np.array([m.jnt_qposadr[m.joint(n).id] for n in jn])
|
|
dadr = np.array([m.jnt_dofadr[m.joint(n).id] for n in jn])
|
|
rng = np.array([m.jnt_range[m.joint(n).id] for n in jn])
|
|
sid = m.site(f"arm_{side}_tool0").id
|
|
q = np.zeros(7)
|
|
jacp, jacr = np.zeros((3, m.nv)), np.zeros((3, m.nv))
|
|
for _ in range(300):
|
|
d.qpos[qadr] = q
|
|
mujoco.mj_kinematics(m, d)
|
|
mujoco.mj_comPos(m, d)
|
|
e = target - d.site_xpos[sid]
|
|
if np.linalg.norm(e) < 5e-4:
|
|
break
|
|
mujoco.mj_jacSite(m, d, jacp, jacr, sid)
|
|
J = jacp[:, dadr]
|
|
q = np.clip(q + 0.6 * (J.T @ np.linalg.solve(J @ J.T + 1e-4 * np.eye(3), e)),
|
|
rng[:, 0], rng[:, 1])
|
|
return q.tolist(), qadr.tolist(), float(np.linalg.norm(target - d.site_xpos[sid]))
|
|
|
|
|
|
def build_p7(cfg: Cfg):
|
|
"""linker 双臂工作站 + 操作台 + 物体。"""
|
|
src = WS_ROOT / cfg.workstation / "workstation.mjcf"
|
|
assert src.exists(), f"找不到工作站:{src}"
|
|
# MjSpec 不认 .mjcf 后缀,也不会跟着 ../../ 去找网格。
|
|
# 读成字符串 + 把引用到的网格按**原样的相对路径**塞进 assets 字典。
|
|
import re as _re
|
|
text = src.read_text()
|
|
assets = {}
|
|
for ref in sorted(set(_re.findall(r'file="([^"]+)"', text))):
|
|
f = (src.parent / ref).resolve()
|
|
if f.is_file():
|
|
assets[ref] = f.read_bytes()
|
|
else:
|
|
print(f" [警告] 网格找不到:{ref}")
|
|
print(f" 载入网格 {len(assets)} 个")
|
|
spec = mujoco.MjSpec.from_string(text, assets=assets)
|
|
if cfg.gripper == "so101":
|
|
swap_gripper(spec)
|
|
print(" 末端已换成 SO-101 二指夹爪")
|
|
info = add_scene(spec, cfg)
|
|
spec.option.timestep = 0.002
|
|
spec.option.integrator = mujoco.mjtIntegrator.mjINT_IMPLICITFAST
|
|
spec.option.iterations = 20
|
|
spec.option.ls_iterations = 20
|
|
return spec, info
|
|
|
|
|
|
def main() -> None:
|
|
cfg = tyro.cli(Cfg)
|
|
out = HERE / "scenes" / cfg.name
|
|
out.mkdir(parents=True, exist_ok=True)
|
|
|
|
if cfg.rig == "so101_pair":
|
|
spec, info = build_so101_pair(cfg)
|
|
info["objects"] = add_objects(spec, spec.worldbody, PAIR["table_top"],
|
|
PAIR["obj_x"], list(PAIR["obj_ys"]))
|
|
print(f" 双 SO-101:台面 z={PAIR['table_top']} 底座 y=±{PAIR['mount_y']}")
|
|
else:
|
|
spec, info = build_p7(cfg)
|
|
m = spec.compile()
|
|
d = mujoco.MjData(m)
|
|
|
|
# 就绪位形:末端停在台面上方
|
|
ready = {}
|
|
if cfg.rig == "so101_pair":
|
|
P = PAIR
|
|
for side, sy in (("left", 1), ("right", -1)):
|
|
jn = [f"arm_{side}_{j}" for j in SO101_JOINTS]
|
|
qadr = np.array([m.jnt_qposadr[m.joint(n).id] for n in jn])
|
|
rng = np.array([m.jnt_range[m.joint(n).id] for n in jn])
|
|
ga = [g for g in range(m.ngeom)
|
|
if (m.geom(g).name or "").startswith(f"arm_{side}_fixed_jaw_sph_tip")]
|
|
gb = [g for g in range(m.ngeom)
|
|
if (m.geom(g).name or "").startswith(f"arm_{side}_moving_jaw_sph_tip")]
|
|
|
|
def jaw_mid(dd, ga=ga, gb=gb):
|
|
return (dd.geom_xpos[ga].mean(0) + dd.geom_xpos[gb].mean(0)) / 2
|
|
|
|
sid = m.site(f"arm_{side}_gripperframe").id
|
|
|
|
def jaw_axis(dd, sid=sid):
|
|
return dd.site_xmat[sid].reshape(3, 3)[:, 0] # 爪口进近轴 = site 局部 +x
|
|
|
|
# 爪口朝下的可达高度有限:离台面 6cm 残差 0.3mm、10cm 9mm、14cm 就 38mm 了
|
|
tgt = np.array([0.22, sy * (P["mount_y"] - 0.05), P["table_top"] + 0.07])
|
|
q, err = ik_pos(m, d, qadr, rng, jaw_mid, tgt,
|
|
axis_fn=jaw_axis, axis_w=[0, 0, -1])
|
|
ready[side] = {"q": q.tolist(), "qadr": qadr.tolist(), "err": err}
|
|
print(f" 就绪位形 {side:5s} 残差 {err * 1000:.1f} mm")
|
|
else:
|
|
for side, ty in (("left", 0.18), ("right", -0.18)):
|
|
q, qadr, err = solve_ready(m, d, side,
|
|
np.array([cfg.table_x - 0.06, ty, cfg.table_top + 0.16]))
|
|
ready[side] = {"q": q, "qadr": qadr, "err": err}
|
|
print(f" 就绪位形 {side:5s} 残差 {err * 1000:.1f} mm")
|
|
|
|
qpos0 = np.array(m.qpos0, dtype=float).copy()
|
|
for side in ("left", "right"):
|
|
qpos0[ready[side]["qadr"]] = ready[side]["q"]
|
|
key = spec.add_key()
|
|
key.name = "ready"
|
|
key.qpos = qpos0.tolist()
|
|
key.ctrl = [0.0] * m.nu
|
|
|
|
inline_meshes(spec, m, cfg.decimate)
|
|
xml = spec.to_xml().replace(' class=""', "")
|
|
(out / "model.xml").write_text(xml)
|
|
mb = (out / "model.xml").stat().st_size / 1e6
|
|
|
|
# 重新编译一遍,确认自包含且几何没跑偏,同时取出前端要的索引
|
|
m2 = mujoco.MjModel.from_xml_path(str(out / "model.xml"))
|
|
d2 = mujoco.MjData(m2)
|
|
mujoco.mj_resetDataKeyframe(m2, d2, 0)
|
|
mujoco.mj_forward(m2, d2)
|
|
|
|
def act_idx(prefix):
|
|
return [i for i in range(m2.nu) if m2.actuator(i).name.startswith(prefix)]
|
|
|
|
arms, hands = {}, {}
|
|
gname = lambda g: m2.geom(g).name or "" # noqa: E731
|
|
for side in ("left", "right"):
|
|
if cfg.rig == "so101_pair":
|
|
pre = f"arm_{side}_"
|
|
a = [i for i in range(m2.nu)
|
|
if m2.actuator(i).name in [f"{pre}{j}" for j in SO101_JOINTS]]
|
|
arms[side] = {
|
|
"act": a,
|
|
"names": [m2.actuator(i).name for i in a],
|
|
"qadr": [int(m2.jnt_qposadr[int(m2.actuator_trnid[i][0])]) for i in a],
|
|
"range": [[float(x) for x in m2.actuator_ctrlrange[i]] for i in a],
|
|
"site": f"{pre}gripperframe",
|
|
"site_id": int(m2.site(f"{pre}gripperframe").id),
|
|
"ready": ready[side]["q"],
|
|
}
|
|
gi = [i for i in range(m2.nu) if m2.actuator(i).name == f"{pre}gripper"]
|
|
hands[side] = {
|
|
"kind": "so101",
|
|
"act": gi,
|
|
"names": [m2.actuator(i).name for i in gi],
|
|
"range": [[float(x) for x in m2.actuator_ctrlrange[i]] for i in gi],
|
|
"geoms": [g for g in range(m2.ngeom)
|
|
if (m2.body(int(m2.geom_bodyid[g])).name or "").startswith(pre)
|
|
and (m2.geom_contype[g] or m2.geom_conaffinity[g])],
|
|
# 夹持中心 = 两爪指尖球的中点(Go2 抓放 demo 验证过的算法)
|
|
"jaw_a": [g for g in range(m2.ngeom) if gname(g).startswith(f"{pre}fixed_jaw_sph_tip")],
|
|
"jaw_b": [g for g in range(m2.ngeom) if gname(g).startswith(f"{pre}moving_jaw_sph_tip")],
|
|
"palm_body": int(m2.body(f"{pre}gripper").id),
|
|
"site_id": int(m2.site(f"{pre}gripperframe").id),
|
|
}
|
|
continue
|
|
|
|
a = act_idx(f"arm_{side}_")
|
|
h = act_idx(f"hand_{side}_")
|
|
arms[side] = {
|
|
"act": a,
|
|
"names": [m2.actuator(i).name for i in a],
|
|
"qadr": [int(m2.jnt_qposadr[int(m2.actuator_trnid[i][0])]) for i in a],
|
|
"range": [[float(x) for x in m2.actuator_ctrlrange[i]] for i in a],
|
|
"site": f"arm_{side}_tool0",
|
|
"site_id": int(m2.site(f"arm_{side}_tool0").id),
|
|
"ready": ready[side]["q"],
|
|
}
|
|
if cfg.gripper == "so101":
|
|
pre = f"grip_{side}_"
|
|
h = [i for i in range(m2.nu) if m2.actuator(i).name.startswith(pre)]
|
|
gname = lambda g: m2.geom(g).name or "" # noqa: E731
|
|
hands[side] = {
|
|
"kind": "so101",
|
|
"act": h,
|
|
"names": [m2.actuator(i).name for i in h],
|
|
"range": [[float(x) for x in m2.actuator_ctrlrange[i]] for i in h],
|
|
"geoms": [g for g in range(m2.ngeom)
|
|
if (m2.body(int(m2.geom_bodyid[g])).name or "").startswith(pre)
|
|
and (m2.geom_contype[g] or m2.geom_conaffinity[g])],
|
|
# 夹持中心 = 两爪指尖球的中点(Go2 上验证过的算法)
|
|
"jaw_a": [g for g in range(m2.ngeom) if gname(g).startswith(f"{pre}fixed_jaw_sph_tip")],
|
|
"jaw_b": [g for g in range(m2.ngeom) if gname(g).startswith(f"{pre}moving_jaw_sph_tip")],
|
|
"palm_body": int(m2.body(f"{pre}gripper").id),
|
|
"site": f"{pre}gripperframe",
|
|
"site_id": int(m2.site(f"{pre}gripperframe").id),
|
|
}
|
|
else:
|
|
pre = f"hand_{side}_{'l' if side == 'left' else 'r'}h_"
|
|
hands[side] = {
|
|
"kind": "l6",
|
|
"act": h,
|
|
"names": [m2.actuator(i).name for i in h],
|
|
"range": [[float(x) for x in m2.actuator_ctrlrange[i]] for i in h],
|
|
# 原始资产的 geom 是靠 class 配的、大多没名字,只能按所属 body 来筛
|
|
"geoms": [g for g in range(m2.ngeom)
|
|
if (m2.body(int(m2.geom_bodyid[g])).name or "").startswith(f"hand_{side}_")
|
|
and (m2.geom_contype[g] or m2.geom_conaffinity[g])],
|
|
"tip_bodies": [int(m2.body(f"{pre}{t}_tip").id)
|
|
for t in ("thumb", "index", "middle", "ring", "pinky")],
|
|
"palm_body": int(m2.body(f"{pre}hand_base_link").id),
|
|
}
|
|
|
|
meta = {
|
|
"name": cfg.name,
|
|
"workstation": cfg.workstation,
|
|
"gripper": cfg.gripper,
|
|
"sim": {"timestep": float(m2.opt.timestep), "substeps": 10},
|
|
"nq": int(m2.nq), "nu": int(m2.nu), "ngeom": int(m2.ngeom),
|
|
"kps": [float(m2.actuator_gainprm[i][0]) for i in range(m2.nu)],
|
|
"arms": arms, "hands": hands,
|
|
"scene": info,
|
|
"xmlMB": round(mb, 2),
|
|
}
|
|
# 物体的 qpos 地址(自由关节),前端要用来改位置
|
|
for o in meta["scene"]["objects"]:
|
|
bid = m2.body(o["body"]).id
|
|
o["qadr"] = int(m2.jnt_qposadr[m2.body_jntadr[bid]])
|
|
o["body_id"] = int(bid)
|
|
o["geom_id"] = int(m2.geom(o["geom"]).id)
|
|
(out / "meta.json").write_text(json.dumps(meta, ensure_ascii=False))
|
|
|
|
print(f"已导出 {cfg.name} → {out}")
|
|
print(f" model.xml {mb:.2f} MB | nq={m2.nq} nu={m2.nu} ngeom={m2.ngeom} nmesh={m2.nmesh}")
|
|
print(f" 桌面 z={info['table']['top']} 中心 x={info['table']['x']} | 物体 "
|
|
+ ", ".join(o["name"] for o in meta["scene"]["objects"]))
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|