Implement limits for revolute joints.

3 files changed, 233 insertions, 6 deletions

diff --git a/src/dynamics/solver/joint_constraint/prismatic_velocity_constraint.rs b/src/dynamics/solver/joint_constraint/prismatic_velocity_constraint.rs
index 1e7570e..4e94c79 100644
--- a/src/dynamics/solver/joint_constraint/prismatic_velocity_constraint.rs
+++ b/src/dynamics/solver/joint_constraint/prismatic_velocity_constraint.rs
@@ -256,8 +256,8 @@ impl PrismaticVelocityConstraint {
 
             limits_impulse_limits.0 = if max_enabled { -Real::INFINITY } else { 0.0 };
             limits_impulse_limits.1 = if min_enabled { Real::INFINITY } else { 0.0 };
-
             limits_active = min_enabled || max_enabled;
+
             if limits_active {
                 limits_rhs = (anchor_linvel2.dot(&axis2) - anchor_linvel1.dot(&axis1))
                     * params.velocity_solve_fraction;
diff --git a/src/dynamics/solver/joint_constraint/revolute_position_constraint.rs b/src/dynamics/solver/joint_constraint/revolute_position_constraint.rs
index 731d0e2..ad470a4 100644
--- a/src/dynamics/solver/joint_constraint/revolute_position_constraint.rs
+++ b/src/dynamics/solver/joint_constraint/revolute_position_constraint.rs
@@ -21,6 +21,11 @@ pub(crate) struct RevolutePositionConstraint {
 
     ang_inv_lhs: AngularInertia<Real>,
 
+    limits_enabled: bool,
+    limits: [Real; 2],
+
+    motor_last_angle: Real,
+
     local_anchor1: Point<Real>,
     local_anchor2: Point<Real>,
 
@@ -61,6 +66,9 @@ impl RevolutePositionConstraint {
             position2: ids2.active_set_offset,
             local_basis1: cparams.basis1,
             local_basis2: cparams.basis2,
+            limits_enabled: cparams.limits_enabled,
+            limits: cparams.limits,
+            motor_last_angle: cparams.motor_last_angle,
         }
     }
 
@@ -119,6 +127,31 @@ impl RevolutePositionConstraint {
                 Rotation::new(self.ii2.transform_vector(-ang_impulse)) * position2.rotation;
         }
 
+        /*
+         * Limits part.
+         */
+        if self.limits_enabled {
+            let angle = RevoluteJoint::estimate_motor_angle_from_params(
+                &(position1 * self.local_axis1),
+                &(position1 * self.local_basis1[0]),
+                &(position2 * self.local_basis2[0]),
+                self.motor_last_angle,
+            );
+            let ang_error = (angle - self.limits[1]).max(0.0) - (self.limits[0] - angle).max(0.0);
+
+            if ang_error != 0.0 {
+                let axis2 = position2 * self.local_axis2;
+                let ang_impulse = self
+                    .ang_inv_lhs
+                    .transform_vector(*axis2 * ang_error * params.joint_erp);
+
+                position1.rotation =
+                    Rotation::new(self.ii1.transform_vector(ang_impulse)) * position1.rotation;
+                position2.rotation =
+                    Rotation::new(self.ii2.transform_vector(-ang_impulse)) * position2.rotation;
+            }
+        }
+
         positions[self.position1 as usize] = position1;
         positions[self.position2 as usize] = position2;
     }
@@ -133,9 +166,14 @@ pub(crate) struct RevolutePositionGroundConstraint {
     anchor1: Point<Real>,
     local_anchor2: Point<Real>,
     axis1: Unit<Vector<Real>>,
-    local_axis2: Unit<Vector<Real>>,
-
     basis1: [Vector<Real>; 2],
+
+    limits_enabled: bool,
+    limits: [Real; 2],
+
+    motor_last_angle: Real,
+
+    local_axis2: Unit<Vector<Real>>,
     local_basis2: [Vector<Real>; 2],
 }
 
@@ -189,6 +227,9 @@ impl RevolutePositionGroundConstraint {
             position2: ids2.active_set_offset,
             basis1,
             local_basis2,
+            limits_enabled: cparams.limits_enabled,
+            limits: cparams.limits,
+            motor_last_angle: cparams.motor_last_angle,
         }
     }
 
@@ -230,6 +271,25 @@ impl RevolutePositionGroundConstraint {
             position2.rotation = Rotation::new(-ang_error) * position2.rotation;
         }
 
+        /*
+         * Limits part.
+         */
+        if self.limits_enabled {
+            let angle = RevoluteJoint::estimate_motor_angle_from_params(
+                &self.axis1,
+                &self.basis1[0],
+                &(position2 * self.local_basis2[0]),
+                self.motor_last_angle,
+            );
+            let ang_error = (angle - self.limits[1]).max(0.0) - (self.limits[0] - angle).max(0.0);
+
+            if ang_error != 0.0 {
+                let axis2 = position2 * self.local_axis2;
+                position2.rotation =
+                    Rotation::new(*axis2 * -ang_error * params.joint_erp) * position2.rotation;
+            }
+        }
+
         positions[self.position2 as usize] = position2;
     }
 }
diff --git a/src/dynamics/solver/joint_constraint/revolute_velocity_constraint.rs b/src/dynamics/solver/joint_constraint/revolute_velocity_constraint.rs
index bf45c11..46114c5 100644
--- a/src/dynamics/solver/joint_constraint/revolute_velocity_constraint.rs
+++ b/src/dynamics/solver/joint_constraint/revolute_velocity_constraint.rs
@@ -30,6 +30,12 @@ pub(crate) struct RevoluteVelocityConstraint {
     motor_axis1: Vector<Real>,
     motor_axis2: Vector<Real>,
 
+    limits_active: bool,
+    limits_impulse: Real,
+    limits_rhs: Real,
+    limits_inv_lhs: Real,
+    limits_impulse_limits: (Real, Real),
+
     basis1: Matrix3x2<Real>,
     basis2: Matrix3x2<Real>,
 
@@ -145,8 +151,11 @@ impl RevoluteVelocityConstraint {
             joint.motor_damping,
         );
 
-        if stiffness != 0.0 {
+        if stiffness != 0.0 || joint.limits_enabled {
             motor_angle = joint.estimate_motor_angle(&poss1.position, &poss2.position);
+        }
+
+        if stiffness != 0.0 {
             motor_rhs += (motor_angle - joint.motor_target_pos) * stiffness;
         }
 
@@ -168,6 +177,47 @@ impl RevoluteVelocityConstraint {
         }
 
         /*
+         * Setup limit constraint.
+         */
+        let mut limits_active = false;
+        let mut limits_rhs = 0.0;
+        let mut limits_inv_lhs = 0.0;
+        let mut limits_impulse_limits = (0.0, 0.0);
+        let mut limits_impulse = 0.0;
+
+        if joint.limits_enabled {
+            // TODO: we should allow predictive constraint activation.
+
+            let (min_limit, max_limit) = (joint.limits[0], joint.limits[1]);
+            let min_enabled = motor_angle < min_limit;
+            let max_enabled = max_limit < motor_angle;
+
+            limits_impulse_limits.0 = if max_enabled { -Real::INFINITY } else { 0.0 };
+            limits_impulse_limits.1 = if min_enabled { Real::INFINITY } else { 0.0 };
+            limits_active = min_enabled || max_enabled;
+
+            if limits_active {
+                limits_rhs = (vels2.angvel.dot(&motor_axis2) - vels1.angvel.dot(&motor_axis1))
+                    * params.velocity_solve_fraction;
+
+                limits_rhs += ((motor_angle - max_limit).max(0.0)
+                    - (min_limit - motor_angle).max(0.0))
+                    * velocity_based_erp_inv_dt;
+
+                limits_inv_lhs = crate::utils::inv(
+                    motor_axis2.dot(&ii2.transform_vector(motor_axis2))
+                        + motor_axis1.dot(&ii1.transform_vector(motor_axis1)),
+                );
+
+                limits_impulse = joint
+                    .limits_impulse
+                    .max(limits_impulse_limits.0)
+                    .min(limits_impulse_limits.1)
+                    * params.warmstart_coeff;
+            }
+        }
+
+        /*
          * Adjust the warmstart impulse.
          * If the velocity along the free axis is somewhat high,
          * we need to adjust the angular warmstart impulse because it
@@ -205,6 +255,11 @@ impl RevoluteVelocityConstraint {
             motor_axis2,
             motor_impulse,
             motor_angle,
+            limits_impulse,
+            limits_impulse_limits,
+            limits_active,
+            limits_inv_lhs,
+            limits_rhs,
         }
     }
 
@@ -228,7 +283,7 @@ impl RevoluteVelocityConstraint {
             .transform_vector(ang_impulse2 + self.r2.gcross(lin_impulse2));
 
         /*
-         * Motor
+         * Warmstart motor.
          */
         if self.motor_inv_lhs != 0.0 {
             mj_lambda1.angular += self
@@ -239,6 +294,14 @@ impl RevoluteVelocityConstraint {
                 .transform_vector(self.motor_axis2 * self.motor_impulse);
         }
 
+        // Warmstart limits.
+        if self.limits_active {
+            let limit_impulse1 = -self.motor_axis2 * self.limits_impulse;
+            let limit_impulse2 = self.motor_axis2 * self.limits_impulse;
+            mj_lambda1.angular += self.ii1_sqrt.transform_vector(limit_impulse1);
+            mj_lambda2.angular += self.ii2_sqrt.transform_vector(limit_impulse2);
+        }
+
         mj_lambdas[self.mj_lambda1 as usize] = mj_lambda1;
         mj_lambdas[self.mj_lambda2 as usize] = mj_lambda2;
     }
@@ -270,6 +333,31 @@ impl RevoluteVelocityConstraint {
             .transform_vector(ang_impulse2 + self.r2.gcross(lin_impulse2));
     }
 
+    fn solve_limits(&mut self, mj_lambda1: &mut DeltaVel<Real>, mj_lambda2: &mut DeltaVel<Real>) {
+        if self.limits_active {
+            let limits_torquedir1 = -self.motor_axis2;
+            let limits_torquedir2 = self.motor_axis2;
+
+            let ang_vel1 = self.ii1_sqrt.transform_vector(mj_lambda1.angular);
+            let ang_vel2 = self.ii2_sqrt.transform_vector(mj_lambda2.angular);
+
+            let ang_dvel = limits_torquedir1.dot(&ang_vel1)
+                + limits_torquedir2.dot(&ang_vel2)
+                + self.limits_rhs;
+            let new_impulse = (self.limits_impulse - ang_dvel * self.limits_inv_lhs)
+                .max(self.limits_impulse_limits.0)
+                .min(self.limits_impulse_limits.1);
+            let dimpulse = new_impulse - self.limits_impulse;
+            self.limits_impulse = new_impulse;
+
+            let ang_impulse1 = limits_torquedir1 * dimpulse;
+            let ang_impulse2 = limits_torquedir2 * dimpulse;
+
+            mj_lambda1.angular += self.ii1_sqrt.transform_vector(ang_impulse1);
+            mj_lambda2.angular += self.ii2_sqrt.transform_vector(ang_impulse2);
+        }
+    }
+
     fn solve_motors(&mut self, mj_lambda1: &mut DeltaVel<Real>, mj_lambda2: &mut DeltaVel<Real>) {
         if self.motor_inv_lhs != 0.0 {
             let ang_vel1 = self.ii1_sqrt.transform_vector(mj_lambda1.angular);
@@ -294,6 +382,7 @@ impl RevoluteVelocityConstraint {
         let mut mj_lambda1 = mj_lambdas[self.mj_lambda1 as usize];
         let mut mj_lambda2 = mj_lambdas[self.mj_lambda2 as usize];
 
+        self.solve_limits(&mut mj_lambda1, &mut mj_lambda2);
         self.solve_dofs(&mut mj_lambda1, &mut mj_lambda2);
         self.solve_motors(&mut mj_lambda1, &mut mj_lambda2);
 
@@ -310,6 +399,7 @@ impl RevoluteVelocityConstraint {
             revolute.prev_axis1 = self.motor_axis1;
             revolute.motor_last_angle = self.motor_angle;
             revolute.motor_impulse = self.motor_impulse;
+            revolute.limits_impulse = self.limits_impulse;
         }
     }
 }
@@ -333,6 +423,12 @@ pub(crate) struct RevoluteVelocityGroundConstraint {
     motor_max_impulse: Real,
     motor_angle: Real, // Exists just for writing it into the joint.
 
+    limits_active: bool,
+    limits_impulse: Real,
+    limits_rhs: Real,
+    limits_inv_lhs: Real,
+    limits_impulse_limits: (Real, Real),
+
     basis2: Matrix3x2<Real>,
 
     im2: Real,
@@ -458,8 +554,11 @@ impl RevoluteVelocityGroundConstraint {
             joint.motor_damping,
         );
 
-        if stiffness != 0.0 {
+        if stiffness != 0.0 || joint.limits_enabled {
             motor_angle = joint.estimate_motor_angle(&poss1.position, &poss2.position);
+        }
+
+        if stiffness != 0.0 {
             motor_rhs += (motor_angle - joint.motor_target_pos) * stiffness;
         }
 
@@ -480,6 +579,43 @@ impl RevoluteVelocityGroundConstraint {
         let motor_impulse = na::clamp(joint.motor_impulse, -motor_max_impulse, motor_max_impulse)
             * params.warmstart_coeff;
 
+        /*
+         * Setup limit constraint.
+         */
+        let mut limits_active = false;
+        let mut limits_rhs = 0.0;
+        let mut limits_inv_lhs = 0.0;
+        let mut limits_impulse_limits = (0.0, 0.0);
+        let mut limits_impulse = 0.0;
+
+        if joint.limits_enabled {
+            // TODO: we should allow predictive constraint activation.
+            let (min_limit, max_limit) = (joint.limits[0], joint.limits[1]);
+            let min_enabled = motor_angle < min_limit;
+            let max_enabled = max_limit < motor_angle;
+
+            limits_impulse_limits.0 = if max_enabled { -Real::INFINITY } else { 0.0 };
+            limits_impulse_limits.1 = if min_enabled { Real::INFINITY } else { 0.0 };
+            limits_active = min_enabled || max_enabled;
+
+            if limits_active {
+                limits_rhs = (vels2.angvel.dot(&axis2) - vels1.angvel.dot(&axis1))
+                    * params.velocity_solve_fraction;
+
+                limits_rhs += ((motor_angle - max_limit).max(0.0)
+                    - (min_limit - motor_angle).max(0.0))
+                    * velocity_based_erp_inv_dt;
+
+                limits_inv_lhs = crate::utils::inv(axis2.dot(&ii2.transform_vector(axis2)));
+
+                limits_impulse = joint
+                    .limits_impulse
+                    .max(limits_impulse_limits.0)
+                    .min(limits_impulse_limits.1)
+                    * params.warmstart_coeff;
+            }
+        }
+
         let result = RevoluteVelocityGroundConstraint {
             joint_id,
             mj_lambda2: ids2.active_set_offset,
@@ -496,6 +632,11 @@ impl RevoluteVelocityGroundConstraint {
             motor_max_impulse,
             motor_rhs,
             motor_angle,
+            limits_impulse,
+            limits_impulse_limits,
+            limits_active,
+            limits_inv_lhs,
+            limits_rhs,
         };
 
         AnyJointVelocityConstraint::RevoluteGroundConstraint(result)
@@ -521,6 +662,12 @@ impl RevoluteVelocityGroundConstraint {
                 .transform_vector(self.motor_axis2 * self.motor_impulse);
         }
 
+        // Warmstart limits.
+        if self.limits_active {
+            let limit_impulse2 = self.motor_axis2 * self.limits_impulse;
+            mj_lambda2.angular += self.ii2_sqrt.transform_vector(limit_impulse2);
+        }
+
         mj_lambdas[self.mj_lambda2 as usize] = mj_lambda2;
     }
 
@@ -542,6 +689,24 @@ impl RevoluteVelocityGroundConstraint {
             .transform_vector(ang_impulse + self.r2.gcross(lin_impulse));
     }
 
+    fn solve_limits(&mut self, mj_lambda2: &mut DeltaVel<Real>) {
+        if self.limits_active {
+            let limits_torquedir2 = self.motor_axis2;
+
+            let ang_vel2 = self.ii2_sqrt.transform_vector(mj_lambda2.angular);
+
+            let ang_dvel = limits_torquedir2.dot(&ang_vel2) + self.limits_rhs;
+            let new_impulse = (self.limits_impulse - ang_dvel * self.limits_inv_lhs)
+                .max(self.limits_impulse_limits.0)
+                .min(self.limits_impulse_limits.1);
+            let dimpulse = new_impulse - self.limits_impulse;
+            self.limits_impulse = new_impulse;
+
+            let ang_impulse2 = limits_torquedir2 * dimpulse;
+            mj_lambda2.angular += self.ii2_sqrt.transform_vector(ang_impulse2);
+        }
+    }
+
     fn solve_motors(&mut self, mj_lambda2: &mut DeltaVel<Real>) {
         if self.motor_inv_lhs != 0.0 {
             let ang_vel2 = self.ii2_sqrt.transform_vector(mj_lambda2.angular);
@@ -563,6 +728,7 @@ impl RevoluteVelocityGroundConstraint {
     pub fn solve(&mut self, mj_lambdas: &mut [DeltaVel<Real>]) {
         let mut mj_lambda2 = mj_lambdas[self.mj_lambda2 as usize];
 
+        self.solve_limits(&mut mj_lambda2);
         self.solve_dofs(&mut mj_lambda2);
         self.solve_motors(&mut mj_lambda2);
 
@@ -576,6 +742,7 @@ impl RevoluteVelocityGroundConstraint {
             revolute.impulse = self.impulse;
             revolute.motor_impulse = self.motor_impulse;
             revolute.motor_last_angle = self.motor_angle;
+            revolute.limits_impulse = self.limits_impulse;
         }
     }
 }