150 lines
7.8 KiB
Python
150 lines
7.8 KiB
Python
"""Offline whole-sequence temporal regularization with explicit distortion budgets."""
|
|
import argparse
|
|
import json
|
|
import shutil
|
|
from pathlib import Path
|
|
import xml.etree.ElementTree as ET
|
|
|
|
import mujoco
|
|
import numpy as np
|
|
from scipy import sparse
|
|
from scipy.sparse.linalg import splu
|
|
from scipy.spatial.transform import Rotation
|
|
|
|
ROOT = Path(__file__).resolve().parents[1]
|
|
|
|
|
|
def main():
|
|
p = argparse.ArgumentParser()
|
|
p.add_argument('--source', type=Path, default=ROOT/'output/l20_15886123')
|
|
p.add_argument('--output', type=Path, default=ROOT/'output/l20_15886123_stable')
|
|
a = p.parse_args()
|
|
a.output.mkdir(parents=True, exist_ok=True)
|
|
with np.load(a.source/'motion.npz') as f:
|
|
m = {k:f[k] for k in f.files}
|
|
n = len(m['qpos'])
|
|
D2 = sparse.diags([np.ones(n-2), -2*np.ones(n-2), np.ones(n-2)], [0,1,2], shape=(n-2,n),format='csc')
|
|
D3 = sparse.diags([-np.ones(n-3),3*np.ones(n-3),-3*np.ones(n-3),np.ones(n-3)], [0,1,2,3],shape=(n-3,n),format='csc')
|
|
eye = sparse.eye(n,format='csc')
|
|
names, active = list(m['joint_names']), list(m['active_joint_names'])
|
|
ai = [names.index(k) for k in active]
|
|
original = ET.parse(ROOT/'third_party/l20_assets/L20/RIGHT/linkerhand_g20_right.urdf').getroot()
|
|
limits = {j.get('name'):np.array([float(j.find('limit').get(k)) for k in ['lower','upper']]) for j in original.findall('joint') if j.get('type')!='fixed'}
|
|
low, high = np.array([limits[k] for k in active]).T.copy()
|
|
mimic = {}
|
|
for j in original.findall('joint'):
|
|
mi = j.find('mimic')
|
|
if mi is None:
|
|
continue
|
|
par, sc, off = mi.get('joint'),float(mi.get('multiplier','1')),float(mi.get('offset','0'))
|
|
mimic[j.get('name')] = par,sc,off
|
|
lo,hi = np.sort((limits[j.get('name')]-off)/sc)
|
|
k = active.index(par)
|
|
low[k],high[k] = max(low[k],lo),min(high[k],hi)
|
|
|
|
def smooth(x, weight, bounds=None):
|
|
# Sample-domain objective ||x-y||² + ||D²x||² + weight*||D³x||².
|
|
A = eye + D2.T@D2 + weight*(D3.T@D3)
|
|
result = splu(A).solve(x)
|
|
if bounds is not None:
|
|
lo,hi = bounds
|
|
# Explicit post-solve projection; evaluate smoothness after projection.
|
|
# This is not claimed to solve the box-constrained temporal QP exactly.
|
|
result = np.clip(result,lo,hi)
|
|
return result
|
|
|
|
def d2(x):
|
|
return float(np.sqrt(np.mean(np.diff(x[30:],n=2,axis=0)**2)))
|
|
|
|
rot = Rotation.from_quat(m['wrist_quat_wxyz'][:,[1,2,3,0]])
|
|
rv = (rot[0].inv()*rot).as_rotvec()
|
|
assert np.linalg.norm(rv,axis=1).max() < np.pi/2
|
|
choices, sweep = {}, {}
|
|
# Budgets are engineering choices relative to prior output, not true-motion tolerances.
|
|
components = dict(fingers=(m['qpos'][:,ai],(low,high),np.deg2rad(8)),
|
|
wrist_position=(m['wrist_pos'],None,.020),
|
|
wrist_rotation=(rv,None,np.deg2rad(6)))
|
|
for key,(x,bounds,budget) in components.items():
|
|
rows, candidates = [], []
|
|
for weight in [100.,300.,1000.,3000.,10000.,30000.]:
|
|
y = smooth(x,weight,bounds)
|
|
if key=='fingers':
|
|
dev = float(np.abs(y-x).max())
|
|
elif key=='wrist_rotation':
|
|
dev = float(np.linalg.norm((rot.inv()*(rot[0]*Rotation.from_rotvec(y))).as_rotvec(),axis=1).max())
|
|
else:
|
|
dev = float(np.linalg.norm(y-x,axis=1).max())
|
|
ratio = d2(y)/d2(x)
|
|
row = dict(weight_d2=1.,weight_d3=weight,max_deviation=dev,budget=float(budget),
|
|
second_difference_ratio=ratio,within_budget=bool(dev<=budget))
|
|
rows.append(row)
|
|
if dev<=budget:
|
|
candidates.append((ratio,y,row))
|
|
assert candidates, key
|
|
_,chosen,row = min(candidates,key=lambda v:v[0])
|
|
choices[key] = chosen
|
|
sweep[key] = dict(selected=row,candidates=rows)
|
|
print(key,row,flush=True)
|
|
before = {k:m[k].copy() for k in ['qpos','wrist_pos','wrist_quat_wxyz','actual']}
|
|
q = m['qpos'].copy()
|
|
q[:,ai] = choices['fingers']
|
|
for name,(parent,sc,off) in mimic.items():
|
|
q[:,names.index(name)] = sc*q[:,names.index(parent)]+off
|
|
m['qpos'] = q
|
|
m['active_qpos'] = q[:,ai]
|
|
m['wrist_pos'] = choices['wrist_position']
|
|
newrot = rot[0]*Rotation.from_rotvec(choices['wrist_rotation'])
|
|
m['wrist_quat_wxyz'] = newrot.as_quat()[:,[3,0,1,2]]
|
|
for k,v in before.items():
|
|
m['before_stabilization_'+k] = v
|
|
m['stabilization_method'] = 'whole-sequence data fidelity + second/third difference regularization; component distortion budgets'
|
|
m['stabilization_source'] = str(a.source.resolve())
|
|
for filename in ['l20_moving.xml','l20_right.xml','l20_dex.urdf','dex_config.json']:
|
|
shutil.copy2(a.source/filename,a.output/filename)
|
|
model = mujoco.MjModel.from_xml_path(str(a.output/'l20_right.xml'))
|
|
data = mujoco.MjData(model)
|
|
addr = [model.jnt_qposadr[model.joint(k).id] for k in names]
|
|
sites = [model.site(f'landmark_{i:02d}').id for i in range(21)]
|
|
actual = []
|
|
for row in q:
|
|
data.qpos[addr] = row
|
|
mujoco.mj_forward(model,data)
|
|
actual.append(data.site_xpos[sites].copy())
|
|
m['actual'] = np.array(actual)
|
|
np.savez_compressed(a.output/'motion.npz',**m)
|
|
np.savetxt(a.output/'trajectory.csv',np.c_[m['time'],m['wrist_pos'],m['wrist_quat_wxyz'],q],delimiter=',',
|
|
header=','.join(['time_s','wrist_x','wrist_y','wrist_z','qw','qx','qy','qz']+names),comments='')
|
|
lim = np.array([limits[k] for k in names])
|
|
violation = float(max(np.maximum(lim[:,0]-q,0).max(),np.maximum(q-lim[:,1],0).max()))
|
|
assert violation < 1e-9
|
|
mimic_error = max(float(np.abs(q[:,names.index(k)]-sc*q[:,names.index(par)]-off).max()) for k,(par,sc,off) in mimic.items())
|
|
assert mimic_error < 1e-10
|
|
|
|
def metrics(q,pos,quat,actual):
|
|
r = Rotation.from_quat(quat[:,[1,2,3,0]])[30:]
|
|
w = (r[:-1].inv()*r[1:]).as_rotvec()
|
|
tips = np.linalg.norm(actual[:,[4,8,12,16,20]]-m['targets'][:,[4,8,12,16,20]],axis=-1)*1000
|
|
return dict(joint_second_difference_rms_deg=np.rad2deg(d2(q)),
|
|
joint_max_step_deg=float(np.rad2deg(np.abs(np.diff(q[30:],axis=0)).max())),
|
|
wrist_second_difference_rms_mm=d2(pos)*1000,
|
|
wrist_max_step_mm=float(np.linalg.norm(np.diff(pos[30:],axis=0),axis=1).max()*1000),
|
|
wrist_rotation_velocity_difference_rms_deg=float(np.rad2deg(np.sqrt(np.mean(np.diff(w,axis=0)**2)))),
|
|
wrist_rotation_max_step_deg=float(np.rad2deg(np.linalg.norm(w,axis=1).max())),
|
|
wrist_range_m=np.ptp(pos,axis=0).tolist(),
|
|
tip_error_mean_mm=float(tips.mean()),tip_error_p95_mm=float(np.percentile(tips,95)))
|
|
report = dict(frames=n,fps=float(m['fps']),evaluation_start_frame=30,
|
|
before=metrics(**dict(q=before['qpos'],pos=before['wrist_pos'],quat=before['wrist_quat_wxyz'],actual=before['actual'])),
|
|
after=metrics(q,m['wrist_pos'],m['wrist_quat_wxyz'],m['actual']),
|
|
selection=sweep,urdf_limit_violation_rad=violation,linear_mimic_error_rad=mimic_error,
|
|
max_joint_change_deg=float(np.rad2deg(np.abs(q-before['qpos']).max())),
|
|
max_wrist_change_mm=float(np.linalg.norm(m['wrist_pos']-before['wrist_pos'],axis=1).max()*1000),
|
|
max_wrist_rotation_change_deg=float(np.rad2deg(np.linalg.norm((rot.inv()*newrot).as_rotvec(),axis=1).max())),
|
|
uses_future_frames=True,boundary='Lower temporal derivatives do not prove true-motion accuracy; no contact/dynamics validation')
|
|
assert all(np.isfinite(m[k]).all() for k in ['qpos','wrist_pos','wrist_quat_wxyz','actual'])
|
|
(a.output/'stabilization_validation.json').write_text(json.dumps(report,indent=2))
|
|
print(json.dumps({k:v for k,v in report.items() if k!='selection'},indent=2))
|
|
|
|
|
|
if __name__=='__main__':
|
|
main()
|