rust_robotics_control/
mpc_control.rs

1#![allow(dead_code, clippy::needless_borrows_for_generic_args)]
2
3//
4// Inverted Pendulum MPC control
5// author: Atsushi Sakai
6// Ported to Rust by: rust_robotics team
7//
8
9#[cfg(feature = "viz")]
10use gnuplot::{AxesCommon, Caption, Color, Figure, PointSymbol};
11use nalgebra::{Matrix1, Matrix1x4, Matrix4, Vector4};
12#[cfg(feature = "viz")]
13use std::f64::consts::PI;
14
15// Model parameters
16const L_BAR: f64 = 2.0; // length of bar
17const M: f64 = 1.0; // [kg] cart mass
18const MASS: f64 = 0.3; // [kg] pendulum mass
19const G: f64 = 9.8; // [m/s^2] gravity
20
21const T: usize = 30; // Horizon length
22const DELTA_T: f64 = 0.1; // time tick [s]
23const SIM_TIME: f64 = 5.0; // simulation time [s]
24const OPTIMIZATION_ITERS: usize = 50;
25
26pub struct InvertedPendulumMPC {
27    pub q: Matrix4<f64>,                            // state cost matrix
28    pub r: Matrix1<f64>,                            // input cost matrix
29    pub trajectory: Vec<(f64, Vector4<f64>)>,       // time, state history
30    pub prediction_history: Vec<Vec<Vector4<f64>>>, // predicted trajectories for animation
31}
32
33impl InvertedPendulumMPC {
34    pub fn new() -> Self {
35        let mut q = Matrix4::<f64>::zeros();
36        q[(1, 1)] = 1.0; // velocity cost
37        q[(2, 2)] = 1.0; // angle cost
38
39        let r = Matrix1::<f64>::from_element(0.01); // input cost
40
41        InvertedPendulumMPC {
42            q,
43            r,
44            trajectory: Vec::new(),
45            prediction_history: Vec::new(),
46        }
47    }
48
49    pub fn simulate(&mut self, x0: Vector4<f64>) -> bool {
50        self.simulate_with_runtime(x0, SIM_TIME, OPTIMIZATION_ITERS)
51    }
52
53    fn simulate_with_runtime(
54        &mut self,
55        x0: Vector4<f64>,
56        sim_time: f64,
57        optimization_iters: usize,
58    ) -> bool {
59        let mut x = x0;
60        let mut time = 0.0;
61        let sim_time = sim_time.max(DELTA_T);
62        let optimization_iters = optimization_iters.max(1);
63
64        self.trajectory.clear();
65        self.prediction_history.clear();
66        self.trajectory.push((time, x));
67
68        while time < sim_time {
69            time += DELTA_T;
70
71            // Calculate MPC control input and predicted trajectory
72            let (predicted_states, u) = self.mpc_control_with_iterations(&x, optimization_iters);
73
74            // Store prediction for animation
75            self.prediction_history.push(predicted_states);
76
77            // Simulate inverted pendulum cart
78            x = self.simulation_step(&x, u);
79
80            self.trajectory.push((time, x));
81        }
82
83        println!("MPC simulation finished");
84        println!(
85            "Final state: x={:.2} [m], theta={:.2} [deg]",
86            x[0],
87            x[2].to_degrees()
88        );
89
90        true
91    }
92
93    fn simulation_step(&self, x: &Vector4<f64>, u: Matrix1<f64>) -> Vector4<f64> {
94        let (a, b) = self.get_model_matrix();
95        a * x + b * u[0]
96    }
97
98    fn mpc_control(&self, x0: &Vector4<f64>) -> (Vec<Vector4<f64>>, Matrix1<f64>) {
99        self.mpc_control_with_iterations(x0, OPTIMIZATION_ITERS)
100    }
101
102    fn mpc_control_with_iterations(
103        &self,
104        x0: &Vector4<f64>,
105        _optimization_iters: usize,
106    ) -> (Vec<Vector4<f64>>, Matrix1<f64>) {
107        let (a, b) = self.get_model_matrix();
108        let gains = self.finite_horizon_gains(&a, &b);
109
110        let mut x = *x0;
111        let mut predicted_states = Vec::with_capacity(T + 1);
112        predicted_states.push(x);
113
114        let mut first_u = 0.0;
115        for (step, gain) in gains.iter().enumerate() {
116            let u = -(gain * x)[0];
117            if step == 0 {
118                first_u = u;
119            }
120            x = a * x + b * u;
121            predicted_states.push(x);
122        }
123
124        (predicted_states, Matrix1::from_element(first_u))
125    }
126
127    fn finite_horizon_gains(&self, a: &Matrix4<f64>, b: &Vector4<f64>) -> Vec<Matrix1x4<f64>> {
128        let mut gains = vec![Matrix1x4::<f64>::zeros(); T];
129        let mut p = Matrix4::<f64>::zeros();
130
131        for step in (0..T).rev() {
132            // PythonRobotics penalizes x[t + 1], so use Q + P_{t+1} here.
133            let state_cost = self.q + p;
134            let s = self.r[0] + (b.transpose() * state_cost * b)[0];
135            let gain = (b.transpose() * state_cost * a) / s;
136
137            gains[step] = gain;
138            p = a.transpose() * state_cost * a - (a.transpose() * state_cost * b) * gain;
139        }
140
141        gains
142    }
143
144    fn simulate_prediction(
145        &self,
146        x0: &Vector4<f64>,
147        u_seq: &[f64],
148        a: &Matrix4<f64>,
149        b: &Vector4<f64>,
150    ) -> (Vec<Vector4<f64>>, f64) {
151        let mut x = *x0;
152        let mut states = vec![x];
153        let mut cost = 0.0;
154
155        for &u_t in u_seq.iter().take(T) {
156            // State cost
157            cost += (x.transpose() * self.q * x)[0];
158
159            // Control cost
160            cost += u_t * self.r[0] * u_t;
161
162            // Update state
163            x = a * x + b * u_t;
164            states.push(x);
165        }
166
167        (states, cost)
168    }
169
170    fn get_model_matrix(&self) -> (Matrix4<f64>, Vector4<f64>) {
171        let mut a = Matrix4::<f64>::new(
172            0.0,
173            1.0,
174            0.0,
175            0.0,
176            0.0,
177            0.0,
178            MASS * G / M,
179            0.0,
180            0.0,
181            0.0,
182            0.0,
183            1.0,
184            0.0,
185            0.0,
186            G * (M + MASS) / (L_BAR * M),
187            0.0,
188        );
189        a = Matrix4::<f64>::identity() + DELTA_T * a;
190
191        let b = Vector4::new(0.0, 1.0 / M, 0.0, 1.0 / (L_BAR * M)) * DELTA_T;
192
193        (a, b)
194    }
195}
196
197// Visualization methods gated behind the "viz" feature
198#[cfg(feature = "viz")]
199impl InvertedPendulumMPC {
200    fn save_animation_frame(
201        &self,
202        frame: usize,
203        time: f64,
204        current_state: &Vector4<f64>,
205        predicted_states: &[Vector4<f64>],
206    ) {
207        let mut fg = Figure::new();
208        {
209            let axes = fg
210                .axes2d()
211                .set_title(
212                    &format!(
213                        "Inverted Pendulum MPC Control - Frame {} (t={:.1}s)",
214                        frame, time
215                    ),
216                    &[],
217                )
218                .set_x_label("Position [m]", &[])
219                .set_y_label("Height [m]", &[])
220                .set_x_range(gnuplot::Fix(-6.0), gnuplot::Fix(3.0))
221                .set_y_range(gnuplot::Fix(-0.5), gnuplot::Fix(3.0))
222                .set_aspect_ratio(gnuplot::Fix(1.0));
223
224            // Draw cart and pendulum
225            self.draw_cart_pendulum(axes, current_state[0], current_state[2]);
226
227            // Draw predicted trajectory
228            if predicted_states.len() > 1 {
229                let pred_x: Vec<f64> = predicted_states.iter().map(|s| s[0]).collect();
230                // Draw predicted cart positions
231                axes.points(
232                    &pred_x,
233                    vec![0.2; pred_x.len()],
234                    &[Caption("Predicted Path"), Color("red"), PointSymbol('.')],
235                );
236
237                // Draw predicted pendulum tips
238                let pred_tip_x: Vec<f64> = predicted_states
239                    .iter()
240                    .map(|s| s[0] + L_BAR * s[2].sin())
241                    .collect();
242                let pred_tip_y: Vec<f64> = predicted_states
243                    .iter()
244                    .map(|s| 0.5 + L_BAR * s[2].cos())
245                    .collect();
246
247                axes.points(
248                    &pred_tip_x,
249                    &pred_tip_y,
250                    &[
251                        Caption("Predicted Pendulum"),
252                        Color("orange"),
253                        PointSymbol('x'),
254                    ],
255                );
256            }
257        }
258
259        let output_path = format!("img/inverted_pendulum/mpc/mpc_frame_{:04}.png", frame);
260        std::fs::create_dir_all("img/inverted_pendulum/mpc").unwrap();
261        fg.set_terminal("pngcairo", &output_path);
262        fg.show().unwrap();
263    }
264
265    fn draw_cart_pendulum(&self, axes: &mut gnuplot::Axes2D, x: f64, theta: f64) {
266        let cart_w = 1.0;
267        let cart_h = 0.5;
268        let radius = 0.1;
269
270        // Cart body
271        let cx = vec![
272            x - cart_w / 2.0,
273            x + cart_w / 2.0,
274            x + cart_w / 2.0,
275            x - cart_w / 2.0,
276            x - cart_w / 2.0,
277        ];
278        let cy = vec![
279            radius * 2.0,
280            radius * 2.0,
281            cart_h + radius * 2.0,
282            cart_h + radius * 2.0,
283            radius * 2.0,
284        ];
285        axes.lines(&cx, &cy, &[Caption("Cart"), Color("blue")]);
286
287        // Pendulum rod
288        let bx = vec![x, x + L_BAR * theta.sin()];
289        let by = vec![
290            cart_h + radius * 2.0,
291            cart_h + radius * 2.0 + L_BAR * theta.cos(),
292        ];
293        axes.lines(&bx, &by, &[Caption("Pendulum"), Color("black")]);
294
295        // Wheels
296        let angles: Vec<f64> = (0..120).map(|i| i as f64 * PI / 60.0).collect();
297        let wheel_x: Vec<f64> = angles
298            .iter()
299            .map(|a| x - cart_w / 4.0 + radius * a.cos())
300            .collect();
301        let wheel_y: Vec<f64> = angles.iter().map(|a| radius + radius * a.sin()).collect();
302        axes.lines(&wheel_x, &wheel_y, &[Color("black")]);
303
304        let wheel_x2: Vec<f64> = angles
305            .iter()
306            .map(|a| x + cart_w / 4.0 + radius * a.cos())
307            .collect();
308        axes.lines(&wheel_x2, &wheel_y, &[Color("black")]);
309
310        // Pendulum mass
311        let mass_x: Vec<f64> = angles.iter().map(|a| bx[1] + radius * a.cos()).collect();
312        let mass_y: Vec<f64> = angles.iter().map(|a| by[1] + radius * a.sin()).collect();
313        axes.lines(&mass_x, &mass_y, &[Color("red")]);
314    }
315
316    pub fn create_summary_plot(&self) {
317        let mut fg = Figure::new();
318        {
319            let axes = fg
320                .axes2d()
321                .set_title("Inverted Pendulum MPC Control - Summary", &[])
322                .set_x_label("Time [s]", &[])
323                .set_y_label("State", &[]);
324
325            let time: Vec<f64> = self.trajectory.iter().map(|(t, _)| *t).collect();
326            let position: Vec<f64> = self.trajectory.iter().map(|(_, x)| x[0]).collect();
327            let angle: Vec<f64> = self
328                .trajectory
329                .iter()
330                .map(|(_, x)| x[2].to_degrees())
331                .collect();
332
333            axes.lines(&time, &position, &[Caption("Position [m]"), Color("blue")]);
334            axes.lines(&time, &angle, &[Caption("Angle [deg]"), Color("red")]);
335        }
336
337        let output_path = "img/inverted_pendulum/mpc/mpc_summary.png";
338        fg.set_terminal("pngcairo", output_path);
339        fg.show().unwrap();
340        println!("MPC summary plot saved to: {}", output_path);
341    }
342}
343
344impl Default for InvertedPendulumMPC {
345    fn default() -> Self {
346        Self::new()
347    }
348}
349
350#[cfg(test)]
351mod tests {
352    use super::*;
353
354    fn assert_vector4_close(actual: &Vector4<f64>, expected: [f64; 4], tol: f64) {
355        for (i, expected_value) in expected.iter().enumerate() {
356            assert!(
357                (actual[i] - expected_value).abs() <= tol,
358                "index {}: actual={} expected={} tol={}",
359                i,
360                actual[i],
361                expected_value,
362                tol
363            );
364        }
365    }
366
367    #[test]
368    fn test_mpc_first_control_matches_reference() {
369        let controller = InvertedPendulumMPC::new();
370        let x0 = Vector4::new(0.0, 0.0, 0.3, 0.0);
371        let (_predicted_states, u) = controller.mpc_control(&x0);
372        let expected_first_u = -21.202_780_205_052_658;
373
374        assert!((u[0] - expected_first_u).abs() <= 1e-9);
375    }
376
377    #[test]
378    fn test_mpc_simulation_completes() {
379        let mut controller = InvertedPendulumMPC::new();
380        let x0 = Vector4::new(0.0, 0.0, 0.3, 0.0);
381        assert!(controller.simulate_with_runtime(x0, 0.6, 4));
382        assert!(!controller.trajectory.is_empty());
383
384        let final_state = controller.trajectory.last().unwrap().1;
385        assert_vector4_close(
386            &final_state,
387            [
388                -1.037_901_493_867_440_5,
389                -1.383_752_006_586_625_8,
390                -0.041_149_711_402_044_535,
391                -0.213_983_985_795_986_85,
392            ],
393            1e-9,
394        );
395    }
396
397    #[test]
398    #[ignore = "long-running regression scenario"]
399    fn test_mpc_simulation_completes_full_runtime() {
400        let mut controller = InvertedPendulumMPC::new();
401        let x0 = Vector4::new(0.0, 0.0, 0.3, 0.0);
402        assert!(controller.simulate(x0));
403        assert!(!controller.trajectory.is_empty());
404
405        let final_state = controller.trajectory.last().unwrap().1;
406        assert_vector4_close(
407            &final_state,
408            [
409                -1.830_471_874_707_321_4,
410                -0.000_333_385_633_250_359_55,
411                -0.000_136_734_168_845_974_78,
412                0.000_230_095_328_469_559_24,
413            ],
414            1e-4,
415        );
416    }
417}
rust_robotics_control/mpc_control.rs

rust_robotics_control/
mpc_control.rs
