diff --git a/mujoco_ros2_control/CMakeLists.txt b/mujoco_ros2_control/CMakeLists.txt index b192059..7e1b241 100644 --- a/mujoco_ros2_control/CMakeLists.txt +++ b/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}) diff --git a/mujoco_ros2_control/patches/mujoco_simulate_managed_user_scn.patch b/mujoco_ros2_control/patches/mujoco_simulate_managed_user_scn.patch new file mode 100644 index 0000000..0830b44 --- /dev/null +++ b/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) diff --git a/mujoco_ros2_control/src/mujoco_system_interface.cpp b/mujoco_ros2_control/src/mujoco_system_interface.cpp index d7666b7..5145cd3 100644 --- a/mujoco_ros2_control/src/mujoco_system_interface.cpp +++ b/mujoco_ros2_control/src/mujoco_system_interface.cpp @@ -78,6 +78,59 @@ using Seconds = std::chrono::duration; 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> + plugin_visualization_scene_states; + +PluginVisualizationSceneState& +get_plugin_visualization_scene_state(const mujoco_ros2_control::MujocoSystemInterface* system_interface) +{ + std::lock_guard lock(plugin_visualization_scene_states_mutex); + auto& state = plugin_visualization_scene_states[system_interface]; + if (!state) + { + state = std::make_unique(); + } + return *state; +} + +PluginVisualizationSceneState* +find_plugin_visualization_scene_state(const mujoco_ros2_control::MujocoSystemInterface* system_interface) +{ + std::lock_guard 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 lock(plugin_visualization_scene_states_mutex); + plugin_visualization_scene_states.erase(system_interface); +} + std::optional 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 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 diff --git a/mujoco_ros2_control_plugins/CMakeLists.txt b/mujoco_ros2_control_plugins/CMakeLists.txt index 6210d0a..84cf69c 100644 --- a/mujoco_ros2_control_plugins/CMakeLists.txt +++ b/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() diff --git a/mujoco_ros2_control_plugins/include/mujoco_ros2_control_plugins/mujoco_ros2_control_plugins_base.hpp b/mujoco_ros2_control_plugins/include/mujoco_ros2_control_plugins/mujoco_ros2_control_plugins_base.hpp index a6ad2bc..60eb862 100644 --- a/mujoco_ros2_control_plugins/include/mujoco_ros2_control_plugins/mujoco_ros2_control_plugins_base.hpp +++ b/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). diff --git a/mujoco_ros2_control_plugins/include/mujoco_ros2_control_plugins/virtual_gantry_plugin.hpp b/mujoco_ros2_control_plugins/include/mujoco_ros2_control_plugins/virtual_gantry_plugin.hpp index b2c44a9..14c89c2 100644 --- a/mujoco_ros2_control_plugins/include/mujoco_ros2_control_plugins/virtual_gantry_plugin.hpp +++ b/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). diff --git a/mujoco_ros2_control_plugins/src/virtual_gantry_plugin.cpp b/mujoco_ros2_control_plugins/src/virtual_gantry_plugin.cpp index 931caac..aa09091 100644 --- a/mujoco_ros2_control_plugins/src/virtual_gantry_plugin.cpp +++ b/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 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). diff --git a/mujoco_ros2_control_plugins/test/test_plugin_visualization.cpp b/mujoco_ros2_control_plugins/test/test_plugin_visualization.cpp new file mode 100644 index 0000000..6b746f2 --- /dev/null +++ b/mujoco_ros2_control_plugins/test/test_plugin_visualization.cpp @@ -0,0 +1,159 @@ +// Copyright 2026 NVIDIA CORPORATION & AFFILIATES +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#include + +#include + +#include + +#include "mujoco_ros2_control_plugins/mujoco_ros2_control_plugins_base.hpp" +#include "mujoco_ros2_control_plugins/virtual_gantry_plugin.hpp" + +namespace mujoco_ros2_control_plugins +{ +class VirtualGantryPluginTestAccessor +{ +public: + static void configure_captured(VirtualGantryPlugin& plugin) + { + plugin.body_id_ = 0; + plugin.body_offset_ = { { 0.0, 0.0, 0.3 } }; + plugin.anchor_pos_ = { { 1.0, 2.0, 1.5 } }; + plugin.enabled_ = true; + plugin.spawn_pos_captured_ = true; + } + + static void set_enabled(VirtualGantryPlugin& plugin, bool enabled) + { + plugin.enabled_ = enabled; + } + + static void set_spawn_pos_captured(VirtualGantryPlugin& plugin, bool captured) + { + plugin.spawn_pos_captured_ = captured; + } +}; +} // namespace mujoco_ros2_control_plugins + +namespace +{ + +class MinimalPlugin : public mujoco_ros2_control_plugins::MuJoCoROS2ControlPluginBase +{ +public: + bool init(rclcpp::Node::SharedPtr /*node*/, const mjModel* /*model*/, mjData* /*data*/) override + { + return true; + } + + void update(const mjModel* /*model*/, mjData* /*data*/) override + { + } + + void cleanup() override + { + } +}; + +TEST(PluginVisualizationTest, BaseHookDoesNotAddVisualizationGeoms) +{ + MinimalPlugin plugin; + mjvScene scene; + mjv_defaultScene(&scene); + scene.ngeom = 7; + + plugin.update_visualization(nullptr, nullptr, &scene); + + EXPECT_EQ(scene.ngeom, 7); +} + +struct GantryVisualizationFixture +{ + mujoco_ros2_control_plugins::VirtualGantryPlugin plugin; + mjModel model{}; + mjData data{}; + std::array xpos{ { 1.0, 2.0, 0.8 } }; + std::array xmat{ { 1.0, 0.0, 0.0, 0.0, 1.0, 0.0, 0.0, 0.0, 1.0 } }; + std::array geoms{}; + std::array geomorder{}; + mjvScene scene{}; + + GantryVisualizationFixture() + { + data.xpos = xpos.data(); + data.xmat = xmat.data(); + + mjv_defaultScene(&scene); + scene.maxgeom = static_cast(geoms.size()); + scene.geoms = geoms.data(); + scene.geomorder = geomorder.data(); + + mujoco_ros2_control_plugins::VirtualGantryPluginTestAccessor::configure_captured(plugin); + } +}; + +TEST(PluginVisualizationTest, GantryDoesNotDrawWhenDisabled) +{ + GantryVisualizationFixture fixture; + mujoco_ros2_control_plugins::VirtualGantryPluginTestAccessor::set_enabled(fixture.plugin, false); + + fixture.plugin.update_visualization(&fixture.model, &fixture.data, &fixture.scene); + + EXPECT_EQ(fixture.scene.ngeom, 0); +} + +TEST(PluginVisualizationTest, GantryDoesNotDrawBeforeAnchorCapture) +{ + GantryVisualizationFixture fixture; + mujoco_ros2_control_plugins::VirtualGantryPluginTestAccessor::set_spawn_pos_captured(fixture.plugin, false); + + fixture.plugin.update_visualization(&fixture.model, &fixture.data, &fixture.scene); + + EXPECT_EQ(fixture.scene.ngeom, 0); +} + +TEST(PluginVisualizationTest, GantryDrawsRopeAndAnchorWhenEnabledAndCaptured) +{ + GantryVisualizationFixture fixture; + + fixture.plugin.update_visualization(&fixture.model, &fixture.data, &fixture.scene); + + ASSERT_EQ(fixture.scene.ngeom, 2); + EXPECT_EQ(fixture.scene.geoms[0].type, mjGEOM_CAPSULE); + EXPECT_EQ(fixture.scene.geoms[1].type, mjGEOM_SPHERE); + EXPECT_NEAR(fixture.scene.geoms[0].pos[0], 1.0, 1e-6); + EXPECT_NEAR(fixture.scene.geoms[0].pos[1], 2.0, 1e-6); + EXPECT_NEAR(fixture.scene.geoms[0].size[0], 0.0125, 1e-6); + EXPECT_NEAR(fixture.scene.geoms[1].pos[0], 1.0, 1e-6); + EXPECT_NEAR(fixture.scene.geoms[1].pos[1], 2.0, 1e-6); + EXPECT_NEAR(fixture.scene.geoms[1].size[0], 0.04, 1e-6); + EXPECT_NEAR(fixture.scene.geoms[1].rgba[0], 1.0, 1e-6); + EXPECT_NEAR(fixture.scene.geoms[1].rgba[1], 0.0, 1e-6); + EXPECT_NEAR(fixture.scene.geoms[1].rgba[2], 0.0, 1e-6); + EXPECT_NEAR(fixture.scene.geoms[1].rgba[3], 1.0, 1e-6); +} + +TEST(PluginVisualizationTest, GantryStillDrawsRopeWhenOnlyOneVisualizationGeomIsAvailable) +{ + GantryVisualizationFixture fixture; + fixture.scene.maxgeom = 1; + + fixture.plugin.update_visualization(&fixture.model, &fixture.data, &fixture.scene); + + ASSERT_EQ(fixture.scene.ngeom, 1); + EXPECT_EQ(fixture.scene.geoms[0].type, mjGEOM_CAPSULE); +} + +} // namespace