unitree_rl_gym/deploy/deploy_real/ikctrl.py

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():
    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()