use rand::SeedableRng;
use rayon::prelude::*;
pub mod config;
use nalgebra::{
DMatrix, DVector, Matrix3, Quaternion, Rotation3, SMatrix, UnitQuaternion, Vector3,
};
use osqp::{CscMatrix, Problem, Settings};
use rand_chacha::ChaCha8Rng;
use rand_distr::{Distribution, Normal};
use std::f32::consts::PI;
use std::ops::AddAssign;
#[derive(thiserror::Error, Debug)]
pub enum SimulationError {
#[error("Rerun error: {0}")]
RerunError(#[from] rerun::RecordingStreamError),
#[error("Rerun spawn error: {0}")]
RerunSpawnError(#[from] rerun::SpawnError),
#[error("Rerun SetLogger error: {0}")]
SetLoggerError(#[from] rerun::external::log::SetLoggerError),
#[error("Nalgebra error: {0}")]
NalgebraError(String),
#[error("OSQP error: {0}")]
OSQPError(String),
#[error("Normal error: {0}")]
NormalError(#[from] rand_distr::NormalError),
#[error("Other error: {0}")]
OtherError(String),
}
pub struct Quadrotor {
pub position: Vector3<f32>,
pub velocity: Vector3<f32>,
pub acceleration: Vector3<f32>,
pub orientation: UnitQuaternion<f32>,
pub angular_velocity: Vector3<f32>,
pub mass: f32,
pub gravity: f32,
pub time_step: f32,
pub drag_coefficient: f32,
pub inertia_matrix: Matrix3<f32>,
pub inertia_matrix_inv: Matrix3<f32>,
}
impl Quadrotor {
pub fn new(
time_step: f32,
mass: f32,
gravity: f32,
drag_coefficient: f32,
inertia_matrix: [f32; 9],
) -> Result<Self, SimulationError> {
let inertia_matrix = Matrix3::from_row_slice(&inertia_matrix);
let inertia_matrix_inv =
inertia_matrix
.try_inverse()
.ok_or(SimulationError::NalgebraError(
"Failed to invert inertia matrix".to_string(),
))?;
Ok(Self {
position: Vector3::zeros(),
velocity: Vector3::zeros(),
acceleration: Vector3::zeros(),
orientation: UnitQuaternion::identity(),
angular_velocity: Vector3::zeros(),
mass,
gravity,
time_step,
drag_coefficient,
inertia_matrix,
inertia_matrix_inv,
})
}
pub fn update_dynamics_with_controls_euler(
&mut self,
control_thrust: f32,
control_torque: &Vector3<f32>,
) {
let gravity_force = Vector3::new(0.0, 0.0, -self.mass * self.gravity);
let drag_force = -self.drag_coefficient * self.velocity.norm() * self.velocity;
let thrust_world = self.orientation * Vector3::new(0.0, 0.0, control_thrust);
self.acceleration = (thrust_world + gravity_force + drag_force) / self.mass;
self.velocity += self.acceleration * self.time_step;
self.position += self.velocity * self.time_step;
let inertia_angular_velocity = self.inertia_matrix * self.angular_velocity;
let gyroscopic_torque = self.angular_velocity.cross(&inertia_angular_velocity);
let angular_acceleration = self.inertia_matrix_inv * (control_torque - gyroscopic_torque);
self.angular_velocity += angular_acceleration * self.time_step;
self.orientation *=
UnitQuaternion::from_scaled_axis(self.angular_velocity * self.time_step);
}
pub fn update_dynamics_with_controls_rk4(
&mut self,
control_thrust: f32,
control_torque: &Vector3<f32>,
) {
let h = self.time_step;
let state = self.get_state();
let k1 = self.state_derivative(&state, control_thrust, control_torque);
let mut temp_state = [0.0; 13];
for i in 0..13 {
temp_state[i] = state[i] + 0.5 * h * k1[i];
}
let k2 = self.state_derivative(&temp_state, control_thrust, control_torque);
for i in 0..13 {
temp_state[i] = state[i] + 0.5 * h * k2[i];
}
let k3 = self.state_derivative(&temp_state, control_thrust, control_torque);
for i in 0..13 {
temp_state[i] = state[i] + h * k3[i];
}
let k4 = self.state_derivative(&temp_state, control_thrust, control_torque);
let mut new_state = [0.0; 13];
for i in 0..13 {
new_state[i] = state[i] + (h / 6.0) * (k1[i] + 2.0 * k2[i] + 2.0 * k3[i] + k4[i]);
}
let mut q = Quaternion::new(new_state[9], new_state[6], new_state[7], new_state[8]);
q = q.normalize();
new_state[6..10].copy_from_slice(q.coords.as_slice());
self.set_state(&new_state);
}
pub fn get_state(&self) -> [f32; 13] {
let mut state = [0.0; 13];
state[0..3].copy_from_slice(self.position.as_slice());
state[3..6].copy_from_slice(self.velocity.as_slice());
state[6..10].copy_from_slice(self.orientation.coords.as_slice());
state[10..13].copy_from_slice(self.angular_velocity.as_slice());
state
}
pub fn set_state(&mut self, state: &[f32; 13]) {
self.position = Vector3::from_column_slice(&state[0..3]);
self.velocity = Vector3::from_column_slice(&state[3..6]);
self.orientation = UnitQuaternion::from_quaternion(Quaternion::new(
state[9], state[6], state[7], state[8],
));
self.angular_velocity = Vector3::from_column_slice(&state[10..13]);
}
pub fn state_derivative(
&mut self,
state: &[f32],
control_thrust: f32,
control_torque: &Vector3<f32>,
) -> [f32; 13] {
let velocity = Vector3::from_column_slice(&state[3..6]);
let orientation = UnitQuaternion::from_quaternion(Quaternion::new(
state[9], state[6], state[7], state[8],
));
let omega_quat = Quaternion::new(0.0, state[10], state[11], state[12]);
let q_dot = orientation.into_inner() * omega_quat * 0.5;
let angular_velocity = Vector3::from_column_slice(&state[10..13]);
let gravity_force = Vector3::new(0.0, 0.0, -self.mass * self.gravity);
let drag_force = -self.drag_coefficient * velocity.norm() * velocity;
let thrust_world = orientation * Vector3::new(0.0, 0.0, control_thrust);
self.acceleration = (thrust_world + gravity_force + drag_force) / self.mass;
let inertia_angular_velocity = self.inertia_matrix * angular_velocity;
let gyroscopic_torque = angular_velocity.cross(&inertia_angular_velocity);
let angular_acceleration = self.inertia_matrix_inv * (control_torque - gyroscopic_torque);
let mut derivative = [0.0; 13];
derivative[0..3].copy_from_slice(velocity.as_slice());
derivative[3..6].copy_from_slice(self.acceleration.as_slice());
derivative[6..10].copy_from_slice(q_dot.coords.as_slice());
derivative[10..13].copy_from_slice(angular_acceleration.as_slice());
derivative
}
pub fn read_imu(&self) -> Result<(Vector3<f32>, Vector3<f32>), SimulationError> {
Ok((self.acceleration, self.angular_velocity))
}
}
pub struct Imu {
pub accel_bias: Vector3<f32>,
pub gyro_bias: Vector3<f32>,
pub accel_noise_std: f32,
pub gyro_noise_std: f32,
pub accel_bias_std: f32,
pub gyro_bias_std: f32,
accel_noise: Normal<f32>,
gyro_noise: Normal<f32>,
accel_bias_drift: Normal<f32>,
gyro_bias_drift: Normal<f32>,
rng: ChaCha8Rng,
}
impl Imu {
pub fn new(
accel_noise_std: f32,
gyro_noise_std: f32,
accel_bias_std: f32,
gyro_bias_std: f32,
) -> Result<Self, SimulationError> {
Ok(Self {
accel_bias: Vector3::zeros(),
gyro_bias: Vector3::zeros(),
accel_noise_std,
gyro_noise_std,
accel_bias_std,
gyro_bias_std,
accel_noise: Normal::new(0.0, accel_noise_std)?,
gyro_noise: Normal::new(0.0, gyro_noise_std)?,
accel_bias_drift: Normal::new(0.0, accel_bias_std)?,
gyro_bias_drift: Normal::new(0.0, gyro_bias_std)?,
rng: ChaCha8Rng::from_os_rng(),
})
}
pub fn update(&mut self, dt: f32) -> Result<(), SimulationError> {
let dt_sqrt = fast_sqrt(dt);
let accel_drift = self.accel_bias_drift.sample(&mut self.rng) * dt_sqrt;
let gyro_drift = self.gyro_bias_drift.sample(&mut self.rng) * dt_sqrt;
self.accel_bias += Vector3::from_iterator((0..3).map(|_| accel_drift));
self.gyro_bias += Vector3::from_iterator((0..3).map(|_| gyro_drift));
Ok(())
}
pub fn read(
&mut self,
true_acceleration: Vector3<f32>,
true_angular_velocity: Vector3<f32>,
) -> Result<(Vector3<f32>, Vector3<f32>), SimulationError> {
let accel_noise_sample =
Vector3::from_iterator((0..3).map(|_| self.accel_noise.sample(&mut self.rng)));
let gyro_noise_sample =
Vector3::from_iterator((0..3).map(|_| self.gyro_noise.sample(&mut self.rng)));
let measured_acceleration = true_acceleration + self.accel_bias + accel_noise_sample;
let measured_ang_velocity = true_angular_velocity + self.gyro_bias + gyro_noise_sample;
Ok((measured_acceleration, measured_ang_velocity))
}
}
pub struct PIDController {
pub kpid_pos: [Vector3<f32>; 3],
pub kpid_att: [Vector3<f32>; 3],
pub integral_pos_error: Vector3<f32>,
pub integral_att_error: Vector3<f32>,
pub max_integral_pos: Vector3<f32>,
pub max_integral_att: Vector3<f32>,
pub mass: f32,
pub gravity: f32,
}
impl PIDController {
pub fn new(
_kpid_pos: [[f32; 3]; 3],
_kpid_att: [[f32; 3]; 3],
_max_integral_pos: [f32; 3],
_max_integral_att: [f32; 3],
_mass: f32,
_gravity: f32,
) -> Self {
Self {
kpid_pos: _kpid_pos.map(Vector3::from),
kpid_att: _kpid_att.map(Vector3::from),
integral_pos_error: Vector3::zeros(),
integral_att_error: Vector3::zeros(),
max_integral_pos: Vector3::from(_max_integral_pos),
max_integral_att: Vector3::from(_max_integral_att),
mass: _mass,
gravity: _gravity,
}
}
pub fn compute_attitude_control(
&mut self,
desired_orientation: &UnitQuaternion<f32>,
current_orientation: &UnitQuaternion<f32>,
current_angular_velocity: &Vector3<f32>,
dt: f32,
) -> Vector3<f32> {
let error_orientation = current_orientation.inverse() * desired_orientation;
let (roll_error, pitch_error, yaw_error) = error_orientation.euler_angles();
let error_angles = Vector3::new(roll_error, pitch_error, yaw_error);
self.integral_att_error += error_angles * dt;
self.integral_att_error = self
.integral_att_error
.zip_map(&self.max_integral_att, |int, max| int.clamp(-max, max));
let error_angular_velocity = -current_angular_velocity; self.kpid_att[0].component_mul(&error_angles)
+ self.kpid_att[1].component_mul(&error_angular_velocity)
+ self.kpid_att[2].component_mul(&self.integral_att_error)
}
pub fn compute_position_control(
&mut self,
desired_position: &Vector3<f32>,
desired_velocity: &Vector3<f32>,
desired_yaw: f32,
current_position: &Vector3<f32>,
current_velocity: &Vector3<f32>,
dt: f32,
) -> (f32, UnitQuaternion<f32>) {
let error_position = desired_position - current_position;
let error_velocity = desired_velocity - current_velocity;
self.integral_pos_error += error_position * dt;
self.integral_pos_error = self
.integral_pos_error
.zip_map(&self.max_integral_pos, |int, max| int.clamp(-max, max));
let acceleration = self.kpid_pos[0].component_mul(&error_position)
+ self.kpid_pos[1].component_mul(&error_velocity)
+ self.kpid_pos[2].component_mul(&self.integral_pos_error);
let gravity_compensation = Vector3::new(0.0, 0.0, self.gravity);
let total_acceleration = acceleration + gravity_compensation;
let total_acc_norm = total_acceleration.norm();
let thrust = self.mass * total_acc_norm;
let desired_orientation = if total_acc_norm > 1e-6 {
let z_body = total_acceleration / total_acc_norm;
let yaw_rotation = UnitQuaternion::from_euler_angles(0.0, 0.0, desired_yaw);
let x_body_horizontal = yaw_rotation * Vector3::new(1.0, 0.0, 0.0);
let y_body = z_body.cross(&x_body_horizontal).normalize();
let x_body = y_body.cross(&z_body);
UnitQuaternion::from_rotation_matrix(&Rotation3::from_matrix_unchecked(
Matrix3::from_columns(&[x_body, y_body, z_body]),
))
} else {
UnitQuaternion::from_euler_angles(0.0, 0.0, desired_yaw)
};
(thrust, desired_orientation)
}
}
pub enum PlannerType {
Hover(HoverPlanner),
MinimumJerkLine(MinimumJerkLinePlanner),
Lissajous(LissajousPlanner),
Circle(CirclePlanner),
Landing(LandingPlanner),
ObstacleAvoidance(ObstacleAvoidancePlanner),
MinimumSnapWaypoint(MinimumSnapWaypointPlanner),
QPpolyTraj(QPpolyTrajPlanner),
}
impl PlannerType {
pub fn plan(
&self,
current_position: Vector3<f32>,
current_velocity: Vector3<f32>,
time: f32,
) -> (Vector3<f32>, Vector3<f32>, f32) {
match self {
PlannerType::Hover(p) => p.plan(current_position, current_velocity, time),
PlannerType::MinimumJerkLine(p) => p.plan(current_position, current_velocity, time),
PlannerType::Lissajous(p) => p.plan(current_position, current_velocity, time),
PlannerType::Circle(p) => p.plan(current_position, current_velocity, time),
PlannerType::Landing(p) => p.plan(current_position, current_velocity, time),
PlannerType::ObstacleAvoidance(p) => p.plan(current_position, current_velocity, time),
PlannerType::MinimumSnapWaypoint(p) => p.plan(current_position, current_velocity, time),
PlannerType::QPpolyTraj(p) => p.plan(current_position, current_velocity, time),
}
}
pub fn is_finished(
&self,
current_position: Vector3<f32>,
time: f32,
) -> Result<bool, SimulationError> {
match self {
PlannerType::Hover(p) => p.is_finished(current_position, time),
PlannerType::MinimumJerkLine(p) => p.is_finished(current_position, time),
PlannerType::Lissajous(p) => p.is_finished(current_position, time),
PlannerType::Circle(p) => p.is_finished(current_position, time),
PlannerType::Landing(p) => p.is_finished(current_position, time),
PlannerType::ObstacleAvoidance(p) => p.is_finished(current_position, time),
PlannerType::MinimumSnapWaypoint(p) => p.is_finished(current_position, time),
PlannerType::QPpolyTraj(p) => p.is_finished(current_position, time),
}
}
}
pub trait Planner {
fn plan(
&self,
current_position: Vector3<f32>,
current_velocity: Vector3<f32>,
time: f32,
) -> (Vector3<f32>, Vector3<f32>, f32);
fn is_finished(
&self,
current_position: Vector3<f32>,
time: f32,
) -> Result<bool, SimulationError>;
}
pub struct HoverPlanner {
pub target_position: Vector3<f32>,
pub target_yaw: f32,
}
impl Planner for HoverPlanner {
fn plan(
&self,
_current_position: Vector3<f32>,
_current_velocity: Vector3<f32>,
_time: f32,
) -> (Vector3<f32>, Vector3<f32>, f32) {
(self.target_position, Vector3::zeros(), self.target_yaw)
}
fn is_finished(
&self,
_current_position: Vector3<f32>,
_time: f32,
) -> Result<bool, SimulationError> {
Ok(false) }
}
pub struct MinimumJerkLinePlanner {
pub start_position: Vector3<f32>,
pub end_position: Vector3<f32>,
pub start_yaw: f32,
pub end_yaw: f32,
pub start_time: f32,
pub duration: f32,
}
impl Planner for MinimumJerkLinePlanner {
fn plan(
&self,
_current_position: Vector3<f32>,
_current_velocity: Vector3<f32>,
time: f32,
) -> (Vector3<f32>, Vector3<f32>, f32) {
let t = ((time - self.start_time) / self.duration).clamp(0.0, 1.0);
let t2 = t * t;
let t3 = t2 * t;
let t4 = t3 * t;
let s = 10.0 * t2 - 15.0 * t3 + 6.0 * t4;
let s_dot = (30.0 * t.powi(2) - 60.0 * t.powi(3) + 30.0 * t.powi(4)) / self.duration;
let position = self.start_position + (self.end_position - self.start_position) * s;
let velocity = (self.end_position - self.start_position) * s_dot;
let yaw = self.start_yaw + (self.end_yaw - self.start_yaw) * s;
(position, velocity, yaw)
}
fn is_finished(
&self,
_current_position: Vector3<f32>,
_time: f32,
) -> Result<bool, SimulationError> {
Ok((_current_position - self.end_position).norm() < 0.01
&& _time >= self.start_time + self.duration)
}
}
pub struct LissajousPlanner {
pub start_position: Vector3<f32>,
pub center: Vector3<f32>,
pub amplitude: Vector3<f32>,
pub frequency: Vector3<f32>,
pub phase: Vector3<f32>,
pub start_time: f32,
pub duration: f32,
pub start_yaw: f32,
pub end_yaw: f32,
pub ramp_time: f32,
}
impl Planner for LissajousPlanner {
fn plan(
&self,
_current_position: Vector3<f32>,
_current_velocity: Vector3<f32>,
time: f32,
) -> (Vector3<f32>, Vector3<f32>, f32) {
let t = ((time - self.start_time) / self.duration).clamp(0.0, 1.0);
let smooth_start = if t < self.ramp_time / self.duration {
let t_ramp = t / (self.ramp_time / self.duration);
t_ramp * t_ramp * (3.0 - 2.0 * t_ramp)
} else {
1.0
};
let velocity_ramp = if t < self.ramp_time / self.duration {
smooth_start
} else if t > 1.0 - self.ramp_time / self.duration {
let t_down = (1.0 - t) / (self.ramp_time / self.duration);
t_down * t_down * (3.0 - 2.0 * t_down)
} else {
1.0
};
let ang_pos = self.frequency * t * 2.0 * PI + self.phase;
let lissajous = self.amplitude.component_mul(&ang_pos.map(f32::sin));
let position =
self.start_position + smooth_start * ((self.center + lissajous) - self.start_position);
let mut velocity = Vector3::new(
self.amplitude.x * self.frequency.x * 2.0 * PI * ang_pos.x.cos(),
self.amplitude.y * self.frequency.y * 2.0 * PI * ang_pos.y.cos(),
self.amplitude.z * self.frequency.z * 2.0 * PI * ang_pos.z.cos(),
) * velocity_ramp
/ self.duration;
if t < self.ramp_time / self.duration {
let transition_velocity = (self.center - self.start_position)
* (2.0 * t / self.ramp_time - 2.0 * t * t / (self.ramp_time * self.ramp_time))
/ self.duration;
velocity += transition_velocity;
}
let yaw = self.start_yaw + (self.end_yaw - self.start_yaw) * t;
(position, velocity, yaw)
}
fn is_finished(
&self,
_current_position: Vector3<f32>,
time: f32,
) -> Result<bool, SimulationError> {
Ok(time >= self.start_time + self.duration)
}
}
pub struct CirclePlanner {
pub center: Vector3<f32>,
pub radius: f32,
pub angular_velocity: f32,
pub start_position: Vector3<f32>,
pub start_time: f32,
pub duration: f32,
pub start_yaw: f32,
pub end_yaw: f32,
pub ramp_time: f32,
}
impl Planner for CirclePlanner {
fn plan(
&self,
_current_position: Vector3<f32>,
_current_velocity: Vector3<f32>,
time: f32,
) -> (Vector3<f32>, Vector3<f32>, f32) {
let t = (time - self.start_time) / self.duration;
let t = t.clamp(0.0, 1.0);
let smooth_start = if t < self.ramp_time / self.duration {
let t_ramp = t / (self.ramp_time / self.duration);
t_ramp * t_ramp * (3.0 - 2.0 * t_ramp)
} else {
1.0
};
let velocity_ramp = if t < self.ramp_time / self.duration {
smooth_start
} else if t > 1.0 - self.ramp_time / self.duration {
let t_down = (1.0 - t) / (self.ramp_time / self.duration);
t_down * t_down * (3.0 - 2.0 * t_down)
} else {
1.0
};
let angle = self.angular_velocity * t * self.duration;
let circle_offset = Vector3::new(self.radius * angle.cos(), self.radius * angle.sin(), 0.0);
let position = self.start_position
+ smooth_start * ((self.center + circle_offset) - self.start_position);
let tangential_velocity = Vector3::new(
-self.radius * self.angular_velocity * angle.sin(),
self.radius * self.angular_velocity * angle.cos(),
0.0,
);
let velocity = tangential_velocity * velocity_ramp;
let yaw = self.start_yaw + (self.end_yaw - self.start_yaw) * t;
(position, velocity, yaw)
}
fn is_finished(
&self,
_current_position: Vector3<f32>,
time: f32,
) -> Result<bool, SimulationError> {
Ok(time >= self.start_time + self.duration)
}
}
pub struct LandingPlanner {
pub start_position: Vector3<f32>,
pub start_time: f32,
pub duration: f32,
pub start_yaw: f32,
}
impl Planner for LandingPlanner {
fn plan(
&self,
_current_position: Vector3<f32>,
_current_velocity: Vector3<f32>,
time: f32,
) -> (Vector3<f32>, Vector3<f32>, f32) {
let t = ((time - self.start_time) / self.duration).clamp(0.0, 1.0);
let target_z = self.start_position.z * (1.0 - t);
let target_position = Vector3::new(self.start_position.x, self.start_position.y, target_z);
let target_velocity = Vector3::new(0.0, 0.0, -self.start_position.z / self.duration);
(target_position, target_velocity, self.start_yaw)
}
fn is_finished(
&self,
current_position: Vector3<f32>,
time: f32,
) -> Result<bool, SimulationError> {
Ok(current_position.z < 0.05 || time >= self.start_time + self.duration)
}
}
pub struct PlannerManager {
pub current_planner: PlannerType,
}
impl PlannerManager {
pub fn new(initial_position: Vector3<f32>, initial_yaw: f32) -> Self {
let hover_planner = HoverPlanner {
target_position: initial_position,
target_yaw: initial_yaw,
};
Self {
current_planner: PlannerType::Hover(hover_planner),
}
}
pub fn set_planner(&mut self, new_planner: PlannerType) {
self.current_planner = new_planner;
}
pub fn update(
&mut self,
current_position: Vector3<f32>,
current_orientation: UnitQuaternion<f32>,
current_velocity: Vector3<f32>,
time: f32,
obstacles: &[Obstacle],
) -> Result<(Vector3<f32>, Vector3<f32>, f32), SimulationError> {
if self.current_planner.is_finished(current_position, time)? {
log::info!("Time: {time:.2} s,\tSwitch Hover");
self.current_planner = PlannerType::Hover(HoverPlanner {
target_position: current_position,
target_yaw: current_orientation.euler_angles().2,
});
}
if let PlannerType::ObstacleAvoidance(ref mut planner) = self.current_planner {
planner.obstacles = obstacles.to_owned();
}
Ok(self
.current_planner
.plan(current_position, current_velocity, time))
}
}
pub struct ObstacleAvoidancePlanner {
pub target_position: Vector3<f32>,
pub start_time: f32,
pub duration: f32,
pub start_yaw: f32,
pub end_yaw: f32,
pub obstacles: Vec<Obstacle>,
pub k_att: f32,
pub k_rep: f32,
pub k_vortex: f32,
pub d0: f32,
pub d_target: f32,
pub max_speed: f32,
}
impl Planner for ObstacleAvoidancePlanner {
fn plan(
&self,
current_position: Vector3<f32>,
current_velocity: Vector3<f32>,
time: f32,
) -> (Vector3<f32>, Vector3<f32>, f32) {
let t = ((time - self.start_time) / self.duration).clamp(0.0, 1.0);
let distance_to_target = (self.target_position - current_position).norm();
let f_att = self.k_att
* self.smooth_attractive_force(distance_to_target)
* (self.target_position - current_position).normalize();
let mut f_rep = Vector3::zeros();
let mut f_vortex = Vector3::zeros();
for obstacle in &self.obstacles {
let diff = current_position - obstacle.position;
let distance = diff.norm();
if distance < self.d0 {
f_rep += self.k_rep
* (1.0 / distance - 1.0 / self.d0)
* (1.0 / distance.powi(2))
* diff.normalize();
f_vortex +=
self.k_vortex * current_velocity.cross(&diff).normalize() / distance.powi(2);
}
}
let f_total = f_att + f_rep + f_vortex;
let desired_velocity = f_total.normalize() * self.max_speed.min(f_total.norm());
let desired_position = current_position + desired_velocity * self.duration * (1.0 - t);
let desired_yaw = self.start_yaw + (self.end_yaw - self.start_yaw) * t;
(desired_position, desired_velocity, desired_yaw)
}
fn is_finished(
&self,
current_position: Vector3<f32>,
time: f32,
) -> Result<bool, SimulationError> {
Ok((current_position - self.target_position).norm() < 0.1
&& time >= self.start_time + self.duration)
}
}
impl ObstacleAvoidancePlanner {
#[inline]
pub fn smooth_attractive_force(&self, distance: f32) -> f32 {
if distance <= self.d_target {
distance
} else {
self.d_target + (distance - self.d_target).tanh()
}
}
}
pub struct MinimumSnapWaypointPlanner {
pub waypoints: Vec<Vector3<f32>>,
pub yaws: Vec<f32>,
pub times: Vec<f32>,
pub coefficients: Vec<Vec<Vector3<f32>>>,
pub yaw_coefficients: Vec<Vec<f32>>,
pub start_time: f32,
}
impl MinimumSnapWaypointPlanner {
pub fn new(
waypoints: Vec<Vector3<f32>>,
yaws: Vec<f32>,
segment_times: Vec<f32>,
start_time: f32,
) -> Result<Self, SimulationError> {
if waypoints.len() < 2 {
return Err(SimulationError::OtherError(
"At least two waypoints are required".to_string(),
));
}
if waypoints.len() != segment_times.len() + 1 || waypoints.len() != yaws.len() {
return Err(SimulationError::OtherError("Number of segment times must be one less than number of waypoints, and yaws must match waypoints".to_string()));
}
let mut planner = Self {
waypoints,
yaws,
times: segment_times,
coefficients: Vec::new(),
yaw_coefficients: Vec::new(),
start_time,
};
planner.compute_minimum_snap_trajectories()?;
planner.compute_minimum_acceleration_yaw_trajectories()?;
Ok(planner)
}
pub fn compute_minimum_snap_trajectories(&mut self) -> Result<(), SimulationError> {
let n = self.waypoints.len() - 1;
for i in 0..n {
let duration = self.times[i];
let (start, end) = (self.waypoints[i], self.waypoints[i + 1]);
let mut a = SMatrix::<f32, 8, 8>::zeros();
let mut b = SMatrix::<f32, 8, 3>::zeros();
a.fixed_view_mut::<4, 4>(0, 0).fill_with_identity();
b.fixed_view_mut::<1, 3>(0, 0).copy_from(&start.transpose());
b.fixed_view_mut::<1, 3>(4, 0).copy_from(&end.transpose());
for j in 0..8 {
a[(4, j)] = duration.powi(j as i32);
if j > 0 {
a[(5, j)] = j as f32 * duration.powi(j as i32 - 1);
}
if j > 1 {
a[(6, j)] = j as f32 * (j - 1) as f32 * duration.powi(j as i32 - 2);
}
if j > 2 {
a[(7, j)] =
j as f32 * (j - 1) as f32 * (j - 2) as f32 * duration.powi(j as i32 - 3);
}
}
let coeffs = a.lu().solve(&b).ok_or(SimulationError::NalgebraError(
"Failed to solve for coefficients in MinimumSnapWaypointPlanner".to_string(),
))?;
self.coefficients.push(
(0..8)
.map(|j| Vector3::new(coeffs[(j, 0)], coeffs[(j, 1)], coeffs[(j, 2)]))
.collect(),
);
}
Ok(())
}
pub fn compute_minimum_acceleration_yaw_trajectories(&mut self) -> Result<(), SimulationError> {
let n = self.yaws.len() - 1; for i in 0..n {
let (duration, start_yaw, end_yaw) = (self.times[i], self.yaws[i], self.yaws[i + 1]);
let mut a = SMatrix::<f32, 4, 4>::zeros();
let mut b = SMatrix::<f32, 4, 1>::zeros();
(a[(0, 0)], a[(1, 1)]) = (1.0, 1.0);
(b[0], b[2]) = (start_yaw, end_yaw);
for j in 0..4 {
a[(2, j)] = duration.powi(j as i32);
if j > 0 {
a[(3, j)] = j as f32 * duration.powi(j as i32 - 1);
}
}
let yaw_coeffs = a.lu().solve(&b).ok_or(SimulationError::NalgebraError(
"Failed to solve for yaw coefficients in MinimumSnapWaypointPlanner".to_string(),
))?;
self.yaw_coefficients.push(yaw_coeffs.as_slice().to_vec());
}
Ok(())
}
pub fn evaluate_polynomial(
&self,
t: f32,
coeffs: &[Vector3<f32>],
yaw_coeffs: &[f32],
) -> (Vector3<f32>, Vector3<f32>, f32, f32) {
let mut position = Vector3::zeros();
let mut velocity = Vector3::zeros();
let mut yaw = 0.0;
let mut yaw_rate = 0.0;
for (i, coeff) in coeffs.iter().enumerate() {
let ti = t.powi(i as i32);
position += coeff * ti;
if i > 0 {
velocity += coeff * (i as f32) * t.powi(i as i32 - 1);
}
}
for (i, &coeff) in yaw_coeffs.iter().enumerate() {
let ti = t.powi(i as i32);
yaw += coeff * ti;
if i > 0 {
yaw_rate += coeff * (i as f32) * t.powi(i as i32 - 1);
}
}
(position, velocity, yaw, yaw_rate)
}
}
impl Planner for MinimumSnapWaypointPlanner {
fn plan(
&self,
_current_position: Vector3<f32>,
_current_velocity: Vector3<f32>,
time: f32,
) -> (Vector3<f32>, Vector3<f32>, f32) {
let relative_time = time - self.start_time;
let mut segment_start_time = 0.0;
let mut current_segment = 0;
for (i, &segment_duration) in self.times.iter().enumerate() {
if relative_time < segment_start_time + segment_duration {
current_segment = i;
break;
}
segment_start_time += segment_duration;
}
let segment_time = relative_time - segment_start_time;
let (position, velocity, yaw, _yaw_rate) = self.evaluate_polynomial(
segment_time,
&self.coefficients[current_segment],
&self.yaw_coefficients[current_segment],
);
(position, velocity, yaw)
}
fn is_finished(
&self,
current_position: Vector3<f32>,
time: f32,
) -> Result<bool, SimulationError> {
let last_waypoint = self.waypoints.last().ok_or(SimulationError::OtherError(
"No waypoints available".to_string(),
))?;
Ok(time >= self.start_time + self.times.iter().sum::<f32>()
&& (current_position - last_waypoint).norm() < 0.1)
}
}
pub struct QPpolyTrajPlanner {
pub coeff: DMatrix<f64>,
pub polyorder: usize,
pub min_deriv: usize,
pub smooth_upto: usize,
pub segment_times: Vec<f32>,
pub waypoints: Vec<Vec<f32>>,
pub max_velocity: f32,
pub max_acceleration: f32,
pub start_time: f32,
pub dt: f32,
}
impl QPpolyTrajPlanner {
#[allow(clippy::too_many_arguments)]
pub fn new(
waypoints: Vec<Vec<f32>>,
segment_times: Vec<f32>,
polyorder: usize,
min_deriv: usize,
smooth_upto: usize,
max_velocity: f32,
max_acceleration: f32,
start_time: f32,
dt: f32,
) -> Result<Self, SimulationError> {
if waypoints.len() < 2 {
return Err(SimulationError::OtherError(
"At least two waypoints are required".to_string(),
));
}
if waypoints.len() != segment_times.len() + 1 || waypoints[0].len() != 4 {
return Err(SimulationError::OtherError("Number of segment times must be one less than number of waypoints, and waypoints must contain x, y, z, yaw values".to_string()));
}
if smooth_upto > polyorder {
return Err(SimulationError::OtherError(format!(
"smooth_upto ({smooth_upto}) cannot be greater than polynomial order ({polyorder})"
)));
}
let mut planner = Self {
coeff: DMatrix::zeros(segment_times.len() * polyorder, 4),
polyorder,
min_deriv,
smooth_upto,
segment_times,
waypoints,
max_velocity,
max_acceleration,
start_time,
dt,
};
planner.compute_coefficients()?;
Ok(planner)
}
pub fn evaluate_polynomial(
&self,
t: f32,
segment: usize,
) -> (Vector3<f32>, Vector3<f32>, f32, f32) {
let mut position = Vector3::zeros();
let mut velocity = Vector3::zeros();
let pos_basis = self.basis(t, 0);
let vel_basis = self.basis(t, 1);
let row_offset = self.polyorder * segment;
let coeff_segment = &self.coeff.view((row_offset, 0), (self.polyorder, 4));
let pose_4d = pos_basis.transpose() * coeff_segment;
let vel_4d = vel_basis.transpose() * coeff_segment;
let pose_4d_f32 = pose_4d.map(|x| x as f32);
let vel_4d_f32 = vel_4d.map(|x| x as f32);
for i in 0..3 {
position[i] = pose_4d_f32[(0, i)];
velocity[i] = vel_4d_f32[(0, i)];
}
let yaw = pose_4d_f32[(0, 3)];
let yaw_rate = vel_4d_f32[(0, 3)];
(position, velocity, yaw, yaw_rate)
}
fn compute_coefficients(&mut self) -> Result<(), SimulationError> {
let settings: Settings = Settings::default()
.max_iter(10000000)
.eps_abs(1e-6)
.eps_rel(5e-6)
.verbose(false);
let (p, q, a, l, u) = self.prepare_qp_problem();
let mut problem = Problem::new(p, q.as_slice(), a, l.as_slice(), u.as_slice(), &settings)
.map_err(|_| {
SimulationError::OSQPError("Failed to set up OSQP problem".to_string())
})?;
let result = problem.solve();
let coeff_vector = result
.x()
.map(|solution| DVector::from_iterator(solution.len(), solution.iter().cloned()))
.ok_or(SimulationError::OSQPError(
"Unable to solve OSQP problem".to_string(),
))?;
self.coeff = DMatrix::from_column_slice(
self.segment_times.len() * self.polyorder,
4,
coeff_vector.as_slice(),
);
Ok(())
}
fn prepare_qp_problem(
&self,
) -> (
CscMatrix,
DVector<f64>,
CscMatrix,
DVector<f64>,
DVector<f64>,
) {
let num_dims = 4;
let num_vars_per_dim = self.polyorder * self.segment_times.len();
let num_vars = num_vars_per_dim * num_dims;
let q_x = self.generate_obj_function();
let mut q = DMatrix::zeros(num_vars, num_vars);
for dim in 0..num_dims {
let row_offset = dim * num_vars_per_dim;
let col_offset = dim * num_vars_per_dim;
q.view_mut((row_offset, col_offset), (q_x.nrows(), q_x.ncols()))
.copy_from(&q_x);
}
let num_constraints_per_block = (2 + self.smooth_upto) * (self.segment_times.len() + 1) - 2;
let mut a_eq = DMatrix::zeros(num_constraints_per_block * num_dims, num_vars);
let mut b_eq = DVector::zeros(num_constraints_per_block * num_dims);
for dim in 0..num_dims {
let row_offset = dim * num_constraints_per_block;
let col_offset = dim * num_vars_per_dim;
let (small_a_eq, small_b_eq) = self.generate_equality_constraints(dim);
a_eq.view_mut(
(row_offset, col_offset),
(num_constraints_per_block, num_vars_per_dim),
)
.copy_from(&small_a_eq);
b_eq.rows_mut(row_offset, num_constraints_per_block)
.copy_from(&small_b_eq);
}
let mut a_ineq: DMatrix<f64>;
let mut d: DVector<f64>;
let mut f: DVector<f64>;
if self.max_acceleration > 0.0 && self.max_velocity > 0.0 {
let (a_ineq_x, d_x, f_x) = self.generate_inequality_constraints();
a_ineq = DMatrix::zeros(a_ineq_x.nrows() * num_dims, num_vars);
d = DVector::zeros(d_x.len() * num_dims);
f = DVector::zeros(f_x.len() * num_dims);
for dim in 0..num_dims {
let row_offset = dim * a_ineq_x.nrows();
let col_offset = dim * num_vars_per_dim;
a_ineq
.view_mut(
(row_offset, col_offset),
(a_ineq_x.nrows(), a_ineq_x.ncols()),
)
.copy_from(&a_ineq_x);
d.rows_mut(row_offset, d_x.len()).copy_from(&d_x);
f.rows_mut(row_offset, f_x.len()).copy_from(&f_x);
}
} else {
a_ineq = DMatrix::zeros(0, num_vars);
d = DVector::zeros(0);
f = DVector::zeros(0);
}
let total_constraints = a_eq.nrows() + a_ineq.nrows();
let mut a = DMatrix::zeros(total_constraints, num_vars);
a.view_mut((0, 0), (a_eq.nrows(), a_eq.ncols()))
.copy_from(&a_eq);
let mut l = DVector::zeros(total_constraints);
let mut u = DVector::zeros(total_constraints);
l.rows_mut(0, b_eq.len()).copy_from(&b_eq);
u.rows_mut(0, b_eq.len()).copy_from(&b_eq);
if a_ineq.nrows() > 0 {
a.view_mut((a_eq.nrows(), 0), (a_ineq.nrows(), a_ineq.ncols()))
.copy_from(&a_ineq);
l.rows_mut(b_eq.len(), d.len()).copy_from(&d);
u.rows_mut(b_eq.len(), f.len()).copy_from(&f);
}
let q_csc = self.convert_dense_to_sparse(&q).into_upper_tri();
let a_csc = self.convert_dense_to_sparse(&a);
(q_csc, DVector::zeros(num_vars), a_csc, l, u)
}
fn convert_dense_to_sparse(&self, a: &DMatrix<f64>) -> CscMatrix {
let (rows, cols) = a.shape();
let column_major_iter: Vec<f64> = a
.column_iter()
.flat_map(|col| col.iter().copied().collect::<Vec<f64>>())
.collect();
CscMatrix::from_column_iter(rows, cols, column_major_iter)
}
fn generate_obj_function(&self) -> DMatrix<f64> {
let num_segments = self.segment_times.len();
let coeff_num = self.polyorder * num_segments;
let mut q_total: DMatrix<f64> = DMatrix::zeros(coeff_num, coeff_num);
for (i, &segment_time) in self.segment_times.iter().enumerate() {
let q_segment = self.generate_q(self.min_deriv, segment_time);
q_total
.view_mut(
(i * self.polyorder, i * self.polyorder),
(self.polyorder, self.polyorder),
)
.copy_from(&q_segment);
}
q_total
}
fn generate_q(&self, min_deriv: usize, time: f32) -> DMatrix<f64> {
let mut q: DMatrix<f64> = DMatrix::zeros(self.polyorder, self.polyorder);
let poly = self.basis(time, min_deriv);
for i in 0..self.polyorder {
for j in 0..self.polyorder {
let power_order = ((i + j) as f64 - (2 * min_deriv) as f64 + 1.0).max(1.0); q[(i, j)] = poly[i] * poly[j] * time as f64 / power_order;
}
}
q
}
fn generate_equality_constraints(&self, dimension: usize) -> (DMatrix<f64>, DVector<f64>) {
let num_segments = self.segment_times.len();
let num_coeff = self.polyorder * num_segments;
let num_constraints = (2 + self.smooth_upto) * (num_segments + 1) - 2;
let mut a: DMatrix<f64> = DMatrix::zeros(num_constraints, num_coeff);
let mut b: DVector<f64> = DVector::zeros(num_constraints);
let mut row_index: usize = 0;
for i in 0..=num_segments {
let col_offset = i * self.polyorder;
match i {
0 => {
let row = self.basis(0.0, 0);
a.row_mut(row_index)
.columns_mut(col_offset, self.polyorder)
.copy_from(&row.transpose());
b[row_index] = self.waypoints[i][dimension] as f64;
row_index += 1;
for j in 1..=self.smooth_upto {
let row = self.basis(0.0, j);
a.row_mut(row_index)
.columns_mut(col_offset, self.polyorder)
.copy_from(&row.transpose());
b[row_index] = 0.0;
row_index += 1;
}
}
n if n == num_segments => {
let time = self.segment_times[num_segments - 1];
let row = self.basis(time, 0);
a.row_mut(row_index)
.columns_mut(col_offset - self.polyorder, self.polyorder)
.copy_from(&row.transpose());
b[row_index] = self.waypoints[i][dimension] as f64;
row_index += 1;
for j in 1..=self.smooth_upto {
let row = self.basis(time, j);
a.row_mut(row_index)
.columns_mut(col_offset - self.polyorder, self.polyorder)
.copy_from(&row.transpose());
b[row_index] = 0.0;
row_index += 1;
}
}
_ => {
{
let segment_prev = self.basis(self.segment_times[i - 1], 0);
let segment_next = self.basis(0.0, 0);
let prev_col_offset = (i - 1) * self.polyorder;
a.row_mut(row_index)
.columns_mut(prev_col_offset, self.polyorder)
.copy_from(&segment_prev.transpose());
b[row_index] = self.waypoints[i][dimension] as f64;
row_index += 1;
a.row_mut(row_index)
.columns_mut(col_offset, self.polyorder)
.copy_from(&segment_next.transpose());
b[row_index] = self.waypoints[i][dimension] as f64;
row_index += 1;
for j in 1..=self.smooth_upto {
let segment_prev = self.basis(self.segment_times[i - 1], j);
let segment_next = self.basis(0.0, j);
a.row_mut(row_index)
.columns_mut(prev_col_offset, self.polyorder)
.copy_from(&segment_prev.transpose());
a.row_mut(row_index)
.columns_mut(col_offset, self.polyorder)
.add_assign(-segment_next.transpose());
b[row_index] = 0.0;
row_index += 1;
}
}
}
}
}
(a, b)
}
fn compute_num_constraints(&self) -> usize {
if self.max_velocity == 0.0 || self.max_acceleration == 0.0 {
return 0;
}
let total_constraints: usize = self
.segment_times
.iter()
.map(|&segment_time| {
let num_timesteps = (segment_time / self.dt).floor() as usize;
num_timesteps * 2
})
.sum();
total_constraints
}
fn generate_inequality_constraints(&self) -> (DMatrix<f64>, DVector<f64>, DVector<f64>) {
let num_segments = self.segment_times.len();
let coeff_num = self.polyorder * num_segments;
let num_constraints = self.compute_num_constraints();
let mut c: DMatrix<f64> = DMatrix::zeros(num_constraints, coeff_num);
let mut d = DVector::from_element(num_constraints, 0.0); let mut f = DVector::from_element(num_constraints, 0.0);
let mut row_index = 0;
for (i, &segment_time) in self.segment_times.iter().enumerate() {
let col_offset = i * self.polyorder;
let mut t = self.dt; while t < segment_time {
let row_velocity = self.basis(t, 1);
c.row_mut(row_index)
.columns_mut(col_offset, self.polyorder)
.copy_from(&row_velocity.transpose());
f[row_index] = self.max_velocity as f64;
d[row_index] = -self.max_velocity as f64;
row_index += 1;
let row_acceleration = self.basis(t, 2);
c.row_mut(row_index)
.columns_mut(col_offset, self.polyorder)
.copy_from(&row_acceleration.transpose());
f[row_index] = self.max_acceleration as f64;
d[row_index] = -self.max_acceleration as f64;
row_index += 1;
t += self.dt; }
}
(c, d, f)
}
fn basis(&self, time: f32, derivative: usize) -> DVector<f64> {
assert!(derivative <= 4, "Derivative order must be less than 4");
let time_f64 = time as f64;
let t: Vec<f64> = (0..self.polyorder)
.map(|i| time_f64.powi(i as i32))
.collect();
let mut b = DVector::from_element(self.polyorder, 0.0);
for i in 0..self.polyorder {
let power = i as isize - derivative as isize;
let mut coeff = 1.0;
if power < 0 {
b[i] = 0.0;
} else {
for j in (i - derivative + 1..=i).rev() {
coeff *= j as f64;
}
b[i] = coeff * t[power as usize];
}
}
b
}
}
impl Planner for QPpolyTrajPlanner {
fn plan(
&self,
_current_position: Vector3<f32>,
_current_velocity: Vector3<f32>,
time: f32,
) -> (Vector3<f32>, Vector3<f32>, f32) {
let relative_time = time - self.start_time;
let mut segment_start_time = 0.0;
let mut current_segment = 0;
for (i, &segment_duration) in self.segment_times.iter().enumerate() {
if relative_time < segment_start_time + segment_duration {
current_segment = i;
break;
}
segment_start_time += segment_duration;
}
let segment_time = relative_time - segment_start_time;
let (position, velocity, yaw, _yaw_rate) =
self.evaluate_polynomial(segment_time, current_segment);
(position, velocity, yaw)
}
fn is_finished(
&self,
current_position: Vector3<f32>,
time: f32,
) -> Result<bool, SimulationError> {
let last_waypoint = self.waypoints.last().ok_or(SimulationError::OtherError(
"No waypoints available".to_string(),
))?;
let total_time: f32 = self.segment_times.iter().sum();
let last_position = Vector3::new(last_waypoint[0], last_waypoint[1], last_waypoint[2]);
Ok(time >= self.start_time + total_time && (current_position - last_position).norm() < 0.1)
}
}
pub struct PlannerStepConfig {
pub step: usize,
pub planner_type: String,
pub params: serde_yaml::Value,
}
pub fn update_planner(
planner_manager: &mut PlannerManager,
step: usize,
time: f32,
simulation_frequency: usize,
quad: &Quadrotor,
obstacles: &[Obstacle],
planner_config: &[PlannerStepConfig],
) -> Result<(), SimulationError> {
if let Some(planner_step) = planner_config
.iter()
.find(|s| s.step * simulation_frequency == step * 1000)
{
log::info!("Time: {:.2} s,\tSwitch {}", time, planner_step.planner_type);
planner_manager.set_planner(create_planner(planner_step, quad, time, obstacles)?);
}
Ok(())
}
pub fn create_planner(
step: &PlannerStepConfig,
quad: &Quadrotor,
time: f32,
obstacles: &[Obstacle],
) -> Result<PlannerType, SimulationError> {
let params = &step.params;
match step.planner_type.as_str() {
"MinimumJerkLine" => Ok(PlannerType::MinimumJerkLine(MinimumJerkLinePlanner {
start_position: quad.position,
end_position: parse_vector3(params, "end_position")?,
start_yaw: quad.orientation.euler_angles().2,
end_yaw: parse_f32(params, "end_yaw")?,
start_time: time,
duration: parse_f32(params, "duration")?,
})),
"Lissajous" => Ok(PlannerType::Lissajous(LissajousPlanner {
start_position: quad.position,
center: parse_vector3(params, "center")?,
amplitude: parse_vector3(params, "amplitude")?,
frequency: parse_vector3(params, "frequency")?,
phase: parse_vector3(params, "phase")?,
start_time: time,
duration: parse_f32(params, "duration")?,
start_yaw: quad.orientation.euler_angles().2,
end_yaw: parse_f32(params, "end_yaw")?,
ramp_time: parse_f32(params, "ramp_time")?,
})),
"Circle" => Ok(PlannerType::Circle(CirclePlanner {
center: parse_vector3(params, "center")?,
radius: parse_f32(params, "radius")?,
angular_velocity: parse_f32(params, "angular_velocity")?,
start_position: quad.position,
start_time: time,
duration: parse_f32(params, "duration")?,
start_yaw: quad.orientation.euler_angles().2,
end_yaw: quad.orientation.euler_angles().2,
ramp_time: parse_f32(params, "ramp_time")?,
})),
"ObstacleAvoidance" => Ok(PlannerType::ObstacleAvoidance(ObstacleAvoidancePlanner {
target_position: parse_vector3(params, "target_position")?,
start_time: time,
duration: parse_f32(params, "duration")?,
start_yaw: quad.orientation.euler_angles().2,
end_yaw: parse_f32(params, "end_yaw")?,
obstacles: obstacles.to_owned(),
k_att: parse_f32(params, "k_att")?,
k_rep: parse_f32(params, "k_rep")?,
k_vortex: parse_f32(params, "k_vortex")?,
d0: parse_f32(params, "d0")?,
d_target: parse_f32(params, "d_target")?,
max_speed: parse_f32(params, "max_speed")?,
})),
"MinimumSnapWaypoint" => {
let mut waypoints = vec![quad.position];
waypoints.extend(
params["waypoints"]
.as_sequence()
.ok_or_else(|| SimulationError::OtherError("Invalid waypoints".to_string()))?
.iter()
.map(|w| {
w.as_sequence()
.and_then(|coords| {
Some(Vector3::new(
coords[0].as_f64()? as f32,
coords[1].as_f64()? as f32,
coords[2].as_f64()? as f32,
))
})
.ok_or(SimulationError::OtherError("Invalid waypoint".to_string()))
})
.collect::<Result<Vec<Vector3<f32>>, SimulationError>>()?,
);
let mut yaws = vec![quad.orientation.euler_angles().2];
yaws.extend(
params["yaws"]
.as_sequence()
.ok_or(SimulationError::OtherError("Invalid yaws".to_string()))?
.iter()
.map(|y| {
y.as_f64()
.map(|v| v as f32)
.ok_or(SimulationError::OtherError("Invalid yaw".to_string()))
})
.collect::<Result<Vec<f32>, SimulationError>>()?,
);
let segment_times = params["segment_times"]
.as_sequence()
.ok_or_else(|| SimulationError::OtherError("Invalid segment_times".to_string()))?
.iter()
.map(|t| {
t.as_f64().map(|v| v as f32).ok_or_else(|| {
SimulationError::OtherError("Invalid segment time".to_string())
})
})
.collect::<Result<Vec<f32>, SimulationError>>()?;
MinimumSnapWaypointPlanner::new(waypoints, yaws, segment_times, time)
.map(PlannerType::MinimumSnapWaypoint)
}
"QPpolyTraj" => {
let mut waypoints = vec![vec![
quad.position.x,
quad.position.y,
quad.position.z,
quad.orientation.euler_angles().2,
]];
waypoints.extend(
params["waypoints"]
.as_sequence()
.ok_or_else(|| SimulationError::OtherError("Invalid waypoints".to_string()))?
.iter()
.map(|w| {
w.as_sequence()
.and_then(|coords| {
Some(vec![
coords[0].as_f64()? as f32,
coords[1].as_f64()? as f32,
coords[2].as_f64()? as f32,
coords[3].as_f64()? as f32,
])
})
.ok_or(SimulationError::OtherError("Invalid waypoint".to_string()))
})
.collect::<Result<Vec<Vec<f32>>, SimulationError>>()?,
);
let segment_times = params["segment_times"]
.as_sequence()
.ok_or_else(|| SimulationError::OtherError("Invalid segment_times".to_string()))?
.iter()
.map(|t| {
t.as_f64().map(|v| v as f32).ok_or_else(|| {
SimulationError::OtherError("Invalid segment time".to_string())
})
})
.collect::<Result<Vec<f32>, SimulationError>>()?;
QPpolyTrajPlanner::new(
waypoints,
segment_times,
parse_usize(params, "polynomial_order")?,
parse_usize(params, "minimize_derivative")?,
parse_usize(params, "smooth_upto")?,
parse_f32(params, "max_velocity")?,
parse_f32(params, "max_acceleration")?,
time,
parse_f32(params, "dt")?,
)
.map(PlannerType::QPpolyTraj)
}
"Landing" => Ok(PlannerType::Landing(LandingPlanner {
start_position: quad.position,
start_time: time,
duration: parse_f32(params, "duration")?,
start_yaw: quad.orientation.euler_angles().2,
})),
_ => Err(SimulationError::OtherError(format!(
"Unknown planner type: {}",
step.planner_type
))),
}
}
pub fn parse_vector3(
value: &serde_yaml::Value,
key: &str,
) -> Result<Vector3<f32>, SimulationError> {
value[key]
.as_sequence()
.and_then(|seq| {
if seq.len() == 3 {
Some(Vector3::new(
seq[0].as_f64()? as f32,
seq[1].as_f64()? as f32,
seq[2].as_f64()? as f32,
))
} else {
None
}
})
.ok_or_else(|| SimulationError::OtherError(format!("Invalid {key} vector")))
}
pub fn parse_f32(value: &serde_yaml::Value, key: &str) -> Result<f32, SimulationError> {
value[key]
.as_f64()
.map(|v| v as f32)
.ok_or_else(|| SimulationError::OtherError(format!("Invalid {key}")))
}
pub fn parse_usize(value: &serde_yaml::Value, key: &str) -> Result<usize, SimulationError> {
value[key]
.as_i64()
.and_then(|v| if v >= 0 { Some(v as usize) } else { None }) .ok_or_else(|| SimulationError::OtherError(format!("Invalid {key}")))
}
#[derive(Clone)]
pub struct Obstacle {
pub position: Vector3<f32>,
pub velocity: Vector3<f32>,
pub radius: f32,
}
impl Obstacle {
pub fn new(position: Vector3<f32>, velocity: Vector3<f32>, radius: f32) -> Self {
Self {
position,
velocity,
radius,
}
}
}
pub struct Maze {
pub lower_bounds: [f32; 3],
pub upper_bounds: [f32; 3],
pub obstacles: Vec<Obstacle>,
pub obstacles_velocity_bounds: [f32; 3],
pub obstacles_radius_bounds: [f32; 2],
pub rng: ChaCha8Rng,
}
impl Maze {
pub fn new(
lower_bounds: [f32; 3],
upper_bounds: [f32; 3],
num_obstacles: usize,
obstacles_velocity_bounds: [f32; 3],
obstacles_radius_bounds: [f32; 2],
) -> Self {
let mut maze = Maze {
lower_bounds,
upper_bounds,
obstacles: Vec::new(),
obstacles_velocity_bounds,
obstacles_radius_bounds,
rng: ChaCha8Rng::from_os_rng(),
};
maze.generate_obstacles(num_obstacles);
maze
}
pub fn generate_obstacles(&mut self, num_obstacles: usize) {
self.obstacles = (0..num_obstacles)
.map(|_| {
let position = Vector3::new(
rand::Rng::random_range(
&mut self.rng,
self.lower_bounds[0]..self.upper_bounds[0],
),
rand::Rng::random_range(
&mut self.rng,
self.lower_bounds[1]..self.upper_bounds[1],
),
rand::Rng::random_range(
&mut self.rng,
self.lower_bounds[2]..self.upper_bounds[2],
),
);
let v_bounds = self.obstacles_velocity_bounds;
let r_bounds = self.obstacles_radius_bounds;
let velocity = Vector3::new(
rand::Rng::random_range(&mut self.rng, -v_bounds[0]..v_bounds[0]),
rand::Rng::random_range(&mut self.rng, -v_bounds[1]..v_bounds[1]),
rand::Rng::random_range(&mut self.rng, -v_bounds[2]..v_bounds[2]),
);
let radius = rand::Rng::random_range(&mut self.rng, r_bounds[0]..r_bounds[1]);
Obstacle::new(position, velocity, radius)
})
.collect();
}
pub fn update_obstacles(&mut self, dt: f32) {
self.obstacles.iter_mut().for_each(|obstacle| {
obstacle.position += obstacle.velocity * dt;
for i in 0..3 {
if obstacle.position[i] - obstacle.radius < self.lower_bounds[i]
|| obstacle.position[i] + obstacle.radius > self.upper_bounds[i]
{
obstacle.velocity[i] *= -1.0;
}
}
});
}
}
pub struct Camera {
pub resolution: (usize, usize),
pub fov_vertical: f32,
pub fov_horizontal: f32,
pub vertical_focal_length: f32,
pub horizontal_focal_length: f32,
pub near: f32,
pub far: f32,
pub aspect_ratio: f32,
pub ray_directions: Vec<Vector3<f32>>,
pub depth_buffer: Vec<f32>,
}
impl Camera {
pub fn new(resolution: (usize, usize), fov_vertical: f32, near: f32, far: f32) -> Self {
let (width, height) = resolution;
let (aspect_ratio, tan_half_fov) =
(width as f32 / height as f32, (fov_vertical / 2.0).tan());
let mut ray_directions = Vec::with_capacity(width * height);
for y in 0..height {
for x in 0..width {
let x_ndc = (2.0 * x as f32 / width as f32 - 1.0) * aspect_ratio * tan_half_fov;
let y_ndc = (1.0 - 2.0 * y as f32 / height as f32) * tan_half_fov;
ray_directions.push(Vector3::new(1.0, x_ndc, y_ndc).normalize());
}
}
let fov_horizontal =
(width as f32 / height as f32 * (fov_vertical / 2.0).tan()).atan() * 2.0;
let horizontal_focal_length = (width as f32 / 2.0) / ((fov_horizontal / 2.0).tan());
let vertical_focal_length = (height as f32 / 2.0) / ((fov_vertical / 2.0).tan());
let depth_buffer = vec![0.0; width * height];
Self {
resolution,
fov_vertical,
fov_horizontal,
vertical_focal_length,
horizontal_focal_length,
near,
far,
aspect_ratio,
ray_directions,
depth_buffer,
}
}
pub fn render_depth(
&mut self,
quad_position: &Vector3<f32>,
quad_orientation: &UnitQuaternion<f32>,
maze: &Maze,
use_multi_threading: bool,
) -> Result<(), SimulationError> {
let (width, height) = self.resolution;
let total_pixels = width * height;
let rotation_camera_to_world = quad_orientation.to_rotation_matrix().matrix()
* Matrix3::new(1.0, 0.0, 0.0, 0.0, -1.0, 0.0, 0.0, 0.0, 1.0);
let rotation_world_to_camera = rotation_camera_to_world.transpose();
const CHUNK_SIZE: usize = 64;
if use_multi_threading {
self.depth_buffer
.par_chunks_mut(CHUNK_SIZE)
.enumerate()
.try_for_each(|(chunk_idx, chunk)| {
let start_idx = chunk_idx * CHUNK_SIZE;
for (i, depth) in chunk.iter_mut().enumerate() {
let ray_idx = start_idx + i;
if ray_idx >= total_pixels {
break;
}
*depth = ray_cast(
quad_position,
&rotation_world_to_camera,
&(rotation_camera_to_world * self.ray_directions[ray_idx]),
maze,
self.near,
self.far,
)?;
}
Ok::<(), SimulationError>(())
})?;
} else {
for i in 0..total_pixels {
self.depth_buffer[i] = ray_cast(
quad_position,
&rotation_world_to_camera,
&(rotation_camera_to_world * self.ray_directions[i]),
maze,
self.near,
self.far,
)?;
}
}
Ok(())
}
}
pub fn ray_cast(
origin: &Vector3<f32>,
rotation_world_to_camera: &Matrix3<f32>,
direction: &Vector3<f32>,
maze: &Maze,
near: f32,
far: f32,
) -> Result<f32, SimulationError> {
let mut closest_hit = far;
for axis in 0..3 {
let t_near = if direction[axis] > 0.0 {
(maze.upper_bounds[axis] - origin[axis]) / direction[axis]
} else {
(maze.lower_bounds[axis] - origin[axis]) / direction[axis]
};
if t_near > near && t_near < closest_hit {
let intersection_point = origin + direction * t_near;
let mut valid = true;
for i in 0..3 {
if i != axis
&& (intersection_point[i] < maze.lower_bounds[i]
|| intersection_point[i] > maze.upper_bounds[i])
{
valid = false;
break;
}
}
if valid {
closest_hit = t_near;
}
}
}
if closest_hit <= near {
return Ok(f32::INFINITY);
}
for obstacle in &maze.obstacles {
let oc = origin - obstacle.position;
let b = oc.dot(direction);
let c = oc.dot(&oc) - obstacle.radius * obstacle.radius;
let discriminant = b * b - c;
if discriminant >= 0.0 {
let t = -b - fast_sqrt(discriminant);
if t > near && t < closest_hit {
closest_hit = t;
}
}
}
if closest_hit < far {
Ok((rotation_world_to_camera * direction * closest_hit).x)
} else {
Ok(f32::INFINITY)
}
}
pub fn log_data(
rec: &rerun::RecordingStream,
quad: &Quadrotor,
desired_position: &Vector3<f32>,
desired_velocity: &Vector3<f32>,
measured_accel: &Vector3<f32>,
measured_gyro: &Vector3<f32>,
) -> Result<(), SimulationError> {
rec.log(
"world/quad/desired_position",
&rerun::Points3D::new([(desired_position.x, desired_position.y, desired_position.z)])
.with_radii([0.1])
.with_colors([rerun::Color::from_rgb(255, 255, 255)]),
)?;
rec.log(
"world/quad/base_link",
&rerun::Transform3D::from_translation_rotation(
rerun::Vec3D::new(quad.position.x, quad.position.y, quad.position.z),
rerun::Quaternion::from_xyzw([
quad.orientation.i,
quad.orientation.j,
quad.orientation.k,
quad.orientation.w,
]),
)
.with_axis_length(0.7),
)?;
let (quad_roll, quad_pitch, quad_yaw) = quad.orientation.euler_angles();
let quad_euler_angles: Vector3<f32> = Vector3::new(quad_roll, quad_pitch, quad_yaw);
for (pre, vec) in [
("position", &quad.position),
("velocity", &quad.velocity),
("accel", measured_accel),
("orientation", &quad_euler_angles),
("gyro", measured_gyro),
("desired_position", desired_position),
("desired_velocity", desired_velocity),
] {
for (i, a) in ["x", "y", "z"].iter().enumerate() {
rec.log(format!("{pre}/{a}"), &rerun::Scalars::single(vec[i] as f64))?;
}
}
Ok(())
}
pub fn log_maze_tube(rec: &rerun::RecordingStream, maze: &Maze) -> Result<(), SimulationError> {
let (lower_bounds, upper_bounds) = (maze.lower_bounds, maze.upper_bounds);
let center_position = rerun::external::glam::Vec3::new(
(lower_bounds[0] + upper_bounds[0]) / 2.0,
(lower_bounds[1] + upper_bounds[1]) / 2.0,
(lower_bounds[2] + upper_bounds[2]) / 2.0,
);
let half_sizes = rerun::external::glam::Vec3::new(
(upper_bounds[0] - lower_bounds[0]) / 2.0,
(upper_bounds[1] - lower_bounds[1]) / 2.0,
(upper_bounds[2] - lower_bounds[2]) / 2.0,
);
rec.log(
"world/maze/tube",
&rerun::Boxes3D::from_centers_and_half_sizes([center_position], [half_sizes])
.with_colors([rerun::Color::from_rgb(128, 128, 255)]),
)?;
Ok(())
}
pub fn log_maze_obstacles(
rec: &rerun::RecordingStream,
maze: &Maze,
) -> Result<(), SimulationError> {
let (positions, radii): (Vec<(f32, f32, f32)>, Vec<f32>) = maze
.obstacles
.iter()
.map(|obstacle| {
let pos = obstacle.position;
((pos.x, pos.y, pos.z), obstacle.radius)
})
.unzip();
rec.log(
"world/maze/obstacles",
&rerun::Points3D::new(positions)
.with_radii(radii)
.with_colors([rerun::Color::from_rgb(255, 128, 128)]),
)?;
Ok(())
}
pub struct Trajectory {
pub points: Vec<Vector3<f32>>,
pub last_logged_point: Vector3<f32>,
pub min_distance_threadhold: f32,
}
impl Trajectory {
pub fn new(initial_point: Vector3<f32>) -> Self {
Self {
points: vec![initial_point],
last_logged_point: initial_point,
min_distance_threadhold: 0.05,
}
}
pub fn add_point(&mut self, point: Vector3<f32>) -> bool {
if (point - self.last_logged_point).norm() > self.min_distance_threadhold {
self.points.push(point);
self.last_logged_point = point;
true
} else {
false
}
}
}
pub fn log_trajectory(
rec: &rerun::RecordingStream,
trajectory: &Trajectory,
) -> Result<(), SimulationError> {
let path = trajectory
.points
.iter()
.map(|p| (p.x, p.y, p.z))
.collect::<Vec<(f32, f32, f32)>>();
rec.log(
"world/quad/path",
&rerun::LineStrips3D::new([path]).with_colors([rerun::Color::from_rgb(0, 255, 255)]),
)?;
Ok(())
}
pub fn log_depth_image(
rec: &rerun::RecordingStream,
cam: &Camera,
use_multi_threading: bool,
) -> Result<(), SimulationError> {
let (width, height) = (cam.resolution.0, cam.resolution.1);
let (min_depth, max_depth) = (cam.near, cam.far);
let depth_image = &cam.depth_buffer;
let mut image: rerun::external::ndarray::Array<u8, _> =
rerun::external::ndarray::Array::zeros((height, width, 3));
let depth_range = max_depth - min_depth;
let scale_factor = 255.0 / depth_range;
if use_multi_threading {
const CHUNK_SIZE: usize = 32;
image
.as_slice_mut()
.expect("Failed to get mutable slice of image")
.par_chunks_exact_mut(CHUNK_SIZE * 3)
.enumerate()
.for_each(|(chunk_index, chunk)| {
let start_index = chunk_index * 32;
for (i, pixel) in chunk.chunks_exact_mut(3).enumerate() {
let idx = start_index + i;
let depth = depth_image[idx];
let color = if depth.is_finite() {
let normalized_depth =
((depth - min_depth) * scale_factor).clamp(0.0, 255.0);
color_map_fn(normalized_depth)
} else {
(0, 0, 0)
};
(pixel[0], pixel[1], pixel[2]) = color;
}
});
} else {
for (index, pixel) in image
.as_slice_mut()
.expect("Failed to get mutable slice of image")
.chunks_exact_mut(3)
.enumerate()
{
let depth = depth_image[index];
let color = if depth.is_finite() {
let normalized_depth = ((depth - min_depth) * scale_factor).clamp(0.0, 255.0);
color_map_fn(normalized_depth)
} else {
(0, 0, 0)
};
(pixel[0], pixel[1], pixel[2]) = color;
}
}
let rerun_image = rerun::Image::from_color_model_and_tensor(rerun::ColorModel::RGB, image)
.map_err(|e| SimulationError::OtherError(format!("Failed to create rerun image: {e}")))?;
rec.log("world/quad/cam/depth", &rerun_image)?;
Ok(())
}
pub fn log_pinhole_depth(
rec: &rerun::RecordingStream,
cam: &Camera,
cam_position: Vector3<f32>,
cam_orientation: UnitQuaternion<f32>,
cam_transform: [f32; 9],
) -> Result<(), SimulationError> {
let depth_image = &cam.depth_buffer;
let (width, height) = cam.resolution;
let pinhole_camera = rerun::Pinhole::from_focal_length_and_resolution(
(cam.horizontal_focal_length, cam.vertical_focal_length),
(width as f32, height as f32),
)
.with_camera_xyz(rerun::components::ViewCoordinates::RDF)
.with_resolution((width as f32, height as f32))
.with_principal_point((width as f32 / 2.0, height as f32 / 2.0));
let rotated_camera_orientation = UnitQuaternion::from_rotation_matrix(
&(cam_orientation.to_rotation_matrix()
* Rotation3::from_matrix_unchecked(Matrix3::from_row_slice(&cam_transform))),
);
let cam_transform = rerun::Transform3D::from_translation_rotation(
rerun::Vec3D::new(cam_position.x, cam_position.y, cam_position.z),
rerun::Quaternion::from_xyzw([
rotated_camera_orientation.i,
rotated_camera_orientation.j,
rotated_camera_orientation.k,
rotated_camera_orientation.w,
]),
);
rec.log("world/quad/cam", &cam_transform)?;
rec.log("world/quad/cam", &pinhole_camera)?;
let depth_image_rerun =
rerun::external::ndarray::Array::from_shape_vec((height, width), depth_image.to_vec())
.unwrap();
rec.log(
"world/quad/cam/rerun_depth",
&rerun::DepthImage::try_from(depth_image_rerun)
.unwrap()
.with_meter(1.0),
)?;
Ok(())
}
#[inline]
pub fn color_map_fn(gray: f32) -> (u8, u8, u8) {
let x = gray / 255.0;
let r = (34.61
+ x * (1172.33 - x * (10793.56 - x * (33300.12 - x * (38394.49 - x * 14825.05)))))
.clamp(0.0, 255.0) as u8;
let g = (23.31 + x * (557.33 + x * (1225.33 - x * (3574.96 - x * (1073.77 + x * 707.56)))))
.clamp(0.0, 255.0) as u8;
let b = (27.2 + x * (3211.1 - x * (15327.97 - x * (27814.0 - x * (22569.18 - x * 6838.66)))))
.clamp(0.0, 255.0) as u8;
(r, g, b)
}
#[inline(always)]
fn fast_sqrt(x: f32) -> f32 {
let i = x.to_bits();
let i = 0x1fbd1df5 + (i >> 1);
f32::from_bits(i)
}