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/isaaclab/assets/articulation/articulation_cfg.py b/source/isaaclab/isaaclab/assets/articulation/articulation_cfg.py index b7514fb8a91c..2124e02aee77 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.configclass import configclass @@ -74,3 +74,25 @@ 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. 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.actuators is MISSING: + return + from isaaclab.sim.schemas.schemas_actuators import define_actuator_properties # noqa: PLC0415 + + # 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/isaaclab/assets/articulation/base_articulation.py b/source/isaaclab/isaaclab/assets/articulation/base_articulation.py index ff3f99e8a024..76e831bdad17 100644 --- a/source/isaaclab/isaaclab/assets/articulation/base_articulation.py +++ b/source/isaaclab/isaaclab/assets/articulation/base_articulation.py @@ -16,6 +16,7 @@ import torch import warp as wp +from ...sim import SimulationContext from ...utils.leapp.leapp_semantics import OutputKindEnum, joint_names_resolver, leapp_tensor_semantics from ..asset_base import AssetBase @@ -95,6 +96,8 @@ def __init__(self, cfg: ArticulationCfg): cfg: A configuration instance. """ 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/assets/asset_base.py b/source/isaaclab/isaaclab/assets/asset_base.py index b72788083ad3..3031debbca2c 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 @@ -79,7 +81,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 @@ -96,6 +98,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 095bb6afd9cd..a284ba394125 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.configclass 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/envs/direct_marl_env.py b/source/isaaclab/isaaclab/envs/direct_marl_env.py index 56d3a38cdd1c..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: @@ -197,6 +198,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) @@ -433,21 +438,32 @@ 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. + # 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) - # 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 c03ca1f73596..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: @@ -203,6 +204,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) @@ -424,21 +429,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._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. + # 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) - # 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 620cf2895718..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: @@ -216,6 +217,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() @@ -530,21 +535,32 @@ 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. + # 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) - # 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 9060e7874c32..48966261a577 100644 --- a/source/isaaclab/isaaclab/envs/manager_based_rl_env.py +++ b/source/isaaclab/isaaclab/envs/manager_based_rl_env.py @@ -199,22 +199,34 @@ 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. + # 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) - # 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 9483d2d2d0cd..f430c105ee78 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,6 +1230,50 @@ def randomize(data: torch.Tensor, params: tuple[float, float]) -> torch.Tensor: damping=damping, joint_ids=actuator.joint_indices, env_ids=env_ids ) + # 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 + + if isinstance(self.asset_cfg.joint_ids, slice): + joint_ids = torch.arange(self.asset.num_joints, device=self.asset.device, dtype=torch.long) + else: + 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( + 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, + ) + 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, + ) + self.asset.write_actuator_damping_to_sim( + damping=new_damping, + env_ids=env_ids, + joint_ids=joint_ids, + ) + class randomize_joint_parameters(ManagerTermBase): """Randomize the simulated joint parameters of an articulation by adding, scaling, or setting random values. diff --git a/source/isaaclab/isaaclab/physics/physics_manager.py b/source/isaaclab/isaaclab/physics/physics_manager.py index 6727d87174b6..3f63dfa01c60 100644 --- a/source/isaaclab/isaaclab/physics/physics_manager.py +++ b/source/isaaclab/isaaclab/physics/physics_manager.py @@ -346,6 +346,29 @@ 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. + """ + pass + + @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/schemas/__init__.pyi b/source/isaaclab/isaaclab/sim/schemas/__init__.pyi index fd0c882a5f35..a9306afdf414 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", @@ -58,6 +59,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/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/isaaclab/sim/simulation_cfg.py b/source/isaaclab/isaaclab/sim/simulation_cfg.py index 8abed4cbfc17..2c71c724ad3d 100644 --- a/source/isaaclab/isaaclab/sim/simulation_cfg.py +++ b/source/isaaclab/isaaclab/sim/simulation_cfg.py @@ -288,6 +288,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/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_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_newton/isaaclab_newton/actuators/adapter.py b/source/isaaclab_newton/isaaclab_newton/actuators/adapter.py new file mode 100644 index 000000000000..0d28cf576dfb --- /dev/null +++ b/source/isaaclab_newton/isaaclab_newton/actuators/adapter.py @@ -0,0 +1,456 @@ +# 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. + +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 Sequence +from typing import Any + +import numpy as np +import torch +import warp as wp +from newton.actuators import Actuator, Clamping, Delay + +from .kernels import ( + build_per_dof_env_mask_kernel, + scatter_gain_kernel, + set_mask_kernel, + zero_at_indices_kernel, +) + +# --------------------------------------------------------------------------- +# Abstract base — backend-independent logic +# --------------------------------------------------------------------------- + + +class NewtonActuatorAdapter: + """Adapter that wraps a list of :class:`newton.actuators.Actuator`. + + 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__( + 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 + + # 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(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] + + # 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" + + def finalize(self, sim_control: Any) -> None: + """Bind the pre-clamp computed-effort buffer onto ``sim_control``. + + Args: + 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. + """ + 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. + + 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]. + """ + # Zero before scatter-add (actuators accumulate into this buffer). + self._computed_effort.zero_() + 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] | torch.Tensor | None = None) -> None: + """Reset actuator states for the given environments. + + Args: + 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): + 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: + 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 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(per_dof_mask) + if sb is not None: + 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) + + @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*. + + 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: + 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. + """ + 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( + 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], + inputs=[ctrl.kd, flat_damping, act.indices, dof_offset, num_joints], + device=wp_device, + ) + 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 + + +# --------------------------------------------------------------------------- +# PhysX-only USD parsing +# --------------------------------------------------------------------------- + + +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 _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/kernels.py b/source/isaaclab_newton/isaaclab_newton/actuators/kernels.py new file mode 100644 index 000000000000..8205773b640a --- /dev/null +++ b/source/isaaclab_newton/isaaclab_newton/actuators/kernels.py @@ -0,0 +1,197 @@ +# 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 torch +import warp as wp + +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`), +# 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. +# --------------------------------------------------------------------------- + + +@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 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), +): + """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 + out_mask[i] = env_mask[env] + + +@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 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 + local_dof = global_dof % num_joints + dst[env * num_joints + local_dof] = src[i] + + +@wp.kernel(enable_backward=False) +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), +): + """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] + + +# --------------------------------------------------------------------------- +# Articulation-level kernels: in-graph post-actuator hook. +# --------------------------------------------------------------------------- + + +@wp.kernel(enable_backward=False) +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_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: 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: + 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: + 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, +) -> 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(): + 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), modes 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..b3f48a2f9dee --- /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 ff04b96c63ca..690065980945 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 @@ -25,6 +26,9 @@ from isaaclab.actuators import ActuatorBase, ActuatorBaseCfg, ImplicitActuator from isaaclab.assets.articulation.base_articulation import BaseArticulation + +_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 from isaaclab.utils.string import resolve_matching_names, resolve_matching_names_values @@ -41,6 +45,7 @@ if TYPE_CHECKING: from isaaclab.assets.articulation.articulation_cfg import ArticulationCfg + # import logger logger = logging.getLogger(__name__) @@ -113,8 +118,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 """ @@ -236,9 +246,18 @@ 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. + # ``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) self._permanent_wrench_composer.reset(env_ids, env_mask) @@ -275,14 +294,26 @@ 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 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 + # (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() + 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. @@ -1472,6 +1503,85 @@ 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, *, @@ -2099,6 +2209,24 @@ 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. + """ + + @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(self._build_env_mask_kernel, dim=env_ids.shape[0], inputs=[mask, env_ids], device=self.device) + return mask + """ Operations - Setters. """ @@ -3428,123 +3556,160 @@ 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 implicit/explicit mask consumed by the in-graph kernel + # ``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) + + 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." + ) - # 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}." + 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 + # 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() + + # 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, ) - # 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=self._data.joint_stiffness.torch[:, joint_ids], - damping=self._data.joint_damping.torch[:, joint_ids], - armature=self._data.joint_armature.torch[:, joint_ids], - friction=self._data.joint_friction_coeff.torch[:, joint_ids], - effort_limit=self._data.joint_effort_limits.torch[:, joint_ids].clone(), - velocity_limit=self._data.joint_vel_limits.torch[:, joint_ids], + 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 + is_implicit = ( + "ImplicitActuator" in cls_type + if isinstance(cls_type, str) + else issubclass(cls_type, ImplicitActuator) + ) + if is_implicit: + self._create_lab_actuator(actuator_name, actuator_cfg) + else: + self._create_lab_actuator(actuator_name, actuator_cfg, properties_only=True) + + # ``_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, ) - # 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) + + # 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``. + + # 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: - # 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._data._sim_bind_joint_computed_effort = wp.zeros( + (self.num_instances, self.num_joints), + dtype=wp.float32, + device=self.device, + ) - # 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, - ) + def _post_actuator() -> None: + wp.launch( + 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_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._computed_torque, + self._data._applied_torque, + ], + device=self.device, + ) + + SimulationManager.register_post_actuator_callback(_post_actuator) + + 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()) @@ -3555,8 +3720,14 @@ 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 NewtonActuatorAdapter is not None 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): @@ -3567,6 +3738,112 @@ 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, + *, + 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( + 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], + ) + + # 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, + ) + + 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 + 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._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._actuator_damping], + 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 0a2a84392619..2ddb3a13d3e4 100644 --- a/source/isaaclab_newton/isaaclab_newton/assets/articulation/articulation_data.py +++ b/source/isaaclab_newton/isaaclab_newton/assets/articulation/articulation_data.py @@ -1507,6 +1507,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_position_target = self._root_view.get_attribute( "joint_target_pos", SimulationManager.get_control() )[:, 0] @@ -1538,6 +1539,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/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)") diff --git a/source/isaaclab_newton/isaaclab_newton/physics/newton_manager.py b/source/isaaclab_newton/isaaclab_newton/physics/newton_manager.py index 7f112af4bacf..fb8188ff5909 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, SceneDataBackend, SceneDataFormat from isaaclab.sim.utils.newton_model_utils import replace_newton_shape_colors @@ -45,6 +45,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__) @@ -166,6 +168,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 @@ -197,6 +200,13 @@ class NewtonManager(PhysicsManager): _world_reset_mask: wp.array | None = None # (num_envs,) wp.int32 — for SolverKamino.reset(world_mask=...) _fk_reset_mask: wp.array | None = None # (articulation_count,) wp.bool — for eval_fk(mask=...) + # Newton actuator adapter (owns actuators and double-buffered states) + _adapter: NewtonActuatorAdapter | 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 _graph_capture_pending: bool = False @@ -407,7 +417,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 @@ -419,9 +452,7 @@ def step(cls) -> None: cls._solver.notify_model_changed(change) NewtonManager._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] @@ -443,20 +474,37 @@ def step(cls) -> None: NewtonManager._world_reset_mask.zero_() NewtonManager._fk_reset_mask.zero_() - # 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) + for cb in cls._post_actuator_callbacks: + cb() + + 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() # Launch solver-specific debug logging after stepping. cls._log_solver_debug() - PhysicsManager._sim_time += cls._solver_dt * cls._num_substeps - @classmethod def close(cls) -> None: """Clean up Newton physics resources.""" @@ -522,6 +570,14 @@ def clear(cls): NewtonManager._newton_frame_transform_sensors = [] NewtonManager._newton_imu_sensors = [] NewtonManager._report_contacts = False + NewtonManager._adapter = None + NewtonManager._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 + # a CUDA graph (see :meth:`_is_all_graphable`). + NewtonManager._use_newton_actuators_active = False + NewtonManager._decimation = 1 # Per-world reset masks NewtonManager._world_reset_mask = None NewtonManager._fk_reset_mask = None @@ -826,6 +882,14 @@ def start_simulation(cls) -> None: NewtonManager._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 the first + # call to ``activate_newton_actuator_path`` from any Newton-fast-path + # 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) NewtonManager._fk_reset_mask = wp.zeros(cls._model.articulation_count, dtype=wp.bool, device=device) @@ -1070,12 +1134,18 @@ def initialize_solver(cls) -> None: if cls._usdrt_stage is not None: cls._setup_cubric_bindings() - device = PhysicsManager._device - use_cuda_graph = cfg.use_cuda_graph and "cuda" in device # type: ignore[union-attr] - if use_cuda_graph: - cls._capture_or_defer_cuda_graph() - else: - NewtonManager._graph = None + # 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: @@ -1096,23 +1166,50 @@ def _setup_cubric_bindings(cls) -> None: logger.warning("cubric bindings init failed; falling back to update_world_xforms()") @classmethod - def _capture_or_defer_cuda_graph(cls) -> None: - """Capture the physics CUDA graph, or defer if RTX is initializing.""" - with Timer(name="newton_cuda_graph", msg="CUDA graph took:"): - if cls._usdrt_stage is None: - # No RTX active — use standard Warp capture (cudaStreamCaptureModeGlobal). - with wp.ScopedCapture() as capture: - cls._simulate() - NewtonManager._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). - NewtonManager._graph = None - NewtonManager._graph_capture_pending = True - logger.info("Newton CUDA graph capture deferred until first step() (RTX active)") + 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 + + use_cuda_graph = cfg.use_cuda_graph and "cuda" in device + if use_cuda_graph: + with Timer(name="newton_cuda_graph", msg="CUDA graph took:"): + 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() + NewtonManager._graph = capture.graph + logger.info("Newton CUDA graph captured (standard Warp mode)") + + # Kamino: StateKamino.from_newton() lazily allocates body_f_total, + # joint_q_prev, and joint_lambdas via wp.clone/wp.zeros during the + # first step() inside graph capture. Replay once to pin those + # memory-pool addresses before any eager solver.reset() call. + if isinstance(cls._solver, SolverKamino): + wp.capture_launch(cls._graph) + 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). + NewtonManager._graph = None + NewtonManager._graph_capture_pending = True + logger.info("Newton CUDA graph capture deferred until first step() (RTX active)") + else: + NewtonManager._graph = None @classmethod def _capture_relaxed_graph(cls, device: str): @@ -1143,7 +1240,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. @@ -1161,8 +1258,9 @@ def _capture_relaxed_graph(cls, device: str): # Warmup: pre-allocate all solver 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). @@ -1195,7 +1293,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 @@ -1235,62 +1333,85 @@ 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) + # ------------------------------------------------------------------ + # 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 _ in range(cls._num_substeps): - cls._step_solver(cls._state_0, cls._state_0, cls._control, cls._solver_dt) + 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): - cls._step_solver(cls._state_0, cls._state_1, cls._control, cls._solver_dt) - 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: - NewtonManager._state_0, NewtonManager._state_1 = cls._state_1, cls._state_0 + 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 = cls._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) + for cb in cls._post_actuator_callbacks: + cb() + + 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]: @@ -1605,6 +1726,97 @@ def get_solver_dt(cls) -> float: """Get the solver substep timestep.""" return cls._solver_dt + @classmethod + def _is_all_graphable(cls) -> bool: + """``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: + """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. + """ + # 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 + if cls._model is None or not cls._model.actuators: + return + from isaaclab_newton.actuators import NewtonActuatorAdapter # noqa: PLC0415 + + dofs_per_env = cls._model.joint_dof_count // cls._num_envs + NewtonManager._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(cls._control) + + @classmethod + def register_post_actuator_callback(cls, callback: Callable[[], None]) -> None: + """Append a hook to the list invoked after the actuator step on every iteration. + + 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. Multiple articulations register + their own implicit-DOF telemetry / FF-routing kernels here; all + registered callbacks fire in registration order each step. + """ + cls._post_actuator_callbacks.append(callback) + + @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/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/setup.py b/source/isaaclab_newton/setup.py index 37dd583df959..40a6b996397f 100644 --- a/source/isaaclab_newton/setup.py +++ b/source/isaaclab_newton/setup.py @@ -60,6 +60,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_newton_actuators_newton.py b/source/isaaclab_newton/test/assets/test_newton_actuators_newton.py new file mode 100644 index 000000000000..49d4da05234e --- /dev/null +++ b/source/isaaclab_newton/test/assets/test_newton_actuators_newton.py @@ -0,0 +1,1438 @@ +# 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_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, 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 + +# --------------------------------------------------------------------------- +# 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, + feedforward: float | None = None, +) -> dict: + """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. + 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'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``, ``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) + 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) + 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() + 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()) + 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 +# --------------------------------------------------------------------------- + + +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 = {} + 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 = 1e-2 + vel_rtol: float = 1e-2 + torque_atol: float = 1e-3 + torque_rtol: float = 1e-3 + + @classmethod + def setUpClass(cls): + 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, + 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_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}", + ) + + +# --------------------------------------------------------------------------- +# 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 + + +# --------------------------------------------------------------------------- +# 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 + + +# --------------------------------------------------------------------------- +# Implicit + non-zero feedforward effort target +# --------------------------------------------------------------------------- + + +class TestImplicitWithFeedforwardEquivalence(_EquivalenceTestBase): + """Implicit-only actuators with a non-zero feedforward effort target. + + 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. + """ + + __test__ = True + actuators = IMPLICIT_ONLY_ACTUATORS + feedforward = 2.0 + torque_atol = 0.5 + + +# --------------------------------------------------------------------------- +# 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}", + ) + + +# --------------------------------------------------------------------------- +# 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 +# --------------------------------------------------------------------------- + +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") + + +# --------------------------------------------------------------------------- +# Decimation tests: re-run equivalence with decimation > 1 + CUDA graph capture +# --------------------------------------------------------------------------- + +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 _DecimationMixin: + """Common knobs for decimation/CUDA-graph variants of equivalence classes.""" + + __test__ = True + dt = 1.0 / 100.0 + newton_cfg = NEWTON_CFG_DEC + num_steps = 5 + decimation = 2 + + +class TestDecimationDCMotor(_DecimationMixin, TestDCMotorEquivalence): + """DCMotor — same equivalence checks, with decimation=2 + CUDA graph.""" + + +class TestDecimationIdealPD(_DecimationMixin, TestIdealPDEquivalence): + """IdealPD — decimation=2 + CUDA graph.""" + + +class TestDecimationDelayedPD(_DecimationMixin, TestDelayedPDEquivalence): + """DelayedPD — decimation=2 + CUDA graph (delay queue stepped inside the captured graph).""" + + +class TestDecimationMixed(_DecimationMixin, TestMixedActuatorEquivalence): + """Mixed (IdealPD + DCMotor) — decimation=2 + CUDA graph.""" + + +# --------------------------------------------------------------------------- +# Per-env reset: actuator state isolation +# --------------------------------------------------------------------------- + +RESET_WARMUP_STEPS = 3 + + +class TestActuatorStateReset(unittest.TestCase): + """Reset must clear the actuator state buffers for the requested envs only. + + Inspects ``adapter.actuators[i].state.delay_state.num_pushes`` directly: + + * 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``. + + 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. + """ + + 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)) + art_cfg = ANYMAL_C_CFG.replace( + actuators=DELAYED_PD_ACTUATORS, + prim_path="/World/Env_.*/Robot", + ) + articulation = Articulation(art_cfg) + sim.reset() + + 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) + 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 = 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", + ) + + 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", + ) + + articulation.reset(env_ids=torch.tensor([self.RESET_ENV], device=articulation.device, dtype=torch.long)) + + self.assertTrue( + torch.all(buf[:, self.RESET_ENV] == 0).item(), + f"Lab: env {self.RESET_ENV} buffer not zeroed after reset.", + ) + self.assertTrue( + (buf[:, self.UNCHANGED_ENV] != 0).any().item(), + f"Lab: env {self.UNCHANGED_ENV} buffer was zeroed — reset leaked into an unselected env.", + ) + finally: + ctx.__exit__(None, None, None) + + +# --------------------------------------------------------------------------- +# 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") + + +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.""" + + +class TestManagerBasedSceneNewtonActuatorAuthoring(unittest.TestCase): + """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() + 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_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 + 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 +# --------------------------------------------------------------------------- + + +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) + + 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, + } + ) + } + torch.jit.save(scripted, tmp_path, _extra_files=extra) + return tmp_path + + +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) + + 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_path, _extra_files=extra) + return tmp_path + + +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"]) + + +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"]) + + +if __name__ == "__main__": + unittest.main() 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. diff --git a/source/isaaclab_physx/isaaclab_physx/assets/articulation/articulation.py b/source/isaaclab_physx/isaaclab_physx/assets/articulation/articulation.py index ea2c85e4c3bd..78fdb387aa04 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 @@ -22,6 +23,10 @@ from isaaclab.actuators import ActuatorBase, ActuatorBaseCfg, ImplicitActuator from isaaclab.assets.articulation.base_articulation import BaseArticulation + +_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 from isaaclab.utils.string import resolve_matching_names, resolve_matching_names_values from isaaclab.utils.types import ArticulationActions @@ -35,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 @@ -111,8 +118,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 """ @@ -217,9 +229,15 @@ 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 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``. + # ``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) self._permanent_wrench_composer.reset(env_ids, env_mask) @@ -239,30 +257,40 @@ def write_data_to_sim(self): if self._instantaneous_wrench_composer.active: composer = self._instantaneous_wrench_composer composer.add_raw_buffers_from(self._permanent_wrench_composer) - get_force_data = self._get_inst_wrench_force_f32 - get_torque_data = self._get_inst_wrench_torque_f32 else: composer = self._permanent_wrench_composer - get_force_data = self._get_perm_wrench_force_f32 - get_torque_data = self._get_perm_wrench_torque_f32 composer.compose_to_body_frame() self.root_view.apply_forces_and_torques_at_position( - force_data=get_force_data(), - torque_data=get_torque_data(), + force_data=composer.out_force_b.warp.flatten().view(wp.float32), + torque_data=composer.out_torque_b.warp.flatten().view(wp.float32), position_data=None, indices=self._ALL_INDICES, is_global=False, ) 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 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 + # (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) + else: + # 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: + 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. @@ -469,7 +497,7 @@ def write_root_link_pose_to_sim_index( self.data._body_com_jacobian_w.timestamp = -1.0 self.data._gravity_compensation_forces.timestamp = -1.0 # set into simulation - self.root_view.set_root_transforms(self._get_root_link_pose_w_f32(), indices=env_ids) + self.root_view.set_root_transforms(self.data._root_link_pose_w.data.view(wp.float32), indices=env_ids) def write_root_link_pose_to_sim_mask( self, @@ -562,7 +590,7 @@ def write_root_com_pose_to_sim_index( self.data._body_com_jacobian_w.timestamp = -1.0 self.data._gravity_compensation_forces.timestamp = -1.0 # set into simulation - self.root_view.set_root_transforms(self._get_root_link_pose_w_f32(), indices=env_ids) + self.root_view.set_root_transforms(self.data._root_link_pose_w.data.view(wp.float32), indices=env_ids) def write_root_com_pose_to_sim_mask( self, @@ -699,7 +727,7 @@ def write_root_com_velocity_to_sim_index( self.data._root_state_w.timestamp = -1.0 self.data._root_com_state_w.timestamp = -1.0 # set into simulation - self.root_view.set_root_velocities(self._get_root_com_vel_w_f32(), indices=env_ids) + self.root_view.set_root_velocities(self.data._root_com_vel_w.data.view(wp.float32), indices=env_ids) def write_root_com_velocity_to_sim_mask( self, @@ -789,7 +817,7 @@ def write_root_link_velocity_to_sim_index( self.data._root_state_w.timestamp = -1.0 self.data._root_com_state_w.timestamp = -1.0 # set into simulation - self.root_view.set_root_velocities(self._get_root_link_vel_w_f32(), indices=env_ids) + self.root_view.set_root_velocities(self.data._root_link_vel_w.data.view(wp.float32), indices=env_ids) def write_root_link_velocity_to_sim_mask( self, @@ -913,12 +941,9 @@ def write_joint_state_to_sim_mask( joint_mask: Joint mask. If None, then all joints are used. env_mask: Environment mask. If None, then all the instances are updated. Shape is (num_instances,). """ - # resolve masks to indices (PhysX only supports index-based TensorAPI) - env_ids = self._resolve_env_mask(env_mask) - joint_ids = self._resolve_joint_mask(joint_mask) - self.write_joint_state_to_sim_index( - position=position, velocity=velocity, joint_ids=joint_ids, env_ids=env_ids, full_data=True - ) + # set into simulation + self.write_joint_position_to_sim_mask(position=position, env_mask=env_mask, joint_mask=joint_mask) + self.write_joint_velocity_to_sim_mask(velocity=velocity, env_mask=env_mask, joint_mask=joint_mask) def write_joint_position_to_sim_index( self, @@ -1152,8 +1177,7 @@ def write_joint_stiffness_to_sim_index( ) # Set into simulation, note that when updating "model" properties with PhysX we need to do it on CPU. cpu_env_ids = self._get_cpu_env_ids(env_ids) - wp.copy(self._cpu_joint_stiffness, self.data._joint_stiffness) - self.root_view.set_dof_stiffnesses(self._cpu_joint_stiffness, indices=cpu_env_ids) + self.root_view.set_dof_stiffnesses(wp.clone(self.data._joint_stiffness, device="cpu"), indices=cpu_env_ids) def write_joint_stiffness_to_sim_mask( self, @@ -1247,8 +1271,104 @@ def write_joint_damping_to_sim_index( ) # Set into simulation, note that when updating "model" properties with PhysX we need to do it on CPU. cpu_env_ids = self._get_cpu_env_ids(env_ids) - wp.copy(self._cpu_joint_damping, self.data._joint_damping) - self.root_view.set_dof_dampings(self._cpu_joint_damping, indices=cpu_env_ids) + self.root_view.set_dof_dampings(wp.clone(self.data._joint_damping, device="cpu"), indices=cpu_env_ids) + + 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 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. + """ + 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`.""" + adapter = self.newton_actuator_adapter + 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, + 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, @@ -1348,8 +1468,7 @@ def write_joint_position_limit_to_sim_index( logger.info(violation_message) # Set into simulation, note that when updating "model" properties with PhysX we need to do it on CPU. cpu_env_ids = self._get_cpu_env_ids(env_ids) - wp.copy(self._cpu_joint_pos_limits, self.data._joint_pos_limits) - self.root_view.set_dof_limits(self._cpu_joint_pos_limits, indices=cpu_env_ids) + self.root_view.set_dof_limits(wp.clone(self.data._joint_pos_limits, device="cpu"), indices=cpu_env_ids) def write_joint_position_limit_to_sim_mask( self, @@ -1453,8 +1572,7 @@ def write_joint_velocity_limit_to_sim_index( ) # Set into simulation, note that when updating "model" properties with PhysX we need to do it on CPU. cpu_env_ids = self._get_cpu_env_ids(env_ids) - wp.copy(self._cpu_joint_vel_limits, self.data._joint_vel_limits) - self.root_view.set_dof_max_velocities(self._cpu_joint_vel_limits, indices=cpu_env_ids) + self.root_view.set_dof_max_velocities(wp.clone(self.data._joint_vel_limits, device="cpu"), indices=cpu_env_ids) def write_joint_velocity_limit_to_sim_mask( self, @@ -1555,8 +1673,7 @@ def write_joint_effort_limit_to_sim_index( ) # Set into simulation, note that when updating "model" properties with PhysX we need to do it on CPU. cpu_env_ids = self._get_cpu_env_ids(env_ids) - wp.copy(self._cpu_joint_effort_limits, self.data._joint_effort_limits) - self.root_view.set_dof_max_forces(self._cpu_joint_effort_limits, indices=cpu_env_ids) + self.root_view.set_dof_max_forces(wp.clone(self.data._joint_effort_limits, device="cpu"), indices=cpu_env_ids) def write_joint_effort_limit_to_sim_mask( self, @@ -1655,8 +1772,7 @@ def write_joint_armature_to_sim_index( if isinstance(env_ids, torch.Tensor): env_ids = wp.from_torch(env_ids, dtype=wp.int32) cpu_env_ids = self._get_cpu_env_ids(env_ids) - wp.copy(self._cpu_joint_armature, self.data._joint_armature) - self.root_view.set_dof_armatures(self._cpu_joint_armature, indices=cpu_env_ids) + self.root_view.set_dof_armatures(wp.clone(self.data._joint_armature, device="cpu"), indices=cpu_env_ids) def write_joint_armature_to_sim_mask( self, @@ -1700,23 +1816,16 @@ def write_joint_friction_coefficient_to_sim_index( ): r"""Write joint friction coefficients over selected environment indices into the simulation. - For Isaac Sim versions below 5.0, only the legacy unitless joint friction coefficient is set. - This limits the resisting force or torque up to a maximum proportional to the transmitted spatial force: - :math:`\|F_{resist}\| \leq \mu_s \, \|F_{spatial}\|`. - - For Isaac Sim versions 5.0 and above, the PhysX joint friction parameter model is used. It combines - Coulomb (static and dynamic) friction with a viscous term: + For Isaac Sim versions below 5.0, only the static friction coefficient is set. + This limits the resisting force or torque up to a maximum proportional to the transmitted + spatial force: :math:`\|F_{resist}\| \leq \mu_s \, \|F_{spatial}\|`. - - Static friction effort defines the maximum effort that prevents motion at rest [N or N·m, depending on - joint type]. - - Dynamic friction effort applies once motion begins and remains constant during motion [N or N·m, - depending on joint type]. - - Viscous friction coefficient is a velocity-proportional resistive term [N·s/m or N·m·s/rad, depending - on joint type]. + For Isaac Sim versions 5.0 and above, the static, dynamic, and viscous friction coefficients + are set. The model combines Coulomb (static & dynamic) friction with a viscous term: - .. warning:: - For Isaac Sim versions 5.0 and above, the static friction effort must be greater than or equal to the - dynamic friction effort. + - Static friction :math:`\mu_s` defines the maximum effort that prevents motion at rest. + - Dynamic friction :math:`\mu_d` applies once motion begins and remains constant during motion. + - Viscous friction :math:`c_v` is a velocity-proportional resistive term. .. note:: This method expects partial data or full data. @@ -1726,12 +1835,11 @@ def write_joint_friction_coefficient_to_sim_index( is only supporting indexing, hence masks need to be converted to indices. Args: - joint_friction_coeff: Legacy unitless coefficient for Isaac Sim versions below 5.0, or static friction - effort [N or N·m, depending on joint type] for Isaac Sim versions 5.0 and above. Shape is - (len(env_ids), len(joint_ids)) or (num_instances, num_joints). - joint_dynamic_friction_coeff: Dynamic friction effort [N or N·m, depending on joint type]. + joint_friction_coeff: Static friction coefficient :math:`\mu_s`. + Shape is (len(env_ids), len(joint_ids)) or (num_instances, num_joints). + joint_dynamic_friction_coeff: Dynamic (Coulomb) friction coefficient :math:`\mu_d`. Same shape as above. If None, the dynamic coefficient is not updated. - joint_viscous_friction_coeff: Viscous friction coefficient [N·s/m or N·m·s/rad, depending on joint type]. + joint_viscous_friction_coeff: Viscous friction coefficient :math:`c_v`. Same shape as above. If None, the viscous coefficient is not updated. joint_ids: Joint indices. If None, then all joints are used. env_ids: Environment indices. If None, then all indices are used. @@ -1803,8 +1911,7 @@ def write_joint_friction_coefficient_to_sim_index( ) # Set into simulation, note that when updating "model" properties with PhysX we need to do it on CPU. cpu_env_ids = self._get_cpu_env_ids(env_ids) - wp.copy(self._cpu_joint_friction_props, friction_props) - self.root_view.set_dof_friction_properties(self._cpu_joint_friction_props, indices=cpu_env_ids) + self.root_view.set_dof_friction_properties(wp.clone(friction_props, device="cpu"), indices=cpu_env_ids) def write_joint_friction_coefficient_to_sim_mask( self, @@ -1817,23 +1924,16 @@ def write_joint_friction_coefficient_to_sim_mask( ): r"""Write joint friction coefficients over selected environment mask into the simulation. - For Isaac Sim versions below 5.0, only the legacy unitless joint friction coefficient is set. - This limits the resisting force or torque up to a maximum proportional to the transmitted spatial force: - :math:`\|F_{resist}\| \leq \mu_s \, \|F_{spatial}\|`. - - For Isaac Sim versions 5.0 and above, the PhysX joint friction parameter model is used. It combines - Coulomb (static and dynamic) friction with a viscous term: + For Isaac Sim versions below 5.0, only the static friction coefficient is set. + This limits the resisting force or torque up to a maximum proportional to the transmitted + spatial force: :math:`\|F_{resist}\| \leq \mu_s \, \|F_{spatial}\|`. - - Static friction effort defines the maximum effort that prevents motion at rest [N or N·m, depending on - joint type]. - - Dynamic friction effort applies once motion begins and remains constant during motion [N or N·m, - depending on joint type]. - - Viscous friction coefficient is a velocity-proportional resistive term [N·s/m or N·m·s/rad, depending - on joint type]. + For Isaac Sim versions 5.0 and above, the static, dynamic, and viscous friction coefficients + are set. The model combines Coulomb (static & dynamic) friction with a viscous term: - .. warning:: - For Isaac Sim versions 5.0 and above, the static friction effort must be greater than or equal to the - dynamic friction effort. + - Static friction :math:`\mu_s` defines the maximum effort that prevents motion at rest. + - Dynamic friction :math:`\mu_d` applies once motion begins and remains constant during motion. + - Viscous friction :math:`c_v` is a velocity-proportional resistive term. .. note:: This method expects full data. @@ -1843,12 +1943,11 @@ def write_joint_friction_coefficient_to_sim_mask( is only supporting indexing, hence masks need to be converted to indices. Args: - joint_friction_coeff: Legacy unitless coefficient for Isaac Sim versions below 5.0, or static friction - effort [N or N·m, depending on joint type] for Isaac Sim versions 5.0 and above. Shape is - (num_instances, num_joints). - joint_dynamic_friction_coeff: Dynamic friction effort [N or N·m, depending on joint type]. + joint_friction_coeff: Static friction coefficient :math:`\mu_s`. + Shape is (num_instances, num_joints). + joint_dynamic_friction_coeff: Dynamic (Coulomb) friction coefficient :math:`\mu_d`. Same shape as above. If None, the dynamic coefficient is not updated. - joint_viscous_friction_coeff: Viscous friction coefficient [N·s/m or N·m·s/rad, depending on joint type]. + joint_viscous_friction_coeff: Viscous friction coefficient :math:`c_v`. Same shape as above. If None, the viscous coefficient is not updated. joint_mask: Joint mask. If None, then all joints are used. env_mask: Environment mask. If None, then all the instances are updated. Shape is (num_instances,). @@ -1874,9 +1973,7 @@ def write_joint_dynamic_friction_coefficient_to_sim_index( env_ids: Sequence[int] | torch.Tensor | wp.array | None = None, full_data: bool = False, ) -> None: - """Write joint dynamic friction effort over selected environment indices into the simulation. - - The dynamic friction effort is [N or N·m, depending on joint type]. + """Write joint dynamic friction coefficient over selected environment indices into the simulation. .. note:: This method expects partial data or full data. @@ -1886,8 +1983,8 @@ def write_joint_dynamic_friction_coefficient_to_sim_index( is only supporting indexing, hence masks need to be converted to indices. Args: - joint_dynamic_friction_coeff: Dynamic friction effort [N or N·m, depending on joint type]. Shape is - (len(env_ids), len(joint_ids)) or (num_instances, num_joints) if full_data. + joint_dynamic_friction_coeff: Joint dynamic friction coefficient. Shape is (len(env_ids), len(joint_ids)) + or (num_instances, num_joints) if full_data. joint_ids: Joint indices. If None, then all joints are used. env_ids: Environment indices. If None, then all indices are used. full_data: Whether to expect full data. Defaults to False. @@ -1931,8 +2028,7 @@ def write_joint_dynamic_friction_coefficient_to_sim_index( ) # Set into simulation, note that when updating "model" properties with PhysX we need to do it on CPU. cpu_env_ids = self._get_cpu_env_ids(env_ids) - wp.copy(self._cpu_joint_friction_props, friction_props) - self.root_view.set_dof_friction_properties(self._cpu_joint_friction_props, indices=cpu_env_ids) + self.root_view.set_dof_friction_properties(wp.clone(friction_props, device="cpu"), indices=cpu_env_ids) def write_joint_dynamic_friction_coefficient_to_sim_mask( self, @@ -1941,9 +2037,7 @@ def write_joint_dynamic_friction_coefficient_to_sim_mask( joint_mask: wp.array | None = None, env_mask: wp.array | None = None, ) -> None: - """Write joint dynamic friction effort over selected environment mask into the simulation. - - The dynamic friction effort is [N or N·m, depending on joint type]. + """Write joint dynamic friction coefficient over selected environment mask into the simulation. .. note:: This method expects full data. @@ -1953,8 +2047,7 @@ def write_joint_dynamic_friction_coefficient_to_sim_mask( is only supporting indexing, hence masks need to be converted to indices. Args: - joint_dynamic_friction_coeff: Dynamic friction effort [N or N·m, depending on joint type]. Shape is - (num_instances, num_joints). + joint_dynamic_friction_coeff: Joint dynamic friction coefficient. Shape is (num_instances, num_joints). joint_mask: Joint mask. If None, then all joints are used. env_mask: Environment mask. If None, then all the instances are updated. Shape is (num_instances,). """ @@ -1979,8 +2072,6 @@ def write_joint_viscous_friction_coefficient_to_sim_index( ) -> None: """Write joint viscous friction coefficient over selected environment indices into the simulation. - The coefficient is [N·s/m or N·m·s/rad, depending on joint type]. - .. note:: This method expects partial data or full data. @@ -1989,8 +2080,8 @@ def write_joint_viscous_friction_coefficient_to_sim_index( is only supporting indexing, hence masks need to be converted to indices. Args: - joint_viscous_friction_coeff: Viscous friction coefficient [N·s/m or N·m·s/rad, depending on joint type]. - Shape is (len(env_ids), len(joint_ids)) or (num_instances, num_joints) if full_data. + joint_viscous_friction_coeff: Joint viscous friction coefficient. Shape is (len(env_ids), len(joint_ids)) + or (num_instances, num_joints) if full_data. joint_ids: Joint indices. If None, then all joints are used. env_ids: Environment indices. If None, then all indices are used. full_data: Whether to expect full data. Defaults to False. @@ -2037,8 +2128,7 @@ def write_joint_viscous_friction_coefficient_to_sim_index( ) # Set into simulation, note that when updating "model" properties with PhysX we need to do it on CPU. cpu_env_ids = self._get_cpu_env_ids(env_ids) - wp.copy(self._cpu_joint_friction_props, friction_props) - self.root_view.set_dof_friction_properties(self._cpu_joint_friction_props, indices=cpu_env_ids) + self.root_view.set_dof_friction_properties(wp.clone(friction_props, device="cpu"), indices=cpu_env_ids) def write_joint_viscous_friction_coefficient_to_sim_mask( self, @@ -2049,8 +2139,6 @@ def write_joint_viscous_friction_coefficient_to_sim_mask( ) -> None: """Write joint viscous friction coefficient over selected environment mask into the simulation. - The coefficient is [N·s/m or N·m·s/rad, depending on joint type]. - .. note:: This method expects full data. @@ -2059,8 +2147,7 @@ def write_joint_viscous_friction_coefficient_to_sim_mask( is only supporting indexing, hence masks need to be converted to indices. Args: - joint_viscous_friction_coeff: Viscous friction coefficient [N·s/m or N·m·s/rad, depending on joint type]. - Shape is (num_instances, num_joints). + joint_viscous_friction_coeff: Joint viscous friction coefficient. Shape is (num_instances, num_joints). joint_mask: Joint mask. If None, then all joints are used. env_mask: Environment mask. If None, then all the instances are updated. Shape is (num_instances,). """ @@ -2128,8 +2215,7 @@ def set_masses_index( # Set into simulation, note that when updating "model" properties with PhysX we need to do it on CPU. cpu_env_ids = self._get_cpu_env_ids(env_ids) - wp.copy(self._cpu_body_mass, self.data._body_mass) - self.root_view.set_masses(self._cpu_body_mass, indices=cpu_env_ids) + self.root_view.set_masses(wp.clone(self.data._body_mass, device="cpu"), indices=cpu_env_ids) def set_masses_mask( self, @@ -2208,11 +2294,12 @@ def set_coms_index( # Set into simulation, note that when updating "model" properties with PhysX we need to do it on CPU. # Convert from wp.transformf to flat (N, M, 7) array for PhysX cpu_env_ids = self._get_cpu_env_ids(env_ids) - wp.copy( - self._cpu_body_coms, - self.data._body_com_pose_b.data.view(wp.float32).reshape((self.num_instances, self.num_bodies, 7)), + body_com_flat = ( + wp.clone(self.data._body_com_pose_b.data, device="cpu") + .view(wp.float32) + .reshape((self.num_instances, self.num_bodies, 7)) ) - self.root_view.set_coms(self._cpu_body_coms, indices=cpu_env_ids) + self.root_view.set_coms(body_com_flat, indices=cpu_env_ids) def set_coms_mask( self, @@ -2290,8 +2377,7 @@ def set_inertias_index( ) # Set into simulation, note that when updating "model" properties with PhysX we need to do it on CPU. cpu_env_ids = self._get_cpu_env_ids(env_ids) - wp.copy(self._cpu_body_inertia, self.data._body_inertia) - self.root_view.set_inertias(self._cpu_body_inertia, indices=cpu_env_ids) + self.root_view.set_inertias(wp.clone(self.data._body_inertia, device="cpu"), indices=cpu_env_ids) def set_inertias_mask( self, @@ -3782,35 +3868,6 @@ def _create_buffers(self): device=self.device, ) - # Cached .view(wp.float32) wrappers for structured warp arrays. - # These avoid per-call wp.array metadata allocation in writers. - # Reset to None each time _create_buffers runs (during initialization). - self._root_link_pose_w_f32: wp.array | None = None - self._root_com_vel_w_f32: wp.array | None = None - self._root_link_vel_w_f32: wp.array | None = None - # Cached wrench views for write_data_to_sim - self._inst_wrench_force_f32: wp.array | None = None - self._inst_wrench_torque_f32: wp.array | None = None - self._perm_wrench_force_f32: wp.array | None = None - self._perm_wrench_torque_f32: wp.array | None = None - - # Pre-allocated pinned CPU buffers for PhysX TensorAPI writes. - # PhysX requires CPU arrays for "model" property updates (stiffness, damping, etc.). - # Pinned memory enables DMA fast path and avoids per-call malloc. - N, J, B = self.num_instances, self.num_joints, self.num_bodies - self._cpu_env_ids_all = wp.zeros(N, dtype=wp.int32, device="cpu", pinned=True) - wp.copy(self._cpu_env_ids_all, self._ALL_INDICES) - self._cpu_joint_stiffness = wp.zeros((N, J), dtype=wp.float32, device="cpu", pinned=True) - self._cpu_joint_damping = wp.zeros((N, J), dtype=wp.float32, device="cpu", pinned=True) - self._cpu_joint_pos_limits = wp.zeros((N, J, 2), dtype=wp.float32, device="cpu", pinned=True) - self._cpu_joint_vel_limits = wp.zeros((N, J), dtype=wp.float32, device="cpu", pinned=True) - self._cpu_joint_effort_limits = wp.zeros((N, J), dtype=wp.float32, device="cpu", pinned=True) - self._cpu_joint_armature = wp.zeros((N, J), dtype=wp.float32, device="cpu", pinned=True) - self._cpu_joint_friction_props = wp.zeros((N, J, 3), dtype=wp.float32, device="cpu", pinned=True) - self._cpu_body_mass = wp.zeros((N, B), dtype=wp.float32, device="cpu", pinned=True) - self._cpu_body_coms = wp.zeros((N, B, 7), dtype=wp.float32, device="cpu", pinned=True) - self._cpu_body_inertia = wp.zeros((N, B, 9), dtype=wp.float32, device="cpu", pinned=True) - def _process_cfg(self): """Post processing of configuration parameters.""" # default state @@ -3869,165 +3926,143 @@ def _process_actuators_cfg(self): """Process and apply articulation joint properties.""" # 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 + # ``sync_torque_telemetry`` kernel. ``None`` when no Newton fast path + # is active. + self._implicit_dof_mask: wp.array | None = None - # 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=self._data.joint_stiffness.torch[:, joint_ids], - damping=self._data.joint_damping.torch[:, joint_ids], - armature=self._data.joint_armature.torch[:, joint_ids], - friction=self._data.joint_friction_coeff.torch[:, joint_ids], - dynamic_friction=self._data.joint_dynamic_friction_coeff.torch[:, joint_ids], - viscous_friction=self._data.joint_viscous_friction_coeff.torch[:, joint_ids], - effort_limit=self._data.joint_effort_limits.torch[:, joint_ids].clone(), - velocity_limit=self._data.joint_vel_limits.torch[:, 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, + _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." ) - 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, + + 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, ) - 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, + + from isaaclab.sim.utils.stage import get_current_stage # noqa: PLC0415 + + # Enable the fast path even for all-implicit articulations: + # PhysX runs PD internally; Lab only forwards targets. + self._has_newton_actuators = True + + # 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() ) - 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, - ], + + # 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, ) - 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, + + 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, + ) + + # 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(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) + + 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 is_implicit: + self._create_lab_actuator(actuator_name, actuator_cfg) + else: + self._create_lab_actuator(actuator_name, actuator_cfg, properties_only=True) + + # ``_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, ) - # 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) + # 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: - # 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._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 --- + 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()) @@ -4038,8 +4073,14 @@ 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 NewtonActuatorAdapter is not None 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): @@ -4050,6 +4091,106 @@ 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, + *, + 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( + 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], + ) + + # 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 + 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) + def _process_tendons(self): """Process fixed and spatial tendons.""" # create a list to store the fixed tendon names @@ -4143,6 +4284,47 @@ def _apply_actuator_model(self): device=self.device, ) + def _apply_actuator_model_newton(self): + """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`. + """ + 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: + self.newton_actuator_adapter.step(w, w, SimulationManager.get_physics_dt()) + + wp.launch( + 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_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=[ + self._data._computed_torque, + self._data._applied_torque, + ], + device=self.device, + ) + """ Internal helpers -- Debugging. """ @@ -4355,23 +4537,17 @@ def format_limits(_, v: tuple[float, float]) -> str: ) def _get_cpu_env_ids(self, env_ids: wp.array | torch.Tensor) -> wp.array: - """Get the CPU environment indices. - - For the full-index case (all environments), returns the pre-allocated - pinned CPU buffer. For partial indices (e.g. during partial resets), clones to CPU. + """ + Get the CPU environment indices. Args: env_ids: Environment indices. Returns: - A warp array of environment indices on CPU. + A warp array of environment indices. """ if isinstance(env_ids, torch.Tensor): env_ids = wp.from_torch(env_ids, dtype=wp.int32) - # Fast path: if these are all indices, use pre-allocated pinned buffer - if env_ids.ptr == self._ALL_INDICES.ptr: - return self._cpu_env_ids_all - # Slow path: partial indices (reset), clone to CPU return wp.clone(env_ids, device="cpu") def _resolve_env_mask(self, env_mask: wp.array | None) -> torch.Tensor | wp.array: @@ -4393,55 +4569,12 @@ def _resolve_env_mask(self, env_mask: wp.array | None) -> torch.Tensor | wp.arra env_ids = self._ALL_INDICES return env_ids - def _get_root_link_pose_w_f32(self) -> wp.array: - """Get a cached float32 view of root_link_pose_w for PhysX TensorAPI. Invalidated in ``_create_buffers``.""" - if self._root_link_pose_w_f32 is None: - self._root_link_pose_w_f32 = self.data._root_link_pose_w.data.view(wp.float32) - return self._root_link_pose_w_f32 - - def _get_root_com_vel_w_f32(self) -> wp.array: - """Get a cached float32 view of root_com_vel_w for PhysX TensorAPI. Invalidated in ``_create_buffers``.""" - if self._root_com_vel_w_f32 is None: - self._root_com_vel_w_f32 = self.data._root_com_vel_w.data.view(wp.float32) - return self._root_com_vel_w_f32 - - def _get_root_link_vel_w_f32(self) -> wp.array: - """Get a cached float32 view of root_link_vel_w for PhysX TensorAPI. Invalidated in ``_create_buffers``.""" - if self._root_link_vel_w_f32 is None: - self._root_link_vel_w_f32 = self.data._root_link_vel_w.data.view(wp.float32) - return self._root_link_vel_w_f32 - - def _get_inst_wrench_force_f32(self) -> wp.array: - """Get a cached flattened float32 view of instantaneous wrench force. Invalidated in ``_create_buffers``.""" - if self._inst_wrench_force_f32 is None: - self._inst_wrench_force_f32 = self._instantaneous_wrench_composer.out_force_b.warp.flatten().view( - wp.float32 - ) - return self._inst_wrench_force_f32 - - def _get_inst_wrench_torque_f32(self) -> wp.array: - """Get a cached flattened float32 view of instantaneous wrench torque. Invalidated in ``_create_buffers``.""" - if self._inst_wrench_torque_f32 is None: - self._inst_wrench_torque_f32 = self._instantaneous_wrench_composer.out_torque_b.warp.flatten().view( - wp.float32 - ) - return self._inst_wrench_torque_f32 - - def _get_perm_wrench_force_f32(self) -> wp.array: - """Get a cached flattened float32 view of permanent wrench force. Invalidated in ``_create_buffers``.""" - if self._perm_wrench_force_f32 is None: - self._perm_wrench_force_f32 = self._permanent_wrench_composer.out_force_b.warp.flatten().view(wp.float32) - return self._perm_wrench_force_f32 - - def _get_perm_wrench_torque_f32(self) -> wp.array: - """Get a cached flattened float32 view of permanent wrench torque. Invalidated in ``_create_buffers``.""" - if self._perm_wrench_torque_f32 is None: - self._perm_wrench_torque_f32 = self._permanent_wrench_composer.out_torque_b.warp.flatten().view(wp.float32) - return self._perm_wrench_torque_f32 - def _resolve_env_ids(self, env_ids: Sequence[int] | torch.Tensor | wp.array | None) -> wp.array: """Resolve environment indices to a warp array. + .. note:: + We need to convert torch tensors to warp arrays since the TensorAPI views only support warp arrays. + Args: env_ids: Environment indices. If None, then all indices are used. @@ -4451,6 +4584,7 @@ def _resolve_env_ids(self, env_ids: Sequence[int] | torch.Tensor | wp.array | No if (env_ids is None) or (env_ids == slice(None)): return self._ALL_INDICES if isinstance(env_ids, torch.Tensor): + # Convert int64 to int32 if needed, as warp expects int32 if env_ids.dtype == torch.int64: env_ids = env_ids.to(torch.int32) return wp.from_torch(env_ids, dtype=wp.int32) @@ -4478,6 +4612,9 @@ def _resolve_joint_mask(self, joint_mask: wp.array | None) -> torch.Tensor | wp. def _resolve_joint_ids(self, joint_ids: Sequence[int] | torch.Tensor | wp.array | None) -> wp.array | torch.Tensor: """Resolve joint indices to a warp array or tensor. + .. note:: + We do not need to convert torch tensors to warp arrays since they never get passed to the TensorAPI views. + Args: joint_ids: Joint indices. If None, then all indices are used. @@ -4488,10 +4625,6 @@ def _resolve_joint_ids(self, joint_ids: Sequence[int] | torch.Tensor | wp.array return wp.array(joint_ids, dtype=wp.int32, device=self.device) if (joint_ids is None) or (joint_ids == slice(None)): return self._ALL_JOINT_INDICES - if isinstance(joint_ids, torch.Tensor): - if joint_ids.dtype == torch.int64: - joint_ids = joint_ids.to(torch.int32) - return wp.from_torch(joint_ids, dtype=wp.int32) return joint_ids def _resolve_body_mask(self, body_mask: wp.array | None) -> torch.Tensor | wp.array: @@ -4524,10 +4657,6 @@ def _resolve_body_ids(self, body_ids: Sequence[int] | torch.Tensor | wp.array | return wp.array(body_ids, dtype=wp.int32, device=self.device) if (body_ids is None) or (body_ids == slice(None)): return self._ALL_BODY_INDICES - if isinstance(body_ids, torch.Tensor): - if body_ids.dtype == torch.int64: - body_ids = body_ids.to(torch.int32) - return wp.from_torch(body_ids, dtype=wp.int32) return body_ids def _resolve_fixed_tendon_mask(self, fixed_tendon_mask: wp.array | None) -> torch.Tensor | wp.array: @@ -4740,11 +4869,14 @@ def write_joint_state_to_sim( joint_ids: Sequence[int] | torch.Tensor | wp.array | None = None, env_ids: Sequence[int] | torch.Tensor | wp.array | None = None, ): - """Deprecated, same as :meth:`write_joint_state_to_sim_index`.""" + """Deprecated, same as :meth:`write_joint_position_to_sim_index` and + :meth:`write_joint_velocity_to_sim_index`.""" warnings.warn( "The function 'write_joint_state_to_sim' will be deprecated in a future release. Please" - " use 'write_joint_state_to_sim_index' instead.", + " use 'write_joint_position_to_sim_index' and 'write_joint_velocity_to_sim_index' instead.", DeprecationWarning, stacklevel=2, ) - self.write_joint_state_to_sim_index(position=position, velocity=velocity, joint_ids=joint_ids, env_ids=env_ids) + # set into simulation + self.write_joint_position_to_sim_index(position=position, joint_ids=joint_ids, env_ids=env_ids) + self.write_joint_velocity_to_sim_index(velocity=velocity, joint_ids=joint_ids, env_ids=env_ids) diff --git a/source/isaaclab_physx/test/assets/test_newton_actuators_physx.py b/source/isaaclab_physx/test/assets/test_newton_actuators_physx.py new file mode 100644 index 000000000000..8e8d2ea134c4 --- /dev/null +++ b/source/isaaclab_physx/test/assets/test_newton_actuators_physx.py @@ -0,0 +1,1100 @@ +# 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_physx.assets import Articulation +from isaaclab_physx.physics import PhysxCfg + +import isaaclab.sim as sim_utils +from isaaclab.actuators import DCMotorCfg, DelayedPDActuatorCfg, IdealPDActuatorCfg, ImplicitActuatorCfg +from isaaclab.sim import SimulationCfg, build_simulation_context + +from isaaclab_assets import ANYMAL_C_CFG + +# --------------------------------------------------------------------------- +# 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_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"], + 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, + feedforward: float | None = None, +) -> dict: + """Run ANYmal-C on PhysX and return recorded trajectories + telemetry. + + 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) + 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) + 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, + "computed_torque": recorded_computed, + "applied_torque": recorded_applied, + "target_pos": target_pos.clone(), + "target_vel": target_vel.clone(), + } + + +# --------------------------------------------------------------------------- +# 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 = {} + feedforward: float | None = None + pos_atol: float = 2e-3 + pos_rtol: float = 1e-3 + 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, + 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, + 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_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}", + ) + + +# --------------------------------------------------------------------------- +# 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 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). + + Verifies that implicit actuators (handled by PhysX joint drives) + coexist correctly with explicit Newton actuators via PhysxActuatorWrapper. + """ + + __test__ = True + 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 + + +class TestImplicitWithFeedforwardEquivalencePhysx(_EquivalenceTestBase): + """Implicit-only actuators with a non-zero feedforward effort target on PhysX.""" + + __test__ = True + actuators = IMPLICIT_ONLY_ACTUATORS + feedforward = 5.0 + + +# --------------------------------------------------------------------------- +# 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}", + ) + + +# --------------------------------------------------------------------------- +# Domain randomization via events.py — PhysX 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.the 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 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. + """ + + @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, + ) 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 = 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) + + +# --------------------------------------------------------------------------- +# Per-env reset: actuator state isolation +# --------------------------------------------------------------------------- + +RESET_WARMUP_STEPS = 3 + + +class TestActuatorStateReset(unittest.TestCase): + """Reset must clear the actuator state buffers for the requested envs only. + + Inspects ``adapter.actuators[i].state.delay_state.num_pushes`` directly: + + * 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``. + + Done independently on Lab and Newton paths. PhysX-side adapter is + per-articulation, available via ``articulation.newton_actuator_adapter``. + """ + + 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)) + art_cfg = ANYMAL_C_CFG.replace( + actuators=DELAYED_PD_ACTUATORS, + prim_path="/World/Env_.*/Robot", + ) + articulation = Articulation(art_cfg) + sim.reset() + + 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) + 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") + + 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", + ) + + 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", + ) + + articulation.reset(env_ids=torch.tensor([self.RESET_ENV], device=articulation.device, dtype=torch.long)) + + self.assertTrue( + torch.all(buf[:, self.RESET_ENV] == 0).item(), + f"Lab: env {self.RESET_ENV} buffer not zeroed after reset.", + ) + self.assertTrue( + (buf[:, self.UNCHANGED_ENV] != 0).any().item(), + f"Lab: env {self.UNCHANGED_ENV} buffer was zeroed — reset leaked into an unselected env.", + ) + finally: + ctx.__exit__(None, None, None) + + +# --------------------------------------------------------------------------- +# 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 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. + + 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_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) + + 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, + } + ) + } + torch.jit.save(scripted, tmp_path, _extra_files=extra) + return tmp_path + + +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) + + 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_path, _extra_files=extra) + return tmp_path + + +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_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_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()