[BUG] SensorIMU world-frame sites ignore per-world gravity

[pietis]

Description

SensorIMU applies the gravity of world 0 to every world-frame site (body=-1) in a multi-world model, even when the site belongs to another world.

On current main, two sites assigned to worlds 0 and 1 with Z gravity values of -1 and -5 m/s² both report a Z specific force of +1:

shape_body: [-1, -1]
shape_world: [0, 1]
gravity_z: [-1.0, -5.0]
accelerometer_z: [1.0, 1.0]

The expected accelerometer values are [1.0, 5.0].

SensorIMU was introduced in PR #1161 with both the body-attached path and the path for sites not attached to a body reading gravity[0]. PR #1381 later added per-world gravity support and updated the body-attached path to select gravity through body_world, while the early-return path for sites with body=-1 continued to read gravity[0].

The affected branch in compute_sensor_imu_kernel therefore returns before selecting the gravity associated with the site's world.

This causes fixed IMUs represented by world-local sites to report incorrect specific force in multi-world simulations with different gravity values.

Minimal reproduction

Two world-frame sites, no bodies, with distinct gravity values assigned to their respective worlds. Each IMU should report the specific force corresponding to its own world.

import newton
from newton.sensors import SensorIMU

builder = newton.ModelBuilder()

builder.begin_world(gravity=(0.0, 0.0, -1.0))
site_0 = builder.add_site(-1)
builder.end_world()

builder.begin_world(gravity=(0.0, 0.0, -5.0))
site_1 = builder.add_site(-1)
builder.end_world()

model = builder.finalize(device="cpu")
sensor = SensorIMU(model, sites=[site_0, site_1])
state = model.state()

state.body_qdd.zero_()
sensor.update(state)

print("shape_body:", model.shape_body.numpy().tolist())
print("shape_world:", model.shape_world.numpy().tolist())
print("gravity_z:", model.gravity.numpy()[:, 2].tolist())
print("accelerometer_z:", sensor.accelerometer.numpy()[:, 2].tolist())

Output on main (825f0c9):

shape_body: [-1, -1]
shape_world: [0, 1]
gravity_z: [-1.0, -5.0]
accelerometer_z: [1.0, 1.0]

Expected behavior

A world-local site not attached to a body should use the gravity assigned to its shape_world. For the reproduction above, the Z accelerometer values should be:

accelerometer_z: [1.0, 5.0]

Global sites (shape_world == -1) may retain the existing convention of using gravity from world 0.

Validation

• Add a regression with two world-local sites not attached to bodies and distinct gravity values.
• Confirm that each site reports the specific force corresponding to its own world.
• Confirm that a global site (shape_world == -1) retains the world-0 gravity fallback.
• Confirm that body-attached IMUs continue to select gravity through body_world.
• Run the existing SensorIMU test suite on the supported Warp backends.

System information

• Newton: 1.5.0.dev0 (825f0c9)
• Warp: 1.16.0.dev20260716
• Python 3.12.13, macOS 26.5.1 arm64
• Device: CPU

CUDA execution was not tested. The failing branch unconditionally indexes gravity[0] without a device-specific code path.

[eric-heiden]

Thanks @pietis for reporting this and for putting together the IMU fix in #3716! We’re now unifying the world_gravity array so that the global world (-1) is stored in the final entry, and we’re incorporating your changes into #3724.