191 lines
12 KiB
Python
191 lines
12 KiB
Python
"""Actuated L20 + freely simulated metric bottle, with a prescribed moving wrist."""
|
|
import os
|
|
os.environ.setdefault('MUJOCO_GL','osmesa')
|
|
import argparse,json
|
|
import xml.etree.ElementTree as ET
|
|
from pathlib import Path
|
|
import numpy as np
|
|
import mujoco
|
|
from scipy.spatial.transform import Rotation
|
|
from scipy.optimize import lsq_linear
|
|
from l20_calibrated import ROOT,Calibration
|
|
|
|
BASE=ROOT/'output/l20_2047635068'
|
|
OUT=BASE/'bottle_grasp'
|
|
|
|
def feedback_torque(model,data,cal,addr,objects,object_body,target_pos,target_velocity):
|
|
selected=[c for c in data.contact if any(g in objects for g in c.geom)]
|
|
if not selected:return None,None
|
|
center=data.xipos[object_body].copy();velocity=np.zeros(6)
|
|
mujoco.mj_objectVelocity(model,data,mujoco.mjtObj.mjOBJ_BODY,object_body,velocity,0)
|
|
desiredforce=np.array([0,0,model.body_mass[object_body]*9.81])+100*(target_pos-center)+2*(target_velocity-velocity[3:])
|
|
rotation=Rotation.from_matrix(data.xmat[object_body].reshape(3,3))
|
|
desiredtorque=-.08*rotation.as_rotvec()-.008*velocity[:3]
|
|
target=np.r_[desiredforce,desiredtorque/.03]
|
|
columns=[];jacobians=[];rays_all=[];preferred=[]
|
|
for c in selected:
|
|
object_first=c.geom[0] in objects
|
|
normal=-c.frame[:3]*(1 if object_first else -1)
|
|
tangent=np.cross(normal,[0,0,1])
|
|
if np.linalg.norm(tangent)<.01:tangent=np.cross(normal,[1,0,0])
|
|
tangent/=np.linalg.norm(tangent);second=np.cross(normal,tangent)
|
|
rays=np.array([normal+.6*(np.cos(a)*tangent+np.sin(a)*second) for a in np.arange(4)*np.pi/2]).T
|
|
rays_all.append(rays);columns.append(np.vstack([rays,np.cross((c.pos-center)[None,:],rays.T).T/.03]))
|
|
body=int(model.geom_bodyid[c.geom[1] if object_first else c.geom[0]])
|
|
preferred.append(6 if model.body(body).name.startswith('thumb') else 1.5)
|
|
jac=np.zeros((3,model.nv));mujoco.mj_jac(model,data,jac,None,c.pos,body)
|
|
ja=np.stack([jac[:,model.jnt_dofadr[model.joint(n).id]] for n in cal.active],axis=1)
|
|
for n,(parent,_,_) in cal.curves.items():
|
|
slope=cal.passive_value(n,data.qpos[addr[parent]])[1]
|
|
ja[:,cal.active.index(parent)]+=slope*jac[:,model.jnt_dofadr[model.joint(n).id]]
|
|
jacobians.append(ja)
|
|
A=np.concatenate(columns,axis=1);reg=np.zeros((len(selected),len(selected)*4))
|
|
for i in range(len(selected)):reg[i,i*4:i*4+4]=.025
|
|
sol=lsq_linear(np.vstack([A,reg]),np.r_[target,.025*np.array(preferred)],bounds=(0,10),tol=1e-6,max_iter=50)
|
|
tau=sum(j.T@(r@sol.x[i*4:i*4+4]) for i,(j,r) in enumerate(zip(jacobians,rays_all)))
|
|
for i,n in enumerate(cal.active):tau[i]+=data.qfrc_bias[model.jnt_dofadr[model.joint(n).id]]
|
|
return tau,float(np.linalg.norm(A@sol.x-target))
|
|
|
|
def build_physics(kp=30):
|
|
tree=ET.parse(BASE/'contact_object/contact_scene.xml');root=tree.getroot()
|
|
hand=root.find("worldbody/body[@name='hand_base_link']")
|
|
hand.remove(hand.find('freejoint'));hand.set('mocap','true')
|
|
option=root.find('option');option.set('timestep','.001');option.set('integrator','implicitfast');option.set('iterations','80')
|
|
root.find('default/joint').set('damping','.08')
|
|
eq=root.find('equality')
|
|
if eq is None:eq=ET.SubElement(root,'equality')
|
|
actuator=ET.SubElement(root,'actuator')
|
|
cal=Calibration()
|
|
for n in cal.active:
|
|
ET.SubElement(actuator,'position',name='motor_'+n,joint=n,kp=str(kp),kv='.3',
|
|
forcelimited='true',forcerange='-1 1')
|
|
for n,(parent,x,y) in cal.curves.items():
|
|
ET.SubElement(eq,'joint',name='lut_'+n,joint1=n,joint2=parent,
|
|
polycoef='0 1 0 0 0',solref='.003 1',solimp='.95 .99 .001')
|
|
path=OUT/'physics_scene.xml';tree.write(path)
|
|
return path
|
|
|
|
def run(seconds=3.,lift=False,kp=30.,preload=0.,force_preload=False,force_control=False,feedback=False):
|
|
path=build_physics(kp);model=mujoco.MjModel.from_xml_path(str(path));data=mujoco.MjData(model)
|
|
cal=Calibration();old=mujoco.MjModel.from_xml_path(str(BASE/'contact_object/contact_scene.xml'))
|
|
fitted=np.load(BASE/'contact_object/contact_result.npz')['qpos']
|
|
names=list(cal.tables)
|
|
addr={n:model.jnt_qposadr[mujoco.mj_name2id(model,mujoco.mjtObj.mjOBJ_JOINT,n)] for n in names}
|
|
for n in names:
|
|
oldaddr=old.jnt_qposadr[mujoco.mj_name2id(old,mujoco.mjtObj.mjOBJ_JOINT,n)]
|
|
data.qpos[addr[n]]=fitted[oldaddr]
|
|
ow=old.jnt_qposadr[mujoco.mj_name2id(old,mujoco.mjtObj.mjOBJ_JOINT,'wrist_free')]
|
|
oo=old.jnt_qposadr[mujoco.mj_name2id(old,mujoco.mjtObj.mjOBJ_JOINT,'object_free')]
|
|
obj=model.jnt_qposadr[mujoco.mj_name2id(model,mujoco.mjtObj.mjOBJ_JOINT,'object_free')]
|
|
# Rigidly reorient the entire registration so bottle +Z is world up.
|
|
rotation=Rotation.from_quat(fitted[oo+3:oo+7][[1,2,3,0]]).as_matrix().T
|
|
handpos=rotation@(fitted[ow:ow+3]-fitted[oo:oo+3])+[0,0,.3]
|
|
handrot=rotation@Rotation.from_quat(fitted[ow+3:ow+7][[1,2,3,0]]).as_matrix()
|
|
data.mocap_pos[0]=handpos;data.mocap_quat[0]=Rotation.from_matrix(handrot).as_quat()[[3,0,1,2]]
|
|
data.qpos[obj:obj+7]=[0,0,.3,1,0,0,0]
|
|
controls=np.array([data.qpos[addr[n]] for n in cal.active])
|
|
reference_controls=controls.copy()
|
|
for i,n in enumerate(cal.active):
|
|
if n.endswith(('pitch','pip','mcp')):controls[i]+=preload
|
|
controls[i]=np.clip(controls[i],*cal.bounds[n])
|
|
data.ctrl[:]=controls
|
|
objects=np.flatnonzero(model.geom_contype==2)
|
|
object_body=mujoco.mj_name2id(model,mujoco.mjtObj.mjOBJ_BODY,'grasp_object')
|
|
preload_report={}
|
|
if force_preload:
|
|
for n,(parent,x,y) in cal.curves.items():
|
|
a=data.qpos[addr[parent]];v,s=cal.passive_value(n,a)
|
|
ei=mujoco.mj_name2id(model,mujoco.mjtObj.mjOBJ_EQUALITY,'lut_'+n)
|
|
model.eq_data[ei,:5]=[v-s*a,s,0,0,0]
|
|
mujoco.mj_forward(model,data)
|
|
selected=[c for c in data.contact if any(g in objects for g in c.geom)]
|
|
columns=[];jacobians=[];cone=[]
|
|
center=data.xipos[object_body].copy()
|
|
for c in selected:
|
|
object_first=c.geom[0] in objects
|
|
outward=c.frame[:3]*(1 if object_first else -1)
|
|
normal=-outward
|
|
tangent=np.cross(normal,[0,0,1])
|
|
if np.linalg.norm(tangent)<.01:tangent=np.cross(normal,[1,0,0])
|
|
tangent/=np.linalg.norm(tangent);second=np.cross(normal,tangent)
|
|
rays=np.array([normal+.6*(np.cos(a)*tangent+np.sin(a)*second) for a in np.arange(4)*np.pi/2]).T
|
|
cone.append(rays)
|
|
columns.append(np.vstack([rays,np.cross((c.pos-center)[None,:],rays.T).T/.03]))
|
|
body=model.geom_bodyid[c.geom[1] if object_first else c.geom[0]]
|
|
jac=np.zeros((3,model.nv));mujoco.mj_jac(model,data,jac,None,c.pos,int(body))
|
|
activejac=np.stack([jac[:,model.jnt_dofadr[mujoco.mj_name2id(model,mujoco.mjtObj.mjOBJ_JOINT,n)]] for n in cal.active],axis=1)
|
|
for n,(parent,x,y) in cal.curves.items():
|
|
slope=cal.passive_value(n,data.qpos[addr[parent]])[1]
|
|
di=model.jnt_dofadr[mujoco.mj_name2id(model,mujoco.mjtObj.mjOBJ_JOINT,n)]
|
|
activejac[:,cal.active.index(parent)]+=slope*jac[:,di]
|
|
jacobians.append(activejac)
|
|
if not selected:raise RuntimeError('No contacts for force preload')
|
|
A=np.concatenate(columns,axis=1)
|
|
regular=np.zeros((len(selected),len(selected)*4))
|
|
for i in range(len(selected)):regular[i,i*4:i*4+4]=.025
|
|
target=np.r_[0,0,model.body_mass[object_body]*9.81,0,0,0]
|
|
preferred=np.array([6.]+[1.5]*(len(selected)-1))
|
|
solution=lsq_linear(np.vstack([A,regular]),np.r_[target,.025*preferred],bounds=(0,10),tol=1e-10)
|
|
tau=np.zeros(16)
|
|
for i,jac in enumerate(jacobians):tau+=jac.T@(cone[i]@solution.x[i*4:i*4+4])
|
|
for i,n in enumerate(cal.active):
|
|
dof=model.jnt_dofadr[mujoco.mj_name2id(model,mujoco.mjtObj.mjOBJ_JOINT,n)]
|
|
tau[i]+=data.qfrc_bias[dof]
|
|
controls+=np.clip(tau,-1,1)/kp
|
|
for i,n in enumerate(cal.active):controls[i]=np.clip(controls[i],*cal.bounds[n])
|
|
data.ctrl[:]=controls
|
|
preload_report=dict(wrench_residual=float(np.linalg.norm(A@solution.x-target)),desired_motor_torque_Nm=tau.tolist(),contact_normal_loads_N=(regular@solution.x/.025).tolist())
|
|
rows=[];forces=[];contacts=[];relative=[];torques=[];wrist_path=[]
|
|
initial_rel=np.array([0,0,.3])-handpos
|
|
feedback_errors=[]
|
|
for k in range(int(seconds/model.opt.timestep)):
|
|
t=k*model.opt.timestep
|
|
for n,(parent,x,y) in cal.curves.items():
|
|
active_angle=data.qpos[addr[parent]];passive,slope=cal.passive_value(n,active_angle)
|
|
eid=mujoco.mj_name2id(model,mujoco.mjtObj.mjOBJ_EQUALITY,'lut_'+n)
|
|
model.eq_data[eid,:5]=[passive-slope*active_angle,slope,0,0,0]
|
|
displacement=np.zeros(3)
|
|
if lift:
|
|
phase=np.clip((t-.5)/1.,0,1);smooth=phase*phase*(3-2*phase)
|
|
displacement=np.array([.02*smooth,0,.10*smooth])
|
|
data.mocap_pos[0]=handpos+displacement
|
|
if feedback and k%10==0:
|
|
phase=np.clip((t-.5),0,1)
|
|
target_velocity=np.array([.02,0,.10])*6*phase*(1-phase) if lift else np.zeros(3)
|
|
target_center=np.array([0,0,.3663])+displacement
|
|
newtau,residual=feedback_torque(model,data,cal,addr,objects,object_body,target_center,target_velocity)
|
|
if newtau is not None:tau=newtau;feedback_errors.append(residual)
|
|
if force_control:
|
|
if not force_preload:raise ValueError('force-control requires force-preload')
|
|
angles=np.array([data.qpos[addr[n]] for n in cal.active])
|
|
data.ctrl[:]=angles+(tau+1.0*(reference_controls-angles))/kp
|
|
mujoco.mj_step(model,data)
|
|
if not np.isfinite(data.qpos).all():raise RuntimeError('Nonfinite simulation')
|
|
if k%33==0:
|
|
rows.append(data.qpos.copy());wrist_path.append(np.r_[data.mocap_pos[0],data.mocap_quat[0]])
|
|
selected=[i for i,c in enumerate(data.contact) if any(g in objects for g in c.geom)]
|
|
normalforce=0
|
|
for ci in selected:
|
|
f=np.zeros(6);mujoco.mj_contactForce(model,data,ci,f);normalforce+=max(0,f[0])
|
|
forces.append(normalforce);contacts.append(len(selected));torques.append(data.actuator_force.copy())
|
|
relative.append(np.linalg.norm(data.qpos[obj:obj+3]-data.mocap_pos[0]-initial_rel))
|
|
rows=np.array(rows);relative=np.array(relative)
|
|
report=dict(seconds=seconds,lift=lift,kp=kp,preload_rad=preload,force_preload=preload_report,force_control=force_control,feedback=feedback,
|
|
feedback_wrench_residual_mean=float(np.mean(feedback_errors)) if feedback_errors else None,
|
|
max_relative_drift_m=float(relative.max()),final_relative_drift_m=float(relative[-1]),
|
|
object_vertical_displacement_m=float(data.qpos[obj+2]-.3),
|
|
contact_frames=int(np.sum(np.array(contacts)>0)),total_frames=len(rows),
|
|
normal_force_mean_N=float(np.mean(forces)),peak_actuator_torque_Nm=float(np.max(np.abs(torques))),
|
|
retained=bool(relative.max()<.02),
|
|
mass_kg=float(model.body_mass[object_body]),friction=.6,
|
|
boundaries=['Prescribed moving wrist, no arm dynamics','16 force-limited active joints; 5 nonlinear LUT equality constraints','No object weld or position reset','Object registration is estimated from hand pose, not tri-camera object tracking'])
|
|
np.savez_compressed(OUT/'physics_motion.npz',qpos=rows,wrist=np.array(wrist_path),fps=1/(33*model.opt.timestep),
|
|
object_qpos_address=obj,contact_count=contacts,normal_force=forces,relative_drift=relative,actuator_force=torques,
|
|
source_to_upright_rotation=rotation,controls=controls)
|
|
(OUT/'physics_validation.json').write_text(json.dumps(report,indent=2));print(json.dumps(report,indent=2),flush=True)
|
|
return report
|
|
|
|
if __name__=='__main__':
|
|
p=argparse.ArgumentParser();p.add_argument('--seconds',type=float,default=3);p.add_argument('--lift',action='store_true');p.add_argument('--kp',type=float,default=30);p.add_argument('--preload',type=float,default=0);p.add_argument('--force-preload',action='store_true');p.add_argument('--force-control',action='store_true');p.add_argument('--feedback',action='store_true');a=p.parse_args()
|
|
run(a.seconds,a.lift,a.kp,a.preload,a.force_preload,a.force_control,a.feedback)
|