Skip to main content

bevy_rapier3d/plugin/systems/
pid_controller.rs

1use crate::control::{PdController, PidController};
2use crate::dynamics::RapierRigidBodyHandle;
3use crate::plugin::context::{
4    RapierContextEntityLink, RapierContextSimulation, RapierRigidBodySet,
5};
6use crate::plugin::RapierConfiguration;
7use bevy::prelude::*;
8use rapier::dynamics::{RigidBody as RapierRigidBody, RigidBodyVelocity};
9
10fn apply_correction(
11    body: &mut RapierRigidBody,
12    linvel: crate::math::Vect,
13    angvel: crate::math::AngVector,
14) {
15    let correction = RigidBodyVelocity { linvel, angvel };
16    if correction != RigidBodyVelocity::zero() {
17        let vels = *body.vels() + correction;
18        body.set_vels(vels, true);
19    }
20}
21
22/// System responsible for driving rigid-bodies with a [`PdController`] or a [`PidController`]
23/// toward their target, by adding the controllers' corrections to their velocities.
24pub fn update_pid_controllers(
25    config: Query<&RapierConfiguration>,
26    mut context_access: Query<(&RapierContextSimulation, &mut RapierRigidBodySet)>,
27    pd_controllers: Query<(
28        &RapierContextEntityLink,
29        &RapierRigidBodyHandle,
30        &PdController,
31    )>,
32    mut pid_controllers: Query<(
33        &RapierContextEntityLink,
34        &RapierRigidBodyHandle,
35        &mut PidController,
36    )>,
37) {
38    let is_active = |link: &RapierContextEntityLink| {
39        config
40            .get(link.0)
41            .is_ok_and(|config| config.physics_pipeline_active)
42    };
43
44    for (link, handle, controller) in pd_controllers.iter() {
45        if !is_active(link) {
46            continue;
47        }
48        let Ok((_, mut rigidbody_set)) = context_access.get_mut(link.0) else {
49            continue;
50        };
51        let Some(body) = rigidbody_set.bodies.get_mut(handle.0) else {
52            continue;
53        };
54        if body.is_fixed() || !body.is_enabled() {
55            continue;
56        }
57        let correction = controller.rigid_body_correction(body);
58        apply_correction(body, correction.linear, correction.angular);
59    }
60
61    for (link, handle, mut controller) in pid_controllers.iter_mut() {
62        if !is_active(link) {
63            continue;
64        }
65        let Ok((context, mut rigidbody_set)) = context_access.get_mut(link.0) else {
66            continue;
67        };
68        let dt = context.integration_parameters.dt;
69        let Some(body) = rigidbody_set.bodies.get_mut(handle.0) else {
70            continue;
71        };
72        if body.is_fixed() || !body.is_enabled() {
73            continue;
74        }
75        let correction = controller.rigid_body_correction(dt, body);
76        apply_correction(body, correction.linear, correction.angular);
77    }
78}
79
80#[cfg(test)]
81mod tests {
82    use super::*;
83    use crate::control::{AxesMask, PidTarget};
84    use crate::geometry::Collider;
85    use crate::math::Vect;
86    use crate::plugin::{NoUserData, RapierPhysicsPlugin};
87    use crate::prelude::RigidBody;
88    use bevy::time::{TimePlugin, TimeUpdateStrategy};
89
90    #[test]
91    fn pid_controller_drives_body_toward_target() {
92        let mut app = App::new();
93        app.add_plugins((
94            TransformPlugin,
95            TimePlugin,
96            RapierPhysicsPlugin::<NoUserData>::default(),
97        ))
98        .insert_resource(TimeUpdateStrategy::ManualDuration(
99            std::time::Duration::from_secs_f32(1.0 / 60.0),
100        ));
101        app.finish();
102
103        let target = Vect::splat(3.0);
104        let controlled = app
105            .world_mut()
106            .spawn((
107                Transform::default(),
108                RigidBody::Dynamic,
109                Collider::ball(0.5),
110                PidController::new(10.0, 1.0, 0.5, AxesMask::LIN_AXES)
111                    .with_target(PidTarget::from_translation(target)),
112            ))
113            .id();
114        let pd_controlled = app
115            .world_mut()
116            .spawn((
117                Transform::from_xyz(10.0, 0.0, 0.0),
118                RigidBody::Dynamic,
119                Collider::ball(0.5),
120                PdController::new(10.0, 0.5, AxesMask::LIN_AXES)
121                    .with_target(PidTarget::from_translation(Vect::X * 10.0 + target)),
122            ))
123            .id();
124
125        for _ in 0..300 {
126            app.update();
127        }
128
129        let translation = |entity: Entity| {
130            let t = app.world().get::<Transform>(entity).unwrap().translation;
131            #[cfg(feature = "dim2")]
132            return t.truncate();
133            #[cfg(feature = "dim3")]
134            return t;
135        };
136        let pid_error = (translation(controlled) - target).length();
137        let pd_error = (translation(pd_controlled) - Vect::X * 10.0 - target).length();
138        // Both bodies start about 5.2 units away from their target and fight gravity.
139        assert!(
140            pid_error < 0.05,
141            "PID controller error too large: {pid_error}"
142        );
143        assert!(pd_error < 0.05, "PD controller error too large: {pd_error}");
144        assert_ne!(
145            app.world()
146                .get::<PidController>(controlled)
147                .unwrap()
148                .lin_integral,
149            Vect::ZERO,
150            "The integral error should be written back"
151        );
152    }
153}