From 517d2e26dbbcb805d81bcd34afc5642ca2b24a7e Mon Sep 17 00:00:00 2001 From: Agon Serifi Date: Mon, 8 Jun 2026 15:00:37 +0200 Subject: [PATCH 1/2] Add maximal-coordinate solver hooks to core env and actuators --- .../changelog.d/aserifi-kamino-core-hooks.minor.rst | 13 +++++++++++++ source/isaaclab/isaaclab/actuators/actuator_base.py | 12 +++++++++++- .../isaaclab/isaaclab/envs/manager_based_rl_env.py | 7 +++++++ 3 files changed, 31 insertions(+), 1 deletion(-) create mode 100644 source/isaaclab/changelog.d/aserifi-kamino-core-hooks.minor.rst diff --git a/source/isaaclab/changelog.d/aserifi-kamino-core-hooks.minor.rst b/source/isaaclab/changelog.d/aserifi-kamino-core-hooks.minor.rst new file mode 100644 index 000000000000..124ed4514a31 --- /dev/null +++ b/source/isaaclab/changelog.d/aserifi-kamino-core-hooks.minor.rst @@ -0,0 +1,13 @@ +Added +^^^^^ + +* Added :attr:`~isaaclab.actuators.ActuatorBase.route_torque_to` to select whether + actuator torque is written to Newton ``Control.joint_f`` or ``Control.joint_act``, + enabling backward-Euler joint treatment under the Kamino solver. + +Fixed +^^^^^ + +* Fixed stale reset poses for maximal-coordinate solvers (e.g. Kamino) by refreshing + kinematics and derived buffers after in-step resets in + :meth:`~isaaclab.envs.ManagerBasedRLEnv.step` (no-op for other backends). diff --git a/source/isaaclab/isaaclab/actuators/actuator_base.py b/source/isaaclab/isaaclab/actuators/actuator_base.py index 4489983366d3..b7e1c6759c6e 100644 --- a/source/isaaclab/isaaclab/actuators/actuator_base.py +++ b/source/isaaclab/isaaclab/actuators/actuator_base.py @@ -7,7 +7,7 @@ from abc import ABC, abstractmethod from collections.abc import Sequence -from typing import TYPE_CHECKING, ClassVar +from typing import TYPE_CHECKING, ClassVar, Literal import torch @@ -43,6 +43,16 @@ class ActuatorBase(ABC): If a class inherits from :class:`ImplicitActuator`, then this flag should be set to :obj:`True`. """ + route_torque_to: ClassVar[Literal["joint_f", "joint_act"]] = "joint_f" + """Where the per-step actuator-computed torque is written into Newton's ``Control``. + + - ``"joint_f"`` (default): torque is written to ``Control.joint_f`` and enters the + body wrench, integrated semi-explicitly. This matches PhysX/MuJoCo semantics. + - ``"joint_act"``: torque is written to ``Control.joint_act`` instead. For the Kamino + solver this enters the joint dynamic-constraint row of the PADMM solve, getting + backward-Euler treatment of joint armature and damping. + """ + computed_effort: torch.Tensor """The computed effort for the actuator group. Shape is (num_envs, num_joints).""" diff --git a/source/isaaclab/isaaclab/envs/manager_based_rl_env.py b/source/isaaclab/isaaclab/envs/manager_based_rl_env.py index dfdbef5091f9..8fc589134bb1 100644 --- a/source/isaaclab/isaaclab/envs/manager_based_rl_env.py +++ b/source/isaaclab/isaaclab/envs/manager_based_rl_env.py @@ -252,6 +252,13 @@ def step(self, action: torch.Tensor) -> VecEnvStepReturn: self._reset_idx(reset_env_ids) + # Refresh kinematics and derived buffers for the just-reset envs so the + # observations / terminations below reflect the reset pose. Maximal-coordinate + # solvers (e.g. Kamino) apply reset FK one step late; for others this is a no-op. + self.scene.write_data_to_sim() + self.sim.forward() + self.scene.update(dt=self.physics_dt) + # if sensors are added to the scene, make sure we render to reflect changes in reset if self.render_enabled and is_rendering and self.has_rtx_sensors and self.cfg.num_rerenders_on_reset > 0: for _ in range(self.cfg.num_rerenders_on_reset): From 9e230b03dd98feab38ac983d4af030c01d49bbd2 Mon Sep 17 00:00:00 2001 From: Agon Serifi Date: Mon, 8 Jun 2026 15:03:53 +0200 Subject: [PATCH 2/2] Add Kamino closed-loop articulation support to Newton backend --- ...serifi-newton-kamino-closed-loop.minor.rst | 14 + .../assets/articulation/articulation.py | 85 ++++- .../assets/articulation/articulation_data.py | 16 + .../assets/articulation/closed_loop_view.py | 349 ++++++++++++++++++ .../cloner/newton_replicate.py | 25 +- .../isaaclab_newton/physics/kamino_manager.py | 171 ++++++++- .../physics/kamino_manager_cfg.py | 10 + .../isaaclab_newton/physics/newton_manager.py | 77 +++- .../sensors/contact_sensor/contact_sensor.py | 5 +- 9 files changed, 716 insertions(+), 36 deletions(-) create mode 100644 source/isaaclab_newton/changelog.d/aserifi-newton-kamino-closed-loop.minor.rst create mode 100644 source/isaaclab_newton/isaaclab_newton/assets/articulation/closed_loop_view.py diff --git a/source/isaaclab_newton/changelog.d/aserifi-newton-kamino-closed-loop.minor.rst b/source/isaaclab_newton/changelog.d/aserifi-newton-kamino-closed-loop.minor.rst new file mode 100644 index 000000000000..b5e76be55a76 --- /dev/null +++ b/source/isaaclab_newton/changelog.d/aserifi-newton-kamino-closed-loop.minor.rst @@ -0,0 +1,14 @@ +Added +^^^^^ + +* Added :class:`~isaaclab_newton.assets.articulation.closed_loop_view.ClosedLoopView` + to simulate closed kinematic chains (parallel linkages) in maximal coordinates with + the Kamino solver. +* Added ``max_contacts_per_world`` to :class:`~isaaclab_newton.physics.KaminoSolverCfg` + to bound per-world contact allocation for the Kamino solver. + +Fixed +^^^^^ + +* Fixed contact-sensor forces for the Kamino solver by routing contact aggregation + through a unified, solver-agnostic path shared with the MuJoCo-Warp backend. diff --git a/source/isaaclab_newton/isaaclab_newton/assets/articulation/articulation.py b/source/isaaclab_newton/isaaclab_newton/assets/articulation/articulation.py index a5ebaffd09fe..533ee26fb11a 100644 --- a/source/isaaclab_newton/isaaclab_newton/assets/articulation/articulation.py +++ b/source/isaaclab_newton/isaaclab_newton/assets/articulation/articulation.py @@ -311,6 +311,7 @@ def write_data_to_sim(self): # Standard Lab actuator path self._apply_actuator_model() self.data._sim_bind_joint_effort.assign(self._joint_effort_target_sim) + self.data._sim_bind_joint_act.assign(self._joint_act_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) @@ -510,7 +511,7 @@ def write_root_link_pose_to_sim_index( if self.data._root_state_w is not None: self.data._root_state_w.timestamp = -1.0 self.data._fk_timestamp = -1.0 # Forces a kinematic update to get the latest body link poses. - SimulationManager.invalidate_fk(env_ids=env_ids, articulation_ids=self._root_view.articulation_ids) + SimulationManager.invalidate_fk(env_ids=env_ids, articulation_ids=self._get_root_view_articulation_ids()) if self.data._body_com_pose_w is not None: self.data._body_com_pose_w.timestamp = -1.0 if self.data._body_state_w is not None: @@ -567,7 +568,7 @@ def write_root_link_pose_to_sim_mask( if self.data._root_state_w is not None: self.data._root_state_w.timestamp = -1.0 self.data._fk_timestamp = -1.0 # Forces a kinematic update to get the latest body link poses. - SimulationManager.invalidate_fk(env_mask=env_mask, articulation_ids=self._root_view.articulation_ids) + SimulationManager.invalidate_fk(env_mask=env_mask, articulation_ids=self._get_root_view_articulation_ids()) if self.data._body_com_pose_w is not None: self.data._body_com_pose_w.timestamp = -1.0 if self.data._body_state_w is not None: @@ -632,7 +633,7 @@ def write_root_com_pose_to_sim_index( if self.data._root_state_w is not None: self.data._root_state_w.timestamp = -1.0 self.data._fk_timestamp = -1.0 # Forces a kinematic update to get the latest body link poses. - SimulationManager.invalidate_fk(env_ids=env_ids, articulation_ids=self._root_view.articulation_ids) + SimulationManager.invalidate_fk(env_ids=env_ids, articulation_ids=self._get_root_view_articulation_ids()) if self.data._body_com_pose_w is not None: self.data._body_com_pose_w.timestamp = -1.0 if self.data._body_state_w is not None: @@ -693,7 +694,7 @@ def write_root_com_pose_to_sim_mask( if self.data._root_state_w is not None: self.data._root_state_w.timestamp = -1.0 self.data._fk_timestamp = -1.0 # Forces a kinematic update to get the latest body link poses. - SimulationManager.invalidate_fk(env_mask=env_mask, articulation_ids=self._root_view.articulation_ids) + SimulationManager.invalidate_fk(env_mask=env_mask, articulation_ids=self._get_root_view_articulation_ids()) if self.data._body_com_pose_w is not None: self.data._body_com_pose_w.timestamp = -1.0 if self.data._body_state_w is not None: @@ -1124,7 +1125,7 @@ def write_joint_position_to_sim_index( ) # Invalidate FK timestamp so body poses are recomputed on next access. self.data._fk_timestamp = -1.0 - SimulationManager.invalidate_fk(env_ids=env_ids, articulation_ids=self._root_view.articulation_ids) + SimulationManager.invalidate_fk(env_ids=env_ids, articulation_ids=self._get_root_view_articulation_ids()) # Need to invalidate the buffer to trigger the update with the new root pose. # Only invalidate if the buffer has been accessed (not None). if self.data._body_link_vel_w is not None: @@ -1182,7 +1183,7 @@ def write_joint_position_to_sim_mask( ) # Invalidate FK timestamp so body poses are recomputed on next access. self.data._fk_timestamp = -1.0 - SimulationManager.invalidate_fk(env_mask=env_mask, articulation_ids=self._root_view.articulation_ids) + SimulationManager.invalidate_fk(env_mask=env_mask, articulation_ids=self._get_root_view_articulation_ids()) # Need to invalidate the buffer to trigger the update with the new root pose. # Only invalidate if the buffer has been accessed (not None). if self.data._body_link_vel_w is not None: @@ -3501,8 +3502,59 @@ def has_articulation_root_api(prim) -> bool: asset_prim, root_expr = resolve_matching_prims_from_source(self.cfg.prim_path)[0] walk_root = asset_prim.GetPath().pathString - root_prims = get_all_matching_child_prims(walk_root, has_articulation_root_api, expected_num_matches=1) - root_prim_path_expr = root_expr + root_prims[0].GetPath().pathString[len(walk_root) :] + + # Find all articulation root prims in the first environment. + first_env_root_prims = get_all_matching_child_prims( + walk_root, + predicate=has_articulation_root_api, + traverse_instance_prims=False, + ) + from isaaclab_newton.cloner.newton_replicate import _prim_has_closed_kinematic_loops + + has_closed_loops = _prim_has_closed_kinematic_loops(asset_prim) + if len(first_env_root_prims) == 0 or has_closed_loops: + # Closed-loop robots (e.g. DR Legs) have cyclic joint graphs that cannot + # form a tree-structured Newton ``ArticulationView``. Use ``ClosedLoopView`` + # which provides strided views into the global Newton model arrays instead. + from .closed_loop_view import ClosedLoopView + + if has_closed_loops and len(first_env_root_prims) > 0: + logger.info( + "Closed kinematic loops detected under '%s'; using ClosedLoopView despite" + " ArticulationRootAPI presence.", + walk_root, + ) + else: + logger.info( + "No ArticulationRootAPI found under '%s'. Using ClosedLoopView for closed-loop robot support.", + walk_root, + ) + self._root_view = ClosedLoopView( + SimulationManager.get_model(), + self.cfg.prim_path.replace(".*", "*"), + ) + SimulationManager.get_physics_sim_view().append(self._root_view) + self._data = ArticulationData(self.root_view, self.device) + self._physics_ready_handle = SimulationManager.register_callback( + lambda _: self._data._create_simulation_bindings(), + PhysicsEvent.PHYSICS_READY, + name=f"articulation_rebind_{self.cfg.prim_path}", + ) + self._create_buffers() + self._process_cfg() + self._process_actuators_cfg() + self._process_tendons() + # Let the articulation data know that it is fully instantiated and ready to use. + self.data.is_primed = True + return + if len(first_env_root_prims) > 1: + raise RuntimeError( + f"Failed to find a single articulation when resolving '{walk_root}'." + f" Found multiple '{first_env_root_prims}' under '{walk_root}'." + " Please ensure that there is only one articulation in the prim path tree." + ) + + root_prim_path_expr = root_expr + first_env_root_prims[0].GetPath().pathString[len(walk_root) :] # -- articulation self._root_view = ArticulationView( SimulationManager.get_model(), @@ -3538,6 +3590,17 @@ def has_articulation_root_api(prim) -> bool: # Let the articulation data know that it is fully instantiated and ready to use. self.data.is_primed = True + def _get_root_view_articulation_ids(self) -> wp.array | None: + """Return the root view's ``articulation_ids``, or ``None`` for closed-loop assets. + + :class:`ClosedLoopView` exposes an empty ``(world_count, 0)`` array, so returning + ``None`` makes the reset machinery scope by ``env_ids`` / ``env_mask`` instead. + """ + art_ids = getattr(self._root_view, "articulation_ids", None) + if art_ids is None or art_ids.ndim < 2 or art_ids.shape[1] == 0: + return None + return art_ids + def _clear_callbacks(self) -> None: """Clears all registered callbacks, including the physics-ready rebind handle.""" super()._clear_callbacks() @@ -3572,6 +3635,7 @@ def _create_buffers(self): self._joint_pos_target_sim = wp.zeros_like(self.data.joint_pos_target.warp, device=self.device) self._joint_vel_target_sim = wp.zeros_like(self.data.joint_pos_target.warp, device=self.device) self._joint_effort_target_sim = wp.zeros_like(self.data.joint_pos_target.warp, device=self.device) + self._joint_act_target_sim = wp.zeros_like(self.data.joint_pos_target.warp, device=self.device) # soft joint position limits (recommended not to be too close to limits). wp.launch( @@ -3975,6 +4039,9 @@ def _apply_actuator_model(self): gear_ratio = actuator.gear_ratio else: gear_ratio = None + target_torque_buf = ( + self._joint_act_target_sim if actuator.route_torque_to == "joint_act" else self._joint_effort_target_sim + ) wp.launch( articulation_kernels.update_targets, dim=(self.num_instances, joint_indices.shape[0]), @@ -3987,7 +4054,7 @@ def _apply_actuator_model(self): outputs=[ self._joint_pos_target_sim, self._joint_vel_target_sim, - self._joint_effort_target_sim, + target_torque_buf, ], device=self.device, ) 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 6eb7020f370b..94bba411441a 100644 --- a/source/isaaclab_newton/isaaclab_newton/assets/articulation/articulation_data.py +++ b/source/isaaclab_newton/isaaclab_newton/assets/articulation/articulation_data.py @@ -315,6 +315,20 @@ def applied_torque(self) -> ProxyArray: """ return self._applied_torque_ta + @property + def joint_act_target(self) -> ProxyArray: + """Per-DOF feedforward actuation written to Newton's ``Control.joint_act``. + + Shape is (num_instances, num_joints), dtype = wp.float32. In torch this resolves to + (num_instances, num_joints). + + Populated each ``write_data_to_sim`` from the per-actuator computed torques whose + :attr:`~isaaclab.actuators.ActuatorBase.route_torque_to` is set to ``"joint_act"``. + For the Kamino solver this becomes ``pd_tau_j_ff`` and feeds the implicit joint + dynamic-constraint row. + """ + return self._joint_act_target_ta + """ Joint properties """ @@ -1806,6 +1820,7 @@ def _pin_proxy_arrays(self) -> None: self._joint_pos_limits_upper_ta = ProxyArray(self._sim_bind_joint_pos_limits_upper) self._joint_vel_limits_ta = ProxyArray(self._sim_bind_joint_vel_limits_sim) self._joint_effort_limits_ta = ProxyArray(self._sim_bind_joint_effort_limits_sim) + self._joint_act_target_ta = ProxyArray(self._sim_bind_joint_act) self._body_mass_ta = ProxyArray(self._sim_bind_body_mass) self._body_inertia_ta = ProxyArray(self._sim_bind_body_inertia) self._body_com_pos_b_ta = ProxyArray(self._sim_bind_body_com_pos_b) @@ -1837,6 +1852,7 @@ def _pin_proxy_arrays(self) -> None: self._joint_pos_limits_upper_ta = ProxyArray(self._sim_bind_joint_pos_limits_upper) self._joint_vel_limits_ta = ProxyArray(self._sim_bind_joint_vel_limits_sim) self._joint_effort_limits_ta = ProxyArray(self._sim_bind_joint_effort_limits_sim) + self._joint_act_target_ta = ProxyArray(self._sim_bind_joint_act) self._soft_joint_pos_limits_ta = ProxyArray(self._soft_joint_pos_limits) self._soft_joint_vel_limits_ta = ProxyArray(self._soft_joint_vel_limits) self._gear_ratio_ta = ProxyArray(self._gear_ratio) diff --git a/source/isaaclab_newton/isaaclab_newton/assets/articulation/closed_loop_view.py b/source/isaaclab_newton/isaaclab_newton/assets/articulation/closed_loop_view.py new file mode 100644 index 000000000000..17529af66bde --- /dev/null +++ b/source/isaaclab_newton/isaaclab_newton/assets/articulation/closed_loop_view.py @@ -0,0 +1,349 @@ +# 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 + +"""View adapter for closed-loop (non-articulated) robots in Newton. + +Closed-loop robots like DR Legs have cyclic joint graphs that cannot form a +tree-structured articulation. Newton's ``ArticulationView`` requires a tree, so this +module provides ``ClosedLoopView`` -- a duck-typed replacement that accesses the global +Newton ``Model`` / ``State`` / ``Control`` arrays directly. +""" + +from __future__ import annotations + +import logging +from typing import TYPE_CHECKING + +import warp as wp + +if TYPE_CHECKING: + from newton import Model, State + +logger = logging.getLogger(__name__) + + +class ClosedLoopView: + """A view adapter for robots without a formal Newton articulation. + + Implements the read subset of ``newton.selection.ArticulationView``'s interface + that ``ArticulationData`` depends on, by providing strided views into the global + Newton arrays. One "instance" corresponds to one world. + + Args: + model: Finalized Newton ``Model``. + pattern: Body-label pattern used to identify the robot (currently + unused -- all bodies/joints in each world are assumed to belong + to the robot). + """ + + def __init__(self, model: Model, pattern: str = "*") -> None: + self.model = model + self.device = model.device + self._pattern = pattern + + self._world_count = model.world_count if model.world_count > 0 else 1 + self._count_per_world = 1 + + # Per-world counts from the Newton model totals (Kamino's ``size`` struct + # doesn't expose per-world joint_dof/coord counts). + nw = max(self._world_count, 1) + self._bodies_per_world = model.body_count // nw + self._joints_per_world = model.joint_count // nw + self._joint_dofs_per_world = getattr(model, "joint_dof_count", model.joint_count) // nw + self._joint_coords_per_world = getattr(model, "joint_coord_count", model.joint_count) // nw + self._shapes_per_world = model.shape_count // nw + # Free-base coord/dof offsets; controllable = per_world - root. + self._root_coord_count = 0 + self._root_dof_count = 0 + + # Extract names from first world only + self._joint_names = self._extract_names(model.joint_label, self._joints_per_world) + self._link_names = self._extract_names(model.body_label, self._bodies_per_world) + + # Build per-body shape counts (first world) + shape_body = model.shape_body.numpy() + self._link_shapes: list[list[int]] = [[] for _ in range(self._bodies_per_world)] + for s in range(self._shapes_per_world): + bid = int(shape_body[s]) + if 0 <= bid < self._bodies_per_world: + self._link_shapes[bid].append(s) + + # Root base type: FREE/DISTANCE -> joint_q[:root_coord_count]; + # FIXED -> joint_X_p[0]; implicit free base (body 0 has no joint to world, + # common for closed-loop USDs) -> body_q[0]. + from newton import JointType + + self._is_implicit_free_base = False + if model.joint_count > 0: + jtype = int(model.joint_type.numpy()[0]) + self._is_fixed_base = jtype == int(JointType.FIXED) + self._is_floating_base = jtype in (int(JointType.FREE), int(JointType.DISTANCE)) + if not self._is_floating_base and not self._is_fixed_base: + joint_child = model.joint_child.numpy()[: self._joints_per_world] + children = set(int(c) for c in joint_child if int(c) >= 0) + if 0 not in children: + self._is_implicit_free_base = True + if self._is_floating_base: + self._root_coord_count = 7 if jtype == int(JointType.FREE) else 1 + self._root_dof_count = 6 if jtype == int(JointType.FREE) else 1 + else: + self._is_fixed_base = True + self._is_floating_base = False + + logger.info( + "ClosedLoopView: %d worlds, %d bodies, %d joints (%d DOFs, %d coords) per world " + "(fixed_base=%s floating_base=%s implicit_free_base=%s, root_coords=%d root_dofs=%d)", + self._world_count, + self._bodies_per_world, + self._joints_per_world, + self._joint_dofs_per_world, + self._joint_coords_per_world, + self._is_fixed_base, + self._is_floating_base, + self._is_implicit_free_base, + self._root_coord_count, + self._root_dof_count, + ) + + @staticmethod + def _extract_names(labels, count: int) -> list[str]: + """Extract short names from the first ``count`` labels.""" + return [lbl.rsplit("/", 1)[-1] for lbl in labels[:count]] + + # ------------------------------------------------------------------ + # Properties matching ArticulationView interface + # ------------------------------------------------------------------ + + @property + def count(self) -> int: + return self._world_count + + @property + def world_count(self) -> int: + return self._world_count + + @property + def count_per_world(self) -> int: + return self._count_per_world + + @property + def joint_dof_count(self) -> int: + return self._joint_dofs_per_world - self._root_dof_count + + @property + def joint_coord_count(self) -> int: + return self._joint_coords_per_world - self._root_coord_count + + @property + def link_count(self) -> int: + return self._bodies_per_world + + @property + def tendon_count(self) -> int: + """Number of fixed tendons (always 0 for closed-loop assets). + + Closed-loop robots are accessed via the global Newton model arrays and do not + expose MuJoCo-style fixed tendons, so ``ArticulationData`` allocates empty tendon + bindings for them. + """ + return 0 + + @property + def articulation_ids(self) -> wp.array: + """Empty ``(world_count, 0)`` int32 array (closed-loop assets have no tree articulations). + + Sizing buffers from ``articulation_ids`` (e.g. + ``ArticulationData._create_jacobian_buffers``) yields zero rows, while + :meth:`Articulation._get_root_view_articulation_ids` treats the empty + array as "no articulations" so reset scoping falls back to ``env_ids`` / + ``env_mask``. + """ + return wp.zeros((self._world_count, 0), dtype=wp.int32, device=self.device) + + @property + def is_fixed_base(self) -> bool: + return self._is_fixed_base + + @property + def is_floating_base(self) -> bool: + return self._is_floating_base + + @property + def joint_dof_names(self) -> list[str]: + if self._is_floating_base: + return self._joint_names[1:] + return self._joint_names + + @property + def link_names(self) -> list[str]: + return self._link_names + + @property + def body_names(self) -> list[str]: + return self._link_names + + @property + def body_shapes(self) -> list[list[int]]: + return self._link_shapes + + @property + def link_shapes(self) -> list[list[int]]: + return self._link_shapes + + # ------------------------------------------------------------------ + # Core array access + # ------------------------------------------------------------------ + + def _get_strided_view( + self, + attrib: wp.array, + items_per_world: int, + _slice: slice | int | None = None, + ) -> wp.array: + """Create a strided view into a global Newton array shaped + ``(world_count, count_per_world=1, val_count, ...)``. A non-zero slice start + is folded into the view's base pointer. Passing ``_slice`` as an ``int`` drops + the ``val_count`` axis (NumPy-style), matching ``articulation_data.py``'s + shape convention for root transforms. + """ + value_stride = attrib.strides[0] + trailing_shape = attrib.shape[1:] + trailing_strides = attrib.strides[1:] + + squeeze_val = False + if _slice is None: + val_count = items_per_world + val_offset = 0 + elif isinstance(_slice, int): + val_count = 1 + val_offset = _slice + squeeze_val = True + elif isinstance(_slice, slice): + start = _slice.start or 0 + stop = _slice.stop if _slice.stop is not None else items_per_world + val_count = stop - start + val_offset = start + else: + val_count = items_per_world + val_offset = 0 + + if squeeze_val: + shape = (self._world_count, self._count_per_world, *trailing_shape) + strides = ( + items_per_world * value_stride, # stride between worlds + items_per_world * value_stride, # stride within worlds (1 robot) + *trailing_strides, + ) + else: + shape = (self._world_count, self._count_per_world, val_count, *trailing_shape) + strides = ( + items_per_world * value_stride, # stride between worlds + items_per_world * value_stride, # stride within worlds (1 robot) + value_stride, # stride between items + *trailing_strides, + ) + + if attrib.ptr is None: + result = wp.empty(shape, dtype=attrib.dtype, device=attrib.device) + result.ptr = None + return result + + base_ptr = int(attrib.ptr) + val_offset * int(value_stride) + result = wp.array( + ptr=base_ptr, + dtype=attrib.dtype, + shape=shape, + strides=strides, + device=attrib.device, + copy=False, + ) + return result + + def _resolve_frequency(self, name: str) -> tuple[int, int]: + """Return (items-per-world, per-world-offset) for an attribute. Offset skips + free-base DOF/coord rows to match isaaclab ArticulationData's controllable-DOF view. + """ + freq = self.model.get_attribute_frequency(name) + if isinstance(freq, str): + return self._joints_per_world, 0 + freq_name = freq.name.lower() if hasattr(freq, "name") else str(freq).lower() + if "body" in freq_name or "link" in freq_name: + return self._bodies_per_world, 0 + if "joint_coord" in freq_name: + return self._joint_coords_per_world, self._root_coord_count + if "joint_dof" in freq_name: + return self._joint_dofs_per_world, self._root_dof_count + if "joint" in freq_name: + return self._joints_per_world, (1 if self._is_floating_base else 0) + if "shape" in freq_name: + return self._shapes_per_world, 0 + return self._joints_per_world, 0 + + # ------------------------------------------------------------------ + # Read API (matches the ArticulationView subset ArticulationData uses) + # ------------------------------------------------------------------ + + def get_attribute(self, name: str, source) -> wp.array: + """Get a strided view of a model/state/control attribute.""" + if "." in name: + parts = name.split(".") + attrib = source + for part in parts: + attrib = getattr(attrib, part) + freq_name = ":".join(parts) + else: + attrib = getattr(source, name) + freq_name = name + + items, offset = self._resolve_frequency(freq_name) + if offset > 0: + return self._get_strided_view(attrib, items, _slice=slice(offset, items)) + return self._get_strided_view(attrib, items) + + def get_root_transforms(self, source: Model | State) -> wp.array: + if self.is_floating_base: + attrib = self._get_strided_view( + source.joint_q, self._joint_coords_per_world, _slice=slice(0, self._root_coord_count) + ) + elif self._is_implicit_free_base and hasattr(source, "body_q"): + attrib = self._get_strided_view(source.body_q, self._bodies_per_world, _slice=0) + else: + attrib = self._get_strided_view(self.model.joint_X_p, self._joints_per_world, _slice=0) + if attrib.dtype is wp.transformf: + return attrib + return wp.array(attrib, dtype=wp.transformf, device=self.device, copy=False) + + def get_root_velocities(self, source: Model | State): + if self.is_floating_base: + attrib = self._get_strided_view( + source.joint_qd, self._joint_dofs_per_world, _slice=slice(0, self._root_dof_count) + ) + if attrib.dtype is wp.spatial_vectorf: + return attrib + return wp.array(attrib, dtype=wp.spatial_vectorf, device=self.device, copy=False) + if self._is_implicit_free_base and hasattr(source, "body_qd"): + attrib = self._get_strided_view(source.body_qd, self._bodies_per_world, _slice=0) + if attrib.dtype is wp.spatial_vectorf: + return attrib + return wp.array(attrib, dtype=wp.spatial_vectorf, device=self.device, copy=False) + return None + + def get_link_transforms(self, source: Model | State) -> wp.array: + return self._get_strided_view(source.body_q, self._bodies_per_world) + + def get_link_velocities(self, source: Model | State) -> wp.array: + return self._get_strided_view(source.body_qd, self._bodies_per_world) + + def _dof_slice(self, is_coord: bool) -> slice: + """Slice of joint_q/joint_qd for controllable DOFs (excludes free-base coords/DOFs).""" + offset = self._root_coord_count if is_coord else self._root_dof_count + per_world = self._joint_coords_per_world if is_coord else self._joint_dofs_per_world + return slice(offset, per_world) + + def get_dof_positions(self, source: Model | State) -> wp.array: + return self._get_strided_view(source.joint_q, self._joint_coords_per_world, _slice=self._dof_slice(True)) + + def get_dof_velocities(self, source: Model | State) -> wp.array: + return self._get_strided_view(source.joint_qd, self._joint_dofs_per_world, _slice=self._dof_slice(False)) diff --git a/source/isaaclab_newton/isaaclab_newton/cloner/newton_replicate.py b/source/isaaclab_newton/isaaclab_newton/cloner/newton_replicate.py index e868607c05c5..e032ffd17a2d 100644 --- a/source/isaaclab_newton/isaaclab_newton/cloner/newton_replicate.py +++ b/source/isaaclab_newton/isaaclab_newton/cloner/newton_replicate.py @@ -12,11 +12,33 @@ from newton import ModelBuilder, solvers from newton._src.usd.schemas import SchemaResolverNewton, SchemaResolverPhysx -from pxr import Usd +from pxr import Usd, UsdPhysics from isaaclab_newton.physics import NewtonManager +def _prim_has_closed_kinematic_loops(prim: Usd.Prim) -> bool: + """True if any body under *prim* is the child of more than one joint (cyclic graph).""" + child_counts: dict[str, int] = {} + for descendant in Usd.PrimRange(prim): + if not descendant.IsA(UsdPhysics.Joint): + continue + try: + joint = UsdPhysics.Joint(descendant) + targets = joint.GetBody1Rel().GetTargets() + except Exception: + continue + if not targets: + continue + child_path = targets[0].pathString + if not child_path: + continue + child_counts[child_path] = child_counts.get(child_path, 0) + 1 + if child_counts[child_path] > 1: + return True + return False + + def _build_newton_builder_from_mapping( stage: Usd.Stage, sources: Sequence[str], @@ -85,6 +107,7 @@ def _build_newton_builder_from_mapping( for src_path in sources: p = NewtonManager.create_builder(up_axis=up_axis) solvers.SolverMuJoCo.register_custom_attributes(p) + solvers.SolverKamino.register_custom_attributes(p) p.add_usd( stage, root_path=src_path, diff --git a/source/isaaclab_newton/isaaclab_newton/physics/kamino_manager.py b/source/isaaclab_newton/isaaclab_newton/physics/kamino_manager.py index ba2d407d2a9a..bc05671a7c3f 100644 --- a/source/isaaclab_newton/isaaclab_newton/physics/kamino_manager.py +++ b/source/isaaclab_newton/isaaclab_newton/physics/kamino_manager.py @@ -10,7 +10,7 @@ import logging import warp as wp -from newton import Model, eval_fk +from newton import JointType, Model, eval_fk from newton.solvers import SolverKamino from isaaclab.physics import PhysicsManager @@ -32,21 +32,147 @@ class NewtonKaminoManager(NewtonManager): @classmethod def _forward_kamino(cls, world_mask: wp.array | None = None) -> None: - """Kamino-specific forward kinematics via ``solver.reset()``. + """Reset Kamino solver state to match the current joint configuration. - Kamino's ``joint_q`` / ``joint_u`` include coordinates for **all** joints - (including free joints), so we pass Newton's full state arrays directly. + Two reset paths are selected at solver-construction time via + :attr:`KaminoSolverCfg.use_fk_solver`: + + * ``use_fk_solver=True`` (default): ``solver.reset(joint_q=..., joint_u=..., + base_q=..., base_u=...)`` runs Kamino's Gauss-Newton FK so ``body_q`` is + consistent with the written ``joint_q`` (needed for non-trivial joint reset + targets on closed-loop assets). + * ``use_fk_solver=False``: ``solver.reset(base_q=..., base_u=...)`` routes to + ``_reset_to_base_state`` -- bodies are placed by transforming the model's + reference body poses by ``base_q`` with **no FK solver**. This is the fast + path for tasks that pin every joint coord to ``0`` (the assembled + closed-loop-valid configuration) and sidesteps an FK-solver buffer overflow + at high ``num_envs``. It is only correct when the written ``joint_q`` equals + the model default (all zeros). Args: world_mask: Per-world mask indicating which worlds to reset. Shape ``(num_worlds,)``, dtype ``wp.int32``. If None, resets all worlds. """ - cls._solver.reset( - state_out=cls._state_0, - joint_q=cls._state_0.joint_q, - joint_u=cls._state_0.joint_qd, - world_mask=world_mask, - ) + _model = cls._model + _nw = max(int(getattr(_model, "world_count", 0) or 0), 1) + _first_joint_is_free = _model.joint_count > 0 and int(_model.joint_type.numpy()[0]) == int(JointType.FREE) + + _base_q = None + _base_u = None + if _first_joint_is_free: + _coord_count = int(getattr(_model, "joint_coord_count", _model.joint_count)) + _dof_count = int(getattr(_model, "joint_dof_count", _model.joint_count)) + if _nw > 0 and _coord_count % _nw == 0 and _dof_count % _nw == 0: + _coords_per_world = _coord_count // _nw + _dofs_per_world = _dof_count // _nw + if _coords_per_world >= 7 and _dofs_per_world >= 6: + _jq = wp.to_torch(cls._state_0.joint_q).reshape(_nw, _coords_per_world) + _ju = wp.to_torch(cls._state_0.joint_qd).reshape(_nw, _dofs_per_world) + _base_q = wp.from_torch(_jq[:, :7].contiguous(), dtype=wp.transformf) + _base_u = wp.from_torch(_ju[:, :6].contiguous(), dtype=wp.spatial_vectorf) + else: + _body_count = int(getattr(_model, "body_count", 0) or 0) + if _body_count > 0 and _nw > 0 and _body_count % _nw == 0: + _bodies_per_world = _body_count // _nw + if _model.joint_count > 0: + _joints_per_world = _model.joint_count // _nw if _nw > 0 else _model.joint_count + _joint_child_np = _model.joint_child.numpy()[: max(_joints_per_world, 0)] + _per_world_children = {int(c) for c in _joint_child_np if int(c) >= 0} + else: + _per_world_children = set() + if 0 not in _per_world_children: + _bq = wp.to_torch(cls._state_0.body_q).reshape(_nw, _bodies_per_world, 7) + _bu = wp.to_torch(cls._state_0.body_qd).reshape(_nw, _bodies_per_world, 6) + _base_q = wp.from_torch(_bq[:, 0, :].contiguous(), dtype=wp.transformf) + _base_u = wp.from_torch(_bu[:, 0, :].contiguous(), dtype=wp.spatial_vectorf) + + _use_fk = bool(getattr(getattr(cls._solver, "_config", None), "use_fk_solver", True)) + if _use_fk: + cls._solver.reset( + state=cls._state_0, + world_mask=world_mask, + joint_q=cls._state_0.joint_q, + joint_u=cls._state_0.joint_qd, + base_q=_base_q, + base_u=_base_u, + ) + elif _base_q is not None: + # No-FK reset honoring the per-world target base pose: ``base_q`` routes Kamino to + # ``_reset_to_base_state``, which rigidly transforms the assembled reference + # configuration (valid at ``joint_q == 0``) so each world keeps its origin offset + # and the closed-loop constraints stay satisfied. ``base_u`` (zero on reset) zeroes + # all body velocities. + cls._solver.reset( + state=cls._state_0, + world_mask=world_mask, + base_q=_base_q, + base_u=_base_u, + ) + else: + # Fallback when the base body could not be identified: restore the assembled + # model-default state (consistent, but ignores the per-world target pose). + cls._solver.reset( + state=cls._state_0, + world_mask=world_mask, + ) + + # Snap ``joint_q_prev`` to the reset joint coords for the reset worlds: the no-FK + # reset leaves it stale, so the Moreau integrator would otherwise emit a spurious + # ``(q_new - q_prev)/dt`` velocity on the first post-reset step. The mask keeps + # other worlds' history intact. + _jqp = getattr(cls._state_0, "joint_q_prev", None) + if _jqp is not None and world_mask is not None: + _jq_t = wp.to_torch(cls._state_0.joint_q) + _jqp_t = wp.to_torch(_jqp) + _mask = wp.to_torch(world_mask).bool() + _nw_m = _mask.shape[0] + if _nw_m > 0 and _jq_t.shape[0] % _nw_m == 0: + _c = _jq_t.shape[0] // _nw_m + _jqp_t.view(_nw_m, _c)[_mask] = _jq_t.view(_nw_m, _c)[_mask] + + # Overwrite body_q via Newton's eval_fk for a consistent frame convention. For + # closed-loop assets the fk mask is zero-length, so this is a no-op and body poses + # come from the base-state reset above. + eval_fk(cls._model, cls._state_0.joint_q, cls._state_0.joint_qd, cls._state_0, cls._fk_reset_mask) + + # Sync state_1 to match state_0 so both states are consistent for the dual-state + # stepping scheme. + if cls._state_1 is not None and cls._state_1 is not cls._state_0: + for attr in ("joint_q", "joint_qd", "body_q", "body_qd"): + _src = getattr(cls._state_0, attr, None) + _dst = getattr(cls._state_1, attr, None) + if _src is not None and _dst is not None: + wp.copy(_dst, _src) + for attr in ("joint_q_prev", "joint_lambdas"): + _src = getattr(cls._state_0, attr, None) + _dst = getattr(cls._state_1, attr, None) + if _src is not None and ( + _dst is None or _dst.shape != _src.shape or _dst.dtype != _src.dtype or _dst.device != _src.device + ): + setattr(cls._state_1, attr, wp.clone(_src)) + elif _src is not None and _dst is not None: + wp.copy(_dst, _src) + + @classmethod + def forward(cls) -> None: + """Update kinematics without stepping physics. + + For Kamino, consume any pending reset flagged by :meth:`invalidate_fk` so the + explicit-reset path (``env.reset()`` -> ``sim.forward()``) makes ``body_q`` + consistent with the reset ``joint_q`` before observations are read. Falls back + to a full ``eval_fk`` when no reset is pending. + """ + if cls._kamino_needs_reset: + cls._forward_kamino(world_mask=cls._world_reset_mask) + # Clear the flag on the BASE class: assigning through ``cls`` would shadow the + # base attribute and make every later ``invalidate_fk`` go unobserved. + NewtonManager._kamino_needs_reset = False + if cls._world_reset_mask is not None: + cls._world_reset_mask.zero_() + if cls._fk_reset_mask is not None: + cls._fk_reset_mask.zero_() + return + eval_fk(cls._model, cls._state_0.joint_q, cls._state_0.joint_qd, cls._state_0, None) @classmethod def step(cls) -> None: @@ -55,12 +181,12 @@ def step(cls) -> None: if sim is None or not sim.is_playing(): return - # Kamino: run solver.reset() with the accumulated world mask to reinitialise - # internal state (warm-start containers, constraint multipliers) for reset worlds. - # Note: runs every step. solver.reset() with an all-False world_mask is a no-op - # (kernels check mask per-world and skip). The cost of a no-op launch is negligible - # compared to the complexity of maintaining a separate boolean guard. - cls._forward_kamino(world_mask=cls._world_reset_mask) + # Run solver.reset() only when invalidate_fk() flagged a pending reset: it rewrites + # joint_q_prev and re-snaps body_q, so it must not run on no-reset steps. + if cls._kamino_needs_reset: + cls._forward_kamino(world_mask=cls._world_reset_mask) + # Clear on the BASE class (see ``forward`` for the shadowing rationale). + NewtonManager._kamino_needs_reset = False # Notify solver of model changes if cls._model_changes: @@ -119,7 +245,20 @@ def _build_solver(cls, model: Model, solver_cfg: KaminoSolverCfg) -> None: Sets :attr:`NewtonManager._needs_collision_pipeline` to ``True`` only when ``use_collision_detector=False`` (Kamino's internal detector handles contacts otherwise). + + Applies :attr:`KaminoSolverCfg.max_contacts_per_world`, when set, by overriding + ``model.rigid_contact_max`` before solver construction. This bounds GPU memory + for contact-rich multi-env training that would otherwise over-allocate from + ``geoms.world_minimum_contacts``. """ + if solver_cfg.max_contacts_per_world is not None: + model.rigid_contact_max = int(solver_cfg.max_contacts_per_world) * model.world_count + logger.info( + "[KAMINO] Capping rigid_contact_max to %d (%d/world * %d worlds)", + model.rigid_contact_max, + solver_cfg.max_contacts_per_world, + model.world_count, + ) NewtonManager._solver = SolverKamino(model, solver_cfg.to_solver_config()) NewtonManager._use_single_state = False NewtonManager._needs_collision_pipeline = not solver_cfg.use_collision_detector diff --git a/source/isaaclab_newton/isaaclab_newton/physics/kamino_manager_cfg.py b/source/isaaclab_newton/isaaclab_newton/physics/kamino_manager_cfg.py index e88e54e59ab3..5ed701d0e5d9 100644 --- a/source/isaaclab_newton/isaaclab_newton/physics/kamino_manager_cfg.py +++ b/source/isaaclab_newton/isaaclab_newton/physics/kamino_manager_cfg.py @@ -145,6 +145,16 @@ class KaminoSolverCfg(NewtonSolverCfg): Only used when :attr:`use_collision_detector` is ``True``. If ``None``, Newton's default is used. """ + max_contacts_per_world: int | None = None + """Cap the per-world contact pre-allocation handed to Kamino. + + When ``None``, Kamino falls back to ``geoms.world_minimum_contacts`` derived from the + collision pipeline, which over-allocates dramatically for contact-rich assets. Set this + to bound GPU memory for multi-env training of contact-heavy tasks (e.g. legged + locomotion or manipulation). The total ``model.rigid_contact_max`` is computed as + ``max_contacts_per_world * model.world_count`` before solver construction. + """ + dynamics_preconditioning: bool = True """Whether to use preconditioning in the constrained dynamics solver. diff --git a/source/isaaclab_newton/isaaclab_newton/physics/newton_manager.py b/source/isaaclab_newton/isaaclab_newton/physics/newton_manager.py index 59464374e1f9..c725862c837e 100644 --- a/source/isaaclab_newton/isaaclab_newton/physics/newton_manager.py +++ b/source/isaaclab_newton/isaaclab_newton/physics/newton_manager.py @@ -108,13 +108,13 @@ def _sync_particle_points( 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.int32), + 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] = wp.int32(1) + world_mask[world] = True fk_mask[articulation_ids[world, arti]] = True @@ -122,16 +122,39 @@ def _or_reset_masks_from_mask( def _scatter_reset_masks_from_ids( env_ids: wp.array(dtype=int), articulation_ids: wp.array2d(dtype=int), - world_mask: wp.array(dtype=wp.int32), + 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 = env_ids[i] - world_mask[world] = wp.int32(1) + world_mask[world] = True fk_mask[articulation_ids[world, arti]] = True +@wp.kernel(enable_backward=False) +def _or_world_mask_from_env_mask( + env_mask: wp.array(dtype=wp.bool), + world_mask: wp.array(dtype=wp.bool), +): + """OR env_mask into world_mask. Used when articulation_ids is unavailable + (e.g. closed-loop assets whose root view does not expose them).""" + world = wp.tid() + if env_mask[world]: + world_mask[world] = True + + +@wp.kernel(enable_backward=False) +def _scatter_world_mask_from_ids( + env_ids: wp.array(dtype=int), + world_mask: wp.array(dtype=wp.bool), +): + """Scatter-set world_mask from sparse env_ids. Used when articulation_ids + is unavailable (e.g. closed-loop assets whose root view does not expose them).""" + i = wp.tid() + world_mask[env_ids[i]] = True + + class NewtonSceneDataBackend(SceneDataBackend): """Scene data backend that reads rigid body transforms from Newton's simulation state. @@ -228,8 +251,11 @@ class NewtonManager(PhysicsManager): _pending_extended_contact_attributes: set[str] = set() _report_contacts: bool = False # Per-world reset masks (allocated in start_simulation, consumed in step) - _world_reset_mask: wp.array | None = None # (num_envs,) wp.int32 — for SolverKamino.reset(world_mask=...) + _world_reset_mask: wp.array | None = None # (num_envs,) wp.bool — for SolverKamino.reset(world_mask=...) _fk_reset_mask: wp.array | None = None # (articulation_count,) wp.bool — for eval_fk(mask=...) + # Set by ``invalidate_fk`` and consumed by ``NewtonKaminoManager.step`` / ``forward`` + # so the Kamino solver reset only runs on steps where a reset actually occurred. + _kamino_needs_reset: bool = False # Newton actuator adapter (owns actuators and double-buffered states) _adapter: NewtonActuatorAdapter | None = None @@ -679,6 +705,7 @@ def clear(cls): # Per-world reset masks NewtonManager._world_reset_mask = None NewtonManager._fk_reset_mask = None + NewtonManager._kamino_needs_reset = False NewtonManager._graph = None NewtonManager._graph_capture_pending = False NewtonManager._newton_stage_path = None @@ -933,6 +960,9 @@ def invalidate_fk( if cls._world_reset_mask is None or cls._fk_reset_mask is None: return + # Flag a pending Kamino solver reset; consumed by NewtonKaminoManager. + NewtonManager._kamino_needs_reset = True + if articulation_ids is not None and env_mask is not None: wp.launch( _or_reset_masks_from_mask, @@ -949,9 +979,33 @@ def invalidate_fk( outputs=[NewtonManager._world_reset_mask, NewtonManager._fk_reset_mask], device=PhysicsManager._device, ) + elif env_mask is not None: + # No articulation_ids (e.g. closed-loop assets whose root view does not expose + # them): scope the world reset to the dirty envs only. Do NOT fall back to + # ``fill_(1)`` here — that would snap every env's body_q to the base state on the + # next Kamino reset, leaking one env's reset into its peers. + wp.launch( + _or_world_mask_from_env_mask, + dim=env_mask.shape, + inputs=[env_mask], + outputs=[NewtonManager._world_reset_mask], + device=PhysicsManager._device, + ) + # Closed-loop assets have articulation_count == 0, so this is a no-op on a + # zero-length array; for tree assets it conservatively refreshes all FK. + NewtonManager._fk_reset_mask.fill_(True) + elif env_ids is not None: + wp.launch( + _scatter_world_mask_from_ids, + dim=env_ids.shape, + inputs=[env_ids], + outputs=[NewtonManager._world_reset_mask], + device=PhysicsManager._device, + ) + NewtonManager._fk_reset_mask.fill_(True) else: - # Fallback: no topology info — mark everything dirty - NewtonManager._world_reset_mask.fill_(1) + # True fallback: no topology AND no env info — mark everything dirty. + NewtonManager._world_reset_mask.fill_(True) NewtonManager._fk_reset_mask.fill_(True) @classmethod @@ -982,8 +1036,13 @@ def start_simulation(cls) -> None: if cls._pending_extended_state_attributes: cls._builder.request_state_attributes(*cls._pending_extended_state_attributes) NewtonManager._pending_extended_state_attributes = set() + # Kamino works in maximal coordinates with closed loops, so assets without an + # ArticulationRootAPI (e.g. DR Legs) legitimately have joints outside any + # articulation; skip the tree-joint validation for Kamino only. + _solver_cfg = getattr(PhysicsManager._cfg, "solver_cfg", None) + _skip_joint_validation = getattr(_solver_cfg, "solver_type", "") == "kamino" with Timer(name="newton_finalize_builder", msg="Finalize builder took:"): - NewtonManager._model = cls._builder.finalize(device=device) + NewtonManager._model = cls._builder.finalize(device=device, skip_validation_joints=_skip_joint_validation) cls._model.set_gravity(cls._gravity_vector) cls._model.num_envs = cls._num_envs @@ -1007,7 +1066,7 @@ def start_simulation(cls) -> 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._world_reset_mask = wp.zeros(cls._model.world_count, dtype=wp.bool, device=device) NewtonManager._fk_reset_mask = wp.zeros(cls._model.articulation_count, dtype=wp.bool, device=device) logger.info("Dispatching PHYSICS_READY callbacks") diff --git a/source/isaaclab_newton/isaaclab_newton/sensors/contact_sensor/contact_sensor.py b/source/isaaclab_newton/isaaclab_newton/sensors/contact_sensor/contact_sensor.py index cee596d25e09..43f34695a475 100644 --- a/source/isaaclab_newton/isaaclab_newton/sensors/contact_sensor/contact_sensor.py +++ b/source/isaaclab_newton/isaaclab_newton/sensors/contact_sensor/contact_sensor.py @@ -401,7 +401,10 @@ def _update_buffers_impl(self, env_mask: wp.array): Args: env_mask: Mask of the environments to update. None: update all environments. """ - # Copy data from Newton into owned buffers (respecting env_mask) + # Copy data from Newton into owned buffers (respecting env_mask). This single path + # works for every solver backend: the Newton :class:`~newton.sensors.SensorContact` + # accumulates per-contact forces from ``Contacts.force``, which the MuJoCo-Warp and + # Kamino solvers both populate in their ``update_contacts`` conversions. # Launch with 3D for coalescing: dim=(num_envs, num_sensors, max(num_filter_objects, 1)) wp.launch( copy_from_newton_kernel,