From 34226bf52d07dc0d14d775d665999944d04570b2 Mon Sep 17 00:00:00 2001
From: Saran Tunyasuvunakool <stunya@google.com>
Date: Mon, 10 Apr 2023 14:35:12 -0700
Subject: [PATCH] Fix a race condition in `launch_passive` and `launch_repl`.

Before this change, these functions can return before an internal call to `mj_forward`, which results in a race condition if the user script then tries to modify the physics state.

Fixes #783
Fixes #790

PiperOrigin-RevId: 523214378
Change-Id: I8a4c6cef0f5d3ec781fd2c5cae3c20dd4b8be67f
---
 doc/changelog.rst       | 5 +++++
 python/mujoco/viewer.py | 7 ++++++-
 2 files changed, 11 insertions(+), 1 deletion(-)

diff --git a/doc/changelog.rst b/doc/changelog.rst
index c92fc350..31921d43 100644
--- a/doc/changelog.rst
+++ b/doc/changelog.rst
@@ -26,6 +26,11 @@ Python bindings
 - Offscreen rendering on macOS is no longer restricted to the main thread. This is achieved by using the low-level
   Core OpenGL (CGL) API to create the OpenGL context, rather than going via GLFW which relies on Cocoa's NSOpenGL.
   The resulting context is not tied to a Cocoa window, and is therefore not tied to the main thread.
+- Fixed a race condition in ``viewer.launch_passive`` and  ``viewer.launch_repl``. These functions could previously
+  return before an internal call to ``mj_forward``. This allows user code to continue and potentially modify physics
+  state concurrently with the internal ``mj_forward``, resulting in e.g.
+  `MuJoCo stack overflow error<https://github.com/deepmind/mujoco/issues/783>`_
+  or `segmentation fault<https://github.com/deepmind/mujoco/issues/790>`_.
 
 Bug fixes
 ^^^^^^^^^
diff --git a/python/mujoco/viewer.py b/python/mujoco/viewer.py
index 9c00bd88..8511ae10 100644
--- a/python/mujoco/viewer.py
+++ b/python/mujoco/viewer.py
@@ -101,7 +101,6 @@ def _reload(
 
     path = load_tuple[2] if len(load_tuple) == 3 else ''
     simulate.load(path, m, d)
-    mujoco.mj_forward(m, d)
 
     return m, d
 
@@ -285,6 +284,8 @@ def launch(model: Optional[mujoco.MjModel] = None,
            run_physics_thread: bool = True,
            loader: Optional[LoaderType] = None) -> None:
   """Launches the Simulate GUI."""
+  if not run_physics_thread:
+    mujoco.mj_forward(model, data)
   _launch_internal(
       model, data, run_physics_thread=run_physics_thread, loader=loader)
 
@@ -300,6 +301,9 @@ def launch_passive(model: mujoco.MjModel, data: mujoco.MjData) -> None:
     raise ValueError(f'`model` is not a mujoco.MjModel: got {model!r}')
   if not isinstance(data, mujoco.MjData):
     raise ValueError(f'`data` is not a mujoco.MjData: got {data!r}')
+
+  mujoco.mj_forward(model, data)
+
   if sys.platform != 'darwin':
     thread = threading.Thread(
         target=_launch_internal,
@@ -379,6 +383,7 @@ def launch_repl(model: mujoco.MjModel, data: mujoco.MjData) -> None:
     repl_thread.start()
 
     # Launch the viewer on the main thread.
+    mujoco.mj_forward(model, data)
     _launch_internal(
         model, data, run_physics_thread=False, simulate=simulate)
     simulate = None