rust_robotics_control/
racing_mppi_3d.rs

1//! Reference-free racing MPPI through 3-D gate planes.
2//!
3//! This extends the 2-D gate-progress objective in [`crate::mppi`] to a 3-D
4//! drone-racing setting:
5//!
6//! - [`RacingGatePlane3D`] is a rectangular gate aperture in 3-D defined by a
7//!   center, a race-direction normal, and an in-plane up axis. A rollout passes
8//!   the gate when a segment crosses the gate plane from behind to in front and
9//!   the crossing point lies inside the rectangular aperture.
10//! - [`RacingDroneDynamics3D`] is a point-mass drone with linear aerodynamic
11//!   drag, optional gravity, a speed cap, and an acceleration-magnitude cap —
12//!   richer than the pure 2-D double integrator.
13//! - [`RacingGateLap3D`] arranges gates into an open course or a closed lap and
14//!   scores reference-free gate progress.
15//! - [`simulate_lap_race`] drives a deterministic, seeded MPPI controller around
16//!   the lap and reports [`RacingLapReport3D`] lap-progress metrics.
17
18use rand::{rngs::StdRng, SeedableRng};
19use rand_distr::{Distribution, Normal};
20use rust_robotics_core::{RoboticsError, RoboticsResult};
21
22/// State for a 3-D point-mass drone: position \[m\] and velocity \[m/s\].
23#[derive(Debug, Clone, Copy, PartialEq)]
24pub struct RacingDroneState3D {
25    pub x: f64,
26    pub y: f64,
27    pub z: f64,
28    pub vx: f64,
29    pub vy: f64,
30    pub vz: f64,
31}
32
33impl RacingDroneState3D {
34    pub fn new(x: f64, y: f64, z: f64, vx: f64, vy: f64, vz: f64) -> Self {
35        Self {
36            x,
37            y,
38            z,
39            vx,
40            vy,
41            vz,
42        }
43    }
44
45    /// Drone at rest at the given position.
46    pub fn at(x: f64, y: f64, z: f64) -> Self {
47        Self::new(x, y, z, 0.0, 0.0, 0.0)
48    }
49
50    pub fn position(self) -> [f64; 3] {
51        [self.x, self.y, self.z]
52    }
53
54    pub fn velocity(self) -> [f64; 3] {
55        [self.vx, self.vy, self.vz]
56    }
57
58    pub fn speed(self) -> f64 {
59        (self.vx * self.vx + self.vy * self.vy + self.vz * self.vz).sqrt()
60    }
61}
62
63/// Commanded acceleration \[m/s^2\] for the drone model.
64#[derive(Debug, Clone, Copy, PartialEq)]
65pub struct RacingDroneControl3D {
66    pub ax: f64,
67    pub ay: f64,
68    pub az: f64,
69}
70
71impl RacingDroneControl3D {
72    pub fn new(ax: f64, ay: f64, az: f64) -> Self {
73        Self { ax, ay, az }
74    }
75
76    pub fn zero() -> Self {
77        Self::new(0.0, 0.0, 0.0)
78    }
79
80    pub fn magnitude(self) -> f64 {
81        (self.ax * self.ax + self.ay * self.ay + self.az * self.az).sqrt()
82    }
83}
84
85/// Point-mass drone dynamics with linear drag, gravity, and actuation limits.
86///
87/// The discrete update is semi-implicit: the commanded acceleration (capped to
88/// `accel_limit` in magnitude) and gravity update the velocity, linear drag
89/// decays it, the speed is capped to `max_speed`, and the position integrates
90/// the new velocity. This stays stable for the step sizes used by the racing
91/// MPPI rollouts while modeling drag and saturation effects a pure double
92/// integrator ignores.
93#[derive(Debug, Clone, Copy, PartialEq)]
94pub struct RacingDroneDynamics3D {
95    /// Linear drag coefficient \[1/s\].
96    pub drag: f64,
97    /// Downward gravitational acceleration \[m/s^2\] applied to `vz`.
98    pub gravity: f64,
99    /// Maximum speed \[m/s\] the drone body can reach.
100    pub max_speed: f64,
101    /// Maximum commanded acceleration magnitude \[m/s^2\].
102    pub accel_limit: f64,
103}
104
105impl Default for RacingDroneDynamics3D {
106    fn default() -> Self {
107        Self {
108            drag: 0.4,
109            gravity: 0.0,
110            max_speed: 5.5,
111            accel_limit: 9.0,
112        }
113    }
114}
115
116impl RacingDroneDynamics3D {
117    pub fn new(drag: f64, gravity: f64, max_speed: f64, accel_limit: f64) -> RoboticsResult<Self> {
118        if !drag.is_finite() || drag < 0.0 {
119            return Err(RoboticsError::InvalidParameter(
120                "drone drag must be finite and non-negative".to_string(),
121            ));
122        }
123        if !gravity.is_finite() || gravity < 0.0 {
124            return Err(RoboticsError::InvalidParameter(
125                "drone gravity must be finite and non-negative".to_string(),
126            ));
127        }
128        if !max_speed.is_finite() || max_speed <= 0.0 {
129            return Err(RoboticsError::InvalidParameter(
130                "drone max_speed must be finite and positive".to_string(),
131            ));
132        }
133        if !accel_limit.is_finite() || accel_limit <= 0.0 {
134            return Err(RoboticsError::InvalidParameter(
135                "drone accel_limit must be finite and positive".to_string(),
136            ));
137        }
138        Ok(Self {
139            drag,
140            gravity,
141            max_speed,
142            accel_limit,
143        })
144    }
145
146    /// Clamp a control's acceleration magnitude to `accel_limit`.
147    pub fn saturate(self, control: RacingDroneControl3D) -> RacingDroneControl3D {
148        let magnitude = control.magnitude();
149        if magnitude <= self.accel_limit || magnitude == 0.0 {
150            return control;
151        }
152        let scale = self.accel_limit / magnitude;
153        RacingDroneControl3D::new(control.ax * scale, control.ay * scale, control.az * scale)
154    }
155
156    /// Advance the drone one step under the commanded acceleration.
157    pub fn step(
158        self,
159        state: RacingDroneState3D,
160        control: RacingDroneControl3D,
161        dt: f64,
162    ) -> RacingDroneState3D {
163        let control = self.saturate(control);
164        let mut vx = state.vx + control.ax * dt;
165        let mut vy = state.vy + control.ay * dt;
166        let mut vz = state.vz + (control.az - self.gravity) * dt;
167
168        let decay = (1.0 - self.drag * dt).clamp(0.0, 1.0);
169        vx *= decay;
170        vy *= decay;
171        vz *= decay;
172
173        let speed = (vx * vx + vy * vy + vz * vz).sqrt();
174        if speed > self.max_speed && speed > 0.0 {
175            let scale = self.max_speed / speed;
176            vx *= scale;
177            vy *= scale;
178            vz *= scale;
179        }
180
181        RacingDroneState3D {
182            x: state.x + vx * dt,
183            y: state.y + vy * dt,
184            z: state.z + vz * dt,
185            vx,
186            vy,
187            vz,
188        }
189    }
190}
191
192/// A rectangular racing gate plane in 3-D.
193///
194/// The gate is defined by its `center`, a unit `normal` pointing in the race
195/// direction, and an in-plane `up` axis. The `right` axis is derived as
196/// `normal x up`. A point is inside the aperture when its in-plane offset along
197/// `right`/`up` is within `half_width`/`half_height` respectively.
198#[derive(Debug, Clone, Copy, PartialEq)]
199pub struct RacingGatePlane3D {
200    center: [f64; 3],
201    normal: [f64; 3],
202    up: [f64; 3],
203    right: [f64; 3],
204    half_width: f64,
205    half_height: f64,
206}
207
208impl RacingGatePlane3D {
209    /// Build a gate from a center, race-direction normal, and up hint.
210    ///
211    /// `up_hint` need not be perpendicular to the normal or unit length; it is
212    /// orthonormalized against the normal. The width axis is `normal x up`.
213    pub fn new(
214        center: [f64; 3],
215        normal: [f64; 3],
216        up_hint: [f64; 3],
217        half_width: f64,
218        half_height: f64,
219    ) -> RoboticsResult<Self> {
220        if !center.iter().all(|v| v.is_finite()) {
221            return Err(RoboticsError::InvalidParameter(
222                "racing gate center must be finite".to_string(),
223            ));
224        }
225        let normal = normalize(normal).ok_or_else(|| {
226            RoboticsError::InvalidParameter(
227                "racing gate normal must be finite and non-zero".to_string(),
228            )
229        })?;
230        // Remove the normal component from the up hint, then normalize.
231        let up_projected = sub(up_hint, scale(normal, dot(up_hint, normal)));
232        let up = normalize(up_projected).ok_or_else(|| {
233            RoboticsError::InvalidParameter(
234                "racing gate up hint must not be parallel to the normal".to_string(),
235            )
236        })?;
237        let right = normalize(cross(normal, up)).ok_or_else(|| {
238            RoboticsError::InvalidParameter(
239                "racing gate axes must be linearly independent".to_string(),
240            )
241        })?;
242        if !half_width.is_finite() || half_width <= 0.0 {
243            return Err(RoboticsError::InvalidParameter(
244                "racing gate half_width must be finite and positive".to_string(),
245            ));
246        }
247        if !half_height.is_finite() || half_height <= 0.0 {
248            return Err(RoboticsError::InvalidParameter(
249                "racing gate half_height must be finite and positive".to_string(),
250            ));
251        }
252        Ok(Self {
253            center,
254            normal,
255            up,
256            right,
257            half_width,
258            half_height,
259        })
260    }
261
262    pub fn center(&self) -> [f64; 3] {
263        self.center
264    }
265
266    pub fn normal(&self) -> [f64; 3] {
267        self.normal
268    }
269
270    pub fn half_width(&self) -> f64 {
271        self.half_width
272    }
273
274    pub fn half_height(&self) -> f64 {
275        self.half_height
276    }
277
278    /// Signed distance to the gate plane; positive is in front (race direction).
279    pub fn signed_distance(&self, point: [f64; 3]) -> f64 {
280        dot(sub(point, self.center), self.normal)
281    }
282
283    /// In-plane offsets `(right, up)` of a point relative to the gate center.
284    pub fn plane_offsets(&self, point: [f64; 3]) -> (f64, f64) {
285        let delta = sub(point, self.center);
286        (dot(delta, self.right), dot(delta, self.up))
287    }
288
289    pub fn squared_distance(&self, point: [f64; 3]) -> f64 {
290        let delta = sub(point, self.center);
291        dot(delta, delta)
292    }
293
294    /// Whether a point projects inside the rectangular aperture (with tolerance).
295    pub fn within_aperture(&self, point: [f64; 3], tolerance: f64) -> bool {
296        let (right_off, up_off) = self.plane_offsets(point);
297        right_off.abs() <= self.half_width + tolerance
298            && up_off.abs() <= self.half_height + tolerance
299    }
300
301    /// Whether the segment `from -> to` passes through the gate aperture.
302    ///
303    /// The segment must cross the plane from behind (negative side) to in front
304    /// (non-negative side), and the plane-crossing point must lie inside the
305    /// aperture.
306    pub fn crosses(&self, from: [f64; 3], to: [f64; 3], tolerance: f64) -> bool {
307        let from_signed = self.signed_distance(from);
308        let to_signed = self.signed_distance(to);
309        if from_signed > tolerance || to_signed < -tolerance {
310            return false;
311        }
312        let denom = to_signed - from_signed;
313        let t = if denom.abs() > 1e-9 {
314            (-from_signed / denom).clamp(0.0, 1.0)
315        } else {
316            0.0
317        };
318        let crossing = [
319            from[0] + t * (to[0] - from[0]),
320            from[1] + t * (to[1] - from[1]),
321            from[2] + t * (to[2] - from[2]),
322        ];
323        self.within_aperture(crossing, tolerance)
324    }
325
326    /// Normalized aperture margin of a crossing point in \[0, 1\]: 1 at the gate
327    /// center, 0 at the aperture edge. Used as a clearance metric.
328    pub fn aperture_margin(&self, point: [f64; 3]) -> f64 {
329        let (right_off, up_off) = self.plane_offsets(point);
330        let right_margin = 1.0 - right_off.abs() / self.half_width;
331        let up_margin = 1.0 - up_off.abs() / self.half_height;
332        right_margin.min(up_margin)
333    }
334}
335
336/// A sequence of gates forming an open course or a closed racing lap.
337///
338/// Holds the reference-free gate-progress objective weights, mirroring
339/// [`crate::mppi::MppiGateRace2D`] but over 3-D gate planes.
340#[derive(Debug, Clone, PartialEq)]
341pub struct RacingGateLap3D {
342    pub gates: Vec<RacingGatePlane3D>,
343    /// When true, after the last gate the race wraps back to the first gate.
344    pub closed: bool,
345    pub progress_weight: f64,
346    pub lateral_weight: f64,
347    pub pass_bonus: f64,
348    pub miss_penalty: f64,
349    pub terminal_gate_weight: f64,
350    pub crossing_tolerance: f64,
351}
352
353/// Per-step outcome of scoring one rollout transition against the active gate.
354#[derive(Debug, Clone, Copy, PartialEq)]
355struct GateTransition {
356    cost: f64,
357    advanced: bool,
358}
359
360impl RacingGateLap3D {
361    fn with_defaults(gates: Vec<RacingGatePlane3D>, closed: bool) -> RoboticsResult<Self> {
362        if gates.is_empty() {
363            return Err(RoboticsError::InvalidParameter(
364                "racing lap must contain at least one gate".to_string(),
365            ));
366        }
367        Ok(Self {
368            gates,
369            closed,
370            progress_weight: 6.0,
371            lateral_weight: 0.4,
372            pass_bonus: 9.0,
373            miss_penalty: 36.0,
374            terminal_gate_weight: 1.4,
375            crossing_tolerance: 0.04,
376        })
377    }
378
379    /// Open course: the race ends after the last gate is passed.
380    pub fn open(gates: Vec<RacingGatePlane3D>) -> RoboticsResult<Self> {
381        Self::with_defaults(gates, false)
382    }
383
384    /// Closed lap: after the last gate the race wraps back to the first gate.
385    pub fn closed_loop(gates: Vec<RacingGatePlane3D>) -> RoboticsResult<Self> {
386        Self::with_defaults(gates, true)
387    }
388
389    pub fn gate_count(&self) -> usize {
390        self.gates.len()
391    }
392
393    /// Gate the race is chasing given how many gates have already been passed.
394    ///
395    /// For a closed lap this wraps modulo the gate count; for an open course it
396    /// saturates at the last gate.
397    pub fn gate_for_pass_count(&self, passed: usize) -> &RacingGatePlane3D {
398        let index = if self.closed {
399            passed % self.gates.len()
400        } else {
401            passed.min(self.gates.len() - 1)
402        };
403        &self.gates[index]
404    }
405
406    /// Length of the polyline through the gate centers; closed laps add the
407    /// return leg from the last gate back to the first.
408    pub fn centerline_length(&self) -> f64 {
409        let mut length = 0.0;
410        for pair in self.gates.windows(2) {
411            length += distance(pair[0].center, pair[1].center);
412        }
413        if self.closed && self.gates.len() > 1 {
414            length += distance(
415                self.gates[self.gates.len() - 1].center,
416                self.gates[0].center,
417            );
418        }
419        length
420    }
421
422    fn transition_cost(&self, from: [f64; 3], to: [f64; 3], passed: usize) -> GateTransition {
423        let gate = self.gate_for_pass_count(passed);
424        let before_sq = gate.squared_distance(from);
425        let after_sq = gate.squared_distance(to);
426        let (right_off, up_off) = gate.plane_offsets(to);
427        let after_distance = after_sq.sqrt();
428        let lateral_proximity = 1.0 / (1.0 + after_distance);
429        let lateral = right_off * right_off + up_off * up_off;
430        let mut cost = self.progress_weight * (after_sq - before_sq)
431            + self.lateral_weight * lateral_proximity * lateral;
432        let mut advanced = false;
433
434        if gate.crosses(from, to, self.crossing_tolerance) {
435            cost -= self.pass_bonus;
436            advanced = true;
437        } else if gate.signed_distance(to) >= 0.0 && !gate.within_aperture(to, 0.0) {
438            let right_miss = (right_off.abs() - gate.half_width).max(0.0);
439            let up_miss = (up_off.abs() - gate.half_height).max(0.0);
440            cost += self.miss_penalty * (right_miss * right_miss + up_miss * up_miss);
441        }
442
443        GateTransition { cost, advanced }
444    }
445
446    fn terminal_cost(&self, point: [f64; 3], passed: usize) -> f64 {
447        let gate = self.gate_for_pass_count(passed);
448        let remaining = if self.closed {
449            // Encourage progress within the current lap.
450            (self.gates.len() - passed % self.gates.len()) as f64
451        } else {
452            (self.gates.len() - passed.min(self.gates.len())) as f64
453        };
454        self.terminal_gate_weight * gate.squared_distance(point) + self.pass_bonus * remaining
455    }
456
457    /// Reference-free gate-progress cost of a position rollout starting from a
458    /// pass count, returning `(cost, gates_passed_in_rollout)`.
459    pub fn score_positions(&self, positions: &[[f64; 3]], start_passed: usize) -> (f64, usize) {
460        if positions.len() < 2 {
461            return (0.0, 0);
462        }
463        let mut passed = start_passed;
464        let mut passed_here = 0;
465        let mut cost = 0.0;
466        for pair in positions.windows(2) {
467            let transition = self.transition_cost(pair[0], pair[1], passed);
468            cost += transition.cost;
469            if transition.advanced {
470                passed += 1;
471                passed_here += 1;
472            }
473        }
474        cost += self.terminal_cost(positions[positions.len() - 1], passed);
475        (cost, passed_here)
476    }
477}
478
479/// Lap-progress metrics from a closed-loop racing rollout.
480#[derive(Debug, Clone, PartialEq)]
481pub struct RacingLapReport3D {
482    /// Control steps executed.
483    pub steps: usize,
484    /// Total gate passes across the whole rollout (counts repeated laps).
485    pub gates_passed: usize,
486    /// Whole laps completed.
487    pub laps_completed: usize,
488    /// Fraction of the current (partial) lap completed, in \[0, 1).
489    pub lap_fraction: f64,
490    /// Mean drone speed \[m/s\] over the rollout.
491    pub mean_speed: f64,
492    /// Peak drone speed \[m/s\] over the rollout.
493    pub max_speed: f64,
494    /// Executed path length \[m\].
495    pub path_length: f64,
496    /// Distance \[m\] to the currently active gate at the end of the rollout.
497    pub gate_distance: f64,
498    /// Smallest normalized aperture margin seen at any gate crossing (1 = dead
499    /// center, 0 = edge); `INFINITY` if no gate was crossed.
500    pub min_aperture_margin: f64,
501    /// Mean commanded acceleration magnitude \[m/s^2\].
502    pub mean_control_effort: f64,
503    /// Time \[s\] of the first completed lap, or `None` if no lap finished.
504    pub first_lap_time: Option<f64>,
505    /// Executed drone positions, including the start.
506    pub path: Vec<[f64; 3]>,
507}
508
509/// Configuration for the deterministic 3-D racing MPPI controller.
510#[derive(Debug, Clone, PartialEq)]
511pub struct RacingMppi3DConfig {
512    pub horizon: usize,
513    pub samples: usize,
514    pub dt: f64,
515    pub noise_sigma: f64,
516    pub velocity_weight: f64,
517    pub control_weight: f64,
518    pub lambda: f64,
519    pub seed: u64,
520}
521
522impl Default for RacingMppi3DConfig {
523    fn default() -> Self {
524        Self {
525            horizon: 20,
526            samples: 480,
527            dt: 0.1,
528            noise_sigma: 3.2,
529            velocity_weight: 0.01,
530            control_weight: 0.01,
531            lambda: 1.0,
532            seed: 7,
533        }
534    }
535}
536
537fn validate_config(config: &RacingMppi3DConfig) -> RoboticsResult<()> {
538    if config.horizon == 0 {
539        return Err(RoboticsError::InvalidParameter(
540            "racing MPPI horizon must be positive".to_string(),
541        ));
542    }
543    if config.samples == 0 {
544        return Err(RoboticsError::InvalidParameter(
545            "racing MPPI samples must be positive".to_string(),
546        ));
547    }
548    if !config.dt.is_finite() || config.dt <= 0.0 {
549        return Err(RoboticsError::InvalidParameter(
550            "racing MPPI dt must be finite and positive".to_string(),
551        ));
552    }
553    if !config.noise_sigma.is_finite() || config.noise_sigma <= 0.0 {
554        return Err(RoboticsError::InvalidParameter(
555            "racing MPPI noise_sigma must be finite and positive".to_string(),
556        ));
557    }
558    if !config.lambda.is_finite() || config.lambda <= 0.0 {
559        return Err(RoboticsError::InvalidParameter(
560            "racing MPPI lambda must be finite and positive".to_string(),
561        ));
562    }
563    if config.velocity_weight < 0.0 || config.control_weight < 0.0 {
564        return Err(RoboticsError::InvalidParameter(
565            "racing MPPI weights must be non-negative".to_string(),
566        ));
567    }
568    Ok(())
569}
570
571/// Deterministic, seeded MPPI controller for 3-D gate racing.
572///
573/// Each `plan` call samples Gaussian acceleration perturbations around the
574/// warm-started nominal control sequence, rolls them out through the drone
575/// dynamics, scores reference-free gate progress, and returns the softmin-
576/// weighted first control. The controller owns no obstacles — the gate
577/// apertures themselves are the constraints.
578#[derive(Debug, Clone)]
579pub struct RacingMppi3DController {
580    config: RacingMppi3DConfig,
581    dynamics: RacingDroneDynamics3D,
582    nominal: Vec<RacingDroneControl3D>,
583    rng: StdRng,
584}
585
586/// Result of one `plan` call.
587#[derive(Debug, Clone, PartialEq)]
588pub struct RacingMppi3DPlan {
589    pub control: RacingDroneControl3D,
590    /// Lowest rollout cost among the samples (diagnostic).
591    pub best_cost: f64,
592    /// Normalized effective sample size in \[0, 1\] (diagnostic).
593    pub normalized_effective_sample_size: f64,
594}
595
596impl RacingMppi3DController {
597    pub fn new(
598        config: RacingMppi3DConfig,
599        dynamics: RacingDroneDynamics3D,
600    ) -> RoboticsResult<Self> {
601        validate_config(&config)?;
602        let nominal = vec![RacingDroneControl3D::zero(); config.horizon];
603        let rng = StdRng::seed_from_u64(config.seed);
604        Ok(Self {
605            config,
606            dynamics,
607            nominal,
608            rng,
609        })
610    }
611
612    pub fn config(&self) -> &RacingMppi3DConfig {
613        &self.config
614    }
615
616    fn rollout_positions(
617        &self,
618        start: RacingDroneState3D,
619        controls: &[RacingDroneControl3D],
620    ) -> (Vec<[f64; 3]>, f64) {
621        let mut state = start;
622        let mut positions = Vec::with_capacity(controls.len() + 1);
623        positions.push(state.position());
624        let mut regularizer = 0.0;
625        for &control in controls {
626            state = self.dynamics.step(state, control, self.config.dt);
627            positions.push(state.position());
628            regularizer += self.config.velocity_weight * state.speed() * state.speed()
629                + self.config.control_weight * control.magnitude() * control.magnitude();
630        }
631        (positions, regularizer)
632    }
633
634    /// Plan the next control given the current state, lap, and pass count.
635    pub fn plan(
636        &mut self,
637        start: RacingDroneState3D,
638        lap: &RacingGateLap3D,
639        passed: usize,
640    ) -> RoboticsResult<RacingMppi3DPlan> {
641        let normal = Normal::new(0.0, self.config.noise_sigma).map_err(|_| {
642            RoboticsError::InvalidParameter("invalid racing MPPI noise distribution".to_string())
643        })?;
644
645        let mut costs = Vec::with_capacity(self.config.samples);
646        let mut perturbed_sequences = Vec::with_capacity(self.config.samples);
647        let mut best_cost = f64::INFINITY;
648
649        for _ in 0..self.config.samples {
650            let mut controls = Vec::with_capacity(self.config.horizon);
651            for &base in &self.nominal {
652                let noisy = RacingDroneControl3D::new(
653                    base.ax + normal.sample(&mut self.rng),
654                    base.ay + normal.sample(&mut self.rng),
655                    base.az + normal.sample(&mut self.rng),
656                );
657                controls.push(self.dynamics.saturate(noisy));
658            }
659            let (positions, regularizer) = self.rollout_positions(start, &controls);
660            let (gate_cost, _) = lap.score_positions(&positions, passed);
661            let cost = gate_cost + regularizer;
662            best_cost = best_cost.min(cost);
663            costs.push(cost);
664            perturbed_sequences.push(controls);
665        }
666
667        let min_cost = costs.iter().copied().fold(f64::INFINITY, f64::min);
668        let mut weights = Vec::with_capacity(costs.len());
669        let mut weight_sum = 0.0;
670        for &cost in &costs {
671            let weight = (-(cost - min_cost) / self.config.lambda).exp();
672            weight_sum += weight;
673            weights.push(weight);
674        }
675        if weight_sum <= 0.0 || !weight_sum.is_finite() {
676            // Degenerate weighting: fall back to the lowest-cost sample.
677            let best_index = costs
678                .iter()
679                .enumerate()
680                .min_by(|a, b| a.1.partial_cmp(b.1).unwrap_or(std::cmp::Ordering::Equal))
681                .map(|(index, _)| index)
682                .unwrap_or(0);
683            let control = perturbed_sequences[best_index][0];
684            return Ok(RacingMppi3DPlan {
685                control,
686                best_cost,
687                normalized_effective_sample_size: 0.0,
688            });
689        }
690
691        let mut updated = vec![RacingDroneControl3D::zero(); self.config.horizon];
692        let mut sum_sq = 0.0;
693        for (weight, controls) in weights.iter().zip(&perturbed_sequences) {
694            let normalized = weight / weight_sum;
695            sum_sq += normalized * normalized;
696            for (acc, control) in updated.iter_mut().zip(controls) {
697                acc.ax += normalized * control.ax;
698                acc.ay += normalized * control.ay;
699                acc.az += normalized * control.az;
700            }
701        }
702        for control in &mut updated {
703            *control = self.dynamics.saturate(*control);
704        }
705
706        let first_control = updated[0];
707        // Warm-start: shift the nominal sequence forward by one step.
708        self.nominal.clear();
709        self.nominal.extend_from_slice(&updated[1..]);
710        self.nominal.push(*updated.last().unwrap());
711
712        let normalized_effective_sample_size = if sum_sq > 0.0 {
713            1.0 / (sum_sq * self.config.samples as f64)
714        } else {
715            0.0
716        };
717
718        Ok(RacingMppi3DPlan {
719            control: first_control,
720            best_cost,
721            normalized_effective_sample_size,
722        })
723    }
724}
725
726/// Drive the racing MPPI controller around a lap and report lap-progress
727/// metrics.
728///
729/// The rollout runs for at most `max_steps` control steps and stops early once
730/// `target_laps` whole laps have been completed. Everything is deterministic
731/// for a fixed `config.seed`.
732pub fn simulate_lap_race(
733    config: RacingMppi3DConfig,
734    dynamics: RacingDroneDynamics3D,
735    lap: &RacingGateLap3D,
736    start: RacingDroneState3D,
737    max_steps: usize,
738    target_laps: usize,
739) -> RoboticsResult<RacingLapReport3D> {
740    let dt = config.dt;
741    let gate_count = lap.gate_count();
742    let mut controller = RacingMppi3DController::new(config, dynamics)?;
743
744    let mut state = start;
745    let mut passed = 0usize;
746    let mut path = vec![state.position()];
747    let mut sum_speed = 0.0;
748    let mut max_speed = 0.0_f64;
749    let mut path_length = 0.0;
750    let mut control_effort = 0.0;
751    let mut min_aperture_margin = f64::INFINITY;
752    let mut first_lap_time = None;
753    let mut executed_steps = 0usize;
754
755    for step in 0..max_steps {
756        let plan = controller.plan(state, lap, passed)?;
757        let next = dynamics.step(state, plan.control, dt);
758        let from = state.position();
759        let to = next.position();
760
761        // Resolve any gate crossings produced by this segment (usually one).
762        loop {
763            let gate = lap.gate_for_pass_count(passed);
764            if !gate.crosses(from, to, lap.crossing_tolerance) {
765                break;
766            }
767            let margin = aperture_crossing_margin(gate, from, to);
768            min_aperture_margin = min_aperture_margin.min(margin);
769            passed += 1;
770            if passed == gate_count && first_lap_time.is_none() {
771                first_lap_time = Some((step + 1) as f64 * dt);
772            }
773            // Open courses only have one gate sequence; stop once exhausted.
774            if !lap.closed && passed >= gate_count {
775                break;
776            }
777        }
778
779        path_length += distance(from, to);
780        control_effort += plan.control.magnitude();
781        let speed = next.speed();
782        sum_speed += speed;
783        max_speed = max_speed.max(speed);
784        state = next;
785        path.push(to);
786        executed_steps += 1;
787
788        if target_laps > 0 && passed >= target_laps * gate_count {
789            break;
790        }
791        if !lap.closed && passed >= gate_count {
792            break;
793        }
794    }
795
796    let laps_completed = passed / gate_count;
797    let lap_fraction = (passed % gate_count) as f64 / gate_count as f64;
798    let active_gate = lap.gate_for_pass_count(passed);
799    let gate_distance = active_gate.squared_distance(state.position()).sqrt();
800    let mean_speed = if executed_steps > 0 {
801        sum_speed / executed_steps as f64
802    } else {
803        0.0
804    };
805    let mean_control_effort = if executed_steps > 0 {
806        control_effort / executed_steps as f64
807    } else {
808        0.0
809    };
810
811    Ok(RacingLapReport3D {
812        steps: executed_steps,
813        gates_passed: passed,
814        laps_completed,
815        lap_fraction,
816        mean_speed,
817        max_speed,
818        path_length,
819        gate_distance,
820        min_aperture_margin,
821        mean_control_effort,
822        first_lap_time,
823        path,
824    })
825}
826
827fn aperture_crossing_margin(gate: &RacingGatePlane3D, from: [f64; 3], to: [f64; 3]) -> f64 {
828    let from_signed = gate.signed_distance(from);
829    let to_signed = gate.signed_distance(to);
830    let denom = to_signed - from_signed;
831    let t = if denom.abs() > 1e-9 {
832        (-from_signed / denom).clamp(0.0, 1.0)
833    } else {
834        0.0
835    };
836    let crossing = [
837        from[0] + t * (to[0] - from[0]),
838        from[1] + t * (to[1] - from[1]),
839        from[2] + t * (to[2] - from[2]),
840    ];
841    gate.aperture_margin(crossing)
842}
843
844// --- Small 3-D vector helpers (kept private to avoid a math dependency). ---
845
846fn dot(a: [f64; 3], b: [f64; 3]) -> f64 {
847    a[0] * b[0] + a[1] * b[1] + a[2] * b[2]
848}
849
850fn sub(a: [f64; 3], b: [f64; 3]) -> [f64; 3] {
851    [a[0] - b[0], a[1] - b[1], a[2] - b[2]]
852}
853
854fn scale(a: [f64; 3], s: f64) -> [f64; 3] {
855    [a[0] * s, a[1] * s, a[2] * s]
856}
857
858fn cross(a: [f64; 3], b: [f64; 3]) -> [f64; 3] {
859    [
860        a[1] * b[2] - a[2] * b[1],
861        a[2] * b[0] - a[0] * b[2],
862        a[0] * b[1] - a[1] * b[0],
863    ]
864}
865
866fn normalize(a: [f64; 3]) -> Option<[f64; 3]> {
867    let norm = dot(a, a).sqrt();
868    if !norm.is_finite() || norm <= 0.0 {
869        return None;
870    }
871    Some(scale(a, 1.0 / norm))
872}
873
874fn distance(a: [f64; 3], b: [f64; 3]) -> f64 {
875    let d = sub(a, b);
876    dot(d, d).sqrt()
877}
878
879#[cfg(test)]
880mod tests {
881    use super::*;
882
883    fn axis_gate(center: [f64; 3]) -> RacingGatePlane3D {
884        RacingGatePlane3D::new(center, [1.0, 0.0, 0.0], [0.0, 0.0, 1.0], 0.5, 0.5).unwrap()
885    }
886
887    #[test]
888    fn gate_orthonormalizes_axes() {
889        let gate =
890            RacingGatePlane3D::new([0.0, 0.0, 0.0], [1.0, 0.0, 0.0], [0.3, 0.0, 1.0], 0.5, 0.4)
891                .unwrap();
892        // up hint had a normal component; it must be removed.
893        assert!(dot(gate.normal(), [1.0, 0.0, 0.0]) > 0.999);
894        let (right, up) = (
895            gate.plane_offsets([0.0, 1.0, 0.0]),
896            gate.signed_distance([0.0, 1.0, 0.0]),
897        );
898        assert!(up.abs() < 1e-9); // point in the plane has zero signed distance
899        assert!(right.0.abs() > 0.0 || right.1.abs() > 0.0);
900    }
901
902    #[test]
903    fn gate_rejects_parallel_up_hint() {
904        let result =
905            RacingGatePlane3D::new([0.0, 0.0, 0.0], [0.0, 0.0, 1.0], [0.0, 0.0, 2.0], 0.5, 0.5);
906        assert!(result.is_err());
907    }
908
909    #[test]
910    fn gate_detects_central_crossing_and_rejects_miss() {
911        let gate = axis_gate([1.0, 0.0, 0.0]);
912        assert!(gate.crosses([0.5, 0.0, 0.0], [1.5, 0.0, 0.0], 0.04));
913        // Crosses the plane but far outside the aperture in y.
914        assert!(!gate.crosses([0.5, 3.0, 0.0], [1.5, 3.0, 0.0], 0.04));
915        // Moving away from the gate (wrong direction) is not a crossing.
916        assert!(!gate.crosses([1.5, 0.0, 0.0], [0.5, 0.0, 0.0], 0.04));
917    }
918
919    #[test]
920    fn aperture_margin_is_one_at_center() {
921        let gate = axis_gate([1.0, 0.0, 0.0]);
922        assert!((gate.aperture_margin([1.0, 0.0, 0.0]) - 1.0).abs() < 1e-9);
923        assert!(gate.aperture_margin([1.0, 0.5, 0.0]).abs() < 1e-9);
924    }
925
926    #[test]
927    fn drag_decays_velocity_without_control() {
928        let dynamics = RacingDroneDynamics3D::default();
929        let state = RacingDroneState3D::new(0.0, 0.0, 0.0, 4.0, 0.0, 0.0);
930        let next = dynamics.step(state, RacingDroneControl3D::zero(), 0.1);
931        assert!(next.vx < state.vx);
932        assert!(next.vx > 0.0);
933    }
934
935    #[test]
936    fn dynamics_caps_speed() {
937        let dynamics = RacingDroneDynamics3D::new(0.0, 0.0, 3.0, 100.0).unwrap();
938        let state = RacingDroneState3D::at(0.0, 0.0, 0.0);
939        let next = dynamics.step(state, RacingDroneControl3D::new(100.0, 0.0, 0.0), 0.1);
940        assert!(next.speed() <= 3.0 + 1e-9);
941    }
942
943    #[test]
944    fn closed_lap_centerline_includes_return_leg() {
945        let gates = vec![
946            axis_gate([0.0, 0.0, 0.0]),
947            axis_gate([1.0, 0.0, 0.0]),
948            axis_gate([1.0, 1.0, 0.0]),
949        ];
950        let open = RacingGateLap3D::open(gates.clone()).unwrap();
951        let closed = RacingGateLap3D::closed_loop(gates).unwrap();
952        assert!(closed.centerline_length() > open.centerline_length());
953    }
954
955    #[test]
956    fn closed_lap_wraps_active_gate() {
957        let gates = vec![axis_gate([0.0, 0.0, 0.0]), axis_gate([1.0, 0.0, 0.0])];
958        let lap = RacingGateLap3D::closed_loop(gates).unwrap();
959        assert_eq!(lap.gate_for_pass_count(0).center(), [0.0, 0.0, 0.0]);
960        assert_eq!(lap.gate_for_pass_count(2).center(), [0.0, 0.0, 0.0]);
961        assert_eq!(lap.gate_for_pass_count(3).center(), [1.0, 0.0, 0.0]);
962    }
963
964    fn straight_lap() -> RacingGateLap3D {
965        let gates = vec![
966            axis_gate([1.0, 0.0, 0.0]),
967            axis_gate([2.0, 0.0, 0.3]),
968            axis_gate([3.0, 0.0, 0.6]),
969        ];
970        RacingGateLap3D::open(gates).unwrap()
971    }
972
973    #[test]
974    fn controller_makes_gate_progress() {
975        let lap = straight_lap();
976        let report = simulate_lap_race(
977            RacingMppi3DConfig::default(),
978            RacingDroneDynamics3D::default(),
979            &lap,
980            RacingDroneState3D::at(0.0, 0.0, 0.0),
981            60,
982            1,
983        )
984        .unwrap();
985        assert!(
986            report.gates_passed >= 1,
987            "expected at least one gate, got {}",
988            report.gates_passed
989        );
990        assert!(report.path_length > 0.0);
991    }
992
993    #[test]
994    fn simulation_is_deterministic() {
995        let lap = straight_lap();
996        let run = || {
997            simulate_lap_race(
998                RacingMppi3DConfig::default(),
999                RacingDroneDynamics3D::default(),
1000                &lap,
1001                RacingDroneState3D::at(0.0, 0.0, 0.0),
1002                60,
1003                1,
1004            )
1005            .unwrap()
1006        };
1007        let first = run();
1008        let second = run();
1009        assert_eq!(first.path, second.path);
1010        assert_eq!(first.gates_passed, second.gates_passed);
1011    }
1012
1013    #[test]
1014    fn invalid_config_is_rejected() {
1015        let config = RacingMppi3DConfig {
1016            horizon: 0,
1017            ..RacingMppi3DConfig::default()
1018        };
1019        assert!(RacingMppi3DController::new(config, RacingDroneDynamics3D::default()).is_err());
1020    }
1021}
rust_robotics_control/racing_mppi_3d.rs

rust_robotics_control/
racing_mppi_3d.rs
