Refactor Viewer to be a base class owning the communication with the simulation, handler registry, core visualization objects and the render function

ViewerApp is simplified, interacting with the endpoint and handlers through the Viewer instance. This change makes it possible to write simulation viewers without using the UI/UX provided by ViewerApp (which should be renamed StudioApp) PiperOrigin-RevId: 943642660 Change-Id: Id8eb3d4e6a27ceda93833367fec4699fa6bc838e
2026-07-06 20:49:58 -07:00
parent ae58855c15
commit 8079ab3d4d
9 changed files with 384 additions and 322 deletions
@@ -0,0 +1,49 @@
+# Copyright 2026 DeepMind Technologies Limited
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Utility functions for Studio viewer implementations."""
+
+import mujoco
+from mujoco.experimental.studio import messages
+from mujoco.experimental.studio import viewer_protocol
+import numpy as np
+
+
+def apply_perturb(
+    viewer: viewer_protocol.Viewer,
+    model: mujoco.MjModel | None,
+    data: mujoco.MjData | None,
+    is_paused: bool = False,
+) -> None:
+  """Apply perturbation to the model and send updated forces to the sim thread."""
+  if model is None or data is None:
+    return
+  perturb = viewer.perturb
+  if not is_paused:
+    sig = int(mujoco.mjtState.mjSTATE_XFRC_APPLIED)
+    size = mujoco.mj_stateSize(model, sig)
+    zero_state = np.zeros(size, np.float64)
+    mujoco.mj_setState(model, data, zero_state, sig)
+    mujoco.mjv_applyPerturbPose(model, data, perturb, 0)
+    mujoco.mjv_applyPerturbForce(model, data, perturb)
+  else:
+    mujoco.mjv_applyPerturbPose(model, data, perturb, 1)
+
+  xfrc_sig: int = int(mujoco.mjtState.mjSTATE_XFRC_APPLIED)
+  xfrc_size: int = mujoco.mj_stateSize(model, xfrc_sig)
+  xfrc_state: np.ndarray = np.zeros(xfrc_size, np.float64)
+
+  mujoco.mj_getState(model, data, xfrc_state, xfrc_sig)
+  viewer.send_to_sim(
+      messages.PerturbEvent(state=xfrc_state, state_sig=xfrc_sig)
+  )