-
Notifications
You must be signed in to change notification settings - Fork 4
Commit
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
Showing
5 changed files
with
367 additions
and
156 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
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 | ||
} | ||
} |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
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; | ||
} | ||
} | ||
} | ||
} |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters