scripts/yam_diag.py

"""YAM diagnostic: action layout, IK command frame, gripper convention."""
import argparse, sys, os
from isaaclab.app import AppLauncher
parser = argparse.ArgumentParser(); AppLauncher.add_app_launcher_args(parser)
args = parser.parse_args(); args.headless=True; args.enable_cameras=True
app = AppLauncher(args).app
import numpy as np, torch, 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  # noqa
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
obs,_=env.reset()
dev="cuda:0"
# action terms
am=u.action_manager
print("[d] action terms:", [(t, am.get_term(t).action_dim) for t in am.active_terms], flush=True)
print("[d] total action dim:", am.total_action_dim, flush=True)
robot=u.scene["right_robot"]; bn=list(robot.data.body_names); origin=u.scene.env_origins[0].cpu().numpy()
def bw(link):
    i=bn.index(link); return robot.data.body_pos_w[0,i].cpu().numpy()-origin, robot.data.body_quat_w[0,i].cpu().numpy()
l6p,l6q=bw("link_6"); lfp,_=bw("left_finger"); rfp,_=bw("right_finger")
root=robot.data.root_pos_w[0].cpu().numpy()-origin; rootq=robot.data.root_quat_w[0].cpu().numpy()
print(f"[d] right root={np.round(root,3)} quat={np.round(rootq,3)}", flush=True)
print(f"[d] link_6={np.round(l6p,3)} quat={np.round(l6q,3)}", flush=True)
print(f"[d] left_finger={np.round(lfp,3)} right_finger={np.round(rfp,3)} sep={np.linalg.norm(lfp-rfp):.4f}", flush=True)
def R(q):
    w,x,y,z=q; return np.array([[1-2*(y*y+z*z),2*(x*y-z*w),2*(x*z+y*w)],[2*(x*y+z*w),1-2*(x*x+z*z),2*(y*z-x*w)],[2*(x*z-y*w),2*(y*z+x*w),1-2*(x*x+y*y)]])
Rl6=R(l6q)
close_w=(lfp-rfp)/(np.linalg.norm(lfp-rfp)+1e-9)
mid=0.5*(lfp+rfp); approach_w=mid-l6p; approach_w=approach_w/(np.linalg.norm(approach_w)+1e-9)
print(f"[d] CLOSE axis in link6 frame: {np.round(Rl6.T@close_w,2)} -> {'xyz'[int(np.argmax(abs(Rl6.T@close_w)))]}", flush=True)
print(f"[d] APPROACH axis in link6 frame: {np.round(Rl6.T@approach_w,2)} -> {'xyz'[int(np.argmax(abs(Rl6.T@approach_w)))]}", flush=True)
# read the diff-ik action term to learn command frame: many IsaacLab diffik use root frame.
# stay test: feed current EEF pose (in root frame) as right-arm command
from isaaclab.utils.math import subtract_frame_transforms, quat_from_matrix
# EEF (link_6 + body_offset 0.13 along its z) in root frame
eef_w = l6p + Rl6@np.array([0,0,0.13]); eef_q = l6q
# to root frame
Rroot=R(rootq)
eef_root_p = Rroot.T@(eef_w - root); 
print(f"[d] EEF(world)={np.round(eef_w,3)} EEF(root)={np.round(eef_root_p,3)}", flush=True)
def act(rp, rq, rg, lp, lq, lg):
    # order from action terms print; assume [L_arm7, L_grip1, R_arm7, R_grip1]
    return torch.tensor(np.concatenate([lp,lq,[lg], rp,rq,[rg]]),dtype=torch.float32,device=dev).view(1,-1)
# hold both at current EEF-root poses
l6pL,l6qL=None,None
rl=u.scene["left_robot"]; bnl=list(rl.data.body_names)
il=bnl.index("link_6"); l6pL=rl.data.body_pos_w[0,il].cpu().numpy()-origin; l6qL=rl.data.body_quat_w[0,il].cpu().numpy()
rootL=rl.data.root_pos_w[0].cpu().numpy()-origin; rootqL=rl.data.root_quat_w[0].cpu().numpy()
eefL_w=l6pL+R(l6qL)@np.array([0,0,0.13]); eefL_root=R(rootqL).T@(eefL_w-rootL)
for _ in range(30):
    a=act(eef_root_p.astype(np.float32),eef_q.astype(np.float32),-1.0, eefL_root.astype(np.float32),l6qL.astype(np.float32),-1.0)
    obs,_,_,_,_=env.step(a)
l6p2,_=bw("link_6"); eef2=l6p2+R(bw('link_6')[1])@np.array([0,0,0.13])
print(f"[d] STAY test: EEF now(world)={np.round(eef2,3)} (was {np.round(eef_w,3)}) err={np.linalg.norm(eef2-eef_w):.3f}", flush=True)
# move down 0.1 in root frame
for _ in range(40):
    a=act((eef_root_p+np.array([0,0,-0.1])).astype(np.float32),eef_q.astype(np.float32),-1.0, eefL_root.astype(np.float32),l6qL.astype(np.float32),-1.0)
    obs,_,_,_,_=env.step(a)
l6p3,_=bw("link_6"); eef3=l6p3+R(bw('link_6')[1])@np.array([0,0,0.13])
print(f"[d] MOVE -0.1z(root) test: EEF now(world)={np.round(eef3,3)} delta={np.round(eef3-eef_w,3)}", flush=True)
env.close(); app.close(); print("YAM_DIAG_OK", flush=True)