Support impedance control on MuJoCo POSITION actuators

mujoco_ros2_control/include/mujoco_ros2_control/actuator_parameters.hpp

@@ -0,0 +1,38 @@
+/**
+ * Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+#pragma once
+
+namespace mujoco_ros2_control
+{
+
+/** The four coefficients used by MuJoCo's affine actuator gain/bias model. */
+struct ActuatorAffineParameters
+{
+  double gain{ 0.0 };
+  double bias_constant{ 0.0 };
+  double bias_position{ 0.0 };
+  double bias_velocity{ 0.0 };
+
+  bool operator==(const ActuatorAffineParameters& other) const
+  {
+    return gain == other.gain && bias_constant == other.bias_constant && bias_position == other.bias_position &&
+           bias_velocity == other.bias_velocity;
+  }
+
+  bool operator!=(const ActuatorAffineParameters& other) const
+  {
+    return !(*this == other);
+  }
+};
+
+/** A complete affine-parameter update for one MuJoCo actuator. */
+struct ActuatorParameterUpdate
+{
+  int actuator_id;
+  ActuatorAffineParameters parameters;
+};
+
+}  // namespace mujoco_ros2_control

mujoco_ros2_control/include/mujoco_ros2_control/data.hpp

@@ -27,7 +27,10 @@
 #include <Eigen/Geometry>
 #include <hardware_interface/types/hardware_interface_type_values.hpp>
 #include "control_toolbox/pid_ros.hpp"
+#include "mujoco_ros2_control/actuator_parameters.hpp"
 
+#include <limits>
+#include <optional>
 #include <string>
 #include <vector>
 
@@ -121,9 +124,21 @@ struct MuJoCoActuatorData
   bool is_velocity_control_enabled{ false };
   bool is_effort_control_enabled{ false };
   bool is_impedance_control_enabled{ false };  // Uses commanded kp/kd for PD control
+  // Impedance on a MuJoCo POSITION actuator is realised by the actuator itself
+  // rather than by summing a torque in write(): see MujocoSystemInterface::write.
+  bool uses_native_position_impedance{ false };
   bool has_pos_pid{ false };
   bool has_vel_pid{ false };
 
+  // Gains as authored in the MJCF, captured at registration for POSITION
+  // actuators so that native-impedance mode can overwrite them and restore
+  // them when the joint is released.
+  std::optional<ActuatorAffineParameters> authored_parameters;
+
+  // Last native-impedance parameters staged for the physics thread, used to skip
+  // redundant updates.
+  std::optional<ActuatorAffineParameters> last_requested_parameters;
+
   void copy_state_to_transmission()
   {
     position_interface.transmission_passthrough_ = position_interface.state_;

mujoco_ros2_control/include/mujoco_ros2_control/mujoco_simulation.hpp

@@ -49,6 +49,8 @@
 #include <mujoco_ros2_control_plugins/mujoco_ros2_control_plugins_base.hpp>
 #include <sensor_msgs/msg/joint_state.hpp>
 
+#include "mujoco_ros2_control/actuator_parameters.hpp"
+
 namespace mujoco_ros2_control
 {
 
@@ -307,10 +309,22 @@ class MujocoSimulation
    * from the simulate app. `control_data->qvel` is copied into `qvel_override_staged_`;
    * entries that are not NaN are plugin-requested velocity overrides (see
    * `MuJoCoROS2ControlPluginBase::update()`'s doc comment for the NaN convention this
-   * relies on) and are applied as direct `qvel` writes before the next `mj_step`. This does
-   * not lock the sim mutex and never waits on stepping.
+   * relies on) and are applied as direct `qvel` writes before the next `mj_step`.
+   * Complete affine actuator-parameter sets are staged through the optional second argument
+   * and applied under the same lock as `ctrl`, so a step cannot observe a new setpoint with
+   * old coefficients (or vice versa). This does not lock the sim mutex and never waits on
+   * stepping.
    */
-  void apply_control_data(mjData* control_data);
+  void apply_control_data(mjData* control_data,
+                          const std::vector<ActuatorParameterUpdate>& actuator_parameter_updates = {});
+
+  /**
+   * @brief Apply a cold-path actuator parameter change synchronously.
+   *
+   * Used for mode/lifecycle transitions that must be observable before the
+   * call returns. Any pending staged update for the actuator is discarded.
+   */
+  void set_actuator_parameters(const ActuatorParameterUpdate& update);
 
   /**
    * @brief Accessor for the mutex which locks access to the data and model.
@@ -339,6 +353,9 @@ class MujocoSimulation
    */
   void apply_staged_control_inputs();
 
+  /** Compose viewer and plugin Cartesian forces without applying other staged inputs. */
+  void compose_cartesian_forces();
+
   /**
    * @brief Publishes the per-step control state from `mj_data_`.
    *
@@ -508,6 +525,10 @@ class MujocoSimulation
   // update()
   std::vector<mjtNum> qvel_override_staged_;
 
+  // Complete affine-parameter sets waiting to be applied with the staged ctrl
+  // immediately before the next physics step.
+  std::vector<ActuatorParameterUpdate> actuator_parameter_updates_staged_;
+
   // Guards only the snapshot pointer swap and snapshot_ready_ flag.
   // Lock order: sim_mutex_ (if needed) is always taken before this one.
   std::mutex data_exchange_mutex_;
@@ -518,9 +539,9 @@ class MujocoSimulation
   std::atomic<bool> snapshot_refresh_requested_{ true };
 
   // Guards the staged control inputs (ctrl_staged_, qfrc_applied_staged_, xfrc_plugin_desired_,
-  // control_inputs_staged_, qvel_override_staged_). Separate from data_exchange_mutex_ so
-  // that staging commands in write() and applying them before each physics step never queue
-  // behind a full mjData copy.
+  // control_inputs_staged_, qvel_override_staged_, actuator_parameter_updates_staged_).
+  // Separate from data_exchange_mutex_ so that staging commands in write() and applying them
+  // before each physics step never queue behind a full mjData copy.
   // Critical sections are all small buffer copies.
   // Lock order: sim_mutex_ (if needed) before this one; never held with data_exchange_mutex_.
   std::mutex control_staging_mutex_;

mujoco_ros2_control/src/mujoco_simulation.cpp

@@ -22,6 +22,7 @@
 #include "mujoco_ros2_control/sim_display_text.hpp"
 
 #include <unistd.h>
+#include <algorithm>
 #include <cerrno>
 #include <chrono>
 #include <cmath>
@@ -957,6 +958,7 @@ void MujocoSimulation::reset_world_state(bool fill_initial_state,
     std::fill(qfrc_applied_staged_.begin(), qfrc_applied_staged_.end(), 0.0);
     std::fill(xfrc_plugin_desired_.begin(), xfrc_plugin_desired_.end(), 0.0);
     std::fill(qvel_override_staged_.begin(), qvel_override_staged_.end(), std::numeric_limits<mjtNum>::quiet_NaN());
+    actuator_parameter_updates_staged_.clear();
   }
   std::fill(xfrc_viewer_capture_.begin(), xfrc_viewer_capture_.end(), 0.0);
   std::fill(xfrc_last_written_.begin(), xfrc_last_written_.end(), 0.0);
@@ -975,18 +977,16 @@ void MujocoSimulation::reset_world_state(bool fill_initial_state,
     apply_free_joint_states(state_overrides.free_joints);
   }
 
-  // Run forward dynamics to update derived quantities
-  mj_forward(mj_model_, mj_data_);
-
-  // Make the reset state visible to snapshot readers before HW-side bookkeeping runs
-  refresh_data_snapshot();
-  publish_control_state();
-
-  // Delegate HW-side bookkeeping (PID resets, command/state interface sync, etc.)
+  // Restore hardware-owned model state (including authored actuator parameters) before
+  // deriving forces and sensors from the reset state.
   if (reset_callback_)
   {
     reset_callback_(fill_initial_state);
   }
+
+  mj_forward(mj_model_, mj_data_);
+  refresh_data_snapshot();
+  publish_control_state();
 }
 
 void MujocoSimulation::reset_world_callback(
@@ -1430,16 +1430,69 @@ mjData* MujocoSimulation::acquire_data_snapshot()
   return snapshot_read_;
 }
 
-void MujocoSimulation::apply_control_data(mjData* control_data)
+namespace
 {
+void validate_actuator_id(const mjModel* model, int actuator_id)
+{
+  if (actuator_id < 0 || actuator_id >= model->nu)
+  {
+    throw std::out_of_range("Actuator parameter update ID is outside the MuJoCo model actuator range");
+  }
+}
+
+void apply_actuator_parameters(mjModel* model, const ActuatorParameterUpdate& update)
+{
+  const int gain_offset = update.actuator_id * mjNGAIN;
+  const int bias_offset = update.actuator_id * mjNBIAS;
+  model->actuator_gainprm[gain_offset] = update.parameters.gain;
+  model->actuator_biasprm[bias_offset] = update.parameters.bias_constant;
+  model->actuator_biasprm[bias_offset + 1] = update.parameters.bias_position;
+  model->actuator_biasprm[bias_offset + 2] = update.parameters.bias_velocity;
+}
+}  // namespace
+
+void MujocoSimulation::apply_control_data(mjData* control_data,
+                                          const std::vector<ActuatorParameterUpdate>& actuator_parameter_updates)
+{
+  for (const auto& update : actuator_parameter_updates)
+  {
+    validate_actuator_id(mj_model_, update.actuator_id);
+  }
+
   const std::lock_guard<std::mutex> lock(control_staging_mutex_);
   mju_copy(ctrl_staged_.data(), control_data->ctrl, static_cast<int>(mj_model_->nu));
   mju_copy(qfrc_applied_staged_.data(), control_data->qfrc_applied, static_cast<int>(mj_model_->nv));
   mju_copy(xfrc_plugin_desired_.data(), control_data->xfrc_applied, 6 * static_cast<int>(mj_model_->nbody));
   mju_copy(qvel_override_staged_.data(), control_data->qvel, static_cast<int>(mj_model_->nv));
+  for (const auto& update : actuator_parameter_updates)
+  {
+    auto pending =
+        std::find_if(actuator_parameter_updates_staged_.begin(), actuator_parameter_updates_staged_.end(),
+                     [&update](const auto& candidate) { return candidate.actuator_id == update.actuator_id; });
+    if (pending == actuator_parameter_updates_staged_.end())
+    {
+      actuator_parameter_updates_staged_.push_back(update);
+    }
+    else
+    {
+      *pending = update;
+    }
+  }
   control_inputs_staged_ = true;
 }
 
+void MujocoSimulation::set_actuator_parameters(const ActuatorParameterUpdate& update)
+{
+  validate_actuator_id(mj_model_, update.actuator_id);
+  const std::lock_guard<std::recursive_mutex> sim_lock(*sim_mutex_);
+  const std::lock_guard<std::mutex> staging_lock(control_staging_mutex_);
+  auto& staged = actuator_parameter_updates_staged_;
+  staged.erase(std::remove_if(staged.begin(), staged.end(),
+                              [&update](const auto& pending) { return pending.actuator_id == update.actuator_id; }),
+               staged.end());
+  apply_actuator_parameters(mj_model_, update);
+}
+
 void MujocoSimulation::refresh_data_snapshot()
 {
   {
@@ -1477,16 +1530,29 @@ void MujocoSimulation::copy_control_state(ControlState& destination)
 
 void MujocoSimulation::apply_staged_control_inputs()
 {
-  const std::lock_guard<std::mutex> lock(control_staging_mutex_);
-
-  // Only apply ctrl / qfrc_applied once the hw interface has staged control inputs,
-  // so that initial or reset values in mj_data_ are not overwritten with zeros.
-  if (control_inputs_staged_)
   {
-    mju_copy(mj_data_->ctrl, ctrl_staged_.data(), static_cast<int>(mj_model_->nu));
-    mju_copy(mj_data_->qfrc_applied, qfrc_applied_staged_.data(), static_cast<int>(mj_model_->nv));
+    const std::lock_guard<std::mutex> lock(control_staging_mutex_);
+
+    // Only apply ctrl / qfrc_applied once the hw interface has staged control inputs,
+    // so that initial or reset values in mj_data_ are not overwritten with zeros.
+    if (control_inputs_staged_)
+    {
+      for (const auto& update : actuator_parameter_updates_staged_)
+      {
+        apply_actuator_parameters(mj_model_, update);
+      }
+      actuator_parameter_updates_staged_.clear();
+      mju_copy(mj_data_->ctrl, ctrl_staged_.data(), static_cast<int>(mj_model_->nu));
+      mju_copy(mj_data_->qfrc_applied, qfrc_applied_staged_.data(), static_cast<int>(mj_model_->nv));
+    }
   }
 
+  compose_cartesian_forces();
+}
+
+void MujocoSimulation::compose_cartesian_forces()
+{
+  const std::lock_guard<std::mutex> lock(control_staging_mutex_);
   const int nbody6 = 6 * static_cast<int>(mj_model_->nbody);
   mju_copy(mj_data_->xfrc_applied, xfrc_viewer_capture_.data(), nbody6);
   mju_addTo(mj_data_->xfrc_applied, xfrc_plugin_desired_.data(), nbody6);
@@ -1711,7 +1777,6 @@ void MujocoSimulation::physics_loop()
           // save current state to history buffer
           if (stepped)
           {
-            apply_staged_control_inputs();
             sim_->AddToHistory();
             update_sim_display();
           }
@@ -1727,14 +1792,11 @@ void MujocoSimulation::physics_loop()
             pending_steps_.fetch_add(1);
           }
 
-          // Record so the next iteration can detect render thread changes, only necessary once
-          // when paused
-          apply_staged_control_inputs();
-
           // Execute one pending step per physics loop iteration so the clock publisher
           // (try_publish) has time to flush between steps, matching play mode behavior.
           if (pending_steps_.load() > 0)
           {
+            apply_staged_control_inputs();
             mj_step(mj_model_, mj_data_);
             publish_control_state();
             publish_clock();
@@ -1763,6 +1825,9 @@ void MujocoSimulation::physics_loop()
           }
           else
           {
+            // Keep viewer/plugin Cartesian forces visible to paused rendering without
+            // making new controls or actuator coefficients observable before a step.
+            compose_cartesian_forces();
             mj_forward(mj_model_, mj_data_);
             sim_->speed_changed = true;
             update_sim_display();