Support multi-articulation selections #1100
Conversation
📝 WalkthroughWalkthroughArticulationView was refactored to be world-aware using a new FrequencyLayout abstraction (per-attribute offsets/strides/contiguity); masking and kernels now use 2D world×articulation mappings. A multi-world selection example and comprehensive squeeze-axis tests were added. Changes
Sequence Diagram(s)sequenceDiagram
participant User
participant ArtView as ArticulationView
participant Freq as FrequencyLayout
participant Kernel as MaskingKernel
participant Model
User->>ArtView: construct(selection, squeeze_axes)
ArtView->>ArtView: group articulations by world
ArtView->>Freq: create per-attribute FrequencyLayout(s)
note right of Freq: rgba(200,200,255,0.5) offset, outer/inner stride,\nvalue_count, indices/slice, contiguity
User->>ArtView: get_model_articulation_mask(optional_mask)
ArtView->>ArtView: build 2D view_mask & view_to_model_map
ArtView->>Kernel: invoke(view_mask, view_to_model_map)
Kernel->>Model: apply per-world, per-articulation masking
Kernel-->>ArtView: masked result
User->>ArtView: get_attribute(name)
ArtView->>Freq: lookup layout for attribute
Freq-->>ArtView: provide 2D indices/slice
ArtView->>ArtView: gather & apply squeeze_axes
ArtView-->>User: attribute array
Estimated code review effort🎯 4 (Complex) | ⏱️ ~45 minutes Suggested reviewers
🚥 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
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 |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 3
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
newton/_src/utils/selection.py (1)
984-1003: Fix kernel launch dimensions in get_model_articulation_mask.Line 1000 uses
dim=mask.size, butmaskis 1D with shape(world_count,)while the kernelset_model_articulation_mask_kernelexpects 2D indexing(world, arti). The launch should usedim=(self.world_count, self.count_per_world)to match the kernel signature.Apply this diff:
assert mask.shape == (self.world_count,) articulation_mask = wp.zeros(self.model.articulation_count, dtype=bool, device=self.device) wp.launch( set_model_articulation_mask_kernel, - dim=mask.size, + dim=(self.world_count, self.count_per_world), inputs=[mask, self.articulation_ids, articulation_mask], ) return articulation_mask
🧹 Nitpick comments (2)
newton/examples/selection/example_selection_multiple.py (2)
149-187: Remove unused noqa directives.Lines 150, 227, and 254 have
# noqa: PLC0415comments that are flagged as unused by Ruff. These can be removed for cleaner code.Apply this diff:
if USE_TORCH: - import torch # noqa: PLC0415 + import torch # default ant root states self.default_root_transforms = wp.to_torch(self.ants.get_root_transforms(self.model)).clone()Similar changes for lines 227 and 254.
280-286: Mark unused lambda parameter with underscore.The
qdparameter in the test lambda is unused. Consider using_to indicate it's intentionally ignored.Apply this diff:
newton.examples.test_body_state( self.model, self.state_0, "all bodies are above the ground", - lambda q, qd: q[2] > 0.01, + lambda q, _: q[2] > 0.01, )
📜 Review details
Configuration used: Path: .coderabbit.yml
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (3)
newton/_src/utils/selection.py(12 hunks)newton/examples/selection/example_selection_multiple.py(1 hunks)newton/tests/test_examples.py(1 hunks)
🧰 Additional context used
🧠 Learnings (3)
📚 Learning: 2025-08-27T19:05:44.697Z
Learnt from: Milad-Rakhsha-NV
Repo: newton-physics/newton PR: 535
File: newton/tests/test_examples.py:320-414
Timestamp: 2025-08-27T19:05:44.697Z
Learning: In newton/examples/__init__.py, the robot policy example is registered with the key "robot_policy" (not "robot.example_robot_policy"), so tests should reference it as name="robot_policy".
Applied to files:
newton/tests/test_examples.py
📚 Learning: 2025-09-22T21:08:31.901Z
Learnt from: dylanturpin
Repo: newton-physics/newton PR: 806
File: newton/examples/ik/example_ik_franka.py:121-123
Timestamp: 2025-09-22T21:08:31.901Z
Learning: In the newton physics framework, when creating warp arrays for IK solver joint variables using wp.array(self.model.joint_q, shape=(1, coord_count)), the resulting array acts as a reference/pointer to the original model's joint coordinates, so updates from the IK solver automatically reflect in the model's joint_q buffer used for rendering.
Applied to files:
newton/_src/utils/selection.py
📚 Learning: 2025-09-22T21:08:31.901Z
Learnt from: dylanturpin
Repo: newton-physics/newton PR: 806
File: newton/examples/ik/example_ik_franka.py:121-123
Timestamp: 2025-09-22T21:08:31.901Z
Learning: In the newton physics framework, when creating warp arrays for IK solver joint variables using wp.array(self.model.joint_q, shape=(1, coord_count)), the resulting array acts as a reference/pointer to the original model's joint coordinates, so updates from the IK solver automatically reflect in the model's joint_q buffer used for rendering, despite the general warp documentation suggesting copies are made by default.
Applied to files:
newton/_src/utils/selection.py
🧬 Code graph analysis (2)
newton/examples/selection/example_selection_multiple.py (3)
newton/_src/utils/selection.py (11)
ArticulationView(196-1015)get_root_transforms(818-836)get_root_velocities(853-872)get_attribute(785-796)get_dof_forces(958-968)set_dof_forces(970-979)set_root_transforms(838-851)set_root_velocities(874-886)set_dof_positions(924-933)set_dof_velocities(947-956)eval_fk(1005-1015)newton/_src/sim/builder.py (5)
ModelBuilder(70-5272)add_mjcf(1185-1275)collapse_fixed_joints(2535-2829)add_ground_plane(3036-3056)replicate(939-968)newton/examples/__init__.py (4)
get_asset(34-35)test_body_state(38-113)create_parser(271-328)init(331-374)
newton/_src/utils/selection.py (3)
newton/_src/sim/joints.py (1)
JointType(20-44)newton/_src/sim/state.py (3)
joint_dof_count(108-112)joint_coord_count(101-105)State(21-112)newton/_src/sim/model.py (3)
ModelAttributeFrequency(49-70)Model(96-798)get_attribute_frequency(782-798)
🪛 GitHub Actions: Pull Request
newton/_src/utils/selection.py
[error] 727-727: SyntaxError: invalid syntax. The expression 'attrib = attrib[*leading_slices, _slice, *trailing_slices]' uses Python syntax not supported in this environment.
🪛 Ruff (0.14.5)
newton/examples/selection/example_selection_multiple.py
150-150: Unused noqa directive (non-enabled: PLC0415)
Remove unused noqa directive
(RUF100)
227-227: Unused noqa directive (non-enabled: PLC0415)
Remove unused noqa directive
(RUF100)
254-254: Unused noqa directive (non-enabled: PLC0415)
Remove unused noqa directive
(RUF100)
285-285: Unused lambda argument: qd
(ARG005)
newton/_src/utils/selection.py
168-168: Avoid specifying long messages outside the exception class
(TRY003)
245-245: Avoid specifying long messages outside the exception class
(TRY003)
269-269: Avoid specifying long messages outside the exception class
(TRY003)
272-272: Avoid specifying long messages outside the exception class
(TRY003)
367-367: Avoid specifying long messages outside the exception class
(TRY003)
403-403: Avoid specifying long messages outside the exception class
(TRY003)
434-434: Avoid specifying long messages outside the exception class
(TRY003)
674-674: Avoid specifying long messages outside the exception class
(TRY003)
720-720: Avoid specifying long messages outside the exception class
(TRY003)
727-727: Cannot use star expression in index on Python 3.10 (syntax was added in Python 3.11)
(invalid-syntax)
727-727: Cannot use star expression in index on Python 3.10 (syntax was added in Python 3.11)
(invalid-syntax)
734-734: Cannot use star expression in index on Python 3.10 (syntax was added in Python 3.11)
(invalid-syntax)
734-734: Cannot use star expression in index on Python 3.10 (syntax was added in Python 3.11)
(invalid-syntax)
⏰ Context from checks skipped due to timeout of 900000ms. You can increase the timeout in your CodeRabbit configuration to a maximum of 15 minutes (900000ms). (2)
- GitHub Check: Run GPU Benchmarks / Run GPU Benchmarks on AWS EC2
- GitHub Check: Run GPU Tests / Run GPU Unit Tests on AWS EC2
🔇 Additional comments (5)
newton/tests/test_examples.py (1)
489-496: LGTM!Test registration is consistent with other selection examples and properly configured.
newton/_src/utils/selection.py (4)
28-38: LGTM!Kernel correctly updated to handle 2D (world, articulation) indexing for multi-articulation support.
126-150: LGTM!
FrequencyLayoutprovides a clean abstraction for articulation attribute layout metadata with clear separation of contiguous vs. indexed cases.
153-194: LGTM!Helper functions provide clear, well-scoped utilities for world-aware articulation grouping and layout computations.
216-659: Excellent refactor for multi-articulation support.The initialization logic properly validates uniform articulation topology across worlds, computes per-world and per-articulation offsets/strides, and constructs
FrequencyLayoutentries for efficient attribute access. Thesqueeze_axesparameter provides flexible dimensionality control.
There was a problem hiding this comment.
Actionable comments posted: 2
🧹 Nitpick comments (1)
newton/examples/selection/example_selection_multiple.py (1)
150-150: Remove unused noqa directives.The
# noqa: PLC0415comments on lines 150, 227, and 254 are unnecessary because the PLC0415 rule (import-outside-toplevel) is not enabled in your linter configuration.Apply this diff to clean up the noqa directives:
if USE_TORCH: - import torch # noqa: PLC0415 + import torchApply the same change to lines 227 and 254.
Also applies to: 227-227, 254-254
📜 Review details
Configuration used: Path: .coderabbit.yml
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (2)
newton/_src/utils/selection.py(12 hunks)newton/examples/selection/example_selection_multiple.py(1 hunks)
🧰 Additional context used
🧠 Learnings (5)
📚 Learning: 2025-08-25T20:20:35.899Z
Learnt from: dylanturpin
Repo: newton-physics/newton PR: 635
File: newton/_src/sim/builder.py:3191-3195
Timestamp: 2025-08-25T20:20:35.899Z
Learning: When reviewing PRs that update Python version requirements (e.g., dropping Python 3.9 support), focus on syntax modernization rather than suggesting behavioral changes. Adding explicit `strict=False` to zip() calls is appropriate for Python 3.10+ compatibility and maintains existing behavior while satisfying linters.
Applied to files:
newton/_src/utils/selection.py
📚 Learning: 2025-09-22T21:08:31.901Z
Learnt from: dylanturpin
Repo: newton-physics/newton PR: 806
File: newton/examples/ik/example_ik_franka.py:121-123
Timestamp: 2025-09-22T21:08:31.901Z
Learning: In the newton physics framework, when creating warp arrays for IK solver joint variables using wp.array(self.model.joint_q, shape=(1, coord_count)), the resulting array acts as a reference/pointer to the original model's joint coordinates, so updates from the IK solver automatically reflect in the model's joint_q buffer used for rendering.
Applied to files:
newton/_src/utils/selection.py
📚 Learning: 2025-09-22T21:08:31.901Z
Learnt from: dylanturpin
Repo: newton-physics/newton PR: 806
File: newton/examples/ik/example_ik_franka.py:121-123
Timestamp: 2025-09-22T21:08:31.901Z
Learning: In the newton physics framework, when creating warp arrays for IK solver joint variables using wp.array(self.model.joint_q, shape=(1, coord_count)), the resulting array acts as a reference/pointer to the original model's joint coordinates, so updates from the IK solver automatically reflect in the model's joint_q buffer used for rendering, despite the general warp documentation suggesting copies are made by default.
Applied to files:
newton/_src/utils/selection.py
📚 Learning: 2025-09-22T21:03:39.624Z
Learnt from: dylanturpin
Repo: newton-physics/newton PR: 806
File: newton/_src/sim/ik/ik_lbfgs_optimizer.py:739-752
Timestamp: 2025-09-22T21:03:39.624Z
Learning: The L-BFGS optimizer in newton/_src/sim/ik/ik_lbfgs_optimizer.py currently intentionally only supports additive updates (assuming n_coords == n_dofs). Velocity space integration for joints with mismatched coordinate/DOF dimensions (like free/ball joints) is planned for future work and should not be flagged as an issue in current reviews.
Applied to files:
newton/examples/selection/example_selection_multiple.py
📚 Learning: 2025-09-22T21:03:18.367Z
Learnt from: dylanturpin
Repo: newton-physics/newton PR: 806
File: newton/_src/sim/ik/ik_lbfgs_optimizer.py:638-648
Timestamp: 2025-09-22T21:03:18.367Z
Learning: The L-BFGS optimizer in newton/_src/sim/ik/ik_lbfgs_optimizer.py currently only supports additive updates, meaning it assumes n_coords == n_dofs. Velocity space integration for joints with mismatched coordinate/DOF dimensions (like free/ball joints) is planned for future work.
Applied to files:
newton/examples/selection/example_selection_multiple.py
🧬 Code graph analysis (2)
newton/_src/utils/selection.py (3)
newton/_src/sim/joints.py (1)
JointType(20-44)newton/_src/sim/state.py (3)
joint_dof_count(108-112)joint_coord_count(101-105)State(21-112)newton/_src/sim/model.py (3)
ModelAttributeFrequency(49-70)Model(96-798)get_attribute_frequency(782-798)
newton/examples/selection/example_selection_multiple.py (3)
newton/_src/utils/selection.py (12)
ArticulationView(203-1053)get_root_transforms(856-874)get_root_velocities(891-910)get_attribute(823-834)get_dof_velocities(973-983)get_dof_forces(996-1006)set_dof_forces(1008-1017)set_root_transforms(876-889)set_root_velocities(912-924)set_dof_positions(962-971)set_dof_velocities(985-994)eval_fk(1043-1053)newton/_src/sim/builder.py (5)
ModelBuilder(70-5272)add_mjcf(1185-1275)collapse_fixed_joints(2535-2829)add_ground_plane(3036-3056)replicate(939-968)newton/examples/__init__.py (4)
get_asset(34-35)test_body_state(38-113)create_parser(271-328)init(331-374)
🪛 GitHub Actions: Pull Request
newton/_src/utils/selection.py
[error] 276-276: KeyError: 'No matching articulations' raised during ArticulationView initialization.
🪛 GitHub Actions: Pull Request - AWS GPU
newton/_src/utils/selection.py
[error] 276-276: KeyError: 'No matching articulations' raised during ArticulationView initialization (No matching articulations).
🪛 Ruff (0.14.5)
newton/_src/utils/selection.py
175-175: Avoid specifying long messages outside the exception class
(TRY003)
252-252: Avoid specifying long messages outside the exception class
(TRY003)
276-276: Avoid specifying long messages outside the exception class
(TRY003)
279-279: Avoid specifying long messages outside the exception class
(TRY003)
374-374: Avoid specifying long messages outside the exception class
(TRY003)
410-410: Avoid specifying long messages outside the exception class
(TRY003)
445-445: Avoid specifying long messages outside the exception class
(TRY003)
704-706: Avoid specifying long messages outside the exception class
(TRY003)
756-756: Avoid specifying long messages outside the exception class
(TRY003)
newton/examples/selection/example_selection_multiple.py
150-150: Unused noqa directive (non-enabled: PLC0415)
Remove unused noqa directive
(RUF100)
227-227: Unused noqa directive (non-enabled: PLC0415)
Remove unused noqa directive
(RUF100)
254-254: Unused noqa directive (non-enabled: PLC0415)
Remove unused noqa directive
(RUF100)
287-287: Unused lambda argument: qd
(ARG005)
There was a problem hiding this comment.
Actionable comments posted: 1
🧹 Nitpick comments (1)
newton/_src/utils/selection.py (1)
294-482: Stride computation logic is thorough and correct.The systematic computation of offsets and strides for joints, DOFs, coordinates, links, and shapes across worlds properly validates homogeneity and builds the foundation for 2D world-aware attribute access.
Consider enhancing error messages at lines 388, 428, and 469 to include specific details about which counts or strides differ, e.g.:
- raise ValueError("Articulations are not identical") + raise ValueError(f"Articulations are not identical: joints={joint_counts[0]} vs others, dofs={joint_dof_counts[0]}, etc.")
📜 Review details
Configuration used: Path: .coderabbit.yml
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (1)
newton/_src/utils/selection.py(12 hunks)
🧰 Additional context used
🧠 Learnings (4)
📚 Learning: 2025-08-25T20:20:35.899Z
Learnt from: dylanturpin
Repo: newton-physics/newton PR: 635
File: newton/_src/sim/builder.py:3191-3195
Timestamp: 2025-08-25T20:20:35.899Z
Learning: When reviewing PRs that update Python version requirements (e.g., dropping Python 3.9 support), focus on syntax modernization rather than suggesting behavioral changes. Adding explicit `strict=False` to zip() calls is appropriate for Python 3.10+ compatibility and maintains existing behavior while satisfying linters.
Applied to files:
newton/_src/utils/selection.py
📚 Learning: 2025-09-10T13:03:58.591Z
Learnt from: gdaviet
Repo: newton-physics/newton PR: 750
File: newton/_src/solvers/implicit_mpm/solve_rheology.py:799-804
Timestamp: 2025-09-10T13:03:58.591Z
Learning: In Warp's tiled launch API, the launch dimension should include the tile size as the last dimension, matching the CUDA block dimension. For tiled operations, use dim=(num_blocks, tile_size) where tile_size equals block_dim, not just dim=num_blocks like regular kernel launches.
Applied to files:
newton/_src/utils/selection.py
📚 Learning: 2025-09-22T21:08:31.901Z
Learnt from: dylanturpin
Repo: newton-physics/newton PR: 806
File: newton/examples/ik/example_ik_franka.py:121-123
Timestamp: 2025-09-22T21:08:31.901Z
Learning: In the newton physics framework, when creating warp arrays for IK solver joint variables using wp.array(self.model.joint_q, shape=(1, coord_count)), the resulting array acts as a reference/pointer to the original model's joint coordinates, so updates from the IK solver automatically reflect in the model's joint_q buffer used for rendering.
Applied to files:
newton/_src/utils/selection.py
📚 Learning: 2025-09-22T21:08:31.901Z
Learnt from: dylanturpin
Repo: newton-physics/newton PR: 806
File: newton/examples/ik/example_ik_franka.py:121-123
Timestamp: 2025-09-22T21:08:31.901Z
Learning: In the newton physics framework, when creating warp arrays for IK solver joint variables using wp.array(self.model.joint_q, shape=(1, coord_count)), the resulting array acts as a reference/pointer to the original model's joint coordinates, so updates from the IK solver automatically reflect in the model's joint_q buffer used for rendering, despite the general warp documentation suggesting copies are made by default.
Applied to files:
newton/_src/utils/selection.py
🔇 Additional comments (4)
newton/_src/utils/selection.py (4)
128-158: LGTM! FrequencyLayout abstraction is well-designed.The new FrequencyLayout class cleanly encapsulates per-attribute layout metadata and correctly distinguishes between contiguous (slice-based) and non-contiguous (index-based) access patterns.
160-201: Helper functions are correctly implemented.The utility functions properly handle edge cases (empty lists, single elements, out-of-range world IDs) and support the new world-aware layout logic.
263-292: World-aware articulation grouping logic is sound.The code correctly groups articulations by world, handles global articulations as a special case, and validates that all worlds have the same articulation count.
718-799: Attribute array construction with layout and squeezing is well-implemented.The complex logic to compute shapes, strides, apply squeezing, and create sliced or indexed arrays correctly handles all combinations of singleton axes and maintains Python 3.10 compatibility.
| include_joint_types: list[int] | None = None, | ||
| exclude_joint_types: list[int] | None = None, | ||
| verbose: bool | None = None, | ||
| squeeze_axes: bool | tuple[int] = (1,), |
There was a problem hiding this comment.
Curious to hear people's opinions on what the default should be. The PR description has notes about the benefits and drawbacks of different options.
There was a problem hiding this comment.
Personally I'm all for squeezing - I think it's reasonable to expect users to figure out the shape based on the contents of the scene. We can document this clearly.
On the other hand, I wouldn't be too surprised if IsaacLab and RL wrappers in general want to expose this in a generic way, and I think at that point squeezing is not really an option anymore? We should for sure check with them to see if there is even a decision to be made here.
There was a problem hiding this comment.
Ok I think I'm reversing my opinion - after thinking some more about it it's way more consistent to keep the singleton axes and do this the same way all the time. Users can still reshape or just pass a slice, right?
One thing that I'm coming back to as well is the question of whether we should always have an explicit world dimension here. I can imagine once you're having multiple articulations per world, maybe you don't care too much about worlds anymore anyway? Was that a hard requirement from IsaacLab or did it just make sense because of how the API works?
There was a problem hiding this comment.
it's way more consistent to keep the singleton axes
Yep, consistency is a good thing for API :)
whether we should always have an explicit world dimension here
I think it is still useful, since in RL we often reset some subset of worlds. We currently use a mask over world indices to control which worlds should be reset. But we could also mask over individual articulations. Maybe @AntoineRichard can chime in here?
Codecov Report❌ Patch coverage is
📢 Thoughts on this report? Let us know! |
| test_options={"num-frames": 100}, | ||
| test_options_cpu={"num-frames": 10}, | ||
| use_viewer=True, | ||
| ) |
There was a problem hiding this comment.
would be great to have explicit test coverage for some of the special cases as well, not just the example.
There was a problem hiding this comment.
Added tests for selections with single/multiple worlds and single/multiple articulations per world, with a focus on the new functionality introduced in this PR (world, arti, and value axes, squeezing options, etc.)
| for world_id in range(world_count): | ||
| count = len(articulation_ids[world_id]) | ||
| counts_per_world[world_id] += count | ||
| articulation_count += count |
There was a problem hiding this comment.
I'm not sure if I just missed this, but where are you handling global articulations in case there are both global and per-world articulations?
There was a problem hiding this comment.
Good point, I'm still thinking about that. The current code assumes that either
- There are matching articulations per-world, but no global matches.
- There are matching global articulations, but no matches per world.
This covers the most common use cases, where we either have multiple replicated worlds or just a single scene with everything global.
If there are both, it's a pickle. Chances are that we won't be able to construct a single attribute array, because the strides will be funky. Also, it's not clear how to pack the global stuff with the per-world stuff in a single array. We would need to add an extra world index for the globals either at the beginning or at the end, neither of which feels quite right. One alternative would be to have different getters and setters for the globals, which is also not great. Or the getters and setters could take a flag. Or the getters could return two arrays, one for globals and one for per-world data, and the setters could take two arrays as well. But that would complicate the API for the sake of a use case that seems pretty rare to me.
My current inclination is to raise an exception, though that's not in the code yet. The guidance for users would be to use different key patterns for the global stuff, to avoid selections that span both global and per-world objects.
Maybe the exception is the way to go for now, and we can revisit later?
There was a problem hiding this comment.
I think these assumptions are fair for the moment, I think it's ok to raise an exception. A typical use-case would probably special-case the global articulations anyway because their likely part of some world/scene and not the actual robot you're focusing on.
There was a problem hiding this comment.
Done, raising an exception when both global and per-world articulations are found
newton/_src/utils/selection.py
Outdated
| self.is_fixed_base = root_joint_type == JointType.FIXED | ||
| # floating base means that all linear and angular degrees of freedom are unlocked at the root | ||
| # (though there might be constraints like distance) | ||
| self.is_floating_base = root_joint_type in (JointType.FREE, JointType.DISTANCE) |
There was a problem hiding this comment.
couldn't this be basically any joint except fixed here?
There was a problem hiding this comment.
This is a little tricky. Currently, I define "floating" as having all linear and angular DOFs unlocked, which matches FREE and DISTANCE joints.
The is_floating_base flag is used when getting and setting root transforms and velocities. For FREE and DISTANCE joints, we do that through joint_q and joint_qd. For other types of joints, we use joint_X_p. This only matters for the "helper" API like get_root_transforms() and get_dof_positions(). Users who explicitly get attributes by name are in full control and take the responsibility for reading and writing the appropriate attributes.
The most common articulation root joints are FREE and FIXED, so the current setup works fine.
With FIXED joints, get_root_transforms() returns joint_X_p and get_dof_positions() returns the internal articulation joints in joint_q.
With FREE (and DISTANCE) joints, get_root_transforms() returns the root joint slice of joint_q and get_dof_positions() returns the internal joint slice of joint_q so that there's no overlap.
With other types of joints (like REVOLUTE or PRISMATIC), get_root_transforms() returns the "fixed" transform in joint_X_p, and get_dof_positions() returns all the DOFs in joint_q, including the root joint DOFs.
I struggled with this quite a bit at the beginning, and I'm still not entirely happy with the solution. But this works fine for Isaac Lab, so I'm happy enough for now. Maybe we'll need to revisit that for maximum coordinate solvers, but it's not clear to me yet. Always open to suggestions for improving the API :)
There was a problem hiding this comment.
That thinking and explanation makes total sense to me. I'm wondering if we should promote this floating_base property/hint to the main newton API as I think it could be useful to unify the handling of this.
There was a problem hiding this comment.
I always thought fixed base means the same as not floating base? So any joint except free joints. Is there some code anywhere that relies on this narrower definition of is_fixed_base? So far it was just used for printing debug info.
There was a problem hiding this comment.
To me, having this three-way distinction makes sense. "Fixed base" implies that the articulation root is immovable and "floating base" implies that it is free to move along all degrees of freedom. Articulations with prismatic or revolute joints at the root are somewhere in between - "partially fixed"?
I can see how this might be surprising to some users though, so maybe we could get rid of these is_fixed_base and is_floating_base flags and let users query the root joint type instead. I've exposed ArticulationView.root_joint_type for this now.
We can revisit this, but maybe after this PR is merged? This is not a new thing in this PR. It's possible that Isaac Lab is relying on these is_fixed_base and is_floating_base flags, so we should consult with them before removing or modifying them.
|
@nvlukasz do you have the time to resurrect this, or should we pick this up from you and finish? Would be great to get this in. I think on the technical side it's good, maybe it needs a few more tests. What do you think? |
newton/_src/utils/selection.py
Outdated
| self.body_names = [] | ||
| self.shape_names = [] | ||
| self.body_shapes = [] | ||
| self.link_names = self.body_names # alias, FIXME: pick one? |
There was a problem hiding this comment.
We call these "keys" rather than "names" in ModelBuilder and Model, maybe they should be renamed here as well?
There was a problem hiding this comment.
I would also stick to "body" instead of links, we use this term also in Model.body_key, etc.
There was a problem hiding this comment.
About "key" vs "name": the key is a full path, like /World/envs/env42/ant/torso, but we only want the names here (i.e., torso), since it's shared across all selected articulations. So I would keep the term "names". It also mirrors what we do in the ArticulationView from the physics tensor API.
About "body" vs "link", we have model.body_*, but also builder.add_link() for articulations... So both terms are used in the API, but "link" is used specifically when constructing articulations. So "link" seems to fit with ArticulationView. But I have no strong preference here, happy to be swayed either way :)
There was a problem hiding this comment.
I settled on link_names and link_shapes, but provided alias properties body_names and body_shapes to avoid breaking existing code. We could deprecate those in the future.
| self.articulation_ids = wp.array(articulation_ids, dtype=int, device=self.device) | ||
|
|
||
| # TODO: this mask is per world, should we support 2d masks (world, arti)? | ||
| self.full_mask = wp.full(world_count, True, dtype=bool, device=self.device) |
There was a problem hiding this comment.
Does this mean that it is only possible to mask articulation attributes now per world? I wonder if this change from before where the mask was per articulation(?) is still needed. One example is if we select multiple articulations that are all in the same world. So in that case having a 2D mask might make sense.
There was a problem hiding this comment.
I think ideally the setter kernels can take 1- or 2-dimensional masks to give the flexibility to set per-world, or per-world and per-articulation attributes.
There was a problem hiding this comment.
We can do that. So far, the main use case with RL is to reset individual envs, but we can support 2d articulation masks too.
There was a problem hiding this comment.
Hmmm, ran into a complication here. We need to be careful when squeezing attribute shapes (e.g., removing the world and/or articulation dimensions) - masks would need to be squeezed in the same way. There are a couple of different ways to deal with this, but things get complicated.
I'm half-tempted to get rid of the squeezing functionality because of this, but at the same time I think it's a nice feature.
Adding tests :) |
nvlukasz
force-pushed
the
selection-multiple-articulations
branch
from
206edfc to
1191b7e
Compare
There was a problem hiding this comment.
Actionable comments posted: 2
🤖 Fix all issues with AI agents
In `@newton/_src/utils/selection.py`:
- Around line 459-467: The current stride check only ensures all worlds have
identical stride lists but not that each world's list is internally uniform;
update the validation around the existing block that inspects
inner_joint_strides, inner_joint_dof_strides, inner_joint_coord_strides,
inner_link_strides, and inner_shape_strides to also verify each sublist is
uniform (e.g., for each sublist in inner_joint_strides: all_equal(sublist) or
len(set(sublist)) == 1), and do the same for the other four stride variables; if
any sublist is non-uniform, raise the existing ValueError("Non-uniform strides
within worlds are not supported") so later code using inner_joint_stride =
inner_joint_strides[0][0] is valid.
- Around line 836-841: Kernels index mask by the first thread component
(mask[i]) which is incorrect when the world axis was squeezed (world_count==1)
because i then iterates over per-world entries; change kernel mask accesses to
use the world/thread's world index (e.g., mask[w] or mask[world_id]) instead of
mask[i]; if kernels only receive a single linear index, compute the world index
explicitly (world = linear_index // count_per_world or world = 0 when
world_count==1) and use that to index mask, leaving mask creation (wp.array(...,
shape=(self.world_count,), ...)) unchanged.
🧹 Nitpick comments (2)
newton/examples/selection/example_selection_multiple.py (1)
149-151: Remove unusednoqadirectives.The
# noqa: PLC0415comments are flagged as unused since PLC0415 is not enabled. These can be safely removed.♻️ Proposed fix
if USE_TORCH: - import torch # noqa: PLC0415 + import torchApply similar changes at lines 227 and 254.
newton/_src/utils/selection.py (1)
185-189: Remove unusedget_world_offsetfunction: This function is defined at lines 185-189 but is never called anywhere in the codebase. Consider removing it to reduce dead code.
📜 Review details
Configuration used: Path: .coderabbit.yml
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (3)
newton/_src/utils/selection.pynewton/examples/selection/example_selection_multiple.pynewton/tests/test_examples.py
🚧 Files skipped from review as they are similar to previous changes (1)
- newton/tests/test_examples.py
🧰 Additional context used
📓 Path-based instructions (5)
newton/examples/**/*.py
📄 CodeRabbit inference engine (AGENTS.md)
newton/examples/**/*.py: Examples must not import fromnewton._src, only from the public Newton API
Examples must follow the format with anExampleclass and must be registered inREADME.mdwith auv runcommand and screenshot
Examples must implement atest_final()method that verifies the simulation state is valid after the example has run
Examples may optionally implement atest_post_step()method that is evaluated after everystep()call
Files:
newton/examples/selection/example_selection_multiple.py
newton/**/*.py
📄 CodeRabbit inference engine (AGENTS.md)
Any user-facing class/function/object added under
_srcmust be exposed via the public Newton API through re-exports
Files:
newton/examples/selection/example_selection_multiple.pynewton/_src/utils/selection.py
**/*.py
📄 CodeRabbit inference engine (AGENTS.md)
**/*.py: Use prefix-first naming for classes:ActuatorPD,ActuatorPID(notPDActuator,PIDActuator)
Use prefix-first naming for methods:add_shape_sphere()(notadd_sphere_shape())
Method names must usesnake_case
Prefer nested classes when self-contained: if a helper type or enum is only meaningful inside one parent class and doesn't need a public identity, define it as a nested class instead of a top-level class/module
Follow PEP 8 for Python code
Use Google-style docstrings with clear, concise explanations of what functions do, their parameters, and return values
Files:
newton/examples/selection/example_selection_multiple.pynewton/_src/utils/selection.py
**/*
📄 CodeRabbit inference engine (AGENTS.md)
CLI arguments must use
kebab-case(e.g.,--use-cuda-graph, not--use_cuda_graph)
Files:
newton/examples/selection/example_selection_multiple.pynewton/_src/utils/selection.py
newton/_src/**/*.py
📄 CodeRabbit inference engine (AGENTS.md)
The
newton/_src/directory is internal library implementation only and must not be imported by user code (examples, documentation, or external users)
Files:
newton/_src/utils/selection.py
🧠 Learnings (13)
📓 Common learnings
Learnt from: nvtw
Repo: newton-physics/newton PR: 1314
File: newton/_src/sim/builder.py:6061-6091
Timestamp: 2026-01-12T10:11:47.988Z
Learning: Repo: newton-physics/newton PR: 1314
File: newton/_src/sim/builder.py (ModelBuilder.finalize)
Learning: The project prefers to keep finalize() validation lightweight; do not add union-find/connected-components validation for loop joints. Current behavior (allow non-articulated joints when the child is in any articulation) is acceptable; avoid overengineering here.
📚 Learning: 2025-09-22T21:03:39.624Z
Learnt from: dylanturpin
Repo: newton-physics/newton PR: 806
File: newton/_src/sim/ik/ik_lbfgs_optimizer.py:739-752
Timestamp: 2025-09-22T21:03:39.624Z
Learning: The L-BFGS optimizer in newton/_src/sim/ik/ik_lbfgs_optimizer.py currently intentionally only supports additive updates (assuming n_coords == n_dofs). Velocity space integration for joints with mismatched coordinate/DOF dimensions (like free/ball joints) is planned for future work and should not be flagged as an issue in current reviews.
Applied to files:
newton/examples/selection/example_selection_multiple.pynewton/_src/utils/selection.py
📚 Learning: 2025-12-12T15:23:30.014Z
Learnt from: vastsoun
Repo: newton-physics/newton PR: 1019
File: newton/_src/solvers/kamino/examples/sim/example_sim_basics_boxes_fourbar.py:149-163
Timestamp: 2025-12-12T15:23:30.014Z
Learning: In Kamino simulation examples under newton/_src/solvers/kamino/examples/sim/, the pattern where logging=True and use_cuda_graph=True creates a mismatch (SimulationLogger constructed but settings.compute_metrics and logger.log() calls guarded by `if not self.use_cuda_graph`) should not be flagged. These examples will be refactored in a future commit, and this pattern is intentionally deferred.
Applied to files:
newton/examples/selection/example_selection_multiple.py
📚 Learning: 2025-09-22T21:03:18.367Z
Learnt from: dylanturpin
Repo: newton-physics/newton PR: 806
File: newton/_src/sim/ik/ik_lbfgs_optimizer.py:638-648
Timestamp: 2025-09-22T21:03:18.367Z
Learning: The L-BFGS optimizer in newton/_src/sim/ik/ik_lbfgs_optimizer.py currently only supports additive updates, meaning it assumes n_coords == n_dofs. Velocity space integration for joints with mismatched coordinate/DOF dimensions (like free/ball joints) is planned for future work.
Applied to files:
newton/examples/selection/example_selection_multiple.py
📚 Learning: 2026-01-12T10:11:47.988Z
Learnt from: nvtw
Repo: newton-physics/newton PR: 1314
File: newton/_src/sim/builder.py:6061-6091
Timestamp: 2026-01-12T10:11:47.988Z
Learning: Repo: newton-physics/newton PR: 1314
File: newton/_src/sim/builder.py (ModelBuilder.finalize)
Learning: The project prefers to keep finalize() validation lightweight; do not add union-find/connected-components validation for loop joints. Current behavior (allow non-articulated joints when the child is in any articulation) is acceptable; avoid overengineering here.
Applied to files:
newton/examples/selection/example_selection_multiple.pynewton/_src/utils/selection.py
📚 Learning: 2025-12-12T08:45:43.428Z
Learnt from: nvtw
Repo: newton-physics/newton PR: 1221
File: newton/examples/example_sdf.py:277-287
Timestamp: 2025-12-12T08:45:43.428Z
Learning: In Newtown (Newton) example code, specifically files under newton/examples, computing contacts once per frame and reusing them across all substeps is an intentional design choice, not a bug. Reviewers should verify that contacts are computed before the substep loop and reused for every substep within the same frame. This pattern reduces redundant work and preserves frame-consistency; do not flag as a regression unless the behavior is changed for correctness or performance reasons.
Applied to files:
newton/examples/selection/example_selection_multiple.py
📚 Learning: 2026-01-13T03:11:40.556Z
Learnt from: jumyungc
Repo: newton-physics/newton PR: 1333
File: newton/_src/sim/builder.py:0-0
Timestamp: 2026-01-13T03:11:40.556Z
Learning: Ensure all VBD damping semantics follow Rayleigh-style damping and are unitless. This convention, demonstrated in SolverVBD, should be applied consistently across all VBD-related constraints in the codebase. In newton/_src/sim/builder.py, validate that any damping terms use the same normalization and are treated as dimensionless; add comments and unit tests to enforce consistency across VBD components.
Applied to files:
newton/examples/selection/example_selection_multiple.pynewton/_src/utils/selection.py
📚 Learning: 2025-08-25T20:20:35.899Z
Learnt from: dylanturpin
Repo: newton-physics/newton PR: 635
File: newton/_src/sim/builder.py:3191-3195
Timestamp: 2025-08-25T20:20:35.899Z
Learning: When reviewing PRs that update Python version requirements (e.g., dropping Python 3.9 support), focus on syntax modernization rather than suggesting behavioral changes. Adding explicit `strict=False` to zip() calls is appropriate for Python 3.10+ compatibility and maintains existing behavior while satisfying linters.
Applied to files:
newton/_src/utils/selection.py
📚 Learning: 2025-08-25T11:07:47.818Z
Learnt from: gdaviet
Repo: newton-physics/newton PR: 631
File: newton/examples/mpm/example_anymal_c_walk_on_sand.py:164-169
Timestamp: 2025-08-25T11:07:47.818Z
Learning: In Newton's MuJoCo solver implementation, setting TARGET_POSITION mode on floating base DOFs (first 6 DOFs) does not overconstrain the base - the solver handles this appropriately and maintains expected behavior for floating robots.
Applied to files:
newton/_src/utils/selection.py
📚 Learning: 2025-12-01T16:21:36.581Z
Learnt from: jvonmuralt
Repo: newton-physics/newton PR: 1160
File: newton/_src/sim/articulation.py:0-0
Timestamp: 2025-12-01T16:21:36.581Z
Learning: In Newton's Featherstone solver, during forward kinematics for FREE/DISTANCE joints, the computed spatial velocity v_wc follows Featherstone's spatial twist convention where the linear component (spatial_top) represents the velocity at the world origin (0,0,0), not at the joint anchor or COM. To convert to State.body_qd (which stores COM velocity), use v_com = v_origin + ω × x_com, where x_com is the absolute world position of the COM.
Applied to files:
newton/_src/utils/selection.py
📚 Learning: 2025-07-21T19:11:04.077Z
Learnt from: dylanturpin
Repo: newton-physics/newton PR: 450
File: newton/sim/ik.py:1286-1296
Timestamp: 2025-07-21T19:11:04.077Z
Learning: In Newton's IK system, the JointLimitObjective class intentionally only handles joints where DOF count equals coordinate count (1-1 mapping). This is appropriate because joint limits are typically only meaningful for simple joint types like revolute and prismatic joints, not complex joints like ball or free joints where DOF count != coordinate count.
Applied to files:
newton/_src/utils/selection.py
📚 Learning: 2025-09-22T21:08:31.901Z
Learnt from: dylanturpin
Repo: newton-physics/newton PR: 806
File: newton/examples/ik/example_ik_franka.py:121-123
Timestamp: 2025-09-22T21:08:31.901Z
Learning: In the newton physics framework, when creating warp arrays for IK solver joint variables using wp.array(self.model.joint_q, shape=(1, coord_count)), the resulting array acts as a reference/pointer to the original model's joint coordinates, so updates from the IK solver automatically reflect in the model's joint_q buffer used for rendering.
Applied to files:
newton/_src/utils/selection.py
📚 Learning: 2025-09-22T21:08:31.901Z
Learnt from: dylanturpin
Repo: newton-physics/newton PR: 806
File: newton/examples/ik/example_ik_franka.py:121-123
Timestamp: 2025-09-22T21:08:31.901Z
Learning: In the newton physics framework, when creating warp arrays for IK solver joint variables using wp.array(self.model.joint_q, shape=(1, coord_count)), the resulting array acts as a reference/pointer to the original model's joint coordinates, so updates from the IK solver automatically reflect in the model's joint_q buffer used for rendering, despite the general warp documentation suggesting copies are made by default.
Applied to files:
newton/_src/utils/selection.py
🧬 Code graph analysis (2)
newton/examples/selection/example_selection_multiple.py (2)
newton/_src/utils/selection.py (1)
ArticulationView(201-1106)newton/examples/__init__.py (3)
get_asset(34-35)test_body_state(38-113)init(350-395)
newton/_src/utils/selection.py (2)
newton/_src/sim/state.py (2)
joint_dof_count(212-216)joint_coord_count(205-209)newton/_src/sim/model.py (2)
ModelAttributeFrequency(49-74)Model(100-909)
🪛 Ruff (0.14.11)
newton/examples/selection/example_selection_multiple.py
150-150: Unused noqa directive (non-enabled: PLC0415)
Remove unused noqa directive
(RUF100)
227-227: Unused noqa directive (non-enabled: PLC0415)
Remove unused noqa directive
(RUF100)
254-254: Unused noqa directive (non-enabled: PLC0415)
Remove unused noqa directive
(RUF100)
287-287: Unused lambda argument: qd
(ARG005)
newton/_src/utils/selection.py
173-173: Avoid specifying long messages outside the exception class
(TRY003)
250-250: Avoid specifying long messages outside the exception class
(TRY003)
282-282: Avoid specifying long messages outside the exception class
(TRY003)
285-285: Avoid specifying long messages outside the exception class
(TRY003)
386-386: Avoid specifying long messages outside the exception class
(TRY003)
426-426: Avoid specifying long messages outside the exception class
(TRY003)
467-467: Avoid specifying long messages outside the exception class
(TRY003)
736-738: Avoid specifying long messages outside the exception class
(TRY003)
788-788: Avoid specifying long messages outside the exception class
(TRY003)
⏰ Context from checks skipped due to timeout of 900000ms. You can increase the timeout in your CodeRabbit configuration to a maximum of 15 minutes (900000ms). (4)
- GitHub Check: Run GPU Tests / Run GPU Unit Tests on AWS EC2
- GitHub Check: Run GPU Benchmarks / Run GPU Benchmarks on AWS EC2
- GitHub Check: run-newton-tests / newton-unittests (ubuntu-latest)
- GitHub Check: run-newton-tests / newton-unittests (windows-latest)
🔇 Additional comments (12)
newton/_src/utils/selection.py (6)
27-39: LGTM - Kernel correctly updated for 2D world-aware mapping.The kernel now properly uses 2D thread indices
(world, arti)and indexesview_to_model_mapas a 2D array, consistent with the multi-articulation per world design.
126-156: LGTM - FrequencyLayout encapsulates layout metadata cleanly.The class correctly determines whether indices form a contiguous slice or require an index array, supporting efficient zero-copy views when possible.
240-251: LGTM - squeeze_axes parameter handling is correct.The three-way logic (True/False/tuple) provides flexible control over output array dimensionality.
274-286: LGTM - Global articulation fallback logic is sound.The code correctly handles scenes with only global articulations (world=-1) by treating them as a single world, maintaining API consistency.
748-764: LGTM - Squeeze logic correctly handles all axis combinations.The squeeze implementation properly handles all combinations of singleton axes while ensuring at least one dimension remains, preventing scalar outputs.
1074-1094: LGTM - Mask conversion to model articulation indices is correct.The kernel launch dimensions match the 2D articulation_ids shape, and the per-world mask is correctly expanded to all articulations within each world.
newton/examples/selection/example_selection_multiple.py (6)
25-35: LGTM - Imports use public Newton API only.The example correctly imports from
newton.selection(public API) rather thannewton._src, following the coding guidelines.
84-94: LGTM - Random forces kernel correctly uses shape[2] for num_dofs.The previous bug using
shape[0]instead ofshape[2]for the DOF count has been correctly addressed.
217-246: LGTM - Step method correctly implements alternating resets and random forces.The logic properly alternates between mask subsets and applies randomized DOF forces using either Torch or Warp paths.
248-276: LGTM - Reset method properly handles both Torch and Warp paths.The conditional FK evaluation for non-MuJoCo solvers is correct, and the debug condition from previous review has been fixed.
282-289: LGTM - test_final method follows example guidelines.The unused
qdparameter in the lambda is required by thetest_body_stateinterface signature (which expectsCallable[[wp.transform, wp.spatial_vectorf], bool]), so this is not a concern.
291-304: LGTM - Main block follows example guidelines with kebab-case CLI argument.The example registration and CLI argument
--num-worldscorrectly use kebab-case as per coding guidelines.
✏️ Tip: You can disable this entire section by setting review_details to false in your review settings.
| # make sure inner strides are uniform | ||
| if not ( | ||
| all_equal(inner_joint_strides) | ||
| and all_equal(inner_joint_dof_strides) | ||
| and all_equal(inner_joint_coord_strides) | ||
| and all_equal(inner_link_strides) | ||
| and all_equal(inner_shape_strides) | ||
| ): | ||
| raise ValueError("Non-uniform strides within worlds are not supported") |
There was a problem hiding this comment.
🧩 Analysis chain
🏁 Script executed:
# First, find the all_equal function definition
rg "def all_equal" newton/_src/utils/selection.py -A 5Repository: newton-physics/newton
Length of output: 192
🏁 Script executed:
# Also search for all_equal imports/definitions in the entire codebase
rg "def all_equal" --type py -A 5Repository: newton-physics/newton
Length of output: 378
🏁 Script executed:
# Check the context around lines 404-479 to understand inner_*_strides structure
sed -n '404,479p' newton/_src/utils/selection.pyRepository: newton-physics/newton
Length of output: 4185
🏁 Script executed:
# Look for additional validation before or after the stride checks
sed -n '440,490p' newton/_src/utils/selection.py | cat -nRepository: newton-physics/newton
Length of output: 3130
🏁 Script executed:
# Search for any other all_equal calls that might validate within-world uniformity
rg "all_equal" newton/_src/utils/selection.py -B 2 -A 2Repository: newton-physics/newton
Length of output: 1610
🏁 Script executed:
# Verify the bug by checking what happens after validation
# Check if within-world uniformity is actually required by the later code logic
sed -n '459,500p' newton/_src/utils/selection.py | head -30Repository: newton-physics/newton
Length of output: 1449
🏁 Script executed:
# Look for any flattening or additional checks on the inner strides after validation
rg "inner_joint_stride" newton/_src/utils/selection.py -A 3 -B 1Repository: newton-physics/newton
Length of output: 1598
🏁 Script executed:
# Check if there's any downstream validation or usage that would catch within-world stride non-uniformity
rg "inner_joint_stride" newton/_src/utils/selection.py -B 5 -A 5 | head -40Repository: newton-physics/newton
Length of output: 2218
🏁 Script executed:
# Look for tests that might validate this behavior
find . -name "*test*.py" -o -name "*_test.py" | xargs grep -l "inner.*stride" 2>/dev/null | head -5Repository: newton-physics/newton
Length of output: 47
The stride validation is incomplete and could miss non-uniform within-world strides.
all_equal(inner_joint_strides) correctly works on nested lists (Python's == compares element-wise), but the validation only checks that all worlds have identical stride patterns. It does not verify that strides are uniform within each world. The error message states "Non-uniform strides within worlds are not supported," but the code later uses inner_joint_stride = inner_joint_strides[0][0], extracting only the first stride value. This assumes all strides are identical, which the validation doesn't enforce.
Example: [[1, 2, 3], [1, 2, 3]] would pass (worlds are identical) but violates the semantic requirement. Add a check to verify all elements within each world's stride list are equal to each other.
🧰 Tools
🪛 Ruff (0.14.11)
467-467: Avoid specifying long messages outside the exception class
(TRY003)
🤖 Prompt for AI Agents
In `@newton/_src/utils/selection.py` around lines 459 - 467, The current stride
check only ensures all worlds have identical stride lists but not that each
world's list is internally uniform; update the validation around the existing
block that inspects inner_joint_strides, inner_joint_dof_strides,
inner_joint_coord_strides, inner_link_strides, and inner_shape_strides to also
verify each sublist is uniform (e.g., for each sublist in inner_joint_strides:
all_equal(sublist) or len(set(sublist)) == 1), and do the same for the other
four stride variables; if any sublist is non-uniform, raise the existing
ValueError("Non-uniform strides within worlds are not supported") so later code
using inner_joint_stride = inner_joint_strides[0][0] is valid.
| if mask is None: | ||
| mask = self.full_mask | ||
| else: | ||
| if not isinstance(mask, wp.array): | ||
| mask = wp.array(mask, dtype=bool, shape=(self.count,), device=self.device, copy=False) | ||
| assert mask.shape == (self.count,) | ||
| elif not isinstance(mask, wp.array): | ||
| mask = wp.array(mask, dtype=bool, shape=(self.world_count,), device=self.device, copy=False) | ||
| assert mask.shape == (self.world_count,) | ||
|
|
There was a problem hiding this comment.
🧩 Analysis chain
🏁 Script executed:
# First, get an overview of the file size and locate the method
wc -l newton/_src/utils/selection.pyRepository: newton-physics/newton
Length of output: 100
🏁 Script executed:
# Read the _set_attribute_values method and surrounding context
sed -n '810,880p' newton/_src/utils/selection.py | cat -nRepository: newton-physics/newton
Length of output: 3331
🏁 Script executed:
# Search for the kernel definitions mentioned in the concern
rg -A 20 "def set_articulation_attribute.*_d_kernel" newton/_src/utils/selection.py | head -100Repository: newton-physics/newton
Length of output: 47
🏁 Script executed:
# Check how squeeze parameter affects the logic
rg -B 5 -A 5 "squeeze\[0\]" newton/_src/utils/selection.pyRepository: newton-physics/newton
Length of output: 530
🏁 Script executed:
# Search for set_articulation_attribute kernels across the file
rg "set_articulation_attribute.*kernel" newton/_src/utils/selection.pyRepository: newton-physics/newton
Length of output: 664
🏁 Script executed:
# Search for where these kernels are defined - they might be module-level
rg -B 5 -A 15 "@wp.kernel" newton/_src/utils/selection.py | grep -A 20 "set_articulation_attribute"Repository: newton-physics/newton
Length of output: 1477
🏁 Script executed:
# Look for full_mask definition
rg -B 3 -A 3 "self.full_mask" newton/_src/utils/selection.py | head -40Repository: newton-physics/newton
Length of output: 981
🏁 Script executed:
# Get more context around the squeeze logic to understand the full picture
sed -n '700,850p' newton/_src/utils/selection.py | cat -nRepository: newton-physics/newton
Length of output: 7376
🏁 Script executed:
# Check if there's any special handling for the squeeze case in the kernels or elsewhere
rg "world_count.*1" newton/_src/utils/selection.py | head -20Repository: newton-physics/newton
Length of output: 193
🏁 Script executed:
# Check squeeze parameter default and usage
rg -B 3 -A 3 "self.squeeze" newton/_src/utils/selection.py | head -50Repository: newton-physics/newton
Length of output: 1321
🏁 Script executed:
# Get the full context of when mask is used - any special cases?
sed -n '815,875p' newton/_src/utils/selection.py | cat -nRepository: newton-physics/newton
Length of output: 2883
🏁 Script executed:
# Check if there's any masking alternative or workaround for the squeeze case
rg -B 5 -A 10 "set_articulation_attribute.*kernel" newton/_src/utils/selection.py | grep -A 15 "@wp.kernel" | head -80Repository: newton-physics/newton
Length of output: 1488
🏁 Script executed:
# Verify: when squeeze[0]=True, what's the actual first dimension of attrib after squeezing?
# Let's trace through the squeeze logic more carefully
python3 << 'EOF'
# Simulate the squeeze logic
def get_shape_after_squeeze(world_count, count_per_world, value_count, squeeze):
"""Simulate _get_attribute_array shape transformation"""
shape = (world_count, count_per_world, value_count)
# squeeze value and articulation axes
if squeeze[2] and value_count == 1:
if squeeze[1] and count_per_world == 1:
shape = shape[:1] # (world_count,)
else:
shape = shape[:2] # (world_count, count_per_world)
elif squeeze[1] and count_per_world == 1:
shape = (shape[0], shape[2]) # (world_count, value_count)
# squeeze world axis
if squeeze[0] and world_count == 1 and len(shape) > 1:
shape = shape[1:] # Remove first dimension
return shape
# Test the problematic case
world_count = 1
count_per_world = 5 # Multiple articulations per world
value_count = 3 # Multiple values per articulation
squeeze = [True, False, False] # squeeze world axis only
result_shape = get_shape_after_squeeze(world_count, count_per_world, value_count, squeeze)
print(f"world_count={world_count}, count_per_world={count_per_world}, value_count={value_count}")
print(f"squeeze={squeeze}")
print(f"Resulting attribute shape: {result_shape}")
print(f"Mask shape: ({world_count},)")
print(f"First kernel thread index 'i' ranges: [0, {result_shape[0]})")
print(f"But mask[i] only valid for i in [0, {world_count})")
print(f"Issue: Kernel accesses mask[i] where i can be > mask size!")
EOFRepository: newton-physics/newton
Length of output: 151
Fix mask indexing in kernels when world axis is squeezed with world_count=1.
When squeeze[0]=True, world_count=1, and the attribute has multiple dimensions, the world axis is removed from the attribute shape (lines 62-64 in _get_attribute_array). However, the mask remains shape (world_count,) = (1,) (line 140). All kernels index mask by the first thread component i, which now corresponds to the articulation (or value) dimension rather than world, causing out-of-bounds access when the first attribute dimension exceeds 1.
Example: With squeeze[0]=True, world_count=1, count_per_world=5, a 3D attribute becomes shape (5, 3) after squeezing. The kernel spawns threads with i ∈ [0, 5) and attempts mask[i] where mask is shape (1,).
🤖 Prompt for AI Agents
In `@newton/_src/utils/selection.py` around lines 836 - 841, Kernels index mask by
the first thread component (mask[i]) which is incorrect when the world axis was
squeezed (world_count==1) because i then iterates over per-world entries; change
kernel mask accesses to use the world/thread's world index (e.g., mask[w] or
mask[world_id]) instead of mask[i]; if kernels only receive a single linear
index, compute the world index explicitly (world = linear_index //
count_per_world or world = 0 when world_count==1) and use that to index mask,
leaving mask creation (wp.array(..., shape=(self.world_count,), ...)) unchanged.
3 participants
Description
ArticulationViewnow supports multiple articulations per world.Example
This example creates 16 worlds with 3 articulations each:
Multi-articulation selections are supported out-of-the-box with no special construction needed:
The optional joint and link filtering options are the same as before. The main difference is the dimensionality of the attributes. To handle multiple articulations per world, we introduce a new dimension so that values in attribute arrays can be accessed using a triple of indices
(world, arti, value), wherevaluestands for links, joints, shapes, etc.Multi-articulation views allow us to write code that handles all of the articulations together, instead of using multiple views to handle one articulation per world.
Here's how random forces can be applied:
Attribute Dimensions
With multiple articulations per world, we need 3D attribute arrays indexed by (
world, arti, value). With a single articulation per world, we only need 2D arrays (world, value), since the singleton articulation dimension can be dropped.For single-body articulations, we can also drop the
valueindex. For example, an attribute with multiple balls per world can be accessed using two indices (world, arti). An attribute with a single ball per world can be accessed using just theworldindex.Here we have a design decision to make.
ArticulationViewcan compute the minimal set of dimensions required for an attribute by "squeezing" (removing) singleton dimensions. Benefits:world, value) rather than a 3D array where the articulation index is always 0.However, there are some drawbacks. Most importantly, the dimensionality of the attribute arrays can change when the number of objects in the scene changes (one per world vs multiple per world).
If code generality is important, it is better to consistently return attribute shapes that won't change with the number of objects.
Since user preferences or requirements may vary, we allow controlling the behaviour using an optional
squeeze_axesargument to theArticulationViewconstructor:The
squeeze_axesargument can take the following values:True: squeeze all singleton axes, return arrays with the lowest dimensionality possible.False: don't squeeze any singleton axes, always return 3D arrays.tuple[int]: explicitly specify squeezable axes, where 0 is the world axis, 1 is the articulation axis, and 2 is the value axis. This is the most fine-grained way to control the dimensionality of the attribute arrays. For example,squeeze_axes=(2,)says that the value axis can be squeezed (for single-body articulations), but the articulation or world axes should not be squeezed.Default behaviour
Currently, the default behaviour is to squeeze the articulation axis if possible (
squeeze_axes=(1,)). The world and value axes won't be squeezed. This was chosen to avoid breaking changes with existing code that worked with a single articulation per world. But we can pick a different default, opinions welcome :)Limitations
There are some limitations currently:
The
ArticulationViewconstructor verifies the assumptions and will raise errors when they are not met.We should be able to relax these limitations in the future to make the
ArticulationViewmore flexible, but we need to be careful about introducing overhead when zero-copy access is not possible.Newton Migration Guide
Please ensure the migration guide for warp.sim users is up-to-date with the changes made in this PR.
docs/migration.rstis up-to dateBefore your PR is "Ready for review"
newton/tests/test_examples.py)pre-commit run -aSummary by CodeRabbit
New Features
Documentation / Examples
Tests
✏️ Tip: You can customize this high-level summary in your review settings.