Support multi-articulation selections

[nvlukasz]

Description

ArticulationView now supports multiple articulations per world.

Example

This example creates 16 worlds with 3 articulations each:

# load articulation
arti = newton.ModelBuilder()
arti.add_mjcf(newton.examples.get_asset("nv_ant.xml"), ...)

# create world with multiple articulations
world = newton.ModelBuilder()
world.add_builder(arti, ...)
world.add_builder(arti, ...)
world.add_builder(arti, ...)

# create scene
scene = newton.ModelBuilder()
scene.add_ground_plane()
scene.replicate(world, num_worlds=16)

# finalize model
model = scene.finalize()

Multi-articulation selections are supported out-of-the-box with no special construction needed:

ants = ArticulationView(model, "ant", ...)

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), where value stands 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:

@wp.kernel
def random_forces_kernel(
    dof_forces: wp.array3d(dtype=float),  # dof forces (output)
    max_magnitude: float,  # maximum force magnitude
    seed: int,  # random seed
):
    world, arti, dof = wp.tid()
    num_artis, num_dofs = dof_forces.shape[1], dof_forces.shape[2]
    rng = wp.rand_init(seed, num_dofs * (world * num_artis + arti) + dof)
    dof_forces[world, arti, dof] = max_magnitude * (1.0 - 2.0 * wp.randf(rng))

...

dof_forces = ants.get_dof_forces(control)
wp.launch(random_forces_kernel, dim=dof_forces.shape, inputs=[dof_forces, 2.0, seed])
ants.set_dof_forces(control, dof_forces)

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 value index. 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 the world index.

Here we have a design decision to make. ArticulationView can compute the minimal set of dimensions required for an attribute by "squeezing" (removing) singleton dimensions. Benefits:

• This seems like the most intuitive behaviour for simple experiments. The user generally knows the contents of the scene and can anticipate the dimensionality of the attribute arrays. With a single Franka arm per world, it's more intuitive to have a 2D attribute array indexed with (world, value) rather than a 3D array where the articulation index is always 0.
• Using lower dimensional arrays is generally more efficient, since the stride/offset computations are simpler.
• No breaking change with views that have a single articulation per world.
• We are less likely to exceed Warp's maximum of four array dimensions. This could happen if the Model attribute has more than two dimensions.

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).

• It's harder to write generic code that can handle any number of objects.
• Different kernel variations may be needed to handle 3D, 2D, or 1D attribute views, depending on the scene contents. This is definitely a pain point for users.

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_axes argument to the ArticulationView constructor:

ants = ArticulationView(model, "ant", squeeze_axes=False, ...)

The squeeze_axes argument 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:

• All articulations in a view must have the same topology (same number of joints/dofs/links/shapes).
• All articulations must have the same root joint type (e.g. floating or fixed).
• The number of articulations per world must be uniform.
• The stride between articulations in the same world must be uniform and the stride between worlds must be uniform. That's what allows us to view the 1D attribute arrays (in Model, State, or Control) as 3D arrays without a copy, just some stride magic.

The ArticulationView constructor 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 ArticulationView more 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.

• The migration guide in docs/migration.rst is up-to date

Before your PR is "Ready for review"

• Necessary tests have been added and new examples are tested (see newton/tests/test_examples.py)
• Documentation is up-to-date
• Code passes formatting and linting checks with pre-commit run -a

Summary by CodeRabbit

• New Features

  • Multi-articulation, world-aware selection and per-world masking for articulation collections.
  • Option to squeeze returned attribute arrays for simpler shapes.
  • New layout abstraction distinguishing contiguous vs non-contiguous attribute storage.

• Documentation

  • Added a runnable example demonstrating multi-articulation control, per-world masking, and simulation drivers.

• Tests

  • Added automated test to exercise the new multi-articulation example.

[coderabbitai]

📝 Walkthrough

Walkthrough

ArticulationView was refactored to be world-aware using a new FrequencyLayout abstraction for per-attribute layout; mask/kernel APIs now use 2D world×articulation indexing. A multi-world selection example and a test registering it were added.

Changes

Cohort / File(s)	Change Summary
Selection core newton/_src/utils/selection.py	Added FrequencyLayout; reworked ArticulationView to use world-aware 2D layouts, per-world counts/strides, and squeeze_axes. Introduced helpers (get_world_offset, is_contiguous_slice, all_equal, list_of_lists, find_matching_ids, get_name_from_key). Updated masking APIs and changed set_model_articulation_mask_kernel to accept a 2D view_to_model_map.
Examples newton/examples/selection/example_selection_multiple.py	New example demonstrating multi-world articulation control. Adds Warp kernels (compute_middle_kernel, init_masks, check_mask, reset_kernel, random_forces_kernel) and an Example class that builds multi-world simulations, creates masks, resets, applies forces, steps, and renders.
Tests newton/tests/test_examples.py	Registered the new selection.example_selection_multiple example in the examples test harness with devices/test options and viewer enabled.

Sequence Diagram(s)

sequenceDiagram
    participant User
    participant ArtView as ArticulationView
    participant Freq as FrequencyLayout
    participant Kernel as Masking Kernel
    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 `#D3E4CD`: 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

Loading

Estimated code review effort

🎯 4 (Complex) | ⏱️ ~45 minutes

Points to review closely:

• FrequencyLayout offsets/stride calculations and contiguity detection.
• ArticulationView initialization, per-world grouping, and squeeze_axes edge cases.
• All call sites and kernels updated to accept a 2D view_to_model_map and view_mask shapes.
• New example: CPU/GPU branches, mask construction, and random/reset kernel correctness.

Possibly related PRs

• Fix selection slice shapes #792 — Modifies selection/slicing logic in newton/_src/utils/selection.py; likely related to slice/int/empty-slice handling that intersects this refactor.

Suggested reviewers

• eric-heiden

Pre-merge checks and finishing touches

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 30.00% which is insufficient. The required threshold is 80.00%.	You can run @coderabbitai generate docstrings to improve docstring coverage.

✅ Passed checks (2 passed)

Check name	Status	Explanation
Description Check	✅ Passed	Check skipped - CodeRabbit’s high-level summary is enabled.
Title check	✅ Passed	The title 'Support multi-articulation selections' directly captures the main objective of extending ArticulationView to support multiple articulations per world, which is the primary feature introduced in this PR.

✨ Finishing touches

• 📝 Generate docstrings

🧪 Generate unit tests (beta)

• Create PR with unit tests
• Post copyable unit tests in a comment

---

📜 Recent review details

Configuration used: Path: .coderabbit.yml

Review profile: CHILL

Plan: Pro

📥 Commits

Reviewing files that changed from the base of the PR and between 3e22190 and 206edfc.

📒 Files selected for processing (1)

• newton/_src/utils/selection.py (12 hunks)

🧰 Additional context used

🧠 Learnings (3)

📚 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

🧬 Code graph analysis (1)

newton/_src/utils/selection.py (2)

newton/_src/sim/joints.py (1)

• JointType (20-44)

newton/_src/sim/model.py (3)

• ModelAttributeFrequency (49-70)
• Model (96-798)
• get_attribute_frequency (782-798)

🪛 Ruff (0.14.5)

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)

---

726-728: Avoid specifying long messages outside the exception class

(TRY003)

---

778-778: 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). (1)

• GitHub Check: Run GPU Benchmarks / Run GPU Benchmarks on AWS EC2

🔇 Additional comments (9)

newton/_src/utils/selection.py (9)

28-38: LGTM: Kernel correctly implements 2D world-aware mapping.

The updated kernel signature and 2D indexing properly map view masks to model articulation masks using the (world, arti) layout. The per-world masking approach is appropriate for the use case.

---

126-156: LGTM: Clean abstraction for per-attribute layout metadata.

The FrequencyLayout class effectively encapsulates stride information and handles both contiguous (slice-based) and non-contiguous (index-based) layouts. The empty indices case (slice(0, 0)) is handled appropriately.

---

240-250: LGTM: Squeeze axes logic correctly handles all three modes.

The validation properly handles True (squeeze all), False (squeeze none), and tuple-based selective squeezing with appropriate bounds checking.

---

261-386: Comprehensive validation ensures uniform articulation topology.

The world-aware grouping correctly:

• Separates articulations by world and handles global articulations
• Validates uniform per-world counts (line 284-285)
• Extracts template-based topology using the first articulation
• Enforces homogeneity across all articulations (line 377-386)

The error messages are clear and actionable.

---

404-479: LGTM: Stride computation correctly handles multi-world and degenerate cases.

The outer (between-world) and inner (within-world) stride calculations properly:

• Compute strides only when needed (multi-world or multi-articulation-per-world)
• Validate uniformity to ensure zero-copy views are possible
• Fall back to articulation counts for degenerate single-world or single-articulation cases

---

633-669: LGTM: FrequencyLayout instances correctly configured for each frequency type.

All five frequency types (JOINT, JOINT_DOF, JOINT_COORD, BODY, SHAPE) are properly instantiated with the computed offsets, strides, value counts, and filtered indices.

---

716-797: LGTM: Attribute array construction correctly implements 3D world-aware layout with configurable squeezing.

The method properly:

• Constructs base 3D (world, arti, value) shape and strides from FrequencyLayout
• Applies squeezing logic for singleton dimensions based on self.squeeze configuration
• Preserves at least 1D arrays (line 751 check)
• Handles both contiguous (slice) and non-contiguous (indexed) layouts
• Appends trailing dimensions for multi-dimensional attributes

The squeezing logic correctly handles all combinations of singleton axes.

---

826-830: LGTM: Mask handling correctly validates per-world shape.

The mask validation properly defaults to full_mask and ensures the provided mask has shape (world_count,), consistent with the per-world masking approach used throughout.

---

177-178: Fix edge case: all_equal fails on empty list.

The function attempts to access values[0] without checking if the list is empty, which would raise an IndexError.

Apply this diff to handle the empty case:

 def all_equal(values):
+    if not values:
+        return True
     return all(x == values[0] for x in values)

Likely an incorrect or invalid review comment.

---

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.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

_{Comment @coderabbitai help to get the list of available commands and usage tips.}

[coderabbitai]

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, but mask is 1D with shape (world_count,) while the kernel set_model_articulation_mask_kernel expects 2D indexing (world, arti). The launch should use dim=(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: PLC0415 comments 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 qd parameter 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

📥 Commits

Reviewing files that changed from the base of the PR and between a546871 and ef201a8.

📒 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!

FrequencyLayout provides 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 FrequencyLayout entries for efficient attribute access. The squeeze_axes parameter provides flexible dimensionality control.

[coderabbitai]

Actionable comments posted: 2

🧹 Nitpick comments (1)

newton/examples/selection/example_selection_multiple.py (1)

150-150: Remove unused noqa directives.

The # noqa: PLC0415 comments 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 torch

Apply 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

📥 Commits

Reviewing files that changed from the base of the PR and between ef201a8 and 55c7ef2.

📒 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)

[coderabbitai]

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

📥 Commits

Reviewing files that changed from the base of the PR and between 55c7ef2 and 3e22190.

📒 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.

[nvlukasz]

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.

[adenzler-nvidia]

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.

[adenzler-nvidia]

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?

[nvlukasz]

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]

Codecov Report

❌ Patch coverage is 89.29766% with 32 lines in your changes missing coverage. Please review.
✅ All tests successful. No failed tests found.

Files with missing lines	Patch %	Lines
newton/_src/utils/selection.py	89.29%	32 Missing ⚠️

📢 Thoughts on this report? Let us know!

[adenzler-nvidia]

would be great to have explicit test coverage for some of the special cases as well, not just the example.

[adenzler-nvidia]

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?

[nvlukasz]

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?

[adenzler-nvidia]

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.

[adenzler-nvidia]

couldn't this be basically any joint except fixed here?

[nvlukasz]

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 :)

[adenzler-nvidia]

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.