Add motors to ball joints.

1 files changed, 196 insertions, 16 deletions

diff --git a/src/dynamics/solver/joint_constraint/ball_velocity_constraint.rs b/src/dynamics/solver/joint_constraint/ball_velocity_constraint.rs
index e75f978..8c19ee6 100644
--- a/src/dynamics/solver/joint_constraint/ball_velocity_constraint.rs
+++ b/src/dynamics/solver/joint_constraint/ball_velocity_constraint.rs
@@ -2,7 +2,7 @@ use crate::dynamics::solver::DeltaVel;
 use crate::dynamics::{
     BallJoint, IntegrationParameters, JointGraphEdge, JointIndex, JointParams, RigidBody,
 };
-use crate::math::{AngularInertia, Real, SdpMatrix, Vector};
+use crate::math::{AngVector, AngularInertia, Real, SdpMatrix, Vector};
 use crate::utils::{WAngularInertia, WCross, WCrossMatrix};
 
 #[derive(Debug)]
@@ -13,13 +13,17 @@ pub(crate) struct BallVelocityConstraint {
     joint_id: JointIndex,
 
     rhs: Vector<Real>,
-    pub(crate) impulse: Vector<Real>,
+    impulse: Vector<Real>,
 
     r1: Vector<Real>,
     r2: Vector<Real>,
 
     inv_lhs: SdpMatrix<Real>,
 
+    motor_rhs: AngVector<Real>,
+    motor_impulse: AngVector<Real>,
+    motor_inv_lhs: Option<AngularInertia<Real>>,
+
     im1: Real,
     im2: Real,
 
@@ -33,10 +37,10 @@ impl BallVelocityConstraint {
         joint_id: JointIndex,
         rb1: &RigidBody,
         rb2: &RigidBody,
-        cparams: &BallJoint,
+        joint: &BallJoint,
     ) -> Self {
-        let anchor1 = rb1.position * cparams.local_anchor1 - rb1.world_com;
-        let anchor2 = rb2.position * cparams.local_anchor2 - rb2.world_com;
+        let anchor1 = rb1.position * joint.local_anchor1 - rb1.world_com;
+        let anchor2 = rb2.position * joint.local_anchor2 - rb2.world_com;
 
         let vel1 = rb1.linvel + rb1.angvel.gcross(anchor1);
         let vel2 = rb2.linvel + rb2.angvel.gcross(anchor2);
@@ -77,17 +81,86 @@ impl BallVelocityConstraint {
 
         let inv_lhs = lhs.inverse_unchecked();
 
+        /*
+         * Motor part.
+         */
+        let mut motor_rhs = na::zero();
+        let mut motor_inv_lhs = None;
+        let motor_max_impulse = joint.motor_max_impulse;
+
+        let (stiffness, damping, gamma, keep_lhs) = joint.motor_model.combine_coefficients(
+            params.dt,
+            joint.motor_stiffness,
+            joint.motor_damping,
+        );
+
+        if stiffness != 0.0 {
+            let dpos =
+                rb2.position.rotation * (rb1.position.rotation * joint.motor_target_pos).inverse();
+            #[cfg(feature = "dim2")]
+            {
+                motor_rhs += dpos.angle() * stiffness;
+            }
+            #[cfg(feature = "dim3")]
+            {
+                motor_rhs += dpos.scaled_axis() * stiffness;
+            }
+        }
+
+        if damping != 0.0 {
+            let curr_vel = rb2.angvel - rb1.angvel;
+            motor_rhs += (curr_vel - joint.motor_target_vel) * damping;
+        }
+
+        #[cfg(feature = "dim2")]
+        if stiffness != 0.0 || damping != 0.0 {
+            motor_inv_lhs = if keep_lhs {
+                let ii1 = rb1.effective_world_inv_inertia_sqrt.squared();
+                let ii2 = rb2.effective_world_inv_inertia_sqrt.squared();
+                Some(gamma / (ii1 + ii2))
+            } else {
+                Some(gamma)
+            };
+            motor_rhs /= gamma;
+        }
+
+        #[cfg(feature = "dim3")]
+        if stiffness != 0.0 || damping != 0.0 {
+            motor_inv_lhs = if keep_lhs {
+                let ii1 = rb1.effective_world_inv_inertia_sqrt.squared();
+                let ii2 = rb2.effective_world_inv_inertia_sqrt.squared();
+                Some((ii1 + ii2).inverse_unchecked() * gamma)
+            } else {
+                Some(SdpMatrix::diagonal(gamma))
+            };
+            motor_rhs /= gamma;
+        }
+
+        #[cfg(feature = "dim2")]
+        let motor_impulse = na::clamp(joint.motor_impulse, -motor_max_impulse, motor_max_impulse)
+            * params.warmstart_coeff;
+
+        #[cfg(feature = "dim3")]
+        let motor_impulse = joint
+            .motor_impulse
+            .try_clamp_magnitude(-motor_max_impulse, motor_max_impulse, 1.0e-6)
+            .unwrap_or_else(na::zero)
+            * params.warmstart_coeff;
+
         BallVelocityConstraint {
             joint_id,
             mj_lambda1: rb1.active_set_offset,
             mj_lambda2: rb2.active_set_offset,
             im1,
             im2,
-            impulse: cparams.impulse * params.warmstart_coeff,
+            impulse: joint.impulse * params.warmstart_coeff,
             r1: anchor1,
             r2: anchor2,
             rhs,
             inv_lhs,
+            motor_rhs,
+            motor_impulse,
+            motor_inv_lhs,
             ii1_sqrt: rb1.effective_world_inv_inertia_sqrt,
             ii2_sqrt: rb2.effective_world_inv_inertia_sqrt,
         }
@@ -98,9 +171,13 @@ impl BallVelocityConstraint {
         let mut mj_lambda2 = mj_lambdas[self.mj_lambda2 as usize];
 
         mj_lambda1.linear += self.im1 * self.impulse;
-        mj_lambda1.angular += self.ii1_sqrt.transform_vector(self.r1.gcross(self.impulse));
+        mj_lambda1.angular += self
+            .ii1_sqrt
+            .transform_vector(self.r1.gcross(self.impulse) + self.motor_impulse);
         mj_lambda2.linear -= self.im2 * self.impulse;
-        mj_lambda2.angular -= self.ii2_sqrt.transform_vector(self.r2.gcross(self.impulse));
+        mj_lambda2.angular -= self
+            .ii2_sqrt
+            .transform_vector(self.r2.gcross(self.impulse) + self.motor_impulse);
 
         mj_lambdas[self.mj_lambda1 as usize] = mj_lambda1;
         mj_lambdas[self.mj_lambda2 as usize] = mj_lambda2;
@@ -125,6 +202,21 @@ impl BallVelocityConstraint {
         mj_lambda2.linear -= self.im2 * impulse;
         mj_lambda2.angular -= self.ii2_sqrt.transform_vector(self.r2.gcross(impulse));
 
+        /*
+         * Motor part.
+         */
+        if let Some(motor_inv_lhs) = &self.motor_inv_lhs {
+            let ang_vel1 = self.ii1_sqrt.transform_vector(mj_lambda1.angular);
+            let ang_vel2 = self.ii2_sqrt.transform_vector(mj_lambda2.angular);
+
+            let dangvel = (ang_vel2 - ang_vel1) + self.motor_rhs;
+            let impulse = motor_inv_lhs.transform_vector(dangvel);
+            self.motor_impulse += impulse;
+
+            mj_lambda1.angular += self.ii1_sqrt.transform_vector(impulse);
+            mj_lambda2.angular -= self.ii2_sqrt.transform_vector(impulse);
+        }
+
         mj_lambdas[self.mj_lambda1 as usize] = mj_lambda1;
         mj_lambdas[self.mj_lambda2 as usize] = mj_lambda2;
     }
@@ -141,10 +233,16 @@ impl BallVelocityConstraint {
 pub(crate) struct BallVelocityGroundConstraint {
     mj_lambda2: usize,
     joint_id: JointIndex,
+    r2: Vector<Real>,
+
     rhs: Vector<Real>,
     impulse: Vector<Real>,
-    r2: Vector<Real>,
     inv_lhs: SdpMatrix<Real>,
+
+    motor_rhs: AngVector<Real>,
+    motor_impulse: AngVector<Real>,
+    motor_inv_lhs: Option<AngularInertia<Real>>,
+
     im2: Real,
     ii2_sqrt: AngularInertia<Real>,
 }
@@ -155,18 +253,18 @@ impl BallVelocityGroundConstraint {
         joint_id: JointIndex,
         rb1: &RigidBody,
         rb2: &RigidBody,
-        cparams: &BallJoint,
+        joint: &BallJoint,
         flipped: bool,
     ) -> Self {
         let (anchor1, anchor2) = if flipped {
             (
-                rb1.position * cparams.local_anchor2 - rb1.world_com,
-                rb2.position * cparams.local_anchor1 - rb2.world_com,
+                rb1.position * joint.local_anchor2 - rb1.world_com,
+                rb2.position * joint.local_anchor1 - rb2.world_com,
             )
         } else {
             (
-                rb1.position * cparams.local_anchor1 - rb1.world_com,
-                rb2.position * cparams.local_anchor2 - rb2.world_com,
+                rb1.position * joint.local_anchor1 - rb1.world_com,
+                rb2.position * joint.local_anchor2 - rb2.world_com,
             )
         };
 
@@ -199,14 +297,81 @@ impl BallVelocityGroundConstraint {
 
         let inv_lhs = lhs.inverse_unchecked();
 
+        /*
+         * Motor part.
+         */
+        let mut motor_rhs = na::zero();
+        let mut motor_inv_lhs = None;
+        let motor_max_impulse = joint.motor_max_impulse;
+
+        let (stiffness, damping, gamma, keep_lhs) = joint.motor_model.combine_coefficients(
+            params.dt,
+            joint.motor_stiffness,
+            joint.motor_damping,
+        );
+
+        if stiffness != 0.0 {
+            let dpos =
+                rb2.position.rotation * (rb1.position.rotation * joint.motor_target_pos).inverse();
+            #[cfg(feature = "dim2")]
+            {
+                motor_rhs += dpos.angle() * stiffness;
+            }
+            #[cfg(feature = "dim3")]
+            {
+                motor_rhs += dpos.scaled_axis() * stiffness;
+            }
+        }
+
+        if damping != 0.0 {
+            let curr_vel = rb2.angvel - rb1.angvel;
+            motor_rhs += (curr_vel - joint.motor_target_vel) * damping;
+        }
+
+        #[cfg(feature = "dim2")]
+        if stiffness != 0.0 || damping != 0.0 {
+            motor_inv_lhs = if keep_lhs {
+                let ii2 = rb2.effective_world_inv_inertia_sqrt.squared();
+                Some(gamma / ii2)
+            } else {
+                Some(gamma)
+            };
+            motor_rhs /= gamma;
+        }
+
+        #[cfg(feature = "dim3")]
+        if stiffness != 0.0 || damping != 0.0 {
+            motor_inv_lhs = if keep_lhs {
+                let ii2 = rb2.effective_world_inv_inertia_sqrt.squared();
+                Some(ii2.inverse_unchecked() * gamma)
+            } else {
+                Some(SdpMatrix::diagonal(gamma))
+            };
+            motor_rhs /= gamma;
+        }
+
+        #[cfg(feature = "dim2")]
+        let motor_impulse = na::clamp(joint.motor_impulse, -motor_max_impulse, motor_max_impulse)
+            * params.warmstart_coeff;
+
+        #[cfg(feature = "dim3")]
+        let motor_impulse = joint
+            .motor_impulse
+            .try_clamp_magnitude(-motor_max_impulse, motor_max_impulse, 1.0e-6)
+            .unwrap_or_else(na::zero)
+            * params.warmstart_coeff;
+
         BallVelocityGroundConstraint {
             joint_id,
             mj_lambda2: rb2.active_set_offset,
             im2,
-            impulse: cparams.impulse * params.warmstart_coeff,
+            impulse: joint.impulse * params.warmstart_coeff,
             r2: anchor2,
             rhs,
             inv_lhs,
+            motor_rhs,
+            motor_impulse,
+            motor_inv_lhs,
             ii2_sqrt: rb2.effective_world_inv_inertia_sqrt,
         }
     }
@@ -214,7 +379,9 @@ impl BallVelocityGroundConstraint {
     pub fn warmstart(&self, mj_lambdas: &mut [DeltaVel<Real>]) {
         let mut mj_lambda2 = mj_lambdas[self.mj_lambda2 as usize];
         mj_lambda2.linear -= self.im2 * self.impulse;
-        mj_lambda2.angular -= self.ii2_sqrt.transform_vector(self.r2.gcross(self.impulse));
+        mj_lambda2.angular -= self
+            .ii2_sqrt
+            .transform_vector(self.r2.gcross(self.impulse) + self.motor_impulse);
         mj_lambdas[self.mj_lambda2 as usize] = mj_lambda2;
     }
 
@@ -231,6 +398,19 @@ impl BallVelocityGroundConstraint {
         mj_lambda2.linear -= self.im2 * impulse;
         mj_lambda2.angular -= self.ii2_sqrt.transform_vector(self.r2.gcross(impulse));
 
+        /*
+         * Motor part.
+         */
+        if let Some(motor_inv_lhs) = &self.motor_inv_lhs {
+            let ang_vel2 = self.ii2_sqrt.transform_vector(mj_lambda2.angular);
+
+            let dangvel = ang_vel2 + self.motor_rhs;
+            let impulse = motor_inv_lhs.transform_vector(dangvel);
+            self.motor_impulse += impulse;
+
+            mj_lambda2.angular -= self.ii2_sqrt.transform_vector(impulse);
+        }
+
         mj_lambdas[self.mj_lambda2 as usize] = mj_lambda2;
     }