Skip to main content

bevy_rapier3d/control/
pid_controller.rs

1use crate::dynamics::Velocity;
2use crate::math::{AngVector, Real, Rot, Vect};
3use crate::utils;
4use bevy::prelude::*;
5use rapier::dynamics::{
6    AxesMask as RapierAxesMask, RigidBody as RapierRigidBody, RigidBodyVelocity,
7};
8
9pub use rapier::control::PdErrors;
10
11/// The axes affected by a [`PdController`] or a [`PidController`].
12///
13/// Only the axes with their flag set are taken into account when computing errors and
14/// corrections.
15#[derive(Copy, Clone, Debug, Default, PartialEq, Eq, Hash, Reflect)]
16#[cfg_attr(feature = "serde-serialize", derive(Serialize, Deserialize))]
17#[reflect(Default, PartialEq, Debug, Hash)]
18pub struct AxesMask(u8);
19
20bitflags::bitflags! {
21    impl AxesMask: u8 {
22        /// The translational X axis.
23        const LIN_X = 1 << 0;
24        /// The translational Y axis.
25        const LIN_Y = 1 << 1;
26        /// The translational Z axis.
27        #[cfg(feature = "dim3")]
28        const LIN_Z = 1 << 2;
29        /// The rotational X axis.
30        #[cfg(feature = "dim3")]
31        const ANG_X = 1 << 3;
32        /// The rotational Y axis.
33        #[cfg(feature = "dim3")]
34        const ANG_Y = 1 << 4;
35        /// The rotational Z axis.
36        const ANG_Z = 1 << 5;
37    }
38}
39
40impl AxesMask {
41    /// All the translational axes.
42    #[cfg(feature = "dim2")]
43    pub const LIN_AXES: Self = Self::LIN_X.union(Self::LIN_Y);
44    /// All the translational axes.
45    #[cfg(feature = "dim3")]
46    pub const LIN_AXES: Self = Self::LIN_X.union(Self::LIN_Y).union(Self::LIN_Z);
47    /// All the rotational axes.
48    #[cfg(feature = "dim2")]
49    pub const ANG_AXES: Self = Self::ANG_Z;
50    /// All the rotational axes.
51    #[cfg(feature = "dim3")]
52    pub const ANG_AXES: Self = Self::ANG_X.union(Self::ANG_Y).union(Self::ANG_Z);
53}
54
55impl From<AxesMask> for RapierAxesMask {
56    fn from(axes: AxesMask) -> RapierAxesMask {
57        RapierAxesMask::from_bits_truncate(axes.bits())
58    }
59}
60
61impl From<RapierAxesMask> for AxesMask {
62    fn from(axes: RapierAxesMask) -> AxesMask {
63        AxesMask::from_bits_truncate(axes.bits())
64    }
65}
66
67/// The world-space pose and velocities a [`PdController`] or [`PidController`] drives its
68/// rigid-body toward.
69///
70/// The target pose is the pose of the rigid-body itself (i.e. its `GlobalTransform` translation
71/// and rotation), not the pose of its center of mass.
72#[derive(Copy, Clone, Debug, Default, PartialEq, Reflect)]
73#[cfg_attr(feature = "serde-serialize", derive(Serialize, Deserialize))]
74#[reflect(Default, PartialEq, Debug)]
75pub struct PidTarget {
76    /// The target world-space translation.
77    pub translation: Vect,
78    /// The target world-space rotation.
79    pub rotation: Rot,
80    /// The target linear velocity.
81    pub linvel: Vect,
82    /// The target angular velocity.
83    pub angvel: AngVector,
84}
85
86impl PidTarget {
87    /// A target with the given pose and zero velocities.
88    pub fn new(translation: Vect, rotation: Rot) -> Self {
89        Self {
90            translation,
91            rotation,
92            ..Default::default()
93        }
94    }
95
96    /// A target with the given translation, an identity rotation and zero velocities.
97    pub fn from_translation(translation: Vect) -> Self {
98        Self {
99            translation,
100            ..Default::default()
101        }
102    }
103
104    /// Sets the target velocities.
105    pub fn with_velocities(mut self, linvel: Vect, angvel: AngVector) -> Self {
106        self.linvel = linvel;
107        self.angvel = angvel;
108        self
109    }
110
111    fn to_rapier(self) -> (rapier::math::Pose, RigidBodyVelocity<Real>) {
112        (
113            utils::pose_from(self.translation, self.rotation),
114            RigidBodyVelocity {
115                linvel: self.linvel,
116                angvel: self.angvel,
117            },
118        )
119    }
120}
121
122fn to_velocity(vels: RigidBodyVelocity<Real>) -> Velocity {
123    Velocity {
124        linear: vels.linvel,
125        angular: vels.angvel,
126    }
127}
128
129#[cfg(feature = "dim2")]
130fn splat_ang(value: Real) -> AngVector {
131    value
132}
133
134#[cfg(feature = "dim3")]
135fn splat_ang(value: Real) -> AngVector {
136    AngVector::splat(value)
137}
138
139/// A Proportional-Derivative (PD) controller driving a rigid-body toward a [`PidTarget`].
140///
141/// When attached to an entity with a non-fixed rigid-body, the plugin computes a velocity
142/// correction before each simulation step and adds it to the rigid-body's velocity. This is
143/// the PID controller without its Integral part, which is generally sufficient for games.
144/// Use either this or a [`PidController`] on a given entity, not both.
145#[derive(Copy, Clone, Debug, PartialEq, Component, Reflect)]
146#[cfg_attr(feature = "serde-serialize", derive(Serialize, Deserialize))]
147#[reflect(Component, Default, PartialEq, Debug)]
148pub struct PdController {
149    /// The pose and velocities the rigid-body is driven toward.
150    pub target: PidTarget,
151    /// The Proportional gain applied to the instantaneous linear position errors.
152    ///
153    /// This is usually set to a multiple of the inverse of the simulation timestep
154    /// (e.g. `60` if the timestep is `1.0 / 60.0`).
155    pub lin_kp: Vect,
156    /// The Derivative gain applied to the instantaneous linear velocity errors.
157    ///
158    /// This is usually in `[0.0, 1.0]`, where `1.0` corrects velocity errors in a single step.
159    pub lin_kd: Vect,
160    /// The Proportional gain applied to the instantaneous angular position errors.
161    pub ang_kp: AngVector,
162    /// The Derivative gain applied to the instantaneous angular velocity errors.
163    pub ang_kd: AngVector,
164    /// The axes affected by this controller.
165    pub axes: AxesMask,
166}
167
168impl Default for PdController {
169    fn default() -> Self {
170        Self::from_raw(&rapier::control::PdController::default())
171    }
172}
173
174impl PdController {
175    /// Initializes a PD controller with the same gains on all axes, and a default target.
176    ///
177    /// Only the axes specified in `axes` are controlled.
178    pub fn new(kp: Real, kd: Real, axes: AxesMask) -> Self {
179        Self::from_raw(&rapier::control::PdController::new(kp, kd, axes.into()))
180    }
181
182    /// Sets the target of this controller.
183    pub fn with_target(mut self, target: PidTarget) -> Self {
184        self.target = target;
185        self
186    }
187
188    /// Converts a Rapier PD controller into this component, with a default target.
189    pub fn from_raw(raw: &rapier::control::PdController) -> Self {
190        Self {
191            target: PidTarget::default(),
192            lin_kp: raw.lin_kp,
193            lin_kd: raw.lin_kd,
194            ang_kp: raw.ang_kp,
195            ang_kd: raw.ang_kd,
196            axes: raw.axes.into(),
197        }
198    }
199
200    /// Converts this component into a Rapier PD controller (the target is not part of it).
201    pub fn to_raw(&self) -> rapier::control::PdController {
202        rapier::control::PdController {
203            lin_kp: self.lin_kp,
204            lin_kd: self.lin_kd,
205            ang_kp: self.ang_kp,
206            ang_kd: self.ang_kd,
207            axes: self.axes.into(),
208        }
209    }
210
211    /// Computes the velocity correction from the given position and velocity errors.
212    pub fn correction(&self, pose_errors: &PdErrors, vel_errors: &PdErrors) -> Velocity {
213        to_velocity(self.to_raw().correction(pose_errors, vel_errors))
214    }
215
216    /// Computes the velocity correction driving the given Rapier rigid-body toward [`Self::target`].
217    pub fn rigid_body_correction(&self, body: &RapierRigidBody) -> Velocity {
218        let (pose, vels) = self.target.to_rapier();
219        to_velocity(self.to_raw().rigid_body_correction(body, pose, vels))
220    }
221}
222
223/// A Proportional-Integral-Derivative (PID) controller driving a rigid-body toward a [`PidTarget`].
224///
225/// When attached to an entity with a non-fixed rigid-body, the plugin computes a velocity
226/// correction before each simulation step and adds it to the rigid-body's velocity. The
227/// accumulated errors of the Integral part are updated at each step. Use either this or a
228/// [`PdController`] on a given entity, not both.
229#[derive(Copy, Clone, Debug, PartialEq, Component, Reflect)]
230#[cfg_attr(feature = "serde-serialize", derive(Serialize, Deserialize))]
231#[reflect(Component, Default, PartialEq, Debug)]
232pub struct PidController {
233    /// The pose and velocities the rigid-body is driven toward.
234    pub target: PidTarget,
235    /// The Proportional gain applied to the instantaneous linear position errors.
236    ///
237    /// This is usually set to a multiple of the inverse of the simulation timestep
238    /// (e.g. `60` if the timestep is `1.0 / 60.0`).
239    pub lin_kp: Vect,
240    /// The Integral gain applied to the accumulated linear position errors.
241    pub lin_ki: Vect,
242    /// The Derivative gain applied to the instantaneous linear velocity errors.
243    ///
244    /// This is usually in `[0.0, 1.0]`, where `1.0` corrects velocity errors in a single step.
245    pub lin_kd: Vect,
246    /// The Proportional gain applied to the instantaneous angular position errors.
247    pub ang_kp: AngVector,
248    /// The Integral gain applied to the accumulated angular position errors.
249    pub ang_ki: AngVector,
250    /// The Derivative gain applied to the instantaneous angular velocity errors.
251    pub ang_kd: AngVector,
252    /// The axes affected by this controller.
253    pub axes: AxesMask,
254    /// The linear position error accumulated through time (updated by the plugin).
255    pub lin_integral: Vect,
256    /// The angular position error accumulated through time (updated by the plugin).
257    pub ang_integral: AngVector,
258}
259
260impl Default for PidController {
261    fn default() -> Self {
262        Self::from_raw(&rapier::control::PidController::default())
263    }
264}
265
266impl PidController {
267    /// Initializes a PID controller with the same gains on all axes, and a default target.
268    ///
269    /// Only the axes specified in `axes` are controlled.
270    pub fn new(kp: Real, ki: Real, kd: Real, axes: AxesMask) -> Self {
271        Self::from_raw(&rapier::control::PidController::new(
272            kp,
273            ki,
274            kd,
275            axes.into(),
276        ))
277    }
278
279    /// Sets the target of this controller.
280    pub fn with_target(mut self, target: PidTarget) -> Self {
281        self.target = target;
282        self
283    }
284
285    /// Resets to zero the accumulated errors used by the Integral part of the controller.
286    pub fn reset_integrals(&mut self) {
287        self.lin_integral = Vect::ZERO;
288        self.ang_integral = splat_ang(0.0);
289    }
290
291    /// Converts a Rapier PID controller into this component, with a default target.
292    pub fn from_raw(raw: &rapier::control::PidController) -> Self {
293        Self {
294            target: PidTarget::default(),
295            lin_kp: raw.pd.lin_kp,
296            lin_ki: raw.lin_ki,
297            lin_kd: raw.pd.lin_kd,
298            ang_kp: raw.pd.ang_kp,
299            ang_ki: raw.ang_ki,
300            ang_kd: raw.pd.ang_kd,
301            axes: raw.pd.axes.into(),
302            lin_integral: raw.lin_integral,
303            ang_integral: raw.ang_integral,
304        }
305    }
306
307    /// Converts this component into a Rapier PID controller (the target is not part of it).
308    pub fn to_raw(&self) -> rapier::control::PidController {
309        rapier::control::PidController {
310            pd: rapier::control::PdController {
311                lin_kp: self.lin_kp,
312                lin_kd: self.lin_kd,
313                ang_kp: self.ang_kp,
314                ang_kd: self.ang_kd,
315                axes: self.axes.into(),
316            },
317            lin_integral: self.lin_integral,
318            ang_integral: self.ang_integral,
319            lin_ki: self.lin_ki,
320            ang_ki: self.ang_ki,
321        }
322    }
323
324    fn read_integrals(&mut self, raw: &rapier::control::PidController) {
325        self.lin_integral = raw.lin_integral;
326        self.ang_integral = raw.ang_integral;
327    }
328
329    /// Computes the velocity correction from the given position and velocity errors, and
330    /// accumulates the position errors for the Integral part.
331    pub fn correction(
332        &mut self,
333        dt: Real,
334        pose_errors: &PdErrors,
335        vel_errors: &PdErrors,
336    ) -> Velocity {
337        let mut raw = self.to_raw();
338        let result = raw.correction(dt, pose_errors, vel_errors);
339        self.read_integrals(&raw);
340        to_velocity(result)
341    }
342
343    /// Computes the velocity correction driving the given Rapier rigid-body toward
344    /// [`Self::target`], and accumulates the position errors for the Integral part.
345    pub fn rigid_body_correction(&mut self, dt: Real, body: &RapierRigidBody) -> Velocity {
346        let (pose, vels) = self.target.to_rapier();
347        let mut raw = self.to_raw();
348        let result = raw.rigid_body_correction(dt, body, pose, vels);
349        self.read_integrals(&raw);
350        to_velocity(result)
351    }
352}