"""Build Franka lift replay bundles (Study D) for the existing trace viewer. Usage: python replay_export.py --recordings ROOT --panda PART_ONE_GRASP_BUNDLE --output DIR Poses are the native body link poses recorded at the 30 Hz policy rate (lift_eval.py --record_envs), expressed in each environment's frame and rotated 180 degrees about z so the robot faces +x as in part one's replays. Display geometry: part one's Panda link meshes (including the finger pads) and each episode's object as the exact primitive the task spawned (cuboid, sphere or capsule), identified by matching the extent of points sampled on its simulated mesh to the task's eight graspable shapes. The translucent sphere marks the commanded target with the 5 cm success radius. Green arrows: net contact force on each finger (drawn near the fingertip). Purple arrows: the policy's relative joint command for each arm joint (radians added to the current angle), drawn at the joint along the motion it drives the hand, 0.3 m per radian. An episode that ends early holds its last recorded pose. """ import argparse import gzip import hashlib import json import math from pathlib import Path import numpy as np COORDINATES = dict(frame='world', up='z', quaternion='wxyz', units='SI') COLORS = dict(isaacsim_physx=0x315d9e, newton_mjwarp=0xb65335) NAMES = {**{f'panda_link{i}': f'link{i}' for i in range(8)}, 'panda_hand': 'hand', 'panda_leftfinger': 'left_finger', 'panda_rightfinger': 'right_finger'} COMMAND_GAIN_M_PER_RAD = .3 FINGERTIP_OFFSET_M = (0., 0., .045) # along the finger frame's z axis, near the pad SUCCESS_RADIUS_M = .05 # Isaac-Lift-Franka's graspable shapes (franka_env_cfg.py), centred on the object origin, capsules along z. SHAPES = [('box', (.05, .05, .05)), ('box', (.025, .05, .05)), ('box', (.025, .025, .05)), ('box', (.01, .05, .05)), ('sphere', .02), ('capsule', (.025, .1)), ('capsule', (.025, .2)), ('capsule', (.01, .2))] def shape_extent(kind, size): if kind == 'box': return size if kind == 'sphere': return (2*size,)*3 r, h = size return (2*r, 2*r, h+2*r) def capsule_mesh(r, h, n=32, rings=8): """Capsule along z (cylinder length h plus hemispherical caps), outward-wound triangles.""" rows = [(math.pi/2*i/rings, h/2) for i in range(rings+1)]+[(math.pi/2*(1+i/rings), -h/2) for i in range(rings+1)] v = [[r*math.sin(phi)*math.cos(2*math.pi*j/n), r*math.sin(phi)*math.sin(2*math.pi*j/n), r*math.cos(phi)+off] for phi, off in rows for j in range(n)] f = [] for i in range(len(rows)-1): for j in range(n): a, b, c, d = i*n+j, i*n+(j+1) % n, (i+1)*n+j, (i+1)*n+(j+1) % n f += [[a, c, b], [b, c, d]] return [[round(x, 6) for x in p] for p in v], f def object_body(extent): kind, size = min(SHAPES, key=lambda s: max(abs(a-b) for a, b in zip(shape_extent(*s), extent))) if max(abs(a-b) for a, b in zip(shape_extent(kind, size), extent)) > .005: raise ValueError(f'object extent {extent} matches no known shape') if kind == 'box': return dict(id='cube', shape='box', size=list(size)), (kind, size) if kind == 'sphere': return dict(id='cube', shape='sphere', radius=size), (kind, size) v, f = capsule_mesh(*size) return dict(id='cube', shape='mesh', vertices=v, faces=f), (kind, size) def quat_matrix(w, x, y, z): 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)]]) def turn(p): """Rotate a point or vector 180 degrees about z.""" return [-p[0], -p[1], p[2]] def turn_quat(q): """Left-multiply a (w, x, y, z) quaternion by a 180-degree rotation about z.""" w, x, y, z = q return [-z, -y, x, w] def trace(rec, env, engine, label, panda, reference): data = rec['envs'][env] names = rec['body_names'] hull = rec['object_hulls'][env] table_pos = turn(rec['table']['pose']) obj_body, obj_shape = object_body(hull['extent_m']) bodies = [panda[NAMES[n]] for n in names] + [ obj_body, dict(id='table', shape='box', size=rec['table']['size_m']), dict(id='target', shape='sphere', radius=SUCCESS_RADIUS_M)] scene = dict(task='panda_lift', floor=True, dt=rec['dt_s'], duration=data['frames'][-1]['time_s'], bodies=bodies, object_extent_m=hull['extent_m'], object_shape=dict(kind=obj_shape[0], size=obj_shape[1])) scene_sha = hashlib.sha256(json.dumps(scene, sort_keys=True).encode()).hexdigest() arm = [j for j in rec['joint_names'] if j.startswith('panda_joint')] frames, previous = [], None for raw in data['frames']: poses = np.array(raw['robot'], dtype=float) poses[:, :3] = [turn(p) for p in poses[:, :3]] poses[:, 3:] = [turn_quat(q/np.linalg.norm(q)) for q in poses[:, 3:]] obj = np.array(raw['object'], dtype=float) obj_p, obj_q = np.array(turn(obj[:3])), np.array(turn_quat(obj[3:]/np.linalg.norm(obj[3:]))) hand = poses[names.index('panda_hand'), :3] forces = {} if raw['time_s'] > 0: for jn, delta in zip(rec['joint_names'], raw['joint_target']): if jn not in arm: continue child = f'panda_link{jn[-1]}' ci = names.index(child) axis = quat_matrix(*poses[ci, 3:]) @ np.array([0., 0., 1.]) # Franka joints turn about the child's z axis tangent = np.cross(axis, hand - poses[ci, :3]) if np.linalg.norm(tangent) < 1e-6: continue vector = COMMAND_GAIN_M_PER_RAD*delta*tangent/np.linalg.norm(tangent) forces.setdefault(child, []).append(dict(on_body=NAMES[child], kind='command', point_m=poses[ci, :3].tolist(), vector_N=vector.tolist(), joint=jn, command_rad=delta)) for finger, f in zip(('panda_leftfinger', 'panda_rightfinger'), raw['finger_force_N']): if np.linalg.norm(f) > 1e-6: fi = names.index(finger) tip = poses[fi, :3] + quat_matrix(*poses[fi, 3:]) @ np.array(FINGERTIP_OFFSET_M) forces.setdefault(finger, []).append(dict(on_body=NAMES[finger], kind='contact', point_m=tip.tolist(), vector_N=turn(f))) out = [] for i, n in enumerate(names): out.append(dict(id=NAMES[n], position_m=poses[i, :3].tolist(), quaternion_wxyz=poses[i, 3:].tolist(), velocity_m_s=[0., 0., 0.], angular_velocity_rad_s=[0., 0., 0.], forces=forces.get(n, []))) v = [0., 0., 0.] if previous is None else ((obj_p-previous[0])/(raw['time_s']-previous[1])).tolist() out += [dict(id='cube', position_m=obj_p.tolist(), quaternion_wxyz=obj_q.tolist(), velocity_m_s=v, angular_velocity_rad_s=[0., 0., 0.], forces=[]), dict(id='table', position_m=table_pos, quaternion_wxyz=[1., 0., 0., 0.], velocity_m_s=[0., 0., 0.], angular_velocity_rad_s=[0., 0., 0.], forces=[]), dict(id='target', position_m=turn(raw['target']), quaternion_wxyz=[1., 0., 0., 0.], velocity_m_s=[0., 0., 0.], angular_velocity_rad_s=[0., 0., 0.], forces=[])] frames.append(dict(time_s=raw['time_s'], bodies=out, succeeded=raw['succeeded'])) previous = (obj_p, raw['time_s']) return dict(schema='simcompare.trace.v2', engine=engine, label=label, color=COLORS[engine], coordinates=COORDINATES, scene=scene, scene_sha256=scene_sha, reference_sha256=reference, frames=frames, status='complete' if data['termination'] == 'time_out' else 'stopped', stop_time_s=frames[-1]['time_s'], outcome=dict(success=data['success'], termination=data['termination'], first_success_s=data['first_success_s']), display_note='Native body link poses at 30 Hz; Panda display meshes from part one, object as the exact primitive the task spawned.') def build(recordings, panda, output, name, panels, notice): """panels: list of (recording file stem, env, engine, label).""" traces, sources = [], [] for stem, env, engine, label in panels: path = recordings/f'{stem}-recording.json.gz' meta = json.loads((recordings/f'{stem}.json').read_text()) rec = json.loads(gzip.decompress(path.read_bytes())) ref = hashlib.sha256(meta['checkpoint'].encode()).hexdigest() traces.append(trace(rec, str(env), engine, label, panda, ref)) sources.append(dict(recording=path.name, env=env, checkpoint=meta['checkpoint'], condition=meta['condition'], hydra=meta['hydra'], seed=meta['seed'], recording_sha256=hashlib.sha256(path.read_bytes()).hexdigest(), outcome=traces[-1]['outcome'], batch_success_rate=meta['result']['success_rate'], batch_n=meta['result']['n'])) out = output/name out.write_bytes(gzip.compress(json.dumps(dict(schema='simcompare.bundle.v1', traces=traces, notice=notice), separators=(',', ':')).encode(), mtime=0)) return dict(sha256=hashlib.sha256(out.read_bytes()).hexdigest(), notice=notice, sources=sources) def load_panda(path): bundle = json.loads(gzip.decompress(Path(path).read_bytes())) return {b['id']: b for b in bundle['traces'][0]['scene']['bodies'] if b['id'] in NAMES.values()} if __name__ == '__main__': p = argparse.ArgumentParser() p.add_argument('--recordings', type=Path, required=True) p.add_argument('--panda', type=Path, required=True) p.add_argument('--output', type=Path, required=True) p.add_argument('--spec', type=Path, required=True, help='JSON: {bundle name: {notice, panels: [[stem, env, engine, label], ...]}}') a = p.parse_args() panda = load_panda(a.panda) spec = json.loads(a.spec.read_text()) a.output.mkdir(parents=True, exist_ok=True) provenance = dict(schema='lift.study-d.replay-provenance.v1', panda_meshes=a.panda.name, panda_meshes_sha256=hashlib.sha256(a.panda.read_bytes()).hexdigest(), bundles={}) for name, s in spec.items(): provenance['bundles'][name] = build(a.recordings, panda, a.output, name, [tuple(x) for x in s['panels']], s['notice']) (a.output/'replay-provenance.json').write_text(json.dumps(provenance, indent=1)+'\n')