rust_robotics_control/
racing_mppi_quadrotor.rs

1//! Reference-free racing MPPI with a full quadrotor attitude model.
2//!
3//! This deepens [`crate::racing_mppi_3d`] from a point-mass drone to a quadrotor
4//! whose orientation matters. The control input is the standard low-level agile-
5//! racing abstraction used by differential-flatness controllers: a mass-
6//! normalized collective thrust along the body z-axis plus three body rates.
7//!
8//! - [`QuadrotorState`] carries position, velocity, and a unit-quaternion
9//!   attitude.
10//! - [`QuadrotorParams`] integrates the rigid-body translational dynamics
11//!   (`a = thrust * body_z - gravity - drag * v`) together with quaternion
12//!   attitude kinematics driven by the commanded body rates.
13//! - [`QuadrotorMppiController`] samples thrust/body-rate perturbations around a
14//!   hover nominal, rolls them through the attitude dynamics, and scores the
15//!   resulting positions with the reference-free gate-progress objective from
16//!   [`crate::racing_mppi_3d`]. Because horizontal motion can only come from
17//!   tilting the thrust vector, the position objective drives the *attitude*:
18//!   the drone learns to pitch and roll toward the next gate.
19//! - [`simulate_quadrotor_race`] flies a deterministic, seeded controller around
20//!   a lap and reports [`QuadrotorLapReport`], which adds attitude metrics (tilt
21//!   angle and body-rate effort) to the lap-progress metrics.
22//!
23//! The gate geometry ([`RacingGatePlane3D`], [`RacingGateLap3D`]) is reused
24//! directly from [`crate::racing_mppi_3d`].
25
26use crate::racing_mppi_3d::{RacingGateLap3D, RacingGatePlane3D};
27use rand::{rngs::StdRng, SeedableRng};
28use rand_distr::{Distribution, Normal};
29use rust_robotics_core::{RoboticsError, RoboticsResult};
30
31/// State of the quadrotor: position \[m\], velocity \[m/s\], and a unit
32/// quaternion attitude `[w, x, y, z]` (body-to-world rotation).
33#[derive(Debug, Clone, Copy, PartialEq)]
34pub struct QuadrotorState {
35    pub position: [f64; 3],
36    pub velocity: [f64; 3],
37    pub attitude: [f64; 4],
38}
39
40impl QuadrotorState {
41    /// Level hover state at the given position (identity attitude, zero
42    /// velocity).
43    pub fn at(x: f64, y: f64, z: f64) -> Self {
44        Self {
45            position: [x, y, z],
46            velocity: [0.0, 0.0, 0.0],
47            attitude: [1.0, 0.0, 0.0, 0.0],
48        }
49    }
50
51    pub fn speed(self) -> f64 {
52        norm(self.velocity)
53    }
54
55    /// The body z-axis (thrust direction) expressed in the world frame.
56    pub fn thrust_axis(self) -> [f64; 3] {
57        rotate(self.attitude, [0.0, 0.0, 1.0])
58    }
59
60    /// Tilt angle \[rad\] between the thrust axis and world up, in `[0, pi]`.
61    pub fn tilt_angle(self) -> f64 {
62        self.thrust_axis()[2].clamp(-1.0, 1.0).acos()
63    }
64}
65
66/// Commanded collective thrust and body rates for the quadrotor.
67#[derive(Debug, Clone, Copy, PartialEq)]
68pub struct QuadrotorControl {
69    /// Mass-normalized collective thrust \[m/s^2\] along the body z-axis.
70    pub thrust: f64,
71    /// Body angular rates `[wx, wy, wz]` \[rad/s\].
72    pub body_rates: [f64; 3],
73}
74
75impl QuadrotorControl {
76    pub fn new(thrust: f64, body_rates: [f64; 3]) -> Self {
77        Self { thrust, body_rates }
78    }
79
80    /// Hover command: thrust balancing gravity, zero body rates.
81    pub fn hover(gravity: f64) -> Self {
82        Self {
83            thrust: gravity,
84            body_rates: [0.0, 0.0, 0.0],
85        }
86    }
87
88    /// Body-rate magnitude \[rad/s\].
89    pub fn body_rate_magnitude(self) -> f64 {
90        norm(self.body_rates)
91    }
92}
93
94/// Quadrotor rigid-body parameters and actuation limits.
95///
96/// Translational dynamics are semi-implicit: the collective thrust along the
97/// current body z-axis, minus gravity and linear drag, updates the velocity,
98/// the speed is capped, and the position integrates the new velocity. Attitude
99/// is integrated from the commanded body rates and renormalized each step.
100#[derive(Debug, Clone, Copy, PartialEq)]
101pub struct QuadrotorParams {
102    /// Downward gravitational acceleration \[m/s^2\].
103    pub gravity: f64,
104    /// Linear aerodynamic drag coefficient \[1/s\].
105    pub drag: f64,
106    /// Minimum mass-normalized collective thrust \[m/s^2\].
107    pub min_thrust: f64,
108    /// Maximum mass-normalized collective thrust \[m/s^2\].
109    pub max_thrust: f64,
110    /// Maximum body-rate magnitude \[rad/s\].
111    pub max_body_rate: f64,
112    /// Maximum speed \[m/s\] the body can reach.
113    pub max_speed: f64,
114}
115
116impl Default for QuadrotorParams {
117    fn default() -> Self {
118        Self {
119            gravity: 9.81,
120            drag: 0.3,
121            min_thrust: 0.0,
122            max_thrust: 24.0,
123            max_body_rate: 6.0,
124            max_speed: 7.0,
125        }
126    }
127}
128
129impl QuadrotorParams {
130    pub fn new(
131        gravity: f64,
132        drag: f64,
133        min_thrust: f64,
134        max_thrust: f64,
135        max_body_rate: f64,
136        max_speed: f64,
137    ) -> RoboticsResult<Self> {
138        if !gravity.is_finite() || gravity < 0.0 {
139            return Err(RoboticsError::InvalidParameter(
140                "quadrotor gravity must be finite and non-negative".to_string(),
141            ));
142        }
143        if !drag.is_finite() || drag < 0.0 {
144            return Err(RoboticsError::InvalidParameter(
145                "quadrotor drag must be finite and non-negative".to_string(),
146            ));
147        }
148        if !min_thrust.is_finite() || min_thrust < 0.0 {
149            return Err(RoboticsError::InvalidParameter(
150                "quadrotor min_thrust must be finite and non-negative".to_string(),
151            ));
152        }
153        if !max_thrust.is_finite() || max_thrust <= min_thrust {
154            return Err(RoboticsError::InvalidParameter(
155                "quadrotor max_thrust must be finite and exceed min_thrust".to_string(),
156            ));
157        }
158        if !max_body_rate.is_finite() || max_body_rate <= 0.0 {
159            return Err(RoboticsError::InvalidParameter(
160                "quadrotor max_body_rate must be finite and positive".to_string(),
161            ));
162        }
163        if !max_speed.is_finite() || max_speed <= 0.0 {
164            return Err(RoboticsError::InvalidParameter(
165                "quadrotor max_speed must be finite and positive".to_string(),
166            ));
167        }
168        Ok(Self {
169            gravity,
170            drag,
171            min_thrust,
172            max_thrust,
173            max_body_rate,
174            max_speed,
175        })
176    }
177
178    /// Clamp a command to the thrust and body-rate limits.
179    pub fn saturate(self, control: QuadrotorControl) -> QuadrotorControl {
180        let thrust = control.thrust.clamp(self.min_thrust, self.max_thrust);
181        let mut rates = control.body_rates;
182        let magnitude = norm(rates);
183        if magnitude > self.max_body_rate && magnitude > 0.0 {
184            let scale = self.max_body_rate / magnitude;
185            rates = [rates[0] * scale, rates[1] * scale, rates[2] * scale];
186        }
187        QuadrotorControl::new(thrust, rates)
188    }
189
190    /// Advance the quadrotor one step under the commanded thrust and body rates.
191    pub fn step(self, state: QuadrotorState, control: QuadrotorControl, dt: f64) -> QuadrotorState {
192        let control = self.saturate(control);
193
194        // Attitude kinematics: q_dot = 0.5 * q (x) [0, wx, wy, wz].
195        let attitude = normalize_quat(integrate_attitude(state.attitude, control.body_rates, dt));
196
197        // Translational dynamics in the world frame.
198        let thrust_axis = rotate(attitude, [0.0, 0.0, 1.0]);
199        let accel = [
200            control.thrust * thrust_axis[0] - self.drag * state.velocity[0],
201            control.thrust * thrust_axis[1] - self.drag * state.velocity[1],
202            control.thrust * thrust_axis[2] - self.gravity - self.drag * state.velocity[2],
203        ];
204        let mut velocity = [
205            state.velocity[0] + accel[0] * dt,
206            state.velocity[1] + accel[1] * dt,
207            state.velocity[2] + accel[2] * dt,
208        ];
209        let speed = norm(velocity);
210        if speed > self.max_speed && speed > 0.0 {
211            let scale = self.max_speed / speed;
212            velocity = [
213                velocity[0] * scale,
214                velocity[1] * scale,
215                velocity[2] * scale,
216            ];
217        }
218        let position = [
219            state.position[0] + velocity[0] * dt,
220            state.position[1] + velocity[1] * dt,
221            state.position[2] + velocity[2] * dt,
222        ];
223
224        QuadrotorState {
225            position,
226            velocity,
227            attitude,
228        }
229    }
230}
231
232/// Lap-progress and attitude metrics from a quadrotor racing rollout.
233#[derive(Debug, Clone, PartialEq)]
234pub struct QuadrotorLapReport {
235    /// Control steps executed.
236    pub steps: usize,
237    /// Total gate passes across the whole rollout (counts repeated laps).
238    pub gates_passed: usize,
239    /// Whole laps completed.
240    pub laps_completed: usize,
241    /// Fraction of the current (partial) lap completed, in \[0, 1).
242    pub lap_fraction: f64,
243    /// Mean speed \[m/s\] over the rollout.
244    pub mean_speed: f64,
245    /// Peak speed \[m/s\] over the rollout.
246    pub max_speed: f64,
247    /// Executed path length \[m\].
248    pub path_length: f64,
249    /// Distance \[m\] to the active gate at the end of the rollout.
250    pub gate_distance: f64,
251    /// Smallest normalized aperture margin at any gate crossing (1 = center,
252    /// 0 = edge); `INFINITY` if no gate was crossed.
253    pub min_aperture_margin: f64,
254    /// Mean tilt angle \[rad\] of the thrust axis from world up.
255    pub mean_tilt: f64,
256    /// Peak tilt angle \[rad\] of the thrust axis from world up.
257    pub max_tilt: f64,
258    /// Mean commanded body-rate magnitude \[rad/s\].
259    pub mean_body_rate: f64,
260    /// Time \[s\] of the first completed lap, or `None`.
261    pub first_lap_time: Option<f64>,
262    /// Executed positions, including the start.
263    pub path: Vec<[f64; 3]>,
264}
265
266/// Configuration for the deterministic quadrotor racing MPPI controller.
267#[derive(Debug, Clone, PartialEq)]
268pub struct QuadrotorMppiConfig {
269    pub horizon: usize,
270    pub samples: usize,
271    pub dt: f64,
272    /// Std-dev of thrust perturbations \[m/s^2\].
273    pub thrust_sigma: f64,
274    /// Std-dev of body-rate perturbations \[rad/s\].
275    pub rate_sigma: f64,
276    /// Weight on the squared speed regularizer.
277    pub velocity_weight: f64,
278    /// Weight on the squared body-rate regularizer.
279    pub rate_weight: f64,
280    /// Weight penalizing tilt (keeps the drone from flipping under noise).
281    pub level_weight: f64,
282    pub lambda: f64,
283    pub seed: u64,
284}
285
286impl Default for QuadrotorMppiConfig {
287    fn default() -> Self {
288        Self {
289            horizon: 28,
290            samples: 900,
291            dt: 0.05,
292            thrust_sigma: 5.0,
293            rate_sigma: 5.0,
294            velocity_weight: 0.01,
295            rate_weight: 0.01,
296            level_weight: 0.6,
297            lambda: 1.0,
298            seed: 7,
299        }
300    }
301}
302
303fn validate_config(config: &QuadrotorMppiConfig) -> RoboticsResult<()> {
304    if config.horizon == 0 {
305        return Err(RoboticsError::InvalidParameter(
306            "quadrotor MPPI horizon must be positive".to_string(),
307        ));
308    }
309    if config.samples == 0 {
310        return Err(RoboticsError::InvalidParameter(
311            "quadrotor MPPI samples must be positive".to_string(),
312        ));
313    }
314    if !config.dt.is_finite() || config.dt <= 0.0 {
315        return Err(RoboticsError::InvalidParameter(
316            "quadrotor MPPI dt must be finite and positive".to_string(),
317        ));
318    }
319    if !config.thrust_sigma.is_finite() || config.thrust_sigma <= 0.0 {
320        return Err(RoboticsError::InvalidParameter(
321            "quadrotor MPPI thrust_sigma must be finite and positive".to_string(),
322        ));
323    }
324    if !config.rate_sigma.is_finite() || config.rate_sigma <= 0.0 {
325        return Err(RoboticsError::InvalidParameter(
326            "quadrotor MPPI rate_sigma must be finite and positive".to_string(),
327        ));
328    }
329    if !config.lambda.is_finite() || config.lambda <= 0.0 {
330        return Err(RoboticsError::InvalidParameter(
331            "quadrotor MPPI lambda must be finite and positive".to_string(),
332        ));
333    }
334    if config.velocity_weight < 0.0 || config.rate_weight < 0.0 || config.level_weight < 0.0 {
335        return Err(RoboticsError::InvalidParameter(
336            "quadrotor MPPI weights must be non-negative".to_string(),
337        ));
338    }
339    Ok(())
340}
341
342/// Result of one `plan` call.
343#[derive(Debug, Clone, PartialEq)]
344pub struct QuadrotorMppiPlan {
345    pub control: QuadrotorControl,
346    /// Lowest rollout cost among the samples (diagnostic).
347    pub best_cost: f64,
348    /// Normalized effective sample size in \[0, 1\] (diagnostic).
349    pub normalized_effective_sample_size: f64,
350}
351
352/// Deterministic, seeded MPPI controller for quadrotor gate racing.
353#[derive(Debug, Clone)]
354pub struct QuadrotorMppiController {
355    config: QuadrotorMppiConfig,
356    params: QuadrotorParams,
357    nominal: Vec<QuadrotorControl>,
358    rng: StdRng,
359}
360
361impl QuadrotorMppiController {
362    pub fn new(config: QuadrotorMppiConfig, params: QuadrotorParams) -> RoboticsResult<Self> {
363        validate_config(&config)?;
364        let nominal = vec![QuadrotorControl::hover(params.gravity); config.horizon];
365        let rng = StdRng::seed_from_u64(config.seed);
366        Ok(Self {
367            config,
368            params,
369            nominal,
370            rng,
371        })
372    }
373
374    pub fn config(&self) -> &QuadrotorMppiConfig {
375        &self.config
376    }
377
378    fn rollout(
379        &self,
380        start: QuadrotorState,
381        controls: &[QuadrotorControl],
382    ) -> (Vec<[f64; 3]>, f64) {
383        let mut state = start;
384        let mut positions = Vec::with_capacity(controls.len() + 1);
385        positions.push(state.position);
386        let mut regularizer = 0.0;
387        for &control in controls {
388            state = self.params.step(state, control, self.config.dt);
389            let speed = state.speed();
390            let rate = control.body_rate_magnitude();
391            // 1 - cos(tilt) is 0 when level and grows as the drone tips over.
392            let tilt_cost = 1.0 - state.thrust_axis()[2];
393            regularizer += self.config.velocity_weight * speed * speed
394                + self.config.rate_weight * rate * rate
395                + self.config.level_weight * tilt_cost;
396            positions.push(state.position);
397        }
398        (positions, regularizer)
399    }
400
401    /// Plan the next control given the current state, lap, and pass count.
402    pub fn plan(
403        &mut self,
404        start: QuadrotorState,
405        lap: &RacingGateLap3D,
406        passed: usize,
407    ) -> RoboticsResult<QuadrotorMppiPlan> {
408        let thrust_noise = Normal::new(0.0, self.config.thrust_sigma).map_err(|_| {
409            RoboticsError::InvalidParameter("invalid quadrotor thrust distribution".to_string())
410        })?;
411        let rate_noise = Normal::new(0.0, self.config.rate_sigma).map_err(|_| {
412            RoboticsError::InvalidParameter("invalid quadrotor rate distribution".to_string())
413        })?;
414
415        let mut costs = Vec::with_capacity(self.config.samples);
416        let mut sequences = Vec::with_capacity(self.config.samples);
417        let mut best_cost = f64::INFINITY;
418
419        for _ in 0..self.config.samples {
420            let mut controls = Vec::with_capacity(self.config.horizon);
421            for &base in &self.nominal {
422                let noisy = QuadrotorControl::new(
423                    base.thrust + thrust_noise.sample(&mut self.rng),
424                    [
425                        base.body_rates[0] + rate_noise.sample(&mut self.rng),
426                        base.body_rates[1] + rate_noise.sample(&mut self.rng),
427                        base.body_rates[2] + rate_noise.sample(&mut self.rng),
428                    ],
429                );
430                controls.push(self.params.saturate(noisy));
431            }
432            let (positions, regularizer) = self.rollout(start, &controls);
433            let (gate_cost, _) = lap.score_positions(&positions, passed);
434            let cost = gate_cost + regularizer;
435            best_cost = best_cost.min(cost);
436            costs.push(cost);
437            sequences.push(controls);
438        }
439
440        let min_cost = costs.iter().copied().fold(f64::INFINITY, f64::min);
441        let mut weights = Vec::with_capacity(costs.len());
442        let mut weight_sum = 0.0;
443        for &cost in &costs {
444            let weight = (-(cost - min_cost) / self.config.lambda).exp();
445            weight_sum += weight;
446            weights.push(weight);
447        }
448
449        if weight_sum <= 0.0 || !weight_sum.is_finite() {
450            let best_index = costs
451                .iter()
452                .enumerate()
453                .min_by(|a, b| a.1.partial_cmp(b.1).unwrap_or(std::cmp::Ordering::Equal))
454                .map(|(index, _)| index)
455                .unwrap_or(0);
456            let control = sequences[best_index][0];
457            return Ok(QuadrotorMppiPlan {
458                control,
459                best_cost,
460                normalized_effective_sample_size: 0.0,
461            });
462        }
463
464        let mut updated = vec![QuadrotorControl::new(0.0, [0.0, 0.0, 0.0]); self.config.horizon];
465        let mut sum_sq = 0.0;
466        for (weight, controls) in weights.iter().zip(&sequences) {
467            let normalized = weight / weight_sum;
468            sum_sq += normalized * normalized;
469            for (acc, control) in updated.iter_mut().zip(controls) {
470                acc.thrust += normalized * control.thrust;
471                acc.body_rates[0] += normalized * control.body_rates[0];
472                acc.body_rates[1] += normalized * control.body_rates[1];
473                acc.body_rates[2] += normalized * control.body_rates[2];
474            }
475        }
476        for control in &mut updated {
477            *control = self.params.saturate(*control);
478        }
479
480        let first_control = updated[0];
481        // Warm-start: shift the nominal forward, repeat the last command.
482        self.nominal.clear();
483        self.nominal.extend_from_slice(&updated[1..]);
484        self.nominal.push(*updated.last().unwrap());
485
486        let normalized_effective_sample_size = if sum_sq > 0.0 {
487            1.0 / (sum_sq * self.config.samples as f64)
488        } else {
489            0.0
490        };
491
492        Ok(QuadrotorMppiPlan {
493            control: first_control,
494            best_cost,
495            normalized_effective_sample_size,
496        })
497    }
498}
499
500/// Drive the quadrotor racing MPPI controller around a lap and report
501/// lap-progress and attitude metrics.
502///
503/// The rollout runs for at most `max_steps` control steps and stops early once
504/// `target_laps` whole laps have been completed. Everything is deterministic for
505/// a fixed `config.seed`.
506pub fn simulate_quadrotor_race(
507    config: QuadrotorMppiConfig,
508    params: QuadrotorParams,
509    lap: &RacingGateLap3D,
510    start: QuadrotorState,
511    max_steps: usize,
512    target_laps: usize,
513) -> RoboticsResult<QuadrotorLapReport> {
514    let dt = config.dt;
515    let gate_count = lap.gate_count();
516    let mut controller = QuadrotorMppiController::new(config, params)?;
517
518    let mut state = start;
519    let mut passed = 0usize;
520    let mut path = vec![state.position];
521    let mut sum_speed = 0.0;
522    let mut max_speed = 0.0_f64;
523    let mut sum_tilt = 0.0;
524    let mut max_tilt = 0.0_f64;
525    let mut sum_body_rate = 0.0;
526    let mut path_length = 0.0;
527    let mut min_aperture_margin = f64::INFINITY;
528    let mut first_lap_time = None;
529    let mut executed_steps = 0usize;
530
531    for step in 0..max_steps {
532        let plan = controller.plan(state, lap, passed)?;
533        let next = params.step(state, plan.control, dt);
534        let from = state.position;
535        let to = next.position;
536
537        loop {
538            let gate = lap.gate_for_pass_count(passed);
539            if !gate.crosses(from, to, lap.crossing_tolerance) {
540                break;
541            }
542            let margin = aperture_crossing_margin(gate, from, to);
543            min_aperture_margin = min_aperture_margin.min(margin);
544            passed += 1;
545            if passed == gate_count && first_lap_time.is_none() {
546                first_lap_time = Some((step + 1) as f64 * dt);
547            }
548            if !lap.closed && passed >= gate_count {
549                break;
550            }
551        }
552
553        path_length += distance(from, to);
554        let speed = next.speed();
555        sum_speed += speed;
556        max_speed = max_speed.max(speed);
557        let tilt = next.tilt_angle();
558        sum_tilt += tilt;
559        max_tilt = max_tilt.max(tilt);
560        sum_body_rate += plan.control.body_rate_magnitude();
561        state = next;
562        path.push(to);
563        executed_steps += 1;
564
565        if target_laps > 0 && passed >= target_laps * gate_count {
566            break;
567        }
568        if !lap.closed && passed >= gate_count {
569            break;
570        }
571    }
572
573    let laps_completed = passed / gate_count;
574    let lap_fraction = (passed % gate_count) as f64 / gate_count as f64;
575    let active_gate = lap.gate_for_pass_count(passed);
576    let gate_distance = distance(active_gate.center(), state.position);
577    let denom = executed_steps.max(1) as f64;
578
579    Ok(QuadrotorLapReport {
580        steps: executed_steps,
581        gates_passed: passed,
582        laps_completed,
583        lap_fraction,
584        mean_speed: sum_speed / denom,
585        max_speed,
586        path_length,
587        gate_distance,
588        min_aperture_margin,
589        mean_tilt: sum_tilt / denom,
590        max_tilt,
591        mean_body_rate: sum_body_rate / denom,
592        first_lap_time,
593        path,
594    })
595}
596
597fn aperture_crossing_margin(gate: &RacingGatePlane3D, from: [f64; 3], to: [f64; 3]) -> f64 {
598    let from_signed = gate.signed_distance(from);
599    let to_signed = gate.signed_distance(to);
600    let denom = to_signed - from_signed;
601    let t = if denom.abs() > 1e-9 {
602        (-from_signed / denom).clamp(0.0, 1.0)
603    } else {
604        0.0
605    };
606    let crossing = [
607        from[0] + t * (to[0] - from[0]),
608        from[1] + t * (to[1] - from[1]),
609        from[2] + t * (to[2] - from[2]),
610    ];
611    gate.aperture_margin(crossing)
612}
613
614// --- Quaternion and vector helpers (kept private to avoid a math dependency).
615
616/// Rotate a vector by a unit quaternion `[w, x, y, z]` (body to world).
617fn rotate(q: [f64; 4], v: [f64; 3]) -> [f64; 3] {
618    let [w, x, y, z] = q;
619    // v + 2 * qv x (qv x v + w v), with qv = (x, y, z).
620    let qv = [x, y, z];
621    let t = scale(cross(qv, v), 2.0);
622    let cross_t = cross(qv, t);
623    [
624        v[0] + w * t[0] + cross_t[0],
625        v[1] + w * t[1] + cross_t[1],
626        v[2] + w * t[2] + cross_t[2],
627    ]
628}
629
630/// Integrate attitude one step from body rates: q + 0.5 * (q (x) [0, w]) * dt.
631fn integrate_attitude(q: [f64; 4], rates: [f64; 3], dt: f64) -> [f64; 4] {
632    let [qw, qx, qy, qz] = q;
633    let [wx, wy, wz] = rates;
634    // Quaternion product q (x) [0, wx, wy, wz].
635    let dw = -(qx * wx + qy * wy + qz * wz);
636    let dx = qw * wx + qy * wz - qz * wy;
637    let dy = qw * wy + qz * wx - qx * wz;
638    let dz = qw * wz + qx * wy - qy * wx;
639    let half_dt = 0.5 * dt;
640    [
641        qw + half_dt * dw,
642        qx + half_dt * dx,
643        qy + half_dt * dy,
644        qz + half_dt * dz,
645    ]
646}
647
648fn normalize_quat(q: [f64; 4]) -> [f64; 4] {
649    let n = (q[0] * q[0] + q[1] * q[1] + q[2] * q[2] + q[3] * q[3]).sqrt();
650    if !n.is_finite() || n <= 0.0 {
651        return [1.0, 0.0, 0.0, 0.0];
652    }
653    [q[0] / n, q[1] / n, q[2] / n, q[3] / n]
654}
655
656fn cross(a: [f64; 3], b: [f64; 3]) -> [f64; 3] {
657    [
658        a[1] * b[2] - a[2] * b[1],
659        a[2] * b[0] - a[0] * b[2],
660        a[0] * b[1] - a[1] * b[0],
661    ]
662}
663
664fn scale(a: [f64; 3], s: f64) -> [f64; 3] {
665    [a[0] * s, a[1] * s, a[2] * s]
666}
667
668fn norm(a: [f64; 3]) -> f64 {
669    (a[0] * a[0] + a[1] * a[1] + a[2] * a[2]).sqrt()
670}
671
672fn distance(a: [f64; 3], b: [f64; 3]) -> f64 {
673    let d = [a[0] - b[0], a[1] - b[1], a[2] - b[2]];
674    norm(d)
675}
676
677#[cfg(test)]
678mod tests {
679    use super::*;
680
681    fn gate(center: [f64; 3], normal: [f64; 3]) -> RacingGatePlane3D {
682        RacingGatePlane3D::new(center, normal, [0.0, 0.0, 1.0], 0.8, 0.8).unwrap()
683    }
684
685    fn straight_lap() -> RacingGateLap3D {
686        let gates = vec![
687            gate([1.5, 0.0, 0.0], [1.0, 0.0, 0.0]),
688            gate([3.0, 0.0, 0.0], [1.0, 0.0, 0.0]),
689            gate([4.5, 0.0, 0.0], [1.0, 0.0, 0.0]),
690        ];
691        RacingGateLap3D::open(gates).unwrap()
692    }
693
694    #[test]
695    fn hover_holds_position() {
696        let params = QuadrotorParams::default();
697        let mut state = QuadrotorState::at(0.0, 0.0, 1.0);
698        for _ in 0..20 {
699            state = params.step(state, QuadrotorControl::hover(params.gravity), 0.05);
700        }
701        // Thrust balances gravity, so it stays put.
702        assert!(state.speed() < 1e-6);
703        assert!((state.position[2] - 1.0).abs() < 1e-6);
704        assert!(state.tilt_angle() < 1e-9);
705    }
706
707    #[test]
708    fn body_rate_tilts_then_thrust_moves_horizontally() {
709        let params = QuadrotorParams::default();
710        let mut state = QuadrotorState::at(0.0, 0.0, 0.0);
711        // Pitch about the body y-axis to tilt the thrust vector toward +x.
712        for _ in 0..6 {
713            state = params.step(
714                state,
715                QuadrotorControl::new(params.gravity, [0.0, 1.0, 0.0]),
716                0.05,
717            );
718        }
719        assert!(state.tilt_angle() > 0.1, "tilt {}", state.tilt_angle());
720        // Now hold attitude and apply extra thrust; horizontal velocity grows.
721        for _ in 0..6 {
722            state = params.step(
723                state,
724                QuadrotorControl::new(params.gravity + 4.0, [0.0, 0.0, 0.0]),
725                0.05,
726            );
727        }
728        assert!(state.velocity[0].abs() > 0.1, "vx {}", state.velocity[0]);
729    }
730
731    #[test]
732    fn attitude_stays_normalized() {
733        let params = QuadrotorParams::default();
734        let mut state = QuadrotorState::at(0.0, 0.0, 0.0);
735        for _ in 0..40 {
736            state = params.step(state, QuadrotorControl::new(11.0, [2.0, -1.5, 0.8]), 0.05);
737        }
738        let q = state.attitude;
739        let n = (q[0] * q[0] + q[1] * q[1] + q[2] * q[2] + q[3] * q[3]).sqrt();
740        assert!((n - 1.0).abs() < 1e-9, "quat norm {n}");
741    }
742
743    #[test]
744    fn saturation_clamps_thrust_and_rates() {
745        let params = QuadrotorParams::default();
746        let saturated = params.saturate(QuadrotorControl::new(1000.0, [100.0, 0.0, 0.0]));
747        assert!((saturated.thrust - params.max_thrust).abs() < 1e-9);
748        assert!((saturated.body_rate_magnitude() - params.max_body_rate).abs() < 1e-9);
749    }
750
751    #[test]
752    fn controller_makes_gate_progress() {
753        let lap = straight_lap();
754        let report = simulate_quadrotor_race(
755            QuadrotorMppiConfig::default(),
756            QuadrotorParams::default(),
757            &lap,
758            QuadrotorState::at(0.0, 0.0, 0.0),
759            120,
760            1,
761        )
762        .unwrap();
763        assert!(
764            report.gates_passed >= 1,
765            "expected at least one gate, got {}",
766            report.gates_passed
767        );
768        // Flying through gates requires tilting the thrust vector.
769        assert!(report.max_tilt > 0.05, "max tilt {}", report.max_tilt);
770        assert!(report.path_length > 0.0);
771    }
772
773    #[test]
774    fn simulation_is_deterministic() {
775        let lap = straight_lap();
776        let run = || {
777            simulate_quadrotor_race(
778                QuadrotorMppiConfig::default(),
779                QuadrotorParams::default(),
780                &lap,
781                QuadrotorState::at(0.0, 0.0, 0.0),
782                80,
783                1,
784            )
785            .unwrap()
786        };
787        let first = run();
788        let second = run();
789        assert_eq!(first.path, second.path);
790        assert_eq!(first.gates_passed, second.gates_passed);
791    }
792
793    #[test]
794    fn invalid_config_is_rejected() {
795        let config = QuadrotorMppiConfig {
796            horizon: 0,
797            ..QuadrotorMppiConfig::default()
798        };
799        assert!(QuadrotorMppiController::new(config, QuadrotorParams::default()).is_err());
800    }
801
802    #[test]
803    fn invalid_params_are_rejected() {
804        assert!(QuadrotorParams::new(9.81, 0.3, 5.0, 4.0, 6.0, 7.0).is_err());
805    }
806}
rust_robotics_control/racing_mppi_quadrotor.rs

rust_robotics_control/
racing_mppi_quadrotor.rs
