From f02fdfe2d1d3149579b29650d75117190c189cdb Mon Sep 17 00:00:00 2001 From: andrew Date: Fri, 14 Mar 2025 21:41:28 -0400 Subject: [PATCH 1/6] add handles for text and image overlays --- python/mujoco/simulate.cc | 50 +++++++++++++++++++++++++++++++++++++++ python/mujoco/viewer.py | 44 +++++++++++++++++++++++++++++++--- simulate/simulate.cc | 18 ++++++++++++++ simulate/simulate.h | 2 ++ 4 files changed, 111 insertions(+), 3 deletions(-) diff --git a/python/mujoco/simulate.cc b/python/mujoco/simulate.cc index 56fe4437..8fdb668c 100644 --- a/python/mujoco/simulate.cc +++ b/python/mujoco/simulate.cc @@ -148,6 +148,50 @@ class SimulateWrapper { void ClearFigures() { simulate_->user_figures_.clear(); } + void SetOverlayText( + const std::vector>& overlay_texts) { + // Collection of [font, gridpos, text1, text2] tuples for overlay text + std::vector> user_overlay_text; + for (const auto& [font, gridpos, text1, text2] : overlay_texts) { + user_overlay_text.push_back(std::make_tuple(font, gridpos, text1, text2)); + } + + // Set them all at once to prevent overlay text flickering. + simulate_->user_text_ = user_overlay_text; + } + + void ClearOverlayText() { simulate_->user_text_.clear(); } + + void SetImages( + const std::vector>>& viewport_images + ) { + // Clear previous images to prevent memory leaks + simulate_->user_images_.clear(); + + for (const auto& [viewport, image] : viewport_images) { + auto buf = image.request(); + if (static_cast(buf.shape[2]) != 3) { + throw std::invalid_argument("image must have 3 channels"); + } + if (buf.ndim != 3) { + throw std::invalid_argument("image must have 3 dimensions (H, W, C)"); + } + + // Calculate size of the image data + size_t height = buf.shape[0]; + size_t width = buf.shape[1]; + size_t size = height * width * 3; + + // Make a copy of the image data to prevent flickering + unsigned char* image_copy = new unsigned char[size]; + std::memcpy(image_copy, buf.ptr, size); + + simulate_->user_images_.push_back(std::make_tuple(viewport, image_copy)); + } + } + + void ClearImages() { simulate_->user_images_.clear(); } + private: mujoco::Simulate* simulate_; std::atomic_int destroyed_ = 0; @@ -249,6 +293,12 @@ PYBIND11_MODULE(_simulate, pymodule) { .def("set_figures", &SimulateWrapper::SetFigures, py::arg("viewports_figures")) .def("clear_figures", &SimulateWrapper::ClearFigures) + .def("overlay_text", &SimulateWrapper::SetOverlayText, + py::arg("overlay_texts")) + .def("clear_overlay_text", &SimulateWrapper::ClearOverlayText) + .def("set_images", &SimulateWrapper::SetImages, + py::arg("viewports_images")) + .def("clear_images", &SimulateWrapper::ClearImages) .def_property_readonly("m", &SimulateWrapper::GetModel) .def_property_readonly("d", &SimulateWrapper::GetData) .def_property_readonly("viewport", &SimulateWrapper::GetViewport) diff --git a/python/mujoco/viewer.py b/python/mujoco/viewer.py index 65852c87..b37355cc 100644 --- a/python/mujoco/viewer.py +++ b/python/mujoco/viewer.py @@ -23,9 +23,8 @@ import queue import sys import threading import time -from typing import Callable, Optional, Tuple, Union +from typing import Callable, List, Optional, Tuple, Union import weakref - import glfw import mujoco from mujoco import _simulate @@ -115,7 +114,7 @@ class Handle: return sim.viewport return None - def set_figures(self, viewports_figures): + def set_figures(self, viewports_figures: List[Tuple[mujoco.MjrRect, mujoco.MjvFigure]]): sim = self._sim() if sim is not None: sim.set_figures(viewports_figures) @@ -125,6 +124,45 @@ class Handle: if sim is not None: sim.clear_figures() + def overlay_text(self, overlay_texts: List[Tuple[int, int, str, str]]): + """ Overlay text on the viewer. + + Args: + overlay_texts: List of tuples of (font, gridpos, text1, text2) + let: + font: Font style from mujoco.mjtFontScale + gridpos: Position of text box from mujoco.mjtGridPos + text1: Left text column + text2: Right text column + """ + sim = self._sim() + if sim is not None: + sim.overlay_text(overlay_texts) + + def clear_overlay_text(self): + sim = self._sim() + if sim is not None: + sim.clear_overlay_text() + + def set_images(self, viewports_images: List[Tuple[mujoco.MjrRect, np.ndarray]]): + sim = self._sim() + if sim is not None: + # Nearest neighbor resize + resize = lambda a, s: a[(np.arange(s[0]) * a.shape[0]) // s[0]][:, (np.arange(s[1]) * a.shape[1]) // s[1]] + resized_viewports_images = [] + for viewport, image in viewports_images: + targ_shape = (viewport.height, viewport.width) + resized = resize(image, targ_shape) + resized = np.flip(resized, axis=0) + resized = np.ascontiguousarray(resized) + resized_viewports_images.append((viewport, resized)) + sim.set_images(resized_viewports_images) + + def clear_images(self): + sim = self._sim() + if sim is not None: + sim.clear_images() + def close(self): sim = self._sim() if sim is not None: diff --git a/simulate/simulate.cc b/simulate/simulate.cc index bf4b4a13..57e3eaaa 100644 --- a/simulate/simulate.cc +++ b/simulate/simulate.cc @@ -546,6 +546,14 @@ void ShowFigure(mj::Simulate* sim, mjrRect viewport, mjvFigure* fig){ mjr_figure(viewport, fig, &sim->platform_ui->mjr_context()); } +void ShowOverlayText(mj::Simulate* sim, mjrRect viewport, int font, int gridpos, std::string text1, std::string text2){ + mjr_overlay(font, gridpos, viewport, text1.c_str(), text2.c_str(), &sim->platform_ui->mjr_context()); +} + +void ShowImage(mj::Simulate* sim, mjrRect viewport, const unsigned char* image) { + mjr_drawPixels(image, nullptr, viewport, &sim->platform_ui->mjr_context()); +} + // load state from history buffer static void LoadScrubState(mj::Simulate* sim) { // get index into circular buffer @@ -2597,6 +2605,16 @@ void Simulate::Render() { ShowFigure(this, viewport, &figure); } + // overlay text + for (auto& [font, gridpos, text1, text2] : this->user_text_) { + ShowOverlayText(this, rect, font, gridpos, text1, text2); + } + + // user images + for (auto& [viewport, image] : this->user_images_) { + ShowImage(this, viewport, image); + } + // finalize this->platform_ui->SwapBuffers(); } diff --git a/simulate/simulate.h b/simulate/simulate.h index cd654192..0bf6ad25 100644 --- a/simulate/simulate.h +++ b/simulate/simulate.h @@ -253,6 +253,8 @@ class Simulate { mjvScene* user_scn = nullptr; mjtByte user_scn_flags_prev_[mjNRNDFLAG]; std::vector> user_figures_; + std::vector> user_text_; + std::vector> user_images_; // OpenGL rendering and UI int refresh_rate = 60; From 86fd31b9b86f1afe1dbf23ef37378fbbb8aca5bb Mon Sep 17 00:00:00 2001 From: andrew Date: Sun, 16 Mar 2025 17:44:24 -0400 Subject: [PATCH 2/6] run pyink and isort --- python/mujoco/viewer.py | 22 +++++++++++++++------- 1 file changed, 15 insertions(+), 7 deletions(-) diff --git a/python/mujoco/viewer.py b/python/mujoco/viewer.py index b37355cc..21a1eeea 100644 --- a/python/mujoco/viewer.py +++ b/python/mujoco/viewer.py @@ -25,10 +25,12 @@ import threading import time from typing import Callable, List, Optional, Tuple, Union import weakref + import glfw +import numpy as np + import mujoco from mujoco import _simulate -import numpy as np if not glfw._glfw: # pylint: disable=protected-access raise RuntimeError('GLFW dynamic library handle is not available') @@ -114,7 +116,9 @@ class Handle: return sim.viewport return None - def set_figures(self, viewports_figures: List[Tuple[mujoco.MjrRect, mujoco.MjvFigure]]): + def set_figures( + self, viewports_figures: List[Tuple[mujoco.MjrRect, mujoco.MjvFigure]] + ): sim = self._sim() if sim is not None: sim.set_figures(viewports_figures) @@ -125,8 +129,8 @@ class Handle: sim.clear_figures() def overlay_text(self, overlay_texts: List[Tuple[int, int, str, str]]): - """ Overlay text on the viewer. - + """Overlay text on the viewer. + Args: overlay_texts: List of tuples of (font, gridpos, text1, text2) let: @@ -138,17 +142,21 @@ class Handle: sim = self._sim() if sim is not None: sim.overlay_text(overlay_texts) - + def clear_overlay_text(self): sim = self._sim() if sim is not None: sim.clear_overlay_text() - def set_images(self, viewports_images: List[Tuple[mujoco.MjrRect, np.ndarray]]): + def set_images( + self, viewports_images: List[Tuple[mujoco.MjrRect, np.ndarray]] + ): sim = self._sim() if sim is not None: # Nearest neighbor resize - resize = lambda a, s: a[(np.arange(s[0]) * a.shape[0]) // s[0]][:, (np.arange(s[1]) * a.shape[1]) // s[1]] + resize = lambda a, s: a[(np.arange(s[0]) * a.shape[0]) // s[0]][ + :, (np.arange(s[1]) * a.shape[1]) // s[1] + ] resized_viewports_images = [] for viewport, image in viewports_images: targ_shape = (viewport.height, viewport.width) From faff1224ba962e1d14570eb790678fc152257faa Mon Sep 17 00:00:00 2001 From: andrew Date: Mon, 17 Mar 2025 10:42:11 -0400 Subject: [PATCH 3/6] fix memory leak --- python/mujoco/simulate.cc | 11 +++++++++-- 1 file changed, 9 insertions(+), 2 deletions(-) diff --git a/python/mujoco/simulate.cc b/python/mujoco/simulate.cc index 8fdb668c..868bec28 100644 --- a/python/mujoco/simulate.cc +++ b/python/mujoco/simulate.cc @@ -91,6 +91,7 @@ class SimulateWrapper { void Destroy() { if (simulate_) { + ClearImages(); delete simulate_; simulate_ = nullptr; destroyed_.store(1); @@ -166,7 +167,7 @@ class SimulateWrapper { const std::vector>>& viewport_images ) { // Clear previous images to prevent memory leaks - simulate_->user_images_.clear(); + ClearImages(); for (const auto& [viewport, image] : viewport_images) { auto buf = image.request(); @@ -190,7 +191,13 @@ class SimulateWrapper { } } - void ClearImages() { simulate_->user_images_.clear(); } + void ClearImages() { + // Free memory for each image before clearing the vector + for (const auto& [viewport, image_ptr] : simulate_->user_images_) { + delete[] image_ptr; + } + simulate_->user_images_.clear(); + } private: mujoco::Simulate* simulate_; From 97092facd639e501f4e3821beb8842aa61dc4dda Mon Sep 17 00:00:00 2001 From: andrew Date: Thu, 27 Mar 2025 19:35:36 -0400 Subject: [PATCH 4/6] resolve PR comments --- python/mujoco/simulate.cc | 19 +++++----- python/mujoco/viewer.py | 75 ++++++++++++++++++++++++++++----------- simulate/simulate.h | 2 +- 3 files changed, 67 insertions(+), 29 deletions(-) diff --git a/python/mujoco/simulate.cc b/python/mujoco/simulate.cc index 868bec28..326879cc 100644 --- a/python/mujoco/simulate.cc +++ b/python/mujoco/simulate.cc @@ -149,7 +149,7 @@ class SimulateWrapper { void ClearFigures() { simulate_->user_figures_.clear(); } - void SetOverlayText( + void SetText( const std::vector>& overlay_texts) { // Collection of [font, gridpos, text1, text2] tuples for overlay text std::vector> user_overlay_text; @@ -161,21 +161,24 @@ class SimulateWrapper { simulate_->user_text_ = user_overlay_text; } - void ClearOverlayText() { simulate_->user_text_.clear(); } + void ClearText() { simulate_->user_text_.clear(); } void SetImages( - const std::vector>>& viewport_images + const std::vector> viewports_images ) { // Clear previous images to prevent memory leaks ClearImages(); - for (const auto& [viewport, image] : viewport_images) { + for (const auto& [viewport, image] : viewports_images) { auto buf = image.request(); + if (buf.ndim != 3) { + throw std::invalid_argument("image must have 3 dimensions (H, W, C)"); + } if (static_cast(buf.shape[2]) != 3) { throw std::invalid_argument("image must have 3 channels"); } - if (buf.ndim != 3) { - throw std::invalid_argument("image must have 3 dimensions (H, W, C)"); + if (buf.itemsize != sizeof(unsigned char)) { + throw std::invalid_argument("image must be uint8 format"); } // Calculate size of the image data @@ -300,9 +303,9 @@ PYBIND11_MODULE(_simulate, pymodule) { .def("set_figures", &SimulateWrapper::SetFigures, py::arg("viewports_figures")) .def("clear_figures", &SimulateWrapper::ClearFigures) - .def("overlay_text", &SimulateWrapper::SetOverlayText, + .def("set_text", &SimulateWrapper::SetText, py::arg("overlay_texts")) - .def("clear_overlay_text", &SimulateWrapper::ClearOverlayText) + .def("clear_text", &SimulateWrapper::ClearText) .def("set_images", &SimulateWrapper::SetImages, py::arg("viewports_images")) .def("clear_images", &SimulateWrapper::ClearImages) diff --git a/python/mujoco/viewer.py b/python/mujoco/viewer.py index 21a1eeea..527fda50 100644 --- a/python/mujoco/viewer.py +++ b/python/mujoco/viewer.py @@ -117,10 +117,21 @@ class Handle: return None def set_figures( - self, viewports_figures: List[Tuple[mujoco.MjrRect, mujoco.MjvFigure]] + self, viewports_figures: Union[Tuple[mujoco.MjrRect, mujoco.MjvFigure], + List[Tuple[mujoco.MjrRect, mujoco.MjvFigure]]] ): + """Overlay figures on the viewer. + + Args: + viewports_figures: Single tuple or list of tuples of (viewport, figure) + viewport: Rectangle defining position and size of the figure + figure: MjvFigure object containing the figure data to display + """ sim = self._sim() if sim is not None: + # Convert single tuple to list if needed + if isinstance(viewports_figures, tuple): + viewports_figures = [viewports_figures] sim.set_figures(viewports_figures) def clear_figures(self): @@ -128,43 +139,67 @@ class Handle: if sim is not None: sim.clear_figures() - def overlay_text(self, overlay_texts: List[Tuple[int, int, str, str]]): + def set_text(self, overlay_texts: Union[Tuple[Optional[int], Optional[int], Optional[str], Optional[str]], + List[Tuple[Optional[int], Optional[int], Optional[str], Optional[str]]]]): """Overlay text on the viewer. Args: - overlay_texts: List of tuples of (font, gridpos, text1, text2) - let: + overlay_texts: Single tuple or list of tuples of (font, gridpos, text1, text2) font: Font style from mujoco.mjtFontScale gridpos: Position of text box from mujoco.mjtGridPos - text1: Left text column - text2: Right text column + text1: Left text column, defaults to empty string if None + text2: Right text column, defaults to empty string if None """ sim = self._sim() if sim is not None: - sim.overlay_text(overlay_texts) + # Convert single tuple to list if needed + if isinstance(overlay_texts, tuple): + overlay_texts = [overlay_texts] + + # Convert None values to empty strings + default_font = mujoco.mjtFontScale.mjFONTSCALE_150 + default_gridpos = mujoco.mjtGridPos.mjGRID_TOPLEFT + processed_texts = [( + default_font if font is None else font, + default_gridpos if gridpos is None else gridpos, + "" if text1 is None else text1, + "" if text2 is None else text2) + for font, gridpos, text1, text2 in overlay_texts] + + sim.set_text(processed_texts) - def clear_overlay_text(self): + def clear_text(self): sim = self._sim() if sim is not None: - sim.clear_overlay_text() + sim.clear_text() def set_images( - self, viewports_images: List[Tuple[mujoco.MjrRect, np.ndarray]] + self, viewports_images: Union[Tuple[mujoco.MjrRect, np.ndarray], + List[Tuple[mujoco.MjrRect, np.ndarray]]] ): + """Overlay images on the viewer. + + Args: + viewports_images: Single tuple or list of tuples of (viewport, image) + viewport: Rectangle defining position and size of the image + image: RGB image with shape (height, width, 3) + """ sim = self._sim() if sim is not None: - # Nearest neighbor resize - resize = lambda a, s: a[(np.arange(s[0]) * a.shape[0]) // s[0]][ - :, (np.arange(s[1]) * a.shape[1]) // s[1] - ] - resized_viewports_images = [] + # Convert single tuple to list if needed + if isinstance(viewports_images, tuple): + viewports_images = [viewports_images] + + processed_images = [] for viewport, image in viewports_images: targ_shape = (viewport.height, viewport.width) - resized = resize(image, targ_shape) - resized = np.flip(resized, axis=0) - resized = np.ascontiguousarray(resized) - resized_viewports_images.append((viewport, resized)) - sim.set_images(resized_viewports_images) + # Check if image is already the correct shape + if image.shape[:2] != targ_shape: + raise ValueError(f"Image shape {image.shape[:2]} does not match target shape {targ_shape}") + flipped = np.flip(image, axis=0) + contiguous = np.ascontiguousarray(flipped) + processed_images.append((viewport, contiguous)) + sim.set_images(processed_images) def clear_images(self): sim = self._sim() diff --git a/simulate/simulate.h b/simulate/simulate.h index 0bf6ad25..00dc5fe5 100644 --- a/simulate/simulate.h +++ b/simulate/simulate.h @@ -249,7 +249,7 @@ class Simulate { mjvFigure figsize = {}; mjvFigure figsensor = {}; - // additional user-defined visualization geoms (used in passive mode) + // additional user-defined visualization mjvScene* user_scn = nullptr; mjtByte user_scn_flags_prev_[mjNRNDFLAG]; std::vector> user_figures_; From 65500cfbdc15244d0cb6e80d831fa6f31eea29bd Mon Sep 17 00:00:00 2001 From: andrew Date: Sun, 30 Mar 2025 22:19:55 -0400 Subject: [PATCH 5/6] improve variable and function names --- python/mujoco/simulate.cc | 20 ++++++++++---------- python/mujoco/viewer.py | 16 ++++++++-------- simulate/simulate.cc | 2 +- simulate/simulate.h | 2 +- 4 files changed, 20 insertions(+), 20 deletions(-) diff --git a/python/mujoco/simulate.cc b/python/mujoco/simulate.cc index 326879cc..9c1e7bec 100644 --- a/python/mujoco/simulate.cc +++ b/python/mujoco/simulate.cc @@ -149,19 +149,19 @@ class SimulateWrapper { void ClearFigures() { simulate_->user_figures_.clear(); } - void SetText( - const std::vector>& overlay_texts) { + void SetTexts( + const std::vector>& texts) { // Collection of [font, gridpos, text1, text2] tuples for overlay text - std::vector> user_overlay_text; - for (const auto& [font, gridpos, text1, text2] : overlay_texts) { - user_overlay_text.push_back(std::make_tuple(font, gridpos, text1, text2)); + std::vector> user_texts; + for (const auto& [font, gridpos, text1, text2] : texts) { + user_texts.push_back(std::make_tuple(font, gridpos, text1, text2)); } - // Set them all at once to prevent overlay text flickering. - simulate_->user_text_ = user_overlay_text; + // Set them all at once to prevent text flickering. + simulate_->user_texts_ = user_texts; } - void ClearText() { simulate_->user_text_.clear(); } + void ClearTexts() { simulate_->user_texts_.clear(); } void SetImages( const std::vector> viewports_images @@ -303,9 +303,9 @@ PYBIND11_MODULE(_simulate, pymodule) { .def("set_figures", &SimulateWrapper::SetFigures, py::arg("viewports_figures")) .def("clear_figures", &SimulateWrapper::ClearFigures) - .def("set_text", &SimulateWrapper::SetText, + .def("set_texts", &SimulateWrapper::SetTexts, py::arg("overlay_texts")) - .def("clear_text", &SimulateWrapper::ClearText) + .def("clear_texts", &SimulateWrapper::ClearTexts) .def("set_images", &SimulateWrapper::SetImages, py::arg("viewports_images")) .def("clear_images", &SimulateWrapper::ClearImages) diff --git a/python/mujoco/viewer.py b/python/mujoco/viewer.py index 527fda50..e78e1229 100644 --- a/python/mujoco/viewer.py +++ b/python/mujoco/viewer.py @@ -139,12 +139,12 @@ class Handle: if sim is not None: sim.clear_figures() - def set_text(self, overlay_texts: Union[Tuple[Optional[int], Optional[int], Optional[str], Optional[str]], + def set_texts(self, texts: Union[Tuple[Optional[int], Optional[int], Optional[str], Optional[str]], List[Tuple[Optional[int], Optional[int], Optional[str], Optional[str]]]]): """Overlay text on the viewer. Args: - overlay_texts: Single tuple or list of tuples of (font, gridpos, text1, text2) + texts: Single tuple or list of tuples of (font, gridpos, text1, text2) font: Font style from mujoco.mjtFontScale gridpos: Position of text box from mujoco.mjtGridPos text1: Left text column, defaults to empty string if None @@ -153,8 +153,8 @@ class Handle: sim = self._sim() if sim is not None: # Convert single tuple to list if needed - if isinstance(overlay_texts, tuple): - overlay_texts = [overlay_texts] + if isinstance(texts, tuple): + texts = [texts] # Convert None values to empty strings default_font = mujoco.mjtFontScale.mjFONTSCALE_150 @@ -164,14 +164,14 @@ class Handle: default_gridpos if gridpos is None else gridpos, "" if text1 is None else text1, "" if text2 is None else text2) - for font, gridpos, text1, text2 in overlay_texts] + for font, gridpos, text1, text2 in texts] - sim.set_text(processed_texts) + sim.set_texts(processed_texts) - def clear_text(self): + def clear_texts(self): sim = self._sim() if sim is not None: - sim.clear_text() + sim.clear_texts() def set_images( self, viewports_images: Union[Tuple[mujoco.MjrRect, np.ndarray], diff --git a/simulate/simulate.cc b/simulate/simulate.cc index 57e3eaaa..87023d50 100644 --- a/simulate/simulate.cc +++ b/simulate/simulate.cc @@ -2606,7 +2606,7 @@ void Simulate::Render() { } // overlay text - for (auto& [font, gridpos, text1, text2] : this->user_text_) { + for (auto& [font, gridpos, text1, text2] : this->user_texts_) { ShowOverlayText(this, rect, font, gridpos, text1, text2); } diff --git a/simulate/simulate.h b/simulate/simulate.h index 00dc5fe5..a8c9da58 100644 --- a/simulate/simulate.h +++ b/simulate/simulate.h @@ -253,7 +253,7 @@ class Simulate { mjvScene* user_scn = nullptr; mjtByte user_scn_flags_prev_[mjNRNDFLAG]; std::vector> user_figures_; - std::vector> user_text_; + std::vector> user_texts_; std::vector> user_images_; // OpenGL rendering and UI From fc516ee3c1eacbb43cf82909434d7e07ae6ec6b7 Mon Sep 17 00:00:00 2001 From: andrew Date: Tue, 8 Apr 2025 14:06:17 -0400 Subject: [PATCH 6/6] revert import order --- python/mujoco/viewer.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/python/mujoco/viewer.py b/python/mujoco/viewer.py index e78e1229..234bf784 100644 --- a/python/mujoco/viewer.py +++ b/python/mujoco/viewer.py @@ -27,10 +27,9 @@ from typing import Callable, List, Optional, Tuple, Union import weakref import glfw -import numpy as np - import mujoco from mujoco import _simulate +import numpy as np if not glfw._glfw: # pylint: disable=protected-access raise RuntimeError('GLFW dynamic library handle is not available')