Skip to main content

nabled_sim/
manipulation.rs

1//! Batch damped least-squares IK over a time grid.
2
3use nabled_core::scalar::NabledReal;
4use nabled_kinematics::chain::{ChainSpec, JointLimits};
5use nabled_kinematics::fk::fk_view;
6use nabled_kinematics::ik::{
7    IkConfig, IkResult, inverse_kinematics_dls_with_limits,
8    inverse_kinematics_tree_dls_with_limits, pose_error,
9};
10use nabled_kinematics::tree::KinematicTreeModel;
11use nabled_linalg::lu::LuProviderScalar;
12use ndarray::{Array1, ArrayView1};
13
14use crate::SimError;
15
16fn error_norm<T: NabledReal>(error: &Array1<T>) -> T {
17    error.iter().map(|v| *v * *v).fold(T::zero(), |acc, v| acc + v).sqrt()
18}
19
20/// IK solver options for trajectory batching.
21#[derive(Debug, Clone, PartialEq)]
22pub struct TrajectoryIkConfig<T> {
23    pub ik_config: IkConfig<T>,
24    pub limits:    Option<Vec<JointLimits<T>>>,
25}
26
27impl<T: NabledReal> Default for TrajectoryIkConfig<T> {
28    fn default() -> Self { Self { ik_config: IkConfig::default(), limits: None } }
29}
30
31/// One time-grid IK solve with pose verification.
32#[derive(Debug, Clone, PartialEq)]
33pub struct TrajectoryIkStep<T> {
34    pub t:               T,
35    pub result:          IkResult<T>,
36    pub pose_error_norm: T,
37}
38
39/// Aggregated trajectory IK output.
40#[derive(Debug, Clone, PartialEq)]
41pub struct TrajectoryIkResult<T> {
42    pub steps:           Vec<TrajectoryIkStep<T>>,
43    pub max_error:       T,
44    pub converged_count: usize,
45}
46
47/// Interpolated joint-space targets over a time grid, solved with serial DLS IK.
48#[derive(Debug, Clone, PartialEq)]
49pub struct TrajectoryIk<T> {
50    pub times:   Vec<T>,
51    pub q_start: Array1<T>,
52    pub q_end:   Array1<T>,
53}
54
55impl<T: NabledReal + LuProviderScalar> TrajectoryIk<T> {
56    /// Solve IK at each time sample by linearly blending `q_start` → `q_end` FK targets.
57    pub fn solve(
58        &self,
59        chain: &ChainSpec<T>,
60        config: &TrajectoryIkConfig<T>,
61    ) -> Result<TrajectoryIkResult<T>, SimError> {
62        if self.times.is_empty() {
63            return Err(SimError::InvalidInput("time grid cannot be empty".to_string()));
64        }
65        if self.q_start.len() != chain.num_joints() || self.q_end.len() != chain.num_joints() {
66            return Err(SimError::DimensionMismatch);
67        }
68
69        let duration = self.times[self.times.len() - 1];
70        if duration.is_zero() {
71            return Err(SimError::InvalidInput("time grid duration must be positive".to_string()));
72        }
73
74        let limits = config.limits.as_deref();
75        let mut q_seed = self.q_start.clone();
76        let mut max_error = T::zero();
77        let mut converged_count = 0_usize;
78        let mut steps = Vec::with_capacity(self.times.len());
79
80        for &t in &self.times {
81            let blend = t / duration;
82            let q_target = &self.q_start + &((&self.q_end - &self.q_start) * blend);
83            let target = fk_view(chain, &q_target.view())?;
84            let result = inverse_kinematics_dls_with_limits(
85                chain,
86                &q_seed,
87                &target,
88                &config.ik_config,
89                limits,
90            )?;
91            let achieved = fk_view(chain, &result.q.view())?;
92            let err = pose_error(&achieved, &target)?;
93            let err_norm = error_norm(&err);
94            if err_norm > max_error {
95                max_error = err_norm;
96            }
97            if result.converged {
98                converged_count += 1;
99            }
100            q_seed.clone_from(&result.q);
101            steps.push(TrajectoryIkStep { t, result, pose_error_norm: err_norm });
102        }
103
104        Ok(TrajectoryIkResult { steps, max_error, converged_count })
105    }
106
107    /// Solve IK from explicit SE(3) targets (one per time sample).
108    pub fn solve_targets(
109        times: &[T],
110        chain: &ChainSpec<T>,
111        q_init: &ArrayView1<'_, T>,
112        targets: &[nabled_linalg::geometry::Transform3<T>],
113        config: &TrajectoryIkConfig<T>,
114    ) -> Result<TrajectoryIkResult<T>, SimError> {
115        if times.len() != targets.len() {
116            return Err(SimError::DimensionMismatch);
117        }
118        if times.is_empty() {
119            return Err(SimError::InvalidInput("time grid cannot be empty".to_string()));
120        }
121        if q_init.len() != chain.num_joints() {
122            return Err(SimError::DimensionMismatch);
123        }
124
125        let limits = config.limits.as_deref();
126        let mut q_seed = q_init.to_owned();
127        let mut max_error = T::zero();
128        let mut converged_count = 0_usize;
129        let mut steps = Vec::with_capacity(times.len());
130
131        for (&t, target) in times.iter().zip(targets.iter()) {
132            let result = inverse_kinematics_dls_with_limits(
133                chain,
134                &q_seed,
135                target,
136                &config.ik_config,
137                limits,
138            )?;
139            let achieved = fk_view(chain, &result.q.view())?;
140            let err = pose_error(&achieved, target)?;
141            let err_norm = error_norm(&err);
142            if err_norm > max_error {
143                max_error = err_norm;
144            }
145            if result.converged {
146                converged_count += 1;
147            }
148            q_seed.clone_from(&result.q);
149            steps.push(TrajectoryIkStep { t, result, pose_error_norm: err_norm });
150        }
151
152        Ok(TrajectoryIkResult { steps, max_error, converged_count })
153    }
154}
155
156/// Tree-model trajectory IK over a time grid (full-model `q` in actuated order).
157#[derive(Debug, Clone, PartialEq)]
158pub struct TrajectoryTreeIk<T> {
159    pub times:     Vec<T>,
160    pub q_start:   Array1<T>,
161    pub q_end:     Array1<T>,
162    pub base_link: String,
163    pub ee_link:   String,
164}
165
166impl<T: NabledReal + LuProviderScalar> TrajectoryTreeIk<T> {
167    /// Solve tree DLS IK at each time sample by blending joint targets then FK-ing to SE(3).
168    pub fn solve<M: KinematicTreeModel<T>>(
169        &self,
170        model: &M,
171        config: &TrajectoryIkConfig<T>,
172    ) -> Result<TrajectoryIkResult<T>, SimError> {
173        if self.times.is_empty() {
174            return Err(SimError::InvalidInput("time grid cannot be empty".to_string()));
175        }
176        if self.q_start.len() != model.dof() || self.q_end.len() != model.dof() {
177            return Err(SimError::DimensionMismatch);
178        }
179
180        let duration = self.times[self.times.len() - 1];
181        if duration.is_zero() {
182            return Err(SimError::InvalidInput("time grid duration must be positive".to_string()));
183        }
184
185        let limits = config.limits.as_deref();
186        let mut q_seed = self.q_start.clone();
187        let mut max_error = T::zero();
188        let mut converged_count = 0_usize;
189        let mut steps = Vec::with_capacity(self.times.len());
190
191        for &t in &self.times {
192            let blend = t / duration;
193            let q_target = &self.q_start + &((&self.q_end - &self.q_start) * blend);
194            let target = nabled_kinematics::tree::end_effector_pose_tree(
195                model,
196                &self.base_link,
197                &self.ee_link,
198                &q_target,
199            )?;
200            let result = inverse_kinematics_tree_dls_with_limits(
201                model,
202                &self.base_link,
203                &self.ee_link,
204                &q_seed,
205                &target,
206                &config.ik_config,
207                limits,
208            )?;
209            let achieved = nabled_kinematics::tree::end_effector_pose_tree(
210                model,
211                &self.base_link,
212                &self.ee_link,
213                &result.q,
214            )?;
215            let err = pose_error(&achieved, &target)?;
216            let err_norm = error_norm(&err);
217            if err_norm > max_error {
218                max_error = err_norm;
219            }
220            if result.converged {
221                converged_count += 1;
222            }
223            q_seed.clone_from(&result.q);
224            steps.push(TrajectoryIkStep { t, result, pose_error_norm: err_norm });
225        }
226
227        Ok(TrajectoryIkResult { steps, max_error, converged_count })
228    }
229}
230
231#[cfg(test)]
232mod tests {
233    use nabled_kinematics::fk::{end_effector_pose, fk_view};
234    use nabled_model::fixture::{load_planar2r_json, load_six_dof_dh_json, load_y_branch_json};
235    use nabled_model::urdf::from_urdf_file;
236    use ndarray::arr1;
237
238    use super::*;
239
240    #[test]
241    fn trajectory_ik_converges_on_grid() {
242        let fixture = load_six_dof_dh_json().expect("fixture");
243        let chain = fixture.to_chain_spec::<f64>().expect("chain");
244        let trajectory = TrajectoryIk {
245            times:   (0..=5).map(|i| f64::from(i) * 0.1).collect(),
246            q_start: arr1(&[0.0, -0.3, 0.5, 0.2, -0.1, 0.4]),
247            q_end:   arr1(&[0.3, -0.2, 0.4, 0.1, 0.0, 0.5]),
248        };
249        let config = TrajectoryIkConfig {
250            ik_config: IkConfig { max_iterations: 200, tolerance: 1e-3, ..IkConfig::default() },
251            limits:    None,
252        };
253        let result = trajectory.solve(&chain, &config).expect("solve");
254        assert_eq!(result.steps.len(), 6);
255        assert!(result.max_error < 1e-2);
256    }
257
258    #[test]
259    fn trajectory_ik_rejects_empty_time_grid() {
260        let fixture = load_six_dof_dh_json().expect("fixture");
261        let chain = fixture.to_chain_spec::<f64>().expect("chain");
262        let trajectory =
263            TrajectoryIk { times: vec![], q_start: arr1(&[0.0; 6]), q_end: arr1(&[0.1; 6]) };
264        let config = TrajectoryIkConfig::default();
265        assert_eq!(
266            trajectory.solve(&chain, &config),
267            Err(SimError::InvalidInput("time grid cannot be empty".to_string()))
268        );
269    }
270
271    #[test]
272    fn trajectory_ik_rejects_zero_duration() {
273        let fixture = load_six_dof_dh_json().expect("fixture");
274        let chain = fixture.to_chain_spec::<f64>().expect("chain");
275        let trajectory = TrajectoryIk {
276            times:   vec![0.0, 0.0],
277            q_start: arr1(&[0.0; 6]),
278            q_end:   arr1(&[0.1; 6]),
279        };
280        let config = TrajectoryIkConfig::default();
281        assert_eq!(
282            trajectory.solve(&chain, &config),
283            Err(SimError::InvalidInput("time grid duration must be positive".to_string()))
284        );
285    }
286
287    #[test]
288    fn trajectory_ik_rejects_dof_mismatch() {
289        let fixture = load_six_dof_dh_json().expect("fixture");
290        let chain = fixture.to_chain_spec::<f64>().expect("chain");
291        let trajectory = TrajectoryIk {
292            times:   vec![0.0, 1.0],
293            q_start: arr1(&[0.0; 5]),
294            q_end:   arr1(&[0.1; 6]),
295        };
296        let config = TrajectoryIkConfig::default();
297        assert_eq!(trajectory.solve(&chain, &config), Err(SimError::DimensionMismatch));
298    }
299
300    #[test]
301    fn trajectory_ik_solve_targets_on_planar2r() {
302        let fixture = load_planar2r_json().expect("fixture");
303        let chain = fixture.to_chain_spec::<f64>().expect("chain");
304        let q_a = arr1(&[0.0_f64, 0.0]);
305        let q_b = arr1(&[0.5, -0.3]);
306        let targets = [&q_a, &q_b]
307            .iter()
308            .map(|q| end_effector_pose(&chain, *q).expect("fk"))
309            .collect::<Vec<_>>();
310        let times = [0.0_f64, 1.0];
311        let config = TrajectoryIkConfig {
312            ik_config: IkConfig { max_iterations: 200, tolerance: 1e-3, ..IkConfig::default() },
313            limits:    None,
314        };
315        let result = TrajectoryIk::solve_targets(&times, &chain, &q_a.view(), &targets, &config)
316            .expect("solve targets");
317        assert_eq!(result.steps.len(), 2);
318        assert!(result.max_error < 1e-2);
319    }
320
321    #[test]
322    fn trajectory_ik_solve_targets_on_six_dof() {
323        let fixture = load_six_dof_dh_json().expect("fixture");
324        let chain = fixture.to_chain_spec::<f64>().expect("chain");
325        let q_init = arr1(&[0.0, -0.3, 0.5, 0.2, -0.1, 0.4]);
326        let q_goal = arr1(&[0.3, -0.2, 0.4, 0.1, 0.0, 0.5]);
327        let targets = [q_init.view(), q_goal.view()]
328            .iter()
329            .map(|q| fk_view(&chain, q).expect("fk"))
330            .collect::<Vec<_>>();
331        let times = [0.0_f64, 1.0];
332        let config = TrajectoryIkConfig {
333            ik_config: IkConfig { max_iterations: 200, tolerance: 1e-3, ..IkConfig::default() },
334            limits:    None,
335        };
336        let result = TrajectoryIk::solve_targets(&times, &chain, &q_init.view(), &targets, &config)
337            .expect("solve targets");
338        assert_eq!(result.steps.len(), 2);
339        assert!(result.max_error < 1e-2);
340    }
341
342    #[test]
343    fn trajectory_ik_solve_targets_rejects_mismatched_grid() {
344        let fixture = load_planar2r_json().expect("fixture");
345        let chain = fixture.to_chain_spec::<f64>().expect("chain");
346        let q = arr1(&[0.0_f64, 0.0]);
347        let target = end_effector_pose(&chain, &q).expect("fk");
348        let config = TrajectoryIkConfig::default();
349        assert_eq!(
350            TrajectoryIk::solve_targets(&[0.0, 1.0], &chain, &q.view(), &[target], &config),
351            Err(SimError::DimensionMismatch)
352        );
353    }
354
355    #[test]
356    fn trajectory_tree_ik_on_y_branch() {
357        let fixture = load_y_branch_json().expect("fixture");
358        let urdf_path = concat!(
359            env!("CARGO_MANIFEST_DIR"),
360            "/../nabled/tests/fixtures/physical_ai/Y_branch.urdf"
361        );
362        let model = from_urdf_file::<f64>(urdf_path).expect("urdf");
363        let case = &fixture.cases[0];
364        let trajectory = TrajectoryTreeIk {
365            times:     vec![0.0_f64, 0.5, 1.0],
366            q_start:   arr1(&case.q),
367            q_end:     arr1(&[0.3, 0.5, -0.2]),
368            base_link: "base".to_string(),
369            ee_link:   "left_ee".to_string(),
370        };
371        let config = TrajectoryIkConfig {
372            ik_config: IkConfig { max_iterations: 300, tolerance: 1e-3, ..IkConfig::default() },
373            limits:    None,
374        };
375        let result = trajectory.solve(&model, &config).expect("tree ik");
376        assert_eq!(result.steps.len(), 3);
377        assert!(result.max_error < 1e-2);
378    }
379
380    #[test]
381    fn trajectory_tree_ik_rejects_empty_time_grid() {
382        let urdf_path = concat!(
383            env!("CARGO_MANIFEST_DIR"),
384            "/../nabled/tests/fixtures/physical_ai/Y_branch.urdf"
385        );
386        let model = from_urdf_file::<f64>(urdf_path).expect("urdf");
387        let trajectory = TrajectoryTreeIk {
388            times:     vec![],
389            q_start:   arr1(&[0.0; 3]),
390            q_end:     arr1(&[0.1; 3]),
391            base_link: "base".to_string(),
392            ee_link:   "right_ee".to_string(),
393        };
394        let config = TrajectoryIkConfig::default();
395        assert_eq!(
396            trajectory.solve(&model, &config),
397            Err(SimError::InvalidInput("time grid cannot be empty".to_string()))
398        );
399    }
400}