| """Export YAM right-arm per-link meshes (in link-local frame) to OBJ for the web viewer.""" |
| import argparse, sys, os |
| from isaaclab.app import AppLauncher |
| p=argparse.ArgumentParser(); AppLauncher.add_app_launcher_args(p); a=p.parse_args(); a.headless=True; a.enable_cameras=False |
| app=AppLauncher(a).app |
| import numpy as np, gymnasium as gym |
| REPO=os.path.dirname(os.path.dirname(os.path.abspath(__file__))); sys.path.insert(0,os.path.join(REPO,"source")) |
| import bimanual.tasks.manager_based.yam |
| from isaaclab_tasks.utils import parse_env_cfg |
| TASK="Template-YAM-Play-v0" |
| env=gym.make(TASK, cfg=parse_env_cfg(TASK, device="cuda:0", num_envs=1)); u=env.unwrapped; env.reset() |
| R=u.scene["right_robot"] |
| import omni.usd |
| from pxr import Usd, UsdGeom, Gf |
| stage=omni.usd.get_context().get_stage() |
| root_path=R.root_physx_view.prim_paths[0] |
| print("[m] root:", root_path, flush=True) |
| xf=UsdGeom.XformCache(Usd.TimeCode.Default()) |
| OUT="/home/horde/project/xiaotong/outputs/yam_web/mesh" |
| os.makedirs(OUT, exist_ok=True) |
|
|
| def find_link_prim(name): |
| for pr in Usd.PrimRange(stage.GetPrimAtPath(root_path)): |
| if pr.GetName()==name: |
| return pr |
| return None |
|
|
| def gfmat_to_np(m): |
| return np.array([[m[i][j] for j in range(4)] for i in range(4)]) |
|
|
| LINKS=["arm","link_1","link_2","link_3","link_4","link_5","link_6","left_finger","right_finger"] |
| linkpath={}; Lwinv={} |
| for n in LINKS: |
| lp=find_link_prim(n) |
| if lp is None: continue |
| linkpath[lp.GetPath().pathString]=n |
| Lwinv[n]=np.linalg.inv(gfmat_to_np(xf.GetLocalToWorldTransform(lp))) |
| |
| |
| FRAME={n:n for n in LINKS} |
| FRAME["left_finger"]="link_6"; FRAME["right_finger"]="link_6" |
| def owning_link(prim): |
| p=prim.GetPath() |
| while not p.isEmpty and p.pathString!="/": |
| if p.pathString in linkpath: return linkpath[p.pathString] |
| p=p.GetParentPath() |
| return None |
| buck={n:{"V":[],"F":[]} for n in LINKS} |
| for pr in Usd.PrimRange(stage.GetPrimAtPath(root_path)): |
| if pr.GetTypeName()!="Mesh": continue |
| ln=owning_link(pr) |
| if ln is None: continue |
| mesh=UsdGeom.Mesh(pr); pts=mesh.GetPointsAttr().Get() |
| if not pts: continue |
| counts=mesh.GetFaceVertexCountsAttr().Get(); idx=mesh.GetFaceVertexIndicesAttr().Get() |
| Mw=gfmat_to_np(xf.GetLocalToWorldTransform(pr)) |
| M=Mw@Lwinv[FRAME[ln]] |
| V=buck[ln]["V"]; F=buck[ln]["F"]; base=len(V) |
| P=np.array([[p[0],p[1],p[2],1.0] for p in pts]); Pl=P@M |
| for q in Pl: V.append((q[0],q[1],q[2])) |
| o=0 |
| for c in counts: |
| for k in range(1,c-1): F.append((base+idx[o]+1, base+idx[o+k]+1, base+idx[o+k+1]+1)) |
| o+=c |
| manifest={} |
| for ln in LINKS: |
| V=buck[ln]["V"]; F=buck[ln]["F"] |
| if not V: print(f"[m] {ln}: no mesh", flush=True); continue |
| fn=os.path.join(OUT, ln+".obj") |
| with open(fn,"w") as f: |
| for v in V: f.write(f"v {v[0]:.5f} {v[1]:.5f} {v[2]:.5f}\n") |
| for t in F: f.write(f"f {t[0]} {t[1]} {t[2]}\n") |
| manifest[ln]={"verts":len(V),"faces":len(F)} |
| print(f"[m] {ln}: {len(V)} verts {len(F)} tris -> {fn}", flush=True) |
| import json; json.dump(manifest, open(os.path.join(OUT,"manifest.json"),"w")) |
| |
| import numpy as _np |
| origin_env=u.scene.env_origins[0].cpu().numpy() |
| bn=list(R.data.body_names) |
| def body_pose(name): |
| i=bn.index(name) |
| p=R.data.body_pos_w[0,i].cpu().numpy()-origin_env |
| q=R.data.body_quat_w[0,i].cpu().numpy() |
| return [float(x) for x in p],[float(x) for x in q] |
| |
| from pxr import UsdPhysics |
| axis_map={} |
| for pr in Usd.PrimRange(stage.GetPrimAtPath(root_path)): |
| if pr.IsA(UsdPhysics.RevoluteJoint): |
| j=UsdPhysics.RevoluteJoint(pr) |
| ax=j.GetAxisAttr().Get() |
| rel=j.GetBody1Rel().GetTargets() |
| child=rel[0].name if rel else pr.GetName() |
| axis_map[child]={"axis":str(ax)} |
| CHAIN=["arm","link_1","link_2","link_3","link_4","link_5","link_6"] |
| pose={} |
| for n in CHAIN: |
| try: pose[n]={"pos":body_pose(n)[0],"quat":body_pose(n)[1]} |
| except Exception as e: print("[m] pose err",n,e,flush=True) |
| out={"pose":pose,"axis":axis_map, |
| "home":[0.0,1.5708,0.7854,0.7854,0.0,0.0], |
| "joint_child":["link_1","link_2","link_3","link_4","link_5","link_6"]} |
| json.dump(out, open(os.path.join(OUT,"chain_truth.json"),"w")) |
| print("[m] pose:", {n:[round(v,4) for v in pose[n]["pos"]] for n in pose}, flush=True) |
| print("[m] axis_map:", axis_map, flush=True) |
| env.close(); app.close(); print("MESH_OK", flush=True) |
|
|
