mujoco: Add support for Universal and Screw Joints

dartsim/src/JointFeatures.cc

@@ -34,27 +34,43 @@ namespace dartsim {
 double JointFeatures::GetJointPosition(
     const Identity &_id, std::size_t _dof) const
 {
+  if (!this->ValidateDofParam(_id, _dof))
+  {
+    return math::NAN_D;
+  }
   return this->ReferenceInterface<JointInfo>(_id)->joint->getPosition(_dof);
 }
 
 /////////////////////////////////////////////////
 double JointFeatures::GetJointVelocity(
     const Identity &_id, std::size_t _dof) const
 {
+  if (!this->ValidateDofParam(_id, _dof))
+  {
+    return math::NAN_D;
+  }
   return this->ReferenceInterface<JointInfo>(_id)->joint->getVelocity(_dof);
 }
 
 /////////////////////////////////////////////////
 double JointFeatures::GetJointAcceleration(
     const Identity &_id, std::size_t _dof) const
 {
+  if (!this->ValidateDofParam(_id, _dof))
+  {
+    return math::NAN_D;
+  }
   return this->ReferenceInterface<JointInfo>(_id)->joint->getAcceleration(_dof);
 }
 
 /////////////////////////////////////////////////
 double JointFeatures::GetJointForce(
     const Identity &_id, std::size_t _dof) const
 {
+  if (!this->ValidateDofParam(_id, _dof))
+  {
+    return math::NAN_D;
+  }
   return this->ReferenceInterface<JointInfo>(_id)->joint->getForce(_dof);
 }
 
@@ -70,6 +86,10 @@ void JointFeatures::SetJointPosition(
     const Identity &_id, std::size_t _dof, double _value)
 {
   auto joint = this->ReferenceInterface<JointInfo>(_id)->joint;
+  if (!this->ValidateDofParam(_id, _dof))
+  {
+    return;
+  }
 
   // Take extra care that the value is finite. A nan can cause the DART
   // constraint solver to fail, which will in turn either cause a crash or
@@ -89,6 +109,10 @@ void JointFeatures::SetJointVelocity(
     const Identity &_id, std::size_t _dof, double _value)
 {
   auto joint = this->ReferenceInterface<JointInfo>(_id)->joint;
+  if (!this->ValidateDofParam(_id, _dof))
+  {
+    return;
+  }
 
   // Take extra care that the value is finite. A nan can cause the DART
   // constraint solver to fail, which will in turn either cause a crash or
@@ -109,6 +133,10 @@ void JointFeatures::SetJointAcceleration(
 {
   auto joint = this->ReferenceInterface<JointInfo>(_id)->joint;
 
+  if (!this->ValidateDofParam(_id, _dof))
+  {
+    return;
+  }
   // Take extra care that the value is finite. A nan can cause the DART
   // constraint solver to fail, which will in turn either cause a crash or
   // collisions to fail
@@ -128,6 +156,10 @@ void JointFeatures::SetJointForce(
 {
   auto joint = this->ReferenceInterface<JointInfo>(_id)->joint;
 
+  if (!this->ValidateDofParam(_id, _dof))
+  {
+    return;
+  }
   // Take extra care that the value is finite. A nan can cause the DART
   // constraint solver to fail, which will in turn either cause a crash or
   // collisions to fail
@@ -409,6 +441,20 @@ void JointFeatures::SetJointSpringReference(
   joint->setRestPosition(_dof, _value);
 }
 
+/////////////////////////////////////////////////
+bool JointFeatures::ValidateDofParam(const Identity &_id,
+                                     std::size_t _dof) const
+{
+  const auto jointInfo = this->ReferenceInterface<JointInfo>(_id);
+  if (_dof >= jointInfo->joint->getNumDofs())
+  {
+    gzerr << "Trying to access an invalid DOF [" << _dof << "] on joint ["
+          << jointInfo->joint->getName() << "]\n";
+    return false;
+  }
+  return true;
+}
+
 /////////////////////////////////////////////////
 std::size_t JointFeatures::GetJointDegreesOfFreedom(const Identity &_id) const
 {

dartsim/src/JointFeatures.hh

@@ -105,6 +105,9 @@ class JointFeatures :
   public: std::size_t GetJointDegreesOfFreedom(
       const Identity &_id) const override;
 
+  private: bool ValidateDofParam(
+      const Identity &_id, std::size_t _dof) const;
+
   public: Pose3d GetJointTransformFromParent(
       const Identity &_id) const override;
 

mujoco/src/Base.cc

@@ -76,6 +76,17 @@ void resolveJointIndices(WorldInfo &_worldInfo)
         continue;
       }
       jointInfo->nv_index = qvelAddr;
+
+      // Resolve screw constraint compiled ID
+      jointInfo->screwEqIndex = std::nullopt;
+      if (jointInfo->screwConstraintSpec)
+      {
+        int eqId = mjs_getId(jointInfo->screwConstraintSpec->element);
+        if (eqId >= 0 && eqId < m->neq)
+        {
+          jointInfo->screwEqIndex = eqId;
+        }
+      }
     }
   }
 

mujoco/src/Base.hh

@@ -168,6 +168,10 @@ struct JointInfo
   std::weak_ptr<ModelInfo> modelInfo;
   WorldInfo* worldInfo{nullptr};
   std::optional<std::size_t> ballJointCacheIndex{std::nullopt};
+  // Pointer to the screw joint equality constraint in the spec
+  mjsEquality* screwConstraintSpec{nullptr};
+  // Compiled screw joint equality constraint index in mjModel (max 1)
+  std::optional<int> screwEqIndex{std::nullopt};
 };
 
 

mujoco/src/JointFeatures.cc

@@ -38,7 +38,7 @@ Eigen::Vector3d *getBallJointPositionImpl(JointInfo *jointInfo)
 
   if (!jointInfo->ballJointCacheIndex)
   {
-    gzerr << "Ball joint cache index is is set for ["
+    gzerr << "Ball joint cache index is not set for ["
           << jointInfo->name << "]\n";
     return nullptr;
   }
@@ -114,6 +114,24 @@ void setJointPositionImpl(JointInfo *jointInfo, std::size_t _dof, double _value)
     d->qpos[jointInfo->nq_index + _dof] = _value;
   }
 }
+
+void updateScrewJointFollower(
+    JointInfo *_jointInfo, double _value, double *_mjDataArray, int _baseIndex)
+{
+  if (!_jointInfo->screwEqIndex.has_value())
+    return;
+
+  auto *m = _jointInfo->worldInfo->mjModelObj;
+  const int eqId = *_jointInfo->screwEqIndex;
+  const double multiplier = m->eq_data[eqId * mjNEQDATA + 1];
+
+  // Since the primary rotational hinge and secondary translational slide
+  // joints are created sequentially on the child body back-to-back in
+  // SDFFeatures.cc, and MuJoCo compiles a body's joints contiguously in
+  // memory, the follower's state index is guaranteed to be exactly
+  // _baseIndex + 1.
+  _mjDataArray[_baseIndex + 1] = _value * multiplier;
+}
 }
 
 /////////////////////////////////////////////////
@@ -260,6 +278,15 @@ void JointFeatures::SetJointPosition(
     return;
   }
   setJointPositionImpl(jointInfo, _dof, _value);
+
+  // If this is the primary rotational hinge joint of a screw joint,
+  // simultaneously update the coupled secondary slide joint position (`qpos`)
+  // by `_value * pitch`. This maintains perfect kinematic consistency and
+  // prevents violent solver impulses during state initialization.
+  updateScrewJointFollower(
+      jointInfo, _value, jointInfo->worldInfo->mjDataObj->qpos,
+      jointInfo->nq_index);
+
   mj_forward(jointInfo->worldInfo->mjModelObj, jointInfo->worldInfo->mjDataObj);
 }
 
@@ -293,6 +320,15 @@ void JointFeatures::SetJointVelocity(
     return;
   }
   jointInfo->worldInfo->mjDataObj->qvel[jointInfo->nv_index + _dof] = _value;
+
+  // If this is the primary rotational hinge joint of a screw joint,
+  // simultaneously update the coupled secondary slide joint velocity (`qvel`)
+  // by `_value * pitch`. This maintains perfect kinematic consistency and
+  // prevents violent solver impulses during state initialization.
+  updateScrewJointFollower(jointInfo, _value,
+                           jointInfo->worldInfo->mjDataObj->qvel,
+                           jointInfo->nv_index);
+
   mj_forward(jointInfo->worldInfo->mjModelObj, jointInfo->worldInfo->mjDataObj);
 }
 
@@ -370,7 +406,13 @@ std::size_t JointFeatures::GetJointDegreesOfFreedom(const Identity &_id) const
   int bodyId = mjs_getId(jointInfo->childBody->element);
   if (bodyId < 0 || bodyId >= m->nbody)
     return 0;
-  return m->body_dofnum[bodyId];
+
+  std::size_t dofs = m->body_dofnum[bodyId];
+  if (jointInfo->screwEqIndex.has_value())
+  {
+    dofs -= 1;
+  }
+  return dofs;
 }
 
 /////////////////////////////////////////////////

mujoco/src/SDFFeatures.cc

@@ -154,6 +154,15 @@ void copyStandardJointAxisProperties(
   }
 }
 /////////////////////////////////////////////////
+/// \brief Convert SDFormat screw thread pitch into MuJoCo coordinate ratio.
+/// \param[in] _pitch Screw thread pitch defined in meters per revolution
+/// (m/rev).
+/// \return Pitch coordinate ratio in meters per radian (m/rad).
+double convertScrewThreadPitch(const double _pitch)
+{
+  return _pitch / (2.0 * GZ_PI);
+}
+/////////////////////////////////////////////////
 struct ModelKinematicStructure
 {
   std::string name;
@@ -287,6 +296,7 @@ struct ModelKinematicStructure
     else
     {
       mjsJoint *joint{nullptr};
+      mjsJoint *joint2{nullptr};
       // It is possible to apply joint forces using `qfrc_applied`, but this
       // makes it harder to retrieve the last applied forces on a joint when
       // implementing GetJoint. Instead, we use actuators and `mjData::ctrl`.
@@ -328,6 +338,82 @@ struct ModelKinematicStructure
           }
         }
       }
+      else if (sdfJoint->Type() == ::sdf::JointType::SCREW)
+      {
+        // Screw joints in MuJoCo are modeled by coupling a hinge joint
+        // (rotation) and a slide joint (translation) along the same axis on
+        // the child body using a joint equality constraint (mjEQ_JOINT).
+        // We store the hinge joint (`joint`) as the primary joint in
+        // JointInfo. This matches DART's choice of using the rotational DOF
+        // as the primary, ensuring both physics plugins expose consistent
+        // angular units (radians and rad/s) for screw joints across the
+        // gz-physics API.
+        // Like the universal joint, `joint` and `joint2` are compiled
+        // contiguously on the same child body.
+        joint = mjs_addJoint(child, nullptr);
+        joint->type = mjJNT_HINGE;
+        const auto *sdfAxis1 = sdfJoint->Axis(0);
+        if (sdfAxis1)
+        {
+          convertJointAxis(sdfAxis1, joint->axis);
+          copyStandardJointAxisProperties(joint, sdfAxis1);
+          // Disable independent position limits on the primary rotational
+          // hinge joint. In MuJoCo's soft constraint solver, enforcing limits
+          // on the hinge joint when coupled with a small thread pitch causes
+          // premature solver clamping and massive numerical damping. By mapping
+          // position limits purely onto the translational slide joint, we
+          // ensure exact kinematic limit enforcement without solver resistance.
+          joint->limited = false;
+        }
+
+        joint2 = mjs_addJoint(child, nullptr);
+        joint2->type = mjJNT_SLIDE;
+        if (sdfAxis1)
+        {
+          convertJointAxis(sdfAxis1, joint2->axis);
+          // We only copy position limits and range to the secondary slide
+          // joint. All passive dynamics (damping, frictionloss, stiffness) and
+          // actuator effort limits are enforced purely on the primary
+          // rotational hinge joint to avoid double-counting and physical unit
+          // mismatches.
+          joint2->limited = static_cast<int>(!std::isinf(sdfAxis1->Lower()) &&
+                                             !std::isinf(sdfAxis1->Upper()));
+          if (joint2->limited)
+          {
+            const double pitch =
+                convertScrewThreadPitch(sdfJoint->ScrewThreadPitch());
+            joint2->range[0] = sdfAxis1->Lower() * pitch;
+            joint2->range[1] = sdfAxis1->Upper() * pitch;
+          }
+        }
+      }
+      else if (sdfJoint->Type() == ::sdf::JointType::UNIVERSAL)
+      {
+        // Universal joints in MuJoCo are modeled as two hinge joints in
+        // series on the child body. We only need to keep a pointer to the
+        // first hinge joint (`joint`) in JointInfo. Since MuJoCo compiles
+        // joints on the same body contiguously in memory, all getters/setters
+        // can safely access the second joint (`joint2`)'s state data using
+        // `nq_index + 1` and `nv_index + 1` without needing to store it
+        // separately in JointInfo.
+        joint = mjs_addJoint(child, nullptr);
+        joint->type = mjJNT_HINGE;
+        const auto *sdfAxis1 = sdfJoint->Axis(0);
+        if (sdfAxis1)
+        {
+          convertJointAxis(sdfAxis1, joint->axis);
+          copyStandardJointAxisProperties(joint, sdfAxis1);
+        }
+
+        joint2 = mjs_addJoint(child, nullptr);
+        joint2->type = mjJNT_HINGE;
+        const auto *sdfAxis2 = sdfJoint->Axis(1);
+        if (sdfAxis2)
+        {
+          convertJointAxis(sdfAxis2, joint2->axis);
+          copyStandardJointAxisProperties(joint2, sdfAxis2);
+        }
+      }
       else if (sdfJoint->Type() != ::sdf::JointType::FIXED)
       {
         gzwarn << "Joint type " << static_cast<int>(sdfJoint->Type())
@@ -339,6 +425,7 @@ struct ModelKinematicStructure
       // it's associated. Note that this body is the child link of the joint
       // in SDF terms.
       auto jointPose = resolveSdfPose(sdfJoint->SemanticPose());
+      mjsEquality *eq = nullptr;
       // Note that no joints will be created when processing a fixed joint.
       if (joint)
       {
@@ -351,6 +438,41 @@ struct ModelKinematicStructure
         mjs_setString(actuator->target, mjJointName.c_str());
 
         copyPos(jointPose.Pos(), joint->pos);
+
+        if (joint2)
+        {
+          // We uniquely name the second axis using a flat `_axis2` suffix.
+          // This flat name avoids indicating any Kinematic/SDF nesting
+          // (`::axis2`) while still satisfying MuJoCo's requirement that
+          // joints must be uniquely named for actuators to target them.
+          const std::string mjJointName2 = mjJointName + "_axis2";
+          mjs_setName(joint2->element, mjJointName2.c_str());
+          mjsActuator *actuator2 = mjs_addActuator(_spec, nullptr);
+          actuator2->trntype = mjtTrn::mjTRN_JOINT;
+          mjs_setString(actuator2->target, mjJointName2.c_str());
+
+          copyPos(jointPose.Pos(), joint2->pos);
+
+          // If this is a screw joint, couple the slide and hinge axes using
+          // a joint equality constraint (mjEQ_JOINT) with the specified pitch.
+          if (sdfJoint->Type() == ::sdf::JointType::SCREW)
+          {
+            eq = mjs_addEquality(_spec, nullptr);
+            eq->type = mjEQ_JOINT;
+            eq->active = 1;
+            mjs_setString(eq->name1, mjJointName2.c_str());
+            mjs_setString(eq->name2, mjJointName.c_str());
+
+            // dif = pos[1] - ref[1] = hinge_pos - hinge_ref
+            // cpos = pos[0] - ref[0] - data[0] - data[1]*dif = 0
+            // enforces: slide_pos - slide_ref =
+            // data[1] * (hinge_pos - hinge_ref)
+            // where data[1] = pitch (meters/rad) = ScrewThreadPitch/2pi
+            std::fill(std::begin(eq->data), std::end(eq->data), 0.0);
+            eq->data[1] =
+                convertScrewThreadPitch(sdfJoint->ScrewThreadPitch());
+          }
+        }
       }
       auto jointInfo =
           std::make_shared<JointInfo>(_base.GetNextEntity(), _modelInfo);
@@ -359,6 +481,10 @@ struct ModelKinematicStructure
       jointInfo->childBody = child;
       jointInfo->actuator = actuator;
       jointInfo->worldInfo = worldInfo;
+      if (sdfJoint->Type() == ::sdf::JointType::SCREW)
+      {
+        jointInfo->screwConstraintSpec = eq;
+      }
       if (sdfJoint->Type() == ::sdf::JointType::BALL)
       {
         jointInfo->worldInfo->ballJointPositionsCache.push_back(std::nullopt);