1use 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#[derive(Debug, Clone, Copy)]
24pub struct ILQRConfig {
25 pub horizon: usize,
27 pub dt: f64,
29 pub max_iterations: usize,
31 pub tolerance: f64,
33 pub q_pos: f64,
35 pub q_yaw: f64,
37 pub r_v: f64,
39 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#[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
68pub struct ILQRPlanner {
70 config: ILQRConfig,
71}
72
73impl ILQRPlanner {
74 pub fn new(config: ILQRConfig) -> Self {
76 Self { config }
77 }
78
79 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}