Moves gravity compensation from schemas to MDP event

[vidurv-nvidia]

Description

Adds gravity compensation support for the Newton/MuJoCo backend as an MDP event, addressing review feedback from @ooctipus on #4847 to keep Newton-specific attributes out of the PhysX schema layer.

New MDP event: set_gravity_compensation() writes mjc:gravcomp (body-level) and mjc:actuatorgravcomp (joint-level) USD attributes at prestartup time. This follows the same pattern as existing USD-modifying events like randomize_rigid_body_scale.

Bug fix: SolverMuJoCo.register_custom_attributes(builder) was missing from NewtonManager.instantiate_builder_from_stage(), causing all mjc: prefixed custom attributes (gravcomp, actuatorgravcomp, etc.) to be silently ignored when parsing USD. The cloning path (newton_replicate.py) already had this call — it was only missing from the normal path.

Usage example:

events = {
    "gravity_compensation": EventTermCfg(
        func=mdp.set_gravity_compensation,
        mode="prestartup",
        params={
            "asset_cfg": SceneEntityCfg("robot"),
            "body_gravity_compensation_scale": 1.0,          # uniform for all bodies
            # or per-body: {".*link[0-3]": 1.0, ".*link4": 0.5}
            "joint_gravity_compensation": [".*"],             # regex patterns for joints
        },
    ),
}

The pipeline: IsaacLab config → USD attributes → Newton ModelBuilder → MuJoCo solver / MJCF XML.

Supersedes #4847

Type of change

• New feature (non-breaking change which adds functionality)
• Bug fix (non-breaking change which fixes an issue)

Checklist

• I have read and understood the contribution guidelines
• I have run the pre-commit checks with ./isaaclab.sh --format
• I have made corresponding changes to the documentation
• My changes generate no new warnings
• I have added tests that prove my fix is effective or that my feature works
• I have updated the changelog and the corresponding version in the extension's config/extension.toml file
• I have added my name to the CONTRIBUTORS.md or my name already exists there

[greptile-apps]

Greptile Summary

This PR moves gravity compensation from the PhysX schema layer (RigidBodyPropertiesCfg / JointDrivePropertiesCfg) into a new set_gravity_compensation MDP event that writes mjc:gravcomp and mjc:actuatorgravcomp USD attributes at prestartup time, keeping Newton/MuJoCo-specific concerns out of the shared schema layer.

• env_ids is silently ignored in set_gravity_compensation (events.py lines 198–207): the function iterates every prim_path unconditionally instead of indexing with env_ids as every other USD-level event does, breaking the documented API contract for callers that pass a subset of environments.
• isaaclab_newton CHANGELOG is not updated: the PR reverses the 0.5.10 schema change but adds no new version entry and does not bump extension.toml, violating the project's changelog guidelines.

Confidence Score: 4/5

Safe to merge after fixing the env_ids logic bug; the changelog omission is a process issue that should also be resolved.

One P1 defect: env_ids is documented as controlling which environments are affected but is never consulted, so any caller passing a subset of env IDs silently gets the wrong behavior. The missing isaaclab_newton changelog entry is a process violation per project guidelines. Both should be addressed before merge.

source/isaaclab/isaaclab/envs/mdp/events.py (env_ids logic) and source/isaaclab_newton/docs/CHANGELOG.rst + config/extension.toml (missing version bump)

Vulnerabilities

No security concerns identified. The change writes custom USD namespace attributes (mjc:gravcomp, mjc:actuatorgravcomp) to prims at prestartup; no user-controlled input is interpolated into prim paths or attribute names without validation.

Important Files Changed

Filename	Overview
source/isaaclab/isaaclab/envs/mdp/events.py	Adds set_gravity_compensation MDP event; env_ids parameter is accepted but never used, always writing to all environments regardless of the caller's intent.
source/isaaclab_newton/docs/CHANGELOG.rst	Missing a new 0.5.11 version entry documenting the removal of gravity compensation from the schema layer; guidelines forbid adding to an existing released version.
source/isaaclab/isaaclab/sim/schemas/schemas_cfg.py	Removes gravity_compensation_scale from RigidBodyPropertiesCfg and gravity_compensation from JointDrivePropertiesCfg; clean removal with no leftover references.
source/isaaclab_newton/test/test_actuator_gravity_comp.py	New E2E test that writes mjc:gravcomp / mjc:actuatorgravcomp USD attributes directly and verifies they propagate through Newton to MJCF XML; well-structured and covers the core contract.
source/isaaclab_newton/config/extension.toml	Version still 0.5.10; should be bumped to 0.5.11 to match the missing changelog entry for this PR's schema removals.

Sequence Diagram

sequenceDiagram
    participant User as User Config
    participant EventMgr as EventManager
    participant Func as set_gravity_compensation()
    participant USD as USD Stage
    participant Newton as Newton ModelBuilder
    participant MuJoCo as SolverMuJoCo / MJCF

    User->>EventMgr: EventTermCfg(func=set_gravity_compensation, mode="prestartup")
    EventMgr->>Func: call(env, env_ids, asset_cfg, ...)
    Func->>USD: safe_set_attribute(body_prim, "mjc:gravcomp", scale)
    Func->>USD: safe_set_attribute(joint_prim, "mjc:actuatorgravcomp", True)
    Note over Func,USD: env_ids currently ignored — writes to ALL prim paths
    USD-->>Func: attributes written
    Note over EventMgr,USD: sim.reset() / model build happens after prestartup
    USD->>Newton: builder.add_usd(stage)
    Newton-->>Newton: parses mjc:gravcomp → model.mujoco.gravcomp
    Newton-->>Newton: parses mjc:actuatorgravcomp → model.mujoco.jnt_actgravcomp
    Newton->>MuJoCo: SolverMuJoCo(model)
    MuJoCo-->>MuJoCo: exports gravcomp / actuatorgravcomp to MJCF XML

Loading

_{Reviews (1): Last reviewed commit: "Move gravity compensation from schemas t..." | Re-trigger Greptile}

[greptile-apps]

env_ids silently ignored — always writes to all environments

env_ids is accepted in the signature and documented as "If None, all environments are used," but the implementation never references it. prim_paths is iterated unconditionally, so the function always writes USD attributes to every environment instance regardless of what the caller passes.

Every analogous USD-level event (e.g. randomize_rigid_body_scale) resolves env_ids first and then indexes prim_paths[env_id]. The same pattern should be applied here:

    asset: Articulation | RigidObject = env.scene[asset_cfg.name]
    prim_paths = sim_utils.find_matching_prim_paths(asset.cfg.prim_path)

    if env_ids is None:
        env_ids = torch.arange(env.scene.num_envs, device="cpu")
    else:
        env_ids = env_ids.cpu()

    from pxr import UsdPhysics  # noqa: PLC0415

    stage = env.sim.stage

    _JOINT_TYPES = {"PhysicsRevoluteJoint", "PhysicsPrismaticJoint", "PhysicsSphericalJoint", "PhysicsD6Joint"}

    for env_id in env_ids:
        prim_path = prim_paths[env_id]

[greptile-apps]

Missing new version entry for isaaclab_newton

The PR removes gravity_compensation from JointDrivePropertiesCfg — functionality that was explicitly added under the 0.5.10 changelog entry ("Moved joint-level gravity compensation to the schemas layer in JointDrivePropertiesCfg"). The guidelines require a new version heading for every change touching a source/<package>/ directory; adding to an existing (released) version is prohibited.

A 0.5.11 entry should be added describing the removal and providing migration guidance (the new set_gravity_compensation MDP event), and source/isaaclab_newton/config/extension.toml should be bumped to 0.5.11 to match.

[greptile-apps]

_JOINT_TYPES recreated on every call

This set is defined inside the function body, so Python recreates it on each invocation. It is a pure constant and should be hoisted to module scope (alongside the other module-level names in this file) so it is constructed once:

Suggested change

	_JOINT_TYPES = {"PhysicsRevoluteJoint", "PhysicsPrismaticJoint", "PhysicsSphericalJoint", "PhysicsD6Joint"}
	_JOINT_TYPES = {"PhysicsRevoluteJoint", "PhysicsPrismaticJoint", "PhysicsSphericalJoint", "PhysicsD6Joint"}

(moved to module level, outside any function)

[isaaclab-review-bot]

🤖 Isaac Lab Review Bot

Summary

This PR cleanly moves gravity compensation from the PhysX schema layer to an MDP event function, addressing prior review feedback. The approach is architecturally sound—it keeps Newton-specific USD attributes (mjc:gravcomp, mjc:actuatorgravcomp) out of the shared schema layer and follows the established prestartup event pattern. The bug fix for register_custom_attributes is correct but incomplete. Two issues need attention before merge.

Design Assessment

Design is sound. An MDP event for writing backend-specific USD attributes at prestartup is the right pattern—it mirrors randomize_rigid_body_scale and keeps Newton concerns separate from PhysX schemas. The API surface (uniform float vs dict regex patterns for bodies, regex list for joints) is flexible and consistent with existing Isaac Lab conventions.

Findings

🔴 Critical: set_gravity_compensation missing from __init__.pyi import block — source/isaaclab/isaaclab/envs/mdp/__init__.pyi

The function is added to __all__ (line 61) but is not added to the from .events import (...) block (around line 168). Since __init__.py uses lazy_export() which parses the .pyi's explicit import statements to determine what to lazy-load, set_gravity_compensation will not be importable as mdp.set_gravity_compensation at runtime. Users following the documented example will get an ImportError.

from .events import (
    apply_external_force_torque,
    push_by_setting_velocity,
    randomize_actuator_gains,
    randomize_fixed_tendon_parameters,
    randomize_joint_parameters,
    randomize_physics_scene_gravity,
    randomize_rigid_body_collider_offsets,
    randomize_rigid_body_com,
    randomize_rigid_body_mass,
    randomize_rigid_body_material,
    randomize_rigid_body_scale,
    set_gravity_compensation,
    randomize_visual_color,
    ...

🔴 Critical: register_custom_attributes missing on proto builder in multi-env path — newton_manager.py:501

The fix correctly adds SolverMuJoCo.register_custom_attributes(builder) for the main builder, but when instantiate_builder_from_stage detects env Xforms and creates a proto ModelBuilder (line ~501), that prototype builder does not get the registration. This means mjc: attributes won't be parsed when building prototypes from per-env subtrees. The cloning path in newton_replicate.py:68 already does this correctly—this path should too.

            proto = ModelBuilder(up_axis=up_axis)
            SolverMuJoCo.register_custom_attributes(proto)
            proto.add_usd(

🟡 Warning: UsdPhysics.RigidBodyAPI(p) does not filter for applied API — events.py:223

The filter [p for p in Usd.PrimRange(root_prim) if UsdPhysics.RigidBodyAPI(p)] constructs a schema wrapper and checks its truthiness. UsdSchemaBase.__bool__ returns GetPrim().IsValid(), which is True for every valid prim in the range—not just prims with RigidBodyAPI applied. This means mjc:gravcomp gets written on collision shapes, joints, xforms, etc.

In practice Newton likely ignores the extra attributes, but it's semantically incorrect and could confuse debugging. The correct check is p.HasAPI(UsdPhysics.RigidBodyAPI).

            body_prims = [p for p in Usd.PrimRange(root_prim) if p.HasAPI(UsdPhysics.RigidBodyAPI)]

🔵 Suggestion: __all__ ordering in __init__.pyi — source/isaaclab/isaaclab/envs/mdp/__init__.pyi:61

set_gravity_compensation is inserted between randomize_rigid_body_scale and randomize_visual_color in __all__, breaking alphabetical ordering within the events group. Should come after all randomize_* entries (after randomize_visual_texture_material).

Test Coverage

• Bug fix (missing register_custom_attributes): The TestGravityCompensationNewtonModel tests serve as an end-to-end regression—they build a Newton model from USD with mjc: attributes, which would fail without registration. ✅
• New feature (MDP event): 6 unit tests covering uniform scale, dict-with-regex, joint regex, env_ids filtering, is_playing guard, and no-op case. Good coverage of the API surface.
• Gap: No test verifies that body_prims correctly filters only rigid bodies (the UsdPhysics.RigidBodyAPI(p) bug above). Adding a test with non-body child prims would catch this.
• Gap: No test covers the multi-env path in instantiate_builder_from_stage (the missing register_custom_attributes on proto builder).

CI Status

Only the labeler check has run (✅ passed). No lint/build/test CI results yet.

Verdict

REQUEST_CHANGES

Two critical issues: the missing import in __init__.pyi makes the function unusable via the public API, and the incomplete register_custom_attributes fix leaves the multi-env builder path broken. Both are straightforward one-line fixes.

---

Review by Isaac Lab Review Bot — [Expert Analysis] + [Error Handling Audit] + [Test Coverage Check]

[ooctipus]

can you rename it to set_newton_gravity_compensation?

[ooctipus]

LGTM just remember to change the name :D

[kellyguo11]

I think I would prefer this logic to live in the schemas config instead of mdps, it doesn't feel like it belongs as an event as it's adding USD attributes, which is what the schemas were designed to do. perhaps the proper way is to redesign the schemas to allow for Newton support and integrate this there.