[FEATURE] #2332 Implement Local Offset in Inverse Kinematics

[alexis779]

Description

Add support for local offset in Inverse Kinematics

Related Issue

Resolves #2332

Motivation and Context

It lets you target a local offset from the link origin, when running Inverse Kinematics.

How Has This Been / Can This Be Tested?

pytest tests/test_rigid_physics.py -k "test_inverse_kinematics_local_point and cpu"
pytest tests/test_rigid_physics.py -k "test_inverse_kinematics_multilink_local_points and cpu"

Full script

The target position on the following 2 images was target_pos = torch.tensor([0.1323, -0.1686, 0.1374]). It compares the robot configuration returned by IK by default, vs setting an arbitrary local offset in the link local frame.

	Before	After
configuration
tcp offset	[0, 0, 0]	[-1.4e-2, -9e-2, 0]

Download MJCF configuration of SO-ARM-100 robot as well

git clone https://github.com/google-deepmind/mujoco_menagerie
git clone https://github.com/Genesis-Embodied-AI/Genesis
cd Genesis
pip install -e .
python examples/so_arm_100.py

Here's the content of examples/so_arm_100.py script

import genesis as gs
from genesis.utils.geom import transform_by_quat, xyz_to_quat, transform_quat_by_quat

gs.init()

from genesis.engine.entities.rigid_entity.rigid_entity import RigidEntity
from genesis.engine.entities.rigid_entity.rigid_link import RigidLink
from genesis.engine.entities.rigid_entity.rigid_joint import RigidJoint

import torch
import logging
from functools import cached_property

class SoArm100:
    def __init__(self, config_path):
        # self.tcp_offset = [0, 0, 0]
        self.tcp_offset = [-1.4e-2, -9e-2, 0]
        self.tcp_offset = torch.tensor(self.tcp_offset)

        res = (860, 540)
        camera_pos = (-0.125, -1, 0.25)

        lookat = (0, 0, 0)

        lights = [
            {"type": "directional", "dir": (1, 1, -1), "color": (1.0, 1.0, 1.0), "intensity": 5.0},
        ]

        self.max_FPS = 24
        dt = 1/self.max_FPS

        substeps = 40

        self.scene = gs.Scene(
            sim_options=gs.options.SimOptions(
                dt=dt,
                substeps=substeps,
            ),
            viewer_options=gs.options.ViewerOptions(
                res=res,
                camera_lookat=lookat,
                camera_pos=camera_pos,
                max_FPS=self.max_FPS,
            ),
            vis_options=gs.options.VisOptions(
                lights=lights,
            ),
            profiling_options=gs.options.ProfilingOptions(
                show_FPS=False,
            ),
            show_viewer=True,
            rigid_options=gs.options.RigidOptions(
                noslip_iterations=substeps,
            ),
        )

        vis_mode = 'visual'

        plane = gs.morphs.Plane()
        self.scene.add_entity(plane, vis_mode=vis_mode)

        arm_morph = self.parse_morph(config_path)
        self.entity: RigidEntity = self.scene.add_entity(arm_morph, vis_mode=vis_mode)

        link_name = 'Fixed_Jaw'
        self.fixed_jaw: RigidLink = self.entity.get_link(link_name)

        joint_name = 'Jaw'
        self.joint: RigidJoint = self.entity.get_joint(joint_name)

        self.golf_ball_radius = 4.27e-2 / 2

        '''
        sphere_morph = gs.morphs.Mesh(
            file="meshes/sphere.obj",
            scale=self.golf_ball_radius
        )
        '''
        sphere_morph = gs.morphs.Sphere(
            radius=self.golf_ball_radius,
        )

        friction = 1.0
        self.sphere = self.scene.add_entity(
            sphere_morph,
            material=gs.materials.Rigid(friction=friction),
            vis_mode=vis_mode,
        )

        self.camera = self.scene.add_camera(
            res=res,
            pos=camera_pos,
            lookat=lookat,
        )

        self.scene.build()

        max_force = 14
        max_force = torch.full((self.entity.n_dofs,), max_force)
        min_force = -max_force
        self.entity.set_dofs_force_range(min_force, max_force)

        force_range = self.entity.get_dofs_force_range()
        self.logger.info(f"force_range={force_range}")

        Kp = self.entity.get_dofs_kp()
        self.logger.info(f"Kp={Kp}")
        Kv = self.entity.get_dofs_kv()
        self.logger.info(f"Kv={Kv}")

        init_qpos = self.entity.init_qpos
        self.logger.info(f"init_qpos={init_qpos}")

        armature = self.entity.get_dofs_armature()
        self.logger.info(f"armature={armature}")

        qpos_limit = self.entity.get_dofs_limit()
        self.logger.info(f"qpos_limit={qpos_limit}")

    @cached_property
    def logger(self) -> logging.Logger:
        # Initialize logger
        logger = logging.getLogger(__name__)
        logger.setLevel(logging.INFO)
        
        # Create console handler if not already present
        if not logger.handlers:
            handler = logging.StreamHandler()
            handler.setLevel(logging.INFO)
            formatter = logging.Formatter('%(asctime)s - %(name)s - %(levelname)s - %(message)s')
            handler.setFormatter(formatter)
            logger.addHandler(handler)

        return logger

    def parse_morph(self, config_path):
        if config_path.endswith(".xml"):
            return gs.morphs.MJCF(
                file=config_path,
            )
        elif config_path.endswith(".urdf"):
            return gs.morphs.URDF(
                file=config_path,
                fixed=True,
            )
        else:
            raise ValueError(f"Unsupported config file format: {config_path}")

    def run_simulation(self):
        self.camera.start_recording()
        self.target_random_pos()
        self.camera.stop_recording("./so_arm_100.mp4", fps=self.max_FPS)
    
    def hold(self):
        while True:
            self.scene.step()

    def target_random_pos(self):
        # 1. generate random roll, pitch, yaw
        # random yaw: [-pi, 0] (y < 0)
        azimuth = torch.rand(1) * torch.pi - torch.pi
        # random pitch: [-pi/2, 0] (negative pitch is up)
        elevation = torch.rand(1) * (-torch.pi / 2)
        # random roll: 0 (doesn't affect position of [r, 0, 0])
        roll = torch.tensor(0.0)

        # 2. generate random radius in the min, max range
        min_radius = 0.05
        max_radius = 0.3
        radius = min_radius + torch.rand(1) * (max_radius - min_radius)

        # 3. convert to x, y, z
        # We construct the quaternion by applying rotations in order: Roll -> Pitch -> Yaw
        # Note: Roll around X doesn't change [r, 0, 0], so only Pitch and Yaw matter for position.
        # We use separate quaternions to ensure Pitch corresponds to elevation (rotation around local Y after Yaw? No, intrinsic vs extrinsic).
        
        # q_pitch rotates around Y axis (elevation). 
        q_pitch = xyz_to_quat(torch.tensor([0.0, elevation, 0.0]))
        # q_yaw rotates around Z axis (azimuth).
        q_yaw = xyz_to_quat(torch.tensor([0.0, 0.0, azimuth]))

        # Combine rotations: apply Pitch then Yaw. 
        # q_combined = q_yaw * q_pitch
        quat = transform_quat_by_quat(q_pitch, q_yaw)

        # We start with a vector along X axis with length=radius, then rotate it
        target_pos = transform_by_quat(torch.tensor([radius, 0.0, 0.0]), quat)
        # target_pos = torch.tensor([0.1323, -0.1686,  0.1374])

        self.logger.info(f"target task space position={target_pos}")

        qpos, err = self.entity.inverse_kinematics(link=self.fixed_jaw, pos=target_pos, local_point=self.tcp_offset, return_error=True)
        self.logger.info(f"IK error={err}")

        self.logger.info(f"joint space solution={qpos}")
        self.entity.control_dofs_position(qpos)

        duration = 5 # seconds

        for step in range(self.max_FPS * duration):
            self.draw_debug_objects(target_pos)

        link_quat = self.fixed_jaw.get_quat()
        tcp_offset_world = transform_by_quat(self.tcp_offset, link_quat)

        link_pos = self.fixed_jaw.get_pos()

        tcp_pos = link_pos + tcp_offset_world

        pos_error = tcp_pos - target_pos
        self.logger.info(f"task space error={pos_error}")
        pos_error_norm = torch.norm(pos_error)
        self.logger.info(f"task space error norm={pos_error_norm}")

        tol = 1e-2
        if pos_error_norm > tol:
            self.logger.warning("Task space error is not within tolerance")


    def draw_debug_objects(self, target_pos):
        self.scene.clear_debug_objects()
        self.draw_link_arrow()
        self.draw_target_arrow(target_pos)
        self.scene.step()
        self.camera.render()

    def draw_link_arrow(self):
        link_pos = self.fixed_jaw.get_pos()

        link_quat = self.fixed_jaw.get_quat()
        tcp_offset_world = transform_by_quat(self.tcp_offset, link_quat)

        self.scene.draw_debug_arrow(link_pos, tcp_offset_world, color=(0, 1, 0), radius=0.005)

    def draw_target_arrow(self, target_pos):
        self.scene.draw_debug_arrow(torch.tensor([0, 0, 0]), target_pos)


mjcf_path = "../mujoco_menagerie/trs_so_arm100/so_arm100.xml"
so_arm_100 = SoArm100(mjcf_path)
so_arm_100.run_simulation()
#so_arm_100.hold()

Screenshots (if appropriate):

Checklist:

• I read the CONTRIBUTING document.
• I followed the Submitting Code Changes section of CONTRIBUTING document.
• I tagged the title correctly (including BUG FIX/FEATURE/MISC/BREAKING)
• I updated the documentation accordingly or no change is needed.
• I tested my changes and added instructions on how to test it for reviewers.

• I have added tests to cover my changes.
• All new and existing tests passed.