Upload 21 files

0ab50fe verified 5 months ago

3.44 kB

# Copyright (c) 2022-2026, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md).
# All rights reserved.
#
# SPDX-License-Identifier: BSD-3-Clause

from __future__ import annotations

from typing import TYPE_CHECKING

import torch

from isaaclab.assets import RigidObject
from isaaclab.managers import SceneEntityCfg
from isaaclab.sensors import FrameTransformer
from isaaclab.utils.math import combine_frame_transforms

if TYPE_CHECKING:
    from isaaclab.envs import ManagerBasedRLEnv


def object_is_lifted(
    env: ManagerBasedRLEnv, minimal_height: float, object_cfg: SceneEntityCfg = SceneEntityCfg("object")
) -> torch.Tensor:
    """Reward the agent for lifting the object above the minimal height."""
    object: RigidObject = env.scene[object_cfg.name]
    return torch.where(object.data.root_pos_w[:, 2] > minimal_height, 1.0, 0.0)


def object_ee_distance(
    env: ManagerBasedRLEnv,
    std: float,
    object_cfg: SceneEntityCfg = SceneEntityCfg("object"),
    ee_frame_cfg: SceneEntityCfg = SceneEntityCfg("ee_frame"),
) -> torch.Tensor:
    """Reward the agent for reaching the object using tanh-kernel."""
    # extract the used quantities (to enable type-hinting)
    object: RigidObject = env.scene[object_cfg.name]
    ee_frame: FrameTransformer = env.scene[ee_frame_cfg.name]
    # Target object position: (num_envs, 3)
    cube_pos_w = object.data.root_pos_w
    # End-effector position: (num_envs, 3)
    ee_w = ee_frame.data.target_pos_w[..., 0, :]
    # Distance of the end-effector to the object: (num_envs,)
    object_ee_distance = torch.norm(cube_pos_w - ee_w, dim=1)

    return 1 - torch.tanh(object_ee_distance / std)


def object_goal_distance(
    env: ManagerBasedRLEnv,
    std: float,
    minimal_height: float,
    command_name: str,
    robot_cfg: SceneEntityCfg = SceneEntityCfg("robot"),
    object_cfg: SceneEntityCfg = SceneEntityCfg("object"),
) -> torch.Tensor:
    """Reward the agent for tracking the goal pose using tanh-kernel."""
    # extract the used quantities (to enable type-hinting)
    robot: RigidObject = env.scene[robot_cfg.name]
    object: RigidObject = env.scene[object_cfg.name]
    command = env.command_manager.get_command(command_name)
    # compute the desired position in the world frame
    des_pos_b = command[:, :3]
    des_pos_w, _ = combine_frame_transforms(robot.data.root_pos_w, robot.data.root_quat_w, des_pos_b)
    # distance of the end-effector to the object: (num_envs,)
    distance = torch.norm(des_pos_w - object.data.root_pos_w, dim=1)
    # rewarded if the object is lifted above the threshold
    return (object.data.root_pos_w[:, 2] > minimal_height) * (1 - torch.tanh(distance / std))

def object_goal_reached_bonus(
    env,
    threshold: float,
    command_name: str,
    robot_cfg=SceneEntityCfg("robot"),
    object_cfg=SceneEntityCfg("object"),
):
    """Sparse bonus: 1.0 when the object is within `threshold` of the goal, else 0."""
    import torch
    from isaaclab.utils.math import combine_frame_transforms
    robot = env.scene[robot_cfg.name]
    object = env.scene[object_cfg.name]
    command = env.command_manager.get_command(command_name)
    des_pos_b = command[:, :3]
    des_pos_w, _ = combine_frame_transforms(robot.data.root_pos_w, robot.data.root_quat_w, des_pos_b)
    distance = torch.norm(des_pos_w - object.data.root_pos_w[:, :3], dim=1)
    return (distance < threshold).float()

	# Copyright (c) 2022-2026, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md).
	# All rights reserved.
	#
	# SPDX-License-Identifier: BSD-3-Clause

	from __future__ import annotations

	from typing import TYPE_CHECKING

	import torch

	from isaaclab.assets import RigidObject
	from isaaclab.managers import SceneEntityCfg
	from isaaclab.sensors import FrameTransformer
	from isaaclab.utils.math import combine_frame_transforms

	if TYPE_CHECKING:
	from isaaclab.envs import ManagerBasedRLEnv


	def object_is_lifted(
	env: ManagerBasedRLEnv, minimal_height: float, object_cfg: SceneEntityCfg = SceneEntityCfg("object")
	) -> torch.Tensor:
	"""Reward the agent for lifting the object above the minimal height."""
	object: RigidObject = env.scene[object_cfg.name]
	return torch.where(object.data.root_pos_w[:, 2] > minimal_height, 1.0, 0.0)


	def object_ee_distance(
	env: ManagerBasedRLEnv,
	std: float,
	object_cfg: SceneEntityCfg = SceneEntityCfg("object"),
	ee_frame_cfg: SceneEntityCfg = SceneEntityCfg("ee_frame"),
	) -> torch.Tensor:
	"""Reward the agent for reaching the object using tanh-kernel."""
	# extract the used quantities (to enable type-hinting)
	object: RigidObject = env.scene[object_cfg.name]
	ee_frame: FrameTransformer = env.scene[ee_frame_cfg.name]
	# Target object position: (num_envs, 3)
	cube_pos_w = object.data.root_pos_w
	# End-effector position: (num_envs, 3)
	ee_w = ee_frame.data.target_pos_w[..., 0, :]
	# Distance of the end-effector to the object: (num_envs,)
	object_ee_distance = torch.norm(cube_pos_w - ee_w, dim=1)

	return 1 - torch.tanh(object_ee_distance / std)


	def object_goal_distance(
	env: ManagerBasedRLEnv,
	std: float,
	minimal_height: float,
	command_name: str,
	robot_cfg: SceneEntityCfg = SceneEntityCfg("robot"),
	object_cfg: SceneEntityCfg = SceneEntityCfg("object"),
	) -> torch.Tensor:
	"""Reward the agent for tracking the goal pose using tanh-kernel."""
	# extract the used quantities (to enable type-hinting)
	robot: RigidObject = env.scene[robot_cfg.name]
	object: RigidObject = env.scene[object_cfg.name]
	command = env.command_manager.get_command(command_name)
	# compute the desired position in the world frame
	des_pos_b = command[:, :3]
	des_pos_w, _ = combine_frame_transforms(robot.data.root_pos_w, robot.data.root_quat_w, des_pos_b)
	# distance of the end-effector to the object: (num_envs,)
	distance = torch.norm(des_pos_w - object.data.root_pos_w, dim=1)
	# rewarded if the object is lifted above the threshold
	return (object.data.root_pos_w[:, 2] > minimal_height) * (1 - torch.tanh(distance / std))

	def object_goal_reached_bonus(
	env,
	threshold: float,
	command_name: str,
	robot_cfg=SceneEntityCfg("robot"),
	object_cfg=SceneEntityCfg("object"),
	):
	"""Sparse bonus: 1.0 when the object is within `threshold` of the goal, else 0."""
	import torch
	from isaaclab.utils.math import combine_frame_transforms
	robot = env.scene[robot_cfg.name]
	object = env.scene[object_cfg.name]
	command = env.command_manager.get_command(command_name)
	des_pos_b = command[:, :3]
	des_pos_w, _ = combine_frame_transforms(robot.data.root_pos_w, robot.data.root_quat_w, des_pos_b)
	distance = torch.norm(des_pos_w - object.data.root_pos_w[:, :3], dim=1)
	return (distance < threshold).float()