diff --git a/source/isaaclab/isaaclab/actuators_warp/__init__.py b/source/isaaclab/isaaclab/actuators_warp/__init__.py new file mode 100644 index 00000000000..70e093ba065 --- /dev/null +++ b/source/isaaclab/isaaclab/actuators_warp/__init__.py @@ -0,0 +1,35 @@ +# 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 ActuatorBaseWarp +from .actuator_cfg import ( # ActuatorNetLSTMCfg,; ActuatorNetMLPCfg,; DelayedPDActuatorCfg,; RemotizedPDActuatorCfg, + ActuatorBaseWarpCfg, + DCMotorWarpCfg, + IdealPDActuatorWarpCfg, + ImplicitActuatorWarpCfg, +) + +# from .actuator_net import ActuatorNetLSTM, ActuatorNetMLP +# from .actuator_pd import DCMotor, DelayedPDActuator, IdealPDActuator, ImplicitActuator, RemotizedPDActuator +from .actuator_pd import DCMotorWarp, IdealPDActuatorWarp, ImplicitActuatorWarp diff --git a/source/isaaclab/isaaclab/actuators_warp/actuator_base.py b/source/isaaclab/isaaclab/actuators_warp/actuator_base.py new file mode 100644 index 00000000000..31ec9ff5bb0 --- /dev/null +++ b/source/isaaclab/isaaclab/actuators_warp/actuator_base.py @@ -0,0 +1,296 @@ +# 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.actuators_warp.kernels import clip_efforts_with_limits +from isaaclab.assets.articulation_warp.kernels import ( + populate_empty_array, + update_joint_array_with_value, + update_joint_array_with_value_array, + update_joint_array_with_value_int, +) + +if TYPE_CHECKING: + from isaaclab.assets.articulation_warp.articulation import ArticulationDataWarp + + from .actuator_cfg import ActuatorBaseWarpCfg + + +class ActuatorBaseWarp(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. + """ + + 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`. + """ + + data: ArticulationDataWarp + """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: ActuatorBaseWarpCfg, + joint_names: list[str], + joint_mask: wp.array, + env_mask: wp.array, + articulation_data: ArticulationDataWarp, + 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.data = articulation_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 + + # Parse the joint commands: + if self.cfg.control_mode == "position": + self.cfg.control_mode = 1 + elif self.cfg.control_mode == "velocity": + self.cfg.control_mode = 2 + elif self.cfg.control_mode == "none": + self.cfg.control_mode = 0 + + # 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.data.sim_bind_joint_vel_limits_sim), + ("effort_limit_sim", self.data.sim_bind_joint_effort_limits_sim), + ("stiffness", self.data.joint_stiffness), + ("damping", self.data.joint_damping), + ("armature", self.data.sim_bind_joint_armature), + ("friction", self.data.sim_bind_joint_friction_coeff), + ("dynamic_friction", self.data.joint_dynamic_friction), + ("viscous_friction", self.data.joint_viscous_friction), + ("control_mode", self.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(dtype=wp.int32)): + """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.data.joint_effort_limits_sim, + effort, + clipped_effort, + self._env_mask, + self.joint_mask, + ], + ) diff --git a/source/isaaclab/isaaclab/actuators_warp/actuator_cfg.py b/source/isaaclab/isaaclab/actuators_warp/actuator_cfg.py new file mode 100644 index 00000000000..01b2c34be5f --- /dev/null +++ b/source/isaaclab/isaaclab/actuators_warp/actuator_cfg.py @@ -0,0 +1,295 @@ +# 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 collections.abc import Iterable +from dataclasses import MISSING +from typing import Literal + +from isaaclab.utils import configclass + +from . import actuator_pd +from .actuator_base import ActuatorBaseWarp + + +@configclass +class ActuatorBaseWarpCfg: + """Configuration for default actuators in an articulation.""" + + class_type: type[ActuatorBaseWarp] = MISSING + """The associated actuator class. + + The class should inherit from :class:`isaaclab.actuators.ActuatorBase`. + """ + + joint_names_expr: list[str] = MISSING + """Articulation's joint names that are part of the group. + + Note: + This can be a list of joint names or a list of regex expressions (e.g. ".*"). + """ + + effort_limit: dict[str, float] | float | None = None + """Force/Torque limit of the joints in the group. Defaults to None. + + This limit is used to clip the computed torque sent to the simulation. If None, the + limit is set to the value specified in the USD joint prim. + + .. attention:: + + The :attr:`effort_limit_sim` attribute should be used to set the effort limit for + the simulation physics solver. + + The :attr:`effort_limit` attribute is used for clipping the effort output of the + actuator model **only** in the case of explicit actuators, such as the + :class:`~isaaclab.actuators.IdealPDActuator`. + + .. note:: + + For implicit actuators, the attributes :attr:`effort_limit` and :attr:`effort_limit_sim` + are equivalent. However, we suggest using the :attr:`effort_limit_sim` attribute because + it is more intuitive. + + """ + + velocity_limit: dict[str, float] | float | None = None + """Velocity limit of the joints in the group. Defaults to None. + + This limit is used by the actuator model. If None, the limit is set to the value specified + in the USD joint prim. + + .. attention:: + + The :attr:`velocity_limit_sim` attribute should be used to set the velocity limit for + the simulation physics solver. + + The :attr:`velocity_limit` attribute is used for clipping the effort output of the + actuator model **only** in the case of explicit actuators, such as the + :class:`~isaaclab.actuators.IdealPDActuator`. + + .. note:: + + For implicit actuators, the attribute :attr:`velocity_limit` is not used. This is to stay + backwards compatible with previous versions of the Isaac Lab, where this parameter was + unused since PhysX did not support setting the velocity limit for the joints using the + PhysX Tensor API. + """ + + effort_limit_sim: dict[str, float] | float | None = None + """Effort limit of the joints in the group applied to the simulation physics solver. Defaults to None. + + 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. + + Since explicit actuators (e.g. DC motor), compute and clip the effort in the actuator model, this + limit is by default set to a large value to prevent the physics engine from any additional clipping. + However, at times, it may be necessary to set this limit to a smaller value as a safety measure. + + If None, the limit is resolved based on the type of actuator model: + + * For implicit actuators, the limit is set to the value specified in the USD joint prim. + * For explicit actuators, the limit is set to 1.0e9. + + """ + + velocity_limit_sim: dict[str, float] | float | None = None + """Velocity limit of the joints in the group applied to the simulation physics solver. Defaults to None. + + 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. + + If None, the limit is set to the value specified in the USD joint prim for both implicit and explicit actuators. + + .. tip:: + If the velocity limit is too tight, the physics engine may have trouble converging to a solution. + In such cases, we recommend either keeping this value sufficiently large or tuning the stiffness and + damping parameters of the joint to ensure the limits are not violated. + + """ + + 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. + + The behavior of the stiffness is different for implicit and explicit actuators. For implicit actuators, + the stiffness gets set into the physics engine directly. For explicit actuators, the stiffness is used + by the actuator model to compute the joint efforts. + + If None, the stiffness is set to the value from the USD joint prim. + """ + + damping: dict[str, float] | float | None = MISSING + """Damping gains (also known as d-gain) of the joints in the group. + + The behavior of the damping is different for implicit and explicit actuators. For implicit actuators, + the damping gets set into the physics engine directly. For explicit actuators, the damping gain is used + by the actuator model to compute the joint efforts. + + If None, the damping is set to the value from the USD joint prim. + """ + + armature: dict[str, float] | float | None = None + """Armature of the joints in the group. Defaults to None. + + The armature is directly added to the corresponding joint-space inertia. It helps improve the + simulation stability by reducing the joint velocities. + + It is a physics engine solver parameter that gets set into the simulation. + + If None, the armature is set to the value from the USD joint prim. + """ + + friction: dict[str, float] | float | None = None + r"""The 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 + 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. + + If None, the joint 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. + """ + + +""" +Implicit Actuator Models. +""" + + +@configclass +class ImplicitActuatorWarpCfg(ActuatorBaseWarpCfg): + """Configuration for an implicit actuator. + + Note: + The PD control is handled implicitly by the simulation. + """ + + class_type: type = actuator_pd.ImplicitActuatorWarp + + +""" +Explicit Actuator Models. +""" + + +@configclass +class IdealPDActuatorWarpCfg(ActuatorBaseWarpCfg): + """Configuration for an ideal PD actuator.""" + + class_type: type = actuator_pd.IdealPDActuatorWarp + + +@configclass +class DCMotorWarpCfg(IdealPDActuatorWarpCfg): + """Configuration for direct control (DC) motor actuator model.""" + + class_type: type = actuator_pd.DCMotorWarp + + saturation_effort: float = MISSING + """Peak motor force/torque of the electric DC motor (in N-m).""" + + +# @configclass +# class ActuatorNetLSTMCfg(DCMotorCfg): +# """Configuration for LSTM-based actuator model.""" +# +# class_type: type = actuator_net.ActuatorNetLSTM +# # we don't use stiffness and damping for actuator net +# stiffness = None +# damping = None +# +# network_file: str = MISSING +# """Path to the file containing network weights.""" +# +# +# @configclass +# class ActuatorNetMLPCfg(DCMotorCfg): +# """Configuration for MLP-based actuator model.""" +# +# class_type: type = actuator_net.ActuatorNetMLP +# # we don't use stiffness and damping for actuator net +# stiffness = None +# damping = None +# +# network_file: str = MISSING +# """Path to the file containing network weights.""" +# +# pos_scale: float = MISSING +# """Scaling of the joint position errors input to the network.""" +# vel_scale: float = MISSING +# """Scaling of the joint velocities input to the network.""" +# torque_scale: float = MISSING +# """Scaling of the joint efforts output from the network.""" +# +# input_order: Literal["pos_vel", "vel_pos"] = MISSING +# """Order of the inputs to the network. +# +# The order can be one of the following: +# +# * ``"pos_vel"``: joint position errors followed by joint velocities +# * ``"vel_pos"``: joint velocities followed by joint position errors +# """ +# +# input_idx: Iterable[int] = MISSING +# """ +# Indices of the actuator history buffer passed as inputs to the network. +# +# The index *0* corresponds to current time-step, while *n* corresponds to n-th +# time-step in the past. The allocated history length is `max(input_idx) + 1`. +# """ +# +# +# @configclass +# class DelayedPDActuatorCfg(IdealPDActuatorCfg): +# """Configuration for a delayed PD actuator.""" +# +# class_type: type = actuator_pd.DelayedPDActuator +# +# min_delay: int = 0 +# """Minimum number of physics time-steps with which the actuator command may be delayed. Defaults to 0.""" +# +# max_delay: int = 0 +# """Maximum number of physics time-steps with which the actuator command may be delayed. Defaults to 0.""" +# +# +# @configclass +# class RemotizedPDActuatorCfg(DelayedPDActuatorCfg): +# """Configuration for a remotized PD actuator. +# +# Note: +# The torque output limits for this actuator is derived from a linear interpolation of a lookup table +# in :attr:`joint_parameter_lookup`. This table describes the relationship between joint angles and +# the output torques. +# """ +# +# class_type: type = actuator_pd.RemotizedPDActuator +# +# joint_parameter_lookup: list[list[float]] = MISSING +# """Joint parameter lookup table. Shape is (num_lookup_points, 3). +# +# This tensor describes the relationship between the joint angle (rad), the transmission ratio (in/out), +# and the output torque (N*m). The table is used to interpolate the output torque based on the joint angle. +# """ diff --git a/source/isaaclab/isaaclab/actuators_warp/actuator_net.py b/source/isaaclab/isaaclab/actuators_warp/actuator_net.py new file mode 100644 index 00000000000..4d45abc3bf0 --- /dev/null +++ b/source/isaaclab/isaaclab/actuators_warp/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 DCMotorWarp + +if TYPE_CHECKING: + from .actuator_cfg import ActuatorNetLSTMCfg, ActuatorNetMLPCfg + + +class ActuatorNetLSTM(DCMotorWarp): + """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(DCMotorWarp): + """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/actuators_warp/actuator_pd.py b/source/isaaclab/isaaclab/actuators_warp/actuator_pd.py new file mode 100644 index 00000000000..b1219a741d5 --- /dev/null +++ b/source/isaaclab/isaaclab/actuators_warp/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 ActuatorBaseWarp +from .kernels import clip_efforts_dc_motor, compute_pd_actuator + +# from isaaclab.utils import DelayBuffer, LinearInterpolation + + +if TYPE_CHECKING: + from .actuator_cfg import ( # DelayedPDActuatorCfg,; RemotizedPDActuatorCfg, + DCMotorWarpCfg, + IdealPDActuatorWarpCfg, + ImplicitActuatorWarpCfg, + ) + + +""" +Implicit Actuator Models. +""" + + +class ImplicitActuatorWarp(ActuatorBaseWarp): + """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: ImplicitActuatorWarpCfg + """The configuration for the actuator model.""" + + def __init__(self, cfg: ImplicitActuatorWarpCfg, *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 + ImplicitActuatorWarp.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.data.joint_target, + self.data.joint_effort_target, + self.data.sim_bind_joint_pos, + self.data.sim_bind_joint_vel, + self.data.joint_stiffness, + self.data.joint_damping, + self.data.joint_control_mode, + self.data.computed_effort, + self._env_mask, + self._joint_mask, + ], + ) + self._clip_effort(self.data.computed_effort, self.data.applied_effort) + + +""" +Explicit Actuator Models. +""" + + +class IdealPDActuatorWarp(ActuatorBaseWarp): + 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: IdealPDActuatorWarpCfg + """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.data.joint_target, + self.data.joint_effort_target, + self.data.sim_bind_joint_pos, + self.data.sim_bind_joint_vel, + self.data.joint_stiffness, + self.data.joint_damping, + self.data.joint_control_mode, + self.data.computed_effort, + self._env_mask, + self._joint_mask, + ], + ) + self._clip_effort(self.data.computed_effort, self.data.sim_bind_joint_effort) + + +class DCMotorWarp(IdealPDActuatorWarp): + 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: DCMotorWarpCfg + """The configuration for the actuator model.""" + + def __init__(self, cfg: DCMotorWarpCfg, *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.data.sim_bind_joint_effort_limits_sim, + self.data.sim_bind_joint_vel_limits_sim, + self.data.sim_bind_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/actuators_warp/kernels.py b/source/isaaclab/isaaclab/actuators_warp/kernels.py new file mode 100644 index 00000000000..0b2391db415 --- /dev/null +++ b/source/isaaclab/isaaclab/actuators_warp/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/assets/__init__.py b/source/isaaclab/isaaclab/assets/__init__.py index 57dafb20569..12a73b806f9 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 .articulation_warp import ArticulationDataWarp, ArticulationWarp, ArticulationWarpCfg from .asset_base import AssetBase from .asset_base_cfg import AssetBaseCfg diff --git a/source/isaaclab/isaaclab/assets/articulation/articulation_data.py b/source/isaaclab/isaaclab/assets/articulation/articulation_data.py index 28e75a48431..dbea88ac217 100644 --- a/source/isaaclab/isaaclab/assets/articulation/articulation_data.py +++ b/source/isaaclab/isaaclab/assets/articulation/articulation_data.py @@ -50,7 +50,7 @@ def __init__(self, root_newton_view, device: str): self._sim_timestamp = 0.0 # obtain global simulation view - gravity = NewtonManager.get_model().gravity + gravity = wp.to_torch(NewtonManager.get_model().gravity)[0] # 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) diff --git a/source/isaaclab/isaaclab/assets/articulation_new/__init__.py b/source/isaaclab/isaaclab/assets/articulation_new/__init__.py new file mode 100644 index 00000000000..f324a56339f --- /dev/null +++ b/source/isaaclab/isaaclab/assets/articulation_new/__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 articulated assets.""" + +from .articulation import ArticulationWarp +from .articulation_cfg import ArticulationWarpCfg +from .articulation_data import ArticulationDataWarp diff --git a/source/isaaclab/isaaclab/assets/articulation_new/articulation_cfg.py b/source/isaaclab/isaaclab/assets/articulation_new/articulation_cfg.py new file mode 100644 index 00000000000..5ef675ad22d --- /dev/null +++ b/source/isaaclab/isaaclab/assets/articulation_new/articulation_cfg.py @@ -0,0 +1,61 @@ +# 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.actuators_warp import ActuatorBaseWarpCfg +from isaaclab.utils import configclass + +from ..asset_base_cfg import AssetBaseCfg +from .articulation import ArticulationWarp + + +@configclass +class ArticulationWarpCfg(AssetBaseCfg): + """Configuration parameters for an articulation.""" + + @configclass + class InitialStateCfg(AssetBaseCfg.InitialStateCfg): + """Initial state of the articulation.""" + + # root velocity + 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).""" + + # joint state + joint_pos: dict[str, float] = {".*": 0.0} + """Joint positions of the joints. Defaults to 0.0 for all joints.""" + joint_vel: dict[str, float] = {".*": 0.0} + """Joint velocities of the joints. Defaults to 0.0 for all joints.""" + + ## + # Initialize configurations. + ## + + class_type: type = ArticulationWarp + + articulation_root_prim_path: str | None = None + """Path to the articulation root prim in the USD file. + + If not provided will search for a prim with the ArticulationRootAPI. Should start with a slash. + """ + + init_state: InitialStateCfg = InitialStateCfg() + """Initial state of the articulated object. Defaults to identity pose with zero velocity and zero joint state.""" + + soft_joint_pos_limit_factor: float = 1.0 + """Fraction specifying the range of joint position limits (parsed from the asset) to use. Defaults to 1.0. + + The soft joint position limits are scaled by this factor to specify a safety region within the simulated + joint position limits. This isn't used by the simulation, but is useful for learning agents to prevent the joint + positions from violating the limits, such as for termination conditions. + + The soft joint position limits are accessible through the :attr:`ArticulationData.soft_joint_pos_limits` attribute. + """ + + actuators: dict[str, ActuatorBaseWarpCfg] = MISSING + """Actuators for the robot with corresponding joint names.""" diff --git a/source/isaaclab/isaaclab/assets/articulation_new/articulation_data_torch.py b/source/isaaclab/isaaclab/assets/articulation_new/articulation_data_torch.py new file mode 100644 index 00000000000..efb824d256f --- /dev/null +++ b/source/isaaclab/isaaclab/assets/articulation_new/articulation_data_torch.py @@ -0,0 +1,952 @@ +# 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 +import torch + +from isaaclab.assets.core.root_properties.root_data import RootData +from isaaclab.assets.core.body_properties.body_data import BodyData +from isaaclab.assets.core.joint_properties.joint_data import JointData +from isaaclab.assets.core.kernels import vec13f, combine_pose_and_velocity_to_state_batched +from isaaclab.utils.helpers import deprecated + + +class ArticulationDataTorch: + def __init__(self, root_newton_view, device: str): + # 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) + + # 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) + + @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) -> 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 wp.to_torch(self._root_data.default_root_pose) + + @property + def default_root_vel(self) -> 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 wp.to_torch(self._root_data.default_root_vel) + + @property + def default_joint_pos(self) -> 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 wp.to_torch(self._joint_data.default_joint_pos) + + @property + def default_joint_vel(self) -> 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 wp.to_torch(self._joint_data.default_joint_vel) + + ## + # Joint commands -- From the user to the actuator model. + ## + + @property + def joint_target(self) -> 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 wp.to_torch(self._joint_data.joint_target) + + @property + def joint_effort_target(self) -> 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 wp.to_torch(self._joint_data.joint_effort_target) + + ## + # Joint commands -- Explicit actuators. + ## + + @property + def computed_effort(self) -> 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 wp.to_torch(self._joint_data.computed_effort) + + @property + def applied_effort(self) -> 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 wp.to_torch(self._joint_data.applied_effort) + + @property + def joint_stiffness(self) -> torch.Tensor: + """Joint stiffness. Shape is (num_instances, num_joints).""" + return wp.to_torch(self._joint_data.joint_stiffness) + + @property + def joint_damping(self) -> torch.Tensor: + """Joint damping. Shape is (num_instances, num_joints).""" + return wp.to_torch(self._joint_data.joint_damping) + + @property + def joint_control_mode(self) -> torch.Tensor: + """Joint control mode. Shape is (num_instances, num_joints).""" + return wp.to_torch(self._joint_data.joint_control_mode) + + ### + # Joint commands. (Directly binded to the simulation) + ### + + @property + def joint_target_sim(self) -> torch.Tensor: + """Joint target. Shape is (num_instances, num_joints).""" + return wp.to_torch(self._joint_data.joint_target_sim) + + @property + def joint_effort_sim(self) -> torch.Tensor: + """Joint effort. Shape is (num_instances, num_joints).""" + return wp.to_torch(self._joint_data.joint_effort_sim) + + ## + # Joint properties. (Directly binded to the simulation) + ## + + @property + def joint_control_mode_sim(self) -> 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 wp.to_torch(self._joint_data.joint_control_mode_sim) + + @property + def joint_stiffness_sim(self) -> 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 wp.to_torch(self._joint_data.joint_stiffness_sim) + + @property + def joint_damping_sim(self) -> 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 wp.to_torch(self._joint_data.joint_damping_sim) + + @property + def joint_armature(self) -> torch.Tensor: + """Joint armature provided to the simulation. Shape is (num_instances, num_joints).""" + return wp.to_torch(self._joint_data.joint_armature) + + @property + def joint_friction_coeff(self) -> torch.Tensor: + """Joint friction coefficient provided to the simulation. Shape is (num_instances, num_joints).""" + return wp.to_torch(self._joint_data.joint_friction_coeff) + + @property + def joint_pos_limits_lower(self) -> torch.Tensor: + """Joint position limits lower provided to the simulation. Shape is (num_instances, num_joints).""" + return wp.to_torch(self._joint_data.joint_pos_limits_lower) + + @property + def joint_pos_limits_upper(self) -> torch.Tensor: + """Joint position limits upper provided to the simulation. Shape is (num_instances, num_joints).""" + return wp.to_torch(self._joint_data.joint_pos_limits_upper) + + + @property + def joint_vel_limits(self) -> torch.Tensor: + """Joint maximum velocity provided to the simulation. Shape is (num_instances, num_joints).""" + return wp.to_torch(self._joint_data.joint_vel_limits) + + @property + def joint_effort_limits(self) -> torch.Tensor: + """Joint maximum effort provided to the simulation. Shape is (num_instances, num_joints).""" + return wp.to_torch(self._joint_data.joint_effort_limits) + + ## + # Joint states + ## + + @property + def joint_pos(self) -> torch.Tensor: + """Joint positions. Shape is (num_instances, num_joints).""" + return wp.to_torch(self._joint_data.joint_pos) + + @property + def joint_vel(self) -> torch.Tensor: + """Joint velocities. Shape is (num_instances, num_joints).""" + return wp.to_torch(self._joint_data.joint_vel) + + ## + # Joint properties - Custom. + ## + + @property + def joint_dynamic_friction(self) -> torch.Tensor: + """Joint dynamic friction. Shape is (num_instances, num_joints).""" + return wp.to_torch(self._joint_data.joint_dynamic_friction) + + @property + def joint_viscous_friction(self) -> torch.Tensor: + """Joint viscous friction. Shape is (num_instances, num_joints).""" + return wp.to_torch(self._joint_data.joint_viscous_friction) + + @property + def soft_joint_pos_limits(self) -> 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 wp.to_torch(self._joint_data.soft_joint_pos_limits) + + @property + def soft_joint_vel_limits(self) -> 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 wp.to_torch(self._joint_data.soft_joint_vel_limits) + + @property + def gear_ratio(self) -> torch.Tensor: + """Gear ratio for relating motor torques to applied Joint torques. Shape is (num_instances, num_joints).""" + return wp.to_torch(self._joint_data.gear_ratio) + + ## + # Root state properties. + ## + + @property + def root_mass(self) -> torch.Tensor: + """Root mass ``wp.float32`` in the world frame. Shape is (num_instances,).""" + return wp.to_torch(self._root_data.root_mass) + + @property + def root_inertia(self) -> torch.Tensor: + """Root inertia ``wp.mat33`` in the world frame. Shape is (num_instances, 9).""" + return wp.to_torch(self._root_data.root_inertia) + + @property + def root_link_pose_w(self) -> 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 wp.to_torch(self._root_data.root_link_pose_w) + + @property + def root_link_vel_w(self) -> 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 wp.to_torch(self._root_data.root_link_vel_w) + + @property + def root_com_pose_w(self) -> 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 wp.to_torch(self._root_data.root_com_pose_w) + + @property + def root_com_vel_w(self) -> 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 wp.to_torch(self._root_data.root_com_vel_w) + + @property + def root_state_w(self) -> 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 wp.to_torch(self._root_data.root_state_w) + + @property + def root_link_state_w(self) -> 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 wp.to_torch(self._root_data.root_link_state_w) + + @property + def root_com_state_w(self) -> 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 wp.to_torch(self._root_data.root_com_state_w) + + @property + def root_com_pose_b(self) -> 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 wp.to_torch(self._root_data.root_com_pose_b) + + @property + def root_com_pos_b(self) -> 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 wp.to_torch(self._root_data.root_com_pos_b) + + @property + def root_com_quat_b(self) -> 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 wp.to_torch(self._root_data.root_com_quat_b) + + ## + # Body state properties. + ## + + @property + def body_mass(self) -> torch.Tensor: + """Body mass ``wp.float32`` in the world frame. Shape is (num_instances, num_bodies).""" + return wp.to_torch(self._body_data.body_mass) + + @property + def body_inertia(self) -> torch.Tensor: + """Body inertia ``wp.mat33`` in the world frame. Shape is (num_instances, num_bodies, 9).""" + return wp.to_torch(self._body_data.body_inertia) + + @property + def body_link_pose_w(self) -> 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 wp.to_torch(self._body_data.body_link_pose_w) + + @property + def body_link_vel_w(self) -> 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 wp.to_torch(self._body_data.body_link_vel_w) + + @property + def body_com_pose_w(self) -> 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 wp.to_torch(self._body_data.body_com_pose_w) + + @property + def body_com_vel_w(self) -> 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 wp.to_torch(self._body_data.body_com_vel_w) + + @property + def body_state_w(self) -> 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 wp.to_torch(self._body_data.body_state_w) + + @property + def body_link_state_w(self) -> 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 wp.to_torch(self._body_data.body_link_state_w) + + @property + def body_com_state_w(self) -> 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 wp.to_torch(self._body_data.body_com_state_w) + + @property + def body_com_acc_w(self) -> 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 wp.to_torch(self._body_data.body_com_acc_w) + + @property + def body_com_pos_b(self) -> 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 wp.to_torch(self._body_data.body_com_pos_b) + + @property + def body_com_pose_b(self) -> 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 wp.to_torch(self._body_data.body_com_pose_b) + + @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 `__ . + """ + return wp.to_torch(self._joint_data.body_incoming_joint_wrench_b) + + ## + # Joint state properties. + ## + + @property + def joint_acc(self) -> torch.Tensor: + """Joint acceleration of all joints. Shape is (num_instances, num_joints).""" + return wp.to_torch(self._joint_data.joint_acc) + ## + # Derived Properties. + ## + + # FIXME: USE SIM_BIND_LINK_POSE_W RATHER THAN JUST THE QUATERNION + @property + def projected_gravity_b(self) -> torch.Tensor: + """Projection of the gravity direction on base frame. Shape is (num_instances, 3).""" + return wp.to_torch(self._root_data.projected_gravity_b) + + # FIXME: USE SIM_BIND_LINK_POSE_W RATHER THAN JUST THE QUATERNION + @property + 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)`. + """ + return wp.to_torch(self._root_data.heading_w) + + @property + def root_link_vel_b(self) -> 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 wp.to_torch(self._root_data.root_link_vel_b) + + @property + def root_com_vel_b(self) -> 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 wp.to_torch(self._root_data.root_com_vel_b) + + ## + # Sliced properties. + ## + + @property + def root_link_pos_w(self) -> 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 wp.to_torch(self._root_data.root_link_pos_w) + + @property + def root_link_quat_w(self) -> 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 wp.to_torch(self._root_data.root_link_quat_w) + + @property + def root_link_lin_vel_w(self) -> 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 wp.to_torch(self._root_data.root_link_lin_vel_w) + + @property + def root_link_ang_vel_w(self) -> 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 wp.to_torch(self._root_data.root_link_ang_vel_w) + + @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 wp.to_torch(self._root_data.root_com_pos_w) + + @property + def root_com_quat_w(self) -> 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 wp.to_torch(self._root_data.root_com_quat_w) + + @property + def root_com_lin_vel_w(self) -> 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 wp.to_torch(self._root_data.root_com_lin_vel_w) + + @property + def root_com_ang_vel_w(self) -> 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 wp.to_torch(self._root_data.root_com_ang_vel_w) + + @property + def body_link_pos_w(self) -> 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 wp.to_torch(self._body_data.body_link_pos_w) + + @property + def body_link_quat_w(self) -> 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 wp.to_torch(self._body_data.body_link_quat_w) + + @property + def body_link_lin_vel_w(self) -> 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 wp.to_torch(self._body_data.body_link_lin_vel_w) + + @property + def body_link_ang_vel_w(self) -> 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 relative to the world. + """ + return wp.to_torch(self._body_data.body_link_ang_vel_w) + + @property + def body_com_pos_w(self) -> 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 wp.to_torch(self._body_data.body_com_pos_w) + + @property + def body_com_quat_w(self) -> 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 wp.to_torch(self._body_data.body_com_quat_w) + + @property + def body_com_lin_vel_w(self) -> 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 wp.to_torch(self._body_data.body_com_lin_vel_w) + + @property + def body_com_ang_vel_w(self) -> 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 wp.to_torch(self._body_data.body_com_ang_vel_w) + + @property + def body_com_lin_acc_w(self) -> 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 wp.to_torch(self._body_data.body_com_lin_acc_w) + + @property + def body_com_ang_acc_w(self) -> 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 wp.to_torch(self._body_data.body_com_ang_acc_w) + + @property + def body_com_quat_b(self) -> 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 wp.to_torch(self._body_data.body_com_quat_b) + + @property + def root_link_lin_vel_b(self) -> 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 wp.to_torch(self._root_data.root_link_lin_vel_b) + + @property + def root_link_ang_vel_b(self) -> 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 wp.to_torch(self._root_data.root_link_ang_vel_b) + + @property + def root_com_lin_vel_b(self) -> 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 wp.to_torch(self._root_data.root_com_lin_vel_b) + + @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 wp.to_torch(self._root_data.root_com_ang_vel_b) + + @property + def joint_pos_limits(self) -> 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 wp.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) -> torch.Tensor: + """Same as :attr:`default_root_pose`.""" + state = wp.zeros((self._root_newton_view.count), dtype=vec13f, device=self.device) + wp.launch( + combine_pose_and_velocity_to_state_batched, + dim=(self._root_newton_view.count), + device=self.device, + inputs=[ + self.default_root_pose, + self.default_root_vel, + state, + ], + ) + return wp.to_torch(state) + + @property + @deprecated("root_link_pose_w") + def root_pose_w(self) -> torch.Tensor: + """Same as :attr:`root_link_pose_w`.""" + return wp.to_torch(self.root_link_pose_w) + + @property + @deprecated("root_link_pos_w") + def root_pos_w(self) -> torch.Tensor: + """Same as :attr:`root_link_pos_w`.""" + return wp.to_torch(self.root_link_pos_w) + + @property + @deprecated("root_link_quat_w") + def root_quat_w(self) -> torch.Tensor: + """Same as :attr:`root_link_quat_w`.""" + return wp.to_torch(self.root_link_quat_w) + + @property + @deprecated("root_com_vel_w") + def root_vel_w(self) -> torch.Tensor: + """Same as :attr:`root_com_vel_w`.""" + return wp.to_torch(self.root_com_vel_w) + + @property + @deprecated("root_com_lin_vel_w") + def root_lin_vel_w(self) -> torch.Tensor: + """Same as :attr:`root_com_lin_vel_w`.""" + return wp.to_torch(self.root_com_lin_vel_w) + + @property + @deprecated("root_com_ang_vel_w") + def root_ang_vel_w(self) -> torch.Tensor: + """Same as :attr:`root_com_ang_vel_w`.""" + return wp.to_torch(self.root_com_ang_vel_w) + + @property + @deprecated("root_com_lin_vel_b") + def root_lin_vel_b(self) -> torch.Tensor: + """Same as :attr:`root_com_lin_vel_b`.""" + return wp.to_torch(self.root_com_lin_vel_b) + + @property + @deprecated("root_com_ang_vel_b") + def root_ang_vel_b(self) -> torch.Tensor: + """Same as :attr:`root_com_ang_vel_b`.""" + return wp.to_torch(self.root_com_ang_vel_b) + + @property + @deprecated("body_link_pose_w") + def body_pose_w(self) -> torch.Tensor: + """Same as :attr:`body_link_pose_w`.""" + return wp.to_torch(self.body_link_pose_w) + + @property + @deprecated("body_link_pos_w") + def body_pos_w(self) -> torch.Tensor: + """Same as :attr:`body_link_pos_w`.""" + return wp.to_torch(self.body_link_pos_w) + + @property + @deprecated("body_link_quat_w") + def body_quat_w(self) -> torch.Tensor: + """Same as :attr:`body_link_quat_w`.""" + return wp.to_torch(self.body_link_quat_w) + + @property + @deprecated("body_com_vel_w") + def body_vel_w(self) -> torch.Tensor: + """Same as :attr:`body_com_vel_w`.""" + return wp.to_torch(self.body_com_vel_w) + + @property + @deprecated("body_com_lin_vel_w") + def body_lin_vel_w(self) -> torch.Tensor: + """Same as :attr:`body_com_lin_vel_w`.""" + return wp.to_torch(self.body_com_lin_vel_w) + + @property + @deprecated("body_com_ang_vel_w") + def body_ang_vel_w(self) -> torch.Tensor: + """Same as :attr:`body_com_ang_vel_w`.""" + return wp.to_torch(self.body_com_ang_vel_w) + + @property + @deprecated("body_com_acc_w") + def body_acc_w(self) -> torch.Tensor: + """Same as :attr:`body_com_acc_w`.""" + return wp.to_torch(self.body_com_acc_w) + + @property + @deprecated("body_com_lin_acc_w") + def body_lin_acc_w(self) -> torch.Tensor: + """Same as :attr:`body_com_lin_acc_w`.""" + return wp.to_torch(self.body_com_lin_acc_w) + + @property + @deprecated("body_com_ang_acc_w") + def body_ang_acc_w(self) -> torch.Tensor: + """Same as :attr:`body_com_ang_acc_w`.""" + return wp.to_torch(self.body_com_ang_acc_w) + + @property + @deprecated("body_com_pos_b") + def com_pos_b(self) -> torch.Tensor: + """Same as :attr:`body_com_pos_b`.""" + return wp.to_torch(self.body_com_pos_b) + + @property + @deprecated("body_com_quat_b") + def com_quat_b(self) -> torch.Tensor: + """Same as :attr:`body_com_quat_b`.""" + return wp.to_torch(self.body_com_quat_b) + + @property + @deprecated("joint_pos_limits") + def joint_limits(self) -> torch.Tensor: + """Deprecated property. Please use :attr:`joint_pos_limits` instead.""" + return wp.to_torch(self.joint_pos_limits) + + @property + @deprecated("joint_friction_coeff") + def joint_friction(self) -> torch.Tensor: + """Deprecated property. Please use :attr:`joint_friction_coeff` instead.""" + return wp.to_torch(self.joint_friction_coeff) diff --git a/source/isaaclab/isaaclab/assets/articulation_new/articulation_data_warp.py b/source/isaaclab/isaaclab/assets/articulation_new/articulation_data_warp.py new file mode 100644 index 00000000000..3efc896c9d6 --- /dev/null +++ b/source/isaaclab/isaaclab/assets/articulation_new/articulation_data_warp.py @@ -0,0 +1,952 @@ +# 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.assets.core.root_properties.root_data import RootData +from isaaclab.assets.core.body_properties.body_data import BodyData +from isaaclab.assets.core.joint_properties.joint_data import JointData +from isaaclab.assets.core.kernels import vec13f, combine_pose_and_velocity_to_state_batched +from isaaclab.utils.helpers import deprecated + + +class ArticulationDataWarp: + + def __init__(self, root_newton_view, device: str): + # 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) + + # 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) + + @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: + """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._root_data.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._root_data.default_root_vel + + @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._joint_data.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._joint_data.default_joint_vel + + ## + # Joint commands -- From the user to the actuator model. + ## + + @property + def joint_target(self) -> wp.array: + """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._joint_data.joint_target + + @property + def joint_effort_target(self) -> wp.array: + """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._joint_data.joint_effort_target + + ## + # Joint commands -- Explicit actuators. + ## + + @property + def computed_effort(self) -> wp.array: + """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._joint_data.computed_effort + + @property + def applied_effort(self) -> wp.array: + """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._joint_data.applied_effort + + @property + def joint_stiffness(self) -> wp.array: + """Joint stiffness. Shape is (num_instances, num_joints).""" + return self._joint_data.joint_stiffness + + @property + def joint_damping(self) -> wp.array: + """Joint damping. Shape is (num_instances, num_joints).""" + return self._joint_data.joint_damping + + @property + def joint_control_mode(self) -> wp.array: + """Joint control mode. Shape is (num_instances, num_joints).""" + return self._joint_data.joint_control_mode + + ### + # Joint commands. (Directly binded to the simulation) + ### + + @property + def joint_target_sim(self) -> wp.array: + """Joint target. Shape is (num_instances, num_joints).""" + return self._joint_data.joint_target_sim + + @property + def joint_effort_sim(self) -> wp.array: + """Joint effort. Shape is (num_instances, num_joints).""" + return self._joint_data.joint_effort_sim + + ## + # Joint properties. (Directly binded to the simulation) + ## + + @property + def joint_control_mode_sim(self) -> wp.array: + """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._joint_data.joint_control_mode_sim + + @property + def joint_stiffness_sim(self) -> wp.array: + """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._joint_data.joint_stiffness_sim + + @property + def joint_damping_sim(self) -> wp.array: + """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._joint_data.joint_damping_sim + + @property + def joint_armature(self) -> wp.array: + """Joint armature provided to the simulation. Shape is (num_instances, num_joints).""" + return self._joint_data.joint_armature + + @property + def joint_friction_coeff(self) -> wp.array: + """Joint friction coefficient provided to the simulation. Shape is (num_instances, num_joints).""" + return self._joint_data.joint_friction_coeff + + @property + def joint_pos_limits_lower(self) -> wp.array: + """Joint position limits lower provided to the simulation. Shape is (num_instances, num_joints).""" + return self._joint_data.joint_pos_limits_lower + + @property + def joint_pos_limits_upper(self) -> wp.array: + """Joint position limits upper provided to the simulation. Shape is (num_instances, num_joints).""" + return self._joint_data.joint_pos_limits_upper + + + @property + def joint_vel_limits(self) -> wp.array: + """Joint maximum velocity provided to the simulation. Shape is (num_instances, num_joints).""" + return self._joint_data.joint_vel_limits + + @property + def joint_effort_limits(self) -> wp.array: + """Joint maximum effort provided to the simulation. Shape is (num_instances, num_joints).""" + return self._joint_data.joint_effort_limits + + ## + # Joint states + ## + + @property + def joint_pos(self) -> wp.array: + """Joint positions. Shape is (num_instances, num_joints).""" + return self._joint_data.joint_pos + + @property + def joint_vel(self) -> wp.array: + """Joint velocities. Shape is (num_instances, num_joints).""" + return self._joint_data.joint_vel + + ## + # Joint properties - Custom. + ## + + @property + def joint_dynamic_friction(self) -> wp.array: + """Joint dynamic friction. Shape is (num_instances, num_joints).""" + return self._joint_data.joint_dynamic_friction + + @property + def joint_viscous_friction(self) -> wp.array: + """Joint viscous friction. Shape is (num_instances, num_joints).""" + return self._joint_data.joint_viscous_friction + + @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._joint_data.soft_joint_pos_limits + + @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._joint_data.soft_joint_vel_limits + + @property + def gear_ratio(self) -> wp.array: + """Gear ratio for relating motor torques to applied Joint torques. Shape is (num_instances, num_joints).""" + return self._joint_data.gear_ratio + + ## + # Root state properties. + ## + + @property + def root_mass(self) -> wp.array: + """Root mass ``wp.float32`` in the world frame. Shape is (num_instances,).""" + return self._root_data.root_mass + + @property + def root_inertia(self) -> wp.array: + """Root inertia ``wp.mat33`` in the world frame. Shape is (num_instances, 9).""" + return self._root_data.root_inertia + + @property + def root_link_pose_w(self) -> wp.array: + """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._root_data.root_link_pose_w + + @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. + """ + + return self._root_data.root_link_vel_w + + @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. + The orientation is provided in (x, y, z, w) format. + """ + + return self._root_data.root_com_pose_w + + @property + def root_com_vel_w(self) -> wp.array: + """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._root_data.root_com_vel_w + + @property + def root_state_w(self) -> wp.array: + """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._root_data.root_state_w + + @property + def root_link_state_w(self) -> wp.array: + """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._root_data.root_link_state_w + + @property + def root_com_state_w(self) -> wp.array: + """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._root_data.root_com_state_w + + @property + def root_com_pose_b(self) -> wp.array: + """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._root_data.root_com_pose_b + + @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._root_data.root_com_pos_b + + @property + def root_com_quat_b(self) -> wp.array: + """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._root_data.root_com_quat_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._body_data.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._body_data.body_inertia + + @property + def body_link_pose_w(self) -> wp.array: + """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._body_data.body_link_pose_w + + @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. + """ + return self._body_data.body_link_vel_w + + @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. + """ + return self._body_data.body_com_pose_w + + @property + def body_com_vel_w(self) -> wp.array: + """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._body_data.body_com_vel_w + + @property + def body_state_w(self) -> wp.array: + """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._body_data.body_state_w + + @property + def body_link_state_w(self) -> wp.array: + """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._body_data.body_link_state_w + + @property + 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 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._body_data.body_com_state_w + + @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,). The acceleration is in the form of [wx, wy, wz, vx, vy, vz]. + All values are relative to the world. + """ + return self._body_data.body_com_acc_w + + @property + def body_com_pos_b(self) -> wp.array: + """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._body_data.body_com_pos_b + + @property + 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. + The orientation is provided in (x, y, z, w) format. + """ + return self._body_data.body_com_pose_b + + @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 `__ . + """ + return self._joint_data.body_incoming_joint_wrench_b + + ## + # Joint state properties. + ## + + @property + def joint_acc(self) -> wp.array: + """Joint acceleration of all joints. Shape is (num_instances, num_joints).""" + return self._joint_data.joint_acc + ## + # Derived Properties. + ## + + # FIXME: USE SIM_BIND_LINK_POSE_W RATHER THAN JUST THE QUATERNION + @property + def projected_gravity_b(self): + """Projection of the gravity direction on base frame. Shape is (num_instances, 3).""" + return self._root_data.projected_gravity_b + + # FIXME: USE SIM_BIND_LINK_POSE_W RATHER THAN JUST THE QUATERNION + @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)`. + """ + return self._root_data.heading_w + + @property + def root_link_vel_b(self) -> wp.array: + """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._root_data.root_link_vel_b + + @property + def root_com_vel_b(self) -> wp.array: + """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._root_data.root_com_vel_b + + ## + # Sliced properties. + ## + + @property + 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. + """ + return self._root_data.root_link_pos_w + + @property + def root_link_quat_w(self) -> wp.array: + """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._root_data.root_link_quat_w + + @property + 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. + """ + return self._root_data.root_link_lin_vel_w + + @property + 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. + """ + return self._root_data.root_link_ang_vel_w + + @property + 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. + """ + return self._root_data.root_com_pos_w + + @property + def root_com_quat_w(self) -> wp.array: + """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._root_data.root_com_quat_w + + @property + 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. + """ + return self._root_data.root_com_lin_vel_w + + @property + 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. + """ + return self._root_data.root_com_ang_vel_w + + @property + 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. + """ + return self._body_data.body_link_pos_w + + @property + def body_link_quat_w(self) -> wp.array: + """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._body_data.body_link_quat_w + + @property + 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. + """ + return self._body_data.body_link_lin_vel_w + + @property + 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. + """ + return self._body_data.body_link_ang_vel_w + + @property + 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. + """ + return self._body_data.body_com_pos_w + + @property + def body_com_quat_w(self) -> wp.array: + """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._body_data.body_com_quat_w + + @property + 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. + """ + return self._body_data.body_com_lin_vel_w + + @property + 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. + """ + return self._body_data.body_com_ang_vel_w + + @property + 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. + """ + return self._body_data.body_com_lin_acc_w + + @property + 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. + """ + return self._body_data.body_com_ang_acc_w + + @property + def body_com_quat_b(self) -> wp.array: + """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._body_data.body_com_quat_b + + @property + 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. + """ + return self._root_data.root_link_lin_vel_b + + @property + 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. + """ + return self._root_data.root_link_ang_vel_b + + @property + 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. + """ + return self._root_data.root_com_lin_vel_b + + @property + 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. + """ + return self._root_data.root_com_ang_vel_b + + @property + 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]`. + """ + return 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: + """Same as :attr:`default_root_pose`.""" + state = wp.zeros((self._root_newton_view.count), dtype=vec13f, device=self.device) + wp.launch( + combine_pose_and_velocity_to_state_batched, + dim=(self._root_newton_view.count), + device=self.device, + inputs=[ + self.default_root_pose, + self.default_root_vel, + state, + ], + ) + return state + + @property + @deprecated("root_link_pose_w") + def root_pose_w(self) -> wp.array: + """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: + """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: + """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: + """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: + """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: + """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: + """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: + """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: + """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: + """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: + """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: + """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: + """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: + """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: + """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: + """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: + """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: + """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: + """Same as :attr:`body_com_quat_b`.""" + return self.body_com_quat_b + + @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.joint_friction_coeff diff --git a/source/isaaclab/isaaclab/assets/articulation_new/articulation_torch.py b/source/isaaclab/isaaclab/assets/articulation_new/articulation_torch.py new file mode 100644 index 00000000000..65e0f3999ff --- /dev/null +++ b/source/isaaclab/isaaclab/assets/articulation_new/articulation_torch.py @@ -0,0 +1,1266 @@ +# 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 +from collections.abc import Sequence +from prettytable import PrettyTable +from typing import TYPE_CHECKING + +import omni.log +import warp as wp +from isaacsim.core.simulation_manager import SimulationManager +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.actuators_warp import ActuatorBaseWarp, ActuatorBaseWarpCfg, ImplicitActuatorWarp +from isaaclab.sim._impl.newton_manager import NewtonManager + +from ..asset_base import AssetBase +from .articulation_data_torch import ArticulationDataTorch +from isaaclab.assets.core.joint_properties.joint import Joint +from isaaclab.assets.core.body_properties.body import Body +from isaaclab.assets.core.root_properties.root import Root +from isaaclab.assets.core.kernels import ( + generate_mask_from_ids, + populate_empty_array, + update_batched_array_with_array_masked, + update_array_with_value, +) + +if TYPE_CHECKING: + from .articulation_cfg import ArticulationCfg + +class ArticulationTorch(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, ActuatorBaseWarp] + """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) -> ArticulationDataTorch: + 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._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_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 + + """ + Operations. + """ + + def reset(self, mask: wp.array): + # 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) + + """ + 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. + """ + 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, 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. + """ + 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, 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, 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, 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,). + """ + self._root.write_root_state_to_sim(root_state, env_mask) + + 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,). + """ + self._root.write_root_com_state_to_sim(root_state, env_mask) + + 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,). + """ + self._root.write_root_link_state_to_sim(root_state, 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._root.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,). + """ + self._root.write_root_link_pose_to_sim(pose, env_mask) + + 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,). + """ + self._root.write_root_com_pose_to_sim(root_pose, env_mask) + + 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._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, 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._root.write_root_com_velocity_to_sim(root_velocity, env_mask) + + 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,). + """ + self._root.write_root_link_velocity_to_sim(root_velocity, env_mask) + + 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._joint.write_joint_state_to_sim(position, velocity, joint_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,). + """ + self._joint.write_joint_position_to_sim(position, joint_mask, env_mask) + + 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,). + """ + 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_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. + """ + 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_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,). + """ + self._joint.write_joint_stiffness_to_sim(stiffness, joint_mask, env_mask) + + 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,). + """ + 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_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,). + """ + 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_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,). + """ + 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_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,). + """ + 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_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,). + """ + 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_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,). + """ + 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_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,). + """ + self._body.set_external_force_and_torque(forces, torques, body_mask, env_mask) + + 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,). + """ + self._joint.set_joint_position_target(target, joint_mask, env_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,). + """ + self._joint.set_joint_velocity_target(target, joint_mask, env_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,). + """ + self._joint.set_joint_effort_target(target, joint_mask, env_mask) + + """ + Operations - Tendons. + """ + + def set_fixed_tendon_stiffness( + self, + stiffness: wp.array, + 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 (num_instances, num_fixed_tendons). + 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_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 (num_instances, num_fixed_tendons). + 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_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 (num_instances, num_fixed_tendons). + 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_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 (num_instances, num_fixed_tendons). + 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_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 (num_instances, num_fixed_tendons). + 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: wp.array | Sequence[int] | None = None, + env_ids: wp.array | Sequence[int] | 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)). + 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_mask: wp.array | None = None, + env_mask: wp.array | None = None, + ) -> None: + """Write fixed tendon properties into the simulation. + + Args: + 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_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 (num_instances, num_spatial_tendons). + 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_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 (num_instances, num_spatial_tendons). + 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_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 (num_instances, num_spatial_tendons). + 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_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 (num_instances, num_spatial_tendons). + 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_mask: wp.array | None = None, + env_mask: wp.array | None = None, + ) -> None: + """Write spatial tendon properties into the simulation. + + Args: + 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 + 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(".*", "*") + + # Perf implication when filtering fixed joints. --> Removing the joints from the middle. + # May need to copy stuff. --> DoFs? Careful with joint properties.... + self._root_newton_view = NewtonArticulationView( + NewtonManager.get_model(), prim_path, verbose=True, exclude_joint_types=[JointType.FREE, JointType.FIXED] + ) + + # container for data access + self._data = ArticulationDataWarp(self._root_newton_view, self.device) + + # create backend setters for the data containers + self._joint = Joint(self._root_newton_view, self._data.joint_data, self.cfg.soft_joint_pos_limit_factor, self.device) + self._body = Body(self._root_newton_view, self._data.body_data, self.device) + self._root = Root(self._root_newton_view, self._data.root_data, self.device) + + # 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() + + # 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.default_root_pose).clone() + 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 _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.default_root_pose, + self._root._ALL_ENV_MASK, + ], + ) + # 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.default_root_vel, + self._root._ALL_ENV_MASK, + ], + ) + # -- 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.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.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: ActuatorBaseWarpCfg + # 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: ActuatorBaseWarp = actuator_cfg.class_type( + cfg=actuator_cfg, + joint_names=joint_names, + joint_mask=joint_mask, + env_mask=self._joint._ALL_ENV_MASK, + articulation_data=self._data.joint_data, + device=self.device, + ) + # 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_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, ImplicitActuatorWarp): + 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(self.data.joint_stiffness, joint_mask=actuator._joint_mask) + self.write_joint_damping_to_sim(self.data.joint_damping, joint_mask=actuator._joint_mask) + # Sets the control mode for the implicit actuators + self.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_target = self.data.sim_bind_joint_target + self.data.joint_effort_target = self.data.sim_bind_joint_effort + 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_mask=actuator._joint_mask) + self.write_joint_damping_to_sim(0.0, joint_mask=actuator._joint_mask) + # Set the control mode to None when using explicit actuators + self.write_joint_control_mode_to_sim(0, joint_mask=actuator._joint_mask) + # Bind the applied effort to the simulation effort + self.data.applied_effort = self.data.sim_bind_joint_effort + + # 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(): + 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_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 = wp.to_torch(self._data.default_joint_pos)[0] < joint_pos_limits[:, 0] + out_of_range |= wp.to_torch(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") diff --git a/source/isaaclab/isaaclab/assets/articulation_new/articulation_warp.py b/source/isaaclab/isaaclab/assets/articulation_new/articulation_warp.py new file mode 100644 index 00000000000..17fe9ba374d --- /dev/null +++ b/source/isaaclab/isaaclab/assets/articulation_new/articulation_warp.py @@ -0,0 +1,1266 @@ +# 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 +from collections.abc import Sequence +from prettytable import PrettyTable +from typing import TYPE_CHECKING + +import omni.log +import warp as wp +from isaacsim.core.simulation_manager import SimulationManager +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.actuators_warp import ActuatorBaseWarp, ActuatorBaseWarpCfg, ImplicitActuatorWarp +from isaaclab.sim._impl.newton_manager import NewtonManager + +from ..asset_base import AssetBase +from .articulation_data_warp import ArticulationDataWarp +from isaaclab.assets.core.joint_properties.joint import Joint +from isaaclab.assets.core.body_properties.body import Body +from isaaclab.assets.core.root_properties.root import Root +from isaaclab.assets.core.kernels import ( + generate_mask_from_ids, + populate_empty_array, + update_batched_array_with_array_masked, + update_array_with_value, +) + +if TYPE_CHECKING: + from .articulation_cfg import ArticulationWarpCfg + +class ArticulationWarp(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: ArticulationWarpCfg + """Configuration instance for the articulations.""" + + actuators: dict[str, ActuatorBaseWarp] + """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: ArticulationWarpCfg): + """Initialize the articulation. + + Args: + cfg: A configuration instance. + """ + super().__init__(cfg) + + """ + Properties + """ + + @property + def data(self) -> ArticulationDataWarp: + 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._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_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 + + """ + Operations. + """ + + def reset(self, mask: wp.array): + # 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) + + """ + 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. + """ + 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, 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. + """ + 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, 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, 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, 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,). + """ + self._root.write_root_state_to_sim(root_state, env_mask) + + 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,). + """ + self._root.write_root_com_state_to_sim(root_state, env_mask) + + 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,). + """ + self._root.write_root_link_state_to_sim(root_state, 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._root.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,). + """ + self._root.write_root_link_pose_to_sim(pose, env_mask) + + 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,). + """ + self._root.write_root_com_pose_to_sim(root_pose, env_mask) + + 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._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, 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._root.write_root_com_velocity_to_sim(root_velocity, env_mask) + + 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,). + """ + self._root.write_root_link_velocity_to_sim(root_velocity, env_mask) + + 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._joint.write_joint_state_to_sim(position, velocity, joint_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,). + """ + self._joint.write_joint_position_to_sim(position, joint_mask, env_mask) + + 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,). + """ + 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_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. + """ + 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_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,). + """ + self._joint.write_joint_stiffness_to_sim(stiffness, joint_mask, env_mask) + + 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,). + """ + 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_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,). + """ + 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_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,). + """ + 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_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,). + """ + 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_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,). + """ + 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_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,). + """ + 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_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,). + """ + self._body.set_external_force_and_torque(forces, torques, body_mask, env_mask) + + 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,). + """ + self._joint.set_joint_position_target(target, joint_mask, env_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,). + """ + self._joint.set_joint_velocity_target(target, joint_mask, env_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,). + """ + self._joint.set_joint_effort_target(target, joint_mask, env_mask) + + """ + Operations - Tendons. + """ + + def set_fixed_tendon_stiffness( + self, + stiffness: wp.array, + 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 (num_instances, num_fixed_tendons). + 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_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 (num_instances, num_fixed_tendons). + 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_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 (num_instances, num_fixed_tendons). + 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_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 (num_instances, num_fixed_tendons). + 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_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 (num_instances, num_fixed_tendons). + 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: wp.array | Sequence[int] | None = None, + env_ids: wp.array | Sequence[int] | 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)). + 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_mask: wp.array | None = None, + env_mask: wp.array | None = None, + ) -> None: + """Write fixed tendon properties into the simulation. + + Args: + 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_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 (num_instances, num_spatial_tendons). + 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_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 (num_instances, num_spatial_tendons). + 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_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 (num_instances, num_spatial_tendons). + 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_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 (num_instances, num_spatial_tendons). + 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_mask: wp.array | None = None, + env_mask: wp.array | None = None, + ) -> None: + """Write spatial tendon properties into the simulation. + + Args: + 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 + 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(".*", "*") + + # Perf implication when filtering fixed joints. --> Removing the joints from the middle. + # May need to copy stuff. --> DoFs? Careful with joint properties.... + self._root_newton_view = NewtonArticulationView( + NewtonManager.get_model(), prim_path, verbose=True, exclude_joint_types=[JointType.FREE, JointType.FIXED] + ) + + # container for data access + self._data = ArticulationDataWarp(self._root_newton_view, self.device) + + # create backend setters for the data containers + self._joint = Joint(self._root_newton_view, self._data.joint_data, self.cfg.soft_joint_pos_limit_factor, self.device) + self._body = Body(self._root_newton_view, self._data.body_data, self.device) + self._root = Root(self._root_newton_view, self._data.root_data, self.device) + + # 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() + + # 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.default_root_pose).clone() + 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 _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.default_root_pose, + self._root._ALL_ENV_MASK, + ], + ) + # 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.default_root_vel, + self._root._ALL_ENV_MASK, + ], + ) + # -- 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.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.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: ActuatorBaseWarpCfg + # 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: ActuatorBaseWarp = actuator_cfg.class_type( + cfg=actuator_cfg, + joint_names=joint_names, + joint_mask=joint_mask, + env_mask=self._joint._ALL_ENV_MASK, + articulation_data=self._data.joint_data, + device=self.device, + ) + # 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_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, ImplicitActuatorWarp): + 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(self.data.joint_stiffness, joint_mask=actuator._joint_mask) + self.write_joint_damping_to_sim(self.data.joint_damping, joint_mask=actuator._joint_mask) + # Sets the control mode for the implicit actuators + self.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_target = self.data.sim_bind_joint_target + self.data.joint_effort_target = self.data.sim_bind_joint_effort + 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_mask=actuator._joint_mask) + self.write_joint_damping_to_sim(0.0, joint_mask=actuator._joint_mask) + # Set the control mode to None when using explicit actuators + self.write_joint_control_mode_to_sim(0, joint_mask=actuator._joint_mask) + # Bind the applied effort to the simulation effort + self.data.applied_effort = self.data.sim_bind_joint_effort + + # 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(): + 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_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 = wp.to_torch(self._data.default_joint_pos)[0] < joint_pos_limits[:, 0] + out_of_range |= wp.to_torch(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") diff --git a/source/isaaclab/isaaclab/assets/articulation_warp/__init__.py b/source/isaaclab/isaaclab/assets/articulation_warp/__init__.py new file mode 100644 index 00000000000..f324a56339f --- /dev/null +++ b/source/isaaclab/isaaclab/assets/articulation_warp/__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 articulated assets.""" + +from .articulation import ArticulationWarp +from .articulation_cfg import ArticulationWarpCfg +from .articulation_data import ArticulationDataWarp diff --git a/source/isaaclab/isaaclab/assets/articulation_warp/articulation.py b/source/isaaclab/isaaclab/assets/articulation_warp/articulation.py new file mode 100644 index 00000000000..3f9a849cb4e --- /dev/null +++ b/source/isaaclab/isaaclab/assets/articulation_warp/articulation.py @@ -0,0 +1,2036 @@ +# 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 +from collections.abc import Sequence +from prettytable import PrettyTable +from typing import TYPE_CHECKING + +import omni.log +import warp as wp +from isaacsim.core.simulation_manager import SimulationManager +from newton import JointType, Model +from newton.selection import ArticulationView as NewtonArticulationView +from newton.solvers import SolverMuJoCo, SolverNotifyFlags +from pxr import UsdPhysics + +import isaaclab.sim as sim_utils +import isaaclab.utils.string as string_utils +from isaaclab.actuators_warp import ActuatorBaseWarp, ActuatorBaseWarpCfg, ImplicitActuatorWarp +from isaaclab.sim._impl.newton_manager import NewtonManager +from isaaclab.utils.helpers import warn_overhead_cost + +from ..asset_base import AssetBase +from .articulation_data import ArticulationDataWarp +from .kernels import ( + generate_mask_from_ids, + populate_empty_array, + project_link_velocity_to_com_frame_masked, + split_state, + transform_CoM_pose_to_link_frame, + update_joint_array, + update_joint_array_int, + update_joint_array_with_value, + update_joint_array_with_value_array, + 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, + update_spatial_vector_array_with_value, + update_transforms_array, + update_transforms_array_with_value, + update_velocity_array, + update_wrench_array_with_force, + update_wrench_array_with_torque, + update_wrench_array_with_value, +) + +if TYPE_CHECKING: + from .articulation_cfg import ArticulationWarpCfg + +# Composition vs Inheritance + +#class NewCLass(AssetBase): +# def __init__(self, cfg: ArticulationCfg, original: ArticulationWarp): +# super().__init__(cfg) +# self._original = original +# +# def write_data_to_sim(self): +# wrap the shit out of it +# return the wrapped shit. +# self._original.write_data_to_sim() + +# Not making a rigid body class just for it. (bundle multi and single body assets) + +class ArticulationWarp(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: ArticulationWarpCfg + """Configuration instance for the articulations.""" + + actuators: dict[str, ActuatorBaseWarp] + """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: ArticulationWarpCfg): + """Initialize the articulation. + + Args: + cfg: A configuration instance. + """ + super().__init__(cfg) + + """ + Properties + """ + + @property + def data(self) -> ArticulationDataWarp: + 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 + + """ + Operations. + """ + + def reset(self, mask: wp.array): + # use ellipses object to skip initial indices. + if mask is None: + mask = self._ALL_ENV_MASK + + # reset actuators + for actuator in self.actuators.values(): + actuator.reset(mask) + # reset external wrench + wp.launch( + update_wrench_array_with_value, + dim=(self.num_instances,), + inputs=[ + wp.spatial_vectorf(0.0, 0.0, 0.0, 0.0, 0.0, 0.0), + self._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. + # apply actuator models + self._apply_actuator_model() + + 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 + + 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 + + def find_fixed_tendons( + self, name_keys: str | Sequence[str], tendon_subsets: list[str] | None = None, preserve_order: bool = False + ) -> tuple[wp.array, 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. + """ + if tendon_subsets is None: + # tendons follow the joint names they are attached to + tendon_subsets = self.fixed_tendon_names + # find tendons + indices, names = string_utils.resolve_matching_names(name_keys, tendon_subsets, preserve_order) + self._FIXED_TENDON_MASK.fill_(False) + mask = wp.clone(self._FIXED_TENDON_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 + + def find_spatial_tendons( + self, name_keys: str | Sequence[str], tendon_subsets: list[str] | None = None, preserve_order: bool = False + ) -> tuple[wp.array, 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. + """ + if tendon_subsets is None: + tendon_subsets = self.spatial_tendon_names + # find tendons + indices, names = string_utils.resolve_matching_names(name_keys, tendon_subsets, preserve_order) + self._SPATIAL_TENDON_MASK.fill_(False) + mask = wp.clone(self._SPATIAL_TENDON_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 - State Writers. + """ + + @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_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 + + # set into simulation + target_root_pose = wp.zeros((self.num_instances), dtype=wp.transformf, device=self.device) + target_root_velocity = wp.zeros((self.num_instances), dtype=wp.spatial_vectorf, device=self.device) + + wp.launch( + split_state, + dim=(self.num_instances,), + inputs=[ + root_state, + target_root_pose, + target_root_velocity, + env_mask, + ], + ) + self.write_root_link_pose_to_sim(target_root_pose, env_mask=env_mask) + self.write_root_com_velocity_to_sim(target_root_velocity, 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 + + target_root_pose = wp.zeros((self.num_instances), dtype=wp.transformf, device=self.device) + target_root_velocity = wp.zeros((self.num_instances), dtype=wp.spatial_vectorf, device=self.device) + + wp.launch( + split_state, + dim=(self.num_instances,), + inputs=[ + root_state, + target_root_pose, + target_root_velocity, + env_mask, + ], + ) + self.write_root_com_pose_to_sim(target_root_pose, env_mask=env_mask) + self.write_root_com_velocity_to_sim(target_root_velocity, 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 + + target_root_pose = wp.zeros((self.num_instances), dtype=wp.transformf, device=self.device) + target_root_velocity = wp.zeros((self.num_instances), dtype=wp.spatial_vectorf, device=self.device) + + wp.launch( + split_state, + dim=(self.num_instances,), + inputs=[ + root_state, + target_root_pose, + target_root_velocity, + env_mask, + ], + ) + self.write_root_link_pose_to_sim(target_root_pose, env_mask=env_mask) + self.write_root_link_velocity_to_sim(target_root_velocity, 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_transforms_array, + dim=self._root_newton_view.count, + inputs=[ + pose, + self._data.sim_bind_root_link_pose_w, + env_mask, + ], + ) + # Need to invalidate the buffer to trigger the update with the new state. + self._data._root_com_pose_w.timestamp = -1.0 + self._data._body_com_pose_w.timestamp = -1.0 + + 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_transforms_array, + dim=self._root_newton_view.count, + inputs=[ + root_pose, + self._data._root_com_pose_w.data, + env_mask, + ], + ) + # set link frame poses + wp.launch( + transform_CoM_pose_to_link_frame, + dim=self._root_newton_view.count, + inputs=[ + self._data.root_com_pose_w, + self._data.body_com_pos_b, + self._data.sim_bind_root_link_pose_w, + env_mask, + ], + ) + self._data._body_com_pose_w.timestamp = -1.0 + + 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_velocity_array, + dim=self._root_newton_view.count, + inputs=[ + root_velocity, + self._data.sim_bind_root_com_vel_w, + env_mask, + ], + ) + self._data._root_link_vel_w.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 + # set into internal buffers + wp.launch( + project_link_velocity_to_com_frame_masked, + dim=self._root_newton_view.count, + inputs=[ + root_velocity, + self._data.body_link_pose_w, + self._data.body_com_pos_b, + self._data.sim_bind_root_com_vel_w, + env_mask, + ], + ) + + 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.sim_bind_joint_pos, + env_mask, + joint_mask, + ], + ) + # invalidate buffers to trigger the update with the new root pose. + self._data._body_com_pose_w.timestamp = -1.0 + + 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.sim_bind_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, + ], + ) + # Need to invalidate the buffer to trigger the update with the new root pose. + self._data._body_link_vel_w.timestamp = -1.0 + + """ + 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.sim_bind_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.sim_bind_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.sim_bind_joint_stiffness_sim, + env_mask, + joint_mask, + ], + ) + else: + wp.launch( + update_joint_array, + dim=(self.num_instances, self.num_joints), + inputs=[ + stiffness, + self._data.sim_bind_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.sim_bind_joint_damping_sim, + env_mask, + joint_mask, + ], + ) + else: + wp.launch( + update_joint_array, + dim=(self.num_instances, self.num_joints), + inputs=[ + damping, + self._data.sim_bind_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.sim_bind_joint_pos_limits_lower, + self._data.sim_bind_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.sim_bind_joint_pos_limits_lower, + self._data.sim_bind_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.sim_bind_joint_vel_limits_sim, + env_mask, + joint_mask, + ], + ) + else: + wp.launch( + update_joint_array, + dim=(self.num_instances, self.num_joints), + inputs=[ + limits, + self._data.sim_bind_joint_vel_limits_sim, + 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.sim_bind_joint_effort_limits_sim, + env_mask, + joint_mask, + ], + ) + else: + wp.launch( + update_joint_array, + dim=(self.num_instances, self.num_joints), + inputs=[ + limits, + self._data.sim_bind_joint_effort_limits_sim, + 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.sim_bind_joint_armature, + env_mask, + joint_mask, + ], + ) + else: + wp.launch( + update_joint_array, + dim=(self.num_instances, self.num_joints), + inputs=[ + armature, + self._data.sim_bind_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.sim_bind_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.sim_bind_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_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._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._external_wrench, + env_mask, + body_mask, + ], + ) + + 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, + ], + ) + + """ + Operations - Tendons. + """ + + def set_fixed_tendon_stiffness( + self, + stiffness: wp.array, + 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 (num_instances, num_fixed_tendons). + 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_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 (num_instances, num_fixed_tendons). + 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_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 (num_instances, num_fixed_tendons). + 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_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 (num_instances, num_fixed_tendons). + 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_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 (num_instances, num_fixed_tendons). + 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: wp.array | Sequence[int] | None = None, + env_ids: wp.array | Sequence[int] | 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)). + 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_mask: wp.array | None = None, + env_mask: wp.array | None = None, + ) -> None: + """Write fixed tendon properties into the simulation. + + Args: + 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_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 (num_instances, num_spatial_tendons). + 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_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 (num_instances, num_spatial_tendons). + 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_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 (num_instances, num_spatial_tendons). + 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_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 (num_instances, num_spatial_tendons). + 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_mask: wp.array | None = None, + env_mask: wp.array | None = None, + ) -> None: + """Write spatial tendon properties into the simulation. + + Args: + 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 + 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(".*", "*") + + # Perf implication when filtering fixed joints. --> Removing the joints from the middle. + # May need to copy stuff. --> DoFs? Careful with joint properties.... + self._root_newton_view = NewtonArticulationView( + NewtonManager.get_model(), prim_path, verbose=True, exclude_joint_types=[JointType.FREE, JointType.FIXED] + ) + + # container for data access + self._data = ArticulationDataWarp(self._root_newton_view, self.device) + + # create simulation bindings and buffers + self._create_simulation_bindings() + 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() + # 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.default_root_pose).clone() + 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_simulation_bindings(self): + """Create simulation bindings for the articulation. + + 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. + """ + + # TODO Check if these are pointers or copies. Strided array vs indexed array. + # TODO add a fallback if the array is not strided. + # TODO Send warning to users if arrays are no strided. WRT potential performance implications. + + # TODO Check if zero copy, or not, this will be per frequency (body, joint, join_dof, etc.) + # TODO make a function that will return the cached value, use a function pointer. + + #TODO Try without this, and use the setters and getters and directly. + + # -- external forces and torques + self._external_wrench = self._root_newton_view.get_attribute("body_f", NewtonManager.get_state_0()) + # -- root properties + self._data.sim_bind_root_link_pose_w = self._root_newton_view.get_root_transforms(NewtonManager.get_state_0()) + self._data.sim_bind_root_com_vel_w = self._root_newton_view.get_root_velocities(NewtonManager.get_state_0()) + # -- body properties + self._data.sim_bind_body_link_pose_w = self._root_newton_view.get_link_transforms(NewtonManager.get_state_0()) + self._data.sim_bind_body_com_vel_w = self._root_newton_view.get_link_velocities(NewtonManager.get_state_0()) + self._data.sim_bind_body_com_pos_b = self._root_newton_view.get_attribute("body_com", NewtonManager.get_model()) + # -- joint properties + self._data.sim_bind_joint_pos_limits_lower = self._root_newton_view.get_attribute( + "joint_limit_lower", NewtonManager.get_model() + ) + self._data.sim_bind_joint_pos_limits_upper = self._root_newton_view.get_attribute( + "joint_limit_upper", NewtonManager.get_model() + ) + self._data.sim_bind_joint_stiffness_sim = self._root_newton_view.get_attribute( + "joint_target_ke", NewtonManager.get_model() + ) + self._data.sim_bind_joint_damping_sim = self._root_newton_view.get_attribute( + "joint_target_kd", NewtonManager.get_model() + ) + self._data.sim_bind_joint_armature = self._root_newton_view.get_attribute( + "joint_armature", NewtonManager.get_model() + ) + self._data.sim_bind_joint_friction_coeff = self._root_newton_view.get_attribute( + "joint_friction", NewtonManager.get_model() + ) + self._data.sim_bind_joint_vel_limits_sim = self._root_newton_view.get_attribute( + "joint_velocity_limit", NewtonManager.get_model() + ) + self._data.sim_bind_joint_effort_limits_sim = self._root_newton_view.get_attribute( + "joint_effort_limit", NewtonManager.get_model() + ) + self._data.sim_bind_joint_control_mode_sim = self._root_newton_view.get_attribute( + "joint_dof_mode", NewtonManager.get_model() + ) + # -- joint states + self._data.sim_bind_joint_pos = self._root_newton_view.get_dof_positions(NewtonManager.get_state_0()) + self._data.sim_bind_joint_vel = self._root_newton_view.get_dof_velocities(NewtonManager.get_state_0()) + # -- joint commands (sent to the simulation) + self._data.sim_bind_joint_effort = self._root_newton_view.get_attribute("joint_f", NewtonManager.get_control()) + self._data.sim_bind_joint_target = self._root_newton_view.get_attribute( + "joint_target", NewtonManager.get_control() + ) + + 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) + self._ALL_JOINT_MASK = wp.ones((self.num_joints,), dtype=wp.bool, device=self.device) + self._ALL_FIXED_TENDON_MASK = wp.ones((self.num_fixed_tendons,), dtype=wp.bool, device=self.device) + self._ALL_SPATIAL_TENDON_MASK = wp.ones((self.num_spatial_tendons,), 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) + self._JOINT_MASK = wp.zeros((self.num_joints,), dtype=wp.bool, device=self.device) + self._FIXED_TENDON_MASK = wp.zeros((self.num_fixed_tendons,), dtype=wp.bool, device=self.device) + self._SPATIAL_TENDON_MASK = wp.zeros((self.num_spatial_tendons,), dtype=wp.bool, device=self.device) + # asset named data + self._data.joint_names = self.joint_names + self._data.body_names = self.body_names + # -- joint commands (sent to the actuator from the user) + self._data.joint_target = wp.zeros((self.num_instances, self.num_joints), dtype=wp.float32, device=self.device) + self._data.joint_effort_target = wp.zeros( + (self.num_instances, self.num_joints), dtype=wp.float32, device=self.device + ) + # -- computed joint efforts from the actuator models + self._data.computed_effort = wp.zeros( + (self.num_instances, self.num_joints), dtype=wp.float32, device=self.device + ) + self._data.applied_effort = wp.zeros( + (self.num_instances, self.num_joints), dtype=wp.float32, device=self.device + ) + self._data.joint_stiffness = wp.clone(self._data.sim_bind_joint_stiffness_sim) + self._data.joint_damping = wp.clone(self._data.sim_bind_joint_damping_sim) + self._data.joint_control_mode = wp.clone(self._data.sim_bind_joint_control_mode_sim) + # -- other data that are filled based on explicit actuator models + self._data.joint_dynamic_friction = wp.zeros( + (self.num_instances, self.num_joints), dtype=wp.float32, device=self.device + ) + self._data.joint_viscous_friction = wp.zeros( + (self.num_instances, self.num_joints), dtype=wp.float32, device=self.device + ) + self._data.soft_joint_vel_limits = wp.zeros( + (self.num_instances, self.num_joints), dtype=wp.float32, device=self.device + ) + self._data.gear_ratio = wp.ones((self.num_instances, self.num_joints), dtype=wp.float32, device=self.device) + # -- update the soft joint position limits + self._data.soft_joint_pos_limits = wp.zeros( + (self.num_instances, self.num_joints), dtype=wp.vec2f, device=self.device + ) + wp.launch( + update_soft_joint_pos_limits, + dim=(self.num_instances, self.num_joints), + inputs=[ + self._data.sim_bind_joint_pos_limits_lower, + self._data.sim_bind_joint_pos_limits_upper, + self._data.soft_joint_pos_limits, + self.cfg.soft_joint_pos_limit_factor, + ], + ) + + def _process_cfg(self): + """Post processing of configuration parameters.""" + # -- root state + self._data.default_root_pose = wp.zeros((self.num_instances), dtype=wp.transformf, device=self.device) + self._data.default_root_vel = wp.zeros((self.num_instances), dtype=wp.spatial_vectorf, device=self.device) + # 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_transforms_array_with_value, + dim=(self.num_instances,), + inputs=[ + wp.transformf(*default_root_pose), + self._data.default_root_pose, + self._ALL_ENV_MASK, + ], + ) + # default velocity + default_root_velocity = tuple(self.cfg.init_state.lin_vel) + tuple(self.cfg.init_state.ang_vel) + wp.launch( + update_spatial_vector_array_with_value, + dim=(self.num_instances,), + inputs=[ + wp.spatial_vectorf(*default_root_velocity), + self._data.default_root_vel, + self._ALL_ENV_MASK, + ], + ) + # -- joint state + self._data.default_joint_pos = wp.zeros( + (self.num_instances, self.num_joints), dtype=wp.float32, device=self.device + ) + self._data.default_joint_vel = wp.zeros( + (self.num_instances, self.num_joints), dtype=wp.float32, device=self.device + ) + # -- 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_MASK.fill_(False) + wp.launch( + generate_mask_from_ids, + dim=(self.num_joints,), + inputs=[ + self._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_joint_array_with_value_array, + dim=(self.num_instances, self.num_joints), + inputs=[ + tmp_joint_data, + self._data.default_joint_pos, + self._ALL_ENV_MASK, + self._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_joint_array_with_value_array, + dim=(self.num_instances, self.num_joints), + inputs=[ + tmp_joint_data, + self._data.default_joint_vel, + self._ALL_ENV_MASK, + self._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: ActuatorBaseWarpCfg + # 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: ActuatorBaseWarp = actuator_cfg.class_type( + cfg=actuator_cfg, + joint_names=joint_names, + joint_mask=joint_mask, + env_mask=self._ALL_ENV_MASK, + articulation_data=self._data, + device=self.device, + ) + # 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_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, ImplicitActuatorWarp): + 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(self.data.joint_stiffness, joint_mask=actuator._joint_mask) + self.write_joint_damping_to_sim(self.data.joint_damping, joint_mask=actuator._joint_mask) + # Sets the control mode for the implicit actuators + self.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_target = self.data.sim_bind_joint_target + self.data.joint_effort_target = self.data.sim_bind_joint_effort + 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_mask=actuator._joint_mask) + self.write_joint_damping_to_sim(0.0, joint_mask=actuator._joint_mask) + # Set the control mode to None when using explicit actuators + self.write_joint_control_mode_to_sim(0, joint_mask=actuator._joint_mask) + # Bind the applied effort to the simulation effort + self.data.applied_effort = self.data.sim_bind_joint_effort + + # 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(): + 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_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 = wp.to_torch(self._data.default_joint_pos)[0] < joint_pos_limits[:, 0] + out_of_range |= wp.to_torch(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") diff --git a/source/isaaclab/isaaclab/assets/articulation_warp/articulation_cfg.py b/source/isaaclab/isaaclab/assets/articulation_warp/articulation_cfg.py new file mode 100644 index 00000000000..5ef675ad22d --- /dev/null +++ b/source/isaaclab/isaaclab/assets/articulation_warp/articulation_cfg.py @@ -0,0 +1,61 @@ +# 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.actuators_warp import ActuatorBaseWarpCfg +from isaaclab.utils import configclass + +from ..asset_base_cfg import AssetBaseCfg +from .articulation import ArticulationWarp + + +@configclass +class ArticulationWarpCfg(AssetBaseCfg): + """Configuration parameters for an articulation.""" + + @configclass + class InitialStateCfg(AssetBaseCfg.InitialStateCfg): + """Initial state of the articulation.""" + + # root velocity + 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).""" + + # joint state + joint_pos: dict[str, float] = {".*": 0.0} + """Joint positions of the joints. Defaults to 0.0 for all joints.""" + joint_vel: dict[str, float] = {".*": 0.0} + """Joint velocities of the joints. Defaults to 0.0 for all joints.""" + + ## + # Initialize configurations. + ## + + class_type: type = ArticulationWarp + + articulation_root_prim_path: str | None = None + """Path to the articulation root prim in the USD file. + + If not provided will search for a prim with the ArticulationRootAPI. Should start with a slash. + """ + + init_state: InitialStateCfg = InitialStateCfg() + """Initial state of the articulated object. Defaults to identity pose with zero velocity and zero joint state.""" + + soft_joint_pos_limit_factor: float = 1.0 + """Fraction specifying the range of joint position limits (parsed from the asset) to use. Defaults to 1.0. + + The soft joint position limits are scaled by this factor to specify a safety region within the simulated + joint position limits. This isn't used by the simulation, but is useful for learning agents to prevent the joint + positions from violating the limits, such as for termination conditions. + + The soft joint position limits are accessible through the :attr:`ArticulationData.soft_joint_pos_limits` attribute. + """ + + actuators: dict[str, ActuatorBaseWarpCfg] = MISSING + """Actuators for the robot with corresponding joint names.""" diff --git a/source/isaaclab/isaaclab/assets/articulation_warp/articulation_data.py b/source/isaaclab/isaaclab/assets/articulation_warp/articulation_data.py new file mode 100644 index 00000000000..b1bc664a6fa --- /dev/null +++ b/source/isaaclab/isaaclab/assets/articulation_warp/articulation_data.py @@ -0,0 +1,1544 @@ +# 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 isaaclab.utils.buffers import TimestampedWarpBuffer +from isaaclab.utils.helpers import deprecated, warn_overhead_cost + +from isaaclab.assets.core.kernels import ( + combine_frame_transforms_partial, + combine_frame_transforms_partial_batch, + combine_pose_and_velocity_to_state, + combine_pose_and_velocity_to_state_batched, + compute_heading, + derive_body_acceleration_from_velocity, + derive_joint_acceleration_from_velocity, + generate_pose_from_position_with_unit_quaternion, + get_angular_velocity, + get_linear_velocity, + make_joint_pos_limits_from_lower_and_upper_limits, + project_com_velocity_to_link_frame_batch, + project_vec_from_quat_single, + project_velocities_to_frame, + vec13f, +) + + +class ArticulationDataWarp: + """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 = 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)) + + # Initialize history for finite differencing + self._previous_body_com_vel = wp.clone(self._root_newton_view.get_link_velocities(NewtonManager.get_state_0())) + self._previous_joint_vel = wp.clone(self._root_newton_view.get_dof_velocities(NewtonManager.get_state_0())) + + # 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._body_link_vel_w = TimestampedWarpBuffer( + shape=(self._root_newton_view.count, self._root_newton_view.link_count), dtype=wp.spatial_vectorf + ) + self._projected_gravity_b = TimestampedWarpBuffer(shape=(self._root_newton_view.count, 3), 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._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 + ) + # -- 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() + + def get_from_newton(self, name: str, container): + return self._root_newton_view.get_attribute(name, container) + + 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 + self.body_com_acc_w + + ## + # 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_pose: wp.array = None + """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. + """ + + default_root_vel: wp.array = None + """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. + """ + + default_joint_pos: 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: 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. + """ + + # TODO: Make them protected with an _ prefix. + A blurb in the doc: DON'T FUCK WITH THESE. + ## + # Root state properties. (Directly binded to the simulation) + ## + + sim_bind_root_link_pose_w: wp.array = None + """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. + """ + + sim_bind_root_com_vel_w: wp.array = None + """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]. + """ + + ## + # Body state properties. (Directly binded to the simulation) + ## + + sim_bind_body_link_pose_w: wp.array = None + """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. + """ + + sim_bind_body_com_vel_w: wp.array = None + """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]. + """ + + sim_bind_body_com_pos_b: wp.array = None + """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. + """ + + ## + # Joint commands -- From the user to the actuator model. + ## + + joint_target: 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 efforts (see :attr:`applied_torque`), + which are then set into the simulation. + """ + + joint_effort_target: 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 efforts (see :attr:`applied_torque`), + which are then set into the simulation. + """ + + ## + # Joint commands -- Explicit actuators. + ## + + computed_effort: wp.array = None + """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. + """ + + applied_effort: wp.array = None + """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. + """ + + joint_stiffness: wp.array = None + """Joint stiffness. Shape is (num_instances, num_joints).""" + + joint_damping: wp.array = None + """Joint damping. Shape is (num_instances, num_joints).""" + + joint_control_mode: wp.array = None + """Joint control mode. Shape is (num_instances, num_joints).""" + + ### + # Joint commands. (Directly binded to the simulation) + ### + + sim_bind_joint_target: wp.array = None + """Joint target. Shape is (num_instances, num_joints).""" + + sim_bind_joint_effort: wp.array = None + """Joint effort. Shape is (num_instances, num_joints).""" + + ## + # Joint properties. (Directly binded to the simulation) + ## + + sim_bind_joint_control_mode_sim: wp.array = 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. + """ + + sim_bind_joint_stiffness_sim: 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. + """ + + sim_bind_joint_damping_sim: 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. + """ + + sim_bind_joint_armature: wp.array = None + """Joint armature provided to the simulation. Shape is (num_instances, num_joints).""" + + sim_bind_joint_friction_coeff: wp.array = None + """Joint friction coefficient provided to the simulation. Shape is (num_instances, num_joints).""" + + sim_bind_joint_pos_limits_lower: wp.array = None + """Joint position limits lower provided to the simulation. Shape is (num_instances, num_joints).""" + + sim_bind_joint_pos_limits_upper: wp.array = None + """Joint position limits upper provided to the simulation. Shape is (num_instances, num_joints).""" + + sim_bind_joint_vel_limits_sim: wp.array = None + """Joint maximum velocity provided to the simulation. Shape is (num_instances, num_joints).""" + + sim_bind_joint_effort_limits_sim: wp.array = None + """Joint maximum effort provided to the simulation. Shape is (num_instances, num_joints).""" + + ## + # Joint states + ## + + sim_bind_joint_pos: wp.array = None + """Joint positions. Shape is (num_instances, num_joints).""" + + sim_bind_joint_vel: wp.array = None + """Joint velocities. Shape is (num_instances, num_joints).""" + + ## + # Joint properties - Custom. + ## + + joint_effort_limits: wp.array = None + """Joint effort limits. Shape is (num_instances, num_joints).""" + + joint_vel_limits: wp.array = None + """Joint velocity limits. Shape is (num_instances, num_joints).""" + + joint_dynamic_friction: wp.array = None + """Joint dynamic friction. Shape is (num_instances, num_joints).""" + + joint_viscous_friction: wp.array = None + """Joint viscous friction. Shape is (num_instances, num_joints).""" + + soft_joint_pos_limits: 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: 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: 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: wp.array = None + """Fixed tendon stiffness provided to the simulation. Shape is (num_instances, num_fixed_tendons).""" + + fixed_tendon_damping: wp.array = None + """Fixed tendon damping provided to the simulation. Shape is (num_instances, num_fixed_tendons).""" + + fixed_tendon_limit_stiffness: wp.array = None + """Fixed tendon limit stiffness provided to the simulation. Shape is (num_instances, num_fixed_tendons).""" + + fixed_tendon_rest_length: wp.array = None + """Fixed tendon rest length provided to the simulation. Shape is (num_instances, num_fixed_tendons).""" + + fixed_tendon_offset: wp.array = None + """Fixed tendon offset provided to the simulation. Shape is (num_instances, num_fixed_tendons).""" + + fixed_tendon_pos_limits: 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: wp.array = None + """Spatial tendon stiffness provided to the simulation. Shape is (num_instances, num_spatial_tendons).""" + + spatial_tendon_damping: wp.array = None + """Spatial tendon damping provided to the simulation. Shape is (num_instances, num_spatial_tendons).""" + + spatial_tendon_limit_stiffness: wp.array = None + """Spatial tendon limit stiffness provided to the simulation. Shape is (num_instances, num_spatial_tendons).""" + + spatial_tendon_offset: wp.array = None + """Spatial tendon offset provided to the simulation. Shape is (num_instances, num_spatial_tendons).""" + + ## + # Direct simulation bindings accessors. + ## + + @property + def root_link_pose_w(self) -> wp.array: + return self.sim_bind_root_link_pose_w + + @property + def root_com_vel_w(self) -> wp.array: + return self.sim_bind_root_com_vel_w + + @property + def body_link_pose_w(self) -> wp.array: + return self.sim_bind_body_link_pose_w + + @property + def body_com_vel_w(self) -> wp.array: + return self.sim_bind_body_com_vel_w + + @property + def body_com_pos_b(self) -> wp.array: + return self.sim_bind_body_com_pos_b + + @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: + return self.sim_bind_joint_stiffness_sim + + @property + def joint_damping_sim(self) -> wp.array: + return self.sim_bind_joint_damping_sim + + @property + def joint_armature(self) -> wp.array: + return self.sim_bind_joint_armature + + @property + def joint_friction_coeff(self) -> wp.array: + return self.sim_bind_joint_friction_coeff + + @property + def joint_pos_limits_lower(self) -> wp.array: + return self.sim_bind_joint_pos_limits_lower + + @property + def joint_pos_limits_upper(self) -> wp.array: + return self.sim_bind_joint_pos_limits_upper + + @property + def joint_vel_limits_sim(self) -> wp.array: + return self.sim_bind_joint_vel_limits_sim + + @property + def joint_effort_limits_sim(self) -> wp.array: + return self.sim_bind_joint_effort_limits_sim + + @property + def joint_pos(self) -> wp.array: + return self.sim_bind_joint_pos + + @property + def joint_vel(self) -> wp.array: + return self.sim_bind_joint_vel + + @property + def joint_target_sim(self) -> wp.array: + return self.sim_bind_joint_target + + @property + def joint_effort_sim(self) -> wp.array: + return self.sim_bind_joint_effort + + ## + # Root state properties. + ## + + @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_batch, + dim=(self._root_newton_view.count), + device=self.device, + inputs=[ + self.root_com_vel_w, + self.body_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. + The orientation is provided in (x, y, z, w) format. + """ + 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 ``[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. + """ + 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 ``[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. + """ + 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 ``[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. + """ + 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 + + ## + # Body state properties. + ## + + @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 ``[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. + """ + + state = wp.zeros( + (self._root_newton_view.count, self._root_newton_view.link_count, 13), dtype=wp.float32, 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 ``[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. + """ + + state = wp.zeros( + (self._root_newton_view.count, self._root_newton_view.link_count, 13), dtype=wp.float32, 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 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. + """ + + state = wp.zeros( + (self._root_newton_view.count, self._root_newton_view.link_count, 13), dtype=wp.float32, 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,). The acceleration is in the form of [wx, wy, wz, vx, vy, vz]. + All values are relative to the world. + """ + if self._body_com_acc_w.timestamp < self._sim_timestamp: + wp.launch( + derive_body_acceleration_from_velocity, + 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 [1, 0, 0, 0]. 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. + The orientation is provided in (x, y, z, w) format. + """ + out = wp.zeros( + (self._root_newton_view.count, 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, self._root_newton_view.link_count), + inputs=[ + self.body_com_pos_b, + out, + ], + ) + return out + + @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 + + ## + # 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 + + ## + # Derived Properties. + ## + + # FIXME: USE SIM_BIND_LINK_POSE_W RATHER THAN JUST THE QUATERNION + @property + def projected_gravity_b(self): + """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): + """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). + + Velocity is provided in the form of [wx, wy, wz, vx, vy, vz]. + """ + 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). + + Velocity is provided in the form of [wx, wy, wz, vx, vy, vz]. + """ + 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 + + ## + # Sliced properties. + ## + + @property + @warn_overhead_cost( + "root_link_pose_w", + "This function extracts the position from the pose (containing both position and orientation). Consider using" + " the whole of the pose instead.", + ) + 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. + """ + return self.sim_bind_root_link_pose_w[:, :3] + + @property + @warn_overhead_cost( + "root_link_pose_w", + "This function extracts the position from the pose (containing both position and orientation). Consider using" + " the whole of the pose instead.", + ) + def root_link_quat_w(self) -> wp.array: + """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.sim_bind_root_link_pose_w[:, 3:] + + @property + @warn_overhead_cost( + "root_link_vel_w", + "This function extracts the linear velocity from the velocities (containing both linear and angular" + " velocities). Consider using the whole of the velocities instead.", + ) + 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( + get_linear_velocity, + dim=self._root_newton_view.count, + inputs=[ + self.root_link_vel_w, + out, + ], + ) + return out + + @property + @warn_overhead_cost( + "root_link_vel_w", + "This function extracts the angular velocity from the velocities (containing both linear and angular" + " velocities). Consider using the whole of the velocities instead.", + ) + 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( + get_angular_velocity, + dim=self._root_newton_view.count, + inputs=[ + self.root_link_vel_w, + out, + ], + ) + return out + + @property + @warn_overhead_cost( + "root_com_pose_w", + "This function extracts the position from the pose (containing both position and orientation). Consider using" + " the whole of the pose instead.", + ) + 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. + """ + return self.root_com_pose_w[:, :3] + + @property + @warn_overhead_cost( + "root_com_pose_w", + "This function extracts the orientation from the pose (containing both position and orientation). Consider" + " using the whole of the pose instead.", + ) + def root_com_quat_w(self) -> wp.array: + """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.root_com_pose_w[:, 3:] + + @property + @warn_overhead_cost( + "root_com_vel_w", + "This function extracts the linear velocity from the velocities (containing both linear and angular" + " velocities). Consider using the whole of the velocities instead.", + ) + 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( + get_linear_velocity, + 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", + "This function extracts the angular velocity from the velocities (containing both linear and angular" + " velocities). Consider using the whole of the velocities instead.", + ) + 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( + get_angular_velocity, + dim=self._root_newton_view.count, + inputs=[ + self.sim_bind_root_com_vel_w, + out, + ], + ) + return out + + @property + @warn_overhead_cost( + "body_link_pose_w", + "This function extracts the position from the pose (containing both position and orientation). 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. + """ + return self.body_link_pose_w[:, :, :3] + + @property + @warn_overhead_cost( + "body_link_pose_w", + "This function extracts the orientation from the pose (containing both position and orientation). 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). + + Format is ``(w, x, y, z)``. + This quantity is the orientation of the articulation bodies' actor frame relative to the world. + """ + return self.body_link_pose_w[:, :, 3:] + + @property + @warn_overhead_cost( + "body_link_vel_w", + "This function extracts the linear velocity from the velocities (containing both linear and angular" + " velocities). Consider using the whole of the velocities 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. + """ + return self.body_link_vel_w[:, :, 3:] + + @property + @warn_overhead_cost( + "body_link_vel_w", + "This function extracts the angular velocity from the velocities (containing both linear and angular" + " velocities). Consider using the whole of the velocities 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( + get_angular_velocity, + 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( + "body_com_pose_w", + "This function extracts the position from the pose (containing both position and orientation). 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. + """ + return self.body_com_pose_w[:, :, 3:] + + @property + @warn_overhead_cost( + "body_com_pose_w", + "This function extracts the orientation from the pose (containing both position and orientation). 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). + + Format is ``(w, x, y, z)``. + This quantity is the orientation of the articulation bodies' actor frame. + """ + return self.body_com_vel_w[:, :, 3:] + + @property + @warn_overhead_cost( + "body_com_vel_w", + "This function extracts the linear velocity from the velocities (containing both linear and angular" + " velocities). Consider using the whole of the velocities 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( + get_linear_velocity, + 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( + "body_com_vel_w", + "This function extracts the angular velocity from the velocities (containing both linear and angular" + " velocities). Consider using the whole of the velocities 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( + get_angular_velocity, + 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( + "body_com_acc_w", + "This function extracts the linear acceleration from the accelerations (containing both linear and angular" + " accelerations). Consider using the whole of the accelerations 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( + get_linear_velocity, + 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( + "body_com_acc_w", + "This function extracts the angular acceleration from the accelerations (containing both linear and angular" + " accelerations). Consider using the whole of the accelerations 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( + get_angular_velocity, + 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( + "body_com_pose_b", + "This function extracts the orientation from the pose (containing both position and orientation). Consider" + " using the whole of the pose instead.", + ) + def body_com_quat_b(self) -> wp.array: + """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.body_com_pose_w[:, :, 3:] + + @property + @warn_overhead_cost( + "root_link_vel_b", + "This function extracts the linear velocity from the velocities (containing both linear and angular" + " velocities). Consider using the whole of the velocities instead.", + ) + 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( + get_linear_velocity, + dim=self._root_newton_view.count, + inputs=[ + self.root_link_vel_b, + out, + ], + ) + return out + + @property + @warn_overhead_cost( + "root_link_vel_b", + "This function extracts the angular velocity from the velocities (containing both linear and angular" + " velocities). Consider using the whole of the velocities instead.", + ) + 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( + get_angular_velocity, + dim=self._root_newton_view.count, + inputs=[ + self.root_link_vel_b, + out, + ], + ) + return out + + @property + @warn_overhead_cost( + "root_com_vel_b", + "This function extracts the linear velocity from the velocities (containing both linear and angular" + " velocities). Consider using the whole of the velocities instead.", + ) + 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( + get_linear_velocity, + dim=self._root_newton_view.count, + inputs=[ + self.root_com_vel_b, + out, + ], + ) + return out + + @property + @warn_overhead_cost( + "root_com_vel_b", + "This function extracts the angular velocity from the velocities (containing both linear and angular" + " velocities). Consider using the whole of the velocities instead.", + ) + 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( + get_angular_velocity, + dim=self._root_newton_view.count, + inputs=[ + self.root_com_vel_b, + out, + ], + ) + return out + + @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.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.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("default_root_pose") + def default_root_state(self) -> wp.array: + """Same as :attr:`default_root_pose`.""" + state = wp.zeros((self._root_newton_view.count), dtype=vec13f, device=self.device) + wp.launch( + combine_pose_and_velocity_to_state_batched, + dim=(self._root_newton_view.count), + device=self.device, + inputs=[ + self.default_root_pose, + self.default_root_vel, + state, + ], + ) + return state + + @property + @deprecated("root_link_pose_w") + def root_pose_w(self) -> wp.array: + """Same as :attr:`root_link_pose_w`.""" + return self.sim_bind_root_link_pose_w + + @property + @deprecated("root_link_pos_w") + def root_pos_w(self) -> wp.array: + """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: + """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: + """Same as :attr:`root_com_vel_w`.""" + return self.sim_bind_root_com_vel_w + + @property + @deprecated("root_com_lin_vel_w") + def root_lin_vel_w(self) -> wp.array: + """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: + """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: + """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: + """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: + """Same as :attr:`body_link_pose_w`.""" + return self.sim_bind_body_link_pose_w + + @property + @deprecated("body_link_pos_w") + def body_pos_w(self) -> wp.array: + """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: + """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: + """Same as :attr:`body_com_vel_w`.""" + return self.sim_bind_body_com_vel_w + + @property + @deprecated("body_com_lin_vel_w") + def body_lin_vel_w(self) -> wp.array: + """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: + """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: + """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: + """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: + """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: + """Same as :attr:`body_com_pos_b`.""" + return self.sim_bind_body_com_pos_b + + @property + @deprecated("body_com_quat_b") + def com_quat_b(self) -> wp.array: + """Same as :attr:`body_com_quat_b`.""" + return self.body_com_quat_b + + @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 + + @property + @deprecated("fixed_tendon_pos_limits") + def fixed_tendon_limit(self) -> wp.array: + return self.fixed_tendon_pos_limits diff --git a/source/isaaclab/isaaclab/assets/core/__init__.py b/source/isaaclab/isaaclab/assets/core/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/source/isaaclab/isaaclab/assets/core/body_properties/body.py b/source/isaaclab/isaaclab/assets/core/body_properties/body.py new file mode 100644 index 00000000000..062b198a77b --- /dev/null +++ b/source/isaaclab/isaaclab/assets/core/body_properties/body.py @@ -0,0 +1,201 @@ +# 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.assets.core.body_properties.body_data import BodyData +from isaaclab.assets.core.kernels import ( + generate_mask_from_ids, + update_wrench_array_with_force, + update_wrench_array_with_torque, + update_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 + + """ + 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_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._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/assets/core/body_properties/body_data.py b/source/isaaclab/isaaclab/assets/core/body_properties/body_data.py new file mode 100644 index 00000000000..35441b58b73 --- /dev/null +++ b/source/isaaclab/isaaclab/assets/core/body_properties/body_data.py @@ -0,0 +1,459 @@ +# 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.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, +) + + +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("mass", NewtonManager.get_model()) + self._sim_bind_body_inertia = self._root_newton_view.get_attribute("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. + """ + return self._sim_bind_body_link_pose_w[:, :, 3:] + + @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]. + """ + return self._sim_bind_body_link_pose_w[:, :, 3:] + + @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. + """ + return self.body_link_vel_w[:, :, 3:] + + @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. + """ + return self.body_link_vel_w[:, :, 3:] + + @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. + """ + return self.body_com_pose_w[:, :, 3:] + + @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]. + """ + return self.body_com_pose_w[:, :, 3:] + + @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. + """ + return self._sim_bind_body_com_vel_w[:, :, 3:] + + @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. + """ + return self._sim_bind_body_com_vel_w[:, :, 3:] + + @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. + """ + return self.body_com_acc_w[:, :, :3] + + @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. + """ + return self.body_com_acc_w[:, :, 3:] + + @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]. + """ + return self.body_com_pose_b[:, :, 3:] \ No newline at end of file diff --git a/source/isaaclab/isaaclab/assets/core/joint_properties/__init__.py b/source/isaaclab/isaaclab/assets/core/joint_properties/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/source/isaaclab/isaaclab/assets/core/joint_properties/joint.py b/source/isaaclab/isaaclab/assets/core/joint_properties/joint.py new file mode 100644 index 00000000000..5c6cd665d62 --- /dev/null +++ b/source/isaaclab/isaaclab/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.assets.core.joint_properties.joint_data import JointData +from isaaclab.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.cfg.soft_joint_pos_limit_factor, + ], + ) \ No newline at end of file diff --git a/source/isaaclab/isaaclab/assets/core/joint_properties/joint_data.py b/source/isaaclab/isaaclab/assets/core/joint_properties/joint_data.py new file mode 100644 index 00000000000..16047f4ad65 --- /dev/null +++ b/source/isaaclab/isaaclab/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.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/assets/core/kernels/__init__.py b/source/isaaclab/isaaclab/assets/core/kernels/__init__.py new file mode 100644 index 00000000000..e4ffab13983 --- /dev/null +++ b/source/isaaclab/isaaclab/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/assets/core/kernels/joint_kernels.py b/source/isaaclab/isaaclab/assets/core/kernels/joint_kernels.py new file mode 100644 index 00000000000..be983ebaca8 --- /dev/null +++ b/source/isaaclab/isaaclab/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/assets/core/kernels/other_kernels.py b/source/isaaclab/isaaclab/assets/core/kernels/other_kernels.py new file mode 100644 index 00000000000..1c11f77c4a9 --- /dev/null +++ b/source/isaaclab/isaaclab/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/assets/core/kernels/state_kernels.py b/source/isaaclab/isaaclab/assets/core/kernels/state_kernels.py new file mode 100644 index 00000000000..41f9a12ba34 --- /dev/null +++ b/source/isaaclab/isaaclab/assets/core/kernels/state_kernels.py @@ -0,0 +1,521 @@ +# 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 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). + + .. caution:: The velocity is given with angular velocity first and linear velocity second. + """ + 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). + + .. caution:: The velocity is given with angular velocity first and linear velocity second. + """ + 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). + + .. caution:: The velocity is given with angular velocity first and linear velocity second. + + 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). + + .. caution:: The velocity is given with angular velocity first and linear velocity second. + + 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). + + .. caution:: The velocity is given with angular velocity first and linear velocity second. + + 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). + + .. caution:: The velocity is given with angular velocity first and linear velocity second. + + 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). + + .. caution:: The velocity is given with angular velocity first and linear velocity second. + + 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/assets/core/kernels/update.py b/source/isaaclab/isaaclab/assets/core/kernels/update.py new file mode 100644 index 00000000000..05bb1adbecf --- /dev/null +++ b/source/isaaclab/isaaclab/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/assets/core/kernels/velocity_kernels.py b/source/isaaclab/isaaclab/assets/core/kernels/velocity_kernels.py new file mode 100644 index 00000000000..3da68d6341f --- /dev/null +++ b/source/isaaclab/isaaclab/assets/core/kernels/velocity_kernels.py @@ -0,0 +1,469 @@ +# 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 get_linear_velocity(velocity: wp.array(dtype=wp.spatial_vectorf), linear_velocity: wp.array(dtype=wp.vec3f)): + """ + Get the linear velocity from a spatial 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: + velocity: The spatial velocity. Shape is (6,). + linear_velocity: The linear velocity. Shape is (3,). (modified) + """ + index = wp.tid() + linear_velocity[index] = wp.spatial_bottom(velocity[index]) + + +@wp.kernel +def get_angular_velocity(velocity: wp.array(dtype=wp.spatial_vectorf), angular_velocity: wp.array(dtype=wp.vec3f)): + """ + Get the angular velocity from a spatial 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: + velocity: The spatial velocity. Shape is (6,). + angular_velocity: The angular velocity. Shape is (3,). (modified) + """ + index = wp.tid() + angular_velocity[index] = wp.spatial_bottom(velocity[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/assets/core/root_properties/__init__.py b/source/isaaclab/isaaclab/assets/core/root_properties/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/source/isaaclab/isaaclab/assets/core/root_properties/root.py b/source/isaaclab/isaaclab/assets/core/root_properties/root.py new file mode 100644 index 00000000000..b074a25aaf5 --- /dev/null +++ b/source/isaaclab/isaaclab/assets/core/root_properties/root.py @@ -0,0 +1,389 @@ +# 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.assets.core.root_properties.root_data import RootData +from isaaclab.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 + + """ + 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 + self._body_data._body_com_pose_w.timestamp = -1.0 + + 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, + ], + ) + # Need to invalidate the buffer to trigger the update with the new state. + self._body_data._body_com_pose_w.timestamp = -1.0 + + 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/assets/core/root_properties/root_data.py b/source/isaaclab/isaaclab/assets/core/root_properties/root_data.py new file mode 100644 index 00000000000..c3a26cc8336 --- /dev/null +++ b/source/isaaclab/isaaclab/assets/core/root_properties/root_data.py @@ -0,0 +1,596 @@ +# 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.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, +) + + +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("mass", NewtonManager.get_model())[:, 0] + self._sim_bind_root_inertia = self._root_newton_view.get_attribute("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 + self._previous_root_com_vel = wp.clone(self._root_newton_view.get_root_velocities(NewtonManager.get_state_0())) + + # -- 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 (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.root_com_pose_b[:, 3:] + + ## + # 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", "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 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. + """ + return self._sim_bind_root_link_pose_w[:, :3] + + @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 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]. + """ + return self._sim_bind_root_link_pose_w[:, 3:] + + @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 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. + """ + return self.root_link_vel_w[:, :3] + + @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 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. + """ + return self.root_link_vel_w[:, 3:] + + @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 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. + """ + return self.root_com_pose_w[:, :3] + + @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 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]. + """ + return self.root_com_pose_w[:, 3:] + + @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 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. + """ + return self._sim_bind_root_com_vel_w[:, :3] + + @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 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. + """ + return self._sim_bind_root_com_vel_w[:, 3:] + + @property + @warn_overhead_cost("root_link_vel_b", "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 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. + """ + return self.root_link_vel_b[:, :3] + + @property + @warn_overhead_cost("root_link_vel_b", "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 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. + """ + return self.root_link_vel_b[:, 3:] + + @property + @warn_overhead_cost("root_com_vel_b", "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 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. + """ + return self.root_com_vel_b[:, :3] + + @property + @warn_overhead_cost("root_com_vel_b", "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 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. + """ + return self.root_com_vel_b[:, 3:] \ No newline at end of file diff --git a/source/isaaclab/isaaclab/assets/rigid_object_new/__init__.py b/source/isaaclab/isaaclab/assets/rigid_object_new/__init__.py new file mode 100644 index 00000000000..94c39677672 --- /dev/null +++ b/source/isaaclab/isaaclab/assets/rigid_object_new/__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_new/rigid_object.py b/source/isaaclab/isaaclab/assets/rigid_object_new/rigid_object.py new file mode 100644 index 00000000000..02185a4bc04 --- /dev/null +++ b/source/isaaclab/isaaclab/assets/rigid_object_new/rigid_object.py @@ -0,0 +1,260 @@ +# 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 omni.physics.tensors.impl.api as physx +from isaacsim.core.simulation_manager import SimulationManager +from pxr import UsdPhysics + +import isaaclab.sim as sim_utils + +from ..asset_base import AssetBase +from .rigid_object_data import RigidObjectData + +if TYPE_CHECKING: + from .rigid_object_cfg import RigidObjectCfg + + +class RigidObject(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.""" + + def __init__(self, cfg: RigidObjectCfg): + """Initialize the rigid object. + + Args: + cfg: A configuration instance. + """ + super().__init__(cfg) + + """ + Properties + """ + + @property + def data(self) -> RigidObjectData: + pass + + @property + def num_instances(self) -> int: + pass + + @property + def num_bodies(self) -> int: + pass + + @property + def body_names(self) -> list[str]: + pass + + @property + def root_physx_view(self): + pass + + """ + Operations. + """ + + def reset(self, env_ids: Sequence[int] | None = None): + pass + + def write_data_to_sim(self): + pass + + def update(self, dt: float): + pass + + """ + Operations - Finders. + """ + + def find_bodies(self, name_keys: str | Sequence[str], preserve_order: bool = False) -> tuple[list[int], list[str]]: + pass + + """ + Operations - Write to simulation. + """ + + def write_root_state_to_sim(self, root_state: torch.Tensor, env_ids: Sequence[int] | None = None): + pass + + def write_root_com_state_to_sim(self, root_state: torch.Tensor, env_ids: Sequence[int] | None = None): + pass + + def write_root_link_state_to_sim(self, root_state: torch.Tensor, env_ids: Sequence[int] | None = None): + pass + + def write_root_pose_to_sim(self, root_pose: torch.Tensor, env_ids: Sequence[int] | None = None): + pass + + def write_root_link_pose_to_sim(self, root_pose: torch.Tensor, env_ids: Sequence[int] | None = None): + pass + + def write_root_com_pose_to_sim(self, root_pose: torch.Tensor, env_ids: Sequence[int] | None = None): + pass + + def write_root_velocity_to_sim(self, root_velocity: torch.Tensor, env_ids: Sequence[int] | None = None): + pass + + def write_root_com_velocity_to_sim(self, root_velocity: torch.Tensor, env_ids: Sequence[int] | None = None): + pass + + def write_root_link_velocity_to_sim(self, root_velocity: torch.Tensor, env_ids: Sequence[int] | None = None): + pass + + """ + Operations - Setters. + """ + + def set_external_force_and_torque( + self, + forces: torch.Tensor, + torques: torch.Tensor, + positions: torch.Tensor | None = None, + body_ids: Sequence[int] | slice | None = None, + env_ids: Sequence[int] | None = None, + is_global: bool = False, + ): + pass + + """ + Internal helper. + """ + + def _initialize_impl(self): + # obtain global simulation view + self._physics_sim_view = SimulationManager.get_physics_sim_view() + # obtain the first prim in the regex expression (all others are assumed to be a copy of this) + template_prim = sim_utils.find_first_matching_prim(self.cfg.prim_path) + if template_prim is None: + raise RuntimeError(f"Failed to find prim for expression: '{self.cfg.prim_path}'.") + template_prim_path = template_prim.GetPath().pathString + + # find rigid root prims + root_prims = sim_utils.get_all_matching_child_prims( + template_prim_path, + predicate=lambda prim: prim.HasAPI(UsdPhysics.RigidBodyAPI), + traverse_instance_prims=False, + ) + if len(root_prims) == 0: + raise RuntimeError( + f"Failed to find a rigid body when resolving '{self.cfg.prim_path}'." + " Please ensure that the prim has 'USD RigidBodyAPI' applied." + ) + if len(root_prims) > 1: + raise RuntimeError( + f"Failed to find a single rigid body when resolving '{self.cfg.prim_path}'." + f" Found multiple '{root_prims}' under '{template_prim_path}'." + " Please ensure that there is only one rigid body in the prim path tree." + ) + + articulation_prims = sim_utils.get_all_matching_child_prims( + template_prim_path, + predicate=lambda prim: prim.HasAPI(UsdPhysics.ArticulationRootAPI), + traverse_instance_prims=False, + ) + if len(articulation_prims) != 0: + if articulation_prims[0].GetAttribute("physxArticulation:articulationEnabled").Get(): + raise RuntimeError( + f"Found an articulation root when resolving '{self.cfg.prim_path}' for rigid objects. These are" + f" located at: '{articulation_prims}' under '{template_prim_path}'. Please disable the articulation" + " root in the USD or from code by setting the parameter" + " 'ArticulationRootPropertiesCfg.articulation_enabled' to False in the spawn configuration." + ) + + # resolve root prim back into regex expression + root_prim_path = root_prims[0].GetPath().pathString + root_prim_path_expr = self.cfg.prim_path + root_prim_path[len(template_prim_path) :] + # -- object view + self._root_physx_view = self._physics_sim_view.create_rigid_body_view(root_prim_path_expr.replace(".*", "*")) + + # check if the rigid body was created + if self._root_physx_view._backend is None: + raise RuntimeError(f"Failed to create rigid body at: {self.cfg.prim_path}. Please check PhysX logs.") + + # log information about the rigid body + omni.log.info(f"Rigid body initialized at: {self.cfg.prim_path} with root '{root_prim_path_expr}'.") + omni.log.info(f"Number of instances: {self.num_instances}") + omni.log.info(f"Number of bodies: {self.num_bodies}") + omni.log.info(f"Body names: {self.body_names}") + + # container for data access + self._data = RigidObjectData(self.root_physx_view, self.device) + + # create buffers + self._create_buffers() + # process configuration + self._process_cfg() + # update the rigid body data + self.update(0.0) + + def _create_buffers(self): + """Create buffers for storing data.""" + # 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) + self.uses_external_wrench_positions = False + self._external_wrench_positions_b = torch.zeros_like(self._external_force_b) + self._use_global_wrench_frame = False + + # set information about rigid body into data + self._data.body_names = self.body_names + self._data.default_mass = self.root_physx_view.get_masses().clone() + self._data.default_inertia = self.root_physx_view.get_inertias().clone() + + 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) + + """ + 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._root_physx_view = None diff --git a/source/isaaclab/isaaclab/assets/rigid_object_new/rigid_object_cfg.py b/source/isaaclab/isaaclab/assets/rigid_object_new/rigid_object_cfg.py new file mode 100644 index 00000000000..1cd24bcc918 --- /dev/null +++ b/source/isaaclab/isaaclab/assets/rigid_object_new/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_new/rigid_object_data.py b/source/isaaclab/isaaclab/assets/rigid_object_new/rigid_object_data.py new file mode 100644 index 00000000000..d446504a94d --- /dev/null +++ b/source/isaaclab/isaaclab/assets/rigid_object_new/rigid_object_data.py @@ -0,0 +1,376 @@ +# 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.assets.core.root_properties.root_data import RootData +from isaaclab.utils.helpers import deprecated + + +class RigidObjectData: + """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_newton_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 + # 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) + + # Initialize the data containers + self._root_data = RootData(root_newton_view, device) + + def update(self, dt: float): + self._root_data.update(dt) + + ## + # Names. + ## + + body_names: list[str] = None + """Body names in the order parsed by the simulation view.""" + + ## + # Defaults. + ## + + @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._root_data.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._root_data.default_root_vel + + ## + # Root state properties. + ## + + @property + def root_mass(self) -> wp.array: + """Root mass ``wp.float32`` in the world frame. Shape is (num_instances,).""" + return self._root_data.root_mass + + @property + def root_inertia(self) -> wp.array: + """Root inertia ``wp.mat33`` in the world frame. Shape is (num_instances, 9).""" + return self._root_data.root_inertia + + @property + def root_link_pose_w(self) -> wp.array: + """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. + """ + return self._root_data.root_link_pose_w + + + @property + def root_link_vel_w(self) -> 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 actor frame of the root + rigid body relative to the world. + """ + return self._root_data.root_link_vel_w + + @property + def root_com_pose_w(self) -> 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 center of mass frame of the root rigid body relative to the world. + The orientation is provided in (w, x, y, z) format. + """ + return self._root_data.root_com_pose_w + + @property + def root_com_vel_w(self) -> 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 root rigid body's center of mass frame + relative to the world. + """ + return self._root_data.root_com_vel_w + + @property + def root_state_w(self) -> wp.array: + """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. + """ + return self._root_data.root_state_w + + @property + def root_link_state_w(self) -> 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 rigid body root frame relative to the + world. The orientation is provided in (w, x, y, z) format. + """ + return self._root_data.root_link_state_w + + @property + def root_com_state_w(self) -> 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 rigid body's center of mass frame + relative to the world. Center of mass frame is the orientation principle axes of inertia. + """ + return self._root_data.root_com_state_w + + ## + # Derived Properties. + ## + + @property + def projected_gravity_b(self) -> wp.array: + """Projection of the gravity direction on base frame. Shape is (num_instances, 3).""" + return self._root_data.projected_gravity_b + + @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)`. + """ + return self._root_data.heading_w + + @property + def root_link_lin_vel_b(self) -> wp.array: + """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. + """ + return self._root_data.root_link_lin_vel_b + + @property + def root_link_ang_vel_b(self) -> wp.array: + """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. + """ + return self._root_data.root_link_ang_vel_b + + @property + def root_com_lin_vel_b(self) -> wp.array: + """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. + """ + return self._root_data.root_com_lin_vel_b + + @property + 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 root rigid body's center of mass frame with respect to the + rigid body's actor frame. + """ + return self._root_data.root_com_ang_vel_b + + @property + def root_com_pos_b(self) -> wp.array: + """Root center of mass position 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._root_data.root_com_pos_b + + @property + def root_com_quat_b(self) -> wp.array: + """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._root_data.root_com_quat_b + + ## + # Sliced properties. + ## + + @property + def root_link_pos_w(self) -> 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. + """ + return self._root_data.root_link_pos_w + + @property + def root_link_quat_w(self) -> 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. + """ + return self._root_data.root_link_quat_w + + @property + def root_link_lin_vel_w(self) -> 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. + """ + return self._root_data.root_link_lin_vel_w + + @property + def root_link_ang_vel_w(self) -> 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. + """ + return self._root_data.root_link_ang_vel_w + + @property + 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. + """ + return self._root_data.root_com_pos_w + + @property + def root_com_quat_w(self) -> 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. + """ + return self._root_data.root_com_quat_w + + @property + def root_com_lin_vel_w(self) -> 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. + """ + return self._root_data.root_com_lin_vel_w + + @property + def root_com_ang_vel_w(self) -> 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. + """ + return self._root_data.root_com_ang_vel_w + + @property + def root_com_quat_b(self) -> wp.array: + """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._root_data.root_com_pose_b[:, 3:] + + ## + # Properties for backwards compatibility. + ## + + @property + @deprecated("root_link_pose_w") + def root_pose_w(self) -> wp.array: + """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: + """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: + """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: + """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: + """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: + """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: + """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: + """Same as :attr:`root_com_ang_vel_b`.""" + return self.root_com_ang_vel_b + + @property + @deprecated("body_com_pos_b") + def com_pos_b(self) -> wp.array: + """Same as :attr:`body_com_pos_b`.""" + return self.root_com_pos_b + + @property + @deprecated("body_com_quat_b") + def com_quat_b(self) -> wp.array: + """Same as :attr:`body_com_quat_b`.""" + return self.root_com_quat_b 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..8157bdc5581 --- /dev/null +++ b/source/isaaclab/isaaclab/envs/direct_rl_env_warp.py @@ -0,0 +1,765 @@ +# 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 + +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 InteractiveSceneWarp +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 = InteractiveSceneWarp(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(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 | None = None): + """Reset environments based on specified indices. + + Args: + env_ids: List of environment ids which must be reset + """ + if mask is None: + mask = self._ALL_ENV_MASK + + self.scene.reset(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/scene/__init__.py b/source/isaaclab/isaaclab/scene/__init__.py index 8a3e818559a..3f5c4b4aa5e 100644 --- a/source/isaaclab/isaaclab/scene/__init__.py +++ b/source/isaaclab/isaaclab/scene/__init__.py @@ -27,3 +27,4 @@ from .interactive_scene import InteractiveScene from .interactive_scene_cfg import InteractiveSceneCfg +from .interactive_scene_warp import InteractiveSceneWarp diff --git a/source/isaaclab/isaaclab/scene/interactive_scene_warp.py b/source/isaaclab/isaaclab/scene/interactive_scene_warp.py new file mode 100644 index 00000000000..3983420b96a --- /dev/null +++ b/source/isaaclab/isaaclab/scene/interactive_scene_warp.py @@ -0,0 +1,674 @@ +# 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 collections.abc import Sequence +from typing import Any + +import carb +import omni.log +import omni.usd +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 +from pxr import PhysxSchema + +import isaaclab.sim as sim_utils +from isaaclab.assets import ArticulationWarp, ArticulationWarpCfg, AssetBaseCfg +from isaaclab.cloner import GridCloner +from isaaclab.sensors import ContactSensorCfg, SensorBase, SensorBaseCfg +from isaaclab.sim import SimulationContext +from isaaclab.sim.utils import get_current_stage_id +from isaaclab.terrains import TerrainImporter, TerrainImporterCfg + +from .interactive_scene_cfg import InteractiveSceneCfg + + +class InteractiveSceneWarp: + """A scene that contains entities added to the simulation. + + The interactive scene parses the :class:`InteractiveSceneCfg` class to create the scene. + Based on the specified number of environments, it clones the entities and groups them into different + categories (e.g., articulations, sensors, etc.). + + Cloning can be performed in two ways: + + * For tasks where all environments contain the same assets, a more performant cloning paradigm + can be used to allow for faster environment creation. This is specified by the ``replicate_physics`` flag. + + .. code-block:: python + + scene = InteractiveScene(cfg=InteractiveSceneCfg(replicate_physics=True)) + + * For tasks that require having separate assets in the environments, ``replicate_physics`` would have to + be set to False, which will add some costs to the overall startup time. + + .. code-block:: python + + scene = InteractiveScene(cfg=InteractiveSceneCfg(replicate_physics=False)) + + Each entity is registered to scene based on its name in the configuration class. For example, if the user + specifies a robot in the configuration class as follows: + + .. code-block:: python + + from isaaclab.scene import InteractiveSceneCfg + from isaaclab.utils import configclass + + from isaaclab_assets.robots.anymal import ANYMAL_C_CFG + + @configclass + class MySceneCfg(InteractiveSceneCfg): + + robot = ANYMAL_C_CFG.replace(prim_path="{ENV_REGEX_NS}/Robot") + + Then the robot can be accessed from the scene as follows: + + .. code-block:: python + + from isaaclab.scene import InteractiveScene + + # create 128 environments + scene = InteractiveScene(cfg=MySceneCfg(num_envs=128)) + + # access the robot from the scene + robot = scene["robot"] + # access the robot based on its type + robot = scene.articulations["robot"] + + If the :class:`InteractiveSceneCfg` class does not include asset entities, the cloning process + can still be triggered if assets were added to the stage outside of the :class:`InteractiveScene` class: + + .. code-block:: python + + scene = InteractiveScene(cfg=InteractiveSceneCfg(num_envs=128, replicate_physics=True)) + scene.clone_environments() + + .. note:: + It is important to note that the scene only performs common operations on the entities. For example, + resetting the internal buffers, writing the buffers to the simulation and updating the buffers from the + simulation. The scene does not perform any task specific to the entity. For example, it does not apply + actions to the robot or compute observations from the robot. These tasks are handled by different + modules called "managers" in the framework. Please refer to the :mod:`isaaclab.managers` sub-package + for more details. + """ + + def __init__(self, cfg: InteractiveSceneCfg): + """Initializes the scene. + + Args: + cfg: The configuration class for the scene. + """ + # check that the config is valid + cfg.validate() + # store inputs + self.cfg = cfg + # initialize scene elements + self._terrain = None + self._articulations = dict() + self._deformable_objects = dict() + self._rigid_objects = dict() + self._rigid_object_collections = dict() + self._sensors = dict() + self._surface_grippers = dict() + self._extras = dict() + # get stage handle + self.sim = SimulationContext.instance() + self.stage = get_current_stage() + self.stage_id = get_current_stage_id() + # physics scene path + self._physics_scene_path = None + # prepare cloner for environment replication + self.cloner = GridCloner(spacing=self.cfg.env_spacing, stage=self.stage) + self.cloner.define_base_env(self.env_ns) + self.env_prim_paths = self.cloner.generate_paths(f"{self.env_ns}/env", self.cfg.num_envs) + # create source prim + self.stage.DefinePrim(self.env_prim_paths[0], "Xform") + + # when replicate_physics=False, we assume heterogeneous environments and clone the xforms first. + # this triggers per-object level cloning in the spawner. + if not self.cfg.replicate_physics: + # check version of Isaac Sim to determine whether clone_in_fabric is valid + isaac_sim_version = float(".".join(get_version()[2])) + if isaac_sim_version < 5: + # clone the env xform + env_origins = self.cloner.clone( + source_prim_path=self.env_prim_paths[0], + prim_paths=self.env_prim_paths, + replicate_physics=False, + copy_from_source=True, + enable_env_ids=( + self.cfg.filter_collisions if self.device != "cpu" else False + ), # this won't do anything because we are not replicating physics + ) + else: + # clone the env xform + env_origins = self.cloner.clone( + source_prim_path=self.env_prim_paths[0], + prim_paths=self.env_prim_paths, + replicate_physics=False, + copy_from_source=True, + enable_env_ids=( + self.cfg.filter_collisions if self.device != "cpu" else False + ), # this won't do anything because we are not replicating physics + clone_in_fabric=self.cfg.clone_in_fabric, + ) + self._default_env_origins = torch.tensor(env_origins, device=self.device, dtype=torch.float32) + else: + # otherwise, environment origins will be initialized during cloning at the end of environment creation + self._default_env_origins = None + + self._global_prim_paths = list() + if self._is_scene_setup_from_cfg(): + # add entities from config + self._add_entities_from_cfg() + # clone environments on a global scope if environment is homogeneous + if self.cfg.replicate_physics: + self.clone_environments(copy_from_source=False) + # replicate physics if we have more than one environment + # this is done to make scene initialization faster at play time + # since env_ids is only applicable when replicating physics, we have to fallback to the previous method + # to filter collisions if replicate_physics is not enabled + # additionally, env_ids is only supported in GPU simulation + if (not self.cfg.replicate_physics and self.cfg.filter_collisions) or self.device == "cpu": + self.filter_collisions(self._global_prim_paths) + + def clone_environments(self, copy_from_source: bool = False): + """Creates clones of the environment ``/World/envs/env_0``. + + Args: + copy_from_source: (bool): If set to False, clones inherit from /World/envs/env_0 and mirror its changes. + If True, clones are independent copies of the source prim and won't reflect its changes (start-up time + may increase). Defaults to False. + """ + # check if user spawned different assets in individual environments + # this flag will be None if no multi asset is spawned + 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( + "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." + ) + + # check version of Isaac Sim to determine whether clone_in_fabric is valid + isaac_sim_version = float(".".join(get_version()[2])) + if isaac_sim_version < 5: + # clone the environment + env_origins = self.cloner.clone( + source_prim_path=self.env_prim_paths[0], + prim_paths=self.env_prim_paths, + replicate_physics=self.cfg.replicate_physics, + copy_from_source=copy_from_source, + enable_env_ids=( + self.cfg.filter_collisions if self.device != "cpu" else False + ), # this automatically filters collisions between environments + ) + else: + # clone the environment + env_origins = self.cloner.clone( + source_prim_path=self.env_prim_paths[0], + prim_paths=self.env_prim_paths, + replicate_physics=self.cfg.replicate_physics, + copy_from_source=copy_from_source, + enable_env_ids=( + self.cfg.filter_collisions if self.device != "cpu" else False + ), # this automatically filters collisions between environments + clone_in_fabric=self.cfg.clone_in_fabric, + ) + + # since env_ids is only applicable when replicating physics, we have to fallback to the previous method + # to filter collisions if replicate_physics is not enabled + # additionally, env_ids is only supported in GPU simulation + 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( + "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." + ) + + # in case of heterogeneous cloning, the env origins is specified at init + if self._default_env_origins is None: + self._default_env_origins = torch.tensor(env_origins, device=self.device, dtype=torch.float32) + + def filter_collisions(self, global_prim_paths: list[str] | None = None): + """Filter environments collisions. + + Disables collisions between the environments in ``/World/envs/env_.*`` and enables collisions with the prims + in global prim paths (e.g. ground plane). + + Args: + global_prim_paths: A list of global prim paths to enable collisions with. + Defaults to None, in which case no global prim paths are considered. + """ + # validate paths in global prim paths + if global_prim_paths is None: + global_prim_paths = [] + else: + # remove duplicates in paths + global_prim_paths = list(set(global_prim_paths)) + + # if "/World/collisions" already exists in the stage, we don't filter again + if self.stage.GetPrimAtPath("/World/collisions"): + return + + # set global prim paths list if not previously defined + if len(self._global_prim_paths) < 1: + self._global_prim_paths += global_prim_paths + + # filter collisions within each environment instance + self.cloner.filter_collisions( + self.physics_scene_path, + "/World/collisions", + self.env_prim_paths, + global_paths=self._global_prim_paths, + ) + + def __str__(self) -> str: + """Returns a string representation of the scene.""" + msg = f"\n" + msg += f"\tNumber of environments: {self.cfg.num_envs}\n" + msg += f"\tEnvironment spacing : {self.cfg.env_spacing}\n" + msg += f"\tSource prim name : {self.env_prim_paths[0]}\n" + msg += f"\tGlobal prim paths : {self._global_prim_paths}\n" + msg += f"\tReplicate physics : {self.cfg.replicate_physics}" + return msg + + """ + Properties. + """ + + @property + def physics_scene_path(self) -> str: + """The path to the USD Physics Scene.""" + if self._physics_scene_path is None: + 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}") + break + if self._physics_scene_path is None: + raise RuntimeError("No physics scene found! Please make sure one exists.") + return self._physics_scene_path + + @property + def physics_dt(self) -> float: + """The physics timestep of the scene.""" + return sim_utils.SimulationContext.instance().get_physics_dt() # pyright: ignore [reportOptionalMemberAccess] + + @property + def device(self) -> str: + """The device on which the scene is created.""" + return sim_utils.SimulationContext.instance().device # pyright: ignore [reportOptionalMemberAccess] + + @property + def env_ns(self) -> str: + """The namespace ``/World/envs`` in which all environments created. + + The environments are present w.r.t. this namespace under "env_{N}" prim, + where N is a natural number. + """ + return "/World/envs" + + @property + def env_regex_ns(self) -> str: + """The namespace ``/World/envs/env_.*`` in which all environments created.""" + return f"{self.env_ns}/env_.*" + + @property + def num_envs(self) -> int: + """The number of environments handled by the scene.""" + return self.cfg.num_envs + + @property + def env_origins(self) -> torch.Tensor: + """The origins of the environments in the scene. Shape is (num_envs, 3).""" + if self._terrain is not None: + return self._terrain.env_origins + else: + return self._default_env_origins + + @property + def terrain(self) -> TerrainImporter | None: + """The terrain in the scene. If None, then the scene has no terrain. + + Note: + We treat terrain separate from :attr:`extras` since terrains define environment origins and are + handled differently from other miscellaneous entities. + """ + return self._terrain + + @property + def articulations(self) -> dict[str, ArticulationWarp]: + """A dictionary of articulations in the scene.""" + return self._articulations + + @property + def deformable_objects(self) -> dict: + """A dictionary of deformable objects in the scene.""" + raise NotImplementedError("Deformable objects are not supported in IsaacLab for Newton.") + + @property + def rigid_objects(self) -> dict: + """A dictionary of rigid objects in the scene.""" + raise NotImplementedError("Rigid objects are not supported in IsaacLab for Newton.") + + @property + def rigid_object_collections(self) -> dict: + """A dictionary of rigid object collections in the scene.""" + raise NotImplementedError("Rigid object collections are not supported in IsaacLab for Newton.") + + @property + def sensors(self) -> dict[str, SensorBase]: + """A dictionary of the sensors in the scene, such as cameras and contact reporters.""" + return self._sensors + + @property + def surface_grippers(self) -> dict: + """A dictionary of the surface grippers in the scene.""" + raise NotImplementedError("Surface grippers are not supported in IsaacLab for Newton.") + + @property + def extras(self) -> dict[str, XFormPrim]: + """A dictionary of miscellaneous simulation objects that neither inherit from assets nor sensors. + + The keys are the names of the miscellaneous objects, and the values are the `XFormPrim`_ + of the corresponding prims. + + As an example, lights or other props in the scene that do not have any attributes or properties that you + want to alter at runtime can be added to this dictionary. + + Note: + These are not reset or updated by the scene. They are mainly other prims that are not necessarily + handled by the interactive scene, but are useful to be accessed by the user. + + .. _XFormPrim: https://docs.omniverse.nvidia.com/py/isaacsim/source/isaacsim.core/docs/index.html#isaacsim.core.prims.XFormPrim + + """ + return self._extras + + @property + def state(self) -> dict[str, dict[str, dict[str, torch.Tensor]]]: + """A dictionary of the state of the scene entities in the simulation world frame. + + Please refer to :meth:`get_state` for the format. + """ + return self.get_state(is_relative=False) + + """ + Operations. + """ + + def reset(self, mask: wp.array | None = None): + """Resets the scene entities. + + Args: + env_ids: The indices of the environments to reset. + Defaults to None (all instances). + """ + # -- assets + for articulation in self._articulations.values(): + articulation.reset(mask) + # -- sensors + # for sensor in self._sensors.values(): + # sensor.reset(mask) + + def write_data_to_sim(self): + """Writes the data of the scene entities to the simulation.""" + # -- assets + for articulation in self._articulations.values(): + articulation.write_data_to_sim() + + def update(self, dt: float) -> None: + """Update the scene entities. + + Args: + dt: The amount of time passed from last :meth:`update` call. + """ + # -- assets + for articulation in self._articulations.values(): + articulation.update(dt) + # -- sensors + for sensor in self._sensors.values(): + sensor.update(dt, force_recompute=not self.cfg.lazy_sensor_update) + + """ + Operations: Scene State. + """ + + def reset_to( + self, + state: dict[str, dict[str, dict[str, torch.Tensor]]], + env_ids: Sequence[int] | None = None, + is_relative: bool = False, + ): + """Resets the entities in the scene to the provided state. + + Args: + state: The state to reset the scene entities to. Please refer to :meth:`get_state` for the format. + env_ids: The indices of the environments to reset. Defaults to None, in which case + all environment instances are reset. + is_relative: If set to True, the state is considered relative to the environment origins. + Defaults to False. + """ + # resolve env_ids + if env_ids is None: + env_ids = slice(None) + # articulations + for asset_name, articulation in self._articulations.items(): + asset_state = state["articulation"][asset_name] + # root state + root_pose = asset_state["root_pose"].clone() + if is_relative: + root_pose[:, :3] += self.env_origins[env_ids] + root_velocity = asset_state["root_velocity"].clone() + articulation.write_root_pose_to_sim(root_pose, env_ids=env_ids) + articulation.write_root_velocity_to_sim(root_velocity, env_ids=env_ids) + # joint state + joint_position = asset_state["joint_position"].clone() + joint_velocity = asset_state["joint_velocity"].clone() + articulation.write_joint_state_to_sim(joint_position, joint_velocity, env_ids=env_ids) + # FIXME: This is not generic as it assumes PD control over the joints. + # This assumption does not hold for effort controlled joints. + articulation.set_joint_position_target(joint_position, env_ids=env_ids) + articulation.set_joint_velocity_target(joint_velocity, env_ids=env_ids) + + # write data to simulation to make sure initial state is set + # this propagates the joint targets to the simulation + self.write_data_to_sim() + + def get_state(self, is_relative: bool = False) -> dict[str, dict[str, dict[str, torch.Tensor]]]: + """Returns the state of the scene entities. + + Based on the type of the entity, the state comprises of different components. + + * For an articulation, the state comprises of the root pose, root velocity, and joint position and velocity. + * For a deformable object, the state comprises of the nodal position and velocity. + * For a rigid object, the state comprises of the root pose and root velocity. + + The returned state is a dictionary with the following format: + + .. code-block:: python + + { + "articulation": { + "entity_1_name": { + "root_pose": torch.Tensor, + "root_velocity": torch.Tensor, + "joint_position": torch.Tensor, + "joint_velocity": torch.Tensor, + }, + "entity_2_name": { + "root_pose": torch.Tensor, + "root_velocity": torch.Tensor, + "joint_position": torch.Tensor, + "joint_velocity": torch.Tensor, + }, + }, + "deformable_object": { + "entity_3_name": { + "nodal_position": torch.Tensor, + "nodal_velocity": torch.Tensor, + } + }, + "rigid_object": { + "entity_4_name": { + "root_pose": torch.Tensor, + "root_velocity": torch.Tensor, + } + }, + } + + where ``entity_N_name`` is the name of the entity registered in the scene. + + Args: + is_relative: If set to True, the state is considered relative to the environment origins. + Defaults to False. + + Returns: + A dictionary of the state of the scene entities. + """ + state = dict() + # articulations + state["articulation"] = dict() + for asset_name, articulation in self._articulations.items(): + asset_state = dict() + asset_state["root_pose"] = articulation.data.root_pose_w.clone() + if is_relative: + asset_state["root_pose"][:, :3] -= self.env_origins + asset_state["root_velocity"] = articulation.data.root_vel_w.clone() + asset_state["joint_position"] = articulation.data.joint_pos.clone() + asset_state["joint_velocity"] = articulation.data.joint_vel.clone() + state["articulation"][asset_name] = asset_state + return state + + """ + Operations: Iteration. + """ + + def keys(self) -> list[str]: + """Returns the keys of the scene entities. + + Returns: + The keys of the scene entities. + """ + all_keys = ["terrain"] + for asset_family in [ + self._articulations, + self._sensors, + self._extras, + ]: + all_keys += list(asset_family.keys()) + return all_keys + + def __getitem__(self, key: str) -> Any: + """Returns the scene entity with the given key. + + Args: + key: The key of the scene entity. + + Returns: + The scene entity. + """ + # check if it is a terrain + if key == "terrain": + return self._terrain + + all_keys = ["terrain"] + # check if it is in other dictionaries + for asset_family in [ + self._articulations, + self._sensors, + self._extras, + ]: + out = asset_family.get(key) + # if found, return + if out is not None: + return out + all_keys += list(asset_family.keys()) + # if not found, raise error + raise KeyError(f"Scene entity with key '{key}' not found. Available Entities: '{all_keys}'") + + """ + Internal methods. + """ + + def _is_scene_setup_from_cfg(self) -> bool: + """Check if scene entities are setup from the config or not. + + Returns: + True if scene entities are setup from the config, False otherwise. + """ + return any( + not (asset_name in InteractiveSceneCfg.__dataclass_fields__ or asset_cfg is None) + for asset_name, asset_cfg in self.cfg.__dict__.items() + ) + + def _add_entities_from_cfg(self): + """Add scene entities from the config.""" + # store paths that are in global collision filter + self._global_prim_paths = list() + # parse the entire scene config and resolve regex + for asset_name, asset_cfg in self.cfg.__dict__.items(): + # skip keywords + # note: easier than writing a list of keywords: [num_envs, env_spacing, lazy_sensor_update] + if asset_name in InteractiveSceneCfg.__dataclass_fields__ or asset_cfg is None: + continue + # resolve regex + if hasattr(asset_cfg, "prim_path"): + asset_cfg.prim_path = asset_cfg.prim_path.format(ENV_REGEX_NS=self.env_regex_ns) + # create asset + if isinstance(asset_cfg, TerrainImporterCfg): + # terrains are special entities since they define environment origins + asset_cfg.num_envs = self.cfg.num_envs + asset_cfg.env_spacing = self.cfg.env_spacing + self._terrain = asset_cfg.class_type(asset_cfg) + elif isinstance(asset_cfg, ArticulationWarpCfg): + self._articulations[asset_name] = asset_cfg.class_type(asset_cfg) + elif isinstance(asset_cfg, SensorBaseCfg): + if isinstance(asset_cfg, ContactSensorCfg): + if asset_cfg.shape_path is not None: + updated_shape_names_expr = [] + for filter_prim_path in asset_cfg.shape_path: + updated_shape_names_expr.append(filter_prim_path.format(ENV_REGEX_NS=self.env_regex_ns)) + asset_cfg.shape_path = updated_shape_names_expr + if asset_cfg.filter_prim_paths_expr is not None: + updated_contact_partners_body_expr = [] + for contact_partners_body_expr in asset_cfg.filter_prim_paths_expr: + updated_contact_partners_body_expr.append( + contact_partners_body_expr.format(ENV_REGEX_NS=self.env_regex_ns) + ) + asset_cfg.filter_prim_paths_expr = updated_contact_partners_body_expr + if asset_cfg.filter_shape_paths_expr is not None: + updated_contact_partners_shape_expr = [] + for contact_partners_shape_expr in asset_cfg.filter_shape_paths_expr: + updated_contact_partners_shape_expr.append( + contact_partners_shape_expr.format(ENV_REGEX_NS=self.env_regex_ns) + ) + asset_cfg.filter_shape_paths_expr = updated_contact_partners_shape_expr + + self._sensors[asset_name] = asset_cfg.class_type(asset_cfg) + elif isinstance(asset_cfg, AssetBaseCfg): + # manually spawn asset + if asset_cfg.spawn is not None: + asset_cfg.spawn.func( + asset_cfg.prim_path, + asset_cfg.spawn, + translation=asset_cfg.init_state.pos, + orientation=asset_cfg.init_state.rot, + ) + # store xform prim view corresponding to this asset + # all prims in the scene are Xform prims (i.e. have a transform component) + self._extras[asset_name] = XFormPrim(asset_cfg.prim_path, reset_xform_properties=False) + else: + raise ValueError(f"Unknown asset config type for {asset_name}: {asset_cfg}") + # store global collision paths + if hasattr(asset_cfg, "collision_group") and asset_cfg.collision_group == -1: + asset_paths = sim_utils.find_matching_prim_paths(asset_cfg.prim_path) + self._global_prim_paths += asset_paths diff --git a/source/isaaclab/isaaclab/sim/_impl/newton_manager.py b/source/isaaclab/isaaclab/sim/_impl/newton_manager.py index ed150b4c3bf..c0a7f9ff413 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) @@ -271,6 +277,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 +289,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 +425,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..6fc8fb29ded 100644 --- a/source/isaaclab/isaaclab/utils/__init__.py +++ b/source/isaaclab/isaaclab/utils/__init__.py @@ -9,6 +9,7 @@ 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 * 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..37c51a80c54 --- /dev/null +++ b/source/isaaclab/isaaclab/utils/helpers.py @@ -0,0 +1,292 @@ +# 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 omni.log +import warp as wp + + +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. + omni.log.warn( + deprecation_message, + ) + # 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. + omni.log.warn( + warning_message, + ) + # 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/ant.py b/source/isaaclab_assets/isaaclab_assets/robots/ant.py index 69b59ecab04..e9f1c143a8a 100644 --- a/source/isaaclab_assets/isaaclab_assets/robots/ant.py +++ b/source/isaaclab_assets/isaaclab_assets/robots/ant.py @@ -9,7 +9,8 @@ import isaaclab.sim as sim_utils from isaaclab.actuators import ImplicitActuatorCfg -from isaaclab.assets import ArticulationCfg +from isaaclab.actuators_warp import ImplicitActuatorWarpCfg +from isaaclab.assets import ArticulationCfg, ArticulationWarpCfg from isaaclab.utils.assets import ISAAC_NUCLEUS_DIR ## @@ -44,3 +45,32 @@ }, ) """Configuration for the Mujoco Ant robot.""" + +ANT_CFG_WARP = ArticulationWarpCfg( + prim_path="{ENV_REGEX_NS}/Robot", + spawn=sim_utils.UsdFileCfg( + usd_path=f"{ISAAC_NUCLEUS_DIR}/Robots/IsaacSim/Ant/ant_instanceable.usd", + copy_from_source=False, + ), + init_state=ArticulationWarpCfg.InitialStateCfg( + pos=(0.0, 0.0, 0.5), + joint_pos={ + ".*_leg": 0.0, + "front_left_foot": 0.785398, # 45 degrees + "front_right_foot": -0.785398, + "left_back_foot": -0.785398, + "right_back_foot": 0.785398, + }, + ), + actuators={ + "body": ImplicitActuatorWarpCfg( + joint_names_expr=[".*"], + control_mode="none", + stiffness=0.0, + damping=0.0, + effort_limit_sim=30.0, + friction=0.0001, + ), + }, +) +"""Configuration for the Mujoco Ant robot.""" diff --git a/source/isaaclab_assets/isaaclab_assets/robots/cartpole.py b/source/isaaclab_assets/isaaclab_assets/robots/cartpole.py index ad577ae6bba..68d1c2d1d96 100644 --- a/source/isaaclab_assets/isaaclab_assets/robots/cartpole.py +++ b/source/isaaclab_assets/isaaclab_assets/robots/cartpole.py @@ -8,7 +8,8 @@ import isaaclab.sim as sim_utils from isaaclab.actuators import ImplicitActuatorCfg -from isaaclab.assets import ArticulationCfg +from isaaclab.actuators_warp import ImplicitActuatorWarpCfg +from isaaclab.assets import ArticulationCfg, ArticulationWarpCfg from isaaclab.utils.assets import ISAACLAB_NUCLEUS_DIR ## @@ -45,3 +46,34 @@ }, ) """Configuration for a simple Cartpole robot.""" + +CARTPOLE_CFG_WARP = ArticulationWarpCfg( + spawn=sim_utils.UsdFileCfg( + usd_path=f"{ISAACLAB_NUCLEUS_DIR}/Robots/Classic/Cartpole/cartpole.usd", + articulation_props=sim_utils.ArticulationRootPropertiesCfg( + enabled_self_collisions=False, + ), + ), + init_state=ArticulationWarpCfg.InitialStateCfg( + pos=(0.0, 0.0, 2.0), joint_pos={"slider_to_cart": 0.0, "cart_to_pole": 0.0} + ), + actuators={ + "cart_actuator": ImplicitActuatorWarpCfg( + joint_names_expr=["slider_to_cart"], + control_mode="position", + effort_limit=400.0, + velocity_limit=100.0, + stiffness=0.0, + damping=10.0, + ), + "pole_actuator": ImplicitActuatorWarpCfg( + joint_names_expr=["cart_to_pole"], + control_mode="none", + effort_limit=400.0, + velocity_limit=100.0, + stiffness=0.0, + damping=0.0, + ), + }, +) +"""Configuration for a simple Cartpole robot.""" diff --git a/source/isaaclab_assets/isaaclab_assets/robots/humanoid.py b/source/isaaclab_assets/isaaclab_assets/robots/humanoid.py index 0d1a156f470..ce4a24d71dc 100644 --- a/source/isaaclab_assets/isaaclab_assets/robots/humanoid.py +++ b/source/isaaclab_assets/isaaclab_assets/robots/humanoid.py @@ -9,7 +9,8 @@ import isaaclab.sim as sim_utils from isaaclab.actuators import ImplicitActuatorCfg -from isaaclab.assets import ArticulationCfg +from isaaclab.actuators_warp import ImplicitActuatorWarpCfg +from isaaclab.assets import ArticulationCfg, ArticulationWarpCfg from isaaclab.utils.assets import ISAAC_NUCLEUS_DIR ## @@ -32,7 +33,55 @@ actuators={ "body": ImplicitActuatorCfg( joint_names_expr=[".*"], - control_mode="none", + control_mode="position", + stiffness={ + ".*_waist.*": 20.0, + ".*_upper_arm.*": 10.0, + "pelvis": 10.0, + ".*_lower_arm": 2.0, + ".*_thigh:0": 10.0, + ".*_thigh:1": 20.0, + ".*_thigh:2": 10.0, + ".*_shin": 5.0, + ".*_foot.*": 2.0, + }, + damping={ + ".*_waist.*": 5.0, + ".*_upper_arm.*": 5.0, + "pelvis": 5.0, + ".*_lower_arm": 1.0, + ".*_thigh:0": 5.0, + ".*_thigh:1": 5.0, + ".*_thigh:2": 5.0, + ".*_shin": 0.1, + ".*_foot.*": 1.0, + }, + velocity_limit_sim={".*": 100.0}, + armature={".*": 0.01}, + friction=0.00001, + effort_limit_sim=150.0, + ), + }, +) +"""Configuration for the Mujoco Humanoid robot.""" + +HUMANOID_CFG_WARP = ArticulationWarpCfg( + prim_path="{ENV_REGEX_NS}/Robot", + spawn=sim_utils.UsdFileCfg( + usd_path=f"{ISAAC_NUCLEUS_DIR}/Robots/IsaacSim/Humanoid/humanoid_instanceable.usd", + articulation_props=sim_utils.ArticulationRootPropertiesCfg( + enabled_self_collisions=True, + ), + copy_from_source=False, + ), + init_state=ArticulationWarpCfg.InitialStateCfg( + pos=(0.0, 0.0, 1.34), + joint_pos={".*": 0.0}, + ), + actuators={ + "body": ImplicitActuatorWarpCfg( + joint_names_expr=[".*"], + control_mode="position", stiffness={ ".*_waist.*": 20.0, ".*_upper_arm.*": 10.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..355ac6dedc7 --- /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_WARP + +import isaaclab.sim as sim_utils +from isaaclab.assets import ArticulationWarpCfg +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: ArticulationWarpCfg = ANT_CFG_WARP.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..a2f5f0df208 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 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..81d1ad9d8c5 --- /dev/null +++ b/source/isaaclab_tasks/isaaclab_tasks/direct/cartpole/cartpole_warp_env.py @@ -0,0 +1,355 @@ +# 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_WARP + +import isaaclab.sim as sim_utils +from isaaclab.assets import ArticulationWarp, ArticulationWarpCfg +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: ArticulationWarpCfg = CARTPOLE_CFG_WARP.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 = ArticulationWarp(self.cfg.robot_cfg) + # 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..6c0b83e0739 --- /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_WARP + +import isaaclab.sim as sim_utils +from isaaclab.assets import ArticulationWarpCfg +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: ArticulationWarpCfg = HUMANOID_CFG_WARP.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..ffd56db07f2 100644 --- a/source/isaaclab_tasks/isaaclab_tasks/direct/locomotion/locomotion_env.py +++ b/source/isaaclab_tasks/isaaclab_tasks/direct/locomotion/locomotion_env.py @@ -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..811f20a14b5 --- /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 ArticulationWarp +from isaaclab.envs import DirectRLEnvCfg, DirectRLEnvWarp + + +@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.sim_bind_joint_pos_limits_lower)[0]) + print("joint_pos_limits_upper", wp.to_torch(self.robot.data.sim_bind_joint_pos_limits_upper)[0]) + print("joint_stiffness_sim", wp.to_torch(self.robot.data.sim_bind_joint_stiffness_sim)[0]) + print("joint_damping_sim", wp.to_torch(self.robot.data.sim_bind_joint_damping_sim)[0]) + print("joint_armature", wp.to_torch(self.robot.data.sim_bind_joint_armature)[0]) + print("joint_friction_coeff", wp.to_torch(self.robot.data.sim_bind_joint_friction_coeff)[0]) + print("joint_vel_limits_sim", wp.to_torch(self.robot.data.sim_bind_joint_vel_limits_sim)[0]) + print("joint_effort_limits_sim", wp.to_torch(self.robot.data.sim_bind_joint_effort_limits_sim)[0]) + print("joint_control_mode_sim", wp.to_torch(self.robot.data.sim_bind_joint_control_mode_sim)[0]) + print("joint_pos", wp.to_torch(self.robot.data.sim_bind_joint_pos)[0]) + print("joint_vel", wp.to_torch(self.robot.data.sim_bind_joint_vel)[0]) + print("joint_effort", wp.to_torch(self.robot.data.sim_bind_joint_effort)[0]) + print("joint_target", wp.to_torch(self.robot.data.sim_bind_joint_target)[0]) + print("soft_joint_pos_limits", wp.to_torch(self.robot.data.soft_joint_pos_limits)[0]) + + def _setup_scene(self) -> None: + self.robot = ArticulationWarp(self.cfg.robot) + # 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: