From 8ec6832491779cc127680f27c76ad18abdc34b06 Mon Sep 17 00:00:00 2001 From: Julia von Muralt Date: Wed, 29 Apr 2026 14:31:36 +0200 Subject: [PATCH 01/35] add integration --- .../assets/articulation/base_articulation.py | 45 + .../isaaclab/isaaclab/envs/direct_marl_env.py | 35 +- .../isaaclab/isaaclab/envs/direct_rl_env.py | 35 +- .../isaaclab/envs/manager_based_env.py | 35 +- .../isaaclab/envs/manager_based_rl_env.py | 32 +- source/isaaclab/isaaclab/envs/mdp/events.py | 8 + .../isaaclab/physics/physics_manager.py | 22 + .../isaaclab/isaaclab/sim/simulation_cfg.py | 16 + .../actuators/newton_actuator_utils.py | 776 +++++++++++++++++ .../assets/articulation/articulation.py | 372 +++++--- .../assets/articulation/articulation_data.py | 4 + .../isaaclab_newton/physics/newton_manager.py | 268 ++++-- source/isaaclab_newton/setup.py | 3 +- .../test/assets/test_pd_equivalence.py | 818 ++++++++++++++++++ .../assets/articulation/articulation.py | 357 ++++---- .../test/assets/test_pd_equivalence.py | 554 ++++++++++++ 16 files changed, 2982 insertions(+), 398 deletions(-) create mode 100644 source/isaaclab_newton/isaaclab_newton/actuators/newton_actuator_utils.py create mode 100644 source/isaaclab_newton/test/assets/test_pd_equivalence.py create mode 100644 source/isaaclab_physx/test/assets/test_pd_equivalence.py diff --git a/source/isaaclab/isaaclab/assets/articulation/base_articulation.py b/source/isaaclab/isaaclab/assets/articulation/base_articulation.py index 14aa592ad103..12cba9b23e78 100644 --- a/source/isaaclab/isaaclab/assets/articulation/base_articulation.py +++ b/source/isaaclab/isaaclab/assets/articulation/base_articulation.py @@ -93,8 +93,53 @@ def __init__(self, cfg: ArticulationCfg): Args: cfg: A configuration instance. """ + from isaaclab.sim import SimulationContext # noqa: PLC0415 + super().__init__(cfg) + sim_ctx = SimulationContext.instance() + self._sim_cfg = sim_ctx.cfg if sim_ctx is not None else None + + self._author_newton_actuator_prims() + + # -- USD authoring helpers ------------------------------------------------ + + def _author_newton_actuator_prims(self) -> None: + """Author Newton actuator prims on the USD stage from Lab configs. + + Called from ``__init__`` after ``super().__init__()`` has spawned the + articulation prototype. When ``use_newton_actuators`` is enabled, + this translates explicit Lab actuator configs into ``NewtonActuator`` + USD prims. + + For every joint covered by a Lab config, any existing + ``NewtonActuator`` prim targeting that joint is replaced. Joints + not covered by any config keep their USD-authored actuators. + """ + if self._sim_cfg is None: + return + if not getattr(self._sim_cfg, "use_newton_actuators", False): + return + + try: + from isaaclab_newton.actuators.newton_actuator_utils import ( # noqa: PLC0415 + author_newton_actuator_prims, + ) + except ImportError: + return + + from isaaclab.sim.utils.queries import find_first_matching_prim # noqa: PLC0415 + + first_prim = find_first_matching_prim(self.cfg.prim_path) + if first_prim is None: + return + + author_newton_actuator_prims( + stage=self.stage, + articulation_prim_path=str(first_prim.GetPath()), + actuator_cfgs=self.cfg.actuators, + ) + """ Properties """ diff --git a/source/isaaclab/isaaclab/envs/direct_marl_env.py b/source/isaaclab/isaaclab/envs/direct_marl_env.py index e73b84628e1f..9b9e4f11c674 100644 --- a/source/isaaclab/isaaclab/envs/direct_marl_env.py +++ b/source/isaaclab/isaaclab/envs/direct_marl_env.py @@ -196,6 +196,10 @@ def _init_sim(self, render_mode: str | None = None, **kwargs): # this shouldn't cause an issue since later on, users do a reset over all the environments # so the lazy buffers would be reset. self.scene.update(dt=self.physics_dt) + # let the physics backend know about the env decimation so it can + # fold the full loop into a single step() when possible + self.sim.physics_manager.set_decimation(self.cfg.decimation) + self._physics_handles_decimation = self.sim.physics_manager.handles_decimation() # check if debug visualization is has been implemented by the environment source_code = inspect.getsource(self._set_debug_vis_impl) @@ -432,21 +436,30 @@ def step(self, actions: dict[AgentID, ActionType]) -> EnvStepReturn: is_rendering = self.sim.is_rendering # perform physics stepping - for _ in range(self.cfg.decimation): - self._sim_step_counter += 1 - # set actions into buffers + if self._physics_handles_decimation: + self._sim_step_counter += self.cfg.decimation self._apply_action() - # set actions into simulator self.scene.write_data_to_sim() - # simulate self.sim.step(render=False) - # render between steps only if the GUI or an RTX sensor needs it. - # When render_enabled is False, Kit visualizer (camera/GUI) is skipped - # but standalone visualizers (Newton, Rerun, Viser) still update. - if self._sim_step_counter % self.cfg.sim.render_interval == 0 and is_rendering: + if is_rendering: self.sim.render(skip_app_pumping=not self.render_enabled) - # update buffers at sim dt - self.scene.update(dt=self.physics_dt) + self.scene.update(dt=self.step_dt) + else: + for _ in range(self.cfg.decimation): + self._sim_step_counter += 1 + # set actions into buffers + self._apply_action() + # set actions into simulator + self.scene.write_data_to_sim() + # simulate + self.sim.step(render=False) + # render between steps only if the GUI or an RTX sensor needs it. + # When render_enabled is False, Kit visualizer (camera/GUI) is skipped + # but standalone visualizers (Newton, Rerun, Viser) still update. + if self._sim_step_counter % self.cfg.sim.render_interval == 0 and is_rendering: + self.sim.render(skip_app_pumping=not self.render_enabled) + # update buffers at sim dt + self.scene.update(dt=self.physics_dt) # post-step: # -- update env counters (used for curriculum generation) diff --git a/source/isaaclab/isaaclab/envs/direct_rl_env.py b/source/isaaclab/isaaclab/envs/direct_rl_env.py index 57442516bebf..312e95532d21 100644 --- a/source/isaaclab/isaaclab/envs/direct_rl_env.py +++ b/source/isaaclab/isaaclab/envs/direct_rl_env.py @@ -202,6 +202,10 @@ def _init_sim(self, render_mode: str | None = None, **kwargs): # this shouldn't cause an issue since later on, users do a reset over all the environments # so the lazy buffers would be reset. self.scene.update(dt=self.physics_dt) + # let the physics backend know about the env decimation so it can + # fold the full loop into a single step() when possible + self.sim.physics_manager.set_decimation(self.cfg.decimation) + self._physics_handles_decimation = self.sim.physics_manager.handles_decimation() # check if debug visualization is has been implemented by the environment source_code = inspect.getsource(self._set_debug_vis_impl) @@ -423,21 +427,30 @@ def step(self, action: torch.Tensor) -> VecEnvStepReturn: is_rendering = self.sim.is_rendering # perform physics stepping - for _ in range(self.cfg.decimation): - self._sim_step_counter += 1 - # set actions into buffers + if self._physics_handles_decimation: + self._sim_step_counter += self.cfg.decimation self._apply_action() - # set actions into simulator self.scene.write_data_to_sim() - # simulate self.sim.step(render=False) - # render between steps only if the GUI or an RTX sensor needs it. - # When render_enabled is False, Kit visualizer (camera/GUI) is skipped - # but standalone visualizers (Newton, Rerun, Viser) still update. - if self._sim_step_counter % self.cfg.sim.render_interval == 0 and is_rendering: + if is_rendering: self.sim.render(skip_app_pumping=not self.render_enabled) - # update buffers at sim dt - self.scene.update(dt=self.physics_dt) + self.scene.update(dt=self.step_dt) + else: + for _ in range(self.cfg.decimation): + self._sim_step_counter += 1 + # set actions into buffers + self._apply_action() + # set actions into simulator + self.scene.write_data_to_sim() + # simulate + self.sim.step(render=False) + # render between steps only if the GUI or an RTX sensor needs it. + # When render_enabled is False, Kit visualizer (camera/GUI) is skipped + # but standalone visualizers (Newton, Rerun, Viser) still update. + if self._sim_step_counter % self.cfg.sim.render_interval == 0 and is_rendering: + self.sim.render(skip_app_pumping=not self.render_enabled) + # update buffers at sim dt + self.scene.update(dt=self.physics_dt) # post-step: # -- update env counters (used for curriculum generation) diff --git a/source/isaaclab/isaaclab/envs/manager_based_env.py b/source/isaaclab/isaaclab/envs/manager_based_env.py index 4209e6a5c53a..a3039dc696eb 100644 --- a/source/isaaclab/isaaclab/envs/manager_based_env.py +++ b/source/isaaclab/isaaclab/envs/manager_based_env.py @@ -215,6 +215,10 @@ def _init_sim(self): # this shouldn't cause an issue since later on, users do a reset over all the environments # so the lazy buffers would be reset. self.scene.update(dt=self.physics_dt) + # let the physics backend know about the env decimation so it can + # fold the full loop into a single step() when possible + self.sim.physics_manager.set_decimation(self.cfg.decimation) + self._physics_handles_decimation = self.sim.physics_manager.handles_decimation() # add timeline event to load managers self.load_managers() @@ -529,21 +533,30 @@ def step(self, action: torch.Tensor) -> tuple[VecEnvObs, dict]: is_rendering = self.sim.is_rendering # perform physics stepping - for _ in range(self.cfg.decimation): - self._sim_step_counter += 1 - # set actions into buffers + if self._physics_handles_decimation: + self._sim_step_counter += self.cfg.decimation self.action_manager.apply_action() - # set actions into simulator self.scene.write_data_to_sim() - # simulate self.sim.step(render=False) - # render between steps only if the GUI or an RTX sensor needs it. - # When render_enabled is False, Kit visualizer (camera/GUI) is skipped - # but standalone visualizers (Newton, Rerun, Viser) still update. - if self._sim_step_counter % self.cfg.sim.render_interval == 0 and is_rendering: + if is_rendering: self.sim.render(skip_app_pumping=not self.render_enabled) - # update buffers at sim dt - self.scene.update(dt=self.physics_dt) + self.scene.update(dt=self.step_dt) + else: + for _ in range(self.cfg.decimation): + self._sim_step_counter += 1 + # set actions into buffers + self.action_manager.apply_action() + # set actions into simulator + self.scene.write_data_to_sim() + # simulate + self.sim.step(render=False) + # render between steps only if the GUI or an RTX sensor needs it. + # When render_enabled is False, Kit visualizer (camera/GUI) is skipped + # but standalone visualizers (Newton, Rerun, Viser) still update. + if self._sim_step_counter % self.cfg.sim.render_interval == 0 and is_rendering: + self.sim.render(skip_app_pumping=not self.render_enabled) + # update buffers at sim dt + self.scene.update(dt=self.physics_dt) # post-step: step interval event if "interval" in self.event_manager.available_modes: diff --git a/source/isaaclab/isaaclab/envs/manager_based_rl_env.py b/source/isaaclab/isaaclab/envs/manager_based_rl_env.py index 4c3b329b0387..0a4c85caf49f 100644 --- a/source/isaaclab/isaaclab/envs/manager_based_rl_env.py +++ b/source/isaaclab/isaaclab/envs/manager_based_rl_env.py @@ -199,22 +199,32 @@ def step(self, action: torch.Tensor) -> VecEnvStepReturn: is_rendering = self.sim.is_rendering # perform physics stepping - for _ in range(self.cfg.decimation): - self._sim_step_counter += 1 - # set actions into buffers + if self._physics_handles_decimation: + self._sim_step_counter += self.cfg.decimation self.action_manager.apply_action() - # set actions into simulator self.scene.write_data_to_sim() - # simulate self.sim.step(render=False) self.recorder_manager.record_post_physics_decimation_step() - # render between steps only if the GUI or an RTX sensor needs it. - # When render_enabled is False, Kit visualizer (camera/GUI) is skipped - # but standalone visualizers (Newton, Rerun, Viser) still update. - if self._sim_step_counter % self.cfg.sim.render_interval == 0 and is_rendering: + if is_rendering: self.sim.render(skip_app_pumping=not self.render_enabled) - # update buffers at sim dt - self.scene.update(dt=self.physics_dt) + self.scene.update(dt=self.step_dt) + else: + for _ in range(self.cfg.decimation): + self._sim_step_counter += 1 + # set actions into buffers + self.action_manager.apply_action() + # set actions into simulator + self.scene.write_data_to_sim() + # simulate + self.sim.step(render=False) + self.recorder_manager.record_post_physics_decimation_step() + # render between steps only if the GUI or an RTX sensor needs it. + # When render_enabled is False, Kit visualizer (camera/GUI) is skipped + # but standalone visualizers (Newton, Rerun, Viser) still update. + if self._sim_step_counter % self.cfg.sim.render_interval == 0 and is_rendering: + self.sim.render(skip_app_pumping=not self.render_enabled) + # update buffers at sim dt + self.scene.update(dt=self.physics_dt) # post-step: # -- update env counters (used for curriculum generation) diff --git a/source/isaaclab/isaaclab/envs/mdp/events.py b/source/isaaclab/isaaclab/envs/mdp/events.py index 2209ea5a28f5..235da296b2c5 100644 --- a/source/isaaclab/isaaclab/envs/mdp/events.py +++ b/source/isaaclab/isaaclab/envs/mdp/events.py @@ -1210,6 +1210,10 @@ def randomize(data: torch.Tensor, params: tuple[float, float]) -> torch.Tensor: self.asset.write_joint_stiffness_to_sim_index( stiffness=stiffness, joint_ids=actuator.joint_indices, env_ids=env_ids ) + elif hasattr(self.asset, "write_actuator_stiffness_to_sim"): + self.asset.write_actuator_stiffness_to_sim( + actuator, stiffness=stiffness, env_ids=env_ids + ) # Randomize damping if damping_distribution_params is not None: damping = actuator.damping[env_ids].clone() @@ -1220,6 +1224,10 @@ def randomize(data: torch.Tensor, params: tuple[float, float]) -> torch.Tensor: self.asset.write_joint_damping_to_sim_index( damping=damping, joint_ids=actuator.joint_indices, env_ids=env_ids ) + elif hasattr(self.asset, "write_actuator_damping_to_sim"): + self.asset.write_actuator_damping_to_sim( + actuator, damping=damping, env_ids=env_ids + ) class randomize_joint_parameters(ManagerTermBase): diff --git a/source/isaaclab/isaaclab/physics/physics_manager.py b/source/isaaclab/isaaclab/physics/physics_manager.py index 7a4cdfe84403..ca333673bb48 100644 --- a/source/isaaclab/isaaclab/physics/physics_manager.py +++ b/source/isaaclab/isaaclab/physics/physics_manager.py @@ -339,6 +339,28 @@ def wait_for_playing(cls) -> None: """Block until the timeline is playing. Default is no-op.""" pass + @classmethod + def set_decimation(cls, decimation: int) -> None: + """Inform the physics backend how many substeps the environment runs per policy step. + + Backends that can fold the full decimation loop into a single + :meth:`step` call (e.g. Newton with all-graphable actuators) use this + to size their internal loop / CUDA graph. The default implementation + is a no-op. + + Args: + decimation: Number of physics steps per environment step. + """ + + @classmethod + def handles_decimation(cls) -> bool: + """``True`` when :meth:`step` executes the full decimation loop internally. + + When this returns ``True`` the environment should call :meth:`step` + once per policy step instead of looping ``decimation`` times. + """ + return False + @classmethod def get_backend(cls) -> str: """Get the tensor backend being used ("numpy" or "torch").""" diff --git a/source/isaaclab/isaaclab/sim/simulation_cfg.py b/source/isaaclab/isaaclab/sim/simulation_cfg.py index 6c2e97bb7e82..3812ff3b1b50 100644 --- a/source/isaaclab/isaaclab/sim/simulation_cfg.py +++ b/source/isaaclab/isaaclab/sim/simulation_cfg.py @@ -231,6 +231,22 @@ class SimulationCfg: with the GUI enabled. This is to allow certain GUI features to work properly. """ + use_newton_actuators: bool = False + """Use Newton-native actuators instead of IsaacLab explicit actuator models. + + When ``True``, explicit actuator configs (e.g. :class:`IdealPDActuatorCfg`, + :class:`DCMotorCfg`) are translated into ``NewtonActuator`` USD prims and + stepped by the physics engine. The Lab config values (stiffness, damping, + effort_limit, etc.) take precedence: for every joint covered by a Lab + actuator config, any existing ``NewtonActuator`` prim targeting that joint + is replaced by one synthesised from the config. Joints that are *not* + covered by a Lab config keep their USD-authored actuators (if any). + + :class:`ImplicitActuatorCfg` entries are still instantiated normally and + their gains are written to the simulation, so joints that use implicit + actuation continue to work as expected. + """ + physics: PhysicsCfg | None = None """Physics manager configuration. Default is None (uses PhysxCfg()). diff --git a/source/isaaclab_newton/isaaclab_newton/actuators/newton_actuator_utils.py b/source/isaaclab_newton/isaaclab_newton/actuators/newton_actuator_utils.py new file mode 100644 index 000000000000..159779a8addc --- /dev/null +++ b/source/isaaclab_newton/isaaclab_newton/actuators/newton_actuator_utils.py @@ -0,0 +1,776 @@ +# Copyright (c) 2022-2026, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +"""Utilities for Newton-native actuators in Isaac Lab. + +This module is organised into four sections: + +1. **Warp kernels** — low-level device kernels for zeroing, masking, + and scattering per-DOF data used by the adapter. +2. **PhysX stepping helper** — :class:`PhysxActuatorWrapper`, a + duck-typed wrapper that exposes flat Warp arrays as the + ``sim_state`` / ``sim_control`` protocol expected by + :meth:`Actuator.step` on the PhysX backend. +3. **Adapter** — :class:`NewtonActuatorAdapter` manages actuator + creation, DOF-to-actuator mapping, config overrides, stepping, + reset, and gain reading for domain randomisation. Used identically + by both Newton and PhysX backends. +4. **USD authoring** — :func:`author_newton_actuator_prims` creates + ``NewtonActuator`` USD prims from IsaacLab actuator configs so that + both the Newton ``ModelBuilder`` (during ``add_usd``) and the PhysX + adapter (via :meth:`NewtonActuatorAdapter.from_usd`) can construct + :class:`Actuator` objects with correct parameters. + +Why :class:`PhysxActuatorWrapper` exists only for PhysX +~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + +Newton's :meth:`Actuator.step` requires a ``sim_state`` / ``sim_control`` +pair that exposes flat 1-D Warp arrays (``joint_q``, ``joint_qd``, +``joint_target_pos``, ``joint_f``, …). On the **Newton backend** these +are the ``State`` and ``Control`` objects that the solver already owns — +no wrapper is needed because: + +* The solver manages double-buffered ``State`` objects for CUDA-graph + capture, and actuators are stepped inside the solver's own simulation + loop where states are already available. +* Wrapping them would add indirection with no benefit; the Newton + articulation code that calls :meth:`Actuator.step` lives in + ``newton_manager.py`` and has direct access to the model's state. + +On the **PhysX backend**, no Newton solver exists — the actuators are +stepped manually from the Lab articulation's ``write_data_to_sim`` +path. Isaac Lab stores joint data as 2-D tensors (``num_envs × +num_joints``), so :class:`PhysxActuatorWrapper` provides zero-copy flat +views that satisfy the protocol without allocating new memory. The +PhysX articulation code that calls :meth:`Actuator.step` lives in +``isaaclab_physx/.../articulation.py``, completely separate from the +Newton solver path, so sharing a single wrapper type across both +backends would not reduce code — it would only add coupling. +""" + +from __future__ import annotations + +from collections.abc import Sequence +from dataclasses import dataclass +from typing import Any + +import numpy as np +import torch +import warp as wp + +from newton.actuators import ( + Actuator, + Clamping, + Delay, +) + + +# =========================================================================== +# 1. Warp kernels +# =========================================================================== + + +@wp.kernel(enable_backward=False) +def _zero_at_indices_kernel(data: wp.array(dtype=wp.float32), indices: wp.array(dtype=wp.uint32)): + i = wp.tid() + data[indices[i]] = 0.0 + + +@wp.kernel(enable_backward=False) +def _set_mask_kernel(mask: wp.array(dtype=wp.bool), indices: wp.array(dtype=wp.int32)): + i = wp.tid() + mask[indices[i]] = True + + +@wp.kernel(enable_backward=False) +def _scatter_gain_kernel( + src: wp.array(dtype=wp.float32), + dst: wp.array(dtype=wp.float32), + indices: wp.array(dtype=wp.uint32), + dof_offset: int, + num_joints: int, +): + i = wp.tid() + global_dof = int(indices[i]) - dof_offset + env = global_dof // num_joints + local_dof = global_dof % num_joints + dst[env * num_joints + local_dof] = src[i] + + +# =========================================================================== +# 2. PhysX stepping helper +# =========================================================================== + + +@dataclass +class PhysxActuatorWrapper: + """Flat-array wrapper serving as ``sim_state`` / ``sim_control`` for + :meth:`Actuator.step` on the PhysX backend. + + Only needed on PhysX — on the Newton backend the solver's own + ``State`` / ``Control`` objects fulfil this role directly (see the + module docstring for rationale). + + Attributes are reassigned each frame to zero-copy flat views of the + existing 2-D Isaac Lab data via ``wp.array.reshape(-1)``. Only + ``joint_f`` is a persistent allocation (zeroed before each step). + """ + + joint_q: wp.array | None = None + joint_qd: wp.array | None = None + joint_target_pos: wp.array | None = None + joint_target_vel: wp.array | None = None + joint_act: wp.array | None = None + joint_f: wp.array | None = None + + +# =========================================================================== +# 3. Adapter — creation, config overrides, stepping, domain randomisation +# =========================================================================== + + +class NewtonActuatorAdapter: + """Manages Newton-native actuators for both Newton and PhysX backends. + + Handles actuator creation (from USD or an existing list), + DOF-to-actuator mapping, stepping, reset, and gain reading for + domain randomisation. + + Construction: + + * **Newton backend** — pass actuators from the Newton model directly:: + + adapter = NewtonActuatorAdapter(model.actuators, num_envs, + num_joints, dof_offset, device) + + * **PhysX backend** — create actuators from USD prims:: + + adapter = NewtonActuatorAdapter.from_usd(stage, joint_names, + num_envs, num_joints, + device) + + Then finalise:: + + adapter.finalize() + + After :meth:`finalize`, the adapter exposes ``.stiffness``, + ``.damping``, and ``.joint_indices`` for ``randomize_actuator_gains``. + """ + + def __init__( + self, + actuators: list[Actuator], + num_envs: int, + num_joints: int, + dof_offset: int, + device: str, + ): + self.actuators = actuators + self.num_joints = num_joints + + self._num_envs = num_envs + self._dof_offset = dof_offset + self._device = device + self._dof_to_actuator = self._build_dof_map() + + managed = [i for i, act_idx in enumerate(self._dof_to_actuator) if act_idx >= 0] + if len(managed) == num_joints: + self.joint_indices: torch.Tensor | slice = slice(None) + else: + self.joint_indices = torch.tensor(managed, dtype=torch.int32, device=device) + + self._states_a = [act.state() for act in actuators] + self._states_b = [act.state() for act in actuators] + + self.stiffness: torch.Tensor | None = None + self.damping: torch.Tensor | None = None + + # -- construction (PhysX path) ------------------------------------------- + + @classmethod + def from_usd( + cls, + stage: Any, + joint_names: list[str], + num_envs: int, + num_joints: int, + device: str, + ) -> NewtonActuatorAdapter: + """Create an adapter by parsing ``NewtonActuator`` prims from USD. + + This is the PhysX-backend counterpart of what Newton's + ``ModelBuilder.add_usd`` does for the Newton backend. Both paths + read the same ``NewtonActuator`` USD prims (authored by + :func:`author_newton_actuator_prims`) and construct + :class:`~newton.actuators.Actuator` objects with matching + controllers, clampings, and delays. + + The key difference is that PhysX uses a **flat per-DOF layout** + where joint position coordinates and velocity DOFs always have the + same count and ordering — there are no free joints or ball joints + that cause coordinate/DOF count divergence. Therefore a single + ``indices`` array is used for all index roles (``indices``, + ``pos_indices``, ``target_pos_indices``), unlike the Newton + builder which computes separate ``pos_indices`` from + ``joint_q_start`` and separate ``target_pos_indices`` from + ``joint_qd_start`` to handle floating-base articulations. + + Joints whose prims resolve to the same controller type, gains, + clamping chain, and delay configuration are merged into a single + :class:`Actuator` with combined index arrays, mirroring the + grouping the Newton builder performs internally. + + Args: + stage: The USD stage containing ``NewtonActuator`` prims. + joint_names: All joint names in the articulation. + num_envs: Number of environments. + num_joints: Joints per environment in the articulation. + device: Warp device string (e.g. ``"cuda:0"``). + + Returns: + A fully constructed adapter ready for :meth:`finalize`. + """ + actuators = cls._create_actuators_from_usd( + stage, joint_names, num_envs, num_joints, device, + ) + return cls(actuators, num_envs, num_joints, dof_offset=0, device=device) + + # -- public API ---------------------------------------------------------- + + def finalize(self) -> None: + """Read actuator gains and store as PyTorch tensors for DR.""" + wp_device = wp.get_device(self._device) + flat_stiffness = wp.zeros(self._num_envs * self.num_joints, dtype=wp.float32, device=wp_device) + flat_damping = wp.zeros(self._num_envs * self.num_joints, dtype=wp.float32, device=wp_device) + + for act in self.actuators: + ctrl = act.controller + if hasattr(ctrl, "kp"): + wp.launch( + _scatter_gain_kernel, dim=act.indices.shape[0], + inputs=[ctrl.kp, flat_stiffness, act.indices, self._dof_offset, self.num_joints], + device=wp_device, + ) + if hasattr(ctrl, "kd"): + wp.launch( + _scatter_gain_kernel, dim=act.indices.shape[0], + inputs=[ctrl.kd, flat_damping, act.indices, self._dof_offset, self.num_joints], + device=wp_device, + ) + + self.stiffness = wp.to_torch(flat_stiffness.reshape((self._num_envs, self.num_joints))) + self.damping = wp.to_torch(flat_damping.reshape((self._num_envs, self.num_joints))) + + def step(self, sim_state: Any, sim_control: Any, dt: float) -> None: + """Zero actuated DOFs, step all actuators, and swap state buffers. + + Args: + sim_state: Object with ``joint_q``, ``joint_qd``, etc. + Newton ``State`` on the Newton backend, + :class:`PhysxActuatorWrapper` on the PhysX backend. + sim_control: Object with ``joint_f``, ``joint_target_pos``, etc. + Newton ``Control`` on the Newton backend, + :class:`PhysxActuatorWrapper` on the PhysX backend. + dt: Physics timestep [s]. + """ + for act in self.actuators: + wp.launch( + _zero_at_indices_kernel, + dim=act.indices.shape[0], + inputs=[sim_control.joint_f, act.indices], + ) + for act, sa, sb in zip(self.actuators, self._states_a, self._states_b): + act.step(sim_state, sim_control, sa, sb, dt=dt) + self._states_a, self._states_b = self._states_b, self._states_a + + def reset(self, env_ids: Sequence[int] | slice | None = None) -> None: + """Reset actuator states for the given environments. + + Args: + env_ids: Environment indices to reset. ``None`` or + ``slice(None)`` resets all environments. + """ + if env_ids is None or env_ids == slice(None): + mask = None + else: + mask = wp.zeros(self._num_envs, dtype=wp.bool, device=self._device) + idx = wp.array(list(env_ids), dtype=wp.int32, device=self._device) + wp.launch(_set_mask_kernel, dim=len(idx), inputs=[mask, idx], device=self._device) + + for sa, sb in zip(self._states_a, self._states_b): + if sa is not None: + sa.reset(mask) + if sb is not None: + sb.reset(mask) + + @property + def is_all_graphable(self) -> bool: + """``True`` when all actuators are CUDA-graph-safe.""" + return len(self.actuators) > 0 and all(a.is_graphable() for a in self.actuators) + + # -- config helpers (used by both adapter and authoring) ----------------- + + @staticmethod + def _resolve_per_dof( + value: dict[str, float | int] | float | int | None, + joint_names: list[str], + cast: type = float, + ) -> dict[str, float | int]: + """Expand a scalar or dict config value into a per-DOF dict. + + When *value* is a dict, keys are treated as regex patterns and + matched against *joint_names* via :func:`re.fullmatch`. + """ + import re # noqa: PLC0415 + + if value is None: + return {} + if isinstance(value, (int, float)): + return {name: cast(value) for name in joint_names} + if isinstance(value, dict): + result: dict[str, float | int] = {} + for name in joint_names: + for pattern, v in value.items(): + if re.fullmatch(pattern, name): + result[name] = cast(v) + break + return result + return {} + + # -- private helpers ----------------------------------------------------- + + def _build_dof_map(self) -> list[int]: + """Build a per-DOF lookup: local DOF index -> actuator list index.""" + dof_to_actuator: list[int] = [-1] * self.num_joints + + for act_idx, act in enumerate(self.actuators): + all_indices = act.indices.numpy() + num_per_act = len(all_indices) // self._num_envs + env0_indices = all_indices[:num_per_act] + for global_dof in env0_indices: + local_dof = global_dof - self._dof_offset + if 0 <= local_dof < self.num_joints: + dof_to_actuator[local_dof] = act_idx + + return dof_to_actuator + + @staticmethod + def _actuator_signature(parsed: Any) -> tuple: + """Build a hashable key from a parsed actuator spec for grouping. + + Joints whose prims resolve to the same signature share identical + controller type, gains, clamping chain, and delay configuration and + can therefore be merged into a single :class:`Actuator` with combined + index arrays. + """ + ctrl_resolved = parsed.controller_class.resolve_arguments( + dict(parsed.controller_kwargs), + ) + ctrl_key = (parsed.controller_class, tuple(sorted(ctrl_resolved.items()))) + + comp_keys: list[tuple] = [] + for comp_cls, comp_kwargs in parsed.component_specs: + resolved = comp_cls.resolve_arguments(comp_kwargs) + comp_keys.append((comp_cls, tuple(sorted(resolved.items())))) + comp_keys.sort(key=lambda t: t[0].__name__) + + return (ctrl_key, tuple(comp_keys)) + + @staticmethod + def _create_actuators_from_usd( + stage: Any, + joint_names: list[str], + num_envs: int, + num_total_joints: int, + device: str, + ) -> list[Actuator]: + """Parse ``NewtonActuator`` prims and instantiate standalone actuators. + + This mirrors the actuator construction that Newton's + ``ModelBuilder.add_usd`` performs, but operates independently of a + Newton ``Model``. It is used on the PhysX backend where there is no + Newton simulation — actuators are stepped manually via the adapter. + + Because PhysX articulations have no free or ball joints, every + joint's coordinate count equals its DOF count. A single + ``indices`` array is therefore sufficient for all index roles + (``indices``, ``pos_indices``, ``target_pos_indices``). + + Joints with identical controller type, gains, clamping chain, and + delay are merged into one :class:`Actuator` with combined indices. + + Each per-DOF scalar parameter (``kp``, ``kd``, ``saturation_effort``, + etc.) is broadcast via :func:`wp.full` to match the group size. + Parameters marked as ``SHARED_PARAMS`` on the controller or clamping + class (e.g. ``model_path``, ``lookup_positions``) are passed through + directly without broadcast. + """ + from collections import defaultdict # noqa: PLC0415 + + from newton.actuators import parse_actuator_prim # noqa: PLC0415 + from pxr import Usd # noqa: PLC0415 + + wp_device = wp.get_device(device) + + joint_name_to_idx: dict[str, int] = {name: i for i, name in enumerate(joint_names)} + + parsed_per_joint: dict[int, Any] = {} + for prim in Usd.PrimRange(stage.GetPseudoRoot()): + parsed = parse_actuator_prim(prim) + if parsed is None: + continue + target_name = parsed.target_path.rsplit("/", 1)[-1] + if target_name in joint_name_to_idx: + parsed_per_joint[joint_name_to_idx[target_name]] = parsed + + if not parsed_per_joint: + raise ValueError( + f"No NewtonActuator prims found targeting any of: {joint_names}" + ) + + groups: dict[tuple, list[int]] = defaultdict(list) + sig_to_parsed: dict[tuple, Any] = {} + for local_idx, parsed in sorted(parsed_per_joint.items()): + sig = NewtonActuatorAdapter._actuator_signature(parsed) + groups[sig].append(local_idx) + if sig not in sig_to_parsed: + sig_to_parsed[sig] = parsed + + actuators = [] + for sig, local_indices in groups.items(): + parsed = sig_to_parsed[sig] + + flat_indices = np.array( + [idx + e * num_total_joints for e in range(num_envs) for idx in local_indices], + dtype=np.uint32, + ) + indices = wp.array(flat_indices, device=wp_device) + num_dofs_in_group = len(local_indices) * num_envs + + # Controller + ctrl_kwargs = dict(parsed.controller_kwargs) + resolved = parsed.controller_class.resolve_arguments(ctrl_kwargs) + shared_ctrl = getattr(parsed.controller_class, "SHARED_PARAMS", set()) + ctrl_arrays = {} + for key, val in resolved.items(): + if key in shared_ctrl: + ctrl_arrays[key] = val + else: + ctrl_arrays[key] = wp.full(num_dofs_in_group, float(val), dtype=wp.float32, device=wp_device) + controller = parsed.controller_class(**ctrl_arrays) + + # Components (delay + clampings) + clampings = [] + delay = None + for comp_cls, comp_kwargs in parsed.component_specs: + if issubclass(comp_cls, Delay): + resolved_kw = Delay.resolve_arguments(comp_kwargs) + delay_steps = int(resolved_kw.get("delay_steps", 0)) + if delay_steps > 0: + delay_arr = wp.full(num_dofs_in_group, delay_steps, dtype=wp.int32, device=wp_device) + delay = Delay(delay_steps=delay_arr, max_delay=delay_steps) + elif issubclass(comp_cls, Clamping): + resolved_kw = comp_cls.resolve_arguments(comp_kwargs) + shared_clamp = getattr(comp_cls, "SHARED_PARAMS", set()) + clamp_arrays = {} + for k, v in resolved_kw.items(): + if k in shared_clamp: + clamp_arrays[k] = v + else: + clamp_arrays[k] = wp.full( + num_dofs_in_group, float(v), dtype=wp.float32, device=wp_device, + ) + clampings.append(comp_cls(**clamp_arrays)) + + actuator = Actuator( + indices=indices, + controller=controller, + delay=delay, + clamping=clampings if clampings else None, + ) + actuators.append(actuator) + + return actuators + + +# =========================================================================== +# 4. USD authoring — create NewtonActuator prims from Lab actuator configs +# =========================================================================== + + +def author_newton_actuator_prims( + stage: Any, + articulation_prim_path: str, + actuator_cfgs: dict[str, Any], +) -> None: + """Author ``NewtonActuator`` USD prims from IsaacLab actuator configs. + + For every joint covered by an explicit (non-implicit) Lab actuator config, + any existing ``NewtonActuator`` prim targeting that joint is replaced by a + new one created from the config values. Joints **not** covered by any + Lab config keep their USD-authored actuators unchanged. + + The supported config-to-schema mapping is: + + * :class:`~isaaclab.actuators.IdealPDActuatorCfg` -> + ``NewtonPDControlAPI`` + ``NewtonMaxEffortClampingAPI`` + * :class:`~isaaclab.actuators.DCMotorCfg` -> + ``NewtonPDControlAPI`` + ``NewtonDCMotorClampingAPI`` + * :class:`~isaaclab.actuators.DelayedPDActuatorCfg` -> + same as ``IdealPDActuatorCfg`` + ``NewtonActuatorDelayAPI`` + * :class:`~isaaclab.actuators.RemotizedPDActuatorCfg` -> + same as ``DelayedPDActuatorCfg`` + ``NewtonPositionBasedClampingAPI`` + * :class:`~isaaclab.actuators.ActuatorNetMLPCfg` / + :class:`~isaaclab.actuators.ActuatorNetLSTMCfg` -> + ``NewtonNeuralControlAPI`` (+ ``NewtonDCMotorClampingAPI``) + + Must be called **after** the articulation is spawned (joint prims exist + on stage) and **before** the cloner / ``ModelBuilder.add_usd`` reads + the stage. + + Args: + stage: The USD stage to author prims on. + articulation_prim_path: Root prim path of the articulation + (e.g. ``"/World/Env_0/Robot"``). Must not contain wildcards. + actuator_cfgs: Mapping of group name to ``ActuatorBaseCfg``. + """ + from pxr import Sdf # noqa: PLC0415 + + from isaaclab.actuators import ImplicitActuator # noqa: PLC0415 + from isaaclab.utils.string import resolve_matching_names # noqa: PLC0415 + + art_prim = stage.GetPrimAtPath(articulation_prim_path) + if not art_prim.IsValid(): + raise ValueError(f"Articulation prim not found: {articulation_prim_path}") + + joint_inventory = _collect_joint_prims(art_prim) + all_joint_names = list(joint_inventory.keys()) + + covered_joint_paths: set[str] = set() + resolve = NewtonActuatorAdapter._resolve_per_dof + + cfg_entries: list[tuple[str, Any, list[str]]] = [] + for group_name, cfg in actuator_cfgs.items(): + cls_type = cfg.class_type + is_implicit = ( + "ImplicitActuator" in cls_type + if isinstance(cls_type, str) + else issubclass(cls_type, ImplicitActuator) + ) + if is_implicit: + continue + + _ids, joint_names = resolve_matching_names(cfg.joint_names_expr, all_joint_names) + if not joint_names: + continue + + cfg_entries.append((group_name, cfg, joint_names)) + for jname in joint_names: + covered_joint_paths.add(joint_inventory[jname]) + + _remove_actuator_prims_for_joints(art_prim, covered_joint_paths) + + from isaaclab.actuators import DCMotorCfg, DelayedPDActuatorCfg # noqa: PLC0415 + from isaaclab.actuators.actuator_net_cfg import ActuatorNetLSTMCfg, ActuatorNetMLPCfg # noqa: PLC0415 + from isaaclab.actuators.actuator_pd_cfg import RemotizedPDActuatorCfg # noqa: PLC0415 + + for group_name, cfg, joint_names in cfg_entries: + stiffness_map = resolve(getattr(cfg, "stiffness", None), joint_names) + damping_map = resolve(getattr(cfg, "damping", None), joint_names) + effort_map = resolve(getattr(cfg, "effort_limit", None), joint_names) + + is_neural = isinstance(cfg, (ActuatorNetMLPCfg, ActuatorNetLSTMCfg)) + is_remotized = isinstance(cfg, RemotizedPDActuatorCfg) + is_dc_motor = isinstance(cfg, DCMotorCfg) + is_delayed = isinstance(cfg, DelayedPDActuatorCfg) + + vel_limit_map = resolve(getattr(cfg, "velocity_limit", None), joint_names) if is_dc_motor else {} + sat_effort_map = resolve(getattr(cfg, "saturation_effort", None), joint_names) if is_dc_motor else {} + + raw_delay = getattr(cfg, "max_delay", 0) if is_delayed else 0 + delay_map = resolve(raw_delay, joint_names, cast=int) if raw_delay else {} + + patched_model_path: str | None = None + if is_neural: + meta: dict[str, Any] = {} + if isinstance(cfg, ActuatorNetMLPCfg): + meta["model_type"] = "mlp" + meta["input_order"] = cfg.input_order + meta["input_idx"] = list(cfg.input_idx) + meta["pos_scale"] = cfg.pos_scale + meta["vel_scale"] = cfg.vel_scale + meta["torque_scale"] = cfg.torque_scale + else: + meta["model_type"] = "lstm" + patched_model_path = _resave_checkpoint_with_metadata(cfg.network_file, meta) + + for jname in joint_names: + joint_prim_path = joint_inventory[jname] + + schemas: list[str] = [] + attrs: dict[str, float | int] = {} + array_attrs: dict[str, list[float]] = {} + + if is_neural: + schemas.append("NewtonNeuralControlAPI") + else: + schemas.append("NewtonPDControlAPI") + attrs["kp"] = stiffness_map.get(jname, 0.0) + attrs["kd"] = damping_map.get(jname, 0.0) + + if is_dc_motor: + schemas.append("NewtonDCMotorClampingAPI") + attrs["saturation_effort"] = sat_effort_map.get(jname, 0.0) + if jname in vel_limit_map: + attrs["velocity_limit"] = vel_limit_map[jname] + if jname in effort_map: + attrs["max_motor_effort"] = effort_map[jname] + elif jname in effort_map: + schemas.append("NewtonMaxEffortClampingAPI") + attrs["max_effort"] = effort_map[jname] + + if is_remotized and isinstance(cfg, RemotizedPDActuatorCfg): + lookup = cfg.joint_parameter_lookup + schemas.append("NewtonPositionBasedClampingAPI") + array_attrs["lookup_positions"] = [row[0] for row in lookup] + array_attrs["lookup_efforts"] = [row[2] for row in lookup] + + delay_steps = delay_map.get(jname, 0) + if delay_steps > 0: + schemas.append("NewtonActuatorDelayAPI") + attrs["delay_steps"] = delay_steps + attrs["max_delay"] = delay_steps + + act_prim_path = f"{articulation_prim_path}/{group_name}_{jname}_actuator" + act_prim = stage.DefinePrim(act_prim_path, "NewtonActuator") + + existing = act_prim.GetMetadata("apiSchemas") or Sdf.TokenListOp() + existing.prependedItems = list(schemas) + act_prim.SetMetadata("apiSchemas", existing) + + rel = act_prim.CreateRelationship("newton:targets") + rel.SetTargets([Sdf.Path(joint_prim_path)]) + + if patched_model_path is not None: + act_prim.CreateAttribute("newton:modelPath", Sdf.ValueTypeNames.Asset).Set( + Sdf.AssetPath(patched_model_path) + ) + + for attr_name, attr_val in attrs.items(): + usd_name = f"newton:{_snake_to_camel(attr_name)}" + if isinstance(attr_val, int): + act_prim.CreateAttribute(usd_name, Sdf.ValueTypeNames.Int).Set(attr_val) + else: + act_prim.CreateAttribute(usd_name, Sdf.ValueTypeNames.Float).Set(float(attr_val)) + + for attr_name, attr_val in array_attrs.items(): + usd_name = f"newton:{_snake_to_camel(attr_name)}" + act_prim.CreateAttribute(usd_name, Sdf.ValueTypeNames.FloatArray).Set(attr_val) + + +# --------------------------------------------------------------------------- +# USD authoring — private helpers +# --------------------------------------------------------------------------- + +_SNAKE_TO_CAMEL_RE = __import__("re").compile(r"_([a-z])") + + +def _snake_to_camel(name: str) -> str: + """Convert a snake_case name to camelCase.""" + return _SNAKE_TO_CAMEL_RE.sub(lambda m: m.group(1).upper(), name) + + +def _collect_joint_prims(art_prim: Any) -> dict[str, str]: + """Collect all joint prims under an articulation subtree. + + Returns: + Ordered mapping of joint name to full prim path. + """ + from pxr import Usd # noqa: PLC0415 + + _JOINT_TYPES = {"PhysicsRevoluteJoint", "PhysicsPrismaticJoint"} + + joints: dict[str, str] = {} + for prim in Usd.PrimRange(art_prim): + if prim.GetTypeName() in _JOINT_TYPES: + joints[prim.GetName()] = str(prim.GetPath()) + return joints + + +def _remove_actuator_prims_for_joints( + art_prim: Any, + joint_paths: set[str], +) -> None: + """Deactivate ``NewtonActuator`` prims whose target is in *joint_paths*. + + Deactivated prims are invisible to ``Usd.PrimRange`` and therefore + ignored by ``ModelBuilder.add_usd``. Using ``SetActive(False)`` + instead of ``RemovePrim`` works correctly when the prim originates + from a USD reference or payload. + + Only prims under the *art_prim* subtree are considered. + """ + from pxr import Usd # noqa: PLC0415 + + to_deactivate: list = [] + for prim in Usd.PrimRange(art_prim): + if prim.GetTypeName() != "NewtonActuator": + continue + rel = prim.GetRelationship("newton:targets") + if rel and rel.IsValid(): + for target in rel.GetTargets(): + if str(target) in joint_paths: + to_deactivate.append(prim) + break + + for prim in to_deactivate: + prim.SetActive(False) + + +def _resave_checkpoint_with_metadata( + original_path: str, + metadata: dict[str, Any], +) -> str: + """Re-save a neural-network checkpoint with updated metadata. + + Loads the original TorchScript or dict checkpoint, merges *metadata* + into any existing metadata (Lab config values take precedence), and + writes the result to a temporary ``.pt`` file that persists for the + lifetime of the process. + + Returns: + Path to the temporary checkpoint file. + """ + import json + import tempfile + + import torch + + extra_files: dict[str, str] = {"metadata.json": ""} + is_torchscript = True + try: + net = torch.jit.load(original_path, map_location="cpu", _extra_files=extra_files) + existing_meta = json.loads(extra_files["metadata.json"]) if extra_files["metadata.json"] else {} + except Exception: + is_torchscript = False + checkpoint = torch.load(original_path, map_location="cpu", weights_only=False) + if not isinstance(checkpoint, dict) or "model" not in checkpoint: + raise ValueError( + f"Cannot load checkpoint at '{original_path}'; " + "expected a TorchScript archive or a dict with a 'model' key" + ) + net = checkpoint["model"] + existing_meta = checkpoint.get("metadata", {}) + + merged = {**existing_meta, **metadata} + + tmp = tempfile.NamedTemporaryFile(suffix=".pt", delete=False) + if is_torchscript: + extra_out = {"metadata.json": json.dumps(merged)} + torch.jit.save(net, tmp.name, _extra_files=extra_out) + else: + torch.save({"model": net, "metadata": merged}, tmp.name) + + return tmp.name diff --git a/source/isaaclab_newton/isaaclab_newton/assets/articulation/articulation.py b/source/isaaclab_newton/isaaclab_newton/assets/articulation/articulation.py index 515176352490..485118128dd9 100644 --- a/source/isaaclab_newton/isaaclab_newton/assets/articulation/articulation.py +++ b/source/isaaclab_newton/isaaclab_newton/assets/articulation/articulation.py @@ -25,6 +25,8 @@ from isaaclab.actuators import ActuatorBase, ActuatorBaseCfg, ImplicitActuator from isaaclab.assets.articulation.base_articulation import BaseArticulation + +from isaaclab_newton.actuators.newton_actuator_utils import NewtonActuatorAdapter from isaaclab.physics import PhysicsEvent from isaaclab.sim.utils.queries import find_first_matching_prim, get_all_matching_child_prims from isaaclab.utils.string import resolve_matching_names, resolve_matching_names_values @@ -39,12 +41,20 @@ from .articulation_data import ArticulationData if TYPE_CHECKING: + from newton.actuators import Actuator + from isaaclab.assets.articulation.articulation_cfg import ArticulationCfg # import logger logger = logging.getLogger(__name__) +@wp.kernel(enable_backward=False) +def _build_env_mask_kernel(mask: wp.array(dtype=wp.bool), indices: wp.array(dtype=wp.int32)): + i = wp.tid() + mask[indices[i]] = True + + class Articulation(BaseArticulation): """An articulation asset class. @@ -275,14 +285,23 @@ def write_data_to_sim(self): ) self._instantaneous_wrench_composer.reset() - # apply actuator models - self._apply_actuator_model() - # write actions into simulation via Newton bindings - self.data._sim_bind_joint_effort.assign(self._joint_effort_target_sim) - # position and velocity targets only for implicit actuators - if self._has_implicit_actuators: - self.data._sim_bind_joint_position_target.assign(self._joint_pos_target_sim) - self.data._sim_bind_joint_velocity_target.assign(self._joint_vel_target_sim) + if self._has_newton_actuators: + # Raw targets go directly to Newton's control object. Newton PD + # consumes them for Newton-managed joints; the solver's built-in + # PD uses them for any implicit joints (whose stiffness/damping + # are non-zero in sim). + # Feedforward effort goes to control.joint_act (not joint_f, + # which is the actuator output buffer that Newton overwrites). + self.data._sim_bind_joint_position_target.assign(self._data._joint_pos_target) + self.data._sim_bind_joint_velocity_target.assign(self._data._joint_vel_target) + self.data._sim_bind_joint_act.assign(self._data._joint_effort_target) + else: + # Standard Lab actuator path + self._apply_actuator_model() + self.data._sim_bind_joint_effort.assign(self._joint_effort_target_sim) + if self._has_implicit_actuators: + self.data._sim_bind_joint_position_target.assign(self._joint_pos_target_sim) + self.data._sim_bind_joint_velocity_target.assign(self._joint_vel_target_sim) def update(self, dt: float): """Updates the simulation data. @@ -1984,6 +2003,108 @@ def write_joint_friction_coefficient_to_sim_mask( # tell the physics engine that some of the joint properties have been updated SimulationManager.add_model_change(SolverNotifyFlags.JOINT_DOF_PROPERTIES) + """ + Operations - Newton Actuator Parameter Writers. + """ + + def write_actuator_stiffness_to_sim( + self, + adapter: NewtonActuatorAdapter, + *, + stiffness: torch.Tensor | wp.array | float, + env_ids: Sequence[int] | torch.Tensor | wp.array | None = None, + ) -> None: + """Write actuator stiffness (kp) into the Newton controller arrays. + + Analogous to :meth:`write_joint_stiffness_to_sim_index` but targets the + Newton actuator controller ``kp`` instead of the solver-level + ``joint_target_ke``. + + Args: + adapter: The :class:`NewtonActuatorAdapter` whose gains to update. + stiffness: Stiffness values [N/m or N*m/rad, depending on joint + type]. Shape is ``(len(env_ids), num_joints)``. + env_ids: Environment indices. If ``None``, all environments. + """ + env_ids = self._resolve_env_ids(env_ids) + stiffness_wp = wp.from_torch(adapter.stiffness, dtype=wp.float32) + if isinstance(stiffness, float): + wp.launch( + articulation_kernels.float_data_to_buffer_with_indices, + dim=(env_ids.shape[0], self._ALL_JOINT_INDICES.shape[0]), + inputs=[stiffness, env_ids, self._ALL_JOINT_INDICES], + outputs=[stiffness_wp], + device=self.device, + ) + else: + wp.launch( + shared_kernels.write_2d_data_to_buffer_with_indices, + dim=(env_ids.shape[0], self._ALL_JOINT_INDICES.shape[0]), + inputs=[stiffness, env_ids, self._ALL_JOINT_INDICES], + outputs=[stiffness_wp], + device=self.device, + ) + env_mask = self._env_ids_to_mask(env_ids) + for newton_act in adapter.actuators: + if hasattr(newton_act.controller, "kp"): + self._root_view.set_actuator_parameter( + actuator=newton_act, component=newton_act.controller, name="kp", + values=stiffness_wp, mask=env_mask, + ) + + def write_actuator_damping_to_sim( + self, + adapter: NewtonActuatorAdapter, + *, + damping: torch.Tensor | wp.array | float, + env_ids: Sequence[int] | torch.Tensor | wp.array | None = None, + ) -> None: + """Write actuator damping (kd) into the Newton controller arrays. + + Analogous to :meth:`write_joint_damping_to_sim_index` but targets the + Newton actuator controller ``kd`` instead of the solver-level + ``joint_target_kd``. + + Args: + adapter: The :class:`NewtonActuatorAdapter` whose gains to update. + damping: Damping values [N*s/m or N*m*s/rad, depending on joint + type]. Shape is ``(len(env_ids), num_joints)``. + env_ids: Environment indices. If ``None``, all environments. + """ + env_ids = self._resolve_env_ids(env_ids) + damping_wp = wp.from_torch(adapter.damping, dtype=wp.float32) + if isinstance(damping, float): + wp.launch( + articulation_kernels.float_data_to_buffer_with_indices, + dim=(env_ids.shape[0], self._ALL_JOINT_INDICES.shape[0]), + inputs=[damping, env_ids, self._ALL_JOINT_INDICES], + outputs=[damping_wp], + device=self.device, + ) + else: + wp.launch( + shared_kernels.write_2d_data_to_buffer_with_indices, + dim=(env_ids.shape[0], self._ALL_JOINT_INDICES.shape[0]), + inputs=[damping, env_ids, self._ALL_JOINT_INDICES], + outputs=[damping_wp], + device=self.device, + ) + env_mask = self._env_ids_to_mask(env_ids) + for newton_act in adapter.actuators: + if hasattr(newton_act.controller, "kd"): + self._root_view.set_actuator_parameter( + actuator=newton_act, component=newton_act.controller, name="kd", + values=damping_wp, mask=env_mask, + ) + + def _env_ids_to_mask(self, env_ids: wp.array) -> wp.array: + """Convert warp env_ids to a boolean Warp mask.""" + if env_ids is self._ALL_INDICES: + return self._ALL_ENV_MASK + mask = wp.zeros(self.num_instances, dtype=wp.bool, device=self.device) + wp.launch(_build_env_mask_kernel, dim=env_ids.shape[0], inputs=[mask, env_ids], device=self.device) + return mask + """ Operations - Setters. """ @@ -3313,123 +3434,50 @@ def _process_actuators_cfg(self): # flag for implicit actuators # if this is false, we by-pass certain checks when doing actuator-related operations self._has_implicit_actuators = False + self._has_newton_actuators = False - # iterate over all actuator configurations - for actuator_name, actuator_cfg in self.cfg.actuators.items(): - # type annotation for type checkers - actuator_cfg: ActuatorBaseCfg - # create actuator group - joint_ids, joint_names = self.find_joints(actuator_cfg.joint_names_expr) - # check if any joints are found - if len(joint_names) == 0: - raise ValueError( - f"No joints found for actuator group: {actuator_name} with joint name expression:" - f" {actuator_cfg.joint_names_expr}." - ) - # resolve joint indices - # we pass a slice if all joints are selected to avoid indexing overhead - if len(joint_names) == self.num_joints: - joint_ids = slice(None) - else: - joint_ids = torch.tensor(joint_ids, device=self.device, dtype=torch.int32) - # create actuator collection - # note: for efficiency avoid indexing when over all indices - actuator: ActuatorBase = actuator_cfg.class_type( - cfg=actuator_cfg, - joint_names=joint_names, - joint_ids=joint_ids, - num_envs=self.num_instances, - device=self.device, - stiffness=wp.to_torch(self._data.joint_stiffness)[:, joint_ids], - damping=wp.to_torch(self._data.joint_damping)[:, joint_ids], - armature=wp.to_torch(self._data.joint_armature)[:, joint_ids], - friction=wp.to_torch(self._data.joint_friction_coeff)[:, joint_ids], - effort_limit=wp.to_torch(self._data.joint_effort_limits)[:, joint_ids].clone(), - velocity_limit=wp.to_torch(self._data.joint_vel_limits)[:, joint_ids], - ) - # store actuator group - self.actuators[actuator_name] = actuator - # set the passed gains and limits into the simulation - if isinstance(actuator, ImplicitActuator): - self._has_implicit_actuators = True - # the gains and limits are set into the simulation since actuator model is implicit - self.write_joint_stiffness_to_sim_index(stiffness=actuator.stiffness, joint_ids=actuator.joint_indices) - self.write_joint_damping_to_sim_index(damping=actuator.damping, joint_ids=actuator.joint_indices) + _use_newton_actuators = getattr(self._sim_cfg, "use_newton_actuators", False) + + if _use_newton_actuators: + from newton import Model as NewtonModel # noqa: PLC0415 + + model = SimulationManager.get_model() + + dof_layout = self._root_view.frequency_layouts[NewtonModel.AttributeFrequency.JOINT_DOF] + if dof_layout.slice is not None: + dof_offset = dof_layout.slice.start + elif dof_layout.indices is not None: + dof_offset = int(dof_layout.indices.numpy()[0]) else: - # the gains and limits are processed by the actuator model - # we set gains to zero, and torque limit to a high value in simulation to avoid any interference - self.write_joint_stiffness_to_sim_index(stiffness=0.0, joint_ids=actuator.joint_indices) - self.write_joint_damping_to_sim_index(damping=0.0, joint_ids=actuator.joint_indices) - - # Set common properties into the simulation - self.write_joint_effort_limit_to_sim_index( - limits=actuator.effort_limit_sim, joint_ids=actuator.joint_indices - ) - self.write_joint_velocity_limit_to_sim_index( - limits=actuator.velocity_limit_sim, joint_ids=actuator.joint_indices - ) - self.write_joint_armature_to_sim_index(armature=actuator.armature, joint_ids=actuator.joint_indices) - self.write_joint_friction_coefficient_to_sim_index( - joint_friction_coeff=actuator.friction, joint_ids=actuator.joint_indices - ) + dof_offset = 0 - # Store the configured values from the actuator model - # note: this is the value configured in the actuator model (for implicit and explicit actuators) - joint_ids = actuator.joint_indices - if joint_ids == slice(None): - joint_ids = self._ALL_JOINT_INDICES - wp.launch( - shared_kernels.write_2d_data_to_buffer_with_indices, - dim=(self.num_instances, joint_ids.shape[0]), - inputs=[ - actuator.stiffness, - self._ALL_INDICES, - joint_ids, - ], - outputs=[ - self.data._sim_bind_joint_stiffness_sim, - ], - device=self.device, - ) - wp.launch( - shared_kernels.write_2d_data_to_buffer_with_indices, - dim=(self.num_instances, joint_ids.shape[0]), - inputs=[ - actuator.damping, - self._ALL_INDICES, - joint_ids, - ], - outputs=[ - self.data._sim_bind_joint_damping_sim, - ], - device=self.device, - ) - wp.launch( - shared_kernels.write_2d_data_to_buffer_with_indices, - dim=(self.num_instances, joint_ids.shape[0]), - inputs=[ - actuator.armature, - self._ALL_INDICES, - joint_ids, - ], - outputs=[ - self.data._sim_bind_joint_armature, - ], - device=self.device, - ) - wp.launch( - shared_kernels.write_2d_data_to_buffer_with_indices, - dim=(self.num_instances, joint_ids.shape[0]), - inputs=[ - actuator.friction, - self._ALL_INDICES, - joint_ids, - ], - outputs=[ - self.data._sim_bind_joint_friction_coeff, - ], - device=self.device, + adapter = NewtonActuatorAdapter( + model.actuators, self.num_instances, self.num_joints, dof_offset, self.device, ) + adapter.finalize() + self.actuators["newton"] = adapter + SimulationManager.register_adapter(adapter) + self.write_joint_stiffness_to_sim_index(stiffness=0.0, joint_ids=adapter.joint_indices) + self.write_joint_damping_to_sim_index(damping=0.0, joint_ids=adapter.joint_indices) + self._has_newton_actuators = True + + # Also create any implicit actuators defined in the configs + for actuator_name, actuator_cfg in self.cfg.actuators.items(): + cls_type = actuator_cfg.class_type + is_implicit = ( + "ImplicitActuator" in cls_type + if isinstance(cls_type, str) + else issubclass(cls_type, ImplicitActuator) + ) + if not is_implicit: + continue + self._create_lab_actuator(actuator_name, actuator_cfg) + + return + + # --- Standard Isaac Lab actuator path --- + for actuator_name, actuator_cfg in self.cfg.actuators.items(): + self._create_lab_actuator(actuator_name, actuator_cfg) # perform some sanity checks to ensure actuators are prepared correctly total_act_joints = sum(actuator.num_joints for actuator in self.actuators.values()) @@ -3442,6 +3490,8 @@ def _process_actuators_cfg(self): if self.cfg.actuator_value_resolution_debug_print: t = PrettyTable(["Group", "Property", "Name", "ID", "USD Value", "ActutatorCfg Value", "Applied"]) for actuator_group, actuator in self.actuators.items(): + if isinstance(actuator, NewtonActuatorAdapter): + continue group_count = 0 for property, resolution_details in actuator.joint_property_resolution_table.items(): for prop_idx, resolution_detail in enumerate(resolution_details): @@ -3452,6 +3502,86 @@ def _process_actuators_cfg(self): group_count += 1 logger.warning(f"\nActuatorCfg-USD Value Discrepancy Resolution (matching values are skipped): \n{t}") + def _create_lab_actuator(self, actuator_name: str, actuator_cfg: ActuatorBaseCfg) -> None: + """Instantiate a single Lab actuator from its config and write properties to sim.""" + joint_ids, joint_names = self.find_joints(actuator_cfg.joint_names_expr) + if len(joint_names) == 0: + raise ValueError( + f"No joints found for actuator group: {actuator_name} with joint name expression:" + f" {actuator_cfg.joint_names_expr}." + ) + if len(joint_names) == self.num_joints: + joint_ids = slice(None) + else: + joint_ids = torch.tensor(joint_ids, device=self.device, dtype=torch.int32) + + actuator: ActuatorBase = actuator_cfg.class_type( + cfg=actuator_cfg, + joint_names=joint_names, + joint_ids=joint_ids, + num_envs=self.num_instances, + device=self.device, + stiffness=wp.to_torch(self._data.joint_stiffness)[:, joint_ids], + damping=wp.to_torch(self._data.joint_damping)[:, joint_ids], + armature=wp.to_torch(self._data.joint_armature)[:, joint_ids], + friction=wp.to_torch(self._data.joint_friction_coeff)[:, joint_ids], + effort_limit=wp.to_torch(self._data.joint_effort_limits)[:, joint_ids].clone(), + velocity_limit=wp.to_torch(self._data.joint_vel_limits)[:, joint_ids], + ) + self.actuators[actuator_name] = actuator + + if isinstance(actuator, ImplicitActuator): + self._has_implicit_actuators = True + self.write_joint_stiffness_to_sim_index(stiffness=actuator.stiffness, joint_ids=actuator.joint_indices) + self.write_joint_damping_to_sim_index(damping=actuator.damping, joint_ids=actuator.joint_indices) + else: + self.write_joint_stiffness_to_sim_index(stiffness=0.0, joint_ids=actuator.joint_indices) + self.write_joint_damping_to_sim_index(damping=0.0, joint_ids=actuator.joint_indices) + + self.write_joint_effort_limit_to_sim_index( + limits=actuator.effort_limit_sim, joint_ids=actuator.joint_indices, + ) + self.write_joint_velocity_limit_to_sim_index( + limits=actuator.velocity_limit_sim, joint_ids=actuator.joint_indices, + ) + self.write_joint_armature_to_sim_index(armature=actuator.armature, joint_ids=actuator.joint_indices) + self.write_joint_friction_coefficient_to_sim_index( + joint_friction_coeff=actuator.friction, joint_ids=actuator.joint_indices, + ) + + # Store the configured values from the actuator model + j_ids = actuator.joint_indices + if j_ids == slice(None): + j_ids = self._ALL_JOINT_INDICES + wp.launch( + shared_kernels.write_2d_data_to_buffer_with_indices, + dim=(self.num_instances, j_ids.shape[0]), + inputs=[actuator.stiffness, self._ALL_INDICES, j_ids], + outputs=[self.data._sim_bind_joint_stiffness_sim], + device=self.device, + ) + wp.launch( + shared_kernels.write_2d_data_to_buffer_with_indices, + dim=(self.num_instances, j_ids.shape[0]), + inputs=[actuator.damping, self._ALL_INDICES, j_ids], + outputs=[self.data._sim_bind_joint_damping_sim], + device=self.device, + ) + wp.launch( + shared_kernels.write_2d_data_to_buffer_with_indices, + dim=(self.num_instances, j_ids.shape[0]), + inputs=[actuator.armature, self._ALL_INDICES, j_ids], + outputs=[self.data._sim_bind_joint_armature], + device=self.device, + ) + wp.launch( + shared_kernels.write_2d_data_to_buffer_with_indices, + dim=(self.num_instances, j_ids.shape[0]), + inputs=[actuator.friction, self._ALL_INDICES, j_ids], + outputs=[self.data._sim_bind_joint_friction_coeff], + device=self.device, + ) + def _process_tendons(self): """Process fixed and spatial tendons.""" # create a list to store the fixed tendon names diff --git a/source/isaaclab_newton/isaaclab_newton/assets/articulation/articulation_data.py b/source/isaaclab_newton/isaaclab_newton/assets/articulation/articulation_data.py index 4dc9507dbe59..f298a274f00a 100644 --- a/source/isaaclab_newton/isaaclab_newton/assets/articulation/articulation_data.py +++ b/source/isaaclab_newton/isaaclab_newton/assets/articulation/articulation_data.py @@ -1286,6 +1286,9 @@ def _create_simulation_bindings(self) -> None: self._sim_bind_joint_effort = self._root_view.get_attribute("joint_f", SimulationManager.get_control())[ :, 0 ] + self._sim_bind_joint_act = self._root_view.get_attribute("joint_act", SimulationManager.get_control())[ + :, 0 + ] self._sim_bind_joint_position_target = self._root_view.get_attribute( "joint_target_pos", SimulationManager.get_control() )[:, 0] @@ -1317,6 +1320,7 @@ def _create_simulation_bindings(self) -> None: self._sim_bind_joint_pos = wp.zeros((self._num_instances, 0), dtype=wp.float32, device=self.device) self._sim_bind_joint_vel = wp.zeros((self._num_instances, 0), dtype=wp.float32, device=self.device) self._sim_bind_joint_effort = wp.zeros((self._num_instances, 0), dtype=wp.float32, device=self.device) + self._sim_bind_joint_act = wp.zeros((self._num_instances, 0), dtype=wp.float32, device=self.device) self._sim_bind_joint_position_target = wp.zeros( (self._num_instances, 0), dtype=wp.float32, device=self.device ) diff --git a/source/isaaclab_newton/isaaclab_newton/physics/newton_manager.py b/source/isaaclab_newton/isaaclab_newton/physics/newton_manager.py index 8a752c8d4c8e..42a50d0f27e5 100644 --- a/source/isaaclab_newton/isaaclab_newton/physics/newton_manager.py +++ b/source/isaaclab_newton/isaaclab_newton/physics/newton_manager.py @@ -42,6 +42,8 @@ if TYPE_CHECKING: from isaaclab.sim.simulation_context import SimulationContext + from isaaclab_newton.actuators import NewtonActuatorAdapter + from .newton_collision_cfg import NewtonCollisionPipelineCfg logger = logging.getLogger(__name__) @@ -110,6 +112,10 @@ class NewtonManager(PhysicsManager): _report_contacts: bool = False _fk_dirty: bool = False + # Newton actuator adapter (owns actuators and double-buffered states) + _adapter: NewtonActuatorAdapter | None = None + _decimation: int = 1 + # CUDA graphing _graph = None _graph_capture_pending: bool = False @@ -303,7 +309,30 @@ def _mark_transforms_dirty(cls) -> None: @classmethod def step(cls) -> None: - """Step the physics simulation.""" + """Step the physics simulation. + + The stepping logic follows one of two paths depending on whether + **all** actuators are CUDA-graph-safe: + + **All-graphable path** (:meth:`_simulate_full`): + + Actuators and solver substeps are captured together in a single + CUDA graph containing the full + ``decimation x (actuators + solver substeps)`` loop. + + **Eager-actuator path** (fallback, some actuators not graph-safe): + + Actuators are stepped eagerly on the CPU timeline (outside the + graph), then a graph containing only the solver substeps is + launched via :meth:`_simulate_physics_only`. + + In both paths the sequence within one physics step is:: + + zero actuated DOFs in control.joint_f + -> actuator.step (computes effort, writes to control.joint_f) + -> solver.step x num_substeps (integrates, reads control.joint_f) + -> sensors.update + """ sim = PhysicsManager._sim if sim is None or not sim.is_playing(): return @@ -315,9 +344,7 @@ def step(cls) -> None: cls._solver.notify_model_changed(change) cls._model_changes = set() - # Lazy CUDA graph capture: deferred from initialize_solver() when RTX was active. - # By the time step() is first called, RTX has fully initialized (all cudaImportExternalMemory - # calls are done) and is idle between render frames — giving us a clean capture window. + # Lazy CUDA graph capture cfg = PhysicsManager._cfg device = PhysicsManager._device if cls._graph_capture_pending and cfg is not None and cfg.use_cuda_graph and "cuda" in device: # type: ignore[union-attr] @@ -328,21 +355,36 @@ def step(cls) -> None: else: logger.warning("Newton deferred CUDA graph capture failed; using eager execution") - # Ensure body_q is up-to-date before collision detection. - # After env resets, joint_q is written but body_q (used by - # broadphase/narrowphase) is stale until FK runs. + # FK update for collision pipeline if cls._fk_dirty and cls._needs_collision_pipeline: eval_fk(cls._model, cls._state_0.joint_q, cls._state_0.joint_qd, cls._state_0, None) cls._fk_dirty = False - # Step simulation (graphed or not; _graph is None when capture is disabled or failed) - if cfg is not None and cfg.use_cuda_graph and cls._graph is not None and "cuda" in device: # type: ignore[union-attr] - wp.capture_launch(cls._graph) - if cls._usdrt_stage is not None: - cls._mark_transforms_dirty() + physics_dt = cls._solver_dt * cls._num_substeps + use_graph = cfg is not None and cfg.use_cuda_graph and cls._graph is not None and "cuda" in device # type: ignore[union-attr] + + if cls._is_all_graphable(): + # --- All actuators are graph-safe: actuators + solver in one graph --- + if use_graph: + wp.capture_launch(cls._graph) + else: + with wp.ScopedDevice(device): + cls._simulate_full() + PhysicsManager._sim_time += physics_dt * cls._decimation else: - with wp.ScopedDevice(device): - cls._simulate() + # --- Some actuators not graph-safe: step them eagerly, graph solver only --- + if cls._adapter is not None: + cls._adapter.step(cls._state_0, cls._control, physics_dt) + + if use_graph: + wp.capture_launch(cls._graph) + else: + with wp.ScopedDevice(device): + cls._simulate_physics_only() + PhysicsManager._sim_time += physics_dt + + if cls._usdrt_stage is not None: + cls._mark_transforms_dirty() # Debug convergence info if cfg is not None and cfg.debug_mode: # type: ignore[union-attr] @@ -350,7 +392,6 @@ def step(cls) -> None: logger.info(f"Solver convergence data: {convergence_data}") if convergence_data["max"] == cls._solver.mjw_model.opt.iterations: logger.warning(f"Solver didn't converge! max_iter={convergence_data['max']}") - PhysicsManager._sim_time += cls._solver_dt * cls._num_substeps @classmethod def close(cls) -> None: @@ -399,6 +440,8 @@ def clear(cls): cls._newton_imu_sensors = [] cls._report_contacts = False cls._fk_dirty = False + cls._adapter = None + cls._decimation = 1 cls._graph = None cls._graph_capture_pending = False cls._newton_stage_path = None @@ -651,6 +694,10 @@ def start_simulation(cls) -> None: cls._control = cls._model.control() eval_fk(cls._model, cls._state_0.joint_q, cls._state_0.joint_qd, cls._state_0, None) + # Actuator registration is deferred to register_adapter(), + # called by _process_actuators_cfg after the adapter is constructed. + cls._adapter = None + logger.info("Dispatching PHYSICS_READY callbacks") cls.dispatch_event(PhysicsEvent.PHYSICS_READY) @@ -870,28 +917,46 @@ def initialize_solver(cls) -> None: logger.warning("cubric bindings init failed; falling back to update_world_xforms()") cls._cubric = None - device = PhysicsManager._device + with Timer(name="newton_cuda_graph", msg="CUDA graph took:"): + cls._capture_or_defer_graph() - use_cuda_graph = cfg.use_cuda_graph and "cuda" in device # type: ignore[union-attr] + @classmethod + def _capture_or_defer_graph(cls) -> None: + """Capture (or schedule deferred capture of) the CUDA graph. + + Called by :meth:`start_simulation` and :meth:`set_decimation` + whenever the graph needs to be (re-)captured. + + * **No USDRT / headless**: captures immediately via + ``wp.ScopedCapture``. + * **RTX active**: defers capture to the first :meth:`step` call + via :meth:`_capture_relaxed_graph`, because RTX background + streams are not yet idle during initialisation. + * **CUDA graphs disabled**: clears the graph reference. + """ + cfg = PhysicsManager._cfg + device = PhysicsManager._device + if cfg is None or device is None: + return - with Timer(name="newton_cuda_graph", msg="CUDA graph took:"): - if use_cuda_graph: - if cls._usdrt_stage is None: - # No RTX active — use standard Warp capture (cudaStreamCaptureModeGlobal). - with wp.ScopedCapture() as capture: - cls._simulate() - cls._graph = capture.graph - logger.info("Newton CUDA graph captured (standard Warp mode)") - else: - # RTX is active during initialization — cudaImportExternalMemory and other - # non-capturable RTX ops run on background CUDA streams right now. - # Defer capture to the first step() call, after RTX is fully initialized - # and idle between render frames (clean capture window). - cls._graph = None - cls._graph_capture_pending = True - logger.info("Newton CUDA graph capture deferred until first step() (RTX active)") + use_cuda_graph = cfg.use_cuda_graph and "cuda" in device + if use_cuda_graph: + if cls._usdrt_stage is None: + simulate = cls._simulate_full if cls._is_all_graphable() else cls._simulate_physics_only + with wp.ScopedCapture() as capture: + simulate() + cls._graph = capture.graph + logger.info("Newton CUDA graph captured (standard Warp mode)") else: + # RTX is active during initialization — cudaImportExternalMemory and other + # non-capturable RTX ops run on background CUDA streams right now. + # Defer capture to the first step() call, after RTX is fully initialized + # and idle between render frames (clean capture window). cls._graph = None + cls._graph_capture_pending = True + logger.info("Newton CUDA graph capture deferred until first step() (RTX active)") + else: + cls._graph = None @classmethod def _capture_relaxed_graph(cls, device: str): @@ -922,7 +987,7 @@ def _capture_relaxed_graph(cls, device: str): this registers the capture in Warp's ``device.captures`` *without* calling ``cudaStreamBeginCapture`` (already done) and *without* changing device-wide memory pool attributes (avoids error 900 in RTX's ``cudaMallocAsync``). - - Run ``_simulate_physics_only()`` inside ``ScopedStream(fresh_stream)``: + - Run the simulate function inside ``ScopedStream(fresh_stream)``: kernels dispatch to ``fresh_stream`` and are captured; ``wp.capture_while`` finds the active capture and inserts a conditional graph node instead of synchronising. - Call ``wp.capture_end(stream=fresh_stream)`` to finalise the Warp-level capture. @@ -940,8 +1005,9 @@ def _capture_relaxed_graph(cls, device: str): # Warmup: pre-allocate all MuJoCo-Warp scratch buffers so the capture window has # no new cudaMalloc calls (which are forbidden inside graph capture). + simulate = cls._simulate_full if cls._is_all_graphable() else cls._simulate_physics_only with wp.ScopedDevice(device): - cls._simulate_physics_only() + simulate() wp.synchronize_stream(wp.get_stream(device)) # Create a non-blocking stream (cudaStreamNonBlocking = 0x01). @@ -974,7 +1040,7 @@ def _capture_relaxed_graph(cls, device: str): err_during_capture = None with wp.ScopedStream(fresh_stream, sync_enter=False): try: - cls._simulate_physics_only() + simulate() except Exception as exc: err_during_capture = exc @@ -1014,68 +1080,83 @@ def _capture_relaxed_graph(cls, device: str): graph.graph_exec = None return graph - @classmethod - def _simulate_physics_only(cls) -> None: - """Run one physics step without Fabric/USD sync — safe for CUDA graph capture. - - Used by :meth:`_capture_relaxed_graph` to capture only the pure physics kernels. - ``sync_transforms_to_usd`` is excluded because it calls ``wp.synchronize_device`` - (forbidden inside graph capture) and ``wp.fabricarray`` (device-wide allocation). - The caller (``step()``) is responsible for calling ``sync_transforms_to_usd()`` - eagerly after ``wp.capture_launch``. - """ - if cls._needs_collision_pipeline: - cls._collision_pipeline.collide(cls._state_0, cls._contacts) - contacts = cls._contacts - else: - contacts = None - - def step_fn(state_0, state_1): - cls._solver.step(state_0, state_1, cls._control, contacts, cls._solver_dt) + # ------------------------------------------------------------------ + # Building blocks — used by _simulate_full / _simulate_physics_only + # ------------------------------------------------------------------ + @classmethod + def _run_solver_substeps(cls, contacts) -> None: + """Run ``num_substeps`` solver iterations, handling double-buffered state swap.""" if cls._use_single_state: - for i in range(cls._num_substeps): - step_fn(cls._state_0, cls._state_0) + for _ in range(cls._num_substeps): + cls._solver.step(cls._state_0, cls._state_0, cls._control, contacts, cls._solver_dt) cls._state_0.clear_forces() else: cfg = PhysicsManager._cfg - need_copy_on_last_substep = (cfg is not None and cfg.use_cuda_graph) and cls._num_substeps % 2 == 1 # type: ignore[union-attr] + need_copy_on_last = (cfg is not None and cfg.use_cuda_graph) and cls._num_substeps % 2 == 1 # type: ignore[union-attr] for i in range(cls._num_substeps): - step_fn(cls._state_0, cls._state_1) - if need_copy_on_last_substep and i == cls._num_substeps - 1: + cls._solver.step(cls._state_0, cls._state_1, cls._control, contacts, cls._solver_dt) + if need_copy_on_last and i == cls._num_substeps - 1: cls._state_0.assign(cls._state_1) else: cls._state_0, cls._state_1 = cls._state_1, cls._state_0 cls._state_0.clear_forces() - # Update frame transform sensors + @classmethod + def _update_sensors(cls, contacts) -> None: + """Push latest state to all registered Newton sensors.""" if cls._newton_frame_transform_sensors: for sensor in cls._newton_frame_transform_sensors: sensor.update(cls._state_0) - - # Update IMU sensors if cls._newton_imu_sensors: for sensor in cls._newton_imu_sensors: sensor.update(cls._state_0) - - # Populate contacts for contact sensors if cls._report_contacts: eval_contacts = contacts if contacts is not None else cls._contacts cls._solver.update_contacts(eval_contacts, cls._state_0) for sensor in cls._newton_contact_sensors.values(): sensor.update(cls._state_0, eval_contacts) + # ------------------------------------------------------------------ + # Composite stepping routines + # ------------------------------------------------------------------ + @classmethod - def _simulate(cls) -> None: - """Run one simulation step with substeps and USD sync. + def _simulate_full(cls) -> None: + """Run ``decimation x (actuators + solver substeps)``, then sensors. - Delegates physics work to :meth:`_simulate_physics_only` and then - marks transforms dirty for the next render-cadence sync. + Works for any decimation count (including 1). All actuators must be + graph-safe so the entire loop can be captured as a single CUDA graph. """ - cls._simulate_physics_only() + physics_dt = cls._solver_dt * cls._num_substeps + contacts = cls._contacts if cls._needs_collision_pipeline else None - if cls._usdrt_stage is not None: - cls._mark_transforms_dirty() + for _ in range(cls._decimation): + if cls._needs_collision_pipeline: + cls._collision_pipeline.collide(cls._state_0, cls._contacts) + + if cls._adapter is not None: + cls._adapter.step(cls._state_0, cls._control, physics_dt) + + cls._run_solver_substeps(contacts) + + cls._update_sensors(contacts) + + @classmethod + def _simulate_physics_only(cls) -> None: + """Collision + solver substeps + sensors (no actuators, no USD sync). + + Used when actuators are stepped eagerly outside the graph, or when + there are no actuators at all. + """ + if cls._needs_collision_pipeline: + cls._collision_pipeline.collide(cls._state_0, cls._contacts) + contacts = cls._contacts + else: + contacts = None + + cls._run_solver_substeps(contacts) + cls._update_sensors(contacts) @classmethod def get_solver_convergence_steps(cls) -> dict[str, float | int]: @@ -1119,6 +1200,51 @@ def get_solver_dt(cls) -> float: """Get the solver substep timestep.""" return cls._solver_dt + @classmethod + def _is_all_graphable(cls) -> bool: + """``True`` when an adapter is registered and all its actuators are CUDA-graph-safe.""" + return cls._adapter is not None and cls._adapter.is_all_graphable + + @classmethod + def register_adapter(cls, adapter: NewtonActuatorAdapter) -> None: + """Register the actuator adapter that owns states and stepping. + + Called by the articulation's ``_process_actuators_cfg`` after + the :class:`NewtonActuatorAdapter` is constructed and finalised. + + Args: + adapter: The actuator adapter instance. + """ + cls._adapter = adapter + + @classmethod + def set_decimation(cls, decimation: int) -> None: + """Set the decimation count and re-capture the CUDA graph. + + When all actuators are graphable the entire decimation loop + (actuators + solver substeps, repeated *decimation* times) + is captured as a single CUDA graph. + + If a CUDA graph was previously captured, it is automatically + re-captured with the new decimation count using the same + strategy as :meth:`start_simulation`: standard + ``wp.ScopedCapture`` when no USDRT stage is active, or + deferred relaxed capture when RTX is running. + """ + cls._decimation = max(1, decimation) + if cls._is_all_graphable(): + cls._capture_or_defer_graph() + + @classmethod + def handles_decimation(cls) -> bool: + """``True`` when :meth:`step` executes the full decimation loop internally. + + This is the case when all Newton actuators are CUDA-graph-safe. + The full decimation loop (including the trivial ``decimation=1`` case) + is folded into a single :meth:`step` call. + """ + return cls._is_all_graphable() + @classmethod def add_contact_sensor( cls, diff --git a/source/isaaclab_newton/setup.py b/source/isaaclab_newton/setup.py index 2e0b87f17543..14e533b69d13 100644 --- a/source/isaaclab_newton/setup.py +++ b/source/isaaclab_newton/setup.py @@ -41,7 +41,7 @@ def run(self): "mujoco==3.6.0", "mujoco-warp==3.6.0", "PyOpenGL-accelerate==3.1.10", - "newton @ git+https://github.com/newton-physics/newton.git@a27277ed49d6f307b8a1e4c394be7e1d14965a62", + "newton @ git+https://github.com/newton-physics/newton.git@b859659c7f3df8b688a5989a8f1719b8dc15aa86", ], } @@ -62,6 +62,7 @@ def run(self): extras_require=EXTRAS_REQUIRE, packages=[ "isaaclab_newton", + "isaaclab_newton.actuators", "isaaclab_newton.assets", "isaaclab_newton.assets.articulation", "isaaclab_newton.assets.rigid_object", diff --git a/source/isaaclab_newton/test/assets/test_pd_equivalence.py b/source/isaaclab_newton/test/assets/test_pd_equivalence.py new file mode 100644 index 000000000000..f9e211891c7e --- /dev/null +++ b/source/isaaclab_newton/test/assets/test_pd_equivalence.py @@ -0,0 +1,818 @@ +# Copyright (c) 2022-2026, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +"""PD actuator equivalence tests on ANYmal-C (floating-base quadruped). + +Compares IsaacLab-native actuators against Newton-native actuators (created +from the same Lab configs via USD authoring) on the Newton physics backend. +Both paths must produce identical joint trajectories within tolerance. + +Using ANYmal-C — a 12-DOF quadruped on a floating base — exercises the +coordinate-vs-DOF index separation that is critical when free joints shift +the mapping between ``joint_q`` (coordinate layout) and ``joint_qd`` +(DOF layout). + +Each test class overrides ANYmal's default actuators with a specific Lab +config (IdealPD, DCMotor, or mixed) and verifies Lab vs Newton equivalence. +""" + +from isaaclab.app import AppLauncher + +simulation_app = AppLauncher(headless=True).app + +import json +import os +import tempfile +import unittest + +import torch +import warp as wp +from isaaclab_assets import ANYMAL_C_CFG +from isaaclab_newton.assets import Articulation +from isaaclab_newton.physics import MJWarpSolverCfg, NewtonCfg +from isaaclab_newton.physics import NewtonManager as SimulationManager + +import isaaclab.sim as sim_utils +from isaaclab.actuators import DCMotorCfg, IdealPDActuatorCfg, ImplicitActuatorCfg +from isaaclab.sim import SimulationCfg, build_simulation_context + +# --------------------------------------------------------------------------- +# Constants +# --------------------------------------------------------------------------- + +NUM_ENVS = 2 +NUM_STEPS = 10 +DT = 1.0 / 120.0 +TARGET_OFFSET = 0.1 # [rad] added to initial joint positions + +NEWTON_CFG = NewtonCfg( + solver_cfg=MJWarpSolverCfg( + njmax=500, + nconmax=500, + ls_iterations=20, + cone="pyramidal", + impratio=1, + ls_parallel=False, + integrator="implicitfast", + ), + num_substeps=1, + debug_mode=False, + use_cuda_graph=False, +) + +# --------------------------------------------------------------------------- +# Actuator configurations under test +# --------------------------------------------------------------------------- + +IDEAL_PD_ACTUATORS = { + "legs": IdealPDActuatorCfg( + joint_names_expr=[".*HAA", ".*HFE", ".*KFE"], + stiffness=40.0, + damping=5.0, + effort_limit=80.0, + ), +} + +DC_MOTOR_ACTUATORS = { + "legs": DCMotorCfg( + joint_names_expr=[".*HAA", ".*HFE", ".*KFE"], + saturation_effort=120.0, + effort_limit=80.0, + velocity_limit=7.5, + stiffness={".*": 40.0}, + damping={".*": 5.0}, + ), +} + +MIXED_ACTUATORS = { + "hips": IdealPDActuatorCfg( + joint_names_expr=[".*HAA"], + stiffness=40.0, + damping=5.0, + effort_limit=80.0, + ), + "knees": DCMotorCfg( + joint_names_expr=[".*HFE", ".*KFE"], + saturation_effort=120.0, + effort_limit=80.0, + velocity_limit=7.5, + stiffness={".*": 40.0}, + damping={".*": 5.0}, + ), +} + +# --------------------------------------------------------------------------- +# Simulation runner +# --------------------------------------------------------------------------- + + +def _run_simulation( + actuators: dict, + use_newton_actuators: bool, + *, + dt: float = DT, + newton_cfg: NewtonCfg = NEWTON_CFG, + num_steps: int = NUM_STEPS, + decimation: int = 1, +) -> dict: + """Run ANYmal-C and return recorded joint trajectories. + + Args: + actuators: Actuator config dict overriding ANYmal's defaults. + use_newton_actuators: Use Newton-native actuators when ``True``. + dt: Physics timestep [s]. + newton_cfg: Newton physics configuration. + num_steps: Number of policy-level steps. + decimation: Actuator steps per policy step (Newton decimation loop). + + Returns: + Dict with ``joint_pos`` and ``joint_vel``, each a list of + ``(NUM_ENVS, num_joints)`` tensors. + """ + sim_cfg = SimulationCfg( + dt=dt, + physics=newton_cfg, + use_newton_actuators=use_newton_actuators, + ) + + with build_simulation_context( + device="cuda:0", + gravity_enabled=True, + add_ground_plane=True, + sim_cfg=sim_cfg, + ) as sim: + sim._app_control_on_stop_handle = None + + for i in range(NUM_ENVS): + sim_utils.create_prim( + f"/World/Env_{i}", "Xform", translation=(i * 3.0, 0, 0) + ) + + art_cfg = ANYMAL_C_CFG.replace( + actuators=actuators, + prim_path="/World/Env_.*/Robot", + ) + articulation = Articulation(art_cfg) + sim.reset() + assert articulation.is_initialized + + if use_newton_actuators and decimation > 1: + SimulationManager.set_decimation(decimation) + + handles_dec = ( + use_newton_actuators + and decimation > 1 + and SimulationManager._is_all_graphable() + and SimulationManager._decimation > 1 + ) + + init_pos = wp.to_torch(articulation.data.joint_pos).clone() + target_pos = init_pos + TARGET_OFFSET + target_vel = torch.zeros_like(init_pos) + + articulation.set_joint_position_target_index(target=target_pos) + articulation.set_joint_velocity_target_index(target=target_vel) + + recorded_pos, recorded_vel = [], [] + for _ in range(num_steps): + if handles_dec: + articulation.write_data_to_sim() + sim.step() + articulation.update(dt * decimation) + else: + for _ in range(decimation): + articulation.write_data_to_sim() + sim.step() + articulation.update(dt) + + recorded_pos.append(wp.to_torch(articulation.data.joint_pos).clone()) + recorded_vel.append(wp.to_torch(articulation.data.joint_vel).clone()) + + return {"joint_pos": recorded_pos, "joint_vel": recorded_vel} + + +# --------------------------------------------------------------------------- +# Base test class +# --------------------------------------------------------------------------- + + +class _EquivalenceTestBase(unittest.TestCase): + """Base for Lab-vs-Newton equivalence tests. + + Subclasses set ``actuators`` to the config under test. ``setUpClass`` + runs the simulation with both ``use_newton_actuators=False`` (Lab path) + and ``True`` (Newton path) and stores the results. + """ + + __test__ = False + actuators: dict = {} + pos_atol: float = 2e-3 + pos_rtol: float = 1e-3 + vel_atol: float = 0.1 + vel_rtol: float = 1e-2 + + @classmethod + def setUpClass(cls): + cls.lab_result = _run_simulation(cls.actuators, use_newton_actuators=False) + cls.newton_result = _run_simulation(cls.actuators, use_newton_actuators=True) + + def test_joint_positions_match(self): + for step_i, (lab, newton) in enumerate( + zip(self.lab_result["joint_pos"], self.newton_result["joint_pos"]) + ): + torch.testing.assert_close( + lab, + newton, + atol=self.pos_atol, + rtol=self.pos_rtol, + msg=f"Joint positions diverged at step {step_i}", + ) + + def test_joint_velocities_match(self): + for step_i, (lab, newton) in enumerate( + zip(self.lab_result["joint_vel"], self.newton_result["joint_vel"]) + ): + torch.testing.assert_close( + lab, + newton, + atol=self.vel_atol, + rtol=self.vel_rtol, + msg=f"Joint velocities diverged at step {step_i}", + ) + + def test_trajectories_not_trivial(self): + first = self.lab_result["joint_pos"][0] + last = self.lab_result["joint_pos"][-1] + diff = (last - first).abs().max().item() + self.assertGreater(diff, 0.01, "Joints did not move — test is trivial") + + +# --------------------------------------------------------------------------- +# Equivalence tests with different actuator types +# --------------------------------------------------------------------------- + + +class TestIdealPDEquivalence(_EquivalenceTestBase): + """IdealPDActuator on all 12 joints: Lab vs Newton.""" + + __test__ = True + actuators = IDEAL_PD_ACTUATORS + + +class TestDCMotorEquivalence(_EquivalenceTestBase): + """DCMotor actuator on all 12 joints: Lab vs Newton.""" + + __test__ = True + actuators = DC_MOTOR_ACTUATORS + + +class TestMixedActuatorEquivalence(_EquivalenceTestBase): + """Mixed actuators (IdealPD on HAA, DCMotor on HFE/KFE): Lab vs Newton.""" + + __test__ = True + actuators = MIXED_ACTUATORS + + +MIXED_WITH_IMPLICIT_ACTUATORS = { + "hips": ImplicitActuatorCfg( + joint_names_expr=[".*HAA"], + stiffness=40.0, + damping=5.0, + ), + "thighs": IdealPDActuatorCfg( + joint_names_expr=[".*HFE"], + stiffness=40.0, + damping=5.0, + effort_limit=80.0, + ), + "knees": DCMotorCfg( + joint_names_expr=[".*KFE"], + saturation_effort=120.0, + effort_limit=80.0, + velocity_limit=7.5, + stiffness=40.0, + damping=5.0, + ), +} + + +class TestMixedWithImplicitEquivalence(_EquivalenceTestBase): + """Implicit HAA + IdealPD HFE + DCMotor KFE: Lab vs Newton. + + Verifies that implicit actuators (handled by the physics engine's + built-in joint drives) coexist correctly with explicit Newton actuators. + """ + + __test__ = True + actuators = MIXED_WITH_IMPLICIT_ACTUATORS + + +# --------------------------------------------------------------------------- +# Decimation test: CUDA graph capture with actuator decimation loop +# --------------------------------------------------------------------------- + +DT_DEC = 1.0 / 100.0 +DECIMATION = 2 +NUM_POLICY_STEPS_DEC = 5 + +NEWTON_CFG_DEC = NewtonCfg( + solver_cfg=MJWarpSolverCfg( + njmax=500, + nconmax=500, + ls_iterations=20, + cone="pyramidal", + impratio=1, + ls_parallel=False, + integrator="implicitfast", + ), + num_substeps=2, + debug_mode=False, + use_cuda_graph=True, +) + + +class TestDecimation(unittest.TestCase): + """Lab vs Newton with decimation=2 and CUDA graph capture. + + Policy runs at 50 Hz, actuators at 100 Hz, physics at 200 Hz. + The Newton path captures the full decimation loop as a CUDA graph; + the Lab path runs an explicit per-substep loop. + """ + + @classmethod + def setUpClass(cls): + cls.lab_result = _run_simulation( + DC_MOTOR_ACTUATORS, + use_newton_actuators=False, + dt=DT_DEC, + newton_cfg=NEWTON_CFG_DEC, + num_steps=NUM_POLICY_STEPS_DEC, + decimation=DECIMATION, + ) + cls.newton_result = _run_simulation( + DC_MOTOR_ACTUATORS, + use_newton_actuators=True, + dt=DT_DEC, + newton_cfg=NEWTON_CFG_DEC, + num_steps=NUM_POLICY_STEPS_DEC, + decimation=DECIMATION, + ) + + def test_joint_positions_match(self): + for step_i, (lab, newton) in enumerate( + zip(self.lab_result["joint_pos"], self.newton_result["joint_pos"]) + ): + torch.testing.assert_close( + lab, + newton, + atol=2e-3, + rtol=1e-3, + msg=f"Positions diverged at policy step {step_i}", + ) + + def test_joint_velocities_match(self): + for step_i, (lab, newton) in enumerate( + zip(self.lab_result["joint_vel"], self.newton_result["joint_vel"]) + ): + torch.testing.assert_close( + lab, + newton, + atol=0.1, + rtol=1e-2, + msg=f"Velocities diverged at policy step {step_i}", + ) + + def test_trajectories_not_trivial(self): + first = self.lab_result["joint_pos"][0] + last = self.lab_result["joint_pos"][-1] + diff = (last - first).abs().max().item() + self.assertGreater(diff, 0.01, "Joints did not move — test is trivial") + + +# --------------------------------------------------------------------------- +# RemotizedPD actuator: PD + delay + position-based clamping lookup table +# --------------------------------------------------------------------------- + +SPOT_KNEE_LOOKUP = [ + [-2.792900, -24.776718, 37.165077], + [-2.767442, -26.290108, 39.435162], + [-2.741984, -27.793369, 41.690054], + [-2.716526, -29.285997, 43.928996], + [-2.691068, -30.767536, 46.151304], + [-2.665610, -32.237423, 48.356134], + [-2.640152, -33.695168, 50.542751], + [-2.614694, -35.140221, 52.710331], + [-2.589236, -36.572052, 54.858078], + [-2.563778, -37.990086, 56.985128], + [-2.538320, -39.393730, 59.090595], + [-2.512862, -40.782406, 61.173609], + [-2.487404, -42.155487, 63.233231], + [-2.461946, -43.512371, 65.268557], + [-2.436488, -44.852371, 67.278557], + [-2.411030, -46.174873, 69.262310], + [-2.385572, -47.479156, 71.218735], + [-2.360114, -48.764549, 73.146824], + [-2.334656, -50.030334, 75.045502], + [-2.309198, -51.275761, 76.913641], + [-2.283740, -52.500103, 78.750154], + [-2.258282, -53.702587, 80.553881], + [-2.232824, -54.882442, 82.323664], + [-2.207366, -56.038860, 84.058290], + [-2.181908, -57.171028, 85.756542], + [-2.156450, -58.278133, 87.417200], + [-2.130992, -59.359314, 89.038971], + [-2.105534, -60.413738, 90.620607], + [-2.080076, -61.440529, 92.160793], + [-2.054618, -62.438812, 93.658218], + [-2.029160, -63.407692, 95.111538], + [-2.003702, -64.346268, 96.519402], + [-1.978244, -65.253670, 97.880505], + [-1.952786, -66.128944, 99.193417], + [-1.927328, -66.971176, 100.456764], + [-1.901870, -67.779457, 101.669186], + [-1.876412, -68.552864, 102.829296], + [-1.850954, -69.290451, 103.935677], + [-1.825496, -69.991325, 104.986988], + [-1.800038, -70.654541, 105.981812], + [-1.774580, -71.279190, 106.918785], + [-1.749122, -71.864319, 107.796478], + [-1.723664, -72.409088, 108.613632], + [-1.698206, -72.912567, 109.368851], + [-1.672748, -73.373871, 110.060806], + [-1.647290, -73.792130, 110.688194], + [-1.621832, -74.166512, 111.249767], + [-1.596374, -74.496147, 111.744221], + [-1.570916, -74.780251, 112.170376], + [-1.545458, -75.017998, 112.526997], + [-1.520000, -75.208656, 112.812984], + [-1.494542, -75.351448, 113.027172], + [-1.469084, -75.445686, 113.168530], + [-1.443626, -75.490677, 113.236015], + [-1.418168, -75.485771, 113.228657], + [-1.392710, -75.430344, 113.145515], + [-1.367252, -75.323830, 112.985744], + [-1.341794, -75.165688, 112.748531], + [-1.316336, -74.955406, 112.433109], + [-1.290878, -74.692551, 112.038826], + [-1.265420, -74.376694, 111.565041], + [-1.239962, -74.007477, 111.011215], + [-1.214504, -73.584579, 110.376869], + [-1.189046, -73.107742, 109.661613], + [-1.163588, -72.576752, 108.865128], + [-1.138130, -71.991455, 107.987183], + [-1.112672, -71.351707, 107.027561], + [-1.087214, -70.657486, 105.986229], + [-1.061756, -69.908813, 104.863220], + [-1.036298, -69.105721, 103.658581], + [-1.010840, -68.248337, 102.372505], + [-0.985382, -67.336861, 101.005291], + [-0.959924, -66.371513, 99.557270], + [-0.934466, -65.352615, 98.028923], + [-0.909008, -64.280533, 96.420799], + [-0.883550, -63.155693, 94.733540], + [-0.858092, -61.978588, 92.967882], + [-0.832634, -60.749775, 91.124662], + [-0.807176, -59.469845, 89.204767], + [-0.781718, -58.139503, 87.209255], + [-0.756260, -56.759487, 85.139231], + [-0.730802, -55.330616, 82.995924], + [-0.705344, -53.853729, 80.780594], + [-0.679886, -52.329796, 78.494694], + [-0.654428, -50.759762, 76.139643], + [-0.628970, -49.144699, 73.717049], + [-0.603512, -47.485737, 71.228605], + [-0.578054, -45.784004, 68.676006], + [-0.552596, -44.040764, 66.061146], + [-0.527138, -42.257267, 63.385900], + [-0.501680, -40.434883, 60.652325], + [-0.476222, -38.574947, 57.862421], + [-0.450764, -36.678982, 55.018473], + [-0.425306, -34.748432, 52.122648], + [-0.399848, -32.784836, 49.177254], + [-0.374390, -30.789810, 46.184715], + [-0.348932, -28.764952, 43.147428], + [-0.323474, -26.711969, 40.067954], + [-0.298016, -24.632576, 36.948864], + [-0.272558, -22.528547, 33.792821], + [-0.247100, -20.401667, 30.602500], +] +"""Spot knee joint parameter lookup table (102 entries). + +Columns: joint angle [rad], transmission ratio, output torque [N*m]. +Sourced from :mod:`isaaclab_assets.robots.spot`. +""" + + +def _run_authoring_introspection(actuator_cfgs: dict) -> dict: + """Instantiate Newton simulation, return Newton actuator introspection. + + Verifies that Lab configs are correctly authored to Newton USD schemas + and that Newton creates the expected controller/clamping/delay objects. + + Returns: + Dict with ``num_actuators``, ``actuator_info`` (list of per-actuator + dicts), and ``joint_pos`` (recorded trajectories). + """ + sim_cfg = SimulationCfg(dt=DT, physics=NEWTON_CFG, use_newton_actuators=True) + + with build_simulation_context( + device="cuda:0", + gravity_enabled=True, + add_ground_plane=True, + sim_cfg=sim_cfg, + ) as sim: + sim._app_control_on_stop_handle = None + + for i in range(NUM_ENVS): + sim_utils.create_prim( + f"/World/Env_{i}", "Xform", translation=(i * 3.0, 0, 0) + ) + + art_cfg = ANYMAL_C_CFG.replace( + actuators=actuator_cfgs, + prim_path="/World/Env_.*/Robot", + ) + articulation = Articulation(art_cfg) + sim.reset() + assert articulation.is_initialized + + model = SimulationManager.get_model() + + actuator_info = [] + for act in model.actuators: + ctrl_type = type(act.controller).__name__ + clamp_types = sorted( + type(c).__name__ for c in (act.clamping or []) + ) + actuator_info.append({ + "controller_type": ctrl_type, + "clamping_types": clamp_types, + "has_delay": act.delay is not None, + "num_indices": len(act.indices), + }) + + init_pos = wp.to_torch(articulation.data.joint_pos).clone() + target_pos = init_pos + TARGET_OFFSET + target_vel = torch.zeros_like(init_pos) + articulation.set_joint_position_target_index(target=target_pos) + articulation.set_joint_velocity_target_index(target=target_vel) + + recorded_pos = [] + for _ in range(NUM_STEPS): + articulation.write_data_to_sim() + sim.step() + articulation.update(DT) + recorded_pos.append(wp.to_torch(articulation.data.joint_pos).clone()) + + return { + "num_actuators": len(model.actuators), + "actuator_info": actuator_info, + "joint_pos": recorded_pos, + } + + +class TestRemotizedPDAuthoring(unittest.TestCase): + """Verify RemotizedPDActuatorCfg is authored as Newton PD + delay + + position-based clamping. + + Uses the Spot knee lookup table (102 entries) on ANYmal's KFE joints, + with IdealPD on HAA and HFE joints. + """ + + @classmethod + def setUpClass(cls): + from isaaclab.actuators.actuator_pd_cfg import RemotizedPDActuatorCfg # noqa: PLC0415 + + cls.result = _run_authoring_introspection({ + "hips": IdealPDActuatorCfg( + joint_names_expr=[".*HAA", ".*HFE"], + stiffness=40.0, + damping=5.0, + effort_limit=80.0, + ), + "knees": RemotizedPDActuatorCfg( + joint_names_expr=[".*KFE"], + stiffness=60.0, + damping=1.5, + effort_limit=80.0, + max_delay=3, + joint_parameter_lookup=SPOT_KNEE_LOOKUP, + ), + }) + + def test_num_actuators(self): + self.assertGreaterEqual(self.result["num_actuators"], 2) + + def test_kfe_controller_is_pd(self): + kfe_acts = [ + a for a in self.result["actuator_info"] + if "ClampingPositionBased" in a["clamping_types"] + ] + self.assertTrue(len(kfe_acts) > 0, "No actuator with position-based clamping found") + for a in kfe_acts: + self.assertEqual(a["controller_type"], "ControllerPD") + + def test_kfe_has_position_based_clamping(self): + kfe_acts = [ + a for a in self.result["actuator_info"] + if "ClampingPositionBased" in a["clamping_types"] + ] + self.assertTrue(len(kfe_acts) > 0, "Position-based clamping not found") + + def test_kfe_has_delay(self): + kfe_acts = [ + a for a in self.result["actuator_info"] + if "ClampingPositionBased" in a["clamping_types"] + ] + for a in kfe_acts: + self.assertTrue(a["has_delay"], "Delay not found on remotized KFE actuator") + + def test_trajectories_not_trivial(self): + first = self.result["joint_pos"][0] + last = self.result["joint_pos"][-1] + diff = (last - first).abs().max().item() + self.assertGreater(diff, 0.01, "Joints did not move — test is trivial") + + +# --------------------------------------------------------------------------- +# Neural network actuator authoring: MLP and LSTM +# --------------------------------------------------------------------------- + + +def _make_dummy_mlp_checkpoint(device: str = "cpu") -> str: + """Create a minimal TorchScript MLP checkpoint with metadata. + + The network accepts 6 inputs (3 history steps x 2 features per step + in pos_vel order) and outputs 1 effort. + """ + torch.manual_seed(42) + net = torch.nn.Sequential( + torch.nn.Linear(6, 8), + torch.nn.ELU(), + torch.nn.Linear(8, 1), + ).to(device).eval() + scripted = torch.jit.script(net) + + tmp = tempfile.NamedTemporaryFile(suffix=".pt", delete=False) + extra = { + "metadata.json": json.dumps({ + "model_type": "mlp", + "input_order": "pos_vel", + "input_idx": [0, 1, 2], + "pos_scale": 1.0, + "vel_scale": 0.5, + "torque_scale": 2.0, + }) + } + torch.jit.save(scripted, tmp.name, _extra_files=extra) + return tmp.name + + +class _DummyLSTM(torch.nn.Module): + """Minimal LSTM network for actuator testing.""" + + def __init__(self): + super().__init__() + self.lstm = torch.nn.LSTM(input_size=2, hidden_size=4, num_layers=1, batch_first=True) + self.fc = torch.nn.Linear(4, 1) + + def forward( + self, + x: torch.Tensor, + hc: tuple[torch.Tensor, torch.Tensor], + ) -> tuple[torch.Tensor, tuple[torch.Tensor, torch.Tensor]]: + out, hc_new = self.lstm(x, hc) + return self.fc(out[:, -1, :]), hc_new + + +def _make_dummy_lstm_checkpoint(device: str = "cpu") -> str: + """Create a minimal TorchScript LSTM checkpoint with metadata.""" + torch.manual_seed(42) + net = _DummyLSTM().to(device).eval() + scripted = torch.jit.script(net) + + tmp = tempfile.NamedTemporaryFile(suffix=".pt", delete=False) + extra = {"metadata.json": json.dumps({"model_type": "lstm"})} + torch.jit.save(scripted, tmp.name, _extra_files=extra) + return tmp.name + + +class TestNeuralMLPAuthoring(unittest.TestCase): + """Verify ActuatorNetMLPCfg is authored as Newton NeuralMLP controller + with DC motor clamping. + """ + + @classmethod + def setUpClass(cls): + from isaaclab.actuators.actuator_net_cfg import ActuatorNetMLPCfg # noqa: PLC0415 + + cls.mlp_path = _make_dummy_mlp_checkpoint() + cls.result = _run_authoring_introspection({ + "mlp_legs": ActuatorNetMLPCfg( + joint_names_expr=[".*HAA"], + network_file=cls.mlp_path, + saturation_effort=120.0, + effort_limit=80.0, + velocity_limit=7.5, + pos_scale=-1.0, + vel_scale=1.0, + torque_scale=1.0, + input_order="pos_vel", + input_idx=[0, 1, 2], + ), + "pd_legs": IdealPDActuatorCfg( + joint_names_expr=[".*HFE", ".*KFE"], + stiffness=40.0, + damping=5.0, + effort_limit=80.0, + ), + }) + + @classmethod + def tearDownClass(cls): + os.unlink(cls.mlp_path) + + def test_num_actuators(self): + self.assertGreaterEqual(self.result["num_actuators"], 2) + + def test_has_neural_mlp_controller(self): + mlp_acts = [ + a for a in self.result["actuator_info"] + if a["controller_type"] == "ControllerNeuralMLP" + ] + self.assertTrue(len(mlp_acts) > 0, "No NeuralMLP controller found") + + def test_mlp_has_dc_motor_clamping(self): + mlp_acts = [ + a for a in self.result["actuator_info"] + if a["controller_type"] == "ControllerNeuralMLP" + ] + for a in mlp_acts: + self.assertIn("ClampingDCMotor", a["clamping_types"]) + + def test_trajectories_not_trivial(self): + first = self.result["joint_pos"][0] + last = self.result["joint_pos"][-1] + diff = (last - first).abs().max().item() + self.assertGreater(diff, 0.01, "Joints did not move — test is trivial") + + +class TestNeuralLSTMAuthoring(unittest.TestCase): + """Verify ActuatorNetLSTMCfg is authored as Newton NeuralLSTM controller + with DC motor clamping. + """ + + @classmethod + def setUpClass(cls): + from isaaclab.actuators.actuator_net_cfg import ActuatorNetLSTMCfg # noqa: PLC0415 + + cls.lstm_path = _make_dummy_lstm_checkpoint() + cls.result = _run_authoring_introspection({ + "lstm_legs": ActuatorNetLSTMCfg( + joint_names_expr=[".*HAA"], + network_file=cls.lstm_path, + saturation_effort=120.0, + effort_limit=80.0, + velocity_limit=7.5, + ), + "pd_legs": IdealPDActuatorCfg( + joint_names_expr=[".*HFE", ".*KFE"], + stiffness=40.0, + damping=5.0, + effort_limit=80.0, + ), + }) + + @classmethod + def tearDownClass(cls): + os.unlink(cls.lstm_path) + + def test_num_actuators(self): + self.assertGreaterEqual(self.result["num_actuators"], 2) + + def test_has_neural_lstm_controller(self): + lstm_acts = [ + a for a in self.result["actuator_info"] + if a["controller_type"] == "ControllerNeuralLSTM" + ] + self.assertTrue(len(lstm_acts) > 0, "No NeuralLSTM controller found") + + def test_lstm_has_dc_motor_clamping(self): + lstm_acts = [ + a for a in self.result["actuator_info"] + if a["controller_type"] == "ControllerNeuralLSTM" + ] + for a in lstm_acts: + self.assertIn("ClampingDCMotor", a["clamping_types"]) + + def test_trajectories_not_trivial(self): + first = self.result["joint_pos"][0] + last = self.result["joint_pos"][-1] + diff = (last - first).abs().max().item() + self.assertGreater(diff, 0.01, "Joints did not move — test is trivial") + + +if __name__ == "__main__": + unittest.main() diff --git a/source/isaaclab_physx/isaaclab_physx/assets/articulation/articulation.py b/source/isaaclab_physx/isaaclab_physx/assets/articulation/articulation.py index 3b403ee8c6d4..dac70836127e 100644 --- a/source/isaaclab_physx/isaaclab_physx/assets/articulation/articulation.py +++ b/source/isaaclab_physx/isaaclab_physx/assets/articulation/articulation.py @@ -22,6 +22,16 @@ from isaaclab.actuators import ActuatorBase, ActuatorBaseCfg, ImplicitActuator from isaaclab.assets.articulation.base_articulation import BaseArticulation + +try: + from isaaclab_newton.actuators.newton_actuator_utils import ( + NewtonActuatorAdapter, + PhysxActuatorWrapper, + ) + + _HAS_NEWTON_ACTUATORS = True +except ImportError: + _HAS_NEWTON_ACTUATORS = False from isaaclab.sim.utils.queries import find_first_matching_prim, get_all_matching_child_prims from isaaclab.utils.string import resolve_matching_names, resolve_matching_names_values from isaaclab.utils.types import ArticulationActions @@ -217,7 +227,7 @@ def reset(self, env_ids: Sequence[int] | None = None, env_mask: wp.array | None # use ellipses object to skip initial indices. if (env_ids is None) or (env_ids == slice(None)): env_ids = slice(None) - # reset actuators + # reset actuators (including Newton-native adapter which owns its states) for actuator in self.actuators.values(): actuator.reset(env_ids) # reset external wrenches. @@ -251,14 +261,22 @@ def write_data_to_sim(self): ) self._instantaneous_wrench_composer.reset() - # apply actuator models - self._apply_actuator_model() - # write actions into simulation - self.root_view.set_dof_actuation_forces(self._joint_effort_target_sim, self._ALL_INDICES) - # position and velocity targets only for implicit actuators - if self._has_implicit_actuators: - self.root_view.set_dof_position_targets(self._joint_pos_target_sim, self._ALL_INDICES) - self.root_view.set_dof_velocity_targets(self._joint_vel_target_sim, self._ALL_INDICES) + if self._physx_actuator_wrapper is not None: + # Newton actuator path: zero effort, run implicit actuators + # (if any) which write directly into the staging buffer, then + # step Newton actuators which add their effort in-place. + self._apply_actuator_model_newton() + self.root_view.set_dof_actuation_forces(self._joint_effort_target_sim, self._ALL_INDICES) + if self._has_implicit_actuators: + self.root_view.set_dof_position_targets(self._joint_pos_target_sim, self._ALL_INDICES) + self.root_view.set_dof_velocity_targets(self._joint_vel_target_sim, self._ALL_INDICES) + else: + # Standard Lab actuator path + self._apply_actuator_model() + self.root_view.set_dof_actuation_forces(self._joint_effort_target_sim, self._ALL_INDICES) + if self._has_implicit_actuators: + self.root_view.set_dof_position_targets(self._joint_pos_target_sim, self._ALL_INDICES) + self.root_view.set_dof_velocity_targets(self._joint_vel_target_sim, self._ALL_INDICES) def update(self, dt: float): """Updates the simulation data. @@ -3726,166 +3744,48 @@ def _process_actuators_cfg(self): """Process and apply articulation joint properties.""" # create actuators self.actuators = dict() + self._physx_actuator_wrapper = None # flag for implicit actuators # if this is false, we by-pass certain checks when doing actuator-related operations self._has_implicit_actuators = False - # iterate over all actuator configurations - for actuator_name, actuator_cfg in self.cfg.actuators.items(): - # type annotation for type checkers - actuator_cfg: ActuatorBaseCfg - # create actuator group - joint_ids, joint_names = self.find_joints(actuator_cfg.joint_names_expr) - # check if any joints are found - if len(joint_names) == 0: - raise ValueError( - f"No joints found for actuator group: {actuator_name} with joint name expression:" - f" {actuator_cfg.joint_names_expr}." - ) - # resolve joint indices - # we pass a slice if all joints are selected to avoid indexing overhead - if len(joint_names) == self.num_joints: - joint_ids = slice(None) - else: - joint_ids = torch.tensor(joint_ids, device=self.device, dtype=torch.int32) - # create actuator collection - # note: for efficiency avoid indexing when over all indices - actuator: ActuatorBase = actuator_cfg.class_type( - cfg=actuator_cfg, - joint_names=joint_names, - joint_ids=joint_ids, + _use_newton_actuators = getattr(self._sim_cfg, "use_newton_actuators", False) + + if _HAS_NEWTON_ACTUATORS and _use_newton_actuators: + from isaaclab.sim.utils.stage import get_current_stage # noqa: PLC0415 + + adapter = NewtonActuatorAdapter.from_usd( + stage=get_current_stage(), + joint_names=self.joint_names, num_envs=self.num_instances, + num_joints=self.num_joints, device=self.device, - stiffness=wp.to_torch(self._data.joint_stiffness)[:, joint_ids], - damping=wp.to_torch(self._data.joint_damping)[:, joint_ids], - armature=wp.to_torch(self._data.joint_armature)[:, joint_ids], - friction=wp.to_torch(self._data.joint_friction_coeff)[:, joint_ids], - dynamic_friction=wp.to_torch(self._data.joint_dynamic_friction_coeff)[:, joint_ids], - viscous_friction=wp.to_torch(self._data.joint_viscous_friction_coeff)[:, joint_ids], - effort_limit=wp.to_torch(self._data.joint_effort_limits)[:, joint_ids].clone(), - velocity_limit=wp.to_torch(self._data.joint_vel_limits)[:, joint_ids], - ) - # store actuator group - self.actuators[actuator_name] = actuator - # Store the configured values from the actuator model - # note: this is the value configured in the actuator model (for implicit and explicit actuators) - joint_ids = actuator.joint_indices - if joint_ids == slice(None): - joint_ids = self._ALL_JOINT_INDICES - wp.launch( - shared_kernels.write_2d_data_to_buffer_with_indices, - dim=(self.num_instances, joint_ids.shape[0]), - inputs=[ - actuator.stiffness, - self._ALL_INDICES, - joint_ids, - False, - ], - outputs=[ - self.data._joint_stiffness, - ], - device=self.device, - ) - wp.launch( - shared_kernels.write_2d_data_to_buffer_with_indices, - dim=(self.num_instances, joint_ids.shape[0]), - inputs=[ - actuator.damping, - self._ALL_INDICES, - joint_ids, - False, - ], - outputs=[ - self.data._joint_damping, - ], - device=self.device, - ) - wp.launch( - shared_kernels.write_2d_data_to_buffer_with_indices, - dim=(self.num_instances, joint_ids.shape[0]), - inputs=[ - actuator.armature, - self._ALL_INDICES, - joint_ids, - False, - ], - outputs=[ - self.data._joint_armature, - ], - device=self.device, - ) - wp.launch( - shared_kernels.write_2d_data_to_buffer_with_indices, - dim=(self.num_instances, joint_ids.shape[0]), - inputs=[ - actuator.friction, - self._ALL_INDICES, - joint_ids, - False, - ], - outputs=[ - self.data._joint_friction_coeff, - ], - device=self.device, - ) - wp.launch( - shared_kernels.write_2d_data_to_buffer_with_indices, - dim=(self.num_instances, joint_ids.shape[0]), - inputs=[ - actuator.dynamic_friction, - self._ALL_INDICES, - joint_ids, - False, - ], - outputs=[ - self.data._joint_dynamic_friction_coeff, - ], - device=self.device, - ) - wp.launch( - shared_kernels.write_2d_data_to_buffer_with_indices, - dim=(self.num_instances, joint_ids.shape[0]), - inputs=[ - actuator.viscous_friction, - self._ALL_INDICES, - joint_ids, - False, - ], - outputs=[ - self.data._joint_viscous_friction_coeff, - ], - device=self.device, - ) - # set the passed gains and limits into the simulation - if isinstance(actuator, ImplicitActuator): - self._has_implicit_actuators = True - # the gains and limits are set into the simulation since actuator model is implicit - self.write_joint_stiffness_to_sim_index(stiffness=actuator.stiffness, joint_ids=actuator.joint_indices) - self.write_joint_damping_to_sim_index(damping=actuator.damping, joint_ids=actuator.joint_indices) - else: - # the gains and limits are processed by the actuator model - # we set gains to zero, and torque limit to a high value in simulation to avoid any interference - self.write_joint_stiffness_to_sim_index(stiffness=0.0, joint_ids=actuator.joint_indices) - self.write_joint_damping_to_sim_index(damping=0.0, joint_ids=actuator.joint_indices) - - # Set common properties into the simulation - self.write_joint_effort_limit_to_sim_index( - limits=actuator.effort_limit_sim, joint_ids=actuator.joint_indices - ) - self.write_joint_velocity_limit_to_sim_index( - limits=actuator.velocity_limit_sim, joint_ids=actuator.joint_indices - ) - self.write_joint_armature_to_sim_index(armature=actuator.armature, joint_ids=actuator.joint_indices) - self.write_joint_friction_coefficient_to_sim_index( - joint_friction_coeff=actuator.friction, joint_ids=actuator.joint_indices - ) - self.write_joint_dynamic_friction_coefficient_to_sim_index( - joint_dynamic_friction_coeff=actuator.dynamic_friction, joint_ids=actuator.joint_indices - ) - self.write_joint_viscous_friction_coefficient_to_sim_index( - joint_viscous_friction_coeff=actuator.viscous_friction, joint_ids=actuator.joint_indices ) + self._physx_actuator_wrapper = PhysxActuatorWrapper() + adapter.finalize() + self.actuators["newton"] = adapter + self.write_joint_stiffness_to_sim_index(stiffness=0.0, joint_ids=adapter.joint_indices) + self.write_joint_damping_to_sim_index(damping=0.0, joint_ids=adapter.joint_indices) + + # Also create any implicit actuators defined in the configs + for actuator_name, actuator_cfg in self.cfg.actuators.items(): + cls_type = actuator_cfg.class_type + is_implicit = ( + "ImplicitActuator" in cls_type + if isinstance(cls_type, str) + else issubclass(cls_type, ImplicitActuator) + ) + if not is_implicit: + continue + self._create_lab_actuator(actuator_name, actuator_cfg) + + return + + # --- Standard Isaac Lab actuator path --- + for actuator_name, actuator_cfg in self.cfg.actuators.items(): + self._create_lab_actuator(actuator_name, actuator_cfg) + # perform some sanity checks to ensure actuators are prepared correctly total_act_joints = sum(actuator.num_joints for actuator in self.actuators.values()) if total_act_joints != (self.num_joints - self.num_fixed_tendons): @@ -3897,6 +3797,8 @@ def _process_actuators_cfg(self): if self.cfg.actuator_value_resolution_debug_print: t = PrettyTable(["Group", "Property", "Name", "ID", "USD Value", "ActutatorCfg Value", "Applied"]) for actuator_group, actuator in self.actuators.items(): + if _HAS_NEWTON_ACTUATORS and isinstance(actuator, NewtonActuatorAdapter): + continue group_count = 0 for property, resolution_details in actuator.joint_property_resolution_table.items(): for prop_idx, resolution_detail in enumerate(resolution_details): @@ -3907,6 +3809,81 @@ def _process_actuators_cfg(self): group_count += 1 logger.warning(f"\nActuatorCfg-USD Value Discrepancy Resolution (matching values are skipped): \n{t}") + def _create_lab_actuator(self, actuator_name: str, actuator_cfg: ActuatorBaseCfg) -> None: + """Instantiate a single Lab actuator from its config and write properties to sim.""" + joint_ids, joint_names = self.find_joints(actuator_cfg.joint_names_expr) + if len(joint_names) == 0: + raise ValueError( + f"No joints found for actuator group: {actuator_name} with joint name expression:" + f" {actuator_cfg.joint_names_expr}." + ) + if len(joint_names) == self.num_joints: + joint_ids = slice(None) + else: + joint_ids = torch.tensor(joint_ids, device=self.device, dtype=torch.int32) + + actuator: ActuatorBase = actuator_cfg.class_type( + cfg=actuator_cfg, + joint_names=joint_names, + joint_ids=joint_ids, + num_envs=self.num_instances, + device=self.device, + stiffness=wp.to_torch(self._data.joint_stiffness)[:, joint_ids], + damping=wp.to_torch(self._data.joint_damping)[:, joint_ids], + armature=wp.to_torch(self._data.joint_armature)[:, joint_ids], + friction=wp.to_torch(self._data.joint_friction_coeff)[:, joint_ids], + dynamic_friction=wp.to_torch(self._data.joint_dynamic_friction_coeff)[:, joint_ids], + viscous_friction=wp.to_torch(self._data.joint_viscous_friction_coeff)[:, joint_ids], + effort_limit=wp.to_torch(self._data.joint_effort_limits)[:, joint_ids].clone(), + velocity_limit=wp.to_torch(self._data.joint_vel_limits)[:, joint_ids], + ) + self.actuators[actuator_name] = actuator + + # Store the configured values from the actuator model + j_ids = actuator.joint_indices + if j_ids == slice(None): + j_ids = self._ALL_JOINT_INDICES + for attr, buf in ( + (actuator.stiffness, self.data._joint_stiffness), + (actuator.damping, self.data._joint_damping), + (actuator.armature, self.data._joint_armature), + (actuator.friction, self.data._joint_friction_coeff), + (actuator.dynamic_friction, self.data._joint_dynamic_friction_coeff), + (actuator.viscous_friction, self.data._joint_viscous_friction_coeff), + ): + wp.launch( + shared_kernels.write_2d_data_to_buffer_with_indices, + dim=(self.num_instances, j_ids.shape[0]), + inputs=[attr, self._ALL_INDICES, j_ids, False], + outputs=[buf], + device=self.device, + ) + + if isinstance(actuator, ImplicitActuator): + self._has_implicit_actuators = True + self.write_joint_stiffness_to_sim_index(stiffness=actuator.stiffness, joint_ids=actuator.joint_indices) + self.write_joint_damping_to_sim_index(damping=actuator.damping, joint_ids=actuator.joint_indices) + else: + self.write_joint_stiffness_to_sim_index(stiffness=0.0, joint_ids=actuator.joint_indices) + self.write_joint_damping_to_sim_index(damping=0.0, joint_ids=actuator.joint_indices) + + self.write_joint_effort_limit_to_sim_index( + limits=actuator.effort_limit_sim, joint_ids=actuator.joint_indices, + ) + self.write_joint_velocity_limit_to_sim_index( + limits=actuator.velocity_limit_sim, joint_ids=actuator.joint_indices, + ) + self.write_joint_armature_to_sim_index(armature=actuator.armature, joint_ids=actuator.joint_indices) + self.write_joint_friction_coefficient_to_sim_index( + joint_friction_coeff=actuator.friction, joint_ids=actuator.joint_indices, + ) + self.write_joint_dynamic_friction_coefficient_to_sim_index( + joint_dynamic_friction_coeff=actuator.dynamic_friction, joint_ids=actuator.joint_indices, + ) + self.write_joint_viscous_friction_coefficient_to_sim_index( + joint_viscous_friction_coeff=actuator.viscous_friction, joint_ids=actuator.joint_indices, + ) + def _process_tendons(self): """Process fixed and spatial tendons.""" # create a list to store the fixed tendon names @@ -4000,6 +3977,64 @@ def _apply_actuator_model(self): device=self.device, ) + def _apply_actuator_model_newton(self): + """Process actuators when Newton-native actuators are active. + + Zeros the effort staging buffer, processes any implicit Lab actuators + (which write directly into it), then steps Newton actuators whose + ``joint_f`` is a flat view of the same buffer — no scatter needed. + """ + # Zero the effort buffer; implicit actuators and Newton both write into it + self._joint_effort_target_sim.zero_() + + # Process implicit Lab actuators (if any) + for actuator in self.actuators.values(): + if _HAS_NEWTON_ACTUATORS and isinstance(actuator, NewtonActuatorAdapter): + continue + control_action = ArticulationActions( + joint_positions=wp.to_torch(self._data.joint_pos_target)[:, actuator.joint_indices], + joint_velocities=wp.to_torch(self._data.joint_vel_target)[:, actuator.joint_indices], + joint_efforts=wp.to_torch(self._data.joint_effort_target)[:, actuator.joint_indices], + joint_indices=actuator.joint_indices, + ) + control_action = actuator.compute( + control_action, + joint_pos=wp.to_torch(self._data.joint_pos)[:, actuator.joint_indices], + joint_vel=wp.to_torch(self._data.joint_vel)[:, actuator.joint_indices], + ) + joint_indices = actuator.joint_indices + if actuator.joint_indices == slice(None) or actuator.joint_indices is None: + joint_indices = self._ALL_JOINT_INDICES + wp.launch( + articulation_kernels.update_targets, + dim=(self.num_instances, joint_indices.shape[0]), + inputs=[ + control_action.joint_positions, + control_action.joint_velocities, + control_action.joint_efforts, + joint_indices, + ], + outputs=[ + self._joint_pos_target_sim, + self._joint_vel_target_sim, + self._joint_effort_target_sim, + ], + device=self.device, + ) + + # Step Newton actuators — joint_f is a flat view of _joint_effort_target_sim + # so Newton writes effort directly into the staging buffer. + newton_adapter = self.actuators.get("newton") + w = self._physx_actuator_wrapper + w.joint_q = self._data.joint_pos.reshape(-1) + w.joint_qd = self._data.joint_vel.reshape(-1) + w.joint_target_pos = self._data.joint_pos_target.reshape(-1) + w.joint_target_vel = self._data.joint_vel_target.reshape(-1) + w.joint_act = self._data.joint_effort_target.reshape(-1) + w.joint_f = self._joint_effort_target_sim.reshape(-1) + + newton_adapter.step(w, w, SimulationManager.get_physics_dt()) + """ Internal helpers -- Debugging. """ diff --git a/source/isaaclab_physx/test/assets/test_pd_equivalence.py b/source/isaaclab_physx/test/assets/test_pd_equivalence.py new file mode 100644 index 000000000000..c07bf6c9fdb5 --- /dev/null +++ b/source/isaaclab_physx/test/assets/test_pd_equivalence.py @@ -0,0 +1,554 @@ +# Copyright (c) 2022-2026, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +"""PD actuator equivalence tests on ANYmal-C (floating-base quadruped) — PhysX backend. + +Compares IsaacLab-native actuators against Newton-native actuators (created +from the same Lab configs via USD authoring, stepped via +:class:`PhysxActuatorWrapper`) on the PhysX physics backend. Both paths +must produce identical joint trajectories within tolerance. + +Using ANYmal-C — a 12-DOF quadruped on a floating base — exercises the +full Lab-to-Newton config translation pipeline on a real-world robot. +""" + +from isaaclab.app import AppLauncher + +simulation_app = AppLauncher(headless=True).app + +import json +import os +import tempfile +import unittest + +import torch +import warp as wp +from isaaclab_assets import ANYMAL_C_CFG +from isaaclab_physx.assets import Articulation +from isaaclab_physx.physics import PhysxCfg + +import isaaclab.sim as sim_utils +from isaaclab.actuators import DCMotorCfg, IdealPDActuatorCfg, ImplicitActuatorCfg +from isaaclab.sim import SimulationCfg, build_simulation_context + +# --------------------------------------------------------------------------- +# Constants +# --------------------------------------------------------------------------- + +NUM_ENVS = 2 +NUM_STEPS = 10 +DT = 1.0 / 120.0 +TARGET_OFFSET = 0.1 # [rad] added to initial joint positions + +# --------------------------------------------------------------------------- +# Actuator configurations under test +# --------------------------------------------------------------------------- + +IDEAL_PD_ACTUATORS = { + "legs": IdealPDActuatorCfg( + joint_names_expr=[".*HAA", ".*HFE", ".*KFE"], + stiffness=40.0, + damping=5.0, + effort_limit=80.0, + ), +} + +DC_MOTOR_ACTUATORS = { + "legs": DCMotorCfg( + joint_names_expr=[".*HAA", ".*HFE", ".*KFE"], + saturation_effort=120.0, + effort_limit=80.0, + velocity_limit=7.5, + stiffness={".*": 40.0}, + damping={".*": 5.0}, + ), +} + +MIXED_WITH_IMPLICIT_ACTUATORS = { + "hips": ImplicitActuatorCfg( + joint_names_expr=[".*HAA"], + stiffness=40.0, + damping=5.0, + ), + "thighs": IdealPDActuatorCfg( + joint_names_expr=[".*HFE"], + stiffness=40.0, + damping=5.0, + effort_limit=80.0, + ), + "knees": DCMotorCfg( + joint_names_expr=[".*KFE"], + saturation_effort=120.0, + effort_limit=80.0, + velocity_limit=7.5, + stiffness=40.0, + damping=5.0, + ), +} + +# --------------------------------------------------------------------------- +# Simulation runner +# --------------------------------------------------------------------------- + + +def _run_simulation( + actuators: dict, + use_newton_actuators: bool, + *, + num_steps: int = NUM_STEPS, +) -> dict: + """Run ANYmal-C on PhysX and return recorded joint trajectories. + + Args: + actuators: Actuator config dict overriding ANYmal's defaults. + use_newton_actuators: Use Newton-native actuators (via + :class:`PhysxActuatorWrapper`) when ``True``. + num_steps: Number of simulation steps. + + Returns: + Dict with ``joint_pos`` and ``joint_vel``, each a list of + ``(NUM_ENVS, num_joints)`` tensors. + """ + sim_cfg = SimulationCfg( + dt=DT, + physics=PhysxCfg(), + use_newton_actuators=use_newton_actuators, + ) + + with build_simulation_context( + device="cuda:0", + gravity_enabled=True, + add_ground_plane=True, + sim_cfg=sim_cfg, + ) as sim: + sim._app_control_on_stop_handle = None + + for i in range(NUM_ENVS): + sim_utils.create_prim( + f"/World/Env_{i}", "Xform", translation=(i * 3.0, 0, 0) + ) + + art_cfg = ANYMAL_C_CFG.replace( + actuators=actuators, + prim_path="/World/Env_.*/Robot", + ) + articulation = Articulation(art_cfg) + sim.reset() + assert articulation.is_initialized + + init_pos = wp.to_torch(articulation.data.joint_pos).clone() + target_pos = init_pos + TARGET_OFFSET + target_vel = torch.zeros_like(init_pos) + + articulation.set_joint_position_target_index(target=target_pos) + articulation.set_joint_velocity_target_index(target=target_vel) + + recorded_pos, recorded_vel = [], [] + for _ in range(num_steps): + articulation.write_data_to_sim() + sim.step() + articulation.update(DT) + + recorded_pos.append(wp.to_torch(articulation.data.joint_pos).clone()) + recorded_vel.append(wp.to_torch(articulation.data.joint_vel).clone()) + + return {"joint_pos": recorded_pos, "joint_vel": recorded_vel} + + +# --------------------------------------------------------------------------- +# Base test class +# --------------------------------------------------------------------------- + + +class _EquivalenceTestBase(unittest.TestCase): + """Base for Lab-vs-Newton equivalence tests on the PhysX backend. + + Subclasses set ``actuators`` to the config under test. ``setUpClass`` + runs the simulation with both ``use_newton_actuators=False`` (Lab path) + and ``True`` (Newton via PhysxActuatorWrapper) and stores the results. + """ + + __test__ = False + actuators: dict = {} + pos_atol: float = 2e-3 + pos_rtol: float = 1e-3 + vel_atol: float = 0.1 + vel_rtol: float = 1e-2 + + @classmethod + def setUpClass(cls): + cls.lab_result = _run_simulation(cls.actuators, use_newton_actuators=False) + cls.newton_result = _run_simulation(cls.actuators, use_newton_actuators=True) + + def test_joint_positions_match(self): + for step_i, (lab, newton) in enumerate( + zip(self.lab_result["joint_pos"], self.newton_result["joint_pos"]) + ): + torch.testing.assert_close( + lab, + newton, + atol=self.pos_atol, + rtol=self.pos_rtol, + msg=f"Joint positions diverged at step {step_i}", + ) + + def test_joint_velocities_match(self): + for step_i, (lab, newton) in enumerate( + zip(self.lab_result["joint_vel"], self.newton_result["joint_vel"]) + ): + torch.testing.assert_close( + lab, + newton, + atol=self.vel_atol, + rtol=self.vel_rtol, + msg=f"Joint velocities diverged at step {step_i}", + ) + + def test_trajectories_not_trivial(self): + first = self.lab_result["joint_pos"][0] + last = self.lab_result["joint_pos"][-1] + diff = (last - first).abs().max().item() + self.assertGreater(diff, 0.01, "Joints did not move — test is trivial") + + +# --------------------------------------------------------------------------- +# Equivalence tests with different actuator types +# --------------------------------------------------------------------------- + + +class TestIdealPDEquivalence(_EquivalenceTestBase): + """IdealPDActuator on all 12 joints: Lab vs Newton (PhysX backend).""" + + __test__ = True + actuators = IDEAL_PD_ACTUATORS + + +class TestDCMotorEquivalence(_EquivalenceTestBase): + """DCMotor actuator on all 12 joints: Lab vs Newton (PhysX backend).""" + + __test__ = True + actuators = DC_MOTOR_ACTUATORS + + +class TestMixedWithImplicitEquivalence(_EquivalenceTestBase): + """Implicit HAA + IdealPD HFE + DCMotor KFE: Lab vs Newton (PhysX). + + Verifies that implicit actuators (handled by PhysX joint drives) + coexist correctly with explicit Newton actuators via PhysxActuatorWrapper. + """ + + __test__ = True + actuators = MIXED_WITH_IMPLICIT_ACTUATORS + + +# --------------------------------------------------------------------------- +# RemotizedPD authoring: PD + delay + position-based clamping lookup table +# --------------------------------------------------------------------------- + +SPOT_KNEE_LOOKUP = [ + [-2.792900, -24.776718, 37.165077], + [-2.767442, -26.290108, 39.435162], + [-2.741984, -27.793369, 41.690054], + [-2.716526, -29.285997, 43.928996], + [-2.691068, -30.767536, 46.151304], + [-2.665610, -32.237423, 48.356134], + [-2.640152, -33.695168, 50.542751], + [-2.614694, -35.140221, 52.710331], + [-2.589236, -36.572052, 54.858078], + [-2.563778, -37.990086, 56.985128], + [-2.538320, -39.393730, 59.090595], + [-2.512862, -40.782406, 61.173609], + [-2.487404, -42.155487, 63.233231], + [-2.461946, -43.512371, 65.268557], + [-2.436488, -44.852371, 67.278557], + [-2.411030, -46.174873, 69.262310], + [-2.385572, -47.479156, 71.218735], + [-2.360114, -48.764549, 73.146824], + [-2.334656, -50.030334, 75.045502], + [-2.309198, -51.275761, 76.913641], + [-2.283740, -52.500103, 78.750154], + [-2.258282, -53.702587, 80.553881], + [-2.232824, -54.882442, 82.323664], + [-2.207366, -56.038860, 84.058290], + [-2.181908, -57.171028, 85.756542], + [-2.156450, -58.278133, 87.417200], + [-2.130992, -59.359314, 89.038971], + [-2.105534, -60.413738, 90.620607], + [-2.080076, -61.440529, 92.160793], + [-2.054618, -62.438812, 93.658218], + [-2.029160, -63.407692, 95.111538], + [-2.003702, -64.346268, 96.519402], + [-1.978244, -65.253670, 97.880505], + [-1.952786, -66.128944, 99.193417], + [-1.927328, -66.971176, 100.456764], + [-1.901870, -67.779457, 101.669186], + [-1.876412, -68.552864, 102.829296], + [-1.850954, -69.290451, 103.935677], + [-1.825496, -69.991325, 104.986988], + [-1.800038, -70.654541, 105.981812], + [-1.774580, -71.279190, 106.918785], + [-1.749122, -71.864319, 107.796478], + [-1.723664, -72.409088, 108.613632], + [-1.698206, -72.912567, 109.368851], + [-1.672748, -73.373871, 110.060806], + [-1.647290, -73.792130, 110.688194], + [-1.621832, -74.166512, 111.249767], + [-1.596374, -74.496147, 111.744221], + [-1.570916, -74.780251, 112.170376], + [-1.545458, -75.017998, 112.526997], + [-1.520000, -75.208656, 112.812984], + [-1.494542, -75.351448, 113.027172], + [-1.469084, -75.445686, 113.168530], + [-1.443626, -75.490677, 113.236015], + [-1.418168, -75.485771, 113.228657], + [-1.392710, -75.430344, 113.145515], + [-1.367252, -75.323830, 112.985744], + [-1.341794, -75.165688, 112.748531], + [-1.316336, -74.955406, 112.433109], + [-1.290878, -74.692551, 112.038826], + [-1.265420, -74.376694, 111.565041], + [-1.239962, -74.007477, 111.011215], + [-1.214504, -73.584579, 110.376869], + [-1.189046, -73.107742, 109.661613], + [-1.163588, -72.576752, 108.865128], + [-1.138130, -71.991455, 107.987183], + [-1.112672, -71.351707, 107.027561], + [-1.087214, -70.657486, 105.986229], + [-1.061756, -69.908813, 104.863220], + [-1.036298, -69.105721, 103.658581], + [-1.010840, -68.248337, 102.372505], + [-0.985382, -67.336861, 101.005291], + [-0.959924, -66.371513, 99.557270], + [-0.934466, -65.352615, 98.028923], + [-0.909008, -64.280533, 96.420799], + [-0.883550, -63.155693, 94.733540], + [-0.858092, -61.978588, 92.967882], + [-0.832634, -60.749775, 91.124662], + [-0.807176, -59.469845, 89.204767], + [-0.781718, -58.139503, 87.209255], + [-0.756260, -56.759487, 85.139231], + [-0.730802, -55.330616, 82.995924], + [-0.705344, -53.853729, 80.780594], + [-0.679886, -52.329796, 78.494694], + [-0.654428, -50.759762, 76.139643], + [-0.628970, -49.144699, 73.717049], + [-0.603512, -47.485737, 71.228605], + [-0.578054, -45.784004, 68.676006], + [-0.552596, -44.040764, 66.061146], + [-0.527138, -42.257267, 63.385900], + [-0.501680, -40.434883, 60.652325], + [-0.476222, -38.574947, 57.862421], + [-0.450764, -36.678982, 55.018473], + [-0.425306, -34.748432, 52.122648], + [-0.399848, -32.784836, 49.177254], + [-0.374390, -30.789810, 46.184715], + [-0.348932, -28.764952, 43.147428], + [-0.323474, -26.711969, 40.067954], + [-0.298016, -24.632576, 36.948864], + [-0.272558, -22.528547, 33.792821], + [-0.247100, -20.401667, 30.602500], +] + + +class TestRemotizedPDFunctional(unittest.TestCase): + """Verify RemotizedPDActuatorCfg runs correctly on PhysX with Newton actuators. + + Uses the Spot knee lookup table (102 entries) on ANYmal's KFE joints. + """ + + @classmethod + def setUpClass(cls): + from isaaclab.actuators.actuator_pd_cfg import RemotizedPDActuatorCfg # noqa: PLC0415 + + cls.result = _run_simulation( + { + "hips": IdealPDActuatorCfg( + joint_names_expr=[".*HAA", ".*HFE"], + stiffness=40.0, + damping=5.0, + effort_limit=80.0, + ), + "knees": RemotizedPDActuatorCfg( + joint_names_expr=[".*KFE"], + stiffness=60.0, + damping=1.5, + effort_limit=80.0, + max_delay=3, + joint_parameter_lookup=SPOT_KNEE_LOOKUP, + ), + }, + use_newton_actuators=True, + ) + + def test_trajectories_not_trivial(self): + first = self.result["joint_pos"][0] + last = self.result["joint_pos"][-1] + diff = (last - first).abs().max().item() + self.assertGreater(diff, 0.01, "Joints did not move — test is trivial") + + def test_positions_finite(self): + for step_i, pos in enumerate(self.result["joint_pos"]): + self.assertTrue( + torch.isfinite(pos).all(), + f"Non-finite positions at step {step_i}", + ) + + +# --------------------------------------------------------------------------- +# Neural network actuator authoring: MLP and LSTM +# --------------------------------------------------------------------------- + + +def _make_dummy_mlp_checkpoint(device: str = "cpu") -> str: + """Create a minimal TorchScript MLP checkpoint with metadata.""" + torch.manual_seed(42) + net = torch.nn.Sequential( + torch.nn.Linear(6, 8), + torch.nn.ELU(), + torch.nn.Linear(8, 1), + ).to(device).eval() + scripted = torch.jit.script(net) + + tmp = tempfile.NamedTemporaryFile(suffix=".pt", delete=False) + extra = { + "metadata.json": json.dumps({ + "model_type": "mlp", + "input_order": "pos_vel", + "input_idx": [0, 1, 2], + "pos_scale": 1.0, + "vel_scale": 0.5, + "torque_scale": 2.0, + }) + } + torch.jit.save(scripted, tmp.name, _extra_files=extra) + return tmp.name + + +class _DummyLSTM(torch.nn.Module): + """Minimal LSTM network for actuator testing.""" + + def __init__(self): + super().__init__() + self.lstm = torch.nn.LSTM(input_size=2, hidden_size=4, num_layers=1, batch_first=True) + self.fc = torch.nn.Linear(4, 1) + + def forward( + self, + x: torch.Tensor, + hc: tuple[torch.Tensor, torch.Tensor], + ) -> tuple[torch.Tensor, tuple[torch.Tensor, torch.Tensor]]: + out, hc_new = self.lstm(x, hc) + return self.fc(out[:, -1, :]), hc_new + + +def _make_dummy_lstm_checkpoint(device: str = "cpu") -> str: + """Create a minimal TorchScript LSTM checkpoint with metadata.""" + torch.manual_seed(42) + net = _DummyLSTM().to(device).eval() + scripted = torch.jit.script(net) + + tmp = tempfile.NamedTemporaryFile(suffix=".pt", delete=False) + extra = {"metadata.json": json.dumps({"model_type": "lstm"})} + torch.jit.save(scripted, tmp.name, _extra_files=extra) + return tmp.name + + +class TestNeuralMLPFunctional(unittest.TestCase): + """Verify ActuatorNetMLPCfg runs on PhysX with Newton actuators.""" + + @classmethod + def setUpClass(cls): + from isaaclab.actuators.actuator_net_cfg import ActuatorNetMLPCfg # noqa: PLC0415 + + cls.mlp_path = _make_dummy_mlp_checkpoint() + cls.result = _run_simulation( + { + "mlp_legs": ActuatorNetMLPCfg( + joint_names_expr=[".*HAA"], + network_file=cls.mlp_path, + saturation_effort=120.0, + effort_limit=80.0, + velocity_limit=7.5, + pos_scale=-1.0, + vel_scale=1.0, + torque_scale=1.0, + input_order="pos_vel", + input_idx=[0, 1, 2], + ), + "pd_legs": IdealPDActuatorCfg( + joint_names_expr=[".*HFE", ".*KFE"], + stiffness=40.0, + damping=5.0, + effort_limit=80.0, + ), + }, + use_newton_actuators=True, + ) + + @classmethod + def tearDownClass(cls): + os.unlink(cls.mlp_path) + + def test_trajectories_not_trivial(self): + first = self.result["joint_pos"][0] + last = self.result["joint_pos"][-1] + diff = (last - first).abs().max().item() + self.assertGreater(diff, 0.01, "Joints did not move — test is trivial") + + def test_positions_finite(self): + for step_i, pos in enumerate(self.result["joint_pos"]): + self.assertTrue( + torch.isfinite(pos).all(), + f"Non-finite positions at step {step_i}", + ) + + +class TestNeuralLSTMFunctional(unittest.TestCase): + """Verify ActuatorNetLSTMCfg runs on PhysX with Newton actuators.""" + + @classmethod + def setUpClass(cls): + from isaaclab.actuators.actuator_net_cfg import ActuatorNetLSTMCfg # noqa: PLC0415 + + cls.lstm_path = _make_dummy_lstm_checkpoint() + cls.result = _run_simulation( + { + "lstm_legs": ActuatorNetLSTMCfg( + joint_names_expr=[".*HAA"], + network_file=cls.lstm_path, + saturation_effort=120.0, + effort_limit=80.0, + velocity_limit=7.5, + ), + "pd_legs": IdealPDActuatorCfg( + joint_names_expr=[".*HFE", ".*KFE"], + stiffness=40.0, + damping=5.0, + effort_limit=80.0, + ), + }, + use_newton_actuators=True, + ) + + @classmethod + def tearDownClass(cls): + os.unlink(cls.lstm_path) + + def test_trajectories_not_trivial(self): + first = self.result["joint_pos"][0] + last = self.result["joint_pos"][-1] + diff = (last - first).abs().max().item() + self.assertGreater(diff, 0.01, "Joints did not move — test is trivial") + + def test_positions_finite(self): + for step_i, pos in enumerate(self.result["joint_pos"]): + self.assertTrue( + torch.isfinite(pos).all(), + f"Non-finite positions at step {step_i}", + ) + + +if __name__ == "__main__": + unittest.main() From 3a867223f8ba645c3b86d700292e08573b941d56 Mon Sep 17 00:00:00 2001 From: Julia von Muralt Date: Wed, 29 Apr 2026 16:19:42 +0200 Subject: [PATCH 02/35] improve tests --- .../actuators/newton_actuator_utils.py | 16 +++++- ...lence.py => test_actuators_equivalence.py} | 51 ++++++++++++++++++- .../assets/articulation/articulation.py | 4 ++ ...lence.py => test_actuators_equivalence.py} | 48 ++++++++++++++++- 4 files changed, 116 insertions(+), 3 deletions(-) rename source/isaaclab_newton/test/assets/{test_pd_equivalence.py => test_actuators_equivalence.py} (94%) rename source/isaaclab_physx/test/assets/{test_pd_equivalence.py => test_actuators_equivalence.py} (93%) diff --git a/source/isaaclab_newton/isaaclab_newton/actuators/newton_actuator_utils.py b/source/isaaclab_newton/isaaclab_newton/actuators/newton_actuator_utils.py index 159779a8addc..eef438c60a87 100644 --- a/source/isaaclab_newton/isaaclab_newton/actuators/newton_actuator_utils.py +++ b/source/isaaclab_newton/isaaclab_newton/actuators/newton_actuator_utils.py @@ -197,6 +197,7 @@ def from_usd( num_envs: int, num_joints: int, device: str, + articulation_prim_path: str | None = None, ) -> NewtonActuatorAdapter: """Create an adapter by parsing ``NewtonActuator`` prims from USD. @@ -228,12 +229,19 @@ def from_usd( num_envs: Number of environments. num_joints: Joints per environment in the articulation. device: Warp device string (e.g. ``"cuda:0"``). + articulation_prim_path: Root prim path of the first + environment's articulation (e.g. ``"/World/Env_0/Robot"``). + When provided, only ``NewtonActuator`` prims under this + subtree are considered — matching the scoped traversal + that Newton's ``ModelBuilder.add_usd`` performs. When + ``None``, the entire stage is scanned (legacy behaviour). Returns: A fully constructed adapter ready for :meth:`finalize`. """ actuators = cls._create_actuators_from_usd( stage, joint_names, num_envs, num_joints, device, + articulation_prim_path=articulation_prim_path, ) return cls(actuators, num_envs, num_joints, dof_offset=0, device=device) @@ -385,6 +393,7 @@ def _create_actuators_from_usd( num_envs: int, num_total_joints: int, device: str, + articulation_prim_path: str | None = None, ) -> list[Actuator]: """Parse ``NewtonActuator`` prims and instantiate standalone actuators. @@ -416,8 +425,13 @@ class (e.g. ``model_path``, ``lookup_positions``) are passed through joint_name_to_idx: dict[str, int] = {name: i for i, name in enumerate(joint_names)} + if articulation_prim_path is not None: + root_prim = stage.GetPrimAtPath(articulation_prim_path) + else: + root_prim = stage.GetPseudoRoot() + parsed_per_joint: dict[int, Any] = {} - for prim in Usd.PrimRange(stage.GetPseudoRoot()): + for prim in Usd.PrimRange(root_prim): parsed = parse_actuator_prim(prim) if parsed is None: continue diff --git a/source/isaaclab_newton/test/assets/test_pd_equivalence.py b/source/isaaclab_newton/test/assets/test_actuators_equivalence.py similarity index 94% rename from source/isaaclab_newton/test/assets/test_pd_equivalence.py rename to source/isaaclab_newton/test/assets/test_actuators_equivalence.py index f9e211891c7e..d959a264699d 100644 --- a/source/isaaclab_newton/test/assets/test_pd_equivalence.py +++ b/source/isaaclab_newton/test/assets/test_actuators_equivalence.py @@ -35,7 +35,7 @@ from isaaclab_newton.physics import NewtonManager as SimulationManager import isaaclab.sim as sim_utils -from isaaclab.actuators import DCMotorCfg, IdealPDActuatorCfg, ImplicitActuatorCfg +from isaaclab.actuators import DCMotorCfg, DelayedPDActuatorCfg, IdealPDActuatorCfg, ImplicitActuatorCfg from isaaclab.sim import SimulationCfg, build_simulation_context # --------------------------------------------------------------------------- @@ -309,6 +309,55 @@ class TestMixedWithImplicitEquivalence(_EquivalenceTestBase): actuators = MIXED_WITH_IMPLICIT_ACTUATORS +# --------------------------------------------------------------------------- +# DelayedPD equivalence: PD with actuator command delay +# --------------------------------------------------------------------------- + +DELAYED_PD_ACTUATORS = { + "legs": DelayedPDActuatorCfg( + joint_names_expr=[".*HAA", ".*HFE", ".*KFE"], + stiffness=40.0, + damping=5.0, + effort_limit=80.0, + min_delay=2, + max_delay=4, + ), +} + + +class TestDelayedPDEquivalence(_EquivalenceTestBase): + """DelayedPDActuator on all 12 joints: Lab vs Newton. + + Verifies that actuator command delays are correctly authored as + ``NewtonActuatorDelayAPI`` and produce matching trajectories. + """ + + __test__ = True + actuators = DELAYED_PD_ACTUATORS + + +class TestDelayedPDAuthoring(unittest.TestCase): + """Verify DelayedPDActuatorCfg is authored with NewtonActuatorDelayAPI.""" + + @classmethod + def setUpClass(cls): + cls.result = _run_authoring_introspection(DELAYED_PD_ACTUATORS) + + def test_has_delay(self): + for a in self.result["actuator_info"]: + self.assertTrue(a["has_delay"], "Delay not found on delayed PD actuator") + + def test_controller_is_pd(self): + for a in self.result["actuator_info"]: + self.assertEqual(a["controller_type"], "ControllerPD") + + def test_trajectories_not_trivial(self): + first = self.result["joint_pos"][0] + last = self.result["joint_pos"][-1] + diff = (last - first).abs().max().item() + self.assertGreater(diff, 0.01, "Joints did not move — test is trivial") + + # --------------------------------------------------------------------------- # Decimation test: CUDA graph capture with actuator decimation loop # --------------------------------------------------------------------------- diff --git a/source/isaaclab_physx/isaaclab_physx/assets/articulation/articulation.py b/source/isaaclab_physx/isaaclab_physx/assets/articulation/articulation.py index dac70836127e..37f0058c8a68 100644 --- a/source/isaaclab_physx/isaaclab_physx/assets/articulation/articulation.py +++ b/source/isaaclab_physx/isaaclab_physx/assets/articulation/articulation.py @@ -3754,12 +3754,16 @@ def _process_actuators_cfg(self): if _HAS_NEWTON_ACTUATORS and _use_newton_actuators: from isaaclab.sim.utils.stage import get_current_stage # noqa: PLC0415 + first_prim = find_first_matching_prim(self.cfg.prim_path) + art_prim_path = str(first_prim.GetPath()) if first_prim is not None else None + adapter = NewtonActuatorAdapter.from_usd( stage=get_current_stage(), joint_names=self.joint_names, num_envs=self.num_instances, num_joints=self.num_joints, device=self.device, + articulation_prim_path=art_prim_path, ) self._physx_actuator_wrapper = PhysxActuatorWrapper() diff --git a/source/isaaclab_physx/test/assets/test_pd_equivalence.py b/source/isaaclab_physx/test/assets/test_actuators_equivalence.py similarity index 93% rename from source/isaaclab_physx/test/assets/test_pd_equivalence.py rename to source/isaaclab_physx/test/assets/test_actuators_equivalence.py index c07bf6c9fdb5..84cbbfaf4986 100644 --- a/source/isaaclab_physx/test/assets/test_pd_equivalence.py +++ b/source/isaaclab_physx/test/assets/test_actuators_equivalence.py @@ -30,7 +30,7 @@ from isaaclab_physx.physics import PhysxCfg import isaaclab.sim as sim_utils -from isaaclab.actuators import DCMotorCfg, IdealPDActuatorCfg, ImplicitActuatorCfg +from isaaclab.actuators import DCMotorCfg, DelayedPDActuatorCfg, IdealPDActuatorCfg, ImplicitActuatorCfg from isaaclab.sim import SimulationCfg, build_simulation_context # --------------------------------------------------------------------------- @@ -66,6 +66,34 @@ ), } +MIXED_ACTUATORS = { + "hips": IdealPDActuatorCfg( + joint_names_expr=[".*HAA"], + stiffness=40.0, + damping=5.0, + effort_limit=80.0, + ), + "knees": DCMotorCfg( + joint_names_expr=[".*HFE", ".*KFE"], + saturation_effort=120.0, + effort_limit=80.0, + velocity_limit=7.5, + stiffness={".*": 40.0}, + damping={".*": 5.0}, + ), +} + +DELAYED_PD_ACTUATORS = { + "legs": DelayedPDActuatorCfg( + joint_names_expr=[".*HAA", ".*HFE", ".*KFE"], + stiffness=40.0, + damping=5.0, + effort_limit=80.0, + min_delay=2, + max_delay=4, + ), +} + MIXED_WITH_IMPLICIT_ACTUATORS = { "hips": ImplicitActuatorCfg( joint_names_expr=[".*HAA"], @@ -232,6 +260,24 @@ class TestDCMotorEquivalence(_EquivalenceTestBase): actuators = DC_MOTOR_ACTUATORS +class TestMixedActuatorEquivalence(_EquivalenceTestBase): + """Mixed actuators (IdealPD on HAA, DCMotor on HFE/KFE): Lab vs Newton (PhysX).""" + + __test__ = True + actuators = MIXED_ACTUATORS + + +class TestDelayedPDEquivalence(_EquivalenceTestBase): + """DelayedPDActuator on all 12 joints: Lab vs Newton (PhysX). + + Verifies that actuator command delays are correctly authored and + produce matching trajectories on the PhysX backend. + """ + + __test__ = True + actuators = DELAYED_PD_ACTUATORS + + class TestMixedWithImplicitEquivalence(_EquivalenceTestBase): """Implicit HAA + IdealPD HFE + DCMotor KFE: Lab vs Newton (PhysX). From 88f325f90085b4e7362de50d271008a86d7790bd Mon Sep 17 00:00:00 2001 From: Julia von Muralt Date: Wed, 29 Apr 2026 16:53:32 +0200 Subject: [PATCH 03/35] rename --- ...quivalence.py => test_newton_actuators.py} | 152 ++++++++++++++++++ ...quivalence.py => test_newton_actuators.py} | 147 +++++++++++++++++ 2 files changed, 299 insertions(+) rename source/isaaclab_newton/test/assets/{test_actuators_equivalence.py => test_newton_actuators.py} (84%) rename source/isaaclab_physx/test/assets/{test_actuators_equivalence.py => test_newton_actuators.py} (79%) diff --git a/source/isaaclab_newton/test/assets/test_actuators_equivalence.py b/source/isaaclab_newton/test/assets/test_newton_actuators.py similarity index 84% rename from source/isaaclab_newton/test/assets/test_actuators_equivalence.py rename to source/isaaclab_newton/test/assets/test_newton_actuators.py index d959a264699d..2e6ae9c33641 100644 --- a/source/isaaclab_newton/test/assets/test_actuators_equivalence.py +++ b/source/isaaclab_newton/test/assets/test_newton_actuators.py @@ -440,6 +440,158 @@ def test_trajectories_not_trivial(self): self.assertGreater(diff, 0.01, "Joints did not move — test is trivial") +# --------------------------------------------------------------------------- +# Partial environment reset: verify per-env reset equivalence +# --------------------------------------------------------------------------- + +RESET_WARMUP_STEPS = 3 +RESET_TOTAL_STEPS = 10 + + +def _run_simulation_with_reset( + actuators: dict, + use_newton_actuators: bool, + *, + dt: float = DT, + newton_cfg: NewtonCfg = NEWTON_CFG, +) -> dict: + """Run ANYmal-C with a mid-simulation reset of env 0 only. + + Steps ``RESET_WARMUP_STEPS``, then resets env 0 to its initial joint state + (zeroing velocity), then steps ``RESET_TOTAL_STEPS - RESET_WARMUP_STEPS`` + more. Returns per-step joint positions and velocities. + + This exercises the actuator state reset path (delay buffers, neural + hidden states, etc.) for a subset of environments. + + Args: + actuators: Actuator config dict overriding ANYmal's defaults. + use_newton_actuators: Use Newton-native actuators when ``True``. + dt: Physics timestep [s]. + newton_cfg: Newton physics configuration. + + Returns: + Dict with ``joint_pos`` and ``joint_vel``, each a list of + ``(NUM_ENVS, num_joints)`` tensors. + """ + sim_cfg = SimulationCfg( + dt=dt, + physics=newton_cfg, + use_newton_actuators=use_newton_actuators, + ) + + with build_simulation_context( + device="cuda:0", + gravity_enabled=True, + add_ground_plane=True, + sim_cfg=sim_cfg, + ) as sim: + sim._app_control_on_stop_handle = None + + for i in range(NUM_ENVS): + sim_utils.create_prim( + f"/World/Env_{i}", "Xform", translation=(i * 3.0, 0, 0) + ) + + art_cfg = ANYMAL_C_CFG.replace( + actuators=actuators, + prim_path="/World/Env_.*/Robot", + ) + articulation = Articulation(art_cfg) + sim.reset() + assert articulation.is_initialized + + init_pos = wp.to_torch(articulation.data.joint_pos).clone() + target_pos = init_pos + TARGET_OFFSET + target_vel = torch.zeros_like(init_pos) + + articulation.set_joint_position_target_index(target=target_pos) + articulation.set_joint_velocity_target_index(target=target_vel) + + recorded_pos, recorded_vel = [], [] + + for step_i in range(RESET_TOTAL_STEPS): + if step_i == RESET_WARMUP_STEPS: + env_ids = torch.tensor([0], device="cuda:0") + articulation.write_joint_position_to_sim_index( + position=init_pos[0:1], env_ids=env_ids, + ) + articulation.write_joint_velocity_to_sim_index( + velocity=torch.zeros_like(init_pos[0:1]), env_ids=env_ids, + ) + articulation.reset(env_ids=[0]) + + articulation.write_data_to_sim() + sim.step() + articulation.update(dt) + + recorded_pos.append(wp.to_torch(articulation.data.joint_pos).clone()) + recorded_vel.append(wp.to_torch(articulation.data.joint_vel).clone()) + + return {"joint_pos": recorded_pos, "joint_vel": recorded_vel} + + +class TestPartialResetEquivalence(unittest.TestCase): + """Per-environment reset with DelayedPD actuators: Lab vs Newton. + + Resets env 0 mid-simulation while env 1 continues uninterrupted. + Uses DelayedPD actuators because they carry internal state (delay + buffers) that must be properly reset per environment. + + Verifies: + - Lab and Newton paths produce matching trajectories after partial reset. + - The two environments diverge after the reset (proving it took effect). + """ + + @classmethod + def setUpClass(cls): + cls.lab_result = _run_simulation_with_reset( + DELAYED_PD_ACTUATORS, use_newton_actuators=False, + ) + cls.newton_result = _run_simulation_with_reset( + DELAYED_PD_ACTUATORS, use_newton_actuators=True, + ) + + def test_joint_positions_match(self): + for step_i, (lab, newton) in enumerate( + zip(self.lab_result["joint_pos"], self.newton_result["joint_pos"]) + ): + torch.testing.assert_close( + lab, + newton, + atol=2e-3, + rtol=1e-3, + msg=f"Positions diverged at step {step_i}", + ) + + def test_joint_velocities_match(self): + for step_i, (lab, newton) in enumerate( + zip(self.lab_result["joint_vel"], self.newton_result["joint_vel"]) + ): + torch.testing.assert_close( + lab, + newton, + atol=0.1, + rtol=1e-2, + msg=f"Velocities diverged at step {step_i}", + ) + + def test_envs_diverge_after_reset(self): + """After resetting env 0, the two envs must have different states.""" + post_reset_pos = self.lab_result["joint_pos"][RESET_WARMUP_STEPS + 1] + diff = (post_reset_pos[0] - post_reset_pos[1]).abs().max().item() + self.assertGreater( + diff, 0.001, + "Env 0 and env 1 are identical after partial reset — reset had no effect", + ) + + def test_trajectories_not_trivial(self): + first = self.lab_result["joint_pos"][0] + last = self.lab_result["joint_pos"][-1] + diff = (last - first).abs().max().item() + self.assertGreater(diff, 0.01, "Joints did not move — test is trivial") + + # --------------------------------------------------------------------------- # RemotizedPD actuator: PD + delay + position-based clamping lookup table # --------------------------------------------------------------------------- diff --git a/source/isaaclab_physx/test/assets/test_actuators_equivalence.py b/source/isaaclab_physx/test/assets/test_newton_actuators.py similarity index 79% rename from source/isaaclab_physx/test/assets/test_actuators_equivalence.py rename to source/isaaclab_physx/test/assets/test_newton_actuators.py index 84cbbfaf4986..08afb6383222 100644 --- a/source/isaaclab_physx/test/assets/test_actuators_equivalence.py +++ b/source/isaaclab_physx/test/assets/test_newton_actuators.py @@ -289,6 +289,153 @@ class TestMixedWithImplicitEquivalence(_EquivalenceTestBase): actuators = MIXED_WITH_IMPLICIT_ACTUATORS +# --------------------------------------------------------------------------- +# Partial environment reset: verify per-env reset equivalence +# --------------------------------------------------------------------------- + +RESET_WARMUP_STEPS = 3 +RESET_TOTAL_STEPS = 10 + + +def _run_simulation_with_reset( + actuators: dict, + use_newton_actuators: bool, +) -> dict: + """Run ANYmal-C on PhysX with a mid-simulation reset of env 0 only. + + Steps ``RESET_WARMUP_STEPS``, then resets env 0 to its initial joint state + (zeroing velocity), then steps ``RESET_TOTAL_STEPS - RESET_WARMUP_STEPS`` + more. Returns per-step joint positions and velocities. + + This exercises the actuator state reset path (delay buffers, neural + hidden states, etc.) for a subset of environments. + + Args: + actuators: Actuator config dict overriding ANYmal's defaults. + use_newton_actuators: Use Newton-native actuators when ``True``. + + Returns: + Dict with ``joint_pos`` and ``joint_vel``, each a list of + ``(NUM_ENVS, num_joints)`` tensors. + """ + sim_cfg = SimulationCfg( + dt=DT, + physics=PhysxCfg(), + use_newton_actuators=use_newton_actuators, + ) + + with build_simulation_context( + device="cuda:0", + gravity_enabled=True, + add_ground_plane=True, + sim_cfg=sim_cfg, + ) as sim: + sim._app_control_on_stop_handle = None + + for i in range(NUM_ENVS): + sim_utils.create_prim( + f"/World/Env_{i}", "Xform", translation=(i * 3.0, 0, 0) + ) + + art_cfg = ANYMAL_C_CFG.replace( + actuators=actuators, + prim_path="/World/Env_.*/Robot", + ) + articulation = Articulation(art_cfg) + sim.reset() + assert articulation.is_initialized + + init_pos = wp.to_torch(articulation.data.joint_pos).clone() + target_pos = init_pos + TARGET_OFFSET + target_vel = torch.zeros_like(init_pos) + + articulation.set_joint_position_target_index(target=target_pos) + articulation.set_joint_velocity_target_index(target=target_vel) + + recorded_pos, recorded_vel = [], [] + + for step_i in range(RESET_TOTAL_STEPS): + if step_i == RESET_WARMUP_STEPS: + env_ids = torch.tensor([0], device="cuda:0") + articulation.write_joint_position_to_sim_index( + position=init_pos[0:1], env_ids=env_ids, + ) + articulation.write_joint_velocity_to_sim_index( + velocity=torch.zeros_like(init_pos[0:1]), env_ids=env_ids, + ) + articulation.reset(env_ids=[0]) + + articulation.write_data_to_sim() + sim.step() + articulation.update(DT) + + recorded_pos.append(wp.to_torch(articulation.data.joint_pos).clone()) + recorded_vel.append(wp.to_torch(articulation.data.joint_vel).clone()) + + return {"joint_pos": recorded_pos, "joint_vel": recorded_vel} + + +class TestPartialResetEquivalence(unittest.TestCase): + """Per-environment reset with DelayedPD actuators: Lab vs Newton (PhysX). + + Resets env 0 mid-simulation while env 1 continues uninterrupted. + Uses DelayedPD actuators because they carry internal state (delay + buffers) that must be properly reset per environment. + + Verifies: + - Lab and Newton paths produce matching trajectories after partial reset. + - The two environments diverge after the reset (proving it took effect). + """ + + @classmethod + def setUpClass(cls): + cls.lab_result = _run_simulation_with_reset( + DELAYED_PD_ACTUATORS, use_newton_actuators=False, + ) + cls.newton_result = _run_simulation_with_reset( + DELAYED_PD_ACTUATORS, use_newton_actuators=True, + ) + + def test_joint_positions_match(self): + for step_i, (lab, newton) in enumerate( + zip(self.lab_result["joint_pos"], self.newton_result["joint_pos"]) + ): + torch.testing.assert_close( + lab, + newton, + atol=2e-3, + rtol=1e-3, + msg=f"Positions diverged at step {step_i}", + ) + + def test_joint_velocities_match(self): + for step_i, (lab, newton) in enumerate( + zip(self.lab_result["joint_vel"], self.newton_result["joint_vel"]) + ): + torch.testing.assert_close( + lab, + newton, + atol=0.1, + rtol=1e-2, + msg=f"Velocities diverged at step {step_i}", + ) + + def test_envs_diverge_after_reset(self): + """After resetting env 0, the two envs must have different states.""" + post_reset_pos = self.lab_result["joint_pos"][RESET_WARMUP_STEPS + 1] + diff = (post_reset_pos[0] - post_reset_pos[1]).abs().max().item() + self.assertGreater( + diff, 0.001, + "Env 0 and env 1 are identical after partial reset — reset had no effect", + ) + + def test_trajectories_not_trivial(self): + first = self.lab_result["joint_pos"][0] + last = self.lab_result["joint_pos"][-1] + diff = (last - first).abs().max().item() + self.assertGreater(diff, 0.01, "Joints did not move — test is trivial") + + # --------------------------------------------------------------------------- # RemotizedPD authoring: PD + delay + position-based clamping lookup table # --------------------------------------------------------------------------- From 5632e2e516e7ed603dc109f50e64fb0e923030ed Mon Sep 17 00:00:00 2001 From: Julia von Muralt Date: Wed, 29 Apr 2026 21:41:15 +0200 Subject: [PATCH 04/35] update version --- source/isaaclab/setup.py | 4 ++-- source/isaaclab_newton/setup.py | 6 +++--- ..._newton_actuators.py => test_newton_actuators_newton.py} | 0 ...t_newton_actuators.py => test_newton_actuators_physx.py} | 0 4 files changed, 5 insertions(+), 5 deletions(-) rename source/isaaclab_newton/test/assets/{test_newton_actuators.py => test_newton_actuators_newton.py} (100%) rename source/isaaclab_physx/test/assets/{test_newton_actuators.py => test_newton_actuators_physx.py} (100%) diff --git a/source/isaaclab/setup.py b/source/isaaclab/setup.py index 7d14504f48a3..3d3d327cc9c0 100644 --- a/source/isaaclab/setup.py +++ b/source/isaaclab/setup.py @@ -30,8 +30,8 @@ # procedural-generation "trimesh", "pyglet>=2.1.6,<3", - "mujoco==3.6.0", - "mujoco-warp==3.6.0", + "mujoco==3.7.0", + "mujoco-warp==3.7.0.1", # image processing "transformers==4.57.6", "einops", # needed for transformers, doesn't always auto-install diff --git a/source/isaaclab_newton/setup.py b/source/isaaclab_newton/setup.py index 14e533b69d13..f7a36b177d37 100644 --- a/source/isaaclab_newton/setup.py +++ b/source/isaaclab_newton/setup.py @@ -38,10 +38,10 @@ def run(self): EXTRAS_REQUIRE = { "all": [ "prettytable==3.3.0", - "mujoco==3.6.0", - "mujoco-warp==3.6.0", + "mujoco==3.7.0", + "mujoco-warp==3.7.0.1", "PyOpenGL-accelerate==3.1.10", - "newton @ git+https://github.com/newton-physics/newton.git@b859659c7f3df8b688a5989a8f1719b8dc15aa86", + "newton @ git+https://github.com/newton-physics/newton.git@dda12f8259221e4068845e38835dcff517622d92", ], } diff --git a/source/isaaclab_newton/test/assets/test_newton_actuators.py b/source/isaaclab_newton/test/assets/test_newton_actuators_newton.py similarity index 100% rename from source/isaaclab_newton/test/assets/test_newton_actuators.py rename to source/isaaclab_newton/test/assets/test_newton_actuators_newton.py diff --git a/source/isaaclab_physx/test/assets/test_newton_actuators.py b/source/isaaclab_physx/test/assets/test_newton_actuators_physx.py similarity index 100% rename from source/isaaclab_physx/test/assets/test_newton_actuators.py rename to source/isaaclab_physx/test/assets/test_newton_actuators_physx.py From 4751440c8dc1caa18dba1dd5104cc8e8bcc65880 Mon Sep 17 00:00:00 2001 From: Julia von Muralt Date: Thu, 30 Apr 2026 11:38:33 +0200 Subject: [PATCH 05/35] fix --- .../isaaclab/assets/articulation/base_articulation.py | 4 +++- .../isaaclab_newton/actuators/newton_actuator_utils.py | 6 +++++- 2 files changed, 8 insertions(+), 2 deletions(-) diff --git a/source/isaaclab/isaaclab/assets/articulation/base_articulation.py b/source/isaaclab/isaaclab/assets/articulation/base_articulation.py index 12cba9b23e78..3e6a7fe459c4 100644 --- a/source/isaaclab/isaaclab/assets/articulation/base_articulation.py +++ b/source/isaaclab/isaaclab/assets/articulation/base_articulation.py @@ -130,7 +130,9 @@ def _author_newton_actuator_prims(self) -> None: from isaaclab.sim.utils.queries import find_first_matching_prim # noqa: PLC0415 - first_prim = find_first_matching_prim(self.cfg.prim_path) + spawn_path = getattr(self.cfg.spawn, "spawn_path", None) if self.cfg.spawn is not None else None + search_path = spawn_path if spawn_path is not None else self.cfg.prim_path + first_prim = find_first_matching_prim(search_path) if first_prim is None: return diff --git a/source/isaaclab_newton/isaaclab_newton/actuators/newton_actuator_utils.py b/source/isaaclab_newton/isaaclab_newton/actuators/newton_actuator_utils.py index eef438c60a87..0bf7b12b0ba1 100644 --- a/source/isaaclab_newton/isaaclab_newton/actuators/newton_actuator_utils.py +++ b/source/isaaclab_newton/isaaclab_newton/actuators/newton_actuator_utils.py @@ -304,7 +304,11 @@ def reset(self, env_ids: Sequence[int] | slice | None = None) -> None: mask = None else: mask = wp.zeros(self._num_envs, dtype=wp.bool, device=self._device) - idx = wp.array(list(env_ids), dtype=wp.int32, device=self._device) + import torch # noqa: PLC0415 + if isinstance(env_ids, torch.Tensor): + idx = wp.from_torch(env_ids.contiguous().to(torch.int32), dtype=wp.int32) + else: + idx = wp.array(list(env_ids), dtype=wp.int32, device=self._device) wp.launch(_set_mask_kernel, dim=len(idx), inputs=[mask, idx], device=self._device) for sa, sb in zip(self._states_a, self._states_b): From b2718e7a643f30b85537aed4076169ed19900df5 Mon Sep 17 00:00:00 2001 From: Julia von Muralt Date: Thu, 30 Apr 2026 14:45:51 +0200 Subject: [PATCH 06/35] fix --- .../actuators/newton_actuator_utils.py | 24 +++++++++++++++++-- .../assets/articulation/articulation.py | 3 +++ .../assets/articulation/articulation.py | 10 ++++---- 3 files changed, 30 insertions(+), 7 deletions(-) diff --git a/source/isaaclab_newton/isaaclab_newton/actuators/newton_actuator_utils.py b/source/isaaclab_newton/isaaclab_newton/actuators/newton_actuator_utils.py index 0bf7b12b0ba1..e04f73034732 100644 --- a/source/isaaclab_newton/isaaclab_newton/actuators/newton_actuator_utils.py +++ b/source/isaaclab_newton/isaaclab_newton/actuators/newton_actuator_utils.py @@ -52,6 +52,7 @@ from __future__ import annotations +import logging from collections.abc import Sequence from dataclasses import dataclass from typing import Any @@ -60,6 +61,8 @@ import torch import warp as wp +logger = logging.getLogger(__name__) + from newton.actuators import ( Actuator, Clamping, @@ -306,7 +309,7 @@ def reset(self, env_ids: Sequence[int] | slice | None = None) -> None: mask = wp.zeros(self._num_envs, dtype=wp.bool, device=self._device) import torch # noqa: PLC0415 if isinstance(env_ids, torch.Tensor): - idx = wp.from_torch(env_ids.contiguous().to(torch.int32), dtype=wp.int32) + idx = wp.from_torch(env_ids.to(device=self._device).contiguous().to(torch.int32), dtype=wp.int32) else: idx = wp.array(list(env_ids), dtype=wp.int32, device=self._device) wp.launch(_set_mask_kernel, dim=len(idx), inputs=[mask, idx], device=self._device) @@ -592,9 +595,26 @@ def author_newton_actuator_prims( from isaaclab.actuators import DCMotorCfg, DelayedPDActuatorCfg # noqa: PLC0415 from isaaclab.actuators.actuator_net_cfg import ActuatorNetLSTMCfg, ActuatorNetMLPCfg # noqa: PLC0415 - from isaaclab.actuators.actuator_pd_cfg import RemotizedPDActuatorCfg # noqa: PLC0415 + from isaaclab.actuators.actuator_pd_cfg import IdealPDActuatorCfg, RemotizedPDActuatorCfg # noqa: PLC0415 + + _SUPPORTED_CFG_TYPES = ( + IdealPDActuatorCfg, + DCMotorCfg, + DelayedPDActuatorCfg, + RemotizedPDActuatorCfg, + ActuatorNetMLPCfg, + ActuatorNetLSTMCfg, + ) for group_name, cfg, joint_names in cfg_entries: + if not isinstance(cfg, _SUPPORTED_CFG_TYPES): + logger.warning( + "Actuator group '%s' uses config type '%s' which is not supported by Newton-native" + " actuator authoring. The group will be skipped.", + group_name, + type(cfg).__name__, + ) + continue stiffness_map = resolve(getattr(cfg, "stiffness", None), joint_names) damping_map = resolve(getattr(cfg, "damping", None), joint_names) effort_map = resolve(getattr(cfg, "effort_limit", None), joint_names) diff --git a/source/isaaclab_newton/isaaclab_newton/assets/articulation/articulation.py b/source/isaaclab_newton/isaaclab_newton/assets/articulation/articulation.py index f9da44f2461f..6dcf552d8ea6 100644 --- a/source/isaaclab_newton/isaaclab_newton/assets/articulation/articulation.py +++ b/source/isaaclab_newton/isaaclab_newton/assets/articulation/articulation.py @@ -3461,6 +3461,9 @@ def _process_actuators_cfg(self): model = SimulationManager.get_model() + if not model.actuators: + return + dof_layout = self._root_view.frequency_layouts[NewtonModel.AttributeFrequency.JOINT_DOF] if dof_layout.slice is not None: dof_offset = dof_layout.slice.start diff --git a/source/isaaclab_physx/isaaclab_physx/assets/articulation/articulation.py b/source/isaaclab_physx/isaaclab_physx/assets/articulation/articulation.py index 9d4303ff3d51..5af2f8dd9c6f 100644 --- a/source/isaaclab_physx/isaaclab_physx/assets/articulation/articulation.py +++ b/source/isaaclab_physx/isaaclab_physx/assets/articulation/articulation.py @@ -4030,11 +4030,11 @@ def _apply_actuator_model_newton(self): # so Newton writes effort directly into the staging buffer. newton_adapter = self.actuators.get("newton") w = self._physx_actuator_wrapper - w.joint_q = self._data.joint_pos.reshape(-1) - w.joint_qd = self._data.joint_vel.reshape(-1) - w.joint_target_pos = self._data.joint_pos_target.reshape(-1) - w.joint_target_vel = self._data.joint_vel_target.reshape(-1) - w.joint_act = self._data.joint_effort_target.reshape(-1) + w.joint_q = self._data.joint_pos.warp.reshape(-1) + w.joint_qd = self._data.joint_vel.warp.reshape(-1) + w.joint_target_pos = self._data.joint_pos_target.warp.reshape(-1) + w.joint_target_vel = self._data.joint_vel_target.warp.reshape(-1) + w.joint_act = self._data.joint_effort_target.warp.reshape(-1) w.joint_f = self._joint_effort_target_sim.reshape(-1) newton_adapter.step(w, w, SimulationManager.get_physics_dt()) From bf56a74c75afc8d36c07df802074ac1c4331d516 Mon Sep 17 00:00:00 2001 From: Julia von Muralt Date: Mon, 4 May 2026 10:09:04 +0200 Subject: [PATCH 07/35] add has newton actuators check --- .../assets/articulation/articulation.py | 18 +++++++++++++++--- source/isaaclab_newton/setup.py | 6 +++--- 2 files changed, 18 insertions(+), 6 deletions(-) diff --git a/source/isaaclab_newton/isaaclab_newton/assets/articulation/articulation.py b/source/isaaclab_newton/isaaclab_newton/assets/articulation/articulation.py index 6dcf552d8ea6..4787fd4edf5e 100644 --- a/source/isaaclab_newton/isaaclab_newton/assets/articulation/articulation.py +++ b/source/isaaclab_newton/isaaclab_newton/assets/articulation/articulation.py @@ -26,7 +26,13 @@ from isaaclab.actuators import ActuatorBase, ActuatorBaseCfg, ImplicitActuator from isaaclab.assets.articulation.base_articulation import BaseArticulation -from isaaclab_newton.actuators.newton_actuator_utils import NewtonActuatorAdapter +try: + from isaaclab_newton.actuators.newton_actuator_utils import NewtonActuatorAdapter + + _HAS_NEWTON_ACTUATORS = True +except ImportError: + _HAS_NEWTON_ACTUATORS = False + from isaaclab.physics import PhysicsEvent from isaaclab.sim.utils.queries import find_first_matching_prim, get_all_matching_child_prims from isaaclab.utils.string import resolve_matching_names, resolve_matching_names_values @@ -3456,7 +3462,13 @@ def _process_actuators_cfg(self): _use_newton_actuators = getattr(self._sim_cfg, "use_newton_actuators", False) - if _use_newton_actuators: + if _use_newton_actuators and not _HAS_NEWTON_ACTUATORS: + logger.warning( + "use_newton_actuators is enabled but 'newton.actuators' is not available. " + "Newton-native actuators will be disabled. Upgrade Newton to >= 1.2.0rc1." + ) + + if _use_newton_actuators and _HAS_NEWTON_ACTUATORS: from newton import Model as NewtonModel # noqa: PLC0415 model = SimulationManager.get_model() @@ -3511,7 +3523,7 @@ def _process_actuators_cfg(self): if self.cfg.actuator_value_resolution_debug_print: t = PrettyTable(["Group", "Property", "Name", "ID", "USD Value", "ActutatorCfg Value", "Applied"]) for actuator_group, actuator in self.actuators.items(): - if isinstance(actuator, NewtonActuatorAdapter): + if _HAS_NEWTON_ACTUATORS and isinstance(actuator, NewtonActuatorAdapter): continue group_count = 0 for property, resolution_details in actuator.joint_property_resolution_table.items(): diff --git a/source/isaaclab_newton/setup.py b/source/isaaclab_newton/setup.py index f7a36b177d37..0ae1d8275ca4 100644 --- a/source/isaaclab_newton/setup.py +++ b/source/isaaclab_newton/setup.py @@ -38,10 +38,10 @@ def run(self): EXTRAS_REQUIRE = { "all": [ "prettytable==3.3.0", - "mujoco==3.7.0", - "mujoco-warp==3.7.0.1", + "mujoco~=3.8.0", + "mujoco-warp>=3.8.0.1,~=3.8.0", "PyOpenGL-accelerate==3.1.10", - "newton @ git+https://github.com/newton-physics/newton.git@dda12f8259221e4068845e38835dcff517622d92", + "newton @ git+https://github.com/newton-physics/newton.git@d89a76bd94542952301e570a3cc36156e165052e", ], } From c1915c6623607fb24490cd8da37debc78f4413ca Mon Sep 17 00:00:00 2001 From: Julia von Muralt Date: Mon, 4 May 2026 14:20:58 +0200 Subject: [PATCH 08/35] try force install --- docker/Dockerfile.base | 12 ++++++++++++ source/isaaclab/setup.py | 4 ++-- source/isaaclab_newton/setup.py | 2 +- source/isaaclab_visualizers/setup.py | 6 +++--- 4 files changed, 18 insertions(+), 6 deletions(-) diff --git a/docker/Dockerfile.base b/docker/Dockerfile.base index c10c036ce0a6..44088d689e77 100644 --- a/docker/Dockerfile.base +++ b/docker/Dockerfile.base @@ -117,6 +117,18 @@ RUN --mount=type=cache,target=${DOCKER_USER_HOME}/.cache/pip \ RUN --mount=type=cache,target=${DOCKER_USER_HOME}/.cache/pip \ ${ISAACLAB_PATH}/isaaclab.sh --install +# HACK: Force-install Newton to match the version required by isaaclab_newton. +# The Docker base image ships an older pre-built copy whose version string +# collides with the new one (both report 1.2.0.dev0), so pip skips the +# upgrade during `isaaclab.sh --install`. --no-deps avoids pulling +# Newton's warp-lang>=1.13 dep (we keep the bundled Warp 1.12 for Isaac Sim +# compatibility). +RUN --mount=type=cache,target=${DOCKER_USER_HOME}/.cache/pip \ + ${ISAACLAB_PATH}/isaaclab.sh -p -m pip install --force-reinstall --no-deps \ + "newton @ git+https://github.com/jvonmuralt/newton.git@ignore_path" && \ + ${ISAACLAB_PATH}/isaaclab.sh -p -m pip install \ + "mujoco~=3.8.0" "mujoco-warp>=3.8.0.1,~=3.8.0" + # HACK: Remove install of quadprog dependency RUN ${ISAACLAB_PATH}/isaaclab.sh -p -m pip uninstall -y quadprog diff --git a/source/isaaclab/setup.py b/source/isaaclab/setup.py index 3d3d327cc9c0..0511db8ec439 100644 --- a/source/isaaclab/setup.py +++ b/source/isaaclab/setup.py @@ -30,8 +30,8 @@ # procedural-generation "trimesh", "pyglet>=2.1.6,<3", - "mujoco==3.7.0", - "mujoco-warp==3.7.0.1", + "mujoco~=3.8.0", + "mujoco-warp>=3.8.0.1,~=3.8.0", # image processing "transformers==4.57.6", "einops", # needed for transformers, doesn't always auto-install diff --git a/source/isaaclab_newton/setup.py b/source/isaaclab_newton/setup.py index 0ae1d8275ca4..88d9c07b5171 100644 --- a/source/isaaclab_newton/setup.py +++ b/source/isaaclab_newton/setup.py @@ -41,7 +41,7 @@ def run(self): "mujoco~=3.8.0", "mujoco-warp>=3.8.0.1,~=3.8.0", "PyOpenGL-accelerate==3.1.10", - "newton @ git+https://github.com/newton-physics/newton.git@d89a76bd94542952301e570a3cc36156e165052e", + "newton @ git+https://github.com/jvonmuralt/newton.git@ignore_path", ], } diff --git a/source/isaaclab_visualizers/setup.py b/source/isaaclab_visualizers/setup.py index fc120619787b..4aef33c223e7 100644 --- a/source/isaaclab_visualizers/setup.py +++ b/source/isaaclab_visualizers/setup.py @@ -17,16 +17,16 @@ "kit": [], "newton": [ "warp-lang", - "newton @ git+https://github.com/newton-physics/newton.git@a27277ed49d6f307b8a1e4c394be7e1d14965a62", + "newton @ git+https://github.com/jvonmuralt/newton.git@ignore_path", "PyOpenGL-accelerate", "imgui-bundle>=1.92.5", ], "rerun": [ - "newton @ git+https://github.com/newton-physics/newton.git@a27277ed49d6f307b8a1e4c394be7e1d14965a62", + "newton @ git+https://github.com/jvonmuralt/newton.git@ignore_path", "rerun-sdk>=0.29.0", ], "viser": [ - "newton @ git+https://github.com/newton-physics/newton.git@a27277ed49d6f307b8a1e4c394be7e1d14965a62", + "newton @ git+https://github.com/jvonmuralt/newton.git@ignore_path", "viser>=1.0.16", ], } From 879a71566f36affab8835886d9196f88e86910b1 Mon Sep 17 00:00:00 2001 From: Julia von Muralt Date: Mon, 4 May 2026 14:56:43 +0200 Subject: [PATCH 09/35] try --- docker/Dockerfile.base | 23 ++++++++++++----------- source/isaaclab_visualizers/setup.py | 3 --- 2 files changed, 12 insertions(+), 14 deletions(-) diff --git a/docker/Dockerfile.base b/docker/Dockerfile.base index 44088d689e77..536c37860274 100644 --- a/docker/Dockerfile.base +++ b/docker/Dockerfile.base @@ -112,23 +112,24 @@ RUN --mount=type=cache,target=${DOCKER_USER_HOME}/.cache/pip \ ${ISAACLAB_PATH}/isaaclab.sh -p -m pip install --no-build-isolation nlopt==2.6.2; \ fi -# installing Isaac Lab dependencies -# use pip caching to avoid reinstalling large packages -RUN --mount=type=cache,target=${DOCKER_USER_HOME}/.cache/pip \ - ${ISAACLAB_PATH}/isaaclab.sh --install - -# HACK: Force-install Newton to match the version required by isaaclab_newton. -# The Docker base image ships an older pre-built copy whose version string -# collides with the new one (both report 1.2.0.dev0), so pip skips the -# upgrade during `isaaclab.sh --install`. --no-deps avoids pulling -# Newton's warp-lang>=1.13 dep (we keep the bundled Warp 1.12 for Isaac Sim -# compatibility). +# HACK: Pre-install Newton before `isaaclab.sh --install`. +# Newton's own setup.py requires warp-lang>=1.13 but isaaclab pins +# warp-lang==1.12 (bundled with Isaac Sim). Installing Newton via pip's +# dependency resolver would fail, so we force-install it with --no-deps +# first. The Docker base image also ships an older pre-built Newton whose +# version string collides (both report 1.2.0.dev0), so --force-reinstall +# ensures the correct commit is used. RUN --mount=type=cache,target=${DOCKER_USER_HOME}/.cache/pip \ ${ISAACLAB_PATH}/isaaclab.sh -p -m pip install --force-reinstall --no-deps \ "newton @ git+https://github.com/jvonmuralt/newton.git@ignore_path" && \ ${ISAACLAB_PATH}/isaaclab.sh -p -m pip install \ "mujoco~=3.8.0" "mujoco-warp>=3.8.0.1,~=3.8.0" +# installing Isaac Lab dependencies +# use pip caching to avoid reinstalling large packages +RUN --mount=type=cache,target=${DOCKER_USER_HOME}/.cache/pip \ + ${ISAACLAB_PATH}/isaaclab.sh --install + # HACK: Remove install of quadprog dependency RUN ${ISAACLAB_PATH}/isaaclab.sh -p -m pip uninstall -y quadprog diff --git a/source/isaaclab_visualizers/setup.py b/source/isaaclab_visualizers/setup.py index 4aef33c223e7..0ad6b9fec0fa 100644 --- a/source/isaaclab_visualizers/setup.py +++ b/source/isaaclab_visualizers/setup.py @@ -17,16 +17,13 @@ "kit": [], "newton": [ "warp-lang", - "newton @ git+https://github.com/jvonmuralt/newton.git@ignore_path", "PyOpenGL-accelerate", "imgui-bundle>=1.92.5", ], "rerun": [ - "newton @ git+https://github.com/jvonmuralt/newton.git@ignore_path", "rerun-sdk>=0.29.0", ], "viser": [ - "newton @ git+https://github.com/jvonmuralt/newton.git@ignore_path", "viser>=1.0.16", ], } From d12e610828393a3b5cb577e7dcca5c52a0ed0d6c Mon Sep 17 00:00:00 2001 From: Julia von Muralt Date: Wed, 6 May 2026 09:07:44 +0200 Subject: [PATCH 10/35] fix and improve tests --- .../assets/articulation/articulation.py | 55 ++++++++---- .../assets/test_newton_actuators_newton.py | 84 +++++++++++++++++-- .../assets/articulation/articulation.py | 62 +++++++++----- .../assets/test_newton_actuators_physx.py | 32 +++++++ 4 files changed, 184 insertions(+), 49 deletions(-) diff --git a/source/isaaclab_newton/isaaclab_newton/assets/articulation/articulation.py b/source/isaaclab_newton/isaaclab_newton/assets/articulation/articulation.py index 4787fd4edf5e..d2a6ea590be1 100644 --- a/source/isaaclab_newton/isaaclab_newton/assets/articulation/articulation.py +++ b/source/isaaclab_newton/isaaclab_newton/assets/articulation/articulation.py @@ -3494,7 +3494,6 @@ def _process_actuators_cfg(self): self.write_joint_damping_to_sim_index(damping=0.0, joint_ids=adapter.joint_indices) self._has_newton_actuators = True - # Also create any implicit actuators defined in the configs for actuator_name, actuator_cfg in self.cfg.actuators.items(): cls_type = actuator_cfg.class_type is_implicit = ( @@ -3502,9 +3501,10 @@ def _process_actuators_cfg(self): if isinstance(cls_type, str) else issubclass(cls_type, ImplicitActuator) ) - if not is_implicit: - continue - self._create_lab_actuator(actuator_name, actuator_cfg) + if is_implicit: + self._create_lab_actuator(actuator_name, actuator_cfg) + else: + self._create_lab_actuator(actuator_name, actuator_cfg, properties_only=True) return @@ -3535,8 +3535,19 @@ def _process_actuators_cfg(self): group_count += 1 logger.warning(f"\nActuatorCfg-USD Value Discrepancy Resolution (matching values are skipped): \n{t}") - def _create_lab_actuator(self, actuator_name: str, actuator_cfg: ActuatorBaseCfg) -> None: - """Instantiate a single Lab actuator from its config and write properties to sim.""" + def _create_lab_actuator( + self, actuator_name: str, actuator_cfg: ActuatorBaseCfg, *, properties_only: bool = False, + ) -> None: + """Instantiate a single Lab actuator from its config and write properties to sim. + + Args: + actuator_name: Name for the actuator group. + actuator_cfg: Configuration for the actuator. + properties_only: When ``True``, only write physical joint properties + (armature, limits, friction) without registering the actuator or + writing stiffness/damping. Used for explicit joints managed by + Newton actuators. + """ joint_ids, joint_names = self.find_joints(actuator_cfg.joint_names_expr) if len(joint_names) == 0: raise ValueError( @@ -3561,16 +3572,8 @@ def _create_lab_actuator(self, actuator_name: str, actuator_cfg: ActuatorBaseCfg effort_limit=wp.to_torch(self._data.joint_effort_limits)[:, joint_ids].clone(), velocity_limit=wp.to_torch(self._data.joint_vel_limits)[:, joint_ids], ) - self.actuators[actuator_name] = actuator - - if isinstance(actuator, ImplicitActuator): - self._has_implicit_actuators = True - self.write_joint_stiffness_to_sim_index(stiffness=actuator.stiffness, joint_ids=actuator.joint_indices) - self.write_joint_damping_to_sim_index(damping=actuator.damping, joint_ids=actuator.joint_indices) - else: - self.write_joint_stiffness_to_sim_index(stiffness=0.0, joint_ids=actuator.joint_indices) - self.write_joint_damping_to_sim_index(damping=0.0, joint_ids=actuator.joint_indices) + # Write physical joint properties (armature, limits, friction) — always needed. self.write_joint_effort_limit_to_sim_index( limits=actuator.effort_limit_sim, joint_ids=actuator.joint_indices, ) @@ -3582,7 +3585,23 @@ def _create_lab_actuator(self, actuator_name: str, actuator_cfg: ActuatorBaseCfg joint_friction_coeff=actuator.friction, joint_ids=actuator.joint_indices, ) - # Store the configured values from the actuator model + if properties_only: + return + + self.actuators[actuator_name] = actuator + + if isinstance(actuator, ImplicitActuator): + self._has_implicit_actuators = True + self.write_joint_stiffness_to_sim_index(stiffness=actuator.stiffness, joint_ids=actuator.joint_indices) + self.write_joint_damping_to_sim_index(damping=actuator.damping, joint_ids=actuator.joint_indices) + else: + self.write_joint_stiffness_to_sim_index(stiffness=0.0, joint_ids=actuator.joint_indices) + self.write_joint_damping_to_sim_index(damping=0.0, joint_ids=actuator.joint_indices) + + # Store the actuator-configured values in Lab-internal buffers. + # These are separate from the sim-bound model arrays so that + # write_joint_stiffness_to_sim_index(0.0) for explicit actuators + # is not overwritten (the solver must see ke=0 for explicit joints). j_ids = actuator.joint_indices if j_ids == slice(None): j_ids = self._ALL_JOINT_INDICES @@ -3590,14 +3609,14 @@ def _create_lab_actuator(self, actuator_name: str, actuator_cfg: ActuatorBaseCfg shared_kernels.write_2d_data_to_buffer_with_indices, dim=(self.num_instances, j_ids.shape[0]), inputs=[actuator.stiffness, self._ALL_INDICES, j_ids], - outputs=[self.data._sim_bind_joint_stiffness_sim], + outputs=[self.data._actuator_stiffness], device=self.device, ) wp.launch( shared_kernels.write_2d_data_to_buffer_with_indices, dim=(self.num_instances, j_ids.shape[0]), inputs=[actuator.damping, self._ALL_INDICES, j_ids], - outputs=[self.data._sim_bind_joint_damping_sim], + outputs=[self.data._actuator_damping], device=self.device, ) wp.launch( diff --git a/source/isaaclab_newton/test/assets/test_newton_actuators_newton.py b/source/isaaclab_newton/test/assets/test_newton_actuators_newton.py index 2e6ae9c33641..4e028288d654 100644 --- a/source/isaaclab_newton/test/assets/test_newton_actuators_newton.py +++ b/source/isaaclab_newton/test/assets/test_newton_actuators_newton.py @@ -382,18 +382,27 @@ def test_trajectories_not_trivial(self): ) -class TestDecimation(unittest.TestCase): - """Lab vs Newton with decimation=2 and CUDA graph capture. +class _DecimationTestBase(unittest.TestCase): + """Base for decimation tests with CUDA graph capture. Policy runs at 50 Hz, actuators at 100 Hz, physics at 200 Hz. The Newton path captures the full decimation loop as a CUDA graph; the Lab path runs an explicit per-substep loop. + + Subclasses set ``actuators`` to the config under test. """ + __test__ = False + actuators: dict = {} + pos_atol: float = 2e-3 + pos_rtol: float = 1e-3 + vel_atol: float = 0.1 + vel_rtol: float = 1e-2 + @classmethod def setUpClass(cls): cls.lab_result = _run_simulation( - DC_MOTOR_ACTUATORS, + cls.actuators, use_newton_actuators=False, dt=DT_DEC, newton_cfg=NEWTON_CFG_DEC, @@ -401,7 +410,7 @@ def setUpClass(cls): decimation=DECIMATION, ) cls.newton_result = _run_simulation( - DC_MOTOR_ACTUATORS, + cls.actuators, use_newton_actuators=True, dt=DT_DEC, newton_cfg=NEWTON_CFG_DEC, @@ -416,8 +425,8 @@ def test_joint_positions_match(self): torch.testing.assert_close( lab, newton, - atol=2e-3, - rtol=1e-3, + atol=self.pos_atol, + rtol=self.pos_rtol, msg=f"Positions diverged at policy step {step_i}", ) @@ -428,8 +437,8 @@ def test_joint_velocities_match(self): torch.testing.assert_close( lab, newton, - atol=0.1, - rtol=1e-2, + atol=self.vel_atol, + rtol=self.vel_rtol, msg=f"Velocities diverged at policy step {step_i}", ) @@ -440,6 +449,65 @@ def test_trajectories_not_trivial(self): self.assertGreater(diff, 0.01, "Joints did not move — test is trivial") +class TestDecimationDCMotor(_DecimationTestBase): + """DCMotor with decimation=2 and CUDA graph capture.""" + + __test__ = True + actuators = DC_MOTOR_ACTUATORS + + +class TestDecimationIdealPD(_DecimationTestBase): + """IdealPD with decimation=2 and CUDA graph capture.""" + + __test__ = True + actuators = IDEAL_PD_ACTUATORS + + +class TestDecimationDelayedPD(_DecimationTestBase): + """DelayedPD with decimation=2 and CUDA graph capture. + + Delay buffers must be correctly stepped inside the captured graph. + """ + + __test__ = True + actuators = DELAYED_PD_ACTUATORS + + +class TestDecimationMixed(_DecimationTestBase): + """Mixed actuators (IdealPD + DCMotor) with decimation=2 and CUDA graph.""" + + __test__ = True + actuators = MIXED_ACTUATORS + + +class TestDecimationRemotizedPD(_DecimationTestBase): + """RemotizedPD (PD + delay + position-based clamping) with decimation=2.""" + + __test__ = True + + @classmethod + def setUpClass(cls): + from isaaclab.actuators.actuator_pd_cfg import RemotizedPDActuatorCfg # noqa: PLC0415 + + cls.actuators = { + "hips": IdealPDActuatorCfg( + joint_names_expr=[".*HAA", ".*HFE"], + stiffness=40.0, + damping=5.0, + effort_limit=80.0, + ), + "knees": RemotizedPDActuatorCfg( + joint_names_expr=[".*KFE"], + stiffness=60.0, + damping=1.5, + effort_limit=80.0, + max_delay=3, + joint_parameter_lookup=SPOT_KNEE_LOOKUP, + ), + } + super().setUpClass() + + # --------------------------------------------------------------------------- # Partial environment reset: verify per-env reset equivalence # --------------------------------------------------------------------------- diff --git a/source/isaaclab_physx/isaaclab_physx/assets/articulation/articulation.py b/source/isaaclab_physx/isaaclab_physx/assets/articulation/articulation.py index 5af2f8dd9c6f..578e7cfed5fa 100644 --- a/source/isaaclab_physx/isaaclab_physx/assets/articulation/articulation.py +++ b/source/isaaclab_physx/isaaclab_physx/assets/articulation/articulation.py @@ -3772,7 +3772,6 @@ def _process_actuators_cfg(self): self.write_joint_stiffness_to_sim_index(stiffness=0.0, joint_ids=adapter.joint_indices) self.write_joint_damping_to_sim_index(damping=0.0, joint_ids=adapter.joint_indices) - # Also create any implicit actuators defined in the configs for actuator_name, actuator_cfg in self.cfg.actuators.items(): cls_type = actuator_cfg.class_type is_implicit = ( @@ -3780,9 +3779,10 @@ def _process_actuators_cfg(self): if isinstance(cls_type, str) else issubclass(cls_type, ImplicitActuator) ) - if not is_implicit: - continue - self._create_lab_actuator(actuator_name, actuator_cfg) + if is_implicit: + self._create_lab_actuator(actuator_name, actuator_cfg) + else: + self._create_lab_actuator(actuator_name, actuator_cfg, properties_only=True) return @@ -3813,8 +3813,19 @@ def _process_actuators_cfg(self): group_count += 1 logger.warning(f"\nActuatorCfg-USD Value Discrepancy Resolution (matching values are skipped): \n{t}") - def _create_lab_actuator(self, actuator_name: str, actuator_cfg: ActuatorBaseCfg) -> None: - """Instantiate a single Lab actuator from its config and write properties to sim.""" + def _create_lab_actuator( + self, actuator_name: str, actuator_cfg: ActuatorBaseCfg, *, properties_only: bool = False, + ) -> None: + """Instantiate a single Lab actuator from its config and write properties to sim. + + Args: + actuator_name: Name for the actuator group. + actuator_cfg: Configuration for the actuator. + properties_only: When ``True``, only write physical joint properties + (armature, limits, friction) without registering the actuator or + writing stiffness/damping. Used for explicit joints managed by + Newton actuators. + """ joint_ids, joint_names = self.find_joints(actuator_cfg.joint_names_expr) if len(joint_names) == 0: raise ValueError( @@ -3841,6 +3852,28 @@ def _create_lab_actuator(self, actuator_name: str, actuator_cfg: ActuatorBaseCfg effort_limit=wp.to_torch(self._data.joint_effort_limits)[:, joint_ids].clone(), velocity_limit=wp.to_torch(self._data.joint_vel_limits)[:, joint_ids], ) + + # Write physical joint properties (armature, limits, friction) — always needed. + self.write_joint_effort_limit_to_sim_index( + limits=actuator.effort_limit_sim, joint_ids=actuator.joint_indices, + ) + self.write_joint_velocity_limit_to_sim_index( + limits=actuator.velocity_limit_sim, joint_ids=actuator.joint_indices, + ) + self.write_joint_armature_to_sim_index(armature=actuator.armature, joint_ids=actuator.joint_indices) + self.write_joint_friction_coefficient_to_sim_index( + joint_friction_coeff=actuator.friction, joint_ids=actuator.joint_indices, + ) + self.write_joint_dynamic_friction_coefficient_to_sim_index( + joint_dynamic_friction_coeff=actuator.dynamic_friction, joint_ids=actuator.joint_indices, + ) + self.write_joint_viscous_friction_coefficient_to_sim_index( + joint_viscous_friction_coeff=actuator.viscous_friction, joint_ids=actuator.joint_indices, + ) + + if properties_only: + return + self.actuators[actuator_name] = actuator # Store the configured values from the actuator model @@ -3871,23 +3904,6 @@ def _create_lab_actuator(self, actuator_name: str, actuator_cfg: ActuatorBaseCfg self.write_joint_stiffness_to_sim_index(stiffness=0.0, joint_ids=actuator.joint_indices) self.write_joint_damping_to_sim_index(damping=0.0, joint_ids=actuator.joint_indices) - self.write_joint_effort_limit_to_sim_index( - limits=actuator.effort_limit_sim, joint_ids=actuator.joint_indices, - ) - self.write_joint_velocity_limit_to_sim_index( - limits=actuator.velocity_limit_sim, joint_ids=actuator.joint_indices, - ) - self.write_joint_armature_to_sim_index(armature=actuator.armature, joint_ids=actuator.joint_indices) - self.write_joint_friction_coefficient_to_sim_index( - joint_friction_coeff=actuator.friction, joint_ids=actuator.joint_indices, - ) - self.write_joint_dynamic_friction_coefficient_to_sim_index( - joint_dynamic_friction_coeff=actuator.dynamic_friction, joint_ids=actuator.joint_indices, - ) - self.write_joint_viscous_friction_coefficient_to_sim_index( - joint_viscous_friction_coeff=actuator.viscous_friction, joint_ids=actuator.joint_indices, - ) - def _process_tendons(self): """Process fixed and spatial tendons.""" # create a list to store the fixed tendon names diff --git a/source/isaaclab_physx/test/assets/test_newton_actuators_physx.py b/source/isaaclab_physx/test/assets/test_newton_actuators_physx.py index 08afb6383222..88f8e07954bd 100644 --- a/source/isaaclab_physx/test/assets/test_newton_actuators_physx.py +++ b/source/isaaclab_physx/test/assets/test_newton_actuators_physx.py @@ -545,6 +545,38 @@ def test_trajectories_not_trivial(self): ] +class TestRemotizedPDEquivalence(_EquivalenceTestBase): + """RemotizedPD (PD + delay + position-based clamping): Lab vs Newton (PhysX). + + Uses the Spot knee lookup table on ANYmal's KFE joints with IdealPD + on HAA and HFE. + """ + + __test__ = True + + @classmethod + def setUpClass(cls): + from isaaclab.actuators.actuator_pd_cfg import RemotizedPDActuatorCfg # noqa: PLC0415 + + cls.actuators = { + "hips": IdealPDActuatorCfg( + joint_names_expr=[".*HAA", ".*HFE"], + stiffness=40.0, + damping=5.0, + effort_limit=80.0, + ), + "knees": RemotizedPDActuatorCfg( + joint_names_expr=[".*KFE"], + stiffness=60.0, + damping=1.5, + effort_limit=80.0, + max_delay=3, + joint_parameter_lookup=SPOT_KNEE_LOOKUP, + ), + } + super().setUpClass() + + class TestRemotizedPDFunctional(unittest.TestCase): """Verify RemotizedPDActuatorCfg runs correctly on PhysX with Newton actuators. From 21d937bbe4cc20522b586971db50e83d9dafe719 Mon Sep 17 00:00:00 2001 From: Julia von Muralt Date: Thu, 7 May 2026 00:03:47 +0200 Subject: [PATCH 11/35] Copy actuator effort to the policy --- .../isaaclab_newton/actuators/kernels.py | 87 +++++++++ .../actuators/newton_actuator_utils.py | 60 +++++- .../assets/articulation/articulation.py | 79 ++++++-- .../assets/test_newton_actuators_newton.py | 173 ++++++++++++++++++ .../assets/articulation/articulation.py | 167 ++++++++++------- .../assets/test_newton_actuators_physx.py | 164 +++++++++++++++++ 6 files changed, 638 insertions(+), 92 deletions(-) create mode 100644 source/isaaclab_newton/isaaclab_newton/actuators/kernels.py diff --git a/source/isaaclab_newton/isaaclab_newton/actuators/kernels.py b/source/isaaclab_newton/isaaclab_newton/actuators/kernels.py new file mode 100644 index 000000000000..47b5ce8e92d5 --- /dev/null +++ b/source/isaaclab_newton/isaaclab_newton/actuators/kernels.py @@ -0,0 +1,87 @@ +# Copyright (c) 2022-2026, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +"""Shared Warp kernels for the Newton actuator fast path.""" + +import warp as wp + + +@wp.kernel +def combine_actuation_force( + user_ff: wp.array2d(dtype=wp.float32), + newton_output: wp.array2d(dtype=wp.float32), + joint_modes: wp.array(dtype=wp.int32), + out: wp.array2d(dtype=wp.float32), +): + """Per-DOF select between user FF (``mode==1``, implicit) and Newton actuator output (``mode==0``). + + For explicit DOFs the user FF was already folded into + :paramref:`newton_output` via Newton's ``joint_act``, so taking it + from there avoids double-counting. + + Args: + user_ff: User-commanded feedforward effort [N·m or N], shape ``(num_envs, num_joints)``. + newton_output: Post-clamp Newton actuator output [N·m or N], shape ``(num_envs, num_joints)``. + joint_modes: Per-DOF mode (``0`` = explicit, ``1`` = implicit), shape ``(num_joints,)``. + out: Output actuation force buffer [N·m or N], shape ``(num_envs, num_joints)``. + """ + i, j = wp.tid() + if joint_modes[j] == 1: + out[i, j] = user_ff[i, j] + else: + out[i, j] = newton_output[i, j] + + +@wp.kernel +def compute_actuator_telemetry( + joint_pos: wp.array2d(dtype=wp.float32), + joint_vel: wp.array2d(dtype=wp.float32), + joint_pos_target: wp.array2d(dtype=wp.float32), + joint_vel_target: wp.array2d(dtype=wp.float32), + joint_effort_target: wp.array2d(dtype=wp.float32), + joint_effort_actual: wp.array2d(dtype=wp.float32), + stiffness: wp.array2d(dtype=wp.float32), + damping: wp.array2d(dtype=wp.float32), + effort_limit: wp.array2d(dtype=wp.float32), + joint_indices: wp.array(dtype=wp.int32), + joint_modes: wp.array(dtype=wp.int32), + target_computed_effort: wp.array2d(dtype=wp.float32), + target_applied_effort: wp.array2d(dtype=wp.float32), +): + """Per-DOF actuator torque telemetry. + + For ``mode==0`` (explicit) copies :paramref:`joint_effort_actual` into both outputs. + For ``mode==1`` (implicit) reproduces :meth:`isaaclab.actuators.ImplicitActuator.compute` — + ``kp*(q_des-q) + kd*(v_des-v) + ff`` with optional clip — as a shadow computation; + the simulator runs the real PD. + + Args: + joint_pos: Current positions [rad or m, depending on joint type], shape ``(num_envs, num_joints)``. + joint_vel: Current velocities [rad/s or m/s], shape ``(num_envs, num_joints)``. + joint_pos_target: Position targets [rad or m, depending on joint type], shape ``(num_envs, num_joints)``. + joint_vel_target: Velocity targets [rad/s or m/s], shape ``(num_envs, num_joints)``. + joint_effort_target: Feedforward efforts [N·m or N], shape ``(num_envs, num_joints)``. + joint_effort_actual: Post-step actuator output [N·m or N], shape ``(num_envs, num_joints)``. + stiffness: Per-joint kp [N·m/rad], shape ``(num_envs, num_joints)``. + damping: Per-joint kd [N·m·s/rad], shape ``(num_envs, num_joints)``. + effort_limit: Absolute effort limit [N·m or N]; use ``inf`` to disable clipping. Shape ``(num_envs, num_joints)``. + joint_indices: DOF indices to process, shape ``(num_actuated_joints,)``. + joint_modes: Per-entry mode (``0`` = copy, ``1`` = compute), shape ``(num_actuated_joints,)``. + target_computed_effort: Output unclipped effort [N·m or N], shape ``(num_envs, num_joints)``. + target_applied_effort: Output post-clip effort [N·m or N], shape ``(num_envs, num_joints)``. + """ + i, j = wp.tid() + dof = joint_indices[j] + if joint_modes[j] == 0: + actual = joint_effort_actual[i, dof] + target_computed_effort[i, dof] = actual + target_applied_effort[i, dof] = actual + else: + err_p = joint_pos_target[i, dof] - joint_pos[i, dof] + err_v = joint_vel_target[i, dof] - joint_vel[i, dof] + computed = stiffness[i, dof] * err_p + damping[i, dof] * err_v + joint_effort_target[i, dof] + target_computed_effort[i, dof] = computed + limit = effort_limit[i, dof] + target_applied_effort[i, dof] = wp.clamp(computed, -limit, limit) diff --git a/source/isaaclab_newton/isaaclab_newton/actuators/newton_actuator_utils.py b/source/isaaclab_newton/isaaclab_newton/actuators/newton_actuator_utils.py index e04f73034732..a8870f22860d 100644 --- a/source/isaaclab_newton/isaaclab_newton/actuators/newton_actuator_utils.py +++ b/source/isaaclab_newton/isaaclab_newton/actuators/newton_actuator_utils.py @@ -69,6 +69,8 @@ Delay, ) +from isaaclab.actuators import ActuatorBase, ImplicitActuator + # =========================================================================== # 1. Warp kernels @@ -112,13 +114,10 @@ class PhysxActuatorWrapper: """Flat-array wrapper serving as ``sim_state`` / ``sim_control`` for :meth:`Actuator.step` on the PhysX backend. - Only needed on PhysX — on the Newton backend the solver's own - ``State`` / ``Control`` objects fulfil this role directly (see the - module docstring for rationale). - - Attributes are reassigned each frame to zero-copy flat views of the - existing 2-D Isaac Lab data via ``wp.array.reshape(-1)``. Only - ``joint_f`` is a persistent allocation (zeroed before each step). + Most attributes are reassigned each frame to zero-copy flat views of + Isaac Lab's 2-D buffers. ``joint_f_2d`` is the only persistent + allocation, sized via :meth:`create`; ``joint_f`` is its flat alias + consumed by the Newton actuator step. """ joint_q: wp.array | None = None @@ -127,6 +126,53 @@ class PhysxActuatorWrapper: joint_target_vel: wp.array | None = None joint_act: wp.array | None = None joint_f: wp.array | None = None + joint_f_2d: wp.array | None = None + + @classmethod + def create(cls, num_envs: int, num_joints: int, device: str) -> PhysxActuatorWrapper: + """Allocate the persistent ``joint_f`` buffer for the given articulation shape.""" + w = cls() + w.joint_f_2d = wp.zeros((num_envs, num_joints), dtype=wp.float32, device=device) + w.joint_f = w.joint_f_2d.reshape(-1) + return w + + +def build_actuator_telemetry( + actuators: dict[str, ActuatorBase], + num_envs: int, + num_joints: int, + device: str, +) -> tuple[wp.array, wp.array, wp.array]: + """Build per-DOF telemetry tables. + + Per-DOF ``modes`` is ``1`` for joints covered by an + :class:`~isaaclab.actuators.ImplicitActuator` group (shadow-PD) and + ``0`` otherwise (copy from the simulator's actuator output). + ``effort_limit`` carries the implicit-clip absolute limit (``inf`` + elsewhere). + + Returns: + ``(indices, modes, effort_limit)`` Warp arrays. + """ + modes = torch.zeros(num_joints, dtype=torch.int32, device=device) + effort_limit = torch.full( + (num_envs, num_joints), float("inf"), device=device, dtype=torch.float32 + ) + for actuator in actuators.values(): + if not isinstance(actuator, ImplicitActuator): + continue + j_ids = actuator.joint_indices + if j_ids == slice(None) or j_ids is None: + modes[:] = 1 + effort_limit[:] = actuator.effort_limit + else: + modes[j_ids.long()] = 1 + effort_limit[:, j_ids.long()] = actuator.effort_limit + + indices = wp.from_torch( + torch.arange(num_joints, dtype=torch.int32, device=device), dtype=wp.int32 + ) + return indices, wp.from_torch(modes, dtype=wp.int32), wp.from_torch(effort_limit, dtype=wp.float32) # =========================================================================== diff --git a/source/isaaclab_newton/isaaclab_newton/assets/articulation/articulation.py b/source/isaaclab_newton/isaaclab_newton/assets/articulation/articulation.py index d2a6ea590be1..dcad788a6b1a 100644 --- a/source/isaaclab_newton/isaaclab_newton/assets/articulation/articulation.py +++ b/source/isaaclab_newton/isaaclab_newton/assets/articulation/articulation.py @@ -27,7 +27,11 @@ from isaaclab.assets.articulation.base_articulation import BaseArticulation try: - from isaaclab_newton.actuators.newton_actuator_utils import NewtonActuatorAdapter + from isaaclab_newton.actuators import kernels as actuator_kernels + from isaaclab_newton.actuators.newton_actuator_utils import ( + NewtonActuatorAdapter, + build_actuator_telemetry, + ) _HAS_NEWTON_ACTUATORS = True except ImportError: @@ -316,6 +320,30 @@ def update(self, dt: float): dt: The time step size in seconds. """ self.data.update(dt) + # Shadow PD telemetry for implicit DOFs (the simulator runs the real PD). + if self._has_newton_actuators and self._telemetry_indices is not None: + wp.launch( + actuator_kernels.compute_actuator_telemetry, + dim=(self.num_instances, self._telemetry_indices.shape[0]), + inputs=[ + self._data.joint_pos.warp, + self._data.joint_vel.warp, + self._data._joint_pos_target, + self._data._joint_vel_target, + self._data._joint_effort_target, + self._data._sim_bind_joint_effort, + self._data.joint_stiffness.warp, + self._data.joint_damping.warp, + self._telemetry_effort_limit, + self._telemetry_indices, + self._telemetry_modes, + ], + outputs=[ + self._data._computed_torque, + self._data._applied_torque, + ], + device=self.device, + ) """ Operations - Finders. @@ -3459,6 +3487,10 @@ def _process_actuators_cfg(self): # if this is false, we by-pass certain checks when doing actuator-related operations self._has_implicit_actuators = False self._has_newton_actuators = False + # Per-DOF telemetry tables; ``None`` when no Newton fast path is active. + self._telemetry_indices: wp.array | None = None + self._telemetry_modes: wp.array | None = None + self._telemetry_effort_limit: wp.array | None = None _use_newton_actuators = getattr(self._sim_cfg, "use_newton_actuators", False) @@ -3473,26 +3505,27 @@ def _process_actuators_cfg(self): model = SimulationManager.get_model() - if not model.actuators: - return + # Enable the fast path even for all-implicit articulations: + # the solver runs PD internally; Lab only forwards targets. + self._has_newton_actuators = True - dof_layout = self._root_view.frequency_layouts[NewtonModel.AttributeFrequency.JOINT_DOF] - if dof_layout.slice is not None: - dof_offset = dof_layout.slice.start - elif dof_layout.indices is not None: - dof_offset = int(dof_layout.indices.numpy()[0]) - else: - dof_offset = 0 + if model.actuators: + dof_layout = self._root_view.frequency_layouts[NewtonModel.AttributeFrequency.JOINT_DOF] + if dof_layout.slice is not None: + dof_offset = dof_layout.slice.start + elif dof_layout.indices is not None: + dof_offset = int(dof_layout.indices.numpy()[0]) + else: + dof_offset = 0 - adapter = NewtonActuatorAdapter( - model.actuators, self.num_instances, self.num_joints, dof_offset, self.device, - ) - adapter.finalize() - self.actuators["newton"] = adapter - SimulationManager.register_adapter(adapter) - self.write_joint_stiffness_to_sim_index(stiffness=0.0, joint_ids=adapter.joint_indices) - self.write_joint_damping_to_sim_index(damping=0.0, joint_ids=adapter.joint_indices) - self._has_newton_actuators = True + adapter = NewtonActuatorAdapter( + model.actuators, self.num_instances, self.num_joints, dof_offset, self.device, + ) + adapter.finalize() + self.actuators["newton"] = adapter + SimulationManager.register_adapter(adapter) + self.write_joint_stiffness_to_sim_index(stiffness=0.0, joint_ids=adapter.joint_indices) + self.write_joint_damping_to_sim_index(damping=0.0, joint_ids=adapter.joint_indices) for actuator_name, actuator_cfg in self.cfg.actuators.items(): cls_type = actuator_cfg.class_type @@ -3506,6 +3539,14 @@ def _process_actuators_cfg(self): else: self._create_lab_actuator(actuator_name, actuator_cfg, properties_only=True) + ( + self._telemetry_indices, + self._telemetry_modes, + self._telemetry_effort_limit, + ) = build_actuator_telemetry( + self.actuators, self.num_instances, self.num_joints, self.device, + ) + return # --- Standard Isaac Lab actuator path --- diff --git a/source/isaaclab_newton/test/assets/test_newton_actuators_newton.py b/source/isaaclab_newton/test/assets/test_newton_actuators_newton.py index 4e028288d654..53935de210dc 100644 --- a/source/isaaclab_newton/test/assets/test_newton_actuators_newton.py +++ b/source/isaaclab_newton/test/assets/test_newton_actuators_newton.py @@ -193,6 +193,76 @@ def _run_simulation( return {"joint_pos": recorded_pos, "joint_vel": recorded_vel} +def _run_simulation_with_telemetry( + actuators: dict, + use_newton_actuators: bool, + *, + dt: float = DT, + newton_cfg: NewtonCfg = NEWTON_CFG, + num_steps: int = NUM_STEPS, + decimation: int = 1, +) -> dict: + """Like :func:`_run_simulation` but also records ``computed_torque`` / ``applied_torque``.""" + sim_cfg = SimulationCfg( + dt=dt, + physics=newton_cfg, + use_newton_actuators=use_newton_actuators, + ) + + with build_simulation_context( + device="cuda:0", + gravity_enabled=True, + add_ground_plane=True, + sim_cfg=sim_cfg, + ) as sim: + sim._app_control_on_stop_handle = None + + for i in range(NUM_ENVS): + sim_utils.create_prim( + f"/World/Env_{i}", "Xform", translation=(i * 3.0, 0, 0) + ) + + art_cfg = ANYMAL_C_CFG.replace( + actuators=actuators, + prim_path="/World/Env_.*/Robot", + ) + articulation = Articulation(art_cfg) + sim.reset() + assert articulation.is_initialized + + if use_newton_actuators and decimation > 1: + SimulationManager.set_decimation(decimation) + + init_pos = wp.to_torch(articulation.data.joint_pos).clone() + target_pos = init_pos + TARGET_OFFSET + target_vel = torch.zeros_like(init_pos) + + articulation.set_joint_position_target_index(target=target_pos) + articulation.set_joint_velocity_target_index(target=target_vel) + + recorded_pos, recorded_vel = [], [] + recorded_computed, recorded_applied = [], [] + for _ in range(num_steps): + for _ in range(decimation): + articulation.write_data_to_sim() + sim.step() + articulation.update(dt) + + recorded_pos.append(wp.to_torch(articulation.data.joint_pos).clone()) + recorded_vel.append(wp.to_torch(articulation.data.joint_vel).clone()) + recorded_computed.append(wp.to_torch(articulation.data.computed_torque).clone()) + recorded_applied.append(wp.to_torch(articulation.data.applied_torque).clone()) + + return { + "joint_pos": recorded_pos, + "joint_vel": recorded_vel, + "computed_torque": recorded_computed, + "applied_torque": recorded_applied, + "target_pos": target_pos.clone(), + "target_vel": target_vel.clone(), + } + + # --------------------------------------------------------------------------- # Base test class # --------------------------------------------------------------------------- @@ -309,6 +379,109 @@ class TestMixedWithImplicitEquivalence(_EquivalenceTestBase): actuators = MIXED_WITH_IMPLICIT_ACTUATORS +# --------------------------------------------------------------------------- +# Implicit-only fast-path: enable Newton actuator branch with no explicit groups +# --------------------------------------------------------------------------- + +IMPLICIT_ONLY_ACTUATORS = { + "legs": ImplicitActuatorCfg( + joint_names_expr=[".*HAA", ".*HFE", ".*KFE"], + stiffness=40.0, + damping=5.0, + ), +} + + +class TestImplicitOnlyEquivalence(_EquivalenceTestBase): + """All-implicit articulation with ``use_newton_actuators=True``: Lab vs fast-path.""" + + __test__ = True + actuators = IMPLICIT_ONLY_ACTUATORS + + +class TestExplicitOnlyTelemetry(unittest.TestCase): + """Explicit-only Newton actuators: telemetry copies from ``joint_f``.""" + + @classmethod + def setUpClass(cls): + cls.result = _run_simulation_with_telemetry( + DC_MOTOR_ACTUATORS, use_newton_actuators=True, + ) + + def test_telemetry_is_nonzero(self): + last = self.result["applied_torque"][-1] + self.assertGreater( + last.abs().max().item(), + 1e-3, + "applied_torque is all-zero — explicit-DOF telemetry path did not run", + ) + + def test_computed_equals_applied_explicit(self): + """For Newton-managed DOFs, both buffers carry ``joint_f`` (post-clamp).""" + for step_i, (comp, app) in enumerate( + zip(self.result["computed_torque"], self.result["applied_torque"]) + ): + torch.testing.assert_close( + comp, + app, + atol=1e-5, + rtol=1e-5, + msg=f"computed != applied for explicit-only at step {step_i}", + ) + + +class TestImplicitOnlyTelemetry(unittest.TestCase): + """Implicit-only fast path: shadow-PD telemetry matches the Lab formula.""" + + @classmethod + def setUpClass(cls): + cls.result = _run_simulation_with_telemetry( + IMPLICIT_ONLY_ACTUATORS, use_newton_actuators=True, + ) + cls.kp = 40.0 + cls.kd = 5.0 + cls.target_offset = TARGET_OFFSET + + def test_telemetry_is_nonzero(self): + """Telemetry kernel actually fired (not stuck at default zeros).""" + last = self.result["computed_torque"][-1] + max_abs = last.abs().max().item() + self.assertGreater(max_abs, 1e-3, "computed_torque is all-zero — telemetry kernel did not run") + + def test_telemetry_matches_pd_formula(self): + """Telemetry value matches kp*(q_des-q) + kd*(v_des-v) within tolerance.""" + target_q = self.result["target_pos"] + target_v = self.result["target_vel"] + for step_i, (q, qd, comp) in enumerate( + zip( + self.result["joint_pos"], + self.result["joint_vel"], + self.result["computed_torque"], + ) + ): + expected = self.kp * (target_q - q) + self.kd * (target_v - qd) + torch.testing.assert_close( + comp, + expected, + atol=5e-2, + rtol=1e-2, + msg=f"Telemetry diverged from PD formula at step {step_i}", + ) + + def test_applied_equals_computed_when_no_clip(self): + """Implicit cfg has no effort_limit → applied == computed.""" + for step_i, (comp, app) in enumerate( + zip(self.result["computed_torque"], self.result["applied_torque"]) + ): + torch.testing.assert_close( + app, + comp, + atol=1e-5, + rtol=1e-5, + msg=f"applied_torque != computed_torque at step {step_i} (no clip expected)", + ) + + # --------------------------------------------------------------------------- # DelayedPD equivalence: PD with actuator command delay # --------------------------------------------------------------------------- diff --git a/source/isaaclab_physx/isaaclab_physx/assets/articulation/articulation.py b/source/isaaclab_physx/isaaclab_physx/assets/articulation/articulation.py index 578e7cfed5fa..fd279b616041 100644 --- a/source/isaaclab_physx/isaaclab_physx/assets/articulation/articulation.py +++ b/source/isaaclab_physx/isaaclab_physx/assets/articulation/articulation.py @@ -24,9 +24,11 @@ from isaaclab.assets.articulation.base_articulation import BaseArticulation try: + from isaaclab_newton.actuators import kernels as actuator_kernels from isaaclab_newton.actuators.newton_actuator_utils import ( NewtonActuatorAdapter, PhysxActuatorWrapper, + build_actuator_telemetry, ) _HAS_NEWTON_ACTUATORS = True @@ -261,15 +263,19 @@ def write_data_to_sim(self): ) self._instantaneous_wrench_composer.reset() - if self._physx_actuator_wrapper is not None: - # Newton actuator path: zero effort, run implicit actuators - # (if any) which write directly into the staging buffer, then - # step Newton actuators which add their effort in-place. + if self._has_newton_actuators and self._physx_actuator_wrapper is None: + # Implicit-only: ImplicitActuator.compute is a no-op, hand user buffers to PhysX directly. + self.root_view.set_dof_actuation_forces(self._data._joint_effort_target, self._ALL_INDICES) + if self._has_implicit_actuators: + self.root_view.set_dof_position_targets(self._data._joint_pos_target, self._ALL_INDICES) + self.root_view.set_dof_velocity_targets(self._data._joint_vel_target, self._ALL_INDICES) + elif self._has_newton_actuators: + # Mixed: Newton step fills explicit DOFs, combine kernel merges with user FF. self._apply_actuator_model_newton() self.root_view.set_dof_actuation_forces(self._joint_effort_target_sim, self._ALL_INDICES) if self._has_implicit_actuators: - self.root_view.set_dof_position_targets(self._joint_pos_target_sim, self._ALL_INDICES) - self.root_view.set_dof_velocity_targets(self._joint_vel_target_sim, self._ALL_INDICES) + self.root_view.set_dof_position_targets(self._data._joint_pos_target, self._ALL_INDICES) + self.root_view.set_dof_velocity_targets(self._data._joint_vel_target, self._ALL_INDICES) else: # Standard Lab actuator path self._apply_actuator_model() @@ -285,6 +291,30 @@ def update(self, dt: float): dt: The time step size in seconds. """ self.data.update(dt) + # Shadow PD telemetry for implicit DOFs (the simulator runs the real PD). + if self._has_newton_actuators and self._telemetry_indices is not None: + wp.launch( + actuator_kernels.compute_actuator_telemetry, + dim=(self.num_instances, self._telemetry_indices.shape[0]), + inputs=[ + self._data.joint_pos.warp, + self._data.joint_vel.warp, + self._data._joint_pos_target, + self._data._joint_vel_target, + self._data._joint_effort_target, + self._joint_effort_target_sim, + self._data.joint_stiffness.warp, + self._data.joint_damping.warp, + self._telemetry_effort_limit, + self._telemetry_indices, + self._telemetry_modes, + ], + outputs=[ + self._data._computed_torque, + self._data._applied_torque, + ], + device=self.device, + ) """ Operations - Finders. @@ -3748,29 +3778,53 @@ def _process_actuators_cfg(self): # flag for implicit actuators # if this is false, we by-pass certain checks when doing actuator-related operations self._has_implicit_actuators = False + self._has_newton_actuators = False + # Per-DOF telemetry tables; ``None`` when no Newton fast path is active. + self._telemetry_indices: wp.array | None = None + self._telemetry_modes: wp.array | None = None + self._telemetry_effort_limit: wp.array | None = None _use_newton_actuators = getattr(self._sim_cfg, "use_newton_actuators", False) if _HAS_NEWTON_ACTUATORS and _use_newton_actuators: from isaaclab.sim.utils.stage import get_current_stage # noqa: PLC0415 - first_prim = find_first_matching_prim(self.cfg.prim_path) - art_prim_path = str(first_prim.GetPath()) if first_prim is not None else None + # Enable the fast path even for all-implicit articulations: + # PhysX runs PD internally; Lab only forwards targets. + self._has_newton_actuators = True - adapter = NewtonActuatorAdapter.from_usd( - stage=get_current_stage(), - joint_names=self.joint_names, - num_envs=self.num_instances, - num_joints=self.num_joints, - device=self.device, - articulation_prim_path=art_prim_path, + # Author Newton actuator prims only if any explicit Lab group exists. + has_explicit = any( + not ( + "ImplicitActuator" in actuator_cfg.class_type + if isinstance(actuator_cfg.class_type, str) + else issubclass(actuator_cfg.class_type, ImplicitActuator) + ) + for actuator_cfg in self.cfg.actuators.values() ) - self._physx_actuator_wrapper = PhysxActuatorWrapper() - adapter.finalize() - self.actuators["newton"] = adapter - self.write_joint_stiffness_to_sim_index(stiffness=0.0, joint_ids=adapter.joint_indices) - self.write_joint_damping_to_sim_index(damping=0.0, joint_ids=adapter.joint_indices) + if has_explicit: + first_prim = find_first_matching_prim(self.cfg.prim_path) + art_prim_path = str(first_prim.GetPath()) if first_prim is not None else None + + adapter = NewtonActuatorAdapter.from_usd( + stage=get_current_stage(), + joint_names=self.joint_names, + num_envs=self.num_instances, + num_joints=self.num_joints, + device=self.device, + articulation_prim_path=art_prim_path, + ) + + self._physx_actuator_wrapper = PhysxActuatorWrapper.create( + num_envs=self.num_instances, + num_joints=self.num_joints, + device=self.device, + ) + adapter.finalize() + self.actuators["newton"] = adapter + self.write_joint_stiffness_to_sim_index(stiffness=0.0, joint_ids=adapter.joint_indices) + self.write_joint_damping_to_sim_index(damping=0.0, joint_ids=adapter.joint_indices) for actuator_name, actuator_cfg in self.cfg.actuators.items(): cls_type = actuator_cfg.class_type @@ -3784,6 +3838,13 @@ def _process_actuators_cfg(self): else: self._create_lab_actuator(actuator_name, actuator_cfg, properties_only=True) + ( + self._telemetry_indices, + self._telemetry_modes, + self._telemetry_effort_limit, + ) = build_actuator_telemetry( + self.actuators, self.num_instances, self.num_joints, self.device, + ) return # --- Standard Isaac Lab actuator path --- @@ -3998,63 +4059,37 @@ def _apply_actuator_model(self): ) def _apply_actuator_model_newton(self): - """Process actuators when Newton-native actuators are active. + """Mixed implicit + Newton path: step Newton actuators, then combine with user FF. - Zeros the effort staging buffer, processes any implicit Lab actuators - (which write directly into it), then steps Newton actuators whose - ``joint_f`` is a flat view of the same buffer — no scatter needed. + Newton writes explicit-DOF torques to ``joint_f_2d``; the combine kernel + merges them with user FF (for implicit DOFs) into ``_joint_effort_target_sim``. + Implicit-only takes the direct path in :meth:`write_data_to_sim`. """ - # Zero the effort buffer; implicit actuators and Newton both write into it - self._joint_effort_target_sim.zero_() - - # Process implicit Lab actuators (if any) - for actuator in self.actuators.values(): - if _HAS_NEWTON_ACTUATORS and isinstance(actuator, NewtonActuatorAdapter): - continue - control_action = ArticulationActions( - joint_positions=wp.to_torch(self._data.joint_pos_target)[:, actuator.joint_indices], - joint_velocities=wp.to_torch(self._data.joint_vel_target)[:, actuator.joint_indices], - joint_efforts=wp.to_torch(self._data.joint_effort_target)[:, actuator.joint_indices], - joint_indices=actuator.joint_indices, - ) - control_action = actuator.compute( - control_action, - joint_pos=wp.to_torch(self._data.joint_pos)[:, actuator.joint_indices], - joint_vel=wp.to_torch(self._data.joint_vel)[:, actuator.joint_indices], - ) - joint_indices = actuator.joint_indices - if actuator.joint_indices == slice(None) or actuator.joint_indices is None: - joint_indices = self._ALL_JOINT_INDICES - wp.launch( - articulation_kernels.update_targets, - dim=(self.num_instances, joint_indices.shape[0]), - inputs=[ - control_action.joint_positions, - control_action.joint_velocities, - control_action.joint_efforts, - joint_indices, - ], - outputs=[ - self._joint_pos_target_sim, - self._joint_vel_target_sim, - self._joint_effort_target_sim, - ], - device=self.device, - ) - - # Step Newton actuators — joint_f is a flat view of _joint_effort_target_sim - # so Newton writes effort directly into the staging buffer. - newton_adapter = self.actuators.get("newton") + newton_adapter = self.actuators["newton"] w = self._physx_actuator_wrapper + + # Newton actuators scatter-add into joint_f_2d; zero it first. + w.joint_f_2d.zero_() w.joint_q = self._data.joint_pos.warp.reshape(-1) w.joint_qd = self._data.joint_vel.warp.reshape(-1) w.joint_target_pos = self._data.joint_pos_target.warp.reshape(-1) w.joint_target_vel = self._data.joint_vel_target.warp.reshape(-1) w.joint_act = self._data.joint_effort_target.warp.reshape(-1) - w.joint_f = self._joint_effort_target_sim.reshape(-1) newton_adapter.step(w, w, SimulationManager.get_physics_dt()) + wp.launch( + actuator_kernels.combine_actuation_force, + dim=(self.num_instances, self.num_joints), + inputs=[ + self._data._joint_effort_target, + w.joint_f_2d, + self._telemetry_modes, + ], + outputs=[self._joint_effort_target_sim], + device=self.device, + ) + """ Internal helpers -- Debugging. """ diff --git a/source/isaaclab_physx/test/assets/test_newton_actuators_physx.py b/source/isaaclab_physx/test/assets/test_newton_actuators_physx.py index 88f8e07954bd..4a7aaa7169d7 100644 --- a/source/isaaclab_physx/test/assets/test_newton_actuators_physx.py +++ b/source/isaaclab_physx/test/assets/test_newton_actuators_physx.py @@ -289,6 +289,170 @@ class TestMixedWithImplicitEquivalence(_EquivalenceTestBase): actuators = MIXED_WITH_IMPLICIT_ACTUATORS +# --------------------------------------------------------------------------- +# Implicit-only fast-path: enable Newton-actuator branch on PhysX with no explicit groups +# --------------------------------------------------------------------------- + +IMPLICIT_ONLY_ACTUATORS = { + "legs": ImplicitActuatorCfg( + joint_names_expr=[".*HAA", ".*HFE", ".*KFE"], + stiffness=40.0, + damping=5.0, + ), +} + + +class TestImplicitOnlyEquivalencePhysx(_EquivalenceTestBase): + """All-implicit articulation on PhysX with ``use_newton_actuators=True``: Lab vs fast-path.""" + + __test__ = True + actuators = IMPLICIT_ONLY_ACTUATORS + + +def _run_simulation_with_telemetry( + actuators: dict, + use_newton_actuators: bool, + *, + num_steps: int = NUM_STEPS, +) -> dict: + """Like :func:`_run_simulation` but also records ``computed_torque`` / ``applied_torque``.""" + sim_cfg = SimulationCfg( + dt=DT, + physics=PhysxCfg(), + use_newton_actuators=use_newton_actuators, + ) + + with build_simulation_context( + device="cuda:0", + gravity_enabled=True, + add_ground_plane=True, + sim_cfg=sim_cfg, + ) as sim: + sim._app_control_on_stop_handle = None + + for i in range(NUM_ENVS): + sim_utils.create_prim( + f"/World/Env_{i}", "Xform", translation=(i * 3.0, 0, 0) + ) + + art_cfg = ANYMAL_C_CFG.replace( + actuators=actuators, + prim_path="/World/Env_.*/Robot", + ) + articulation = Articulation(art_cfg) + sim.reset() + assert articulation.is_initialized + + init_pos = wp.to_torch(articulation.data.joint_pos).clone() + target_pos = init_pos + TARGET_OFFSET + target_vel = torch.zeros_like(init_pos) + + articulation.set_joint_position_target_index(target=target_pos) + articulation.set_joint_velocity_target_index(target=target_vel) + + recorded_pos, recorded_vel = [], [] + recorded_computed, recorded_applied = [], [] + for _ in range(num_steps): + articulation.write_data_to_sim() + sim.step() + articulation.update(DT) + + recorded_pos.append(wp.to_torch(articulation.data.joint_pos).clone()) + recorded_vel.append(wp.to_torch(articulation.data.joint_vel).clone()) + recorded_computed.append(wp.to_torch(articulation.data.computed_torque).clone()) + recorded_applied.append(wp.to_torch(articulation.data.applied_torque).clone()) + + return { + "joint_pos": recorded_pos, + "joint_vel": recorded_vel, + "computed_torque": recorded_computed, + "applied_torque": recorded_applied, + "target_pos": target_pos.clone(), + "target_vel": target_vel.clone(), + } + + +class TestImplicitOnlyTelemetryPhysx(unittest.TestCase): + """Implicit-only fast path on PhysX: shadow-PD telemetry matches the Lab formula.""" + + @classmethod + def setUpClass(cls): + cls.result = _run_simulation_with_telemetry( + IMPLICIT_ONLY_ACTUATORS, use_newton_actuators=True, + ) + cls.kp = 40.0 + cls.kd = 5.0 + + def test_telemetry_is_nonzero(self): + last = self.result["computed_torque"][-1] + self.assertGreater( + last.abs().max().item(), + 1e-3, + "computed_torque is all-zero — implicit telemetry kernel did not run", + ) + + def test_telemetry_matches_pd_formula(self): + target_q = self.result["target_pos"] + target_v = self.result["target_vel"] + for step_i, (q, qd, comp) in enumerate( + zip( + self.result["joint_pos"], + self.result["joint_vel"], + self.result["computed_torque"], + ) + ): + expected = self.kp * (target_q - q) + self.kd * (target_v - qd) + torch.testing.assert_close( + comp, + expected, + atol=5e-2, + rtol=1e-2, + msg=f"Telemetry diverged from PD formula at step {step_i}", + ) + + def test_applied_equals_computed_when_no_clip(self): + for step_i, (comp, app) in enumerate( + zip(self.result["computed_torque"], self.result["applied_torque"]) + ): + torch.testing.assert_close( + app, + comp, + atol=1e-5, + rtol=1e-5, + msg=f"applied_torque != computed_torque at step {step_i} (no clip expected)", + ) + + +class TestExplicitOnlyTelemetryPhysx(unittest.TestCase): + """Explicit-only Newton actuators on PhysX: telemetry copies from the staging effort buffer.""" + + @classmethod + def setUpClass(cls): + cls.result = _run_simulation_with_telemetry( + DC_MOTOR_ACTUATORS, use_newton_actuators=True, + ) + + def test_telemetry_is_nonzero(self): + last = self.result["applied_torque"][-1] + self.assertGreater( + last.abs().max().item(), + 1e-3, + "applied_torque is all-zero — explicit-DOF telemetry path did not run", + ) + + def test_computed_equals_applied_explicit(self): + for step_i, (comp, app) in enumerate( + zip(self.result["computed_torque"], self.result["applied_torque"]) + ): + torch.testing.assert_close( + comp, + app, + atol=1e-5, + rtol=1e-5, + msg=f"computed != applied for explicit-only at step {step_i}", + ) + + # --------------------------------------------------------------------------- # Partial environment reset: verify per-env reset equivalence # --------------------------------------------------------------------------- From 01229144644c104b5a13feb54a91755fbfe5b074 Mon Sep 17 00:00:00 2001 From: Julia von Muralt Date: Thu, 7 May 2026 15:17:58 +0200 Subject: [PATCH 12/35] versions --- docker/Dockerfile.base | 12 ------- source/isaaclab/setup.py | 2 +- .../isaaclab_newton/physics/newton_manager.py | 34 +++++++++++++++++-- source/isaaclab_newton/setup.py | 2 +- .../renderers/ovrtx_renderer_kernels.py | 2 +- source/isaaclab_physx/setup.py | 2 +- source/isaaclab_visualizers/setup.py | 3 ++ tools/wheel_builder/res/python_packages.toml | 10 +++--- 8 files changed, 43 insertions(+), 24 deletions(-) diff --git a/docker/Dockerfile.base b/docker/Dockerfile.base index 536c37860274..c7fce487e63e 100644 --- a/docker/Dockerfile.base +++ b/docker/Dockerfile.base @@ -112,18 +112,6 @@ RUN --mount=type=cache,target=${DOCKER_USER_HOME}/.cache/pip \ ${ISAACLAB_PATH}/isaaclab.sh -p -m pip install --no-build-isolation nlopt==2.6.2; \ fi -# HACK: Pre-install Newton before `isaaclab.sh --install`. -# Newton's own setup.py requires warp-lang>=1.13 but isaaclab pins -# warp-lang==1.12 (bundled with Isaac Sim). Installing Newton via pip's -# dependency resolver would fail, so we force-install it with --no-deps -# first. The Docker base image also ships an older pre-built Newton whose -# version string collides (both report 1.2.0.dev0), so --force-reinstall -# ensures the correct commit is used. -RUN --mount=type=cache,target=${DOCKER_USER_HOME}/.cache/pip \ - ${ISAACLAB_PATH}/isaaclab.sh -p -m pip install --force-reinstall --no-deps \ - "newton @ git+https://github.com/jvonmuralt/newton.git@ignore_path" && \ - ${ISAACLAB_PATH}/isaaclab.sh -p -m pip install \ - "mujoco~=3.8.0" "mujoco-warp>=3.8.0.1,~=3.8.0" # installing Isaac Lab dependencies # use pip caching to avoid reinstalling large packages diff --git a/source/isaaclab/setup.py b/source/isaaclab/setup.py index 0511db8ec439..61439d0fc0ed 100644 --- a/source/isaaclab/setup.py +++ b/source/isaaclab/setup.py @@ -35,7 +35,7 @@ # image processing "transformers==4.57.6", "einops", # needed for transformers, doesn't always auto-install - "warp-lang==1.12.0", + "warp-lang==1.13.0", "matplotlib>=3.10.3", # minimum version for Python 3.12 support # make sure this is consistent with isaac sim version "pillow==12.1.1", diff --git a/source/isaaclab_newton/isaaclab_newton/physics/newton_manager.py b/source/isaaclab_newton/isaaclab_newton/physics/newton_manager.py index 7193c2eaf293..01eec0026a1e 100644 --- a/source/isaaclab_newton/isaaclab_newton/physics/newton_manager.py +++ b/source/isaaclab_newton/isaaclab_newton/physics/newton_manager.py @@ -71,7 +71,7 @@ def _set_fabric_transforms( i = int(wp.tid()) idx = int(newton_indices[i]) transform = newton_body_q[idx] - fabric_transforms[i] = wp.transpose(wp.mat44d(wp.math.transform_to_matrix(transform))) + fabric_transforms[i] = wp.transpose(wp.mat44d(wp.transform_to_matrix(transform))) @wp.kernel(enable_backward=False) @@ -517,6 +517,11 @@ def clear(cls): cls._newton_imu_sensors = [] cls._report_contacts = False cls._adapter = None + # Set by an articulation that took the ``use_newton_actuators=True`` + # branch in ``_process_actuators_cfg``. Together with the adapter + # check, this gates whether the decimation loop can be captured into + # a CUDA graph (see :meth:`_is_all_graphable`). + cls._use_newton_actuators_active = False cls._decimation = 1 # Per-world reset masks cls._world_reset_mask = None @@ -822,6 +827,7 @@ def start_simulation(cls) -> None: # Actuator registration is deferred to register_adapter(), # called by _process_actuators_cfg after the adapter is constructed. cls._adapter = None + cls._use_newton_actuators_active = False # Allocate per-world reset masks (used by all solvers for masked FK, and by Kamino for masked reset) cls._world_reset_mask = wp.zeros(cls._model.world_count, dtype=wp.int32, device=device) @@ -1345,8 +1351,30 @@ def get_solver_dt(cls) -> float: @classmethod def _is_all_graphable(cls) -> bool: - """``True`` when an adapter is registered and all its actuators are CUDA-graph-safe.""" - return cls._adapter is not None and cls._adapter.is_all_graphable + """``True`` when the decimation loop can be captured into a CUDA graph. + + Requires: + 1. An articulation took the ``use_newton_actuators=True`` branch + (signalled via :meth:`activate_newton_actuator_path`). + 2. Either no actuator adapter was needed (all-implicit) or every + actuator in the adapter is CUDA-graph-safe. + """ + if not cls._use_newton_actuators_active: + return False + return cls._adapter is None or cls._adapter.is_all_graphable + + @classmethod + def activate_newton_actuator_path(cls) -> None: + """Signal that an articulation has opted into the Newton actuator fast path. + + Called by the articulation's ``_process_actuators_cfg`` whenever it + enters the ``use_newton_actuators=True`` branch — including the + all-implicit case where no :class:`NewtonActuatorAdapter` is built. + Required because :meth:`_is_all_graphable` can no longer rely on + adapter presence alone to distinguish the fast path from the + standard Lab path (which also has ``_adapter is None``). + """ + cls._use_newton_actuators_active = True @classmethod def register_adapter(cls, adapter: NewtonActuatorAdapter) -> None: diff --git a/source/isaaclab_newton/setup.py b/source/isaaclab_newton/setup.py index 88d9c07b5171..4f80a1af3b34 100644 --- a/source/isaaclab_newton/setup.py +++ b/source/isaaclab_newton/setup.py @@ -41,7 +41,7 @@ def run(self): "mujoco~=3.8.0", "mujoco-warp>=3.8.0.1,~=3.8.0", "PyOpenGL-accelerate==3.1.10", - "newton @ git+https://github.com/jvonmuralt/newton.git@ignore_path", + "newton @ git+https://github.com/newton-physics/newton.git@v1.2.0rc2", ], } diff --git a/source/isaaclab_ov/isaaclab_ov/renderers/ovrtx_renderer_kernels.py b/source/isaaclab_ov/isaaclab_ov/renderers/ovrtx_renderer_kernels.py index d7647df4c3d7..207d6dc7f264 100644 --- a/source/isaaclab_ov/isaaclab_ov/renderers/ovrtx_renderer_kernels.py +++ b/source/isaaclab_ov/isaaclab_ov/renderers/ovrtx_renderer_kernels.py @@ -264,4 +264,4 @@ def sync_newton_transforms_kernel( i = wp.tid() body_idx = newton_body_indices[i] transform = newton_body_q[body_idx] - ovrtx_transforms[i] = wp.transpose(wp.mat44d(wp.math.transform_to_matrix(transform))) + ovrtx_transforms[i] = wp.transpose(wp.mat44d(wp.transform_to_matrix(transform))) diff --git a/source/isaaclab_physx/setup.py b/source/isaaclab_physx/setup.py index 1e917e938c2b..9cc172addf50 100644 --- a/source/isaaclab_physx/setup.py +++ b/source/isaaclab_physx/setup.py @@ -20,7 +20,7 @@ EXTRAS_REQUIRE = { "newton": [ - "newton @ git+https://github.com/newton-physics/newton.git@2684d75bfa4bb8b058a93b81c458a74b7701c997", + "newton @ git+https://github.com/newton-physics/newton.git@v1.2.0rc2", ], } diff --git a/source/isaaclab_visualizers/setup.py b/source/isaaclab_visualizers/setup.py index 0ad6b9fec0fa..fba3ac1d7f46 100644 --- a/source/isaaclab_visualizers/setup.py +++ b/source/isaaclab_visualizers/setup.py @@ -18,12 +18,15 @@ "newton": [ "warp-lang", "PyOpenGL-accelerate", + "newton @ git+https://github.com/newton-physics/newton.git@v1.2.0rc2", "imgui-bundle>=1.92.5", ], "rerun": [ + "newton @ git+https://github.com/newton-physics/newton.git@v1.2.0rc2", "rerun-sdk>=0.29.0", ], "viser": [ + "newton @ git+https://github.com/newton-physics/newton.git@v1.2.0rc2", "viser>=1.0.16", ], } diff --git a/tools/wheel_builder/res/python_packages.toml b/tools/wheel_builder/res/python_packages.toml index f6a42b90a1bc..7a415ef4de62 100644 --- a/tools/wheel_builder/res/python_packages.toml +++ b/tools/wheel_builder/res/python_packages.toml @@ -22,7 +22,7 @@ pyproject.dependencies.all = [ # image processing "transformers==4.57.6", "einops", # needed for transformers, doesn't always auto-install - "warp-lang==1.12.0", + "warp-lang==1.13.0", "matplotlib>=3.10.3", # make sure this is consistent with isaac sim version "pillow==12.1.1", @@ -82,10 +82,10 @@ pyproject.optional-dependencies.all = [ # https://github.com/isaac-sim/IsaacLab/blob/main/source/isaaclab_newton/setup.py # ================================================================================ { "newton" = [ - "warp-lang==1.12.0", - "mujoco==3.6.0", - "mujoco-warp==3.6.0", - "newton @ git+https://github.com/newton-physics/newton.git@a27277ed49d6f307b8a1e4c394be7e1d14965a62", + "warp-lang==1.13.0", + "mujoco==3.8.0", + "mujoco-warp==3.8.0.1", + "newton @ git+https://github.com/newton-physics/newton.git@v1.2.0rc2", "PyOpenGL-accelerate==3.1.10" ] }, # ================================================================================ From 79891437a11f17a7f1dfe0a6f927e49d7440eec1 Mon Sep 17 00:00:00 2001 From: Julia von Muralt Date: Thu, 7 May 2026 15:20:18 +0200 Subject: [PATCH 13/35] refactor --- .../actuators/newton_actuator_utils.py | 130 +++++++++++++++++- .../assets/articulation/articulation.py | 93 +++---------- .../assets/test_newton_actuators_newton.py | 25 +++- .../assets/articulation/articulation.py | 116 +++++++++++++--- .../assets/test_newton_actuators_physx.py | 70 ++++++++++ 5 files changed, 340 insertions(+), 94 deletions(-) diff --git a/source/isaaclab_newton/isaaclab_newton/actuators/newton_actuator_utils.py b/source/isaaclab_newton/isaaclab_newton/actuators/newton_actuator_utils.py index a8870f22860d..bcf2521c492d 100644 --- a/source/isaaclab_newton/isaaclab_newton/actuators/newton_actuator_utils.py +++ b/source/isaaclab_newton/isaaclab_newton/actuators/newton_actuator_utils.py @@ -53,7 +53,7 @@ from __future__ import annotations import logging -from collections.abc import Sequence +from collections.abc import Callable, Sequence from dataclasses import dataclass from typing import Any @@ -104,6 +104,47 @@ def _scatter_gain_kernel( dst[env * num_joints + local_dof] = src[i] +@wp.kernel(enable_backward=False) +def _gather_gain_kernel( + flat_src: wp.array(dtype=wp.float32), + dst: wp.array(dtype=wp.float32), + indices: wp.array(dtype=wp.uint32), + env_mask: wp.array(dtype=wp.bool), + dof_offset: int, + num_joints: int, +): + """Gather from flat ``(num_envs * num_joints)`` layout into a per-actuator + controller array, only for envs where ``env_mask`` is ``True``.""" + i = wp.tid() + global_dof = int(indices[i]) - dof_offset + env = global_dof // num_joints + if env_mask[env]: + local_dof = global_dof % num_joints + dst[i] = flat_src[env * num_joints + local_dof] + + +@wp.kernel(enable_backward=False) +def _scatter_gain_at_envs_kernel( + in_data: wp.array2d(dtype=wp.float32), + env_ids: wp.array(dtype=wp.int32), + out_data: wp.array2d(dtype=wp.float32), +): + """Scatter ``in_data[i, j]`` into ``out_data[env_ids[i], j]`` for all (i, j).""" + i, j = wp.tid() + out_data[env_ids[i], j] = in_data[i, j] + + +@wp.kernel(enable_backward=False) +def _fill_gain_at_envs_kernel( + value: float, + env_ids: wp.array(dtype=wp.int32), + out_data: wp.array2d(dtype=wp.float32), +): + """Set ``out_data[env_ids[i], j] = value`` for all (i, j).""" + i, j = wp.tid() + out_data[env_ids[i], j] = value + + # =========================================================================== # 2. PhysX stepping helper # =========================================================================== @@ -371,6 +412,93 @@ def is_all_graphable(self) -> bool: """``True`` when all actuators are CUDA-graph-safe.""" return len(self.actuators) > 0 and all(a.is_graphable() for a in self.actuators) + def update_gain_at_env_ids( + self, + gain: str, + values: torch.Tensor | wp.array | float, + env_ids: wp.array, + ) -> wp.array: + """Scatter ``values`` into :attr:`stiffness` or :attr:`damping` at *env_ids*. + + Shared backend-independent step used by ``write_actuator_*_to_sim`` to + keep the kp/kd buffer the adapter exposes for DR in sync with what + each Newton actuator's controller will receive. + + Args: + gain: ``"stiffness"`` or ``"damping"``. + values: Per-env-per-joint values shape ``(len(env_ids), num_joints)``, + or a scalar to broadcast to every (env, joint). + env_ids: Warp int32 array of env indices to update. + + Returns: + Warp view of the updated ``(num_envs, num_joints)`` buffer. + """ + if gain == "stiffness": + buf = self.stiffness + elif gain == "damping": + buf = self.damping + else: + raise ValueError(f"gain must be 'stiffness' or 'damping', got {gain!r}") + buf_wp = wp.from_torch(buf, dtype=wp.float32) + if isinstance(values, float): + wp.launch( + _fill_gain_at_envs_kernel, + dim=(env_ids.shape[0], self.num_joints), + inputs=[values, env_ids], + outputs=[buf_wp], + device=self._device, + ) + else: + wp.launch( + _scatter_gain_at_envs_kernel, + dim=(env_ids.shape[0], self.num_joints), + inputs=[values, env_ids], + outputs=[buf_wp], + device=self._device, + ) + return buf_wp + + def write_stiffness_to_sim( + self, + stiffness: torch.Tensor | wp.array | float, + env_ids: wp.array, + env_mask: wp.array, + propagate_fn: Callable[["NewtonActuatorAdapter", Actuator, Any, str, wp.array, wp.array], None], + ) -> None: + """Update the kp buffer at *env_ids* and push the new values into each Newton controller. + + The per-actuator propagation step is backend-specific (Newton uses + :meth:`ArticulationView.set_actuator_parameter`; PhysX scatters via a + local Warp kernel), so the caller injects it as *propagate_fn*. + """ + self._write_gain_to_sim("stiffness", "kp", stiffness, env_ids, env_mask, propagate_fn) + + def write_damping_to_sim( + self, + damping: torch.Tensor | wp.array | float, + env_ids: wp.array, + env_mask: wp.array, + propagate_fn: Callable[["NewtonActuatorAdapter", Actuator, Any, str, wp.array, wp.array], None], + ) -> None: + """Update the kd buffer at *env_ids* and push the new values into each Newton controller.""" + self._write_gain_to_sim("damping", "kd", damping, env_ids, env_mask, propagate_fn) + + def _write_gain_to_sim( + self, + gain: str, + attr: str, + values: torch.Tensor | wp.array | float, + env_ids: wp.array, + env_mask: wp.array, + propagate_fn: Callable[["NewtonActuatorAdapter", Actuator, Any, str, wp.array, wp.array], None], + ) -> None: + """Shared body for :meth:`write_stiffness_to_sim` / :meth:`write_damping_to_sim`.""" + buf = self.update_gain_at_env_ids(gain, values, env_ids) + for newton_act in self.actuators: + ctrl = newton_act.controller + if hasattr(ctrl, attr): + propagate_fn(self, newton_act, ctrl, attr, buf, env_mask) + # -- config helpers (used by both adapter and authoring) ----------------- @staticmethod diff --git a/source/isaaclab_newton/isaaclab_newton/assets/articulation/articulation.py b/source/isaaclab_newton/isaaclab_newton/assets/articulation/articulation.py index dcad788a6b1a..6e4c3a80778e 100644 --- a/source/isaaclab_newton/isaaclab_newton/assets/articulation/articulation.py +++ b/source/isaaclab_newton/isaaclab_newton/assets/articulation/articulation.py @@ -51,8 +51,6 @@ from .articulation_data import ArticulationData if TYPE_CHECKING: - from newton.actuators import Actuator - from isaaclab.assets.articulation.articulation_cfg import ArticulationCfg # import logger @@ -2066,43 +2064,10 @@ def write_actuator_stiffness_to_sim( stiffness: torch.Tensor | wp.array | float, env_ids: Sequence[int] | torch.Tensor | wp.array | None = None, ) -> None: - """Write actuator stiffness (kp) into the Newton controller arrays. - - Analogous to :meth:`write_joint_stiffness_to_sim_index` but targets the - Newton actuator controller ``kp`` instead of the solver-level - ``joint_target_ke``. - - Args: - adapter: The :class:`NewtonActuatorAdapter` whose gains to update. - stiffness: Stiffness values [N/m or N*m/rad, depending on joint - type]. Shape is ``(len(env_ids), num_joints)``. - env_ids: Environment indices. If ``None``, all environments. - """ + """Write actuator stiffness (``kp``) into each Newton controller for *env_ids*.""" env_ids = self._resolve_env_ids(env_ids) - stiffness_wp = wp.from_torch(adapter.stiffness, dtype=wp.float32) - if isinstance(stiffness, float): - wp.launch( - articulation_kernels.float_data_to_buffer_with_indices, - dim=(env_ids.shape[0], self._ALL_JOINT_INDICES.shape[0]), - inputs=[stiffness, env_ids, self._ALL_JOINT_INDICES], - outputs=[stiffness_wp], - device=self.device, - ) - else: - wp.launch( - shared_kernels.write_2d_data_to_buffer_with_indices, - dim=(env_ids.shape[0], self._ALL_JOINT_INDICES.shape[0]), - inputs=[stiffness, env_ids, self._ALL_JOINT_INDICES], - outputs=[stiffness_wp], - device=self.device, - ) env_mask = self._env_ids_to_mask(env_ids) - for newton_act in adapter.actuators: - if hasattr(newton_act.controller, "kp"): - self._root_view.set_actuator_parameter( - actuator=newton_act, component=newton_act.controller, name="kp", - values=stiffness_wp, mask=env_mask, - ) + adapter.write_stiffness_to_sim(stiffness, env_ids, env_mask, self._propagate_gain_via_view) def write_actuator_damping_to_sim( self, @@ -2111,43 +2076,26 @@ def write_actuator_damping_to_sim( damping: torch.Tensor | wp.array | float, env_ids: Sequence[int] | torch.Tensor | wp.array | None = None, ) -> None: - """Write actuator damping (kd) into the Newton controller arrays. - - Analogous to :meth:`write_joint_damping_to_sim_index` but targets the - Newton actuator controller ``kd`` instead of the solver-level - ``joint_target_kd``. - - Args: - adapter: The :class:`NewtonActuatorAdapter` whose gains to update. - damping: Damping values [N*s/m or N*m*s/rad, depending on joint - type]. Shape is ``(len(env_ids), num_joints)``. - env_ids: Environment indices. If ``None``, all environments. - """ + """Write actuator damping (``kd``) into each Newton controller for *env_ids*.""" env_ids = self._resolve_env_ids(env_ids) - damping_wp = wp.from_torch(adapter.damping, dtype=wp.float32) - if isinstance(damping, float): - wp.launch( - articulation_kernels.float_data_to_buffer_with_indices, - dim=(env_ids.shape[0], self._ALL_JOINT_INDICES.shape[0]), - inputs=[damping, env_ids, self._ALL_JOINT_INDICES], - outputs=[damping_wp], - device=self.device, - ) - else: - wp.launch( - shared_kernels.write_2d_data_to_buffer_with_indices, - dim=(env_ids.shape[0], self._ALL_JOINT_INDICES.shape[0]), - inputs=[damping, env_ids, self._ALL_JOINT_INDICES], - outputs=[damping_wp], - device=self.device, - ) env_mask = self._env_ids_to_mask(env_ids) - for newton_act in adapter.actuators: - if hasattr(newton_act.controller, "kd"): - self._root_view.set_actuator_parameter( - actuator=newton_act, component=newton_act.controller, name="kd", - values=damping_wp, mask=env_mask, - ) + adapter.write_damping_to_sim(damping, env_ids, env_mask, self._propagate_gain_via_view) + + def _propagate_gain_via_view( + self, + adapter: NewtonActuatorAdapter, + actuator, + controller, + attr: str, + values: wp.array, + env_mask: wp.array, + ) -> None: + """Per-actuator gain propagation using Newton's simulator-side scatter API.""" + del adapter # Newton path needs only the view + per-actuator info. + self._root_view.set_actuator_parameter( + actuator=actuator, component=controller, name=attr, + values=values, mask=env_mask, + ) def _env_ids_to_mask(self, env_ids: wp.array) -> wp.array: """Convert warp env_ids to a boolean Warp mask.""" @@ -3508,6 +3456,7 @@ def _process_actuators_cfg(self): # Enable the fast path even for all-implicit articulations: # the solver runs PD internally; Lab only forwards targets. self._has_newton_actuators = True + SimulationManager.activate_newton_actuator_path() if model.actuators: dof_layout = self._root_view.frequency_layouts[NewtonModel.AttributeFrequency.JOINT_DOF] diff --git a/source/isaaclab_newton/test/assets/test_newton_actuators_newton.py b/source/isaaclab_newton/test/assets/test_newton_actuators_newton.py index 53935de210dc..7f39e894e455 100644 --- a/source/isaaclab_newton/test/assets/test_newton_actuators_newton.py +++ b/source/isaaclab_newton/test/assets/test_newton_actuators_newton.py @@ -282,11 +282,18 @@ class _EquivalenceTestBase(unittest.TestCase): pos_rtol: float = 1e-3 vel_atol: float = 0.1 vel_rtol: float = 1e-2 + # Torque equivalence tolerates a one-physics-step timing offset: + # standard path computes applied_torque pre-step, fast path post-step. + # The difference scales with kp*v*dt + kd*a*dt — largest in the first + # transient steps before targets are tracked, so skip those. + torque_atol: float = 5.0 + torque_rtol: float = 0.15 + torque_skip_steps: int = 3 @classmethod def setUpClass(cls): - cls.lab_result = _run_simulation(cls.actuators, use_newton_actuators=False) - cls.newton_result = _run_simulation(cls.actuators, use_newton_actuators=True) + cls.lab_result = _run_simulation_with_telemetry(cls.actuators, use_newton_actuators=False) + cls.newton_result = _run_simulation_with_telemetry(cls.actuators, use_newton_actuators=True) def test_joint_positions_match(self): for step_i, (lab, newton) in enumerate( @@ -312,6 +319,20 @@ def test_joint_velocities_match(self): msg=f"Joint velocities diverged at step {step_i}", ) + def test_applied_torque_match(self): + for step_i, (lab, newton) in enumerate( + zip(self.lab_result["applied_torque"], self.newton_result["applied_torque"]) + ): + if step_i < self.torque_skip_steps: + continue + torch.testing.assert_close( + lab, + newton, + atol=self.torque_atol, + rtol=self.torque_rtol, + msg=f"applied_torque diverged at step {step_i}", + ) + def test_trajectories_not_trivial(self): first = self.lab_result["joint_pos"][0] last = self.lab_result["joint_pos"][-1] diff --git a/source/isaaclab_physx/isaaclab_physx/assets/articulation/articulation.py b/source/isaaclab_physx/isaaclab_physx/assets/articulation/articulation.py index fd279b616041..84d6dc130c10 100644 --- a/source/isaaclab_physx/isaaclab_physx/assets/articulation/articulation.py +++ b/source/isaaclab_physx/isaaclab_physx/assets/articulation/articulation.py @@ -28,12 +28,19 @@ from isaaclab_newton.actuators.newton_actuator_utils import ( NewtonActuatorAdapter, PhysxActuatorWrapper, + _gather_gain_kernel, build_actuator_telemetry, ) _HAS_NEWTON_ACTUATORS = True except ImportError: _HAS_NEWTON_ACTUATORS = False + + +@wp.kernel(enable_backward=False) +def _build_env_mask_kernel(mask: wp.array(dtype=wp.bool), indices: wp.array(dtype=wp.int32)): + i = wp.tid() + mask[indices[i]] = True from isaaclab.sim.utils.queries import find_first_matching_prim, get_all_matching_child_prims from isaaclab.utils.string import resolve_matching_names, resolve_matching_names_values from isaaclab.utils.types import ArticulationActions @@ -263,21 +270,22 @@ def write_data_to_sim(self): ) self._instantaneous_wrench_composer.reset() - if self._has_newton_actuators and self._physx_actuator_wrapper is None: - # Implicit-only: ImplicitActuator.compute is a no-op, hand user buffers to PhysX directly. - self.root_view.set_dof_actuation_forces(self._data._joint_effort_target, self._ALL_INDICES) - if self._has_implicit_actuators: - self.root_view.set_dof_position_targets(self._data._joint_pos_target, self._ALL_INDICES) - self.root_view.set_dof_velocity_targets(self._data._joint_vel_target, self._ALL_INDICES) - elif self._has_newton_actuators: - # Mixed: Newton step fills explicit DOFs, combine kernel merges with user FF. - self._apply_actuator_model_newton() - self.root_view.set_dof_actuation_forces(self._joint_effort_target_sim, self._ALL_INDICES) + if self._has_newton_actuators: + # Newton fast path: pos/vel targets pass straight through; the + # actuation buffer is either user FF (implicit-only) or the + # combine-kernel result of Newton output + user FF (mixed). + if self._physx_actuator_wrapper is not None: + self._apply_actuator_model_newton() + actuation_buf = self._joint_effort_target_sim + else: + actuation_buf = self._data._joint_effort_target + self.root_view.set_dof_actuation_forces(actuation_buf, self._ALL_INDICES) if self._has_implicit_actuators: self.root_view.set_dof_position_targets(self._data._joint_pos_target, self._ALL_INDICES) self.root_view.set_dof_velocity_targets(self._data._joint_vel_target, self._ALL_INDICES) else: - # Standard Lab actuator path + # Standard Lab actuator path: per-group ``actuator.compute()`` may + # transform targets, so we push the staging buffers PhysX-side. self._apply_actuator_model() self.root_view.set_dof_actuation_forces(self._joint_effort_target_sim, self._ALL_INDICES) if self._has_implicit_actuators: @@ -1218,6 +1226,69 @@ def write_joint_damping_to_sim_index( cpu_env_ids = self._get_cpu_env_ids(env_ids) self.root_view.set_dof_dampings(wp.clone(self.data._joint_damping, device="cpu"), indices=cpu_env_ids) + """ + Operations - Newton Actuator Parameter Writers. + + Mirror of the writers on :class:`isaaclab_newton.assets.Articulation`. + Required so that domain randomization terms (``randomize_actuator_gains``) + propagate kp/kd updates to explicit Newton actuators on the PhysX backend. + """ + + def write_actuator_stiffness_to_sim( + self, + adapter: NewtonActuatorAdapter, + *, + stiffness: torch.Tensor | wp.array | float, + env_ids: Sequence[int] | torch.Tensor | wp.array | None = None, + ) -> None: + """Write actuator stiffness (``kp``) into each Newton controller for *env_ids*.""" + env_ids = self._resolve_env_ids(env_ids) + env_mask = self._env_ids_to_mask(env_ids) + adapter.write_stiffness_to_sim(stiffness, env_ids, env_mask, self._propagate_gain_via_kernel) + + def write_actuator_damping_to_sim( + self, + adapter: NewtonActuatorAdapter, + *, + damping: torch.Tensor | wp.array | float, + env_ids: Sequence[int] | torch.Tensor | wp.array | None = None, + ) -> None: + """Write actuator damping (``kd``) into each Newton controller for *env_ids*.""" + env_ids = self._resolve_env_ids(env_ids) + env_mask = self._env_ids_to_mask(env_ids) + adapter.write_damping_to_sim(damping, env_ids, env_mask, self._propagate_gain_via_kernel) + + def _propagate_gain_via_kernel( + self, + adapter: NewtonActuatorAdapter, + actuator, + controller, + attr: str, + values: wp.array, + env_mask: wp.array, + ) -> None: + """Per-actuator gain propagation using a local Warp scatter kernel. + + Newton has :meth:`ArticulationView.set_actuator_parameter` for this; + on PhysX we use :data:`_gather_gain_kernel` since no equivalent + simulator-side API exists. + """ + wp.launch( + _gather_gain_kernel, + dim=actuator.indices.shape[0], + inputs=[ + values.flatten(), getattr(controller, attr), actuator.indices, env_mask, + adapter._dof_offset, adapter.num_joints, + ], + device=self.device, + ) + + def _env_ids_to_mask(self, env_ids: wp.array) -> wp.array: + """Convert warp ``env_ids`` to a boolean Warp mask of length ``num_instances``.""" + mask = wp.zeros(self.num_instances, dtype=wp.bool, device=self.device) + wp.launch(_build_env_mask_kernel, dim=env_ids.shape[0], inputs=[mask, env_ids], device=self.device) + return mask + def write_joint_damping_to_sim_mask( self, *, @@ -3821,6 +3892,16 @@ def _process_actuators_cfg(self): num_joints=self.num_joints, device=self.device, ) + # Bind the wrapper's flat aliases of state/input buffers once. + # The underlying wp.arrays alias stable PhysX-owned GPU memory + # whose device pointer is fixed for the articulation's lifetime, + # so the views remain valid for every subsequent step. + w = self._physx_actuator_wrapper + w.joint_q = self._data.joint_pos.warp.reshape(-1) + w.joint_qd = self._data.joint_vel.warp.reshape(-1) + w.joint_target_pos = self._data.joint_pos_target.warp.reshape(-1) + w.joint_target_vel = self._data.joint_vel_target.warp.reshape(-1) + w.joint_act = self._data.joint_effort_target.warp.reshape(-1) adapter.finalize() self.actuators["newton"] = adapter self.write_joint_stiffness_to_sim_index(stiffness=0.0, joint_ids=adapter.joint_indices) @@ -4068,14 +4149,11 @@ def _apply_actuator_model_newton(self): newton_adapter = self.actuators["newton"] w = self._physx_actuator_wrapper - # Newton actuators scatter-add into joint_f_2d; zero it first. - w.joint_f_2d.zero_() - w.joint_q = self._data.joint_pos.warp.reshape(-1) - w.joint_qd = self._data.joint_vel.warp.reshape(-1) - w.joint_target_pos = self._data.joint_pos_target.warp.reshape(-1) - w.joint_target_vel = self._data.joint_vel_target.warp.reshape(-1) - w.joint_act = self._data.joint_effort_target.warp.reshape(-1) - + # Newton actuators scatter-add into joint_f at their own indices; + # ``adapter.step`` pre-zeros exactly those slots. Slots outside any + # actuator's indices are implicit DOFs that the combine kernel reads + # from ``user_ff`` instead of ``joint_f_2d``, so they don't need + # clearing here. newton_adapter.step(w, w, SimulationManager.get_physics_dt()) wp.launch( diff --git a/source/isaaclab_physx/test/assets/test_newton_actuators_physx.py b/source/isaaclab_physx/test/assets/test_newton_actuators_physx.py index 4a7aaa7169d7..6ea74df0972e 100644 --- a/source/isaaclab_physx/test/assets/test_newton_actuators_physx.py +++ b/source/isaaclab_physx/test/assets/test_newton_actuators_physx.py @@ -453,6 +453,76 @@ def test_computed_equals_applied_explicit(self): ) +class TestWriteActuatorGainsPhysx(unittest.TestCase): + """``write_actuator_*_to_sim`` propagates kp/kd into Newton controllers (PhysX backend). + + Catches the silent no-op that ``randomize_actuator_gains`` would suffer + from on PhysX with ``use_newton_actuators=True`` if these writers were + missing (the events.py term falls back to ``hasattr`` and skips + explicit actuators). + """ + + def test_writers_exist(self): + # Guards against the missing-writer regression. + from isaaclab_physx.assets import Articulation + self.assertTrue( + hasattr(Articulation, "write_actuator_stiffness_to_sim"), + "Articulation is missing write_actuator_stiffness_to_sim — DR will silently no-op", + ) + self.assertTrue( + hasattr(Articulation, "write_actuator_damping_to_sim"), + "Articulation is missing write_actuator_damping_to_sim — DR will silently no-op", + ) + + def test_writers_propagate_to_controller(self): + sim_cfg = SimulationCfg(dt=DT, physics=PhysxCfg(), use_newton_actuators=True) + with build_simulation_context( + device="cuda:0", gravity_enabled=True, add_ground_plane=True, sim_cfg=sim_cfg, + ) as sim: + sim._app_control_on_stop_handle = None + for i in range(NUM_ENVS): + sim_utils.create_prim(f"/World/Env_{i}", "Xform", translation=(i * 3.0, 0, 0)) + art_cfg = ANYMAL_C_CFG.replace( + actuators=DC_MOTOR_ACTUATORS, prim_path="/World/Env_.*/Robot", + ) + articulation = Articulation(art_cfg) + sim.reset() + adapter = articulation.actuators["newton"] + # Snapshot initial controller gains. + kp_before = [ + wp.to_torch(a.controller.kp).clone() for a in adapter.actuators if hasattr(a.controller, "kp") + ] + kd_before = [ + wp.to_torch(a.controller.kd).clone() for a in adapter.actuators if hasattr(a.controller, "kd") + ] + self.assertGreater(len(kp_before), 0, "expected at least one PD controller in adapter") + new_kp = adapter.stiffness.clone() * 2.0 + new_kd = adapter.damping.clone() * 3.0 + articulation.write_actuator_stiffness_to_sim(adapter, stiffness=new_kp) + articulation.write_actuator_damping_to_sim(adapter, damping=new_kd) + # Verify each controller's kp/kd actually changed (and roughly doubled/tripled). + kp_idx = 0 + kd_idx = 0 + for newton_act in adapter.actuators: + ctrl = newton_act.controller + if hasattr(ctrl, "kp"): + after = wp.to_torch(ctrl.kp) + self.assertFalse( + torch.equal(after, kp_before[kp_idx]), + "controller.kp unchanged after write_actuator_stiffness_to_sim", + ) + torch.testing.assert_close(after, kp_before[kp_idx] * 2.0, atol=1e-4, rtol=1e-4) + kp_idx += 1 + if hasattr(ctrl, "kd"): + after = wp.to_torch(ctrl.kd) + self.assertFalse( + torch.equal(after, kd_before[kd_idx]), + "controller.kd unchanged after write_actuator_damping_to_sim", + ) + torch.testing.assert_close(after, kd_before[kd_idx] * 3.0, atol=1e-4, rtol=1e-4) + kd_idx += 1 + + # --------------------------------------------------------------------------- # Partial environment reset: verify per-env reset equivalence # --------------------------------------------------------------------------- From 905297871bf5cea215deb2f6dd6039c71d210623 Mon Sep 17 00:00:00 2001 From: Julia von Muralt Date: Thu, 7 May 2026 21:16:01 +0200 Subject: [PATCH 14/35] refactor shadowing --- docker/Dockerfile.base | 1 - .../isaaclab_newton/actuators/kernels.py | 105 +++++-------- .../actuators/newton_actuator_utils.py | 28 +--- .../assets/articulation/articulation.py | 81 +++++----- .../isaaclab_newton/physics/newton_manager.py | 29 +++- .../assets/test_newton_actuators_newton.py | 146 +++++++++++++++--- .../assets/articulation/articulation.py | 114 ++++++-------- .../assets/test_newton_actuators_physx.py | 142 ++++++++++++++--- 8 files changed, 412 insertions(+), 234 deletions(-) diff --git a/docker/Dockerfile.base b/docker/Dockerfile.base index c7fce487e63e..c10c036ce0a6 100644 --- a/docker/Dockerfile.base +++ b/docker/Dockerfile.base @@ -112,7 +112,6 @@ RUN --mount=type=cache,target=${DOCKER_USER_HOME}/.cache/pip \ ${ISAACLAB_PATH}/isaaclab.sh -p -m pip install --no-build-isolation nlopt==2.6.2; \ fi - # installing Isaac Lab dependencies # use pip caching to avoid reinstalling large packages RUN --mount=type=cache,target=${DOCKER_USER_HOME}/.cache/pip \ diff --git a/source/isaaclab_newton/isaaclab_newton/actuators/kernels.py b/source/isaaclab_newton/isaaclab_newton/actuators/kernels.py index 47b5ce8e92d5..467cc340e6c6 100644 --- a/source/isaaclab_newton/isaaclab_newton/actuators/kernels.py +++ b/source/isaaclab_newton/isaaclab_newton/actuators/kernels.py @@ -8,80 +8,53 @@ import warp as wp -@wp.kernel -def combine_actuation_force( - user_ff: wp.array2d(dtype=wp.float32), - newton_output: wp.array2d(dtype=wp.float32), - joint_modes: wp.array(dtype=wp.int32), - out: wp.array2d(dtype=wp.float32), -): - """Per-DOF select between user FF (``mode==1``, implicit) and Newton actuator output (``mode==0``). - - For explicit DOFs the user FF was already folded into - :paramref:`newton_output` via Newton's ``joint_act``, so taking it - from there avoids double-counting. - - Args: - user_ff: User-commanded feedforward effort [N·m or N], shape ``(num_envs, num_joints)``. - newton_output: Post-clamp Newton actuator output [N·m or N], shape ``(num_envs, num_joints)``. - joint_modes: Per-DOF mode (``0`` = explicit, ``1`` = implicit), shape ``(num_joints,)``. - out: Output actuation force buffer [N·m or N], shape ``(num_envs, num_joints)``. - """ - i, j = wp.tid() - if joint_modes[j] == 1: - out[i, j] = user_ff[i, j] - else: - out[i, j] = newton_output[i, j] - - -@wp.kernel -def compute_actuator_telemetry( +@wp.kernel(enable_backward=False) +def synch_torque_and_apply_implicit_feedforwards( joint_pos: wp.array2d(dtype=wp.float32), joint_vel: wp.array2d(dtype=wp.float32), joint_pos_target: wp.array2d(dtype=wp.float32), joint_vel_target: wp.array2d(dtype=wp.float32), joint_effort_target: wp.array2d(dtype=wp.float32), - joint_effort_actual: wp.array2d(dtype=wp.float32), - stiffness: wp.array2d(dtype=wp.float32), - damping: wp.array2d(dtype=wp.float32), + joint_stiffness: wp.array2d(dtype=wp.float32), + joint_damping: wp.array2d(dtype=wp.float32), effort_limit: wp.array2d(dtype=wp.float32), - joint_indices: wp.array(dtype=wp.int32), joint_modes: wp.array(dtype=wp.int32), - target_computed_effort: wp.array2d(dtype=wp.float32), - target_applied_effort: wp.array2d(dtype=wp.float32), + sim_bind_joint_effort: wp.array2d(dtype=wp.float32), + computed: wp.array2d(dtype=wp.float32), + applied: wp.array2d(dtype=wp.float32), ): - """Per-DOF actuator torque telemetry. - - For ``mode==0`` (explicit) copies :paramref:`joint_effort_actual` into both outputs. - For ``mode==1`` (implicit) reproduces :meth:`isaaclab.actuators.ImplicitActuator.compute` — - ``kp*(q_des-q) + kd*(v_des-v) + ff`` with optional clip — as a shadow computation; - the simulator runs the real PD. - - Args: - joint_pos: Current positions [rad or m, depending on joint type], shape ``(num_envs, num_joints)``. - joint_vel: Current velocities [rad/s or m/s], shape ``(num_envs, num_joints)``. - joint_pos_target: Position targets [rad or m, depending on joint type], shape ``(num_envs, num_joints)``. - joint_vel_target: Velocity targets [rad/s or m/s], shape ``(num_envs, num_joints)``. - joint_effort_target: Feedforward efforts [N·m or N], shape ``(num_envs, num_joints)``. - joint_effort_actual: Post-step actuator output [N·m or N], shape ``(num_envs, num_joints)``. - stiffness: Per-joint kp [N·m/rad], shape ``(num_envs, num_joints)``. - damping: Per-joint kd [N·m·s/rad], shape ``(num_envs, num_joints)``. - effort_limit: Absolute effort limit [N·m or N]; use ``inf`` to disable clipping. Shape ``(num_envs, num_joints)``. - joint_indices: DOF indices to process, shape ``(num_actuated_joints,)``. - joint_modes: Per-entry mode (``0`` = copy, ``1`` = compute), shape ``(num_actuated_joints,)``. - target_computed_effort: Output unclipped effort [N·m or N], shape ``(num_envs, num_joints)``. - target_applied_effort: Output post-clip effort [N·m or N], shape ``(num_envs, num_joints)``. + """In-graph post-actuator hook: route implicit FF and sync telemetry. + + For each (env, dof): + * Implicit DOF: write user FF to ``joint_f`` (sim integrates it + alongside the joint-drive PD), clamp the shadow PD ``kp*err_p + + kd*err_v`` to ``±effort_limit``, and write ``computed = applied + = PD_clipped + FF``. + * Explicit DOF: mirror Newton's post-actuator ``joint_f`` into + ``computed`` / ``applied``. + + Limitation: ``effort_limit`` here only clamps the **PD shadow** used + for telemetry. The simulator's joint drive applies its max-force + only to the PD term, so user feedforward effort can exceed + ``effort_limit`` once it lands in ``joint_f``. Limiting the *total* + applied effort would require Newton's motor-actuator path + (configurable via the Newton team), which may have negative perf + implications. """ i, j = wp.tid() - dof = joint_indices[j] - if joint_modes[j] == 0: - actual = joint_effort_actual[i, dof] - target_computed_effort[i, dof] = actual - target_applied_effort[i, dof] = actual + if joint_modes[j] == 1: + sim_bind_joint_effort[i, j] = joint_effort_target[i, j] + err_p = joint_pos_target[i, j] - joint_pos[i, j] + err_v = joint_vel_target[i, j] - joint_vel[i, j] + pd = joint_stiffness[i, j] * err_p + joint_damping[i, j] * err_v + limit = effort_limit[i, j] + pd_clipped = wp.clamp(pd, -limit, limit) + total = pd_clipped + sim_bind_joint_effort[i, j] + computed[i, j] = total + applied[i, j] = total else: - err_p = joint_pos_target[i, dof] - joint_pos[i, dof] - err_v = joint_vel_target[i, dof] - joint_vel[i, dof] - computed = stiffness[i, dof] * err_p + damping[i, dof] * err_v + joint_effort_target[i, dof] - target_computed_effort[i, dof] = computed - limit = effort_limit[i, dof] - target_applied_effort[i, dof] = wp.clamp(computed, -limit, limit) + val = sim_bind_joint_effort[i, j] + computed[i, j] = val + applied[i, j] = val + + diff --git a/source/isaaclab_newton/isaaclab_newton/actuators/newton_actuator_utils.py b/source/isaaclab_newton/isaaclab_newton/actuators/newton_actuator_utils.py index bcf2521c492d..a886c44b3277 100644 --- a/source/isaaclab_newton/isaaclab_newton/actuators/newton_actuator_utils.py +++ b/source/isaaclab_newton/isaaclab_newton/actuators/newton_actuator_utils.py @@ -178,42 +178,26 @@ def create(cls, num_envs: int, num_joints: int, device: str) -> PhysxActuatorWra return w -def build_actuator_telemetry( +def build_implicit_dof_mask( actuators: dict[str, ActuatorBase], - num_envs: int, num_joints: int, device: str, -) -> tuple[wp.array, wp.array, wp.array]: - """Build per-DOF telemetry tables. +) -> wp.array: + """Per-DOF mask consumed by the in-graph implicit-FF kernel. - Per-DOF ``modes`` is ``1`` for joints covered by an - :class:`~isaaclab.actuators.ImplicitActuator` group (shadow-PD) and - ``0`` otherwise (copy from the simulator's actuator output). - ``effort_limit`` carries the implicit-clip absolute limit (``inf`` - elsewhere). - - Returns: - ``(indices, modes, effort_limit)`` Warp arrays. + Entry is ``1`` for DOFs covered by an + :class:`~isaaclab.actuators.ImplicitActuator` group, ``0`` otherwise. """ modes = torch.zeros(num_joints, dtype=torch.int32, device=device) - effort_limit = torch.full( - (num_envs, num_joints), float("inf"), device=device, dtype=torch.float32 - ) for actuator in actuators.values(): if not isinstance(actuator, ImplicitActuator): continue j_ids = actuator.joint_indices if j_ids == slice(None) or j_ids is None: modes[:] = 1 - effort_limit[:] = actuator.effort_limit else: modes[j_ids.long()] = 1 - effort_limit[:, j_ids.long()] = actuator.effort_limit - - indices = wp.from_torch( - torch.arange(num_joints, dtype=torch.int32, device=device), dtype=wp.int32 - ) - return indices, wp.from_torch(modes, dtype=wp.int32), wp.from_torch(effort_limit, dtype=wp.float32) + return wp.from_torch(modes, dtype=wp.int32) # =========================================================================== diff --git a/source/isaaclab_newton/isaaclab_newton/assets/articulation/articulation.py b/source/isaaclab_newton/isaaclab_newton/assets/articulation/articulation.py index 6e4c3a80778e..1e8069501e77 100644 --- a/source/isaaclab_newton/isaaclab_newton/assets/articulation/articulation.py +++ b/source/isaaclab_newton/isaaclab_newton/assets/articulation/articulation.py @@ -30,7 +30,7 @@ from isaaclab_newton.actuators import kernels as actuator_kernels from isaaclab_newton.actuators.newton_actuator_utils import ( NewtonActuatorAdapter, - build_actuator_telemetry, + build_implicit_dof_mask, ) _HAS_NEWTON_ACTUATORS = True @@ -295,11 +295,12 @@ def write_data_to_sim(self): if self._has_newton_actuators: # Raw targets go directly to Newton's control object. Newton PD - # consumes them for Newton-managed joints; the solver's built-in - # PD uses them for any implicit joints (whose stiffness/damping - # are non-zero in sim). - # Feedforward effort goes to control.joint_act (not joint_f, - # which is the actuator output buffer that Newton overwrites). + # consumes them for explicit (Newton-managed) joints; the + # solver's built-in joint drive does the PD for implicit joints + # (whose stiffness/damping are non-zero in sim). FF for explicit + # joints goes to control.joint_act; for implicit joints the + # in-graph post-actuator kernel writes the FF directly into + # control.joint_f. self.data._sim_bind_joint_position_target.assign(self._data._joint_pos_target) self.data._sim_bind_joint_velocity_target.assign(self._data._joint_vel_target) self.data._sim_bind_joint_act.assign(self._data._joint_effort_target) @@ -318,30 +319,6 @@ def update(self, dt: float): dt: The time step size in seconds. """ self.data.update(dt) - # Shadow PD telemetry for implicit DOFs (the simulator runs the real PD). - if self._has_newton_actuators and self._telemetry_indices is not None: - wp.launch( - actuator_kernels.compute_actuator_telemetry, - dim=(self.num_instances, self._telemetry_indices.shape[0]), - inputs=[ - self._data.joint_pos.warp, - self._data.joint_vel.warp, - self._data._joint_pos_target, - self._data._joint_vel_target, - self._data._joint_effort_target, - self._data._sim_bind_joint_effort, - self._data.joint_stiffness.warp, - self._data.joint_damping.warp, - self._telemetry_effort_limit, - self._telemetry_indices, - self._telemetry_modes, - ], - outputs=[ - self._data._computed_torque, - self._data._applied_torque, - ], - device=self.device, - ) """ Operations - Finders. @@ -3435,10 +3412,10 @@ def _process_actuators_cfg(self): # if this is false, we by-pass certain checks when doing actuator-related operations self._has_implicit_actuators = False self._has_newton_actuators = False - # Per-DOF telemetry tables; ``None`` when no Newton fast path is active. - self._telemetry_indices: wp.array | None = None - self._telemetry_modes: wp.array | None = None - self._telemetry_effort_limit: wp.array | None = None + # Per-DOF implicit/explicit mask consumed by the in-graph kernel + # ``synch_torque_and_apply_implicit_feedforwards``. ``None`` when + # no Newton fast path is active. + self._implicit_dof_mask: wp.array | None = None _use_newton_actuators = getattr(self._sim_cfg, "use_newton_actuators", False) @@ -3488,14 +3465,38 @@ def _process_actuators_cfg(self): else: self._create_lab_actuator(actuator_name, actuator_cfg, properties_only=True) - ( - self._telemetry_indices, - self._telemetry_modes, - self._telemetry_effort_limit, - ) = build_actuator_telemetry( - self.actuators, self.num_instances, self.num_joints, self.device, + self._implicit_dof_mask = build_implicit_dof_mask( + self.actuators, self.num_joints, self.device, ) + # Run the implicit-DOF FF-routing + telemetry kernel inside the + # captured graph, right after the actuator step. Closure captures + # the buffers we need via ``self._data``. + def _post_actuator() -> None: + wp.launch( + actuator_kernels.synch_torque_and_apply_implicit_feedforwards, + dim=(self.num_instances, self.num_joints), + inputs=[ + self._data.joint_pos.warp, + self._data.joint_vel.warp, + self._data._joint_pos_target, + self._data._joint_vel_target, + self._data._joint_effort_target, + self._data.joint_stiffness.warp, + self._data.joint_damping.warp, + self._data.joint_effort_limits.warp, + self._implicit_dof_mask, + ], + outputs=[ + self._data._sim_bind_joint_effort, + self._data._computed_torque, + self._data._applied_torque, + ], + device=self.device, + ) + + SimulationManager.register_post_actuator_callback(_post_actuator) + return # --- Standard Isaac Lab actuator path --- diff --git a/source/isaaclab_newton/isaaclab_newton/physics/newton_manager.py b/source/isaaclab_newton/isaaclab_newton/physics/newton_manager.py index 01eec0026a1e..fc96a3f219d6 100644 --- a/source/isaaclab_newton/isaaclab_newton/physics/newton_manager.py +++ b/source/isaaclab_newton/isaaclab_newton/physics/newton_manager.py @@ -11,7 +11,8 @@ import ctypes import inspect import logging -from typing import TYPE_CHECKING +from collections.abc import Callable +from typing import TYPE_CHECKING, Any import numpy as np import warp as wp @@ -148,6 +149,10 @@ class NewtonManager(PhysicsManager): # Newton actuator adapter (owns actuators and double-buffered states) _adapter: NewtonActuatorAdapter | None = None _decimation: int = 1 + # Optional in-graph hook invoked after the actuator step and before the + # solver substeps. Used by the articulation's implicit-DOF telemetry / + # FF-routing kernel (see isaaclab_newton.assets.articulation). + _post_actuator_callback: Callable[[], None] | None = None # CUDA graphing _graph = None @@ -452,6 +457,8 @@ def step(cls) -> None: # --- Some actuators not graph-safe: step them eagerly, graph solver only --- if cls._adapter is not None: cls._adapter.step(cls._state_0, cls._control, physics_dt) + if cls._post_actuator_callback is not None: + cls._post_actuator_callback() if use_graph: wp.capture_launch(cls._graph) @@ -517,6 +524,7 @@ def clear(cls): cls._newton_imu_sensors = [] cls._report_contacts = False cls._adapter = None + cls._post_actuator_callback = None # Set by an articulation that took the ``use_newton_actuators=True`` # branch in ``_process_actuators_cfg``. Together with the adapter # check, this gates whether the decimation loop can be captured into @@ -1286,6 +1294,8 @@ def _simulate_full(cls) -> None: if cls._adapter is not None: cls._adapter.step(cls._state_0, cls._control, physics_dt) + if cls._post_actuator_callback is not None: + cls._post_actuator_callback() cls._run_solver_substeps(contacts) @@ -1388,6 +1398,23 @@ def register_adapter(cls, adapter: NewtonActuatorAdapter) -> None: """ cls._adapter = adapter + @classmethod + def register_post_actuator_callback( + cls, callback: Callable[[], None] | None, + ) -> None: + """Register a hook invoked after the actuator step on every iteration. + + The callback runs inside the captured CUDA graph (when + :meth:`_is_all_graphable` is ``True``) right after + :meth:`NewtonActuatorAdapter.step` and before the solver substeps, + so kernel writes to ``state``/``control`` are visible to the + integrator on the same iteration. Used by the articulation to + run its implicit-DOF telemetry / FF-routing kernel. + + Pass ``None`` to clear the registration. + """ + cls._post_actuator_callback = callback + @classmethod def set_decimation(cls, decimation: int) -> None: """Set the decimation count and re-capture the CUDA graph. diff --git a/source/isaaclab_newton/test/assets/test_newton_actuators_newton.py b/source/isaaclab_newton/test/assets/test_newton_actuators_newton.py index 7f39e894e455..263796ddf1b7 100644 --- a/source/isaaclab_newton/test/assets/test_newton_actuators_newton.py +++ b/source/isaaclab_newton/test/assets/test_newton_actuators_newton.py @@ -452,47 +452,40 @@ def test_computed_equals_applied_explicit(self): class TestImplicitOnlyTelemetry(unittest.TestCase): - """Implicit-only fast path: shadow-PD telemetry matches the Lab formula.""" + """Implicit-only fast path: shadow-PD telemetry matches the Lab actuator path.""" @classmethod def setUpClass(cls): - cls.result = _run_simulation_with_telemetry( + cls.lab_result = _run_simulation_with_telemetry( + IMPLICIT_ONLY_ACTUATORS, use_newton_actuators=False, + ) + cls.newton_result = _run_simulation_with_telemetry( IMPLICIT_ONLY_ACTUATORS, use_newton_actuators=True, ) - cls.kp = 40.0 - cls.kd = 5.0 - cls.target_offset = TARGET_OFFSET def test_telemetry_is_nonzero(self): """Telemetry kernel actually fired (not stuck at default zeros).""" - last = self.result["computed_torque"][-1] + last = self.newton_result["computed_torque"][-1] max_abs = last.abs().max().item() self.assertGreater(max_abs, 1e-3, "computed_torque is all-zero — telemetry kernel did not run") - def test_telemetry_matches_pd_formula(self): - """Telemetry value matches kp*(q_des-q) + kd*(v_des-v) within tolerance.""" - target_q = self.result["target_pos"] - target_v = self.result["target_vel"] - for step_i, (q, qd, comp) in enumerate( - zip( - self.result["joint_pos"], - self.result["joint_vel"], - self.result["computed_torque"], - ) + def test_telemetry_matches_lab_path(self): + """Newton fast-path telemetry agrees with the Lab actuator-path telemetry.""" + for step_i, (lab_comp, newton_comp) in enumerate( + zip(self.lab_result["computed_torque"], self.newton_result["computed_torque"]) ): - expected = self.kp * (target_q - q) + self.kd * (target_v - qd) torch.testing.assert_close( - comp, - expected, + newton_comp, + lab_comp, atol=5e-2, rtol=1e-2, - msg=f"Telemetry diverged from PD formula at step {step_i}", + msg=f"computed_torque diverged from Lab path at step {step_i}", ) def test_applied_equals_computed_when_no_clip(self): """Implicit cfg has no effort_limit → applied == computed.""" for step_i, (comp, app) in enumerate( - zip(self.result["computed_torque"], self.result["applied_torque"]) + zip(self.newton_result["computed_torque"], self.newton_result["applied_torque"]) ): torch.testing.assert_close( app, @@ -503,6 +496,117 @@ def test_applied_equals_computed_when_no_clip(self): ) +# --------------------------------------------------------------------------- +# Implicit + non-zero feedforward effort target +# --------------------------------------------------------------------------- + + +def _run_simulation_with_ff( + actuators: dict, + use_newton_actuators: bool, + *, + feedforward: float, + dt: float = DT, + newton_cfg: NewtonCfg = NEWTON_CFG, + num_steps: int = NUM_STEPS, +) -> dict: + """Run with both a position target and a non-zero feedforward effort target.""" + sim_cfg = SimulationCfg( + dt=dt, + physics=newton_cfg, + use_newton_actuators=use_newton_actuators, + ) + + with build_simulation_context( + device="cuda:0", + gravity_enabled=True, + add_ground_plane=True, + sim_cfg=sim_cfg, + ) as sim: + sim._app_control_on_stop_handle = None + + for i in range(NUM_ENVS): + sim_utils.create_prim( + f"/World/Env_{i}", "Xform", translation=(i * 3.0, 0, 0) + ) + + art_cfg = ANYMAL_C_CFG.replace( + actuators=actuators, + prim_path="/World/Env_.*/Robot", + ) + articulation = Articulation(art_cfg) + sim.reset() + assert articulation.is_initialized + + init_pos = wp.to_torch(articulation.data.joint_pos).clone() + target_pos = init_pos + TARGET_OFFSET + target_vel = torch.zeros_like(init_pos) + target_eff = torch.full_like(init_pos, feedforward) + + articulation.set_joint_position_target_index(target=target_pos) + articulation.set_joint_velocity_target_index(target=target_vel) + articulation.set_joint_effort_target_index(target=target_eff) + + recorded_pos, recorded_vel = [], [] + for _ in range(num_steps): + articulation.write_data_to_sim() + sim.step() + articulation.update(dt) + recorded_pos.append(wp.to_torch(articulation.data.joint_pos).clone()) + recorded_vel.append(wp.to_torch(articulation.data.joint_vel).clone()) + + return {"joint_pos": recorded_pos, "joint_vel": recorded_vel} + + +class TestImplicitWithFeedforwardEquivalence(unittest.TestCase): + """Implicit-only actuators with a non-zero feedforward effort target. + + For implicit actuators, the user's feedforward effort must be additive + on top of the simulator's PD on both the Lab path and the Newton fast + path. This test commands a constant FF effort and compares joint + trajectories between the two paths. + """ + + FEEDFORWARD = 5.0 + pos_atol = 2e-3 + pos_rtol = 1e-3 + vel_atol = 0.1 + vel_rtol = 1e-2 + + @classmethod + def setUpClass(cls): + cls.lab_result = _run_simulation_with_ff( + IMPLICIT_ONLY_ACTUATORS, use_newton_actuators=False, feedforward=cls.FEEDFORWARD, + ) + cls.newton_result = _run_simulation_with_ff( + IMPLICIT_ONLY_ACTUATORS, use_newton_actuators=True, feedforward=cls.FEEDFORWARD, + ) + + def test_joint_positions_match(self): + for step_i, (lab, newton) in enumerate( + zip(self.lab_result["joint_pos"], self.newton_result["joint_pos"]) + ): + torch.testing.assert_close( + lab, newton, atol=self.pos_atol, rtol=self.pos_rtol, + msg=f"Joint positions diverged at step {step_i}", + ) + + def test_joint_velocities_match(self): + for step_i, (lab, newton) in enumerate( + zip(self.lab_result["joint_vel"], self.newton_result["joint_vel"]) + ): + torch.testing.assert_close( + lab, newton, atol=self.vel_atol, rtol=self.vel_rtol, + msg=f"Joint velocities diverged at step {step_i}", + ) + + def test_trajectories_not_trivial(self): + first = self.lab_result["joint_pos"][0] + last = self.lab_result["joint_pos"][-1] + diff = (last - first).abs().max().item() + self.assertGreater(diff, 0.01, "Joints did not move — test is trivial") + + # --------------------------------------------------------------------------- # DelayedPD equivalence: PD with actuator command delay # --------------------------------------------------------------------------- diff --git a/source/isaaclab_physx/isaaclab_physx/assets/articulation/articulation.py b/source/isaaclab_physx/isaaclab_physx/assets/articulation/articulation.py index 84d6dc130c10..2eb09206a89e 100644 --- a/source/isaaclab_physx/isaaclab_physx/assets/articulation/articulation.py +++ b/source/isaaclab_physx/isaaclab_physx/assets/articulation/articulation.py @@ -29,7 +29,7 @@ NewtonActuatorAdapter, PhysxActuatorWrapper, _gather_gain_kernel, - build_actuator_telemetry, + build_implicit_dof_mask, ) _HAS_NEWTON_ACTUATORS = True @@ -272,14 +272,14 @@ def write_data_to_sim(self): if self._has_newton_actuators: # Newton fast path: pos/vel targets pass straight through; the - # actuation buffer is either user FF (implicit-only) or the - # combine-kernel result of Newton output + user FF (mixed). - if self._physx_actuator_wrapper is not None: - self._apply_actuator_model_newton() - actuation_buf = self._joint_effort_target_sim - else: - actuation_buf = self._data._joint_effort_target - self.root_view.set_dof_actuation_forces(actuation_buf, self._ALL_INDICES) + # in-graph kernel inside ``_apply_actuator_model_newton`` merges + # Newton's actuator output (explicit DOFs) with user FF + # (implicit DOFs) into ``w.joint_f_2d``, which is what we push + # to PhysX as the actuation force. + self._apply_actuator_model_newton() + self.root_view.set_dof_actuation_forces( + self._physx_actuator_wrapper.joint_f_2d, self._ALL_INDICES, + ) if self._has_implicit_actuators: self.root_view.set_dof_position_targets(self._data._joint_pos_target, self._ALL_INDICES) self.root_view.set_dof_velocity_targets(self._data._joint_vel_target, self._ALL_INDICES) @@ -299,30 +299,6 @@ def update(self, dt: float): dt: The time step size in seconds. """ self.data.update(dt) - # Shadow PD telemetry for implicit DOFs (the simulator runs the real PD). - if self._has_newton_actuators and self._telemetry_indices is not None: - wp.launch( - actuator_kernels.compute_actuator_telemetry, - dim=(self.num_instances, self._telemetry_indices.shape[0]), - inputs=[ - self._data.joint_pos.warp, - self._data.joint_vel.warp, - self._data._joint_pos_target, - self._data._joint_vel_target, - self._data._joint_effort_target, - self._joint_effort_target_sim, - self._data.joint_stiffness.warp, - self._data.joint_damping.warp, - self._telemetry_effort_limit, - self._telemetry_indices, - self._telemetry_modes, - ], - outputs=[ - self._data._computed_torque, - self._data._applied_torque, - ], - device=self.device, - ) """ Operations - Finders. @@ -3850,10 +3826,10 @@ def _process_actuators_cfg(self): # if this is false, we by-pass certain checks when doing actuator-related operations self._has_implicit_actuators = False self._has_newton_actuators = False - # Per-DOF telemetry tables; ``None`` when no Newton fast path is active. - self._telemetry_indices: wp.array | None = None - self._telemetry_modes: wp.array | None = None - self._telemetry_effort_limit: wp.array | None = None + # Per-DOF implicit/explicit mask consumed by the + # ``synch_torque_and_apply_implicit_feedforwards`` kernel. ``None`` + # when no Newton fast path is active. + self._implicit_dof_mask: wp.array | None = None _use_newton_actuators = getattr(self._sim_cfg, "use_newton_actuators", False) @@ -3874,6 +3850,15 @@ def _process_actuators_cfg(self): for actuator_cfg in self.cfg.actuators.values() ) + # Always allocate the wrapper so ``_apply_actuator_model_newton`` + # has a ``joint_f_2d`` buffer to merge effort into, even when + # there are no explicit Newton actuators (implicit-only case). + self._physx_actuator_wrapper = PhysxActuatorWrapper.create( + num_envs=self.num_instances, + num_joints=self.num_joints, + device=self.device, + ) + if has_explicit: first_prim = find_first_matching_prim(self.cfg.prim_path) art_prim_path = str(first_prim.GetPath()) if first_prim is not None else None @@ -3887,11 +3872,6 @@ def _process_actuators_cfg(self): articulation_prim_path=art_prim_path, ) - self._physx_actuator_wrapper = PhysxActuatorWrapper.create( - num_envs=self.num_instances, - num_joints=self.num_joints, - device=self.device, - ) # Bind the wrapper's flat aliases of state/input buffers once. # The underlying wp.arrays alias stable PhysX-owned GPU memory # whose device pointer is fixed for the articulation's lifetime, @@ -3919,12 +3899,8 @@ def _process_actuators_cfg(self): else: self._create_lab_actuator(actuator_name, actuator_cfg, properties_only=True) - ( - self._telemetry_indices, - self._telemetry_modes, - self._telemetry_effort_limit, - ) = build_actuator_telemetry( - self.actuators, self.num_instances, self.num_joints, self.device, + self._implicit_dof_mask = build_implicit_dof_mask( + self.actuators, self.num_joints, self.device, ) return @@ -4140,31 +4116,41 @@ def _apply_actuator_model(self): ) def _apply_actuator_model_newton(self): - """Mixed implicit + Newton path: step Newton actuators, then combine with user FF. - - Newton writes explicit-DOF torques to ``joint_f_2d``; the combine kernel - merges them with user FF (for implicit DOFs) into ``_joint_effort_target_sim``. - Implicit-only takes the direct path in :meth:`write_data_to_sim`. + """Step Newton actuators (when present) then route FF + sync telemetry. + + After ``newton_adapter.step`` (no-op if no explicit Newton actuators + exist), ``w.joint_f_2d`` holds the actuator output for explicit DOFs + and zero elsewhere. The + :func:`synch_torque_and_apply_implicit_feedforwards` kernel then + writes the user FF into ``joint_f_2d`` for implicit DOFs and fills + ``_data._computed_torque`` / ``_data._applied_torque``. The + resulting ``joint_f_2d`` is what gets pushed to PhysX as the + actuation force in :meth:`write_data_to_sim`. """ - newton_adapter = self.actuators["newton"] w = self._physx_actuator_wrapper - - # Newton actuators scatter-add into joint_f at their own indices; - # ``adapter.step`` pre-zeros exactly those slots. Slots outside any - # actuator's indices are implicit DOFs that the combine kernel reads - # from ``user_ff`` instead of ``joint_f_2d``, so they don't need - # clearing here. - newton_adapter.step(w, w, SimulationManager.get_physics_dt()) + newton_adapter = self.actuators.get("newton") + if newton_adapter is not None: + newton_adapter.step(w, w, SimulationManager.get_physics_dt()) wp.launch( - actuator_kernels.combine_actuation_force, + actuator_kernels.synch_torque_and_apply_implicit_feedforwards, dim=(self.num_instances, self.num_joints), inputs=[ + self._data.joint_pos.warp, + self._data.joint_vel.warp, + self._data._joint_pos_target, + self._data._joint_vel_target, self._data._joint_effort_target, + self._data.joint_stiffness.warp, + self._data.joint_damping.warp, + self._data.joint_effort_limits.warp, + self._implicit_dof_mask, + ], + outputs=[ w.joint_f_2d, - self._telemetry_modes, + self._data._computed_torque, + self._data._applied_torque, ], - outputs=[self._joint_effort_target_sim], device=self.device, ) diff --git a/source/isaaclab_physx/test/assets/test_newton_actuators_physx.py b/source/isaaclab_physx/test/assets/test_newton_actuators_physx.py index 6ea74df0972e..8feb9916ea5f 100644 --- a/source/isaaclab_physx/test/assets/test_newton_actuators_physx.py +++ b/source/isaaclab_physx/test/assets/test_newton_actuators_physx.py @@ -373,46 +373,41 @@ def _run_simulation_with_telemetry( class TestImplicitOnlyTelemetryPhysx(unittest.TestCase): - """Implicit-only fast path on PhysX: shadow-PD telemetry matches the Lab formula.""" + """Implicit-only fast path on PhysX: shadow-PD telemetry matches the Lab actuator path.""" @classmethod def setUpClass(cls): - cls.result = _run_simulation_with_telemetry( + cls.lab_result = _run_simulation_with_telemetry( + IMPLICIT_ONLY_ACTUATORS, use_newton_actuators=False, + ) + cls.newton_result = _run_simulation_with_telemetry( IMPLICIT_ONLY_ACTUATORS, use_newton_actuators=True, ) - cls.kp = 40.0 - cls.kd = 5.0 def test_telemetry_is_nonzero(self): - last = self.result["computed_torque"][-1] + last = self.newton_result["computed_torque"][-1] self.assertGreater( last.abs().max().item(), 1e-3, "computed_torque is all-zero — implicit telemetry kernel did not run", ) - def test_telemetry_matches_pd_formula(self): - target_q = self.result["target_pos"] - target_v = self.result["target_vel"] - for step_i, (q, qd, comp) in enumerate( - zip( - self.result["joint_pos"], - self.result["joint_vel"], - self.result["computed_torque"], - ) + def test_telemetry_matches_lab_path(self): + """Newton fast-path telemetry agrees with the Lab actuator-path telemetry.""" + for step_i, (lab_comp, newton_comp) in enumerate( + zip(self.lab_result["computed_torque"], self.newton_result["computed_torque"]) ): - expected = self.kp * (target_q - q) + self.kd * (target_v - qd) torch.testing.assert_close( - comp, - expected, + newton_comp, + lab_comp, atol=5e-2, rtol=1e-2, - msg=f"Telemetry diverged from PD formula at step {step_i}", + msg=f"computed_torque diverged from Lab path at step {step_i}", ) def test_applied_equals_computed_when_no_clip(self): for step_i, (comp, app) in enumerate( - zip(self.result["computed_torque"], self.result["applied_torque"]) + zip(self.newton_result["computed_torque"], self.newton_result["applied_torque"]) ): torch.testing.assert_close( app, @@ -423,6 +418,115 @@ def test_applied_equals_computed_when_no_clip(self): ) +# --------------------------------------------------------------------------- +# Implicit + non-zero feedforward effort target +# --------------------------------------------------------------------------- + + +def _run_simulation_with_ff( + actuators: dict, + use_newton_actuators: bool, + *, + feedforward: float, + num_steps: int = NUM_STEPS, +) -> dict: + """Run with both a position target and a non-zero feedforward effort target.""" + sim_cfg = SimulationCfg( + dt=DT, + physics=PhysxCfg(), + use_newton_actuators=use_newton_actuators, + ) + + with build_simulation_context( + device="cuda:0", + gravity_enabled=True, + add_ground_plane=True, + sim_cfg=sim_cfg, + ) as sim: + sim._app_control_on_stop_handle = None + + for i in range(NUM_ENVS): + sim_utils.create_prim( + f"/World/Env_{i}", "Xform", translation=(i * 3.0, 0, 0) + ) + + art_cfg = ANYMAL_C_CFG.replace( + actuators=actuators, + prim_path="/World/Env_.*/Robot", + ) + articulation = Articulation(art_cfg) + sim.reset() + assert articulation.is_initialized + + init_pos = wp.to_torch(articulation.data.joint_pos).clone() + target_pos = init_pos + TARGET_OFFSET + target_vel = torch.zeros_like(init_pos) + target_eff = torch.full_like(init_pos, feedforward) + + articulation.set_joint_position_target_index(target=target_pos) + articulation.set_joint_velocity_target_index(target=target_vel) + articulation.set_joint_effort_target_index(target=target_eff) + + recorded_pos, recorded_vel = [], [] + for _ in range(num_steps): + articulation.write_data_to_sim() + sim.step() + articulation.update(DT) + recorded_pos.append(wp.to_torch(articulation.data.joint_pos).clone()) + recorded_vel.append(wp.to_torch(articulation.data.joint_vel).clone()) + + return {"joint_pos": recorded_pos, "joint_vel": recorded_vel} + + +class TestImplicitWithFeedforwardEquivalencePhysx(unittest.TestCase): + """Implicit-only actuators with a non-zero feedforward effort target on PhysX. + + For implicit actuators, the user's feedforward effort must be additive + on top of the simulator's PD on both the Lab path and the Newton fast + path. This test commands a constant FF effort and compares joint + trajectories between the two paths. + """ + + FEEDFORWARD = 5.0 + pos_atol = 2e-3 + pos_rtol = 1e-3 + vel_atol = 0.1 + vel_rtol = 1e-2 + + @classmethod + def setUpClass(cls): + cls.lab_result = _run_simulation_with_ff( + IMPLICIT_ONLY_ACTUATORS, use_newton_actuators=False, feedforward=cls.FEEDFORWARD, + ) + cls.newton_result = _run_simulation_with_ff( + IMPLICIT_ONLY_ACTUATORS, use_newton_actuators=True, feedforward=cls.FEEDFORWARD, + ) + + def test_joint_positions_match(self): + for step_i, (lab, newton) in enumerate( + zip(self.lab_result["joint_pos"], self.newton_result["joint_pos"]) + ): + torch.testing.assert_close( + lab, newton, atol=self.pos_atol, rtol=self.pos_rtol, + msg=f"Joint positions diverged at step {step_i}", + ) + + def test_joint_velocities_match(self): + for step_i, (lab, newton) in enumerate( + zip(self.lab_result["joint_vel"], self.newton_result["joint_vel"]) + ): + torch.testing.assert_close( + lab, newton, atol=self.vel_atol, rtol=self.vel_rtol, + msg=f"Joint velocities diverged at step {step_i}", + ) + + def test_trajectories_not_trivial(self): + first = self.lab_result["joint_pos"][0] + last = self.lab_result["joint_pos"][-1] + diff = (last - first).abs().max().item() + self.assertGreater(diff, 0.01, "Joints did not move — test is trivial") + + class TestExplicitOnlyTelemetryPhysx(unittest.TestCase): """Explicit-only Newton actuators on PhysX: telemetry copies from the staging effort buffer.""" From 2bd037b5f143b104c125202ab13497c7787282eb Mon Sep 17 00:00:00 2001 From: Julia von Muralt Date: Fri, 8 May 2026 00:13:28 +0200 Subject: [PATCH 15/35] refactor for multi arti --- source/isaaclab/isaaclab/envs/mdp/events.py | 48 +++++++-- .../assets/articulation/articulation.py | 67 +++++-------- .../isaaclab_newton/physics/newton_manager.py | 76 ++++++++------ .../assets/test_newton_actuators_newton.py | 96 ++++++++++++++++++ .../assets/articulation/articulation.py | 15 ++- .../assets/test_newton_actuators_physx.py | 98 ++++++++++++++++++- 6 files changed, 314 insertions(+), 86 deletions(-) diff --git a/source/isaaclab/isaaclab/envs/mdp/events.py b/source/isaaclab/isaaclab/envs/mdp/events.py index 5c788b1c1d65..18f5816a974a 100644 --- a/source/isaaclab/isaaclab/envs/mdp/events.py +++ b/source/isaaclab/isaaclab/envs/mdp/events.py @@ -1210,10 +1210,6 @@ def randomize(data: torch.Tensor, params: tuple[float, float]) -> torch.Tensor: self.asset.write_joint_stiffness_to_sim_index( stiffness=stiffness, joint_ids=actuator.joint_indices, env_ids=env_ids ) - elif hasattr(self.asset, "write_actuator_stiffness_to_sim"): - self.asset.write_actuator_stiffness_to_sim( - actuator, stiffness=stiffness, env_ids=env_ids - ) # Randomize damping if damping_distribution_params is not None: damping = actuator.damping[env_ids].clone() @@ -1224,10 +1220,46 @@ def randomize(data: torch.Tensor, params: tuple[float, float]) -> torch.Tensor: self.asset.write_joint_damping_to_sim_index( damping=damping, joint_ids=actuator.joint_indices, env_ids=env_ids ) - elif hasattr(self.asset, "write_actuator_damping_to_sim"): - self.asset.write_actuator_damping_to_sim( - actuator, damping=damping, env_ids=env_ids - ) + + # Direct write to the global Newton actuator adapter + adapter = getattr(env.sim.physics_manager, "_adapter", None) + if adapter is None: + return + from newton import Model as NewtonModel # noqa: PLC0415 + + dof_layout = self.asset._root_view.frequency_layouts[NewtonModel.AttributeFrequency.JOINT_DOF] + if dof_layout.slice is not None: + dof_offset = dof_layout.slice.start + elif dof_layout.indices is not None: + dof_offset = int(dof_layout.indices.numpy()[0]) + else: + dof_offset = 0 + if isinstance(self.asset_cfg.joint_ids, slice): + arti_local_ids = torch.arange(self.asset.num_joints, device=self.asset.device, dtype=torch.long) + else: + arti_local_ids = torch.tensor(self.asset_cfg.joint_ids, device=self.asset.device, dtype=torch.long) + # Column indices into the adapter's model-wide gain buffer for the + # joints this DR term should randomize. + adapter_indices = arti_local_ids + dof_offset + env_ids_wp = self.asset._resolve_env_ids(env_ids) + env_mask = self.asset._env_ids_to_mask(env_ids_wp) + + def _randomize_at(data: torch.Tensor, params: tuple[float, float]) -> None: + _randomize_prop_by_op( + data, params, dim_0_ids=None, dim_1_ids=adapter_indices, + operation=operation, distribution=distribution, + ) + + if stiffness_distribution_params is not None: + stiffness = adapter.stiffness[env_ids].clone() + stiffness[:, adapter_indices] = self.default_joint_stiffness[env_ids][:, arti_local_ids].clone() + _randomize_at(stiffness, stiffness_distribution_params) + adapter.write_stiffness_to_sim(stiffness, env_ids_wp, env_mask, self.asset._propagate_gain_via_view) + if damping_distribution_params is not None: + damping = adapter.damping[env_ids].clone() + damping[:, adapter_indices] = self.default_joint_damping[env_ids][:, arti_local_ids].clone() + _randomize_at(damping, damping_distribution_params) + adapter.write_damping_to_sim(damping, env_ids_wp, env_mask, self.asset._propagate_gain_via_view) class randomize_joint_parameters(ManagerTermBase): diff --git a/source/isaaclab_newton/isaaclab_newton/assets/articulation/articulation.py b/source/isaaclab_newton/isaaclab_newton/assets/articulation/articulation.py index 1e8069501e77..fbabbcfa49fd 100644 --- a/source/isaaclab_newton/isaaclab_newton/assets/articulation/articulation.py +++ b/source/isaaclab_newton/isaaclab_newton/assets/articulation/articulation.py @@ -2034,30 +2034,6 @@ def write_joint_friction_coefficient_to_sim_mask( Operations - Newton Actuator Parameter Writers. """ - def write_actuator_stiffness_to_sim( - self, - adapter: NewtonActuatorAdapter, - *, - stiffness: torch.Tensor | wp.array | float, - env_ids: Sequence[int] | torch.Tensor | wp.array | None = None, - ) -> None: - """Write actuator stiffness (``kp``) into each Newton controller for *env_ids*.""" - env_ids = self._resolve_env_ids(env_ids) - env_mask = self._env_ids_to_mask(env_ids) - adapter.write_stiffness_to_sim(stiffness, env_ids, env_mask, self._propagate_gain_via_view) - - def write_actuator_damping_to_sim( - self, - adapter: NewtonActuatorAdapter, - *, - damping: torch.Tensor | wp.array | float, - env_ids: Sequence[int] | torch.Tensor | wp.array | None = None, - ) -> None: - """Write actuator damping (``kd``) into each Newton controller for *env_ids*.""" - env_ids = self._resolve_env_ids(env_ids) - env_mask = self._env_ids_to_mask(env_ids) - adapter.write_damping_to_sim(damping, env_ids, env_mask, self._propagate_gain_via_view) - def _propagate_gain_via_view( self, adapter: NewtonActuatorAdapter, @@ -3428,30 +3404,37 @@ def _process_actuators_cfg(self): if _use_newton_actuators and _HAS_NEWTON_ACTUATORS: from newton import Model as NewtonModel # noqa: PLC0415 - model = SimulationManager.get_model() - # Enable the fast path even for all-implicit articulations: # the solver runs PD internally; Lab only forwards targets. self._has_newton_actuators = True SimulationManager.activate_newton_actuator_path() - if model.actuators: - dof_layout = self._root_view.frequency_layouts[NewtonModel.AttributeFrequency.JOINT_DOF] - if dof_layout.slice is not None: - dof_offset = dof_layout.slice.start - elif dof_layout.indices is not None: - dof_offset = int(dof_layout.indices.numpy()[0]) - else: - dof_offset = 0 - - adapter = NewtonActuatorAdapter( - model.actuators, self.num_instances, self.num_joints, dof_offset, self.device, + # Build (or share) the single sim-level actuator adapter. Idempotent — + # the first articulation to call this constructs it from + # ``model.actuators``; subsequent articulations reuse it. + SimulationManager.ensure_global_actuator_adapter() + + # Zero the simulator's joint-drive PD on DOFs covered by an explicit + # Lab actuator config in *this* articulation. The global Newton + # adapter's actuator step writes their effort to ``joint_f`` + # directly; the joint drive shouldn't add its own PD on top. + explicit_joint_ids: list[int] = [] + for actuator_cfg in self.cfg.actuators.values(): + cls_type = actuator_cfg.class_type + if ( + "ImplicitActuator" in cls_type + if isinstance(cls_type, str) + else issubclass(cls_type, ImplicitActuator) + ): + continue + joint_ids, _ = self.find_joints(actuator_cfg.joint_names_expr) + explicit_joint_ids.extend(int(j) for j in joint_ids) + if explicit_joint_ids: + explicit_ids_t = torch.tensor( + sorted(set(explicit_joint_ids)), dtype=torch.int32, device=self.device, ) - adapter.finalize() - self.actuators["newton"] = adapter - SimulationManager.register_adapter(adapter) - self.write_joint_stiffness_to_sim_index(stiffness=0.0, joint_ids=adapter.joint_indices) - self.write_joint_damping_to_sim_index(damping=0.0, joint_ids=adapter.joint_indices) + self.write_joint_stiffness_to_sim_index(stiffness=0.0, joint_ids=explicit_ids_t) + self.write_joint_damping_to_sim_index(damping=0.0, joint_ids=explicit_ids_t) for actuator_name, actuator_cfg in self.cfg.actuators.items(): cls_type = actuator_cfg.class_type diff --git a/source/isaaclab_newton/isaaclab_newton/physics/newton_manager.py b/source/isaaclab_newton/isaaclab_newton/physics/newton_manager.py index fc96a3f219d6..dff968a306e2 100644 --- a/source/isaaclab_newton/isaaclab_newton/physics/newton_manager.py +++ b/source/isaaclab_newton/isaaclab_newton/physics/newton_manager.py @@ -149,10 +149,10 @@ class NewtonManager(PhysicsManager): # Newton actuator adapter (owns actuators and double-buffered states) _adapter: NewtonActuatorAdapter | None = None _decimation: int = 1 - # Optional in-graph hook invoked after the actuator step and before the - # solver substeps. Used by the articulation's implicit-DOF telemetry / - # FF-routing kernel (see isaaclab_newton.assets.articulation). - _post_actuator_callback: Callable[[], None] | None = None + # In-graph hooks invoked after the actuator step and before the solver + # substeps, in registration order. Multiple articulations register their + # implicit-DOF telemetry / FF-routing kernels here. + _post_actuator_callbacks: list[Callable[[], None]] = [] # CUDA graphing _graph = None @@ -457,8 +457,8 @@ def step(cls) -> None: # --- Some actuators not graph-safe: step them eagerly, graph solver only --- if cls._adapter is not None: cls._adapter.step(cls._state_0, cls._control, physics_dt) - if cls._post_actuator_callback is not None: - cls._post_actuator_callback() + for cb in cls._post_actuator_callbacks: + cb() if use_graph: wp.capture_launch(cls._graph) @@ -524,7 +524,7 @@ def clear(cls): cls._newton_imu_sensors = [] cls._report_contacts = False cls._adapter = None - cls._post_actuator_callback = None + cls._post_actuator_callbacks = [] # Set by an articulation that took the ``use_newton_actuators=True`` # branch in ``_process_actuators_cfg``. Together with the adapter # check, this gates whether the decimation loop can be captured into @@ -832,8 +832,9 @@ def start_simulation(cls) -> None: cls._control = cls._model.control() eval_fk(cls._model, cls._state_0.joint_q, cls._state_0.joint_qd, cls._state_0, None) - # Actuator registration is deferred to register_adapter(), - # called by _process_actuators_cfg after the adapter is constructed. + # The single global actuator adapter is built lazily on first + # call to ``ensure_global_actuator_adapter`` from any + # Newton-fast-path articulation after this point. cls._adapter = None cls._use_newton_actuators_active = False @@ -1294,8 +1295,8 @@ def _simulate_full(cls) -> None: if cls._adapter is not None: cls._adapter.step(cls._state_0, cls._control, physics_dt) - if cls._post_actuator_callback is not None: - cls._post_actuator_callback() + for cb in cls._post_actuator_callbacks: + cb() cls._run_solver_substeps(contacts) @@ -1387,33 +1388,48 @@ def activate_newton_actuator_path(cls) -> None: cls._use_newton_actuators_active = True @classmethod - def register_adapter(cls, adapter: NewtonActuatorAdapter) -> None: - """Register the actuator adapter that owns states and stepping. - - Called by the articulation's ``_process_actuators_cfg`` after - the :class:`NewtonActuatorAdapter` is constructed and finalised. - - Args: - adapter: The actuator adapter instance. + def ensure_global_actuator_adapter(cls) -> None: + """Build the single sim-level :class:`NewtonActuatorAdapter` if not yet built. + + Idempotent — each Newton-fast-path articulation calls this from + ``_process_actuators_cfg`` after the model is finalized; the first + call constructs the adapter from ``cls._model.actuators`` and + subsequent calls are no-ops. A single global adapter avoids the + multi-articulation clobbering that the per-articulation + ``register_adapter`` design suffered from. The adapter operates on + the full flat DOF layout (``dof_offset=0``, ``num_joints`` = + per-env total DOF count), so it iterates and steps every actuator + in the model in one pass. """ - cls._adapter = adapter + if cls._adapter is not None: + return + if cls._model is None or not cls._model.actuators: + return + from isaaclab_newton.actuators.newton_actuator_utils import NewtonActuatorAdapter # noqa: PLC0415 + + dofs_per_env = cls._model.joint_dof_count // cls._num_envs + cls._adapter = NewtonActuatorAdapter( + actuators=list(cls._model.actuators), + num_envs=cls._num_envs, + num_joints=dofs_per_env, + dof_offset=0, + device=PhysicsManager._device, + ) + cls._adapter.finalize() @classmethod - def register_post_actuator_callback( - cls, callback: Callable[[], None] | None, - ) -> None: - """Register a hook invoked after the actuator step on every iteration. + def register_post_actuator_callback(cls, callback: Callable[[], None]) -> None: + """Append a hook to the list invoked after the actuator step on every iteration. - The callback runs inside the captured CUDA graph (when + Each callback runs inside the captured CUDA graph (when :meth:`_is_all_graphable` is ``True``) right after :meth:`NewtonActuatorAdapter.step` and before the solver substeps, so kernel writes to ``state``/``control`` are visible to the - integrator on the same iteration. Used by the articulation to - run its implicit-DOF telemetry / FF-routing kernel. - - Pass ``None`` to clear the registration. + integrator on the same iteration. Multiple articulations register + their own implicit-DOF telemetry / FF-routing kernels here; all + registered callbacks fire in registration order each step. """ - cls._post_actuator_callback = callback + cls._post_actuator_callbacks.append(callback) @classmethod def set_decimation(cls, decimation: int) -> None: diff --git a/source/isaaclab_newton/test/assets/test_newton_actuators_newton.py b/source/isaaclab_newton/test/assets/test_newton_actuators_newton.py index 263796ddf1b7..50818005f6c7 100644 --- a/source/isaaclab_newton/test/assets/test_newton_actuators_newton.py +++ b/source/isaaclab_newton/test/assets/test_newton_actuators_newton.py @@ -607,6 +607,102 @@ def test_trajectories_not_trivial(self): self.assertGreater(diff, 0.01, "Joints did not move — test is trivial") +# --------------------------------------------------------------------------- +# Multi-articulation Newton scene (regression test for class-attr clobber) +# --------------------------------------------------------------------------- + + +CARTPOLE_EXPLICIT_ACTUATORS = { + "all_joints": IdealPDActuatorCfg( + joint_names_expr=["slider_to_cart", "cart_to_pole"], + stiffness=10.0, + damping=1.0, + effort_limit=100.0, + ), +} + + +def _run_anymal_and_cartpole(use_newton_actuators: bool, *, num_steps: int = NUM_STEPS) -> dict: + """Spawn ANYmal-C + Cartpole per env (different DOF counts, different base types).""" + from isaaclab_assets import CARTPOLE_CFG # noqa: PLC0415 + + sim_cfg = SimulationCfg(dt=DT, physics=NEWTON_CFG, use_newton_actuators=use_newton_actuators) + with build_simulation_context( + device="cuda:0", gravity_enabled=True, add_ground_plane=True, sim_cfg=sim_cfg, + ) as sim: + sim._app_control_on_stop_handle = None + + for i in range(NUM_ENVS): + sim_utils.create_prim(f"/World/Env_{i}", "Xform", translation=(i * 6.0, 0, 0)) + + anymal_cfg = ANYMAL_C_CFG.replace(actuators=IDEAL_PD_ACTUATORS, prim_path="/World/Env_.*/Anymal") + cartpole_cfg = CARTPOLE_CFG.replace( + actuators=CARTPOLE_EXPLICIT_ACTUATORS, prim_path="/World/Env_.*/Cartpole", + ) + # Stand the cartpole well clear of the anymal. + cartpole_cfg.init_state = cartpole_cfg.init_state.replace(pos=(0.0, 3.0, 2.0)) + + anymal = Articulation(anymal_cfg) + cartpole = Articulation(cartpole_cfg) + sim.reset() + assert anymal.is_initialized and cartpole.is_initialized + + init_anymal = wp.to_torch(anymal.data.joint_pos).clone() + init_cartpole = wp.to_torch(cartpole.data.joint_pos).clone() + anymal.set_joint_position_target_index(target=init_anymal + TARGET_OFFSET) + anymal.set_joint_velocity_target_index(target=torch.zeros_like(init_anymal)) + cartpole.set_joint_position_target_index(target=init_cartpole + TARGET_OFFSET) + cartpole.set_joint_velocity_target_index(target=torch.zeros_like(init_cartpole)) + + pos_anymal, pos_cartpole = [], [] + for _ in range(num_steps): + anymal.write_data_to_sim() + cartpole.write_data_to_sim() + sim.step() + anymal.update(DT) + cartpole.update(DT) + pos_anymal.append(wp.to_torch(anymal.data.joint_pos).clone()) + pos_cartpole.append(wp.to_torch(cartpole.data.joint_pos).clone()) + + return {"joint_pos_anymal": pos_anymal, "joint_pos_cartpole": pos_cartpole} + + +class TestHeterogeneousMultiArticulationNewton(unittest.TestCase): + """Two structurally-different articulations (ANYmal floating + Cartpole fixed) on Newton. + + Regression for the singleton-clobber bug in ``NewtonManager._adapter`` + / ``_post_actuator_callback`` — fixed by the global-adapter refactor + + callback-list multiplexing. Heterogeneous DOF counts (12 vs 2) and + base types (floating vs fixed) stress the global adapter's handling + of varied actuator index patterns. Equivalence against the Lab + actuator path is the meaningful end-to-end check: divergence on + either robot would indicate broken stepping. + """ + + @classmethod + def setUpClass(cls): + cls.lab_result = _run_anymal_and_cartpole(use_newton_actuators=False) + cls.newton_result = _run_anymal_and_cartpole(use_newton_actuators=True) + + def test_anymal_matches_lab(self): + for step_i, (lab, newton) in enumerate( + zip(self.lab_result["joint_pos_anymal"], self.newton_result["joint_pos_anymal"]) + ): + torch.testing.assert_close( + newton, lab, atol=2e-3, rtol=1e-3, + msg=f"ANYmal joint_pos diverged from Lab path at step {step_i}", + ) + + def test_cartpole_matches_lab(self): + for step_i, (lab, newton) in enumerate( + zip(self.lab_result["joint_pos_cartpole"], self.newton_result["joint_pos_cartpole"]) + ): + torch.testing.assert_close( + newton, lab, atol=2e-3, rtol=1e-3, + msg=f"Cartpole joint_pos diverged from Lab path at step {step_i}", + ) + + # --------------------------------------------------------------------------- # DelayedPD equivalence: PD with actuator command delay # --------------------------------------------------------------------------- diff --git a/source/isaaclab_physx/isaaclab_physx/assets/articulation/articulation.py b/source/isaaclab_physx/isaaclab_physx/assets/articulation/articulation.py index 2eb09206a89e..80f60b3e04de 100644 --- a/source/isaaclab_physx/isaaclab_physx/assets/articulation/articulation.py +++ b/source/isaaclab_physx/isaaclab_physx/assets/articulation/articulation.py @@ -1212,24 +1212,31 @@ def write_joint_damping_to_sim_index( def write_actuator_stiffness_to_sim( self, - adapter: NewtonActuatorAdapter, *, stiffness: torch.Tensor | wp.array | float, env_ids: Sequence[int] | torch.Tensor | wp.array | None = None, ) -> None: - """Write actuator stiffness (``kp``) into each Newton controller for *env_ids*.""" + """Write actuator stiffness (``kp``) into this articulation's Newton controllers. + + No-op when no Newton actuators are registered for this articulation. + """ + adapter = self.actuators.get("newton") + if adapter is None: + return env_ids = self._resolve_env_ids(env_ids) env_mask = self._env_ids_to_mask(env_ids) adapter.write_stiffness_to_sim(stiffness, env_ids, env_mask, self._propagate_gain_via_kernel) def write_actuator_damping_to_sim( self, - adapter: NewtonActuatorAdapter, *, damping: torch.Tensor | wp.array | float, env_ids: Sequence[int] | torch.Tensor | wp.array | None = None, ) -> None: - """Write actuator damping (``kd``) into each Newton controller for *env_ids*.""" + """Write actuator damping (``kd``) into this articulation's Newton controllers.""" + adapter = self.actuators.get("newton") + if adapter is None: + return env_ids = self._resolve_env_ids(env_ids) env_mask = self._env_ids_to_mask(env_ids) adapter.write_damping_to_sim(damping, env_ids, env_mask, self._propagate_gain_via_kernel) diff --git a/source/isaaclab_physx/test/assets/test_newton_actuators_physx.py b/source/isaaclab_physx/test/assets/test_newton_actuators_physx.py index 8feb9916ea5f..29727b7a195d 100644 --- a/source/isaaclab_physx/test/assets/test_newton_actuators_physx.py +++ b/source/isaaclab_physx/test/assets/test_newton_actuators_physx.py @@ -527,6 +527,100 @@ def test_trajectories_not_trivial(self): self.assertGreater(diff, 0.01, "Joints did not move — test is trivial") +# --------------------------------------------------------------------------- +# Heterogeneous multi-articulation (ANYmal floating-base + Cartpole fixed-base) +# --------------------------------------------------------------------------- + + +CARTPOLE_EXPLICIT_ACTUATORS = { + "all_joints": IdealPDActuatorCfg( + joint_names_expr=["slider_to_cart", "cart_to_pole"], + stiffness=10.0, + damping=1.0, + effort_limit=100.0, + ), +} + + +def _run_anymal_and_cartpole(use_newton_actuators: bool, *, num_steps: int = NUM_STEPS) -> dict: + """Spawn ANYmal-C + Cartpole per env on PhysX (different DOF counts, base types).""" + from isaaclab_assets import CARTPOLE_CFG # noqa: PLC0415 + + sim_cfg = SimulationCfg(dt=DT, physics=PhysxCfg(), use_newton_actuators=use_newton_actuators) + with build_simulation_context( + device="cuda:0", gravity_enabled=True, add_ground_plane=True, sim_cfg=sim_cfg, + ) as sim: + sim._app_control_on_stop_handle = None + + for i in range(NUM_ENVS): + sim_utils.create_prim(f"/World/Env_{i}", "Xform", translation=(i * 6.0, 0, 0)) + + anymal_cfg = ANYMAL_C_CFG.replace(actuators=IDEAL_PD_ACTUATORS, prim_path="/World/Env_.*/Anymal") + cartpole_cfg = CARTPOLE_CFG.replace( + actuators=CARTPOLE_EXPLICIT_ACTUATORS, prim_path="/World/Env_.*/Cartpole", + ) + cartpole_cfg.init_state = cartpole_cfg.init_state.replace(pos=(0.0, 3.0, 2.0)) + + anymal = Articulation(anymal_cfg) + cartpole = Articulation(cartpole_cfg) + sim.reset() + assert anymal.is_initialized and cartpole.is_initialized + + init_anymal = wp.to_torch(anymal.data.joint_pos).clone() + init_cartpole = wp.to_torch(cartpole.data.joint_pos).clone() + anymal.set_joint_position_target_index(target=init_anymal + TARGET_OFFSET) + anymal.set_joint_velocity_target_index(target=torch.zeros_like(init_anymal)) + cartpole.set_joint_position_target_index(target=init_cartpole + TARGET_OFFSET) + cartpole.set_joint_velocity_target_index(target=torch.zeros_like(init_cartpole)) + + pos_anymal, pos_cartpole = [], [] + for _ in range(num_steps): + anymal.write_data_to_sim() + cartpole.write_data_to_sim() + sim.step() + anymal.update(DT) + cartpole.update(DT) + pos_anymal.append(wp.to_torch(anymal.data.joint_pos).clone()) + pos_cartpole.append(wp.to_torch(cartpole.data.joint_pos).clone()) + + return {"joint_pos_anymal": pos_anymal, "joint_pos_cartpole": pos_cartpole} + + +class TestHeterogeneousMultiArticulationPhysx(unittest.TestCase): + """Two structurally-different articulations (ANYmal floating + Cartpole fixed) on PhysX. + + Each PhysX articulation owns its own :class:`PhysxActuatorWrapper` + and per-art :class:`NewtonActuatorAdapter`. Heterogeneous DOF counts + (12 vs 2) and base types (floating vs fixed) verify the + per-articulation authoring + adapter construction works for varied + structures. Equivalence against the Lab actuator path is the + meaningful end-to-end check. + """ + + @classmethod + def setUpClass(cls): + cls.lab_result = _run_anymal_and_cartpole(use_newton_actuators=False) + cls.newton_result = _run_anymal_and_cartpole(use_newton_actuators=True) + + def test_anymal_matches_lab(self): + for step_i, (lab, newton) in enumerate( + zip(self.lab_result["joint_pos_anymal"], self.newton_result["joint_pos_anymal"]) + ): + torch.testing.assert_close( + newton, lab, atol=2e-3, rtol=1e-3, + msg=f"ANYmal joint_pos diverged from Lab path at step {step_i}", + ) + + def test_cartpole_matches_lab(self): + for step_i, (lab, newton) in enumerate( + zip(self.lab_result["joint_pos_cartpole"], self.newton_result["joint_pos_cartpole"]) + ): + torch.testing.assert_close( + newton, lab, atol=2e-3, rtol=1e-3, + msg=f"Cartpole joint_pos diverged from Lab path at step {step_i}", + ) + + class TestExplicitOnlyTelemetryPhysx(unittest.TestCase): """Explicit-only Newton actuators on PhysX: telemetry copies from the staging effort buffer.""" @@ -602,8 +696,8 @@ def test_writers_propagate_to_controller(self): self.assertGreater(len(kp_before), 0, "expected at least one PD controller in adapter") new_kp = adapter.stiffness.clone() * 2.0 new_kd = adapter.damping.clone() * 3.0 - articulation.write_actuator_stiffness_to_sim(adapter, stiffness=new_kp) - articulation.write_actuator_damping_to_sim(adapter, damping=new_kd) + articulation.write_actuator_stiffness_to_sim(stiffness=new_kp) + articulation.write_actuator_damping_to_sim(damping=new_kd) # Verify each controller's kp/kd actually changed (and roughly doubled/tripled). kp_idx = 0 kd_idx = 0 From 354ced7943849e381b129b562543d25df6bf01b7 Mon Sep 17 00:00:00 2001 From: Julia von Muralt Date: Fri, 8 May 2026 08:38:52 +0200 Subject: [PATCH 16/35] refactor --- .../assets/articulation/base_articulation.py | 42 +- source/isaaclab/isaaclab/envs/mdp/events.py | 4 +- .../isaaclab_newton/actuators/adapter.py | 551 ++++++++++ .../isaaclab_newton/actuators/authoring.py | 391 +++++++ .../isaaclab_newton/actuators/kernels.py | 85 ++ .../actuators/newton_actuator_utils.py | 972 ------------------ .../actuators/physx_wrapper.py | 60 ++ .../assets/articulation/articulation.py | 35 +- .../isaaclab_newton/physics/newton_manager.py | 25 +- .../assets/articulation/articulation.py | 35 +- 10 files changed, 1121 insertions(+), 1079 deletions(-) create mode 100644 source/isaaclab_newton/isaaclab_newton/actuators/adapter.py create mode 100644 source/isaaclab_newton/isaaclab_newton/actuators/authoring.py delete mode 100644 source/isaaclab_newton/isaaclab_newton/actuators/newton_actuator_utils.py create mode 100644 source/isaaclab_newton/isaaclab_newton/actuators/physx_wrapper.py diff --git a/source/isaaclab/isaaclab/assets/articulation/base_articulation.py b/source/isaaclab/isaaclab/assets/articulation/base_articulation.py index 3e6a7fe459c4..f275bdc3e984 100644 --- a/source/isaaclab/isaaclab/assets/articulation/base_articulation.py +++ b/source/isaaclab/isaaclab/assets/articulation/base_articulation.py @@ -100,47 +100,13 @@ def __init__(self, cfg: ArticulationCfg): sim_ctx = SimulationContext.instance() self._sim_cfg = sim_ctx.cfg if sim_ctx is not None else None - self._author_newton_actuator_prims() - - # -- USD authoring helpers ------------------------------------------------ - - def _author_newton_actuator_prims(self) -> None: - """Author Newton actuator prims on the USD stage from Lab configs. - - Called from ``__init__`` after ``super().__init__()`` has spawned the - articulation prototype. When ``use_newton_actuators`` is enabled, - this translates explicit Lab actuator configs into ``NewtonActuator`` - USD prims. - - For every joint covered by a Lab config, any existing - ``NewtonActuator`` prim targeting that joint is replaced. Joints - not covered by any config keep their USD-authored actuators. - """ - if self._sim_cfg is None: - return - if not getattr(self._sim_cfg, "use_newton_actuators", False): - return try: - from isaaclab_newton.actuators.newton_actuator_utils import ( # noqa: PLC0415 - author_newton_actuator_prims, - ) + from isaaclab_newton.actuators import author_actuator_prims_for_articulation # noqa: PLC0415 except ImportError: - return - - from isaaclab.sim.utils.queries import find_first_matching_prim # noqa: PLC0415 - - spawn_path = getattr(self.cfg.spawn, "spawn_path", None) if self.cfg.spawn is not None else None - search_path = spawn_path if spawn_path is not None else self.cfg.prim_path - first_prim = find_first_matching_prim(search_path) - if first_prim is None: - return - - author_newton_actuator_prims( - stage=self.stage, - articulation_prim_path=str(first_prim.GetPath()), - actuator_cfgs=self.cfg.actuators, - ) + pass + else: + author_actuator_prims_for_articulation(self.cfg, self._sim_cfg, self.stage) """ Properties diff --git a/source/isaaclab/isaaclab/envs/mdp/events.py b/source/isaaclab/isaaclab/envs/mdp/events.py index 18f5816a974a..52b24d17bcaf 100644 --- a/source/isaaclab/isaaclab/envs/mdp/events.py +++ b/source/isaaclab/isaaclab/envs/mdp/events.py @@ -1254,12 +1254,12 @@ def _randomize_at(data: torch.Tensor, params: tuple[float, float]) -> None: stiffness = adapter.stiffness[env_ids].clone() stiffness[:, adapter_indices] = self.default_joint_stiffness[env_ids][:, arti_local_ids].clone() _randomize_at(stiffness, stiffness_distribution_params) - adapter.write_stiffness_to_sim(stiffness, env_ids_wp, env_mask, self.asset._propagate_gain_via_view) + adapter.write_stiffness_to_sim(stiffness, env_ids_wp, env_mask) if damping_distribution_params is not None: damping = adapter.damping[env_ids].clone() damping[:, adapter_indices] = self.default_joint_damping[env_ids][:, arti_local_ids].clone() _randomize_at(damping, damping_distribution_params) - adapter.write_damping_to_sim(damping, env_ids_wp, env_mask, self.asset._propagate_gain_via_view) + adapter.write_damping_to_sim(damping, env_ids_wp, env_mask) class randomize_joint_parameters(ManagerTermBase): diff --git a/source/isaaclab_newton/isaaclab_newton/actuators/adapter.py b/source/isaaclab_newton/isaaclab_newton/actuators/adapter.py new file mode 100644 index 000000000000..4a76e440dc82 --- /dev/null +++ b/source/isaaclab_newton/isaaclab_newton/actuators/adapter.py @@ -0,0 +1,551 @@ +# Copyright (c) 2022-2026, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +"""Newton-actuator adapter shared by the Newton and PhysX backends. + +:class:`NewtonActuatorAdapter` manages actuator creation, DOF-to-actuator +mapping, stepping, reset, and gain reading for domain randomisation. +Used identically by both Newton and PhysX backends. + +The companion helper :func:`build_implicit_dof_mask` is consumed by the +in-graph post-actuator kernel +(:func:`~isaaclab_newton.actuators.kernels.synch_torque_and_apply_implicit_feedforwards`). +""" + +from __future__ import annotations + +from collections.abc import Callable, Sequence +from typing import Any + +import numpy as np +import torch +import warp as wp + +from newton.actuators import Actuator, Clamping, Delay + +from isaaclab.actuators import ActuatorBase, ImplicitActuator + +from .kernels import ( + fill_gain_at_envs_kernel, + gather_gain_kernel, + scatter_gain_at_envs_kernel, + scatter_gain_kernel, + set_mask_kernel, + zero_at_indices_kernel, +) + + +def _actuator_signature(parsed: Any) -> tuple: + """Build a hashable key from a parsed actuator spec for grouping. + + Joints whose prims resolve to the same signature share identical + controller type, gains, clamping chain, and delay configuration and + can therefore be merged into a single :class:`~newton.actuators.Actuator` + with combined index arrays. + """ + ctrl_resolved = parsed.controller_class.resolve_arguments( + dict(parsed.controller_kwargs), + ) + ctrl_key = (parsed.controller_class, tuple(sorted(ctrl_resolved.items()))) + + comp_keys: list[tuple] = [] + for comp_cls, comp_kwargs in parsed.component_specs: + resolved = comp_cls.resolve_arguments(comp_kwargs) + comp_keys.append((comp_cls, tuple(sorted(resolved.items())))) + comp_keys.sort(key=lambda t: t[0].__name__) + + return (ctrl_key, tuple(comp_keys)) + + +def build_implicit_dof_mask( + actuators: dict[str, ActuatorBase], + num_joints: int, + device: str, +) -> wp.array: + """Per-DOF mask consumed by the in-graph implicit-FF kernel. + + Entry is ``1`` for DOFs covered by an + :class:`~isaaclab.actuators.ImplicitActuator` group, ``0`` otherwise. + """ + modes = torch.zeros(num_joints, dtype=torch.int32, device=device) + for actuator in actuators.values(): + if not isinstance(actuator, ImplicitActuator): + continue + j_ids = actuator.joint_indices + if j_ids == slice(None) or j_ids is None: + modes[:] = 1 + else: + modes[j_ids.long()] = 1 + return wp.from_torch(modes, dtype=wp.int32) + + +class NewtonActuatorAdapter: + """Manages Newton-native actuators for both Newton and PhysX backends. + + Handles actuator creation (from USD or an existing list), + DOF-to-actuator mapping, stepping, reset, and gain reading for + domain randomisation. + + Construction: + + * **Newton backend** — pass actuators from the Newton model directly:: + + adapter = NewtonActuatorAdapter(model.actuators, num_envs, + num_joints, dof_offset, device) + + * **PhysX backend** — create actuators from USD prims:: + + adapter = NewtonActuatorAdapter.from_usd(stage, joint_names, + num_envs, num_joints, + device) + + Then finalise:: + + adapter.finalize() + + After :meth:`finalize`, the adapter exposes ``.stiffness``, + ``.damping``, and ``.joint_indices`` for ``randomize_actuator_gains``. + """ + + def __init__( + self, + actuators: list[Actuator], + num_envs: int, + num_joints: int, + dof_offset: int, + device: str, + ): + self.actuators = actuators + self.num_joints = num_joints + + self._num_envs = num_envs + self._dof_offset = dof_offset + self._device = device + self._dof_to_actuator = self._build_dof_map() + + managed = [i for i, act_idx in enumerate(self._dof_to_actuator) if act_idx >= 0] + if len(managed) == num_joints: + self.joint_indices: torch.Tensor | slice = slice(None) + else: + self.joint_indices = torch.tensor(managed, dtype=torch.int32, device=device) + + self._states_a = [act.state() for act in actuators] + self._states_b = [act.state() for act in actuators] + + self.stiffness: torch.Tensor | None = None + self.damping: torch.Tensor | None = None + + # Per-actuator gain propagator. Configured once at adapter setup time + # via :meth:`set_view_propagator` (Newton: simulator-side scatter API) + # or :meth:`set_kernel_propagator` (PhysX: local Warp scatter kernel). + # Stays ``None`` if gain DR is not used; ``write_*_to_sim`` then + # only updates the adapter's gain buffer without pushing to controllers. + self._propagator: Callable[[Actuator, Any, str, wp.array, wp.array], None] | None = None + + # -- construction (PhysX path) ------------------------------------------- + + @classmethod + def from_usd( + cls, + stage: Any, + joint_names: list[str], + num_envs: int, + num_joints: int, + device: str, + articulation_prim_path: str | None = None, + ) -> "NewtonActuatorAdapter": + """Create an adapter by parsing ``NewtonActuator`` prims from USD. + + This is the PhysX-backend counterpart of what Newton's + ``ModelBuilder.add_usd`` does for the Newton backend. Both paths + read the same ``NewtonActuator`` USD prims (authored by + :func:`~isaaclab_newton.actuators.authoring.author_newton_actuator_prims`) + and construct :class:`~newton.actuators.Actuator` objects with + matching controllers, clampings, and delays. + + The key difference is that PhysX uses a **flat per-DOF layout** + where joint position coordinates and velocity DOFs always have the + same count and ordering — there are no free joints or ball joints + that cause coordinate/DOF count divergence. Therefore a single + ``indices`` array is used for all index roles (``indices``, + ``pos_indices``, ``target_pos_indices``), unlike the Newton + builder which computes separate ``pos_indices`` from + ``joint_q_start`` and separate ``target_pos_indices`` from + ``joint_qd_start`` to handle floating-base articulations. + + Joints whose prims resolve to the same controller type, gains, + clamping chain, and delay configuration are merged into a single + :class:`Actuator` with combined index arrays, mirroring the + grouping the Newton builder performs internally. + + Args: + stage: The USD stage containing ``NewtonActuator`` prims. + joint_names: All joint names in the articulation. + num_envs: Number of environments. + num_joints: Joints per environment in the articulation. + device: Warp device string (e.g. ``"cuda:0"``). + articulation_prim_path: Root prim path of the first + environment's articulation (e.g. ``"/World/Env_0/Robot"``). + When provided, only ``NewtonActuator`` prims under this + subtree are considered — matching the scoped traversal + that Newton's ``ModelBuilder.add_usd`` performs. When + ``None``, the entire stage is scanned (legacy behaviour). + + Returns: + A fully constructed adapter ready for :meth:`finalize`. + """ + actuators = _create_actuators_from_usd( + stage, joint_names, num_envs, num_joints, device, + articulation_prim_path=articulation_prim_path, + ) + return cls(actuators, num_envs, num_joints, dof_offset=0, device=device) + + # -- public API ---------------------------------------------------------- + + def finalize(self) -> None: + """Read actuator gains and store as PyTorch tensors for DR.""" + wp_device = wp.get_device(self._device) + flat_stiffness = wp.zeros(self._num_envs * self.num_joints, dtype=wp.float32, device=wp_device) + flat_damping = wp.zeros(self._num_envs * self.num_joints, dtype=wp.float32, device=wp_device) + + for act in self.actuators: + ctrl = act.controller + if hasattr(ctrl, "kp"): + wp.launch( + scatter_gain_kernel, dim=act.indices.shape[0], + inputs=[ctrl.kp, flat_stiffness, act.indices, self._dof_offset, self.num_joints], + device=wp_device, + ) + if hasattr(ctrl, "kd"): + wp.launch( + scatter_gain_kernel, dim=act.indices.shape[0], + inputs=[ctrl.kd, flat_damping, act.indices, self._dof_offset, self.num_joints], + device=wp_device, + ) + + self.stiffness = wp.to_torch(flat_stiffness.reshape((self._num_envs, self.num_joints))) + self.damping = wp.to_torch(flat_damping.reshape((self._num_envs, self.num_joints))) + + def step(self, sim_state: Any, sim_control: Any, dt: float) -> None: + """Zero actuated DOFs, step all actuators, and swap state buffers. + + Args: + sim_state: Object with ``joint_q``, ``joint_qd``, etc. + Newton ``State`` on the Newton backend, + :class:`~isaaclab_newton.actuators.physx_wrapper.PhysxActuatorWrapper` + on the PhysX backend. + sim_control: Object with ``joint_f``, ``joint_target_pos``, etc. + Newton ``Control`` on the Newton backend, + :class:`~isaaclab_newton.actuators.physx_wrapper.PhysxActuatorWrapper` + on the PhysX backend. + dt: Physics timestep [s]. + """ + for act in self.actuators: + wp.launch( + zero_at_indices_kernel, + dim=act.indices.shape[0], + inputs=[sim_control.joint_f, act.indices], + ) + for act, sa, sb in zip(self.actuators, self._states_a, self._states_b): + act.step(sim_state, sim_control, sa, sb, dt=dt) + self._states_a, self._states_b = self._states_b, self._states_a + + def reset(self, env_ids: Sequence[int] | slice | None = None) -> None: + """Reset actuator states for the given environments. + + Args: + env_ids: Environment indices to reset. ``None`` or + ``slice(None)`` resets all environments. A partial slice + (e.g. ``slice(0, 5)``) is materialized to explicit indices. + """ + if env_ids is None or env_ids == slice(None): + mask = None + else: + # Normalize a partial slice to an explicit index list before + # building the wp.array — slices aren't iterable. + if isinstance(env_ids, slice): + env_ids = list(range(*env_ids.indices(self._num_envs))) + if isinstance(env_ids, torch.Tensor): + if env_ids.numel() == 0: + return + idx = wp.from_torch(env_ids.to(device=self._device).contiguous().to(torch.int32), dtype=wp.int32) + else: + if len(env_ids) == 0: + return + idx = wp.array(list(env_ids), dtype=wp.int32, device=self._device) + mask = wp.zeros(self._num_envs, dtype=wp.bool, device=self._device) + wp.launch(set_mask_kernel, dim=idx.shape[0], inputs=[mask, idx], device=self._device) + + for sa, sb in zip(self._states_a, self._states_b): + if sa is not None: + sa.reset(mask) + if sb is not None: + sb.reset(mask) + + @property + def is_all_graphable(self) -> bool: + """``True`` when all actuators are CUDA-graph-safe.""" + return len(self.actuators) > 0 and all(a.is_graphable() for a in self.actuators) + + def update_gain_at_env_ids( + self, + gain: str, + values: torch.Tensor | wp.array | float, + env_ids: wp.array, + ) -> wp.array: + """Scatter ``values`` into :attr:`stiffness` or :attr:`damping` at *env_ids*. + + Shared backend-independent step used by ``write_*_to_sim`` to + keep the kp/kd buffer the adapter exposes for DR in sync with what + each Newton actuator's controller will receive. + + Args: + gain: ``"stiffness"`` or ``"damping"``. + values: Per-env-per-joint values shape ``(len(env_ids), num_joints)``, + or a scalar to broadcast to every (env, joint). + env_ids: Warp int32 array of env indices to update. + + Returns: + Warp view of the updated ``(num_envs, num_joints)`` buffer. + """ + if gain == "stiffness": + buf = self.stiffness + elif gain == "damping": + buf = self.damping + else: + raise ValueError(f"gain must be 'stiffness' or 'damping', got {gain!r}") + buf_wp = wp.from_torch(buf, dtype=wp.float32) + if isinstance(values, float): + wp.launch( + fill_gain_at_envs_kernel, + dim=(env_ids.shape[0], self.num_joints), + inputs=[values, env_ids], + outputs=[buf_wp], + device=self._device, + ) + else: + wp.launch( + scatter_gain_at_envs_kernel, + dim=(env_ids.shape[0], self.num_joints), + inputs=[values, env_ids], + outputs=[buf_wp], + device=self._device, + ) + return buf_wp + + def set_view_propagator(self, root_view: Any) -> None: + """Configure gain propagation via Newton's simulator-side scatter API. + + Used on the Newton backend, where each per-actuator gain push goes + through ``ArticulationView.set_actuator_parameter``. Any one + articulation's view works since the call dispatches on the actuator + object (model-scoped), not on the view's articulation. + """ + def _push( + actuator: Actuator, controller: Any, attr: str, + values: wp.array, env_mask: wp.array, + ) -> None: + root_view.set_actuator_parameter( + actuator=actuator, component=controller, name=attr, + values=values, mask=env_mask, + ) + self._propagator = _push + + def set_kernel_propagator(self) -> None: + """Configure gain propagation via the local ``gather_gain_kernel``. + + Used on the PhysX backend, where there is no simulator-side scatter + API. The kernel reads the adapter's per-DOF gain buffer at each + actuator's flat indices and writes into ``controller.kp`` / + ``controller.kd``. + """ + def _push( + actuator: Actuator, controller: Any, attr: str, + values: wp.array, env_mask: wp.array, + ) -> None: + wp.launch( + gather_gain_kernel, + dim=actuator.indices.shape[0], + inputs=[ + values.flatten(), getattr(controller, attr), actuator.indices, + env_mask, self._dof_offset, self.num_joints, + ], + device=self._device, + ) + self._propagator = _push + + def write_stiffness_to_sim( + self, + stiffness: torch.Tensor | wp.array | float, + env_ids: wp.array, + env_mask: wp.array, + ) -> None: + """Update the kp buffer at *env_ids* and push the new values into each Newton controller.""" + self._write_gain_to_sim("stiffness", "kp", stiffness, env_ids, env_mask) + + def write_damping_to_sim( + self, + damping: torch.Tensor | wp.array | float, + env_ids: wp.array, + env_mask: wp.array, + ) -> None: + """Update the kd buffer at *env_ids* and push the new values into each Newton controller.""" + self._write_gain_to_sim("damping", "kd", damping, env_ids, env_mask) + + def _write_gain_to_sim( + self, + gain: str, + attr: str, + values: torch.Tensor | wp.array | float, + env_ids: wp.array, + env_mask: wp.array, + ) -> None: + """Shared body for :meth:`write_stiffness_to_sim` / :meth:`write_damping_to_sim`.""" + buf = self.update_gain_at_env_ids(gain, values, env_ids) + if self._propagator is None: + return + for newton_act in self.actuators: + ctrl = newton_act.controller + if hasattr(ctrl, attr): + self._propagator(newton_act, ctrl, attr, buf, env_mask) + + # -- private helpers ----------------------------------------------------- + + def _build_dof_map(self) -> list[int]: + """Build a per-DOF lookup: local DOF index -> actuator list index.""" + dof_to_actuator: list[int] = [-1] * self.num_joints + + for act_idx, act in enumerate(self.actuators): + all_indices = act.indices.numpy() + num_per_act = len(all_indices) // self._num_envs + env0_indices = all_indices[:num_per_act] + for global_dof in env0_indices: + local_dof = global_dof - self._dof_offset + if 0 <= local_dof < self.num_joints: + dof_to_actuator[local_dof] = act_idx + + return dof_to_actuator + + +def _create_actuators_from_usd( + stage: Any, + joint_names: list[str], + num_envs: int, + num_total_joints: int, + device: str, + articulation_prim_path: str | None = None, +) -> list[Actuator]: + """Parse ``NewtonActuator`` prims and instantiate standalone actuators. + + This mirrors the actuator construction that Newton's + ``ModelBuilder.add_usd`` performs, but operates independently of a + Newton ``Model``. It is used on the PhysX backend where there is no + Newton simulation — actuators are stepped manually via the adapter. + + Because PhysX articulations have no free or ball joints, every + joint's coordinate count equals its DOF count. A single + ``indices`` array is therefore sufficient for all index roles + (``indices``, ``pos_indices``, ``target_pos_indices``). + + Joints with identical controller type, gains, clamping chain, and + delay are merged into one :class:`Actuator` with combined indices. + + Each per-DOF scalar parameter (``kp``, ``kd``, ``saturation_effort``, + etc.) is broadcast via :func:`wp.full` to match the group size. + Parameters marked as ``SHARED_PARAMS`` on the controller or clamping + class (e.g. ``model_path``, ``lookup_positions``) are passed through + directly without broadcast. + """ + from collections import defaultdict # noqa: PLC0415 + + from newton.actuators import parse_actuator_prim # noqa: PLC0415 + from pxr import Usd # noqa: PLC0415 + + wp_device = wp.get_device(device) + + joint_name_to_idx: dict[str, int] = {name: i for i, name in enumerate(joint_names)} + + if articulation_prim_path is not None: + root_prim = stage.GetPrimAtPath(articulation_prim_path) + else: + root_prim = stage.GetPseudoRoot() + + parsed_per_joint: dict[int, Any] = {} + for prim in Usd.PrimRange(root_prim): + parsed = parse_actuator_prim(prim) + if parsed is None: + continue + target_name = parsed.target_path.rsplit("/", 1)[-1] + if target_name in joint_name_to_idx: + parsed_per_joint[joint_name_to_idx[target_name]] = parsed + + if not parsed_per_joint: + raise ValueError( + f"No NewtonActuator prims found targeting any of: {joint_names}" + ) + + groups: dict[tuple, list[int]] = defaultdict(list) + sig_to_parsed: dict[tuple, Any] = {} + for local_idx, parsed in sorted(parsed_per_joint.items()): + sig = _actuator_signature(parsed) + groups[sig].append(local_idx) + if sig not in sig_to_parsed: + sig_to_parsed[sig] = parsed + + actuators = [] + for sig, local_indices in groups.items(): + parsed = sig_to_parsed[sig] + + flat_indices = np.array( + [idx + e * num_total_joints for e in range(num_envs) for idx in local_indices], + dtype=np.uint32, + ) + indices = wp.array(flat_indices, device=wp_device) + num_dofs_in_group = len(local_indices) * num_envs + + # Controller + ctrl_kwargs = dict(parsed.controller_kwargs) + resolved = parsed.controller_class.resolve_arguments(ctrl_kwargs) + shared_ctrl = getattr(parsed.controller_class, "SHARED_PARAMS", set()) + ctrl_arrays = {} + for key, val in resolved.items(): + if key in shared_ctrl: + ctrl_arrays[key] = val + else: + ctrl_arrays[key] = wp.full(num_dofs_in_group, float(val), dtype=wp.float32, device=wp_device) + controller = parsed.controller_class(**ctrl_arrays) + + # Components (delay + clampings) + clampings = [] + delay = None + for comp_cls, comp_kwargs in parsed.component_specs: + if issubclass(comp_cls, Delay): + resolved_kw = Delay.resolve_arguments(comp_kwargs) + delay_steps = int(resolved_kw.get("delay_steps", 0)) + if delay_steps > 0: + delay_arr = wp.full(num_dofs_in_group, delay_steps, dtype=wp.int32, device=wp_device) + delay = Delay(delay_steps=delay_arr, max_delay=delay_steps) + elif issubclass(comp_cls, Clamping): + resolved_kw = comp_cls.resolve_arguments(comp_kwargs) + shared_clamp = getattr(comp_cls, "SHARED_PARAMS", set()) + clamp_arrays = {} + for k, v in resolved_kw.items(): + if k in shared_clamp: + clamp_arrays[k] = v + else: + clamp_arrays[k] = wp.full( + num_dofs_in_group, float(v), dtype=wp.float32, device=wp_device, + ) + clampings.append(comp_cls(**clamp_arrays)) + + actuator = Actuator( + indices=indices, + controller=controller, + delay=delay, + clamping=clampings if clampings else None, + ) + actuators.append(actuator) + + return actuators diff --git a/source/isaaclab_newton/isaaclab_newton/actuators/authoring.py b/source/isaaclab_newton/isaaclab_newton/actuators/authoring.py new file mode 100644 index 000000000000..999870e526f1 --- /dev/null +++ b/source/isaaclab_newton/isaaclab_newton/actuators/authoring.py @@ -0,0 +1,391 @@ +# Copyright (c) 2022-2026, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +"""USD authoring for Newton-native actuators. + +:func:`author_newton_actuator_prims` translates IsaacLab actuator configs +into ``NewtonActuator`` USD prims. Both the Newton ``ModelBuilder.add_usd`` +path and the PhysX adapter's +:meth:`~isaaclab_newton.actuators.adapter.NewtonActuatorAdapter.from_usd` +read the same authored prims, ensuring both backends construct +:class:`~newton.actuators.Actuator` instances with matching parameters. + +:func:`author_actuator_prims_for_articulation` is a thin wrapper that +resolves per-articulation pre-conditions (sim cfg gating, prim lookup) +and dispatches into :func:`author_newton_actuator_prims`. It is invoked +from the schema-side articulation initialisation path. +""" + +from __future__ import annotations + +import logging +import re +from typing import Any + +logger = logging.getLogger(__name__) + + +def resolve_per_dof( + value: dict[str, float | int] | float | int | None, + joint_names: list[str], + cast: type = float, +) -> dict[str, float | int]: + """Expand a scalar or regex-keyed dict cfg value into a per-joint mapping. + + Used by :func:`author_newton_actuator_prims` to flatten the various + accepted forms of a per-DOF config field (``stiffness``, ``damping``, + ``effort_limit``, …) into a single ``{joint_name: value}`` dict that + the authoring loop can ``.get(jname, default)`` against. + + Accepted forms of *value*: + + * ``None`` → empty dict. + * scalar (``int`` / ``float``) → broadcast to every joint name. + * dict → keys are treated as regex patterns and matched against + *joint_names* via :func:`re.fullmatch`. The first matching pattern + wins per joint name. + """ + if value is None: + return {} + if isinstance(value, (int, float)): + return {name: cast(value) for name in joint_names} + if isinstance(value, dict): + result: dict[str, float | int] = {} + for name in joint_names: + for pattern, v in value.items(): + if re.fullmatch(pattern, name): + result[name] = cast(v) + break + return result + return {} + + +def author_actuator_prims_for_articulation(cfg: Any, sim_cfg: Any, stage: Any) -> None: + """Resolve per-articulation authoring args from ``cfg`` + ``sim_cfg`` and dispatch. + + Higher-level wrapper around :func:`author_newton_actuator_prims` that + handles the pre-conditions previously inlined in + ``BaseArticulation._author_newton_actuator_prims``: + + * No-op when the simulation is not configured for the Newton actuator + fast path (``sim_cfg is None`` or ``use_newton_actuators=False``). + * Falls back from ``cfg.spawn.spawn_path`` to ``cfg.prim_path`` when the + former is not set, then resolves the first matching USD prim. + * No-op when no matching prim exists on the stage. + + Keeping this resolution on the schema side lets articulations call + authoring with a single line, and lets the rules be tested without + constructing an articulation instance. + """ + if sim_cfg is None or not getattr(sim_cfg, "use_newton_actuators", False): + return + + from isaaclab.sim.utils.queries import find_first_matching_prim # noqa: PLC0415 + + spawn_path = getattr(cfg.spawn, "spawn_path", None) if cfg.spawn is not None else None + search_path = spawn_path if spawn_path is not None else cfg.prim_path + first_prim = find_first_matching_prim(search_path) + if first_prim is None: + return + + author_newton_actuator_prims( + stage=stage, + articulation_prim_path=str(first_prim.GetPath()), + actuator_cfgs=cfg.actuators, + ) + + +def author_newton_actuator_prims( + stage: Any, + articulation_prim_path: str, + actuator_cfgs: dict[str, Any], +) -> None: + """Author ``NewtonActuator`` USD prims from IsaacLab actuator configs. + + For every joint covered by an explicit (non-implicit) Lab actuator config, + any existing ``NewtonActuator`` prim targeting that joint is replaced by a + new one created from the config values. Joints **not** covered by any + Lab config keep their USD-authored actuators unchanged. + + The supported config-to-schema mapping is: + + * :class:`~isaaclab.actuators.IdealPDActuatorCfg` -> + ``NewtonPDControlAPI`` + ``NewtonMaxEffortClampingAPI`` + * :class:`~isaaclab.actuators.DCMotorCfg` -> + ``NewtonPDControlAPI`` + ``NewtonDCMotorClampingAPI`` + * :class:`~isaaclab.actuators.DelayedPDActuatorCfg` -> + same as ``IdealPDActuatorCfg`` + ``NewtonActuatorDelayAPI`` + * :class:`~isaaclab.actuators.RemotizedPDActuatorCfg` -> + same as ``DelayedPDActuatorCfg`` + ``NewtonPositionBasedClampingAPI`` + * :class:`~isaaclab.actuators.ActuatorNetMLPCfg` / + :class:`~isaaclab.actuators.ActuatorNetLSTMCfg` -> + ``NewtonNeuralControlAPI`` (+ ``NewtonDCMotorClampingAPI``) + + Must be called **after** the articulation is spawned (joint prims exist + on stage) and **before** the cloner / ``ModelBuilder.add_usd`` reads + the stage. + + Args: + stage: The USD stage to author prims on. + articulation_prim_path: Root prim path of the articulation + (e.g. ``"/World/Env_0/Robot"``). Must not contain wildcards. + actuator_cfgs: Mapping of group name to ``ActuatorBaseCfg``. + """ + from pxr import Sdf # noqa: PLC0415 + + from isaaclab.actuators import ImplicitActuator # noqa: PLC0415 + from isaaclab.utils.string import resolve_matching_names # noqa: PLC0415 + + art_prim = stage.GetPrimAtPath(articulation_prim_path) + if not art_prim.IsValid(): + raise ValueError(f"Articulation prim not found: {articulation_prim_path}") + + joint_inventory = _collect_joint_prims(art_prim) + all_joint_names = list(joint_inventory.keys()) + + covered_joint_paths: set[str] = set() + + cfg_entries: list[tuple[str, Any, list[str]]] = [] + for group_name, cfg in actuator_cfgs.items(): + cls_type = cfg.class_type + is_implicit = ( + "ImplicitActuator" in cls_type + if isinstance(cls_type, str) + else issubclass(cls_type, ImplicitActuator) + ) + if is_implicit: + continue + + _ids, joint_names = resolve_matching_names(cfg.joint_names_expr, all_joint_names) + if not joint_names: + continue + + cfg_entries.append((group_name, cfg, joint_names)) + for jname in joint_names: + covered_joint_paths.add(joint_inventory[jname]) + + _remove_actuator_prims_for_joints(art_prim, covered_joint_paths) + + from isaaclab.actuators import DCMotorCfg, DelayedPDActuatorCfg # noqa: PLC0415 + from isaaclab.actuators.actuator_net_cfg import ActuatorNetLSTMCfg, ActuatorNetMLPCfg # noqa: PLC0415 + from isaaclab.actuators.actuator_pd_cfg import IdealPDActuatorCfg, RemotizedPDActuatorCfg # noqa: PLC0415 + + _SUPPORTED_CFG_TYPES = ( + IdealPDActuatorCfg, + DCMotorCfg, + DelayedPDActuatorCfg, + RemotizedPDActuatorCfg, + ActuatorNetMLPCfg, + ActuatorNetLSTMCfg, + ) + + for group_name, cfg, joint_names in cfg_entries: + if not isinstance(cfg, _SUPPORTED_CFG_TYPES): + logger.warning( + "Actuator group '%s' uses config type '%s' which is not supported by Newton-native" + " actuator authoring. The group will be skipped.", + group_name, + type(cfg).__name__, + ) + continue + stiffness_map = resolve_per_dof(getattr(cfg, "stiffness", None), joint_names) + damping_map = resolve_per_dof(getattr(cfg, "damping", None), joint_names) + effort_map = resolve_per_dof(getattr(cfg, "effort_limit", None), joint_names) + + is_neural = isinstance(cfg, (ActuatorNetMLPCfg, ActuatorNetLSTMCfg)) + is_remotized = isinstance(cfg, RemotizedPDActuatorCfg) + is_dc_motor = isinstance(cfg, DCMotorCfg) + is_delayed = isinstance(cfg, DelayedPDActuatorCfg) + + vel_limit_map = resolve_per_dof(getattr(cfg, "velocity_limit", None), joint_names) if is_dc_motor else {} + sat_effort_map = ( + resolve_per_dof(getattr(cfg, "saturation_effort", None), joint_names) if is_dc_motor else {} + ) + + raw_delay = getattr(cfg, "max_delay", 0) if is_delayed else 0 + delay_map = resolve_per_dof(raw_delay, joint_names, cast=int) if raw_delay else {} + + patched_model_path: str | None = None + if is_neural: + meta: dict[str, Any] = {} + if isinstance(cfg, ActuatorNetMLPCfg): + meta["model_type"] = "mlp" + meta["input_order"] = cfg.input_order + meta["input_idx"] = list(cfg.input_idx) + meta["pos_scale"] = cfg.pos_scale + meta["vel_scale"] = cfg.vel_scale + meta["torque_scale"] = cfg.torque_scale + else: + meta["model_type"] = "lstm" + patched_model_path = _resave_checkpoint_with_metadata(cfg.network_file, meta) + + for jname in joint_names: + joint_prim_path = joint_inventory[jname] + + schemas: list[str] = [] + attrs: dict[str, float | int] = {} + array_attrs: dict[str, list[float]] = {} + + if is_neural: + schemas.append("NewtonNeuralControlAPI") + else: + schemas.append("NewtonPDControlAPI") + attrs["kp"] = stiffness_map.get(jname, 0.0) + attrs["kd"] = damping_map.get(jname, 0.0) + + if is_dc_motor: + schemas.append("NewtonDCMotorClampingAPI") + attrs["saturation_effort"] = sat_effort_map.get(jname, 0.0) + if jname in vel_limit_map: + attrs["velocity_limit"] = vel_limit_map[jname] + if jname in effort_map: + attrs["max_motor_effort"] = effort_map[jname] + elif jname in effort_map: + schemas.append("NewtonMaxEffortClampingAPI") + attrs["max_effort"] = effort_map[jname] + + if is_remotized and isinstance(cfg, RemotizedPDActuatorCfg): + lookup = cfg.joint_parameter_lookup + schemas.append("NewtonPositionBasedClampingAPI") + array_attrs["lookup_positions"] = [row[0] for row in lookup] + array_attrs["lookup_efforts"] = [row[2] for row in lookup] + + delay_steps = delay_map.get(jname, 0) + if delay_steps > 0: + schemas.append("NewtonActuatorDelayAPI") + attrs["delay_steps"] = delay_steps + attrs["max_delay"] = delay_steps + + act_prim_path = f"{articulation_prim_path}/{group_name}_{jname}_actuator" + act_prim = stage.DefinePrim(act_prim_path, "NewtonActuator") + + existing = act_prim.GetMetadata("apiSchemas") or Sdf.TokenListOp() + existing.prependedItems = list(schemas) + act_prim.SetMetadata("apiSchemas", existing) + + rel = act_prim.CreateRelationship("newton:targets") + rel.SetTargets([Sdf.Path(joint_prim_path)]) + + if patched_model_path is not None: + act_prim.CreateAttribute("newton:modelPath", Sdf.ValueTypeNames.Asset).Set( + Sdf.AssetPath(patched_model_path) + ) + + for attr_name, attr_val in attrs.items(): + usd_name = f"newton:{_snake_to_camel(attr_name)}" + if isinstance(attr_val, int): + act_prim.CreateAttribute(usd_name, Sdf.ValueTypeNames.Int).Set(attr_val) + else: + act_prim.CreateAttribute(usd_name, Sdf.ValueTypeNames.Float).Set(float(attr_val)) + + for attr_name, attr_val in array_attrs.items(): + usd_name = f"newton:{_snake_to_camel(attr_name)}" + act_prim.CreateAttribute(usd_name, Sdf.ValueTypeNames.FloatArray).Set(attr_val) + + +# --------------------------------------------------------------------------- +# Private helpers +# --------------------------------------------------------------------------- + +_SNAKE_TO_CAMEL_RE = re.compile(r"_([a-z])") + + +def _snake_to_camel(name: str) -> str: + """Convert a snake_case name to camelCase.""" + return _SNAKE_TO_CAMEL_RE.sub(lambda m: m.group(1).upper(), name) + + +def _collect_joint_prims(art_prim: Any) -> dict[str, str]: + """Collect all joint prims under an articulation subtree. + + Returns: + Ordered mapping of joint name to full prim path. + """ + from pxr import Usd # noqa: PLC0415 + + _JOINT_TYPES = {"PhysicsRevoluteJoint", "PhysicsPrismaticJoint"} + + joints: dict[str, str] = {} + for prim in Usd.PrimRange(art_prim): + if prim.GetTypeName() in _JOINT_TYPES: + joints[prim.GetName()] = str(prim.GetPath()) + return joints + + +def _remove_actuator_prims_for_joints( + art_prim: Any, + joint_paths: set[str], +) -> None: + """Deactivate ``NewtonActuator`` prims whose target is in *joint_paths*. + + Deactivated prims are invisible to ``Usd.PrimRange`` and therefore + ignored by ``ModelBuilder.add_usd``. Using ``SetActive(False)`` + instead of ``RemovePrim`` works correctly when the prim originates + from a USD reference or payload. + + Only prims under the *art_prim* subtree are considered. + """ + from pxr import Usd # noqa: PLC0415 + + to_deactivate: list = [] + for prim in Usd.PrimRange(art_prim): + if prim.GetTypeName() != "NewtonActuator": + continue + rel = prim.GetRelationship("newton:targets") + if rel and rel.IsValid(): + for target in rel.GetTargets(): + if str(target) in joint_paths: + to_deactivate.append(prim) + break + + for prim in to_deactivate: + prim.SetActive(False) + + +def _resave_checkpoint_with_metadata( + original_path: str, + metadata: dict[str, Any], +) -> str: + """Re-save a neural-network checkpoint with updated metadata. + + Loads the original TorchScript or dict checkpoint, merges *metadata* + into any existing metadata (Lab config values take precedence), and + writes the result to a temporary ``.pt`` file that persists for the + lifetime of the process. + + Returns: + Path to the temporary checkpoint file. + """ + import json # noqa: PLC0415 + import tempfile # noqa: PLC0415 + + import torch # noqa: PLC0415 + + extra_files: dict[str, str] = {"metadata.json": ""} + is_torchscript = True + try: + net = torch.jit.load(original_path, map_location="cpu", _extra_files=extra_files) + existing_meta = json.loads(extra_files["metadata.json"]) if extra_files["metadata.json"] else {} + except Exception: + is_torchscript = False + checkpoint = torch.load(original_path, map_location="cpu", weights_only=False) + if not isinstance(checkpoint, dict) or "model" not in checkpoint: + raise ValueError( + f"Cannot load checkpoint at '{original_path}'; " + "expected a TorchScript archive or a dict with a 'model' key" + ) + net = checkpoint["model"] + existing_meta = checkpoint.get("metadata", {}) + + merged = {**existing_meta, **metadata} + + tmp = tempfile.NamedTemporaryFile(suffix=".pt", delete=False) + if is_torchscript: + extra_out = {"metadata.json": json.dumps(merged)} + torch.jit.save(net, tmp.name, _extra_files=extra_out) + else: + torch.save({"model": net, "metadata": merged}, tmp.name) + + return tmp.name diff --git a/source/isaaclab_newton/isaaclab_newton/actuators/kernels.py b/source/isaaclab_newton/isaaclab_newton/actuators/kernels.py index 467cc340e6c6..b850365c41ec 100644 --- a/source/isaaclab_newton/isaaclab_newton/actuators/kernels.py +++ b/source/isaaclab_newton/isaaclab_newton/actuators/kernels.py @@ -8,6 +8,91 @@ import warp as wp +# --------------------------------------------------------------------------- +# Adapter-internal kernels: per-DOF zeroing, env-mask building, and gain +# scatter/gather between the adapter's flat per-env-per-DOF buffer and the +# per-actuator controller arrays. Used by :class:`NewtonActuatorAdapter` +# (stepping, finalize, gain DR) and by the kernel-based propagator. +# --------------------------------------------------------------------------- + + +@wp.kernel(enable_backward=False) +def zero_at_indices_kernel(data: wp.array(dtype=wp.float32), indices: wp.array(dtype=wp.uint32)): + """Zero a flat ``data`` buffer at the given flat ``indices``.""" + i = wp.tid() + data[indices[i]] = 0.0 + + +@wp.kernel(enable_backward=False) +def set_mask_kernel(mask: wp.array(dtype=wp.bool), indices: wp.array(dtype=wp.int32)): + """Set ``mask[indices[i]] = True`` for each ``i``. The mask must be pre-zeroed.""" + i = wp.tid() + mask[indices[i]] = True + + +@wp.kernel(enable_backward=False) +def scatter_gain_kernel( + src: wp.array(dtype=wp.float32), + dst: wp.array(dtype=wp.float32), + indices: wp.array(dtype=wp.uint32), + dof_offset: int, + num_joints: int, +): + """Scatter per-actuator ``src`` values into the adapter's flat per-env-per-DOF ``dst``.""" + i = wp.tid() + global_dof = int(indices[i]) - dof_offset + env = global_dof // num_joints + local_dof = global_dof % num_joints + dst[env * num_joints + local_dof] = src[i] + + +@wp.kernel(enable_backward=False) +def gather_gain_kernel( + flat_src: wp.array(dtype=wp.float32), + dst: wp.array(dtype=wp.float32), + indices: wp.array(dtype=wp.uint32), + env_mask: wp.array(dtype=wp.bool), + dof_offset: int, + num_joints: int, +): + """Gather from the adapter's flat ``(num_envs * num_joints)`` layout into a + per-actuator controller array, only for envs where ``env_mask`` is ``True``. + """ + i = wp.tid() + global_dof = int(indices[i]) - dof_offset + env = global_dof // num_joints + if env_mask[env]: + local_dof = global_dof % num_joints + dst[i] = flat_src[env * num_joints + local_dof] + + +@wp.kernel(enable_backward=False) +def scatter_gain_at_envs_kernel( + in_data: wp.array2d(dtype=wp.float32), + env_ids: wp.array(dtype=wp.int32), + out_data: wp.array2d(dtype=wp.float32), +): + """Scatter ``in_data[i, j]`` into ``out_data[env_ids[i], j]`` for all (i, j).""" + i, j = wp.tid() + out_data[env_ids[i], j] = in_data[i, j] + + +@wp.kernel(enable_backward=False) +def fill_gain_at_envs_kernel( + value: float, + env_ids: wp.array(dtype=wp.int32), + out_data: wp.array2d(dtype=wp.float32), +): + """Set ``out_data[env_ids[i], j] = value`` for all (i, j).""" + i, j = wp.tid() + out_data[env_ids[i], j] = value + + +# --------------------------------------------------------------------------- +# Articulation-level kernels: in-graph post-actuator hook. +# --------------------------------------------------------------------------- + + @wp.kernel(enable_backward=False) def synch_torque_and_apply_implicit_feedforwards( joint_pos: wp.array2d(dtype=wp.float32), diff --git a/source/isaaclab_newton/isaaclab_newton/actuators/newton_actuator_utils.py b/source/isaaclab_newton/isaaclab_newton/actuators/newton_actuator_utils.py deleted file mode 100644 index a886c44b3277..000000000000 --- a/source/isaaclab_newton/isaaclab_newton/actuators/newton_actuator_utils.py +++ /dev/null @@ -1,972 +0,0 @@ -# Copyright (c) 2022-2026, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md). -# All rights reserved. -# -# SPDX-License-Identifier: BSD-3-Clause - -"""Utilities for Newton-native actuators in Isaac Lab. - -This module is organised into four sections: - -1. **Warp kernels** — low-level device kernels for zeroing, masking, - and scattering per-DOF data used by the adapter. -2. **PhysX stepping helper** — :class:`PhysxActuatorWrapper`, a - duck-typed wrapper that exposes flat Warp arrays as the - ``sim_state`` / ``sim_control`` protocol expected by - :meth:`Actuator.step` on the PhysX backend. -3. **Adapter** — :class:`NewtonActuatorAdapter` manages actuator - creation, DOF-to-actuator mapping, config overrides, stepping, - reset, and gain reading for domain randomisation. Used identically - by both Newton and PhysX backends. -4. **USD authoring** — :func:`author_newton_actuator_prims` creates - ``NewtonActuator`` USD prims from IsaacLab actuator configs so that - both the Newton ``ModelBuilder`` (during ``add_usd``) and the PhysX - adapter (via :meth:`NewtonActuatorAdapter.from_usd`) can construct - :class:`Actuator` objects with correct parameters. - -Why :class:`PhysxActuatorWrapper` exists only for PhysX -~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ - -Newton's :meth:`Actuator.step` requires a ``sim_state`` / ``sim_control`` -pair that exposes flat 1-D Warp arrays (``joint_q``, ``joint_qd``, -``joint_target_pos``, ``joint_f``, …). On the **Newton backend** these -are the ``State`` and ``Control`` objects that the solver already owns — -no wrapper is needed because: - -* The solver manages double-buffered ``State`` objects for CUDA-graph - capture, and actuators are stepped inside the solver's own simulation - loop where states are already available. -* Wrapping them would add indirection with no benefit; the Newton - articulation code that calls :meth:`Actuator.step` lives in - ``newton_manager.py`` and has direct access to the model's state. - -On the **PhysX backend**, no Newton solver exists — the actuators are -stepped manually from the Lab articulation's ``write_data_to_sim`` -path. Isaac Lab stores joint data as 2-D tensors (``num_envs × -num_joints``), so :class:`PhysxActuatorWrapper` provides zero-copy flat -views that satisfy the protocol without allocating new memory. The -PhysX articulation code that calls :meth:`Actuator.step` lives in -``isaaclab_physx/.../articulation.py``, completely separate from the -Newton solver path, so sharing a single wrapper type across both -backends would not reduce code — it would only add coupling. -""" - -from __future__ import annotations - -import logging -from collections.abc import Callable, Sequence -from dataclasses import dataclass -from typing import Any - -import numpy as np -import torch -import warp as wp - -logger = logging.getLogger(__name__) - -from newton.actuators import ( - Actuator, - Clamping, - Delay, -) - -from isaaclab.actuators import ActuatorBase, ImplicitActuator - - -# =========================================================================== -# 1. Warp kernels -# =========================================================================== - - -@wp.kernel(enable_backward=False) -def _zero_at_indices_kernel(data: wp.array(dtype=wp.float32), indices: wp.array(dtype=wp.uint32)): - i = wp.tid() - data[indices[i]] = 0.0 - - -@wp.kernel(enable_backward=False) -def _set_mask_kernel(mask: wp.array(dtype=wp.bool), indices: wp.array(dtype=wp.int32)): - i = wp.tid() - mask[indices[i]] = True - - -@wp.kernel(enable_backward=False) -def _scatter_gain_kernel( - src: wp.array(dtype=wp.float32), - dst: wp.array(dtype=wp.float32), - indices: wp.array(dtype=wp.uint32), - dof_offset: int, - num_joints: int, -): - i = wp.tid() - global_dof = int(indices[i]) - dof_offset - env = global_dof // num_joints - local_dof = global_dof % num_joints - dst[env * num_joints + local_dof] = src[i] - - -@wp.kernel(enable_backward=False) -def _gather_gain_kernel( - flat_src: wp.array(dtype=wp.float32), - dst: wp.array(dtype=wp.float32), - indices: wp.array(dtype=wp.uint32), - env_mask: wp.array(dtype=wp.bool), - dof_offset: int, - num_joints: int, -): - """Gather from flat ``(num_envs * num_joints)`` layout into a per-actuator - controller array, only for envs where ``env_mask`` is ``True``.""" - i = wp.tid() - global_dof = int(indices[i]) - dof_offset - env = global_dof // num_joints - if env_mask[env]: - local_dof = global_dof % num_joints - dst[i] = flat_src[env * num_joints + local_dof] - - -@wp.kernel(enable_backward=False) -def _scatter_gain_at_envs_kernel( - in_data: wp.array2d(dtype=wp.float32), - env_ids: wp.array(dtype=wp.int32), - out_data: wp.array2d(dtype=wp.float32), -): - """Scatter ``in_data[i, j]`` into ``out_data[env_ids[i], j]`` for all (i, j).""" - i, j = wp.tid() - out_data[env_ids[i], j] = in_data[i, j] - - -@wp.kernel(enable_backward=False) -def _fill_gain_at_envs_kernel( - value: float, - env_ids: wp.array(dtype=wp.int32), - out_data: wp.array2d(dtype=wp.float32), -): - """Set ``out_data[env_ids[i], j] = value`` for all (i, j).""" - i, j = wp.tid() - out_data[env_ids[i], j] = value - - -# =========================================================================== -# 2. PhysX stepping helper -# =========================================================================== - - -@dataclass -class PhysxActuatorWrapper: - """Flat-array wrapper serving as ``sim_state`` / ``sim_control`` for - :meth:`Actuator.step` on the PhysX backend. - - Most attributes are reassigned each frame to zero-copy flat views of - Isaac Lab's 2-D buffers. ``joint_f_2d`` is the only persistent - allocation, sized via :meth:`create`; ``joint_f`` is its flat alias - consumed by the Newton actuator step. - """ - - joint_q: wp.array | None = None - joint_qd: wp.array | None = None - joint_target_pos: wp.array | None = None - joint_target_vel: wp.array | None = None - joint_act: wp.array | None = None - joint_f: wp.array | None = None - joint_f_2d: wp.array | None = None - - @classmethod - def create(cls, num_envs: int, num_joints: int, device: str) -> PhysxActuatorWrapper: - """Allocate the persistent ``joint_f`` buffer for the given articulation shape.""" - w = cls() - w.joint_f_2d = wp.zeros((num_envs, num_joints), dtype=wp.float32, device=device) - w.joint_f = w.joint_f_2d.reshape(-1) - return w - - -def build_implicit_dof_mask( - actuators: dict[str, ActuatorBase], - num_joints: int, - device: str, -) -> wp.array: - """Per-DOF mask consumed by the in-graph implicit-FF kernel. - - Entry is ``1`` for DOFs covered by an - :class:`~isaaclab.actuators.ImplicitActuator` group, ``0`` otherwise. - """ - modes = torch.zeros(num_joints, dtype=torch.int32, device=device) - for actuator in actuators.values(): - if not isinstance(actuator, ImplicitActuator): - continue - j_ids = actuator.joint_indices - if j_ids == slice(None) or j_ids is None: - modes[:] = 1 - else: - modes[j_ids.long()] = 1 - return wp.from_torch(modes, dtype=wp.int32) - - -# =========================================================================== -# 3. Adapter — creation, config overrides, stepping, domain randomisation -# =========================================================================== - - -class NewtonActuatorAdapter: - """Manages Newton-native actuators for both Newton and PhysX backends. - - Handles actuator creation (from USD or an existing list), - DOF-to-actuator mapping, stepping, reset, and gain reading for - domain randomisation. - - Construction: - - * **Newton backend** — pass actuators from the Newton model directly:: - - adapter = NewtonActuatorAdapter(model.actuators, num_envs, - num_joints, dof_offset, device) - - * **PhysX backend** — create actuators from USD prims:: - - adapter = NewtonActuatorAdapter.from_usd(stage, joint_names, - num_envs, num_joints, - device) - - Then finalise:: - - adapter.finalize() - - After :meth:`finalize`, the adapter exposes ``.stiffness``, - ``.damping``, and ``.joint_indices`` for ``randomize_actuator_gains``. - """ - - def __init__( - self, - actuators: list[Actuator], - num_envs: int, - num_joints: int, - dof_offset: int, - device: str, - ): - self.actuators = actuators - self.num_joints = num_joints - - self._num_envs = num_envs - self._dof_offset = dof_offset - self._device = device - self._dof_to_actuator = self._build_dof_map() - - managed = [i for i, act_idx in enumerate(self._dof_to_actuator) if act_idx >= 0] - if len(managed) == num_joints: - self.joint_indices: torch.Tensor | slice = slice(None) - else: - self.joint_indices = torch.tensor(managed, dtype=torch.int32, device=device) - - self._states_a = [act.state() for act in actuators] - self._states_b = [act.state() for act in actuators] - - self.stiffness: torch.Tensor | None = None - self.damping: torch.Tensor | None = None - - # -- construction (PhysX path) ------------------------------------------- - - @classmethod - def from_usd( - cls, - stage: Any, - joint_names: list[str], - num_envs: int, - num_joints: int, - device: str, - articulation_prim_path: str | None = None, - ) -> NewtonActuatorAdapter: - """Create an adapter by parsing ``NewtonActuator`` prims from USD. - - This is the PhysX-backend counterpart of what Newton's - ``ModelBuilder.add_usd`` does for the Newton backend. Both paths - read the same ``NewtonActuator`` USD prims (authored by - :func:`author_newton_actuator_prims`) and construct - :class:`~newton.actuators.Actuator` objects with matching - controllers, clampings, and delays. - - The key difference is that PhysX uses a **flat per-DOF layout** - where joint position coordinates and velocity DOFs always have the - same count and ordering — there are no free joints or ball joints - that cause coordinate/DOF count divergence. Therefore a single - ``indices`` array is used for all index roles (``indices``, - ``pos_indices``, ``target_pos_indices``), unlike the Newton - builder which computes separate ``pos_indices`` from - ``joint_q_start`` and separate ``target_pos_indices`` from - ``joint_qd_start`` to handle floating-base articulations. - - Joints whose prims resolve to the same controller type, gains, - clamping chain, and delay configuration are merged into a single - :class:`Actuator` with combined index arrays, mirroring the - grouping the Newton builder performs internally. - - Args: - stage: The USD stage containing ``NewtonActuator`` prims. - joint_names: All joint names in the articulation. - num_envs: Number of environments. - num_joints: Joints per environment in the articulation. - device: Warp device string (e.g. ``"cuda:0"``). - articulation_prim_path: Root prim path of the first - environment's articulation (e.g. ``"/World/Env_0/Robot"``). - When provided, only ``NewtonActuator`` prims under this - subtree are considered — matching the scoped traversal - that Newton's ``ModelBuilder.add_usd`` performs. When - ``None``, the entire stage is scanned (legacy behaviour). - - Returns: - A fully constructed adapter ready for :meth:`finalize`. - """ - actuators = cls._create_actuators_from_usd( - stage, joint_names, num_envs, num_joints, device, - articulation_prim_path=articulation_prim_path, - ) - return cls(actuators, num_envs, num_joints, dof_offset=0, device=device) - - # -- public API ---------------------------------------------------------- - - def finalize(self) -> None: - """Read actuator gains and store as PyTorch tensors for DR.""" - wp_device = wp.get_device(self._device) - flat_stiffness = wp.zeros(self._num_envs * self.num_joints, dtype=wp.float32, device=wp_device) - flat_damping = wp.zeros(self._num_envs * self.num_joints, dtype=wp.float32, device=wp_device) - - for act in self.actuators: - ctrl = act.controller - if hasattr(ctrl, "kp"): - wp.launch( - _scatter_gain_kernel, dim=act.indices.shape[0], - inputs=[ctrl.kp, flat_stiffness, act.indices, self._dof_offset, self.num_joints], - device=wp_device, - ) - if hasattr(ctrl, "kd"): - wp.launch( - _scatter_gain_kernel, dim=act.indices.shape[0], - inputs=[ctrl.kd, flat_damping, act.indices, self._dof_offset, self.num_joints], - device=wp_device, - ) - - self.stiffness = wp.to_torch(flat_stiffness.reshape((self._num_envs, self.num_joints))) - self.damping = wp.to_torch(flat_damping.reshape((self._num_envs, self.num_joints))) - - def step(self, sim_state: Any, sim_control: Any, dt: float) -> None: - """Zero actuated DOFs, step all actuators, and swap state buffers. - - Args: - sim_state: Object with ``joint_q``, ``joint_qd``, etc. - Newton ``State`` on the Newton backend, - :class:`PhysxActuatorWrapper` on the PhysX backend. - sim_control: Object with ``joint_f``, ``joint_target_pos``, etc. - Newton ``Control`` on the Newton backend, - :class:`PhysxActuatorWrapper` on the PhysX backend. - dt: Physics timestep [s]. - """ - for act in self.actuators: - wp.launch( - _zero_at_indices_kernel, - dim=act.indices.shape[0], - inputs=[sim_control.joint_f, act.indices], - ) - for act, sa, sb in zip(self.actuators, self._states_a, self._states_b): - act.step(sim_state, sim_control, sa, sb, dt=dt) - self._states_a, self._states_b = self._states_b, self._states_a - - def reset(self, env_ids: Sequence[int] | slice | None = None) -> None: - """Reset actuator states for the given environments. - - Args: - env_ids: Environment indices to reset. ``None`` or - ``slice(None)`` resets all environments. - """ - if env_ids is None or env_ids == slice(None): - mask = None - else: - mask = wp.zeros(self._num_envs, dtype=wp.bool, device=self._device) - import torch # noqa: PLC0415 - if isinstance(env_ids, torch.Tensor): - idx = wp.from_torch(env_ids.to(device=self._device).contiguous().to(torch.int32), dtype=wp.int32) - else: - idx = wp.array(list(env_ids), dtype=wp.int32, device=self._device) - wp.launch(_set_mask_kernel, dim=len(idx), inputs=[mask, idx], device=self._device) - - for sa, sb in zip(self._states_a, self._states_b): - if sa is not None: - sa.reset(mask) - if sb is not None: - sb.reset(mask) - - @property - def is_all_graphable(self) -> bool: - """``True`` when all actuators are CUDA-graph-safe.""" - return len(self.actuators) > 0 and all(a.is_graphable() for a in self.actuators) - - def update_gain_at_env_ids( - self, - gain: str, - values: torch.Tensor | wp.array | float, - env_ids: wp.array, - ) -> wp.array: - """Scatter ``values`` into :attr:`stiffness` or :attr:`damping` at *env_ids*. - - Shared backend-independent step used by ``write_actuator_*_to_sim`` to - keep the kp/kd buffer the adapter exposes for DR in sync with what - each Newton actuator's controller will receive. - - Args: - gain: ``"stiffness"`` or ``"damping"``. - values: Per-env-per-joint values shape ``(len(env_ids), num_joints)``, - or a scalar to broadcast to every (env, joint). - env_ids: Warp int32 array of env indices to update. - - Returns: - Warp view of the updated ``(num_envs, num_joints)`` buffer. - """ - if gain == "stiffness": - buf = self.stiffness - elif gain == "damping": - buf = self.damping - else: - raise ValueError(f"gain must be 'stiffness' or 'damping', got {gain!r}") - buf_wp = wp.from_torch(buf, dtype=wp.float32) - if isinstance(values, float): - wp.launch( - _fill_gain_at_envs_kernel, - dim=(env_ids.shape[0], self.num_joints), - inputs=[values, env_ids], - outputs=[buf_wp], - device=self._device, - ) - else: - wp.launch( - _scatter_gain_at_envs_kernel, - dim=(env_ids.shape[0], self.num_joints), - inputs=[values, env_ids], - outputs=[buf_wp], - device=self._device, - ) - return buf_wp - - def write_stiffness_to_sim( - self, - stiffness: torch.Tensor | wp.array | float, - env_ids: wp.array, - env_mask: wp.array, - propagate_fn: Callable[["NewtonActuatorAdapter", Actuator, Any, str, wp.array, wp.array], None], - ) -> None: - """Update the kp buffer at *env_ids* and push the new values into each Newton controller. - - The per-actuator propagation step is backend-specific (Newton uses - :meth:`ArticulationView.set_actuator_parameter`; PhysX scatters via a - local Warp kernel), so the caller injects it as *propagate_fn*. - """ - self._write_gain_to_sim("stiffness", "kp", stiffness, env_ids, env_mask, propagate_fn) - - def write_damping_to_sim( - self, - damping: torch.Tensor | wp.array | float, - env_ids: wp.array, - env_mask: wp.array, - propagate_fn: Callable[["NewtonActuatorAdapter", Actuator, Any, str, wp.array, wp.array], None], - ) -> None: - """Update the kd buffer at *env_ids* and push the new values into each Newton controller.""" - self._write_gain_to_sim("damping", "kd", damping, env_ids, env_mask, propagate_fn) - - def _write_gain_to_sim( - self, - gain: str, - attr: str, - values: torch.Tensor | wp.array | float, - env_ids: wp.array, - env_mask: wp.array, - propagate_fn: Callable[["NewtonActuatorAdapter", Actuator, Any, str, wp.array, wp.array], None], - ) -> None: - """Shared body for :meth:`write_stiffness_to_sim` / :meth:`write_damping_to_sim`.""" - buf = self.update_gain_at_env_ids(gain, values, env_ids) - for newton_act in self.actuators: - ctrl = newton_act.controller - if hasattr(ctrl, attr): - propagate_fn(self, newton_act, ctrl, attr, buf, env_mask) - - # -- config helpers (used by both adapter and authoring) ----------------- - - @staticmethod - def _resolve_per_dof( - value: dict[str, float | int] | float | int | None, - joint_names: list[str], - cast: type = float, - ) -> dict[str, float | int]: - """Expand a scalar or dict config value into a per-DOF dict. - - When *value* is a dict, keys are treated as regex patterns and - matched against *joint_names* via :func:`re.fullmatch`. - """ - import re # noqa: PLC0415 - - if value is None: - return {} - if isinstance(value, (int, float)): - return {name: cast(value) for name in joint_names} - if isinstance(value, dict): - result: dict[str, float | int] = {} - for name in joint_names: - for pattern, v in value.items(): - if re.fullmatch(pattern, name): - result[name] = cast(v) - break - return result - return {} - - # -- private helpers ----------------------------------------------------- - - def _build_dof_map(self) -> list[int]: - """Build a per-DOF lookup: local DOF index -> actuator list index.""" - dof_to_actuator: list[int] = [-1] * self.num_joints - - for act_idx, act in enumerate(self.actuators): - all_indices = act.indices.numpy() - num_per_act = len(all_indices) // self._num_envs - env0_indices = all_indices[:num_per_act] - for global_dof in env0_indices: - local_dof = global_dof - self._dof_offset - if 0 <= local_dof < self.num_joints: - dof_to_actuator[local_dof] = act_idx - - return dof_to_actuator - - @staticmethod - def _actuator_signature(parsed: Any) -> tuple: - """Build a hashable key from a parsed actuator spec for grouping. - - Joints whose prims resolve to the same signature share identical - controller type, gains, clamping chain, and delay configuration and - can therefore be merged into a single :class:`Actuator` with combined - index arrays. - """ - ctrl_resolved = parsed.controller_class.resolve_arguments( - dict(parsed.controller_kwargs), - ) - ctrl_key = (parsed.controller_class, tuple(sorted(ctrl_resolved.items()))) - - comp_keys: list[tuple] = [] - for comp_cls, comp_kwargs in parsed.component_specs: - resolved = comp_cls.resolve_arguments(comp_kwargs) - comp_keys.append((comp_cls, tuple(sorted(resolved.items())))) - comp_keys.sort(key=lambda t: t[0].__name__) - - return (ctrl_key, tuple(comp_keys)) - - @staticmethod - def _create_actuators_from_usd( - stage: Any, - joint_names: list[str], - num_envs: int, - num_total_joints: int, - device: str, - articulation_prim_path: str | None = None, - ) -> list[Actuator]: - """Parse ``NewtonActuator`` prims and instantiate standalone actuators. - - This mirrors the actuator construction that Newton's - ``ModelBuilder.add_usd`` performs, but operates independently of a - Newton ``Model``. It is used on the PhysX backend where there is no - Newton simulation — actuators are stepped manually via the adapter. - - Because PhysX articulations have no free or ball joints, every - joint's coordinate count equals its DOF count. A single - ``indices`` array is therefore sufficient for all index roles - (``indices``, ``pos_indices``, ``target_pos_indices``). - - Joints with identical controller type, gains, clamping chain, and - delay are merged into one :class:`Actuator` with combined indices. - - Each per-DOF scalar parameter (``kp``, ``kd``, ``saturation_effort``, - etc.) is broadcast via :func:`wp.full` to match the group size. - Parameters marked as ``SHARED_PARAMS`` on the controller or clamping - class (e.g. ``model_path``, ``lookup_positions``) are passed through - directly without broadcast. - """ - from collections import defaultdict # noqa: PLC0415 - - from newton.actuators import parse_actuator_prim # noqa: PLC0415 - from pxr import Usd # noqa: PLC0415 - - wp_device = wp.get_device(device) - - joint_name_to_idx: dict[str, int] = {name: i for i, name in enumerate(joint_names)} - - if articulation_prim_path is not None: - root_prim = stage.GetPrimAtPath(articulation_prim_path) - else: - root_prim = stage.GetPseudoRoot() - - parsed_per_joint: dict[int, Any] = {} - for prim in Usd.PrimRange(root_prim): - parsed = parse_actuator_prim(prim) - if parsed is None: - continue - target_name = parsed.target_path.rsplit("/", 1)[-1] - if target_name in joint_name_to_idx: - parsed_per_joint[joint_name_to_idx[target_name]] = parsed - - if not parsed_per_joint: - raise ValueError( - f"No NewtonActuator prims found targeting any of: {joint_names}" - ) - - groups: dict[tuple, list[int]] = defaultdict(list) - sig_to_parsed: dict[tuple, Any] = {} - for local_idx, parsed in sorted(parsed_per_joint.items()): - sig = NewtonActuatorAdapter._actuator_signature(parsed) - groups[sig].append(local_idx) - if sig not in sig_to_parsed: - sig_to_parsed[sig] = parsed - - actuators = [] - for sig, local_indices in groups.items(): - parsed = sig_to_parsed[sig] - - flat_indices = np.array( - [idx + e * num_total_joints for e in range(num_envs) for idx in local_indices], - dtype=np.uint32, - ) - indices = wp.array(flat_indices, device=wp_device) - num_dofs_in_group = len(local_indices) * num_envs - - # Controller - ctrl_kwargs = dict(parsed.controller_kwargs) - resolved = parsed.controller_class.resolve_arguments(ctrl_kwargs) - shared_ctrl = getattr(parsed.controller_class, "SHARED_PARAMS", set()) - ctrl_arrays = {} - for key, val in resolved.items(): - if key in shared_ctrl: - ctrl_arrays[key] = val - else: - ctrl_arrays[key] = wp.full(num_dofs_in_group, float(val), dtype=wp.float32, device=wp_device) - controller = parsed.controller_class(**ctrl_arrays) - - # Components (delay + clampings) - clampings = [] - delay = None - for comp_cls, comp_kwargs in parsed.component_specs: - if issubclass(comp_cls, Delay): - resolved_kw = Delay.resolve_arguments(comp_kwargs) - delay_steps = int(resolved_kw.get("delay_steps", 0)) - if delay_steps > 0: - delay_arr = wp.full(num_dofs_in_group, delay_steps, dtype=wp.int32, device=wp_device) - delay = Delay(delay_steps=delay_arr, max_delay=delay_steps) - elif issubclass(comp_cls, Clamping): - resolved_kw = comp_cls.resolve_arguments(comp_kwargs) - shared_clamp = getattr(comp_cls, "SHARED_PARAMS", set()) - clamp_arrays = {} - for k, v in resolved_kw.items(): - if k in shared_clamp: - clamp_arrays[k] = v - else: - clamp_arrays[k] = wp.full( - num_dofs_in_group, float(v), dtype=wp.float32, device=wp_device, - ) - clampings.append(comp_cls(**clamp_arrays)) - - actuator = Actuator( - indices=indices, - controller=controller, - delay=delay, - clamping=clampings if clampings else None, - ) - actuators.append(actuator) - - return actuators - - -# =========================================================================== -# 4. USD authoring — create NewtonActuator prims from Lab actuator configs -# =========================================================================== - - -def author_newton_actuator_prims( - stage: Any, - articulation_prim_path: str, - actuator_cfgs: dict[str, Any], -) -> None: - """Author ``NewtonActuator`` USD prims from IsaacLab actuator configs. - - For every joint covered by an explicit (non-implicit) Lab actuator config, - any existing ``NewtonActuator`` prim targeting that joint is replaced by a - new one created from the config values. Joints **not** covered by any - Lab config keep their USD-authored actuators unchanged. - - The supported config-to-schema mapping is: - - * :class:`~isaaclab.actuators.IdealPDActuatorCfg` -> - ``NewtonPDControlAPI`` + ``NewtonMaxEffortClampingAPI`` - * :class:`~isaaclab.actuators.DCMotorCfg` -> - ``NewtonPDControlAPI`` + ``NewtonDCMotorClampingAPI`` - * :class:`~isaaclab.actuators.DelayedPDActuatorCfg` -> - same as ``IdealPDActuatorCfg`` + ``NewtonActuatorDelayAPI`` - * :class:`~isaaclab.actuators.RemotizedPDActuatorCfg` -> - same as ``DelayedPDActuatorCfg`` + ``NewtonPositionBasedClampingAPI`` - * :class:`~isaaclab.actuators.ActuatorNetMLPCfg` / - :class:`~isaaclab.actuators.ActuatorNetLSTMCfg` -> - ``NewtonNeuralControlAPI`` (+ ``NewtonDCMotorClampingAPI``) - - Must be called **after** the articulation is spawned (joint prims exist - on stage) and **before** the cloner / ``ModelBuilder.add_usd`` reads - the stage. - - Args: - stage: The USD stage to author prims on. - articulation_prim_path: Root prim path of the articulation - (e.g. ``"/World/Env_0/Robot"``). Must not contain wildcards. - actuator_cfgs: Mapping of group name to ``ActuatorBaseCfg``. - """ - from pxr import Sdf # noqa: PLC0415 - - from isaaclab.actuators import ImplicitActuator # noqa: PLC0415 - from isaaclab.utils.string import resolve_matching_names # noqa: PLC0415 - - art_prim = stage.GetPrimAtPath(articulation_prim_path) - if not art_prim.IsValid(): - raise ValueError(f"Articulation prim not found: {articulation_prim_path}") - - joint_inventory = _collect_joint_prims(art_prim) - all_joint_names = list(joint_inventory.keys()) - - covered_joint_paths: set[str] = set() - resolve = NewtonActuatorAdapter._resolve_per_dof - - cfg_entries: list[tuple[str, Any, list[str]]] = [] - for group_name, cfg in actuator_cfgs.items(): - cls_type = cfg.class_type - is_implicit = ( - "ImplicitActuator" in cls_type - if isinstance(cls_type, str) - else issubclass(cls_type, ImplicitActuator) - ) - if is_implicit: - continue - - _ids, joint_names = resolve_matching_names(cfg.joint_names_expr, all_joint_names) - if not joint_names: - continue - - cfg_entries.append((group_name, cfg, joint_names)) - for jname in joint_names: - covered_joint_paths.add(joint_inventory[jname]) - - _remove_actuator_prims_for_joints(art_prim, covered_joint_paths) - - from isaaclab.actuators import DCMotorCfg, DelayedPDActuatorCfg # noqa: PLC0415 - from isaaclab.actuators.actuator_net_cfg import ActuatorNetLSTMCfg, ActuatorNetMLPCfg # noqa: PLC0415 - from isaaclab.actuators.actuator_pd_cfg import IdealPDActuatorCfg, RemotizedPDActuatorCfg # noqa: PLC0415 - - _SUPPORTED_CFG_TYPES = ( - IdealPDActuatorCfg, - DCMotorCfg, - DelayedPDActuatorCfg, - RemotizedPDActuatorCfg, - ActuatorNetMLPCfg, - ActuatorNetLSTMCfg, - ) - - for group_name, cfg, joint_names in cfg_entries: - if not isinstance(cfg, _SUPPORTED_CFG_TYPES): - logger.warning( - "Actuator group '%s' uses config type '%s' which is not supported by Newton-native" - " actuator authoring. The group will be skipped.", - group_name, - type(cfg).__name__, - ) - continue - stiffness_map = resolve(getattr(cfg, "stiffness", None), joint_names) - damping_map = resolve(getattr(cfg, "damping", None), joint_names) - effort_map = resolve(getattr(cfg, "effort_limit", None), joint_names) - - is_neural = isinstance(cfg, (ActuatorNetMLPCfg, ActuatorNetLSTMCfg)) - is_remotized = isinstance(cfg, RemotizedPDActuatorCfg) - is_dc_motor = isinstance(cfg, DCMotorCfg) - is_delayed = isinstance(cfg, DelayedPDActuatorCfg) - - vel_limit_map = resolve(getattr(cfg, "velocity_limit", None), joint_names) if is_dc_motor else {} - sat_effort_map = resolve(getattr(cfg, "saturation_effort", None), joint_names) if is_dc_motor else {} - - raw_delay = getattr(cfg, "max_delay", 0) if is_delayed else 0 - delay_map = resolve(raw_delay, joint_names, cast=int) if raw_delay else {} - - patched_model_path: str | None = None - if is_neural: - meta: dict[str, Any] = {} - if isinstance(cfg, ActuatorNetMLPCfg): - meta["model_type"] = "mlp" - meta["input_order"] = cfg.input_order - meta["input_idx"] = list(cfg.input_idx) - meta["pos_scale"] = cfg.pos_scale - meta["vel_scale"] = cfg.vel_scale - meta["torque_scale"] = cfg.torque_scale - else: - meta["model_type"] = "lstm" - patched_model_path = _resave_checkpoint_with_metadata(cfg.network_file, meta) - - for jname in joint_names: - joint_prim_path = joint_inventory[jname] - - schemas: list[str] = [] - attrs: dict[str, float | int] = {} - array_attrs: dict[str, list[float]] = {} - - if is_neural: - schemas.append("NewtonNeuralControlAPI") - else: - schemas.append("NewtonPDControlAPI") - attrs["kp"] = stiffness_map.get(jname, 0.0) - attrs["kd"] = damping_map.get(jname, 0.0) - - if is_dc_motor: - schemas.append("NewtonDCMotorClampingAPI") - attrs["saturation_effort"] = sat_effort_map.get(jname, 0.0) - if jname in vel_limit_map: - attrs["velocity_limit"] = vel_limit_map[jname] - if jname in effort_map: - attrs["max_motor_effort"] = effort_map[jname] - elif jname in effort_map: - schemas.append("NewtonMaxEffortClampingAPI") - attrs["max_effort"] = effort_map[jname] - - if is_remotized and isinstance(cfg, RemotizedPDActuatorCfg): - lookup = cfg.joint_parameter_lookup - schemas.append("NewtonPositionBasedClampingAPI") - array_attrs["lookup_positions"] = [row[0] for row in lookup] - array_attrs["lookup_efforts"] = [row[2] for row in lookup] - - delay_steps = delay_map.get(jname, 0) - if delay_steps > 0: - schemas.append("NewtonActuatorDelayAPI") - attrs["delay_steps"] = delay_steps - attrs["max_delay"] = delay_steps - - act_prim_path = f"{articulation_prim_path}/{group_name}_{jname}_actuator" - act_prim = stage.DefinePrim(act_prim_path, "NewtonActuator") - - existing = act_prim.GetMetadata("apiSchemas") or Sdf.TokenListOp() - existing.prependedItems = list(schemas) - act_prim.SetMetadata("apiSchemas", existing) - - rel = act_prim.CreateRelationship("newton:targets") - rel.SetTargets([Sdf.Path(joint_prim_path)]) - - if patched_model_path is not None: - act_prim.CreateAttribute("newton:modelPath", Sdf.ValueTypeNames.Asset).Set( - Sdf.AssetPath(patched_model_path) - ) - - for attr_name, attr_val in attrs.items(): - usd_name = f"newton:{_snake_to_camel(attr_name)}" - if isinstance(attr_val, int): - act_prim.CreateAttribute(usd_name, Sdf.ValueTypeNames.Int).Set(attr_val) - else: - act_prim.CreateAttribute(usd_name, Sdf.ValueTypeNames.Float).Set(float(attr_val)) - - for attr_name, attr_val in array_attrs.items(): - usd_name = f"newton:{_snake_to_camel(attr_name)}" - act_prim.CreateAttribute(usd_name, Sdf.ValueTypeNames.FloatArray).Set(attr_val) - - -# --------------------------------------------------------------------------- -# USD authoring — private helpers -# --------------------------------------------------------------------------- - -_SNAKE_TO_CAMEL_RE = __import__("re").compile(r"_([a-z])") - - -def _snake_to_camel(name: str) -> str: - """Convert a snake_case name to camelCase.""" - return _SNAKE_TO_CAMEL_RE.sub(lambda m: m.group(1).upper(), name) - - -def _collect_joint_prims(art_prim: Any) -> dict[str, str]: - """Collect all joint prims under an articulation subtree. - - Returns: - Ordered mapping of joint name to full prim path. - """ - from pxr import Usd # noqa: PLC0415 - - _JOINT_TYPES = {"PhysicsRevoluteJoint", "PhysicsPrismaticJoint"} - - joints: dict[str, str] = {} - for prim in Usd.PrimRange(art_prim): - if prim.GetTypeName() in _JOINT_TYPES: - joints[prim.GetName()] = str(prim.GetPath()) - return joints - - -def _remove_actuator_prims_for_joints( - art_prim: Any, - joint_paths: set[str], -) -> None: - """Deactivate ``NewtonActuator`` prims whose target is in *joint_paths*. - - Deactivated prims are invisible to ``Usd.PrimRange`` and therefore - ignored by ``ModelBuilder.add_usd``. Using ``SetActive(False)`` - instead of ``RemovePrim`` works correctly when the prim originates - from a USD reference or payload. - - Only prims under the *art_prim* subtree are considered. - """ - from pxr import Usd # noqa: PLC0415 - - to_deactivate: list = [] - for prim in Usd.PrimRange(art_prim): - if prim.GetTypeName() != "NewtonActuator": - continue - rel = prim.GetRelationship("newton:targets") - if rel and rel.IsValid(): - for target in rel.GetTargets(): - if str(target) in joint_paths: - to_deactivate.append(prim) - break - - for prim in to_deactivate: - prim.SetActive(False) - - -def _resave_checkpoint_with_metadata( - original_path: str, - metadata: dict[str, Any], -) -> str: - """Re-save a neural-network checkpoint with updated metadata. - - Loads the original TorchScript or dict checkpoint, merges *metadata* - into any existing metadata (Lab config values take precedence), and - writes the result to a temporary ``.pt`` file that persists for the - lifetime of the process. - - Returns: - Path to the temporary checkpoint file. - """ - import json - import tempfile - - import torch - - extra_files: dict[str, str] = {"metadata.json": ""} - is_torchscript = True - try: - net = torch.jit.load(original_path, map_location="cpu", _extra_files=extra_files) - existing_meta = json.loads(extra_files["metadata.json"]) if extra_files["metadata.json"] else {} - except Exception: - is_torchscript = False - checkpoint = torch.load(original_path, map_location="cpu", weights_only=False) - if not isinstance(checkpoint, dict) or "model" not in checkpoint: - raise ValueError( - f"Cannot load checkpoint at '{original_path}'; " - "expected a TorchScript archive or a dict with a 'model' key" - ) - net = checkpoint["model"] - existing_meta = checkpoint.get("metadata", {}) - - merged = {**existing_meta, **metadata} - - tmp = tempfile.NamedTemporaryFile(suffix=".pt", delete=False) - if is_torchscript: - extra_out = {"metadata.json": json.dumps(merged)} - torch.jit.save(net, tmp.name, _extra_files=extra_out) - else: - torch.save({"model": net, "metadata": merged}, tmp.name) - - return tmp.name diff --git a/source/isaaclab_newton/isaaclab_newton/actuators/physx_wrapper.py b/source/isaaclab_newton/isaaclab_newton/actuators/physx_wrapper.py new file mode 100644 index 000000000000..35dadc684524 --- /dev/null +++ b/source/isaaclab_newton/isaaclab_newton/actuators/physx_wrapper.py @@ -0,0 +1,60 @@ +# Copyright (c) 2022-2026, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +"""PhysX-only stepping helper for :class:`~newton.actuators.Actuator`. + +Newton's :meth:`Actuator.step` requires a ``sim_state`` / ``sim_control`` +pair that exposes flat 1-D Warp arrays (``joint_q``, ``joint_qd``, +``joint_target_pos``, ``joint_f``, …). On the **Newton backend** these +are the ``State`` and ``Control`` objects that the solver already owns — +no wrapper is needed because: + +* The solver manages double-buffered ``State`` objects for CUDA-graph + capture, and actuators are stepped inside the solver's own simulation + loop where states are already available. +* Wrapping them would add indirection with no benefit; the Newton + articulation code that calls :meth:`Actuator.step` lives in + ``newton_manager.py`` and has direct access to the model's state. + +On the **PhysX backend**, no Newton solver exists — the actuators are +stepped manually from the Lab articulation's ``write_data_to_sim`` +path. Isaac Lab stores joint data as 2-D tensors (``num_envs × +num_joints``), so :class:`PhysxActuatorWrapper` provides zero-copy flat +views that satisfy the protocol without allocating new memory. +""" + +from __future__ import annotations + +from dataclasses import dataclass + +import warp as wp + + +@dataclass +class PhysxActuatorWrapper: + """Flat-array wrapper serving as ``sim_state`` / ``sim_control`` for + :meth:`Actuator.step` on the PhysX backend. + + Most attributes are bound once at articulation init to zero-copy flat + views of Isaac Lab's 2-D buffers. ``joint_f_2d`` is the only persistent + allocation, sized via :meth:`create`; ``joint_f`` is its flat alias + consumed by the Newton actuator step. + """ + + joint_q: wp.array | None = None + joint_qd: wp.array | None = None + joint_target_pos: wp.array | None = None + joint_target_vel: wp.array | None = None + joint_act: wp.array | None = None + joint_f: wp.array | None = None + joint_f_2d: wp.array | None = None + + @classmethod + def create(cls, num_envs: int, num_joints: int, device: str) -> "PhysxActuatorWrapper": + """Allocate the persistent ``joint_f`` buffer for the given articulation shape.""" + w = cls() + w.joint_f_2d = wp.zeros((num_envs, num_joints), dtype=wp.float32, device=device) + w.joint_f = w.joint_f_2d.reshape(-1) + return w diff --git a/source/isaaclab_newton/isaaclab_newton/assets/articulation/articulation.py b/source/isaaclab_newton/isaaclab_newton/assets/articulation/articulation.py index fbabbcfa49fd..b4ff8c998d19 100644 --- a/source/isaaclab_newton/isaaclab_newton/assets/articulation/articulation.py +++ b/source/isaaclab_newton/isaaclab_newton/assets/articulation/articulation.py @@ -27,11 +27,8 @@ from isaaclab.assets.articulation.base_articulation import BaseArticulation try: + from isaaclab_newton.actuators import NewtonActuatorAdapter, build_implicit_dof_mask from isaaclab_newton.actuators import kernels as actuator_kernels - from isaaclab_newton.actuators.newton_actuator_utils import ( - NewtonActuatorAdapter, - build_implicit_dof_mask, - ) _HAS_NEWTON_ACTUATORS = True except ImportError: @@ -57,12 +54,6 @@ logger = logging.getLogger(__name__) -@wp.kernel(enable_backward=False) -def _build_env_mask_kernel(mask: wp.array(dtype=wp.bool), indices: wp.array(dtype=wp.int32)): - i = wp.tid() - mask[indices[i]] = True - - class Articulation(BaseArticulation): """An articulation asset class. @@ -2034,28 +2025,18 @@ def write_joint_friction_coefficient_to_sim_mask( Operations - Newton Actuator Parameter Writers. """ - def _propagate_gain_via_view( - self, - adapter: NewtonActuatorAdapter, - actuator, - controller, - attr: str, - values: wp.array, - env_mask: wp.array, - ) -> None: - """Per-actuator gain propagation using Newton's simulator-side scatter API.""" - del adapter # Newton path needs only the view + per-actuator info. - self._root_view.set_actuator_parameter( - actuator=actuator, component=controller, name=attr, - values=values, mask=env_mask, - ) + @staticmethod + @wp.kernel(enable_backward=False) + def _build_env_mask_kernel(mask: wp.array(dtype=wp.bool), indices: wp.array(dtype=wp.int32)): + i = wp.tid() + mask[indices[i]] = True def _env_ids_to_mask(self, env_ids: wp.array) -> wp.array: """Convert warp env_ids to a boolean Warp mask.""" if env_ids is self._ALL_INDICES: return self._ALL_ENV_MASK mask = wp.zeros(self.num_instances, dtype=wp.bool, device=self.device) - wp.launch(_build_env_mask_kernel, dim=env_ids.shape[0], inputs=[mask, env_ids], device=self.device) + wp.launch(self._build_env_mask_kernel, dim=env_ids.shape[0], inputs=[mask, env_ids], device=self.device) return mask """ @@ -3412,7 +3393,7 @@ def _process_actuators_cfg(self): # Build (or share) the single sim-level actuator adapter. Idempotent — # the first articulation to call this constructs it from # ``model.actuators``; subsequent articulations reuse it. - SimulationManager.ensure_global_actuator_adapter() + SimulationManager.ensure_global_actuator_adapter(self._root_view) # Zero the simulator's joint-drive PD on DOFs covered by an explicit # Lab actuator config in *this* articulation. The global Newton diff --git a/source/isaaclab_newton/isaaclab_newton/physics/newton_manager.py b/source/isaaclab_newton/isaaclab_newton/physics/newton_manager.py index dff968a306e2..f656e69c982a 100644 --- a/source/isaaclab_newton/isaaclab_newton/physics/newton_manager.py +++ b/source/isaaclab_newton/isaaclab_newton/physics/newton_manager.py @@ -1388,24 +1388,25 @@ def activate_newton_actuator_path(cls) -> None: cls._use_newton_actuators_active = True @classmethod - def ensure_global_actuator_adapter(cls) -> None: + def ensure_global_actuator_adapter(cls, root_view: Any) -> None: """Build the single sim-level :class:`NewtonActuatorAdapter` if not yet built. Idempotent — each Newton-fast-path articulation calls this from - ``_process_actuators_cfg`` after the model is finalized; the first - call constructs the adapter from ``cls._model.actuators`` and - subsequent calls are no-ops. A single global adapter avoids the - multi-articulation clobbering that the per-articulation - ``register_adapter`` design suffered from. The adapter operates on - the full flat DOF layout (``dof_offset=0``, ``num_joints`` = - per-env total DOF count), so it iterates and steps every actuator - in the model in one pass. + ``_process_actuators_cfg`` after the model is finalized, passing + its own ``root_view``; the first call constructs the adapter from + ``cls._model.actuators`` and configures the gain propagator with + the supplied view, subsequent calls are no-ops. A single global + adapter avoids the multi-articulation clobbering that the + per-articulation ``register_adapter`` design suffered from. The + adapter operates on the full flat DOF layout (``dof_offset=0``, + ``num_joints`` = per-env total DOF count), so it iterates and + steps every actuator in the model in one pass. """ if cls._adapter is not None: return if cls._model is None or not cls._model.actuators: return - from isaaclab_newton.actuators.newton_actuator_utils import NewtonActuatorAdapter # noqa: PLC0415 + from isaaclab_newton.actuators import NewtonActuatorAdapter # noqa: PLC0415 dofs_per_env = cls._model.joint_dof_count // cls._num_envs cls._adapter = NewtonActuatorAdapter( @@ -1416,6 +1417,10 @@ def ensure_global_actuator_adapter(cls) -> None: device=PhysicsManager._device, ) cls._adapter.finalize() + # ``set_actuator_parameter`` dispatches on the actuator object (which + # is model-scoped), so any one articulation's view works for the + # global adapter spanning multiple articulations. + cls._adapter.set_view_propagator(root_view) @classmethod def register_post_actuator_callback(cls, callback: Callable[[], None]) -> None: diff --git a/source/isaaclab_physx/isaaclab_physx/assets/articulation/articulation.py b/source/isaaclab_physx/isaaclab_physx/assets/articulation/articulation.py index 80f60b3e04de..17d17267d498 100644 --- a/source/isaaclab_physx/isaaclab_physx/assets/articulation/articulation.py +++ b/source/isaaclab_physx/isaaclab_physx/assets/articulation/articulation.py @@ -24,13 +24,12 @@ from isaaclab.assets.articulation.base_articulation import BaseArticulation try: - from isaaclab_newton.actuators import kernels as actuator_kernels - from isaaclab_newton.actuators.newton_actuator_utils import ( + from isaaclab_newton.actuators import ( NewtonActuatorAdapter, PhysxActuatorWrapper, - _gather_gain_kernel, build_implicit_dof_mask, ) + from isaaclab_newton.actuators import kernels as actuator_kernels _HAS_NEWTON_ACTUATORS = True except ImportError: @@ -1225,7 +1224,7 @@ def write_actuator_stiffness_to_sim( return env_ids = self._resolve_env_ids(env_ids) env_mask = self._env_ids_to_mask(env_ids) - adapter.write_stiffness_to_sim(stiffness, env_ids, env_mask, self._propagate_gain_via_kernel) + adapter.write_stiffness_to_sim(stiffness, env_ids, env_mask) def write_actuator_damping_to_sim( self, @@ -1239,32 +1238,7 @@ def write_actuator_damping_to_sim( return env_ids = self._resolve_env_ids(env_ids) env_mask = self._env_ids_to_mask(env_ids) - adapter.write_damping_to_sim(damping, env_ids, env_mask, self._propagate_gain_via_kernel) - - def _propagate_gain_via_kernel( - self, - adapter: NewtonActuatorAdapter, - actuator, - controller, - attr: str, - values: wp.array, - env_mask: wp.array, - ) -> None: - """Per-actuator gain propagation using a local Warp scatter kernel. - - Newton has :meth:`ArticulationView.set_actuator_parameter` for this; - on PhysX we use :data:`_gather_gain_kernel` since no equivalent - simulator-side API exists. - """ - wp.launch( - _gather_gain_kernel, - dim=actuator.indices.shape[0], - inputs=[ - values.flatten(), getattr(controller, attr), actuator.indices, env_mask, - adapter._dof_offset, adapter.num_joints, - ], - device=self.device, - ) + adapter.write_damping_to_sim(damping, env_ids, env_mask) def _env_ids_to_mask(self, env_ids: wp.array) -> wp.array: """Convert warp ``env_ids`` to a boolean Warp mask of length ``num_instances``.""" @@ -3890,6 +3864,7 @@ def _process_actuators_cfg(self): w.joint_target_vel = self._data.joint_vel_target.warp.reshape(-1) w.joint_act = self._data.joint_effort_target.warp.reshape(-1) adapter.finalize() + adapter.set_kernel_propagator() self.actuators["newton"] = adapter self.write_joint_stiffness_to_sim_index(stiffness=0.0, joint_ids=adapter.joint_indices) self.write_joint_damping_to_sim_index(damping=0.0, joint_ids=adapter.joint_indices) From 0d8b79421e2f37b06a0a2d5c91f1e57143468948 Mon Sep 17 00:00:00 2001 From: Antoine Richard Date: Fri, 8 May 2026 15:48:20 +0200 Subject: [PATCH 17/35] Apply trivial review nits to actuator integration MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Address three small review comments on PR #5455: - PhysicsManager.set_decimation now has an explicit ``pass`` body to match the other no-op classmethods in the base. - _decimation declaration moves next to _num_substeps in NewtonManager for consistency with related solver-stepping configuration. - Reword FF-routing comments in PhysX Articulation to talk about actuated DOFs rather than implicit DOFs — the surrounding kernel operates on the full actuated DOF set. --- .../isaaclab/physics/physics_manager.py | 1 + .../isaaclab_newton/physics/newton_manager.py | 2 +- .../assets/articulation/articulation.py | 42 ++++++++++++------- 3 files changed, 30 insertions(+), 15 deletions(-) diff --git a/source/isaaclab/isaaclab/physics/physics_manager.py b/source/isaaclab/isaaclab/physics/physics_manager.py index ca333673bb48..e6157177ded4 100644 --- a/source/isaaclab/isaaclab/physics/physics_manager.py +++ b/source/isaaclab/isaaclab/physics/physics_manager.py @@ -351,6 +351,7 @@ def set_decimation(cls, decimation: int) -> None: Args: decimation: Number of physics steps per environment step. """ + pass @classmethod def handles_decimation(cls) -> bool: diff --git a/source/isaaclab_newton/isaaclab_newton/physics/newton_manager.py b/source/isaaclab_newton/isaaclab_newton/physics/newton_manager.py index f656e69c982a..521205383c40 100644 --- a/source/isaaclab_newton/isaaclab_newton/physics/newton_manager.py +++ b/source/isaaclab_newton/isaaclab_newton/physics/newton_manager.py @@ -114,6 +114,7 @@ class NewtonManager(PhysicsManager): _solver_dt: float = 1.0 / 200.0 _num_substeps: int = 1 + _decimation: int = 1 _num_envs: int | None = None # Newton model and state @@ -148,7 +149,6 @@ class NewtonManager(PhysicsManager): # Newton actuator adapter (owns actuators and double-buffered states) _adapter: NewtonActuatorAdapter | None = None - _decimation: int = 1 # In-graph hooks invoked after the actuator step and before the solver # substeps, in registration order. Multiple articulations register their # implicit-DOF telemetry / FF-routing kernels here. diff --git a/source/isaaclab_physx/isaaclab_physx/assets/articulation/articulation.py b/source/isaaclab_physx/isaaclab_physx/assets/articulation/articulation.py index 17d17267d498..30c53c4fd02e 100644 --- a/source/isaaclab_physx/isaaclab_physx/assets/articulation/articulation.py +++ b/source/isaaclab_physx/isaaclab_physx/assets/articulation/articulation.py @@ -40,6 +40,8 @@ def _build_env_mask_kernel(mask: wp.array(dtype=wp.bool), indices: wp.array(dtype=wp.int32)): i = wp.tid() mask[indices[i]] = True + + from isaaclab.sim.utils.queries import find_first_matching_prim, get_all_matching_child_prims from isaaclab.utils.string import resolve_matching_names, resolve_matching_names_values from isaaclab.utils.types import ArticulationActions @@ -272,12 +274,13 @@ def write_data_to_sim(self): if self._has_newton_actuators: # Newton fast path: pos/vel targets pass straight through; the # in-graph kernel inside ``_apply_actuator_model_newton`` merges - # Newton's actuator output (explicit DOFs) with user FF - # (implicit DOFs) into ``w.joint_f_2d``, which is what we push - # to PhysX as the actuation force. + # Newton's actuator output with the user FF for actuated DOFs + # into ``w.joint_f_2d``, which is what we push to PhysX as the + # actuation force. self._apply_actuator_model_newton() self.root_view.set_dof_actuation_forces( - self._physx_actuator_wrapper.joint_f_2d, self._ALL_INDICES, + self._physx_actuator_wrapper.joint_f_2d, + self._ALL_INDICES, ) if self._has_implicit_actuators: self.root_view.set_dof_position_targets(self._data._joint_pos_target, self._ALL_INDICES) @@ -3882,7 +3885,9 @@ def _process_actuators_cfg(self): self._create_lab_actuator(actuator_name, actuator_cfg, properties_only=True) self._implicit_dof_mask = build_implicit_dof_mask( - self.actuators, self.num_joints, self.device, + self.actuators, + self.num_joints, + self.device, ) return @@ -3914,7 +3919,11 @@ def _process_actuators_cfg(self): logger.warning(f"\nActuatorCfg-USD Value Discrepancy Resolution (matching values are skipped): \n{t}") def _create_lab_actuator( - self, actuator_name: str, actuator_cfg: ActuatorBaseCfg, *, properties_only: bool = False, + self, + actuator_name: str, + actuator_cfg: ActuatorBaseCfg, + *, + properties_only: bool = False, ) -> None: """Instantiate a single Lab actuator from its config and write properties to sim. @@ -3955,20 +3964,25 @@ def _create_lab_actuator( # Write physical joint properties (armature, limits, friction) — always needed. self.write_joint_effort_limit_to_sim_index( - limits=actuator.effort_limit_sim, joint_ids=actuator.joint_indices, + limits=actuator.effort_limit_sim, + joint_ids=actuator.joint_indices, ) self.write_joint_velocity_limit_to_sim_index( - limits=actuator.velocity_limit_sim, joint_ids=actuator.joint_indices, + limits=actuator.velocity_limit_sim, + joint_ids=actuator.joint_indices, ) self.write_joint_armature_to_sim_index(armature=actuator.armature, joint_ids=actuator.joint_indices) self.write_joint_friction_coefficient_to_sim_index( - joint_friction_coeff=actuator.friction, joint_ids=actuator.joint_indices, + joint_friction_coeff=actuator.friction, + joint_ids=actuator.joint_indices, ) self.write_joint_dynamic_friction_coefficient_to_sim_index( - joint_dynamic_friction_coeff=actuator.dynamic_friction, joint_ids=actuator.joint_indices, + joint_dynamic_friction_coeff=actuator.dynamic_friction, + joint_ids=actuator.joint_indices, ) self.write_joint_viscous_friction_coefficient_to_sim_index( - joint_viscous_friction_coeff=actuator.viscous_friction, joint_ids=actuator.joint_indices, + joint_viscous_friction_coeff=actuator.viscous_friction, + joint_ids=actuator.joint_indices, ) if properties_only: @@ -4100,11 +4114,11 @@ def _apply_actuator_model(self): def _apply_actuator_model_newton(self): """Step Newton actuators (when present) then route FF + sync telemetry. - After ``newton_adapter.step`` (no-op if no explicit Newton actuators - exist), ``w.joint_f_2d`` holds the actuator output for explicit DOFs + After ``newton_adapter.step`` (no-op if no Newton actuators exist), + ``w.joint_f_2d`` holds the actuator output for the actuated DOFs and zero elsewhere. The :func:`synch_torque_and_apply_implicit_feedforwards` kernel then - writes the user FF into ``joint_f_2d`` for implicit DOFs and fills + writes the user FF into ``joint_f_2d`` for actuated DOFs and fills ``_data._computed_torque`` / ``_data._applied_torque``. The resulting ``joint_f_2d`` is what gets pushed to PhysX as the actuation force in :meth:`write_data_to_sim`. From 35938c16163646a93077fbd8b02ae9bfc1e6f7b5 Mon Sep 17 00:00:00 2001 From: Antoine Richard Date: Fri, 8 May 2026 16:11:28 +0200 Subject: [PATCH 18/35] Restore SolverKamino import after develop merge The develop branch dropped solver-specific imports from newton_manager in favour of subclass-defined hooks, but the actuator-integration branch still references SolverKamino in _capture_or_defer_graph for the Kamino-specific replay-once after CUDA graph capture. --- .../isaaclab_newton/isaaclab_newton/physics/newton_manager.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/source/isaaclab_newton/isaaclab_newton/physics/newton_manager.py b/source/isaaclab_newton/isaaclab_newton/physics/newton_manager.py index ac7e5cdfd1e1..41b36f6ccca7 100644 --- a/source/isaaclab_newton/isaaclab_newton/physics/newton_manager.py +++ b/source/isaaclab_newton/isaaclab_newton/physics/newton_manager.py @@ -31,7 +31,7 @@ from newton.sensors import SensorContact as NewtonContactSensor from newton.sensors import SensorFrameTransform from newton.sensors import SensorIMU as NewtonSensorIMU -from newton.solvers import SolverBase, SolverNotifyFlags +from newton.solvers import SolverBase, SolverKamino, SolverNotifyFlags from isaaclab.physics import CallbackHandle, PhysicsEvent, PhysicsManager from isaaclab.sim.utils.newton_model_utils import replace_newton_shape_colors From ec957c962b41ea3d2daec3323087ce535a5a6ff2 Mon Sep 17 00:00:00 2001 From: Antoine Richard Date: Fri, 8 May 2026 16:27:54 +0200 Subject: [PATCH 19/35] Lift actuator USD authoring into the schema layer Moves NewtonActuator USD authoring from the actuator extension into isaaclab/sim/schemas/, exposed as `define_actuator_properties` next to the other `define_*_properties` writers, and routes invocation through the spawner so that BaseArticulation no longer reaches into isaaclab_newton.actuators. * New `isaaclab.sim.schemas.define_actuator_properties` (ex `author_newton_actuator_prims`); the old `isaaclab_newton/actuators/authoring.py` is removed. * New `actuator_props` field on `FileCfg`, auto-populated from `ArticulationCfg.actuators` via `__post_init__`. * The from-files spawner authors actuator prims as part of its regular schema-modification pipeline. * `BaseArticulation.__init__` drops the lazy isaaclab_newton import and authoring call. Addresses review comments #3188200527, #3188205400, #3188207898 on PR isaac-sim/IsaacLab#5455. --- .../assets/articulation/articulation_cfg.py | 17 +++ .../assets/articulation/base_articulation.py | 8 -- .../isaaclab/sim/schemas/__init__.pyi | 4 + .../sim/schemas/schemas_actuators.py} | 133 +++++++++--------- .../sim/spawners/from_files/from_files.py | 4 + .../sim/spawners/from_files/from_files_cfg.py | 14 ++ .../isaaclab_newton/actuators/adapter.py | 71 ++++++---- 7 files changed, 150 insertions(+), 101 deletions(-) rename source/{isaaclab_newton/isaaclab_newton/actuators/authoring.py => isaaclab/isaaclab/sim/schemas/schemas_actuators.py} (80%) diff --git a/source/isaaclab/isaaclab/assets/articulation/articulation_cfg.py b/source/isaaclab/isaaclab/assets/articulation/articulation_cfg.py index d3d94a9b23e0..b7cb3cafd9c7 100644 --- a/source/isaaclab/isaaclab/assets/articulation/articulation_cfg.py +++ b/source/isaaclab/isaaclab/assets/articulation/articulation_cfg.py @@ -74,3 +74,20 @@ class InitialStateCfg(AssetBaseCfg.InitialStateCfg): actuator_value_resolution_debug_print = False """Print the resolution of actuator final value when input cfg is different from USD value, Defaults to False """ + + def __post_init__(self) -> None: + """Propagate :attr:`actuators` to the spawn cfg's ``actuator_props`` field. + + The spawner uses ``actuator_props`` to author ``NewtonActuator`` USD prims + via :func:`~isaaclab.sim.schemas.define_actuator_properties` (a no-op when + ``use_newton_actuators`` is disabled). Keeping the propagation here means + users only declare actuators once, on the asset cfg. + + User-provided ``spawn.actuator_props`` is preserved and takes precedence. + """ + if ( + self.spawn is not None + and self.actuators is not MISSING + and getattr(self.spawn, "actuator_props", None) is None + ): + self.spawn.actuator_props = dict(self.actuators) diff --git a/source/isaaclab/isaaclab/assets/articulation/base_articulation.py b/source/isaaclab/isaaclab/assets/articulation/base_articulation.py index 786d2223f95e..1b9888d744bc 100644 --- a/source/isaaclab/isaaclab/assets/articulation/base_articulation.py +++ b/source/isaaclab/isaaclab/assets/articulation/base_articulation.py @@ -101,14 +101,6 @@ def __init__(self, cfg: ArticulationCfg): sim_ctx = SimulationContext.instance() self._sim_cfg = sim_ctx.cfg if sim_ctx is not None else None - - try: - from isaaclab_newton.actuators import author_actuator_prims_for_articulation # noqa: PLC0415 - except ImportError: - pass - else: - author_actuator_prims_for_articulation(self.cfg, self._sim_cfg, self.stage) - """ Properties """ diff --git a/source/isaaclab/isaaclab/sim/schemas/__init__.pyi b/source/isaaclab/isaaclab/sim/schemas/__init__.pyi index 9a90ed0d810d..6820e8fd87ef 100644 --- a/source/isaaclab/isaaclab/sim/schemas/__init__.pyi +++ b/source/isaaclab/isaaclab/sim/schemas/__init__.pyi @@ -8,6 +8,7 @@ __all__ = [ "PHYSX_MESH_COLLISION_CFGS", "USD_MESH_COLLISION_CFGS", "activate_contact_sensors", + "define_actuator_properties", "define_articulation_root_properties", "define_collision_properties", "define_mass_properties", @@ -50,6 +51,9 @@ from .schemas import ( modify_rigid_body_properties, modify_spatial_tendon_properties, ) +from .schemas_actuators import ( + define_actuator_properties, +) from .schemas_cfg import ( ArticulationRootBaseCfg, BoundingCubePropertiesCfg, diff --git a/source/isaaclab_newton/isaaclab_newton/actuators/authoring.py b/source/isaaclab/isaaclab/sim/schemas/schemas_actuators.py similarity index 80% rename from source/isaaclab_newton/isaaclab_newton/actuators/authoring.py rename to source/isaaclab/isaaclab/sim/schemas/schemas_actuators.py index 999870e526f1..4f8c8e60602f 100644 --- a/source/isaaclab_newton/isaaclab_newton/actuators/authoring.py +++ b/source/isaaclab/isaaclab/sim/schemas/schemas_actuators.py @@ -3,19 +3,19 @@ # # SPDX-License-Identifier: BSD-3-Clause -"""USD authoring for Newton-native actuators. +"""USD schema authoring for Newton-native actuators. -:func:`author_newton_actuator_prims` translates IsaacLab actuator configs +:func:`define_actuator_properties` translates IsaacLab actuator configs into ``NewtonActuator`` USD prims. Both the Newton ``ModelBuilder.add_usd`` path and the PhysX adapter's :meth:`~isaaclab_newton.actuators.adapter.NewtonActuatorAdapter.from_usd` read the same authored prims, ensuring both backends construct :class:`~newton.actuators.Actuator` instances with matching parameters. -:func:`author_actuator_prims_for_articulation` is a thin wrapper that -resolves per-articulation pre-conditions (sim cfg gating, prim lookup) -and dispatches into :func:`author_newton_actuator_prims`. It is invoked -from the schema-side articulation initialisation path. +This module lives on the schema side so that authoring is a regular +``define_*_properties`` step in the spawner pipeline, alongside +:func:`define_articulation_root_properties` and friends, rather than a +side effect of asset construction. """ from __future__ import annotations @@ -34,7 +34,7 @@ def resolve_per_dof( ) -> dict[str, float | int]: """Expand a scalar or regex-keyed dict cfg value into a per-joint mapping. - Used by :func:`author_newton_actuator_prims` to flatten the various + Used by :func:`define_actuator_properties` to flatten the various accepted forms of a per-DOF config field (``stiffness``, ``damping``, ``effort_limit``, …) into a single ``{joint_name: value}`` dict that the authoring loop can ``.get(jname, default)`` against. @@ -62,77 +62,78 @@ def resolve_per_dof( return {} -def author_actuator_prims_for_articulation(cfg: Any, sim_cfg: Any, stage: Any) -> None: - """Resolve per-articulation authoring args from ``cfg`` + ``sim_cfg`` and dispatch. +def define_actuator_properties( + prim_path: str, + actuator_cfgs: dict[str, Any], + stage: Any | None = None, +) -> None: + """Author ``NewtonActuator`` USD prims under an articulation root. + + For every joint covered by an explicit (non-implicit) Lab actuator + config, any existing ``NewtonActuator`` prim targeting that joint is + replaced by a new one created from the config values. Joints **not** + covered by any Lab config keep their USD-authored actuators unchanged. + + The supported config-to-schema mapping is: + + * :class:`~isaaclab.actuators.IdealPDActuatorCfg` → + ``NewtonPDControlAPI`` + ``NewtonMaxEffortClampingAPI`` + * :class:`~isaaclab.actuators.DCMotorCfg` → + ``NewtonPDControlAPI`` + ``NewtonDCMotorClampingAPI`` + * :class:`~isaaclab.actuators.DelayedPDActuatorCfg` → + same as ``IdealPDActuatorCfg`` + ``NewtonActuatorDelayAPI`` + * :class:`~isaaclab.actuators.RemotizedPDActuatorCfg` → + same as ``DelayedPDActuatorCfg`` + ``NewtonPositionBasedClampingAPI`` + * :class:`~isaaclab.actuators.ActuatorNetMLPCfg` / + :class:`~isaaclab.actuators.ActuatorNetLSTMCfg` → + ``NewtonNeuralControlAPI`` (+ ``NewtonDCMotorClampingAPI``) - Higher-level wrapper around :func:`author_newton_actuator_prims` that - handles the pre-conditions previously inlined in - ``BaseArticulation._author_newton_actuator_prims``: + No-ops (returns immediately) when: - * No-op when the simulation is not configured for the Newton actuator - fast path (``sim_cfg is None`` or ``use_newton_actuators=False``). - * Falls back from ``cfg.spawn.spawn_path`` to ``cfg.prim_path`` when the - former is not set, then resolves the first matching USD prim. - * No-op when no matching prim exists on the stage. + * the active :class:`~isaaclab.sim.SimulationContext` was configured + with ``use_newton_actuators=False`` (or no context is active), or + * *prim_path* does not resolve to a valid prim on the stage. - Keeping this resolution on the schema side lets articulations call - authoring with a single line, and lets the rules be tested without - constructing an articulation instance. + Must be called **after** the articulation is spawned (joint prims + exist on stage) and **before** the cloner / ``ModelBuilder.add_usd`` + reads the stage. + + Args: + prim_path: Root prim path of the articulation (e.g. + ``"/World/Env_0/Robot"``). May contain a regex pattern; the + first matching prim is used. + actuator_cfgs: Mapping of group name to + :class:`~isaaclab.actuators.ActuatorBaseCfg`. + stage: USD stage to author on. When ``None``, the current stage + is used. """ + from isaaclab.sim import SimulationContext # noqa: PLC0415 + + sim_ctx = SimulationContext.instance() + sim_cfg = sim_ctx.cfg if sim_ctx is not None else None if sim_cfg is None or not getattr(sim_cfg, "use_newton_actuators", False): return from isaaclab.sim.utils.queries import find_first_matching_prim # noqa: PLC0415 + from isaaclab.sim.utils.stage import get_current_stage # noqa: PLC0415 + + if stage is None: + stage = get_current_stage() - spawn_path = getattr(cfg.spawn, "spawn_path", None) if cfg.spawn is not None else None - search_path = spawn_path if spawn_path is not None else cfg.prim_path - first_prim = find_first_matching_prim(search_path) + first_prim = find_first_matching_prim(prim_path) if first_prim is None: return + articulation_prim_path = str(first_prim.GetPath()) - author_newton_actuator_prims( - stage=stage, - articulation_prim_path=str(first_prim.GetPath()), - actuator_cfgs=cfg.actuators, - ) + _author_actuator_prims(stage, articulation_prim_path, actuator_cfgs) -def author_newton_actuator_prims( +def _author_actuator_prims( stage: Any, articulation_prim_path: str, actuator_cfgs: dict[str, Any], ) -> None: - """Author ``NewtonActuator`` USD prims from IsaacLab actuator configs. - - For every joint covered by an explicit (non-implicit) Lab actuator config, - any existing ``NewtonActuator`` prim targeting that joint is replaced by a - new one created from the config values. Joints **not** covered by any - Lab config keep their USD-authored actuators unchanged. - - The supported config-to-schema mapping is: - - * :class:`~isaaclab.actuators.IdealPDActuatorCfg` -> - ``NewtonPDControlAPI`` + ``NewtonMaxEffortClampingAPI`` - * :class:`~isaaclab.actuators.DCMotorCfg` -> - ``NewtonPDControlAPI`` + ``NewtonDCMotorClampingAPI`` - * :class:`~isaaclab.actuators.DelayedPDActuatorCfg` -> - same as ``IdealPDActuatorCfg`` + ``NewtonActuatorDelayAPI`` - * :class:`~isaaclab.actuators.RemotizedPDActuatorCfg` -> - same as ``DelayedPDActuatorCfg`` + ``NewtonPositionBasedClampingAPI`` - * :class:`~isaaclab.actuators.ActuatorNetMLPCfg` / - :class:`~isaaclab.actuators.ActuatorNetLSTMCfg` -> - ``NewtonNeuralControlAPI`` (+ ``NewtonDCMotorClampingAPI``) - - Must be called **after** the articulation is spawned (joint prims exist - on stage) and **before** the cloner / ``ModelBuilder.add_usd`` reads - the stage. - - Args: - stage: The USD stage to author prims on. - articulation_prim_path: Root prim path of the articulation - (e.g. ``"/World/Env_0/Robot"``). Must not contain wildcards. - actuator_cfgs: Mapping of group name to ``ActuatorBaseCfg``. - """ + """Inner authoring routine; exposed separately for test fixtures.""" from pxr import Sdf # noqa: PLC0415 from isaaclab.actuators import ImplicitActuator # noqa: PLC0415 @@ -151,9 +152,7 @@ def author_newton_actuator_prims( for group_name, cfg in actuator_cfgs.items(): cls_type = cfg.class_type is_implicit = ( - "ImplicitActuator" in cls_type - if isinstance(cls_type, str) - else issubclass(cls_type, ImplicitActuator) + "ImplicitActuator" in cls_type if isinstance(cls_type, str) else issubclass(cls_type, ImplicitActuator) ) if is_implicit: continue @@ -200,9 +199,7 @@ def author_newton_actuator_prims( is_delayed = isinstance(cfg, DelayedPDActuatorCfg) vel_limit_map = resolve_per_dof(getattr(cfg, "velocity_limit", None), joint_names) if is_dc_motor else {} - sat_effort_map = ( - resolve_per_dof(getattr(cfg, "saturation_effort", None), joint_names) if is_dc_motor else {} - ) + sat_effort_map = resolve_per_dof(getattr(cfg, "saturation_effort", None), joint_names) if is_dc_motor else {} raw_delay = getattr(cfg, "max_delay", 0) if is_delayed else 0 delay_map = resolve_per_dof(raw_delay, joint_names, cast=int) if raw_delay else {} @@ -321,7 +318,7 @@ def _remove_actuator_prims_for_joints( """Deactivate ``NewtonActuator`` prims whose target is in *joint_paths*. Deactivated prims are invisible to ``Usd.PrimRange`` and therefore - ignored by ``ModelBuilder.add_usd``. Using ``SetActive(False)`` + ignored by ``ModelBuilder.add_usd``. Using ``SetActive(False)`` instead of ``RemovePrim`` works correctly when the prim originates from a USD reference or payload. @@ -381,7 +378,7 @@ def _resave_checkpoint_with_metadata( merged = {**existing_meta, **metadata} - tmp = tempfile.NamedTemporaryFile(suffix=".pt", delete=False) + tmp = tempfile.NamedTemporaryFile(suffix=".pt", delete=False) # noqa: SIM115 if is_torchscript: extra_out = {"metadata.json": json.dumps(merged)} torch.jit.save(net, tmp.name, _extra_files=extra_out) diff --git a/source/isaaclab/isaaclab/sim/spawners/from_files/from_files.py b/source/isaaclab/isaaclab/sim/spawners/from_files/from_files.py index be7bd6e7074e..ecb6d5fa2052 100644 --- a/source/isaaclab/isaaclab/sim/spawners/from_files/from_files.py +++ b/source/isaaclab/isaaclab/sim/spawners/from_files/from_files.py @@ -367,6 +367,10 @@ def _spawn_from_usd_file( if cfg.joint_drive_props is not None: schemas.modify_joint_drive_properties(prim_path, cfg.joint_drive_props) + # author Newton-native actuator USD prims (no-op when use_newton_actuators=False) + if cfg.actuator_props is not None: + schemas.define_actuator_properties(prim_path, cfg.actuator_props, stage=stage) + # define deformable body properties, or modify if deformable body API is present (PhysX only) if cfg.deformable_props is not None: prim = stage.GetPrimAtPath(prim_path) diff --git a/source/isaaclab/isaaclab/sim/spawners/from_files/from_files_cfg.py b/source/isaaclab/isaaclab/sim/spawners/from_files/from_files_cfg.py index 87fc48de4d30..1d321b9e141e 100644 --- a/source/isaaclab/isaaclab/sim/spawners/from_files/from_files_cfg.py +++ b/source/isaaclab/isaaclab/sim/spawners/from_files/from_files_cfg.py @@ -7,6 +7,7 @@ from collections.abc import Callable from dataclasses import MISSING +from typing import TYPE_CHECKING # deformables only supported on PhysX backend from isaaclab_physx.sim.spawners.spawner_cfg import DeformableObjectSpawnerCfg @@ -17,6 +18,9 @@ from isaaclab.utils import configclass from isaaclab.utils.assets import ISAAC_NUCLEUS_DIR +if TYPE_CHECKING: + from isaaclab.actuators import ActuatorBaseCfg + @configclass class FileCfg(RigidObjectSpawnerCfg, DeformableObjectSpawnerCfg): @@ -56,6 +60,16 @@ class FileCfg(RigidObjectSpawnerCfg, DeformableObjectSpawnerCfg): for specific joints in an articulation. """ + actuator_props: dict[str, ActuatorBaseCfg] | None = None + """Newton-native actuator configs to author as ``NewtonActuator`` USD prims. + + Auto-populated from :attr:`~isaaclab.assets.ArticulationCfg.actuators` when an + articulation cfg owns this spawn cfg. The spawner authors NewtonActuator prims + via :func:`~isaaclab.sim.schemas.define_actuator_properties` after the asset is + spawned. The function is a no-op unless + :attr:`~isaaclab.sim.SimulationCfg.use_newton_actuators` is enabled. + """ + visual_material_path: str = "material" """Path to the visual material to use for the prim. Defaults to "material". diff --git a/source/isaaclab_newton/isaaclab_newton/actuators/adapter.py b/source/isaaclab_newton/isaaclab_newton/actuators/adapter.py index 4a76e440dc82..d777a4c4e32d 100644 --- a/source/isaaclab_newton/isaaclab_newton/actuators/adapter.py +++ b/source/isaaclab_newton/isaaclab_newton/actuators/adapter.py @@ -22,7 +22,6 @@ import numpy as np import torch import warp as wp - from newton.actuators import Actuator, Clamping, Delay from isaaclab.actuators import ActuatorBase, ImplicitActuator @@ -92,14 +91,11 @@ class NewtonActuatorAdapter: * **Newton backend** — pass actuators from the Newton model directly:: - adapter = NewtonActuatorAdapter(model.actuators, num_envs, - num_joints, dof_offset, device) + adapter = NewtonActuatorAdapter(model.actuators, num_envs, num_joints, dof_offset, device) * **PhysX backend** — create actuators from USD prims:: - adapter = NewtonActuatorAdapter.from_usd(stage, joint_names, - num_envs, num_joints, - device) + adapter = NewtonActuatorAdapter.from_usd(stage, joint_names, num_envs, num_joints, device) Then finalise:: @@ -155,15 +151,15 @@ def from_usd( num_joints: int, device: str, articulation_prim_path: str | None = None, - ) -> "NewtonActuatorAdapter": + ) -> NewtonActuatorAdapter: """Create an adapter by parsing ``NewtonActuator`` prims from USD. This is the PhysX-backend counterpart of what Newton's ``ModelBuilder.add_usd`` does for the Newton backend. Both paths read the same ``NewtonActuator`` USD prims (authored by - :func:`~isaaclab_newton.actuators.authoring.author_newton_actuator_prims`) - and construct :class:`~newton.actuators.Actuator` objects with - matching controllers, clampings, and delays. + :func:`~isaaclab.sim.schemas.define_actuator_properties`) and + construct :class:`~newton.actuators.Actuator` objects with matching + controllers, clampings, and delays. The key difference is that PhysX uses a **flat per-DOF layout** where joint position coordinates and velocity DOFs always have the @@ -197,7 +193,11 @@ def from_usd( A fully constructed adapter ready for :meth:`finalize`. """ actuators = _create_actuators_from_usd( - stage, joint_names, num_envs, num_joints, device, + stage, + joint_names, + num_envs, + num_joints, + device, articulation_prim_path=articulation_prim_path, ) return cls(actuators, num_envs, num_joints, dof_offset=0, device=device) @@ -214,13 +214,15 @@ def finalize(self) -> None: ctrl = act.controller if hasattr(ctrl, "kp"): wp.launch( - scatter_gain_kernel, dim=act.indices.shape[0], + scatter_gain_kernel, + dim=act.indices.shape[0], inputs=[ctrl.kp, flat_stiffness, act.indices, self._dof_offset, self.num_joints], device=wp_device, ) if hasattr(ctrl, "kd"): wp.launch( - scatter_gain_kernel, dim=act.indices.shape[0], + scatter_gain_kernel, + dim=act.indices.shape[0], inputs=[ctrl.kd, flat_damping, act.indices, self._dof_offset, self.num_joints], device=wp_device, ) @@ -343,14 +345,22 @@ def set_view_propagator(self, root_view: Any) -> None: articulation's view works since the call dispatches on the actuator object (model-scoped), not on the view's articulation. """ + def _push( - actuator: Actuator, controller: Any, attr: str, - values: wp.array, env_mask: wp.array, + actuator: Actuator, + controller: Any, + attr: str, + values: wp.array, + env_mask: wp.array, ) -> None: root_view.set_actuator_parameter( - actuator=actuator, component=controller, name=attr, - values=values, mask=env_mask, + actuator=actuator, + component=controller, + name=attr, + values=values, + mask=env_mask, ) + self._propagator = _push def set_kernel_propagator(self) -> None: @@ -361,19 +371,28 @@ def set_kernel_propagator(self) -> None: actuator's flat indices and writes into ``controller.kp`` / ``controller.kd``. """ + def _push( - actuator: Actuator, controller: Any, attr: str, - values: wp.array, env_mask: wp.array, + actuator: Actuator, + controller: Any, + attr: str, + values: wp.array, + env_mask: wp.array, ) -> None: wp.launch( gather_gain_kernel, dim=actuator.indices.shape[0], inputs=[ - values.flatten(), getattr(controller, attr), actuator.indices, - env_mask, self._dof_offset, self.num_joints, + values.flatten(), + getattr(controller, attr), + actuator.indices, + env_mask, + self._dof_offset, + self.num_joints, ], device=self._device, ) + self._propagator = _push def write_stiffness_to_sim( @@ -461,6 +480,7 @@ class (e.g. ``model_path``, ``lookup_positions``) are passed through from collections import defaultdict # noqa: PLC0415 from newton.actuators import parse_actuator_prim # noqa: PLC0415 + from pxr import Usd # noqa: PLC0415 wp_device = wp.get_device(device) @@ -482,9 +502,7 @@ class (e.g. ``model_path``, ``lookup_positions``) are passed through parsed_per_joint[joint_name_to_idx[target_name]] = parsed if not parsed_per_joint: - raise ValueError( - f"No NewtonActuator prims found targeting any of: {joint_names}" - ) + raise ValueError(f"No NewtonActuator prims found targeting any of: {joint_names}") groups: dict[tuple, list[int]] = defaultdict(list) sig_to_parsed: dict[tuple, Any] = {} @@ -536,7 +554,10 @@ class (e.g. ``model_path``, ``lookup_positions``) are passed through clamp_arrays[k] = v else: clamp_arrays[k] = wp.full( - num_dofs_in_group, float(v), dtype=wp.float32, device=wp_device, + num_dofs_in_group, + float(v), + dtype=wp.float32, + device=wp_device, ) clampings.append(comp_cls(**clamp_arrays)) From f9d11209ad9a6caf6b8d7d576a8770c491122ebe Mon Sep 17 00:00:00 2001 From: Julia von Muralt Date: Fri, 8 May 2026 18:06:42 +0200 Subject: [PATCH 20/35] move authoring logic --- .../assets/articulation/articulation_cfg.py | 22 ++++++++++++++++++- .../assets/articulation/base_articulation.py | 13 ----------- source/isaaclab/isaaclab/assets/asset_base.py | 2 ++ .../isaaclab/assets/asset_base_cfg.py | 16 +++++++++++++- .../isaaclab/physics/physics_manager.py | 1 + .../assets/articulation/articulation.py | 5 +++++ .../assets/articulation/articulation.py | 5 +++++ 7 files changed, 49 insertions(+), 15 deletions(-) diff --git a/source/isaaclab/isaaclab/assets/articulation/articulation_cfg.py b/source/isaaclab/isaaclab/assets/articulation/articulation_cfg.py index d3d94a9b23e0..3bcd54b3c906 100644 --- a/source/isaaclab/isaaclab/assets/articulation/articulation_cfg.py +++ b/source/isaaclab/isaaclab/assets/articulation/articulation_cfg.py @@ -6,7 +6,7 @@ from __future__ import annotations from dataclasses import MISSING -from typing import TYPE_CHECKING +from typing import TYPE_CHECKING, Any from isaaclab.actuators import ActuatorBaseCfg from isaaclab.utils import configclass @@ -74,3 +74,23 @@ class InitialStateCfg(AssetBaseCfg.InitialStateCfg): actuator_value_resolution_debug_print = False """Print the resolution of actuator final value when input cfg is different from USD value, Defaults to False """ + + def _post_spawn(self, stage: Any) -> None: + """Author ``NewtonActuator`` USD prims from :attr:`actuators` after spawn. + + Invoked by :class:`~isaaclab.assets.AssetBase` once the articulation's prims + exist on the stage. The actual authoring logic lives in + :func:`~isaaclab_newton.actuators.author_actuator_prims_for_articulation`, + which gates itself on ``sim_cfg.use_newton_actuators`` and silently no-ops when + the Newton actuator package is unavailable. + """ + from isaaclab.sim import SimulationContext # noqa: PLC0415 + + try: + from isaaclab_newton.actuators import author_actuator_prims_for_articulation # noqa: PLC0415 + except ImportError: + return + + sim_ctx = SimulationContext.instance() + sim_cfg = sim_ctx.cfg if sim_ctx is not None else None + author_actuator_prims_for_articulation(self, sim_cfg, stage) diff --git a/source/isaaclab/isaaclab/assets/articulation/base_articulation.py b/source/isaaclab/isaaclab/assets/articulation/base_articulation.py index f275bdc3e984..14aa592ad103 100644 --- a/source/isaaclab/isaaclab/assets/articulation/base_articulation.py +++ b/source/isaaclab/isaaclab/assets/articulation/base_articulation.py @@ -93,21 +93,8 @@ def __init__(self, cfg: ArticulationCfg): Args: cfg: A configuration instance. """ - from isaaclab.sim import SimulationContext # noqa: PLC0415 - super().__init__(cfg) - sim_ctx = SimulationContext.instance() - self._sim_cfg = sim_ctx.cfg if sim_ctx is not None else None - - - try: - from isaaclab_newton.actuators import author_actuator_prims_for_articulation # noqa: PLC0415 - except ImportError: - pass - else: - author_actuator_prims_for_articulation(self.cfg, self._sim_cfg, self.stage) - """ Properties """ diff --git a/source/isaaclab/isaaclab/assets/asset_base.py b/source/isaaclab/isaaclab/assets/asset_base.py index f5f121c5ad2b..cccc97ff7162 100644 --- a/source/isaaclab/isaaclab/assets/asset_base.py +++ b/source/isaaclab/isaaclab/assets/asset_base.py @@ -86,6 +86,8 @@ def __init__(self, cfg: AssetBaseCfg): matching_prims = sim_utils.find_matching_prims(check_path) if len(matching_prims) == 0: raise RuntimeError(f"Could not find prim with path {check_path}.") + # schema-side post-spawn hook (e.g. ArticulationCfg authors NewtonActuator prims here) + self.cfg._post_spawn(self.stage) else: # asset should exist at run time check_path = self.cfg.prim_path diff --git a/source/isaaclab/isaaclab/assets/asset_base_cfg.py b/source/isaaclab/isaaclab/assets/asset_base_cfg.py index 37d551ddec7d..f82dbd74b674 100644 --- a/source/isaaclab/isaaclab/assets/asset_base_cfg.py +++ b/source/isaaclab/isaaclab/assets/asset_base_cfg.py @@ -6,7 +6,7 @@ from __future__ import annotations from dataclasses import MISSING -from typing import Literal +from typing import Any, Literal from isaaclab.sim import SpawnerCfg from isaaclab.utils import configclass @@ -75,3 +75,17 @@ class InitialStateCfg: debug_vis: bool = False """Whether to enable debug visualization for the asset. Defaults to ``False``.""" + + def _post_spawn(self, stage: Any) -> None: + """Hook invoked by :class:`~isaaclab.assets.AssetBase` after the asset's prims are + spawned and verified to exist on the stage. + + The default implementation is a no-op. Subclasses that need to author additional + USD schemas tied to this asset (for example, :class:`~isaaclab.assets.ArticulationCfg` + which authors ``NewtonActuator`` prims from its ``actuators`` mapping) should + override this method. + + Args: + stage: The USD stage on which the asset was spawned. + """ + pass diff --git a/source/isaaclab/isaaclab/physics/physics_manager.py b/source/isaaclab/isaaclab/physics/physics_manager.py index ca333673bb48..e6157177ded4 100644 --- a/source/isaaclab/isaaclab/physics/physics_manager.py +++ b/source/isaaclab/isaaclab/physics/physics_manager.py @@ -351,6 +351,7 @@ def set_decimation(cls, decimation: int) -> None: Args: decimation: Number of physics steps per environment step. """ + pass @classmethod def handles_decimation(cls) -> bool: diff --git a/source/isaaclab_newton/isaaclab_newton/assets/articulation/articulation.py b/source/isaaclab_newton/isaaclab_newton/assets/articulation/articulation.py index b4ff8c998d19..727f1a4d2395 100644 --- a/source/isaaclab_newton/isaaclab_newton/assets/articulation/articulation.py +++ b/source/isaaclab_newton/isaaclab_newton/assets/articulation/articulation.py @@ -122,8 +122,13 @@ def __init__(self, cfg: ArticulationCfg): Args: cfg: A configuration instance. """ + from isaaclab.sim import SimulationContext # noqa: PLC0415 + super().__init__(cfg) + sim_ctx = SimulationContext.instance() + self._sim_cfg = sim_ctx.cfg if sim_ctx is not None else None + """ Properties """ diff --git a/source/isaaclab_physx/isaaclab_physx/assets/articulation/articulation.py b/source/isaaclab_physx/isaaclab_physx/assets/articulation/articulation.py index 17d17267d498..a83a83aaa11f 100644 --- a/source/isaaclab_physx/isaaclab_physx/assets/articulation/articulation.py +++ b/source/isaaclab_physx/isaaclab_physx/assets/articulation/articulation.py @@ -129,8 +129,13 @@ def __init__(self, cfg: ArticulationCfg): Args: cfg: A configuration instance. """ + from isaaclab.sim import SimulationContext # noqa: PLC0415 + super().__init__(cfg) + sim_ctx = SimulationContext.instance() + self._sim_cfg = sim_ctx.cfg if sim_ctx is not None else None + """ Properties """ From 16138b0e265800659ccaee89dfbd3ee99c76cad5 Mon Sep 17 00:00:00 2001 From: Julia von Muralt Date: Sun, 10 May 2026 17:38:41 +0200 Subject: [PATCH 21/35] refactor --- .../assets/articulation/articulation_cfg.py | 6 +- source/isaaclab/isaaclab/assets/asset_base.py | 4 +- .../isaaclab/assets/asset_base_cfg.py | 6 +- source/isaaclab/isaaclab/envs/mdp/events.py | 69 +- .../isaaclab_newton/actuators/adapter.py | 532 ++++------ .../isaaclab_newton/actuators/authoring.py | 9 +- .../isaaclab_newton/actuators/kernels.py | 162 +-- .../assets/articulation/articulation.py | 157 ++- .../isaaclab_newton/physics/newton_manager.py | 46 +- .../views/mock_articulation_view.py | 11 +- .../assets/test_newton_actuators_newton.py | 954 ++++++++---------- .../assets/articulation/articulation.py | 172 +++- .../assets/test_newton_actuators_physx.py | 794 ++++++--------- 13 files changed, 1363 insertions(+), 1559 deletions(-) diff --git a/source/isaaclab/isaaclab/assets/articulation/articulation_cfg.py b/source/isaaclab/isaaclab/assets/articulation/articulation_cfg.py index 3bcd54b3c906..0ca196f58bb3 100644 --- a/source/isaaclab/isaaclab/assets/articulation/articulation_cfg.py +++ b/source/isaaclab/isaaclab/assets/articulation/articulation_cfg.py @@ -6,7 +6,9 @@ from __future__ import annotations from dataclasses import MISSING -from typing import TYPE_CHECKING, Any +from typing import TYPE_CHECKING + +from pxr import Usd from isaaclab.actuators import ActuatorBaseCfg from isaaclab.utils import configclass @@ -75,7 +77,7 @@ class InitialStateCfg(AssetBaseCfg.InitialStateCfg): """Print the resolution of actuator final value when input cfg is different from USD value, Defaults to False """ - def _post_spawn(self, stage: Any) -> None: + def _post_spawn(self, stage: Usd.Stage) -> None: """Author ``NewtonActuator`` USD prims from :attr:`actuators` after spawn. Invoked by :class:`~isaaclab.assets.AssetBase` once the articulation's prims diff --git a/source/isaaclab/isaaclab/assets/asset_base.py b/source/isaaclab/isaaclab/assets/asset_base.py index cccc97ff7162..3c95b78617fa 100644 --- a/source/isaaclab/isaaclab/assets/asset_base.py +++ b/source/isaaclab/isaaclab/assets/asset_base.py @@ -15,6 +15,8 @@ import torch import warp as wp +from pxr import Usd + import isaaclab.sim as sim_utils from isaaclab.physics import PhysicsEvent, PhysicsManager from isaaclab.sim.simulation_context import SimulationContext @@ -69,7 +71,7 @@ def __init__(self, cfg: AssetBaseCfg): # flag for whether the asset is initialized self._is_initialized = False # get stage handle - self.stage = get_current_stage() + self.stage: Usd.Stage = get_current_stage() # spawn the asset # determine path where prims should exist after spawn diff --git a/source/isaaclab/isaaclab/assets/asset_base_cfg.py b/source/isaaclab/isaaclab/assets/asset_base_cfg.py index f82dbd74b674..d936982c52d1 100644 --- a/source/isaaclab/isaaclab/assets/asset_base_cfg.py +++ b/source/isaaclab/isaaclab/assets/asset_base_cfg.py @@ -6,7 +6,9 @@ from __future__ import annotations from dataclasses import MISSING -from typing import Any, Literal +from typing import Literal + +from pxr import Usd from isaaclab.sim import SpawnerCfg from isaaclab.utils import configclass @@ -76,7 +78,7 @@ class InitialStateCfg: debug_vis: bool = False """Whether to enable debug visualization for the asset. Defaults to ``False``.""" - def _post_spawn(self, stage: Any) -> None: + def _post_spawn(self, stage: Usd.Stage) -> None: """Hook invoked by :class:`~isaaclab.assets.AssetBase` after the asset's prims are spawned and verified to exist on the stage. diff --git a/source/isaaclab/isaaclab/envs/mdp/events.py b/source/isaaclab/isaaclab/envs/mdp/events.py index 52b24d17bcaf..36612af65f49 100644 --- a/source/isaaclab/isaaclab/envs/mdp/events.py +++ b/source/isaaclab/isaaclab/envs/mdp/events.py @@ -1135,13 +1135,22 @@ def __init__(self, cfg: EventTermCfg, env: ManagerBasedEnv): self.default_joint_stiffness = self.asset.data.joint_stiffness.torch.clone() self.default_joint_damping = self.asset.data.joint_damping.torch.clone() - # For explicit actuators the sim-level stiffness/damping is zeroed out, so patch - # the defaults with the actual actuator PD gains. + # For explicit Lab actuators the sim-level stiffness/damping is zeroed out, + # so patch the defaults with the actual actuator PD gains. for actuator in self.asset.actuators.values(): if not isinstance(actuator, ImplicitActuator): joint_ids = actuator.joint_indices self.default_joint_stiffness[:, joint_ids] = actuator.stiffness self.default_joint_damping[:, joint_ids] = actuator.damping + # Same for explicit Newton actuators on either backend — their kp/kd + # live on the per-actuator controller arrays (not on a Lab Actuator + # object), so the asset exposes a per-articulation snapshot taken + # at articulation init time. + newton_default_stiffness = getattr(self.asset, "newton_default_stiffness", None) + if newton_default_stiffness is not None: + joint_ids = self.asset.newton_managed_local_joints + self.default_joint_stiffness[:, joint_ids] = newton_default_stiffness[:, joint_ids] + self.default_joint_damping[:, joint_ids] = self.asset.newton_default_damping[:, joint_ids] # check for valid operation if cfg.params["operation"] == "scale": @@ -1221,45 +1230,39 @@ def randomize(data: torch.Tensor, params: tuple[float, float]) -> torch.Tensor: damping=damping, joint_ids=actuator.joint_indices, env_ids=env_ids ) - # Direct write to the global Newton actuator adapter - adapter = getattr(env.sim.physics_manager, "_adapter", None) - if adapter is None: + # Push DR updates to explicit Newton-actuator controllers via the asset's + # own write methods. Each backend's articulation iterates the adapter's + # actuators and propagates per actuator, using the appropriate backend + # mechanism (Newton ``ArticulationView`` on the Newton backend, an + # in-place scatter kernel on PhysX). + if not hasattr(self.asset, "write_actuator_stiffness_to_sim"): return - from newton import Model as NewtonModel # noqa: PLC0415 - dof_layout = self.asset._root_view.frequency_layouts[NewtonModel.AttributeFrequency.JOINT_DOF] - if dof_layout.slice is not None: - dof_offset = dof_layout.slice.start - elif dof_layout.indices is not None: - dof_offset = int(dof_layout.indices.numpy()[0]) - else: - dof_offset = 0 if isinstance(self.asset_cfg.joint_ids, slice): - arti_local_ids = torch.arange(self.asset.num_joints, device=self.asset.device, dtype=torch.long) + joint_ids = torch.arange(self.asset.num_joints, device=self.asset.device, dtype=torch.long) else: - arti_local_ids = torch.tensor(self.asset_cfg.joint_ids, device=self.asset.device, dtype=torch.long) - # Column indices into the adapter's model-wide gain buffer for the - # joints this DR term should randomize. - adapter_indices = arti_local_ids + dof_offset - env_ids_wp = self.asset._resolve_env_ids(env_ids) - env_mask = self.asset._env_ids_to_mask(env_ids_wp) - - def _randomize_at(data: torch.Tensor, params: tuple[float, float]) -> None: + joint_ids = torch.tensor(self.asset_cfg.joint_ids, device=self.asset.device, dtype=torch.long) + + if stiffness_distribution_params is not None: + new_stiffness = self.default_joint_stiffness[env_ids][:, joint_ids].clone() _randomize_prop_by_op( - data, params, dim_0_ids=None, dim_1_ids=adapter_indices, + new_stiffness, stiffness_distribution_params, + dim_0_ids=None, dim_1_ids=slice(None), operation=operation, distribution=distribution, ) - - if stiffness_distribution_params is not None: - stiffness = adapter.stiffness[env_ids].clone() - stiffness[:, adapter_indices] = self.default_joint_stiffness[env_ids][:, arti_local_ids].clone() - _randomize_at(stiffness, stiffness_distribution_params) - adapter.write_stiffness_to_sim(stiffness, env_ids_wp, env_mask) + self.asset.write_actuator_stiffness_to_sim( + stiffness=new_stiffness, env_ids=env_ids, joint_ids=joint_ids, + ) if damping_distribution_params is not None: - damping = adapter.damping[env_ids].clone() - damping[:, adapter_indices] = self.default_joint_damping[env_ids][:, arti_local_ids].clone() - _randomize_at(damping, damping_distribution_params) - adapter.write_damping_to_sim(damping, env_ids_wp, env_mask) + new_damping = self.default_joint_damping[env_ids][:, joint_ids].clone() + _randomize_prop_by_op( + new_damping, damping_distribution_params, + dim_0_ids=None, dim_1_ids=slice(None), + operation=operation, distribution=distribution, + ) + self.asset.write_actuator_damping_to_sim( + damping=new_damping, env_ids=env_ids, joint_ids=joint_ids, + ) class randomize_joint_parameters(ManagerTermBase): diff --git a/source/isaaclab_newton/isaaclab_newton/actuators/adapter.py b/source/isaaclab_newton/isaaclab_newton/actuators/adapter.py index 4a76e440dc82..a328b0e0b617 100644 --- a/source/isaaclab_newton/isaaclab_newton/actuators/adapter.py +++ b/source/isaaclab_newton/isaaclab_newton/actuators/adapter.py @@ -5,18 +5,19 @@ """Newton-actuator adapter shared by the Newton and PhysX backends. -:class:`NewtonActuatorAdapter` manages actuator creation, DOF-to-actuator -mapping, stepping, reset, and gain reading for domain randomisation. -Used identically by both Newton and PhysX backends. - -The companion helper :func:`build_implicit_dof_mask` is consumed by the -in-graph post-actuator kernel -(:func:`~isaaclab_newton.actuators.kernels.synch_torque_and_apply_implicit_feedforwards`). +Owns the actuator-state lifecycle, the pre-clamp computed-effort buffer, +and the per-step ``step`` / ``reset`` / ``finalize`` calls. The +:meth:`~NewtonActuatorAdapter.from_usd` classmethod parses +``NewtonActuator`` USD prims on the PhysX backend (Newton populates +``model.actuators`` itself). + +DR gain updates bypass the adapter — the articulation writes straight +to controller arrays. """ from __future__ import annotations -from collections.abc import Callable, Sequence +from collections.abc import Sequence from typing import Any import numpy as np @@ -25,88 +26,25 @@ from newton.actuators import Actuator, Clamping, Delay -from isaaclab.actuators import ActuatorBase, ImplicitActuator - from .kernels import ( - fill_gain_at_envs_kernel, - gather_gain_kernel, - scatter_gain_at_envs_kernel, + build_per_dof_env_mask_kernel, scatter_gain_kernel, set_mask_kernel, zero_at_indices_kernel, ) -def _actuator_signature(parsed: Any) -> tuple: - """Build a hashable key from a parsed actuator spec for grouping. - - Joints whose prims resolve to the same signature share identical - controller type, gains, clamping chain, and delay configuration and - can therefore be merged into a single :class:`~newton.actuators.Actuator` - with combined index arrays. - """ - ctrl_resolved = parsed.controller_class.resolve_arguments( - dict(parsed.controller_kwargs), - ) - ctrl_key = (parsed.controller_class, tuple(sorted(ctrl_resolved.items()))) - - comp_keys: list[tuple] = [] - for comp_cls, comp_kwargs in parsed.component_specs: - resolved = comp_cls.resolve_arguments(comp_kwargs) - comp_keys.append((comp_cls, tuple(sorted(resolved.items())))) - comp_keys.sort(key=lambda t: t[0].__name__) - - return (ctrl_key, tuple(comp_keys)) - - -def build_implicit_dof_mask( - actuators: dict[str, ActuatorBase], - num_joints: int, - device: str, -) -> wp.array: - """Per-DOF mask consumed by the in-graph implicit-FF kernel. - - Entry is ``1`` for DOFs covered by an - :class:`~isaaclab.actuators.ImplicitActuator` group, ``0`` otherwise. - """ - modes = torch.zeros(num_joints, dtype=torch.int32, device=device) - for actuator in actuators.values(): - if not isinstance(actuator, ImplicitActuator): - continue - j_ids = actuator.joint_indices - if j_ids == slice(None) or j_ids is None: - modes[:] = 1 - else: - modes[j_ids.long()] = 1 - return wp.from_torch(modes, dtype=wp.int32) +# --------------------------------------------------------------------------- +# Abstract base — backend-independent logic +# --------------------------------------------------------------------------- class NewtonActuatorAdapter: - """Manages Newton-native actuators for both Newton and PhysX backends. - - Handles actuator creation (from USD or an existing list), - DOF-to-actuator mapping, stepping, reset, and gain reading for - domain randomisation. - - Construction: - - * **Newton backend** — pass actuators from the Newton model directly:: - - adapter = NewtonActuatorAdapter(model.actuators, num_envs, - num_joints, dof_offset, device) - - * **PhysX backend** — create actuators from USD prims:: - - adapter = NewtonActuatorAdapter.from_usd(stage, joint_names, - num_envs, num_joints, - device) + """Adapter that wraps a list of :class:`newton.actuators.Actuator`. - Then finalise:: - - adapter.finalize() - - After :meth:`finalize`, the adapter exposes ``.stiffness``, - ``.damping``, and ``.joint_indices`` for ``randomize_actuator_gains``. + Owns the actuator-state lifecycle, DOF-to-actuator bookkeeping, + stepping, reset, and the pre-clamp computed-effort buffer the + in-graph telemetry kernel reads on the post-actuator hook. """ def __init__( @@ -123,110 +61,51 @@ def __init__( self._num_envs = num_envs self._dof_offset = dof_offset self._device = device - self._dof_to_actuator = self._build_dof_map() - managed = [i for i, act_idx in enumerate(self._dof_to_actuator) if act_idx >= 0] + # Collect the set of local DOFs covered by some actuator. Only the + # env-0 slice of each actuator's flat ``indices`` array is needed — + # later envs are repeats with a constant ``num_joints`` stride. + managed: set[int] = set() + for act in actuators: + all_indices = act.indices.numpy() + num_per_act = len(all_indices) // num_envs + for global_dof in all_indices[:num_per_act]: + local_dof = global_dof - dof_offset + if 0 <= local_dof < num_joints: + managed.add(local_dof) + if len(managed) == num_joints: self.joint_indices: torch.Tensor | slice = slice(None) else: - self.joint_indices = torch.tensor(managed, dtype=torch.int32, device=device) + self.joint_indices = torch.tensor(sorted(managed), dtype=torch.int32, device=device) self._states_a = [act.state() for act in actuators] self._states_b = [act.state() for act in actuators] - self.stiffness: torch.Tensor | None = None - self.damping: torch.Tensor | None = None - - # Per-actuator gain propagator. Configured once at adapter setup time - # via :meth:`set_view_propagator` (Newton: simulator-side scatter API) - # or :meth:`set_kernel_propagator` (PhysX: local Warp scatter kernel). - # Stays ``None`` if gain DR is not used; ``write_*_to_sim`` then - # only updates the adapter's gain buffer without pushing to controllers. - self._propagator: Callable[[Actuator, Any, str, wp.array, wp.array], None] | None = None - - # -- construction (PhysX path) ------------------------------------------- + # Pre-clamp computed effort buffer. Each Newton actuator scatter-adds + # its raw controller output to ``sim_control.joint_computed_f`` when + # ``control_computed_output_attr`` is set; we route that to this + # buffer so the post-actuator telemetry kernel can report the actual + # computed (pre-clamp) effort instead of mirroring ``joint_f``. The + # binding onto ``sim_control`` happens in :meth:`finalize`. + self._computed_effort = wp.zeros( + num_envs * num_joints, dtype=wp.float32, device=device, + ) + self.computed_effort_2d = self._computed_effort.reshape((num_envs, num_joints)) + for act in actuators: + act.control_computed_output_attr = "joint_computed_f" - @classmethod - def from_usd( - cls, - stage: Any, - joint_names: list[str], - num_envs: int, - num_joints: int, - device: str, - articulation_prim_path: str | None = None, - ) -> "NewtonActuatorAdapter": - """Create an adapter by parsing ``NewtonActuator`` prims from USD. - - This is the PhysX-backend counterpart of what Newton's - ``ModelBuilder.add_usd`` does for the Newton backend. Both paths - read the same ``NewtonActuator`` USD prims (authored by - :func:`~isaaclab_newton.actuators.authoring.author_newton_actuator_prims`) - and construct :class:`~newton.actuators.Actuator` objects with - matching controllers, clampings, and delays. - - The key difference is that PhysX uses a **flat per-DOF layout** - where joint position coordinates and velocity DOFs always have the - same count and ordering — there are no free joints or ball joints - that cause coordinate/DOF count divergence. Therefore a single - ``indices`` array is used for all index roles (``indices``, - ``pos_indices``, ``target_pos_indices``), unlike the Newton - builder which computes separate ``pos_indices`` from - ``joint_q_start`` and separate ``target_pos_indices`` from - ``joint_qd_start`` to handle floating-base articulations. - - Joints whose prims resolve to the same controller type, gains, - clamping chain, and delay configuration are merged into a single - :class:`Actuator` with combined index arrays, mirroring the - grouping the Newton builder performs internally. + def finalize(self, sim_control: Any) -> None: + """Bind the pre-clamp computed-effort buffer onto ``sim_control``. Args: - stage: The USD stage containing ``NewtonActuator`` prims. - joint_names: All joint names in the articulation. - num_envs: Number of environments. - num_joints: Joints per environment in the articulation. - device: Warp device string (e.g. ``"cuda:0"``). - articulation_prim_path: Root prim path of the first - environment's articulation (e.g. ``"/World/Env_0/Robot"``). - When provided, only ``NewtonActuator`` prims under this - subtree are considered — matching the scoped traversal - that Newton's ``ModelBuilder.add_usd`` performs. When - ``None``, the entire stage is scanned (legacy behaviour). - - Returns: - A fully constructed adapter ready for :meth:`finalize`. + sim_control: The ``sim_control`` object that will be passed + to :meth:`step` for this adapter's lifetime. Newton's + ``Control`` on the Newton backend, an + :class:`~isaaclab_newton.actuators.physx_wrapper.PhysxActuatorWrapper` + on the PhysX backend. """ - actuators = _create_actuators_from_usd( - stage, joint_names, num_envs, num_joints, device, - articulation_prim_path=articulation_prim_path, - ) - return cls(actuators, num_envs, num_joints, dof_offset=0, device=device) - - # -- public API ---------------------------------------------------------- - - def finalize(self) -> None: - """Read actuator gains and store as PyTorch tensors for DR.""" - wp_device = wp.get_device(self._device) - flat_stiffness = wp.zeros(self._num_envs * self.num_joints, dtype=wp.float32, device=wp_device) - flat_damping = wp.zeros(self._num_envs * self.num_joints, dtype=wp.float32, device=wp_device) - - for act in self.actuators: - ctrl = act.controller - if hasattr(ctrl, "kp"): - wp.launch( - scatter_gain_kernel, dim=act.indices.shape[0], - inputs=[ctrl.kp, flat_stiffness, act.indices, self._dof_offset, self.num_joints], - device=wp_device, - ) - if hasattr(ctrl, "kd"): - wp.launch( - scatter_gain_kernel, dim=act.indices.shape[0], - inputs=[ctrl.kd, flat_damping, act.indices, self._dof_offset, self.num_joints], - device=wp_device, - ) - - self.stiffness = wp.to_torch(flat_stiffness.reshape((self._num_envs, self.num_joints))) - self.damping = wp.to_torch(flat_damping.reshape((self._num_envs, self.num_joints))) + sim_control.joint_computed_f = self._computed_effort def step(self, sim_state: Any, sim_control: Any, dt: float) -> None: """Zero actuated DOFs, step all actuators, and swap state buffers. @@ -242,6 +121,8 @@ def step(self, sim_state: Any, sim_control: Any, dt: float) -> None: on the PhysX backend. dt: Physics timestep [s]. """ + # Zero before scatter-add (actuators accumulate into this buffer). + self._computed_effort.zero_() for act in self.actuators: wp.launch( zero_at_indices_kernel, @@ -252,181 +133,200 @@ def step(self, sim_state: Any, sim_control: Any, dt: float) -> None: act.step(sim_state, sim_control, sa, sb, dt=dt) self._states_a, self._states_b = self._states_b, self._states_a - def reset(self, env_ids: Sequence[int] | slice | None = None) -> None: + def reset(self, env_ids: Sequence[int] | torch.Tensor | None = None) -> None: """Reset actuator states for the given environments. Args: - env_ids: Environment indices to reset. ``None`` or - ``slice(None)`` resets all environments. A partial slice - (e.g. ``slice(0, 5)``) is materialized to explicit indices. + env_ids: Environment indices to reset. ``None`` (or + ``slice(None)``, which IsaacLab callers sometimes pass) + resets all environments. Otherwise expects a torch tensor + or sequence of int indices. + + Newton's :meth:`Actuator.State.reset` expects a per-DOF boolean + mask of length ``num_actuators`` (= ``num_envs * dofs_per_actuator``), + not a per-env mask — each entry gates the corresponding column of + the actuator's state buffers (delay queue, controller integral, + etc.). We therefore build a per-actuator per-DOF mask from the + env mask before delegating to each state. """ if env_ids is None or env_ids == slice(None): - mask = None + for sa, sb in zip(self._states_a, self._states_b): + if sa is not None: + sa.reset(None) + if sb is not None: + sb.reset(None) + return + + if isinstance(env_ids, torch.Tensor): + if env_ids.numel() == 0: + return + idx = wp.from_torch(env_ids.to(device=self._device).contiguous().to(torch.int32), dtype=wp.int32) else: - # Normalize a partial slice to an explicit index list before - # building the wp.array — slices aren't iterable. - if isinstance(env_ids, slice): - env_ids = list(range(*env_ids.indices(self._num_envs))) - if isinstance(env_ids, torch.Tensor): - if env_ids.numel() == 0: - return - idx = wp.from_torch(env_ids.to(device=self._device).contiguous().to(torch.int32), dtype=wp.int32) - else: - if len(env_ids) == 0: - return - idx = wp.array(list(env_ids), dtype=wp.int32, device=self._device) - mask = wp.zeros(self._num_envs, dtype=wp.bool, device=self._device) - wp.launch(set_mask_kernel, dim=idx.shape[0], inputs=[mask, idx], device=self._device) + if len(env_ids) == 0: + return + idx = wp.array(list(env_ids), dtype=wp.int32, device=self._device) + env_mask = wp.zeros(self._num_envs, dtype=wp.bool, device=self._device) + wp.launch(set_mask_kernel, dim=idx.shape[0], inputs=[env_mask, idx], device=self._device) - for sa, sb in zip(self._states_a, self._states_b): + for act, sa, sb in zip(self.actuators, self._states_a, self._states_b): + per_dof_mask = wp.zeros(act.indices.shape[0], dtype=wp.bool, device=self._device) + wp.launch( + build_per_dof_env_mask_kernel, + dim=act.indices.shape[0], + inputs=[act.indices, env_mask, self._dof_offset, self.num_joints, per_dof_mask], + device=self._device, + ) if sa is not None: - sa.reset(mask) + sa.reset(per_dof_mask) if sb is not None: - sb.reset(mask) + sb.reset(per_dof_mask) @property def is_all_graphable(self) -> bool: """``True`` when all actuators are CUDA-graph-safe.""" return len(self.actuators) > 0 and all(a.is_graphable() for a in self.actuators) - def update_gain_at_env_ids( - self, - gain: str, - values: torch.Tensor | wp.array | float, - env_ids: wp.array, - ) -> wp.array: - """Scatter ``values`` into :attr:`stiffness` or :attr:`damping` at *env_ids*. + @classmethod + def from_usd( + cls, + stage: Any, + joint_names: list[str], + num_envs: int, + num_joints: int, + device: str, + articulation_prim_path: str | None = None, + ) -> "NewtonActuatorAdapter": + """Build an adapter from ``NewtonActuator`` prims authored on *stage*. - Shared backend-independent step used by ``write_*_to_sim`` to - keep the kp/kd buffer the adapter exposes for DR in sync with what - each Newton actuator's controller will receive. + PhysX-side counterpart of Newton's ``ModelBuilder.add_usd``: reads + the same prims and constructs matching + :class:`~newton.actuators.Actuator` objects. Joints with the same + controller, gains, clamping, and delay are merged into one + Actuator with combined indices. Used on the PhysX backend only — + Newton populates ``model.actuators`` itself. Args: - gain: ``"stiffness"`` or ``"damping"``. - values: Per-env-per-joint values shape ``(len(env_ids), num_joints)``, - or a scalar to broadcast to every (env, joint). - env_ids: Warp int32 array of env indices to update. - - Returns: - Warp view of the updated ``(num_envs, num_joints)`` buffer. + stage: The USD stage containing ``NewtonActuator`` prims. + joint_names: All joint names in the articulation. + num_envs: Number of environments. + num_joints: Joints per environment. + device: Warp device string (e.g. ``"cuda:0"``). + articulation_prim_path: Root prim path of env 0's + articulation. When set, only prims under this subtree are + considered; otherwise the whole stage is scanned. """ - if gain == "stiffness": - buf = self.stiffness - elif gain == "damping": - buf = self.damping - else: - raise ValueError(f"gain must be 'stiffness' or 'damping', got {gain!r}") - buf_wp = wp.from_torch(buf, dtype=wp.float32) - if isinstance(values, float): + actuators = _create_actuators_from_usd( + stage, joint_names, num_envs, num_joints, device, + articulation_prim_path=articulation_prim_path, + ) + return cls(actuators, num_envs, num_joints, dof_offset=0, device=device) + + +# --------------------------------------------------------------------------- +# Per-articulation initial-gain snapshot — consumed by +# ``randomize_actuator_gains`` to seed ``default_joint_*`` baselines. +# --------------------------------------------------------------------------- + + +def build_newton_actuator_defaults( + actuators: list[Actuator], + num_envs: int, + num_joints: int, + dof_offset: int, + device: str, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor | slice]: + """Snapshot the initial kp/kd of every Newton actuator owned by one articulation. + + Filters *actuators* to those whose env-0 DOF lives in + ``[dof_offset, dof_offset + num_joints)`` (a no-op on PhysX where the + adapter is already per-articulation; meaningful on Newton where the + global adapter holds actuators from every articulation), then + scatter-gathers their ``controller.kp`` / ``controller.kd`` into + contiguous ``(num_envs, num_joints)`` torch tensors and records which + articulation-local joints they cover. + + Args: + actuators: All Newton actuators visible to this articulation. + num_envs: Number of environments. + num_joints: Articulation-local joint count. + dof_offset: Offset of this articulation's DOFs in the env-major + global index space (``0`` on PhysX, view-dependent on Newton). + device: Warp device string (e.g. ``"cuda:0"``). + + Returns: + Tuple of ``(stiffness, damping, joint_indices)``: + + * ``stiffness``: Initial kp values, ``(num_envs, num_joints)``, articulation-local. + * ``damping``: Initial kd values, ``(num_envs, num_joints)``, articulation-local. + * ``joint_indices``: Articulation-local joint positions covered by + the adapter's actuators. ``slice(None)`` when every joint is + covered, otherwise an int32 tensor of column indices. + """ + arti_actuators = [ + act for act in actuators + if dof_offset <= int(act.indices.numpy()[0]) < dof_offset + num_joints + ] + + managed_local: set[int] = set() + for act in arti_actuators: + per_act = act.indices.shape[0] // num_envs + for global_dof in act.indices.numpy()[:per_act]: + local = int(global_dof) - dof_offset + if 0 <= local < num_joints: + managed_local.add(local) + joint_indices: torch.Tensor | slice + if len(managed_local) == num_joints: + joint_indices = slice(None) + else: + joint_indices = torch.tensor(sorted(managed_local), dtype=torch.int32, device=device) + + wp_device = wp.get_device(device) + flat_stiffness = wp.zeros(num_envs * num_joints, dtype=wp.float32, device=wp_device) + flat_damping = wp.zeros(num_envs * num_joints, dtype=wp.float32, device=wp_device) + for act in arti_actuators: + ctrl = act.controller + if hasattr(ctrl, "kp"): wp.launch( - fill_gain_at_envs_kernel, - dim=(env_ids.shape[0], self.num_joints), - inputs=[values, env_ids], - outputs=[buf_wp], - device=self._device, + scatter_gain_kernel, dim=act.indices.shape[0], + inputs=[ctrl.kp, flat_stiffness, act.indices, dof_offset, num_joints], + device=wp_device, ) - else: + if hasattr(ctrl, "kd"): wp.launch( - scatter_gain_at_envs_kernel, - dim=(env_ids.shape[0], self.num_joints), - inputs=[values, env_ids], - outputs=[buf_wp], - device=self._device, + scatter_gain_kernel, dim=act.indices.shape[0], + inputs=[ctrl.kd, flat_damping, act.indices, dof_offset, num_joints], + device=wp_device, ) - return buf_wp - - def set_view_propagator(self, root_view: Any) -> None: - """Configure gain propagation via Newton's simulator-side scatter API. + stiffness = wp.to_torch(flat_stiffness.reshape((num_envs, num_joints))) + damping = wp.to_torch(flat_damping.reshape((num_envs, num_joints))) + return stiffness, damping, joint_indices - Used on the Newton backend, where each per-actuator gain push goes - through ``ArticulationView.set_actuator_parameter``. Any one - articulation's view works since the call dispatches on the actuator - object (model-scoped), not on the view's articulation. - """ - def _push( - actuator: Actuator, controller: Any, attr: str, - values: wp.array, env_mask: wp.array, - ) -> None: - root_view.set_actuator_parameter( - actuator=actuator, component=controller, name=attr, - values=values, mask=env_mask, - ) - self._propagator = _push - def set_kernel_propagator(self) -> None: - """Configure gain propagation via the local ``gather_gain_kernel``. +# --------------------------------------------------------------------------- +# PhysX-only USD parsing +# --------------------------------------------------------------------------- - Used on the PhysX backend, where there is no simulator-side scatter - API. The kernel reads the adapter's per-DOF gain buffer at each - actuator's flat indices and writes into ``controller.kp`` / - ``controller.kd``. - """ - def _push( - actuator: Actuator, controller: Any, attr: str, - values: wp.array, env_mask: wp.array, - ) -> None: - wp.launch( - gather_gain_kernel, - dim=actuator.indices.shape[0], - inputs=[ - values.flatten(), getattr(controller, attr), actuator.indices, - env_mask, self._dof_offset, self.num_joints, - ], - device=self._device, - ) - self._propagator = _push - def write_stiffness_to_sim( - self, - stiffness: torch.Tensor | wp.array | float, - env_ids: wp.array, - env_mask: wp.array, - ) -> None: - """Update the kp buffer at *env_ids* and push the new values into each Newton controller.""" - self._write_gain_to_sim("stiffness", "kp", stiffness, env_ids, env_mask) - - def write_damping_to_sim( - self, - damping: torch.Tensor | wp.array | float, - env_ids: wp.array, - env_mask: wp.array, - ) -> None: - """Update the kd buffer at *env_ids* and push the new values into each Newton controller.""" - self._write_gain_to_sim("damping", "kd", damping, env_ids, env_mask) - - def _write_gain_to_sim( - self, - gain: str, - attr: str, - values: torch.Tensor | wp.array | float, - env_ids: wp.array, - env_mask: wp.array, - ) -> None: - """Shared body for :meth:`write_stiffness_to_sim` / :meth:`write_damping_to_sim`.""" - buf = self.update_gain_at_env_ids(gain, values, env_ids) - if self._propagator is None: - return - for newton_act in self.actuators: - ctrl = newton_act.controller - if hasattr(ctrl, attr): - self._propagator(newton_act, ctrl, attr, buf, env_mask) +def _actuator_signature(parsed: Any) -> tuple: + """Build a hashable key from a parsed actuator spec for grouping. - # -- private helpers ----------------------------------------------------- + Joints whose prims resolve to the same signature share identical + controller type, gains, clamping chain, and delay configuration and + can therefore be merged into a single :class:`~newton.actuators.Actuator` + with combined index arrays. + """ + ctrl_resolved = parsed.controller_class.resolve_arguments( + dict(parsed.controller_kwargs), + ) + ctrl_key = (parsed.controller_class, tuple(sorted(ctrl_resolved.items()))) - def _build_dof_map(self) -> list[int]: - """Build a per-DOF lookup: local DOF index -> actuator list index.""" - dof_to_actuator: list[int] = [-1] * self.num_joints + comp_keys: list[tuple] = [] + for comp_cls, comp_kwargs in parsed.component_specs: + resolved = comp_cls.resolve_arguments(comp_kwargs) + comp_keys.append((comp_cls, tuple(sorted(resolved.items())))) + comp_keys.sort(key=lambda t: t[0].__name__) - for act_idx, act in enumerate(self.actuators): - all_indices = act.indices.numpy() - num_per_act = len(all_indices) // self._num_envs - env0_indices = all_indices[:num_per_act] - for global_dof in env0_indices: - local_dof = global_dof - self._dof_offset - if 0 <= local_dof < self.num_joints: - dof_to_actuator[local_dof] = act_idx - - return dof_to_actuator + return (ctrl_key, tuple(comp_keys)) def _create_actuators_from_usd( diff --git a/source/isaaclab_newton/isaaclab_newton/actuators/authoring.py b/source/isaaclab_newton/isaaclab_newton/actuators/authoring.py index 999870e526f1..34a4e4539b64 100644 --- a/source/isaaclab_newton/isaaclab_newton/actuators/authoring.py +++ b/source/isaaclab_newton/isaaclab_newton/actuators/authoring.py @@ -381,11 +381,12 @@ def _resave_checkpoint_with_metadata( merged = {**existing_meta, **metadata} - tmp = tempfile.NamedTemporaryFile(suffix=".pt", delete=False) + with tempfile.NamedTemporaryFile(suffix=".pt", delete=False) as tmp: + tmp_path = tmp.name if is_torchscript: extra_out = {"metadata.json": json.dumps(merged)} - torch.jit.save(net, tmp.name, _extra_files=extra_out) + torch.jit.save(net, tmp_path, _extra_files=extra_out) else: - torch.save({"model": net, "metadata": merged}, tmp.name) + torch.save({"model": net, "metadata": merged}, tmp_path) - return tmp.name + return tmp_path diff --git a/source/isaaclab_newton/isaaclab_newton/actuators/kernels.py b/source/isaaclab_newton/isaaclab_newton/actuators/kernels.py index b850365c41ec..e710d97ccfd6 100644 --- a/source/isaaclab_newton/isaaclab_newton/actuators/kernels.py +++ b/source/isaaclab_newton/isaaclab_newton/actuators/kernels.py @@ -5,14 +5,19 @@ """Shared Warp kernels for the Newton actuator fast path.""" +import torch import warp as wp +from isaaclab.actuators import ActuatorBase, ImplicitActuator + # --------------------------------------------------------------------------- -# Adapter-internal kernels: per-DOF zeroing, env-mask building, and gain -# scatter/gather between the adapter's flat per-env-per-DOF buffer and the -# per-actuator controller arrays. Used by :class:`NewtonActuatorAdapter` -# (stepping, finalize, gain DR) and by the kernel-based propagator. +# Adapter / per-actuator helper kernels: per-DOF zeroing, env-mask building, +# per-DOF env-mask projection (used by :meth:`NewtonActuatorAdapter.reset`), +# and a partial scatter for DR gain updates that overwrites only the cells +# in a (env_ids × joint_ids) sub-grid of a Newton ``Actuator``'s controller +# parameter array. Used on the PhysX backend (no Newton view available); +# the Newton backend uses ``ArticulationView.set_actuator_parameter`` instead. # --------------------------------------------------------------------------- @@ -31,61 +36,89 @@ def set_mask_kernel(mask: wp.array(dtype=wp.bool), indices: wp.array(dtype=wp.in @wp.kernel(enable_backward=False) -def scatter_gain_kernel( - src: wp.array(dtype=wp.float32), - dst: wp.array(dtype=wp.float32), +def build_per_dof_env_mask_kernel( indices: wp.array(dtype=wp.uint32), + env_mask: wp.array(dtype=wp.bool), dof_offset: int, num_joints: int, + out_mask: wp.array(dtype=wp.bool), ): - """Scatter per-actuator ``src`` values into the adapter's flat per-env-per-DOF ``dst``.""" + """Build a per-DOF mask from a per-env mask, for one Newton actuator. + + Newton's :meth:`Actuator.State.reset` expects a mask of length + ``num_actuators`` (= ``num_envs * dofs_per_actuator``). Each entry + gates the corresponding column of the actuator's state buffers. This + kernel maps a per-env boolean mask onto that per-DOF layout via the + actuator's flat ``indices``. + """ i = wp.tid() global_dof = int(indices[i]) - dof_offset env = global_dof // num_joints - local_dof = global_dof % num_joints - dst[env * num_joints + local_dof] = src[i] + out_mask[i] = env_mask[env] @wp.kernel(enable_backward=False) -def gather_gain_kernel( - flat_src: wp.array(dtype=wp.float32), +def scatter_gain_kernel( + src: wp.array(dtype=wp.float32), dst: wp.array(dtype=wp.float32), indices: wp.array(dtype=wp.uint32), - env_mask: wp.array(dtype=wp.bool), dof_offset: int, num_joints: int, ): - """Gather from the adapter's flat ``(num_envs * num_joints)`` layout into a - per-actuator controller array, only for envs where ``env_mask`` is ``True``. + """Scatter per-actuator ``src`` values into a flat per-env-per-DOF ``dst``. + + Used at adapter finalize to snapshot each ``controller.kp`` / + ``controller.kd`` into the ``(num_envs, num_joints)`` torch tensor + that ``randomize_actuator_gains`` reads as + ``actuator.stiffness`` / ``.damping`` for its + ``default_joint_stiffness`` / ``default_joint_damping`` baseline. """ i = wp.tid() global_dof = int(indices[i]) - dof_offset env = global_dof // num_joints - if env_mask[env]: - local_dof = global_dof % num_joints - dst[i] = flat_src[env * num_joints + local_dof] - - -@wp.kernel(enable_backward=False) -def scatter_gain_at_envs_kernel( - in_data: wp.array2d(dtype=wp.float32), - env_ids: wp.array(dtype=wp.int32), - out_data: wp.array2d(dtype=wp.float32), -): - """Scatter ``in_data[i, j]`` into ``out_data[env_ids[i], j]`` for all (i, j).""" - i, j = wp.tid() - out_data[env_ids[i], j] = in_data[i, j] + local_dof = global_dof % num_joints + dst[env * num_joints + local_dof] = src[i] @wp.kernel(enable_backward=False) -def fill_gain_at_envs_kernel( - value: float, - env_ids: wp.array(dtype=wp.int32), - out_data: wp.array2d(dtype=wp.float32), +def patch_actuator_param_kernel( + indices: wp.array(dtype=wp.uint32), + env_id_pos: wp.array(dtype=wp.int32), + joint_id_pos: wp.array(dtype=wp.int32), + values: wp.array2d(dtype=wp.float32), + dof_offset: int, + num_joints: int, + dst: wp.array(dtype=wp.float32), ): - """Set ``out_data[env_ids[i], j] = value`` for all (i, j).""" - i, j = wp.tid() - out_data[env_ids[i], j] = value + """Per-actuator scatter for partial DR gain updates. + + For each slot ``i`` in the actuator's flat env-major ``indices``, derive + the (env, local-joint) pair, look it up against the dense position + arrays, and — when both axes are in the DR sub-grid — overwrite + ``dst[i]`` (the controller parameter) with ``values[e_pos, j_pos]``. + Cells outside the sub-grid are left untouched. + + Args: + indices: Actuator's flat indices into the (env-major) DOF layout. + env_id_pos: ``env_id_pos[env]`` gives the row in ``values`` for + envs being updated, ``-1`` otherwise. Length ``num_envs``. + joint_id_pos: ``joint_id_pos[joint]`` gives the column in + ``values`` for joints being updated, ``-1`` otherwise. + Length ``num_joints`` (articulation-local). + values: New parameter values shaped ``(len(env_ids), len(joint_ids))``. + dof_offset: Offset of this articulation's DOFs in the env-major + global index space (``0`` on PhysX, view-dependent on Newton). + num_joints: Articulation-local joint count. + dst: Per-actuator controller parameter array (e.g. ``controller.kp``). + """ + i = wp.tid() + global_dof = int(indices[i]) - dof_offset + env = global_dof // num_joints + joint = global_dof % num_joints + e_pos = env_id_pos[env] + j_pos = joint_id_pos[joint] + if e_pos >= 0 and j_pos >= 0: + dst[i] = values[e_pos, j_pos] # --------------------------------------------------------------------------- @@ -94,41 +127,32 @@ def fill_gain_at_envs_kernel( @wp.kernel(enable_backward=False) -def synch_torque_and_apply_implicit_feedforwards( +def sync_torque_telemetry( joint_pos: wp.array2d(dtype=wp.float32), joint_vel: wp.array2d(dtype=wp.float32), joint_pos_target: wp.array2d(dtype=wp.float32), joint_vel_target: wp.array2d(dtype=wp.float32), - joint_effort_target: wp.array2d(dtype=wp.float32), joint_stiffness: wp.array2d(dtype=wp.float32), joint_damping: wp.array2d(dtype=wp.float32), effort_limit: wp.array2d(dtype=wp.float32), joint_modes: wp.array(dtype=wp.int32), sim_bind_joint_effort: wp.array2d(dtype=wp.float32), + actuator_computed_effort: wp.array2d(dtype=wp.float32), computed: wp.array2d(dtype=wp.float32), applied: wp.array2d(dtype=wp.float32), ): - """In-graph post-actuator hook: route implicit FF and sync telemetry. - - For each (env, dof): - * Implicit DOF: write user FF to ``joint_f`` (sim integrates it - alongside the joint-drive PD), clamp the shadow PD ``kp*err_p + - kd*err_v`` to ``±effort_limit``, and write ``computed = applied - = PD_clipped + FF``. - * Explicit DOF: mirror Newton's post-actuator ``joint_f`` into - ``computed`` / ``applied``. - - Limitation: ``effort_limit`` here only clamps the **PD shadow** used - for telemetry. The simulator's joint drive applies its max-force - only to the PD term, so user feedforward effort can exceed - ``effort_limit`` once it lands in ``joint_f``. Limiting the *total* - applied effort would require Newton's motor-actuator path - (configurable via the Newton team), which may have negative perf - implications. + """In-graph post-actuator hook: fill ``computed`` / ``applied`` torque telemetry. + + For implicit DOFs we compute the shadow PD locally (no Newton actuator + runs on these); for explicit DOFs we read the pre-clamp effort the + actuators just scatter-added into ``actuator_computed_effort`` and the + post-clamp effort already in ``sim_bind_joint_effort`` (= ``joint_f``). + + Note: ``effort_limit`` clamps only the PD shadow used for implicit-DOF + telemetry; the FF written into ``joint_f`` is not bounded by it. """ i, j = wp.tid() if joint_modes[j] == 1: - sim_bind_joint_effort[i, j] = joint_effort_target[i, j] err_p = joint_pos_target[i, j] - joint_pos[i, j] err_v = joint_vel_target[i, j] - joint_vel[i, j] pd = joint_stiffness[i, j] * err_p + joint_damping[i, j] * err_v @@ -138,8 +162,28 @@ def synch_torque_and_apply_implicit_feedforwards( computed[i, j] = total applied[i, j] = total else: - val = sim_bind_joint_effort[i, j] - computed[i, j] = val - applied[i, j] = val + computed[i, j] = actuator_computed_effort[i, j] + applied[i, j] = sim_bind_joint_effort[i, j] +def build_implicit_dof_mask( + actuators: dict[str, ActuatorBase], + num_joints: int, + device: str, +) -> wp.array: + """Per-DOF mask consumed by :func:`sync_torque_telemetry`. + + Entry is ``1`` for DOFs covered by an + :class:`~isaaclab.actuators.ImplicitActuator` group, ``0`` otherwise. + """ + modes = torch.zeros(num_joints, dtype=torch.int32, device=device) + for actuator in actuators.values(): + if not isinstance(actuator, ImplicitActuator): + continue + j_ids = actuator.joint_indices + if j_ids == slice(None) or j_ids is None: + modes[:] = 1 + else: + modes[j_ids.long()] = 1 + return wp.from_torch(modes, dtype=wp.int32) + diff --git a/source/isaaclab_newton/isaaclab_newton/assets/articulation/articulation.py b/source/isaaclab_newton/isaaclab_newton/assets/articulation/articulation.py index 727f1a4d2395..d678b56c9c52 100644 --- a/source/isaaclab_newton/isaaclab_newton/assets/articulation/articulation.py +++ b/source/isaaclab_newton/isaaclab_newton/assets/articulation/articulation.py @@ -27,7 +27,7 @@ from isaaclab.assets.articulation.base_articulation import BaseArticulation try: - from isaaclab_newton.actuators import NewtonActuatorAdapter, build_implicit_dof_mask + from isaaclab_newton.actuators import NewtonActuatorAdapter, build_implicit_dof_mask, build_newton_actuator_defaults from isaaclab_newton.actuators import kernels as actuator_kernels _HAS_NEWTON_ACTUATORS = True @@ -250,9 +250,16 @@ def reset(self, env_ids: Sequence[int] | None = None, env_mask: wp.array | None # use ellipses object to skip initial indices. if (env_ids is None) or (env_ids == slice(None)): env_ids = slice(None) - # reset actuators + # reset Lab actuators registered on this articulation for actuator in self.actuators.values(): actuator.reset(env_ids) + # reset the global Newton actuator adapter (its ``_states_a/_b`` buffers + # carry per-env state — delay queues, neural hidden states — that must + # be cleared for the resetting envs). The adapter spans the whole model, + # so calling reset here resets state for every articulation that shares + # this env id; that's correct because env ids are world-scoped. + if self._has_newton_actuators and SimulationManager._adapter is not None: + SimulationManager._adapter.reset(env_ids) # reset external wrenches. self._instantaneous_wrench_composer.reset(env_ids, env_mask) self._permanent_wrench_composer.reset(env_ids, env_mask) @@ -291,15 +298,17 @@ def write_data_to_sim(self): if self._has_newton_actuators: # Raw targets go directly to Newton's control object. Newton PD - # consumes them for explicit (Newton-managed) joints; the + # consumes ``joint_act`` for explicit (Newton-managed) joints; the # solver's built-in joint drive does the PD for implicit joints - # (whose stiffness/damping are non-zero in sim). FF for explicit - # joints goes to control.joint_act; for implicit joints the - # in-graph post-actuator kernel writes the FF directly into - # control.joint_f. + # (whose stiffness/damping are non-zero in sim) and adds whatever + # is in ``joint_f`` as feedforward. We pre-fill ``joint_f`` with + # the user's effort target across all DOFs here; the adapter step + # will zero it at explicit DOFs and overwrite them with each + # actuator's computed effort, while implicit DOFs keep the FF. self.data._sim_bind_joint_position_target.assign(self._data._joint_pos_target) self.data._sim_bind_joint_velocity_target.assign(self._data._joint_vel_target) self.data._sim_bind_joint_act.assign(self._data._joint_effort_target) + self.data._sim_bind_joint_effort.assign(self._data._joint_effort_target) else: # Standard Lab actuator path self._apply_actuator_model() @@ -1423,6 +1432,79 @@ def write_joint_damping_to_sim_mask( # tell the physics engine that some of the joint properties have been updated SimulationManager.add_model_change(SolverNotifyFlags.JOINT_DOF_PROPERTIES) + def write_actuator_stiffness_to_sim( + self, + *, + stiffness: torch.Tensor, + env_ids: torch.Tensor, + joint_ids: torch.Tensor, + ) -> None: + """Write actuator kp at the (env_ids, joint_ids) sub-grid and propagate to controllers. + + Iterates the global adapter's Newton actuators and uses + :meth:`ArticulationView.get_actuator_parameter` / + :meth:`~ArticulationView.set_actuator_parameter` to patch each + controller's ``kp`` array. Actuators belonging to a different + articulation are no-ops because the view's per-DOF mapping + returns ``-1`` for DOFs outside this articulation's range. + + Args: + stiffness: Sub-grid of new kp values, shape ``(len(env_ids), len(joint_ids))``. + env_ids: 1D torch tensor of env indices. + joint_ids: 1D torch tensor of articulation-local joint indices. + + No-op when the Newton fast path is not active. + """ + self._write_actuator_param("kp", stiffness, env_ids, joint_ids) + + def write_actuator_damping_to_sim( + self, + *, + damping: torch.Tensor, + env_ids: torch.Tensor, + joint_ids: torch.Tensor, + ) -> None: + """Write actuator kd at the (env_ids, joint_ids) sub-grid and propagate to controllers.""" + self._write_actuator_param("kd", damping, env_ids, joint_ids) + + def _write_actuator_param( + self, + attr: str, + values: torch.Tensor, + env_ids: torch.Tensor, + joint_ids: torch.Tensor, + ) -> None: + """Shared body for :meth:`write_actuator_stiffness_to_sim` / :meth:`write_actuator_damping_to_sim`.""" + from isaaclab_newton.actuators import kernels as actuator_kernels # noqa: PLC0415 + + adapter = self.newton_actuator_adapter + if adapter is None: + return + + env_ids_wp = wp.from_torch( + env_ids.to(self.device, dtype=torch.int32).contiguous(), dtype=wp.int32, + ) + env_mask = wp.zeros(self.num_instances, dtype=wp.bool, device=self.device) + wp.launch( + actuator_kernels.set_mask_kernel, + dim=env_ids_wp.shape[0], inputs=[env_mask, env_ids_wp], device=self.device, + ) + + env_ids_long = env_ids.to(self.device, dtype=torch.long).unsqueeze(1) + joint_ids_long = joint_ids.to(self.device, dtype=torch.long).unsqueeze(0) + + for act in adapter.actuators: + ctrl = act.controller + if not hasattr(ctrl, attr): + continue + cur_wp = self._root_view.get_actuator_parameter(act, ctrl, attr) + cur_torch = wp.to_torch(cur_wp) + cur_torch[env_ids_long, joint_ids_long] = values.to(cur_torch.device, dtype=cur_torch.dtype) + self._root_view.set_actuator_parameter( + actuator=act, component=ctrl, name=attr, + values=cur_wp, mask=env_mask, + ) + def write_joint_position_limit_to_sim_index( self, *, @@ -3375,9 +3457,16 @@ def _process_actuators_cfg(self): self._has_implicit_actuators = False self._has_newton_actuators = False # Per-DOF implicit/explicit mask consumed by the in-graph kernel - # ``synch_torque_and_apply_implicit_feedforwards``. ``None`` when - # no Newton fast path is active. + # ``sync_torque_telemetry``. ``None`` when no Newton fast path is active. self._implicit_dof_mask: wp.array | None = None + # Reference to the global Newton actuator adapter (or ``None`` + # when this articulation has no explicit Newton actuators) and a + # per-articulation kp/kd snapshot consumed by + # ``randomize_actuator_gains`` to seed its DR baselines. + self.newton_actuator_adapter = None + self.newton_default_stiffness: torch.Tensor | None = None + self.newton_default_damping: torch.Tensor | None = None + self.newton_managed_local_joints: torch.Tensor | slice | None = None _use_newton_actuators = getattr(self._sim_cfg, "use_newton_actuators", False) @@ -3393,13 +3482,10 @@ def _process_actuators_cfg(self): # Enable the fast path even for all-implicit articulations: # the solver runs PD internally; Lab only forwards targets. self._has_newton_actuators = True + # Opt this articulation into the Newton fast path and (idempotently) + # build the single sim-level actuator adapter from ``model.actuators``. SimulationManager.activate_newton_actuator_path() - # Build (or share) the single sim-level actuator adapter. Idempotent — - # the first articulation to call this constructs it from - # ``model.actuators``; subsequent articulations reuse it. - SimulationManager.ensure_global_actuator_adapter(self._root_view) - # Zero the simulator's joint-drive PD on DOFs covered by an explicit # Lab actuator config in *this* articulation. The global Newton # adapter's actuator step writes their effort to ``joint_f`` @@ -3441,23 +3527,60 @@ def _process_actuators_cfg(self): # Run the implicit-DOF FF-routing + telemetry kernel inside the # captured graph, right after the actuator step. Closure captures # the buffers we need via ``self._data``. + from newton import Model as NewtonModel # noqa: PLC0415 + + # Per-articulation view of the global adapter's pre-clamp + # computed-effort buffer. Set up once here (the adapter is + # already built by ``activate_newton_actuator_path``) so the + # callback below has nothing to resolve. Falls back to a zero + # buffer for all-implicit scenes where no global adapter + # exists — the kernel only reads it on explicit DOFs. + adapter = SimulationManager._adapter + if adapter is not None: + dof_layout = self._root_view.frequency_layouts[NewtonModel.AttributeFrequency.JOINT_DOF] + if dof_layout.slice is not None: + arti_start = dof_layout.slice.start + elif dof_layout.indices is not None: + arti_start = int(dof_layout.indices.numpy()[0]) + else: + arti_start = 0 + self._data._sim_bind_joint_computed_effort = adapter.computed_effort_2d[ + :, arti_start : arti_start + self.num_joints + ] + self.newton_actuator_adapter = adapter + ( + self.newton_default_stiffness, + self.newton_default_damping, + self.newton_managed_local_joints, + ) = build_newton_actuator_defaults( + actuators=adapter.actuators, + num_envs=self.num_instances, + num_joints=self.num_joints, + dof_offset=arti_start, + device=self.device, + ) + else: + self._data._sim_bind_joint_computed_effort = wp.zeros( + (self.num_instances, self.num_joints), dtype=wp.float32, device=self.device, + ) + def _post_actuator() -> None: wp.launch( - actuator_kernels.synch_torque_and_apply_implicit_feedforwards, + actuator_kernels.sync_torque_telemetry, dim=(self.num_instances, self.num_joints), inputs=[ self._data.joint_pos.warp, self._data.joint_vel.warp, self._data._joint_pos_target, self._data._joint_vel_target, - self._data._joint_effort_target, self._data.joint_stiffness.warp, self._data.joint_damping.warp, self._data.joint_effort_limits.warp, self._implicit_dof_mask, + self._data._sim_bind_joint_effort, + self._data._sim_bind_joint_computed_effort, ], outputs=[ - self._data._sim_bind_joint_effort, self._data._computed_torque, self._data._applied_torque, ], diff --git a/source/isaaclab_newton/isaaclab_newton/physics/newton_manager.py b/source/isaaclab_newton/isaaclab_newton/physics/newton_manager.py index f656e69c982a..049521a24443 100644 --- a/source/isaaclab_newton/isaaclab_newton/physics/newton_manager.py +++ b/source/isaaclab_newton/isaaclab_newton/physics/newton_manager.py @@ -832,9 +832,9 @@ def start_simulation(cls) -> None: cls._control = cls._model.control() eval_fk(cls._model, cls._state_0.joint_q, cls._state_0.joint_qd, cls._state_0, None) - # The single global actuator adapter is built lazily on first - # call to ``ensure_global_actuator_adapter`` from any - # Newton-fast-path articulation after this point. + # The single global actuator adapter is built lazily on the first + # call to ``activate_newton_actuator_path`` from any Newton-fast-path + # articulation after this point. cls._adapter = None cls._use_newton_actuators_active = False @@ -1376,32 +1376,20 @@ def _is_all_graphable(cls) -> bool: @classmethod def activate_newton_actuator_path(cls) -> None: - """Signal that an articulation has opted into the Newton actuator fast path. - - Called by the articulation's ``_process_actuators_cfg`` whenever it - enters the ``use_newton_actuators=True`` branch — including the - all-implicit case where no :class:`NewtonActuatorAdapter` is built. - Required because :meth:`_is_all_graphable` can no longer rely on - adapter presence alone to distinguish the fast path from the - standard Lab path (which also has ``_adapter is None``). + """Opt an articulation into the Newton actuator fast path. + + Idempotent — called by every Newton-fast-path articulation's + ``_process_actuators_cfg``: + + 1. Sets :attr:`_use_newton_actuators_active`, which + :meth:`_is_all_graphable` checks (adapter presence alone + cannot distinguish the fast path from the standard Lab path). + 2. On first call, builds the single sim-level + :class:`NewtonActuatorAdapter` over the full flat DOF layout; + later calls reuse it. """ cls._use_newton_actuators_active = True - @classmethod - def ensure_global_actuator_adapter(cls, root_view: Any) -> None: - """Build the single sim-level :class:`NewtonActuatorAdapter` if not yet built. - - Idempotent — each Newton-fast-path articulation calls this from - ``_process_actuators_cfg`` after the model is finalized, passing - its own ``root_view``; the first call constructs the adapter from - ``cls._model.actuators`` and configures the gain propagator with - the supplied view, subsequent calls are no-ops. A single global - adapter avoids the multi-articulation clobbering that the - per-articulation ``register_adapter`` design suffered from. The - adapter operates on the full flat DOF layout (``dof_offset=0``, - ``num_joints`` = per-env total DOF count), so it iterates and - steps every actuator in the model in one pass. - """ if cls._adapter is not None: return if cls._model is None or not cls._model.actuators: @@ -1416,11 +1404,7 @@ def ensure_global_actuator_adapter(cls, root_view: Any) -> None: dof_offset=0, device=PhysicsManager._device, ) - cls._adapter.finalize() - # ``set_actuator_parameter`` dispatches on the actuator object (which - # is model-scoped), so any one articulation's view works for the - # global adapter spanning multiple articulations. - cls._adapter.set_view_propagator(root_view) + cls._adapter.finalize(cls._control) @classmethod def register_post_actuator_callback(cls, callback: Callable[[], None]) -> None: diff --git a/source/isaaclab_newton/isaaclab_newton/test/mock_interfaces/views/mock_articulation_view.py b/source/isaaclab_newton/isaaclab_newton/test/mock_interfaces/views/mock_articulation_view.py index 26761937667e..a75647624e5b 100644 --- a/source/isaaclab_newton/isaaclab_newton/test/mock_interfaces/views/mock_articulation_view.py +++ b/source/isaaclab_newton/isaaclab_newton/test/mock_interfaces/views/mock_articulation_view.py @@ -97,7 +97,9 @@ def _ensure_root_velocities(self) -> wp.array: def _ensure_attribute(self, name: str) -> wp.array: if self._attributes[name] is None: self._attributes[name] = self._create_default_attribute(name) - return self._attributes[name] + value = self._attributes[name] + assert value is not None + return value def _create_default_attribute(self, name: str) -> wp.array: N, B = self._num_envs, self._num_bodies @@ -250,6 +252,7 @@ def __init__( "joint_effort_limit": None, "body_f": None, "joint_f": None, + "joint_act": None, "joint_target_pos": None, "joint_target_vel": None, "joint_limit_ke": None, @@ -358,7 +361,9 @@ def _ensure_attribute(self, name: str) -> wp.array: """Lazily create an attribute array.""" if self._attributes[name] is None: self._attributes[name] = self._create_default_attribute(name) - return self._attributes[name] + value = self._attributes[name] + assert value is not None + return value def _create_default_attribute(self, name: str) -> wp.array: """Create a default attribute array based on name.""" @@ -383,6 +388,7 @@ def _create_default_attribute(self, name: str) -> wp.array: "joint_velocity_limit", "joint_effort_limit", "joint_f", + "joint_act", "joint_target_pos", "joint_target_vel", "joint_limit_ke", @@ -658,6 +664,7 @@ def set_random_mock_data(self) -> None: "joint_velocity_limit", "joint_effort_limit", "joint_f", + "joint_act", "joint_target_pos", "joint_target_vel", "joint_limit_ke", diff --git a/source/isaaclab_newton/test/assets/test_newton_actuators_newton.py b/source/isaaclab_newton/test/assets/test_newton_actuators_newton.py index 50818005f6c7..264481e81796 100644 --- a/source/isaaclab_newton/test/assets/test_newton_actuators_newton.py +++ b/source/isaaclab_newton/test/assets/test_newton_actuators_newton.py @@ -29,7 +29,6 @@ import torch import warp as wp -from isaaclab_assets import ANYMAL_C_CFG from isaaclab_newton.assets import Articulation from isaaclab_newton.physics import MJWarpSolverCfg, NewtonCfg from isaaclab_newton.physics import NewtonManager as SimulationManager @@ -38,6 +37,8 @@ from isaaclab.actuators import DCMotorCfg, DelayedPDActuatorCfg, IdealPDActuatorCfg, ImplicitActuatorCfg from isaaclab.sim import SimulationCfg, build_simulation_context +from isaaclab_assets import ANYMAL_C_CFG + # --------------------------------------------------------------------------- # Constants # --------------------------------------------------------------------------- @@ -116,8 +117,13 @@ def _run_simulation( newton_cfg: NewtonCfg = NEWTON_CFG, num_steps: int = NUM_STEPS, decimation: int = 1, + feedforward: float | None = None, ) -> dict: - """Run ANYmal-C and return recorded joint trajectories. + """Run ANYmal-C and return recorded trajectories + telemetry. + + Always records ``joint_pos``, ``joint_vel``, ``computed_torque``, and + ``applied_torque`` so callers don't need a separate "with telemetry" + runner. Optionally applies a constant per-DOF feedforward effort target. Args: actuators: Actuator config dict overriding ANYmal's defaults. @@ -125,35 +131,25 @@ def _run_simulation( dt: Physics timestep [s]. newton_cfg: Newton physics configuration. num_steps: Number of policy-level steps. - decimation: Actuator steps per policy step (Newton decimation loop). + decimation: Actuator steps per policy step (Newton's CUDA-graph + d-loop is used when all-graphable; otherwise an explicit Python + inner loop). + feedforward: When not ``None``, set a constant per-DOF feedforward + effort target. Used by the implicit-FF equivalence test. Returns: - Dict with ``joint_pos`` and ``joint_vel``, each a list of - ``(NUM_ENVS, num_joints)`` tensors. + Dict with ``joint_pos``, ``joint_vel``, ``computed_torque``, + ``applied_torque`` (lists of per-step tensors) plus ``target_pos`` + and ``target_vel`` snapshots. """ - sim_cfg = SimulationCfg( - dt=dt, - physics=newton_cfg, - use_newton_actuators=use_newton_actuators, - ) - + sim_cfg = SimulationCfg(dt=dt, physics=newton_cfg, use_newton_actuators=use_newton_actuators) with build_simulation_context( - device="cuda:0", - gravity_enabled=True, - add_ground_plane=True, - sim_cfg=sim_cfg, + device="cuda:0", gravity_enabled=True, add_ground_plane=True, sim_cfg=sim_cfg, ) as sim: sim._app_control_on_stop_handle = None - for i in range(NUM_ENVS): - sim_utils.create_prim( - f"/World/Env_{i}", "Xform", translation=(i * 3.0, 0, 0) - ) - - art_cfg = ANYMAL_C_CFG.replace( - actuators=actuators, - prim_path="/World/Env_.*/Robot", - ) + sim_utils.create_prim(f"/World/Env_{i}", "Xform", translation=(i * 3.0, 0, 0)) + art_cfg = ANYMAL_C_CFG.replace(actuators=actuators, prim_path="/World/Env_.*/Robot") articulation = Articulation(art_cfg) sim.reset() assert articulation.is_initialized @@ -171,11 +167,15 @@ def _run_simulation( init_pos = wp.to_torch(articulation.data.joint_pos).clone() target_pos = init_pos + TARGET_OFFSET target_vel = torch.zeros_like(init_pos) - articulation.set_joint_position_target_index(target=target_pos) articulation.set_joint_velocity_target_index(target=target_vel) + if feedforward is not None: + articulation.set_joint_effort_target_index( + target=torch.full_like(init_pos, feedforward), + ) recorded_pos, recorded_vel = [], [] + recorded_computed, recorded_applied = [], [] for _ in range(num_steps): if handles_dec: articulation.write_data_to_sim() @@ -186,68 +186,6 @@ def _run_simulation( articulation.write_data_to_sim() sim.step() articulation.update(dt) - - recorded_pos.append(wp.to_torch(articulation.data.joint_pos).clone()) - recorded_vel.append(wp.to_torch(articulation.data.joint_vel).clone()) - - return {"joint_pos": recorded_pos, "joint_vel": recorded_vel} - - -def _run_simulation_with_telemetry( - actuators: dict, - use_newton_actuators: bool, - *, - dt: float = DT, - newton_cfg: NewtonCfg = NEWTON_CFG, - num_steps: int = NUM_STEPS, - decimation: int = 1, -) -> dict: - """Like :func:`_run_simulation` but also records ``computed_torque`` / ``applied_torque``.""" - sim_cfg = SimulationCfg( - dt=dt, - physics=newton_cfg, - use_newton_actuators=use_newton_actuators, - ) - - with build_simulation_context( - device="cuda:0", - gravity_enabled=True, - add_ground_plane=True, - sim_cfg=sim_cfg, - ) as sim: - sim._app_control_on_stop_handle = None - - for i in range(NUM_ENVS): - sim_utils.create_prim( - f"/World/Env_{i}", "Xform", translation=(i * 3.0, 0, 0) - ) - - art_cfg = ANYMAL_C_CFG.replace( - actuators=actuators, - prim_path="/World/Env_.*/Robot", - ) - articulation = Articulation(art_cfg) - sim.reset() - assert articulation.is_initialized - - if use_newton_actuators and decimation > 1: - SimulationManager.set_decimation(decimation) - - init_pos = wp.to_torch(articulation.data.joint_pos).clone() - target_pos = init_pos + TARGET_OFFSET - target_vel = torch.zeros_like(init_pos) - - articulation.set_joint_position_target_index(target=target_pos) - articulation.set_joint_velocity_target_index(target=target_vel) - - recorded_pos, recorded_vel = [], [] - recorded_computed, recorded_applied = [], [] - for _ in range(num_steps): - for _ in range(decimation): - articulation.write_data_to_sim() - sim.step() - articulation.update(dt) - recorded_pos.append(wp.to_torch(articulation.data.joint_pos).clone()) recorded_vel.append(wp.to_torch(articulation.data.joint_vel).clone()) recorded_computed.append(wp.to_torch(articulation.data.computed_torque).clone()) @@ -278,32 +216,36 @@ class _EquivalenceTestBase(unittest.TestCase): __test__ = False actuators: dict = {} + feedforward: float | None = None + dt: float = DT + newton_cfg: NewtonCfg = NEWTON_CFG + num_steps: int = NUM_STEPS + decimation: int = 1 pos_atol: float = 2e-3 pos_rtol: float = 1e-3 - vel_atol: float = 0.1 + vel_atol: float = 1e-2 vel_rtol: float = 1e-2 - # Torque equivalence tolerates a one-physics-step timing offset: - # standard path computes applied_torque pre-step, fast path post-step. - # The difference scales with kp*v*dt + kd*a*dt — largest in the first - # transient steps before targets are tracked, so skip those. - torque_atol: float = 5.0 - torque_rtol: float = 0.15 - torque_skip_steps: int = 3 + torque_atol: float = 1e-3 + torque_rtol: float = 1e-3 @classmethod def setUpClass(cls): - cls.lab_result = _run_simulation_with_telemetry(cls.actuators, use_newton_actuators=False) - cls.newton_result = _run_simulation_with_telemetry(cls.actuators, use_newton_actuators=True) + kwargs = dict( + feedforward=cls.feedforward, + dt=cls.dt, + newton_cfg=cls.newton_cfg, + num_steps=cls.num_steps, + decimation=cls.decimation, + ) + cls.lab_result = _run_simulation(cls.actuators, use_newton_actuators=False, **kwargs) + cls.newton_result = _run_simulation(cls.actuators, use_newton_actuators=True, **kwargs) def test_joint_positions_match(self): for step_i, (lab, newton) in enumerate( zip(self.lab_result["joint_pos"], self.newton_result["joint_pos"]) ): torch.testing.assert_close( - lab, - newton, - atol=self.pos_atol, - rtol=self.pos_rtol, + lab, newton, atol=self.pos_atol, rtol=self.pos_rtol, msg=f"Joint positions diverged at step {step_i}", ) @@ -312,10 +254,7 @@ def test_joint_velocities_match(self): zip(self.lab_result["joint_vel"], self.newton_result["joint_vel"]) ): torch.testing.assert_close( - lab, - newton, - atol=self.vel_atol, - rtol=self.vel_rtol, + lab, newton, atol=self.vel_atol, rtol=self.vel_rtol, msg=f"Joint velocities diverged at step {step_i}", ) @@ -323,21 +262,19 @@ def test_applied_torque_match(self): for step_i, (lab, newton) in enumerate( zip(self.lab_result["applied_torque"], self.newton_result["applied_torque"]) ): - if step_i < self.torque_skip_steps: - continue torch.testing.assert_close( - lab, - newton, - atol=self.torque_atol, - rtol=self.torque_rtol, + lab, newton, atol=self.torque_atol, rtol=self.torque_rtol, msg=f"applied_torque diverged at step {step_i}", ) - def test_trajectories_not_trivial(self): - first = self.lab_result["joint_pos"][0] - last = self.lab_result["joint_pos"][-1] - diff = (last - first).abs().max().item() - self.assertGreater(diff, 0.01, "Joints did not move — test is trivial") + def test_computed_torque_match(self): + for step_i, (lab, newton) in enumerate( + zip(self.lab_result["computed_torque"], self.newton_result["computed_torque"]) + ): + torch.testing.assert_close( + lab, newton, atol=self.torque_atol, rtol=self.torque_rtol, + msg=f"computed_torque diverged at step {step_i}", + ) # --------------------------------------------------------------------------- @@ -420,191 +357,22 @@ class TestImplicitOnlyEquivalence(_EquivalenceTestBase): actuators = IMPLICIT_ONLY_ACTUATORS -class TestExplicitOnlyTelemetry(unittest.TestCase): - """Explicit-only Newton actuators: telemetry copies from ``joint_f``.""" - - @classmethod - def setUpClass(cls): - cls.result = _run_simulation_with_telemetry( - DC_MOTOR_ACTUATORS, use_newton_actuators=True, - ) - - def test_telemetry_is_nonzero(self): - last = self.result["applied_torque"][-1] - self.assertGreater( - last.abs().max().item(), - 1e-3, - "applied_torque is all-zero — explicit-DOF telemetry path did not run", - ) - - def test_computed_equals_applied_explicit(self): - """For Newton-managed DOFs, both buffers carry ``joint_f`` (post-clamp).""" - for step_i, (comp, app) in enumerate( - zip(self.result["computed_torque"], self.result["applied_torque"]) - ): - torch.testing.assert_close( - comp, - app, - atol=1e-5, - rtol=1e-5, - msg=f"computed != applied for explicit-only at step {step_i}", - ) - - -class TestImplicitOnlyTelemetry(unittest.TestCase): - """Implicit-only fast path: shadow-PD telemetry matches the Lab actuator path.""" - - @classmethod - def setUpClass(cls): - cls.lab_result = _run_simulation_with_telemetry( - IMPLICIT_ONLY_ACTUATORS, use_newton_actuators=False, - ) - cls.newton_result = _run_simulation_with_telemetry( - IMPLICIT_ONLY_ACTUATORS, use_newton_actuators=True, - ) - - def test_telemetry_is_nonzero(self): - """Telemetry kernel actually fired (not stuck at default zeros).""" - last = self.newton_result["computed_torque"][-1] - max_abs = last.abs().max().item() - self.assertGreater(max_abs, 1e-3, "computed_torque is all-zero — telemetry kernel did not run") - - def test_telemetry_matches_lab_path(self): - """Newton fast-path telemetry agrees with the Lab actuator-path telemetry.""" - for step_i, (lab_comp, newton_comp) in enumerate( - zip(self.lab_result["computed_torque"], self.newton_result["computed_torque"]) - ): - torch.testing.assert_close( - newton_comp, - lab_comp, - atol=5e-2, - rtol=1e-2, - msg=f"computed_torque diverged from Lab path at step {step_i}", - ) - - def test_applied_equals_computed_when_no_clip(self): - """Implicit cfg has no effort_limit → applied == computed.""" - for step_i, (comp, app) in enumerate( - zip(self.newton_result["computed_torque"], self.newton_result["applied_torque"]) - ): - torch.testing.assert_close( - app, - comp, - atol=1e-5, - rtol=1e-5, - msg=f"applied_torque != computed_torque at step {step_i} (no clip expected)", - ) - - # --------------------------------------------------------------------------- # Implicit + non-zero feedforward effort target # --------------------------------------------------------------------------- -def _run_simulation_with_ff( - actuators: dict, - use_newton_actuators: bool, - *, - feedforward: float, - dt: float = DT, - newton_cfg: NewtonCfg = NEWTON_CFG, - num_steps: int = NUM_STEPS, -) -> dict: - """Run with both a position target and a non-zero feedforward effort target.""" - sim_cfg = SimulationCfg( - dt=dt, - physics=newton_cfg, - use_newton_actuators=use_newton_actuators, - ) - - with build_simulation_context( - device="cuda:0", - gravity_enabled=True, - add_ground_plane=True, - sim_cfg=sim_cfg, - ) as sim: - sim._app_control_on_stop_handle = None - - for i in range(NUM_ENVS): - sim_utils.create_prim( - f"/World/Env_{i}", "Xform", translation=(i * 3.0, 0, 0) - ) - - art_cfg = ANYMAL_C_CFG.replace( - actuators=actuators, - prim_path="/World/Env_.*/Robot", - ) - articulation = Articulation(art_cfg) - sim.reset() - assert articulation.is_initialized - - init_pos = wp.to_torch(articulation.data.joint_pos).clone() - target_pos = init_pos + TARGET_OFFSET - target_vel = torch.zeros_like(init_pos) - target_eff = torch.full_like(init_pos, feedforward) - - articulation.set_joint_position_target_index(target=target_pos) - articulation.set_joint_velocity_target_index(target=target_vel) - articulation.set_joint_effort_target_index(target=target_eff) - - recorded_pos, recorded_vel = [], [] - for _ in range(num_steps): - articulation.write_data_to_sim() - sim.step() - articulation.update(dt) - recorded_pos.append(wp.to_torch(articulation.data.joint_pos).clone()) - recorded_vel.append(wp.to_torch(articulation.data.joint_vel).clone()) - - return {"joint_pos": recorded_pos, "joint_vel": recorded_vel} - - -class TestImplicitWithFeedforwardEquivalence(unittest.TestCase): +class TestImplicitWithFeedforwardEquivalence(_EquivalenceTestBase): """Implicit-only actuators with a non-zero feedforward effort target. - For implicit actuators, the user's feedforward effort must be additive - on top of the simulator's PD on both the Lab path and the Newton fast - path. This test commands a constant FF effort and compares joint - trajectories between the two paths. + Verifies that the user's FF effort lands additively on top of the + simulator's joint-drive PD identically on both Lab and Newton paths. """ - FEEDFORWARD = 5.0 - pos_atol = 2e-3 - pos_rtol = 1e-3 - vel_atol = 0.1 - vel_rtol = 1e-2 - - @classmethod - def setUpClass(cls): - cls.lab_result = _run_simulation_with_ff( - IMPLICIT_ONLY_ACTUATORS, use_newton_actuators=False, feedforward=cls.FEEDFORWARD, - ) - cls.newton_result = _run_simulation_with_ff( - IMPLICIT_ONLY_ACTUATORS, use_newton_actuators=True, feedforward=cls.FEEDFORWARD, - ) - - def test_joint_positions_match(self): - for step_i, (lab, newton) in enumerate( - zip(self.lab_result["joint_pos"], self.newton_result["joint_pos"]) - ): - torch.testing.assert_close( - lab, newton, atol=self.pos_atol, rtol=self.pos_rtol, - msg=f"Joint positions diverged at step {step_i}", - ) - - def test_joint_velocities_match(self): - for step_i, (lab, newton) in enumerate( - zip(self.lab_result["joint_vel"], self.newton_result["joint_vel"]) - ): - torch.testing.assert_close( - lab, newton, atol=self.vel_atol, rtol=self.vel_rtol, - msg=f"Joint velocities diverged at step {step_i}", - ) - - def test_trajectories_not_trivial(self): - first = self.lab_result["joint_pos"][0] - last = self.lab_result["joint_pos"][-1] - diff = (last - first).abs().max().item() - self.assertGreater(diff, 0.01, "Joints did not move — test is trivial") + __test__ = True + actuators = IMPLICIT_ONLY_ACTUATORS + feedforward = 2.0 + torque_atol = 0.5 # --------------------------------------------------------------------------- @@ -703,6 +471,190 @@ def test_cartpole_matches_lab(self): ) +# --------------------------------------------------------------------------- +# Domain randomization via events.py — Newton backend +# --------------------------------------------------------------------------- + + +class _MockScene: + """Minimal stand-in for ``InteractiveScene`` accepted by ``ManagerTermBase``.""" + + def __init__(self, assets: dict, num_envs: int): + self._assets = assets + self.num_envs = num_envs + + def __getitem__(self, name: str): + return self._assets[name] + + +class _MockEnv: + """Minimal stand-in for ``ManagerBasedEnv`` for invoking DR terms. + + ``randomize_actuator_gains`` only reads ``env.scene[name]`` and + ``env.scene.num_envs`` (plus ``env.num_envs`` / ``env.device`` from the + ``ManagerTermBase`` properties). No simulator access is needed because + the DR term reaches the actuator adapter via ``self.asset.newton_actuator_adapter``. + """ + + def __init__(self, assets: dict, num_envs: int, device: str): + self.scene = _MockScene(assets, num_envs) + self.num_envs = num_envs + self.device = device + + +def _build_dr_term(env, asset_name, joint_ids=None): + from isaaclab.envs.mdp.events import randomize_actuator_gains # noqa: PLC0415 + from isaaclab.managers import EventTermCfg, SceneEntityCfg # noqa: PLC0415 + + asset_cfg = SceneEntityCfg(asset_name) + if joint_ids is not None: + asset_cfg.joint_ids = joint_ids + cfg = EventTermCfg( + func=randomize_actuator_gains, + params={ + "asset_cfg": asset_cfg, + "stiffness_distribution_params": (100.0, 100.0), + "damping_distribution_params": (5.0, 5.0), + "operation": "abs", + "distribution": "uniform", + }, + ) + return randomize_actuator_gains(cfg, env), asset_cfg + + +class TestRandomizeActuatorGainsViaEventsNewton(unittest.TestCase): + """End-to-end DR test for the Newton backend. + + Drives ``randomize_actuator_gains`` (events.py) and verifies the new + kp/kd values land on the controllers of the articulation's Newton + actuators — exercising the full path: events → + ``write_actuator_stiffness_to_sim`` → per-actuator + ``ArticulationView.set_actuator_parameter`` (with the per-DOF mapping + silently skipping actuators that belong to other articulations). + + With ``operation="abs"`` and ``distribution="uniform"`` over a + degenerate range ``(K, K)``, every randomized cell is set to exactly + ``K`` — so the assertions are deterministic. + """ + + @staticmethod + def _gather_param(articulation, attr) -> torch.Tensor: + """Read ``controller.`` for every Newton actuator via the view. + + Iterates the global adapter's actuator list. ``get_actuator_parameter`` + returns zeros for DOFs that don't belong to this articulation's + view (the per-DOF mapping skips them), so summing across all + actuators yields a clean ``(num_envs, num_joints)`` snapshot for + this articulation. + """ + n_env = articulation.num_instances + n_j = articulation.num_joints + out = torch.zeros((n_env, n_j), device=articulation.device) + adapter = SimulationManager._adapter + if adapter is None: + return out + for act in adapter.actuators: + ctrl = act.controller + if not hasattr(ctrl, attr): + continue + cur_wp = articulation._root_view.get_actuator_parameter(act, ctrl, attr) + out += wp.to_torch(cur_wp) + return out + + def test_single_articulation(self): + sim_cfg = SimulationCfg(dt=DT, physics=NEWTON_CFG, use_newton_actuators=True) + with build_simulation_context( + device="cuda:0", gravity_enabled=True, add_ground_plane=True, sim_cfg=sim_cfg, + ) as sim: + sim._app_control_on_stop_handle = None + for i in range(NUM_ENVS): + sim_utils.create_prim(f"/World/Env_{i}", "Xform", translation=(i * 3.0, 0, 0)) + art_cfg = ANYMAL_C_CFG.replace( + actuators=IDEAL_PD_ACTUATORS, prim_path="/World/Env_.*/Robot", + ) + anymal = Articulation(art_cfg) + sim.reset() + + adapter = SimulationManager._adapter + self.assertIsNotNone(adapter, "Newton adapter should exist with use_newton_actuators=True") + kp_before = self._gather_param(anymal, "kp").clone() + kd_before = self._gather_param(anymal, "kd").clone() + + env = _MockEnv({"robot": anymal}, NUM_ENVS, anymal.device) + term, asset_cfg = _build_dr_term(env, "robot") + env_ids = torch.tensor([0], device=anymal.device, dtype=torch.long) + + term( + env, env_ids=env_ids, asset_cfg=asset_cfg, + stiffness_distribution_params=(100.0, 100.0), + damping_distribution_params=(5.0, 5.0), + operation="abs", distribution="uniform", + ) + + kp_after = self._gather_param(anymal, "kp") + kd_after = self._gather_param(anymal, "kd") + n = anymal.num_joints + torch.testing.assert_close(kp_after[0], torch.full((n,), 100.0, device=anymal.device)) + torch.testing.assert_close(kd_after[0], torch.full((n,), 5.0, device=anymal.device)) + # Other envs untouched. + for env_idx in range(1, NUM_ENVS): + torch.testing.assert_close(kp_after[env_idx], kp_before[env_idx]) + torch.testing.assert_close(kd_after[env_idx], kd_before[env_idx]) + + def test_two_articulations(self): + from isaaclab_assets import CARTPOLE_CFG # noqa: PLC0415 + + sim_cfg = SimulationCfg(dt=DT, physics=NEWTON_CFG, use_newton_actuators=True) + with build_simulation_context( + device="cuda:0", gravity_enabled=True, add_ground_plane=True, sim_cfg=sim_cfg, + ) as sim: + sim._app_control_on_stop_handle = None + for i in range(NUM_ENVS): + sim_utils.create_prim(f"/World/Env_{i}", "Xform", translation=(i * 6.0, 0, 0)) + + anymal_cfg = ANYMAL_C_CFG.replace(actuators=IDEAL_PD_ACTUATORS, prim_path="/World/Env_.*/Anymal") + cartpole_cfg = CARTPOLE_CFG.replace( + actuators=CARTPOLE_EXPLICIT_ACTUATORS, prim_path="/World/Env_.*/Cartpole", + ) + cartpole_cfg.init_state = cartpole_cfg.init_state.replace(pos=(0.0, 3.0, 2.0)) + anymal = Articulation(anymal_cfg) + cartpole = Articulation(cartpole_cfg) + sim.reset() + + self.assertIsNotNone(SimulationManager._adapter) + + anymal_kp_before = self._gather_param(anymal, "kp").clone() + anymal_kd_before = self._gather_param(anymal, "kd").clone() + cp_kp_before = self._gather_param(cartpole, "kp").clone() + cp_kd_before = self._gather_param(cartpole, "kd").clone() + + env = _MockEnv({"anymal": anymal, "cartpole": cartpole}, NUM_ENVS, anymal.device) + term, asset_cfg = _build_dr_term(env, "cartpole") + env_ids = torch.tensor([0], device=anymal.device, dtype=torch.long) + + term( + env, env_ids=env_ids, asset_cfg=asset_cfg, + stiffness_distribution_params=(100.0, 100.0), + damping_distribution_params=(5.0, 5.0), + operation="abs", distribution="uniform", + ) + + cp_kp_after = self._gather_param(cartpole, "kp") + cp_kd_after = self._gather_param(cartpole, "kd") + n_cp = cartpole.num_joints + torch.testing.assert_close(cp_kp_after[0], torch.full((n_cp,), 100.0, device=anymal.device)) + torch.testing.assert_close(cp_kd_after[0], torch.full((n_cp,), 5.0, device=anymal.device)) + + # ANYmal is untouched (DR was scoped to cartpole). + torch.testing.assert_close(self._gather_param(anymal, "kp"), anymal_kp_before) + torch.testing.assert_close(self._gather_param(anymal, "kd"), anymal_kd_before) + + # Cartpole's other envs are also untouched (env_ids=[0] only). + for env_idx in range(1, NUM_ENVS): + torch.testing.assert_close(cp_kp_after[env_idx], cp_kp_before[env_idx]) + torch.testing.assert_close(cp_kd_after[env_idx], cp_kd_before[env_idx]) + + # --------------------------------------------------------------------------- # DelayedPD equivalence: PD with actuator command delay # --------------------------------------------------------------------------- @@ -745,21 +697,11 @@ def test_controller_is_pd(self): for a in self.result["actuator_info"]: self.assertEqual(a["controller_type"], "ControllerPD") - def test_trajectories_not_trivial(self): - first = self.result["joint_pos"][0] - last = self.result["joint_pos"][-1] - diff = (last - first).abs().max().item() - self.assertGreater(diff, 0.01, "Joints did not move — test is trivial") - # --------------------------------------------------------------------------- -# Decimation test: CUDA graph capture with actuator decimation loop +# Decimation tests: re-run equivalence with decimation > 1 + CUDA graph capture # --------------------------------------------------------------------------- -DT_DEC = 1.0 / 100.0 -DECIMATION = 2 -NUM_POLICY_STEPS_DEC = 5 - NEWTON_CFG_DEC = NewtonCfg( solver_cfg=MJWarpSolverCfg( njmax=500, @@ -776,282 +718,160 @@ def test_trajectories_not_trivial(self): ) -class _DecimationTestBase(unittest.TestCase): - """Base for decimation tests with CUDA graph capture. - - Policy runs at 50 Hz, actuators at 100 Hz, physics at 200 Hz. - The Newton path captures the full decimation loop as a CUDA graph; - the Lab path runs an explicit per-substep loop. - - Subclasses set ``actuators`` to the config under test. - """ - - __test__ = False - actuators: dict = {} - pos_atol: float = 2e-3 - pos_rtol: float = 1e-3 - vel_atol: float = 0.1 - vel_rtol: float = 1e-2 - - @classmethod - def setUpClass(cls): - cls.lab_result = _run_simulation( - cls.actuators, - use_newton_actuators=False, - dt=DT_DEC, - newton_cfg=NEWTON_CFG_DEC, - num_steps=NUM_POLICY_STEPS_DEC, - decimation=DECIMATION, - ) - cls.newton_result = _run_simulation( - cls.actuators, - use_newton_actuators=True, - dt=DT_DEC, - newton_cfg=NEWTON_CFG_DEC, - num_steps=NUM_POLICY_STEPS_DEC, - decimation=DECIMATION, - ) - - def test_joint_positions_match(self): - for step_i, (lab, newton) in enumerate( - zip(self.lab_result["joint_pos"], self.newton_result["joint_pos"]) - ): - torch.testing.assert_close( - lab, - newton, - atol=self.pos_atol, - rtol=self.pos_rtol, - msg=f"Positions diverged at policy step {step_i}", - ) - - def test_joint_velocities_match(self): - for step_i, (lab, newton) in enumerate( - zip(self.lab_result["joint_vel"], self.newton_result["joint_vel"]) - ): - torch.testing.assert_close( - lab, - newton, - atol=self.vel_atol, - rtol=self.vel_rtol, - msg=f"Velocities diverged at policy step {step_i}", - ) - - def test_trajectories_not_trivial(self): - first = self.lab_result["joint_pos"][0] - last = self.lab_result["joint_pos"][-1] - diff = (last - first).abs().max().item() - self.assertGreater(diff, 0.01, "Joints did not move — test is trivial") - - -class TestDecimationDCMotor(_DecimationTestBase): - """DCMotor with decimation=2 and CUDA graph capture.""" - - __test__ = True - actuators = DC_MOTOR_ACTUATORS - - -class TestDecimationIdealPD(_DecimationTestBase): - """IdealPD with decimation=2 and CUDA graph capture.""" +class _DecimationMixin: + """Common knobs for decimation/CUDA-graph variants of equivalence classes.""" __test__ = True - actuators = IDEAL_PD_ACTUATORS - + dt = 1.0 / 100.0 + newton_cfg = NEWTON_CFG_DEC + num_steps = 5 + decimation = 2 -class TestDecimationDelayedPD(_DecimationTestBase): - """DelayedPD with decimation=2 and CUDA graph capture. - Delay buffers must be correctly stepped inside the captured graph. - """ +class TestDecimationDCMotor(_DecimationMixin, TestDCMotorEquivalence): + """DCMotor — same equivalence checks, with decimation=2 + CUDA graph.""" - __test__ = True - actuators = DELAYED_PD_ACTUATORS +class TestDecimationIdealPD(_DecimationMixin, TestIdealPDEquivalence): + """IdealPD — decimation=2 + CUDA graph.""" -class TestDecimationMixed(_DecimationTestBase): - """Mixed actuators (IdealPD + DCMotor) with decimation=2 and CUDA graph.""" - __test__ = True - actuators = MIXED_ACTUATORS +class TestDecimationDelayedPD(_DecimationMixin, TestDelayedPDEquivalence): + """DelayedPD — decimation=2 + CUDA graph (delay queue stepped inside the captured graph).""" -class TestDecimationRemotizedPD(_DecimationTestBase): - """RemotizedPD (PD + delay + position-based clamping) with decimation=2.""" - - __test__ = True - - @classmethod - def setUpClass(cls): - from isaaclab.actuators.actuator_pd_cfg import RemotizedPDActuatorCfg # noqa: PLC0415 - - cls.actuators = { - "hips": IdealPDActuatorCfg( - joint_names_expr=[".*HAA", ".*HFE"], - stiffness=40.0, - damping=5.0, - effort_limit=80.0, - ), - "knees": RemotizedPDActuatorCfg( - joint_names_expr=[".*KFE"], - stiffness=60.0, - damping=1.5, - effort_limit=80.0, - max_delay=3, - joint_parameter_lookup=SPOT_KNEE_LOOKUP, - ), - } - super().setUpClass() +class TestDecimationMixed(_DecimationMixin, TestMixedActuatorEquivalence): + """Mixed (IdealPD + DCMotor) — decimation=2 + CUDA graph.""" # --------------------------------------------------------------------------- -# Partial environment reset: verify per-env reset equivalence +# Per-env reset: actuator state isolation # --------------------------------------------------------------------------- RESET_WARMUP_STEPS = 3 -RESET_TOTAL_STEPS = 10 -def _run_simulation_with_reset( - actuators: dict, - use_newton_actuators: bool, - *, - dt: float = DT, - newton_cfg: NewtonCfg = NEWTON_CFG, -) -> dict: - """Run ANYmal-C with a mid-simulation reset of env 0 only. - Steps ``RESET_WARMUP_STEPS``, then resets env 0 to its initial joint state - (zeroing velocity), then steps ``RESET_TOTAL_STEPS - RESET_WARMUP_STEPS`` - more. Returns per-step joint positions and velocities. +class TestActuatorStateReset(unittest.TestCase): + """Reset must clear the actuator state buffers for the requested envs only. - This exercises the actuator state reset path (delay buffers, neural - hidden states, etc.) for a subset of environments. + Inspects ``adapter.actuators[i].state.delay_state.num_pushes`` directly: - Args: - actuators: Actuator config dict overriding ANYmal's defaults. - use_newton_actuators: Use Newton-native actuators when ``True``. - dt: Physics timestep [s]. - newton_cfg: Newton physics configuration. + * After warmup, ``num_pushes > 0`` for every DOF (buffer was populated). + * After ``articulation.reset(env_ids=[0])``, the entries for env 0's DOFs + must be ``0`` and the entries for env 1's DOFs must remain ``> 0``. - Returns: - Dict with ``joint_pos`` and ``joint_vel``, each a list of - ``(NUM_ENVS, num_joints)`` tensors. + Done independently on Lab and Newton paths. Lab inspects the + ``positions_delay_buffer._circular_buffer`` of its DelayedPDActuator; + Newton inspects the model-wide adapter's per-actuator state. """ - sim_cfg = SimulationCfg( - dt=dt, - physics=newton_cfg, - use_newton_actuators=use_newton_actuators, - ) - with build_simulation_context( - device="cuda:0", - gravity_enabled=True, - add_ground_plane=True, - sim_cfg=sim_cfg, - ) as sim: - sim._app_control_on_stop_handle = None + RESET_ENV: int = 0 + UNCHANGED_ENV: int = 1 + def _build_and_warm(self, *, use_newton_actuators: bool): + sim_cfg = SimulationCfg( + dt=DT, physics=NEWTON_CFG, use_newton_actuators=use_newton_actuators, + ) + ctx = build_simulation_context( + device="cuda:0", gravity_enabled=True, add_ground_plane=True, sim_cfg=sim_cfg, + ) + sim = ctx.__enter__() + sim._app_control_on_stop_handle = None for i in range(NUM_ENVS): - sim_utils.create_prim( - f"/World/Env_{i}", "Xform", translation=(i * 3.0, 0, 0) - ) - + sim_utils.create_prim(f"/World/Env_{i}", "Xform", translation=(i * 3.0, 0, 0)) art_cfg = ANYMAL_C_CFG.replace( - actuators=actuators, - prim_path="/World/Env_.*/Robot", + actuators=DELAYED_PD_ACTUATORS, prim_path="/World/Env_.*/Robot", ) articulation = Articulation(art_cfg) sim.reset() - assert articulation.is_initialized init_pos = wp.to_torch(articulation.data.joint_pos).clone() target_pos = init_pos + TARGET_OFFSET target_vel = torch.zeros_like(init_pos) - articulation.set_joint_position_target_index(target=target_pos) articulation.set_joint_velocity_target_index(target=target_vel) - - recorded_pos, recorded_vel = [], [] - - for step_i in range(RESET_TOTAL_STEPS): - if step_i == RESET_WARMUP_STEPS: - env_ids = torch.tensor([0], device="cuda:0") - articulation.write_joint_position_to_sim_index( - position=init_pos[0:1], env_ids=env_ids, - ) - articulation.write_joint_velocity_to_sim_index( - velocity=torch.zeros_like(init_pos[0:1]), env_ids=env_ids, - ) - articulation.reset(env_ids=[0]) - + for _ in range(RESET_WARMUP_STEPS): articulation.write_data_to_sim() sim.step() - articulation.update(dt) - - recorded_pos.append(wp.to_torch(articulation.data.joint_pos).clone()) - recorded_vel.append(wp.to_torch(articulation.data.joint_vel).clone()) - - return {"joint_pos": recorded_pos, "joint_vel": recorded_vel} - - -class TestPartialResetEquivalence(unittest.TestCase): - """Per-environment reset with DelayedPD actuators: Lab vs Newton. + articulation.update(DT) + return ctx, sim, articulation + + def test_newton_state_reset_isolated_to_reset_env(self): + """Newton: ``num_pushes`` zeroes for env 0's DOFs only after reset of [0].""" + ctx, sim, articulation = self._build_and_warm(use_newton_actuators=True) + try: + adapter = SimulationManager._adapter + self.assertIsNotNone(adapter) + # Find a DelayedPD actuator (it's the only one with delay_state). + stateful_pairs = [ + (act, st) for act, st in zip(adapter.actuators, adapter._states_a) + if st is not None and getattr(st, "delay_state", None) is not None + ] + self.assertGreater(len(stateful_pairs), 0, "expected at least one DelayedPD actuator with delay_state") + + # Per-DOF entry layout inside each actuator's state: ``act.indices`` + # is the flat global DOF id; envs are stacked so env 0's DOFs come first. + for act, state in stateful_pairs: + pushes_before = state.delay_state.num_pushes.numpy() + self.assertTrue( + (pushes_before > 0).all(), + "expected non-zero num_pushes for all DOFs after warmup", + ) - Resets env 0 mid-simulation while env 1 continues uninterrupted. - Uses DelayedPD actuators because they carry internal state (delay - buffers) that must be properly reset per environment. + articulation.reset(env_ids=torch.tensor([self.RESET_ENV], device=articulation.device, dtype=torch.long)) + + # Map each entry of ``act.indices`` to its env via the adapter's full + # per-env DOF count (model.joint_dof_count // num_envs — includes free + # joint DOFs on floating-base articulations, unlike articulation.num_joints + # which counts only actuated DOFs). + for act, state in stateful_pairs: + pushes_after = state.delay_state.num_pushes.numpy() + indices_np = act.indices.numpy() + for i, global_dof in enumerate(indices_np): + env = int(global_dof) // adapter.num_joints + if env == self.RESET_ENV: + self.assertEqual( + int(pushes_after[i]), 0, + f"DOF {i} (env {env}) should be reset to 0, got {pushes_after[i]}", + ) + else: + self.assertGreater( + int(pushes_after[i]), 0, + f"DOF {i} (env {env}) was NOT in reset env_ids but num_pushes is 0", + ) + finally: + ctx.__exit__(None, None, None) + + def test_lab_state_reset_isolated_to_reset_env(self): + """Lab: DelayedPDActuator circular buffer zeroed for env 0 only.""" + ctx, sim, articulation = self._build_and_warm(use_newton_actuators=False) + try: + from isaaclab.actuators import DelayedPDActuator # noqa: PLC0415 + + delayed = [a for a in articulation.actuators.values() if isinstance(a, DelayedPDActuator)] + self.assertGreater(len(delayed), 0, "expected at least one Lab DelayedPDActuator") + actuator = delayed[0] + buf = actuator.positions_delay_buffer._circular_buffer._buffer + # ``_buffer`` shape: (max_length, batch_size, num_joints). + self.assertIsNotNone(buf, "delay buffer should be populated after warmup") + self.assertTrue( + (buf[:, self.UNCHANGED_ENV] != 0).any().item(), + "expected non-zero buffer entries for env 1 after warmup", + ) - Verifies: - - Lab and Newton paths produce matching trajectories after partial reset. - - The two environments diverge after the reset (proving it took effect). - """ + articulation.reset(env_ids=torch.tensor([self.RESET_ENV], device=articulation.device, dtype=torch.long)) - @classmethod - def setUpClass(cls): - cls.lab_result = _run_simulation_with_reset( - DELAYED_PD_ACTUATORS, use_newton_actuators=False, - ) - cls.newton_result = _run_simulation_with_reset( - DELAYED_PD_ACTUATORS, use_newton_actuators=True, - ) - - def test_joint_positions_match(self): - for step_i, (lab, newton) in enumerate( - zip(self.lab_result["joint_pos"], self.newton_result["joint_pos"]) - ): - torch.testing.assert_close( - lab, - newton, - atol=2e-3, - rtol=1e-3, - msg=f"Positions diverged at step {step_i}", + self.assertTrue( + torch.all(buf[:, self.RESET_ENV] == 0).item(), + f"Lab: env {self.RESET_ENV} buffer not zeroed after reset.", ) - - def test_joint_velocities_match(self): - for step_i, (lab, newton) in enumerate( - zip(self.lab_result["joint_vel"], self.newton_result["joint_vel"]) - ): - torch.testing.assert_close( - lab, - newton, - atol=0.1, - rtol=1e-2, - msg=f"Velocities diverged at step {step_i}", + self.assertTrue( + (buf[:, self.UNCHANGED_ENV] != 0).any().item(), + f"Lab: env {self.UNCHANGED_ENV} buffer was zeroed — reset leaked into an unselected env.", ) - - def test_envs_diverge_after_reset(self): - """After resetting env 0, the two envs must have different states.""" - post_reset_pos = self.lab_result["joint_pos"][RESET_WARMUP_STEPS + 1] - diff = (post_reset_pos[0] - post_reset_pos[1]).abs().max().item() - self.assertGreater( - diff, 0.001, - "Env 0 and env 1 are identical after partial reset — reset had no effect", - ) - - def test_trajectories_not_trivial(self): - first = self.lab_result["joint_pos"][0] - last = self.lab_result["joint_pos"][-1] - diff = (last - first).abs().max().item() - self.assertGreater(diff, 0.01, "Joints did not move — test is trivial") + finally: + ctx.__exit__(None, None, None) # --------------------------------------------------------------------------- @@ -1292,11 +1112,37 @@ def test_kfe_has_delay(self): for a in kfe_acts: self.assertTrue(a["has_delay"], "Delay not found on remotized KFE actuator") - def test_trajectories_not_trivial(self): - first = self.result["joint_pos"][0] - last = self.result["joint_pos"][-1] - diff = (last - first).abs().max().item() - self.assertGreater(diff, 0.01, "Joints did not move — test is trivial") + +class TestRemotizedPDEquivalence(_EquivalenceTestBase): + """RemotizedPD (PD + delay + position-based clamping): Lab vs Newton.""" + + __test__ = True + + @classmethod + def setUpClass(cls): + from isaaclab.actuators.actuator_pd_cfg import RemotizedPDActuatorCfg # noqa: PLC0415 + + cls.actuators = { + "hips": IdealPDActuatorCfg( + joint_names_expr=[".*HAA", ".*HFE"], + stiffness=40.0, + damping=5.0, + effort_limit=80.0, + ), + "knees": RemotizedPDActuatorCfg( + joint_names_expr=[".*KFE"], + stiffness=60.0, + damping=1.5, + effort_limit=80.0, + max_delay=3, + joint_parameter_lookup=SPOT_KNEE_LOOKUP, + ), + } + super().setUpClass() + + +class TestDecimationRemotizedPD(_DecimationMixin, TestRemotizedPDEquivalence): + """RemotizedPD — decimation=2 + CUDA graph.""" # --------------------------------------------------------------------------- @@ -1318,7 +1164,8 @@ def _make_dummy_mlp_checkpoint(device: str = "cpu") -> str: ).to(device).eval() scripted = torch.jit.script(net) - tmp = tempfile.NamedTemporaryFile(suffix=".pt", delete=False) + with tempfile.NamedTemporaryFile(suffix=".pt", delete=False) as tmp: + tmp_path = tmp.name extra = { "metadata.json": json.dumps({ "model_type": "mlp", @@ -1329,8 +1176,8 @@ def _make_dummy_mlp_checkpoint(device: str = "cpu") -> str: "torque_scale": 2.0, }) } - torch.jit.save(scripted, tmp.name, _extra_files=extra) - return tmp.name + torch.jit.save(scripted, tmp_path, _extra_files=extra) + return tmp_path class _DummyLSTM(torch.nn.Module): @@ -1356,10 +1203,11 @@ def _make_dummy_lstm_checkpoint(device: str = "cpu") -> str: net = _DummyLSTM().to(device).eval() scripted = torch.jit.script(net) - tmp = tempfile.NamedTemporaryFile(suffix=".pt", delete=False) + with tempfile.NamedTemporaryFile(suffix=".pt", delete=False) as tmp: + tmp_path = tmp.name extra = {"metadata.json": json.dumps({"model_type": "lstm"})} - torch.jit.save(scripted, tmp.name, _extra_files=extra) - return tmp.name + torch.jit.save(scripted, tmp_path, _extra_files=extra) + return tmp_path class TestNeuralMLPAuthoring(unittest.TestCase): @@ -1415,12 +1263,6 @@ def test_mlp_has_dc_motor_clamping(self): for a in mlp_acts: self.assertIn("ClampingDCMotor", a["clamping_types"]) - def test_trajectories_not_trivial(self): - first = self.result["joint_pos"][0] - last = self.result["joint_pos"][-1] - diff = (last - first).abs().max().item() - self.assertGreater(diff, 0.01, "Joints did not move — test is trivial") - class TestNeuralLSTMAuthoring(unittest.TestCase): """Verify ActuatorNetLSTMCfg is authored as Newton NeuralLSTM controller @@ -1470,12 +1312,6 @@ def test_lstm_has_dc_motor_clamping(self): for a in lstm_acts: self.assertIn("ClampingDCMotor", a["clamping_types"]) - def test_trajectories_not_trivial(self): - first = self.result["joint_pos"][0] - last = self.result["joint_pos"][-1] - diff = (last - first).abs().max().item() - self.assertGreater(diff, 0.01, "Joints did not move — test is trivial") - if __name__ == "__main__": unittest.main() diff --git a/source/isaaclab_physx/isaaclab_physx/assets/articulation/articulation.py b/source/isaaclab_physx/isaaclab_physx/assets/articulation/articulation.py index a83a83aaa11f..f55eb78a6da5 100644 --- a/source/isaaclab_physx/isaaclab_physx/assets/articulation/articulation.py +++ b/source/isaaclab_physx/isaaclab_physx/assets/articulation/articulation.py @@ -28,6 +28,7 @@ NewtonActuatorAdapter, PhysxActuatorWrapper, build_implicit_dof_mask, + build_newton_actuator_defaults, ) from isaaclab_newton.actuators import kernels as actuator_kernels @@ -36,10 +37,6 @@ _HAS_NEWTON_ACTUATORS = False -@wp.kernel(enable_backward=False) -def _build_env_mask_kernel(mask: wp.array(dtype=wp.bool), indices: wp.array(dtype=wp.int32)): - i = wp.tid() - mask[indices[i]] = True from isaaclab.sim.utils.queries import find_first_matching_prim, get_all_matching_child_prims from isaaclab.utils.string import resolve_matching_names, resolve_matching_names_values from isaaclab.utils.types import ArticulationActions @@ -243,6 +240,10 @@ def reset(self, env_ids: Sequence[int] | None = None, env_mask: wp.array | None # reset actuators (including Newton-native adapter which owns its states) for actuator in self.actuators.values(): actuator.reset(env_ids) + # Reset Newton-actuator per-env states (delay queues, neural hidden state, etc.). + # The adapter is per-articulation on PhysX and is not part of ``self.actuators``. + if self._has_newton_actuators and self.newton_actuator_adapter is not None: + self.newton_actuator_adapter.reset(env_ids) # reset external wrenches. self._instantaneous_wrench_composer.reset(env_ids, env_mask) self._permanent_wrench_composer.reset(env_ids, env_mask) @@ -1206,50 +1207,85 @@ def write_joint_damping_to_sim_index( cpu_env_ids = self._get_cpu_env_ids(env_ids) self.root_view.set_dof_dampings(wp.clone(self.data._joint_damping, device="cpu"), indices=cpu_env_ids) - """ - Operations - Newton Actuator Parameter Writers. - - Mirror of the writers on :class:`isaaclab_newton.assets.Articulation`. - Required so that domain randomization terms (``randomize_actuator_gains``) - propagate kp/kd updates to explicit Newton actuators on the PhysX backend. - """ - def write_actuator_stiffness_to_sim( self, *, - stiffness: torch.Tensor | wp.array | float, - env_ids: Sequence[int] | torch.Tensor | wp.array | None = None, + stiffness: torch.Tensor, + env_ids: torch.Tensor, + joint_ids: torch.Tensor, ) -> None: - """Write actuator stiffness (``kp``) into this articulation's Newton controllers. + """Write actuator kp at the (env_ids, joint_ids) sub-grid and propagate to controllers. + + Iterates the per-articulation adapter's Newton actuators and uses + :data:`patch_actuator_param_kernel` to overwrite each + controller's ``kp`` array at the ``(env_ids × joint_ids)`` + cells. DOFs not owned by an actuator are skipped by the kernel's + per-slot index mapping. + + Args: + stiffness: Sub-grid of new kp values, shape ``(len(env_ids), len(joint_ids))``. + env_ids: 1D torch tensor of env indices. + joint_ids: 1D torch tensor of articulation-local joint indices. No-op when no Newton actuators are registered for this articulation. """ - adapter = self.actuators.get("newton") - if adapter is None: - return - env_ids = self._resolve_env_ids(env_ids) - env_mask = self._env_ids_to_mask(env_ids) - adapter.write_stiffness_to_sim(stiffness, env_ids, env_mask) + self._write_actuator_param("kp", stiffness, env_ids, joint_ids) def write_actuator_damping_to_sim( self, *, - damping: torch.Tensor | wp.array | float, - env_ids: Sequence[int] | torch.Tensor | wp.array | None = None, + damping: torch.Tensor, + env_ids: torch.Tensor, + joint_ids: torch.Tensor, ) -> None: - """Write actuator damping (``kd``) into this articulation's Newton controllers.""" - adapter = self.actuators.get("newton") + """Write actuator kd at the (env_ids, joint_ids) sub-grid and propagate to controllers.""" + self._write_actuator_param("kd", damping, env_ids, joint_ids) + + def _write_actuator_param( + self, + attr: str, + values: torch.Tensor, + env_ids: torch.Tensor, + joint_ids: torch.Tensor, + ) -> None: + """Shared body for :meth:`write_actuator_stiffness_to_sim` / :meth:`write_actuator_damping_to_sim`.""" + adapter = self.newton_actuator_adapter if adapter is None: return - env_ids = self._resolve_env_ids(env_ids) - env_mask = self._env_ids_to_mask(env_ids) - adapter.write_damping_to_sim(damping, env_ids, env_mask) - def _env_ids_to_mask(self, env_ids: wp.array) -> wp.array: - """Convert warp ``env_ids`` to a boolean Warp mask of length ``num_instances``.""" - mask = wp.zeros(self.num_instances, dtype=wp.bool, device=self.device) - wp.launch(_build_env_mask_kernel, dim=env_ids.shape[0], inputs=[mask, env_ids], device=self.device) - return mask + env_id_pos = torch.full( + (self.num_instances,), -1, dtype=torch.int32, device=self.device, + ) + env_id_pos[env_ids.to(self.device, dtype=torch.long)] = torch.arange( + env_ids.shape[0], dtype=torch.int32, device=self.device, + ) + joint_id_pos = torch.full( + (self.num_joints,), -1, dtype=torch.int32, device=self.device, + ) + joint_id_pos[joint_ids.to(self.device, dtype=torch.long)] = torch.arange( + joint_ids.shape[0], dtype=torch.int32, device=self.device, + ) + + values_wp = wp.from_torch( + values.to(self.device, dtype=torch.float32).contiguous(), dtype=wp.float32, + ) + env_id_pos_wp = wp.from_torch(env_id_pos, dtype=wp.int32) + joint_id_pos_wp = wp.from_torch(joint_id_pos, dtype=wp.int32) + + for act in adapter.actuators: + ctrl = act.controller + if not hasattr(ctrl, attr): + continue + wp.launch( + actuator_kernels.patch_actuator_param_kernel, + dim=act.indices.shape[0], + inputs=[ + act.indices, env_id_pos_wp, joint_id_pos_wp, values_wp, + 0, self.num_joints, + ], + outputs=[getattr(ctrl, attr)], + device=self.device, + ) def write_joint_damping_to_sim_mask( self, @@ -3808,13 +3844,21 @@ def _process_actuators_cfg(self): # create actuators self.actuators = dict() self._physx_actuator_wrapper = None + # Per-articulation Newton actuator adapter and the frozen kp/kd + # snapshot consumed by ``randomize_actuator_gains``. ``None`` for + # articulations not on the Newton fast path or with only implicit + # Lab actuators. + self.newton_actuator_adapter: NewtonActuatorAdapter | None = None + self.newton_default_stiffness: torch.Tensor | None = None + self.newton_default_damping: torch.Tensor | None = None + self.newton_managed_local_joints: torch.Tensor | slice | None = None # flag for implicit actuators # if this is false, we by-pass certain checks when doing actuator-related operations self._has_implicit_actuators = False self._has_newton_actuators = False # Per-DOF implicit/explicit mask consumed by the - # ``synch_torque_and_apply_implicit_feedforwards`` kernel. ``None`` - # when no Newton fast path is active. + # ``sync_torque_telemetry`` kernel. ``None`` when no Newton fast path + # is active. self._implicit_dof_mask: wp.array | None = None _use_newton_actuators = getattr(self._sim_cfg, "use_newton_actuators", False) @@ -3868,9 +3912,19 @@ def _process_actuators_cfg(self): w.joint_target_pos = self._data.joint_pos_target.warp.reshape(-1) w.joint_target_vel = self._data.joint_vel_target.warp.reshape(-1) w.joint_act = self._data.joint_effort_target.warp.reshape(-1) - adapter.finalize() - adapter.set_kernel_propagator() - self.actuators["newton"] = adapter + adapter.finalize(w) + self.newton_actuator_adapter = adapter + ( + self.newton_default_stiffness, + self.newton_default_damping, + self.newton_managed_local_joints, + ) = build_newton_actuator_defaults( + actuators=adapter.actuators, + num_envs=self.num_instances, + num_joints=self.num_joints, + dof_offset=0, + device=self.device, + ) self.write_joint_stiffness_to_sim_index(stiffness=0.0, joint_ids=adapter.joint_indices) self.write_joint_damping_to_sim_index(damping=0.0, joint_ids=adapter.joint_indices) @@ -3889,6 +3943,16 @@ def _process_actuators_cfg(self): self._implicit_dof_mask = build_implicit_dof_mask( self.actuators, self.num_joints, self.device, ) + # Per-articulation view of the adapter's pre-clamp computed-effort + # buffer (or zero fallback when there are no explicit Newton + # actuators). Set once here so ``_apply_actuator_model_newton`` + # passes it straight to ``sync_torque_telemetry``. + if self.newton_actuator_adapter is not None: + self._data._sim_bind_joint_computed_effort = self.newton_actuator_adapter.computed_effort_2d + else: + self._data._sim_bind_joint_computed_effort = wp.zeros( + (self.num_instances, self.num_joints), dtype=wp.float32, device=self.device, + ) return # --- Standard Isaac Lab actuator path --- @@ -4103,38 +4167,38 @@ def _apply_actuator_model(self): ) def _apply_actuator_model_newton(self): - """Step Newton actuators (when present) then route FF + sync telemetry. - - After ``newton_adapter.step`` (no-op if no explicit Newton actuators - exist), ``w.joint_f_2d`` holds the actuator output for explicit DOFs - and zero elsewhere. The - :func:`synch_torque_and_apply_implicit_feedforwards` kernel then - writes the user FF into ``joint_f_2d`` for implicit DOFs and fills - ``_data._computed_torque`` / ``_data._applied_torque``. The - resulting ``joint_f_2d`` is what gets pushed to PhysX as the - actuation force in :meth:`write_data_to_sim`. + """Pre-fill effort buffer with FF, step Newton actuators, sync telemetry. + + Pre-fills ``w.joint_f_2d`` with the user's effort target across all + DOFs. ``newton_adapter.step`` (no-op if no explicit Newton actuators + exist) then zeroes ``joint_f_2d`` at explicit DOFs and overwrites + them with each actuator's computed effort, while implicit DOFs keep + the FF. The :func:`sync_torque_telemetry` kernel then fills + ``_data._computed_torque`` / ``_data._applied_torque`` from the + resulting buffer. The final ``joint_f_2d`` is what gets pushed to + PhysX as the actuation force in :meth:`write_data_to_sim`. """ w = self._physx_actuator_wrapper - newton_adapter = self.actuators.get("newton") - if newton_adapter is not None: - newton_adapter.step(w, w, SimulationManager.get_physics_dt()) + w.joint_f_2d.assign(self._data._joint_effort_target) + if self.newton_actuator_adapter is not None: + self.newton_actuator_adapter.step(w, w, SimulationManager.get_physics_dt()) wp.launch( - actuator_kernels.synch_torque_and_apply_implicit_feedforwards, + actuator_kernels.sync_torque_telemetry, dim=(self.num_instances, self.num_joints), inputs=[ self._data.joint_pos.warp, self._data.joint_vel.warp, self._data._joint_pos_target, self._data._joint_vel_target, - self._data._joint_effort_target, self._data.joint_stiffness.warp, self._data.joint_damping.warp, self._data.joint_effort_limits.warp, self._implicit_dof_mask, + w.joint_f_2d, + self._data._sim_bind_joint_computed_effort, ], outputs=[ - w.joint_f_2d, self._data._computed_torque, self._data._applied_torque, ], diff --git a/source/isaaclab_physx/test/assets/test_newton_actuators_physx.py b/source/isaaclab_physx/test/assets/test_newton_actuators_physx.py index 29727b7a195d..e91e74de52dc 100644 --- a/source/isaaclab_physx/test/assets/test_newton_actuators_physx.py +++ b/source/isaaclab_physx/test/assets/test_newton_actuators_physx.py @@ -25,7 +25,6 @@ import torch import warp as wp -from isaaclab_assets import ANYMAL_C_CFG from isaaclab_physx.assets import Articulation from isaaclab_physx.physics import PhysxCfg @@ -33,6 +32,8 @@ from isaaclab.actuators import DCMotorCfg, DelayedPDActuatorCfg, IdealPDActuatorCfg, ImplicitActuatorCfg from isaaclab.sim import SimulationCfg, build_simulation_context +from isaaclab_assets import ANYMAL_C_CFG + # --------------------------------------------------------------------------- # Constants # --------------------------------------------------------------------------- @@ -126,42 +127,22 @@ def _run_simulation( use_newton_actuators: bool, *, num_steps: int = NUM_STEPS, + feedforward: float | None = None, ) -> dict: - """Run ANYmal-C on PhysX and return recorded joint trajectories. + """Run ANYmal-C on PhysX and return recorded trajectories + telemetry. - Args: - actuators: Actuator config dict overriding ANYmal's defaults. - use_newton_actuators: Use Newton-native actuators (via - :class:`PhysxActuatorWrapper`) when ``True``. - num_steps: Number of simulation steps. - - Returns: - Dict with ``joint_pos`` and ``joint_vel``, each a list of - ``(NUM_ENVS, num_joints)`` tensors. + Always records ``joint_pos``, ``joint_vel``, ``computed_torque``, and + ``applied_torque``. Optionally applies a constant per-DOF feedforward + effort target. """ - sim_cfg = SimulationCfg( - dt=DT, - physics=PhysxCfg(), - use_newton_actuators=use_newton_actuators, - ) - + sim_cfg = SimulationCfg(dt=DT, physics=PhysxCfg(), use_newton_actuators=use_newton_actuators) with build_simulation_context( - device="cuda:0", - gravity_enabled=True, - add_ground_plane=True, - sim_cfg=sim_cfg, + device="cuda:0", gravity_enabled=True, add_ground_plane=True, sim_cfg=sim_cfg, ) as sim: sim._app_control_on_stop_handle = None - for i in range(NUM_ENVS): - sim_utils.create_prim( - f"/World/Env_{i}", "Xform", translation=(i * 3.0, 0, 0) - ) - - art_cfg = ANYMAL_C_CFG.replace( - actuators=actuators, - prim_path="/World/Env_.*/Robot", - ) + sim_utils.create_prim(f"/World/Env_{i}", "Xform", translation=(i * 3.0, 0, 0)) + art_cfg = ANYMAL_C_CFG.replace(actuators=actuators, prim_path="/World/Env_.*/Robot") articulation = Articulation(art_cfg) sim.reset() assert articulation.is_initialized @@ -169,20 +150,32 @@ def _run_simulation( init_pos = wp.to_torch(articulation.data.joint_pos).clone() target_pos = init_pos + TARGET_OFFSET target_vel = torch.zeros_like(init_pos) - articulation.set_joint_position_target_index(target=target_pos) articulation.set_joint_velocity_target_index(target=target_vel) + if feedforward is not None: + articulation.set_joint_effort_target_index( + target=torch.full_like(init_pos, feedforward), + ) recorded_pos, recorded_vel = [], [] + recorded_computed, recorded_applied = [], [] for _ in range(num_steps): articulation.write_data_to_sim() sim.step() articulation.update(DT) - recorded_pos.append(wp.to_torch(articulation.data.joint_pos).clone()) recorded_vel.append(wp.to_torch(articulation.data.joint_vel).clone()) + recorded_computed.append(wp.to_torch(articulation.data.computed_torque).clone()) + recorded_applied.append(wp.to_torch(articulation.data.applied_torque).clone()) - return {"joint_pos": recorded_pos, "joint_vel": recorded_vel} + return { + "joint_pos": recorded_pos, + "joint_vel": recorded_vel, + "computed_torque": recorded_computed, + "applied_torque": recorded_applied, + "target_pos": target_pos.clone(), + "target_vel": target_vel.clone(), + } # --------------------------------------------------------------------------- @@ -200,25 +193,29 @@ class _EquivalenceTestBase(unittest.TestCase): __test__ = False actuators: dict = {} + feedforward: float | None = None pos_atol: float = 2e-3 pos_rtol: float = 1e-3 - vel_atol: float = 0.1 + vel_atol: float = 1e-2 vel_rtol: float = 1e-2 + torque_atol: float = 1e-3 + torque_rtol: float = 1e-3 @classmethod def setUpClass(cls): - cls.lab_result = _run_simulation(cls.actuators, use_newton_actuators=False) - cls.newton_result = _run_simulation(cls.actuators, use_newton_actuators=True) + cls.lab_result = _run_simulation( + cls.actuators, use_newton_actuators=False, feedforward=cls.feedforward, + ) + cls.newton_result = _run_simulation( + cls.actuators, use_newton_actuators=True, feedforward=cls.feedforward, + ) def test_joint_positions_match(self): for step_i, (lab, newton) in enumerate( zip(self.lab_result["joint_pos"], self.newton_result["joint_pos"]) ): torch.testing.assert_close( - lab, - newton, - atol=self.pos_atol, - rtol=self.pos_rtol, + lab, newton, atol=self.pos_atol, rtol=self.pos_rtol, msg=f"Joint positions diverged at step {step_i}", ) @@ -227,18 +224,27 @@ def test_joint_velocities_match(self): zip(self.lab_result["joint_vel"], self.newton_result["joint_vel"]) ): torch.testing.assert_close( - lab, - newton, - atol=self.vel_atol, - rtol=self.vel_rtol, + lab, newton, atol=self.vel_atol, rtol=self.vel_rtol, msg=f"Joint velocities diverged at step {step_i}", ) - def test_trajectories_not_trivial(self): - first = self.lab_result["joint_pos"][0] - last = self.lab_result["joint_pos"][-1] - diff = (last - first).abs().max().item() - self.assertGreater(diff, 0.01, "Joints did not move — test is trivial") + def test_applied_torque_match(self): + for step_i, (lab, newton) in enumerate( + zip(self.lab_result["applied_torque"], self.newton_result["applied_torque"]) + ): + torch.testing.assert_close( + lab, newton, atol=self.torque_atol, rtol=self.torque_rtol, + msg=f"applied_torque diverged at step {step_i}", + ) + + def test_computed_torque_match(self): + for step_i, (lab, newton) in enumerate( + zip(self.lab_result["computed_torque"], self.newton_result["computed_torque"]) + ): + torch.testing.assert_close( + lab, newton, atol=self.torque_atol, rtol=self.torque_rtol, + msg=f"computed_torque diverged at step {step_i}", + ) # --------------------------------------------------------------------------- @@ -309,222 +315,12 @@ class TestImplicitOnlyEquivalencePhysx(_EquivalenceTestBase): actuators = IMPLICIT_ONLY_ACTUATORS -def _run_simulation_with_telemetry( - actuators: dict, - use_newton_actuators: bool, - *, - num_steps: int = NUM_STEPS, -) -> dict: - """Like :func:`_run_simulation` but also records ``computed_torque`` / ``applied_torque``.""" - sim_cfg = SimulationCfg( - dt=DT, - physics=PhysxCfg(), - use_newton_actuators=use_newton_actuators, - ) - - with build_simulation_context( - device="cuda:0", - gravity_enabled=True, - add_ground_plane=True, - sim_cfg=sim_cfg, - ) as sim: - sim._app_control_on_stop_handle = None - - for i in range(NUM_ENVS): - sim_utils.create_prim( - f"/World/Env_{i}", "Xform", translation=(i * 3.0, 0, 0) - ) - - art_cfg = ANYMAL_C_CFG.replace( - actuators=actuators, - prim_path="/World/Env_.*/Robot", - ) - articulation = Articulation(art_cfg) - sim.reset() - assert articulation.is_initialized - - init_pos = wp.to_torch(articulation.data.joint_pos).clone() - target_pos = init_pos + TARGET_OFFSET - target_vel = torch.zeros_like(init_pos) - - articulation.set_joint_position_target_index(target=target_pos) - articulation.set_joint_velocity_target_index(target=target_vel) - - recorded_pos, recorded_vel = [], [] - recorded_computed, recorded_applied = [], [] - for _ in range(num_steps): - articulation.write_data_to_sim() - sim.step() - articulation.update(DT) - - recorded_pos.append(wp.to_torch(articulation.data.joint_pos).clone()) - recorded_vel.append(wp.to_torch(articulation.data.joint_vel).clone()) - recorded_computed.append(wp.to_torch(articulation.data.computed_torque).clone()) - recorded_applied.append(wp.to_torch(articulation.data.applied_torque).clone()) - - return { - "joint_pos": recorded_pos, - "joint_vel": recorded_vel, - "computed_torque": recorded_computed, - "applied_torque": recorded_applied, - "target_pos": target_pos.clone(), - "target_vel": target_vel.clone(), - } - - -class TestImplicitOnlyTelemetryPhysx(unittest.TestCase): - """Implicit-only fast path on PhysX: shadow-PD telemetry matches the Lab actuator path.""" - - @classmethod - def setUpClass(cls): - cls.lab_result = _run_simulation_with_telemetry( - IMPLICIT_ONLY_ACTUATORS, use_newton_actuators=False, - ) - cls.newton_result = _run_simulation_with_telemetry( - IMPLICIT_ONLY_ACTUATORS, use_newton_actuators=True, - ) - - def test_telemetry_is_nonzero(self): - last = self.newton_result["computed_torque"][-1] - self.assertGreater( - last.abs().max().item(), - 1e-3, - "computed_torque is all-zero — implicit telemetry kernel did not run", - ) - - def test_telemetry_matches_lab_path(self): - """Newton fast-path telemetry agrees with the Lab actuator-path telemetry.""" - for step_i, (lab_comp, newton_comp) in enumerate( - zip(self.lab_result["computed_torque"], self.newton_result["computed_torque"]) - ): - torch.testing.assert_close( - newton_comp, - lab_comp, - atol=5e-2, - rtol=1e-2, - msg=f"computed_torque diverged from Lab path at step {step_i}", - ) - - def test_applied_equals_computed_when_no_clip(self): - for step_i, (comp, app) in enumerate( - zip(self.newton_result["computed_torque"], self.newton_result["applied_torque"]) - ): - torch.testing.assert_close( - app, - comp, - atol=1e-5, - rtol=1e-5, - msg=f"applied_torque != computed_torque at step {step_i} (no clip expected)", - ) - - -# --------------------------------------------------------------------------- -# Implicit + non-zero feedforward effort target -# --------------------------------------------------------------------------- - - -def _run_simulation_with_ff( - actuators: dict, - use_newton_actuators: bool, - *, - feedforward: float, - num_steps: int = NUM_STEPS, -) -> dict: - """Run with both a position target and a non-zero feedforward effort target.""" - sim_cfg = SimulationCfg( - dt=DT, - physics=PhysxCfg(), - use_newton_actuators=use_newton_actuators, - ) - - with build_simulation_context( - device="cuda:0", - gravity_enabled=True, - add_ground_plane=True, - sim_cfg=sim_cfg, - ) as sim: - sim._app_control_on_stop_handle = None - - for i in range(NUM_ENVS): - sim_utils.create_prim( - f"/World/Env_{i}", "Xform", translation=(i * 3.0, 0, 0) - ) - - art_cfg = ANYMAL_C_CFG.replace( - actuators=actuators, - prim_path="/World/Env_.*/Robot", - ) - articulation = Articulation(art_cfg) - sim.reset() - assert articulation.is_initialized - - init_pos = wp.to_torch(articulation.data.joint_pos).clone() - target_pos = init_pos + TARGET_OFFSET - target_vel = torch.zeros_like(init_pos) - target_eff = torch.full_like(init_pos, feedforward) - - articulation.set_joint_position_target_index(target=target_pos) - articulation.set_joint_velocity_target_index(target=target_vel) - articulation.set_joint_effort_target_index(target=target_eff) - - recorded_pos, recorded_vel = [], [] - for _ in range(num_steps): - articulation.write_data_to_sim() - sim.step() - articulation.update(DT) - recorded_pos.append(wp.to_torch(articulation.data.joint_pos).clone()) - recorded_vel.append(wp.to_torch(articulation.data.joint_vel).clone()) - - return {"joint_pos": recorded_pos, "joint_vel": recorded_vel} - - -class TestImplicitWithFeedforwardEquivalencePhysx(unittest.TestCase): - """Implicit-only actuators with a non-zero feedforward effort target on PhysX. - - For implicit actuators, the user's feedforward effort must be additive - on top of the simulator's PD on both the Lab path and the Newton fast - path. This test commands a constant FF effort and compares joint - trajectories between the two paths. - """ - - FEEDFORWARD = 5.0 - pos_atol = 2e-3 - pos_rtol = 1e-3 - vel_atol = 0.1 - vel_rtol = 1e-2 - - @classmethod - def setUpClass(cls): - cls.lab_result = _run_simulation_with_ff( - IMPLICIT_ONLY_ACTUATORS, use_newton_actuators=False, feedforward=cls.FEEDFORWARD, - ) - cls.newton_result = _run_simulation_with_ff( - IMPLICIT_ONLY_ACTUATORS, use_newton_actuators=True, feedforward=cls.FEEDFORWARD, - ) - - def test_joint_positions_match(self): - for step_i, (lab, newton) in enumerate( - zip(self.lab_result["joint_pos"], self.newton_result["joint_pos"]) - ): - torch.testing.assert_close( - lab, newton, atol=self.pos_atol, rtol=self.pos_rtol, - msg=f"Joint positions diverged at step {step_i}", - ) - - def test_joint_velocities_match(self): - for step_i, (lab, newton) in enumerate( - zip(self.lab_result["joint_vel"], self.newton_result["joint_vel"]) - ): - torch.testing.assert_close( - lab, newton, atol=self.vel_atol, rtol=self.vel_rtol, - msg=f"Joint velocities diverged at step {step_i}", - ) +class TestImplicitWithFeedforwardEquivalencePhysx(_EquivalenceTestBase): + """Implicit-only actuators with a non-zero feedforward effort target on PhysX.""" - def test_trajectories_not_trivial(self): - first = self.lab_result["joint_pos"][0] - last = self.lab_result["joint_pos"][-1] - diff = (last - first).abs().max().item() - self.assertGreater(diff, 0.01, "Joints did not move — test is trivial") + __test__ = True + actuators = IMPLICIT_ONLY_ACTUATORS + feedforward = 5.0 # --------------------------------------------------------------------------- @@ -621,58 +417,87 @@ def test_cartpole_matches_lab(self): ) -class TestExplicitOnlyTelemetryPhysx(unittest.TestCase): - """Explicit-only Newton actuators on PhysX: telemetry copies from the staging effort buffer.""" +# --------------------------------------------------------------------------- +# Domain randomization via events.py — PhysX backend +# --------------------------------------------------------------------------- - @classmethod - def setUpClass(cls): - cls.result = _run_simulation_with_telemetry( - DC_MOTOR_ACTUATORS, use_newton_actuators=True, - ) - def test_telemetry_is_nonzero(self): - last = self.result["applied_torque"][-1] - self.assertGreater( - last.abs().max().item(), - 1e-3, - "applied_torque is all-zero — explicit-DOF telemetry path did not run", - ) +class _MockScene: + """Minimal stand-in for ``InteractiveScene`` accepted by ``ManagerTermBase``.""" - def test_computed_equals_applied_explicit(self): - for step_i, (comp, app) in enumerate( - zip(self.result["computed_torque"], self.result["applied_torque"]) - ): - torch.testing.assert_close( - comp, - app, - atol=1e-5, - rtol=1e-5, - msg=f"computed != applied for explicit-only at step {step_i}", - ) + def __init__(self, assets: dict, num_envs: int): + self._assets = assets + self.num_envs = num_envs + def __getitem__(self, name: str): + return self._assets[name] -class TestWriteActuatorGainsPhysx(unittest.TestCase): - """``write_actuator_*_to_sim`` propagates kp/kd into Newton controllers (PhysX backend). - Catches the silent no-op that ``randomize_actuator_gains`` would suffer - from on PhysX with ``use_newton_actuators=True`` if these writers were - missing (the events.py term falls back to ``hasattr`` and skips - explicit actuators). +class _MockEnv: + """Minimal stand-in for ``ManagerBasedEnv`` for invoking DR terms. + + ``randomize_actuator_gains`` only reads ``env.scene[name]`` and + ``env.scene.num_envs`` (plus ``env.num_envs`` / ``env.device`` from the + ``ManagerTermBase`` properties). No simulator access is needed because + the DR term reaches the actuator adapter via ``self.the actuator adapter``. """ - def test_writers_exist(self): - # Guards against the missing-writer regression. - from isaaclab_physx.assets import Articulation - self.assertTrue( - hasattr(Articulation, "write_actuator_stiffness_to_sim"), - "Articulation is missing write_actuator_stiffness_to_sim — DR will silently no-op", - ) - self.assertTrue( - hasattr(Articulation, "write_actuator_damping_to_sim"), - "Articulation is missing write_actuator_damping_to_sim — DR will silently no-op", - ) + def __init__(self, assets: dict, num_envs: int, device: str): + self.scene = _MockScene(assets, num_envs) + self.num_envs = num_envs + self.device = device + + +def _build_dr_term(env, asset_name, joint_ids=None): + from isaaclab.envs.mdp.events import randomize_actuator_gains # noqa: PLC0415 + from isaaclab.managers import EventTermCfg, SceneEntityCfg # noqa: PLC0415 + + asset_cfg = SceneEntityCfg(asset_name) + if joint_ids is not None: + asset_cfg.joint_ids = joint_ids + cfg = EventTermCfg( + func=randomize_actuator_gains, + params={ + "asset_cfg": asset_cfg, + "stiffness_distribution_params": (100.0, 100.0), + "damping_distribution_params": (5.0, 5.0), + "operation": "abs", + "distribution": "uniform", + }, + ) + return randomize_actuator_gains(cfg, env), asset_cfg + + +class TestRandomizeActuatorGainsViaEventsPhysx(unittest.TestCase): + """End-to-end DR test for the PhysX backend. + + Drives ``randomize_actuator_gains`` (events.py) and verifies the new + kp/kd values land in the per-articulation adapter's buffer at the + right cells — exercising the full path: events → + the actuator adapter → write_stiffness/damping → propagation + to controllers. + + With ``operation="abs"`` and ``distribution="uniform"`` over a + degenerate range ``(K, K)``, every randomized cell is set to exactly + ``K`` — so the assertions are deterministic. + """ - def test_writers_propagate_to_controller(self): + @staticmethod + def _gather_param(adapter, num_envs, num_joints, attr, device): + """Reconstruct a ``(num_envs, num_joints)`` view of ``controller.`` across all actuators.""" + out = torch.zeros((num_envs, num_joints), device=device) + for act in adapter.actuators: + ctrl = act.controller + if not hasattr(ctrl, attr): + continue + flat_t = wp.to_torch(getattr(ctrl, attr)) + idx_np = act.indices.numpy() + envs = torch.from_numpy((idx_np // num_joints).astype("int64")).to(device) + locals_ = torch.from_numpy((idx_np % num_joints).astype("int64")).to(device) + out[envs, locals_] = flat_t + return out + + def test_single_articulation(self): sim_cfg = SimulationCfg(dt=DT, physics=PhysxCfg(), use_newton_actuators=True) with build_simulation_context( device="cuda:0", gravity_enabled=True, add_ground_plane=True, sim_cfg=sim_cfg, @@ -681,191 +506,218 @@ def test_writers_propagate_to_controller(self): for i in range(NUM_ENVS): sim_utils.create_prim(f"/World/Env_{i}", "Xform", translation=(i * 3.0, 0, 0)) art_cfg = ANYMAL_C_CFG.replace( - actuators=DC_MOTOR_ACTUATORS, prim_path="/World/Env_.*/Robot", + actuators=IDEAL_PD_ACTUATORS, prim_path="/World/Env_.*/Robot", ) - articulation = Articulation(art_cfg) + anymal = Articulation(art_cfg) sim.reset() - adapter = articulation.actuators["newton"] - # Snapshot initial controller gains. - kp_before = [ - wp.to_torch(a.controller.kp).clone() for a in adapter.actuators if hasattr(a.controller, "kp") - ] - kd_before = [ - wp.to_torch(a.controller.kd).clone() for a in adapter.actuators if hasattr(a.controller, "kd") - ] - self.assertGreater(len(kp_before), 0, "expected at least one PD controller in adapter") - new_kp = adapter.stiffness.clone() * 2.0 - new_kd = adapter.damping.clone() * 3.0 - articulation.write_actuator_stiffness_to_sim(stiffness=new_kp) - articulation.write_actuator_damping_to_sim(damping=new_kd) - # Verify each controller's kp/kd actually changed (and roughly doubled/tripled). - kp_idx = 0 - kd_idx = 0 - for newton_act in adapter.actuators: - ctrl = newton_act.controller - if hasattr(ctrl, "kp"): - after = wp.to_torch(ctrl.kp) - self.assertFalse( - torch.equal(after, kp_before[kp_idx]), - "controller.kp unchanged after write_actuator_stiffness_to_sim", - ) - torch.testing.assert_close(after, kp_before[kp_idx] * 2.0, atol=1e-4, rtol=1e-4) - kp_idx += 1 - if hasattr(ctrl, "kd"): - after = wp.to_torch(ctrl.kd) - self.assertFalse( - torch.equal(after, kd_before[kd_idx]), - "controller.kd unchanged after write_actuator_damping_to_sim", - ) - torch.testing.assert_close(after, kd_before[kd_idx] * 3.0, atol=1e-4, rtol=1e-4) - kd_idx += 1 + + adapter = anymal.newton_actuator_adapter + self.assertIsNotNone(adapter, "PhysX per-articulation adapter should exist") + n = anymal.num_joints + kp_before = self._gather_param(adapter, NUM_ENVS, n, "kp", anymal.device).clone() + kd_before = self._gather_param(adapter, NUM_ENVS, n, "kd", anymal.device).clone() + + env = _MockEnv({"robot": anymal}, NUM_ENVS, anymal.device) + term, asset_cfg = _build_dr_term(env, "robot") + env_ids = torch.tensor([0], device=anymal.device, dtype=torch.long) + + term( + env, env_ids=env_ids, asset_cfg=asset_cfg, + stiffness_distribution_params=(100.0, 100.0), + damping_distribution_params=(5.0, 5.0), + operation="abs", distribution="uniform", + ) + + kp_after = self._gather_param(adapter, NUM_ENVS, n, "kp", anymal.device) + kd_after = self._gather_param(adapter, NUM_ENVS, n, "kd", anymal.device) + torch.testing.assert_close(kp_after[0], torch.full((n,), 100.0, device=anymal.device)) + torch.testing.assert_close(kd_after[0], torch.full((n,), 5.0, device=anymal.device)) + for env_idx in range(1, NUM_ENVS): + torch.testing.assert_close(kp_after[env_idx], kp_before[env_idx]) + torch.testing.assert_close(kd_after[env_idx], kd_before[env_idx]) + + def test_two_articulations(self): + from isaaclab_assets import CARTPOLE_CFG # noqa: PLC0415 + + sim_cfg = SimulationCfg(dt=DT, physics=PhysxCfg(), use_newton_actuators=True) + with build_simulation_context( + device="cuda:0", gravity_enabled=True, add_ground_plane=True, sim_cfg=sim_cfg, + ) as sim: + sim._app_control_on_stop_handle = None + for i in range(NUM_ENVS): + sim_utils.create_prim(f"/World/Env_{i}", "Xform", translation=(i * 6.0, 0, 0)) + + anymal_cfg = ANYMAL_C_CFG.replace(actuators=IDEAL_PD_ACTUATORS, prim_path="/World/Env_.*/Anymal") + cartpole_cfg = CARTPOLE_CFG.replace( + actuators=CARTPOLE_EXPLICIT_ACTUATORS, prim_path="/World/Env_.*/Cartpole", + ) + cartpole_cfg.init_state = cartpole_cfg.init_state.replace(pos=(0.0, 3.0, 2.0)) + anymal = Articulation(anymal_cfg) + cartpole = Articulation(cartpole_cfg) + sim.reset() + + # On PhysX each articulation owns its own adapter — they are distinct objects. + anymal_adapter = anymal.newton_actuator_adapter + cartpole_adapter = cartpole.newton_actuator_adapter + self.assertIsNotNone(anymal_adapter) + self.assertIsNotNone(cartpole_adapter) + self.assertIsNot(anymal_adapter, cartpole_adapter) + + n_anymal = anymal.num_joints + n_cp = cartpole.num_joints + anymal_kp_before = self._gather_param(anymal_adapter, NUM_ENVS, n_anymal, "kp", anymal.device).clone() + anymal_kd_before = self._gather_param(anymal_adapter, NUM_ENVS, n_anymal, "kd", anymal.device).clone() + cp_kp_before = self._gather_param(cartpole_adapter, NUM_ENVS, n_cp, "kp", anymal.device).clone() + cp_kd_before = self._gather_param(cartpole_adapter, NUM_ENVS, n_cp, "kd", anymal.device).clone() + + env = _MockEnv({"anymal": anymal, "cartpole": cartpole}, NUM_ENVS, anymal.device) + term, asset_cfg = _build_dr_term(env, "cartpole") + env_ids = torch.tensor([0], device=anymal.device, dtype=torch.long) + + term( + env, env_ids=env_ids, asset_cfg=asset_cfg, + stiffness_distribution_params=(100.0, 100.0), + damping_distribution_params=(5.0, 5.0), + operation="abs", distribution="uniform", + ) + + cp_kp_after = self._gather_param(cartpole_adapter, NUM_ENVS, n_cp, "kp", anymal.device) + cp_kd_after = self._gather_param(cartpole_adapter, NUM_ENVS, n_cp, "kd", anymal.device) + torch.testing.assert_close(cp_kp_after[0], torch.full((n_cp,), 100.0, device=anymal.device)) + torch.testing.assert_close(cp_kd_after[0], torch.full((n_cp,), 5.0, device=anymal.device)) + for env_idx in range(1, NUM_ENVS): + torch.testing.assert_close(cp_kp_after[env_idx], cp_kp_before[env_idx]) + torch.testing.assert_close(cp_kd_after[env_idx], cp_kd_before[env_idx]) + + # ANYmal's controllers are fully untouched — DR was scoped to cartpole. + anymal_kp_after = self._gather_param(anymal_adapter, NUM_ENVS, n_anymal, "kp", anymal.device) + anymal_kd_after = self._gather_param(anymal_adapter, NUM_ENVS, n_anymal, "kd", anymal.device) + torch.testing.assert_close(anymal_kp_after, anymal_kp_before) + torch.testing.assert_close(anymal_kd_after, anymal_kd_before) # --------------------------------------------------------------------------- -# Partial environment reset: verify per-env reset equivalence +# Per-env reset: actuator state isolation # --------------------------------------------------------------------------- RESET_WARMUP_STEPS = 3 -RESET_TOTAL_STEPS = 10 -def _run_simulation_with_reset( - actuators: dict, - use_newton_actuators: bool, -) -> dict: - """Run ANYmal-C on PhysX with a mid-simulation reset of env 0 only. - Steps ``RESET_WARMUP_STEPS``, then resets env 0 to its initial joint state - (zeroing velocity), then steps ``RESET_TOTAL_STEPS - RESET_WARMUP_STEPS`` - more. Returns per-step joint positions and velocities. +class TestActuatorStateReset(unittest.TestCase): + """Reset must clear the actuator state buffers for the requested envs only. - This exercises the actuator state reset path (delay buffers, neural - hidden states, etc.) for a subset of environments. + Inspects ``adapter.actuators[i].state.delay_state.num_pushes`` directly: - Args: - actuators: Actuator config dict overriding ANYmal's defaults. - use_newton_actuators: Use Newton-native actuators when ``True``. + * After warmup, ``num_pushes > 0`` for every DOF (buffer was populated). + * After ``articulation.reset(env_ids=[0])``, the entries for env 0's DOFs + must be ``0`` and the entries for env 1's DOFs must remain ``> 0``. - Returns: - Dict with ``joint_pos`` and ``joint_vel``, each a list of - ``(NUM_ENVS, num_joints)`` tensors. + Done independently on Lab and Newton paths. PhysX-side adapter is + per-articulation, available via ``articulation.newton_actuator_adapter``. """ - sim_cfg = SimulationCfg( - dt=DT, - physics=PhysxCfg(), - use_newton_actuators=use_newton_actuators, - ) - with build_simulation_context( - device="cuda:0", - gravity_enabled=True, - add_ground_plane=True, - sim_cfg=sim_cfg, - ) as sim: - sim._app_control_on_stop_handle = None + RESET_ENV: int = 0 + UNCHANGED_ENV: int = 1 + def _build_and_warm(self, *, use_newton_actuators: bool): + sim_cfg = SimulationCfg( + dt=DT, physics=PhysxCfg(), use_newton_actuators=use_newton_actuators, + ) + ctx = build_simulation_context( + device="cuda:0", gravity_enabled=True, add_ground_plane=True, sim_cfg=sim_cfg, + ) + sim = ctx.__enter__() + sim._app_control_on_stop_handle = None for i in range(NUM_ENVS): - sim_utils.create_prim( - f"/World/Env_{i}", "Xform", translation=(i * 3.0, 0, 0) - ) - + sim_utils.create_prim(f"/World/Env_{i}", "Xform", translation=(i * 3.0, 0, 0)) art_cfg = ANYMAL_C_CFG.replace( - actuators=actuators, - prim_path="/World/Env_.*/Robot", + actuators=DELAYED_PD_ACTUATORS, prim_path="/World/Env_.*/Robot", ) articulation = Articulation(art_cfg) sim.reset() - assert articulation.is_initialized init_pos = wp.to_torch(articulation.data.joint_pos).clone() target_pos = init_pos + TARGET_OFFSET target_vel = torch.zeros_like(init_pos) - articulation.set_joint_position_target_index(target=target_pos) articulation.set_joint_velocity_target_index(target=target_vel) - - recorded_pos, recorded_vel = [], [] - - for step_i in range(RESET_TOTAL_STEPS): - if step_i == RESET_WARMUP_STEPS: - env_ids = torch.tensor([0], device="cuda:0") - articulation.write_joint_position_to_sim_index( - position=init_pos[0:1], env_ids=env_ids, - ) - articulation.write_joint_velocity_to_sim_index( - velocity=torch.zeros_like(init_pos[0:1]), env_ids=env_ids, - ) - articulation.reset(env_ids=[0]) - + for _ in range(RESET_WARMUP_STEPS): articulation.write_data_to_sim() sim.step() articulation.update(DT) + return ctx, sim, articulation + + def test_newton_state_reset_isolated_to_reset_env(self): + """Newton: ``num_pushes`` zeroes for env 0's DOFs only after reset of [0].""" + ctx, sim, articulation = self._build_and_warm(use_newton_actuators=True) + try: + adapter = articulation.newton_actuator_adapter + self.assertIsNotNone(adapter) + stateful_pairs = [ + (act, st) for act, st in zip(adapter.actuators, adapter._states_a) + if st is not None and getattr(st, "delay_state", None) is not None + ] + self.assertGreater(len(stateful_pairs), 0, "expected at least one DelayedPD actuator with delay_state") - recorded_pos.append(wp.to_torch(articulation.data.joint_pos).clone()) - recorded_vel.append(wp.to_torch(articulation.data.joint_vel).clone()) - - return {"joint_pos": recorded_pos, "joint_vel": recorded_vel} - - -class TestPartialResetEquivalence(unittest.TestCase): - """Per-environment reset with DelayedPD actuators: Lab vs Newton (PhysX). - - Resets env 0 mid-simulation while env 1 continues uninterrupted. - Uses DelayedPD actuators because they carry internal state (delay - buffers) that must be properly reset per environment. + for act, state in stateful_pairs: + pushes_before = state.delay_state.num_pushes.numpy() + self.assertTrue( + (pushes_before > 0).all(), + "expected non-zero num_pushes for all DOFs after warmup", + ) - Verifies: - - Lab and Newton paths produce matching trajectories after partial reset. - - The two environments diverge after the reset (proving it took effect). - """ + articulation.reset(env_ids=torch.tensor([self.RESET_ENV], device=articulation.device, dtype=torch.long)) + + # Map each entry of ``act.indices`` to its env via ``adapter.num_joints`` + # (PhysX adapter is per-articulation so this equals articulation.num_joints — + # using adapter.num_joints keeps the test symmetric with the Newton path). + for act, state in stateful_pairs: + pushes_after = state.delay_state.num_pushes.numpy() + indices_np = act.indices.numpy() + for i, global_dof in enumerate(indices_np): + env = int(global_dof) // adapter.num_joints + if env == self.RESET_ENV: + self.assertEqual( + int(pushes_after[i]), 0, + f"DOF {i} (env {env}) should be reset to 0, got {pushes_after[i]}", + ) + else: + self.assertGreater( + int(pushes_after[i]), 0, + f"DOF {i} (env {env}) was NOT in reset env_ids but num_pushes is 0", + ) + finally: + ctx.__exit__(None, None, None) + + def test_lab_state_reset_isolated_to_reset_env(self): + """Lab: DelayedPDActuator circular buffer zeroed for env 0 only.""" + ctx, sim, articulation = self._build_and_warm(use_newton_actuators=False) + try: + from isaaclab.actuators import DelayedPDActuator # noqa: PLC0415 + + delayed = [a for a in articulation.actuators.values() if isinstance(a, DelayedPDActuator)] + self.assertGreater(len(delayed), 0, "expected at least one Lab DelayedPDActuator") + actuator = delayed[0] + buf = actuator.positions_delay_buffer._circular_buffer._buffer + self.assertIsNotNone(buf, "delay buffer should be populated after warmup") + self.assertTrue( + (buf[:, self.UNCHANGED_ENV] != 0).any().item(), + "expected non-zero buffer entries for env 1 after warmup", + ) - @classmethod - def setUpClass(cls): - cls.lab_result = _run_simulation_with_reset( - DELAYED_PD_ACTUATORS, use_newton_actuators=False, - ) - cls.newton_result = _run_simulation_with_reset( - DELAYED_PD_ACTUATORS, use_newton_actuators=True, - ) + articulation.reset(env_ids=torch.tensor([self.RESET_ENV], device=articulation.device, dtype=torch.long)) - def test_joint_positions_match(self): - for step_i, (lab, newton) in enumerate( - zip(self.lab_result["joint_pos"], self.newton_result["joint_pos"]) - ): - torch.testing.assert_close( - lab, - newton, - atol=2e-3, - rtol=1e-3, - msg=f"Positions diverged at step {step_i}", + self.assertTrue( + torch.all(buf[:, self.RESET_ENV] == 0).item(), + f"Lab: env {self.RESET_ENV} buffer not zeroed after reset.", ) - - def test_joint_velocities_match(self): - for step_i, (lab, newton) in enumerate( - zip(self.lab_result["joint_vel"], self.newton_result["joint_vel"]) - ): - torch.testing.assert_close( - lab, - newton, - atol=0.1, - rtol=1e-2, - msg=f"Velocities diverged at step {step_i}", + self.assertTrue( + (buf[:, self.UNCHANGED_ENV] != 0).any().item(), + f"Lab: env {self.UNCHANGED_ENV} buffer was zeroed — reset leaked into an unselected env.", ) - - def test_envs_diverge_after_reset(self): - """After resetting env 0, the two envs must have different states.""" - post_reset_pos = self.lab_result["joint_pos"][RESET_WARMUP_STEPS + 1] - diff = (post_reset_pos[0] - post_reset_pos[1]).abs().max().item() - self.assertGreater( - diff, 0.001, - "Env 0 and env 1 are identical after partial reset — reset had no effect", - ) - - def test_trajectories_not_trivial(self): - first = self.lab_result["joint_pos"][0] - last = self.lab_result["joint_pos"][-1] - diff = (last - first).abs().max().item() - self.assertGreater(diff, 0.01, "Joints did not move — test is trivial") + finally: + ctx.__exit__(None, None, None) # --------------------------------------------------------------------------- @@ -1039,12 +891,6 @@ def setUpClass(cls): use_newton_actuators=True, ) - def test_trajectories_not_trivial(self): - first = self.result["joint_pos"][0] - last = self.result["joint_pos"][-1] - diff = (last - first).abs().max().item() - self.assertGreater(diff, 0.01, "Joints did not move — test is trivial") - def test_positions_finite(self): for step_i, pos in enumerate(self.result["joint_pos"]): self.assertTrue( @@ -1068,7 +914,8 @@ def _make_dummy_mlp_checkpoint(device: str = "cpu") -> str: ).to(device).eval() scripted = torch.jit.script(net) - tmp = tempfile.NamedTemporaryFile(suffix=".pt", delete=False) + with tempfile.NamedTemporaryFile(suffix=".pt", delete=False) as tmp: + tmp_path = tmp.name extra = { "metadata.json": json.dumps({ "model_type": "mlp", @@ -1079,8 +926,8 @@ def _make_dummy_mlp_checkpoint(device: str = "cpu") -> str: "torque_scale": 2.0, }) } - torch.jit.save(scripted, tmp.name, _extra_files=extra) - return tmp.name + torch.jit.save(scripted, tmp_path, _extra_files=extra) + return tmp_path class _DummyLSTM(torch.nn.Module): @@ -1106,10 +953,11 @@ def _make_dummy_lstm_checkpoint(device: str = "cpu") -> str: net = _DummyLSTM().to(device).eval() scripted = torch.jit.script(net) - tmp = tempfile.NamedTemporaryFile(suffix=".pt", delete=False) + with tempfile.NamedTemporaryFile(suffix=".pt", delete=False) as tmp: + tmp_path = tmp.name extra = {"metadata.json": json.dumps({"model_type": "lstm"})} - torch.jit.save(scripted, tmp.name, _extra_files=extra) - return tmp.name + torch.jit.save(scripted, tmp_path, _extra_files=extra) + return tmp_path class TestNeuralMLPFunctional(unittest.TestCase): @@ -1148,12 +996,6 @@ def setUpClass(cls): def tearDownClass(cls): os.unlink(cls.mlp_path) - def test_trajectories_not_trivial(self): - first = self.result["joint_pos"][0] - last = self.result["joint_pos"][-1] - diff = (last - first).abs().max().item() - self.assertGreater(diff, 0.01, "Joints did not move — test is trivial") - def test_positions_finite(self): for step_i, pos in enumerate(self.result["joint_pos"]): self.assertTrue( @@ -1193,12 +1035,6 @@ def setUpClass(cls): def tearDownClass(cls): os.unlink(cls.lstm_path) - def test_trajectories_not_trivial(self): - first = self.result["joint_pos"][0] - last = self.result["joint_pos"][-1] - diff = (last - first).abs().max().item() - self.assertGreater(diff, 0.01, "Joints did not move — test is trivial") - def test_positions_finite(self): for step_i, pos in enumerate(self.result["joint_pos"]): self.assertTrue( From e9e337f562238fbececc58bd8a98e99c594c1ff1 Mon Sep 17 00:00:00 2001 From: Antoine Richard Date: Mon, 11 May 2026 17:55:14 +0200 Subject: [PATCH 22/35] Wire schema-side actuator authoring through AssetBaseCfg._post_spawn Replaces the spawn-cfg ``actuator_props`` field + spawner integration with a generic ``_post_spawn(stage)`` hook on :class:`AssetBaseCfg`, which :class:`ArticulationCfg` overrides to call :func:`schemas.define_actuator_properties`. ``AssetBase.__init__`` invokes the hook after spawn has produced the asset's prims on the stage. Same goal as the previous commit (USD authoring lives on the schema side; ``BaseArticulation`` no longer touches ``isaaclab_newton``), but matches the call mechanism jvonmuralt landed on PR #5455 head: simpler, reusable for other post-spawn schema work, and avoids leaking actuator plumbing into ``FileCfg``. --- .../assets/articulation/articulation_cfg.py | 28 +++++++++---------- source/isaaclab/isaaclab/assets/asset_base.py | 2 ++ .../isaaclab/assets/asset_base_cfg.py | 16 ++++++++++- .../sim/spawners/from_files/from_files.py | 4 --- .../sim/spawners/from_files/from_files_cfg.py | 14 ---------- 5 files changed, 30 insertions(+), 34 deletions(-) diff --git a/source/isaaclab/isaaclab/assets/articulation/articulation_cfg.py b/source/isaaclab/isaaclab/assets/articulation/articulation_cfg.py index b7cb3cafd9c7..f5e14e2def02 100644 --- a/source/isaaclab/isaaclab/assets/articulation/articulation_cfg.py +++ b/source/isaaclab/isaaclab/assets/articulation/articulation_cfg.py @@ -6,7 +6,7 @@ from __future__ import annotations from dataclasses import MISSING -from typing import TYPE_CHECKING +from typing import TYPE_CHECKING, Any from isaaclab.actuators import ActuatorBaseCfg from isaaclab.utils import configclass @@ -75,19 +75,17 @@ class InitialStateCfg(AssetBaseCfg.InitialStateCfg): """Print the resolution of actuator final value when input cfg is different from USD value, Defaults to False """ - def __post_init__(self) -> None: - """Propagate :attr:`actuators` to the spawn cfg's ``actuator_props`` field. + def _post_spawn(self, stage: Any) -> None: + """Author ``NewtonActuator`` USD prims from :attr:`actuators` after spawn. - The spawner uses ``actuator_props`` to author ``NewtonActuator`` USD prims - via :func:`~isaaclab.sim.schemas.define_actuator_properties` (a no-op when - ``use_newton_actuators`` is disabled). Keeping the propagation here means - users only declare actuators once, on the asset cfg. - - User-provided ``spawn.actuator_props`` is preserved and takes precedence. + Invoked by :class:`~isaaclab.assets.AssetBase` once the articulation's prims + exist on the stage. Delegates to + :func:`~isaaclab.sim.schemas.define_actuator_properties`, which gates itself + on ``sim_cfg.use_newton_actuators`` and silently no-ops when the simulation + is not configured for Newton-native actuators. """ - if ( - self.spawn is not None - and self.actuators is not MISSING - and getattr(self.spawn, "actuator_props", None) is None - ): - self.spawn.actuator_props = dict(self.actuators) + if self.actuators is MISSING: + return + from isaaclab.sim.schemas import define_actuator_properties # noqa: PLC0415 + + define_actuator_properties(self.prim_path, self.actuators, stage=stage) diff --git a/source/isaaclab/isaaclab/assets/asset_base.py b/source/isaaclab/isaaclab/assets/asset_base.py index b72788083ad3..7d403d49ac5a 100644 --- a/source/isaaclab/isaaclab/assets/asset_base.py +++ b/source/isaaclab/isaaclab/assets/asset_base.py @@ -96,6 +96,8 @@ def __init__(self, cfg: AssetBaseCfg): matching_prims = sim_utils.find_matching_prims(check_path) if len(matching_prims) == 0: raise RuntimeError(f"Could not find prim with path {check_path}.") + # schema-side post-spawn hook (e.g. ArticulationCfg authors NewtonActuator prims here) + self.cfg._post_spawn(self.stage) else: # asset should exist at run time check_path = self.cfg.prim_path diff --git a/source/isaaclab/isaaclab/assets/asset_base_cfg.py b/source/isaaclab/isaaclab/assets/asset_base_cfg.py index 4575acc08452..ddf0ec8bb875 100644 --- a/source/isaaclab/isaaclab/assets/asset_base_cfg.py +++ b/source/isaaclab/isaaclab/assets/asset_base_cfg.py @@ -6,7 +6,7 @@ from __future__ import annotations from dataclasses import MISSING -from typing import Literal +from typing import Any, Literal from isaaclab.sim import SpawnerCfg from isaaclab.utils import configclass @@ -88,3 +88,17 @@ class InitialStateCfg: When ``None`` (the default), shape checks follow Python's ``__debug__`` flag — enabled in normal mode, disabled with ``python -O``. """ + + def _post_spawn(self, stage: Any) -> None: + """Hook invoked by :class:`~isaaclab.assets.AssetBase` after the asset's prims are + spawned and verified to exist on the stage. + + The default implementation is a no-op. Subclasses that need to author additional + USD schemas tied to this asset (for example, :class:`~isaaclab.assets.ArticulationCfg` + which authors ``NewtonActuator`` prims from its ``actuators`` mapping) should + override this method. + + Args: + stage: The USD stage on which the asset was spawned. + """ + pass diff --git a/source/isaaclab/isaaclab/sim/spawners/from_files/from_files.py b/source/isaaclab/isaaclab/sim/spawners/from_files/from_files.py index ecb6d5fa2052..be7bd6e7074e 100644 --- a/source/isaaclab/isaaclab/sim/spawners/from_files/from_files.py +++ b/source/isaaclab/isaaclab/sim/spawners/from_files/from_files.py @@ -367,10 +367,6 @@ def _spawn_from_usd_file( if cfg.joint_drive_props is not None: schemas.modify_joint_drive_properties(prim_path, cfg.joint_drive_props) - # author Newton-native actuator USD prims (no-op when use_newton_actuators=False) - if cfg.actuator_props is not None: - schemas.define_actuator_properties(prim_path, cfg.actuator_props, stage=stage) - # define deformable body properties, or modify if deformable body API is present (PhysX only) if cfg.deformable_props is not None: prim = stage.GetPrimAtPath(prim_path) diff --git a/source/isaaclab/isaaclab/sim/spawners/from_files/from_files_cfg.py b/source/isaaclab/isaaclab/sim/spawners/from_files/from_files_cfg.py index 1d321b9e141e..87fc48de4d30 100644 --- a/source/isaaclab/isaaclab/sim/spawners/from_files/from_files_cfg.py +++ b/source/isaaclab/isaaclab/sim/spawners/from_files/from_files_cfg.py @@ -7,7 +7,6 @@ from collections.abc import Callable from dataclasses import MISSING -from typing import TYPE_CHECKING # deformables only supported on PhysX backend from isaaclab_physx.sim.spawners.spawner_cfg import DeformableObjectSpawnerCfg @@ -18,9 +17,6 @@ from isaaclab.utils import configclass from isaaclab.utils.assets import ISAAC_NUCLEUS_DIR -if TYPE_CHECKING: - from isaaclab.actuators import ActuatorBaseCfg - @configclass class FileCfg(RigidObjectSpawnerCfg, DeformableObjectSpawnerCfg): @@ -60,16 +56,6 @@ class FileCfg(RigidObjectSpawnerCfg, DeformableObjectSpawnerCfg): for specific joints in an articulation. """ - actuator_props: dict[str, ActuatorBaseCfg] | None = None - """Newton-native actuator configs to author as ``NewtonActuator`` USD prims. - - Auto-populated from :attr:`~isaaclab.assets.ArticulationCfg.actuators` when an - articulation cfg owns this spawn cfg. The spawner authors NewtonActuator prims - via :func:`~isaaclab.sim.schemas.define_actuator_properties` after the asset is - spawned. The function is a no-op unless - :attr:`~isaaclab.sim.SimulationCfg.use_newton_actuators` is enabled. - """ - visual_material_path: str = "material" """Path to the visual material to use for the prim. Defaults to "material". From 82a03d9105b9acb6f308486ba2d8daa289b9ea9e Mon Sep 17 00:00:00 2001 From: Antoine Richard Date: Mon, 11 May 2026 19:11:50 +0200 Subject: [PATCH 23/35] Add changelog fragments for pr5455-followups Three patch/minor fragments covering Phase A (review nits) and Phase B (schema-side actuator authoring): * isaaclab.minor: new define_actuator_properties writer + _post_spawn hook on AssetBaseCfg. * isaaclab_newton: removed isaaclab_newton.actuators.authoring module + decimation declaration regrouping. * isaaclab_physx: FF-routing comment wording fix. --- ...toiner-refactor-pr5455-followups.minor.rst | 22 +++++++++++++++++++ .../antoiner-refactor-pr5455-followups.rst | 11 ++++++++++ .../antoiner-refactor-pr5455-followups.rst | 8 +++++++ 3 files changed, 41 insertions(+) create mode 100644 source/isaaclab/changelog.d/antoiner-refactor-pr5455-followups.minor.rst create mode 100644 source/isaaclab_newton/changelog.d/antoiner-refactor-pr5455-followups.rst create mode 100644 source/isaaclab_physx/changelog.d/antoiner-refactor-pr5455-followups.rst diff --git a/source/isaaclab/changelog.d/antoiner-refactor-pr5455-followups.minor.rst b/source/isaaclab/changelog.d/antoiner-refactor-pr5455-followups.minor.rst new file mode 100644 index 000000000000..b1dff35671d8 --- /dev/null +++ b/source/isaaclab/changelog.d/antoiner-refactor-pr5455-followups.minor.rst @@ -0,0 +1,22 @@ +Added +^^^^^ + +* Added :func:`~isaaclab.sim.schemas.define_actuator_properties` to author + ``NewtonActuator`` USD prims from IsaacLab actuator configs. Lives alongside + the other ``define_*_properties`` schema writers and is invoked from the + schema-side post-spawn hook below. +* Added :meth:`~isaaclab.assets.AssetBaseCfg._post_spawn` hook (no-op by + default) invoked by :class:`~isaaclab.assets.AssetBase` after spawning the + asset. :class:`~isaaclab.assets.ArticulationCfg` overrides it to author + Newton-native actuator prims from :attr:`~isaaclab.assets.ArticulationCfg.actuators`. + +Changed +^^^^^^^ + +* :class:`~isaaclab.assets.BaseArticulation` no longer imports + ``isaaclab_newton`` from its ``__init__``. Newton-native actuator authoring + now flows through the generic ``_post_spawn`` hook on + :class:`~isaaclab.assets.AssetBaseCfg`. +* :meth:`~isaaclab.physics.PhysicsManager.set_decimation` now has an explicit + ``pass`` body so the base class is consistent with the other no-op + classmethods. diff --git a/source/isaaclab_newton/changelog.d/antoiner-refactor-pr5455-followups.rst b/source/isaaclab_newton/changelog.d/antoiner-refactor-pr5455-followups.rst new file mode 100644 index 000000000000..d340938c97b4 --- /dev/null +++ b/source/isaaclab_newton/changelog.d/antoiner-refactor-pr5455-followups.rst @@ -0,0 +1,11 @@ +Changed +^^^^^^^ + +* Moved Newton-native actuator USD authoring out of + ``isaaclab_newton.actuators.authoring`` (now deleted) into + :func:`~isaaclab.sim.schemas.define_actuator_properties`. The authoring step + is now invoked via the schema-side ``_post_spawn`` hook on + :class:`~isaaclab.assets.AssetBaseCfg`. +* Grouped :attr:`~isaaclab_newton.physics.NewtonManager._decimation` next to + :attr:`~isaaclab_newton.physics.NewtonManager._num_substeps` for consistency + with related solver-stepping configuration. diff --git a/source/isaaclab_physx/changelog.d/antoiner-refactor-pr5455-followups.rst b/source/isaaclab_physx/changelog.d/antoiner-refactor-pr5455-followups.rst new file mode 100644 index 000000000000..5d20d3888780 --- /dev/null +++ b/source/isaaclab_physx/changelog.d/antoiner-refactor-pr5455-followups.rst @@ -0,0 +1,8 @@ +Changed +^^^^^^^ + +* Reworded the FF-routing comments in + :class:`~isaaclab_physx.assets.Articulation` to refer to "actuated DOFs" + rather than splitting on implicit vs explicit, since the + ``synch_torque_and_apply_implicit_feedforwards`` kernel operates on the full + actuated DOF set. From e7e2619532abf072fc36e8b5e38cdafb435df71d Mon Sep 17 00:00:00 2001 From: Julia von Muralt Date: Tue, 12 May 2026 09:14:06 +0200 Subject: [PATCH 24/35] tmp --- .../assets/articulation/articulation_cfg.py | 25 +- .../isaaclab/assets/asset_base_cfg.py | 19 +- .../isaaclab/sim/schemas/schemas_actuators.py | 389 ++++++++++++++++++ .../isaaclab_newton/actuators/authoring.py | 2 +- 4 files changed, 414 insertions(+), 21 deletions(-) create mode 100644 source/isaaclab/isaaclab/sim/schemas/schemas_actuators.py diff --git a/source/isaaclab/isaaclab/assets/articulation/articulation_cfg.py b/source/isaaclab/isaaclab/assets/articulation/articulation_cfg.py index 0ca196f58bb3..240c0927af12 100644 --- a/source/isaaclab/isaaclab/assets/articulation/articulation_cfg.py +++ b/source/isaaclab/isaaclab/assets/articulation/articulation_cfg.py @@ -6,9 +6,7 @@ from __future__ import annotations from dataclasses import MISSING -from typing import TYPE_CHECKING - -from pxr import Usd +from typing import TYPE_CHECKING, Any from isaaclab.actuators import ActuatorBaseCfg from isaaclab.utils import configclass @@ -77,22 +75,17 @@ class InitialStateCfg(AssetBaseCfg.InitialStateCfg): """Print the resolution of actuator final value when input cfg is different from USD value, Defaults to False """ - def _post_spawn(self, stage: Usd.Stage) -> None: + def _post_spawn(self, stage: Any) -> None: """Author ``NewtonActuator`` USD prims from :attr:`actuators` after spawn. Invoked by :class:`~isaaclab.assets.AssetBase` once the articulation's prims - exist on the stage. The actual authoring logic lives in - :func:`~isaaclab_newton.actuators.author_actuator_prims_for_articulation`, - which gates itself on ``sim_cfg.use_newton_actuators`` and silently no-ops when - the Newton actuator package is unavailable. + exist on the stage. Delegates to + :func:`~isaaclab.sim.schemas.define_actuator_properties`, which gates itself + on ``sim_cfg.use_newton_actuators`` and silently no-ops when the simulation + is not configured for Newton-native actuators. """ - from isaaclab.sim import SimulationContext # noqa: PLC0415 - - try: - from isaaclab_newton.actuators import author_actuator_prims_for_articulation # noqa: PLC0415 - except ImportError: + if self.actuators is MISSING: return + from isaaclab.sim.schemas.schemas_actuators import define_actuator_properties # noqa: PLC0415 - sim_ctx = SimulationContext.instance() - sim_cfg = sim_ctx.cfg if sim_ctx is not None else None - author_actuator_prims_for_articulation(self, sim_cfg, stage) + define_actuator_properties(self.prim_path, self.actuators, stage=stage) diff --git a/source/isaaclab/isaaclab/assets/asset_base_cfg.py b/source/isaaclab/isaaclab/assets/asset_base_cfg.py index d936982c52d1..ddf0ec8bb875 100644 --- a/source/isaaclab/isaaclab/assets/asset_base_cfg.py +++ b/source/isaaclab/isaaclab/assets/asset_base_cfg.py @@ -6,9 +6,7 @@ from __future__ import annotations from dataclasses import MISSING -from typing import Literal - -from pxr import Usd +from typing import Any, Literal from isaaclab.sim import SpawnerCfg from isaaclab.utils import configclass @@ -78,7 +76,20 @@ class InitialStateCfg: debug_vis: bool = False """Whether to enable debug visualization for the asset. Defaults to ``False``.""" - def _post_spawn(self, stage: Usd.Stage) -> None: + disable_shape_checks: bool | None = None + """Disable shape/dtype validation in setter and writer methods. + + When ``True``, :meth:`~AssetBase.assert_shape_and_dtype` and + :meth:`~AssetBase.assert_shape_and_dtype_mask` become no-ops, + eliminating per-call assertion overhead. + + When ``False``, shape checks are always enabled, even under ``python -O``. + + When ``None`` (the default), shape checks follow Python's ``__debug__`` + flag — enabled in normal mode, disabled with ``python -O``. + """ + + def _post_spawn(self, stage: Any) -> None: """Hook invoked by :class:`~isaaclab.assets.AssetBase` after the asset's prims are spawned and verified to exist on the stage. diff --git a/source/isaaclab/isaaclab/sim/schemas/schemas_actuators.py b/source/isaaclab/isaaclab/sim/schemas/schemas_actuators.py new file mode 100644 index 000000000000..40456384e473 --- /dev/null +++ b/source/isaaclab/isaaclab/sim/schemas/schemas_actuators.py @@ -0,0 +1,389 @@ +# Copyright (c) 2022-2026, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +"""USD schema authoring for Newton-native actuators. + +:func:`define_actuator_properties` translates IsaacLab actuator configs +into ``NewtonActuator`` USD prims. Both the Newton ``ModelBuilder.add_usd`` +path and the PhysX adapter's +:meth:`~isaaclab_newton.actuators.adapter.NewtonActuatorAdapter.from_usd` +read the same authored prims, ensuring both backends construct +:class:`~newton.actuators.Actuator` instances with matching parameters. + +This module lives on the schema side so that authoring is a regular +``define_*_properties`` step in the spawner pipeline, alongside +:func:`define_articulation_root_properties` and friends, rather than a +side effect of asset construction. +""" + +from __future__ import annotations + +import logging +import re +from typing import Any + +logger = logging.getLogger(__name__) + + +def resolve_per_dof( + value: dict[str, float | int] | float | int | None, + joint_names: list[str], + cast: type = float, +) -> dict[str, float | int]: + """Expand a scalar or regex-keyed dict cfg value into a per-joint mapping. + + Used by :func:`define_actuator_properties` to flatten the various + accepted forms of a per-DOF config field (``stiffness``, ``damping``, + ``effort_limit``, …) into a single ``{joint_name: value}`` dict that + the authoring loop can ``.get(jname, default)`` against. + + Accepted forms of *value*: + + * ``None`` → empty dict. + * scalar (``int`` / ``float``) → broadcast to every joint name. + * dict → keys are treated as regex patterns and matched against + *joint_names* via :func:`re.fullmatch`. The first matching pattern + wins per joint name. + """ + if value is None: + return {} + if isinstance(value, (int, float)): + return {name: cast(value) for name in joint_names} + if isinstance(value, dict): + result: dict[str, float | int] = {} + for name in joint_names: + for pattern, v in value.items(): + if re.fullmatch(pattern, name): + result[name] = cast(v) + break + return result + return {} + + +def define_actuator_properties( + prim_path: str, + actuator_cfgs: dict[str, Any], + stage: Any | None = None, +) -> None: + """Author ``NewtonActuator`` USD prims under an articulation root. + + For every joint covered by an explicit (non-implicit) Lab actuator + config, any existing ``NewtonActuator`` prim targeting that joint is + replaced by a new one created from the config values. Joints **not** + covered by any Lab config keep their USD-authored actuators unchanged. + + The supported config-to-schema mapping is: + + * :class:`~isaaclab.actuators.IdealPDActuatorCfg` → + ``NewtonPDControlAPI`` + ``NewtonMaxEffortClampingAPI`` + * :class:`~isaaclab.actuators.DCMotorCfg` → + ``NewtonPDControlAPI`` + ``NewtonDCMotorClampingAPI`` + * :class:`~isaaclab.actuators.DelayedPDActuatorCfg` → + same as ``IdealPDActuatorCfg`` + ``NewtonActuatorDelayAPI`` + * :class:`~isaaclab.actuators.RemotizedPDActuatorCfg` → + same as ``DelayedPDActuatorCfg`` + ``NewtonPositionBasedClampingAPI`` + * :class:`~isaaclab.actuators.ActuatorNetMLPCfg` / + :class:`~isaaclab.actuators.ActuatorNetLSTMCfg` → + ``NewtonNeuralControlAPI`` (+ ``NewtonDCMotorClampingAPI``) + + No-ops (returns immediately) when: + + * the active :class:`~isaaclab.sim.SimulationContext` was configured + with ``use_newton_actuators=False`` (or no context is active), or + * *prim_path* does not resolve to a valid prim on the stage. + + Must be called **after** the articulation is spawned (joint prims + exist on stage) and **before** the cloner / ``ModelBuilder.add_usd`` + reads the stage. + + Args: + prim_path: Root prim path of the articulation (e.g. + ``"/World/Env_0/Robot"``). May contain a regex pattern; the + first matching prim is used. + actuator_cfgs: Mapping of group name to + :class:`~isaaclab.actuators.ActuatorBaseCfg`. + stage: USD stage to author on. When ``None``, the current stage + is used. + """ + from isaaclab.sim import SimulationContext # noqa: PLC0415 + + sim_ctx = SimulationContext.instance() + sim_cfg = sim_ctx.cfg if sim_ctx is not None else None + if sim_cfg is None or not getattr(sim_cfg, "use_newton_actuators", False): + return + + from isaaclab.sim.utils.queries import find_first_matching_prim # noqa: PLC0415 + from isaaclab.sim.utils.stage import get_current_stage # noqa: PLC0415 + + if stage is None: + stage = get_current_stage() + + first_prim = find_first_matching_prim(prim_path) + if first_prim is None: + return + articulation_prim_path = str(first_prim.GetPath()) + + _author_actuator_prims(stage, articulation_prim_path, actuator_cfgs) + + +def _author_actuator_prims( + stage: Any, + articulation_prim_path: str, + actuator_cfgs: dict[str, Any], +) -> None: + """Inner authoring routine; exposed separately for test fixtures.""" + from pxr import Sdf # noqa: PLC0415 + + from isaaclab.actuators import ImplicitActuator # noqa: PLC0415 + from isaaclab.utils.string import resolve_matching_names # noqa: PLC0415 + + art_prim = stage.GetPrimAtPath(articulation_prim_path) + if not art_prim.IsValid(): + raise ValueError(f"Articulation prim not found: {articulation_prim_path}") + + joint_inventory = _collect_joint_prims(art_prim) + all_joint_names = list(joint_inventory.keys()) + + covered_joint_paths: set[str] = set() + + cfg_entries: list[tuple[str, Any, list[str]]] = [] + for group_name, cfg in actuator_cfgs.items(): + cls_type = cfg.class_type + is_implicit = ( + "ImplicitActuator" in cls_type if isinstance(cls_type, str) else issubclass(cls_type, ImplicitActuator) + ) + if is_implicit: + continue + + _ids, joint_names = resolve_matching_names(cfg.joint_names_expr, all_joint_names) + if not joint_names: + continue + + cfg_entries.append((group_name, cfg, joint_names)) + for jname in joint_names: + covered_joint_paths.add(joint_inventory[jname]) + + _remove_actuator_prims_for_joints(art_prim, covered_joint_paths) + + from isaaclab.actuators import DCMotorCfg, DelayedPDActuatorCfg # noqa: PLC0415 + from isaaclab.actuators.actuator_net_cfg import ActuatorNetLSTMCfg, ActuatorNetMLPCfg # noqa: PLC0415 + from isaaclab.actuators.actuator_pd_cfg import IdealPDActuatorCfg, RemotizedPDActuatorCfg # noqa: PLC0415 + + _SUPPORTED_CFG_TYPES = ( + IdealPDActuatorCfg, + DCMotorCfg, + DelayedPDActuatorCfg, + RemotizedPDActuatorCfg, + ActuatorNetMLPCfg, + ActuatorNetLSTMCfg, + ) + + for group_name, cfg, joint_names in cfg_entries: + if not isinstance(cfg, _SUPPORTED_CFG_TYPES): + logger.warning( + "Actuator group '%s' uses config type '%s' which is not supported by Newton-native" + " actuator authoring. The group will be skipped.", + group_name, + type(cfg).__name__, + ) + continue + stiffness_map = resolve_per_dof(getattr(cfg, "stiffness", None), joint_names) + damping_map = resolve_per_dof(getattr(cfg, "damping", None), joint_names) + effort_map = resolve_per_dof(getattr(cfg, "effort_limit", None), joint_names) + + is_neural = isinstance(cfg, (ActuatorNetMLPCfg, ActuatorNetLSTMCfg)) + is_remotized = isinstance(cfg, RemotizedPDActuatorCfg) + is_dc_motor = isinstance(cfg, DCMotorCfg) + is_delayed = isinstance(cfg, DelayedPDActuatorCfg) + + vel_limit_map = resolve_per_dof(getattr(cfg, "velocity_limit", None), joint_names) if is_dc_motor else {} + sat_effort_map = resolve_per_dof(getattr(cfg, "saturation_effort", None), joint_names) if is_dc_motor else {} + + raw_delay = getattr(cfg, "max_delay", 0) if is_delayed else 0 + delay_map = resolve_per_dof(raw_delay, joint_names, cast=int) if raw_delay else {} + + patched_model_path: str | None = None + if is_neural: + meta: dict[str, Any] = {} + if isinstance(cfg, ActuatorNetMLPCfg): + meta["model_type"] = "mlp" + meta["input_order"] = cfg.input_order + meta["input_idx"] = list(cfg.input_idx) + meta["pos_scale"] = cfg.pos_scale + meta["vel_scale"] = cfg.vel_scale + meta["torque_scale"] = cfg.torque_scale + else: + meta["model_type"] = "lstm" + patched_model_path = _resave_checkpoint_with_metadata(cfg.network_file, meta) + + for jname in joint_names: + joint_prim_path = joint_inventory[jname] + + schemas: list[str] = [] + attrs: dict[str, float | int] = {} + array_attrs: dict[str, list[float]] = {} + + if is_neural: + schemas.append("NewtonNeuralControlAPI") + else: + schemas.append("NewtonPDControlAPI") + attrs["kp"] = stiffness_map.get(jname, 0.0) + attrs["kd"] = damping_map.get(jname, 0.0) + + if is_dc_motor: + schemas.append("NewtonDCMotorClampingAPI") + attrs["saturation_effort"] = sat_effort_map.get(jname, 0.0) + if jname in vel_limit_map: + attrs["velocity_limit"] = vel_limit_map[jname] + if jname in effort_map: + attrs["max_motor_effort"] = effort_map[jname] + elif jname in effort_map: + schemas.append("NewtonMaxEffortClampingAPI") + attrs["max_effort"] = effort_map[jname] + + if is_remotized and isinstance(cfg, RemotizedPDActuatorCfg): + lookup = cfg.joint_parameter_lookup + schemas.append("NewtonPositionBasedClampingAPI") + array_attrs["lookup_positions"] = [row[0] for row in lookup] + array_attrs["lookup_efforts"] = [row[2] for row in lookup] + + delay_steps = delay_map.get(jname, 0) + if delay_steps > 0: + schemas.append("NewtonActuatorDelayAPI") + attrs["delay_steps"] = delay_steps + attrs["max_delay"] = delay_steps + + act_prim_path = f"{articulation_prim_path}/{group_name}_{jname}_actuator" + act_prim = stage.DefinePrim(act_prim_path, "NewtonActuator") + + existing = act_prim.GetMetadata("apiSchemas") or Sdf.TokenListOp() + existing.prependedItems = list(schemas) + act_prim.SetMetadata("apiSchemas", existing) + + rel = act_prim.CreateRelationship("newton:targets") + rel.SetTargets([Sdf.Path(joint_prim_path)]) + + if patched_model_path is not None: + act_prim.CreateAttribute("newton:modelPath", Sdf.ValueTypeNames.Asset).Set( + Sdf.AssetPath(patched_model_path) + ) + + for attr_name, attr_val in attrs.items(): + usd_name = f"newton:{_snake_to_camel(attr_name)}" + if isinstance(attr_val, int): + act_prim.CreateAttribute(usd_name, Sdf.ValueTypeNames.Int).Set(attr_val) + else: + act_prim.CreateAttribute(usd_name, Sdf.ValueTypeNames.Float).Set(float(attr_val)) + + for attr_name, attr_val in array_attrs.items(): + usd_name = f"newton:{_snake_to_camel(attr_name)}" + act_prim.CreateAttribute(usd_name, Sdf.ValueTypeNames.FloatArray).Set(attr_val) + + +# --------------------------------------------------------------------------- +# Private helpers +# --------------------------------------------------------------------------- + +_SNAKE_TO_CAMEL_RE = re.compile(r"_([a-z])") + + +def _snake_to_camel(name: str) -> str: + """Convert a snake_case name to camelCase.""" + return _SNAKE_TO_CAMEL_RE.sub(lambda m: m.group(1).upper(), name) + + +def _collect_joint_prims(art_prim: Any) -> dict[str, str]: + """Collect all joint prims under an articulation subtree. + + Returns: + Ordered mapping of joint name to full prim path. + """ + from pxr import Usd # noqa: PLC0415 + + _JOINT_TYPES = {"PhysicsRevoluteJoint", "PhysicsPrismaticJoint"} + + joints: dict[str, str] = {} + for prim in Usd.PrimRange(art_prim): + if prim.GetTypeName() in _JOINT_TYPES: + joints[prim.GetName()] = str(prim.GetPath()) + return joints + + +def _remove_actuator_prims_for_joints( + art_prim: Any, + joint_paths: set[str], +) -> None: + """Deactivate ``NewtonActuator`` prims whose target is in *joint_paths*. + + Deactivated prims are invisible to ``Usd.PrimRange`` and therefore + ignored by ``ModelBuilder.add_usd``. Using ``SetActive(False)`` + instead of ``RemovePrim`` works correctly when the prim originates + from a USD reference or payload. + + Only prims under the *art_prim* subtree are considered. + """ + from pxr import Usd # noqa: PLC0415 + + to_deactivate: list = [] + for prim in Usd.PrimRange(art_prim): + if prim.GetTypeName() != "NewtonActuator": + continue + rel = prim.GetRelationship("newton:targets") + if rel and rel.IsValid(): + for target in rel.GetTargets(): + if str(target) in joint_paths: + to_deactivate.append(prim) + break + + for prim in to_deactivate: + prim.SetActive(False) + + +def _resave_checkpoint_with_metadata( + original_path: str, + metadata: dict[str, Any], +) -> str: + """Re-save a neural-network checkpoint with updated metadata. + + Loads the original TorchScript or dict checkpoint, merges *metadata* + into any existing metadata (Lab config values take precedence), and + writes the result to a temporary ``.pt`` file that persists for the + lifetime of the process. + + Returns: + Path to the temporary checkpoint file. + """ + import json # noqa: PLC0415 + import tempfile # noqa: PLC0415 + + import torch # noqa: PLC0415 + + extra_files: dict[str, str] = {"metadata.json": ""} + is_torchscript = True + try: + net = torch.jit.load(original_path, map_location="cpu", _extra_files=extra_files) + existing_meta = json.loads(extra_files["metadata.json"]) if extra_files["metadata.json"] else {} + except Exception: + is_torchscript = False + checkpoint = torch.load(original_path, map_location="cpu", weights_only=False) + if not isinstance(checkpoint, dict) or "model" not in checkpoint: + raise ValueError( + f"Cannot load checkpoint at '{original_path}'; " + "expected a TorchScript archive or a dict with a 'model' key" + ) + net = checkpoint["model"] + existing_meta = checkpoint.get("metadata", {}) + + merged = {**existing_meta, **metadata} + + with tempfile.NamedTemporaryFile(suffix=".pt", delete=False) as tmp: + tmp_path = tmp.name + if is_torchscript: + extra_out = {"metadata.json": json.dumps(merged)} + torch.jit.save(net, tmp_path, _extra_files=extra_out) + else: + torch.save({"model": net, "metadata": merged}, tmp_path) + + return tmp_path diff --git a/source/isaaclab_newton/isaaclab_newton/actuators/authoring.py b/source/isaaclab_newton/isaaclab_newton/actuators/authoring.py index 34a4e4539b64..fd5abe9da9b9 100644 --- a/source/isaaclab_newton/isaaclab_newton/actuators/authoring.py +++ b/source/isaaclab_newton/isaaclab_newton/actuators/authoring.py @@ -79,7 +79,7 @@ def author_actuator_prims_for_articulation(cfg: Any, sim_cfg: Any, stage: Any) - authoring with a single line, and lets the rules be tested without constructing an articulation instance. """ - if sim_cfg is None or not getattr(sim_cfg, "use_newton_actuators", False): + if sim_cfg is not None and not getattr(sim_cfg, "use_newton_actuators", False): return from isaaclab.sim.utils.queries import find_first_matching_prim # noqa: PLC0415 From 03f8e1d314b3833e7ea7ade172d49f117003f3b1 Mon Sep 17 00:00:00 2001 From: Julia von Muralt Date: Tue, 12 May 2026 09:58:24 +0200 Subject: [PATCH 25/35] fix --- .../isaaclab_newton/physics/newton_manager.py | 15 ++++++++++----- 1 file changed, 10 insertions(+), 5 deletions(-) diff --git a/source/isaaclab_newton/isaaclab_newton/physics/newton_manager.py b/source/isaaclab_newton/isaaclab_newton/physics/newton_manager.py index e748b288cbe7..77ba356234b4 100644 --- a/source/isaaclab_newton/isaaclab_newton/physics/newton_manager.py +++ b/source/isaaclab_newton/isaaclab_newton/physics/newton_manager.py @@ -827,9 +827,11 @@ def start_simulation(cls) -> None: # The single global actuator adapter is built lazily on the first # call to ``activate_newton_actuator_path`` from any Newton-fast-path - # articulation after this point. - cls._adapter = None - cls._use_newton_actuators_active = False + # articulation after this point. Assign through the explicit base + # class so external readers (which import ``NewtonManager`` directly) + # observe the canonical state regardless of which subclass is active. + NewtonManager._adapter = None + NewtonManager._use_newton_actuators_active = False # Allocate per-world reset masks (used by all solvers for masked FK, and by Kamino for masked reset) NewtonManager._world_reset_mask = wp.zeros(cls._model.world_count, dtype=wp.int32, device=device) @@ -1410,7 +1412,10 @@ def activate_newton_actuator_path(cls) -> None: :class:`NewtonActuatorAdapter` over the full flat DOF layout; later calls reuse it. """ - cls._use_newton_actuators_active = True + # Shared state lives on the base class so all readers (including + # framework code that imports ``NewtonManager`` directly) see the + # same flag regardless of which solver subclass is active. + NewtonManager._use_newton_actuators_active = True if cls._adapter is not None: return @@ -1419,7 +1424,7 @@ def activate_newton_actuator_path(cls) -> None: from isaaclab_newton.actuators import NewtonActuatorAdapter # noqa: PLC0415 dofs_per_env = cls._model.joint_dof_count // cls._num_envs - cls._adapter = NewtonActuatorAdapter( + NewtonManager._adapter = NewtonActuatorAdapter( actuators=list(cls._model.actuators), num_envs=cls._num_envs, num_joints=dofs_per_env, From 8e571c7dc232e27051ffbc092960ff509c00bdd3 Mon Sep 17 00:00:00 2001 From: Julia von Muralt Date: Tue, 12 May 2026 10:07:30 +0200 Subject: [PATCH 26/35] format --- .../assets/articulation/base_articulation.py | 1 + source/isaaclab/isaaclab/envs/mdp/events.py | 26 +- .../isaaclab_newton/actuators/adapter.py | 35 +- .../isaaclab_newton/actuators/kernels.py | 2 - .../actuators/physx_wrapper.py | 2 +- .../assets/articulation/articulation.py | 42 ++- .../assets/articulation/articulation_data.py | 4 +- .../isaaclab_newton/physics/newton_manager.py | 2 +- .../assets/test_newton_actuators_newton.py | 298 ++++++++++-------- .../assets/articulation/articulation.py | 61 +++- .../assets/test_newton_actuators_physx.py | 143 ++++++--- 11 files changed, 380 insertions(+), 236 deletions(-) diff --git a/source/isaaclab/isaaclab/assets/articulation/base_articulation.py b/source/isaaclab/isaaclab/assets/articulation/base_articulation.py index fb07050f334c..f57d54aa74c5 100644 --- a/source/isaaclab/isaaclab/assets/articulation/base_articulation.py +++ b/source/isaaclab/isaaclab/assets/articulation/base_articulation.py @@ -98,6 +98,7 @@ def __init__(self, cfg: ArticulationCfg): super().__init__(cfg) sim_ctx = SimulationContext.instance() self._sim_cfg = sim_ctx.cfg if sim_ctx is not None else None + """ Properties """ diff --git a/source/isaaclab/isaaclab/envs/mdp/events.py b/source/isaaclab/isaaclab/envs/mdp/events.py index 36612af65f49..f430c105ee78 100644 --- a/source/isaaclab/isaaclab/envs/mdp/events.py +++ b/source/isaaclab/isaaclab/envs/mdp/events.py @@ -1246,22 +1246,32 @@ def randomize(data: torch.Tensor, params: tuple[float, float]) -> torch.Tensor: if stiffness_distribution_params is not None: new_stiffness = self.default_joint_stiffness[env_ids][:, joint_ids].clone() _randomize_prop_by_op( - new_stiffness, stiffness_distribution_params, - dim_0_ids=None, dim_1_ids=slice(None), - operation=operation, distribution=distribution, + new_stiffness, + stiffness_distribution_params, + dim_0_ids=None, + dim_1_ids=slice(None), + operation=operation, + distribution=distribution, ) self.asset.write_actuator_stiffness_to_sim( - stiffness=new_stiffness, env_ids=env_ids, joint_ids=joint_ids, + stiffness=new_stiffness, + env_ids=env_ids, + joint_ids=joint_ids, ) if damping_distribution_params is not None: new_damping = self.default_joint_damping[env_ids][:, joint_ids].clone() _randomize_prop_by_op( - new_damping, damping_distribution_params, - dim_0_ids=None, dim_1_ids=slice(None), - operation=operation, distribution=distribution, + new_damping, + damping_distribution_params, + dim_0_ids=None, + dim_1_ids=slice(None), + operation=operation, + distribution=distribution, ) self.asset.write_actuator_damping_to_sim( - damping=new_damping, env_ids=env_ids, joint_ids=joint_ids, + damping=new_damping, + env_ids=env_ids, + joint_ids=joint_ids, ) diff --git a/source/isaaclab_newton/isaaclab_newton/actuators/adapter.py b/source/isaaclab_newton/isaaclab_newton/actuators/adapter.py index a328b0e0b617..0d28cf576dfb 100644 --- a/source/isaaclab_newton/isaaclab_newton/actuators/adapter.py +++ b/source/isaaclab_newton/isaaclab_newton/actuators/adapter.py @@ -23,7 +23,6 @@ import numpy as np import torch import warp as wp - from newton.actuators import Actuator, Clamping, Delay from .kernels import ( @@ -33,7 +32,6 @@ zero_at_indices_kernel, ) - # --------------------------------------------------------------------------- # Abstract base — backend-independent logic # --------------------------------------------------------------------------- @@ -89,7 +87,9 @@ def __init__( # computed (pre-clamp) effort instead of mirroring ``joint_f``. The # binding onto ``sim_control`` happens in :meth:`finalize`. self._computed_effort = wp.zeros( - num_envs * num_joints, dtype=wp.float32, device=device, + num_envs * num_joints, + dtype=wp.float32, + device=device, ) self.computed_effort_2d = self._computed_effort.reshape((num_envs, num_joints)) for act in actuators: @@ -195,7 +195,7 @@ def from_usd( num_joints: int, device: str, articulation_prim_path: str | None = None, - ) -> "NewtonActuatorAdapter": + ) -> NewtonActuatorAdapter: """Build an adapter from ``NewtonActuator`` prims authored on *stage*. PhysX-side counterpart of Newton's ``ModelBuilder.add_usd``: reads @@ -216,7 +216,11 @@ def from_usd( considered; otherwise the whole stage is scanned. """ actuators = _create_actuators_from_usd( - stage, joint_names, num_envs, num_joints, device, + stage, + joint_names, + num_envs, + num_joints, + device, articulation_prim_path=articulation_prim_path, ) return cls(actuators, num_envs, num_joints, dof_offset=0, device=device) @@ -262,10 +266,7 @@ def build_newton_actuator_defaults( the adapter's actuators. ``slice(None)`` when every joint is covered, otherwise an int32 tensor of column indices. """ - arti_actuators = [ - act for act in actuators - if dof_offset <= int(act.indices.numpy()[0]) < dof_offset + num_joints - ] + arti_actuators = [act for act in actuators if dof_offset <= int(act.indices.numpy()[0]) < dof_offset + num_joints] managed_local: set[int] = set() for act in arti_actuators: @@ -287,13 +288,15 @@ def build_newton_actuator_defaults( ctrl = act.controller if hasattr(ctrl, "kp"): wp.launch( - scatter_gain_kernel, dim=act.indices.shape[0], + scatter_gain_kernel, + dim=act.indices.shape[0], inputs=[ctrl.kp, flat_stiffness, act.indices, dof_offset, num_joints], device=wp_device, ) if hasattr(ctrl, "kd"): wp.launch( - scatter_gain_kernel, dim=act.indices.shape[0], + scatter_gain_kernel, + dim=act.indices.shape[0], inputs=[ctrl.kd, flat_damping, act.indices, dof_offset, num_joints], device=wp_device, ) @@ -361,6 +364,7 @@ class (e.g. ``model_path``, ``lookup_positions``) are passed through from collections import defaultdict # noqa: PLC0415 from newton.actuators import parse_actuator_prim # noqa: PLC0415 + from pxr import Usd # noqa: PLC0415 wp_device = wp.get_device(device) @@ -382,9 +386,7 @@ class (e.g. ``model_path``, ``lookup_positions``) are passed through parsed_per_joint[joint_name_to_idx[target_name]] = parsed if not parsed_per_joint: - raise ValueError( - f"No NewtonActuator prims found targeting any of: {joint_names}" - ) + raise ValueError(f"No NewtonActuator prims found targeting any of: {joint_names}") groups: dict[tuple, list[int]] = defaultdict(list) sig_to_parsed: dict[tuple, Any] = {} @@ -436,7 +438,10 @@ class (e.g. ``model_path``, ``lookup_positions``) are passed through clamp_arrays[k] = v else: clamp_arrays[k] = wp.full( - num_dofs_in_group, float(v), dtype=wp.float32, device=wp_device, + num_dofs_in_group, + float(v), + dtype=wp.float32, + device=wp_device, ) clampings.append(comp_cls(**clamp_arrays)) diff --git a/source/isaaclab_newton/isaaclab_newton/actuators/kernels.py b/source/isaaclab_newton/isaaclab_newton/actuators/kernels.py index e710d97ccfd6..be2994adac1c 100644 --- a/source/isaaclab_newton/isaaclab_newton/actuators/kernels.py +++ b/source/isaaclab_newton/isaaclab_newton/actuators/kernels.py @@ -10,7 +10,6 @@ from isaaclab.actuators import ActuatorBase, ImplicitActuator - # --------------------------------------------------------------------------- # Adapter / per-actuator helper kernels: per-DOF zeroing, env-mask building, # per-DOF env-mask projection (used by :meth:`NewtonActuatorAdapter.reset`), @@ -186,4 +185,3 @@ def build_implicit_dof_mask( else: modes[j_ids.long()] = 1 return wp.from_torch(modes, dtype=wp.int32) - diff --git a/source/isaaclab_newton/isaaclab_newton/actuators/physx_wrapper.py b/source/isaaclab_newton/isaaclab_newton/actuators/physx_wrapper.py index 35dadc684524..b3f48a2f9dee 100644 --- a/source/isaaclab_newton/isaaclab_newton/actuators/physx_wrapper.py +++ b/source/isaaclab_newton/isaaclab_newton/actuators/physx_wrapper.py @@ -52,7 +52,7 @@ class PhysxActuatorWrapper: joint_f_2d: wp.array | None = None @classmethod - def create(cls, num_envs: int, num_joints: int, device: str) -> "PhysxActuatorWrapper": + def create(cls, num_envs: int, num_joints: int, device: str) -> PhysxActuatorWrapper: """Allocate the persistent ``joint_f`` buffer for the given articulation shape.""" w = cls() w.joint_f_2d = wp.zeros((num_envs, num_joints), dtype=wp.float32, device=device) diff --git a/source/isaaclab_newton/isaaclab_newton/assets/articulation/articulation.py b/source/isaaclab_newton/isaaclab_newton/assets/articulation/articulation.py index 7140d5e12d40..f743dcb72b89 100644 --- a/source/isaaclab_newton/isaaclab_newton/assets/articulation/articulation.py +++ b/source/isaaclab_newton/isaaclab_newton/assets/articulation/articulation.py @@ -1555,12 +1555,15 @@ def _write_actuator_param( return env_ids_wp = wp.from_torch( - env_ids.to(self.device, dtype=torch.int32).contiguous(), dtype=wp.int32, + env_ids.to(self.device, dtype=torch.int32).contiguous(), + dtype=wp.int32, ) env_mask = wp.zeros(self.num_instances, dtype=wp.bool, device=self.device) wp.launch( actuator_kernels.set_mask_kernel, - dim=env_ids_wp.shape[0], inputs=[env_mask, env_ids_wp], device=self.device, + dim=env_ids_wp.shape[0], + inputs=[env_mask, env_ids_wp], + device=self.device, ) env_ids_long = env_ids.to(self.device, dtype=torch.long).unsqueeze(1) @@ -1574,8 +1577,11 @@ def _write_actuator_param( cur_torch = wp.to_torch(cur_wp) cur_torch[env_ids_long, joint_ids_long] = values.to(cur_torch.device, dtype=cur_torch.dtype) self._root_view.set_actuator_parameter( - actuator=act, component=ctrl, name=attr, - values=cur_wp, mask=env_mask, + actuator=act, + component=ctrl, + name=attr, + values=cur_wp, + mask=env_mask, ) def write_joint_position_limit_to_sim_index( @@ -3600,7 +3606,9 @@ def _process_actuators_cfg(self): explicit_joint_ids.extend(int(j) for j in joint_ids) if explicit_joint_ids: explicit_ids_t = torch.tensor( - sorted(set(explicit_joint_ids)), dtype=torch.int32, device=self.device, + sorted(set(explicit_joint_ids)), + dtype=torch.int32, + device=self.device, ) self.write_joint_stiffness_to_sim_index(stiffness=0.0, joint_ids=explicit_ids_t) self.write_joint_damping_to_sim_index(damping=0.0, joint_ids=explicit_ids_t) @@ -3618,13 +3626,14 @@ def _process_actuators_cfg(self): self._create_lab_actuator(actuator_name, actuator_cfg, properties_only=True) self._implicit_dof_mask = build_implicit_dof_mask( - self.actuators, self.num_joints, self.device, + self.actuators, + self.num_joints, + self.device, ) # Run the implicit-DOF FF-routing + telemetry kernel inside the # captured graph, right after the actuator step. Closure captures # the buffers we need via ``self._data``. - from newton import Model as NewtonModel # noqa: PLC0415 # Per-articulation view of the global adapter's pre-clamp # computed-effort buffer. Set up once here (the adapter is @@ -3658,7 +3667,9 @@ def _process_actuators_cfg(self): ) else: self._data._sim_bind_joint_computed_effort = wp.zeros( - (self.num_instances, self.num_joints), dtype=wp.float32, device=self.device, + (self.num_instances, self.num_joints), + dtype=wp.float32, + device=self.device, ) def _post_actuator() -> None: @@ -3716,7 +3727,11 @@ def _post_actuator() -> None: logger.warning(f"\nActuatorCfg-USD Value Discrepancy Resolution (matching values are skipped): \n{t}") def _create_lab_actuator( - self, actuator_name: str, actuator_cfg: ActuatorBaseCfg, *, properties_only: bool = False, + self, + actuator_name: str, + actuator_cfg: ActuatorBaseCfg, + *, + properties_only: bool = False, ) -> None: """Instantiate a single Lab actuator from its config and write properties to sim. @@ -3755,14 +3770,17 @@ def _create_lab_actuator( # Write physical joint properties (armature, limits, friction) — always needed. self.write_joint_effort_limit_to_sim_index( - limits=actuator.effort_limit_sim, joint_ids=actuator.joint_indices, + limits=actuator.effort_limit_sim, + joint_ids=actuator.joint_indices, ) self.write_joint_velocity_limit_to_sim_index( - limits=actuator.velocity_limit_sim, joint_ids=actuator.joint_indices, + limits=actuator.velocity_limit_sim, + joint_ids=actuator.joint_indices, ) self.write_joint_armature_to_sim_index(armature=actuator.armature, joint_ids=actuator.joint_indices) self.write_joint_friction_coefficient_to_sim_index( - joint_friction_coeff=actuator.friction, joint_ids=actuator.joint_indices, + joint_friction_coeff=actuator.friction, + joint_ids=actuator.joint_indices, ) if properties_only: diff --git a/source/isaaclab_newton/isaaclab_newton/assets/articulation/articulation_data.py b/source/isaaclab_newton/isaaclab_newton/assets/articulation/articulation_data.py index 5eae244b2a8e..26d458f17c18 100644 --- a/source/isaaclab_newton/isaaclab_newton/assets/articulation/articulation_data.py +++ b/source/isaaclab_newton/isaaclab_newton/assets/articulation/articulation_data.py @@ -1369,9 +1369,7 @@ def _create_simulation_bindings(self) -> None: self._sim_bind_joint_effort = self._root_view.get_attribute("joint_f", SimulationManager.get_control())[ :, 0 ] - self._sim_bind_joint_act = self._root_view.get_attribute("joint_act", SimulationManager.get_control())[ - :, 0 - ] + self._sim_bind_joint_act = self._root_view.get_attribute("joint_act", SimulationManager.get_control())[:, 0] self._sim_bind_joint_position_target = self._root_view.get_attribute( "joint_target_pos", SimulationManager.get_control() )[:, 0] diff --git a/source/isaaclab_newton/isaaclab_newton/physics/newton_manager.py b/source/isaaclab_newton/isaaclab_newton/physics/newton_manager.py index 77ba356234b4..6fcf61caf761 100644 --- a/source/isaaclab_newton/isaaclab_newton/physics/newton_manager.py +++ b/source/isaaclab_newton/isaaclab_newton/physics/newton_manager.py @@ -12,7 +12,7 @@ import logging from abc import abstractmethod from collections.abc import Callable -from typing import TYPE_CHECKING, Any +from typing import TYPE_CHECKING import warp as wp diff --git a/source/isaaclab_newton/test/assets/test_newton_actuators_newton.py b/source/isaaclab_newton/test/assets/test_newton_actuators_newton.py index 264481e81796..7e34a4622788 100644 --- a/source/isaaclab_newton/test/assets/test_newton_actuators_newton.py +++ b/source/isaaclab_newton/test/assets/test_newton_actuators_newton.py @@ -144,7 +144,10 @@ def _run_simulation( """ sim_cfg = SimulationCfg(dt=dt, physics=newton_cfg, use_newton_actuators=use_newton_actuators) with build_simulation_context( - device="cuda:0", gravity_enabled=True, add_ground_plane=True, sim_cfg=sim_cfg, + device="cuda:0", + gravity_enabled=True, + add_ground_plane=True, + sim_cfg=sim_cfg, ) as sim: sim._app_control_on_stop_handle = None for i in range(NUM_ENVS): @@ -241,20 +244,22 @@ def setUpClass(cls): cls.newton_result = _run_simulation(cls.actuators, use_newton_actuators=True, **kwargs) def test_joint_positions_match(self): - for step_i, (lab, newton) in enumerate( - zip(self.lab_result["joint_pos"], self.newton_result["joint_pos"]) - ): + for step_i, (lab, newton) in enumerate(zip(self.lab_result["joint_pos"], self.newton_result["joint_pos"])): torch.testing.assert_close( - lab, newton, atol=self.pos_atol, rtol=self.pos_rtol, + lab, + newton, + atol=self.pos_atol, + rtol=self.pos_rtol, msg=f"Joint positions diverged at step {step_i}", ) def test_joint_velocities_match(self): - for step_i, (lab, newton) in enumerate( - zip(self.lab_result["joint_vel"], self.newton_result["joint_vel"]) - ): + for step_i, (lab, newton) in enumerate(zip(self.lab_result["joint_vel"], self.newton_result["joint_vel"])): torch.testing.assert_close( - lab, newton, atol=self.vel_atol, rtol=self.vel_rtol, + lab, + newton, + atol=self.vel_atol, + rtol=self.vel_rtol, msg=f"Joint velocities diverged at step {step_i}", ) @@ -263,7 +268,10 @@ def test_applied_torque_match(self): zip(self.lab_result["applied_torque"], self.newton_result["applied_torque"]) ): torch.testing.assert_close( - lab, newton, atol=self.torque_atol, rtol=self.torque_rtol, + lab, + newton, + atol=self.torque_atol, + rtol=self.torque_rtol, msg=f"applied_torque diverged at step {step_i}", ) @@ -272,7 +280,10 @@ def test_computed_torque_match(self): zip(self.lab_result["computed_torque"], self.newton_result["computed_torque"]) ): torch.testing.assert_close( - lab, newton, atol=self.torque_atol, rtol=self.torque_rtol, + lab, + newton, + atol=self.torque_atol, + rtol=self.torque_rtol, msg=f"computed_torque diverged at step {step_i}", ) @@ -396,7 +407,10 @@ def _run_anymal_and_cartpole(use_newton_actuators: bool, *, num_steps: int = NUM sim_cfg = SimulationCfg(dt=DT, physics=NEWTON_CFG, use_newton_actuators=use_newton_actuators) with build_simulation_context( - device="cuda:0", gravity_enabled=True, add_ground_plane=True, sim_cfg=sim_cfg, + device="cuda:0", + gravity_enabled=True, + add_ground_plane=True, + sim_cfg=sim_cfg, ) as sim: sim._app_control_on_stop_handle = None @@ -405,7 +419,8 @@ def _run_anymal_and_cartpole(use_newton_actuators: bool, *, num_steps: int = NUM anymal_cfg = ANYMAL_C_CFG.replace(actuators=IDEAL_PD_ACTUATORS, prim_path="/World/Env_.*/Anymal") cartpole_cfg = CARTPOLE_CFG.replace( - actuators=CARTPOLE_EXPLICIT_ACTUATORS, prim_path="/World/Env_.*/Cartpole", + actuators=CARTPOLE_EXPLICIT_ACTUATORS, + prim_path="/World/Env_.*/Cartpole", ) # Stand the cartpole well clear of the anymal. cartpole_cfg.init_state = cartpole_cfg.init_state.replace(pos=(0.0, 3.0, 2.0)) @@ -457,7 +472,10 @@ def test_anymal_matches_lab(self): zip(self.lab_result["joint_pos_anymal"], self.newton_result["joint_pos_anymal"]) ): torch.testing.assert_close( - newton, lab, atol=2e-3, rtol=1e-3, + newton, + lab, + atol=2e-3, + rtol=1e-3, msg=f"ANYmal joint_pos diverged from Lab path at step {step_i}", ) @@ -466,7 +484,10 @@ def test_cartpole_matches_lab(self): zip(self.lab_result["joint_pos_cartpole"], self.newton_result["joint_pos_cartpole"]) ): torch.testing.assert_close( - newton, lab, atol=2e-3, rtol=1e-3, + newton, + lab, + atol=2e-3, + rtol=1e-3, msg=f"Cartpole joint_pos diverged from Lab path at step {step_i}", ) @@ -564,13 +585,17 @@ def _gather_param(articulation, attr) -> torch.Tensor: def test_single_articulation(self): sim_cfg = SimulationCfg(dt=DT, physics=NEWTON_CFG, use_newton_actuators=True) with build_simulation_context( - device="cuda:0", gravity_enabled=True, add_ground_plane=True, sim_cfg=sim_cfg, + device="cuda:0", + gravity_enabled=True, + add_ground_plane=True, + sim_cfg=sim_cfg, ) as sim: sim._app_control_on_stop_handle = None for i in range(NUM_ENVS): sim_utils.create_prim(f"/World/Env_{i}", "Xform", translation=(i * 3.0, 0, 0)) art_cfg = ANYMAL_C_CFG.replace( - actuators=IDEAL_PD_ACTUATORS, prim_path="/World/Env_.*/Robot", + actuators=IDEAL_PD_ACTUATORS, + prim_path="/World/Env_.*/Robot", ) anymal = Articulation(art_cfg) sim.reset() @@ -585,10 +610,13 @@ def test_single_articulation(self): env_ids = torch.tensor([0], device=anymal.device, dtype=torch.long) term( - env, env_ids=env_ids, asset_cfg=asset_cfg, + env, + env_ids=env_ids, + asset_cfg=asset_cfg, stiffness_distribution_params=(100.0, 100.0), damping_distribution_params=(5.0, 5.0), - operation="abs", distribution="uniform", + operation="abs", + distribution="uniform", ) kp_after = self._gather_param(anymal, "kp") @@ -606,7 +634,10 @@ def test_two_articulations(self): sim_cfg = SimulationCfg(dt=DT, physics=NEWTON_CFG, use_newton_actuators=True) with build_simulation_context( - device="cuda:0", gravity_enabled=True, add_ground_plane=True, sim_cfg=sim_cfg, + device="cuda:0", + gravity_enabled=True, + add_ground_plane=True, + sim_cfg=sim_cfg, ) as sim: sim._app_control_on_stop_handle = None for i in range(NUM_ENVS): @@ -614,7 +645,8 @@ def test_two_articulations(self): anymal_cfg = ANYMAL_C_CFG.replace(actuators=IDEAL_PD_ACTUATORS, prim_path="/World/Env_.*/Anymal") cartpole_cfg = CARTPOLE_CFG.replace( - actuators=CARTPOLE_EXPLICIT_ACTUATORS, prim_path="/World/Env_.*/Cartpole", + actuators=CARTPOLE_EXPLICIT_ACTUATORS, + prim_path="/World/Env_.*/Cartpole", ) cartpole_cfg.init_state = cartpole_cfg.init_state.replace(pos=(0.0, 3.0, 2.0)) anymal = Articulation(anymal_cfg) @@ -633,10 +665,13 @@ def test_two_articulations(self): env_ids = torch.tensor([0], device=anymal.device, dtype=torch.long) term( - env, env_ids=env_ids, asset_cfg=asset_cfg, + env, + env_ids=env_ids, + asset_cfg=asset_cfg, stiffness_distribution_params=(100.0, 100.0), damping_distribution_params=(5.0, 5.0), - operation="abs", distribution="uniform", + operation="abs", + distribution="uniform", ) cp_kp_after = self._gather_param(cartpole, "kp") @@ -751,7 +786,6 @@ class TestDecimationMixed(_DecimationMixin, TestMixedActuatorEquivalence): RESET_WARMUP_STEPS = 3 - class TestActuatorStateReset(unittest.TestCase): """Reset must clear the actuator state buffers for the requested envs only. @@ -771,17 +805,23 @@ class TestActuatorStateReset(unittest.TestCase): def _build_and_warm(self, *, use_newton_actuators: bool): sim_cfg = SimulationCfg( - dt=DT, physics=NEWTON_CFG, use_newton_actuators=use_newton_actuators, + dt=DT, + physics=NEWTON_CFG, + use_newton_actuators=use_newton_actuators, ) ctx = build_simulation_context( - device="cuda:0", gravity_enabled=True, add_ground_plane=True, sim_cfg=sim_cfg, + device="cuda:0", + gravity_enabled=True, + add_ground_plane=True, + sim_cfg=sim_cfg, ) sim = ctx.__enter__() sim._app_control_on_stop_handle = None for i in range(NUM_ENVS): sim_utils.create_prim(f"/World/Env_{i}", "Xform", translation=(i * 3.0, 0, 0)) art_cfg = ANYMAL_C_CFG.replace( - actuators=DELAYED_PD_ACTUATORS, prim_path="/World/Env_.*/Robot", + actuators=DELAYED_PD_ACTUATORS, + prim_path="/World/Env_.*/Robot", ) articulation = Articulation(art_cfg) sim.reset() @@ -805,7 +845,8 @@ def test_newton_state_reset_isolated_to_reset_env(self): self.assertIsNotNone(adapter) # Find a DelayedPD actuator (it's the only one with delay_state). stateful_pairs = [ - (act, st) for act, st in zip(adapter.actuators, adapter._states_a) + (act, st) + for act, st in zip(adapter.actuators, adapter._states_a) if st is not None and getattr(st, "delay_state", None) is not None ] self.assertGreater(len(stateful_pairs), 0, "expected at least one DelayedPD actuator with delay_state") @@ -832,12 +873,14 @@ def test_newton_state_reset_isolated_to_reset_env(self): env = int(global_dof) // adapter.num_joints if env == self.RESET_ENV: self.assertEqual( - int(pushes_after[i]), 0, + int(pushes_after[i]), + 0, f"DOF {i} (env {env}) should be reset to 0, got {pushes_after[i]}", ) else: self.assertGreater( - int(pushes_after[i]), 0, + int(pushes_after[i]), + 0, f"DOF {i} (env {env}) was NOT in reset env_ids but num_pushes is 0", ) finally: @@ -1009,9 +1052,7 @@ def _run_authoring_introspection(actuator_cfgs: dict) -> dict: sim._app_control_on_stop_handle = None for i in range(NUM_ENVS): - sim_utils.create_prim( - f"/World/Env_{i}", "Xform", translation=(i * 3.0, 0, 0) - ) + sim_utils.create_prim(f"/World/Env_{i}", "Xform", translation=(i * 3.0, 0, 0)) art_cfg = ANYMAL_C_CFG.replace( actuators=actuator_cfgs, @@ -1026,15 +1067,15 @@ def _run_authoring_introspection(actuator_cfgs: dict) -> dict: actuator_info = [] for act in model.actuators: ctrl_type = type(act.controller).__name__ - clamp_types = sorted( - type(c).__name__ for c in (act.clamping or []) + clamp_types = sorted(type(c).__name__ for c in (act.clamping or [])) + actuator_info.append( + { + "controller_type": ctrl_type, + "clamping_types": clamp_types, + "has_delay": act.delay is not None, + "num_indices": len(act.indices), + } ) - actuator_info.append({ - "controller_type": ctrl_type, - "clamping_types": clamp_types, - "has_delay": act.delay is not None, - "num_indices": len(act.indices), - }) init_pos = wp.to_torch(articulation.data.joint_pos).clone() target_pos = init_pos + TARGET_OFFSET @@ -1068,47 +1109,40 @@ class TestRemotizedPDAuthoring(unittest.TestCase): def setUpClass(cls): from isaaclab.actuators.actuator_pd_cfg import RemotizedPDActuatorCfg # noqa: PLC0415 - cls.result = _run_authoring_introspection({ - "hips": IdealPDActuatorCfg( - joint_names_expr=[".*HAA", ".*HFE"], - stiffness=40.0, - damping=5.0, - effort_limit=80.0, - ), - "knees": RemotizedPDActuatorCfg( - joint_names_expr=[".*KFE"], - stiffness=60.0, - damping=1.5, - effort_limit=80.0, - max_delay=3, - joint_parameter_lookup=SPOT_KNEE_LOOKUP, - ), - }) + cls.result = _run_authoring_introspection( + { + "hips": IdealPDActuatorCfg( + joint_names_expr=[".*HAA", ".*HFE"], + stiffness=40.0, + damping=5.0, + effort_limit=80.0, + ), + "knees": RemotizedPDActuatorCfg( + joint_names_expr=[".*KFE"], + stiffness=60.0, + damping=1.5, + effort_limit=80.0, + max_delay=3, + joint_parameter_lookup=SPOT_KNEE_LOOKUP, + ), + } + ) def test_num_actuators(self): self.assertGreaterEqual(self.result["num_actuators"], 2) def test_kfe_controller_is_pd(self): - kfe_acts = [ - a for a in self.result["actuator_info"] - if "ClampingPositionBased" in a["clamping_types"] - ] + kfe_acts = [a for a in self.result["actuator_info"] if "ClampingPositionBased" in a["clamping_types"]] self.assertTrue(len(kfe_acts) > 0, "No actuator with position-based clamping found") for a in kfe_acts: self.assertEqual(a["controller_type"], "ControllerPD") def test_kfe_has_position_based_clamping(self): - kfe_acts = [ - a for a in self.result["actuator_info"] - if "ClampingPositionBased" in a["clamping_types"] - ] + kfe_acts = [a for a in self.result["actuator_info"] if "ClampingPositionBased" in a["clamping_types"]] self.assertTrue(len(kfe_acts) > 0, "Position-based clamping not found") def test_kfe_has_delay(self): - kfe_acts = [ - a for a in self.result["actuator_info"] - if "ClampingPositionBased" in a["clamping_types"] - ] + kfe_acts = [a for a in self.result["actuator_info"] if "ClampingPositionBased" in a["clamping_types"]] for a in kfe_acts: self.assertTrue(a["has_delay"], "Delay not found on remotized KFE actuator") @@ -1157,24 +1191,30 @@ def _make_dummy_mlp_checkpoint(device: str = "cpu") -> str: in pos_vel order) and outputs 1 effort. """ torch.manual_seed(42) - net = torch.nn.Sequential( - torch.nn.Linear(6, 8), - torch.nn.ELU(), - torch.nn.Linear(8, 1), - ).to(device).eval() + net = ( + torch.nn.Sequential( + torch.nn.Linear(6, 8), + torch.nn.ELU(), + torch.nn.Linear(8, 1), + ) + .to(device) + .eval() + ) scripted = torch.jit.script(net) with tempfile.NamedTemporaryFile(suffix=".pt", delete=False) as tmp: tmp_path = tmp.name extra = { - "metadata.json": json.dumps({ - "model_type": "mlp", - "input_order": "pos_vel", - "input_idx": [0, 1, 2], - "pos_scale": 1.0, - "vel_scale": 0.5, - "torque_scale": 2.0, - }) + "metadata.json": json.dumps( + { + "model_type": "mlp", + "input_order": "pos_vel", + "input_idx": [0, 1, 2], + "pos_scale": 1.0, + "vel_scale": 0.5, + "torque_scale": 2.0, + } + ) } torch.jit.save(scripted, tmp_path, _extra_files=extra) return tmp_path @@ -1220,26 +1260,28 @@ def setUpClass(cls): from isaaclab.actuators.actuator_net_cfg import ActuatorNetMLPCfg # noqa: PLC0415 cls.mlp_path = _make_dummy_mlp_checkpoint() - cls.result = _run_authoring_introspection({ - "mlp_legs": ActuatorNetMLPCfg( - joint_names_expr=[".*HAA"], - network_file=cls.mlp_path, - saturation_effort=120.0, - effort_limit=80.0, - velocity_limit=7.5, - pos_scale=-1.0, - vel_scale=1.0, - torque_scale=1.0, - input_order="pos_vel", - input_idx=[0, 1, 2], - ), - "pd_legs": IdealPDActuatorCfg( - joint_names_expr=[".*HFE", ".*KFE"], - stiffness=40.0, - damping=5.0, - effort_limit=80.0, - ), - }) + cls.result = _run_authoring_introspection( + { + "mlp_legs": ActuatorNetMLPCfg( + joint_names_expr=[".*HAA"], + network_file=cls.mlp_path, + saturation_effort=120.0, + effort_limit=80.0, + velocity_limit=7.5, + pos_scale=-1.0, + vel_scale=1.0, + torque_scale=1.0, + input_order="pos_vel", + input_idx=[0, 1, 2], + ), + "pd_legs": IdealPDActuatorCfg( + joint_names_expr=[".*HFE", ".*KFE"], + stiffness=40.0, + damping=5.0, + effort_limit=80.0, + ), + } + ) @classmethod def tearDownClass(cls): @@ -1249,17 +1291,11 @@ def test_num_actuators(self): self.assertGreaterEqual(self.result["num_actuators"], 2) def test_has_neural_mlp_controller(self): - mlp_acts = [ - a for a in self.result["actuator_info"] - if a["controller_type"] == "ControllerNeuralMLP" - ] + mlp_acts = [a for a in self.result["actuator_info"] if a["controller_type"] == "ControllerNeuralMLP"] self.assertTrue(len(mlp_acts) > 0, "No NeuralMLP controller found") def test_mlp_has_dc_motor_clamping(self): - mlp_acts = [ - a for a in self.result["actuator_info"] - if a["controller_type"] == "ControllerNeuralMLP" - ] + mlp_acts = [a for a in self.result["actuator_info"] if a["controller_type"] == "ControllerNeuralMLP"] for a in mlp_acts: self.assertIn("ClampingDCMotor", a["clamping_types"]) @@ -1274,21 +1310,23 @@ def setUpClass(cls): from isaaclab.actuators.actuator_net_cfg import ActuatorNetLSTMCfg # noqa: PLC0415 cls.lstm_path = _make_dummy_lstm_checkpoint() - cls.result = _run_authoring_introspection({ - "lstm_legs": ActuatorNetLSTMCfg( - joint_names_expr=[".*HAA"], - network_file=cls.lstm_path, - saturation_effort=120.0, - effort_limit=80.0, - velocity_limit=7.5, - ), - "pd_legs": IdealPDActuatorCfg( - joint_names_expr=[".*HFE", ".*KFE"], - stiffness=40.0, - damping=5.0, - effort_limit=80.0, - ), - }) + cls.result = _run_authoring_introspection( + { + "lstm_legs": ActuatorNetLSTMCfg( + joint_names_expr=[".*HAA"], + network_file=cls.lstm_path, + saturation_effort=120.0, + effort_limit=80.0, + velocity_limit=7.5, + ), + "pd_legs": IdealPDActuatorCfg( + joint_names_expr=[".*HFE", ".*KFE"], + stiffness=40.0, + damping=5.0, + effort_limit=80.0, + ), + } + ) @classmethod def tearDownClass(cls): @@ -1298,17 +1336,11 @@ def test_num_actuators(self): self.assertGreaterEqual(self.result["num_actuators"], 2) def test_has_neural_lstm_controller(self): - lstm_acts = [ - a for a in self.result["actuator_info"] - if a["controller_type"] == "ControllerNeuralLSTM" - ] + lstm_acts = [a for a in self.result["actuator_info"] if a["controller_type"] == "ControllerNeuralLSTM"] self.assertTrue(len(lstm_acts) > 0, "No NeuralLSTM controller found") def test_lstm_has_dc_motor_clamping(self): - lstm_acts = [ - a for a in self.result["actuator_info"] - if a["controller_type"] == "ControllerNeuralLSTM" - ] + lstm_acts = [a for a in self.result["actuator_info"] if a["controller_type"] == "ControllerNeuralLSTM"] for a in lstm_acts: self.assertIn("ClampingDCMotor", a["clamping_types"]) diff --git a/source/isaaclab_physx/isaaclab_physx/assets/articulation/articulation.py b/source/isaaclab_physx/isaaclab_physx/assets/articulation/articulation.py index f55eb78a6da5..fa2eca5ccde6 100644 --- a/source/isaaclab_physx/isaaclab_physx/assets/articulation/articulation.py +++ b/source/isaaclab_physx/isaaclab_physx/assets/articulation/articulation.py @@ -283,7 +283,8 @@ def write_data_to_sim(self): # to PhysX as the actuation force. self._apply_actuator_model_newton() self.root_view.set_dof_actuation_forces( - self._physx_actuator_wrapper.joint_f_2d, self._ALL_INDICES, + self._physx_actuator_wrapper.joint_f_2d, + self._ALL_INDICES, ) if self._has_implicit_actuators: self.root_view.set_dof_position_targets(self._data._joint_pos_target, self._ALL_INDICES) @@ -1254,20 +1255,31 @@ def _write_actuator_param( return env_id_pos = torch.full( - (self.num_instances,), -1, dtype=torch.int32, device=self.device, + (self.num_instances,), + -1, + dtype=torch.int32, + device=self.device, ) env_id_pos[env_ids.to(self.device, dtype=torch.long)] = torch.arange( - env_ids.shape[0], dtype=torch.int32, device=self.device, + env_ids.shape[0], + dtype=torch.int32, + device=self.device, ) joint_id_pos = torch.full( - (self.num_joints,), -1, dtype=torch.int32, device=self.device, + (self.num_joints,), + -1, + dtype=torch.int32, + device=self.device, ) joint_id_pos[joint_ids.to(self.device, dtype=torch.long)] = torch.arange( - joint_ids.shape[0], dtype=torch.int32, device=self.device, + joint_ids.shape[0], + dtype=torch.int32, + device=self.device, ) values_wp = wp.from_torch( - values.to(self.device, dtype=torch.float32).contiguous(), dtype=wp.float32, + values.to(self.device, dtype=torch.float32).contiguous(), + dtype=wp.float32, ) env_id_pos_wp = wp.from_torch(env_id_pos, dtype=wp.int32) joint_id_pos_wp = wp.from_torch(joint_id_pos, dtype=wp.int32) @@ -1280,8 +1292,12 @@ def _write_actuator_param( actuator_kernels.patch_actuator_param_kernel, dim=act.indices.shape[0], inputs=[ - act.indices, env_id_pos_wp, joint_id_pos_wp, values_wp, - 0, self.num_joints, + act.indices, + env_id_pos_wp, + joint_id_pos_wp, + values_wp, + 0, + self.num_joints, ], outputs=[getattr(ctrl, attr)], device=self.device, @@ -3941,7 +3957,9 @@ def _process_actuators_cfg(self): self._create_lab_actuator(actuator_name, actuator_cfg, properties_only=True) self._implicit_dof_mask = build_implicit_dof_mask( - self.actuators, self.num_joints, self.device, + self.actuators, + self.num_joints, + self.device, ) # Per-articulation view of the adapter's pre-clamp computed-effort # buffer (or zero fallback when there are no explicit Newton @@ -3951,7 +3969,9 @@ def _process_actuators_cfg(self): self._data._sim_bind_joint_computed_effort = self.newton_actuator_adapter.computed_effort_2d else: self._data._sim_bind_joint_computed_effort = wp.zeros( - (self.num_instances, self.num_joints), dtype=wp.float32, device=self.device, + (self.num_instances, self.num_joints), + dtype=wp.float32, + device=self.device, ) return @@ -3983,7 +4003,11 @@ def _process_actuators_cfg(self): logger.warning(f"\nActuatorCfg-USD Value Discrepancy Resolution (matching values are skipped): \n{t}") def _create_lab_actuator( - self, actuator_name: str, actuator_cfg: ActuatorBaseCfg, *, properties_only: bool = False, + self, + actuator_name: str, + actuator_cfg: ActuatorBaseCfg, + *, + properties_only: bool = False, ) -> None: """Instantiate a single Lab actuator from its config and write properties to sim. @@ -4024,20 +4048,25 @@ def _create_lab_actuator( # Write physical joint properties (armature, limits, friction) — always needed. self.write_joint_effort_limit_to_sim_index( - limits=actuator.effort_limit_sim, joint_ids=actuator.joint_indices, + limits=actuator.effort_limit_sim, + joint_ids=actuator.joint_indices, ) self.write_joint_velocity_limit_to_sim_index( - limits=actuator.velocity_limit_sim, joint_ids=actuator.joint_indices, + limits=actuator.velocity_limit_sim, + joint_ids=actuator.joint_indices, ) self.write_joint_armature_to_sim_index(armature=actuator.armature, joint_ids=actuator.joint_indices) self.write_joint_friction_coefficient_to_sim_index( - joint_friction_coeff=actuator.friction, joint_ids=actuator.joint_indices, + joint_friction_coeff=actuator.friction, + joint_ids=actuator.joint_indices, ) self.write_joint_dynamic_friction_coefficient_to_sim_index( - joint_dynamic_friction_coeff=actuator.dynamic_friction, joint_ids=actuator.joint_indices, + joint_dynamic_friction_coeff=actuator.dynamic_friction, + joint_ids=actuator.joint_indices, ) self.write_joint_viscous_friction_coefficient_to_sim_index( - joint_viscous_friction_coeff=actuator.viscous_friction, joint_ids=actuator.joint_indices, + joint_viscous_friction_coeff=actuator.viscous_friction, + joint_ids=actuator.joint_indices, ) if properties_only: diff --git a/source/isaaclab_physx/test/assets/test_newton_actuators_physx.py b/source/isaaclab_physx/test/assets/test_newton_actuators_physx.py index e91e74de52dc..8e8d2ea134c4 100644 --- a/source/isaaclab_physx/test/assets/test_newton_actuators_physx.py +++ b/source/isaaclab_physx/test/assets/test_newton_actuators_physx.py @@ -137,7 +137,10 @@ def _run_simulation( """ sim_cfg = SimulationCfg(dt=DT, physics=PhysxCfg(), use_newton_actuators=use_newton_actuators) with build_simulation_context( - device="cuda:0", gravity_enabled=True, add_ground_plane=True, sim_cfg=sim_cfg, + device="cuda:0", + gravity_enabled=True, + add_ground_plane=True, + sim_cfg=sim_cfg, ) as sim: sim._app_control_on_stop_handle = None for i in range(NUM_ENVS): @@ -204,27 +207,33 @@ class _EquivalenceTestBase(unittest.TestCase): @classmethod def setUpClass(cls): cls.lab_result = _run_simulation( - cls.actuators, use_newton_actuators=False, feedforward=cls.feedforward, + cls.actuators, + use_newton_actuators=False, + feedforward=cls.feedforward, ) cls.newton_result = _run_simulation( - cls.actuators, use_newton_actuators=True, feedforward=cls.feedforward, + cls.actuators, + use_newton_actuators=True, + feedforward=cls.feedforward, ) def test_joint_positions_match(self): - for step_i, (lab, newton) in enumerate( - zip(self.lab_result["joint_pos"], self.newton_result["joint_pos"]) - ): + for step_i, (lab, newton) in enumerate(zip(self.lab_result["joint_pos"], self.newton_result["joint_pos"])): torch.testing.assert_close( - lab, newton, atol=self.pos_atol, rtol=self.pos_rtol, + lab, + newton, + atol=self.pos_atol, + rtol=self.pos_rtol, msg=f"Joint positions diverged at step {step_i}", ) def test_joint_velocities_match(self): - for step_i, (lab, newton) in enumerate( - zip(self.lab_result["joint_vel"], self.newton_result["joint_vel"]) - ): + for step_i, (lab, newton) in enumerate(zip(self.lab_result["joint_vel"], self.newton_result["joint_vel"])): torch.testing.assert_close( - lab, newton, atol=self.vel_atol, rtol=self.vel_rtol, + lab, + newton, + atol=self.vel_atol, + rtol=self.vel_rtol, msg=f"Joint velocities diverged at step {step_i}", ) @@ -233,7 +242,10 @@ def test_applied_torque_match(self): zip(self.lab_result["applied_torque"], self.newton_result["applied_torque"]) ): torch.testing.assert_close( - lab, newton, atol=self.torque_atol, rtol=self.torque_rtol, + lab, + newton, + atol=self.torque_atol, + rtol=self.torque_rtol, msg=f"applied_torque diverged at step {step_i}", ) @@ -242,7 +254,10 @@ def test_computed_torque_match(self): zip(self.lab_result["computed_torque"], self.newton_result["computed_torque"]) ): torch.testing.assert_close( - lab, newton, atol=self.torque_atol, rtol=self.torque_rtol, + lab, + newton, + atol=self.torque_atol, + rtol=self.torque_rtol, msg=f"computed_torque diverged at step {step_i}", ) @@ -344,7 +359,10 @@ def _run_anymal_and_cartpole(use_newton_actuators: bool, *, num_steps: int = NUM sim_cfg = SimulationCfg(dt=DT, physics=PhysxCfg(), use_newton_actuators=use_newton_actuators) with build_simulation_context( - device="cuda:0", gravity_enabled=True, add_ground_plane=True, sim_cfg=sim_cfg, + device="cuda:0", + gravity_enabled=True, + add_ground_plane=True, + sim_cfg=sim_cfg, ) as sim: sim._app_control_on_stop_handle = None @@ -353,7 +371,8 @@ def _run_anymal_and_cartpole(use_newton_actuators: bool, *, num_steps: int = NUM anymal_cfg = ANYMAL_C_CFG.replace(actuators=IDEAL_PD_ACTUATORS, prim_path="/World/Env_.*/Anymal") cartpole_cfg = CARTPOLE_CFG.replace( - actuators=CARTPOLE_EXPLICIT_ACTUATORS, prim_path="/World/Env_.*/Cartpole", + actuators=CARTPOLE_EXPLICIT_ACTUATORS, + prim_path="/World/Env_.*/Cartpole", ) cartpole_cfg.init_state = cartpole_cfg.init_state.replace(pos=(0.0, 3.0, 2.0)) @@ -403,7 +422,10 @@ def test_anymal_matches_lab(self): zip(self.lab_result["joint_pos_anymal"], self.newton_result["joint_pos_anymal"]) ): torch.testing.assert_close( - newton, lab, atol=2e-3, rtol=1e-3, + newton, + lab, + atol=2e-3, + rtol=1e-3, msg=f"ANYmal joint_pos diverged from Lab path at step {step_i}", ) @@ -412,7 +434,10 @@ def test_cartpole_matches_lab(self): zip(self.lab_result["joint_pos_cartpole"], self.newton_result["joint_pos_cartpole"]) ): torch.testing.assert_close( - newton, lab, atol=2e-3, rtol=1e-3, + newton, + lab, + atol=2e-3, + rtol=1e-3, msg=f"Cartpole joint_pos diverged from Lab path at step {step_i}", ) @@ -500,13 +525,17 @@ def _gather_param(adapter, num_envs, num_joints, attr, device): def test_single_articulation(self): sim_cfg = SimulationCfg(dt=DT, physics=PhysxCfg(), use_newton_actuators=True) with build_simulation_context( - device="cuda:0", gravity_enabled=True, add_ground_plane=True, sim_cfg=sim_cfg, + device="cuda:0", + gravity_enabled=True, + add_ground_plane=True, + sim_cfg=sim_cfg, ) as sim: sim._app_control_on_stop_handle = None for i in range(NUM_ENVS): sim_utils.create_prim(f"/World/Env_{i}", "Xform", translation=(i * 3.0, 0, 0)) art_cfg = ANYMAL_C_CFG.replace( - actuators=IDEAL_PD_ACTUATORS, prim_path="/World/Env_.*/Robot", + actuators=IDEAL_PD_ACTUATORS, + prim_path="/World/Env_.*/Robot", ) anymal = Articulation(art_cfg) sim.reset() @@ -522,10 +551,13 @@ def test_single_articulation(self): env_ids = torch.tensor([0], device=anymal.device, dtype=torch.long) term( - env, env_ids=env_ids, asset_cfg=asset_cfg, + env, + env_ids=env_ids, + asset_cfg=asset_cfg, stiffness_distribution_params=(100.0, 100.0), damping_distribution_params=(5.0, 5.0), - operation="abs", distribution="uniform", + operation="abs", + distribution="uniform", ) kp_after = self._gather_param(adapter, NUM_ENVS, n, "kp", anymal.device) @@ -541,7 +573,10 @@ def test_two_articulations(self): sim_cfg = SimulationCfg(dt=DT, physics=PhysxCfg(), use_newton_actuators=True) with build_simulation_context( - device="cuda:0", gravity_enabled=True, add_ground_plane=True, sim_cfg=sim_cfg, + device="cuda:0", + gravity_enabled=True, + add_ground_plane=True, + sim_cfg=sim_cfg, ) as sim: sim._app_control_on_stop_handle = None for i in range(NUM_ENVS): @@ -549,7 +584,8 @@ def test_two_articulations(self): anymal_cfg = ANYMAL_C_CFG.replace(actuators=IDEAL_PD_ACTUATORS, prim_path="/World/Env_.*/Anymal") cartpole_cfg = CARTPOLE_CFG.replace( - actuators=CARTPOLE_EXPLICIT_ACTUATORS, prim_path="/World/Env_.*/Cartpole", + actuators=CARTPOLE_EXPLICIT_ACTUATORS, + prim_path="/World/Env_.*/Cartpole", ) cartpole_cfg.init_state = cartpole_cfg.init_state.replace(pos=(0.0, 3.0, 2.0)) anymal = Articulation(anymal_cfg) @@ -575,10 +611,13 @@ def test_two_articulations(self): env_ids = torch.tensor([0], device=anymal.device, dtype=torch.long) term( - env, env_ids=env_ids, asset_cfg=asset_cfg, + env, + env_ids=env_ids, + asset_cfg=asset_cfg, stiffness_distribution_params=(100.0, 100.0), damping_distribution_params=(5.0, 5.0), - operation="abs", distribution="uniform", + operation="abs", + distribution="uniform", ) cp_kp_after = self._gather_param(cartpole_adapter, NUM_ENVS, n_cp, "kp", anymal.device) @@ -603,7 +642,6 @@ def test_two_articulations(self): RESET_WARMUP_STEPS = 3 - class TestActuatorStateReset(unittest.TestCase): """Reset must clear the actuator state buffers for the requested envs only. @@ -622,17 +660,23 @@ class TestActuatorStateReset(unittest.TestCase): def _build_and_warm(self, *, use_newton_actuators: bool): sim_cfg = SimulationCfg( - dt=DT, physics=PhysxCfg(), use_newton_actuators=use_newton_actuators, + dt=DT, + physics=PhysxCfg(), + use_newton_actuators=use_newton_actuators, ) ctx = build_simulation_context( - device="cuda:0", gravity_enabled=True, add_ground_plane=True, sim_cfg=sim_cfg, + device="cuda:0", + gravity_enabled=True, + add_ground_plane=True, + sim_cfg=sim_cfg, ) sim = ctx.__enter__() sim._app_control_on_stop_handle = None for i in range(NUM_ENVS): sim_utils.create_prim(f"/World/Env_{i}", "Xform", translation=(i * 3.0, 0, 0)) art_cfg = ANYMAL_C_CFG.replace( - actuators=DELAYED_PD_ACTUATORS, prim_path="/World/Env_.*/Robot", + actuators=DELAYED_PD_ACTUATORS, + prim_path="/World/Env_.*/Robot", ) articulation = Articulation(art_cfg) sim.reset() @@ -655,7 +699,8 @@ def test_newton_state_reset_isolated_to_reset_env(self): adapter = articulation.newton_actuator_adapter self.assertIsNotNone(adapter) stateful_pairs = [ - (act, st) for act, st in zip(adapter.actuators, adapter._states_a) + (act, st) + for act, st in zip(adapter.actuators, adapter._states_a) if st is not None and getattr(st, "delay_state", None) is not None ] self.assertGreater(len(stateful_pairs), 0, "expected at least one DelayedPD actuator with delay_state") @@ -679,12 +724,14 @@ def test_newton_state_reset_isolated_to_reset_env(self): env = int(global_dof) // adapter.num_joints if env == self.RESET_ENV: self.assertEqual( - int(pushes_after[i]), 0, + int(pushes_after[i]), + 0, f"DOF {i} (env {env}) should be reset to 0, got {pushes_after[i]}", ) else: self.assertGreater( - int(pushes_after[i]), 0, + int(pushes_after[i]), + 0, f"DOF {i} (env {env}) was NOT in reset env_ids but num_pushes is 0", ) finally: @@ -907,24 +954,30 @@ def test_positions_finite(self): def _make_dummy_mlp_checkpoint(device: str = "cpu") -> str: """Create a minimal TorchScript MLP checkpoint with metadata.""" torch.manual_seed(42) - net = torch.nn.Sequential( - torch.nn.Linear(6, 8), - torch.nn.ELU(), - torch.nn.Linear(8, 1), - ).to(device).eval() + net = ( + torch.nn.Sequential( + torch.nn.Linear(6, 8), + torch.nn.ELU(), + torch.nn.Linear(8, 1), + ) + .to(device) + .eval() + ) scripted = torch.jit.script(net) with tempfile.NamedTemporaryFile(suffix=".pt", delete=False) as tmp: tmp_path = tmp.name extra = { - "metadata.json": json.dumps({ - "model_type": "mlp", - "input_order": "pos_vel", - "input_idx": [0, 1, 2], - "pos_scale": 1.0, - "vel_scale": 0.5, - "torque_scale": 2.0, - }) + "metadata.json": json.dumps( + { + "model_type": "mlp", + "input_order": "pos_vel", + "input_idx": [0, 1, 2], + "pos_scale": 1.0, + "vel_scale": 0.5, + "torque_scale": 2.0, + } + ) } torch.jit.save(scripted, tmp_path, _extra_files=extra) return tmp_path From 97f7491c4972233c723fe903f6f4901f7dedfdb6 Mon Sep 17 00:00:00 2001 From: Julia von Muralt Date: Tue, 12 May 2026 12:27:25 +0200 Subject: [PATCH 27/35] fix --- .../isaaclab_physx/assets/articulation/articulation.py | 10 ---------- 1 file changed, 10 deletions(-) diff --git a/source/isaaclab_physx/isaaclab_physx/assets/articulation/articulation.py b/source/isaaclab_physx/isaaclab_physx/assets/articulation/articulation.py index fa2eca5ccde6..58c38d12fa76 100644 --- a/source/isaaclab_physx/isaaclab_physx/assets/articulation/articulation.py +++ b/source/isaaclab_physx/isaaclab_physx/assets/articulation/articulation.py @@ -486,8 +486,6 @@ def write_root_link_pose_to_sim_index( ], outputs=[ self.data.root_link_pose_w, - None, # self.data._root_link_state_w.data, - None, # self.data._root_state_w.data, ], device=self.device, ) @@ -576,9 +574,6 @@ def write_root_com_pose_to_sim_index( outputs=[ self.data.root_com_pose_w, self.data.root_link_pose_w, - None, # self.data._root_com_state_w.data, - None, # self.data._root_link_state_w.data, - None, # self.data._root_state_w.data, ], device=self.device, ) @@ -723,8 +718,6 @@ def write_root_com_velocity_to_sim_index( outputs=[ self.data.root_com_vel_w, self.data.body_com_acc_w, - None, # self.data._root_state_w.data, - None, # self.data._root_com_state_w.data, ], device=self.device, ) @@ -814,9 +807,6 @@ def write_root_link_velocity_to_sim_index( self.data.root_link_vel_w, self.data.root_com_vel_w, self.data.body_com_acc_w, - None, # self.data._root_link_state_w.data, - None, # self.data._root_state_w.data, - None, # self.data._root_com_state_w.data, ], device=self.device, ) From 48fe2d737bcb7f95a2d935899d4ac9f3edde3530 Mon Sep 17 00:00:00 2001 From: Julia von Muralt Date: Tue, 12 May 2026 16:57:26 +0200 Subject: [PATCH 28/35] fix --- .../isaaclab_newton/actuators/kernels.py | 14 ++++++++++++-- .../assets/articulation/articulation.py | 6 +++++- .../isaaclab_newton/physics/newton_manager.py | 13 ++++++++++++- .../assets/articulation/articulation.py | 6 +++++- 4 files changed, 34 insertions(+), 5 deletions(-) diff --git a/source/isaaclab_newton/isaaclab_newton/actuators/kernels.py b/source/isaaclab_newton/isaaclab_newton/actuators/kernels.py index be2994adac1c..8205773b640a 100644 --- a/source/isaaclab_newton/isaaclab_newton/actuators/kernels.py +++ b/source/isaaclab_newton/isaaclab_newton/actuators/kernels.py @@ -169,11 +169,21 @@ def build_implicit_dof_mask( actuators: dict[str, ActuatorBase], num_joints: int, device: str, -) -> wp.array: +) -> tuple[wp.array, torch.Tensor]: """Per-DOF mask consumed by :func:`sync_torque_telemetry`. Entry is ``1`` for DOFs covered by an :class:`~isaaclab.actuators.ImplicitActuator` group, ``0`` otherwise. + + Returns: + Tuple of ``(wp_mask, torch_owner)``. ``wp_mask`` is the Warp + view used by the kernel; ``torch_owner`` is the underlying + :class:`torch.Tensor` whose GPU memory ``wp_mask`` aliases. The + caller **must keep a reference to** ``torch_owner`` for the + Warp view's lifetime — otherwise the torch refcount drops to + zero, the memory becomes eligible for reallocation by the + caching allocator, and any captured CUDA graph that baked in + ``wp_mask``'s device pointer will read garbage at replay time. """ modes = torch.zeros(num_joints, dtype=torch.int32, device=device) for actuator in actuators.values(): @@ -184,4 +194,4 @@ def build_implicit_dof_mask( modes[:] = 1 else: modes[j_ids.long()] = 1 - return wp.from_torch(modes, dtype=wp.int32) + return wp.from_torch(modes, dtype=wp.int32), modes diff --git a/source/isaaclab_newton/isaaclab_newton/assets/articulation/articulation.py b/source/isaaclab_newton/isaaclab_newton/assets/articulation/articulation.py index f743dcb72b89..f1a15399e014 100644 --- a/source/isaaclab_newton/isaaclab_newton/assets/articulation/articulation.py +++ b/source/isaaclab_newton/isaaclab_newton/assets/articulation/articulation.py @@ -3625,7 +3625,11 @@ def _process_actuators_cfg(self): else: self._create_lab_actuator(actuator_name, actuator_cfg, properties_only=True) - self._implicit_dof_mask = build_implicit_dof_mask( + # ``_implicit_dof_mask_owner`` is the underlying torch tensor that owns + # the GPU memory aliased by ``_implicit_dof_mask``. We keep it as an + # instance attribute so the memory isn't freed while a CUDA graph + # holds a captured pointer into it. + self._implicit_dof_mask, self._implicit_dof_mask_owner = build_implicit_dof_mask( self.actuators, self.num_joints, self.device, diff --git a/source/isaaclab_newton/isaaclab_newton/physics/newton_manager.py b/source/isaaclab_newton/isaaclab_newton/physics/newton_manager.py index 6fcf61caf761..34383d31fa6e 100644 --- a/source/isaaclab_newton/isaaclab_newton/physics/newton_manager.py +++ b/source/isaaclab_newton/isaaclab_newton/physics/newton_manager.py @@ -1077,7 +1077,18 @@ def initialize_solver(cls) -> None: if cls._usdrt_stage is not None: cls._setup_cubric_bindings() - cls._capture_or_defer_graph() + # Skip the initial graph capture when the Newton actuator fast path is + # active. Capturing here would use ``cls._decimation`` (still its default + # of 1, because the env's ``set_decimation`` hasn't run yet); a second + # capture from ``set_decimation`` then triggers an illegal-memory-access + # CUDA fault inside the captured ``_simulate_full`` graph (back-to-back + # captures of the contact + actuator pipeline don't survive re-capture + # — root cause is in Newton's collision/actuator buffer handling, not + # Lab code). For non-Newton-actuator paths this branch is unaffected: + # ``set_decimation`` is a no-op for them (``_is_all_graphable`` is False), + # so we still need the start-time capture below. + if not cls._use_newton_actuators_active: + cls._capture_or_defer_graph() @classmethod def _setup_cubric_bindings(cls) -> None: diff --git a/source/isaaclab_physx/isaaclab_physx/assets/articulation/articulation.py b/source/isaaclab_physx/isaaclab_physx/assets/articulation/articulation.py index 58c38d12fa76..6ab5403f8755 100644 --- a/source/isaaclab_physx/isaaclab_physx/assets/articulation/articulation.py +++ b/source/isaaclab_physx/isaaclab_physx/assets/articulation/articulation.py @@ -3946,7 +3946,11 @@ def _process_actuators_cfg(self): else: self._create_lab_actuator(actuator_name, actuator_cfg, properties_only=True) - self._implicit_dof_mask = build_implicit_dof_mask( + # ``_implicit_dof_mask_owner`` is the underlying torch tensor that owns + # the GPU memory aliased by ``_implicit_dof_mask``. We keep it as an + # instance attribute so the memory isn't freed while a CUDA graph + # holds a captured pointer into it. + self._implicit_dof_mask, self._implicit_dof_mask_owner = build_implicit_dof_mask( self.actuators, self.num_joints, self.device, From 1009239f43d075cddf7f767bb78ed3d69caa6075 Mon Sep 17 00:00:00 2001 From: Julia von Muralt Date: Wed, 13 May 2026 00:11:05 +0200 Subject: [PATCH 29/35] bug fix in authoring --- .../assets/articulation/articulation_cfg.py | 9 +++- .../assets/test_newton_actuators_newton.py | 45 +++++++++++++++++++ 2 files changed, 53 insertions(+), 1 deletion(-) diff --git a/source/isaaclab/isaaclab/assets/articulation/articulation_cfg.py b/source/isaaclab/isaaclab/assets/articulation/articulation_cfg.py index 240c0927af12..33eb948833d6 100644 --- a/source/isaaclab/isaaclab/assets/articulation/articulation_cfg.py +++ b/source/isaaclab/isaaclab/assets/articulation/articulation_cfg.py @@ -88,4 +88,11 @@ def _post_spawn(self, stage: Any) -> None: return from isaaclab.sim.schemas.schemas_actuators import define_actuator_properties # noqa: PLC0415 - define_actuator_properties(self.prim_path, self.actuators, stage=stage) + # In InteractiveScene, articulated assets are often spawned first under + # a template path (for example ``/World/template/Robot``) and cloned + # into ``{ENV_REGEX_NS}`` later. Author NewtonActuator prims on the + # actual spawned source prim so clones inherit them. + author_prim_path = ( + self.spawn.spawn_path if self.spawn is not None and self.spawn.spawn_path is not None else self.prim_path + ) + define_actuator_properties(author_prim_path, self.actuators, stage=stage) diff --git a/source/isaaclab_newton/test/assets/test_newton_actuators_newton.py b/source/isaaclab_newton/test/assets/test_newton_actuators_newton.py index 7e34a4622788..594a6746a43a 100644 --- a/source/isaaclab_newton/test/assets/test_newton_actuators_newton.py +++ b/source/isaaclab_newton/test/assets/test_newton_actuators_newton.py @@ -35,8 +35,12 @@ import isaaclab.sim as sim_utils from isaaclab.actuators import DCMotorCfg, DelayedPDActuatorCfg, IdealPDActuatorCfg, ImplicitActuatorCfg +from isaaclab.envs import ManagerBasedRLEnv from isaaclab.sim import SimulationCfg, build_simulation_context +import isaaclab_tasks # noqa: F401 +from isaaclab_tasks.manager_based.locomotion.velocity.config.g1.flat_env_cfg import G1FlatEnvCfg + from isaaclab_assets import ANYMAL_C_CFG # --------------------------------------------------------------------------- @@ -1179,6 +1183,47 @@ class TestDecimationRemotizedPD(_DecimationMixin, TestRemotizedPDEquivalence): """RemotizedPD — decimation=2 + CUDA graph.""" +class TestManagerBasedSceneNewtonActuatorAuthoring(unittest.TestCase): + """Regression test for Newton actuator authoring in manager-based clone paths.""" + + def test_newton_actuators_present_for_g1_manager_env(self): + env_cfg = G1FlatEnvCfg() + env_cfg.scene.num_envs = 1 + env_cfg.decimation = 1 + env_cfg.scene.contact_forces = None + env_cfg.rewards.feet_air_time = None + env_cfg.rewards.feet_slide = None + env_cfg.terminations.base_contact = None + env_cfg.sim.physics = NewtonCfg( + solver_cfg=MJWarpSolverCfg( + njmax=95, + nconmax=10, + cone="pyramidal", + impratio=1, + integrator="implicitfast", + ), + num_substeps=1, + debug_mode=False, + ) + env_cfg.sim.use_newton_actuators = True + env = ManagerBasedRLEnv(cfg=env_cfg) + try: + stage = env.unwrapped.sim.stage + actuator_prim_count = sum(1 for prim in stage.Traverse() if prim.GetTypeName() == "NewtonActuator") + self.assertGreater( + actuator_prim_count, + 0, + "Expected authored NewtonActuator prims in manager-based scene workflow.", + ) + self.assertGreater( + len(SimulationManager.get_model().actuators), + 0, + "Expected Newton model actuators to be non-empty with use_newton_actuators=True.", + ) + finally: + env.close() + + # --------------------------------------------------------------------------- # Neural network actuator authoring: MLP and LSTM # --------------------------------------------------------------------------- From 1a767d45f6d2eb1ecd2bf9ca1fb1accaaefd17ee Mon Sep 17 00:00:00 2001 From: Julia von Muralt Date: Wed, 13 May 2026 11:32:35 +0200 Subject: [PATCH 30/35] add init --- .../isaaclab_newton/actuators/__init__.py | 36 ++++++++++++++++++ .../assets/articulation/articulation.py | 38 +++++++++---------- 2 files changed, 55 insertions(+), 19 deletions(-) create mode 100644 source/isaaclab_newton/isaaclab_newton/actuators/__init__.py diff --git a/source/isaaclab_newton/isaaclab_newton/actuators/__init__.py b/source/isaaclab_newton/isaaclab_newton/actuators/__init__.py new file mode 100644 index 000000000000..aacc07f96666 --- /dev/null +++ b/source/isaaclab_newton/isaaclab_newton/actuators/__init__.py @@ -0,0 +1,36 @@ +# Copyright (c) 2022-2026, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +"""Newton-native actuator integration for Isaac Lab. + +Public API surface: + +* :class:`~isaaclab_newton.actuators.adapter.NewtonActuatorAdapter` — + the actuator adapter used by both backends. The Newton backend + constructs it directly from ``model.actuators``; the PhysX backend + uses :meth:`~NewtonActuatorAdapter.from_usd` to build the actuators + from authored ``NewtonActuator`` USD prims. +* :class:`~isaaclab_newton.actuators.physx_wrapper.PhysxActuatorWrapper` + — flat-array wrapper that satisfies the Newton actuator + ``sim_state`` / ``sim_control`` protocol on the PhysX backend. +* :func:`~isaaclab_newton.actuators.kernels.build_implicit_dof_mask` — + builds the per-DOF implicit-actuator mask consumed by the in-graph + post-actuator kernel. + +USD authoring lives on the schema side as +:func:`~isaaclab.sim.schemas.define_actuator_properties`; both backends +call into it via :meth:`ArticulationCfg._post_spawn`. +""" + +from .adapter import NewtonActuatorAdapter, build_newton_actuator_defaults +from .kernels import build_implicit_dof_mask +from .physx_wrapper import PhysxActuatorWrapper + +__all__ = [ + "NewtonActuatorAdapter", + "PhysxActuatorWrapper", + "build_implicit_dof_mask", + "build_newton_actuator_defaults", +] diff --git a/source/isaaclab_physx/isaaclab_physx/assets/articulation/articulation.py b/source/isaaclab_physx/isaaclab_physx/assets/articulation/articulation.py index 6ab5403f8755..84af2ded183f 100644 --- a/source/isaaclab_physx/isaaclab_physx/assets/articulation/articulation.py +++ b/source/isaaclab_physx/isaaclab_physx/assets/articulation/articulation.py @@ -408,7 +408,7 @@ def write_root_pose_to_sim_index( The root pose comprises of the cartesian position and quaternion orientation in (x, y, z, w). .. note:: - This method expect partial data. + This method expects partial data. .. tip:: For maximum performance we recommend using the index method. This is because in PhysX, the tensor API @@ -432,7 +432,7 @@ def write_root_pose_to_sim_mask( The root pose comprises of the cartesian position and quaternion orientation in (x, y, z, w). .. note:: - This method expect full data. + This method expects full data. .. tip:: For maximum performance we recommend using the index method. This is because in PhysX, the tensor API @@ -457,7 +457,7 @@ def write_root_link_pose_to_sim_index( The root pose comprises of the cartesian position and quaternion orientation in (x, y, z, w). .. note:: - This method expect partial data or full data. + This method expects partial data or full data. .. tip:: For maximum performance we recommend using the index method. This is because in PhysX, the tensor API @@ -512,7 +512,7 @@ def write_root_link_pose_to_sim_mask( The root pose comprises of the cartesian position and quaternion orientation in (x, y, z, w). .. note:: - This method expect full data. + This method expects full data. .. tip:: For maximum performance we recommend using the index method. This is because in PhysX, the tensor API @@ -541,7 +541,7 @@ def write_root_com_pose_to_sim_index( The orientation is the orientation of the principal axes of inertia. .. note:: - This method expect partial data or full data. + This method expects partial data or full data. .. tip:: For maximum performance we recommend using the index method. This is because in PhysX, the tensor API @@ -602,7 +602,7 @@ def write_root_com_pose_to_sim_mask( The orientation is the orientation of the principal axes of inertia. .. note:: - This method expect full data. + This method expects full data. .. tip:: For maximum performance we recommend using the index method. This is because in PhysX, the tensor API @@ -632,7 +632,7 @@ def write_root_velocity_to_sim_index( This sets the velocity of the root's center of mass rather than the root's frame. .. note:: - This method expect partial data. + This method expects partial data. .. tip:: For maximum performance we recommend using the index method. This is because in PhysX, the tensor API @@ -659,7 +659,7 @@ def write_root_velocity_to_sim_mask( This sets the velocity of the root's center of mass rather than the root's frame. .. note:: - This method expect full data. + This method expects full data. .. tip:: For maximum performance we recommend using the index method. This is because in PhysX, the tensor API @@ -687,7 +687,7 @@ def write_root_com_velocity_to_sim_index( This sets the velocity of the root's center of mass rather than the root's frame. .. note:: - This method expect partial data or full data. + This method expects partial data or full data. .. tip:: For maximum performance we recommend using the index method. This is because in PhysX, the tensor API @@ -741,7 +741,7 @@ def write_root_com_velocity_to_sim_mask( This sets the velocity of the root's center of mass rather than the root's frame. .. note:: - This method expect full data. + This method expects full data. .. tip:: For maximum performance we recommend using the index method. This is because in PhysX, the tensor API @@ -772,7 +772,7 @@ def write_root_link_velocity_to_sim_index( This sets the velocity of the root's frame rather than the root's center of mass. .. note:: - This method expect partial data or full data. + This method expects partial data or full data. .. tip:: For maximum performance we recommend using the index method. This is because in PhysX, the tensor API @@ -831,7 +831,7 @@ def write_root_link_velocity_to_sim_mask( This sets the velocity of the root's frame rather than the root's center of mass. .. note:: - This method expect full data. + This method expects full data. .. tip:: For maximum performance we recommend using the index method. This is because in PhysX, the tensor API @@ -858,7 +858,7 @@ def write_joint_state_to_sim_mask( """Write joint positions and velocities over selected environment mask into the simulation. .. note:: - This method expect full data. + This method expects full data. .. tip:: For maximum performance we recommend using the index method. This is because in PhysX, the tensor API @@ -885,7 +885,7 @@ def write_joint_position_to_sim_index( """Write joint positions over selected environment indices into the simulation. .. note:: - This method expect partial data or full data. + This method expects partial data or full data. .. tip:: For maximum performance we recommend using the index method. This is because in PhysX, the tensor API @@ -942,7 +942,7 @@ def write_joint_position_to_sim_mask( """Write joint positions over selected environment mask into the simulation. .. note:: - This method expect full data. + This method expects full data. .. tip:: For maximum performance we recommend using the index method. This is because in PhysX, the tensor API @@ -970,7 +970,7 @@ def write_joint_velocity_to_sim_index( """Write joint velocities to the simulation. .. note:: - This method expect partial data or full data. + This method expects partial data or full data. .. tip:: For maximum performance we recommend using the index method. This is because in PhysX, the tensor API @@ -1019,7 +1019,7 @@ def write_joint_velocity_to_sim_mask( """Write joint velocities over selected environment mask into the simulation. .. note:: - This method expect full data. + This method expects full data. .. tip:: For maximum performance we recommend using the index method. This is because in PhysX, the tensor API @@ -1051,7 +1051,7 @@ def write_joint_stiffness_to_sim_index( """Write joint stiffness over selected environment indices into the simulation. .. note:: - This method expect partial data or full data. + This method expects partial data or full data. .. tip:: For maximum performance we recommend using the index method. This is because in PhysX, the tensor API @@ -1114,7 +1114,7 @@ def write_joint_stiffness_to_sim_mask( """Write joint stiffness over selected environment mask into the simulation. .. note:: - This method expect full data. + This method expects full data. .. tip:: For maximum performance we recommend using the index method. This is because in PhysX, the tensor API From 198e56225f752b9455ceba3181cb3f4096b648ac Mon Sep 17 00:00:00 2001 From: Julia von Muralt Date: Wed, 13 May 2026 12:12:37 +0200 Subject: [PATCH 31/35] layz import and test fix --- .../assets/articulation/articulation.py | 22 +++++---- .../assets/test_newton_actuators_newton.py | 46 ++++++++++++++++++- .../assets/articulation/articulation.py | 40 ++++++++++------ 3 files changed, 86 insertions(+), 22 deletions(-) diff --git a/source/isaaclab_newton/isaaclab_newton/assets/articulation/articulation.py b/source/isaaclab_newton/isaaclab_newton/assets/articulation/articulation.py index f1a15399e014..4c7c2af3d29c 100644 --- a/source/isaaclab_newton/isaaclab_newton/assets/articulation/articulation.py +++ b/source/isaaclab_newton/isaaclab_newton/assets/articulation/articulation.py @@ -8,6 +8,7 @@ from __future__ import annotations +import importlib.util import logging import warnings from collections.abc import Sequence @@ -26,13 +27,7 @@ from isaaclab.actuators import ActuatorBase, ActuatorBaseCfg, ImplicitActuator from isaaclab.assets.articulation.base_articulation import BaseArticulation -try: - from isaaclab_newton.actuators import NewtonActuatorAdapter, build_implicit_dof_mask, build_newton_actuator_defaults - from isaaclab_newton.actuators import kernels as actuator_kernels - - _HAS_NEWTON_ACTUATORS = True -except ImportError: - _HAS_NEWTON_ACTUATORS = False +_HAS_NEWTON_ACTUATORS = importlib.util.find_spec("isaaclab_newton.actuators") is not None from isaaclab.physics import PhysicsEvent from isaaclab.sim.utils.queries import find_first_matching_prim, get_all_matching_child_prims @@ -50,6 +45,7 @@ if TYPE_CHECKING: from isaaclab.assets.articulation.articulation_cfg import ArticulationCfg + # import logger logger = logging.getLogger(__name__) @@ -3582,6 +3578,12 @@ def _process_actuators_cfg(self): if _use_newton_actuators and _HAS_NEWTON_ACTUATORS: from newton import Model as NewtonModel # noqa: PLC0415 + from isaaclab_newton.actuators import ( # noqa: PLC0415 + build_implicit_dof_mask, + build_newton_actuator_defaults, + ) + from isaaclab_newton.actuators import kernels as actuator_kernels # noqa: PLC0415 + # Enable the fast path even for all-implicit articulations: # the solver runs PD internally; Lab only forwards targets. self._has_newton_actuators = True @@ -3716,9 +3718,13 @@ def _post_actuator() -> None: ) if self.cfg.actuator_value_resolution_debug_print: + if _HAS_NEWTON_ACTUATORS: + from isaaclab_newton.actuators import NewtonActuatorAdapter # noqa: PLC0415 + else: + NewtonActuatorAdapter = None # type: ignore[assignment] t = PrettyTable(["Group", "Property", "Name", "ID", "USD Value", "ActutatorCfg Value", "Applied"]) for actuator_group, actuator in self.actuators.items(): - if _HAS_NEWTON_ACTUATORS and isinstance(actuator, NewtonActuatorAdapter): + if NewtonActuatorAdapter is not None and isinstance(actuator, NewtonActuatorAdapter): continue group_count = 0 for property, resolution_details in actuator.joint_property_resolution_table.items(): diff --git a/source/isaaclab_newton/test/assets/test_newton_actuators_newton.py b/source/isaaclab_newton/test/assets/test_newton_actuators_newton.py index 594a6746a43a..49d4da05234e 100644 --- a/source/isaaclab_newton/test/assets/test_newton_actuators_newton.py +++ b/source/isaaclab_newton/test/assets/test_newton_actuators_newton.py @@ -1184,7 +1184,13 @@ class TestDecimationRemotizedPD(_DecimationMixin, TestRemotizedPDEquivalence): class TestManagerBasedSceneNewtonActuatorAuthoring(unittest.TestCase): - """Regression test for Newton actuator authoring in manager-based clone paths.""" + """Regression test for Newton actuator authoring in manager-based clone paths. + + The default G1 config uses ``ImplicitActuatorCfg`` for every group, which + intentionally skips ``NewtonActuator`` USD authoring. To exercise the + authoring path we override the scene's robot actuators with explicit + ``DCMotorCfg`` groups covering the same joint patterns. + """ def test_newton_actuators_present_for_g1_manager_env(self): env_cfg = G1FlatEnvCfg() @@ -1206,6 +1212,44 @@ def test_newton_actuators_present_for_g1_manager_env(self): debug_mode=False, ) env_cfg.sim.use_newton_actuators = True + env_cfg.scene.robot.actuators = { + "legs": DCMotorCfg( + joint_names_expr=[ + ".*_hip_yaw_joint", + ".*_hip_roll_joint", + ".*_hip_pitch_joint", + ".*_knee_joint", + "torso_joint", + ], + saturation_effort=300.0, + effort_limit=300.0, + velocity_limit=20.0, + stiffness=150.0, + damping=5.0, + ), + "feet": DCMotorCfg( + joint_names_expr=[".*_ankle_pitch_joint", ".*_ankle_roll_joint"], + saturation_effort=20.0, + effort_limit=20.0, + velocity_limit=20.0, + stiffness=20.0, + damping=2.0, + ), + "arms": DCMotorCfg( + joint_names_expr=[ + ".*_shoulder_pitch_joint", + ".*_shoulder_roll_joint", + ".*_shoulder_yaw_joint", + ".*_elbow_pitch_joint", + ".*_elbow_roll_joint", + ], + saturation_effort=300.0, + effort_limit=300.0, + velocity_limit=20.0, + stiffness=40.0, + damping=10.0, + ), + } env = ManagerBasedRLEnv(cfg=env_cfg) try: stage = env.unwrapped.sim.stage diff --git a/source/isaaclab_physx/isaaclab_physx/assets/articulation/articulation.py b/source/isaaclab_physx/isaaclab_physx/assets/articulation/articulation.py index 84af2ded183f..1d7c14abba6d 100644 --- a/source/isaaclab_physx/isaaclab_physx/assets/articulation/articulation.py +++ b/source/isaaclab_physx/isaaclab_physx/assets/articulation/articulation.py @@ -8,6 +8,7 @@ from __future__ import annotations +import importlib.util import logging import warnings from collections.abc import Sequence @@ -23,18 +24,7 @@ from isaaclab.actuators import ActuatorBase, ActuatorBaseCfg, ImplicitActuator from isaaclab.assets.articulation.base_articulation import BaseArticulation -try: - from isaaclab_newton.actuators import ( - NewtonActuatorAdapter, - PhysxActuatorWrapper, - build_implicit_dof_mask, - build_newton_actuator_defaults, - ) - from isaaclab_newton.actuators import kernels as actuator_kernels - - _HAS_NEWTON_ACTUATORS = True -except ImportError: - _HAS_NEWTON_ACTUATORS = False +_HAS_NEWTON_ACTUATORS = importlib.util.find_spec("isaaclab_newton.actuators") is not None from isaaclab.sim.utils.queries import find_first_matching_prim, get_all_matching_child_prims @@ -50,6 +40,8 @@ from .articulation_data import ArticulationData if TYPE_CHECKING: + from isaaclab_newton.actuators import NewtonActuatorAdapter + import omni.physics.tensors.api as physx from isaaclab.assets.articulation.articulation_cfg import ArticulationCfg @@ -1244,6 +1236,8 @@ def _write_actuator_param( if adapter is None: return + from isaaclab_newton.actuators import kernels as actuator_kernels # noqa: PLC0415 + env_id_pos = torch.full( (self.num_instances,), -1, @@ -3869,7 +3863,21 @@ def _process_actuators_cfg(self): _use_newton_actuators = getattr(self._sim_cfg, "use_newton_actuators", False) + if _use_newton_actuators and not _HAS_NEWTON_ACTUATORS: + logger.warning( + "use_newton_actuators is enabled but 'isaaclab_newton.actuators' is not available." + " Newton-native actuators will be disabled and the simulation will fall back to the" + " Isaac Lab actuator path. Install the isaaclab_newton extension to enable the fast path." + ) + if _HAS_NEWTON_ACTUATORS and _use_newton_actuators: + from isaaclab_newton.actuators import ( # noqa: PLC0415 + NewtonActuatorAdapter, + PhysxActuatorWrapper, + build_implicit_dof_mask, + build_newton_actuator_defaults, + ) + from isaaclab.sim.utils.stage import get_current_stage # noqa: PLC0415 # Enable the fast path even for all-implicit articulations: @@ -3982,9 +3990,13 @@ def _process_actuators_cfg(self): ) if self.cfg.actuator_value_resolution_debug_print: + if _HAS_NEWTON_ACTUATORS: + from isaaclab_newton.actuators import NewtonActuatorAdapter # noqa: PLC0415 + else: + NewtonActuatorAdapter = None # type: ignore[assignment] t = PrettyTable(["Group", "Property", "Name", "ID", "USD Value", "ActutatorCfg Value", "Applied"]) for actuator_group, actuator in self.actuators.items(): - if _HAS_NEWTON_ACTUATORS and isinstance(actuator, NewtonActuatorAdapter): + if NewtonActuatorAdapter is not None and isinstance(actuator, NewtonActuatorAdapter): continue group_count = 0 for property, resolution_details in actuator.joint_property_resolution_table.items(): @@ -4201,6 +4213,8 @@ def _apply_actuator_model_newton(self): resulting buffer. The final ``joint_f_2d`` is what gets pushed to PhysX as the actuation force in :meth:`write_data_to_sim`. """ + from isaaclab_newton.actuators import kernels as actuator_kernels # noqa: PLC0415 + w = self._physx_actuator_wrapper w.joint_f_2d.assign(self._data._joint_effort_target) if self.newton_actuator_adapter is not None: From c5675d344def026a4ff9b158f3654f70a8c33795 Mon Sep 17 00:00:00 2001 From: Julia von Muralt Date: Wed, 13 May 2026 13:38:57 +0200 Subject: [PATCH 32/35] fix rendering logic in the fast track --- source/isaaclab/isaaclab/envs/direct_marl_env.py | 4 +++- source/isaaclab/isaaclab/envs/direct_rl_env.py | 4 +++- source/isaaclab/isaaclab/envs/manager_based_env.py | 4 +++- source/isaaclab/isaaclab/envs/manager_based_rl_env.py | 4 +++- 4 files changed, 12 insertions(+), 4 deletions(-) diff --git a/source/isaaclab/isaaclab/envs/direct_marl_env.py b/source/isaaclab/isaaclab/envs/direct_marl_env.py index b846971eb18c..fa49336eaf12 100644 --- a/source/isaaclab/isaaclab/envs/direct_marl_env.py +++ b/source/isaaclab/isaaclab/envs/direct_marl_env.py @@ -442,7 +442,9 @@ def step(self, actions: dict[AgentID, ActionType]) -> EnvStepReturn: self._apply_action() self.scene.write_data_to_sim() self.sim.step(render=False) - if is_rendering: + # render only when a render_interval boundary falls within this decimation block, + # mirroring the per-sub-step check in the else branch. + if self._sim_step_counter % self.cfg.sim.render_interval == 0 and is_rendering: self.sim.render(skip_app_pumping=not self.render_enabled) self.scene.update(dt=self.step_dt) else: diff --git a/source/isaaclab/isaaclab/envs/direct_rl_env.py b/source/isaaclab/isaaclab/envs/direct_rl_env.py index baded7bd5bed..68146b10e3d9 100644 --- a/source/isaaclab/isaaclab/envs/direct_rl_env.py +++ b/source/isaaclab/isaaclab/envs/direct_rl_env.py @@ -433,7 +433,9 @@ def step(self, action: torch.Tensor) -> VecEnvStepReturn: self._apply_action() self.scene.write_data_to_sim() self.sim.step(render=False) - if is_rendering: + # render only when a render_interval boundary falls within this decimation block, + # mirroring the per-sub-step check in the else branch. + if self._sim_step_counter % self.cfg.sim.render_interval == 0 and is_rendering: self.sim.render(skip_app_pumping=not self.render_enabled) self.scene.update(dt=self.step_dt) else: diff --git a/source/isaaclab/isaaclab/envs/manager_based_env.py b/source/isaaclab/isaaclab/envs/manager_based_env.py index 342f6c6e2b66..63042ae078cb 100644 --- a/source/isaaclab/isaaclab/envs/manager_based_env.py +++ b/source/isaaclab/isaaclab/envs/manager_based_env.py @@ -539,7 +539,9 @@ def step(self, action: torch.Tensor) -> tuple[VecEnvObs, dict]: self.action_manager.apply_action() self.scene.write_data_to_sim() self.sim.step(render=False) - if is_rendering: + # render only when a render_interval boundary falls within this decimation block, + # mirroring the per-sub-step check in the else branch. + if self._sim_step_counter % self.cfg.sim.render_interval == 0 and is_rendering: self.sim.render(skip_app_pumping=not self.render_enabled) self.scene.update(dt=self.step_dt) else: diff --git a/source/isaaclab/isaaclab/envs/manager_based_rl_env.py b/source/isaaclab/isaaclab/envs/manager_based_rl_env.py index 9bb9e21e6f5c..58c0a0eaef42 100644 --- a/source/isaaclab/isaaclab/envs/manager_based_rl_env.py +++ b/source/isaaclab/isaaclab/envs/manager_based_rl_env.py @@ -205,7 +205,9 @@ def step(self, action: torch.Tensor) -> VecEnvStepReturn: self.scene.write_data_to_sim() self.sim.step(render=False) self.recorder_manager.record_post_physics_decimation_step() - if is_rendering: + # render only when a render_interval boundary falls within this decimation block, + # mirroring the per-sub-step check in the else branch. + if self._sim_step_counter % self.cfg.sim.render_interval == 0 and is_rendering: self.sim.render(skip_app_pumping=not self.render_enabled) self.scene.update(dt=self.step_dt) else: From 232b19df67a7a03c9a07fa76ef720a7186088cb4 Mon Sep 17 00:00:00 2001 From: Antoine Richard Date: Wed, 13 May 2026 15:33:25 +0200 Subject: [PATCH 33/35] Fix NewtonKaminoManager call to renamed _simulate The actuator integration refactor in this PR split NewtonManager._simulate into _simulate_full (decimation x actuators + solver) and _simulate_physics_only (solver substeps + sensors, no actuators). The NewtonKaminoManager override of step() and the CUDA graph capture helper still called the removed _simulate, which raised AttributeError on every Kamino step(). Route both call sites to _simulate_physics_only (the semantic equivalent of the prior _simulate -- no actuator stepping, no decimation loop) and consolidate _mark_transforms_dirty() to run unconditionally after both branches of step(), matching the new parent step() idiom. Fixes the kamino_internal_contacts and kamino_newton_pipeline cases of test_initialize_solver_populates_canonical_state. --- .../isaaclab_newton/physics/kamino_manager.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/source/isaaclab_newton/isaaclab_newton/physics/kamino_manager.py b/source/isaaclab_newton/isaaclab_newton/physics/kamino_manager.py index 39d0fafbc1bb..ba2d407d2a9a 100644 --- a/source/isaaclab_newton/isaaclab_newton/physics/kamino_manager.py +++ b/source/isaaclab_newton/isaaclab_newton/physics/kamino_manager.py @@ -101,11 +101,11 @@ def step(cls) -> None: # Step simulation (graphed or not; _graph is None when capture is disabled or failed) if cfg is not None and cfg.use_cuda_graph and cls._graph is not None and "cuda" in device: # type: ignore[union-attr] wp.capture_launch(cls._graph) - if cls._usdrt_stage is not None: - cls._mark_transforms_dirty() else: with wp.ScopedDevice(device): - cls._simulate() + cls._simulate_physics_only() + if cls._usdrt_stage is not None: + cls._mark_transforms_dirty() # Launch solver-specific debug logging after stepping. cls._log_solver_debug() @@ -138,7 +138,7 @@ def _capture_or_defer_cuda_graph(cls) -> None: if cls._usdrt_stage is None: # No RTX active — use standard Warp capture (cudaStreamCaptureModeGlobal). with wp.ScopedCapture() as capture: - cls._simulate() + cls._simulate_physics_only() NewtonManager._graph = capture.graph logger.info("Newton CUDA graph captured (standard Warp mode)") From f215bbd9f403ddaedd29cc697691ed5050ad1179 Mon Sep 17 00:00:00 2001 From: Julia von Muralt Date: Tue, 19 May 2026 10:17:02 +0200 Subject: [PATCH 34/35] fix --- .../isaaclab_newton/assets/articulation/articulation.py | 6 ++++-- .../isaaclab_physx/assets/articulation/articulation.py | 6 ++++-- 2 files changed, 8 insertions(+), 4 deletions(-) diff --git a/source/isaaclab_newton/isaaclab_newton/assets/articulation/articulation.py b/source/isaaclab_newton/isaaclab_newton/assets/articulation/articulation.py index 4c7c2af3d29c..690065980945 100644 --- a/source/isaaclab_newton/isaaclab_newton/assets/articulation/articulation.py +++ b/source/isaaclab_newton/isaaclab_newton/assets/articulation/articulation.py @@ -254,7 +254,9 @@ def reset(self, env_ids: Sequence[int] | None = None, env_mask: wp.array | None # be cleared for the resetting envs). The adapter spans the whole model, # so calling reset here resets state for every articulation that shares # this env id; that's correct because env ids are world-scoped. - if self._has_newton_actuators and SimulationManager._adapter is not None: + # ``getattr`` guards subclasses (e.g. ``Multirotor``) that override + # ``_process_actuators_cfg`` and never initialize ``_has_newton_actuators``. + if getattr(self, "_has_newton_actuators", False) and SimulationManager._adapter is not None: SimulationManager._adapter.reset(env_ids) # reset external wrenches. self._instantaneous_wrench_composer.reset(env_ids, env_mask) @@ -292,7 +294,7 @@ def write_data_to_sim(self): ) self._instantaneous_wrench_composer.reset() - if self._has_newton_actuators: + if getattr(self, "_has_newton_actuators", False): # Raw targets go directly to Newton's control object. Newton PD # consumes ``joint_act`` for explicit (Newton-managed) joints; the # solver's built-in joint drive does the PD for implicit joints diff --git a/source/isaaclab_physx/isaaclab_physx/assets/articulation/articulation.py b/source/isaaclab_physx/isaaclab_physx/assets/articulation/articulation.py index cea8d6817eec..78fdb387aa04 100644 --- a/source/isaaclab_physx/isaaclab_physx/assets/articulation/articulation.py +++ b/source/isaaclab_physx/isaaclab_physx/assets/articulation/articulation.py @@ -234,7 +234,9 @@ def reset(self, env_ids: Sequence[int] | None = None, env_mask: wp.array | None actuator.reset(env_ids) # Reset Newton-actuator per-env states (delay queues, neural hidden state, etc.). # The adapter is per-articulation on PhysX and is not part of ``self.actuators``. - if self._has_newton_actuators and self.newton_actuator_adapter is not None: + # ``getattr`` guards subclasses (e.g. ``Multirotor``) that override + # ``_process_actuators_cfg`` and never initialize these attributes. + if getattr(self, "_has_newton_actuators", False) and getattr(self, "newton_actuator_adapter", None) is not None: self.newton_actuator_adapter.reset(env_ids) # reset external wrenches. self._instantaneous_wrench_composer.reset(env_ids, env_mask) @@ -267,7 +269,7 @@ def write_data_to_sim(self): ) self._instantaneous_wrench_composer.reset() - if self._has_newton_actuators: + if getattr(self, "_has_newton_actuators", False): # Newton fast path: pos/vel targets pass straight through; the # in-graph kernel inside ``_apply_actuator_model_newton`` merges # Newton's actuator output (explicit DOFs) with user FF From a48d3bd22151a5f3a23abcab54b2bfe730958984 Mon Sep 17 00:00:00 2001 From: Julia von Muralt Date: Tue, 19 May 2026 10:33:53 +0200 Subject: [PATCH 35/35] fix --- source/isaaclab/isaaclab/envs/direct_marl_env.py | 1 + source/isaaclab/isaaclab/envs/direct_rl_env.py | 1 + source/isaaclab/isaaclab/envs/manager_based_env.py | 1 + 3 files changed, 3 insertions(+) diff --git a/source/isaaclab/isaaclab/envs/direct_marl_env.py b/source/isaaclab/isaaclab/envs/direct_marl_env.py index fa49336eaf12..824dda6d0150 100644 --- a/source/isaaclab/isaaclab/envs/direct_marl_env.py +++ b/source/isaaclab/isaaclab/envs/direct_marl_env.py @@ -92,6 +92,7 @@ def __init__(self, cfg: DirectMARLEnvCfg, render_mode: str | None = None, **kwar self.render_mode = render_mode # initialize internal variables self._is_closed = False + self._physics_handles_decimation = False # set the seed for the environment if self.cfg.seed is not None: diff --git a/source/isaaclab/isaaclab/envs/direct_rl_env.py b/source/isaaclab/isaaclab/envs/direct_rl_env.py index 68146b10e3d9..650568bbb1e2 100644 --- a/source/isaaclab/isaaclab/envs/direct_rl_env.py +++ b/source/isaaclab/isaaclab/envs/direct_rl_env.py @@ -97,6 +97,7 @@ def __init__(self, cfg: DirectRLEnvCfg, render_mode: str | None = None, **kwargs self.render_mode = render_mode # initialize internal variables self._is_closed = False + self._physics_handles_decimation = False # set the seed for the environment if self.cfg.seed is not None: diff --git a/source/isaaclab/isaaclab/envs/manager_based_env.py b/source/isaaclab/isaaclab/envs/manager_based_env.py index 63042ae078cb..f89101dae072 100644 --- a/source/isaaclab/isaaclab/envs/manager_based_env.py +++ b/source/isaaclab/isaaclab/envs/manager_based_env.py @@ -94,6 +94,7 @@ def __init__(self, cfg: ManagerBasedEnvCfg): self.cfg = cfg # initialize internal variables self._is_closed = False + self._physics_handles_decimation = False # set the seed for the environment if self.cfg.seed is not None: