rust_robotics_control/
two_joint_arm_control.rs

1#![allow(
2    dead_code,
3    clippy::needless_borrows_for_generic_args,
4    clippy::new_without_default
5)]
6
7//
8// Inverse kinematics of a two-joint arm
9// Author: Daniel Ingram (daniel-s-ingram)
10//         Atsushi Sakai (@Atsushi_twi)
11// Ported to Rust by: rust_robotics team
12//
13
14#[cfg(feature = "viz")]
15use gnuplot::{AxesCommon, Caption, Color, Figure, PointSymbol};
16
17use std::f64::consts::PI;
18
19// Simulation parameters
20const KP: f64 = 15.0;
21const DT: f64 = 0.01;
22
23// Link lengths
24const L1: f64 = 1.0;
25const L2: f64 = 1.0;
26
27pub struct TwoJointArm {
28    pub theta1: f64,
29    pub theta2: f64,
30    pub target_x: f64,
31    pub target_y: f64,
32    pub trajectory: Vec<(f64, f64, f64, f64)>, // (theta1, theta2, end_x, end_y)
33}
34
35impl TwoJointArm {
36    pub fn new() -> Self {
37        TwoJointArm {
38            theta1: 0.0,
39            theta2: 0.0,
40            target_x: 2.0,
41            target_y: 0.0,
42            trajectory: Vec::new(),
43        }
44    }
45
46    pub fn set_target(&mut self, x: f64, y: f64) {
47        self.target_x = x;
48        self.target_y = y;
49    }
50
51    pub fn solve_inverse_kinematics(&mut self, goal_threshold: f64) -> bool {
52        let mut iterations = 0;
53        const MAX_ITERATIONS: usize = 1000;
54
55        self.trajectory.clear();
56
57        while iterations < MAX_ITERATIONS {
58            let (end_x, end_y) = self.forward_kinematics();
59            self.trajectory
60                .push((self.theta1, self.theta2, end_x, end_y));
61
62            // Check if goal is reached
63            let distance_to_goal =
64                ((end_x - self.target_x).powi(2) + (end_y - self.target_y).powi(2)).sqrt();
65            if distance_to_goal < goal_threshold {
66                return true;
67            }
68
69            // Calculate target joint angles using inverse kinematics
70            let (theta1_goal, theta2_goal) = match self.calculate_target_angles() {
71                Some(angles) => angles,
72                None => continue, // Target unreachable, skip this iteration
73            };
74
75            // Update joint angles with proportional control
76            self.theta1 += KP * self.angle_diff(theta1_goal, self.theta1) * DT;
77            self.theta2 += KP * self.angle_diff(theta2_goal, self.theta2) * DT;
78
79            iterations += 1;
80        }
81
82        false // Failed to reach target within max iterations
83    }
84
85    fn calculate_target_angles(&self) -> Option<(f64, f64)> {
86        let x = self.target_x;
87        let y = self.target_y;
88
89        // Check if target is reachable
90        let distance = (x.powi(2) + y.powi(2)).sqrt();
91        if distance > (L1 + L2) {
92            return None; // Target too far
93        }
94
95        // Calculate theta2 using law of cosines
96        let cos_theta2 = (x.powi(2) + y.powi(2) - L1.powi(2) - L2.powi(2)) / (2.0 * L1 * L2);
97
98        // Clamp to valid range to avoid numerical errors
99        let cos_theta2 = cos_theta2.clamp(-1.0, 1.0);
100        let mut theta2_goal = cos_theta2.acos();
101
102        // Calculate theta1
103        let tmp = (L2 * theta2_goal.sin()).atan2(L1 + L2 * theta2_goal.cos());
104        let mut theta1_goal = y.atan2(x) - tmp;
105
106        // Try the other solution if theta1 is negative
107        if theta1_goal < 0.0 {
108            theta2_goal = -theta2_goal;
109            let tmp = (L2 * theta2_goal.sin()).atan2(L1 + L2 * theta2_goal.cos());
110            theta1_goal = y.atan2(x) - tmp;
111        }
112
113        Some((theta1_goal, theta2_goal))
114    }
115
116    pub fn forward_kinematics(&self) -> (f64, f64) {
117        let elbow_x = L1 * self.theta1.cos();
118        let elbow_y = L1 * self.theta1.sin();
119
120        let end_x = elbow_x + L2 * (self.theta1 + self.theta2).cos();
121        let end_y = elbow_y + L2 * (self.theta1 + self.theta2).sin();
122
123        (end_x, end_y)
124    }
125
126    fn angle_diff(&self, theta1: f64, theta2: f64) -> f64 {
127        let mut diff = theta1 - theta2;
128        while diff > PI {
129            diff -= 2.0 * PI;
130        }
131        while diff < -PI {
132            diff += 2.0 * PI;
133        }
134        diff
135    }
136}
137
138// Visualization methods gated behind the "viz" feature
139#[cfg(feature = "viz")]
140impl TwoJointArm {
141    pub fn save_animation_frames(&self, target_name: &str) {
142        let output_dir = format!("img/arm_navigation/{}", target_name);
143        std::fs::create_dir_all(&output_dir).unwrap();
144
145        for (frame, &(theta1, theta2, end_x, end_y)) in self.trajectory.iter().enumerate() {
146            if frame % 5 == 0 {
147                // Save every 5th frame
148                self.save_frame(frame / 5, theta1, theta2, end_x, end_y, &output_dir);
149            }
150        }
151
152        println!("Animation frames saved in {}/frame_*.png", output_dir);
153        println!("Create video with: ffmpeg -r 10 -i {}/frame_%04d.png -c:v libx264 -pix_fmt yuv420p {}_animation.mp4", output_dir, target_name);
154    }
155
156    fn save_frame(
157        &self,
158        frame: usize,
159        theta1: f64,
160        _theta2: f64,
161        end_x: f64,
162        end_y: f64,
163        output_dir: &str,
164    ) {
165        let mut fg = Figure::new();
166        {
167            let axes = fg
168                .axes2d()
169                .set_title(
170                    &format!(
171                        "Two Joint Arm Control - Frame {} (Target: {:.2}, {:.2})",
172                        frame, self.target_x, self.target_y
173                    ),
174                    &[],
175                )
176                .set_x_label("X [m]", &[])
177                .set_y_label("Y [m]", &[])
178                .set_x_range(gnuplot::Fix(-2.5), gnuplot::Fix(2.5))
179                .set_y_range(gnuplot::Fix(-2.5), gnuplot::Fix(2.5))
180                .set_aspect_ratio(gnuplot::Fix(1.0));
181
182            // Calculate joint positions
183            let shoulder = (0.0, 0.0);
184            let elbow = (L1 * theta1.cos(), L1 * theta1.sin());
185            let wrist = (end_x, end_y);
186
187            // Draw arm links (black lines like Python version)
188            axes.lines(
189                &[shoulder.0, elbow.0],
190                &[shoulder.1, elbow.1],
191                &[Color("black")],
192            );
193            axes.lines(&[elbow.0, wrist.0], &[elbow.1, wrist.1], &[Color("black")]);
194
195            // Draw joints (red circles like Python version)
196            axes.points(
197                &[shoulder.0, elbow.0, wrist.0],
198                &[shoulder.1, elbow.1, wrist.1],
199                &[Color("red"), PointSymbol('O')],
200            );
201
202            // Draw target (green star like Python version)
203            axes.points(
204                &[self.target_x],
205                &[self.target_y],
206                &[Color("green"), PointSymbol('*')],
207            );
208
209            // Draw line to target (green dashed line like Python version)
210            axes.lines(
211                &[wrist.0, self.target_x],
212                &[wrist.1, self.target_y],
213                &[Color("green")],
214            );
215        }
216
217        let output_path = format!("{}/frame_{:04}.png", output_dir, frame);
218        fg.set_terminal("pngcairo", &output_path);
219        fg.show().unwrap();
220    }
221
222    /// Save a single arm visualization image (like PythonRobotics)
223    pub fn save_arm_plot(&self, output_path: &str) {
224        let (end_x, end_y) = self.forward_kinematics();
225
226        let mut fg = Figure::new();
227        {
228            let axes = fg
229                .axes2d()
230                .set_title("Two Joint Arm to Point Control", &[])
231                .set_x_label("X [m]", &[])
232                .set_y_label("Y [m]", &[])
233                .set_x_range(gnuplot::Fix(-2.5), gnuplot::Fix(2.5))
234                .set_y_range(gnuplot::Fix(-2.5), gnuplot::Fix(2.5))
235                .set_aspect_ratio(gnuplot::Fix(1.0));
236
237            // Calculate joint positions
238            let shoulder = (0.0, 0.0);
239            let elbow = (L1 * self.theta1.cos(), L1 * self.theta1.sin());
240            let wrist = (end_x, end_y);
241
242            // Draw arm links (black lines like Python version)
243            axes.lines(
244                &[shoulder.0, elbow.0],
245                &[shoulder.1, elbow.1],
246                &[Color("black")],
247            );
248            axes.lines(&[elbow.0, wrist.0], &[elbow.1, wrist.1], &[Color("black")]);
249
250            // Draw joints (red circles like Python version)
251            axes.points(
252                &[shoulder.0, elbow.0, wrist.0],
253                &[shoulder.1, elbow.1, wrist.1],
254                &[Color("red"), PointSymbol('O')],
255            );
256
257            // Draw target (green star like Python version)
258            axes.points(
259                &[self.target_x],
260                &[self.target_y],
261                &[Color("green"), PointSymbol('*')],
262            );
263
264            // Draw line to target (green dashed line like Python version)
265            axes.lines(
266                &[wrist.0, self.target_x],
267                &[wrist.1, self.target_y],
268                &[Color("green")],
269            );
270        }
271
272        fg.set_terminal("pngcairo", output_path);
273        fg.show().unwrap();
274        println!("Arm plot saved to: {}", output_path);
275    }
276
277    pub fn create_summary_plot(&self, target_name: &str) {
278        let mut fg = Figure::new();
279        {
280            let axes = fg
281                .axes2d()
282                .set_title("Two Joint Arm - Joint Angles vs Time", &[])
283                .set_x_label("Iteration", &[])
284                .set_y_label("Angle [rad]", &[]);
285
286            let iterations: Vec<f64> = (0..self.trajectory.len()).map(|i| i as f64).collect();
287            let theta1_data: Vec<f64> = self.trajectory.iter().map(|(t1, _, _, _)| *t1).collect();
288            let theta2_data: Vec<f64> = self.trajectory.iter().map(|(_, t2, _, _)| *t2).collect();
289
290            axes.lines(
291                &iterations,
292                &theta1_data,
293                &[Caption("Theta1 [rad]"), Color("blue")],
294            );
295            axes.lines(
296                &iterations,
297                &theta2_data,
298                &[Caption("Theta2 [rad]"), Color("red")],
299            );
300        }
301
302        let output_path = format!("img/arm_navigation/{}_summary.png", target_name);
303        std::fs::create_dir_all("img/arm_navigation").unwrap();
304        fg.set_terminal("pngcairo", &output_path);
305        fg.show().unwrap();
306        println!("Summary plot saved to: {}", output_path);
307    }
308
309    pub fn run_random_targets_demo(&mut self, num_targets: usize) {
310        use rand::Rng;
311        let mut rng = rand::rng();
312
313        println!("Starting Two Joint Arm random targets demo...");
314
315        for i in 0..num_targets {
316            // Generate random target within workspace
317            let angle = rng.random::<f64>() * 2.0 * PI;
318            let radius = rng.random::<f64>() * (L1 + L2 - 0.1); // Slightly inside workspace
319            let target_x = radius * angle.cos();
320            let target_y = radius * angle.sin();
321
322            self.set_target(target_x, target_y);
323
324            println!("Target {}: ({:.2}, {:.2})", i + 1, target_x, target_y);
325
326            if self.solve_inverse_kinematics(0.01) {
327                println!("  Reached target in {} iterations", self.trajectory.len());
328                let target_name = format!("target_{:02}", i + 1);
329                self.save_animation_frames(&target_name);
330            } else {
331                println!("  Failed to reach target");
332            }
333        }
334
335        self.create_summary_plot("random_demo");
336        println!("Two Joint Arm demo complete!");
337    }
338}
339
340#[cfg(test)]
341mod tests {
342    use super::*;
343
344    #[test]
345    fn test_forward_kinematics_at_origin() {
346        let arm = TwoJointArm::new();
347        let (x, y) = arm.forward_kinematics();
348        assert!((x - 2.0).abs() < 1e-10); // L1 + L2
349        assert!(y.abs() < 1e-10);
350    }
351
352    #[test]
353    fn test_solve_ik_reaches_target() {
354        let mut arm = TwoJointArm::new();
355        arm.set_target(1.5, 0.5);
356        assert!(arm.solve_inverse_kinematics(0.01));
357
358        let (end_x, end_y) = arm.forward_kinematics();
359        let distance = ((end_x - 1.5).powi(2) + (end_y - 0.5).powi(2)).sqrt();
360        assert!(distance < 0.01);
361    }
362}
rust_robotics_control/two_joint_arm_control.rs

rust_robotics_control/
two_joint_arm_control.rs
