Skip to content

[Lab] Support runtime randomization of gravity #597#804

Merged
eric-heiden merged 15 commits into
newton-physics:mainfrom
vreutskyy:597-lab-support-runtime-randomization-of-gravity
Sep 25, 2025
Merged

[Lab] Support runtime randomization of gravity #597#804
eric-heiden merged 15 commits into
newton-physics:mainfrom
vreutskyy:597-lab-support-runtime-randomization-of-gravity

Conversation

@vreutskyy

@vreutskyy vreutskyy commented Sep 22, 2025

Copy link
Copy Markdown
Member

Description

This PR adds support for runtime gravity modification in Newton simulations. Closes #597

Key Changes:

  • Gravity is now stored as warp.array(dtype=wp.vec3, shape=[1]) in the Model class to enable runtime modification
  • Added Model.set_gravity() method to update gravity during simulation
  • Added SolverNotifyFlags.MODEL_PROPERTIES flag for notifying solvers of model property changes
  • Updated all solver kernels to accept gravity as a warp.array parameter instead of a value
  • Implemented optimized GPU-to-GPU transfer for MuJoCo solver gravity updates

Usage:

# Create model with initial gravity
builder = newton.ModelBuilder(gravity=-9.81)
model = builder.finalize()
solver = newton.solvers.SolverXPBD(model)

# Change gravity at runtime
model.set_gravity((0.0, 0.0, -4.9))  # Half gravity
solver.notify_model_changed(newton.solvers.SolverNotifyFlags.MODEL_PROPERTIES)

# Also works with CUDA graph capture
wp.capture_begin()
solver.step(state_0, state_1, control, None, dt)
graph = wp.capture_end()

model.set_gravity((0.0, 0.0, 0.0))  # Zero gravity
wp.capture_launch(graph)  # Uses updated gravity value

Technical Details:

  • Gravity array reference is passed to kernels, allowing them to read current values even during CUDA graph replay
  • All solvers (XPBD, VBD, Featherstone, SemiImplicit, Style3D, ImplicitMPM, MuJoCo) have been updated

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

    • Model.set_gravity enables runtime gravity updates and synchronizes gravity across solvers/backends.
  • Behavior / API

    • Gravity is now stored/passed as a 1-element vector array; callers should use the new setter and updated gravity representation.
  • Documentation

    • Solver notifications document model/global property updates (MODEL_PROPERTIES).
  • Tests

    • Added comprehensive runtime-gravity tests (particles, bodies, multi-solver/device, CUDA graph); one flaky MuJoCo test skipped.
  • Examples

    • Examples updated to use the public gravity setter API.

Signed-off-by: Viktor Reutskyy <vreutskyy@nvidia.com>
@coderabbitai

coderabbitai Bot commented Sep 22, 2025

Copy link
Copy Markdown
Contributor
📝 Walkthrough

Walkthrough

Model gravity storage changed to a 1-element wp.array(dtype=wp.vec3) with a new Model.set_gravity(); solver kernels and public solver APIs updated to accept gravity arrays and index gravity[0]. MuJoCo solver gained gravity-sync paths. XPBD restitution flow now includes a rigid-body restitution step. Runtime-gravity tests and examples updated.

Changes

Cohort / File(s) Summary
Gravity API: kernels & solver signatures
newton/_src/solvers/solver.py, newton/_src/solvers/implicit_mpm/solver_implicit_mpm.py, newton/_src/solvers/vbd/solver_vbd.py, newton/_src/solvers/style3d/kernels.py, newton/_src/solvers/featherstone/kernels.py, newton/_src/solvers/xpbd/kernels.py
Changed gravity parameter types from wp.vec3 to wp.array(dtype=wp.vec3) across kernels and solver APIs; internal uses updated to index gravity[0].
Model gravity storage & setter
newton/_src/sim/model.py, newton/_src/sim/builder.py
Model.gravity is now a 1-element Warp array; added Model.set_gravity(...) for runtime updates; builder/finalize creates a device-backed gravity array instead of a raw vec3.
MuJoCo gravity synchronization
newton/_src/solvers/mujoco/solver_mujoco.py
Added update_model_properties_kernel and SolverMuJoCo.update_model_properties() to copy/synchronize gravity into MuJoCo CPU/Warp structures; notify/initialization paths dispatch MODEL_PROPERTIES.
Notification flags
newton/_src/solvers/flags.py
Added MODEL_PROPERTIES enum member and included it in ALL to signal model-global property updates (e.g., gravity).
XPBD restitution flow
newton/_src/solvers/xpbd/solver_xpbd.py, newton/_src/solvers/xpbd/kernels.py
Inserted a rigid-body restitution kernel launch in the restitution-enabled path and a subsequent apply_body_delta_velocities call; updated kernel signature(s) to accept gravity array.
Runtime gravity tests
newton/tests/test_runtime_gravity.py
New test module validating runtime gravity changes for particles and rigid bodies across multiple solvers/devices, including CUDA-graph capture and gravity-fallback scenarios.
Examples & tests updated to use setter
newton/examples/cloth/example_cloth_franka.py, newton/examples/cloth/example_cloth_style3d.py, newton/examples/mpm/example_mpm_granular.py, newton/examples/robot/example_robot_policy.py, newton/tests/test_cloth.py
Replaced direct model.gravity = … assignments with model.set_gravity(...) calls; tests updated to read gravity from the warp array (e.g., .numpy()[0]).
Featherstone minor formatting
newton/_src/solvers/featherstone/solver_featherstone.py
Single blank-line insertion; no functional change.

Sequence Diagram(s)

sequenceDiagram
  autonumber
  actor User
  participant Model
  participant Solver
  participant Backend as MuJoCo/Backend
  participant Kernels

  User->>Model: set_gravity((gx,gy,gz))
  Model-->>User: gravity[0] updated

  User->>Solver: notify(MODEL_PROPERTIES)
  Solver->>Backend: update_model_properties()
  alt MuJoCo CPU
    Backend->>Backend: mj_model.opt.gravity = gravity.numpy()[0]
  else MuJoCo Warp
    Backend->>Kernels: launch update_model_properties_kernel(gravity_src, gravity_dst)
  end

  loop per simulation step
    Solver->>Kernels: launch integrate_* / apply_* (pass gravity array)
    Kernels-->>Solver: read/use gravity[0] in computations
  end
Loading
sequenceDiagram
  autonumber
  participant SolverXPBD as Solver XPBD
  participant BodyBuffers as Body State/Buffers
  participant Kernels

  rect rgb(245,250,255)
  note right of SolverXPBD: Restitution-enabled path (rigid bodies present)
  SolverXPBD->>BodyBuffers: zero/prepare body_deltas
  SolverXPBD->>Kernels: apply_rigid_restitution(..., gravity, ...)
  Kernels-->>BodyBuffers: rigid contact impulses/deltas (use gravity[0])
  SolverXPBD->>Kernels: apply_body_delta_velocities(...)
  Kernels-->>BodyBuffers: velocities updated
  end
Loading

Estimated code review effort

🎯 3 (Moderate) | ⏱️ ~30 minutes

Possibly related PRs

Suggested reviewers

  • eric-heiden
  • mmacklin

Pre-merge checks and finishing touches

❌ Failed checks (1 warning)
Check name Status Explanation Resolution
Out of Scope Changes Check ⚠️ Warning The pull request includes unrelated changes such as a new rigid restitution path in the XPBD solver and the skipping of a noncontiguous joint indexing test in the MuJoCo solver suite, neither of which pertains to runtime gravity modification. These functional modifications fall outside the scope defined by issue #597. Separate the restitution enhancement and test‐skipping changes into a distinct pull request or revert them here so that this PR remains focused solely on runtime gravity support.
✅ Passed checks (4 passed)
Check name Status Explanation
Description Check ✅ Passed Check skipped - CodeRabbit’s high-level summary is enabled.
Title Check ✅ Passed The title clearly describes the main enhancement—support for runtime randomization of gravity—and directly reflects the primary change set without introducing unrelated concepts. Although it includes the issue number and a “[Lab]” prefix, these do not obscure the intent or clarity of the title.
Linked Issues Check ✅ Passed All objectives from issue #597 are met: gravity is stored as a one‐element Warp array, Model.set_gravity allows runtime updates, SolverNotifyFlags.MODEL_PROPERTIES informs solvers of changes, all solver kernels now accept the updated gravity representation, and CUDA graph capture tests confirm the new behavior.
Docstring Coverage ✅ Passed No functions found in the changes. Docstring coverage check skipped.
✨ Finishing touches
  • 📝 Generate Docstrings
🧪 Generate unit tests
  • Create PR with unit tests
  • Post copyable unit tests in a 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

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

Signed-off-by: Viktor Reutskyy <vreutskyy@nvidia.com>

@eric-heiden eric-heiden left a comment

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I don't know if this is useful because when gravity is kept as a vec3 that means changing it won't update the simulation if it has been captured in a CUDA graph, you would need a Warp array to store gravity for that.
Also I don't think it makes sense to store it in the State, do we expect gravity to change over time? Why not keep in the Model?

@coderabbitai coderabbitai Bot left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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 (3)
newton/_src/solvers/solver.py (1)

47-49: Bug: gravity incorrectly gated by inv_mass sign (likely disables gravity).

wp.step(-inv_mass) returns 0 for positive inv_mass (typical dynamic particles), so gravity won’t apply. Use wp.nonzero(inv_mass) (or wp.step(inv_mass)) to enable gravity for non‑kinematic particles.

Apply this diff:

-    v1 = v0 + (f0 * inv_mass + gravity * wp.step(-inv_mass)) * dt
+    v1 = v0 + (f0 * inv_mass + gravity * wp.nonzero(inv_mass)) * dt
newton/_src/solvers/vbd/solver_vbd.py (2)

1219-1225: Apply gravity only to dynamic particles.

forward_step adds gravity unconditionally; pinned/kinematic particles (inv_mass == 0) will accelerate. Gate gravity on nonzero inv_mass (align with SolverBase kernels).

-    vel_new = vel[particle] + (gravity + external_force[particle] * inv_mass[particle]) * dt
+    vel_new = vel[particle] + (
+        gravity * wp.nonzero(inv_mass[particle]) + external_force[particle] * inv_mass[particle]
+    ) * dt

1248-1251: Same gating needed in penetration-free path.

Mirror the fix in forward_step_penetration_free.

-    vel_new = vel[particle_index] + (gravity + external_force[particle_index] * inv_mass[particle_index]) * dt
+    vel_new = vel[particle_index] + (
+        gravity * wp.nonzero(inv_mass[particle_index]) + external_force[particle_index] * inv_mass[particle_index]
+    ) * dt
🧹 Nitpick comments (4)
newton/_src/solvers/solver.py (1)

178-221: Avoid duplication: centralize gravity selection in SolverBase.

Multiple solvers repeat the same “state or model gravity” pattern. Consider a small helper (private) on SolverBase to reduce drift.

 class SolverBase:
@@
     def integrate_bodies(
@@
-        if model.body_count:
-            # Use gravity from state if available, otherwise use model gravity
-            gravity = state_in.gravity if state_in.gravity is not None else model.gravity
+        if model.body_count:
+            gravity = self._select_gravity(state_in)
@@
     def integrate_particles(
@@
-        if model.particle_count:
-            # Use gravity from state if available, otherwise use model gravity
-            gravity = state_in.gravity if state_in.gravity is not None else model.gravity
+        if model.particle_count:
+            gravity = self._select_gravity(state_in)
+
+    def _select_gravity(self, state: State) -> wp.vec3:
+        return state.gravity if state.gravity is not None else self.model.gravity

Also applies to: 231-257

newton/_src/solvers/vbd/solver_vbd.py (1)

2513-2516: Optional: persist gravity across state swaps.

If desired, set state_out.gravity = state_in.gravity at the end of each step to ease the common state_in, state_out = state_out, state_in pattern. Keep as optional to avoid surprising users who rely on explicit setting.

Would you like this persistence added (guarded behind a flag)?

Also applies to: 2677-2680

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

64-66: API addition looks good; consider a brief note on state swapping.

Add a short note that when swapping states each step, gravity must be set on both states to persist.

newton/tests/test_runtime_gravity.py (1)

160-165: Guard MuJoCo tests when MuJoCo/mujoco_warp are not installed.

CI/workstations without MuJoCo will hard-fail. Skip these tests conditionally.

Apply this diff:

@@
-solvers_bodies = {
-    "xpbd": lambda model: SolverXPBD(model),
-    "semi_implicit": lambda model: SolverSemiImplicit(model),
-    "mujoco_cpu": lambda model: newton.solvers.SolverMuJoCo(model, use_mujoco_cpu=True, update_data_interval=0),
-    "mujoco_warp": lambda model: newton.solvers.SolverMuJoCo(model, use_mujoco_cpu=False, update_data_interval=0),
-}
+solvers_bodies = {
+    "xpbd": lambda model: SolverXPBD(model),
+    "semi_implicit": lambda model: SolverSemiImplicit(model),
+}
+try:
+    import mujoco  # noqa: F401
+    import mujoco_warp  # noqa: F401
+    MUJOCO_AVAILABLE = True
+except Exception:
+    MUJOCO_AVAILABLE = False
+if MUJOCO_AVAILABLE:
+    solvers_bodies.update(
+        {
+            "mujoco_cpu": lambda model: newton.solvers.SolverMuJoCo(
+                model, use_mujoco_cpu=True, update_data_interval=0
+            ),
+            "mujoco_warp": lambda model: newton.solvers.SolverMuJoCo(
+                model, use_mujoco_cpu=False, update_data_interval=0
+            ),
+        }
+    )
@@
-    for solver_name, solver_fn in solvers_bodies.items():
+    for solver_name, solver_fn in solvers_bodies.items():
         # Skip CPU MuJoCo on CUDA devices
-        if device.is_cuda and solver_name == "mujoco_cpu":
+        if device.is_cuda and solver_name == "mujoco_cpu":
             continue

Also applies to: 167-191

📜 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 791a0ee and e737fa8.

📒 Files selected for processing (9)
  • newton/_src/sim/state.py (2 hunks)
  • newton/_src/solvers/featherstone/solver_featherstone.py (2 hunks)
  • newton/_src/solvers/implicit_mpm/solver_implicit_mpm.py (2 hunks)
  • newton/_src/solvers/mujoco/solver_mujoco.py (1 hunks)
  • newton/_src/solvers/solver.py (4 hunks)
  • newton/_src/solvers/style3d/solver_style3d.py (1 hunks)
  • newton/_src/solvers/vbd/solver_vbd.py (2 hunks)
  • newton/_src/solvers/xpbd/solver_xpbd.py (2 hunks)
  • newton/tests/test_runtime_gravity.py (1 hunks)
🧰 Additional context used
🧠 Learnings (2)
📚 Learning: 2025-08-12T05:17:34.423Z
Learnt from: shi-eric
PR: newton-physics/newton#521
File: newton/examples/example_cloth_hanging.py:36-36
Timestamp: 2025-08-12T05:17:34.423Z
Learning: The STYLE3D solver in Newton is a new solver that did not exist in warp.sim, so it does not require any migration documentation.

Applied to files:

  • newton/_src/solvers/style3d/solver_style3d.py
📚 Learning: 2025-08-19T12:27:30.629Z
Learnt from: preist-nvidia
PR: newton-physics/newton#579
File: newton/examples/example_mujoco.py:350-354
Timestamp: 2025-08-19T12:27:30.629Z
Learning: In Newton examples, Example.use_mujoco is a high-level attribute that controls whether to use MuJoCo solver vs other solvers (like XPBD), while SolverMuJoCo.use_mujoco_cpu controls the backend within MuJoCo (CPU vs Warp). These are separate concepts serving different purposes - the PR rename only applies to the solver parameter, not the Example class attributes.

Applied to files:

  • newton/_src/solvers/mujoco/solver_mujoco.py
🧬 Code graph analysis (2)
newton/_src/solvers/style3d/solver_style3d.py (1)
newton/_src/solvers/style3d/kernels.py (1)
  • init_step_kernel (162-193)
newton/tests/test_runtime_gravity.py (5)
newton/_src/solvers/xpbd/solver_xpbd.py (1)
  • step (206-662)
newton/tests/unittest_utils.py (2)
  • add_function_test (280-299)
  • get_test_devices (97-131)
newton/_src/sim/state.py (2)
  • clear_forces (67-78)
  • set_gravity (80-91)
newton/_src/solvers/featherstone/solver_featherstone.py (1)
  • step (278-669)
newton/_src/solvers/mujoco/solver_mujoco.py (1)
  • step (1330-1361)
⏰ 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 (Pull Request)
  • GitHub Check: Run GPU Unit Tests on AWS EC2 (Pull Request)
  • GitHub Check: run-newton-tests / newton-unittests (ubuntu-latest)
  • GitHub Check: run-newton-tests / newton-unittests (windows-latest)
🔇 Additional comments (9)
newton/_src/solvers/style3d/solver_style3d.py (1)

116-118: Gravity override correctly sourced from State.

Looks good and consistent with the PR’s runtime override model.

newton/_src/solvers/solver.py (2)

198-200: Per-call gravity selection for bodies: LGTM.

Also applies to: 213-213


238-240: Per-call gravity selection for particles: LGTM.

Also applies to: 250-250

newton/_src/solvers/vbd/solver_vbd.py (2)

2513-2515: Per-call gravity selection/pass-through: LGTM.

Also applies to: 2520-2520


2677-2679: Per-call gravity selection/pass-through (penetration-free): LGTM.

Also applies to: 2684-2684

newton/_src/solvers/featherstone/solver_featherstone.py (1)

365-366: Gravity override plumbed through ID: LGTM.

Also applies to: 383-384

newton/_src/solvers/implicit_mpm/solver_implicit_mpm.py (1)

1457-1459: MPM: gravity override integrated into velocity RHS: LGTM.

Also applies to: 1468-1469

newton/_src/solvers/xpbd/solver_xpbd.py (1)

616-618: Cast gravity to wp.vec3f and ensure integrate paths honor state_in.gravity

Explicitly convert the selected gravity to wp.vec3f before passing to kernels; audit integrate_particles/integrate_bodies code paths to ensure they also respect state_in.gravity.

Location: newton/_src/solvers/xpbd/solver_xpbd.py lines 616–618 (also check 643–645).

@@
-                    # use gravity from state if available, otherwise use model gravity
-                    gravity = state_in.gravity if state_in.gravity is not None else model.gravity
+                    # use gravity from state if available, otherwise use model gravity
+                    _g = state_in.gravity if state_in.gravity is not None else model.gravity
+                    gravity = wp.vec3f(float(_g[0]), float(_g[1]), float(_g[2]))
newton/_src/solvers/mujoco/solver_mujoco.py (1)

1331-1338: Make gravity assignment type-safe for both backends (CPU and Warp)

state_in.gravity is a wp.vec3; convert to matching dtypes before assigning to avoid dtype/runtime errors.

File: newton/_src/solvers/mujoco/solver_mujoco.py (lines 1331-1338)

@@
-        # Update gravity if it's set in state
-        if state_in.gravity is not None:
-            if self.use_mujoco_cpu:
-                self.mj_model.opt.gravity[:] = state_in.gravity
-            else:
-                # MuJoCo Warp uses a warp array for gravity
-                self.mjw_model.opt.gravity.fill_(state_in.gravity)
+        # Update gravity if it's set in state
+        if state_in.gravity is not None:
+            g = state_in.gravity
+            if self.use_mujoco_cpu:
+                # match MuJoCo's underlying dtype (usually float64)
+                self.mj_model.opt.gravity[:] = np.asarray(
+                    [float(g[0]), float(g[1]), float(g[2])],
+                    dtype=self.mj_model.opt.gravity.dtype,
+                )
+            else:
+                # MuJoCo Warp uses a vec3 array; fill_ expects a vec3f value
+                self.mjw_model.opt.gravity.fill_(wp.vec3f(float(g[0]), float(g[1]), float(g[2])))

Runtime sanity checks not possible here — run tests/test_runtime_gravity.py (both mujoco_cpu and mujoco_warp) to confirm no dtype warnings and that gravity changes acceleration.

Comment thread newton/_src/sim/state.py Outdated

@coderabbitai coderabbitai Bot left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Actionable comments posted: 2

Caution

Some comments are outside the diff and can’t be posted inline due to platform limitations.

⚠️ Outside diff range comments (1)
newton/_src/solvers/solver.py (1)

48-51: Likely bug: gravity masked for dynamic particles.

v1 = v0 + (f0 * inv_mass + gravity * wp.step(-inv_mass)) * dt applies gravity only when inv_mass <= 0. For dynamic particles (inv_mass > 0), this zeros gravity. Use wp.nonzero(inv_mass) (as done in rigid integration) or remove the mask.

Apply this diff:

-    v1 = v0 + (f0 * inv_mass + gravity * wp.step(-inv_mass)) * dt
+    v1 = v0 + (f0 * inv_mass + gravity * wp.nonzero(inv_mass)) * dt
🧹 Nitpick comments (3)
newton/_src/solvers/vbd/solver_vbd.py (2)

2513-2515: Gravity override wired correctly; consider explicit wp.vec3 to avoid dtype ambiguity.

The fallback to model gravity is right and the kernel argument ordering matches forward_step(dt, gravity, ...). To prevent any device/dtype surprises when model.gravity is a NumPy array, consider explicitly casting once at call sites.

Apply:

-        gravity = state_in.gravity if state_in.gravity is not None else model.gravity
+        gravity = state_in.gravity if state_in.gravity is not None else wp.vec3(*model.gravity.tolist())

Also applies to: 2520-2520


2677-2679: Same here for penetration‑free path.

Mirror the explicit wp.vec3 cast to keep both paths consistent.

-        gravity = state_in.gravity if state_in.gravity is not None else model.gravity
+        gravity = state_in.gravity if state_in.gravity is not None else wp.vec3(*model.gravity.tolist())

Also applies to: 2684-2684

newton/tests/test_runtime_gravity.py (1)

29-79: Particle test flow looks good and exercises runtime changes across signs.

As-is this validates the override well. Consider adding a VBD coverage variant since this PR modifies VBD; VBD requires coloring. If you want, I can provide a small patch that does builder.color() and adds SolverVBD to solvers_particles.

📜 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 791a0ee and e737fa8.

📒 Files selected for processing (9)
  • newton/_src/sim/state.py (2 hunks)
  • newton/_src/solvers/featherstone/solver_featherstone.py (2 hunks)
  • newton/_src/solvers/implicit_mpm/solver_implicit_mpm.py (2 hunks)
  • newton/_src/solvers/mujoco/solver_mujoco.py (1 hunks)
  • newton/_src/solvers/solver.py (4 hunks)
  • newton/_src/solvers/style3d/solver_style3d.py (1 hunks)
  • newton/_src/solvers/vbd/solver_vbd.py (2 hunks)
  • newton/_src/solvers/xpbd/solver_xpbd.py (2 hunks)
  • newton/tests/test_runtime_gravity.py (1 hunks)
🧰 Additional context used
🧬 Code graph analysis (2)
newton/tests/test_runtime_gravity.py (2)
newton/tests/unittest_utils.py (2)
  • add_function_test (280-299)
  • get_test_devices (97-131)
newton/_src/sim/state.py (2)
  • clear_forces (67-78)
  • set_gravity (80-91)
newton/_src/solvers/style3d/solver_style3d.py (1)
newton/_src/solvers/style3d/kernels.py (1)
  • init_step_kernel (162-193)
⏰ 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 (Pull Request)
  • GitHub Check: Run GPU Unit Tests on AWS EC2 (Pull Request)
🔇 Additional comments (10)
newton/_src/sim/state.py (1)

64-66: LGTM: per-state gravity override added.

Attribute + docstring are clear and align with solver usage (fallback to model.gravity).

newton/_src/solvers/featherstone/solver_featherstone.py (1)

365-366: LGTM: consistent gravity selection and propagation.

Correctly prefers state_in.gravity, falling back to model.gravity, and wires through to eval_rigid_id.

Also applies to: 383-384

newton/_src/solvers/implicit_mpm/solver_implicit_mpm.py (1)

1457-1469: LGTM: runtime gravity plumbed into integrand.

Local gravity selection and passing to integrate_velocity match the new State override.

newton/_src/solvers/style3d/solver_style3d.py (1)

116-125: LGTM: gravity override used in init_step.

Consistent with kernel signature; no further issues.

newton/_src/solvers/solver.py (2)

199-214: LGTM: bodies integration honors state-level gravity.


239-251: LGTM: particles integration honors state-level gravity.

newton/_src/solvers/xpbd/solver_xpbd.py (1)

616-644: LGTM: restitution path now respects per-state gravity.

Local gravity selection and threading into apply_rigid_restitution are consistent.

newton/tests/test_runtime_gravity.py (3)

160-165: MuJoCo parametrization guard is thoughtful; double-check CPU/GPU matrix.

You skip MuJoCo-CPU on CUDA devices. If any CI shard runs CPU-only, mujoco_warp should still behave; if not, add a symmetric skip for warp-on-CPU.

Also applies to: 182-189


124-150: Fallback test is appropriate.

Confirms None → model.gravity path. No change needed.


54-56: Replace PEP-604 unions used in isinstance(...) with tuple-of-types (fix TypeError)

Using expressions like isinstance(x, tuple | list) will raise TypeError at runtime — change to isinstance(x, (tuple, list)).

  • Apply in newton/_src/sim/state.py:
-        if isinstance(gravity, tuple | list):
+        if isinstance(gravity, (tuple, list)):
             self.gravity = wp.vec3(gravity[0], gravity[1], gravity[2])
         else:
             self.gravity = gravity
  • Also update:
    • newton/_src/viewer/viewer.py — change
      if isinstance(value, tuple | list | np.ndarray)
      to
      if isinstance(value, (tuple, list, np.ndarray))
    • newton/_src/viewer/viewer_gl.py — change
      if isinstance(colors, tuple | list)
      to
      if isinstance(colors, (tuple, list))
⛔ Skipped due to learnings
Learnt from: dylanturpin
PR: newton-physics/newton#635
File: docs/generate_api.py:64-66
Timestamp: 2025-08-25T20:18:47.535Z
Learning: PEP 604 (Python 3.10+) allows using the union operator | as the second argument to isinstance() and issubclass(). For example, isinstance(obj, int | float | str) is valid syntax in Python 3.10+. The traditional tuple form isinstance(obj, (int, float, str)) still works but is not required.

Comment thread newton/_src/sim/state.py Outdated
Comment thread newton/_src/solvers/mujoco/solver_mujoco.py Outdated
@vreutskyy

Copy link
Copy Markdown
Member Author

@eric-heiden
As I understand it, the way to change the gravity in runtime is what was requested by #597. Or is it about something else?

@eric-heiden

Copy link
Copy Markdown
Member

This is supposed to be another model update, as part of solver.notify_model_changed(SolverNotifyFlags.MODEL_PROPERTIES).
I think we need another flag like SolverNotifyFlags.MODEL_PROPERTIES for this.
Also, I don't think the current approach works when the MuJoCo step has been captured in a graph, can you test this? We may have to move Model.gravity to a warp.array(dtype=vec3) and potentially leave the option open to have per-environment gravity in Newton eventually, but that part is not needed right now (can be just a 1-element Warp array for now).

Signed-off-by: Viktor Reutskyy <vreutskyy@nvidia.com>

@coderabbitai coderabbitai Bot left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Actionable comments posted: 5

🧹 Nitpick comments (6)
newton/_src/solvers/xpbd/kernels.py (1)

2407-2423: Tiny load reduction: cache gravity[0] locally.

Avoid two global loads by reading gravity[0] once.

Apply:

@@
-    if body_a >= 0:
-        v_a = velocity_at_point(body_qd_prev[body_a], r_a) + gravity[0] * dt
+    g = gravity[0]
+    if body_a >= 0:
+        v_a = velocity_at_point(body_qd_prev[body_a], r_a) + g * dt
@@
-    if body_b >= 0:
-        v_b = velocity_at_point(body_qd_prev[body_b], r_b) + gravity[0] * dt
+    if body_b >= 0:
+        v_b = velocity_at_point(body_qd_prev[body_b], r_b) + g * dt
newton/_src/sim/builder.py (1)

4408-4414: Model gravity as 1‑element array: LGTM; consider not tracking grads.

No strong need for gradients on gravity; disabling saves memory/overhead.

Apply:

             m.gravity = wp.array(
                 [wp.vec3(*(g * self.gravity for g in self.up_vector))],
                 dtype=wp.vec3,
                 device=device,
-                requires_grad=requires_grad,
+                requires_grad=False,
             )
newton/_src/solvers/solver.py (2)

197-216: Prefer state gravity override (fallback to model).

To honor runtime per‑state gravity, pass state_in.gravity when present.

-            wp.launch(
+            gravity_arr = getattr(state_in, "gravity", None) or model.gravity
+            wp.launch(
                 kernel=integrate_bodies,
                 dim=model.body_count,
                 inputs=[
                     state_in.body_q,
                     state_in.body_qd,
                     state_in.body_f,
                     model.body_com,
                     model.body_mass,
                     model.body_inertia,
                     model.body_inv_mass,
                     model.body_inv_inertia,
-                    model.gravity,
+                    gravity_arr,
                     angular_damping,
                     dt,
                 ],

234-250: Apply state gravity override for particles too.

-            wp.launch(
+            gravity_arr = getattr(state_in, "gravity", None) or model.gravity
+            wp.launch(
                 kernel=integrate_particles,
                 dim=model.particle_count,
                 inputs=[
                     state_in.particle_q,
                     state_in.particle_qd,
                     state_in.particle_f,
                     model.particle_inv_mass,
                     model.particle_flags,
-                    model.gravity,
+                    gravity_arr,
                     dt,
                     model.particle_max_velocity,
                 ],
newton/_src/sim/model.py (1)

523-526: Optional: shorten error text (ruff TRY003).

Consider a shorter message, e.g., “Gravity not initialized; finalize Model with ModelBuilder first.”

newton/tests/test_runtime_gravity.py (1)

218-221: Re-raise original exception to preserve traceback.

-    except Exception as e:
+    except Exception:
         # Make sure to end capture if something goes wrong
         wp.capture_end()
-        raise e
+        raise
📜 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 e737fa8 and dd041b9.

📒 Files selected for processing (13)
  • newton/_src/sim/builder.py (1 hunks)
  • newton/_src/sim/model.py (2 hunks)
  • newton/_src/solvers/featherstone/kernels.py (3 hunks)
  • newton/_src/solvers/featherstone/solver_featherstone.py (1 hunks)
  • newton/_src/solvers/flags.py (1 hunks)
  • newton/_src/solvers/implicit_mpm/solver_implicit_mpm.py (1 hunks)
  • newton/_src/solvers/mujoco/solver_mujoco.py (4 hunks)
  • newton/_src/solvers/solver.py (6 hunks)
  • newton/_src/solvers/style3d/kernels.py (2 hunks)
  • newton/_src/solvers/vbd/solver_vbd.py (3 hunks)
  • newton/_src/solvers/xpbd/kernels.py (3 hunks)
  • newton/_src/solvers/xpbd/solver_xpbd.py (1 hunks)
  • newton/tests/test_runtime_gravity.py (1 hunks)
🚧 Files skipped from review as they are similar to previous changes (4)
  • newton/_src/solvers/vbd/solver_vbd.py
  • newton/_src/solvers/implicit_mpm/solver_implicit_mpm.py
  • newton/_src/solvers/featherstone/solver_featherstone.py
  • newton/_src/solvers/xpbd/solver_xpbd.py
🧰 Additional context used
🧠 Learnings (5)
📚 Learning: 2025-08-12T17:51:37.474Z
Learnt from: nvlukasz
PR: newton-physics/newton#519
File: newton/geometry.py:16-22
Timestamp: 2025-08-12T17:51:37.474Z
Learning: In the Newton public API refactor, common geometry symbols like ParticleFlags and ShapeFlags are now exposed at the top level in newton/__init__.py rather than through newton.geometry. The newton/geometry.py module is intentionally focused only on broad-phase collision detection classes (BroadPhaseAllPairs, BroadPhaseExplicit, BroadPhaseSAP).

Applied to files:

  • newton/_src/solvers/flags.py
  • newton/_src/solvers/solver.py
📚 Learning: 2025-09-22T21:08:31.878Z
Learnt from: dylanturpin
PR: newton-physics/newton#806
File: newton/examples/ik/example_ik_franka.py:121-123
Timestamp: 2025-09-22T21:08:31.878Z
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/solvers/featherstone/kernels.py
📚 Learning: 2025-09-22T21:08:31.878Z
Learnt from: dylanturpin
PR: newton-physics/newton#806
File: newton/examples/ik/example_ik_franka.py:121-123
Timestamp: 2025-09-22T21:08:31.878Z
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/solvers/featherstone/kernels.py
📚 Learning: 2025-08-18T15:56:26.587Z
Learnt from: adenzler-nvidia
PR: newton-physics/newton#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/_src/solvers/featherstone/kernels.py
📚 Learning: 2025-08-12T05:17:34.423Z
Learnt from: shi-eric
PR: newton-physics/newton#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.

Applied to files:

  • newton/_src/solvers/mujoco/solver_mujoco.py
🧬 Code graph analysis (5)
newton/_src/solvers/xpbd/kernels.py (1)
newton/_src/core/spatial.py (1)
  • velocity_at_point (22-34)
newton/_src/sim/builder.py (4)
newton/_src/solvers/solver.py (1)
  • device (169-176)
newton/_src/sim/collide.py (1)
  • device (331-338)
newton/_src/sim/contacts.py (1)
  • device (89-93)
newton/_src/sim/state.py (1)
  • requires_grad (78-84)
newton/_src/solvers/mujoco/solver_mujoco.py (2)
newton/_src/solvers/flags.py (1)
  • SolverNotifyFlags (22-56)
newton/_src/solvers/solver.py (1)
  • device (169-176)
newton/_src/solvers/solver.py (3)
newton/_src/solvers/featherstone/solver_featherstone.py (1)
  • step (278-667)
newton/_src/solvers/xpbd/solver_xpbd.py (1)
  • step (206-660)
newton/_src/solvers/style3d/solver_style3d.py (1)
  • step (112-275)
newton/tests/test_runtime_gravity.py (5)
newton/tests/unittest_utils.py (2)
  • add_function_test (280-299)
  • get_test_devices (97-131)
newton/_src/solvers/solver.py (2)
  • device (169-176)
  • notify_model_changed (267-294)
newton/_src/sim/builder.py (2)
  • ModelBuilder (67-4471)
  • finalize (4058-4416)
newton/_src/sim/model.py (1)
  • set_gravity (511-531)
newton/_src/solvers/mujoco/solver_mujoco.py (1)
  • notify_model_changed (1416-1426)
🪛 Ruff (0.13.1)
newton/_src/sim/model.py

524-526: Avoid specifying long messages outside the exception class

(TRY003)

newton/tests/test_runtime_gravity.py

221-221: Use raise without specifying exception name

Remove exception name

(TRY201)

⏰ 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 (Pull Request)
🔇 Additional comments (8)
newton/_src/solvers/xpbd/kernels.py (1)

2328-2329: Gravity as wp.array(vec3): API change looks good; verify all launch sites.

Switching to a 1‑element wp.array(vec3) is consistent with the new model/state gravity representation.

Run to find call sites and quickly inspect their argument for gravity (should pass a wp.array(dtype=wp.vec3), e.g., model.gravity or state.gravity):

newton/_src/solvers/flags.py (1)

45-55: Approve: MODEL_PROPERTIES flag added and wired through solvers.
Bits don't collide and ALL includes the new flag; verified usage in newton/_src/solvers/mujoco/solver_mujoco.py (handles MODEL_PROPERTIES), newton/tests/test_runtime_gravity.py (calls solver.notify_model_changed), and newton/_src/sim/model.py (docstring advises notifying solvers).

newton/_src/solvers/style3d/kernels.py (1)

164-165: Signature update to gravity array — OK.

init_step_kernel is launched with self.model.gravity (wp.array(dtype=wp.vec3)) in newton/_src/solvers/style3d/solver_style3d.py (wp.launch inputs include self.model.gravity).

newton/_src/solvers/featherstone/kernels.py (1)

681-688: Gravity-as-array usage looks correct.

Passing gravity as a 1-element wp.array and indexing gravity[0] is consistent and safe here. Call sites must pass dtype=wp.vec3 arrays. No issues found.

Please confirm all invocations of eval_rigid_id pass a wp.array(dtype=wp.vec3) and not a bare vec3.

Also applies to: 736-739, 766-776

newton/_src/solvers/mujoco/solver_mujoco.py (2)

1425-1427: Good: handles MODEL_PROPERTIES updates.

Hooking MODEL_PROPERTIES into notify path ensures gravity sync without recreating the solver.


1886-1886: OK: CPU spec gravity initialization.

Using model.gravity.numpy()[0] to seed spec.option.gravity is robust.

newton/_src/solvers/solver.py (1)

286-286: Docs updated appropriately.

Adding MODEL_PROPERTIES to the docstring aligns with the new notify flow.

newton/tests/test_runtime_gravity.py (1)

197-201: Nice coverage across solvers and devices.

Good assertions and parameterization; exercises zero/upward gravity and CUDA graph capture paths.

Also applies to: 224-237, 252-263

Comment thread newton/_src/sim/model.py Outdated
Comment thread newton/_src/solvers/mujoco/solver_mujoco.py
Comment thread newton/_src/solvers/mujoco/solver_mujoco.py
Comment thread newton/_src/solvers/solver.py
Comment thread newton/_src/solvers/style3d/kernels.py
vreutskyy and others added 2 commits September 23, 2025 12:19
Signed-off-by: Viktor Reutskyy <vreutskyy@nvidia.com>

@coderabbitai coderabbitai Bot left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Actionable comments posted: 0

Caution

Some comments are outside the diff and can’t be posted inline due to platform limitations.

⚠️ Outside diff range comments (3)
newton/tests/test_cloth.py (3)

1049-1057: Fix gravity usage to handle wp.array(vec3[1]) → NumPy vector conversion.

Root cause of CI failure: example.model.gravity is now a 1‑element wp.array(vec3). Converting it with np.array(...) yields a shape mismatch ((1,3) vs (N,3)) and triggers invalid indexing on wp.array elsewhere.

Update this block to extract a 3‑vector first:

-    gravity_vector = np.array(example.model.gravity)
-    gravity_direction = gravity_vector / np.linalg.norm(gravity_vector)
-    gravity_displacement_per_vertex = np.dot(position_diff, gravity_direction)
-    gravity_component_per_vertex = gravity_displacement_per_vertex[:, np.newaxis] * gravity_direction[np.newaxis, :]
+    gravity_vec = example.model.gravity.numpy()[0]
+    gravity_dir = gravity_vec / np.linalg.norm(gravity_vec)
+    gravity_displacement_per_vertex = np.dot(position_diff, gravity_dir)
+    gravity_component_per_vertex = gravity_displacement_per_vertex[:, np.newaxis] * gravity_dir[np.newaxis, :]

1107-1114: Replace direct np.array(gravity) usage to avoid wp.array indexing error.

This fixes “Item indexing is not supported on wp.array objects” and aligns shapes for matmul.

-    gravity = np.array(example.model.gravity)
+    gravity_vec = example.model.gravity.numpy()[0]
     diff = final_pos - initial_pos
-    vertical_translation_norm = diff @ gravity[..., None] / (np.linalg.norm(gravity) ** 2)
-    horizontal_move = diff - (vertical_translation_norm * gravity)
+    vertical_translation_norm = diff @ (gravity_vec / np.linalg.norm(gravity_vec))[..., None]
+    horizontal_move = diff - (vertical_translation_norm * gravity_vec)

1049-1057: Replace direct uses of model.gravity in tests with model.gravity.numpy()[0]

  • newton/tests/test_cloth.py:1050 — replace
    gravity_vector = np.array(example.model.gravity)
    with
    gravity_vector = example.model.gravity.numpy()[0]

  • newton/tests/test_cloth.py:1107 — replace
    gravity = np.array(example.model.gravity)
    with
    gravity = example.model.gravity.numpy()[0]

🧹 Nitpick comments (3)
newton/_src/solvers/mujoco/solver_mujoco.py (2)

1069-1077: Explicitly cast vec3 -> vec3f in kernel to avoid implicit dtype conversion.

Avoid relying on implicit casting from wp.vec3 to wp.vec3f. Cast once per-thread and write the result.

 @wp.kernel
 def update_model_properties_kernel(
     # Newton model properties
     gravity_src: wp.array(dtype=wp.vec3),
     # MuJoCo model properties
     gravity_dst: wp.array(dtype=wp.vec3f),
 ):
-    world_idx = wp.tid()
-    gravity_dst[world_idx] = gravity_src[0]
+    wid = wp.tid()
+    g = gravity_src[0]
+    gravity_dst[wid] = wp.vec3f(g[0], g[1], g[2])

2752-2769: Launch gravity update without relying on mjw_data and size by dst shape.

  • Gate on hasattr(self, "mjw_data") is unnecessary and skips updates before put_data(). Use the destination’s first-dimension length instead.
  • Keeps behavior correct for single/multi-world and works pre/post data creation.
 def update_model_properties(self):
     """Update model properties including gravity in the MuJoCo model."""
     if self.use_mujoco_cpu:
         self.mj_model.opt.gravity[:] = np.array([*self.model.gravity.numpy()[0]])
     else:
-        if hasattr(self, "mjw_data"):
-            wp.launch(
-                kernel=update_model_properties_kernel,
-                dim=self.mjw_data.nworld,
-                inputs=[
-                    self.model.gravity,
-                ],
-                outputs=[
-                    self.mjw_model.opt.gravity,
-                ],
-                device=self.model.device,
-            )
+        nworld = self.mjw_model.opt.gravity.shape[0]
+        wp.launch(
+            kernel=update_model_properties_kernel,
+            dim=nworld,
+            inputs=[self.model.gravity],
+            outputs=[self.mjw_model.opt.gravity],
+            device=self.model.device,
+        )
newton/tests/test_cloth.py (1)

1049-1057: Optional: provide a convenience getter to simplify tests.

Consider adding Model.get_gravity() returning a tuple or NumPy 3‑vector to avoid repeated .numpy()[0] patterns in client code/tests.

📜 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 dd041b9 and 75a5219.

📒 Files selected for processing (4)
  • newton/_src/solvers/mujoco/solver_mujoco.py (4 hunks)
  • newton/examples/cloth/example_cloth_franka.py (2 hunks)
  • newton/examples/cloth/example_cloth_style3d.py (1 hunks)
  • newton/tests/test_cloth.py (2 hunks)
🧰 Additional context used
🧠 Learnings (1)
📚 Learning: 2025-08-12T05:17:34.423Z
Learnt from: shi-eric
PR: newton-physics/newton#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.

Applied to files:

  • newton/_src/solvers/mujoco/solver_mujoco.py
🧬 Code graph analysis (4)
newton/_src/solvers/mujoco/solver_mujoco.py (2)
newton/_src/solvers/flags.py (1)
  • SolverNotifyFlags (22-56)
newton/_src/solvers/solver.py (1)
  • device (169-176)
newton/tests/test_cloth.py (1)
newton/_src/sim/model.py (1)
  • set_gravity (511-531)
newton/examples/cloth/example_cloth_style3d.py (1)
newton/_src/sim/model.py (1)
  • set_gravity (511-531)
newton/examples/cloth/example_cloth_franka.py (1)
newton/_src/sim/model.py (1)
  • set_gravity (511-531)
🪛 GitHub Actions: Pull Request
newton/tests/test_cloth.py

[error] 1054-1054: ValueError: shapes (14,3) and (1,3) not aligned during gravity displacement calculation in test_cloth_free_fall_with_internal_forces_and_damping (gravity_displacement_per_vertex = np.dot(position_diff, gravity_direction)).


[error] 1054-1054: Originating failure occurs for test_cloth_free_fall_with_internal_forces_and_damping_vbd_cpu.


[error] 1054-1054: Originating failure occurs for test_cloth_free_fall_with_internal_forces_and_damping_xpbd_cpu.

🪛 GitHub Actions: GPU Unit Tests on AWS EC2
newton/tests/test_cloth.py

[error] 1107-1107: RuntimeError: Item indexing is not supported on wp.array objects.

⏰ 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 (Pull Request)
🔇 Additional comments (6)
newton/_src/solvers/mujoco/solver_mujoco.py (2)

1425-1426: LGTM: model property change hook added.

Registering MODEL_PROPERTIES and delegating to update_model_properties() looks correct.


1886-1886: LGTM: initialize MuJoCo gravity from model gravity.

Host-side conversion is fine here; shape and dtype are appropriate for MjSpec.

newton/examples/cloth/example_cloth_style3d.py (1)

127-127: Good API migration to runtime gravity setter.

Setting gravity via set_gravity((0.0, 0.0, -9.81)) prior to solver creation is correct and future‑proof.

newton/examples/cloth/example_cloth_franka.py (1)

521-534: Setter usage looks correct; consider notifying solvers only if they cache model props.

Using set_gravity when toggling gravity during the substep loop is fine. If any solver involved caches model properties (e.g., MuJoCo), call solver.notify_model_changed(SolverNotifyFlags.MODEL_PROPERTIES) after updates. For VBD/Featherstone reading the array each step, this is usually not required.

newton/tests/test_cloth.py (2)

489-489: Correct migration to the public gravity setter.

Using self.model.set_gravity((0.0, -1000.0, 0.0)) aligns tests with the new API.


835-835: Correct migration to the public gravity setter.

Passing a tuple with conditional Y component is correct.

vreutskyy and others added 3 commits September 23, 2025 17:06
Signed-off-by: Viktor Reutskyy <vreutskyy@nvidia.com>
…vity

Signed-off-by: Viktor Reutskyy <vreutskyy@nvidia.com>

@coderabbitai coderabbitai Bot left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Actionable comments posted: 1

🧹 Nitpick comments (3)
newton/examples/robot/example_robot_policy.py (2)

253-254: Clarify gravity comment; avoid leaving dead code.

The issue isn’t the vec3 type—it’s that builder-time gravity won’t reflect runtime updates or CUDA graph replay. Suggest tightening the comment and removing the commented-out assignment.

-        # builder's gravity isn't a vec3. use model.set_gravity()
-        # builder.gravity = wp.vec3(0.0, 0.0, -9.81)
+        # Gravity is a runtime model property; set it via model.set_gravity() after finalize.
+        # Avoid setting builder.gravity here—runtime updates (and CUDA graph replay) won't see it.

269-270: Good: set gravity via model after finalize; add a note about runtime updates.

This is the correct placement (before solver creation). Add a brief note to show how to notify solvers if gravity changes later at runtime.

+        # Initial gravity (Z-up). If you modify gravity later at runtime, also call:
+        #   self.solver.notify_model_changed(SolverNotifyFlags.MODEL_PROPERTIES)
         self.model.set_gravity((0.0, 0.0, -9.81))
newton/_src/solvers/implicit_mpm/solver_implicit_mpm.py (1)

1211-1215: Acknowledge MODEL_PROPERTIES in notify hook (no-op now, future‑proof)

Handle the new flag explicitly to keep solver hooks consistent across the codebase and future-proof for per-env gravity without changing behavior.

Apply this diff:

     @override
     def notify_model_changed(self, flags: int):
         if flags & newton.SolverNotifyFlags.PARTICLE_PROPERTIES:
             self.mpm_model.notify_particle_material_changed()
+        if flags & newton.SolverNotifyFlags.MODEL_PROPERTIES:
+            # Gravity is read from `model.gravity` at assembly time; nothing to rebuild.
+            # Keep this branch to acknowledge the flag and for future extensions.
+            pass
📜 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 475ccd3 and e186e28.

📒 Files selected for processing (3)
  • newton/_src/solvers/implicit_mpm/solver_implicit_mpm.py (2 hunks)
  • newton/examples/mpm/example_mpm_granular.py (1 hunks)
  • newton/examples/robot/example_robot_policy.py (2 hunks)
🧰 Additional context used
🧠 Learnings (2)
📚 Learning: 2025-08-27T19:05:44.697Z
Learnt from: Milad-Rakhsha-NV
PR: newton-physics/newton#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/examples/robot/example_robot_policy.py
📚 Learning: 2025-09-18T07:05:56.836Z
Learnt from: gdaviet
PR: newton-physics/newton#750
File: newton/_src/solvers/implicit_mpm/solve_rheology.py:969-986
Timestamp: 2025-09-18T07:05:56.836Z
Learning: In newton/_src/solvers/implicit_mpm/solve_rheology.py, transposed_strain_mat parameter cannot be None - the type signature was corrected to reflect this guarantee, eliminating the need for None checks when accessing transposed_strain_mat.offsets.

Applied to files:

  • newton/_src/solvers/implicit_mpm/solver_implicit_mpm.py
🧬 Code graph analysis (2)
newton/examples/robot/example_robot_policy.py (1)
newton/_src/sim/model.py (1)
  • set_gravity (511-531)
newton/examples/mpm/example_mpm_granular.py (2)
newton/_src/sim/builder.py (1)
  • finalize (4058-4416)
newton/_src/sim/model.py (1)
  • set_gravity (511-531)
⏰ 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 Unit Tests on AWS EC2 (Pull Request)
  • GitHub Check: Run GPU Benchmarks (Pull Request)
🔇 Additional comments (4)
newton/_src/solvers/implicit_mpm/solver_implicit_mpm.py (2)

124-124: Adopt gravity as wp.array(vec3): LGTM

Signature change to gravity: wp.array(dtype=wp.vec3) is correct and aligns with runtime updates/CUDA-graph requirements.


131-135: Use gravity[0] in momentum RHS: LGTM

particle_density * (vel + dt * gravity[0]) correctly injects gravitational impulse while preserving existing APIC flow.

newton/examples/mpm/example_mpm_granular.py (2)

52-53: LGTM: Switching to Model.set_gravity() is correct

Avoids the incorrect vec3 assignment on the builder (which expects a scalar). Good move.


57-57: Initialize gravity before solver creation (correct placement)

Setting gravity immediately after finalize and before constructing the solver is appropriate. If you later randomize gravity at runtime, remember to call solver.notify_model_changed(SolverNotifyFlags.MODEL_PROPERTIES).

Comment thread newton/examples/mpm/example_mpm_granular.py
vreutskyy and others added 2 commits September 24, 2025 22:56
Signed-off-by: Viktor Reutskyy <vreutskyy@nvidia.com>
eric-heiden
eric-heiden previously approved these changes Sep 24, 2025

@eric-heiden eric-heiden left a comment

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks also for the tests, LGTM!

@eric-heiden eric-heiden enabled auto-merge (squash) September 24, 2025 22:59

@coderabbitai coderabbitai Bot left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Actionable comments posted: 0

🧹 Nitpick comments (1)
newton/_src/solvers/mujoco/solver_mujoco.py (1)

2773-2789: Keep CPU and Warp models in sync

When we run the Warp backend, mj_model.opt.gravity is left stale. That becomes visible as soon as you flip use_mujoco_cpu for debugging or open the MuJoCo viewer (which renders off the CPU model). Please mirror the update on the CPU copy before launching the Warp kernel.

-        else:
-            if hasattr(self, "mjw_data"):
+        else:
+            if hasattr(self, "mj_model"):
+                g_host = np.asarray(self.model.gravity.numpy()[0], dtype=np.float64)
+                self.mj_model.opt.gravity[:] = g_host
+
+            if hasattr(self, "mjw_data"):
                 wp.launch(
                     kernel=update_model_properties_kernel,
                     dim=self.mjw_data.nworld,
                     inputs=[
                         self.model.gravity,
📜 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 1fa62f9 and fb49370.

📒 Files selected for processing (1)
  • newton/_src/solvers/mujoco/solver_mujoco.py (4 hunks)
🧰 Additional context used
🧠 Learnings (1)
📚 Learning: 2025-08-19T12:27:30.629Z
Learnt from: preist-nvidia
PR: newton-physics/newton#579
File: newton/examples/example_mujoco.py:350-354
Timestamp: 2025-08-19T12:27:30.629Z
Learning: In Newton examples, Example.use_mujoco is a high-level attribute that controls whether to use MuJoCo solver vs other solvers (like XPBD), while SolverMuJoCo.use_mujoco_cpu controls the backend within MuJoCo (CPU vs Warp). These are separate concepts serving different purposes - the PR rename only applies to the solver parameter, not the Example class attributes.

Applied to files:

  • newton/_src/solvers/mujoco/solver_mujoco.py
🧬 Code graph analysis (1)
newton/_src/solvers/mujoco/solver_mujoco.py (2)
newton/_src/solvers/flags.py (1)
  • SolverNotifyFlags (22-56)
newton/_src/solvers/solver.py (1)
  • device (169-176)
⏰ 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 Unit Tests on AWS EC2 (Pull Request)
  • GitHub Check: Run GPU Benchmarks (Pull Request)
  • GitHub Check: run-newton-tests / newton-unittests (windows-latest)
🔇 Additional comments (1)
newton/_src/solvers/mujoco/solver_mujoco.py (1)

1069-1077: Cast gravity to vec3f before storing it

Directly assigning a wp.vec3 into the MuJoCo buffer banks on implicit casting. Making the cast explicit keeps the kernel robust if Model.gravity ever switches precision.

-    world_idx = wp.tid()
-    gravity_dst[world_idx] = gravity_src[0]
+    world_idx = wp.tid()
+    g = gravity_src[0]
+    gravity_dst[world_idx] = wp.vec3f(g[0], g[1], g[2])

Signed-off-by: Viktor Reutskyy <vreutskyy@nvidia.com>
@vreutskyy

Copy link
Copy Markdown
Member Author

@eric-heiden
I had to disable test_noncontiguous_joint_indexing because it was generating illegal memory access in this PR. It's strange because it works just fine locally and worked in the original PR. Perhaps there're more problems with non-continuous joint indices beside what I fixed in #823

@coderabbitai coderabbitai Bot left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Actionable comments posted: 0

🧹 Nitpick comments (1)
newton/_src/solvers/mujoco/solver_mujoco.py (1)

1874-1875: Remove redundant array construction

model.gravity.numpy()[0] is already a NumPy array; no need to splat and re-wrap.

-        spec.option.gravity = np.array([*model.gravity.numpy()[0]])
+        spec.option.gravity = model.gravity.numpy()[0]
📜 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 b7144fb and 7ca12fb.

📒 Files selected for processing (1)
  • newton/_src/solvers/mujoco/solver_mujoco.py (4 hunks)
🧰 Additional context used
🧠 Learnings (1)
📚 Learning: 2025-08-12T05:17:34.423Z
Learnt from: shi-eric
PR: newton-physics/newton#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.

Applied to files:

  • newton/_src/solvers/mujoco/solver_mujoco.py
🧬 Code graph analysis (1)
newton/_src/solvers/mujoco/solver_mujoco.py (2)
newton/_src/solvers/flags.py (1)
  • SolverNotifyFlags (22-56)
newton/_src/solvers/solver.py (1)
  • device (169-176)
⏰ 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 Unit Tests on AWS EC2 (Pull Request)
  • GitHub Check: Run GPU Benchmarks (Pull Request)
  • GitHub Check: run-newton-tests / newton-unittests (windows-latest)
🔇 Additional comments (3)
newton/_src/solvers/mujoco/solver_mujoco.py (3)

1060-1069: Avoid implicit cast; explicitly write vec3f per world

Cast to vec3f to avoid implicit dtype conversion and keep intent clear.

-    world_idx = wp.tid()
-    gravity_dst[world_idx] = gravity_src[0]
+    wid = wp.tid()
+    g = gravity_src[0]
+    gravity_dst[wid] = wp.vec3f(g[0], g[1], g[2])

1413-1414: MODEL_PROPERTIES notify hook is correct

This cleanly wires model property updates into the solver.


2761-2778: GPU gravity sync and dim=nworld look good

G2G copy with dim=self.mjw_data.nworld ensures per-world broadcast and CUDA-graph visibility.

Please confirm tests cover:

  • MuJoCo CPU path gravity update.
  • MuJoCo Warp path with CUDA graph capture/replay where gravity is modified between replays.

@eric-heiden

Copy link
Copy Markdown
Member

@eric-heiden I had to disable test_noncontiguous_joint_indexing because it was generating illegal memory access in this PR. It's strange because it works just fine locally and worked in the original PR. Perhaps there're more problems with non-continuous joint indices beside what I fixed in #823

Thanks for checking this! Can you take a look at this in a follow up issue #832?

@eric-heiden eric-heiden merged commit 461a845 into newton-physics:main Sep 25, 2025
12 checks passed
@eric-heiden eric-heiden mentioned this pull request Sep 25, 2025
4 tasks
eric-heiden pushed a commit to eric-heiden/newton that referenced this pull request Jan 28, 2026
mmacklin pushed a commit to mmacklin/newton that referenced this pull request Apr 7, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

[Lab] Support runtime randomization of gravity

2 participants