newton/sim/builder.py (1)

3703-3713: Fix required: add_particles must update particle_group to match particle_count

Short: add_particles() extends particle arrays but does not update self.particle_group, so finalize()'s conversion (m.particle_group = wp.array(self.particle_group, ...)) can produce a length mismatch.

Places to fix:

• newton/sim/builder.py
  • add_particles (def at line ~2673): after self.particle_flags.extend(flags) add:
    self.particle_group.extend([self.current_env_group] * len(pos))
  • finalize (m.particle_group assignment at line ~3710): add a pre-finalize sanity check:
    assert len(self.particle_group) == self.particle_count,
    f"particle_group length {len(self.particle_group)} != particle_count {self.particle_count}"

Suggested diff (minimal):

•    self.particle_flags.extend(flags)
  

•    self.particle_group.extend([self.current_env_group] * len(pos))
  

• before converting to Warp arrays in finalize:

• assert len(self.particle_group) == self.particle_count, \

•    f"particle_group length {len(self.particle_group)} != particle_count {self.particle_count}"
  

Tag:

newton/sim/builder.py (3)

120-146: Docs clarify env-grouping semantics; add one note about scope and explicit env behavior

The “Environment Grouping” section is clear and directly addresses previous confusion about add_builder overriding sub-builder groups.

Two small doc nits to prevent confusion:

• Explicitly call out that grouping applies only to core entities: particles, bodies, shapes, joints, articulations (not springs/constraints/etc.).
• Clarify what happens to num_envs when a specific environment is provided to add_builder(environment=...).

Proposed wording tweak (insert after Line 128):

Applies to: particles, bodies, shapes, joints, articulations. Other elements (e.g., springs, constraints) are currently ungrouped.

When environment is specified in add_builder(environment=...), the builder will use that index for all copied entities. num_envs will be updated to cover the highest environment index seen (environment + 1).

---

415-416: Body group state added — LGTM

Inline comment is helpful and resolves past request to document.

---

565-592: Docs explicitly stating override-on-merge behavior — LGTM

This addresses prior confusion called out in earlier review threads. Thanks for spelling out the two usage patterns and examples.

newton/sim/builder.py (2)

561-562: add_builder gains environment parameter — consider validating input

Signature extension makes sense. Consider validating environment >= -1 to catch accidental negative values other than -1.

Example:

if environment is not None and environment < -1:
    raise ValueError(f"Invalid environment index {environment}; expected -1 or >= 0.")

---

597-607: Optional: validate environment index early

Defensive check to avoid unexpected negative indices beyond -1.

Diff (within this block):

         # Set the environment group for entities being added
-        if environment is None:
+        if environment is None:
             # Use the current environment count as the group index if not specified
             group_idx = self.num_envs if update_num_env_count else self.current_env_group
         else:
+            if environment < -1:
+                raise ValueError(f"Invalid environment index {environment}; expected -1 or >= 0.")
             group_idx = environment

Configuration used: .coderabbit.yml
Review profile: CHILL
Plan: Pro

💡 Knowledge Base configuration:

• MCP integration is disabled by default for public repositories
• Jira integration is disabled by default for public repositories
• Linear integration is disabled by default for public repositories

You can enable these settings in your CodeRabbit configuration.

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

3698-3699: Fix required: ModelBuilder.add_particles doesn't populate particle_group — causes length mismatch at finalize

Confirmed: ModelBuilder.add_particles (list form) extends particle arrays but does not add entries to self.particle_group. finalize builds the model from self.particle_group (so a shorter particle_group will produce shape/length mismatches at runtime).

Places to fix / verify:

• newton/_src/sim/builder.py — add_particles (def at ~line 2661): missing particle_group extension.
• newton/_src/sim/builder.py — add_particle (single, ~lines 2628–2656): correctly does self.particle_group.append(self.current_env_group).
• newton/_src/sim/builder.py — add_builder merge path (~lines 708–714): correctly extends particle_group for merged builders.
• newton/_src/sim/builder.py — finalize (def at ~line 3655; m.particle_group assigned ~lines 3696–3699): consumes self.particle_group to create the model array.

Suggested minimal fix (apply inside ModelBuilder.add_particles, after self.particle_flags.extend(flags)):

# ensure environment group index is set for each new particle
self.particle_group.extend([self.current_env_group] * len(pos))

Optional fail-fast sanity check to add in finalize (before creating model arrays):

assert len(self.particle_group) == len(self.particle_q), "particle_group length mismatch"

Please apply the small change above in add_particles and add the assertion in finalize to prevent regressions.

newton/_src/sim/builder.py (2)

1674-1949: collapse_fixed_joints must remap group arrays (body_group, joint_group, articulation_group) — fix required

Verified: collapse_fixed_joints (newton/_src/sim/builder.py, def at ~line 1674) rebuilds bodies/joints and updates articulation_start but does not update self.body_group, self.joint_group or self.articulation_group. finalize() later emits these arrays to the Model (m.body_group @ ~3896, m.joint_group @ ~3909, m.articulation_group @ ~3949), so the groups become stale/misaligned. Tests and import paths call collapse_fixed_joints (e.g. newton/tests/test_model.py lines ~162–166), so this is a correctness/regression risk.

Locations to change

• newton/_src/sim/builder.py
  • collapse_fixed_joints definition (starts ~1674)
  • body list clearing & repopulation (clears at ~1839; retained_bodies loop ~1852–1875)
  • retained_joints handling and joint repopulation (~1877–1917)
  • articulation_start update and the problematic list(set(...)) usage (~1883–1892)
  • finalize() uses of group arrays (~3896, ~3909, ~3949)

Suggested minimal fix

• Snapshot old group arrays before clearing:
  • old_body_group = list(self.body_group)
  • old_joint_group = list(self.joint_group)
  • old_articulation_group = list(self.articulation_group)
  • old_articulation_start = list(self.articulation_start)
• Remap body_group while rebuilding bodies:
  • In the retained_bodies loop, alongside body_remap/new_id append new_body_group.append(old_body_group[body["original_id"]]).
• Remap joint_group while rebuilding joints:
  • When iterating retained_joints (in original order), append new_joint_group.append(old_joint_group[joint["original_id"]]).
• Rebuild articulation_start and articulation_group together (do NOT use list(set(...))):
  • Map each old articulation start via joint_remap (or keep start if unmapped).
  • Build list of (mapped_start, old_articulation_index, old_group), sort by mapped_start, remove duplicates where mapped_start collapses with previous articulation, then set self.articulation_start = [mapped_start ...] and self.articulation_group = [corresponding old_group ...].
• Replace the current articulation_start = list(set(...)) line with the deterministic remapping above to preserve ordering and group alignment.

Would you like me to prepare a targeted patch/PR that implements the above remapping and replaces the list(set(...)) use with a deterministic reconstruction of articulation_group?

---

2681-2690: Fix required: add_particles must populate particle_group to preserve particle_count invariant

add_particles extends particle arrays but does not add entries to self.particle_group; add_particle (single) does. This causes len(particle_group) != particle_count and can break finalization.

Files / locations to change:

• newton/_src/sim/builder.py — add_particles (around lines 2661–2691): after self.particle_flags.extend(flags) add group entries.
• Note: add_particle (around lines 2625–2660) already does self.particle_group.append(self.current_env_group) — use the same behavior for vectorized creation.

Suggested diff:

         self.particle_radius.extend(radius)
         self.particle_flags.extend(flags)
+        # Populate environment group indices for vectorized particle creation
+        self.particle_group.extend([self.current_env_group] * len(pos))

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

94-96: New per-entity group fields: API shape looks good

Adding nullable int32 arrays per-entity type is consistent and non-breaking. One optional improvement: consider registering these fields in attribute_frequency for introspection parity (e.g., "particle", "shape", "body", "joint", "articulation").

If you want, I can add the attribute_frequency mappings in the right section for you.

Also applies to: 150-152, 224-226, 285-287, 291-293

---

382-441: Consider exposing group fields in attribute_frequency for consistency

Downstream utilities that rely on attribute_frequency might benefit from knowing the frequency of the new group attributes (e.g., for generic attribute iteration/validation).

Happy to add:

• "particle_group": "particle"
• "shape_group": "shape"
• "body_group": "body"
• "joint_group": "joint"
• "articulation_group": "articulation"

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

816-818: Avoid incrementing num_envs for global (-1) “environment”

Only increment when assigning to a non-negative environment. This also aligns with the doc intent and typical per-env allocations downstream.

Apply this diff:

-        if update_num_env_count:
-            self.num_envs += 1
+        if update_num_env_count and group_idx >= 0:
+            self.num_envs += 1

Configuration used: .coderabbit.yml
Review profile: CHILL
Plan: Pro

💡 Knowledge Base configuration:

• MCP integration is disabled by default for public repositories
• Jira integration is disabled by default for public repositories
• Linear integration is disabled by default for public repositories

You can enable these sources in your CodeRabbit configuration.

newton/_src/sim/builder.py (8)

329-330: Introduce and maintain invariants for group arrays

Defining particle_group is good. Please ensure a strict invariant throughout the builder: for every entity type, len(*_group) must always equal the corresponding count. A later comment flags a concrete violation in add_particles.

---

449-453: Default current_env_group (-1) is sensible

The default aligns with “global entities.” Consider documenting that users who set current_env_group directly should also ensure num_envs is large enough (see finalize comment).

---

589-599: Exception safety: restore current_env_group via try/finally

current_env_group is saved and restored at the end. If an exception occurs in-between (e.g., in transform_mul or deep copies), current_env_group will remain set to the temporary group. Wrap the body of add_builder in a try/finally to guarantee restoration.

Apply this minimal restructuring (pseudo-diff, indentation omitted for brevity):

- prev_env_group = self.current_env_group
- self.current_env_group = group_idx
+ prev_env_group = self.current_env_group
+ self.current_env_group = group_idx
+ try:
     # ... existing add_builder logic ...
- # Restore the previous environment group
- self.current_env_group = prev_env_group
+ finally:
+   self.current_env_group = prev_env_group

---

823-825: State restoration acknowledged

Restoring current_env_group is good; see earlier note about using try/finally for exception safety.

---

3675-3691: Make num_envs consistent with assigned group indices

Currently, finalize forces num_envs = max(1, self.num_envs). If users set current_env_group manually (e.g., to 1 or higher) without using add_builder(update_num_env_count=True), num_envs may be too small. Suggest ensuring num_envs is at least max(nonnegative group index) + 1 across all entity types.

Apply this refinement:

-        # ensure the env count is set correctly
-        self.num_envs = max(1, self.num_envs)
+        # ensure the env count covers all non-negative group indices
+        def _max_nonneg(xs):
+            return max((g for g in xs if g >= 0), default=-1)
+        max_group = max(
+            _max_nonneg(self.particle_group),
+            _max_nonneg(self.shape_group),
+            _max_nonneg(self.body_group),
+            _max_nonneg(self.joint_group),
+            _max_nonneg(self.articulation_group),
+        )
+        self.num_envs = max(1, self.num_envs, max_group + 1)

This keeps backward compatibility and hardens finalize for manually assigned groups.

---

3703-3703: Serialize particle_group — ensure size matches particle_count

This will error or misbehave if add_particles hasn’t updated particle_group accordingly. After fixing add_particles, consider a lightweight runtime check before creating arrays:

assert len(self.particle_group) == self.particle_count

(Or raise a ValueError with a clear message.)

---

3952-3954: Serialize articulation_group — potential mismatch with articulation_start updates

collapse_fixed_joints mutates articulation_start (including a non-stable set() dedup), but articulation_group is not updated. Consider keeping articulation_start order stable and updating articulation_group in lockstep.

Would you like a concrete patch to keep articulation_start stable and rebuild articulation_group appropriately?

---

544-584: Optional: Instrumentation to detect grouping mismatches early

Given the new per-entity group arrays, adding lightweight validation at the end of add_builder could save debugging time:

• particle_count == len(particle_group) - pre and post
• shape_count == len(shape_group)
• body_count == len(body_group)
• joint_count == len(joint_group)
• articulation_count == len(articulation_group)

This can be gated behind a debug flag or wp.config.verbose.

Also applies to: 589-599, 713-739

Configuration used: .coderabbit.yml
Review profile: CHILL
Plan: Pro

💡 Knowledge Base configuration:

• MCP integration is disabled by default for public repositories
• Jira integration is disabled by default for public repositories
• Linear integration is disabled by default for public repositories

You can enable these sources in your CodeRabbit configuration.

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

3618-3623: Free-joint auto-add assigns joint_group based on current_env_group, not the body’s env

When adding free joints for floating bodies, the joint’s group should inherit the child body’s environment (so world->child joints remain in the same env). Today it uses whatever current_env_group happens to be, which can mislabel joints in mixed-env scenes.

Apply this change:

     def add_free_joints_to_floating_bodies(self, new_bodies: Iterable[int] | None = None):
         """
         Adds a free joint to every body that is not a child in any joint and has mass > 0.
         Should be called after all other joints have been added.
         """
         # set(self.joint_child) is connected_bodies
         floating_bodies = set(new_bodies) - set(self.joint_child)
         for body_id in floating_bodies:
             if self.body_mass[body_id] > 0:
-                self.add_joint_free(child=body_id)
+                prev_env = self.current_env_group
+                try:
+                    # Inherit the child's environment for the world->child free joint
+                    body_env = self.body_group[body_id] if body_id < len(self.body_group) else -1
+                    self.current_env_group = body_env
+                    self.add_joint_free(child=body_id)
+                finally:
+                    self.current_env_group = prev_env

This aligns joint_group with the body’s env group.

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

2714-2716: Bulk particle add now maintains particle_group — thanks for fixing the invariant

This resolves length mismatches between particle_group and particle_* arrays reported earlier.

newton/tests/test_model.py (3)

270-291: add_particles grouping test covers the core path; consider tightening assertions

This validates default (-1) and non-default (0) environments and exercises finalize->Model propagation. Optionally assert particle_count == 5 and that dtype is int32 to catch serialization regressions early.

Apply this small augmentation:

         model = builder.finalize()
         particle_groups = model.particle_group.numpy()
 
         # First 3 particles should be in group -1
         self.assertTrue(np.all(particle_groups[0:3] == -1))
         # Next 2 particles should be in group 0
         self.assertTrue(np.all(particle_groups[3:5] == 0))
+        self.assertEqual(model.particle_count, 5)
+        self.assertEqual(particle_groups.dtype, np.int32)

---

292-396: Environment grouping test is comprehensive; make indexing assertions less brittle

Great coverage across particles, bodies, shapes, joints, articulations and auto-assign (env=2). The slice-based assertions assume insertion order; future refactors might change ordering without breaking grouping semantics. Prefer checking counts per group to reduce brittleness.

For example:

-        if len(body_groups) > 0:
-            self.assertEqual(body_groups[0], -1)  # ground body
-            self.assertTrue(np.all(body_groups[1:3] == 0))
-            self.assertTrue(np.all(body_groups[3:5] == 1))
-            self.assertTrue(np.all(body_groups[5:7] == 2))
+        if len(body_groups) > 0:
+            # counts by group
+            unique, counts = np.unique(body_groups, return_counts=True)
+            counts_by_group = dict(zip(unique.tolist(), counts.tolist()))
+            self.assertEqual(counts_by_group.get(-1, 0), 1)
+            self.assertEqual(counts_by_group.get(0, 0), 2)
+            self.assertEqual(counts_by_group.get(1, 0), 2)
+            self.assertEqual(counts_by_group.get(2, 0), 2)

Do similarly for particle_groups and shape_groups.

---

397-494: collapse_fixed_joints grouping checks: solid; consider asserting articulation grouping too

The pre-/post-collapse assertions for body_group and joint_group are spot on and validate the rebuild logic. As an extra guardrail, you could assert articulation_count and articulation_group remain consistent with the retained joint segments.

For example:

         model = builder.finalize()
         body_groups = model.body_group.numpy()
         joint_groups = model.joint_group.numpy()
+        if model.articulation_count > 0 and model.articulation_group is not None:
+            art_groups = model.articulation_group.numpy()
+            self.assertTrue(np.all((art_groups == 0) | (art_groups == 1)))

newton/_src/sim/builder.py (2)

544-585: add_builder: docs correctly state group override behavior; add an explicit note on env=-1 counting

The behavior “ALL entities from the source builder are assigned to the specified group” is clearly documented. Consider adding that passing environment=-1 produces global entities and should not increase num_envs (see code suggestion below).

---

3724-3725: *Model serialization of _group arrays — LGTM; optionally validate lengths before writing

Serializing group arrays as wp.int32 (or None when empty) matches the public API. To catch regressions earlier, consider validating lengths before serialization.

Add a lightweight validation helper:

     def finalize(self, device: Devicelike | None = None, requires_grad: bool = False) -> Model:
+        # Sanity checks for group array lengths
+        def _check_group_lengths():
+            assert len(self.particle_group) == self.particle_count or self.particle_count == 0
+            assert len(self.shape_group) == self.shape_count or self.shape_count == 0
+            assert len(self.body_group) == self.body_count or self.body_count == 0
+            assert len(self.joint_group) == self.joint_count or self.joint_count == 0
+            assert len(self.articulation_group) == self.articulation_count or self.articulation_count == 0
+        _check_group_lengths()

If asserts are too strict for end-users, log warnings instead of asserting.

Also applies to: 3767-3768, 3922-3923, 3935-3936, 3975-3976

Configuration used: .coderabbit.yml
Review profile: CHILL
Plan: Pro

💡 Knowledge Base configuration:

• MCP integration is disabled by default for public repositories
• Jira integration is disabled by default for public repositories
• Linear integration is disabled by default for public repositories

You can enable these sources in your CodeRabbit configuration.

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

823-832: Num envs update logic fixed (no increment for globals; reconcile explicit env indices)

This addresses the previously raised mismatch and ensures num_envs reflects the highest non-negative group used.

---

2724-2726: Bulk particle creation now updates particle_group

This resolves the earlier mismatch risk for add_particles.

---

1853-1900: Rebuild body_group in collapse_fixed_joints — nice catch

Saving the original mapping and rebuilding body_group for retained bodies prevents post-collapse misalignment.

newton/_src/sim/builder.py (2)

1918-1964: Rebuild joint_group in collapse_fixed_joints — correct; consider doing the same for articulations

Reconstructing joint_group to match retained_joints resolves joint-group desyncs. Consider verifying whether articulation_start changes (set/dedup) can desynchronize articulation_group. If collapse affects articulations or changes their ordering/count, rebuild articulation_group accordingly to keep it 1:1 with articulation_start.

Would you like a patch sketch to rebuild articulation_group based on the new articulation starts?

---

3734-3735: Serialize and Guard Per-Entity Group Arrays — Approved

All per-entity group arrays (particle_group, shape_group, body_group, joint_group, articulation_group) are correctly exported to the Model and guarded for empty counts in newton/_src/sim/builder.py at:

• Lines 3734–3735 (particle_group)
• Lines 3777–3778 (shape_group)
• Lines 3932–3933 (body_group)
• Lines 3945–3946 (joint_group)
• Lines 3984–3986 (articulation_group)

Optional Safety Hardening: To prevent future regressions, you can reconcile self.num_envs with the highest non-negative group index in use. Replace the existing

-        # ensure the env count is set correctly
-        self.num_envs = max(1, self.num_envs)

with:

+        # ensure the env count is set correctly
+        # Optional safety: reconcile with the highest non-negative group index in use
+        groups_max = max(
+            max(self.particle_group)      if self.particle_group      else -1,
+            max(self.shape_group)         if self.shape_group         else -1,
+            max(self.body_group)          if self.body_group          else -1,
+            max(self.joint_group)         if self.joint_group         else -1,
+            max(self.articulation_group)  if self.articulation_group  else -1,
+        )
+        computed_envs = groups_max + 1 if groups_max >= 0 else 0
+        self.num_envs = max(1, max(self.num_envs, computed_envs))

Apply this change at the location where self.num_envs is currently set (e.g., at the end of add_builder or in your finalize routine).

Group-invariant verification (already confirmed):

rg -n -C2 'self\.(particle|shape|body|joint|articulation)_group\.(append|extend)\(' newton/_src/sim/builder.py
rg -n -C2 'm\.(particle|shape|body|joint|articulation)_group\s*='   newton/_src/sim/builder.py

Configuration used: .coderabbit.yml
Review profile: CHILL
Plan: Pro

💡 Knowledge Base configuration:

• MCP integration is disabled by default for public repositories
• Jira integration is disabled by default for public repositories
• Linear integration is disabled by default for public repositories

You can enable these sources in your CodeRabbit configuration.