Inverse Kinematics Problem

Consider a robot with $n$ degrees of freedom and a forward kinematics functions $FK(\textbf{x})$, where $\textbf{x} \in \R^n$ represents a robot joint state.

In many robotic settings, we are given some $\textbf{T}_i = (q_i, \textbf{t}i)\in SE(3){\mathbb{H}_1}$, a pose state for our $i^{th}$ link, which we need our robot to match. For example, it is common for us to require our robot to match some end-effector pose in order to perform a task (like grasping or moving an object). This requires us to find the right joint state $\textbf{x}$ that will place our $i^{th}$ link at the desired position. This can be formulated as an optimization problem.

An initial approach may involve minimizing the distance between the positions of our links, as specified by the forward kinematics of some joint state, versus our target position $\textbf{t}_i$ (just the position information in our desired pose state $\textbf{T}_i$). This can be formulated as:

$$ \argmin_{\textbf{x} \in \R^n} \| FK(\textbf{x})[i]_{position} - \textbf{t}_i \| $$

However, a pose $\textbf{T}i \in SE(3){\mathbb{H}_1}$ also specifies the rotational information $q_i$ of a link, which is neglected in the optimization problem above. From Lie Theory, we can incorporate this information by minimizing the distance between $FK(\textbf{x})[i]$ and $\textbf{T}i$ as elements of $SE(3){\mathbb{H}_1}$ (and not just as positional vectors). This gives the following problem:

$$ \argmin_{\textbf{x} \in \R^n} \| ln(disp(FK(\mathbf{x})[i], \mathbf{T}_i))^\vee \|_2^2 $$

Below, we provide the implementation of these optimization problems using the Apollo Toolbox.

use std::sync::Arc;
use bevy::app::App;
use apollo_rust_bevy::ApolloBevyTrait;
use apollo_rust_differentiation::{FunctionEngine, FunctionNalgebraTrait};
use apollo_rust_differentiation::derivative_methods::DerivativeMethodFD;
use apollo_rust_interpolation::InterpolatorTrait;
use apollo_rust_interpolation::splines::{InterpolatingSpline, InterpolatingSplineType};
use apollo_rust_lie::{EuclideanSpaceElement, LieAlgebraElement, LieGroupElement};
use apollo_rust_linalg::{ApolloDVectorTrait, V};
use apollo_rust_modules::ResourcesRootDirectory;
use apollo_rust_robotics::ToChainNalgebra;
use apollo_rust_robotics_core::ChainNalgebra;
use apollo_rust_spatial::isometry3::{ApolloIsometry3Trait, I3};
use apollo_rust_spatial::lie::h1::ApolloUnitQuaternionH1LieTrait;
use apollo_rust_spatial::lie::se3_implicit_quaternion::ISE3q;

// Matching Position Only
pub struct IKObjectiveFunction1 {
    pub c: Arc<ChainNalgebra>,
    pub goal_link_idx: usize,
    pub ik_goal: ISE3q
}
impl FunctionNalgebraTrait for IKObjectiveFunction1 {
    fn call_raw(&self, x: &V) -> V {
		    // Predicted link pose from forward kinematics
        let fk_res = self.c.fk(x);
        let link_frame = &fk_res[self.goal_link_idx];
        
        // Objective function
        let diff = &link_frame.0.translation.vector - &self.ik_goal.0.translation.vector;

        V::new(&[diff.norm().powi(2)])
    }

    fn input_dim(&self) -> usize {
        self.c.num_dofs()
    }

    fn output_dim(&self) -> usize {
        1
    }
}

// Matching as Implicit SE(3) elements
pub struct IKObjectiveFunction2 {
    pub c: Arc<ChainNalgebra>,
    pub goal_link_idx: usize,
    pub ik_goal: ISE3q
}
impl FunctionNalgebraTrait for IKObjectiveFunction2 {
    fn call_raw(&self, x: &V) -> V {
		    // Predicted link pose from forward kinematics
        let fk_res = self.c.fk(x);
        let link_frame = &fk_res[self.goal_link_idx];
        
        // Equivalent to objective function described above
        let position_diff = &link_frame.0.translation.vector - &self.ik_goal.0.translation.vector;
        let orientation_diff = link_frame.0.rotation.to_lie_group_h1().displacement(&self.ik_goal.0.rotation.to_lie_group_h1()).ln().vee();
        let res = position_diff.norm().powi(2) + orientation_diff.norm().powi(2);
        
        V::new(&[res])
    }

    fn input_dim(&self) -> usize {
        self.c.num_dofs()
    }

    fn output_dim(&self) -> usize {
        1
    }
}

fn main() {

    let r = ResourcesRootDirectory::new_from_default_apollo_robots_dir();
    let c = r.get_subdirectory("ur5").to_chain_nalgebra();
    let ac = c.clone().to_arc_chain();

		// Randomly specified goal for link 9 (the end-effector) of UR5 robot
    let ik_goal = ISE3q::new(I3::from_slices_euler_angles(&[0.2, 0.0, 0.5], &[0.,0.,0.]));

    let objective_function = IKObjectiveFunction2 {
        c: ac.clone(),
        goal_link_idx: 9,
        ik_goal: ik_goal.clone(),
    };

		// For derivative computation of objective function
    let function_engine = FunctionEngine::new(objective_function, DerivativeMethodFD::new(0.000001));

    let mut output_states = vec![];
    let mut curr_state = V::new(&[0.2; 6]);
    output_states.push(curr_state.clone());
    let lambda = 0.1;
    let max_iters = 1000;

		// Vanilla gradient descent to find optimal state
    for _ in 0..max_iters {
        let (_f, g) = function_engine.derivative(&curr_state);
        curr_state += -lambda*&g.transpose();
        output_states.push(curr_state.clone());
    }

		// Based on stored output states (in discrete steps), interpolate the continuous movement
    let linear_spline = InterpolatingSpline::new(output_states, InterpolatingSplineType::Linear).to_timed_interpolator(5.0);

		// Visualization
    let mut app = App::new()
        .apollo_bevy_robotics_base(true) // spawns robot scene 
        .apollo_bevy_spawn_chain_default(&c, 0, ISE3q::identity()) // places robot at origin
        .apollo_bevy_chain_display(&c) // displays robot chain
        .apollo_bevy_draw_frame(&ik_goal, 0.1) // pinpoints IK Goal
        .apollo_bevy_chain_motion_playback(0, linear_spline); // continuous playback of joint states during optimization

    app.run();
}