lqr/
lib.rs

1use nalgebra as na;
2extern crate rand;
3use core::fmt;
4use na::allocator::Allocator;
5use na::base::DefaultAllocator;
6use na::core::dimension::{Dim, DimMin, DimName};
7use std::error::Error;
8
9#[derive(Clone, Debug)]
10/// Generic LQR optimal feedback controller
11///
12/// Templated via T,S,C which are numeric type, state and control dimensions
13///
14/// Riccatti equation solved via
15/// https://www.tandfonline.com/doi/abs/10.1080/00207170410001714988
16/// https://scicomp.stackexchange.com/questions/30757/discrete-time-algebraic-riccati-equation-dare-solver-in-c
17pub struct LQRController<T, S, C>
18where
19    T: na::RealField,
20    S: Dim + DimName + DimMin<S>, // State dimensions
21    C: Dim + DimName + DimMin<C>, // Control dimensions
22    DefaultAllocator:
23        Allocator<T, S, S> + Allocator<T, C, C> + Allocator<T, S, C> + Allocator<T, C, S>,
24{
25    /// state cost
26    pub q: Option<na::OMatrix<T, S, S>>,
27    /// control cost
28    pub r: Option<na::OMatrix<T, C, C>>,
29    /// optimal gain
30    pub k: Option<na::OMatrix<T, C, S>>,
31
32    /// the controller also has an i-controller for removing steady state error in y dimension
33    /// i-controller coefficient
34    ki: T,
35    /// accumulating integral error for the i-controller
36    integral_error: T,
37}
38
39impl<T, S, C> LQRController<T, S, C>
40where
41    T: na::RealField + Copy,
42    S: Dim + DimName + DimMin<S>,
43    C: Dim + DimName + DimMin<C>,
44    DefaultAllocator:
45        Allocator<T, S, S> + Allocator<T, C, C> + Allocator<T, S, C> + Allocator<T, C, S>,
46    rand::distributions::Standard: rand::distributions::Distribution<T>,
47{
48    // Instantiate controller in default mode without the i-controller
49    pub fn new() -> Result<LQRController<T, S, C>, &'static str> {
50        Ok(LQRController {
51            q: None,
52            r: None,
53            k: None,
54            ki: T::zero(),
55            integral_error: T::zero(),
56        })
57    }
58
59    // Initialise the i-controller with its ki coefficients. This enables the controller
60    pub fn setup_i_controller(&mut self, ki: T) -> Result<(), Box<dyn Error>> {
61        if ki != T::zero() {
62            self.ki = ki;
63        } else {
64            Err("You tried setting ki to 0.0 which is not allowed")?;
65        }
66        Ok(())
67    }
68
69    /// Computes and returns the optimal gain matrix K for the LQR controller
70    ///
71    /// # Arguments
72    ///
73    /// * `a` - state matrix of shape SxS
74    /// * `b` - control matrix of shape SxC
75    /// * `q` - state cost matrix of shape SxS
76    /// * `r` - control cost amtrix of shape CxC
77    /// * `epsilon` - small value to avoid division by 0; 1e-6 works nicely
78    ///
79    /// # Returns
80    /// optimal feedback gain matrix `k` of shape CxS
81    pub fn compute_gain(
82        &mut self,
83        a: &na::OMatrix<T, S, S>,
84        b: &na::OMatrix<T, S, C>,
85        q: &na::OMatrix<T, S, S>,
86        r: &na::OMatrix<T, C, C>,
87        epsilon: T,
88    ) -> Result<na::OMatrix<T, C, S>, &'static str> {
89        let mut a_k = a.clone(); // copy to here
90        let mut g_k = b.clone()
91            * r.clone().try_inverse().expect("Couldn't compute inverse")
92            * b.clone().transpose();
93
94        let mut h_k_1 = q.clone();
95        let mut h_k = na::OMatrix::<T, S, S>::from_fn(|_, _| rand::random());
96
97        while {
98            let error = (h_k_1.clone() - h_k).norm() / h_k_1.norm();
99            h_k = h_k_1.clone();
100            error >= epsilon
101        } {
102            let temp = (na::OMatrix::<T, S, S>::identity() + &g_k * &h_k)
103                .try_inverse()
104                .expect("Couldn't compute inverse");
105            let a_k_1 = &a_k * &temp * &a_k;
106            let g_k_1 = &g_k + &a_k * &temp * &g_k * &a_k.transpose();
107            h_k_1 = &h_k + &a_k.transpose() * &h_k * &temp * &a_k;
108            a_k = a_k_1;
109            g_k = g_k_1;
110        }
111
112        // calculate final gain matrix
113        self.k =
114            Some(r.clone().try_inverse().expect("Couldn't compute inverse") * b.transpose() * h_k);
115        return Ok(self.k.clone().unwrap());
116    }
117
118    /// Returns the optimal feedback control based on the desired and current state vectors.
119    /// This should  be called only after compute_gain() has already been called.
120    ///
121    /// # Arguments
122    ///
123    /// * `current_state` - vector of length S
124    /// * `desired_state` - vector of length S which we want to get to
125    ///
126    /// # Returns
127    /// The feedback control gains. These are insufficient to control anything and have to be
128    /// combined with the feedforward controls. Check examples
129    pub fn compute_optimal_controls(
130        &mut self,
131        current_state: &na::OVector<T, S>,
132        desired_state: &na::OVector<T, S>,
133    ) -> Result<na::OVector<T, C>, &'static str>
134    where
135        T: na::RealField,
136        DefaultAllocator: Allocator<T, S> + Allocator<T, C> + Allocator<T, C, S>,
137    {
138        let error = desired_state - current_state;
139        let y_error = error[1];
140        if y_error.abs() < T::from_f64(1e-2).unwrap() {
141            self.integral_error = T::zero();
142        } else {
143            self.integral_error += y_error * self.ki.clone();
144        }
145
146        let mut controls = &self.k.clone().unwrap() * error;
147        controls[0] -= self.integral_error.clone();
148
149        Ok(controls)
150    }
151}
152
153impl<T, S, C> fmt::Display for LQRController<T, S, C>
154where
155    T: na::RealField,
156    S: Dim + DimName + DimMin<S>,
157    C: Dim + DimName + DimMin<C>,
158    DefaultAllocator:
159        Allocator<T, S, S> + Allocator<T, C, C> + Allocator<T, S, C> + Allocator<T, C, S>,
160    rand::distributions::Standard: rand::distributions::Distribution<T>,
161{
162    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
163        write!(
164            f,
165            "LQR controller
166            Q: {:?}
167            R: {:?}
168            ki: {:}",
169            self.q, self.r, self.ki,
170        )
171    }
172}
173
174pub struct ClosestIndexError;
175
176/// Finds the two closest indices to the current pose along with their respective
177/// Euclidean errors. The idea is that the lower index is behind the car and the upper
178/// index is in front of the car. The opposite is true when driving backwards which
179/// is controlled via the `forwards` parameter
180/// Return format is ((lower_index, lower_error), (upper_index, upper_error)
181pub fn find_closest_indices<T, S, N>(
182    current_state: &na::OVector<T, S>,
183    trajectory: &na::OMatrix<T, S, N>,
184) -> Result<((usize, T), (usize, T)), ClosestIndexError>
185where
186    T: na::RealField + PartialOrd + Copy,
187    S: Dim + DimName + DimMin<S>, // State dimensions
188    N: Dim + DimName,             // Trajectory length
189    DefaultAllocator: Allocator<T, S>
190        + Allocator<T, S, N>
191        + Allocator<T, na::Const<1_usize>, N>
192        + nalgebra::allocator::Allocator<T, N>,
193{
194    if trajectory.ncols() < 2 {
195        return Err(ClosestIndexError);
196    }
197
198    let mut state = Vec::with_capacity(current_state.len());
199    for i in current_state.iter() {
200        state.push(i.clone());
201    }
202    let state = vec![state; trajectory.ncols()].concat();
203
204    let current_state_broadcasted = na::OMatrix::<T, S, N>::from_vec(state);
205    let errors = trajectory - current_state_broadcasted;
206    let errors = errors.map(|x| x.powi(2));
207    let errors = errors.row_sum_tr();
208    let (idx, error) = errors.argmin();
209
210    // figure out lower and upper indices and their errors
211    let (lower, upper) = if idx == 0 {
212        ((0, error.sqrt()), (0, error.sqrt()))
213    } else if idx == errors.len() - 1 {
214        (
215            (errors.len() - 1, error.sqrt()),
216            (errors.len() - 1, error.sqrt()),
217        )
218    } else {
219        if errors[idx + 1] > errors[idx - 1] {
220            ((idx, error.sqrt()), (idx + 1, errors[idx + 1].sqrt()))
221        } else {
222            ((idx - 1, errors[idx - 1].sqrt()), (idx, error.sqrt()))
223        }
224    };
225
226    // Sanity check
227    if lower.0 != 0 && upper.0 != 0 && lower.0 > upper.0 {
228        // Incorrectly detected trajectory indices
229        return Err(ClosestIndexError);
230    }
231
232    Ok((lower, upper))
233}
234
235/// Compute a pose target between two different ones via linear interpolation
236pub fn compute_target<T, S, N>(
237    trajectory: &na::OMatrix<T, S, N>,
238    lower: (usize, T),
239    upper: (usize, T),
240) -> na::OVector<T, S>
241where
242    T: na::RealField + Copy,
243    S: Dim + DimName + DimMin<S>, // State dimensions
244    N: Dim + DimName,             // Trajectory length
245    DefaultAllocator: Allocator<T, S> + Allocator<T, S, N>,
246{
247    trajectory
248        .column(lower.0)
249        .component_mul(&na::OVector::from_element(upper.1 / (lower.1 + upper.1)))
250        + trajectory
251            .column(upper.0)
252            .component_mul(&na::OVector::from_element(lower.1 / (lower.1 + upper.1)))
253}
lqr/lib.rs

lqr/
lib.rs
