Add handles for overlaying user text and images to the Passive Viewer #2503

Andrew-Luo1 · 2025-03-17T02:37:36Z

See issue #2494

Added the text and image overlay handles in the style of commit 37d759.

This supports:
Brax PR 584
Playground PR 83

aftersomemath

@yuvaltassa asked me to take a look.

This is a great feature to have! My comments are mostly about minor things.

However, there is a race condition introduced with the previous commit for set_figures 37d7591 that is being duplicated in this PR. In my testing this is particular easy to trigger with set_images and it results in a crash/segfault.

That commit and this PR introduce the vectors user_text_, user_images_, and user_figures_ to the Simulate class. These vectors are updated on the Python side and the data in them is used to render the viewers UI. To avoid undefined behavior, there needs to be some synchronization to ensure these vectors are not updated while the viewer is rendering.

In addition to the crashing I mentioned above, I believe this race condition is why there are several comments in the Python side of the code about doing things a certain way to avoid flickering.

To resolve the situation, set_figures, overlay_text and set_images, should be refactored to work like update_hfield, update_mesh, and update_mesh. The latter methods sync the data with the viewer's render thread using a condition variable as can be seen here and here.

Note: It is in theory not enough to simply acquire the Simulate lock with viewer.lock() prior to calling set_figures, overlay_text, or set_images (though in practice I did find it helped with the segfaulting when using set_images). This is because Simulate acquires its lock only to stage the data for rendering. Subsequently, it releases the lock and renders. This was done on purpose helps ensure the viewer runs smoothly, and we should probably keep it that way, thus my suggestion to stage the data like update_hfield, update_mesh, and update_mesh.

aftersomemath · 2025-03-20T18:04:41Z

python/mujoco/viewer.py

@@ -125,6 +128,49 @@ def clear_figures(self):
    if sim is not None:
      sim.clear_figures()

+  def overlay_text(self, overlay_texts: List[Tuple[int, int, str, str]]):


It may be best to name this method set_text to be consistent with set_figures and set_images.

aftersomemath · 2025-03-20T18:05:04Z

python/mujoco/viewer.py

+    if sim is not None:
+      sim.overlay_text(overlay_texts)
+
+  def clear_overlay_text(self):


Similarly, this would be clear_text.

aftersomemath · 2025-03-20T18:05:28Z

python/mujoco/viewer.py

+    if sim is not None:
+      sim.clear_overlay_text()
+
+  def set_images(


This method should get a docstring as overlay_text does.

aftersomemath · 2025-03-20T18:07:06Z

python/mujoco/viewer.py

+    """Overlay text on the viewer.
+
+    Args:
+      overlay_texts: List of tuples of (font, gridpos, text1, text2)


Would it be easy to accept a single tuple or a list of tuples? For example, we do the equivalent of this in rollout.

aftersomemath · 2025-03-20T18:07:29Z

python/mujoco/viewer.py

+      sim.clear_overlay_text()
+
+  def set_images(
+      self, viewports_images: List[Tuple[mujoco.MjrRect, np.ndarray]]


Similarly, can this accept a single instance in addition to a list?

aftersomemath · 2025-03-20T18:55:39Z

python/mujoco/simulate.cc

+      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));


user_images_ is being modified while it is potentially used by Simulate::Render. See top level review comment about how we can resolve this issue in general.

aftersomemath · 2025-03-20T18:55:47Z

python/mujoco/simulate.cc

+    }
+
+    // Set them all at once to prevent overlay text flickering.
+    simulate_->user_text_ = user_overlay_text;


user_text_ is being modified while it is potentially used by Simulate::Render. See top level review comment about how we can resolve this issue in general. The fix should also fix the "flickering" mentioned in the comment above.

aftersomemath · 2025-03-20T18:57:27Z

python/mujoco/simulate.cc

+      size_t size = height * width * 3;
+
+      // Make a copy of the image data to prevent flickering
+      unsigned char* image_copy = new unsigned char[size];


See top level review comment about how we should handle locking. It should prevent the flickering and the need to make this copy.

aftersomemath · 2025-03-20T18:58:29Z

python/mujoco/simulate.cc

+        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)");


Can we check the type of the image data here? It is assumed to be a uint8 below.

aftersomemath · 2025-03-20T19:05:58Z

simulate/simulate.cc

@@ -2597,6 +2605,16 @@ void Simulate::Render() {
    ShowFigure(this, viewport, &figure);


These calls to ShowFigure, ShowOverlayText and ShowImage are being made from Simulate's Render function. Render is called from the RenderLoop after unlocking Simulate's lock (this helps with parallelism and makes the viewer run smoother).

However, ShowFigure, ShowOverlayText and ShowImage use vectors, user_figures_, user_text_, and user_images_ which can be modified concurrently by the user code. The result is that undefined behavior can occur.

See top level review comment about how we can resolve this issue in general.

aftersomemath · 2025-03-20T19:45:56Z

python/mujoco/viewer.py

+    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]][


Does resizing by slicing have good performance? It might be a good idea to check. If not, I would just force the user to pass in images of the correct size.

TODO: will check if resize > 1m fps, else check sizes match

Turns out it's only about 7k FPS. Will remove

josechenf · 2025-03-29T08:23:24Z

We'd be super keen to get this in and use this!
@aftersomemath Let us know if this works for you!

aftersomemath

Looks good to me. My comments are just nits.

Thanks for adding a docstring, typing, etc to set_figures.

I think merging this with the race condition is fine because set_figures is already on main with the race condition. I will submit a PR for fixing the race condition in a few days.

After reviewing this a second time, I came to the conclusion that signatures of the set_X functions should be more pythonic (maybe each argument should be its own list like with rollout). But we don't need to fix that in this PR.

aftersomemath · 2025-03-30T12:34:56Z

python/mujoco/viewer.py

      sim.set_figures(viewports_figures)

  def clear_figures(self):
    sim = self._sim()
    if sim is not None:
      sim.clear_figures()

+  def set_text(self, overlay_texts: Union[Tuple[Optional[int], Optional[int], Optional[str], Optional[str]], 


nit: if the other functions are plural (set_images, set_figures) then this should be set_texts.

aftersomemath · 2025-03-30T12:45:59Z

python/mujoco/viewer.py

+      for viewport, image in viewports_images:
+        targ_shape = (viewport.height, viewport.width)
+        # Check if image is already the correct shape
+        if image.shape[:2] != targ_shape:


Since we are no longer resizing the image, passing a viewport with the image is no longer required, the size is implicit and this check is not needed. However, I'm ok with leaving the check here to avoid breaking the API in the future if resizing is brought back somehow.

aftersomemath · 2025-03-30T12:48:10Z

python/mujoco/viewer.py

+    """Overlay text on the viewer.
+
+    Args:
+      overlay_texts: Single tuple or list of tuples of (font, gridpos, text1, text2)


nit: since set_figures argument is viewports_figures, this argument should be fonts_gridpos_texts, but that is cumbersome, how about texts? We don't need to include the overlay prefix.

Andrew-Luo1 · 2025-03-31T02:20:35Z

Updated!

Andrew-Luo1 added 2 commits March 14, 2025 21:41

add handles for text and image overlays

f02fdfe

run pyink and isort

86fd31b

This was referenced Mar 17, 2025

Rscope: MJX + Brax Reinforcement Learning Debugger google/brax#584

Closed

Rscope: MJX + Brax Reinforcement Learning Debugger google-deepmind/mujoco_playground#83

Open

fix memory leak

faff122

aftersomemath reviewed Mar 20, 2025

View reviewed changes

resolve PR comments

97092fa

aftersomemath approved these changes Mar 30, 2025

View reviewed changes

improve variable and function names

65500cf

yuvaltassa approved these changes Apr 8, 2025

View reviewed changes

revert import order

fc516ee

yuvaltassa approved these changes Apr 8, 2025

View reviewed changes

copybara-service bot merged commit 699ad13 into google-deepmind:main Apr 8, 2025
11 of 14 checks passed

Andrew-Luo1 mentioned this pull request Apr 9, 2025

Add live visualizer support similar to Isaac Lab/Gym google-deepmind/mujoco_playground#59

Open

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Add handles for overlaying user text and images to the Passive Viewer #2503

Add handles for overlaying user text and images to the Passive Viewer #2503

Andrew-Luo1 commented Mar 17, 2025 •

edited

Loading

aftersomemath left a comment

aftersomemath Mar 20, 2025

aftersomemath Mar 20, 2025

aftersomemath Mar 20, 2025

aftersomemath Mar 20, 2025

aftersomemath Mar 20, 2025

aftersomemath Mar 20, 2025

aftersomemath Mar 20, 2025

aftersomemath Mar 20, 2025

aftersomemath Mar 20, 2025

aftersomemath Mar 20, 2025

aftersomemath Mar 20, 2025

Andrew-Luo1 Mar 27, 2025

Andrew-Luo1 Mar 27, 2025

josechenf commented Mar 29, 2025

aftersomemath left a comment

aftersomemath Mar 30, 2025

aftersomemath Mar 30, 2025

aftersomemath Mar 30, 2025

Andrew-Luo1 commented Mar 31, 2025

		@@ -2597,6 +2605,16 @@ void Simulate::Render() {
		ShowFigure(this, viewport, &figure);

Add handles for overlaying user text and images to the Passive Viewer #2503

Add handles for overlaying user text and images to the Passive Viewer #2503

Conversation

Andrew-Luo1 commented Mar 17, 2025 • edited Loading

aftersomemath left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

josechenf commented Mar 29, 2025

aftersomemath left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Andrew-Luo1 commented Mar 31, 2025

Andrew-Luo1 commented Mar 17, 2025 •

edited

Loading