Skip to content

Commit

Permalink
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
WIP: Implement custom cached forward kinematics
Browse files Browse the repository at this point in the history
kylc committed Dec 17, 2023
1 parent 58ac804 commit 8fa7ec5
Showing 5 changed files with 367 additions and 156 deletions.
58 changes: 40 additions & 18 deletions crates/optik/benches/bench.rs
Original file line number Diff line number Diff line change
@@ -1,6 +1,8 @@
use criterion::{black_box, criterion_group, criterion_main, Criterion};
use nalgebra::Isometry3;

use optik::kinematics::KinematicsCache;
use optik::objective::ObjectiveArgs;
use optik::*;

const BENCH_MODEL_STR: &str = include_str!("../tests/data/ur3e.urdf");
@@ -9,61 +11,81 @@ fn load_benchmark_model() -> Robot {
Robot::from_urdf_str(BENCH_MODEL_STR, "ur_base_link", "ur_ee_link")
}

fn bench_jacobian(c: &mut Criterion) {
fn bench_jacobian(_c: &mut Criterion) {
let robot = load_benchmark_model();
let q = robot.random_configuration(&mut rand::thread_rng());
let _q = robot.random_configuration(&mut rand::thread_rng());

c.bench_function("jacobian", |b| {
b.iter(|| robot.jacobian_local(black_box(&q)))
});
// c.bench_function("jacobian", |b| {
// b.iter(|| robot.jacobian_local(black_box(&q)))
// });
}

fn bench_gradient(c: &mut Criterion) {
use optik::objective::objective_grad;

let robot = load_benchmark_model();

let q = robot.random_configuration(&mut rand::thread_rng());
let mut g = vec![0.0; q.len()];
let tfm_target = Isometry3::identity();
let args = ObjectiveArgs {
robot,
config: SolverConfig::default(),
tfm_target,
robot: &robot,
config: &SolverConfig::default(),
tfm_target: &tfm_target,
};

c.bench_function("gradient_analytical", |b| {
let args = ObjectiveArgs {
config: SolverConfig {
config: &SolverConfig {
gradient_mode: GradientMode::Analytical,
..args.config
..*args.config
},
..args.clone()
};
b.iter(|| objective_grad(black_box(&q), &mut g, &args))
b.iter(|| {
objective_grad(
black_box(&q),
&mut g,
&args,
&mut KinematicsCache::default(),
)
})
});

c.bench_function("gradient_numerical", |b| {
let args = ObjectiveArgs {
config: SolverConfig {
config: &SolverConfig {
gradient_mode: GradientMode::Numerical,
..args.config
..*args.config
},
..args.clone()
};
b.iter(|| objective_grad(black_box(&q), &mut g, &args))
b.iter(|| {
objective_grad(
black_box(&q),
&mut g,
&args,
&mut KinematicsCache::default(),
)
})
});
}

fn bench_objective(c: &mut Criterion) {
use optik::objective::objective;

let robot = load_benchmark_model();

let q = robot.random_configuration(&mut rand::thread_rng());
let tfm_target = Isometry3::identity();
let args = ObjectiveArgs {
robot,
config: SolverConfig::default(),
tfm_target,
robot: &robot,
config: &SolverConfig::default(),
tfm_target: &tfm_target,
};
c.bench_function("objective", |b| b.iter(|| objective(black_box(&q), &args)));
c.bench_function("objective", |b| {
b.iter(|| objective(black_box(&q), &args, &mut KinematicsCache::default()))
});
}

fn bench_ik(c: &mut Criterion) {
132 changes: 132 additions & 0 deletions crates/optik/src/kinematics.rs
Original file line number Diff line number Diff line change
@@ -0,0 +1,132 @@
use nalgebra::{Isometry3, Matrix6xX, RealField, Translation3, Unit, UnitQuaternion, Vector3};

#[derive(Debug, Clone)]
pub struct Joint<T: RealField> {
offset: Isometry3<T>,
typ: JointType<T>,
}

#[derive(Debug, Clone)]
enum JointType<T: RealField> {
Revolute { axis: Unit<Vector3<T>> },
Prismatic { axis: Unit<Vector3<T>> },
}

impl<T: RealField + Copy> Joint<T> {
pub fn revolute(offset: Isometry3<T>, axis: Unit<Vector3<T>>) -> Joint<T> {
Joint {
offset,
typ: JointType::Revolute { axis },
}
}

pub fn prismatic(offset: Isometry3<T>, axis: Unit<Vector3<T>>) -> Joint<T> {
Joint {
offset,
typ: JointType::Prismatic { axis },
}
}

fn local_transform(&self, q: &[T]) -> Isometry3<T> {
match self.typ {
JointType::Revolute { axis } => {
// let s = q[0].sin();
// let c = q[0].cos();

// let v = Vector3::x();
// let k = &axis;
// let v_rot =
// v * c + k.cross(&v) * s + k.scale(k.dot(&v) * (nalgebra::one::<T>() - c));
self.offset * UnitQuaternion::from_axis_angle(&axis, q[0])
}
JointType::Prismatic { axis } => self.offset * Translation3::from(axis.scale(q[0])),
}
}

pub fn nq(&self) -> usize {
match self.typ {
JointType::Revolute { .. } => 1,
JointType::Prismatic { .. } => 1,
}
}
}

pub fn joint_kinematics<'a>(
chain: &'a [Joint<f64>],
q: &'a [f64],
) -> impl Iterator<Item = Isometry3<f64>> + 'a {
struct State {
qidx: usize,
tfm: Isometry3<f64>,
}

chain.iter().scan(
State {
tfm: Isometry3::identity(),
qidx: 0,
},
|state, joint| {
state.tfm *= joint.local_transform(&q[state.qidx..(state.qidx + joint.nq())]);
state.qidx += joint.nq();

Some(state.tfm)
},
)
}

pub fn joint_jacobian(
arm: &[Joint<f64>],
ee_offset: &Isometry3<f64>,
tfms: &[Isometry3<f64>],
) -> Matrix6xX<f64> {
let dof = arm.len();
let tfm_w_ee = tfms.last().unwrap() * ee_offset;

let mut m = Matrix6xX::zeros(dof);
for (col, joint) in arm.iter().enumerate() {
let tfm_w_i = tfms[col];

match joint.typ {
JointType::Revolute { axis, .. } => {
let a_i = tfm_w_i.rotation * axis;
let dp_i = a_i.cross(&(tfm_w_ee.translation.vector - tfm_w_i.translation.vector));

// in local frame
let a_i = tfm_w_ee.inverse_transform_vector(&a_i);
let dp_i = tfm_w_ee.inverse_transform_vector(&dp_i);

m.fixed_slice_mut::<3, 1>(0, col).copy_from(&dp_i);
m.fixed_slice_mut::<3, 1>(3, col).copy_from(&a_i);
}
JointType::Prismatic { .. } => todo!(),
};
}

m
}

/// A cache structure which stores the expensive forward kinematics computation
/// for its respective generalized joint configuration.
#[derive(Clone, Default)]
pub struct KinematicsCache {
q: Vec<f64>,
frames: Vec<Isometry3<f64>>,
}

impl KinematicsCache {
/// If a cache entry exists for the given joint configuration vector `q`,
/// return that. Otherwise, compute the forward kinematics, replace the the
/// cache with the results, and return them.
pub fn get_or_update(&mut self, joints: &[Joint<f64>], q: &[f64]) -> &[Isometry3<f64>] {
if self.q != q {
// Clear and extend the cached vectors to prevent any allocations.
self.q.clear();
self.q.extend_from_slice(q);

self.frames.clear();
self.frames.extend(joint_kinematics(joints, q));
}

&self.frames
}
}
200 changes: 71 additions & 129 deletions crates/optik/src/lib.rs
Original file line number Diff line number Diff line change
@@ -1,4 +1,5 @@
use std::{
cell::RefCell,
path::Path,
sync::{
atomic::{AtomicBool, Ordering},
@@ -7,8 +8,8 @@ use std::{
time::Instant,
};

use k::{joint::Range, Chain, SerialChain};
use nalgebra::{DMatrix, DVector, DVectorSlice, Isometry3};
use k::{joint::Range, Chain, JointType, SerialChain};
use nalgebra::{DVector, DVectorSlice, Isometry3, Matrix6xX};
use ordered_float::OrderedFloat;
use rand::SeedableRng;
use rand_chacha::ChaCha8Rng;
@@ -19,10 +20,13 @@ use rayon::{
use slsqp_sys::*;

mod config;
pub mod kinematics;
pub mod math;
pub mod objective;

pub use config::*;
use math::*;
use kinematics::*;
use objective::*;

pub fn set_parallelism(n: usize) {
ThreadPoolBuilder::new()
@@ -33,12 +37,37 @@ pub fn set_parallelism(n: usize) {

#[derive(Clone)]
pub struct Robot {
pub chain: SerialChain<f64>,
chain: SerialChain<f64>,
joints: Vec<Joint<f64>>,
ee_offset: Isometry3<f64>,
}

impl Robot {
pub fn new(chain: SerialChain<f64>) -> Self {
Self { chain }
let mut joints = vec![];
let mut fixed_offset = Isometry3::identity();
for node in chain.iter() {
let joint = node.joint();
match joint.joint_type {
JointType::Rotational { axis } => {
joints.push(Joint::revolute(joint.origin() * fixed_offset, axis));
fixed_offset = Isometry3::identity();
}
JointType::Linear { axis } => {
joints.push(Joint::prismatic(joint.origin() * fixed_offset, axis));
fixed_offset = Isometry3::identity();
}
JointType::Fixed {} => {
fixed_offset *= joint.origin();
}
}
}

Self {
chain,
joints,
ee_offset: fixed_offset,
}
}

pub fn from_urdf_file(path: impl AsRef<Path>, base_link: &str, ee_link: &str) -> Self {
@@ -65,31 +94,8 @@ impl Robot {
Robot::new(serial)
}

pub fn jacobian_local(&self, q: &[f64]) -> DMatrix<f64> {
let t_n = self.fk(q); // this updates the joint positions

let mut m = k::jacobian(&self.chain);

// K computes a Jacobian J(q) in Pinocchio's terms as
// LOCAL_WORLD_ALIGNED. Because we compute the right Jacobian of log
// SE(3) J_r(X), we prefer to work in body frame. Convert J(q) into the
// local body frame (Pinocchio calls this LOCAL frame).
let w_inv = t_n.rotation.inverse();
for mut col in m.column_iter_mut() {
let mut linear = col.fixed_rows_mut::<3>(0);
let linear_w = w_inv * &linear;
linear.copy_from(&linear_w);

let mut angular = col.fixed_rows_mut::<3>(3);
let angular_w = w_inv * &angular;
angular.copy_from(&angular_w);
}

m
}

pub fn num_positions(&self) -> usize {
self.chain.dof()
self.joints.len()
}

pub fn joint_limits(&self) -> (DVector<f64>, DVector<f64>) {
@@ -107,14 +113,20 @@ impl Robot {

pub fn random_configuration(&self, rng: &mut impl rand::Rng) -> Vec<f64> {
let (lb, ub) = self.joint_limits();
(0..self.chain.dof())
(0..self.num_positions())
.map(|i| rng.gen_range(lb[i]..=ub[i]))
.collect()
}

pub fn joint_jacobian(&self, frames: &[Isometry3<f64>]) -> Matrix6xX<f64> {
joint_jacobian(&self.joints, &self.ee_offset, frames)
}

pub fn fk(&self, q: &[f64]) -> Isometry3<f64> {
self.chain.set_joint_positions_unchecked(q);
self.chain.end_transform()
joint_kinematics(&self.joints, q)
.last()
.unwrap_or_else(Isometry3::identity)
* self.ee_offset
}

pub fn ik(
@@ -163,10 +175,11 @@ impl Robot {
let mut rng = ChaCha8Rng::seed_from_u64(RNG_SEED);
rng.set_stream(i);

let cache = RefCell::new(KinematicsCache::default());
let args = ObjectiveArgs {
robot: self.clone(),
config: config.clone(),
tfm_target: *tfm_target,
robot: self,
config,
tfm_target,
};

// The first attempt gets the initial seed provided by the caller.
@@ -186,17 +199,17 @@ impl Robot {

// Bookkeeping for stopping criteria.
let mut x_prev = x.clone();
let mut f_prev = objective(&x, &args);
let mut f_prev = objective(&x, &args, &mut cache.borrow_mut());

// Iterate the soler until any of:
// - The solver converges within the tolerance
// - Another thread signals that it has converged
// - The timeout expires
// Iterate the solver until any of:
// - The solver converges within the active convergence criteria
// - Another thread arrives at a solution (in speed mode)
// - The specified timeout expires
while !should_exit.load(Ordering::Relaxed) && !is_timed_out() {
match solver.iterate(
&mut x,
|x| objective(x, &args),
|x, g| objective_grad(x, g, &args),
|x| objective(x, &args, &mut cache.borrow_mut()),
|x, g| objective_grad(x, g, &args, &mut cache.borrow_mut()),
) {
IterationResult::Continue => {
x_prev.copy_from_slice(&x);
@@ -206,19 +219,31 @@ impl Robot {
}
IterationResult::Converged => {
// The SLSQP solver can report convergence
// regardless of whether our tol_f, tol_df, nor
// regardless of whether our tol_f, tol_df, or
// tol_dx conditions are met. If we verify that this
// solution does meet the criteria then it can be
// returned.
//
// This is due to an internal `accuracy` parameter,
// which reports convergence if the change in
// objective function falls below it.
let df = solver.cost() - f_prev;
let dx = DVector::from_row_slice(&x) - DVector::from_row_slice(&x_prev);
let dx = DVectorSlice::from_slice(&x, x.len())
- DVectorSlice::from_slice(&x_prev, x_prev.len());

// Report convergence if _any_ of the active
// convergence criteria are met.
if solver.cost().abs() < config.tol_f
|| (config.tol_df > 0.0 && df.abs() < config.tol_df)
|| (config.tol_dx > 0.0 && dx.norm().abs() < config.tol_dx)
{
// Short-circuit any other threads for a modest
// performance increase.
//
// In `Speed` mode, if the current thread has
// converged on a satisfactory solution then we
// don't care what the others are going to
// produce.
if config.solution_mode == SolutionMode::Speed {
should_exit.store(true, Ordering::Relaxed);
}
@@ -243,8 +268,8 @@ impl Robot {
// best of all solutions. In this case, the cost of a given
// solution is computed as its distance from the seed.
solution_stream.min_by_key(|(x, _)| {
let x = DVectorSlice::from_slice(x, self.num_positions());
let x0 = DVectorSlice::from_slice(&x0, self.num_positions());
let x = DVectorSlice::from_slice(x, x.len());
let x0 = DVectorSlice::from_slice(&x0, x0.len());

OrderedFloat(x.metric_distance(&x0))
})
@@ -256,86 +281,3 @@ impl Robot {
}
}
}

#[derive(Clone)]
pub struct ObjectiveArgs {
pub robot: Robot,
pub config: SolverConfig,
pub tfm_target: Isometry3<f64>,
}

pub fn objective(q: &[f64], args: &ObjectiveArgs) -> f64 {
// Compute the pose error w.r.t. the actual pose:
//
// X_AT = X_WA^1 * X_WT
let tfm_actual = args.robot.fk(q);
let tfm_target = args.tfm_target;
let tfm_error = tfm_target.inv_mul(&tfm_actual);

// Minimize the sqaure Euclidean distance of the log pose error. We choose
// to use the square distance due to its smoothness.
se3::log(&tfm_error).norm_squared()
}

/// Compute the gradient `g` w.r.t. the local parameterization.
pub fn objective_grad(q: &[f64], g: &mut [f64], args: &ObjectiveArgs) {
let robot = &args.robot;
let tfm_actual = &args.robot.fk(q);
let tfm_target = &args.tfm_target;

match args.config.gradient_mode {
GradientMode::Analytical => {
// Pose error is computed as in the objective function.
let tfm_error = tfm_target.inv_mul(tfm_actual);

// We compute the Jacobian of our task w.r.t. the joint angles.
//
// - Jqdot: the local (body-coordinate) frame Jacobian of X w.r.t. q
// - Jlog6: the right (body-coordinate) Jacobian of log6(X)
// - Jtask: the Jacobian of our pose tracking task
//
// Jtask(q) = Jlog6(X) * J(q)
let j_qdot = robot.jacobian_local(q);
let j_log6 = se3::right_jacobian(&tfm_error);
let j_task = j_log6 * j_qdot;

// We must compute the objective function gradient:
//
// ∇h = [ ∂h/∂x1 ... ∂h/∂xn ]
//
// Given the pose task Jacobian of the form:
//
// J = ⎡ ∂f1/∂x1 ... ∂f1/∂xn ⎤
// ⎢ . . ⎥
// ⎢ . . ⎥
// ⎣ ∂fm/∂x1 ... ∂fm/∂xn ⎦
//
// Apply the chain rule to compute the derivative:
//
// g = log(X)
// f = || g ||^2
// h' = (f' ∘ g) * g' = (2.0 * log6(X)) * J
let fdot_g = 2.0 * se3::log(&tfm_error).transpose();
let grad_h = fdot_g * j_task;

g.copy_from_slice(grad_h.as_slice());
}
GradientMode::Numerical => {
// Central finite difference:
//
// ∇h = (f(x + Δx) - f(x - Δx)) / (2 * Δx)
let n = q.len();
let mut x0 = q.to_vec();
let eps = f64::EPSILON.powf(1.0 / 3.0);
for i in 0..n {
let x0i = x0[i];
x0[i] = x0i - eps;
let fl = objective(&x0, args);
x0[i] = x0i + eps;
let fh = objective(&x0, args);
g[i] = (fh - fl) / (2.0 * eps);
x0[i] = x0i;
}
}
}
}
89 changes: 89 additions & 0 deletions crates/optik/src/objective.rs
Original file line number Diff line number Diff line change
@@ -0,0 +1,89 @@
use nalgebra::Isometry3;

use crate::{kinematics::KinematicsCache, math::se3, GradientMode, Robot, SolverConfig};

#[derive(Clone)]
pub struct ObjectiveArgs<'a> {
pub robot: &'a Robot,
pub config: &'a SolverConfig,
pub tfm_target: &'a Isometry3<f64>,
}

pub fn objective(x: &[f64], args: &ObjectiveArgs, cache: &mut KinematicsCache) -> f64 {
// Compute the pose error w.r.t. the actual pose:
//
// X_AT = X_WA^1 * X_WT
let frames = cache.get_or_update(&args.robot.joints, x);
let tfm_actual = frames.last().unwrap() * args.robot.ee_offset;
let tfm_target = args.tfm_target;
let tfm_error = tfm_target.inv_mul(&tfm_actual);

// Minimize the sqaure Euclidean distance of the log pose error. We choose
// to use the square distance due to its smoothness.
se3::log(&tfm_error).norm_squared()
}

/// Compute the gradient `g` w.r.t. the local parameterization.
pub fn objective_grad(q: &[f64], g: &mut [f64], args: &ObjectiveArgs, cache: &mut KinematicsCache) {
let robot = &args.robot;

match args.config.gradient_mode {
GradientMode::Analytical => {
let frames = cache.get_or_update(&args.robot.joints, q);

// Pose error is computed as in the objective function.
let tfm_actual = frames.last().unwrap() * robot.ee_offset;
let tfm_target = &args.tfm_target;
let tfm_error = tfm_target.inv_mul(&tfm_actual);

// We compute the Jacobian of our task w.r.t. the joint angles.
//
// - Jqdot: the local (body-coordinate) frame Jacobian of X w.r.t. q
// - Jlog6: the right (body-coordinate) Jacobian of log6(X)
// - Jtask: the Jacobian of our pose tracking task
//
// Jtask(q) = Jlog6(X) * J(q)
let j_qdot = robot.joint_jacobian(frames);
let j_log6 = se3::right_jacobian(&tfm_error);
let j_task = j_log6 * j_qdot;

// We must compute the objective function gradient:
//
// ∇h = [ ∂h/∂x1 ... ∂h/∂xn ]
//
// Given the pose task Jacobian of the form:
//
// J = ⎡ ∂f1/∂x1 ... ∂f1/∂xn ⎤
// ⎢ . . ⎥
// ⎢ . . ⎥
// ⎣ ∂fm/∂x1 ... ∂fm/∂xn ⎦
//
// Apply the chain rule to compute the derivative:
//
// g = log(X)
// f = || g ||^2
// h' = (f' ∘ g) * g' = (2.0 * log6(X)) * J
let fdot_g = 2.0 * se3::log(&tfm_error).transpose();
let grad_h = fdot_g * j_task;

g.copy_from_slice(grad_h.as_slice());
}
GradientMode::Numerical => {
// Central finite difference:
//
// ∇h = (f(x + Δx) - f(x - Δx)) / (2 * Δx)
let n = q.len();
let mut x0 = q.to_vec();
let eps = f64::EPSILON.powf(1.0 / 3.0);
for i in 0..n {
let x0i = x0[i];
x0[i] = x0i - eps;
let fl = objective(&x0, args, cache);
x0[i] = x0i + eps;
let fh = objective(&x0, args, cache);
g[i] = (fh - fl) / (2.0 * eps);
x0[i] = x0i;
}
}
}
}
44 changes: 35 additions & 9 deletions crates/optik/tests/test_gradient.rs
Original file line number Diff line number Diff line change
@@ -1,6 +1,10 @@
use approx::assert_abs_diff_eq;
use nalgebra::Vector6;
use optik::{objective_grad, GradientMode, ObjectiveArgs, Robot, SolverConfig};
use optik::{
kinematics::KinematicsCache,
objective::{objective_grad, ObjectiveArgs},
GradientMode, Robot, SolverConfig,
};
use rand::{rngs::StdRng, Rng, SeedableRng};

const TEST_MODEL_STR: &str = include_str!("data/ur3e.urdf");
@@ -14,19 +18,41 @@ fn test_gradient_analytical_vs_numerical() {
let x0: Vector6<f64> = rng.gen();
let tfm_target = rng.gen();

let mut args = ObjectiveArgs {
robot: robot.clone(),
config: SolverConfig::default(),
tfm_target,
// Analytical gradient
let args = ObjectiveArgs {
robot: &robot,
config: &SolverConfig {
gradient_mode: GradientMode::Analytical,
..Default::default()
},
tfm_target: &tfm_target,
};

let mut g_a = Vector6::zeros();
args.config.gradient_mode = GradientMode::Analytical;
objective_grad(x0.as_slice(), g_a.as_mut_slice(), &args);
objective_grad(
x0.as_slice(),
g_a.as_mut_slice(),
&args,
&mut KinematicsCache::default(),
);

// Numerical gradient
let args = ObjectiveArgs {
robot: &robot,
config: &SolverConfig {
gradient_mode: GradientMode::Numerical,
..Default::default()
},
tfm_target: &tfm_target,
};

let mut g_n = Vector6::zeros();
args.config.gradient_mode = GradientMode::Numerical;
objective_grad(x0.as_slice(), g_n.as_mut_slice(), &args);
objective_grad(
x0.as_slice(),
g_n.as_mut_slice(),
&args,
&mut KinematicsCache::default(),
);

assert_abs_diff_eq!(g_a, g_n, epsilon = 1e-6);
}

0 comments on commit 8fa7ec5

Please sign in to comment.