Add built-in simulation reset: Model.reset_state() + viewer R key

docs/api/newton.rst

@@ -69,6 +69,7 @@ newton
    eval_ik
    eval_jacobian
    eval_mass_matrix
+   reset_state
 
 .. rubric:: Constants
 

newton/__init__.py

@@ -61,6 +61,7 @@
     eval_ik,
     eval_jacobian,
     eval_mass_matrix,
+    reset_state,
 )
 
 __all__ += [
@@ -78,6 +79,7 @@
     "eval_ik",
     "eval_jacobian",
     "eval_mass_matrix",
+    "reset_state",
 ]
 
 # ==================================================================================

newton/_src/sim/__init__.py

@@ -30,4 +30,19 @@
     "eval_ik",
     "eval_jacobian",
     "eval_mass_matrix",
+    "reset_state",
 ]
+
+
+def reset_state(model: Model, state: State, eval_fk: bool = True) -> None:
+    """Reset a state to the model's initial configuration.
+
+    Convenience wrapper for :meth:`Model.reset_state`. See that method for
+    full documentation.
+
+    Args:
+        model: The model whose initial configuration to restore.
+        state: The state object to reset.
+        eval_fk: Whether to re-evaluate forward kinematics.
+    """
+    model.reset_state(state, eval_fk=eval_fk)

newton/_src/sim/model.py

@@ -849,6 +849,48 @@ def state(self, requires_grad: bool | None = None) -> State:
 
         return s
 
+    def reset_state(self, state: State, eval_fk: bool = True) -> None:
+        """
+        Reset a :class:`State` to this model's initial configuration in-place.
+
+        Copies the model's initial position and velocity arrays into ``state``
+        and zeroes all force arrays. Unlike :meth:`state`, this reuses the
+        existing GPU allocations -- no new arrays are created.
+
+        Args:
+            state: The state object to reset (must have been created by this model).
+            eval_fk: If True and the model has joints, re-evaluate forward
+                kinematics so that :attr:`State.body_q` and :attr:`State.body_qd`
+                are consistent with the restored joint coordinates.
+        """
+        if self.particle_count:
+            wp.copy(state.particle_q, self.particle_q)
+            wp.copy(state.particle_qd, self.particle_qd)
+            state.particle_f.zero_()
+
+        if self.body_count:
+            wp.copy(state.body_q, self.body_q)
+            wp.copy(state.body_qd, self.body_qd)
+            state.body_f.zero_()
+            if getattr(state, "body_q_prev", None) is not None:
+                wp.copy(state.body_q_prev, self.body_q)
+            if getattr(state, "body_qdd", None) is not None:
+                state.body_qdd.zero_()
+            if getattr(state, "body_parent_f", None) is not None:
+                state.body_parent_f.zero_()
+
+        if self.joint_count:
+            wp.copy(state.joint_q, self.joint_q)
+            wp.copy(state.joint_qd, self.joint_qd)
+            mujoco_ns = getattr(state, "mujoco", None)
+            if mujoco_ns is not None and getattr(mujoco_ns, "qfrc_actuator", None) is not None:
+                mujoco_ns.qfrc_actuator.zero_()
+
+        if eval_fk and self.joint_count:
+            from .articulation import eval_fk as _eval_fk  # noqa: PLC0415
+
+            _eval_fk(self, self.joint_q, self.joint_qd, state)
+
     def control(self, requires_grad: bool | None = None, clone_variables: bool = True) -> Control:
         """
         Create and return a new :class:`Control` object for this model.

newton/_src/viewer/viewer.py

@@ -59,6 +59,32 @@ def is_paused(self) -> bool:
         """
         return False
 
+    def is_reset_requested(self) -> bool:
+        """Report whether a simulation reset has been requested.
+
+        The flag is set by viewer UI controls (e.g. the *R* key or
+        *Reset* button in :class:`ViewerGL`).  Callers should check this
+        once per frame before stepping the simulation and call
+        :meth:`clear_reset_request` after handling the reset.
+
+        Returns:
+            bool: True when a reset has been requested.
+        """
+        return self._reset_requested
+
+    def clear_reset_request(self) -> None:
+        """Clear the reset-requested flag after the reset has been handled."""
+        self._reset_requested = False
+
+    def request_reset(self) -> None:
+        """Request a simulation reset.
+
+        Sets the internal flag that :meth:`is_reset_requested` queries.
+        The next frame's run-loop iteration (or user code) should detect
+        the flag, perform the reset, and call :meth:`clear_reset_request`.
+        """
+        self._reset_requested = True
+
     def is_key_down(self, key: str | int) -> bool:
         """Default key query API. Concrete viewers can override.
 
@@ -111,6 +137,9 @@ def clear_model(self) -> None:
         # Picking
         self.picking_enabled = True
 
+        # Reset signal
+        self._reset_requested = False
+
         # Display options
         self.show_joints = False
         self.show_com = False

newton/_src/viewer/viewer_gl.py

@@ -1637,6 +1637,9 @@ def on_key_press(self, symbol: int, modifiers: int):
         elif symbol == pyglet.window.key.F:
             # Frame camera around model bounds
             self._frame_camera_on_model()
+        elif symbol == pyglet.window.key.R:
+            # Request simulation reset
+            self.request_reset()
         elif symbol == pyglet.window.key.ESCAPE:
             # Exit with Escape key
             self.renderer.close()
@@ -1967,6 +1970,8 @@ def _render_left_panel(self):
 
                     # Pause simulation checkbox
                     changed, self._paused = imgui.checkbox("Pause", self._paused)
+                    if imgui.button("Reset"):
+                        self.request_reset()
 
                 # Visualization Controls section
                 imgui.set_next_item_open(True, imgui.Cond_.appearing)

newton/examples/__init__.py

@@ -187,7 +187,6 @@ class _ExampleBrowser:
     def __init__(self, viewer):
         self.viewer = viewer
         self.switch_target: str | None = None
-        self._reset_requested = False
         self.callback = None
         self._tree: dict[str, list[tuple[str, str]]] = {}
 
@@ -214,7 +213,7 @@ def _browser_ui(imgui):
                         imgui.tree_pop()
                 imgui.separator()
                 if imgui.button("Reset"):
-                    self._reset_requested = True
+                    self.viewer.request_reset()
 
         self.callback = _browser_ui
         viewer.register_ui_callback(_browser_ui, position="panel")
@@ -240,7 +239,6 @@ def switch(self, example_class):
 
     def reset(self, example_class):
         """Reset the current example by re-creating it. Returns the new example or None."""
-        self._reset_requested = False
         self.viewer.clear_model()
         try:
             parser = getattr(example_class, "create_parser", create_parser)()
@@ -279,8 +277,17 @@ def run(example, args):
             example, example_class = browser.switch(example_class)
             continue
 
-        if browser is not None and browser._reset_requested:
-            example = browser.reset(example_class)
+        if viewer.is_reset_requested():
+            viewer.clear_reset_request()
+            if hasattr(example, "reset"):
+                example.reset()
+            elif hasattr(example, "model"):
+                for attr in ("state_0", "state_1"):
+                    s = getattr(example, attr, None)
+                    if s is not None:
+                        example.model.reset_state(s)
+                if hasattr(example, "sim_time"):
+                    example.sim_time = 0.0
             continue
 
         if example is None: