211 lines
6.3 KiB
Python
211 lines
6.3 KiB
Python
import os
|
|
from typing import Tuple
|
|
from contextlib import contextmanager
|
|
from pathlib import Path
|
|
import time
|
|
|
|
import yaml
|
|
import numpy as np
|
|
import torch as th
|
|
# from yourdfpy import URDF
|
|
|
|
import pinocchio as pin
|
|
import pink
|
|
from pink.tasks import FrameTask
|
|
|
|
|
|
@contextmanager
|
|
def with_dir(d):
|
|
d0 = os.getcwd()
|
|
try:
|
|
os.chdir(d)
|
|
yield
|
|
finally:
|
|
os.chdir(d0)
|
|
|
|
|
|
def xyzw2wxyz(q_xyzw: th.Tensor, dim: int = -1):
|
|
if isinstance(q_xyzw, np.ndarray):
|
|
return np.roll(q_xyzw, 1, axis=dim)
|
|
return th.roll(q_xyzw, 1, dims=dim)
|
|
|
|
|
|
def wxyz2xyzw(q_wxyz: th.Tensor, dim: int = -1):
|
|
if isinstance(q_wxyz, np.ndarray):
|
|
return np.roll(q_wxyz, -1, axis=dim)
|
|
return th.roll(q_wxyz, -1, dims=dim)
|
|
|
|
|
|
def dls_ik(
|
|
dpose: np.ndarray,
|
|
jac: np.ndarray,
|
|
sqlmda: float
|
|
):
|
|
"""
|
|
Arg:
|
|
dpose: task-space error (A[..., err])
|
|
jac: jacobian (A[..., err, dof])
|
|
sqlmda: DLS damping factor.
|
|
Return:
|
|
joint residual (A[..., dof])
|
|
"""
|
|
if isinstance(dpose, tuple):
|
|
dpose = np.concatenate([dpose[0], dpose[1]], axis=-1)
|
|
J = jac
|
|
A = J @ J.T
|
|
# NOTE(ycho): add to view of diagonal
|
|
a = np.einsum('...ii->...i', A)
|
|
a += sqlmda
|
|
dq = (J.T @ np.linalg.solve(A, dpose[..., None]))[..., 0]
|
|
return dq
|
|
|
|
|
|
class IKCtrl:
|
|
def __init__(self,
|
|
urdf_path: str,
|
|
act_joints: Tuple[str, ...],
|
|
frame: str = 'left_hand_palm_link',
|
|
sqlmda: float = 0.05**2):
|
|
path = Path(urdf_path)
|
|
with with_dir(path.parent):
|
|
robot = pin.RobotWrapper.BuildFromURDF(filename=path.name,
|
|
package_dirs=["."],
|
|
root_joint=None)
|
|
self.robot = robot
|
|
# NOTE(ycho): we skip joint#0(="universe")
|
|
joint_names = list(self.robot.model.names)
|
|
assert (joint_names[0] == 'universe')
|
|
self.joint_names = joint_names[1:]
|
|
|
|
# NOTE(ycho): build index map between pin.q and other set(s) of ordered
|
|
# joints.
|
|
act_from_pin = []
|
|
for j in act_joints:
|
|
act_from_pin.append(robot.index(j) - 1)
|
|
self.frame = frame
|
|
self.act_from_pin = np.asarray(act_from_pin, dtype=np.int32)
|
|
self.task = FrameTask(frame, position_cost=1.0, orientation_cost=0.0)
|
|
self.sqlmda = sqlmda
|
|
self.cfg = pink.Configuration(robot.model, robot.data,
|
|
np.zeros_like(robot.q0))
|
|
|
|
def fk(self, q: np.ndarray):
|
|
robot = self.robot
|
|
return pink.Configuration(
|
|
robot.model, robot.data, q).get_transform_frame_to_world(
|
|
self.frame)
|
|
|
|
def __call__(self,
|
|
q0: np.ndarray,
|
|
target_pose: np.ndarray,
|
|
rel: bool = False
|
|
):
|
|
"""
|
|
Arg:
|
|
q0: Current robot joints; A[..., 43?]
|
|
target_pose:
|
|
Policy output. A[..., 7] formatted as (xyz, q_{wxyz})
|
|
Given as world frame absolute pose, for some reason.
|
|
Return:
|
|
joint residual: A[..., 7]
|
|
"""
|
|
robot = self.robot
|
|
|
|
# source pose
|
|
self.cfg.update(q0)
|
|
T0 = self.cfg.get_transform_frame_to_world(self.frame)
|
|
|
|
# target pose
|
|
dst_xyz = target_pose[..., 0:3]
|
|
dst_quat = pin.Quaternion(wxyz2xyzw(target_pose[..., 3:7]))
|
|
T1 = pin.SE3(dst_quat, dst_xyz)
|
|
if rel:
|
|
TL = pin.SE3.Identity()
|
|
TL.translation = dst_xyz
|
|
TR = pin.SE3.Identity()
|
|
TR.rotation = dst_quat.toRotationMatrix()
|
|
T1 = TL * T0 * TR
|
|
|
|
# jacobian
|
|
self.task.set_target(T0)
|
|
jac = self.task.compute_jacobian(self.cfg)
|
|
jac = jac[:, self.act_from_pin]
|
|
|
|
# error&ik
|
|
dT = T1.actInv(T0)
|
|
dpose = pin.log(dT).vector
|
|
dq = dls_ik(dpose, jac, self.sqlmda)
|
|
return dq
|
|
|
|
|
|
def main():
|
|
from yourdfpy import URDF
|
|
urdf_path = '../../resources/robots/g1_description/g1_29dof_with_hand_rev_1_0.urdf'
|
|
|
|
# == init yourdfpy robot ==
|
|
path = Path(urdf_path)
|
|
with with_dir(path.parent):
|
|
viz = URDF.load(path.name,
|
|
build_collision_scene_graph=True,
|
|
load_meshes=False,
|
|
load_collision_meshes=True)
|
|
|
|
# == populate with defaults ==
|
|
with open('./configs/ik.yaml', 'r') as fp:
|
|
data = yaml.safe_load(fp)
|
|
q_mot = np.asarray(data['default_angles'])
|
|
q_viz = np.zeros((len(viz.actuated_joint_names)))
|
|
|
|
viz_from_mot = np.zeros(len(data['motor_joint']),
|
|
dtype=np.int32)
|
|
for i_mot, j in enumerate(data['motor_joint']):
|
|
i_viz = viz.actuated_joint_names.index(j)
|
|
viz_from_mot[i_mot] = i_viz
|
|
q_viz[viz_from_mot] = q_mot
|
|
|
|
ctrl = IKCtrl(urdf_path, data['act_joint'])
|
|
|
|
q_pin = np.zeros_like(ctrl.cfg.q)
|
|
pin_from_mot = np.zeros(len(data['motor_joint']),
|
|
dtype=np.int32)
|
|
for i_mot, j in enumerate(data['motor_joint']):
|
|
i_pin = ctrl.joint_names.index(j)
|
|
pin_from_mot[i_mot] = i_pin
|
|
|
|
mot_from_act = np.zeros(len(data['act_joint']),
|
|
dtype=np.int32)
|
|
for i_act, j in enumerate(data['act_joint']):
|
|
i_mot = data['motor_joint'].index(j)
|
|
mot_from_act[i_act] = i_mot
|
|
|
|
viz.update_cfg(q_viz)
|
|
# viz.show(collision_geometry=True)
|
|
|
|
if True:
|
|
current_pose = viz.get_transform(ctrl.frame)
|
|
print('curpose (viz)', current_pose)
|
|
print('curpose (pin)', ctrl.fk(q_pin).homogeneous)
|
|
|
|
def callback(scene):
|
|
if True:
|
|
current_pose = viz.get_transform(ctrl.frame)
|
|
T = pin.SE3()
|
|
T.translation = (current_pose[..., :3, 3]
|
|
+ [0.01, 0.0, 0.0])
|
|
T.rotation = current_pose[..., :3, :3]
|
|
target_pose = pin.SE3ToXYZQUAT(T)
|
|
target_pose[..., 3:7] = xyzw2wxyz(target_pose[..., 3:7])
|
|
q_pin[pin_from_mot] = q_mot
|
|
dq = ctrl(q_pin, target_pose, rel=False)
|
|
q_mot[mot_from_act] += dq
|
|
|
|
q_viz[viz_from_mot] = q_mot
|
|
viz.update_cfg(q_viz)
|
|
|
|
viz.show(collision_geometry=True,
|
|
callback=callback)
|
|
|
|
|
|
if __name__ == '__main__':
|
|
main()
|