Joint API and joint motors improvements

3 files changed, 48 insertions, 23 deletions

diff --git a/src/dynamics/joint/multibody_joint/multibody_joint.rs b/src/dynamics/joint/multibody_joint/multibody_joint.rs
index 1d14680..2f7a71e 100644
--- a/src/dynamics/joint/multibody_joint/multibody_joint.rs
+++ b/src/dynamics/joint/multibody_joint/multibody_joint.rs
@@ -1,6 +1,6 @@
 use crate::dynamics::solver::AnyJointVelocityConstraint;
 use crate::dynamics::{
-    joint, FixedJoint, IntegrationParameters, JointData, Multibody, MultibodyLink,
+    joint, FixedJointBuilder, GenericJoint, IntegrationParameters, Multibody, MultibodyLink,
     RigidBodyVelocity,
 };
 use crate::math::{
@@ -14,13 +14,13 @@ use na::{UnitQuaternion, Vector3};
 #[cfg_attr(feature = "serde-serialize", derive(Serialize, Deserialize))]
 #[derive(Copy, Clone, Debug)]
 pub struct MultibodyJoint {
-    pub data: JointData,
+    pub data: GenericJoint,
     pub(crate) coords: SpacialVector<Real>,
     pub(crate) joint_rot: Rotation<Real>,
 }
 
 impl MultibodyJoint {
-    pub fn new(data: JointData) -> Self {
+    pub fn new(data: GenericJoint) -> Self {
         Self {
             data,
             coords: na::zero(),
@@ -29,13 +29,13 @@ impl MultibodyJoint {
     }
 
     pub(crate) fn free(pos: Isometry<Real>) -> Self {
-        let mut result = Self::new(JointData::default());
+        let mut result = Self::new(GenericJoint::default());
         result.set_free_pos(pos);
         result
     }
 
     pub(crate) fn fixed(pos: Isometry<Real>) -> Self {
-        Self::new(FixedJoint::new().local_frame1(pos).into())
+        Self::new(FixedJointBuilder::new().local_frame1(pos).build().into())
     }
 
     pub(crate) fn set_free_pos(&mut self, pos: Isometry<Real>) {
@@ -263,13 +263,20 @@ impl MultibodyJoint {
 
         for i in 0..DIM {
             if (locked_bits & (1 << i)) == 0 {
-                if (limit_bits & (1 << i)) != 0 {
-                    joint::unit_joint_limit_constraint(
+                let limits = if (limit_bits & (1 << i)) != 0 {
+                    Some([self.data.limits[i].min, self.data.limits[i].max])
+                } else {
+                    None
+                };
+
+                if (motor_bits & (1 << i)) != 0 {
+                    joint::unit_joint_motor_constraint(
                         params,
                         multibody,
                         link,
-                        [self.data.limits[i].min, self.data.limits[i].max],
+                        &self.data.motors[i],
                         self.coords[i],
+                        limits,
                         dof_id + curr_free_dof,
                         j_id,
                         jacobians,
@@ -277,12 +284,12 @@ impl MultibodyJoint {
                     );
                 }
 
-                if (motor_bits & (1 << i)) != 0 {
-                    joint::unit_joint_motor_constraint(
+                if (limit_bits & (1 << i)) != 0 {
+                    joint::unit_joint_limit_constraint(
                         params,
                         multibody,
                         link,
-                        &self.data.motors[i],
+                        [self.data.limits[i].min, self.data.limits[i].max],
                         self.coords[i],
                         dof_id + curr_free_dof,
                         j_id,
@@ -310,19 +317,23 @@ impl MultibodyJoint {
         // TODO: we should make special cases for multi-angular-dofs limits/motors
         for i in DIM..SPATIAL_DIM {
             if (locked_bits & (1 << i)) == 0 {
-                if (limit_bits & (1 << i)) != 0 {
+                let limits = if (limit_bits & (1 << i)) != 0 {
+                    let limits = [self.data.limits[i].min, self.data.limits[i].max];
                     joint::unit_joint_limit_constraint(
                         params,
                         multibody,
                         link,
-                        [self.data.limits[i].min, self.data.limits[i].max],
+                        limits,
                         self.coords[i],
                         dof_id + curr_free_dof,
                         j_id,
                         jacobians,
                         constraints,
                     );
-                }
+                    Some(limits)
+                } else {
+                    None
+                };
 
                 if (motor_bits & (1 << i)) != 0 {
                     joint::unit_joint_motor_constraint(
@@ -331,6 +342,7 @@ impl MultibodyJoint {
                         link,
                         &self.data.motors[i],
                         self.coords[i],
+                        limits,
                         dof_id + curr_free_dof,
                         j_id,
                         jacobians,
diff --git a/src/dynamics/joint/multibody_joint/multibody_joint_set.rs b/src/dynamics/joint/multibody_joint/multibody_joint_set.rs
index 3e786e2..7e512a8 100644
--- a/src/dynamics/joint/multibody_joint/multibody_joint_set.rs
+++ b/src/dynamics/joint/multibody_joint/multibody_joint_set.rs
@@ -1,7 +1,7 @@
 use crate::data::{Arena, Coarena, ComponentSet, ComponentSetMut, Index};
 use crate::dynamics::joint::MultibodyLink;
 use crate::dynamics::{
-    IslandManager, JointData, Multibody, MultibodyJoint, RigidBodyActivation, RigidBodyHandle,
+    GenericJoint, IslandManager, Multibody, MultibodyJoint, RigidBodyActivation, RigidBodyHandle,
     RigidBodyIds, RigidBodyType,
 };
 use crate::geometry::{InteractionGraph, RigidBodyGraphIndex};
@@ -112,7 +112,7 @@ impl MultibodyJointSet {
         &mut self,
         body1: RigidBodyHandle,
         body2: RigidBodyHandle,
-        data: impl Into<JointData>,
+        data: impl Into<GenericJoint>,
     ) -> Option<MultibodyJointHandle> {
         let data = data.into();
         let link1 = self.rb2mb.get(body1.0).copied().unwrap_or_else(|| {
diff --git a/src/dynamics/joint/multibody_joint/unit_multibody_joint.rs b/src/dynamics/joint/multibody_joint/unit_multibody_joint.rs
index 3367108..42212be 100644
--- a/src/dynamics/joint/multibody_joint/unit_multibody_joint.rs
+++ b/src/dynamics/joint/multibody_joint/unit_multibody_joint.rs
@@ -26,7 +26,8 @@ pub fn unit_joint_limit_constraint(
 
     let min_enabled = curr_pos < limits[0];
     let max_enabled = limits[1] < curr_pos;
-    let erp_inv_dt = params.erp_inv_dt();
+    let erp_inv_dt = params.joint_erp_inv_dt();
+    let cfm_coeff = params.joint_cfm_coeff();
     let rhs_bias = ((curr_pos - limits[1]).max(0.0) - (limits[0] - curr_pos).max(0.0)) * erp_inv_dt;
     let rhs_wo_bias = joint_velocity[dof_id];
 
@@ -54,6 +55,8 @@ pub fn unit_joint_limit_constraint(
         inv_lhs: crate::utils::inv(lhs),
         rhs: rhs_wo_bias + rhs_bias,
         rhs_wo_bias,
+        cfm_coeff,
+        cfm_gain: 0.0,
         writeback_id: WritebackId::Limit(dof_id),
     };
 
@@ -71,11 +74,13 @@ pub fn unit_joint_motor_constraint(
     link: &MultibodyLink,
     motor: &JointMotor,
     curr_pos: Real,
+    limits: Option<[Real; 2]>,
     dof_id: usize,
     j_id: &mut usize,
     jacobians: &mut DVector<Real>,
     constraints: &mut Vec<AnyJointVelocityConstraint>,
 ) {
+    let inv_dt = params.inv_dt();
     let ndofs = multibody.ndofs();
     let joint_velocity = multibody.joint_velocity(link);
 
@@ -93,14 +98,20 @@ pub fn unit_joint_motor_constraint(
     let impulse_bounds = [-motor_params.max_impulse, motor_params.max_impulse];
 
     let mut rhs_wo_bias = 0.0;
-    if motor_params.stiffness != 0.0 {
-        rhs_wo_bias += (curr_pos - motor_params.target_pos) * motor_params.stiffness;
+    if motor_params.erp_inv_dt != 0.0 {
+        rhs_wo_bias += (curr_pos - motor_params.target_pos) * motor_params.erp_inv_dt;
     }
 
-    if motor_params.damping != 0.0 {
-        let dvel = joint_velocity[dof_id];
-        rhs_wo_bias += (dvel - motor_params.target_vel) * motor_params.damping;
-    }
+    let mut target_vel = motor_params.target_vel;
+    if let Some(limits) = limits {
+        target_vel = target_vel.clamp(
+            (limits[0] - curr_pos) * inv_dt,
+            (limits[1] - curr_pos) * inv_dt,
+        );
+    };
+
+    let dvel = joint_velocity[dof_id];
+    rhs_wo_bias += dvel - target_vel;
 
     let constraint = JointGenericVelocityGroundConstraint {
         mj_lambda2: multibody.solver_id,
@@ -109,6 +120,8 @@ pub fn unit_joint_motor_constraint(
         joint_id: usize::MAX,
         impulse: 0.0,
         impulse_bounds,
+        cfm_coeff: motor_params.cfm_coeff,
+        cfm_gain: motor_params.cfm_gain,
         inv_lhs: crate::utils::inv(lhs),
         rhs: rhs_wo_bias,
         rhs_wo_bias,