rust_robotics_control/
n_joint_arm_3d.rs

1#![allow(dead_code, clippy::new_without_default)]
2
3//! 3D N-joint arm control
4//!
5//! Extension of the planar N-link arm to 3D space. Uses alternating
6//! yaw/pitch joints to reach arbitrary 3D positions.
7//!
8//! Joint convention:
9//! - Even-indexed joints (0, 2, 4, ...) rotate around the Z axis (yaw)
10//! - Odd-indexed joints (1, 3, 5, ...) rotate around the Y axis (pitch)
11//!
12//! Each link extends along the local X direction after the accumulated
13//! rotations.
14
15use nalgebra::{Matrix3, Matrix3xX, Vector3};
16
17/// Epsilon for numerical Jacobian computation.
18const JACOBIAN_EPSILON: f64 = 1e-6;
19
20/// Damping factor for Levenberg-Marquardt IK.
21const DAMPING_LAMBDA: f64 = 0.5;
22
23/// A 3D N-link arm with alternating yaw/pitch joints.
24#[derive(Debug, Clone)]
25pub struct NLinkArm3D {
26    link_lengths: Vec<f64>,
27    joint_angles: Vec<f64>,
28}
29
30impl NLinkArm3D {
31    /// Creates a new arm with all joint angles initialized to zero.
32    pub fn new(link_lengths: Vec<f64>) -> Self {
33        let n = link_lengths.len();
34        Self {
35            link_lengths,
36            joint_angles: vec![0.0; n],
37        }
38    }
39
40    /// Creates a new arm with specified joint angles.
41    ///
42    /// # Panics
43    ///
44    /// Panics if `link_lengths` and `joint_angles` have different lengths.
45    pub fn with_angles(link_lengths: Vec<f64>, joint_angles: Vec<f64>) -> Self {
46        assert_eq!(
47            link_lengths.len(),
48            joint_angles.len(),
49            "link_lengths and joint_angles must have the same length"
50        );
51        Self {
52            link_lengths,
53            joint_angles,
54        }
55    }
56
57    /// Computes forward kinematics, returning positions of base, all joints,
58    /// and the end effector (n+1 points total).
59    ///
60    /// For each joint i:
61    /// - If i is even: rotate around Z by `joint_angles\[i\]` (yaw)
62    /// - If i is odd: rotate around Y by `joint_angles\[i\]` (pitch)
63    ///
64    /// Then advance by `link_lengths\[i\]` along the current local X direction.
65    pub fn forward_kinematics(&self) -> Vec<Vector3<f64>> {
66        let n = self.link_lengths.len();
67        let mut points = Vec::with_capacity(n + 1);
68        let mut pos = Vector3::zeros();
69        let mut rot = Matrix3::identity();
70
71        points.push(pos);
72
73        for i in 0..n {
74            let angle = self.joint_angles[i];
75            let local_rot = if i % 2 == 0 {
76                rotation_z(angle)
77            } else {
78                rotation_y(angle)
79            };
80            rot *= local_rot;
81            let direction = rot * Vector3::new(self.link_lengths[i], 0.0, 0.0);
82            pos += direction;
83            points.push(pos);
84        }
85
86        points
87    }
88
89    /// Returns the end-effector position.
90    pub fn end_effector(&self) -> Vector3<f64> {
91        let points = self.forward_kinematics();
92        // forward_kinematics always pushes at least the base position
93        *points
94            .last()
95            .expect("forward_kinematics always returns at least one point")
96    }
97
98    /// Computes the 3-by-n numerical Jacobian via finite differences.
99    ///
100    /// Each column j is `(end_effector(q + eps*e_j) - end_effector(q - eps*e_j)) / (2*eps)`.
101    pub fn jacobian(&self) -> Matrix3xX<f64> {
102        let n = self.joint_angles.len();
103        let mut jac = Matrix3xX::zeros(n);
104
105        for j in 0..n {
106            let mut angles_plus = self.joint_angles.clone();
107            let mut angles_minus = self.joint_angles.clone();
108            angles_plus[j] += JACOBIAN_EPSILON;
109            angles_minus[j] -= JACOBIAN_EPSILON;
110
111            let arm_plus = NLinkArm3D {
112                link_lengths: self.link_lengths.clone(),
113                joint_angles: angles_plus,
114            };
115            let arm_minus = NLinkArm3D {
116                link_lengths: self.link_lengths.clone(),
117                joint_angles: angles_minus,
118            };
119
120            let diff =
121                (arm_plus.end_effector() - arm_minus.end_effector()) / (2.0 * JACOBIAN_EPSILON);
122            jac.set_column(j, &diff);
123        }
124
125        jac
126    }
127
128    /// Solves inverse kinematics using damped least squares (Levenberg-Marquardt).
129    ///
130    /// `dq = J^T * (J * J^T + lambda * I)^{-1} * error`
131    ///
132    /// Returns `true` if the end effector converged to within `tolerance` of
133    /// the target.
134    pub fn inverse_kinematics(
135        &mut self,
136        target: Vector3<f64>,
137        max_iter: usize,
138        tolerance: f64,
139    ) -> bool {
140        for _ in 0..max_iter {
141            let ee = self.end_effector();
142            let error = target - ee;
143
144            if error.norm() < tolerance {
145                return true;
146            }
147
148            let j = self.jacobian();
149            // J * J^T is 3x3
150            let jjt = &j * j.transpose();
151            let damped = jjt + Matrix3::identity() * DAMPING_LAMBDA;
152
153            // Solve (J*J^T + lambda*I) * x = error for x, then dq = J^T * x
154            if let Some(decomp) = damped.try_inverse() {
155                let x = decomp * error;
156                let dq = j.transpose() * x;
157
158                for i in 0..self.joint_angles.len() {
159                    self.joint_angles[i] += dq[i];
160                }
161            } else {
162                // Singular; cannot proceed
163                return false;
164            }
165        }
166
167        // Check final convergence
168        let ee = self.end_effector();
169        (target - ee).norm() < tolerance
170    }
171
172    /// Returns a reference to the current joint angles.
173    pub fn joint_angles(&self) -> &[f64] {
174        &self.joint_angles
175    }
176
177    /// Returns a reference to the link lengths.
178    pub fn link_lengths(&self) -> &[f64] {
179        &self.link_lengths
180    }
181}
182
183/// Rotation matrix around the Z axis by angle theta.
184fn rotation_z(theta: f64) -> Matrix3<f64> {
185    let (s, c) = theta.sin_cos();
186    Matrix3::new(c, -s, 0.0, s, c, 0.0, 0.0, 0.0, 1.0)
187}
188
189/// Rotation matrix around the Y axis by angle theta.
190fn rotation_y(theta: f64) -> Matrix3<f64> {
191    let (s, c) = theta.sin_cos();
192    Matrix3::new(c, 0.0, s, 0.0, 1.0, 0.0, -s, 0.0, c)
193}
194
195#[cfg(test)]
196mod tests {
197    use super::*;
198
199    #[test]
200    fn test_3d_arm_creation() {
201        let arm = NLinkArm3D::new(vec![1.0, 1.0, 1.0]);
202        assert_eq!(arm.link_lengths().len(), 3);
203        assert_eq!(arm.joint_angles().len(), 3);
204        assert!(arm.joint_angles().iter().all(|&a| a == 0.0));
205
206        let arm2 = NLinkArm3D::with_angles(vec![1.0, 2.0], vec![0.5, -0.3]);
207        assert_eq!(arm2.link_lengths(), &[1.0, 2.0]);
208        assert_eq!(arm2.joint_angles(), &[0.5, -0.3]);
209    }
210
211    #[test]
212    fn test_forward_kinematics_straight() {
213        // All angles zero: every link extends along X.
214        let arm = NLinkArm3D::new(vec![1.0, 1.5, 0.5]);
215        let ee = arm.end_effector();
216        let total: f64 = arm.link_lengths().iter().sum();
217        assert!(
218            (ee.x - total).abs() < 1e-10,
219            "Expected x={total}, got x={}",
220            ee.x
221        );
222        assert!(ee.y.abs() < 1e-10, "Expected y=0, got y={}", ee.y);
223        assert!(ee.z.abs() < 1e-10, "Expected z=0, got z={}", ee.z);
224
225        // Also check that we get n+1 points
226        let points = arm.forward_kinematics();
227        assert_eq!(points.len(), 4);
228        assert!((points[0] - Vector3::zeros()).norm() < 1e-10);
229    }
230
231    #[test]
232    fn test_inverse_kinematics_reachable() {
233        // 4-link arm, total reach = 4.0. Target on X-axis at distance 3.0.
234        let mut arm = NLinkArm3D::with_angles(vec![1.0; 4], vec![0.1, 0.1, 0.1, 0.1]);
235        let target = Vector3::new(3.0, 0.0, 0.0);
236        let converged = arm.inverse_kinematics(target, 500, 0.05);
237        assert!(
238            converged,
239            "IK should converge for reachable in-plane target"
240        );
241
242        let ee = arm.end_effector();
243        let dist = (ee - target).norm();
244        assert!(
245            dist < 0.05,
246            "End effector should be within tolerance, got dist={dist}"
247        );
248    }
249
250    #[test]
251    fn test_inverse_kinematics_3d_target() {
252        // 6-link arm reaching a point above the base plane.
253        // Total reach = 6.0. Target at (2, 1, 2) has distance ~3.0.
254        let mut arm = NLinkArm3D::with_angles(vec![1.0; 6], vec![0.1; 6]);
255        let target = Vector3::new(2.0, 1.0, 2.0);
256        let converged = arm.inverse_kinematics(target, 1000, 0.05);
257        assert!(converged, "IK should converge for reachable 3D target");
258
259        let ee = arm.end_effector();
260        let dist = (ee - target).norm();
261        assert!(
262            dist < 0.05,
263            "End effector should reach 3D target, got dist={dist}"
264        );
265    }
266}
rust_robotics_control/n_joint_arm_3d.rs

rust_robotics_control/
n_joint_arm_3d.rs
