# Copyright (c) 2022-2024, The ORBIT Project Developers.
# All rights reserved.
#
# SPDX-License-Identifier: BSD-3-Clause
"""Common functions that can be used to create observation terms.
The functions can be passed to the :class:`omni.isaac.orbit.managers.ObservationTermCfg` object to enable
the observation introduced by the function.
"""
from __future__ import annotations
import torch
from typing import TYPE_CHECKING
import omni.isaac.orbit.utils.math as math_utils
from omni.isaac.orbit.assets import Articulation, RigidObject
from omni.isaac.orbit.managers import SceneEntityCfg
from omni.isaac.orbit.sensors import RayCaster
if TYPE_CHECKING:
from omni.isaac.orbit.envs import BaseEnv, RLTaskEnv
"""
Root state.
"""
[docs]def base_pos_z(env: BaseEnv, asset_cfg: SceneEntityCfg = SceneEntityCfg("robot")) -> torch.Tensor:
"""Root height in the simulation world frame."""
# extract the used quantities (to enable type-hinting)
asset: Articulation = env.scene[asset_cfg.name]
return asset.data.root_pos_w[:, 2].unsqueeze(-1)
[docs]def base_lin_vel(env: BaseEnv, asset_cfg: SceneEntityCfg = SceneEntityCfg("robot")) -> torch.Tensor:
"""Root linear velocity in the asset's root frame."""
# extract the used quantities (to enable type-hinting)
asset: RigidObject = env.scene[asset_cfg.name]
return asset.data.root_lin_vel_b
[docs]def base_ang_vel(env: BaseEnv, asset_cfg: SceneEntityCfg = SceneEntityCfg("robot")) -> torch.Tensor:
"""Root angular velocity in the asset's root frame."""
# extract the used quantities (to enable type-hinting)
asset: RigidObject = env.scene[asset_cfg.name]
return asset.data.root_ang_vel_b
[docs]def projected_gravity(env: BaseEnv, asset_cfg: SceneEntityCfg = SceneEntityCfg("robot")) -> torch.Tensor:
"""Gravity projection on the asset's root frame."""
# extract the used quantities (to enable type-hinting)
asset: RigidObject = env.scene[asset_cfg.name]
return asset.data.projected_gravity_b
[docs]def root_pos_w(env: BaseEnv, asset_cfg: SceneEntityCfg = SceneEntityCfg("robot")) -> torch.Tensor:
"""Asset root position in the environment frame."""
# extract the used quantities (to enable type-hinting)
asset: RigidObject = env.scene[asset_cfg.name]
return asset.data.root_pos_w - env.scene.env_origins
[docs]def root_quat_w(
env: BaseEnv, make_quat_unique: bool = False, asset_cfg: SceneEntityCfg = SceneEntityCfg("robot")
) -> torch.Tensor:
"""Asset root orientation (w, x, y, z) in the environment frame.
If :attr:`make_quat_unique` is True, then returned quaternion is made unique by ensuring
the quaternion has non-negative real component. This is because both ``q`` and ``-q`` represent
the same orientation.
"""
# extract the used quantities (to enable type-hinting)
asset: RigidObject = env.scene[asset_cfg.name]
quat = asset.data.root_quat_w
# make the quaternion real-part positive if configured
return math_utils.quat_unique(quat) if make_quat_unique else quat
[docs]def root_lin_vel_w(env: BaseEnv, asset_cfg: SceneEntityCfg = SceneEntityCfg("robot")) -> torch.Tensor:
"""Asset root linear velocity in the environment frame."""
# extract the used quantities (to enable type-hinting)
asset: RigidObject = env.scene[asset_cfg.name]
return asset.data.root_lin_vel_w
[docs]def root_ang_vel_w(env: BaseEnv, asset_cfg: SceneEntityCfg = SceneEntityCfg("robot")) -> torch.Tensor:
"""Asset root angular velocity in the environment frame."""
# extract the used quantities (to enable type-hinting)
asset: RigidObject = env.scene[asset_cfg.name]
return asset.data.root_ang_vel_w
"""
Joint state.
"""
[docs]def joint_pos(env: BaseEnv, asset_cfg: SceneEntityCfg = SceneEntityCfg("robot")) -> torch.Tensor:
"""The joint positions of the asset.
Note: Only the joints configured in :attr:`asset_cfg.joint_ids` will have their positions returned.
"""
# extract the used quantities (to enable type-hinting)
asset: Articulation = env.scene[asset_cfg.name]
return asset.data.joint_pos[:, asset_cfg.joint_ids]
[docs]def joint_pos_rel(env: BaseEnv, asset_cfg: SceneEntityCfg = SceneEntityCfg("robot")) -> torch.Tensor:
"""The joint positions of the asset w.r.t. the default joint positions.
Note: Only the joints configured in :attr:`asset_cfg.joint_ids` will have their positions returned.
"""
# extract the used quantities (to enable type-hinting)
asset: Articulation = env.scene[asset_cfg.name]
return asset.data.joint_pos[:, asset_cfg.joint_ids] - asset.data.default_joint_pos[:, asset_cfg.joint_ids]
[docs]def joint_pos_limit_normalized(env: BaseEnv, asset_cfg: SceneEntityCfg = SceneEntityCfg("robot")) -> torch.Tensor:
"""The joint positions of the asset normalized with the asset's joint limits.
Note: Only the joints configured in :attr:`asset_cfg.joint_ids` will have their normalized positions returned.
"""
# extract the used quantities (to enable type-hinting)
asset: Articulation = env.scene[asset_cfg.name]
return math_utils.scale_transform(
asset.data.joint_pos[:, asset_cfg.joint_ids],
asset.data.soft_joint_pos_limits[:, asset_cfg.joint_ids, 0],
asset.data.soft_joint_pos_limits[:, asset_cfg.joint_ids, 1],
)
[docs]def joint_vel(env: BaseEnv, asset_cfg: SceneEntityCfg = SceneEntityCfg("robot")):
"""The joint velocities of the asset.
Note: Only the joints configured in :attr:`asset_cfg.joint_ids` will have their velocities returned.
"""
# extract the used quantities (to enable type-hinting)
asset: Articulation = env.scene[asset_cfg.name]
return asset.data.joint_vel[:, asset_cfg.joint_ids]
[docs]def joint_vel_rel(env: BaseEnv, asset_cfg: SceneEntityCfg = SceneEntityCfg("robot")):
"""The joint velocities of the asset w.r.t. the default joint velocities.
Note: Only the joints configured in :attr:`asset_cfg.joint_ids` will have their velocities returned.
"""
# extract the used quantities (to enable type-hinting)
asset: Articulation = env.scene[asset_cfg.name]
return asset.data.joint_vel[:, asset_cfg.joint_ids] - asset.data.default_joint_vel[:, asset_cfg.joint_ids]
"""
Sensors.
"""
[docs]def height_scan(env: BaseEnv, sensor_cfg: SceneEntityCfg, offset: float = 0.5) -> torch.Tensor:
"""Height scan from the given sensor w.r.t. the sensor's frame.
The provided offset (Defaults to 0.5) is subtracted from the returned values.
"""
# extract the used quantities (to enable type-hinting)
sensor: RayCaster = env.scene.sensors[sensor_cfg.name]
# height scan: height = sensor_height - hit_point_z - offset
return sensor.data.pos_w[:, 2].unsqueeze(1) - sensor.data.ray_hits_w[..., 2] - offset
[docs]def body_incoming_wrench(env: BaseEnv, asset_cfg: SceneEntityCfg) -> torch.Tensor:
"""Incoming spatial wrench on bodies of an articulation in the simulation world frame.
This is the 6-D wrench (force and torque) applied to the body link by the incoming joint force.
"""
# extract the used quantities (to enable type-hinting)
asset: Articulation = env.scene[asset_cfg.name]
# obtain the link incoming forces in world frame
link_incoming_forces = asset.root_physx_view.get_link_incoming_joint_force()[:, asset_cfg.body_ids]
return link_incoming_forces.view(env.num_envs, -1)
"""
Actions.
"""
[docs]def last_action(env: BaseEnv, action_name: str | None = None) -> torch.Tensor:
"""The last input action to the environment.
The name of the action term for which the action is required. If None, the
entire action tensor is returned.
"""
if action_name is None:
return env.action_manager.action
else:
return env.action_manager.get_term(action_name).raw_actions
"""
Commands.
"""
[docs]def generated_commands(env: RLTaskEnv, command_name: str) -> torch.Tensor:
"""The generated command from command term in the command manager with the given name."""
return env.command_manager.get_command(command_name)
