rust_robotics_control/
n_joint_arm_control.rs

1#![allow(
2    dead_code,
3    clippy::needless_borrows_for_generic_args,
4    clippy::new_without_default
5)]
6
7//! Generalized N-joint planar arm inverse kinematics and control.
8//!
9//! Uses the Jacobian pseudo-inverse method to solve inverse kinematics
10//! for a planar arm with an arbitrary number of revolute joints.
11//!
12//! Reference: PythonRobotics ArmNavigation/n\_joint\_arm\_to\_point\_control
13
14use std::f64::consts::PI;
15
16/// Simulation parameters
17const KP: f64 = 2.0;
18const DT: f64 = 0.1;
19const DEFAULT_N_ITERATIONS: usize = 10000;
20const DEFAULT_GOAL_THRESHOLD: f64 = 0.1;
21
22/// A planar N-link arm with forward and inverse kinematics.
23#[derive(Debug, Clone)]
24pub struct NLinkArm {
25    /// Length of each link
26    pub link_lengths: Vec<f64>,
27    /// Current joint angles \[rad\]
28    pub joint_angles: Vec<f64>,
29    /// Computed joint positions (n\_links + 1 points, including the base at origin)
30    pub points: Vec<[f64; 2]>,
31    /// End-effector position
32    pub end_effector: [f64; 2],
33    /// Number of links
34    pub n_links: usize,
35}
36
37impl NLinkArm {
38    /// Creates a new N-link arm.
39    ///
40    /// # Panics
41    ///
42    /// Panics if `link_lengths` and `joint_angles` have different lengths.
43    pub fn new(link_lengths: Vec<f64>, joint_angles: Vec<f64>) -> Self {
44        assert_eq!(
45            link_lengths.len(),
46            joint_angles.len(),
47            "link_lengths and joint_angles must have the same length"
48        );
49        let n_links = link_lengths.len();
50        let mut arm = NLinkArm {
51            link_lengths,
52            joint_angles,
53            points: vec![[0.0, 0.0]; n_links + 1],
54            end_effector: [0.0, 0.0],
55            n_links,
56        };
57        arm.update_points();
58        arm
59    }
60
61    /// Updates joint angles and recomputes link positions.
62    pub fn update_joints(&mut self, joint_angles: &[f64]) {
63        assert_eq!(joint_angles.len(), self.n_links);
64        self.joint_angles = joint_angles.to_vec();
65        self.update_points();
66    }
67
68    /// Recomputes all joint positions from current angles.
69    pub fn update_points(&mut self) {
70        self.points[0] = [0.0, 0.0];
71        for i in 1..=self.n_links {
72            let cumulative_angle: f64 = self.joint_angles[..i].iter().sum();
73            self.points[i][0] =
74                self.points[i - 1][0] + self.link_lengths[i - 1] * cumulative_angle.cos();
75            self.points[i][1] =
76                self.points[i - 1][1] + self.link_lengths[i - 1] * cumulative_angle.sin();
77        }
78        self.end_effector = self.points[self.n_links];
79    }
80
81    /// Total workspace radius (sum of all link lengths).
82    pub fn workspace_radius(&self) -> f64 {
83        self.link_lengths.iter().sum()
84    }
85}
86
87/// Computes the forward kinematics for given link lengths and joint angles.
88///
89/// Returns the end-effector position \[x, y\].
90pub fn forward_kinematics(link_lengths: &[f64], joint_angles: &[f64]) -> [f64; 2] {
91    let n = link_lengths.len();
92    let mut x = 0.0;
93    let mut y = 0.0;
94    for i in 1..=n {
95        let cumulative_angle: f64 = joint_angles[..i].iter().sum();
96        x += link_lengths[i - 1] * cumulative_angle.cos();
97        y += link_lengths[i - 1] * cumulative_angle.sin();
98    }
99    [x, y]
100}
101
102/// Computes the Jacobian matrix for the planar arm.
103///
104/// Returns a 2-by-n matrix stored as `[row0, row1]` where each row has `n` elements.
105/// The Jacobian relates joint velocities to end-effector velocity.
106fn jacobian(link_lengths: &[f64], joint_angles: &[f64]) -> Vec<Vec<f64>> {
107    let n = link_lengths.len();
108    let mut j0 = vec![0.0; n];
109    let mut j1 = vec![0.0; n];
110    for (i, (j0_val, j1_val)) in j0.iter_mut().zip(j1.iter_mut()).enumerate() {
111        for k in i..n {
112            let cumulative_angle: f64 = joint_angles[..=k].iter().sum();
113            *j0_val -= link_lengths[k] * cumulative_angle.sin();
114            *j1_val += link_lengths[k] * cumulative_angle.cos();
115        }
116    }
117    vec![j0, j1]
118}
119
120/// Computes the pseudo-inverse of a 2-by-n matrix using the formula:
121/// `J^+ = J^T (J J^T)^{-1}`
122fn pseudo_inverse_2xn(j: &[Vec<f64>]) -> Vec<Vec<f64>> {
123    let n = j[0].len();
124
125    // Compute J * J^T (2x2 matrix)
126    let mut jjt = [[0.0f64; 2]; 2];
127    for row in 0..2 {
128        for col in 0..2 {
129            for (jr, jc) in j[row].iter().zip(j[col].iter()) {
130                jjt[row][col] += jr * jc;
131            }
132        }
133    }
134
135    // Invert the 2x2 matrix
136    let det = jjt[0][0] * jjt[1][1] - jjt[0][1] * jjt[1][0];
137    if det.abs() < 1e-12 {
138        // Near-singular: return zero matrix (damped least squares fallback)
139        return vec![vec![0.0; 2]; n];
140    }
141    let inv_det = 1.0 / det;
142    let inv = [
143        [jjt[1][1] * inv_det, -jjt[0][1] * inv_det],
144        [-jjt[1][0] * inv_det, jjt[0][0] * inv_det],
145    ];
146
147    // J^T * (J J^T)^{-1}  => n x 2
148    let mut result = vec![vec![0.0; 2]; n];
149    for (i, result_row) in result.iter_mut().enumerate().take(n) {
150        for (col, result_cell) in result_row.iter_mut().enumerate().take(2) {
151            for k in 0..2 {
152                *result_cell += j[k][i] * inv[k][col];
153            }
154        }
155    }
156    result
157}
158
159/// Computes the distance and error vector from current position to goal.
160pub fn distance_to_goal(current: &[f64; 2], goal: &[f64; 2]) -> ([f64; 2], f64) {
161    let dx = goal[0] - current[0];
162    let dy = goal[1] - current[1];
163    ([dx, dy], (dx * dx + dy * dy).sqrt())
164}
165
166/// Normalizes an angle to the range \[-pi, pi\].
167pub fn angle_mod(angle: f64) -> f64 {
168    let mut a = angle % (2.0 * PI);
169    if a > PI {
170        a -= 2.0 * PI;
171    } else if a < -PI {
172        a += 2.0 * PI;
173    }
174    a
175}
176
177/// Computes element-wise angle difference, each wrapped to \[-pi, pi\].
178fn ang_diff(theta1: &[f64], theta2: &[f64]) -> Vec<f64> {
179    theta1
180        .iter()
181        .zip(theta2.iter())
182        .map(|(a, b)| angle_mod(a - b))
183        .collect()
184}
185
186/// Solves inverse kinematics using the Jacobian pseudo-inverse method.
187///
188/// Returns `Some(joint_angles)` if a solution is found within tolerance,
189/// or `None` if the maximum number of iterations is exceeded.
190pub fn inverse_kinematics(
191    link_lengths: &[f64],
192    initial_joint_angles: &[f64],
193    goal_pos: &[f64; 2],
194) -> Option<Vec<f64>> {
195    inverse_kinematics_with_params(
196        link_lengths,
197        initial_joint_angles,
198        goal_pos,
199        DEFAULT_N_ITERATIONS,
200        DEFAULT_GOAL_THRESHOLD,
201    )
202}
203
204/// Solves inverse kinematics with configurable iteration limit and threshold.
205pub fn inverse_kinematics_with_params(
206    link_lengths: &[f64],
207    initial_joint_angles: &[f64],
208    goal_pos: &[f64; 2],
209    max_iterations: usize,
210    goal_threshold: f64,
211) -> Option<Vec<f64>> {
212    let mut joint_angles = initial_joint_angles.to_vec();
213
214    for _iteration in 0..max_iterations {
215        let current_pos = forward_kinematics(link_lengths, &joint_angles);
216        let (errors, distance) = distance_to_goal(&current_pos, goal_pos);
217        if distance < goal_threshold {
218            return Some(joint_angles);
219        }
220
221        let j = jacobian(link_lengths, &joint_angles);
222        let j_pinv = pseudo_inverse_2xn(&j);
223
224        // joint_angles += J_pinv * errors
225        for i in 0..joint_angles.len() {
226            joint_angles[i] += j_pinv[i][0] * errors[0] + j_pinv[i][1] * errors[1];
227        }
228    }
229    None
230}
231
232/// Runs the proportional-control IK loop: smoothly moves from current angles
233/// toward goal angles.
234///
235/// Returns the trajectory of joint angle snapshots.
236pub fn move_to_goal(
237    arm: &mut NLinkArm,
238    goal_angles: &[f64],
239    goal_pos: &[f64; 2],
240    goal_threshold: f64,
241) -> Vec<Vec<f64>> {
242    let mut trajectory = Vec::new();
243
244    for _ in 0..10000 {
245        let (_, distance) = distance_to_goal(&arm.end_effector, goal_pos);
246        if distance < goal_threshold {
247            break;
248        }
249        let diff = ang_diff(goal_angles, &arm.joint_angles);
250        let new_angles: Vec<f64> = arm
251            .joint_angles
252            .iter()
253            .zip(diff.iter())
254            .map(|(a, d)| a + KP * d * DT)
255            .collect();
256        arm.update_joints(&new_angles);
257        trajectory.push(arm.joint_angles.clone());
258    }
259    trajectory
260}
261
262#[cfg(test)]
263mod tests {
264    use super::*;
265    use std::f64::consts::FRAC_PI_4;
266
267    #[test]
268    fn test_forward_kinematics_straight() {
269        let links = vec![1.0, 1.0, 1.0];
270        let angles = vec![0.0, 0.0, 0.0];
271        let [x, y] = forward_kinematics(&links, &angles);
272        assert!((x - 3.0).abs() < 1e-10);
273        assert!(y.abs() < 1e-10);
274    }
275
276    #[test]
277    fn test_forward_kinematics_bent() {
278        let links = vec![1.0, 1.0];
279        let angles = vec![0.0, FRAC_PI_4];
280        let [x, y] = forward_kinematics(&links, &angles);
281        // link1 at angle 0: (1, 0)
282        // link2 at cumulative angle pi/4: (1 + cos(pi/4), sin(pi/4))
283        let expected_x = 1.0 + (FRAC_PI_4).cos();
284        let expected_y = (FRAC_PI_4).sin();
285        assert!((x - expected_x).abs() < 1e-10);
286        assert!((y - expected_y).abs() < 1e-10);
287    }
288
289    #[test]
290    fn test_nlink_arm_new() {
291        let arm = NLinkArm::new(vec![1.0, 1.0], vec![0.0, 0.0]);
292        assert_eq!(arm.n_links, 2);
293        assert!((arm.end_effector[0] - 2.0).abs() < 1e-10);
294        assert!(arm.end_effector[1].abs() < 1e-10);
295        assert_eq!(arm.points.len(), 3);
296    }
297
298    #[test]
299    fn test_nlink_arm_update_joints() {
300        let mut arm = NLinkArm::new(vec![1.0, 1.0], vec![0.0, 0.0]);
301        arm.update_joints(&[PI / 2.0, 0.0]);
302        // Both links point up
303        assert!(arm.end_effector[0].abs() < 1e-10);
304        assert!((arm.end_effector[1] - 2.0).abs() < 1e-10);
305    }
306
307    #[test]
308    fn test_inverse_kinematics_reachable() {
309        // Use a 2-link arm with initial angles that break symmetry
310        let links = vec![1.0, 1.0];
311        let initial = vec![0.3, 0.3]; // avoid singular straight config
312        let goal = [1.0, 1.0]; // well within reach (max reach = 2.0)
313
314        let result = inverse_kinematics(&links, &initial, &goal);
315        assert!(result.is_some(), "IK should converge for reachable target");
316
317        let angles = result.unwrap();
318        let pos = forward_kinematics(&links, &angles);
319        let dist = ((pos[0] - goal[0]).powi(2) + (pos[1] - goal[1]).powi(2)).sqrt();
320        assert!(dist < DEFAULT_GOAL_THRESHOLD);
321    }
322
323    #[test]
324    fn test_inverse_kinematics_unreachable() {
325        let links = vec![1.0, 1.0];
326        let initial = vec![0.0, 0.0];
327        // Way outside workspace radius of 2.0
328        let goal = [10.0, 10.0];
329
330        let result = inverse_kinematics_with_params(&links, &initial, &goal, 500, 0.01);
331        assert!(result.is_none());
332    }
333
334    #[test]
335    fn test_angle_mod() {
336        assert!((angle_mod(0.0)).abs() < 1e-10);
337        assert!((angle_mod(2.0 * PI) - 0.0).abs() < 1e-10);
338        assert!((angle_mod(PI + 0.1) - (-PI + 0.1)).abs() < 1e-10);
339        assert!((angle_mod(-PI - 0.1) - (PI - 0.1)).abs() < 1e-10);
340    }
341
342    #[test]
343    fn test_distance_to_goal() {
344        let (errors, dist) = distance_to_goal(&[0.0, 0.0], &[3.0, 4.0]);
345        assert!((errors[0] - 3.0).abs() < 1e-10);
346        assert!((errors[1] - 4.0).abs() < 1e-10);
347        assert!((dist - 5.0).abs() < 1e-10);
348    }
349
350    #[test]
351    fn test_jacobian_shape() {
352        let links = vec![1.0, 1.0, 1.0];
353        let angles = vec![0.0, 0.0, 0.0];
354        let j = jacobian(&links, &angles);
355        assert_eq!(j.len(), 2);
356        assert_eq!(j[0].len(), 3);
357        assert_eq!(j[1].len(), 3);
358    }
359
360    #[test]
361    fn test_move_to_goal_converges() {
362        let mut arm = NLinkArm::new(vec![1.0, 1.0], vec![0.5, 0.5]);
363        let goal = [1.5, 0.5];
364        let goal_angles =
365            inverse_kinematics(&arm.link_lengths.clone(), &arm.joint_angles.clone(), &goal)
366                .unwrap();
367        let traj = move_to_goal(&mut arm, &goal_angles, &goal, 0.1);
368        assert!(!traj.is_empty());
369        let (_, dist) = distance_to_goal(&arm.end_effector, &goal);
370        assert!(dist < 0.2);
371    }
372
373    #[test]
374    fn test_10_link_arm() {
375        let n = 10;
376        let links = vec![0.5; n]; // total reach = 5.0
377        let initial = vec![0.1; n]; // slight spread to avoid singular configuration
378        let goal = [3.0, 1.0]; // well within reach
379
380        let result = inverse_kinematics(&links, &initial, &goal);
381        assert!(result.is_some(), "10-link IK should converge");
382
383        let angles = result.unwrap();
384        let pos = forward_kinematics(&links, &angles);
385        let dist = ((pos[0] - goal[0]).powi(2) + (pos[1] - goal[1]).powi(2)).sqrt();
386        assert!(dist < DEFAULT_GOAL_THRESHOLD);
387    }
388
389    #[test]
390    fn test_workspace_radius() {
391        let arm = NLinkArm::new(vec![1.0, 2.0, 0.5], vec![0.0, 0.0, 0.0]);
392        assert!((arm.workspace_radius() - 3.5).abs() < 1e-10);
393    }
394
395    #[test]
396    fn test_pseudo_inverse_identity_like() {
397        // For a 2x2 non-singular Jacobian, pinv should be close to the true inverse
398        let j = vec![vec![1.0, 0.0], vec![0.0, 1.0]];
399        let pinv = pseudo_inverse_2xn(&j);
400        assert!((pinv[0][0] - 1.0).abs() < 1e-10);
401        assert!((pinv[0][1] - 0.0).abs() < 1e-10);
402        assert!((pinv[1][0] - 0.0).abs() < 1e-10);
403        assert!((pinv[1][1] - 1.0).abs() < 1e-10);
404    }
405}
rust_robotics_control/n_joint_arm_control.rs

rust_robotics_control/
n_joint_arm_control.rs
