Add com points in viewer #1313
Add com points in viewer #1313
Conversation
📝 WalkthroughWalkthroughAdds center-of-mass (COM) visualization: a new Warp kernel computing world-space COM positions from body transforms and local COMs, plus viewer integrations that allocate buffers, invoke the kernel, and log COM points during rendering. Changes
Sequence Diagram(s)sequenceDiagram
participant RenderLoop as Render Loop
participant Viewer as ViewerBase
participant Kernel as compute_com_positions
participant Viz as Visualization System
RenderLoop->>Viewer: log_state(state)
Viewer->>Viewer: _log_com(state) (lazy init buffers)
Viewer->>Kernel: launch(body_q, body_com, body_world, world_offsets)
Kernel->>Kernel: transform local COM -> world COM (per-body)
Kernel-->>Viewer: com_positions
Viewer->>Viz: log_points("/model/com", positions, colors, radii)
Viz-->>RenderLoop: render COM points
Estimated code review effort🎯 3 (Moderate) | ⏱️ ~20 minutes Possibly related PRs
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
📜 Recent review detailsConfiguration used: Path: .coderabbit.yml Review profile: CHILL Plan: Pro 📒 Files selected for processing (2)
✅ Files skipped from review due to trivial changes (1)
🧰 Additional context used🧠 Learnings (1)📓 Common learnings⏰ 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)
🔇 Additional comments (1)
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 |
Codecov Report✅ All modified and coverable lines are covered by tests. 📢 Thoughts on this report? Let us know! |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 0
🧹 Nitpick comments (2)
newton/_src/viewer/kernels.py (1)
444-446: Consider adding bounds check forworld_offsetsaccess.The kernel accesses
world_offsets[world_idx]without first checking ifworld_offsetsis valid. Whenself.world_offsetsisNone(possible per viewer.py initialization), this could cause issues.For reference,
update_shape_xforms(lines 220-224) uses a safer pattern:if world_offsets: if shape_world >= 0 and shape_world < world_offsets.shape[0]: offset = world_offsets[shape_world]That said,
compute_joint_basis_linesuses the same pattern as your code, so this may be acceptable in practice. Worth verifying that the caller always ensuresworld_offsetsis allocated when any body has a validworld_idx.♻️ Suggested safer pattern
world_idx = body_world[tid] - if world_idx >= 0: + if world_offsets and world_idx >= 0 and world_idx < world_offsets.shape[0]: world_com = world_com + world_offsets[world_idx]newton/_src/viewer/viewer.py (1)
1220-1245: Implementation follows established patterns; minor optimization possible.The method correctly implements COM visualization:
- Lazy buffer allocation matches
_log_jointspattern (line 1185-1188)- Using
hidden=not self.show_commatches_log_particlespattern (line 1271)- Kernel launch and inputs are correct
Optional optimization: Unlike
_log_jointswhich returns early whenshow_jointsis False (line 1171-1174), this method always runs the kernel. You could add an early return for performance when COM visualization is disabled:♻️ Optional early return for performance
def _log_com(self, state): num_bodies = self.model.body_count if num_bodies == 0: return + if not self.show_com: + self.log_points("/model/com", None, None, None, hidden=True) + return + if self._com_positions is None or len(self._com_positions) < num_bodies:Note on static analysis hint: The
# noqa: PLC0415on line 1230 follows the existing file convention (see lines 378, 762, 1192), so keeping it for consistency is reasonable even if not strictly required by the current ruff configuration.
📜 Review details
Configuration used: Path: .coderabbit.yml
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (2)
newton/_src/viewer/kernels.pynewton/_src/viewer/viewer.py
🧰 Additional context used
🧠 Learnings (2)
📓 Common learnings
Learnt from: nvtw
Repo: newton-physics/newton PR: 747
File: newton/_src/viewer/kernels.py:0-0
Timestamp: 2025-09-10T10:09:05.676Z
Learning: In Newton physics picking implementation (newton/_src/viewer/kernels.py), damping is intentionally applied using center of mass velocity (vel_com) rather than velocity at the pick point. This is a deliberate design choice for stability improvements.
📚 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/viewer/kernels.py
🧬 Code graph analysis (1)
newton/_src/viewer/viewer.py (3)
newton/_src/viewer/kernels.py (1)
compute_com_positions(431-447)newton/_src/viewer/viewer_gl.py (1)
log_points(367-392)newton/_src/viewer/viewer_file.py (1)
log_points(114-116)
🪛 Ruff (0.14.10)
newton/_src/viewer/viewer.py
1230-1230: Unused noqa directive (non-enabled: PLC0415)
Remove unused noqa directive
(RUF100)
⏰ 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-newton-tests / newton-unittests (windows-latest)
🔇 Additional comments (3)
newton/_src/viewer/kernels.py (1)
430-448: LGTM on kernel logic.The COM world-space transformation is correct: extracting body position and rotation, then computing
body_pos + quat_rotate(body_rot, local_com). This matches the established pattern used elsewhere (e.g., line 142-143 inapply_picking_force_kernel).newton/_src/viewer/viewer.py (2)
66-69: LGTM on buffer initialization.The lazy allocation pattern with
Noneinitialization is consistent with existing buffers like_joint_points0and_joint_colors.
349-349: LGTM on integration point.The call to
_log_com(state)is correctly placed alongside the other visualization logging methods.
There was a problem hiding this comment.
Actionable comments posted: 1
🤖 Fix all issues with AI agents
In @newton/_src/viewer/kernels.py:
- Around line 430-448: In compute_com_positions, add defensive checks before
indexing world_offsets: first verify world_offsets is not null (e.g., ensure
world_offsets is valid) and then check that world_idx is within range (world_idx
>= 0 and world_idx < world_offsets.shape[0]) before doing world_com = world_com
+ world_offsets[world_idx]; if either check fails, skip adding the offset so
behavior matches the safe pattern used elsewhere (see similar checks at line
222/404).
📜 Review details
Configuration used: Path: .coderabbit.yml
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (1)
newton/_src/viewer/kernels.py
🧰 Additional context used
🧠 Learnings (3)
📓 Common learnings
Learnt from: nvtw
Repo: newton-physics/newton PR: 747
File: newton/_src/viewer/kernels.py:0-0
Timestamp: 2025-09-10T10:09:05.676Z
Learning: In Newton physics picking implementation (newton/_src/viewer/kernels.py), damping is intentionally applied using center of mass velocity (vel_com) rather than velocity at the pick point. This is a deliberate design choice for stability improvements.
📚 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/viewer/kernels.py
📚 Learning: 2025-09-10T10:09:05.676Z
Learnt from: nvtw
Repo: newton-physics/newton PR: 747
File: newton/_src/viewer/kernels.py:0-0
Timestamp: 2025-09-10T10:09:05.676Z
Learning: In Newton physics picking implementation (newton/_src/viewer/kernels.py), damping is intentionally applied using center of mass velocity (vel_com) rather than velocity at the pick point. This is a deliberate design choice for stability improvements.
Applied to files:
newton/_src/viewer/kernels.py
🧬 Code graph analysis (1)
newton/_src/viewer/kernels.py (1)
newton/_src/utils/import_utils.py (1)
transform(139-140)
⏰ 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 Benchmarks / Run GPU Benchmarks on AWS EC2
- GitHub Check: Run GPU Tests / Run GPU Unit Tests on AWS EC2
- GitHub Check: run-newton-tests / newton-unittests (windows-latest)
- GitHub Check: run-newton-tests / newton-unittests (ubuntu-latest)
There was a problem hiding this comment.
Actionable comments posted: 0
🧹 Nitpick comments (1)
newton/_src/viewer/kernels.py (1)
430-448: Kernel logic is correct; consider minor polish for consistency.The core transformation logic correctly computes world-space COM positions by rotating the local COM and adding the body position, with proper bounds-checked world offset application.
Note: The AI summary claims this kernel is "added in two places in the file, effectively duplicating the same functionality," but only one instance appears in the provided code.
Optional improvements for consistency and conciseness:
Use
wp.transform_pointfor clarity—it's more idiomatic and matches the pattern inapply_picking_force_kernel(line 142):Add a docstring explaining the kernel's purpose, similar to
compute_pick_state_kernel(lines 35-43),compute_contact_lines(line 304), andcompute_joint_basis_lines(lines 365-368).♻️ Suggested refinements
@wp.kernel def compute_com_positions( body_q: wp.array(dtype=wp.transform), body_com: wp.array(dtype=wp.vec3), body_world: wp.array(dtype=int), world_offsets: wp.array(dtype=wp.vec3), com_positions: wp.array(dtype=wp.vec3), ): + """ + Compute world-space center-of-mass positions for all bodies. + + Transforms each body's local COM to world space using the body's transform, + then applies the world offset if the body is assigned to a valid world. + """ tid = wp.tid() body_tf = body_q[tid] - body_pos = wp.transform_get_translation(body_tf) - body_rot = wp.transform_get_rotation(body_tf) local_com = body_com[tid] - world_com = body_pos + wp.quat_rotate(body_rot, local_com) + world_com = wp.transform_point(body_tf, local_com) world_idx = body_world[tid] if world_offsets and world_idx >= 0 and world_idx < world_offsets.shape[0]: world_com = world_com + world_offsets[world_idx] com_positions[tid] = world_com
📜 Review details
Configuration used: Path: .coderabbit.yml
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (1)
newton/_src/viewer/kernels.py
🧰 Additional context used
🧠 Learnings (5)
📓 Common learnings
Learnt from: nvtw
Repo: newton-physics/newton PR: 747
File: newton/_src/viewer/kernels.py:0-0
Timestamp: 2025-09-10T10:09:05.676Z
Learning: In Newton physics picking implementation (newton/_src/viewer/kernels.py), damping is intentionally applied using center of mass velocity (vel_com) rather than velocity at the pick point. This is a deliberate design choice for stability improvements.
📚 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/viewer/kernels.py
📚 Learning: 2025-09-10T10:09:05.676Z
Learnt from: nvtw
Repo: newton-physics/newton PR: 747
File: newton/_src/viewer/kernels.py:0-0
Timestamp: 2025-09-10T10:09:05.676Z
Learning: In Newton physics picking implementation (newton/_src/viewer/kernels.py), damping is intentionally applied using center of mass velocity (vel_com) rather than velocity at the pick point. This is a deliberate design choice for stability improvements.
Applied to files:
newton/_src/viewer/kernels.py
📚 Learning: 2025-09-25T16:14:22.033Z
Learnt from: dylanturpin
Repo: newton-physics/newton PR: 806
File: newton/_src/sim/ik/ik_objectives.py:0-0
Timestamp: 2025-09-25T16:14:22.033Z
Learning: wp.transform in NVIDIA Warp supports indexing using numerical indices (e.g., body_tf[0], body_tf[3], etc.) to access translation and rotation components. The indexing approach used in the Newton codebase is valid and compiles correctly.
Applied to files:
newton/_src/viewer/kernels.py
📚 Learning: 2025-09-26T05:58:21.013Z
Learnt from: WenchaoHuang
Repo: newton-physics/newton PR: 835
File: newton/_src/solvers/style3d/collision/collision.py:49-53
Timestamp: 2025-09-26T05:58:21.013Z
Learning: In newton/_src/solvers/style3d/collision/collision.py, the broad_phase buffers (broad_phase_ee, broad_phase_ef, broad_phase_vf) created with wp.array() don't need zero-initialization because they follow a write-first-read-later pattern: BVH query kernels in frame_begin() populate these buffers completely (write-only operations), and collision kernels only read from them afterward in accumulate_contact_force(), ensuring no undefined reads occur.
Applied to files:
newton/_src/viewer/kernels.py
🧬 Code graph analysis (1)
newton/_src/viewer/kernels.py (1)
newton/_src/utils/import_utils.py (1)
transform(139-140)
⏰ 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). (3)
- GitHub Check: Run GPU Benchmarks / Run GPU Benchmarks on AWS EC2
- GitHub Check: Run GPU Tests / Run GPU Unit Tests on AWS EC2
- GitHub Check: run-newton-tests / newton-unittests (windows-latest)
There was a problem hiding this comment.
Actionable comments posted: 0
🧹 Nitpick comments (1)
newton/_src/viewer/kernels.py (1)
430-444: LGTM! Implementation is correct and follows existing patterns.The kernel correctly transforms local body COM to world space and applies world offsets consistently with other kernels like
update_shape_xformsandcompute_joint_basis_lines.Consider adding a brief docstring for consistency with other kernels in this file:
📝 Optional: Add docstring
@wp.kernel def compute_com_positions( body_q: wp.array(dtype=wp.transform), body_com: wp.array(dtype=wp.vec3), body_world: wp.array(dtype=int), world_offsets: wp.array(dtype=wp.vec3), com_positions: wp.array(dtype=wp.vec3), ): + """Compute world-space center of mass positions for all bodies.""" tid = wp.tid()
📜 Review details
Configuration used: Path: .coderabbit.yml
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (1)
newton/_src/viewer/kernels.py
🧰 Additional context used
🧠 Learnings (5)
📓 Common learnings
Learnt from: nvtw
Repo: newton-physics/newton PR: 747
File: newton/_src/viewer/kernels.py:0-0
Timestamp: 2025-09-10T10:09:05.676Z
Learning: In Newton physics picking implementation (newton/_src/viewer/kernels.py), damping is intentionally applied using center of mass velocity (vel_com) rather than velocity at the pick point. This is a deliberate design choice for stability improvements.
📚 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/viewer/kernels.py
📚 Learning: 2025-09-10T10:09:05.676Z
Learnt from: nvtw
Repo: newton-physics/newton PR: 747
File: newton/_src/viewer/kernels.py:0-0
Timestamp: 2025-09-10T10:09:05.676Z
Learning: In Newton physics picking implementation (newton/_src/viewer/kernels.py), damping is intentionally applied using center of mass velocity (vel_com) rather than velocity at the pick point. This is a deliberate design choice for stability improvements.
Applied to files:
newton/_src/viewer/kernels.py
📚 Learning: 2025-09-25T16:14:22.033Z
Learnt from: dylanturpin
Repo: newton-physics/newton PR: 806
File: newton/_src/sim/ik/ik_objectives.py:0-0
Timestamp: 2025-09-25T16:14:22.033Z
Learning: In NVIDIA Warp's Newton physics library, wp.transform supports direct numerical indexing (e.g., body_tf[0], body_tf[1], body_tf[2] for position and body_tf[3], body_tf[4], body_tf[5], body_tf[6] for quaternion components) to access its elements. This is the standard API used throughout the codebase.
Applied to files:
newton/_src/viewer/kernels.py
📚 Learning: 2025-09-25T16:14:22.033Z
Learnt from: dylanturpin
Repo: newton-physics/newton PR: 806
File: newton/_src/sim/ik/ik_objectives.py:0-0
Timestamp: 2025-09-25T16:14:22.033Z
Learning: wp.transform in NVIDIA Warp supports indexing using numerical indices (e.g., body_tf[0], body_tf[3], etc.) to access translation and rotation components. The indexing approach used in the Newton codebase is valid and compiles correctly.
Applied to files:
newton/_src/viewer/kernels.py
🧬 Code graph analysis (1)
newton/_src/viewer/kernels.py (1)
newton/_src/utils/import_utils.py (1)
transform(139-140)
mzamoramora-nvidia
left a comment
mzamoramora-nvidia
left a comment
There was a problem hiding this comment.
Thanks for the contribution. It seems like newton-actuators was a added as a submodule to this PR. I guess that was just an accident. Otherwise the PR looks good to me.
github-merge-queue
bot
removed this pull request from the merge queue due to failed status checks
github-merge-queue
bot
removed this pull request from the merge queue due to failed status checks
github-merge-queue
bot
removed this pull request from the merge queue due to failed status checks
3 participants
Description
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
Bug Fixes / Behavior
✏️ Tip: You can customize this high-level summary in your review settings.