From 49211a05c58ddaa7af0459d8c30cc5b5c3128a17 Mon Sep 17 00:00:00 2001
From: Matija Kecman <matijak@google.com>
Date: Tue, 2 Jun 2026 10:27:01 -0700
Subject: [PATCH] Add support for custom physics step functions in StudioApp.

This change allows users to provide an optional `step_fn` callable to `StudioApp.update` and `StudioApp.update_from_viewer`. When provided, this function is called to advance the physics simulation instead of the default `step_control.advance`.

PiperOrigin-RevId: 925439890
Change-Id: Ia833503f6dd0c22fb8d75af6c5f06757496853e5
---
 .../mujoco/experimental/studio/studio_app.py  | 56 ++++++++++++++++---
 1 file changed, 47 insertions(+), 9 deletions(-)

diff --git a/python/mujoco/experimental/studio/studio_app.py b/python/mujoco/experimental/studio/studio_app.py
index 4b50ea25..fe030122 100644
--- a/python/mujoco/experimental/studio/studio_app.py
+++ b/python/mujoco/experimental/studio/studio_app.py
@@ -34,6 +34,7 @@ See the sample/ folder for usage examples.
 
 import os
 import sys
+import typing
 
 import mujoco
 from mujoco.experimental.studio import parser
@@ -45,6 +46,9 @@ import numpy as np
 
 from mujoco.experimental.dear_imgui import dear_imgui as imgui
 
+# Type alias for a custom physics step function.
+StepFn = typing.Callable[[mujoco.MjModel, mujoco.MjData], None]
+
 
 def load_model_from_file(
     model_path: str,
@@ -246,13 +250,28 @@ class StudioApp:
     else:
       mujoco.mjv_applyPerturbPose(self.model, self.data, perturb, 1)
 
-  def update_physics(self, perturb: mujoco.MjvPerturb) -> None:
-    """Applies the purturbations and advances the physics."""
+  def update_physics(
+      self,
+      perturb: mujoco.MjvPerturb,
+      *,
+      step_fn: StepFn | None = None,
+  ) -> None:
+    """Applies the perturbations and advances the physics.
+
+    Args:
+      perturb: The MuJoCo perturbation object.
+      step_fn: Optional custom physics step function.  When provided, it is
+        called instead of ``step_control.advance``.  The function receives
+        ``(model, data)`` and should step the simulation in-place.
+    """
     self.apply_perturb(perturb)
 
-    advance_status = self.step_control.advance(self.model, self.data)
-    if advance_status == sim.StepStatus.AUTO_RESET:
-      self.reset_physics()
+    if step_fn is not None:
+      step_fn(self.model, self.data)
+    else:
+      advance_status = self.step_control.advance(self.model, self.data)
+      if advance_status == sim.StepStatus.AUTO_RESET:
+        self.reset_physics()
 
   def reset_physics_gui(self) -> None:
     """GUI to Reset the physics i.e., the reset button."""
@@ -272,13 +291,15 @@ class StudioApp:
       camera: mujoco.MjvCamera,
       vis_options: mujoco.MjvOption,
       perturb: mujoco.MjvPerturb,
+      *,
       drop_file: str = '',
+      step_fn: StepFn | None = None,
   ) -> bool:
     """Update the simulation and handle user input.
 
     Handles mouse input to compute perturbations or camera motion.
     Handles keyboard input e.g., for keybindings or camera motion.
-    Applies the purturbations and advances the physics.
+    Applies the perturbations and advances the physics.
     The argument objects are provided by the viewer.
 
     Args:
@@ -287,6 +308,8 @@ class StudioApp:
       perturb: The MuJoCo perturbation object.
       drop_file: Path of a file dropped into the viewer window. If non-empty the
         current model is replaced with the dropped file.
+      step_fn: Optional custom physics step function.  When provided, it is
+        called instead of ``step_control.advance``.
 
     Returns:
       Whether the application should continue running, this is a
@@ -297,16 +320,31 @@ class StudioApp:
 
     self.handle_mouse_events(camera, vis_options, perturb)
     self.handle_keyboard_events(camera, vis_options)
-    self.update_physics(perturb)
+    self.update_physics(perturb, step_fn=step_fn)
     return self.is_running()
 
-  def update_from_viewer(self, viewer: viewer_protocol.Viewer) -> bool:
-    """Convenience wrapper around update() that unpacks viewer attributes."""
+  def update_from_viewer(
+      self,
+      viewer: viewer_protocol.Viewer,
+      *,
+      step_fn: StepFn | None = None,
+  ) -> bool:
+    """Convenience wrapper around update() that unpacks viewer attributes.
+
+    Args:
+      viewer: A viewer conforming to the Viewer protocol.
+      step_fn: Optional custom physics step function.  When provided, it is
+        called instead of ``step_control.advance``.
+
+    Returns:
+      Whether the application should continue running.
+    """
     return self.update(
         viewer.camera,
         viewer.vis_options,
         viewer.perturb,
         drop_file=viewer.get_drop_file(),
+        step_fn=step_fn,
     )
 
   def build_gui(