Fix per-env gravity randomization on Newton backend

[ooctipus]

Summary

GRAVITY_VEC_W on Newton's ArticulationData, RigidObjectData, and RigidObjectCollectionData was snapshotted from env 0 at construction, normalized to a unit direction, and broadcast to every env (and every body for the collection). Any in-place mutation of Newton's per-world model.gravity array — including the randomize_physics_scene_gravity event term, which writes per-env values via model.gravity[env_ids] = ... — was invisible to every consumer of GRAVITY_VEC_W and to the lazily-recomputed projected_gravity_b.

This PR binds GRAVITY_VEC_W directly to SimulationManager.get_model().gravity so it is a live view of the source of truth. The field now carries world-frame m/s² (per-env) instead of a unit direction; downstream consumers that need a unit projection normalize on read.

• New dedicated projected_gravity_b_kernel / projected_gravity_b_2D_kernel in isaaclab_newton.assets.kernels replace the generic quat_apply_inverse_* calls for gravity and normalize internally.
• body_projected_gravity_b in isaaclab.envs.mdp.observations normalizes the torch view (no-op for PhysX/OvPhysX, which still pre-normalize at construction).
• Warp kernels in isaaclab_experimental (_projected_gravity_kernel, _flat_orientation_l2_kernel) and humanoid task kernels (_base_up_proj_kernel, _upright_posture_bonus_kernel) switch from gravity_w[0] to wp.normalize(gravity_w[i]) so per-env randomized gravity is respected.
• The parity-test mock provides a per-env GRAVITY_VEC_W shaped to match Newton's model.gravity.

PhysX and OvPhysX backends are unchanged: their GRAVITY_VEC_W remains a pre-normalized unit vector, and the added F.normalize on the torch path is a no-op there.

Behavior change to note

isaaclab_newton.assets.{ArticulationData,RigidObjectData,RigidObjectCollectionData}.GRAVITY_VEC_W now exposes the raw m/s² gravity vector instead of a unit direction. Code that consumed it directly (rather than going through projected_gravity_b) needs to normalize before use. RigidObjectCollectionData.GRAVITY_VEC_W is now shaped (num_instances, 3) instead of (num_instances, num_bodies, 3); projected_gravity_b retains its (num_instances, num_bodies, 3) shape via the dedicated 2D kernel.

Test plan

Added test_gravity_vec_w_tracks_model_gravity to the three Newton asset test suites:

• source/isaaclab_newton/test/assets/test_articulation.py
• source/isaaclab_newton/test/assets/test_rigid_object.py
• source/isaaclab_newton/test/assets/test_rigid_object_collection.py

Each test pointer-equality-checks GRAVITY_VEC_W.warp.ptr against SimulationManager.get_model().gravity.ptr, mutates model.gravity per env in place, then asserts both GRAVITY_VEC_W and projected_gravity_b reflect the new values without sim.step() (avoiding orientation drift).

• ./isaaclab.sh -p -m pytest source/isaaclab_newton/test/assets/test_articulation.py::test_gravity_vec_w_tracks_model_gravity
• ./isaaclab.sh -p -m pytest source/isaaclab_newton/test/assets/test_rigid_object.py::test_gravity_vec_w_tracks_model_gravity
• ./isaaclab.sh -p -m pytest source/isaaclab_newton/test/assets/test_rigid_object_collection.py::test_gravity_vec_w_tracks_model_gravity
• Sanity-check randomize_physics_scene_gravity is observable in a quick locomotion config (e.g. shadow_hand).
• Verify regression tests fail when the fix is reverted (per AGENTS.md).

[greptile-apps]

Greptile Summary

This PR fixes per-env gravity randomization on the Newton backend by replacing the snapshotted, broadcast GRAVITY_VEC_W with a zero-copy ProxyArray bind directly to SimulationManager.get_model().gravity, so in-place mutations (e.g. randomize_physics_scene_gravity) are immediately visible to all consumers.

• ArticulationData, RigidObjectData, RigidObjectCollectionData: construction-time snapshot removed; GRAVITY_VEC_W is now a live view of Newton's per-env m/s\u00b2 gravity array, and two new dedicated Warp kernels (projected_gravity_b_kernel, projected_gravity_b_2D_kernel) normalize internally before projecting into the body frame.
• Experimental/humanoid MDP kernels: all six Warp kernels that previously indexed gravity_w[0] and skipped normalization are updated to wp.normalize(gravity_w[i]), enabling per-env gravity randomization throughout the observation and reward pipeline.
• Tests: three new regression tests verify pointer-equality between GRAVITY_VEC_W and model.gravity, then assert that an in-place gravity mutation propagates into projected_gravity_b without a physics step.

Confidence Score: 3/5

Safe to merge for the normal gravity-enabled path; the disabled-gravity configuration now silently produces NaN from wp.normalize on a zero vector wherever projected_gravity_b is consumed.

The live-bind approach is correct and the new regression tests are thorough for the happy path. However, every Warp kernel added or modified by this PR calls wp.normalize(gravity_w[i]) unconditionally. When Newton gravity is disabled, gravity_w[i] is (0,0,0) and IEEE 754 division-by-zero turns the entire projected_gravity_b output and the flat-orientation and upright-posture reward signals to NaN. The previous quat_apply_inverse kernels rotated the zero vector directly and produced (0,0,0), a well-defined result. No test covers this regression, so it could go undetected in workflows that include a gravity-off phase or toggled-gravity curriculum.

source/isaaclab_newton/isaaclab_newton/assets/kernels.py — both new kernels; also the six Warp kernels in isaaclab_experimental and isaaclab_tasks_experimental that call wp.normalize(gravity_w[i]).

Important Files Changed

Filename	Overview
source/isaaclab_newton/isaaclab_newton/assets/kernels.py	Adds projected_gravity_b_kernel and projected_gravity_b_2D_kernel that normalize gravity_w[i] before rotating into body frame; wp.normalize on a zero vector (disabled gravity) produces NaN.
source/isaaclab_newton/isaaclab_newton/assets/articulation/articulation_data.py	Replaces snapshot-and-broadcast construction of GRAVITY_VEC_W with a direct ProxyArray bind to model.gravity; projected-gravity kernel swapped to the new normalizing variant.
source/isaaclab_newton/isaaclab_newton/assets/rigid_object_collection/rigid_object_collection_data.py	Live-bind fix; GRAVITY_VEC_W shape changes from (num_instances, num_bodies) to (num_instances,), and per-body broadcast now happens inside projected_gravity_b_2D_kernel.
source/isaaclab/isaaclab/envs/mdp/observations.py	Adds F.normalize on the torch path before quat_apply_inverse; no-op for PhysX (already unit), correct for Newton (m/s² → unit direction).
source/isaaclab_experimental/isaaclab_experimental/envs/mdp/observations.py	Switches Warp kernels from gravity_w[0] to wp.normalize(gravity_w[i]); correct for per-env randomization but inherits the zero-gravity NaN risk.
source/isaaclab_newton/test/assets/test_rigid_object_collection.py	Existing test updated for new shape/magnitude; new regression test covers the live-bind contract but is missing @pytest.mark.isaacsim_ci unlike the sibling test in test_rigid_object.py.

Flowchart

%%{init: {'theme': 'neutral'}}%%
flowchart TD
    A["Newton model.gravity\n(num_envs, vec3f, m/s²)"] -->|"ProxyArray — zero-copy bind"| B["GRAVITY_VEC_W\n(ArticulationData / RigidObjectData /\nRigidObjectCollectionData)"]

    B -->|"1D: projected_gravity_b_kernel\nwp.normalize → quat_rotate_inv"| C["projected_gravity_b\n(unit vector, body frame)"]
    B -->|"2D: projected_gravity_b_2D_kernel\nbroadcast per env → per body\nwp.normalize → quat_rotate_inv"| D["projected_gravity_b\n(RigidObjectCollection,\nnum_envs × num_bodies)"]

    B -->|"torch path: F.normalize"| E["body_projected_gravity_b\n(isaaclab observations.py)"]

    B -->|"wp.normalize(gravity_w[i])"| F["_projected_gravity_kernel\n_flat_orientation_l2_kernel\n_base_up_proj_kernel\n_upright_posture_bonus_kernel\n(experimental/humanoid mdp)"]

    G["randomize_physics_scene_gravity\nmodel.gravity[env_ids] = ..."] -->|"in-place write"| A

    style A fill:#4CAF50,color:#fff
    style G fill:#2196F3,color:#fff
    style C fill:#FF9800,color:#fff
    style D fill:#FF9800,color:#fff

Loading

_{Reviews (1): Last reviewed commit: "Fix per-env gravity randomization on New..." | Re-trigger Greptile}

[greptile-apps]

Zero-gravity input produces NaN

When gravity is disabled on Newton, model.gravity is (0, 0, 0) for every env. wp.normalize((0,0,0)) performs a IEEE 754 0 / 0 division, producing a NaN vec3f. Every downstream consumer then receives NaN — projected_gravity_b, the flat-orientation reward, base_up_proj, and the humanoid upright-posture bonus — silently corrupting training.

The previous quat_apply_inverse_1D_kernel path rotated the zero vector directly, correctly yielding (0,0,0). The same edge case affects projected_gravity_b_2D_kernel (line 482) and all six Warp kernels in isaaclab_experimental / isaaclab_tasks_experimental that call wp.normalize(gravity_w[i]). A length-guard (or delegating to wp.normalize only when length > 0) would restore the old well-defined behavior for disabled-gravity configurations.

[greptile-apps]

Missing @pytest.mark.isaacsim_ci marker

test_gravity_vec_w_tracks_model_gravity here is missing @pytest.mark.isaacsim_ci, while the analogous test in test_rigid_object.py carries it, and the pre-existing test_gravity_vec_w in this same file also uses @pytest.mark.isaacsim_ci. Without the marker, this test may run in CI environments that lack an Isaac Sim licence or the required GPU, causing the suite to fail for the wrong reason.

Note: If this suggestion doesn't match your team's coding style, reply to this and let me know. I'll remember it for next time!

[isaaclab-review-bot]

Code Review: Fix per-env gravity randomization on Newton backend

Thanks for this well-structured fix, @ooctipus! The root cause analysis and implementation approach are both solid.

Summary

This PR correctly addresses a subtle but important bug where GRAVITY_VEC_W on Newton's data classes was snapshotted at construction time from env 0, normalized, and broadcast to all environments. This made per-env gravity randomization (via randomize_physics_scene_gravity) invisible to downstream consumers.

Key changes:

1. Bind GRAVITY_VEC_W directly to SimulationManager.get_model().gravity (a live view)
2. Add dedicated projected_gravity_b_kernel / projected_gravity_b_2D_kernel that normalize internally
3. Update all consumers to normalize before use (torch path via F.normalize, warp kernels via wp.normalize)

Strengths ✅

Design:

• Clean separation: GRAVITY_VEC_W exposes raw m/s² data, normalization happens at point-of-use
• Live binding to model.gravity eliminates cache invalidation complexity
• PhysX/OvPhysX compatibility preserved (their GRAVITY_VEC_W is already unit-normalized, so F.normalize is a no-op)

Implementation:

• Warp kernels are properly documented with clear docstrings
• The 2D kernel correctly broadcasts per-env gravity across bodies
• Consistent fix pattern across all three data classes (ArticulationData, RigidObjectData, RigidObjectCollectionData)

Testing:

• Regression tests verify pointer equality (GRAVITY_VEC_W.warp.ptr == model_gravity_arr.ptr) — this is the right check
• Tests mutate model.gravity per-env and verify both GRAVITY_VEC_W and projected_gravity_b reflect changes
• No sim.step() in tests avoids orientation drift — smart approach

Documentation:

• Excellent changelog entry explaining the before/after behavior
• PR description clearly documents the breaking change to GRAVITY_VEC_W shape for collections

Minor Observations (Non-blocking)

1. Shape change for RigidObjectCollectionData.GRAVITY_VEC_W: The shape changes from (num_instances, num_bodies, 3) to (num_instances, 3). This is correctly documented in the PR description and changelog. Any downstream code directly accessing GRAVITY_VEC_W on collections will need adjustment — the PR description calls this out appropriately.

2. wp.normalize on zero vector: If gravity_w is exactly zero (gravity disabled), wp.normalize(gravity_w[i]) could produce NaN/undefined behavior depending on Warp's implementation. However, looking at the test test_gravity_vec_w which tests gravity_enabled=False, this case is handled (gravity vector is (0,0,0) and doesn't need normalization for comparison). The kernel would still be called, but the test doesn't validate projected_gravity_b in the zero-gravity case. This is edge-case and likely fine for the intended use cases.

3. Index change in experimental kernels: The change from gravity_w[0] to gravity_w[i] in _projected_gravity_kernel and _flat_orientation_l2_kernel is correct — just noting this is the crux of the per-env fix.

Verdict

This is a clean, well-tested bug fix with appropriate attention to backward compatibility documentation. The approach of binding directly to the source array rather than maintaining a copy that needs invalidation is the right architectural choice.

LGTM — the implementation correctly addresses the bug and the tests cover the regression scenario effectively.

[isaaclab-review-bot]

Code Review: Fix per-env gravity randomization on Newton backend

Thanks for fixing this subtle but important bug in the Newton backend gravity handling, @ooctipus! The core issue—GRAVITY_VEC_W being a stale snapshot instead of a live view—was a significant problem that would have silently broken any workflows using per-environment gravity randomization.

Summary

The fix correctly binds GRAVITY_VEC_W directly to Newton's model.gravity array, making it a live view of the per-env gravity values. The approach is sound and the implementation is clean.

Findings

CI: Missing Changelog Fragments (action required)

The Check changelog fragments CI job is failing because two modified packages are missing their changelog entries:

• source/isaaclab/changelog.d/<slug>.rst — for the body_projected_gravity_b() normalization change
• source/isaaclab_experimental/changelog.d/<slug>.rst — for the kernel changes in observations/rewards

The isaaclab_newton changelog is present and well-written.

Documentation: Consider noting breaking change for direct consumers

The PR description clearly documents the behavior change: GRAVITY_VEC_W now returns m/s² vectors instead of unit directions. This is a breaking change for any code that consumed GRAVITY_VEC_W directly (rather than through projected_gravity_b).

Consider whether this warrants an entry in the main isaaclab changelog noting the contract change, even if the change is technically Newton-only. Users with cross-backend code relying on GRAVITY_VEC_W being unit-normalized may be surprised.

Code Quality: Consistent normalization across backends

The fix adds torch.nn.functional.normalize() calls in the torch path and wp.normalize() in Warp kernels. This is the right approach—normalizing on read preserves the semantic contract (unit-direction projection) while allowing the source data to carry magnitude.

One minor note: torch.nn.functional.normalize() with a zero vector returns [0, 0, 0] (no NaN), and wp.normalize() should behave similarly. This handles the gravity_enabled=False case gracefully, which the existing test_gravity_vec_w test with gravity_enabled=False parameter would exercise. However, you may want to verify this edge case is covered by your new tests or add a brief comment noting the behavior.

Test Coverage: Strong regression tests

The new test_gravity_vec_w_tracks_model_gravity tests are well-designed:

• They verify pointer equality between GRAVITY_VEC_W.warp.ptr and model.gravity.ptr
• They mutate gravity per-env and verify propagation
• They avoid sim.step() to prevent orientation drift—good practice for isolated testing

The tests cover ArticulationData, RigidObjectData, and RigidObjectCollectionData. The collection test appropriately verifies the shape change from (num_envs, num_bodies, 3) to (num_envs, 3).

Minor: RigidObjectCollectionData.GRAVITY_VEC_W shape change

The shape change from (num_instances, num_bodies, 3) to (num_instances, 3) is documented in the PR description. The new projected_gravity_b_2D_kernel correctly broadcasts the per-env gravity across bodies. This is cleaner—gravity doesn't vary per-body within an env.

Overall Assessment

This is a well-executed fix for a real bug. The implementation is sound, the tests are thorough, and the approach (live binding + normalize-on-read) is the right one. Once the missing changelog fragments are added, this should be ready to merge.

Suggested priority: Address the CI failure by adding the missing changelog fragments.

---

Update (24a3bd0): ✅ The missing changelog fragments have been added for both isaaclab and isaaclab_experimental. The follow-up commits also include nice cleanup: docstring condensation, switching FORWARD_VEC_B initialization from torch to numpy (cleaner, avoids torch import for a static array), and adding @pytest.mark.isaacsim_ci markers to relevant tests. No new concerns—ready for merge pending CI.

---

Update (4111a04): Reviewed the incremental diff. The latest push adds no new code changes—the PR remains clean. All previously noted concerns have been addressed. Looks good to merge! ✅

[StafaH]

Thanks @ooctipus, left some minor comments.