Example cube stacking#1256
Conversation
📝 WalkthroughWalkthroughAdds a new multi-world inverse-kinematics cube-stacking example for the Franka robot, updates README example links/images to include Cloth Twist and replace an IK reference with ik_cube_stacking, and registers the new example in tests with CUDA-specific options. Changes
Sequence Diagram(s)sequenceDiagram
participant CLI as User/CLI
participant App as Example (python)
participant Kernels as CUDA Kernels
participant IK as IK Solver
participant Sim as Simulator
participant Scene as Scene (worlds/cubes/table)
CLI->>App: start (args: num-worlds, cube-count)
App->>Scene: build multi-world scene
App->>Kernels: launch set_target_pose_kernel (per-world)
Kernels->>App: ee_pos/rot targets, gripper targets
App->>IK: set IK objectives & solve
IK->>App: joint targets
App->>Sim: apply joint targets / step simulation
Sim->>Scene: update transforms
App->>Kernels: launch advance_task_kernel (per-world)
Kernels->>App: update task_idx, task_time_elapsed
loop repeat per frame
App->>Kernels: set_target_pose_kernel
Kernels->>IK: targets
IK->>Sim: joint commands via App
end
App->>App: test_final validates cube placements
App-->>CLI: report success/failure
Estimated code review effort🎯 4 (Complex) | ⏱️ ~45 minutes
Suggested reviewers
Pre-merge checks and finishing touches❌ Failed checks (1 warning)
✅ Passed checks (2 passed)
✨ Finishing touches
🧪 Generate unit tests (beta)
📜 Recent review detailsConfiguration used: Path: .coderabbit.yml Review profile: CHILL Plan: Pro 📒 Files selected for processing (2)
🚧 Files skipped from review as they are similar to previous changes (1)
🧰 Additional context used🧠 Learnings (7)📓 Common learnings📚 Learning: 2025-09-22T21:08:31.901ZApplied to files:
📚 Learning: 2025-08-25T20:10:59.536ZApplied to files:
📚 Learning: 2025-09-22T21:08:31.901ZApplied to files:
📚 Learning: 2025-08-18T15:56:26.587ZApplied to files:
📚 Learning: 2025-08-12T17:52:15.009ZApplied to files:
📚 Learning: 2025-12-12T08:45:43.428ZApplied to files:
🧬 Code graph analysis (1)newton/examples/ik/example_ik_cube_stacking.py (3)
🪛 Ruff (0.14.8)newton/examples/ik/example_ik_cube_stacking.py476-476: Avoid specifying long messages outside the exception class (TRY003) 679-679: 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). (2)
🔇 Additional comments (6)
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! |
There was a problem hiding this comment.
Actionable comments posted: 1
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
README.md (1)
290-303: Update example command name to match registry.The command
ik_stack_cubesin the README is incorrect. The fileexample_ik_cube_stacking.pyis registered asik_cube_stacking(registration removes the "example_" prefix and ".py" extension). Update the README to useuv run -m newton.examples ik_cube_stacking.
🧹 Nitpick comments (3)
newton/examples/ik/example_ik_cube_stacking.py (3)
50-141: Consider documenting the magic number 14.The kernel logic is well-structured. However, the hardcoded offset
14(line 81:cube_index = cube_body_index - 14) assumes a fixed robot body count. This value is also hardcoded insetup_tasks(line 560) andtest_final(line 652). Consider extracting this as a named constant or kernel parameter to improve maintainability if the robot configuration changes.+ROBOT_BODY_COUNT = 14 # Number of bodies in the Franka robot (before cubes) + @wp.kernel(enable_backward=False) def set_target_pose_kernel( ... + robot_body_count: int, ... ): ... - cube_index = cube_body_index - 14 + cube_index = cube_body_index - robot_body_count
462-469: Add a maximum iteration guard to prevent potential infinite loop.The rejection sampling for cube placement could theoretically loop indefinitely if
min_distanceis too large relative to the sampling region or ifcube_countis high. Consider adding a maximum iteration limit with a fallback or warning.if len(cube_pos) > 0: # Check if the new cube is too close to the existing cubes. l2_dists_too_close = [wp.norm_l2(new_pos - pos) < min_distance for pos in cube_pos] + max_attempts = 1000 + attempts = 0 while any(l2_dists_too_close): new_pos = get_random_pos() l2_dists_too_close = [wp.norm_l2(new_pos - pos) < min_distance for pos in cube_pos] - print("Generating new random position...") + attempts += 1 + if attempts >= max_attempts: + raise RuntimeError(f"Failed to place cube {i} after {max_attempts} attempts")
569-607: Minor optimization: Avoid allocating arrays every frame.The target arrays (
ee_pos_target,ee_rot_target,gripper_target, etc.) are allocated on everyset_joint_targetscall. Consider pre-allocating these in__init__orsetup_tasksand reusing them.def setup_tasks(self): ... self.task_time_elapsed = wp.zeros(self.num_worlds, dtype=wp.float32) + + # Pre-allocate target arrays for set_joint_targets + self.ee_pos_target = wp.zeros(self.num_worlds, dtype=wp.vec3) + self.ee_pos_target_interpolated = wp.zeros(self.num_worlds, dtype=wp.vec3) + self.ee_rot_target = wp.zeros(self.num_worlds, dtype=wp.vec4) + self.ee_rot_target_interpolated = wp.zeros(self.num_worlds, dtype=wp.vec4) + self.gripper_target = wp.zeros(self.num_worlds, dtype=wp.vec2)
📜 Review details
Configuration used: Path: .coderabbit.yml
Review profile: CHILL
Plan: Pro
⛔ Files ignored due to path filters (1)
docs/images/examples/example_ik_cube_stacking.jpgis excluded by!**/*.jpg
📒 Files selected for processing (3)
README.md(3 hunks)newton/examples/ik/example_ik_cube_stacking.py(1 hunks)newton/tests/test_examples.py(1 hunks)
🧰 Additional context used
🧠 Learnings (4)
📓 Common learnings
Learnt from: shi-eric
Repo: newton-physics/newton PR: 521
File: newton/examples/example_cloth_hanging.py:36-36
Timestamp: 2025-08-12T05:17:34.423Z
Learning: The Newton migration guide (docs/migration.rst) is specifically for documenting how to migrate existing warp.sim functionality to Newton equivalents. New Newton-only features that didn't exist in warp.sim do not need migration documentation.
📚 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-08-19T12:27:30.629Z
Learnt from: preist-nvidia
Repo: newton-physics/newton PR: 579
File: newton/examples/example_mujoco.py:350-354
Timestamp: 2025-08-19T12:27:30.629Z
Learning: In Newton examples, there's a distinction between solver parameters and Example class attributes. The Example class can have its own use_mujoco attribute for controlling example-level behavior (like CUDA graphs, rendering logic), while the solver uses use_mujoco_cpu for backend selection. These serve different purposes and should not be conflated during API renames.
Applied to files:
newton/examples/ik/example_ik_cube_stacking.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/ik/example_ik_cube_stacking.py
🪛 Ruff (0.14.8)
newton/examples/ik/example_ik_cube_stacking.py
671-671: 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). (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 (7)
newton/tests/test_examples.py (1)
462-468: LGTM!The test registration follows the established pattern for IK examples. CUDA-only device restriction is appropriate for multi-world simulation, and the reduced
cube-count: 2(vs. the commented3) helps keep test runtime reasonable.README.md (1)
229-242: LGTM!The Cloth Twist example is correctly added to the Cloth Examples section, following the established table layout pattern.
newton/examples/ik/example_ik_cube_stacking.py (5)
38-47: LGTM!The
TaskTypeenum provides a clear state machine for the pick-and-place task sequence.
144-190: LGTM!The task advancement kernel correctly checks position/rotation error thresholds before transitioning states, and properly snapshots body transforms for the next task's interpolation reference.
193-280: LGTM!The initialization follows Newton example conventions. The collision pipeline is correctly created from args, and the CUDA graph capture for simulation is appropriately set up.
301-338: LGTM!The simulation loop correctly computes contacts once per frame and reuses them across substeps. Based on learnings, this is an intentional design pattern in Newton examples.
643-673: LGTM with minor suggestion.The validation logic correctly checks cube positions and orientations against tolerances. The 80% success rate threshold is reasonable for multi-world stochastic scenarios.
Per the static analysis hint (TRY003), consider using a custom exception class for the failure case, though this is a minor style preference.
There was a problem hiding this comment.
Actionable comments posted: 0
🧹 Nitpick comments (2)
newton/examples/ik/example_ik_cube_stacking.py (2)
439-496: LGTM: Robust cube placement with collision avoidance.The method correctly generates random cube positions while ensuring minimum separation distance. The max-attempts limit prevents infinite loops if placement fails.
Optional: Line 475 could use a custom exception class instead of a long message (per Ruff TRY003), but this is a minor style preference in example code.
650-680: LGTM: Validation logic with reasonable tolerances.The test correctly validates that cubes are stacked at target positions with 5cm position tolerance and 2-degree rotation tolerance, requiring 80% success rate across worlds.
Optional: Line 678 could use a custom exception class instead of a long message (per Ruff TRY003), but this is acceptable in test/example code where detailed messages aid debugging.
📜 Review details
Configuration used: Path: .coderabbit.yml
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (2)
README.md(3 hunks)newton/examples/ik/example_ik_cube_stacking.py(1 hunks)
🧰 Additional context used
🧠 Learnings (7)
📓 Common learnings
Learnt from: shi-eric
Repo: newton-physics/newton PR: 521
File: newton/examples/example_cloth_hanging.py:36-36
Timestamp: 2025-08-12T05:17:34.423Z
Learning: The Newton migration guide (docs/migration.rst) is specifically for documenting how to migrate existing warp.sim functionality to Newton equivalents. New Newton-only features that didn't exist in warp.sim do not need migration documentation.
📚 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/examples/ik/example_ik_cube_stacking.py
📚 Learning: 2025-08-25T20:10:59.536Z
Learnt from: dylanturpin
Repo: newton-physics/newton PR: 634
File: newton/examples/ik/example_ik_franka.py:0-0
Timestamp: 2025-08-25T20:10:59.536Z
Learning: In Newton's IK examples, when creating solver arrays with `wp.array(source_array, shape=new_shape)`, this creates a view into the same underlying memory rather than a copy. This means updates to the reshaped array are automatically reflected in the original array without needing explicit copy operations like `wp.copy()`.
Applied to files:
newton/examples/ik/example_ik_cube_stacking.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/examples/ik/example_ik_cube_stacking.py
📚 Learning: 2025-08-18T15:56:26.587Z
Learnt from: adenzler-nvidia
Repo: newton-physics/newton PR: 552
File: newton/_src/solvers/mujoco/solver_mujoco.py:0-0
Timestamp: 2025-08-18T15:56:26.587Z
Learning: In Newton's MuJoCo solver, when transforming joint axes from Newton's internal frame to MuJoCo's expected frame, use wp.quat_rotate(joint_rot, axis) not wp.quat_rotate_inv(joint_rot, axis). The joint_rot represents rotation from joint-local to body frame, so forward rotation is correct.
Applied to files:
newton/examples/ik/example_ik_cube_stacking.py
📚 Learning: 2025-08-12T17:52:15.009Z
Learnt from: nvlukasz
Repo: newton-physics/newton PR: 519
File: newton/_src/utils/import_mjcf.py:226-231
Timestamp: 2025-08-12T17:52:15.009Z
Learning: In the Newton codebase, when passing array-like objects (numpy arrays, lists, tuples) to wp.vec3(), the consistent pattern is to use the unpacking operator: wp.vec3(*array) rather than wp.vec3(array). This pattern is used throughout newton/_src/utils/import_mjcf.py and newton/_src/core/types.py.
Applied to files:
newton/examples/ik/example_ik_cube_stacking.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/ik/example_ik_cube_stacking.py
🧬 Code graph analysis (1)
newton/examples/ik/example_ik_cube_stacking.py (3)
newton/_src/sim/ik/ik_solver.py (2)
joint_q(362-363)IKSolver(204-456)newton/_src/sim/ik/ik_objectives.py (2)
IKPositionObjective(260-408)IKRotationObjective(706-863)newton/_src/sim/ik/ik_common.py (1)
IKJacobianMode(25-37)
🪛 Ruff (0.14.8)
newton/examples/ik/example_ik_cube_stacking.py
475-475: Avoid specifying long messages outside the exception class
(TRY003)
678-678: 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). (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 (12)
newton/examples/ik/example_ik_cube_stacking.py (10)
1-47: LGTM: Clean structure for task-based IK control.The license, imports, and TaskType enum are well-organized. The task sequence (approach, grasp, lift, move, release, retract, home) forms a complete pick-and-place workflow suitable for multi-world cube stacking.
50-143: LGTM: Task-based target pose computation is well-structured.The kernel correctly computes end-effector target poses and gripper positions based on the current task type. The use of
wp.quat_slerpfor rotation interpolation is appropriate, and the task-specific logic handles approach, grasp, lift, move, release, and retract phases correctly.
145-192: LGTM: Task advancement logic with appropriate error thresholds.The kernel correctly advances to the next task when position error is below 1cm and rotation error is below 1 degree, with a time-based soft limit. The snapshot of body transforms at task transitions ensures smooth interpolation for the next task phase.
195-283: LGTM: Well-structured initialization.The constructor properly initializes the multi-world simulation, builds the scene with the Franka robot and table, configures the MuJoCo solver, and sets up the IK objectives and task scheduling. The sequence of FK evaluation before IK setup is correct.
284-346: LGTM: Follows Newton simulation patterns correctly.The CUDA graph capture for performance optimization, once-per-frame contact computation reused across substeps (consistent with Newton's design), and the step/render loop are all implemented correctly.
Based on learnings, computing contacts once per frame is intentional.
348-437: LGTM: Scene construction follows Newton patterns.The Franka robot loading with proper URDF import, control parameter configuration, and multi-world scene building with proper
begin_world()/end_world()calls are all correctly implemented.
498-541: LGTM: IK solver setup with appropriate objectives.The IK configuration correctly uses position, rotation, and joint limit objectives. The use of
wp.clone()forjoint_q_ikensures independent arrays for the solver, andIKJacobianMode.ANALYTICis an appropriate choice for performance.
543-582: LGTM: Task scheduling setup is well-organized.The method correctly creates a task schedule with 9 phases per cube (approach, refine, grasp, lift, move, refine drop-off, release, retract, home), initializes task state tracking, and allocates target arrays for pose and gripper control.
683-692: LGTM: Standard Newton example entry point.The main block correctly follows Newton's example pattern with command-line argument parsing for
--num-worldsand--cube-count, proper initialization, and the standard run loop.
584-648: No action needed: array conversion is correct for Warp's type system.The pattern at line 630 using
wp.array(self.gripper_target_interpolated, dtype=wp.float32)is intentional and correct. Thegripper_target_interpolatedarray has dtypewp.vec2with shape(num_worlds,), storing two floats per element. Converting tofloat32dtype expands the vector components correctly to the expected(num_worlds, 2)shape for the copy destination. This is standard Warp usage and works as intended; the example passes all tests.README.md (2)
229-242: LGTM: Cloth Twist example properly documented.The new Cloth Twist entry follows the established README pattern with correct links, image reference, and run command.
290-303: LGTM: Cube stacking example properly integrated.The IK section correctly references the new
example_ik_cube_stacking.pyfile with appropriate links, image, and run command. The documentation aligns with the new example file added in this PR.
eric-heiden
left a comment
There was a problem hiding this comment.
Thanks, this looks great!
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
Documentation
Tests
✏️ Tip: You can customize this high-level summary in your review settings.