[Newton] Fixes find joints / body masks

source/isaaclab/isaaclab/assets/articulation/articulation.py

@@ -17,7 +17,7 @@
     from isaaclab_newton.assets.articulation import ArticulationData as NewtonArticulationData
 
 
-class Articulation(FactoryBase):
+class Articulation(FactoryBase, BaseArticulation):
     """Factory for creating articulation instances."""
 
     data: BaseArticulationData | NewtonArticulationData

source/isaaclab/isaaclab/assets/articulation/base_articulation.py

@@ -18,6 +18,8 @@
 from ..asset_base import AssetBase
 
 if TYPE_CHECKING:
+    from isaaclab.utils.wrench_composer import WrenchComposer
+
     from .articulation_cfg import ArticulationCfg
     from .articulation_data import ArticulationData
 
@@ -176,12 +178,33 @@ def root_view(self):
         """
         raise NotImplementedError()
 
+    @property
+    @abstractmethod
+    def instantaneous_wrench_composer(self) -> WrenchComposer:
+        """Instantaneous wrench composer for the articulation."""
+        raise NotImplementedError()
+
+    @property
+    @abstractmethod
+    def permanent_wrench_composer(self) -> WrenchComposer:
+        """Permanent wrench composer for the articulation."""
+        raise NotImplementedError()
+
     """
     Operations.
     """
 
     @abstractmethod
-    def reset(self, env_ids: Sequence[int] | None = None, mask: wp.array | torch.Tensor | None = None):
+    def reset(self, env_ids: Sequence[int] | None = None, env_mask: wp.array | torch.Tensor | None = None):
+        """Reset the articulation.
+
+        Note: If both env_ids and env_mask are provided, then env_mask will be used. For performance reasons, it is
+        recommended to use the env_mask instead of env_ids.
+
+        Args:
+            env_ids: Environment indices. If None, then all indices are used.
+            env_mask: Environment mask. Shape is (num_instances,).
+        """
         raise NotImplementedError()
 
     @abstractmethod
@@ -831,6 +854,8 @@ def write_joint_armature_to_sim(
     def write_joint_friction_coefficient_to_sim(
         self,
         joint_friction_coeff: torch.Tensor | wp.array | float,
+        joint_dynamic_friction_coeff: torch.Tensor | wp.array | float | None = None,
+        joint_viscous_friction_coeff: torch.Tensor | wp.array | float | None = None,
         joint_ids: Sequence[int] | slice | None = None,
         env_ids: Sequence[int] | None = None,
         joint_mask: torch.Tensor | wp.array | None = None,
@@ -860,6 +885,8 @@ def write_joint_friction_coefficient_to_sim(
 
         Args:
             joint_friction_coeff: Joint static friction coefficient. Shape is (len(env_ids), len(joint_ids)) or (num_instances, num_joints).
+            joint_dynamic_friction_coeff: Joint dynamic friction coefficient. Shape is (len(env_ids), len(joint_ids)) or (num_instances, num_joints).
+            joint_viscous_friction_coeff: Joint viscous friction coefficient. Shape is (len(env_ids), len(joint_ids)) or (num_instances, num_joints).
             joint_ids: The joint indices to set the joint torque limits for. Defaults to None (all joints).
             env_ids: The environment indices to set the joint torque limits for. Defaults to None (all environments).
             joint_mask: The joint mask. Shape is (num_joints).
@@ -878,6 +905,17 @@ def write_joint_dynamic_friction_coefficient_to_sim(
     ):
         raise NotImplementedError()
 
+    @abstractmethod
+    def write_joint_viscous_friction_coefficient_to_sim(
+        self,
+        joint_viscous_friction_coeff: torch.Tensor | wp.array | float,
+        joint_ids: Sequence[int] | slice | None = None,
+        env_ids: Sequence[int] | None = None,
+        joint_mask: torch.Tensor | wp.array | None = None,
+        env_mask: torch.Tensor | wp.array | None = None,
+    ):
+        raise NotImplementedError()
+
     """
     Operations - Setters.
     """

source/isaaclab/isaaclab/assets/asset_base.py

@@ -15,6 +15,8 @@
 from collections.abc import Sequence
 from typing import TYPE_CHECKING, Any
 
+import warp as wp
+
 import isaaclab.sim as sim_utils
 import isaaclab.sim.utils.prims as prim_utils
 from isaaclab.sim import SimulationContext
@@ -243,11 +245,12 @@ def set_debug_vis(self, debug_vis: bool) -> bool:
         return True
 
     @abstractmethod
-    def reset(self, env_ids: Sequence[int] | None = None):
+    def reset(self, env_ids: Sequence[int] | None = None, env_mask: wp.array | torch.Tensor | None = None):
         """Resets all internal buffers of selected environments.
 
         Args:
             env_ids: The indices of the object to reset. Defaults to None (all instances).
+            env_mask: The mask of the object to reset. Defaults to None (all instances).
         """
         raise NotImplementedError
 

source/isaaclab/isaaclab/assets/rigid_object/base_rigid_object.py

@@ -12,6 +12,8 @@
 
 import warp as wp
 
+from isaaclab.utils.wrench_composer import WrenchComposer
+
 from ..asset_base import AssetBase
 
 if TYPE_CHECKING:
@@ -99,12 +101,33 @@ def root_view(self):
         """
         raise NotImplementedError()
 
+    @property
+    @abstractmethod
+    def instantaneous_wrench_composer(self) -> WrenchComposer:
+        """Instantaneous wrench composer for the rigid object."""
+        raise NotImplementedError()
+
+    @property
+    @abstractmethod
+    def permanent_wrench_composer(self) -> WrenchComposer:
+        """Permanent wrench composer for the rigid object."""
+        raise NotImplementedError()
+
     """
     Operations.
     """
 
     @abstractmethod
     def reset(self, env_ids: Sequence[int] | None = None, mask: wp.array | torch.Tensor | None = None):
+        """Reset the rigid object.
+
+        Note: If both env_ids and env_mask are provided, then env_mask will be used. For performance reasons, it is
+        recommended to use the env_mask instead of env_ids.
+
+        Args:
+            env_ids: Environment indices. If None, then all indices are used.
+            env_mask: Environment mask. Shape is (num_instances,).
+        """
         raise NotImplementedError()
 
     @abstractmethod
@@ -402,7 +425,7 @@ def set_masses(
         body_mask: wp.array | None = None,
         env_mask: wp.array | None = None,
     ):
-        """Set masses of all bodies in the simulation world frame.
+        """Set masses of all bodies.
 
         Args:
             masses: Masses of all bodies. Shape is (num_instances, num_bodies).
@@ -413,6 +436,26 @@ def set_masses(
         """
         raise NotImplementedError()
 
+    @abstractmethod
+    def set_coms(
+        self,
+        coms: torch.Tensor | wp.array,
+        body_ids: Sequence[int] | None = None,
+        env_ids: Sequence[int] | None = None,
+        body_mask: wp.array | None = None,
+        env_mask: wp.array | None = None,
+    ):
+        """Set center of mass positions of all bodies.
+
+        Args:
+            coms: Center of mass positions of all bodies. Shape is (num_instances, num_bodies, 3).
+            body_ids: The body indices to set the center of mass positions for. Defaults to None (all bodies).
+            env_ids: The environment indices to set the center of mass positions for. Defaults to None (all environments).
+            body_mask: The body mask. Shape is (num_bodies).
+            env_mask: The environment mask. Shape is (num_instances,).
+        """
+        raise NotImplementedError()
+
     @abstractmethod
     def set_inertias(
         self,
@@ -422,7 +465,7 @@ def set_inertias(
         body_mask: wp.array | None = None,
         env_mask: wp.array | None = None,
     ):
-        """Set inertias of all bodies in the simulation world frame.
+        """Set inertias of all bodies.
 
         Args:
             inertias: Inertias of all bodies. Shape is (num_instances, num_bodies, 3, 3).

source/isaaclab/isaaclab/assets/rigid_object/rigid_object.py

@@ -17,14 +17,14 @@
     from isaaclab_newton.assets.rigid_object import RigidObjectData as NewtonRigidObjectData
 
 
-class RigidObject(FactoryBase):
-    """Factory for creating articulation instances."""
+class RigidObject(FactoryBase, BaseRigidObject):
+    """Factory for creating rigid object instances."""
 
     data: BaseRigidObjectData | NewtonRigidObjectData
 
     def __new__(cls, *args, **kwargs) -> BaseRigidObject | NewtonRigidObject:
-        """Create a new instance of an articulation based on the backend."""
+        """Create a new instance of a rigid object based on the backend."""
         # The `FactoryBase` __new__ method will handle the logic and return
-        # an instance of the correct backend-specific articulation class,
+        # an instance of the correct backend-specific rigid object class,
         # which is guaranteed to be a subclass of `BaseArticulation` by convention.
         return super().__new__(cls, *args, **kwargs)

source/isaaclab/isaaclab/scene/interactive_scene.py

@@ -15,7 +15,7 @@
 
 import isaaclab.sim as sim_utils
 from isaaclab import cloner
-from isaaclab.assets import Articulation, ArticulationCfg, AssetBaseCfg
+from isaaclab.assets import Articulation, ArticulationCfg, AssetBaseCfg, RigidObject, RigidObjectCfg
 from isaaclab.sensors import ContactSensorCfg, SensorBase, SensorBaseCfg
 from isaaclab.sim import SimulationContext
 from isaaclab.sim.prims import XFormPrim
@@ -301,9 +301,9 @@ def deformable_objects(self) -> dict:
         raise NotImplementedError("Deformable objects are not supported in IsaacLab for Newton.")
 
     @property
-    def rigid_objects(self) -> dict:
+    def rigid_objects(self) -> dict[str, RigidObject]:
         """A dictionary of rigid objects in the scene."""
-        raise NotImplementedError("Rigid objects are not supported in IsaacLab for Newton.")
+        return self._rigid_objects
 
     @property
     def rigid_object_collections(self) -> dict:
@@ -348,26 +348,30 @@ def state(self) -> dict[str, dict[str, dict[str, torch.Tensor]]]:
     Operations.
     """
 
-    def reset(self, env_ids: Sequence[int] | None = None, mask: wp.array | torch.Tensor | None = None):
+    def reset(self, env_ids: Sequence[int] | None = None, env_mask: wp.array | torch.Tensor | None = None):
         """Resets the scene entities.
 
         Args:
-            env_ids: The indices of the environments to reset.
-                Defaults to None (all instances).
+            env_ids: The indices of the environments to reset. Defaults to None (all instances).
+            env_mask: The mask of the environments to reset. Defaults to None (all instances).
         """
         # FIXME: Homogenize the API for env_ids and env_mask.
         # -- assets
         for articulation in self._articulations.values():
-            articulation.reset(ids=env_ids, mask=mask)
+            articulation.reset(env_ids=env_ids, env_mask=env_mask)
+        for rigid_object in self._rigid_objects.values():
+            rigid_object.reset(env_ids=env_ids, env_mask=env_mask)
         # -- sensors
         for sensor in self._sensors.values():
-            sensor.reset(env_ids=env_ids, env_mask=mask)
+            sensor.reset(env_ids=env_ids, env_mask=env_mask)
 
     def write_data_to_sim(self):
         """Writes the data of the scene entities to the simulation."""
         # -- assets
         for articulation in self._articulations.values():
             articulation.write_data_to_sim()
+        for rigid_object in self._rigid_objects.values():
+            rigid_object.write_data_to_sim()
 
     def update(self, dt: float) -> None:
         """Update the scene entities.
@@ -378,6 +382,8 @@ def update(self, dt: float) -> None:
         # -- assets
         for articulation in self._articulations.values():
             articulation.update(dt)
+        for rigid_object in self._rigid_objects.values():
+            rigid_object.update(dt)
         # -- sensors
         for sensor in self._sensors.values():
             sensor.update(dt, force_recompute=not self.cfg.lazy_sensor_update)
@@ -422,7 +428,15 @@ def reset_to(
             #   This assumption does not hold for effort controlled joints.
             articulation.set_joint_position_target(joint_position, env_ids=env_ids)
             articulation.set_joint_velocity_target(joint_velocity, env_ids=env_ids)
-
+        # rigid objects
+        for asset_name, rigid_object in self._rigid_objects.items():
+            asset_state = state["rigid_object"][asset_name]
+            root_pose = asset_state["root_pose"].clone()
+            if is_relative:
+                root_pose[:, :3] += self.env_origins[env_ids]
+            root_velocity = asset_state["root_velocity"].clone()
+            rigid_object.write_root_pose_to_sim(root_pose, env_ids=env_ids)
+            rigid_object.write_root_velocity_to_sim(root_velocity, env_ids=env_ids)
         # write data to simulation to make sure initial state is set
         # this propagates the joint targets to the simulation
         self.write_data_to_sim()
@@ -490,6 +504,15 @@ def get_state(self, is_relative: bool = False) -> dict[str, dict[str, dict[str,
             asset_state["joint_position"] = articulation.data.joint_pos.clone()
             asset_state["joint_velocity"] = articulation.data.joint_vel.clone()
             state["articulation"][asset_name] = asset_state
+        # rigid objects
+        state["rigid_object"] = dict()
+        for asset_name, rigid_object in self._rigid_objects.items():
+            asset_state = dict()
+            asset_state["root_pose"] = rigid_object.data.root_pose_w.clone()
+            if is_relative:
+                asset_state["root_pose"][:, :3] -= self.env_origins
+            asset_state["root_velocity"] = rigid_object.data.root_vel_w.clone()
+            state["rigid_object"][asset_name] = asset_state
         return state
 
     """
@@ -505,6 +528,7 @@ def keys(self) -> list[str]:
         all_keys = ["terrain"]
         for asset_family in [
             self._articulations,
+            self._rigid_objects,
             self._sensors,
             self._extras,
         ]:
@@ -528,6 +552,7 @@ def __getitem__(self, key: str) -> Any:
         # check if it is in other dictionaries
         for asset_family in [
             self._articulations,
+            self._rigid_objects,
             self._sensors,
             self._extras,
         ]:
@@ -584,6 +609,8 @@ def _add_entities_from_cfg(self):
                 self._terrain = asset_cfg.class_type(asset_cfg)
             elif isinstance(asset_cfg, ArticulationCfg):
                 self._articulations[asset_name] = asset_cfg.class_type(asset_cfg)
+            elif isinstance(asset_cfg, RigidObjectCfg):
+                self._rigid_objects[asset_name] = asset_cfg.class_type(asset_cfg)
             elif isinstance(asset_cfg, SensorBaseCfg):
                 if isinstance(asset_cfg, ContactSensorCfg):
                     if asset_cfg.shape_path is not None:

source/isaaclab/isaaclab/sensors/sensor_base.py

@@ -207,7 +207,8 @@ def reset(self, env_ids: Sequence[int] | None = None, env_mask: wp.array | None
         """Resets the sensor internals.
 
         Args:
-            env_ids: The sensor ids to reset. Defaults to None.
+            env_ids: The sensor ids to reset. Defaults to None (all instances).
+            env_mask: The sensor mask to reset. Defaults to None (all instances).
         """
         # Resolve sensor ids
         if env_ids is None:

source/isaaclab/isaaclab/sim/_impl/newton_manager.py

@@ -5,20 +5,23 @@
 
 from __future__ import annotations
 
+import logging
 import numpy as np
 import re
 
 import warp as wp
 from newton import Axis, Contacts, Control, Model, ModelBuilder, State, eval_fk
 from newton.examples import create_collision_pipeline
-from newton.sensors import ContactSensor as NewtonContactSensor
+from newton.sensors import SensorContact as NewtonContactSensor
 from newton.sensors import populate_contacts
 from newton.solvers import SolverBase, SolverFeatherstone, SolverMuJoCo, SolverNotifyFlags, SolverXPBD
 
 from isaaclab.sim._impl.newton_manager_cfg import NewtonCfg
 from isaaclab.sim.utils.stage import get_current_stage
 from isaaclab.utils.timer import Timer
 
+logger = logging.getLogger(__name__)
+
 
 def flipped_match(x: str, y: str) -> re.Match | None:
     """Flipped match function.
@@ -197,15 +200,15 @@ def initialize_solver(cls):
             else:
                 NewtonManager._needs_collision_pipeline = True
 
-        # Ensure we are using a CUDA enabled device
-        assert NewtonManager._device.startswith("cuda"), "NewtonManager only supports CUDA enabled devices"
-
         # Capture the graph if CUDA is enabled
         with Timer(name="newton_cuda_graph", msg="CUDA graph took:", enable=True, format="ms"):
-            if NewtonManager._cfg.use_cuda_graph:
+            if NewtonManager._cfg.use_cuda_graph and NewtonManager._device.startswith("cuda"):
                 with wp.ScopedCapture() as capture:
                     NewtonManager.simulate()
                 NewtonManager._graph = capture.graph
+            elif NewtonManager._cfg.use_cuda_graph and not NewtonManager._device.startswith("cuda"):
+                logger.warning("CUDA graphs requested but device is CPU. Disabling CUDA graphs.")
+                NewtonManager._cfg.use_cuda_graph = False
 
     @classmethod
     def simulate(cls) -> None: