Skip to content

Commit

Permalink
WIP: Implement custom cached forward kinematics
Browse files Browse the repository at this point in the history
  • Loading branch information
kylc committed Nov 26, 2023
1 parent dfbde06 commit 3f32456
Show file tree
Hide file tree
Showing 3 changed files with 249 additions and 72 deletions.
63 changes: 40 additions & 23 deletions crates/optik/benches/bench.rs
Original file line number Diff line number Diff line change
@@ -1,26 +1,27 @@
use criterion::{black_box, criterion_group, criterion_main, Criterion};
use k::Chain;
use nalgebra::Isometry3;

use optik::*;

const BENCHMARK_MODEL_PATH: &str = "../../models/ur3e.urdf";
const BENCHMARK_MODEL_BASE_NAME: &str = "ur_base_link";
const BENCHMARK_MODEL_EE_NAME: &str = "ur_ee_link";

fn load_benchmark_model() -> Robot {
let chain = Chain::<f64>::from_urdf_file(BENCHMARK_MODEL_PATH).unwrap();
let serial = k::SerialChain::from_end(chain.find_link(BENCHMARK_MODEL_EE_NAME).unwrap());

Robot::new(serial)
Robot::from_urdf_file(
BENCHMARK_MODEL_PATH,
BENCHMARK_MODEL_BASE_NAME,
BENCHMARK_MODEL_EE_NAME,
)
}

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) {
Expand All @@ -30,31 +31,45 @@ fn bench_gradient(c: &mut Criterion) {
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(),
)
})
});
}

Expand All @@ -64,11 +79,13 @@ fn bench_objective(c: &mut Criterion) {
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) {
Expand Down
120 changes: 120 additions & 0 deletions crates/optik/src/kinematics.rs
Original file line number Diff line number Diff line change
@@ -0,0 +1,120 @@
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 {
// let mut tfms = vec![];

// let mut tfm = Isometry3::identity();
// let mut qidx = 0;
// for joint in chain.iter() {
// tfm *= joint.local_transform(&q[qidx..(qidx + joint.nq())]);
// tfms.push(tfm);

// qidx += joint.nq();
// }

struct State {
qidx: usize,
tfm: Isometry3<f64>,
}

// let mut qidx = 0;
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)
},
)

// (tfms, tfm * ee_tfm)
}

pub fn 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
}
Loading

0 comments on commit 3f32456

Please sign in to comment.