diff --git a/source/isaaclab/isaaclab/__init__.py b/source/isaaclab/isaaclab/__init__.py index 7bb50a0eab2..13518a1f1cb 100644 --- a/source/isaaclab/isaaclab/__init__.py +++ b/source/isaaclab/isaaclab/__init__.py @@ -5,6 +5,7 @@ """Package containing the core framework.""" +from enum import IntEnum import os import toml @@ -17,3 +18,7 @@ # Configure the module-level variables __version__ = ISAACLAB_METADATA["package"]["version"] + +class Backend(IntEnum): + NEWTON = 0 + PHYSX = 1 \ No newline at end of file diff --git a/source/isaaclab/isaaclab/actuators/__init__.py b/source/isaaclab/isaaclab/actuators/__init__.py index 5ccf5d7b082..e4940f06a66 100644 --- a/source/isaaclab/isaaclab/actuators/__init__.py +++ b/source/isaaclab/isaaclab/actuators/__init__.py @@ -21,6 +21,12 @@ which defines the common interface for all actuator models. The actuator models are handled and called by the :class:`isaaclab.assets.Articulation` class. """ +from enum import IntEnum +class ControlMode(IntEnum): + """Enum for the control mode of the actuator.""" + NONE = 0 + POSITION = 1 + VELOCITY = 2 from .actuator_base import ActuatorBase from .actuator_cfg import ( diff --git a/source/isaaclab/isaaclab/actuators/actuator_base.py b/source/isaaclab/isaaclab/actuators/actuator_base.py index 60d0c46f1f1..e4c47f27e4b 100644 --- a/source/isaaclab/isaaclab/actuators/actuator_base.py +++ b/source/isaaclab/isaaclab/actuators/actuator_base.py @@ -8,9 +8,8 @@ import torch from abc import ABC, abstractmethod from collections.abc import Sequence -from typing import TYPE_CHECKING, ClassVar, Literal +from typing import TYPE_CHECKING, ClassVar -import isaaclab.utils.string as string_utils from isaaclab.utils.types import ArticulationActions if TYPE_CHECKING: @@ -36,68 +35,15 @@ class ActuatorBase(ABC): To see how the class is used, check the :class:`isaaclab.assets.Articulation` class. """ + __backend_name__: str = "base" + """The name of the backend for the actuator.""" + is_implicit_model: ClassVar[bool] = False """Flag indicating if the actuator is an implicit or explicit actuator model. If a class inherits from :class:`ImplicitActuator`, then this flag should be set to :obj:`True`. """ - computed_effort: torch.Tensor - """The computed effort for the actuator group. Shape is (num_envs, num_joints).""" - - applied_effort: torch.Tensor - """The applied effort for the actuator group. Shape is (num_envs, num_joints). - - This is the effort obtained after clipping the :attr:`computed_effort` based on the - actuator characteristics. - """ - - effort_limit: torch.Tensor - """The effort limit for the actuator group. Shape is (num_envs, num_joints). - - For implicit actuators, the :attr:`effort_limit` and :attr:`effort_limit_sim` are the same. - """ - - effort_limit_sim: torch.Tensor - """The effort limit for the actuator group in the simulation. Shape is (num_envs, num_joints). - - For implicit actuators, the :attr:`effort_limit` and :attr:`effort_limit_sim` are the same. - """ - - velocity_limit: torch.Tensor - """The velocity limit for the actuator group. Shape is (num_envs, num_joints). - - For implicit actuators, the :attr:`velocity_limit` and :attr:`velocity_limit_sim` are the same. - """ - - velocity_limit_sim: torch.Tensor - """The velocity limit for the actuator group in the simulation. Shape is (num_envs, num_joints). - - For implicit actuators, the :attr:`velocity_limit` and :attr:`velocity_limit_sim` are the same. - """ - - control_mode: Literal["position", "velocity", "none"] - """The control mode of the actuator. - - The control mode can be one of the following: - - * ``"position"``: Position control - * ``"velocity"``: Velocity control - * ``"none"``: No control (used for explicit actuators or direct effort control) - """ - - stiffness: torch.Tensor - """The stiffness (P gain) of the PD controller. Shape is (num_envs, num_joints).""" - - damping: torch.Tensor - """The damping (D gain) of the PD controller. Shape is (num_envs, num_joints).""" - - armature: torch.Tensor - """The armature of the actuator joints. Shape is (num_envs, num_joints).""" - - friction: torch.Tensor - """The joint friction of the actuator joints. Shape is (num_envs, num_joints).""" - _DEFAULT_MAX_EFFORT_SIM: ClassVar[float] = 1.0e9 """The default maximum effort for the actuator joints in the simulation. Defaults to 1.0e9. @@ -108,17 +54,6 @@ class ActuatorBase(ABC): def __init__( self, cfg: ActuatorBaseCfg, - joint_names: list[str], - joint_ids: slice | torch.Tensor, - num_envs: int, - device: str, - control_mode: Literal["position", "velocity"] = "position", - stiffness: torch.Tensor | float = 0.0, - damping: torch.Tensor | float = 0.0, - armature: torch.Tensor | float = 0.0, - friction: torch.Tensor | float = 0.0, - effort_limit: torch.Tensor | float = torch.inf, - velocity_limit: torch.Tensor | float = torch.inf, ): """Initialize the actuator. @@ -131,56 +66,9 @@ def __init__( Args: cfg: The configuration of the actuator model. - joint_names: The joint names in the articulation. - joint_ids: The joint indices in the articulation. If :obj:`slice(None)`, then all - the joints in the articulation are part of the group. - num_envs: Number of articulations in the view. - device: Device used for processing. - control_mode: The default control mode. Defaults to "position". - stiffness: The default joint stiffness (P gain). Defaults to 0.0. - If a tensor, then the shape is (num_envs, num_joints). - damping: The default joint damping (D gain). Defaults to 0.0. - If a tensor, then the shape is (num_envs, num_joints). - armature: The default joint armature. Defaults to 0.0. - If a tensor, then the shape is (num_envs, num_joints). - friction: The default joint friction. Defaults to 0.0. - If a tensor, then the shape is (num_envs, num_joints). - effort_limit: The default effort limit. Defaults to infinity. - If a tensor, then the shape is (num_envs, num_joints). - velocity_limit: The default velocity limit. Defaults to infinity. - If a tensor, then the shape is (num_envs, num_joints). """ # save parameters self.cfg = cfg - self._num_envs = num_envs - self._device = device - self._joint_names = joint_names - self._joint_indices = joint_ids - - # For explicit models, we do not want to enforce the effort limit through the solver - # (unless it is explicitly set) - if not self.is_implicit_model and self.cfg.effort_limit_sim is None: - self.cfg.effort_limit_sim = self._DEFAULT_MAX_EFFORT_SIM - - # parse control mode - self.control_mode = control_mode - # parse joint stiffness and damping - self.stiffness = self._parse_joint_parameter(self.cfg.stiffness, stiffness) - self.damping = self._parse_joint_parameter(self.cfg.damping, damping) - # parse joint armature and friction - self.armature = self._parse_joint_parameter(self.cfg.armature, armature) - self.friction = self._parse_joint_parameter(self.cfg.friction, friction) - # parse joint limits - # -- velocity - self.velocity_limit_sim = self._parse_joint_parameter(self.cfg.velocity_limit_sim, velocity_limit) - self.velocity_limit = self._parse_joint_parameter(self.cfg.velocity_limit, self.velocity_limit_sim) - # -- effort - self.effort_limit_sim = self._parse_joint_parameter(self.cfg.effort_limit_sim, effort_limit) - self.effort_limit = self._parse_joint_parameter(self.cfg.effort_limit, self.effort_limit_sim) - - # create commands buffers for allocation - self.computed_effort = torch.zeros(self._num_envs, self.num_joints, device=self._device) - self.applied_effort = torch.zeros_like(self.computed_effort) def __str__(self) -> str: """Returns: A string representation of the actuator group.""" @@ -205,16 +93,19 @@ def __str__(self) -> str: """ @property + @abstractmethod def num_joints(self) -> int: """Number of actuators in the group.""" - return len(self._joint_names) + raise NotImplementedError @property + @abstractmethod def joint_names(self) -> list[str]: """Articulation's joint names that are part of the group.""" - return self._joint_names + raise NotImplementedError @property + @abstractmethod def joint_indices(self) -> slice | torch.Tensor: """Articulation's joint indices that are part of the group. @@ -222,7 +113,7 @@ def joint_indices(self) -> slice | torch.Tensor: If :obj:`slice(None)` is returned, then the group contains all the joints in the articulation. We do this to avoid unnecessary indexing of the joints for performance reasons. """ - return self._joint_indices + raise NotImplementedError """ Operations. @@ -260,6 +151,11 @@ def compute( Helper functions. """ + @abstractmethod + def _record_actuator_resolution(self, cfg_val, new_val, usd_val, joint_names, joint_ids, actuator_param: str): + raise NotImplementedError + + @abstractmethod def _parse_joint_parameter( self, cfg_value: float | dict[str, float] | None, default_value: float | torch.Tensor | None ) -> torch.Tensor: @@ -279,47 +175,9 @@ def _parse_joint_parameter( ValueError: If the parameter value is None and no default value is provided. ValueError: If the default value tensor is the wrong shape. """ - # create parameter buffer - param = torch.zeros(self._num_envs, self.num_joints, device=self._device) - # parse the parameter - if cfg_value is not None: - if isinstance(cfg_value, (float, int)): - # if float, then use the same value for all joints - param[:] = float(cfg_value) - elif isinstance(cfg_value, dict): - # if dict, then parse the regular expression - indices, _, values = string_utils.resolve_matching_names_values(cfg_value, self.joint_names) - # note: need to specify type to be safe (e.g. values are ints, but we want floats) - param[:, indices] = torch.tensor(values, dtype=torch.float, device=self._device) - else: - raise TypeError( - f"Invalid type for parameter value: {type(cfg_value)} for " - + f"actuator on joints {self.joint_names}. Expected float or dict." - ) - elif default_value is not None: - if isinstance(default_value, (float, int)): - # if float, then use the same value for all joints - param[:] = float(default_value) - elif isinstance(default_value, torch.Tensor): - # if tensor, then use the same tensor for all joints - if default_value.shape == (self._num_envs, self.num_joints): - param = default_value.float() - else: - raise ValueError( - "Invalid default value tensor shape.\n" - f"Got: {default_value.shape}\n" - f"Expected: {(self._num_envs, self.num_joints)}" - ) - else: - raise TypeError( - f"Invalid type for default value: {type(default_value)} for " - + f"actuator on joints {self.joint_names}. Expected float or Tensor." - ) - else: - raise ValueError("The parameter value is None and no default value is provided.") - - return param + raise NotImplementedError + @abstractmethod def _clip_effort(self, effort: torch.Tensor) -> torch.Tensor: """Clip the desired torques based on the motor limits. @@ -329,4 +187,4 @@ def _clip_effort(self, effort: torch.Tensor) -> torch.Tensor: Returns: The clipped torques. """ - return torch.clip(effort, min=-self.effort_limit, max=self.effort_limit) + raise NotImplementedError diff --git a/source/isaaclab/isaaclab/actuators/actuator_cfg.py b/source/isaaclab/isaaclab/actuators/actuator_cfg.py index 98202234c25..ba25223c3d7 100644 --- a/source/isaaclab/isaaclab/actuators/actuator_cfg.py +++ b/source/isaaclab/isaaclab/actuators/actuator_cfg.py @@ -10,14 +10,14 @@ from isaaclab.utils import configclass from . import actuator_net, actuator_pd -from .actuator_base import ActuatorBase +from . import ControlMode @configclass class ActuatorBaseCfg: """Configuration for default actuators in an articulation.""" - class_type: type[ActuatorBase] = MISSING + class_type: type = MISSING """The associated actuator class. The class should inherit from :class:`isaaclab.actuators.ActuatorBase`. @@ -109,16 +109,6 @@ class ActuatorBaseCfg: """ - control_mode: Literal["position", "velocity", "none"] = "position" - """Control mode of the actuator. Defaults to "position". - - The control mode can be one of the following: - - * ``"position"``: Position control - * ``"velocity"``: Velocity control - * ``"none"``: No control (used for explicit actuators or direct effort control) - """ - stiffness: dict[str, float] | float | None = MISSING """Stiffness gains (also known as p-gain) of the joints in the group. @@ -151,18 +141,36 @@ class ActuatorBaseCfg: """ friction: dict[str, float] | float | None = None - r"""The friction coefficient of the joints in the group. Defaults to None. + r"""The static friction coefficient of the joints in the group. Defaults to None. - The joint friction is a unitless quantity. It relates the magnitude of the spatial force transmitted - from the parent body to the child body to the maximal friction force that may be applied by the solver + The joint static friction is a unitless quantity. It relates the magnitude of the spatial force transmitted + from the parent body to the child body to the maximal static friction force that may be applied by the solver to resist the joint motion. Mathematically, this means that: :math:`F_{resist} \leq \mu F_{spatial}`, where :math:`F_{resist}` is the resisting force applied by the solver and :math:`F_{spatial}` is the spatial force - transmitted from the parent body to the child body. The simulated friction effect is therefore - similar to static and Coulomb friction. + transmitted from the parent body to the child body. The simulated static friction effect is therefore + similar to static and Coulomb static friction. + + If None, the joint static friction is set to the value from the USD joint prim. + """ + + dynamic_friction: dict[str, float] | float | None = None + """The dynamic friction coefficient of the joints in the group. Defaults to None. + """ + + viscous_friction: dict[str, float] | float | None = None + """The viscous friction coefficient of the joints in the group. Defaults to None. + """ + + control_mode: ControlMode = ControlMode.POSITION + """The control mode of the actuator. Defaults to :obj:`ControlMode.POSITION`. + + The control mode can be one of the following: - If None, the joint friction is set to the value from the USD joint prim. + * ``"position"``: position control + * ``"velocity"``: velocity control + * ``"effort"``: effort control """ diff --git a/source/isaaclab/isaaclab/actuators/actuator_net.py b/source/isaaclab/isaaclab/actuators/actuator_net.py index e3dbc765b13..b6ba0a04946 100644 --- a/source/isaaclab/isaaclab/actuators/actuator_net.py +++ b/source/isaaclab/isaaclab/actuators/actuator_net.py @@ -12,180 +12,19 @@ """ -from __future__ import annotations +from isaaclab.utils.backend_utils import FactoryBase +from .actuator_base import ActuatorBase -import torch -from collections.abc import Sequence -from typing import TYPE_CHECKING +class ActuatorNetLSTM(FactoryBase): + """Factory for creating LSTM-based actuator models.""" -from isaaclab.utils.assets import read_file -from isaaclab.utils.types import ArticulationActions + def __new__(cls, backend: str, *args, **kwargs) -> ActuatorBase: + """Create a new instance of an LSTM-based actuator model based on the backend.""" + return super().__new__(cls, backend, *args, **kwargs) -from .actuator_pd import DCMotor +class ActuatorNetMLP(FactoryBase): + """Factory for creating MLP-based actuator models.""" -if TYPE_CHECKING: - from .actuator_cfg import ActuatorNetLSTMCfg, ActuatorNetMLPCfg - - -class ActuatorNetLSTM(DCMotor): - """Actuator model based on recurrent neural network (LSTM). - - Unlike the MLP implementation :cite:t:`hwangbo2019learning`, this class implements - the learned model as a temporal neural network (LSTM) based on the work from - :cite:t:`rudin2022learning`. This removes the need of storing a history as the - hidden states of the recurrent network captures the history. - - Note: - Only the desired joint positions are used as inputs to the network. - """ - - cfg: ActuatorNetLSTMCfg - """The configuration of the actuator model.""" - - def __init__(self, cfg: ActuatorNetLSTMCfg, *args, **kwargs): - super().__init__(cfg, *args, **kwargs) - - assert self.cfg.control_mode == "position", "ActuatorNetLSTM only supports position control" - - # load the model from JIT file - file_bytes = read_file(self.cfg.network_file) - self.network = torch.jit.load(file_bytes, map_location=self._device).eval() - - # extract number of lstm layers and hidden dim from the shape of weights - num_layers = len(self.network.lstm.state_dict()) // 4 - hidden_dim = self.network.lstm.state_dict()["weight_hh_l0"].shape[1] - # create buffers for storing LSTM inputs - self.sea_input = torch.zeros(self._num_envs * self.num_joints, 1, 2, device=self._device) - self.sea_hidden_state = torch.zeros( - num_layers, self._num_envs * self.num_joints, hidden_dim, device=self._device - ) - self.sea_cell_state = torch.zeros(num_layers, self._num_envs * self.num_joints, hidden_dim, device=self._device) - # reshape via views (doesn't change the actual memory layout) - layer_shape_per_env = (num_layers, self._num_envs, self.num_joints, hidden_dim) - self.sea_hidden_state_per_env = self.sea_hidden_state.view(layer_shape_per_env) - self.sea_cell_state_per_env = self.sea_cell_state.view(layer_shape_per_env) - - """ - Operations. - """ - - def reset(self, env_ids: Sequence[int]): - # reset the hidden and cell states for the specified environments - with torch.no_grad(): - self.sea_hidden_state_per_env[:, env_ids] = 0.0 - self.sea_cell_state_per_env[:, env_ids] = 0.0 - - def compute( - self, control_action: ArticulationActions, joint_pos: torch.Tensor, joint_vel: torch.Tensor - ) -> ArticulationActions: - # compute network inputs - self.sea_input[:, 0, 0] = (control_action.joint_targets - joint_pos).flatten() - self.sea_input[:, 0, 1] = joint_vel.flatten() - # save current joint vel for dc-motor clipping - self._joint_vel[:] = joint_vel - - # run network inference - with torch.inference_mode(): - torques, (self.sea_hidden_state[:], self.sea_cell_state[:]) = self.network( - self.sea_input, (self.sea_hidden_state, self.sea_cell_state) - ) - self.computed_effort = torques.reshape(self._num_envs, self.num_joints) - - # clip the computed effort based on the motor limits - self.applied_effort = self._clip_effort(self.computed_effort) - - # return torques - control_action.joint_efforts = self.applied_effort - control_action.joint_targets = None - return control_action - - -class ActuatorNetMLP(DCMotor): - """Actuator model based on multi-layer perceptron and joint history. - - Many times the analytical model is not sufficient to capture the actuator dynamics, the - delay in the actuator response, or the non-linearities in the actuator. In these cases, - a neural network model can be used to approximate the actuator dynamics. This model is - trained using data collected from the physical actuator and maps the joint state and the - desired joint command to the produced torque by the actuator. - - This class implements the learned model as a neural network based on the work from - :cite:t:`hwangbo2019learning`. The class stores the history of the joint positions errors - and velocities which are used to provide input to the neural network. The model is loaded - as a TorchScript. - - Note: - Only the desired joint positions are used as inputs to the network. - - """ - - cfg: ActuatorNetMLPCfg - """The configuration of the actuator model.""" - - def __init__(self, cfg: ActuatorNetMLPCfg, *args, **kwargs): - super().__init__(cfg, *args, **kwargs) - - assert self.cfg.control_mode == "position", "ActuatorNetLSTM only supports position control" - - # load the model from JIT file - file_bytes = read_file(self.cfg.network_file) - self.network = torch.jit.load(file_bytes, map_location=self._device).eval() - - # create buffers for MLP history - history_length = max(self.cfg.input_idx) + 1 - self._joint_pos_error_history = torch.zeros( - self._num_envs, history_length, self.num_joints, device=self._device - ) - self._joint_vel_history = torch.zeros(self._num_envs, history_length, self.num_joints, device=self._device) - - """ - Operations. - """ - - def reset(self, env_ids: Sequence[int]): - # reset the history for the specified environments - self._joint_pos_error_history[env_ids] = 0.0 - self._joint_vel_history[env_ids] = 0.0 - - def compute( - self, control_action: ArticulationActions, joint_pos: torch.Tensor, joint_vel: torch.Tensor - ) -> ArticulationActions: - # move history queue by 1 and update top of history - # -- positions - self._joint_pos_error_history = self._joint_pos_error_history.roll(1, 1) - self._joint_pos_error_history[:, 0] = control_action.joint_targets - joint_pos - # -- velocity - self._joint_vel_history = self._joint_vel_history.roll(1, 1) - self._joint_vel_history[:, 0] = joint_vel - # save current joint vel for dc-motor clipping - self._joint_vel[:] = joint_vel - - # compute network inputs - # -- positions - pos_input = torch.cat([self._joint_pos_error_history[:, i].unsqueeze(2) for i in self.cfg.input_idx], dim=2) - pos_input = pos_input.view(self._num_envs * self.num_joints, -1) - # -- velocity - vel_input = torch.cat([self._joint_vel_history[:, i].unsqueeze(2) for i in self.cfg.input_idx], dim=2) - vel_input = vel_input.view(self._num_envs * self.num_joints, -1) - # -- scale and concatenate inputs - if self.cfg.input_order == "pos_vel": - network_input = torch.cat([pos_input * self.cfg.pos_scale, vel_input * self.cfg.vel_scale], dim=1) - elif self.cfg.input_order == "vel_pos": - network_input = torch.cat([vel_input * self.cfg.vel_scale, pos_input * self.cfg.pos_scale], dim=1) - else: - raise ValueError( - f"Invalid input order for MLP actuator net: {self.cfg.input_order}. Must be 'pos_vel' or 'vel_pos'." - ) - - # run network inference - with torch.inference_mode(): - torques = self.network(network_input).view(self._num_envs, self.num_joints) - self.computed_effort = torques.view(self._num_envs, self.num_joints) * self.cfg.torque_scale - - # clip the computed effort based on the motor limits - self.applied_effort = self._clip_effort(self.computed_effort) - - # return torques - control_action.joint_efforts = self.applied_effort - control_action.joint_targets = None - return control_action + def __new__(cls, backend: str, *args, **kwargs) -> ActuatorBase: + """Create a new instance of an MLP-based actuator model based on the backend.""" + return super().__new__(cls, backend, *args, **kwargs) diff --git a/source/isaaclab/isaaclab/actuators/actuator_pd.py b/source/isaaclab/isaaclab/actuators/actuator_pd.py index d8e0699294b..540877783bc 100644 --- a/source/isaaclab/isaaclab/actuators/actuator_pd.py +++ b/source/isaaclab/isaaclab/actuators/actuator_pd.py @@ -2,430 +2,58 @@ # All rights reserved. # # SPDX-License-Identifier: BSD-3-Clause - from __future__ import annotations - -import torch -from collections.abc import Sequence +from isaaclab.utils.backend_utils import FactoryBase from typing import TYPE_CHECKING -import omni.log - -from isaaclab.utils import DelayBuffer, LinearInterpolation -from isaaclab.utils.types import ArticulationActions - from .actuator_base import ActuatorBase if TYPE_CHECKING: - from .actuator_cfg import ( - DCMotorCfg, - DelayedPDActuatorCfg, - IdealPDActuatorCfg, - ImplicitActuatorCfg, - RemotizedPDActuatorCfg, - ) - + from isaaclab.newton.actuators.actuator_pd import ImplicitActuator as NewtonImplicitActuator + from isaaclab.newton.actuators.actuator_pd import IdealPDActuator as NewtonIdealPDActuator + from isaaclab.newton.actuators.actuator_pd import DCMotor as NewtonDCMotor + #from isaaclab.newton.actuators.actuator_pd import DelayedPDActuator as NewtonDelayedPDActuator + #from isaaclab.newton.actuators.actuator_pd import RemotizedPDActuator as NewtonRemotizedPDActuator """ Implicit Actuator Models. """ +class ImplicitActuator(FactoryBase): + """Factory for creating implicit actuator models.""" -class ImplicitActuator(ActuatorBase): - """Implicit actuator model that is handled by the simulation. - - This performs a similar function as the :class:`IdealPDActuator` class. However, the PD control is handled - implicitly by the simulation which performs continuous-time integration of the PD control law. This is - generally more accurate than the explicit PD control law used in :class:`IdealPDActuator` when the simulation - time-step is large. - - The articulation class sets the stiffness and damping parameters from the implicit actuator configuration - into the simulation. Thus, the class does not perform its own computations on the joint action that - needs to be applied to the simulation. However, it computes the approximate torques for the actuated joint - since PhysX does not expose this quantity explicitly. - - .. caution:: - - The class is only provided for consistency with the other actuator models. It does not implement any - functionality and should not be used. All values should be set to the simulation directly. - """ - - cfg: ImplicitActuatorCfg - """The configuration for the actuator model.""" - - def __init__(self, cfg: ImplicitActuatorCfg, *args, **kwargs): - # effort limits - if cfg.effort_limit_sim is None and cfg.effort_limit is not None: - # throw a warning that we have a replacement for the deprecated parameter - omni.log.warn( - "The object has a value for 'effort_limit'." - " This parameter will be removed in the future." - " To set the effort limit, please use 'effort_limit_sim' instead." - ) - cfg.effort_limit_sim = cfg.effort_limit - elif cfg.effort_limit_sim is not None and cfg.effort_limit is None: - # TODO: Eventually we want to get rid of 'effort_limit' for implicit actuators. - # We should do this once all parameters have an "_sim" suffix. - cfg.effort_limit = cfg.effort_limit_sim - elif cfg.effort_limit_sim is not None and cfg.effort_limit is not None: - if cfg.effort_limit_sim != cfg.effort_limit: - raise ValueError( - "The object has set both 'effort_limit_sim' and 'effort_limit'" - f" and they have different values {cfg.effort_limit_sim} != {cfg.effort_limit}." - " Please only set 'effort_limit_sim' for implicit actuators." - ) - - # velocity limits - if cfg.velocity_limit_sim is None and cfg.velocity_limit is not None: - # throw a warning that previously this was not set - # it leads to different simulation behavior so we want to remain backwards compatible - omni.log.warn( - "The object has a value for 'velocity_limit'." - " Previously, although this value was specified, it was not getting used by implicit" - " actuators. Since this parameter affects the simulation behavior, we continue to not" - " use it. This parameter will be removed in the future." - " To set the velocity limit, please use 'velocity_limit_sim' instead." - ) - cfg.velocity_limit = None - elif cfg.velocity_limit_sim is not None and cfg.velocity_limit is None: - # TODO: Eventually we want to get rid of 'velocity_limit' for implicit actuators. - # We should do this once all parameters have an "_sim" suffix. - cfg.velocity_limit = cfg.velocity_limit_sim - elif cfg.velocity_limit_sim is not None and cfg.velocity_limit is not None: - if cfg.velocity_limit_sim != cfg.velocity_limit: - raise ValueError( - "The object has set both 'velocity_limit_sim' and 'velocity_limit'" - f" and they have different values {cfg.velocity_limit_sim} != {cfg.velocity_limit}." - " Please only set 'velocity_limit_sim' for implicit actuators." - ) - - # set implicit actuator model flag - ImplicitActuator.is_implicit_model = True - # call the base class - super().__init__(cfg, *args, **kwargs) - - """ - Operations. - """ - - def reset(self, *args, **kwargs): - # This is a no-op. There is no state to reset for implicit actuators. - pass - - def compute( - self, control_action: ArticulationActions, joint_pos: torch.Tensor, joint_vel: torch.Tensor - ) -> ArticulationActions: - """Process the actuator group actions and compute the articulation actions. - - In case of implicit actuator, the control action is directly returned as the computed action. - This function is a no-op and does not perform any computation on the input control action. - However, it computes the approximate torques for the actuated joint since PhysX does not compute - this quantity explicitly. - - Args: - control_action: The joint action instance comprising of the desired joint positions, joint velocities - and (feed-forward) joint efforts. - joint_pos: The current joint positions of the joints in the group. Shape is (num_envs, num_joints). - joint_vel: The current joint velocities of the joints in the group. Shape is (num_envs, num_joints). - - Returns: - The computed desired joint positions, joint velocities and joint efforts. - """ - # store approximate torques for reward computation - if self.control_mode == "position": - error_pos = control_action.joint_targets - joint_pos - error_vel = torch.zeros_like(joint_vel) - joint_vel - elif self.control_mode == "velocity": - error_pos = torch.zeros_like(joint_pos) - joint_pos - error_vel = control_action.joint_targets - joint_vel - self.computed_effort = self.stiffness * error_pos + self.damping * error_vel + control_action.joint_efforts - # clip the torques based on the motor limits - self.applied_effort = self._clip_effort(self.computed_effort) - return control_action - + def __new__(cls, *args, **kwargs) -> ActuatorBase | NewtonImplicitActuator: + """Create a new instance of an implicit actuator model based on the backend.""" + return super().__new__(cls, *args, **kwargs) """ Explicit Actuator Models. """ +class IdealPDActuator(FactoryBase): + """Factory for creating ideal PD actuator models.""" -class IdealPDActuator(ActuatorBase): - r"""Ideal torque-controlled actuator model with a simple saturation model. - - It employs the following model for computing torques for the actuated joint :math:`j`: - - .. math:: - - \tau_{j, computed} = k_p * (q - q_{des}) + k_d * (\dot{q} - \dot{q}_{des}) + \tau_{ff} - - where, :math:`k_p` and :math:`k_d` are joint stiffness and damping gains, :math:`q` and :math:`\dot{q}` - are the current joint positions and velocities, :math:`q_{des}`, :math:`\dot{q}_{des}` and :math:`\tau_{ff}` - are the desired joint positions, velocities and torques commands. - - The clipping model is based on the maximum torque applied by the motor. It is implemented as: - - .. math:: - - \tau_{j, max} & = \gamma \times \tau_{motor, max} \\ - \tau_{j, applied} & = clip(\tau_{computed}, -\tau_{j, max}, \tau_{j, max}) - - where the clipping function is defined as :math:`clip(x, x_{min}, x_{max}) = min(max(x, x_{min}), x_{max})`. - The parameters :math:`\gamma` is the gear ratio of the gear box connecting the motor and the actuated joint ends, - and :math:`\tau_{motor, max}` is the maximum motor effort possible. These parameters are read from - the configuration instance passed to the class. - """ - - cfg: IdealPDActuatorCfg - """The configuration for the actuator model.""" - - """ - Operations. - """ - - def reset(self, env_ids: Sequence[int]): - pass - - def compute( - self, control_action: ArticulationActions, joint_pos: torch.Tensor, joint_vel: torch.Tensor - ) -> ArticulationActions: - # compute errors - if self.control_mode == "position": - error_pos = control_action.joint_targets - joint_pos - error_vel = torch.zeros_like(joint_vel) - joint_vel - elif self.control_mode == "velocity": - error_pos = torch.zeros_like(joint_pos) - joint_pos - error_vel = control_action.joint_targets - joint_vel - # calculate the desired joint torques - self.computed_effort = self.stiffness * error_pos + self.damping * error_vel + control_action.joint_efforts - # clip the torques based on the motor limits - self.applied_effort = self._clip_effort(self.computed_effort) - # set the computed actions back into the control action - control_action.joint_efforts = self.applied_effort - control_action.joint_targets = None - return control_action - - -class DCMotor(IdealPDActuator): - r"""Direct control (DC) motor actuator model with velocity-based saturation model. - - It uses the same model as the :class:`IdealActuator` for computing the torques from input commands. - However, it implements a saturation model defined by DC motor characteristics. - - A DC motor is a type of electric motor that is powered by direct current electricity. In most cases, - the motor is connected to a constant source of voltage supply, and the current is controlled by a rheostat. - Depending on various design factors such as windings and materials, the motor can draw a limited maximum power - from the electronic source, which limits the produced motor torque and speed. - - A DC motor characteristics are defined by the following parameters: - - * Continuous-rated speed (:math:`\dot{q}_{motor, max}`) : The maximum-rated speed of the motor. - * Continuous-stall torque (:math:`\tau_{motor, max}`): The maximum-rated torque produced at 0 speed. - * Saturation torque (:math:`\tau_{motor, sat}`): The maximum torque that can be outputted for a short period. - - Based on these parameters, the instantaneous minimum and maximum torques are defined as follows: - - .. math:: - - \tau_{j, max}(\dot{q}) & = clip \left (\tau_{j, sat} \times \left(1 - - \frac{\dot{q}}{\dot{q}_{j, max}}\right), 0.0, \tau_{j, max} \right) \\ - \tau_{j, min}(\dot{q}) & = clip \left (\tau_{j, sat} \times \left( -1 - - \frac{\dot{q}}{\dot{q}_{j, max}}\right), - \tau_{j, max}, 0.0 \right) - - where :math:`\gamma` is the gear ratio of the gear box connecting the motor and the actuated joint ends, - :math:`\dot{q}_{j, max} = \gamma^{-1} \times \dot{q}_{motor, max}`, :math:`\tau_{j, max} = - \gamma \times \tau_{motor, max}` and :math:`\tau_{j, peak} = \gamma \times \tau_{motor, peak}` - are the maximum joint velocity, maximum joint torque and peak torque, respectively. These parameters - are read from the configuration instance passed to the class. - - Using these values, the computed torques are clipped to the minimum and maximum values based on the - instantaneous joint velocity: - - .. math:: - - \tau_{j, applied} = clip(\tau_{computed}, \tau_{j, min}(\dot{q}), \tau_{j, max}(\dot{q})) - - """ - - cfg: DCMotorCfg - """The configuration for the actuator model.""" - - def __init__(self, cfg: DCMotorCfg, *args, **kwargs): - super().__init__(cfg, *args, **kwargs) - # parse configuration - if self.cfg.saturation_effort is not None: - self._saturation_effort = self.cfg.saturation_effort - else: - self._saturation_effort = torch.inf - # prepare joint vel buffer for max effort computation - self._joint_vel = torch.zeros_like(self.computed_effort) - # create buffer for zeros effort - self._zeros_effort = torch.zeros_like(self.computed_effort) - # check that quantities are provided - if self.cfg.velocity_limit is None: - raise ValueError("The velocity limit must be provided for the DC motor actuator model.") - - """ - Operations. - """ - - def compute( - self, control_action: ArticulationActions, joint_pos: torch.Tensor, joint_vel: torch.Tensor - ) -> ArticulationActions: - # save current joint vel - self._joint_vel[:] = joint_vel - # calculate the desired joint torques - return super().compute(control_action, joint_pos, joint_vel) - - """ - Helper functions. - """ - - def _clip_effort(self, effort: torch.Tensor) -> torch.Tensor: - # compute torque limits - # -- max limit - max_effort = self._saturation_effort * (1.0 - self._joint_vel / self.velocity_limit) - max_effort = torch.clip(max_effort, min=self._zeros_effort, max=self.effort_limit) - # -- min limit - min_effort = self._saturation_effort * (-1.0 - self._joint_vel / self.velocity_limit) - min_effort = torch.clip(min_effort, min=-self.effort_limit, max=self._zeros_effort) - - # clip the torques based on the motor limits - return torch.clip(effort, min=min_effort, max=max_effort) - - -class DelayedPDActuator(IdealPDActuator): - """Ideal PD actuator with delayed command application. - - This class extends the :class:`IdealPDActuator` class by adding a delay to the actuator commands. The delay - is implemented using a circular buffer that stores the actuator commands for a certain number of physics steps. - The most recent actuation value is pushed to the buffer at every physics step, but the final actuation value - applied to the simulation is lagged by a certain number of physics steps. - - The amount of time lag is configurable and can be set to a random value between the minimum and maximum time - lag bounds at every reset. The minimum and maximum time lag values are set in the configuration instance passed - to the class. - """ - - cfg: DelayedPDActuatorCfg - """The configuration for the actuator model.""" - - def __init__(self, cfg: DelayedPDActuatorCfg, *args, **kwargs): - super().__init__(cfg, *args, **kwargs) - # instantiate the delay buffers - self.joint_targets_delay_buffer = DelayBuffer(cfg.max_delay, self._num_envs, device=self._device) - self.efforts_delay_buffer = DelayBuffer(cfg.max_delay, self._num_envs, device=self._device) - # all of the envs - self._ALL_INDICES = torch.arange(self._num_envs, dtype=torch.long, device=self._device) - - def reset(self, env_ids: Sequence[int]): - super().reset(env_ids) - # number of environments (since env_ids can be a slice) - if env_ids is None or env_ids == slice(None): - num_envs = self._num_envs - else: - num_envs = len(env_ids) - # set a new random delay for environments in env_ids - time_lags = torch.randint( - low=self.cfg.min_delay, - high=self.cfg.max_delay + 1, - size=(num_envs,), - dtype=torch.int, - device=self._device, - ) - # set delays - self.joint_targets_delay_buffer.set_time_lag(time_lags, env_ids) - self.efforts_delay_buffer.set_time_lag(time_lags, env_ids) - # reset buffers - self.joint_targets_delay_buffer.reset(env_ids) - self.efforts_delay_buffer.reset(env_ids) - - def compute( - self, control_action: ArticulationActions, joint_pos: torch.Tensor, joint_vel: torch.Tensor - ) -> ArticulationActions: - # apply delay based on the delay the model for all the setpoints - control_action.joint_targets = self.joint_targets_delay_buffer.compute(control_action.joint_targets) - control_action.joint_efforts = self.efforts_delay_buffer.compute(control_action.joint_efforts) - # compte actuator model - return super().compute(control_action, joint_pos, joint_vel) - - -class RemotizedPDActuator(DelayedPDActuator): - """Ideal PD actuator with angle-dependent torque limits. - - This class extends the :class:`DelayedPDActuator` class by adding angle-dependent torque limits to the actuator. - The torque limits are applied by querying a lookup table describing the relationship between the joint angle - and the maximum output torque. The lookup table is provided in the configuration instance passed to the class. - - The torque limits are interpolated based on the current joint positions and applied to the actuator commands. - """ - - def __init__( - self, - cfg: RemotizedPDActuatorCfg, - joint_names: list[str], - joint_ids: Sequence[int], - num_envs: int, - device: str, - control_mode: str = "position", - stiffness: torch.Tensor | float = 0.0, - damping: torch.Tensor | float = 0.0, - armature: torch.Tensor | float = 0.0, - friction: torch.Tensor | float = 0.0, - effort_limit: torch.Tensor | float = torch.inf, - velocity_limit: torch.Tensor | float = torch.inf, - ): - # remove effort and velocity box constraints from the base class - cfg.effort_limit = torch.inf - cfg.velocity_limit = torch.inf - # call the base method and set default effort_limit and velocity_limit to inf - super().__init__( - cfg, - joint_names, - joint_ids, - num_envs, - device, - control_mode, - stiffness, - damping, - armature, - friction, - torch.inf, - torch.inf, - ) - self._joint_parameter_lookup = torch.tensor(cfg.joint_parameter_lookup, device=device) - # define remotized joint torque limit - self._torque_limit = LinearInterpolation(self.angle_samples, self.max_torque_samples, device=device) + def __new__(cls, *args, **kwargs) -> ActuatorBase | NewtonIdealPDActuator: + """Create a new instance of an ideal PD actuator model based on the backend.""" + return super().__new__(cls, *args, **kwargs) - """ - Properties. - """ +class DCMotor(FactoryBase): + """Factory for creating DC motor actuator models.""" - @property - def angle_samples(self) -> torch.Tensor: - return self._joint_parameter_lookup[:, 0] + def __new__(cls, *args, **kwargs) -> ActuatorBase | NewtonDCMotor: + """Create a new instance of a DC motor actuator model based on the backend.""" + return super().__new__(cls, *args, **kwargs) - @property - def transmission_ratio_samples(self) -> torch.Tensor: - return self._joint_parameter_lookup[:, 1] +class DelayedPDActuator(FactoryBase): + """Factory for creating delayed PD actuator models.""" - @property - def max_torque_samples(self) -> torch.Tensor: - return self._joint_parameter_lookup[:, 2] + def __new__(cls, *args, **kwargs) -> ActuatorBase: + """Create a new instance of a delayed PD actuator model based on the backend.""" + return super().__new__(cls, *args, **kwargs) - """ - Operations. - """ +class RemotizedPDActuator(FactoryBase): + """Factory for creating remotized PD actuator models.""" - def compute( - self, control_action: ArticulationActions, joint_pos: torch.Tensor, joint_vel: torch.Tensor - ) -> ArticulationActions: - # call the base method - control_action = super().compute(control_action, joint_pos, joint_vel) - # compute the absolute torque limits for the current joint positions - abs_torque_limits = self._torque_limit.compute(joint_pos) - # apply the limits - control_action.joint_efforts = torch.clamp( - control_action.joint_efforts, min=-abs_torque_limits, max=abs_torque_limits - ) - self.applied_effort = control_action.joint_efforts - return control_action + def __new__(cls, *args, **kwargs) -> ActuatorBase: + """Create a new instance of a remotized PD actuator model based on the backend.""" + return super().__new__(cls, *args, **kwargs) \ No newline at end of file diff --git a/source/isaaclab/isaaclab/assets/__init__.py b/source/isaaclab/isaaclab/assets/__init__.py index 57dafb20569..91ab3d5292c 100644 --- a/source/isaaclab/isaaclab/assets/__init__.py +++ b/source/isaaclab/isaaclab/assets/__init__.py @@ -39,5 +39,6 @@ """ from .articulation import Articulation, ArticulationCfg, ArticulationData +from .rigid_object import RigidObject, RigidObjectCfg, RigidObjectData from .asset_base import AssetBase from .asset_base_cfg import AssetBaseCfg diff --git a/source/isaaclab/isaaclab/assets/articulation/__init__.py b/source/isaaclab/isaaclab/assets/articulation/__init__.py index 9429f0fec31..a9f9f8435ed 100644 --- a/source/isaaclab/isaaclab/assets/articulation/__init__.py +++ b/source/isaaclab/isaaclab/assets/articulation/__init__.py @@ -7,4 +7,4 @@ from .articulation import Articulation from .articulation_cfg import ArticulationCfg -from .articulation_data import ArticulationData +from .articulation_data import ArticulationData \ No newline at end of file diff --git a/source/isaaclab/isaaclab/assets/articulation/articulation.py b/source/isaaclab/isaaclab/assets/articulation/articulation.py index ac853bbb7f8..ff4199ad53d 100644 --- a/source/isaaclab/isaaclab/assets/articulation/articulation.py +++ b/source/isaaclab/isaaclab/assets/articulation/articulation.py @@ -1,1813 +1,18 @@ -# Copyright (c) 2022-2025, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md). -# All rights reserved. -# -# SPDX-License-Identifier: BSD-3-Clause - -# Flag for pyright to ignore type errors in this file. -# pyright: reportPrivateUsage=false - from __future__ import annotations +from .base_articulation import BaseArticulation +from isaaclab.utils.backend_utils import FactoryBase -import torch -from collections.abc import Sequence -from prettytable import PrettyTable -from typing import TYPE_CHECKING, Literal - -import omni.log -import warp as wp -from isaacsim.core.simulation_manager import SimulationManager -from newton import JointMode, JointType, Model -from newton.selection import ArticulationView as NewtonArticulationView -from newton.solvers import SolverMuJoCo -from pxr import UsdPhysics - -import isaaclab.sim as sim_utils -import isaaclab.utils.math as math_utils -import isaaclab.utils.string as string_utils -from isaaclab.actuators import ActuatorBase, ActuatorBaseCfg, ImplicitActuator -from isaaclab.sim._impl.newton_manager import NewtonManager -from isaaclab.utils.types import ArticulationActions - -from ..asset_base import AssetBase -from .articulation_data import ArticulationData +from typing import TYPE_CHECKING if TYPE_CHECKING: - from .articulation_cfg import ArticulationCfg - - -class Articulation(AssetBase): - """An articulation asset class. - - An articulation is a collection of rigid bodies connected by joints. The joints can be either - fixed or actuated. The joints can be of different types, such as revolute, prismatic, D-6, etc. - However, the articulation class has currently been tested with revolute and prismatic joints. - The class supports both floating-base and fixed-base articulations. The type of articulation - is determined based on the root joint of the articulation. If the root joint is fixed, then - the articulation is considered a fixed-base system. Otherwise, it is considered a floating-base - system. This can be checked using the :attr:`Articulation.is_fixed_base` attribute. - - For an asset to be considered an articulation, the root prim of the asset must have the - `USD ArticulationRootAPI`_. This API is used to define the sub-tree of the articulation using - the reduced coordinate formulation. On playing the simulation, the physics engine parses the - articulation root prim and creates the corresponding articulation in the physics engine. The - articulation root prim can be specified using the :attr:`AssetBaseCfg.prim_path` attribute. - - The articulation class also provides the functionality to augment the simulation of an articulated - system with custom actuator models. These models can either be explicit or implicit, as detailed in - the :mod:`isaaclab.actuators` module. The actuator models are specified using the - :attr:`ArticulationCfg.actuators` attribute. These are then parsed and used to initialize the - corresponding actuator models, when the simulation is played. - - During the simulation step, the articulation class first applies the actuator models to compute - the joint commands based on the user-specified targets. These joint commands are then applied - into the simulation. The joint commands can be either position, velocity, or effort commands. - As an example, the following snippet shows how this can be used for position commands: - - .. code-block:: python - - # an example instance of the articulation class - my_articulation = Articulation(cfg) - - # set joint position targets - my_articulation.set_joint_position_target(position) - # propagate the actuator models and apply the computed commands into the simulation - my_articulation.write_data_to_sim() - - # step the simulation using the simulation context - sim_context.step() - - # update the articulation state, where dt is the simulation time step - my_articulation.update(dt) - - .. _`USD ArticulationRootAPI`: https://openusd.org/dev/api/class_usd_physics_articulation_root_a_p_i.html - - """ - - cfg: ArticulationCfg - """Configuration instance for the articulations.""" - - actuators: dict[str, ActuatorBase] - """Dictionary of actuator instances for the articulation. - - The keys are the actuator names and the values are the actuator instances. The actuator instances - are initialized based on the actuator configurations specified in the :attr:`ArticulationCfg.actuators` - attribute. They are used to compute the joint commands during the :meth:`write_data_to_sim` function. - """ - - def __init__(self, cfg: ArticulationCfg): - """Initialize the articulation. - - Args: - cfg: A configuration instance. - """ - super().__init__(cfg) - - """ - Properties - """ - - @property - def data(self) -> ArticulationData: - return self._data - - @property - def num_instances(self) -> int: - return self._root_newton_view.count - - @property - def is_fixed_base(self) -> bool: - """Whether the articulation is a fixed-base or floating-base system.""" - return self._root_newton_view.is_fixed_base - - @property - def num_joints(self) -> int: - """Number of joints in articulation.""" - return self._root_newton_view.joint_dof_count - - @property - def num_fixed_tendons(self) -> int: - """Number of fixed tendons in articulation.""" - return 0 - - @property - def num_spatial_tendons(self) -> int: - """Number of spatial tendons in articulation.""" - return 0 - - @property - def num_bodies(self) -> int: - """Number of bodies in articulation.""" - return self._root_newton_view.link_count - - @property - def joint_names(self) -> list[str]: - """Ordered names of joints in articulation.""" - return self._root_newton_view.joint_dof_names - - @property - def fixed_tendon_names(self) -> list[str]: - """Ordered names of fixed tendons in articulation.""" - # TODO: check if the articulation has fixed tendons - return [] - - @property - def spatial_tendon_names(self) -> list[str]: - """Ordered names of spatial tendons in articulation.""" - # TODO: check if the articulation has spatial tendons - return [] - - @property - def body_names(self) -> list[str]: - """Ordered names of bodies in articulation.""" - return self._root_newton_view.body_names - - @property - def root_newton_view(self) -> NewtonArticulationView: - """Articulation view for the asset (Newton). - - Note: - Use this view with caution. It requires handling of tensors in a specific way. - """ - return self._root_newton_view - - @property - def root_newton_model(self) -> Model: - """Newton model for the asset.""" - return self._root_newton_view.model - - @property - def num_shapes_per_body(self) -> list[int]: - """Number of collision shapes per body in the articulation. - - This property returns a list where each element represents the number of collision - shapes for the corresponding body in the articulation. This is cached for efficient - access during material property randomization and other operations. - - Returns: - List of integers representing the number of shapes per body. - """ - if not hasattr(self, "_num_shapes_per_body"): - self._num_shapes_per_body = [] - for shapes in self._root_newton_view.body_shapes: - self._num_shapes_per_body.append(len(shapes)) - return self._num_shapes_per_body - - """ - Operations. - """ - - def reset(self, env_ids: Sequence[int] | None = None): - # use ellipses object to skip initial indices. - if env_ids is None: - env_ids = slice(None) - # reset actuators - for actuator in self.actuators.values(): - actuator.reset(env_ids) - # reset external wrench - self._external_force_b[env_ids] = 0.0 - self._external_torque_b[env_ids] = 0.0 - - def write_data_to_sim(self): - """Write external wrenches and joint commands to the simulation. - - If any explicit actuators are present, then the actuator models are used to compute the - joint commands. Otherwise, the joint commands are directly set into the simulation. - - Note: - We write external wrench to the simulation here since this function is called before the simulation step. - This ensures that the external wrench is applied at every simulation step. - """ - # write external wrench - if self.has_external_wrench: - wrenches_b = torch.cat([self._external_force_b, self._external_torque_b], dim=-1) - self._root_newton_view.set_attribute("body_f", NewtonManager.get_state_0(), wp.from_torch(wrenches_b)) - - # apply actuator models - self._apply_actuator_model() - # write actions into simulation - self._root_newton_view.set_attribute("joint_f", NewtonManager.get_control(), self._joint_effort_target_sim) - # position and velocity targets only for implicit actuators - if self._has_implicit_actuators: - # Sets the position or velocity target for the implicit actuators depending on the actuator type. - self._root_newton_view.set_attribute("joint_target", NewtonManager.get_control(), self._joint_target_sim) - - def update(self, dt: float): - self._data.update(dt) - - """ - Operations - Finders. - """ - - def find_bodies(self, name_keys: str | Sequence[str], preserve_order: bool = False) -> tuple[list[int], list[str]]: - """Find bodies in the articulation based on the name keys. - - Please check the :meth:`isaaclab.utils.string_utils.resolve_matching_names` function for more - information on the name matching. - - Args: - name_keys: A regular expression or a list of regular expressions to match the body names. - preserve_order: Whether to preserve the order of the name keys in the output. Defaults to False. - - Returns: - A tuple of lists containing the body indices and names. - """ - return string_utils.resolve_matching_names(name_keys, self.body_names, preserve_order) - - def find_joints( - self, name_keys: str | Sequence[str], joint_subset: list[str] | None = None, preserve_order: bool = False - ) -> tuple[list[int], list[str]]: - """Find joints in the articulation based on the name keys. - - Please see the :func:`isaaclab.utils.string.resolve_matching_names` function for more information - on the name matching. - - Args: - name_keys: A regular expression or a list of regular expressions to match the joint names. - joint_subset: A subset of joints to search for. Defaults to None, which means all joints - in the articulation are searched. - preserve_order: Whether to preserve the order of the name keys in the output. Defaults to False. - - Returns: - A tuple of lists containing the joint indices and names. - """ - if joint_subset is None: - joint_subset = self.joint_names - # find joints - return string_utils.resolve_matching_names(name_keys, joint_subset, preserve_order) - - def find_fixed_tendons( - self, name_keys: str | Sequence[str], tendon_subsets: list[str] | None = None, preserve_order: bool = False - ) -> tuple[list[int], list[str]]: - """Find fixed tendons in the articulation based on the name keys. - - Please see the :func:`isaaclab.utils.string.resolve_matching_names` function for more information - on the name matching. - - Args: - name_keys: A regular expression or a list of regular expressions to match the joint - names with fixed tendons. - tendon_subsets: A subset of joints with fixed tendons to search for. Defaults to None, which means - all joints in the articulation are searched. - preserve_order: Whether to preserve the order of the name keys in the output. Defaults to False. - - Returns: - A tuple of lists containing the tendon indices and names. - """ - if tendon_subsets is None: - # tendons follow the joint names they are attached to - tendon_subsets = self.fixed_tendon_names - # find tendons - return string_utils.resolve_matching_names(name_keys, tendon_subsets, preserve_order) - - def find_spatial_tendons( - self, name_keys: str | Sequence[str], tendon_subsets: list[str] | None = None, preserve_order: bool = False - ) -> tuple[list[int], list[str]]: - """Find spatial tendons in the articulation based on the name keys. - - Please see the :func:`isaaclab.utils.string.resolve_matching_names` function for more information - on the name matching. - - Args: - name_keys: A regular expression or a list of regular expressions to match the tendon names. - tendon_subsets: A subset of tendons to search for. Defaults to None, which means all tendons - in the articulation are searched. - preserve_order: Whether to preserve the order of the name keys in the output. Defaults to False. - - Returns: - A tuple of lists containing the tendon indices and names. - """ - if tendon_subsets is None: - tendon_subsets = self.spatial_tendon_names - # find tendons - return string_utils.resolve_matching_names(name_keys, tendon_subsets, preserve_order) - - """ - Operations - State Writers. - """ - - def write_root_state_to_sim(self, root_state: torch.Tensor, env_ids: Sequence[int] | None = None): - """Set the root state over selected environment indices into the simulation. - - The root state comprises of the cartesian position, quaternion orientation in (w, x, y, z), and linear - and angular velocity. All the quantities are in the simulation frame. - - Args: - root_state: Root state in simulation frame. Shape is (len(env_ids), 13). - env_ids: Environment indices. If None, then all indices are used. - """ - - # set into simulation - self.write_root_link_pose_to_sim(root_state[:, :7], env_ids=env_ids) - self.write_root_com_velocity_to_sim(root_state[:, 7:], env_ids=env_ids) - - def write_root_com_state_to_sim(self, root_state: torch.Tensor, env_ids: Sequence[int] | None = None): - """Set the root center of mass state over selected environment indices into the simulation. - - The root state comprises of the cartesian position, quaternion orientation in (w, x, y, z), and linear - and angular velocity. All the quantities are in the simulation frame. - - Args: - root_state: Root state in simulation frame. Shape is (len(env_ids), 13). - env_ids: Environment indices. If None, then all indices are used. - """ - self.write_root_com_pose_to_sim(root_state[:, :7], env_ids=env_ids) - self.write_root_com_velocity_to_sim(root_state[:, 7:], env_ids=env_ids) - - def write_root_link_state_to_sim(self, root_state: torch.Tensor, env_ids: Sequence[int] | None = None): - """Set the root link state over selected environment indices into the simulation. - - The root state comprises of the cartesian position, quaternion orientation in (w, x, y, z), and linear - and angular velocity. All the quantities are in the simulation frame. - - Args: - root_state: Root state in simulation frame. Shape is (len(env_ids), 13). - env_ids: Environment indices. If None, then all indices are used. - """ - self.write_root_link_pose_to_sim(root_state[:, :7], env_ids=env_ids) - self.write_root_link_velocity_to_sim(root_state[:, 7:], env_ids=env_ids) - - def write_root_pose_to_sim(self, root_pose: torch.Tensor, env_ids: Sequence[int] | None = None): - """Set the root pose over selected environment indices into the simulation. - - The root pose comprises of the cartesian position and quaternion orientation in (w, x, y, z). - - Args: - root_pose: Root poses in simulation frame. Shape is (len(env_ids), 7). - env_ids: Environment indices. If None, then all indices are used. - """ - self.write_root_link_pose_to_sim(root_pose, env_ids=env_ids) - - def write_root_link_pose_to_sim(self, root_pose: torch.Tensor, env_ids: Sequence[int] | None = None): - """Set the root link pose over selected environment indices into the simulation. - - The root pose comprises of the cartesian position and quaternion orientation in (w, x, y, z). - - Args: - root_pose: Root poses in simulation frame. Shape is (len(env_ids), 7). - env_ids: Environment indices. If None, then all indices are used. - """ - # resolve all indices - physx_env_ids = env_ids - if env_ids is None: - env_ids = slice(None) - physx_env_ids = self._ALL_INDICES - - # note: we need to do this here since tensors are not set into simulation until step. - # set into internal buffers - self._data.root_link_pose_w[env_ids] = root_pose.clone() - # update these buffers only if the user is using them. Otherwise this adds to overhead. - if self._data._root_link_state_w.data is not None: - self._data.root_link_state_w[env_ids, :7] = self._data.root_link_pose_w[env_ids] - if self._data._root_state_w.data is not None: - self._data.root_state_w[env_ids, :7] = self._data.root_link_pose_w[env_ids] - - # convert root quaternion from wxyz to xyzw - root_poses_xyzw = self._data.root_link_pose_w.clone() - root_poses_xyzw[:, 3:] = math_utils.convert_quat(root_poses_xyzw[:, 3:], to="xyzw") - - # Need to invalidate the buffer to trigger the update with the new state. - self._data._body_link_pose_w.timestamp = -1.0 - self._data._body_com_pose_w.timestamp = -1.0 - self._data._body_state_w.timestamp = -1.0 - self._data._body_link_state_w.timestamp = -1.0 - self._data._body_com_state_w.timestamp = -1.0 - - # set into simulation - self._mask.fill_(False) - self._mask[physx_env_ids] = True - self._root_newton_view.set_root_transforms(NewtonManager.get_state_0(), root_poses_xyzw, mask=self._mask) - - def write_root_com_pose_to_sim(self, root_pose: torch.Tensor, env_ids: Sequence[int] | None = None): - """Set the root center of mass pose over selected environment indices into the simulation. - - The root pose comprises of the cartesian position and quaternion orientation in (w, x, y, z). - The orientation is the orientation of the principle axes of inertia. - - Args: - root_pose: Root center of mass poses in simulation frame. Shape is (len(env_ids), 7). - env_ids: Environment indices. If None, then all indices are used. - """ - # resolve all indices - if env_ids is None: - local_env_ids = slice(env_ids) - else: - local_env_ids = env_ids - - # set into internal buffers - self._data.root_com_pose_w[local_env_ids] = root_pose.clone() - # update these buffers only if the user is using them. Otherwise this adds to overhead. - if self._data._root_com_state_w.data is not None: - self._data.root_com_state_w[local_env_ids, :7] = self._data.root_com_pose_w[local_env_ids] - - # get CoM pose in link frame - com_pos_b = self.data.body_com_pos_b[local_env_ids, 0, :] - com_quat_b = self.data.body_com_quat_b[local_env_ids, 0, :] - # transform input CoM pose to link frame - root_link_pos, root_link_quat = math_utils.combine_frame_transforms( - root_pose[..., :3], - root_pose[..., 3:7], - math_utils.quat_apply(math_utils.quat_inv(com_quat_b), -com_pos_b), - math_utils.quat_inv(com_quat_b), - ) - root_link_pose = torch.cat((root_link_pos, root_link_quat), dim=-1) - - # write transformed pose in link frame to sim - self.write_root_link_pose_to_sim(root_pose=root_link_pose, env_ids=env_ids) - - def write_root_velocity_to_sim(self, root_velocity: torch.Tensor, env_ids: Sequence[int] | None = None): - """Set the root center of mass velocity over selected environment indices into the simulation. - - The velocity comprises linear velocity (x, y, z) and angular velocity (x, y, z) in that order. - NOTE: This sets the velocity of the root's center of mass rather than the roots frame. - - Args: - root_velocity: Root center of mass velocities in simulation world frame. Shape is (len(env_ids), 6). - env_ids: Environment indices. If None, then all indices are used. - """ - self.write_root_com_velocity_to_sim(root_velocity=root_velocity, env_ids=env_ids) - - def write_root_com_velocity_to_sim(self, root_velocity: torch.Tensor, env_ids: Sequence[int] | None = None): - """Set the root center of mass velocity over selected environment indices into the simulation. - - The velocity comprises linear velocity (x, y, z) and angular velocity (x, y, z) in that order. - NOTE: This sets the velocity of the root's center of mass rather than the roots frame. - - Args: - root_velocity: Root center of mass velocities in simulation world frame. Shape is (len(env_ids), 6). - env_ids: Environment indices. If None, then all indices are used. - """ - # resolve all indices - physx_env_ids = env_ids - if env_ids is None: - env_ids = slice(None) - physx_env_ids = self._ALL_INDICES - - # note: we need to do this here since tensors are not set into simulation until step. - # set into internal buffers - self._data.root_com_vel_w[env_ids] = root_velocity.clone() - # update these buffers only if the user is using them. Otherwise this adds to overhead. - if self._data._root_com_state_w.data is not None: - self._data.root_com_state_w[env_ids, 7:] = self._data.root_com_vel_w[env_ids] - if self._data._root_state_w.data is not None: - self._data.root_state_w[env_ids, 7:] = self._data.root_com_vel_w[env_ids] - # make the acceleration zero to prevent reporting old values - self._data.body_acc_w[env_ids] = 0.0 - - # set into simulation - self._mask.fill_(False) - self._mask[physx_env_ids] = True - self._root_newton_view.set_root_velocities( - NewtonManager.get_state_0(), self._data.root_state_w[:, 7:].clone(), mask=self._mask - ) - - def write_root_link_velocity_to_sim(self, root_velocity: torch.Tensor, env_ids: Sequence[int] | None = None): - """Set the root link velocity over selected environment indices into the simulation. - - The velocity comprises linear velocity (x, y, z) and angular velocity (x, y, z) in that order. - NOTE: This sets the velocity of the root's frame rather than the roots center of mass. - - Args: - root_velocity: Root frame velocities in simulation world frame. Shape is (len(env_ids), 6). - env_ids: Environment indices. If None, then all indices are used. - """ - # resolve all indices - if env_ids is None: - local_env_ids = slice(env_ids) - else: - local_env_ids = env_ids - - # set into internal buffers - self._data.root_link_vel_w[local_env_ids] = root_velocity.clone() - # update these buffers only if the user is using them. Otherwise this adds to overhead. - if self._data._root_link_state_w.data is not None: - self._data.root_link_state_w[local_env_ids, 7:] = self._data.root_link_vel_w[local_env_ids] - - # get CoM pose in link frame - quat = self.data.root_link_quat_w[local_env_ids] - com_pos_b = self.data.body_com_pos_b[local_env_ids, 0, :] - # transform input velocity to center of mass frame - root_com_velocity = root_velocity.clone() - root_com_velocity[:, :3] += torch.linalg.cross( - root_com_velocity[:, 3:], math_utils.quat_apply(quat, com_pos_b), dim=-1 - ) - - # write transformed velocity in CoM frame to sim - self.write_root_com_velocity_to_sim(root_velocity=root_com_velocity, env_ids=env_ids) - - def write_joint_state_to_sim( - self, - position: torch.Tensor, - velocity: torch.Tensor, - joint_ids: Sequence[int] | slice | None = None, - env_ids: Sequence[int] | slice | None = None, - ): - """Write joint positions and velocities to the simulation. - - Args: - position: Joint positions. Shape is (len(env_ids), len(joint_ids)). - velocity: Joint velocities. Shape is (len(env_ids), len(joint_ids)). - joint_ids: The joint indices to set the targets for. Defaults to None (all joints). - env_ids: The environment indices to set the targets for. Defaults to None (all environments). - """ - # set into simulation - self.write_joint_position_to_sim(position, joint_ids=joint_ids, env_ids=env_ids) - self.write_joint_velocity_to_sim(velocity, joint_ids=joint_ids, env_ids=env_ids) - - def write_joint_position_to_sim( - self, - position: torch.Tensor, - joint_ids: Sequence[int] | slice | None = None, - env_ids: Sequence[int] | slice | None = None, - ): - """Write joint positions to the simulation. - - Args: - position: Joint positions. Shape is (len(env_ids), len(joint_ids)). - joint_ids: The joint indices to set the targets for. Defaults to None (all joints). - env_ids: The environment indices to set the targets for. Defaults to None (all environments). - """ - # resolve indices - physx_env_ids = env_ids - if env_ids is None: - env_ids = slice(None) - physx_env_ids = self._ALL_INDICES - if joint_ids is None: - joint_ids = slice(None) - # broadcast env_ids if needed to allow double indexing - if env_ids != slice(None) and joint_ids != slice(None): - env_ids = env_ids[:, None] - # set into internal buffers - self._data.joint_pos[env_ids, joint_ids] = position - # Need to invalidate the buffer to trigger the update with the new root pose. - self._data._body_com_vel_w.timestamp = -1.0 - self._data._body_link_vel_w.timestamp = -1.0 - self._data._body_com_pose_b.timestamp = -1.0 - self._data._body_com_pose_w.timestamp = -1.0 - self._data._body_link_pose_w.timestamp = -1.0 - - self._data._body_state_w.timestamp = -1.0 - self._data._body_link_state_w.timestamp = -1.0 - self._data._body_com_state_w.timestamp = -1.0 - # set into simulation - self._mask.fill_(False) - self._mask[physx_env_ids] = True - self._root_newton_view.set_dof_positions( - NewtonManager.get_state_0(), self._data.joint_pos.clone(), mask=self._mask - ) - - def write_joint_velocity_to_sim( - self, - velocity: torch.Tensor, - joint_ids: Sequence[int] | slice | None = None, - env_ids: Sequence[int] | slice | None = None, - ): - """Write joint velocities to the simulation. - - Args: - velocity: Joint velocities. Shape is (len(env_ids), len(joint_ids)). - joint_ids: The joint indices to set the targets for. Defaults to None (all joints). - env_ids: The environment indices to set the targets for. Defaults to None (all environments). - """ - # resolve indices - physx_env_ids = env_ids - if env_ids is None: - env_ids = slice(None) - physx_env_ids = self._ALL_INDICES - if joint_ids is None: - joint_ids = slice(None) - # broadcast env_ids if needed to allow double indexing - if env_ids != slice(None) and joint_ids != slice(None): - env_ids = env_ids[:, None] - # set into internal buffers - self._data.joint_vel[env_ids, joint_ids] = velocity - self._data._previous_joint_vel[env_ids, joint_ids] = velocity - self._data.joint_acc[env_ids, joint_ids] = 0.0 - # set into simulation - self._mask.fill_(False) - self._mask[physx_env_ids] = True - self._root_newton_view.set_dof_velocities( - NewtonManager.get_state_0(), self._data.joint_vel.clone(), mask=self._mask - ) - - """ - Operations - Simulation Parameters Writers. - """ - - def write_joint_control_mode_to_sim( - self, - control_mode: Literal["position", "velocity", "none"] | None, - joint_ids: Sequence[int] | slice | None = None, - env_ids: Sequence[int] | None = None, - ): - """Write joint control mode into the simulation. - - Args: - control_mode: Joint control mode. Shape is (len(env_ids), len(joint_ids)). - joint_ids: The joint indices to set the control mode for. Defaults to None (all joints). - env_ids: The environment indices to set the control mode for. Defaults to None (all environments). - - Raises: - ValueError: If the control mode is invalid. - """ - # resolve indices - physx_env_ids = env_ids - if env_ids is None: - env_ids = slice(None) - physx_env_ids = self._ALL_INDICES - if joint_ids is None: - joint_ids = slice(None) - # broadcast env_ids if needed to allow double indexing - if env_ids != slice(None) and joint_ids != slice(None): - env_ids = env_ids[:, None] - # set into internal buffers - if control_mode == "position": - self._data.joint_control_mode[env_ids, joint_ids] = JointMode.TARGET_POSITION - elif control_mode == "velocity": - self._data.joint_control_mode[env_ids, joint_ids] = JointMode.TARGET_VELOCITY - elif (control_mode is None) or (control_mode == "none"): - # Set the control mode to None when using explicit actuators - self._data.joint_control_mode[env_ids, joint_ids] = JointMode.NONE - else: - raise ValueError(f"Invalid control mode: {control_mode}") - - # set into simulation - self._mask.fill_(False) - self._mask[physx_env_ids] = True - self._root_newton_view.set_attribute( - "joint_dof_mode", NewtonManager.get_model(), self._data.joint_control_mode, mask=self._mask - ) - - def write_joint_stiffness_to_sim( - self, - stiffness: torch.Tensor | float, - joint_ids: Sequence[int] | slice | None = None, - env_ids: Sequence[int] | None = None, - ): - """Write joint stiffness into the simulation. - - Args: - stiffness: Joint stiffness. Shape is (len(env_ids), len(joint_ids)). - joint_ids: The joint indices to set the stiffness for. Defaults to None (all joints). - env_ids: The environment indices to set the stiffness for. Defaults to None (all environments). - """ - # note: This function isn't setting the values for actuator models. (#128) - # resolve indices - physx_env_ids = env_ids - if env_ids is None: - env_ids = slice(None) - physx_env_ids = self._ALL_INDICES - if joint_ids is None: - joint_ids = slice(None) - # broadcast env_ids if needed to allow double indexing - if env_ids != slice(None) and joint_ids != slice(None): - env_ids = env_ids[:, None] - # set into internal buffers - self._data.joint_stiffness[env_ids, joint_ids] = stiffness - # set into simulation - self._mask.fill_(False) - self._mask[physx_env_ids] = True - self._root_newton_view.set_attribute( - "joint_target_ke", NewtonManager.get_model(), self._data.joint_stiffness, mask=self._mask - ) - - def write_joint_damping_to_sim( - self, - damping: torch.Tensor | float, - joint_ids: Sequence[int] | slice | None = None, - env_ids: Sequence[int] | None = None, - ): - """Write joint damping into the simulation. - - Args: - damping: Joint damping. Shape is (len(env_ids), len(joint_ids)). - joint_ids: The joint indices to set the damping for. Defaults to None (all joints). - env_ids: The environment indices to set the damping for. Defaults to None (all environments). - """ - # note: This function isn't setting the values for actuator models. (#128) - # resolve indices - physx_env_ids = env_ids - if env_ids is None: - env_ids = slice(None) - physx_env_ids = self._ALL_INDICES - if joint_ids is None: - joint_ids = slice(None) - # broadcast env_ids if needed to allow double indexing - if env_ids != slice(None) and joint_ids != slice(None): - env_ids = env_ids[:, None] - # set into internal buffers - self._data.joint_damping[env_ids, joint_ids] = damping - # set into simulation - self._mask.fill_(False) - self._mask[physx_env_ids] = True - self._root_newton_view.set_attribute( - "joint_target_kd", NewtonManager.get_model(), self._data.joint_damping, mask=self._mask - ) - - def write_joint_position_limit_to_sim( - self, - limits: torch.Tensor | float, - joint_ids: Sequence[int] | slice | None = None, - env_ids: Sequence[int] | None = None, - warn_limit_violation: bool = True, - ): - """Write joint position limits into the simulation. - - Args: - limits: Joint limits. Shape is (len(env_ids), len(joint_ids), 2). - joint_ids: The joint indices to set the limits for. Defaults to None (all joints). - env_ids: The environment indices to set the limits for. Defaults to None (all environments). - warn_limit_violation: Whether to use warning or info level logging when default joint positions - exceed the new limits. Defaults to True. - """ - # note: This function isn't setting the values for actuator models. (#128) - # resolve indices - physx_env_ids = env_ids - if env_ids is None: - env_ids = slice(None) - physx_env_ids = self._ALL_INDICES - if joint_ids is None: - joint_ids = slice(None) - # broadcast env_ids if needed to allow double indexing - if env_ids != slice(None) and joint_ids != slice(None): - env_ids = env_ids[:, None] - # set into internal buffers - self._data.joint_pos_limits[env_ids, joint_ids] = limits - # update default joint pos to stay within the new limits - if torch.any( - (self._data.default_joint_pos[env_ids, joint_ids] < limits[..., 0]) - | (self._data.default_joint_pos[env_ids, joint_ids] > limits[..., 1]) - ): - self._data.default_joint_pos[env_ids, joint_ids] = torch.clamp( - self._data.default_joint_pos[env_ids, joint_ids], limits[..., 0], limits[..., 1] - ) - violation_message = ( - "Some default joint positions are outside of the range of the new joint limits. Default joint positions" - " will be clamped to be within the new joint limits." - ) - if warn_limit_violation: - # warn level will show in console - omni.log.warn(violation_message) - else: - # info level is only written to log file - omni.log.info(violation_message) - # set into simulation - self._mask.fill_(False) - self._mask[physx_env_ids] = True - self._root_newton_view.set_attribute( - "joint_limit_lower", NewtonManager.get_model(), self._data.joint_pos_limits[..., 0], mask=self._mask - ) - self._root_newton_view.set_attribute( - "joint_limit_upper", NewtonManager.get_model(), self._data.joint_pos_limits[..., 1], mask=self._mask - ) - - # compute the soft limits based on the joint limits - # TODO: Optimize this computation for only selected joints - # soft joint position limits (recommended not to be too close to limits). - joint_pos_mean = (self._data.joint_pos_limits[..., 0] + self._data.joint_pos_limits[..., 1]) / 2 - joint_pos_range = self._data.joint_pos_limits[..., 1] - self._data.joint_pos_limits[..., 0] - soft_limit_factor = self.cfg.soft_joint_pos_limit_factor - # add to data - self._data.soft_joint_pos_limits[..., 0] = joint_pos_mean - 0.5 * joint_pos_range * soft_limit_factor - self._data.soft_joint_pos_limits[..., 1] = joint_pos_mean + 0.5 * joint_pos_range * soft_limit_factor - - def write_joint_velocity_limit_to_sim( - self, - limits: torch.Tensor | float, - joint_ids: Sequence[int] | slice | None = None, - env_ids: Sequence[int] | None = None, - ): - """Write joint max velocity to the simulation. - - The velocity limit is used to constrain the joint velocities in the physics engine. The joint will only - be able to reach this velocity if the joint's effort limit is sufficiently large. If the joint is moving - faster than this velocity, the physics engine will actually try to brake the joint to reach this velocity. - - .. warn:: This function is ignored when using the Mujoco solver. - - Args: - limits: Joint max velocity. Shape is (len(env_ids), len(joint_ids)). - joint_ids: The joint indices to set the max velocity for. Defaults to None (all joints). - env_ids: The environment indices to set the max velocity for. Defaults to None (all environments). - """ - # Warn if using Mujoco solver - if isinstance(NewtonManager._solver, SolverMuJoCo): - omni.log.warn("write_joint_velocity_limit_to_sim is ignored when using the Mujoco solver.") - - # resolve indices - physx_env_ids = env_ids - if env_ids is None: - env_ids = slice(None) - physx_env_ids = self._ALL_INDICES - if joint_ids is None: - joint_ids = slice(None) - # broadcast env_ids if needed to allow double indexing - if env_ids != slice(None) and joint_ids != slice(None): - env_ids = env_ids[:, None] - # move tensor to cpu if needed - if isinstance(limits, torch.Tensor): - limits = limits.to(self.device) - # set into internal buffers - self._data.joint_vel_limits[env_ids, joint_ids] = limits - # set into simulation - self._mask.fill_(False) - self._mask[physx_env_ids] = True - self._root_newton_view.set_attribute( - "joint_velocity_limit", NewtonManager.get_model(), self._data.joint_vel_limits, mask=self._mask - ) - - def write_joint_effort_limit_to_sim( - self, - limits: torch.Tensor | float, - joint_ids: Sequence[int] | slice | None = None, - env_ids: Sequence[int] | None = None, - ): - """Write joint effort limits into the simulation. - - The effort limit is used to constrain the computed joint efforts in the physics engine. If the - computed effort exceeds this limit, the physics engine will clip the effort to this value. - - Args: - limits: Joint torque limits. Shape is (len(env_ids), len(joint_ids)). - 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). - """ - # note: This function isn't setting the values for actuator models. (#128) - # resolve indices - physx_env_ids = env_ids - if env_ids is None: - env_ids = slice(None) - physx_env_ids = self._ALL_INDICES - if joint_ids is None: - joint_ids = slice(None) - # broadcast env_ids if needed to allow double indexing - if env_ids != slice(None) and joint_ids != slice(None): - env_ids = env_ids[:, None] - # move tensor to cpu if needed - if isinstance(limits, torch.Tensor): - limits = limits.to(self.device) - # set into internal buffers - self._data.joint_effort_limits[env_ids, joint_ids] = limits - # set into simulation - self._mask.fill_(False) - self._mask[physx_env_ids] = True - self._root_newton_view.set_attribute( - "joint_effort_limit", NewtonManager.get_model(), self._data.joint_effort_limits, mask=self._mask - ) - - def write_joint_armature_to_sim( - self, - armature: torch.Tensor | float, - joint_ids: Sequence[int] | slice | None = None, - env_ids: Sequence[int] | None = None, - ): - """Write joint armature into the simulation. - - The armature is directly added to the corresponding joint-space inertia. It helps improve the - simulation stability by reducing the joint velocities. - - Args: - armature: Joint armature. Shape is (len(env_ids), len(joint_ids)). - 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). - """ - # resolve indices - physx_env_ids = env_ids - if env_ids is None: - env_ids = slice(None) - physx_env_ids = self._ALL_INDICES - if joint_ids is None: - joint_ids = slice(None) - # broadcast env_ids if needed to allow double indexing - if env_ids != slice(None) and joint_ids != slice(None): - env_ids = env_ids[:, None] - # set into internal buffers - self._data.joint_armature[env_ids, joint_ids] = armature - # set into simulation: Only used by the Featherstone solver - self._mask.fill_(False) - self._mask[physx_env_ids] = True - self._root_newton_view.set_attribute( - "joint_armature", NewtonManager.get_model(), self._data.joint_armature, mask=self._mask - ) - - def write_joint_friction_coefficient_to_sim( - self, - joint_friction_coeff: torch.Tensor | float, - joint_ids: Sequence[int] | slice | None = None, - env_ids: Sequence[int] | None = None, - ): - r"""Write joint friction coefficients into the simulation. - - The joint friction is a unitless quantity. It relates the magnitude of the spatial force transmitted - from the parent body to the child body to the maximal friction force that may be applied by the solver - to resist the joint motion. - - Mathematically, this means that: :math:`F_{resist} \leq \mu F_{spatial}`, where :math:`F_{resist}` - is the resisting force applied by the solver and :math:`F_{spatial}` is the spatial force - transmitted from the parent body to the child body. The simulated friction effect is therefore - similar to static and Coulomb friction. - - Args: - joint_friction: Joint friction. Shape is (len(env_ids), len(joint_ids)). - 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). - """ - # resolve indices - physx_env_ids = env_ids - if env_ids is None: - env_ids = slice(None) - physx_env_ids = self._ALL_INDICES - if joint_ids is None: - joint_ids = slice(None) - # broadcast env_ids if needed to allow double indexing - if env_ids != slice(None) and joint_ids != slice(None): - env_ids = env_ids[:, None] - # set into internal buffers - self._data.joint_friction_coeff[env_ids, joint_ids] = joint_friction_coeff - # set into simulation - self._mask.fill_(False) - self._mask[physx_env_ids] = True - self._root_newton_view.set_attribute( - "joint_friction", NewtonManager.get_model(), self._data.joint_friction_coeff, mask=self._mask - ) - - """ - Operations - Setters. - """ - - def set_external_force_and_torque( - self, - forces: torch.Tensor, - torques: torch.Tensor, - body_ids: Sequence[int] | slice | None = None, - env_ids: Sequence[int] | None = None, - ): - """Set external force and torque to apply on the asset's bodies in their local frame. - - For many applications, we want to keep the applied external force on rigid bodies constant over a period of - time (for instance, during the policy control). This function allows us to store the external force and torque - into buffers which are then applied to the simulation at every step. - - .. caution:: - If the function is called with empty forces and torques, then this function disables the application - of external wrench to the simulation. - - .. code-block:: python - - # example of disabling external wrench - asset.set_external_force_and_torque(forces=torch.zeros(0, 3), torques=torch.zeros(0, 3)) - - .. note:: - This function does not apply the external wrench to the simulation. It only fills the buffers with - the desired values. To apply the external wrench, call the :meth:`write_data_to_sim` function - right before the simulation step. - - Args: - forces: External forces in bodies' local frame. Shape is (len(env_ids), len(body_ids), 3). - torques: External torques in bodies' local frame. Shape is (len(env_ids), len(body_ids), 3). - body_ids: Body indices to apply external wrench to. Defaults to None (all bodies). - env_ids: Environment indices to apply external wrench to. Defaults to None (all instances). - """ - if forces.any() or torques.any(): - self.has_external_wrench = True - else: - self.has_external_wrench = False - - # resolve all indices - # -- env_ids - if env_ids is None: - env_ids = self._ALL_INDICES - elif not isinstance(env_ids, torch.Tensor): - env_ids = torch.tensor(env_ids, dtype=torch.long, device=self.device) - # -- body_ids - if body_ids is None: - body_ids = torch.arange(self.num_bodies, dtype=torch.long, device=self.device) - elif isinstance(body_ids, slice): - body_ids = torch.arange(self.num_bodies, dtype=torch.long, device=self.device)[body_ids] - elif not isinstance(body_ids, torch.Tensor): - body_ids = torch.tensor(body_ids, dtype=torch.long, device=self.device) - - # note: we need to do this complicated indexing since torch doesn't support multi-indexing - # create global body indices from env_ids and env_body_ids - # (env_id * total_bodies_per_env) + body_id - indices = body_ids.repeat(len(env_ids), 1) + env_ids.unsqueeze(1) * self.num_bodies - indices = indices.view(-1) - # set into internal buffers - # note: these are applied in the write_to_sim function - self._external_force_b.flatten(0, 1)[indices] = forces.flatten(0, 1) - self._external_torque_b.flatten(0, 1)[indices] = torques.flatten(0, 1) - - def set_joint_position_target( - self, target: torch.Tensor, joint_ids: Sequence[int] | slice | None = None, env_ids: Sequence[int] | None = None - ): - """Set joint position targets into internal buffers. - - This function does not apply the joint targets to the simulation. It only fills the buffers with - the desired values. To apply the joint targets, call the :meth:`write_data_to_sim` function. - - Args: - target: Joint position targets. Shape is (len(env_ids), len(joint_ids)). - joint_ids: The joint indices to set the targets for. Defaults to None (all joints). - env_ids: The environment indices to set the targets for. Defaults to None (all environments). - """ - # resolve indices - if env_ids is None: - env_ids = slice(None) - if joint_ids is None: - joint_ids = slice(None) - # broadcast env_ids if needed to allow double indexing - if env_ids != slice(None) and joint_ids != slice(None): - env_ids = env_ids[:, None] - # set targets - self._data.joint_target[env_ids, joint_ids] = target - - def set_joint_velocity_target( - self, target: torch.Tensor, joint_ids: Sequence[int] | slice | None = None, env_ids: Sequence[int] | None = None - ): - """Set joint velocity targets into internal buffers. - - This function does not apply the joint targets to the simulation. It only fills the buffers with - the desired values. To apply the joint targets, call the :meth:`write_data_to_sim` function. - - Args: - target: Joint velocity targets. Shape is (len(env_ids), len(joint_ids)). - joint_ids: The joint indices to set the targets for. Defaults to None (all joints). - env_ids: The environment indices to set the targets for. Defaults to None (all environments). - """ - # resolve indices - if env_ids is None: - env_ids = slice(None) - if joint_ids is None: - joint_ids = slice(None) - # broadcast env_ids if needed to allow double indexing - if env_ids != slice(None) and joint_ids != slice(None): - env_ids = env_ids[:, None] - # set targets - self._data.joint_target[env_ids, joint_ids] = target - - def set_joint_effort_target( - self, target: torch.Tensor, joint_ids: Sequence[int] | slice | None = None, env_ids: Sequence[int] | None = None - ): - """Set joint efforts into internal buffers. - - This function does not apply the joint targets to the simulation. It only fills the buffers with - the desired values. To apply the joint targets, call the :meth:`write_data_to_sim` function. - - Args: - target: Joint effort targets. Shape is (len(env_ids), len(joint_ids)). - joint_ids: The joint indices to set the targets for. Defaults to None (all joints). - env_ids: The environment indices to set the targets for. Defaults to None (all environments). - """ - # resolve indices - if env_ids is None: - env_ids = slice(None) - if joint_ids is None: - joint_ids = slice(None) - # broadcast env_ids if needed to allow double indexing - if env_ids != slice(None) and joint_ids != slice(None): - env_ids = env_ids[:, None] - # set targets - self._data.joint_effort_target[env_ids, joint_ids] = target - - """ - Operations - Tendons. - """ - - def set_fixed_tendon_stiffness( - self, - stiffness: torch.Tensor, - fixed_tendon_ids: Sequence[int] | slice | None = None, - env_ids: Sequence[int] | None = None, - ): - """Set fixed tendon stiffness into internal buffers. - - This function does not apply the tendon stiffness to the simulation. It only fills the buffers with - the desired values. To apply the tendon stiffness, call the :meth:`write_fixed_tendon_properties_to_sim` function. - - Args: - stiffness: Fixed tendon stiffness. Shape is (len(env_ids), len(fixed_tendon_ids)). - fixed_tendon_ids: The tendon indices to set the stiffness for. Defaults to None (all fixed tendons). - env_ids: The environment indices to set the stiffness for. Defaults to None (all environments). - """ - raise NotImplementedError("Fixed tendon stiffness is not supported in Newton.") - - def set_fixed_tendon_damping( - self, - damping: torch.Tensor, - fixed_tendon_ids: Sequence[int] | slice | None = None, - env_ids: Sequence[int] | None = None, - ): - """Set fixed tendon damping into internal buffers. - - This function does not apply the tendon damping to the simulation. It only fills the buffers with - the desired values. To apply the tendon damping, call the :meth:`write_fixed_tendon_properties_to_sim` function. - - Args: - damping: Fixed tendon damping. Shape is (len(env_ids), len(fixed_tendon_ids)). - fixed_tendon_ids: The tendon indices to set the damping for. Defaults to None (all fixed tendons). - env_ids: The environment indices to set the damping for. Defaults to None (all environments). - """ - raise NotImplementedError("Fixed tendon damping is not supported in Newton.") - - def set_fixed_tendon_limit_stiffness( - self, - limit_stiffness: torch.Tensor, - fixed_tendon_ids: Sequence[int] | slice | None = None, - env_ids: Sequence[int] | None = None, - ): - """Set fixed tendon limit stiffness efforts into internal buffers. - - This function does not apply the tendon limit stiffness to the simulation. It only fills the buffers with - the desired values. To apply the tendon limit stiffness, call the :meth:`write_fixed_tendon_properties_to_sim` function. - - Args: - limit_stiffness: Fixed tendon limit stiffness. Shape is (len(env_ids), len(fixed_tendon_ids)). - fixed_tendon_ids: The tendon indices to set the limit stiffness for. Defaults to None (all fixed tendons). - env_ids: The environment indices to set the limit stiffness for. Defaults to None (all environments). - """ - raise NotImplementedError("Fixed tendon limit stiffness is not supported in Newton.") - - def set_fixed_tendon_position_limit( - self, - limit: torch.Tensor, - fixed_tendon_ids: Sequence[int] | slice | None = None, - env_ids: Sequence[int] | None = None, - ): - """Set fixed tendon limit efforts into internal buffers. - - This function does not apply the tendon limit to the simulation. It only fills the buffers with - the desired values. To apply the tendon limit, call the :meth:`write_fixed_tendon_properties_to_sim` function. - - Args: - limit: Fixed tendon limit. Shape is (len(env_ids), len(fixed_tendon_ids)). - fixed_tendon_ids: The tendon indices to set the limit for. Defaults to None (all fixed tendons). - env_ids: The environment indices to set the limit for. Defaults to None (all environments). - """ - raise NotImplementedError("Fixed tendon position limit is not supported in Newton.") - - def set_fixed_tendon_rest_length( - self, - rest_length: torch.Tensor, - fixed_tendon_ids: Sequence[int] | slice | None = None, - env_ids: Sequence[int] | None = None, - ): - """Set fixed tendon rest length efforts into internal buffers. - - This function does not apply the tendon rest length to the simulation. It only fills the buffers with - the desired values. To apply the tendon rest length, call the :meth:`write_fixed_tendon_properties_to_sim` function. - - Args: - rest_length: Fixed tendon rest length. Shape is (len(env_ids), len(fixed_tendon_ids)). - fixed_tendon_ids: The tendon indices to set the rest length for. Defaults to None (all fixed tendons). - env_ids: The environment indices to set the rest length for. Defaults to None (all environments). - """ - raise NotImplementedError("Fixed tendon rest length is not supported in Newton.") - - def set_fixed_tendon_offset( - self, - offset: torch.Tensor, - fixed_tendon_ids: Sequence[int] | slice | None = None, - env_ids: Sequence[int] | None = None, - ): - """Set fixed tendon offset efforts into internal buffers. - - This function does not apply the tendon offset to the simulation. It only fills the buffers with - the desired values. To apply the tendon offset, call the :meth:`write_fixed_tendon_properties_to_sim` function. - - Args: - offset: Fixed tendon offset. Shape is (len(env_ids), len(fixed_tendon_ids)). - fixed_tendon_ids: The tendon indices to set the offset for. Defaults to None (all fixed tendons). - env_ids: The environment indices to set the offset for. Defaults to None (all environments). - """ - raise NotImplementedError("Fixed tendon offset is not supported in Newton.") - - def write_fixed_tendon_properties_to_sim( - self, - fixed_tendon_ids: Sequence[int] | slice | None = None, - env_ids: Sequence[int] | None = None, - ): - """Write fixed tendon properties into the simulation. - - Args: - fixed_tendon_ids: The fixed tendon indices to set the limits for. Defaults to None (all fixed tendons). - env_ids: The environment indices to set the limits for. Defaults to None (all environments). - """ - raise NotImplementedError("Fixed tendon properties are not supported in Newton.") - - def set_spatial_tendon_stiffness( - self, - stiffness: torch.Tensor, - spatial_tendon_ids: Sequence[int] | slice | None = None, - env_ids: Sequence[int] | None = None, - ): - """Set spatial tendon stiffness into internal buffers. - - This function does not apply the tendon stiffness to the simulation. It only fills the buffers with - the desired values. To apply the tendon stiffness, call the :meth:`write_spatial_tendon_properties_to_sim` function. - - Args: - stiffness: Spatial tendon stiffness. Shape is (len(env_ids), len(spatial_tendon_ids)). - spatial_tendon_ids: The tendon indices to set the stiffness for. Defaults to None (all spatial tendons). - env_ids: The environment indices to set the stiffness for. Defaults to None (all environments). - """ - raise NotImplementedError("Spatial tendon stiffness is not supported in Newton.") - - def set_spatial_tendon_damping( - self, - damping: torch.Tensor, - spatial_tendon_ids: Sequence[int] | slice | None = None, - env_ids: Sequence[int] | None = None, - ): - """Set spatial tendon damping into internal buffers. - - This function does not apply the tendon damping to the simulation. It only fills the buffers with - the desired values. To apply the tendon damping, call the :meth:`write_spatial_tendon_properties_to_sim` function. - - Args: - damping: Spatial tendon damping. Shape is (len(env_ids), len(spatial_tendon_ids)). - spatial_tendon_ids: The tendon indices to set the damping for. Defaults to None (all spatial tendons). - env_ids: The environment indices to set the damping for. Defaults to None (all environments). - """ - raise NotImplementedError("Spatial tendon damping is not supported in Newton.") - - def set_spatial_tendon_limit_stiffness( - self, - limit_stiffness: torch.Tensor, - spatial_tendon_ids: Sequence[int] | slice | None = None, - env_ids: Sequence[int] | None = None, - ): - """Set spatial tendon limit stiffness into internal buffers. - - This function does not apply the tendon limit stiffness to the simulation. It only fills the buffers with - the desired values. To apply the tendon limit stiffness, call the :meth:`write_spatial_tendon_properties_to_sim` function. - - Args: - limit_stiffness: Spatial tendon limit stiffness. Shape is (len(env_ids), len(spatial_tendon_ids)). - spatial_tendon_ids: The tendon indices to set the limit stiffness for. Defaults to None (all spatial tendons). - env_ids: The environment indices to set the limit stiffness for. Defaults to None (all environments). - """ - raise NotImplementedError("Spatial tendon limit stiffness is not supported in Newton.") - - def set_spatial_tendon_offset( - self, - offset: torch.Tensor, - spatial_tendon_ids: Sequence[int] | slice | None = None, - env_ids: Sequence[int] | None = None, - ): - """Set spatial tendon offset efforts into internal buffers. - - This function does not apply the tendon offset to the simulation. It only fills the buffers with - the desired values. To apply the tendon offset, call the :meth:`write_spatial_tendon_properties_to_sim` function. - - Args: - offset: Spatial tendon offset. Shape is (len(env_ids), len(spatial_tendon_ids)). - spatial_tendon_ids: The tendon indices to set the offset for. Defaults to None (all spatial tendons). - env_ids: The environment indices to set the offset for. Defaults to None (all environments). - """ - raise NotImplementedError("Spatial tendon offset is not supported in Newton.") - - def write_spatial_tendon_properties_to_sim( - self, - spatial_tendon_ids: Sequence[int] | slice | None = None, - env_ids: Sequence[int] | None = None, - ): - """Write spatial tendon properties into the simulation. - - Args: - spatial_tendon_ids: The spatial tendon indices to set the properties for. Defaults to None (all spatial tendons). - env_ids: The environment indices to set the properties for. Defaults to None (all environments). - """ - raise NotImplementedError("Spatial tendon properties are not supported in Newton.") - - """ - Internal helper. - """ - - def _initialize_impl(self): - # obtain global simulation view - self._physics_sim_view = SimulationManager.get_physics_sim_view() - - if self.cfg.articulation_root_prim_path is not None: - # The articulation root prim path is specified explicitly, so we can just use this. - root_prim_path_expr = self.cfg.prim_path + self.cfg.articulation_root_prim_path - else: - # No articulation root prim path was specified, so we need to search - # for it. We search for this in the first environment and then - # create a regex that matches all environments. - first_env_matching_prim = sim_utils.find_first_matching_prim(self.cfg.prim_path) - if first_env_matching_prim is None: - raise RuntimeError(f"Failed to find prim for expression: '{self.cfg.prim_path}'.") - first_env_matching_prim_path = first_env_matching_prim.GetPath().pathString - - # Find all articulation root prims in the first environment. - first_env_root_prims = sim_utils.get_all_matching_child_prims( - first_env_matching_prim_path, - predicate=lambda prim: prim.HasAPI(UsdPhysics.ArticulationRootAPI), - ) - if len(first_env_root_prims) == 0: - raise RuntimeError( - f"Failed to find an articulation when resolving '{first_env_matching_prim_path}'." - " Please ensure that the prim has 'USD ArticulationRootAPI' applied." - ) - if len(first_env_root_prims) > 1: - raise RuntimeError( - f"Failed to find a single articulation when resolving '{first_env_matching_prim_path}'." - f" Found multiple '{first_env_root_prims}' under '{first_env_matching_prim_path}'." - " Please ensure that there is only one articulation in the prim path tree." - ) - - # resolve articulation root prim back into regex expression - first_env_root_prim_path = first_env_root_prims[0].GetPath().pathString - root_prim_path_relative_to_prim_path = first_env_root_prim_path[len(first_env_matching_prim_path) :] - root_prim_path_expr = self.cfg.prim_path + root_prim_path_relative_to_prim_path - - prim_path = root_prim_path_expr.replace(".*", "*") - - self._root_newton_view = NewtonArticulationView( - NewtonManager.get_model(), prim_path, verbose=True, exclude_joint_types=[JointType.FREE, JointType.FIXED] - ) - - # log information about the articulation - print(f"[INFO]:Articulation initialized at: {self.cfg.prim_path} with root '{prim_path}'.") - print(f"[INFO]:Is fixed root: {self.is_fixed_base}") - print(f"[INFO]:Number of bodies: {self.num_bodies}") - print(f"[INFO]:Body names: {self.body_names}") - print(f"[INFO]:Number of joints: {self.num_joints}") - print(f"[INFO]:Joint names: {self.joint_names}") - print(f"[INFO]:Number of fixed tendons: {self.num_fixed_tendons}") - - # container for data access - self._data = ArticulationData(self._root_newton_view, self.device) - - # create buffers - self._create_buffers() - # process configuration - self._process_cfg() - self._process_actuators_cfg() - self._process_tendons() - # validate configuration - self._validate_cfg() - # update the robot data - self.update(0.0) - # log joint information - self._log_articulation_info() - - # Moves the articulation to its default pose before the solver is initialized - generated_pose = self._data.default_root_state[:, :7].clone() - generated_pose[:, 3:] = math_utils.convert_quat(generated_pose[:, 3:], to="xyzw") - generated_pose[:, :2] += wp.to_torch(self._root_newton_view.get_root_transforms(NewtonManager.get_model()))[ - :, :2 - ] - self._root_newton_view.set_root_transforms(NewtonManager.get_state_0(), generated_pose) - self._root_newton_view.set_root_transforms(NewtonManager.get_model(), generated_pose) - - def _create_buffers(self): - # constants - self._ALL_INDICES = torch.arange(self.num_instances, dtype=torch.long, device=self.device) - - # external forces and torques - self.has_external_wrench = False - self._external_force_b = torch.zeros((self.num_instances, self.num_bodies, 3), device=self.device) - self._external_torque_b = torch.zeros_like(self._external_force_b) - - # mask for writing to simulation - self._mask = torch.zeros(self.num_instances, dtype=torch.bool, device=self.device) - - # asset named data - self._data.joint_names = self.joint_names - self._data.body_names = self.body_names - # tendon names are set in _process_tendons function - - # -- joint properties - self._data.default_joint_pos_limits = torch.stack( - ( - wp.to_torch(self._root_newton_view.get_attribute("joint_limit_lower", NewtonManager.get_model())), - wp.to_torch(self._root_newton_view.get_attribute("joint_limit_upper", NewtonManager.get_model())), - ), - dim=2, - ).clone() - self._data.default_joint_stiffness = wp.to_torch( - self._root_newton_view.get_attribute("joint_target_ke", NewtonManager.get_model()) - ).clone() - self._data.default_joint_damping = wp.to_torch( - self._root_newton_view.get_attribute("joint_target_kd", NewtonManager.get_model()) - ).clone() - self._data.default_joint_armature = wp.to_torch( - self._root_newton_view.get_attribute("joint_armature", NewtonManager.get_model()) - ).clone() - self._data.default_joint_friction_coeff = wp.to_torch( - self._root_newton_view.get_attribute("joint_friction", NewtonManager.get_model()) - ).clone() - - self._data.joint_pos_limits = self._data.default_joint_pos_limits.clone() - self._data.joint_vel_limits = wp.to_torch( - self._root_newton_view.get_attribute("joint_velocity_limit", NewtonManager.get_model()) - ).clone() - self._data.joint_effort_limits = wp.to_torch( - self._root_newton_view.get_attribute("joint_effort_limit", NewtonManager.get_model()) - ).clone() - self._data.joint_stiffness = self._data.default_joint_stiffness.clone() - self._data.joint_damping = self._data.default_joint_damping.clone() - self._data.joint_armature = self._data.default_joint_armature.clone() - self._data.joint_friction_coeff = self._data.default_joint_friction_coeff.clone() - - # -- body properties - self._data.default_mass = wp.to_torch( - self._root_newton_view.get_attribute("body_mass", NewtonManager.get_model()) - ).clone() - self._data.default_inertia = ( - wp.to_torch(self._root_newton_view.get_attribute("body_inertia", NewtonManager.get_model())) - .clone() - .reshape(self.num_instances, self.num_bodies, 9) - ) - - # -- joint control mode - self._data.joint_control_mode = torch.zeros( - self.num_instances, self.num_joints, dtype=torch.int32, device=self.device - ) - # -- joint commands (sent to the actuator from the user) - self._data.joint_target = torch.zeros(self.num_instances, self.num_joints, device=self.device) - self._data.joint_effort_target = torch.zeros_like(self._data.joint_target) - # -- joint commands (sent to the simulation after actuator processing) - self._joint_target_sim = torch.zeros_like(self._data.joint_target) - self._joint_effort_target_sim = torch.zeros_like(self._data.joint_target) - - # -- computed joint efforts from the actuator models - self._data.computed_torque = torch.zeros_like(self._data.joint_target) - self._data.applied_torque = torch.zeros_like(self._data.joint_target) - - # -- other data that are filled based on explicit actuator models - self._data.soft_joint_vel_limits = torch.zeros(self.num_instances, self.num_joints, device=self.device) - self._data.gear_ratio = torch.ones(self.num_instances, self.num_joints, device=self.device) - - # soft joint position limits (recommended not to be too close to limits). - joint_pos_mean = (self._data.joint_pos_limits[..., 0] + self._data.joint_pos_limits[..., 1]) / 2 - joint_pos_range = self._data.joint_pos_limits[..., 1] - self._data.joint_pos_limits[..., 0] - soft_limit_factor = self.cfg.soft_joint_pos_limit_factor - # add to data - self._data.soft_joint_pos_limits = torch.zeros(self.num_instances, self.num_joints, 2, device=self.device) - self._data.soft_joint_pos_limits[..., 0] = joint_pos_mean - 0.5 * joint_pos_range * soft_limit_factor - self._data.soft_joint_pos_limits[..., 1] = joint_pos_mean + 0.5 * joint_pos_range * soft_limit_factor - - def _process_cfg(self): - """Post processing of configuration parameters.""" - # default state - # -- root state - # note: we cast to tuple to avoid torch/numpy type mismatch. - default_root_state = ( - tuple(self.cfg.init_state.pos) - + tuple(self.cfg.init_state.rot) - + tuple(self.cfg.init_state.lin_vel) - + tuple(self.cfg.init_state.ang_vel) - ) - default_root_state = torch.tensor(default_root_state, dtype=torch.float, device=self.device) - self._data.default_root_state = default_root_state.repeat(self.num_instances, 1) - - # -- joint state - self._data.default_joint_pos = torch.zeros(self.num_instances, self.num_joints, device=self.device) - self._data.default_joint_vel = torch.zeros_like(self._data.default_joint_pos) - - if self.num_joints > 0: - # joint pos - indices_list, _, values_list = string_utils.resolve_matching_names_values( - self.cfg.init_state.joint_pos, self.joint_names - ) - self._data.default_joint_pos[:, indices_list] = torch.tensor(values_list, device=self.device) - - # joint vel - indices_list, _, values_list = string_utils.resolve_matching_names_values( - self.cfg.init_state.joint_vel, self.joint_names - ) - self._data.default_joint_vel[:, indices_list] = torch.tensor(values_list, device=self.device) - - """ - Internal simulation callbacks. - """ - - def _invalidate_initialize_callback(self, event): - """Invalidates the scene elements.""" - # call parent - super()._invalidate_initialize_callback(event) - self._root_physx_view = None - - """ - Internal helpers -- Actuators. - """ - - def _process_actuators_cfg(self): - """Process and apply articulation joint properties.""" - # create actuators - self.actuators = dict() - # flag for implicit actuators - # if this is false, we by-pass certain checks when doing actuator-related operations - self._has_implicit_actuators = False - - # iterate over all actuator configurations - for actuator_name, actuator_cfg in self.cfg.actuators.items(): - # type annotation for type checkers - actuator_cfg: ActuatorBaseCfg - # create actuator group - joint_ids, joint_names = self.find_joints(actuator_cfg.joint_names_expr) - # check if any joints are found - if len(joint_names) == 0: - raise ValueError( - f"No joints found for actuator group: {actuator_name} with joint name expression:" - f" {actuator_cfg.joint_names_expr}." - ) - # resolve joint indices - # we pass a slice if all joints are selected to avoid indexing overhead - if len(joint_names) == self.num_joints: - joint_ids = slice(None) - else: - joint_ids = torch.tensor(joint_ids, device=self.device) - # create actuator collection - # note: for efficiency avoid indexing when over all indices - actuator: ActuatorBase = actuator_cfg.class_type( - cfg=actuator_cfg, - joint_names=joint_names, - joint_ids=joint_ids, - num_envs=self.num_instances, - device=self.device, - stiffness=self._data.default_joint_stiffness[:, joint_ids], - damping=self._data.default_joint_damping[:, joint_ids], - armature=self._data.default_joint_armature[:, joint_ids], - friction=self._data.default_joint_friction_coeff[:, joint_ids], - effort_limit=self._data.joint_effort_limits[:, joint_ids], - velocity_limit=self._data.joint_vel_limits[:, joint_ids], - ) - # log information on actuator groups - model_type = "implicit" if actuator.is_implicit_model else "explicit" - omni.log.info( - f"Actuator collection: {actuator_name} with model '{actuator_cfg.class_type.__name__}'" - f" (type: {model_type}) and joint names: {joint_names} [{joint_ids}]." - ) - # store actuator group - self.actuators[actuator_name] = actuator - # set the passed gains and limits into the simulation - # TODO: write out all joint parameters from simulation - if isinstance(actuator, ImplicitActuator): - self._has_implicit_actuators = True - # the gains and limits are set into the simulation since actuator model is implicit - self.write_joint_stiffness_to_sim(actuator.stiffness, joint_ids=actuator.joint_indices) - self.write_joint_damping_to_sim(actuator.damping, joint_ids=actuator.joint_indices) - # Sets the control mode for the implicit actuators - self.write_joint_control_mode_to_sim(actuator.control_mode, joint_ids=actuator.joint_indices) - else: - # the gains and limits are processed by the actuator model - # we set gains to zero, and torque limit to a high value in simulation to avoid any interference - self.write_joint_stiffness_to_sim(0.0, joint_ids=actuator.joint_indices) - self.write_joint_damping_to_sim(0.0, joint_ids=actuator.joint_indices) - # Set the control mode to None when using explicit actuators - self.write_joint_control_mode_to_sim(None, joint_ids=actuator.joint_indices) - - # Set common properties into the simulation - self.write_joint_effort_limit_to_sim(actuator.effort_limit_sim, joint_ids=actuator.joint_indices) - self.write_joint_velocity_limit_to_sim(actuator.velocity_limit_sim, joint_ids=actuator.joint_indices) - self.write_joint_armature_to_sim(actuator.armature, joint_ids=actuator.joint_indices) - self.write_joint_friction_coefficient_to_sim(actuator.friction, joint_ids=actuator.joint_indices) - - # Store the configured values from the actuator model - # note: this is the value configured in the actuator model (for implicit and explicit actuators) - self._data.default_joint_stiffness[:, actuator.joint_indices] = actuator.stiffness - self._data.default_joint_damping[:, actuator.joint_indices] = actuator.damping - self._data.default_joint_armature[:, actuator.joint_indices] = actuator.armature - self._data.default_joint_friction_coeff[:, actuator.joint_indices] = actuator.friction - - # perform some sanity checks to ensure actuators are prepared correctly - total_act_joints = sum(actuator.num_joints for actuator in self.actuators.values()) - if total_act_joints != (self.num_joints - self.num_fixed_tendons): - omni.log.warn( - "Not all actuators are configured! Total number of actuated joints not equal to number of" - f" joints available: {total_act_joints} != {self.num_joints - self.num_fixed_tendons}." - ) - - def _process_tendons(self): - """Process fixed and spatialtendons.""" - # create a list to store the fixed tendon names - self._fixed_tendon_names = list() - self._spatial_tendon_names = list() - # parse fixed tendons properties if they exist - if self.num_fixed_tendons > 0: - raise NotImplementedError("Tendons are not implemented yet") - - def _apply_actuator_model(self): - """Processes joint commands for the articulation by forwarding them to the actuators. - - The actions are first processed using actuator models. Depending on the robot configuration, - the actuator models compute the joint level simulation commands and sets them into the PhysX buffers. - """ - # process actions per group - for actuator in self.actuators.values(): - # prepare input for actuator model based on cached data - # TODO : A tensor dict would be nice to do the indexing of all tensors together - control_action = ArticulationActions( - joint_targets=self._data.joint_target[:, actuator.joint_indices], - joint_efforts=self._data.joint_effort_target[:, actuator.joint_indices], - joint_indices=actuator.joint_indices, - ) - # compute joint command from the actuator model - control_action = actuator.compute( - control_action, - joint_pos=self._data.joint_pos[:, actuator.joint_indices], - joint_vel=self._data.joint_vel[:, actuator.joint_indices], - ) - # update targets (these are set into the simulation) - if control_action.joint_targets is not None: - self._joint_target_sim[:, actuator.joint_indices] = control_action.joint_targets - if control_action.joint_efforts is not None: - self._joint_effort_target_sim[:, actuator.joint_indices] = control_action.joint_efforts - # update state of the actuator model - # -- torques - self._data.computed_torque[:, actuator.joint_indices] = actuator.computed_effort - self._data.applied_torque[:, actuator.joint_indices] = actuator.applied_effort - # -- actuator data - self._data.soft_joint_vel_limits[:, actuator.joint_indices] = actuator.velocity_limit - # TODO: find a cleaner way to handle gear ratio. Only needed for variable gear ratio actuators. - if hasattr(actuator, "gear_ratio"): - self._data.gear_ratio[:, actuator.joint_indices] = actuator.gear_ratio - - """ - Internal helpers -- Debugging. - """ - - def _validate_cfg(self): - """Validate the configuration after processing. - - Note: - This function should be called only after the configuration has been processed and the buffers have been - created. Otherwise, some settings that are altered during processing may not be validated. - For instance, the actuator models may change the joint max velocity limits. - """ - # check that the default values are within the limits - joint_pos_limits = torch.stack( - ( - wp.to_torch(self._root_newton_view.get_attribute("joint_limit_lower", NewtonManager.get_model())), - wp.to_torch(self._root_newton_view.get_attribute("joint_limit_upper", NewtonManager.get_model())), - ), - dim=2, - )[0].to(self.device) - out_of_range = self._data.default_joint_pos[0] < joint_pos_limits[:, 0] - out_of_range |= self._data.default_joint_pos[0] > joint_pos_limits[:, 1] - violated_indices = torch.nonzero(out_of_range, as_tuple=False).squeeze(-1) - # throw error if any of the default joint positions are out of the limits - if len(violated_indices) > 0: - # prepare message for violated joints - msg = "The following joints have default positions out of the limits: \n" - for idx in violated_indices: - joint_name = self.data.joint_names[idx] - joint_limit = joint_pos_limits[idx] - joint_pos = self.data.default_joint_pos[0, idx] - # add to message - msg += f"\t- '{joint_name}': {joint_pos:.3f} not in [{joint_limit[0]:.3f}, {joint_limit[1]:.3f}]\n" - raise ValueError(msg) - - def _log_articulation_info(self): - """Log information about the articulation. - - Note: We purposefully read the values from the simulator to ensure that the values are configured as expected. - """ - # TODO: read out all joint parameters from simulation - # read out all joint parameters from simulation - # -- gains - stiffnesses = torch.zeros([self.num_joints], dtype=torch.float32, device=self.device).tolist() - dampings = torch.zeros([self.num_joints], dtype=torch.float32, device=self.device).tolist() - # -- properties - armatures = torch.zeros([self.num_joints], dtype=torch.float32, device=self.device).tolist() - frictions = torch.zeros([self.num_joints], dtype=torch.float32, device=self.device).tolist() - # -- limits - position_limits = torch.zeros([self.num_joints, 2], dtype=torch.float32, device=self.device).tolist() - velocity_limits = torch.zeros([self.num_joints], dtype=torch.float32, device=self.device).tolist() - effort_limits = torch.zeros([self.num_joints], dtype=torch.float32, device=self.device).tolist() - # create table for term information - joint_table = PrettyTable() - joint_table.title = f"Simulation Joint Information (Prim path: {self.cfg.prim_path})" - joint_table.field_names = [ - "Index", - "Name", - "Stiffness", - "Damping", - "Armature", - "Friction", - "Position Limits", - "Velocity Limits", - "Effort Limits", - ] - joint_table.float_format = ".3" - joint_table.custom_format["Position Limits"] = lambda f, v: f"[{v[0]:.3f}, {v[1]:.3f}]" - # set alignment of table columns - joint_table.align["Name"] = "l" - # add info on each term - for index, name in enumerate(self.joint_names): - joint_table.add_row([ - index, - name, - stiffnesses[index], - dampings[index], - armatures[index], - frictions[index], - position_limits[index], - velocity_limits[index], - effort_limits[index], - ]) - # convert table to string - omni.log.info(f"Simulation parameters for joints in {self.cfg.prim_path}:\n" + joint_table.get_string()) - - # read out all fixed tendon parameters from simulation - if self.num_fixed_tendons > 0: - raise NotImplementedError("Tendons are not implemented yet") - - if self.num_spatial_tendons > 0: - raise NotImplementedError("Tendons are not implemented yet") - - """ - Deprecated methods. - """ - - def write_joint_friction_to_sim( - self, - joint_friction: torch.Tensor | float, - joint_ids: Sequence[int] | slice | None = None, - env_ids: Sequence[int] | None = None, - ): - """Write joint friction coefficients into the simulation. - - .. deprecated:: 2.1.0 - Please use :meth:`write_joint_friction_coefficient_to_sim` instead. - """ - omni.log.warn( - "The function 'write_joint_friction_to_sim' will be deprecated in a future release. Please" - " use 'write_joint_friction_coefficient_to_sim' instead." - ) - self.write_joint_friction_coefficient_to_sim(joint_friction, joint_ids=joint_ids, env_ids=env_ids) - - def write_joint_limits_to_sim( - self, - limits: torch.Tensor | float, - joint_ids: Sequence[int] | slice | None = None, - env_ids: Sequence[int] | None = None, - warn_limit_violation: bool = True, - ): - """Write joint limits into the simulation. - - .. deprecated:: 2.1.0 - Please use :meth:`write_joint_position_limit_to_sim` instead. - """ - omni.log.warn( - "The function 'write_joint_limits_to_sim' will be deprecated in a future release. Please" - " use 'write_joint_position_limit_to_sim' instead." - ) - self.write_joint_position_limit_to_sim( - limits, joint_ids=joint_ids, env_ids=env_ids, warn_limit_violation=warn_limit_violation - ) + from isaaclab.newton.assets.articulation import Articulation as NewtonArticulation - def set_fixed_tendon_limit( - self, - limit: torch.Tensor, - fixed_tendon_ids: Sequence[int] | slice | None = None, - env_ids: Sequence[int] | None = None, - ): - """Set fixed tendon position limits into internal buffers. +class Articulation(FactoryBase): + """Factory for creating articulation instances.""" - .. deprecated:: 2.1.0 - Please use :meth:`set_fixed_tendon_position_limit` instead. - """ - raise NotImplementedError("This method is not implemented for Newton.") + def __new__(cls, *args, **kwargs) -> BaseArticulation | NewtonArticulation: + """Create a new instance of an articulation based on the backend.""" + # The `FactoryBase` __new__ method will handle the logic and return + # an instance of the correct backend-specific articulation class, + # which is guaranteed to be a subclass of `BaseArticulation` by convention. + return super().__new__(cls, *args, **kwargs) diff --git a/source/isaaclab/isaaclab/assets/articulation/articulation_data.py b/source/isaaclab/isaaclab/assets/articulation/articulation_data.py index 28e75a48431..96f69d1d0e8 100644 --- a/source/isaaclab/isaaclab/assets/articulation/articulation_data.py +++ b/source/isaaclab/isaaclab/assets/articulation/articulation_data.py @@ -1,1137 +1,14 @@ -# Copyright (c) 2022-2025, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md). -# All rights reserved. -# -# SPDX-License-Identifier: BSD-3-Clause +from __future__ import annotations +from .base_articulation_data import BaseArticulationData +from isaaclab.utils.backend_utils import FactoryBase +from typing import TYPE_CHECKING -import torch -import weakref +if TYPE_CHECKING: + from isaaclab.newton.assets.articulation.articulation_data import ArticulationData as NewtonArticulationData -import omni.log -import warp as wp +class ArticulationData(FactoryBase): + """Factory for creating articulation data instances.""" -import isaaclab.utils.math as math_utils -from isaaclab.sim._impl.newton_manager import NewtonManager -from isaaclab.utils.buffers import TimestampedBuffer - - -class ArticulationData: - """Data container for an articulation. - - This class contains the data for an articulation in the simulation. The data includes the state of - the root rigid body, the state of all the bodies in the articulation, and the joint state. The data is - stored in the simulation world frame unless otherwise specified. - - An articulation is comprised of multiple rigid bodies or links. For a rigid body, there are two frames - of reference that are used: - - - Actor frame: The frame of reference of the rigid body prim. This typically corresponds to the Xform prim - with the rigid body schema. - - Center of mass frame: The frame of reference of the center of mass of the rigid body. - - Depending on the settings, the two frames may not coincide with each other. In the robotics sense, the actor frame - can be interpreted as the link frame. - """ - - def __init__(self, root_newton_view, device: str): - """Initializes the articulation data. - - Args: - root_newton_view: The root articulation view. - device: The device used for processing. - """ - # Set the parameters - self.device = device - # Set the root articulation view - # note: this is stored as a weak reference to avoid circular references between the asset class - # and the data container. This is important to avoid memory leaks. - self._root_newton_view = weakref.proxy(root_newton_view) - - # Set initial time stamp - self._sim_timestamp = 0.0 - - # obtain global simulation view - gravity = NewtonManager.get_model().gravity - # Convert to direction vector - gravity_dir = torch.tensor((gravity[0], gravity[1], gravity[2]), device=self.device) - gravity_dir = math_utils.normalize(gravity_dir.unsqueeze(0)).squeeze(0) - - # Initialize constants - self.GRAVITY_VEC_W = gravity_dir.repeat(self._root_newton_view.count, 1) - self.FORWARD_VEC_B = torch.tensor((1.0, 0.0, 0.0), device=self.device).repeat(self._root_newton_view.count, 1) - - # Initialize history for finite differencing - velocity = wp.to_torch(self._root_newton_view.get_link_velocities(NewtonManager.get_state_0())).clone() - self._previous_body_com_vel = torch.cat((velocity[:, :, 3:], velocity[:, :, :3]), dim=-1) - self._previous_joint_vel = wp.to_torch( - self._root_newton_view.get_dof_velocities(NewtonManager.get_state_0()) - ).clone() - - # Initialize the lazy buffers. - # -- link frame w.r.t. world frame - self._root_link_pose_w = TimestampedBuffer() - self._root_link_vel_w = TimestampedBuffer() - self._body_link_pose_w = TimestampedBuffer() - self._body_link_vel_w = TimestampedBuffer() - # -- com frame w.r.t. link frame - self._body_com_pose_b = TimestampedBuffer() - # -- com frame w.r.t. world frame - self._root_com_pose_w = TimestampedBuffer() - self._root_com_vel_w = TimestampedBuffer() - self._body_com_pose_w = TimestampedBuffer() - self._body_com_vel_w = TimestampedBuffer() - self._body_com_acc_w = TimestampedBuffer() - # -- combined state (these are cached as they concatenate) - self._root_state_w = TimestampedBuffer() - self._root_link_state_w = TimestampedBuffer() - self._root_com_state_w = TimestampedBuffer() - self._body_state_w = TimestampedBuffer() - self._body_link_state_w = TimestampedBuffer() - self._body_com_state_w = TimestampedBuffer() - # -- joint state - self._joint_pos = TimestampedBuffer() - self._joint_vel = TimestampedBuffer() - self._joint_acc = TimestampedBuffer() - self._body_incoming_joint_wrench_b = TimestampedBuffer() - - def update(self, dt: float): - # update the simulation timestamp - self._sim_timestamp += dt - # Trigger an update of the joint acceleration buffer at a higher frequency - # since we do finite differencing. - self.joint_acc - - ## - # Names. - ## - - body_names: list[str] = None - """Body names in the order parsed by the simulation view.""" - - joint_names: list[str] = None - """Joint names in the order parsed by the simulation view.""" - - fixed_tendon_names: list[str] = None - """Fixed tendon names in the order parsed by the simulation view.""" - - spatial_tendon_names: list[str] = None - """Spatial tendon names in the order parsed by the simulation view.""" - - ## - # Defaults - Initial state. - ## - - default_root_state: torch.Tensor = None - """Default root state ``[pos, quat, lin_vel, ang_vel]`` in the local environment frame. Shape is (num_instances, 13). - - The position and quaternion are of the articulation root's actor frame. Meanwhile, the linear and angular - velocities are of its center of mass frame. - - This quantity is configured through the :attr:`isaaclab.assets.ArticulationCfg.init_state` parameter. - """ - - default_joint_pos: torch.Tensor = None - """Default joint positions of all joints. Shape is (num_instances, num_joints). - - This quantity is configured through the :attr:`isaaclab.assets.ArticulationCfg.init_state` parameter. - """ - - default_joint_vel: torch.Tensor = None - """Default joint velocities of all joints. Shape is (num_instances, num_joints). - - This quantity is configured through the :attr:`isaaclab.assets.ArticulationCfg.init_state` parameter. - """ - - ## - # Defaults - Physical properties. - ## - - default_mass: torch.Tensor = None - """Default mass for all the bodies in the articulation. Shape is (num_instances, num_bodies). - - This quantity is parsed from the USD schema at the time of initialization. - """ - - default_inertia: torch.Tensor = None - """Default inertia for all the bodies in the articulation. Shape is (num_instances, num_bodies, 9). - - The inertia is the inertia tensor relative to the center of mass frame. The values are stored in - the order :math:`[I_{xx}, I_{xy}, I_{xz}, I_{yx}, I_{yy}, I_{yz}, I_{zx}, I_{zy}, I_{zz}]`. - - This quantity is parsed from the USD schema at the time of initialization. - """ - - default_joint_stiffness: torch.Tensor = None - """Default joint stiffness of all joints. Shape is (num_instances, num_joints). - - This quantity is configured through the actuator model's :attr:`isaaclab.actuators.ActuatorBaseCfg.stiffness` - parameter. If the parameter's value is None, the value parsed from the USD schema, at the time of initialization, - is used. - - .. attention:: - The default stiffness is the value configured by the user or the value parsed from the USD schema. - It should not be confused with :attr:`joint_stiffness`, which is the value set into the simulation. - """ - - default_joint_damping: torch.Tensor = None - """Default joint damping of all joints. Shape is (num_instances, num_joints). - - This quantity is configured through the actuator model's :attr:`isaaclab.actuators.ActuatorBaseCfg.damping` - parameter. If the parameter's value is None, the value parsed from the USD schema, at the time of initialization, - is used. - - .. attention:: - The default stiffness is the value configured by the user or the value parsed from the USD schema. - It should not be confused with :attr:`joint_damping`, which is the value set into the simulation. - """ - - default_joint_armature: torch.Tensor = None - """Default joint armature of all joints. Shape is (num_instances, num_joints). - - This quantity is configured through the actuator model's :attr:`isaaclab.actuators.ActuatorBaseCfg.armature` - parameter. If the parameter's value is None, the value parsed from the USD schema, at the time of initialization, - is used. - """ - - default_joint_friction_coeff: torch.Tensor = None - """Default joint friction coefficient of all joints. Shape is (num_instances, num_joints). - - This quantity is configured through the actuator model's :attr:`isaaclab.actuators.ActuatorBaseCfg.friction` - parameter. If the parameter's value is None, the value parsed from the USD schema, at the time of initialization, - is used. - """ - - default_joint_pos_limits: torch.Tensor = None - """Default joint position limits of all joints. Shape is (num_instances, num_joints, 2). - - The limits are in the order :math:`[lower, upper]`. They are parsed from the USD schema at the time of initialization. - """ - default_fixed_tendon_stiffness: torch.Tensor = None - """Default tendon stiffness of all fixed tendons. Shape is (num_instances, num_fixed_tendons). - - This quantity is parsed from the USD schema at the time of initialization. - """ - - default_fixed_tendon_damping: torch.Tensor = None - """Default tendon damping of all fixed tendons. Shape is (num_instances, num_fixed_tendons). - - This quantity is parsed from the USD schema at the time of initialization. - """ - - default_fixed_tendon_limit_stiffness: torch.Tensor = None - """Default tendon limit stiffness of all fixed tendons. Shape is (num_instances, num_fixed_tendons). - - This quantity is parsed from the USD schema at the time of initialization. - """ - - default_fixed_tendon_rest_length: torch.Tensor = None - """Default tendon rest length of all fixed tendons. Shape is (num_instances, num_fixed_tendons). - - This quantity is parsed from the USD schema at the time of initialization. - """ - - default_fixed_tendon_offset: torch.Tensor = None - """Default tendon offset of all fixed tendons. Shape is (num_instances, num_fixed_tendons). - - This quantity is parsed from the USD schema at the time of initialization. - """ - - default_fixed_tendon_pos_limits: torch.Tensor = None - """Default tendon position limits of all fixed tendons. Shape is (num_instances, num_fixed_tendons, 2). - - The position limits are in the order :math:`[lower, upper]`. They are parsed from the USD schema at the time of - initialization. - """ - - default_spatial_tendon_stiffness: torch.Tensor = None - """Default tendon stiffness of all spatial tendons. Shape is (num_instances, num_spatial_tendons). - - This quantity is parsed from the USD schema at the time of initialization. - """ - - default_spatial_tendon_damping: torch.Tensor = None - """Default tendon damping of all spatial tendons. Shape is (num_instances, num_spatial_tendons). - - This quantity is parsed from the USD schema at the time of initialization. - """ - - default_spatial_tendon_limit_stiffness: torch.Tensor = None - """Default tendon limit stiffness of all spatial tendons. Shape is (num_instances, num_spatial_tendons). - - This quantity is parsed from the USD schema at the time of initialization. - """ - - default_spatial_tendon_offset: torch.Tensor = None - """Default tendon offset of all spatial tendons. Shape is (num_instances, num_spatial_tendons). - - This quantity is parsed from the USD schema at the time of initialization. - """ - - ## - # Joint commands -- Set into simulation. - ## - - joint_control_mode: torch.Tensor = None - """Joint control mode. Shape is (num_instances, num_joints). - - When using implicit actuator models Newton needs to know how the joints are controlled. - The control mode can be one of the following: - - * None: 0 - * Position control: 1 - * Velocity control: 2 - - This quantity is set by the :meth:`Articulation.write_joint_control_mode_to_sim` method. - """ - - joint_target: torch.Tensor = None - """Joint position targets commanded by the user. Shape is (num_instances, num_joints). - - For an implicit actuator model, the targets are directly set into the simulation. - For an explicit actuator model, the targets are used to compute the joint torques (see :attr:`applied_torque`), - which are then set into the simulation. - """ - - joint_effort_target: torch.Tensor = None - """Joint effort targets commanded by the user. Shape is (num_instances, num_joints). - - For an implicit actuator model, the targets are directly set into the simulation. - For an explicit actuator model, the targets are used to compute the joint torques (see :attr:`applied_torque`), - which are then set into the simulation. - """ - - ## - # Joint commands -- Explicit actuators. - ## - - computed_torque: torch.Tensor = None - """Joint torques computed from the actuator model (before clipping). Shape is (num_instances, num_joints). - - This quantity is the raw torque output from the actuator mode, before any clipping is applied. - It is exposed for users who want to inspect the computations inside the actuator model. - For instance, to penalize the learning agent for a difference between the computed and applied torques. - """ - - applied_torque: torch.Tensor = None - """Joint torques applied from the actuator model (after clipping). Shape is (num_instances, num_joints). - - These torques are set into the simulation, after clipping the :attr:`computed_torque` based on the - actuator model. - """ - - ## - # Joint properties. - ## - - joint_stiffness: torch.Tensor = None - """Joint stiffness provided to the simulation. Shape is (num_instances, num_joints). - - In the case of explicit actuators, the value for the corresponding joints is zero. - """ - - joint_damping: torch.Tensor = None - """Joint damping provided to the simulation. Shape is (num_instances, num_joints) - - In the case of explicit actuators, the value for the corresponding joints is zero. - """ - - joint_armature: torch.Tensor = None - """Joint armature provided to the simulation. Shape is (num_instances, num_joints).""" - - joint_friction_coeff: torch.Tensor = None - """Joint friction coefficient provided to the simulation. Shape is (num_instances, num_joints).""" - - joint_pos_limits: torch.Tensor = None - """Joint position limits provided to the simulation. Shape is (num_instances, num_joints, 2). - - The limits are in the order :math:`[lower, upper]`. - """ - - joint_vel_limits: torch.Tensor = None - """Joint maximum velocity provided to the simulation. Shape is (num_instances, num_joints).""" - - joint_effort_limits: torch.Tensor = None - """Joint maximum effort provided to the simulation. Shape is (num_instances, num_joints).""" - - ## - # Joint properties - Custom. - ## - - soft_joint_pos_limits: torch.Tensor = None - r"""Soft joint positions limits for all joints. Shape is (num_instances, num_joints, 2). - - The limits are in the order :math:`[lower, upper]`.The soft joint position limits are computed as - a sub-region of the :attr:`joint_pos_limits` based on the - :attr:`~isaaclab.assets.ArticulationCfg.soft_joint_pos_limit_factor` parameter. - - Consider the joint position limits :math:`[lower, upper]` and the soft joint position limits - :math:`[soft_lower, soft_upper]`. The soft joint position limits are computed as: - - .. math:: - - soft\_lower = (lower + upper) / 2 - factor * (upper - lower) / 2 - soft\_upper = (lower + upper) / 2 + factor * (upper - lower) / 2 - - The soft joint position limits help specify a safety region around the joint limits. It isn't used by the - simulation, but is useful for learning agents to prevent the joint positions from violating the limits. - """ - - soft_joint_vel_limits: torch.Tensor = None - """Soft joint velocity limits for all joints. Shape is (num_instances, num_joints). - - These are obtained from the actuator model. It may differ from :attr:`joint_vel_limits` if the actuator model - has a variable velocity limit model. For instance, in a variable gear ratio actuator model. - """ - - gear_ratio: torch.Tensor = None - """Gear ratio for relating motor torques to applied Joint torques. Shape is (num_instances, num_joints).""" - - ## - # Fixed tendon properties. - ## - - fixed_tendon_stiffness: torch.Tensor = None - """Fixed tendon stiffness provided to the simulation. Shape is (num_instances, num_fixed_tendons).""" - - fixed_tendon_damping: torch.Tensor = None - """Fixed tendon damping provided to the simulation. Shape is (num_instances, num_fixed_tendons).""" - - fixed_tendon_limit_stiffness: torch.Tensor = None - """Fixed tendon limit stiffness provided to the simulation. Shape is (num_instances, num_fixed_tendons).""" - - fixed_tendon_rest_length: torch.Tensor = None - """Fixed tendon rest length provided to the simulation. Shape is (num_instances, num_fixed_tendons).""" - - fixed_tendon_offset: torch.Tensor = None - """Fixed tendon offset provided to the simulation. Shape is (num_instances, num_fixed_tendons).""" - - fixed_tendon_pos_limits: torch.Tensor = None - """Fixed tendon position limits provided to the simulation. Shape is (num_instances, num_fixed_tendons, 2).""" - - ## - # Spatial tendon properties. - ## - - spatial_tendon_stiffness: torch.Tensor = None - """Spatial tendon stiffness provided to the simulation. Shape is (num_instances, num_spatial_tendons).""" - - spatial_tendon_damping: torch.Tensor = None - """Spatial tendon damping provided to the simulation. Shape is (num_instances, num_spatial_tendons).""" - - spatial_tendon_limit_stiffness: torch.Tensor = None - """Spatial tendon limit stiffness provided to the simulation. Shape is (num_instances, num_spatial_tendons).""" - - spatial_tendon_offset: torch.Tensor = None - """Spatial tendon offset provided to the simulation. Shape is (num_instances, num_spatial_tendons).""" - - ## - # Root state properties. - ## - - @property - def root_link_pose_w(self) -> torch.Tensor: - """Root link pose ``[pos, quat]`` in simulation world frame. Shape is (num_instances, 7). - - This quantity is the pose of the articulation root's actor frame relative to the world. - The orientation is provided in (w, x, y, z) format. - """ - if self._root_link_pose_w.timestamp < self._sim_timestamp: - # read data from simulation - # Newton reads poses as [x, y, z, qx, qy, qz, qw] Isaac reads as [x, y, z, qw, qx, qy, qz] - pose = wp.to_torch(self._root_newton_view.get_root_transforms(NewtonManager.get_state_0())).clone() - pose[:, 3:7] = math_utils.convert_quat(pose[:, 3:7], to="wxyz") - # set the buffer data and timestamp - self._root_link_pose_w.data = pose - self._root_link_pose_w.timestamp = self._sim_timestamp - - return self._root_link_pose_w.data - - @property - def root_link_vel_w(self) -> torch.Tensor: - """Root link velocity ``[lin_vel, ang_vel]`` in simulation world frame. Shape is (num_instances, 6). - - This quantity contains the linear and angular velocities of the articulation root's actor frame - relative to the world. - """ - if self._root_link_vel_w.timestamp < self._sim_timestamp: - # read the CoM velocity - vel = self.root_com_vel_w.clone() - # adjust linear velocity to link from center of mass - vel[:, :3] += torch.linalg.cross( - vel[:, 3:], math_utils.quat_apply(self.root_link_quat_w, -self.body_com_pos_b[:, 0]), dim=-1 - ) - # set the buffer data and timestamp - self._root_link_vel_w.data = vel - self._root_link_vel_w.timestamp = self._sim_timestamp - - return self._root_link_vel_w.data - - @property - def root_com_pose_w(self) -> torch.Tensor: - """Root center of mass pose ``[pos, quat]`` in simulation world frame. Shape is (num_instances, 7). - - This quantity is the pose of the articulation root's center of mass frame relative to the world. - The orientation is provided in (w, x, y, z) format. - """ - if self._root_com_pose_w.timestamp < self._sim_timestamp: - # apply local transform to center of mass frame - pos, quat = math_utils.combine_frame_transforms( - self.root_link_pos_w, self.root_link_quat_w, self.body_com_pos_b[:, 0], self.body_com_quat_b[:, 0] - ) - # set the buffer data and timestamp - self._root_com_pose_w.data = torch.cat((pos, quat), dim=-1) - self._root_com_pose_w.timestamp = self._sim_timestamp - - return self._root_com_pose_w.data - - @property - def root_com_vel_w(self) -> torch.Tensor: - """Root center of mass velocity ``[lin_vel, ang_vel]`` in simulation world frame. Shape is (num_instances, 6). - - This quantity contains the linear and angular velocities of the articulation root's center of mass frame - relative to the world. - """ - if self._root_com_vel_w.timestamp < self._sim_timestamp: - # Newton and Isaac: [vx, vy, vz, wx, wy, wz] - velocity = self._root_newton_view.get_root_velocities(NewtonManager.get_state_0()) - if velocity is None: - velocity = torch.zeros((self._root_newton_view.count, 6), device=self.device) - else: - velocity = wp.to_torch(velocity).clone() - self._root_com_vel_w.data = velocity - self._root_com_vel_w.timestamp = self._sim_timestamp - - return self._root_com_vel_w.data - - @property - def root_state_w(self): - """Root state ``[pos, quat, lin_vel, ang_vel]`` in simulation world frame. Shape is (num_instances, 13). - - The position and quaternion are of the articulation root's actor frame relative to the world. Meanwhile, - the linear and angular velocities are of the articulation root's center of mass frame. - """ - if self._root_state_w.timestamp < self._sim_timestamp: - self._root_state_w.data = torch.cat((self.root_link_pose_w, self.root_com_vel_w), dim=-1) - self._root_state_w.timestamp = self._sim_timestamp - - return self._root_state_w.data - - @property - def root_link_state_w(self): - """Root state ``[pos, quat, lin_vel, ang_vel]`` in simulation world frame. Shape is (num_instances, 13). - - The position, quaternion, and linear/angular velocity are of the articulation root's actor frame relative to the - world. - """ - if self._root_link_state_w.timestamp < self._sim_timestamp: - self._root_link_state_w.data = torch.cat((self.root_link_pose_w, self.root_link_vel_w), dim=-1) - self._root_link_state_w.timestamp = self._sim_timestamp - - return self._root_link_state_w.data - - @property - def root_com_state_w(self): - """Root center of mass state ``[pos, quat, lin_vel, ang_vel]`` in simulation world frame. Shape is (num_instances, 13). - - The position, quaternion, and linear/angular velocity are of the articulation root link's center of mass frame - relative to the world. Center of mass frame is assumed to be the same orientation as the link rather than the - orientation of the principle inertia. - """ - if self._root_com_state_w.timestamp < self._sim_timestamp: - self._root_com_state_w.data = torch.cat((self.root_com_pose_w, self.root_com_vel_w), dim=-1) - self._root_com_state_w.timestamp = self._sim_timestamp - - return self._root_com_state_w.data - - ## - # Body state properties. - ## - - @property - def body_link_pose_w(self) -> torch.Tensor: - """Body link pose ``[pos, quat]`` in simulation world frame. - Shape is (num_instances, num_bodies, 7). - - This quantity is the pose of the articulation links' actor frame relative to the world. - The orientation is provided in (w, x, y, z) format. - """ - if self._body_link_pose_w.timestamp < self._sim_timestamp: - # perform forward kinematics (shouldn't cause overhead if it happened already) - # self._physics_sim_view.update_articulations_kinematic() - # read data from simulation - # Newton reads poses as [x, y, z, qx, qy, qz, qw] Isaac reads as [x, y, z, qw, qx, qy, qz] - poses = wp.to_torch(self._root_newton_view.get_link_transforms(NewtonManager.get_state_0())).clone() - poses[..., 3:7] = math_utils.convert_quat(poses[..., 3:7], to="wxyz") - # set the buffer data and timestamp - self._body_link_pose_w.data = poses - self._body_link_pose_w.timestamp = self._sim_timestamp - - return self._body_link_pose_w.data - - @property - def body_link_vel_w(self) -> torch.Tensor: - """Body link velocity ``[lin_vel, ang_vel]`` in simulation world frame. - Shape is (num_instances, num_bodies, 6). - - This quantity contains the linear and angular velocities of the articulation links' actor frame - relative to the world. - """ - if self._body_link_vel_w.timestamp < self._sim_timestamp: - # read data from simulation - velocities = self.body_com_vel_w.clone() - # adjust linear velocity to link from center of mass - velocities[..., :3] += torch.linalg.cross( - velocities[..., 3:], math_utils.quat_apply(self.body_link_quat_w, -self.body_com_pos_b), dim=-1 - ) - # set the buffer data and timestamp - self._body_link_vel_w.data = velocities - self._body_link_vel_w.timestamp = self._sim_timestamp - - return self._body_link_vel_w.data - - @property - def body_com_pose_w(self) -> torch.Tensor: - """Body center of mass pose ``[pos, quat]`` in simulation world frame. - Shape is (num_instances, num_bodies, 7). - - This quantity is the pose of the center of mass frame of the articulation links relative to the world. - The orientation is provided in (w, x, y, z) format. - """ - if self._body_com_pose_w.timestamp < self._sim_timestamp: - # apply local transform to center of mass frame - pos, quat = math_utils.combine_frame_transforms( - self.body_link_pos_w, self.body_link_quat_w, self.body_com_pos_b, self.body_com_quat_b - ) - # set the buffer data and timestamp - self._body_com_pose_w.data = torch.cat((pos, quat), dim=-1) - self._body_com_pose_w.timestamp = self._sim_timestamp - - return self._body_com_pose_w.data - - @property - def body_com_vel_w(self) -> torch.Tensor: - """Body center of mass velocity ``[lin_vel, ang_vel]`` in simulation world frame. - Shape is (num_instances, num_bodies, 6). - - This quantity contains the linear and angular velocities of the articulation links' center of mass frame - relative to the world. - """ - if self._body_com_vel_w.timestamp < self._sim_timestamp: - # Newton reads velocities as [wx, wy, wz, vx, vy, vz] Isaac reads as [vx, vy, vz, wx, wy, wz] - velocity = wp.to_torch(self._root_newton_view.get_link_velocities(NewtonManager.get_state_0())).clone() - self._body_com_vel_w.data = torch.cat((velocity[:, :, 3:], velocity[:, :, :3]), dim=-1) - self._body_com_vel_w.timestamp = self._sim_timestamp - - return self._body_com_vel_w.data - - @property - def body_state_w(self): - """State of all bodies `[pos, quat, lin_vel, ang_vel]` in simulation world frame. - Shape is (num_instances, num_bodies, 13). - - The position and quaternion are of all the articulation links' actor frame. Meanwhile, the linear and angular - velocities are of the articulation links's center of mass frame. - """ - if self._body_state_w.timestamp < self._sim_timestamp: - self._body_state_w.data = torch.cat((self.body_link_pose_w, self.body_com_vel_w), dim=-1) - self._body_state_w.timestamp = self._sim_timestamp - - return self._body_state_w.data - - @property - def body_link_state_w(self): - """State of all bodies' link frame`[pos, quat, lin_vel, ang_vel]` in simulation world frame. - Shape is (num_instances, num_bodies, 13). - - The position, quaternion, and linear/angular velocity are of the body's link frame relative to the world. - """ - if self._body_link_state_w.timestamp < self._sim_timestamp: - self._body_link_state_w.data = torch.cat((self.body_link_pose_w, self.body_link_vel_w), dim=-1) - self._body_link_state_w.timestamp = self._sim_timestamp - - return self._body_link_state_w.data - - @property - def body_com_state_w(self): - """State of all bodies center of mass `[pos, quat, lin_vel, ang_vel]` in simulation world frame. - Shape is (num_instances, num_bodies, 13). - - The position, quaternion, and linear/angular velocity are of the body's center of mass frame relative to the - world. Center of mass frame is assumed to be the same orientation as the link rather than the orientation of the - principle inertia. - """ - if self._body_com_state_w.timestamp < self._sim_timestamp: - self._body_com_state_w.data = torch.cat((self.body_com_pose_w, self.body_com_vel_w), dim=-1) - self._body_com_state_w.timestamp = self._sim_timestamp - - return self._body_com_state_w.data - - @property - def body_com_acc_w(self): - """Acceleration of all bodies center of mass ``[lin_acc, ang_acc]``. - Shape is (num_instances, num_bodies, 6). - - All values are relative to the world. - """ - if self._body_com_acc_w.timestamp < self._sim_timestamp: - # read data from simulation and set the buffer data and timestamp - # self._body_com_acc_w.data = self._root_physx_view.get_link_accelerations() - time_elapsed = self._sim_timestamp - self._joint_acc.timestamp - self._body_com_acc_w.data = (self.body_com_vel_w - self._previous_body_com_vel) / time_elapsed - self._previous_body_com_vel[:] = self.body_com_vel_w - self._body_com_acc_w.timestamp = self._sim_timestamp - return self._body_com_acc_w.data - - @property - def body_com_pose_b(self) -> torch.Tensor: - """Center of mass pose ``[pos, quat]`` of all bodies in their respective body's link frames. - Shape is (num_instances, 1, 7). - - This quantity is the pose of the center of mass frame of the rigid body relative to the body's link frame. - The orientation is provided in (w, x, y, z) format. - """ - if self._body_com_pose_b.timestamp < self._sim_timestamp: - # read data from simulation - - position = wp.to_torch(self._root_newton_view.get_attribute("body_com", NewtonManager.get_model())).clone() - quat = torch.zeros((position.shape[0], position.shape[1], 4), device=self.device) - quat[:, :, 0] = 1.0 - pose = torch.cat((position, quat), dim=-1) - # set the buffer data and timestamp - self._body_com_pose_b.data = pose - self._body_com_pose_b.timestamp = self._sim_timestamp - - return self._body_com_pose_b.data - - @property - def body_incoming_joint_wrench_b(self) -> torch.Tensor: - """Joint reaction wrench applied from body parent to child body in parent body frame. - - Shape is (num_instances, num_bodies, 6). All body reaction wrenches are provided including the root body to the - world of an articulation. - - For more information on joint wrenches, please check the`PhysX documentation `__ - and the underlying `PhysX Tensor API `__ . - """ - raise NotImplementedError("Body incoming joint wrench in body frame is not implemented for Newton.") - if self._body_incoming_joint_wrench_b.timestamp < self._sim_timestamp: - self._body_incoming_joint_wrench_b.data = self._root_physx_view.get_link_incoming_joint_force() - self._body_incoming_joint_wrench_b.time_stamp = self._sim_timestamp - return self._body_incoming_joint_wrench_b.data - - ## - # Joint state properties. - ## - - @property - def joint_pos(self): - """Joint positions of all joints. Shape is (num_instances, num_joints).""" - if self._joint_pos.timestamp < self._sim_timestamp: - # read data from simulation and set the buffer data and timestamp - self._joint_pos.data = wp.to_torch(self._root_newton_view.get_dof_positions(NewtonManager.get_state_0())) - self._joint_pos.timestamp = self._sim_timestamp - return self._joint_pos.data - - @property - def joint_vel(self): - """Joint velocities of all joints. Shape is (num_instances, num_joints).""" - if self._joint_vel.timestamp < self._sim_timestamp: - # read data from simulation and set the buffer data and timestamp - self._joint_vel.data = wp.to_torch(self._root_newton_view.get_dof_velocities(NewtonManager.get_state_0())) - self._joint_vel.timestamp = self._sim_timestamp - return self._joint_vel.data - - @property - def joint_acc(self): - """Joint acceleration of all joints. Shape is (num_instances, num_joints).""" - if self._joint_acc.timestamp < self._sim_timestamp: - # note: we use finite differencing to compute acceleration - time_elapsed = self._sim_timestamp - self._joint_acc.timestamp - self._joint_acc.data = (self.joint_vel - self._previous_joint_vel) / time_elapsed - self._joint_acc.timestamp = self._sim_timestamp - # update the previous joint velocity - self._previous_joint_vel[:] = self.joint_vel - return self._joint_acc.data - - ## - # Derived Properties. - ## - - @property - def projected_gravity_b(self): - """Projection of the gravity direction on base frame. Shape is (num_instances, 3).""" - return math_utils.quat_apply_inverse(self.root_link_quat_w, self.GRAVITY_VEC_W) - - @property - def heading_w(self): - """Yaw heading of the base frame (in radians). Shape is (num_instances,). - - Note: - This quantity is computed by assuming that the forward-direction of the base - frame is along x-direction, i.e. :math:`(1, 0, 0)`. - """ - forward_w = math_utils.quat_apply(self.root_link_quat_w, self.FORWARD_VEC_B) - return torch.atan2(forward_w[:, 1], forward_w[:, 0]) - - @property - def root_link_lin_vel_b(self) -> torch.Tensor: - """Root link linear velocity in base frame. Shape is (num_instances, 3). - - This quantity is the linear velocity of the articulation root's actor frame with respect to the - its actor frame. - """ - return math_utils.quat_apply_inverse(self.root_link_quat_w, self.root_link_lin_vel_w) - - @property - def root_link_ang_vel_b(self) -> torch.Tensor: - """Root link angular velocity in base world frame. Shape is (num_instances, 3). - - This quantity is the angular velocity of the articulation root's actor frame with respect to the - its actor frame. - """ - return math_utils.quat_apply_inverse(self.root_link_quat_w, self.root_link_ang_vel_w) - - @property - def root_com_lin_vel_b(self) -> torch.Tensor: - """Root center of mass linear velocity in base frame. Shape is (num_instances, 3). - - This quantity is the linear velocity of the articulation root's center of mass frame with respect to the - its actor frame. - """ - return math_utils.quat_apply_inverse(self.root_link_quat_w, self.root_com_lin_vel_w) - - @property - def root_com_ang_vel_b(self) -> torch.Tensor: - """Root center of mass angular velocity in base world frame. Shape is (num_instances, 3). - - This quantity is the angular velocity of the articulation root's center of mass frame with respect to the - its actor frame. - """ - return math_utils.quat_apply_inverse(self.root_link_quat_w, self.root_com_ang_vel_w) - - ## - # Sliced properties. - ## - - @property - def root_link_pos_w(self) -> torch.Tensor: - """Root link position in simulation world frame. Shape is (num_instances, 3). - - This quantity is the position of the actor frame of the root rigid body relative to the world. - """ - return self.root_link_pose_w[:, :3] - - @property - def root_link_quat_w(self) -> torch.Tensor: - """Root link orientation (w, x, y, z) in simulation world frame. Shape is (num_instances, 4). - - This quantity is the orientation of the actor frame of the root rigid body. - """ - return self.root_link_pose_w[:, 3:7] - - @property - def root_link_lin_vel_w(self) -> torch.Tensor: - """Root linear velocity in simulation world frame. Shape is (num_instances, 3). - - This quantity is the linear velocity of the root rigid body's actor frame relative to the world. - """ - return self.root_link_vel_w[:, :3] - - @property - def root_link_ang_vel_w(self) -> torch.Tensor: - """Root link angular velocity in simulation world frame. Shape is (num_instances, 3). - - This quantity is the angular velocity of the actor frame of the root rigid body relative to the world. - """ - return self.root_link_vel_w[:, 3:6] - - @property - def root_com_pos_w(self) -> torch.Tensor: - """Root center of mass position in simulation world frame. Shape is (num_instances, 3). - - This quantity is the position of the actor frame of the root rigid body relative to the world. - """ - return self.root_com_pose_w[:, :3] - - @property - def root_com_quat_w(self) -> torch.Tensor: - """Root center of mass orientation (w, x, y, z) in simulation world frame. Shape is (num_instances, 4). - - This quantity is the orientation of the actor frame of the root rigid body relative to the world. - """ - return self.root_com_pose_w[:, 3:7] - - @property - def root_com_lin_vel_w(self) -> torch.Tensor: - """Root center of mass linear velocity in simulation world frame. Shape is (num_instances, 3). - - This quantity is the linear velocity of the root rigid body's center of mass frame relative to the world. - """ - return self.root_com_vel_w[:, :3] - - @property - def root_com_ang_vel_w(self) -> torch.Tensor: - """Root center of mass angular velocity in simulation world frame. Shape is (num_instances, 3). - - This quantity is the angular velocity of the root rigid body's center of mass frame relative to the world. - """ - return self.root_com_vel_w[:, 3:6] - - @property - def body_link_pos_w(self) -> torch.Tensor: - """Positions of all bodies in simulation world frame. Shape is (num_instances, num_bodies, 3). - - This quantity is the position of the articulation bodies' actor frame relative to the world. - """ - return self.body_link_pose_w[..., :3] - - @property - def body_link_quat_w(self) -> torch.Tensor: - """Orientation (w, x, y, z) of all bodies in simulation world frame. Shape is (num_instances, num_bodies, 4). - - This quantity is the orientation of the articulation bodies' actor frame relative to the world. - """ - return self.body_link_pose_w[..., 3:7] - - @property - def body_link_lin_vel_w(self) -> torch.Tensor: - """Linear velocity of all bodies in simulation world frame. Shape is (num_instances, num_bodies, 3). - - This quantity is the linear velocity of the articulation bodies' center of mass frame relative to the world. - """ - return self.body_link_vel_w[..., :3] - - @property - def body_link_ang_vel_w(self) -> torch.Tensor: - """Angular velocity of all bodies in simulation world frame. Shape is (num_instances, num_bodies, 3). - - This quantity is the angular velocity of the articulation bodies' center of mass frame relative to the world. - """ - return self.body_link_vel_w[..., 3:6] - - @property - def body_com_pos_w(self) -> torch.Tensor: - """Positions of all bodies in simulation world frame. Shape is (num_instances, num_bodies, 3). - - This quantity is the position of the articulation bodies' actor frame. - """ - return self.body_com_pose_w[..., :3] - - @property - def body_com_quat_w(self) -> torch.Tensor: - """Orientation (w, x, y, z) of the principle axis of inertia of all bodies in simulation world frame. - Shape is (num_instances, num_bodies, 4). - - This quantity is the orientation of the articulation bodies' actor frame. - """ - return self.body_com_pose_w[..., 3:7] - - @property - def body_com_lin_vel_w(self) -> torch.Tensor: - """Linear velocity of all bodies in simulation world frame. Shape is (num_instances, num_bodies, 3). - - This quantity is the linear velocity of the articulation bodies' center of mass frame. - """ - return self.body_com_vel_w[..., :3] - - @property - def body_com_ang_vel_w(self) -> torch.Tensor: - """Angular velocity of all bodies in simulation world frame. Shape is (num_instances, num_bodies, 3). - - This quantity is the angular velocity of the articulation bodies' center of mass frame. - """ - return self.body_com_vel_w[..., 3:6] - - @property - def body_com_lin_acc_w(self) -> torch.Tensor: - """Linear acceleration of all bodies in simulation world frame. Shape is (num_instances, num_bodies, 3). - - This quantity is the linear acceleration of the articulation bodies' center of mass frame. - """ - return self.body_com_acc_w[..., :3] - - @property - def body_com_ang_acc_w(self) -> torch.Tensor: - """Angular acceleration of all bodies in simulation world frame. Shape is (num_instances, num_bodies, 3). - - This quantity is the angular acceleration of the articulation bodies' center of mass frame. - """ - return self.body_com_acc_w[..., 3:6] - - @property - def body_com_pos_b(self) -> torch.Tensor: - """Center of mass position of all of the bodies in their respective link frames. - Shape is (num_instances, num_bodies, 3). - - This quantity is the center of mass location relative to its body'slink frame. - """ - return self.body_com_pose_b[..., :3] - - @property - def body_com_quat_b(self) -> torch.Tensor: - """Orientation (w, x, y, z) of the principle axis of inertia of all of the bodies in their - respective link frames. Shape is (num_instances, num_bodies, 4). - - This quantity is the orientation of the principles axes of inertia relative to its body's link frame. - """ - return self.body_com_pose_b[..., 3:7] - - ## - # Backward compatibility. - ## - - @property - def root_pose_w(self) -> torch.Tensor: - """Same as :attr:`root_link_pose_w`.""" - return self.root_link_pose_w - - @property - def root_pos_w(self) -> torch.Tensor: - """Same as :attr:`root_link_pos_w`.""" - return self.root_link_pos_w - - @property - def root_quat_w(self) -> torch.Tensor: - """Same as :attr:`root_link_quat_w`.""" - return self.root_link_quat_w - - @property - def root_vel_w(self) -> torch.Tensor: - """Same as :attr:`root_com_vel_w`.""" - return self.root_com_vel_w - - @property - def root_lin_vel_w(self) -> torch.Tensor: - """Same as :attr:`root_com_lin_vel_w`.""" - return self.root_com_lin_vel_w - - @property - def root_ang_vel_w(self) -> torch.Tensor: - """Same as :attr:`root_com_ang_vel_w`.""" - return self.root_com_ang_vel_w - - @property - def root_lin_vel_b(self) -> torch.Tensor: - """Same as :attr:`root_com_lin_vel_b`.""" - return self.root_com_lin_vel_b - - @property - def root_ang_vel_b(self) -> torch.Tensor: - """Same as :attr:`root_com_ang_vel_b`.""" - return self.root_com_ang_vel_b - - @property - def body_pose_w(self) -> torch.Tensor: - """Same as :attr:`body_link_pose_w`.""" - return self.body_link_pose_w - - @property - def body_pos_w(self) -> torch.Tensor: - """Same as :attr:`body_link_pos_w`.""" - return self.body_link_pos_w - - @property - def body_quat_w(self) -> torch.Tensor: - """Same as :attr:`body_link_quat_w`.""" - return self.body_link_quat_w - - @property - def body_vel_w(self) -> torch.Tensor: - """Same as :attr:`body_com_vel_w`.""" - return self.body_com_vel_w - - @property - def body_lin_vel_w(self) -> torch.Tensor: - """Same as :attr:`body_com_lin_vel_w`.""" - return self.body_com_lin_vel_w - - @property - def body_ang_vel_w(self) -> torch.Tensor: - """Same as :attr:`body_com_ang_vel_w`.""" - return self.body_com_ang_vel_w - - @property - def body_acc_w(self) -> torch.Tensor: - """Same as :attr:`body_com_acc_w`.""" - return self.body_com_acc_w - - @property - def body_lin_acc_w(self) -> torch.Tensor: - """Same as :attr:`body_com_lin_acc_w`.""" - return self.body_com_lin_acc_w - - @property - def body_ang_acc_w(self) -> torch.Tensor: - """Same as :attr:`body_com_ang_acc_w`.""" - return self.body_com_ang_acc_w - - @property - def com_pos_b(self) -> torch.Tensor: - """Same as :attr:`body_com_pos_b`.""" - return self.body_com_pos_b - - @property - def com_quat_b(self) -> torch.Tensor: - """Same as :attr:`body_com_quat_b`.""" - return self.body_com_quat_b - - @property - def joint_limits(self) -> torch.Tensor: - """Deprecated property. Please use :attr:`joint_pos_limits` instead.""" - omni.log.warn( - "The `joint_limits` property will be deprecated in a future release. Please use `joint_pos_limits` instead." - ) - return self.joint_pos_limits - - @property - def default_joint_limits(self) -> torch.Tensor: - """Deprecated property. Please use :attr:`default_joint_pos_limits` instead.""" - omni.log.warn( - "The `default_joint_limits` property will be deprecated in a future release. Please use" - " `default_joint_pos_limits` instead." - ) - return self.default_joint_pos_limits - - @property - def joint_velocity_limits(self) -> torch.Tensor: - """Deprecated property. Please use :attr:`joint_vel_limits` instead.""" - omni.log.warn( - "The `joint_velocity_limits` property will be deprecated in a future release. Please use" - " `joint_vel_limits` instead." - ) - return self.joint_vel_limits - - @property - def joint_friction(self) -> torch.Tensor: - """Deprecated property. Please use :attr:`joint_friction_coeff` instead.""" - omni.log.warn( - "The `joint_friction` property will be deprecated in a future release. Please use" - " `joint_friction_coeff` instead." - ) - return self.joint_friction_coeff - - @property - def default_joint_friction(self) -> torch.Tensor: - """Deprecated property. Please use :attr:`default_joint_friction_coeff` instead.""" - omni.log.warn( - "The `default_joint_friction` property will be deprecated in a future release. Please use" - " `default_joint_friction_coeff` instead." - ) - return self.default_joint_friction_coeff - - @property - def fixed_tendon_limit(self) -> torch.Tensor: - """Deprecated property. Please use :attr:`fixed_tendon_pos_limits` instead.""" - omni.log.warn( - "The `fixed_tendon_limit` property will be deprecated in a future release. Please use" - " `fixed_tendon_pos_limits` instead." - ) - return self.fixed_tendon_pos_limits - - @property - def default_fixed_tendon_limit(self) -> torch.Tensor: - """Deprecated property. Please use :attr:`default_fixed_tendon_pos_limits` instead.""" - omni.log.warn( - "The `default_fixed_tendon_limit` property will be deprecated in a future release. Please use" - " `default_fixed_tendon_pos_limits` instead." - ) - return self.default_fixed_tendon_pos_limits + def __new__(cls, *args, **kwargs) -> BaseArticulationData | NewtonArticulationData: + """Create a new instance of an articulation data based on the backend.""" + return super().__new__(cls, *args, **kwargs) \ No newline at end of file diff --git a/source/isaaclab/isaaclab/assets/articulation/base_articulation.py b/source/isaaclab/isaaclab/assets/articulation/base_articulation.py new file mode 100644 index 00000000000..a7a9be3e60c --- /dev/null +++ b/source/isaaclab/isaaclab/assets/articulation/base_articulation.py @@ -0,0 +1,1549 @@ +# Copyright (c) 2022-2025, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +# Flag for pyright to ignore type errors in this file. +# pyright: reportPrivateUsage=false + +from __future__ import annotations + +import torch +import warp as wp +from collections.abc import Sequence +from typing import TYPE_CHECKING +from abc import abstractmethod + +from ..asset_base import AssetBase + +if TYPE_CHECKING: + from .articulation_cfg import ArticulationCfg + from .articulation_data import ArticulationData + + +class BaseArticulation(AssetBase): + """An articulation asset class. + + An articulation is a collection of rigid bodies connected by joints. The joints can be either + fixed or actuated. The joints can be of different types, such as revolute, prismatic, D-6, etc. + However, the articulation class has currently been tested with revolute and prismatic joints. + The class supports both floating-base and fixed-base articulations. The type of articulation + is determined based on the root joint of the articulation. If the root joint is fixed, then + the articulation is considered a fixed-base system. Otherwise, it is considered a floating-base + system. This can be checked using the :attr:`Articulation.is_fixed_base` attribute. + + For an asset to be considered an articulation, the root prim of the asset must have the + `USD ArticulationRootAPI`_. This API is used to define the sub-tree of the articulation using + the reduced coordinate formulation. On playing the simulation, the physics engine parses the + articulation root prim and creates the corresponding articulation in the physics engine. The + articulation root prim can be specified using the :attr:`AssetBaseCfg.prim_path` attribute. + + The articulation class also provides the functionality to augment the simulation of an articulated + system with custom actuator models. These models can either be explicit or implicit, as detailed in + the :mod:`isaaclab.actuators` module. The actuator models are specified using the + :attr:`ArticulationCfg.actuators` attribute. These are then parsed and used to initialize the + corresponding actuator models, when the simulation is played. + + During the simulation step, the articulation class first applies the actuator models to compute + the joint commands based on the user-specified targets. These joint commands are then applied + into the simulation. The joint commands can be either position, velocity, or effort commands. + As an example, the following snippet shows how this can be used for position commands: + + .. code-block:: python + + # an example instance of the articulation class + my_articulation = Articulation(cfg) + + # set joint position targets + my_articulation.set_joint_position_target(position) + # propagate the actuator models and apply the computed commands into the simulation + my_articulation.write_data_to_sim() + + # step the simulation using the simulation context + sim_context.step() + + # update the articulation state, where dt is the simulation time step + my_articulation.update(dt) + + .. _`USD ArticulationRootAPI`: https://openusd.org/dev/api/class_usd_physics_articulation_root_a_p_i.html + + """ + + cfg: ArticulationCfg + """Configuration instance for the articulations.""" + + __backend_name__: str = "base" + """The name of the backend for the articulation.""" + + actuators: dict + """Dictionary of actuator instances for the articulation. + + The keys are the actuator names and the values are the actuator instances. The actuator instances + are initialized based on the actuator configurations specified in the :attr:`ArticulationCfg.actuators` + attribute. They are used to compute the joint commands during the :meth:`write_data_to_sim` function. + """ + + def __init__(self, cfg: ArticulationCfg): + """Initialize the articulation. + + Args: + cfg: A configuration instance. + """ + super().__init__(cfg) + + """ + Properties + """ + + @property + @abstractmethod + def data(self) -> ArticulationData: + raise NotImplementedError() + + @property + @abstractmethod + def num_instances(self) -> int: + raise NotImplementedError() + + @property + @abstractmethod + def is_fixed_base(self) -> bool: + """Whether the articulation is a fixed-base or floating-base system.""" + raise NotImplementedError() + + @property + @abstractmethod + def num_joints(self) -> int: + """Number of joints in articulation.""" + raise NotImplementedError() + + @property + @abstractmethod + def num_fixed_tendons(self) -> int: + """Number of fixed tendons in articulation.""" + raise NotImplementedError() + + @property + @abstractmethod + def num_spatial_tendons(self) -> int: + """Number of spatial tendons in articulation.""" + raise NotImplementedError() + + @property + @abstractmethod + def num_bodies(self) -> int: + """Number of bodies in articulation.""" + raise NotImplementedError() + + @property + @abstractmethod + def joint_names(self) -> list[str]: + """Ordered names of joints in articulation.""" + raise NotImplementedError() + + @property + @abstractmethod + def fixed_tendon_names(self) -> list[str]: + """Ordered names of fixed tendons in articulation.""" + raise NotImplementedError() + + @property + @abstractmethod + def spatial_tendon_names(self) -> list[str]: + """Ordered names of spatial tendons in articulation.""" + raise NotImplementedError() + + @property + @abstractmethod + def body_names(self) -> list[str]: + """Ordered names of bodies in articulation.""" + raise NotImplementedError() + + @property + @abstractmethod + def root_view(self): + """Root view for the asset. + + Note: + Use this view with caution. It requires handling of tensors in a specific way. + """ + raise NotImplementedError() + + """ + Operations. + """ + + @abstractmethod + def reset(self, env_ids: Sequence[int] | None = None, mask: wp.array | torch.Tensor | None = None): + raise NotImplementedError() + + @abstractmethod + def write_data_to_sim(self): + """Write external wrenches and joint commands to the simulation. + + If any explicit actuators are present, then the actuator models are used to compute the + joint commands. Otherwise, the joint commands are directly set into the simulation. + + Note: + We write external wrench to the simulation here since this function is called before the simulation step. + This ensures that the external wrench is applied at every simulation step. + """ + raise NotImplementedError() + + @abstractmethod + def update(self, dt: float): + raise NotImplementedError() + + """ + Operations - Finders. + """ + + @abstractmethod + def find_bodies( + self, name_keys: str | Sequence[str], preserve_order: bool = False + ) -> tuple[torch.Tensor | wp.array, list[int], list[str]]: + """Find bodies in the articulation based on the name keys. + + Please check the :meth:`isaaclab.utils.string_utils.resolve_matching_names` function for more + information on the name matching. + + Args: + name_keys: A regular expression or a list of regular expressions to match the body names. + preserve_order: Whether to preserve the order of the name keys in the output. Defaults to False. + + Returns: + A tuple of lists containing the body indices and names. + """ + raise NotImplementedError() + + @abstractmethod + def find_joints( + self, name_keys: str | Sequence[str], joint_subset: list[str] | None = None, preserve_order: bool = False + ) -> tuple[torch.Tensor | wp.array, list[int], list[str]]: + """Find joints in the articulation based on the name keys. + + Please see the :func:`isaaclab.utils.string.resolve_matching_names` function for more information + on the name matching. + + Args: + name_keys: A regular expression or a list of regular expressions to match the joint names. + joint_subset: A subset of joints to search for. Defaults to None, which means all joints + in the articulation are searched. + preserve_order: Whether to preserve the order of the name keys in the output. Defaults to False. + + Returns: + A tuple of lists containing the joint mask, names, and indices. + """ + raise NotImplementedError() + + @abstractmethod + def find_fixed_tendons( + self, name_keys: str | Sequence[str], tendon_subsets: list[str] | None = None, preserve_order: bool = False + ) -> tuple[torch.Tensor | wp.array, list[int], list[str]]: + """Find fixed tendons in the articulation based on the name keys. + + Please see the :func:`isaaclab.utils.string.resolve_matching_names` function for more information + on the name matching. + + Args: + name_keys: A regular expression or a list of regular expressions to match the joint + names with fixed tendons. + tendon_subsets: A subset of joints with fixed tendons to search for. Defaults to None, which means + all joints in the articulation are searched. + preserve_order: Whether to preserve the order of the name keys in the output. Defaults to False. + + Returns: + A tuple of lists containing the tendon mask, names, and indices. + """ + raise NotImplementedError() + + @abstractmethod + def find_spatial_tendons( + self, name_keys: str | Sequence[str], tendon_subsets: list[str] | None = None, preserve_order: bool = False + ) -> tuple[torch.Tensor | wp.array, list[int], list[str]]: + """Find spatial tendons in the articulation based on the name keys. + + Please see the :func:`isaaclab.utils.string.resolve_matching_names` function for more information + on the name matching. + + Args: + name_keys: A regular expression or a list of regular expressions to match the tendon names. + tendon_subsets: A subset of tendons to search for. Defaults to None, which means all tendons + in the articulation are searched. + preserve_order: Whether to preserve the order of the name keys in the output. Defaults to False. + + Returns: + A tuple of lists containing the tendon mask, names, and indices. + """ + raise NotImplementedError() + + """ + Operations - State Writers. + """ + + @abstractmethod + def write_root_state_to_sim( + self, + root_state: torch.Tensor | wp.array, + env_ids: Sequence[int] | None = None, + env_mask: torch.Tensor | wp.array | None = None + ) -> None: + """Set the root state over selected environment indices into the simulation. + + The root state comprises of the cartesian position, quaternion orientation in (w, x, y, z), and linear + and angular velocity. All the quantities are in the simulation frame. + + ..note:: When ids are provided, then partial data is expected. When masks are provided the whole of the data + is expected. For example, if env_ids is provided, then root_state should be of shape (len(env_ids), 13). If + env_mask is provided, then root_state should be of shape (num_instances, 13). + + ..caution:: If both env_mask and env_ids are provided, then env_mask will be used. The function will thus + expect the whole of the data to be provided. If none of them are provided, then the function expects the whole + of the data to be provided. + + Args: + root_state: Root state in simulation frame. Shape is (len(env_ids), 13) or (num_instances, 13). + env_ids: Environment indices. If None, then all indices are used. + env_mask: Environment mask. Shape is (num_instances,). + """ + raise NotImplementedError() + + @abstractmethod + def write_root_com_state_to_sim( + self, + root_state: torch.Tensor | wp.array, + env_ids: Sequence[int] | None = None, + env_mask: torch.Tensor | wp.array | None = None + ) -> None: + """Set the root center of mass state over selected environment indices into the simulation. + + The root state comprises of the cartesian position, quaternion orientation in (w, x, y, z), and linear + and angular velocity. All the quantities are in the simulation frame. + + ..note:: When ids are provided, then partial data is expected. When masks are provided the whole of the data + is expected. For example, if env_ids is provided, then root_state should be of shape (len(env_ids), 13). If + env_mask is provided, then root_state should be of shape (num_instances, 13). + + ..caution:: If both env_mask and env_ids are provided, then env_mask will be used. The function will thus + expect the whole of the data to be provided. If none of them are provided, then the function expects the whole + of the data to be provided. + + Args: + root_state: Root state in simulation frame. Shape is (len(env_ids), 13) or (num_instances, 13). + env_ids: Environment indices. If None, then all indices are used. + env_mask: Environment mask. Shape is (num_instances,). + """ + raise NotImplementedError() + + @abstractmethod + def write_root_link_state_to_sim( + self, + root_state: torch.Tensor | wp.array, + env_ids: Sequence[int] | None = None, + env_mask: torch.Tensor | wp.array | None = None + ) -> None: + """Set the root link state over selected environment indices into the simulation. + + The root state comprises of the cartesian position, quaternion orientation in (w, x, y, z), and linear + and angular velocity. All the quantities are in the simulation frame. + + ..note:: When ids are provided, then partial data is expected. When masks are provided the whole of the data + is expected. For example, if env_ids is provided, then root_state should be of shape (len(env_ids), 13). If + env_mask is provided, then root_state should be of shape (num_instances, 13). + + ..caution:: If both env_mask and env_ids are provided, then env_mask will be used. The function will thus + expect the whole of the data to be provided. If none of them are provided, then the function expects the whole + of the data to be provided. + + Args: + root_state: Root state in simulation frame. Shape is (len(env_ids), 13) or (num_instances, 13). + env_ids: Environment indices. If None, then all indices are used. + """ + raise NotImplementedError() + + @abstractmethod + def write_root_pose_to_sim( + self, + root_pose: torch.Tensor | wp.array, + env_ids: Sequence[int] | None = None, + env_mask: torch.Tensor | wp.array | None = None + ) -> None: + """Set the root pose over selected environment indices into the simulation. + + The root pose comprises of the cartesian position and quaternion orientation in (w, x, y, z). + + ..note:: When ids are provided, then partial data is expected. When masks are provided the whole of the data + is expected. For example, if env_ids is provided, then root_pose should be of shape (len(env_ids), 7). If + env_mask is provided, then root_pose should be of shape (num_instances, 7). + + ..caution:: If both env_mask and env_ids are provided, then env_mask will be used. The function will thus + expect the whole of the data to be provided. If none of them are provided, then the function expects the whole + of the data to be provided. + + Args: + root_pose: Root poses in simulation frame. Shape is (len(env_ids), 7) or (num_instances, 7). + env_ids: Environment indices. If None, then all indices are used. + env_mask: Environment mask. Shape is (num_instances,). + """ + raise NotImplementedError() + + @abstractmethod + def write_root_link_pose_to_sim( + self, + root_pose: torch.Tensor | wp.array, + env_ids: Sequence[int] | None = None, + env_mask: torch.Tensor | wp.array | None = None + ) -> None: + """Set the root link pose over selected environment indices into the simulation. + + The root pose comprises of the cartesian position and quaternion orientation in (w, x, y, z). + + ..note:: When ids are provided, then partial data is expected. When masks are provided the whole of the data + is expected. For example, if env_ids is provided, then root_pose should be of shape (len(env_ids), 7). If + env_mask is provided, then root_pose should be of shape (num_instances, 7). + + ..caution:: If both env_mask and env_ids are provided, then env_mask will be used. The function will thus + expect the whole of the data to be provided. If none of them are provided, then the function expects the whole + of the data to be provided. + + Args: + root_pose: Root poses in simulation frame. Shape is (len(env_ids), 7) or (num_instances, 7). + env_ids: Environment indices. If None, then all indices are used. + env_mask: Environment mask. Shape is (num_instances,). + """ + raise NotImplementedError() + + @abstractmethod + def write_root_com_pose_to_sim( + self, + root_pose: torch.Tensor | wp.array, + env_ids: Sequence[int] | None = None, + env_mask: torch.Tensor | wp.array | None = None + ) -> None: + """Set the root center of mass pose over selected environment indices into the simulation. + + The root pose comprises of the cartesian position and quaternion orientation in (w, x, y, z). + The orientation is the orientation of the principle axes of inertia. + + ..note:: When ids are provided, then partial data is expected. When masks are provided the whole of the data + is expected. For example, if env_ids is provided, then root_pose should be of shape (len(env_ids), 7). If + env_mask is provided, then root_pose should be of shape (num_instances, 7). + + ..caution:: If both env_mask and env_ids are provided, then env_mask will be used. The function will thus + expect the whole of the data to be provided. If none of them are provided, then the function expects the whole + of the data to be provided. + + Args: + root_pose: Root center of mass poses in simulation frame. Shape is (len(env_ids), 7) or (num_instances, 7). + env_ids: Environment indices. If None, then all indices are used. + env_mask: Environment mask. Shape is (num_instances,). + """ + raise NotImplementedError() + + @abstractmethod + def write_root_velocity_to_sim( + self, + root_velocity: torch.Tensor | wp.array, + env_ids: Sequence[int] | None = None, + env_mask: torch.Tensor | wp.array | None = None + ) -> None: + """Set the root center of mass velocity over selected environment indices into the simulation. + + The velocity comprises linear velocity (x, y, z) and angular velocity (x, y, z) in that order. + ..note:: This sets the velocity of the root's center of mass rather than the roots frame. + + ..note:: When ids are provided, then partial data is expected. When masks are provided the whole of the data + is expected. For example, if env_ids is provided, then root_velocity should be of shape (len(env_ids), 6). If + env_mask is provided, then root_velocity should be of shape (num_instances, 6). + + ..caution:: If both env_mask and env_ids are provided, then env_mask will be used. The function will thus + expect the whole of the data to be provided. If none of them are provided, then the function expects the whole + of the data to be provided. + + Args: + root_velocity: Root center of mass velocities in simulation world frame. Shape is (len(env_ids), 6) or (num_instances, 6). + env_ids: Environment indices. If None, then all indices are used. + env_mask: Environment mask. Shape is (num_instances,). + """ + raise NotImplementedError() + + @abstractmethod + def write_root_com_velocity_to_sim( + self, + root_velocity: torch.Tensor | wp.array, + env_ids: Sequence[int] | None = None, + env_mask: torch.Tensor | wp.array | None = None + ) -> None: + """Set the root center of mass velocity over selected environment indices into the simulation. + + The velocity comprises linear velocity (x, y, z) and angular velocity (x, y, z) in that order. + ..note:: This sets the velocity of the root's center of mass rather than the roots frame. + + ..note:: When ids are provided, then partial data is expected. When masks are provided the whole of the data + is expected. For example, if env_ids is provided, then root_velocity should be of shape (len(env_ids), 6). If + env_mask is provided, then root_velocity should be of shape (num_instances, 6). + + ..caution:: If both env_mask and env_ids are provided, then env_mask will be used. The function will thus + expect the whole of the data to be provided. If none of them are provided, then the function expects the whole + of the data to be provided. + + Args: + root_velocity: Root center of mass velocities in simulation world frame. Shape is (len(env_ids), 6) or (num_instances, 6). + env_ids: Environment indices. If None, then all indices are used. + env_mask: Environment mask. Shape is (num_instances,). + """ + raise NotImplementedError() + + @abstractmethod + def write_root_link_velocity_to_sim( + self, + root_velocity: torch.Tensor | wp.array, + env_ids: Sequence[int] | None = None, + env_mask: torch.Tensor | wp.array | None = None + ) -> None: + """Set the root link velocity over selected environment indices into the simulation. + + The velocity comprises linear velocity (x, y, z) and angular velocity (x, y, z) in that order. + ..note:: This sets the velocity of the root's frame rather than the roots center of mass. + + ..note:: When ids are provided, then partial data is expected. When masks are provided the whole of the data + is expected. For example, if env_ids is provided, then root_velocity should be of shape (len(env_ids), 6). If + env_mask is provided, then root_velocity should be of shape (num_instances, 6). + + ..caution:: If both env_mask and env_ids are provided, then env_mask will be used. The function will thus + expect the whole of the data to be provided. If none of them are provided, then the function expects the whole + of the data to be provided. + + Args: + root_velocity: Root frame velocities in simulation world frame. Shape is (len(env_ids), 6) or (num_instances, 6). + env_ids: Environment indices. If None, then all indices are used. + env_mask: Environment mask. Shape is (num_instances,). + """ + raise NotImplementedError() + + @abstractmethod + def write_joint_state_to_sim( + self, + position: torch.Tensor | wp.array, + velocity: torch.Tensor | wp.array, + joint_ids: Sequence[int] | slice | None = None, + env_ids: Sequence[int] | slice | None = None, + joint_mask: torch.Tensor | wp.array | None = None, + env_mask: torch.Tensor | wp.array | None = None, + ): + """Write joint positions and velocities to the simulation. + + ..note:: When ids are provided, then partial data is expected. When masks are provided the whole of the data + is expected. For example, if env_ids and joint_ids are provided, then position should be of shape + (len(env_ids), len(joint_ids)). If env_mask is provided, then position should be of shape + (num_instances, num_joints). + + ..caution:: If both env_mask and env_ids are provided, then env_mask will be used. The function will thus + expect the whole of the data to be provided. If none of them are provided, then the function expects the whole + of the data to be provided. + + ..caution:: Do not mix ids and masks. If a mask is provided, then ids will be ignored. + + Args: + position: Joint positions. Shape is (len(env_ids), len(joint_ids)) or (num_instances, num_joints). + velocity: Joint velocities. Shape is (len(env_ids), len(joint_ids)) or (num_instances, num_joints). + joint_ids: The joint indices to set the targets for. Defaults to None (all joints). + env_ids: The environment indices to set the targets for. Defaults to None (all environments). + joint_mask: The joint mask. Shape is (num_joints). + env_mask: The environment mask. Shape is (num_instances,). + """ + raise NotImplementedError() + + @abstractmethod + def write_joint_position_to_sim( + self, + position: torch.Tensor | wp.array, + joint_ids: Sequence[int] | slice | None = None, + env_ids: Sequence[int] | slice | None = None, + joint_mask: torch.Tensor | wp.array | None = None, + env_mask: torch.Tensor | wp.array | None = None, + ): + """Write joint positions to the simulation. + + ..note:: When ids are provided, then partial data is expected. When masks are provided the whole of the data + is expected. For example, if env_ids and joint_ids are provided, then position should be of shape + (len(env_ids), len(joint_ids)). If env_mask is provided, then position should be of shape + (num_instances, num_joints). + + ..caution:: If both env_mask and env_ids are provided, then env_mask will be used. The function will thus + expect the whole of the data to be provided. If none of them are provided, then the function expects the whole + of the data to be provided. + + ..caution:: Do not mix ids and masks. If a mask is provided, then ids will be ignored. + + Args: + position: Joint positions. Shape is (len(env_ids), len(joint_ids)) or (num_instances, num_joints). + joint_ids: The joint indices to set the targets for. Defaults to None (all joints). + env_ids: The environment indices to set the targets for. Defaults to None (all environments). + joint_mask: The joint mask. Shape is (num_joints). + env_mask: The environment mask. Shape is (num_instances,). + """ + raise NotImplementedError() + + @abstractmethod + def write_joint_velocity_to_sim( + self, + velocity: torch.Tensor | wp.array, + joint_ids: Sequence[int] | slice | None = None, + env_ids: Sequence[int] | slice | None = None, + joint_mask: torch.Tensor | wp.array | None = None, + env_mask: torch.Tensor | wp.array | None = None, + ): + """Write joint velocities to the simulation. + + ..note:: When ids are provided, then partial data is expected. When masks are provided the whole of the data + is expected. For example, if env_ids and joint_ids are provided, then velocity should be of shape + (len(env_ids), len(joint_ids)). If env_mask is provided, then velocity should be of shape + (num_instances, num_joints). + + ..caution:: If both env_mask and env_ids are provided, then env_mask will be used. The function will thus + expect the whole of the data to be provided. If none of them are provided, then the function expects the whole + of the data to be provided. + + ..caution:: Do not mix ids and masks. If a mask is provided, then ids will be ignored. + + Args: + velocity: Joint velocities. Shape is (len(env_ids), len(joint_ids)) or (num_instances, num_joints). + joint_ids: The joint indices to set the targets for. Defaults to None (all joints). + env_ids: The environment indices to set the targets for. Defaults to None (all environments). + joint_mask: The joint mask. Shape is (num_joints). + env_mask: The environment mask. Shape is (num_instances,). + """ + raise NotImplementedError() + + """ + Operations - Simulation Parameters Writers. + """ + + @abstractmethod + def write_joint_stiffness_to_sim( + self, + stiffness: 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, + ): + """Write joint stiffness into the simulation. + + ..note:: When ids are provided, then partial data is expected. When masks are provided the whole of the data + is expected. For example, if env_ids and joint_ids are provided, then stiffness should be of shape + (len(env_ids), len(joint_ids)). If env_mask is provided, then stiffness should be of shape + (num_instances, num_joints). + + ..caution:: If both env_mask and env_ids are provided, then env_mask will be used. The function will thus + expect the whole of the data to be provided. If none of them are provided, then the function expects the whole + of the data to be provided. + + ..caution:: Do not mix ids and masks. If a mask is provided, then ids will be ignored. + + Args: + stiffness: Joint stiffness. Shape is (len(env_ids), len(joint_ids)) or (num_instances, num_joints). + joint_ids: The joint indices to set the stiffness for. Defaults to None (all joints). + env_ids: The environment indices to set the stiffness for. Defaults to None (all environments). + joint_mask: The joint mask. Shape is (num_joints). + env_mask: The environment mask. Shape is (num_instances,). + """ + # note: This function isn't setting the values for actuator models. (#128) + raise NotImplementedError() + + @abstractmethod + def write_joint_damping_to_sim( + self, + damping: 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, + ): + """Write joint damping into the simulation. + + ..note:: When ids are provided, then partial data is expected. When masks are provided the whole of the data + is expected. For example, if env_ids and joint_ids are provided, then damping should be of shape + (len(env_ids), len(joint_ids)). If env_mask is provided, then damping should be of shape + (num_instances, num_joints). + + ..caution:: If both env_mask and env_ids are provided, then env_mask will be used. The function will thus + expect the whole of the data to be provided. If none of them are provided, then the function expects the whole + of the data to be provided. + + ..caution:: Do not mix ids and masks. If a mask is provided, then ids will be ignored. + + Args: + damping: Joint damping. Shape is (len(env_ids), len(joint_ids)) or (num_instances, num_joints). + joint_ids: The joint indices to set the damping for. Defaults to None (all joints). + env_ids: The environment indices to set the damping for. Defaults to None (all environments). + joint_mask: The joint mask. Shape is (num_joints). + env_mask: The environment mask. Shape is (num_instances,). + """ + raise NotImplementedError() + + @abstractmethod + def write_joint_position_limit_to_sim( + self, + limits: 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, + warn_limit_violation: bool = True, + ): + """Write joint position limits into the simulation. + + ..note:: When ids are provided, then partial data is expected. When masks are provided the whole of the data + is expected. For example, if env_ids and joint_ids are provided, then limits should be of shape + (len(env_ids), len(joint_ids), 2). If env_mask is provided, then limits should be of shape + (num_instances, num_joints, 2). + + ..caution:: If both env_mask and env_ids are provided, then env_mask will be used. The function will thus + expect the whole of the data to be provided. If none of them are provided, then the function expects the whole + of the data to be provided. + + ..caution:: Do not mix ids and masks. If a mask is provided, then ids will be ignored. + + Args: + limits: Joint limits. Shape is (len(env_ids), len(joint_ids), 2) or (num_instances, num_joints, 2). + joint_ids: The joint indices to set the limits for. Defaults to None (all joints). + env_ids: The environment indices to set the limits for. Defaults to None (all environments). + joint_mask: The joint mask. Shape is (num_joints). + env_mask: The environment mask. Shape is (num_instances,). + warn_limit_violation: Whether to use warning or info level logging when default joint positions + exceed the new limits. Defaults to True. + """ + raise NotImplementedError() + + @abstractmethod + def write_joint_velocity_limit_to_sim( + self, + limits: 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, + ): + """Write joint max velocity to the simulation. + + The velocity limit is used to constrain the joint velocities in the physics engine. The joint will only + be able to reach this velocity if the joint's effort limit is sufficiently large. If the joint is moving + faster than this velocity, the physics engine will actually try to brake the joint to reach this velocity. + + ..note:: When ids are provided, then partial data is expected. When masks are provided the whole of the data + is expected. For example, if env_ids and joint_ids are provided, then limits should be of shape + (len(env_ids), len(joint_ids)). If env_mask is provided, then limits should be of shape + (num_instances, num_joints). + + ..caution:: If both env_mask and env_ids are provided, then env_mask will be used. The function will thus + expect the whole of the data to be provided. If none of them are provided, then the function expects the whole + of the data to be provided. + + ..caution:: Do not mix ids and masks. If a mask is provided, then ids will be ignored. + + Args: + limits: Joint max velocity. Shape is (len(env_ids), len(joint_ids)) or (num_instances, num_joints). + joint_ids: The joint indices to set the max velocity for. Defaults to None (all joints). + env_ids: The environment indices to set the max velocity for. Defaults to None (all environments). + joint_mask: The joint mask. Shape is (num_joints). + env_mask: The environment mask. Shape is (num_instances,). + """ + raise NotImplementedError() + + @abstractmethod + def write_joint_effort_limit_to_sim( + self, + limits: 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, + ): + """Write joint effort limits into the simulation. + + The effort limit is used to constrain the computed joint efforts in the physics engine. If the + computed effort exceeds this limit, the physics engine will clip the effort to this value. + + ..note:: When ids are provided, then partial data is expected. When masks are provided the whole of the data + is expected. For example, if env_ids and joint_ids are provided, then limits should be of shape + (len(env_ids), len(joint_ids)). If env_mask is provided, then limits should be of shape + (num_instances, num_joints). + + ..caution:: If both env_mask and env_ids are provided, then env_mask will be used. The function will thus + expect the whole of the data to be provided. If none of them are provided, then the function expects the whole + of the data to be provided. + + ..caution:: Do not mix ids and masks. If a mask is provided, then ids will be ignored. + + Args: + limits: Joint torque limits. 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). + env_mask: The environment mask. Shape is (num_instances,). + """ + raise NotImplementedError() + + @abstractmethod + def write_joint_armature_to_sim( + self, + armature: 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, + ): + """Write joint armature into the simulation. + + The armature is directly added to the corresponding joint-space inertia. It helps improve the + simulation stability by reducing the joint velocities. + + ..note:: When ids are provided, then partial data is expected. When masks are provided the whole of the data + is expected. For example, if env_ids and joint_ids are provided, then armature should be of shape + (len(env_ids), len(joint_ids)). If env_mask is provided, then armature should be of shape + (num_instances, num_joints). + + ..caution:: If both env_mask and env_ids are provided, then env_mask will be used. The function will thus + expect the whole of the data to be provided. If none of them are provided, then the function expects the whole + of the data to be provided. + + ..caution:: Do not mix ids and masks. If a mask is provided, then ids will be ignored. + + Args: + armature: Joint armature. 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). + env_mask: The environment mask. Shape is (num_instances,). + """ + raise NotImplementedError() + + @abstractmethod + def write_joint_friction_coefficient_to_sim( + self, + joint_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, + ): + r"""Write joint static friction coefficients into the simulation. + + The joint static friction is a unitless quantity. It relates the magnitude of the spatial force transmitted + from the parent body to the child body to the maximal static friction force that may be applied by the solver + to resist the joint motion. + + Mathematically, this means that: :math:`F_{resist} \leq \mu F_{spatial}`, where :math:`F_{resist}` + is the resisting force applied by the solver and :math:`F_{spatial}` is the spatial force + transmitted from the parent body to the child body. The simulated static friction effect is therefore + similar to static and Coulomb static friction. + + ..note:: When ids are provided, then partial data is expected. When masks are provided the whole of the data + is expected. For example, if env_ids and joint_ids are provided, then joint_friction_coeff should be of shape + (len(env_ids), len(joint_ids)). If env_mask is provided, then joint_friction_coeff should be of shape + (num_instances, num_joints). + + ..caution:: If both env_mask and env_ids are provided, then env_mask will be used. The function will thus + expect the whole of the data to be provided. If none of them are provided, then the function expects the whole + of the data to be provided. + + ..caution:: Do not mix ids and masks. If a mask is provided, then ids will be ignored. + + Args: + joint_friction_coeff: Joint static 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). + env_mask: The environment mask. Shape is (num_instances,). + """ + raise NotImplementedError() + + @abstractmethod + def write_joint_dynamic_friction_coefficient_to_sim( + self, + joint_dynamic_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. + """ + + @abstractmethod + def set_external_force_and_torque( + self, + forces: torch.Tensor | wp.array, + torques: torch.Tensor | wp.array, + positions: torch.Tensor | None = None, + body_ids: Sequence[int] | slice | None = None, + env_ids: Sequence[int] | None = None, + body_mask: torch.Tensor | wp.array | None = None, + env_mask: torch.Tensor | wp.array | None = None, + is_global: bool = False, + ): + """Set external force and torque to apply on the asset's bodies in their local frame. + + For many applications, we want to keep the applied external force on rigid bodies constant over a period of + time (for instance, during the policy control). This function allows us to store the external force and torque + into buffers which are then applied to the simulation at every step. Optionally, set the position to apply the + external wrench at (in the local link frame of the bodies). + + .. caution:: + If the function is called with empty forces and torques, then this function disables the application + of external wrench to the simulation. + + .. code-block:: python + + # example of disabling external wrench + asset.set_external_force_and_torque(forces=torch.zeros(0, 3), torques=torch.zeros(0, 3)) + + .. caution:: + If the function is called consecutively with and with different values for ``is_global``, then the + all the external wrenches will be applied in the frame specified by the last call. + + .. code-block:: python + + # example of setting external wrench in the global frame + asset.set_external_force_and_torque(forces=torch.ones(1, 1, 3), env_ids=[0], is_global=True) + # example of setting external wrench in the link frame + asset.set_external_force_and_torque(forces=torch.ones(1, 1, 3), env_ids=[1], is_global=False) + # Both environments will have the external wrenches applied in the link frame + + .. note:: + This function does not apply the external wrench to the simulation. It only fills the buffers with + the desired values. To apply the external wrench, call the :meth:`write_data_to_sim` function + right before the simulation step. + + ..note:: When ids are provided, then partial data is expected. When masks are provided the whole of the data + is expected. For example, if env_ids is provided, then root_state should be of shape (len(env_ids), 13). If + env_mask is provided, then root_state should be of shape (num_instances, 13). + + ..caution:: If both env_mask and env_ids are provided, then env_mask will be used. The function will thus + expect the whole of the data to be provided. If none of them are provided, then the function expects the whole + of the data to be provided. + + ..caution:: Do not mix ids and masks. If a mask is provided, then ids will be ignored. + + Args: + forces: External forces in bodies' local frame. Shape is (len(env_ids), len(body_ids), 3) or (num_instances, num_bodies, 3). + torques: External torques in bodies' local frame. Shape is (len(env_ids), len(body_ids), 3) or (num_instances, num_bodies, 3). + positions: Positions to apply external wrench. Shape is (len(env_ids), len(body_ids), 3) or (num_instances, num_bodies, 3). Defaults to None. + body_ids: Body indices to apply external wrench to. Defaults to None (all bodies). + env_ids: Environment indices to apply external wrench to. Defaults to None (all instances). + body_mask: The body mask. Shape is (num_bodies). + env_mask: The environment mask. Shape is (num_instances,). + is_global: Whether to apply the external wrench in the global frame. Defaults to False. If set to False, + the external wrench is applied in the link frame of the articulations' bodies. + """ + raise NotImplementedError() + + @abstractmethod + def set_joint_position_target( + self, + target: torch.Tensor | wp.array, + 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, + ): + """Set joint position targets into internal buffers. + + This function does not apply the joint targets to the simulation. It only fills the buffers with + the desired values. To apply the joint targets, call the :meth:`write_data_to_sim` function. + + ..note:: When ids are provided, then partial data is expected. When masks are provided the whole of the data + is expected. For example, if env_ids is provided, then root_state should be of shape (len(env_ids), 13). If + env_mask is provided, then root_state should be of shape (num_instances, 13). + + ..caution:: If both env_mask and env_ids are provided, then env_mask will be used. The function will thus + expect the whole of the data to be provided. If none of them are provided, then the function expects the whole + of the data to be provided. + + ..caution:: Do not mix ids and masks. If a mask is provided, then ids will be ignored. + + Args: + target: Joint position targets. Shape is (len(env_ids), len(joint_ids)). + joint_ids: The joint indices to set the targets for. Defaults to None (all joints). + env_ids: The environment indices to set the targets for. Defaults to None (all environments). + """ + raise NotImplementedError() + + @abstractmethod + def set_joint_velocity_target( + self, + target: torch.Tensor | wp.array, + 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, + ): + """Set joint velocity targets into internal buffers. + + This function does not apply the joint targets to the simulation. It only fills the buffers with + the desired values. To apply the joint targets, call the :meth:`write_data_to_sim` function. + + ..note:: When ids are provided, then partial data is expected. When masks are provided the whole of the data + is expected. For example, if env_ids is provided, then root_state should be of shape (len(env_ids), 13). If + env_mask is provided, then root_state should be of shape (num_instances, 13). + + ..caution:: If both env_mask and env_ids are provided, then env_mask will be used. The function will thus + expect the whole of the data to be provided. If none of them are provided, then the function expects the whole + of the data to be provided. + + ..caution:: Do not mix ids and masks. If a mask is provided, then ids will be ignored. + + Args: + target: Joint velocity targets. Shape is (len(env_ids), len(joint_ids)) or (num_instances, num_joints). + joint_ids: The joint indices to set the targets for. Defaults to None (all joints). + env_ids: The environment indices to set the targets for. Defaults to None (all environments). + joint_mask: The joint mask. Shape is (num_joints). + env_mask: The environment mask. Shape is (num_instances,). + """ + raise NotImplementedError() + + @abstractmethod + def set_joint_effort_target( + self, + target: torch.Tensor | wp.array, + 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, + ): + """Set joint efforts into internal buffers. + + This function does not apply the joint targets to the simulation. It only fills the buffers with + the desired values. To apply the joint targets, call the :meth:`write_data_to_sim` function. + + ..note:: When ids are provided, then partial data is expected. When masks are provided the whole of the data + is expected. For example, if env_ids is provided, then root_state should be of shape (len(env_ids), 13). If + env_mask is provided, then root_state should be of shape (num_instances, 13). + + ..caution:: If both env_mask and env_ids are provided, then env_mask will be used. The function will thus + expect the whole of the data to be provided. If none of them are provided, then the function expects the whole + of the data to be provided. + + ..caution:: Do not mix ids and masks. If a mask is provided, then ids will be ignored. + + Args: + target: Joint effort targets. Shape is (len(env_ids), len(joint_ids)) or (num_instances, num_joints). + joint_ids: The joint indices to set the targets for. Defaults to None (all joints). + env_ids: The environment indices to set the targets for. Defaults to None (all environments). + joint_mask: The joint mask. Shape is (num_joints). + env_mask: The environment mask. Shape is (num_instances,). + """ + raise NotImplementedError() + + """ + Operations - Tendons. + """ + + @abstractmethod + def set_fixed_tendon_stiffness( + self, + stiffness: torch.Tensor | wp.array, + fixed_tendon_ids: Sequence[int] | slice | None = None, + env_ids: Sequence[int] | None = None, + fixed_tendon_mask: torch.Tensor | wp.array | None = None, + env_mask: torch.Tensor | wp.array | None = None, + ): + """Set fixed tendon stiffness into internal buffers. + + This function does not apply the tendon stiffness to the simulation. It only fills the buffers with + the desired values. To apply the tendon stiffness, call the :meth:`write_fixed_tendon_properties_to_sim` function. + + ..note:: When ids are provided, then partial data is expected. When masks are provided the whole of the data + is expected. For example, if env_ids is provided, then root_state should be of shape (len(env_ids), 13). If + env_mask is provided, then root_state should be of shape (num_instances, 13). + + ..caution:: If both env_mask and env_ids are provided, then env_mask will be used. The function will thus + expect the whole of the data to be provided. If none of them are provided, then the function expects the whole + of the data to be provided. + + ..caution:: Do not mix ids and masks. If a mask is provided, then ids will be ignored. + + Args: + stiffness: Fixed tendon stiffness. Shape is (len(env_ids), len(fixed_tendon_ids)) or (num_instances, num_fixed_tendons). + fixed_tendon_ids: The tendon indices to set the stiffness for. Defaults to None (all fixed tendons). + env_ids: The environment indices to set the stiffness for. Defaults to None (all environments). + fixed_tendon_mask: The fixed tendon mask. Shape is (num_fixed_tendons). + env_mask: The environment mask. Shape is (num_instances,). + """ + raise NotImplementedError() + + @abstractmethod + def set_fixed_tendon_damping( + self, + damping: torch.Tensor | wp.array, + fixed_tendon_ids: Sequence[int] | slice | None = None, + env_ids: Sequence[int] | None = None, + fixed_tendon_mask: torch.Tensor | wp.array | None = None, + env_mask: torch.Tensor | wp.array | None = None, + ): + """Set fixed tendon damping into internal buffers. + + This function does not apply the tendon damping to the simulation. It only fills the buffers with + the desired values. To apply the tendon damping, call the :meth:`write_fixed_tendon_properties_to_sim` function. + + ..note:: When ids are provided, then partial data is expected. When masks are provided the whole of the data + is expected. For example, if env_ids is provided, then root_state should be of shape (len(env_ids), 13). If + env_mask is provided, then root_state should be of shape (num_instances, 13). + + ..caution:: If both env_mask and env_ids are provided, then env_mask will be used. The function will thus + expect the whole of the data to be provided. If none of them are provided, then the function expects the whole + of the data to be provided. + + ..caution:: Do not mix ids and masks. If a mask is provided, then ids will be ignored. + + Args: + damping: Fixed tendon damping. Shape is (len(env_ids), len(fixed_tendon_ids)) or (num_instances, num_fixed_tendons). + fixed_tendon_ids: The tendon indices to set the damping for. Defaults to None (all fixed tendons). + env_ids: The environment indices to set the damping for. Defaults to None (all environments). + fixed_tendon_mask: The fixed tendon mask. Shape is (num_fixed_tendons). + env_mask: The environment mask. Shape is (num_instances,). + """ + raise NotImplementedError() + + @abstractmethod + def set_fixed_tendon_limit_stiffness( + self, + limit_stiffness: torch.Tensor | wp.array, + fixed_tendon_ids: Sequence[int] | slice | None = None, + env_ids: Sequence[int] | None = None, + fixed_tendon_mask: torch.Tensor | wp.array | None = None, + env_mask: torch.Tensor | wp.array | None = None, + ): + """Set fixed tendon limit stiffness efforts into internal buffers. + + This function does not apply the tendon limit stiffness to the simulation. It only fills the buffers with + the desired values. To apply the tendon limit stiffness, call the :meth:`write_fixed_tendon_properties_to_sim` function. + + ..note:: When ids are provided, then partial data is expected. When masks are provided the whole of the data + is expected. For example, if env_ids is provided, then root_state should be of shape (len(env_ids), 13). If + env_mask is provided, then root_state should be of shape (num_instances, 13). + + ..caution:: If both env_mask and env_ids are provided, then env_mask will be used. The function will thus + expect the whole of the data to be provided. If none of them are provided, then the function expects the whole + of the data to be provided. + + ..caution:: Do not mix ids and masks. If a mask is provided, then ids will be ignored. + + Args: + limit_stiffness: Fixed tendon limit stiffness. Shape is (len(env_ids), len(fixed_tendon_ids)) or (num_instances, num_fixed_tendons). + fixed_tendon_ids: The tendon indices to set the limit stiffness for. Defaults to None (all fixed tendons). + env_ids: The environment indices to set the limit stiffness for. Defaults to None (all environments). + fixed_tendon_mask: The fixed tendon mask. Shape is (num_fixed_tendons). + env_mask: The environment mask. Shape is (num_instances,). + """ + raise NotImplementedError() + + @abstractmethod + def set_fixed_tendon_position_limit( + self, + limit: torch.Tensor | wp.array, + fixed_tendon_ids: Sequence[int] | slice | None = None, + env_ids: Sequence[int] | None = None, + fixed_tendon_mask: torch.Tensor | wp.array | None = None, + env_mask: torch.Tensor | wp.array | None = None, + ): + """Set fixed tendon limit efforts into internal buffers. + + This function does not apply the tendon limit to the simulation. It only fills the buffers with + the desired values. To apply the tendon limit, call the :meth:`write_fixed_tendon_properties_to_sim` function. + + ..note:: When ids are provided, then partial data is expected. When masks are provided the whole of the data + is expected. For example, if env_ids is provided, then root_state should be of shape (len(env_ids), 13). If + env_mask is provided, then root_state should be of shape (num_instances, 13). + + ..caution:: If both env_mask and env_ids are provided, then env_mask will be used. The function will thus + expect the whole of the data to be provided. If none of them are provided, then the function expects the whole + of the data to be provided. + + ..caution:: Do not mix ids and masks. If a mask is provided, then ids will be ignored. + + Args: + limit: Fixed tendon limit. Shape is (len(env_ids), len(fixed_tendon_ids)) or (num_instances, num_fixed_tendons). + fixed_tendon_ids: The tendon indices to set the limit for. Defaults to None (all fixed tendons). + env_ids: The environment indices to set the limit for. Defaults to None (all environments). + fixed_tendon_mask: The fixed tendon mask. Shape is (num_fixed_tendons). + env_mask: The environment mask. Shape is (num_instances,). + """ + raise NotImplementedError() + + @abstractmethod + def set_fixed_tendon_rest_length( + self, + rest_length: torch.Tensor | wp.array, + fixed_tendon_ids: Sequence[int] | slice | None = None, + env_ids: Sequence[int] | None = None, + fixed_tendon_mask: torch.Tensor | wp.array | None = None, + env_mask: torch.Tensor | wp.array | None = None, + ): + """Set fixed tendon rest length efforts into internal buffers. + + This function does not apply the tendon rest length to the simulation. It only fills the buffers with + the desired values. To apply the tendon rest length, call the :meth:`write_fixed_tendon_properties_to_sim` function. + + ..note:: When ids are provided, then partial data is expected. When masks are provided the whole of the data + is expected. For example, if env_ids is provided, then root_state should be of shape (len(env_ids), 13). If + env_mask is provided, then root_state should be of shape (num_instances, 13). + + ..caution:: If both env_mask and env_ids are provided, then env_mask will be used. The function will thus + expect the whole of the data to be provided. If none of them are provided, then the function expects the whole + of the data to be provided. + + ..caution:: Do not mix ids and masks. If a mask is provided, then ids will be ignored. + + Args: + rest_length: Fixed tendon rest length. Shape is (len(env_ids), len(fixed_tendon_ids)) or (num_instances, num_fixed_tendons). + fixed_tendon_ids: The tendon indices to set the rest length for. Defaults to None (all fixed tendons). + env_ids: The environment indices to set the rest length for. Defaults to None (all environments). + fixed_tendon_mask: The fixed tendon mask. Shape is (num_fixed_tendons). + env_mask: The environment mask. Shape is (num_instances,). + """ + raise NotImplementedError() + + @abstractmethod + def set_fixed_tendon_offset( + self, + offset: torch.Tensor | wp.array, + fixed_tendon_ids: Sequence[int] | slice | None = None, + env_ids: Sequence[int] | None = None, + fixed_tendon_mask: torch.Tensor | wp.array | None = None, + env_mask: torch.Tensor | wp.array | None = None, + ): + """Set fixed tendon offset efforts into internal buffers. + + This function does not apply the tendon offset to the simulation. It only fills the buffers with + the desired values. To apply the tendon offset, call the :meth:`write_fixed_tendon_properties_to_sim` function. + + ..note:: When ids are provided, then partial data is expected. When masks are provided the whole of the data + is expected. For example, if env_ids is provided, then root_state should be of shape (len(env_ids), 13). If + env_mask is provided, then root_state should be of shape (num_instances, 13). + + ..caution:: If both env_mask and env_ids are provided, then env_mask will be used. The function will thus + expect the whole of the data to be provided. If none of them are provided, then the function expects the whole + of the data to be provided. + + ..caution:: Do not mix ids and masks. If a mask is provided, then ids will be ignored. + + Args: + offset: Fixed tendon offset. Shape is (len(env_ids), len(fixed_tendon_ids)) or (num_instances, num_fixed_tendons). + fixed_tendon_ids: The tendon indices to set the offset for. Defaults to None (all fixed tendons). + env_ids: The environment indices to set the offset for. Defaults to None (all environments). + fixed_tendon_mask: The fixed tendon mask. Shape is (num_fixed_tendons). + env_mask: The environment mask. Shape is (num_instances,). + """ + raise NotImplementedError() + + @abstractmethod + def write_fixed_tendon_properties_to_sim( + self, + fixed_tendon_ids: Sequence[int] | slice | None = None, + env_ids: Sequence[int] | None = None, + fixed_tendon_mask: torch.Tensor | wp.array | None = None, + env_mask: torch.Tensor | wp.array | None = None, + ): + """Write fixed tendon properties into the simulation. + + ..note:: When ids are provided, then partial data is expected. When masks are provided the whole of the data + is expected. For example, if env_ids is provided, then root_state should be of shape (len(env_ids), 13). If + env_mask is provided, then root_state should be of shape (num_instances, 13). + + ..caution:: If both env_mask and env_ids are provided, then env_mask will be used. The function will thus + expect the whole of the data to be provided. If none of them are provided, then the function expects the whole + of the data to be provided. + + ..caution:: Do not mix ids and masks. If a mask is provided, then ids will be ignored. + + Args: + fixed_tendon_ids: The fixed tendon indices to set the limits for. Defaults to None (all fixed tendons). + env_ids: The environment indices to set the limits for. Defaults to None (all environments). + fixed_tendon_mask: The fixed tendon mask. Shape is (num_fixed_tendons). + env_mask: The environment mask. Shape is (num_instances,). + """ + raise NotImplementedError() + + @abstractmethod + def set_spatial_tendon_stiffness( + self, + stiffness: torch.Tensor | wp.array, + spatial_tendon_ids: Sequence[int] | slice | None = None, + env_ids: Sequence[int] | None = None, + spatial_tendon_mask: torch.Tensor | wp.array | None = None, + env_mask: torch.Tensor | wp.array | None = None, + ): + """Set spatial tendon stiffness into internal buffers. + + This function does not apply the tendon stiffness to the simulation. It only fills the buffers with + the desired values. To apply the tendon stiffness, call the :meth:`write_spatial_tendon_properties_to_sim` function. + + ..note:: When ids are provided, then partial data is expected. When masks are provided the whole of the data + is expected. For example, if env_ids is provided, then root_state should be of shape (len(env_ids), 13). If + env_mask is provided, then root_state should be of shape (num_instances, 13). + + ..caution:: If both env_mask and env_ids are provided, then env_mask will be used. The function will thus + expect the whole of the data to be provided. If none of them are provided, then the function expects the whole + of the data to be provided. + + ..caution:: Do not mix ids and masks. If a mask is provided, then ids will be ignored. + + Args: + stiffness: Spatial tendon stiffness. Shape is (len(env_ids), len(spatial_tendon_ids)) or (num_instances, num_spatial_tendons). + spatial_tendon_ids: The tendon indices to set the stiffness for. Defaults to None (all spatial tendons). + env_ids: The environment indices to set the stiffness for. Defaults to None (all environments). + spatial_tendon_mask: The spatial tendon mask. Shape is (num_spatial_tendons). + env_mask: The environment mask. Shape is (num_instances,). + """ + raise NotImplementedError() + + @abstractmethod + def set_spatial_tendon_damping( + self, + damping: torch.Tensor | wp.array, + spatial_tendon_ids: Sequence[int] | slice | None = None, + env_ids: Sequence[int] | None = None, + spatial_tendon_mask: torch.Tensor | wp.array | None = None, + env_mask: torch.Tensor | wp.array | None = None, + ): + """Set spatial tendon damping into internal buffers. + + This function does not apply the tendon damping to the simulation. It only fills the buffers with + the desired values. To apply the tendon damping, call the :meth:`write_spatial_tendon_properties_to_sim` function. + + ..note:: When ids are provided, then partial data is expected. When masks are provided the whole of the data + is expected. For example, if env_ids is provided, then root_state should be of shape (len(env_ids), 13). If + env_mask is provided, then root_state should be of shape (num_instances, 13). + + ..caution:: If both env_mask and env_ids are provided, then env_mask will be used. The function will thus + expect the whole of the data to be provided. If none of them are provided, then the function expects the whole + of the data to be provided. + + ..caution:: Do not mix ids and masks. If a mask is provided, then ids will be ignored. + + Args: + damping: Spatial tendon damping. Shape is (len(env_ids), len(spatial_tendon_ids)) or (num_instances, num_spatial_tendons). + spatial_tendon_ids: The tendon indices to set the damping for. Defaults to None (all spatial tendons). + env_ids: The environment indices to set the damping for. Defaults to None (all environments). + spatial_tendon_mask: The spatial tendon mask. Shape is (num_spatial_tendons). + env_mask: The environment mask. Shape is (num_instances,). + """ + raise NotImplementedError() + + @abstractmethod + def set_spatial_tendon_limit_stiffness( + self, + limit_stiffness: torch.Tensor, + spatial_tendon_ids: Sequence[int] | slice | None = None, + env_ids: Sequence[int] | None = None, + spatial_tendon_mask: torch.Tensor | wp.array | None = None, + env_mask: torch.Tensor | wp.array | None = None, + ): + """Set spatial tendon limit stiffness into internal buffers. + + This function does not apply the tendon limit stiffness to the simulation. It only fills the buffers with + the desired values. To apply the tendon limit stiffness, call the :meth:`write_spatial_tendon_properties_to_sim` function. + + ..note:: When ids are provided, then partial data is expected. When masks are provided the whole of the data + is expected. For example, if env_ids is provided, then root_state should be of shape (len(env_ids), 13). If + env_mask is provided, then root_state should be of shape (num_instances, 13). + + ..caution:: If both env_mask and env_ids are provided, then env_mask will be used. The function will thus + expect the whole of the data to be provided. If none of them are provided, then the function expects the whole + of the data to be provided. + + ..caution:: Do not mix ids and masks. If a mask is provided, then ids will be ignored. + + Args: + limit_stiffness: Spatial tendon limit stiffness. Shape is (len(env_ids), len(spatial_tendon_ids)) or (num_instances, num_spatial_tendons). + spatial_tendon_ids: The tendon indices to set the limit stiffness for. Defaults to None (all spatial tendons). + env_ids: The environment indices to set the limit stiffness for. Defaults to None (all environments). + spatial_tendon_mask: The spatial tendon mask. Shape is (num_spatial_tendons). + env_mask: The environment mask. Shape is (num_instances,). + """ + raise NotImplementedError() + + @abstractmethod + def set_spatial_tendon_offset( + self, + offset: torch.Tensor | wp.array, + spatial_tendon_ids: Sequence[int] | slice | None = None, + env_ids: Sequence[int] | None = None, + spatial_tendon_mask: torch.Tensor | wp.array | None = None, + env_mask: torch.Tensor | wp.array | None = None, + ): + """Set spatial tendon offset efforts into internal buffers. + + This function does not apply the tendon offset to the simulation. It only fills the buffers with + the desired values. To apply the tendon offset, call the :meth:`write_spatial_tendon_properties_to_sim` function. + + ..note:: When ids are provided, then partial data is expected. When masks are provided the whole of the data + is expected. For example, if env_ids and spatial_tendon_ids are provided, then offset should be of shape + (len(env_ids), len(spatial_tendon_ids)). If env_mask is provided, then offset should be of shape + (num_instances, num_spatial_tendons). + + ..caution:: If both env_mask and env_ids are provided, then env_mask will be used. The function will thus + expect the whole of the data to be provided. If none of them are provided, then the function expects the whole + of the data to be provided. + + ..caution:: Do not mix ids and masks. If a mask is provided, then ids will be ignored. + + Args: + offset: Spatial tendon offset. Shape is (len(env_ids), len(spatial_tendon_ids)) or (num_instances, num_spatial_tendons). + spatial_tendon_ids: The tendon indices to set the offset for. Defaults to None (all spatial tendons). + env_ids: The environment indices to set the offset for. Defaults to None (all environments). + spatial_tendon_mask: The spatial tendon mask. Shape is (num_spatial_tendons). + env_mask: The environment mask. Shape is (num_instances,). + """ + raise NotImplementedError() + + @abstractmethod + def write_spatial_tendon_properties_to_sim( + self, + spatial_tendon_ids: Sequence[int] | slice | None = None, + env_ids: Sequence[int] | None = None, + spatial_tendon_mask: torch.Tensor | wp.array | None = None, + env_mask: torch.Tensor | wp.array | None = None, + ): + """Write spatial tendon properties into the simulation. + + ..caution:: Do not mix ids and masks. If a mask is provided, then ids will be ignored. + + Args: + spatial_tendon_ids: The spatial tendon indices to set the properties for. Defaults to None (all spatial tendons). + env_ids: The environment indices to set the properties for. Defaults to None (all environments). + spatial_tendon_mask: The spatial tendon mask. Shape is (num_spatial_tendons). + env_mask: The environment mask. Shape is (num_instances,). + """ + raise NotImplementedError() + + """ + Internal helper. + """ + + @abstractmethod + def _initialize_impl(self): + raise NotImplementedError() + + @abstractmethod + def _create_buffers(self): + raise NotImplementedError() + + @abstractmethod + def _process_cfg(self): + """Post processing of configuration parameters.""" + raise NotImplementedError() + + """ + Internal simulation callbacks. + """ + + @abstractmethod + def _invalidate_initialize_callback(self, event): + """Invalidates the scene elements.""" + raise NotImplementedError() + + """ + Internal helpers -- Actuators. + """ + + @abstractmethod + def _process_actuators_cfg(self): + """Process and apply articulation joint properties.""" + raise NotImplementedError() + + @abstractmethod + def _process_tendons(self): + """Process fixed and spatial tendons.""" + raise NotImplementedError() + + @abstractmethod + def _apply_actuator_model(self): + """Processes joint commands for the articulation by forwarding them to the actuators. + + The actions are first processed using actuator models. Depending on the robot configuration, + the actuator models compute the joint level simulation commands and sets them into the PhysX buffers. + """ + raise NotImplementedError() + + """ + Internal helpers -- Debugging. + """ + + @abstractmethod + def _validate_cfg(self): + """Validate the configuration after processing. + + Note: + This function should be called only after the configuration has been processed and the buffers have been + created. Otherwise, some settings that are altered during processing may not be validated. + For instance, the actuator models may change the joint max velocity limits. + """ + raise NotImplementedError() + + @abstractmethod + def _log_articulation_info(self): + """Log information about the articulation. + + Note: We purposefully read the values from the simulator to ensure that the values are configured as expected. + """ + raise NotImplementedError() + + """ + Deprecated methods. + """ + + @abstractmethod + def write_joint_friction_to_sim( + self, + joint_friction: 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, + ): + """Write joint friction coefficients into the simulation. + + .. deprecated:: 2.1.0 + Please use :meth:`write_joint_friction_coefficient_to_sim` instead. + """ + raise NotImplementedError() + + @abstractmethod + def write_joint_limits_to_sim( + self, + limits: torch.Tensor | wp.array | float, + joint_ids: Sequence[int] | slice | None = None, + env_ids: Sequence[int] | None = None, + warn_limit_violation: bool = True, + joint_mask: torch.Tensor | wp.array | None = None, + env_mask: torch.Tensor | wp.array | None = None, + ): + """Write joint limits into the simulation. + + .. deprecated:: 2.1.0 + Please use :meth:`write_joint_position_limit_to_sim` instead. + """ + raise NotImplementedError() + + @abstractmethod + def set_fixed_tendon_limit( + self, + limit: torch.Tensor | wp.array, + fixed_tendon_ids: Sequence[int] | slice | None = None, + env_ids: Sequence[int] | None = None, + fixed_tendon_mask: torch.Tensor | wp.array | None = None, + env_mask: torch.Tensor | wp.array | None = None, + ): + """Set fixed tendon position limits into internal buffers. + + .. deprecated:: 2.1.0 + Please use :meth:`set_fixed_tendon_position_limit` instead. + """ + raise NotImplementedError() diff --git a/source/isaaclab/isaaclab/assets/articulation/base_articulation_data.py b/source/isaaclab/isaaclab/assets/articulation/base_articulation_data.py new file mode 100644 index 00000000000..eeab90f04b8 --- /dev/null +++ b/source/isaaclab/isaaclab/assets/articulation/base_articulation_data.py @@ -0,0 +1,1002 @@ +# Copyright (c) 2022-2025, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +import torch +import warp as wp +from abc import ABC, abstractmethod + + +class BaseArticulationData(ABC): + """Data container for an articulation. + + This class contains the data for an articulation in the simulation. The data includes the state of + the root rigid body, the state of all the bodies in the articulation, and the joint state. The data is + stored in the simulation world frame unless otherwise specified. + + An articulation is comprised of multiple rigid bodies or links. For a rigid body, there are two frames + of reference that are used: + + - Actor frame: The frame of reference of the rigid body prim. This typically corresponds to the Xform prim + with the rigid body schema. + - Center of mass frame: The frame of reference of the center of mass of the rigid body. + + Depending on the settings, the two frames may not coincide with each other. In the robotics sense, the actor frame + can be interpreted as the link frame. + """ + + def __init__(self, root_view, device: str): + """Initializes the articulation data. + + Args: + root_view: The root articulation view. + device: The device used for processing. + """ + # Set the parameters + self.device = device + + @abstractmethod + def update(self, dt: float) -> None: + raise NotImplementedError + + ## + # Names. + ## + + body_names: list[str] = None + """Body names in the order parsed by the simulation view.""" + + joint_names: list[str] = None + """Joint names in the order parsed by the simulation view.""" + + fixed_tendon_names: list[str] = None + """Fixed tendon names in the order parsed by the simulation view.""" + + spatial_tendon_names: list[str] = None + """Spatial tendon names in the order parsed by the simulation view.""" + + ## + # Defaults - Initial state. + ## + + default_root_state: torch.Tensor | wp.array = None + """Default root state ``[pos, quat, lin_vel, ang_vel]`` in the local environment frame. Shape is (num_instances, 13). + + The position and quaternion are of the articulation root's actor frame. Meanwhile, the linear and angular + velocities are of its center of mass frame. + + This quantity is configured through the :attr:`isaaclab.assets.ArticulationCfg.init_state` parameter. + """ + + default_joint_pos: torch.Tensor | wp.array = None + """Default joint positions of all joints. Shape is (num_instances, num_joints). + + This quantity is configured through the :attr:`isaaclab.assets.ArticulationCfg.init_state` parameter. + """ + + default_joint_vel: torch.Tensor | wp.array = None + """Default joint velocities of all joints. Shape is (num_instances, num_joints). + + This quantity is configured through the :attr:`isaaclab.assets.ArticulationCfg.init_state` parameter. + """ + + ## + # Defaults - Physical properties. + ## + + default_mass: torch.Tensor | wp.array = None + """Default mass for all the bodies in the articulation. Shape is (num_instances, num_bodies). + + This quantity is parsed from the USD schema at the time of initialization. + """ + + default_inertia: torch.Tensor | wp.array = None + """Default inertia for all the bodies in the articulation. Shape is (num_instances, num_bodies, 9). + + The inertia tensor should be given with respect to the center of mass, expressed in the articulation links' actor frame. + The values are stored in the order :math:`[I_{xx}, I_{yx}, I_{zx}, I_{xy}, I_{yy}, I_{zy}, I_{xz}, I_{yz}, I_{zz}]`. + However, due to the symmetry of inertia tensors, row- and column-major orders are equivalent. + + This quantity is parsed from the USD schema at the time of initialization. + """ + + default_joint_stiffness: torch.Tensor | wp.array = None + """Default joint stiffness of all joints. Shape is (num_instances, num_joints). + + This quantity is configured through the actuator model's :attr:`isaaclab.actuators.ActuatorBaseCfg.stiffness` + parameter. If the parameter's value is None, the value parsed from the USD schema, at the time of initialization, + is used. + + .. attention:: + The default stiffness is the value configured by the user or the value parsed from the USD schema. + It should not be confused with :attr:`joint_stiffness`, which is the value set into the simulation. + """ + + default_joint_damping: torch.Tensor | wp.array = None + """Default joint damping of all joints. Shape is (num_instances, num_joints). + + This quantity is configured through the actuator model's :attr:`isaaclab.actuators.ActuatorBaseCfg.damping` + parameter. If the parameter's value is None, the value parsed from the USD schema, at the time of initialization, + is used. + + .. attention:: + The default stiffness is the value configured by the user or the value parsed from the USD schema. + It should not be confused with :attr:`joint_damping`, which is the value set into the simulation. + """ + + default_joint_armature: torch.Tensor | wp.array = None + """Default joint armature of all joints. Shape is (num_instances, num_joints). + + This quantity is configured through the actuator model's :attr:`isaaclab.actuators.ActuatorBaseCfg.armature` + parameter. If the parameter's value is None, the value parsed from the USD schema, at the time of initialization, + is used. + """ + + default_joint_friction_coeff: torch.Tensor | wp.array = None + """Default joint static friction coefficient of all joints. Shape is (num_instances, num_joints). + + This quantity is configured through the actuator model's :attr:`isaaclab.actuators.ActuatorBaseCfg.friction` + parameter. If the parameter's value is None, the value parsed from the USD schema, at the time of initialization, + is used. + """ + + default_joint_dynamic_friction_coeff: torch.Tensor | wp.array = None + """Default joint dynamic friction coefficient of all joints. Shape is (num_instances, num_joints). + + This quantity is configured through the actuator model's :attr:`isaaclab.actuators.ActuatorBaseCfg.dynamic_friction` + parameter. If the parameter's value is None, the value parsed from the USD schema, at the time of initialization, + is used. + """ + + default_joint_viscous_friction_coeff: torch.Tensor | wp.array = None + """Default joint viscous friction coefficient of all joints. Shape is (num_instances, num_joints). + + This quantity is configured through the actuator model's :attr:`isaaclab.actuators.ActuatorBaseCfg.viscous_friction` + parameter. If the parameter's value is None, the value parsed from the USD schema, at the time of initialization, + is used. + """ + + default_joint_pos_limits: torch.Tensor | wp.array = None + """Default joint position limits of all joints. Shape is (num_instances, num_joints, 2). + + The limits are in the order :math:`[lower, upper]`. They are parsed from the USD schema at the time of initialization. + """ + default_fixed_tendon_stiffness: torch.Tensor | wp.array = None + """Default tendon stiffness of all fixed tendons. Shape is (num_instances, num_fixed_tendons). + + This quantity is parsed from the USD schema at the time of initialization. + """ + + default_fixed_tendon_damping: torch.Tensor | wp.array = None + """Default tendon damping of all fixed tendons. Shape is (num_instances, num_fixed_tendons). + + This quantity is parsed from the USD schema at the time of initialization. + """ + + default_fixed_tendon_limit_stiffness: torch.Tensor | wp.array = None + """Default tendon limit stiffness of all fixed tendons. Shape is (num_instances, num_fixed_tendons). + + This quantity is parsed from the USD schema at the time of initialization. + """ + + default_fixed_tendon_rest_length: torch.Tensor | wp.array = None + """Default tendon rest length of all fixed tendons. Shape is (num_instances, num_fixed_tendons). + + This quantity is parsed from the USD schema at the time of initialization. + """ + + default_fixed_tendon_offset: torch.Tensor | wp.array = None + """Default tendon offset of all fixed tendons. Shape is (num_instances, num_fixed_tendons). + + This quantity is parsed from the USD schema at the time of initialization. + """ + + default_fixed_tendon_pos_limits: torch.Tensor | wp.array = None + """Default tendon position limits of all fixed tendons. Shape is (num_instances, num_fixed_tendons, 2). + + The position limits are in the order :math:`[lower, upper]`. They are parsed from the USD schema at the time of + initialization. + """ + + default_spatial_tendon_stiffness: torch.Tensor | wp.array = None + """Default tendon stiffness of all spatial tendons. Shape is (num_instances, num_spatial_tendons). + + This quantity is parsed from the USD schema at the time of initialization. + """ + + default_spatial_tendon_damping: torch.Tensor | wp.array = None + """Default tendon damping of all spatial tendons. Shape is (num_instances, num_spatial_tendons). + + This quantity is parsed from the USD schema at the time of initialization. + """ + + default_spatial_tendon_limit_stiffness: torch.Tensor | wp.array = None + """Default tendon limit stiffness of all spatial tendons. Shape is (num_instances, num_spatial_tendons). + + This quantity is parsed from the USD schema at the time of initialization. + """ + + default_spatial_tendon_offset: torch.Tensor | wp.array = None + """Default tendon offset of all spatial tendons. Shape is (num_instances, num_spatial_tendons). + + This quantity is parsed from the USD schema at the time of initialization. + """ + + ## + # Joint commands -- Set into simulation. + ## + + joint_pos_target: torch.Tensor | wp.array = None + """Joint position targets commanded by the user. Shape is (num_instances, num_joints). + + For an implicit actuator model, the targets are directly set into the simulation. + For an explicit actuator model, the targets are used to compute the joint torques (see :attr:`applied_torque`), + which are then set into the simulation. + """ + + joint_vel_target: torch.Tensor | wp.array = None + """Joint velocity targets commanded by the user. Shape is (num_instances, num_joints). + + For an implicit actuator model, the targets are directly set into the simulation. + For an explicit actuator model, the targets are used to compute the joint torques (see :attr:`applied_torque`), + which are then set into the simulation. + """ + + joint_effort_target: torch.Tensor | wp.array = None + """Joint effort targets commanded by the user. Shape is (num_instances, num_joints). + + For an implicit actuator model, the targets are directly set into the simulation. + For an explicit actuator model, the targets are used to compute the joint torques (see :attr:`applied_torque`), + which are then set into the simulation. + """ + + ## + # Joint commands -- Explicit actuators. + ## + + computed_torque: torch.Tensor | wp.array = None + """Joint torques computed from the actuator model (before clipping). Shape is (num_instances, num_joints). + + This quantity is the raw torque output from the actuator mode, before any clipping is applied. + It is exposed for users who want to inspect the computations inside the actuator model. + For instance, to penalize the learning agent for a difference between the computed and applied torques. + """ + + applied_torque: torch.Tensor | wp.array = None + """Joint torques applied from the actuator model (after clipping). Shape is (num_instances, num_joints). + + These torques are set into the simulation, after clipping the :attr:`computed_torque` based on the + actuator model. + """ + + ## + # Joint properties. + ## + + joint_stiffness: torch.Tensor | wp.array = None + """Joint stiffness provided to the simulation. Shape is (num_instances, num_joints). + + In the case of explicit actuators, the value for the corresponding joints is zero. + """ + + joint_damping: torch.Tensor | wp.array = None + """Joint damping provided to the simulation. Shape is (num_instances, num_joints) + + In the case of explicit actuators, the value for the corresponding joints is zero. + """ + + joint_armature: torch.Tensor | wp.array = None + """Joint armature provided to the simulation. Shape is (num_instances, num_joints).""" + + joint_friction_coeff: torch.Tensor | wp.array = None + """Joint static friction coefficient provided to the simulation. Shape is (num_instances, num_joints).""" + + joint_dynamic_friction_coeff: torch.Tensor | wp.array = None + """Joint dynamic friction coefficient provided to the simulation. Shape is (num_instances, num_joints).""" + + joint_viscous_friction_coeff: torch.Tensor | wp.array = None + """Joint viscous friction coefficient provided to the simulation. Shape is (num_instances, num_joints).""" + + joint_pos_limits: torch.Tensor | wp.array = None + """Joint position limits provided to the simulation. Shape is (num_instances, num_joints, 2). + + The limits are in the order :math:`[lower, upper]`. + """ + + joint_vel_limits: torch.Tensor | wp.array = None + """Joint maximum velocity provided to the simulation. Shape is (num_instances, num_joints).""" + + joint_effort_limits: torch.Tensor | wp.array = None + """Joint maximum effort provided to the simulation. Shape is (num_instances, num_joints).""" + + ## + # Joint properties - Custom. + ## + + soft_joint_pos_limits: torch.Tensor | wp.array = None + r"""Soft joint positions limits for all joints. Shape is (num_instances, num_joints, 2). + + The limits are in the order :math:`[lower, upper]`.The soft joint position limits are computed as + a sub-region of the :attr:`joint_pos_limits` based on the + :attr:`~isaaclab.assets.ArticulationCfg.soft_joint_pos_limit_factor` parameter. + + Consider the joint position limits :math:`[lower, upper]` and the soft joint position limits + :math:`[soft_lower, soft_upper]`. The soft joint position limits are computed as: + + .. math:: + + soft\_lower = (lower + upper) / 2 - factor * (upper - lower) / 2 + soft\_upper = (lower + upper) / 2 + factor * (upper - lower) / 2 + + The soft joint position limits help specify a safety region around the joint limits. It isn't used by the + simulation, but is useful for learning agents to prevent the joint positions from violating the limits. + """ + + soft_joint_vel_limits: torch.Tensor | wp.array = None + """Soft joint velocity limits for all joints. Shape is (num_instances, num_joints). + + These are obtained from the actuator model. It may differ from :attr:`joint_vel_limits` if the actuator model + has a variable velocity limit model. For instance, in a variable gear ratio actuator model. + """ + + gear_ratio: torch.Tensor | wp.array = None + """Gear ratio for relating motor torques to applied Joint torques. Shape is (num_instances, num_joints).""" + + ## + # Fixed tendon properties. + ## + + fixed_tendon_stiffness: torch.Tensor | wp.array = None + """Fixed tendon stiffness provided to the simulation. Shape is (num_instances, num_fixed_tendons).""" + + fixed_tendon_damping: torch.Tensor | wp.array = None + """Fixed tendon damping provided to the simulation. Shape is (num_instances, num_fixed_tendons).""" + + fixed_tendon_limit_stiffness: torch.Tensor | wp.array = None + """Fixed tendon limit stiffness provided to the simulation. Shape is (num_instances, num_fixed_tendons).""" + + fixed_tendon_rest_length: torch.Tensor | wp.array = None + """Fixed tendon rest length provided to the simulation. Shape is (num_instances, num_fixed_tendons).""" + + fixed_tendon_offset: torch.Tensor | wp.array = None + """Fixed tendon offset provided to the simulation. Shape is (num_instances, num_fixed_tendons).""" + + fixed_tendon_pos_limits: torch.Tensor | wp.array = None + """Fixed tendon position limits provided to the simulation. Shape is (num_instances, num_fixed_tendons, 2).""" + + ## + # Spatial tendon properties. + ## + + spatial_tendon_stiffness: torch.Tensor | wp.array = None + """Spatial tendon stiffness provided to the simulation. Shape is (num_instances, num_spatial_tendons).""" + + spatial_tendon_damping: torch.Tensor | wp.array = None + """Spatial tendon damping provided to the simulation. Shape is (num_instances, num_spatial_tendons).""" + + spatial_tendon_limit_stiffness: torch.Tensor | wp.array = None + """Spatial tendon limit stiffness provided to the simulation. Shape is (num_instances, num_spatial_tendons).""" + + spatial_tendon_offset: torch.Tensor | wp.array = None + """Spatial tendon offset provided to the simulation. Shape is (num_instances, num_spatial_tendons).""" + + ## + # Root state properties. + ## + + @property + @abstractmethod + def root_link_pose_w(self) -> torch.Tensor | wp.array: + """Root link pose ``[pos, quat]`` in simulation world frame. Shape is (num_instances, 7). + + This quantity is the pose of the articulation root's actor frame relative to the world. + The orientation is provided in (w, x, y, z) format. + """ + raise NotImplementedError + + @property + @abstractmethod + def root_link_vel_w(self) -> torch.Tensor | wp.array: + """Root link velocity ``[lin_vel, ang_vel]`` in simulation world frame. Shape is (num_instances, 6). + + This quantity contains the linear and angular velocities of the articulation root's actor frame + relative to the world. + """ + raise NotImplementedError + + @property + @abstractmethod + def root_com_pose_w(self) -> torch.Tensor | wp.array: + """Root center of mass pose ``[pos, quat]`` in simulation world frame. Shape is (num_instances, 7). + + This quantity is the pose of the articulation root's center of mass frame relative to the world. + The orientation is provided in (w, x, y, z) format. + """ + raise NotImplementedError + + @property + @abstractmethod + def root_com_vel_w(self) -> torch.Tensor | wp.array: + """Root center of mass velocity ``[lin_vel, ang_vel]`` in simulation world frame. Shape is (num_instances, 6). + + This quantity contains the linear and angular velocities of the articulation root's center of mass frame + relative to the world. + """ + raise NotImplementedError + + @property + @abstractmethod + def root_state_w(self) -> torch.Tensor | wp.array: + """Root state ``[pos, quat, lin_vel, ang_vel]`` in simulation world frame. Shape is (num_instances, 13). + + The position and quaternion are of the articulation root's actor frame relative to the world. Meanwhile, + the linear and angular velocities are of the articulation root's center of mass frame. + """ + raise NotImplementedError + + @property + @abstractmethod + def root_link_state_w(self) -> torch.Tensor | wp.array: + """Root state ``[pos, quat, lin_vel, ang_vel]`` in simulation world frame. Shape is (num_instances, 13). + + The position, quaternion, and linear/angular velocity are of the articulation root's actor frame relative to the + world. + """ + raise NotImplementedError + + @property + @abstractmethod + def root_com_state_w(self) -> torch.Tensor | wp.array: + """Root center of mass state ``[pos, quat, lin_vel, ang_vel]`` in simulation world frame. Shape is (num_instances, 13). + + The position, quaternion, and linear/angular velocity are of the articulation root link's center of mass frame + relative to the world. Center of mass frame is assumed to be the same orientation as the link rather than the + orientation of the principle inertia. + """ + raise NotImplementedError + + ## + # Body state properties. + ## + + @property + @abstractmethod + def body_link_pose_w(self) -> torch.Tensor | wp.array: + """Body link pose ``[pos, quat]`` in simulation world frame. + Shape is (num_instances, num_bodies, 7). + + This quantity is the pose of the articulation links' actor frame relative to the world. + The orientation is provided in (w, x, y, z) format. + """ + raise NotImplementedError + + @property + @abstractmethod + def body_link_vel_w(self) -> torch.Tensor | wp.array: + """Body link velocity ``[lin_vel, ang_vel]`` in simulation world frame. + Shape is (num_instances, num_bodies, 6). + + This quantity contains the linear and angular velocities of the articulation links' actor frame + relative to the world. + """ + raise NotImplementedError + + @property + @abstractmethod + def body_com_pose_w(self) -> torch.Tensor | wp.array: + """Body center of mass pose ``[pos, quat]`` in simulation world frame. + Shape is (num_instances, num_bodies, 7). + + This quantity is the pose of the center of mass frame of the articulation links relative to the world. + The orientation is provided in (w, x, y, z) format. + """ + raise NotImplementedError + + @property + @abstractmethod + def body_com_vel_w(self) -> torch.Tensor | wp.array: + """Body center of mass velocity ``[lin_vel, ang_vel]`` in simulation world frame. + Shape is (num_instances, num_bodies, 6). + + This quantity contains the linear and angular velocities of the articulation links' center of mass frame + relative to the world. + """ + raise NotImplementedError + + @property + @abstractmethod + def body_state_w(self) -> torch.Tensor | wp.array: + """State of all bodies `[pos, quat, lin_vel, ang_vel]` in simulation world frame. + Shape is (num_instances, num_bodies, 13). + + The position and quaternion are of all the articulation links' actor frame. Meanwhile, the linear and angular + velocities are of the articulation links's center of mass frame. + """ + raise NotImplementedError + + @property + @abstractmethod + def body_link_state_w(self) -> torch.Tensor | wp.array: + """State of all bodies' link frame`[pos, quat, lin_vel, ang_vel]` in simulation world frame. + Shape is (num_instances, num_bodies, 13). + + The position, quaternion, and linear/angular velocity are of the body's link frame relative to the world. + """ + raise NotImplementedError + + @property + @abstractmethod + def body_com_state_w(self) -> torch.Tensor | wp.array: + """State of all bodies center of mass `[pos, quat, lin_vel, ang_vel]` in simulation world frame. + Shape is (num_instances, num_bodies, 13). + + The position, quaternion, and linear/angular velocity are of the body's center of mass frame relative to the + world. Center of mass frame is assumed to be the same orientation as the link rather than the orientation of the + principle inertia. + """ + raise NotImplementedError + + @property + @abstractmethod + def body_com_acc_w(self) -> torch.Tensor | wp.array: + """Acceleration of all bodies center of mass ``[lin_acc, ang_acc]``. + Shape is (num_instances, num_bodies, 6). + + All values are relative to the world. + """ + raise NotImplementedError + + @property + @abstractmethod + def body_com_pose_b(self) -> torch.Tensor | wp.array: + """Center of mass pose ``[pos, quat]`` of all bodies in their respective body's link frames. + Shape is (num_instances, 1, 7). + + This quantity is the pose of the center of mass frame of the rigid body relative to the body's link frame. + The orientation is provided in (w, x, y, z) format. + """ + raise NotImplementedError + + @property + @abstractmethod + def body_incoming_joint_wrench_b(self) -> torch.Tensor | wp.array: + """Joint reaction wrench applied from body parent to child body in parent body frame. + + Shape is (num_instances, num_bodies, 6). All body reaction wrenches are provided including the root body to the + world of an articulation. + + For more information on joint wrenches, please check the`PhysX documentation `__ + and the underlying `PhysX Tensor API `__ . + """ + raise NotImplementedError + + ## + # Joint state properties. + ## + + @property + @abstractmethod + def joint_pos(self) -> torch.Tensor | wp.array: + """Joint positions of all joints. Shape is (num_instances, num_joints).""" + raise NotImplementedError + + @property + @abstractmethod + def joint_vel(self) -> torch.Tensor | wp.array: + """Joint velocities of all joints. Shape is (num_instances, num_joints).""" + raise NotImplementedError + + @property + @abstractmethod + def joint_acc(self) -> torch.Tensor | wp.array: + """Joint acceleration of all joints. Shape is (num_instances, num_joints).""" + raise NotImplementedError + + ## + # Derived Properties. + ## + + @property + @abstractmethod + def projected_gravity_b(self) -> torch.Tensor | wp.array: + """Projection of the gravity direction on base frame. Shape is (num_instances, 3).""" + raise NotImplementedError + + @property + @abstractmethod + def heading_w(self) -> torch.Tensor | wp.array: + """Yaw heading of the base frame (in radians). Shape is (num_instances,). + + Note: + This quantity is computed by assuming that the forward-direction of the base + frame is along x-direction, i.e. :math:`(1, 0, 0)`. + """ + raise NotImplementedError + + @property + @abstractmethod + def root_link_lin_vel_b(self) -> torch.Tensor | wp.array: + """Root link linear velocity in base frame. Shape is (num_instances, 3). + + This quantity is the linear velocity of the articulation root's actor frame with respect to the + its actor frame. + """ + raise NotImplementedError + + @property + @abstractmethod + def root_link_ang_vel_b(self) -> torch.Tensor | wp.array: + """Root link angular velocity in base world frame. Shape is (num_instances, 3). + + This quantity is the angular velocity of the articulation root's actor frame with respect to the + its actor frame. + """ + raise NotImplementedError + + @property + @abstractmethod + def root_com_lin_vel_b(self) -> torch.Tensor | wp.array: + """Root center of mass linear velocity in base frame. Shape is (num_instances, 3). + + This quantity is the linear velocity of the articulation root's center of mass frame with respect to the + its actor frame. + """ + raise NotImplementedError + + @property + @abstractmethod + def root_com_ang_vel_b(self) -> torch.Tensor | wp.array: + """Root center of mass angular velocity in base world frame. Shape is (num_instances, 3). + + This quantity is the angular velocity of the articulation root's center of mass frame with respect to the + its actor frame. + """ + raise NotImplementedError + + ## + # Sliced properties. + ## + + @property + @abstractmethod + def root_link_pos_w(self) -> torch.Tensor | wp.array: + """Root link position in simulation world frame. Shape is (num_instances, 3). + + This quantity is the position of the actor frame of the root rigid body relative to the world. + """ + raise NotImplementedError + + @property + @abstractmethod + def root_link_quat_w(self) -> torch.Tensor | wp.array: + """Root link orientation (w, x, y, z) in simulation world frame. Shape is (num_instances, 4). + + This quantity is the orientation of the actor frame of the root rigid body. + """ + raise NotImplementedError + + @property + @abstractmethod + def root_link_lin_vel_w(self) -> torch.Tensor | wp.array: + """Root linear velocity in simulation world frame. Shape is (num_instances, 3). + + This quantity is the linear velocity of the root rigid body's actor frame relative to the world. + """ + raise NotImplementedError + + @property + @abstractmethod + def root_link_ang_vel_w(self) -> torch.Tensor | wp.array: + """Root link angular velocity in simulation world frame. Shape is (num_instances, 3). + + This quantity is the angular velocity of the actor frame of the root rigid body relative to the world. + """ + raise NotImplementedError + + @property + @abstractmethod + def root_com_pos_w(self) -> torch.Tensor | wp.array: + """Root center of mass position in simulation world frame. Shape is (num_instances, 3). + + This quantity is the position of the actor frame of the root rigid body relative to the world. + """ + raise NotImplementedError + + @property + @abstractmethod + def root_com_quat_w(self) -> torch.Tensor | wp.array: + """Root center of mass orientation (w, x, y, z) in simulation world frame. Shape is (num_instances, 4). + + This quantity is the orientation of the actor frame of the root rigid body relative to the world. + """ + raise NotImplementedError + + @property + @abstractmethod + def root_com_lin_vel_w(self) -> torch.Tensor | wp.array: + """Root center of mass linear velocity in simulation world frame. Shape is (num_instances, 3). + + This quantity is the linear velocity of the root rigid body's center of mass frame relative to the world. + """ + raise NotImplementedError + + @property + @abstractmethod + def root_com_ang_vel_w(self) -> torch.Tensor | wp.array: + """Root center of mass angular velocity in simulation world frame. Shape is (num_instances, 3). + + This quantity is the angular velocity of the root rigid body's center of mass frame relative to the world. + """ + raise NotImplementedError + + @property + @abstractmethod + def body_link_pos_w(self) -> torch.Tensor | wp.array: + """Positions of all bodies in simulation world frame. Shape is (num_instances, num_bodies, 3). + + This quantity is the position of the articulation bodies' actor frame relative to the world. + """ + raise NotImplementedError + + @property + @abstractmethod + def body_link_quat_w(self) -> torch.Tensor | wp.array: + """Orientation (w, x, y, z) of all bodies in simulation world frame. Shape is (num_instances, num_bodies, 4). + + This quantity is the orientation of the articulation bodies' actor frame relative to the world. + """ + raise NotImplementedError + + @property + @abstractmethod + def body_link_lin_vel_w(self) -> torch.Tensor | wp.array: + """Linear velocity of all bodies in simulation world frame. Shape is (num_instances, num_bodies, 3). + + This quantity is the linear velocity of the articulation bodies' center of mass frame relative to the world. + """ + raise NotImplementedError + + @property + @abstractmethod + def body_link_ang_vel_w(self) -> torch.Tensor | wp.array: + """Angular velocity of all bodies in simulation world frame. Shape is (num_instances, num_bodies, 3). + + This quantity is the angular velocity of the articulation bodies' center of mass frame relative to the world. + """ + raise NotImplementedError + + @property + @abstractmethod + def body_com_pos_w(self) -> torch.Tensor | wp.array: + """Positions of all bodies in simulation world frame. Shape is (num_instances, num_bodies, 3). + + This quantity is the position of the articulation bodies' actor frame. + """ + raise NotImplementedError + + @property + @abstractmethod + def body_com_quat_w(self) -> torch.Tensor | wp.array: + """Orientation (w, x, y, z) of the principle axis of inertia of all bodies in simulation world frame. + Shape is (num_instances, num_bodies, 4). + + This quantity is the orientation of the articulation bodies' actor frame. + """ + raise NotImplementedError + + @property + @abstractmethod + def body_com_lin_vel_w(self) -> torch.Tensor | wp.array: + """Linear velocity of all bodies in simulation world frame. Shape is (num_instances, num_bodies, 3). + + This quantity is the linear velocity of the articulation bodies' center of mass frame. + """ + raise NotImplementedError + + @property + @abstractmethod + def body_com_ang_vel_w(self) -> torch.Tensor | wp.array: + """Angular velocity of all bodies in simulation world frame. Shape is (num_instances, num_bodies, 3). + + This quantity is the angular velocity of the articulation bodies' center of mass frame. + """ + raise NotImplementedError + + @property + @abstractmethod + def body_com_lin_acc_w(self) -> torch.Tensor | wp.array: + """Linear acceleration of all bodies in simulation world frame. Shape is (num_instances, num_bodies, 3). + + This quantity is the linear acceleration of the articulation bodies' center of mass frame. + """ + raise NotImplementedError + + @property + @abstractmethod + def body_com_ang_acc_w(self) -> torch.Tensor | wp.array: + """Angular acceleration of all bodies in simulation world frame. Shape is (num_instances, num_bodies, 3). + + This quantity is the angular acceleration of the articulation bodies' center of mass frame. + """ + raise NotImplementedError + + @property + @abstractmethod + def body_com_pos_b(self) -> torch.Tensor | wp.array: + """Center of mass position of all of the bodies in their respective link frames. + Shape is (num_instances, num_bodies, 3). + + This quantity is the center of mass location relative to its body'slink frame. + """ + raise NotImplementedError + + @property + @abstractmethod + def body_com_quat_b(self) -> torch.Tensor | wp.array: + """Orientation (w, x, y, z) of the principle axis of inertia of all of the bodies in their + respective link frames. Shape is (num_instances, num_bodies, 4). + + This quantity is the orientation of the principles axes of inertia relative to its body's link frame. + """ + raise NotImplementedError + + ## + # Backward compatibility. + ## + + @property + @abstractmethod + def root_pose_w(self) -> torch.Tensor | wp.array: + """Same as :attr:`root_link_pose_w`.""" + raise NotImplementedError + + @property + @abstractmethod + def root_pos_w(self) -> torch.Tensor | wp.array: + """Same as :attr:`root_link_pos_w`.""" + raise NotImplementedError + + @property + @abstractmethod + def root_quat_w(self) -> torch.Tensor | wp.array: + """Same as :attr:`root_link_quat_w`.""" + raise NotImplementedError + + @property + @abstractmethod + def root_vel_w(self) -> torch.Tensor | wp.array: + """Same as :attr:`root_com_vel_w`.""" + raise NotImplementedError + + @property + @abstractmethod + def root_lin_vel_w(self) -> torch.Tensor | wp.array: + """Same as :attr:`root_com_lin_vel_w`.""" + raise NotImplementedError + + @property + @abstractmethod + def root_ang_vel_w(self) -> torch.Tensor | wp.array: + """Same as :attr:`root_com_ang_vel_w`.""" + raise NotImplementedError + + @property + @abstractmethod + def root_lin_vel_b(self) -> torch.Tensor | wp.array: + """Same as :attr:`root_com_lin_vel_b`.""" + raise NotImplementedError + + @property + @abstractmethod + def root_ang_vel_b(self) -> torch.Tensor | wp.array: + """Same as :attr:`root_com_ang_vel_b`.""" + raise NotImplementedError + + @property + @abstractmethod + def body_pose_w(self) -> torch.Tensor | wp.array: + """Same as :attr:`body_link_pose_w`.""" + raise NotImplementedError + + @property + @abstractmethod + def body_pos_w(self) -> torch.Tensor | wp.array: + """Same as :attr:`body_link_pos_w`.""" + raise NotImplementedError + + @property + @abstractmethod + def body_quat_w(self) -> torch.Tensor | wp.array: + """Same as :attr:`body_link_quat_w`.""" + raise NotImplementedError + + @property + @abstractmethod + def body_vel_w(self) -> torch.Tensor | wp.array: + """Same as :attr:`body_com_vel_w`.""" + raise NotImplementedError + + @property + @abstractmethod + def body_lin_vel_w(self) -> torch.Tensor | wp.array: + """Same as :attr:`body_com_lin_vel_w`.""" + raise NotImplementedError + + @property + @abstractmethod + def body_ang_vel_w(self) -> torch.Tensor | wp.array: + """Same as :attr:`body_com_ang_vel_w`.""" + raise NotImplementedError + + @property + @abstractmethod + def body_acc_w(self) -> torch.Tensor | wp.array: + """Same as :attr:`body_com_acc_w`.""" + raise NotImplementedError + + @property + @abstractmethod + def body_lin_acc_w(self) -> torch.Tensor | wp.array: + """Same as :attr:`body_com_lin_acc_w`.""" + raise NotImplementedError + + @property + @abstractmethod + def body_ang_acc_w(self) -> torch.Tensor | wp.array: + """Same as :attr:`body_com_ang_acc_w`.""" + raise NotImplementedError + + @property + @abstractmethod + def com_pos_b(self) -> torch.Tensor | wp.array: + """Same as :attr:`body_com_pos_b`.""" + raise NotImplementedError + + @property + @abstractmethod + def com_quat_b(self) -> torch.Tensor | wp.array: + """Same as :attr:`body_com_quat_b`.""" + raise NotImplementedError + + @property + @abstractmethod + def joint_limits(self) -> torch.Tensor | wp.array: + """Deprecated property. Please use :attr:`joint_pos_limits` instead.""" + raise NotImplementedError + + @property + @abstractmethod + def default_joint_limits(self) -> torch.Tensor | wp.array: + """Deprecated property. Please use :attr:`default_joint_pos_limits` instead.""" + raise NotImplementedError + + @property + @abstractmethod + def joint_velocity_limits(self) -> torch.Tensor | wp.array: + """Deprecated property. Please use :attr:`joint_vel_limits` instead.""" + raise NotImplementedError + + @property + @abstractmethod + def joint_friction(self) -> torch.Tensor | wp.array: + """Deprecated property. Please use :attr:`joint_friction_coeff` instead.""" + raise NotImplementedError + + @property + @abstractmethod + def default_joint_friction(self) -> torch.Tensor | wp.array: + """Deprecated property. Please use :attr:`default_joint_friction_coeff` instead.""" + raise NotImplementedError + + @property + @abstractmethod + def fixed_tendon_limit(self) -> torch.Tensor | wp.array: + """Deprecated property. Please use :attr:`fixed_tendon_pos_limits` instead.""" + raise NotImplementedError + + @property + @abstractmethod + def default_fixed_tendon_limit(self) -> torch.Tensor | wp.array: + """Deprecated property. Please use :attr:`default_fixed_tendon_pos_limits` instead.""" + raise NotImplementedError diff --git a/source/isaaclab/isaaclab/assets/rigid_object/__init__.py b/source/isaaclab/isaaclab/assets/rigid_object/__init__.py new file mode 100644 index 00000000000..94c39677672 --- /dev/null +++ b/source/isaaclab/isaaclab/assets/rigid_object/__init__.py @@ -0,0 +1,10 @@ +# Copyright (c) 2022-2025, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +"""Sub-module for rigid object assets.""" + +from .rigid_object import RigidObject +from .rigid_object_cfg import RigidObjectCfg +from .rigid_object_data import RigidObjectData diff --git a/source/isaaclab/isaaclab/assets/rigid_object/base_rigid_object.py b/source/isaaclab/isaaclab/assets/rigid_object/base_rigid_object.py new file mode 100644 index 00000000000..d06f62c723a --- /dev/null +++ b/source/isaaclab/isaaclab/assets/rigid_object/base_rigid_object.py @@ -0,0 +1,470 @@ +# Copyright (c) 2022-2025, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +from __future__ import annotations + +import torch +import warp as wp +from collections.abc import Sequence +from typing import TYPE_CHECKING +from abc import abstractmethod + +from ..asset_base import AssetBase + +if TYPE_CHECKING: + from .rigid_object_cfg import RigidObjectCfg + from .rigid_object_data import RigidObjectData + + +class BaseRigidObject(AssetBase): + """A rigid object asset class. + + Rigid objects are assets comprising of rigid bodies. They can be used to represent dynamic objects + such as boxes, spheres, etc. A rigid body is described by its pose, velocity and mass distribution. + + For an asset to be considered a rigid object, the root prim of the asset must have the `USD RigidBodyAPI`_ + applied to it. This API is used to define the simulation properties of the rigid body. On playing the + simulation, the physics engine will automatically register the rigid body and create a corresponding + rigid body handle. This handle can be accessed using the :attr:`root_physx_view` attribute. + + .. note:: + + For users familiar with Isaac Sim, the PhysX view class API is not the exactly same as Isaac Sim view + class API. Similar to Isaac Lab, Isaac Sim wraps around the PhysX view API. However, as of now (2023.1 release), + we see a large difference in initializing the view classes in Isaac Sim. This is because the view classes + in Isaac Sim perform additional USD-related operations which are slow and also not required. + + .. _`USD RigidBodyAPI`: https://openusd.org/dev/api/class_usd_physics_rigid_body_a_p_i.html + """ + + cfg: RigidObjectCfg + """Configuration instance for the rigid object.""" + + __backend_name__: str = "base" + """The name of the backend for the rigid object.""" + + def __init__(self, cfg: RigidObjectCfg): + """Initialize the rigid object. + + Args: + cfg: A configuration instance. + """ + super().__init__(cfg) + + """ + Properties + """ + + @property + @abstractmethod + def data(self) -> RigidObjectData: + raise NotImplementedError() + + @property + @abstractmethod + def num_instances(self) -> int: + raise NotImplementedError() + + @property + @abstractmethod + def num_bodies(self) -> int: + """Number of bodies in the asset. + + This is always 1 since each object is a single rigid body. + """ + raise NotImplementedError() + + @property + @abstractmethod + def body_names(self) -> list[str]: + """Ordered names of bodies in the rigid object.""" + raise NotImplementedError() + + @property + @abstractmethod + def root_view(self): + """Root view for the asset. + + Note: + Use this view with caution. It requires handling of tensors in a specific way. + """ + raise NotImplementedError() + + """ + Operations. + """ + + @abstractmethod + def reset(self, env_ids: Sequence[int] | None = None, mask: wp.array | torch.Tensor | None = None): + raise NotImplementedError() + + @abstractmethod + def write_data_to_sim(self) -> None: + """Write external wrench to the simulation. + + Note: + We write external wrench to the simulation here since this function is called before the simulation step. + This ensures that the external wrench is applied at every simulation step. + """ + raise NotImplementedError() + + @abstractmethod + def update(self, dt: float) -> None: + raise NotImplementedError() + + """ + Operations - Finders. + """ + + @abstractmethod + def find_bodies( + self, + name_keys: str | Sequence[str], + preserve_order: bool = False + ) -> tuple[wp.array | torch.Tensor, list[str], list[int]]: + """Find bodies in the rigid body based on the name keys. + + Please check the :meth:`isaaclab.utils.string_utils.resolve_matching_names` function for more + information on the name matching. + + Args: + name_keys: A regular expression or a list of regular expressions to match the body names. + preserve_order: Whether to preserve the order of the name keys in the output. Defaults to False. + + Returns: + A tuple of lists containing the body mask, names and indices. + """ + raise NotImplementedError() + + """ + Operations - Write to simulation. + """ + + @abstractmethod + def write_root_state_to_sim( + self, + root_state: torch.Tensor | wp.array, + env_ids: Sequence[int] | None = None, + env_mask: torch.Tensor | wp.array | None = None, + ) -> None: + """Set the root state over selected environment indices into the simulation. + + The root state comprises of the cartesian position, quaternion orientation in (x, y, z, w), and linear + and angular velocity. All the quantities are in the simulation frame. + + ..note:: When ids are provided, then partial data is expected. When masks are provided the whole of the data + is expected. For example, if env_ids is provided, then root_state should be of shape (len(env_ids), 13). If + env_mask is provided, then root_state should be of shape (num_instances, 13). + + ..caution:: If both env_mask and env_ids are provided, then env_mask will be used. The function will thus + expect the whole of the data to be provided. If none of them are provided, then the function expects the whole + of the data to be provided. + + Args: + root_state: Root state in simulation frame. Shape is (len(env_ids), 13) or (num_instances, 13). + env_ids: Environment indices. If None, then all indices are used. + env_mask: Environment mask. Shape is (num_instances,). + """ + raise NotImplementedError() + + @abstractmethod + def write_root_com_state_to_sim( + self, + root_state: torch.Tensor | wp.array, + env_ids: Sequence[int] | None = None, + env_mask: torch.Tensor | wp.array | None = None, + ) -> None: + """Set the root center of mass state over selected environment indices into the simulation. + + The root state comprises of the cartesian position, quaternion orientation in (x, y, z, w), and linear + and angular velocity. All the quantities are in the simulation frame. + + ..note:: When ids are provided, then partial data is expected. When masks are provided the whole of the data + is expected. For example, if env_ids is provided, then root_state should be of shape (len(env_ids), 13). If + env_mask is provided, then root_state should be of shape (num_instances, 13). + + ..caution:: If both env_mask and env_ids are provided, then env_mask will be used. The function will thus + expect the whole of the data to be provided. If none of them are provided, then the function expects the whole + of the data to be provided. + + Args: + root_state: Root state in simulation frame. Shape is (len(env_ids), 13) or (num_instances, 13). + env_ids: Environment indices. If None, then all indices are used. + env_mask: Environment mask. Shape is (num_instances,). + """ + raise NotImplementedError() + + @abstractmethod + def write_root_link_state_to_sim( + self, + root_state: torch.Tensor | wp.array, + env_ids: Sequence[int] | None = None, + env_mask: torch.Tensor | wp.array | None = None, + ) -> None: + """Set the root link state over selected environment indices into the simulation. + + The root state comprises of the cartesian position, quaternion orientation in (x, y, z, w), and linear + and angular velocity. All the quantities are in the simulation frame. + + ..note:: When ids are provided, then partial data is expected. When masks are provided the whole of the data + is expected. For example, if env_ids is provided, then root_state should be of shape (len(env_ids), 13). If + env_mask is provided, then root_state should be of shape (num_instances, 13). + + ..caution:: If both env_mask and env_ids are provided, then env_mask will be used. The function will thus + expect the whole of the data to be provided. If none of them are provided, then the function expects the whole + of the data to be provided. + + Args: + root_state: Root state in simulation frame. Shape is (len(env_ids), 13) or (num_instances, 13). + env_mask: Environment mask. Shape is (num_instances,). + env_ids: Environment indices. If None, then all indices are used. + """ + raise NotImplementedError() + + @abstractmethod + def write_root_pose_to_sim( + self, + root_pose: torch.Tensor | wp.array, + env_ids: Sequence[int] | None = None, + env_mask: torch.Tensor | wp.array | None = None, + ) -> None: + """Set the root pose over selected environment indices into the simulation. + + The root pose comprises of the cartesian position and quaternion orientation in (x, y, z, w). + + ..note:: When ids are provided, then partial data is expected. When masks are provided the whole of the data + is expected. For example, if env_ids is provided, then root_pose should be of shape (len(env_ids), 7). If + env_mask is provided, then root_pose should be of shape (num_instances, 7). + + ..caution:: If both env_mask and env_ids are provided, then env_mask will be used. The function will thus + expect the whole of the data to be provided. If none of them are provided, then the function expects the whole + of the data to be provided. + + Args: + root_pose: Root link poses in simulation frame. Shape is (len(env_ids), 7) or (num_instances, 7). + env_ids: Environment indices. If None, then all indices are used. + env_mask: Environment mask. Shape is (num_instances,). + """ + raise NotImplementedError() + + @abstractmethod + def write_root_link_pose_to_sim( + self, + root_pose: torch.Tensor | wp.array, + env_ids: Sequence[int] | None = None, + env_mask: torch.Tensor | wp.array | None = None, + ) -> None: + """Set the root link pose over selected environment indices into the simulation. + + The root pose comprises of the cartesian position and quaternion orientation in (x, y, z, w). + + ..note:: When ids are provided, then partial data is expected. When masks are provided the whole of the data + is expected. For example, if env_ids is provided, then root_pose should be of shape (len(env_ids), 7). If + env_mask is provided, then root_pose should be of shape (num_instances, 7). + + ..caution:: If both env_mask and env_ids are provided, then env_mask will be used. The function will thus + expect the whole of the data to be provided. If none of them are provided, then the function expects the whole + of the data to be provided. + + Args: + root_pose: Root link poses in simulation frame. Shape is (len(env_ids), 7) or (num_instances, 7). + env_mask: Environment mask. Shape is (num_instances,). + env_ids: Environment indices. If None, then all indices are used. + """ + raise NotImplementedError() + + @abstractmethod + def write_root_com_pose_to_sim( + self, + root_pose: torch.Tensor | wp.array, + env_ids: Sequence[int] | None = None, + env_mask: torch.Tensor | wp.array | None = None, + ) -> None: + """Set the root center of mass pose over selected environment indices into the simulation. + + The root pose comprises of the cartesian position and quaternion orientation in (x, y, z, w). + The orientation is the orientation of the principle axes of inertia. + + ..note:: When ids are provided, then partial data is expected. When masks are provided the whole of the data + is expected. For example, if env_ids is provided, then root_pose should be of shape (len(env_ids), 7). If + env_mask is provided, then root_pose should be of shape (num_instances, 7). + + ..caution:: If both env_mask and env_ids are provided, then env_mask will be used. The function will thus + expect the whole of the data to be provided. If none of them are provided, then the function expects the whole + of the data to be provided. + + Args: + root_pose: Root center of mass poses in simulation frame. Shape is (len(env_ids), 7) or (num_instances, 7). + env_mask: Environment mask. Shape is (num_instances,). + env_ids: Environment indices. If None, then all indices are used. + """ + # resolve all indices + raise NotImplementedError() + + @abstractmethod + def write_root_velocity_to_sim( + self, + root_velocity: torch.Tensor | wp.array, + env_ids: Sequence[int] | None = None, + env_mask: torch.Tensor | wp.array | None = None, + ) -> None: + """Set the root center of mass velocity over selected environment indices into the simulation. + + The velocity comprises linear velocity (x, y, z) and angular velocity (x, y, z) in that order. + ..note:: This sets the velocity of the root's center of mass rather than the roots frame. + + ..note:: When ids are provided, then partial data is expected. When masks are provided the whole of the data + is expected. For example, if env_ids is provided, then root_velocity should be of shape (len(env_ids), 6). If + env_mask is provided, then root_velocity should be of shape (num_instances, 6). + + ..caution:: If both env_mask and env_ids are provided, then env_mask will be used. The function will thus + expect the whole of the data to be provided. If none of them are provided, then the function expects the whole + of the data to be provided. + + Args: + root_velocity: Root center of mass velocities in simulation world frame. Shape is (len(env_ids), 6) or (num_instances, 6). + env_mask: Environment mask. Shape is (num_instances,). + env_ids: Environment indices. If None, then all indices are used. + """ + raise NotImplementedError() + + @abstractmethod + def write_root_com_velocity_to_sim( + self, root_velocity: torch.Tensor | wp.array, + env_ids: Sequence[int] | None = None, + env_mask: torch.Tensor | wp.array | None = None, + ) -> None: + """Set the root center of mass velocity over selected environment indices into the simulation. + + The velocity comprises linear velocity (x, y, z) and angular velocity (x, y, z) in that order. + ..note:: This sets the velocity of the root's center of mass rather than the roots frame. + + ..note:: When ids are provided, then partial data is expected. When masks are provided the whole of the data + is expected. For example, if env_ids is provided, then root_velocity should be of shape (len(env_ids), 6). If + env_mask is provided, then root_velocity should be of shape (num_instances, 6). + + ..caution:: If both env_mask and env_ids are provided, then env_mask will be used. The function will thus + expect the whole of the data to be provided. If none of them are provided, then the function expects the whole + of the data to be provided. + + Args: + root_velocity: Root center of mass velocities in simulation world frame. Shape is (len(env_ids), 6) or (num_instances, 6). + env_mask: Environment mask. Shape is (num_instances,). + env_ids: Environment indices. If None, then all indices are used. + """ + raise NotImplementedError() + + @abstractmethod + def write_root_link_velocity_to_sim( + self, + root_velocity: torch.Tensor | wp.array, + env_ids: Sequence[int] | None = None, + env_mask: torch.Tensor | wp.array | None = None, + ) -> None: + """Set the root link velocity over selected environment indices into the simulation. + + The velocity comprises linear velocity (x, y, z) and angular velocity (x, y, z) in that order. + ..note:: This sets the velocity of the root's frame rather than the roots center of mass. + + ..note:: When ids are provided, then partial data is expected. When masks are provided the whole of the data + is expected. For example, if env_ids is provided, then root_velocity should be of shape (len(env_ids), 6). If + env_mask is provided, then root_velocity should be of shape (num_instances, 6). + + ..caution:: If both env_mask and env_ids are provided, then env_mask will be used. The function will thus + expect the whole of the data to be provided. If none of them are provided, then the function expects the whole + of the data to be provided. + + Args: + root_velocity: Root frame velocities in simulation world frame. Shape is (len(env_ids), 6) or (num_instances, 6). + env_mask: Environment mask. Shape is (num_instances,). + env_ids: Environment indices. If None, then all indices are used. + """ + raise NotImplementedError() + + """ + Operations - Setters. + """ + + @abstractmethod + def set_external_force_and_torque( + self, + forces: torch.Tensor | wp.array, + torques: torch.Tensor | wp.array, + positions: torch.Tensor | wp.array | None = None, + body_ids: Sequence[int] | slice | None = None, + env_ids: Sequence[int] | None = None, + body_mask: torch.Tensor | wp.array | None = None, + env_mask: torch.Tensor | wp.array | None = None, + is_global: bool = False, + ) -> None: + """Set external force and torque to apply on the asset's bodies in their local frame. + + For many applications, we want to keep the applied external force on rigid bodies constant over a period of + time (for instance, during the policy control). This function allows us to store the external force and torque + into buffers which are then applied to the simulation at every step. Optionally, set the position to apply the + external wrench at (in the local link frame of the bodies). + + .. caution:: + If the function is called with empty forces and torques, then this function disables the application + of external wrench to the simulation. + + .. code-block:: python + + # example of disabling external wrench + asset.set_external_force_and_torque(forces=torch.zeros(0, 3), torques=torch.zeros(0, 3)) + + .. caution:: + If the function is called consecutively with and with different values for ``is_global``, then the + all the external wrenches will be applied in the frame specified by the last call. + + .. code-block:: python + + # example of setting external wrench in the global frame + asset.set_external_force_and_torque(forces=torch.ones(1, 1, 3), env_ids=[0], is_global=True) + # example of setting external wrench in the link frame + asset.set_external_force_and_torque(forces=torch.ones(1, 1, 3), env_ids=[1], is_global=False) + # Both environments will have the external wrenches applied in the link frame + + .. note:: + This function does not apply the external wrench to the simulation. It only fills the buffers with + the desired values. To apply the external wrench, call the :meth:`write_data_to_sim` function + right before the simulation step. + + Args: + forces: External forces in bodies' local frame. Shape is (len(env_ids), len(body_ids), 3). + torques: External torques in bodies' local frame. Shape is (len(env_ids), len(body_ids), 3). + positions: External wrench positions in bodies' local frame. Shape is (len(env_ids), len(body_ids), 3). Defaults to None. + body_ids: Body indices to apply external wrench to. Defaults to None (all bodies). + env_ids: Environment indices to apply external wrench to. Defaults to None (all instances). + is_global: Whether to apply the external wrench in the global frame. Defaults to False. If set to False, + the external wrench is applied in the link frame of the bodies. + """ + raise NotImplementedError() + + """ + Internal helper. + """ + + @abstractmethod + def _initialize_impl(self): + raise NotImplementedError() + + @abstractmethod + def _create_buffers(self): + raise NotImplementedError() + + @abstractmethod + def _process_cfg(self) -> None: + """Post processing of configuration parameters.""" + raise NotImplementedError() + + """ + Internal simulation callbacks. + """ + + @abstractmethod + def _invalidate_initialize_callback(self, event) -> None: + """Invalidates the scene elements.""" + raise NotImplementedError() diff --git a/source/isaaclab/isaaclab/assets/rigid_object/base_rigid_object_data.py b/source/isaaclab/isaaclab/assets/rigid_object/base_rigid_object_data.py new file mode 100644 index 00000000000..86cf591853f --- /dev/null +++ b/source/isaaclab/isaaclab/assets/rigid_object/base_rigid_object_data.py @@ -0,0 +1,615 @@ +# Copyright (c) 2022-2025, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +import weakref + +import omni.physics.tensors.impl.api as physx + +import isaaclab.utils.math as math_utils +from isaaclab.sim.utils import get_current_stage_id +from isaaclab.utils.buffers import TimestampedBuffer + +import torch +import warp as wp +from abc import ABC, abstractmethod + +class BaseRigidObjectData(ABC): + """Data container for a rigid object. + + This class contains the data for a rigid object in the simulation. The data includes the state of + the root rigid body and the state of all the bodies in the object. The data is stored in the simulation + world frame unless otherwise specified. + + For a rigid body, there are two frames of reference that are used: + + - Actor frame: The frame of reference of the rigid body prim. This typically corresponds to the Xform prim + with the rigid body schema. + - Center of mass frame: The frame of reference of the center of mass of the rigid body. + + Depending on the settings of the simulation, the actor frame and the center of mass frame may be the same. + This needs to be taken into account when interpreting the data. + + The data is lazily updated, meaning that the data is only updated when it is accessed. This is useful + when the data is expensive to compute or retrieve. The data is updated when the timestamp of the buffer + is older than the current simulation timestamp. The timestamp is updated whenever the data is updated. + """ + + def __init__(self, root_view, device: str): + """Initializes the rigid object data. + + Args: + root_physx_view: The root rigid body view. + device: The device used for processing. + """ + # Set the parameters + self.device = device + + @abstractmethod + def update(self, dt: float) -> None: + """Updates the data for the rigid object. + + Args: + dt: The time step for the update. This must be a positive value. + """ + raise NotImplementedError() + + ## + # Names. + ## + + body_names: list[str] = None + """Body names in the order parsed by the simulation view.""" + + ## + # Defaults. + ## + + default_root_state: torch.Tensor = None + """Default root state ``[pos, quat, lin_vel, ang_vel]`` in local environment frame. Shape is (num_instances, 13). + + The position and quaternion are of the rigid body's actor frame. Meanwhile, the linear and angular velocities are + of the center of mass frame. + """ + + ## + # Root state properties. + ## + + @property + @abstractmethod + def root_link_pose_w(self) -> torch.Tensor: + """Root link pose ``[pos, quat]`` in simulation world frame. Shape is (num_instances, 7). + + This quantity is the pose of the actor frame of the root rigid body relative to the world. + The orientation is provided in (w, x, y, z) format. + """ + raise NotImplementedError() + + @property + @abstractmethod + def root_link_vel_w(self) -> torch.Tensor: + """Root link velocity ``[lin_vel, ang_vel]`` in simulation world frame. Shape is (num_instances, 6). + + This quantity contains the linear and angular velocities of the actor frame of the root + rigid body relative to the world. + """ + raise NotImplementedError() + + @property + @abstractmethod + def root_com_pose_w(self) -> torch.Tensor: + """Root center of mass pose ``[pos, quat]`` in simulation world frame. Shape is (num_instances, 7). + + This quantity is the pose of the center of mass frame of the root rigid body relative to the world. + The orientation is provided in (w, x, y, z) format. + """ + raise NotImplementedError() + + @property + @abstractmethod + def root_com_vel_w(self) -> torch.Tensor: + """Root center of mass velocity ``[lin_vel, ang_vel]`` in simulation world frame. Shape is (num_instances, 6). + + This quantity contains the linear and angular velocities of the root rigid body's center of mass frame + relative to the world. + """ + raise NotImplementedError() + + @property + @abstractmethod + def root_state_w(self) -> torch.Tensor: + """Root state ``[pos, quat, lin_vel, ang_vel]`` in simulation world frame. Shape is (num_instances, 13). + + The position and orientation are of the rigid body's actor frame. Meanwhile, the linear and angular + velocities are of the rigid body's center of mass frame. + """ + raise NotImplementedError() + + @property + @abstractmethod + def root_link_state_w(self) -> torch.Tensor: + """Root state ``[pos, quat, lin_vel, ang_vel]`` in simulation world frame. Shape is (num_instances, 13). + + The position, quaternion, and linear/angular velocity are of the rigid body root frame relative to the + world. The orientation is provided in (w, x, y, z) format. + """ + raise NotImplementedError() + + @property + @abstractmethod + def root_com_state_w(self) -> torch.Tensor: + """Root center of mass state ``[pos, quat, lin_vel, ang_vel]`` in simulation world frame. + Shape is (num_instances, 13). + + The position, quaternion, and linear/angular velocity are of the rigid body's center of mass frame + relative to the world. Center of mass frame is the orientation principle axes of inertia. + """ + raise NotImplementedError() + + ## + # Body state properties. + ## + + @property + @abstractmethod + def body_link_pose_w(self) -> torch.Tensor: + """Body link pose ``[pos, quat]`` in simulation world frame. Shape is (num_instances, 1, 7). + + This quantity is the pose of the actor frame of the rigid body relative to the world. + The orientation is provided in (w, x, y, z) format. + """ + raise NotImplementedError() + + @property + @abstractmethod + def body_link_vel_w(self) -> torch.Tensor: + """Body link velocity ``[lin_vel, ang_vel]`` in simulation world frame. Shape is (num_instances, 1, 6). + + This quantity contains the linear and angular velocities of the actor frame of the root + rigid body relative to the world. + """ + raise NotImplementedError() + + @property + @abstractmethod + def body_com_pose_w(self) -> torch.Tensor: + """Body center of mass pose ``[pos, quat]`` in simulation world frame. Shape is (num_instances, 1, 7). + + This quantity is the pose of the center of mass frame of the rigid body relative to the world. + The orientation is provided in (w, x, y, z) format. + """ + raise NotImplementedError() + + @property + @abstractmethod + def body_com_vel_w(self) -> torch.Tensor: + """Body center of mass velocity ``[lin_vel, ang_vel]`` in simulation world frame. + Shape is (num_instances, 1, 6). + + This quantity contains the linear and angular velocities of the root rigid body's center of mass frame + relative to the world. + """ + raise NotImplementedError() + + @property + @abstractmethod + def body_state_w(self) -> torch.Tensor: + """State of all bodies `[pos, quat, lin_vel, ang_vel]` in simulation world frame. + Shape is (num_instances, 1, 13). + + The position and orientation are of the rigid bodies' actor frame. Meanwhile, the linear and angular + velocities are of the rigid bodies' center of mass frame. + """ + raise NotImplementedError() + + @property + @abstractmethod + def body_link_state_w(self) -> torch.Tensor: + """State of all bodies ``[pos, quat, lin_vel, ang_vel]`` in simulation world frame. + Shape is (num_instances, 1, 13). + + The position, quaternion, and linear/angular velocity are of the body's link frame relative to the world. + The orientation is provided in (w, x, y, z) format. + """ + raise NotImplementedError() + + @property + @abstractmethod + def body_com_state_w(self) -> torch.Tensor: + """State of all bodies ``[pos, quat, lin_vel, ang_vel]`` in simulation world frame. + Shape is (num_instances, num_bodies, 13). + + The position, quaternion, and linear/angular velocity are of the body's center of mass frame relative to the + world. Center of mass frame is assumed to be the same orientation as the link rather than the orientation of the + principle inertia. The orientation is provided in (w, x, y, z) format. + """ + raise NotImplementedError() + + @property + @abstractmethod + def body_com_acc_w(self) -> torch.Tensor: + """Acceleration of all bodies ``[lin_acc, ang_acc]`` in the simulation world frame. + Shape is (num_instances, 1, 6). + + This quantity is the acceleration of the rigid bodies' center of mass frame relative to the world. + """ + raise NotImplementedError() + + @property + @abstractmethod + def body_com_pose_b(self) -> torch.Tensor: + """Center of mass pose ``[pos, quat]`` of all bodies in their respective body's link frames. + Shape is (num_instances, 1, 7). + + This quantity is the pose of the center of mass frame of the rigid body relative to the body's link frame. + The orientation is provided in (w, x, y, z) format. + """ + raise NotImplementedError() + + ## + # Derived Properties. + ## + + @property + @abstractmethod + def projected_gravity_b(self) -> torch.Tensor: + """Projection of the gravity direction on base frame. Shape is (num_instances, 3).""" + raise NotImplementedError() + + @property + @abstractmethod + def heading_w(self) -> torch.Tensor: + """Yaw heading of the base frame (in radians). Shape is (num_instances,). + + Note: + This quantity is computed by assuming that the forward-direction of the base + frame is along x-direction, i.e. :math:`(1, 0, 0)`. + """ + raise NotImplementedError() + + @property + @abstractmethod + def root_link_lin_vel_b(self) -> torch.Tensor: + """Root link linear velocity in base frame. Shape is (num_instances, 3). + + This quantity is the linear velocity of the actor frame of the root rigid body frame with respect to the + rigid body's actor frame. + """ + raise NotImplementedError() + + @property + @abstractmethod + def root_link_ang_vel_b(self) -> torch.Tensor: + """Root link angular velocity in base world frame. Shape is (num_instances, 3). + + This quantity is the angular velocity of the actor frame of the root rigid body frame with respect to the + rigid body's actor frame. + """ + raise NotImplementedError() + + @property + @abstractmethod + def root_com_lin_vel_b(self) -> torch.Tensor: + """Root center of mass linear velocity in base frame. Shape is (num_instances, 3). + + This quantity is the linear velocity of the root rigid body's center of mass frame with respect to the + rigid body's actor frame. + """ + raise NotImplementedError() + + @property + @abstractmethod + def root_com_ang_vel_b(self) -> torch.Tensor: + """Root center of mass angular velocity in base world frame. Shape is (num_instances, 3). + + This quantity is the angular velocity of the root rigid body's center of mass frame with respect to the + rigid body's actor frame. + """ + raise NotImplementedError() + + ## + # Sliced properties. + ## + + @property + @abstractmethod + def root_link_pos_w(self) -> torch.Tensor: + """Root link position in simulation world frame. Shape is (num_instances, 3). + + This quantity is the position of the actor frame of the root rigid body relative to the world. + """ + raise NotImplementedError() + + @property + @abstractmethod + def root_link_quat_w(self) -> torch.Tensor: + """Root link orientation (w, x, y, z) in simulation world frame. Shape is (num_instances, 4). + + This quantity is the orientation of the actor frame of the root rigid body. + """ + raise NotImplementedError() + + @property + @abstractmethod + def root_link_lin_vel_w(self) -> torch.Tensor: + """Root linear velocity in simulation world frame. Shape is (num_instances, 3). + + This quantity is the linear velocity of the root rigid body's actor frame relative to the world. + """ + raise NotImplementedError() + + @property + @abstractmethod + def root_link_ang_vel_w(self) -> torch.Tensor: + """Root link angular velocity in simulation world frame. Shape is (num_instances, 3). + + This quantity is the angular velocity of the actor frame of the root rigid body relative to the world. + """ + raise NotImplementedError() + + @property + @abstractmethod + def root_com_pos_w(self) -> torch.Tensor: + """Root center of mass position in simulation world frame. Shape is (num_instances, 3). + + This quantity is the position of the actor frame of the root rigid body relative to the world. + """ + raise NotImplementedError() + + @property + @abstractmethod + def root_com_quat_w(self) -> torch.Tensor: + """Root center of mass orientation (w, x, y, z) in simulation world frame. Shape is (num_instances, 4). + + This quantity is the orientation of the actor frame of the root rigid body relative to the world. + """ + raise NotImplementedError() + + @property + @abstractmethod + def root_com_lin_vel_w(self) -> torch.Tensor: + """Root center of mass linear velocity in simulation world frame. Shape is (num_instances, 3). + + This quantity is the linear velocity of the root rigid body's center of mass frame relative to the world. + """ + raise NotImplementedError() + + @property + @abstractmethod + def root_com_ang_vel_w(self) -> torch.Tensor: + """Root center of mass angular velocity in simulation world frame. Shape is (num_instances, 3). + + This quantity is the angular velocity of the root rigid body's center of mass frame relative to the world. + """ + raise NotImplementedError() + + @property + @abstractmethod + def body_link_pos_w(self) -> torch.Tensor: + """Positions of all bodies in simulation world frame. Shape is (num_instances, 1, 3). + + This quantity is the position of the rigid bodies' actor frame relative to the world. + """ + raise NotImplementedError() + + @property + @abstractmethod + def body_link_quat_w(self) -> torch.Tensor: + """Orientation (w, x, y, z) of all bodies in simulation world frame. Shape is (num_instances, 1, 4). + + This quantity is the orientation of the rigid bodies' actor frame relative to the world. + """ + raise NotImplementedError() + + @property + @abstractmethod + def body_link_lin_vel_w(self) -> torch.Tensor: + """Linear velocity of all bodies in simulation world frame. Shape is (num_instances, 1, 3). + + This quantity is the linear velocity of the rigid bodies' center of mass frame relative to the world. + """ + raise NotImplementedError() + + @property + @abstractmethod + def body_link_ang_vel_w(self) -> torch.Tensor: + """Angular velocity of all bodies in simulation world frame. Shape is (num_instances, 1, 3). + + This quantity is the angular velocity of the rigid bodies' center of mass frame relative to the world. + """ + raise NotImplementedError() + + @property + @abstractmethod + def body_com_pos_w(self) -> torch.Tensor: + """Positions of all bodies in simulation world frame. Shape is (num_instances, 1, 3). + + This quantity is the position of the rigid bodies' actor frame. + """ + raise NotImplementedError() + + @property + @abstractmethod + def body_com_quat_w(self) -> torch.Tensor: + """Orientation (w, x, y, z) of the principle axis of inertia of all bodies in simulation world frame. + + Shape is (num_instances, 1, 4). This quantity is the orientation of the rigid bodies' actor frame. + """ + raise NotImplementedError() + + @property + @abstractmethod + def body_com_lin_vel_w(self) -> torch.Tensor: + """Linear velocity of all bodies in simulation world frame. Shape is (num_instances, 1, 3). + + This quantity is the linear velocity of the rigid bodies' center of mass frame. + """ + raise NotImplementedError() + + @property + @abstractmethod + def body_com_ang_vel_w(self) -> torch.Tensor: + """Angular velocity of all bodies in simulation world frame. Shape is (num_instances, 1, 3). + + This quantity is the angular velocity of the rigid bodies' center of mass frame. + """ + raise NotImplementedError() + + @property + @abstractmethod + def body_com_lin_acc_w(self) -> torch.Tensor: + """Linear acceleration of all bodies in simulation world frame. Shape is (num_instances, 1, 3). + + This quantity is the linear acceleration of the rigid bodies' center of mass frame. + """ + raise NotImplementedError() + + @property + @abstractmethod + def body_com_ang_acc_w(self) -> torch.Tensor: + """Angular acceleration of all bodies in simulation world frame. Shape is (num_instances, 1, 3). + + This quantity is the angular acceleration of the rigid bodies' center of mass frame. + """ + raise NotImplementedError() + + @property + @abstractmethod + def body_com_pos_b(self) -> torch.Tensor: + """Center of mass position of all of the bodies in their respective link frames. + Shape is (num_instances, 1, 3). + + This quantity is the center of mass location relative to its body'slink frame. + """ + raise NotImplementedError() + + @property + @abstractmethod + def body_com_quat_b(self) -> torch.Tensor: + """Orientation (w, x, y, z) of the principle axis of inertia of all of the bodies in their + respective link frames. Shape is (num_instances, 1, 4). + + This quantity is the orientation of the principles axes of inertia relative to its body's link frame. + """ + raise NotImplementedError() + + ## + # Properties for backwards compatibility. + ## + + @property + @abstractmethod + def root_pose_w(self) -> torch.Tensor: + """Same as :attr:`root_link_pose_w`.""" + raise NotImplementedError() + + @property + @abstractmethod + def root_pos_w(self) -> torch.Tensor: + """Same as :attr:`root_link_pos_w`.""" + raise NotImplementedError() + + @property + @abstractmethod + def root_quat_w(self) -> torch.Tensor: + """Same as :attr:`root_link_quat_w`.""" + raise NotImplementedError() + + @property + @abstractmethod + def root_vel_w(self) -> torch.Tensor: + """Same as :attr:`root_com_vel_w`.""" + raise NotImplementedError() + + @property + @abstractmethod + def root_lin_vel_w(self) -> torch.Tensor: + """Same as :attr:`root_com_lin_vel_w`.""" + raise NotImplementedError() + + @property + @abstractmethod + def root_ang_vel_w(self) -> torch.Tensor: + """Same as :attr:`root_com_ang_vel_w`.""" + raise NotImplementedError() + + @property + @abstractmethod + def root_lin_vel_b(self) -> torch.Tensor: + """Same as :attr:`root_com_lin_vel_b`.""" + raise NotImplementedError() + + @property + @abstractmethod + def root_ang_vel_b(self) -> torch.Tensor: + """Same as :attr:`root_com_ang_vel_b`.""" + raise NotImplementedError() + + @property + @abstractmethod + def body_pose_w(self) -> torch.Tensor: + """Same as :attr:`body_link_pose_w`.""" + raise NotImplementedError() + + @property + @abstractmethod + def body_pos_w(self) -> torch.Tensor: + """Same as :attr:`body_link_pos_w`.""" + raise NotImplementedError() + + @property + @abstractmethod + def body_quat_w(self) -> torch.Tensor: + """Same as :attr:`body_link_quat_w`.""" + raise NotImplementedError() + + @property + @abstractmethod + def body_vel_w(self) -> torch.Tensor: + """Same as :attr:`body_com_vel_w`.""" + raise NotImplementedError() + + @property + @abstractmethod + def body_lin_vel_w(self) -> torch.Tensor: + """Same as :attr:`body_com_lin_vel_w`.""" + raise NotImplementedError() + + @property + @abstractmethod + def body_ang_vel_w(self) -> torch.Tensor: + """Same as :attr:`body_com_ang_vel_w`.""" + raise NotImplementedError() + + @property + @abstractmethod + def body_acc_w(self) -> torch.Tensor: + """Same as :attr:`body_com_acc_w`.""" + raise NotImplementedError() + + @property + @abstractmethod + def body_lin_acc_w(self) -> torch.Tensor: + """Same as :attr:`body_com_lin_acc_w`.""" + raise NotImplementedError() + + @property + @abstractmethod + def body_ang_acc_w(self) -> torch.Tensor: + """Same as :attr:`body_com_ang_acc_w`.""" + raise NotImplementedError() + + @property + @abstractmethod + def com_pos_b(self) -> torch.Tensor: + """Same as :attr:`body_com_pos_b`.""" + raise NotImplementedError() + + @property + @abstractmethod + def com_quat_b(self) -> torch.Tensor: + """Same as :attr:`body_com_quat_b`.""" + + raise NotImplementedError() \ No newline at end of file diff --git a/source/isaaclab/isaaclab/assets/rigid_object/rigid_object.py b/source/isaaclab/isaaclab/assets/rigid_object/rigid_object.py new file mode 100644 index 00000000000..704cd9ec930 --- /dev/null +++ b/source/isaaclab/isaaclab/assets/rigid_object/rigid_object.py @@ -0,0 +1,12 @@ +from .base_rigid_object import BaseRigidObject +from isaaclab.utils.backend_utils import FactoryBase + +class RigidObject(FactoryBase): + """Factory for creating articulation instances.""" + + def __new__(cls, *args, **kwargs) -> BaseRigidObject: + """Create a new instance of an articulation based on the backend.""" + # The `FactoryBase` __new__ method will handle the logic and return + # an instance of the correct backend-specific articulation class, + # which is guaranteed to be a subclass of `BaseArticulation` by convention. + return super().__new__(cls, *args, **kwargs) \ No newline at end of file diff --git a/source/isaaclab/isaaclab/assets/rigid_object/rigid_object_cfg.py b/source/isaaclab/isaaclab/assets/rigid_object/rigid_object_cfg.py new file mode 100644 index 00000000000..1cd24bcc918 --- /dev/null +++ b/source/isaaclab/isaaclab/assets/rigid_object/rigid_object_cfg.py @@ -0,0 +1,32 @@ +# Copyright (c) 2022-2025, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +from isaaclab.utils import configclass + +from ..asset_base_cfg import AssetBaseCfg +from .rigid_object import RigidObject + + +@configclass +class RigidObjectCfg(AssetBaseCfg): + """Configuration parameters for a rigid object.""" + + @configclass + class InitialStateCfg(AssetBaseCfg.InitialStateCfg): + """Initial state of the rigid body.""" + + lin_vel: tuple[float, float, float] = (0.0, 0.0, 0.0) + """Linear velocity of the root in simulation world frame. Defaults to (0.0, 0.0, 0.0).""" + ang_vel: tuple[float, float, float] = (0.0, 0.0, 0.0) + """Angular velocity of the root in simulation world frame. Defaults to (0.0, 0.0, 0.0).""" + + ## + # Initialize configurations. + ## + + class_type: type = RigidObject + + init_state: InitialStateCfg = InitialStateCfg() + """Initial state of the rigid object. Defaults to identity pose with zero velocity.""" diff --git a/source/isaaclab/isaaclab/assets/rigid_object/rigid_object_data.py b/source/isaaclab/isaaclab/assets/rigid_object/rigid_object_data.py new file mode 100644 index 00000000000..778be3a4cb1 --- /dev/null +++ b/source/isaaclab/isaaclab/assets/rigid_object/rigid_object_data.py @@ -0,0 +1,9 @@ +from .base_rigid_object_data import BaseRigidObjectData +from isaaclab.utils.backend_utils import FactoryBase + +class RigidObjectData(FactoryBase): + """Factory for creating rigid object data instances.""" + + def __new__(cls, *args, **kwargs) -> BaseRigidObjectData: + """Create a new instance of a rigid object data based on the backend.""" + return super().__new__(cls, *args, **kwargs) \ No newline at end of file diff --git a/source/isaaclab/isaaclab/cloner/grid_cloner.py b/source/isaaclab/isaaclab/cloner/grid_cloner.py index a7feb6842a6..9e529463a65 100644 --- a/source/isaaclab/isaaclab/cloner/grid_cloner.py +++ b/source/isaaclab/isaaclab/cloner/grid_cloner.py @@ -129,7 +129,6 @@ def get_clone_transforms( return positions, orientations - @Timer(name="newton_clone", msg="Clone took:", enable=True, format="ms") def clone( self, source_prim_path: str, diff --git a/source/isaaclab/isaaclab/envs/__init__.py b/source/isaaclab/isaaclab/envs/__init__.py index 12190ff04c3..5360f1ae5d1 100644 --- a/source/isaaclab/isaaclab/envs/__init__.py +++ b/source/isaaclab/isaaclab/envs/__init__.py @@ -46,6 +46,7 @@ from .common import VecEnvObs, VecEnvStepReturn, ViewerCfg from .direct_rl_env import DirectRLEnv from .direct_rl_env_cfg import DirectRLEnvCfg +from .direct_rl_env_warp import DirectRLEnvWarp from .manager_based_env import ManagerBasedEnv from .manager_based_env_cfg import ManagerBasedEnvCfg from .manager_based_rl_env import ManagerBasedRLEnv diff --git a/source/isaaclab/isaaclab/envs/direct_rl_env.py b/source/isaaclab/isaaclab/envs/direct_rl_env.py index 84d7d287515..743f5c61d57 100644 --- a/source/isaaclab/isaaclab/envs/direct_rl_env.py +++ b/source/isaaclab/isaaclab/envs/direct_rl_env.py @@ -350,9 +350,10 @@ def step(self, action: torch.Tensor) -> VecEnvStepReturn: for _ in range(self.cfg.decimation): self._sim_step_counter += 1 # set actions into buffers - self._apply_action() - # set actions into simulator - self.scene.write_data_to_sim() + with Timer(name="apply_action", msg="Action processing step took:", enable=True, format="us"): + self._apply_action() + # set actions into simulator + self.scene.write_data_to_sim() # simulate with Timer(name="simulate", msg="Newton simulation step took:", enable=True, format="us"): self.sim.step(render=False) @@ -364,39 +365,40 @@ def step(self, action: torch.Tensor) -> VecEnvStepReturn: # update buffers at sim dt self.scene.update(dt=self.physics_dt) - # post-step: - # -- update env counters (used for curriculum generation) - self.episode_length_buf += 1 # step in current episode (per env) - self.common_step_counter += 1 # total step (common for all envs) - - self.reset_terminated[:], self.reset_time_outs[:] = self._get_dones() - self.reset_buf = self.reset_terminated | self.reset_time_outs - self.reward_buf = self._get_rewards() - - # -- reset envs that terminated/timed-out and log the episode information - reset_env_ids = self.reset_buf.nonzero(as_tuple=False).squeeze(-1) - if len(reset_env_ids) > 0: - self._reset_idx(reset_env_ids) - # update articulation kinematics - self.scene.write_data_to_sim() - # if sensors are added to the scene, make sure we render to reflect changes in reset - if self.sim.has_rtx_sensors() and self.cfg.rerender_on_reset: - self.sim.render() - - # post-step: step interval event - if self.cfg.events: - if "interval" in self.event_manager.available_modes: - self.event_manager.apply(mode="interval", dt=self.step_dt) + with Timer(name="post_processing", msg="Post-Processing step took:", enable=True, format="us"): + # post-step: + # -- update env counters (used for curriculum generation) + self.episode_length_buf += 1 # step in current episode (per env) + self.common_step_counter += 1 # total step (common for all envs) + + self.reset_terminated[:], self.reset_time_outs[:] = self._get_dones() + self.reset_buf = self.reset_terminated | self.reset_time_outs + self.reward_buf = self._get_rewards() + + # -- reset envs that terminated/timed-out and log the episode information + reset_env_ids = self.reset_buf.nonzero(as_tuple=False).squeeze(-1) + if len(reset_env_ids) > 0: + self._reset_idx(reset_env_ids) + # update articulation kinematics + self.scene.write_data_to_sim() + # if sensors are added to the scene, make sure we render to reflect changes in reset + if self.sim.has_rtx_sensors() and self.cfg.rerender_on_reset: + self.sim.render() + + # post-step: step interval event + if self.cfg.events: + if "interval" in self.event_manager.available_modes: + self.event_manager.apply(mode="interval", dt=self.step_dt) - # update observations - self.obs_buf = self._get_observations() + # update observations + self.obs_buf = self._get_observations() - # add observation noise - # note: we apply no noise to the state space (since it is used for critic networks) - if self.cfg.observation_noise_model: - self.obs_buf["policy"] = self._observation_noise_model(self.obs_buf["policy"]) + # add observation noise + # note: we apply no noise to the state space (since it is used for critic networks) + if self.cfg.observation_noise_model: + self.obs_buf["policy"] = self._observation_noise_model(self.obs_buf["policy"]) - # return observations, rewards, resets and extras + # return observations, rewards, resets and extras return self.obs_buf, self.reward_buf, self.reset_terminated, self.reset_time_outs, self.extras @staticmethod diff --git a/source/isaaclab/isaaclab/envs/direct_rl_env_warp.py b/source/isaaclab/isaaclab/envs/direct_rl_env_warp.py new file mode 100644 index 00000000000..2d1595000e9 --- /dev/null +++ b/source/isaaclab/isaaclab/envs/direct_rl_env_warp.py @@ -0,0 +1,763 @@ +# Copyright (c) 2022-2025, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +from __future__ import annotations + +import builtins +import gymnasium as gym +import inspect +import math +import numpy as np +import torch +import weakref +from abc import abstractmethod +from dataclasses import MISSING +from typing import Any, ClassVar, Sequence + +import isaacsim.core.utils.torch as torch_utils +import omni.kit.app +import omni.log +import omni.physx +import warp as wp +from isaacsim.core.simulation_manager import SimulationManager +from isaacsim.core.version import get_version + +from isaaclab.managers import EventManager +from isaaclab.scene import InteractiveScene +from isaaclab.sim import SimulationContext +from isaaclab.sim.utils import attach_stage_to_usd_context, use_stage +from isaaclab.utils.noise import NoiseModel +from isaaclab.utils.timer import Timer + +from .common import VecEnvObs, VecEnvStepReturn +from .direct_rl_env_cfg import DirectRLEnvCfg +from .ui import ViewportCameraController +from .utils.spaces import sample_space, spec_to_gym_space + + +@wp.kernel +def zero_mask_int32( + mask: wp.array(dtype=wp.bool), + data: wp.array(dtype=wp.int32), +): + env_index = wp.tid() + if mask[env_index]: + data[env_index] = 0 + + +@wp.kernel +def add_to_env( + data: wp.array(dtype=wp.int32), + value: wp.int32, +): + env_index = wp.tid() + data[env_index] += value + + +class DirectRLEnvWarp(gym.Env): + """The superclass for the direct workflow to design environments. + + This class implements the core functionality for reinforcement learning (RL) + environments. It is designed to be used with any RL library. The class is designed + to be used with vectorized environments, i.e., the environment is expected to be run + in parallel with multiple sub-environments. + + While the environment itself is implemented as a vectorized environment, we do not + inherit from :class:`gym.vector.VectorEnv`. This is mainly because the class adds + various methods (for wait and asynchronous updates) which are not required. + Additionally, each RL library typically has its own definition for a vectorized + environment. Thus, to reduce complexity, we directly use the :class:`gym.Env` over + here and leave it up to library-defined wrappers to take care of wrapping this + environment for their agents. + + Note: + For vectorized environments, it is recommended to **only** call the :meth:`reset` + method once before the first call to :meth:`step`, i.e. after the environment is created. + After that, the :meth:`step` function handles the reset of terminated sub-environments. + in a vectorized environment. + + """ + + is_vector_env: ClassVar[bool] = True + """Whether the environment is a vectorized environment.""" + metadata: ClassVar[dict[str, Any]] = { + "render_modes": [None, "human", "rgb_array"], + "isaac_sim_version": get_version(), + } + """Metadata for the environment.""" + + def __init__(self, cfg: DirectRLEnvCfg, render_mode: str | None = None, **kwargs): + """Initialize the environment. + + Args: + cfg: The configuration object for the environment. + render_mode: The render mode for the environment. Defaults to None, which + is similar to ``"human"``. + + Raises: + RuntimeError: If a simulation context already exists. The environment must always create one + since it configures the simulation context and controls the simulation. + """ + # check that the config is valid + cfg.validate() + # store inputs to class + self.cfg = cfg + # store the render mode + self.render_mode = render_mode + # initialize internal variables + self._is_closed = False + + # set the seed for the environment + if self.cfg.seed is not None: + self.cfg.seed = self.seed(self.cfg.seed) + else: + omni.log.warn("Seed not set for the environment. The environment creation may not be deterministic.") + + # create a simulation context to control the simulator + if SimulationContext.instance() is None: + self.sim: SimulationContext = SimulationContext(self.cfg.sim) + else: + raise RuntimeError("Simulation context already exists. Cannot create a new one.") + + # make sure torch is running on the correct device + if "cuda" in self.device: + torch.cuda.set_device(self.device) + + # print useful information + print("[INFO]: Base environment:") + print(f"\tEnvironment device : {self.device}") + print(f"\tEnvironment seed : {self.cfg.seed}") + print(f"\tPhysics step-size : {self.physics_dt}") + print(f"\tRendering step-size : {self.physics_dt * self.cfg.sim.render_interval}") + print(f"\tEnvironment step-size : {self.step_dt}") + + if self.cfg.sim.render_interval < self.cfg.decimation: + msg = ( + f"The render interval ({self.cfg.sim.render_interval}) is smaller than the decimation " + f"({self.cfg.decimation}). Multiple render calls will happen for each environment step." + "If this is not intended, set the render interval to be equal to the decimation." + ) + omni.log.warn(msg) + + # generate scene + with Timer("[INFO]: Time taken for scene creation", "scene_creation"): + # set the stage context for scene creation steps which use the stage + with use_stage(self.sim.get_initial_stage()): + self.scene = InteractiveScene(self.cfg.scene) + self._setup_scene() + attach_stage_to_usd_context() + print("[INFO]: Scene manager: ", self.scene) + + # set up camera viewport controller + # viewport is not available in other rendering modes so the function will throw a warning + # FIXME: This needs to be fixed in the future when we unify the UI functionalities even for + # non-rendering modes. + if self.sim.render_mode >= self.sim.RenderMode.PARTIAL_RENDERING: + self.viewport_camera_controller = ViewportCameraController(self, self.cfg.viewer) + else: + self.viewport_camera_controller = None + + # create event manager + # note: this is needed here (rather than after simulation play) to allow USD-related randomization events + # that must happen before the simulation starts. Example: randomizing mesh scale + if self.cfg.events: + self.event_manager = EventManager(self.cfg.events, self) + + # apply USD-related randomization events + if "prestartup" in self.event_manager.available_modes: + self.event_manager.apply(mode="prestartup") + + # play the simulator to activate physics handles + # note: this activates the physics simulation view that exposes TensorAPIs + # note: when started in extension mode, first call sim.reset_async() and then initialize the managers + if builtins.ISAAC_LAUNCHED_FROM_TERMINAL is False: + print("[INFO]: Starting the simulation. This may take a few seconds. Please wait...") + with Timer("[INFO]: Time taken for simulation start", "simulation_start"): + # since the reset can trigger callbacks which use the stage, + # we need to set the stage context here + with use_stage(self.sim.get_initial_stage()): + self.sim.reset() + # update scene to pre populate data buffers for assets and sensors. + # this is needed for the observation manager to get valid tensors for initialization. + # this shouldn't cause an issue since later on, users do a reset over all the environments so the lazy buffers would be reset. + self.scene.update(dt=self.physics_dt) + + # check if debug visualization is has been implemented by the environment + source_code = inspect.getsource(self._set_debug_vis_impl) + self.has_debug_vis_implementation = "NotImplementedError" not in source_code + self._debug_vis_handle = None + + # extend UI elements + # we need to do this here after all the managers are initialized + # this is because they dictate the sensors and commands right now + if self.sim.has_gui() and self.cfg.ui_window_class_type is not None: + self._window = self.cfg.ui_window_class_type(self, window_name="IsaacLab") + else: + # if no window, then we don't need to store the window + self._window = None + + # allocate dictionary to store metrics + self.extras = {} + + # initialize data and constants + # -- counter for simulation steps + self._sim_step_counter = 0 + # -- counter for curriculum + self.common_step_counter = 0 + # -- init buffers + self.episode_length_buf = wp.zeros(self.num_envs, dtype=wp.int32, device=self.device) + # self.episode_length_buf = torch.zeros(self.num_envs, device=self.device, dtype=torch.long) + self.reset_terminated = wp.zeros(self.num_envs, dtype=wp.bool, device=self.device) + self.reset_time_outs = wp.zeros(self.num_envs, dtype=wp.bool, device=self.device) + self.reset_buf = wp.zeros(self.num_envs, dtype=wp.bool, device=self.device) + self._ALL_ENV_MASK = wp.ones(self.num_envs, dtype=wp.bool, device=self.device) + + # Expected bindings: + self.torch_obs_buf: torch.Tensor = None + self.torch_reward_buf: torch.Tensor = None + self.torch_reset_terminated: torch.Tensor = None + self.torch_reset_time_outs: torch.Tensor = None + self.torch_episode_length_buf: torch.Tensor = None + + # Graph captured by torch + self.graph_end_step = None + self.graph_action_step = None + + # setup the action and observation spaces for Gym + self._configure_gym_env_spaces() + + # setup noise cfg for adding action and observation noise + if self.cfg.action_noise_model: + self._action_noise_model: NoiseModel = self.cfg.action_noise_model.class_type( + self.cfg.action_noise_model, num_envs=self.num_envs, device=self.device + ) + if self.cfg.observation_noise_model: + self._observation_noise_model: NoiseModel = self.cfg.observation_noise_model.class_type( + self.cfg.observation_noise_model, num_envs=self.num_envs, device=self.device + ) + + # perform events at the start of the simulation + if self.cfg.events: + # we print it here to make the logging consistent + print("[INFO] Event Manager: ", self.event_manager) + + if "startup" in self.event_manager.available_modes: + self.event_manager.apply(mode="startup") + + # set the framerate of the gym video recorder wrapper so that the playback speed of the produced + # video matches the simulation + self.metadata["render_fps"] = 1 / self.step_dt + + # print the environment information + print("[INFO]: Completed setting up the environment...") + + def __del__(self): + """Cleanup for the environment.""" + self.close() + + """ + Properties. + """ + + @property + def num_envs(self) -> int: + """The number of instances of the environment that are running.""" + return self.scene.num_envs + + @property + def physics_dt(self) -> float: + """The physics time-step (in s). + + This is the lowest time-decimation at which the simulation is happening. + """ + return self.cfg.sim.dt + + @property + def step_dt(self) -> float: + """The environment stepping time-step (in s). + + This is the time-step at which the environment steps forward. + """ + return self.cfg.sim.dt * self.cfg.decimation + + @property + def device(self): + """The device on which the environment is running.""" + return self.sim.device + + @property + def max_episode_length_s(self) -> float: + """Maximum episode length in seconds.""" + return self.cfg.episode_length_s + + @property + def max_episode_length(self): + """The maximum episode length in steps adjusted from s.""" + return math.ceil(self.max_episode_length_s / (self.cfg.sim.dt * self.cfg.decimation)) + + """ + Operations. + """ + + def reset(self, seed: int | None = None, options: dict[str, Any] | None = None) -> tuple[VecEnvObs, dict]: + """Resets all the environments and returns observations. + + This function calls the :meth:`_reset_idx` function to reset all the environments. + However, certain operations, such as procedural terrain generation, that happened during initialization + are not repeated. + + Args: + seed: The seed to use for randomization. Defaults to None, in which case the seed is not set. + options: Additional information to specify how the environment is reset. Defaults to None. + + Note: + This argument is used for compatibility with Gymnasium environment definition. + + Returns: + A tuple containing the observations and extras. + """ + # set the seed + if seed is not None: + self.seed(seed) + + # reset state of scene + self._reset_idx(self._ALL_ENV_MASK) + + # update articulation kinematics + self.scene.write_data_to_sim() + + # if sensors are added to the scene, make sure we render to reflect changes in reset + if self.sim.has_rtx_sensors() and self.cfg.rerender_on_reset: + self.sim.render() + + if self.cfg.wait_for_textures and self.sim.has_rtx_sensors(): + while SimulationManager.assets_loading(): + self.sim.render() + + # return observations + self._get_observations() + return self.torch_obs_buf.clone(), self.extras + + @Timer(name="env_step", msg="Step took:", enable=True, format="us") + def step(self, action: torch.Tensor) -> VecEnvStepReturn: + """Execute one time-step of the environment's dynamics. + + The environment steps forward at a fixed time-step, while the physics simulation is decimated at a + lower time-step. This is to ensure that the simulation is stable. These two time-steps can be configured + independently using the :attr:`DirectRLEnvCfg.decimation` (number of simulation steps per environment step) + and the :attr:`DirectRLEnvCfg.sim.physics_dt` (physics time-step). Based on these parameters, the environment + time-step is computed as the product of the two. + + This function performs the following steps: + + 1. Pre-process the actions before stepping through the physics. + 2. Apply the actions to the simulator and step through the physics in a decimated manner. + 3. Compute the reward and done signals. + 4. Reset environments that have terminated or reached the maximum episode length. + 5. Apply interval events if they are enabled. + 6. Compute observations. + + Args: + action: The actions to apply on the environment. Shape is (num_envs, action_dim). + + Returns: + A tuple containing the observations, rewards, resets (terminated and truncated) and extras. + """ + + action = action.to(self.device) + # add action noise + if self.cfg.action_noise_model: + action = self._action_noise_model(action) + + # process actions, #TODO pass the torch tensor directly. + self._pre_physics_step(wp.from_torch(action)) # Creates a tensor and throws it away. --> Not graphable unless the training loop is using the same pointer. + + # check if we need to do rendering within the physics loop + # note: checked here once to avoid multiple checks within the loop + is_rendering = self.sim.has_gui() or self.sim.has_rtx_sensors() + + # perform physics stepping + for _ in range(self.cfg.decimation): + self._sim_step_counter += 1 + # set actions into buffers + # simulate + with Timer(name="apply_action", msg="Action processing step took:", enable=True, format="us"): + if self.graph_action_step is None: + with wp.ScopedCapture() as capture: + self.step_warp_action() + self.graph_action_step = capture.graph + else: + wp.capture_launch(self.graph_action_step) + + with Timer(name="simulate", msg="Newton simulation step took:", enable=True, format="us"): + self.sim.step(render=False) + # render between steps only if the GUI or an RTX sensor needs it + # note: we assume the render interval to be the shortest accepted rendering interval. + # If a camera needs rendering at a faster frequency, this will lead to unexpected behavior. + if self._sim_step_counter % self.cfg.sim.render_interval == 0 and is_rendering: + self.sim.render() + # update buffers at sim dt + self.scene.update(dt=self.physics_dt) + + self.common_step_counter += 1 # total step (common for all envs) + with Timer(name="post_processing", msg="Post-Processing step took:", enable=True, format="us"): + if self.graph_end_step is None: + with wp.ScopedCapture() as capture: + self.step_warp_end() + self.graph_end_step = capture.graph + else: + wp.capture_launch(self.graph_end_step) + + # return observations, rewards, resets and extras + return ( + {"policy": self.torch_obs_buf.clone()}, + self.torch_reward_buf, + self.torch_reset_terminated, + self.torch_reset_time_outs, + self.extras, + ) + + def step_warp_action(self) -> None: + self._apply_action() + # set actions into simulator + self.scene.write_data_to_sim() + + def step_warp_end(self) -> None: + wp.launch( + add_to_env, + dim=self.num_envs, + inputs=[ + self.episode_length_buf, + 1, + ], + ) + self._get_dones() + self._get_rewards() + + # -- reset envs that terminated/timed-out and log the episode information + self._reset_idx(mask=self.reset_buf) + # update articulation kinematics + self.scene.write_data_to_sim() + # if sensors are added to the scene, make sure we render to reflect changes in reset + # if self.sim.has_rtx_sensors() and self.cfg.rerender_on_reset: + # self.sim.render() + + # TODO We could split it out. + # post-step: step interval event + # if self.cfg.events: + # if "interval" in self.event_manager.available_modes: + # self.event_manager.apply(mode="interval", dt=self.step_dt) + + # update observations + self._get_observations() + + # add observation noise + # note: we apply no noise to the state space (since it is used for critic networks) + # if self.cfg.observation_noise_model: + # self.obs_buf["policy"] = self._observation_noise_model(self.obs_buf["policy"]) + + @staticmethod + def seed(seed: int = -1) -> int: + """Set the seed for the environment. + + Args: + seed: The seed for random generator. Defaults to -1. + + Returns: + The seed used for random generator. + """ + # set seed for replicator + try: + import omni.replicator.core as rep + + rep.set_global_seed(seed) + except ModuleNotFoundError: + pass + # set seed for torch and other libraries + return torch_utils.set_seed(seed) + + def render(self, recompute: bool = False) -> np.ndarray | None: + """Run rendering without stepping through the physics. + + By convention, if mode is: + + - **human**: Render to the current display and return nothing. Usually for human consumption. + - **rgb_array**: Return an numpy.ndarray with shape (x, y, 3), representing RGB values for an + x-by-y pixel image, suitable for turning into a video. + + Args: + recompute: Whether to force a render even if the simulator has already rendered the scene. + Defaults to False. + + Returns: + The rendered image as a numpy array if mode is "rgb_array". Otherwise, returns None. + + Raises: + RuntimeError: If mode is set to "rgb_data" and simulation render mode does not support it. + In this case, the simulation render mode must be set to ``RenderMode.PARTIAL_RENDERING`` + or ``RenderMode.FULL_RENDERING``. + NotImplementedError: If an unsupported rendering mode is specified. + """ + # run a rendering step of the simulator + # if we have rtx sensors, we do not need to render again sim + if not self.sim.has_rtx_sensors() and not recompute: + self.sim.render() + # decide the rendering mode + if self.render_mode == "human" or self.render_mode is None: + return None + elif self.render_mode == "rgb_array": + # check that if any render could have happened + if self.sim.render_mode.value < self.sim.RenderMode.PARTIAL_RENDERING.value: + raise RuntimeError( + f"Cannot render '{self.render_mode}' when the simulation render mode is" + f" '{self.sim.render_mode.name}'. Please set the simulation render mode to:" + f"'{self.sim.RenderMode.PARTIAL_RENDERING.name}' or '{self.sim.RenderMode.FULL_RENDERING.name}'." + " If running headless, make sure --enable_cameras is set." + ) + # create the annotator if it does not exist + if not hasattr(self, "_rgb_annotator"): + import omni.replicator.core as rep + + # create render product + self._render_product = rep.create.render_product( + self.cfg.viewer.cam_prim_path, self.cfg.viewer.resolution + ) + # create rgb annotator -- used to read data from the render product + self._rgb_annotator = rep.AnnotatorRegistry.get_annotator("rgb", device="cpu") + self._rgb_annotator.attach([self._render_product]) + # obtain the rgb data + rgb_data = self._rgb_annotator.get_data() + # convert to numpy array + rgb_data = np.frombuffer(rgb_data, dtype=np.uint8).reshape(*rgb_data.shape) + # return the rgb data + # note: initially the renerer is warming up and returns empty data + if rgb_data.size == 0: + return np.zeros((self.cfg.viewer.resolution[1], self.cfg.viewer.resolution[0], 3), dtype=np.uint8) + else: + return rgb_data[:, :, :3] + else: + raise NotImplementedError( + f"Render mode '{self.render_mode}' is not supported. Please use: {self.metadata['render_modes']}." + ) + + def close(self): + """Cleanup for the environment.""" + if not self._is_closed: + # close entities related to the environment + # note: this is order-sensitive to avoid any dangling references + if self.cfg.events: + del self.event_manager + del self.scene + if self.viewport_camera_controller is not None: + del self.viewport_camera_controller + + # clear callbacks and instance + if float(".".join(get_version()[2])) >= 5: + if self.cfg.sim.create_stage_in_memory: + # detach physx stage + omni.physx.get_physx_simulation_interface().detach_stage() + self.sim.stop() + self.sim.clear() + + self.sim.clear_all_callbacks() + self.sim.clear_instance() + + # destroy the window + if self._window is not None: + self._window = None + # update closing status + self._is_closed = True + + """ + Operations - Debug Visualization. + """ + + def set_debug_vis(self, debug_vis: bool) -> bool: + """Toggles the environment debug visualization. + + Args: + debug_vis: Whether to visualize the environment debug visualization. + + Returns: + Whether the debug visualization was successfully set. False if the environment + does not support debug visualization. + """ + # check if debug visualization is supported + if not self.has_debug_vis_implementation: + return False + # toggle debug visualization objects + self._set_debug_vis_impl(debug_vis) + # toggle debug visualization handles + if debug_vis: + # create a subscriber for the post update event if it doesn't exist + if self._debug_vis_handle is None: + app_interface = omni.kit.app.get_app_interface() + self._debug_vis_handle = app_interface.get_post_update_event_stream().create_subscription_to_pop( + lambda event, obj=weakref.proxy(self): obj._debug_vis_callback(event) + ) + else: + # remove the subscriber if it exists + if self._debug_vis_handle is not None: + self._debug_vis_handle.unsubscribe() + self._debug_vis_handle = None + # return success + return True + + """ + Helper functions. + """ + + def _configure_gym_env_spaces(self): + """Configure the action and observation spaces for the Gym environment.""" + # show deprecation message and overwrite configuration + if self.cfg.num_actions is not None: + omni.log.warn("DirectRLEnvCfg.num_actions is deprecated. Use DirectRLEnvCfg.action_space instead.") + if isinstance(self.cfg.action_space, type(MISSING)): + self.cfg.action_space = self.cfg.num_actions + if self.cfg.num_observations is not None: + omni.log.warn( + "DirectRLEnvCfg.num_observations is deprecated. Use DirectRLEnvCfg.observation_space instead." + ) + if isinstance(self.cfg.observation_space, type(MISSING)): + self.cfg.observation_space = self.cfg.num_observations + if self.cfg.num_states is not None: + omni.log.warn("DirectRLEnvCfg.num_states is deprecated. Use DirectRLEnvCfg.state_space instead.") + if isinstance(self.cfg.state_space, type(MISSING)): + self.cfg.state_space = self.cfg.num_states + + # set up spaces + self.single_observation_space = gym.spaces.Dict() + self.single_observation_space["policy"] = spec_to_gym_space(self.cfg.observation_space) + self.single_action_space = spec_to_gym_space(self.cfg.action_space) + + # batch the spaces for vectorized environments + self.observation_space = gym.vector.utils.batch_space(self.single_observation_space["policy"], self.num_envs) + self.action_space = gym.vector.utils.batch_space(self.single_action_space, self.num_envs) + + # optional state space for asymmetric actor-critic architectures + self.state_space = None + if self.cfg.state_space: + self.single_observation_space["critic"] = spec_to_gym_space(self.cfg.state_space) + self.state_space = gym.vector.utils.batch_space(self.single_observation_space["critic"], self.num_envs) + + # instantiate actions (needed for tasks for which the observations computation is dependent on the actions) + self.actions = sample_space(self.single_action_space, self.sim.device, batch_size=self.num_envs, fill_value=0) + + def _reset_idx(self, mask: wp.array | torch.Tensor | None = None): + """Reset environments based on specified indices. + + Args: + env_ids: List of environment ids which must be reset + """ + + self.scene.reset(env_ids=None, mask=mask) + + # apply events such as randomization for environments that need a reset + # if self.cfg.events: + # if "reset" in self.event_manager.available_modes: + # env_step_count = self._sim_step_counter // self.cfg.decimation + # self.event_manager.apply(mode="reset", env_ids=env_ids, global_env_step_count=env_step_count) + + # reset noise models + # if self.cfg.action_noise_model: + # self._action_noise_model.reset(env_ids) + # if self.cfg.observation_noise_model: + # self._observation_noise_model.reset(env_ids) + + # reset the episode length buffer + wp.launch( + zero_mask_int32, + dim=self.num_envs, + inputs=[ + mask, + self.episode_length_buf, + ], + ) + + """ + Implementation-specific functions. + """ + + def _setup_scene(self): + """Setup the scene for the environment. + + This function is responsible for creating the scene objects and setting up the scene for the environment. + The scene creation can happen through :class:`isaaclab.scene.InteractiveSceneCfg` or through + directly creating the scene objects and registering them with the scene manager. + + We leave the implementation of this function to the derived classes. If the environment does not require + any explicit scene setup, the function can be left empty. + """ + pass + + @abstractmethod + def _pre_physics_step(self, actions: torch.Tensor): + """Pre-process actions before stepping through the physics. + + This function is responsible for pre-processing the actions before stepping through the physics. + It is called before the physics stepping (which is decimated). + + Args: + actions: The actions to apply on the environment. Shape is (num_envs, action_dim). + """ + raise NotImplementedError(f"Please implement the '_pre_physics_step' method for {self.__class__.__name__}.") + + @abstractmethod + def _apply_action(self): + """Apply actions to the simulator. + + This function is responsible for applying the actions to the simulator. It is called at each + physics time-step. + """ + raise NotImplementedError(f"Please implement the '_apply_action' method for {self.__class__.__name__}.") + + @abstractmethod + def _get_observations(self) -> VecEnvObs: + """Compute and return the observations for the environment. + + Returns: + The observations for the environment. + """ + raise NotImplementedError(f"Please implement the '_get_observations' method for {self.__class__.__name__}.") + + def _get_states(self) -> VecEnvObs | None: + """Compute and return the states for the environment. + + The state-space is used for asymmetric actor-critic architectures. It is configured + using the :attr:`DirectRLEnvCfg.state_space` parameter. + + Returns: + The states for the environment. If the environment does not have a state-space, the function + returns a None. + """ + return None # noqa: R501 + + @abstractmethod + def _get_rewards(self) -> torch.Tensor: + """Compute and return the rewards for the environment. + + Returns: + The rewards for the environment. Shape is (num_envs,). + """ + raise NotImplementedError(f"Please implement the '_get_rewards' method for {self.__class__.__name__}.") + + @abstractmethod + def _get_dones(self) -> tuple[torch.Tensor, torch.Tensor]: + """Compute and return the done flags for the environment. + + Returns: + A tuple containing the done flags for termination and time-out. + Shape of individual tensors is (num_envs,). + """ + raise NotImplementedError(f"Please implement the '_get_dones' method for {self.__class__.__name__}.") + + def _set_debug_vis_impl(self, debug_vis: bool): + """Set debug visualization into visualization objects. + + This function is responsible for creating the visualization objects if they don't exist + and input ``debug_vis`` is True. If the visualization objects exist, the function should + set their visibility into the stage. + """ + raise NotImplementedError(f"Debug visualization is not implemented for {self.__class__.__name__}.") diff --git a/source/isaaclab/isaaclab/envs/mdp/actions/actions_cfg.py b/source/isaaclab/isaaclab/envs/mdp/actions/actions_cfg.py index 31419e55517..8d21eabf425 100644 --- a/source/isaaclab/isaaclab/envs/mdp/actions/actions_cfg.py +++ b/source/isaaclab/isaaclab/envs/mdp/actions/actions_cfg.py @@ -213,11 +213,138 @@ class NonHolonomicActionCfg(ActionTermCfg): ## @configclass class DifferentialInverseKinematicsActionCfg(ActionTermCfg): - def __post_init__(self): - raise NotImplementedError("Not implemented") + """Configuration for inverse differential kinematics action term. + + See :class:`DifferentialInverseKinematicsAction` for more details. + """ + + pass + # @configclass + # class OffsetCfg: + # """The offset pose from parent frame to child frame. + + # On many robots, end-effector frames are fictitious frames that do not have a corresponding + # rigid body. In such cases, it is easier to define this transform w.r.t. their parent rigid body. + # For instance, for the Franka Emika arm, the end-effector is defined at an offset to the the + # "panda_hand" frame. + # """ + + # pos: tuple[float, float, float] = (0.0, 0.0, 0.0) + # """Translation w.r.t. the parent frame. Defaults to (0.0, 0.0, 0.0).""" + # rot: tuple[float, float, float, float] = (1.0, 0.0, 0.0, 0.0) + # """Quaternion rotation ``(w, x, y, z)`` w.r.t. the parent frame. Defaults to (1.0, 0.0, 0.0, 0.0).""" + + # class_type: type[ActionTerm] = task_space_actions.DifferentialInverseKinematicsAction + + # joint_names: list[str] = MISSING + # """List of joint names or regex expressions that the action will be mapped to.""" + # body_name: str = MISSING + # """Name of the body or frame for which IK is performed.""" + # body_offset: OffsetCfg | None = None + # """Offset of target frame w.r.t. to the body frame. Defaults to None, in which case no offset is applied.""" + # scale: float | tuple[float, ...] = 1.0 + # """Scale factor for the action. Defaults to 1.0.""" + # controller: DifferentialIKControllerCfg = MISSING + # """The configuration for the differential IK controller.""" + + +@configclass +class DifferentialInverseKinematicsNewtonActionCfg(ActionTermCfg): + """Configuration for inverse differential kinematics action term. + + See :class:`DifferentialInverseKinematicsAction` for more details. + """ + + pass + + # @configclass + # class OffsetCfg: + # """The offset pose from parent frame to child frame. + + # On many robots, end-effector frames are fictitious frames that do not have a corresponding + # rigid body. In such cases, it is easier to define this transform w.r.t. their parent rigid body. + # For instance, for the Franka Emika arm, the end-effector is defined at an offset to the the + # "panda_hand" frame. + # """ + + # pos: tuple[float, float, float] = (0.0, 0.0, 0.0) + # """Translation w.r.t. the parent frame. Defaults to (0.0, 0.0, 0.0).""" + # rot: tuple[float, float, float, float] = (1.0, 0.0, 0.0, 0.0) + # """Quaternion rotation ``(w, x, y, z)`` w.r.t. the parent frame. Defaults to (1.0, 0.0, 0.0, 0.0).""" + + # class_type: type[ActionTerm] = task_space_actions.DifferentialInverseKinematicsNewtonAction + + # joint_names: list[str] = MISSING + # """List of joint names or regex expressions that the action will be mapped to.""" + # body_name: str = MISSING + # """Name of the body or frame for which IK is performed.""" + # body_offset: OffsetCfg | None = None + # """Offset of target frame w.r.t. to the body frame. Defaults to None, in which case no offset is applied.""" + # scale: float | tuple[float, ...] = 1.0 + # """Scale factor for the action. Defaults to 1.0.""" + # controller: DifferentialIKControllerCfg = MISSING + # """The configuration for the differential IK controller.""" @configclass class OperationalSpaceControllerActionCfg(ActionTermCfg): - def __post_init__(self): - raise NotImplementedError("Not implemented") + """Configuration for operational space controller action term. + + See :class:`OperationalSpaceControllerAction` for more details. + """ + + pass + + # @configclass + # class OffsetCfg: + # """The offset pose from parent frame to child frame. + + # On many robots, end-effector frames are fictitious frames that do not have a corresponding + # rigid body. In such cases, it is easier to define this transform w.r.t. their parent rigid body. + # For instance, for the Franka Emika arm, the end-effector is defined at an offset to the the + # "panda_hand" frame. + # """ + + # pos: tuple[float, float, float] = (0.0, 0.0, 0.0) + # """Translation w.r.t. the parent frame. Defaults to (0.0, 0.0, 0.0).""" + # rot: tuple[float, float, float, float] = (1.0, 0.0, 0.0, 0.0) + # """Quaternion rotation ``(w, x, y, z)`` w.r.t. the parent frame. Defaults to (1.0, 0.0, 0.0, 0.0).""" + + # class_type: type[ActionTerm] = task_space_actions.OperationalSpaceControllerAction + + # joint_names: list[str] = MISSING + # """List of joint names or regex expressions that the action will be mapped to.""" + + # body_name: str = MISSING + # """Name of the body or frame for which motion/force control is performed.""" + + # body_offset: OffsetCfg | None = None + # """Offset of target frame w.r.t. to the body frame. Defaults to None, in which case no offset is applied.""" + + # task_frame_rel_path: str = None + # """The path of a ``RigidObject``, relative to the sub-environment, representing task frame. Defaults to None.""" + + # controller_cfg: OperationalSpaceControllerCfg = MISSING + # """The configuration for the operational space controller.""" + + # position_scale: float = 1.0 + # """Scale factor for the position targets. Defaults to 1.0.""" + + # orientation_scale: float = 1.0 + # """Scale factor for the orientation (quad for ``pose_abs`` or axis-angle for ``pose_rel``). Defaults to 1.0.""" + + # wrench_scale: float = 1.0 + # """Scale factor for the wrench targets. Defaults to 1.0.""" + + # stiffness_scale: float = 1.0 + # """Scale factor for the stiffness commands. Defaults to 1.0.""" + + # damping_ratio_scale: float = 1.0 + # """Scale factor for the damping ratio commands. Defaults to 1.0.""" + + # nullspace_joint_pos_target: str = "none" + # """The joint targets for the null-space control: ``"none"``, ``"zero"``, ``"default"``, ``"center"``. + + # Note: Functional only when ``nullspace_control`` is set to ``"position"`` within the + # ``OperationalSpaceControllerCfg``. + # """ diff --git a/source/isaaclab/isaaclab/newton/__init__.py b/source/isaaclab/isaaclab/newton/__init__.py new file mode 100644 index 00000000000..a8bccef2c57 --- /dev/null +++ b/source/isaaclab/isaaclab/newton/__init__.py @@ -0,0 +1,5 @@ +from enum import IntEnum + +class Frontend(IntEnum): + WARP = 0 + TORCH = 1 \ No newline at end of file diff --git a/source/isaaclab/isaaclab/newton/actuators/__init__.py b/source/isaaclab/isaaclab/newton/actuators/__init__.py new file mode 100644 index 00000000000..127c784cfa9 --- /dev/null +++ b/source/isaaclab/isaaclab/newton/actuators/__init__.py @@ -0,0 +1,28 @@ +# Copyright (c) 2022-2025, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +"""Sub-package for different actuator models. + +Actuator models are used to model the behavior of the actuators in an articulation. These +are usually meant to be used in simulation to model different actuator dynamics and delays. + +There are two main categories of actuator models that are supported: + +- **Implicit**: Motor model with ideal PD from the physics engine. This is similar to having a continuous time + PD controller. The motor model is implicit in the sense that the motor model is not explicitly defined by the user. +- **Explicit**: Motor models based on physical drive models. + + - **Physics-based**: Derives the motor models based on first-principles. + - **Neural Network-based**: Learned motor models from actuator data. + +Every actuator model inherits from the :class:`isaaclab.actuators.ActuatorBase` class, +which defines the common interface for all actuator models. The actuator models are handled +and called by the :class:`isaaclab.assets.Articulation` class. +""" + +from .actuator_base import ActuatorBase +# from .actuator_net import ActuatorNetLSTM, ActuatorNetMLP +# from .actuator_pd import DCMotor, DelayedPDActuator, IdealPDActuator, ImplicitActuator, RemotizedPDActuator +from .actuator_pd import DCMotor, IdealPDActuator, ImplicitActuator diff --git a/source/isaaclab/isaaclab/newton/actuators/actuator_base.py b/source/isaaclab/isaaclab/newton/actuators/actuator_base.py new file mode 100644 index 00000000000..8e65fa202bb --- /dev/null +++ b/source/isaaclab/isaaclab/newton/actuators/actuator_base.py @@ -0,0 +1,290 @@ +# Copyright (c) 2022-2025, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +from __future__ import annotations + +from abc import ABC, abstractmethod +from typing import TYPE_CHECKING, ClassVar + +import warp as wp + +import isaaclab.utils.string as string_utils +from isaaclab.newton.actuators.kernels import clip_efforts_with_limits +from isaaclab.newton.assets.core.kernels import ( + populate_empty_array, + update_joint_array_with_value, + update_joint_array_with_value_array, + update_joint_array_with_value_int, +) +from isaaclab.actuators.actuator_cfg import ActuatorBaseCfg + +if TYPE_CHECKING: + from isaaclab.newton.assets.core.joint_properties.joint_data import JointData + + +class ActuatorBase(ABC): + """Base class for actuator models over a collection of actuated joints in an articulation. + + Actuator models augment the simulated articulation joints with an external drive dynamics model. + The model is used to convert the user-provided joint commands (positions, velocities and efforts) + into the desired joint positions, velocities and efforts that are applied to the simulated articulation. + + The base class provides the interface for the actuator models. It is responsible for parsing the + actuator parameters from the configuration and storing them as buffers. It also provides the + interface for resetting the actuator state and computing the desired joint commands for the simulation. + + For each actuator model, a corresponding configuration class is provided. The configuration class + is used to parse the actuator parameters from the configuration. It also specifies the joint names + for which the actuator model is applied. These names can be specified as regular expressions, which + are matched against the joint names in the articulation. + + To see how the class is used, check the :class:`isaaclab.assets.Articulation` class. + """ + + __backend_name__: str = "newton" + """The name of the backend for the actuator.""" + + is_implicit_model: ClassVar[bool] = False + """Flag indicating if the actuator is an implicit or explicit actuator model. + + If a class inherits from :class:`ImplicitActuator`, then this flag should be set to :obj:`True`. + """ + + joint_data: JointData + """The data of the articulation.""" + + _DEFAULT_MAX_EFFORT_SIM: ClassVar[float] = 1.0e9 + """The default maximum effort for the actuator joints in the simulation. Defaults to 1.0e9. + + If the :attr:`ActuatorBaseCfg.effort_limit_sim` is not specified and the actuator is an explicit + actuator, then this value is used. + """ + + def __init__( + self, + cfg: ActuatorBaseCfg, + joint_names: list[str], + joint_mask: wp.array, + env_mask: wp.array, + joint_data: JointData, + device: str, + ): + """Initialize the actuator. + + The actuator parameters are parsed from the configuration and stored as buffers. If the parameters + are not specified in the configuration, then their values provided in the constructor are used. + + .. note:: + The values in the constructor are typically obtained through the USD schemas corresponding + to the joints in the actuator model. + + Args: + cfg: The configuration of the actuator model. + joint_names: The joint names in the articulation. + joint_mask: The mask of joints to use. + env_mask: The mask of environments to use. + articulation_data: The data for the articulation. + device: Device used for processing. + """ + # save parameters + self.cfg = cfg + self._device = device + self._joint_names = joint_names + self._joint_mask = joint_mask + self._env_mask = env_mask + # Get the number of environments and joints from the articulation data + self._num_envs = env_mask.shape[0] + self._num_joints = joint_mask.shape[0] + # Get the data from the articulation + self.joint_data = joint_data + + # For explicit models, we do not want to enforce the effort limit through the solver + # (unless it is explicitly set) + if not self.is_implicit_model and self.cfg.effort_limit_sim is None: + self.cfg.effort_limit_sim = self._DEFAULT_MAX_EFFORT_SIM + + # resolve usd, actuator configuration values + # case 1: if usd_value == actuator_cfg_value: all good, + # case 2: if usd_value != actuator_cfg_value: we use actuator_cfg_value + # case 3: if actuator_cfg_value is None: we use usd_value + + to_check = [ + ("velocity_limit_sim", self.joint_data.joint_vel_limits), + ("effort_limit_sim", self.joint_data.joint_effort_limits), + ("stiffness", self.joint_data.joint_stiffness), + ("damping", self.joint_data.joint_damping), + ("armature", self.joint_data.joint_armature), + ("friction", self.joint_data.joint_friction_coeff), + ("dynamic_friction", self.joint_data.joint_dynamic_friction), + ("viscous_friction", self.joint_data.joint_viscous_friction), + ("control_mode", self.joint_data.joint_control_mode), + ] + for param_name, newton_val in to_check: + cfg_val = getattr(self.cfg, param_name) + self._parse_joint_parameter(cfg_val, newton_val) + + def __str__(self) -> str: + """Returns: A string representation of the actuator group.""" + # resolve model type (implicit or explicit) + model_type = "implicit" if self.is_implicit_model else "explicit" + + return ( + f" object:\n" + f"\tModel type : {model_type}\n" + f"\tNumber of joints : {self.num_joints}\n" + f"\tJoint names expression: {self.cfg.joint_names_expr}\n" + f"\tJoint names : {self.joint_names}\n" + f"\tJoint mask : {self.joint_mask}\n" + ) + + """ + Properties. + """ + + @property + def num_joints(self) -> int: + """Number of actuators in the group.""" + return len(self._joint_names) + + @property + def joint_names(self) -> list[str]: + """Articulation's joint names that are part of the group.""" + return self._joint_names + + @property + def joint_mask(self) -> wp.array: + """Articulation's masked indices that denote which joints are part of the group.""" + return self._joint_mask + + """ + Operations. + """ + + @abstractmethod + def reset(self, env_mask: wp.array): + """Reset the internals within the group. + + Args: + env_mask: Mask of environments to reset. + """ + raise NotImplementedError + + @abstractmethod + def compute(self): + """Process the actuator group actions and compute the articulation actions. + + It computes the articulation actions based on the actuator model type. To do so, it reads + the following quantities from the articulation data: + - sim_bind_joint_pos, the current joint positions + - sim_bind_joint_vel, the current joint velocities + - joint_control_mode, the current joint control mode + - joint_stiffness, the current joint stiffness + - joint_damping, the current joint damping + + With these, it updates the following quantities: + + """ + raise NotImplementedError + + """ + Helper functions. + """ + + def _parse_joint_parameter(self, cfg_value: float | dict[str, float] | None, original_value: wp.array | None): + """Parse the joint parameter from the configuration. + + Args: + cfg_value: The parameter value from the configuration. If None, then use the default value. + default_value: The default value to use if the parameter is None. If it is also None, + then an error is raised. + + Returns: + The parsed parameter value. + + Raises: + TypeError: If the parameter value is not of the expected type. + TypeError: If the default value is not of the expected type. + ValueError: If the parameter value is None and no default value is provided. + ValueError: If the default value tensor is the wrong shape. + """ + # parse the parameter + if cfg_value is not None: + if isinstance(cfg_value, float): + # if float, then use the same value for all joints + wp.launch( + update_joint_array_with_value, + dim=(self._num_envs, self._num_joints), + inputs=[ + float(cfg_value), + original_value, + self._env_mask, + self._joint_mask, + ], + ) + elif isinstance(cfg_value, int): + # if int, then use the same value for all joints + wp.launch( + update_joint_array_with_value_int, + dim=(self._num_envs, self._num_joints), + inputs=[ + cfg_value, + original_value, + self._env_mask, + self._joint_mask, + ], + ) + elif isinstance(cfg_value, dict): + # if dict, then parse the regular expression + indices, _, values = string_utils.resolve_matching_names_values(cfg_value, self.joint_names) + tmp_param = wp.zeros((self._num_joints,), dtype=wp.float32, device=self._device) + wp.launch( + populate_empty_array, + dim=(self._num_joints,), + inputs=[ + wp.array(values, dtype=wp.float32, device=self._device), + tmp_param, + wp.array(indices, dtype=wp.int32, device=self._device), + ], + ) + wp.launch( + update_joint_array_with_value_array, + dim=(self._num_envs, self._num_joints), + inputs=[ + tmp_param, + original_value, + self._env_mask, + self._joint_mask, + ], + ) + else: + raise TypeError( + f"Invalid type for parameter value: {type(cfg_value)} for " + + f"actuator on joints {self.joint_names}. Expected float or dict." + ) + elif original_value is None: + raise ValueError("The parameter value is None and no newton value is provided.") + + def _clip_effort(self, effort: wp.array, clipped_effort: wp.array) -> None: + """Clip the desired torques based on the motor limits. + + .. note:: The array is modified in place. + + Args: + desired_torques: The desired torques to clip. Expected shape is (num_envs, num_joints). + + Returns: + The clipped torques. + """ + wp.launch( + clip_efforts_with_limits, + dim=(self._num_envs, self._num_joints), + inputs=[ + self.joint_data.joint_effort_limits, + effort, + clipped_effort, + self._env_mask, + self.joint_mask, + ], + ) diff --git a/source/isaaclab/isaaclab/newton/actuators/actuator_net.py b/source/isaaclab/isaaclab/newton/actuators/actuator_net.py new file mode 100644 index 00000000000..8461a3b5e4e --- /dev/null +++ b/source/isaaclab/isaaclab/newton/actuators/actuator_net.py @@ -0,0 +1,191 @@ +# Copyright (c) 2022-2025, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +"""Neural network models for actuators. + +Currently, the following models are supported: + +* Multi-Layer Perceptron (MLP) +* Long Short-Term Memory (LSTM) + +""" + +from __future__ import annotations + +import torch +from collections.abc import Sequence +from typing import TYPE_CHECKING + +from isaaclab.utils.assets import read_file +from isaaclab.utils.types import ArticulationActions + +from .actuator_pd import DCMotor + +if TYPE_CHECKING: + from isaaclab.actuators.actuator_cfg import ActuatorNetLSTMCfg, ActuatorNetMLPCfg + + +class ActuatorNetLSTM(DCMotor): + """Actuator model based on recurrent neural network (LSTM). + + Unlike the MLP implementation :cite:t:`hwangbo2019learning`, this class implements + the learned model as a temporal neural network (LSTM) based on the work from + :cite:t:`rudin2022learning`. This removes the need of storing a history as the + hidden states of the recurrent network captures the history. + + Note: + Only the desired joint positions are used as inputs to the network. + """ + + cfg: ActuatorNetLSTMCfg + """The configuration of the actuator model.""" + + def __init__(self, cfg: ActuatorNetLSTMCfg, *args, **kwargs): + super().__init__(cfg, *args, **kwargs) + + assert self.cfg.control_mode == "position", "ActuatorNetLSTM only supports position control" + + # load the model from JIT file + file_bytes = read_file(self.cfg.network_file) + self.network = torch.jit.load(file_bytes, map_location=self._device).eval() + + # extract number of lstm layers and hidden dim from the shape of weights + num_layers = len(self.network.lstm.state_dict()) // 4 + hidden_dim = self.network.lstm.state_dict()["weight_hh_l0"].shape[1] + # create buffers for storing LSTM inputs + self.sea_input = torch.zeros(self._num_envs * self.num_joints, 1, 2, device=self._device) + self.sea_hidden_state = torch.zeros( + num_layers, self._num_envs * self.num_joints, hidden_dim, device=self._device + ) + self.sea_cell_state = torch.zeros(num_layers, self._num_envs * self.num_joints, hidden_dim, device=self._device) + # reshape via views (doesn't change the actual memory layout) + layer_shape_per_env = (num_layers, self._num_envs, self.num_joints, hidden_dim) + self.sea_hidden_state_per_env = self.sea_hidden_state.view(layer_shape_per_env) + self.sea_cell_state_per_env = self.sea_cell_state.view(layer_shape_per_env) + + """ + Operations. + """ + + def reset(self, env_ids: Sequence[int]): + # reset the hidden and cell states for the specified environments + with torch.no_grad(): + self.sea_hidden_state_per_env[:, env_ids] = 0.0 + self.sea_cell_state_per_env[:, env_ids] = 0.0 + + def compute( + self, control_action: ArticulationActions, joint_pos: torch.Tensor, joint_vel: torch.Tensor + ) -> ArticulationActions: + # compute network inputs + self.sea_input[:, 0, 0] = (control_action.joint_targets - joint_pos).flatten() + self.sea_input[:, 0, 1] = joint_vel.flatten() + # save current joint vel for dc-motor clipping + self._joint_vel[:] = joint_vel + + # run network inference + with torch.inference_mode(): + torques, (self.sea_hidden_state[:], self.sea_cell_state[:]) = self.network( + self.sea_input, (self.sea_hidden_state, self.sea_cell_state) + ) + self.computed_effort = torques.reshape(self._num_envs, self.num_joints) + + # clip the computed effort based on the motor limits + self.applied_effort = self._clip_effort(self.computed_effort) + + # return torques + control_action.joint_efforts = self.applied_effort + control_action.joint_targets = None + return control_action + + +class ActuatorNetMLP(DCMotor): + """Actuator model based on multi-layer perceptron and joint history. + + Many times the analytical model is not sufficient to capture the actuator dynamics, the + delay in the actuator response, or the non-linearities in the actuator. In these cases, + a neural network model can be used to approximate the actuator dynamics. This model is + trained using data collected from the physical actuator and maps the joint state and the + desired joint command to the produced torque by the actuator. + + This class implements the learned model as a neural network based on the work from + :cite:t:`hwangbo2019learning`. The class stores the history of the joint positions errors + and velocities which are used to provide input to the neural network. The model is loaded + as a TorchScript. + + Note: + Only the desired joint positions are used as inputs to the network. + + """ + + cfg: ActuatorNetMLPCfg + """The configuration of the actuator model.""" + + def __init__(self, cfg: ActuatorNetMLPCfg, *args, **kwargs): + super().__init__(cfg, *args, **kwargs) + + assert self.cfg.control_mode == "position", "ActuatorNetLSTM only supports position control" + + # load the model from JIT file + file_bytes = read_file(self.cfg.network_file) + self.network = torch.jit.load(file_bytes, map_location=self._device).eval() + + # create buffers for MLP history + history_length = max(self.cfg.input_idx) + 1 + self._joint_pos_error_history = torch.zeros( + self._num_envs, history_length, self.num_joints, device=self._device + ) + self._joint_vel_history = torch.zeros(self._num_envs, history_length, self.num_joints, device=self._device) + + """ + Operations. + """ + + def reset(self, env_ids: Sequence[int]): + # reset the history for the specified environments + self._joint_pos_error_history[env_ids] = 0.0 + self._joint_vel_history[env_ids] = 0.0 + + def compute( + self, control_action: ArticulationActions, joint_pos: torch.Tensor, joint_vel: torch.Tensor + ) -> ArticulationActions: + # move history queue by 1 and update top of history + # -- positions + self._joint_pos_error_history = self._joint_pos_error_history.roll(1, 1) + self._joint_pos_error_history[:, 0] = control_action.joint_targets - joint_pos + # -- velocity + self._joint_vel_history = self._joint_vel_history.roll(1, 1) + self._joint_vel_history[:, 0] = joint_vel + # save current joint vel for dc-motor clipping + self._joint_vel[:] = joint_vel + + # compute network inputs + # -- positions + pos_input = torch.cat([self._joint_pos_error_history[:, i].unsqueeze(2) for i in self.cfg.input_idx], dim=2) + pos_input = pos_input.view(self._num_envs * self.num_joints, -1) + # -- velocity + vel_input = torch.cat([self._joint_vel_history[:, i].unsqueeze(2) for i in self.cfg.input_idx], dim=2) + vel_input = vel_input.view(self._num_envs * self.num_joints, -1) + # -- scale and concatenate inputs + if self.cfg.input_order == "pos_vel": + network_input = torch.cat([pos_input * self.cfg.pos_scale, vel_input * self.cfg.vel_scale], dim=1) + elif self.cfg.input_order == "vel_pos": + network_input = torch.cat([vel_input * self.cfg.vel_scale, pos_input * self.cfg.pos_scale], dim=1) + else: + raise ValueError( + f"Invalid input order for MLP actuator net: {self.cfg.input_order}. Must be 'pos_vel' or 'vel_pos'." + ) + + # run network inference + with torch.inference_mode(): + torques = self.network(network_input).view(self._num_envs, self.num_joints) + self.computed_effort = torques.view(self._num_envs, self.num_joints) * self.cfg.torque_scale + + # clip the computed effort based on the motor limits + self.applied_effort = self._clip_effort(self.computed_effort) + + # return torques + control_action.joint_efforts = self.applied_effort + control_action.joint_targets = None + return control_action diff --git a/source/isaaclab/isaaclab/newton/actuators/actuator_pd.py b/source/isaaclab/isaaclab/newton/actuators/actuator_pd.py new file mode 100644 index 00000000000..237adcf9cac --- /dev/null +++ b/source/isaaclab/isaaclab/newton/actuators/actuator_pd.py @@ -0,0 +1,431 @@ +# Copyright (c) 2022-2025, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +from __future__ import annotations + +import torch +from collections.abc import Sequence +from typing import TYPE_CHECKING + +import omni.log +import warp as wp + +from .actuator_base import ActuatorBase +from .kernels import clip_efforts_dc_motor, compute_pd_actuator + +# from isaaclab.utils import DelayBuffer, LinearInterpolation + + +if TYPE_CHECKING: + from isaaclab.actuators.actuator_cfg import ( # DelayedPDActuatorCfg,; RemotizedPDActuatorCfg, + DCMotorCfg, + IdealPDActuatorCfg, + ImplicitActuatorCfg, + ) + + +""" +Implicit Actuator Models. +""" + + +class ImplicitActuator(ActuatorBase): + """Implicit actuator model that is handled by the simulation. + + This performs a similar function as the :class:`IdealPDActuator` class. However, the PD control is handled + implicitly by the simulation which performs continuous-time integration of the PD control law. This is + generally more accurate than the explicit PD control law used in :class:`IdealPDActuator` when the simulation + time-step is large. + + The articulation class sets the stiffness and damping parameters from the implicit actuator configuration + into the simulation. Thus, the class does not perform its own computations on the joint action that + needs to be applied to the simulation. However, it computes the approximate torques for the actuated joint + since PhysX does not expose this quantity explicitly. + + .. caution:: + + The class is only provided for consistency with the other actuator models. It does not implement any + functionality and should not be used. All values should be set to the simulation directly. + """ + + cfg: ImplicitActuatorCfg + """The configuration for the actuator model.""" + + def __init__(self, cfg: ImplicitActuatorCfg, *args, **kwargs): + # effort limits + if cfg.effort_limit_sim is None and cfg.effort_limit is not None: + # throw a warning that we have a replacement for the deprecated parameter + omni.log.warn( + "The object has a value for 'effort_limit'." + " This parameter will be removed in the future." + " To set the effort limit, please use 'effort_limit_sim' instead." + ) + cfg.effort_limit_sim = cfg.effort_limit + elif cfg.effort_limit_sim is not None and cfg.effort_limit is None: + # TODO: Eventually we want to get rid of 'effort_limit' for implicit actuators. + # We should do this once all parameters have an "_sim" suffix. + cfg.effort_limit = cfg.effort_limit_sim + elif cfg.effort_limit_sim is not None and cfg.effort_limit is not None: + if cfg.effort_limit_sim != cfg.effort_limit: + raise ValueError( + "The object has set both 'effort_limit_sim' and 'effort_limit'" + f" and they have different values {cfg.effort_limit_sim} != {cfg.effort_limit}." + " Please only set 'effort_limit_sim' for implicit actuators." + ) + + # velocity limits + if cfg.velocity_limit_sim is None and cfg.velocity_limit is not None: + # throw a warning that previously this was not set + # it leads to different simulation behavior so we want to remain backwards compatible + omni.log.warn( + "The object has a value for 'velocity_limit'." + " Previously, although this value was specified, it was not getting used by implicit" + " actuators. Since this parameter affects the simulation behavior, we continue to not" + " use it. This parameter will be removed in the future." + " To set the velocity limit, please use 'velocity_limit_sim' instead." + ) + cfg.velocity_limit = None + elif cfg.velocity_limit_sim is not None and cfg.velocity_limit is None: + # TODO: Eventually we want to get rid of 'velocity_limit' for implicit actuators. + # We should do this once all parameters have an "_sim" suffix. + cfg.velocity_limit = cfg.velocity_limit_sim + elif cfg.velocity_limit_sim is not None and cfg.velocity_limit is not None: + if cfg.velocity_limit_sim != cfg.velocity_limit: + raise ValueError( + "The object has set both 'velocity_limit_sim' and 'velocity_limit'" + f" and they have different values {cfg.velocity_limit_sim} != {cfg.velocity_limit}." + " Please only set 'velocity_limit_sim' for implicit actuators." + ) + + # set implicit actuator model flag + ImplicitActuator.is_implicit_model = True + # call the base class + super().__init__(cfg, *args, **kwargs) + + """ + Operations. + """ + + def reset(self, *args, **kwargs): + # This is a no-op. There is no state to reset for implicit actuators. + pass + + def compute(self): + """Process the actuator group actions and compute the articulation actions. + + In case of implicit actuator, the control action is directly returned as the computed action. + This function is a no-op and does not perform any computation on the input control action. + However, it computes the approximate torques for the actuated joint since PhysX does not compute + this quantity explicitly. + + Args: + control_action: The joint action instance comprising of the desired joint positions, joint velocities + and (feed-forward) joint efforts. + joint_pos: The current joint positions of the joints in the group. Shape is (num_envs, num_joints). + joint_vel: The current joint velocities of the joints in the group. Shape is (num_envs, num_joints). + + Returns: + The computed desired joint positions, joint velocities and joint efforts. + """ + wp.launch( + compute_pd_actuator, + dim=(self._num_envs, self._num_joints), + inputs=[ + self.joint_data.joint_target, + self.joint_data.joint_effort_target, + self.joint_data.joint_pos, + self.joint_data.joint_vel, + self.joint_data.joint_stiffness, + self.joint_data.joint_damping, + self.joint_data.joint_control_mode, + self.joint_data.computed_effort, + self._env_mask, + self._joint_mask, + ], + ) + self._clip_effort(self.joint_data.computed_effort, self.joint_data.applied_effort) + + +""" +Explicit Actuator Models. +""" + + +class IdealPDActuator(ActuatorBase): + r"""Ideal torque-controlled actuator model with a simple saturation model. + + It employs the following model for computing torques for the actuated joint :math:`j`: + + .. math:: + + \tau_{j, computed} = k_p * (q - q_{des}) + k_d * (\dot{q} - \dot{q}_{des}) + \tau_{ff} + + where, :math:`k_p` and :math:`k_d` are joint stiffness and damping gains, :math:`q` and :math:`\dot{q}` + are the current joint positions and velocities, :math:`q_{des}`, :math:`\dot{q}_{des}` and :math:`\tau_{ff}` + are the desired joint positions, velocities and torques commands. + + The clipping model is based on the maximum torque applied by the motor. It is implemented as: + + .. math:: + + \tau_{j, max} & = \gamma \times \tau_{motor, max} \\ + \tau_{j, applied} & = clip(\tau_{computed}, -\tau_{j, max}, \tau_{j, max}) + + where the clipping function is defined as :math:`clip(x, x_{min}, x_{max}) = min(max(x, x_{min}), x_{max})`. + The parameters :math:`\gamma` is the gear ratio of the gear box connecting the motor and the actuated joint ends, + and :math:`\tau_{motor, max}` is the maximum motor effort possible. These parameters are read from + the configuration instance passed to the class. + """ + + cfg: IdealPDActuatorCfg + """The configuration for the actuator model.""" + + """ + Operations. + """ + + def reset(self, env_ids: Sequence[int]): + pass + + def compute(self): + wp.launch( + compute_pd_actuator, + dim=(self._num_envs, self._num_joints), + inputs=[ + self.joint_data.joint_target, + self.joint_data.joint_effort_target, + self.joint_data.joint_pos, + self.joint_data.joint_vel, + self.joint_data.joint_stiffness, + self.joint_data.joint_damping, + self.joint_data.joint_control_mode, + self.joint_data.computed_effort, + self._env_mask, + self._joint_mask, + ], + ) + self._clip_effort(self.joint_data.computed_effort, self.joint_data.applied_effort) + + +class DCMotor(IdealPDActuator): + r"""Direct control (DC) motor actuator model with velocity-based saturation model. + + It uses the same model as the :class:`IdealActuator` for computing the torques from input commands. + However, it implements a saturation model defined by DC motor characteristics. + + A DC motor is a type of electric motor that is powered by direct current electricity. In most cases, + the motor is connected to a constant source of voltage supply, and the current is controlled by a rheostat. + Depending on various design factors such as windings and materials, the motor can draw a limited maximum power + from the electronic source, which limits the produced motor torque and speed. + + A DC motor characteristics are defined by the following parameters: + + * Continuous-rated speed (:math:`\dot{q}_{motor, max}`) : The maximum-rated speed of the motor. + * Continuous-stall torque (:math:`\tau_{motor, max}`): The maximum-rated torque produced at 0 speed. + * Saturation torque (:math:`\tau_{motor, sat}`): The maximum torque that can be outputted for a short period. + + Based on these parameters, the instantaneous minimum and maximum torques are defined as follows: + + .. math:: + + \tau_{j, max}(\dot{q}) & = clip \left (\tau_{j, sat} \times \left(1 - + \frac{\dot{q}}{\dot{q}_{j, max}}\right), 0.0, \tau_{j, max} \right) \\ + \tau_{j, min}(\dot{q}) & = clip \left (\tau_{j, sat} \times \left( -1 - + \frac{\dot{q}}{\dot{q}_{j, max}}\right), - \tau_{j, max}, 0.0 \right) + + where :math:`\gamma` is the gear ratio of the gear box connecting the motor and the actuated joint ends, + :math:`\dot{q}_{j, max} = \gamma^{-1} \times \dot{q}_{motor, max}`, :math:`\tau_{j, max} = + \gamma \times \tau_{motor, max}` and :math:`\tau_{j, peak} = \gamma \times \tau_{motor, peak}` + are the maximum joint velocity, maximum joint torque and peak torque, respectively. These parameters + are read from the configuration instance passed to the class. + + Using these values, the computed torques are clipped to the minimum and maximum values based on the + instantaneous joint velocity: + + .. math:: + + \tau_{j, applied} = clip(\tau_{computed}, \tau_{j, min}(\dot{q}), \tau_{j, max}(\dot{q})) + + """ + + cfg: DCMotorCfg + """The configuration for the actuator model.""" + + def __init__(self, cfg: DCMotorCfg, *args, **kwargs): + super().__init__(cfg, *args, **kwargs) + # parse configuration + if self.cfg.saturation_effort is not None: + self._saturation_effort = self.cfg.saturation_effort + else: + self._saturation_effort = torch.inf + # check that quantities are provided + if self.cfg.velocity_limit is None: + raise ValueError("The velocity limit must be provided for the DC motor actuator model.") + + """ + Operations. + """ + + def compute(self): + # calculate the desired joint torques + return super().compute() + + """ + Helper functions. + """ + + def _clip_effort(self, effort: wp.array, clipped_effort: wp.array) -> None: + wp.launch( + clip_efforts_dc_motor, + dim=(self._num_envs, self._num_joints), + inputs=[ + self._saturation_effort, + self.joint_data.joint_effort_limits, + self.joint_data.joint_vel_limits, + self.joint_data.joint_vel, + effort, + clipped_effort, + self._env_mask, + self._joint_mask, + ], + ) + + +# class DelayedPDActuator(IdealPDActuator): +# """Ideal PD actuator with delayed command application. +# +# This class extends the :class:`IdealPDActuator` class by adding a delay to the actuator commands. The delay +# is implemented using a circular buffer that stores the actuator commands for a certain number of physics steps. +# The most recent actuation value is pushed to the buffer at every physics step, but the final actuation value +# applied to the simulation is lagged by a certain number of physics steps. +# +# The amount of time lag is configurable and can be set to a random value between the minimum and maximum time +# lag bounds at every reset. The minimum and maximum time lag values are set in the configuration instance passed +# to the class. +# """ +# +# cfg: DelayedPDActuatorCfg +# """The configuration for the actuator model.""" +# +# def __init__(self, cfg: DelayedPDActuatorCfg, *args, **kwargs): +# super().__init__(cfg, *args, **kwargs) +# # instantiate the delay buffers +# self.joint_targets_delay_buffer = DelayBuffer(cfg.max_delay, self._num_envs, device=self._device) +# self.efforts_delay_buffer = DelayBuffer(cfg.max_delay, self._num_envs, device=self._device) +# # all of the envs +# self._ALL_INDICES = torch.arange(self._num_envs, dtype=torch.long, device=self._device) +# +# def reset(self, env_ids: Sequence[int]): +# super().reset(env_ids) +# # number of environments (since env_ids can be a slice) +# if env_ids is None or env_ids == slice(None): +# num_envs = self._num_envs +# else: +# num_envs = len(env_ids) +# # set a new random delay for environments in env_ids +# time_lags = torch.randint( +# low=self.cfg.min_delay, +# high=self.cfg.max_delay + 1, +# size=(num_envs,), +# dtype=torch.int, +# device=self._device, +# ) +# # set delays +# self.joint_targets_delay_buffer.set_time_lag(time_lags, env_ids) +# self.efforts_delay_buffer.set_time_lag(time_lags, env_ids) +# # reset buffers +# self.joint_targets_delay_buffer.reset(env_ids) +# self.efforts_delay_buffer.reset(env_ids) +# +# def compute( +# self, control_action: ArticulationActions, joint_pos: torch.Tensor, joint_vel: torch.Tensor +# ) -> ArticulationActions: +# # apply delay based on the delay the model for all the setpoints +# control_action.joint_targets = self.joint_targets_delay_buffer.compute(control_action.joint_targets) +# control_action.joint_efforts = self.efforts_delay_buffer.compute(control_action.joint_efforts) +# # compte actuator model +# return super().compute(control_action, joint_pos, joint_vel) +# +# +# class RemotizedPDActuator(DelayedPDActuator): +# """Ideal PD actuator with angle-dependent torque limits. +# +# This class extends the :class:`DelayedPDActuator` class by adding angle-dependent torque limits to the actuator. +# The torque limits are applied by querying a lookup table describing the relationship between the joint angle +# and the maximum output torque. The lookup table is provided in the configuration instance passed to the class. +# +# The torque limits are interpolated based on the current joint positions and applied to the actuator commands. +# """ +# +# def __init__( +# self, +# cfg: RemotizedPDActuatorCfg, +# joint_names: list[str], +# joint_ids: Sequence[int], +# num_envs: int, +# device: str, +# control_mode: str = "position", +# stiffness: torch.Tensor | float = 0.0, +# damping: torch.Tensor | float = 0.0, +# armature: torch.Tensor | float = 0.0, +# friction: torch.Tensor | float = 0.0, +# effort_limit: torch.Tensor | float = torch.inf, +# velocity_limit: torch.Tensor | float = torch.inf, +# ): +# # remove effort and velocity box constraints from the base class +# cfg.effort_limit = torch.inf +# cfg.velocity_limit = torch.inf +# # call the base method and set default effort_limit and velocity_limit to inf +# super().__init__( +# cfg, +# joint_names, +# joint_ids, +# num_envs, +# device, +# control_mode, +# stiffness, +# damping, +# armature, +# friction, +# torch.inf, +# torch.inf, +# ) +# self._joint_parameter_lookup = torch.tensor(cfg.joint_parameter_lookup, device=device) +# # define remotized joint torque limit +# self._torque_limit = LinearInterpolation(self.angle_samples, self.max_torque_samples, device=device) +# +# """ +# Properties. +# """ +# +# @property +# def angle_samples(self) -> torch.Tensor: +# return self._joint_parameter_lookup[:, 0] +# +# @property +# def transmission_ratio_samples(self) -> torch.Tensor: +# return self._joint_parameter_lookup[:, 1] +# +# @property +# def max_torque_samples(self) -> torch.Tensor: +# return self._joint_parameter_lookup[:, 2] +# +# """ +# Operations. +# """ +# +# def compute( +# self, control_action: ArticulationActions, joint_pos: torch.Tensor, joint_vel: torch.Tensor +# ) -> ArticulationActions: +# # call the base method +# control_action = super().compute(control_action, joint_pos, joint_vel) +# # compute the absolute torque limits for the current joint positions +# abs_torque_limits = self._torque_limit.compute(joint_pos) +# # apply the limits +# control_action.joint_efforts = torch.clamp( +# control_action.joint_efforts, min=-abs_torque_limits, max=abs_torque_limits +# ) +# self.applied_effort = control_action.joint_efforts +# return control_action diff --git a/source/isaaclab/isaaclab/newton/actuators/kernels.py b/source/isaaclab/isaaclab/newton/actuators/kernels.py new file mode 100644 index 00000000000..0b2391db415 --- /dev/null +++ b/source/isaaclab/isaaclab/newton/actuators/kernels.py @@ -0,0 +1,98 @@ +# Copyright (c) 2022-2025, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +import warp as wp + + +@wp.kernel +def compute_pd_actuator( + joint_targets: wp.array2d(dtype=wp.float32), + added_effort: wp.array2d(dtype=wp.float32), + joint_pos: wp.array2d(dtype=wp.float32), + joint_vel: wp.array2d(dtype=wp.float32), + stiffness: wp.array2d(dtype=wp.float32), + damping: wp.array2d(dtype=wp.float32), + control_mode: wp.array2d(dtype=wp.int32), + computed_effort: wp.array2d(dtype=wp.float32), + env_mask: wp.array(dtype=wp.bool), + joint_mask: wp.array(dtype=wp.bool), +) -> None: + env_index, joint_index = wp.tid() + if env_mask[env_index] and joint_mask[joint_index]: + # No control + if control_mode[env_index, joint_index] == 0: + computed_effort[env_index, joint_index] = ( + stiffness[env_index, joint_index] * (-joint_pos[env_index, joint_index]) + + damping[env_index, joint_index] * (-joint_vel[env_index, joint_index]) + + added_effort[env_index, joint_index] + ) + # Position control + elif control_mode[env_index, joint_index] == 1: + computed_effort[env_index, joint_index] = ( + stiffness[env_index, joint_index] + * (joint_targets[env_index, joint_index] - joint_pos[env_index, joint_index]) + + damping[env_index, joint_index] * (-joint_vel[env_index, joint_index]) + + added_effort[env_index, joint_index] + ) + # Velocity control + elif control_mode[env_index, joint_index] == 2: + computed_effort[env_index, joint_index] = ( + stiffness[env_index, joint_index] * (-joint_pos[env_index, joint_index]) + + damping[env_index, joint_index] + * (joint_targets[env_index, joint_index] - joint_vel[env_index, joint_index]) + + added_effort[env_index, joint_index] + ) + + +@wp.kernel +def clip_efforts_with_limits( + limits: wp.array2d(dtype=wp.float32), + joint_array: wp.array2d(dtype=wp.float32), + clipped_joint_array: wp.array2d(dtype=wp.float32), + env_mask: wp.array(dtype=wp.bool), + joint_mask: wp.array(dtype=wp.bool), +): + env_index, joint_index = wp.tid() + if env_mask[env_index] and joint_mask[joint_index]: + clipped_joint_array[env_index, joint_index] = wp.clamp( + joint_array[env_index, joint_index], -limits[env_index, joint_index], limits[env_index, joint_index] + ) + + +@wp.func +def clip_effort_dc_motor( + saturation_effort: float, + vel_limit: float, + effort_limit: float, + joint_vel: float, + effort: float, +): + max_effort = saturation_effort * (1.0 - joint_vel) / vel_limit + min_effort = saturation_effort * (-1.0 - joint_vel) / vel_limit + max_effort = wp.clamp(max_effort, 0.0, effort_limit) + min_effort = wp.clamp(min_effort, -effort_limit, 0.0) + return wp.clamp(effort, min_effort, max_effort) + + +@wp.kernel +def clip_efforts_dc_motor( + saturation_effort: wp.array2d(dtype=wp.float32), + effort_limit: wp.array2d(dtype=wp.float32), + vel_limit: wp.array2d(dtype=wp.float32), + joint_vel: wp.array2d(dtype=wp.float32), + joint_array: wp.array2d(dtype=wp.float32), + clipped_joint_array: wp.array2d(dtype=wp.float32), + env_mask: wp.array(dtype=wp.bool), + joint_mask: wp.array(dtype=wp.bool), +): + env_index, joint_index = wp.tid() + if env_mask[env_index] and joint_mask[joint_index]: + clipped_joint_array[env_index, joint_index] = clip_effort_dc_motor( + saturation_effort[env_index, joint_index], + vel_limit[env_index, joint_index], + effort_limit[env_index, joint_index], + joint_vel[env_index, joint_index], + joint_array[env_index, joint_index], + ) diff --git a/source/isaaclab/isaaclab/newton/assets/articulation/__init__.py b/source/isaaclab/isaaclab/newton/assets/articulation/__init__.py new file mode 100644 index 00000000000..e2d97161ee7 --- /dev/null +++ b/source/isaaclab/isaaclab/newton/assets/articulation/__init__.py @@ -0,0 +1,8 @@ +# Copyright (c) 2022-2025, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +"""Sub-module for rigid articulated assets.""" +from .articulation import Articulation +from .articulation_data import ArticulationData \ No newline at end of file diff --git a/source/isaaclab/isaaclab/newton/assets/articulation/articulation.py b/source/isaaclab/isaaclab/newton/assets/articulation/articulation.py new file mode 100644 index 00000000000..7876265cbd0 --- /dev/null +++ b/source/isaaclab/isaaclab/newton/assets/articulation/articulation.py @@ -0,0 +1,1655 @@ +# Copyright (c) 2022-2025, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +# Flag for pyright to ignore type errors in this file. +# pyright: reportPrivateUsage=false + +from __future__ import annotations +import torch +import logging +import warnings +from collections.abc import Sequence +from prettytable import PrettyTable +from typing import TYPE_CHECKING + +import warp as wp +from newton import JointType, Model +from newton.selection import ArticulationView as NewtonArticulationView +from pxr import UsdPhysics + +import isaaclab.sim as sim_utils +import isaaclab.utils.string as string_utils +from isaaclab.sim._impl.newton_manager import NewtonManager +from isaaclab.assets.articulation.base_articulation import BaseArticulation + +from isaaclab.newton.actuators import ActuatorBase, ImplicitActuator +from isaaclab.newton.assets.articulation.articulation_data import ArticulationData +from isaaclab.newton.assets.core.joint_properties.joint import Joint +from isaaclab.newton.assets.core.body_properties.body import Body +from isaaclab.newton.assets.core.root_properties.root import Root +from isaaclab.newton.assets.core.kernels import ( + generate_mask_from_ids, + populate_empty_array, + update_batched_array_with_array_masked, + update_array_with_value, + vec13f, +) +from isaaclab.newton import Frontend + +if TYPE_CHECKING: + from isaaclab.assets.articulation.articulation_cfg import ArticulationCfg + from isaaclab.actuators.actuator_cfg import ActuatorBaseCfg + +logger = logging.getLogger(__name__) +warnings.simplefilter("once", UserWarning) +logging.captureWarnings(True) + +class Articulation(BaseArticulation): + """An articulation asset class. + + An articulation is a collection of rigid bodies connected by joints. The joints can be either + fixed or actuated. The joints can be of different types, such as revolute, prismatic, D-6, etc. + However, the articulation class has currently been tested with revolute and prismatic joints. + The class supports both floating-base and fixed-base articulations. The type of articulation + is determined based on the root joint of the articulation. If the root joint is fixed, then + the articulation is considered a fixed-base system. Otherwise, it is considered a floating-base + system. This can be checked using the :attr:`Articulation.is_fixed_base` attribute. + + For an asset to be considered an articulation, the root prim of the asset must have the + `USD ArticulationRootAPI`_. This API is used to define the sub-tree of the articulation using + the reduced coordinate formulation. On playing the simulation, the physics engine parses the + articulation root prim and creates the corresponding articulation in the physics engine. The + articulation root prim can be specified using the :attr:`AssetBaseCfg.prim_path` attribute. + + The articulation class also provides the functionality to augment the simulation of an articulated + system with custom actuator models. These models can either be explicit or implicit, as detailed in + the :mod:`isaaclab.actuators` module. The actuator models are specified using the + :attr:`ArticulationCfg.actuators` attribute. These are then parsed and used to initialize the + corresponding actuator models, when the simulation is played. + + During the simulation step, the articulation class first applies the actuator models to compute + the joint commands based on the user-specified targets. These joint commands are then applied + into the simulation. The joint commands can be either position, velocity, or effort commands. + As an example, the following snippet shows how this can be used for position commands: + + .. code-block:: python + + # an example instance of the articulation class + my_articulation = Articulation(cfg) + + # set joint position targets + my_articulation.set_joint_position_target(position) + # propagate the actuator models and apply the computed commands into the simulation + my_articulation.write_data_to_sim() + + # step the simulation using the simulation context + sim_context.step() + + # update the articulation state, where dt is the simulation time step + my_articulation.update(dt) + + .. _`USD ArticulationRootAPI`: https://openusd.org/dev/api/class_usd_physics_articulation_root_a_p_i.html + + """ + + cfg: ArticulationCfg + """Configuration instance for the articulations.""" + + __backend_name__: str = "newton" + """The name of the backend for the articulation.""" + + actuators: dict[str, ActuatorBase] + """Dictionary of actuator instances for the articulation. + + The keys are the actuator names and the values are the actuator instances. The actuator instances + are initialized based on the actuator configurations specified in the :attr:`ArticulationCfg.actuators` + attribute. They are used to compute the joint commands during the :meth:`write_data_to_sim` function. + """ + + def __init__(self, cfg: ArticulationCfg, frontend: Frontend = Frontend.TORCH): + """Initialize the articulation. + + Args: + cfg: A configuration instance. + """ + super().__init__(cfg) + self._frontend = frontend + + """ + Properties + """ + + @property + def data(self) -> ArticulationData: + return self._data + + @property + def num_instances(self) -> int: + return self._root_view.count + + @property + def is_fixed_base(self) -> bool: + """Whether the articulation is a fixed-base or floating-base system.""" + return self._root_view.is_fixed_base + + @property + def num_joints(self) -> int: + """Number of joints in articulation.""" + return self._joint.num_joints + + @property + def num_fixed_tendons(self) -> int: + """Number of fixed tendons in articulation.""" + return 0 + + @property + def num_spatial_tendons(self) -> int: + """Number of spatial tendons in articulation.""" + return 0 + + @property + def num_bodies(self) -> int: + """Number of bodies in articulation.""" + return self._body.num_bodies + + @property + def joint_names(self) -> list[str]: + """Ordered names of joints in articulation.""" + return self._joint.joint_names + + @property + def fixed_tendon_names(self) -> list[str]: + """Ordered names of fixed tendons in articulation.""" + # TODO: check if the articulation has fixed tendons + return [] + + @property + def spatial_tendon_names(self) -> list[str]: + """Ordered names of spatial tendons in articulation.""" + # TODO: check if the articulation has spatial tendons + return [] + + @property + def body_names(self) -> list[str]: + """Ordered names of bodies in articulation.""" + return self._body.body_names + + @property + def root_view(self) -> NewtonArticulationView: + """Articulation view for the asset (Newton). + + Note: + Use this view with caution. It requires handling of tensors in a specific way. + """ + return self._root_view + + @property + def root_newton_model(self) -> Model: + """Newton model for the asset.""" + return self._root_view.model + + """ + Operations. + """ + + def reset(self, ids: Sequence[int] | None = None, mask: wp.array | None = None): + if self._frontend == Frontend.TORCH: + if ids is not None: + mask = torch.zeros(self.num_instances, dtype=torch.bool, device=self.device) + mask[ids] = True + mask = wp.from_torch(mask, dtype=wp.bool) + # use ellipses object to skip initial indices. + self._root.reset(mask) + self._joint.reset(mask) + self._body.reset(mask) + + def write_data_to_sim(self): + """Write external wrenches and joint commands to the simulation. + + If any explicit actuators are present, then the actuator models are used to compute the + joint commands. Otherwise, the joint commands are directly set into the simulation. + + Note: + We write external wrench to the simulation here since this function is called before the simulation step. + This ensures that the external wrench is applied at every simulation step. + """ + # Wrenches are automatically applied by set_external_force_and_torque. + # apply actuator models + self._apply_actuator_model() + + def update(self, dt: float): + self._data.update(dt) + + """ + Frontend conversions - Torch to Warp. + """ + + def _torch_to_warp_single_index( + self, + value: torch.Tensor, + N: int, + ids: Sequence[int] | None = None, + mask: torch.Tensor | None = None, + dtype: type = wp.float32, + ) -> tuple[wp.array, wp.array | None]: + """ Converts any Torch frontend data into warp data with single index support. + + Args: + value: The value to convert. Shape is (N,). + N: The number of elements in the value. + ids: The index ids. + mask: The index mask. + + Returns: + A tuple of warp data with its mask. + """ + if mask is None: + if ids is not None: + # Create a mask from scratch + env_mask = torch.zeros(N, dtype=torch.bool, device=self.device) + env_mask[ids] = True + env_mask = wp.from_torch(env_mask, dtype=wp.bool) + # Create a complete data buffer from scratch + complete = torch.zeros((N, *value.shape[1:]), dtype=value.dtype, device=self.device) + complete[ids] = value + value = wp.from_torch(complete, dtype=dtype) + else: + value = wp.from_torch(value, dtype=dtype) + else: + if ids is not None: + warnings.warn( + "ids is not None, but mask is provided. Ignoring ids. Please make sure you are providing complete data buffers.", + UserWarning, + ) + env_mask = wp.from_torch(mask, dtype=wp.bool) + value = wp.from_torch(value, dtype=dtype) + return value, env_mask + + def _torch_to_warp_dual_index( + self, + value: torch.Tensor, + N: int, + M: int, + first_ids: Sequence[int] | None = None, + second_ids: Sequence[int] | None = None, + first_mask: torch.Tensor | None = None, + second_mask: torch.Tensor | None = None, + dtype: type = wp.float32, + ) -> tuple[wp.array, wp.array | None, wp.array | None]: + """ Converts any Torch frontend data into warp data with dual index support. + + Args: + value: The value to convert. Shape is (N, M) or (len(first_ids), len(second_ids)). + first_ids: The first index ids. + second_ids: The second index ids. + first_mask: The first index mask. + second_mask: The second index mask. + dtype: The dtype of the value. + + Returns: + A tuple of warp data with its two masks. + """ + if first_mask is None: + if (first_ids is not None) or (second_ids is not None): + # If we are provided with either first_ids or second_ids, we need to create a mask from scratch and + # we are expecting partial values. + if first_ids is not None: + # Create a mask from scratch + first_mask = torch.zeros(N, dtype=torch.bool, device=self.device) + first_mask[first_ids] = True + first_mask = wp.from_torch(first_mask, dtype=wp.bool) + else: + first_ids = slice(None) + if second_ids is not None: + # Create a mask from scratch + second_mask = torch.zeros(M, dtype=torch.bool, device=self.device) + second_mask[second_ids] = True + second_mask = wp.from_torch(second_mask, dtype=wp.bool) + else: + second_ids = slice(None) + if first_ids != slice(None) and second_ids != slice(None): + first_ids = first_ids[:, None] + + # Create a complete data buffer from scratch + complete_value = torch.zeros(N, M, dtype=value.dtype, device=self.device) + complete_value[first_ids, second_ids] = value + value = wp.from_torch(complete_value, dtype=dtype) + elif second_mask is not None: + second_mask = wp.from_torch(second_mask, dtype=wp.bool) + value = wp.from_torch(value, dtype=dtype) + else: + value = wp.from_torch(value, dtype=dtype) + else: + if (first_ids is not None) or (second_ids is not None): + warnings.warn( + "Mask and ids are provided. Ignoring ids. Please make sure you are providing complete data buffers.", + UserWarning, + ) + first_mask = wp.from_torch(first_mask, dtype=wp.bool) + if second_mask is not None: + second_mask = wp.from_torch(second_mask, dtype=wp.bool) + else: + value = wp.from_torch(value, dtype=dtype) + + return value, first_mask, second_mask + + """ + Operations - Finders. + """ + + def find_bodies( + self, name_keys: str | Sequence[str], preserve_order: bool = False + ) -> tuple[wp.array | torch.Tensor, list[str], list[int]]: + """Find bodies in the articulation based on the name keys. + + Please check the :meth:`isaaclab.utils.string_utils.resolve_matching_names` function for more + information on the name matching. + + Args: + name_keys: A regular expression or a list of regular expressions to match the body names. + preserve_order: Whether to preserve the order of the name keys in the output. Defaults to False. + + Returns: + A tuple of lists containing the body mask, names, and indices. + """ + if self._frontend == Frontend.TORCH: + mask, names, indices = self._body.find_bodies(name_keys, preserve_order) + return wp.to_torch(mask), names, indices + else: + return self._body.find_bodies(name_keys, preserve_order) + + def find_joints( + self, name_keys: str | Sequence[str], joint_subset: list[str] | None = None, preserve_order: bool = False + ) -> tuple[wp.array | torch.Tensor, list[str], list[int]]: + """Find joints in the articulation based on the name keys. + + Please see the :func:`isaaclab.utils.string.resolve_matching_names` function for more information + on the name matching. + + Args: + name_keys: A regular expression or a list of regular expressions to match the joint names. + joint_subset: A subset of joints to search for. Defaults to None, which means all joints + in the articulation are searched. + preserve_order: Whether to preserve the order of the name keys in the output. Defaults to False. + + Returns: + A tuple of lists containing the joint mask, names, and indices. + """ + if self._frontend == Frontend.TORCH: + mask, names, indices = self._joint.find_joints(name_keys, joint_subset, preserve_order) + return wp.to_torch(mask), names, indices + else: + return self._joint.find_joints(name_keys, joint_subset, preserve_order) + + def find_fixed_tendons( + self, name_keys: str | Sequence[str], tendon_subsets: list[str] | None = None, preserve_order: bool = False + ) -> tuple[wp.array | torch.Tensor, list[str], list[int]]: + """Find fixed tendons in the articulation based on the name keys. + + Please see the :func:`isaaclab.utils.string.resolve_matching_names` function for more information + on the name matching. + + Args: + name_keys: A regular expression or a list of regular expressions to match the joint + names with fixed tendons. + tendon_subsets: A subset of joints with fixed tendons to search for. Defaults to None, which means + all joints in the articulation are searched. + preserve_order: Whether to preserve the order of the name keys in the output. Defaults to False. + + Returns: + A tuple of lists containing the tendon mask, names, and indices. + """ + raise NotImplementedError("Fixed tendons are not supported in Newton.") + + def find_spatial_tendons( + self, name_keys: str | Sequence[str], tendon_subsets: list[str] | None = None, preserve_order: bool = False + ) -> tuple[wp.array | torch.Tensor, list[str], list[int]]: + """Find spatial tendons in the articulation based on the name keys. + + Please see the :func:`isaaclab.utils.string.resolve_matching_names` function for more information + on the name matching. + + Args: + name_keys: A regular expression or a list of regular expressions to match the tendon names. + tendon_subsets: A subset of tendons to search for. Defaults to None, which means all tendons + in the articulation are searched. + preserve_order: Whether to preserve the order of the name keys in the output. Defaults to False. + + Returns: + A tuple of lists containing the tendon mask, names, and indices. + """ + raise NotImplementedError("Spatial tendons are not supported in Newton.") + """ + Operations - State Writers. + """ + + def write_root_state_to_sim( + self, + root_state: wp.array | torch.Tensor, + env_ids: Sequence[int] | None = None, + env_mask: wp.array | torch.Tensor | None = None, + ) -> None: + """Set the root state over selected environment indices into the simulation. + + The root state comprises of the cartesian position, quaternion orientation in (x, y, z, w), and angular + and linear velocity. All the quantities are in the simulation frame. + + Args: + root_state: Root state in simulation frame. Shape is (len(env_ids), 13) or (num_instances, 13). + env_ids: Environment indices. If None, then all indices are used. + env_mask: Environment mask. Shape is (num_instances,). + """ + # Resolve indices into mask, convert from partial data to complete data, handles the conversion to warp. + if self._frontend == Frontend.TORCH: + root_state, env_mask = self._torch_to_warp_single_index(root_state, self.num_instances, env_ids, env_mask, dtype=vec13f) + # Write to Newton using warp + self._root.write_root_state_to_sim(root_state, env_mask=env_mask) + + def write_root_com_state_to_sim( + self, + root_state: wp.array | torch.Tensor, + env_ids: Sequence[int] | None = None, + env_mask: wp.array | torch.Tensor | None = None, + ) -> None: + """Set the root center of mass state over selected environment indices into the simulation. + + The root state comprises of the cartesian position, quaternion orientation in (x, y, z, w), and angular + and linear velocity. All the quantities are in the simulation frame. + + Args: + root_state: Root state in simulation frame. Shape is (len(env_ids), 13) or (num_instances, 13). + env_ids: Environment indices. If None, then all indices are used. + env_mask: Environment mask. Shape is (num_instances,). + """ + # Resolve indices into mask, convert from partial data to complete data, handles the conversion to warp. + if self._frontend == Frontend.TORCH: + root_state, env_mask = self._torch_to_warp_single_index(root_state, self.num_instances, env_ids, env_mask, dtype=vec13f) + # Write to Newton using warp + self._root.write_root_com_state_to_sim(root_state, env_mask=env_mask) + + def write_root_link_state_to_sim( + self, + root_state: wp.array | torch.Tensor, + env_ids: Sequence[int] | None = None, + env_mask: wp.array | torch.Tensor | None = None, + ) -> None: + """Set the root link state over selected environment indices into the simulation. + + The root state comprises of the cartesian position, quaternion orientation in (x, y, z, w), and angular + and linear velocity. All the quantities are in the simulation frame. + + Args: + root_state: Root state in simulation frame. Shape is (len(env_ids), 13) or (num_instances, 13). + env_ids: Environment indices. If None, then all indices are used. + env_mask: Environment mask. Shape is (num_instances,). + """ + # Resolve indices into mask, convert from partial data to complete data, handles the conversion to warp. + if self._frontend == Frontend.TORCH: + root_state, env_mask = self._torch_to_warp_single_index(root_state, self.num_instances, env_ids, env_mask, dtype=vec13f) + # Write to Newton using warp + self._root.write_root_link_state_to_sim(root_state, env_mask=env_mask) + + def write_root_pose_to_sim( + self, + root_pose: wp.array | torch.Tensor, + env_ids: Sequence[int] | None = None, + env_mask: wp.array | torch.Tensor | None = None + ) -> None: + """Set the root pose over selected environment indices into the simulation. + + The root pose comprises of the cartesian position and quaternion orientation in (x, y, z, w). + + Args: + root_pose: Root poses in simulation frame. Shape is (len(env_ids), 7) or (num_instances, 7). + env_ids: Environment indices. If None, then all indices are used. + env_mask: Environment mask. Shape is (num_instances,). + """ + # Resolve indices into mask, convert from partial data to complete data, handles the conversion to warp. + if self._frontend == Frontend.TORCH: + root_pose, env_mask = self._torch_to_warp_single_index(root_pose, self.num_instances, env_ids, env_mask, dtype=wp.transformf) + # Write to Newton using warp + self._root.write_root_link_pose_to_sim(root_pose, env_mask=env_mask) + + def write_root_link_pose_to_sim( + self, + pose: wp.array | torch.Tensor, + env_ids: Sequence[int] | None = None, + env_mask: wp.array | torch.Tensor | None = None, + ) -> None: + """Set the root link pose over selected environment indices into the simulation. + + + The root pose ``wp.transformf`` comprises of the cartesian position and quaternion orientation in (x, y, z, w). + + Args: + root_pose: Root poses in simulation frame. Shape is (len(env_ids), 7) or (num_instances, 7). + env_ids: Environment indices. If None, then all indices are used. + env_mask: Environment mask. Shape is (num_instances,). + """ + # Resolve indices into mask, convert from partial data to complete data, handles the conversion to warp. + if self._frontend == Frontend.TORCH: + pose, env_mask = self._torch_to_warp_single_index(pose, self.num_instances, env_ids, env_mask, dtype=wp.transformf) + # Write to Newton using warp + self._root.write_root_link_pose_to_sim(pose, env_mask=env_mask) + + def write_root_com_pose_to_sim( + self, + root_pose: wp.array | torch.Tensor, + env_ids: Sequence[int] | None = None, + env_mask: wp.array | torch.Tensor | None = None, + ) -> None: + """Set the root center of mass pose over selected environment indices into the simulation. + + The root pose comprises of the cartesian position and quaternion orientation in (w, x, y, z). + The orientation is the orientation of the principle axes of inertia. + + Args: + root_pose: Root center of mass poses in simulation frame. Shape is (len(env_ids), 7) or (num_instances, 7). + env_ids: Environment indices. If None, then all indices are used. + env_mask: Environment mask. Shape is (num_instances,). + """ + # Resolve indices into mask, convert from partial data to complete data, handles the conversion to warp. + if self._frontend == Frontend.TORCH: + root_pose, env_mask = self._torch_to_warp_single_index(root_pose, self.num_instances, env_ids, env_mask, dtype=wp.transformf) + # Write to Newton using warp + self._root.write_root_com_pose_to_sim(root_pose, env_mask=env_mask) + + def write_root_velocity_to_sim( + self, + root_velocity: wp.array | torch.Tensor, + env_ids: Sequence[int] | None = None, + env_mask: wp.array | torch.Tensor | None = None, + ) -> None: + """Set the root center of mass velocity over selected environment indices into the simulation. + + The velocity comprises linear velocity (x, y, z) and angular velocity (x, y, z) in that order. + NOTE: This sets the velocity of the root's center of mass rather than the roots frame. + + Args: + root_velocity: Root center of mass velocities in simulation world frame. Shape is (len(env_ids), 6) or (num_instances, 6). + env_ids: Environment indices. If None, then all indices are used. + env_mask: Environment mask. Shape is (num_instances,). + """ + # Resolve indices into mask, convert from partial data to complete data, handles the conversion to warp. + if self._frontend == Frontend.TORCH: + root_velocity, env_mask = self._torch_to_warp_single_index(root_velocity, self.num_instances, env_ids, env_mask, dtype=wp.spatial_vectorf) + # Write to Newton using warp + self._root.write_root_com_velocity_to_sim(root_velocity=root_velocity, env_mask=env_mask) + + def write_root_com_velocity_to_sim( + self, + root_velocity: wp.array | torch.Tensor, + env_ids: Sequence[int] | None = None, + env_mask: wp.array | torch.Tensor | None = None, + ) -> None: + """Set the root center of mass velocity over selected environment indices into the simulation. + + The velocity comprises linear velocity (x, y, z) and angular velocity (x, y, z) in that order. + NOTE: This sets the velocity of the root's center of mass rather than the roots frame. + + Args: + root_velocity: Root center of mass velocities in simulation world frame. Shape is (len(env_ids), 6) or (num_instances, 6). + env_ids: Environment indices. If None, then all indices are used. + env_mask: Environment mask. Shape is (num_instances,). + """ + # Resolve indices into mask, convert from partial data to complete data, handles the conversion to warp. + if self._frontend == Frontend.TORCH: + root_velocity, env_mask = self._torch_to_warp_single_index(root_velocity, self.num_instances, env_ids, env_mask, dtype=wp.spatial_vectorf) + # Write to Newton using warp + self._root.write_root_com_velocity_to_sim(root_velocity, env_mask=env_mask) + + def write_root_link_velocity_to_sim(self, root_velocity: wp.array, env_ids: Sequence[int] | None = None, env_mask: wp.array | None = None) -> None: + """Set the root link velocity over selected environment indices into the simulation. + + The velocity comprises linear velocity (x, y, z) and angular velocity (x, y, z) in that order. + NOTE: This sets the velocity of the root's frame rather than the roots center of mass. + + Args: + root_velocity: Root frame velocities in simulation world frame. Shape is (len(env_ids), 6) or (num_instances, 6). + env_ids: Environment indices. If None, then all indices are used. + env_mask: Environment mask. Shape is (num_instances,). + """ + # Resolve indices into mask, convert from partial data to complete data, handles the conversion to warp. + if self._frontend == Frontend.TORCH: + root_velocity, env_mask = self._torch_to_warp_single_index(root_velocity, self.num_instances, env_ids, env_mask, dtype=wp.spatial_vectorf) + # Write to Newton using warp + self._root.write_root_link_velocity_to_sim(root_velocity, env_mask=env_mask) + + def write_joint_state_to_sim( + self, + position: wp.array, + velocity: wp.array, + joint_ids: Sequence[int] | None = None, + env_ids: Sequence[int] | None = None, + joint_mask: wp.array | None = None, + env_mask: wp.array | None = None, + ): + """Write joint positions and velocities to the simulation. + + Args: + position: Joint positions. Shape is (len(env_ids), len(joint_ids)) or (num_instances, num_joints). + velocity: Joint velocities. Shape is (len(env_ids), len(joint_ids)) or (num_instances, num_joints). + joint_ids: The joint indices to set the targets for. Defaults to None (all joints). + env_ids: The environment indices to set the targets for. Defaults to None (all environments). + joint_mask: The joint mask. Shape is (num_joints). + env_mask: The environment mask. Shape is (num_instances,). + """ + # Resolve indices into mask, convert from partial data to complete data, handles the conversion to warp. + if self._frontend == Frontend.TORCH: + position, _, _ = self._torch_to_warp_dual_index(position, self.num_instances, self.num_joints, env_ids, joint_ids, env_mask, joint_mask, dtype=wp.float32) + velocity, env_mask, joint_mask = self._torch_to_warp_dual_index(velocity, self.num_instances, self.num_joints, env_ids, joint_ids, env_mask, joint_mask, dtype=wp.float32) + # set into simulation + self._joint.write_joint_state_to_sim(position, velocity, joint_mask, env_mask) + + def write_joint_position_to_sim( + self, + position: wp.array, + joint_ids: Sequence[int] | None = None, + env_ids: Sequence[int] | None = None, + joint_mask: wp.array | None = None, + env_mask: wp.array | None = None, + ): + """Write joint positions to the simulation. + + Args: + position: Joint positions. Shape is (len(env_ids), len(joint_ids)) or (num_instances, num_joints). + joint_ids: The joint indices to set the targets for. Defaults to None (all joints). + env_ids: The environment indices to set the targets for. Defaults to None (all environments). + joint_mask: The joint mask. Shape is (num_joints). + env_mask: The environment mask. Shape is (num_instances,). + """ + # Resolve indices into mask, convert from partial data to complete data, handles the conversion to warp. + if self._frontend == Frontend.TORCH: + position, env_mask, joint_mask = self._torch_to_warp_dual_index(position, self.num_instances, self.num_joints, env_ids, joint_ids, env_mask, joint_mask, dtype=wp.float32) + # set into simulation + self._joint.write_joint_position_to_sim(position, joint_mask, env_mask) + + def write_joint_velocity_to_sim( + self, + velocity: wp.array, + joint_ids: Sequence[int] | None = None, + env_ids: Sequence[int] | None = None, + joint_mask: wp.array | None = None, + env_mask: wp.array | None = None, + ): + """Write joint velocities to the simulation. + + Args: + velocity: Joint velocities. Shape is (len(env_ids), len(joint_ids)) or (num_instances, num_joints). + joint_ids: The joint indices to set the targets for. Defaults to None (all joints). + env_ids: The environment indices to set the targets for. Defaults to None (all environments). + joint_mask: The joint mask. Shape is (num_joints). + env_mask: The environment mask. Shape is (num_instances,). + """ + # Resolve indices into mask, convert from partial data to complete data, handles the conversion to warp. + if self._frontend == Frontend.TORCH: + velocity, env_mask, joint_mask = self._torch_to_warp_dual_index(velocity, self.num_instances, self.num_joints, env_ids, joint_ids, env_mask, joint_mask, dtype=wp.float32) + # set into simulation + self._joint.write_joint_velocity_to_sim(velocity, joint_mask, env_mask) + + """ + Operations - Simulation Parameters Writers. + """ + + def write_joint_control_mode_to_sim( + self, + control_mode: wp.array | int, + joint_ids: Sequence[int] | None = None, + env_ids: Sequence[int] | None = None, + joint_mask: wp.array | None = None, + env_mask: wp.array | None = None, + ): + """Write joint control mode into the simulation. + + Args: + control_mode: Joint control mode. Shape is (len(env_ids), len(joint_ids)) or (num_instances, num_joints). + joint_ids: The joint indices to set the targets for. Defaults to None (all joints). + env_ids: The environment indices to set the targets for. Defaults to None (all environments). + joint_mask: The joint mask. Shape is (num_joints). + env_mask: The environment mask. Shape is (num_instances,). + + Raises: + ValueError: If the control mode is invalid. + """ + # Resolve indices into mask, convert from partial data to complete data, handles the conversion to warp. + if self._frontend == Frontend.TORCH: + control_mode, env_mask, joint_mask = self._torch_to_warp_dual_index(control_mode, self.num_instances, self.num_joints, env_ids, joint_ids, env_mask, joint_mask, dtype=wp.float32) + self._joint.write_joint_control_mode_to_sim(control_mode, joint_mask, env_mask) + + def write_joint_stiffness_to_sim( + self, + stiffness: wp.array | float, + joint_ids: Sequence[int] | None = None, + env_ids: Sequence[int] | None = None, + joint_mask: wp.array | None = None, + env_mask: wp.array | None = None, + ) -> None: + """Write joint stiffness into the simulation. + + Args: + stiffness: Joint stiffness. Shape is (len(env_ids), len(joint_ids)) or (num_instances, num_joints). + joint_ids: The joint indices to set the targets for. Defaults to None (all joints). + env_ids: The environment indices to set the targets for. Defaults to None (all environments). + joint_mask: The joint mask. Shape is (num_joints). + env_mask: The environment mask. Shape is (num_instances,). + """ + # Resolve indices into mask, convert from partial data to complete data, handles the conversion to warp. + if self._frontend == Frontend.TORCH: + stiffness, env_mask, joint_mask = self._torch_to_warp_dual_index(stiffness, self.num_instances, self.num_joints, env_ids, joint_ids, env_mask, joint_mask, dtype=wp.float32) + self._joint.write_joint_stiffness_to_sim(stiffness, joint_mask, env_mask) + + def write_joint_damping_to_sim( + self, + damping: wp.array | float, + joint_ids: Sequence[int] | None = None, + env_ids: Sequence[int] | None = None, + joint_mask: wp.array | None = None, + env_mask: wp.array | None = None, + ) -> None: + """Write joint damping into the simulation. + + Args: + damping: Joint damping. Shape is (len(env_ids), len(joint_ids)) or (num_instances, num_joints). + joint_ids: The joint indices to set the targets for. Defaults to None (all joints). + env_ids: The environment indices to set the targets for. Defaults to None (all environments). + joint_mask: The joint mask. Shape is (num_joints). + env_mask: The environment mask. Shape is (num_instances,). + """ + # Resolve indices into mask, convert from partial data to complete data, handles the conversion to warp. + if self._frontend == Frontend.TORCH: + damping, env_mask, joint_mask = self._torch_to_warp_dual_index(damping, self.num_instances, self.num_joints, env_ids, joint_ids, env_mask, joint_mask, dtype=wp.float32) + self._joint.write_joint_damping_to_sim(damping, joint_mask, env_mask) + + def write_joint_position_limit_to_sim( + self, + upper_limits: wp.array | float, + lower_limits: wp.array | float, + joint_ids: Sequence[int] | None = None, + env_ids: Sequence[int] | None = None, + joint_mask: wp.array | None = None, + env_mask: wp.array | None = None, + ) -> None: + """Write joint position limits into the simulation. + + Args: + upper_limits: Joint upper limits. Shape is (len(env_ids), len(joint_ids)) or (num_instances, num_joints). + lower_limits: Joint lower limits. Shape is (num_instances, num_joints). + joint_ids: The joint indices to set the targets for. Defaults to None (all joints). + env_ids: The environment indices to set the targets for. Defaults to None (all environments). + joint_mask: The joint mask. Shape is (num_joints). + env_mask: The environment mask. Shape is (num_instances,). + """ + # Resolve indices into mask, convert from partial data to complete data, handles the conversion to warp. + if self._frontend == Frontend.TORCH: + upper_limits, _, _ = self._torch_to_warp_dual_index(upper_limits, self.num_instances, self.num_joints, env_ids, joint_ids, env_mask, joint_mask, dtype=wp.float32) + lower_limits, env_mask, joint_mask = self._torch_to_warp_dual_index(lower_limits, self.num_instances, self.num_joints, env_ids, joint_ids, env_mask, joint_mask, dtype=wp.float32) + self._joint.write_joint_position_limit_to_sim(upper_limits, lower_limits, joint_mask, env_mask) + + def write_joint_velocity_limit_to_sim( + self, + limits: wp.array | float, + joint_ids: Sequence[int] | None = None, + env_ids: Sequence[int] | None = None, + joint_mask: wp.array | None = None, + env_mask: wp.array | None = None, + ) -> None: + """Write joint max velocity to the simulation. + + The velocity limit is used to constrain the joint velocities in the physics engine. The joint will only + be able to reach this velocity if the joint's effort limit is sufficiently large. If the joint is moving + faster than this velocity, the physics engine will actually try to brake the joint to reach this velocity. + + .. warn:: This function is ignored when using the Mujoco solver. + + Args: + limits: Joint max velocity. Shape is (len(env_ids), len(joint_ids)) or (num_instances, num_joints). + joint_ids: The joint indices to set the targets for. Defaults to None (all joints). + env_ids: The environment indices to set the targets for. Defaults to None (all environments). + joint_mask: The joint mask. Shape is (num_joints). + env_mask: The environment mask. Shape is (num_instances,). + """ + # Resolve indices into mask, convert from partial data to complete data, handles the conversion to warp. + if self._frontend == Frontend.TORCH: + limits, env_mask, joint_mask = self._torch_to_warp_dual_index(limits, self.num_instances, self.num_joints, env_ids, joint_ids, env_mask, joint_mask, dtype=wp.float32) + self._joint.write_joint_velocity_limit_to_sim(limits, joint_mask, env_mask) + + def write_joint_effort_limit_to_sim( + self, + limits: wp.array | float, + joint_ids: Sequence[int] | None = None, + env_ids: Sequence[int] | None = None, + joint_mask: wp.array | None = None, + env_mask: wp.array | None = None, + ) -> None: + """Write joint effort limits into the simulation. + + The effort limit is used to constrain the computed joint efforts in the physics engine. If the + computed effort exceeds this limit, the physics engine will clip the effort to this value. + + Args: + limits: Joint torque limits. Shape is (len(env_ids), len(joint_ids)) or (num_instances, num_joints). + joint_ids: The joint indices to set the targets for. Defaults to None (all joints). + env_ids: The environment indices to set the targets for. Defaults to None (all environments). + joint_mask: The joint mask. Shape is (num_joints). + env_mask: The environment mask. Shape is (num_instances,). + """ + # Resolve indices into mask, convert from partial data to complete data, handles the conversion to warp. + if self._frontend == Frontend.TORCH: + limits, env_mask, joint_mask = self._torch_to_warp_dual_index(limits, self.num_instances, self.num_joints, env_ids, joint_ids, env_mask, joint_mask, dtype=wp.float32) + self._joint.write_joint_effort_limit_to_sim(limits, joint_mask, env_mask) + + def write_joint_armature_to_sim( + self, + armature: wp.array | float, + joint_ids: Sequence[int] | None = None, + env_ids: Sequence[int] | None = None, + joint_mask: wp.array | None = None, + env_mask: wp.array | None = None, + ) -> None: + """Write joint armature into the simulation. + + The armature is directly added to the corresponding joint-space inertia. It helps improve the + simulation stability by reducing the joint velocities. + + Args: + armature: Joint armature. Shape is (len(env_ids), len(joint_ids)) or (num_instances, num_joints). + joint_ids: The joint indices to set the targets for. Defaults to None (all joints). + env_ids: The environment indices to set the targets for. Defaults to None (all environments). + joint_mask: The joint mask. Shape is (num_joints). + env_mask: The environment mask. Shape is (num_instances,). + """ + # Resolve indices into mask, convert from partial data to complete data, handles the conversion to warp. + if self._frontend == Frontend.TORCH: + armature, env_mask, joint_mask = self._torch_to_warp_dual_index(armature, self.num_instances, self.num_joints, env_ids, joint_ids, env_mask, joint_mask, dtype=wp.float32) + self._joint.write_joint_armature_to_sim(armature, joint_mask, env_mask) + + def write_joint_friction_coefficient_to_sim( + self, + joint_friction_coeff: wp.array | float, + joint_ids: Sequence[int] | None = None, + env_ids: Sequence[int] | None = None, + joint_mask: wp.array | None = None, + env_mask: wp.array | None = None, + ) -> None: + r"""Write joint friction coefficients into the simulation. + + The joint friction is a unitless quantity. It relates the magnitude of the spatial force transmitted + from the parent body to the child body to the maximal friction force that may be applied by the solver + to resist the joint motion. + + Mathematically, this means that: :math:`F_{resist} \leq \mu F_{spatial}`, where :math:`F_{resist}` + is the resisting force applied by the solver and :math:`F_{spatial}` is the spatial force + transmitted from the parent body to the child body. The simulated friction effect is therefore + similar to static and Coulomb friction. + + Args: + joint_friction_coeff: Joint friction coefficients. Shape is (len(env_ids), len(joint_ids)) or (num_instances, num_joints). + joint_ids: The joint indices to set the targets for. Defaults to None (all joints). + env_ids: The environment indices to set the targets for. Defaults to None (all environments). + joint_mask: The joint mask. Shape is (num_joints). + env_mask: The environment mask. Shape is (num_instances,). + """ + # Resolve indices into mask, convert from partial data to complete data, handles the conversion to warp. + if self._frontend == Frontend.TORCH: + joint_friction_coeff, env_mask, joint_mask = self._torch_to_warp_dual_index(joint_friction_coeff, self.num_instances, self.num_joints, env_ids, joint_ids, env_mask, joint_mask, dtype=wp.float32) + self._joint.write_joint_friction_coefficient_to_sim(joint_friction_coeff, joint_mask, env_mask) + + """ + Operations - Setters. + """ + + def set_external_force_and_torque( + self, + forces: wp.array, + torques: 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, + ) -> None: + """Set external force and torque to apply on the asset's bodies in their local frame. + + For many applications, we want to keep the applied external force on rigid bodies constant over a period of + time (for instance, during the policy control). This function allows us to store the external force and torque + into buffers which are then applied to the simulation at every step. + + .. caution:: + If the function is called with empty forces and torques, then this function disables the application + of external wrench to the simulation. + + .. code-block:: python + + # example of disabling external wrench + asset.set_external_force_and_torque(forces=wp.zeros(0, 3), torques=wp.zeros(0, 3)) + + .. note:: + This function does not apply the external wrench to the simulation. It only fills the buffers with + the desired values. To apply the external wrench, call the :meth:`write_data_to_sim` function + right before the simulation step. + + Args: + forces: External forces in bodies' local frame. Shape is (len(env_ids), len(body_ids), 3) or (num_instances, num_bodies, 3). + torques: External torques in bodies' local frame. Shape is (len(env_ids), len(body_ids), 3) or (num_instances, num_bodies, 3). + body_ids: The body indices to set the targets for. Defaults to None (all bodies). + env_ids: The environment indices to set the targets for. Defaults to None (all environments). + body_mask: The body mask. Shape is (num_bodies). + env_mask: The environment mask. Shape is (num_instances,). + """ + # Resolve indices into mask, convert from partial data to complete data, handles the conversion to warp. + if self._frontend == Frontend.TORCH: + if forces is not None: + forces, env_mask, body_mask = self._torch_to_warp_dual_index(forces, self.num_instances, self.num_bodies, env_ids, body_ids, env_mask, body_mask, dtype=wp.float32) + if torques is not None: + torques, env_mask, body_mask = self._torch_to_warp_dual_index(torques, self.num_instances, self.num_bodies, env_ids, body_ids, env_mask, body_mask, dtype=wp.float32) + self._body.set_external_force_and_torque(forces, torques, body_mask, env_mask) + + def set_joint_position_target( + self, + target: wp.array, + joint_ids: Sequence[int] | None = None, + env_ids: Sequence[int] | None = None, + joint_mask: wp.array | None = None, + env_mask: wp.array | None = None, + ) -> None: + """Set joint position targets into internal buffers. + + This function does not apply the joint targets to the simulation. It only fills the buffers with + the desired values. To apply the joint targets, call the :meth:`write_data_to_sim` function. + + Args: + target: Joint position targets. Shape is (len(env_ids), len(joint_ids)) or (num_instances, num_joints). + joint_ids: The joint indices to set the targets for. Defaults to None (all joints). + env_ids: The environment indices to set the targets for. Defaults to None (all environments). + joint_mask: The joint mask. Shape is (num_joints). + env_mask: The environment mask. Shape is (num_instances,). + """ + # Resolve indices into mask, convert from partial data to complete data, handles the conversion to warp. + if self._frontend == Frontend.TORCH: + target, env_mask, joint_mask = self._torch_to_warp_dual_index(target, self.num_instances, self.num_joints, env_ids, joint_ids, env_mask, joint_mask, dtype=wp.float32) + self._joint.set_joint_position_target(target, joint_mask, env_mask) + + def set_joint_velocity_target( + self, + target: wp.array, + joint_ids: Sequence[int] | None = None, + env_ids: Sequence[int] | None = None, + joint_mask: wp.array | None = None, + env_mask: wp.array | None = None, + ) -> None: + """Set joint velocity targets into internal buffers. + + This function does not apply the joint targets to the simulation. It only fills the buffers with + the desired values. To apply the joint targets, call the :meth:`write_data_to_sim` function. + + Args: + target: Joint velocity targets. Shape is (len(env_ids), len(joint_ids)) or (num_instances, num_joints). + joint_ids: The joint indices to set the targets for. Defaults to None (all joints). + env_ids: The environment indices to set the targets for. Defaults to None (all environments). + joint_mask: The joint mask. Shape is (num_joints). + env_mask: The environment mask. Shape is (num_instances,). + """ + # Resolve indices into mask, convert from partial data to complete data, handles the conversion to warp. + if self._frontend == Frontend.TORCH: + target, env_mask, joint_mask = self._torch_to_warp_dual_index(target, self.num_instances, self.num_joints, env_ids, joint_ids, env_mask, joint_mask, dtype=wp.float32) + self._joint.set_joint_velocity_target(target, joint_mask, env_mask) + + def set_joint_effort_target( + self, + target: wp.array, + joint_ids: Sequence[int] | None = None, + env_ids: Sequence[int] | None = None, + joint_mask: wp.array | None = None, + env_mask: wp.array | None = None, + ) -> None: + """Set joint efforts into internal buffers. + + This function does not apply the joint targets to the simulation. It only fills the buffers with + the desired values. To apply the joint targets, call the :meth:`write_data_to_sim` function. + + Args: + target: Joint effort targets. Shape is (len(env_ids), len(joint_ids)) or (num_instances, num_joints). + joint_ids: The joint indices to set the targets for. Defaults to None (all joints). + env_ids: The environment indices to set the targets for. Defaults to None (all environments). + joint_mask: The joint mask. Shape is (num_joints). + env_mask: The environment mask. Shape is (num_instances,). + """ + # Resolve indices into mask, convert from partial data to complete data, handles the conversion to warp. + if self._frontend == Frontend.TORCH: + target, env_mask, joint_mask = self._torch_to_warp_dual_index(target, self.num_instances, self.num_joints, env_ids, joint_ids, env_mask, joint_mask, dtype=wp.float32) + self._joint.set_joint_effort_target(target, joint_mask, env_mask) + + """ + Operations - Tendons. + """ + + def set_fixed_tendon_stiffness( + self, + stiffness: wp.array, + fixed_tendon_ids: Sequence[int] | None = None, + env_ids: Sequence[int] | None = None, + fixed_tendon_mask: wp.array | None = None, + env_mask: wp.array | None = None, + ) -> None: + """Set fixed tendon stiffness into internal buffers. + + This function does not apply the tendon stiffness to the simulation. It only fills the buffers with + the desired values. To apply the tendon stiffness, call the :meth:`write_fixed_tendon_properties_to_sim` function. + + Args: + stiffness: Fixed tendon stiffness. Shape is (len(env_ids), len(fixed_tendon_ids)) or (num_instances, num_fixed_tendons). + fixed_tendon_ids: The fixed tendon indices to set the stiffness for. Defaults to None (all fixed tendons). + env_ids: The environment indices to set the stiffness for. Defaults to None (all environments). + fixed_tendon_mask: The fixed tendon mask. Shape is (num_fixed_tendons). + env_mask: The environment mask. Shape is (num_instances,). + """ + raise NotImplementedError("Fixed tendon stiffness is not supported in Newton.") + + def set_fixed_tendon_damping( + self, + damping: wp.array, + fixed_tendon_ids: Sequence[int] | None = None, + env_ids: Sequence[int] | None = None, + fixed_tendon_mask: wp.array | None = None, + env_mask: wp.array | None = None, + ) -> None: + """Set fixed tendon damping into internal buffers. + + This function does not apply the tendon damping to the simulation. It only fills the buffers with + the desired values. To apply the tendon damping, call the :meth:`write_fixed_tendon_properties_to_sim` function. + + Args: + damping: Fixed tendon damping. Shape is (len(env_ids), len(fixed_tendon_ids)) or (num_instances, num_fixed_tendons). + fixed_tendon_ids: The fixed tendon indices to set the damping for. Defaults to None (all fixed tendons). + env_ids: The environment indices to set the damping for. Defaults to None (all environments). + fixed_tendon_mask: The fixed tendon mask. Shape is (num_fixed_tendons). + env_mask: The environment mask. Shape is (num_instances,). + """ + raise NotImplementedError("Fixed tendon damping is not supported in Newton.") + + def set_fixed_tendon_limit_stiffness( + self, + limit_stiffness: wp.array, + fixed_tendon_ids: Sequence[int] | None = None, + env_ids: Sequence[int] | None = None, + fixed_tendon_mask: wp.array | None = None, + env_mask: wp.array | None = None, + ) -> None: + """Set fixed tendon limit stiffness efforts into internal buffers. + + This function does not apply the tendon limit stiffness to the simulation. It only fills the buffers with + the desired values. To apply the tendon limit stiffness, call the :meth:`write_fixed_tendon_properties_to_sim` function. + + Args: + limit_stiffness: Fixed tendon limit stiffness. Shape is (len(env_ids), len(fixed_tendon_ids)) or (num_instances, num_fixed_tendons). + fixed_tendon_ids: The fixed tendon indices to set the limit stiffness for. Defaults to None (all fixed tendons). + env_ids: The environment indices to set the limit stiffness for. Defaults to None (all environments). + fixed_tendon_mask: The fixed tendon mask. Shape is (num_fixed_tendons). + env_mask: The environment mask. Shape is (num_instances,). + """ + raise NotImplementedError("Fixed tendon limit stiffness is not supported in Newton.") + + def set_fixed_tendon_position_limit( + self, + limit: wp.array, + fixed_tendon_ids: Sequence[int] | None = None, + env_ids: Sequence[int] | None = None, + fixed_tendon_mask: wp.array | None = None, + env_mask: wp.array | None = None, + ) -> None: + """Set fixed tendon limit efforts into internal buffers. + + This function does not apply the tendon limit to the simulation. It only fills the buffers with + the desired values. To apply the tendon limit, call the :meth:`write_fixed_tendon_properties_to_sim` function. + + Args: + limit: Fixed tendon limit. Shape is (len(env_ids), len(fixed_tendon_ids)) or (num_instances, num_fixed_tendons). + fixed_tendon_ids: The fixed tendon indices to set the limit for. Defaults to None (all fixed tendons). + env_ids: The environment indices to set the limit for. Defaults to None (all environments). + fixed_tendon_mask: The fixed tendon mask. Shape is (num_fixed_tendons). + env_mask: The environment mask. Shape is (num_instances,). + """ + raise NotImplementedError("Fixed tendon position limit is not supported in Newton.") + + def set_fixed_tendon_rest_length( + self, + rest_length: wp.array, + fixed_tendon_ids: Sequence[int] | None = None, + env_ids: Sequence[int] | None = None, + fixed_tendon_mask: wp.array | None = None, + env_mask: wp.array | None = None, + ) -> None: + """Set fixed tendon rest length efforts into internal buffers. + + This function does not apply the tendon rest length to the simulation. It only fills the buffers with + the desired values. To apply the tendon rest length, call the :meth:`write_fixed_tendon_properties_to_sim` function. + + Args: + rest_length: Fixed tendon rest length. Shape is (len(env_ids), len(fixed_tendon_ids)) or (num_instances, num_fixed_tendons). + fixed_tendon_ids: The fixed tendon indices to set the rest length for. Defaults to None (all fixed tendons). + env_ids: The environment indices to set the rest length for. Defaults to None (all environments). + fixed_tendon_mask: The fixed tendon mask. Shape is (num_fixed_tendons). + env_mask: The environment mask. Shape is (num_instances,). + """ + raise NotImplementedError("Fixed tendon rest length is not supported in Newton.") + + def set_fixed_tendon_offset( + self, + offset: wp.array, + fixed_tendon_ids: Sequence[int] | None = None, + env_ids: Sequence[int] | None = None, + fixed_tendon_mask: wp.array | None = None, + env_mask: wp.array | None = None, + ) -> None: + """Set fixed tendon offset efforts into internal buffers. + + This function does not apply the tendon offset to the simulation. It only fills the buffers with + the desired values. To apply the tendon offset, call the :meth:`write_fixed_tendon_properties_to_sim` function. + + Args: + offset: Fixed tendon offset. Shape is (len(env_ids), len(fixed_tendon_ids)) or (num_instances, num_fixed_tendons). + fixed_tendon_ids: The tendon indices to set the offset for. Defaults to None (all fixed tendons). + env_ids: The environment indices to set the offset for. Defaults to None (all environments). + fixed_tendon_mask: The fixed tendon mask. Shape is (num_fixed_tendons). + env_mask: The environment mask. Shape is (num_instances,). + """ + raise NotImplementedError("Fixed tendon offset is not supported in Newton.") + + def write_fixed_tendon_properties_to_sim( + self, + fixed_tendon_ids: Sequence[int] | None = None, + env_ids: Sequence[int] | None = None, + fixed_tendon_mask: wp.array | None = None, + env_mask: wp.array | None = None, + ) -> None: + """Write fixed tendon properties into the simulation. + + Args: + fixed_tendon_ids: The fixed tendon indices to set the properties for. Defaults to None (all fixed tendons). + env_ids: The environment indices to set the properties for. Defaults to None (all environments). + fixed_tendon_mask: The fixed tendon mask. Shape is (num_fixed_tendons,). + env_mask: The environment mask. Shape is (num_instances,). + """ + raise NotImplementedError("Fixed tendon properties are not supported in Newton.") + + def set_spatial_tendon_stiffness( + self, + stiffness: wp.array, + spatial_tendon_ids: Sequence[int] | None = None, + env_ids: Sequence[int] | None = None, + spatial_tendon_mask: wp.array | None = None, + env_mask: wp.array | None = None, + ) -> None: + """Set spatial tendon stiffness into internal buffers. + + This function does not apply the tendon stiffness to the simulation. It only fills the buffers with + the desired values. To apply the tendon stiffness, call the :meth:`write_spatial_tendon_properties_to_sim` function. + + Args: + stiffness: Spatial tendon stiffness. Shape is (len(env_ids), len(spatial_tendon_ids)) or (num_instances, num_spatial_tendons). + spatial_tendon_ids: The spatial tendon indices to set the stiffness for. Defaults to None (all spatial tendons). + env_ids: The environment indices to set the stiffness for. Defaults to None (all environments). + spatial_tendon_mask: The spatial tendon mask. Shape is (num_spatial_tendons,). + env_mask: The environment mask. Shape is (num_instances,). + """ + raise NotImplementedError("Spatial tendon stiffness is not supported in Newton.") + + def set_spatial_tendon_damping( + self, + damping: wp.array, + spatial_tendon_ids: Sequence[int] | None = None, + env_ids: Sequence[int] | None = None, + spatial_tendon_mask: wp.array | None = None, + env_mask: wp.array | None = None, + ) -> None: + """Set spatial tendon damping into internal buffers. + + This function does not apply the tendon damping to the simulation. It only fills the buffers with + the desired values. To apply the tendon damping, call the :meth:`write_spatial_tendon_properties_to_sim` function. + + Args: + damping: Spatial tendon damping. Shape is (len(env_ids), len(spatial_tendon_ids)) or (num_instances, num_spatial_tendons). + spatial_tendon_ids: The spatial tendon indices to set the damping for. Defaults to None (all spatial tendons). + env_ids: The environment indices to set the damping for. Defaults to None (all environments). + spatial_tendon_mask: The spatial tendon mask. Shape is (num_spatial_tendons,). + env_mask: The environment mask. Shape is (num_instances,). + """ + raise NotImplementedError("Spatial tendon damping is not supported in Newton.") + + def set_spatial_tendon_limit_stiffness( + self, + limit_stiffness: wp.array, + spatial_tendon_ids: Sequence[int] | None = None, + env_ids: Sequence[int] | None = None, + spatial_tendon_mask: wp.array | None = None, + env_mask: wp.array | None = None, + ) -> None: + """Set spatial tendon limit stiffness into internal buffers. + + This function does not apply the tendon limit stiffness to the simulation. It only fills the buffers with + the desired values. To apply the tendon limit stiffness, call the :meth:`write_spatial_tendon_properties_to_sim` function. + + Args: + limit_stiffness: Spatial tendon limit stiffness. Shape is (len(env_ids), len(spatial_tendon_ids)) or (num_instances, num_spatial_tendons). + spatial_tendon_ids: The spatial tendon indices to set the limit stiffness for. Defaults to None (all spatial tendons). + env_ids: The environment indices to set the limit stiffness for. Defaults to None (all environments). + spatial_tendon_mask: The spatial tendon mask. Shape is (num_spatial_tendons,). + env_mask: The environment mask. Shape is (num_instances,). + """ + raise NotImplementedError("Spatial tendon limit stiffness is not supported in Newton.") + + def set_spatial_tendon_offset( + self, + offset: wp.array, + spatial_tendon_ids: Sequence[int] | None = None, + env_ids: Sequence[int] | None = None, + spatial_tendon_mask: wp.array | None = None, + env_mask: wp.array | None = None, + ) -> None: + """Set spatial tendon offset efforts into internal buffers. + + This function does not apply the tendon offset to the simulation. It only fills the buffers with + the desired values. To apply the tendon offset, call the :meth:`write_spatial_tendon_properties_to_sim` function. + + Args: + offset: Spatial tendon offset. Shape is (len(env_ids), len(spatial_tendon_ids)) or (num_instances, num_spatial_tendons). + spatial_tendon_ids: The spatial tendon indices to set the offset for. Defaults to None (all spatial tendons). + env_ids: The environment indices to set the offset for. Defaults to None (all environments). + spatial_tendon_mask: The spatial tendon mask. Shape is (num_spatial_tendons,). + env_mask: The environment mask. Shape is (num_instances,). + """ + raise NotImplementedError("Spatial tendon offset is not supported in Newton.") + + def write_spatial_tendon_properties_to_sim( + self, + spatial_tendon_ids: Sequence[int] | None = None, + env_ids: Sequence[int] | None = None, + spatial_tendon_mask: wp.array | None = None, + env_mask: wp.array | None = None, + ) -> None: + """Write spatial tendon properties into the simulation. + + Args: + spatial_tendon_ids: The spatial tendon indices to set the properties for. Defaults to None (all spatial tendons). + env_ids: The environment indices to set the properties for. Defaults to None (all environments). + spatial_tendon_mask: The spatial tendon mask. Shape is (num_spatial_tendons,). + env_mask: The environment mask. Shape is (num_instances,). + """ + raise NotImplementedError("Spatial tendon properties are not supported in Newton.") + + """ + Internal helper. + """ + + def _initialize_impl(self): + # obtain global simulation view + if self.cfg.articulation_root_prim_path is not None: + # The articulation root prim path is specified explicitly, so we can just use this. + root_prim_path_expr = self.cfg.prim_path + self.cfg.articulation_root_prim_path + else: + # No articulation root prim path was specified, so we need to search + # for it. We search for this in the first environment and then + # create a regex that matches all environments. + first_env_matching_prim = sim_utils.find_first_matching_prim(self.cfg.prim_path) + if first_env_matching_prim is None: + raise RuntimeError(f"Failed to find prim for expression: '{self.cfg.prim_path}'.") + first_env_matching_prim_path = first_env_matching_prim.GetPath().pathString + + # Find all articulation root prims in the first environment. + first_env_root_prims = sim_utils.get_all_matching_child_prims( + first_env_matching_prim_path, + predicate=lambda prim: prim.HasAPI(UsdPhysics.ArticulationRootAPI), + ) + if len(first_env_root_prims) == 0: + raise RuntimeError( + f"Failed to find an articulation when resolving '{first_env_matching_prim_path}'." + " Please ensure that the prim has 'USD ArticulationRootAPI' applied." + ) + if len(first_env_root_prims) > 1: + raise RuntimeError( + f"Failed to find a single articulation when resolving '{first_env_matching_prim_path}'." + f" Found multiple '{first_env_root_prims}' under '{first_env_matching_prim_path}'." + " Please ensure that there is only one articulation in the prim path tree." + ) + + # resolve articulation root prim back into regex expression + first_env_root_prim_path = first_env_root_prims[0].GetPath().pathString + root_prim_path_relative_to_prim_path = first_env_root_prim_path[len(first_env_matching_prim_path) :] + root_prim_path_expr = self.cfg.prim_path + root_prim_path_relative_to_prim_path + + prim_path = root_prim_path_expr.replace(".*", "*") + + self._root_view = NewtonArticulationView( + NewtonManager.get_model(), prim_path, verbose=True, exclude_joint_types=[JointType.FREE, JointType.FIXED] + ) + + # container for data access + self._data = ArticulationData(self._root_view, self.device, frontend=self._frontend) + + # create backend setters for the data containers + self._joint = Joint(self._root_view, self._data.joint_data, self.cfg.soft_joint_pos_limit_factor, self.device) + self._body = Body(self._root_view, self._data.body_data, self.device) + self._root = Root(self._root_view, self._data.root_data, self.device) + + # process configuration + self._create_buffers() + self._process_cfg() + self._process_actuators_cfg() + self._process_tendons() + # validate configuration + self._validate_cfg() + # update the robot data + self.update(0.0) + # log joint information + self._log_articulation_info() + + # Offsets the spawned pose by the default root pose prior to initializing the solver. This ensures that the + # solver is initialized at the correct pose, avoiding potential miscalculations in the maximum number of + # constraints or contact required to run the simulation. + # TODO: Do this is warp directly? + generated_pose = wp.to_torch(self._data.root_data.default_root_pose).clone() + generated_pose[:, :2] += wp.to_torch(self._root_view.get_root_transforms(NewtonManager.get_model()))[ + :, :2 + ] + self._root_view.set_root_transforms(NewtonManager.get_state_0(), generated_pose) + self._root_view.set_root_transforms(NewtonManager.get_model(), generated_pose) + + + def _process_cfg(self): + """Post processing of configuration parameters.""" + # default pose with quaternion given as (w, x, y, z) --> (x, y, z, w) + default_root_pose = tuple(self.cfg.init_state.pos) + ( + self.cfg.init_state.rot[1], + self.cfg.init_state.rot[2], + self.cfg.init_state.rot[3], + self.cfg.init_state.rot[0], + ) + # update the default root pose + wp.launch( + update_array_with_value, + dim=(self.num_instances,), + inputs=[ + wp.transformf(*default_root_pose), + self._data.root_data.default_root_pose, + ], + ) + # default velocity + default_root_velocity = tuple(self.cfg.init_state.lin_vel) + tuple(self.cfg.init_state.ang_vel) + wp.launch( + update_array_with_value, + dim=(self.num_instances,), + inputs=[ + wp.spatial_vectorf(*default_root_velocity), + self._data.root_data.default_root_vel, + ], + ) + # -- joint pos + # joint pos + indices_list, _, values_list = string_utils.resolve_matching_names_values( + self.cfg.init_state.joint_pos, self.joint_names + ) + # Compute the mask once and use it for all joint operations + self._joint._JOINT_MASK.fill_(False) + wp.launch( + generate_mask_from_ids, + dim=(self.num_joints,), + inputs=[ + self._joint._JOINT_MASK, + wp.array(indices_list, dtype=wp.int32, device=self.device), + ], + ) + tmp_joint_data = wp.zeros((self.num_joints,), dtype=wp.float32, device=self.device) + wp.launch( + populate_empty_array, + dim=(self.num_joints,), + inputs=[ + wp.array(values_list, dtype=wp.float32, device=self.device), + tmp_joint_data, + wp.array(indices_list, dtype=wp.int32, device=self.device), + ], + ) + wp.launch( + update_batched_array_with_array_masked, + dim=(self.num_instances, self.num_joints), + inputs=[ + tmp_joint_data, + self._data.joint_data.default_joint_pos, + self._joint._ALL_ENV_MASK, + self._joint._JOINT_MASK, + ], + ) + # joint vel + indices_list, _, values_list = string_utils.resolve_matching_names_values( + self.cfg.init_state.joint_vel, self.joint_names + ) + wp.launch( + populate_empty_array, + dim=(self.num_joints,), + inputs=[ + wp.array(values_list, dtype=wp.float32, device=self.device), + tmp_joint_data, + wp.array(indices_list, dtype=wp.int32, device=self.device), + ], + ) + wp.launch( + update_batched_array_with_array_masked, + dim=(self.num_instances, self.num_joints), + inputs=[ + tmp_joint_data, + self._data.joint_data.default_joint_vel, + self._joint._ALL_ENV_MASK, + self._joint._JOINT_MASK, + ], + ) + + """ + Internal simulation callbacks. + """ + + def _invalidate_initialize_callback(self, event): + """Invalidates the scene elements.""" + # call parent + super()._invalidate_initialize_callback(event) + + """ + Internal helpers -- Actuators. + """ + + def _process_actuators_cfg(self): + """Process and apply articulation joint properties.""" + # create actuators + self.actuators = dict() + # flag for implicit actuators + # if this is false, we by-pass certain checks when doing actuator-related operations + self._has_implicit_actuators = False + + # iterate over all actuator configurations + for actuator_name, actuator_cfg in self.cfg.actuators.items(): + # type annotation for type checkers + actuator_cfg: ActuatorBaseCfg + # create actuator group + joint_mask, joint_names, joint_indices = self.find_joints(actuator_cfg.joint_names_expr) + # check if any joints are found + if len(joint_names) == 0: + raise ValueError( + f"No joints found for actuator group: {actuator_name} with joint name expression:" + f" {actuator_cfg.joint_names_expr}." + ) + # create actuator collection + # note: for efficiency avoid indexing when over all indices + actuator: ActuatorBase = actuator_cfg.class_type( + cfg=actuator_cfg, + joint_names=joint_names, + joint_mask=joint_mask, + env_mask=self._joint._ALL_ENV_MASK, + joint_data=self._data.joint_data, + device=self.device, + ) + # log information on actuator groups + model_type = "implicit" if actuator.is_implicit_model else "explicit" + logger.info( + f"Actuator collection: {actuator_name} with model '{actuator_cfg.class_type.__name__}'" + f" (type: {model_type}) and joint names: {joint_names} [{joint_mask}]." + ) + # store actuator group + self.actuators[actuator_name] = actuator + # set the passed gains and limits into the simulation + # TODO: write out all joint parameters from simulation + if isinstance(actuator, ImplicitActuator): + self._has_implicit_actuators = True + # the gains and limits are set into the simulation since actuator model is implicit + self._joint.write_joint_stiffness_to_sim(self.data.joint_stiffness, joint_mask=actuator._joint_mask) + self._joint.write_joint_damping_to_sim(self.data.joint_damping, joint_mask=actuator._joint_mask) + # Sets the control mode for the implicit actuators + self._joint.write_joint_control_mode_to_sim(self.data.joint_control_mode, joint_mask=actuator._joint_mask) + + # When using implicit actuators, we bind the commands sent from the user to the simulation. + # We only run the actuator model to compute the estimated joint efforts. + self.data._joint_data._joint_target = self.data._joint_data.joint_target_sim + self.data._joint_data._joint_effort_target = self.data._joint_data.joint_effort_sim + else: + # the gains and limits are processed by the actuator model + # we set gains to zero, and torque limit to a high value in simulation to avoid any interference + self._joint.write_joint_stiffness_to_sim(0.0, joint_mask=actuator._joint_mask) + self._joint.write_joint_damping_to_sim(0.0, joint_mask=actuator._joint_mask) + # Set the control mode to None when using explicit actuators + self._joint.write_joint_control_mode_to_sim(0, joint_mask=actuator._joint_mask) + # Bind the applied effort to the simulation effort + self.data._joint_data._applied_effort = self.data.joint_effort_sim + + # perform some sanity checks to ensure actuators are prepared correctly + total_act_joints = sum(actuator.num_joints for actuator in self.actuators.values()) + if total_act_joints != (self.num_joints - self.num_fixed_tendons): + warnings.warn( + "Not all actuators are configured! Total number of actuated joints not equal to number of" + f" joints available: {total_act_joints} != {self.num_joints - self.num_fixed_tendons}.", + UserWarning, + ) + + def _process_tendons(self): + """Process fixed and spatialtendons.""" + # create a list to store the fixed tendon names + self._fixed_tendon_names = list() + self._spatial_tendon_names = list() + # parse fixed tendons properties if they exist + if self.num_fixed_tendons > 0: + raise NotImplementedError("Tendons are not implemented yet") + + def _apply_actuator_model(self): + """Processes joint commands for the articulation by forwarding them to the actuators. + + The actions are first processed using actuator models. Depending on the robot configuration, + the actuator models compute the joint level simulation commands and sets them into the PhysX buffers. + """ + # process actions per group + for actuator in self.actuators.values(): + actuator.compute() + # TODO: find a cleaner way to handle gear ratio. Only needed for variable gear ratio actuators. + # if hasattr(actuator, "gear_ratio"): + # self._data.gear_ratio[:, actuator.joint_indices] = actuator.gear_ratio + + """ + Internal helpers -- Debugging. + """ + + def _validate_cfg(self): + """Validate the configuration after processing. + + Note: + This function should be called only after the configuration has been processed and the buffers have been + created. Otherwise, some settings that are altered during processing may not be validated. + For instance, the actuator models may change the joint max velocity limits. + """ + # check that the default values are within the limits + joint_pos_limits = torch.stack( + ( + wp.to_torch(self._root_view.get_attribute("joint_limit_lower", NewtonManager.get_model())), + wp.to_torch(self._root_view.get_attribute("joint_limit_upper", NewtonManager.get_model())), + ), + dim=2, + )[0].to(self.device) + out_of_range = wp.to_torch(self._data.joint_data.default_joint_pos)[0] < joint_pos_limits[:, 0] + out_of_range |= wp.to_torch(self._data.joint_data.default_joint_pos)[0] > joint_pos_limits[:, 1] + violated_indices = torch.nonzero(out_of_range, as_tuple=False).squeeze(-1) + # throw error if any of the default joint positions are out of the limits + if len(violated_indices) > 0: + # prepare message for violated joints + msg = "The following joints have default positions out of the limits: \n" + for idx in violated_indices: + joint_name = self.data.joint_names[idx] + joint_limit = joint_pos_limits[idx] + joint_pos = self.data.default_joint_pos[0, idx] + # add to message + msg += f"\t- '{joint_name}': {joint_pos:.3f} not in [{joint_limit[0]:.3f}, {joint_limit[1]:.3f}]\n" + raise ValueError(msg) + + def _log_articulation_info(self): + """Log information about the articulation. + + Note: We purposefully read the values from the simulator to ensure that the values are configured as expected. + """ + # TODO: read out all joint parameters from simulation + # read out all joint parameters from simulation + # -- gains + stiffnesses = torch.zeros([self.num_joints], dtype=torch.float32, device=self.device).tolist() + dampings = torch.zeros([self.num_joints], dtype=torch.float32, device=self.device).tolist() + # -- properties + armatures = torch.zeros([self.num_joints], dtype=torch.float32, device=self.device).tolist() + frictions = torch.zeros([self.num_joints], dtype=torch.float32, device=self.device).tolist() + # -- limits + position_limits = torch.zeros([self.num_joints, 2], dtype=torch.float32, device=self.device).tolist() + velocity_limits = torch.zeros([self.num_joints], dtype=torch.float32, device=self.device).tolist() + effort_limits = torch.zeros([self.num_joints], dtype=torch.float32, device=self.device).tolist() + # create table for term information + joint_table = PrettyTable() + joint_table.title = f"Simulation Joint Information (Prim path: {self.cfg.prim_path})" + joint_table.field_names = [ + "Index", + "Name", + "Stiffness", + "Damping", + "Armature", + "Friction", + "Position Limits", + "Velocity Limits", + "Effort Limits", + ] + joint_table.float_format = ".3" + joint_table.custom_format["Position Limits"] = lambda f, v: f"[{v[0]:.3f}, {v[1]:.3f}]" + # set alignment of table columns + joint_table.align["Name"] = "l" + # add info on each term + for index, name in enumerate(self.joint_names): + joint_table.add_row([ + index, + name, + stiffnesses[index], + dampings[index], + armatures[index], + frictions[index], + position_limits[index], + velocity_limits[index], + effort_limits[index], + ]) + # convert table to string + logger.info(f"Simulation parameters for joints in {self.cfg.prim_path}:\n" + joint_table.get_string()) + + # read out all fixed tendon parameters from simulation + if self.num_fixed_tendons > 0: + raise NotImplementedError("Tendons are not implemented yet") + + if self.num_spatial_tendons > 0: + raise NotImplementedError("Tendons are not implemented yet") + + def _create_buffers(self, *args, **kwargs): + self._ALL_INDICES = torch.arange(self.num_instances, dtype=torch.long, device=self.device) + + def set_fixed_tendon_limit(self, *args, **kwargs): + pass + def write_joint_dynamic_friction_coefficient_to_sim(self, *args, **kwargs): + pass + def write_joint_friction_to_sim(self, *args, **kwargs): + pass + def write_joint_limits_to_sim(self, *args, **kwargs): + pass \ No newline at end of file diff --git a/source/isaaclab/isaaclab/newton/assets/articulation/articulation_data.py b/source/isaaclab/isaaclab/newton/assets/articulation/articulation_data.py new file mode 100644 index 00000000000..fb660eb6762 --- /dev/null +++ b/source/isaaclab/isaaclab/newton/assets/articulation/articulation_data.py @@ -0,0 +1,971 @@ +# Copyright (c) 2022-2025, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause +from __future__ import annotations + +import weakref + +import warp as wp +from typing import TYPE_CHECKING + +from isaaclab.assets.articulation.base_articulation_data import BaseArticulationData +from isaaclab.newton.assets.core.root_properties.root_data import RootData +from isaaclab.newton.assets.core.body_properties.body_data import BodyData +from isaaclab.newton.assets.core.joint_properties.joint_data import JointData +from isaaclab.newton.assets.core.kernels import vec13f, combine_pose_and_velocity_to_state +from isaaclab.utils.helpers import deprecated + +from isaaclab.newton import Frontend + +if TYPE_CHECKING: + import torch + +class ArticulationData(BaseArticulationData): + def __init__(self, root_newton_view, device: str, frontend: Frontend = Frontend.TORCH): + # Set the parameters + self.device = device + # Set the root articulation view + # note: this is stored as a weak reference to avoid circular references between the asset class + # and the data container. This is important to avoid memory leaks. + self._root_view = weakref.proxy(root_newton_view) + + # Initialize the data containers + self._root_data = RootData(root_newton_view, device) + self._body_data = BodyData(root_newton_view, device) + self._joint_data = JointData(root_newton_view, device) + + self.frontend = frontend + + def maybe_to_torch(self, data: wp.array) -> wp.array | torch.Tensor: + return wp.to_torch(data) if self.frontend == Frontend.TORCH else data + + @property + def joint_data(self) -> JointData: + return self._joint_data + + @property + def body_data(self) -> BodyData: + return self._body_data + + @property + def root_data(self) -> RootData: + return self._root_data + + def update(self, dt: float): + self._root_data.update(dt) + self._body_data.update(dt) + self._joint_data.update(dt) + + ## + # Names. + ## + + body_names: list[str] = None + """Body names in the order parsed by the simulation view.""" + + joint_names: list[str] = None + """Joint names in the order parsed by the simulation view.""" + + fixed_tendon_names: list[str] = None + """Fixed tendon names in the order parsed by the simulation view.""" + + spatial_tendon_names: list[str] = None + """Spatial tendon names in the order parsed by the simulation view.""" + + ## + # Defaults. + ## + + @property + def default_root_pose(self) -> wp.array | torch.Tensor: + """Default root pose ``[pos, quat]`` in the local environment frame. Shape is (num_instances, 7). + + The position and quaternion are of the articulation root's actor frame. + + This quantity is configured through the :attr:`isaaclab.assets.ArticulationCfg.init_state` parameter. + """ + return self.maybe_to_torch(self._root_data.default_root_pose) + + @property + def default_root_vel(self) -> wp.array | torch.Tensor: + """Default root velocity ``[lin_vel, ang_vel]`` in the local environment frame. Shape is (num_instances, 6). + + The linear and angular velocities are of the articulation root's center of mass frame. + + This quantity is configured through the :attr:`isaaclab.assets.ArticulationCfg.init_state` parameter. + """ + return self.maybe_to_torch(self._root_data.default_root_vel) + + @property + def default_joint_pos(self) -> wp.array | torch.Tensor: + """Default joint positions of all joints. Shape is (num_instances, num_joints). + + This quantity is configured through the :attr:`isaaclab.assets.ArticulationCfg.init_state` parameter. + """ + return self.maybe_to_torch(self._joint_data.default_joint_pos) + + @property + def default_joint_vel(self) -> wp.array | torch.Tensor: + """Default joint velocities of all joints. Shape is (num_instances, num_joints). + + This quantity is configured through the :attr:`isaaclab.assets.ArticulationCfg.init_state` parameter. + """ + return self.maybe_to_torch(self._joint_data.default_joint_vel) + + ## + # Joint commands -- From the user to the actuator model. + ## + + @property + def joint_target(self) -> wp.array | torch.Tensor: + """Joint position targets commanded by the user. Shape is (num_instances, num_joints). + + For an implicit actuator model, the targets are directly set into the simulation. + For an explicit actuator model, the targets are used to compute the joint efforts (see :attr:`applied_torque`), + which are then set into the simulation. + """ + return self.maybe_to_torch(self._joint_data.joint_target) + + @property + def joint_effort_target(self) -> wp.array | torch.Tensor: + """Joint effort targets commanded by the user. Shape is (num_instances, num_joints). + + For an implicit actuator model, the targets are directly set into the simulation. + For an explicit actuator model, the targets are used to compute the joint efforts (see :attr:`applied_torque`), + which are then set into the simulation. + """ + return self.maybe_to_torch(self._joint_data.joint_effort_target) + + ## + # Joint commands -- Explicit actuators. + ## + + @property + def computed_effort(self) -> wp.array | torch.Tensor: + """Joint efforts computed from the actuator model (before clipping). Shape is (num_instances, num_joints). + + This quantity is the raw effort output from the actuator mode, before any clipping is applied. + It is exposed for users who want to inspect the computations inside the actuator model. + For instance, to penalize the learning agent for a difference between the computed and applied torques. + """ + return self.maybe_to_torch(self._joint_data.computed_effort) + + @property + def applied_effort(self) -> wp.array | torch.Tensor: + """Joint efforts applied from the actuator model (after clipping). Shape is (num_instances, num_joints). + + These efforts are set into the simulation, after clipping the :attr:`computed_effort` based on the + actuator model. + """ + return self.maybe_to_torch(self._joint_data.applied_effort) + + @property + def joint_stiffness(self) -> wp.array | torch.Tensor: + """Joint stiffness. Shape is (num_instances, num_joints).""" + return self.maybe_to_torch(self._joint_data.joint_stiffness) + + @property + def joint_damping(self) -> wp.array | torch.Tensor: + """Joint damping. Shape is (num_instances, num_joints).""" + return self.maybe_to_torch(self._joint_data.joint_damping) + + @property + def joint_control_mode(self) -> wp.array | torch.Tensor: + """Joint control mode. Shape is (num_instances, num_joints).""" + return self.maybe_to_torch(self._joint_data.joint_control_mode) + + ### + # Joint commands. (Directly binded to the simulation) + ### + + @property + def joint_target_sim(self) -> wp.array | torch.Tensor: + """Joint target. Shape is (num_instances, num_joints).""" + return self.maybe_to_torch(self._joint_data.joint_target_sim) + + @property + def joint_effort_sim(self) -> wp.array | torch.Tensor: + """Joint effort. Shape is (num_instances, num_joints).""" + return self.maybe_to_torch(self._joint_data.joint_effort_sim) + + ## + # Joint properties. (Directly binded to the simulation) + ## + + @property + def joint_control_mode_sim(self) -> wp.array | torch.Tensor: + """Joint control mode. Shape is (num_instances, num_joints). + + When using implicit actuator models Newton needs to know how the joints are controlled. + The control mode can be one of the following: + + * None: 0 + * Position control: 1 + * Velocity control: 2 + + This quantity is set by the :meth:`Articulation.write_joint_control_mode_to_sim` method. + """ + return self.maybe_to_torch(self._joint_data.joint_control_mode_sim) + + @property + def joint_stiffness_sim(self) -> wp.array | torch.Tensor: + """Joint stiffness provided to the simulation. Shape is (num_instances, num_joints). + + In the case of explicit actuators, the value for the corresponding joints is zero. + """ + return self.maybe_to_torch(self._joint_data.joint_stiffness_sim) + + @property + def joint_damping_sim(self) -> wp.array | torch.Tensor: + """Joint damping provided to the simulation. Shape is (num_instances, num_joints) + + In the case of explicit actuators, the value for the corresponding joints is zero. + """ + return self.maybe_to_torch(self._joint_data.joint_damping_sim) + + @property + def joint_armature(self) -> wp.array | torch.Tensor: + """Joint armature provided to the simulation. Shape is (num_instances, num_joints).""" + return self.maybe_to_torch(self._joint_data.joint_armature) + + @property + def joint_friction_coeff(self) -> wp.array | torch.Tensor: + """Joint friction coefficient provided to the simulation. Shape is (num_instances, num_joints).""" + return self.maybe_to_torch(self._joint_data.joint_friction_coeff) + + @property + def joint_pos_limits_lower(self) -> wp.array | torch.Tensor: + """Joint position limits lower provided to the simulation. Shape is (num_instances, num_joints).""" + return self.maybe_to_torch(self._joint_data.joint_pos_limits_lower) + + @property + def joint_pos_limits_upper(self) -> wp.array | torch.Tensor: + """Joint position limits upper provided to the simulation. Shape is (num_instances, num_joints).""" + return self.maybe_to_torch(self._joint_data.joint_pos_limits_upper) + + @property + def joint_vel_limits(self) -> wp.array | torch.Tensor: + """Joint maximum velocity provided to the simulation. Shape is (num_instances, num_joints).""" + return self.maybe_to_torch(self._joint_data.joint_vel_limits) + + @property + def joint_effort_limits(self) -> wp.array | torch.Tensor: + """Joint maximum effort provided to the simulation. Shape is (num_instances, num_joints).""" + return self.maybe_to_torch(self._joint_data.joint_effort_limits) + + ## + # Joint states + ## + + @property + def joint_pos(self) -> wp.array | torch.Tensor: + """Joint positions. Shape is (num_instances, num_joints).""" + return self.maybe_to_torch(self._joint_data.joint_pos) + + @property + def joint_vel(self) -> wp.array | torch.Tensor: + """Joint velocities. Shape is (num_instances, num_joints).""" + return self.maybe_to_torch(self._joint_data.joint_vel) + + ## + # Joint properties - Custom. + ## + + @property + def joint_dynamic_friction(self) -> wp.array | torch.Tensor: + """Joint dynamic friction. Shape is (num_instances, num_joints).""" + return self.maybe_to_torch(self._joint_data.joint_dynamic_friction) + + @property + def joint_viscous_friction(self) -> wp.array | torch.Tensor: + """Joint viscous friction. Shape is (num_instances, num_joints).""" + return self.maybe_to_torch(self._joint_data.joint_viscous_friction) + + @property + def soft_joint_pos_limits(self) -> wp.array | torch.Tensor: + r"""Soft joint positions limits for all joints. Shape is (num_instances, num_joints, 2). + + The limits are in the order :math:`[lower, upper]`.The soft joint position limits are computed as + a sub-region of the :attr:`joint_pos_limits` based on the + :attr:`~isaaclab.assets.ArticulationCfg.soft_joint_pos_limit_factor` parameter. + + Consider the joint position limits :math:`[lower, upper]` and the soft joint position limits + :math:`[soft_lower, soft_upper]`. The soft joint position limits are computed as: + + .. math:: + + soft\_lower = (lower + upper) / 2 - factor * (upper - lower) / 2 + soft\_upper = (lower + upper) / 2 + factor * (upper - lower) / 2 + + The soft joint position limits help specify a safety region around the joint limits. It isn't used by the + simulation, but is useful for learning agents to prevent the joint positions from violating the limits. + """ + return self.maybe_to_torch(self._joint_data.soft_joint_pos_limits) + + @property + def soft_joint_vel_limits(self) -> wp.array | torch.Tensor: + """Soft joint velocity limits for all joints. Shape is (num_instances, num_joints). + + These are obtained from the actuator model. It may differ from :attr:`joint_vel_limits` if the actuator model + has a variable velocity limit model. For instance, in a variable gear ratio actuator model. + """ + return self.maybe_to_torch(self._joint_data.soft_joint_vel_limits) + + @property + def gear_ratio(self) -> wp.array | torch.Tensor: #TODO: Mayank got some comments + """Gear ratio for relating motor torques to applied Joint torques. Shape is (num_instances, num_joints).""" + return self.maybe_to_torch(self._joint_data.gear_ratio) + + ## + # Root state properties. + ## + + @property + def root_mass(self) -> wp.array | torch.Tensor: + """Root mass ``wp.float32`` in the world frame. Shape is (num_instances,).""" + return self.maybe_to_torch(self._root_data.root_mass) + + @property + def root_inertia(self) -> wp.array | torch.Tensor: + """Root inertia ``wp.mat33`` in the world frame. Shape is (num_instances, 9).""" + return self.maybe_to_torch(self._root_data.root_inertia) + + @property + def root_link_pose_w(self) -> wp.array | torch.Tensor: + """Root link pose ``wp.transformf`` in simulation world frame. + + Shapes are (num_instances,). The pose is in the form of [pos, quat]. + This quantity is the pose of the articulation root's actor frame relative to the world. + The orientation is provided in (x, y, z, w) format. + """ + return self.maybe_to_torch(self._root_data.root_link_pose_w) + + @property + def root_link_vel_w(self) -> wp.array | torch.Tensor: + """Root link velocity ``wp.spatial_vectorf`` in simulation world frame. + + Shapes are (num_instances,). Velocities are in the form of [wx, wy, wz, vx, vy, vz]. + This quantity contains the linear and angular velocities of the articulation root's actor frame + relative to the world. + """ + + return self.maybe_to_torch(self._root_data.root_link_vel_w) + + @property + def root_com_pose_w(self) -> wp.array | torch.Tensor: + """Root center of mass pose ``wp.transformf`` in simulation world frame. + + Shapes are (num_instances,). The pose is in the form of [pos, quat]. + This quantity is the pose of the articulation root's center of mass frame relative to the world. + The orientation is provided in (x, y, z, w) format. + """ + + return self.maybe_to_torch(self._root_data.root_com_pose_w) + + @property + def root_com_vel_w(self) -> wp.array | torch.Tensor: + """Root center of mass velocity ``wp.spatial_vectorf`` in simulation world frame. + + Shapes are (num_instances,). The velocity is in the form of [wx, wy, wz, vx, vy, vz]. + This quantity contains the linear and angular velocities of the articulation root's center of mass frame + relative to the world. + """ + return self.maybe_to_torch(self._root_data.root_com_vel_w) + + @property + def root_state_w(self) -> wp.array | torch.Tensor: + """Root state ``[wp.transformf, wp.spatial_vectorf]`` in simulation world frame. + + Shapes are (num_instances,), (num_instances,). The pose is in the form of [pos, quat]. + The velocity is in the form of [wx, wy, wz, vx, vy, vz]. + The pose is of the articulation root's actor frame relative to the world. + The velocity is of the articulation root's center of mass frame. + """ + return self.maybe_to_torch(self._root_data.root_state_w) + + @property + def root_link_state_w(self) -> wp.array | torch.Tensor: + """Root link state ``[wp.transformf, wp.spatial_vectorf]`` in simulation world frame. + + Shapes are (num_instances,), (num_instances,). The pose is in the form of [pos, quat]. + The velocity is in the form of [wx, wy, wz, vx, vy, vz]. + The pose is of the articulation root's actor frame relative to the world. + The velocity is of the articulation root's actor frame. + """ + return self.maybe_to_torch(self._root_data.root_link_state_w) + + @property + def root_com_state_w(self) -> wp.array | torch.Tensor: + """Root center of mass state ``[wp.transformf, wp.spatial_vectorf]`` in simulation world frame. + + Shapes are (num_instances,), (num_instances,). The pose is in the form of [pos, quat]. + The velocity is in the form of [wx, wy, wz, vx, vy, vz]. + The pose is of the articulation root's center of mass frame relative to the world. + The velocity is of the articulation root's center of mass frame. + """ + return self.maybe_to_torch(self._root_data.root_com_state_w) + + @property + def root_com_pose_b(self) -> wp.array | torch.Tensor: + """Root center of mass pose ``wp.transformf`` in base frame. Shape is (num_instances,). + + This quantity is the pose of the root rigid body's center of mass frame with respect to the + rigid body's actor frame. + """ + return self.maybe_to_torch(self._root_data.root_com_pose_b) + + @property + def root_com_pos_b(self) -> wp.array | torch.Tensor: + """Root center of mass position ``wp.vec3f`` in base frame. Shape is (num_instances, 3). + + This quantity is the position of the root rigid body's center of mass frame with respect to the + rigid body's actor frame. + """ + return self.maybe_to_torch(self._root_data.root_com_pos_b) + + @property + def root_com_quat_b(self) -> wp.array | torch.Tensor: + """Root center of mass orientation (w, x, y, z) in base frame. Shape is (num_instances, 4). + + This quantity is the orientation of the root rigid body's center of mass frame with respect to the + rigid body's actor frame. + """ + return self.maybe_to_torch(self._root_data.root_com_quat_b) + + ## + # Body state properties. + ## + + @property + def body_mass(self) -> wp.array | torch.Tensor: + """Body mass ``wp.float32`` in the world frame. Shape is (num_instances, num_bodies).""" + return self.maybe_to_torch(self._body_data.body_mass) + + @property + def body_inertia(self) -> wp.array | torch.Tensor: + """Body inertia ``wp.mat33`` in the world frame. Shape is (num_instances, num_bodies, 9).""" + return self.maybe_to_torch(self._body_data.body_inertia) + + @property + def body_link_pose_w(self) -> wp.array | torch.Tensor: + """Body link pose ``wp.transformf`` in simulation world frame. + + Shapes are (num_instances, num_bodies,). The pose is in the form of [pos, quat]. + This quantity is the pose of the articulation links' actor frame relative to the world. + The orientation is provided in (x, y, z, w) format. + """ + return self.maybe_to_torch(self._body_data.body_link_pose_w) + + @property + def body_link_vel_w(self) -> wp.array | torch.Tensor: + """Body link velocity ``wp.spatial_vectorf`` in simulation world frame. + + Shapes are (num_instances, num_bodies,). Velocities are in the form of [wx, wy, wz, vx, vy, vz]. + This quantity contains the linear and angular velocities of the articulation links' actor frame + relative to the world. + """ + return self.maybe_to_torch(self._body_data.body_link_vel_w) + + @property + def body_com_pose_w(self) -> wp.array | torch.Tensor: + """Body center of mass pose ``wp.transformf`` in simulation world frame. + + Shapes are (num_instances, num_bodies,). The pose is in the form of [pos, quat]. + This quantity is the pose of the center of mass frame of the articulation links relative to the world. + The orientation is provided in (x, y, z, w) format. + """ + return self.maybe_to_torch(self._body_data.body_com_pose_w) + + @property + def body_com_vel_w(self) -> wp.array | torch.Tensor: + """Body center of mass velocity ``wp.spatial_vectorf`` in simulation world frame. + + Shapes are (num_instances, num_bodies,). Velocities are in the form of [wx, wy, wz, vx, vy, vz]. + This quantity contains the linear and angular velocities of the articulation links' center of mass frame + relative to the world. + """ + return self.maybe_to_torch(self._body_data.body_com_vel_w) + + @property + def body_state_w(self) -> wp.array | torch.Tensor: + """State of all bodies ``[wp.transformf, wp.spatial_vectorf]`` in simulation world frame. + + Shapes are (num_instances, num_bodies,), (num_instances, num_bodies,). The pose is in the form of [pos, quat]. + The velocity is in the form of [wx, wy, wz, vx, vy, vz]. + The pose is of the articulation links' actor frame relative to the world. + The velocity is of the articulation links' center of mass frame. + """ + + return self.maybe_to_torch(self._body_data.body_state_w) + + @property + def body_link_state_w(self) -> wp.array | torch.Tensor: + """State of all bodies' link frame ``[wp.transformf, wp.spatial_vectorf]`` in simulation world frame. + + Shapes are (num_instances, num_bodies,), (num_instances, num_bodies,). The pose is in the form of [pos, quat]. + The velocity is in the form of [wx, wy, wz, vx, vy, vz]. + The position, quaternion, and linear/angular velocity are of the body's link frame relative to the world. + """ + + return self.maybe_to_torch(self._body_data.body_link_state_w) + + @property + def body_com_state_w(self) -> wp.array | torch.Tensor: + """State of all bodies center of mass ``[wp.transformf, wp.spatial_vectorf]`` in simulation world frame. + + Shapes are (num_instances, num_bodies,), (num_instances, num_bodies,). The pose is in the form of [pos, quat]. + The velocity is in the form of [wx, wy, wz, vx, vy, vz]. + + The position, quaternion, and linear/angular velocity are of the body's center of mass frame relative to the + world. Center of mass frame is assumed to be the same orientation as the link rather than the orientation of the + principle inertia. + """ + + return self.maybe_to_torch(self._body_data.body_com_state_w) + + @property + def body_com_acc_w(self) -> wp.array | torch.Tensor: + """Acceleration of all bodies center of mass ``wp.spatial_vectorf`` in simulation world frame. + + Shapes are (num_instances, num_bodies,). The acceleration is in the form of [wx, wy, wz, vx, vy, vz]. + All values are relative to the world. + """ + return self.maybe_to_torch(self._body_data.body_com_acc_w) + + @property + def body_com_pos_b(self) -> wp.array | torch.Tensor: + """Center of mass position ``wp.vec3f`` of all bodies in their respective body's link frames. + + Shapes are (num_instances, num_bodies,). The position is in the form of [x, y, z]. + This quantity is the position of the center of mass frame of the rigid body relative to the body's link frame. + """ + return self.maybe_to_torch(self._body_data.body_com_pos_b) + + @property + def body_com_pose_b(self) -> wp.array | torch.Tensor: + """Center of mass pose ``wp.transformf`` of all bodies in their respective body's link frames. + + Shapes are (num_instances, num_bodies,). The pose is in the form of [pos, quat]. + This quantity is the pose of the center of mass frame of the rigid body relative to the body's link frame. + The orientation is provided in (x, y, z, w) format. + """ + return self.maybe_to_torch(self._body_data.body_com_pose_b) + + @property + def body_incoming_joint_wrench_b(self) -> wp.array | torch.Tensor: + """Joint reaction wrench applied from body parent to child body in parent body frame. + + Shape is (num_instances, num_bodies, 6). All body reaction wrenches are provided including the root body to the + world of an articulation. + + For more information on joint wrenches, please check the`PhysX documentation `__ + and the underlying `PhysX Tensor API `__ . + """ + return self.maybe_to_torch(self._joint_data.body_incoming_joint_wrench_b) + + ## + # Joint state properties. + ## + + @property + def joint_acc(self) -> wp.array | torch.Tensor: + """Joint acceleration of all joints. Shape is (num_instances, num_joints).""" + return self.maybe_to_torch(self._joint_data.joint_acc) + ## + # Derived Properties. + ## + + @property + def projected_gravity_b(self) -> wp.array | torch.Tensor: + """Projection of the gravity direction on base frame. Shape is (num_instances, 3).""" + return self.maybe_to_torch(self._root_data.projected_gravity_b) + + @property + def heading_w(self) -> wp.array | torch.Tensor: + """Yaw heading of the base frame (in radians). Shape is (num_instances,). + + Note: + This quantity is computed by assuming that the forward-direction of the base + frame is along x-direction, i.e. :math:`(1, 0, 0)`. + """ + return self.maybe_to_torch(self._root_data.heading_w) + + @property + def root_link_vel_b(self) -> wp.array | torch.Tensor: + """Root link velocity ``wp.spatial_vectorf`` in base frame. Shape is (num_instances). + + Velocity is provided in the form of [wx, wy, wz, vx, vy, vz]. + """ + return self.maybe_to_torch(self._root_data.root_link_vel_b) + + @property + def root_com_vel_b(self) -> wp.array | torch.Tensor: + """Root center of mass velocity ``wp.spatial_vectorf`` in base frame. Shape is (num_instances). + + Velocity is provided in the form of [wx, wy, wz, vx, vy, vz]. + """ + return self.maybe_to_torch(self._root_data.root_com_vel_b) + + ## + # Sliced properties. + ## + + @property + def root_link_pos_w(self) -> wp.array | torch.Tensor: + """Root link position ``wp.vec3f`` in simulation world frame. Shape is (num_instances). + + This quantity is the position of the actor frame of the root rigid body relative to the world. + """ + return self.maybe_to_torch(self._root_data.root_link_pos_w) + + @property + def root_link_quat_w(self) -> wp.array | torch.Tensor: + """Root link orientation ``wp.quatf`` in simulation world frame. Shape is (num_instances,). + + Format is ``(w, x, y, z)``. + This quantity is the orientation of the actor frame of the root rigid body. + """ + return self.maybe_to_torch(self._root_data.root_link_quat_w) + + @property + def root_link_lin_vel_w(self) -> wp.array | torch.Tensor: + """Root linear velocity ``wp.vec3f`` in simulation world frame. Shape is (num_instances). + + This quantity is the linear velocity of the root rigid body's actor frame relative to the world. + """ + return self.maybe_to_torch(self._root_data.root_link_lin_vel_w) + + @property + def root_link_ang_vel_w(self) -> wp.array | torch.Tensor: + """Root link angular velocity ``wp.vec3f`` in simulation world frame. Shape is (num_instances). + + This quantity is the angular velocity of the actor frame of the root rigid body relative to the world. + """ + return self.maybe_to_torch(self._root_data.root_link_ang_vel_w) + + @property + def root_com_pos_w(self) -> wp.array | torch.Tensor: + """Root center of mass position in simulation world frame. Shape is (num_instances, 3). + + This quantity is the position of the actor frame of the root rigid body relative to the world. + """ + return self.maybe_to_torch(self._root_data.root_com_pos_w) + + @property + def root_com_quat_w(self) -> wp.array | torch.Tensor: + """Root center of mass orientation ``wp.quatf`` in simulation world frame. Shape is (num_instances,). + + Format is ``(w, x, y, z)``. + This quantity is the orientation of the root rigid body's center of mass frame. + """ + return self.maybe_to_torch(self._root_data.root_com_quat_w) + + @property + def root_com_lin_vel_w(self) -> wp.array | torch.Tensor: + """Root center of mass linear velocity ``wp.vec3f`` in simulation world frame. Shape is (num_instances,). + + This quantity is the linear velocity of the root rigid body's center of mass frame relative to the world. + """ + return self.maybe_to_torch(self._root_data.root_com_lin_vel_w) + + @property + def root_com_ang_vel_w(self) -> wp.array | torch.Tensor: + """Root center of mass angular velocity ``wp.vec3f`` in simulation world frame. Shape is (num_instances). + + This quantity is the angular velocity of the root rigid body's center of mass frame relative to the world. + """ + return self.maybe_to_torch(self._root_data.root_com_ang_vel_w) + + @property + def body_link_pos_w(self) -> wp.array | torch.Tensor: + """Positions of all bodies in simulation world frame ``wp.vec3f``. Shape is (num_instances, num_bodies). + + This quantity is the position of the articulation bodies' actor frame relative to the world. + """ + return self.maybe_to_torch(self._body_data.body_link_pos_w) + + @property + def body_link_quat_w(self) -> wp.array | torch.Tensor: + """Orientation ``wp.quatf`` of all bodies in simulation world frame. Shape is (num_instances, num_bodies). + + Format is ``(w, x, y, z)``. + This quantity is the orientation of the articulation bodies' actor frame relative to the world. + """ + return self.maybe_to_torch(self._body_data.body_link_quat_w) + + @property + def body_link_lin_vel_w(self) -> wp.array | torch.Tensor: + """Linear velocity ``wp.vec3f`` of all bodies in simulation world frame. Shape is (num_instances, num_bodies). + + This quantity is the linear velocity of the articulation bodies' center of mass frame relative to the world. + """ + return self.maybe_to_torch(self._body_data.body_link_lin_vel_w) + + @property + def body_link_ang_vel_w(self) -> wp.array | torch.Tensor: + """Angular velocity ``wp.vec3f`` of all bodies in simulation world frame. Shape is (num_instances, num_bodies). + + This quantity is the angular velocity of the articulation bodies' center of mass frame. + """ + return self.maybe_to_torch(self._body_data.body_link_ang_vel_w) + + @property + def body_com_pos_w(self) -> wp.array | torch.Tensor: + """Positions of all bodies in simulation world frame ``wp.vec3f``. Shape is (num_instances, num_bodies). + + This quantity is the position of the articulation bodies' actor frame. + """ + return self.maybe_to_torch(self._body_data.body_com_pos_w) + + @property + def body_com_quat_w(self) -> wp.array | torch.Tensor: + """Orientation ``wp.quatf`` of all bodies in simulation world frame. Shape is (num_instances, num_bodies). + + Format is ``(w, x, y, z)``. + This quantity is the orientation of the articulation bodies' actor frame. + """ + return self.maybe_to_torch(self._body_data.body_com_quat_w) + + @property + def body_com_lin_vel_w(self) -> wp.array | torch.Tensor: + """Linear velocity ``wp.vec3f`` of all bodies in simulation world frame. Shape is (num_instances, num_bodies). + + This quantity is the linear velocity of the articulation bodies' center of mass frame. + """ + return self.maybe_to_torch(self._body_data.body_com_lin_vel_w) + + @property + def body_com_ang_vel_w(self) -> wp.array | torch.Tensor: + """Angular velocity ``wp.vec3f`` of all bodies in simulation world frame. Shape is (num_instances, num_bodies). + + This quantity is the angular velocity of the articulation bodies' center of mass frame. + """ + return self.maybe_to_torch(self._body_data.body_com_ang_vel_w) + + @property + def body_com_lin_acc_w(self) -> wp.array | torch.Tensor: + """Linear acceleration ``wp.vec3f`` of all bodies in simulation world frame. Shape is (num_instances, num_bodies). + + This quantity is the linear acceleration of the articulation bodies' center of mass frame. + """ + return self.maybe_to_torch(self._body_data.body_com_lin_acc_w) + + @property + def body_com_ang_acc_w(self) -> wp.array | torch.Tensor: + """Angular acceleration ``wp.vec3f`` of all bodies in simulation world frame. Shape is (num_instances, num_bodies). + + This quantity is the angular acceleration of the articulation bodies' center of mass frame. + """ + return self.maybe_to_torch(self._body_data.body_com_ang_acc_w) + + @property + def body_com_quat_b(self) -> wp.array | torch.Tensor: + """Orientation (x, y, z, w) of the principle axis of inertia of all of the bodies in their + respective link frames. Shape is (num_instances, num_bodies, 4). + + This quantity is the orientation of the principles axes of inertia relative to its body's link frame. + """ + return self.maybe_to_torch(self._body_data.body_com_quat_b) + + @property + def root_link_lin_vel_b(self) -> wp.array | torch.Tensor: + """Root link linear velocity ``wp.vec3f`` in base frame. Shape is (num_instances). + + This quantity is the linear velocity of the articulation root's actor frame with respect to the + its actor frame. + """ + return self.maybe_to_torch(self._root_data.root_link_lin_vel_b) + + @property + def root_link_ang_vel_b(self) -> wp.array | torch.Tensor: + """Root link angular velocity ``wp.vec3f`` in base world frame. Shape is (num_instances). + + This quantity is the angular velocity of the articulation root's actor frame with respect to the + its actor frame. + """ + return self.maybe_to_torch(self._root_data.root_link_ang_vel_b) + + @property + def root_com_lin_vel_b(self) -> wp.array | torch.Tensor: + """Root center of mass linear velocity ``wp.vec3f`` in base frame. Shape is (num_instances). + + This quantity is the linear velocity of the articulation root's center of mass frame with respect to the + its actor frame. + """ + return self.maybe_to_torch(self._root_data.root_com_lin_vel_b) + + @property + def root_com_ang_vel_b(self) -> wp.array | torch.Tensor: + """Root center of mass angular velocity in base world frame. Shape is (num_instances, 3). + + This quantity is the angular velocity of the articulation root's center of mass frame with respect to the + its actor frame. + """ + return self.maybe_to_torch(self._root_data.root_com_ang_vel_b) + + @property + def joint_pos_limits(self) -> wp.array | torch.Tensor: + """Joint position limits provided to the simulation. Shape is (num_instances, num_joints, 2). + + The limits are in the order :math:`[lower, upper]`. + """ + return self.maybe_to_torch(self._joint_data.joint_pos_limits) + + ## + # Backward compatibility. Need to nuke these properties in a future release. + ## + + @property + @deprecated("default_root_pose") + def default_root_state(self) -> wp.array | torch.Tensor: + """Same as :attr:`default_root_pose`.""" + state = wp.zeros((self._root_view.count), dtype=vec13f, device=self.device) + wp.launch( + combine_pose_and_velocity_to_state, + dim=(self._root_view.count), + device=self.device, + inputs=[ + self.default_root_pose, + self.default_root_vel, + state, + ], + ) + return self.maybe_to_torch(state) + + @property + @deprecated("root_link_pose_w") + def root_pose_w(self) -> wp.array | torch.Tensor: + """Same as :attr:`root_link_pose_w`.""" + return self.root_link_pose_w + + @property + @deprecated("root_link_pos_w") + def root_pos_w(self) -> wp.array | torch.Tensor: + """Same as :attr:`root_link_pos_w`.""" + return self.root_link_pos_w + + @property + @deprecated("root_link_quat_w") + def root_quat_w(self) -> wp.array | torch.Tensor: + """Same as :attr:`root_link_quat_w`.""" + return self.root_link_quat_w + + @property + @deprecated("root_com_vel_w") + def root_vel_w(self) -> wp.array | torch.Tensor: + """Same as :attr:`root_com_vel_w`.""" + return self.root_com_vel_w + + @property + @deprecated("root_com_lin_vel_w") + def root_lin_vel_w(self) -> wp.array | torch.Tensor: + """Same as :attr:`root_com_lin_vel_w`.""" + return self.root_com_lin_vel_w + + @property + @deprecated("root_com_ang_vel_w") + def root_ang_vel_w(self) -> wp.array | torch.Tensor: + """Same as :attr:`root_com_ang_vel_w`.""" + return self.root_com_ang_vel_w + + @property + @deprecated("root_com_lin_vel_b") + def root_lin_vel_b(self) -> wp.array | torch.Tensor: + """Same as :attr:`root_com_lin_vel_b`.""" + return self.root_com_lin_vel_b + + @property + @deprecated("root_com_ang_vel_b") + def root_ang_vel_b(self) -> wp.array | torch.Tensor: + """Same as :attr:`root_com_ang_vel_b`.""" + return self.root_com_ang_vel_b + + @property + @deprecated("body_link_pose_w") + def body_pose_w(self) -> wp.array | torch.Tensor: + """Same as :attr:`body_link_pose_w`.""" + return self.body_link_pose_w + + @property + @deprecated("body_link_pos_w") + def body_pos_w(self) -> wp.array | torch.Tensor: + """Same as :attr:`body_link_pos_w`.""" + return self.body_link_pos_w + + @property + @deprecated("body_link_quat_w") + def body_quat_w(self) -> wp.array | torch.Tensor: + """Same as :attr:`body_link_quat_w`.""" + return self.body_link_quat_w + + @property + @deprecated("body_com_vel_w") + def body_vel_w(self) -> wp.array | torch.Tensor: + """Same as :attr:`body_com_vel_w`.""" + return self.body_com_vel_w + + @property + @deprecated("body_com_lin_vel_w") + def body_lin_vel_w(self) -> wp.array | torch.Tensor: + """Same as :attr:`body_com_lin_vel_w`.""" + return self.body_com_lin_vel_w + + @property + @deprecated("body_com_ang_vel_w") + def body_ang_vel_w(self) -> wp.array | torch.Tensor: + """Same as :attr:`body_com_ang_vel_w`.""" + return self.body_com_ang_vel_w + + @property + @deprecated("body_com_acc_w") + def body_acc_w(self) -> wp.array | torch.Tensor: + """Same as :attr:`body_com_acc_w`.""" + return self.body_com_acc_w + + @property + @deprecated("body_com_lin_acc_w") + def body_lin_acc_w(self) -> wp.array | torch.Tensor: + """Same as :attr:`body_com_lin_acc_w`.""" + return self.body_com_lin_acc_w + + @property + @deprecated("body_com_ang_acc_w") + def body_ang_acc_w(self) -> wp.array | torch.Tensor: + """Same as :attr:`body_com_ang_acc_w`.""" + return self.body_com_ang_acc_w + + @property + @deprecated("body_com_pos_b") + def com_pos_b(self) -> wp.array | torch.Tensor: + """Same as :attr:`body_com_pos_b`.""" + return self.body_com_pos_b + + @property + @deprecated("body_com_quat_b") + def com_quat_b(self) -> wp.array | torch.Tensor: + """Same as :attr:`body_com_quat_b`.""" + return self.body_com_quat_b + + @property + @deprecated("joint_pos_limits") + def joint_limits(self) -> wp.array | torch.Tensor: + """Deprecated property. Please use :attr:`joint_pos_limits` instead.""" + return self.joint_pos_limits + + @property + @deprecated("joint_friction_coeff") + def joint_friction(self) -> wp.array | torch.Tensor: + """Deprecated property. Please use :attr:`joint_friction_coeff` instead.""" + return self.joint_friction_coeff + + def default_fixed_tendon_limit(self): + return None + def default_joint_friction(self): + return None + def default_joint_limits(self): + return None + def fixed_tendon_limit(self): + return None + def joint_velocity_limits(self): + return None \ No newline at end of file diff --git a/source/isaaclab/isaaclab/newton/assets/core/__init__.py b/source/isaaclab/isaaclab/newton/assets/core/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/source/isaaclab/isaaclab/newton/assets/core/body_properties/body.py b/source/isaaclab/isaaclab/newton/assets/core/body_properties/body.py new file mode 100644 index 00000000000..1c07c7b2a2d --- /dev/null +++ b/source/isaaclab/isaaclab/newton/assets/core/body_properties/body.py @@ -0,0 +1,203 @@ +# Copyright (c) 2022-2025, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +# Flag for pyright to ignore type errors in this file. +# pyright: reportPrivateUsage=false + + +from collections.abc import Sequence +import weakref +import warp as wp + +import isaaclab.utils.string as string_utils +from newton.selection import ArticulationView as NewtonArticulationView +from isaaclab.newton.assets.core.body_properties.body_data import BodyData +from isaaclab.newton.assets.core.kernels import ( + generate_mask_from_ids, + update_wrench_array_with_force, + update_wrench_array_with_torque, + update_batched_array_with_value_masked, +) + + +class Body: + def __init__(self, root_newton_view: NewtonArticulationView, body_data: BodyData, device: str): + self._root_newton_view: NewtonArticulationView = weakref.proxy(root_newton_view) + self._data = body_data + self._device = device + + self._create_buffers() + + """ + Properties + """ + + @property + def data(self) -> BodyData: + return self._data + + @property + def num_instances(self) -> int: + return self._root_newton_view.count + + @property + def num_bodies(self) -> int: + """Number of bodies in articulation.""" + return self._root_newton_view.link_count + + @property + def body_names(self) -> list[str]: + """Ordered names of bodies in articulation.""" + return self._root_newton_view.body_names + + @property + def root_newton_view(self) -> NewtonArticulationView: + """Articulation view for the asset (Newton). + + Note: + Use this view with caution. It requires handling of tensors in a specific way. + """ + return self._root_newton_view + + """ + Operations. + """ + + def reset(self, mask: wp.array): + # reset external wrench + wp.launch( + update_batched_array_with_value_masked, + dim=(self.num_instances, self.num_bodies), + inputs=[ + wp.spatial_vectorf(0.0, 0.0, 0.0, 0.0, 0.0, 0.0), + self._data.external_wrench, + mask, + self._ALL_BODY_MASK, + ], + ) + + def write_data_to_sim(self): + """Write external wrenches and joint commands to the simulation. + + If any explicit actuators are present, then the actuator models are used to compute the + joint commands. Otherwise, the joint commands are directly set into the simulation. + + Note: + We write external wrench to the simulation here since this function is called before the simulation step. + This ensures that the external wrench is applied at every simulation step. + """ + # Wrenches are automatically applied by set_external_force_and_torque. + pass + + def update(self, dt: float): + self._data.update(dt) + + """ + Operations - Finders. + """ + + def find_bodies( + self, name_keys: str | Sequence[str], preserve_order: bool = False + ) -> tuple[wp.array, list[str], list[int]]: + """Find bodies in the articulation based on the name keys. + + Please check the :meth:`isaaclab.utils.string_utils.resolve_matching_names` function for more + information on the name matching. + + Args: + name_keys: A regular expression or a list of regular expressions to match the body names. + preserve_order: Whether to preserve the order of the name keys in the output. Defaults to False. + + Returns: + A tuple of lists containing the body mask, names, and indices. + """ + indices, names = string_utils.resolve_matching_names(name_keys, self.body_names, preserve_order) + self._BODY_MASK.fill_(False) + mask = wp.clone(self._BODY_MASK) + wp.launch( + generate_mask_from_ids, + dim=(len(indices),), + inputs=[ + mask, + wp.array(indices, dtype=wp.int32, device=self._device), + ], + ) + return mask, names, indices + + """ + Operations - Setters. + """ + + def set_external_force_and_torque( + self, + forces: wp.array, + torques: wp.array, + body_mask: wp.array | None = None, + env_mask: wp.array | None = None, + ) -> None: + """Set external force and torque to apply on the asset's bodies in their local frame. + + For many applications, we want to keep the applied external force on rigid bodies constant over a period of + time (for instance, during the policy control). This function allows us to store the external force and torque + into buffers which are then applied to the simulation at every step. + + .. caution:: + If the function is called with empty forces and torques, then this function disables the application + of external wrench to the simulation. + + .. code-block:: python + + # example of disabling external wrench + asset.set_external_force_and_torque(forces=wp.zeros(0, 3), torques=wp.zeros(0, 3)) + + .. note:: + This function does not apply the external wrench to the simulation. It only fills the buffers with + the desired values. To apply the external wrench, call the :meth:`write_data_to_sim` function + right before the simulation step. + + Args: + forces: External forces in bodies' local frame. Shape is (num_instances, num_bodies, 3). + torques: External torques in bodies' local frame. Shape is (num_instances, num_bodies, 3). + body_mask: The body mask. Shape is (num_bodies). + env_mask: The environment mask. Shape is (num_instances,). + """ + # resolve indices + if env_mask is None: + env_mask = self._ALL_ENV_MASK + if body_mask is None: + body_mask = self._ALL_BODY_MASK + # Check if there are any external forces or torques + if (forces is not None) or (torques is not None): + self.has_external_wrench = True + if forces is not None: + wp.launch( + update_wrench_array_with_force, + dim=(self.num_instances, self.num_bodies), + inputs=[ + forces, + self._data.external_wrench, + env_mask, + body_mask, + ], + ) + if torques is not None: + wp.launch( + update_wrench_array_with_torque, + dim=(self.num_instances, self.num_bodies), + inputs=[ + torques, + self._data.external_wrench, + env_mask, + body_mask, + ], + ) + + def _create_buffers(self): + # constants + self._ALL_ENV_MASK = wp.ones((self.num_instances,), dtype=wp.bool, device=self._device) + self._ALL_BODY_MASK = wp.ones((self.num_bodies,), dtype=wp.bool, device=self._device) + # masks + self._ENV_MASK = wp.zeros((self.num_instances,), dtype=wp.bool, device=self._device) + self._BODY_MASK = wp.zeros((self.num_bodies,), dtype=wp.bool, device=self._device) \ No newline at end of file diff --git a/source/isaaclab/isaaclab/newton/assets/core/body_properties/body_data.py b/source/isaaclab/isaaclab/newton/assets/core/body_properties/body_data.py new file mode 100644 index 00000000000..dbb9f4d7f13 --- /dev/null +++ b/source/isaaclab/isaaclab/newton/assets/core/body_properties/body_data.py @@ -0,0 +1,562 @@ +# Copyright (c) 2022-2025, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +import weakref + +import warp as wp + +from isaaclab.sim._impl.newton_manager import NewtonManager +from newton.selection import ArticulationView as NewtonArticulationView +from isaaclab.utils.buffers import TimestampedWarpBuffer +from isaaclab.utils.helpers import deprecated, warn_overhead_cost + +from isaaclab.newton.assets.core.kernels import ( + combine_frame_transforms_partial_batch, + combine_pose_and_velocity_to_state_batched, + derive_body_acceleration_from_velocity_batched, + generate_pose_from_position_with_unit_quaternion_batched, + project_com_velocity_to_link_frame_batch, + vec13f, + split_transform_batched_array_to_position_batched_array, + split_transform_batched_array_to_quaternion_batched_array, + split_spatial_vectory_batched_array_to_linear_velocity_batched_array, + split_spatial_vectory_batched_array_to_angular_velocity_batched_array, +) + + +class BodyData: + def __init__(self, root_newton_view: NewtonArticulationView, device: str) -> None: + """Initializes the container for body data. + + Usually, bodies are all the rigid bodies belonging to an articulation. + + Args: + root_newton_view: The root articulation view. + device: The device used for processing. + """ + # Set the parameters + self.device = device + # Set the root articulation view + # note: this is stored as a weak reference to avoid circular references between the asset class + # and the data container. This is important to avoid memory leaks. + self._root_newton_view: NewtonArticulationView = weakref.proxy(root_newton_view) + + # Set initial time stamp + self._sim_timestamp = 0.0 + + self._create_simulation_bindings() + self._create_buffers() + + def update(self, dt: float): + # update the simulation timestamp + self._sim_timestamp += dt + # Trigger an update of the body acceleration buffer at a higher frequency + # since we do finite differencing. + self.body_com_acc_w + + def _create_simulation_bindings(self): + """Create simulation bindings for the body data. + + Direct simulation bindings are pointers to the simulation data, their data is not copied, and should + only be updated using warp kernels. Any modifications made to the bindings will be reflected in the simulation. + Hence we encourage users to carefully think about the data they modify and in which order it should be updated. + + .. caution:: This is possible if and only if the properties that we access are strided from newton and not + indexed. Newton willing this is the case all the time, but we should pay attention to this if things look off. + """ + self._sim_bind_body_link_pose_w = self._root_newton_view.get_link_transforms(NewtonManager.get_state_0()) + self._sim_bind_body_com_vel_w = self._root_newton_view.get_link_velocities(NewtonManager.get_state_0()) + self._sim_bind_body_com_pos_b = self._root_newton_view.get_attribute("body_com", NewtonManager.get_model()) + self._sim_bind_body_mass = self._root_newton_view.get_attribute("body_mass", NewtonManager.get_model()) + self._sim_bind_body_inertia = self._root_newton_view.get_attribute("body_inertia", NewtonManager.get_model()) + self._sim_bind_body_external_wrench = self._root_newton_view.get_attribute("body_f", NewtonManager.get_state_0()) + + def _create_buffers(self): + """Create buffers for the body data.""" + # Initialize history for finite differencing + self._previous_body_com_vel = wp.clone(self._root_newton_view.get_link_velocities(NewtonManager.get_state_0())) + + # Initialize the lazy buffers. + # -- link frame w.r.t. world frame + self._body_link_vel_w = TimestampedWarpBuffer( + shape=(self._root_newton_view.count, self._root_newton_view.link_count), dtype=wp.spatial_vectorf + ) + # -- com frame w.r.t. world frame + self._body_com_pose_w = TimestampedWarpBuffer( + shape=(self._root_newton_view.count, self._root_newton_view.link_count), dtype=wp.transformf + ) + self._body_com_acc_w = TimestampedWarpBuffer( + shape=(self._root_newton_view.count, self._root_newton_view.link_count), dtype=wp.spatial_vectorf + ) + + ## + # Direct simulation bindings accessors. + ## + + @property + def body_link_pose_w(self) -> wp.array: + """Body link pose ``wp.transformf`` in the world frame. Shape is (num_instances, num_bodies). + + The pose is in the form of [pos, quat]. The orientation is provided in (x, y, z, w) format. + """ + return self._sim_bind_body_link_pose_w + + @property + def body_com_vel_w(self) -> wp.array: + """Body center of mass velocity ``wp.spatial_vectorf`` in the world frame. Shape is (num_instances, num_bodies). + + The velocity is in the form of [ang_vel, lin_vel]. + """ + return self._sim_bind_body_com_vel_w + + @property + def body_com_pos_b(self) -> wp.array: + """Center of mass pose ``wp.transformf`` of all bodies in their respective body's link frames. + + Shapes are (num_instances, num_bodies,). The pose is in the form of [pos, quat]. + This quantity is the pose of the center of mass frame of the rigid body relative to the body's link frame. + The orientation is provided in (x, y, z, w) format. + """ + return self._sim_bind_body_com_pos_b + + ## + # Body state properties. + ## + + @property + def body_mass(self) -> wp.array: + """Body mass ``wp.float32`` in the world frame. Shape is (num_instances, num_bodies).""" + return self._sim_bind_body_mass + + @property + def body_inertia(self) -> wp.array: + """Body inertia ``wp.mat33`` in the world frame. Shape is (num_instances, num_bodies, 9).""" + return self._sim_bind_body_inertia + + @property + def external_wrench(self) -> wp.array: + """External wrench ``wp.spatial_vectorf`` in the world frame. Shape is (num_instances, num_bodies, 6).""" + return self._sim_bind_body_external_wrench + + @property + def body_link_vel_w(self) -> wp.array: + """Body link velocity ``wp.spatial_vectorf`` in simulation world frame. + + Shapes are (num_instances, num_bodies,). Velocities are in the form of [wx, wy, wz, vx, vy, vz]. + This quantity contains the linear and angular velocities of the articulation links' actor frame + relative to the world. + """ + if self._body_link_vel_w.timestamp < self._sim_timestamp: + # Project the velocity from the center of mass frame to the link frame + wp.launch( + project_com_velocity_to_link_frame_batch, + dim=(self._root_newton_view.count, self._root_newton_view.link_count), + device=self.device, + inputs=[ + self._sim_bind_body_com_vel_w, + self._sim_bind_body_link_pose_w, + self._sim_bind_body_com_pos_b, + self._body_link_vel_w.data, + ], + ) + # set the buffer data and timestamp + self._body_link_vel_w.timestamp = self._sim_timestamp + return self._body_link_vel_w.data + + @property + def body_com_pose_w(self) -> wp.array: + """Body center of mass pose ``wp.transformf`` in simulation world frame. + + Shapes are (num_instances, num_bodies,). The pose is in the form of [pos, quat]. + This quantity is the pose of the center of mass frame of the articulation links relative to the world. + The orientation is provided in (x, y, z, w) format. + """ + if self._body_com_pose_w.timestamp < self._sim_timestamp: + # Apply local transform to center of mass frame + wp.launch( + combine_frame_transforms_partial_batch, + dim=(self._root_newton_view.count, self._root_newton_view.link_count), + device=self.device, + inputs=[ + self._sim_bind_body_link_pose_w, + self._sim_bind_body_com_pos_b, + self._body_com_pose_w.data, + ], + ) + # set the buffer data and timestamp + self._body_com_pose_w.timestamp = self._sim_timestamp + return self._body_com_pose_w.data + + @property + @warn_overhead_cost( + "body_link_pose_w and/or body_com_vel_w", + "This function outputs the state of the link frame, containing both pose and velocity. However, Newton outputs" + " pose and velocities separately. Consider using only one of them instead. If both are required, use both" + " body_link_pose_w and body_com_vel_w instead.", + ) + def body_state_w(self) -> wp.array: + """State of all bodies ``vec13f`` in simulation world frame. + + Shapes are (num_instances, num_bodies, 13). The pose is in the form of [pos, quat]. + + The pose is of the articulation links' actor frame relative to the world. + The velocity is of the articulation links' center of mass frame. + + .. note:: The velocity is in the form of [vx, vy, vz, wx, wy, wz]. + .. note:: The quaternion is in the form of [x, y, z, w]. + """ + + state = wp.zeros( + (self._root_newton_view.count, self._root_newton_view.link_count), dtype=vec13f, device=self.device + ) + wp.launch( + combine_pose_and_velocity_to_state_batched, + dim=(self._root_newton_view.count, self._root_newton_view.link_count), + device=self.device, + inputs=[ + self._sim_bind_body_link_pose_w, + self._sim_bind_body_com_vel_w, + state, + ], + ) + return state + + @property + @warn_overhead_cost( + "body_link_pose_w and/or body_link_vel_w", + "This function outputs the state of the link frame, containing both pose and velocity. However, Newton outputs" + " pose and velocities separately. Consider using only one of them instead. If both are required, use both" + " body_link_pose_w and body_link_vel_w instead.", + ) + def body_link_state_w(self) -> wp.array: + """State of all bodies' link frame ``vec13f`` in simulation world frame. + + Shapes are (num_instances, num_bodies, 13). The pose is in the form of [pos, quat]. + The position, quaternion, and linear/angular velocity are of the body's link frame relative to the world. + + .. note:: The velocity is in the form of [vx, vy, vz, wx, wy, wz]. + .. note:: The quaternion is in the form of [x, y, z, w]. + """ + + state = wp.zeros( + (self._root_newton_view.count, self._root_newton_view.link_count), dtype=vec13f, device=self.device + ) + wp.launch( + combine_pose_and_velocity_to_state_batched, + dim=(self._root_newton_view.count, self._root_newton_view.link_count), + device=self.device, + inputs=[ + self._sim_bind_body_link_pose_w, + self.body_link_vel_w, + state, + ], + ) + return state + + @property + @warn_overhead_cost( + "body_com_pose_w and/or body_com_vel_w", + "This function outputs the state of the CoM, containing both pose and velocity. However, Newton outputs pose" + " and velocities separately. Consider using only one of them instead. If both are required, use both" + " body_com_pose_w and body_com_vel_w instead.", + ) + def body_com_state_w(self) -> wp.array: + """State of all bodies center of mass ``[wp.transformf, wp.spatial_vectorf]`` in simulation world frame. + + Shapes are (num_instances, num_bodies,), (num_instances, num_bodies,). The pose is in the form of [pos, quat]. + + The position, quaternion, and linear/angular velocity are of the body's center of mass frame relative to the + world. Center of mass frame is assumed to be the same orientation as the link rather than the orientation of the + principle inertia. + + .. note:: The velocity is in the form of [vx, vy, vz, wx, wy, wz]. + .. note:: The quaternion is in the form of [x, y, z, w]. + """ + + state = wp.zeros( + (self._root_newton_view.count, self._root_newton_view.link_count), dtype=vec13f, device=self.device + ) + wp.launch( + combine_pose_and_velocity_to_state_batched, + dim=(self._root_newton_view.count, self._root_newton_view.link_count), + device=self.device, + inputs=[ + self.body_com_pose_w, + self._sim_bind_body_com_vel_w, + state, + ], + ) + return state + + @property + def body_com_acc_w(self) -> wp.array: + """Acceleration of all bodies center of mass ``wp.spatial_vectorf`` in simulation world frame. + + Shapes are (num_instances, num_bodies,). All values are relative to the world. + + .. note:: The acceleration is in the form of [vx, vy, vz, wx, wy, wz]. + """ + if self._body_com_acc_w.timestamp < self._sim_timestamp: + wp.launch( + derive_body_acceleration_from_velocity_batched, + dim=(self._root_newton_view.count, self._root_newton_view.link_count), + inputs=[ + self._sim_bind_body_com_vel_w, + self._previous_body_com_vel, + NewtonManager.get_dt(), + self._body_com_acc_w.data, + ], + ) + # set the buffer data and timestamp + self._body_com_acc_w.timestamp = self._sim_timestamp + return self._body_com_acc_w.data + + @property + @warn_overhead_cost( + "body_com_pose_b", + "This function outputs the pose of the CoM, containing both position and orientation. However, in Newton, the" + " CoM is always aligned with the link frame. This means that the quaternion is always [0, 0, 0, 1]. Consider" + " using the position only instead.", + ) + def body_com_pose_b(self) -> wp.array: + """Center of mass pose ``wp.transformf`` of all bodies in their respective body's link frames. + + Shapes are (num_instances, num_bodies,). The pose is in the form of [pos, quat]. + This quantity is the pose of the center of mass frame of the rigid body relative to the body's link frame. + + .. note:: The quaternion is in the form of [x, y, z, w]. + """ + out = wp.zeros( + (self._root_newton_view.count, self._root_newton_view.link_count), dtype=wp.transformf, device=self.device + ) + wp.launch( + generate_pose_from_position_with_unit_quaternion_batched, + dim=(self._root_newton_view.count, self._root_newton_view.link_count), + inputs=[ + self.body_com_pos_b, + out, + ], + ) + return out + + ## + # Strided properties. + ## + + @property + @warn_overhead_cost("root_link_pose_w", "Returns a strided array, this is fine if this results is called only once at the beginning of the simulation." + "However, if this is called multiple times, consider using the whole of the pose instead.") + def body_link_pos_w(self) -> wp.array: + """Positions of all bodies in simulation world frame ``wp.vec3f``. Shape is (num_instances, num_bodies). + + This quantity is the position of the articulation bodies' actor frame relative to the world. + """ + out = wp.zeros((self._root_newton_view.count, self._root_newton_view.link_count), dtype=wp.vec3f, device=self.device) + wp.launch( + split_transform_batched_array_to_position_batched_array, + dim=(self._root_newton_view.count, self._root_newton_view.link_count), + inputs=[ + self._sim_bind_body_link_pose_w, + out, + ], + ) + return out + + @property + @warn_overhead_cost("root_link_pose_w", "Returns a strided array, this is fine if this results is called only once at the beginning of the simulation." + "However, if this is called multiple times, consider using the whole of the pose instead.") + def body_link_quat_w(self) -> wp.array: + """Orientation ``wp.quatf`` of all bodies in simulation world frame. Shape is (num_instances, num_bodies). + + This quantity is the orientation of the articulation bodies' actor frame relative to the world. + + .. note:: The quaternion is in the form of [x, y, z, w]. + """ + out = wp.zeros((self._root_newton_view.count, self._root_newton_view.link_count), dtype=wp.quatf, device=self.device) + wp.launch( + split_transform_batched_array_to_quaternion_batched_array, + dim=(self._root_newton_view.count, self._root_newton_view.link_count), + inputs=[ + self._sim_bind_body_link_pose_w, + out, + ], + ) + return out + + @property + @warn_overhead_cost("root_link_vel_w", "Returns a strided array, this is fine if this results is called only once at the beginning of the simulation." + "However, if this is called multiple times, consider using the whole of the velocity instead.") + def body_link_lin_vel_w(self) -> wp.array: + """Linear velocity ``wp.vec3f`` of all bodies in simulation world frame. Shape is (num_instances, num_bodies). + + This quantity is the linear velocity of the articulation bodies' center of mass frame relative to the world. + """ + out = wp.zeros((self._root_newton_view.count, self._root_newton_view.link_count), dtype=wp.vec3f, device=self.device) + wp.launch( + split_spatial_vectory_batched_array_to_linear_velocity_batched_array, + dim=(self._root_newton_view.count, self._root_newton_view.link_count), + inputs=[ + self.body_link_vel_w, + out, + ], + ) + return out + + @property + @warn_overhead_cost("root_link_vel_w", "Returns a strided array, this is fine if this results is called only once at the beginning of the simulation." + "However, if this is called multiple times, consider using the whole of the velocity instead.") + def body_link_ang_vel_w(self) -> wp.array: + """Angular velocity ``wp.vec3f`` of all bodies in simulation world frame. Shape is (num_instances, num_bodies). + + This quantity is the angular velocity of the articulation bodies' center of mass frame relative to the world. + """ + out = wp.zeros((self._root_newton_view.count, self._root_newton_view.link_count), dtype=wp.vec3f, device=self.device) + wp.launch( + split_spatial_vectory_batched_array_to_angular_velocity_batched_array, + dim=(self._root_newton_view.count, self._root_newton_view.link_count), + inputs=[ + self.body_link_vel_w, + out, + ], + ) + return out + + @property + @warn_overhead_cost("root_com_pose_w", "Returns a strided array, this is fine if this results is called only once at the beginning of the simulation." + "However, if this is called multiple times, consider using the whole of the pose instead.") + def body_com_pos_w(self) -> wp.array: + """Positions of all bodies in simulation world frame ``wp.vec3f``. Shape is (num_instances, num_bodies). + + This quantity is the position of the articulation bodies' actor frame. + """ + out = wp.zeros((self._root_newton_view.count, self._root_newton_view.link_count), dtype=wp.vec3f, device=self.device) + wp.launch( + split_transform_batched_array_to_position_batched_array, + dim=(self._root_newton_view.count, self._root_newton_view.link_count), + inputs=[ + self.body_com_pose_w, + out, + ], + ) + return out + + @property + @warn_overhead_cost("root_com_pose_w", "Returns a strided array, this is fine if this results is called only once at the beginning of the simulation." + "However, if this is called multiple times, consider using the whole of the pose instead.") + def body_com_quat_w(self) -> wp.array: + """Orientation ``wp.quatf`` of all bodies in simulation world frame. Shape is (num_instances, num_bodies). + + This quantity is the orientation of the articulation bodies' actor frame. + + .. note:: The quaternion is in the form of [x, y, z, w]. + """ + out = wp.zeros((self._root_newton_view.count, self._root_newton_view.link_count), dtype=wp.quatf, device=self.device) + wp.launch( + split_transform_batched_array_to_quaternion_batched_array, + dim=(self._root_newton_view.count, self._root_newton_view.link_count), + inputs=[ + self.body_com_pose_w, + out, + ], + ) + return out + + @property + @warn_overhead_cost("root_com_vel_w", "Returns a strided array, this is fine if this results is called only once at the beginning of the simulation." + "However, if this is called multiple times, consider using the whole of the velocity instead.") + def body_com_lin_vel_w(self) -> wp.array: + """Linear velocity ``wp.vec3f`` of all bodies in simulation world frame. Shape is (num_instances, num_bodies). + + This quantity is the linear velocity of the articulation bodies' center of mass frame. + """ + out = wp.zeros((self._root_newton_view.count, self._root_newton_view.link_count), dtype=wp.vec3f, device=self.device) + wp.launch( + split_spatial_vectory_batched_array_to_linear_velocity_batched_array, + dim=(self._root_newton_view.count, self._root_newton_view.link_count), + inputs=[ + self._sim_bind_body_com_vel_w, + out, + ], + ) + return out + + @property + @warn_overhead_cost("root_com_vel_w", "Returns a strided array, this is fine if this results is called only once at the beginning of the simulation." + "However, if this is called multiple times, consider using the whole of the velocity instead.") + def body_com_ang_vel_w(self) -> wp.array: + """Angular velocity ``wp.vec3f`` of all bodies in simulation world frame. Shape is (num_instances, num_bodies). + + This quantity is the angular velocity of the articulation bodies' center of mass frame. + """ + out = wp.zeros((self._root_newton_view.count, self._root_newton_view.link_count), dtype=wp.vec3f, device=self.device) + wp.launch( + split_spatial_vectory_batched_array_to_angular_velocity_batched_array, + dim=(self._root_newton_view.count, self._root_newton_view.link_count), + inputs=[ + self._sim_bind_body_com_vel_w, + out, + ], + ) + return out + + @property + @warn_overhead_cost("root_com_vel_w", "Returns a strided array, this is fine if this results is called only once at the beginning of the simulation." + "However, if this is called multiple times, consider using the whole of the velocity instead.") + def body_com_lin_acc_w(self) -> wp.array: + """Linear acceleration ``wp.vec3f`` of all bodies in simulation world frame. Shape is (num_instances, num_bodies). + + This quantity is the linear acceleration of the articulation bodies' center of mass frame. + """ + out = wp.zeros((self._root_newton_view.count, self._root_newton_view.link_count), dtype=wp.vec3f, device=self.device) + wp.launch( + split_spatial_vectory_batched_array_to_linear_velocity_batched_array, + dim=(self._root_newton_view.count, self._root_newton_view.link_count), + inputs=[ + self.body_com_acc_w, + out, + ], + ) + return out + + @property + @warn_overhead_cost("root_com_vel_w", "Returns a strided array, this is fine if this results is called only once at the beginning of the simulation." + "However, if this is called multiple times, consider using the whole of the velocity instead.") + def body_com_ang_acc_w(self) -> wp.array: + """Angular acceleration ``wp.vec3f`` of all bodies in simulation world frame. Shape is (num_instances, num_bodies). + + This quantity is the angular acceleration of the articulation bodies' center of mass frame. + """ + out = wp.zeros((self._root_newton_view.count, self._root_newton_view.link_count), dtype=wp.vec3f, device=self.device) + wp.launch( + split_spatial_vectory_batched_array_to_angular_velocity_batched_array, + dim=(self._root_newton_view.count, self._root_newton_view.link_count), + inputs=[ + self.body_com_acc_w, + out, + ], + ) + return out + + @property + @warn_overhead_cost("root_com_pose_w", "Returns a strided array, this is fine if this results is called only once at the beginning of the simulation." + "However, if this is called multiple times, consider using the whole of the pose instead.") + def body_com_quat_b(self) -> wp.array: + """Orientation ``wp.quatf`` of the principle axis of inertia of all of the bodies in their + respective link frames. Shape is (num_instances, num_bodies, 4). + + This quantity is the orientation of the principles axes of inertia relative to its body's link frame. + + .. note:: The quaternion is in the form of [x, y, z, w]. + """ + out = wp.zeros((self._root_newton_view.count, self._root_newton_view.link_count), dtype=wp.quatf, device=self.device) + wp.launch( + split_transform_batched_array_to_quaternion_batched_array, + dim=(self._root_newton_view.count, self._root_newton_view.link_count), + inputs=[ + self.body_com_pose_b, + out, + ], + ) + return out \ No newline at end of file diff --git a/source/isaaclab/isaaclab/newton/assets/core/joint_properties/__init__.py b/source/isaaclab/isaaclab/newton/assets/core/joint_properties/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/source/isaaclab/isaaclab/newton/assets/core/joint_properties/joint.py b/source/isaaclab/isaaclab/newton/assets/core/joint_properties/joint.py new file mode 100644 index 00000000000..6a9e4f172b7 --- /dev/null +++ b/source/isaaclab/isaaclab/newton/assets/core/joint_properties/joint.py @@ -0,0 +1,781 @@ +# Copyright (c) 2022-2025, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +# Flag for pyright to ignore type errors in this file. +# pyright: reportPrivateUsage=false + + +from collections.abc import Sequence + +import omni.log +import weakref +import warp as wp +from newton.selection import ArticulationView as NewtonArticulationView +from newton.solvers import SolverMuJoCo, SolverNotifyFlags + +import isaaclab.utils.string as string_utils +from isaaclab.sim._impl.newton_manager import NewtonManager + +from isaaclab.newton.assets.core.joint_properties.joint_data import JointData +from isaaclab.newton.assets.core.kernels import ( + generate_mask_from_ids, + update_joint_array, + update_joint_array_int, + update_joint_array_with_value, + update_joint_array_with_value_int, + update_joint_limits, + update_joint_limits_value_vec2f, + update_joint_pos_with_limits, + update_joint_pos_with_limits_value_vec2f, + update_soft_joint_pos_limits, +) + + +class Joint: + """Core implementation of the joint setters in warp. + + This class is a zero copy, mask only, implementation of the joint setters in warp. + It relies on the buffers in the joint data class, and updates them based on the users inputs. + Note that since most of the buffers in the joint data class are direct simulation bindings, there + is no need to explicitly call the selection API in most cases. + """ + + def __init__(self, root_newton_view: NewtonArticulationView, joint_data: JointData, soft_joint_pos_limit_factor: float, device: str): + self._root_newton_view: NewtonArticulationView = weakref.proxy(root_newton_view) + self._data = joint_data + self._soft_joint_pos_limit_factor = soft_joint_pos_limit_factor + self._device = device + + self._create_buffers() + + @property + def data(self) -> JointData: + return self._data + + @property + def num_instances(self) -> int: + return self._root_newton_view.count + + @property + def num_joints(self) -> int: + """Number of joints in articulation.""" + return self._root_newton_view.joint_dof_count + + @property + def joint_names(self) -> list[str]: + """Ordered names of joints in articulation.""" + return self._root_newton_view.joint_dof_names + + @property + def root_newton_view(self) -> NewtonArticulationView: + """Articulation view for the asset (Newton). + + Note: + Use this view with caution. It requires handling of tensors in a specific way. + """ + return self._root_newton_view + + def update(self, dt: float) -> None: + self._data.update(dt) + + """ + Operations - Finders. + """ + + def find_joints( + self, name_keys: str | Sequence[str], joint_subset: list[str] | None = None, preserve_order: bool = False + ) -> tuple[wp.array, list[str], list[int]]: + """Find joints in the articulation based on the name keys. + + Please see the :func:`isaaclab.utils.string.resolve_matching_names` function for more information + on the name matching. + + Args: + name_keys: A regular expression or a list of regular expressions to match the joint names. + joint_subset: A subset of joints to search for. Defaults to None, which means all joints + in the articulation are searched. + preserve_order: Whether to preserve the order of the name keys in the output. Defaults to False. + + Returns: + A tuple of lists containing the joint mask, names, and indices. + """ + if joint_subset is None: + joint_subset = self.joint_names + # find joints + indices, names = string_utils.resolve_matching_names(name_keys, joint_subset, preserve_order) + self._JOINT_MASK.fill_(False) + mask = wp.clone(self._JOINT_MASK) + wp.launch( + generate_mask_from_ids, + dim=(len(indices),), + inputs=[ + mask, + wp.array(indices, dtype=wp.int32, device=self._device), + ], + ) + return mask, names, indices + + """ + Operations. + """ + + def reset(self, mask: wp.array): + pass + + """ + Operations - State Writers. + """ + + def write_joint_state_to_sim( + self, + position: wp.array, + velocity: wp.array, + joint_mask: wp.array | None = None, + env_mask: wp.array | None = None, + ): + """Write joint positions and velocities to the simulation. + + Args: + position: Joint positions. Shape is (num_instances, num_joints). + velocity: Joint velocities. Shape is (num_instances, num_joints). + joint_mask: The joint mask. Shape is (num_joints). + env_mask: The environment mask. Shape is (num_instances,). + """ + # set into simulation + self.write_joint_position_to_sim(position, joint_mask=joint_mask, env_mask=env_mask) + self.write_joint_velocity_to_sim(velocity, joint_mask=joint_mask, env_mask=env_mask) + + def write_joint_position_to_sim( + self, + position: wp.array, + joint_mask: wp.array | None = None, + env_mask: wp.array | None = None, + ): + """Write joint positions to the simulation. + + Args: + position: Joint positions. Shape is (num_instances, num_joints). + joint_mask: The joint mask. Shape is (num_joints). + env_mask: The environment mask. Shape is (num_instances,). + """ + # resolve indices + if env_mask is None: + env_mask = self._ALL_ENV_MASK + if joint_mask is None: + joint_mask = self._ALL_JOINT_MASK + # set into internal buffers + wp.launch( + update_joint_array, + dim=(self.num_instances, self.num_joints), + inputs=[ + position, + self._data.joint_pos, + env_mask, + joint_mask, + ], + ) + # We do not invalidate the buffers related to the body poses since they will only be updated by the simulation + # At the next step. If this is needed then we need to call forward_kinematics first, and then invalidate the + # buffers related to the body poses. This is not cheap, and we should only do it if strictly necessary. + + def write_joint_velocity_to_sim( + self, + velocity: wp.array, + joint_mask: wp.array | None = None, + env_mask: wp.array | None = None, + ): + """Write joint velocities to the simulation. + + Args: + velocity: Joint velocities. Shape is (num_instances, num_joints). + joint_mask: The joint mask. Shape is (num_joints). + env_mask: The environment mask. Shape is (num_instances,). + """ + if env_mask is None: + env_mask = self._ALL_ENV_MASK + if joint_mask is None: + joint_mask = self._ALL_JOINT_MASK + # update joint velocity + wp.launch( + update_joint_array, + dim=(self.num_instances, self.num_joints), + inputs=[ + velocity, + self._data.joint_vel, + env_mask, + joint_mask, + ], + ) + # update previous joint velocity + wp.launch( + update_joint_array, + dim=(self.num_instances, self.num_joints), + inputs=[ + velocity, + self._data._previous_joint_vel, + env_mask, + joint_mask, + ], + ) + # Set joint acceleration to 0.0 + wp.launch( + update_joint_array_with_value, + dim=(self.num_instances, self.num_joints), + inputs=[ + 0.0, + self._data._joint_acc.data, + env_mask, + joint_mask, + ], + ) + # We do not invalidate the buffers related to the body velocities since they will only be updated by the + # simulation At the next step. If this is needed then we need to call forward_kinematics first, and then + # invalidate the buffers related to the body velocities. This is not cheap, and we should only do it if + # strictly necessary. + + """ + Operations - Simulation Parameters Writers. + """ + + def write_joint_control_mode_to_sim( + self, + control_mode: wp.array | int, + joint_mask: wp.array | None = None, + env_mask: wp.array | None = None, + ): + """Write joint control mode into the simulation. + + Args: + control_mode: Joint control mode. Shape is (num_instances, num_joints). + joint_mask: The joint mask. Shape is (num_joints). + env_mask: The environment mask. Shape is (num_instances,). + + Raises: + ValueError: If the control mode is invalid. + """ + # resolve indices + if env_mask is None: + env_mask = self._ALL_ENV_MASK + if joint_mask is None: + joint_mask = self._ALL_JOINT_MASK + # set to simulation + if isinstance(control_mode, int): + wp.launch( + update_joint_array_with_value_int, + dim=(self.num_instances, self.num_joints), + inputs=[ + control_mode, + self._data.joint_control_mode_sim, + env_mask, + joint_mask, + ], + ) + else: + wp.launch( + update_joint_array_int, + dim=(self.num_instances, self.num_joints), + inputs=[ + control_mode, + self._data.joint_control_mode_sim, + env_mask, + joint_mask, + ], + ) + # tell the physics engine to use the new control mode + NewtonManager.add_model_change(SolverNotifyFlags.JOINT_DOF_PROPERTIES) + + def write_joint_stiffness_to_sim( + self, + stiffness: wp.array | float, + joint_mask: wp.array | None = None, + env_mask: wp.array | None = None, + ) -> None: + """Write joint stiffness into the simulation. + + Args: + stiffness: Joint stiffness. Shape is (num_instances, num_joints). + joint_mask: The joint mask. Shape is (num_joints). + env_mask: The environment mask. Shape is (num_instances,). + """ + # note: This function isn't setting the values for actuator models. (#128) + # resolve indices + if env_mask is None: + env_mask = self._ALL_ENV_MASK + if joint_mask is None: + joint_mask = self._ALL_JOINT_MASK + # set into internal buffers + if isinstance(stiffness, float): + wp.launch( + update_joint_array_with_value, + dim=(self.num_instances, self.num_joints), + inputs=[ + stiffness, + self._data.joint_stiffness_sim, + env_mask, + joint_mask, + ], + ) + else: + wp.launch( + update_joint_array, + dim=(self.num_instances, self.num_joints), + inputs=[ + stiffness, + self._data.joint_stiffness_sim, + env_mask, + joint_mask, + ], + ) + # tell the physics engine to use the new stiffness + NewtonManager.add_model_change(SolverNotifyFlags.JOINT_DOF_PROPERTIES) + + def write_joint_damping_to_sim( + self, + damping: wp.array | float, + joint_mask: wp.array | None = None, + env_mask: wp.array | None = None, + ) -> None: + """Write joint damping into the simulation. + + Args: + damping: Joint damping. Shape is (num_instances, num_joints). + joint_mask: The joint mask. Shape is (num_joints). + env_mask: The environment mask. Shape is (num_instances,). + """ + # note: This function isn't setting the values for actuator models. (#128) + # resolve indices + if env_mask is None: + env_mask = self._ALL_ENV_MASK + if joint_mask is None: + joint_mask = self._ALL_JOINT_MASK + # set into internal buffers + if isinstance(damping, float): + wp.launch( + update_joint_array_with_value, + dim=(self.num_instances, self.num_joints), + inputs=[ + damping, + self._data.joint_damping_sim, + env_mask, + joint_mask, + ], + ) + else: + wp.launch( + update_joint_array, + dim=(self.num_instances, self.num_joints), + inputs=[ + damping, + self._data.joint_damping_sim, + env_mask, + joint_mask, + ], + ) + # tell the physics engine to use the new damping + NewtonManager.add_model_change(SolverNotifyFlags.JOINT_DOF_PROPERTIES) + + def write_joint_position_limit_to_sim( + self, + upper_limits: wp.array | float, + lower_limits: wp.array | float, + joint_mask: wp.array | None = None, + env_mask: wp.array | None = None, + ) -> None: + """Write joint position limits into the simulation. + + Args: + upper_limits: Joint upper limits. Shape is (num_instances, num_joints). + lower_limits: Joint lower limits. Shape is (num_instances, num_joints). + joint_mask: The joint mask. Shape is (num_joints). + env_mask: The environment mask. Shape is (num_instances,). + """ + # note: This function isn't setting the values for actuator models. (#128) + # resolve indices + if env_mask is None: + env_mask = self._ALL_ENV_MASK + if joint_mask is None: + joint_mask = self._ALL_JOINT_MASK + if isinstance(upper_limits, float) and isinstance(lower_limits, float): + # update default joint pos to stay within the new limits + wp.launch( + update_joint_pos_with_limits_value_vec2f, + dim=(self.num_instances, self.num_joints), + inputs=[ + wp.vec2f(upper_limits, lower_limits), + self._data.default_joint_pos, + env_mask, + joint_mask, + ], + ) + # set into simulation + wp.launch( + update_joint_limits_value_vec2f, + dim=(self.num_instances, self.num_joints), + inputs=[ + wp.vec2f(upper_limits, lower_limits), + self.cfg.soft_joint_pos_limit_factor, + self._data.joint_pos_limits_lower, + self._data.joint_pos_limits_upper, + self._data.soft_joint_pos_limits, + env_mask, + joint_mask, + ], + ) + elif isinstance(upper_limits, wp.array) and isinstance(lower_limits, wp.array): + # update default joint pos to stay within the new limits + wp.launch( + update_joint_pos_with_limits, + dim=(self.num_instances, self.num_joints), + inputs=[ + lower_limits, + upper_limits, + self._data.default_joint_pos, + env_mask, + joint_mask, + ], + ) + # set into simulation + wp.launch( + update_joint_limits, + dim=(self.num_instances, self.num_joints), + inputs=[ + lower_limits, + upper_limits, + self.cfg.soft_joint_pos_limit_factor, + self._data.joint_pos_limits_lower, + self._data.joint_pos_limits_upper, + self._data.soft_joint_pos_limits, + env_mask, + joint_mask, + ], + ) + else: + raise NotImplementedError("Only float or wp.array of float is supported for upper and lower limits.") + # tell the physics engine to use the new limits + NewtonManager.add_model_change(SolverNotifyFlags.JOINT_DOF_PROPERTIES) + + def write_joint_velocity_limit_to_sim( + self, + limits: wp.array | float, + joint_mask: wp.array | None = None, + env_mask: wp.array | None = None, + ) -> None: + """Write joint max velocity to the simulation. + + The velocity limit is used to constrain the joint velocities in the physics engine. The joint will only + be able to reach this velocity if the joint's effort limit is sufficiently large. If the joint is moving + faster than this velocity, the physics engine will actually try to brake the joint to reach this velocity. + + .. warn:: This function is ignored when using the Mujoco solver. + + Args: + limits: Joint max velocity. Shape is (num_instances, num_joints). + joint_mask: The joint mask. Shape is (num_joints). + env_mask: The environment mask. Shape is (num_instances,). + """ + # Warn if using Mujoco solver + if isinstance(NewtonManager._solver, SolverMuJoCo): + omni.log.warn("write_joint_velocity_limit_to_sim is ignored when using the Mujoco solver.") + + # resolve indices + if env_mask is None: + env_mask = self._ALL_ENV_MASK + if joint_mask is None: + joint_mask = self._ALL_JOINT_MASK + # set into internal buffers + if isinstance(limits, float): + wp.launch( + update_joint_array_with_value, + dim=(self.num_instances, self.num_joints), + inputs=[ + limits, + self._data.joint_vel_limits, + env_mask, + joint_mask, + ], + ) + else: + wp.launch( + update_joint_array, + dim=(self.num_instances, self.num_joints), + inputs=[ + limits, + self._data.joint_vel_limits, + env_mask, + joint_mask, + ], + ) + # tell the physics engine to use the new limits + NewtonManager.add_model_change(SolverNotifyFlags.JOINT_DOF_PROPERTIES) + + def write_joint_effort_limit_to_sim( + self, + limits: wp.array | float, + joint_mask: wp.array | None = None, + env_mask: wp.array | None = None, + ) -> None: + """Write joint effort limits into the simulation. + + The effort limit is used to constrain the computed joint efforts in the physics engine. If the + computed effort exceeds this limit, the physics engine will clip the effort to this value. + + Args: + limits: Joint torque limits. Shape is (num_instances, num_joints). + joint_mask: The joint mask. Shape is (num_joints). + env_mask: The environment mask. Shape is (num_instances,). + """ + # note: This function isn't setting the values for actuator models. (#128) + # resolve indices + if env_mask is None: + env_mask = self._ALL_ENV_MASK + if joint_mask is None: + joint_mask = self._ALL_JOINT_MASK + # set into internal buffers + if isinstance(limits, float): + wp.launch( + update_joint_array_with_value, + dim=(self.num_instances, self.num_joints), + inputs=[ + limits, + self._data.joint_effort_limits, + env_mask, + joint_mask, + ], + ) + else: + wp.launch( + update_joint_array, + dim=(self.num_instances, self.num_joints), + inputs=[ + limits, + self._data.joint_effort_limits, + env_mask, + joint_mask, + ], + ) + # tell the physics engine to use the new limits + NewtonManager.add_model_change(SolverNotifyFlags.JOINT_DOF_PROPERTIES) + + def write_joint_armature_to_sim( + self, + armature: wp.array | float, + joint_mask: wp.array | None = None, + env_mask: wp.array | None = None, + ) -> None: + """Write joint armature into the simulation. + + The armature is directly added to the corresponding joint-space inertia. It helps improve the + simulation stability by reducing the joint velocities. + + Args: + armature: Joint armature. Shape is (num_instances, num_joints). + joint_mask: The joint mask. Shape is (num_joints). + env_mask: The environment mask. Shape is (num_instances,). + """ + # resolve indices + if env_mask is None: + env_mask = self._ALL_ENV_MASK + if joint_mask is None: + joint_mask = self._ALL_JOINT_MASK + # set into internal buffers + if isinstance(armature, float): + wp.launch( + update_joint_array_with_value, + dim=(self.num_instances, self.num_joints), + inputs=[ + armature, + self._data.joint_armature, + env_mask, + joint_mask, + ], + ) + else: + wp.launch( + update_joint_array, + dim=(self.num_instances, self.num_joints), + inputs=[ + armature, + self._data.joint_armature, + env_mask, + joint_mask, + ], + ) + # tell the physics engine to use the new armature + NewtonManager.add_model_change(SolverNotifyFlags.JOINT_DOF_PROPERTIES) + + def write_joint_friction_coefficient_to_sim( + self, + joint_friction_coeff: wp.array | float, + joint_mask: wp.array | None = None, + env_mask: wp.array | None = None, + ) -> None: + r"""Write joint friction coefficients into the simulation. + + The joint friction is a unitless quantity. It relates the magnitude of the spatial force transmitted + from the parent body to the child body to the maximal friction force that may be applied by the solver + to resist the joint motion. + + Mathematically, this means that: :math:`F_{resist} \leq \mu F_{spatial}`, where :math:`F_{resist}` + is the resisting force applied by the solver and :math:`F_{spatial}` is the spatial force + transmitted from the parent body to the child body. The simulated friction effect is therefore + similar to static and Coulomb friction. + + Args: + joint_friction_coeff: Joint friction coefficients. Shape is (num_instances, num_joints). + joint_mask: The joint mask. Shape is (num_joints). + env_mask: The environment mask. Shape is (num_instances,). + """ + # resolve indices + if env_mask is None: + env_mask = self._ALL_ENV_MASK + if joint_mask is None: + joint_mask = self._ALL_JOINT_MASK + # set into internal buffers + if isinstance(joint_friction_coeff, float): + wp.launch( + update_joint_array_with_value, + dim=(self.num_instances, self.num_joints), + inputs=[ + joint_friction_coeff, + self._data.joint_friction_coeff, + env_mask, + joint_mask, + ], + ) + else: + wp.launch( + update_joint_array, + dim=(self.num_instances, self.num_joints), + inputs=[ + joint_friction_coeff, + self._data.joint_friction_coeff, + env_mask, + joint_mask, + ], + ) + # tell the physics engine to use the new friction + NewtonManager.add_model_change(SolverNotifyFlags.JOINT_DOF_PROPERTIES) + + """ + Operations - Setters. + """ + + def set_joint_position_target( + self, + target: wp.array, + joint_mask: wp.array | None = None, + env_mask: wp.array | None = None, + ) -> None: + """Set joint position targets into internal buffers. + + This function does not apply the joint targets to the simulation. It only fills the buffers with + the desired values. To apply the joint targets, call the :meth:`write_data_to_sim` function. + + Args: + target: Joint position targets. Shape is (num_instances, num_joints). + joint_mask: The joint mask. Shape is (num_joints). + env_mask: The environment mask. Shape is (num_instances,). + """ + # resolve indices + if env_mask is None: + env_mask = self._ALL_ENV_MASK + if joint_mask is None: + joint_mask = self._ALL_JOINT_MASK + # set targets + wp.launch( + update_joint_array, + dim=(self.num_instances, self.num_joints), + inputs=[ + target, + self._data.joint_target, + env_mask, + joint_mask, + ], + ) + + def set_joint_velocity_target( + self, + target: wp.array, + joint_mask: wp.array | None = None, + env_mask: wp.array | None = None, + ) -> None: + """Set joint velocity targets into internal buffers. + + This function does not apply the joint targets to the simulation. It only fills the buffers with + the desired values. To apply the joint targets, call the :meth:`write_data_to_sim` function. + + Args: + target: Joint velocity targets. Shape is (num_instances, num_joints). + joint_mask: The joint mask. Shape is (num_joints). + env_mask: The environment mask. Shape is (num_instances,). + """ + # resolve indice + if env_mask is None: + env_mask = self._ALL_ENV_MASK + if joint_mask is None: + joint_mask = self._ALL_JOINT_MASK + # set targets + wp.launch( + update_joint_array, + dim=(self.num_instances, self.num_joints), + inputs=[ + target, + self._data.joint_target, + env_mask, + joint_mask, + ], + ) + + def set_joint_effort_target( + self, + target: wp.array, + joint_mask: wp.array | None = None, + env_mask: wp.array | None = None, + ) -> None: + """Set joint efforts into internal buffers. + + This function does not apply the joint targets to the simulation. It only fills the buffers with + the desired values. To apply the joint targets, call the :meth:`write_data_to_sim` function. + + Args: + target: Joint effort targets. Shape is (num_instances, num_joints). + joint_mask: The joint mask. Shape is (num_joints). + env_mask: The environment mask. Shape is (num_instances,). + """ + # resolve indices + if env_mask is None: + env_mask = self._ALL_ENV_MASK + if joint_mask is None: + joint_mask = self._ALL_JOINT_MASK + # set targets + wp.launch( + update_joint_array, + dim=(self.num_instances, self.num_joints), + inputs=[ + target, + self._data.joint_effort_target, + env_mask, + joint_mask, + ], + ) + + def _create_buffers(self): + # constants + self._ALL_ENV_MASK = wp.ones((self.num_instances,), dtype=wp.bool, device=self._device) + self._ALL_JOINT_MASK = wp.ones((self.num_joints,), dtype=wp.bool, device=self._device) + # masks + self._ENV_MASK = wp.zeros((self.num_instances,), dtype=wp.bool, device=self._device) + self._JOINT_MASK = wp.zeros((self.num_joints,), dtype=wp.bool, device=self._device) + + wp.launch( + update_soft_joint_pos_limits, + dim=(self.num_instances, self.num_joints), + inputs=[ + self.data.joint_pos_limits_lower, + self.data.joint_pos_limits_upper, + self.data.soft_joint_pos_limits, + self._soft_joint_pos_limit_factor, + ], + ) \ No newline at end of file diff --git a/source/isaaclab/isaaclab/newton/assets/core/joint_properties/joint_data.py b/source/isaaclab/isaaclab/newton/assets/core/joint_properties/joint_data.py new file mode 100644 index 00000000000..15af86cb7a8 --- /dev/null +++ b/source/isaaclab/isaaclab/newton/assets/core/joint_properties/joint_data.py @@ -0,0 +1,411 @@ +# Copyright (c) 2022-2025, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +import weakref + +import warp as wp + +from newton.selection import ArticulationView as NewtonArticulationView +from isaaclab.sim._impl.newton_manager import NewtonManager +from isaaclab.utils.buffers import TimestampedWarpBuffer +from isaaclab.utils.helpers import deprecated, warn_overhead_cost + +from isaaclab.newton.assets.core.kernels import ( + derive_joint_acceleration_from_velocity, + make_joint_pos_limits_from_lower_and_upper_limits, +) + + +class JointData: + def __init__(self, root_newton_view: NewtonArticulationView, device: str) -> None: + """Initializes the container for joint data. + + Args: + root_newton_view: The root articulation view. + device: The device used for processing. + """ + # Set the parameters + self.device = device + # Set the root articulation view + # note: this is stored as a weak reference to avoid circular references between the asset class + # and the data container. This is important to avoid memory leaks. + self._root_newton_view: NewtonArticulationView = weakref.proxy(root_newton_view) + + # Set initial time stamp + self._sim_timestamp = 0.0 + + self._create_simulation_bindings() + self._create_buffers() + + + def update(self, dt: float): + # update the simulation timestamp + self._sim_timestamp += dt + # Trigger an update of the joint acceleration buffer at a higher frequency + # since we do finite differencing. + self.joint_acc + + def _create_simulation_bindings(self): + """Create simulation bindings for the joint data. + + Direct simulation bindings are pointers to the simulation data, their data is not copied, and should + only be updated using warp kernels. Any modifications made to the bindings will be reflected in the simulation. + Hence we encourage users to carefully think about the data they modify and in which order it should be updated. + + .. caution:: This is possible if and only if the properties that we access are strided from newton and not + indexed. Newton willing this is the case all the time, but we should pay attention to this if things look off. + """ + + # -- joint properties + self._sim_bind_joint_pos_limits_lower = self._root_newton_view.get_attribute( + "joint_limit_lower", NewtonManager.get_model() + ) + self._sim_bind_joint_pos_limits_upper = self._root_newton_view.get_attribute( + "joint_limit_upper", NewtonManager.get_model() + ) + self._sim_bind_joint_stiffness_sim = self._root_newton_view.get_attribute( + "joint_target_ke", NewtonManager.get_model() + ) + self._sim_bind_joint_damping_sim = self._root_newton_view.get_attribute( + "joint_target_kd", NewtonManager.get_model() + ) + self._sim_bind_joint_armature = self._root_newton_view.get_attribute( + "joint_armature", NewtonManager.get_model() + ) + self._sim_bind_joint_friction_coeff = self._root_newton_view.get_attribute( + "joint_friction", NewtonManager.get_model() + ) + self._sim_bind_joint_vel_limits_sim = self._root_newton_view.get_attribute( + "joint_velocity_limit", NewtonManager.get_model() + ) + self._sim_bind_joint_effort_limits_sim = self._root_newton_view.get_attribute( + "joint_effort_limit", NewtonManager.get_model() + ) + self._sim_bind_joint_control_mode_sim = self._root_newton_view.get_attribute( + "joint_dof_mode", NewtonManager.get_model() + ) + # -- joint states + self._sim_bind_joint_pos = self._root_newton_view.get_dof_positions(NewtonManager.get_state_0()) + self._sim_bind_joint_vel = self._root_newton_view.get_dof_velocities(NewtonManager.get_state_0()) + # -- joint commands (sent to the simulation) + self._sim_bind_joint_effort = self._root_newton_view.get_attribute("joint_f", NewtonManager.get_control()) + self._sim_bind_joint_target = self._root_newton_view.get_attribute( + "joint_target", NewtonManager.get_control() + ) + + def _create_buffers(self): + """Create buffers for the joint data.""" + # -- default joint positions and velocities + self._default_joint_pos = wp.zeros( + (self._root_newton_view.count, self._root_newton_view.joint_dof_count), dtype=wp.float32, device=self.device + ) + self._default_joint_vel = wp.zeros( + (self._root_newton_view.count, self._root_newton_view.joint_dof_count), dtype=wp.float32, device=self.device + ) + # -- joint commands (sent to the actuator from the user) + self._joint_target = wp.zeros( + (self._root_newton_view.count, self._root_newton_view.joint_dof_count), dtype=wp.float32, device=self.device + ) + self._joint_effort_target = wp.zeros( + (self._root_newton_view.count, self._root_newton_view.joint_dof_count), dtype=wp.float32, device=self.device + ) + # -- computed joint efforts from the actuator models + self._computed_effort = wp.zeros( + (self._root_newton_view.count, self._root_newton_view.joint_dof_count), dtype=wp.float32, device=self.device + ) + self._applied_effort = wp.zeros( + (self._root_newton_view.count, self._root_newton_view.joint_dof_count), dtype=wp.float32, device=self.device + ) + # -- joint properties for the actuator models + self._joint_stiffness = wp.clone(self._sim_bind_joint_stiffness_sim) + self._joint_damping = wp.clone(self._sim_bind_joint_damping_sim) + self._joint_control_mode = wp.clone(self._sim_bind_joint_control_mode_sim) + # -- other data that are filled based on explicit actuator models + self._joint_dynamic_friction = wp.zeros( + (self._root_newton_view.count, self._root_newton_view.joint_dof_count), dtype=wp.float32, device=self.device + ) + self._joint_viscous_friction = wp.zeros( + (self._root_newton_view.count, self._root_newton_view.joint_dof_count), dtype=wp.float32, device=self.device + ) + self._soft_joint_vel_limits = wp.zeros( + (self._root_newton_view.count, self._root_newton_view.joint_dof_count), dtype=wp.float32, device=self.device + ) + self._gear_ratio = wp.ones( + (self._root_newton_view.count, self._root_newton_view.joint_dof_count), dtype=wp.float32, device=self.device + ) + # -- update the soft joint position limits + self._soft_joint_pos_limits = wp.zeros( + (self._root_newton_view.count, self._root_newton_view.joint_dof_count), dtype=wp.vec2f, device=self.device + ) + + # Initialize history for finite differencing + self._previous_joint_vel = wp.clone(self._root_newton_view.get_dof_velocities(NewtonManager.get_state_0())) + + # Initialize the lazy buffers. + # -- joint state + self._joint_acc = TimestampedWarpBuffer( + shape=(self._root_newton_view.count, self._root_newton_view.joint_dof_count), dtype=wp.float32 + ) + # self._body_incoming_joint_wrench_b = TimestampedWarpBuffer() + + ## + # Direct simulation bindings accessors. + ## + + @property + def joint_control_mode_sim(self) -> wp.array: + return self._sim_bind_joint_control_mode_sim + + @property + def joint_stiffness_sim(self) -> wp.array: + """Joint stiffness as defined in the simulation. Shape is (num_instances, num_joints).""" + return self._sim_bind_joint_stiffness_sim + + @property + def joint_damping_sim(self) -> wp.array: + """Joint damping as defined in the simulation. Shape is (num_instances, num_joints).""" + return self._sim_bind_joint_damping_sim + + @property + def joint_armature(self) -> wp.array: + """Joint armature as defined in the simulation. Shape is (num_instances, num_joints).""" + return self._sim_bind_joint_armature + + @property + def joint_friction_coeff(self) -> wp.array: + """Joint friction coefficient as defined in the simulation. Shape is (num_instances, num_joints).""" + return self._sim_bind_joint_friction_coeff + + @property + def joint_pos_limits_lower(self) -> wp.array: + """Joint position limits lower as defined in the simulation. Shape is (num_instances, num_joints).""" + return self._sim_bind_joint_pos_limits_lower + + @property + def joint_pos_limits_upper(self) -> wp.array: + """Joint position limits upper as defined in the simulation. Shape is (num_instances, num_joints).""" + return self._sim_bind_joint_pos_limits_upper + + @property + def joint_vel_limits(self) -> wp.array: + """Joint velocity limits as defined in the simulation. Shape is (num_instances, num_joints).""" + return self._sim_bind_joint_vel_limits_sim + + @property + def joint_effort_limits(self) -> wp.array: + """Joint effort limits as defined in the simulation. Shape is (num_instances, num_joints).""" + return self._sim_bind_joint_effort_limits_sim + + @property + def joint_pos(self) -> wp.array: + """Joint posiitons. Shape is (num_instances, num_joints).""" + return self._sim_bind_joint_pos + + @property + def joint_vel(self) -> wp.array: + """Joint velocities. Shape is (num_instances, num_joints).""" + return self._sim_bind_joint_vel + + @property + def joint_target_sim(self) -> wp.array: + """Joint targets in the simulation. Shape is (num_instances, num_joints).""" + return self._sim_bind_joint_target + + @property + def joint_effort_sim(self) -> wp.array: + """Joint effort in the simulation. Shape is (num_instances, num_joints).""" + return self._sim_bind_joint_effort + + ### + # Buffers accessors. + ## + + @property + def joint_target(self) -> wp.array: + """Joint targets commanded by the user. Shape is (num_instances, num_joints).""" + return self._joint_target + + @property + def joint_effort_target(self) -> wp.array: + """Joint effort targets commanded by the user. Shape is (num_instances, num_joints).""" + return self._joint_effort_target + + @property + def computed_effort(self) -> wp.array: + """Joint efforts computed from the actuator model (before clipping). Shape is (num_instances, num_joints).""" + return self._computed_effort + + @property + def applied_effort(self) -> wp.array: + """Joint efforts applied from the actuator model (after clipping). Shape is (num_instances, num_joints).""" + return self._applied_effort + + @property + def joint_stiffness(self) -> wp.array: + """Joint stiffness as defined in the actuator model. Shape is (num_instances, num_joints).""" + return self._joint_stiffness + + @property + def joint_damping(self) -> wp.array: + """Joint damping as defined in the actuator model. Shape is (num_instances, num_joints).""" + return self._joint_damping + + @property + def joint_control_mode(self) -> wp.array: + """Joint control mode as defined in the actuator model. Shape is (num_instances, num_joints).""" + return self._joint_control_mode + + @property + def joint_dynamic_friction(self) -> wp.array: + """Joint dynamic friction as defined in the actuator model. Shape is (num_instances, num_joints).""" + return self._joint_dynamic_friction + + @property + def joint_viscous_friction(self) -> wp.array: + """Joint viscous friction as defined in the actuator model. Shape is (num_instances, num_joints).""" + return self._joint_viscous_friction + + @property + def soft_joint_vel_limits(self) -> wp.array: + """Soft joint velocity limits for all joints. Shape is (num_instances, num_joints). + + These are obtained from the actuator model. It may differ from :attr:`joint_vel_limits` if the actuator model + has a variable velocity limit model. For instance, in a variable gear ratio actuator model. + """ + return self._soft_joint_vel_limits + + @property + def gear_ratio(self) -> wp.array: + """Gear ratio as defined in the actuator model. Shape is (num_instances, num_joints).""" + return self._gear_ratio + + @property + def soft_joint_pos_limits(self) -> wp.array: + r"""Soft joint positions limits for all joints. Shape is (num_instances, num_joints, 2). + + The limits are in the order :math:`[lower, upper]`.The soft joint position limits are computed as + a sub-region of the :attr:`joint_pos_limits` based on the + :attr:`~isaaclab.assets.ArticulationCfg.soft_joint_pos_limit_factor` parameter. + + Consider the joint position limits :math:`[lower, upper]` and the soft joint position limits + :math:`[soft_lower, soft_upper]`. The soft joint position limits are computed as: + + .. math:: + + soft\_lower = (lower + upper) / 2 - factor * (upper - lower) / 2 + soft\_upper = (lower + upper) / 2 + factor * (upper - lower) / 2 + + The soft joint position limits help specify a safety region around the joint limits. It isn't used by the + simulation, but is useful for learning agents to prevent the joint positions from violating the limits. + """ + return self._soft_joint_pos_limits + + ## + # Default accessors. + ## + + @property + def default_joint_pos(self) -> wp.array: + """Default joint positions of all joints. Shape is (num_instances, num_joints). + + This quantity is configured through the :attr:`isaaclab.assets.ArticulationCfg.init_state` parameter. + """ + return self._default_joint_pos + + @property + def default_joint_vel(self) -> wp.array: + """Default joint velocities of all joints. Shape is (num_instances, num_joints). + + This quantity is configured through the :attr:`isaaclab.assets.ArticulationCfg.init_state` parameter. + """ + return self._default_joint_vel + + @default_joint_pos.setter + def default_joint_pos(self, value: wp.array) -> None: + self._default_joint_pos = value + + @default_joint_vel.setter + def default_joint_vel(self, value: wp.array) -> None: + self._default_joint_vel = value + + ## + # Joint state properties. + ## + + @property + def joint_acc(self) -> wp.array: + """Joint acceleration of all joints. Shape is (num_instances, num_joints).""" + if self._joint_acc.timestamp < self._sim_timestamp: + # note: we use finite differencing to compute acceleration + wp.launch( + derive_joint_acceleration_from_velocity, + dim=(self._root_newton_view.count, self._root_newton_view.joint_dof_count), + inputs=[ + self._sim_bind_joint_vel, + self._previous_joint_vel, + NewtonManager.get_dt(), + self._joint_acc.data, + ], + ) + self._joint_acc.timestamp = self._sim_timestamp + return self._joint_acc.data + + @property + def body_incoming_joint_wrench_b(self) -> wp.array: + """Joint reaction wrench applied from body parent to child body in parent body frame. + + Shape is (num_instances, num_bodies, 6). All body reaction wrenches are provided including the root body to the + world of an articulation. + + For more information on joint wrenches, please check the`PhysX documentation `__ + and the underlying `PhysX Tensor API `__ . + """ + raise NotImplementedError("Body incoming joint wrench in body frame is not implemented for Newton.") + if self._body_incoming_joint_wrench_b.timestamp < self._sim_timestamp: + self._body_incoming_joint_wrench_b.data = self._root_physx_view.get_link_incoming_joint_force() + self._body_incoming_joint_wrench_b.time_stamp = self._sim_timestamp + return self._body_incoming_joint_wrench_b.data + + ## + # Sliced properties. + ## + + @property + @warn_overhead_cost( + "joint_pos_limits_lower or joint_pos_limits_upper", + "This function combines both the lower and upper limits into a single array, use it only if necessary.", + ) + def joint_pos_limits(self) -> wp.array: + """Joint position limits provided to the simulation. Shape is (num_instances, num_joints, 2). + + The limits are in the order :math:`[lower, upper]`. + """ + out = wp.zeros((self._root_newton_view.count, self._root_newton_view.joint_dof_count), dtype=wp.vec2f, device=self.device) + wp.launch( + make_joint_pos_limits_from_lower_and_upper_limits, + dim=(self._root_newton_view.count, self._root_newton_view.joint_dof_count), + inputs=[ + self._sim_bind_joint_pos_limits_lower, + self._sim_bind_joint_pos_limits_upper, + out, + ], + ) + return out + + ## + # Backward compatibility. Need to nuke these properties in a future release. + ## + + @property + @deprecated("joint_pos_limits") + def joint_limits(self) -> wp.array: + """Deprecated property. Please use :attr:`joint_pos_limits` instead.""" + return self.joint_pos_limits + + @property + @deprecated("joint_friction_coeff") + def joint_friction(self) -> wp.array: + """Deprecated property. Please use :attr:`joint_friction_coeff` instead.""" + return self._sim_bind_joint_friction_coeff diff --git a/source/isaaclab/isaaclab/newton/assets/core/kernels/__init__.py b/source/isaaclab/isaaclab/newton/assets/core/kernels/__init__.py new file mode 100644 index 00000000000..e4ffab13983 --- /dev/null +++ b/source/isaaclab/isaaclab/newton/assets/core/kernels/__init__.py @@ -0,0 +1,10 @@ +# Copyright (c) 2022-2025, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +from .joint_kernels import * # noqa: F401, F403 +from .other_kernels import * # noqa: F401, F403 +from .state_kernels import * # noqa: F401, F403 +from .velocity_kernels import * # noqa: F401, F403 +from .update import * # noqa: F401, F403 \ No newline at end of file diff --git a/source/isaaclab/isaaclab/newton/assets/core/kernels/joint_kernels.py b/source/isaaclab/isaaclab/newton/assets/core/kernels/joint_kernels.py new file mode 100644 index 00000000000..be983ebaca8 --- /dev/null +++ b/source/isaaclab/isaaclab/newton/assets/core/kernels/joint_kernels.py @@ -0,0 +1,447 @@ +# Copyright (c) 2022-2025, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +import warp as wp + +""" +Helper kernels for updating joint data. +""" + + +@wp.kernel +def update_joint_array( + new_data: wp.array2d(dtype=wp.float32), + joint_data: wp.array2d(dtype=wp.float32), + env_mask: wp.array(dtype=bool), + joint_mask: wp.array(dtype=bool), +): + """ + Update the joint data for the given environment and joint indices from the newton data. + + .. note:: The :arg:`env_mask` length must be equal to the number of instances in the newton data. + The :arg:`joint_mask` length must be equal to the number of joints in the newton data. The :arg:`new_data` shape + must match the :arg:`joint_data` shape. + + Args: + new_data: The new data to update the joint data with. Shape is (num_instances, num_joints). + joint_data: The joint data to update. Shape is (num_instances, num_joints). (modified) + env_mask: The environment mask to update the joint data for. Shape is (num_instances,). + joint_mask: The joint mask to update the joint data for. Shape is (num_joints,). + """ + env_index, joint_index = wp.tid() + if env_mask[env_index] and joint_mask[joint_index]: + joint_data[env_index, joint_index] = new_data[env_index, joint_index] + + +@wp.kernel +def update_joint_array_int( + new_data: wp.array2d(dtype=wp.int32), + joint_data: wp.array2d(dtype=wp.int32), + env_mask: wp.array(dtype=bool), + joint_mask: wp.array(dtype=bool), +): + """ + Update the joint data for the given environment and joint indices from the newton data. + + .. note:: The :arg:`env_mask` length must be equal to the number of instances in the newton data. + The :arg:`joint_mask` length must be equal to the number of joints in the newton data. The :arg:`new_data` shape + must match the :arg:`joint_data` shape. + + Args: + new_data: The new data to update the joint data with. Shape is (num_instances, num_joints). + joint_data: The joint data to update. Shape is (num_instances, num_joints). (modified) + env_mask: The environment mask to update the joint data for. Shape is (num_instances,). + joint_mask: The joint mask to update the joint data for. Shape is (num_joints,). + """ + env_index, joint_index = wp.tid() + if env_mask[env_index] and joint_mask[joint_index]: + joint_data[env_index, joint_index] = new_data[env_index, joint_index] + + +@wp.kernel +def update_joint_array_with_value_array( + value: wp.array(dtype=wp.float32), + joint_data: wp.array2d(dtype=wp.float32), + env_mask: wp.array(dtype=bool), + joint_mask: wp.array(dtype=bool), +): + """Update the joint data for the given environment and joint indices with a value array. + + .. note:: The :arg:`env_mask` length must be equal to the number of instances in the newton data. + The :arg:`joint_mask` length must be equal to the number of joints in the newton data. The :arg:`value` shape + must be (num_joints,). + + Args: + value: The value array to update the joint data with. Shape is (num_joints,). + joint_data: The joint data to update. Shape is (num_instances, num_joints). (modified) + env_mask: The environment mask to update the joint data for. Shape is (num_instances,). + joint_mask: The joint mask to update the joint data for. Shape is (num_joints,). + """ + env_index, joint_index = wp.tid() + if env_mask[env_index] and joint_mask[joint_index]: + joint_data[env_index, joint_index] = value[joint_index] + + +@wp.kernel +def update_joint_array_with_value( + value: wp.float32, + joint_data: wp.array2d(dtype=wp.float32), + env_mask: wp.array(dtype=bool), + joint_mask: wp.array(dtype=bool), +): + """Update the joint data for the given environment and joint indices with a value. + + .. note:: The :arg:`env_mask` length must be equal to the number of instances in the newton data. + The :arg:`joint_mask` length must be equal to the number of joints in the newton data. The :arg:`joint_data` shape + must be (num_instances, num_joints). + + Args: + value: The value to update the joint data with. + joint_data: The joint data to update. Shape is (num_instances, num_joints). (modified) + env_mask: The environment mask to update the joint data for. Shape is (num_instances,). + joint_mask: The joint mask to update the joint data for. Shape is (num_joints,). + """ + env_index, joint_index = wp.tid() + if env_mask[env_index] and joint_mask[joint_index]: + joint_data[env_index, joint_index] = value + + +@wp.kernel +def update_joint_array_with_value_int( + value: wp.int32, + joint_data: wp.array2d(dtype=wp.int32), + env_mask: wp.array(dtype=bool), + joint_mask: wp.array(dtype=bool), +): + """Update the joint data for the given environment and joint indices with a value. + + .. note:: The :arg:`env_mask` length must be equal to the number of instances in the newton data. + The :arg:`joint_mask` length must be equal to the number of joints in the newton data. The :arg:`joint_data` shape + must be (num_instances, num_joints). + + Args: + value: The value to update the joint data with. + joint_data: The joint data to update. Shape is (num_instances, num_joints). (modified) + env_mask: The environment mask to update the joint data for. Shape is (num_instances,). + joint_mask: The joint mask to update the joint data for. Shape is (num_joints,). + """ + env_index, joint_index = wp.tid() + if env_mask[env_index] and joint_mask[joint_index]: + joint_data[env_index, joint_index] = value + + +""" +Kernels to update joint limits. +""" + + +@wp.func +def get_soft_joint_limits(lower_limit: float, upper_limit: float, soft_factor: float) -> wp.vec2f: + """Get the soft joint limits for the given lower and upper limits and soft factor. + + Args: + lower_limit: The lower limit of the joint. + upper_limit: The upper limit of the joint. + soft_factor: The soft factor to use for the joint limits. + + Returns: + The soft joint limits. Shape is (2,). + """ + mean = (lower_limit + upper_limit) / 2.0 + range = upper_limit - lower_limit + lower_limit = mean - 0.5 * range * soft_factor + upper_limit = mean + 0.5 * range * soft_factor + return wp.vec2f(lower_limit, upper_limit) + + +@wp.kernel +def update_joint_limits( + new_limits_lower: wp.array2d(dtype=wp.float32), + new_limits_upper: wp.array2d(dtype=wp.float32), + soft_factor: float, + lower_limits: wp.array2d(dtype=wp.float32), + upper_limits: wp.array2d(dtype=wp.float32), + soft_joint_limits: wp.array2d(dtype=wp.vec2f), + env_mask: wp.array(dtype=bool), + joint_mask: wp.array(dtype=bool), +): + """Update the joint limits for the given environment and joint indices. + + .. note:: The :arg:`env_mask` length must be equal to the number of instances in the newton data. + The :arg:`joint_mask` length must be equal to the number of joints in the newton data. + + Args: + new_limits_lower: The new lower limits to update the joint limits with. Shape is (num_instances, num_joints). + new_limits_upper: The new upper limits to update the joint limits with. Shape is (num_instances, num_joints). + soft_factor: The soft factor to use for the soft joint limits. + lower_limits: The lower limits to update the joint limits with. Shape is (num_instances, num_joints). (modified) + upper_limits: The upper limits to update the joint limits with. Shape is (num_instances, num_joints). (modified) + soft_joint_limits: The soft joint limits to update. Shape is (num_instances, num_joints). (modified) + env_mask: The environment mask to update the joint limits for. Shape is (num_instances,). + joint_mask: The joint mask to update the joint limits for. Shape is (num_joints,). + """ + env_index, joint_index = wp.tid() + if env_mask[env_index] and joint_mask[joint_index]: + lower_limits[env_index, joint_index] = new_limits_lower[env_index, joint_index] + upper_limits[env_index, joint_index] = new_limits_upper[env_index, joint_index] + + soft_joint_limits[env_index, joint_index] = get_soft_joint_limits( + lower_limits[env_index, joint_index], upper_limits[env_index, joint_index], soft_factor + ) + + +@wp.kernel +def update_joint_limits_with_value( + new_limits: float, + soft_factor: float, + lower_limits: wp.array2d(dtype=wp.float32), + upper_limits: wp.array2d(dtype=wp.float32), + soft_joint_limits: wp.array2d(dtype=wp.vec2f), + env_mask: wp.array(dtype=bool), + joint_mask: wp.array(dtype=bool), +): + """Update the joint limits for the given environment and joint indices with a value. + + .. note:: The :arg:`env_mask` length must be equal to the number of instances in the newton data. + The :arg:`joint_mask` length must be equal to the number of joints in the newton data. + + Args: + new_limits: The new limits to update the joint limits with. + soft_factor: The soft factor to use for the soft joint limits. + lower_limits: The lower limits to update the joint limits with. Shape is (num_instances, num_joints). (modified) + upper_limits: The upper limits to update the joint limits with. Shape is (num_instances, num_joints). (modified) + soft_joint_limits: The soft joint limits to update. Shape is (num_instances, num_joints). (modified) + env_mask: The environment mask to update the joint limits for. Shape is (num_instances,). + joint_mask: The joint mask to update the joint limits for. Shape is (num_joints,). + """ + env_index, joint_index = wp.tid() + if env_mask[env_index] and joint_mask[joint_index]: + lower_limits[env_index, joint_index] = new_limits + upper_limits[env_index, joint_index] = new_limits + + soft_joint_limits[env_index, joint_index] = get_soft_joint_limits( + lower_limits[env_index, joint_index], upper_limits[env_index, joint_index], soft_factor + ) + + +@wp.kernel +def update_joint_limits_value_vec2f( + new_limits: wp.vec2f, + soft_factor: float, + lower_limits: wp.array2d(dtype=wp.float32), + upper_limits: wp.array2d(dtype=wp.float32), + soft_joint_limits: wp.array2d(dtype=wp.vec2f), + env_mask: wp.array(dtype=bool), + joint_mask: wp.array(dtype=bool), +): + """Update the joint limits for the given environment and joint indices with a value. + + Args: + new_limits: The new limits to update the joint limits with. + soft_factor: The soft factor to use for the soft joint limits. + lower_limits: The lower limits to update the joint limits with. Shape is (num_instances, num_joints). (modified) + upper_limits: The upper limits to update the joint limits with. Shape is (num_instances, num_joints). (modified) + soft_joint_limits: The soft joint limits to update. Shape is (num_instances, num_joints). (modified) + env_mask: The environment mask to update the joint limits for. Shape is (num_instances,). + joint_mask: The joint mask to update the joint limits for. Shape is (num_joints,). + """ + env_index, joint_index = wp.tid() + if env_mask[env_index] and joint_mask[joint_index]: + lower_limits[env_index, joint_index] = new_limits[0] + upper_limits[env_index, joint_index] = new_limits[1] + + soft_joint_limits[env_index, joint_index] = get_soft_joint_limits( + lower_limits[env_index, joint_index], upper_limits[env_index, joint_index], soft_factor + ) + + +""" +Kernels to update joint position from joint limits. +""" + + +@wp.kernel +def update_joint_pos_with_limits( + joint_pos_limits_lower: wp.array2d(dtype=wp.float32), + joint_pos_limits_upper: wp.array2d(dtype=wp.float32), + joint_pos: wp.array2d(dtype=wp.float32), + env_mask: wp.array(dtype=bool), + joint_mask: wp.array(dtype=bool), +): + """Update the joint position for the given environment and joint indices with the limits. + + .. note:: The :arg:`env_mask` length must be equal to the number of instances in the newton data. + The :arg:`joint_mask` length must be equal to the number of joints in the newton data. + + Args: + joint_pos_limits_lower: The lower limits to update the joint position with. Shape is (num_instances, num_joints). + joint_pos_limits_upper: The upper limits to update the joint position with. Shape is (num_instances, num_joints). + joint_pos: The joint position to update. Shape is (num_instances, num_joints). (modified) + env_mask: The environment mask to update the joint position for. Shape is (num_instances,). + joint_mask: The joint mask to update the joint position for. Shape is (num_joints,). + """ + env_index, joint_index = wp.tid() + if env_mask[env_index] and joint_mask[joint_index]: + joint_pos[env_index, joint_index] = wp.clamp( + joint_pos[env_index, joint_index], + joint_pos_limits_lower[env_index, joint_index], + joint_pos_limits_upper[env_index, joint_index], + ) + + +@wp.kernel +def update_joint_pos_with_limits_value( + joint_pos_limits: float, + joint_pos: wp.array2d(dtype=wp.float32), + env_mask: wp.array(dtype=bool), + joint_mask: wp.array(dtype=bool), +): + """Update the joint position for the given environment and joint indices with the limits. + + .. note:: The :arg:`env_mask` length must be equal to the number of instances in the newton data. + The :arg:`joint_mask` length must be equal to the number of joints in the newton data. + + Args: + joint_pos_limits: The joint position limits to update. + joint_pos: The joint position to update. Shape is (num_instances, num_joints). (modified) + env_mask: The environment mask to update the joint position for. Shape is (num_instances,). + joint_mask: The joint mask to update the joint position for. Shape is (num_joints,). + """ + env_index, joint_index = wp.tid() + if env_mask[env_index] and joint_mask[joint_index]: + joint_pos[env_index, joint_index] = wp.clamp( + joint_pos[env_index, joint_index], joint_pos_limits, joint_pos_limits + ) + + +@wp.kernel +def update_joint_pos_with_limits_value_vec2f( + joint_pos_limits: wp.vec2f, + joint_pos: wp.array2d(dtype=wp.float32), + env_mask: wp.array(dtype=bool), + joint_mask: wp.array(dtype=bool), +): + """Update the joint position for the given environment and joint indices with the limits. + + .. note:: The :arg:`env_mask` length must be equal to the number of instances in the newton data. + The :arg:`joint_mask` length must be equal to the number of joints in the newton data. + + Args: + joint_pos_limits: The joint position limits to update. Shape is (2,) + joint_pos: The joint position to update. Shape is (num_instances, num_joints). (modified) + env_mask: The environment mask to update the joint position for. Shape is (num_instances,). + joint_mask: The joint mask to update the joint position for. Shape is (num_joints,). + """ + env_index, joint_index = wp.tid() + if env_mask[env_index] and joint_mask[joint_index]: + joint_pos[env_index, joint_index] = wp.clamp( + joint_pos[env_index, joint_index], joint_pos_limits[0], joint_pos_limits[1] + ) + + +""" +Helper kernel to reconstruct limits +""" + + +@wp.kernel +def make_joint_pos_limits_from_lower_and_upper_limits( + lower_limits: wp.array2d(dtype=wp.float32), + upper_limits: wp.array2d(dtype=wp.float32), + joint_pos_limits: wp.array2d(dtype=wp.vec2f), +): + """Make the joint position limits from the lower and upper limits. + + Args: + lower_limits: The lower limits to make the joint position limits from. Shape is (num_instances, num_joints). + upper_limits: The upper limits to make the joint position limits from. Shape is (num_instances, num_joints). + joint_pos_limits: The joint position limits to make. Shape is (num_instances, num_joints, 2). (destination) + """ + env_index, joint_index = wp.tid() + joint_pos_limits[env_index, joint_index] = wp.vec2f( + lower_limits[env_index, joint_index], upper_limits[env_index, joint_index] + ) + + +""" +Helper kernel to update soft joint position limits. +""" + + +@wp.kernel +def update_soft_joint_pos_limits( + joint_pos_limits_lower: wp.array2d(dtype=wp.float32), + joint_pos_limits_upper: wp.array2d(dtype=wp.float32), + soft_joint_pos_limits: wp.array2d(dtype=wp.vec2f), + soft_factor: float, +): + """Update the soft joint position limits for the given environment and joint indices. + + Args: + joint_pos_limits_lower: The lower limits to update the soft joint position limits with. Shape is (num_instances, num_joints). + joint_pos_limits_upper: The upper limits to update the soft joint position limits with. Shape is (num_instances, num_joints). + soft_joint_pos_limits: The soft joint position limits to update. Shape is (num_instances, num_joints). (modified) + soft_factor: The soft factor to use for the soft joint position limits. + """ + env_index, joint_index = wp.tid() + soft_joint_pos_limits[env_index, joint_index] = get_soft_joint_limits( + joint_pos_limits_lower[env_index, joint_index], joint_pos_limits_upper[env_index, joint_index], soft_factor + ) + + +""" +Kernels to derive joint acceleration from velocity. +""" + + +@wp.kernel +def derive_joint_acceleration_from_velocity( + joint_velocity: wp.array2d(dtype=wp.float32), + previous_joint_velocity: wp.array2d(dtype=wp.float32), + dt: float, + joint_acceleration: wp.array2d(dtype=wp.float32), +): + """ + Derive the joint acceleration from the velocity. + + Args: + joint_velocity: The joint velocity. Shape is (num_instances, num_joints). + previous_joint_velocity: The previous joint velocity. Shape is (num_instances, num_joints). + dt: The time step. + joint_acceleration: The joint acceleration. Shape is (num_instances, num_joints). (modified) + """ + env_index, joint_index = wp.tid() + # compute acceleration + joint_acceleration[env_index, joint_index] = ( + joint_velocity[env_index, joint_index] - previous_joint_velocity[env_index, joint_index] + ) / dt + + # update previous velocity + previous_joint_velocity[env_index, joint_index] = joint_velocity[env_index, joint_index] + + +@wp.kernel +def clip_joint_array_with_limits_masked( + lower_limits: wp.array(dtype=wp.float32), + upper_limits: wp.array(dtype=wp.float32), + joint_array: wp.array(dtype=wp.float32), + env_mask: wp.array(dtype=wp.bool), + joint_mask: wp.array(dtype=wp.bool), +): + joint_index = wp.tid() + if env_mask[joint_index] and joint_mask[joint_index]: + joint_array[joint_index] = wp.clamp( + joint_array[joint_index], lower_limits[joint_index], upper_limits[joint_index] + ) + + +@wp.kernel +def clip_joint_array_with_limits( + lower_limits: wp.array(dtype=wp.float32), + upper_limits: wp.array(dtype=wp.float32), + joint_array: wp.array(dtype=wp.float32), +): + index = wp.tid() + joint_array[index] = wp.clamp(joint_array[index], lower_limits[index], upper_limits[index]) diff --git a/source/isaaclab/isaaclab/newton/assets/core/kernels/other_kernels.py b/source/isaaclab/isaaclab/newton/assets/core/kernels/other_kernels.py new file mode 100644 index 00000000000..1c11f77c4a9 --- /dev/null +++ b/source/isaaclab/isaaclab/newton/assets/core/kernels/other_kernels.py @@ -0,0 +1,87 @@ +# Copyright (c) 2022-2025, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +import warp as wp + + +@wp.kernel +def update_wrench_array( + new_value: wp.array2d(dtype=wp.spatial_vectorf), + wrench: wp.array2d(dtype=wp.spatial_vectorf), + env_ids: wp.array(dtype=wp.bool), + body_ids: wp.array(dtype=wp.bool), +): + env_index, body_index = wp.tid() + if env_ids[env_index] and body_ids[body_index]: + wrench[env_index, body_index] = new_value[env_index, body_index] + + +@wp.kernel +def update_wrench_array_with_value( + value: wp.spatial_vectorf, + wrench: wp.array2d(dtype=wp.spatial_vectorf), + env_ids: wp.array(dtype=wp.bool), + body_ids: wp.array(dtype=wp.bool), +): + env_index, body_index = wp.tid() + if env_ids[env_index] and body_ids[body_index]: + wrench[env_index, body_index] = value + + +@wp.func +def update_wrench_with_force( + force: wp.vec3f, +) -> wp.spatial_vectorf: + return wp.spatial_vectorf(0.0, 0.0, 0.0, force[0], force[1], force[2]) + + +@wp.func +def update_wrench_with_torque( + torque: wp.vec3f, +) -> wp.spatial_vectorf: + return wp.spatial_vectorf(torque[0], torque[1], torque[2], 0.0, 0.0, 0.0) + + +@wp.kernel +def update_wrench_array_with_force( + forces: wp.array2d(dtype=wp.vec3f), + wrench: wp.array2d(dtype=wp.spatial_vectorf), + env_ids: wp.array(dtype=wp.bool), + body_ids: wp.array(dtype=wp.bool), +): + env_index, body_index = wp.tid() + if env_ids[env_index] and body_ids[body_index]: + wrench[env_index, body_index] = update_wrench_with_force(forces[env_index, body_index]) + + +@wp.kernel +def update_wrench_array_with_torque( + torques: wp.array2d(dtype=wp.vec3f), + wrench: wp.array2d(dtype=wp.spatial_vectorf), + env_ids: wp.array(dtype=wp.bool), + body_ids: wp.array(dtype=wp.bool), +): + env_index, body_index = wp.tid() + if env_ids[env_index] and body_ids[body_index]: + wrench[env_index, body_index] = update_wrench_with_torque(torques[env_index, body_index]) + + +@wp.kernel +def generate_mask_from_ids( + mask: wp.array(dtype=wp.bool), + ids: wp.array(dtype=wp.int32), +): + index = wp.tid() + mask[ids[index]] = True + + +@wp.kernel +def populate_empty_array( + input_array: wp.array(dtype=wp.float32), + output_array: wp.array(dtype=wp.float32), + indices: wp.array(dtype=wp.int32), +): + index = wp.tid() + output_array[indices[index]] = input_array[index] diff --git a/source/isaaclab/isaaclab/newton/assets/core/kernels/state_kernels.py b/source/isaaclab/isaaclab/newton/assets/core/kernels/state_kernels.py new file mode 100644 index 00000000000..235a37b0f97 --- /dev/null +++ b/source/isaaclab/isaaclab/newton/assets/core/kernels/state_kernels.py @@ -0,0 +1,563 @@ +# Copyright (c) 2022-2025, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +import warp as wp + +""" +Split/Combine pose kernels +""" + +vec13f = wp.types.vector(length=13, dtype=wp.float32) + + +@wp.kernel +def split_transform_array_to_position_array( + transform: wp.array(dtype=wp.transformf), + position: wp.array(dtype=wp.vec3f), +): + """ + Split a transform array to a position array. + + Args: + transform: The transform. Shape is (num_instances, 7). + position: The position. Shape is (num_instances, 3). (modified) + """ + index = wp.tid() + position[index] = wp.transform_get_translation(transform[index]) + +@wp.kernel +def split_transform_array_to_quaternion_array( + transform: wp.array(dtype=wp.transformf), + quaternion: wp.array(dtype=wp.quatf), +): + + """ + Split a transform array to a quaternion array. + + Args: + transform: The transform. Shape is (num_instances, 7). + quaternion: The quaternion. Shape is (num_instances, 4). (modified) + """ + index = wp.tid() + quaternion[index] = wp.transform_get_rotation(transform[index]) + +@wp.kernel +def split_transform_batched_array_to_position_batched_array( + transform: wp.array2d(dtype=wp.transformf), + position: wp.array2d(dtype=wp.vec3f), +): + """ + Split a transform batched array to a position batched array. + """ + index, body_index = wp.tid() + position[index, body_index] = wp.transform_get_translation(transform[index, body_index]) + +@wp.kernel +def split_transform_batched_array_to_quaternion_batched_array( + transform: wp.array2d(dtype=wp.transformf), + quaternion: wp.array2d(dtype=wp.quatf), +): + """ + Split a transform batched array to a quaternion batched array. + """ + index, body_index = wp.tid() + quaternion[index, body_index] = wp.transform_get_rotation(transform[index, body_index]) + +@wp.kernel +def generate_pose_from_position_with_unit_quaternion( + position: wp.array(dtype=wp.vec3f), + pose: wp.array(dtype=wp.transformf), +): + """ + Generate a pose from a position with a unit quaternion. + + Args: + position: The position. Shape is (num_instances, 3). + pose: The pose. Shape is (num_instances, 7). (modified) + """ + index = wp.tid() + pose[index] = wp.transformf(position[index], wp.quatf(0.0, 0.0, 0.0, 1.0)) + +@wp.kernel +def generate_pose_from_position_with_unit_quaternion_batched( + position: wp.array2d(dtype=wp.vec3f), + pose: wp.array2d(dtype=wp.transformf), +): + """ + Generate a pose from a position with a unit quaternion. + + Args: + position: The position. Shape is (num_instances, num_bodies, 3). + pose: The pose. Shape is (num_instances, num_bodies, 7). (modified) + """ + index, body_index = wp.tid() + pose[index, body_index] = wp.transformf(position[index, body_index], wp.quatf(0.0, 0.0, 0.0, 1.0)) + + +""" +Split/Combine state kernels +""" + + +@wp.func +def split_state_to_pose( + state: vec13f, +) -> wp.transformf: + """ + Split a state into a pose. + + The state is given in the following format: (x, y, z, qx, qy, qz, qw, wx, wy, wz, vx, vy, vz). + + .. note:: The quaternion is given in the following format: (qx, qy, qz, qw). + + """ + return wp.transformf(wp.vec3f(state[0], state[1], state[2]), wp.quatf(state[3], state[4], state[5], state[6])) + + +@wp.func +def split_state_to_velocity( + state: vec13f, +) -> wp.spatial_vectorf: + """ + Split a state into a velocity. + + The state is given in the following format: (x, y, z, qx, qy, qz, qw, wx, wy, wz, vx, vy, vz). + + .. note:: The quaternion is given in the following format: (qx, qy, qz, qw). + + """ + return wp.spatial_vectorf(state[7], state[8], state[9], state[10], state[11], state[12]) + +@wp.func +def combine_state( + pose: wp.transformf, + velocity: wp.spatial_vectorf, +) -> vec13f: + """ + Combine a pose and a velocity into a state. + + The state is given in the following format: (x, y, z, qx, qy, qz, qw, wx, wy, wz, vx, vy, vz). + + .. note:: The quaternion is given in the following format: (qx, qy, qz, qw). + + Args: + pose: The pose. Shape is (1, 7). + velocity: The velocity. Shape is (1, 6). + + Returns: + The state. Shape is (1, 13). + """ + return vec13f( + pose[0], + pose[1], + pose[2], + pose[3], + pose[4], + pose[5], + pose[6], + velocity[0], + velocity[1], + velocity[2], + velocity[3], + velocity[4], + velocity[5], + ) + + +@wp.kernel +def combine_pose_and_velocity_to_state( + root_pose: wp.array(dtype=wp.transformf), + root_velocity: wp.array(dtype=wp.spatial_vectorf), + root_state: wp.array(dtype=vec13f), +): + """ + Combine a pose and a velocity into a state. + + The state is given in the following format: (x, y, z, qx, qy, qz, qw, wx, wy, wz, vx, vy, vz). + + .. note:: The quaternion is given in the following format: (qx, qy, qz, qw). + + Args: + pose: The pose. Shape is (num_instances, 7). + velocity: The velocity. Shape is (num_instances, 6). + state: The state. Shape is (num_instances, 13). (modified) + """ + env_index = wp.tid() + root_state[env_index] = combine_state(root_pose[env_index], root_velocity[env_index]) + + +@wp.kernel +def combine_pose_and_velocity_to_state_masked( + root_pose: wp.array(dtype=wp.transformf), + root_velocity: wp.array(dtype=wp.spatial_vectorf), + root_state: wp.array(dtype=vec13f), + env_mask: wp.array(dtype=wp.bool), +): + """ + Combine a pose and a velocity into a state. + + The state is given in the following format: (x, y, z, qx, qy, qz, qw, wx, wy, wz, vx, vy, vz). + + .. note:: The quaternion is given in the following format: (qx, qy, qz, qw). + + Args: + pose: The pose. Shape is (num_instances, 7). + velocity: The velocity. Shape is (num_instances, 6). + state: The state. Shape is (num_instances, 13). (modified) + env_mask: The mask of the environments to combine the state for. Shape is (num_instances,). + """ + env_index = wp.tid() + if env_mask[env_index]: + root_state[env_index] = combine_state(root_pose[env_index], root_velocity[env_index]) + + +@wp.kernel +def combine_pose_and_velocity_to_state_batched( + root_pose: wp.array2d(dtype=wp.transformf), + root_velocity: wp.array2d(dtype=wp.spatial_vectorf), + root_state: wp.array2d(dtype=vec13f), +): + """ + Combine a pose and a velocity into a state. + + The state is given in the following format: (x, y, z, qx, qy, qz, qw, wx, wy, wz, vx, vy, vz). + + .. note:: The quaternion is given in the following format: (qx, qy, qz, qw). + + Args: + pose: The pose. Shape is (num_instances, num_bodies, 7). + velocity: The velocity. Shape is (num_instances, num_bodies, 6). + state: The state. Shape is (num_instances, num_bodies, 13). (modified) + """ + env_index, body_index = wp.tid() + root_state[env_index, body_index] = combine_state( + root_pose[env_index, body_index], root_velocity[env_index, body_index] + ) + + +@wp.kernel +def combine_pose_and_velocity_to_state_batched_masked( + root_pose: wp.array2d(dtype=wp.transformf), + root_velocity: wp.array2d(dtype=wp.spatial_vectorf), + root_state: wp.array2d(dtype=vec13f), + env_mask: wp.array(dtype=wp.bool), + body_mask: wp.array(dtype=wp.bool), +): + """ + Combine a pose and a velocity into a state. + + The state is given in the following format: (x, y, z, qx, qy, qz, qw, wx, wy, wz, vx, vy, vz). + + .. note:: The quaternion is given in the following format: (qx, qy, qz, qw). + + Args: + pose: The pose. Shape is (num_instances, num_bodies, 7). + velocity: The velocity. Shape is (num_instances, num_bodies, 6). + state: The state. Shape is (num_instances, num_bodies, 13). (modified) + env_mask: The mask of the environments to combine the state for. Shape is (num_instances,). + body_mask: The mask of the bodies to combine the state for. Shape is (num_bodies,). + """ + env_index, body_index = wp.tid() + if env_mask[env_index] and body_mask[body_index]: + root_state[env_index, body_index] = combine_state( + root_pose[env_index, body_index], root_velocity[env_index, body_index] + ) + + +""" +Frame combination kernels +""" + + +@wp.func +def combine_transforms(p1: wp.vec3f, q1: wp.quatf, p2: wp.vec3f, q2: wp.quatf) -> wp.transformf: + """ + Combine two transforms. + + Args: + p1: The position of the first transform. Shape is (3,). + q1: The quaternion of the first transform. Shape is (4,). + p2: The position of the second transform. Shape is (3,). + q2: The quaternion of the second transform. Shape is (4,). + + Returns: + The combined transform. Shape is (1, 7). + """ + return wp.transformf(p1 + wp.quat_rotate(q1, p2), q1 * q2) + + +@wp.kernel +def combine_frame_transforms_partial( + pose_1: wp.array(dtype=wp.transformf), + position_2: wp.array(dtype=wp.vec3f), + resulting_pose: wp.array(dtype=wp.transformf), +): + """ + Combine a frame transform with a position. + + Args: + pose_1: The frame transform. Shape is (1, 7). + position_2: The position. Shape is (1, 3). + resulting_pose: The resulting pose. Shape is (1, 7). (modified) + """ + index = wp.tid() + resulting_pose[index] = combine_transforms( + wp.transform_get_translation(pose_1[index]), + wp.transform_get_rotation(pose_1[index]), + position_2[index], + wp.quatf(0.0, 0.0, 0.0, 1.0), + ) + + +@wp.kernel +def combine_frame_transforms_partial_batch( + pose_1: wp.array2d(dtype=wp.transformf), + position_2: wp.array2d(dtype=wp.vec3f), + resulting_pose: wp.array2d(dtype=wp.transformf), +): + """ + Combine a frame transform with a position. + + Args: + pose_1: The frame transform. Shape is (num_instances, 7). + position_2: The position. Shape is (num_instances, 3). + resulting_pose: The resulting pose. Shape is (num_instances, 7). (modified) + """ + env_idx, body_idx = wp.tid() + resulting_pose[env_idx, body_idx] = combine_transforms( + wp.transform_get_translation(pose_1[env_idx, body_idx]), + wp.transform_get_rotation(pose_1[env_idx, body_idx]), + position_2[env_idx, body_idx], + wp.quatf(0.0, 0.0, 0.0, 1.0), + ) + + +@wp.kernel +def combine_frame_transforms( + pose_1: wp.array(dtype=wp.transformf), + pose_2: wp.array(dtype=wp.transformf), + resulting_pose: wp.array(dtype=wp.transformf), +): + """ + Combine two transforms. + + Args: + pose_1: The first transform. Shape is (1, 7). + pose_2: The second transform. Shape is (1, 7). + resulting_pose: The resulting pose. Shape is (1, 7). (modified) + """ + index = wp.tid() + resulting_pose[index] = combine_transforms( + wp.transform_get_translation(pose_1[index]), + wp.transform_get_rotation(pose_1[index]), + wp.transform_get_translation(pose_2[index]), + wp.transform_get_rotation(pose_2[index]), + ) + + +@wp.kernel +def combine_frame_transforms_batch( + pose_1: wp.array2d(dtype=wp.transformf), + pose_2: wp.array2d(dtype=wp.transformf), + resulting_pose: wp.array2d(dtype=wp.transformf), +): + """ + Combine two transforms. + + Args: + pose_1: The first transform. Shape is (num_instances, 7). + pose_2: The second transform. Shape is (num_instances, 7). + resulting_pose: The resulting pose. Shape is (num_instances, 7). (modified) + """ + env_idx, body_idx = wp.tid() + resulting_pose[env_idx, body_idx] = combine_transforms( + wp.transform_get_translation(pose_1[env_idx, body_idx]), + wp.transform_get_rotation(pose_1[env_idx, body_idx]), + wp.transform_get_translation(pose_2[env_idx, body_idx]), + wp.transform_get_rotation(pose_2[env_idx, body_idx]), + ) + + +@wp.kernel +def project_vec_from_quat_single( + vec: wp.vec3f, quat: wp.array(dtype=wp.quatf), resulting_vec: wp.array(dtype=wp.vec3f) +): + """ + Project a vector from a quaternion. + + Args: + vec: The vector. Shape is (3,). + quat: The quaternion. Shape is (4,). + resulting_vec: The resulting vector. Shape is (3,). (modified) + """ + index = wp.tid() + resulting_vec[index] = wp.quat_rotate(quat[index], vec) + + +@wp.func +def project_velocity_to_frame( + velocity: wp.spatial_vectorf, + pose: wp.transformf, +) -> wp.spatial_vectorf: + """ + Project a velocity to a frame. + + Args: + velocity: The velocity. Shape is (6,). + pose: The pose. Shape is (1, 7). + resulting_velocity: The resulting velocity. Shape is (6,). (modified) + """ + w = wp.quat_rotate_inv(wp.transform_get_rotation(pose), wp.spatial_top(velocity)) + v = wp.quat_rotate_inv(wp.transform_get_rotation(pose), wp.spatial_bottom(velocity)) + return wp.spatial_vectorf(w[0], w[1], w[2], v[0], v[1], v[2]) + + +@wp.kernel +def project_velocities_to_frame( + velocity: wp.array(dtype=wp.spatial_vectorf), + pose: wp.array(dtype=wp.transformf), + resulting_velocity: wp.array(dtype=wp.spatial_vectorf), +): + """ + Project a velocity to a frame. + + Args: + velocity: The velocity. Shape is (num_instances, 6). + pose: The pose. Shape is (num_instances, 7). + resulting_velocity: The resulting velocity. Shape is (num_instances, 6). (modified) + """ + index = wp.tid() + resulting_velocity[index] = project_velocity_to_frame(velocity[index], pose[index]) + + +""" +Heading utility kernels +""" + + +@wp.func +def heading_vec_b(quat: wp.quatf, vec: wp.vec3f) -> float: + quat_rot = wp.quat_rotate(quat, vec) + return wp.atan2(quat_rot[0], quat_rot[3]) + + +@wp.kernel +def compute_heading(forward_vec_b: wp.vec3f, quat_w: wp.array(dtype=wp.quatf), heading: wp.array(dtype=wp.float32)): + index = wp.tid() + heading[index] = heading_vec_b(quat_w[index], forward_vec_b) + + +""" +Update kernels +""" + + +@wp.kernel +def update_transforms_array( + new_pose: wp.array(dtype=wp.transformf), + pose: wp.array(dtype=wp.transformf), + env_mask: wp.array(dtype=wp.bool), +): + """ + Update a transforms array. + + Args: + new_pose: The new pose. Shape is (num_instances, 7). + pose: The pose. Shape is (num_instances, 7). (modified) + env_mask: The mask of the environments to update the pose for. Shape is (num_instances,). + """ + index = wp.tid() + if env_mask[index]: + pose[index] = new_pose[index] + + +@wp.kernel +def update_transforms_array_with_value( + value: wp.transformf, + pose: wp.array(dtype=wp.transformf), + env_mask: wp.array(dtype=wp.bool), +): + """ + Update a transforms array with a value. + + Args: + value: The value. Shape is (7,). + pose: The pose. Shape is (num_instances, 7). (modified) + env_mask: The mask of the environments to update the pose for. Shape is (num_instances,). + """ + index = wp.tid() + if env_mask[index]: + pose[index] = value + + +@wp.kernel +def update_spatial_vector_array( + velocity: wp.array(dtype=wp.spatial_vectorf), + new_velocity: wp.array(dtype=wp.spatial_vectorf), + env_mask: wp.array(dtype=wp.bool), +): + """ + Update a spatial vector array. + + Args: + new_velocity: The new velocity. Shape is (num_instances, 6). + velocity: The velocity. Shape is (num_instances, 6). (modified) + env_mask: The mask of the environments to update the velocity for. Shape is (num_instances,). + """ + index = wp.tid() + if env_mask[index]: + velocity[index] = new_velocity[index] + + +@wp.kernel +def update_spatial_vector_array_with_value( + value: wp.spatial_vectorf, + velocity: wp.array(dtype=wp.spatial_vectorf), + env_mask: wp.array(dtype=wp.bool), +): + """ + Update a spatial vector array with a value. + + Args: + value: The value. Shape is (6,). + velocity: The velocity. Shape is (num_instances, 6). (modified) + env_mask: The mask of the environments to update the velocity for. Shape is (num_instances,). + """ + index = wp.tid() + if env_mask[index]: + velocity[index] = value + + +""" +Transform kernels +""" + + +@wp.kernel +def transform_CoM_pose_to_link_frame_masked( + com_pose_w: wp.array(dtype=wp.transformf), + com_pose_link_frame: wp.array(dtype=wp.transformf), + link_pose_w: wp.array(dtype=wp.transformf), + env_mask: wp.array(dtype=wp.bool), +): + """ + Transform a CoM pose to a link frame. + + + + Args: + com_pose_w: The CoM pose in the world frame. Shape is (num_instances, 7). + com_pose_link_frame: The CoM pose in the link frame. Shape is (num_instances, 7). + link_pose_w: The link pose in the world frame. Shape is (num_instances, 7). (modified) + env_mask: The mask of the environments to transform the CoM pose to the link frame for. Shape is (num_instances,). + """ + index = wp.tid() + if env_mask[index]: + link_pose_w[index] = combine_transforms( + wp.transform_get_translation(com_pose_w[index]), + wp.transform_get_rotation(com_pose_w[index]), + wp.transform_get_translation(com_pose_link_frame[index]), + wp.quatf(0.0, 0.0, 0.0, 1.0), + ) diff --git a/source/isaaclab/isaaclab/newton/assets/core/kernels/update.py b/source/isaaclab/isaaclab/newton/assets/core/kernels/update.py new file mode 100644 index 00000000000..05bb1adbecf --- /dev/null +++ b/source/isaaclab/isaaclab/newton/assets/core/kernels/update.py @@ -0,0 +1,210 @@ +import warp as wp +from typing import Any + +@wp.kernel +def update_array_with_value( + new_value: Any, + array: Any, +): + """ + Assigns a value to all the elements of the array. + + Args: + new_value: The new value. + + Modifies: + array: The array to update. + """ + index = wp.tid() + array[index] = new_value + + +@wp.kernel +def update_array_with_value_masked( + new_value: Any, + array: Any, + mask: wp.array(dtype=wp.bool), +): + """ + Assigns a value to all the elements of the array where the mask is true. + + Args: + new_value: The new value. + mask: The mask to use. Shape is (N,). + + Modifies: + array: The array to update. Shape is (N,). + """ + index = wp.tid() + if mask[index]: + array[index] = new_value + + +@wp.kernel +def update_batched_array_with_value( + new_value: Any, + array_2d: Any, +): + """ + Assigns a value to all the elements of the batched array. + + Args: + new_value: The new value. + + Modifies: + array_2d: The array to update. Shape is (N, M). + """ + index_1, index_2 = wp.tid() + array_2d[index_1, index_2] = new_value + + +@wp.kernel +def update_batched_array_with_value_masked( + new_value: Any, + array_2d: Any, + mask_1: wp.array(dtype=wp.bool), + mask_2: wp.array(dtype=wp.bool), +): + """ + Assigns a value to all the elements of the batched array where the masks are true. + + Args: + new_value: The new value. + mask_1: The mask to use. Shape is (N,). + mask_2: The mask to use. Shape is (M,). + + Modifies: + array_2d: The array to update. Shape is (N, M). + """ + index_1, index_2 = wp.tid() + if mask_1[index_1] and mask_2[index_2]: + array_2d[index_1, index_2] = new_value + +@wp.kernel +def update_array_with_array( + new_array: Any, + array: Any, +): + """ + Assigns the elements of the new array to the elements of the array. + + Args: + new_array: The new array. Shape is (N,). + + Modifies: + array: The array to update. Shape is (N,). + """ + index = wp.tid() + array[index] = new_array[index] + +@wp.kernel +def update_array_with_array_masked( + new_array: Any, + array: Any, + mask: wp.array(dtype=wp.bool), +): + """ + Assigns the elements of the new array to the elements of the array where the mask is true. + + Args: + new_array: The new array. Shape is (N,). + mask: The mask to use. Shape is (N,). + + Modifies: + array: The array to update. Shape is (N,). + """ + index = wp.tid() + if mask[index]: + array[index] = new_array[index] + +@wp.kernel +def update_batched_array_with_array( + new_array: Any, + array_2d: Any, +): + """ + Assigns the elements of the new array to the elements of the batched array. + + Args: + new_array: The new array. Shape is (M,). + + Modifies: + array_2d: The array to update. Shape is (N, M). + """ + index_1, index_2 = wp.tid() + array_2d[index_1, index_2] = new_array[index_2] + +@wp.kernel +def update_batched_array_with_array_masked( + new_array: Any, + array_2d: Any, + mask_1: wp.array(dtype=wp.bool), + mask_2: wp.array(dtype=wp.bool), +): + """ + Assigns the elements of the new array to the elements of the batched array where the masks are true. + + Args: + new_array: The new array. Shape is (M,). + mask_1: The mask to use. Shape is (N,). + mask_2: The mask to use. Shape is (M,). + + Modifies: + array_2d: The array to update. Shape is (N, M). + """ + index_1, index_2 = wp.tid() + if mask_1[index_1] and mask_2[index_2]: + array_2d[index_1, index_2] = new_array[index_2] + +@wp.kernel +def update_batched_array_with_batched_array( + new_array_2d: Any, + array_2d: Any, +): + """ + Assigns the elements of the new array to the elements of the batched array. + + Args: + new_array_2d: The new array. Shape is (N, M). + + Modifies: + array_2d: The array to update. Shape is (N, M). + """ + index_1, index_2 = wp.tid() + array_2d[index_1, index_2] = new_array_2d[index_1, index_2] + +@wp.kernel +def update_batched_array_with_batched_array_masked( + new_array_2d: Any, + array_2d: Any, + mask_1: wp.array(dtype=wp.bool), + mask_2: wp.array(dtype=wp.bool), +): + """ + Assigns the elements of the new array to the elements of the batched array where the masks are true. + + Args: + new_array_2d: The new array. Shape is (N, M). + mask_1: The mask to use. Shape is (N,). + mask_2: The mask to use. Shape is (M,). + + Modifies: + array_2d: The array to update. Shape is (N, M). + """ + index_1, index_2 = wp.tid() + if mask_1[index_1] and mask_2[index_2]: + array_2d[index_1, index_2] = new_array_2d[index_1, index_2] + + +for dtype in [wp.float32, wp.int32, wp.bool, wp.vec2f, wp.vec3f, wp.quatf, wp.transformf, wp.spatial_vectorf]: + wp.overload(update_array_with_value, {"new_value": dtype, "array": wp.array(dtype=dtype)}) + wp.overload(update_array_with_value_masked, {"new_value": dtype, "array": wp.array(dtype=dtype)}) + wp.overload(update_batched_array_with_value, {"new_value": dtype, "array_2d": wp.array2d(dtype=dtype)}) + wp.overload(update_batched_array_with_value_masked, {"new_value": dtype, "array_2d": wp.array2d(dtype=dtype)}) + wp.overload(update_array_with_array, {"new_array": wp.array(dtype=dtype), "array": wp.array(dtype=dtype)}) + wp.overload(update_array_with_array_masked, {"new_array": wp.array(dtype=dtype), "array": wp.array(dtype=dtype)}) + wp.overload(update_batched_array_with_array, {"new_array": wp.array(dtype=dtype), "array_2d": wp.array2d(dtype=dtype)}) + wp.overload(update_batched_array_with_array_masked, {"new_array": wp.array(dtype=dtype), "array_2d": wp.array2d(dtype=dtype)}) + wp.overload(update_batched_array_with_batched_array, {"new_array_2d": wp.array2d(dtype=dtype), "array_2d": wp.array2d(dtype=dtype)}) + wp.overload(update_batched_array_with_batched_array_masked, {"new_array_2d": wp.array2d(dtype=dtype), "array_2d": wp.array2d(dtype=dtype)}) + \ No newline at end of file diff --git a/source/isaaclab/isaaclab/newton/assets/core/kernels/velocity_kernels.py b/source/isaaclab/isaaclab/newton/assets/core/kernels/velocity_kernels.py new file mode 100644 index 00000000000..901005bf7e9 --- /dev/null +++ b/source/isaaclab/isaaclab/newton/assets/core/kernels/velocity_kernels.py @@ -0,0 +1,492 @@ +# Copyright (c) 2022-2025, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +import warp as wp + +""" +Casters from spatial velocity to linear and angular velocity +""" + +@wp.kernel +def split_spatial_vectory_array_to_linear_velocity_array( + velocity: wp.array(dtype=wp.spatial_vectorf), + linear_velocity: wp.array(dtype=wp.vec3f) +): + """ + Split a spatial velocity array to a linear velocity array. + + Velocities are given in the following format: (vx, vy, vz, wx, wy, wz). + + Args: + velocity: The spatial velocity. Shape is (6,). + linear_velocity: The linear velocity. Shape is (3,). (modified) + """ + index = wp.tid() + linear_velocity[index] = wp.spatial_top(velocity[index]) + + +@wp.kernel +def split_spatial_vectory_array_to_angular_velocity_array( + velocity: wp.array(dtype=wp.spatial_vectorf), + angular_velocity: wp.array(dtype=wp.vec3f) +): + """ + Split a spatial velocity array to an angular velocity array. + + Velocities are given in the following format: (vx, vy, vz, wx, wy, wz). + + Args: + velocity: The spatial velocity. Shape is (num_instances, 6). + angular_velocity: The angular velocity. Shape is (num_instances, 3). (modified) + """ + index = wp.tid() + angular_velocity[index] = wp.spatial_bottom(velocity[index]) + +@wp.kernel +def split_spatial_vectory_batched_array_to_linear_velocity_batched_array( + velocity: wp.array2d(dtype=wp.spatial_vectorf), + linear_velocity: wp.array2d(dtype=wp.vec3f) +): + """ + Split a spatial velocity batched array to a linear velocity batched array. + """ + index, body_index = wp.tid() + linear_velocity[index, body_index] = wp.spatial_top(velocity[index, body_index]) + +@wp.kernel +def split_spatial_vectory_batched_array_to_angular_velocity_batched_array( + velocity: wp.array2d(dtype=wp.spatial_vectorf), + angular_velocity: wp.array2d(dtype=wp.vec3f) +): + """ + Split a spatial velocity batched array to an angular velocity batched array. + """ + + index, body_index = wp.tid() + angular_velocity[index, body_index] = wp.spatial_bottom(velocity[index, body_index]) + +""" +Projectors from com frame to link frame and vice versa +""" + + +@wp.func +def velocity_projector( + com_velocity: wp.spatial_vectorf, + link_pose: wp.transformf, + com_position: wp.vec3f, +) -> wp.spatial_vectorf: + """ + Project a velocity from the com frame to the link frame. + + Velocities are given in the following format: (wx, wy, wz, vx, vy, vz). + + .. caution:: Velocities are given with angular velocity first and linear velocity second. + + .. note:: Only :arg:`com_position` is needed as in Newton, the CoM orientation is always aligned with the + link frame. + + Args: + com_velocity: The velocity in the com frame. Shape is (6,). + link_pose: The link pose in the world frame. Shape is (7,). + com_position: The position of the com in the link frame. Shape is (3,). + + Returns: + wp.spatial_vectorf: The projected velocity in the link frame. Shape is (6,). + """ + u = wp.spatial_top(com_velocity) + w = wp.spatial_bottom(com_velocity) + wp.cross( + wp.spatial_bottom(com_velocity), + wp.quat_rotate(wp.transform_get_rotation(link_pose), -com_position), + ) + return wp.spatial_vectorf(u[0], u[1], u[2], w[0], w[1], w[2]) + #return wp.spatial_vector(u, w) --> Do it like that. + + +@wp.func +def velocity_projector_inv( + com_velocity: wp.spatial_vectorf, + link_pose: wp.transformf, + com_position: wp.vec3f, +) -> wp.spatial_vectorf: + """ + Project a velocity from the link frame to the com frame. + + Velocities are given in the following format: (wx, wy, wz, vx, vy, vz). + + .. caution:: Velocities are given with angular velocity first and linear velocity second. + + .. note:: Only :arg:`com_position` is needed as in Newton, the CoM orientation is always aligned with the + link frame. + + Args: + com_velocity: The velocity in the link frame. Shape is (6,). + link_pose: The link pose in the world frame. Shape is (7,). + com_position: The position of the com in the link frame. Shape is (3,). + + Returns: + wp.spatial_vectorf: The projected velocity in the com frame. Shape is (6,). + """ + u = wp.spatial_top(com_velocity) + w = wp.spatial_bottom(com_velocity) + wp.cross( + wp.spatial_bottom(com_velocity), + wp.quat_rotate(wp.transform_get_rotation(link_pose), com_position), + ) + return wp.spatial_vectorf(u[0], u[1], u[2], w[0], w[1], w[2]) + + +""" +Kernels to project velocities to and from the com frame +""" + + +@wp.kernel +def project_com_velocity_to_link_frame( + com_velocity: wp.array(dtype=wp.spatial_vectorf), + link_pose: wp.array(dtype=wp.transformf), + com_position: wp.array(dtype=wp.vec3f), + link_velocity: wp.array(dtype=wp.spatial_vectorf), +): + """ + Project a velocity from the com frame to the link frame. + + Velocities are given in the following format: (wx, wy, wz, vx, vy, vz). + + .. caution:: Velocities are given with angular velocity first and linear velocity second. + + .. note:: Only :arg:`com_position` is needed as in Newton, the CoM orientation is always aligned with the + link frame. + + Args: + com_velocity: The com velocity in the world frame. Shape is (num_links, 6). + link_pose: The link pose in the world frame. Shape is (num_links, 7). + com_position: The com position in link frame. Shape is (num_links, 3). + link_velocity: The link velocity. Shape is (num_links, 6). (modified) + """ + index = wp.tid() + link_velocity[index] = velocity_projector(com_velocity[index], link_pose[index], com_position[index]) + + +@wp.kernel +def project_com_velocity_to_link_frame_masked( + com_velocity: wp.array(dtype=wp.spatial_vectorf), + link_pose: wp.array(dtype=wp.transformf), + com_position: wp.array(dtype=wp.vec3f), + link_velocity: wp.array(dtype=wp.spatial_vectorf), + mask: wp.array(dtype=wp.bool), +): + """ + Project a velocity from the com frame to the link frame. + + Velocities are given in the following format: (wx, wy, wz, vx, vy, vz). + + .. caution:: Velocities are given with angular velocity first and linear velocity second. + + .. note:: Only :arg:`com_position` is needed as in Newton, the CoM orientation is always aligned with the + link frame. + + Args: + com_velocity: The com velocity in the world frame. Shape is (num_links, 6). + link_pose: The link pose in the world frame. Shape is (num_links, 7). + com_position: The com position in link frame. Shape is (num_links, 3). + link_velocity: The link velocity in the world frame. Shape is (num_links, 6). (modified) + mask: The mask of the links to project the velocity to. Shape is (num_links,). + """ + index = wp.tid() + if mask[index]: + link_velocity[index] = velocity_projector( + com_velocity[index], + link_pose[index], + com_position[index], + ) + + +@wp.kernel +def project_com_velocity_to_link_frame_batch( + com_velocity: wp.array2d(dtype=wp.spatial_vectorf), + link_pose: wp.array2d(dtype=wp.transformf), + com_position: wp.array2d(dtype=wp.vec3f), + link_velocity: wp.array2d(dtype=wp.spatial_vectorf), +): + """ + Project a velocity from the com frame to the link frame. + + Velocities are given in the following format: (wx, wy, wz, vx, vy, vz). + + .. caution:: Velocities are given with angular velocity first and linear velocity second. + + .. note:: Only :attr:`com_position` is needed as in Newton, the CoM orientation is always aligned with the + link frame. + + Args: + com_velocity: The com velocity in the world frame. Shape is (num_links, 6). + link_pose: The link pose in the world frame. Shape is (num_links, 7). + com_position: The com position in link frame. Shape is (num_links, 3). + link_velocity: The link velocity in the world frame. Shape is (num_links, 6). (modified) + """ + env_idx, body_idx = wp.tid() + link_velocity[env_idx, body_idx] = velocity_projector( + com_velocity[env_idx, body_idx], + link_pose[env_idx, body_idx], + com_position[env_idx, body_idx], + ) + + +@wp.kernel +def project_com_velocity_to_link_frame_batch_masked( + com_velocity: wp.array2d(dtype=wp.spatial_vectorf), + link_pose: wp.array2d(dtype=wp.transformf), + com_position: wp.array2d(dtype=wp.vec3f), + link_velocity: wp.array2d(dtype=wp.spatial_vectorf), + env_mask: wp.array(dtype=wp.bool), + body_mask: wp.array(dtype=wp.bool), +): + """ + Project a velocity from the com frame to the link frame. + + Velocities are given in the following format: (wx, wy, wz, vx, vy, vz). + + .. caution:: Velocities are given with angular velocity first and linear velocity second. + + .. note:: Only :arg:`com_position` is needed as in Newton, the CoM orientation is always aligned with the + link frame. + + Args: + com_velocity: The com velocity in the world frame. Shape is (num_links, 6). + link_pose: The link pose in the world frame. Shape is (num_links, 7). + com_position: The com position in link frame. Shape is (num_links, 3). + link_velocity: The link velocity in the world frame. Shape is (num_links, 6). (modified) + env_mask: The mask of the environments to project the velocity to. Shape is (num_links,). + body_mask: The mask of the bodies to project the velocity to. Shape is (num_links,). + """ + env_idx, body_idx = wp.tid() + if env_mask[env_idx] and body_mask[body_idx]: + link_velocity[env_idx, body_idx] = velocity_projector( + com_velocity[env_idx, body_idx], + link_pose[env_idx, body_idx], + com_position[env_idx, body_idx], + ) + + +@wp.kernel +def project_link_velocity_to_com_frame( + link_velocity: wp.array(dtype=wp.spatial_vectorf), + link_pose: wp.array(dtype=wp.transformf), + com_position: wp.array(dtype=wp.vec3f), + com_velocity: wp.array(dtype=wp.spatial_vectorf), +): + """ + Project a velocity from the link frame to the com frame. + + Velocities are given in the following format: (wx, wy, wz, vx, vy, vz). + + .. caution:: Velocities are given with angular velocity first and linear velocity second. + + .. note:: Only :arg:`com_position` is needed as in Newton, the CoM orientation is always aligned with the + link frame. + + Args: + link_velocity: The link velocity in the world frame. Shape is (num_links, 6). + link_pose: The link pose in the world frame. Shape is (num_links, 7). + com_position: The com position in link frame. Shape is (num_links, 3). + com_velocity: The com velocity in the world frame. Shape is (num_links, 6). (modified) + """ + index = wp.tid() + com_velocity[index] = velocity_projector_inv(link_velocity[index], link_pose[index], com_position[index]) + + +@wp.kernel +def project_link_velocity_to_com_frame_masked( + link_velocity: wp.array(dtype=wp.spatial_vectorf), + link_pose: wp.array(dtype=wp.transformf), + com_position: wp.array(dtype=wp.vec3f), + com_velocity: wp.array(dtype=wp.spatial_vectorf), + mask: wp.array(dtype=wp.bool), +): + """ + Project a velocity from the link frame to the com frame. + + Velocities are given in the following format: (wx, wy, wz, vx, vy, vz). + + .. caution:: Velocities are given with angular velocity first and linear velocity second. + + .. note:: Only :arg:`com_position` is needed as in Newton, the CoM orientation is always aligned with the + link frame. + + Args: + link_velocity: The link velocity in the world frame. Shape is (num_links, 6). + link_pose: The link pose in the world frame. Shape is (num_links, 7). + com_position: The com position in link frame. Shape is (num_links, 3). + com_velocity: The com velocity in the world frame. Shape is (num_links, 6). (modified) + mask: The mask of the links to project the velocity to. Shape is (num_links,). + """ + index = wp.tid() + if mask[index]: + com_velocity[index] = velocity_projector_inv( + link_velocity[index], + link_pose[index], + com_position[index], + ) + + +@wp.kernel +def project_link_velocity_to_com_frame_batch( + link_velocity: wp.array2d(dtype=wp.spatial_vectorf), + link_pose: wp.array2d(dtype=wp.transformf), + com_position: wp.array2d(dtype=wp.vec3f), + com_velocity: wp.array2d(dtype=wp.spatial_vectorf), +): + """ + Project a velocity from the link frame to the com frame. + + Velocities are given in the following format: (wx, wy, wz, vx, vy, vz). + + .. caution:: Velocities are given with angular velocity first and linear velocity second. + + .. note:: Only :arg:`com_position` is needed as in Newton, the CoM orientation is always aligned with the + link frame. + + Args: + link_velocity (wp.array2d(dtype=wp.spatial_vectorf)): The link velocity in the world frame. + link_pose (wp.array2d(dtype=wp.transformf)): The link pose in the world frame. + com_position (wp.array2d(dtype=wp.vec3f)): The com position in link frame. + com_velocity (wp.array2d(dtype=wp.spatial_vectorf)): The com velocity in the world frame. (destination) + """ + env_idx, body_idx = wp.tid() + com_velocity[env_idx, body_idx] = velocity_projector_inv( + link_velocity[env_idx, body_idx], link_pose[env_idx, body_idx], com_position[env_idx, body_idx] + ) + + +@wp.kernel +def project_link_velocity_to_com_frame_batch_masked( + link_velocity: wp.array2d(dtype=wp.spatial_vectorf), + link_pose: wp.array2d(dtype=wp.transformf), + com_position: wp.array2d(dtype=wp.vec3f), + com_velocity: wp.array2d(dtype=wp.spatial_vectorf), + env_mask: wp.array(dtype=wp.bool), + body_mask: wp.array(dtype=wp.bool), +): + """ + Project a velocity from the link frame to the com frame. + + Velocities are given in the following format: (wx, wy, wz, vx, vy, vz). + + .. caution:: Velocities are given with angular velocity first and linear velocity second. + + .. note:: Only :arg:`com_position` is needed as in Newton, the CoM orientation is always aligned with the + link frame. + + Args: + link_velocity: The link velocity in the world frame. Shape is (num_links, 6). + link_pose: The link pose in the world frame. Shape is (num_links, 7). + com_position: The com position in link frame. Shape is (num_links, 3). + com_velocity: The com velocity in the world frame. Shape is (num_links, 6). (modified) + env_mask: The mask of the environments to project the velocity to. Shape is (num_links,). + body_mask: The mask of the bodies to project the velocity to. Shape is (num_links,). + """ + env_idx, body_idx = wp.tid() + if env_mask[env_idx] and body_mask[body_idx]: + com_velocity[env_idx, body_idx] = velocity_projector_inv( + link_velocity[env_idx, body_idx], + link_pose[env_idx, body_idx], + com_position[env_idx, body_idx], + ) + + +""" +Kernels to update spatial vector arrays +""" + +@wp.kernel +def update_spatial_vector_array_masked( + new_velocity: wp.array(dtype=wp.spatial_vectorf), + velocity: wp.array(dtype=wp.spatial_vectorf), + mask: wp.array(dtype=wp.bool), +): + """ + Update a velocity array with a new velocity. + + Velocities are given in the following format: (wx, wy, wz, vx, vy, vz). + + .. caution:: Velocities are given with angular velocity first and linear velocity second. + + Args: + new_velocity: The new velocity. Shape is (num_links, 6). + velocity: The velocity array. Shape is (num_links, 6). (modified) + mask: The mask of the velocities to update. Shape is (num_links,). + """ + index = wp.tid() + if mask[index]: + velocity[index] = new_velocity[index] + + +@wp.kernel +def update_spatial_vector_array_batch_masked( + new_velocity: wp.array2d(dtype=wp.spatial_vectorf), + velocity: wp.array2d(dtype=wp.spatial_vectorf), + env_mask: wp.array(dtype=wp.bool), + body_mask: wp.array(dtype=wp.bool), +): + """ + Update a velocity array with a new velocity. + + Velocities are given in the following format: (wx, wy, wz, vx, vy, vz). + + .. caution:: Velocities are given with angular velocity first and linear velocity second. + + Args: + new_velocity: The new velocity. Shape is (num_links, 6). + velocity: The velocity array. Shape is (num_links, 6). (modified) + env_mask: The mask of the environments to update. Shape is (num_links,). + body_mask: The mask of the bodies to update. Shape is (num_links,). + """ + env_idx, body_idx = wp.tid() + if env_mask[env_idx] and body_mask[body_idx]: + velocity[env_idx, body_idx] = new_velocity[env_idx, body_idx] + + +""" +Kernels to derive body acceleration from velocity. +""" + +@wp.kernel +def derive_body_acceleration_from_velocity( + velocity: wp.array(dtype=wp.spatial_vectorf), + previous_velocity: wp.array(dtype=wp.spatial_vectorf), + dt: float, + acceleration: wp.array(dtype=wp.spatial_vectorf), +): + """ + Derive the body acceleration from the velocity. + + Args: + velocity: The velocity. Shape is (num_instances, 6). + previous_velocity: The previous velocity. Shape is (num_instances, 6). + dt: The time step. + acceleration: The acceleration. Shape is (num_instances, 6). (modified) + """ + env_idx = wp.tid() + acceleration[env_idx] = (velocity[env_idx] - previous_velocity[env_idx]) / dt + + +@wp.kernel +def derive_body_acceleration_from_velocity_batched( + velocity: wp.array2d(dtype=wp.spatial_vectorf), + previous_velocity: wp.array2d(dtype=wp.spatial_vectorf), + dt: float, + acceleration: wp.array2d(dtype=wp.spatial_vectorf), +): + """ + Derive the body acceleration from the velocity. + + Args: + velocity: The velocity. Shape is (num_instances, num_bodies, 6). + previous_velocity: The previous velocity. Shape is (num_instances, num_bodies, 6). + dt: The time step. + acceleration: The acceleration. Shape is (num_instances, num_bodies, 6). (modified) + """ + env_idx, body_idx = wp.tid() + acceleration[env_idx, body_idx] = (velocity[env_idx, body_idx] - previous_velocity[env_idx, body_idx]) / dt diff --git a/source/isaaclab/isaaclab/newton/assets/core/root_properties/__init__.py b/source/isaaclab/isaaclab/newton/assets/core/root_properties/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/source/isaaclab/isaaclab/newton/assets/core/root_properties/root.py b/source/isaaclab/isaaclab/newton/assets/core/root_properties/root.py new file mode 100644 index 00000000000..8f56bef680e --- /dev/null +++ b/source/isaaclab/isaaclab/newton/assets/core/root_properties/root.py @@ -0,0 +1,390 @@ +# Copyright (c) 2022-2025, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +# Flag for pyright to ignore type errors in this file. +# pyright: reportPrivateUsage=false + + +import warp as wp +import weakref + +from newton.selection import ArticulationView as NewtonArticulationView +from isaaclab.utils.helpers import warn_overhead_cost +from isaaclab.newton.assets.core.root_properties.root_data import RootData +from isaaclab.newton.assets.core.kernels import ( + project_link_velocity_to_com_frame_masked, + transform_CoM_pose_to_link_frame_masked, + update_wrench_array_with_force, + update_wrench_array_with_torque, + update_array_with_value_masked, + update_array_with_array_masked, +) + + +class Root: + def __init__(self, root_newton_view: NewtonArticulationView, root_data: RootData, device: str): + self._root_newton_view: NewtonArticulationView = weakref.proxy(root_newton_view) + self._root_data = root_data + self._device = device + + self._create_buffers() + + """ + Properties + """ + + @property + def data(self) -> RootData: + return self._root_data + + @property + def num_instances(self) -> int: + return self._root_newton_view.count + + @property + def root_body_names(self) -> list[str]: + return self._root_newton_view.body_names[0] + + @property + def root_newton_view(self) -> NewtonArticulationView: + """Articulation view for the asset (Newton). + + Note: + Use this view with caution. It requires handling of tensors in a specific way. + """ + return self._root_newton_view + + """ + Operations. + """ + + def reset(self, mask: wp.array): + # use ellipses object to skip initial indices. + if mask is None: + mask = self._ALL_ENV_MASK + + # reset external wrench + wp.launch( + update_array_with_value_masked, + dim=(self.num_instances,), + inputs=[ + wp.spatial_vectorf(0.0, 0.0, 0.0, 0.0, 0.0, 0.0), + self._external_wrench, + mask, + ], + ) + + def write_data_to_sim(self): + """Write external wrenches and joint commands to the simulation. + + If any explicit actuators are present, then the actuator models are used to compute the + joint commands. Otherwise, the joint commands are directly set into the simulation. + + Note: + We write external wrench to the simulation here since this function is called before the simulation step. + This ensures that the external wrench is applied at every simulation step. + """ + # Wrenches are automatically applied by set_external_force_and_torque. + # apply actuator models + pass + + def update(self, dt: float): + self._root_data.update(dt) + + """ + Operations - State Writers. + """ + + @warn_overhead_cost( + "write_root_link_pose_to_sim and or write_root_com_velocity_to_sim", + "This function splits the root state into pose and velocity. Consider using write_root_link_pose_to_sim and" + " write_root_com_velocity_to_sim instead. In general there is no good reasons to be using states with Newton.", + ) + def write_root_state_to_sim(self, root_state: wp.array, env_mask: wp.array | None = None): + """Set the root state over selected environment indices into the simulation. + + The root state comprises of the cartesian position, quaternion orientation in (x, y, z, w), and angular + and linear velocity. All the quantities are in the simulation frame. + + Args: + root_state: Root state in simulation frame. Shape is (num_instances, 13). + env_mask: Environment mask. Shape is (num_instances,). + """ + + if env_mask is None: + env_mask = self._ALL_ENV_MASK + + self.write_root_link_pose_to_sim(root_state[:, :7], env_mask=env_mask) + self.write_root_com_velocity_to_sim(root_state[:, 7:], env_mask=env_mask) + + @warn_overhead_cost( + "write_root_state_to_sim", + "This function splits the root state into pose and velocity. Consider using write_root_link_pose_to_sim and" + " write_root_com_velocity_to_sim instead. In general there is no good reasons to be using states with Newton.", + ) + def write_root_com_state_to_sim(self, root_state: wp.array, env_mask: wp.array | None = None): + """Set the root center of mass state over selected environment indices into the simulation. + + The root state comprises of the cartesian position, quaternion orientation in (x, y, z, w), and angular + and linear velocity. All the quantities are in the simulation frame. + + Args: + root_state: Root state in simulation frame. Shape is (num_instances, 13). + env_mask: Environment mask. Shape is (num_instances,). + """ + + if env_mask is None: + env_mask = self._ALL_ENV_MASK + + self.write_root_com_pose_to_sim(root_state[:, :7], env_mask=env_mask) + self.write_root_com_velocity_to_sim(root_state[:, 7:], env_mask=env_mask) + + @warn_overhead_cost( + "write_root_state_to_sim", + "This function splits the root state into pose and velocity. Consider using write_root_link_pose_to_sim and" + " write_root_com_velocity_to_sim instead. In general there is no good reasons to be using states with Newton.", + ) + def write_root_link_state_to_sim(self, root_state: wp.array, env_mask: wp.array | None = None): + """Set the root link state over selected environment indices into the simulation. + + The root state comprises of the cartesian position, quaternion orientation in (x, y, z, w), and angular + and linear velocity. All the quantities are in the simulation frame. + + Args: + root_state: Root state in simulation frame. Shape is (num_instances, 13). + env_mask: Environment mask. Shape is (num_instances,). + """ + + if env_mask is None: + env_mask = self._ALL_ENV_MASK + + self.write_root_link_pose_to_sim(root_state[:, :7], env_mask=env_mask) + self.write_root_link_velocity_to_sim(root_state[:, 7:], env_mask=env_mask) + + def write_root_pose_to_sim(self, root_pose: wp.array, env_mask: wp.array | None = None): + """Set the root pose over selected environment indices into the simulation. + + The root pose comprises of the cartesian position and quaternion orientation in (x, y, z, w). + + Args: + root_pose: Root poses in simulation frame. Shape is (num_instances, 7). + env_mask: Environment mask. Shape is (num_instances,). + """ + self.write_root_link_pose_to_sim(root_pose, env_mask=env_mask) + + def write_root_link_pose_to_sim(self, pose: wp.array, env_mask: wp.array | None = None): + """Set the root link pose over selected environment indices into the simulation. + + + The root pose ``wp.transformf`` comprises of the cartesian position and quaternion orientation in (x, y, z, w). + + Args: + root_pose: Root poses in simulation frame. Shape is (num_instances, 7). + env_mask: Environment mask. Shape is (num_instances,). + """ + if env_mask is None: + env_mask = self._ALL_ENV_MASK + + # set into internal buffers + wp.launch( + update_array_with_array_masked, + dim=self.num_instances, + inputs=[ + pose, + self._root_data.root_link_pose_w, + env_mask, + ], + ) + # Need to invalidate the buffer to trigger the update with the new state. + self._root_data._root_com_pose_w.timestamp = -1.0 + # Don't need to invalidate the body pose since they will not change unless we call forward kinematics. + + def write_root_com_pose_to_sim(self, root_pose: wp.array, env_mask: wp.array | None = None) -> None: + """Set the root center of mass pose over selected environment indices into the simulation. + + The root pose comprises of the cartesian position and quaternion orientation in (w, x, y, z). + The orientation is the orientation of the principle axes of inertia. + + Args: + root_pose: Root center of mass poses in simulation frame. Shape is (num_instances, 7). + env_mask: Environment mask. Shape is (num_instances,). + """ + # resolve all indices + if env_mask is None: + env_mask = self._ALL_ENV_MASK + + # set into internal buffers + wp.launch( + update_array_with_array_masked, + dim=self.num_instances, + inputs=[ + root_pose, + self._root_data.root_com_pose_w, + env_mask, + ], + ) + # set link frame poses + wp.launch( + transform_CoM_pose_to_link_frame_masked, + dim=self.num_instances, + inputs=[ + self._root_data.root_com_pose_w, + self._root_data.root_com_pos_b, + self._root_data.root_link_pose_w, + env_mask, + ], + ) + # Don't need to invalidate the body pose since they will not change unless we call forward kinematics. + + def write_root_velocity_to_sim(self, root_velocity: wp.array, env_mask: wp.array | None = None) -> None: + """Set the root center of mass velocity over selected environment indices into the simulation. + + The velocity comprises angular velocity (x, y, z) and linear velocity (x, y, z) in that order. + NOTE: This sets the velocity of the root's center of mass rather than the roots frame. + + Args: + root_velocity: Root center of mass velocities in simulation world frame. Shape is (num_instances, 6). + env_mask: Environment mask. Shape is (num_instances,). + """ + self.write_root_com_velocity_to_sim(root_velocity=root_velocity, env_mask=env_mask) + + def write_root_com_velocity_to_sim(self, root_velocity: wp.array, env_mask: wp.array | None = None) -> None: + """Set the root center of mass velocity over selected environment indices into the simulation. + + The velocity comprises angular velocity (x, y, z) and linear velocity (x, y, z) in that order. + NOTE: This sets the velocity of the root's center of mass rather than the roots frame. + + Args: + root_velocity: Root center of mass velocities in simulation world frame. Shape is (num_instances, 6). + env_mask: Environment mask. Shape is (num_instances,). + """ + # resolve all indices + if env_mask is None: + env_mask = self._ALL_ENV_MASK + + # set into internal buffers + wp.launch( + update_array_with_array_masked, + dim=self.num_instances, + inputs=[ + root_velocity, + self._root_data.root_com_vel_w, + env_mask, + ], + ) + # Need to invalidate the buffer to trigger the update with the new state. + self._root_data._root_link_vel_w.timestamp = -1.0 + self._root_data._root_link_vel_b.timestamp = -1.0 + self._root_data._root_com_vel_b.timestamp = -1.0 + + def write_root_link_velocity_to_sim(self, root_velocity: wp.array, env_mask: wp.array | None = None) -> None: + """Set the root link velocity over selected environment indices into the simulation. + + The velocity comprises angular velocity (x, y, z) and linear velocity (x, y, z) in that order. + NOTE: This sets the velocity of the root's frame rather than the roots center of mass. + + Args: + root_velocity: Root frame velocities in simulation world frame. Shape is (num_instances, 6). + env_mask: Environment mask. Shape is (num_instances,). + """ + # resolve all indices + if env_mask is None: + env_mask = self._ALL_ENV_MASK + # update the root link velocity + wp.launch( + update_array_with_array_masked, + dim=self.num_instances, + inputs=[ + root_velocity, + self._root_data.root_link_vel_w, + env_mask, + ], + ) + # set into internal buffers + wp.launch( + project_link_velocity_to_com_frame_masked, + dim=self.num_instances, + inputs=[ + root_velocity, + self._root_data.root_link_pose_w, + self._root_data.root_com_pos_b, + self._root_data.root_com_vel_w, + env_mask, + ], + ) + # Need to invalidate the buffer to trigger the update with the new state. + self._root_data._root_link_vel_b.timestamp = -1.0 + self._root_data._root_com_vel_b.timestamp = -1.0 + + """ + Operations - Setters. + """ + + def set_external_force_and_torque( + self, + forces: wp.array, + torques: wp.array, + env_mask: wp.array | None = None, + ) -> None: + """Set external force and torque to apply on the asset's bodies in their local frame. + + For many applications, we want to keep the applied external force on rigid bodies constant over a period of + time (for instance, during the policy control). This function allows us to store the external force and torque + into buffers which are then applied to the simulation at every step. + + .. caution:: + If the function is called with empty forces and torques, then this function disables the application + of external wrench to the simulation. + + .. code-block:: python + + # example of disabling external wrench + asset.set_external_force_and_torque(forces=wp.zeros(0, 3), torques=wp.zeros(0, 3)) + + .. note:: + This function does not apply the external wrench to the simulation. It only fills the buffers with + the desired values. To apply the external wrench, call the :meth:`write_data_to_sim` function + right before the simulation step. + + Args: + forces: External forces in bodies' local frame. Shape is (num_instances, num_bodies, 3). + torques: External torques in bodies' local frame. Shape is (num_instances, num_bodies, 3). + body_mask: The body mask. Shape is (num_bodies). + env_mask: The environment mask. Shape is (num_instances,). + """ + # resolve indices + if env_mask is None: + env_mask = self._ALL_ENV_MASK + # Check if there are any external forces or torques + if (forces is not None) or (torques is not None): + self.has_external_wrench = True + if forces is not None: + wp.launch( + update_wrench_array_with_force, + dim=self.num_instances, + inputs=[ + forces, + self._external_wrench, + env_mask, + ], + ) + if torques is not None: + wp.launch( + update_wrench_array_with_torque, + dim=self.num_instances, + inputs=[ + torques, + self._external_wrench, + env_mask, + ], + ) + + def _create_buffers(self): + # constants + self._ALL_ENV_MASK = wp.ones((self.num_instances,), dtype=wp.bool, device=self._device) + # masks + self._ENV_MASK = wp.zeros((self.num_instances,), dtype=wp.bool, device=self._device) + # external wrench + self._external_wrench = wp.zeros((self.num_instances), dtype=wp.spatial_vectorf, device=self._device) \ No newline at end of file diff --git a/source/isaaclab/isaaclab/newton/assets/core/root_properties/root_data.py b/source/isaaclab/isaaclab/newton/assets/core/root_properties/root_data.py new file mode 100644 index 00000000000..1916e0a3bf4 --- /dev/null +++ b/source/isaaclab/isaaclab/newton/assets/core/root_properties/root_data.py @@ -0,0 +1,706 @@ +# Copyright (c) 2022-2025, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +import weakref + +import warp as wp + +from newton.selection import ArticulationView as NewtonArticulationView +from isaaclab.sim._impl.newton_manager import NewtonManager +from isaaclab.utils.buffers import TimestampedWarpBuffer +from isaaclab.utils.helpers import deprecated, warn_overhead_cost + +from isaaclab.newton.assets.core.kernels import ( + combine_frame_transforms_partial, + combine_pose_and_velocity_to_state, + compute_heading, + project_com_velocity_to_link_frame, + project_vec_from_quat_single, + project_velocities_to_frame, + derive_body_acceleration_from_velocity, + generate_pose_from_position_with_unit_quaternion, + vec13f, + split_transform_array_to_position_array, + split_transform_array_to_quaternion_array, + split_spatial_vectory_array_to_linear_velocity_array, + split_spatial_vectory_array_to_angular_velocity_array, +) + + +class RootData: + def __init__(self, root_newton_view: NewtonArticulationView, device: str): + """Initializes the articulation data. + + Args: + root_newton_view: The root articulation view. + device: The device used for processing. + """ + # Set the parameters + self.device = device + # Set the root articulation view + # note: this is stored as a weak reference to avoid circular references between the asset class + # and the data container. This is important to avoid memory leaks. + self._root_newton_view: NewtonArticulationView = weakref.proxy(root_newton_view) + + # Set initial time stamp + self._sim_timestamp = 0.0 + + # obtain global simulation view + gravity = wp.to_torch(NewtonManager.get_model().gravity)[0] + gravity_dir = [float(i) / sum(gravity) for i in gravity] + # Initialize constants + self.GRAVITY_VEC_W = wp.vec3f(gravity_dir[0], gravity_dir[1], gravity_dir[2]) + self.FORWARD_VEC_B = wp.vec3f((1.0, 0.0, 0.0)) + + self._create_simulation_bindings() + self._create_buffers() + + def update(self, dt: float): + # update the simulation timestamp + self._sim_timestamp += dt + # Trigger an update of the root acceleration buffer at a higher frequency + # since we do finite differencing. + #self.root_com_acc_w + + def _create_simulation_bindings(self): + """Create simulation bindings for the root data. + + Direct simulation bindings are pointers to the simulation data, their data is not copied, and should + only be updated using warp kernels. Any modifications made to the bindings will be reflected in the simulation. + Hence we encourage users to carefully think about the data they modify and in which order it should be updated. + + .. caution:: This is possible if and only if the properties that we access are strided from newton and not + indexed. Newton willing this is the case all the time, but we should pay attention to this if things look off. + """ + self._sim_bind_root_link_pose_w = self._root_newton_view.get_root_transforms(NewtonManager.get_state_0()) + self._sim_bind_root_com_vel_w = self._root_newton_view.get_root_velocities(NewtonManager.get_state_0()) + self._sim_bind_body_com_pos_b = self._root_newton_view.get_attribute("body_com", NewtonManager.get_model())[:, 0] + self._sim_bind_root_mass = self._root_newton_view.get_attribute("body_mass", NewtonManager.get_model())[:, 0] + self._sim_bind_root_inertia = self._root_newton_view.get_attribute("body_inertia", NewtonManager.get_model())[:, 0] + self._sim_bind_root_external_wrench = self._root_newton_view.get_attribute("body_f", NewtonManager.get_state_0())[:, 0] + + def _create_buffers(self): + """Create buffers for the root data.""" + # Initialize history for finite differencing + try: + self._previous_root_com_vel = wp.clone(self._root_newton_view.get_root_velocities(NewtonManager.get_state_0())) + except: + self._previous_root_com_vel = wp.zeros((self._root_newton_view.count, 1), dtype=wp.spatial_vectorf, device=self.device) + # -- default root pose and velocity + self._default_root_pose = wp.zeros((self._root_newton_view.count), dtype=wp.transformf, device=self.device) + self._default_root_vel = wp.zeros((self._root_newton_view.count), dtype=wp.spatial_vectorf, device=self.device) + + # Initialize the lazy buffers. + # -- link frame w.r.t. world frame + self._root_link_vel_w = TimestampedWarpBuffer(shape=(self._root_newton_view.count), dtype=wp.spatial_vectorf) + self._root_link_vel_b = TimestampedWarpBuffer(shape=(self._root_newton_view.count), dtype=wp.spatial_vectorf) + self._projected_gravity_b = TimestampedWarpBuffer(shape=(self._root_newton_view.count), dtype=wp.vec3f) + self._heading_w = TimestampedWarpBuffer(shape=(self._root_newton_view.count), dtype=wp.float32) + # -- com frame w.r.t. world frame + self._root_com_pose_w = TimestampedWarpBuffer(shape=(self._root_newton_view.count), dtype=wp.transformf) + self._root_com_vel_b = TimestampedWarpBuffer(shape=(self._root_newton_view.count), dtype=wp.spatial_vectorf) + self._root_com_acc_w = TimestampedWarpBuffer(shape=(self._root_newton_view.count), dtype=wp.spatial_vectorf) + + ## + # Direct simulation bindings accessors. + ## + + @property + def root_link_pose_w(self) -> wp.array: + """Root link pose ``wp.transformf`` in the world frame. Shape is (num_instances,). + + The pose is in the form of [pos, quat]. The orientation is provided in (x, y, z, w) format. + """ + return self._sim_bind_root_link_pose_w + + @property + def root_com_vel_w(self) -> wp.array: + """Root center of mass velocity ``wp.spatial_vectorf`` in the world frame. Shape is (num_instances,). + + The velocity is in the form of [ang_vel, lin_vel]. + """ + return self._sim_bind_root_com_vel_w + + ## + # Default accessors. + ## + + @property + def default_root_pose(self) -> wp.array: + """Default root pose ``[pos, quat]`` in the local environment frame. Shape is (num_instances, 7). + + The position and quaternion are of the articulation root's actor frame. + + This quantity is configured through the :attr:`isaaclab.assets.ArticulationCfg.init_state` parameter. + """ + return self._default_root_pose + + @property + def default_root_vel(self) -> wp.array: + """Default root velocity ``[lin_vel, ang_vel]`` in the local environment frame. Shape is (num_instances, 6). + + The linear and angular velocities are of the articulation root's center of mass frame. + + This quantity is configured through the :attr:`isaaclab.assets.ArticulationCfg.init_state` parameter. + """ + return self._default_root_vel + + @default_root_pose.setter + def default_root_pose(self, value: wp.array) -> None: + self._default_root_pose = value + + @default_root_vel.setter + def default_root_vel(self, value: wp.array) -> None: + self._default_root_vel = value + + ## + # Root state properties. + ## + + @property + def root_mass(self) -> wp.array: + """Root mass ``wp.float32`` in the world frame. Shape is (num_instances,).""" + return self._sim_bind_root_mass + + @property + def root_inertia(self) -> wp.array: + """Root inertia ``wp.mat33`` in the world frame. Shape is (num_instances, 9).""" + return self._sim_bind_root_inertia + + @property + def root_external_wrench(self) -> wp.array: + """Root external wrench ``wp.spatial_vectorf`` in the world frame. Shape is (num_instances, 6).""" + return self._sim_bind_root_external_wrench + + @property + def root_link_vel_w(self) -> wp.array: + """Root link velocity ``wp.spatial_vectorf`` in simulation world frame. + + Shapes are (num_instances,). Velocities are in the form of [wx, wy, wz, vx, vy, vz]. + This quantity contains the linear and angular velocities of the articulation root's actor frame + relative to the world. + """ + if self._root_link_vel_w.timestamp < self._sim_timestamp: + wp.launch( + project_com_velocity_to_link_frame, + dim=(self._root_newton_view.count), + device=self.device, + inputs=[ + self.root_com_vel_w, + self._sim_bind_root_link_pose_w, + self._sim_bind_body_com_pos_b, + self._root_link_vel_w.data, + ], + ) + # set the buffer data and timestamp + self._root_link_vel_w.timestamp = self._sim_timestamp + + return self._root_link_vel_w.data + + @property + def root_com_pose_w(self) -> wp.array: + """Root center of mass pose ``wp.transformf`` in simulation world frame. + + Shapes are (num_instances,). The pose is in the form of [pos, quat]. + This quantity is the pose of the articulation root's center of mass frame relative to the world. + + .. note:: The quaternion is in the form of [x, y, z, w]. + """ + if self._root_com_pose_w.timestamp < self._sim_timestamp: + # apply local transform to center of mass frame + wp.launch( + combine_frame_transforms_partial, + dim=(self._root_newton_view.count), + device=self.device, + inputs=[ + self._sim_bind_root_link_pose_w, + self._sim_bind_body_com_pos_b, + self._root_com_pose_w.data, + ], + ) + # set the buffer data and timestamp + self._root_com_pose_w.timestamp = self._sim_timestamp + + return self._root_com_pose_w.data + + # TODO: Pre-allocate the state array. + @property + @warn_overhead_cost( + "root_link_pose_w and root_com_vel_w", + "This function combines the root link pose and root center of mass velocity into a single state. However, Newton" + " outputs pose and velocities separately. Consider using only one of them instead. If both are required, use both" + " root_link_pose_w and root_com_vel_w instead.", + ) + def root_state_w(self) -> wp.array: + """Root state ``vec13f`` in simulation world frame. + + Shapes are (num_instances, 13). The pose is in the form of [pos, quat]. + The pose is of the articulation root's actor frame relative to the world. + The velocity is of the articulation root's center of mass frame. + + .. note:: The velocity is in the form of [vx, vy, vz, wx, wy, wz]. + .. note:: The quaternion is in the form of [x, y, z, w]. + """ + state = wp.zeros((self._root_newton_view.count, 13), dtype=wp.float32, device=self.device) + wp.launch( + combine_pose_and_velocity_to_state, + dim=(self._root_newton_view.count,), + device=self.device, + inputs=[ + self._sim_bind_root_link_pose_w, + self._sim_bind_root_com_vel_w, + state, + ], + ) + return state + + @property + @warn_overhead_cost( + "root_link_pose_w and root_link_vel_w", + "This function combines the root link pose and root link velocity into a single state. However, Newton" + " outputs pose and velocities separately. Consider using only one of them instead. If both are required, use both" + " root_link_pose_w and root_link_vel_w instead.", + ) + def root_link_state_w(self) -> wp.array: + """Root link state ``vec13f`` in simulation world frame. + + Shapes are (num_instances, 13). The pose is in the form of [pos, quat]. + The pose is of the articulation root's actor frame relative to the world. + The velocity is of the articulation root's actor frame. + + .. note:: The velocity is in the form of [wx, wy, wz, vx, vy, vz]. + .. note:: The quaternion is in the form of [x, y, z, w]. + """ + state = wp.zeros((self._root_newton_view.count, 13), dtype=wp.float32, device=self.device) + wp.launch( + combine_pose_and_velocity_to_state, + dim=(self._root_newton_view.count,), + device=self.device, + inputs=[ + self._sim_bind_root_link_pose_w, + self.root_link_vel_w, + state, + ], + ) + return state + + @property + @warn_overhead_cost( + "root_com_pose_w and root_com_vel_w", + "This function combines the root center of mass pose and root center of mass velocity into a single state. However, Newton" + " outputs pose and velocities separately. Consider using only one of them instead. If both are required, use both" + " root_com_pose_w and root_com_vel_w instead.", + ) + def root_com_state_w(self) -> wp.array: + """Root center of mass state ``vec13f`` in simulation world frame. + + Shapes are (num_instances, 13). The pose is in the form of [pos, quat]. + The pose is of the articulation root's center of mass frame relative to the world. + The velocity is of the articulation root's center of mass frame. + + .. note:: The velocity is in the form of [vx, vy, vz, wx, wy, wz]. + .. note:: The quaternion is in the form of [x, y, z, w]. + """ + state = wp.zeros((self._root_newton_view.count, 13), dtype=wp.float32, device=self.device) + wp.launch( + combine_pose_and_velocity_to_state, + dim=(self._root_newton_view.count,), + device=self.device, + inputs=[ + self.root_com_pose_w, + self._sim_bind_root_com_vel_w, + state, + ], + ) + return state + + @property + def root_com_acc_w(self) -> wp.array: + """Acceleration of all bodies center of mass ``wp.spatial_vectorf`` in simulation world frame. + + Shapes are (num_instances, num_bodies,). + All values are relative to the world. + + .. note:: The acceleration is in the form of [vx, vy, vz, wx, wy, wz]. + """ + if self._root_com_acc_w.timestamp < self._sim_timestamp: + wp.launch( + derive_body_acceleration_from_velocity, + dim=(self._root_newton_view.count), + inputs=[ + self._sim_bind_root_com_vel_w, + self._previous_root_com_vel, + NewtonManager.get_dt(), + self._root_com_acc_w.data, + ], + ) + # set the buffer data and timestamp + self._root_com_acc_w.timestamp = self._sim_timestamp + return self._root_com_acc_w.data + + @property + def root_com_pos_b(self) -> wp.array: + """Root center of mass position ``wp.vec3f`` in base frame. Shape is (num_instances, 3). + + This quantity is the position of the root rigid body's center of mass frame with respect to the + rigid body's actor frame. + """ + return self._sim_bind_body_com_pos_b + + @property + @warn_overhead_cost( + "body_com_pose_b", + "This function outputs the pose of the CoM, containing both position and orientation. However, in Newton, the" + " CoM is always aligned with the link frame. This means that the quaternion is always [0, 0, 0, 1]. Consider" + " using the position only instead.", + ) + def root_com_pose_b(self) -> wp.array: + """Center of mass pose ``wp.transformf`` of all bodies in their respective body's link frames. + + Shapes are (num_instances, num_bodies,). The pose is in the form of [pos, quat]. + This quantity is the pose of the center of mass frame of the rigid body relative to the body's link frame. + + .. note:: The quaternion is in the form of [x, y, z, w]. + """ + out = wp.zeros( + (self._root_newton_view.count,), dtype=wp.transformf, device=self.device + ) + wp.launch( + generate_pose_from_position_with_unit_quaternion, + dim=self._root_newton_view.count, + inputs=[ + self._sim_bind_body_com_pos_b, + out, + ], + ) + return out + + @property + def root_com_quat_b(self) -> wp.array: + """Root center of mass orientation (x, y, z, w) in base frame. Shape is (num_instances, 4). + + This quantity is the orientation of the root rigid body's center of mass frame with respect to the + rigid body's actor frame. + """ + out = wp.zeros((self._root_newton_view.count, 4), dtype=wp.quatf, device=self.device) + wp.launch( + split_transform_array_to_quaternion_array, + dim=self._root_newton_view.count, + inputs=[ + self.root_com_pose_b, + out, + ], + ) + return out + + ## + # Derived Properties. + ## + + # FIXME: USE SIM_BIND_LINK_POSE_W RATHER THAN JUST THE QUATERNION + @property + def projected_gravity_b(self) -> wp.array: + """Projection of the gravity direction on base frame. Shape is (num_instances, 3).""" + if self._projected_gravity_b.timestamp < self._sim_timestamp: + wp.launch( + project_vec_from_quat_single, + dim=self._root_newton_view.count, + inputs=[ + self.GRAVITY_VEC_W, + self.root_link_quat_w, + self._projected_gravity_b.data, + ], + ) + # set the buffer data and timestamp + self._projected_gravity_b.timestamp = self._sim_timestamp + return self._projected_gravity_b.data + + # FIXME: USE SIM_BIND_LINK_POSE_W RATHER THAN JUST THE QUATERNION + @property + def heading_w(self) -> wp.array: + """Yaw heading of the base frame (in radians). Shape is (num_instances,). + + Note: + This quantity is computed by assuming that the forward-direction of the base + frame is along x-direction, i.e. :math:`(1, 0, 0)`. + """ + if self._heading_w.timestamp < self._sim_timestamp: + wp.launch( + compute_heading, + dim=self._root_newton_view.count, + inputs=[ + self.FORWARD_VEC_B, + self.root_link_quat_w, + self._heading_w.data, + ], + ) + # set the buffer data and timestamp + self._heading_w.timestamp = self._sim_timestamp + return self._heading_w.data + + @property + def root_link_vel_b(self) -> wp.array: + """Root link velocity ``wp.spatial_vectorf`` in base frame. Shape is (num_instances). + + .. note:: The velocity is in the form of [vx, vy, vz, wx, wy, wz]. + """ + if self._root_link_vel_b.timestamp < self._sim_timestamp: + wp.launch( + project_velocities_to_frame, + dim=self._root_newton_view.count, + inputs=[ + self.root_link_vel_w, + self._sim_bind_root_link_pose_w, + self._root_link_vel_b.data, + ], + ) + self._root_link_vel_b.timestamp = self._sim_timestamp + return self._root_link_vel_b.data + + @property + def root_com_vel_b(self) -> wp.array: + """Root center of mass velocity ``wp.spatial_vectorf`` in base frame. Shape is (num_instances). + + .. note:: The velocity is in the form of [vx, vy, vz, wx, wy, wz]. + """ + if self._root_com_vel_b.timestamp < self._sim_timestamp: + wp.launch( + project_velocities_to_frame, + dim=self._root_newton_view.count, + inputs=[ + self._sim_bind_root_com_vel_w, + self._sim_bind_root_link_pose_w, + self._root_com_vel_b.data, + ], + ) + self._root_com_vel_b.timestamp = self._sim_timestamp + return self._root_com_vel_b.data + + ## + # Strided properties. + ## + + @property + @warn_overhead_cost("root_link_pose_w", "Launches a kernel to split the transform array to a position array.") + def root_link_pos_w(self) -> wp.array: + """Root link position ``wp.vec3f`` in simulation world frame. Shape is (num_instances). + + This quantity is the position of the actor frame of the root rigid body relative to the world. + """ + + out = wp.zeros((self._root_newton_view.count,), dtype=wp.vec3f, device=self.device) + wp.launch( + split_transform_array_to_position_array, + dim=self._root_newton_view.count, + inputs=[ + self._sim_bind_root_link_pose_w, + out, + ], + ) + return out + + @property + @warn_overhead_cost("root_link_pose_w", "Launches a kernel to split the transform array to a quaternion array.") + def root_link_quat_w(self) -> wp.array: + """Root link orientation ``wp.quatf`` in simulation world frame. Shape is (num_instances,). + + This quantity is the orientation of the actor frame of the root rigid body. + + .. note:: The quaternion is in the form of [x, y, z, w]. + """ + out = wp.zeros((self._root_newton_view.count,), dtype=wp.quatf, device=self.device) + wp.launch( + split_transform_array_to_quaternion_array, + dim=self._root_newton_view.count, + inputs=[ + self._sim_bind_root_link_pose_w, + out, + ], + ) + return out + + @property + @warn_overhead_cost("root_link_vel_w", "Launches a kernel to split the velocity array to a linear velocity array.") + def root_link_lin_vel_w(self) -> wp.array: + """Root linear velocity ``wp.vec3f`` in simulation world frame. Shape is (num_instances). + + This quantity is the linear velocity of the root rigid body's actor frame relative to the world. + """ + out = wp.zeros((self._root_newton_view.count,), dtype=wp.vec3f, device=self.device) + wp.launch( + split_spatial_vectory_array_to_linear_velocity_array, + dim=self._root_newton_view.count, + inputs=[ + self.root_link_vel_w, + out, + ], + ) + return out + + @property + @warn_overhead_cost("root_link_vel_w", "Launches a kernel to split the velocity array to an angular velocity array.") + def root_link_ang_vel_w(self) -> wp.array: + """Root link angular velocity ``wp.vec3f`` in simulation world frame. Shape is (num_instances). + + This quantity is the angular velocity of the actor frame of the root rigid body relative to the world. + """ + out = wp.zeros((self._root_newton_view.count,), dtype=wp.vec3f, device=self.device) + wp.launch( + split_spatial_vectory_array_to_angular_velocity_array, + dim=self._root_newton_view.count, + inputs=[ + self.root_link_vel_w, + out, + ], + ) + + @property + @warn_overhead_cost("root_com_pose_w", "Launches a kernel to split the transform array to a position array.") + def root_com_pos_w(self) -> wp.array: + """Root center of mass position in simulation world frame. Shape is (num_instances, 3). + + This quantity is the position of the actor frame of the root rigid body relative to the world. + """ + out = wp.zeros((self._root_newton_view.count,), dtype=wp.vec3f, device=self.device) + wp.launch( + split_transform_array_to_position_array, + dim=self._root_newton_view.count, + inputs=[ + self.root_com_pose_w, + out, + ], + ) + + @property + @warn_overhead_cost("root_com_pose_w", "Launches a kernel to split the transform array to a quaternion array.") + def root_com_quat_w(self) -> wp.array: + """Root center of mass orientation ``wp.quatf`` in simulation world frame. Shape is (num_instances,). + + This quantity is the orientation of the root rigid body's center of mass frame. + + .. note:: The quaternion is in the form of [x, y, z, w]. + """ + out = wp.zeros((self._root_newton_view.count,), dtype=wp.quatf, device=self.device) + wp.launch( + split_transform_array_to_quaternion_array, + dim=self._root_newton_view.count, + inputs=[ + self.root_com_pose_w, + out, + ], + ) + return out + + @property + @warn_overhead_cost("root_com_vel_w", "Launches a kernel to split the velocity array to a linear velocity array.") + def root_com_lin_vel_w(self) -> wp.array: + """Root center of mass linear velocity ``wp.vec3f`` in simulation world frame. Shape is (num_instances,). + + This quantity is the linear velocity of the root rigid body's center of mass frame relative to the world. + """ + out = wp.zeros((self._root_newton_view.count,), dtype=wp.vec3f, device=self.device) + wp.launch( + split_spatial_vectory_array_to_linear_velocity_array, + dim=self._root_newton_view.count, + inputs=[ + self._sim_bind_root_com_vel_w, + out, + ], + ) + return out + + @property + @warn_overhead_cost("root_com_vel_w", "Launches a kernel to split the velocity array to an angular velocity array.") + def root_com_ang_vel_w(self) -> wp.array: + """Root center of mass angular velocity ``wp.vec3f`` in simulation world frame. Shape is (num_instances). + + This quantity is the angular velocity of the root rigid body's center of mass frame relative to the world. + """ + out = wp.zeros((self._root_newton_view.count,), dtype=wp.vec3f, device=self.device) + wp.launch( + split_spatial_vectory_array_to_angular_velocity_array, + dim=self._root_newton_view.count, + inputs=[ + self._sim_bind_root_com_vel_w, + out, + ], + ) + return out + + @property + @warn_overhead_cost("root_link_vel_b", "Launches a kernel to split the velocity array to a linear velocity array.") + def root_link_lin_vel_b(self) -> wp.array: + """Root link linear velocity ``wp.vec3f`` in base frame. Shape is (num_instances). + + This quantity is the linear velocity of the articulation root's actor frame with respect to the + its actor frame. + """ + out = wp.zeros((self._root_newton_view.count,), dtype=wp.vec3f, device=self.device) + wp.launch( + split_spatial_vectory_array_to_linear_velocity_array, + dim=self._root_newton_view.count, + inputs=[ + self.root_link_vel_b, + out, + ], + ) + return out + + @property + @warn_overhead_cost("root_link_vel_b", "Launches a kernel to split the velocity array to an angular velocity array.") + def root_link_ang_vel_b(self) -> wp.array: + """Root link angular velocity ``wp.vec3f`` in base world frame. Shape is (num_instances). + + This quantity is the angular velocity of the articulation root's actor frame with respect to the + its actor frame. + """ + out = wp.zeros((self._root_newton_view.count,), dtype=wp.vec3f, device=self.device) + wp.launch( + split_spatial_vectory_array_to_angular_velocity_array, + dim=self._root_newton_view.count, + inputs=[ + self.root_link_vel_b, + out, + ], + ) + return out + + @property + @warn_overhead_cost("root_com_vel_b", "Launches a kernel to split the velocity array to a linear velocity array.") + def root_com_lin_vel_b(self) -> wp.array: + """Root center of mass linear velocity ``wp.vec3f`` in base frame. Shape is (num_instances). + + This quantity is the linear velocity of the articulation root's center of mass frame with respect to the + its actor frame. + """ + out = wp.zeros((self._root_newton_view.count,), dtype=wp.vec3f, device=self.device) + wp.launch( + split_spatial_vectory_array_to_linear_velocity_array, + dim=self._root_newton_view.count, + inputs=[ + self.root_com_vel_b, + out, + ], + ) + return out + + @property + @warn_overhead_cost("root_com_vel_b", "Launches a kernel to split the velocity array to a linear velocity array.") + def root_com_ang_vel_b(self) -> wp.array: + """Root center of mass angular velocity in base world frame. Shape is (num_instances, 3). + + This quantity is the angular velocity of the articulation root's center of mass frame with respect to the + its actor frame. + """ + out = wp.zeros((self._root_newton_view.count,), dtype=wp.vec3f, device=self.device) + wp.launch( + split_spatial_vectory_array_to_angular_velocity_array, + dim=self._root_newton_view.count, + inputs=[ + self.root_com_vel_b, + out, + ], + ) + return out \ No newline at end of file diff --git a/source/isaaclab/isaaclab/newton/cloner/__init__.py b/source/isaaclab/isaaclab/newton/cloner/__init__.py new file mode 100644 index 00000000000..634fbcd3231 --- /dev/null +++ b/source/isaaclab/isaaclab/newton/cloner/__init__.py @@ -0,0 +1,16 @@ +# Copyright (c) 2022-2025, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +# SPDX-FileCopyrightText: Copyright (c) 2022-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: LicenseRef-NvidiaProprietary +# +# NVIDIA CORPORATION, its affiliates and licensors retain all intellectual +# property and proprietary rights in and to this material, related +# documentation and any modifications thereto. Any use, reproduction, +# disclosure or distribution of this material and related documentation +# without an express license agreement from NVIDIA CORPORATION or +# its affiliates is strictly prohibited. +from isaaclab.cloner.cloner import Cloner +from isaaclab.cloner.grid_cloner import GridCloner diff --git a/source/isaaclab/isaaclab/newton/cloner/grid_cloner.py b/source/isaaclab/isaaclab/newton/cloner/grid_cloner.py new file mode 100644 index 00000000000..376f4ee795c --- /dev/null +++ b/source/isaaclab/isaaclab/newton/cloner/grid_cloner.py @@ -0,0 +1,179 @@ +# Copyright (c) 2022-2025, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +# SPDX-FileCopyrightText: Copyright (c) 2022-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: LicenseRef-NvidiaProprietary +# +# NVIDIA CORPORATION, its affiliates and licensors retain all intellectual +# property and proprietary rights in and to this material, related +# documentation and any modifications thereto. Any use, reproduction, +# disclosure or distribution of this material and related documentation +# without an express license agreement from NVIDIA CORPORATION or +# its affiliates is strictly prohibited. +import numpy as np +import torch + +import omni.usd +from pxr import Gf, Usd, UsdGeom + +from isaaclab.newton.cloner.utils import replicate_environment +from isaaclab.sim._impl.newton_manager import NewtonManager + + +class GridCloner: + """This is a specialized Cloner class that will automatically generate clones in a grid fashion.""" + + def __init__(self, spacing: float, num_per_row: int = -1, stage: Usd.Stage = None): + """ + Args: + spacing (float): Spacing between clones. + num_per_row (int): Number of clones to place in a row. Defaults to sqrt(num_clones). + stage (Usd.Stage): Usd stage where source prim and clones are added to. + """ + self._spacing = spacing + self._num_per_row = num_per_row + + self._positions = None + self._orientations = None + + self._base_env_path = None + self._root_path = None + self._stage = stage + if stage is None: + self._stage = omni.usd.get_context().get_stage() + + def get_clone_transforms( + self, + num_clones: int, + position_offsets: np.ndarray = None, + orientation_offsets: np.ndarray = None, + ): + """Computes the positions and orientations of clones in a grid. + + Args: + num_clones (int): Number of clones. + position_offsets (np.ndarray): Positions to be applied as local translations on top of computed clone position. + position_offsets (np.ndarray | torch.Tensor): Positions to be applied as local translations on top of computed clone position. + Defaults to None, no offset will be applied. + orientation_offsets (np.ndarray | torch.Tensor): Orientations to be applied as local rotations for each clone. + Defaults to None, no offset will be applied. + Returns: + positions (List): Computed positions of all clones. + orientations (List): Computed orientations of all clones. + """ + # check if inputs are valid + if position_offsets is not None: + if len(position_offsets) != num_clones: + raise ValueError("Dimension mismatch between position_offsets and prim_paths!") + # convert to numpy array + if isinstance(position_offsets, torch.Tensor): + position_offsets = position_offsets.detach().cpu().numpy() + elif not isinstance(position_offsets, np.ndarray): + position_offsets = np.asarray(position_offsets) + if orientation_offsets is not None: + if len(orientation_offsets) != num_clones: + raise ValueError("Dimension mismatch between orientation_offsets and prim_paths!") + # convert to numpy array + if isinstance(orientation_offsets, torch.Tensor): + orientation_offsets = orientation_offsets.detach().cpu().numpy() + elif not isinstance(orientation_offsets, np.ndarray): + orientation_offsets = np.asarray(orientation_offsets) + + if self._positions is not None and self._orientations is not None: + return self._positions, self._orientations + + self._num_per_row = int(np.sqrt(num_clones)) if self._num_per_row == -1 else self._num_per_row + num_rows = np.ceil(num_clones / self._num_per_row) + num_cols = np.ceil(num_clones / num_rows) + + row_offset = 0.5 * self._spacing * (num_rows - 1) + col_offset = 0.5 * self._spacing * (num_cols - 1) + + positions = [] + orientations = [] + + for i in range(num_clones): + # compute transform + row = i // num_cols + col = i % num_cols + y = row_offset - row * self._spacing + x = col * self._spacing - col_offset + + up_axis = UsdGeom.GetStageUpAxis(self._stage) + position = [y, x, 0] if up_axis == UsdGeom.Tokens.z else [x, 0, y] + orientation = Gf.Quatd.GetIdentity() + + if position_offsets is not None: + translation = position_offsets[i] + position + else: + translation = position + + if orientation_offsets is not None: + orientation = ( + Gf.Quatd(orientation_offsets[i][0].item(), Gf.Vec3d(orientation_offsets[i][1:].tolist())) + * orientation + ) + + orientation = [ + orientation.GetReal(), + orientation.GetImaginary()[0], + orientation.GetImaginary()[1], + orientation.GetImaginary()[2], + ] + + positions.append(translation) + orientations.append(orientation) + + self._positions = positions + self._orientations = orientations + + return positions, orientations + + def clone( + self, + source_prim_path: str, + prim_paths: list[str], + position_offsets: np.ndarray = None, + orientation_offsets: np.ndarray = None, + root_path: str = None, + ): + """Creates clones in a grid fashion. Positions of clones are computed automatically. + + Args: + source_prim_path (str): Path of source object. + prim_paths (List[str]): List of destination paths. + position_offsets (np.ndarray): Positions to be applied as local translations on top of computed clone position. + Defaults to None, no offset will be applied. + orientation_offsets (np.ndarray): Orientations to be applied as local rotations for each clone. + Defaults to None, no offset will be applied. + replicate_physics (bool): Uses omni.physics replication. This will replicate physics properties directly for paths beginning with root_path and skip physics parsing for anything under the base_env_path. + base_env_path (str): Path to namespace for all environments. Required if replicate_physics=True and define_base_env() not called. + clone_in_fabric (bool): Not supported in Newton. This is here for compatibility with IL 2.2. + root_path (str): Prefix path for each environment. Required if replicate_physics=True and generate_paths() not called. + copy_from_source: (bool): Setting this to False will inherit all clones from the source prim; any changes made to the source prim will be reflected in the clones. + Setting this to True will make copies of the source prim when creating new clones; changes to the source prim will not be reflected in clones. Defaults to False. Note that setting this to True will take longer to execute. + enable_env_ids (bool): Setting this enables co-location of clones in physics with automatic filtering of collisions between clones. + Returns: + positions (List): Computed positions of all clones. + """ + + num_clones = len(prim_paths) + NewtonManager._num_envs = num_clones + + positions, orientations = self.get_clone_transforms(num_clones, position_offsets, orientation_offsets) + clone_base_path = self._root_path if root_path is None else root_path + builder, _ = replicate_environment( + self._stage, + source_prim_path, + clone_base_path + "{}", + positions, + orientations, + simplify_meshes=True, + collapse_fixed_joints=False, + joint_ordering="dfs", + joint_drive_gains_scaling=1.0, + ) + NewtonManager.set_builder(builder) + return positions diff --git a/source/isaaclab/isaaclab/newton/cloner/utils.py b/source/isaaclab/isaaclab/newton/cloner/utils.py new file mode 100644 index 00000000000..88892b6a70f --- /dev/null +++ b/source/isaaclab/isaaclab/newton/cloner/utils.py @@ -0,0 +1,151 @@ +# Copyright (c) 2022-2025, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +import numpy as np +from typing import Any + +import warp as wp +from newton import AxisType, ModelBuilder +from pxr import Usd + +from isaaclab.utils.timer import Timer + + +@Timer(name="replicate_environment", msg="Replicate environment took:", enable=True, format="ms") +def replicate_environment( + source, + prototype_path: str, + path_pattern: str, + positions: np.ndarray, + orientations: np.ndarray, + up_axis: AxisType = "Z", + simplify_meshes: bool = True, + spawn_offset: tuple[float] = (0.0, 0.0, 20.0), + **usd_kwargs, +) -> tuple[ModelBuilder, dict[str:Any]]: + """ + Replicates a prototype USD environment in Newton. + + Args: + source (str | pxr.UsdStage): The file path to the USD file, or an existing USD stage instance. + prototype_path (str): The USD path where the prototype env is defined, e.g., "/World/envs/env_0". + path_pattern (str): The USD path pattern for replicated envs, e.g., "/World/envs/env_{}". + num_envs (int): Number of replicas to create. + env_spacing (tuple[float]): Environment spacing vector. + up_axis (AxisType): The desired up-vector (should match the USD stage). + simplify_meshes (bool): If True, simplify the meshes to reduce the number of triangles. This is useful when + meshes with complex geometry are used as collision meshes. + spawn_offset (tuple[float]): The offset to apply to the spawned environments. + **usd_kwargs: Keyword arguments to pass to the USD importer (see `newton.utils.parse_usd()`). + + Returns: + (ModelBuilder, dict): The resulting ModelBuilder containing all replicated environments and a dictionary with USD stage information. + """ + + with Timer(name="newton_env_builder", msg="Env Builder took:", enable=True, format="ms"): + builder = ModelBuilder(up_axis=up_axis) + + stage_info = builder.add_usd( + source, + ignore_paths=[prototype_path], + **usd_kwargs, + ) + + # up_axis sanity check + stage_up_axis = stage_info.get("up_axis") + if isinstance(stage_up_axis, str) and stage_up_axis.upper() != up_axis.upper(): + print(f"WARNING: up_axis '{up_axis}' does not match USD stage up_axis '{stage_up_axis}'") + + with Timer(name="newton_prototype_builder", msg="Prototype Builder took:", enable=True, format="ms"): + # Get child xforms from the prototype path + child_xforms = [] + if isinstance(source, str): + # If source is a file path, load the stage + stage = Usd.Stage.Open(source) + else: + # If source is already a stage + stage = source + + # Get the prototype prim + prototype_prim = stage.GetPrimAtPath(prototype_path) + if prototype_prim.IsValid(): + # Get all child prims that are Xforms + for child_prim in prototype_prim.GetAllChildren(): + if child_prim.GetTypeName() == "Xform": + child_xforms.append(child_prim.GetPath().pathString) + + # If no child xforms found, use the prototype path itself + if not child_xforms: + child_xforms = [prototype_path] + + prototype_builder = ModelBuilder(up_axis=up_axis) + for child_path in child_xforms: + prototype_builder.add_usd( + source, + root_path=child_path, + load_non_physics_prims=False, + **usd_kwargs, + ) + prototype_builder.approximate_meshes("convex_hull") + + with Timer(name="newton_multiple_add_to_builder", msg="All add to builder took:", enable=True, format="ms"): + # clone the prototype env with updated paths + for i, (pos, ori) in enumerate(zip(positions, orientations)): + body_start = builder.body_count + shape_start = builder.shape_count + joint_start = builder.joint_count + articulation_start = builder.articulation_count + + builder.add_builder( + prototype_builder, xform=wp.transform(np.array(pos) + np.array(spawn_offset), wp.quat_identity()) + ) + + if i > 0: + update_paths( + builder, + prototype_path, + path_pattern.format(i), + body_start=body_start, + shape_start=shape_start, + joint_start=joint_start, + articulation_start=articulation_start, + ) + + return builder, stage_info + + +def update_paths( + builder: ModelBuilder, + old_root: str, + new_root: str, + body_start: int | None = None, + shape_start: int | None = None, + joint_start: int | None = None, + articulation_start: int | None = None, +) -> None: + """Updates the paths of the builder to match the new root path. + + Args: + builder (ModelBuilder): The builder to update. + old_root (str): The old root path. + new_root (str): The new root path. + body_start (int): The start index of the bodies. + shape_start (int): The start index of the shapes. + joint_start (int): The start index of the joints. + articulation_start (int): The start index of the articulations. + """ + old_len = len(old_root) + if body_start is not None: + for i in range(body_start, builder.body_count): + builder.body_key[i] = f"{new_root}{builder.body_key[i][old_len:]}" + if shape_start is not None: + for i in range(shape_start, builder.shape_count): + builder.shape_key[i] = f"{new_root}{builder.shape_key[i][old_len:]}" + if joint_start is not None: + for i in range(joint_start, builder.joint_count): + builder.joint_key[i] = f"{new_root}{builder.joint_key[i][old_len:]}" + if articulation_start is not None: + for i in range(articulation_start, builder.articulation_count): + builder.articulation_key[i] = f"{new_root}{builder.articulation_key[i][old_len:]}" diff --git a/source/isaaclab/isaaclab/newton/sensors/__init__.py b/source/isaaclab/isaaclab/newton/sensors/__init__.py new file mode 100644 index 00000000000..c69a30415c6 --- /dev/null +++ b/source/isaaclab/isaaclab/newton/sensors/__init__.py @@ -0,0 +1,41 @@ +# Copyright (c) 2022-2025, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +"""Sub-package containing various sensor classes implementations. + +This subpackage contains the sensor classes that are compatible with Isaac Sim. We include both +USD-based and custom sensors: + +* **USD-prim sensors**: Available in Omniverse and require creating a USD prim for them. + For instance, RTX ray tracing camera and lidar sensors. +* **USD-schema sensors**: Available in Omniverse and require creating a USD schema on an existing prim. + For instance, contact sensors and frame transformers. +* **Custom sensors**: Implemented in Python and do not require creating any USD prim or schema. + For instance, warp-based ray-casters. + +Due to the above categorization, the prim paths passed to the sensor's configuration class +are interpreted differently based on the sensor type. The following table summarizes the +interpretation of the prim paths for different sensor types: + ++---------------------+---------------------------+---------------------------------------------------------------+ +| Sensor Type | Example Prim Path | Pre-check | ++=====================+===========================+===============================================================+ +| Camera | /World/robot/base/camera | Leaf is available, and it will spawn a USD camera | ++---------------------+---------------------------+---------------------------------------------------------------+ +| Contact Sensor | /World/robot/feet_* | Leaf is available and checks if the schema exists | ++---------------------+---------------------------+---------------------------------------------------------------+ +| Ray Caster | /World/robot/base | Leaf exists and is a physics body (Articulation / Rigid Body) | ++---------------------+---------------------------+---------------------------------------------------------------+ +| Frame Transformer | /World/robot/base | Leaf exists and is a physics body (Articulation / Rigid Body) | ++---------------------+---------------------------+---------------------------------------------------------------+ +| Imu | /World/robot/base | Leaf exists and is a physics body (Rigid Body) | ++---------------------+---------------------------+---------------------------------------------------------------+ + +""" + +from .camera import * # noqa: F401, F403 +from .contact_sensor import * # noqa: F401, F403 +from .sensor_base import SensorBase # noqa: F401 +from .sensor_base_cfg import SensorBaseCfg # noqa: F401 diff --git a/source/isaaclab/isaaclab/newton/sensors/camera/__init__.py b/source/isaaclab/isaaclab/newton/sensors/camera/__init__.py new file mode 100644 index 00000000000..4b4497916e1 --- /dev/null +++ b/source/isaaclab/isaaclab/newton/sensors/camera/__init__.py @@ -0,0 +1,13 @@ +# Copyright (c) 2022-2025, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +"""Sub-module for camera wrapper around USD camera prim.""" + +from .camera import Camera +from .camera_cfg import CameraCfg +from .camera_data import CameraData +from .tiled_camera import TiledCamera +from .tiled_camera_cfg import TiledCameraCfg +from .utils import * # noqa: F401, F403 diff --git a/source/isaaclab/isaaclab/newton/sensors/camera/camera.py b/source/isaaclab/isaaclab/newton/sensors/camera/camera.py new file mode 100644 index 00000000000..d5e3d2fa439 --- /dev/null +++ b/source/isaaclab/isaaclab/newton/sensors/camera/camera.py @@ -0,0 +1,705 @@ +# Copyright (c) 2022-2025, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +from __future__ import annotations + +import json +import numpy as np +import re +import torch +from collections.abc import Sequence +from typing import TYPE_CHECKING, Any, Literal + +import carb +import isaacsim.core.utils.stage as stage_utils +import omni.kit.commands +import omni.usd +from isaacsim.core.prims import XFormPrim + +# from isaacsim.core.version import get_version +from pxr import UsdGeom + +import isaaclab.sim as sim_utils +import isaaclab.utils.sensors as sensor_utils +from isaaclab.utils import to_camel_case +from isaaclab.utils.array import convert_to_torch +from isaaclab.utils.math import ( + convert_camera_frame_orientation_convention, + create_rotation_matrix_from_view, + quat_from_matrix, +) + +from ..sensor_base import SensorBase +from .camera_data import CameraData + +if TYPE_CHECKING: + from .camera_cfg import CameraCfg + + +class Camera(SensorBase): + r"""The camera sensor for acquiring visual data. + + This class wraps over the `UsdGeom Camera`_ for providing a consistent API for acquiring visual data. + It ensures that the camera follows the ROS convention for the coordinate system. + + Summarizing from the `replicator extension`_, the following sensor types are supported: + + - ``"rgb"``: A 3-channel rendered color image. + - ``"rgba"``: A 4-channel rendered color image with alpha channel. + - ``"distance_to_camera"``: An image containing the distance to camera optical center. + - ``"distance_to_image_plane"``: An image containing distances of 3D points from camera plane along camera's z-axis. + - ``"depth"``: The same as ``"distance_to_image_plane"``. + - ``"normals"``: An image containing the local surface normal vectors at each pixel. + - ``"motion_vectors"``: An image containing the motion vector data at each pixel. + - ``"semantic_segmentation"``: The semantic segmentation data. + - ``"instance_segmentation_fast"``: The instance segmentation data. + - ``"instance_id_segmentation_fast"``: The instance id segmentation data. + + .. note:: + Currently the following sensor types are not supported in a "view" format: + + - ``"instance_segmentation"``: The instance segmentation data. Please use the fast counterparts instead. + - ``"instance_id_segmentation"``: The instance id segmentation data. Please use the fast counterparts instead. + - ``"bounding_box_2d_tight"``: The tight 2D bounding box data (only contains non-occluded regions). + - ``"bounding_box_2d_tight_fast"``: The tight 2D bounding box data (only contains non-occluded regions). + - ``"bounding_box_2d_loose"``: The loose 2D bounding box data (contains occluded regions). + - ``"bounding_box_2d_loose_fast"``: The loose 2D bounding box data (contains occluded regions). + - ``"bounding_box_3d"``: The 3D view space bounding box data. + - ``"bounding_box_3d_fast"``: The 3D view space bounding box data. + + .. _replicator extension: https://docs.omniverse.nvidia.com/extensions/latest/ext_replicator/annotators_details.html#annotator-output + .. _USDGeom Camera: https://graphics.pixar.com/usd/docs/api/class_usd_geom_camera.html + + """ + + cfg: CameraCfg + """The configuration parameters.""" + + UNSUPPORTED_TYPES: set[str] = { + "instance_id_segmentation", + "instance_segmentation", + "bounding_box_2d_tight", + "bounding_box_2d_loose", + "bounding_box_3d", + "bounding_box_2d_tight_fast", + "bounding_box_2d_loose_fast", + "bounding_box_3d_fast", + } + """The set of sensor types that are not supported by the camera class.""" + + def __init__(self, cfg: CameraCfg): + """Initializes the camera sensor. + + Args: + cfg: The configuration parameters. + + Raises: + RuntimeError: If no camera prim is found at the given path. + ValueError: If the provided data types are not supported by the camera. + """ + # check if sensor path is valid + # note: currently we do not handle environment indices if there is a regex pattern in the leaf + # For example, if the prim path is "/World/Sensor_[1,2]". + sensor_path = cfg.prim_path.split("/")[-1] + sensor_path_is_regex = re.match(r"^[a-zA-Z0-9/_]+$", sensor_path) is None + if sensor_path_is_regex: + raise RuntimeError( + f"Invalid prim path for the camera sensor: {self.cfg.prim_path}." + "\n\tHint: Please ensure that the prim path does not contain any regex patterns in the leaf." + ) + # perform check on supported data types + self._check_supported_data_types(cfg) + # initialize base class + super().__init__(cfg) + + # toggle rendering of rtx sensors as True + # this flag is read by SimulationContext to determine if rtx sensors should be rendered + carb_settings_iface = carb.settings.get_settings() + carb_settings_iface.set_bool("/isaaclab/render/rtx_sensors", True) + + # spawn the asset + if self.cfg.spawn is not None: + # compute the rotation offset + rot = torch.tensor(self.cfg.offset.rot, dtype=torch.float32, device="cpu").unsqueeze(0) + rot_offset = convert_camera_frame_orientation_convention( + rot, origin=self.cfg.offset.convention, target="opengl" + ) + rot_offset = rot_offset.squeeze(0).cpu().numpy() + # ensure vertical aperture is set, otherwise replace with default for squared pixels + if self.cfg.spawn.vertical_aperture is None: + self.cfg.spawn.vertical_aperture = self.cfg.spawn.horizontal_aperture * self.cfg.height / self.cfg.width + # spawn the asset + self.cfg.spawn.func( + self.cfg.prim_path, self.cfg.spawn, translation=self.cfg.offset.pos, orientation=rot_offset + ) + # check that spawn was successful + matching_prims = sim_utils.find_matching_prims(self.cfg.prim_path) + if len(matching_prims) == 0: + raise RuntimeError(f"Could not find prim with path {self.cfg.prim_path}.") + + # UsdGeom Camera prim for the sensor + self._sensor_prims: list[UsdGeom.Camera] = list() + # Create empty variables for storing output data + self._data = CameraData() + + def __del__(self): + """Unsubscribes from callbacks and detach from the replicator registry.""" + # unsubscribe callbacks + super().__del__() + # delete from replicator registry + for _, annotators in self._rep_registry.items(): + for annotator, render_product_path in zip(annotators, self._render_product_paths): + annotator.detach([render_product_path]) + annotator = None + + def __str__(self) -> str: + """Returns: A string containing information about the instance.""" + # message for class + return ( + f"Camera @ '{self.cfg.prim_path}': \n" + f"\tdata types : {list(self.data.output.keys())} \n" + f"\tsemantic filter : {self.cfg.semantic_filter}\n" + f"\tcolorize semantic segm. : {self.cfg.colorize_semantic_segmentation}\n" + f"\tcolorize instance segm. : {self.cfg.colorize_instance_segmentation}\n" + f"\tcolorize instance id segm.: {self.cfg.colorize_instance_id_segmentation}\n" + f"\tupdate period (s): {self.cfg.update_period}\n" + f"\tshape : {self.image_shape}\n" + f"\tnumber of sensors : {self._view.count}" + ) + + """ + Properties + """ + + @property + def num_instances(self) -> int: + return self._view.count + + @property + def data(self) -> CameraData: + # update sensors if needed + self._update_outdated_buffers() + # return the data + return self._data + + @property + def frame(self) -> torch.tensor: + """Frame number when the measurement took place.""" + return self._frame + + @property + def render_product_paths(self) -> list[str]: + """The path of the render products for the cameras. + + This can be used via replicator interfaces to attach to writes or external annotator registry. + """ + return self._render_product_paths + + @property + def image_shape(self) -> tuple[int, int]: + """A tuple containing (height, width) of the camera sensor.""" + return (self.cfg.height, self.cfg.width) + + """ + Configuration + """ + + def set_intrinsic_matrices( + self, matrices: torch.Tensor, focal_length: float | None = None, env_ids: Sequence[int] | None = None + ): + """Set parameters of the USD camera from its intrinsic matrix. + + The intrinsic matrix is used to set the following parameters to the USD camera: + + - ``focal_length``: The focal length of the camera. + - ``horizontal_aperture``: The horizontal aperture of the camera. + - ``vertical_aperture``: The vertical aperture of the camera. + - ``horizontal_aperture_offset``: The horizontal offset of the camera. + - ``vertical_aperture_offset``: The vertical offset of the camera. + + .. warning:: + + Due to limitations of Omniverse camera, we need to assume that the camera is a spherical lens, + i.e. has square pixels, and the optical center is centered at the camera eye. If this assumption + is not true in the input intrinsic matrix, then the camera will not set up correctly. + + Args: + matrices: The intrinsic matrices for the camera. Shape is (N, 3, 3). + focal_length: Perspective focal length (in cm) used to calculate pixel size. Defaults to None. If None, + focal_length will be calculated 1 / width. + env_ids: A sensor ids to manipulate. Defaults to None, which means all sensor indices. + """ + # resolve env_ids + if env_ids is None: + env_ids = self._ALL_INDICES + # convert matrices to numpy tensors + if isinstance(matrices, torch.Tensor): + matrices = matrices.cpu().numpy() + else: + matrices = np.asarray(matrices, dtype=float) + # iterate over env_ids + for i, intrinsic_matrix in zip(env_ids, matrices): + + height, width = self.image_shape + + params = sensor_utils.convert_camera_intrinsics_to_usd( + intrinsic_matrix=intrinsic_matrix.reshape(-1), height=height, width=width, focal_length=focal_length + ) + + # change data for corresponding camera index + sensor_prim = self._sensor_prims[i] + # set parameters for camera + for param_name, param_value in params.items(): + # convert to camel case (CC) + param_name = to_camel_case(param_name, to="CC") + # get attribute from the class + param_attr = getattr(sensor_prim, f"Get{param_name}Attr") + # set value + # note: We have to do it this way because the camera might be on a different + # layer (default cameras are on session layer), and this is the simplest + # way to set the property on the right layer. + omni.usd.set_prop_val(param_attr(), param_value) + # update the internal buffers + self._update_intrinsic_matrices(env_ids) + + """ + Operations - Set pose. + """ + + def set_world_poses( + self, + positions: torch.Tensor | None = None, + orientations: torch.Tensor | None = None, + env_ids: Sequence[int] | None = None, + convention: Literal["opengl", "ros", "world"] = "ros", + ): + r"""Set the pose of the camera w.r.t. the world frame using specified convention. + + Since different fields use different conventions for camera orientations, the method allows users to + set the camera poses in the specified convention. Possible conventions are: + + - :obj:`"opengl"` - forward axis: -Z - up axis +Y - Offset is applied in the OpenGL (Usd.Camera) convention + - :obj:`"ros"` - forward axis: +Z - up axis -Y - Offset is applied in the ROS convention + - :obj:`"world"` - forward axis: +X - up axis +Z - Offset is applied in the World Frame convention + + See :meth:`isaaclab.sensors.camera.utils.convert_camera_frame_orientation_convention` for more details + on the conventions. + + Args: + positions: The cartesian coordinates (in meters). Shape is (N, 3). + Defaults to None, in which case the camera position in not changed. + orientations: The quaternion orientation in (w, x, y, z). Shape is (N, 4). + Defaults to None, in which case the camera orientation in not changed. + env_ids: A sensor ids to manipulate. Defaults to None, which means all sensor indices. + convention: The convention in which the poses are fed. Defaults to "ros". + + Raises: + RuntimeError: If the camera prim is not set. Need to call :meth:`initialize` method first. + """ + # resolve env_ids + if env_ids is None: + env_ids = self._ALL_INDICES + # convert to backend tensor + if positions is not None: + if isinstance(positions, np.ndarray): + positions = torch.from_numpy(positions).to(device=self._device) + elif not isinstance(positions, torch.Tensor): + positions = torch.tensor(positions, device=self._device) + # convert rotation matrix from input convention to OpenGL + if orientations is not None: + if isinstance(orientations, np.ndarray): + orientations = torch.from_numpy(orientations).to(device=self._device) + elif not isinstance(orientations, torch.Tensor): + orientations = torch.tensor(orientations, device=self._device) + orientations = convert_camera_frame_orientation_convention(orientations, origin=convention, target="opengl") + # set the pose + self._view.set_world_poses(positions, orientations, env_ids) + + def set_world_poses_from_view( + self, eyes: torch.Tensor, targets: torch.Tensor, env_ids: Sequence[int] | None = None + ): + """Set the poses of the camera from the eye position and look-at target position. + + Args: + eyes: The positions of the camera's eye. Shape is (N, 3). + targets: The target locations to look at. Shape is (N, 3). + env_ids: A sensor ids to manipulate. Defaults to None, which means all sensor indices. + + Raises: + RuntimeError: If the camera prim is not set. Need to call :meth:`initialize` method first. + NotImplementedError: If the stage up-axis is not "Y" or "Z". + """ + # resolve env_ids + if env_ids is None: + env_ids = self._ALL_INDICES + # get up axis of current stage + up_axis = stage_utils.get_stage_up_axis() + # set camera poses using the view + orientations = quat_from_matrix(create_rotation_matrix_from_view(eyes, targets, up_axis, device=self._device)) + self._view.set_world_poses(eyes, orientations, env_ids) + + """ + Operations + """ + + def reset(self, env_ids: Sequence[int] | None = None): + if not self._is_initialized: + raise RuntimeError( + "Camera could not be initialized. Please ensure --enable_cameras is used to enable rendering." + ) + # reset the timestamps + super().reset(env_ids) + # resolve None + # note: cannot do smart indexing here since we do a for loop over data. + if env_ids is None: + env_ids = self._ALL_INDICES + # reset the data + # note: this recomputation is useful if one performs events such as randomizations on the camera poses. + self._update_poses(env_ids) + # Reset the frame count + self._frame[env_ids] = 0 + + """ + Implementation. + """ + + def _initialize_impl(self): + """Initializes the sensor handles and internal buffers. + + This function creates handles and registers the provided data types with the replicator registry to + be able to access the data from the sensor. It also initializes the internal buffers to store the data. + + Raises: + RuntimeError: If the number of camera prims in the view does not match the number of environments. + RuntimeError: If replicator was not found. + """ + carb_settings_iface = carb.settings.get_settings() + if not carb_settings_iface.get("/isaaclab/cameras_enabled"): + raise RuntimeError( + "A camera was spawned without the --enable_cameras flag. Please use --enable_cameras to enable" + " rendering." + ) + + import omni.replicator.core as rep + from omni.syntheticdata.scripts.SyntheticData import SyntheticData + + # Initialize parent class + super()._initialize_impl() + # Create a view for the sensor + self._view = XFormPrim(self.cfg.prim_path, reset_xform_properties=False) + self._view.initialize() + # Check that sizes are correct + if self._view.count != self._num_envs: + raise RuntimeError( + f"Number of camera prims in the view ({self._view.count}) does not match" + f" the number of environments ({self._num_envs})." + ) + + # Create all env_ids buffer + self._ALL_INDICES = torch.arange(self._view.count, device=self._device, dtype=torch.long) + # Create frame count buffer + self._frame = torch.zeros(self._view.count, device=self._device, dtype=torch.long) + + # Attach the sensor data types to render node + self._render_product_paths: list[str] = list() + self._rep_registry: dict[str, list[rep.annotators.Annotator]] = {name: list() for name in self.cfg.data_types} + + # Convert all encapsulated prims to Camera + for cam_prim_path in self._view.prim_paths: + # Get camera prim + cam_prim = self.stage.GetPrimAtPath(cam_prim_path) + # Check if prim is a camera + if not cam_prim.IsA(UsdGeom.Camera): + raise RuntimeError(f"Prim at path '{cam_prim_path}' is not a Camera.") + # Add to list + sensor_prim = UsdGeom.Camera(cam_prim) + self._sensor_prims.append(sensor_prim) + + # Get render product + # From Isaac Sim 2023.1 onwards, render product is a HydraTexture so we need to extract the path + render_prod_path = rep.create.render_product(cam_prim_path, resolution=(self.cfg.width, self.cfg.height)) + if not isinstance(render_prod_path, str): + render_prod_path = render_prod_path.path + self._render_product_paths.append(render_prod_path) + + # Check if semantic types or semantic filter predicate is provided + if isinstance(self.cfg.semantic_filter, list): + semantic_filter_predicate = ":*; ".join(self.cfg.semantic_filter) + ":*" + elif isinstance(self.cfg.semantic_filter, str): + semantic_filter_predicate = self.cfg.semantic_filter + else: + raise ValueError(f"Semantic types must be a list or a string. Received: {self.cfg.semantic_filter}.") + # set the semantic filter predicate + # copied from rep.scripts.writes_default.basic_writer.py + SyntheticData.Get().set_instance_mapping_semantic_filter(semantic_filter_predicate) + + # Iterate over each data type and create annotator + # TODO: This will move out of the loop once Replicator supports multiple render products within a single + # annotator, i.e.: rep_annotator.attach(self._render_product_paths) + for name in self.cfg.data_types: + # note: we are verbose here to make it easier to understand the code. + # if colorize is true, the data is mapped to colors and a uint8 4 channel image is returned. + # if colorize is false, the data is returned as a uint32 image with ids as values. + if name == "semantic_segmentation": + init_params = { + "colorize": self.cfg.colorize_semantic_segmentation, + "mapping": json.dumps(self.cfg.semantic_segmentation_mapping), + } + elif name == "instance_segmentation_fast": + init_params = {"colorize": self.cfg.colorize_instance_segmentation} + elif name == "instance_id_segmentation_fast": + init_params = {"colorize": self.cfg.colorize_instance_id_segmentation} + else: + init_params = None + + # Resolve device name + if "cuda" in self._device: + device_name = self._device.split(":")[0] + else: + device_name = "cpu" + + # Map special cases to their corresponding annotator names + special_cases = {"rgba": "rgb", "depth": "distance_to_image_plane"} + # Get the annotator name, falling back to the original name if not a special case + annotator_name = special_cases.get(name, name) + # Create the annotator node + rep_annotator = rep.AnnotatorRegistry.get_annotator(annotator_name, init_params, device=device_name) + + # attach annotator to render product + rep_annotator.attach(render_prod_path) + # add to registry + self._rep_registry[name].append(rep_annotator) + + # Create internal buffers + self._create_buffers() + self._update_intrinsic_matrices(self._ALL_INDICES) + + def _update_buffers_impl(self, env_ids: Sequence[int]): + # Increment frame count + self._frame[env_ids] += 1 + # -- pose + if self.cfg.update_latest_camera_pose: + self._update_poses(env_ids) + # -- read the data from annotator registry + # check if buffer is called for the first time. If so then, allocate the memory + if len(self._data.output) == 0: + # this is the first time buffer is called + # it allocates memory for all the sensors + self._create_annotator_data() + else: + # iterate over all the data types + for name, annotators in self._rep_registry.items(): + # iterate over all the annotators + for index in env_ids: + # get the output + output = annotators[index].get_data() + # process the output + data, info = self._process_annotator_output(name, output) + # add data to output + self._data.output[name][index] = data + # add info to output + self._data.info[index][name] = info + # NOTE: The `distance_to_camera` annotator returns the distance to the camera optical center. However, + # the replicator depth clipping is applied w.r.t. to the image plane which may result in values + # larger than the clipping range in the output. We apply an additional clipping to ensure values + # are within the clipping range for all the annotators. + if name == "distance_to_camera": + self._data.output[name][self._data.output[name] > self.cfg.spawn.clipping_range[1]] = torch.inf + # apply defined clipping behavior + if ( + name == "distance_to_camera" or name == "distance_to_image_plane" + ) and self.cfg.depth_clipping_behavior != "none": + self._data.output[name][torch.isinf(self._data.output[name])] = ( + 0.0 if self.cfg.depth_clipping_behavior == "zero" else self.cfg.spawn.clipping_range[1] + ) + + """ + Private Helpers + """ + + def _check_supported_data_types(self, cfg: CameraCfg): + """Checks if the data types are supported by the ray-caster camera.""" + # check if there is any intersection in unsupported types + # reason: these use np structured data types which we can't yet convert to torch tensor + common_elements = set(cfg.data_types) & Camera.UNSUPPORTED_TYPES + if common_elements: + # provide alternative fast counterparts + fast_common_elements = [] + for item in common_elements: + if "instance_segmentation" in item or "instance_id_segmentation" in item: + fast_common_elements.append(item + "_fast") + # raise error + raise ValueError( + f"Camera class does not support the following sensor types: {common_elements}." + "\n\tThis is because these sensor types output numpy structured data types which" + "can't be converted to torch tensors easily." + "\n\tHint: If you need to work with these sensor types, we recommend using their fast counterparts." + f"\n\t\tFast counterparts: {fast_common_elements}" + ) + + def _create_buffers(self): + """Create buffers for storing data.""" + # create the data object + # -- pose of the cameras + self._data.pos_w = torch.zeros((self._view.count, 3), device=self._device) + self._data.quat_w_world = torch.zeros((self._view.count, 4), device=self._device) + # -- intrinsic matrix + self._data.intrinsic_matrices = torch.zeros((self._view.count, 3, 3), device=self._device) + self._data.image_shape = self.image_shape + # -- output data + # lazy allocation of data dictionary + # since the size of the output data is not known in advance, we leave it as None + # the memory will be allocated when the buffer() function is called for the first time. + self._data.output = {} + self._data.info = [{name: None for name in self.cfg.data_types} for _ in range(self._view.count)] + + def _update_intrinsic_matrices(self, env_ids: Sequence[int]): + """Compute camera's matrix of intrinsic parameters. + + Also called calibration matrix. This matrix works for linear depth images. We assume square pixels. + + Note: + The calibration matrix projects points in the 3D scene onto an imaginary screen of the camera. + The coordinates of points on the image plane are in the homogeneous representation. + """ + # iterate over all cameras + for i in env_ids: + # Get corresponding sensor prim + sensor_prim = self._sensor_prims[i] + # get camera parameters + # currently rendering does not use aperture offsets or vertical aperture + focal_length = sensor_prim.GetFocalLengthAttr().Get() + horiz_aperture = sensor_prim.GetHorizontalApertureAttr().Get() + + # get viewport parameters + height, width = self.image_shape + # extract intrinsic parameters + f_x = (width * focal_length) / horiz_aperture + f_y = f_x + c_x = width * 0.5 + c_y = height * 0.5 + # create intrinsic matrix for depth linear + self._data.intrinsic_matrices[i, 0, 0] = f_x + self._data.intrinsic_matrices[i, 0, 2] = c_x + self._data.intrinsic_matrices[i, 1, 1] = f_y + self._data.intrinsic_matrices[i, 1, 2] = c_y + self._data.intrinsic_matrices[i, 2, 2] = 1 + + def _update_poses(self, env_ids: Sequence[int]): + """Computes the pose of the camera in the world frame with ROS convention. + + This methods uses the ROS convention to resolve the input pose. In this convention, + we assume that the camera front-axis is +Z-axis and up-axis is -Y-axis. + + Returns: + A tuple of the position (in meters) and quaternion (w, x, y, z). + """ + # check camera prim exists + if len(self._sensor_prims) == 0: + raise RuntimeError("Camera prim is None. Please call 'sim.play()' first.") + + # get the poses from the view + poses, quat = self._view.get_world_poses(env_ids) + self._data.pos_w[env_ids] = poses + self._data.quat_w_world[env_ids] = convert_camera_frame_orientation_convention( + quat, origin="opengl", target="world" + ) + + def _create_annotator_data(self): + """Create the buffers to store the annotator data. + + We create a buffer for each annotator and store the data in a dictionary. Since the data + shape is not known beforehand, we create a list of buffers and concatenate them later. + + This is an expensive operation and should be called only once. + """ + # add data from the annotators + for name, annotators in self._rep_registry.items(): + # create a list to store the data for each annotator + data_all_cameras = list() + # iterate over all the annotators + for index in self._ALL_INDICES: + # get the output + output = annotators[index].get_data() + # process the output + data, info = self._process_annotator_output(name, output) + # append the data + data_all_cameras.append(data) + # store the info + self._data.info[index][name] = info + # concatenate the data along the batch dimension + self._data.output[name] = torch.stack(data_all_cameras, dim=0) + # NOTE: `distance_to_camera` and `distance_to_image_plane` are not both clipped to the maximum defined + # in the clipping range. The clipping is applied only to `distance_to_image_plane` and then both + # outputs are only clipped where the values in `distance_to_image_plane` exceed the threshold. To + # have a unified behavior between all cameras, we clip both outputs to the maximum value defined. + if name == "distance_to_camera": + self._data.output[name][self._data.output[name] > self.cfg.spawn.clipping_range[1]] = torch.inf + # clip the data if needed + if ( + name == "distance_to_camera" or name == "distance_to_image_plane" + ) and self.cfg.depth_clipping_behavior != "none": + self._data.output[name][torch.isinf(self._data.output[name])] = ( + 0.0 if self.cfg.depth_clipping_behavior == "zero" else self.cfg.spawn.clipping_range[1] + ) + + def _process_annotator_output(self, name: str, output: Any) -> tuple[torch.tensor, dict | None]: + """Process the annotator output. + + This function is called after the data has been collected from all the cameras. + """ + # extract info and data from the output + if isinstance(output, dict): + data = output["data"] + info = output["info"] + else: + data = output + info = None + # convert data into torch tensor + data = convert_to_torch(data, device=self.device) + + # process data for different segmentation types + # Note: Replicator returns raw buffers of dtype int32 for segmentation types + # so we need to convert them to uint8 4 channel images for colorized types + height, width = self.image_shape + if name == "semantic_segmentation": + if self.cfg.colorize_semantic_segmentation: + data = data.view(torch.uint8).reshape(height, width, -1) + else: + data = data.view(height, width, 1) + elif name == "instance_segmentation_fast": + if self.cfg.colorize_instance_segmentation: + data = data.view(torch.uint8).reshape(height, width, -1) + else: + data = data.view(height, width, 1) + elif name == "instance_id_segmentation_fast": + if self.cfg.colorize_instance_id_segmentation: + data = data.view(torch.uint8).reshape(height, width, -1) + else: + data = data.view(height, width, 1) + # make sure buffer dimensions are consistent as (H, W, C) + elif name == "distance_to_camera" or name == "distance_to_image_plane" or name == "depth": + data = data.view(height, width, 1) + # we only return the RGB channels from the RGBA output if rgb is required + # normals return (x, y, z) in first 3 channels, 4th channel is unused + elif name == "rgb" or name == "normals": + data = data[..., :3] + # motion vectors return (x, y) in first 2 channels, 3rd and 4th channels are unused + elif name == "motion_vectors": + data = data[..., :2] + + # return the data and info + return data, info + + """ + Internal simulation callbacks. + """ + + def _invalidate_initialize_callback(self, event): + """Invalidates the scene elements.""" + # call parent + super()._invalidate_initialize_callback(event) + # set all existing views to None to invalidate them + self._view = None diff --git a/source/isaaclab/isaaclab/newton/sensors/camera/camera_cfg.py b/source/isaaclab/isaaclab/newton/sensors/camera/camera_cfg.py new file mode 100644 index 00000000000..a123bd00b57 --- /dev/null +++ b/source/isaaclab/isaaclab/newton/sensors/camera/camera_cfg.py @@ -0,0 +1,142 @@ +# Copyright (c) 2022-2025, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +from dataclasses import MISSING +from typing import Literal + +from isaaclab.sim import FisheyeCameraCfg, PinholeCameraCfg +from isaaclab.utils import configclass + +from ..sensor_base_cfg import SensorBaseCfg +from .camera import Camera + + +@configclass +class CameraCfg(SensorBaseCfg): + """Configuration for a camera sensor.""" + + @configclass + class OffsetCfg: + """The offset pose of the sensor's frame from the sensor's parent frame.""" + + pos: tuple[float, float, float] = (0.0, 0.0, 0.0) + """Translation w.r.t. the parent frame. Defaults to (0.0, 0.0, 0.0).""" + + rot: tuple[float, float, float, float] = (1.0, 0.0, 0.0, 0.0) + """Quaternion rotation (w, x, y, z) w.r.t. the parent frame. Defaults to (1.0, 0.0, 0.0, 0.0).""" + + convention: Literal["opengl", "ros", "world"] = "ros" + """The convention in which the frame offset is applied. Defaults to "ros". + + - ``"opengl"`` - forward axis: ``-Z`` - up axis: ``+Y`` - Offset is applied in the OpenGL (Usd.Camera) convention. + - ``"ros"`` - forward axis: ``+Z`` - up axis: ``-Y`` - Offset is applied in the ROS convention. + - ``"world"`` - forward axis: ``+X`` - up axis: ``+Z`` - Offset is applied in the World Frame convention. + + """ + + class_type: type = Camera + + offset: OffsetCfg = OffsetCfg() + """The offset pose of the sensor's frame from the sensor's parent frame. Defaults to identity. + + Note: + The parent frame is the frame the sensor attaches to. For example, the parent frame of a + camera at path ``/World/envs/env_0/Robot/Camera`` is ``/World/envs/env_0/Robot``. + """ + + spawn: PinholeCameraCfg | FisheyeCameraCfg | None = MISSING + """Spawn configuration for the asset. + + If None, then the prim is not spawned by the asset. Instead, it is assumed that the + asset is already present in the scene. + """ + + depth_clipping_behavior: Literal["max", "zero", "none"] = "none" + """Clipping behavior for the camera for values exceed the maximum value. Defaults to "none". + + - ``"max"``: Values are clipped to the maximum value. + - ``"zero"``: Values are clipped to zero. + - ``"none``: No clipping is applied. Values will be returned as ``inf``. + """ + + data_types: list[str] = ["rgb"] + """List of sensor names/types to enable for the camera. Defaults to ["rgb"]. + + Please refer to the :class:`Camera` class for a list of available data types. + """ + + width: int = MISSING + """Width of the image in pixels.""" + + height: int = MISSING + """Height of the image in pixels.""" + + update_latest_camera_pose: bool = False + """Whether to update the latest camera pose when fetching the camera's data. Defaults to False. + + If True, the latest camera pose is updated in the camera's data which will slow down performance + due to the use of :class:`XformPrimView`. + If False, the pose of the camera during initialization is returned. + """ + + semantic_filter: str | list[str] = "*:*" + """A string or a list specifying a semantic filter predicate. Defaults to ``"*:*"``. + + If a string, it should be a disjunctive normal form of (semantic type, labels). For examples: + + * ``"typeA : labelA & !labelB | labelC , typeB: labelA ; typeC: labelE"``: + All prims with semantic type "typeA" and label "labelA" but not "labelB" or with label "labelC". + Also, all prims with semantic type "typeB" and label "labelA", or with semantic type "typeC" and label "labelE". + * ``"typeA : * ; * : labelA"``: All prims with semantic type "typeA" or with label "labelA" + + If a list of strings, each string should be a semantic type. The segmentation for prims with + semantics of the specified types will be retrieved. For example, if the list is ["class"], only + the segmentation for prims with semantics of type "class" will be retrieved. + + .. seealso:: + + For more information on the semantics filter, see the documentation on `Replicator Semantics Schema Editor`_. + + .. _Replicator Semantics Schema Editor: https://docs.omniverse.nvidia.com/extensions/latest/ext_replicator/semantics_schema_editor.html#semantics-filtering + """ + + colorize_semantic_segmentation: bool = True + """Whether to colorize the semantic segmentation images. Defaults to True. + + If True, semantic segmentation is converted to an image where semantic IDs are mapped to colors + and returned as a ``uint8`` 4-channel array. If False, the output is returned as a ``int32`` array. + """ + + colorize_instance_id_segmentation: bool = True + """Whether to colorize the instance ID segmentation images. Defaults to True. + + If True, instance id segmentation is converted to an image where instance IDs are mapped to colors. + and returned as a ``uint8`` 4-channel array. If False, the output is returned as a ``int32`` array. + """ + + colorize_instance_segmentation: bool = True + """Whether to colorize the instance ID segmentation images. Defaults to True. + + If True, instance segmentation is converted to an image where instance IDs are mapped to colors. + and returned as a ``uint8`` 4-channel array. If False, the output is returned as a ``int32`` array. + """ + + semantic_segmentation_mapping: dict = {} + """Dictionary mapping semantics to specific colours + + Eg. + + .. code-block:: python + + { + "class:cube_1": (255, 36, 66, 255), + "class:cube_2": (255, 184, 48, 255), + "class:cube_3": (55, 255, 139, 255), + "class:table": (255, 237, 218, 255), + "class:ground": (100, 100, 100, 255), + "class:robot": (61, 178, 255, 255), + } + + """ diff --git a/source/isaaclab/isaaclab/newton/sensors/camera/camera_data.py b/source/isaaclab/isaaclab/newton/sensors/camera/camera_data.py new file mode 100644 index 00000000000..dfcc780b4d1 --- /dev/null +++ b/source/isaaclab/isaaclab/newton/sensors/camera/camera_data.py @@ -0,0 +1,91 @@ +# Copyright (c) 2022-2025, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +import torch +from dataclasses import dataclass +from typing import Any + +from isaaclab.utils.math import convert_camera_frame_orientation_convention + + +@dataclass +class CameraData: + """Data container for the camera sensor.""" + + ## + # Frame state. + ## + + pos_w: torch.Tensor = None + """Position of the sensor origin in world frame, following ROS convention. + + Shape is (N, 3) where N is the number of sensors. + """ + + quat_w_world: torch.Tensor = None + """Quaternion orientation `(w, x, y, z)` of the sensor origin in world frame, following the world coordinate frame + + .. note:: + World frame convention follows the camera aligned with forward axis +X and up axis +Z. + + Shape is (N, 4) where N is the number of sensors. + """ + + ## + # Camera data + ## + + image_shape: tuple[int, int] = None + """A tuple containing (height, width) of the camera sensor.""" + + intrinsic_matrices: torch.Tensor = None + """The intrinsic matrices for the camera. + + Shape is (N, 3, 3) where N is the number of sensors. + """ + + output: dict[str, torch.Tensor] = None + """The retrieved sensor data with sensor types as key. + + The format of the data is available in the `Replicator Documentation`_. For semantic-based data, + this corresponds to the ``"data"`` key in the output of the sensor. + + .. _Replicator Documentation: https://docs.omniverse.nvidia.com/prod_extensions/prod_extensions/ext_replicator/annotators_details.html#annotator-output + """ + + info: list[dict[str, Any]] = None + """The retrieved sensor info with sensor types as key. + + This contains extra information provided by the sensor such as semantic segmentation label mapping, prim paths. + For semantic-based data, this corresponds to the ``"info"`` key in the output of the sensor. For other sensor + types, the info is empty. + """ + + ## + # Additional Frame orientation conventions + ## + + @property + def quat_w_ros(self) -> torch.Tensor: + """Quaternion orientation `(w, x, y, z)` of the sensor origin in the world frame, following ROS convention. + + .. note:: + ROS convention follows the camera aligned with forward axis +Z and up axis -Y. + + Shape is (N, 4) where N is the number of sensors. + """ + return convert_camera_frame_orientation_convention(self.quat_w_world, origin="world", target="ros") + + @property + def quat_w_opengl(self) -> torch.Tensor: + """Quaternion orientation `(w, x, y, z)` of the sensor origin in the world frame, following + Opengl / USD Camera convention. + + .. note:: + OpenGL convention follows the camera aligned with forward axis -Z and up axis +Y. + + Shape is (N, 4) where N is the number of sensors. + """ + return convert_camera_frame_orientation_convention(self.quat_w_world, origin="world", target="opengl") diff --git a/source/isaaclab/isaaclab/newton/sensors/camera/tiled_camera.py b/source/isaaclab/isaaclab/newton/sensors/camera/tiled_camera.py new file mode 100644 index 00000000000..a0600db4b11 --- /dev/null +++ b/source/isaaclab/isaaclab/newton/sensors/camera/tiled_camera.py @@ -0,0 +1,420 @@ +# Copyright (c) 2022-2025, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +from __future__ import annotations + +import json +import math +import numpy as np +import torch +from collections.abc import Sequence +from typing import TYPE_CHECKING, Any + +import carb +import warp as wp +from isaacsim.core.prims import XFormPrim +from pxr import UsdGeom + +from isaaclab.utils.warp.kernels import reshape_tiled_image + +from ..sensor_base import SensorBase +from .camera import Camera + +if TYPE_CHECKING: + from .tiled_camera_cfg import TiledCameraCfg + + +class TiledCamera(Camera): + r"""The tiled rendering based camera sensor for acquiring the same data as the Camera class. + + This class inherits from the :class:`Camera` class but uses the tiled-rendering API to acquire + the visual data. Tiled-rendering concatenates the rendered images from multiple cameras into a single image. + This allows for rendering multiple cameras in parallel and is useful for rendering large scenes with multiple + cameras efficiently. + + The following sensor types are supported: + + - ``"rgb"``: A 3-channel rendered color image. + - ``"rgba"``: A 4-channel rendered color image with alpha channel. + - ``"distance_to_camera"``: An image containing the distance to camera optical center. + - ``"distance_to_image_plane"``: An image containing distances of 3D points from camera plane along camera's z-axis. + - ``"depth"``: Alias for ``"distance_to_image_plane"``. + - ``"normals"``: An image containing the local surface normal vectors at each pixel. + - ``"motion_vectors"``: An image containing the motion vector data at each pixel. + - ``"semantic_segmentation"``: The semantic segmentation data. + - ``"instance_segmentation_fast"``: The instance segmentation data. + - ``"instance_id_segmentation_fast"``: The instance id segmentation data. + + .. note:: + Currently the following sensor types are not supported in a "view" format: + + - ``"instance_segmentation"``: The instance segmentation data. Please use the fast counterparts instead. + - ``"instance_id_segmentation"``: The instance id segmentation data. Please use the fast counterparts instead. + - ``"bounding_box_2d_tight"``: The tight 2D bounding box data (only contains non-occluded regions). + - ``"bounding_box_2d_tight_fast"``: The tight 2D bounding box data (only contains non-occluded regions). + - ``"bounding_box_2d_loose"``: The loose 2D bounding box data (contains occluded regions). + - ``"bounding_box_2d_loose_fast"``: The loose 2D bounding box data (contains occluded regions). + - ``"bounding_box_3d"``: The 3D view space bounding box data. + - ``"bounding_box_3d_fast"``: The 3D view space bounding box data. + + .. _replicator extension: https://docs.omniverse.nvidia.com/extensions/latest/ext_replicator/annotators_details.html#annotator-output + .. _USDGeom Camera: https://graphics.pixar.com/usd/docs/api/class_usd_geom_camera.html + + .. versionadded:: v1.0.0 + + This feature is available starting from Isaac Sim 4.2. Before this version, the tiled rendering APIs + were not available. + + """ + + cfg: TiledCameraCfg + """The configuration parameters.""" + + def __init__(self, cfg: TiledCameraCfg): + """Initializes the tiled camera sensor. + + Args: + cfg: The configuration parameters. + + Raises: + RuntimeError: If no camera prim is found at the given path. + RuntimeError: If Isaac Sim version < 4.2 + ValueError: If the provided data types are not supported by the camera. + """ + super().__init__(cfg) + + def __del__(self): + """Unsubscribes from callbacks and detach from the replicator registry.""" + # unsubscribe from callbacks + SensorBase.__del__(self) + # detach from the replicator registry + for annotator in self._annotators.values(): + annotator.detach(self.render_product_paths) + + def __str__(self) -> str: + """Returns: A string containing information about the instance.""" + # message for class + return ( + f"Tiled Camera @ '{self.cfg.prim_path}': \n" + f"\tdata types : {list(self.data.output.keys())} \n" + f"\tsemantic filter : {self.cfg.semantic_filter}\n" + f"\tcolorize semantic segm. : {self.cfg.colorize_semantic_segmentation}\n" + f"\tcolorize instance segm. : {self.cfg.colorize_instance_segmentation}\n" + f"\tcolorize instance id segm.: {self.cfg.colorize_instance_id_segmentation}\n" + f"\tupdate period (s): {self.cfg.update_period}\n" + f"\tshape : {self.image_shape}\n" + f"\tnumber of sensors : {self._view.count}" + ) + + """ + Operations + """ + + def reset(self, env_ids: Sequence[int] | None = None): + if not self._is_initialized: + raise RuntimeError( + "TiledCamera could not be initialized. Please ensure --enable_cameras is used to enable rendering." + ) + # reset the timestamps + SensorBase.reset(self, env_ids) + # resolve None + if env_ids is None: + env_ids = slice(None) + # reset the frame count + self._frame[env_ids] = 0 + + """ + Implementation. + """ + + def _initialize_impl(self): + """Initializes the sensor handles and internal buffers. + + This function creates handles and registers the provided data types with the replicator registry to + be able to access the data from the sensor. It also initializes the internal buffers to store the data. + + Raises: + RuntimeError: If the number of camera prims in the view does not match the number of environments. + RuntimeError: If replicator was not found. + """ + carb_settings_iface = carb.settings.get_settings() + if not carb_settings_iface.get("/isaaclab/cameras_enabled"): + raise RuntimeError( + "A camera was spawned without the --enable_cameras flag. Please use --enable_cameras to enable" + " rendering." + ) + + import omni.replicator.core as rep + + # Initialize parent class + SensorBase._initialize_impl(self) + # Create a view for the sensor + self._view = XFormPrim(self.cfg.prim_path, reset_xform_properties=False) + self._view.initialize() + # Check that sizes are correct + if self._view.count != self._num_envs: + raise RuntimeError( + f"Number of camera prims in the view ({self._view.count}) does not match" + f" the number of environments ({self._num_envs})." + ) + + # Create all env_ids buffer + self._ALL_INDICES = torch.arange(self._view.count, device=self._device, dtype=torch.long) + # Create frame count buffer + self._frame = torch.zeros(self._view.count, device=self._device, dtype=torch.long) + + # Convert all encapsulated prims to Camera + for cam_prim_path in self._view.prim_paths: + # Get camera prim + cam_prim = self.stage.GetPrimAtPath(cam_prim_path) + # Check if prim is a camera + if not cam_prim.IsA(UsdGeom.Camera): + raise RuntimeError(f"Prim at path '{cam_prim_path}' is not a Camera.") + # Add to list + sensor_prim = UsdGeom.Camera(cam_prim) + self._sensor_prims.append(sensor_prim) + + # Create replicator tiled render product + rp = rep.create.render_product_tiled( + cameras=self._view.prim_paths, tile_resolution=(self.cfg.width, self.cfg.height) + ) + self._render_product_paths = [rp.path] + + # Define the annotators based on requested data types + self._annotators = dict() + for annotator_type in self.cfg.data_types: + if annotator_type == "rgba" or annotator_type == "rgb": + annotator = rep.AnnotatorRegistry.get_annotator("rgb", device=self.device, do_array_copy=False) + self._annotators["rgba"] = annotator + elif annotator_type == "depth" or annotator_type == "distance_to_image_plane": + # keep depth for backwards compatibility + annotator = rep.AnnotatorRegistry.get_annotator( + "distance_to_image_plane", device=self.device, do_array_copy=False + ) + self._annotators[annotator_type] = annotator + # note: we are verbose here to make it easier to understand the code. + # if colorize is true, the data is mapped to colors and a uint8 4 channel image is returned. + # if colorize is false, the data is returned as a uint32 image with ids as values. + else: + init_params = None + if annotator_type == "semantic_segmentation": + init_params = { + "colorize": self.cfg.colorize_semantic_segmentation, + "mapping": json.dumps(self.cfg.semantic_segmentation_mapping), + } + elif annotator_type == "instance_segmentation_fast": + init_params = {"colorize": self.cfg.colorize_instance_segmentation} + elif annotator_type == "instance_id_segmentation_fast": + init_params = {"colorize": self.cfg.colorize_instance_id_segmentation} + + annotator = rep.AnnotatorRegistry.get_annotator( + annotator_type, init_params, device=self.device, do_array_copy=False + ) + self._annotators[annotator_type] = annotator + + # Attach the annotator to the render product + for annotator in self._annotators.values(): + annotator.attach(self._render_product_paths) + + # Create internal buffers + self._create_buffers() + + def _update_buffers_impl(self, env_ids: Sequence[int]): + # Increment frame count + self._frame[env_ids] += 1 + + # update latest camera pose + if self.cfg.update_latest_camera_pose: + self._update_poses(env_ids) + + # Extract the flattened image buffer + for data_type, annotator in self._annotators.items(): + # check whether returned data is a dict (used for segmentation) + output = annotator.get_data() + if isinstance(output, dict): + tiled_data_buffer = output["data"] + self._data.info[data_type] = output["info"] + else: + tiled_data_buffer = output + + # convert data buffer to warp array + if isinstance(tiled_data_buffer, np.ndarray): + tiled_data_buffer = wp.array(tiled_data_buffer, device=self.device, dtype=wp.uint8) + else: + tiled_data_buffer = tiled_data_buffer.to(device=self.device) + + # process data for different segmentation types + # Note: Replicator returns raw buffers of dtype uint32 for segmentation types + # so we need to convert them to uint8 4 channel images for colorized types + if ( + (data_type == "semantic_segmentation" and self.cfg.colorize_semantic_segmentation) + or (data_type == "instance_segmentation_fast" and self.cfg.colorize_instance_segmentation) + or (data_type == "instance_id_segmentation_fast" and self.cfg.colorize_instance_id_segmentation) + ): + tiled_data_buffer = wp.array( + ptr=tiled_data_buffer.ptr, shape=(*tiled_data_buffer.shape, 4), dtype=wp.uint8, device=self.device + ) + + # For motion vectors, we only require the first two channels of the tiled buffer + # Note: Not doing this breaks the alignment of the data (check: https://github.com/isaac-sim/IsaacLab/issues/2003) + if data_type == "motion_vectors": + tiled_data_buffer = tiled_data_buffer[:, :, :2].contiguous() + + wp.launch( + kernel=reshape_tiled_image, + dim=(self._view.count, self.cfg.height, self.cfg.width), + inputs=[ + tiled_data_buffer.flatten(), + wp.from_torch(self._data.output[data_type]), # zero-copy alias + *list(self._data.output[data_type].shape[1:]), # height, width, num_channels + self._tiling_grid_shape()[0], # num_tiles_x + ], + device=self.device, + ) + + # alias rgb as first 3 channels of rgba + if data_type == "rgba" and "rgb" in self.cfg.data_types: + self._data.output["rgb"] = self._data.output["rgba"][..., :3] + + # NOTE: The `distance_to_camera` annotator returns the distance to the camera optical center. However, + # the replicator depth clipping is applied w.r.t. to the image plane which may result in values + # larger than the clipping range in the output. We apply an additional clipping to ensure values + # are within the clipping range for all the annotators. + if data_type == "distance_to_camera": + self._data.output[data_type][ + self._data.output[data_type] > self.cfg.spawn.clipping_range[1] + ] = torch.inf + # apply defined clipping behavior + if ( + data_type == "distance_to_camera" or data_type == "distance_to_image_plane" or data_type == "depth" + ) and self.cfg.depth_clipping_behavior != "none": + self._data.output[data_type][torch.isinf(self._data.output[data_type])] = ( + 0.0 if self.cfg.depth_clipping_behavior == "zero" else self.cfg.spawn.clipping_range[1] + ) + + """ + Private Helpers + """ + + def _check_supported_data_types(self, cfg: TiledCameraCfg): + """Checks if the data types are supported by the ray-caster camera.""" + # check if there is any intersection in unsupported types + # reason: these use np structured data types which we can't yet convert to torch tensor + common_elements = set(cfg.data_types) & Camera.UNSUPPORTED_TYPES + if common_elements: + # provide alternative fast counterparts + fast_common_elements = [] + for item in common_elements: + if "instance_segmentation" in item or "instance_id_segmentation" in item: + fast_common_elements.append(item + "_fast") + # raise error + raise ValueError( + f"TiledCamera class does not support the following sensor types: {common_elements}." + "\n\tThis is because these sensor types output numpy structured data types which" + "can't be converted to torch tensors easily." + "\n\tHint: If you need to work with these sensor types, we recommend using their fast counterparts." + f"\n\t\tFast counterparts: {fast_common_elements}" + ) + + def _create_buffers(self): + """Create buffers for storing data.""" + # create the data object + # -- pose of the cameras + self._data.pos_w = torch.zeros((self._view.count, 3), device=self._device) + self._data.quat_w_world = torch.zeros((self._view.count, 4), device=self._device) + self._update_poses(self._ALL_INDICES) + # -- intrinsic matrix + self._data.intrinsic_matrices = torch.zeros((self._view.count, 3, 3), device=self._device) + self._update_intrinsic_matrices(self._ALL_INDICES) + self._data.image_shape = self.image_shape + # -- output data + data_dict = dict() + if "rgba" in self.cfg.data_types or "rgb" in self.cfg.data_types: + data_dict["rgba"] = torch.zeros( + (self._view.count, self.cfg.height, self.cfg.width, 4), device=self.device, dtype=torch.uint8 + ).contiguous() + if "rgb" in self.cfg.data_types: + # RGB is the first 3 channels of RGBA + data_dict["rgb"] = data_dict["rgba"][..., :3] + if "distance_to_image_plane" in self.cfg.data_types: + data_dict["distance_to_image_plane"] = torch.zeros( + (self._view.count, self.cfg.height, self.cfg.width, 1), device=self.device, dtype=torch.float32 + ).contiguous() + if "depth" in self.cfg.data_types: + data_dict["depth"] = torch.zeros( + (self._view.count, self.cfg.height, self.cfg.width, 1), device=self.device, dtype=torch.float32 + ).contiguous() + if "distance_to_camera" in self.cfg.data_types: + data_dict["distance_to_camera"] = torch.zeros( + (self._view.count, self.cfg.height, self.cfg.width, 1), device=self.device, dtype=torch.float32 + ).contiguous() + if "normals" in self.cfg.data_types: + data_dict["normals"] = torch.zeros( + (self._view.count, self.cfg.height, self.cfg.width, 3), device=self.device, dtype=torch.float32 + ).contiguous() + if "motion_vectors" in self.cfg.data_types: + data_dict["motion_vectors"] = torch.zeros( + (self._view.count, self.cfg.height, self.cfg.width, 2), device=self.device, dtype=torch.float32 + ).contiguous() + if "semantic_segmentation" in self.cfg.data_types: + if self.cfg.colorize_semantic_segmentation: + data_dict["semantic_segmentation"] = torch.zeros( + (self._view.count, self.cfg.height, self.cfg.width, 4), device=self.device, dtype=torch.uint8 + ).contiguous() + else: + data_dict["semantic_segmentation"] = torch.zeros( + (self._view.count, self.cfg.height, self.cfg.width, 1), device=self.device, dtype=torch.int32 + ).contiguous() + if "instance_segmentation_fast" in self.cfg.data_types: + if self.cfg.colorize_instance_segmentation: + data_dict["instance_segmentation_fast"] = torch.zeros( + (self._view.count, self.cfg.height, self.cfg.width, 4), device=self.device, dtype=torch.uint8 + ).contiguous() + else: + data_dict["instance_segmentation_fast"] = torch.zeros( + (self._view.count, self.cfg.height, self.cfg.width, 1), device=self.device, dtype=torch.int32 + ).contiguous() + if "instance_id_segmentation_fast" in self.cfg.data_types: + if self.cfg.colorize_instance_id_segmentation: + data_dict["instance_id_segmentation_fast"] = torch.zeros( + (self._view.count, self.cfg.height, self.cfg.width, 4), device=self.device, dtype=torch.uint8 + ).contiguous() + else: + data_dict["instance_id_segmentation_fast"] = torch.zeros( + (self._view.count, self.cfg.height, self.cfg.width, 1), device=self.device, dtype=torch.int32 + ).contiguous() + + self._data.output = data_dict + self._data.info = dict() + + def _tiled_image_shape(self) -> tuple[int, int]: + """Returns a tuple containing the dimension of the tiled image.""" + cols, rows = self._tiling_grid_shape() + return (self.cfg.width * cols, self.cfg.height * rows) + + def _tiling_grid_shape(self) -> tuple[int, int]: + """Returns a tuple containing the tiling grid dimension.""" + cols = math.ceil(math.sqrt(self._view.count)) + rows = math.ceil(self._view.count / cols) + return (cols, rows) + + def _create_annotator_data(self): + # we do not need to create annotator data for the tiled camera sensor + raise RuntimeError("This function should not be called for the tiled camera sensor.") + + def _process_annotator_output(self, name: str, output: Any) -> tuple[torch.tensor, dict | None]: + # we do not need to process annotator output for the tiled camera sensor + raise RuntimeError("This function should not be called for the tiled camera sensor.") + + """ + Internal simulation callbacks. + """ + + def _invalidate_initialize_callback(self, event): + """Invalidates the scene elements.""" + # call parent + super()._invalidate_initialize_callback(event) + # set all existing views to None to invalidate them + self._view = None diff --git a/source/isaaclab/isaaclab/newton/sensors/camera/tiled_camera_cfg.py b/source/isaaclab/isaaclab/newton/sensors/camera/tiled_camera_cfg.py new file mode 100644 index 00000000000..a14e74d6f8f --- /dev/null +++ b/source/isaaclab/isaaclab/newton/sensors/camera/tiled_camera_cfg.py @@ -0,0 +1,16 @@ +# Copyright (c) 2022-2025, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +from isaaclab.utils import configclass + +from .camera_cfg import CameraCfg +from .tiled_camera import TiledCamera + + +@configclass +class TiledCameraCfg(CameraCfg): + """Configuration for a tiled rendering-based camera sensor.""" + + class_type: type = TiledCamera diff --git a/source/isaaclab/isaaclab/newton/sensors/camera/utils.py b/source/isaaclab/isaaclab/newton/sensors/camera/utils.py new file mode 100644 index 00000000000..70787140bfa --- /dev/null +++ b/source/isaaclab/isaaclab/newton/sensors/camera/utils.py @@ -0,0 +1,271 @@ +# Copyright (c) 2022-2025, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +"""Helper functions to project between pointcloud and depth images.""" + +# needed to import for allowing type-hinting: torch.device | str | None +from __future__ import annotations + +import numpy as np +import torch +from collections.abc import Sequence + +import warp as wp + +import isaaclab.utils.math as math_utils +from isaaclab.utils.array import TensorData, convert_to_torch + +""" +Depth <-> Pointcloud conversions. +""" + + +def transform_points( + points: TensorData, + position: Sequence[float] | None = None, + orientation: Sequence[float] | None = None, + device: torch.device | str | None = None, +) -> np.ndarray | torch.Tensor: + r"""Transform input points in a given frame to a target frame. + + This function transform points from a source frame to a target frame. The transformation is defined by the + position ``t`` and orientation ``R`` of the target frame in the source frame. + + .. math:: + p_{target} = R_{target} \times p_{source} + t_{target} + + If either the inputs `position` and `orientation` are None, the corresponding transformation is not applied. + + Args: + points: a tensor of shape (p, 3) or (n, p, 3) comprising of 3d points in source frame. + position: The position of source frame in target frame. Defaults to None. + orientation: The orientation (w, x, y, z) of source frame in target frame. + Defaults to None. + device: The device for torch where the computation + should be executed. Defaults to None, i.e. takes the device that matches the depth image. + + Returns: + A tensor of shape (N, 3) comprising of 3D points in target frame. + If the input is a numpy array, the output is a numpy array. Otherwise, it is a torch tensor. + """ + # check if numpy + is_numpy = isinstance(points, np.ndarray) + # decide device + if device is None and is_numpy: + device = torch.device("cpu") + # convert to torch + points = convert_to_torch(points, dtype=torch.float32, device=device) + # update the device with the device of the depth image + # note: this is needed since warp does not provide the device directly + device = points.device + # apply rotation + if orientation is not None: + orientation = convert_to_torch(orientation, dtype=torch.float32, device=device) + # apply translation + if position is not None: + position = convert_to_torch(position, dtype=torch.float32, device=device) + # apply transformation + points = math_utils.transform_points(points, position, orientation) + + # return everything according to input type + if is_numpy: + return points.detach().cpu().numpy() + else: + return points + + +def create_pointcloud_from_depth( + intrinsic_matrix: np.ndarray | torch.Tensor | wp.array, + depth: np.ndarray | torch.Tensor | wp.array, + keep_invalid: bool = False, + position: Sequence[float] | None = None, + orientation: Sequence[float] | None = None, + device: torch.device | str | None = None, +) -> np.ndarray | torch.Tensor: + r"""Creates pointcloud from input depth image and camera intrinsic matrix. + + This function creates a pointcloud from a depth image and camera intrinsic matrix. The pointcloud is + computed using the following equation: + + .. math:: + p_{camera} = K^{-1} \times [u, v, 1]^T \times d + + where :math:`K` is the camera intrinsic matrix, :math:`u` and :math:`v` are the pixel coordinates and + :math:`d` is the depth value at the pixel. + + Additionally, the pointcloud can be transformed from the camera frame to a target frame by providing + the position ``t`` and orientation ``R`` of the camera in the target frame: + + .. math:: + p_{target} = R_{target} \times p_{camera} + t_{target} + + Args: + intrinsic_matrix: A (3, 3) array providing camera's calibration matrix. + depth: An array of shape (H, W) with values encoding the depth measurement. + keep_invalid: Whether to keep invalid points in the cloud or not. Invalid points + correspond to pixels with depth values 0.0 or NaN. Defaults to False. + position: The position of the camera in a target frame. Defaults to None. + orientation: The orientation (w, x, y, z) of the camera in a target frame. Defaults to None. + device: The device for torch where the computation should be executed. + Defaults to None, i.e. takes the device that matches the depth image. + + Returns: + An array/tensor of shape (N, 3) comprising of 3D coordinates of points. + The returned datatype is torch if input depth is of type torch.tensor or wp.array. Otherwise, a np.ndarray + is returned. + """ + # We use PyTorch here for matrix multiplication since it is compiled with Intel MKL while numpy + # by default uses OpenBLAS. With PyTorch (CPU), we could process a depth image of size (480, 640) + # in 0.0051 secs, while with numpy it took 0.0292 secs. + + # convert to numpy matrix + is_numpy = isinstance(depth, np.ndarray) + # decide device + if device is None and is_numpy: + device = torch.device("cpu") + # convert depth to torch tensor + depth = convert_to_torch(depth, dtype=torch.float32, device=device) + # update the device with the device of the depth image + # note: this is needed since warp does not provide the device directly + device = depth.device + # convert inputs to torch tensors + intrinsic_matrix = convert_to_torch(intrinsic_matrix, dtype=torch.float32, device=device) + if position is not None: + position = convert_to_torch(position, dtype=torch.float32, device=device) + if orientation is not None: + orientation = convert_to_torch(orientation, dtype=torch.float32, device=device) + # compute pointcloud + depth_cloud = math_utils.unproject_depth(depth, intrinsic_matrix) + # convert 3D points to world frame + depth_cloud = math_utils.transform_points(depth_cloud, position, orientation) + + # keep only valid entries if flag is set + if not keep_invalid: + pts_idx_to_keep = torch.all(torch.logical_and(~torch.isnan(depth_cloud), ~torch.isinf(depth_cloud)), dim=1) + depth_cloud = depth_cloud[pts_idx_to_keep, ...] + + # return everything according to input type + if is_numpy: + return depth_cloud.detach().cpu().numpy() + else: + return depth_cloud + + +def create_pointcloud_from_rgbd( + intrinsic_matrix: torch.Tensor | np.ndarray | wp.array, + depth: torch.Tensor | np.ndarray | wp.array, + rgb: torch.Tensor | wp.array | np.ndarray | tuple[float, float, float] = None, + normalize_rgb: bool = False, + position: Sequence[float] | None = None, + orientation: Sequence[float] | None = None, + device: torch.device | str | None = None, + num_channels: int = 3, +) -> tuple[torch.Tensor, torch.Tensor] | tuple[np.ndarray, np.ndarray]: + """Creates pointcloud from input depth image and camera transformation matrix. + + This function provides the same functionality as :meth:`create_pointcloud_from_depth` but also allows + to provide the RGB values for each point. + + The ``rgb`` attribute is used to resolve the corresponding point's color: + + - If a ``np.array``/``wp.array``/``torch.tensor`` of shape (H, W, 3), then the corresponding channels encode RGB values. + - If a tuple, then the point cloud has a single color specified by the values (r, g, b). + - If None, then default color is white, i.e. (0, 0, 0). + + If the input ``normalize_rgb`` is set to :obj:`True`, then the RGB values are normalized to be in the range [0, 1]. + + Args: + intrinsic_matrix: A (3, 3) array/tensor providing camera's calibration matrix. + depth: An array/tensor of shape (H, W) with values encoding the depth measurement. + rgb: Color for generated point cloud. Defaults to None. + normalize_rgb: Whether to normalize input rgb. Defaults to False. + position: The position of the camera in a target frame. Defaults to None. + orientation: The orientation `(w, x, y, z)` of the camera in a target frame. Defaults to None. + device: The device for torch where the computation should be executed. Defaults to None, in which case + it takes the device that matches the depth image. + num_channels: Number of channels in RGB pointcloud. Defaults to 3. + + Returns: + A tuple of (N, 3) arrays or tensors containing the 3D coordinates of points and their RGB color respectively. + The returned datatype is torch if input depth is of type torch.tensor or wp.array. Otherwise, a np.ndarray + is returned. + + Raises: + ValueError: When rgb image is a numpy array but not of shape (H, W, 3) or (H, W, 4). + """ + # check valid inputs + if rgb is not None and not isinstance(rgb, tuple): + if len(rgb.shape) == 3: + if rgb.shape[2] not in [3, 4]: + raise ValueError(f"Input rgb image of invalid shape: {rgb.shape} != (H, W, 3) or (H, W, 4).") + else: + raise ValueError(f"Input rgb image not three-dimensional. Received shape: {rgb.shape}.") + if num_channels not in [3, 4]: + raise ValueError(f"Invalid number of channels: {num_channels} != 3 or 4.") + + # check if input depth is numpy array + is_numpy = isinstance(depth, np.ndarray) + # decide device + if device is None and is_numpy: + device = torch.device("cpu") + # convert depth to torch tensor + if is_numpy: + depth = torch.from_numpy(depth).to(device=device) + # retrieve XYZ pointcloud + points_xyz = create_pointcloud_from_depth(intrinsic_matrix, depth, True, position, orientation, device=device) + + # get image height and width + im_height, im_width = depth.shape[:2] + # total number of points + num_points = im_height * im_width + # extract color value + if rgb is not None: + if isinstance(rgb, (np.ndarray, torch.Tensor, wp.array)): + # copy numpy array to preserve + rgb = convert_to_torch(rgb, device=device, dtype=torch.float32) + rgb = rgb[:, :, :3] + # convert the matrix to (W, H, 3) from (H, W, 3) since depth processing + # is done in the order (u, v) where u: (0, W-1) and v: (0 - H-1) + points_rgb = rgb.permute(1, 0, 2).reshape(-1, 3) + elif isinstance(rgb, (tuple, list)): + # same color for all points + points_rgb = torch.Tensor((rgb,) * num_points, device=device, dtype=torch.uint8) + else: + # default color is white + points_rgb = torch.Tensor(((0, 0, 0),) * num_points, device=device, dtype=torch.uint8) + else: + points_rgb = torch.Tensor(((0, 0, 0),) * num_points, device=device, dtype=torch.uint8) + # normalize color values + if normalize_rgb: + points_rgb = points_rgb.float() / 255 + + # remove invalid points + pts_idx_to_keep = torch.all(torch.logical_and(~torch.isnan(points_xyz), ~torch.isinf(points_xyz)), dim=1) + points_rgb = points_rgb[pts_idx_to_keep, ...] + points_xyz = points_xyz[pts_idx_to_keep, ...] + + # add additional channels if required + if num_channels == 4: + points_rgb = torch.nn.functional.pad(points_rgb, (0, 1), mode="constant", value=1.0) + + # return everything according to input type + if is_numpy: + return points_xyz.cpu().numpy(), points_rgb.cpu().numpy() + else: + return points_xyz, points_rgb + + +def save_images_to_file(images: torch.Tensor, file_path: str): + """Save images to file. + + Args: + images: A tensor of shape (N, H, W, C) containing the images. + file_path: The path to save the images to. + """ + from torchvision.utils import make_grid, save_image + + save_image( + make_grid(torch.swapaxes(images.unsqueeze(1), 1, -1).squeeze(-1), nrow=round(images.shape[0] ** 0.5)), file_path + ) diff --git a/source/isaaclab/isaaclab/newton/sensors/contact_sensor/__init__.py b/source/isaaclab/isaaclab/newton/sensors/contact_sensor/__init__.py new file mode 100644 index 00000000000..07e91f88344 --- /dev/null +++ b/source/isaaclab/isaaclab/newton/sensors/contact_sensor/__init__.py @@ -0,0 +1,10 @@ +# Copyright (c) 2022-2025, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +"""Sub-module for rigid contact sensor based on :class:`isaacsim.core.prims.RigidContactView`.""" + +from .contact_sensor import ContactSensor +from .contact_sensor_cfg import ContactSensorCfg +from .contact_sensor_data import ContactSensorData diff --git a/source/isaaclab/isaaclab/newton/sensors/contact_sensor/contact_sensor.py b/source/isaaclab/isaaclab/newton/sensors/contact_sensor/contact_sensor.py new file mode 100644 index 00000000000..1343bbdaa44 --- /dev/null +++ b/source/isaaclab/isaaclab/newton/sensors/contact_sensor/contact_sensor.py @@ -0,0 +1,422 @@ +# Copyright (c) 2022-2025, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +# Ignore optional memory usage warning globally +# pyright: reportOptionalSubscript=false + +from __future__ import annotations + +import torch +from collections.abc import Sequence +from typing import TYPE_CHECKING + +import warp as wp +from newton.sensors import ContactSensor as NewtonContactSensor +from newton.sensors import MatchKind + +import isaaclab.utils.string as string_utils +from isaaclab.markers import VisualizationMarkers +from isaaclab.sim._impl.newton_manager import NewtonManager + +from ..sensor_base import SensorBase +from .contact_sensor_data import ContactSensorData + +if TYPE_CHECKING: + from .contact_sensor_cfg import ContactSensorCfg + + +class ContactSensor(SensorBase): + """A contact reporting sensor. + + The contact sensor reports the normal contact forces on a rigid body in the world frame. + It relies on the `PhysX ContactReporter`_ API to be activated on the rigid bodies. + + To enable the contact reporter on a rigid body, please make sure to enable the + :attr:`isaaclab.sim.spawner.RigidObjectSpawnerCfg.activate_contact_sensors` on your + asset spawner configuration. This will enable the contact reporter on all the rigid bodies + in the asset. + + The sensor can be configured to report the contact forces on a set of bodies with a given + filter pattern using the :attr:`ContactSensorCfg.filter_prim_paths_expr`. This is useful + when you want to report the contact forces between the sensor bodies and a specific set of + bodies in the scene. The data can be accessed using the :attr:`ContactSensorData.force_matrix_w`. + Please check the documentation on `RigidContact`_ for more details. + + The reporting of the filtered contact forces is only possible as one-to-many. This means that only one + sensor body in an environment can be filtered against multiple bodies in that environment. If you need to + filter multiple sensor bodies against multiple bodies, you need to create separate sensors for each sensor + body. + + As an example, suppose you want to report the contact forces for all the feet of a robot against an object + exclusively. In that case, setting the :attr:`ContactSensorCfg.prim_path` and + :attr:`ContactSensorCfg.filter_prim_paths_expr` with ``{ENV_REGEX_NS}/Robot/.*_FOOT`` and ``{ENV_REGEX_NS}/Object`` + respectively will not work. Instead, you need to create a separate sensor for each foot and filter + it against the object. + + .. _PhysX ContactReporter: https://docs.omniverse.nvidia.com/kit/docs/omni_usd_schema_physics/104.2/class_physx_schema_physx_contact_report_a_p_i.html + .. _RigidContact: https://docs.omniverse.nvidia.com/py/isaacsim/source/isaacsim.core/docs/index.html#isaacsim.core.prims.RigidContact + """ + + cfg: ContactSensorCfg + """The configuration parameters.""" + + def __init__(self, cfg: ContactSensorCfg): + """Initializes the contact sensor object. + + Args: + cfg: The configuration parameters. + """ + # initialize base class + super().__init__(cfg) + + # Create empty variables for storing output data + self._data: ContactSensorData = ContactSensorData() + + def __str__(self) -> str: + """Returns: A string containing information about the instance.""" + return ( + f"Contact sensor @ '{self.cfg.prim_path}': \n" + f"\tupdate period (s) : {self.cfg.update_period}\n" + f"\tnumber of bodies : {self.num_bodies}\n" + f"\tbody names : {self.body_names}\n" + ) + + """ + Properties + """ + + @property + def num_instances(self) -> int | None: + return self._num_bodies + + @property + def data(self) -> ContactSensorData: + # update sensors if needed + self._update_outdated_buffers() + # return the data + return self._data + + @property + def num_bodies(self) -> int: + """Number of bodies with contact sensors attached.""" + return self._num_bodies + + @property + def body_names(self) -> list[str] | None: + """Ordered names of shapes or bodies with contact sensors attached.""" + return self._body_names + + @property + def contact_partner_names(self) -> list[str] | None: + """Ordered names of shapes or bodies that are selected as contact partners.""" + return self._contact_partner_names + + @property + def contact_newton_view(self) -> NewtonContactSensor: + """View for the contact forces captured (Newton).""" + return NewtonManager._newton_contact_sensor + + """ + Operations + """ + + def reset(self, env_ids: Sequence[int] | None = None): + # reset the timers and counters + super().reset(env_ids) + # resolve None + if env_ids is None: + env_ids = slice(None) + # reset accumulative data buffers + self._data.net_forces_w[env_ids] = 0.0 + self._data.net_forces_w_history[env_ids] = 0.0 + if self.cfg.history_length > 0: + self._data.net_forces_w_history[env_ids] = 0.0 + # reset force matrix + if self.cfg.filter_prim_paths_expr is not None or self.cfg.filter_shape_paths_expr is not None: + self._data.force_matrix_w[env_ids] = 0.0 + # reset the current air time + if self.cfg.track_air_time: + self._data.current_air_time[env_ids] = 0.0 + self._data.last_air_time[env_ids] = 0.0 + self._data.current_contact_time[env_ids] = 0.0 + self._data.last_contact_time[env_ids] = 0.0 + + def find_bodies(self, name_keys: str | Sequence[str], preserve_order: bool = False) -> tuple[list[int], list[str]]: + """Find bodies in the articulation based on the name keys. + + Args: + name_keys: A regular expression or a list of regular expressions to match the body names. + preserve_order: Whether to preserve the order of the name keys in the output. Defaults to False. + + Returns: + A tuple of lists containing the body indices and names. + """ + return string_utils.resolve_matching_names(name_keys, self.body_names, preserve_order) + + def compute_first_contact(self, dt: float, abs_tol: float = 1.0e-8) -> torch.Tensor: + """Checks if bodies that have established contact within the last :attr:`dt` seconds. + + This function checks if the bodies have established contact within the last :attr:`dt` seconds + by comparing the current contact time with the given time period. If the contact time is less + than the given time period, then the bodies are considered to be in contact. + + Note: + The function assumes that :attr:`dt` is a factor of the sensor update time-step. In other + words :math:`dt / dt_sensor = n`, where :math:`n` is a natural number. This is always true + if the sensor is updated by the physics or the environment stepping time-step and the sensor + is read by the environment stepping time-step. + + Args: + dt: The time period since the contact was established. + abs_tol: The absolute tolerance for the comparison. + + Returns: + A boolean tensor indicating the bodies that have established contact within the last + :attr:`dt` seconds. Shape is (N, B), where N is the number of sensors and B is the + number of bodies in each sensor. + + Raises: + RuntimeError: If the sensor is not configured to track contact time. + """ + # check if the sensor is configured to track contact time + if not self.cfg.track_air_time: + raise RuntimeError( + "The contact sensor is not configured to track contact time." + "Please enable the 'track_air_time' in the sensor configuration." + ) + # check if the bodies are in contact + currently_in_contact = self.data.current_contact_time > 0.0 + less_than_dt_in_contact = self.data.current_contact_time < (dt + abs_tol) + return currently_in_contact * less_than_dt_in_contact + + def compute_first_air(self, dt: float, abs_tol: float = 1.0e-8) -> torch.Tensor: + """Checks if bodies that have broken contact within the last :attr:`dt` seconds. + + This function checks if the bodies have broken contact within the last :attr:`dt` seconds + by comparing the current air time with the given time period. If the air time is less + than the given time period, then the bodies are considered to not be in contact. + + Note: + It assumes that :attr:`dt` is a factor of the sensor update time-step. In other words, + :math:`dt / dt_sensor = n`, where :math:`n` is a natural number. This is always true if + the sensor is updated by the physics or the environment stepping time-step and the sensor + is read by the environment stepping time-step. + + Args: + dt: The time period since the contract is broken. + abs_tol: The absolute tolerance for the comparison. + + Returns: + A boolean tensor indicating the bodies that have broken contact within the last :attr:`dt` seconds. + Shape is (N, B), where N is the number of sensors and B is the number of bodies in each sensor. + + Raises: + RuntimeError: If the sensor is not configured to track contact time. + """ + # check if the sensor is configured to track contact time + if not self.cfg.track_air_time: + raise RuntimeError( + "The contact sensor is not configured to track contact time." + "Please enable the 'track_air_time' in the sensor configuration." + ) + # check if the sensor is configured to track contact time + currently_detached = self.data.current_air_time > 0.0 + less_than_dt_detached = self.data.current_air_time < (dt + abs_tol) + return currently_detached * less_than_dt_detached + + """ + Implementation. + """ + + def _initialize_impl(self): + super()._initialize_impl() + """Initializes the sensor-related handles and internal buffers.""" + # construct regex expression for the body names + + if self.cfg.filter_prim_paths_expr is not None or self.cfg.filter_shape_paths_expr is not None: + self._generate_force_matrix = True + else: + self._generate_force_matrix = False + + body_names_regex = self.cfg.prim_path + if self.cfg.shape_path is not None: + shape_names_regex = "(" + "|".join(self.cfg.shape_path) + ")" + else: + shape_names_regex = None + if self.cfg.filter_prim_paths_expr is not None: + contact_partners_body_regex = "(" + "|".join(self.cfg.filter_prim_paths_expr) + ")" + else: + contact_partners_body_regex = None + if self.cfg.filter_shape_paths_expr is not None: + contact_partners_shape_regex = "(" + "|".join(self.cfg.filter_shape_paths_expr) + ")" + else: + contact_partners_shape_regex = None + + NewtonManager.add_contact_sensor( + body_names_expr=body_names_regex, + shape_names_expr=shape_names_regex, + contact_partners_body_expr=contact_partners_body_regex, + contact_partners_shape_expr=contact_partners_shape_regex, + ) + self._create_buffers() + + def _create_buffers(self): + # resolve the true count of bodies + self._num_bodies = self.contact_newton_view.shape[0] // self._num_envs + + # Check that number of bodies is an integer + if self.contact_newton_view.shape[0] % self._num_envs != 0: + raise RuntimeError( + "Number of bodies is not an integer multiple of the number of environments. Received:" + f" {self._num_bodies} bodies and {self._num_envs} environments." + ) + print(f"[INFO] Contact sensor initialized with {self._num_bodies} bodies.") + + # Assume homogeneous envs, i.e. all envs have the same number of bodies / shapes + # Only get the names for the first env. Expected structure: /World/envs/env_.*/... + def get_name(idx, match_kind): + if match_kind == MatchKind.BODY: + return NewtonManager._model.body_key[idx].split("/")[-1] + if match_kind == MatchKind.SHAPE: + return NewtonManager._model.shape_key[idx].split("/")[-1] + return "MATCH_ANY" + + self._body_names = [get_name(idx, kind) for idx, kind in self.contact_newton_view.sensing_objs] + # Assumes the environments are processed in order. + self._body_names = self._body_names[: self._num_bodies] + self._contact_partner_names = [get_name(idx, kind) for idx, kind in self.contact_newton_view.counterparts] + + # prepare data buffers + self._data.net_forces_w = torch.zeros(self._num_envs, self._num_bodies, 3, device=self._device) + # optional buffers + # -- history of net forces + if self.cfg.history_length > 0: + self._data.net_forces_w_history = torch.zeros( + self._num_envs, self.cfg.history_length, self._num_bodies, 3, device=self._device + ) + else: + self._data.net_forces_w_history = self._data.net_forces_w.unsqueeze(1) + # -- pose of sensor origins + if self.cfg.track_pose: + self._data.pos_w = torch.zeros(self._num_envs, self._num_bodies, 3, device=self._device) + self._data.quat_w = torch.zeros(self._num_envs, self._num_bodies, 4, device=self._device) + # -- air/contact time between contacts + if self.cfg.track_air_time: + self._data.last_air_time = torch.zeros(self._num_envs, self._num_bodies, device=self._device) + self._data.current_air_time = torch.zeros(self._num_envs, self._num_bodies, device=self._device) + self._data.last_contact_time = torch.zeros(self._num_envs, self._num_bodies, device=self._device) + self._data.current_contact_time = torch.zeros(self._num_envs, self._num_bodies, device=self._device) + # force matrix: (num_envs, num_bodies, num_filter_shapes, 3) + if self._generate_force_matrix: + num_filters = self.contact_newton_view.shape[1] + self._data.force_matrix_w = torch.zeros( + self._num_envs, self._num_bodies, num_filters, 3, device=self._device + ) + + def _update_buffers_impl(self, env_ids: Sequence[int]): + """Fills the buffers of the sensor data.""" + + # default to all sensors + if len(env_ids) == self._num_envs: + env_ids = slice(None) + + # net_force is a matrix of shape (num_bodies * num_envs, num_filters, 3) + net_forces_w = wp.to_torch(self.contact_newton_view.net_force).clone() + self._data.net_forces_w[env_ids, :, :] = net_forces_w[:, 0, :].reshape(self._num_envs, self._num_bodies, 3)[ + env_ids + ] + + if self.cfg.history_length > 0: + self._data.net_forces_w_history[env_ids, 1:] = self._data.net_forces_w_history[env_ids, :-1].clone() + self._data.net_forces_w_history[env_ids, 0] = self._data.net_forces_w[env_ids] + + # obtain the contact force matrix + if self._generate_force_matrix: + # shape of the filtering matrix: (num_envs, num_bodies, num_filter_shapes, 3) + num_filters = self.contact_newton_view.shape[1] - 1 # -1 for the total force + # acquire and shape the force matrix + self._data.force_matrix_w[env_ids] = net_forces_w[:, 1:, :].reshape( + self._num_envs, self._num_bodies, num_filters, 3 + )[env_ids] + + # FIXME: Re-enable this when we have a non-physx rigid body view? + # obtain the pose of the sensor origin + # if self.cfg.track_pose: + # pose = self.body_physx_view.get_transforms().view(-1, self._num_bodies, 7)[env_ids] + # pose[..., 3:] = convert_quat(pose[..., 3:], to="wxyz") + # self._data.pos_w[env_ids], self._data.quat_w[env_ids] = pose.split([3, 4], dim=-1) + + # obtain the air time + if self.cfg.track_air_time: + # -- time elapsed since last update + # since this function is called every frame, we can use the difference to get the elapsed time + elapsed_time = self._timestamp[env_ids] - self._timestamp_last_update[env_ids] + # -- check contact state of bodies + is_contact = torch.norm(self._data.net_forces_w[env_ids, :, :], dim=-1) > self.cfg.force_threshold + is_first_contact = (self._data.current_air_time[env_ids] > 0) * is_contact + is_first_detached = (self._data.current_contact_time[env_ids] > 0) * ~is_contact + # -- update the last contact time if body has just become in contact + self._data.last_air_time[env_ids] = torch.where( + is_first_contact, + self._data.current_air_time[env_ids] + elapsed_time.unsqueeze(-1), + self._data.last_air_time[env_ids], + ) + # -- increment time for bodies that are not in contact + self._data.current_air_time[env_ids] = torch.where( + ~is_contact, self._data.current_air_time[env_ids] + elapsed_time.unsqueeze(-1), 0.0 + ) + # -- update the last contact time if body has just detached + self._data.last_contact_time[env_ids] = torch.where( + is_first_detached, + self._data.current_contact_time[env_ids] + elapsed_time.unsqueeze(-1), + self._data.last_contact_time[env_ids], + ) + # -- increment time for bodies that are in contact + self._data.current_contact_time[env_ids] = torch.where( + is_contact, self._data.current_contact_time[env_ids] + elapsed_time.unsqueeze(-1), 0.0 + ) + + def _set_debug_vis_impl(self, debug_vis: bool): + # set visibility of markers + # note: parent only deals with callbacks. not their visibility + if debug_vis: + # create markers if necessary for the first tome + if not hasattr(self, "contact_visualizer"): + self.contact_visualizer = VisualizationMarkers(self.cfg.visualizer_cfg) + # set their visibility to true + self.contact_visualizer.set_visibility(True) + else: + if hasattr(self, "contact_visualizer"): + self.contact_visualizer.set_visibility(False) + + def _debug_vis_callback(self, event): + # safely return if view becomes invalid + return + # note: this invalidity happens because of isaac sim view callbacks + # if self.body_physx_view is None: + # return + # marker indices + # 0: contact, 1: no contact + # net_contact_force_w = torch.norm(self._data.net_forces_w, dim=-1) + # marker_indices = torch.where(net_contact_force_w > self.cfg.force_threshold, 0, 1) + # check if prim is visualized + # if self.cfg.track_pose: + # frame_origins: torch.Tensor = self._data.pos_w + # else: + # pose = self.body_physx_view.get_transforms() + # frame_origins = pose.view(-1, self._num_bodies, 7)[:, :, :3] + # visualize + # self.contact_visualizer.visualize(frame_origins.view(-1, 3), marker_indices=marker_indices.view(-1)) + + """ + Internal simulation callbacks. + """ + + def _invalidate_initialize_callback(self, event): + """Invalidates the scene elements.""" + # call parent + super()._invalidate_initialize_callback(event) + # set all existing views to None to invalidate them + # TODO: invalidate NewtonManager if necessary diff --git a/source/isaaclab/isaaclab/newton/sensors/contact_sensor/contact_sensor_cfg.py b/source/isaaclab/isaaclab/newton/sensors/contact_sensor/contact_sensor_cfg.py new file mode 100644 index 00000000000..6d2709468de --- /dev/null +++ b/source/isaaclab/isaaclab/newton/sensors/contact_sensor/contact_sensor_cfg.py @@ -0,0 +1,107 @@ +# Copyright (c) 2022-2025, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +from isaaclab.markers import VisualizationMarkersCfg +from isaaclab.markers.config import CONTACT_SENSOR_MARKER_CFG +from isaaclab.utils import configclass + +from ..sensor_base_cfg import SensorBaseCfg +from .contact_sensor import ContactSensor + + +@configclass +class ContactSensorCfg(SensorBaseCfg): + """Configuration for the contact sensor.""" + + class_type: type = ContactSensor + + track_pose: bool = False + """Whether to track the pose of the sensor's origin. Defaults to False.""" + + track_air_time: bool = False + """Whether to track the air/contact time of the bodies (time between contacts). Defaults to False.""" + + force_threshold: float = 1.0 + """The threshold on the norm of the contact force that determines whether two bodies are in collision or not. + + This value is only used for tracking the mode duration (the time in contact or in air), + if :attr:`track_air_time` is True. + """ + + shape_path: list[str] | None = None + """A list of expressions to filter contacts shapes with. Defaults to None. If both :attr:`body_names_expr` and + :attr:`shape_names_expr` are None, the contact with all bodies/shapes is reported. + + Only one of :attr:`body_names_expr` or :attr:`shape_names_expr` can be provided. + If both are provided, an error will be raised. + + We make an explicit difference between a body and a shape. A body is a rigid body, while a shape is a collision + shape. A body can have multiple shapes. The shape option allows a more fine-grained control over the contact + reporting. + + .. note:: + The expression in the list can contain the environment namespace regex ``{ENV_REGEX_NS}`` which + will be replaced with the environment namespace. + """ + + filter_prim_paths_expr: list[str] | None = None + """A list of expressions to filter contacts bodies with. Defaults to None. If both :attr:`contact_partners_body_expr` and + :attr:`contact_partners_shape_expr` are None, the contact with all bodies/shapes is reported. + + Only one of :attr:`contact_partners_body_expr` or :attr:`contact_partners_shape_expr` can be provided. + If both are provided, an error will be raised. + + The contact sensor allows reporting contacts between the primitive specified with either :attr:`body_names_expr` or + :attr:`shape_names_expr` and other primitives in the scene. For instance, in a scene containing a robot, a ground + plane and an object, you can obtain individual contact reports of the base of the robot with the ground plane and + the object. + + We make an explicit difference between a body and a shape. A body is a rigid body, while a shape is a collision + shape. A body can have multiple shapes. The shape option allows a more fine-grained control over the contact + reporting. + + .. note:: + The expression in the list can contain the environment namespace regex ``{ENV_REGEX_NS}`` which + will be replaced with the environment namespace. + + Example: ``{ENV_REGEX_NS}/Object`` will be replaced with ``/World/envs/env_.*/Object``. + + .. attention:: + The reporting of filtered contacts only works when the sensor primitive :attr:`prim_path` corresponds to a + single primitive in that environment. If the sensor primitive corresponds to multiple primitives, the + filtering will not work as expected. Please check :class:`~isaaclab.sensors.contact_sensor.ContactSensor` + for more details. + """ + + filter_shape_paths_expr: list[str] | None = None + """A list of expressions to filter contacts shapes with. Defaults to None. If both :attr:`contact_partners_body_expr` and + :attr:`contact_partners_shape_expr` are None, the contact with all bodies/shapes is reported. + + Only one of :attr:`contact_partners_body_expr` or :attr:`contact_partners_shape_expr` can be provided. + If both are provided, an error will be raised. + + The contact sensor allows reporting contacts between the primitive specified with either :attr:`body_names_expr` or + :attr:`shape_names_expr` and other primitives in the scene. For instance, in a scene containing a robot, a ground + plane and an object, you can obtain individual contact reports of the base of the robot with the ground plane and + the object. + + + We make an explicit difference between a body and a shape. A body is a rigid body, while a shape is a collision + shape. A body can have multiple shapes. The shape option allows a more fine-grained control over the contact + reporting. + + .. note:: + The expression in the list can contain the environment namespace regex ``{ENV_REGEX_NS}`` which + will be replaced with the environment namespace. + + Example: ``{ENV_REGEX_NS}/Object`` will be replaced with ``/World/envs/env_.*/Object``. + """ + + visualizer_cfg: VisualizationMarkersCfg = CONTACT_SENSOR_MARKER_CFG.replace(prim_path="/Visuals/ContactSensor") + """The configuration object for the visualization markers. Defaults to CONTACT_SENSOR_MARKER_CFG. + + .. note:: + This attribute is only used when debug visualization is enabled. + """ diff --git a/source/isaaclab/isaaclab/newton/sensors/contact_sensor/contact_sensor_data.py b/source/isaaclab/isaaclab/newton/sensors/contact_sensor/contact_sensor_data.py new file mode 100644 index 00000000000..cd01630af61 --- /dev/null +++ b/source/isaaclab/isaaclab/newton/sensors/contact_sensor/contact_sensor_data.py @@ -0,0 +1,102 @@ +# Copyright (c) 2022-2025, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +# needed to import for allowing type-hinting: torch.Tensor | None +from __future__ import annotations + +import torch +from dataclasses import dataclass + + +@dataclass +class ContactSensorData: + """Data container for the contact reporting sensor.""" + + pos_w: torch.Tensor | None = None + """Position of the sensor origin in world frame. + + Shape is (N, 3), where N is the number of sensors. + + Note: + If the :attr:`ContactSensorCfg.track_pose` is False, then this quantity is None. + """ + + quat_w: torch.Tensor | None = None + """Orientation of the sensor origin in quaternion (w, x, y, z) in world frame. + + Shape is (N, 4), where N is the number of sensors. + + Note: + If the :attr:`ContactSensorCfg.track_pose` is False, then this quantity is None. + """ + + net_forces_w: torch.Tensor | None = None + """The net normal contact forces in world frame. + + Shape is (N, B, 3), where N is the number of sensors and B is the number of bodies in each sensor. + + Note: + This quantity is the sum of the normal contact forces acting on the sensor bodies. It must not be confused + with the total contact forces acting on the sensor bodies (which also includes the tangential forces). + """ + + net_forces_w_history: torch.Tensor | None = None + """The net normal contact forces in world frame. + + Shape is (N, T, B, 3), where N is the number of sensors, T is the configured history length + and B is the number of bodies in each sensor. + + In the history dimension, the first index is the most recent and the last index is the oldest. + + Note: + This quantity is the sum of the normal contact forces acting on the sensor bodies. It must not be confused + with the total contact forces acting on the sensor bodies (which also includes the tangential forces). + """ + + force_matrix_w: torch.Tensor | None = None + """The normal contact forces filtered between the sensor bodies and filtered bodies in world frame. + + Shape is (N, B, M, 3), where N is the number of sensors, B is number of bodies in each sensor + and ``M`` is the number of filtered bodies. + + Note: + If the :attr:`ContactSensorCfg.filter_prim_paths_expr` is empty, then this quantity is None. + """ + + last_air_time: torch.Tensor | None = None + """Time spent (in s) in the air before the last contact. + + Shape is (N, B), where N is the number of sensors and B is the number of bodies in each sensor. + + Note: + If the :attr:`ContactSensorCfg.track_air_time` is False, then this quantity is None. + """ + + current_air_time: torch.Tensor | None = None + """Time spent (in s) in the air since the last detach. + + Shape is (N, B), where N is the number of sensors and B is the number of bodies in each sensor. + + Note: + If the :attr:`ContactSensorCfg.track_air_time` is False, then this quantity is None. + """ + + last_contact_time: torch.Tensor | None = None + """Time spent (in s) in contact before the last detach. + + Shape is (N, B), where N is the number of sensors and B is the number of bodies in each sensor. + + Note: + If the :attr:`ContactSensorCfg.track_air_time` is False, then this quantity is None. + """ + + current_contact_time: torch.Tensor | None = None + """Time spent (in s) in contact since the last contact. + + Shape is (N, B), where N is the number of sensors and B is the number of bodies in each sensor. + + Note: + If the :attr:`ContactSensorCfg.track_air_time` is False, then this quantity is None. + """ diff --git a/source/isaaclab/isaaclab/newton/sensors/sensor_base.py b/source/isaaclab/isaaclab/newton/sensors/sensor_base.py new file mode 100644 index 00000000000..6138cbc6f61 --- /dev/null +++ b/source/isaaclab/isaaclab/newton/sensors/sensor_base.py @@ -0,0 +1,346 @@ +# Copyright (c) 2022-2025, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +"""Base class for sensors. + +This class defines an interface for sensors similar to how the :class:`isaaclab.assets.AssetBase` class works. +Each sensor class should inherit from this class and implement the abstract methods. +""" + +from __future__ import annotations + +import builtins +import inspect +import re +import torch +import weakref +from abc import ABC, abstractmethod +from collections.abc import Sequence +from typing import TYPE_CHECKING, Any + +import omni.kit.app +import omni.timeline +from isaacsim.core.simulation_manager import IsaacEvents, SimulationManager +from isaacsim.core.utils.stage import get_current_stage + +import isaaclab.sim as sim_utils +from isaaclab.sim import SimulationContext +from isaaclab.sim._impl.newton_manager import NewtonManager + +if TYPE_CHECKING: + from .sensor_base_cfg import SensorBaseCfg + + +class SensorBase(ABC): + """The base class for implementing a sensor. + + The implementation is based on lazy evaluation. The sensor data is only updated when the user + tries accessing the data through the :attr:`data` property or sets ``force_compute=True`` in + the :meth:`update` method. This is done to avoid unnecessary computation when the sensor data + is not used. + + The sensor is updated at the specified update period. If the update period is zero, then the + sensor is updated at every simulation step. + """ + + def __init__(self, cfg: SensorBaseCfg): + """Initialize the sensor class. + + Args: + cfg: The configuration parameters for the sensor. + """ + # check that config is valid + if cfg.history_length < 0: + raise ValueError(f"History length must be greater than 0! Received: {cfg.history_length}") + # check that the config is valid + cfg.validate() + # store inputs + self.cfg = cfg.copy() + # flag for whether the sensor is initialized + self._is_initialized = False + # flag for whether the sensor is in visualization mode + self._is_visualizing = False + # get stage handle + self.stage = get_current_stage() + + # register simulator callbacks (with weakref safety to avoid crashes on deletion) + def safe_callback(callback_name, event, obj_ref): + """Safely invoke a callback on a weakly-referenced object, ignoring ReferenceError if deleted.""" + try: + obj = obj_ref + getattr(obj, callback_name)(event) + except ReferenceError: + # Object has been deleted; ignore. + pass + + # note: use weakref on callbacks to ensure that this object can be deleted when its destructor is called. + # add callbacks for stage play/stop + obj_ref = weakref.proxy(self) + timeline_event_stream = omni.timeline.get_timeline_interface().get_timeline_event_stream() + + # the order is set to 10 which is arbitrary but should be lower priority than the default order of 0 + NewtonManager.add_on_start_callback(lambda: safe_callback("_initialize_callback", None, obj_ref)) + # register timeline STOP event callback (lower priority with order=10) + self._invalidate_initialize_handle = timeline_event_stream.create_subscription_to_pop_by_type( + int(omni.timeline.TimelineEventType.STOP), + lambda event, obj_ref=obj_ref: safe_callback("_invalidate_initialize_callback", event, obj_ref), + order=10, + ) + # register prim deletion callback + self._prim_deletion_callback_id = SimulationManager.register_callback( + lambda event, obj_ref=obj_ref: safe_callback("_on_prim_deletion", event, obj_ref), + event=IsaacEvents.PRIM_DELETION, + ) + + # add handle for debug visualization (this is set to a valid handle inside set_debug_vis) + self._debug_vis_handle = None + # set initial state of debug visualization + self.set_debug_vis(self.cfg.debug_vis) + + def __del__(self): + """Unsubscribe from the callbacks.""" + # clear physics events handles + self._clear_callbacks() + + """ + Properties + """ + + @property + def is_initialized(self) -> bool: + """Whether the sensor is initialized. + + Returns True if the sensor is initialized, False otherwise. + """ + return self._is_initialized + + @property + def num_instances(self) -> int: + """Number of instances of the sensor. + + This is equal to the number of sensors per environment multiplied by the number of environments. + """ + return self._num_envs + + @property + def device(self) -> str: + """Memory device for computation.""" + return self._device + + @property + @abstractmethod + def data(self) -> Any: + """Data from the sensor. + + This property is only updated when the user tries to access the data. This is done to avoid + unnecessary computation when the sensor data is not used. + + For updating the sensor when this property is accessed, you can use the following + code snippet in your sensor implementation: + + .. code-block:: python + + # update sensors if needed + self._update_outdated_buffers() + # return the data (where `_data` is the data for the sensor) + return self._data + """ + raise NotImplementedError + + @property + def has_debug_vis_implementation(self) -> bool: + """Whether the sensor has a debug visualization implemented.""" + # check if function raises NotImplementedError + source_code = inspect.getsource(self._set_debug_vis_impl) + return "NotImplementedError" not in source_code + + """ + Operations + """ + + def set_debug_vis(self, debug_vis: bool) -> bool: + """Sets whether to visualize the sensor data. + + Args: + debug_vis: Whether to visualize the sensor data. + + Returns: + Whether the debug visualization was successfully set. False if the sensor + does not support debug visualization. + """ + # check if debug visualization is supported + if not self.has_debug_vis_implementation: + return False + # toggle debug visualization objects + self._set_debug_vis_impl(debug_vis) + # toggle debug visualization flag + self._is_visualizing = debug_vis + # toggle debug visualization handles + if debug_vis: + # create a subscriber for the post update event if it doesn't exist + if self._debug_vis_handle is None: + app_interface = omni.kit.app.get_app_interface() + self._debug_vis_handle = app_interface.get_post_update_event_stream().create_subscription_to_pop( + lambda event, obj=weakref.proxy(self): obj._debug_vis_callback(event) + ) + else: + # remove the subscriber if it exists + if self._debug_vis_handle is not None: + self._debug_vis_handle.unsubscribe() + self._debug_vis_handle = None + # return success + return True + + def reset(self, env_ids: Sequence[int] | None = None): + """Resets the sensor internals. + + Args: + env_ids: The sensor ids to reset. Defaults to None. + """ + # Resolve sensor ids + if env_ids is None: + env_ids = slice(None) + # Reset the timestamp for the sensors + self._timestamp[env_ids] = 0.0 + + self._timestamp_last_update[env_ids] = 0.0 + # Set all reset sensors to outdated so that they are updated when data is called the next time. + self._is_outdated[env_ids] = True + + def update(self, dt: float, force_recompute: bool = False): + # Update the timestamp for the sensors + self._timestamp += dt + self._is_outdated |= self._timestamp - self._timestamp_last_update + 1e-6 >= self.cfg.update_period + # Update the buffers + # TODO (from @mayank): Why is there a history length here when it doesn't mean anything in the sensor base?!? + # It is only for the contact sensor but there we should redefine the update function IMO. + if force_recompute or self._is_visualizing or (self.cfg.history_length > 0): + self._update_outdated_buffers() + + """ + Implementation specific. + """ + + @abstractmethod + def _initialize_impl(self): + """Initializes the sensor-related handles and internal buffers.""" + # Obtain Simulation Context + sim = sim_utils.SimulationContext.instance() + if sim is None: + raise RuntimeError("Simulation Context is not initialized!") + # Obtain device and backend + self._device = sim.device + self._backend = sim.backend + self._sim_physics_dt = sim.get_physics_dt() + self._num_envs = NewtonManager._num_envs + # Boolean tensor indicating whether the sensor data has to be refreshed + self._is_outdated = torch.ones(self._num_envs, dtype=torch.bool, device=self._device) + # Current timestamp (in seconds) + self._timestamp = torch.zeros(self._num_envs, device=self._device) + # Timestamp from last update + self._timestamp_last_update = torch.zeros_like(self._timestamp) + + @abstractmethod + def _update_buffers_impl(self, env_ids: Sequence[int]): + """Fills the sensor data for provided environment ids. + + This function does not perform any time-based checks and directly fills the data into the + data container. + + Args: + env_ids: The indices of the sensors that are ready to capture. + """ + raise NotImplementedError + + def _set_debug_vis_impl(self, debug_vis: bool): + """Set debug visualization into visualization objects. + + This function is responsible for creating the visualization objects if they don't exist + and input ``debug_vis`` is True. If the visualization objects exist, the function should + set their visibility into the stage. + """ + raise NotImplementedError(f"Debug visualization is not implemented for {self.__class__.__name__}.") + + def _debug_vis_callback(self, event): + """Callback for debug visualization. + + This function calls the visualization objects and sets the data to visualize into them. + """ + raise NotImplementedError(f"Debug visualization is not implemented for {self.__class__.__name__}.") + + """ + Internal simulation callbacks. + """ + + def _initialize_callback(self, event): + """Initializes the scene elements. + + Note: + PhysX handles are only enabled once the simulator starts playing. Hence, this function needs to be + called whenever the simulator "plays" from a "stop" state. + """ + if not self._is_initialized: + try: + self._initialize_impl() + except Exception as e: + if builtins.ISAACLAB_CALLBACK_EXCEPTION is None: + builtins.ISAACLAB_CALLBACK_EXCEPTION = e + self._is_initialized = True + + def _invalidate_initialize_callback(self, event): + """Invalidates the scene elements.""" + self._is_initialized = False + if self._debug_vis_handle is not None: + self._debug_vis_handle.unsubscribe() + self._debug_vis_handle = None + + def _on_prim_deletion(self, prim_path: str) -> None: + """Invalidates and deletes the callbacks when the prim is deleted. + + Args: + prim_path: The path to the prim that is being deleted. + + Note: + This function is called when the prim is deleted. + """ + # skip callback if required + if getattr(SimulationContext.instance(), "_skip_next_prim_deletion_callback_fn", False): + return + if prim_path == "/": + self._clear_callbacks() + return + result = re.match( + pattern="^" + "/".join(self.cfg.prim_path.split("/")[: prim_path.count("/") + 1]) + "$", string=prim_path + ) + if result: + self._clear_callbacks() + + def _clear_callbacks(self) -> None: + """Clears the callbacks.""" + if self._prim_deletion_callback_id: + SimulationManager.deregister_callback(self._prim_deletion_callback_id) + self._prim_deletion_callback_id = None + if self._invalidate_initialize_handle: + self._invalidate_initialize_handle.unsubscribe() + self._invalidate_initialize_handle = None + # clear debug visualization + if self._debug_vis_handle: + self._debug_vis_handle.unsubscribe() + self._debug_vis_handle = None + + """ + Helper functions. + """ + + def _update_outdated_buffers(self): + """Fills the sensor data for the outdated sensors.""" + outdated_env_ids = self._is_outdated.nonzero().squeeze(-1) + if len(outdated_env_ids) > 0: + # obtain new data + self._update_buffers_impl(outdated_env_ids) + # update the timestamp from last update + self._timestamp_last_update[outdated_env_ids] = self._timestamp[outdated_env_ids] + # set outdated flag to false for the updated sensors + self._is_outdated[outdated_env_ids] = False diff --git a/source/isaaclab/isaaclab/newton/sensors/sensor_base_cfg.py b/source/isaaclab/isaaclab/newton/sensors/sensor_base_cfg.py new file mode 100644 index 00000000000..9b2884a2128 --- /dev/null +++ b/source/isaaclab/isaaclab/newton/sensors/sensor_base_cfg.py @@ -0,0 +1,42 @@ +# Copyright (c) 2022-2025, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +from dataclasses import MISSING + +from isaaclab.utils import configclass + +from .sensor_base import SensorBase + + +@configclass +class SensorBaseCfg: + """Configuration parameters for a sensor.""" + + class_type: type[SensorBase] = MISSING + """The associated sensor class. + + The class should inherit from :class:`isaaclab.sensors.sensor_base.SensorBase`. + """ + + prim_path: str = MISSING + """Prim path (or expression) to the sensor. + + .. note:: + The expression can contain the environment namespace regex ``{ENV_REGEX_NS}`` which + will be replaced with the environment namespace. + + Example: ``{ENV_REGEX_NS}/Robot/sensor`` will be replaced with ``/World/envs/env_.*/Robot/sensor``. + + """ + + update_period: float = 0.0 + """Update period of the sensor buffers (in seconds). Defaults to 0.0 (update every step).""" + + history_length: int = 0 + """Number of past frames to store in the sensor buffers. Defaults to 0, which means that only + the current data is stored (no history).""" + + debug_vis: bool = False + """Whether to visualize the sensor. Defaults to False.""" diff --git a/source/isaaclab/isaaclab/scene/interactive_scene.py b/source/isaaclab/isaaclab/scene/interactive_scene.py index 90dacd95c9c..31ae1496824 100644 --- a/source/isaaclab/isaaclab/scene/interactive_scene.py +++ b/source/isaaclab/isaaclab/scene/interactive_scene.py @@ -8,8 +8,9 @@ from typing import Any import carb -import omni.log -import omni.usd +import logging +import warnings +import warp as wp from isaacsim.core.prims import XFormPrim from isaacsim.core.utils.stage import get_current_stage from isaacsim.core.version import get_version @@ -25,6 +26,9 @@ from .interactive_scene_cfg import InteractiveSceneCfg +logger = logging.getLogger(__name__) +warnings.simplefilter("once", UserWarning) +logging.captureWarnings(True) class InteractiveScene: """A scene that contains entities added to the simulation. @@ -188,7 +192,7 @@ def clone_environments(self, copy_from_source: bool = False): carb_settings_iface = carb.settings.get_settings() has_multi_assets = carb_settings_iface.get("/isaaclab/spawn/multi_assets") if has_multi_assets and self.cfg.replicate_physics: - omni.log.warn( + warnings.warn( "Varying assets might have been spawned under different environments." " However, the replicate physics flag is enabled in the 'InteractiveScene' configuration." " This may adversely affect PhysX parsing. We recommend disabling this property." @@ -226,7 +230,7 @@ def clone_environments(self, copy_from_source: bool = False): if (not self.cfg.replicate_physics and self.cfg.filter_collisions) or self.device == "cpu": # if scene is specified through cfg, this is already taken care of if not self._is_scene_setup_from_cfg(): - omni.log.warn( + warnings.warn( "Collision filtering can only be automatically enabled when replicate_physics=True and using GPU" " simulation. Please call scene.filter_collisions(global_prim_paths) to filter collisions across" " environments." @@ -290,7 +294,7 @@ def physics_scene_path(self) -> str: for prim in self.stage.Traverse(): if prim.HasAPI(PhysxSchema.PhysxSceneAPI): self._physics_scene_path = prim.GetPrimPath().pathString - omni.log.info(f"Physics scene prim path: {self._physics_scene_path}") + logger.info(f"Physics scene prim path: {self._physics_scene_path}") break if self._physics_scene_path is None: raise RuntimeError("No physics scene found! Please make sure one exists.") @@ -404,7 +408,7 @@ def state(self) -> dict[str, dict[str, dict[str, torch.Tensor]]]: Operations. """ - def reset(self, env_ids: Sequence[int] | None = None): + def reset(self, env_ids: Sequence[int] | None = None, mask: wp.array | torch.Tensor | None = None): """Resets the scene entities. Args: @@ -413,10 +417,10 @@ def reset(self, env_ids: Sequence[int] | None = None): """ # -- assets for articulation in self._articulations.values(): - articulation.reset(env_ids) + articulation.reset(ids = env_ids, mask = mask) # -- sensors for sensor in self._sensors.values(): - sensor.reset(env_ids) + sensor.reset(ids = env_ids, mask = mask) def write_data_to_sim(self): """Writes the data of the scene entities to the simulation.""" diff --git a/source/isaaclab/isaaclab/sim/_impl/newton_manager.py b/source/isaaclab/isaaclab/sim/_impl/newton_manager.py index ed150b4c3bf..3071485ab29 100644 --- a/source/isaaclab/isaaclab/sim/_impl/newton_manager.py +++ b/source/isaaclab/isaaclab/sim/_impl/newton_manager.py @@ -13,7 +13,7 @@ from newton import Axis, Contacts, Control, Model, ModelBuilder, State, eval_fk from newton.sensors import ContactSensor as NewtonContactSensor from newton.sensors import populate_contacts -from newton.solvers import SolverBase, SolverFeatherstone, SolverMuJoCo, SolverXPBD +from newton.solvers import SolverBase, SolverFeatherstone, SolverMuJoCo, SolverNotifyFlags, SolverXPBD from isaaclab.sim._impl.newton_manager_cfg import NewtonCfg from isaaclab.sim._impl.newton_viewer import NewtonViewerGL @@ -80,6 +80,7 @@ class NewtonManager: _visualizer_update_frequency: int = 1 # Configurable frequency for all rendering updates _visualizer_train_mode: bool = True # Whether visualizer is in training mode _visualizer_disabled: bool = False # Whether visualizer has been disabled by user + _model_changes: set[int] = set() @classmethod def clear(cls): @@ -106,6 +107,7 @@ def clear(cls): NewtonManager._visualizer_update_counter = 0 NewtonManager._visualizer_disabled = False NewtonManager._visualizer_update_frequency = NewtonManager._cfg.newton_viewer_update_frequency + NewtonManager._model_changes = set() @classmethod def set_builder(cls, builder): @@ -119,6 +121,10 @@ def add_on_init_callback(cls, callback) -> None: def add_on_start_callback(cls, callback) -> None: NewtonManager._on_start_callbacks.append(callback) + @classmethod + def add_model_change(cls, change: SolverNotifyFlags) -> None: + NewtonManager._model_changes.add(change) + @classmethod def start_simulation(cls) -> None: """Starts the simulation. @@ -133,7 +139,7 @@ def start_simulation(cls) -> None: for callback in NewtonManager._on_init_callbacks: callback() print(f"[INFO] Finalizing model on device: {NewtonManager._device}") - NewtonManager._builder.gravity = np.array(NewtonManager._gravity_vector) + NewtonManager._builder.gravity = np.array(NewtonManager._gravity_vector)[-1] NewtonManager._builder.up_axis = Axis.from_string(NewtonManager._up_axis) with Timer(name="newton_finalize_builder", msg="Finalize builder took:", enable=True, format="ms"): NewtonManager._model = NewtonManager._builder.finalize(device=NewtonManager._device) @@ -187,6 +193,7 @@ def initialize_solver(cls): with Timer(name="newton_initialize_solver", msg="Initialize solver took:", enable=True, format="ms"): NewtonManager._num_substeps = NewtonManager._cfg.num_substeps NewtonManager._solver_dt = NewtonManager._dt / NewtonManager._num_substeps + print(NewtonManager._model.gravity) NewtonManager._solver = NewtonManager._get_solver(NewtonManager._model, NewtonManager._cfg.solver_cfg) # Ensure we are using a CUDA enabled device @@ -271,6 +278,11 @@ def step(cls) -> None: This function steps the simulation by the specified time step in the simulation configuration. """ + if NewtonManager._model_changes: + for change in NewtonManager._model_changes: + NewtonManager._solver.notify_model_changed(change) + NewtonManager._model_changes = set() + if NewtonManager._cfg.use_cuda_graph: wp.capture_launch(NewtonManager._graph) else: @@ -278,7 +290,7 @@ def step(cls) -> None: if NewtonManager._cfg.debug_mode: convergence_data = NewtonManager.get_solver_convergence_steps() - # print(f"solver niter: {convergence_data}") + print(f"solver niter: {convergence_data}") if convergence_data["max"] == NewtonManager._solver.mjw_model.opt.iterations: print("solver didn't converge!", convergence_data["max"]) @@ -414,6 +426,14 @@ def get_state_1(cls): def get_control(cls): return NewtonManager._control + @classmethod + def get_dt(cls): + return NewtonManager._dt + + @classmethod + def get_solver_dt(cls): + return NewtonManager._solver_dt + @classmethod def forward_kinematics(cls, mask: wp.array | None = None) -> None: """Evaluates the forward kinematics for the selected articulations. diff --git a/source/isaaclab/isaaclab/sim/_impl/newton_viewer.py b/source/isaaclab/isaaclab/sim/_impl/newton_viewer.py index ddd72859fb0..c8b1ecec546 100644 --- a/source/isaaclab/isaaclab/sim/_impl/newton_viewer.py +++ b/source/isaaclab/isaaclab/sim/_impl/newton_viewer.py @@ -16,6 +16,7 @@ from __future__ import annotations import newton as nt +import warp as wp from newton.viewer import ViewerGL @@ -150,7 +151,7 @@ def _render_left_panel(self): imgui.text(f"Environments: {self.model.num_envs}") axis_names = ["X", "Y", "Z"] imgui.text(f"Up Axis: {axis_names[self.model.up_axis]}") - gravity = self.model.gravity + gravity = wp.to_torch(self.model.gravity).cpu().numpy()[0] gravity_text = f"Gravity: ({gravity[0]:.2f}, {gravity[1]:.2f}, {gravity[2]:.2f})" imgui.text(gravity_text) diff --git a/source/isaaclab/isaaclab/sim/simulation_context.py b/source/isaaclab/isaaclab/sim/simulation_context.py index 33c93873c1c..14fa782a5b2 100644 --- a/source/isaaclab/isaaclab/sim/simulation_context.py +++ b/source/isaaclab/isaaclab/sim/simulation_context.py @@ -636,6 +636,36 @@ def step(self, render: bool = True): if "cuda" in self.device: torch.cuda.set_device(self.device) + def step_warp(self, render: bool = True): + """Steps the simulation. + + .. note:: + This function blocks if the timeline is paused. It only returns when the timeline is playing. + + Args: + render: Whether to render the scene after stepping the physics simulation. + If set to False, the scene is not rendered and only the physics simulation is stepped. + """ + + if render: + # physics dt is zero, no need to step physics, just render + if self.is_playing(): + NewtonManager.step() + if self.get_physics_dt() == 0: # noqa: SIM114 + SimulationContext.render(self) + # rendering dt is zero, but physics is not, call step and then render + elif self.get_rendering_dt() == 0 and self.get_physics_dt() != 0: # noqa: SIM114 + SimulationContext.render(self) + else: + self._app.update() + else: + if self.is_playing(): + NewtonManager.step() + + # Use the NewtonManager to render the scene if enabled + if self.cfg.enable_newton_rendering: + NewtonManager.render() + def render(self, mode: RenderMode | None = None): """Refreshes the rendering components including UI elements and view-ports depending on the render mode. diff --git a/source/isaaclab/isaaclab/utils/__init__.py b/source/isaaclab/isaaclab/utils/__init__.py index a5365948699..a6fc8c8eb71 100644 --- a/source/isaaclab/isaaclab/utils/__init__.py +++ b/source/isaaclab/isaaclab/utils/__init__.py @@ -9,9 +9,11 @@ from .buffers import * from .configclass import configclass from .dict import * +from .helpers import deprecated, warn_overhead_cost from .interpolation import * from .modifiers import * from .string import * from .timer import Timer from .types import * from .version import * +from .backend_utils import * \ No newline at end of file diff --git a/source/isaaclab/isaaclab/utils/backend_utils.py b/source/isaaclab/isaaclab/utils/backend_utils.py new file mode 100644 index 00000000000..1b694c1e63a --- /dev/null +++ b/source/isaaclab/isaaclab/utils/backend_utils.py @@ -0,0 +1,72 @@ +import importlib +import logging + +logger = logging.getLogger(__name__) + +class FactoryBase: + """A generic factory class that dynamically loads backends.""" + + def __init_subclass__(cls, **kwargs): + """Initializes a new factory subclass.""" + super().__init_subclass__(**kwargs) + cls._registry = {} + # Determine the module subpath for dynamic loading. + # e.g., if factory is in 'isaaclab.assets.articulation.articulation', + # the subpath becomes 'assets.articulation'. + module_parts = cls.__module__.split('.') + if module_parts[0] != 'isaaclab': + raise ImportError(f"Factory class {cls.__name__} must be defined within the 'isaaclab' package.") + # The subpath is what comes between 'isaaclab' and the final module name. + cls._module_subpath = ".".join(module_parts[1:-1]) + + @classmethod + def register(cls, name: str, sub_class) -> None: + """Register a new implementation class.""" + if name in cls._registry and cls._registry[name] is not sub_class: + raise ValueError(f"Backend {name!r} already registered with a different class for factory {cls.__name__}.") + cls._registry[name] = sub_class + logger.info(f"Registered backend {name!r} for factory {cls.__name__}.") + + def __new__(cls, *args, **kwargs): + """Create a new instance of an implementation based on the backend.""" + + backend = "newton" + + if cls == FactoryBase: + raise TypeError("FactoryBase cannot be instantiated directly. Please subclass it.") + + # If backend is not in registry, try to import it and register the class. + # This is done to only import the module once. + if backend not in cls._registry: + # Construct the module name from the backend and the determined subpath. + module_name = f"isaaclab.{backend}.{cls._module_subpath}" + try: + module = importlib.import_module(module_name) + module_class = getattr(module, cls.__name__) + # Manually register the class + cls.register(backend, module_class) + + except ImportError as e: + raise ValueError( + f"Could not import module for backend {backend!r} for factory {cls.__name__}. " + f"Attempted to import from '{module_name}'.\n" + f"Original error: {e}" + ) from e + + # Now check registry again. The import should have registered the class. + try: + impl = cls._registry[backend] + except KeyError: + available = list(cls.get_registry_keys()) + raise ValueError( + f"Unknown backend {backend!r} for {cls.__name__}. " + f"A module was found at '{module_name}', but it did not contain a class with the name {cls.__name__}.\n" + f"Currently available backends: {available}." + ) from None + # Return an instance of the chosen class. + return impl(*args, **kwargs) + + @classmethod + def get_registry_keys(cls) -> list[str]: + """Returns a list of registered backend names.""" + return list(cls._registry.keys()) \ No newline at end of file diff --git a/source/isaaclab/isaaclab/utils/buffers/__init__.py b/source/isaaclab/isaaclab/utils/buffers/__init__.py index dc63468b8d5..104b14866c7 100644 --- a/source/isaaclab/isaaclab/utils/buffers/__init__.py +++ b/source/isaaclab/isaaclab/utils/buffers/__init__.py @@ -8,3 +8,4 @@ from .circular_buffer import CircularBuffer from .delay_buffer import DelayBuffer from .timestamped_buffer import TimestampedBuffer +from .timestamped_wp_buffer import TimestampedWarpBuffer diff --git a/source/isaaclab/isaaclab/utils/buffers/timestamped_wp_buffer.py b/source/isaaclab/isaaclab/utils/buffers/timestamped_wp_buffer.py new file mode 100644 index 00000000000..db9eca0fd05 --- /dev/null +++ b/source/isaaclab/isaaclab/utils/buffers/timestamped_wp_buffer.py @@ -0,0 +1,43 @@ +# Copyright (c) 2022-2025, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +from dataclasses import dataclass + +import warp as wp + + +@dataclass +class TimestampedWarpBuffer: + """A buffer class containing data and its timestamp. + + This class is a simple data container that stores a tensor and its timestamp. The timestamp is used to + track the last update of the buffer. The timestamp is set to -1.0 by default, indicating that the buffer + has not been updated yet. The timestamp should be updated whenever the data in the buffer is updated. This + way the buffer can be used to check whether the data is outdated and needs to be refreshed. + + The buffer is useful for creating lazy buffers that only update the data when it is outdated. This can be + useful when the data is expensive to compute or retrieve. For example usage, refer to the data classes in + the :mod:`isaaclab.assets` module. + """ + + data: wp.array = None # type: ignore + """The data stored in the buffer. Default is None, indicating that the buffer is empty.""" + + timestamp: float = -1.0 + """Timestamp at the last update of the buffer. Default is -1.0, indicating that the buffer has not been updated.""" + + shape: tuple[int, ...] | None = None + """Shape of the data stored in the buffer. Default is None, indicating that the buffer is empty.""" + + dtype: type | None = None + """Dtype of the data stored in the buffer. Default is None, indicating that the buffer is empty.""" + + def __post_init__(self): + if self.shape is None: + raise ValueError("Shape of the data stored in the buffer is not set.") + if self.dtype is None: + raise ValueError("Dtype of the data stored in the buffer is not set.") + if self.data is None: + self.data = wp.empty(self.shape, dtype=self.dtype) diff --git a/source/isaaclab/isaaclab/utils/helpers.py b/source/isaaclab/isaaclab/utils/helpers.py new file mode 100644 index 00000000000..4d1d69d6239 --- /dev/null +++ b/source/isaaclab/isaaclab/utils/helpers.py @@ -0,0 +1,297 @@ +# Copyright (c) 2022-2025, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +"""Utility functions for functions.""" + +import functools +import inspect +import torch +from collections.abc import Callable +import logging +import warnings +import warp as wp + +logger = logging.getLogger(__name__) +warnings.simplefilter("once", UserWarning) +logging.captureWarnings(True) + +def deprecated( + *replacement_function_names: str, + message: str = "", + since: str | None = None, + remove_in: str | None = None, +): + """A decorator to mark functions as deprecated. + + It will result in a warning being emitted when the function is used. + + Args: + replacement_function_names: The names of the functions to use instead. + message: A custom message to append to the warning. + since: The version in which the function was deprecated. + remove_in: The version in which the function will be removed. + """ + + def decorator(func): + @functools.wraps(func) + def wrapper(*args, **kwargs): + # Form deprecation message. + deprecation_message = f"Call to deprecated function '{func.__name__}'." + # Add version information if provided. + if since: + deprecation_message += f" It was deprecated in version {since}." + if remove_in: + deprecation_message += f" It will be removed in version {remove_in}." + else: + deprecation_message += " It will be removed in a future version." + # Add replacement function information if provided. + if replacement_function_names: + deprecation_message += f" Use {', '.join(replacement_function_names)} instead." + # Add custom message if provided. + if message: + deprecation_message += f" {message}" + + # Emit warning. + warnings.warn( + deprecation_message, + UserWarning, + ) + # Call the original function. + return func(*args, **kwargs) + + return wrapper + + return decorator + + +def warn_overhead_cost( + *replacement_function_names: str, + message: str = "", +): + """A decorator to mark functions as having a high overhead cost. + + It will result in a warning being emitted when the function is used. + + Args: + replacement_function_names: The names of the functions to use instead. + message: A custom message to append to the warning. + """ + + def decorator(func): + @functools.wraps(func) + def wrapper(*args, **kwargs): + # Form deprecation message. + warning_message = f"Call to '{func.__name__}' which is a high overhead operation." + # Add replacement function information if provided. + if replacement_function_names: + warning_message += f" Use {', '.join(replacement_function_names)} instead." + # Add custom message if provided. + if message: + warning_message += f" {message}" + + # Emit warning. + warnings.warn( + warning_message, + UserWarning, + ) + # Call the original function. + return func(*args, **kwargs) + + return wrapper + + return decorator + + +def _analyze_and_convert_args(self, func, *args, **kwargs): + """A helper to analyze and convert arguments from PyTorch to Warp.""" + sig = inspect.signature(func) + bound_args = sig.bind(self, *args, **kwargs) + bound_args.apply_defaults() + arguments = bound_args.arguments + + # -- Device conversion + device = getattr(self, "device", "cpu") + + # -- Tensor conversion + spec = getattr(func, "_torch_frontend_spec", {}) + tensor_args = spec.get("tensor_args", {}) + first_torch_tensor = None + for arg_name in tensor_args: + arg_value = arguments.get(arg_name) + if isinstance(arg_value, torch.Tensor): + if first_torch_tensor is None: + first_torch_tensor = arg_value + tensor = arguments[arg_name] + dtype = tensor_args[arg_name] + arguments[arg_name] = wp.from_torch(tensor, dtype=dtype) + + # -- Mask conversion + mask_configs = { + "env_mask": {"id_arg": "env_ids", "shape_attrs": ["num_instances", "num_envs"]}, + "joint_mask": {"id_arg": "joint_ids", "shape_attrs": ["num_joints"]}, + "body_mask": {"id_arg": "body_ids", "shape_attrs": ["num_bodies"]}, + } + + for mask_name, config in mask_configs.items(): + id_arg_name = config["id_arg"] + if mask_name in sig.parameters and id_arg_name in arguments and arguments[id_arg_name] is not None: + indices = arguments.pop(id_arg_name) + shape_val = 0 + for attr in config["shape_attrs"]: + val = getattr(self, attr, None) + if val is not None: + shape_val = val + break + if shape_val == 0: + raise ValueError( + f"Cannot convert '{id_arg_name}' to '{mask_name}'. The instance is missing one of the " + f"following attributes: {config['shape_attrs']}." + ) + + mask_torch = torch.zeros(shape_val, dtype=torch.bool, device=device) + + if isinstance(indices, slice): + mask_torch[indices] = True + elif isinstance(indices, (list, tuple, torch.Tensor)): + mask_torch[torch.as_tensor(indices, device=device)] = True + else: + raise TypeError(f"Unsupported type for indices '{type(indices)}'.") + + arguments[mask_name] = wp.from_torch(mask_torch) + + arguments.pop("self") + return arguments + + +def _convert_output_to_torch(data): + """Recursively convert warp arrays in a data structure to torch tensors.""" + if isinstance(data, wp.array): + return wp.to_torch(data) + elif isinstance(data, (list, tuple)): + return type(data)(_convert_output_to_torch(item) for item in data) + elif isinstance(data, dict): + return {key: _convert_output_to_torch(value) for key, value in data.items()} + else: + return data + + +def torch_frontend_method(tensor_args: dict[str, any] | None = None, *, convert_output: bool = False): + """A method decorator to specify tensor conversion rules for the torch frontend. + + This decorator attaches metadata to a method, which is then used by the + `torch_frontend_class` decorator to apply the correct data conversions. + + Args: + tensor_args: A dictionary mapping tensor argument names to their + target `warp.dtype`. Defaults to None. + convert_output: If True, the output of the decorated function will be + converted from warp arrays to torch tensors. Defaults to False. + + Example: + >>> @torch_frontend_class + ... class MyAsset: + ... @torch_frontend_method({"root_state": wp.transformf}) + ... def write_root_state_to_sim(self, root_state: wp.array, env_mask: wp.array | None = None): + ... pass + ... + ... @torch_frontend_method(convert_output=True) + ... def get_root_state(self) -> wp.array: + ... pass + """ + if tensor_args is None: + tensor_args = {} + + def decorator(func: Callable) -> Callable: + setattr(func, "_torch_frontend_spec", {"tensor_args": tensor_args, "convert_output": convert_output}) + return func + + return decorator + + +def torch_frontend_class(cls=None, *, indices_arg: str = "env_ids", mask_arg: str = "env_mask"): + """A class decorator to add a PyTorch frontend to a class that uses a Warp backend. + + This decorator patches the ``__init__`` method of a class. After the original ``__init__`` is called, + it checks for a `frontend` attribute on the instance. If `self.frontend` is "torch", it inspects + the class for methods decorated with `@torch_frontend_method` and wraps them to make them + accept PyTorch tensors. + + The wrapped methods will: + - Convert specified PyTorch tensor arguments to Warp arrays based on the rules defined + in the `@torch_frontend_method` decorator. + - Convert an argument with indices (e.g., `env_ids`) to a boolean mask (e.g., `env_mask`). + + If `self.frontend` is not "torch", the class's methods remain unchanged, ensuring zero + overhead for the default Warp backend. + + Args: + cls: The class to decorate. This is handled automatically by Python. + indices_arg: The name of the argument that may contain indices for conversion to a mask. + Defaults to "env_ids". + mask_arg: The name of the argument that will receive the generated boolean mask. + Defaults to "env_mask". + + Example: + >>> import warp as wp + >>> import torch + >>> + >>> @torch_frontend_class + ... class MyAsset: + ... def __init__(self, num_envs, device, frontend="warp"): + ... self.num_instances = num_envs + ... self.device = device + ... self.frontend = frontend + ... wp.init() + ... + ... @torch_frontend_method({"root_state": wp.transformf}) + ... def write_root_state_to_sim(self, root_state: wp.array, env_mask: wp.array | None = None): + ... print(f"root_state type: {type(root_state)}") + ... if env_mask: + ... print(f"env_mask type: {type(env_mask)}") + ... + >>> # -- Using warp frontend (no overhead) + >>> asset_wp = MyAsset(num_envs=4, device="cpu", frontend="warp") + >>> root_state_wp = wp.zeros(4, dtype=wp.transformf) + >>> asset_wp.write_root_state_to_sim(root_state_wp) + root_state type: + >>> + >>> # -- Using torch frontend (methods are patched) + >>> asset_torch = MyAsset(num_envs=4, device="cpu", frontend="torch") + >>> root_state_torch = torch.rand(4, 7) + >>> asset_torch.write_root_state_to_sim(root_state_torch, env_ids=[0, 2]) + root_state type: + env_mask type: + """ + # This allows using the decorator with or without parentheses: + # @torch_frontend_class or @torch_frontend_class(indices_arg="...") + if cls is None: + return functools.partial(torch_frontend_class, indices_arg=indices_arg, mask_arg=mask_arg) + + original_init = cls.__init__ + + @functools.wraps(original_init) + def new_init(self, *args, **kwargs): + original_init(self, *args, **kwargs) + + if getattr(self, "frontend", "warp") == "torch": + for name, method in inspect.getmembers(cls, predicate=inspect.isfunction): + if hasattr(method, "_torch_frontend_spec"): + spec = getattr(method, "_torch_frontend_spec") + convert_output = spec.get("convert_output", False) + + @functools.wraps(method) + def adapted_method_wrapper(self, *args, method=method, **kwargs): + converted_args = _analyze_and_convert_args(self, method, *args, **kwargs) + output = method(self, **converted_args) + + if convert_output: + return _convert_output_to_torch(output) + else: + return output + + setattr(self, name, adapted_method_wrapper.__get__(self, cls)) + + cls.__init__ = new_init + return cls diff --git a/source/isaaclab_assets/isaaclab_assets/robots/allegro.py b/source/isaaclab_assets/isaaclab_assets/robots/allegro.py index ef5b5491f57..bf342edcc3a 100644 --- a/source/isaaclab_assets/isaaclab_assets/robots/allegro.py +++ b/source/isaaclab_assets/isaaclab_assets/robots/allegro.py @@ -18,6 +18,7 @@ import isaaclab.sim as sim_utils from isaaclab.actuators.actuator_cfg import ImplicitActuatorCfg +from isaaclab.actuators import ControlMode from isaaclab.assets.articulation import ArticulationCfg from isaaclab.utils.assets import ISAAC_NUCLEUS_DIR @@ -39,7 +40,7 @@ ), actuators={ "fingers": ImplicitActuatorCfg( - control_mode="position", + control_mode=ControlMode.POSITION, joint_names_expr=[".*"], effort_limit_sim=0.5, stiffness=1.0, diff --git a/source/isaaclab_assets/isaaclab_assets/robots/ant.py b/source/isaaclab_assets/isaaclab_assets/robots/ant.py index 69b59ecab04..6417971b663 100644 --- a/source/isaaclab_assets/isaaclab_assets/robots/ant.py +++ b/source/isaaclab_assets/isaaclab_assets/robots/ant.py @@ -8,7 +8,7 @@ from __future__ import annotations import isaaclab.sim as sim_utils -from isaaclab.actuators import ImplicitActuatorCfg +from isaaclab.actuators import ImplicitActuatorCfg, ControlMode from isaaclab.assets import ArticulationCfg from isaaclab.utils.assets import ISAAC_NUCLEUS_DIR @@ -35,7 +35,7 @@ actuators={ "body": ImplicitActuatorCfg( joint_names_expr=[".*"], - control_mode="none", + control_mode=ControlMode.NONE, stiffness=0.0, damping=0.0, effort_limit_sim=30.0, @@ -43,4 +43,4 @@ ), }, ) -"""Configuration for the Mujoco Ant robot.""" +"""Configuration for the Mujoco Ant robot.""" \ No newline at end of file diff --git a/source/isaaclab_assets/isaaclab_assets/robots/anymal.py b/source/isaaclab_assets/isaaclab_assets/robots/anymal.py index 48b61c85191..5f9f93a5ea0 100644 --- a/source/isaaclab_assets/isaaclab_assets/robots/anymal.py +++ b/source/isaaclab_assets/isaaclab_assets/robots/anymal.py @@ -21,7 +21,7 @@ import isaaclab.sim as sim_utils -from isaaclab.actuators import ActuatorNetLSTMCfg, DCMotorCfg +from isaaclab.actuators import ActuatorNetLSTMCfg, DCMotorCfg, ControlMode from isaaclab.assets.articulation import ArticulationCfg from isaaclab.utils.assets import ISAACLAB_NUCLEUS_DIR @@ -31,7 +31,7 @@ ANYDRIVE_3_SIMPLE_ACTUATOR_CFG = DCMotorCfg( joint_names_expr=[".*HAA", ".*HFE", ".*KFE"], - control_mode="position", + control_mode=ControlMode.POSITION, saturation_effort=120.0, effort_limit=80.0, effort_limit_sim=80.0, diff --git a/source/isaaclab_assets/isaaclab_assets/robots/cartpole.py b/source/isaaclab_assets/isaaclab_assets/robots/cartpole.py index ad577ae6bba..478cd8b65d1 100644 --- a/source/isaaclab_assets/isaaclab_assets/robots/cartpole.py +++ b/source/isaaclab_assets/isaaclab_assets/robots/cartpole.py @@ -7,7 +7,7 @@ import isaaclab.sim as sim_utils -from isaaclab.actuators import ImplicitActuatorCfg +from isaaclab.actuators import ImplicitActuatorCfg, ControlMode from isaaclab.assets import ArticulationCfg from isaaclab.utils.assets import ISAACLAB_NUCLEUS_DIR @@ -28,7 +28,7 @@ actuators={ "cart_actuator": ImplicitActuatorCfg( joint_names_expr=["slider_to_cart"], - control_mode="position", + control_mode=ControlMode.POSITION, effort_limit=400.0, velocity_limit=100.0, stiffness=0.0, @@ -36,7 +36,7 @@ ), "pole_actuator": ImplicitActuatorCfg( joint_names_expr=["cart_to_pole"], - control_mode="none", + control_mode=ControlMode.NONE, effort_limit=400.0, velocity_limit=100.0, stiffness=0.0, @@ -44,4 +44,4 @@ ), }, ) -"""Configuration for a simple Cartpole robot.""" +"""Configuration for a simple Cartpole robot.""" \ No newline at end of file diff --git a/source/isaaclab_assets/isaaclab_assets/robots/cassie.py b/source/isaaclab_assets/isaaclab_assets/robots/cassie.py index 94d6179c122..15e462b76b6 100644 --- a/source/isaaclab_assets/isaaclab_assets/robots/cassie.py +++ b/source/isaaclab_assets/isaaclab_assets/robots/cassie.py @@ -13,7 +13,7 @@ """ import isaaclab.sim as sim_utils -from isaaclab.actuators import ImplicitActuatorCfg +from isaaclab.actuators import ImplicitActuatorCfg, ControlMode from isaaclab.assets.articulation import ArticulationCfg from isaaclab.utils.assets import ISAACLAB_NUCLEUS_DIR @@ -50,7 +50,7 @@ actuators={ "legs": ImplicitActuatorCfg( joint_names_expr=["hip_.*", "thigh_.*", "ankle_.*"], - control_mode="position", + control_mode=ControlMode.POSITION, effort_limit_sim=200.0, stiffness={ "hip_abduction.*": 100.0, @@ -71,7 +71,7 @@ ), "toes": ImplicitActuatorCfg( joint_names_expr=["toe_.*"], - control_mode="position", + control_mode=ControlMode.POSITION, effort_limit_sim=20.0, stiffness={ "toe_joint.*": 20.0, diff --git a/source/isaaclab_assets/isaaclab_assets/robots/franka.py b/source/isaaclab_assets/isaaclab_assets/robots/franka.py index 5411686d396..73c72dc936b 100644 --- a/source/isaaclab_assets/isaaclab_assets/robots/franka.py +++ b/source/isaaclab_assets/isaaclab_assets/robots/franka.py @@ -14,7 +14,7 @@ """ import isaaclab.sim as sim_utils -from isaaclab.actuators import ImplicitActuatorCfg +from isaaclab.actuators import ImplicitActuatorCfg, ControlMode from isaaclab.assets.articulation import ArticulationCfg from isaaclab.utils.assets import ISAACLAB_NUCLEUS_DIR @@ -43,7 +43,7 @@ ), actuators={ "panda_shoulder": ImplicitActuatorCfg( - control_mode="position", + control_mode=ControlMode.POSITION, joint_names_expr=["panda_joint[1-4]"], effort_limit_sim=87.0, stiffness=80.0, @@ -52,7 +52,7 @@ armature=1e-3, ), "panda_forearm": ImplicitActuatorCfg( - control_mode="position", + control_mode=ControlMode.POSITION, joint_names_expr=["panda_joint[5-7]"], effort_limit_sim=12.0, stiffness=80.0, @@ -61,7 +61,7 @@ armature=1e-3, ), "panda_hand": ImplicitActuatorCfg( - control_mode="position", + control_mode=ControlMode.POSITION, joint_names_expr=["panda_finger_joint.*"], effort_limit_sim=3.0, stiffness=20.0, diff --git a/source/isaaclab_assets/isaaclab_assets/robots/humanoid.py b/source/isaaclab_assets/isaaclab_assets/robots/humanoid.py index 0d1a156f470..99beeea72ce 100644 --- a/source/isaaclab_assets/isaaclab_assets/robots/humanoid.py +++ b/source/isaaclab_assets/isaaclab_assets/robots/humanoid.py @@ -8,7 +8,7 @@ from __future__ import annotations import isaaclab.sim as sim_utils -from isaaclab.actuators import ImplicitActuatorCfg +from isaaclab.actuators import ImplicitActuatorCfg, ControlMode from isaaclab.assets import ArticulationCfg from isaaclab.utils.assets import ISAAC_NUCLEUS_DIR @@ -32,7 +32,7 @@ actuators={ "body": ImplicitActuatorCfg( joint_names_expr=[".*"], - control_mode="none", + control_mode=ControlMode.POSITION, stiffness={ ".*_waist.*": 20.0, ".*_upper_arm.*": 10.0, @@ -62,4 +62,4 @@ ), }, ) -"""Configuration for the Mujoco Humanoid robot.""" +"""Configuration for the Mujoco Humanoid robot.""" \ No newline at end of file diff --git a/source/isaaclab_assets/isaaclab_assets/robots/spot.py b/source/isaaclab_assets/isaaclab_assets/robots/spot.py index c6aaff1eaae..9f48811d71c 100644 --- a/source/isaaclab_assets/isaaclab_assets/robots/spot.py +++ b/source/isaaclab_assets/isaaclab_assets/robots/spot.py @@ -12,7 +12,7 @@ """ import isaaclab.sim as sim_utils -from isaaclab.actuators import DelayedPDActuatorCfg, RemotizedPDActuatorCfg +from isaaclab.actuators import DelayedPDActuatorCfg, RemotizedPDActuatorCfg, ControlMode from isaaclab.assets.articulation import ArticulationCfg from isaaclab.utils.assets import ISAAC_NUCLEUS_DIR diff --git a/source/isaaclab_assets/isaaclab_assets/robots/unitree.py b/source/isaaclab_assets/isaaclab_assets/robots/unitree.py index 6aa34baf16d..7e6379c9bee 100644 --- a/source/isaaclab_assets/isaaclab_assets/robots/unitree.py +++ b/source/isaaclab_assets/isaaclab_assets/robots/unitree.py @@ -19,7 +19,7 @@ """ import isaaclab.sim as sim_utils -from isaaclab.actuators import ActuatorNetMLPCfg, DCMotorCfg, ImplicitActuatorCfg +from isaaclab.actuators import ActuatorNetMLPCfg, DCMotorCfg, ImplicitActuatorCfg, ControlMode from isaaclab.assets.articulation import ArticulationCfg from isaaclab.utils.assets import ISAAC_NUCLEUS_DIR, ISAACLAB_NUCLEUS_DIR @@ -173,7 +173,7 @@ actuators={ "legs": ImplicitActuatorCfg( joint_names_expr=[".*_hip_yaw", ".*_hip_roll", ".*_hip_pitch", ".*_knee", "torso"], - control_mode="position", + control_mode=ControlMode.POSITION, effort_limit=300, velocity_limit=100.0, stiffness={ @@ -194,7 +194,7 @@ ), "feet": ImplicitActuatorCfg( joint_names_expr=[".*_ankle"], - control_mode="position", + control_mode=ControlMode.POSITION, effort_limit=100, velocity_limit=100.0, stiffness={".*_ankle": 20.0}, @@ -203,7 +203,7 @@ ), "arms": ImplicitActuatorCfg( joint_names_expr=[".*_shoulder_pitch", ".*_shoulder_roll", ".*_shoulder_yaw", ".*_elbow"], - control_mode="position", + control_mode=ControlMode.POSITION, effort_limit=300, velocity_limit=100.0, stiffness={ @@ -273,7 +273,7 @@ ".*_knee_joint", "torso_joint", ], - control_mode="position", + control_mode=ControlMode.POSITION, effort_limit=300, velocity_limit=100.0, stiffness={ @@ -298,7 +298,7 @@ friction=0.00001, ), "feet": ImplicitActuatorCfg( - control_mode="position", + control_mode=ControlMode.POSITION, effort_limit=20, joint_names_expr=[".*_ankle_pitch_joint", ".*_ankle_roll_joint"], stiffness=20.0, @@ -321,7 +321,7 @@ ".*_one_joint", ".*_two_joint", ], - control_mode="position", + control_mode=ControlMode.POSITION, effort_limit=300, velocity_limit=100.0, stiffness=40.0, diff --git a/source/isaaclab_assets/isaaclab_assets/robots/universal_robots.py b/source/isaaclab_assets/isaaclab_assets/robots/universal_robots.py index 9eac923a2e3..037a2b7c101 100644 --- a/source/isaaclab_assets/isaaclab_assets/robots/universal_robots.py +++ b/source/isaaclab_assets/isaaclab_assets/robots/universal_robots.py @@ -13,7 +13,7 @@ """ import isaaclab.sim as sim_utils -from isaaclab.actuators import ImplicitActuatorCfg +from isaaclab.actuators import ImplicitActuatorCfg, ControlMode from isaaclab.assets.articulation import ArticulationCfg from isaaclab.utils.assets import ISAACLAB_NUCLEUS_DIR @@ -38,7 +38,7 @@ ), actuators={ "arm": ImplicitActuatorCfg( - control_mode="position", + control_mode=ControlMode.POSITION, joint_names_expr=[".*"], effort_limit_sim=87.0, stiffness=800.0, diff --git a/source/isaaclab_rl/isaaclab_rl/rsl_rl/vecenv_wrapper.py b/source/isaaclab_rl/isaaclab_rl/rsl_rl/vecenv_wrapper.py index 73ceae04693..f980cd6e4a9 100644 --- a/source/isaaclab_rl/isaaclab_rl/rsl_rl/vecenv_wrapper.py +++ b/source/isaaclab_rl/isaaclab_rl/rsl_rl/vecenv_wrapper.py @@ -7,9 +7,10 @@ import torch from tensordict import TensorDict +import warp as wp from rsl_rl.env import VecEnv -from isaaclab.envs import DirectRLEnv, ManagerBasedRLEnv +from isaaclab.envs import DirectRLEnv, DirectRLEnvWarp, ManagerBasedRLEnv class RslRlVecEnvWrapper(VecEnv): @@ -24,7 +25,7 @@ class RslRlVecEnvWrapper(VecEnv): https://github.com/leggedrobotics/rsl_rl/blob/master/rsl_rl/env/vec_env.py """ - def __init__(self, env: ManagerBasedRLEnv | DirectRLEnv, clip_actions: float | None = None): + def __init__(self, env: ManagerBasedRLEnv | DirectRLEnv | DirectRLEnvWarp, clip_actions: float | None = None): """Initializes the wrapper. Note: @@ -35,11 +36,15 @@ def __init__(self, env: ManagerBasedRLEnv | DirectRLEnv, clip_actions: float | N clip_actions: The clipping value for actions. If ``None``, then no clipping is done. Raises: - ValueError: When the environment is not an instance of :class:`ManagerBasedRLEnv` or :class:`DirectRLEnv`. + ValueError: When the environment is not an instance of :class:`ManagerBasedRLEnv` or :class:`DirectRLEnv` or :class:`DirectRLEnvWarp`. """ # check that input is valid - if not isinstance(env.unwrapped, ManagerBasedRLEnv) and not isinstance(env.unwrapped, DirectRLEnv): + if ( + not isinstance(env.unwrapped, ManagerBasedRLEnv) + and not isinstance(env.unwrapped, DirectRLEnv) + and not isinstance(env.unwrapped, DirectRLEnvWarp) + ): raise ValueError( "The environment must be inherited from ManagerBasedRLEnv or DirectRLEnv. Environment type:" f" {type(env)}" @@ -104,7 +109,7 @@ def class_name(cls) -> str: return cls.__name__ @property - def unwrapped(self) -> ManagerBasedRLEnv | DirectRLEnv: + def unwrapped(self) -> ManagerBasedRLEnv | DirectRLEnv | DirectRLEnvWarp: """Returns the base environment of the wrapper. This will be the bare :class:`gymnasium.Env` environment, underneath all layers of wrappers. @@ -118,7 +123,10 @@ def unwrapped(self) -> ManagerBasedRLEnv | DirectRLEnv: @property def episode_length_buf(self) -> torch.Tensor: """The episode length buffer.""" - return self.unwrapped.episode_length_buf + if isinstance(self.unwrapped, DirectRLEnvWarp): + return wp.to_torch(self.unwrapped.episode_length_buf) + else: + return self.unwrapped.episode_length_buf @episode_length_buf.setter def episode_length_buf(self, value: torch.Tensor): @@ -127,7 +135,10 @@ def episode_length_buf(self, value: torch.Tensor): Note: This is needed to perform random initialization of episode lengths in RSL-RL. """ - self.unwrapped.episode_length_buf = value + if isinstance(self.unwrapped, DirectRLEnvWarp): + self.unwrapped.episode_length_buf = wp.from_torch(value) + else: + self.unwrapped.episode_length_buf = value """ Operations - MDP @@ -146,7 +157,11 @@ def get_observations(self) -> TensorDict: if hasattr(self.unwrapped, "observation_manager"): obs_dict = self.unwrapped.observation_manager.compute() else: - obs_dict = self.unwrapped._get_observations() + if isinstance(self.unwrapped, DirectRLEnvWarp): + self.unwrapped._get_observations() + obs_dict = {"policy": self.unwrapped.torch_obs_buf.clone()} + else: + obs_dict = self.unwrapped._get_observations() return TensorDict(obs_dict, batch_size=[self.num_envs]) def step(self, actions: torch.Tensor) -> tuple[TensorDict, torch.Tensor, torch.Tensor, dict]: @@ -162,7 +177,7 @@ def step(self, actions: torch.Tensor) -> tuple[TensorDict, torch.Tensor, torch.T if not self.unwrapped.cfg.is_finite_horizon: extras["time_outs"] = truncated # return the step information - return TensorDict(obs_dict, batch_size=[self.num_envs]), rew, dones, extras + return TensorDict(obs_dict, batch_size=[self.num_envs]), rew.clone(), dones.clone(), extras def close(self): # noqa: D102 return self.env.close() diff --git a/source/isaaclab_tasks/isaaclab_tasks/direct/ant/__init__.py b/source/isaaclab_tasks/isaaclab_tasks/direct/ant/__init__.py index 5f66eda9885..58a5f94b135 100644 --- a/source/isaaclab_tasks/isaaclab_tasks/direct/ant/__init__.py +++ b/source/isaaclab_tasks/isaaclab_tasks/direct/ant/__init__.py @@ -26,3 +26,15 @@ "skrl_cfg_entry_point": f"{agents.__name__}:skrl_ppo_cfg.yaml", }, ) + +gym.register( + id="Isaac-Ant-Warp-v0", + entry_point=f"{__name__}.ant_env_warp:AntWarpEnv", + disable_env_checker=True, + kwargs={ + "env_cfg_entry_point": f"{__name__}.ant_env_warp:AntWarpEnvCfg", + "rl_games_cfg_entry_point": f"{agents.__name__}:rl_games_ppo_cfg.yaml", + "rsl_rl_cfg_entry_point": f"{agents.__name__}.rsl_rl_ppo_cfg:AntPPORunnerCfg", + "skrl_cfg_entry_point": f"{agents.__name__}:skrl_ppo_cfg.yaml", + }, +) diff --git a/source/isaaclab_tasks/isaaclab_tasks/direct/ant/ant_env_warp.py b/source/isaaclab_tasks/isaaclab_tasks/direct/ant/ant_env_warp.py new file mode 100644 index 00000000000..f7704a06e09 --- /dev/null +++ b/source/isaaclab_tasks/isaaclab_tasks/direct/ant/ant_env_warp.py @@ -0,0 +1,93 @@ +# Copyright (c) 2022-2025, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +from __future__ import annotations + +from isaaclab_assets import ANT_CFG + +import isaaclab.sim as sim_utils +from isaaclab.assets import ArticulationCfg +from isaaclab.envs import DirectRLEnvCfg +from isaaclab.scene import InteractiveSceneCfg +from isaaclab.sim import SimulationCfg +from isaaclab.sim._impl.newton_manager_cfg import NewtonCfg +from isaaclab.sim._impl.solvers_cfg import MJWarpSolverCfg +from isaaclab.terrains import TerrainImporterCfg +from isaaclab.utils import configclass + +from isaaclab_tasks.direct.locomotion.locomotion_env_warp import LocomotionWarpEnv + + +@configclass +class AntWarpEnvCfg(DirectRLEnvCfg): + # env + episode_length_s = 15.0 + decimation = 2 + action_scale = 0.5 + action_space = 8 + observation_space = 36 + state_space = 0 + + solver_cfg = MJWarpSolverCfg( + njmax=38, + ncon_per_env=15, + ls_iterations=10, + cone="pyramidal", + ls_parallel=True, + impratio=1, + update_data_interval=1, + ) + newton_cfg = NewtonCfg( + solver_cfg=solver_cfg, + num_substeps=1, + debug_mode=False, + use_cuda_graph=True, + ) + + # simulation + sim: SimulationCfg = SimulationCfg(dt=1 / 120, render_interval=decimation, newton_cfg=newton_cfg) + terrain = TerrainImporterCfg( + prim_path="/World/ground", + terrain_type="plane", + collision_group=-1, + physics_material=sim_utils.RigidBodyMaterialCfg( + friction_combine_mode="average", + restitution_combine_mode="average", + static_friction=1.0, + dynamic_friction=1.0, + restitution=0.0, + ), + debug_vis=False, + ) + + # scene + scene: InteractiveSceneCfg = InteractiveSceneCfg( + num_envs=4096, env_spacing=4.0, replicate_physics=True, clone_in_fabric=True + ) + + # robot + robot: ArticulationCfg = ANT_CFG.replace(prim_path="/World/envs/env_.*/Robot") + joint_gears: list = [15, 15, 15, 15, 15, 15, 15, 15] + + heading_weight: float = 0.5 + up_weight: float = 0.1 + + energy_cost_scale: float = 0.05 + actions_cost_scale: float = 0.005 + alive_reward_scale: float = 0.5 + dof_vel_scale: float = 0.2 + + death_cost: float = -2.0 + termination_height: float = 0.31 + + angular_velocity_scale: float = 1.0 + contact_force_scale: float = 0.1 + + +class AntWarpEnv(LocomotionWarpEnv): + cfg: AntWarpEnvCfg + + def __init__(self, cfg: AntWarpEnvCfg, render_mode: str | None = None, **kwargs): + super().__init__(cfg, render_mode, **kwargs) diff --git a/source/isaaclab_tasks/isaaclab_tasks/direct/cartpole/__init__.py b/source/isaaclab_tasks/isaaclab_tasks/direct/cartpole/__init__.py index 7b2b689c7de..46dc807b1dd 100644 --- a/source/isaaclab_tasks/isaaclab_tasks/direct/cartpole/__init__.py +++ b/source/isaaclab_tasks/isaaclab_tasks/direct/cartpole/__init__.py @@ -28,6 +28,19 @@ }, ) +gym.register( + id="Isaac-Cartpole-Warp-v0", + entry_point=f"{__name__}.cartpole_warp_env:CartpoleWarpEnv", + disable_env_checker=True, + kwargs={ + "env_cfg_entry_point": f"{__name__}.cartpole_warp_env:CartpoleWarpEnvCfg", + "rl_games_cfg_entry_point": f"{agents.__name__}:rl_games_ppo_cfg.yaml", + "rsl_rl_cfg_entry_point": f"{agents.__name__}.rsl_rl_ppo_cfg:CartpolePPORunnerCfg", + "skrl_cfg_entry_point": f"{agents.__name__}:skrl_ppo_cfg.yaml", + "sb3_cfg_entry_point": f"{agents.__name__}:sb3_ppo_cfg.yaml", + }, +) + gym.register( id="Isaac-Cartpole-RGB-Camera-Direct-v0", entry_point=f"{__name__}.cartpole_camera_env:CartpoleCameraEnv", diff --git a/source/isaaclab_tasks/isaaclab_tasks/direct/cartpole/cartpole_env.py b/source/isaaclab_tasks/isaaclab_tasks/direct/cartpole/cartpole_env.py index 911622188ce..85b06d81a92 100644 --- a/source/isaaclab_tasks/isaaclab_tasks/direct/cartpole/cartpole_env.py +++ b/source/isaaclab_tasks/isaaclab_tasks/direct/cartpole/cartpole_env.py @@ -35,6 +35,7 @@ class CartpoleEnvCfg(DirectRLEnvCfg): solver_cfg = MJWarpSolverCfg( njmax=5, + ncon_per_env=3, ls_iterations=3, cone="pyramidal", impratio=1, @@ -44,6 +45,7 @@ class CartpoleEnvCfg(DirectRLEnvCfg): solver_cfg=solver_cfg, num_substeps=1, debug_mode=False, + use_cuda_graph=True, ) # simulation @@ -77,10 +79,12 @@ class CartpoleEnv(DirectRLEnv): def __init__(self, cfg: CartpoleEnvCfg, render_mode: str | None = None, **kwargs): super().__init__(cfg, render_mode, **kwargs) - self._cart_dof_idx, _ = self.cartpole.find_joints(self.cfg.cart_dof_name) - self._pole_dof_idx, _ = self.cartpole.find_joints(self.cfg.pole_dof_name) + _, _, self._cart_dof_idx = self.cartpole.find_joints(self.cfg.cart_dof_name) + _, _, self._pole_dof_idx = self.cartpole.find_joints(self.cfg.pole_dof_name) self.action_scale = self.cfg.action_scale + # Memory views into the Newton simulation data, they should not be modified directly + # They are updated automatically by the simulation, no need to copy or fetch them at each step. self.joint_pos = self.cartpole.data.joint_pos self.joint_vel = self.cartpole.data.joint_vel @@ -115,7 +119,7 @@ def _get_observations(self) -> dict: ), dim=-1, ) - observations = {"policy": obs} + observations = {"policy": obs.clone()} return observations def _get_rewards(self) -> torch.Tensor: @@ -134,9 +138,6 @@ def _get_rewards(self) -> torch.Tensor: return total_reward def _get_dones(self) -> tuple[torch.Tensor, torch.Tensor]: - self.joint_pos = self.cartpole.data.joint_pos - self.joint_vel = self.cartpole.data.joint_vel - time_out = self.episode_length_buf >= self.max_episode_length - 1 out_of_bounds = torch.any(torch.abs(self.joint_pos[:, self._cart_dof_idx]) > self.cfg.max_cart_pos, dim=1) out_of_bounds = out_of_bounds | torch.any(torch.abs(self.joint_pos[:, self._pole_dof_idx]) > math.pi / 2, dim=1) @@ -156,14 +157,6 @@ def _reset_idx(self, env_ids: Sequence[int] | None): ) joint_vel = self.cartpole.data.default_joint_vel[env_ids] - default_root_state = self.cartpole.data.default_root_state[env_ids] - default_root_state[:, :3] += self.scene.env_origins[env_ids] - - self.joint_pos[env_ids] = joint_pos - self.joint_vel[env_ids] = joint_vel - - self.cartpole.write_root_pose_to_sim(default_root_state[:, :7], env_ids) - self.cartpole.write_root_velocity_to_sim(default_root_state[:, 7:], env_ids) self.cartpole.write_joint_state_to_sim(joint_pos, joint_vel, None, env_ids) diff --git a/source/isaaclab_tasks/isaaclab_tasks/direct/cartpole/cartpole_warp_env.py b/source/isaaclab_tasks/isaaclab_tasks/direct/cartpole/cartpole_warp_env.py new file mode 100644 index 00000000000..8d604e2fec1 --- /dev/null +++ b/source/isaaclab_tasks/isaaclab_tasks/direct/cartpole/cartpole_warp_env.py @@ -0,0 +1,356 @@ +# Copyright (c) 2022-2025, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +from __future__ import annotations + +import math + +import warp as wp + +from isaaclab_assets.robots.cartpole import CARTPOLE_CFG + +import isaaclab.sim as sim_utils +from isaaclab.assets import Articulation, ArticulationCfg +from isaaclab.newton import Frontend +from isaaclab.envs import DirectRLEnvCfg, DirectRLEnvWarp +from isaaclab.scene import InteractiveSceneCfg +from isaaclab.sim import SimulationCfg +from isaaclab.sim._impl.newton_manager_cfg import NewtonCfg +from isaaclab.sim._impl.solvers_cfg import MJWarpSolverCfg +from isaaclab.sim.spawners.from_files import GroundPlaneCfg, spawn_ground_plane +from isaaclab.utils import configclass + + +@configclass +class CartpoleWarpEnvCfg(DirectRLEnvCfg): + # env + decimation = 2 + episode_length_s = 5.0 + action_scale = 100.0 # [N] + action_space = 1 + observation_space = 4 + state_space = 0 + + solver_cfg = MJWarpSolverCfg( + njmax=5, + ncon_per_env=3, + ls_iterations=3, + cone="pyramidal", + impratio=1, + update_data_interval=1, + ) + + newton_cfg = NewtonCfg( + solver_cfg=solver_cfg, + num_substeps=1, + debug_mode=False, + use_cuda_graph=True, + ) + + # simulation + sim: SimulationCfg = SimulationCfg(dt=1 / 120, render_interval=decimation, newton_cfg=newton_cfg) + + # robot + robot_cfg: ArticulationCfg = CARTPOLE_CFG.replace(prim_path="/World/envs/env_.*/Robot") + cart_dof_name = "slider_to_cart" + pole_dof_name = "cart_to_pole" + + # scene + scene: InteractiveSceneCfg = InteractiveSceneCfg( + num_envs=4096, env_spacing=4.0, replicate_physics=True, clone_in_fabric=True + ) + + # reset + max_cart_pos = 3.0 # the cart is reset if it exceeds that position [m] + initial_pole_angle_range = [-0.25, 0.25] # the range in which the pole angle is sampled from on reset [rad] + + # reward scales + rew_scale_alive = 1.0 + rew_scale_terminated = -2.0 + rew_scale_pole_pos = -1.0 + rew_scale_cart_vel = -0.01 + rew_scale_pole_vel = -0.005 + + +@wp.kernel +def get_observations( + joint_pos: wp.array2d(dtype=wp.float32), + joint_vel: wp.array2d(dtype=wp.float32), + cart_dof_idx: wp.int32, + pole_dof_idx: wp.int32, + observations: wp.array(dtype=wp.vec4f), +): + env_index = wp.tid() + observations[env_index][0] = joint_pos[env_index, pole_dof_idx] + observations[env_index][1] = joint_vel[env_index, pole_dof_idx] + observations[env_index][2] = joint_pos[env_index, cart_dof_idx] + observations[env_index][3] = joint_vel[env_index, cart_dof_idx] + + +@wp.kernel +def update_actions( + input_actions: wp.array2d(dtype=wp.float32), + actions: wp.array2d(dtype=wp.float32), + action_scale: wp.float32, +): + env_index = wp.tid() + actions[env_index, 0] = action_scale * input_actions[env_index, 0] + + +@wp.kernel +def get_dones( + joint_pos: wp.array2d(dtype=wp.float32), + episode_length_buf: wp.array(dtype=wp.int32), + cart_dof_idx: wp.int32, + pole_dof_idx: wp.int32, + max_episode_length: wp.int32, + max_cart_pos: wp.float32, + out_of_bounds: wp.array(dtype=wp.bool), + time_out: wp.array(dtype=wp.bool), + reset: wp.array(dtype=wp.bool), +): + env_index = wp.tid() + out_of_bounds[env_index] = (wp.abs(joint_pos[env_index, cart_dof_idx]) > max_cart_pos) or ( + wp.abs(joint_pos[env_index, pole_dof_idx]) > math.pi / 2.0 + ) + time_out[env_index] = episode_length_buf[env_index] >= (max_episode_length - 1) + reset[env_index] = out_of_bounds[env_index] or time_out[env_index] + + +@wp.func +def compute_rew_alive(rew_scale_alive: wp.float32, reset_terminated: bool) -> wp.float32: + if reset_terminated: + return wp.float32(0.0) + return rew_scale_alive + + +@wp.func +def compute_rew_termination(rew_scale_terminated: wp.float32, reset_terminated: bool) -> wp.float32: + if reset_terminated: + return rew_scale_terminated + return wp.float32(0.0) + + +@wp.func +def compute_rew_pole_pos( + rew_scale_pole_pos: wp.float32, + pole_pos: wp.float32, +) -> wp.float32: + return rew_scale_pole_pos * pole_pos * pole_pos + + +@wp.func +def compute_rew_cart_vel( + rew_scale_cart_vel: wp.float32, + cart_vel: wp.float32, +) -> wp.float32: + return rew_scale_cart_vel * wp.abs(cart_vel) + + +@wp.func +def compute_rew_pole_vel( + rew_scale_pole_vel: wp.float32, + pole_vel: wp.float32, +) -> wp.float32: + return rew_scale_pole_vel * wp.abs(pole_vel) + + +@wp.kernel +def compute_rewards( + rew_scale_alive: wp.float32, + rew_scale_terminated: wp.float32, + rew_scale_pole_pos: wp.float32, + rew_scale_cart_vel: wp.float32, + rew_scale_pole_vel: wp.float32, + joint_pos: wp.array2d(dtype=wp.float32), + joint_vel: wp.array2d(dtype=wp.float32), + cart_dof_idx: wp.int32, + pole_dof_idx: wp.int32, + reset_terminated: wp.array(dtype=wp.bool), + reward: wp.array(dtype=wp.float32), +): + env_index = wp.tid() + reward[env_index] = ( + compute_rew_alive(rew_scale_alive, reset_terminated[env_index]) + + compute_rew_termination(rew_scale_terminated, reset_terminated[env_index]) + + compute_rew_pole_pos(rew_scale_pole_pos, joint_pos[env_index, pole_dof_idx]) + + compute_rew_cart_vel(rew_scale_cart_vel, joint_vel[env_index, cart_dof_idx]) + + compute_rew_pole_vel(rew_scale_pole_vel, joint_vel[env_index, pole_dof_idx]) + ) + + +@wp.kernel +def reset( + default_joint_pos: wp.array2d(dtype=wp.float32), + default_joint_vel: wp.array2d(dtype=wp.float32), + joint_pos: wp.array2d(dtype=wp.float32), + joint_vel: wp.array2d(dtype=wp.float32), + cart_dof_idx: wp.int32, + pole_dof_idx: wp.int32, + initial_pose_angle_range: wp.vec2f, + env_mask: wp.array(dtype=wp.bool), + state: wp.array(dtype=wp.uint32), +): + env_index = wp.tid() + if env_mask[env_index]: + joint_pos[env_index, cart_dof_idx] = default_joint_pos[env_index, cart_dof_idx] + joint_pos[env_index, pole_dof_idx] = default_joint_pos[env_index, pole_dof_idx] + wp.randf( + state[env_index], initial_pose_angle_range[0] * wp.pi, initial_pose_angle_range[1] * wp.pi + ) + joint_vel[env_index, 0] = default_joint_vel[env_index, 0] + joint_vel[env_index, 1] = default_joint_vel[env_index, 1] + state[env_index] += wp.uint32(1) + + +@wp.kernel +def initialize_state( + state: wp.array(dtype=wp.uint32), + seed: wp.int32, +): + env_index = wp.tid() + state[env_index] = wp.rand_init(seed, env_index) + + +class CartpoleWarpEnv(DirectRLEnvWarp): + cfg: CartpoleWarpEnvCfg + + def __init__(self, cfg: CartpoleWarpEnvCfg, render_mode: str | None = None, **kwargs) -> None: + super().__init__(cfg, render_mode, **kwargs) + + # Get the indices + self._cart_dof_mask, _, self._cart_dof_idx = self.cartpole.find_joints(self.cfg.cart_dof_name) + self._pole_dof_mask, _, self._pole_dof_idx = self.cartpole.find_joints(self.cfg.pole_dof_name) + + self.action_scale = self.cfg.action_scale + + # Simulation bindings + # Note: these are direct memory views into the Newton simulation data, they should not be modified directly + self.joint_pos = self.cartpole.data.joint_pos + self.joint_vel = self.cartpole.data.joint_vel + + # Buffers + self.observations = wp.zeros((self.num_envs), dtype=wp.vec4f, device=self.device) + self.actions = wp.zeros((self.num_envs, 1), dtype=wp.float32, device=self.device) + self.rewards = wp.zeros((self.num_envs), dtype=wp.float32, device=self.device) + self.states = wp.zeros((self.num_envs), dtype=wp.uint32, device=self.device) + + if self.cfg.seed is None: + self.cfg.seed = -1 + + wp.launch( + initialize_state, + dim=self.num_envs, + inputs=[ + self.states, + self.cfg.seed, + ], + ) + + # Bind torch buffers to warp buffers + self.torch_obs_buf = wp.to_torch(self.observations) + self.torch_reward_buf = wp.to_torch(self.rewards) + self.torch_reset_terminated = wp.to_torch(self.reset_terminated) + self.torch_reset_time_outs = wp.to_torch(self.reset_time_outs) + self.torch_episode_length_buf = wp.to_torch(self.episode_length_buf) + + def _setup_scene(self) -> None: + self.cartpole = Articulation(self.cfg.robot_cfg, frontend=Frontend.WARP) + # add ground plane + spawn_ground_plane(prim_path="/World/ground", cfg=GroundPlaneCfg()) + # clone and replicate + self.scene.clone_environments(copy_from_source=False) + # we need to explicitly filter collisions for CPU simulation + if self.device == "cpu": + self.scene.filter_collisions(global_prim_paths=[]) + # add articulation to scene + self.scene.articulations["cartpole"] = self.cartpole + # add lights + light_cfg = sim_utils.DomeLightCfg(intensity=2000.0, color=(0.75, 0.75, 0.75)) + light_cfg.func("/World/Light", light_cfg) + + def _pre_physics_step(self, actions: wp.array) -> None: + wp.launch( + update_actions, + dim=self.num_envs, + inputs=[ + actions, + self.actions, + self.action_scale, + ], + ) + + def _apply_action(self) -> None: + self.cartpole.set_joint_effort_target(self.actions, joint_mask=self._cart_dof_mask) + + def _get_observations(self) -> None: + wp.launch( + get_observations, + dim=self.num_envs, + inputs=[ + self.joint_pos, + self.joint_vel, + self._cart_dof_idx[0], + self._pole_dof_idx[0], + self.observations, + ], + ) + + def _get_rewards(self) -> None: + wp.launch( + compute_rewards, + dim=self.num_envs, + inputs=[ + self.cfg.rew_scale_alive, + self.cfg.rew_scale_terminated, + self.cfg.rew_scale_pole_pos, + self.cfg.rew_scale_cart_vel, + self.cfg.rew_scale_pole_vel, + self.joint_pos, + self.joint_vel, + self._cart_dof_idx[0], + self._pole_dof_idx[0], + self.reset_terminated, + self.rewards, + ], + ) + + def _get_dones(self) -> None: + wp.launch( + get_dones, + dim=self.num_envs, + inputs=[ + self.joint_pos, + self.episode_length_buf, + self._cart_dof_idx[0], + self._pole_dof_idx[0], + self.max_episode_length, + self.cfg.max_cart_pos, + self.reset_terminated, + self.reset_time_outs, + self.reset_buf, + ], + ) + + def _reset_idx(self, mask: wp.array | None = None) -> None: + if mask is None: + mask = self.cartpole._ALL_ENV_MASK + + super()._reset_idx(mask) + + wp.launch( + reset, + dim=self.num_envs, + inputs=[ + self.cartpole.data.default_joint_pos, + self.cartpole.data.default_joint_vel, + self.joint_pos, + self.joint_vel, + self._cart_dof_idx[0], + self._pole_dof_idx[0], + self.cfg.initial_pole_angle_range, + mask, + self.states, + ], + ) diff --git a/source/isaaclab_tasks/isaaclab_tasks/direct/humanoid/__init__.py b/source/isaaclab_tasks/isaaclab_tasks/direct/humanoid/__init__.py index ff38052a5cd..8275c3f49eb 100644 --- a/source/isaaclab_tasks/isaaclab_tasks/direct/humanoid/__init__.py +++ b/source/isaaclab_tasks/isaaclab_tasks/direct/humanoid/__init__.py @@ -26,3 +26,15 @@ "skrl_cfg_entry_point": f"{agents.__name__}:skrl_ppo_cfg.yaml", }, ) + +gym.register( + id="Isaac-Humanoid-Warp-v0", + entry_point=f"{__name__}.humanoid_warp_env:HumanoidWarpEnv", + disable_env_checker=True, + kwargs={ + "env_cfg_entry_point": f"{__name__}.humanoid_warp_env:HumanoidWarpEnvCfg", + "rl_games_cfg_entry_point": f"{agents.__name__}:rl_games_ppo_cfg.yaml", + "rsl_rl_cfg_entry_point": f"{agents.__name__}.rsl_rl_ppo_cfg:HumanoidPPORunnerCfg", + "skrl_cfg_entry_point": f"{agents.__name__}:skrl_ppo_cfg.yaml", + }, +) diff --git a/source/isaaclab_tasks/isaaclab_tasks/direct/humanoid/humanoid_warp_env.py b/source/isaaclab_tasks/isaaclab_tasks/direct/humanoid/humanoid_warp_env.py new file mode 100644 index 00000000000..ef7f5c8addf --- /dev/null +++ b/source/isaaclab_tasks/isaaclab_tasks/direct/humanoid/humanoid_warp_env.py @@ -0,0 +1,114 @@ +# Copyright (c) 2022-2025, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +from __future__ import annotations + +from isaaclab_assets import HUMANOID_CFG + +import isaaclab.sim as sim_utils +from isaaclab.assets import ArticulationCfg +from isaaclab.envs import DirectRLEnvCfg +from isaaclab.scene import InteractiveSceneCfg +from isaaclab.sim import SimulationCfg +from isaaclab.sim._impl.newton_manager_cfg import NewtonCfg +from isaaclab.sim._impl.solvers_cfg import MJWarpSolverCfg +from isaaclab.terrains import TerrainImporterCfg +from isaaclab.utils import configclass + +from isaaclab_tasks.direct.locomotion.locomotion_env_warp import LocomotionWarpEnv + + +@configclass +class HumanoidWarpEnvCfg(DirectRLEnvCfg): + # env + episode_length_s = 15.0 + decimation = 2 + action_scale = 1.0 + action_space = 21 + observation_space = 75 + state_space = 0 + + solver_cfg = MJWarpSolverCfg( + njmax=80, + ncon_per_env=25, + ls_iterations=15, + ls_parallel=True, + cone="pyramidal", + update_data_interval=2, + impratio=1, + ) + newton_cfg = NewtonCfg( + solver_cfg=solver_cfg, + num_substeps=2, + debug_mode=False, + ) + + # simulation + sim: SimulationCfg = SimulationCfg(dt=1 / 120, render_interval=decimation, newton_cfg=newton_cfg) + terrain = TerrainImporterCfg( + prim_path="/World/ground", + terrain_type="plane", + collision_group=-1, + physics_material=sim_utils.RigidBodyMaterialCfg( + friction_combine_mode="average", + restitution_combine_mode="average", + static_friction=1.0, + dynamic_friction=1.0, + restitution=0.0, + ), + debug_vis=False, + ) + + # scene + scene: InteractiveSceneCfg = InteractiveSceneCfg( + num_envs=4096, env_spacing=4.0, replicate_physics=True, clone_in_fabric=True + ) + + # robot + robot: ArticulationCfg = HUMANOID_CFG.replace(prim_path="/World/envs/env_.*/Robot") + joint_gears: list = [ + 67.5000, # left_upper_arm + 67.5000, # left_upper_arm + 45.0000, # left_lower_arm + 67.5000, # lower_waist + 67.5000, # lower_waist + 67.5000, # pelvis + 45.0000, # left_thigh: x + 135.0000, # left_thigh: y + 45.0000, # left_thigh: z + 90.0000, # left_knee + 22.5, # left_foot + 22.5, # left_foot + 45.0000, # right_thigh: x + 135.0000, # right_thigh: y + 45.0000, # right_thigh: z + 90.0000, # right_knee + 22.5, # right_foot + 22.5, # right_foot + 67.5000, # right_upper_arm + 67.5000, # right_upper_arm + 45.0000, # right_lower_arm + ] + + heading_weight: float = 0.5 + up_weight: float = 0.1 + + energy_cost_scale: float = 0.05 + actions_cost_scale: float = 0.01 + alive_reward_scale: float = 2.0 + dof_vel_scale: float = 0.1 + + death_cost: float = -1.0 + termination_height: float = 0.8 + + angular_velocity_scale: float = 0.25 + contact_force_scale: float = 0.01 + + +class HumanoidWarpEnv(LocomotionWarpEnv): + cfg: HumanoidWarpEnvCfg + + def __init__(self, cfg: HumanoidWarpEnvCfg, render_mode: str | None = None, **kwargs): + super().__init__(cfg, render_mode, **kwargs) diff --git a/source/isaaclab_tasks/isaaclab_tasks/direct/locomotion/locomotion_env.py b/source/isaaclab_tasks/isaaclab_tasks/direct/locomotion/locomotion_env.py index 59689242895..30d5b5a7c4b 100644 --- a/source/isaaclab_tasks/isaaclab_tasks/direct/locomotion/locomotion_env.py +++ b/source/isaaclab_tasks/isaaclab_tasks/direct/locomotion/locomotion_env.py @@ -28,7 +28,7 @@ def __init__(self, cfg: DirectRLEnvCfg, render_mode: str | None = None, **kwargs self.action_scale = self.cfg.action_scale self.joint_gears = torch.tensor(self.cfg.joint_gears, dtype=torch.float32, device=self.sim.device) self.motor_effort_ratio = torch.ones_like(self.joint_gears, device=self.sim.device) - self._joint_dof_idx, _ = self.robot.find_joints(".*") + _, _,self._joint_dof_idx = self.robot.find_joints(".*") self.potentials = torch.zeros(self.num_envs, dtype=torch.float32, device=self.sim.device) self.prev_potentials = torch.zeros_like(self.potentials) @@ -45,6 +45,22 @@ def __init__(self, cfg: DirectRLEnvCfg, render_mode: str | None = None, **kwargs self.basis_vec0 = self.heading_vec.clone() self.basis_vec1 = self.up_vec.clone() + # Get the values from the robot data + print("joint_pos_limits_lower", self.robot.data.joint_pos_limits[0, :, 0]) + print("joint_pos_limits_upper", self.robot.data.joint_pos_limits[0, :, 1]) + print("joint_stiffness_sim", self.robot.data.joint_stiffness[0]) + print("joint_damping_sim", self.robot.data.joint_damping[0]) + print("joint_armature", self.robot.data.joint_armature[0]) + print("joint_friction_coeff", self.robot.data.joint_friction_coeff[0]) + print("joint_vel_limits_sim", self.robot.data.joint_vel_limits[0]) + print("joint_effort_limits_sim", self.robot.data.joint_effort_limits[0]) + print("joint_control_mode_sim", self.robot.data.joint_control_mode[0]) + print("joint_pos", self.robot.data.joint_pos[0]) + print("joint_vel", self.robot.data.joint_vel[0]) + print("joint_effort", self.robot.data.joint_effort_target[0]) + print("joint_target", self.robot.data.joint_target[0]) + print("soft_joint_pos_limits", self.robot.data.soft_joint_pos_limits[0]) + def _setup_scene(self): self.robot = Articulation(self.cfg.robot) # add ground plane diff --git a/source/isaaclab_tasks/isaaclab_tasks/direct/locomotion/locomotion_env_warp.py b/source/isaaclab_tasks/isaaclab_tasks/direct/locomotion/locomotion_env_warp.py new file mode 100644 index 00000000000..3c70d33631a --- /dev/null +++ b/source/isaaclab_tasks/isaaclab_tasks/direct/locomotion/locomotion_env_warp.py @@ -0,0 +1,583 @@ +# Copyright (c) 2022-2025, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +from __future__ import annotations + +import warp as wp + +import isaaclab.sim as sim_utils +from isaaclab.assets import Articulation +from isaaclab.envs import DirectRLEnvCfg, DirectRLEnvWarp +from isaaclab.newton import Frontend + +@wp.func +def fmod(x: wp.float32, y: wp.float32) -> wp.float32: + return x - y * wp.floor(x / y) + + +@wp.func +def euler_from_quat(q: wp.quatf) -> wp.vec3f: + sinr_cosp = 2.0 * (q[3] * q[0] + q[1] * q[2]) + cosr_cosp = q[3] * q[3] - q[0] * q[0] - q[1] * q[1] + q[2] * q[2] + sinp = 2.0 * (q[3] * q[1] - q[2] * q[0]) + siny_cosp = 2.0 * (q[3] * q[2] + q[0] * q[1]) + cosy_cosp = q[3] * q[3] + q[0] * q[0] - q[1] * q[1] - q[2] * q[2] + roll = wp.atan2(sinr_cosp, cosr_cosp) + if wp.abs(sinp) >= 1: + pitch = wp.sign(sinp) * wp.pi / 2.0 + else: + pitch = wp.asin(sinp) + yaw = wp.atan2(siny_cosp, cosy_cosp) + + return wp.vec3f( + fmod(roll, wp.static(2.0 * wp.pi)), + fmod(pitch, wp.static(2.0 * wp.pi)), + fmod(yaw, wp.static(2.0 * wp.pi)), + ) + + +@wp.kernel +def get_dones( + episode_length_buf: wp.array(dtype=wp.int32), + torso_pose: wp.array(dtype=wp.transformf), + max_episode_length: wp.int32, + termination_height: wp.float32, + out_of_bounds: wp.array(dtype=wp.bool), + time_out: wp.array(dtype=wp.bool), + reset: wp.array(dtype=wp.bool), +): + env_index = wp.tid() + out_of_bounds[env_index] = wp.abs(torso_pose[env_index][2]) < termination_height + time_out[env_index] = episode_length_buf[env_index] >= (max_episode_length - 1) + reset[env_index] = out_of_bounds[env_index] or time_out[env_index] + + +@wp.kernel +def observations( + torso_pose: wp.array(dtype=wp.transformf), + velocity: wp.array(dtype=wp.spatial_vectorf), + rpy: wp.array(dtype=wp.vec3f), + angle_to_target: wp.array(dtype=wp.float32), + up_proj: wp.array(dtype=wp.float32), + heading_proj: wp.array(dtype=wp.float32), + dof_pos_scaled: wp.array2d(dtype=wp.float32), + dof_vel: wp.array2d(dtype=wp.float32), + actions: wp.array2d(dtype=wp.float32), + observations: wp.array2d(dtype=wp.float32), + dof_vel_scale: wp.float32, + angular_velocity_scale: wp.float32, + num_dof: wp.int32, +): + env_index = wp.tid() + observations[env_index, 0] = torso_pose[env_index][2] + observations[env_index, 1] = velocity[env_index][0] + observations[env_index, 2] = velocity[env_index][1] + observations[env_index, 3] = velocity[env_index][2] + observations[env_index, 4] = velocity[env_index][3] * angular_velocity_scale + observations[env_index, 5] = velocity[env_index][4] * angular_velocity_scale + observations[env_index, 6] = velocity[env_index][5] * angular_velocity_scale + observations[env_index, 7] = wp.atan2(wp.sin(rpy[env_index][2]), wp.cos(rpy[env_index][2])) + observations[env_index, 8] = wp.atan2(wp.sin(rpy[env_index][0]), wp.cos(rpy[env_index][0])) + observations[env_index, 9] = wp.atan2(wp.sin(angle_to_target[env_index]), wp.cos(angle_to_target[env_index])) + observations[env_index, 10] = up_proj[env_index] + observations[env_index, 11] = heading_proj[env_index] + + offset_1 = 12 + num_dof + offset_2 = offset_1 + num_dof + + for i in range(num_dof): + observations[env_index, 12 + i] = dof_pos_scaled[env_index, i] + for i in range(num_dof): + observations[env_index, offset_1 + i] = dof_vel[env_index, i] * dof_vel_scale + for i in range(num_dof): + observations[env_index, offset_2 + i] = actions[env_index, i] + + +@wp.func +def translate_transform( + transform: wp.transformf, + translation: wp.vec3f, +) -> wp.transformf: + return wp.transform( + wp.transform_get_translation(transform) + translation, + wp.transform_get_rotation(transform), + ) + + +@wp.kernel +def reset_root( + default_root_pose: wp.array(dtype=wp.transformf), + default_root_vel: wp.array(dtype=wp.spatial_vectorf), + env_origins: wp.array(dtype=wp.vec3f), + dt: wp.float32, + to_targets: wp.array(dtype=wp.vec3f), + potentials: wp.array(dtype=wp.float32), + root_pose: wp.array(dtype=wp.transformf), + root_vel: wp.array(dtype=wp.spatial_vectorf), + env_mask: wp.array(dtype=wp.bool), +): + env_index = wp.tid() + if env_mask[env_index]: + root_pose[env_index] = default_root_pose[env_index] + root_pose[env_index] = translate_transform(root_pose[env_index], env_origins[env_index]) + root_vel[env_index] = default_root_vel[env_index] + to_targets[env_index] = wp.transform_get_translation(root_pose[env_index]) - wp.transform_get_translation( + default_root_pose[env_index] + ) + to_targets[env_index][2] = 0.0 + potentials[env_index] = -wp.length(to_targets[env_index]) / dt + + +@wp.kernel +def reset_joints( + default_joint_pos: wp.array2d(dtype=wp.float32), + default_joint_vel: wp.array2d(dtype=wp.float32), + joint_pos: wp.array2d(dtype=wp.float32), + joint_vel: wp.array2d(dtype=wp.float32), + env_mask: wp.array(dtype=wp.bool), +): + env_index, joint_index = wp.tid() + if env_mask[env_index]: + joint_pos[env_index, joint_index] = default_joint_pos[env_index, joint_index] + joint_vel[env_index, joint_index] = default_joint_vel[env_index, joint_index] + + +@wp.func +def heading_reward( + heading_proj: wp.float32, + heading_weight: wp.float32, +) -> wp.float32: + if heading_proj > 0.8: + return heading_weight + else: + return heading_weight * heading_proj / 0.8 + + +@wp.func +def up_reward( + up_proj: wp.float32, + up_weight: wp.float32, +) -> wp.float32: + + if up_proj > 0.93: + return up_weight + else: + return 0.0 + + +@wp.func +def progress_reward( + current_value: wp.float32, + prev_value: wp.float32, +) -> wp.float32: + return current_value - prev_value + + +@wp.func +def actions_cost( + actions: wp.array(dtype=wp.float32), +) -> wp.float32: + sum_ = wp.float32(0.0) + for i in range(len(actions)): + sum_ += actions[i] * actions[i] + return sum_ + + +@wp.func +def electricity_cost( + actions: wp.array(dtype=wp.float32), + dof_vel: wp.array(dtype=wp.float32), + dof_vel_scale: wp.float32, + motor_effort_ratio: wp.array(dtype=wp.float32), +) -> wp.float32: + sum_ = wp.float32(0.0) + for i in range(len(actions)): + sum_ += wp.abs(actions[i] * dof_vel[i] * dof_vel_scale) * motor_effort_ratio[i] + return sum_ + + +@wp.func +def dof_at_limit_cost( + dof_pos_scaled: wp.array(dtype=wp.float32), +) -> wp.float32: + sum_ = wp.float32(0.0) + for i in range(len(dof_pos_scaled)): + if dof_pos_scaled[i] > 0.98: + sum_ += 1.0 + return sum_ + + +@wp.kernel +def compute_rewards( + actions: wp.array2d(dtype=wp.float32), + dof_vel: wp.array2d(dtype=wp.float32), + dof_pos_scaled: wp.array2d(dtype=wp.float32), + reset_terminated: wp.array(dtype=wp.bool), + heading_proj: wp.array(dtype=wp.float32), + up_proj: wp.array(dtype=wp.float32), + potentials: wp.array(dtype=wp.float32), + prev_potentials: wp.array(dtype=wp.float32), + motor_effort_ratio: wp.array(dtype=wp.float32), + up_weight: wp.float32, + heading_weight: wp.float32, + actions_cost_scale: wp.float32, + energy_cost_scale: wp.float32, + dof_vel_scale: wp.float32, + death_cost: wp.float32, + alive_reward_scale: wp.float32, + reward: wp.array(dtype=wp.float32), +): + env_index = wp.tid() + if reset_terminated[env_index]: + reward[env_index] = death_cost + else: + reward[env_index] = ( + progress_reward(potentials[env_index], prev_potentials[env_index]) + + alive_reward_scale + + up_reward(up_proj[env_index], up_weight) + + heading_reward(heading_proj[env_index], heading_weight) + - actions_cost_scale * actions_cost(actions[env_index]) + - energy_cost_scale + * electricity_cost(actions[env_index], dof_vel[env_index], dof_vel_scale, motor_effort_ratio) + - dof_at_limit_cost(dof_pos_scaled[env_index]) + ) + + +@wp.kernel +def compute_heading_and_up( + torso_pose: wp.array(dtype=wp.transformf), + targets: wp.array(dtype=wp.vec3f), + dt: wp.float32, + to_targets: wp.array(dtype=wp.vec3f), + up_proj: wp.array(dtype=wp.float32), + heading_proj: wp.array(dtype=wp.float32), + up_vec: wp.array(dtype=wp.vec3f), + heading_vec: wp.array(dtype=wp.vec3f), + potentials: wp.array(dtype=wp.float32), + prev_potentials: wp.array(dtype=wp.float32), +): + env_index = wp.tid() + up_vec[env_index] = wp.quat_rotate(wp.transform_get_rotation(torso_pose[env_index]), wp.static(wp.vec3f(0, 0, 1))) + heading_vec[env_index] = wp.quat_rotate( + wp.transform_get_rotation(torso_pose[env_index]), wp.static(wp.vec3f(1, 0, 0)) + ) + up_proj[env_index] = up_vec[env_index][2] + to_targets[env_index] = targets[env_index] - wp.transform_get_translation(torso_pose[env_index]) + to_targets[env_index][2] = 0.0 + heading_proj[env_index] = wp.dot(heading_vec[env_index], wp.normalize(to_targets[env_index])) + prev_potentials[env_index] = potentials[env_index] + potentials[env_index] = -wp.length(to_targets[env_index]) / dt + + +@wp.func +def spatial_rotate_inv(quat: wp.quatf, vec: wp.spatial_vectorf) -> wp.spatial_vectorf: + return wp.spatial_vector( + wp.quat_rotate_inv(quat, wp.spatial_top(vec)), + wp.quat_rotate_inv(quat, wp.spatial_bottom(vec)), + ) + + +@wp.func +def unscale(x: wp.float32, lower: wp.float32, upper: wp.float32) -> wp.float32: + return (2.0 * x - upper - lower) / (upper - lower) + + +@wp.kernel +def compute_rot( + torso_pose: wp.array(dtype=wp.transformf), + velocity: wp.array(dtype=wp.spatial_vectorf), + targets: wp.array(dtype=wp.vec3f), + vec_loc: wp.array(dtype=wp.spatial_vectorf), + rpy: wp.array(dtype=wp.vec3f), + angle_to_target: wp.array(dtype=wp.float32), +): + env_index = wp.tid() + vec_loc[env_index] = spatial_rotate_inv(wp.transform_get_rotation(torso_pose[env_index]), velocity[env_index]) + rpy[env_index] = euler_from_quat(wp.transform_get_rotation(torso_pose[env_index])) + angle_to_target[env_index] = ( + wp.atan2(targets[env_index][2] - torso_pose[env_index][2], targets[env_index][0] - torso_pose[env_index][0]) + - rpy[env_index][2] + ) + + +@wp.kernel +def scale_dof_pos( + dof_pos: wp.array2d(dtype=wp.float32), + dof_limits: wp.array2d(dtype=wp.vec2f), + dof_pos_scaled: wp.array2d(dtype=wp.float32), +): + env_index, joint_index = wp.tid() + dof_pos_scaled[env_index, joint_index] = unscale( + dof_pos[env_index, joint_index], dof_limits[env_index, joint_index][0], dof_limits[env_index, joint_index][1] + ) + + +@wp.kernel +def update_actions( + input_actions: wp.array2d(dtype=wp.float32), + actions: wp.array2d(dtype=wp.float32), + joint_gears: wp.array(dtype=wp.float32), + action_scale: wp.float32, +): + env_index, joint_index = wp.tid() + actions[env_index, joint_index] = ( + action_scale * joint_gears[joint_index] * wp.clamp(input_actions[env_index, joint_index], -1.0, 1.0) + ) + + +@wp.kernel +def initialize_state( + env_origins: wp.array(dtype=wp.vec3f), + targets: wp.array(dtype=wp.vec3f), + state: wp.array(dtype=wp.uint32), + seed: wp.int32, +): + env_index = wp.tid() + state[env_index] = wp.rand_init(seed, env_index) + targets[env_index] = env_origins[env_index] + targets[env_index] += wp.static(wp.vec3f(1000.0, 0.0, 0.0)) + + +class LocomotionWarpEnv(DirectRLEnvWarp): + cfg: DirectRLEnvCfg + + def __init__(self, cfg: DirectRLEnvCfg, render_mode: str | None = None, **kwargs): + super().__init__(cfg, render_mode, **kwargs) + + self.action_scale = self.cfg.action_scale + self.joint_gears = wp.array(self.cfg.joint_gears, dtype=wp.float32, device=self.sim.device) + self.motor_effort_ratio = wp.ones_like(self.joint_gears, device=self.sim.device) + self._joint_dof_mask, _, self._joint_dof_idx = self.robot.find_joints(".*") + + # Simulation bindings + # Note: these are direct memory views into the Newton simulation data, they should not be modified directly + self.joint_pos = self.robot.data.joint_pos + self.joint_vel = self.robot.data.joint_vel + self.root_pose_w = self.robot.data.root_pose_w + self.root_vel_w = self.robot.data.root_vel_w + self.soft_joint_pos_limits = self.robot.data.soft_joint_pos_limits + + # Buffers + self.observations = wp.zeros( + (self.num_envs, self.cfg.observation_space), dtype=wp.float32, device=self.sim.device + ) + self.rewards = wp.zeros((self.num_envs), dtype=wp.float32, device=self.sim.device) + self.actions = wp.zeros((self.num_envs, self.robot.num_joints), dtype=wp.float32, device=self.sim.device) + self.states = wp.zeros((self.num_envs), dtype=wp.uint32, device=self.sim.device) + self.potentials = wp.zeros(self.num_envs, dtype=wp.float32, device=self.sim.device) + self.prev_potentials = wp.zeros_like(self.potentials) + self.targets = wp.zeros((self.num_envs), dtype=wp.vec3f, device=self.sim.device) + self.up_vec = wp.zeros((self.num_envs), dtype=wp.vec3f, device=self.sim.device) + self.heading_vec = wp.zeros((self.num_envs), dtype=wp.vec3f, device=self.sim.device) + self.to_targets = wp.zeros((self.num_envs), dtype=wp.vec3f, device=self.sim.device) + self.up_proj = wp.zeros((self.num_envs), dtype=wp.float32, device=self.sim.device) + self.heading_proj = wp.zeros((self.num_envs), dtype=wp.float32, device=self.sim.device) + self.vec_loc = wp.zeros((self.num_envs), dtype=wp.spatial_vectorf, device=self.sim.device) + self.rpy = wp.zeros((self.num_envs), dtype=wp.vec3f, device=self.sim.device) + self.angle_to_target = wp.zeros((self.num_envs), dtype=wp.float32, device=self.sim.device) + self.dof_pos_scaled = wp.zeros((self.num_envs, self.robot.num_joints), dtype=wp.float32, device=self.sim.device) + self.env_origins = wp.from_torch(self.scene.env_origins, dtype=wp.vec3f) + self.actions_mapped = wp.zeros((self.num_envs, self.robot.num_joints), dtype=wp.float32, device=self.sim.device) + + # Initial states and targets + if self.cfg.seed is None: + self.cfg.seed = -1 + + wp.launch( + initialize_state, + dim=self.num_envs, + inputs=[ + self.env_origins, + self.targets, + self.states, + self.cfg.seed, + ], + ) + + # Bind torch buffers to warp buffers + self.torch_obs_buf = wp.to_torch(self.observations) + self.torch_reward_buf = wp.to_torch(self.rewards) + self.torch_reset_terminated = wp.to_torch(self.reset_terminated) + self.torch_reset_time_outs = wp.to_torch(self.reset_time_outs) + self.torch_episode_length_buf = wp.to_torch(self.episode_length_buf) + + # Get the values from the robot data + print("joint_pos_limits_lower", wp.to_torch(self.robot.data.joint_pos_limits_lower)[0]) + print("joint_pos_limits_upper", wp.to_torch(self.robot.data.joint_pos_limits_upper)[0]) + print("joint_stiffness_sim", wp.to_torch(self.robot.data.joint_stiffness_sim)[0]) + print("joint_damping_sim", wp.to_torch(self.robot.data.joint_damping_sim)[0]) + print("joint_armature", wp.to_torch(self.robot.data.joint_armature)[0]) + print("joint_friction_coeff", wp.to_torch(self.robot.data.joint_friction_coeff)[0]) + print("joint_vel_limits_sim", wp.to_torch(self.robot.data.joint_vel_limits)[0]) + print("joint_effort_limits_sim", wp.to_torch(self.robot.data.joint_effort_limits)[0]) + print("joint_control_mode_sim", wp.to_torch(self.robot.data.joint_control_mode_sim)[0]) + print("joint_pos", wp.to_torch(self.robot.data.joint_pos)[0]) + print("joint_vel", wp.to_torch(self.robot.data.joint_vel)[0]) + print("joint_effort", wp.to_torch(self.robot.data.joint_effort_sim)[0]) + print("joint_target", wp.to_torch(self.robot.data.joint_target_sim)[0]) + print("soft_joint_pos_limits", wp.to_torch(self.robot.data.soft_joint_pos_limits)[0]) + + def _setup_scene(self) -> None: + self.robot = Articulation(self.cfg.robot, frontend=Frontend.WARP) + # add ground plane + self.cfg.terrain.num_envs = self.scene.cfg.num_envs + self.cfg.terrain.env_spacing = self.scene.cfg.env_spacing + self.terrain = self.cfg.terrain.class_type(self.cfg.terrain) + # clone and replicate + self.scene.clone_environments(copy_from_source=False) + # add articulation to scene + self.scene.articulations["robot"] = self.robot + # add lights + light_cfg = sim_utils.DomeLightCfg(intensity=2000.0, color=(0.75, 0.75, 0.75)) + light_cfg.func("/World/Light", light_cfg) + + def _pre_physics_step(self, actions: wp.array) -> None: + self.actions.assign(actions) + wp.launch( + update_actions, + dim=(self.num_envs, self.robot.num_joints), + inputs=[actions, self.actions_mapped, self.joint_gears, self.action_scale], + ) + + def _apply_action(self) -> None: + self.robot.set_joint_effort_target(self.actions_mapped, joint_mask=self._joint_dof_mask) + + def _compute_intermediate_values(self) -> None: + wp.launch( + compute_heading_and_up, + dim=self.num_envs, + inputs=[ + self.root_pose_w, + self.targets, + self.cfg.sim.dt, + self.to_targets, + self.up_proj, + self.heading_proj, + self.up_vec, + self.heading_vec, + self.potentials, + self.prev_potentials, + ], + ) + + wp.launch( + compute_rot, + dim=self.num_envs, + inputs=[ + self.root_pose_w, + self.root_vel_w, + self.targets, + self.vec_loc, + self.rpy, + self.angle_to_target, + ], + ) + wp.launch( + scale_dof_pos, + dim=(self.num_envs, self.robot.num_joints), + inputs=[ + self.joint_pos, + self.soft_joint_pos_limits, + self.dof_pos_scaled, + ], + ) + + def _get_observations(self) -> None: + wp.launch( + observations, + dim=self.num_envs, + inputs=[ + self.root_pose_w, + self.vec_loc, + self.rpy, + self.angle_to_target, + self.up_proj, + self.heading_proj, + self.dof_pos_scaled, + self.joint_vel, + self.actions, + self.observations, + self.cfg.dof_vel_scale, + self.cfg.angular_velocity_scale, + self.robot.num_joints, + ], + ) + + def _get_rewards(self) -> None: + wp.launch( + compute_rewards, + dim=self.num_envs, + inputs=[ + self.actions, + self.joint_vel, + self.dof_pos_scaled, + self.reset_terminated, + self.heading_proj, + self.up_proj, + self.potentials, + self.prev_potentials, + self.motor_effort_ratio, + self.cfg.up_weight, + self.cfg.heading_weight, + self.cfg.actions_cost_scale, + self.cfg.energy_cost_scale, + self.cfg.dof_vel_scale, + self.cfg.death_cost, + self.cfg.alive_reward_scale, + self.rewards, + ], + ) + + def _get_dones(self) -> None: + self._compute_intermediate_values() + + wp.launch( + get_dones, + dim=self.num_envs, + inputs=[ + self.episode_length_buf, + self.root_pose_w, + self.max_episode_length, + self.cfg.termination_height, + self.reset_terminated, + self.reset_time_outs, + self.reset_buf, + ], + ) + + def _reset_idx(self, mask: wp.array | None = None): + if mask is None: + mask = self.robot._ALL_ENV_MASK + + super()._reset_idx(mask) + + wp.launch( + reset_root, + dim=self.num_envs, + inputs=[ + self.robot.data.default_root_pose, + self.robot.data.default_root_vel, + self.env_origins, + self.cfg.sim.dt, + self.to_targets, + self.potentials, + self.root_pose_w, + self.root_vel_w, + mask, + ], + ) + wp.launch( + reset_joints, + dim=(self.num_envs, self.robot.num_joints), + inputs=[ + self.robot.data.default_joint_pos, + self.robot.data.default_joint_vel, + self.joint_pos, + self.joint_vel, + mask, + ], + ) + + self._compute_intermediate_values() diff --git a/source/isaaclab_tasks/isaaclab_tasks/manager_based/classic/cartpole/cartpole_env_cfg.py b/source/isaaclab_tasks/isaaclab_tasks/manager_based/classic/cartpole/cartpole_env_cfg.py index f452efda276..ffabd312697 100644 --- a/source/isaaclab_tasks/isaaclab_tasks/manager_based/classic/cartpole/cartpole_env_cfg.py +++ b/source/isaaclab_tasks/isaaclab_tasks/manager_based/classic/cartpole/cartpole_env_cfg.py @@ -15,6 +15,9 @@ from isaaclab.managers import SceneEntityCfg from isaaclab.managers import TerminationTermCfg as DoneTerm from isaaclab.scene import InteractiveSceneCfg +from isaaclab.sim import SimulationCfg +from isaaclab.sim._impl.newton_manager_cfg import NewtonCfg +from isaaclab.sim._impl.solvers_cfg import MJWarpSolverCfg from isaaclab.utils import configclass import isaaclab_tasks.manager_based.classic.cartpole.mdp as mdp @@ -167,6 +170,23 @@ class CartpoleEnvCfg(ManagerBasedRLEnvCfg): # MDP settings rewards: RewardsCfg = RewardsCfg() terminations: TerminationsCfg = TerminationsCfg() + # Simulation settings + sim: SimulationCfg = SimulationCfg( + newton_cfg=NewtonCfg( + solver_cfg=MJWarpSolverCfg( + njmax=5, + ncon_per_env=3, + ls_iterations=10, + cone="pyramidal", + impratio=1, + ls_parallel=True, + integrator="implicit", + ), + num_substeps=1, + debug_mode=False, + use_cuda_graph=True, + ) + ) # Post initialization def __post_init__(self) -> None: