simultaneous pos & vel control of joints

src/control.rs

@@ -44,6 +44,14 @@ pub struct ControlMotorPositions {
     pub motor_props: Vec<MotorProps>,
 }
 
+#[derive(Message)]
+pub struct ControlMotors {
+    pub handle: Handle<UrdfAsset>,
+    pub positions: Vec<f32>,
+    pub velocities: Vec<f32>,
+    pub motor_props: Vec<MotorProps>,
+}
+
 #[derive(Message)]
 pub struct ControlThrusters {
     pub handle: Handle<UrdfAsset>,
@@ -207,6 +215,88 @@ pub(crate) fn handle_control_motor_positions(
     }
 }
 
+pub(crate) fn handle_control_motors(
+    mut er_control_motors: MessageReader<ControlMotors>,
+    q_urdf_robots: Query<(Entity, &URDFRobot)>,
+    mut q_rapier_joints: Query<(&mut RapierContextJoints, &mut RapierRigidBodySet)>,
+) {
+    for event in er_control_motors.read() {
+        // Early exit conditions
+        if event.positions.is_empty() {
+            continue;
+        }
+
+        // Find target robot
+        let target_robot = q_urdf_robots
+            .iter()
+            .find(|(_, robot)| robot.handle == event.handle);
+
+        let Some((_, urdf_robot)) = target_robot else {
+            continue;
+        };
+
+        let mut actuator_index = 0;
+
+        for (mut rapier_context_joints, mut rapier_rigid_bodies) in q_rapier_joints.iter_mut() {
+            for joint_link_handle in &urdf_robot.rapier_handles.joints {
+                let Some(joint_handle) = joint_link_handle.joint else {
+                    continue;
+                };
+
+                let Some((multibody, index)) =
+                    rapier_context_joints.multibody_joints.get_mut(joint_handle)
+                else {
+                    continue;
+                };
+
+                let Some(link) = multibody.link_mut(index) else {
+                    continue;
+                };
+
+                let joint = &mut link.joint.data;
+                let Some(revolute) = joint.as_revolute_mut() else {
+                    continue;
+                };
+
+                if revolute.motor().is_none() {
+                    continue;
+                }
+
+                wakeup_links(&mut rapier_rigid_bodies, joint_link_handle);
+
+                if actuator_index >= event.positions.len() {
+                    panic!(
+                        "Not enough positions parameters provided. Required: {}, Available: {}",
+                        actuator_index + 1,
+                        event.positions.len()
+                    );
+                }
+
+                if actuator_index >= event.motor_props.len() {
+                    panic!(
+                        "Not enough motor_props parameters provided. Required: {}, Available: {}",
+                        actuator_index + 1,
+                        event.positions.len()
+                    );
+                }
+
+                let target_position = event.positions[actuator_index];
+                let target_velocity = event.velocities[actuator_index];
+                let motor_props = event.motor_props[actuator_index];
+
+                joint.set_motor(
+                    JointAxis::AngX,
+                    target_position,
+                    target_velocity,
+                    motor_props.stiffness,
+                    motor_props.damping,
+                );
+                actuator_index += 1;
+            }
+        }
+    }
+}
+
 fn wakeup_links(
     rapier_rigid_bodies: &mut Mut<RapierRigidBodySet>,
     joint_link_handle: &UrdfJointHandle<Option<MultibodyJointHandle>>,

src/plugin.rs

@@ -24,9 +24,9 @@ pub fn rapier_to_bevy_rotation() -> Quat {
 
 use crate::{
     control::{
-        handle_control_motor_positions, handle_control_motor_velocities, ControlFins,
-        ControlMotorPositions, ControlMotorVelocities, ControlThrusters, SensorsRead,
-        UAVStateUpdate, UUVStateUpdate,
+        handle_control_motor_positions, handle_control_motor_velocities, handle_control_motors,
+        ControlFins, ControlMotorPositions, ControlMotorVelocities, ControlMotors,
+        ControlThrusters, SensorsRead, UAVStateUpdate, UUVStateUpdate,
     },
     spawn::{
         handle_despawn_robot, handle_load_robot, handle_spawn_robot, handle_wait_robot_loaded,
@@ -66,6 +66,7 @@ where
             PhysicsSet::SyncBackend => (
                 handle_control_motor_velocities,
                 handle_control_motor_positions,
+                handle_control_motors,
                 handle_control_thrusts,
                 handle_control_thrusters,
                 handle_control_fins,
@@ -103,6 +104,7 @@ impl Plugin for URDFPlugin {
         app.init_asset_loader::<urdf_asset_loader::RpyAssetLoader>()
             .add_message::<ControlMotorVelocities>()
             .add_message::<ControlMotorPositions>()
+            .add_message::<ControlMotors>()
             .add_message::<ControlThrusts>()
             .add_message::<ControlThrusters>()
             .add_message::<ControlFins>()