Add motors to prismatic joints.

1 files changed, 179 insertions, 61 deletions

diff --git a/src/dynamics/solver/joint_constraint/prismatic_velocity_constraint.rs b/src/dynamics/solver/joint_constraint/prismatic_velocity_constraint.rs
index a361a37..9aaf172 100644
--- a/src/dynamics/solver/joint_constraint/prismatic_velocity_constraint.rs
+++ b/src/dynamics/solver/joint_constraint/prismatic_velocity_constraint.rs
@@ -41,6 +41,13 @@ pub(crate) struct PrismaticVelocityConstraint {
     #[cfg(feature = "dim2")]
     impulse: Vector2<Real>,
 
+    motor_axis1: Vector<Real>,
+    motor_axis2: Vector<Real>,
+    motor_impulse: Real,
+    motor_rhs: Real,
+    motor_inv_lhs: Real,
+    motor_max_impulse: Real,
+
     limits_impulse: Real,
     limits_forcedirs: Option<(Vector<Real>, Vector<Real>)>,
     limits_rhs: Real,
@@ -63,35 +70,21 @@ impl PrismaticVelocityConstraint {
         joint_id: JointIndex,
         rb1: &RigidBody,
         rb2: &RigidBody,
-        cparams: &PrismaticJoint,
+        joint: &PrismaticJoint,
     ) -> Self {
         // Linear part.
-        let anchor1 = rb1.position * cparams.local_anchor1;
-        let anchor2 = rb2.position * cparams.local_anchor2;
-        let axis1 = rb1.position * cparams.local_axis1;
-        let axis2 = rb2.position * cparams.local_axis2;
+        let anchor1 = rb1.position * joint.local_anchor1;
+        let anchor2 = rb2.position * joint.local_anchor2;
+        let axis1 = rb1.position * joint.local_axis1;
+        let axis2 = rb2.position * joint.local_axis2;
         #[cfg(feature = "dim2")]
-        let basis1 = rb1.position * cparams.basis1[0];
+        let basis1 = rb1.position * joint.basis1[0];
         #[cfg(feature = "dim3")]
         let basis1 = Matrix3x2::from_columns(&[
-            rb1.position * cparams.basis1[0],
-            rb1.position * cparams.basis1[1],
+            rb1.position * joint.basis1[0],
+            rb1.position * joint.basis1[1],
         ]);
 
-        // #[cfg(feature = "dim2")]
-        // let r21 = Rotation::rotation_between_axis(&axis1, &axis2)
-        //     .to_rotation_matrix()
-        //     .into_inner();
-        // #[cfg(feature = "dim3")]
-        // let r21 = Rotation::rotation_between_axis(&axis1, &axis2)
-        //     .unwrap_or_else(Rotation::identity)
-        //     .to_rotation_matrix()
-        //     .into_inner();
-        // let basis2 = r21 * basis1;
-        // NOTE: we use basis2 := basis1 for now is that allows
-        // simplifications of the computation without introducing
-        // much instabilities.
-
         let im1 = rb1.effective_inv_mass;
         let ii1 = rb1.effective_world_inv_inertia_sqrt.squared();
         let r1 = anchor1 - rb1.world_com;
@@ -149,25 +142,57 @@ impl PrismaticVelocityConstraint {
         #[cfg(feature = "dim3")]
         let rhs = Vector5::new(lin_rhs.x, lin_rhs.y, ang_rhs.x, ang_rhs.y, ang_rhs.z);
 
-        // Setup limit constraint.
+        /*
+         * Setup motor.
+         */
+        let mut motor_rhs = 0.0;
+        let mut motor_inv_lhs = 0.0;
+        let mut 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 dist = anchor2.coords.dot(&axis2) - anchor1.coords.dot(&axis1);
+            motor_rhs += (dist - joint.motor_target_pos) * stiffness;
+        }
+
+        if damping != 0.0 {
+            let curr_vel = rb2.linvel.dot(&axis2) - rb1.linvel.dot(&axis1);
+            motor_rhs += (curr_vel - joint.motor_target_vel) * damping;
+        }
+
+        if stiffness != 0.0 || damping != 0.0 {
+            motor_inv_lhs = if keep_lhs { gamma / (im1 + im2) } else { gamma };
+            motor_rhs /= gamma;
+        }
+
+        let motor_impulse = na::clamp(joint.motor_impulse, -motor_max_impulse, motor_max_impulse);
+
+        /*
+         * Setup limit constraint.
+         */
         let mut limits_forcedirs = None;
         let mut limits_rhs = 0.0;
         let mut limits_impulse = 0.0;
 
-        if cparams.limits_enabled {
+        if joint.limits_enabled {
             let danchor = anchor2 - anchor1;
             let dist = danchor.dot(&axis1);
 
-            // FIXME: we should allow both limits to be active at
+            // TODO: we should allow both limits to be active at
             // the same time, and allow predictive constraint activation.
-            if dist < cparams.limits[0] {
+            if dist < joint.limits[0] {
                 limits_forcedirs = Some((-axis1.into_inner(), axis2.into_inner()));
                 limits_rhs = anchor_linvel2.dot(&axis2) - anchor_linvel1.dot(&axis1);
-                limits_impulse = cparams.limits_impulse;
-            } else if dist > cparams.limits[1] {
+                limits_impulse = joint.limits_impulse;
+            } else if dist > joint.limits[1] {
                 limits_forcedirs = Some((axis1.into_inner(), -axis2.into_inner()));
                 limits_rhs = -anchor_linvel2.dot(&axis2) + anchor_linvel1.dot(&axis1);
-                limits_impulse = cparams.limits_impulse;
+                limits_impulse = joint.limits_impulse;
             }
         }
 
@@ -179,10 +204,16 @@ impl PrismaticVelocityConstraint {
             ii1_sqrt: rb1.effective_world_inv_inertia_sqrt,
             im2,
             ii2_sqrt: rb2.effective_world_inv_inertia_sqrt,
-            impulse: cparams.impulse * params.warmstart_coeff,
+            impulse: joint.impulse * params.warmstart_coeff,
             limits_impulse: limits_impulse * params.warmstart_coeff,
             limits_forcedirs,
             limits_rhs,
+            motor_rhs,
+            motor_inv_lhs,
+            motor_impulse,
+            motor_axis1: *axis1,
+            motor_axis2: *axis2,
+            motor_max_impulse,
             basis1,
             inv_lhs,
             rhs,
@@ -211,6 +242,11 @@ impl PrismaticVelocityConstraint {
             .ii2_sqrt
             .transform_vector(ang_impulse + self.r2.gcross(lin_impulse));
 
+        // Warmstart motors.
+        mj_lambda1.linear += self.motor_axis1 * (self.im1 * self.motor_impulse);
+        mj_lambda2.linear -= self.motor_axis2 * (self.im2 * self.motor_impulse);
+
+        // Warmstart limits.
         if let Some((limits_forcedir1, limits_forcedir2)) = self.limits_forcedirs {
             mj_lambda1.linear += limits_forcedir1 * (self.im1 * self.limits_impulse);
             mj_lambda2.linear += limits_forcedir2 * (self.im2 * self.limits_impulse);
@@ -225,7 +261,7 @@ impl PrismaticVelocityConstraint {
         let mut mj_lambda2 = mj_lambdas[self.mj_lambda2 as usize];
 
         /*
-         * Joint consraint.
+         * Joint constraint.
          */
         let ang_vel1 = self.ii1_sqrt.transform_vector(mj_lambda1.angular);
         let ang_vel2 = self.ii2_sqrt.transform_vector(mj_lambda2.angular);
@@ -257,11 +293,28 @@ impl PrismaticVelocityConstraint {
             .transform_vector(ang_impulse + self.r2.gcross(lin_impulse));
 
         /*
+         * Motors.
+         */
+        if self.motor_inv_lhs != 0.0 {
+            let lin_dvel = self.motor_axis2.dot(&mj_lambda2.linear)
+                - self.motor_axis1.dot(&mj_lambda1.linear)
+                + self.motor_rhs;
+            let new_impulse = na::clamp(
+                self.motor_impulse + lin_dvel * self.motor_inv_lhs,
+                -self.motor_max_impulse,
+                self.motor_max_impulse,
+            );
+            let dimpulse = new_impulse - self.motor_impulse;
+            self.motor_impulse = new_impulse;
+
+            mj_lambda1.linear += self.motor_axis1 * (self.im1 * dimpulse);
+            mj_lambda2.linear -= self.motor_axis2 * (self.im2 * dimpulse);
+        }
+
+        /*
          * Joint limits.
          */
         if let Some((limits_forcedir1, limits_forcedir2)) = self.limits_forcedirs {
-            // FIXME: the transformation by ii2_sqrt could be avoided by
-            // reusing some computations above.
             let ang_vel1 = self.ii1_sqrt.transform_vector(mj_lambda1.angular);
             let ang_vel2 = self.ii2_sqrt.transform_vector(mj_lambda2.angular);
 
@@ -284,6 +337,7 @@ impl PrismaticVelocityConstraint {
         let joint = &mut joints_all[self.joint_id].weight;
         if let JointParams::PrismaticJoint(revolute) = &mut joint.params {
             revolute.impulse = self.impulse;
+            revolute.motor_impulse = self.motor_impulse;
             revolute.limits_impulse = self.limits_impulse;
         }
     }
@@ -315,6 +369,11 @@ pub(crate) struct PrismaticVelocityGroundConstraint {
     limits_rhs: Real,
 
     axis2: Vector<Real>,
+    motor_impulse: Real,
+    motor_rhs: Real,
+    motor_inv_lhs: Real,
+    motor_max_impulse: Real,
+
     #[cfg(feature = "dim2")]
     basis1: Vector2<Real>,
     #[cfg(feature = "dim3")]
@@ -331,7 +390,7 @@ impl PrismaticVelocityGroundConstraint {
         joint_id: JointIndex,
         rb1: &RigidBody,
         rb2: &RigidBody,
-        cparams: &PrismaticJoint,
+        joint: &PrismaticJoint,
         flipped: bool,
     ) -> Self {
         let anchor2;
@@ -341,35 +400,35 @@ impl PrismaticVelocityGroundConstraint {
         let basis1;
 
         if flipped {
-            anchor2 = rb2.position * cparams.local_anchor1;
-            anchor1 = rb1.position * cparams.local_anchor2;
-            axis2 = rb2.position * cparams.local_axis1;
-            axis1 = rb1.position * cparams.local_axis2;
+            anchor2 = rb2.position * joint.local_anchor1;
+            anchor1 = rb1.position * joint.local_anchor2;
+            axis2 = rb2.position * joint.local_axis1;
+            axis1 = rb1.position * joint.local_axis2;
             #[cfg(feature = "dim2")]
             {
-                basis1 = rb1.position * cparams.basis2[0];
+                basis1 = rb1.position * joint.basis2[0];
             }
             #[cfg(feature = "dim3")]
             {
                 basis1 = Matrix3x2::from_columns(&[
-                    rb1.position * cparams.basis2[0],
-                    rb1.position * cparams.basis2[1],
+                    rb1.position * joint.basis2[0],
+                    rb1.position * joint.basis2[1],
                 ]);
             }
         } else {
-            anchor2 = rb2.position * cparams.local_anchor2;
-            anchor1 = rb1.position * cparams.local_anchor1;
-            axis2 = rb2.position * cparams.local_axis2;
-            axis1 = rb1.position * cparams.local_axis1;
+            anchor2 = rb2.position * joint.local_anchor2;
+            anchor1 = rb1.position * joint.local_anchor1;
+            axis2 = rb2.position * joint.local_axis2;
+            axis1 = rb1.position * joint.local_axis1;
             #[cfg(feature = "dim2")]
             {
-                basis1 = rb1.position * cparams.basis1[0];
+                basis1 = rb1.position * joint.basis1[0];
             }
             #[cfg(feature = "dim3")]
             {
                 basis1 = Matrix3x2::from_columns(&[
-                    rb1.position * cparams.basis1[0],
-                    rb1.position * cparams.basis1[1],
+                    rb1.position * joint.basis1[0],
+                    rb1.position * joint.basis1[1],
                 ]);
             }
         };
@@ -438,26 +497,62 @@ impl PrismaticVelocityGroundConstraint {
         #[cfg(feature = "dim3")]
         let rhs = Vector5::new(lin_rhs.x, lin_rhs.y, ang_rhs.x, ang_rhs.y, ang_rhs.z);
 
-        // Setup limit constraint.
+        /*
+         * Setup motor.
+         */
+
+        /*
+         * Setup motor.
+         */
+        let mut motor_rhs = 0.0;
+        let mut motor_inv_lhs = 0.0;
+        let mut 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 dist = anchor2.coords.dot(&axis2) - anchor1.coords.dot(&axis1);
+            motor_rhs += (dist - joint.motor_target_pos) * stiffness;
+        }
+
+        if damping != 0.0 {
+            let curr_vel = rb2.linvel.dot(&axis2) - rb1.linvel.dot(&axis1);
+            motor_rhs += (curr_vel - joint.motor_target_vel) * damping;
+        }
+
+        if stiffness != 0.0 || damping != 0.0 {
+            motor_inv_lhs = if keep_lhs { gamma / im2 } else { gamma };
+            motor_rhs /= gamma;
+        }
+
+        let motor_impulse = na::clamp(joint.motor_impulse, -motor_max_impulse, motor_max_impulse);
+
+        /*
+         * Setup limit constraint.
+         */
         let mut limits_forcedir2 = None;
         let mut limits_rhs = 0.0;
         let mut limits_impulse = 0.0;
 
-        if cparams.limits_enabled {
+        if joint.limits_enabled {
             let danchor = anchor2 - anchor1;
             let dist = danchor.dot(&axis1);
 
-            // FIXME: we should allow both limits to be active at
+            // TODO: we should allow both limits to be active at
             // the same time.
-            // FIXME: allow predictive constraint activation.
-            if dist < cparams.limits[0] {
+            // TODO: allow predictive constraint activation.
+            if dist < joint.limits[0] {
                 limits_forcedir2 = Some(axis2.into_inner());
                 limits_rhs = anchor_linvel2.dot(&axis2) - anchor_linvel1.dot(&axis1);
-                limits_impulse = cparams.limits_impulse;
-            } else if dist > cparams.limits[1] {
+                limits_impulse = joint.limits_impulse;
+            } else if dist > joint.limits[1] {
                 limits_forcedir2 = Some(-axis2.into_inner());
                 limits_rhs = -anchor_linvel2.dot(&axis2) + anchor_linvel1.dot(&axis1);
-                limits_impulse = cparams.limits_impulse;
+                limits_impulse = joint.limits_impulse;
             }
         }
 
@@ -466,8 +561,12 @@ impl PrismaticVelocityGroundConstraint {
             mj_lambda2: rb2.active_set_offset,
             im2,
             ii2_sqrt: rb2.effective_world_inv_inertia_sqrt,
-            impulse: cparams.impulse * params.warmstart_coeff,
+            impulse: joint.impulse * params.warmstart_coeff,
             limits_impulse: limits_impulse * params.warmstart_coeff,
+            motor_rhs,
+            motor_inv_lhs,
+            motor_impulse,
+            motor_max_impulse,
             basis1,
             inv_lhs,
             rhs,
@@ -492,6 +591,10 @@ impl PrismaticVelocityGroundConstraint {
             .ii2_sqrt
             .transform_vector(ang_impulse + self.r2.gcross(lin_impulse));
 
+        // Warmstart motors.
+        mj_lambda2.linear -= self.axis2 * (self.im2 * self.motor_impulse);
+
+        // Warmstart limits.
         if let Some(limits_forcedir2) = self.limits_forcedir2 {
             mj_lambda2.linear += limits_forcedir2 * (self.im2 * self.limits_impulse);
         }
@@ -503,7 +606,7 @@ impl PrismaticVelocityGroundConstraint {
         let mut mj_lambda2 = mj_lambdas[self.mj_lambda2 as usize];
 
         /*
-         * Joint consraint.
+         * Joint constraint.
          */
         let ang_vel2 = self.ii2_sqrt.transform_vector(mj_lambda2.angular);
         let lin_vel2 = mj_lambda2.linear + ang_vel2.gcross(self.r2);
@@ -528,11 +631,25 @@ impl PrismaticVelocityGroundConstraint {
             .transform_vector(ang_impulse + self.r2.gcross(lin_impulse));
 
         /*
+         * Motors.
+         */
+        if self.motor_inv_lhs != 0.0 {
+            let lin_dvel = self.axis2.dot(&mj_lambda2.linear) + self.motor_rhs;
+            let new_impulse = na::clamp(
+                self.motor_impulse + lin_dvel * self.motor_inv_lhs,
+                -self.motor_max_impulse,
+                self.motor_max_impulse,
+            );
+            let dimpulse = new_impulse - self.motor_impulse;
+            self.motor_impulse = new_impulse;
+
+            mj_lambda2.linear -= self.axis2 * (self.im2 * dimpulse);
+        }
+
+        /*
          * Joint limits.
          */
         if let Some(limits_forcedir2) = self.limits_forcedir2 {
-            // FIXME: the transformation by ii2_sqrt could be avoided by
-            // reusing some computations above.
             let ang_vel2 = self.ii2_sqrt.transform_vector(mj_lambda2.angular);
 
             let lin_dvel = limits_forcedir2.dot(&(mj_lambda2.linear + ang_vel2.gcross(self.r2)))
@@ -547,11 +664,12 @@ impl PrismaticVelocityGroundConstraint {
         mj_lambdas[self.mj_lambda2 as usize] = mj_lambda2;
     }
 
-    // FIXME: duplicated code with the non-ground constraint.
+    // TODO: duplicated code with the non-ground constraint.
     pub fn writeback_impulses(&self, joints_all: &mut [JointGraphEdge]) {
         let joint = &mut joints_all[self.joint_id].weight;
         if let JointParams::PrismaticJoint(revolute) = &mut joint.params {
             revolute.impulse = self.impulse;
+            revolute.motor_impulse = self.motor_impulse;
             revolute.limits_impulse = self.limits_impulse;
         }
     }