Add built-in simulation reset: Model.reset_state() + viewer R key#2468
Closed
vidurv-nvidia wants to merge 7 commits intonewton-physics:mainfrom
Closed
Add built-in simulation reset: Model.reset_state() + viewer R key#2468vidurv-nvidia wants to merge 7 commits intonewton-physics:mainfrom
vidurv-nvidia wants to merge 7 commits intonewton-physics:mainfrom
Conversation
Add reset_state() method to Model that copies initial configuration arrays back into an existing State without reallocating GPU memory. Supports optional forward kinematics re-evaluation via eval_fk flag. Includes comprehensive tests covering body, particle, joint reset, FK evaluation, and pointer stability. Co-Authored-By: Claude Sonnet 4.6 (1M context) <noreply@anthropic.com>
Co-Authored-By: Claude Sonnet 4.6 (1M context) <noreply@anthropic.com>
Add is_reset_requested() and clear_reset_request() methods to ViewerBase, with _reset_requested flag initialized in clear_model(). This gives simulation loops a way to detect when the viewer UI requests a state reset. Add tests in test_viewer_reset.py. Co-Authored-By: Claude Sonnet 4.6 (1M context) <noreply@anthropic.com>
Co-Authored-By: Claude Sonnet 4.6 (1M context) <noreply@anthropic.com>
Co-Authored-By: Claude Sonnet 4.6 (1M context) <noreply@anthropic.com>
Co-Authored-By: Claude Sonnet 4.6 (1M context) <noreply@anthropic.com>
Contributor
📝 WalkthroughWalkthroughAdds a public reset_state API and Model.reset_state implementation, viewer-level reset signaling and UI/keyboard hooks, example browser integration to perform in-place resets, and tests validating reset behavior and viewer reset signaling. Changes
Sequence DiagramsequenceDiagram
actor User
participant ViewerGL as Viewer (GL)
participant Example as Example Browser / Loop
participant Model
participant State
User->>ViewerGL: Press R key or click Reset
ViewerGL->>ViewerGL: request_reset() (set _reset_requested = True)
Example->>ViewerGL: is_reset_requested()?
ViewerGL-->>Example: True
alt example.reset() exists
Example->>Example: example.reset()
else
Example->>Model: reset_state(model, state, eval_fk=True)
Model->>State: Copy model q/qdot into state, zero force buffers
alt eval_fk=True and joints exist
Model->>Model: call eval_fk()
Model->>State: update body_q/body_qd from FK
end
end
Example->>ViewerGL: clear_reset_request()
ViewerGL->>ViewerGL: clear flag (set _reset_requested = False)
Estimated code review effort🎯 3 (Moderate) | ⏱️ ~25 minutes 🚥 Pre-merge checks | ✅ 2 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (2 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
vidurv-nvidia
had a problem deploying
to
external-pr-approval
GitHub Actions
Error
vidurv-nvidia
had a problem deploying
to
external-pr-approval
GitHub Actions
Error
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
Contributor
There was a problem hiding this comment.
Actionable comments posted: 5
🧹 Nitpick comments (1)
newton/examples/__init__.py (1)
215-216: Avoid toggling the viewer’s private reset flag from example code.
newton/examplesis user-facing code, so reaching intoself.viewer._reset_requestedhard-couples the browser to a private implementation detail. Please add a publicrequest_reset()/set_reset_requested()onViewerBaseand call that here instead.🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@newton/examples/__init__.py` around lines 215 - 216, The example code directly toggles the viewer's private flag (self.viewer._reset_requested); add a public method on ViewerBase—e.g., request_reset() or set_reset_requested(bool)—that sets the internal _reset_requested state and any related behavior, update ViewerBase implementation to encapsulate the flag, and change the example button handler to call self.viewer.request_reset() instead of touching _reset_requested directly.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@newton/__init__.py`:
- Around line 64-65: The public API now exports the symbol reset_state from
newton/__init__.py but the generated API docs are out of sync; run the
documentation generator (docs/generate_api.py) to regenerate docs/api/newton.rst
so reset_state appears in the API reference, then commit the updated docs file
along with the code change.
In `@newton/_src/sim/model.py`:
- Around line 852-884: reset_state currently skips optional extended buffers so
stale values persist; update reset_state to also handle the extended body
buffers by checking for their allocation and resetting them: if the model
provides self.body_q_prev, copy it into state.body_q_prev (similar to
wp.copy(state.body_q, self.body_q)), and for computed/force buffers like
body_qdd and body_parent_f zero the corresponding state buffers
(state.body_qdd.zero_(), state.body_parent_f.zero_()) when they exist; perform
these checks and operations alongside the existing body_count/joint_count
branches to ensure extended state is consistent after reset.
In `@newton/tests/test_reset_state.py`:
- Around line 74-90: The tests only mutate position arrays so they don't verify
velocities are restored; update test_reset_restores_particle_state() and
test_reset_restores_joint_state() to also mutate the corresponding velocity
arrays (state.particle_qd and state.joint_qd) before calling
model.reset_state(state), e.g. create junk arrays matching dtype/shape and
wp.copy into state.particle_qd and state.joint_qd, then after reset assert
state.particle_qd.numpy() equals the saved initial_particle_qd and
state.joint_qd.numpy() equals the saved initial_joint_qd; keep using the
existing model.reset_state, state.*, and the test helper names to locate the
spots to change.
- Around line 43-68: Update the tests to seed the force buffers with a non-zero
sentinel before calling model.reset_state so the zeroing check is meaningful: in
test_reset_restores_body_state (and the analogous test for particles) write a
non-zero array into state.body_f (and state.particle_f) using wp.array/wp.copy,
assert the mutation took effect, then call model.reset_state(state) and finally
assert the force buffers are zeroed; reference model.reset_state, state.body_f,
state.particle_f, and the test names (test_reset_restores_body_state and the
particle equivalent) to locate where to add the pre-reset seeding and
verification.
- Around line 111-145: The tests currently pass even if eval_fk is ignored
because the model's raw body_q/body_qd happen to equal the FK result; modify
both tests so the model's raw base state is deliberately different from FK by
setting model.body_q and model.body_qd to distinct values (e.g., fill with a
sentinel like 999 or zeros) immediately after building the articulation and
before calling eval_fk; then compute FK with newton.eval_fk(model,
state.joint_q, state.joint_qd, state) to get expected_body_q and
expected_body_qd, mutate state.body_q/state.body_qd as done now, call
model.reset_state(state, eval_fk=True) and assert state.body_q and state.body_qd
match expected_body_q/expected_body_qd, and in test_reset_without_eval_fk call
model.reset_state(state, eval_fk=False) and assert state.body_q and
state.body_qd equal the raw model.body_q/model.body_qd you set earlier.
---
Nitpick comments:
In `@newton/examples/__init__.py`:
- Around line 215-216: The example code directly toggles the viewer's private
flag (self.viewer._reset_requested); add a public method on ViewerBase—e.g.,
request_reset() or set_reset_requested(bool)—that sets the internal
_reset_requested state and any related behavior, update ViewerBase
implementation to encapsulate the flag, and change the example button handler to
call self.viewer.request_reset() instead of touching _reset_requested directly.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yml
Review profile: CHILL
Plan: Pro
Run ID: 4225109d-a5fe-437f-b7c4-5fc19ac6319b
📒 Files selected for processing (8)
newton/__init__.pynewton/_src/sim/__init__.pynewton/_src/sim/model.pynewton/_src/viewer/viewer.pynewton/_src/viewer/viewer_gl.pynewton/examples/__init__.pynewton/tests/test_reset_state.pynewton/tests/test_viewer_reset.py
Comment on lines
+64
to
65
| reset_state, | ||
| ) |
Contributor
There was a problem hiding this comment.
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
set -euo pipefail
# Confirm API generation script exists
fd -i "generate_api.py"
# Check whether docs already reference the new public symbol
if [ -d docs ]; then
rg -n --iglob "*.rst" --iglob "*.md" '\breset_state\b' docs || true
fi
# Confirm the symbol is exported in package surfaces
rg -n --iglob "*.py" '"reset_state"|reset_state' newton/__init__.py newton/_src/sim/__init__.pyRepository: newton-physics/newton
Length of output: 504
🏁 Script executed:
#!/bin/bash
set -euo pipefail
# Examine what the generate_api.py script does
head -50 docs/generate_api.py
# Find generated API doc files
find docs -type f \( -name "*.rst" -o -name "*.md" \) | head -20
# Check for API reference files that would list public symbols
find docs -type f -name "*api*" -o -name "*reference*" | head -20
# Look for generated auto-API documentation
find docs -type f \( -name "*.rst" \) -exec grep -l "reset_state\|newton\.reset_state" {} \;Repository: newton-physics/newton
Length of output: 2281
🏁 Script executed:
#!/bin/bash
set -euo pipefail
# Check the generated API doc for the main newton module
cat docs/api/newton.rstRepository: newton-physics/newton
Length of output: 1430
Run docs/generate_api.py to include the new reset_state export in the API documentation.
The symbol was added to the public API in newton/__init__.py but is missing from the generated API reference file (docs/api/newton.rst). Running the script will ensure the documentation stays in sync with the code.
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@newton/__init__.py` around lines 64 - 65, The public API now exports the
symbol reset_state from newton/__init__.py but the generated API docs are out of
sync; run the documentation generator (docs/generate_api.py) to regenerate
docs/api/newton.rst so reset_state appears in the API reference, then commit the
updated docs file along with the code change.
Codecov Report❌ Patch coverage is
📢 Thoughts on this report? Let us know! |
- reset_state now zeroes extended body buffers (body_qdd, body_parent_f, body_q_prev, mujoco:qfrc_actuator) when allocated on the state - Tests mutate velocities (not just positions) and seed forces with non-zero sentinels before verifying zeroing - FK tests use non-trivial joint offsets so raw model.body_q differs from FK-computed values, making eval_fk flag meaningful - Add request_reset() public method on ViewerBase; viewer_gl and example browser use it instead of touching _reset_requested directly - Regenerate API docs to include reset_state Co-Authored-By: Claude Sonnet 4.6 (1M context) <noreply@anthropic.com>
vidurv-nvidia
requested a deployment
to
external-pr-approval
GitHub Actions
Waiting
vidurv-nvidia
requested a deployment
to
external-pr-approval
GitHub Actions
Waiting
Contributor
There was a problem hiding this comment.
Actionable comments posted: 2
♻️ Duplicate comments (1)
newton/_src/sim/model.py (1)
875-892:⚠️ Potential issue | 🟠 MajorMove
body_q_prevrefresh after FK.With
eval_fk=True, Lines 889-892 can overwritestate.body_qfromjoint_q/joint_qd. Copyingbody_q_prevfromself.body_qon Lines 875-876 can therefore leave it out of sync with the final reset pose, which is exactly the value finite-difference velocity code expects.💡 Proposed fix
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 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) + + if self.body_count and getattr(state, "body_q_prev", None) is not None: + wp.copy(state.body_q_prev, state.body_q)🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@newton/_src/sim/model.py` around lines 875 - 892, The copy of state.body_q_prev from self.body_q is done before FK and can be overwritten by the subsequent FK update; move the wp.copy(state.body_q_prev, self.body_q) call so it happens after the eval_fk block (i.e., after the _eval_fk(self, self.joint_q, self.joint_qd, state) call) and keep the same attribute check (getattr(state, "body_q_prev", None) is not None) so body_q_prev is set to the final pose used for finite-difference velocity calculations.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@newton/_src/sim/model.py`:
- Around line 852-893: reset_state currently restores only built-in fields while
state() also clones attributes registered with assignment == STATE; modify
reset_state to iterate the same registration list/state-attribute registry used
by state() and for each registered STATE attribute copy the model's initial
value into the provided State (use wp.copy for Warp arrays and preserve
zero_/in-place semantics for buffers), i.e. locate the registry used by state()
and for each entry set state.<attr_name> = copy of self.<attr_name> (or call
.zero_() where appropriate) before running eval_fk so custom per-state Warp
arrays are reset as well.
In `@newton/examples/__init__.py`:
- Around line 282-290: In the fallback reset branch (where example has no reset
method but has a model) also recreate/reset the example.control to the model's
default by calling example.model.control() and assigning it to example.control;
keep the existing logic that resets example.model.reset_state for attributes
"state_0" and "state_1" and resets example.sim_time to 0.0, but add the
example.control reset so runtime edits to control don't persist across resets.
---
Duplicate comments:
In `@newton/_src/sim/model.py`:
- Around line 875-892: The copy of state.body_q_prev from self.body_q is done
before FK and can be overwritten by the subsequent FK update; move the
wp.copy(state.body_q_prev, self.body_q) call so it happens after the eval_fk
block (i.e., after the _eval_fk(self, self.joint_q, self.joint_qd, state) call)
and keep the same attribute check (getattr(state, "body_q_prev", None) is not
None) so body_q_prev is set to the final pose used for finite-difference
velocity calculations.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yml
Review profile: CHILL
Plan: Pro
Run ID: b78977b5-2a60-4cd5-b6bb-f0cacdfb3b84
📒 Files selected for processing (7)
docs/api/newton.rstnewton/_src/sim/model.pynewton/_src/viewer/viewer.pynewton/_src/viewer/viewer_gl.pynewton/examples/__init__.pynewton/tests/test_reset_state.pynewton/tests/test_viewer_reset.py
✅ Files skipped from review due to trivial changes (2)
- docs/api/newton.rst
- newton/tests/test_viewer_reset.py
🚧 Files skipped from review as they are similar to previous changes (2)
- newton/_src/viewer/viewer.py
- newton/tests/test_reset_state.py
Comment on lines
+852
to
+893
| 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) | ||
|
|
Contributor
There was a problem hiding this comment.
Reset registered custom STATE attributes too.
state() clones all attributes registered with assignment == STATE, but reset_state() only restores the built-in fields. Any model that keeps extra per-state Warp arrays on State will carry their last runtime values across reset, so this new public API does not actually return the full state to its initial configuration.
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@newton/_src/sim/model.py` around lines 852 - 893, reset_state currently
restores only built-in fields while state() also clones attributes registered
with assignment == STATE; modify reset_state to iterate the same registration
list/state-attribute registry used by state() and for each registered STATE
attribute copy the model's initial value into the provided State (use wp.copy
for Warp arrays and preserve zero_/in-place semantics for buffers), i.e. locate
the registry used by state() and for each entry set state.<attr_name> = copy of
self.<attr_name> (or call .zero_() where appropriate) before running eval_fk so
custom per-state Warp arrays are reset as well.
Comment on lines
+282
to
+290
| 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 |
Contributor
There was a problem hiding this comment.
Reset example.control in the fallback path too.
The generic reset only rewinds state_0/state_1 and sim_time. Examples commonly allocate self.control = self.model.control() alongside those states, so any runtime edits to targets/actuation survive the reset and the next step no longer starts from the original setup.
💡 Minimal fix
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 getattr(example, "control", None) is not None:
+ example.control = example.model.control()
if hasattr(example, "sim_time"):
example.sim_time = 0.0🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@newton/examples/__init__.py` around lines 282 - 290, In the fallback reset
branch (where example has no reset method but has a model) also recreate/reset
the example.control to the model's default by calling example.model.control()
and assigning it to example.control; keep the existing logic that resets
example.model.reset_state for attributes "state_0" and "state_1" and resets
example.sim_time to 0.0, but add the example.control reset so runtime edits to
control don't persist across resets.
eric-heiden
reviewed
Apr 15, 2026
Member
eric-heiden
left a comment
eric-heiden
left a comment
There was a problem hiding this comment.
Thanks @vidurv-nvidia! I would like to wait with committing to a reset functionality until we have #1266 implemented which will tackle the solver-specific reset logic, as requested by @vastsoun.
Not every solver needs eval_fk and custom attributes are also not handled here in reset_state yet (which will be the task of solver.reset() to handle).
Member
|
Thanks @vidurv-nvidia. May I remind you to adhere to the policy of searching for open tickets/discussions first in the future before opening a PR? This is especially relevant for any change touching public API. |
Contributor
Author
|
@eric-heiden Thanks for the feedback. I will look for that PR closely. |
Member
|
no worries - can we take the discussion to the other ticket and iterate there, or is there any need to keep this PR around? |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
3 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Description
Add a common reset paradigm to Newton so simulations can be reset to initial state without custom per-simulation logic.
Three components:
Model.reset_state(state, eval_fk=True)— Restores aStateto the model's initial configuration in-place usingwp.copy(no GPU reallocation). Also available asnewton.reset_state(model, state).Viewer reset triggers —
Rkey shortcut and "Reset" button in the Model Information panel (next to Pause). Sets a_reset_requestedflag onViewerBase, withis_reset_requested()/clear_reset_request()API matching the existingis_paused()pattern.Example run loop integration — The
run()loop checksviewer.is_reset_requested()before stepping. Callsexample.reset()if available (lightweight path), otherwise usesmodel.reset_state()onstate_0/state_1automatically.Checklist
CHANGELOG.mdhas been updated (if user-facing change)Test plan
10 tests covering:
Model.reset_state()restores body, particle, and joint stateptr)ViewerBasereset signal lifecycleNew feature / API change
Summary by CodeRabbit
New Features
Tests