WIP: Implement custom cached forward kinematics

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) {

crates/optik/src/kinematics.rs

@@ -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
+    }
+}