diff --git a/source/isaaclab_tasks/changelog.d/task-cleanup-dex-part09.minor.rst b/source/isaaclab_tasks/changelog.d/task-cleanup-dex-part09.minor.rst new file mode 100644 index 000000000000..adb0e7023e1c --- /dev/null +++ b/source/isaaclab_tasks/changelog.d/task-cleanup-dex-part09.minor.rst @@ -0,0 +1,9 @@ +Added +^^^^^ + +* Added a manager-based counterpart for the state-based Shadow Hand + reorientation environment. +* Added manager-based counterparts for the Shadow Hand OpenAI feed-forward and + recurrent reorientation environments. +* Added an RSL-RL LSTM training configuration for the Shadow Hand OpenAI + reorientation environment. diff --git a/source/isaaclab_tasks/isaaclab_tasks/core/reorient/config/shadow_hand/__init__.py b/source/isaaclab_tasks/isaaclab_tasks/core/reorient/config/shadow_hand/__init__.py index 5135565424d8..5b4f21179fbc 100644 --- a/source/isaaclab_tasks/isaaclab_tasks/core/reorient/config/shadow_hand/__init__.py +++ b/source/isaaclab_tasks/isaaclab_tasks/core/reorient/config/shadow_hand/__init__.py @@ -17,6 +17,18 @@ reorient_direct_entry = "isaaclab_tasks.core.reorient.reorient_direct_env:ReorientDirectEnv" +gym.register( + id="Isaac-Reorient-Cube-Shadow", + entry_point="isaaclab.envs:ManagerBasedRLEnv", + disable_env_checker=True, + kwargs={ + "env_cfg_entry_point": f"{__name__}.shadow_hand_manager_env_cfg:ShadowHandManagerEnvCfg", + "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:ShadowHandPPORunnerCfg", + "skrl_cfg_entry_point": f"{agents.__name__}:skrl_ppo_cfg.yaml", + }, +) + gym.register( id="Isaac-Reorient-Cube-Shadow-Direct", entry_point=reorient_direct_entry, @@ -29,6 +41,29 @@ }, ) +gym.register( + id="Isaac-Reorient-Cube-Shadow-OpenAI-FF", + entry_point="isaaclab.envs:ManagerBasedRLEnv", + disable_env_checker=True, + kwargs={ + "env_cfg_entry_point": f"{__name__}.shadow_hand_manager_env_cfg:ShadowHandOpenAIManagerEnvCfg", + "rl_games_cfg_entry_point": f"{agents.__name__}:rl_games_ppo_ff_cfg.yaml", + "rsl_rl_cfg_entry_point": f"{agents.__name__}.rsl_rl_ppo_cfg:ShadowHandAsymFFPPORunnerCfg", + "skrl_cfg_entry_point": f"{agents.__name__}:skrl_ff_ppo_cfg.yaml", + }, +) + +gym.register( + id="Isaac-Reorient-Cube-Shadow-OpenAI-LSTM", + entry_point="isaaclab.envs:ManagerBasedRLEnv", + disable_env_checker=True, + kwargs={ + "env_cfg_entry_point": f"{__name__}.shadow_hand_manager_env_cfg:ShadowHandOpenAIManagerEnvCfg", + "rl_games_cfg_entry_point": f"{agents.__name__}:rl_games_ppo_lstm_cfg.yaml", + "rsl_rl_cfg_entry_point": f"{agents.__name__}.rsl_rl_ppo_cfg:ShadowHandAsymLSTMPPORunnerCfg", + }, +) + gym.register( id="Isaac-Reorient-Cube-Shadow-OpenAI-FF-Direct", entry_point=reorient_direct_entry, @@ -48,6 +83,7 @@ kwargs={ "env_cfg_entry_point": f"{__name__}.shadow_hand_env_cfg:ShadowHandOpenAIEnvCfg", "rl_games_cfg_entry_point": f"{agents.__name__}:rl_games_ppo_lstm_cfg.yaml", + "rsl_rl_cfg_entry_point": f"{agents.__name__}.rsl_rl_ppo_cfg:ShadowHandAsymLSTMPPORunnerCfg", }, ) diff --git a/source/isaaclab_tasks/isaaclab_tasks/core/reorient/config/shadow_hand/agents/rsl_rl_ppo_cfg.py b/source/isaaclab_tasks/isaaclab_tasks/core/reorient/config/shadow_hand/agents/rsl_rl_ppo_cfg.py index f5573ae71190..b5dd6e7c3071 100644 --- a/source/isaaclab_tasks/isaaclab_tasks/core/reorient/config/shadow_hand/agents/rsl_rl_ppo_cfg.py +++ b/source/isaaclab_tasks/isaaclab_tasks/core/reorient/config/shadow_hand/agents/rsl_rl_ppo_cfg.py @@ -5,7 +5,7 @@ from isaaclab.utils.configclass import configclass -from isaaclab_rl.rsl_rl import RslRlMLPModelCfg, RslRlOnPolicyRunnerCfg, RslRlPpoAlgorithmCfg +from isaaclab_rl.rsl_rl import RslRlMLPModelCfg, RslRlOnPolicyRunnerCfg, RslRlPpoAlgorithmCfg, RslRlRNNModelCfg @configclass @@ -47,6 +47,7 @@ class ShadowHandAsymFFPPORunnerCfg(RslRlOnPolicyRunnerCfg): max_iterations = 10000 save_interval = 250 experiment_name = "shadow_hand_openai_ff" + obs_groups = {"actor": ["policy"], "critic": ["critic"]} actor = RslRlMLPModelCfg( hidden_dims=[400, 400, 200, 100], activation="elu", @@ -74,12 +75,37 @@ class ShadowHandAsymFFPPORunnerCfg(RslRlOnPolicyRunnerCfg): ) +@configclass +class ShadowHandAsymLSTMPPORunnerCfg(ShadowHandAsymFFPPORunnerCfg): + """RSL-RL recurrent policy configuration for the asymmetric OpenAI observations.""" + + experiment_name = "shadow_hand_openai_lstm" + actor = RslRlRNNModelCfg( + hidden_dims=[400, 400, 200, 100], + activation="elu", + obs_normalization=True, + distribution_cfg=RslRlMLPModelCfg.GaussianDistributionCfg(init_std=1.0), + rnn_type="lstm", + rnn_hidden_dim=256, + rnn_num_layers=1, + ) + critic = RslRlRNNModelCfg( + hidden_dims=[512, 512, 256, 128], + activation="elu", + obs_normalization=True, + rnn_type="lstm", + rnn_hidden_dim=256, + rnn_num_layers=1, + ) + + @configclass class ShadowHandCameraFFPPORunnerCfg(RslRlOnPolicyRunnerCfg): num_steps_per_env = 64 max_iterations = 50000 save_interval = 250 experiment_name = "shadow_hand_camera" + obs_groups = {"actor": ["policy"], "critic": ["critic"]} actor = RslRlMLPModelCfg( hidden_dims=[1024, 512, 512, 256, 128], activation="elu", diff --git a/source/isaaclab_tasks/isaaclab_tasks/core/reorient/config/shadow_hand/shadow_hand_manager_env_cfg.py b/source/isaaclab_tasks/isaaclab_tasks/core/reorient/config/shadow_hand/shadow_hand_manager_env_cfg.py new file mode 100644 index 000000000000..0157904c95c1 --- /dev/null +++ b/source/isaaclab_tasks/isaaclab_tasks/core/reorient/config/shadow_hand/shadow_hand_manager_env_cfg.py @@ -0,0 +1,435 @@ +# 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 + +"""Manager-based counterpart of the state-based Shadow Hand reorientation task.""" + +import isaaclab.sim as sim_utils +from isaaclab.assets import AssetBaseCfg +from isaaclab.envs import ManagerBasedRLEnvCfg +from isaaclab.managers import EventTermCfg as EventTerm +from isaaclab.managers import ObservationGroupCfg as ObsGroup +from isaaclab.managers import ObservationTermCfg as ObsTerm +from isaaclab.managers import RewardTermCfg as RewTerm +from isaaclab.managers import SceneEntityCfg +from isaaclab.managers import TerminationTermCfg as DoneTerm +from isaaclab.scene import InteractiveSceneCfg +from isaaclab.sensors import JointWrenchSensorCfg +from isaaclab.sim import SimulationCfg +from isaaclab.utils.configclass import configclass + +import isaaclab_tasks.core.reorient.mdp as mdp +from isaaclab_tasks.core.reorient.config.shadow_hand.shadow_hand_env_cfg import ( + GOAL_OBJECT_CFG, + OBJECT_CFG, + OPENAI_ACTION_NOISE_CFG, + OPENAI_OBSERVATION_NOISE_CFG, + ROBOT_CFG, + SHADOW_SIM_CFG, + NewtonEventCfg, + ObjectCfg, + PhysicsCfg, + PhysxEventCfg, +) +from isaaclab_tasks.core.reorient.reorient_task_base import ( + GOAL_MARKER_POSITION, + IN_HAND_POS_OFFSET, + SHADOW_ACTUATED_JOINT_NAMES, + SHADOW_FINGERTIP_BODY_NAMES, +) +from isaaclab_tasks.utils import PresetCfg + + +@configclass +class _ShadowHandManagerSceneCfg(InteractiveSceneCfg): + """Scene shared by the Shadow Hand Manager backend alternatives.""" + + ground = AssetBaseCfg(prim_path="/World/ground", spawn=sim_utils.GroundPlaneCfg()) + robot: PresetCfg = ROBOT_CFG + object: ObjectCfg = OBJECT_CFG + light = AssetBaseCfg( + prim_path="/World/Light", + spawn=sim_utils.DomeLightCfg(intensity=2000.0, color=(0.75, 0.75, 0.75)), + ) + + +@configclass +class ShadowHandManagerSceneCfg(PresetCfg): + """Backend-specific scene cloning settings matching the Direct task.""" + + physx = _ShadowHandManagerSceneCfg( + num_envs=8192, + env_spacing=0.75, + replicate_physics=True, + clone_in_fabric=True, + ) + newton_mjwarp = _ShadowHandManagerSceneCfg( + num_envs=8192, + env_spacing=0.75, + replicate_physics=True, + clone_in_fabric=False, + ) + ovphysx = physx + newton_kamino = newton_mjwarp + default = physx + + +@configclass +class CommandsCfg: + """Object pose goal matching the Direct in-hand target.""" + + object_pose = mdp.ReorientEpisodeCommandCfg( + asset_name="object", + init_pos_offset=IN_HAND_POS_OFFSET, + update_goal_on_success=True, + orientation_success_threshold=0.1, + make_quat_unique=False, + fixed_marker_pos=GOAL_MARKER_POSITION, + goal_pose_visualizer_cfg=GOAL_OBJECT_CFG, + debug_vis=True, + ) + + +@configclass +class ActionsCfg: + """Twenty actuated Shadow Hand joints.""" + + joint_pos = mdp.EMAJointPositionToLimitsActionCfg( + asset_name="robot", + joint_names=SHADOW_ACTUATED_JOINT_NAMES, + alpha=1.0, + rescale_to_limits=True, + ) + + +@configclass +class FullStateWithoutActionCfg(ObsGroup): + """Shared first 137 dimensions of the full Shadow state.""" + + joint_pos = ObsTerm( + func=mdp.joint_pos_limit_normalized, + params={"asset_cfg": SceneEntityCfg("robot", joint_names=".*", preserve_order=False)}, + ) + joint_vel = ObsTerm( + func=mdp.joint_vel, + scale=0.2, + params={"asset_cfg": SceneEntityCfg("robot", joint_names=".*", preserve_order=False)}, + ) + object_pos = ObsTerm(func=mdp.root_pos_w, params={"asset_cfg": SceneEntityCfg("object")}) + object_quat = ObsTerm( + func=mdp.root_quat_w, + params={"asset_cfg": SceneEntityCfg("object"), "make_quat_unique": False}, + ) + object_lin_vel = ObsTerm(func=mdp.root_lin_vel_w, params={"asset_cfg": SceneEntityCfg("object")}) + object_ang_vel = ObsTerm( + func=mdp.root_ang_vel_w, + scale=0.2, + params={"asset_cfg": SceneEntityCfg("object")}, + ) + goal_pose = ObsTerm(func=mdp.generated_commands, params={"command_name": "object_pose"}) + goal_quat_diff = ObsTerm( + func=mdp.goal_quat_diff, + params={"asset_cfg": SceneEntityCfg("object"), "command_name": "object_pose", "make_quat_unique": False}, + ) + fingertip_pos = ObsTerm( + func=mdp.fingertip_pos, + params={"asset_cfg": SceneEntityCfg("robot", body_names=SHADOW_FINGERTIP_BODY_NAMES, preserve_order=False)}, + ) + fingertip_quat = ObsTerm( + func=mdp.fingertip_quat, + params={"asset_cfg": SceneEntityCfg("robot", body_names=SHADOW_FINGERTIP_BODY_NAMES, preserve_order=False)}, + ) + fingertip_vel = ObsTerm( + func=mdp.fingertip_vel, + params={"asset_cfg": SceneEntityCfg("robot", body_names=SHADOW_FINGERTIP_BODY_NAMES, preserve_order=False)}, + ) + + def __post_init__(self): + self.enable_corruption = False + self.concatenate_terms = True + + +@configclass +class ObservationsCfg: + """Full 157-dimensional state observation in Direct order.""" + + @configclass + class PolicyCfg(FullStateWithoutActionCfg): + last_action = ObsTerm(func=mdp.reorient_last_action, params={"action_name": "joint_pos"}) + + policy: PolicyCfg = PolicyCfg() + + +@configclass +class EventCfg: + """Reset distributions matching the Direct task.""" + + reset_state = EventTerm( + func=mdp.reset_reorient_state, + mode="reset", + params={ + "position_noise": 0.01, + "joint_position_noise": 0.2, + "joint_velocity_noise": 0.0, + "action_name": "joint_pos", + }, + ) + + +@configclass +class RewardsCfg: + """Direct-compatible reward and success accounting.""" + + reorient = RewTerm( + func=mdp.DirectReorientReward, + weight=1.0, + params={ + "command_name": "object_pose", + "distance_scale": -10.0, + "rotation_scale": 1.0, + "rotation_epsilon": 0.1, + "action_penalty_scale": -0.0002, + "success_tolerance": 0.1, + "success_bonus": 250.0, + "fall_distance": 0.24, + "fall_penalty": 0.0, + "averaging_factor": 0.1, + "success_count_threshold": 1, + "object_cfg": SceneEntityCfg("object"), + }, + ) + + +@configclass +class TerminationsCfg: + """Termination conditions matching the Direct task.""" + + object_out_of_reach = DoneTerm( + func=mdp.object_reorientation_out_of_reach, + params={ + "threshold": 0.24, + "command_name": "object_pose", + "object_cfg": SceneEntityCfg("object"), + }, + ) + time_out = DoneTerm(func=mdp.time_out, time_out=True) + + +@configclass +class ShadowHandManagerEnvCfg(ManagerBasedRLEnvCfg): + """Manager-based state Shadow Hand task with Direct-compatible semantics.""" + + scene: ShadowHandManagerSceneCfg = ShadowHandManagerSceneCfg() + sim: SimulationCfg = SHADOW_SIM_CFG + observations: ObservationsCfg = ObservationsCfg() + actions: ActionsCfg = ActionsCfg() + commands: CommandsCfg = CommandsCfg() + rewards: RewardsCfg = RewardsCfg() + terminations: TerminationsCfg = TerminationsCfg() + events: EventCfg = EventCfg() + + def __post_init__(self): + self.decimation = 2 + self.episode_length_s = 10.0 + self.sim.render_interval = self.decimation + self.viewer.eye = (2.0, 2.0, 2.0) + + +@configclass +class OpenAICommandsCfg: + """OpenAI goal command with its wider success tolerance.""" + + object_pose = mdp.ReorientEpisodeCommandCfg( + asset_name="object", + init_pos_offset=IN_HAND_POS_OFFSET, + update_goal_on_success=True, + orientation_success_threshold=0.4, + make_quat_unique=False, + fixed_marker_pos=GOAL_MARKER_POSITION, + goal_pose_visualizer_cfg=GOAL_OBJECT_CFG, + debug_vis=True, + ) + + +@configclass +class OpenAIActionsCfg: + """OpenAI actions with Direct-compatible EMA and stateful noise.""" + + joint_pos = mdp.NoisyEMAJointPositionToLimitsActionCfg( + asset_name="robot", + joint_names=SHADOW_ACTUATED_JOINT_NAMES, + alpha=0.3, + rescale_to_limits=True, + noise_model=OPENAI_ACTION_NOISE_CFG, + ) + + +@configclass +class OpenAIObservationsCfg: + """OpenAI 42-dimensional actor and 187-dimensional critic observations.""" + + @configclass + class PolicyCfg(ObsGroup): + openai = ObsTerm( + func=mdp.OpenAIPolicyObservation, + params={ + "command_name": "object_pose", + "action_name": "joint_pos", + "noise_model": OPENAI_OBSERVATION_NOISE_CFG, + "robot_cfg": SceneEntityCfg("robot", body_names=SHADOW_FINGERTIP_BODY_NAMES, preserve_order=False), + "object_cfg": SceneEntityCfg("object"), + }, + ) + + def __post_init__(self): + self.enable_corruption = False + self.concatenate_terms = True + + @configclass + class CriticCfg(FullStateWithoutActionCfg): + fingertip_wrench = ObsTerm( + func=mdp.fingertip_wrench, + scale=10.0, + params={ + "sensor_cfg": SceneEntityCfg( + "joint_wrench", body_names=SHADOW_FINGERTIP_BODY_NAMES, preserve_order=False + ) + }, + ) + last_action = ObsTerm(func=mdp.reorient_last_action, params={"action_name": "joint_pos"}) + + policy: PolicyCfg = PolicyCfg() + critic: CriticCfg = CriticCfg() + + +@configclass +class _ShadowHandOpenAIManagerSceneCfg(_ShadowHandManagerSceneCfg): + """Shadow Hand scene with fingertip joint-wrench sensing.""" + + joint_wrench = JointWrenchSensorCfg(prim_path="{ENV_REGEX_NS}/Robot") + + +@configclass +class ShadowHandOpenAIManagerSceneCfg(PresetCfg): + """Backend-specific OpenAI scene alternatives.""" + + physx = _ShadowHandOpenAIManagerSceneCfg( + num_envs=8192, + env_spacing=0.75, + replicate_physics=True, + clone_in_fabric=True, + ) + newton_mjwarp = _ShadowHandOpenAIManagerSceneCfg( + num_envs=8192, + env_spacing=0.75, + replicate_physics=True, + clone_in_fabric=False, + ) + ovphysx = physx + newton_kamino = newton_mjwarp + default = physx + + +_OPENAI_RESET_PARAMS = { + "position_noise": 0.01, + "joint_position_noise": 0.2, + "joint_velocity_noise": 0.0, + "action_name": "joint_pos", +} + + +@configclass +class OpenAIPhysxEventCfg(PhysxEventCfg): + """PhysX OpenAI randomization and state reset events.""" + + reset_state = EventTerm(func=mdp.reset_reorient_state, mode="reset", params=_OPENAI_RESET_PARAMS) + + +@configclass +class OpenAINewtonEventCfg(NewtonEventCfg): + """Newton OpenAI randomization and state reset events.""" + + reset_state = EventTerm(func=mdp.reset_reorient_state, mode="reset", params=_OPENAI_RESET_PARAMS) + + +@configclass +class OpenAIEventCfg(PresetCfg): + """Backend-specific OpenAI event alternatives.""" + + physx = OpenAIPhysxEventCfg() + newton_mjwarp = OpenAINewtonEventCfg() + ovphysx = physx + newton_kamino = newton_mjwarp + default = physx + + +@configclass +class OpenAIRewardsCfg: + """Direct-compatible OpenAI reward and success accounting.""" + + reorient = RewTerm( + func=mdp.DirectReorientReward, + weight=1.0, + params={ + "command_name": "object_pose", + "distance_scale": -10.0, + "rotation_scale": 1.0, + "rotation_epsilon": 0.1, + "action_penalty_scale": -0.0002, + "success_tolerance": 0.4, + "success_bonus": 250.0, + "fall_distance": 0.24, + "fall_penalty": -50.0, + "averaging_factor": 0.1, + "success_count_threshold": 1, + "action_name": "joint_pos", + "object_cfg": SceneEntityCfg("object"), + }, + ) + + +@configclass +class OpenAITerminationsCfg: + """Direct-compatible OpenAI termination conditions.""" + + object_out_of_reach = DoneTerm( + func=mdp.object_reorientation_out_of_reach, + params={ + "threshold": 0.24, + "command_name": "object_pose", + "object_cfg": SceneEntityCfg("object"), + }, + ) + time_out = DoneTerm( + func=mdp.DirectReorientTimeout, + time_out=True, + params={ + "command_name": "object_pose", + "reward_name": "reorient", + "success_tolerance": 0.4, + "max_successes": 50, + "object_cfg": SceneEntityCfg("object"), + }, + ) + + +@configclass +class ShadowHandOpenAIManagerEnvCfg(ShadowHandManagerEnvCfg): + """Manager counterpart shared by the OpenAI FF and LSTM variants.""" + + scene: ShadowHandOpenAIManagerSceneCfg = ShadowHandOpenAIManagerSceneCfg() + observations: OpenAIObservationsCfg = OpenAIObservationsCfg() + actions: OpenAIActionsCfg = OpenAIActionsCfg() + commands: OpenAICommandsCfg = OpenAICommandsCfg() + rewards: OpenAIRewardsCfg = OpenAIRewardsCfg() + terminations: OpenAITerminationsCfg = OpenAITerminationsCfg() + events: OpenAIEventCfg = OpenAIEventCfg() + + def __post_init__(self): + super().__post_init__() + self.decimation = 3 + self.episode_length_s = 8.0 + self.sim.dt = 1 / 60 + self.sim.render_interval = self.decimation + self.sim.physics = PhysicsCfg() diff --git a/source/isaaclab_tasks/isaaclab_tasks/core/reorient/mdp/__init__.pyi b/source/isaaclab_tasks/isaaclab_tasks/core/reorient/mdp/__init__.pyi index e835f887dd8a..2cd4322b640f 100644 --- a/source/isaaclab_tasks/isaaclab_tasks/core/reorient/mdp/__init__.pyi +++ b/source/isaaclab_tasks/isaaclab_tasks/core/reorient/mdp/__init__.pyi @@ -4,19 +4,63 @@ # SPDX-License-Identifier: BSD-3-Clause __all__ = [ + "NoisyEMAJointPositionToLimitsAction", + "NoisyEMAJointPositionToLimitsActionCfg", "ReorientCommand", "ReorientCommandCfg", + "ReorientEpisodeCommand", + "ReorientEpisodeCommandCfg", + "reset_reorient_state", + "fingertip_pos", + "fingertip_quat", + "fingertip_vel", + "fingertip_wrench", + "reorient_last_action", + "OpenAIPolicyObservation", "goal_quat_diff", "success_bonus", "track_orientation_inv_l2", "track_pos_l2", + "rotation_distance_kernel", + "reorient_success_kernel", + "direct_reorient_reward", + "DirectReorientReward", "max_consecutive_success", "object_away_from_goal", "object_away_from_robot", + "object_reorientation_out_of_reach", + "DirectReorientTimeout", ] -from .commands import ReorientCommand, ReorientCommandCfg -from .observations import goal_quat_diff -from .rewards import success_bonus, track_orientation_inv_l2, track_pos_l2 -from .terminations import max_consecutive_success, object_away_from_goal, object_away_from_robot +from .commands import ReorientCommand, ReorientCommandCfg, ReorientEpisodeCommand, ReorientEpisodeCommandCfg +from .events import reset_reorient_state +from .actions import ( + NoisyEMAJointPositionToLimitsAction, + NoisyEMAJointPositionToLimitsActionCfg, +) +from .observations import ( + OpenAIPolicyObservation, + fingertip_pos, + fingertip_quat, + fingertip_vel, + fingertip_wrench, + goal_quat_diff, + reorient_last_action, +) +from .rewards import ( + DirectReorientReward, + direct_reorient_reward, + reorient_success_kernel, + rotation_distance_kernel, + success_bonus, + track_orientation_inv_l2, + track_pos_l2, +) +from .terminations import ( + DirectReorientTimeout, + max_consecutive_success, + object_away_from_goal, + object_away_from_robot, + object_reorientation_out_of_reach, +) from isaaclab.envs.mdp import * diff --git a/source/isaaclab_tasks/isaaclab_tasks/core/reorient/mdp/actions.py b/source/isaaclab_tasks/isaaclab_tasks/core/reorient/mdp/actions.py new file mode 100644 index 000000000000..ee9fbcea3d78 --- /dev/null +++ b/source/isaaclab_tasks/isaaclab_tasks/core/reorient/mdp/actions.py @@ -0,0 +1,63 @@ +# 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 + +"""Action terms for the reorientation task family.""" + +from __future__ import annotations + +from collections.abc import Sequence +from dataclasses import MISSING +from typing import TYPE_CHECKING + +import torch + +from isaaclab.envs.mdp import EMAJointPositionToLimitsActionCfg +from isaaclab.envs.mdp.actions import EMAJointPositionToLimitsAction +from isaaclab.utils.configclass import configclass +from isaaclab.utils.noise import NoiseModelCfg + +if TYPE_CHECKING: + from isaaclab.envs import ManagerBasedEnv + + +class NoisyEMAJointPositionToLimitsAction(EMAJointPositionToLimitsAction): + """Apply a stateful noise model before EMA joint-position processing.""" + + def __init__(self, cfg: NoisyEMAJointPositionToLimitsActionCfg, env: ManagerBasedEnv): + """Initialize the noisy action term. + + Args: + cfg: Action configuration including the stateful noise model. + env: Manager-based environment containing the hand. + """ + super().__init__(cfg, env) + self._noise_model = cfg.noise_model.class_type(cfg.noise_model, num_envs=self.num_envs, device=self.device) + + def process_actions(self, actions: torch.Tensor) -> None: + """Apply noise to normalized actions before scaling and EMA filtering. + + Args: + actions: Normalized joint actions, shape ``(num_envs, num_actions)``. + """ + super().process_actions(self._noise_model(actions)) + + def reset(self, env_ids: Sequence[int] | None = None) -> None: + """Reset the noise state and standard EMA action buffers. + + Args: + env_ids: Environment indices to reset, or ``None`` for every environment. + """ + self._noise_model.reset(env_ids) + super().reset(env_ids) + + +@configclass +class NoisyEMAJointPositionToLimitsActionCfg(EMAJointPositionToLimitsActionCfg): + """EMA joint action configuration with Direct-compatible stateful noise.""" + + class_type = NoisyEMAJointPositionToLimitsAction + + noise_model: NoiseModelCfg = MISSING + """Stateful noise applied to incoming normalized actions.""" diff --git a/source/isaaclab_tasks/isaaclab_tasks/core/reorient/mdp/observations.py b/source/isaaclab_tasks/isaaclab_tasks/core/reorient/mdp/observations.py index b54b2beeab9f..87c7ba1aee2d 100644 --- a/source/isaaclab_tasks/isaaclab_tasks/core/reorient/mdp/observations.py +++ b/source/isaaclab_tasks/isaaclab_tasks/core/reorient/mdp/observations.py @@ -7,36 +7,318 @@ from __future__ import annotations +from collections.abc import Sequence from typing import TYPE_CHECKING import torch +import warp as wp -import isaaclab.utils.math as math_utils -from isaaclab.managers import SceneEntityCfg +from isaaclab.managers import ManagerTermBase, ObservationTermCfg, SceneEntityCfg +from isaaclab.utils.noise import NoiseModelCfg + +# CUBE_HALF_SIZE, cube_keypoints_from_quat_kernel, and compute_cube_keypoints are +# re-exported for API stability: the cube-keypoint math moved to the family math +# root in :mod:`isaaclab_tasks.core.reorient.reorient_kernels` +from isaaclab_tasks.core.reorient.reorient_kernels import ( # noqa: F401 + CUBE_HALF_SIZE, + compute_cube_keypoints, + cube_keypoints_from_quat_kernel, + fingertip_pos_kernel, + fingertip_quat_kernel, + fingertip_vel_kernel, +) if TYPE_CHECKING: from isaaclab.assets import RigidObject from isaaclab.envs import ManagerBasedRLEnv + from isaaclab.sensors import Camera, JointWrenchSensor + + from isaaclab_tasks.core.reorient.config.shadow_hand.feature_extractor import FeatureExtractorCfg from .commands import ReorientCommand -def goal_quat_diff( - env: ManagerBasedRLEnv, asset_cfg: SceneEntityCfg, command_name: str, make_quat_unique: bool -) -> torch.Tensor: +@wp.kernel +def _fingertip_wrench_kernel( + force: wp.array2d(dtype=wp.vec3f), + torque: wp.array2d(dtype=wp.vec3f), + body_ids: wp.array(dtype=wp.int32), + out: wp.array2d(dtype=wp.float32), +): + i, j = wp.tid() + f = force[i, body_ids[j]] + t = torque[i, body_ids[j]] + out[i, 6 * j + 0] = f[0] + out[i, 6 * j + 1] = f[1] + out[i, 6 * j + 2] = f[2] + out[i, 6 * j + 3] = t[0] + out[i, 6 * j + 4] = t[1] + out[i, 6 * j + 5] = t[2] + + +@wp.kernel +def _goal_quat_error_kernel( + asset_quat: wp.array(dtype=wp.quatf), + goal_quat: wp.array(dtype=wp.quatf), + make_unique: int, + out: wp.array(dtype=wp.quatf), +): + """Per-environment quaternion error ``asset * conjugate(goal)`` in (x, y, z, w) order.""" + i = wp.tid() + # quat_inverse == conjugate for these unit quaternions, matching + # isaaclab.utils.math.quat_mul/quat_conjugate semantics + qe = asset_quat[i] * wp.quat_inverse(goal_quat[i]) + sign = 1.0 + # make_unique keeps the real part non-negative (isaaclab.utils.math.quat_unique) + if make_unique != 0 and qe[3] < 0.0: + sign = -1.0 + out[i] = wp.quatf(sign * qe[0], sign * qe[1], sign * qe[2], sign * qe[3]) + + +class goal_quat_diff(ManagerTermBase): """Goal orientation relative to the asset's root frame. - The quaternion is represented as (w, x, y, z). The real part is always positive. + The real part is always positive when ``make_quat_unique`` is set. """ - # extract useful elements - asset: RigidObject = env.scene[asset_cfg.name] - command_term: ReorientCommand = env.command_manager.get_term(command_name) - - # obtain the orientations - goal_quat_w = command_term.command[:, 3:7] - asset_quat_w = asset.data.root_quat_w.torch - - # compute quaternion difference - quat = math_utils.quat_mul(asset_quat_w, math_utils.quat_conjugate(goal_quat_w)) - # make sure the quaternion real-part is always positive - return math_utils.quat_unique(quat) if make_quat_unique else quat + + def __init__(self, cfg: ObservationTermCfg, env: ManagerBasedRLEnv): + super().__init__(cfg, env) + self._out = torch.empty(env.num_envs, 4, dtype=torch.float32, device=env.device) + # cached Warp views; the hot loop launches the kernel without conversions + self._out_wp = wp.from_torch(self._out, dtype=wp.quatf) + # resolved on first call: the command term does not exist yet during manager construction + self._goal_quat_wp: wp.array | None = None + + def __call__( + self, env: ManagerBasedRLEnv, asset_cfg: SceneEntityCfg, command_name: str, make_quat_unique: bool + ) -> torch.Tensor: + """Return the per-environment quaternion error, shape ``(num_envs, 4)``.""" + asset: RigidObject = env.scene[asset_cfg.name] + if self._goal_quat_wp is None: + command_term: ReorientCommand = env.command_manager.get_term(command_name) + self._goal_quat_wp = wp.from_torch(command_term.quat_command_w, dtype=wp.quatf) + wp.launch( + _goal_quat_error_kernel, + dim=self.num_envs, + inputs=[asset.data.root_quat_w.warp, self._goal_quat_wp, make_quat_unique], + outputs=[self._out_wp], + device=self._out_wp.device, + ) + return self._out + + +class fingertip_pos(ManagerTermBase): + """Flattened fingertip positions in the environment frame [m].""" + + def __init__(self, cfg: ObservationTermCfg, env: ManagerBasedRLEnv): + super().__init__(cfg, env) + body_ids = cfg.params["asset_cfg"].body_ids + self._body_ids_wp = wp.array(body_ids, dtype=wp.int32, device=str(env.device)) + self._out = torch.empty(env.num_envs, len(body_ids) * 3, dtype=torch.float32, device=env.device) + self._out_wp = wp.from_torch(self._out) + self._env_origins_wp = wp.from_torch(env.scene.env_origins, dtype=wp.vec3f) + + def __call__(self, env: ManagerBasedRLEnv, asset_cfg: SceneEntityCfg) -> torch.Tensor: + """Return the flattened per-fingertip block, shape ``(num_envs, num_fingertips * 3)``.""" + asset = env.scene[asset_cfg.name] + wp.launch( + fingertip_pos_kernel, + dim=(self.num_envs, self._body_ids_wp.shape[0]), + inputs=[asset.data.body_pos_w.warp, self._env_origins_wp, self._body_ids_wp], + outputs=[self._out_wp], + device=self._out_wp.device, + ) + return self._out + + +class fingertip_quat(ManagerTermBase): + """Flattened fingertip ``(x, y, z, w)`` orientations.""" + + def __init__(self, cfg: ObservationTermCfg, env: ManagerBasedRLEnv): + super().__init__(cfg, env) + body_ids = cfg.params["asset_cfg"].body_ids + self._body_ids_wp = wp.array(body_ids, dtype=wp.int32, device=str(env.device)) + self._out = torch.empty(env.num_envs, len(body_ids) * 4, dtype=torch.float32, device=env.device) + self._out_wp = wp.from_torch(self._out) + + def __call__(self, env: ManagerBasedRLEnv, asset_cfg: SceneEntityCfg) -> torch.Tensor: + """Return the flattened per-fingertip block, shape ``(num_envs, num_fingertips * 4)``.""" + asset = env.scene[asset_cfg.name] + wp.launch( + fingertip_quat_kernel, + dim=(self.num_envs, self._body_ids_wp.shape[0]), + inputs=[asset.data.body_quat_w.warp, self._body_ids_wp], + outputs=[self._out_wp], + device=self._out_wp.device, + ) + return self._out + + +class fingertip_vel(ManagerTermBase): + """Flattened fingertip spatial velocities [m/s, rad/s].""" + + def __init__(self, cfg: ObservationTermCfg, env: ManagerBasedRLEnv): + super().__init__(cfg, env) + body_ids = cfg.params["asset_cfg"].body_ids + self._body_ids_wp = wp.array(body_ids, dtype=wp.int32, device=str(env.device)) + self._out = torch.empty(env.num_envs, len(body_ids) * 6, dtype=torch.float32, device=env.device) + self._out_wp = wp.from_torch(self._out) + + def __call__(self, env: ManagerBasedRLEnv, asset_cfg: SceneEntityCfg) -> torch.Tensor: + """Return the flattened per-fingertip block, shape ``(num_envs, num_fingertips * 6)``.""" + asset = env.scene[asset_cfg.name] + wp.launch( + fingertip_vel_kernel, + dim=(self.num_envs, self._body_ids_wp.shape[0]), + inputs=[asset.data.body_vel_w.warp, self._body_ids_wp], + outputs=[self._out_wp], + device=self._out_wp.device, + ) + return self._out + + +class fingertip_wrench(ManagerTermBase): + """Fingertip reaction wrenches [N, N·m] with Direct-compatible zero fallback.""" + + def __init__(self, cfg: ObservationTermCfg, env: ManagerBasedRLEnv): + super().__init__(cfg, env) + body_ids = cfg.params["sensor_cfg"].body_ids + self._body_ids_wp = wp.array(body_ids, dtype=wp.int32, device=str(env.device)) + self._out = torch.zeros(env.num_envs, len(body_ids) * 6, dtype=torch.float32, device=env.device) + self._out_wp = wp.from_torch(self._out) + + def __call__(self, env: ManagerBasedRLEnv, sensor_cfg: SceneEntityCfg) -> torch.Tensor: + """Return the flattened wrench block, shape ``(num_envs, num_fingertips * 6)``.""" + sensor: JointWrenchSensor = env.scene.sensors[sensor_cfg.name] + force_data = sensor.data.force + torque_data = sensor.data.torque + if force_data is None or torque_data is None: + # Direct-compatible fallback: report zero wrenches until the sensor produces data + return self._out + wp.launch( + _fingertip_wrench_kernel, + dim=(self.num_envs, self._body_ids_wp.shape[0]), + inputs=[force_data.warp, torque_data.warp, self._body_ids_wp], + outputs=[self._out_wp], + device=self._out_wp.device, + ) + return self._out + + +def reorient_last_action(env: ManagerBasedRLEnv, action_name: str) -> torch.Tensor: + """Return the Direct-compatible last action across same-step autoreset. + + Args: + env: Environment containing the action term and reset buffers. + action_name: Action term whose raw action is observed. + + Returns: + Raw actions, retaining each terminal action in its same-step reset observation. + """ + raw_action = env.action_manager.get_term(action_name).raw_actions + reset_action = getattr(env, "_reorient_reset_action", None) + reset_step = getattr(env, "_reorient_reset_step", None) + common_step_counter = getattr(env, "common_step_counter", None) + if reset_action is None or reset_step is None or common_step_counter is None: + return raw_action + return torch.where((reset_step == common_step_counter).unsqueeze(-1), reset_action, raw_action) + + +class OpenAIPolicyObservation(ManagerTermBase): + """Apply one stateful noise model to the concatenated OpenAI actor observation.""" + + def __init__(self, cfg: ObservationTermCfg, env: ManagerBasedRLEnv): + super().__init__(cfg, env) + noise_model: NoiseModelCfg = cfg.params["noise_model"] + self._noise_model = noise_model.class_type(noise_model, num_envs=self.num_envs, device=self.device) + self._quat_error = torch.empty(env.num_envs, 4, dtype=torch.float32, device=env.device) + robot_body_ids = cfg.params["robot_cfg"].body_ids + self._robot_body_ids_wp = wp.array(robot_body_ids, dtype=wp.int32, device=str(env.device)) + self._fingertip_buf = torch.empty(env.num_envs, len(robot_body_ids) * 3, dtype=torch.float32, device=env.device) + self._fingertip_buf_wp = wp.from_torch(self._fingertip_buf) + self._env_origins_wp = wp.from_torch(env.scene.env_origins, dtype=wp.vec3f) + # cached Warp views; the hot loop launches the kernel without conversions + self._quat_error_wp = wp.from_torch(self._quat_error, dtype=wp.quatf) + # resolved on first call: the command term does not exist yet during manager construction + self._goal_quat_wp: wp.array | None = None + # ObservationManager probes callable terms once for their shape and then + # calls reset. Keep that probe side-effect free so initialization matches + # DirectRLEnv's first noise-model reset and application. + self._shape_probe_pending = True + + def reset(self, env_ids: Sequence[int] | None = None) -> None: + """Reset the actor observation bias for selected environments. + + Args: + env_ids: Environment indices to reset, or ``None`` for every environment. + """ + if self._shape_probe_pending: + self._shape_probe_pending = False + return + self._noise_model.reset(env_ids) + + def __call__( + self, + env: ManagerBasedRLEnv, + command_name: str, + action_name: str, + noise_model: NoiseModelCfg, + robot_cfg: SceneEntityCfg, + object_cfg: SceneEntityCfg, + ) -> torch.Tensor: + """Return the corrupted 42-dimensional actor observation.""" + del noise_model + object_asset: RigidObject = env.scene[object_cfg.name] + object_pos = object_asset.data.root_pos_w.torch - env.scene.env_origins + if self._goal_quat_wp is None: + command_term: ReorientCommand = env.command_manager.get_term(command_name) + self._goal_quat_wp = wp.from_torch(command_term.quat_command_w, dtype=wp.quatf) + wp.launch( + _goal_quat_error_kernel, + dim=self.num_envs, + inputs=[object_asset.data.root_quat_w.warp, self._goal_quat_wp, 0], + outputs=[self._quat_error_wp], + device=self._quat_error_wp.device, + ) + robot = env.scene[robot_cfg.name] + wp.launch( + fingertip_pos_kernel, + dim=(self.num_envs, self._robot_body_ids_wp.shape[0]), + inputs=[robot.data.body_pos_w.warp, self._env_origins_wp, self._robot_body_ids_wp], + outputs=[self._fingertip_buf_wp], + device=self._fingertip_buf_wp.device, + ) + # Direct actor-observation order: fingertips, object position, goal quat error, last action + observation = torch.cat( + (self._fingertip_buf, object_pos, self._quat_error, reorient_last_action(env, action_name)), + dim=-1, + ) + if self._shape_probe_pending: + return observation + return self._noise_model(observation) + + +# --------------------------------------------------------------------------- +# Shadow Hand camera observation terms. +# +# These terms wrap the CNN feature pipeline defined in the shadow-hand config +# package. The config layer imports the mdp layer, so the FeatureExtractor +# machinery is imported lazily at term construction time. +# --------------------------------------------------------------------------- + + +def shadow_hand_camera_cached_features(env: ManagerBasedRLEnv) -> torch.Tensor: + """Return camera features computed by the preceding policy observation group. + + Args: + env: Environment whose policy group cached the current camera embedding. + + Returns: + Detached camera embeddings, shape ``(num_envs, 27)``. + """ + embeddings = getattr(env, "_shadow_hand_camera_embeddings", None) + if embeddings is None: + raise RuntimeError("Shadow Hand camera policy features must be computed before critic observations.") + return embeddings + + diff --git a/source/isaaclab_tasks/test/benchmarking/configs.yaml b/source/isaaclab_tasks/test/benchmarking/configs.yaml index 62a652e995e3..876c8b0212a0 100644 --- a/source/isaaclab_tasks/test/benchmarking/configs.yaml +++ b/source/isaaclab_tasks/test/benchmarking/configs.yaml @@ -240,6 +240,13 @@ full: upper_thresholds: duration: 1500 # Dexterous manipulation - Shadow hand + Isaac-Reorient-Cube-Shadow: + max_iterations: 3000 + lower_thresholds: + reward: 1000 + episode_length: 300 + upper_thresholds: + duration: 10000 Isaac-Reorient-Cube-Shadow-Direct: max_iterations: 3000 lower_thresholds: @@ -254,6 +261,13 @@ full: episode_length: 50 upper_thresholds: duration: 15000 + Isaac-Reorient-Cube-Shadow-OpenAI-FF: + max_iterations: 3000 + lower_thresholds: + reward: 1000 + episode_length: 50 + upper_thresholds: + duration: 15000 Isaac-Reorient-Cube-Shadow-OpenAI-LSTM-Direct: max_iterations: 3000 lower_thresholds: @@ -261,6 +275,13 @@ full: episode_length: 100 upper_thresholds: duration: 30000 + Isaac-Reorient-Cube-Shadow-OpenAI-LSTM: + max_iterations: 3000 + lower_thresholds: + reward: 1000 + episode_length: 100 + upper_thresholds: + duration: 30000 Isaac-Reorient-Cube-Shadow-Camera-Direct: max_iterations: 3000 lower_thresholds: