add is_paused, set_global

tools/machine-learning/mujoco/packages/mujoco-interactive-viewer/src/mujoco_interactive_viewer/context.py

@@ -1,4 +1,9 @@
-from mujoco_interactive_viewer.viewer import Viewer
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from mujoco_interactive_viewer.viewer import Viewer
 
 _viewer: Viewer | None = None
 

tools/machine-learning/mujoco/packages/mujoco-interactive-viewer/src/mujoco_interactive_viewer/viewer.py

@@ -8,6 +8,7 @@
 from mujoco._structs import MjvGeom
 from numpy.typing import ArrayLike, NDArray
 
+from mujoco_interactive_viewer.context import set_global_viewer
 from mujoco_interactive_viewer.figure import Figure
 from mujoco_interactive_viewer.interaction import InteractionState
 from mujoco_interactive_viewer.marker import Marker
@@ -52,11 +53,13 @@ def __init__(
         width: int | None = None,
         height: int | None = None,
         font_scale: mujoco.mjtFontScale = mujoco.mjtFontScale.mjFONTSCALE_100,
+        is_paused: bool = False,
+        set_global: bool = True,
     ) -> None:
         self._gui_lock = Lock()
         self._interaction_state = InteractionState()
         self._visualization_state = VisualizationState()
-        self._render_state = RenderState()
+        self._render_state = RenderState(is_paused=is_paused)
         self.is_alive = True
 
         self.model = model
@@ -105,6 +108,9 @@ def __init__(
         self._figures: dict[str, Figure] = {}
         self._overlay: dict[mujoco.mjtGridPos, Overlay] = {}
 
+        if set_global:
+            set_global_viewer(self)
+
     def add_marker(
         self,
         kind: mujoco.mjtGeom,

tools/machine-learning/mujoco/scripts/mujoco-walking.py

@@ -2,7 +2,7 @@
 
 import click
 import numpy as np
-from mujoco_interactive_viewer import Viewer, set_global_viewer
+from mujoco_interactive_viewer import Viewer
 from nao_env import NaoWalking
 from stable_baselines3 import PPO
 
@@ -31,7 +31,6 @@ def main(*, throw_tomatoes: bool, load_policy: str | None) -> None:
     env.initialize_terrain(max_height=0.1, step_height=0.01)
 
     viewer = Viewer(env.model, env.data)
-    set_global_viewer(viewer)
     rewards_figure = viewer.figure("rewards")
     rewards_figure.set_title("Rewards")
     rewards_figure.set_x_label("Step")