Fix masked coupled solver resets - #3649
Conversation
|
Note Reviews pausedIt looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the Use the following commands to manage reviews:
Use the checkboxes below for quick actions:
📝 WalkthroughWalkthroughMasked reset support centralizes world-mask validation and applies selective clearing across coupled state, contact matching, ADMM history, implicit MPM history, MuJoCo synchronization, proxy state, and solver-specific buffers. Tests cover selected, unselected, global, empty, invalid, and replayed masks. ChangesMasked world reset
Estimated code review effort: 5 (Critical) | ~90 minutes Possibly related issues
Possibly related PRs
Suggested reviewers: 🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✨ 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 |
GitHub Actions
Error
GitHub Actions
Error
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 1
🧹 Nitpick comments (6)
newton/_src/solvers/mujoco/solver_mujoco.py (1)
3765-3780: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueAdd CPU-backend coverage for masked MuJoCo buffer resets.
newton/tests/test_mujoco_reset.py::TestMuJoCoResetexercisesSolverMuJoCo’s default mjwarp path, butreset()’s CPU-only branch takes a separate control-flow path and checkslocal_world_mask[0]. Add ause_mujoco_cpu=Truemulti-world reset test that covers global-only masks, local-world-only masks, and the early-return behavior.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@newton/_src/solvers/mujoco/solver_mujoco.py` around lines 3765 - 3780, Add CPU-backend coverage in TestMuJoCoReset for SolverMuJoCo.reset() with use_mujoco_cpu=True and multiple worlds. Exercise global-only and local-world-only masks, verifying the corresponding MuJoCo buffers reset, and verify early return when local_world_mask[0] is false; preserve existing default mjwarp tests.newton/_src/geometry/contact_match.py (1)
762-779: 🎯 Functional Correctness | 🔵 Trivial | ⚡ Quick win
ContactMatcher.reset()skips the shared mask-validation contract.Unlike
SolverBase._validate_reset_world_mask(dtype, ndim, device, and W/W+1 length checks), thisreset()callsworld_mask.numpy()and OR's it into the host mask with no dtype/device/shape validation. A caller passing a wrong-dtype, wrong-device, or wrong-length mask will hit a raw numpy/Warp error instead of the clear, consistent error messages the rest of this PR establishes for the same contract (e.g.SolverBase,SolverVBD,SolverKamino).♻️ Suggested validation for consistency with the shared contract
def reset(self, world_mask: wp.array[wp.bool] | None = None) -> None: """Clear all or reset-selected cross-frame contact history. ... """ if world_mask is None: self._prev_count.zero_() self._reset_world_mask_host.fill(False) self._reset_world_mask.zero_() return + if not isinstance(world_mask, wp.array) or world_mask.dtype != wp.bool or world_mask.ndim != 1: + raise TypeError("'world_mask' must be a one-dimensional Warp bool array or None.") + if world_mask.shape[0] not in (self._world_count, self._world_count + 1): + raise ValueError( + f"'world_mask' length {world_mask.shape[0]} must equal world_count " + f"({self._world_count}) or world_count + 1 ({self._world_count + 1})." + ) mask = world_mask.numpy() if not bool(mask.any()): return self._reset_world_mask_host[: mask.shape[0]] |= mask self._reset_world_mask.assign(self._reset_world_mask_host)Since
newton/_src/sim/collide.py(whereCollisionPipeline.reset_contact_matching()presumably calls into this) is not part of this review, please confirm whether validation already happens at that call site before deciding whether to duplicate it here.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@newton/_src/geometry/contact_match.py` around lines 762 - 779, Update ContactMatcher.reset to enforce the same world-mask validation contract as SolverBase._validate_reset_world_mask before calling world_mask.numpy() or modifying reset masks, including dtype, ndim, device, and W/W+1 length checks. First verify whether CollisionPipeline.reset_contact_matching already performs this validation; if it does not, add validation in ContactMatcher.reset using the shared validator or equivalent established checks and preserve the existing full-reset behavior.newton/_src/solvers/coupled/interface.py (1)
248-261: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winAdd an
Args:section to the new hook's docstring.
coupling_sync_reset_statehas four parameters but noArgs:block, unlike sibling hooks in this file (e.g.coupling_eval_effective_mass,coupling_eval_effective_mass_block) which document each parameter. Since this is a new public extension point solver authors will override, documentingstate_in,state_out,world_mask, andflagsimproves discoverability.As per coding guidelines: "Use Google-style docstrings, keep types in annotations, format
Args:entries asname: description..."📝 Proposed docstring addition
def coupling_sync_reset_state( self, state_in: State, state_out: State, world_mask: wp.array[wp.bool], flags: StateFlags | int | None, ) -> None: """Synchronize solver-owned state after a masked coupled reset. Public state arrays are synchronized by the coupled solver. Override this hook only for persistent custom state arrays. + + Args: + state_in: Entry-local reset input state. + state_out: Entry-local reset output state. + world_mask: Validated reset mask forwarded unchanged from the coupled reset call. + flags: Reset flags bitmask forwarded unchanged from the coupled reset call. """ del state_in, state_out, world_mask, flags🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@newton/_src/solvers/coupled/interface.py` around lines 248 - 261, Add a Google-style Args section to the coupling_sync_reset_state docstring, documenting state_in, state_out, world_mask, and flags with concise descriptions; keep their type information in the existing annotations and leave the hook behavior unchanged.Source: Coding guidelines
newton/_src/sim/collide.py (1)
1227-1246: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winConsider sharing world_mask validation with
SolverBase._validate_reset_world_mask.This method re-implements the same type/dtype/ndim/device/shape (
world_countorworld_count + 1) checks asSolverBase._validate_reset_world_maskinnewton/_src/solvers/solver.py. Extracting a shared module-level validator (e.g. in a common utils module) both call would prevent the two copies from silently drifting if the mask contract changes.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@newton/_src/sim/collide.py` around lines 1227 - 1246, Extract the duplicated world_mask validation from reset_contact_matching and SolverBase._validate_reset_world_mask into a shared module-level validator, then call it from both methods. Preserve all existing type, dtype, dimensionality, device, and world_count/world_count + 1 shape checks and their current validation behavior.newton/_src/solvers/coupled/solver_coupled.py (1)
2108-2153: 🚀 Performance & Scalability | 🔵 TrivialAvoid host-side
.numpy()roundtrips in_reset_rows.
_reset_rowscopies model entry mapping arrays to host, applies the mask on host, then launches entry-local Warp arrays. Withworld_mask, maskedreset()is intended as the selective RL reset path; on this path it syncs per entry/domain pair, which can scale poorly in repeated per-step resets. Keep the row-selection/result on-device or amortize the selection across entries instead of redoing it inside_reset_rows.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@newton/_src/solvers/coupled/solver_coupled.py` around lines 2108 - 2153, The _reset_rows method currently materializes world mappings with host-side .numpy() calls and repeats masked selection for every entry/domain pair. Replace this path with device-resident row selection, or cache/amortize the shared selection so world_mask resets avoid per-entry/domain synchronization while preserving the returned per-entry domain mappings.newton/_src/solvers/coupled/solver_coupled_proxy.py (1)
1078-1105: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winDocument why masked reset only clears
coupling_forces.The masked branch zeroes only
mapping.coupling_forces, while the full reset also clearscoupling_forces_previous,aitken_residual_previous,aitken_stats,aitken_relaxation,aitken_has_previous, andproxy_qd_before. This is safe today because_stash_proxy_feedback/wp.copy(proxy.proxy_qd_before, ...)unconditionally overwrite those buffers before they're read each step, and_reset_aitken_iteration_state()resets the Aitken scalars every_step_coupledcall regardless ofreset(). However, this asymmetry is non-obvious and risks a future regression (e.g., if the overwrite-before-use invariant is ever broken elsewhere). A short comment explaining why these buffers are intentionally excluded from the masked path would help future maintainers avoid re-introducing the removed per-buffer masked kernels unnecessarily, or worse, missing a case where the invariant doesn't hold.📝 Suggested comment
else: + # coupling_forces_previous / aitken_* / proxy_qd_before are intentionally left + # untouched here: they are unconditionally overwritten (stash/copy) before being + # read on the next step, and Aitken scalars are reset every _step_coupled call + # regardless of reset(), so masking them would be redundant. for mapping_group, entity_world in (🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@newton/_src/solvers/coupled/solver_coupled_proxy.py` around lines 1078 - 1105, Add a concise comment in the masked reset branch near the `_zero_global_proxy_values_masked_kernel` launch explaining that only `coupling_forces` is intentionally cleared because `_stash_proxy_feedback`/`wp.copy` overwrite `proxy_qd_before` and related feedback/Aitken buffers before use, while `_reset_aitken_iteration_state()` resets Aitken state each step. Clarify that the full reset still clears those buffers, preserving the intentional asymmetry.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@newton/tests/test_implicit_mpm_multiworld_sparse.py`:
- Around line 523-530: Extend the reset assertions after solver.reset in the
sparse-grid test to also validate solver._grid_accumulated_status. Assert its
value remains the expected Volume.REBUILD_VOXEL_CAPACITY_EXCEEDED status,
alongside the existing _grid_status assertion, covering both persistent status
buffers.
---
Nitpick comments:
In `@newton/_src/geometry/contact_match.py`:
- Around line 762-779: Update ContactMatcher.reset to enforce the same
world-mask validation contract as SolverBase._validate_reset_world_mask before
calling world_mask.numpy() or modifying reset masks, including dtype, ndim,
device, and W/W+1 length checks. First verify whether
CollisionPipeline.reset_contact_matching already performs this validation; if it
does not, add validation in ContactMatcher.reset using the shared validator or
equivalent established checks and preserve the existing full-reset behavior.
In `@newton/_src/sim/collide.py`:
- Around line 1227-1246: Extract the duplicated world_mask validation from
reset_contact_matching and SolverBase._validate_reset_world_mask into a shared
module-level validator, then call it from both methods. Preserve all existing
type, dtype, dimensionality, device, and world_count/world_count + 1 shape
checks and their current validation behavior.
In `@newton/_src/solvers/coupled/interface.py`:
- Around line 248-261: Add a Google-style Args section to the
coupling_sync_reset_state docstring, documenting state_in, state_out,
world_mask, and flags with concise descriptions; keep their type information in
the existing annotations and leave the hook behavior unchanged.
In `@newton/_src/solvers/coupled/solver_coupled_proxy.py`:
- Around line 1078-1105: Add a concise comment in the masked reset branch near
the `_zero_global_proxy_values_masked_kernel` launch explaining that only
`coupling_forces` is intentionally cleared because
`_stash_proxy_feedback`/`wp.copy` overwrite `proxy_qd_before` and related
feedback/Aitken buffers before use, while `_reset_aitken_iteration_state()`
resets Aitken state each step. Clarify that the full reset still clears those
buffers, preserving the intentional asymmetry.
In `@newton/_src/solvers/coupled/solver_coupled.py`:
- Around line 2108-2153: The _reset_rows method currently materializes world
mappings with host-side .numpy() calls and repeats masked selection for every
entry/domain pair. Replace this path with device-resident row selection, or
cache/amortize the shared selection so world_mask resets avoid per-entry/domain
synchronization while preserving the returned per-entry domain mappings.
In `@newton/_src/solvers/mujoco/solver_mujoco.py`:
- Around line 3765-3780: Add CPU-backend coverage in TestMuJoCoReset for
SolverMuJoCo.reset() with use_mujoco_cpu=True and multiple worlds. Exercise
global-only and local-world-only masks, verifying the corresponding MuJoCo
buffers reset, and verify early return when local_world_mask[0] is false;
preserve existing default mjwarp tests.
🪄 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 Plus
Run ID: df6dbd9b-5351-4edd-ba62-1ffe0eb5279f
📒 Files selected for processing (23)
CHANGELOG.mdnewton/_src/geometry/contact_match.pynewton/_src/sim/collide.pynewton/_src/solvers/coupled/admm_utils.pynewton/_src/solvers/coupled/interface.pynewton/_src/solvers/coupled/solver_coupled.pynewton/_src/solvers/coupled/solver_coupled_admm.pynewton/_src/solvers/coupled/solver_coupled_proxy.pynewton/_src/solvers/implicit_mpm/implicit_mpm_solver_kernels.pynewton/_src/solvers/implicit_mpm/solver_implicit_mpm.pynewton/_src/solvers/kamino/_src/solvers/dvi/solver.pynewton/_src/solvers/kamino/_src/solvers/padmm/solver.pynewton/_src/solvers/kamino/solver_kamino.pynewton/_src/solvers/mujoco/kernels.pynewton/_src/solvers/mujoco/solver_mujoco.pynewton/_src/solvers/solver.pynewton/_src/solvers/vbd/solver_vbd.pynewton/tests/test_admm_coupled_solver.pynewton/tests/test_contact_matching.pynewton/tests/test_coupled_solver.pynewton/tests/test_implicit_mpm_multiworld_sparse.pynewton/tests/test_mujoco_reset.pynewton/tests/test_solver_vbd.py
gdaviet
left a comment
gdaviet
left a comment
There was a problem hiding this comment.
Graph-capturability and attribute-metadata findings.
6704635 to
5dac41a
Compare
GitHub Actions
Inactive
GitHub Actions
Inactive
|
@gdaviet One follow-up from testing this that I’d appreciate your take on. I kept this PR scoped to masked reset instead of adding broad configuration rejections, but two adjacent cases remain:
I left both existing behaviors intact here rather than rejecting otherwise valid configurations. |
There was a problem hiding this comment.
Actionable comments posted: 1
🧹 Nitpick comments (2)
newton/_src/solvers/coupled/solver_coupled.py (1)
490-509: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueUnresolvable frequencies silently drop out of reset.
When
_custom_frequency_row_worldsreturnsNone(custom frequency with no attribute referencingWORLD), the frequency is omitted from_reset_row_world, so_launch_reset_view_kernel/_launch_reset_owned_scatterbecome no-ops for every state attribute at that frequency — including the full-reset path (world_mask=None→_full_reset_world_mask). A registeredSTATEattribute at such a frequency will never be synchronized or cleared, with no diagnostic. Consider a debug-level log so this is discoverable.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@newton/_src/solvers/coupled/solver_coupled.py` around lines 490 - 509, The reset-row-world construction in the loop handling required frequencies silently skips custom frequencies when _custom_frequency_row_worlds returns None. Add a debug-level diagnostic at that branch identifying the unresolved item_frequency and that reset handling is being skipped, while preserving the existing continue behavior and all other reset paths.newton/_src/geometry/contact_match.py (1)
594-612: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winExtract the sort-key shape decode into
make_contact_sort_key.The broken-contact kernel in
newton/_src/geometry/contact_match.py:607-608duplicates the layout (shape_a << 43 | shape_b << 23 | sub_key, with 20-bit masks) fromnewton/_src/geometry/contact_data.py. Add an inversewp.func, e.g._decode_sort_key_shapes, to both the packer and this kernel so the layout stays single-sourced.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@newton/_src/geometry/contact_match.py` around lines 594 - 612, Extract the duplicated shape-bit decoding from the broken-contact kernel into a reusable inverse wp.func alongside make_contact_sort_key in contact_data.py, using the same 43/23 shifts and 20-bit masks. Update the broken-contact kernel to call this decoder for shape0 and shape1, preserving the existing reset_world_selected behavior.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@newton/_src/geometry/contact_match.py`:
- Around line 764-770: Update the ContactMatcher.reset docstring to add a
Google-style Args: entry for world_mask, documenting the accepted W or W + 1
shape contract and that applying a mask accumulates selection rather than
clearing _prev_count; leave the reset behavior unchanged.
---
Nitpick comments:
In `@newton/_src/geometry/contact_match.py`:
- Around line 594-612: Extract the duplicated shape-bit decoding from the
broken-contact kernel into a reusable inverse wp.func alongside
make_contact_sort_key in contact_data.py, using the same 43/23 shifts and 20-bit
masks. Update the broken-contact kernel to call this decoder for shape0 and
shape1, preserving the existing reset_world_selected behavior.
In `@newton/_src/solvers/coupled/solver_coupled.py`:
- Around line 490-509: The reset-row-world construction in the loop handling
required frequencies silently skips custom frequencies when
_custom_frequency_row_worlds returns None. Add a debug-level diagnostic at that
branch identifying the unresolved item_frequency and that reset handling is
being skipped, while preserving the existing continue behavior and all other
reset paths.
🪄 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 Plus
Run ID: 9c7621a9-9078-4e06-bf13-1e2ee8b96160
📒 Files selected for processing (25)
CHANGELOG.mddocs/concepts/coupling.rstnewton/_src/core/reset.pynewton/_src/geometry/contact_match.pynewton/_src/sim/collide.pynewton/_src/solvers/coupled/admm_utils.pynewton/_src/solvers/coupled/interface.pynewton/_src/solvers/coupled/solver_coupled.pynewton/_src/solvers/coupled/solver_coupled_admm.pynewton/_src/solvers/coupled/solver_coupled_proxy.pynewton/_src/solvers/implicit_mpm/implicit_mpm_solver_kernels.pynewton/_src/solvers/implicit_mpm/solver_implicit_mpm.pynewton/_src/solvers/kamino/_src/solvers/dvi/solver.pynewton/_src/solvers/kamino/_src/solvers/padmm/solver.pynewton/_src/solvers/kamino/solver_kamino.pynewton/_src/solvers/mujoco/kernels.pynewton/_src/solvers/mujoco/solver_mujoco.pynewton/_src/solvers/solver.pynewton/_src/solvers/vbd/solver_vbd.pynewton/tests/test_admm_coupled_solver.pynewton/tests/test_contact_matching.pynewton/tests/test_coupled_solver.pynewton/tests/test_implicit_mpm_multiworld_sparse.pynewton/tests/test_mujoco_reset.pynewton/tests/test_solver_vbd.py
🚧 Files skipped from review as they are similar to previous changes (12)
- CHANGELOG.md
- newton/_src/solvers/coupled/interface.py
- newton/_src/solvers/mujoco/kernels.py
- newton/_src/sim/collide.py
- newton/_src/solvers/kamino/_src/solvers/padmm/solver.py
- newton/_src/solvers/kamino/_src/solvers/dvi/solver.py
- newton/_src/solvers/kamino/solver_kamino.py
- newton/_src/solvers/coupled/solver_coupled_admm.py
- newton/tests/test_implicit_mpm_multiworld_sparse.py
- newton/tests/test_admm_coupled_solver.py
- newton/_src/solvers/mujoco/solver_mujoco.py
- newton/_src/solvers/coupled/solver_coupled_proxy.py
I think 1 is fine, at least for the scope of this PR. Typically if |
chschuma-disney
left a comment
chschuma-disney
left a comment
There was a problem hiding this comment.
Thanks for looking into this. The changes on Kamino's side look good to me.
Good catch pointing out the issues with the callbacks. I'll make sure that we track the issue and fix it in a follow-up PR.
I'm not too familiar with the contact match code, so it might make sense to ping someone to review those changes.
jcarius-nv
added
the
IsaacLab
Codecov Report❌ Patch coverage is 📢 Thoughts on this report? Let us know! |
jcarius-nv
left a comment
jcarius-nv
left a comment
There was a problem hiding this comment.
Standards and spec findings from an exact-head review of 5dac41a7. I did not duplicate the existing open ContactMatcher.reset() docstring thread.
|
Fit question for @nvtw: this PR adds |
Use a final mask entry for global world -1 across the public solver reset APIs and coupled proxy state. Keep local-only masks compatible during deprecation and document how world-start offsets represent split global entity ranges.
GitHub Actions
Error
GitHub Actions
Error
had my agent rebase, can you check this matches your idea @eric-heiden ? |
There was a problem hiding this comment.
Actionable comments posted: 2
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
newton/_src/solvers/implicit_mpm/solver_implicit_mpm.py (1)
1566-1572: 📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick winDocstring overstates status clearing for masked resets.
Rebuild status is now cleared only when at least one mask entry is selected (see the masked kernel at lines 141-151); an all-false mask leaves both status buffers untouched, which the new regression asserts. Please align the wording.
📝 Proposed doc tweak
- every warm-start field. Sparse-grid rebuild status is always cleared at - a valid reset boundary, and the previous-collider-pose cache is + every warm-start field. Sparse-grid rebuild status is cleared at a + valid reset boundary when at least one world is selected, and the + previous-collider-pose cache is refreshed from ``state``.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@newton/_src/solvers/implicit_mpm/solver_implicit_mpm.py` around lines 1566 - 1572, Update the reset-operation docstring near the masked reset behavior to state that sparse-grid rebuild status is cleared only when the mask selects at least one entry; explicitly note that an all-false mask leaves both status buffers unchanged. Keep the existing descriptions of full resets and previous-collider-pose cache refresh intact.
🧹 Nitpick comments (3)
newton/_src/solvers/vbd/rigid_vbd_kernels.py (1)
99-110: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winConsolidate with the shared canonical mask-selection helper.
This reimplements the "reset mask contract" selection logic already centralized in
newton/_src/core/reset.py::reset_world_selected, but with slightly different semantics: it treats any negativeworldas the global slot, whereas the shared helper strictly requiresworld == -1.implicit_mpm_solver_kernels.pyalready calls the sharedwp.funcdirectly from another solver module, so the same pattern works here. Delegating avoids two independently-maintained copies of this contract drifting apart.♻️ Proposed consolidation
+from ...core.reset import reset_world_selected + + `@wp.func` def _reset_world_selected( world: int, world_mask: wp.array[wp.bool], reset_all: bool, world_count: int, ): """Query a public reset mask whose final entry selects global entities.""" if reset_all: return True - if world < 0: - world = world_count - return world_mask[world] + return reset_world_selected(world, world_mask, world_count)🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@newton/_src/solvers/vbd/rigid_vbd_kernels.py` around lines 99 - 110, Replace the local selection logic in _reset_world_selected with a call to the canonical reset_world_selected helper from newton._src.core.reset, preserving the existing arguments and return behavior while relying on its strict world == -1 contract; remove the duplicated reset_all, negative-world, and mask-index handling.newton/_src/solvers/coupled/solver_coupled_admm.py (1)
1820-1857: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winConsider a single group-enumeration helper.
The full ADMM group set is now spelled out in
_reset_coupling_state,_reset_admm_history, and the proximal-marking paths. A new group type must be added to every list or masked reset silently stops clearing it. A shared_all_admm_groups()(or per-kind accessors returning the endpoint/world pairs) would keep these in lockstep.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@newton/_src/solvers/coupled/solver_coupled_admm.py` around lines 1820 - 1857, Add a shared ADMM group-enumeration helper, such as _all_admm_groups(), that returns every group with its endpoint IDs and world views, then reuse it in _reset_admm_history, _reset_coupling_state, and proximal-marking paths. Remove the duplicated group lists and preserve each caller’s existing filtering or masking behavior so newly added group types are handled consistently.newton/tests/test_coupled_solver.py (1)
870-885: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winMove the legacy-mask and empty-model assertions into their own tests.
These lines run after the
for mask_values, selected_rowsloop and silently reusemodel,parent,coupled, andentryfrom the last iteration (the all-false subTest), then rebindmodelfor an unrelated fixture. It works, but the coupling to loop-variable leakage makes failures hard to attribute; two small test methods (..._expands_legacy_mask,..._rejects_oversized_mask_on_empty_model) would be self-contained.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@newton/tests/test_coupled_solver.py` around lines 870 - 885, The legacy-mask deprecation assertions and empty-model oversized-mask assertion currently depend on variables left by the preceding mask loop. Move the legacy-mask scenario into a self-contained test named for expanding the legacy mask, and move the empty-model ValueError scenario into a separate test named for rejecting an oversized mask; create each test’s own model, parent, coupled solver, and entry fixtures while preserving the existing assertions.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@docs/concepts/coupling.rst`:
- Around line 211-212: Revise the reset-contract wording in the paragraph
beginning “Masked resets preserve” to explicitly include global entities:
describe matching history selected by the mask, including the optional W + 1
global-entity entry, rather than only selected-world history.
In `@newton/_src/solvers/coupled/solver_coupled.py`:
- Around line 146-157: Update _zero_reset_view_rows_kernel to zero values using
a dtype-compatible default scalar constructor instead of values.dtype(0.0),
ensuring reset and forced input attributes such as state.body_f,
state.particle_f, body_parent_f, and scalar transform-backed fields work across
their supported array dtypes; alternatively, restrict the kernel to dtypes that
safely support the current constructor.
---
Outside diff comments:
In `@newton/_src/solvers/implicit_mpm/solver_implicit_mpm.py`:
- Around line 1566-1572: Update the reset-operation docstring near the masked
reset behavior to state that sparse-grid rebuild status is cleared only when the
mask selects at least one entry; explicitly note that an all-false mask leaves
both status buffers unchanged. Keep the existing descriptions of full resets and
previous-collider-pose cache refresh intact.
---
Nitpick comments:
In `@newton/_src/solvers/coupled/solver_coupled_admm.py`:
- Around line 1820-1857: Add a shared ADMM group-enumeration helper, such as
_all_admm_groups(), that returns every group with its endpoint IDs and world
views, then reuse it in _reset_admm_history, _reset_coupling_state, and
proximal-marking paths. Remove the duplicated group lists and preserve each
caller’s existing filtering or masking behavior so newly added group types are
handled consistently.
In `@newton/_src/solvers/vbd/rigid_vbd_kernels.py`:
- Around line 99-110: Replace the local selection logic in _reset_world_selected
with a call to the canonical reset_world_selected helper from
newton._src.core.reset, preserving the existing arguments and return behavior
while relying on its strict world == -1 contract; remove the duplicated
reset_all, negative-world, and mask-index handling.
In `@newton/tests/test_coupled_solver.py`:
- Around line 870-885: The legacy-mask deprecation assertions and empty-model
oversized-mask assertion currently depend on variables left by the preceding
mask loop. Move the legacy-mask scenario into a self-contained test named for
expanding the legacy mask, and move the empty-model ValueError scenario into a
separate test named for rejecting an oversized mask; create each test’s own
model, parent, coupled solver, and entry fixtures while preserving the existing
assertions.
🪄 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 Plus
Run ID: 6e48bbdd-f5e9-4731-a778-2149503a3aa8
📒 Files selected for processing (30)
CHANGELOG.mddocs/concepts/coupling.rstdocs/concepts/worlds.rstnewton/_src/core/reset.pynewton/_src/geometry/contact_match.pynewton/_src/sim/collide.pynewton/_src/solvers/coupled/admm_utils.pynewton/_src/solvers/coupled/interface.pynewton/_src/solvers/coupled/solver_coupled.pynewton/_src/solvers/coupled/solver_coupled_admm.pynewton/_src/solvers/coupled/solver_coupled_proxy.pynewton/_src/solvers/implicit_mpm/implicit_mpm_solver_kernels.pynewton/_src/solvers/implicit_mpm/solver_implicit_mpm.pynewton/_src/solvers/kamino/_src/solvers/dvi/solver.pynewton/_src/solvers/kamino/_src/solvers/padmm/solver.pynewton/_src/solvers/kamino/solver_kamino.pynewton/_src/solvers/kamino/tests/test_solvers_dvi.pynewton/_src/solvers/mujoco/kernels.pynewton/_src/solvers/mujoco/solver_mujoco.pynewton/_src/solvers/solver.pynewton/_src/solvers/vbd/rigid_vbd_kernels.pynewton/_src/solvers/vbd/solver_vbd.pynewton/tests/test_admm_coupled_solver.pynewton/tests/test_contact_matching.pynewton/tests/test_coupled_solver.pynewton/tests/test_custom_solver.pynewton/tests/test_implicit_mpm_multiworld_sparse.pynewton/tests/test_mujoco_reset.pynewton/tests/test_solver_kamino_dvi.pynewton/tests/test_solver_vbd.py
🚧 Files skipped from review as they are similar to previous changes (1)
- CHANGELOG.md
|
🔵 Nice work making the reset mask handle global entities in world One thing to sort out before merge: the new |
Render the legacy mask deprecation in every solver reset API and make masked row clearing compile for scalar and transform state arrays.
GitHub Actions
Error
GitHub Actions
Error
GitHub Actions
Inactive
GitHub Actions
Inactive
GitHub Actions
Inactive
GitHub Actions
Inactive
GitHub Actions
Error
GitHub Actions
Error
GitHub Actions
Inactive
6 participants
Description
Builds on the canonical reset-mask contract introduced by #3726, now merged into
main: a boolean mask with shape(world_count + 1,), with local-world slots followed by the global-world-1slot. Established public solver reset entry points still accept the deprecated(world_count,)form, warn once, and normalize it immediately; internal reset paths and child-solver dispatch use only the canonical form.SolverCoupled.reset()previously forwarded a mask to child solvers while its surrounding distribution, synchronization, reconciliation, and history clearing could still touch complete entry states. Some child solvers also cleared private history globally, so resetting one world could perturb an unselected world.This PR:
Generic reconciliation covers registered
STATEattributes with explicit coupled-entry ownership. Custom frequencies require coupler-specific ownership logic. Native MuJoCo CPU synchronization remains host-based, and Kamino receives a local-world slice only at its private backend adapter.Checklist
CHANGELOG.mdhas been updatedValidation
uvx pre-commit run -a— passedDeprecationWarnings treated as errors — 96 passedcuda:0), and RTX 4090 (cuda:1)Bug reproduction
(world_count + 1,)mask.Summary by CodeRabbit