# 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 physics manager for Isaac Lab."""
from __future__ import annotations
import contextlib
import gc
import inspect
import logging
import re
from abc import abstractmethod
from collections.abc import Callable, Sequence
from functools import partial
from typing import TYPE_CHECKING, Any, ClassVar
import numpy as np
import torch
import warp as wp
from newton import (
Axis,
CollisionPipeline,
Contacts,
Control,
Heightfield,
Model,
ModelBuilder,
ModelFlags,
ShapeFlags,
State,
eval_fk,
)
from newton.selection import ArticulationView
from newton.sensors import SensorContact as NewtonContactSensor
from newton.sensors import SensorFrameTransform
from newton.sensors import SensorIMU as NewtonSensorIMU
from newton.solvers import SolverBase, SolverMuJoCo
from newton.usd import SchemaResolver, SchemaResolverMjc, SchemaResolverNewton, SchemaResolverPhysx
from pxr import Usd, UsdGeom
import isaaclab.cloner as cloner
from isaaclab.physics import CallbackHandle, PhysicsEvent, PhysicsManager
from isaaclab.scene_data import SceneDataBackend, SceneDataFormat, SceneDataProvider
from isaaclab.sim import SimulationContext
from isaaclab.utils import checked_apply, to_dict
from isaaclab.utils.buffers import TimestampedBuffer
from isaaclab.utils.string import resolve_matching_names
from isaaclab.utils.timer import Timer
from isaaclab.utils.version import has_kit
from isaaclab.utils.warp.index_kernel import IndexKernelDispatcher
from isaaclab_newton.physics.featherstone_manager_cfg import FeatherstoneSolverCfg
from isaaclab_newton.physics.mjwarp_manager_cfg import MJWarpSolverCfg
from isaaclab_newton.physics.newton_manager_cfg import (
NewtonBackendCfg,
NewtonBuilderCfg,
NewtonCfg,
NewtonShapeCfg,
NewtonSolverCfg,
)
from isaaclab_newton.physics.xpbd_manager_cfg import XPBDSolverCfg
from isaaclab_newton.renderers.visual_material import (
VisualMaterialWriter,
VisualShapeColorWriter,
)
if TYPE_CHECKING:
from isaaclab.actuators.newton import NewtonActuatorAdapter
from isaaclab.assets import BaseArticulation
from isaaclab.renderers.base_renderer import VisualMaterialBatch
from isaaclab_newton.physics.newton_collision_cfg import NewtonCollisionPipelineCfg
_SENSORS_BY_STATE_ATTRIBUTE = {
"body_qdd": "the IMU or PVA sensor",
"body_parent_f": "the joint-wrench sensor",
}
"""Which Isaac Lab sensors read each state attribute that MuJoCo's sensor stage fills."""
_SENSOR_STAGE_STATE_ATTRIBUTES = frozenset(_SENSORS_BY_STATE_ATTRIBUTE)
"""Extended state attributes that MuJoCo Warp's sensor stage fills via ``rne_postconstraint``."""
def _compile_label_pattern(expr: str | list[str] | None) -> re.Pattern[str] | None:
"""Compile selector expressions for Newton's full label matching."""
if not expr:
return None
return re.compile("|".join((expr,) if isinstance(expr, str) else expr))
logger = logging.getLogger(__name__)
# Tagged union for entries in _cl_site_index_map.
# _GlobalSite: (global_shape_idx, None) — body_pattern was None
# _LocalSite: (None, [[env0_idx, ...], ...]) — per-world site indices
@wp.kernel(enable_backward=False)
def _or_reset_masks_from_mask(
env_mask: wp.array(dtype=wp.bool),
articulation_ids: wp.array2d(dtype=int),
world_mask: wp.array(dtype=wp.bool),
fk_mask: wp.array(dtype=wp.bool),
):
"""OR env_mask into world_mask and set corresponding articulation bits in fk_mask."""
world, arti = wp.tid()
if env_mask[world]:
world_mask[world] = True
fk_mask[articulation_ids[world, arti]] = True
@wp.kernel(enable_backward=False)
def _scatter_reset_masks_from_ids(
env_ids: wp.array(dtype=Any),
articulation_ids: wp.array2d(dtype=int),
world_mask: wp.array(dtype=wp.bool),
fk_mask: wp.array(dtype=wp.bool),
):
"""Scatter-set world_mask and fk_mask from sparse env_ids."""
i, arti = wp.tid()
world = wp.int32(env_ids[i])
world_mask[world] = True
fk_mask[articulation_ids[world, arti]] = True
_SCATTER_RESET_MASKS_FROM_IDS_DISPATCHER = IndexKernelDispatcher(_scatter_reset_masks_from_ids, ("env_ids",))
def _scatter_reset_masks_from_ids_kernel(env_ids: wp.array | torch.Tensor) -> wp.Kernel:
"""Select the reset-mask writer matching the environment selector dtype."""
return _SCATTER_RESET_MASKS_FROM_IDS_DISPATCHER.select(env_ids)
@wp.kernel(enable_backward=False)
def _or_world_reset_mask_from_mask(env_mask: wp.array(dtype=wp.bool), world_mask: wp.array(dtype=wp.bool)):
"""Mark masked worlds for solver reset without requesting FK."""
world = wp.tid()
if env_mask[world]:
world_mask[world] = True
@wp.kernel(enable_backward=False)
def _scatter_world_reset_mask_from_ids(env_ids: wp.array(dtype=wp.int32), world_mask: wp.array(dtype=wp.bool)):
"""Mark selected worlds for solver reset without requesting FK."""
world_mask[env_ids[wp.tid()]] = True
[docs]
def create_newton_builder(cfg: NewtonBuilderCfg) -> ModelBuilder:
"""Construct a native builder through the selected physics manager's factory.
Args:
cfg: The selected physics configuration.
Returns:
An empty builder with the selected solver schemas and shape/BVH defaults.
"""
if isinstance(cfg.physics_cfg, NewtonCfg):
return cfg.physics_cfg.class_type.create_builder(physics_cfg=cfg.physics_cfg)
return ModelBuilder()
class NewtonBackend:
"""Own one finalized Newton model and its native state and control buffers."""
def __init__(self, cfg: NewtonBackendCfg):
builder = SimulationContext.instance().get_or_create_backend(NewtonBuilderCfg(physics_cfg=cfg.physics_cfg))
self.model = builder.finalize(device=cfg.device)
self.particle_ranges: dict[str, tuple[int, int]] = {}
# Newton 1.6 preserves groups through builder replication but not finalization.
# Remove this snapshot when the pinned Newton includes newton-physics/newton#3326.
self.deformable_ranges = {
label: (start, end - start, kind)
for family, kind in (("cloth", "surface"), ("soft", "volume"))
for label, start, end in zip(
getattr(builder, f"_{family}_label"),
getattr(builder, f"_{family}_particle_start"),
getattr(builder, f"_{family}_particle_end"),
strict=True,
)
}
self.model.num_envs = self.model.world_count
simulation = isinstance(cfg.physics_cfg, NewtonCfg)
soft_contact = cfg.physics_cfg.soft_contact_cfg if simulation else None
if soft_contact is not None:
self.model.soft_contact_ke = float(soft_contact.soft_contact_ke)
self.model.soft_contact_kd = float(soft_contact.soft_contact_kd)
self.model.soft_contact_mu = float(soft_contact.soft_contact_mu)
self.state_0 = self.model.state()
self.state_1 = self.model.state() if simulation else None
self.control = self.model.control() if simulation else None
self.geometry_offsets = (
{} if simulation else {path: bounds[0] for path, bounds in self.deformable_ranges.items()}
)
self.bvh_refit = TimestampedBuffer()
def create_visual_material_writer(self, batches: tuple[VisualMaterialBatch, ...]) -> VisualMaterialWriter:
"""Bind material writes to this resource's native shape-color buffer."""
return VisualMaterialWriter(self.model, batches)
def close(self) -> None:
"""Drop native handles after consumers release their bindings."""
self.control = self.state_1 = self.state_0 = self.model = None
self.deformable_ranges.clear()
self.particle_ranges = {}
self.geometry_offsets = {}
self.bvh_refit = TimestampedBuffer()
class NewtonQueries:
"""Stateless query and capture functions operating on an explicit Newton resource."""
@staticmethod
def run_query(
backend: NewtonBackend,
timestamp: int,
query: Callable[[], None],
graph: tuple[tuple[int, ...], wp.Graph] | None,
*,
use_cuda_graph: bool = True,
) -> tuple[tuple[int, ...], wp.Graph] | None:
"""Refresh shared BVHs once per publication and run a consumer-owned query graph.
Args:
backend: Shared model, state, and acceleration structures.
timestamp: Sum of the monotonic SDP transform and geometry timestamps for this state.
query: Camera or ray-cast work using the supplied backend.
graph: This consumer's previous captured query and pointer layout, or None.
use_cuda_graph: Whether to capture GPU queries. CPU queries always run eagerly.
Returns:
The consumer's captured query and pointer layout, or None for eager execution.
"""
model, state = backend.model, backend.state_0
pointers = tuple(array.ptr if array is not None else 0 for array in (state.body_q, state.particle_q))
cached = backend.bvh_refit
if cached.timestamp != timestamp:
if model.device.is_cuda and use_cuda_graph:
if cached.data is None or cached.data[0] != pointers:
refit = partial(NewtonQueries._refit_bvh, backend)
cached.data = pointers, NewtonQueries.capture_graph(str(model.device), refit, relaxed=has_kit())
wp.capture_launch(cached.data[1])
else:
NewtonQueries._refit_bvh(backend)
cached.timestamp = timestamp
if not model.device.is_cuda or not use_cuda_graph:
query()
return None
if graph is None or graph[0] != pointers:
graph = pointers, NewtonQueries.capture_graph(str(model.device), query, relaxed=has_kit())
wp.capture_launch(graph[1])
return graph
@staticmethod
def capture_graph(device: str, capture_target: Callable[[], None], *, relaxed: bool = False) -> wp.Graph:
"""Record work without an eager warmup or graph replay, using Warp's public capture API."""
stream = wp.get_stream(device)
mode = wp.CaptureMode.THREAD_LOCAL
if relaxed:
# RTX uses the legacy CUDA stream. A nonblocking stream avoids implicit synchronization.
stream = wp.stream_from_torch(torch.cuda.Stream(device=device))
mode = wp.CaptureMode.RELAXED
with NewtonQueries._paused_gc(), wp.ScopedStream(stream):
with wp.ScopedCapture(stream=stream, capture_mode=mode) as capture:
capture_target()
return capture.graph
@staticmethod
def _refit_bvh(backend: NewtonBackend) -> None:
model, state = backend.model, backend.state_0
if model.shape_count:
model.bvh_refit_shapes(state)
if model.particle_count:
model.bvh_refit_particles(state)
@staticmethod
@contextlib.contextmanager
def _paused_gc():
"""Keep collection-driven frees out of CUDA graph capture."""
was_enabled = gc.isenabled()
gc.disable()
try:
yield
finally:
if was_enabled:
gc.enable()
gc.collect()
class NewtonSceneDataBackend(SceneDataBackend):
"""Scene data backend that reads rigid body transforms from Newton's simulation state.
The backend reads ``body_q`` (an array of :class:`wp.transformf`) from
Newton's current state and exposes it as :class:`SceneDataFormat.Transform`.
Body paths come from the model's ``body_label`` attribute.
"""
def __init__(self):
self._transforms = SceneDataFormat.Transform()
self.transforms_timestamp = 0
self.geometry_timestamp = 0
self._geometry_batches = []
def initialize_geometry(self) -> None:
"""Bind imported geometry paths to native particle ranges and capsule endpoints."""
state = NewtonManager.get_state_0()
self._geometry_batches = [
(source, ranges)
for source, ranges in self._geometry_batches
if source._cls is not SceneDataFormat.CapsuleEndpoints
]
endpoints, ranges = [], {}
offset = 0
for path, shapes in NewtonManager._cable_bindings.items():
ids = np.asarray(shapes, dtype=np.int32)
left = np.concatenate((ids[:1], ids))
right = np.concatenate((ids, ids[-1:]))
left_sign = np.ones(len(left), dtype=np.int32)
right_sign = -np.ones(len(right), dtype=np.int32)
left_sign[0], right_sign[-1] = -1, 1
endpoints.append(np.column_stack((left, left_sign, right, right_sign)))
ranges[path] = (offset, len(left))
offset += len(left)
if endpoints:
model = self.model
source = SceneDataFormat.CapsuleEndpoints()
source.transforms = state.body_q
source.shape_body = model.shape_body
source.shape_transform = model.shape_transform
source.shape_scale = model.shape_scale
source.endpoints = wp.array(np.concatenate(endpoints), dtype=wp.vec4i, device=model.device)
self._geometry_batches.append((source, ranges))
@property
def native_geometry_formats(self) -> tuple[type, ...]:
return (SceneDataFormat.Points, SceneDataFormat.WeightedPoints, SceneDataFormat.CapsuleEndpoints)
def get_geometry_batches(self, output_format=SceneDataFormat.Points):
"""Publish native arrays; SDP derives cable endpoints and applies destination layouts."""
state = self.state
for source, _ in self._geometry_batches:
attribute = "transforms" if source._cls is SceneDataFormat.CapsuleEndpoints else "points"
data = state.particle_q if attribute == "points" else state.body_q
if getattr(source, attribute) is not data:
setattr(source, attribute, data)
self.geometry_timestamp += 1
return self._geometry_batches
@property
def transforms(self) -> SceneDataFormat.Transform:
"""Publish the authoritative native pointer, including solver state-buffer swaps."""
transforms = self.state.body_q
if self._transforms.transforms is not transforms:
self._transforms.transforms = transforms
self.transforms_timestamp += 1
return self._transforms
@property
def transform_count(self) -> int:
"""Return the number of rigid body transforms in the Newton sim."""
return self.model.body_count
@property
def transform_paths(self) -> list[str]:
"""Return the prim paths for each rigid body transform."""
if self.model.body_label is not None:
return list(self.model.body_label)
return []
@property
def model(self) -> Model:
return NewtonManager.get_model()
@property
def state(self) -> State:
"""Return native physics state without entering the rendering consumer path."""
if NewtonManager.transforms_may_change_on_graph_replay:
# Raw external graph replays bypass Python invalidation, so these reads must stay conservative.
self.transforms_timestamp += 1
self.geometry_timestamp += 1
NewtonManager.forward()
return NewtonManager.get_state_0()
def _eval_fk_unbound(world_reset_mask: wp.array | None, fk_mask: wp.array | None) -> None:
"""Default :attr:`NewtonManager._eval_fk` value before a solver is initialized.
Raises so a stray ``forward()`` / ``step()`` before ``initialize_solver()`` fails loudly
instead of silently running a wrong (or no) FK.
"""
raise RuntimeError(
"FK hook is not bound. NewtonManager.initialize_solver() must run "
"(via reset()) before forward()/step() can run forward kinematics."
)
def _reset_solver_internals_unbound(world_mask: wp.array | None) -> None:
"""Default reset-hook delegate value before a solver is initialized."""
raise RuntimeError(
"Solver reset hook is not bound. NewtonManager.initialize_solver() must run "
"(via reset()) before forward()/step() can reset solver internals."
)
[docs]
class NewtonManager(PhysicsManager):
"""Abstract Newton physics manager for Isaac Lab.
Class-level manager for physics lifecycle, solvers, contacts, sensors, replication, and
CUDA graphs. Model, state, and control allocations belong to the registered native backend.
Concrete subclasses (one per solver) implement :meth:`_build_solver` and
may extend :meth:`_initialize_contacts`, :meth:`_prepare_builder_for_finalize`,
:meth:`_step_solver`, :meth:`_supports_cuda_graph_capture`,
:meth:`_reset_solver_internals`,
:meth:`_solver_specific_clear`, :meth:`_check_solver_status`, and
:meth:`_log_solver_debug`.
Subclasses are selected via :attr:`NewtonSolverCfg.class_type`, which
:meth:`NewtonCfg.__post_init__` propagates onto :attr:`NewtonCfg.class_type`
so that ``SimulationContext`` resolves the matching subclass automatically.
Lifecycle: ``initialize() -> reset() -> step()`` (repeated) ``-> close()``.
.. note::
Assign shared lifecycle state through :class:`NewtonManager`, not ``cls``.
This keeps :attr:`backend` and solver state visible to every consumer without
shadowing the base attributes on a concrete solver subclass.
"""
@classmethod
def provides_implicit_damping(cls) -> bool:
# Newton's symplectic integrator has no implicit damping.
return False
[docs]
@classmethod
def setup_deformable_body(cls, prim: Any, deformable_type: str, sim_mesh_prim: Any, vis_mesh_prim: Any) -> None:
"""Apply Newton's token deformable anchor schemas and sync the visual mesh geometry."""
sim_mesh_path = sim_mesh_prim.GetPath().pathString
token = "PhysicsVolumeDeformableSimAPI" if deformable_type == "volume" else "PhysicsSurfaceDeformableSimAPI"
if not sim_mesh_prim.AddAppliedSchema(token):
raise RuntimeError(f"Failed to set {deformable_type} deformable sim API on prim '{sim_mesh_path}'.")
# Newton renders the simulation mesh directly: overwrite the visual mesh geometry until
# separate visual/simulation meshes are supported.
vis_mesh = UsdGeom.Mesh(vis_mesh_prim)
if deformable_type == "volume":
tet_mesh = UsdGeom.TetMesh(sim_mesh_prim)
surface_indices = tet_mesh.GetSurfaceFaceVertexIndicesAttr().Get()
if surface_indices is None or len(surface_indices) == 0:
raise ValueError(
f"Deformable body at '{prim.GetPath().pathString}' has no surface indices on its TetMesh"
" prim; cannot sync to visual mesh."
)
vis_mesh.GetPointsAttr().Set(tet_mesh.GetPointsAttr().Get())
vis_mesh.GetFaceVertexIndicesAttr().Set(np.asarray(surface_indices).flatten())
vis_mesh.GetFaceVertexCountsAttr().Set([3] * len(surface_indices))
else:
sim_mesh = UsdGeom.Mesh(sim_mesh_prim)
vis_mesh.GetFaceVertexIndicesAttr().Set(sim_mesh.GetFaceVertexIndicesAttr().Get())
vis_mesh.GetFaceVertexCountsAttr().Set(sim_mesh.GetFaceVertexCountsAttr().Get())
if not prim.AddAppliedSchema("PhysicsDeformableBodyAPI"):
raise RuntimeError(f"Failed to set deformable body API on prim '{prim.GetPath().pathString}'.")
_solver_dt: float = 1.0 / 200.0
_num_substeps: int = 1
_decimation: int = 1
_collision_decimation: int = 0
_deterministic_mode: wp.DeterministicMode = wp.DeterministicMode.NOT_GUARANTEED
_num_envs: int | None = None
_supports_rigid_body_force_input: bool = False
"""Whether the solver consumes applied rigid-body forces from :class:`State`."""
# Native ownership and physics orchestration
backend: ClassVar[NewtonBackend | None] = None
"""Borrowed native resource shared by physics and scene consumers; the simulation registry owns it."""
_solver: SolverBase | None = None
_use_single_state: bool | None = None
"""Use only one state for both input and output for solver stepping. Requires solver support."""
# Physics settings
_gravity_vector: tuple[float, float, float] = (0.0, 0.0, -9.81)
_up_axis: str = "Z"
# Collision and contacts
_contacts: Contacts | None = None
_needs_collision_pipeline: bool = False
_collision_pipeline = None
_collision_cfg: NewtonCollisionPipelineCfg | None = None
_newton_contact_sensors: dict = {} # Maps sensor_key to NewtonContactSensor
_newton_frame_transform_sensors: list = [] # List of SensorFrameTransform
_newton_imu_sensors: list = [] # List of NewtonSensorIMU
_pending_extended_state_attributes: set[str] = set()
_active_extended_state_attributes: set[str] = set()
_pending_extended_contact_attributes: set[str] = set()
_report_contacts: bool = False
_supports_contact_sensors: bool = True
# Per-world reset masks (allocated in start_simulation, consumed in step/forward).
# Newton reserves the final slot for global entities in world -1.
_world_reset_mask: wp.array | None = None # (num_envs + 1,) wp.bool
_fk_reset_mask: wp.array | None = None # (articulation_count,) wp.bool — for eval_fk(mask=...)
kinematics_dirty: bool = False
# Solver-specialized FK delegate. Bound in initialize_solver() to the active subclass's choice of FK implementation.
_eval_fk: Callable[[wp.array | None, wp.array | None], None] = _eval_fk_unbound
# Solver-specialized reset delegate. Like _eval_fk, this must dispatch correctly through the base manager.
_reset_solver_internals_delegate: Callable[[wp.array | None], None] = _reset_solver_internals_unbound
# 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]] = []
# In-graph hooks invoked immediately before every solver substep.
_state_force_callbacks: list[Callable[[State], None]] = []
# In-graph hooks invoked after the last solver substep and before sensors,
# in registration order. Articulations with non-identity ordering register
# their backend-to-user state republish kernels here so the reorders are
# recorded into every captured graph.
_post_step_callbacks: list[Callable[[], None]] = []
# CUDA graphing
_graph = None
_graph_capture_pending: bool = False
# Native scene-data publication.
transforms_may_change_on_graph_replay: bool = False
_cable_bindings: dict[str, list[int]] = {}
# Model changes (callbacks use unified system from PhysicsManager)
_model_changes: set[int] = set()
# Model changes the active solver does not apply, mapped to the reason; each is warned about once
_ignored_model_changes: dict[int, str] = {}
_warned_model_changes: set[int] = set()
# Scene data backend
_scene_data_backend: NewtonSceneDataBackend | None = None
_builder_attribute_solvers: tuple[type[SolverBase], ...] = ()
_particle_ranges: dict[str, tuple[int, int]] = {}
# CL: Cloning / Replication logic
# TODO: These attributes support cloning-specific logic and should be moved into a cloner class
# Pending site requests from sensors.
# Key: (body_pattern, per_world, xform_floats), Value: (label, wp.transform)
# identical (body_pattern, per_world, transform) reuses the same site.
_cl_pending_sites: dict[tuple[str | None, bool, tuple[float, ...]], tuple[str, wp.transform]] = {}
# Maps each site label to its resolved global or local site entry.
_GlobalSite = tuple[int, None]
_LocalSite = tuple[None, list[list[int]]]
_SiteEntry = _GlobalSite | _LocalSite
_cl_site_index_map: dict[str, _SiteEntry] = {}
_world_xforms: list[wp.transform] | None = None
# Per-source builders retained from replication, keyed by clone-plan source
# path. Single-model consumers (e.g. batched Newton IK) finalize a single-env
# model from these using the asset prototype's native source path.
_cl_protos: dict[str, ModelBuilder] = {}
_per_world_builder_hooks: list[Callable[[ModelBuilder, int, np.ndarray, np.ndarray], None]] = []
[docs]
@classmethod
def initialize(cls, sim_context: SimulationContext) -> None:
"""Initialize the manager with simulation context.
Args:
sim_context: Parent simulation context.
"""
super().initialize(sim_context)
# This context imports NewtonManager, so it can only be imported after this module initializes.
from isaaclab_newton.cloner import NewtonReplicateContext # noqa: PLC0415
cls.clone_context_type = NewtonReplicateContext
# Newton-specific setup: get gravity from SimulationCfg (not physics manager cfg)
sim = PhysicsManager._sim
if sim is not None:
NewtonManager._gravity_vector = sim.cfg.gravity # type: ignore[union-attr]
NewtonManager._scene_data_backend = NewtonSceneDataBackend()
[docs]
@classmethod
def reset(cls, soft: bool = False) -> None:
"""Reset physics simulation.
A hard reset (``soft=False``) re-finalizes the Newton model, reallocating
its device arrays. The cached collision pipeline, contacts and any
captured CUDA graph reference the old buffers, so they are released here
and rebuilt against the re-finalized model by :meth:`initialize_solver`.
This avoids the illegal CUDA memory access (CUDA error 700) that would
otherwise occur on the first step after a hard reset.
A soft reset (``soft=True``) skips this full reinitialization and reuses
the existing model, solver, collision pipeline and CUDA graph.
Args:
soft: If True, skip full reinitialization.
"""
if not soft:
if NewtonManager.backend is not None:
cls.dispatch_event(PhysicsEvent.STOP)
# Release the cached collision pipeline, contacts and CUDA graph;
# they point at the old model's freed buffers (CUDA 700 on next step).
NewtonManager._graph = None
NewtonManager._graph_capture_pending = False
NewtonManager._collision_pipeline = None
NewtonManager._contacts = None
NewtonManager._solver = None
NewtonManager._eval_fk = _eval_fk_unbound
NewtonManager._reset_solver_internals_delegate = _reset_solver_internals_unbound
NewtonManager._adapter = None
if NewtonManager.backend is not None:
SimulationContext.instance().close_backend(NewtonManager.backend)
NewtonManager.backend = None
cls.start_simulation()
cls.initialize_solver()
@classmethod
def _eval_fk_impl(cls, world_reset_mask: wp.array | None, fk_mask: wp.array | None) -> None:
"""Update body states from joint coordinates.
Solver-specialized FK implementation. The base implementation runs Newton's generic
``eval_fk`` over the articulations selected by ``fk_mask``. Subclasses may override
this method to use a solver-specific FK.
Args:
world_reset_mask: Per-world mask of environments to reset (``None`` means all).
Unused by the base implementation; consumed by solver-specific overrides such as
:meth:`NewtonKaminoManager._eval_fk_impl`.
fk_mask: Per-articulation mask of articulations to update (``None`` means all).
"""
backend = cls.backend
eval_fk(backend.model, backend.state_0.joint_q, backend.state_0.joint_qd, backend.state_0, fk_mask)
[docs]
@classmethod
def forward(cls) -> None:
"""Update articulation kinematics without stepping physics.
Update body poses from joint coordinates via the solver-specialized FK delegate
(:attr:`_eval_fk`, bound to the active subclass's :meth:`_eval_fk_impl` in
:meth:`initialize_solver`). Only the articulations flagged dirty in
:attr:`_fk_reset_mask` and :attr:`_world_reset_mask` (see :meth:`invalidate_fk`) are
updated. The masks are consumed (zeroed) afterwards so the next :meth:`step` does not
redundantly re-solve them.
Asset and scene-data reads share the same pending work. The bound delegate dispatches
calls on ``NewtonManager`` to the active solver's implementation.
"""
if not (cls.kinematics_dirty or cls.transforms_may_change_on_graph_replay):
return
cls._reset_solver_internals_delegate(cls._world_reset_mask)
cls._eval_fk(cls._world_reset_mask, cls._fk_reset_mask)
if cls._fk_reset_mask is not None:
cls._fk_reset_mask.zero_()
if cls._world_reset_mask is not None:
cls._world_reset_mask.zero_()
NewtonManager.kinematics_dirty = False
[docs]
@classmethod
def video_capture_backend(cls) -> str:
"""Newton GL headless perspective video capture."""
return "newton_gl"
@classmethod
def _mark_transforms_changed(cls) -> None:
"""Publish authored rigid-body changes and invalidate cable geometry."""
if NewtonManager._scene_data_backend is not None:
NewtonManager._scene_data_backend.transforms_timestamp += 1
NewtonManager._scene_data_backend.geometry_timestamp += 1
device = PhysicsManager._device
if device is not None:
device = wp.get_device(device)
if device.is_cuda and device.stream.is_capturing:
NewtonManager.transforms_may_change_on_graph_replay = True
@classmethod
def _mark_particles_dirty(cls) -> None:
"""Invalidate SDP geometry after native particle writes."""
NewtonManager._scene_data_backend.geometry_timestamp += 1
device = wp.get_device(PhysicsManager._device)
if device.is_cuda and device.stream.is_capturing:
NewtonManager.transforms_may_change_on_graph_replay = True
[docs]
@classmethod
def step(cls) -> None:
"""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
cls._reset_solver_internals_delegate(cls._world_reset_mask)
# Notify solver of model changes
if cls._model_changes:
with wp.ScopedDevice(PhysicsManager._device):
for change in cls._model_changes:
if change in cls._ignored_model_changes and change not in cls._warned_model_changes:
logger.warning(cls._ignored_model_changes[change])
cls._warned_model_changes.add(change)
cls._solver.notify_model_changed(change)
NewtonManager._model_changes = set()
# Reset-authored state and persistent solver resources must be ready before capture.
cls.forward()
cfg = PhysicsManager._cfg
device = PhysicsManager._device
if cls._graph_capture_pending and cfg is not None and cfg.use_cuda_graph:
simulate = cls._simulate_full if cls._is_all_graphable() else cls._simulate_physics_only
with Timer(name="newton_cuda_graph", msg="CUDA graph took:"):
NewtonManager._graph = cls._capture_graph(simulate)
NewtonManager._graph_capture_pending = False
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:
# --- Some actuators not graph-safe: step them eagerly, graph solver only ---
if cls._adapter is not None:
cls._adapter.step(cls.backend.state_0, cls.backend.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
cls._mark_transforms_changed()
cls._check_solver_status()
# Launch solver-specific debug logging after stepping.
cls._log_solver_debug()
[docs]
@classmethod
def close(cls) -> None:
"""Clean up Newton physics resources."""
super().close()
cls.clear()
[docs]
@classmethod
def get_scene_data_backend(cls) -> SceneDataBackend | None:
"""Return the SceneDataBackend for the SceneDataProvider."""
return cls._scene_data_backend
[docs]
@classmethod
def register_callback(
cls,
callback: Callable,
event: PhysicsEvent,
order: int = 0,
name: str | None = None,
wrap_weak_ref: bool = True,
) -> CallbackHandle:
"""Register a callback. Passes event to parent class."""
return PhysicsManager.register_callback(callback, event, order, name, wrap_weak_ref)
[docs]
@classmethod
def get_physics_sim_view(cls) -> list:
"""Return the registered articulation views."""
return [view for (manager, _), view in cls.views.items() if manager is NewtonManager]
[docs]
@classmethod
def is_fabric_enabled(cls) -> bool:
"""Check if fabric interface is enabled (not applicable for Newton)."""
return False
[docs]
@classmethod
def clear(cls):
"""Clear all Newton-specific state (callbacks cleared by super().close())."""
NewtonManager._num_envs = None
NewtonManager._solver = None
NewtonManager._use_single_state = None
NewtonManager._supports_rigid_body_force_input = False
NewtonManager._contacts = None
NewtonManager._needs_collision_pipeline = False
NewtonManager._ignored_model_changes = {}
NewtonManager._warned_model_changes = set()
NewtonManager._deterministic_mode = wp.DeterministicMode.NOT_GUARANTEED
NewtonManager._eval_fk = _eval_fk_unbound
NewtonManager._reset_solver_internals_delegate = _reset_solver_internals_unbound
NewtonManager._collision_pipeline = None
NewtonManager._collision_cfg = None
NewtonManager._newton_contact_sensors = {}
NewtonManager._newton_frame_transform_sensors = []
NewtonManager._newton_imu_sensors = []
NewtonManager._report_contacts = False
NewtonManager._supports_contact_sensors = True
NewtonManager._adapter = None
NewtonManager._post_actuator_callbacks = []
NewtonManager._state_force_callbacks = []
NewtonManager._post_step_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
NewtonManager.kinematics_dirty = False
NewtonManager._graph = None
NewtonManager._graph_capture_pending = False
NewtonManager.transforms_may_change_on_graph_replay = False
NewtonManager._cable_bindings = {}
NewtonManager._particle_ranges = {}
NewtonManager._per_world_builder_hooks = []
NewtonManager._up_axis = "Z"
NewtonManager._model_changes = set()
NewtonManager._scene_data_backend = None
NewtonManager._cl_pending_sites = {}
NewtonManager._cl_site_index_map = {}
NewtonManager._world_xforms = None
NewtonManager._cl_protos = {}
NewtonManager._pending_extended_state_attributes = set()
NewtonManager._active_extended_state_attributes = set()
NewtonManager._pending_extended_contact_attributes = set()
for key in [key for key in NewtonManager.views if key[0] is NewtonManager]:
del NewtonManager.views[key]
cls._solver_specific_clear()
if NewtonManager.backend is not None:
SimulationContext.instance().close_backend(NewtonManager.backend)
NewtonManager.backend = None
[docs]
@classmethod
def create_builder(
cls, up_axis: str | None = None, *, physics_cfg: NewtonCfg | None = None, **kwargs
) -> ModelBuilder:
"""Create a :class:`ModelBuilder` configured with default settings.
Forwards :class:`NewtonShapeCfg` defaults onto Newton's upstream
``ModelBuilder.default_shape_cfg`` via :func:`~isaaclab.utils.checked_apply`.
Falls back to wrapper defaults when no Newton config is active so
rough-terrain margin/gap still apply during early construction.
Args:
up_axis: Override for the up-axis. Defaults to ``None``, which uses
the manager's ``_up_axis``.
physics_cfg: Explicit builder settings; None uses the active physics configuration.
**kwargs: Forwarded to :class:`ModelBuilder`.
Returns:
New builder with up-axis and per-shape defaults (gap, margin) applied.
"""
# Resolve which NewtonShapeCfg to apply: user override if active config
# is NewtonCfg, else the wrapper's own defaults so callers from non-Newton
# contexts (tests, early construction) still get the rough-terrain margin.
cfg = PhysicsManager._cfg if physics_cfg is None else physics_cfg
builder = ModelBuilder(up_axis=up_axis or cls._up_axis, **kwargs)
builder.default_bvh_cfg = ModelBuilder.BvhConfig(
mesh_constructor=cfg.bvh_constructor_geometry if isinstance(cfg, NewtonCfg) else None,
gaussian_constructor=cfg.bvh_constructor_gaussian if isinstance(cfg, NewtonCfg) else None,
shape_constructor=cfg.bvh_constructor_scene if isinstance(cfg, NewtonCfg) else None,
shape_flags=ShapeFlags.VISIBLE,
)
cls._register_builder_attributes(builder)
shape_cfg = cfg.default_shape_cfg if isinstance(cfg, NewtonCfg) else NewtonShapeCfg()
checked_apply(shape_cfg, builder.default_shape_cfg)
return builder
@classmethod
def _register_builder_attributes(cls, builder: ModelBuilder) -> None:
"""Register custom attributes required by the active solver."""
for solver_cls in cls._builder_attribute_solvers:
solver_cls.register_custom_attributes(builder)
@classmethod
def _registers_builder_attributes_from_solver(cls, solver_cls: type[SolverBase]) -> bool:
"""Return whether this manager registers custom attributes from ``solver_cls``."""
return solver_cls in cls._builder_attribute_solvers
@classmethod
def _prepare_builder_for_finalize(cls, builder: ModelBuilder) -> None:
"""Subclass hook to normalize *builder* before model finalization.
Override in solver subclasses that need to adapt imported or replicated
builder data before :meth:`ModelBuilder.finalize` allocates model arrays.
The default implementation is a no-op.
"""
[docs]
@classmethod
def cl_register_site(cls, body_pattern: str | None, xform: wp.transform, *, per_world: bool = False) -> str:
"""Register a site request for injection into prototypes before replication.
Sensors call this during ``__init__``. Sites are injected into prototype
builders by :meth:`_cl_inject_sites` (called from ``newton_replicate``)
before ``add_builder``, so they replicate correctly per-world.
Identical ``(body_pattern, per_world, transform)`` registrations share sites.
The *body_pattern* is matched against prototype-local body labels
(e.g. ``"Robot/link.*"``) when replication is active, or against the
flat builder's body labels in the fallback path. Wildcard patterns
that match multiple bodies create one site per matched body.
Args:
body_pattern: Regex pattern matched against body labels in the
prototype builder (e.g. ``"Robot/link0"`` or ``"Robot/finger.*"``
for multi-body wildcards), or ``None`` for global sites
(world-origin reference, etc.).
xform: Site transform relative to body.
per_world: When ``True``, ``body_pattern`` must be ``None`` and one
bodyless site is created in each cloned world's frame.
Returns:
Assigned site label suffix.
"""
if per_world and body_pattern is not None:
raise ValueError("per_world site registration requires body_pattern=None.")
xform_key = tuple(xform)
key = (body_pattern, per_world, xform_key)
if key in cls._cl_pending_sites:
return cls._cl_pending_sites[key][0]
label = f"ft_{len(cls._cl_pending_sites)}"
cls._cl_pending_sites[key] = (label, xform)
return label
[docs]
@classmethod
def request_extended_state_attribute(cls, attr: str) -> None:
"""Request an extended state attribute (e.g. ``"body_qdd"``).
Sensors call this during ``__init__``, before model finalization.
Attributes are forwarded to the builder in :meth:`start_simulation`
so that subsequent ``model.state()`` calls allocate them.
Args:
attr: State attribute name (must be in ``State.EXTENDED_ATTRIBUTES``).
"""
cls._pending_extended_state_attributes.add(attr)
@classmethod
def _cl_inject_sites(
cls,
main_builder: ModelBuilder,
source_builders: dict[str, ModelBuilder],
) -> tuple[dict[str, int], dict[int, dict[str, list[int]]], dict[str, wp.transform]]:
"""Inject registered sites into source builders or an explicitly supplied main builder.
Body sites are matched against source builders, then the main builder containing
shared assets. Bodyless global sites are added to *main_builder* with ``body=-1``.
Returns builder-local shape indices so that ``newton_replicate`` can
compute final indices during replication without a second pattern match.
Pending requests are cleared after processing.
Args:
main_builder: Top-level builder that receives global sites.
source_builders: ``{source_path: ModelBuilder}`` source builders.
Returns:
Tuple of ``(global_site_indices, source_site_indices, env_root_sites)`` where
*global_site_indices* maps ``{label: main_builder_shape_idx}``,
*source_site_indices* maps ``{id(builder): {label: [builder_shape_idx, ...]}}``,
and *env_root_sites* maps ``{label: env_root_relative_transform}``.
"""
global_site_indices: dict[str, int] = {}
source_site_indices: dict[int, dict[str, list[int]]] = {}
env_root_sites: dict[str, wp.transform] = {}
for (body_pattern, per_world, _xform_key), (label, xform) in cls._cl_pending_sites.items():
if per_world:
env_root_sites[label] = xform
continue
if body_pattern is None:
site_idx = main_builder.add_site(body=-1, xform=xform, label=label)
global_site_indices[label] = site_idx
continue
any_matched = False
for source_builder in (*source_builders.values(), main_builder):
if source_builder is main_builder and any_matched:
break
matched_indices, matched_names = resolve_matching_names(
body_pattern, source_builder.body_label, raise_when_no_match=False
)
if not matched_indices: # Pattern has no matches in this source builder
continue
any_matched = True
site_indices: list[int] = []
for body_idx, body_name in zip(matched_indices, matched_names):
site_label = f"{body_name}/{label}"
source_site_idx = source_builder.add_site(body=body_idx, xform=xform, label=site_label)
site_indices.append(source_site_idx)
logger.debug(f"Injected site '{site_label}' into builder")
source_site_indices.setdefault(id(source_builder), {})[label] = site_indices
if not any_matched:
raise ValueError(f"Site '{label}' with body_pattern '{body_pattern}' matched no builder bodies.")
cls._cl_pending_sites.clear()
return global_site_indices, source_site_indices, env_root_sites
[docs]
@classmethod
def add_model_change(cls, change: ModelFlags) -> None:
"""Register a model change to notify the solver."""
cls._model_changes.add(change)
[docs]
@classmethod
def invalidate_fk(
cls, env_mask: wp.array | None = None, env_ids: wp.array | None = None, articulation_ids: wp.array | None = None
) -> None:
"""Mark environments as needing FK recomputation and solver reset.
Called by asset write methods that modify joint coordinates or root
transforms. The masks are consumed by the next forward, raw-state,
rendering, or physics-step boundary.
Args:
env_mask: Boolean mask of dirtied environments. Shape ``(num_envs,)``.
Used by ``_mask`` write methods.
env_ids: Integer indices of dirtied environments.
Used by ``_index`` write methods.
articulation_ids: Mapping from ``(world, arti)`` to model articulation
index. Shape ``(world_count, count_per_world)``. Obtained from
``ArticulationView.articulation_ids``.
"""
cls._mark_transforms_changed()
if cls._world_reset_mask is None or cls._fk_reset_mask is None:
return
NewtonManager.kinematics_dirty = True
if articulation_ids is not None and env_mask is not None:
wp.launch(
_or_reset_masks_from_mask,
dim=articulation_ids.shape,
inputs=[env_mask, articulation_ids],
outputs=[NewtonManager._world_reset_mask, NewtonManager._fk_reset_mask],
device=PhysicsManager._device,
)
elif articulation_ids is not None and env_ids is not None:
wp.launch(
_scatter_reset_masks_from_ids_kernel(env_ids),
dim=(env_ids.shape[0], articulation_ids.shape[1]),
inputs=[env_ids, articulation_ids],
outputs=[NewtonManager._world_reset_mask, NewtonManager._fk_reset_mask],
device=PhysicsManager._device,
)
else:
# Fallback: no topology info — mark everything dirty
NewtonManager._world_reset_mask[: cls.backend.model.world_count].fill_(True)
NewtonManager._fk_reset_mask.fill_(True)
[docs]
@classmethod
def invalidate_body_state(
cls, env_ids: wp.array(dtype=wp.int32) | None = None, env_mask: wp.array(dtype=wp.bool) | None = None
) -> None:
"""Mark selected maximal-coordinate body state as changed without requesting FK.
Args:
env_ids: Integer indices of dirtied environments. Used by index write methods.
env_mask: Boolean mask of dirtied environments. Used by mask write methods.
"""
cls._mark_transforms_changed()
if cls._world_reset_mask is None:
return
NewtonManager.kinematics_dirty = True
if env_mask is not None:
wp.launch(
_or_world_reset_mask_from_mask,
dim=env_mask.shape[0],
inputs=[env_mask],
outputs=[NewtonManager._world_reset_mask],
device=PhysicsManager._device,
)
elif env_ids is not None:
wp.launch(
_scatter_world_reset_mask_from_ids,
dim=env_ids.shape[0],
inputs=[env_ids],
outputs=[NewtonManager._world_reset_mask],
device=PhysicsManager._device,
)
else:
NewtonManager._world_reset_mask[: cls.backend.model.world_count].fill_(True)
@classmethod
def _drain_stale_cuda_error(cls) -> None:
"""Clear a stale CUDA error latched on the device before (re)initialization.
Warp 1.15 leaves the per-thread CUDA error uncleared when
``wp_free_device_async`` fails to add a graph memory free node while a
capture is still registered (its "capture ended" sibling branch clears
the identical error as benign), and the next Warp array copy then
surfaces that stale error as its own failure, aborting simulation
initialization. Draining here keeps a prior simulation lifecycle's
latched error from poisoning this one. Remove once the upstream Warp
fix lands.
"""
device = wp.get_device(str(PhysicsManager._device))
if not device.is_cuda:
return
# Private Warp API: the drain primitives are not exposed publicly; getting the
# device above guarantees the runtime is initialized. Guard the whole private
# interaction so a future Warp internals reshuffle degrades to a skipped drain
# rather than hard-failing simulation start.
try:
from warp._src.context import runtime as _wp_runtime
core = _wp_runtime.core
# wp_cuda_context_check drains via cudaGetLastError() but returns the
# post-sync error state (0 once a non-sticky error was drained), and its
# internal check_cuda() prints the drained error verbatim to stderr;
# suppress the print and diff Warp's error buffer to report the drain.
before = core.wp_get_error_string()
was_enabled = bool(core.wp_is_error_output_enabled())
core.wp_set_error_output_enabled(0)
try:
persistent = core.wp_cuda_context_check(device.context)
finally:
core.wp_set_error_output_enabled(1 if was_enabled else 0)
after = core.wp_get_error_string()
except (ImportError, AttributeError) as exc:
logger.warning("Skipping stale CUDA error drain; Warp internals unavailable: %s", exc)
return
if persistent != 0:
logger.error(
"CUDA error %d persists after drain; the device context is likely unrecoverable: %s",
persistent,
after.decode(errors="replace"),
)
elif after != before:
logger.warning(
"Drained stale CUDA error latched by a prior lifecycle: %s (last Warp error recorded before drain: %s)",
after.decode(errors="replace"),
before.decode(errors="replace") or "<none>",
)
[docs]
@classmethod
def start_simulation(cls) -> None:
"""Start simulation by finalizing model and initializing state.
This function finalizes the model and initializes the simulation state.
Note: Collision pipeline is initialized later in initialize_solver() after
we determine whether the solver needs external collision detection.
"""
sim = SimulationContext.instance()
cfg = NewtonBackendCfg(physics_cfg=sim.cfg.physics, device=sim.device)
builder = sim.get_or_create_backend(NewtonBuilderCfg(physics_cfg=cfg.physics_cfg))
cls._drain_stale_cuda_error()
logger.info("Dispatching MODEL_INIT callbacks")
cls.dispatch_event(PhysicsEvent.MODEL_INIT)
# Explicit builders and MODEL_INIT callbacks use the same site import as clone prototypes.
if cls._cl_pending_sites:
global_sites, body_sites, world_sites = cls._cl_inject_sites(builder, {})
cls._cl_site_index_map.update((label, (index, None)) for label, index in global_sites.items())
cls._cl_site_index_map.update(
(label, (None, [indices])) for label, indices in body_sites.get(id(builder), {}).items()
)
for label, xform in world_sites.items():
index = builder.add_site(body=-1, xform=xform, label=label)
cls._cl_site_index_map[label] = (None, [[index]])
device = PhysicsManager._device
logger.info(f"Finalizing model on device: {device}")
builder.up_axis = Axis.from_string(cls._up_axis)
# Forward pending extended attribute requests to builder and clear them. The requests are
# retained because initialize_solver() runs afterwards and must know which sensors depend
# on state that MuJoCo's sensor stage fills.
if cls._pending_extended_state_attributes:
builder.request_state_attributes(*cls._pending_extended_state_attributes)
NewtonManager._active_extended_state_attributes |= cls._pending_extended_state_attributes
NewtonManager._pending_extended_state_attributes = set()
cls._prepare_builder_for_finalize(builder)
builder.request_contact_attributes(*sorted(cls._pending_extended_contact_attributes))
with Timer(
name="newton_finalize_builder",
msg="Finalize builder took:",
activity="Finalizing physics model",
synchronize="both",
device=device,
):
NewtonManager.backend = sim.get_or_create_backend(cfg)
cls.backend.particle_ranges = cls._particle_ranges
cls.backend.model.set_gravity(cls._gravity_vector)
if cls._num_envs is not None:
cls.backend.model.num_envs = cls._num_envs
NewtonManager._num_envs = cls.backend.model.num_envs
NewtonManager._pending_extended_contact_attributes = set()
# The initial body-state update from joint coordinates is deferred to the tail of
# initialize_solver(), where it runs through the solver-specialized FK delegate after the solver is initialized.
# 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
# Newton's final reset-mask slot selects global entities in world -1.
# Isaac Lab resets local environments only, so that slot remains false.
NewtonManager._world_reset_mask = wp.zeros(cls.backend.model.world_count + 1, dtype=wp.bool, device=device)
NewtonManager._fk_reset_mask = wp.zeros(cls.backend.model.articulation_count, dtype=wp.bool, device=device)
cls._scene_data_backend.initialize_geometry()
logger.info("Dispatching PHYSICS_READY callbacks")
cls.dispatch_event(PhysicsEvent.PHYSICS_READY)
@classmethod
def _inject_terrain_heightfields(
cls, stage: Usd.Stage, builder: ModelBuilder, *, root_paths: Sequence[str]
) -> list[str]:
"""Replace height-field-tagged terrain colliders with Newton heightfields.
Scans the stage for prims carrying the ``newton:heightfield:resolution``
attribute authored by :class:`~isaaclab.terrains.TerrainImporter`. For each,
the collision mesh is rasterized into a :class:`newton.Heightfield` through
:meth:`newton.Heightfield.create_from_mesh` and added to *builder* as a
static heightfield shape. The tagged prim paths are returned so the caller
can exclude them from ``add_usd`` -- otherwise the terrain would be imported
twice (once as a mesh, once as a heightfield).
Heightfields compile on the MuJoCo solver roughly two orders of magnitude
faster than the equivalent multi-hundred-thousand-vertex terrain mesh while
colliding identically at the same horizontal resolution.
Args:
stage: The USD stage being imported.
builder: The Newton model builder receiving the heightfield shapes.
root_paths: Concrete subtree roots to scan.
Returns:
Prim paths of terrain colliders that were converted to heightfields.
"""
ignore_paths: list[str] = []
xform_cache = UsdGeom.XformCache()
for prim in (prim for root_path in root_paths for prim in Usd.PrimRange(stage.GetPrimAtPath(root_path))):
attr = prim.GetAttribute("newton:heightfield:resolution")
if not attr or not attr.HasAuthoredValue():
continue
resolution = float(attr.Get())
# Locate the collision mesh under the tagged prim.
if prim.IsA(UsdGeom.Mesh):
mesh_prim = prim
else:
mesh_prim = next((p for p in Usd.PrimRange(prim) if p.IsA(UsdGeom.Mesh)), None)
if mesh_prim is None:
continue
mesh = UsdGeom.Mesh(mesh_prim)
points = np.asarray(mesh.GetPointsAttr().Get(), dtype=np.float64)
faces = np.asarray(mesh.GetFaceVertexIndicesAttr().Get(), dtype=np.int32)
# Transform vertices into world frame (USD uses row-vector convention).
mat = np.array(xform_cache.GetLocalToWorldTransform(mesh_prim), dtype=np.float64).reshape(4, 4)
world = (points @ mat[:3, :3] + mat[3, :3]).astype(np.float32)
device = str(PhysicsManager._device)
wp_mesh = wp.Mesh(
points=wp.array(world, dtype=wp.vec3, device=device),
indices=wp.array(faces, dtype=wp.int32, device=device),
)
heightfield, xform = Heightfield.create_from_mesh(wp_mesh, resolution)
builder.add_shape_heightfield(heightfield=heightfield, xform=xform)
logger.info(
"Converted terrain collider %s (%d faces) to a %dx%d heightfield.",
prim.GetPath().pathString,
faces.shape[0] // 3,
heightfield.nrow,
heightfield.ncol,
)
ignore_paths.append(prim.GetPath().pathString)
return ignore_paths
@classmethod
def _get_usd_import_schema_resolvers(cls) -> list[SchemaResolver]:
"""Return ordered schema resolvers for physics-model USD imports.
MJC is enabled for managers that register ``SolverMuJoCo`` attributes.
Visualization and articulation-ordering builders keep their fixed pair
because solver attributes do not affect their outputs.
"""
resolvers: list[SchemaResolver] = [SchemaResolverNewton(), SchemaResolverPhysx()]
if cls._registers_builder_attributes_from_solver(SolverMuJoCo):
resolvers.append(SchemaResolverMjc())
return resolvers
[docs]
@classmethod
def instantiate_builder_from_stage(cls):
"""Import the explicitly declared clone plan into the Newton builder."""
cloner.replicate(PhysicsManager._sim.get_clone_plan())
@classmethod
def _initialize_contacts(cls) -> None:
"""Initialize contacts using Newton's :class:`CollisionPipeline`.
This default implementation handles solvers that rely on Newton's
unified collision pipeline (XPBD, Featherstone, and MuJoCo with
``use_mujoco_contacts=False``). Solver subclasses with internal
contact handling (e.g. :class:`NewtonMJWarpManager` when
``use_mujoco_contacts=True``) override this method to allocate a
:class:`Contacts` object sized to the solver's internal contact buffer.
"""
if not cls._needs_collision_pipeline:
return
pipeline_args = {"broad_phase": "explicit"}
if cls._collision_cfg is not None:
pipeline_args = cls._collision_cfg.to_pipeline_args()
pipeline_args["deterministic"] = cls._deterministic_mode != wp.DeterministicMode.NOT_GUARANTEED
if cls._collision_pipeline is None:
NewtonManager._collision_pipeline = CollisionPipeline(cls.backend.model, **pipeline_args)
if cls._contacts is None:
NewtonManager._contacts = cls._collision_pipeline.contacts()
# Grow the collision-pipeline contact buffer to the solver's max when the
# solver (e.g. MuJoCo/mujoco_warp) requires more contacts than the pipeline
# auto-estimate. Without this, the RSL-RL sensor path (use_mujoco_contacts=
# False) sizes _contacts from the pipeline alone and solver.update_contacts()
# raises when naconmax (nconmax * num_envs) exceeds rigid_contact_max.
# Mirrors the mjwarp_manager.py override for the use_mujoco_contacts=True path.
_solver = cls._solver
if _solver is not None and hasattr(_solver, "get_max_contact_count"):
_need = _solver.get_max_contact_count()
if _need > NewtonManager._contacts.rigid_contact_max:
if cls._deterministic_mode != wp.DeterministicMode.NOT_GUARANTEED:
# In deterministic mode, CollisionPipeline sizes _sort_key_array from rigid_contact_max at
# construction. Rebuild so the sort and contact buffers retain matching capacity; replacing
# Contacts alone would leave the sorting buffer undersized.
pipeline_args["rigid_contact_max"] = _need
NewtonManager._collision_pipeline = CollisionPipeline(cls.backend.model, **pipeline_args)
NewtonManager._contacts = cls._collision_pipeline.contacts()
else:
NewtonManager._contacts = Contacts(
rigid_contact_max=_need,
soft_contact_max=0,
device=PhysicsManager._device,
requested_attributes=cls.backend.model.get_requested_contact_attributes(),
)
# ----- Solver construction (subclass contract) ------------------------
@classmethod
def _create_solver(cls, model: Model, solver_cfg) -> SolverBase:
"""Construct a solver without changing the active manager state.
Solver-manager subclasses override this hook so nested consumers can
reuse their typed construction logic through ``solver_cfg.class_type``.
"""
raise NotImplementedError(f"{cls.__name__} does not implement solver construction.")
@classmethod
@abstractmethod
def _build_solver(cls, model: Model, solver_cfg) -> None:
"""Construct the solver this manager owns and assign it onto the base class.
Subclasses must populate the canonical :class:`NewtonManager` slots:
* :attr:`NewtonManager._solver` — the constructed :class:`SolverBase`
instance.
* :attr:`NewtonManager._use_single_state` — ``True`` if the solver
steps in-place on a single :class:`State` (e.g. MuJoCo); ``False``
if it needs separate input/output states (e.g. XPBD, Featherstone,
Kamino).
* :attr:`NewtonManager._needs_collision_pipeline` — ``True`` if the
manager owns Newton's :class:`CollisionPipeline` for contact
generation; ``False`` if the solver runs internal collision
detection (MuJoCo internal contacts, Kamino with its own detector).
* :attr:`NewtonManager._supports_rigid_body_force_input` — ``True`` if
the solver consumes external rigid-body forces from
:attr:`State.body_f`; ``False`` otherwise.
Writing through ``NewtonManager._foo`` (rather than ``cls._foo``)
keeps the canonical state visible to external readers regardless of
which subclass is active.
Args:
model: Finalized Newton model the solver should run on.
solver_cfg: The manager-specific :class:`NewtonSolverCfg`
subclass (i.e. the inner ``cfg.solver_cfg``, not the outer
:class:`NewtonCfg`).
"""
raise NotImplementedError("NewtonManager subclasses must implement _build_solver()")
@staticmethod
def _filter_solver_kwargs(solver_cls: type, solver_cfg) -> dict:
"""Return cfg fields that match ``solver_cls.__init__`` parameters.
Drops keys that the solver constructor doesn't accept (e.g. cfg-only
metadata like ``solver_type`` / ``class_type``). ``self`` and ``model``
are always excluded — ``model`` is passed positionally at construction.
"""
valid = set(inspect.signature(solver_cls.__init__).parameters) - {"self", "model"}
kwargs = {k: v for k, v in to_dict(solver_cfg).items() if k in valid}
if "deterministic" in valid:
kwargs["deterministic"] = NewtonManager._deterministic_mode
return kwargs
@classmethod
def _validate_deterministic_solver_cfg(
cls, solver_cfg: NewtonSolverCfg, deterministic_mode: wp.DeterministicMode
) -> None:
"""Validate that a solver can provide the requested determinism guarantee."""
if deterministic_mode == wp.DeterministicMode.NOT_GUARANTEED:
return
solver_cfg_type = type(solver_cfg).__name__
if not isinstance(solver_cfg, (FeatherstoneSolverCfg, MJWarpSolverCfg, XPBDSolverCfg)):
raise ValueError(
f"Newton deterministic mode {deterministic_mode.name} is not supported by {solver_cfg_type}. "
"Use MJWarp on the GPU, XPBD, or Featherstone, or disable deterministic mode."
)
if getattr(solver_cfg, "use_mujoco_cpu", False):
raise ValueError(
f"Newton deterministic mode {deterministic_mode.name} is not supported by the MuJoCo CPU backend. "
"Set MJWarpSolverCfg.use_mujoco_cpu=False or disable deterministic mode."
)
if isinstance(solver_cfg, MJWarpSolverCfg) and not solver_cfg.disable_sensors:
raise ValueError(
f"Newton deterministic mode {deterministic_mode.name} is not supported while MuJoCo Warp's "
"internal sensor computation is enabled. Set MJWarpSolverCfg.disable_sensors=True or disable "
"deterministic mode."
)
blocked = cls._active_extended_state_attributes & _SENSOR_STAGE_STATE_ATTRIBUTES
if isinstance(solver_cfg, MJWarpSolverCfg) and blocked:
sensors = sorted({_SENSORS_BY_STATE_ATTRIBUTE[attr] for attr in blocked})
raise ValueError(
f"This task does not support deterministic physics: it uses {' and '.join(sensors)},"
f" reading {sorted(blocked)}. Those attributes come from MuJoCo's post-constraint pass,"
" which runs inside the sensor stage that a determinism guarantee must disable, so the"
" values would never be refreshed. Remove the sensors, or drop the determinism request"
f" (deterministic_mode={deterministic_mode.name})."
)
@classmethod
def _apply_deterministic_request(cls, cfg: NewtonCfg) -> wp.DeterministicMode:
"""Translate the backend-agnostic determinism request into Newton settings.
:attr:`~isaaclab.physics.PhysicsCfg.deterministic` is the generic request. An explicitly
set :attr:`~isaaclab_newton.physics.NewtonCfg.deterministic_mode` is the more specific
instruction and wins. MuJoCo on the CPU is already reproducible and Warp's deterministic
mode does not reach it, so no mode is applied there and the request is logged instead.
Args:
cfg: Resolved Newton configuration.
Returns:
The deterministic mode to apply to the solver.
"""
solver_cfg = cfg.solver_cfg
# MuJoCo-C is reproducible on its own and Warp's deterministic mode never reaches it, so
# no Newton setting applies -- report the request rather than dropping it silently.
if getattr(solver_cfg, "use_mujoco_cpu", False):
if cfg.deterministic or cfg.deterministic_mode != "not_guaranteed":
logger.info("MuJoCo CPU backend is already reproducible; Newton's deterministic mode is not applied.")
return wp.DeterministicMode.NOT_GUARANTEED
# Precedence: an explicit mode, then the generic request, then no guarantee. An explicit
# mode's prerequisites stay the caller's responsibility, so it is returned untouched.
if cfg.deterministic_mode != "not_guaranteed":
return cls._resolve_deterministic_mode(cfg.deterministic_mode)
if not cfg.deterministic:
return wp.DeterministicMode.NOT_GUARANTEED
# MuJoCo Warp cannot honour a guarantee while its internal sensor kernels run, so the
# generic request implies the prerequisite rather than failing on it.
if isinstance(solver_cfg, MJWarpSolverCfg):
solver_cfg.disable_sensors = True
return wp.DeterministicMode.RUN_TO_RUN
@staticmethod
def _resolve_deterministic_mode(deterministic_mode: str) -> wp.DeterministicMode:
"""Convert a Newton config value to Warp's deterministic-mode enum."""
return {
"not_guaranteed": wp.DeterministicMode.NOT_GUARANTEED,
"run_to_run": wp.DeterministicMode.RUN_TO_RUN,
"gpu_to_gpu": wp.DeterministicMode.GPU_TO_GPU,
}[deterministic_mode]
@classmethod
def _step_solver(
cls, state_0: State, state_1: State, control: Control, contacts: Contacts | None, substep_dt: float
) -> None:
"""Run one solver substep.
Default invokes :attr:`_solver` once. Subclasses can override to
batch multiple solvers within a single substep.
"""
cls._solver.step(state_0, state_1, control, contacts, substep_dt)
@classmethod
def _solver_specific_clear(cls) -> None:
"""Solver-specific cleanup hook called from :meth:`clear`.
Default no-op. Subclasses override to release sub-solver references
or other solver-specific resources.
"""
@classmethod
def _check_solver_status(cls) -> None:
"""Raise solver-specific asynchronous failures after stepping.
Default no-op. Subclasses override when a solver requires a host-side
status check after CUDA graph replay.
"""
@classmethod
def _log_solver_debug(cls) -> None:
"""Solver-specific debug logging after stepping.
Default no-op. Subclasses override to log solver-specific debug info
(e.g. constraint violations, contact forces, etc.) after stepping.
"""
@classmethod
def _reset_solver_internals(cls, world_mask: wp.array | None) -> None:
"""Clear solver-internal state for environments reset since the last boundary.
The hook runs immediately before reset masks are consumed by :meth:`step`
and :meth:`forward`. The base implementation delegates to
:meth:`SolverBase.reset` with ``flags=0``, preserving the joint state
authored by Isaac Lab while clearing solver-owned buffers. Solvers with
no reset implementation are unaffected.
Args:
world_mask: Per-world reset mask, or ``None`` when no simulation
state is available.
"""
if world_mask is None:
return
cls._solver.reset(cls.backend.state_0, world_mask=world_mask, flags=0)
# ----- Lifecycle orchestration ----------------------------------------
[docs]
@classmethod
def initialize_solver(cls) -> None:
"""Initialize the solver and collision pipeline.
Construct the solver and contacts, establish the initial body state, and schedule
graph capture for the first step after the environment has authored its initial state.
Initialization and capture do not advance physics.
"""
cfg = PhysicsManager._cfg
if cfg is None:
return
with Timer(
name="newton_initialize_solver",
msg="Initialize solver took:",
activity="Initializing solver",
synchronize="both",
device=cls.get_device(),
):
NewtonManager._num_substeps = cfg.num_substeps # type: ignore[union-attr]
NewtonManager._collision_decimation = cfg.collision_decimation # type: ignore[union-attr]
deterministic_mode = cls._apply_deterministic_request(cfg) # type: ignore[arg-type]
cls._validate_deterministic_solver_cfg(cfg.solver_cfg, deterministic_mode) # type: ignore[union-attr]
NewtonManager._deterministic_mode = deterministic_mode
NewtonManager._solver_dt = cls.get_physics_dt() / cls._num_substeps
NewtonManager._collision_cfg = cfg.collision_cfg # type: ignore[union-attr]
cls._build_solver(cls.backend.model, cfg.solver_cfg) # type: ignore[union-attr]
if NewtonManager._solver is None:
raise RuntimeError(
f"{cls.__name__}._build_solver did not assign NewtonManager._solver. "
"Subclasses of NewtonManager must populate NewtonManager._solver, "
"NewtonManager._use_single_state, NewtonManager._needs_collision_pipeline, and "
"NewtonManager._supports_rigid_body_force_input."
)
cls._initialize_contacts()
# Picking callbacks must be registered after the concrete solver has
# published its force-input capability, but before CUDA graph capture.
sim = PhysicsManager._sim
if NewtonManager._supports_rigid_body_force_input and sim is not None:
sim._prepare_newton_visualizer_for_capture()
# Bind the solver-specialized FK delegate to the active subclass's _eval_fk_impl so
# that forward()/step() dispatch correctly even when forward() is invoked through the
# base class (the data layer imports NewtonManager directly). ``cls`` is the concrete
# subclass here, since initialize_solver is reached via sim.physics_manager.reset().
NewtonManager._eval_fk = cls._eval_fk_impl
NewtonManager._reset_solver_internals_delegate = cls._reset_solver_internals
# Establish the initial kinematically-consistent body state through the
# solver-specialized FK delegate, now that the solver and the delegate both exist.
cls._eval_fk(None, None)
cls._mark_transforms_changed()
cls._invalidate_graph()
@classmethod
def _invalidate_graph(cls) -> None:
"""Schedule capture after reset and decimation setup, without advancing physics."""
NewtonManager._graph = None
NewtonManager._graph_capture_pending = False
cfg = PhysicsManager._cfg
device = PhysicsManager._device
if cfg is None or device is None or not cfg.use_cuda_graph or "cuda" not in device:
return
if not cls._supports_cuda_graph_capture():
logger.warning(
"%s does not support CUDA graph capture for the current solver configuration; using eager execution.",
cls.__name__,
)
return
NewtonManager._graph_capture_pending = True
@classmethod
def _supports_cuda_graph_capture(cls) -> bool:
"""Return whether the active solver configuration supports CUDA graph capture."""
return True
@classmethod
def _capture_graph(cls, capture_target: Callable[[], None]) -> wp.Graph:
"""Record physics or scene queries without an eager warmup or graph replay."""
sim = PhysicsManager._sim
relaxed = has_kit() and (sim.has_gui or sim.has_offscreen_render)
return NewtonQueries.capture_graph(PhysicsManager._device, capture_target, relaxed=relaxed)
# ------------------------------------------------------------------
# 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."""
backend = cls.backend
collide_every = cls._collision_decimation
# Last substep is skipped: its contact set would only feed the next tick's
# top-of-loop collide(), not this one.
collide_mid_loop = collide_every > 0 and cls._needs_collision_pipeline and contacts is not None
if cls._use_single_state:
for i in range(cls._num_substeps):
for callback in cls._state_force_callbacks:
callback(backend.state_0)
cls._step_solver(backend.state_0, backend.state_0, backend.control, contacts, cls._solver_dt)
backend.state_0.clear_forces()
if collide_mid_loop and (i + 1) % collide_every == 0 and i + 1 < cls._num_substeps:
cls._collision_pipeline.collide(backend.state_0, contacts)
else:
cfg = PhysicsManager._cfg
need_copy_on_last = cfg is not None and cls._num_substeps % 2 == 1
for i in range(cls._num_substeps):
for callback in cls._state_force_callbacks:
callback(backend.state_0)
cls._step_solver(backend.state_0, backend.state_1, backend.control, contacts, cls._solver_dt)
if need_copy_on_last and i == cls._num_substeps - 1:
backend.state_0.assign(backend.state_1)
else:
backend.state_0, backend.state_1 = backend.state_1, backend.state_0
backend.state_0.clear_forces()
if collide_mid_loop and (i + 1) % collide_every == 0 and i + 1 < cls._num_substeps:
cls._collision_pipeline.collide(backend.state_0, contacts)
@classmethod
def _update_sensors(cls, contacts) -> None:
"""Push latest state to all registered Newton sensors."""
for sensor in cls._newton_frame_transform_sensors:
sensor.update(cls.backend.state_0)
for sensor in cls._newton_imu_sensors:
sensor.update(cls.backend.state_0)
if cls._report_contacts:
eval_contacts = contacts if contacts is not None else cls._contacts
cls._solver.update_contacts(eval_contacts, cls.backend.state_0)
for sensor in cls._newton_contact_sensors.values():
sensor.update(cls.backend.state_0, eval_contacts)
# ------------------------------------------------------------------
# Composite stepping routines
# ------------------------------------------------------------------
@classmethod
def _simulate_full(cls) -> None:
"""Run ``decimation x (actuators + solver substeps)``, then sensors.
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.
"""
physics_dt = cls._solver_dt * cls._num_substeps
contacts = cls._contacts if cls._needs_collision_pipeline else None
for i in range(cls._decimation):
if cls._needs_collision_pipeline:
cls._collision_pipeline.collide(cls.backend.state_0, cls._contacts)
if cls._adapter is not None:
cls._adapter.step(
cls.backend.state_0,
cls.backend.control,
physics_dt,
swap_state=cls._decimation % 2 == 0 or i < cls._decimation - 1,
)
for cb in cls._post_actuator_callbacks:
cb()
cls._run_solver_substeps(contacts)
for cb in cls._post_step_callbacks:
cb()
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.backend.state_0, cls._contacts)
contacts = cls._contacts
else:
contacts = None
cls._run_solver_substeps(contacts)
for cb in cls._post_step_callbacks:
cb()
cls._update_sensors(contacts)
# State accessors (used extensively by articulation/rigid object data)
[docs]
@classmethod
def create_visual_material_writer(cls, batches: tuple[VisualMaterialBatch, ...]) -> VisualMaterialWriter:
"""Compile material-to-shape addresses for the active Newton model."""
return cls.backend.create_visual_material_writer(batches)
[docs]
@classmethod
def create_visual_shape_color_writer(
cls, asset: BaseArticulation, body_names: tuple[str, ...]
) -> VisualShapeColorWriter:
"""Compile selected articulation-body shape addresses for the active Newton model."""
model = cls.get_model()
view = asset.root_view
if not isinstance(view, ArticulationView):
root_expr = asset.cfg.prim_path
root_expr += (
"(?:/.*)?" if asset.cfg.articulation_root_prim_path is None else asset.cfg.articulation_root_prim_path
)
prim_paths = [path for path in model.articulation_label if re.fullmatch(root_expr, path)]
view = ArticulationView(model, prim_paths, verbose=False)
return VisualShapeColorWriter(model, view, body_names)
[docs]
@classmethod
def get_model(cls) -> Model:
"""Return the active physics model. Render consumers acquire their backend from the registry."""
return None if cls.backend is None else cls.backend.model
[docs]
@classmethod
def get_state_0(cls) -> State:
"""Get the current state."""
return None if cls.backend is None else cls.backend.state_0
@classmethod
def get_num_envs(cls) -> int:
return cls._num_envs
[docs]
@classmethod
def get_scene_data_provider(cls) -> SceneDataProvider:
"""Return the active scene data provider."""
return SimulationContext.instance().get_scene_data_provider()
[docs]
@classmethod
def get_state_1(cls) -> State:
"""Get the next state."""
return None if cls.backend is None else cls.backend.state_1
[docs]
@classmethod
def get_control(cls) -> Control:
"""Get the control object."""
return None if cls.backend is None else cls.backend.control
[docs]
@classmethod
def get_dt(cls) -> float:
"""Get the physics timestep. Alias for get_physics_dt()."""
return cls.get_physics_dt()
[docs]
@classmethod
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
[docs]
@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.backend is None or not cls.backend.model.actuators:
return
from isaaclab.actuators.newton import NewtonActuatorAdapter # noqa: PLC0415
dofs_per_env = cls.backend.model.joint_dof_count // cls._num_envs
NewtonManager._adapter = NewtonActuatorAdapter(
actuators=list(cls.backend.model.actuators),
num_envs=cls._num_envs,
num_joints=dofs_per_env,
dof_offset=0,
device=PhysicsManager._device,
)
cls._adapter.finalize(cls.backend.control)
[docs]
@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)
[docs]
@classmethod
def register_state_force_callback(cls, callback: Callable[[State], None]) -> None:
"""Register a graph-safe callback that applies forces before every solver substep.
Callbacks must be registered before solver initialization so they are
included in CUDA graph capture.
Args:
callback: Function that adds forces [N, N·m] to the provided state.
"""
if callback in NewtonManager._state_force_callbacks:
return
NewtonManager._state_force_callbacks.append(callback)
[docs]
@classmethod
def register_post_step_callback(cls, callback: Callable[[], None]) -> None:
"""Append a hook to the list invoked after the last solver substep on every step.
Each callback runs inside the stepped (and, when
:meth:`_is_all_graphable` is ``True``, captured) region right after the
final solver substep of the decimation loop and before
:meth:`_update_sensors`, so the launches it issues are recorded into
every captured CUDA graph and replayed on each tick. The hook fires
exactly once per :meth:`step` call, reflecting the state after all
decimation iterations (and their solver substeps) have completed -- not
once per substep and not once per decimation iteration. Callbacks must be
graph-safe (fixed shapes, no host branching on device data) and must be
registered before capture. Articulations with non-identity ordering
register their backend-to-user state republish here; all registered
callbacks fire in registration order each step.
"""
cls._post_step_callbacks.append(callback)
[docs]
@classmethod
def unregister_post_step_callback(cls, callback: Callable[[], None]) -> None:
"""Remove a previously registered post-step callback.
Symmetric to :meth:`register_post_step_callback`, this lets an
articulation deregister its republish hook when its callbacks are
cleared so the bound method does not linger on the class-level list
after the articulation is gone. Removing a callback that was never
registered (or was already removed) is a safe no-op, matching the
tolerant deregistration of other handles.
"""
with contextlib.suppress(ValueError):
cls._post_step_callbacks.remove(callback)
[docs]
@classmethod
def create_fixed_tendon_control(cls, articulation):
"""Build the solver's fixed-tendon command adapter for ``articulation``.
Tendon *state* is backend-neutral and lives on the articulation; how a target reaches the
solver is not. MuJoCo drives tendons through actuator controls, so only the MJWarp manager
implements this. The articulation stores what it gets, the way it stores its actuator
control, and never needs to know which solver is active.
Only the MuJoCo solver registers the ``mujoco:tendon`` frequency, so an articulation reports
tendons under MJWarp alone and this base is unreachable through the normal path. Reaching it
means a solver gained tendons with no way to command them, which is worth saying rather than
returning nothing -- ``None`` already means "this asset's tendons are all passive".
Args:
articulation: Newton articulation to drive.
Raises:
NotImplementedError: Always -- this solver has no fixed-tendon transmission.
"""
raise NotImplementedError(
f"{cls.__name__} does not drive fixed tendons. Fixed-tendon targets require a solver"
" that transmits to tendons, such as MJWarp."
)
[docs]
@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.
Invalidate the existing graph when the loop changes. Its replacement is captured
immediately before the next requested step, after authored state is reconciled.
"""
cls._decimation = max(1, decimation)
if cls._is_all_graphable():
cls._invalidate_graph()
[docs]
@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()
[docs]
@classmethod
def add_imu_sensor(cls, sites: list[int]) -> int:
"""Add an IMU sensor for measuring acceleration and angular velocity at sites.
Creates a ``newton.sensors.SensorIMU`` from pre-resolved site indices,
appends it to the internal list, and returns its index.
Args:
sites: Ordered list of site indices (one per environment).
Returns:
Index of the newly created sensor in the internal IMU sensor list.
"""
if cls.backend is None:
raise RuntimeError("add_imu_sensor called before model finalization (start_simulation).")
sensor = NewtonSensorIMU(
cls.backend.model,
sites=sites,
request_state_attributes=False, # Already requested via NewtonManager
)
idx = len(cls._newton_imu_sensors)
cls._newton_imu_sensors.append(sensor)
logger.info(f"Added IMU sensor (index={idx}, sites={len(sites)})")
return idx
