From 37d7591ce0d1424f6853dd9293eec1579c8e623f Mon Sep 17 00:00:00 2001 From: Google DeepMind Date: Fri, 14 Feb 2025 10:44:50 -0800 Subject: [PATCH] Add support for user-defined figures in the MuJoCo viewer. PiperOrigin-RevId: 726983516 Change-Id: Ifb821f3cafb163d4f3077d76b0668037b4bc5a9c --- doc/changelog.rst | 2 ++ python/mujoco/simulate.cc | 62 +++++++++++++++++++++++++++------------ python/mujoco/viewer.py | 17 +++++++++++ simulate/simulate.cc | 9 ++++++ simulate/simulate.h | 1 + 5 files changed, 72 insertions(+), 19 deletions(-) diff --git a/doc/changelog.rst b/doc/changelog.rst index 9b88b28a..8ec65806 100644 --- a/doc/changelog.rst +++ b/doc/changelog.rst @@ -25,6 +25,8 @@ Feature promotion General ^^^^^^^ +- Add support for custom plots in the MuJoCo viewer by exposing a ``viewport`` property, a ``set_figures`` method, + and a ``clear_figures`` method. - Separate collision and deformation meshes for :ref:`flex`. This enables a fixed cost for the soft body computations, while preserving the fidelity of high-resolution collisions. - Added :ref:`mjs_setDeepCopy` API function. When the deep copy flag is 0, attaching a model will not copy it to the diff --git a/python/mujoco/simulate.cc b/python/mujoco/simulate.cc index 0ddd5982..56fe4437 100644 --- a/python/mujoco/simulate.cc +++ b/python/mujoco/simulate.cc @@ -28,6 +28,7 @@ #include #include #include +#include namespace mujoco::python { namespace { @@ -69,11 +70,10 @@ class UIAdapterWithPyCallback : public Adapter { class SimulateWrapper { public: SimulateWrapper(std::unique_ptr platform_ui_adapter, - py::object cam, py::object opt, - py::object pert, py::object user_scn, bool is_passive) + py::object cam, py::object opt, py::object pert, + py::object user_scn, bool is_passive) : simulate_(new mujoco::Simulate( - std::move(platform_ui_adapter), - cam.cast().get(), + std::move(platform_ui_adapter), cam.cast().get(), opt.cast().get(), pert.cast().get(), is_passive)), m_(py::none()), @@ -126,6 +126,28 @@ class SimulateWrapper { py::object GetModel() const { return m_; } py::object GetData() const { return d_; } + mjrRect GetViewport() const { + // Return the viewport corresponding to the 3D view, i.e. the viewer window + // without the UI elements. + return simulate_->uistate.rect[3]; + } + + void SetFigures( + const std::vector>& viewports_figures) { + // Pairs of [viewport, figure], where viewport corresponds to the location + // of the figure on the viewer window. + std::vector> user_figures; + for (const auto& [viewport, figure] : viewports_figures) { + mjvFigure casted_figure = *figure.cast().get(); + user_figures.push_back(std::make_pair(viewport, casted_figure)); + } + + // Set them all at once to prevent figure flickering. + simulate_->user_figures_ = user_figures; + } + + void ClearFigures() { simulate_->user_figures_.clear(); } + private: mujoco::Simulate* simulate_; std::atomic_int destroyed_ = 0; @@ -173,14 +195,14 @@ inline auto CallIfNotNull(void (mujoco::Simulate::*func)(Args...)) { } template -inline auto GetIfNotNull(T mujoco::Simulate::*member) { +inline auto GetIfNotNull(T mujoco::Simulate::* member) { return [member](SimulateWrapper& wrapper) -> T& { return SimulateRefOrThrow(wrapper).*member; }; } template -inline auto SetIfNotNull(T mujoco::Simulate::*member) { +inline auto SetIfNotNull(T mujoco::Simulate::* member) { return [member](SimulateWrapper& wrapper, const T& value) -> void { SimulateRefOrThrow(wrapper).*member = value; }; @@ -205,8 +227,7 @@ PYBIND11_MODULE(_simulate, pymodule) { py::object key_callback) { bool is_passive = !run_physics_thread; return std::make_unique( - std::make_unique>( - key_callback), + std::make_unique>(key_callback), scn, cam, opt, pert, is_passive); })) .def("destroy", &SimulateWrapper::Destroy) @@ -225,8 +246,12 @@ PYBIND11_MODULE(_simulate, pymodule) { .def("lock", GetIfNotNull(&mujoco::Simulate::mtx), py::call_guard(), py::return_value_policy::reference_internal) + .def("set_figures", &SimulateWrapper::SetFigures, + py::arg("viewports_figures")) + .def("clear_figures", &SimulateWrapper::ClearFigures) .def_property_readonly("m", &SimulateWrapper::GetModel) .def_property_readonly("d", &SimulateWrapper::GetData) + .def_property_readonly("viewport", &SimulateWrapper::GetViewport) .def_property_readonly("ctrl_noise_std", GetIfNotNull(&mujoco::Simulate::ctrl_noise_std), py::call_guard()) @@ -260,18 +285,17 @@ PYBIND11_MODULE(_simulate, pymodule) { return sim.exitrequest.load(); }), py::call_guard()) - .def( - "exit", - [](SimulateWrapper& wrapper) { - mujoco::Simulate* sim = wrapper.simulate(); - if (!sim) { - return; - } + .def("exit", + [](SimulateWrapper& wrapper) { + mujoco::Simulate* sim = wrapper.simulate(); + if (!sim) { + return; + } - int value = 0; - sim->exitrequest.compare_exchange_strong(value, 1); - wrapper.WaitUntilExit(); - }) + int value = 0; + sim->exitrequest.compare_exchange_strong(value, 1); + wrapper.WaitUntilExit(); + }) .def_property_readonly("uiloadrequest", CallIfNotNull(+[](mujoco::Simulate& sim) { diff --git a/python/mujoco/viewer.py b/python/mujoco/viewer.py index 31ed22a8..65852c87 100644 --- a/python/mujoco/viewer.py +++ b/python/mujoco/viewer.py @@ -108,6 +108,23 @@ class Handle: return sim.d return None + @property + def viewport(self): + sim = self._sim() + if sim is not None: + return sim.viewport + return None + + def set_figures(self, viewports_figures): + sim = self._sim() + if sim is not None: + sim.set_figures(viewports_figures) + + def clear_figures(self): + sim = self._sim() + if sim is not None: + sim.clear_figures() + def close(self): sim = self._sim() if sim is not None: diff --git a/simulate/simulate.cc b/simulate/simulate.cc index 8fc78476..bf4b4a13 100644 --- a/simulate/simulate.cc +++ b/simulate/simulate.cc @@ -542,6 +542,10 @@ void ShowSensor(mj::Simulate* sim, mjrRect rect) { mjr_figure(viewport, &sim->figsensor, &sim->platform_ui->mjr_context()); } +void ShowFigure(mj::Simulate* sim, mjrRect viewport, mjvFigure* fig){ + mjr_figure(viewport, fig, &sim->platform_ui->mjr_context()); +} + // load state from history buffer static void LoadScrubState(mj::Simulate* sim) { // get index into circular buffer @@ -2588,6 +2592,11 @@ void Simulate::Render() { } } + // user figures + for (auto& [viewport, figure] : this->user_figures_) { + ShowFigure(this, viewport, &figure); + } + // finalize this->platform_ui->SwapBuffers(); } diff --git a/simulate/simulate.h b/simulate/simulate.h index bffacb37..cd654192 100644 --- a/simulate/simulate.h +++ b/simulate/simulate.h @@ -252,6 +252,7 @@ class Simulate { // additional user-defined visualization geoms (used in passive mode) mjvScene* user_scn = nullptr; mjtByte user_scn_flags_prev_[mjNRNDFLAG]; + std::vector> user_figures_; // OpenGL rendering and UI int refresh_rate = 60;