Skip to main content

rust_robotics_control/
ilqr.rs

1//! Iterative Linear Quadratic Regulator (iLQR) for a unicycle model.
2//!
3//! The planner optimizes a finite-horizon state and control trajectory to
4//! drive a robot from a start pose to a goal pose.
5//!
6//! Reference:
7//! - Weiwei Li, Emanuel Todorov, "Iterative Linear Quadratic Regulator
8//!   Design for Nonlinear Biological Movement Systems":
9//!   <https://homes.cs.washington.edu/~todorov/papers/LiICINCO04.pdf>
10
11use nalgebra::{Matrix2, Matrix2x3, Matrix3, Matrix3x2, Vector2, Vector3};
12use std::f64::consts::PI;
13
14const MAX_LINEAR_SPEED: f64 = 5.0;
15const MAX_ANGULAR_SPEED: f64 = PI;
16const TERMINAL_WEIGHT_SCALE: f64 = 10.0;
17const REGULARIZATION: f64 = 1.0e-6;
18
19type FeedforwardTerms = Vec<Vector2<f64>>;
20type FeedbackGains = Vec<Matrix2x3<f64>>;
21
22/// Configuration for the iLQR planner.
23#[derive(Debug, Clone, Copy)]
24pub struct ILQRConfig {
25    /// Number of time steps in the optimization horizon.
26    pub horizon: usize,
27    /// Time step \[s\].
28    pub dt: f64,
29    /// Maximum iLQR iterations.
30    pub max_iterations: usize,
31    /// Convergence threshold on cost change.
32    pub tolerance: f64,
33    /// Position cost weight.
34    pub q_pos: f64,
35    /// Heading cost weight.
36    pub q_yaw: f64,
37    /// Linear velocity input cost weight.
38    pub r_v: f64,
39    /// Angular velocity input cost weight.
40    pub r_omega: f64,
41}
42
43impl Default for ILQRConfig {
44    fn default() -> Self {
45        Self {
46            horizon: 50,
47            dt: 0.1,
48            max_iterations: 50,
49            tolerance: 1.0e-6,
50            q_pos: 1.0,
51            q_yaw: 0.1,
52            r_v: 0.01,
53            r_omega: 0.01,
54        }
55    }
56}
57
58/// Optimized iLQR trajectory.
59#[derive(Debug, Clone)]
60pub struct ILQRResult {
61    pub states: Vec<Vector3<f64>>,
62    pub controls: Vec<Vector2<f64>>,
63    pub cost: f64,
64    pub iterations: usize,
65    pub converged: bool,
66}
67
68/// iLQR planner for a unicycle model.
69pub struct ILQRPlanner {
70    config: ILQRConfig,
71}
72
73impl ILQRPlanner {
74    /// Creates a new iLQR planner.
75    pub fn new(config: ILQRConfig) -> Self {
76        Self { config }
77    }
78
79    /// Plans a trajectory from `start` to `goal`.
80    pub fn plan(&self, start: Vector3<f64>, goal: Vector3<f64>) -> ILQRResult {
81        let mut controls = self.initial_controls(start, goal);
82        let mut states = self.rollout(start, &controls);
83        let mut cost = self.total_cost(&states, &controls, goal);
84        let mut iterations = 0;
85        let mut converged = false;
86
87        for iter in 0..self.config.max_iterations {
88            iterations = iter + 1;
89            let Some((feedforward, feedback)) = self.backward_pass(&states, &controls, goal) else {
90                break;
91            };
92
93            let mut accepted = None;
94            for alpha in [1.0, 0.5, 0.25, 0.1, 0.05, 0.01] {
95                let (candidate_states, candidate_controls) =
96                    self.forward_pass(start, &states, &controls, &feedforward, &feedback, alpha);
97                let candidate_cost = self.total_cost(&candidate_states, &candidate_controls, goal);
98
99                if candidate_cost < cost {
100                    accepted = Some((candidate_states, candidate_controls, candidate_cost));
101                    break;
102                }
103            }
104
105            let Some((candidate_states, candidate_controls, candidate_cost)) = accepted else {
106                break;
107            };
108
109            let cost_change = (cost - candidate_cost).abs();
110            states = candidate_states;
111            controls = candidate_controls;
112            cost = candidate_cost;
113
114            if cost_change < self.config.tolerance {
115                converged = true;
116                break;
117            }
118        }
119
120        if !converged {
121            let final_state = states.last().copied().unwrap_or(start);
122            let pos_error =
123                ((goal[0] - final_state[0]).powi(2) + (goal[1] - final_state[1]).powi(2)).sqrt();
124            let yaw_error = normalize_angle(goal[2] - final_state[2]).abs();
125            converged = pos_error < 0.25 && yaw_error < 0.1;
126        }
127
128        ILQRResult {
129            states,
130            controls,
131            cost,
132            iterations,
133            converged,
134        }
135    }
136
137    fn initial_controls(&self, start: Vector3<f64>, goal: Vector3<f64>) -> Vec<Vector2<f64>> {
138        let distance = ((goal[0] - start[0]).powi(2) + (goal[1] - start[1]).powi(2)).sqrt();
139        let total_time = (self.config.horizon as f64 * self.config.dt).max(self.config.dt);
140        let v = (distance / total_time).clamp(-MAX_LINEAR_SPEED, MAX_LINEAR_SPEED);
141        let omega = (normalize_angle(goal[2] - start[2]) / total_time)
142            .clamp(-MAX_ANGULAR_SPEED, MAX_ANGULAR_SPEED);
143        vec![Vector2::new(v, omega); self.config.horizon]
144    }
145
146    fn rollout(&self, start: Vector3<f64>, controls: &[Vector2<f64>]) -> Vec<Vector3<f64>> {
147        let mut states = Vec::with_capacity(controls.len() + 1);
148        let mut state = start;
149        states.push(state);
150
151        for control in controls {
152            state = dynamics(&state, control, self.config.dt);
153            states.push(state);
154        }
155
156        states
157    }
158
159    fn total_cost(
160        &self,
161        states: &[Vector3<f64>],
162        controls: &[Vector2<f64>],
163        goal: Vector3<f64>,
164    ) -> f64 {
165        let stage_cost: f64 = states
166            .iter()
167            .take(controls.len())
168            .zip(controls.iter())
169            .map(|(state, control)| self.stage_cost(state, control, goal))
170            .sum();
171        stage_cost + self.terminal_cost(states.last().expect("trajectory is non-empty"), goal)
172    }
173
174    fn stage_cost(&self, state: &Vector3<f64>, control: &Vector2<f64>, goal: Vector3<f64>) -> f64 {
175        let error = state_error(state, goal);
176        self.config.q_pos * (error[0].powi(2) + error[1].powi(2))
177            + self.config.q_yaw * error[2].powi(2)
178            + self.config.r_v * control[0].powi(2)
179            + self.config.r_omega * control[1].powi(2)
180    }
181
182    fn terminal_cost(&self, state: &Vector3<f64>, goal: Vector3<f64>) -> f64 {
183        let error = state_error(state, goal);
184        TERMINAL_WEIGHT_SCALE
185            * (self.config.q_pos * (error[0].powi(2) + error[1].powi(2))
186                + self.config.q_yaw * error[2].powi(2))
187    }
188
189    fn backward_pass(
190        &self,
191        states: &[Vector3<f64>],
192        controls: &[Vector2<f64>],
193        goal: Vector3<f64>,
194    ) -> Option<(FeedforwardTerms, FeedbackGains)> {
195        let horizon = controls.len();
196        let mut feedforward = vec![Vector2::zeros(); horizon];
197        let mut feedback = vec![Matrix2x3::<f64>::zeros(); horizon];
198
199        let (mut value_x, mut value_xx) =
200            self.terminal_cost_derivatives(states.last().expect("trajectory is non-empty"), goal);
201
202        for t in (0..horizon).rev() {
203            let (a, b) = linearize_dynamics(&states[t], &controls[t], self.config.dt);
204            let (l_x, l_u, l_xx, l_uu) =
205                self.stage_cost_derivatives(&states[t], &controls[t], goal);
206
207            let q_x = l_x + a.transpose() * value_x;
208            let q_u = l_u + b.transpose() * value_x;
209            let q_xx = l_xx + a.transpose() * value_xx * a;
210            let mut q_uu = l_uu + b.transpose() * value_xx * b;
211            let q_ux = b.transpose() * value_xx * a;
212
213            q_uu = symmetrize_2x2(q_uu) + REGULARIZATION * Matrix2::identity();
214            let q_uu_inv = q_uu.try_inverse()?;
215
216            let k = -q_uu_inv * q_u;
217            let k_feedback = -q_uu_inv * q_ux;
218            feedforward[t] = k;
219            feedback[t] = k_feedback;
220
221            value_x = q_x
222                + k_feedback.transpose() * q_uu * k
223                + k_feedback.transpose() * q_u
224                + q_ux.transpose() * k;
225            value_xx = q_xx
226                + k_feedback.transpose() * q_uu * k_feedback
227                + k_feedback.transpose() * q_ux
228                + q_ux.transpose() * k_feedback;
229            value_xx = symmetrize_3x3(value_xx);
230        }
231
232        Some((feedforward, feedback))
233    }
234
235    fn forward_pass(
236        &self,
237        start: Vector3<f64>,
238        nominal_states: &[Vector3<f64>],
239        nominal_controls: &[Vector2<f64>],
240        feedforward: &[Vector2<f64>],
241        feedback: &[Matrix2x3<f64>],
242        alpha: f64,
243    ) -> (Vec<Vector3<f64>>, Vec<Vector2<f64>>) {
244        let mut states = Vec::with_capacity(nominal_controls.len() + 1);
245        let mut controls = Vec::with_capacity(nominal_controls.len());
246        states.push(start);
247
248        for t in 0..nominal_controls.len() {
249            let delta_state = state_error(&states[t], nominal_states[t]);
250            let delta_control = alpha * feedforward[t] + feedback[t] * delta_state;
251            let mut control = nominal_controls[t] + delta_control;
252            control[0] = control[0].clamp(-MAX_LINEAR_SPEED, MAX_LINEAR_SPEED);
253            control[1] = control[1].clamp(-MAX_ANGULAR_SPEED, MAX_ANGULAR_SPEED);
254            controls.push(control);
255
256            let next_state = dynamics(&states[t], &control, self.config.dt);
257            states.push(next_state);
258        }
259
260        (states, controls)
261    }
262
263    fn stage_cost_derivatives(
264        &self,
265        state: &Vector3<f64>,
266        control: &Vector2<f64>,
267        goal: Vector3<f64>,
268    ) -> (Vector3<f64>, Vector2<f64>, Matrix3<f64>, Matrix2<f64>) {
269        let error = state_error(state, goal);
270        let l_x = Vector3::new(
271            2.0 * self.config.q_pos * error[0],
272            2.0 * self.config.q_pos * error[1],
273            2.0 * self.config.q_yaw * error[2],
274        );
275        let l_u = Vector2::new(
276            2.0 * self.config.r_v * control[0],
277            2.0 * self.config.r_omega * control[1],
278        );
279        let l_xx = Matrix3::from_diagonal(&Vector3::new(
280            2.0 * self.config.q_pos,
281            2.0 * self.config.q_pos,
282            2.0 * self.config.q_yaw,
283        ));
284        let l_uu = Matrix2::from_diagonal(&Vector2::new(
285            2.0 * self.config.r_v,
286            2.0 * self.config.r_omega,
287        ));
288
289        (l_x, l_u, l_xx, l_uu)
290    }
291
292    fn terminal_cost_derivatives(
293        &self,
294        state: &Vector3<f64>,
295        goal: Vector3<f64>,
296    ) -> (Vector3<f64>, Matrix3<f64>) {
297        let error = state_error(state, goal);
298        let v_x = Vector3::new(
299            2.0 * TERMINAL_WEIGHT_SCALE * self.config.q_pos * error[0],
300            2.0 * TERMINAL_WEIGHT_SCALE * self.config.q_pos * error[1],
301            2.0 * TERMINAL_WEIGHT_SCALE * self.config.q_yaw * error[2],
302        );
303        let v_xx = Matrix3::from_diagonal(&Vector3::new(
304            2.0 * TERMINAL_WEIGHT_SCALE * self.config.q_pos,
305            2.0 * TERMINAL_WEIGHT_SCALE * self.config.q_pos,
306            2.0 * TERMINAL_WEIGHT_SCALE * self.config.q_yaw,
307        ));
308
309        (v_x, v_xx)
310    }
311}
312
313fn state_error(state: &Vector3<f64>, goal: Vector3<f64>) -> Vector3<f64> {
314    Vector3::new(
315        state[0] - goal[0],
316        state[1] - goal[1],
317        normalize_angle(state[2] - goal[2]),
318    )
319}
320
321fn symmetrize_2x2(matrix: Matrix2<f64>) -> Matrix2<f64> {
322    0.5 * (matrix + matrix.transpose())
323}
324
325fn symmetrize_3x3(matrix: Matrix3<f64>) -> Matrix3<f64> {
326    0.5 * (matrix + matrix.transpose())
327}
328
329fn normalize_angle(mut angle: f64) -> f64 {
330    while angle > PI {
331        angle -= 2.0 * PI;
332    }
333    while angle < -PI {
334        angle += 2.0 * PI;
335    }
336    angle
337}
338
339fn dynamics(state: &Vector3<f64>, control: &Vector2<f64>, dt: f64) -> Vector3<f64> {
340    let x = state[0] + control[0] * state[2].cos() * dt;
341    let y = state[1] + control[0] * state[2].sin() * dt;
342    let yaw = normalize_angle(state[2] + control[1] * dt);
343    Vector3::new(x, y, yaw)
344}
345
346fn linearize_dynamics(
347    state: &Vector3<f64>,
348    control: &Vector2<f64>,
349    dt: f64,
350) -> (Matrix3<f64>, Matrix3x2<f64>) {
351    let yaw = state[2];
352    let v = control[0];
353
354    let a = Matrix3::new(
355        1.0,
356        0.0,
357        -dt * v * yaw.sin(),
358        0.0,
359        1.0,
360        dt * v * yaw.cos(),
361        0.0,
362        0.0,
363        1.0,
364    );
365    let b = Matrix3x2::new(dt * yaw.cos(), 0.0, dt * yaw.sin(), 0.0, 0.0, dt);
366
367    (a, b)
368}
369
370#[cfg(test)]
371mod tests {
372    use super::*;
373
374    fn assert_close(actual: f64, expected: f64, tolerance: f64) {
375        assert!(
376            (actual - expected).abs() <= tolerance,
377            "expected {expected:.12}, got {actual:.12}, tolerance {tolerance:.3e}"
378        );
379    }
380
381    #[test]
382    fn test_ilqr_config_defaults() {
383        let config = ILQRConfig::default();
384        assert_eq!(config.horizon, 50);
385        assert_eq!(config.dt, 0.1);
386        assert_eq!(config.max_iterations, 50);
387        assert_eq!(config.tolerance, 1.0e-6);
388        assert_eq!(config.q_pos, 1.0);
389        assert_eq!(config.q_yaw, 0.1);
390        assert_eq!(config.r_v, 0.01);
391        assert_eq!(config.r_omega, 0.01);
392    }
393
394    #[test]
395    fn test_ilqr_straight_line_converges() {
396        let planner = ILQRPlanner::new(ILQRConfig {
397            horizon: 60,
398            ..ILQRConfig::default()
399        });
400        let start = Vector3::new(0.0, 0.0, 0.0);
401        let goal = Vector3::new(5.0, 0.0, 0.0);
402
403        let result = planner.plan(start, goal);
404        let final_state = result.states.last().expect("trajectory").clone_owned();
405
406        assert!(result.converged);
407        assert!((final_state[0] - goal[0]).abs() < 0.25);
408        assert!(final_state[1].abs() < 0.1);
409        assert!(normalize_angle(final_state[2] - goal[2]).abs() < 0.1);
410    }
411
412    #[test]
413    fn test_ilqr_turn_converges() {
414        let planner = ILQRPlanner::new(ILQRConfig {
415            horizon: 70,
416            ..ILQRConfig::default()
417        });
418        let start = Vector3::new(0.0, 0.0, 0.0);
419        let goal = Vector3::new(3.0, 2.0, PI / 2.0);
420
421        let result = planner.plan(start, goal);
422        let final_state = result.states.last().expect("trajectory").clone_owned();
423        let pos_error =
424            ((final_state[0] - goal[0]).powi(2) + (final_state[1] - goal[1]).powi(2)).sqrt();
425
426        assert!(result.converged);
427        assert!(pos_error < 0.3);
428        assert!(normalize_angle(final_state[2] - goal[2]).abs() < 0.15);
429    }
430
431    #[test]
432    fn test_ilqr_cost_decreases_over_iterations() {
433        let planner = ILQRPlanner::new(ILQRConfig::default());
434        let start = Vector3::new(0.0, 0.0, 0.0);
435        let goal = Vector3::new(4.0, 1.0, 0.2);
436
437        let initial_controls = planner.initial_controls(start, goal);
438        let initial_states = planner.rollout(start, &initial_controls);
439        let initial_cost = planner.total_cost(&initial_states, &initial_controls, goal);
440
441        let result = planner.plan(start, goal);
442
443        assert!(result.cost < initial_cost);
444    }
445
446    #[test]
447    fn test_ilqr_turning_case_regression() {
448        let planner = ILQRPlanner::new(ILQRConfig {
449            horizon: 70,
450            ..ILQRConfig::default()
451        });
452        let start = Vector3::new(0.0, 0.0, 0.0);
453        let goal = Vector3::new(3.0, 2.0, PI / 2.0);
454
455        let result = planner.plan(start, goal);
456        let final_state = result.states.last().expect("trajectory").clone_owned();
457
458        assert!(result.converged);
459        assert_eq!(result.iterations, 10);
460        assert_close(result.cost, 43.962457904378, 1.0e-6);
461        assert_close(final_state[0], 2.991948883770, 1.0e-6);
462        assert_close(final_state[1], 2.000000028451, 1.0e-6);
463        assert_close(final_state[2], 1.570796059231, 1.0e-6);
464        assert_close(result.controls[0][0], 5.0, 1.0e-12);
465        assert_close(result.controls[0][1], PI, 1.0e-12);
466        assert_close(result.controls[1][0], 5.0, 1.0e-12);
467        assert_close(result.controls[1][1], PI, 1.0e-12);
468    }
469}