From 0dd73c8f9216489a125dff60d0133fd20e24e4de Mon Sep 17 00:00:00 2001 From: Kyle Cesare Date: Sun, 26 Nov 2023 15:40:56 -0700 Subject: [PATCH] WIP: Implement custom cached forward kinematics --- crates/optik/benches/bench.rs | 58 +++++--- crates/optik/src/kinematics.rs | 132 ++++++++++++++++++ crates/optik/src/lib.rs | 200 ++++++++++------------------ crates/optik/src/objective.rs | 89 +++++++++++++ crates/optik/tests/test_gradient.rs | 44 ++++-- 5 files changed, 367 insertions(+), 156 deletions(-) create mode 100644 crates/optik/src/kinematics.rs create mode 100644 crates/optik/src/objective.rs diff --git a/crates/optik/benches/bench.rs b/crates/optik/benches/bench.rs index b87fbdf..6ec6ddb 100644 --- a/crates/optik/benches/bench.rs +++ b/crates/optik/benches/bench.rs @@ -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) { diff --git a/crates/optik/src/kinematics.rs b/crates/optik/src/kinematics.rs new file mode 100644 index 0000000..ff5609a --- /dev/null +++ b/crates/optik/src/kinematics.rs @@ -0,0 +1,132 @@ +use nalgebra::{Isometry3, Matrix6xX, RealField, Translation3, Unit, UnitQuaternion, Vector3}; + +#[derive(Debug, Clone)] +pub struct Joint { + offset: Isometry3, + typ: JointType, +} + +#[derive(Debug, Clone)] +enum JointType { + Revolute { axis: Unit> }, + Prismatic { axis: Unit> }, +} + +impl Joint { + pub fn revolute(offset: Isometry3, axis: Unit>) -> Joint { + Joint { + offset, + typ: JointType::Revolute { axis }, + } + } + + pub fn prismatic(offset: Isometry3, axis: Unit>) -> Joint { + Joint { + offset, + typ: JointType::Prismatic { axis }, + } + } + + fn local_transform(&self, q: &[T]) -> Isometry3 { + 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::() - 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], + q: &'a [f64], +) -> impl Iterator> + 'a { + struct State { + qidx: usize, + tfm: Isometry3, + } + + 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], + ee_offset: &Isometry3, + tfms: &[Isometry3], +) -> Matrix6xX { + 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, + frames: Vec>, +} + +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], q: &[f64]) -> &[Isometry3] { + 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 + } +} diff --git a/crates/optik/src/lib.rs b/crates/optik/src/lib.rs index afe5b38..68a0af9 100644 --- a/crates/optik/src/lib.rs +++ b/crates/optik/src/lib.rs @@ -1,4 +1,5 @@ use std::{ + cell::RefCell, path::Path, sync::{ atomic::{AtomicBool, Ordering}, @@ -8,8 +9,8 @@ use std::{ }; use float_ord::FloatOrd; -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 rand::SeedableRng; use rand_chacha::ChaCha8Rng; use rayon::{ @@ -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, + chain: SerialChain, + joints: Vec>, + ee_offset: Isometry3, } impl Robot { pub fn new(chain: SerialChain) -> 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, base_link: &str, ee_link: &str) -> Self { @@ -65,31 +94,8 @@ impl Robot { Robot::new(serial) } - pub fn jacobian_local(&self, q: &[f64]) -> DMatrix { - 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, DVector) { @@ -107,14 +113,20 @@ impl Robot { pub fn random_configuration(&self, rng: &mut impl rand::Rng) -> Vec { 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]) -> Matrix6xX { + joint_jacobian(&self.joints, &self.ee_offset, frames) + } + pub fn fk(&self, q: &[f64]) -> Isometry3 { - 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()); FloatOrd(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, -} - -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; - } - } - } -} diff --git a/crates/optik/src/objective.rs b/crates/optik/src/objective.rs new file mode 100644 index 0000000..d57f6e7 --- /dev/null +++ b/crates/optik/src/objective.rs @@ -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, +} + +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; + } + } + } +} diff --git a/crates/optik/tests/test_gradient.rs b/crates/optik/tests/test_gradient.rs index 4d9599f..4ea41a9 100644 --- a/crates/optik/tests/test_gradient.rs +++ b/crates/optik/tests/test_gradient.rs @@ -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 = 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); }