From a74907b59e2b5129b5d94a9c5931d8186727c2cf Mon Sep 17 00:00:00 2001 From: jichuanh Date: Wed, 8 Jul 2026 11:53:17 +0000 Subject: [PATCH 1/7] Add Shadow state and OpenAI manager counterparts Manager-based counterparts for the state-based and OpenAI (feed-forward and recurrent) Shadow Hand reorientation tasks, built on the shared reorientation MDP terms. Adds the RSL-RL LSTM runner configuration used by both the Direct and manager OpenAI-LSTM tasks, and observation-group mappings so asymmetric critics see the critic observation group. Validated by full manager training on PhysX and Newton (success streaks met; convergence-gap diagnosis vs Direct tracked as follow-up). --- .../task-cleanup-dex-part09.minor.rst | 9 + .../reorient/config/shadow_hand/__init__.py | 36 ++ .../shadow_hand/agents/rsl_rl_ppo_cfg.py | 28 +- .../shadow_hand_manager_actions.py | 51 +++ .../shadow_hand_manager_env_cfg.py | 424 ++++++++++++++++++ .../test/benchmarking/configs.yaml | 21 + 6 files changed, 568 insertions(+), 1 deletion(-) create mode 100644 source/isaaclab_tasks/changelog.d/task-cleanup-dex-part09.minor.rst create mode 100644 source/isaaclab_tasks/isaaclab_tasks/core/reorient/config/shadow_hand/shadow_hand_manager_actions.py create mode 100644 source/isaaclab_tasks/isaaclab_tasks/core/reorient/config/shadow_hand/shadow_hand_manager_env_cfg.py 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_actions.py b/source/isaaclab_tasks/isaaclab_tasks/core/reorient/config/shadow_hand/shadow_hand_manager_actions.py new file mode 100644 index 000000000000..676d646a6d90 --- /dev/null +++ b/source/isaaclab_tasks/isaaclab_tasks/core/reorient/config/shadow_hand/shadow_hand_manager_actions.py @@ -0,0 +1,51 @@ +# 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 + +"""Runtime-only action terms for Shadow Hand Manager environments.""" + +from __future__ import annotations + +from collections.abc import Sequence +from typing import TYPE_CHECKING + +import torch + +from isaaclab.envs.mdp.actions import EMAJointPositionToLimitsAction + +if TYPE_CHECKING: + from isaaclab.envs import ManagerBasedEnv + + from .shadow_hand_manager_env_cfg import NoisyEMAJointPositionToLimitsActionCfg + + +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) 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..ab086dfbb8a8 --- /dev/null +++ b/source/isaaclab_tasks/isaaclab_tasks/core/reorient/config/shadow_hand/shadow_hand_manager_env_cfg.py @@ -0,0 +1,424 @@ +# 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.""" + +from dataclasses import MISSING + +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.sim.spawners.materials import RigidBodyMaterialBaseCfg +from isaaclab.utils.configclass import configclass +from isaaclab.utils.noise import NoiseModelCfg + +import isaaclab_tasks.core.reorient.mdp as mdp +from isaaclab_tasks.core.reorient.config.shadow_hand.shadow_hand_env_cfg import ( + NewtonEventCfg, + ObjectCfg, + PhysicsCfg, + PhysxEventCfg, + ShadowHandEnvCfg, + ShadowHandOpenAIEnvCfg, +) +from isaaclab_tasks.utils import PresetCfg + +_DIRECT_CFG = ShadowHandEnvCfg() +_ACTUATED_JOINT_NAMES = _DIRECT_CFG.actuated_joint_names +_FINGERTIP_BODY_NAMES = _DIRECT_CFG.fingertip_body_names +_OPENAI_DIRECT_CFG = ShadowHandOpenAIEnvCfg() + + +@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 = _DIRECT_CFG.robot_cfg + object: ObjectCfg = _DIRECT_CFG.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, clone_in_fabric=True) + newton_mjwarp = _ShadowHandManagerSceneCfg(num_envs=8192, env_spacing=0.75, clone_in_fabric=False) + ovphysx = _ShadowHandManagerSceneCfg(num_envs=8192, env_spacing=0.75, clone_in_fabric=True) + default = physx + + +@configclass +class CommandsCfg: + """Object pose goal matching the Direct in-hand target.""" + + object_pose = mdp.ReorientEpisodeCommandCfg( + asset_name="object", + init_pos_offset=(0.0, 0.0, -0.04), + update_goal_on_success=True, + orientation_success_threshold=_DIRECT_CFG.success_tolerance, + make_quat_unique=False, + fixed_marker_pos=(-0.2, -0.45, 0.68), + goal_pose_visualizer_cfg=_DIRECT_CFG.goal_object_cfg, + debug_vis=True, + ) + + +@configclass +class ActionsCfg: + """Twenty actuated Shadow Hand joints.""" + + joint_pos = mdp.EMAJointPositionToLimitsActionCfg( + asset_name="robot", + joint_names=_ACTUATED_JOINT_NAMES, + alpha=_DIRECT_CFG.act_moving_average, + 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=_DIRECT_CFG.vel_obs_scale, + 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=_DIRECT_CFG.vel_obs_scale, + 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=_FINGERTIP_BODY_NAMES, preserve_order=False)}, + ) + fingertip_quat = ObsTerm( + func=mdp.fingertip_quat, + params={"asset_cfg": SceneEntityCfg("robot", body_names=_FINGERTIP_BODY_NAMES, preserve_order=False)}, + ) + fingertip_vel = ObsTerm( + func=mdp.fingertip_vel, + params={"asset_cfg": SceneEntityCfg("robot", body_names=_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": _DIRECT_CFG.reset_position_noise, + "joint_position_noise": _DIRECT_CFG.reset_dof_pos_noise, + "joint_velocity_noise": _DIRECT_CFG.reset_dof_vel_noise, + "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": _DIRECT_CFG.dist_reward_scale, + "rotation_scale": _DIRECT_CFG.rot_reward_scale, + "rotation_epsilon": _DIRECT_CFG.rot_eps, + "action_penalty_scale": _DIRECT_CFG.action_penalty_scale, + "success_tolerance": _DIRECT_CFG.success_tolerance, + "success_bonus": _DIRECT_CFG.reach_goal_bonus, + "fall_distance": _DIRECT_CFG.fall_dist, + "fall_penalty": _DIRECT_CFG.fall_penalty, + "averaging_factor": _DIRECT_CFG.av_factor, + "success_count_threshold": _DIRECT_CFG.success_count_threshold, + "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": _DIRECT_CFG.fall_dist, + "command_name": "object_pose", + "object_cfg": SceneEntityCfg("object"), + }, + ) + time_out = DoneTerm(func=mdp.direct_timeout, time_out=True) + + +@configclass +class ShadowHandManagerEnvCfg(ManagerBasedRLEnvCfg): + """Manager-based state Shadow Hand task with Direct-compatible semantics.""" + + scene: ShadowHandManagerSceneCfg = ShadowHandManagerSceneCfg() + sim: SimulationCfg = SimulationCfg( + dt=1.0 / 120.0, + render_interval=2, + physics_material=RigidBodyMaterialBaseCfg(static_friction=1.0, dynamic_friction=1.0), + physics=PhysicsCfg(), + ) + observations: ObservationsCfg = ObservationsCfg() + actions: ActionsCfg = ActionsCfg() + commands: CommandsCfg = CommandsCfg() + rewards: RewardsCfg = RewardsCfg() + terminations: TerminationsCfg = TerminationsCfg() + events: EventCfg = EventCfg() + + def __post_init__(self): + self.decimation = _DIRECT_CFG.decimation + self.episode_length_s = _DIRECT_CFG.episode_length_s + self.sim.render_interval = self.decimation + self.viewer.eye = (2.0, 2.0, 2.0) + + +@configclass +class NoisyEMAJointPositionToLimitsActionCfg(mdp.EMAJointPositionToLimitsActionCfg): + """EMA joint action configuration with Direct-compatible stateful noise.""" + + class_type = f"{__package__}.shadow_hand_manager_actions:NoisyEMAJointPositionToLimitsAction" + + noise_model: NoiseModelCfg = MISSING + """Stateful noise applied to incoming normalized actions.""" + + +@configclass +class OpenAICommandsCfg: + """OpenAI goal command with its wider success tolerance.""" + + object_pose = mdp.ReorientEpisodeCommandCfg( + asset_name="object", + init_pos_offset=(0.0, 0.0, -0.04), + update_goal_on_success=True, + orientation_success_threshold=_OPENAI_DIRECT_CFG.success_tolerance, + make_quat_unique=False, + fixed_marker_pos=(-0.2, -0.45, 0.68), + goal_pose_visualizer_cfg=_OPENAI_DIRECT_CFG.goal_object_cfg, + debug_vis=True, + ) + + +@configclass +class OpenAIActionsCfg: + """OpenAI actions with Direct-compatible EMA and stateful noise.""" + + joint_pos = NoisyEMAJointPositionToLimitsActionCfg( + asset_name="robot", + joint_names=_ACTUATED_JOINT_NAMES, + alpha=_OPENAI_DIRECT_CFG.act_moving_average, + rescale_to_limits=True, + noise_model=_OPENAI_DIRECT_CFG.action_noise_model, + ) + + +@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_DIRECT_CFG.observation_noise_model, + "robot_cfg": SceneEntityCfg("robot", body_names=_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=_OPENAI_DIRECT_CFG.force_torque_obs_scale, + params={ + "sensor_cfg": SceneEntityCfg("joint_wrench", body_names=_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, clone_in_fabric=True) + newton_mjwarp = _ShadowHandOpenAIManagerSceneCfg(num_envs=8192, env_spacing=0.75, clone_in_fabric=False) + ovphysx = _ShadowHandOpenAIManagerSceneCfg(num_envs=8192, env_spacing=0.75, clone_in_fabric=True) + default = physx + + +_OPENAI_RESET_PARAMS = { + "position_noise": _OPENAI_DIRECT_CFG.reset_position_noise, + "joint_position_noise": _OPENAI_DIRECT_CFG.reset_dof_pos_noise, + "joint_velocity_noise": _OPENAI_DIRECT_CFG.reset_dof_vel_noise, + "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 = OpenAIPhysxEventCfg() + 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": _OPENAI_DIRECT_CFG.dist_reward_scale, + "rotation_scale": _OPENAI_DIRECT_CFG.rot_reward_scale, + "rotation_epsilon": _OPENAI_DIRECT_CFG.rot_eps, + "action_penalty_scale": _OPENAI_DIRECT_CFG.action_penalty_scale, + "success_tolerance": _OPENAI_DIRECT_CFG.success_tolerance, + "success_bonus": _OPENAI_DIRECT_CFG.reach_goal_bonus, + "fall_distance": _OPENAI_DIRECT_CFG.fall_dist, + "fall_penalty": _OPENAI_DIRECT_CFG.fall_penalty, + "averaging_factor": _OPENAI_DIRECT_CFG.av_factor, + "success_count_threshold": _OPENAI_DIRECT_CFG.success_count_threshold, + "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": _OPENAI_DIRECT_CFG.fall_dist, + "command_name": "object_pose", + "object_cfg": SceneEntityCfg("object"), + }, + ) + time_out = DoneTerm( + func=mdp.direct_reorient_timeout, + time_out=True, + params={ + "command_name": "object_pose", + "reward_name": "reorient", + "success_tolerance": _OPENAI_DIRECT_CFG.success_tolerance, + "max_successes": _OPENAI_DIRECT_CFG.max_consecutive_success, + "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 = _OPENAI_DIRECT_CFG.decimation + self.episode_length_s = _OPENAI_DIRECT_CFG.episode_length_s + self.sim.dt = _OPENAI_DIRECT_CFG.sim.dt + self.sim.render_interval = self.decimation + self.sim.physics = PhysicsCfg() 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: From 3ac2a7c27ac795b26295591e21d4fe83d56025d6 Mon Sep 17 00:00:00 2001 From: jichuanh Date: Thu, 9 Jul 2026 04:26:45 +0000 Subject: [PATCH 2/7] Apply review updates to the Shadow managers The noisy EMA action term and its cfg move into the task mdp package, replicate_physics is forwarded on the manager scenes, the newton_kamino alias resolves the Kamino backend, and the goal-marker position comes from the shared constants module. --- .../shadow_hand_manager_env_cfg.py | 43 +++-- .../core/reorient/mdp/__init__.pyi | 56 +++++- .../actions.py} | 18 +- .../core/reorient/mdp/observations.py | 163 +++++++++++++++++- 4 files changed, 249 insertions(+), 31 deletions(-) rename source/isaaclab_tasks/isaaclab_tasks/core/reorient/{config/shadow_hand/shadow_hand_manager_actions.py => mdp/actions.py} (74%) 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 index ab086dfbb8a8..dd441b0552fb 100644 --- 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 @@ -5,8 +5,6 @@ """Manager-based counterpart of the state-based Shadow Hand reorientation task.""" -from dataclasses import MISSING - import isaaclab.sim as sim_utils from isaaclab.assets import AssetBaseCfg from isaaclab.envs import ManagerBasedRLEnvCfg @@ -21,7 +19,6 @@ from isaaclab.sim import SimulationCfg from isaaclab.sim.spawners.materials import RigidBodyMaterialBaseCfg from isaaclab.utils.configclass import configclass -from isaaclab.utils.noise import NoiseModelCfg import isaaclab_tasks.core.reorient.mdp as mdp from isaaclab_tasks.core.reorient.config.shadow_hand.shadow_hand_env_cfg import ( @@ -32,6 +29,7 @@ ShadowHandEnvCfg, ShadowHandOpenAIEnvCfg, ) +from isaaclab_tasks.core.reorient.reorient_task_constants import GOAL_MARKER_POSITION from isaaclab_tasks.utils import PresetCfg _DIRECT_CFG = ShadowHandEnvCfg() @@ -57,9 +55,12 @@ class _ShadowHandManagerSceneCfg(InteractiveSceneCfg): class ShadowHandManagerSceneCfg(PresetCfg): """Backend-specific scene cloning settings matching the Direct task.""" - physx = _ShadowHandManagerSceneCfg(num_envs=8192, env_spacing=0.75, clone_in_fabric=True) - newton_mjwarp = _ShadowHandManagerSceneCfg(num_envs=8192, env_spacing=0.75, clone_in_fabric=False) - ovphysx = _ShadowHandManagerSceneCfg(num_envs=8192, env_spacing=0.75, clone_in_fabric=True) + 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 = _ShadowHandManagerSceneCfg(num_envs=8192, env_spacing=0.75, replicate_physics=True, clone_in_fabric=True) + newton_kamino = newton_mjwarp default = physx @@ -73,7 +74,7 @@ class CommandsCfg: update_goal_on_success=True, orientation_success_threshold=_DIRECT_CFG.success_tolerance, make_quat_unique=False, - fixed_marker_pos=(-0.2, -0.45, 0.68), + fixed_marker_pos=GOAL_MARKER_POSITION, goal_pose_visualizer_cfg=_DIRECT_CFG.goal_object_cfg, debug_vis=True, ) @@ -229,16 +230,6 @@ def __post_init__(self): self.viewer.eye = (2.0, 2.0, 2.0) -@configclass -class NoisyEMAJointPositionToLimitsActionCfg(mdp.EMAJointPositionToLimitsActionCfg): - """EMA joint action configuration with Direct-compatible stateful noise.""" - - class_type = f"{__package__}.shadow_hand_manager_actions:NoisyEMAJointPositionToLimitsAction" - - noise_model: NoiseModelCfg = MISSING - """Stateful noise applied to incoming normalized actions.""" - - @configclass class OpenAICommandsCfg: """OpenAI goal command with its wider success tolerance.""" @@ -249,7 +240,7 @@ class OpenAICommandsCfg: update_goal_on_success=True, orientation_success_threshold=_OPENAI_DIRECT_CFG.success_tolerance, make_quat_unique=False, - fixed_marker_pos=(-0.2, -0.45, 0.68), + fixed_marker_pos=GOAL_MARKER_POSITION, goal_pose_visualizer_cfg=_OPENAI_DIRECT_CFG.goal_object_cfg, debug_vis=True, ) @@ -259,7 +250,7 @@ class OpenAICommandsCfg: class OpenAIActionsCfg: """OpenAI actions with Direct-compatible EMA and stateful noise.""" - joint_pos = NoisyEMAJointPositionToLimitsActionCfg( + joint_pos = mdp.NoisyEMAJointPositionToLimitsActionCfg( asset_name="robot", joint_names=_ACTUATED_JOINT_NAMES, alpha=_OPENAI_DIRECT_CFG.act_moving_average, @@ -315,9 +306,16 @@ class _ShadowHandOpenAIManagerSceneCfg(_ShadowHandManagerSceneCfg): class ShadowHandOpenAIManagerSceneCfg(PresetCfg): """Backend-specific OpenAI scene alternatives.""" - physx = _ShadowHandOpenAIManagerSceneCfg(num_envs=8192, env_spacing=0.75, clone_in_fabric=True) - newton_mjwarp = _ShadowHandOpenAIManagerSceneCfg(num_envs=8192, env_spacing=0.75, clone_in_fabric=False) - ovphysx = _ShadowHandOpenAIManagerSceneCfg(num_envs=8192, env_spacing=0.75, clone_in_fabric=True) + 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 = _ShadowHandOpenAIManagerSceneCfg( + num_envs=8192, env_spacing=0.75, replicate_physics=True, clone_in_fabric=True + ) + newton_kamino = newton_mjwarp default = physx @@ -350,6 +348,7 @@ class OpenAIEventCfg(PresetCfg): physx = OpenAIPhysxEventCfg() newton_mjwarp = OpenAINewtonEventCfg() ovphysx = OpenAIPhysxEventCfg() + newton_kamino = newton_mjwarp default = physx 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..4ee0f1873449 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,67 @@ # 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", + "openai_policy_observation", + "OpenAIPolicyObservation", "goal_quat_diff", "success_bonus", "track_orientation_inv_l2", "track_pos_l2", + "direct_reorient_rotation_distance", + "evaluate_reorient_success", + "direct_reorient_reward", + "DirectReorientReward", "max_consecutive_success", "object_away_from_goal", "object_away_from_robot", + "object_reorientation_out_of_reach", + "direct_timeout", + "direct_reorient_timeout", ] -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, + openai_policy_observation, + reorient_last_action, +) +from .rewards import ( + DirectReorientReward, + direct_reorient_reward, + direct_reorient_rotation_distance, + evaluate_reorient_success, + success_bonus, + track_orientation_inv_l2, + track_pos_l2, +) +from .terminations import ( + direct_reorient_timeout, + direct_timeout, + 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/config/shadow_hand/shadow_hand_manager_actions.py b/source/isaaclab_tasks/isaaclab_tasks/core/reorient/mdp/actions.py similarity index 74% rename from source/isaaclab_tasks/isaaclab_tasks/core/reorient/config/shadow_hand/shadow_hand_manager_actions.py rename to source/isaaclab_tasks/isaaclab_tasks/core/reorient/mdp/actions.py index 676d646a6d90..ee9fbcea3d78 100644 --- a/source/isaaclab_tasks/isaaclab_tasks/core/reorient/config/shadow_hand/shadow_hand_manager_actions.py +++ b/source/isaaclab_tasks/isaaclab_tasks/core/reorient/mdp/actions.py @@ -3,22 +3,24 @@ # # SPDX-License-Identifier: BSD-3-Clause -"""Runtime-only action terms for Shadow Hand Manager environments.""" +"""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 - from .shadow_hand_manager_env_cfg import NoisyEMAJointPositionToLimitsActionCfg - class NoisyEMAJointPositionToLimitsAction(EMAJointPositionToLimitsAction): """Apply a stateful noise model before EMA joint-position processing.""" @@ -49,3 +51,13 @@ def reset(self, env_ids: Sequence[int] | None = None) -> None: """ 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..8e7f77303764 100644 --- a/source/isaaclab_tasks/isaaclab_tasks/core/reorient/mdp/observations.py +++ b/source/isaaclab_tasks/isaaclab_tasks/core/reorient/mdp/observations.py @@ -7,16 +7,19 @@ from __future__ import annotations +from collections.abc import Sequence from typing import TYPE_CHECKING import torch 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 if TYPE_CHECKING: from isaaclab.assets import RigidObject from isaaclab.envs import ManagerBasedRLEnv + from isaaclab.sensors import JointWrenchSensor from .commands import ReorientCommand @@ -26,7 +29,7 @@ def goal_quat_diff( ) -> torch.Tensor: """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 quaternion is represented as (x, y, z, w). The real part is always positive. """ # extract useful elements asset: RigidObject = env.scene[asset_cfg.name] @@ -40,3 +43,159 @@ def goal_quat_diff( 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 fingertip_pos(env: ManagerBasedRLEnv, asset_cfg: SceneEntityCfg) -> torch.Tensor: + """Return flattened fingertip positions in the environment frame [m]. + + Args: + env: Environment containing the hand. + asset_cfg: Hand entity with resolved fingertip body indices. + + Returns: + Fingertip positions [m], shape ``(num_envs, num_fingertips * 3)``. + """ + asset = env.scene[asset_cfg.name] + positions = asset.data.body_pos_w.torch[:, asset_cfg.body_ids] + positions = positions - env.scene.env_origins[:, None, :] + return positions.flatten(start_dim=1) + + +def fingertip_quat(env: ManagerBasedRLEnv, asset_cfg: SceneEntityCfg) -> torch.Tensor: + """Return flattened fingertip ``(x, y, z, w)`` orientations. + + Args: + env: Environment containing the hand. + asset_cfg: Hand entity with resolved fingertip body indices. + + Returns: + Unit quaternions, shape ``(num_envs, num_fingertips * 4)``. + """ + asset = env.scene[asset_cfg.name] + return asset.data.body_quat_w.torch[:, asset_cfg.body_ids].flatten(start_dim=1) + + +def fingertip_vel(env: ManagerBasedRLEnv, asset_cfg: SceneEntityCfg) -> torch.Tensor: + """Return flattened fingertip spatial velocities in the world frame. + + Args: + env: Environment containing the hand. + asset_cfg: Hand entity with resolved fingertip body indices. + + Returns: + Spatial velocities [m/s, rad/s], shape ``(num_envs, num_fingertips * 6)``. + """ + asset = env.scene[asset_cfg.name] + return asset.data.body_vel_w.torch[:, asset_cfg.body_ids].flatten(start_dim=1) + + +def fingertip_wrench(env: ManagerBasedRLEnv, sensor_cfg: SceneEntityCfg) -> torch.Tensor: + """Return fingertip reaction wrenches with Direct-compatible zero fallback. + + Args: + env: Environment containing the joint-wrench sensor. + sensor_cfg: Joint-wrench sensor entity with resolved fingertip body indices. + + Returns: + Fingertip reaction wrenches [N, N·m], 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: + body_count = len(sensor_cfg.body_ids) + return torch.zeros(env.num_envs, body_count * 6, device=env.device) + force = force_data.torch[:, sensor_cfg.body_ids] + torque = torque_data.torch[:, sensor_cfg.body_ids] + return torch.cat((force, torque), dim=-1).flatten(start_dim=1) + + +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) + + +def openai_policy_observation( + env: ManagerBasedRLEnv, + command_name: str, + action_name: str, + robot_cfg: SceneEntityCfg, + object_cfg: SceneEntityCfg, +) -> torch.Tensor: + """Build the Direct OpenAI actor observation before corruption. + + Args: + env: Environment containing the hand, object, command, and action term. + command_name: Goal command term name. + action_name: Action term whose raw action is observed. + robot_cfg: Hand entity with resolved fingertip body indices. + object_cfg: Object scene entity. + + Returns: + Actor observation in Direct order, shape ``(num_envs, 42)``. + """ + object_asset: RigidObject = env.scene[object_cfg.name] + object_pos = object_asset.data.root_pos_w.torch - env.scene.env_origins + return torch.cat( + ( + fingertip_pos(env, robot_cfg), + object_pos, + goal_quat_diff(env, object_cfg, command_name, make_quat_unique=False), + reorient_last_action(env, action_name), + ), + dim=-1, + ) + + +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) + # 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 + observation = openai_policy_observation(env, command_name, action_name, robot_cfg, object_cfg) + if self._shape_probe_pending: + return observation + return self._noise_model(observation) From bf14fd89267166ffd04d0fe651f359170d9f16c3 Mon Sep 17 00:00:00 2001 From: jichuanh Date: Thu, 9 Jul 2026 08:23:32 +0000 Subject: [PATCH 3/7] Consume the task constants in the Shadow managers The Shadow and OpenAI manager cfgs read the shared scalar values and noise-model cfgs from the constants module and the Direct cfg module scope, the parity-only direct_timeout custom is replaced by the framework time_out term, and the OpenAI policy observation folds into its term class with a Warp quat-error kernel and caller-owned buffer. --- .../shadow_hand_manager_env_cfg.py | 178 +++++++++++------- .../core/reorient/mdp/__init__.pyi | 10 +- .../core/reorient/mdp/observations.py | 129 ++++++++----- 3 files changed, 194 insertions(+), 123 deletions(-) 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 index dd441b0552fb..06e298925b15 100644 --- 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 @@ -22,29 +22,67 @@ 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, NewtonEventCfg, ObjectCfg, PhysicsCfg, PhysxEventCfg, - ShadowHandEnvCfg, - ShadowHandOpenAIEnvCfg, ) -from isaaclab_tasks.core.reorient.reorient_task_constants import GOAL_MARKER_POSITION +from isaaclab_tasks.core.reorient.reorient_task_constants import ( + GOAL_MARKER_POSITION, + OPENAI_ACT_MOVING_AVERAGE, + OPENAI_ACTION_PENALTY_SCALE, + OPENAI_AV_FACTOR, + OPENAI_DECIMATION, + OPENAI_DIST_REWARD_SCALE, + OPENAI_EPISODE_LENGTH_S, + OPENAI_FALL_DIST, + OPENAI_FALL_PENALTY, + OPENAI_FORCE_TORQUE_OBS_SCALE, + OPENAI_MAX_CONSECUTIVE_SUCCESS, + OPENAI_REACH_GOAL_BONUS, + OPENAI_RESET_DOF_POS_NOISE, + OPENAI_RESET_DOF_VEL_NOISE, + OPENAI_RESET_POSITION_NOISE, + OPENAI_ROT_EPS, + OPENAI_ROT_REWARD_SCALE, + OPENAI_SIM_DT, + OPENAI_SUCCESS_COUNT_THRESHOLD, + OPENAI_SUCCESS_TOLERANCE, + SHADOW_ACT_MOVING_AVERAGE, + SHADOW_ACTION_PENALTY_SCALE, + SHADOW_ACTUATED_JOINT_NAMES, + SHADOW_AV_FACTOR, + SHADOW_DECIMATION, + SHADOW_DIST_REWARD_SCALE, + SHADOW_EPISODE_LENGTH_S, + SHADOW_FALL_DIST, + SHADOW_FALL_PENALTY, + SHADOW_FINGERTIP_BODY_NAMES, + SHADOW_REACH_GOAL_BONUS, + SHADOW_RESET_DOF_POS_NOISE, + SHADOW_RESET_DOF_VEL_NOISE, + SHADOW_RESET_POSITION_NOISE, + SHADOW_ROT_EPS, + SHADOW_ROT_REWARD_SCALE, + SHADOW_SUCCESS_COUNT_THRESHOLD, + SHADOW_SUCCESS_TOLERANCE, + SHADOW_VEL_OBS_SCALE, +) from isaaclab_tasks.utils import PresetCfg -_DIRECT_CFG = ShadowHandEnvCfg() -_ACTUATED_JOINT_NAMES = _DIRECT_CFG.actuated_joint_names -_FINGERTIP_BODY_NAMES = _DIRECT_CFG.fingertip_body_names -_OPENAI_DIRECT_CFG = ShadowHandOpenAIEnvCfg() - @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 = _DIRECT_CFG.robot_cfg - object: ObjectCfg = _DIRECT_CFG.object_cfg + 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)), @@ -59,7 +97,7 @@ class ShadowHandManagerSceneCfg(PresetCfg): newton_mjwarp = _ShadowHandManagerSceneCfg( num_envs=8192, env_spacing=0.75, replicate_physics=True, clone_in_fabric=False ) - ovphysx = _ShadowHandManagerSceneCfg(num_envs=8192, env_spacing=0.75, replicate_physics=True, clone_in_fabric=True) + ovphysx = physx newton_kamino = newton_mjwarp default = physx @@ -72,10 +110,10 @@ class CommandsCfg: asset_name="object", init_pos_offset=(0.0, 0.0, -0.04), update_goal_on_success=True, - orientation_success_threshold=_DIRECT_CFG.success_tolerance, + orientation_success_threshold=SHADOW_SUCCESS_TOLERANCE, make_quat_unique=False, fixed_marker_pos=GOAL_MARKER_POSITION, - goal_pose_visualizer_cfg=_DIRECT_CFG.goal_object_cfg, + goal_pose_visualizer_cfg=GOAL_OBJECT_CFG, debug_vis=True, ) @@ -86,8 +124,8 @@ class ActionsCfg: joint_pos = mdp.EMAJointPositionToLimitsActionCfg( asset_name="robot", - joint_names=_ACTUATED_JOINT_NAMES, - alpha=_DIRECT_CFG.act_moving_average, + joint_names=SHADOW_ACTUATED_JOINT_NAMES, + alpha=SHADOW_ACT_MOVING_AVERAGE, rescale_to_limits=True, ) @@ -102,7 +140,7 @@ class FullStateWithoutActionCfg(ObsGroup): ) joint_vel = ObsTerm( func=mdp.joint_vel, - scale=_DIRECT_CFG.vel_obs_scale, + scale=SHADOW_VEL_OBS_SCALE, params={"asset_cfg": SceneEntityCfg("robot", joint_names=".*", preserve_order=False)}, ) object_pos = ObsTerm(func=mdp.root_pos_w, params={"asset_cfg": SceneEntityCfg("object")}) @@ -113,7 +151,7 @@ class FullStateWithoutActionCfg(ObsGroup): 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=_DIRECT_CFG.vel_obs_scale, + scale=SHADOW_VEL_OBS_SCALE, params={"asset_cfg": SceneEntityCfg("object")}, ) goal_pose = ObsTerm(func=mdp.generated_commands, params={"command_name": "object_pose"}) @@ -123,15 +161,15 @@ class FullStateWithoutActionCfg(ObsGroup): ) fingertip_pos = ObsTerm( func=mdp.fingertip_pos, - params={"asset_cfg": SceneEntityCfg("robot", body_names=_FINGERTIP_BODY_NAMES, preserve_order=False)}, + 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=_FINGERTIP_BODY_NAMES, preserve_order=False)}, + 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=_FINGERTIP_BODY_NAMES, preserve_order=False)}, + params={"asset_cfg": SceneEntityCfg("robot", body_names=SHADOW_FINGERTIP_BODY_NAMES, preserve_order=False)}, ) def __post_init__(self): @@ -158,9 +196,9 @@ class EventCfg: func=mdp.reset_reorient_state, mode="reset", params={ - "position_noise": _DIRECT_CFG.reset_position_noise, - "joint_position_noise": _DIRECT_CFG.reset_dof_pos_noise, - "joint_velocity_noise": _DIRECT_CFG.reset_dof_vel_noise, + "position_noise": SHADOW_RESET_POSITION_NOISE, + "joint_position_noise": SHADOW_RESET_DOF_POS_NOISE, + "joint_velocity_noise": SHADOW_RESET_DOF_VEL_NOISE, "action_name": "joint_pos", }, ) @@ -175,16 +213,16 @@ class RewardsCfg: weight=1.0, params={ "command_name": "object_pose", - "distance_scale": _DIRECT_CFG.dist_reward_scale, - "rotation_scale": _DIRECT_CFG.rot_reward_scale, - "rotation_epsilon": _DIRECT_CFG.rot_eps, - "action_penalty_scale": _DIRECT_CFG.action_penalty_scale, - "success_tolerance": _DIRECT_CFG.success_tolerance, - "success_bonus": _DIRECT_CFG.reach_goal_bonus, - "fall_distance": _DIRECT_CFG.fall_dist, - "fall_penalty": _DIRECT_CFG.fall_penalty, - "averaging_factor": _DIRECT_CFG.av_factor, - "success_count_threshold": _DIRECT_CFG.success_count_threshold, + "distance_scale": SHADOW_DIST_REWARD_SCALE, + "rotation_scale": SHADOW_ROT_REWARD_SCALE, + "rotation_epsilon": SHADOW_ROT_EPS, + "action_penalty_scale": SHADOW_ACTION_PENALTY_SCALE, + "success_tolerance": SHADOW_SUCCESS_TOLERANCE, + "success_bonus": SHADOW_REACH_GOAL_BONUS, + "fall_distance": SHADOW_FALL_DIST, + "fall_penalty": SHADOW_FALL_PENALTY, + "averaging_factor": SHADOW_AV_FACTOR, + "success_count_threshold": SHADOW_SUCCESS_COUNT_THRESHOLD, "object_cfg": SceneEntityCfg("object"), }, ) @@ -197,12 +235,12 @@ class TerminationsCfg: object_out_of_reach = DoneTerm( func=mdp.object_reorientation_out_of_reach, params={ - "threshold": _DIRECT_CFG.fall_dist, + "threshold": SHADOW_FALL_DIST, "command_name": "object_pose", "object_cfg": SceneEntityCfg("object"), }, ) - time_out = DoneTerm(func=mdp.direct_timeout, time_out=True) + time_out = DoneTerm(func=mdp.time_out, time_out=True) @configclass @@ -224,8 +262,8 @@ class ShadowHandManagerEnvCfg(ManagerBasedRLEnvCfg): events: EventCfg = EventCfg() def __post_init__(self): - self.decimation = _DIRECT_CFG.decimation - self.episode_length_s = _DIRECT_CFG.episode_length_s + self.decimation = SHADOW_DECIMATION + self.episode_length_s = SHADOW_EPISODE_LENGTH_S self.sim.render_interval = self.decimation self.viewer.eye = (2.0, 2.0, 2.0) @@ -238,10 +276,10 @@ class OpenAICommandsCfg: asset_name="object", init_pos_offset=(0.0, 0.0, -0.04), update_goal_on_success=True, - orientation_success_threshold=_OPENAI_DIRECT_CFG.success_tolerance, + orientation_success_threshold=OPENAI_SUCCESS_TOLERANCE, make_quat_unique=False, fixed_marker_pos=GOAL_MARKER_POSITION, - goal_pose_visualizer_cfg=_OPENAI_DIRECT_CFG.goal_object_cfg, + goal_pose_visualizer_cfg=GOAL_OBJECT_CFG, debug_vis=True, ) @@ -252,10 +290,10 @@ class OpenAIActionsCfg: joint_pos = mdp.NoisyEMAJointPositionToLimitsActionCfg( asset_name="robot", - joint_names=_ACTUATED_JOINT_NAMES, - alpha=_OPENAI_DIRECT_CFG.act_moving_average, + joint_names=SHADOW_ACTUATED_JOINT_NAMES, + alpha=OPENAI_ACT_MOVING_AVERAGE, rescale_to_limits=True, - noise_model=_OPENAI_DIRECT_CFG.action_noise_model, + noise_model=OPENAI_ACTION_NOISE_CFG, ) @@ -270,8 +308,8 @@ class PolicyCfg(ObsGroup): params={ "command_name": "object_pose", "action_name": "joint_pos", - "noise_model": _OPENAI_DIRECT_CFG.observation_noise_model, - "robot_cfg": SceneEntityCfg("robot", body_names=_FINGERTIP_BODY_NAMES, preserve_order=False), + "noise_model": OPENAI_OBSERVATION_NOISE_CFG, + "robot_cfg": SceneEntityCfg("robot", body_names=SHADOW_FINGERTIP_BODY_NAMES, preserve_order=False), "object_cfg": SceneEntityCfg("object"), }, ) @@ -284,9 +322,11 @@ def __post_init__(self): class CriticCfg(FullStateWithoutActionCfg): fingertip_wrench = ObsTerm( func=mdp.fingertip_wrench, - scale=_OPENAI_DIRECT_CFG.force_torque_obs_scale, + scale=OPENAI_FORCE_TORQUE_OBS_SCALE, params={ - "sensor_cfg": SceneEntityCfg("joint_wrench", body_names=_FINGERTIP_BODY_NAMES, preserve_order=False) + "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"}) @@ -312,17 +352,15 @@ class ShadowHandOpenAIManagerSceneCfg(PresetCfg): newton_mjwarp = _ShadowHandOpenAIManagerSceneCfg( num_envs=8192, env_spacing=0.75, replicate_physics=True, clone_in_fabric=False ) - ovphysx = _ShadowHandOpenAIManagerSceneCfg( - num_envs=8192, env_spacing=0.75, replicate_physics=True, clone_in_fabric=True - ) + ovphysx = physx newton_kamino = newton_mjwarp default = physx _OPENAI_RESET_PARAMS = { - "position_noise": _OPENAI_DIRECT_CFG.reset_position_noise, - "joint_position_noise": _OPENAI_DIRECT_CFG.reset_dof_pos_noise, - "joint_velocity_noise": _OPENAI_DIRECT_CFG.reset_dof_vel_noise, + "position_noise": OPENAI_RESET_POSITION_NOISE, + "joint_position_noise": OPENAI_RESET_DOF_POS_NOISE, + "joint_velocity_noise": OPENAI_RESET_DOF_VEL_NOISE, "action_name": "joint_pos", } @@ -347,7 +385,7 @@ class OpenAIEventCfg(PresetCfg): physx = OpenAIPhysxEventCfg() newton_mjwarp = OpenAINewtonEventCfg() - ovphysx = OpenAIPhysxEventCfg() + ovphysx = physx newton_kamino = newton_mjwarp default = physx @@ -361,16 +399,16 @@ class OpenAIRewardsCfg: weight=1.0, params={ "command_name": "object_pose", - "distance_scale": _OPENAI_DIRECT_CFG.dist_reward_scale, - "rotation_scale": _OPENAI_DIRECT_CFG.rot_reward_scale, - "rotation_epsilon": _OPENAI_DIRECT_CFG.rot_eps, - "action_penalty_scale": _OPENAI_DIRECT_CFG.action_penalty_scale, - "success_tolerance": _OPENAI_DIRECT_CFG.success_tolerance, - "success_bonus": _OPENAI_DIRECT_CFG.reach_goal_bonus, - "fall_distance": _OPENAI_DIRECT_CFG.fall_dist, - "fall_penalty": _OPENAI_DIRECT_CFG.fall_penalty, - "averaging_factor": _OPENAI_DIRECT_CFG.av_factor, - "success_count_threshold": _OPENAI_DIRECT_CFG.success_count_threshold, + "distance_scale": OPENAI_DIST_REWARD_SCALE, + "rotation_scale": OPENAI_ROT_REWARD_SCALE, + "rotation_epsilon": OPENAI_ROT_EPS, + "action_penalty_scale": OPENAI_ACTION_PENALTY_SCALE, + "success_tolerance": OPENAI_SUCCESS_TOLERANCE, + "success_bonus": OPENAI_REACH_GOAL_BONUS, + "fall_distance": OPENAI_FALL_DIST, + "fall_penalty": OPENAI_FALL_PENALTY, + "averaging_factor": OPENAI_AV_FACTOR, + "success_count_threshold": OPENAI_SUCCESS_COUNT_THRESHOLD, "action_name": "joint_pos", "object_cfg": SceneEntityCfg("object"), }, @@ -384,19 +422,19 @@ class OpenAITerminationsCfg: object_out_of_reach = DoneTerm( func=mdp.object_reorientation_out_of_reach, params={ - "threshold": _OPENAI_DIRECT_CFG.fall_dist, + "threshold": OPENAI_FALL_DIST, "command_name": "object_pose", "object_cfg": SceneEntityCfg("object"), }, ) time_out = DoneTerm( - func=mdp.direct_reorient_timeout, + func=mdp.DirectReorientTimeout, time_out=True, params={ "command_name": "object_pose", "reward_name": "reorient", - "success_tolerance": _OPENAI_DIRECT_CFG.success_tolerance, - "max_successes": _OPENAI_DIRECT_CFG.max_consecutive_success, + "success_tolerance": OPENAI_SUCCESS_TOLERANCE, + "max_successes": OPENAI_MAX_CONSECUTIVE_SUCCESS, "object_cfg": SceneEntityCfg("object"), }, ) @@ -416,8 +454,8 @@ class ShadowHandOpenAIManagerEnvCfg(ShadowHandManagerEnvCfg): def __post_init__(self): super().__post_init__() - self.decimation = _OPENAI_DIRECT_CFG.decimation - self.episode_length_s = _OPENAI_DIRECT_CFG.episode_length_s - self.sim.dt = _OPENAI_DIRECT_CFG.sim.dt + self.decimation = OPENAI_DECIMATION + self.episode_length_s = OPENAI_EPISODE_LENGTH_S + self.sim.dt = OPENAI_SIM_DT 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 4ee0f1873449..e8d7aba1b5f4 100644 --- a/source/isaaclab_tasks/isaaclab_tasks/core/reorient/mdp/__init__.pyi +++ b/source/isaaclab_tasks/isaaclab_tasks/core/reorient/mdp/__init__.pyi @@ -16,9 +16,9 @@ __all__ = [ "fingertip_vel", "fingertip_wrench", "reorient_last_action", - "openai_policy_observation", "OpenAIPolicyObservation", "goal_quat_diff", + "compute_goal_quat_error", "success_bonus", "track_orientation_inv_l2", "track_pos_l2", @@ -30,8 +30,7 @@ __all__ = [ "object_away_from_goal", "object_away_from_robot", "object_reorientation_out_of_reach", - "direct_timeout", - "direct_reorient_timeout", + "DirectReorientTimeout", ] from .commands import ReorientCommand, ReorientCommandCfg, ReorientEpisodeCommand, ReorientEpisodeCommandCfg @@ -46,8 +45,8 @@ from .observations import ( fingertip_quat, fingertip_vel, fingertip_wrench, + compute_goal_quat_error, goal_quat_diff, - openai_policy_observation, reorient_last_action, ) from .rewards import ( @@ -60,8 +59,7 @@ from .rewards import ( track_pos_l2, ) from .terminations import ( - direct_reorient_timeout, - direct_timeout, + DirectReorientTimeout, max_consecutive_success, object_away_from_goal, object_away_from_robot, 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 8e7f77303764..a189328181e1 100644 --- a/source/isaaclab_tasks/isaaclab_tasks/core/reorient/mdp/observations.py +++ b/source/isaaclab_tasks/isaaclab_tasks/core/reorient/mdp/observations.py @@ -11,8 +11,8 @@ from typing import TYPE_CHECKING import torch +import warp as wp -import isaaclab.utils.math as math_utils from isaaclab.managers import ManagerTermBase, ObservationTermCfg, SceneEntityCfg from isaaclab.utils.noise import NoiseModelCfg @@ -24,25 +24,81 @@ from .commands import ReorientCommand -def goal_quat_diff( - env: ManagerBasedRLEnv, asset_cfg: SceneEntityCfg, command_name: str, make_quat_unique: bool +wp.init() + + +@wp.kernel +def _goal_quat_error_kernel( + asset_quat: wp.array(dtype=wp.vec4), + goal_quat: wp.array(dtype=wp.vec4), + make_unique: int, + out: wp.array(dtype=wp.vec4), +): + """Per-environment quaternion error ``asset * conjugate(goal)`` in (x, y, z, w) order.""" + i = wp.tid() + q1 = asset_quat[i] + q2 = goal_quat[i] + # Hamilton product against the conjugate, matching isaaclab.utils.math.quat_mul/quat_conjugate; + # quaternions are stored (x, y, z, w) + w = q1[3] * q2[3] + q1[0] * q2[0] + q1[1] * q2[1] + q1[2] * q2[2] + x = q1[0] * q2[3] - q1[3] * q2[0] - q1[1] * q2[2] + q1[2] * q2[1] + y = q1[1] * q2[3] - q1[3] * q2[1] - q1[2] * q2[0] + q1[0] * q2[2] + z = q1[2] * q2[3] - q1[3] * q2[2] - q1[0] * q2[1] + q1[1] * q2[0] + sign = 1.0 + # make_unique keeps the real part non-negative (isaaclab.utils.math.quat_unique) + if make_unique != 0 and w < 0.0: + sign = -1.0 + out[i] = wp.vec4(sign * x, sign * y, sign * z, sign * w) + + +def _as_wp(tensor: torch.Tensor, dtype) -> wp.array: + """View a contiguous float tensor as a Warp array of *dtype*.""" + return wp.from_torch(tensor.contiguous(), dtype=dtype) + + +def compute_goal_quat_error( + asset_quat: torch.Tensor, goal_quat: torch.Tensor, make_quat_unique: bool, out: torch.Tensor ) -> torch.Tensor: + """Compute the quaternion error between asset and goal orientations. + + Args: + asset_quat: Asset ``(x, y, z, w)`` orientations, shape ``(num_envs, 4)``. + goal_quat: Goal ``(x, y, z, w)`` orientations, shape ``(num_envs, 4)``. + make_quat_unique: Flip the sign so the real part is always non-negative. + out: Caller-owned output buffer, shape ``(num_envs, 4)``, float32. + + Returns: + ``out`` filled with per-environment quaternion errors. + """ + wp.launch( + _goal_quat_error_kernel, + dim=out.shape[0], + inputs=[_as_wp(asset_quat, wp.vec4), _as_wp(goal_quat, wp.vec4), int(make_quat_unique)], + outputs=[wp.from_torch(out, dtype=wp.vec4)], + device=wp.device_from_torch(out.device), + ) + return out + + +class goal_quat_diff(ManagerTermBase): """Goal orientation relative to the asset's root frame. - The quaternion is represented as (x, y, z, w). 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 + 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) - # 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 __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] + command_term: ReorientCommand = env.command_manager.get_term(command_name) + return compute_goal_quat_error( + asset.data.root_quat_w.torch, command_term.command[:, 3:7], make_quat_unique, self._out + ) def fingertip_pos(env: ManagerBasedRLEnv, asset_cfg: SceneEntityCfg) -> torch.Tensor: @@ -129,38 +185,6 @@ def reorient_last_action(env: ManagerBasedRLEnv, action_name: str) -> torch.Tens return torch.where((reset_step == common_step_counter).unsqueeze(-1), reset_action, raw_action) -def openai_policy_observation( - env: ManagerBasedRLEnv, - command_name: str, - action_name: str, - robot_cfg: SceneEntityCfg, - object_cfg: SceneEntityCfg, -) -> torch.Tensor: - """Build the Direct OpenAI actor observation before corruption. - - Args: - env: Environment containing the hand, object, command, and action term. - command_name: Goal command term name. - action_name: Action term whose raw action is observed. - robot_cfg: Hand entity with resolved fingertip body indices. - object_cfg: Object scene entity. - - Returns: - Actor observation in Direct order, shape ``(num_envs, 42)``. - """ - object_asset: RigidObject = env.scene[object_cfg.name] - object_pos = object_asset.data.root_pos_w.torch - env.scene.env_origins - return torch.cat( - ( - fingertip_pos(env, robot_cfg), - object_pos, - goal_quat_diff(env, object_cfg, command_name, make_quat_unique=False), - reorient_last_action(env, action_name), - ), - dim=-1, - ) - - class OpenAIPolicyObservation(ManagerTermBase): """Apply one stateful noise model to the concatenated OpenAI actor observation.""" @@ -168,6 +192,7 @@ 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) # 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. @@ -195,7 +220,17 @@ def __call__( ) -> torch.Tensor: """Return the corrupted 42-dimensional actor observation.""" del noise_model - observation = openai_policy_observation(env, command_name, action_name, robot_cfg, object_cfg) + object_asset: RigidObject = env.scene[object_cfg.name] + object_pos = object_asset.data.root_pos_w.torch - env.scene.env_origins + command_term: ReorientCommand = env.command_manager.get_term(command_name) + compute_goal_quat_error( + object_asset.data.root_quat_w.torch, command_term.command[:, 3:7], False, self._quat_error + ) + # Direct actor-observation order: fingertips, object position, goal quat error, last action + observation = torch.cat( + (fingertip_pos(env, robot_cfg), 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) From 8d9050e7c33c43b43ee662ff89e3f232f6a830e5 Mon Sep 17 00:00:00 2001 From: jichuanh Date: Thu, 9 Jul 2026 09:45:40 +0000 Subject: [PATCH 4/7] Consume the task constants in the Shadow managers The Shadow and OpenAI manager cfgs read the shared scalar values and noise-model cfgs from the constants module and the Direct cfg module scope, the parity-only direct_timeout custom is replaced by the framework time_out term, and the OpenAI policy observation folds into its term class with a Warp quat-error kernel and caller-owned buffer. --- .../isaaclab_tasks/core/reorient/mdp/observations.py | 12 ++++++------ 1 file changed, 6 insertions(+), 6 deletions(-) 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 a189328181e1..23256fd2d6a9 100644 --- a/source/isaaclab_tasks/isaaclab_tasks/core/reorient/mdp/observations.py +++ b/source/isaaclab_tasks/isaaclab_tasks/core/reorient/mdp/observations.py @@ -51,11 +51,6 @@ def _goal_quat_error_kernel( out[i] = wp.vec4(sign * x, sign * y, sign * z, sign * w) -def _as_wp(tensor: torch.Tensor, dtype) -> wp.array: - """View a contiguous float tensor as a Warp array of *dtype*.""" - return wp.from_torch(tensor.contiguous(), dtype=dtype) - - def compute_goal_quat_error( asset_quat: torch.Tensor, goal_quat: torch.Tensor, make_quat_unique: bool, out: torch.Tensor ) -> torch.Tensor: @@ -73,7 +68,12 @@ def compute_goal_quat_error( wp.launch( _goal_quat_error_kernel, dim=out.shape[0], - inputs=[_as_wp(asset_quat, wp.vec4), _as_wp(goal_quat, wp.vec4), int(make_quat_unique)], + # .contiguous(): the goal quaternion arrives as a non-contiguous command slice + inputs=[ + wp.from_torch(asset_quat.contiguous(), dtype=wp.vec4), + wp.from_torch(goal_quat.contiguous(), dtype=wp.vec4), + int(make_quat_unique), + ], outputs=[wp.from_torch(out, dtype=wp.vec4)], device=wp.device_from_torch(out.device), ) From bf005583b654f795db0d44ce3080634c43ddd537 Mon Sep 17 00:00:00 2001 From: jichuanh Date: Thu, 9 Jul 2026 10:07:12 +0000 Subject: [PATCH 5/7] Consume the task constants in the Shadow managers The Shadow and OpenAI manager cfgs read the shared scalar values and noise-model cfgs from the constants module and the Direct cfg module scope, the parity-only direct_timeout custom is replaced by the framework time_out term, and the OpenAI policy observation folds into its term class with a Warp quat-error kernel and caller-owned buffer. --- .../core/reorient/mdp/__init__.pyi | 10 ++- .../core/reorient/mdp/observations.py | 70 ++++++++----------- 2 files changed, 35 insertions(+), 45 deletions(-) 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 e8d7aba1b5f4..2cd4322b640f 100644 --- a/source/isaaclab_tasks/isaaclab_tasks/core/reorient/mdp/__init__.pyi +++ b/source/isaaclab_tasks/isaaclab_tasks/core/reorient/mdp/__init__.pyi @@ -18,12 +18,11 @@ __all__ = [ "reorient_last_action", "OpenAIPolicyObservation", "goal_quat_diff", - "compute_goal_quat_error", "success_bonus", "track_orientation_inv_l2", "track_pos_l2", - "direct_reorient_rotation_distance", - "evaluate_reorient_success", + "rotation_distance_kernel", + "reorient_success_kernel", "direct_reorient_reward", "DirectReorientReward", "max_consecutive_success", @@ -45,15 +44,14 @@ from .observations import ( fingertip_quat, fingertip_vel, fingertip_wrench, - compute_goal_quat_error, goal_quat_diff, reorient_last_action, ) from .rewards import ( DirectReorientReward, direct_reorient_reward, - direct_reorient_rotation_distance, - evaluate_reorient_success, + reorient_success_kernel, + rotation_distance_kernel, success_bonus, track_orientation_inv_l2, track_pos_l2, 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 23256fd2d6a9..fc53135945db 100644 --- a/source/isaaclab_tasks/isaaclab_tasks/core/reorient/mdp/observations.py +++ b/source/isaaclab_tasks/isaaclab_tasks/core/reorient/mdp/observations.py @@ -29,10 +29,10 @@ @wp.kernel def _goal_quat_error_kernel( - asset_quat: wp.array(dtype=wp.vec4), - goal_quat: wp.array(dtype=wp.vec4), + asset_quat: wp.array(dtype=wp.quatf), + goal_quat: wp.array(dtype=wp.quatf), make_unique: int, - out: wp.array(dtype=wp.vec4), + out: wp.array(dtype=wp.quatf), ): """Per-environment quaternion error ``asset * conjugate(goal)`` in (x, y, z, w) order.""" i = wp.tid() @@ -48,36 +48,7 @@ def _goal_quat_error_kernel( # make_unique keeps the real part non-negative (isaaclab.utils.math.quat_unique) if make_unique != 0 and w < 0.0: sign = -1.0 - out[i] = wp.vec4(sign * x, sign * y, sign * z, sign * w) - - -def compute_goal_quat_error( - asset_quat: torch.Tensor, goal_quat: torch.Tensor, make_quat_unique: bool, out: torch.Tensor -) -> torch.Tensor: - """Compute the quaternion error between asset and goal orientations. - - Args: - asset_quat: Asset ``(x, y, z, w)`` orientations, shape ``(num_envs, 4)``. - goal_quat: Goal ``(x, y, z, w)`` orientations, shape ``(num_envs, 4)``. - make_quat_unique: Flip the sign so the real part is always non-negative. - out: Caller-owned output buffer, shape ``(num_envs, 4)``, float32. - - Returns: - ``out`` filled with per-environment quaternion errors. - """ - wp.launch( - _goal_quat_error_kernel, - dim=out.shape[0], - # .contiguous(): the goal quaternion arrives as a non-contiguous command slice - inputs=[ - wp.from_torch(asset_quat.contiguous(), dtype=wp.vec4), - wp.from_torch(goal_quat.contiguous(), dtype=wp.vec4), - int(make_quat_unique), - ], - outputs=[wp.from_torch(out, dtype=wp.vec4)], - device=wp.device_from_torch(out.device), - ) - return out + out[i] = wp.quatf(sign * x, sign * y, sign * z, sign * w) class goal_quat_diff(ManagerTermBase): @@ -89,16 +60,27 @@ class goal_quat_diff(ManagerTermBase): 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] - command_term: ReorientCommand = env.command_manager.get_term(command_name) - return compute_goal_quat_error( - asset.data.root_quat_w.torch, command_term.command[:, 3:7], make_quat_unique, self._out + 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, int(make_quat_unique)], + outputs=[self._out_wp], + device=self._out_wp.device, ) + return self._out def fingertip_pos(env: ManagerBasedRLEnv, asset_cfg: SceneEntityCfg) -> torch.Tensor: @@ -193,6 +175,10 @@ def __init__(self, cfg: ObservationTermCfg, env: ManagerBasedRLEnv): 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) + # 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. @@ -222,9 +208,15 @@ def __call__( del noise_model object_asset: RigidObject = env.scene[object_cfg.name] object_pos = object_asset.data.root_pos_w.torch - env.scene.env_origins - command_term: ReorientCommand = env.command_manager.get_term(command_name) - compute_goal_quat_error( - object_asset.data.root_quat_w.torch, command_term.command[:, 3:7], False, self._quat_error + 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, ) # Direct actor-observation order: fingertips, object position, goal quat error, last action observation = torch.cat( From 9a51738d50bcf5008c36367bb6b90719f9ba47e9 Mon Sep 17 00:00:00 2001 From: jichuanh Date: Sat, 11 Jul 2026 00:34:17 +0000 Subject: [PATCH 6/7] Wire SHADOW_SIM_DT into the Shadow manager cfg The Shadow manager cfg reads the simulation timestep from the shared task constants instead of a literal, keeping the manager and Direct variants single-sourced, and the shared observation terms pick up the Warp kernel conversion. --- .../shadow_hand_manager_env_cfg.py | 3 +- .../core/reorient/mdp/observations.py | 212 +++++++++++++----- 2 files changed, 160 insertions(+), 55 deletions(-) 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 index 06e298925b15..5bdba6849326 100644 --- 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 @@ -69,6 +69,7 @@ SHADOW_RESET_POSITION_NOISE, SHADOW_ROT_EPS, SHADOW_ROT_REWARD_SCALE, + SHADOW_SIM_DT, SHADOW_SUCCESS_COUNT_THRESHOLD, SHADOW_SUCCESS_TOLERANCE, SHADOW_VEL_OBS_SCALE, @@ -249,7 +250,7 @@ class ShadowHandManagerEnvCfg(ManagerBasedRLEnvCfg): scene: ShadowHandManagerSceneCfg = ShadowHandManagerSceneCfg() sim: SimulationCfg = SimulationCfg( - dt=1.0 / 120.0, + dt=SHADOW_SIM_DT, render_interval=2, physics_material=RigidBodyMaterialBaseCfg(static_friction=1.0, dynamic_friction=1.0), physics=PhysicsCfg(), 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 fc53135945db..0eac3483cbfd 100644 --- a/source/isaaclab_tasks/isaaclab_tasks/core/reorient/mdp/observations.py +++ b/source/isaaclab_tasks/isaaclab_tasks/core/reorient/mdp/observations.py @@ -27,6 +27,64 @@ wp.init() +@wp.kernel +def _fingertip_pos_kernel( + body_pos_w: wp.array2d(dtype=wp.vec3f), + env_origins: wp.array(dtype=wp.vec3f), + body_ids: wp.array(dtype=wp.int32), + out: wp.array2d(dtype=wp.float32), +): + i, j = wp.tid() + p = body_pos_w[i, body_ids[j]] - env_origins[i] + out[i, 3 * j + 0] = p[0] + out[i, 3 * j + 1] = p[1] + out[i, 3 * j + 2] = p[2] + + +@wp.kernel +def _fingertip_quat_kernel( + body_quat_w: wp.array2d(dtype=wp.quatf), + body_ids: wp.array(dtype=wp.int32), + out: wp.array2d(dtype=wp.float32), +): + i, j = wp.tid() + q = body_quat_w[i, body_ids[j]] + out[i, 4 * j + 0] = q[0] + out[i, 4 * j + 1] = q[1] + out[i, 4 * j + 2] = q[2] + out[i, 4 * j + 3] = q[3] + + +@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 _fingertip_vel_kernel( + body_vel_w: wp.array2d(dtype=wp.spatial_vectorf), + body_ids: wp.array(dtype=wp.int32), + out: wp.array2d(dtype=wp.float32), +): + i, j = wp.tid() + v = body_vel_w[i, body_ids[j]] + for k in range(6): + out[i, 6 * j + k] = v[k] + + @wp.kernel def _goal_quat_error_kernel( asset_quat: wp.array(dtype=wp.quatf), @@ -83,69 +141,102 @@ def __call__( return self._out -def fingertip_pos(env: ManagerBasedRLEnv, asset_cfg: SceneEntityCfg) -> torch.Tensor: - """Return flattened fingertip positions in the environment frame [m]. +class fingertip_pos(ManagerTermBase): + """Flattened fingertip positions in the environment frame [m].""" - Args: - env: Environment containing the hand. - asset_cfg: Hand entity with resolved fingertip body indices. - - Returns: - Fingertip positions [m], shape ``(num_envs, num_fingertips * 3)``. - """ - asset = env.scene[asset_cfg.name] - positions = asset.data.body_pos_w.torch[:, asset_cfg.body_ids] - positions = positions - env.scene.env_origins[:, None, :] - return positions.flatten(start_dim=1) - - -def fingertip_quat(env: ManagerBasedRLEnv, asset_cfg: SceneEntityCfg) -> torch.Tensor: - """Return flattened fingertip ``(x, y, z, w)`` orientations. - - Args: - env: Environment containing the hand. - asset_cfg: Hand entity with resolved fingertip body indices. + 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 - Returns: - Unit quaternions, shape ``(num_envs, num_fingertips * 4)``. - """ - asset = env.scene[asset_cfg.name] - return asset.data.body_quat_w.torch[:, asset_cfg.body_ids].flatten(start_dim=1) +class fingertip_quat(ManagerTermBase): + """Flattened fingertip ``(x, y, z, w)`` orientations.""" -def fingertip_vel(env: ManagerBasedRLEnv, asset_cfg: SceneEntityCfg) -> torch.Tensor: - """Return flattened fingertip spatial velocities in the world frame. + 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 - Args: - env: Environment containing the hand. - asset_cfg: Hand entity with resolved fingertip body indices. - Returns: - Spatial velocities [m/s, rad/s], shape ``(num_envs, num_fingertips * 6)``. - """ - asset = env.scene[asset_cfg.name] - return asset.data.body_vel_w.torch[:, asset_cfg.body_ids].flatten(start_dim=1) +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 -def fingertip_wrench(env: ManagerBasedRLEnv, sensor_cfg: SceneEntityCfg) -> torch.Tensor: - """Return fingertip reaction wrenches with Direct-compatible zero fallback. - Args: - env: Environment containing the joint-wrench sensor. - sensor_cfg: Joint-wrench sensor entity with resolved fingertip body indices. +class fingertip_wrench(ManagerTermBase): + """Fingertip reaction wrenches [N, N·m] with Direct-compatible zero fallback.""" - Returns: - Fingertip reaction wrenches [N, N·m], 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: - body_count = len(sensor_cfg.body_ids) - return torch.zeros(env.num_envs, body_count * 6, device=env.device) - force = force_data.torch[:, sensor_cfg.body_ids] - torque = torque_data.torch[:, sensor_cfg.body_ids] - return torch.cat((force, torque), dim=-1).flatten(start_dim=1) + 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: @@ -175,6 +266,11 @@ def __init__(self, cfg: ObservationTermCfg, env: ManagerBasedRLEnv): 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 @@ -218,9 +314,17 @@ def __call__( 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( - (fingertip_pos(env, robot_cfg), object_pos, self._quat_error, reorient_last_action(env, action_name)), + (self._fingertip_buf, object_pos, self._quat_error, reorient_last_action(env, action_name)), dim=-1, ) if self._shape_probe_pending: From de8bf48079b8ef88772d8cf15a647d1e78358303 Mon Sep 17 00:00:00 2001 From: Jichuan Hu Date: Tue, 14 Jul 2026 02:07:24 -0700 Subject: [PATCH 7/7] Sync the Shadow manager configuration to per-paradigm parameters The Shadow and OpenAI manager configurations declare their scalar task parameters at the term declarations (matching the Direct values), and the observation terms pick up the family kernel module imports. --- .../shadow_hand_manager_env_cfg.py | 159 ++++++++---------- .../core/reorient/mdp/observations.py | 110 ++++++------ 2 files changed, 117 insertions(+), 152 deletions(-) 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 index 5bdba6849326..0157904c95c1 100644 --- 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 @@ -17,7 +17,6 @@ from isaaclab.scene import InteractiveSceneCfg from isaaclab.sensors import JointWrenchSensorCfg from isaaclab.sim import SimulationCfg -from isaaclab.sim.spawners.materials import RigidBodyMaterialBaseCfg from isaaclab.utils.configclass import configclass import isaaclab_tasks.core.reorient.mdp as mdp @@ -27,52 +26,17 @@ OPENAI_ACTION_NOISE_CFG, OPENAI_OBSERVATION_NOISE_CFG, ROBOT_CFG, + SHADOW_SIM_CFG, NewtonEventCfg, ObjectCfg, PhysicsCfg, PhysxEventCfg, ) -from isaaclab_tasks.core.reorient.reorient_task_constants import ( +from isaaclab_tasks.core.reorient.reorient_task_base import ( GOAL_MARKER_POSITION, - OPENAI_ACT_MOVING_AVERAGE, - OPENAI_ACTION_PENALTY_SCALE, - OPENAI_AV_FACTOR, - OPENAI_DECIMATION, - OPENAI_DIST_REWARD_SCALE, - OPENAI_EPISODE_LENGTH_S, - OPENAI_FALL_DIST, - OPENAI_FALL_PENALTY, - OPENAI_FORCE_TORQUE_OBS_SCALE, - OPENAI_MAX_CONSECUTIVE_SUCCESS, - OPENAI_REACH_GOAL_BONUS, - OPENAI_RESET_DOF_POS_NOISE, - OPENAI_RESET_DOF_VEL_NOISE, - OPENAI_RESET_POSITION_NOISE, - OPENAI_ROT_EPS, - OPENAI_ROT_REWARD_SCALE, - OPENAI_SIM_DT, - OPENAI_SUCCESS_COUNT_THRESHOLD, - OPENAI_SUCCESS_TOLERANCE, - SHADOW_ACT_MOVING_AVERAGE, - SHADOW_ACTION_PENALTY_SCALE, + IN_HAND_POS_OFFSET, SHADOW_ACTUATED_JOINT_NAMES, - SHADOW_AV_FACTOR, - SHADOW_DECIMATION, - SHADOW_DIST_REWARD_SCALE, - SHADOW_EPISODE_LENGTH_S, - SHADOW_FALL_DIST, - SHADOW_FALL_PENALTY, SHADOW_FINGERTIP_BODY_NAMES, - SHADOW_REACH_GOAL_BONUS, - SHADOW_RESET_DOF_POS_NOISE, - SHADOW_RESET_DOF_VEL_NOISE, - SHADOW_RESET_POSITION_NOISE, - SHADOW_ROT_EPS, - SHADOW_ROT_REWARD_SCALE, - SHADOW_SIM_DT, - SHADOW_SUCCESS_COUNT_THRESHOLD, - SHADOW_SUCCESS_TOLERANCE, - SHADOW_VEL_OBS_SCALE, ) from isaaclab_tasks.utils import PresetCfg @@ -94,9 +58,17 @@ class _ShadowHandManagerSceneCfg(InteractiveSceneCfg): 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) + 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 + num_envs=8192, + env_spacing=0.75, + replicate_physics=True, + clone_in_fabric=False, ) ovphysx = physx newton_kamino = newton_mjwarp @@ -109,9 +81,9 @@ class CommandsCfg: object_pose = mdp.ReorientEpisodeCommandCfg( asset_name="object", - init_pos_offset=(0.0, 0.0, -0.04), + init_pos_offset=IN_HAND_POS_OFFSET, update_goal_on_success=True, - orientation_success_threshold=SHADOW_SUCCESS_TOLERANCE, + orientation_success_threshold=0.1, make_quat_unique=False, fixed_marker_pos=GOAL_MARKER_POSITION, goal_pose_visualizer_cfg=GOAL_OBJECT_CFG, @@ -126,7 +98,7 @@ class ActionsCfg: joint_pos = mdp.EMAJointPositionToLimitsActionCfg( asset_name="robot", joint_names=SHADOW_ACTUATED_JOINT_NAMES, - alpha=SHADOW_ACT_MOVING_AVERAGE, + alpha=1.0, rescale_to_limits=True, ) @@ -141,7 +113,7 @@ class FullStateWithoutActionCfg(ObsGroup): ) joint_vel = ObsTerm( func=mdp.joint_vel, - scale=SHADOW_VEL_OBS_SCALE, + 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")}) @@ -152,7 +124,7 @@ class FullStateWithoutActionCfg(ObsGroup): 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=SHADOW_VEL_OBS_SCALE, + scale=0.2, params={"asset_cfg": SceneEntityCfg("object")}, ) goal_pose = ObsTerm(func=mdp.generated_commands, params={"command_name": "object_pose"}) @@ -197,9 +169,9 @@ class EventCfg: func=mdp.reset_reorient_state, mode="reset", params={ - "position_noise": SHADOW_RESET_POSITION_NOISE, - "joint_position_noise": SHADOW_RESET_DOF_POS_NOISE, - "joint_velocity_noise": SHADOW_RESET_DOF_VEL_NOISE, + "position_noise": 0.01, + "joint_position_noise": 0.2, + "joint_velocity_noise": 0.0, "action_name": "joint_pos", }, ) @@ -214,16 +186,16 @@ class RewardsCfg: weight=1.0, params={ "command_name": "object_pose", - "distance_scale": SHADOW_DIST_REWARD_SCALE, - "rotation_scale": SHADOW_ROT_REWARD_SCALE, - "rotation_epsilon": SHADOW_ROT_EPS, - "action_penalty_scale": SHADOW_ACTION_PENALTY_SCALE, - "success_tolerance": SHADOW_SUCCESS_TOLERANCE, - "success_bonus": SHADOW_REACH_GOAL_BONUS, - "fall_distance": SHADOW_FALL_DIST, - "fall_penalty": SHADOW_FALL_PENALTY, - "averaging_factor": SHADOW_AV_FACTOR, - "success_count_threshold": SHADOW_SUCCESS_COUNT_THRESHOLD, + "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"), }, ) @@ -236,7 +208,7 @@ class TerminationsCfg: object_out_of_reach = DoneTerm( func=mdp.object_reorientation_out_of_reach, params={ - "threshold": SHADOW_FALL_DIST, + "threshold": 0.24, "command_name": "object_pose", "object_cfg": SceneEntityCfg("object"), }, @@ -249,12 +221,7 @@ class ShadowHandManagerEnvCfg(ManagerBasedRLEnvCfg): """Manager-based state Shadow Hand task with Direct-compatible semantics.""" scene: ShadowHandManagerSceneCfg = ShadowHandManagerSceneCfg() - sim: SimulationCfg = SimulationCfg( - dt=SHADOW_SIM_DT, - render_interval=2, - physics_material=RigidBodyMaterialBaseCfg(static_friction=1.0, dynamic_friction=1.0), - physics=PhysicsCfg(), - ) + sim: SimulationCfg = SHADOW_SIM_CFG observations: ObservationsCfg = ObservationsCfg() actions: ActionsCfg = ActionsCfg() commands: CommandsCfg = CommandsCfg() @@ -263,8 +230,8 @@ class ShadowHandManagerEnvCfg(ManagerBasedRLEnvCfg): events: EventCfg = EventCfg() def __post_init__(self): - self.decimation = SHADOW_DECIMATION - self.episode_length_s = SHADOW_EPISODE_LENGTH_S + self.decimation = 2 + self.episode_length_s = 10.0 self.sim.render_interval = self.decimation self.viewer.eye = (2.0, 2.0, 2.0) @@ -275,9 +242,9 @@ class OpenAICommandsCfg: object_pose = mdp.ReorientEpisodeCommandCfg( asset_name="object", - init_pos_offset=(0.0, 0.0, -0.04), + init_pos_offset=IN_HAND_POS_OFFSET, update_goal_on_success=True, - orientation_success_threshold=OPENAI_SUCCESS_TOLERANCE, + orientation_success_threshold=0.4, make_quat_unique=False, fixed_marker_pos=GOAL_MARKER_POSITION, goal_pose_visualizer_cfg=GOAL_OBJECT_CFG, @@ -292,7 +259,7 @@ class OpenAIActionsCfg: joint_pos = mdp.NoisyEMAJointPositionToLimitsActionCfg( asset_name="robot", joint_names=SHADOW_ACTUATED_JOINT_NAMES, - alpha=OPENAI_ACT_MOVING_AVERAGE, + alpha=0.3, rescale_to_limits=True, noise_model=OPENAI_ACTION_NOISE_CFG, ) @@ -323,7 +290,7 @@ def __post_init__(self): class CriticCfg(FullStateWithoutActionCfg): fingertip_wrench = ObsTerm( func=mdp.fingertip_wrench, - scale=OPENAI_FORCE_TORQUE_OBS_SCALE, + scale=10.0, params={ "sensor_cfg": SceneEntityCfg( "joint_wrench", body_names=SHADOW_FINGERTIP_BODY_NAMES, preserve_order=False @@ -348,10 +315,16 @@ class ShadowHandOpenAIManagerSceneCfg(PresetCfg): """Backend-specific OpenAI scene alternatives.""" physx = _ShadowHandOpenAIManagerSceneCfg( - num_envs=8192, env_spacing=0.75, replicate_physics=True, clone_in_fabric=True + 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 + num_envs=8192, + env_spacing=0.75, + replicate_physics=True, + clone_in_fabric=False, ) ovphysx = physx newton_kamino = newton_mjwarp @@ -359,9 +332,9 @@ class ShadowHandOpenAIManagerSceneCfg(PresetCfg): _OPENAI_RESET_PARAMS = { - "position_noise": OPENAI_RESET_POSITION_NOISE, - "joint_position_noise": OPENAI_RESET_DOF_POS_NOISE, - "joint_velocity_noise": OPENAI_RESET_DOF_VEL_NOISE, + "position_noise": 0.01, + "joint_position_noise": 0.2, + "joint_velocity_noise": 0.0, "action_name": "joint_pos", } @@ -400,16 +373,16 @@ class OpenAIRewardsCfg: weight=1.0, params={ "command_name": "object_pose", - "distance_scale": OPENAI_DIST_REWARD_SCALE, - "rotation_scale": OPENAI_ROT_REWARD_SCALE, - "rotation_epsilon": OPENAI_ROT_EPS, - "action_penalty_scale": OPENAI_ACTION_PENALTY_SCALE, - "success_tolerance": OPENAI_SUCCESS_TOLERANCE, - "success_bonus": OPENAI_REACH_GOAL_BONUS, - "fall_distance": OPENAI_FALL_DIST, - "fall_penalty": OPENAI_FALL_PENALTY, - "averaging_factor": OPENAI_AV_FACTOR, - "success_count_threshold": OPENAI_SUCCESS_COUNT_THRESHOLD, + "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"), }, @@ -423,7 +396,7 @@ class OpenAITerminationsCfg: object_out_of_reach = DoneTerm( func=mdp.object_reorientation_out_of_reach, params={ - "threshold": OPENAI_FALL_DIST, + "threshold": 0.24, "command_name": "object_pose", "object_cfg": SceneEntityCfg("object"), }, @@ -434,8 +407,8 @@ class OpenAITerminationsCfg: params={ "command_name": "object_pose", "reward_name": "reorient", - "success_tolerance": OPENAI_SUCCESS_TOLERANCE, - "max_successes": OPENAI_MAX_CONSECUTIVE_SUCCESS, + "success_tolerance": 0.4, + "max_successes": 50, "object_cfg": SceneEntityCfg("object"), }, ) @@ -455,8 +428,8 @@ class ShadowHandOpenAIManagerEnvCfg(ShadowHandManagerEnvCfg): def __post_init__(self): super().__post_init__() - self.decimation = OPENAI_DECIMATION - self.episode_length_s = OPENAI_EPISODE_LENGTH_S - self.sim.dt = OPENAI_SIM_DT + 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/observations.py b/source/isaaclab_tasks/isaaclab_tasks/core/reorient/mdp/observations.py index 0eac3483cbfd..87c7ba1aee2d 100644 --- a/source/isaaclab_tasks/isaaclab_tasks/core/reorient/mdp/observations.py +++ b/source/isaaclab_tasks/isaaclab_tasks/core/reorient/mdp/observations.py @@ -16,43 +16,26 @@ 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 JointWrenchSensor - - from .commands import ReorientCommand - - -wp.init() + from isaaclab.sensors import Camera, JointWrenchSensor + from isaaclab_tasks.core.reorient.config.shadow_hand.feature_extractor import FeatureExtractorCfg -@wp.kernel -def _fingertip_pos_kernel( - body_pos_w: wp.array2d(dtype=wp.vec3f), - env_origins: wp.array(dtype=wp.vec3f), - body_ids: wp.array(dtype=wp.int32), - out: wp.array2d(dtype=wp.float32), -): - i, j = wp.tid() - p = body_pos_w[i, body_ids[j]] - env_origins[i] - out[i, 3 * j + 0] = p[0] - out[i, 3 * j + 1] = p[1] - out[i, 3 * j + 2] = p[2] - - -@wp.kernel -def _fingertip_quat_kernel( - body_quat_w: wp.array2d(dtype=wp.quatf), - body_ids: wp.array(dtype=wp.int32), - out: wp.array2d(dtype=wp.float32), -): - i, j = wp.tid() - q = body_quat_w[i, body_ids[j]] - out[i, 4 * j + 0] = q[0] - out[i, 4 * j + 1] = q[1] - out[i, 4 * j + 2] = q[2] - out[i, 4 * j + 3] = q[3] + from .commands import ReorientCommand @wp.kernel @@ -73,18 +56,6 @@ def _fingertip_wrench_kernel( out[i, 6 * j + 5] = t[2] -@wp.kernel -def _fingertip_vel_kernel( - body_vel_w: wp.array2d(dtype=wp.spatial_vectorf), - body_ids: wp.array(dtype=wp.int32), - out: wp.array2d(dtype=wp.float32), -): - i, j = wp.tid() - v = body_vel_w[i, body_ids[j]] - for k in range(6): - out[i, 6 * j + k] = v[k] - - @wp.kernel def _goal_quat_error_kernel( asset_quat: wp.array(dtype=wp.quatf), @@ -94,19 +65,14 @@ def _goal_quat_error_kernel( ): """Per-environment quaternion error ``asset * conjugate(goal)`` in (x, y, z, w) order.""" i = wp.tid() - q1 = asset_quat[i] - q2 = goal_quat[i] - # Hamilton product against the conjugate, matching isaaclab.utils.math.quat_mul/quat_conjugate; - # quaternions are stored (x, y, z, w) - w = q1[3] * q2[3] + q1[0] * q2[0] + q1[1] * q2[1] + q1[2] * q2[2] - x = q1[0] * q2[3] - q1[3] * q2[0] - q1[1] * q2[2] + q1[2] * q2[1] - y = q1[1] * q2[3] - q1[3] * q2[1] - q1[2] * q2[0] + q1[0] * q2[2] - z = q1[2] * q2[3] - q1[3] * q2[2] - q1[0] * q2[1] + q1[1] * q2[0] + # 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 w < 0.0: + if make_unique != 0 and qe[3] < 0.0: sign = -1.0 - out[i] = wp.quatf(sign * x, sign * y, sign * z, sign * w) + out[i] = wp.quatf(sign * qe[0], sign * qe[1], sign * qe[2], sign * qe[3]) class goal_quat_diff(ManagerTermBase): @@ -134,7 +100,7 @@ def __call__( wp.launch( _goal_quat_error_kernel, dim=self.num_envs, - inputs=[asset.data.root_quat_w.warp, self._goal_quat_wp, int(make_quat_unique)], + inputs=[asset.data.root_quat_w.warp, self._goal_quat_wp, make_quat_unique], outputs=[self._out_wp], device=self._out_wp.device, ) @@ -156,7 +122,7 @@ def __call__(self, env: ManagerBasedRLEnv, asset_cfg: SceneEntityCfg) -> torch.T """Return the flattened per-fingertip block, shape ``(num_envs, num_fingertips * 3)``.""" asset = env.scene[asset_cfg.name] wp.launch( - _fingertip_pos_kernel, + 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], @@ -179,7 +145,7 @@ def __call__(self, env: ManagerBasedRLEnv, asset_cfg: SceneEntityCfg) -> torch.T """Return the flattened per-fingertip block, shape ``(num_envs, num_fingertips * 4)``.""" asset = env.scene[asset_cfg.name] wp.launch( - _fingertip_quat_kernel, + 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], @@ -202,7 +168,7 @@ def __call__(self, env: ManagerBasedRLEnv, asset_cfg: SceneEntityCfg) -> torch.T """Return the flattened per-fingertip block, shape ``(num_envs, num_fingertips * 6)``.""" asset = env.scene[asset_cfg.name] wp.launch( - _fingertip_vel_kernel, + 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], @@ -316,7 +282,7 @@ def __call__( ) robot = env.scene[robot_cfg.name] wp.launch( - _fingertip_pos_kernel, + 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], @@ -330,3 +296,29 @@ def __call__( 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 + +