[Newton] Adds IMU Sensor

[camevor]

Description

• Adds a Newton-backend IMU sensor (isaaclab_newton.sensors.imu) that wraps Newton's native SensorIMU, providing angular velocity and linear acceleration in the sensor's body frame
• Extends NewtonManager with request_state_attribute() / request_contact_attribute() for sensors to declare required extended attributes before model finalization
• Adds IMU sensor lifecycle management to NewtonManager (add_imu_sensor(), per-step updates in _simulate_physics_only())
• Wires Newton IMU into the base Imu/ImuData factory classes for automatic backend dispatch
• Uses site injection via cl_register_site().

[greptile-apps]

Greptile Summary

This PR wires Newton's native SensorIMU into IsaacLab's Imu/ImuData factory classes, extends NewtonManager with IMU lifecycle management (add_imu_sensor, request_extended_state_attribute/contact_attribute, per-step updates), and adds a _invalidate_initialize_callback to FrameTransformer for close/reinit resilience.

• Missing implementation: __init__.pyi declares from .imu import Imu, ImuData but no isaaclab_newton/sensors/imu/ package (no __init__.py, imu.py, imu_data.py) exists. The lazy-export system will raise ImportError on first access of NewtonImu/NewtonImuData, making the feature completely non-functional as shipped.

Confidence Score: 3/5

• Not safe to merge: the Newton IMU sensor is declared in the public stub but the implementation package is absent, making the feature broken.
• The P1 finding (missing imu subpackage) means the advertised feature is completely non-functional — any code that imports from isaaclab_newton.sensors import Imu will get an ImportError at runtime. The NewtonManager infrastructure and FrameTransformer fix are solid, but the core deliverable of the PR is missing.
• source/isaaclab_newton/isaaclab_newton/sensors/init.pyi needs the corresponding imu/ subpackage added (imu.py, imu_data.py, init.py, init.pyi).

Important Files Changed

Filename	Overview
source/isaaclab_newton/isaaclab_newton/sensors/init.pyi	Declares from .imu import Imu, ImuData but the imu subpackage does not exist anywhere in isaaclab_newton/sensors/, causing an ImportError at access time.
source/isaaclab_newton/isaaclab_newton/physics/newton_manager.py	Adds add_imu_sensor(), request_extended_state_attribute(), request_extended_contact_attribute(), and per-step IMU sensor updates in _simulate_physics_only(). Logic is sound and consistent with the existing FrameTransformer pattern.
source/isaaclab/isaaclab/sensors/imu/imu.py	Wires Newton backend types (NewtonImu, NewtonImuData) into the Imu factory class's type annotations. Correct and minimal change.
source/isaaclab/isaaclab/sensors/imu/imu_data.py	Adds Newton backend type annotation to ImuData factory class. Correct and minimal change.
source/isaaclab_newton/isaaclab_newton/sensors/frame_transformer/frame_transformer.py	Adds _invalidate_initialize_callback to re-register sites after close()/reinit, correctly clearing _newton_transforms and _sensor_index and re-calling cl_register_site. Solid fix.
source/isaaclab_newton/docs/CHANGELOG.rst	Version jump from 0.5.13 to 0.5.15 skips 0.5.14; FrameTransformer invalidation fix is unmentioned.
source/isaaclab_newton/config/extension.toml	Bumped to 0.5.15 to match the CHANGELOG, but should be 0.5.14 (sequential bump from 0.5.13).

Sequence Diagram

sequenceDiagram
    participant User
    participant Imu as isaaclab.sensors.Imu (factory)
    participant NewtonImu as isaaclab_newton.sensors.imu.Imu (MISSING)
    participant NM as NewtonManager
    participant SIMU as newton.sensors.SensorIMU

    User->>Imu: Imu(cfg)
    Imu->>NewtonImu: __new__ dispatches to Newton backend
    NewtonImu-->>User: ImportError (module not found)

    Note over NM,SIMU: If implementation existed:
    NewtonImu->>NM: request_extended_state_attribute("body_qdd")
    NewtonImu->>NM: cl_register_site(body_pattern, xform)
    NM->>NM: start_simulation() → builder.request_state_attributes()
    NewtonImu->>NM: add_imu_sensor(sites)
    NM->>SIMU: "SensorIMU(model, sites=sites, request_state_attributes=False)"
    NM-->>NewtonImu: sensor_index

    loop Each physics step
        NM->>SIMU: sensor.update(state_0)
        SIMU-->>NewtonImu: ang_vel_b, lin_acc_b populated
    end

Loading

_{Reviews (1): Last reviewed commit: "Fix changelog" | Re-trigger Greptile}

[greptile-apps]

Missing CHANGELOG entry for FrameTransformer fix

The PR includes a second commit (6a00eb6a) that adds _invalidate_initialize_callback to FrameTransformer — a bug fix that ensures sites survive a close()/reinit cycle. This fix is absent from the 0.5.15 CHANGELOG entry, violating the project rule that every change to the source directory must be recorded. It should appear under a Fixed section, e.g.:

Fixed
^^^^^

* Fixed :class:`~isaaclab_newton.sensors.FrameTransformer` failing to re-register
  sites after a ``close()``/reinit cycle by adding ``_invalidate_initialize_callback``.

[isaaclab-review-bot]

🤖 Isaac Lab Review Bot

Summary

This PR adds Newton-backend IMU sensor support by: (1) creating a Imu/ImuData wrapper around Newton's native SensorIMU, (2) extending NewtonManager with request_extended_state_attribute() / request_extended_contact_attribute() APIs for pre-finalization attribute declaration, (3) wiring Newton IMU types into the base factory classes, and (4) fixing FrameTransformer's _invalidate_initialize_callback to re-register sites on close/reinit cycles. The implementation is clean and well-tested, with a few design points worth considering.

Design Assessment

Design is sound. The approach correctly follows Isaac Lab's established patterns:

• The Newton Imu class mirrors the FrameTransformer lifecycle: register sites/attributes in __init__ → resolve in _initialize_impl → copy in _update_buffers_impl → clean up in _invalidate_initialize_callback.
• Extended attribute request APIs (request_extended_state_attribute/request_extended_contact_attribute) are a clean, general-purpose mechanism that will benefit future sensors beyond IMU.
• add_imu_sensor() improves on add_frame_transform_sensor() by adding a cls._model is None pre-condition check.
• The Warp kernels are minimal and correct — simple masked copies that map 1:1 sensor-to-environment.

Findings

🟡 Warning: request_extended_state_attribute() silently succeeds after model finalization — newton_manager.py:470

If a sensor calls request_extended_state_attribute() after start_simulation() has already run, the attribute is added to _pending_extended_state_attributes but will never be forwarded to the builder (the builder is already finalized and the pending set was cleared at line 637). The call silently succeeds without effect.

The docstring says "Sensors call this during __init__, before model finalization" which documents the expectation, but a defensive guard would catch misuse:

@classmethod
def request_extended_state_attribute(cls, attr: str) -> None:
    if cls._model is not None:
        raise RuntimeError(
            f"request_extended_state_attribute('{attr}') called after model finalization. "
            "Extended attributes must be requested before start_simulation()."
        )
    cls._pending_extended_state_attributes.add(attr)

Note: This would need to account for the _invalidate_initialize_callback reinit path, where _model is set but clear() has already reset it to None before the callback fires. Based on the docstring for the invalidation callback ("NewtonManager.close() calls clear() before dispatching STOP"), the guard should be safe — but worth verifying.

The same consideration applies to request_extended_contact_attribute().

🟡 Warning: _invalidate_initialize_callback re-registers body_qdd but not contact attributes — imu.py:183

The IMU invalidation callback re-registers the site and the body_qdd state attribute, which is correct. However, request_extended_contact_attribute() exists as a parallel API. If a future sensor uses both state and contact attributes, the invalidation pattern here sets the precedent. Worth confirming that contact attributes are similarly re-registered where needed — or adding a comment clarifying that IMU only needs state attributes (not contacts).

Currently the IMU only uses body_qdd (state), so this is not a bug — just a note for future-proofing.

🔵 Suggestion: Test coverage could verify the close/reinit cycle — test_imu.py

The test suite is thorough for a new feature PR: initialization, shapes, gravity-at-rest, freefall, reset, and sensor print. The freefall test is a particularly good physics correctness check.

Missing coverage that would be valuable:

• Close/reinit cycle: The FrameTransformer invalidation fix and the IMU invalidation callback are critical for multi-reset workflows. A test that calls sim.reset() twice (triggering the close → clear → reinit path) would verify the site re-registration logic works end-to-end.
• Partial reset: imu.reset(env_ids=[0]) with 2+ envs — verifying only the specified env is zeroed while the other retains its data.

🔵 Suggestion: Version 0.5.14 in extension.toml but CHANGELOG shows 0.5.14 — Minor bookkeeping

The version bump from 0.5.13 → 0.5.14 is consistent between extension.toml and CHANGELOG.rst. No issue here — just noting that the gap is clean (no orphaned 0.5.14 entry exists elsewhere).

Test Coverage

• Feature PR: Coverage — Yes, good baseline
• The test file (test_imu.py) covers 7 test cases: initialization, data shapes, gravity at rest, angular velocity at rest, reset, freefall acceleration, and sensor print
• The gravity-at-rest and freefall tests validate physics correctness with appropriate tolerances (±0.5 m/s²)
• Gaps: Close/reinit cycle test and partial-reset test would strengthen coverage but are not blocking

CI Status

• ✅ pre-commit, Check for Broken Links, Detect Changes, Detect Doc Build Type, labeler
• ⏳ Build Base Docker Image, Build Latest Docs, Build cuRobo Docker Image, license-check (pending)
• ⏭️ Build Multi-Version Docs, Deploy Docs (skipping)

Verdict

Minor fixes needed

Clean, well-structured implementation that follows established Isaac Lab patterns. The extended attribute request APIs are a good general-purpose addition. The test suite covers the important physics behaviors. The main suggestion is adding a defensive guard to request_extended_state_attribute() / request_extended_contact_attribute() to catch post-finalization misuse, and optionally adding a close/reinit integration test. No blocking issues.

[isaaclab-review-bot]

🟡 Warning: If this method is called after start_simulation(), the attribute is silently added to the pending set but never forwarded to the builder (it was already finalized and the set cleared). Consider adding a guard:

Suggested change

	cls._pending_extended_state_attributes.add(attr)
	if cls._model is not None:
	raise RuntimeError(
	f"request_extended_state_attribute('{attr}') called after model finalization. "
	"Extended attributes must be requested before start_simulation()."
	)
	cls._pending_extended_state_attributes.add(attr)

The _invalidate_initialize_callback path should be safe since NewtonManager.clear() resets _model to None before callbacks fire.

[AntoineRichard]

I think this is fine, users are not expected to call this, and it should be called in order.

[AntoineRichard]

LGTM, let's wait for CI to come back.

[AntoineRichard]

I think this is fine, users are not expected to call this, and it should be called in order.

[AntoineRichard]

Hum, I see this was required at some point, is there a way to make a test that would help catch bugs related to this?