diff --git a/src/control.rs b/src/control.rs index 905b00f..e8396ea 100644 --- a/src/control.rs +++ b/src/control.rs @@ -44,6 +44,14 @@ pub struct ControlMotorPositions { pub motor_props: Vec, } +#[derive(Message)] +pub struct ControlMotors { + pub handle: Handle, + pub positions: Vec, + pub velocities: Vec, + pub motor_props: Vec, +} + #[derive(Message)] pub struct ControlThrusters { pub handle: Handle, @@ -207,6 +215,88 @@ pub(crate) fn handle_control_motor_positions( } } +pub(crate) fn handle_control_motors( + mut er_control_motors: MessageReader, + 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, joint_link_handle: &UrdfJointHandle>, diff --git a/src/plugin.rs b/src/plugin.rs index 5f54410..12cc22e 100644 --- a/src/plugin.rs +++ b/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::() .add_message::() .add_message::() + .add_message::() .add_message::() .add_message::() .add_message::()