1use 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#[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#[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#[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#[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 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 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#[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 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(×, &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(×, &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}