Visualize virtual gantry in MuJoCo viewer

mujoco_ros2_control/CMakeLists.txt

@@ -133,9 +133,27 @@ else()
     OUTPUT_NAME simulate
     POSITION_INDEPENDENT_CODE ON
   )
+  set(MUJOCO_SIMULATE_PATCHED_SOURCE "${CMAKE_CURRENT_BINARY_DIR}/generated/simulate.cc")
+  set(MUJOCO_SIMULATE_PATCH_FILE "${CMAKE_CURRENT_SOURCE_DIR}/patches/mujoco_simulate_managed_user_scn.patch")
+  find_program(PATCH_EXECUTABLE patch REQUIRED)
+  file(MAKE_DIRECTORY "${CMAKE_CURRENT_BINARY_DIR}/generated")
+  configure_file("${MUJOCO_SIMULATE_DIR}/simulate.cc" "${MUJOCO_SIMULATE_PATCHED_SOURCE}" COPYONLY)
+  execute_process(
+    COMMAND "${PATCH_EXECUTABLE}" -p1 --forward --batch --input "${MUJOCO_SIMULATE_PATCH_FILE}"
+    WORKING_DIRECTORY "${CMAKE_CURRENT_BINARY_DIR}/generated"
+    RESULT_VARIABLE MUJOCO_SIMULATE_PATCH_RESULT
+    OUTPUT_VARIABLE MUJOCO_SIMULATE_PATCH_OUTPUT
+    ERROR_VARIABLE MUJOCO_SIMULATE_PATCH_ERROR
+  )
+  if(NOT MUJOCO_SIMULATE_PATCH_RESULT EQUAL 0)
+    message(FATAL_ERROR
+      "Could not patch MuJoCo Simulate for managed-mode user geoms:\n"
+      "${MUJOCO_SIMULATE_PATCH_OUTPUT}\n${MUJOCO_SIMULATE_PATCH_ERROR}"
+    )
+  endif()
   target_sources(mujoco_simulate
     PRIVATE ${MUJOCO_SIMULATE_DIR}/simulate.h
-    PRIVATE ${MUJOCO_SIMULATE_DIR}/simulate.cc ${MUJOCO_SIMULATE_DIR}/array_safety.h
+    PRIVATE ${MUJOCO_SIMULATE_PATCHED_SOURCE} ${MUJOCO_SIMULATE_DIR}/array_safety.h
   )
   target_include_directories(mujoco_simulate PUBLIC ${MUJOCO_SIMULATE_DIR} ${MUJOCO_INCLUDE_DIR})
   target_link_libraries(mujoco_simulate PRIVATE lodepng mujoco::platform_ui_adapter ${MUJOCO_LIB})

mujoco_ros2_control/patches/mujoco_simulate_managed_user_scn.patch

@@ -0,0 +1,88 @@
+--- a/simulate.cc
++++ b/simulate.cc
+@@ -2161,6 +2161,27 @@ void Simulate::Sync(bool state_only) {
+   // update scene or sync data from user in passive mode
+   if (!is_passive_) {
+     mjv_updateScene(m_, d_, &this->opt, &this->pert, &this->cam, mjCAT_ALL, &this->scn);
++
++    // append geoms from user_scn to scratch space
++    if (user_scn) {
++      user_scn_geoms_.clear();
++      user_scn_geoms_.reserve(user_scn->ngeom);
++      for (int i = 0; i < user_scn->ngeom; ++i) {
++        user_scn_geoms_.push_back(user_scn->geoms[i]);
++      }
++    }
++
++    // pick up rendering flags changed via user_scn
++    if (user_scn) {
++      for (int i = 0; i < mjNRNDFLAG; ++i) {
++        if (user_scn->flags[i] != user_scn_flags_prev_[i]) {
++          scn.flags[i] = user_scn->flags[i];
++          pending_.ui_update_rendering = true;
++        }
++      }
++      Copy(user_scn->flags, scn.flags);
++      Copy(user_scn_flags_prev_, user_scn->flags);
++    }
+   } else {
+     if (state_only) {
+       int state_size = mj_stateSize(m_, mjSTATE_INTEGRATION);
+@@ -2400,6 +2421,27 @@ void Simulate::LoadOnRenderThread() {
+   // update scene in managed mode, in passive mode copy data from user (update in RenderLoop)
+   if (!is_passive_) {
+     mjv_updateScene(this->m_, this->d_, &this->opt, &this->pert, &this->cam, mjCAT_ALL, &this->scn);
++
++    // append geoms from user_scn to scratch space
++    if (user_scn) {
++      user_scn_geoms_.clear();
++      user_scn_geoms_.reserve(user_scn->ngeom);
++      for (int i = 0; i < user_scn->ngeom; ++i) {
++        user_scn_geoms_.push_back(user_scn->geoms[i]);
++      }
++    }
++
++    // pick up rendering flags changed via user_scn
++    if (user_scn) {
++      for (int i = 0; i < mjNRNDFLAG; ++i) {
++        if (user_scn->flags[i] != user_scn_flags_prev_[i]) {
++          scn.flags[i] = user_scn->flags[i];
++          pending_.ui_update_rendering = true;
++        }
++      }
++      Copy(user_scn->flags, scn.flags);
++      Copy(user_scn_flags_prev_, user_scn->flags);
++    }
+   } else {
+     mjopt_prev_ = m_->opt;
+     opt_prev_ = opt;
+@@ -2864,6 +2906,19 @@ void Simulate::RenderLoop() {
+       // update scene, doing a full sync if in fully managed mode
+       if (!is_passive_) {
+         Sync();
++
++        // add user geoms to scene
++        int nusergeom = user_scn_geoms_.size();
++        int ngeom = std::min(nusergeom, this->scn.maxgeom - this->scn.ngeom);
++        if (ngeom < nusergeom) {
++          mj_warning(this->d_, mjWARN_VGEOMFULL, this->scn.maxgeom);
++        }
++        std::memcpy(this->scn.geoms + this->scn.ngeom, user_scn_geoms_.data(),
++                    ngeom * sizeof(mjvGeom));
++        for (int i = 0; i < ngeom; ++i) {
++          this->scn.geomorder[this->scn.ngeom + i] = this->scn.ngeom + i;
++        }
++        this->scn.ngeom += ngeom;
+       } else if (m_passive_ && d_passive_) {
+         // the user has called Sync() in their code
+         mjv_updateScene(m_passive_, d_passive_,
+@@ -2877,6 +2932,9 @@ void Simulate::RenderLoop() {
+         }
+         std::memcpy(this->scn.geoms + this->scn.ngeom, user_scn_geoms_.data(),
+                     ngeom * sizeof(mjvGeom));
++        for (int i = 0; i < ngeom; ++i) {
++          this->scn.geomorder[this->scn.ngeom + i] = this->scn.ngeom + i;
++        }
+         this->scn.ngeom += ngeom;
+       }
+     }  // MutexLock (unblocks simulation thread)

mujoco_ros2_control/src/mujoco_system_interface.cpp

@@ -78,6 +78,59 @@ using Seconds = std::chrono::duration<double>;
 
 namespace
 {
+struct PluginVisualizationSceneState
+{
+  PluginVisualizationSceneState()
+  {
+    mjv_defaultScene(&scene);
+  }
+
+  ~PluginVisualizationSceneState()
+  {
+    if (initialized)
+    {
+      mjv_freeScene(&scene);
+    }
+  }
+
+  mjvScene scene;
+  bool initialized{ false };
+};
+
+std::mutex plugin_visualization_scene_states_mutex;
+std::unordered_map<const mujoco_ros2_control::MujocoSystemInterface*, std::unique_ptr<PluginVisualizationSceneState>>
+    plugin_visualization_scene_states;
+
+PluginVisualizationSceneState&
+get_plugin_visualization_scene_state(const mujoco_ros2_control::MujocoSystemInterface* system_interface)
+{
+  std::lock_guard<std::mutex> lock(plugin_visualization_scene_states_mutex);
+  auto& state = plugin_visualization_scene_states[system_interface];
+  if (!state)
+  {
+    state = std::make_unique<PluginVisualizationSceneState>();
+  }
+  return *state;
+}
+
+PluginVisualizationSceneState*
+find_plugin_visualization_scene_state(const mujoco_ros2_control::MujocoSystemInterface* system_interface)
+{
+  std::lock_guard<std::mutex> lock(plugin_visualization_scene_states_mutex);
+  auto state = plugin_visualization_scene_states.find(system_interface);
+  if (state == plugin_visualization_scene_states.end())
+  {
+    return nullptr;
+  }
+  return state->second.get();
+}
+
+void erase_plugin_visualization_scene_state(const mujoco_ros2_control::MujocoSystemInterface* system_interface)
+{
+  std::lock_guard<std::mutex> lock(plugin_visualization_scene_states_mutex);
+  plugin_visualization_scene_states.erase(system_interface);
+}
+
 std::optional<std::string> get_hardware_parameter(const hardware_interface::HardwareInfo& hardware_info,
                                                   const std::string& key)
 {
@@ -739,6 +792,8 @@ MujocoSystemInterface::~MujocoSystemInterface()
   }
   plugin_instances_.clear();
 
+  erase_plugin_visualization_scene_state(this);
+
   // Cleanup data and the model, if they haven't been
   if (mj_data_)
   {
@@ -951,6 +1006,16 @@ MujocoSystemInterface::on_init(const hardware_interface::HardwareComponentInterf
     std::unique_lock<std::recursive_mutex> lock(*sim_mutex_);
     mj_data_ = mj_makeData(mj_model_);
     mj_data_control_ = mj_makeData(mj_model_);
+    if (!headless_)
+    {
+      auto& plugin_visualization_scene_state = get_plugin_visualization_scene_state(this);
+      if (!plugin_visualization_scene_state.initialized)
+      {
+        mjv_makeScene(mj_model_, &plugin_visualization_scene_state.scene, 32);
+      }
+      plugin_visualization_scene_state.initialized = true;
+      sim_->user_scn = &plugin_visualization_scene_state.scene;
+    }
   }
   if (!mj_data_ || !mj_data_control_)
   {
@@ -3327,6 +3392,17 @@ void MujocoSystemInterface::update_sim_display()
     return;
   }
 
+  auto* plugin_visualization_scene_state = find_plugin_visualization_scene_state(this);
+  if (plugin_visualization_scene_state && plugin_visualization_scene_state->initialized)
+  {
+    mj::MutexLock lock(sim_->mtx);
+    plugin_visualization_scene_state->scene.ngeom = 0;
+    for (auto& plugin : plugin_instances_)
+    {
+      plugin->update_visualization(mj_model_, mj_data_control_, &plugin_visualization_scene_state->scene);
+    }
+  }
+
   // Only write user_texts_new_ when the render thread has consumed the previous
   // update (newtextrequest == 0). Use compare_exchange to atomically claim the
   // slot: if it fails, the render thread hasn't swapped yet, so skip this

mujoco_ros2_control_plugins/CMakeLists.txt

@@ -59,6 +59,10 @@ install(TARGETS mujoco_ros2_control_plugins
 
 if(BUILD_TESTING)
   find_package(ament_cmake_gtest REQUIRED)
+  ament_add_gtest(test_plugin_visualization test/test_plugin_visualization.cpp)
+  target_link_libraries(test_plugin_visualization
+    mujoco_ros2_control_plugins
+  )
   # find_package(ament_lint_auto REQUIRED)
   # ament_lint_auto_find_test_dependencies()
 endif()

mujoco_ros2_control_plugins/include/mujoco_ros2_control_plugins/mujoco_ros2_control_plugins_base.hpp

@@ -73,6 +73,18 @@ class MuJoCoROS2ControlPluginBase
   {
   }
 
+  /**
+   * @brief Called when the MuJoCo viewer display is refreshed.
+   * @param model Pointer to the MuJoCo model
+   * @param data Pointer to the MuJoCo data
+   * @param scene Scene where the plugin can append viewer-only geoms
+   * @note Override to add lightweight viewer visualization. This must not mutate
+   *       physics state.
+   */
+  virtual void update_visualization(const mjModel* /*model*/, const mjData* /*data*/, mjvScene* /*scene*/)
+  {
+  }
+
   /**
    * @brief Called from the GLFW UI thread on every key event.
    * @return true if the event was consumed (caller should not forward to the MuJoCo viewer).

mujoco_ros2_control_plugins/include/mujoco_ros2_control_plugins/virtual_gantry_plugin.hpp

@@ -38,9 +38,12 @@ class VirtualGantryPlugin : public MuJoCoROS2ControlPluginBase
   void update(const mjModel* model, mjData* data) override;
   void reset() override;
   void cleanup() override;
+  void update_visualization(const mjModel* model, const mjData* data, mjvScene* scene) override;
   bool on_key(int key, int scancode, int action, int mods) override;
 
 private:
+  friend class VirtualGantryPluginTestAccessor;
+
   rclcpp::Node::SharedPtr node_;
 
   // Target body (attachment point on the robot).

mujoco_ros2_control_plugins/src/virtual_gantry_plugin.cpp

@@ -219,6 +219,46 @@ void VirtualGantryPlugin::reset()
   last_update_time_ = -1.0;
 }
 
+void VirtualGantryPlugin::update_visualization(const mjModel* /*model*/, const mjData* data, mjvScene* scene)
+{
+  std::lock_guard<std::mutex> lock(state_mutex_);
+
+  if (!enabled_ || !spawn_pos_captured_ || body_id_ < 0 || !data || !data->xpos || !data->xmat || !scene ||
+      !scene->geoms || scene->ngeom >= scene->maxgeom)
+  {
+    return;
+  }
+
+  const double* xpos = &data->xpos[body_id_ * 3];
+  const double* xmat = &data->xmat[body_id_ * 9];
+
+  mjtNum attach_pos[3];
+  for (int i = 0; i < 3; ++i)
+  {
+    attach_pos[i] = xpos[i] + xmat[i * 3 + 0] * body_offset_[0] + xmat[i * 3 + 1] * body_offset_[1] +
+                    xmat[i * 3 + 2] * body_offset_[2];
+  }
+
+  const mjtNum anchor_pos[3] = { anchor_pos_[0], anchor_pos_[1], anchor_pos_[2] };
+  const float rope_rgba[4] = { 0.1f, 0.95f, 0.2f, 1.0f };
+
+  mjvGeom* geom = &scene->geoms[scene->ngeom++];
+  mjv_initGeom(geom, mjGEOM_CAPSULE, nullptr, nullptr, nullptr, rope_rgba);
+  geom->category = mjCAT_DECOR;
+  mjv_connector(geom, mjGEOM_CAPSULE, 0.0125, anchor_pos, attach_pos);
+
+  if (scene->ngeom >= scene->maxgeom)
+  {
+    return;
+  }
+
+  const mjtNum anchor_size[3] = { 0.04, 0.0, 0.0 };
+  const float anchor_rgba[4] = { 1.0f, 0.0f, 0.0f, 1.0f };
+  mjvGeom* anchor_geom = &scene->geoms[scene->ngeom++];
+  mjv_initGeom(anchor_geom, mjGEOM_SPHERE, anchor_size, anchor_pos, nullptr, anchor_rgba);
+  anchor_geom->category = mjCAT_DECOR;
+}
+
 bool VirtualGantryPlugin::on_key(int key, int /*scancode*/, int action, int /*mods*/)
 {
   // GLFW constants (stable values matching glfw3.h; avoids adding glfw as a plugin dependency).