use crate::joint::Joint;
use crate::math::{Fix128, QuatFix, Vec3Fix};
use crate::solver::RigidBody;
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum MotorMode {
Off,
Position,
Velocity,
}
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub struct PdController {
pub kp: Fix128,
pub kd: Fix128,
pub max_force: Fix128,
pub target_position: Fix128,
pub target_velocity: Fix128,
pub mode: MotorMode,
}
impl PdController {
#[inline]
#[must_use]
pub const fn new(kp: Fix128, kd: Fix128, max_force: Fix128) -> Self {
Self {
kp,
kd,
max_force,
target_position: Fix128::ZERO,
target_velocity: Fix128::ZERO,
mode: MotorMode::Off,
}
}
pub fn set_position_target(&mut self, target: Fix128) {
self.target_position = target;
self.mode = MotorMode::Position;
}
pub fn set_velocity_target(&mut self, target: Fix128) {
self.target_velocity = target;
self.mode = MotorMode::Velocity;
}
pub fn disable(&mut self) {
self.mode = MotorMode::Off;
}
#[must_use]
pub fn compute(&self, current_position: Fix128, current_velocity: Fix128) -> Fix128 {
match self.mode {
MotorMode::Off => Fix128::ZERO,
MotorMode::Position => {
let error = self.target_position - current_position;
let vel_error = self.target_velocity - current_velocity;
let output = self.kp * error + self.kd * vel_error;
clamp(output, -self.max_force, self.max_force)
}
MotorMode::Velocity => {
let vel_error = self.target_velocity - current_velocity;
let output = self.kp * vel_error;
clamp(output, -self.max_force, self.max_force)
}
}
}
}
impl Default for PdController {
fn default() -> Self {
Self::new(
Fix128::from_int(100), Fix128::from_int(10), Fix128::from_int(1000), )
}
}
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub struct JointMotor {
pub joint_index: usize,
pub controller: PdController,
}
impl JointMotor {
#[must_use]
pub const fn new(joint_index: usize, controller: PdController) -> Self {
Self {
joint_index,
controller,
}
}
}
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub struct PdController3D {
pub kp: Vec3Fix,
pub kd: Vec3Fix,
pub max_torque: Fix128,
pub target_rotation: QuatFix,
pub target_angular_velocity: Vec3Fix,
pub mode: MotorMode,
}
impl PdController3D {
#[inline]
#[must_use]
pub const fn new(kp: Vec3Fix, kd: Vec3Fix, max_torque: Fix128) -> Self {
Self {
kp,
kd,
max_torque,
target_rotation: QuatFix::IDENTITY,
target_angular_velocity: Vec3Fix::ZERO,
mode: MotorMode::Off,
}
}
pub fn set_rotation_target(&mut self, target: QuatFix) {
self.target_rotation = target;
self.mode = MotorMode::Position;
}
#[must_use]
pub fn compute_torque(
&self,
current_rotation: QuatFix,
current_angular_velocity: Vec3Fix,
) -> Vec3Fix {
match self.mode {
MotorMode::Off => Vec3Fix::ZERO,
MotorMode::Position => {
let error_quat =
shortest_arc(self.target_rotation.mul(current_rotation.conjugate()));
let error_axis = Vec3Fix::new(error_quat.x, error_quat.y, error_quat.z);
let two = Fix128::from_int(2);
let error_scaled =
Vec3Fix::new(error_axis.x * two, error_axis.y * two, error_axis.z * two);
let vel_error = self.target_angular_velocity - current_angular_velocity;
let torque = Vec3Fix::new(
self.kp.x * error_scaled.x + self.kd.x * vel_error.x,
self.kp.y * error_scaled.y + self.kd.y * vel_error.y,
self.kp.z * error_scaled.z + self.kd.z * vel_error.z,
);
let mag = torque.length();
if mag > self.max_torque && !mag.is_zero() {
torque * (self.max_torque / mag)
} else {
torque
}
}
MotorMode::Velocity => {
let vel_error = self.target_angular_velocity - current_angular_velocity;
let torque = Vec3Fix::new(
self.kp.x * vel_error.x,
self.kp.y * vel_error.y,
self.kp.z * vel_error.z,
);
let mag = torque.length();
if mag > self.max_torque && !mag.is_zero() {
torque * (self.max_torque / mag)
} else {
torque
}
}
}
}
}
pub fn apply_motors(motors: &[JointMotor], joints: &[Joint], bodies: &mut [RigidBody], dt: Fix128) {
for motor in motors {
if motor.controller.mode == MotorMode::Off {
continue;
}
if motor.joint_index >= joints.len() {
continue;
}
let joint = &joints[motor.joint_index];
let (body_a_idx, body_b_idx) = joint.bodies();
if body_a_idx >= bodies.len() || body_b_idx >= bodies.len() {
continue;
}
let body_a = bodies[body_a_idx];
let body_b = bodies[body_b_idx];
if let Joint::Hinge(hinge) = joint {
let axis = body_a.rotation.rotate_vec(hinge.local_axis_a);
let rel_quat = body_b.rotation.mul(body_a.rotation.conjugate());
let current_pos = crate::joint::compute_twist_angle(rel_quat, axis);
let current_vel = (body_b.angular_velocity - body_a.angular_velocity).dot(axis);
let torque = motor.controller.compute(current_pos, current_vel);
if torque.is_zero() {
continue;
}
let angular_impulse = axis * (torque * dt);
apply_angular_pair(bodies, body_a_idx, body_b_idx, angular_impulse);
continue;
}
let delta = body_b.position - body_a.position;
let current_pos = delta.length();
let rel_vel = body_b.velocity - body_a.velocity;
let current_vel = if current_pos.is_zero() {
Fix128::ZERO
} else {
rel_vel.dot(delta / current_pos)
};
let force = motor.controller.compute(current_pos, current_vel);
if force.is_zero() || current_pos.is_zero() {
continue;
}
let direction = delta / current_pos;
let impulse = direction * (force * dt);
if !body_a.inv_mass.is_zero() {
bodies[body_a_idx].velocity = bodies[body_a_idx].velocity - impulse * body_a.inv_mass;
}
if !body_b.inv_mass.is_zero() {
bodies[body_b_idx].velocity = bodies[body_b_idx].velocity + impulse * body_b.inv_mass;
}
}
}
pub(crate) fn apply_motors_3d(
motors: &[(usize, PdController3D)],
joints: &[Joint],
bodies: &mut [RigidBody],
dt: Fix128,
) {
for (joint_index, controller) in motors {
if controller.mode == MotorMode::Off || *joint_index >= joints.len() {
continue;
}
let (body_a_idx, body_b_idx) = joints[*joint_index].bodies();
if body_a_idx >= bodies.len() || body_b_idx >= bodies.len() {
continue;
}
let body_a = bodies[body_a_idx];
let body_b = bodies[body_b_idx];
let rel_quat = body_b.rotation.mul(body_a.rotation.conjugate());
let rel_omega = body_b.angular_velocity - body_a.angular_velocity;
let torque = controller.compute_torque(rel_quat, rel_omega);
if torque == Vec3Fix::ZERO {
continue;
}
apply_angular_pair(bodies, body_a_idx, body_b_idx, torque * dt);
}
}
fn apply_angular_pair(
bodies: &mut [RigidBody],
body_a_idx: usize,
body_b_idx: usize,
angular_impulse: Vec3Fix,
) {
let body_a = bodies[body_a_idx];
let body_b = bodies[body_b_idx];
if !body_a.inv_mass.is_zero() {
bodies[body_a_idx].angular_velocity =
body_a.angular_velocity - body_a.world_inv_inertia_apply(angular_impulse);
}
if !body_b.inv_mass.is_zero() {
bodies[body_b_idx].angular_velocity =
bodies[body_b_idx].angular_velocity + body_b.world_inv_inertia_apply(angular_impulse);
}
}
#[inline]
#[must_use]
pub(crate) fn shortest_arc(q: QuatFix) -> QuatFix {
if q.w < Fix128::ZERO {
QuatFix::new(-q.x, -q.y, -q.z, -q.w)
} else {
q
}
}
#[inline]
fn clamp(v: Fix128, min: Fix128, max: Fix128) -> Fix128 {
if v < min {
min
} else if v > max {
max
} else {
v
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::joint::BallJoint;
#[cfg(not(feature = "std"))]
use alloc::{vec, vec::Vec};
#[test]
fn test_pd_position_control() {
let mut pd = PdController::new(
Fix128::from_int(10),
Fix128::from_int(1),
Fix128::from_int(100),
);
pd.set_position_target(Fix128::from_int(5));
let output = pd.compute(Fix128::ZERO, Fix128::ZERO);
assert_eq!(output.hi, 50);
}
#[test]
fn test_pd_clamp() {
let mut pd = PdController::new(
Fix128::from_int(1000),
Fix128::ZERO,
Fix128::from_int(10), );
pd.set_position_target(Fix128::from_int(100));
let output = pd.compute(Fix128::ZERO, Fix128::ZERO);
assert_eq!(output.hi, 10, "Should be clamped to max_force");
}
#[test]
fn test_pd_off() {
let pd = PdController::new(
Fix128::from_int(10),
Fix128::from_int(1),
Fix128::from_int(100),
);
let output = pd.compute(Fix128::ZERO, Fix128::ZERO);
assert!(output.is_zero());
}
#[test]
fn test_velocity_mode() {
let mut pd = PdController::new(Fix128::from_int(10), Fix128::ZERO, Fix128::from_int(1000));
pd.set_velocity_target(Fix128::from_int(5));
let output = pd.compute(Fix128::ZERO, Fix128::ZERO);
assert_eq!(output.hi, 50);
}
#[test]
fn test_3d_controller() {
let mut pd = PdController3D::new(
Vec3Fix::from_int(100, 100, 100),
Vec3Fix::from_int(10, 10, 10),
Fix128::from_int(1000),
);
pd.set_rotation_target(QuatFix::IDENTITY);
pd.mode = MotorMode::Position;
let torque = pd.compute_torque(QuatFix::IDENTITY, Vec3Fix::ZERO);
let mag = torque.length();
assert!(
mag < Fix128::from_ratio(1, 10),
"Zero error should give zero torque"
);
}
fn motor_scene(ia: i64, ib: i64) -> (Vec<RigidBody>, Vec<Joint>) {
let mut a = RigidBody::new_dynamic(Vec3Fix::ZERO, Fix128::ONE);
let mut b = RigidBody::new_dynamic(Vec3Fix::from_int(4, 0, 0), Fix128::ONE);
a.inv_mass = Fix128::from_int(ia);
b.inv_mass = Fix128::from_int(ib);
let joints = vec![Joint::Ball(BallJoint::new(
0,
1,
Vec3Fix::ZERO,
Vec3Fix::ZERO,
))];
(vec![a, b], joints)
}
#[test]
fn apply_motors_position_mode_pushes_along_joint_axis_by_force_dt_inv_mass() {
let (mut bodies, joints) = motor_scene(0, 1);
let mut pd = PdController::new(Fix128::from_int(10), Fix128::ZERO, Fix128::from_int(100));
pd.set_position_target(Fix128::from_int(6));
let motors = [JointMotor::new(0, pd)];
apply_motors(&motors, &joints, &mut bodies, Fix128::from_ratio(1, 4));
assert_eq!(bodies[1].velocity, Vec3Fix::from_int(5, 0, 0));
assert_eq!(bodies[0].velocity, Vec3Fix::ZERO);
let (mut closer, joints) = motor_scene(0, 1);
let mut pd_in = pd;
pd_in.set_position_target(Fix128::from_int(2));
apply_motors(
&[JointMotor::new(0, pd_in)],
&joints,
&mut closer,
Fix128::from_ratio(1, 4),
);
assert_eq!(closer[1].velocity, Vec3Fix::from_int(-5, 0, 0));
let (mut both, joints) = motor_scene(1, 2);
apply_motors(
&[JointMotor::new(0, pd)],
&joints,
&mut both,
Fix128::from_ratio(1, 4),
);
assert_eq!(both[0].velocity, Vec3Fix::from_int(-5, 0, 0));
assert_eq!(both[1].velocity, Vec3Fix::from_int(10, 0, 0));
let (mut clamped, joints) = motor_scene(0, 1);
let mut pd_clamp =
PdController::new(Fix128::from_int(10), Fix128::ZERO, Fix128::from_int(8));
pd_clamp.set_position_target(Fix128::from_int(6));
apply_motors(
&[JointMotor::new(0, pd_clamp)],
&joints,
&mut clamped,
Fix128::from_ratio(1, 4),
);
assert_eq!(clamped[1].velocity, Vec3Fix::from_int(2, 0, 0));
}
#[test]
fn apply_motors_velocity_mode_uses_relative_velocity_along_axis() {
let (mut bodies, joints) = motor_scene(0, 1);
bodies[1].velocity = Vec3Fix::from_int(1, 7, 0);
let mut pd = PdController::new(Fix128::from_int(10), Fix128::ZERO, Fix128::from_int(100));
pd.set_velocity_target(Fix128::from_int(3));
apply_motors(
&[JointMotor::new(0, pd)],
&joints,
&mut bodies,
Fix128::from_ratio(1, 4),
);
assert_eq!(bodies[1].velocity, Vec3Fix::from_int(6, 7, 0));
}
#[test]
fn apply_motors_skips_off_out_of_range_zero_force_and_coincident_bodies() {
let (mut bodies, joints) = motor_scene(0, 1);
let mut pd = PdController::new(Fix128::from_int(10), Fix128::ZERO, Fix128::from_int(100));
pd.set_position_target(Fix128::from_int(6));
let dt = Fix128::from_ratio(1, 4);
let mut off = pd;
off.disable();
apply_motors(&[JointMotor::new(0, off)], &joints, &mut bodies, dt);
assert_eq!(bodies[1].velocity, Vec3Fix::ZERO);
apply_motors(&[JointMotor::new(1, pd)], &joints, &mut bodies, dt);
assert_eq!(bodies[1].velocity, Vec3Fix::ZERO);
let far = [Joint::Ball(BallJoint::new(
0,
5,
Vec3Fix::ZERO,
Vec3Fix::ZERO,
))];
apply_motors(&[JointMotor::new(0, pd)], &far, &mut bodies, dt);
assert_eq!(bodies[1].velocity, Vec3Fix::ZERO);
let mut at_target = pd;
at_target.set_position_target(Fix128::from_int(4));
apply_motors(&[JointMotor::new(0, at_target)], &joints, &mut bodies, dt);
assert_eq!(bodies[1].velocity, Vec3Fix::ZERO);
bodies[1].position = Vec3Fix::ZERO;
apply_motors(&[JointMotor::new(0, pd)], &joints, &mut bodies, dt);
assert_eq!(bodies[1].velocity, Vec3Fix::ZERO);
bodies[1].position = Vec3Fix::from_int(4, 0, 0);
apply_motors(&[JointMotor::new(0, pd)], &joints, &mut bodies, dt);
assert_eq!(bodies[1].velocity, Vec3Fix::from_int(5, 0, 0));
}
fn assert_vec_near(actual: Vec3Fix, expected: Vec3Fix, what: &str) {
let tol = Fix128::from_ratio(1, 1_000_000_000);
let d = actual - expected;
assert!(
d.x.abs() < tol && d.y.abs() < tol && d.z.abs() < tol,
"{what}: {actual:?} vs {expected:?}"
);
}
fn turn_z(sign: i64) -> QuatFix {
let s = Fix128::from_int(sign);
QuatFix::new(
Fix128::ZERO,
Fix128::ZERO,
Fix128::from_ratio(4, 5) * s,
Fix128::from_ratio(3, 5) * s,
)
}
#[test]
fn controller_3d_position_mode_follows_the_shortest_arc_and_clamps_the_magnitude() {
let mut pd = PdController3D::new(
Vec3Fix::from_int(10, 10, 10),
Vec3Fix::from_int(1, 1, 1),
Fix128::from_int(100),
);
pd.set_rotation_target(QuatFix::IDENTITY);
let omega = Vec3Fix::from_int(0, 0, 2);
let expected = Vec3Fix::from_int(0, 0, -18);
assert_vec_near(pd.compute_torque(turn_z(1), omega), expected, "w > 0");
assert_vec_near(pd.compute_torque(turn_z(-1), omega), expected, "w < 0");
pd.max_torque = Fix128::from_int(9);
assert_vec_near(
pd.compute_torque(turn_z(-1), omega),
Vec3Fix::from_int(0, 0, -9),
"clamped",
);
pd.mode = MotorMode::Off;
assert_eq!(pd.compute_torque(turn_z(1), omega), Vec3Fix::ZERO);
}
#[test]
fn controller_3d_velocity_mode_is_per_axis_proportional_and_clamped_by_magnitude() {
let mut pd = PdController3D::new(
Vec3Fix::from_int(2, 3, 4),
Vec3Fix::ZERO,
Fix128::from_int(100),
);
pd.target_angular_velocity = Vec3Fix::from_int(1, 1, 1);
pd.mode = MotorMode::Velocity;
assert_eq!(
pd.compute_torque(QuatFix::IDENTITY, Vec3Fix::from_int(0, 1, 0)),
Vec3Fix::from_int(2, 0, 4)
);
pd.kp = Vec3Fix::from_int(10, 10, 10);
pd.max_torque = Fix128::from_int(25);
pd.target_angular_velocity = Vec3Fix::from_int(3, 4, 0);
assert_vec_near(
pd.compute_torque(QuatFix::IDENTITY, Vec3Fix::ZERO),
Vec3Fix::from_int(15, 20, 0),
"clamped",
);
}
fn spinning_pair() -> (Vec<RigidBody>, Vec<Joint>) {
let (mut bodies, joints) = motor_scene(1, 1);
bodies[0].inv_inertia = Vec3Fix::from_int(1, 1, 1);
bodies[1].inv_inertia = Vec3Fix::from_int(2, 2, 2);
(bodies, joints)
}
#[test]
fn apply_motors_3d_applies_equal_and_opposite_angular_impulses() {
let mut pd = PdController3D::new(
Vec3Fix::from_int(2, 3, 4),
Vec3Fix::ZERO,
Fix128::from_int(100),
);
pd.target_angular_velocity = Vec3Fix::from_int(1, 1, 1);
pd.mode = MotorMode::Velocity;
let dt = Fix128::from_ratio(1, 2);
let (mut bodies, joints) = spinning_pair();
bodies[1].angular_velocity = Vec3Fix::from_int(0, 1, 0);
apply_motors_3d(&[(0, pd)], &joints, &mut bodies, dt);
assert_eq!(bodies[0].angular_velocity, Vec3Fix::from_int(-1, 0, -2));
assert_eq!(bodies[1].angular_velocity, Vec3Fix::from_int(2, 1, 4));
let half = Fix128::from_ratio(1, 2);
let momentum = bodies[0].angular_velocity + bodies[1].angular_velocity * half;
assert_eq!(momentum, Vec3Fix::new(Fix128::ZERO, half, Fix128::ZERO));
let (mut fixed_a, joints) = spinning_pair();
fixed_a[0].inv_mass = Fix128::ZERO;
fixed_a[1].angular_velocity = Vec3Fix::from_int(0, 1, 0);
apply_motors_3d(&[(0, pd)], &joints, &mut fixed_a, dt);
assert_eq!(fixed_a[0].angular_velocity, Vec3Fix::ZERO);
assert_eq!(fixed_a[1].angular_velocity, Vec3Fix::from_int(2, 1, 4));
}
#[test]
fn apply_motors_3d_skips_off_out_of_range_and_zero_torque() {
let mut pd = PdController3D::new(
Vec3Fix::from_int(1, 1, 1),
Vec3Fix::ZERO,
Fix128::from_int(100),
);
pd.target_angular_velocity = Vec3Fix::from_int(0, 0, 1);
pd.mode = MotorMode::Velocity;
let dt = Fix128::ONE;
let (mut bodies, joints) = spinning_pair();
let mut off = pd;
off.mode = MotorMode::Off;
apply_motors_3d(&[(0, off), (1, pd)], &joints, &mut bodies, dt);
let far = [Joint::Ball(BallJoint::new(
0,
5,
Vec3Fix::ZERO,
Vec3Fix::ZERO,
))];
apply_motors_3d(&[(0, pd)], &far, &mut bodies, dt);
assert_eq!(bodies[0].angular_velocity, Vec3Fix::ZERO);
assert_eq!(bodies[1].angular_velocity, Vec3Fix::ZERO);
bodies[1].angular_velocity = Vec3Fix::from_int(0, 0, 1);
apply_motors_3d(&[(0, pd)], &joints, &mut bodies, dt);
assert_eq!(bodies[0].angular_velocity, Vec3Fix::ZERO);
assert_eq!(bodies[1].angular_velocity, Vec3Fix::from_int(0, 0, 1));
bodies[1].angular_velocity = Vec3Fix::ZERO;
apply_motors_3d(&[(0, pd)], &joints, &mut bodies, dt);
assert_eq!(bodies[0].angular_velocity, Vec3Fix::from_int(0, 0, -1));
assert_eq!(bodies[1].angular_velocity, Vec3Fix::from_int(0, 0, 2));
}
#[test]
fn hinge_motor_drives_the_relative_rate_about_the_axis_and_clamps_below() {
let hinge = || {
Joint::Hinge(crate::joint::HingeJoint::new(
0,
1,
Vec3Fix::ZERO,
Vec3Fix::ZERO,
Vec3Fix::UNIT_Y,
Vec3Fix::UNIT_Y,
))
};
let pair = || {
let (mut bodies, _) = motor_scene(1, 1);
bodies[0].inv_inertia = Vec3Fix::from_int(1, 1, 1);
bodies[1].inv_inertia = Vec3Fix::from_int(1, 1, 1);
bodies
};
let dt = Fix128::from_ratio(1, 4);
let mut pd = PdController::new(Fix128::from_int(4), Fix128::ZERO, Fix128::from_int(100));
pd.set_velocity_target(Fix128::ONE);
let mut bodies = pair();
apply_motors(&[JointMotor::new(0, pd)], &[hinge()], &mut bodies, dt);
assert_eq!(bodies[0].angular_velocity, Vec3Fix::from_int(0, -1, 0));
assert_eq!(bodies[1].angular_velocity, Vec3Fix::from_int(0, 1, 0));
let mut slow = pd;
slow.max_force = Fix128::from_int(2);
slow.set_velocity_target(Fix128::from_int(-10));
let mut bodies = pair();
apply_motors(&[JointMotor::new(0, slow)], &[hinge()], &mut bodies, dt);
let half = Fix128::from_ratio(1, 2);
assert_eq!(
bodies[0].angular_velocity,
Vec3Fix::new(Fix128::ZERO, half, Fix128::ZERO)
);
assert_eq!(
bodies[1].angular_velocity,
Vec3Fix::new(Fix128::ZERO, -half, Fix128::ZERO)
);
let mut bodies = pair();
bodies[1].angular_velocity = Vec3Fix::from_int(0, 1, 0);
apply_motors(&[JointMotor::new(0, pd)], &[hinge()], &mut bodies, dt);
assert_eq!(bodies[0].angular_velocity, Vec3Fix::ZERO);
assert_eq!(bodies[1].angular_velocity, Vec3Fix::from_int(0, 1, 0));
}
}