diff --git a/docs/source/overview/environments.rst b/docs/source/overview/environments.rst index 8e612678792f..48206b784dd1 100644 --- a/docs/source/overview/environments.rst +++ b/docs/source/overview/environments.rst @@ -219,8 +219,8 @@ for the lift-cube environment: | | |allegro-direct-link| | | ``ovphysx`` (direct only) | +-------------------------+------------------------------+-----------------------------------------------------------------------------+------------------------------+ | |cube-shadow| | |cube-shadow-link| | In-hand reorientation of a cube using Shadow hand | **physics=** ``physx``, | - | | | | ``newton_mjwarp`` | - | | |cube-shadow-ff-link| | | | + | | | | ``newton_mjwarp``, | + | | |cube-shadow-ff-link| | | ``ovphysx`` | | | | | | | | |cube-shadow-lstm-link| | | | +-------------------------+------------------------------+-----------------------------------------------------------------------------+------------------------------+ @@ -355,10 +355,10 @@ for the lift-cube environment: .. |kuka-allegro-reorient-link| replace:: :isaaclab-source:`Isaac-Reorient-KukaAllegro ` .. |ka-lift-cam-link| replace:: :isaaclab-source:`Isaac-Lift-KukaAllegro-Camera ` .. |ka-reorient-cam-link| replace:: :isaaclab-source:`Isaac-Reorient-KukaAllegro-Camera ` -.. |cube-shadow-link| replace:: :isaaclab-source:`Isaac-Reorient-Cube-Shadow-Direct ` -.. |cube-shadow-ff-link| replace:: :isaaclab-source:`Isaac-Reorient-Cube-Shadow-OpenAI-FF-Direct ` -.. |cube-shadow-lstm-link| replace:: :isaaclab-source:`Isaac-Reorient-Cube-Shadow-OpenAI-LSTM-Direct ` -.. |cube-shadow-vis-link| replace:: :isaaclab-source:`Isaac-Reorient-Cube-Shadow-Camera-Direct ` +.. |cube-shadow-link| replace:: :isaaclab-source:`Isaac-Reorient-Cube-Shadow-Direct ` +.. |cube-shadow-ff-link| replace:: :isaaclab-source:`Isaac-Reorient-Cube-Shadow-OpenAI-FF-Direct ` +.. |cube-shadow-lstm-link| replace:: :isaaclab-source:`Isaac-Reorient-Cube-Shadow-OpenAI-LSTM-Direct ` +.. |cube-shadow-vis-link| replace:: :isaaclab-source:`Isaac-Reorient-Cube-Shadow-Camera-Direct ` .. |agibot_place_mug-link| replace:: :isaaclab-source:`IsaacContrib-Place-Mug-Agibot-Left-Arm-RmpFlow ` .. |agibot_place_toy-link| replace:: :isaaclab-source:`IsaacContrib-Place-Toy2Box-Agibot-Right-Arm-RmpFlow ` .. |reach_openarm_bi-link| replace:: :isaaclab-source:`IsaacContrib-Reach-OpenArmBi ` @@ -1054,7 +1054,7 @@ inferencing, including reading from an already trained checkpoint and disabling - Isaac-Reorient-Cube-Allegro-Play - Manager Based - **rl_games** (PPO), **rsl_rl** (PPO), **skrl** (PPO) - - + - **physics=** ``newton_mjwarp``, ``ovphysx``, ``physx`` * - Isaac-Reorient-Cube-Allegro-Direct - - Direct @@ -1065,28 +1065,50 @@ inferencing, including reading from an already trained checkpoint and disabling - Direct - **rl_games** (PPO), **rsl_rl** (PPO), **skrl** (PPO) - + * - Isaac-Reorient-Cube-Shadow + - + - Manager Based + - **rl_games** (PPO), **rsl_rl** (PPO), **skrl** (PPO) + - **physics=** ``newton_kamino``, ``newton_mjwarp``, ``ovphysx``, ``physx`` + * - Isaac-Reorient-Cube-Shadow-Camera + - Isaac-Reorient-Cube-Shadow-Camera-Play + - Manager Based + - **rl_games** (PPO), **rsl_rl** (PPO) + - | **physics=** ``newton_kamino``, ``newton_mjwarp``, ``ovphysx``, ``physx`` + | **renderer=** ``isaacsim_rtx``, ``newton_renderer``, ``ovrtx``, ``rtx`` + | **presets=** ``albedo``, ``depth``, ``full``, ``rgb``, ``rgb_depth``, ``semantic_segmentation``, ``simple_shading_constant_diffuse``, ``simple_shading_diffuse_mdl``, ``simple_shading_full_mdl`` * - Isaac-Reorient-Cube-Shadow-Camera-Direct - Isaac-Reorient-Cube-Shadow-Camera-Direct-Play - Direct - **rl_games** (PPO), **rsl_rl** (PPO) - - | **physics=** ``newton_kamino``, ``newton_mjwarp``, ``physx`` + - | **physics=** ``newton_kamino``, ``newton_mjwarp``, ``ovphysx``, ``physx`` | **renderer=** ``isaacsim_rtx``, ``newton_renderer``, ``ovrtx``, ``rtx`` - | **presets=** ``albedo``, ``depth``, ``full``, ``rgb``, ``semantic_segmentation``, ``simple_shading_constant_diffuse``, ``simple_shading_diffuse_mdl``, ``simple_shading_full_mdl`` + | **presets=** ``albedo``, ``depth``, ``full``, ``rgb``, ``rgb_depth``, ``semantic_segmentation``, ``simple_shading_constant_diffuse``, ``simple_shading_diffuse_mdl``, ``simple_shading_full_mdl`` * - Isaac-Reorient-Cube-Shadow-Direct - - Direct - **rl_games** (PPO), **rsl_rl** (PPO), **skrl** (PPO) - - **physics=** ``newton_kamino``, ``newton_mjwarp``, ``physx`` + - **physics=** ``newton_kamino``, ``newton_mjwarp``, ``ovphysx``, ``physx`` + * - Isaac-Reorient-Cube-Shadow-OpenAI-FF + - + - Manager Based + - **rl_games** (PPO), **rsl_rl** (PPO), **skrl** (PPO) + - **physics=** ``newton_kamino``, ``newton_mjwarp``, ``ovphysx``, ``physx`` * - Isaac-Reorient-Cube-Shadow-OpenAI-FF-Direct - - Direct - **rl_games** (PPO), **rsl_rl** (PPO), **skrl** (PPO) - - **physics=** ``newton_kamino``, ``newton_mjwarp``, ``physx`` + - **physics=** ``newton_kamino``, ``newton_mjwarp``, ``ovphysx``, ``physx`` + * - Isaac-Reorient-Cube-Shadow-OpenAI-LSTM + - + - Manager Based + - **rl_games** (PPO), **rsl_rl** (PPO) + - **physics=** ``newton_kamino``, ``newton_mjwarp``, ``ovphysx``, ``physx`` * - Isaac-Reorient-Cube-Shadow-OpenAI-LSTM-Direct - - Direct - - **rl_games** (PPO) - - **physics=** ``newton_kamino``, ``newton_mjwarp``, ``physx`` + - **rl_games** (PPO), **rsl_rl** (PPO) + - **physics=** ``newton_kamino``, ``newton_mjwarp``, ``ovphysx``, ``physx`` * - Isaac-Reorient-Franka - Isaac-Reorient-Franka-Play - Manager Based @@ -1106,11 +1128,16 @@ inferencing, including reading from an already trained checkpoint and disabling - | **physics=** ``newton_mjwarp``, ``physx`` | **renderer=** ``isaacsim_rtx``, ``newton_renderer``, ``ovrtx``, ``rtx`` | **presets=** ``albedo128``, ``albedo256``, ``albedo64``, ``cube``, ``depth128``, ``depth256``, ``depth64``, ``duo_camera``, ``raycaster_depth128``, ``raycaster_depth256``, ``raycaster_depth64``, ``rgb128``, ``rgb256``, ``rgb64``, ``semantic_segmentation128``, ``semantic_segmentation256``, ``semantic_segmentation64``, ``shapes``, ``simple_shading_constant_diffuse128``, ``simple_shading_constant_diffuse256``, ``simple_shading_constant_diffuse64``, ``simple_shading_diffuse_mdl128``, ``simple_shading_diffuse_mdl256``, ``simple_shading_diffuse_mdl64``, ``simple_shading_full_mdl128``, ``simple_shading_full_mdl256``, ``simple_shading_full_mdl64``, ``single_camera`` + * - Isaac-Shadow-Handover + - + - Manager Based + - **rsl_rl** (PPO) + - **physics=** ``newton_mjwarp``, ``ovphysx``, ``physx`` * - Isaac-Shadow-Handover-Direct - - Direct - - **rl_games** (PPO), **skrl** (PPO, IPPO, MAPPO) - - **physics=** ``newton_mjwarp``, ``physx`` + - **rl_games** (PPO), **rsl_rl** (PPO), **skrl** (PPO, IPPO, MAPPO) + - **physics=** ``newton_mjwarp``, ``ovphysx``, ``physx`` * - Isaac-Velocity-Flat-AnymalB-Warp-v0 - Isaac-Velocity-Flat-AnymalB-Warp-Play-v0 - Manager Based diff --git a/source/isaaclab/changelog.d/dexterous-env-convergence.rst b/source/isaaclab/changelog.d/dexterous-env-convergence.rst new file mode 100644 index 000000000000..e7023a98d361 --- /dev/null +++ b/source/isaaclab/changelog.d/dexterous-env-convergence.rst @@ -0,0 +1,7 @@ +Fixed +^^^^^ + +* Fixed :meth:`~isaaclab.envs.DirectRLEnv.reset` to store the observation buffer + like :meth:`~isaaclab.envs.DirectRLEnv.step` already does, and exposed the + latest observations on the multi-agent-to-single-agent adapter through the + same public buffer. diff --git a/source/isaaclab/isaaclab/envs/direct_rl_env.py b/source/isaaclab/isaaclab/envs/direct_rl_env.py index fd5623ed5171..afa02618093e 100644 --- a/source/isaaclab/isaaclab/envs/direct_rl_env.py +++ b/source/isaaclab/isaaclab/envs/direct_rl_env.py @@ -382,7 +382,9 @@ def reset(self, seed: int | None = None, options: dict[str, Any] | None = None) self.sim.render() # return observations - return self._get_observations(), self.extras + # store the buffer like step() does, so consumers can read the latest observations + self.obs_buf = self._get_observations() + return self.obs_buf, self.extras def step(self, action: torch.Tensor) -> VecEnvStepReturn: """Execute one time-step of the environment's dynamics. diff --git a/source/isaaclab/isaaclab/envs/utils/marl.py b/source/isaaclab/isaaclab/envs/utils/marl.py index 55c25acd0b8c..b48d68898d50 100644 --- a/source/isaaclab/isaaclab/envs/utils/marl.py +++ b/source/isaaclab/isaaclab/envs/utils/marl.py @@ -81,22 +81,34 @@ def __init__(self, env: DirectMARLEnv) -> None: ) self.action_space = gym.vector.utils.batch_space(self.single_action_space, self.num_envs) - def reset(self, seed: int | None = None, options: dict[str, Any] | None = None) -> tuple[VecEnvObs, dict]: - obs, extras = self.env.reset(seed, options) + @property + def episode_length_buf(self) -> torch.Tensor: + """Episode lengths from the wrapped multi-agent environment.""" + return self.env.episode_length_buf - # use environment state as observation - if self._state_as_observation: - obs = {"policy": self.env.state()} - # concatenate agents' observations + @episode_length_buf.setter + def episode_length_buf(self, value: torch.Tensor) -> None: + self.env.episode_length_buf = value + + @property + def obs_buf(self) -> VecEnvObs: + """Latest observations from the wrapped multi-agent environment.""" + return self._convert_observations(self.env.obs_dict) + + def _convert_observations(self, obs: dict[AgentID, ObsType]) -> VecEnvObs: + """Convert multi-agent observations to the single-agent policy observation.""" # FIXME: This implementation assumes the spaces are fundamental ones. Fix it to support composite spaces - else: - obs = { - "policy": torch.cat( - [obs[agent].reshape(self.num_envs, -1) for agent in self.env.possible_agents], dim=-1 - ) - } + if self._state_as_observation: + return {"policy": self.env.state()} + return { + "policy": torch.cat( + [obs[agent].reshape(self.num_envs, -1) for agent in self.env.possible_agents], dim=-1 + ) + } - return obs, extras + def reset(self, seed: int | None = None, options: dict[str, Any] | None = None) -> tuple[VecEnvObs, dict]: + obs, extras = self.env.reset(seed, options) + return self._convert_observations(obs), extras def step(self, action: torch.Tensor) -> VecEnvStepReturn: # split single-agent actions to build the multi-agent ones @@ -111,17 +123,7 @@ def step(self, action: torch.Tensor) -> VecEnvStepReturn: # step the environment obs, rewards, terminated, time_outs, extras = self.env.step(_actions) - # use environment state as observation - if self._state_as_observation: - obs = {"policy": self.env.state()} - # concatenate agents' observations - # FIXME: This implementation assumes the spaces are fundamental ones. Fix it to support composite spaces - else: - obs = { - "policy": torch.cat( - [obs[agent].reshape(self.num_envs, -1) for agent in self.env.possible_agents], dim=-1 - ) - } + obs = self._convert_observations(obs) # process environment outputs to return single-agent data rewards = sum(rewards.values()) diff --git a/source/isaaclab/test/cli/test_install_command_parsing.py b/source/isaaclab/test/cli/test_install_command_parsing.py index 8a6abe523130..34ec8e915a63 100644 --- a/source/isaaclab/test/cli/test_install_command_parsing.py +++ b/source/isaaclab/test/cli/test_install_command_parsing.py @@ -106,6 +106,11 @@ def test_core_submodules_starts_with_isaaclab(self): "isaaclab must be first so dependents resolve against the local copy" ) + def test_core_submodules_install_contrib_before_tasks(self): + assert CORE_ISAACLAB_SUBMODULES.index("isaaclab_contrib") < CORE_ISAACLAB_SUBMODULES.index("isaaclab_tasks"), ( + "isaaclab_contrib must be installed before isaaclab_tasks so its declared local dependency resolves" + ) + def test_core_submodules_contains_expected_packages(self): expected = { "isaaclab", diff --git a/source/isaaclab/test/envs/test_marl_utils.py b/source/isaaclab/test/envs/test_marl_utils.py new file mode 100644 index 000000000000..25704d0005ad --- /dev/null +++ b/source/isaaclab/test/envs/test_marl_utils.py @@ -0,0 +1,82 @@ +# Copyright (c) 2022-2026, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +from types import SimpleNamespace + +import gymnasium as gym +import torch + +from isaaclab.envs.utils.marl import multi_agent_to_single_agent + + +class _FakeMultiAgentEnv: + possible_agents = ["agent_0", "agent_1"] + observation_spaces = { + "agent_0": gym.spaces.Box(low=-1.0, high=1.0, shape=(2,)), + "agent_1": gym.spaces.Box(low=-1.0, high=1.0, shape=(1,)), + } + action_spaces = { + "agent_0": gym.spaces.Box(low=-1.0, high=1.0, shape=(1,)), + "agent_1": gym.spaces.Box(low=-1.0, high=1.0, shape=(1,)), + } + render_mode = None + + def __init__(self): + self.unwrapped = self + self.cfg = SimpleNamespace(state_space=2) + self.state_space = gym.spaces.Box(low=-1.0, high=1.0, shape=(2,)) + self.sim = object() + self.scene = SimpleNamespace(num_envs=2) + self.episode_length_buf = torch.tensor([1, 2]) + self.obs_dict = { + "agent_0": torch.tensor([[1.0, 2.0], [3.0, 4.0]]), + "agent_1": torch.tensor([[5.0], [6.0]]), + } + + def reset(self, seed=None, options=None): + return self.obs_dict, {} + + def state(self): + return torch.tensor([[7.0, 8.0], [9.0, 10.0]]) + + def close(self): + pass + + +def test_multi_agent_to_single_agent_reset_concatenates_agents(): + """The adapter reset should concatenate the agents' observations.""" + env = multi_agent_to_single_agent(_FakeMultiAgentEnv()) + + observations, _ = env.reset() + + torch.testing.assert_close(observations["policy"], torch.tensor([[1.0, 2.0, 5.0], [3.0, 4.0, 6.0]])) + + +def test_multi_agent_to_single_agent_reset_can_use_state(): + """The adapter reset should support the state-as-observation mode.""" + env = multi_agent_to_single_agent(_FakeMultiAgentEnv(), state_as_observation=True) + + observations, _ = env.reset() + + torch.testing.assert_close(observations["policy"], torch.tensor([[7.0, 8.0], [9.0, 10.0]])) + + +def test_multi_agent_to_single_agent_forwards_episode_lengths(): + """RSL-RL episode randomization should update the wrapped environment buffer.""" + source_env = _FakeMultiAgentEnv() + env = multi_agent_to_single_agent(source_env) + episode_lengths = torch.tensor([3, 4]) + + env.episode_length_buf = episode_lengths + + assert env.episode_length_buf is episode_lengths + assert source_env.episode_length_buf is episode_lengths + + +def test_multi_agent_to_single_agent_exposes_latest_observations(): + """The public observation buffer should reflect the wrapped environment's buffer.""" + env = multi_agent_to_single_agent(_FakeMultiAgentEnv()) + + torch.testing.assert_close(env.obs_buf["policy"], torch.tensor([[1.0, 2.0, 5.0], [3.0, 4.0, 6.0]])) diff --git a/source/isaaclab_assets/changelog.d/task-cleanup-dex-part03.minor.rst b/source/isaaclab_assets/changelog.d/task-cleanup-dex-part03.minor.rst new file mode 100644 index 000000000000..6c8250b50c79 --- /dev/null +++ b/source/isaaclab_assets/changelog.d/task-cleanup-dex-part03.minor.rst @@ -0,0 +1,9 @@ +Added +^^^^^ + +* Added the dexterous-hand actuated-joint and fingertip body-name lists + (:obj:`~isaaclab_assets.robots.shadow_hand.SHADOW_ACTUATED_JOINT_NAMES`, + :obj:`~isaaclab_assets.robots.shadow_hand.SHADOW_FINGERTIP_BODY_NAMES`, + :obj:`~isaaclab_assets.robots.allegro.ALLEGRO_ACTUATED_JOINT_NAMES`, + :obj:`~isaaclab_assets.robots.allegro.ALLEGRO_FINGERTIP_BODY_NAMES`) to the + robot asset modules so tasks can reference them from a single source. diff --git a/source/isaaclab_assets/isaaclab_assets/robots/allegro.py b/source/isaaclab_assets/isaaclab_assets/robots/allegro.py index 10d96e277770..b2669f2dd283 100644 --- a/source/isaaclab_assets/isaaclab_assets/robots/allegro.py +++ b/source/isaaclab_assets/isaaclab_assets/robots/allegro.py @@ -66,3 +66,32 @@ soft_joint_pos_limit_factor=1.0, ) """Configuration of Allegro Hand robot.""" + + +ALLEGRO_FINGERTIP_BODY_NAMES: list[str] = [ + "index_link_3", + "middle_link_3", + "ring_link_3", + "thumb_link_3", +] +"""Allegro Hand fingertip body names.""" + +ALLEGRO_ACTUATED_JOINT_NAMES: list[str] = [ + "index_joint_0", + "middle_joint_0", + "ring_joint_0", + "thumb_joint_0", + "index_joint_1", + "index_joint_2", + "index_joint_3", + "middle_joint_1", + "middle_joint_2", + "middle_joint_3", + "ring_joint_1", + "ring_joint_2", + "ring_joint_3", + "thumb_joint_1", + "thumb_joint_2", + "thumb_joint_3", +] +"""Allegro Hand actuated joint names, in the Direct task's actuation order.""" diff --git a/source/isaaclab_assets/isaaclab_assets/robots/shadow_hand.py b/source/isaaclab_assets/isaaclab_assets/robots/shadow_hand.py index 696b8f9b5a4d..b348cc32a57e 100644 --- a/source/isaaclab_assets/isaaclab_assets/robots/shadow_hand.py +++ b/source/isaaclab_assets/isaaclab_assets/robots/shadow_hand.py @@ -82,3 +82,37 @@ soft_joint_pos_limit_factor=1.0, ) """Configuration of Shadow Hand robot.""" + + +SHADOW_FINGERTIP_BODY_NAMES: list[str] = [ + "robot0_ffdistal", + "robot0_mfdistal", + "robot0_rfdistal", + "robot0_lfdistal", + "robot0_thdistal", +] +"""Shadow Hand fingertip body names (identical on every backend asset).""" + +SHADOW_ACTUATED_JOINT_NAMES: list[str] = [ + "robot0_WRJ1", + "robot0_WRJ0", + "robot0_FFJ3", + "robot0_FFJ2", + "robot0_FFJ1", + "robot0_MFJ3", + "robot0_MFJ2", + "robot0_MFJ1", + "robot0_RFJ3", + "robot0_RFJ2", + "robot0_RFJ1", + "robot0_LFJ4", + "robot0_LFJ3", + "robot0_LFJ2", + "robot0_LFJ1", + "robot0_THJ4", + "robot0_THJ3", + "robot0_THJ2", + "robot0_THJ1", + "robot0_THJ0", +] +"""Shadow Hand actuated joint names, in the Direct task's actuation order.""" diff --git a/source/isaaclab_experimental/changelog.d/dexterous-env-convergence.rst b/source/isaaclab_experimental/changelog.d/dexterous-env-convergence.rst new file mode 100644 index 000000000000..37781a5ec013 --- /dev/null +++ b/source/isaaclab_experimental/changelog.d/dexterous-env-convergence.rst @@ -0,0 +1,7 @@ +Fixed +^^^^^ + +* Fixed :meth:`~isaaclab_experimental.envs.DirectRLEnvWarp.step` and + :meth:`~isaaclab_experimental.envs.DirectRLEnvWarp.reset` to store the returned + observation dictionary in ``obs_buf``, which + :meth:`~isaaclab_rl.rsl_rl.RslRlVecEnvWrapper.get_observations` now reads. diff --git a/source/isaaclab_experimental/isaaclab_experimental/envs/direct_rl_env_warp.py b/source/isaaclab_experimental/isaaclab_experimental/envs/direct_rl_env_warp.py index d69b13a8feec..0a4f52bfbf99 100644 --- a/source/isaaclab_experimental/isaaclab_experimental/envs/direct_rl_env_warp.py +++ b/source/isaaclab_experimental/isaaclab_experimental/envs/direct_rl_env_warp.py @@ -377,7 +377,9 @@ def reset(self, seed: int | None = None, options: dict[str, Any] | None = None) # return observations self._get_observations() - return {"policy": self.torch_obs_buf.clone()}, self.extras + # store the returned buffer so RslRlVecEnvWrapper.get_observations() can read env.obs_buf + self.obs_buf = {"policy": self.torch_obs_buf.clone()} + return self.obs_buf, self.extras @Timer(name="env_step", msg="Step took:", enable=DEBUG_TIMER_STEP or DEBUG_TIMERS) def step(self, action: torch.Tensor) -> VecEnvStepReturn: @@ -460,8 +462,10 @@ def step(self, action: torch.Tensor) -> VecEnvStepReturn: self._post_step_visualize() # return observations, rewards, resets and extras + # store the returned buffer so RslRlVecEnvWrapper.get_observations() can read env.obs_buf + self.obs_buf = {"policy": self.torch_obs_buf.clone()} return ( - {"policy": self.torch_obs_buf.clone()}, + self.obs_buf, self.torch_reward_buf, self.torch_reset_terminated, self.torch_reset_time_outs, diff --git a/source/isaaclab_rl/changelog.d/dexterous-env-convergence.minor.rst b/source/isaaclab_rl/changelog.d/dexterous-env-convergence.minor.rst new file mode 100644 index 000000000000..d9cb7373b999 --- /dev/null +++ b/source/isaaclab_rl/changelog.d/dexterous-env-convergence.minor.rst @@ -0,0 +1,10 @@ +Changed +^^^^^^^ + +* Changed :meth:`~isaaclab_rl.rsl_rl.RslRlVecEnvWrapper.get_observations` to read + the environment-owned observation buffer instead of calling private environment + methods. The returned observations now match the latest reset/step returns, + including observation-noise corruption that the private path skipped, and + multi-agent environments converted with + :func:`~isaaclab.envs.utils.multi_agent_to_single_agent` train with RSL-RL + without environment-side accommodations. diff --git a/source/isaaclab_rl/isaaclab_rl/rsl_rl/vecenv_wrapper.py b/source/isaaclab_rl/isaaclab_rl/rsl_rl/vecenv_wrapper.py index 5a3fdcd47164..88c8a7ddfcee 100644 --- a/source/isaaclab_rl/isaaclab_rl/rsl_rl/vecenv_wrapper.py +++ b/source/isaaclab_rl/isaaclab_rl/rsl_rl/vecenv_wrapper.py @@ -171,11 +171,7 @@ def reset(self) -> tuple[TensorDict, dict]: # noqa: D102 def get_observations(self) -> TensorDict: """Returns the current observations of the environment.""" - if hasattr(self.unwrapped, "observation_manager"): - obs_dict = self.unwrapped.observation_manager.compute() - else: - obs_dict = self.unwrapped._get_observations() - return TensorDict(obs_dict, batch_size=[self.num_envs]) + return TensorDict(self.unwrapped.obs_buf, batch_size=[self.num_envs]) def step(self, actions: torch.Tensor) -> tuple[TensorDict, torch.Tensor, torch.Tensor, dict]: # clip actions diff --git a/source/isaaclab_tasks/changelog.d/task-cleanup-dex-part03.minor.rst b/source/isaaclab_tasks/changelog.d/task-cleanup-dex-part03.minor.rst new file mode 100644 index 000000000000..9a453bb3de5d --- /dev/null +++ b/source/isaaclab_tasks/changelog.d/task-cleanup-dex-part03.minor.rst @@ -0,0 +1,13 @@ +Added +^^^^^ + +* Added behavioral-success metrics and threshold-independent episode-error + diagnostics to the dexterous reorientation environments. + +Fixed +^^^^^ + +* Fixed dexterous hand resets that could initialize joints below their lower + position limits. Reset joint positions now sample uniformly across the full + joint range; previously the distribution was biased toward the lower half of + the range. diff --git a/source/isaaclab_tasks/changelog.d/task-cleanup-dex-part04.minor.rst b/source/isaaclab_tasks/changelog.d/task-cleanup-dex-part04.minor.rst new file mode 100644 index 000000000000..12c88a87b468 --- /dev/null +++ b/source/isaaclab_tasks/changelog.d/task-cleanup-dex-part04.minor.rst @@ -0,0 +1,25 @@ +Added +^^^^^ + +* Added an RSL-RL training configuration and behavioral-success metrics to the + Shadow handover Direct task. +* Added renderer presets and configuration validation to the Shadow camera + Direct task, including an RGB-depth preset for training with the Newton Warp + renderer. +* Added OVPhysX physics presets to the handover and camera Direct + environments. + +Deprecated +^^^^^^^^^^ + +* Deprecated ``shadow_hand_camera_env.compute_keypoints`` in favor of + :func:`~isaaclab_tasks.core.reorient.mdp.observations.compute_cube_keypoints`. +* Deprecated the ``Isaac-Reorient-Cube-Shadow-Camera-Benchmark-Direct`` + registration in favor of the regular camera task with the + ``env.feature_extractor.enabled=False`` override. + +Fixed +^^^^^ + +* Fixed handover construction on Newton, broken by renamed distal joints in + the current Shadow Newton asset. diff --git a/source/isaaclab_tasks/changelog.d/task-cleanup-dex-part05.rst b/source/isaaclab_tasks/changelog.d/task-cleanup-dex-part05.rst new file mode 100644 index 000000000000..fb46bc4fb50c --- /dev/null +++ b/source/isaaclab_tasks/changelog.d/task-cleanup-dex-part05.rst @@ -0,0 +1,11 @@ +Added +^^^^^ + +* Added success-rate reporting to the environment training benchmark + utilities, with unit tests for the benchmark discovery helpers. + +Fixed +^^^^^ + +* Fixed training benchmark discovery to exclude inference-only camera + benchmark registrations. diff --git a/source/isaaclab_tasks/changelog.d/task-cleanup-dex-part08.major.rst b/source/isaaclab_tasks/changelog.d/task-cleanup-dex-part08.major.rst new file mode 100644 index 000000000000..7d5dcbfaaafd --- /dev/null +++ b/source/isaaclab_tasks/changelog.d/task-cleanup-dex-part08.major.rst @@ -0,0 +1,31 @@ +Added +^^^^^ + +* Added manager-based counterparts for the Allegro and Shadow cube + reorientation tasks (state and OpenAI FF/LSTM observation variants), sharing + the Direct tasks' scalar parameters and boolean success metrics through + common MDP terms. +* Added opt-in domain randomization to the manager-based Allegro environment + (``enable_domain_randomization``, disabled by default; enabling requires + retraining). +* Added a Newton physics preset to the manager-based Allegro environment. + +Removed +^^^^^^^ + +* Removed the legacy manager-based reorientation configuration + ``ReorientObjectEnvCfg`` and the manager terms only it consumed + (``success_bonus``, ``track_pos_l2``, ``track_orientation_inv_l2``, + ``max_consecutive_success``, ``object_away_from_goal``, and + ``object_away_from_robot``). Use the Direct-compatible manager + configurations and terms instead (e.g. + :class:`~isaaclab_tasks.core.reorient.mdp.ReorientReward` and + :class:`~isaaclab_tasks.core.reorient.mdp.ReorientTimeout`). + +Changed +^^^^^^^ + +* **Breaking:** Changed the manager-based Allegro reorientation environment to + match the Direct observation, action, reward, reset, termination, success, + asset, and benchmark contracts. Existing manager checkpoints are + incompatible and must be retrained. diff --git a/source/isaaclab_tasks/changelog.d/task-cleanup-dex-part11.minor.rst b/source/isaaclab_tasks/changelog.d/task-cleanup-dex-part11.minor.rst new file mode 100644 index 000000000000..560f84a934db --- /dev/null +++ b/source/isaaclab_tasks/changelog.d/task-cleanup-dex-part11.minor.rst @@ -0,0 +1,8 @@ +Added +^^^^^ + +* Added manager-based counterparts for the Shadow handover and Shadow camera + reorientation tasks, completing the manager coverage of the dexterous task + families. +* Added Direct-vs-manager scalar value-parity checks for the dexterous task + families. diff --git a/source/isaaclab_tasks/isaaclab_tasks/core/handover/__init__.py b/source/isaaclab_tasks/isaaclab_tasks/core/handover/__init__.py index 19a42486b912..444aae7a0aa3 100644 --- a/source/isaaclab_tasks/isaaclab_tasks/core/handover/__init__.py +++ b/source/isaaclab_tasks/isaaclab_tasks/core/handover/__init__.py @@ -15,12 +15,23 @@ # Register Gym environments. ## +gym.register( + id="Isaac-Shadow-Handover", + entry_point="isaaclab.envs:ManagerBasedRLEnv", + disable_env_checker=True, + kwargs={ + "env_cfg_entry_point": f"{__name__}.handover_manager_env_cfg:HandoverManagerEnvCfg", + "rsl_rl_cfg_entry_point": f"{agents.__name__}.rsl_rl_ppo_cfg:HandoverPPORunnerCfg", + }, +) + gym.register( id="Isaac-Shadow-Handover-Direct", entry_point=f"{__name__}.handover_env:HandoverEnv", disable_env_checker=True, kwargs={ "env_cfg_entry_point": f"{__name__}.handover_env_cfg:HandoverEnvCfg", + "rsl_rl_cfg_entry_point": f"{agents.__name__}.rsl_rl_ppo_cfg:HandoverPPORunnerCfg", "rl_games_cfg_entry_point": f"{agents.__name__}:rl_games_ppo_cfg.yaml", "skrl_cfg_entry_point": f"{agents.__name__}:skrl_ppo_cfg.yaml", "skrl_ippo_cfg_entry_point": f"{agents.__name__}:skrl_ippo_cfg.yaml", diff --git a/source/isaaclab_tasks/isaaclab_tasks/core/handover/agents/rsl_rl_ppo_cfg.py b/source/isaaclab_tasks/isaaclab_tasks/core/handover/agents/rsl_rl_ppo_cfg.py new file mode 100644 index 000000000000..75d1b4c0f98f --- /dev/null +++ b/source/isaaclab_tasks/isaaclab_tasks/core/handover/agents/rsl_rl_ppo_cfg.py @@ -0,0 +1,44 @@ +# Copyright (c) 2022-2026, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +from isaaclab.utils.configclass import configclass + +from isaaclab_rl.rsl_rl import RslRlMLPModelCfg, RslRlOnPolicyRunnerCfg, RslRlPpoAlgorithmCfg + + +@configclass +class HandoverPPORunnerCfg(RslRlOnPolicyRunnerCfg): + """RSL-RL PPO configuration for the single-agent view of Shadow Hand handover.""" + + num_steps_per_env = 16 + max_iterations = 5000 + save_interval = 250 + experiment_name = "handover" + obs_groups = {"actor": ["policy"], "critic": ["policy"]} + actor = RslRlMLPModelCfg( + hidden_dims=[512, 512, 256, 128], + activation="elu", + obs_normalization=True, + distribution_cfg=RslRlMLPModelCfg.GaussianDistributionCfg(init_std=1.0), + ) + critic = RslRlMLPModelCfg( + hidden_dims=[512, 512, 256, 128], + activation="elu", + obs_normalization=True, + ) + algorithm = RslRlPpoAlgorithmCfg( + value_loss_coef=1.0, + use_clipped_value_loss=True, + clip_param=0.2, + entropy_coef=0.005, + num_learning_epochs=5, + num_mini_batches=4, + learning_rate=5.0e-4, + schedule="adaptive", + gamma=0.99, + lam=0.95, + desired_kl=0.016, + max_grad_norm=1.0, + ) diff --git a/source/isaaclab_tasks/isaaclab_tasks/core/handover/handover_common.py b/source/isaaclab_tasks/isaaclab_tasks/core/handover/handover_common.py new file mode 100644 index 000000000000..8f9fdbe82b3b --- /dev/null +++ b/source/isaaclab_tasks/isaaclab_tasks/core/handover/handover_common.py @@ -0,0 +1,95 @@ +# 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 + +"""Shared task-defining parameters for the handover family. + +The Direct environment and the manager counterpart inherit the same +shared structural definitions so the two variants cannot drift apart. Every +scalar task parameter reads as a flat ``cfg.`` field on the env +configuration; non-scalar items (joint/body-name lists, the goal offset, and +marker templates) stay module-level constants. +""" + +import isaaclab.sim as sim_utils +from isaaclab.markers import VisualizationMarkersCfg + +from isaaclab_tasks.utils.hydra import preset + +from isaaclab_assets.robots.shadow_hand import ( + SHADOW_ACTUATED_JOINT_NAMES as ACTUATED_JOINT_NAMES, +) +from isaaclab_assets.robots.shadow_hand import ( + SHADOW_FINGERTIP_BODY_NAMES as FINGERTIP_BODY_NAMES, +) + +__all__ = [ + "ACTUATED_JOINT_NAMES", + "ACTUATED_JOINT_NAMES_NEWTON", + "ACTUATED_JOINT_NAMES_PRESET", + "FINGERTIP_BODY_NAMES", + "GOAL_MARKER_CFG", + "GOAL_POSITION_OFFSET", + "OBJECT_RADIUS", +] + + +ACTUATED_JOINT_NAMES_NEWTON: list[str] = [ + "robot0_WRJ1", + "robot0_WRJ0", + "robot0_FFJ4", + "robot0_FFJ3", + "robot0_FFJ2", + "robot0_MFJ4", + "robot0_MFJ3", + "robot0_MFJ2", + "robot0_RFJ4", + "robot0_RFJ3", + "robot0_RFJ2", + "robot0_LFJ5", + "robot0_LFJ4", + "robot0_LFJ3", + "robot0_LFJ2", + "robot0_THJ4", + "robot0_THJ3", + "robot0_THJ2", + "robot0_THJ1", + "robot0_THJ0", +] +"""Actuated joint names on the production Newton Shadow Hand asset (+1 finger renumbering).""" + + +ACTUATED_JOINT_NAMES_PRESET = preset( + physx=ACTUATED_JOINT_NAMES, + newton_mjwarp=ACTUATED_JOINT_NAMES_NEWTON, + ovphysx=ACTUATED_JOINT_NAMES, + default=ACTUATED_JOINT_NAMES, +) +"""Per-backend actuated joint names, resolved by the physics preset key.""" + + +OBJECT_RADIUS: float = 0.0335 +"""Hand-over object sphere radius [m], also used for the goal marker.""" + +GOAL_POSITION_OFFSET: tuple[float, float, float] = (0.0, -0.25, 0.0) +"""Goal-position offset from the object's default position [m]. + +The Direct environment and the manager command derive the same fixed goal +point from the object's default root position plus this offset. +""" + +GOAL_MARKER_CFG = VisualizationMarkersCfg( + prim_path="/Visuals/goal_marker", + markers={ + "goal": sim_utils.SphereCfg( + radius=OBJECT_RADIUS, + visual_material=sim_utils.PreviewSurfaceCfg(diffuse_color=(0.4, 0.3, 1.0)), + ), + }, +) +"""Goal-marker template shared by the Direct environment and the manager command term. + +Consumers relying on a different prim path use ``replace`` +on this template; configclass deep-copies defaults, so sharing the instance is safe. +""" diff --git a/source/isaaclab_tasks/isaaclab_tasks/core/handover/handover_env.py b/source/isaaclab_tasks/isaaclab_tasks/core/handover/handover_env.py index ae68d13d98ab..328ada86c247 100644 --- a/source/isaaclab_tasks/isaaclab_tasks/core/handover/handover_env.py +++ b/source/isaaclab_tasks/isaaclab_tasks/core/handover/handover_env.py @@ -8,7 +8,6 @@ from collections.abc import Sequence -import numpy as np import torch import isaaclab.sim as sim_utils @@ -17,17 +16,13 @@ from isaaclab.envs import DirectMARLEnv from isaaclab.markers import VisualizationMarkers from isaaclab.sim.spawners.from_files import GroundPlaneCfg, spawn_ground_plane -from isaaclab.utils.math import ( - quat_conjugate, - quat_from_angle_axis, - quat_mul, - sample_uniform, - saturate, - scale_transform, - unscale_transform, -) +from isaaclab.utils.math import quat_conjugate, quat_mul, sample_uniform, saturate, scale_transform, unscale_transform +from isaaclab_tasks.core.handover.handover_common import GOAL_POSITION_OFFSET from isaaclab_tasks.core.handover.handover_env_cfg import HandoverEnvCfg +from isaaclab_tasks.core.handover.mdp.rewards import evaluate_handover_success, handover_reward +from isaaclab_tasks.core.reorient.mdp.events import randomize_rotation, sample_joint_positions_within_limits +from isaaclab_tasks.core.reorient.mdp.rewards import EpisodeErrorRecorder class HandoverEnv(DirectMARLEnv): @@ -53,16 +48,18 @@ def __init__(self, cfg: HandoverEnvCfg, render_mode: str | None = None, **kwargs ) # list of actuated joints - self.actuated_dof_indices = list() - for joint_name in cfg.actuated_joint_names: - self.actuated_dof_indices.append(self.right_hand.joint_names.index(joint_name)) - self.actuated_dof_indices.sort() + self.actuated_dof_indices, _ = self.right_hand.find_joints(cfg.actuated_joint_names) + if len(self.actuated_dof_indices) != len(cfg.actuated_joint_names): + raise ValueError( + f"Expected {len(cfg.actuated_joint_names)} actuated joints, found {len(self.actuated_dof_indices)}." + ) # finger bodies - self.finger_bodies = list() - for body_name in self.cfg.fingertip_body_names: - self.finger_bodies.append(self.right_hand.body_names.index(body_name)) - self.finger_bodies.sort() + self.finger_bodies, _ = self.right_hand.find_bodies(self.cfg.fingertip_body_names) + if len(self.finger_bodies) != len(self.cfg.fingertip_body_names): + raise ValueError( + f"Expected {len(self.cfg.fingertip_body_names)} fingertip bodies, found {len(self.finger_bodies)}." + ) self.num_fingertips = len(self.finger_bodies) # joint limits @@ -72,14 +69,18 @@ def __init__(self, cfg: HandoverEnvCfg, render_mode: str | None = None, **kwargs # default goal positions self.goal_rot = torch.zeros((self.num_envs, 4), dtype=torch.float, device=self.device) - self.goal_rot[:, 0] = 1.0 + self.goal_rot[:, 3] = 1.0 # identity quaternion in (x, y, z, w) layout self.goal_pos = torch.zeros((self.num_envs, 3), dtype=torch.float, device=self.device) - self.goal_pos[:, :] = torch.tensor([0.0, -0.64, 0.54], device=self.device) + # goal = object default position + shared offset (mirrors HandoverCommand.__init__) + self.goal_pos[:, :] = self.object.data.default_root_pose.torch[:, :3].to(self.device) + torch.tensor( + GOAL_POSITION_OFFSET, dtype=torch.float, device=self.device + ) # initialize goal marker self.goal_markers = VisualizationMarkers(self.cfg.goal_object_cfg) # Sticky per-env flag: True once the object reached the goal within threshold. self._episode_succeeded = torch.zeros(self.num_envs, dtype=torch.bool, device=self.device) + self._goal_distance = EpisodeErrorRecorder(self.num_envs, self.device) # unit tensors for sampling goal/object rotations about the x and y axes self.x_unit_tensor = torch.tensor([1, 0, 0], dtype=torch.float, device=self.device).repeat((self.num_envs, 1)) @@ -192,17 +193,21 @@ def _get_states(self) -> torch.Tensor: def _get_rewards(self) -> dict[str, torch.Tensor]: # compute reward - goal_dist = torch.linalg.norm(self.object_pos - self.goal_pos, ord=2, dim=-1) - rew_dist = 2 * torch.exp(-self.cfg.dist_reward_scale * goal_dist) + succeeded, goal_dist = evaluate_handover_success( + self.object_pos, self.goal_pos, self.cfg.success_distance_threshold + ) + self._goal_distance.update(goal_dist) + rew_dist = handover_reward(goal_dist, self.cfg.dist_reward_scale) - # log reward components + # log reward components as tensors, not .item(): a per-step host sync stalls the GPU if "log" not in self.extras: self.extras["log"] = dict() + goal_dist_mean = goal_dist.mean() self.extras["log"]["dist_reward"] = rew_dist.mean() - self.extras["log"]["dist_goal"] = goal_dist.mean() - self.extras["log"]["Metrics/goal_distance"] = goal_dist.mean().item() + self.extras["log"]["dist_goal"] = goal_dist_mean + self.extras["log"]["Metrics/goal_distance"] = goal_dist_mean # Sticky per-env success: True once the object reached the goal within threshold. - self._episode_succeeded |= goal_dist < self.cfg.success_distance_threshold + self._episode_succeeded |= succeeded return {"right_hand": rew_dist, "left_hand": rew_dist} @@ -222,9 +227,10 @@ def _reset_idx(self, env_ids: Sequence[int] | torch.Tensor | None): if env_ids is None: env_ids = self.right_hand._ALL_INDICES # Flush per-episode success (sticky binary: object ever reached the goal within threshold). - self.extras.setdefault("log", {})["Metrics/success_rate"] = ( - self._episode_succeeded[env_ids].float().mean().item() - ) + # 0-dim device tensor: avoids a host sync here; consumers read it at logging cadence + self.extras.setdefault("log", {})["Metrics/success_rate"] = self._episode_succeeded[env_ids].float().mean() + for statistic, value in self._goal_distance.reset(env_ids).items(): + self.extras["log"][f"Diagnostics/episode_min_goal_distance_{statistic}"] = value self._episode_succeeded[env_ids] = False # reset articulation and rigid body attributes super()._reset_idx(env_ids) @@ -251,12 +257,9 @@ def _reset_idx(self, env_ids: Sequence[int] | torch.Tensor | None): self.object.write_root_velocity_to_sim_index(root_velocity=object_default_vel, env_ids=env_ids) # reset right hand - delta_max = self.hand_dof_upper_limits[env_ids] - self.right_hand.data.default_joint_pos.torch[env_ids] - delta_min = self.hand_dof_lower_limits[env_ids] - self.right_hand.data.default_joint_pos.torch[env_ids] - - dof_pos_noise = sample_uniform(-1.0, 1.0, (len(env_ids), self.num_hand_dofs), device=self.device) - rand_delta = delta_min + (delta_max - delta_min) * 0.5 * dof_pos_noise - dof_pos = self.right_hand.data.default_joint_pos.torch[env_ids] + self.cfg.reset_dof_pos_noise * rand_delta + default_dof_pos = self.right_hand.data.default_joint_pos.torch[env_ids] + dof_limits = self.right_hand.data.joint_limits.torch[env_ids] + dof_pos = sample_joint_positions_within_limits(default_dof_pos, dof_limits, self.cfg.reset_dof_pos_noise) dof_vel_noise = sample_uniform(-1.0, 1.0, (len(env_ids), self.num_hand_dofs), device=self.device) dof_vel = self.right_hand.data.default_joint_vel.torch[env_ids] + self.cfg.reset_dof_vel_noise * dof_vel_noise @@ -269,12 +272,9 @@ def _reset_idx(self, env_ids: Sequence[int] | torch.Tensor | None): self.right_hand.write_joint_velocity_to_sim_index(velocity=dof_vel, env_ids=env_ids) # reset left hand - delta_max = self.hand_dof_upper_limits[env_ids] - self.left_hand.data.default_joint_pos.torch[env_ids] - delta_min = self.hand_dof_lower_limits[env_ids] - self.left_hand.data.default_joint_pos.torch[env_ids] - - dof_pos_noise = sample_uniform(-1.0, 1.0, (len(env_ids), self.num_hand_dofs), device=self.device) - rand_delta = delta_min + (delta_max - delta_min) * 0.5 * dof_pos_noise - dof_pos = self.left_hand.data.default_joint_pos.torch[env_ids] + self.cfg.reset_dof_pos_noise * rand_delta + default_dof_pos = self.left_hand.data.default_joint_pos.torch[env_ids] + dof_limits = self.left_hand.data.joint_limits.torch[env_ids] + dof_pos = sample_joint_positions_within_limits(default_dof_pos, dof_limits, self.cfg.reset_dof_pos_noise) dof_vel_noise = sample_uniform(-1.0, 1.0, (len(env_ids), self.num_hand_dofs), device=self.device) dof_vel = self.left_hand.data.default_joint_vel.torch[env_ids] + self.cfg.reset_dof_vel_noise * dof_vel_noise @@ -328,10 +328,3 @@ def _compute_intermediate_values(self): self.object_rot = self.object.data.root_quat_w.torch self.object_linvel = self.object.data.root_lin_vel_w.torch self.object_angvel = self.object.data.root_ang_vel_w.torch - - -@torch.jit.script -def randomize_rotation(rand0, rand1, x_unit_tensor, y_unit_tensor): - return quat_mul( - quat_from_angle_axis(rand0 * np.pi, x_unit_tensor), quat_from_angle_axis(rand1 * np.pi, y_unit_tensor) - ) diff --git a/source/isaaclab_tasks/isaaclab_tasks/core/handover/handover_env_cfg.py b/source/isaaclab_tasks/isaaclab_tasks/core/handover/handover_env_cfg.py index a3d7126ce0bd..76661afc736b 100644 --- a/source/isaaclab_tasks/isaaclab_tasks/core/handover/handover_env_cfg.py +++ b/source/isaaclab_tasks/isaaclab_tasks/core/handover/handover_env_cfg.py @@ -3,11 +3,14 @@ # # SPDX-License-Identifier: BSD-3-Clause +import torch from isaaclab_newton.physics import MJWarpSolverCfg, NewtonCfg +from isaaclab_ovphysx.physics import OvPhysxCfg from isaaclab_physx.physics import PhysxCfg import isaaclab.envs.mdp as mdp import isaaclab.sim as sim_utils +import isaaclab.utils.math as math_utils from isaaclab.actuators import ImplicitActuatorCfg from isaaclab.assets import ArticulationCfg, RigidObjectCfg from isaaclab.envs import DirectMARLEnvCfg @@ -16,10 +19,16 @@ from isaaclab.markers import VisualizationMarkersCfg from isaaclab.scene import InteractiveSceneCfg from isaaclab.sim import SimulationCfg -from isaaclab.sim.spawners.materials.physics_materials_cfg import RigidBodyMaterialCfg +from isaaclab.sim.spawners.materials import RigidBodyMaterialBaseCfg from isaaclab.utils.configclass import configclass -from isaaclab_tasks.core.reorient.config.shadow_hand.shadow_hand_env_cfg import ShadowHandRobotCfg +from isaaclab_tasks.core.handover.handover_common import ( + ACTUATED_JOINT_NAMES_PRESET, + FINGERTIP_BODY_NAMES, + GOAL_MARKER_CFG, + OBJECT_RADIUS, +) +from isaaclab_tasks.core.reorient.config.shadow_hand.shadow_hand_common import ShadowHandRobotCfg from isaaclab_tasks.utils import PresetCfg, preset from isaaclab_assets.robots.shadow_hand import SHADOW_HAND_CFG @@ -127,7 +136,7 @@ class EventCfg: # Reuse the single-agent Shadow Hand Newton port (USD path, ``rot`` reapplication # workaround, effort limits, joint regex). The multi-agent variant only diverges # in actuator gains (stiffness/damping bumped for the catch task) and adds a -# ``distal_passive`` override for the J0 USD-baked values. +# ``distal_passive`` override for the J1 USD-baked values. _SHADOW_HAND_NEWTON_CFG = ShadowHandRobotCfg().newton_mjwarp @@ -136,9 +145,9 @@ def _shadow_hand_cfg( init_pos: tuple[float, float, float], init_rot: tuple[float, float, float, float], ) -> PresetCfg: - """Per-hand Shadow Hand preset (PhysX and Newton MJWarp variants). + """Per-hand Shadow Hand preset (PhysX, Newton MJWarp, and OVPhysX variants). - Both variants are placed at *prim_path* with the same init pose; per-hand + All variants are placed at *prim_path* with the same init pose; per-hand differences (right vs left) come from the caller's *prim_path* / *init_pos* / *init_rot* — the gain tuning is identical on both hands. @@ -154,23 +163,64 @@ def _shadow_hand_cfg( authority. ``20.0`` / ``2.0`` is the smallest tested setting at which MAPPO learns the catch (mean reward at iter 200 / 2048 envs goes from ~27 at PhysX-mirrored gains to ~777). - * ``distal_passive`` on the four ``robot0_(FF|MF|RF|LF)J0`` joints with - ``stiffness=10.0`` / ``damping=0.1``. The Newton USD bakes - ``stiffness=286 / damping=57`` on these joints from the MJCF→USD - translation, which fights the ``MjcTendon`` coupling and bounces the - ball. ``stiffness=10`` (~1/3 of PhysX's ``limit_stiffness=30``) keeps - the joints near-passive while the tendon constraint dominates. + * ``distal_passive`` on the four ``robot0_(FF|MF|RF|LF)J1`` distal joints + (named ``J0`` before the current asset release) with ``stiffness=10.0`` / + ``damping=0.1``. The Newton USD bakes ``stiffness=286 / damping=57`` on + these joints from the MJCF→USD translation, which fights the + ``MjcTendon`` coupling and bounces the ball. ``stiffness=10`` (~1/3 of + PhysX's ``limit_stiffness=30``) keeps the joints near-passive while the + tendon constraint dominates. """ physx_cfg = SHADOW_HAND_CFG.replace(prim_path=prim_path).replace( init_state=ArticulationCfg.InitialStateCfg(pos=init_pos, rot=init_rot, joint_pos={".*": 0.0}) ) + # Newton's importer bakes the asset's native root orientation into the + # root joint (see the note on _SHADOW_HAND_NEWTON_CFG.init_state), so the + # task rotation must compose with that base rotation rather than replace + # it — replacing left both palms heading 90 degrees off and the object + # never rested in the right hand. + # Composed in float64 via the shared (x, y, z, w) quaternion product, + # matching the previously used wp.quatd math bit-for-bit. + newton_rot = tuple( + math_utils.quat_mul( + torch.tensor(init_rot, dtype=torch.float64), + torch.tensor(_SHADOW_HAND_NEWTON_CFG.init_state.rot, dtype=torch.float64), + ).tolist() + ) newton_cfg = _SHADOW_HAND_NEWTON_CFG.replace( prim_path=prim_path, - init_state=_SHADOW_HAND_NEWTON_CFG.init_state.replace(pos=init_pos, rot=init_rot), + init_state=_SHADOW_HAND_NEWTON_CFG.init_state.replace(pos=init_pos, rot=newton_rot), actuators={ - "fingers": _SHADOW_HAND_NEWTON_CFG.actuators["fingers"].replace(stiffness=20.0, damping=2.0), + # The inherited "fingers" expression predates the renamed Newton + # asset: on the renumbered chains it drives the tendon-coupled + # distal J1 joints and leaves the J4 knuckles (and the LFJ5 + # metacarpal) without a drive. Redeclare it against the renamed + # joints so the actuated set matches the PhysX hand physically. + "fingers": _SHADOW_HAND_NEWTON_CFG.actuators["fingers"].replace( + joint_names_expr=[ + "robot0_WR.*", + "robot0_(FF|MF|RF)J(4|3|2)", + "robot0_LFJ(5|4|3|2)", + "robot0_THJ[0-4]", + ], + effort_limit_sim={ + "robot0_WRJ1": 4.785, + "robot0_WRJ0": 2.175, + "robot0_(FF|MF|RF|LF)J2": 0.7245, + "robot0_FFJ(4|3)": 0.9, + "robot0_MFJ(4|3)": 0.9, + "robot0_RFJ(4|3)": 0.9, + "robot0_LFJ(5|4|3)": 0.9, + "robot0_THJ4": 2.3722, + "robot0_THJ3": 1.45, + "robot0_THJ(2|1)": 0.99, + "robot0_THJ0": 0.81, + }, + stiffness=20.0, + damping=2.0, + ), "distal_passive": ImplicitActuatorCfg( - joint_names_expr=["robot0_(FF|MF|RF|LF)J0"], + joint_names_expr=["robot0_(FF|MF|RF|LF)J1"], stiffness=10.0, damping=0.1, friction=1e-2, @@ -178,7 +228,26 @@ def _shadow_hand_cfg( ), }, ) - return preset(default=physx_cfg, physx=physx_cfg, newton_mjwarp=newton_cfg) + ovphysx_cfg = SHADOW_HAND_CFG.replace( + prim_path=prim_path, + # OVPhysX does not expose the fixed-tendon runtime API, so spawn without tendon overrides. + spawn=SHADOW_HAND_CFG.spawn.replace(fixed_tendons_props=None), + init_state=SHADOW_HAND_CFG.init_state.replace(pos=init_pos, rot=init_rot), + ) + return preset(default=physx_cfg, physx=physx_cfg, newton_mjwarp=newton_cfg, ovphysx=ovphysx_cfg) + + +# Per-hand presets shared by the Direct environment and the manager scene. +RIGHT_HAND_CFG = _shadow_hand_cfg( + prim_path="/World/envs/env_.*/RightRobot", + init_pos=(0.0, 0.0, 0.5), + init_rot=(0.0, 0.0, 0.0, 1.0), +) +LEFT_HAND_CFG = _shadow_hand_cfg( + prim_path="/World/envs/env_.*/LeftRobot", + init_pos=(0.0, -1.0, 0.5), + init_rot=(0.0, 0.0, 1.0, 0.0), +) @configclass @@ -196,7 +265,7 @@ class ObjectCfg(PresetCfg): physx = RigidObjectCfg( prim_path="/World/envs/env_.*/object", spawn=sim_utils.SphereCfg( - radius=0.0335, + radius=OBJECT_RADIUS, visual_material=sim_utils.PreviewSurfaceCfg(diffuse_color=(0.8, 1.0, 0.0)), physics_material=sim_utils.RigidBodyMaterialCfg(static_friction=0.7), rigid_props=sim_utils.RigidBodyPropertiesCfg( @@ -217,7 +286,7 @@ class ObjectCfg(PresetCfg): newton_mjwarp = RigidObjectCfg( prim_path="/World/envs/env_.*/object", spawn=sim_utils.SphereCfg( - radius=0.0335, + radius=OBJECT_RADIUS, visual_material=sim_utils.PreviewSurfaceCfg(diffuse_color=(0.8, 1.0, 0.0)), rigid_props=sim_utils.RigidBodyPropertiesCfg( kinematic_enabled=False, @@ -229,7 +298,8 @@ class ObjectCfg(PresetCfg): ), init_state=RigidObjectCfg.InitialStateCfg(pos=(0.0, -0.39, 0.54), rot=(0.0, 0.0, 0.0, 1.0)), ) - default = physx + ovphysx = physx # OvPhysX is PhysX-based; reuse the PhysX-tuned rigid sphere + default = newton_mjwarp @configclass @@ -258,10 +328,14 @@ class PhysicsCfg(PresetCfg): update_data_interval=2, ccd_iterations=50, # bumped from default 35 for multi-finger contact geometry ), - num_substeps=2, + # 4 substeps (vs the single-agent port's 2): sustained ball-palm contact + # against the near-passive distal joints explodes ~0.7% of 8192 envs to + # NaN at 2 substeps (zero-action probe, 300 steps); 4 substeps shows none. + num_substeps=4, debug_mode=False, ) - default = physx + ovphysx = OvPhysxCfg() + default = newton_mjwarp @configclass @@ -274,69 +348,24 @@ class HandoverEnvCfg(DirectMARLEnvCfg): observation_spaces = {"right_hand": 157, "left_hand": 157} state_space = 290 - # simulation + # simulation — values mirrored by the manager cfg (guarded by the value-parity test) sim: SimulationCfg = SimulationCfg( dt=1 / 120, - render_interval=decimation, - physics_material=RigidBodyMaterialCfg( - static_friction=1.0, - dynamic_friction=1.0, - ), + render_interval=2, + physics_material=RigidBodyMaterialBaseCfg(static_friction=1.0, dynamic_friction=1.0), physics=PhysicsCfg(), ) + # robot - right_robot_cfg: PresetCfg = _shadow_hand_cfg( - prim_path="/World/envs/env_.*/RightRobot", - init_pos=(0.0, 0.0, 0.5), - init_rot=(0.0, 0.0, 0.0, 1.0), - ) - left_robot_cfg: PresetCfg = _shadow_hand_cfg( - prim_path="/World/envs/env_.*/LeftRobot", - init_pos=(0.0, -1.0, 0.5), - init_rot=(0.0, 0.0, 1.0, 0.0), - ) - actuated_joint_names = [ - "robot0_WRJ1", - "robot0_WRJ0", - "robot0_FFJ3", - "robot0_FFJ2", - "robot0_FFJ1", - "robot0_MFJ3", - "robot0_MFJ2", - "robot0_MFJ1", - "robot0_RFJ3", - "robot0_RFJ2", - "robot0_RFJ1", - "robot0_LFJ4", - "robot0_LFJ3", - "robot0_LFJ2", - "robot0_LFJ1", - "robot0_THJ4", - "robot0_THJ3", - "robot0_THJ2", - "robot0_THJ1", - "robot0_THJ0", - ] - fingertip_body_names = [ - "robot0_ffdistal", - "robot0_mfdistal", - "robot0_rfdistal", - "robot0_lfdistal", - "robot0_thdistal", - ] + right_robot_cfg: PresetCfg = RIGHT_HAND_CFG + left_robot_cfg: PresetCfg = LEFT_HAND_CFG + actuated_joint_names: PresetCfg = ACTUATED_JOINT_NAMES_PRESET + fingertip_body_names = FINGERTIP_BODY_NAMES # in-hand object object_cfg: ObjectCfg = ObjectCfg() # goal object - goal_object_cfg: VisualizationMarkersCfg = VisualizationMarkersCfg( - prim_path="/Visuals/goal_marker", - markers={ - "goal": sim_utils.SphereCfg( - radius=0.0335, - visual_material=sim_utils.PreviewSurfaceCfg(diffuse_color=(0.4, 0.3, 1.0)), - ), - }, - ) + goal_object_cfg: VisualizationMarkersCfg = GOAL_MARKER_CFG # scene scene: InteractiveSceneCfg = InteractiveSceneCfg(num_envs=2048, env_spacing=1.5, replicate_physics=True) diff --git a/source/isaaclab_tasks/isaaclab_tasks/core/handover/handover_manager_env_cfg.py b/source/isaaclab_tasks/isaaclab_tasks/core/handover/handover_manager_env_cfg.py new file mode 100644 index 000000000000..c9514befa4bb --- /dev/null +++ b/source/isaaclab_tasks/isaaclab_tasks/core/handover/handover_manager_env_cfg.py @@ -0,0 +1,200 @@ +# 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 Shadow Hand handover task.""" + +import isaaclab.sim as sim_utils +from isaaclab.assets import AssetBaseCfg +from isaaclab.envs import ManagerBasedRLEnvCfg +from isaaclab.managers import EventTermCfg as EventTerm +from isaaclab.managers import ObservationGroupCfg as ObsGroup +from isaaclab.managers import ObservationTermCfg as ObsTerm +from isaaclab.managers import RewardTermCfg as RewTerm +from isaaclab.managers import SceneEntityCfg +from isaaclab.managers import TerminationTermCfg as DoneTerm +from isaaclab.scene import InteractiveSceneCfg +from isaaclab.sim.spawners.materials import RigidBodyMaterialBaseCfg +from isaaclab.utils.configclass import configclass + +import isaaclab_tasks.core.handover.mdp as mdp +from isaaclab_tasks.core.handover.handover_common import ( + ACTUATED_JOINT_NAMES_PRESET, + FINGERTIP_BODY_NAMES, +) +from isaaclab_tasks.core.handover.handover_env_cfg import ( + LEFT_HAND_CFG, + RIGHT_HAND_CFG, + ObjectCfg, + PhysicsCfg, +) +from isaaclab_tasks.utils import PresetCfg + + +@configclass +class HandoverManagerSceneCfg(PresetCfg): + """Backend-specific scene cloning settings for handover.""" + + @configclass + class SceneCfg(InteractiveSceneCfg): + """Scene shared by the handover Manager backend alternatives.""" + + num_envs = 2048 + env_spacing = 1.5 + replicate_physics = True + + ground = AssetBaseCfg( + prim_path="/World/ground", + spawn=sim_utils.GroundPlaneCfg(), + ) + right_hand: PresetCfg = RIGHT_HAND_CFG + left_hand: PresetCfg = LEFT_HAND_CFG + object: ObjectCfg = ObjectCfg() + light = AssetBaseCfg( + prim_path="/World/Light", + spawn=sim_utils.DomeLightCfg(intensity=2000.0, color=(0.75, 0.75, 0.75)), + ) + + physx = SceneCfg(clone_in_fabric=True) + newton_mjwarp = SceneCfg(clone_in_fabric=False) + ovphysx = physx + default = newton_mjwarp + + +@configclass +class CommandsCfg: + """Handover goal command.""" + + object_pose = mdp.HandoverCommandCfg(asset_name="object", debug_vis=True) + + +@configclass +class ActionsCfg: + """Two-hand action terms, ordered right then left like the Direct adapter.""" + + right_hand = mdp.EMAJointPositionToLimitsActionCfg( + asset_name="right_hand", + joint_names=ACTUATED_JOINT_NAMES_PRESET, + alpha=1.0, + rescale_to_limits=True, + ) + left_hand = mdp.EMAJointPositionToLimitsActionCfg( + asset_name="left_hand", + joint_names=ACTUATED_JOINT_NAMES_PRESET, + alpha=1.0, + rescale_to_limits=True, + ) + + +def _hand_entity(name: str) -> SceneEntityCfg: + return SceneEntityCfg(name, joint_names=".*") + + +def _fingertip_entity(name: str) -> SceneEntityCfg: + return SceneEntityCfg(name, body_names=FINGERTIP_BODY_NAMES) + + +@configclass +class ObservationsCfg: + """Single-agent observations matching the Direct MARL adapter.""" + + @configclass + class PolicyCfg(ObsGroup): + # Right agent: 133 hand dimensions followed by 24 object/goal dimensions. + # soft limits equal the hard limits here: soft_joint_pos_limits_factor defaults to 1.0 + right_joint_pos = ObsTerm(func=mdp.joint_pos_limit_normalized, params={"asset_cfg": _hand_entity("right_hand")}) + right_joint_vel = ObsTerm(func=mdp.joint_vel, scale=0.2, params={"asset_cfg": _hand_entity("right_hand")}) + right_fingertip_pos = ObsTerm(func=mdp.fingertip_pos, params={"asset_cfg": _fingertip_entity("right_hand")}) + right_fingertip_quat = ObsTerm(func=mdp.fingertip_quat, params={"asset_cfg": _fingertip_entity("right_hand")}) + right_fingertip_vel = ObsTerm(func=mdp.fingertip_vel, params={"asset_cfg": _fingertip_entity("right_hand")}) + right_action = ObsTerm(func=mdp.hand_action, params={"action_name": "right_hand"}) + right_object_goal = ObsTerm( + func=mdp.object_goal, + params={"command_name": "object_pose", "object_cfg": SceneEntityCfg("object"), "vel_obs_scale": 0.2}, + ) + + # Left agent: the same 157-dimensional layout. + # soft limits equal the hard limits here: soft_joint_pos_limits_factor defaults to 1.0 + left_joint_pos = ObsTerm(func=mdp.joint_pos_limit_normalized, params={"asset_cfg": _hand_entity("left_hand")}) + left_joint_vel = ObsTerm(func=mdp.joint_vel, scale=0.2, params={"asset_cfg": _hand_entity("left_hand")}) + left_fingertip_pos = ObsTerm(func=mdp.fingertip_pos, params={"asset_cfg": _fingertip_entity("left_hand")}) + left_fingertip_quat = ObsTerm(func=mdp.fingertip_quat, params={"asset_cfg": _fingertip_entity("left_hand")}) + left_fingertip_vel = ObsTerm(func=mdp.fingertip_vel, params={"asset_cfg": _fingertip_entity("left_hand")}) + left_action = ObsTerm(func=mdp.hand_action, params={"action_name": "left_hand"}) + left_object_goal = ObsTerm( + func=mdp.object_goal, + params={"command_name": "object_pose", "object_cfg": SceneEntityCfg("object"), "vel_obs_scale": 0.2}, + ) + + def __post_init__(self): + self.enable_corruption = False + self.concatenate_terms = True + + policy: PolicyCfg = PolicyCfg() + + +@configclass +class EventCfg: + """Reset distributions matching the Direct handover environment.""" + + reset_handover = EventTerm( + func=mdp.reset_handover_state, + mode="reset", + params={ + "position_noise": 0.01, + "joint_position_noise": 0.2, + "joint_velocity_noise": 0.0, + "action_names": ("right_hand", "left_hand"), + }, + ) + + +@configclass +class RewardsCfg: + """Summed two-agent reward exposed by the Direct single-agent adapter.""" + + handover = RewTerm( + func=mdp.HandoverReward, + weight=1.0, + params={ + "command_name": "object_pose", + "distance_scale": 20.0, + "success_distance_threshold": 0.1, + "object_cfg": SceneEntityCfg("object"), + }, + ) + + +@configclass +class TerminationsCfg: + """Termination conditions matching the Direct single-agent adapter.""" + + object_out_of_reach = DoneTerm( + func=mdp.root_height_below_minimum, + params={"minimum_height": 0.24, "asset_cfg": SceneEntityCfg("object")}, + ) + time_out = DoneTerm(func=mdp.time_out, time_out=True) + + +@configclass +class HandoverManagerEnvCfg(ManagerBasedRLEnvCfg): + """Manager-based handover environment matching the Direct RSL-RL view.""" + + scene: HandoverManagerSceneCfg = HandoverManagerSceneCfg() + observations: ObservationsCfg = ObservationsCfg() + actions: ActionsCfg = ActionsCfg() + commands: CommandsCfg = CommandsCfg() + rewards: RewardsCfg = RewardsCfg() + terminations: TerminationsCfg = TerminationsCfg() + events: EventCfg = EventCfg() + + def __post_init__(self): + self.decimation = 2 + self.episode_length_s = 7.5 + # simulation — mirrors the Direct cfg (guarded by the value-parity test) + self.sim.dt = 1 / 120 + self.sim.render_interval = self.decimation + self.sim.physics_material = RigidBodyMaterialBaseCfg(static_friction=1.0, dynamic_friction=1.0) + self.sim.physics = PhysicsCfg() + self.viewer.eye = (2.0, 2.0, 2.0) diff --git a/source/isaaclab_tasks/isaaclab_tasks/core/handover/mdp/__init__.py b/source/isaaclab_tasks/isaaclab_tasks/core/handover/mdp/__init__.py new file mode 100644 index 000000000000..3a693209a889 --- /dev/null +++ b/source/isaaclab_tasks/isaaclab_tasks/core/handover/mdp/__init__.py @@ -0,0 +1,10 @@ +# 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 MDP terms for the Shadow Hand handover task.""" + +from isaaclab.utils.module import lazy_export + +lazy_export() diff --git a/source/isaaclab_tasks/isaaclab_tasks/core/handover/mdp/__init__.pyi b/source/isaaclab_tasks/isaaclab_tasks/core/handover/mdp/__init__.pyi new file mode 100644 index 000000000000..93e16dbd48c4 --- /dev/null +++ b/source/isaaclab_tasks/isaaclab_tasks/core/handover/mdp/__init__.pyi @@ -0,0 +1,30 @@ +# 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 + +__all__ = [ + "HandoverCommand", + "HandoverCommandCfg", + "reset_handover_state", + "fingertip_pos", + "fingertip_quat", + "fingertip_vel", + "hand_action", + "object_goal", + "HandoverReward", + "handover_reward", + "evaluate_handover_success", +] + +from .commands import HandoverCommand, HandoverCommandCfg +from .events import reset_handover_state +from .observations import ( + fingertip_pos, + fingertip_quat, + fingertip_vel, + hand_action, + object_goal, +) +from .rewards import HandoverReward, evaluate_handover_success, handover_reward +from isaaclab.envs.mdp import * diff --git a/source/isaaclab_tasks/isaaclab_tasks/core/handover/mdp/commands.py b/source/isaaclab_tasks/isaaclab_tasks/core/handover/mdp/commands.py new file mode 100644 index 000000000000..ee3a9693fe90 --- /dev/null +++ b/source/isaaclab_tasks/isaaclab_tasks/core/handover/mdp/commands.py @@ -0,0 +1,96 @@ +# 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 + +"""Goal-pose command for the manager-based handover task.""" + +from __future__ import annotations + +from collections.abc import Sequence +from dataclasses import MISSING +from typing import TYPE_CHECKING + +import torch + +import isaaclab.utils.math as math_utils +from isaaclab.managers import CommandTerm, CommandTermCfg +from isaaclab.markers import VisualizationMarkers, VisualizationMarkersCfg +from isaaclab.utils.configclass import configclass + +from isaaclab_tasks.core.handover.handover_common import GOAL_MARKER_CFG, GOAL_POSITION_OFFSET + +if TYPE_CHECKING: + from isaaclab.assets import RigidObject + from isaaclab.envs import ManagerBasedRLEnv + + +class HandoverCommand(CommandTerm): + """Sample the fixed-position, random-orientation handover goal pose.""" + + cfg: HandoverCommandCfg + + def __init__(self, cfg: HandoverCommandCfg, env: ManagerBasedRLEnv): + super().__init__(cfg, env) + self._object: RigidObject = env.scene[cfg.asset_name] + offset = torch.tensor(cfg.position_offset, dtype=torch.float, device=self.device) + self.pos_command_e = self._object.data.default_root_pose.torch[:, :3] + offset + self.quat_command_w = torch.zeros(self.num_envs, 4, device=self.device) + self.quat_command_w[:, 3] = 1.0 # identity quaternion in (x, y, z, w) layout + # persistent (num_envs, 7) pose command: the position half is static and written once + # here, the quaternion half is refreshed by _resample_command; `command` returns this + # buffer directly instead of allocating a torch.cat every call + self._command_buf = torch.cat((self.pos_command_e, self.quat_command_w), dim=-1) + self._x_unit = torch.tensor([1.0, 0.0, 0.0], device=self.device).repeat(self.num_envs, 1) + self._y_unit = torch.tensor([0.0, 1.0, 0.0], device=self.device).repeat(self.num_envs, 1) + + @property + def command(self) -> torch.Tensor: + """Goal pose in the environment frame [m, unit quaternion]. + + The returned tensor is a persistent buffer refreshed in place; consumers that + store it across steps must copy it. + """ + return self._command_buf + + def _update_metrics(self) -> None: + pass + + def _resample_command(self, env_ids: Sequence[int]) -> None: + random_values = 2.0 * torch.rand((len(env_ids), 2), device=self.device) - 1.0 + self.quat_command_w[env_ids] = math_utils.quat_mul( + math_utils.quat_from_angle_axis(random_values[:, 0] * torch.pi, self._x_unit[env_ids]), + math_utils.quat_from_angle_axis(random_values[:, 1] * torch.pi, self._y_unit[env_ids]), + ) + # keep the persistent pose-command buffer current (position half is static) + self._command_buf[env_ids, 3:] = self.quat_command_w[env_ids] + + def _update_command(self) -> None: + pass + + def _set_debug_vis_impl(self, debug_vis: bool) -> None: + if debug_vis: + if not hasattr(self, "_goal_visualizer"): + self._goal_visualizer = VisualizationMarkers(self.cfg.goal_visualizer_cfg) + self._goal_visualizer.set_visibility(True) + elif hasattr(self, "_goal_visualizer"): + self._goal_visualizer.set_visibility(False) + + def _debug_vis_callback(self, event) -> None: + del event + self._goal_visualizer.visualize( + translations=self.pos_command_e + self._env.scene.env_origins, + orientations=self.quat_command_w, + ) + + +@configclass +class HandoverCommandCfg(CommandTermCfg): + """Configuration for :class:`HandoverCommand`.""" + + class_type: type[HandoverCommand] = HandoverCommand + resampling_time_range: tuple[float, float] = (1.0e6, 1.0e6) + asset_name: str = MISSING + position_offset: tuple[float, float, float] = GOAL_POSITION_OFFSET + """Goal-position offset from the object's default position [m].""" + goal_visualizer_cfg: VisualizationMarkersCfg = GOAL_MARKER_CFG.replace(prim_path="/Visuals/Command/goal_marker") diff --git a/source/isaaclab_tasks/isaaclab_tasks/core/handover/mdp/events.py b/source/isaaclab_tasks/isaaclab_tasks/core/handover/mdp/events.py new file mode 100644 index 000000000000..e3cc54a125a9 --- /dev/null +++ b/source/isaaclab_tasks/isaaclab_tasks/core/handover/mdp/events.py @@ -0,0 +1,76 @@ +# 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 + +"""Reset events for the manager-based handover task.""" + +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_tasks.core.reorient.mdp.events import random_xy_rotation, sample_joint_positions_within_limits + +if TYPE_CHECKING: + from isaaclab.assets import Articulation, RigidObject + from isaaclab.envs import ManagerBasedRLEnv + + +def reset_handover_state( + env: ManagerBasedRLEnv, + env_ids: Sequence[int], + position_noise: float, + joint_position_noise: float, + joint_velocity_noise: float, + action_names: tuple[str, ...], + right_hand_cfg: SceneEntityCfg = SceneEntityCfg("right_hand"), + left_hand_cfg: SceneEntityCfg = SceneEntityCfg("left_hand"), + object_cfg: SceneEntityCfg = SceneEntityCfg("object"), +) -> None: + """Reset the object and both hands with the Direct task's distributions. + + Args: + env: Environment containing both hands and the object. + env_ids: Environment indices to reset. + position_noise: Object-position noise half-width [m]. + joint_position_noise: Scale applied to sampled joint-position deltas. + joint_velocity_noise: Joint-velocity noise half-width [rad/s]. + action_names: Action terms whose pre-reset raw actions are retained in reset observations. + right_hand_cfg: Right-hand scene entity. + left_hand_cfg: Left-hand scene entity. + object_cfg: Object scene entity. + """ + if not hasattr(env, "_handover_reset_actions"): + env._handover_reset_actions = {} + for action_name in action_names: + raw_action = env.action_manager.get_term(action_name).raw_actions + if action_name not in env._handover_reset_actions: + env._handover_reset_actions[action_name] = torch.zeros_like(raw_action) + env._handover_reset_actions[action_name][env_ids] = raw_action[env_ids] + + object_asset: RigidObject = env.scene[object_cfg.name] + object_pose = object_asset.data.default_root_pose.torch[env_ids].clone() + object_velocity = torch.zeros_like(object_asset.data.default_root_vel.torch[env_ids]) + position_delta = math_utils.sample_uniform(-1.0, 1.0, (len(env_ids), 3), device=env.device) + object_pose[:, :3] += position_noise * position_delta + env.scene.env_origins[env_ids] + object_pose[:, 3:7] = random_xy_rotation(len(env_ids), env.device) + object_asset.write_root_pose_to_sim_index(root_pose=object_pose, env_ids=env_ids) + object_asset.write_root_velocity_to_sim_index(root_velocity=object_velocity, env_ids=env_ids) + + for hand_cfg in (right_hand_cfg, left_hand_cfg): + hand: Articulation = env.scene[hand_cfg.name] + default_position = hand.data.default_joint_pos.torch[env_ids] + limits = hand.data.joint_limits.torch[env_ids] + joint_position = sample_joint_positions_within_limits(default_position, limits, joint_position_noise) + velocity_sample = math_utils.sample_uniform(-1.0, 1.0, (len(env_ids), hand.num_joints), device=env.device) + joint_velocity = hand.data.default_joint_vel.torch[env_ids] + joint_velocity_noise * velocity_sample + + hand.set_joint_position_target_index(target=joint_position, env_ids=env_ids) + hand.write_joint_position_to_sim_index(position=joint_position, env_ids=env_ids) + hand.write_joint_velocity_to_sim_index(velocity=joint_velocity, env_ids=env_ids) diff --git a/source/isaaclab_tasks/isaaclab_tasks/core/handover/mdp/observations.py b/source/isaaclab_tasks/isaaclab_tasks/core/handover/mdp/observations.py new file mode 100644 index 000000000000..5c1395e9df3b --- /dev/null +++ b/source/isaaclab_tasks/isaaclab_tasks/core/handover/mdp/observations.py @@ -0,0 +1,81 @@ +# 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 + +"""Observation terms for the manager-based handover task.""" + +from __future__ import annotations + +from typing import TYPE_CHECKING + +import torch + +import isaaclab.utils.math as math_utils +from isaaclab.managers import SceneEntityCfg + +# Handover reuses the reorientation fingertip observation terms verbatim. +from isaaclab_tasks.core.reorient.mdp.observations import ( # noqa: F401 + fingertip_pos, + fingertip_quat, + fingertip_vel, +) + +if TYPE_CHECKING: + from isaaclab.assets import RigidObject + from isaaclab.envs import ManagerBasedRLEnv + + +def hand_action(env: ManagerBasedRLEnv, action_name: str) -> torch.Tensor: + """Return one hand's Direct-compatible raw action across resets. + + Args: + env: Environment containing the action term and episode-length buffer. + action_name: Action term whose raw action is observed. + + Returns: + Current raw actions, retaining pre-reset actions while episode length is zero. + """ + raw_action = env.action_manager.get_term(action_name).raw_actions + reset_actions = getattr(env, "_handover_reset_actions", None) + episode_length_buf = getattr(env, "episode_length_buf", None) + if reset_actions is None or action_name not in reset_actions or episode_length_buf is None: + return raw_action + return torch.where((episode_length_buf == 0).unsqueeze(-1), reset_actions[action_name], raw_action) + + +def object_goal( + env: ManagerBasedRLEnv, command_name: str, object_cfg: SceneEntityCfg, vel_obs_scale: float +) -> torch.Tensor: + """Return the 24-dimensional object and handover-goal observation block. + + Position components use [m], linear velocities [m/s], angular velocities + [rad/s], and quaternion components are unitless. The angular-velocity + scale arrives as the ``vel_obs_scale`` term param, wired at declaration. + + Args: + env: Environment containing the object and goal command. + command_name: Goal command term name. + object_cfg: Object scene entity. + vel_obs_scale: Angular-velocity observation scale. + + Returns: + Object pose, spatial velocity, goal pose, and quaternion error, shape ``(num_envs, 24)``. + """ + object_asset: RigidObject = env.scene[object_cfg.name] + command_term = env.command_manager.get_term(command_name) + object_pos_e = object_asset.data.root_pos_w.torch - env.scene.env_origins + object_quat = object_asset.data.root_quat_w.torch + quat_error = math_utils.quat_mul(object_quat, math_utils.quat_conjugate(command_term.quat_command_w)) + return torch.cat( + ( + object_pos_e, + object_quat, + object_asset.data.root_lin_vel_w.torch, + vel_obs_scale * object_asset.data.root_ang_vel_w.torch, + command_term.pos_command_e, + command_term.quat_command_w, + quat_error, + ), + dim=-1, + ) diff --git a/source/isaaclab_tasks/isaaclab_tasks/core/handover/mdp/rewards.py b/source/isaaclab_tasks/isaaclab_tasks/core/handover/mdp/rewards.py new file mode 100644 index 000000000000..c3dbee23591d --- /dev/null +++ b/source/isaaclab_tasks/isaaclab_tasks/core/handover/mdp/rewards.py @@ -0,0 +1,92 @@ +# 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 + +"""Reward terms for the manager-based handover task.""" + +from __future__ import annotations + +from collections.abc import Sequence +from typing import TYPE_CHECKING + +import torch + +from isaaclab.managers import ManagerTermBase, RewardTermCfg, SceneEntityCfg + +from isaaclab_tasks.core.reorient.mdp.rewards import EpisodeErrorRecorder + +if TYPE_CHECKING: + from isaaclab.assets import RigidObject + from isaaclab.envs import ManagerBasedRLEnv + + +def handover_reward(goal_distance: torch.Tensor, distance_scale: float) -> torch.Tensor: + """Return one hand's Direct reward for the current object-goal distance.""" + return 2.0 * torch.exp(-distance_scale * goal_distance) + + +@torch.jit.script +def evaluate_handover_success( + object_position: torch.Tensor, target_position: torch.Tensor, success_distance_threshold: float +) -> tuple[torch.Tensor, torch.Tensor]: + """Evaluate handover success while exposing its physical error. + + Args: + object_position: Object positions [m]. + target_position: Goal positions [m]. + success_distance_threshold: Exclusive successful goal-distance threshold [m]. + + Returns: + Per-environment success flags and object-to-goal distances [m]. + """ + goal_distance = torch.linalg.norm(object_position - target_position, ord=2, dim=-1) + return goal_distance < success_distance_threshold, goal_distance + + +class HandoverReward(ManagerTermBase): + """Compute summed hand rewards and track sticky per-episode success. + + The scalar task parameters arrive as term params, set at the configuration + declaration site to match the Direct environment's values. + """ + + def __init__(self, cfg: RewardTermCfg, env: ManagerBasedRLEnv): + super().__init__(cfg, env) + self._episode_succeeded = torch.zeros(self.num_envs, dtype=torch.bool, device=self.device) + self._goal_distance = EpisodeErrorRecorder(self.num_envs, self.device) + + def reset(self, env_ids: Sequence[int] | None = None) -> None: + if env_ids is None: + env_ids = slice(None) + # 0-dim device tensor: avoids a host sync here; consumers read it at logging cadence + self._env.extras.setdefault("log", {})["Metrics/success_rate"] = self._episode_succeeded[env_ids].float().mean() + for statistic, value in self._goal_distance.reset(env_ids).items(): + self._env.extras["log"][f"Diagnostics/episode_min_goal_distance_{statistic}"] = value + self._episode_succeeded[env_ids] = False + + def __call__( + self, + env: ManagerBasedRLEnv, + command_name: str, + distance_scale: float, + success_distance_threshold: float, + object_cfg: SceneEntityCfg = SceneEntityCfg("object"), + ) -> torch.Tensor: + object_asset: RigidObject = env.scene[object_cfg.name] + object_pos = object_asset.data.root_pos_w.torch - env.scene.env_origins + goal_pos = env.command_manager.get_command(command_name)[:, :3] + succeeded, goal_distance = evaluate_handover_success(object_pos, goal_pos, success_distance_threshold) + self._goal_distance.update(goal_distance) + per_agent_reward = handover_reward(goal_distance, distance_scale) + + # tensors, not .item(): a host sync every step stalls the GPU at large env counts + goal_distance_mean = goal_distance.mean() + env.extras.setdefault("log", {})["dist_reward"] = per_agent_reward.mean() + env.extras["log"]["dist_goal"] = goal_distance_mean + env.extras["log"]["Metrics/goal_distance"] = goal_distance_mean + self._episode_succeeded |= succeeded + + # RewardManager applies step_dt to all terms. Divide here to preserve the Direct + # environment's per-control-step reward exposed after summing both agents. + return 2.0 * per_agent_reward / env.step_dt diff --git a/source/isaaclab_tasks/isaaclab_tasks/core/reorient/__init__.py b/source/isaaclab_tasks/isaaclab_tasks/core/reorient/__init__.py index 81f1425e1bf8..b12e5252b837 100644 --- a/source/isaaclab_tasks/isaaclab_tasks/core/reorient/__init__.py +++ b/source/isaaclab_tasks/isaaclab_tasks/core/reorient/__init__.py @@ -8,10 +8,9 @@ This package consolidates the direct-workflow and manager-based-workflow in-hand manipulation tasks, where a dexterous hand reorients an object to match a goal orientation. The shared direct base environment lives in -:mod:`~isaaclab_tasks.core.reorient.reorient_direct_env` and the shared manager-based -base configuration in :mod:`~isaaclab_tasks.core.reorient.reorient_manager_env_cfg`. -Robot-specific tasks are organized under the ``config`` subpackage -(``config/allegro_hand`` and ``config/shadow_hand``). +:mod:`~isaaclab_tasks.core.reorient.reorient_direct_env`; the manager-based +configurations live with their robot-specific tasks under the ``config`` +subpackage (``config/allegro_hand`` and ``config/shadow_hand``). These environments are based on the `dexterous cube manipulation`_ environments provided in IsaacGymEnvs repository from NVIDIA. However, they contain certain diff --git a/source/isaaclab_tasks/isaaclab_tasks/core/reorient/config/allegro_hand/allegro_hand_common.py b/source/isaaclab_tasks/isaaclab_tasks/core/reorient/config/allegro_hand/allegro_hand_common.py new file mode 100644 index 000000000000..a3ed24dc514d --- /dev/null +++ b/source/isaaclab_tasks/isaaclab_tasks/core/reorient/config/allegro_hand/allegro_hand_common.py @@ -0,0 +1,114 @@ +# 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 + +"""Allegro Hand identity shared by the Direct and manager-based reorientation tasks. + +Asset and marker configurations, joint/body name lists, backend physics +presets, and the sim mixin. No task tunables. +""" + +from isaaclab_newton.physics import MJWarpSolverCfg, NewtonCfg +from isaaclab_ovphysx.physics import OvPhysxCfg +from isaaclab_physx.physics import PhysxCfg + +import isaaclab.sim as sim_utils +from isaaclab.assets import ArticulationCfg, RigidObjectCfg +from isaaclab.markers import VisualizationMarkersCfg +from isaaclab.utils.assets import ISAAC_NUCLEUS_DIR +from isaaclab.utils.configclass import configclass + +from isaaclab_tasks.utils import PresetCfg + +from isaaclab_assets.robots.allegro import ALLEGRO_HAND_CFG + + +@configclass +class ObjectCfg(PresetCfg): + physx = RigidObjectCfg( + prim_path="/World/envs/env_.*/object", + spawn=sim_utils.UsdFileCfg( + usd_path=f"{ISAAC_NUCLEUS_DIR}/Props/Blocks/DexCube/dex_cube_instanceable.usd", + rigid_props=sim_utils.RigidBodyPropertiesCfg( + kinematic_enabled=False, + disable_gravity=False, + enable_gyroscopic_forces=True, + solver_position_iteration_count=8, + solver_velocity_iteration_count=0, + sleep_threshold=0.005, + stabilization_threshold=0.0025, + max_depenetration_velocity=1000.0, + ), + mass_props=sim_utils.MassPropertiesCfg(density=400.0), + scale=(1.2, 1.2, 1.2), + ), + init_state=RigidObjectCfg.InitialStateCfg(pos=(0.0, -0.17, 0.56), rot=(0.0, 0.0, 0.0, 1.0)), + ) + newton_mjwarp = ArticulationCfg( + prim_path="/World/envs/env_.*/object", + spawn=sim_utils.UsdFileCfg( + usd_path=f"{ISAAC_NUCLEUS_DIR}/Props/Blocks/DexCube/dex_cube_instanceable.usd", + mass_props=sim_utils.MassPropertiesCfg(density=400.0), + scale=(1.2, 1.2, 1.2), + ), + init_state=ArticulationCfg.InitialStateCfg( + pos=(0.0, -0.17, 0.565), rot=(0.0, 0.0, 0.0, 1.0), joint_pos={}, joint_vel={} + ), + actuators={}, + articulation_root_prim_path="", + ) + ovphysx = RigidObjectCfg( + prim_path="/World/envs/env_.*/object", + spawn=sim_utils.UsdFileCfg( + usd_path=f"{ISAAC_NUCLEUS_DIR}/Props/Blocks/DexCube/dex_cube_instanceable.usd", + rigid_props=sim_utils.RigidBodyPropertiesCfg( + kinematic_enabled=False, + disable_gravity=False, + enable_gyroscopic_forces=True, + solver_position_iteration_count=8, + solver_velocity_iteration_count=0, + sleep_threshold=0.005, + stabilization_threshold=0.0025, + max_depenetration_velocity=1000.0, + ), + mass_props=sim_utils.MassPropertiesCfg(density=400.0), + scale=(1.2, 1.2, 1.2), + ), + init_state=RigidObjectCfg.InitialStateCfg(pos=(0.0, -0.17, 0.56), rot=(0.0, 0.0, 0.0, 1.0)), + ) + default = newton_mjwarp + + +@configclass +class PhysicsCfg(PresetCfg): + physx = PhysxCfg( + bounce_threshold_velocity=0.2, + ) + newton_mjwarp = NewtonCfg( + solver_cfg=MJWarpSolverCfg( + integrator="implicitfast", + njmax=80, + nconmax=70, + impratio=10.0, + cone="elliptic", + update_data_interval=2, + ), + num_substeps=2, + ) + ovphysx = OvPhysxCfg() + default = newton_mjwarp + + +# Scene pieces shared verbatim by the manager-based variant. +ROBOT_CFG = ALLEGRO_HAND_CFG.replace(prim_path="/World/envs/env_.*/Robot") +OBJECT_CFG = ObjectCfg() +GOAL_OBJECT_CFG = VisualizationMarkersCfg( + prim_path="/Visuals/goal_marker", + markers={ + "goal": sim_utils.UsdFileCfg( + usd_path=f"{ISAAC_NUCLEUS_DIR}/Props/Blocks/DexCube/dex_cube_instanceable.usd", + scale=(1.2, 1.2, 1.2), + ) + }, +) diff --git a/source/isaaclab_tasks/isaaclab_tasks/core/reorient/config/allegro_hand/allegro_hand_direct_env_cfg.py b/source/isaaclab_tasks/isaaclab_tasks/core/reorient/config/allegro_hand/allegro_hand_direct_env_cfg.py index 8f32326763f6..66c9ab696fb8 100644 --- a/source/isaaclab_tasks/isaaclab_tasks/core/reorient/config/allegro_hand/allegro_hand_direct_env_cfg.py +++ b/source/isaaclab_tasks/isaaclab_tasks/core/reorient/config/allegro_hand/allegro_hand_direct_env_cfg.py @@ -3,104 +3,23 @@ # # SPDX-License-Identifier: BSD-3-Clause - -from isaaclab_newton.physics import MJWarpSolverCfg, NewtonCfg -from isaaclab_ovphysx.physics import OvPhysxCfg -from isaaclab_physx.physics import PhysxCfg - -import isaaclab.sim as sim_utils -from isaaclab.assets import ArticulationCfg, RigidObjectCfg +from isaaclab.assets import ArticulationCfg from isaaclab.envs import DirectRLEnvCfg from isaaclab.markers import VisualizationMarkersCfg from isaaclab.scene import InteractiveSceneCfg from isaaclab.sim import SimulationCfg -from isaaclab.sim.spawners.materials.physics_materials_cfg import RigidBodyMaterialCfg -from isaaclab.utils.assets import ISAAC_NUCLEUS_DIR +from isaaclab.sim.spawners.materials import RigidBodyMaterialBaseCfg from isaaclab.utils.configclass import configclass -from isaaclab_tasks.utils import PresetCfg - -from isaaclab_assets.robots.allegro import ALLEGRO_HAND_CFG - +from isaaclab_tasks.core.reorient.config.allegro_hand.allegro_hand_common import ( + GOAL_OBJECT_CFG, + OBJECT_CFG, + ROBOT_CFG, + ObjectCfg, + PhysicsCfg, +) -@configclass -class ObjectCfg(PresetCfg): - physx = RigidObjectCfg( - prim_path="/World/envs/env_.*/object", - spawn=sim_utils.UsdFileCfg( - usd_path=f"{ISAAC_NUCLEUS_DIR}/Props/Blocks/DexCube/dex_cube_instanceable.usd", - rigid_props=sim_utils.RigidBodyPropertiesCfg( - kinematic_enabled=False, - disable_gravity=False, - enable_gyroscopic_forces=True, - solver_position_iteration_count=8, - solver_velocity_iteration_count=0, - sleep_threshold=0.005, - stabilization_threshold=0.0025, - max_depenetration_velocity=1000.0, - ), - mass_props=sim_utils.MassPropertiesCfg(density=400.0), - scale=(1.2, 1.2, 1.2), - ), - init_state=RigidObjectCfg.InitialStateCfg(pos=(0.0, -0.17, 0.56), rot=(0.0, 0.0, 0.0, 1.0)), - ) - newton_mjwarp = ArticulationCfg( - prim_path="/World/envs/env_.*/object", - spawn=sim_utils.UsdFileCfg( - usd_path=f"{ISAAC_NUCLEUS_DIR}/Props/Blocks/DexCube/dex_cube_instanceable.usd", - mass_props=sim_utils.MassPropertiesCfg(density=400.0), - scale=(1.2, 1.2, 1.2), - ), - init_state=ArticulationCfg.InitialStateCfg( - pos=(0.0, -0.17, 0.565), rot=(0.0, 0.0, 0.0, 1.0), joint_pos={}, joint_vel={} - ), - actuators={}, - articulation_root_prim_path="", - ) - ovphysx = RigidObjectCfg( - prim_path="/World/envs/env_.*/object", - spawn=sim_utils.UsdFileCfg( - usd_path=f"{ISAAC_NUCLEUS_DIR}/Props/Blocks/DexCube/dex_cube_instanceable.usd", - rigid_props=sim_utils.RigidBodyPropertiesCfg( - kinematic_enabled=False, - disable_gravity=False, - enable_gyroscopic_forces=True, - solver_position_iteration_count=8, - solver_velocity_iteration_count=0, - sleep_threshold=0.005, - stabilization_threshold=0.0025, - max_depenetration_velocity=1000.0, - ), - mass_props=sim_utils.MassPropertiesCfg(density=400.0), - scale=(1.2, 1.2, 1.2), - ), - init_state=RigidObjectCfg.InitialStateCfg(pos=(0.0, -0.17, 0.56), rot=(0.0, 0.0, 0.0, 1.0)), - ) - default = physx - - -@configclass -class PhysicsCfg(PresetCfg): - physx = PhysxCfg( - bounce_threshold_velocity=0.2, - ) - newton_mjwarp = NewtonCfg( - solver_cfg=MJWarpSolverCfg( - solver="newton", - integrator="implicitfast", - njmax=80, - nconmax=70, - impratio=10.0, - cone="elliptic", - update_data_interval=2, - iterations=100, - # save_to_mjcf="AllegroHand.xml", - ), - num_substeps=2, - debug_mode=False, - ) - ovphysx = OvPhysxCfg() - default = physx +from isaaclab_assets.robots.allegro import ALLEGRO_ACTUATED_JOINT_NAMES, ALLEGRO_FINGERTIP_BODY_NAMES @configclass @@ -113,59 +32,30 @@ class AllegroHandEnvCfg(DirectRLEnvCfg): state_space = 0 asymmetric_obs = False obs_type = "full" - # simulation + + # simulation — values mirrored by the manager cfg (guarded by the value-parity test) sim: SimulationCfg = SimulationCfg( dt=1 / 120, - render_interval=decimation, - physics_material=RigidBodyMaterialCfg( - static_friction=1.0, - dynamic_friction=1.0, - ), + render_interval=4, + physics_material=RigidBodyMaterialBaseCfg(static_friction=1.0, dynamic_friction=1.0), physics=PhysicsCfg(), ) # robot - robot_cfg: ArticulationCfg = ALLEGRO_HAND_CFG.replace(prim_path="/World/envs/env_.*/Robot") + robot_cfg: ArticulationCfg = ROBOT_CFG - actuated_joint_names = [ - "index_joint_0", - "middle_joint_0", - "ring_joint_0", - "thumb_joint_0", - "index_joint_1", - "index_joint_2", - "index_joint_3", - "middle_joint_1", - "middle_joint_2", - "middle_joint_3", - "ring_joint_1", - "ring_joint_2", - "ring_joint_3", - "thumb_joint_1", - "thumb_joint_2", - "thumb_joint_3", - ] - fingertip_body_names = [ - "index_link_3", - "middle_link_3", - "ring_link_3", - "thumb_link_3", - ] + actuated_joint_names = ALLEGRO_ACTUATED_JOINT_NAMES + fingertip_body_names = ALLEGRO_FINGERTIP_BODY_NAMES # in-hand object - object_cfg: ObjectCfg = ObjectCfg() + object_cfg: ObjectCfg = OBJECT_CFG # goal object - goal_object_cfg: VisualizationMarkersCfg = VisualizationMarkersCfg( - prim_path="/Visuals/goal_marker", - markers={ - "goal": sim_utils.UsdFileCfg( - usd_path=f"{ISAAC_NUCLEUS_DIR}/Props/Blocks/DexCube/dex_cube_instanceable.usd", - scale=(1.2, 1.2, 1.2), - ) - }, - ) + goal_object_cfg: VisualizationMarkersCfg = GOAL_OBJECT_CFG # scene scene: InteractiveSceneCfg = InteractiveSceneCfg( - 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, ) # reset reset_position_noise = 0.01 # range of position at reset @@ -176,8 +66,8 @@ class AllegroHandEnvCfg(DirectRLEnvCfg): rot_reward_scale = 1.0 rot_eps = 0.1 action_penalty_scale = -0.0002 - reach_goal_bonus = 250 - fall_penalty = 0 + reach_goal_bonus = 250.0 + fall_penalty = 0.0 fall_dist = 0.24 vel_obs_scale = 0.2 success_tolerance = 0.2 diff --git a/source/isaaclab_tasks/isaaclab_tasks/core/reorient/config/allegro_hand/allegro_hand_manager_env_cfg.py b/source/isaaclab_tasks/isaaclab_tasks/core/reorient/config/allegro_hand/allegro_hand_manager_env_cfg.py index 6ad5db896b0e..c06a22d638df 100644 --- a/source/isaaclab_tasks/isaaclab_tasks/core/reorient/config/allegro_hand/allegro_hand_manager_env_cfg.py +++ b/source/isaaclab_tasks/isaaclab_tasks/core/reorient/config/allegro_hand/allegro_hand_manager_env_cfg.py @@ -3,36 +3,315 @@ # # SPDX-License-Identifier: BSD-3-Clause +"""Manager-based counterpart of the Allegro Hand Direct reorientation task.""" + +import isaaclab.sim as sim_utils +from isaaclab.assets import ArticulationCfg, 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.sim.spawners.materials import RigidBodyMaterialBaseCfg from isaaclab.utils.configclass import configclass -from isaaclab_tasks.core.reorient.reorient_manager_env_cfg import ReorientObjectEnvCfg +import isaaclab_tasks.core.reorient.mdp as mdp +from isaaclab_tasks.core.reorient.config.allegro_hand.allegro_hand_common import ( + GOAL_OBJECT_CFG, + OBJECT_CFG, + ROBOT_CFG, + ObjectCfg, + PhysicsCfg, +) +from isaaclab_tasks.core.reorient.reorient_common import GOAL_MARKER_POSITION, IN_HAND_POS_OFFSET +from isaaclab_tasks.utils import PresetCfg -## -# Pre-defined configs -## -from isaaclab_assets import ALLEGRO_HAND_CFG # isort: skip +from isaaclab_assets.robots.allegro import ALLEGRO_ACTUATED_JOINT_NAMES, ALLEGRO_FINGERTIP_BODY_NAMES @configclass -class AllegroCubeEnvCfg(ReorientObjectEnvCfg): - def __post_init__(self): - # post init of parent - super().__post_init__() +class AllegroCubeSceneCfg(PresetCfg): + """Backend-specific scene cloning settings matching the Direct task.""" + + @configclass + class SceneCfg(InteractiveSceneCfg): + """Allegro scene shared by the backend alternatives.""" + + num_envs = 8192 + env_spacing = 0.75 + replicate_physics = True + + ground = AssetBaseCfg(prim_path="/World/ground", spawn=sim_utils.GroundPlaneCfg()) + robot: ArticulationCfg = 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)), + ) - # switch robot to allegro hand - self.scene.robot = ALLEGRO_HAND_CFG.replace(prim_path="{ENV_REGEX_NS}/Robot") - # enable clone in fabric - self.scene.clone_in_fabric = True + physx = SceneCfg(clone_in_fabric=True) + newton_mjwarp = SceneCfg(clone_in_fabric=False) + ovphysx = physx + default = newton_mjwarp + + def set_num_envs(self, num_envs: int) -> None: + """Set the environment count on every backend alternative.""" + for scene in (self.physx, self.newton_mjwarp, self.ovphysx, self.default): + scene.num_envs = num_envs + + +@configclass +class CommandsCfg: + """Object pose goal matching the Direct in-hand target.""" + + object_pose = mdp.ReorientEpisodeCommandCfg( + asset_name="object", + init_pos_offset=IN_HAND_POS_OFFSET, + update_goal_on_success=True, + orientation_success_threshold=0.2, + make_quat_unique=False, + fixed_marker_pos=GOAL_MARKER_POSITION, + goal_pose_visualizer_cfg=GOAL_OBJECT_CFG, + debug_vis=True, + ) + + +@configclass +class ActionsCfg: + """Sixteen actuated Allegro Hand joints in Direct order.""" + + joint_pos = mdp.EMAJointPositionToLimitsActionCfg( + asset_name="robot", + joint_names=ALLEGRO_ACTUATED_JOINT_NAMES, + alpha=1.0, + rescale_to_limits=True, + ) + + +@configclass +class ObservationsCfg: + """Full 124-dimensional state observation in Direct order.""" + + @configclass + class PolicyCfg(ObsGroup): + joint_pos = ObsTerm( + func=mdp.joint_pos_limit_normalized, + params={"asset_cfg": SceneEntityCfg("robot", joint_names=".*", preserve_order=False)}, + ) + joint_vel = ObsTerm( + func=mdp.joint_vel, + scale=0.2, + params={"asset_cfg": SceneEntityCfg("robot", joint_names=".*", preserve_order=False)}, + ) + object_pos = ObsTerm(func=mdp.root_pos_w, params={"asset_cfg": SceneEntityCfg("object")}) + object_quat = ObsTerm( + func=mdp.root_quat_w, + params={"asset_cfg": SceneEntityCfg("object"), "make_quat_unique": False}, + ) + object_lin_vel = ObsTerm(func=mdp.root_lin_vel_w, params={"asset_cfg": SceneEntityCfg("object")}) + object_ang_vel = ObsTerm( + func=mdp.root_ang_vel_w, + scale=0.2, + params={"asset_cfg": SceneEntityCfg("object")}, + ) + goal_pose = ObsTerm(func=mdp.generated_commands, params={"command_name": "object_pose"}) + goal_quat_diff = ObsTerm( + func=mdp.goal_quat_diff, + params={"asset_cfg": SceneEntityCfg("object"), "command_name": "object_pose", "make_quat_unique": False}, + ) + fingertip_pos = ObsTerm( + func=mdp.fingertip_pos, + params={ + "asset_cfg": SceneEntityCfg("robot", body_names=ALLEGRO_FINGERTIP_BODY_NAMES, preserve_order=False) + }, + ) + fingertip_quat = ObsTerm( + func=mdp.fingertip_quat, + params={ + "asset_cfg": SceneEntityCfg("robot", body_names=ALLEGRO_FINGERTIP_BODY_NAMES, preserve_order=False) + }, + ) + fingertip_vel = ObsTerm( + func=mdp.fingertip_vel, + params={ + "asset_cfg": SceneEntityCfg("robot", body_names=ALLEGRO_FINGERTIP_BODY_NAMES, preserve_order=False) + }, + ) + last_action = ObsTerm(func=mdp.reorient_last_action, params={"action_name": "joint_pos"}) + + def __post_init__(self): + self.enable_corruption = False + self.concatenate_terms = True + + policy: PolicyCfg = PolicyCfg() + + +@configclass +class EventCfg: + """Reset distributions matching the Direct task, plus opt-in domain randomization. + + The domain-randomization terms reproduce the legacy manager recipe. They are + startup-mode terms and are dropped by default (see + :attr:`AllegroCubeEnvCfg.enable_domain_randomization`): the Direct task has no + domain randomization, and the validated benchmark thresholds were calibrated + without it. Enabling them requires retraining. + """ + + # -- opt-in domain randomization (legacy manager recipe parameters) + robot_physics_material = EventTerm( + func=mdp.randomize_rigid_body_material, + mode="startup", + params={ + "asset_cfg": SceneEntityCfg("robot", body_names=".*"), + "static_friction_range": (0.7, 1.3), + "dynamic_friction_range": (0.7, 1.3), + "restitution_range": (0.0, 0.0), + "num_buckets": 250, + }, + ) + robot_scale_mass = EventTerm( + func=mdp.randomize_rigid_body_mass, + mode="startup", + params={ + "asset_cfg": SceneEntityCfg("robot", body_names=".*"), + "mass_distribution_params": (0.95, 1.05), + "operation": "scale", + }, + ) + robot_joint_stiffness_and_damping = EventTerm( + func=mdp.randomize_actuator_gains, + mode="startup", + params={ + "asset_cfg": SceneEntityCfg("robot", joint_names=".*"), + "stiffness_distribution_params": (0.3, 3.0), + "damping_distribution_params": (0.75, 1.5), + "operation": "scale", + "distribution": "log_uniform", + }, + ) + object_physics_material = EventTerm( + func=mdp.randomize_rigid_body_material, + mode="startup", + params={ + "asset_cfg": SceneEntityCfg("object", body_names=".*"), + "static_friction_range": (0.7, 1.3), + "dynamic_friction_range": (0.7, 1.3), + "restitution_range": (0.0, 0.0), + "num_buckets": 250, + }, + ) + object_scale_mass = EventTerm( + func=mdp.randomize_rigid_body_mass, + mode="startup", + params={ + "asset_cfg": SceneEntityCfg("object"), + "mass_distribution_params": (0.4, 1.6), + "operation": "scale", + }, + ) + + reset_state = EventTerm( + func=mdp.reset_reorient_state, + mode="reset", + params={ + "position_noise": 0.01, + "joint_position_noise": 0.2, + "joint_velocity_noise": 0.0, + "action_name": "joint_pos", + }, + ) + + +@configclass +class RewardsCfg: + """Direct-compatible reward and success accounting.""" + + reorient = RewTerm( + func=mdp.ReorientReward, + weight=1.0, + params={ + "command_name": "object_pose", + "distance_scale": -10.0, + "rotation_scale": 1.0, + "rotation_epsilon": 0.1, + "action_penalty_scale": -0.0002, + "success_tolerance": 0.2, + "success_bonus": 250.0, + "fall_distance": 0.24, + "fall_penalty": 0.0, + "averaging_factor": 0.1, + "success_count_threshold": 1, + "object_cfg": SceneEntityCfg("object"), + }, + ) + + +@configclass +class TerminationsCfg: + """Termination conditions matching the Direct task.""" + + object_out_of_reach = DoneTerm( + func=mdp.object_reorientation_out_of_reach, + params={ + "threshold": 0.24, + "command_name": "object_pose", + "object_cfg": SceneEntityCfg("object"), + }, + ) + time_out = DoneTerm(func=mdp.time_out, time_out=True) + + +@configclass +class AllegroCubeEnvCfg(ManagerBasedRLEnvCfg): + """Manager-based Allegro Hand task with Direct-compatible semantics.""" + + scene: AllegroCubeSceneCfg = AllegroCubeSceneCfg() + observations: ObservationsCfg = ObservationsCfg() + actions: ActionsCfg = ActionsCfg() + commands: CommandsCfg = CommandsCfg() + rewards: RewardsCfg = RewardsCfg() + terminations: TerminationsCfg = TerminationsCfg() + events: EventCfg = EventCfg() + + enable_domain_randomization: bool = False + """Enable the legacy startup domain-randomization terms. + + Disabled by default: the validated reference training runs and the benchmark + thresholds were produced without domain randomization, so enabling it + requires retraining and threshold recalibration. + """ + + _DOMAIN_RANDOMIZATION_TERMS = ( + "robot_physics_material", + "robot_scale_mass", + "robot_joint_stiffness_and_damping", + "object_physics_material", + "object_scale_mass", + ) + + def __post_init__(self): + self.decimation = 4 + self.episode_length_s = 10.0 + # simulation — mirrors the Direct cfg (guarded by the value-parity test) + self.sim.dt = 1 / 120 + self.sim.render_interval = self.decimation + self.sim.physics_material = RigidBodyMaterialBaseCfg(static_friction=1.0, dynamic_friction=1.0) + self.sim.physics = PhysicsCfg() + self.viewer.eye = (2.0, 2.0, 2.0) + if not self.enable_domain_randomization: + for term_name in self._DOMAIN_RANDOMIZATION_TERMS: + setattr(self.events, term_name, None) @configclass class AllegroCubeEnvCfg_PLAY(AllegroCubeEnvCfg): + """Reduced, deterministic play configuration.""" + def __post_init__(self): - # post init of parent super().__post_init__() - # make a smaller scene for play - self.scene.num_envs = 50 - # disable randomization for play + self.scene.set_num_envs(50) self.observations.policy.enable_corruption = False - # remove termination due to timeouts self.terminations.time_out = None 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..9483421c5dfa 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,24 +17,59 @@ 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, disable_env_checker=True, kwargs={ - "env_cfg_entry_point": f"{__name__}.shadow_hand_env_cfg:ShadowHandEnvCfg", + "env_cfg_entry_point": f"{__name__}.shadow_hand_direct_env_cfg:ShadowHandEnvCfg", "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-OpenAI-FF", + entry_point="isaaclab.envs:ManagerBasedRLEnv", + disable_env_checker=True, + kwargs={ + "env_cfg_entry_point": f"{__name__}.shadow_hand_openai_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_openai_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, disable_env_checker=True, kwargs={ - "env_cfg_entry_point": f"{__name__}.shadow_hand_env_cfg:ShadowHandOpenAIEnvCfg", + "env_cfg_entry_point": f"{__name__}.shadow_hand_direct_env_cfg:ShadowHandOpenAIEnvCfg", "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", @@ -46,8 +81,9 @@ entry_point=reorient_direct_entry, disable_env_checker=True, kwargs={ - "env_cfg_entry_point": f"{__name__}.shadow_hand_env_cfg:ShadowHandOpenAIEnvCfg", + "env_cfg_entry_point": f"{__name__}.shadow_hand_direct_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", }, ) @@ -55,12 +91,34 @@ # Vision # ------- +gym.register( + id="Isaac-Reorient-Cube-Shadow-Camera", + entry_point="isaaclab.envs:ManagerBasedRLEnv", + disable_env_checker=True, + kwargs={ + "env_cfg_entry_point": f"{__name__}.shadow_hand_camera_manager_env_cfg:ShadowHandCameraManagerEnvCfg", + "rsl_rl_cfg_entry_point": f"{agents.__name__}.rsl_rl_ppo_cfg:ShadowHandCameraFFPPORunnerCfg", + "rl_games_cfg_entry_point": f"{agents.__name__}:rl_games_ppo_camera_cfg.yaml", + }, +) + +gym.register( + id="Isaac-Reorient-Cube-Shadow-Camera-Play", + entry_point="isaaclab.envs:ManagerBasedRLEnv", + disable_env_checker=True, + kwargs={ + "env_cfg_entry_point": f"{__name__}.shadow_hand_camera_manager_env_cfg:ShadowHandCameraManagerPlayEnvCfg", + "rsl_rl_cfg_entry_point": f"{agents.__name__}.rsl_rl_ppo_cfg:ShadowHandCameraFFPPORunnerCfg", + "rl_games_cfg_entry_point": f"{agents.__name__}:rl_games_ppo_camera_cfg.yaml", + }, +) + gym.register( id="Isaac-Reorient-Cube-Shadow-Camera-Direct", - entry_point=f"{__name__}.shadow_hand_camera_env:ShadowHandCameraEnv", + entry_point=f"{__name__}.shadow_hand_direct_camera_env:ShadowHandCameraEnv", disable_env_checker=True, kwargs={ - "env_cfg_entry_point": f"{__name__}.shadow_hand_camera_env_cfg:ShadowHandCameraEnvCfg", + "env_cfg_entry_point": f"{__name__}.shadow_hand_direct_camera_env_cfg:ShadowHandCameraEnvCfg", "rsl_rl_cfg_entry_point": f"{agents.__name__}.rsl_rl_ppo_cfg:ShadowHandCameraFFPPORunnerCfg", "rl_games_cfg_entry_point": f"{agents.__name__}:rl_games_ppo_camera_cfg.yaml", }, @@ -68,10 +126,10 @@ gym.register( id="Isaac-Reorient-Cube-Shadow-Camera-Direct-Play", - entry_point=f"{__name__}.shadow_hand_camera_env:ShadowHandCameraEnv", + entry_point=f"{__name__}.shadow_hand_direct_camera_env:ShadowHandCameraEnv", disable_env_checker=True, kwargs={ - "env_cfg_entry_point": f"{__name__}.shadow_hand_camera_env_cfg:ShadowHandCameraEnvPlayCfg", + "env_cfg_entry_point": f"{__name__}.shadow_hand_direct_camera_env_cfg:ShadowHandCameraEnvPlayCfg", "rsl_rl_cfg_entry_point": f"{agents.__name__}.rsl_rl_ppo_cfg:ShadowHandCameraFFPPORunnerCfg", "rl_games_cfg_entry_point": f"{agents.__name__}:rl_games_ppo_camera_cfg.yaml", }, @@ -79,10 +137,10 @@ gym.register( id="Isaac-Reorient-Cube-Shadow-Camera-Benchmark-Direct", - entry_point=f"{__name__}.shadow_hand_camera_env:ShadowHandCameraEnv", + entry_point=f"{__name__}.shadow_hand_direct_camera_env:ShadowHandCameraEnv", disable_env_checker=True, kwargs={ - "env_cfg_entry_point": f"{__name__}.shadow_hand_camera_env_cfg:ShadowHandCameraBenchmarkEnvCfg", + "env_cfg_entry_point": f"{__name__}.shadow_hand_direct_camera_env_cfg:ShadowHandCameraBenchmarkEnvCfg", "rsl_rl_cfg_entry_point": f"{agents.__name__}.rsl_rl_ppo_cfg:ShadowHandCameraFFPPORunnerCfg", "rl_games_cfg_entry_point": f"{agents.__name__}:rl_games_ppo_camera_cfg.yaml", }, 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/feature_extractor.py b/source/isaaclab_tasks/isaaclab_tasks/core/reorient/config/shadow_hand/feature_extractor.py index 56c159a1446e..aab4c74d8af2 100644 --- a/source/isaaclab_tasks/isaaclab_tasks/core/reorient/config/shadow_hand/feature_extractor.py +++ b/source/isaaclab_tasks/isaaclab_tasks/core/reorient/config/shadow_hand/feature_extractor.py @@ -3,8 +3,11 @@ # # SPDX-License-Identifier: BSD-3-Clause +from __future__ import annotations + import glob import os +from typing import TYPE_CHECKING import torch import torch.nn as nn @@ -13,6 +16,12 @@ from isaaclab.sensors import save_images_to_file from isaaclab.utils.configclass import configclass +# re-exported for backward compatibility; the shared implementation lives in the family math root +from isaaclab_tasks.core.reorient.mdp.observations import compute_cube_keypoints # noqa: F401 + +if TYPE_CHECKING: + pass + # Number of output channels for each supported camera data type. _DATA_TYPE_CHANNELS: dict[str, int] = { "rgb": 3, diff --git a/source/isaaclab_tasks/isaaclab_tasks/core/reorient/config/shadow_hand/shadow_hand_camera_env_cfg.py b/source/isaaclab_tasks/isaaclab_tasks/core/reorient/config/shadow_hand/shadow_hand_camera_env_cfg.py deleted file mode 100644 index ae5e80d9fb60..000000000000 --- a/source/isaaclab_tasks/isaaclab_tasks/core/reorient/config/shadow_hand/shadow_hand_camera_env_cfg.py +++ /dev/null @@ -1,174 +0,0 @@ -# Copyright (c) 2022-2026, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md). -# All rights reserved. -# -# SPDX-License-Identifier: BSD-3-Clause - -from __future__ import annotations - -import isaaclab.sim as sim_utils -from isaaclab.scene import InteractiveSceneCfg -from isaaclab.sensors import CameraCfg -from isaaclab.utils.configclass import configclass - -from isaaclab_tasks.core.reorient.config.shadow_hand.feature_extractor import FeatureExtractorCfg -from isaaclab_tasks.core.reorient.config.shadow_hand.shadow_hand_env_cfg import ShadowHandEnvCfg -from isaaclab_tasks.utils import PresetCfg -from isaaclab_tasks.utils.presets import MultiBackendRendererCfg - - -@configclass -class _ShadowHandBaseTiledCameraCfg(CameraCfg): - """Base camera configuration for the shadow hand vision environment. - - This is an internal config used by :class:`ShadowHandTiledCameraCfg` presets and - by derived env configs that hard-code a specific data type. It embeds - :class:`~isaaclab_tasks.utils.MultiBackendRendererCfg` so the renderer backend can - still be selected via the ``presets`` CLI argument. - """ - - prim_path: str = "/World/envs/env_.*/Camera" - offset: CameraCfg.OffsetCfg = CameraCfg.OffsetCfg( - pos=(0, -0.35, 1.0), rot=(0.0, 0.7071, 0.0, 0.7071), convention="world" - ) - data_types: list[str] = [] - spawn: sim_utils.PinholeCameraCfg = sim_utils.PinholeCameraCfg( - focal_length=24.0, focus_distance=400.0, horizontal_aperture=20.955, clipping_range=(0.1, 20.0) - ) - width: int = 120 - height: int = 120 - renderer_cfg: MultiBackendRendererCfg = MultiBackendRendererCfg() - - -@configclass -class ShadowHandTiledCameraCfg(PresetCfg): - """Camera data-type presets for the shadow hand vision environment. - - Each preset selects which image modalities are captured. The selected data types must - match :attr:`FeatureExtractorCfg.data_types` so the CNN receives the expected channels. - - Select a data-type preset via the ``presets`` CLI argument, e.g.:: - - presets = rgb # RGB only (3 channels) - presets = albedo # albedo (3 channels) - presets = simple_shading_constant_diffuse # simple shading, constant diffuse (3 channels) - - Renderer and data-type presets can be combined:: - - presets = newton_renderer, rgb - """ - - default: _ShadowHandBaseTiledCameraCfg = _ShadowHandBaseTiledCameraCfg( - data_types=["rgb", "depth", "semantic_segmentation"] - ) - """Default: RGB + depth + semantic segmentation (7 CNN input channels).""" - - full: _ShadowHandBaseTiledCameraCfg = _ShadowHandBaseTiledCameraCfg( - data_types=["rgb", "depth", "semantic_segmentation"] - ) - """Full modalities: RGB + depth + semantic segmentation (7 channels). Alias for default.""" - - rgb: _ShadowHandBaseTiledCameraCfg = _ShadowHandBaseTiledCameraCfg(data_types=["rgb"]) - """RGB only (3 CNN input channels).""" - - albedo: _ShadowHandBaseTiledCameraCfg = _ShadowHandBaseTiledCameraCfg(data_types=["albedo"]) - """Albedo (3 CNN input channels).""" - - simple_shading_constant_diffuse: _ShadowHandBaseTiledCameraCfg = _ShadowHandBaseTiledCameraCfg( - data_types=["simple_shading_constant_diffuse"] - ) - """Simple shading with constant diffuse (3 CNN input channels).""" - - simple_shading_diffuse_mdl: _ShadowHandBaseTiledCameraCfg = _ShadowHandBaseTiledCameraCfg( - data_types=["simple_shading_diffuse_mdl"] - ) - """Simple shading with diffuse MDL (3 CNN input channels).""" - - simple_shading_full_mdl: _ShadowHandBaseTiledCameraCfg = _ShadowHandBaseTiledCameraCfg( - data_types=["simple_shading_full_mdl"] - ) - """Simple shading with full MDL (3 CNN input channels).""" - - depth: _ShadowHandBaseTiledCameraCfg = _ShadowHandBaseTiledCameraCfg(data_types=["depth"]) - """Depth only (1 channel). - - .. warning:: - This preset is intended for **benchmarking only**. The keypoint-regression CNN - cannot be meaningfully trained from depth alone. Use it with - :class:`ShadowHandCameraBenchmarkEnvCfg` (``feature_extractor.enabled=False``) - to measure pure depth-rendering throughput, e.g.:: - - presets=depth # depth rendering, default renderer - presets=depth,newton_renderer # depth rendering with Newton renderer - presets=depth,ovrtx # depth rendering with OVRTX renderer - """ - - semantic_segmentation: _ShadowHandBaseTiledCameraCfg = _ShadowHandBaseTiledCameraCfg( - data_types=["semantic_segmentation"] - ) - """Semantic segmentation (3 CNN input channels).""" - - -@configclass -class ShadowHandCameraEnvCfg(ShadowHandEnvCfg): - # scene - scene: InteractiveSceneCfg = InteractiveSceneCfg(num_envs=1225, env_spacing=2.0, replicate_physics=True) - - # camera — data-type and renderer backend selectable via CLI presets - tiled_camera: ShadowHandTiledCameraCfg = ShadowHandTiledCameraCfg() - feature_extractor: FeatureExtractorCfg = FeatureExtractorCfg() - - # env - observation_space = 164 + 27 # state observation + vision CNN embedding - state_space = 187 + 27 # asymmetric states + vision CNN embedding - - def validate_config(self): - """Check renderer/data-type and feature-extractor compatibility.""" - renderer_type = getattr(self.tiled_camera.renderer_cfg, "renderer_type", None) - warp_supported = {"rgb", "depth", "distance_to_camera", "distance_to_image_plane", "normals"} - if renderer_type == "newton_warp": - unsupported = set(self.tiled_camera.data_types) - warp_supported - if unsupported: - raise ValueError( - f"Warp renderer only supports data types {sorted(warp_supported)}, " - f"but the camera is configured with unsupported types: {sorted(unsupported)}. " - "Choose a compatible preset, e.g. presets=newton_renderer,rgb." - ) - - non_depth_data_types = set(self.tiled_camera.data_types).difference( - {"depth", "distance_to_image_plane", "distance_to_camera"} - ) - if self.tiled_camera.data_types and not non_depth_data_types and self.feature_extractor.enabled: - raise ValueError( - "Depth-only camera data type is intended for benchmarking only. " - "The keypoint-regression CNN cannot be meaningfully trained from depth alone. " - "Disable the feature extractor with 'feature_extractor.enabled=False' " - "(e.g. use Isaac-Reorient-Cube-Shadow-Camera-Benchmark-Direct), " - "or choose a data type that includes colour, e.g. presets=rgb." - ) - - -@configclass -class ShadowHandCameraEnvPlayCfg(ShadowHandCameraEnvCfg): - # scene - scene: InteractiveSceneCfg = InteractiveSceneCfg(num_envs=64, env_spacing=2.0, replicate_physics=True) - # inference for CNN - feature_extractor: FeatureExtractorCfg = FeatureExtractorCfg(train=False, load_checkpoint=True) - - -@configclass -class ShadowHandCameraBenchmarkEnvCfg(ShadowHandCameraEnvCfg): - """Benchmark configuration with the feature extractor CNN disabled. - - The tiled camera renders frames each step as normal, but the CNN forward pass is - bypassed — zero embeddings are returned instead. This isolates rendering throughput - from CNN inference overhead when profiling. - - The renderer backend and camera data types can still be selected via ``presets``:: - - presets = newton_renderer # benchmark with Newton renderer - presets = ovrtx # benchmark with OVRTX renderer - presets = rgb # benchmark RGB rendering only - presets = depth, newton_renderer # benchmark depth rendering with Newton - """ - - feature_extractor: FeatureExtractorCfg = FeatureExtractorCfg(enabled=False) diff --git a/source/isaaclab_tasks/isaaclab_tasks/core/reorient/config/shadow_hand/shadow_hand_camera_manager_env_cfg.py b/source/isaaclab_tasks/isaaclab_tasks/core/reorient/config/shadow_hand/shadow_hand_camera_manager_env_cfg.py new file mode 100644 index 000000000000..1f90e3d6dd52 --- /dev/null +++ b/source/isaaclab_tasks/isaaclab_tasks/core/reorient/config/shadow_hand/shadow_hand_camera_manager_env_cfg.py @@ -0,0 +1,141 @@ +# 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 Shadow Hand camera reorientation task.""" + +from isaaclab.managers import ObservationTermCfg as ObsTerm +from isaaclab.managers import SceneEntityCfg +from isaaclab.sensors import JointWrenchSensorCfg +from isaaclab.utils.configclass import configclass + +import isaaclab_tasks.core.reorient.mdp as mdp +from isaaclab_tasks.core.reorient.config.shadow_hand.feature_extractor import FeatureExtractorCfg +from isaaclab_tasks.core.reorient.config.shadow_hand.shadow_hand_direct_camera_env_cfg import ( + ShadowHandTiledCameraCfg, + validate_shadow_hand_camera_settings, +) +from isaaclab_tasks.core.reorient.config.shadow_hand.shadow_hand_manager_env_cfg import ( + FullStateWithoutActionCfg, + ShadowHandManagerEnvCfg, + _ShadowHandManagerSceneCfg, +) +from isaaclab_tasks.core.reorient.reorient_common import CAMERA_GOAL_MARKER_POSITION, CAMERA_PLAY_NUM_ENVS +from isaaclab_tasks.utils import PresetCfg + +from isaaclab_assets.robots.shadow_hand import SHADOW_FINGERTIP_BODY_NAMES + + +@configclass +class _ShadowHandCameraManagerSceneCfg(_ShadowHandManagerSceneCfg): + """State Manager scene augmented with camera and fingertip-wrench sensors.""" + + num_envs = 1225 + env_spacing = 2.0 + + ground = None + tiled_camera: ShadowHandTiledCameraCfg = ShadowHandTiledCameraCfg() + joint_wrench = JointWrenchSensorCfg(prim_path="{ENV_REGEX_NS}/Robot") + + +@configclass +class ShadowHandCameraManagerSceneCfg(PresetCfg): + """Backend-specific camera scene alternatives for training and benchmarking.""" + + physx = _ShadowHandCameraManagerSceneCfg(clone_in_fabric=True) + newton_mjwarp = _ShadowHandCameraManagerSceneCfg(clone_in_fabric=False) + ovphysx = physx + default = newton_mjwarp + + +@configclass +class ShadowHandCameraManagerPlaySceneCfg(PresetCfg): + """Reduced backend-specific camera scenes for checkpoint playback.""" + + physx = _ShadowHandCameraManagerSceneCfg(num_envs=CAMERA_PLAY_NUM_ENVS, clone_in_fabric=True) + newton_mjwarp = _ShadowHandCameraManagerSceneCfg(num_envs=CAMERA_PLAY_NUM_ENVS, clone_in_fabric=False) + ovphysx = _ShadowHandCameraManagerSceneCfg(num_envs=CAMERA_PLAY_NUM_ENVS, clone_in_fabric=True) + default = newton_mjwarp + + +@configclass +class CameraPolicyCfg(FullStateWithoutActionCfg): + """Direct-compatible 191-dimensional camera actor observation.""" + + last_action = ObsTerm(func=mdp.reorient_last_action, params={"action_name": "joint_pos"}) + camera_features = ObsTerm( + func=mdp.ShadowHandCameraFeatures, + params={ + "feature_extractor_cfg": FeatureExtractorCfg(), + "sensor_cfg": SceneEntityCfg("tiled_camera"), + "object_cfg": SceneEntityCfg("object"), + }, + ) + goal_keypoints = ObsTerm(func=mdp.shadow_hand_goal_keypoints, params={"command_name": "object_pose"}) + + def __post_init__(self): + super().__post_init__() + # Camera actor observations infer object state from pixels. These five + # privileged state terms are present only in the critic. + self.object_pos = None + self.object_quat = None + self.object_lin_vel = None + self.object_ang_vel = None + self.goal_quat_diff = None + + +@configclass +class CameraCriticCfg(FullStateWithoutActionCfg): + """Direct-compatible 214-dimensional asymmetric camera critic state.""" + + fingertip_wrench = ObsTerm( + func=mdp.fingertip_wrench, + scale=10.0, + params={ + "sensor_cfg": SceneEntityCfg("joint_wrench", body_names=SHADOW_FINGERTIP_BODY_NAMES, preserve_order=False) + }, + ) + last_action = ObsTerm(func=mdp.reorient_last_action, params={"action_name": "joint_pos"}) + camera_features = ObsTerm(func=mdp.shadow_hand_camera_cached_features) + + +@configclass +class CameraObservationsCfg: + """Camera actor and asymmetric critic observation groups.""" + + policy: CameraPolicyCfg = CameraPolicyCfg() + critic: CameraCriticCfg = CameraCriticCfg() + + +@configclass +class ShadowHandCameraManagerEnvCfg(ShadowHandManagerEnvCfg): + """Manager-based camera task with exact Direct dynamics and observations.""" + + # only the fields that differ from ShadowHandManagerEnvCfg are overridden + scene: ShadowHandCameraManagerSceneCfg = ShadowHandCameraManagerSceneCfg() + observations: CameraObservationsCfg = CameraObservationsCfg() + feature_extractor: FeatureExtractorCfg = FeatureExtractorCfg() + + def __post_init__(self): + super().__post_init__() + # camera tasks display the goal inside the tiled camera's frustum + self.commands.object_pose.fixed_marker_pos = CAMERA_GOAL_MARKER_POSITION + self.observations.policy.camera_features.params["feature_extractor_cfg"] = self.feature_extractor + + def validate_config(self): + """Check every unresolved scene alternative or the selected camera pipeline.""" + if isinstance(self.scene, PresetCfg): + scenes = (self.scene.physx, self.scene.newton_mjwarp, self.scene.ovphysx) + else: + scenes = (self.scene,) + for scene in scenes: + validate_shadow_hand_camera_settings(scene.tiled_camera, self.feature_extractor) + + +@configclass +class ShadowHandCameraManagerPlayEnvCfg(ShadowHandCameraManagerEnvCfg): + """Manager camera task configured for checkpoint playback.""" + + scene: ShadowHandCameraManagerPlaySceneCfg = ShadowHandCameraManagerPlaySceneCfg() + feature_extractor: FeatureExtractorCfg = FeatureExtractorCfg(train=False, load_checkpoint=True) diff --git a/source/isaaclab_tasks/isaaclab_tasks/core/reorient/config/shadow_hand/shadow_hand_env_cfg.py b/source/isaaclab_tasks/isaaclab_tasks/core/reorient/config/shadow_hand/shadow_hand_common.py similarity index 64% rename from source/isaaclab_tasks/isaaclab_tasks/core/reorient/config/shadow_hand/shadow_hand_env_cfg.py rename to source/isaaclab_tasks/isaaclab_tasks/core/reorient/config/shadow_hand/shadow_hand_common.py index f6bdae2bd0d8..88436d5aa672 100644 --- a/source/isaaclab_tasks/isaaclab_tasks/core/reorient/config/shadow_hand/shadow_hand_env_cfg.py +++ b/source/isaaclab_tasks/isaaclab_tasks/core/reorient/config/shadow_hand/shadow_hand_common.py @@ -3,20 +3,24 @@ # # SPDX-License-Identifier: BSD-3-Clause +"""Shadow Hand identity shared by the Direct and manager-based reorientation tasks. + +Asset and marker configurations, joint/body name lists, backend physics and +domain-randomization presets, and the sim mixins. No task tunables: reward +scales and thresholds live inline in the workflow configuration files. +""" + from isaaclab_newton.physics import KaminoSolverCfg, MJWarpSolverCfg, NewtonCfg +from isaaclab_ovphysx.physics import OvPhysxCfg from isaaclab_physx.physics import PhysxCfg import isaaclab.envs.mdp as mdp import isaaclab.sim as sim_utils from isaaclab.actuators import ImplicitActuatorCfg from isaaclab.assets import ArticulationCfg, RigidObjectCfg -from isaaclab.envs import DirectRLEnvCfg from isaaclab.managers import EventTermCfg as EventTerm from isaaclab.managers import SceneEntityCfg from isaaclab.markers import VisualizationMarkersCfg -from isaaclab.scene import InteractiveSceneCfg -from isaaclab.sim import SimulationCfg -from isaaclab.sim.spawners.materials.physics_materials_cfg import RigidBodyMaterialCfg from isaaclab.utils.assets import ISAAC_NUCLEUS_DIR from isaaclab.utils.configclass import configclass from isaaclab.utils.noise import GaussianNoiseCfg, NoiseModelWithAdditiveBiasCfg @@ -134,7 +138,8 @@ class PhysxEventCfg: class ShadowHandEventCfg(PresetCfg): physx = PhysxEventCfg() newton_mjwarp = NewtonEventCfg() - default = physx + ovphysx = physx # OvPhysX is PhysX-based; reuse the PhysX randomization terms + default = newton_mjwarp newton_kamino = newton_mjwarp @@ -207,7 +212,17 @@ class ShadowHandRobotCfg(PresetCfg): }, soft_joint_pos_limit_factor=1.0, ) - default = physx + ovphysx = SHADOW_HAND_CFG.replace( + prim_path="/World/envs/env_.*/Robot", + # OVPhysX does not expose the fixed-tendon runtime API, so spawn without tendon overrides. + spawn=SHADOW_HAND_CFG.spawn.replace(fixed_tendons_props=None), + init_state=ArticulationCfg.InitialStateCfg( + pos=(0.0, 0.0, 0.5), + rot=(0.0, 0.0, 0.0, 1.0), + joint_pos={".*": 0.0}, + ), + ) + default = newton_mjwarp newton_kamino = newton_mjwarp @@ -247,25 +262,8 @@ class ObjectCfg(PresetCfg): actuators={}, articulation_root_prim_path="", ) - default = physx - newton_kamino = newton_mjwarp - - -@configclass -class ShadowHandSceneCfg(PresetCfg): - """Scene configuration presets for the shadow hand environment. - - PhysX supports ``clone_in_fabric=True`` for faster scene cloning via the Fabric layer. - Newton does not support Fabric cloning, so ``clone_in_fabric`` must be ``False``. - """ - - physx: InteractiveSceneCfg = InteractiveSceneCfg( - num_envs=8192, env_spacing=0.75, replicate_physics=True, clone_in_fabric=True - ) - newton_mjwarp: InteractiveSceneCfg = InteractiveSceneCfg( - num_envs=8192, env_spacing=0.75, replicate_physics=True, clone_in_fabric=False - ) - default: InteractiveSceneCfg = physx + ovphysx = physx # OvPhysX is PhysX-based; use the rigid-body cube, not Newton's articulation + default = newton_mjwarp newton_kamino = newton_mjwarp @@ -278,152 +276,40 @@ class PhysicsCfg(PresetCfg): ) newton_mjwarp = NewtonCfg( solver_cfg=MJWarpSolverCfg( - solver="newton", integrator="implicitfast", njmax=200, nconmax=70, impratio=10.0, cone="elliptic", update_data_interval=2, - iterations=100, ), num_substeps=2, - debug_mode=False, ) - default = physx + ovphysx = OvPhysxCfg() + default = newton_mjwarp newton_kamino = NewtonCfg(solver_cfg=KaminoSolverCfg(max_contacts_per_world=128)) -@configclass -class ShadowHandEnvCfg(DirectRLEnvCfg): - # env - decimation = 2 - episode_length_s = 10.0 - action_space = 20 - observation_space = 157 # (full) - state_space = 0 - asymmetric_obs = False - obs_type = "full" - - # simulation - sim: SimulationCfg = SimulationCfg( - dt=1 / 120, - render_interval=decimation, - physics_material=RigidBodyMaterialCfg(static_friction=1.0, dynamic_friction=1.0), - physics=PhysicsCfg(), - ) - # robot - robot_cfg: ShadowHandRobotCfg = ShadowHandRobotCfg() - actuated_joint_names = [ - "robot0_WRJ1", - "robot0_WRJ0", - "robot0_FFJ3", - "robot0_FFJ2", - "robot0_FFJ1", - "robot0_MFJ3", - "robot0_MFJ2", - "robot0_MFJ1", - "robot0_RFJ3", - "robot0_RFJ2", - "robot0_RFJ1", - "robot0_LFJ4", - "robot0_LFJ3", - "robot0_LFJ2", - "robot0_LFJ1", - "robot0_THJ4", - "robot0_THJ3", - "robot0_THJ2", - "robot0_THJ1", - "robot0_THJ0", - ] - fingertip_body_names = [ - "robot0_ffdistal", - "robot0_mfdistal", - "robot0_rfdistal", - "robot0_lfdistal", - "robot0_thdistal", - ] - - # in-hand object - object_cfg: ObjectCfg = ObjectCfg() - # goal object - goal_object_cfg: VisualizationMarkersCfg = VisualizationMarkersCfg( - prim_path="/Visuals/goal_marker", - markers={ - "goal": sim_utils.UsdFileCfg( - usd_path=f"{ISAAC_NUCLEUS_DIR}/Props/Blocks/DexCube/dex_cube_instanceable.usd", - scale=(1.0, 1.0, 1.0), - ) - }, - ) - # scene — use ShadowHandSceneCfg so that presets=newton_mjwarp disables clone_in_fabric automatically - scene: ShadowHandSceneCfg = ShadowHandSceneCfg() - - # reset - reset_position_noise = 0.01 # range of position at reset - reset_dof_pos_noise = 0.2 # range of dof pos at reset - reset_dof_vel_noise = 0.0 # range of dof vel at reset - # reward scales - dist_reward_scale = -10.0 - rot_reward_scale = 1.0 - rot_eps = 0.1 - action_penalty_scale = -0.0002 - reach_goal_bonus = 250 - fall_penalty = 0 - fall_dist = 0.24 - vel_obs_scale = 0.2 - success_tolerance = 0.1 - max_consecutive_success = 0 - success_count_threshold: int = 1 - """Minimum number of goals reached in an episode to count it as a successful episode.""" - av_factor = 0.1 - act_moving_average = 1.0 - force_torque_obs_scale = 10.0 - - -@configclass -class ShadowHandOpenAIEnvCfg(ShadowHandEnvCfg): - # env - decimation = 3 - episode_length_s = 8.0 - action_space = 20 - observation_space = 42 - state_space = 187 - asymmetric_obs = True - obs_type = "openai" - # simulation - sim: SimulationCfg = SimulationCfg( - dt=1 / 60, - render_interval=decimation, - physics_material=RigidBodyMaterialCfg(static_friction=1.0, dynamic_friction=1.0), - physics=PhysicsCfg(), - ) - # reset - reset_position_noise = 0.01 # range of position at reset - reset_dof_pos_noise = 0.2 # range of dof pos at reset - reset_dof_vel_noise = 0.0 # range of dof vel at reset - # reward scales - dist_reward_scale = -10.0 - rot_reward_scale = 1.0 - rot_eps = 0.1 - action_penalty_scale = -0.0002 - reach_goal_bonus = 250 - fall_penalty = -50 - vel_obs_scale = 0.2 - success_tolerance = 0.4 - max_consecutive_success = 50 - av_factor = 0.1 - act_moving_average = 0.3 - force_torque_obs_scale = 10.0 - # domain randomization config - events: ShadowHandEventCfg = ShadowHandEventCfg() - # at every time-step add gaussian noise + bias. The bias is a gaussian sampled at reset - action_noise_model: NoiseModelWithAdditiveBiasCfg = NoiseModelWithAdditiveBiasCfg( - noise_cfg=GaussianNoiseCfg(mean=0.0, std=0.05, operation="add"), - bias_noise_cfg=GaussianNoiseCfg(mean=0.0, std=0.015, operation="abs"), - ) - # at every time-step add gaussian noise + bias. The bias is a gaussian sampled at reset - observation_noise_model: NoiseModelWithAdditiveBiasCfg = NoiseModelWithAdditiveBiasCfg( - noise_cfg=GaussianNoiseCfg(mean=0.0, std=0.002, operation="add"), - bias_noise_cfg=GaussianNoiseCfg(mean=0.0, std=0.0001, operation="abs"), - ) +# Scene pieces shared verbatim by the manager-based variants. +ROBOT_CFG = ShadowHandRobotCfg() +OBJECT_CFG = ObjectCfg() +GOAL_OBJECT_CFG = VisualizationMarkersCfg( + prim_path="/Visuals/goal_marker", + markers={ + "goal": sim_utils.UsdFileCfg( + usd_path=f"{ISAAC_NUCLEUS_DIR}/Props/Blocks/DexCube/dex_cube_instanceable.usd", + scale=(1.0, 1.0, 1.0), + ) + }, +) + + +# Per-step gaussian noise + reset-sampled bias, shared verbatim by the manager-based variant. +OPENAI_ACTION_NOISE_CFG = NoiseModelWithAdditiveBiasCfg( + noise_cfg=GaussianNoiseCfg(mean=0.0, std=0.05, operation="add"), + bias_noise_cfg=GaussianNoiseCfg(mean=0.0, std=0.015, operation="abs"), +) +OPENAI_OBSERVATION_NOISE_CFG = NoiseModelWithAdditiveBiasCfg( + noise_cfg=GaussianNoiseCfg(mean=0.0, std=0.002, operation="add"), + bias_noise_cfg=GaussianNoiseCfg(mean=0.0, std=0.0001, operation="abs"), +) diff --git a/source/isaaclab_tasks/isaaclab_tasks/core/reorient/config/shadow_hand/shadow_hand_camera_env.py b/source/isaaclab_tasks/isaaclab_tasks/core/reorient/config/shadow_hand/shadow_hand_direct_camera_env.py similarity index 80% rename from source/isaaclab_tasks/isaaclab_tasks/core/reorient/config/shadow_hand/shadow_hand_camera_env.py rename to source/isaaclab_tasks/isaaclab_tasks/core/reorient/config/shadow_hand/shadow_hand_direct_camera_env.py index 90de75ba3ee7..5074099270d1 100644 --- a/source/isaaclab_tasks/isaaclab_tasks/core/reorient/config/shadow_hand/shadow_hand_camera_env.py +++ b/source/isaaclab_tasks/isaaclab_tasks/core/reorient/config/shadow_hand/shadow_hand_direct_camera_env.py @@ -6,6 +6,7 @@ from __future__ import annotations +import warnings from typing import TYPE_CHECKING import torch @@ -14,13 +15,15 @@ from isaaclab import cloner from isaaclab.assets import Articulation, RigidObject from isaaclab.sensors import Camera -from isaaclab.utils.math import quat_apply, scale_transform +from isaaclab.utils.math import scale_transform from isaaclab_tasks.core.reorient.config.shadow_hand.feature_extractor import FeatureExtractor +from isaaclab_tasks.core.reorient.mdp.observations import compute_cube_keypoints +from isaaclab_tasks.core.reorient.reorient_common import CAMERA_GOAL_MARKER_POSITION from isaaclab_tasks.core.reorient.reorient_direct_env import ReorientDirectEnv if TYPE_CHECKING: - from isaaclab_tasks.core.reorient.config.shadow_hand.shadow_hand_camera_env_cfg import ShadowHandCameraEnvCfg + from isaaclab_tasks.core.reorient.config.shadow_hand.shadow_hand_direct_camera_env_cfg import ShadowHandCameraEnvCfg class ShadowHandCameraEnv(ReorientDirectEnv): @@ -40,7 +43,7 @@ def __init__(self, cfg: ShadowHandCameraEnvCfg, render_mode: str | None = None, width=self.cfg.tiled_camera.width, ) # hide goal cubes - self.goal_pos[:, :] = torch.tensor([-0.2, 0.1, 0.6], device=self.device) + self.goal_pos[:, :] = torch.tensor(CAMERA_GOAL_MARKER_POSITION, device=self.device) # keypoints buffer self.gt_keypoints = torch.ones(self.num_envs, 8, 3, dtype=torch.float32, device=self.device) self.goal_keypoints = torch.ones(self.num_envs, 8, 3, dtype=torch.float32, device=self.device) @@ -66,7 +69,7 @@ def _setup_scene(self): def _compute_image_observations(self): # generate ground truth keypoints for in-hand cube - compute_keypoints(pose=torch.cat((self.object_pos, self.object_rot), dim=1), out=self.gt_keypoints) + compute_cube_keypoints(pose=torch.cat((self.object_pos, self.object_rot), dim=1), out=self.gt_keypoints) object_pose = torch.cat([self.object_pos, self.gt_keypoints.view(-1, 24)], dim=-1) @@ -78,7 +81,7 @@ def _compute_image_observations(self): self.embeddings = embeddings.clone().detach() # compute keypoints for goal cube - compute_keypoints( + compute_cube_keypoints( pose=torch.cat((torch.zeros_like(self.goal_pos), self.goal_rot), dim=-1), out=self.goal_keypoints ) @@ -138,32 +141,29 @@ def _get_observations(self) -> dict: return observations -@torch.jit.script def compute_keypoints( pose: torch.Tensor, num_keypoints: int = 8, size: tuple[float, float, float] = (2 * 0.03, 2 * 0.03, 2 * 0.03), out: torch.Tensor | None = None, -): - """Computes positions of 8 corner keypoints of a cube. +) -> torch.Tensor: + """Compute cube keypoints using the shared implementation. + + .. deprecated:: 9.0.0 + Use :func:`compute_cube_keypoints` instead. Args: - pose: Position and orientation of the center of the cube. Shape is (N, 7) - num_keypoints: Number of keypoints to compute. Default = 8 - size: Length of X, Y, Z dimensions of cube. Default = [0.06, 0.06, 0.06] - out: Buffer to store keypoints. If None, a new buffer will be created. + pose: Cube center poses ``(x, y, z, qx, qy, qz, qw)`` [m, unit quaternion]. + num_keypoints: Number of binary-sign corners to compute. + size: Cube side lengths along each axis [m]. + out: Optional output buffer [m], shape ``(num_envs, num_keypoints, 3)``. + + Returns: + Cube-corner positions [m], shape ``(num_envs, num_keypoints, 3)``. """ - num_envs = pose.shape[0] - if out is None: - out = torch.ones(num_envs, num_keypoints, 3, dtype=torch.float32, device=pose.device) - else: - out[:] = 1.0 - for i in range(num_keypoints): - # which dimensions to negate - n = [((i >> k) & 1) == 0 for k in range(3)] - corner_loc = ([(1 if n[k] else -1) * s / 2 for k, s in enumerate(size)],) - corner = torch.tensor(corner_loc, dtype=torch.float32, device=pose.device) * out[:, i, :] - # express corner position in the world frame - out[:, i, :] = pose[:, :3] + quat_apply(pose[:, 3:7], corner) - - return out + warnings.warn( + "compute_keypoints() is deprecated; use compute_cube_keypoints() instead.", + DeprecationWarning, + stacklevel=2, + ) + return compute_cube_keypoints(pose, num_keypoints=num_keypoints, size=size, out=out) diff --git a/source/isaaclab_tasks/isaaclab_tasks/core/reorient/config/shadow_hand/shadow_hand_direct_camera_env_cfg.py b/source/isaaclab_tasks/isaaclab_tasks/core/reorient/config/shadow_hand/shadow_hand_direct_camera_env_cfg.py new file mode 100644 index 000000000000..8fa8b3c5089d --- /dev/null +++ b/source/isaaclab_tasks/isaaclab_tasks/core/reorient/config/shadow_hand/shadow_hand_direct_camera_env_cfg.py @@ -0,0 +1,201 @@ +# Copyright (c) 2022-2026, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +from __future__ import annotations + +import isaaclab.sim as sim_utils +from isaaclab.scene import InteractiveSceneCfg +from isaaclab.sensors import CameraCfg +from isaaclab.utils.configclass import configclass + +from isaaclab_tasks.core.reorient.config.shadow_hand.feature_extractor import FeatureExtractorCfg +from isaaclab_tasks.core.reorient.config.shadow_hand.shadow_hand_direct_env_cfg import ShadowHandEnvCfg +from isaaclab_tasks.core.reorient.reorient_common import CAMERA_PLAY_NUM_ENVS +from isaaclab_tasks.utils import PresetCfg +from isaaclab_tasks.utils.presets import MultiBackendRendererCfg + + +def validate_shadow_hand_camera_settings( + tiled_camera: CameraCfg | ShadowHandTiledCameraCfg, + feature_extractor: FeatureExtractorCfg, +) -> None: + """Validate one resolved or defaulted Shadow Hand camera pipeline.""" + while isinstance(tiled_camera, PresetCfg): + tiled_camera = tiled_camera.default + renderer_cfg = tiled_camera.renderer_cfg + while isinstance(renderer_cfg, PresetCfg): + renderer_cfg = renderer_cfg.default + + renderer_type = getattr(renderer_cfg, "renderer_type", None) + warp_supported = {"rgb", "depth", "distance_to_camera", "distance_to_image_plane", "normals"} + if renderer_type == "newton_warp": + unsupported = set(tiled_camera.data_types) - warp_supported + if unsupported: + raise ValueError( + f"Warp renderer only supports data types {sorted(warp_supported)}, " + f"but the camera is configured with unsupported types: {sorted(unsupported)}. " + "Choose a compatible preset, e.g. presets=newton_renderer,rgb." + ) + + non_depth_data_types = set(tiled_camera.data_types).difference( + {"depth", "distance_to_image_plane", "distance_to_camera"} + ) + if tiled_camera.data_types and not non_depth_data_types and feature_extractor.enabled: + raise ValueError( + "Depth-only camera data type is intended for benchmarking only. " + "The keypoint-regression CNN cannot be meaningfully trained from depth alone. " + "Disable the feature extractor with 'feature_extractor.enabled=False' " + "(e.g. use Isaac-Reorient-Cube-Shadow-Camera-Benchmark-Direct), " + "or choose a data type that includes colour, e.g. presets=rgb." + ) + + +@configclass +class ShadowHandTiledCameraCfg(PresetCfg): + """Camera data-type presets for the shadow hand vision environment. + + Each preset selects which image modalities are captured. The selected data types must + match :attr:`FeatureExtractorCfg.data_types` so the CNN receives the expected channels. + + Select a data-type preset via the ``presets`` CLI argument, e.g.:: + + presets = rgb # RGB only (3 channels) + presets = rgb_depth # RGB + depth (4 channels) + presets = albedo # albedo (3 channels) + presets = simple_shading_constant_diffuse # simple shading, constant diffuse (3 channels) + + Renderer and data-type presets can be combined:: + + presets = newton_renderer, rgb + """ + + @configclass + class BaseTiledCameraCfg(CameraCfg): + """Base camera configuration for the shadow hand vision environment. + + This is an internal config used by :class:`ShadowHandTiledCameraCfg` presets and + by derived env configs that hard-code a specific data type. It embeds + :class:`~isaaclab_tasks.utils.MultiBackendRendererCfg` so the renderer backend can + still be selected via the ``presets`` CLI argument. + """ + + prim_path: str = "/World/envs/env_.*/Camera" + offset: CameraCfg.OffsetCfg = CameraCfg.OffsetCfg( + pos=(0, -0.35, 1.0), rot=(0.0, 0.7071, 0.0, 0.7071), convention="world" + ) + data_types: list[str] = [] + spawn: sim_utils.PinholeCameraCfg = sim_utils.PinholeCameraCfg( + focal_length=24.0, focus_distance=400.0, horizontal_aperture=20.955, clipping_range=(0.1, 20.0) + ) + width: int = 120 + height: int = 120 + renderer_cfg: MultiBackendRendererCfg = MultiBackendRendererCfg() + + default: BaseTiledCameraCfg = BaseTiledCameraCfg(data_types=["rgb", "depth", "semantic_segmentation"]) + """Default: RGB + depth + semantic segmentation (7 CNN input channels).""" + + full: BaseTiledCameraCfg = BaseTiledCameraCfg(data_types=["rgb", "depth", "semantic_segmentation"]) + """Full modalities: RGB + depth + semantic segmentation (7 channels). Alias for default.""" + + rgb: BaseTiledCameraCfg = BaseTiledCameraCfg(data_types=["rgb"]) + """RGB only (3 CNN input channels).""" + + rgb_depth: BaseTiledCameraCfg = BaseTiledCameraCfg(data_types=["rgb", "depth"]) + """RGB and depth (4 CNN input channels).""" + + albedo: BaseTiledCameraCfg = BaseTiledCameraCfg(data_types=["albedo"]) + """Albedo (3 CNN input channels).""" + + simple_shading_constant_diffuse: BaseTiledCameraCfg = BaseTiledCameraCfg( + data_types=["simple_shading_constant_diffuse"] + ) + """Simple shading with constant diffuse (3 CNN input channels).""" + + simple_shading_diffuse_mdl: BaseTiledCameraCfg = BaseTiledCameraCfg(data_types=["simple_shading_diffuse_mdl"]) + """Simple shading with diffuse MDL (3 CNN input channels).""" + + simple_shading_full_mdl: BaseTiledCameraCfg = BaseTiledCameraCfg(data_types=["simple_shading_full_mdl"]) + """Simple shading with full MDL (3 CNN input channels).""" + + depth: BaseTiledCameraCfg = BaseTiledCameraCfg(data_types=["depth"]) + """Depth only (1 channel). + + .. warning:: + This preset is intended for **benchmarking only**. The keypoint-regression CNN + cannot be meaningfully trained from depth alone. Use it with + :class:`ShadowHandCameraBenchmarkEnvCfg` (``feature_extractor.enabled=False``) + to measure pure depth-rendering throughput, e.g.:: + + presets=depth # depth rendering, default renderer + presets=depth,newton_renderer # depth rendering with Newton renderer + presets=depth,ovrtx # depth rendering with OVRTX renderer + """ + + semantic_segmentation: BaseTiledCameraCfg = BaseTiledCameraCfg(data_types=["semantic_segmentation"]) + """Semantic segmentation (3 CNN input channels).""" + + +@configclass +class ShadowHandCameraEnvCfg(ShadowHandEnvCfg): + # scene + scene: InteractiveSceneCfg = InteractiveSceneCfg(num_envs=1225, env_spacing=2.0, replicate_physics=True) + + # camera — data-type and renderer backend selectable via CLI presets + tiled_camera: ShadowHandTiledCameraCfg = ShadowHandTiledCameraCfg() + feature_extractor: FeatureExtractorCfg = FeatureExtractorCfg() + + # env + observation_space = 164 + 27 # state observation + vision CNN embedding + state_space = 187 + 27 # asymmetric states + vision CNN embedding + + def __post_init__(self): + # The vision env renders through the Isaac RTX tiled camera, whose render + # products require the Fabric cloning path. The Newton backend disables Fabric + # cloning (see the base env's ``clone_in_fabric`` scene preset), so under Newton + # the ``rgb`` annotator has no render products for ``num_envs > 1`` and the + # default RGB/depth/semantic render fails. Default the vision env to PhysX so it + # renders out of the box; Newton stays selectable via ``physics=newton_mjwarp`` + # for the depth-only Newton-warp-renderer benchmark path (``presets=newton_renderer``). + super().__post_init__() + for backend_cfg in (self.sim.physics, self.robot_cfg, self.object_cfg): + backend_cfg.default = backend_cfg.physx + + def validate_config(self): + """Check renderer/data-type and feature-extractor compatibility.""" + validate_shadow_hand_camera_settings(self.tiled_camera, self.feature_extractor) + + +@configclass +class ShadowHandCameraEnvPlayCfg(ShadowHandCameraEnvCfg): + # scene + scene: InteractiveSceneCfg = InteractiveSceneCfg( + num_envs=CAMERA_PLAY_NUM_ENVS, env_spacing=2.0, replicate_physics=True + ) + # inference for CNN + feature_extractor: FeatureExtractorCfg = FeatureExtractorCfg(train=False, load_checkpoint=True) + + +@configclass +class ShadowHandCameraBenchmarkEnvCfg(ShadowHandCameraEnvCfg): + """Benchmark configuration with the feature extractor CNN disabled. + + .. deprecated:: 9.0.0 + Use the regular camera task with the ``env.feature_extractor.enabled=False`` + override instead. The ``Isaac-Reorient-Cube-Shadow-Camera-Benchmark-Direct`` + registration will be removed in a future release. + + The tiled camera renders frames each step as normal, but the CNN forward pass is + bypassed — zero embeddings are returned instead. This isolates rendering throughput + from CNN inference overhead when profiling. + + The renderer backend and camera data types can still be selected via ``presets``:: + + presets = newton_renderer # benchmark with Newton renderer + presets = ovrtx # benchmark with OVRTX renderer + presets = rgb # benchmark RGB rendering only + presets = depth, newton_renderer # benchmark depth rendering with Newton + """ + + feature_extractor: FeatureExtractorCfg = FeatureExtractorCfg(enabled=False) diff --git a/source/isaaclab_tasks/isaaclab_tasks/core/reorient/config/shadow_hand/shadow_hand_direct_env_cfg.py b/source/isaaclab_tasks/isaaclab_tasks/core/reorient/config/shadow_hand/shadow_hand_direct_env_cfg.py new file mode 100644 index 000000000000..8a65002bca88 --- /dev/null +++ b/source/isaaclab_tasks/isaaclab_tasks/core/reorient/config/shadow_hand/shadow_hand_direct_env_cfg.py @@ -0,0 +1,143 @@ +# Copyright (c) 2022-2026, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +from isaaclab.envs import DirectRLEnvCfg +from isaaclab.markers import VisualizationMarkersCfg +from isaaclab.scene import InteractiveSceneCfg +from isaaclab.sim import SimulationCfg +from isaaclab.sim.spawners.materials import RigidBodyMaterialBaseCfg +from isaaclab.utils.configclass import configclass +from isaaclab.utils.noise import NoiseModelWithAdditiveBiasCfg + +from isaaclab_tasks.core.reorient.config.shadow_hand.shadow_hand_common import ( + GOAL_OBJECT_CFG, + OBJECT_CFG, + OPENAI_ACTION_NOISE_CFG, + OPENAI_OBSERVATION_NOISE_CFG, + ROBOT_CFG, + ObjectCfg, + PhysicsCfg, + ShadowHandEventCfg, + ShadowHandRobotCfg, +) +from isaaclab_tasks.utils import preset + +from isaaclab_assets.robots.shadow_hand import SHADOW_ACTUATED_JOINT_NAMES, SHADOW_FINGERTIP_BODY_NAMES + + +@configclass +class ShadowHandSceneCfg(InteractiveSceneCfg): + """Shadow Direct scene defaults. + + ``clone_in_fabric`` is the only backend-varying field: PhysX/OvPhysX use Fabric + cloning for speed; Newton does not support it. The per-backend value is selected + inline at the scene field via :func:`~isaaclab_tasks.utils.preset`. + """ + + num_envs = 8192 + env_spacing = 0.75 + replicate_physics = True + + +@configclass +class ShadowHandEnvCfg(DirectRLEnvCfg): + # env + decimation = 2 + episode_length_s = 10.0 + action_space = 20 + observation_space = 157 # (full) + state_space = 0 + asymmetric_obs = False + obs_type = "full" + + # simulation — values mirrored by the manager cfg (guarded by the value-parity test) + sim: SimulationCfg = SimulationCfg( + dt=1 / 120, + render_interval=2, + physics_material=RigidBodyMaterialBaseCfg(static_friction=1.0, dynamic_friction=1.0), + physics=PhysicsCfg(), + ) + + # robot + robot_cfg: ShadowHandRobotCfg = ROBOT_CFG + actuated_joint_names = SHADOW_ACTUATED_JOINT_NAMES + fingertip_body_names = SHADOW_FINGERTIP_BODY_NAMES + + # in-hand object + object_cfg: ObjectCfg = OBJECT_CFG + # goal object + goal_object_cfg: VisualizationMarkersCfg = GOAL_OBJECT_CFG + # scene — clone_in_fabric is the only backend-varying field (Newton cannot use Fabric cloning) + scene: InteractiveSceneCfg = preset( + default=ShadowHandSceneCfg(clone_in_fabric=False), + physx=ShadowHandSceneCfg(clone_in_fabric=True), + ovphysx=ShadowHandSceneCfg(clone_in_fabric=True), + newton_mjwarp=ShadowHandSceneCfg(clone_in_fabric=False), + newton_kamino=ShadowHandSceneCfg(clone_in_fabric=False), + ) + + # reset + reset_position_noise = 0.01 # range of position at reset + reset_dof_pos_noise = 0.2 # range of dof pos at reset + reset_dof_vel_noise = 0.0 # range of dof vel at reset + # reward scales + dist_reward_scale = -10.0 + rot_reward_scale = 1.0 + rot_eps = 0.1 + action_penalty_scale = -0.0002 + reach_goal_bonus = 250.0 + fall_penalty = 0.0 + fall_dist = 0.24 + vel_obs_scale = 0.2 + success_tolerance = 0.1 + max_consecutive_success = 0 + success_count_threshold: int = 1 + """Minimum number of goals reached in an episode to count it as a successful episode.""" + av_factor = 0.1 + act_moving_average = 1.0 + force_torque_obs_scale = 10.0 + + +@configclass +class ShadowHandOpenAIEnvCfg(ShadowHandEnvCfg): + # env + decimation = 3 + episode_length_s = 8.0 + action_space = 20 + observation_space = 42 + state_space = 187 + asymmetric_obs = True + obs_type = "openai" + + # simulation — values mirrored by the manager cfg (guarded by the value-parity test) + sim: SimulationCfg = SimulationCfg( + dt=1 / 60, + render_interval=3, + physics_material=RigidBodyMaterialBaseCfg(static_friction=1.0, dynamic_friction=1.0), + physics=PhysicsCfg(), + ) + # reset + reset_position_noise = 0.01 # range of position at reset + reset_dof_pos_noise = 0.2 # range of dof pos at reset + reset_dof_vel_noise = 0.0 # range of dof vel at reset + # reward scales + dist_reward_scale = -10.0 + rot_reward_scale = 1.0 + rot_eps = 0.1 + action_penalty_scale = -0.0002 + reach_goal_bonus = 250.0 + fall_penalty = -50.0 + vel_obs_scale = 0.2 + success_tolerance = 0.4 + max_consecutive_success = 50 + av_factor = 0.1 + act_moving_average = 0.3 + force_torque_obs_scale = 10.0 + # domain randomization config + events: ShadowHandEventCfg = ShadowHandEventCfg() + # at every time-step add gaussian noise + bias. The bias is a gaussian sampled at reset + action_noise_model: NoiseModelWithAdditiveBiasCfg = OPENAI_ACTION_NOISE_CFG + # at every time-step add gaussian noise + bias. The bias is a gaussian sampled at reset + observation_noise_model: NoiseModelWithAdditiveBiasCfg = OPENAI_OBSERVATION_NOISE_CFG 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..b989bcb0c426 --- /dev/null +++ b/source/isaaclab_tasks/isaaclab_tasks/core/reorient/config/shadow_hand/shadow_hand_manager_env_cfg.py @@ -0,0 +1,226 @@ +# Copyright (c) 2022-2026, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +"""Manager-based counterpart of the state-based Shadow Hand reorientation task.""" + +import isaaclab.sim as sim_utils +from isaaclab.assets import AssetBaseCfg +from isaaclab.envs import ManagerBasedRLEnvCfg +from isaaclab.managers import EventTermCfg as EventTerm +from isaaclab.managers import ObservationGroupCfg as ObsGroup +from isaaclab.managers import ObservationTermCfg as ObsTerm +from isaaclab.managers import RewardTermCfg as RewTerm +from isaaclab.managers import SceneEntityCfg +from isaaclab.managers import TerminationTermCfg as DoneTerm +from isaaclab.scene import InteractiveSceneCfg +from isaaclab.sim.spawners.materials import RigidBodyMaterialBaseCfg +from isaaclab.utils.configclass import configclass + +import isaaclab_tasks.core.reorient.mdp as mdp +from isaaclab_tasks.core.reorient.config.shadow_hand.shadow_hand_common import ( + GOAL_OBJECT_CFG, + OBJECT_CFG, + ROBOT_CFG, + ObjectCfg, + PhysicsCfg, +) +from isaaclab_tasks.core.reorient.reorient_common import GOAL_MARKER_POSITION, IN_HAND_POS_OFFSET +from isaaclab_tasks.utils import PresetCfg + +from isaaclab_assets.robots.shadow_hand import SHADOW_ACTUATED_JOINT_NAMES, SHADOW_FINGERTIP_BODY_NAMES + +# ---------------------------------- state task ---------------------------------- + + +@configclass +class _ShadowHandManagerSceneCfg(InteractiveSceneCfg): + """Scene shared by the Shadow Hand Manager backend alternatives.""" + + num_envs = 8192 + env_spacing = 0.75 + replicate_physics = True + + ground = AssetBaseCfg(prim_path="/World/ground", spawn=sim_utils.GroundPlaneCfg()) + robot: PresetCfg = ROBOT_CFG + object: ObjectCfg = OBJECT_CFG + light = AssetBaseCfg( + prim_path="/World/Light", + spawn=sim_utils.DomeLightCfg(intensity=2000.0, color=(0.75, 0.75, 0.75)), + ) + + +@configclass +class ShadowHandManagerSceneCfg(PresetCfg): + """Backend-specific scene cloning settings matching the Direct task.""" + + physx = _ShadowHandManagerSceneCfg(clone_in_fabric=True) + newton_mjwarp = _ShadowHandManagerSceneCfg(clone_in_fabric=False) + ovphysx = physx + newton_kamino = newton_mjwarp + default = newton_mjwarp + + +@configclass +class CommandsCfg: + """Object pose goal matching the Direct in-hand target.""" + + object_pose = mdp.ReorientEpisodeCommandCfg( + asset_name="object", + init_pos_offset=IN_HAND_POS_OFFSET, + update_goal_on_success=True, + orientation_success_threshold=0.1, + make_quat_unique=False, + fixed_marker_pos=GOAL_MARKER_POSITION, + goal_pose_visualizer_cfg=GOAL_OBJECT_CFG, + debug_vis=True, + ) + + +@configclass +class ActionsCfg: + """Twenty actuated Shadow Hand joints.""" + + joint_pos = mdp.EMAJointPositionToLimitsActionCfg( + asset_name="robot", + joint_names=SHADOW_ACTUATED_JOINT_NAMES, + alpha=1.0, + rescale_to_limits=True, + ) + + +@configclass +class FullStateWithoutActionCfg(ObsGroup): + """Shared first 137 dimensions of the full Shadow state.""" + + joint_pos = ObsTerm( + func=mdp.joint_pos_limit_normalized, + params={"asset_cfg": SceneEntityCfg("robot", joint_names=".*", preserve_order=False)}, + ) + joint_vel = ObsTerm( + func=mdp.joint_vel, + scale=0.2, + params={"asset_cfg": SceneEntityCfg("robot", joint_names=".*", preserve_order=False)}, + ) + object_pos = ObsTerm(func=mdp.root_pos_w, params={"asset_cfg": SceneEntityCfg("object")}) + object_quat = ObsTerm( + func=mdp.root_quat_w, + params={"asset_cfg": SceneEntityCfg("object"), "make_quat_unique": False}, + ) + object_lin_vel = ObsTerm(func=mdp.root_lin_vel_w, params={"asset_cfg": SceneEntityCfg("object")}) + object_ang_vel = ObsTerm( + func=mdp.root_ang_vel_w, + scale=0.2, + params={"asset_cfg": SceneEntityCfg("object")}, + ) + goal_pose = ObsTerm(func=mdp.generated_commands, params={"command_name": "object_pose"}) + goal_quat_diff = ObsTerm( + func=mdp.goal_quat_diff, + params={"asset_cfg": SceneEntityCfg("object"), "command_name": "object_pose", "make_quat_unique": False}, + ) + fingertip_pos = ObsTerm( + func=mdp.fingertip_pos, + params={"asset_cfg": SceneEntityCfg("robot", body_names=SHADOW_FINGERTIP_BODY_NAMES, preserve_order=False)}, + ) + fingertip_quat = ObsTerm( + func=mdp.fingertip_quat, + params={"asset_cfg": SceneEntityCfg("robot", body_names=SHADOW_FINGERTIP_BODY_NAMES, preserve_order=False)}, + ) + fingertip_vel = ObsTerm( + func=mdp.fingertip_vel, + params={"asset_cfg": SceneEntityCfg("robot", body_names=SHADOW_FINGERTIP_BODY_NAMES, preserve_order=False)}, + ) + + def __post_init__(self): + self.enable_corruption = False + self.concatenate_terms = True + + +@configclass +class ObservationsCfg: + """Full 157-dimensional state observation in Direct order.""" + + @configclass + class PolicyCfg(FullStateWithoutActionCfg): + last_action = ObsTerm(func=mdp.reorient_last_action, params={"action_name": "joint_pos"}) + + policy: PolicyCfg = PolicyCfg() + + +@configclass +class EventCfg: + """Reset distributions matching the Direct task.""" + + reset_state = EventTerm( + func=mdp.reset_reorient_state, + mode="reset", + params={ + "position_noise": 0.01, + "joint_position_noise": 0.2, + "joint_velocity_noise": 0.0, + "action_name": "joint_pos", + }, + ) + + +@configclass +class RewardsCfg: + """Direct-compatible reward and success accounting.""" + + reorient = RewTerm( + func=mdp.ReorientReward, + weight=1.0, + params={ + "command_name": "object_pose", + "distance_scale": -10.0, + "rotation_scale": 1.0, + "rotation_epsilon": 0.1, + "action_penalty_scale": -0.0002, + "success_tolerance": 0.1, + "success_bonus": 250.0, + "fall_distance": 0.24, + "fall_penalty": 0.0, + "averaging_factor": 0.1, + "success_count_threshold": 1, + "object_cfg": SceneEntityCfg("object"), + }, + ) + + +@configclass +class TerminationsCfg: + """Termination conditions matching the Direct task.""" + + object_out_of_reach = DoneTerm( + func=mdp.object_reorientation_out_of_reach, + params={ + "threshold": 0.24, + "command_name": "object_pose", + "object_cfg": SceneEntityCfg("object"), + }, + ) + time_out = DoneTerm(func=mdp.time_out, time_out=True) + + +@configclass +class ShadowHandManagerEnvCfg(ManagerBasedRLEnvCfg): + """Manager-based state Shadow Hand task with Direct-compatible semantics.""" + + scene: ShadowHandManagerSceneCfg = ShadowHandManagerSceneCfg() + observations: ObservationsCfg = ObservationsCfg() + actions: ActionsCfg = ActionsCfg() + commands: CommandsCfg = CommandsCfg() + rewards: RewardsCfg = RewardsCfg() + terminations: TerminationsCfg = TerminationsCfg() + events: EventCfg = EventCfg() + + def __post_init__(self): + self.decimation = 2 + self.episode_length_s = 10.0 + # simulation — mirrors the Direct cfg (guarded by the value-parity test) + self.sim.dt = 1 / 120 + self.sim.render_interval = self.decimation + self.sim.physics_material = RigidBodyMaterialBaseCfg(static_friction=1.0, dynamic_friction=1.0) + self.sim.physics = PhysicsCfg() + self.viewer.eye = (2.0, 2.0, 2.0) diff --git a/source/isaaclab_tasks/isaaclab_tasks/core/reorient/config/shadow_hand/shadow_hand_openai_manager_env_cfg.py b/source/isaaclab_tasks/isaaclab_tasks/core/reorient/config/shadow_hand/shadow_hand_openai_manager_env_cfg.py new file mode 100644 index 000000000000..7e63bab8cc50 --- /dev/null +++ b/source/isaaclab_tasks/isaaclab_tasks/core/reorient/config/shadow_hand/shadow_hand_openai_manager_env_cfg.py @@ -0,0 +1,230 @@ +# 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 OpenAI Shadow Hand reorientation variants (FF and LSTM).""" + +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.sensors import JointWrenchSensorCfg +from isaaclab.sim.spawners.materials import RigidBodyMaterialBaseCfg +from isaaclab.utils.configclass import configclass + +import isaaclab_tasks.core.reorient.mdp as mdp +from isaaclab_tasks.core.reorient.config.shadow_hand.shadow_hand_common import ( + GOAL_OBJECT_CFG, + OPENAI_ACTION_NOISE_CFG, + OPENAI_OBSERVATION_NOISE_CFG, + NewtonEventCfg, + PhysicsCfg, + PhysxEventCfg, +) +from isaaclab_tasks.core.reorient.config.shadow_hand.shadow_hand_manager_env_cfg import ( + FullStateWithoutActionCfg, + _ShadowHandManagerSceneCfg, +) +from isaaclab_tasks.core.reorient.reorient_common import GOAL_MARKER_POSITION, IN_HAND_POS_OFFSET +from isaaclab_tasks.utils import PresetCfg + +from isaaclab_assets.robots.shadow_hand import SHADOW_ACTUATED_JOINT_NAMES, SHADOW_FINGERTIP_BODY_NAMES + + +@configclass +class OpenAICommandsCfg: + """OpenAI goal command with its wider success tolerance.""" + + object_pose = mdp.ReorientEpisodeCommandCfg( + asset_name="object", + init_pos_offset=IN_HAND_POS_OFFSET, + update_goal_on_success=True, + orientation_success_threshold=0.4, + make_quat_unique=False, + fixed_marker_pos=GOAL_MARKER_POSITION, + goal_pose_visualizer_cfg=GOAL_OBJECT_CFG, + debug_vis=True, + ) + + +@configclass +class OpenAIActionsCfg: + """OpenAI actions with Direct-compatible EMA and stateful noise.""" + + joint_pos = mdp.NoisyEMAJointPositionToLimitsActionCfg( + asset_name="robot", + joint_names=SHADOW_ACTUATED_JOINT_NAMES, + alpha=0.3, + rescale_to_limits=True, + noise_model=OPENAI_ACTION_NOISE_CFG, + ) + + +@configclass +class OpenAIObservationsCfg: + """OpenAI 42-dimensional actor and 187-dimensional critic observations.""" + + @configclass + class PolicyCfg(ObsGroup): + openai = ObsTerm( + func=mdp.OpenAIPolicyObservation, + params={ + "command_name": "object_pose", + "action_name": "joint_pos", + "noise_model": OPENAI_OBSERVATION_NOISE_CFG, + "robot_cfg": SceneEntityCfg("robot", body_names=SHADOW_FINGERTIP_BODY_NAMES, preserve_order=False), + "object_cfg": SceneEntityCfg("object"), + }, + ) + + def __post_init__(self): + self.enable_corruption = False + self.concatenate_terms = True + + @configclass + class CriticCfg(FullStateWithoutActionCfg): + fingertip_wrench = ObsTerm( + func=mdp.fingertip_wrench, + scale=10.0, + params={ + "sensor_cfg": SceneEntityCfg( + "joint_wrench", body_names=SHADOW_FINGERTIP_BODY_NAMES, preserve_order=False + ) + }, + ) + last_action = ObsTerm(func=mdp.reorient_last_action, params={"action_name": "joint_pos"}) + + policy: PolicyCfg = PolicyCfg() + critic: CriticCfg = CriticCfg() + + +@configclass +class ShadowHandOpenAIManagerSceneCfg(PresetCfg): + """Backend-specific OpenAI scene alternatives.""" + + @configclass + class SceneCfg(_ShadowHandManagerSceneCfg): + """Shadow Hand scene with fingertip joint-wrench sensing.""" + + joint_wrench = JointWrenchSensorCfg(prim_path="{ENV_REGEX_NS}/Robot") + + physx = SceneCfg(clone_in_fabric=True) + newton_mjwarp = SceneCfg(clone_in_fabric=False) + ovphysx = physx + newton_kamino = newton_mjwarp + default = newton_mjwarp + + +_OPENAI_RESET_PARAMS = { + "position_noise": 0.01, + "joint_position_noise": 0.2, + "joint_velocity_noise": 0.0, + "action_name": "joint_pos", +} + + +@configclass +class OpenAIPhysxEventCfg(PhysxEventCfg): + """PhysX OpenAI randomization and state reset events.""" + + reset_state = EventTerm(func=mdp.reset_reorient_state, mode="reset", params=_OPENAI_RESET_PARAMS) + + +@configclass +class OpenAINewtonEventCfg(NewtonEventCfg): + """Newton OpenAI randomization and state reset events.""" + + reset_state = EventTerm(func=mdp.reset_reorient_state, mode="reset", params=_OPENAI_RESET_PARAMS) + + +@configclass +class OpenAIEventCfg(PresetCfg): + """Backend-specific OpenAI event alternatives.""" + + physx = OpenAIPhysxEventCfg() + newton_mjwarp = OpenAINewtonEventCfg() + ovphysx = physx + newton_kamino = newton_mjwarp + default = newton_mjwarp + + +@configclass +class OpenAIRewardsCfg: + """Direct-compatible OpenAI reward and success accounting.""" + + reorient = RewTerm( + func=mdp.ReorientReward, + weight=1.0, + params={ + "command_name": "object_pose", + "distance_scale": -10.0, + "rotation_scale": 1.0, + "rotation_epsilon": 0.1, + "action_penalty_scale": -0.0002, + "success_tolerance": 0.4, + "success_bonus": 250.0, + "fall_distance": 0.24, + "fall_penalty": -50.0, + "averaging_factor": 0.1, + "success_count_threshold": 1, + "action_name": "joint_pos", + "object_cfg": SceneEntityCfg("object"), + }, + ) + + +@configclass +class OpenAITerminationsCfg: + """Direct-compatible OpenAI termination conditions.""" + + object_out_of_reach = DoneTerm( + func=mdp.object_reorientation_out_of_reach, + params={ + "threshold": 0.24, + "command_name": "object_pose", + "object_cfg": SceneEntityCfg("object"), + }, + ) + time_out = DoneTerm( + func=mdp.ReorientTimeout, + time_out=True, + params={ + "command_name": "object_pose", + "reward_name": "reorient", + "success_tolerance": 0.4, + "max_successes": 50, + "object_cfg": SceneEntityCfg("object"), + }, + ) + + +@configclass +class ShadowHandOpenAIManagerEnvCfg(ManagerBasedRLEnvCfg): + """Manager counterpart shared by the OpenAI FF and LSTM variants. + + Standalone rather than a subclass of :class:`ShadowHandManagerEnvCfg`: + every section differs from the state task, so this block is the complete + recipe. + """ + + 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): + self.decimation = 3 + self.episode_length_s = 8.0 + # simulation — mirrors the Direct cfg (guarded by the value-parity test) + self.sim.dt = 1 / 60 + self.sim.render_interval = self.decimation + self.sim.physics_material = RigidBodyMaterialBaseCfg(static_friction=1.0, dynamic_friction=1.0) + self.sim.physics = PhysicsCfg() + self.viewer.eye = (2.0, 2.0, 2.0) 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..4f4750cbb049 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,55 @@ # SPDX-License-Identifier: BSD-3-Clause __all__ = [ + "NoisyEMAJointPositionToLimitsAction", + "NoisyEMAJointPositionToLimitsActionCfg", + "ShadowHandCameraFeatures", + "shadow_hand_camera_cached_features", + "shadow_hand_goal_keypoints", "ReorientCommand", "ReorientCommandCfg", + "ReorientEpisodeCommand", + "ReorientEpisodeCommandCfg", + "reset_reorient_state", + "fingertip_pos", + "fingertip_quat", + "fingertip_vel", + "fingertip_wrench", + "reorient_last_action", + "OpenAIPolicyObservation", "goal_quat_diff", - "success_bonus", - "track_orientation_inv_l2", - "track_pos_l2", - "max_consecutive_success", - "object_away_from_goal", - "object_away_from_robot", + "evaluate_reorient_success", + "reorient_reward", + "ReorientReward", + "object_reorientation_out_of_reach", + "ReorientTimeout", ] -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 ( + ShadowHandCameraFeatures, + shadow_hand_camera_cached_features, + shadow_hand_goal_keypoints, + OpenAIPolicyObservation, + fingertip_pos, + fingertip_quat, + fingertip_vel, + fingertip_wrench, + goal_quat_diff, + reorient_last_action, +) +from .rewards import ( + ReorientReward, + reorient_reward, + evaluate_reorient_success, +) +from .terminations import ( + ReorientTimeout, + object_reorientation_out_of_reach, +) from isaaclab.envs.mdp import * diff --git a/source/isaaclab_tasks/isaaclab_tasks/core/reorient/mdp/actions.py b/source/isaaclab_tasks/isaaclab_tasks/core/reorient/mdp/actions.py new file mode 100644 index 000000000000..ee9fbcea3d78 --- /dev/null +++ b/source/isaaclab_tasks/isaaclab_tasks/core/reorient/mdp/actions.py @@ -0,0 +1,63 @@ +# Copyright (c) 2022-2026, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +"""Action terms for the reorientation task family.""" + +from __future__ import annotations + +from collections.abc import Sequence +from dataclasses import MISSING +from typing import TYPE_CHECKING + +import torch + +from isaaclab.envs.mdp import EMAJointPositionToLimitsActionCfg +from isaaclab.envs.mdp.actions import EMAJointPositionToLimitsAction +from isaaclab.utils.configclass import configclass +from isaaclab.utils.noise import NoiseModelCfg + +if TYPE_CHECKING: + from isaaclab.envs import ManagerBasedEnv + + +class NoisyEMAJointPositionToLimitsAction(EMAJointPositionToLimitsAction): + """Apply a stateful noise model before EMA joint-position processing.""" + + def __init__(self, cfg: NoisyEMAJointPositionToLimitsActionCfg, env: ManagerBasedEnv): + """Initialize the noisy action term. + + Args: + cfg: Action configuration including the stateful noise model. + env: Manager-based environment containing the hand. + """ + super().__init__(cfg, env) + self._noise_model = cfg.noise_model.class_type(cfg.noise_model, num_envs=self.num_envs, device=self.device) + + def process_actions(self, actions: torch.Tensor) -> None: + """Apply noise to normalized actions before scaling and EMA filtering. + + Args: + actions: Normalized joint actions, shape ``(num_envs, num_actions)``. + """ + super().process_actions(self._noise_model(actions)) + + def reset(self, env_ids: Sequence[int] | None = None) -> None: + """Reset the noise state and standard EMA action buffers. + + Args: + env_ids: Environment indices to reset, or ``None`` for every environment. + """ + self._noise_model.reset(env_ids) + super().reset(env_ids) + + +@configclass +class NoisyEMAJointPositionToLimitsActionCfg(EMAJointPositionToLimitsActionCfg): + """EMA joint action configuration with Direct-compatible stateful noise.""" + + class_type = NoisyEMAJointPositionToLimitsAction + + noise_model: NoiseModelCfg = MISSING + """Stateful noise applied to incoming normalized actions.""" diff --git a/source/isaaclab_tasks/isaaclab_tasks/core/reorient/mdp/commands.py b/source/isaaclab_tasks/isaaclab_tasks/core/reorient/mdp/commands.py index 77040f635282..d7e05803945d 100644 --- a/source/isaaclab_tasks/isaaclab_tasks/core/reorient/mdp/commands.py +++ b/source/isaaclab_tasks/isaaclab_tasks/core/reorient/mdp/commands.py @@ -21,6 +21,8 @@ from isaaclab.utils.configclass import configclass from isaaclab.utils.leapp import POSE7_ELEMENT_NAMES +from .rewards import evaluate_reorient_success + if TYPE_CHECKING: from isaaclab.assets import RigidObject from isaaclab.envs import ManagerBasedRLEnv @@ -66,6 +68,10 @@ def __init__(self, cfg: ReorientCommandCfg, env: ManagerBasedRLEnv): # -- orientation: (x, y, z, w) self.quat_command_w = torch.zeros(self.num_envs, 4, device=self.device) self.quat_command_w[:, 3] = 1.0 # set the scalar component to 1.0 + # persistent (num_envs, 7) pose command: the position half is static and written once + # here, the quaternion half is refreshed by _resample_command; `command` returns this + # buffer directly instead of allocating a torch.cat every call + self._command_buf = torch.cat((self.pos_command_e, self.quat_command_w), dim=-1) # -- unit vectors self._X_UNIT_VEC = torch.tensor([1.0, 0, 0], device=self.device).repeat((self.num_envs, 1)) @@ -80,6 +86,8 @@ def __init__(self, cfg: ReorientCommandCfg, env: ManagerBasedRLEnv): # -- per-attempt success accounting: each success-driven resample completes one attempt; # the trailing attempt at episode end counts as one unsuccessful attempt. self._completed_attempts = torch.zeros(self.num_envs, device=self.device) + # goal-marker position with the configured offset; built lazily on first render + self._marker_pos_w: torch.Tensor | None = None # adds (optional) cmd kind and element names for leapp export # during export, semantic data about this command will be used to annotate the command input @@ -97,26 +105,28 @@ def __str__(self) -> str: @property def command(self) -> torch.Tensor: - """The desired goal pose in the environment frame. Shape is (num_envs, 7).""" - return torch.cat((self.pos_command_e, self.quat_command_w), dim=-1) + """The desired goal pose in the environment frame. Shape is (num_envs, 7). + + The returned tensor is a persistent buffer refreshed in place; consumers that + store it across steps must copy it. + """ + return self._command_buf """ Implementation specific functions. """ def _update_metrics(self): - # logs data - # -- compute the orientation error - self.metrics["orientation_error"] = math_utils.quat_error_magnitude( - self.object.data.root_quat_w.torch, self.quat_command_w + success_flags, orientation_error = evaluate_reorient_success( + self.object.data.root_quat_w.torch, self.quat_command_w, self.cfg.orientation_success_threshold ) - # -- compute the position error - self.metrics["position_error"] = torch.linalg.norm( - self.object.data.root_pos_w.torch - self.pos_command_w, dim=1 + # write the stable metric buffers in place; the manager holds references to them + self.metrics["orientation_error"][:] = orientation_error + self.metrics["position_error"][:] = torch.linalg.norm( + self.object.data.root_pos_w.torch - self.pos_command_w, ord=2, dim=-1 ) - # -- compute the number of consecutive successes - successes = self.metrics["orientation_error"] < self.cfg.orientation_success_threshold - self.metrics["consecutive_success"] += successes.float() + # bool flags promote to the metric's float dtype; add_ avoids the .float() temporary + self.metrics["consecutive_success"].add_(success_flags) def reset(self, env_ids: Sequence[int] | None = None) -> dict[str, float]: # Snapshot per-attempt success rate BEFORE the base class logs and zeros metrics. @@ -148,6 +158,8 @@ def _resample_command(self, env_ids: Sequence[int]): ) # make sure the quaternion real-part is always positive self.quat_command_w[env_ids] = math_utils.quat_unique(quat) if self.cfg.make_quat_unique else quat + # keep the persistent pose-command buffer current (position half is static) + self._command_buf[env_ids, 3:] = self.quat_command_w[env_ids] def _update_command(self): # update the command if goal is reached @@ -173,10 +185,11 @@ def _set_debug_vis_impl(self, debug_vis: bool): def _debug_vis_callback(self, event): # add an offset to the marker position to visualize the goal - marker_pos = self.pos_command_w + torch.tensor(self.cfg.marker_pos_offset, device=self.device) - marker_quat = self.quat_command_w + if self._marker_pos_w is None: + # constant per run; cached to avoid a host-to-device allocation every render frame + self._marker_pos_w = self.pos_command_w + torch.tensor(self.cfg.marker_pos_offset, device=self.device) # visualize the goal marker - self.goal_pose_visualizer.visualize(translations=marker_pos, orientations=marker_quat) + self.goal_pose_visualizer.visualize(translations=self._marker_pos_w, orientations=self.quat_command_w) @configclass @@ -209,7 +222,10 @@ class ReorientCommandCfg(CommandTermCfg): """ orientation_success_threshold: float = MISSING - """Threshold for the orientation error to consider the goal orientation to be reached.""" + """Threshold [rad] for the orientation error to consider the goal orientation to be reached. + + Set per family at the declaration site, matching the Direct configuration's value. + """ update_goal_on_success: bool = MISSING """Whether to update the goal orientation when the goal orientation is reached.""" @@ -231,3 +247,68 @@ class ReorientCommandCfg(CommandTermCfg): }, ) """The configuration for the goal pose visualization marker. Defaults to a DexCube marker.""" + + +class ReorientEpisodeCommand(ReorientCommand): + """Reorientation command whose success metric is owned by an episode reward term. + + This variant retains success-triggered goal resampling while suppressing the + generic command's per-attempt ``Metrics/success_rate`` value. + """ + + cfg: ReorientEpisodeCommandCfg + + def __init__(self, cfg: ReorientEpisodeCommandCfg, env: ManagerBasedRLEnv): + super().__init__(cfg, env) + self._skip_success_update = torch.zeros(self.num_envs, dtype=torch.bool, device=self.device) + self._fixed_marker_pos_w: torch.Tensor | None = None + + def _update_command(self): + if self.cfg.update_goal_on_success: + goal_reset_ids = ( + ( + (self.metrics["orientation_error"] <= self.cfg.orientation_success_threshold) + & ~self._skip_success_update + ) + .nonzero(as_tuple=False) + .squeeze(-1) + ) + self._resample(goal_reset_ids) + self._skip_success_update[:] = False + + def _debug_vis_callback(self, event): + if self.cfg.fixed_marker_pos is None: + super()._debug_vis_callback(event) + return + if self._fixed_marker_pos_w is None: + # constant per run; cached to avoid a host-to-device allocation every render frame + self._fixed_marker_pos_w = ( + torch.tensor(self.cfg.fixed_marker_pos, device=self.device).repeat(self.num_envs, 1) + + self._env.scene.env_origins + ) + self.goal_pose_visualizer.visualize(translations=self._fixed_marker_pos_w, orientations=self.quat_command_w) + + def reset(self, env_ids: Sequence[int] | None = None) -> dict[str, float]: + if env_ids is None: + env_ids = slice(None) + extras = CommandTerm.reset(self, env_ids) + self._completed_attempts[env_ids] = 0.0 + # Auto-reset happens immediately before CommandManager.compute(). Skip + # success handling for those IDs until one new physics/reward step has run. + reset_buf = getattr(self._env, "reset_buf", None) + if reset_buf is None: + self._skip_success_update[env_ids] = False + else: + self._skip_success_update[env_ids] = reset_buf[env_ids] + extras.pop("success_rate", None) + return extras + + +@configclass +class ReorientEpisodeCommandCfg(ReorientCommandCfg): + """Configuration for episode-accounted reorientation commands.""" + + class_type: type[ReorientEpisodeCommand] = ReorientEpisodeCommand + + fixed_marker_pos: tuple[float, float, float] | None = None + """Fixed goal-marker position [m] in each environment, or ``None`` to use the command position.""" diff --git a/source/isaaclab_tasks/isaaclab_tasks/core/reorient/mdp/events.py b/source/isaaclab_tasks/isaaclab_tasks/core/reorient/mdp/events.py new file mode 100644 index 000000000000..29ea93e5967e --- /dev/null +++ b/source/isaaclab_tasks/isaaclab_tasks/core/reorient/mdp/events.py @@ -0,0 +1,131 @@ +# 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 + +"""Reset events for state-based in-hand reorientation tasks.""" + +from __future__ import annotations + +from collections.abc import Sequence +from typing import TYPE_CHECKING + +import numpy as np +import torch + +import isaaclab.utils.math as math_utils +from isaaclab.managers import SceneEntityCfg +from isaaclab.utils.math import quat_from_angle_axis, quat_mul + +if TYPE_CHECKING: + from isaaclab.assets import Articulation, RigidObject + from isaaclab.envs import ManagerBasedRLEnv + + +def sample_joint_positions_within_limits( + default_position: torch.Tensor, + limits: torch.Tensor, + noise_scale: float, +) -> torch.Tensor: + """Sample reset positions between each joint's default position and limits. + + Args: + default_position: Default joint positions [m or rad, depending on joint type], shape ``(..., J)``. + limits: Lower and upper joint-position limits [m or rad, depending on joint type], shape ``(..., J, 2)``. + noise_scale: Dimensionless interpolation scale from the default position toward the sampled limits. + + Returns: + Sampled joint positions [m or rad, depending on joint type], shape ``(..., J)``. + + Raises: + ValueError: If :paramref:`noise_scale` is outside ``[0, 1]``. + """ + if not 0.0 <= noise_scale <= 1.0: + raise ValueError(f"Expected noise_scale in [0, 1], got {noise_scale}.") + position_sample = math_utils.sample_uniform( + -1.0, + 1.0, + default_position.shape, + device=default_position.device, + ) + position_fraction = 0.5 * (position_sample + 1.0) + position_delta = limits[..., 0] - default_position + position_delta = position_delta + (limits[..., 1] - limits[..., 0]) * position_fraction + joint_position = default_position + noise_scale * position_delta + return torch.clamp(joint_position, min=limits[..., 0], max=limits[..., 1]) + + +def random_xy_rotation(count: int, device: str | torch.device) -> torch.Tensor: + """Sample the Direct tasks' sequential random X/Y rotation. + + Args: + count: Number of rotations to sample. + device: Device on which to sample. + + Returns: + Sampled ``(x, y, z, w)`` unit quaternions, shape ``(count, 4)``. + """ + random_values = math_utils.sample_uniform(-1.0, 1.0, (count, 2), device=device) + x_unit = torch.tensor([1.0, 0.0, 0.0], device=device).repeat(count, 1) + y_unit = torch.tensor([0.0, 1.0, 0.0], device=device).repeat(count, 1) + return math_utils.quat_mul( + math_utils.quat_from_angle_axis(random_values[:, 0] * torch.pi, x_unit), + math_utils.quat_from_angle_axis(random_values[:, 1] * torch.pi, y_unit), + ) + + +@torch.jit.script +def randomize_rotation(rand0, rand1, x_unit_tensor, y_unit_tensor): + """Compose ``[-pi, pi]``-scaled random X- and Y-axis rotations into ``(x, y, z, w)`` quaternions.""" + return quat_mul( + quat_from_angle_axis(rand0 * np.pi, x_unit_tensor), quat_from_angle_axis(rand1 * np.pi, y_unit_tensor) + ) + + +def reset_reorient_state( + env: ManagerBasedRLEnv, + env_ids: Sequence[int], + position_noise: float, + joint_position_noise: float, + joint_velocity_noise: float, + action_name: str, + robot_cfg: SceneEntityCfg = SceneEntityCfg("robot"), + object_cfg: SceneEntityCfg = SceneEntityCfg("object"), +) -> None: + """Reset the object and hand with the Direct task's distributions. + + Args: + env: Environment containing the robot and object. + env_ids: Environment indices to reset. + position_noise: Object-position noise half-width [m]. + joint_position_noise: Scale applied to sampled joint-position deltas. + joint_velocity_noise: Joint-velocity noise half-width [rad/s]. + action_name: Action term whose terminal raw action is retained in the reset observation. + robot_cfg: Robot scene entity. + object_cfg: Object scene entity. + """ + raw_action = env.action_manager.get_term(action_name).raw_actions + if not hasattr(env, "_reorient_reset_action"): + env._reorient_reset_action = torch.zeros_like(raw_action) + env._reorient_reset_step = torch.full((env.num_envs,), -1, dtype=torch.long, device=raw_action.device) + env._reorient_reset_action[env_ids] = raw_action[env_ids] + env._reorient_reset_step[env_ids] = env.common_step_counter + + object_asset: Articulation | RigidObject = env.scene[object_cfg.name] + object_pose = object_asset.data.default_root_pose.torch[env_ids].clone() + object_velocity = torch.zeros_like(object_asset.data.default_root_vel.torch[env_ids]) + position_delta = math_utils.sample_uniform(-1.0, 1.0, (len(env_ids), 3), device=env.device) + object_pose[:, :3] += position_noise * position_delta + env.scene.env_origins[env_ids] + object_pose[:, 3:7] = random_xy_rotation(len(env_ids), env.device) + object_asset.write_root_pose_to_sim_index(root_pose=object_pose, env_ids=env_ids) + object_asset.write_root_velocity_to_sim_index(root_velocity=object_velocity, env_ids=env_ids) + + robot: Articulation = env.scene[robot_cfg.name] + default_position = robot.data.default_joint_pos.torch[env_ids] + limits = robot.data.joint_limits.torch[env_ids] + joint_position = sample_joint_positions_within_limits(default_position, limits, joint_position_noise) + velocity_sample = math_utils.sample_uniform(-1.0, 1.0, (len(env_ids), robot.num_joints), device=env.device) + joint_velocity = robot.data.default_joint_vel.torch[env_ids] + joint_velocity_noise * velocity_sample + robot.set_joint_position_target_index(target=joint_position, env_ids=env_ids) + robot.write_joint_position_to_sim_index(position=joint_position, env_ids=env_ids) + robot.write_joint_velocity_to_sim_index(velocity=joint_velocity, env_ids=env_ids) 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..d93013e0351d 100644 --- a/source/isaaclab_tasks/isaaclab_tasks/core/reorient/mdp/observations.py +++ b/source/isaaclab_tasks/isaaclab_tasks/core/reorient/mdp/observations.py @@ -7,36 +7,342 @@ 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 Camera, JointWrenchSensor + + from isaaclab_tasks.core.reorient.config.shadow_hand.feature_extractor import FeatureExtractorCfg from .commands import ReorientCommand +CUBE_HALF_SIZE: tuple[float, float, float] = (0.03, 0.03, 0.03) +"""Half side lengths [m] of the reorientation cube.""" + + +def _cube_corner_offsets( + size: tuple[float, float, float], num_keypoints: int, device: torch.device | str +) -> torch.Tensor: + """Corner offsets [m] from the cube center; corner index bits select the +/- half side per axis.""" + signs = torch.tensor( + [[1 - 2 * ((corner >> axis) & 1) for axis in range(3)] for corner in range(num_keypoints)], + dtype=torch.float32, + device=device, + ) + half_size = torch.tensor(size, dtype=torch.float32, device=device) / 2.0 + return signs * half_size + + +def compute_cube_keypoints( + pose: torch.Tensor, + num_keypoints: int = 8, + size: tuple[float, float, float] = (2 * 0.03, 2 * 0.03, 2 * 0.03), + out: torch.Tensor | None = None, +) -> torch.Tensor: + """Compute cube-corner positions for batched poses. + + Args: + pose: Cube center poses ``(x, y, z, qx, qy, qz, qw)`` [m, unit quaternion]. + num_keypoints: Number of binary-sign corners to compute. + size: Cube side lengths along each axis [m]. + out: Optional output buffer [m], shape ``(num_envs, num_keypoints, 3)``. + + Returns: + Cube-corner positions [m], shape ``(num_envs, num_keypoints, 3)``. + """ + num_envs = pose.shape[0] + corners = _cube_corner_offsets(size, num_keypoints, pose.device) + rotated = math_utils.quat_apply( + pose[:, None, 3:7].expand(num_envs, num_keypoints, 4), corners.expand(num_envs, num_keypoints, 3) + ) + keypoints = pose[:, None, 0:3] + rotated + if out is None: + return keypoints + out.copy_(keypoints) + return out + + +def cube_keypoints_from_quat( + quat: torch.Tensor, + half_size: tuple[float, float, float] = CUBE_HALF_SIZE, + num_keypoints: int = 8, +) -> torch.Tensor: + """Rotation-only cube-corner offsets [m] from batched ``(x, y, z, w)`` orientations. + + Args: + quat: Cube orientations, shape ``(num_envs, 4)``. + half_size: Cube half side lengths along each axis [m]. + num_keypoints: Number of binary-sign corners to compute. + + Returns: + Flattened corner offsets [m], shape ``(num_envs, num_keypoints * 3)``. + """ + num_envs = quat.shape[0] + size = (2.0 * half_size[0], 2.0 * half_size[1], 2.0 * half_size[2]) + corners = _cube_corner_offsets(size, num_keypoints, quat.device) + rotated = math_utils.quat_apply( + quat[:, None, :].expand(num_envs, num_keypoints, 4), corners.expand(num_envs, num_keypoints, 3) + ) + return rotated.reshape(num_envs, num_keypoints * 3) + + def goal_quat_diff( env: ManagerBasedRLEnv, asset_cfg: SceneEntityCfg, command_name: str, make_quat_unique: bool ) -> 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 real part is always positive when ``make_quat_unique`` is set. + + Args: + env: The environment object. + asset_cfg: The scene entity whose root orientation is compared. + command_name: The command term to be used for extracting the goal. + make_quat_unique: Whether to keep the quaternion real part non-negative. + + Returns: + Per-environment quaternion error ``asset * conjugate(goal)`` in ``(x, y, z, w)`` order. """ - # extract useful elements asset: RigidObject = env.scene[asset_cfg.name] command_term: ReorientCommand = env.command_manager.get_term(command_name) + quat_error = math_utils.quat_mul( + asset.data.root_quat_w.torch, math_utils.quat_conjugate(command_term.quat_command_w) + ) + return math_utils.quat_unique(quat_error) if make_quat_unique else quat_error + + +def fingertip_pos(env: ManagerBasedRLEnv, asset_cfg: SceneEntityCfg) -> torch.Tensor: + """Flattened fingertip positions in the environment frame [m], shape ``(num_envs, num_fingertips * 3)``.""" + asset = env.scene[asset_cfg.name] + positions = asset.data.body_pos_w.torch[:, asset_cfg.body_ids] - env.scene.env_origins.unsqueeze(1) + return positions.reshape(env.num_envs, -1) + + +def fingertip_quat(env: ManagerBasedRLEnv, asset_cfg: SceneEntityCfg) -> torch.Tensor: + """Flattened fingertip ``(x, y, z, w)`` orientations, shape ``(num_envs, num_fingertips * 4)``.""" + asset = env.scene[asset_cfg.name] + return asset.data.body_quat_w.torch[:, asset_cfg.body_ids].reshape(env.num_envs, -1) + + +def fingertip_vel(env: ManagerBasedRLEnv, asset_cfg: SceneEntityCfg) -> torch.Tensor: + """Flattened fingertip 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].reshape(env.num_envs, -1) + + +class fingertip_wrench(ManagerTermBase): + """Fingertip reaction wrenches [N, N·m] with Direct-compatible zero fallback.""" + + def __init__(self, cfg: ObservationTermCfg, env: ManagerBasedRLEnv): + super().__init__(cfg, env) + body_ids = cfg.params["sensor_cfg"].body_ids + # Direct-compatible fallback: report zero wrenches until the sensor produces data + self._zeros = torch.zeros(env.num_envs, len(body_ids) * 6, dtype=torch.float32, device=env.device) + + 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: + return self._zeros + force = force_data.torch[:, sensor_cfg.body_ids] + torque = torque_data.torch[:, sensor_cfg.body_ids] + return torch.cat((force, torque), dim=-1).reshape(env.num_envs, -1) + - # obtain the orientations - goal_quat_w = command_term.command[:, 3:7] - asset_quat_w = asset.data.root_quat_w.torch +def reorient_last_action(env: ManagerBasedRLEnv, action_name: str) -> torch.Tensor: + """Return the Direct-compatible last action across same-step autoreset. - # 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 + Args: + env: Environment containing the action term and reset buffers. + action_name: Action term whose raw action is observed. + + Returns: + Raw actions, retaining each terminal action in its same-step reset observation. + """ + raw_action = env.action_manager.get_term(action_name).raw_actions + reset_action = getattr(env, "_reorient_reset_action", None) + reset_step = getattr(env, "_reorient_reset_step", None) + common_step_counter = getattr(env, "common_step_counter", None) + if reset_action is None or reset_step is None or common_step_counter is None: + return raw_action + return torch.where((reset_step == common_step_counter).unsqueeze(-1), reset_action, raw_action) + + +class OpenAIPolicyObservation(ManagerTermBase): + """Apply one stateful noise model to the concatenated OpenAI actor observation.""" + + def __init__(self, cfg: ObservationTermCfg, env: ManagerBasedRLEnv): + super().__init__(cfg, env) + noise_model: NoiseModelCfg = cfg.params["noise_model"] + self._noise_model = noise_model.class_type(noise_model, num_envs=self.num_envs, device=self.device) + # ObservationManager probes callable terms once for their shape and then + # calls reset. Keep that probe side-effect free so initialization matches + # DirectRLEnv's first noise-model reset and application. + self._shape_probe_pending = True + + def reset(self, env_ids: Sequence[int] | None = None) -> None: + """Reset the actor observation bias for selected environments. + + Args: + env_ids: Environment indices to reset, or ``None`` for every environment. + """ + if self._shape_probe_pending: + self._shape_probe_pending = False + return + self._noise_model.reset(env_ids) + + def __call__( + self, + env: ManagerBasedRLEnv, + command_name: str, + action_name: str, + noise_model: NoiseModelCfg, + robot_cfg: SceneEntityCfg, + object_cfg: SceneEntityCfg, + ) -> torch.Tensor: + """Return the corrupted 42-dimensional actor observation.""" + del noise_model + object_asset: RigidObject = env.scene[object_cfg.name] + object_pos = object_asset.data.root_pos_w.torch - env.scene.env_origins + command_term: ReorientCommand = env.command_manager.get_term(command_name) + quat_error = math_utils.quat_mul( + object_asset.data.root_quat_w.torch, math_utils.quat_conjugate(command_term.quat_command_w) + ) + fingertips = fingertip_pos(env, robot_cfg) + # Direct actor-observation order: fingertips, object position, goal quat error, last action + observation = torch.cat( + (fingertips, object_pos, quat_error, reorient_last_action(env, action_name)), + dim=-1, + ) + if self._shape_probe_pending: + return observation + return self._noise_model(observation) + + +# --------------------------------------------------------------------------- +# Shadow Hand camera observation terms. +# +# These terms wrap the CNN feature pipeline defined in the shadow-hand config +# package. The config layer imports the mdp layer, so the FeatureExtractor +# machinery is imported lazily at term construction time. +# --------------------------------------------------------------------------- + + +class ShadowHandCameraFeatures(ManagerTermBase): + """Run the Direct camera feature pipeline as one Manager observation term.""" + + def __init__(self, cfg: ObservationTermCfg, env: ManagerBasedRLEnv): + super().__init__(cfg, env) + sensor_cfg: SceneEntityCfg = cfg.params["sensor_cfg"] + camera: Camera = env.scene.sensors[sensor_cfg.name] + # Runtime-only import: the mdp layer must not import the task-config layer + # at module load (config modules import mdp; see the layering note above). + from isaaclab_tasks.core.reorient.config.shadow_hand.feature_extractor import FeatureExtractor + + feature_extractor_cfg: FeatureExtractorCfg = env.cfg.feature_extractor + self._feature_extractor = FeatureExtractor( + feature_extractor_cfg, + env.device, + camera.cfg.data_types, + env.cfg.log_dir, + height=camera.cfg.height, + width=camera.cfg.width, + ) + # ObservationManager calls terms once to infer their shape. Do not train + # or save a CNN checkpoint during that initialization probe. + self._shape_probe_pending = True + self._keypoints_buf = torch.empty(env.num_envs, 8, 3, dtype=torch.float32, device=env.device) + + def reset(self, env_ids: Sequence[int] | None = None) -> None: + """Finish the shape-probe phase on the first Manager reset. + + Args: + env_ids: Environment indices being reset. The feature extractor + has no per-environment state, so the indices are unused. + """ + del env_ids + if self._shape_probe_pending: + self._shape_probe_pending = False + + def __call__( + self, + env: ManagerBasedRLEnv, + feature_extractor_cfg: FeatureExtractorCfg, + sensor_cfg: SceneEntityCfg, + object_cfg: SceneEntityCfg, + ) -> torch.Tensor: + """Return the detached 27-dimensional cube-pose embedding. + + Args: + env: Environment containing the object and tiled camera. + feature_extractor_cfg: Feature-extractor configuration captured by + the observation term. The initialized extractor owns its copy. + sensor_cfg: Tiled-camera scene entity. + object_cfg: Reoriented-object scene entity. + + Returns: + Predicted object position and cube keypoints [m], shape + ``(num_envs, 27)``. + """ + del feature_extractor_cfg + if self._shape_probe_pending: + embeddings = torch.zeros(env.num_envs, 27, dtype=torch.float32, device=env.device) + env._shadow_hand_camera_embeddings = embeddings + return embeddings + + camera: Camera = env.scene.sensors[sensor_cfg.name] + object_asset: RigidObject = env.scene[object_cfg.name] + object_pos = object_asset.data.root_pos_w.torch - env.scene.env_origins + object_pose = torch.cat((object_pos, object_asset.data.root_quat_w.torch), dim=-1) + keypoints = compute_cube_keypoints(object_pose, out=self._keypoints_buf) + target = torch.cat((object_pos, keypoints.flatten(start_dim=1)), dim=-1) + camera_output = { + data_type: value if isinstance(value, torch.Tensor) else value.torch + for data_type, value in camera.data.output.items() + } + pose_loss, embeddings = self._feature_extractor.step(camera_output, target) + embeddings = embeddings.clone().detach() + env._shadow_hand_camera_embeddings = embeddings + if pose_loss is not None: + env.extras.setdefault("log", {})["pose_loss"] = pose_loss + return embeddings + + +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 + + +def shadow_hand_goal_keypoints(env: ManagerBasedRLEnv, command_name: str) -> torch.Tensor: + """Flattened zero-origin cube keypoints [m] for the current goal orientation. + + Args: + env: Environment containing the goal command term. + command_name: Goal command term name. + + Returns: + Flattened zero-origin cube keypoints [m], shape ``(num_envs, 24)``. + """ + command_term = env.command_manager.get_term(command_name) + return cube_keypoints_from_quat(command_term.quat_command_w) diff --git a/source/isaaclab_tasks/isaaclab_tasks/core/reorient/mdp/rewards.py b/source/isaaclab_tasks/isaaclab_tasks/core/reorient/mdp/rewards.py index 9974462578f4..6073312d82ce 100644 --- a/source/isaaclab_tasks/isaaclab_tasks/core/reorient/mdp/rewards.py +++ b/source/isaaclab_tasks/isaaclab_tasks/core/reorient/mdp/rewards.py @@ -7,94 +7,255 @@ 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, RewardTermCfg, SceneEntityCfg if TYPE_CHECKING: from isaaclab.assets import RigidObject from isaaclab.envs import ManagerBasedRLEnv - from .commands import ReorientCommand +class EpisodeErrorRecorder: + """Record the minimum physical error reached in each episode. -def success_bonus( - env: ManagerBasedRLEnv, command_name: str, object_cfg: SceneEntityCfg = SceneEntityCfg("object") -) -> torch.Tensor: - """Bonus reward for successfully reaching the goal. + The recorder deliberately contains no success threshold. This keeps the + measured task error separate from the policy that converts it to a success + result. + """ - The object is considered to have reached the goal when the object orientation is within the threshold. - The reward is 1.0 if the object has reached the goal, otherwise 0.0. + def __init__(self, num_envs: int, device: str | torch.device): + """Initialize per-environment error buffers. - Args: - env: The environment object. - command_name: The command term to be used for extracting the goal. - object_cfg: The configuration for the scene entity. Default is "object". - """ - # extract useful elements - asset: RigidObject = env.scene[object_cfg.name] - command_term: ReorientCommand = env.command_manager.get_term(command_name) + Args: + num_envs: Number of parallel environments. + device: Device on which to store the buffers. + """ + self.minimum_error = torch.full((num_envs,), torch.inf, device=device) + self._has_sample = torch.zeros(num_envs, dtype=torch.bool, device=device) + + def update(self, error: torch.Tensor) -> None: + """Record one error sample for every environment. - # obtain the goal orientation - goal_quat_w = command_term.command[:, 3:7] - # obtain the threshold for the orientation error - threshold = command_term.cfg.orientation_success_threshold - # calculate the orientation error - dtheta = math_utils.quat_error_magnitude(asset.data.root_quat_w.torch, goal_quat_w) + Args: + error: Per-environment physical errors, in task-defined units. - return dtheta <= threshold + Raises: + ValueError: If :paramref:`error` does not match the recorder shape. + """ + if error.shape != self.minimum_error.shape: + raise ValueError(f"Expected error shape {self.minimum_error.shape}, got {error.shape}.") + finite = torch.isfinite(error) + # non-finite samples keep the running minimum; boolean advanced indexing would + # force a host synchronization every step, so substitute-and-minimum instead + torch.minimum(self.minimum_error, torch.where(finite, error, self.minimum_error), out=self.minimum_error) + self._has_sample |= finite + def reset(self, env_ids: Sequence[int] | torch.Tensor | slice | None = None) -> dict[str, torch.Tensor]: + """Summarize and clear completed episodes. -def track_pos_l2( - env: ManagerBasedRLEnv, command_name: str, object_cfg: SceneEntityCfg = SceneEntityCfg("object") -) -> torch.Tensor: - """Reward for tracking the object position using the L2 norm. + Args: + env_ids: Environments whose episodes completed, or ``None`` for all. - The reward is the distance between the object position and the goal position. + Returns: + Mean, median, and 90th-percentile episode-minimum errors as 0-dim + device tensors, so logging them does not force a host + synchronization in the reset path. The result is empty when none of + the selected environments has a sample. + """ + if env_ids is None: + env_ids = slice(None) + valid = self._has_sample[env_ids] + values = self.minimum_error[env_ids][valid] + statistics = {} + if values.numel() > 0: + statistics = { + "mean": values.mean(), + "median": values.median(), + "p90": torch.quantile(values, 0.9), + } + self.minimum_error[env_ids] = torch.inf + self._has_sample[env_ids] = False + return statistics + + +@torch.jit.script +def evaluate_reorient_success( + object_quat: torch.Tensor, target_quat: torch.Tensor, success_tolerance: float +) -> tuple[torch.Tensor, torch.Tensor]: + """Evaluate reorientation success while exposing its physical error. + + This is the single per-step success evaluation: callers reuse the returned + flags and errors for the reward, the episode bookkeeping, and the + episode-minimum error tracking instead of recomputing the quaternion math. Args: - env: The environment object. - command_term: The command term to be used for extracting the goal. - object_cfg: The configuration for the scene entity. Default is "object". + object_quat: Object ``(x, y, z, w)`` orientations. + target_quat: Target ``(x, y, z, w)`` orientations. + success_tolerance: Maximum successful orientation error [rad]. + + Returns: + Per-environment success flags and orientation errors [rad]. """ - # extract useful elements - asset: RigidObject = env.scene[object_cfg.name] - command_term: ReorientCommand = env.command_manager.get_term(command_name) + orientation_error = math_utils.quat_error_magnitude(object_quat, target_quat) + return orientation_error <= success_tolerance, orientation_error - # obtain the goal position - goal_pos_e = command_term.command[:, 0:3] - # obtain the object position in the environment frame - object_pos_e = asset.data.root_pos_w.torch - env.scene.env_origins - return torch.linalg.norm(goal_pos_e - object_pos_e, ord=2, dim=-1) +@torch.jit.script +def reorient_reward( + reset_buf: torch.Tensor, + reset_goal_buf: torch.Tensor, + successes: torch.Tensor, + consecutive_successes: torch.Tensor, + object_pos: torch.Tensor, + target_pos: torch.Tensor, + goal_reached: torch.Tensor, + rotation_distance: torch.Tensor, + actions: torch.Tensor, + distance_scale: float, + rotation_scale: float, + rotation_epsilon: float, + action_penalty_scale: float, + success_bonus: float, + fall_distance: float, + fall_penalty: float, + averaging_factor: float, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + """Compute the Direct reorientation reward and success state transition. + The success evaluation is not recomputed here: callers pass the flags and + orientation errors from :func:`evaluate_reorient_success`, computed once + per step. -def track_orientation_inv_l2( - env: ManagerBasedRLEnv, - command_name: str, - object_cfg: SceneEntityCfg = SceneEntityCfg("object"), - rot_eps: float = 1e-3, -) -> torch.Tensor: - """Reward for tracking the object orientation using the inverse of the orientation error. + Args: + reset_buf: Current episode-reset flags. + reset_goal_buf: Current goal-reset flags. + successes: Goals reached in each episode. + consecutive_successes: Moving-average success count. + object_pos: Object positions in the environment frame [m]. + target_pos: Goal positions in the environment frame [m]. + goal_reached: Per-environment success flags for this step. + rotation_distance: Per-environment orientation errors [rad]. + actions: Normalized joint actions. + distance_scale: Position-distance reward scale [1/m]. + rotation_scale: Orientation reward scale [rad]. + rotation_epsilon: Orientation reward regularizer [rad]. + action_penalty_scale: Squared-action reward scale. + success_bonus: Reward added when a goal is reached. + fall_distance: Object-to-goal termination distance [m]. + fall_penalty: Reward added when the object is out of reach. + averaging_factor: Consecutive-success moving-average factor. + + Returns: + Reward, goal-reset flags, episode success counts, and moving-average + consecutive successes. + """ + goal_distance = torch.linalg.norm(object_pos - target_pos, ord=2, dim=-1) + goal_resets = reset_goal_buf | goal_reached + successes = successes + goal_resets + fell = goal_distance >= fall_distance + reward = ( + goal_distance * distance_scale + + rotation_scale / (rotation_distance + rotation_epsilon) + + actions.square().sum(dim=-1) * action_penalty_scale + + goal_resets.to(goal_distance.dtype) * success_bonus + + fell.to(goal_distance.dtype) * fall_penalty + ) + resets = reset_buf | fell + num_resets = resets.sum() + finished_successes = (successes * resets).sum() + mean_successes = finished_successes / num_resets.clamp_min(1) + consecutive_successes = torch.where( + num_resets > 0, + averaging_factor * mean_successes + (1.0 - averaging_factor) * consecutive_successes, + consecutive_successes, + ) + return reward, goal_resets, successes, consecutive_successes - The reward is the inverse of the orientation error between the object orientation and the goal orientation. - Args: - env: The environment object. - command_name: The command term to be used for extracting the goal. - object_cfg: The configuration for the scene entity. Default is "object". - rot_eps: The threshold for the orientation error. Default is 1e-3. +class ReorientReward(ManagerTermBase): + """Compute reorientation rewards with sticky per-episode success accounting. + + Matches the reward semantics of the Direct-workflow implementation + (:class:`~isaaclab_tasks.core.reorient.reorient_direct_env.ReorientDirectEnv`). + The scalar task parameters arrive as term params, set at the configuration + declaration site to match the Direct environment's values. """ - # extract useful elements - asset: RigidObject = env.scene[object_cfg.name] - command_term: ReorientCommand = env.command_manager.get_term(command_name) - # obtain the goal orientation - goal_quat_w = command_term.command[:, 3:7] - # calculate the orientation error - dtheta = math_utils.quat_error_magnitude(asset.data.root_quat_w.torch, goal_quat_w) + def __init__(self, cfg: RewardTermCfg, env: ManagerBasedRLEnv): + super().__init__(cfg, env) + self._successes = torch.zeros(self.num_envs, device=self.device) + self._consecutive_successes = torch.zeros(1, device=self.device) + self._orientation_error = EpisodeErrorRecorder(self.num_envs, self.device) + + @property + def successes(self) -> torch.Tensor: + """Goals reached in each current episode.""" + return self._successes + + def reset(self, env_ids: Sequence[int] | None = None) -> None: + if env_ids is None: + env_ids = slice(None) + threshold = self.cfg.params["success_count_threshold"] + # 0-dim device tensor: avoids a host sync here; consumers read it at logging cadence + self._env.extras.setdefault("log", {})["Metrics/success_rate"] = ( + (self._successes[env_ids] >= threshold).float().mean() + ) + for statistic, value in self._orientation_error.reset(env_ids).items(): + self._env.extras["log"][f"Diagnostics/episode_min_orientation_error_{statistic}"] = value + self._successes[env_ids] = 0.0 - return 1.0 / (dtheta + rot_eps) + def __call__( + self, + env: ManagerBasedRLEnv, + command_name: str, + distance_scale: float, + rotation_scale: float, + rotation_epsilon: float, + action_penalty_scale: float, + success_tolerance: float, + success_bonus: float, + fall_distance: float, + fall_penalty: float, + averaging_factor: float, + success_count_threshold: int, + action_name: str | None = None, + object_cfg: SceneEntityCfg = SceneEntityCfg("object"), + ) -> torch.Tensor: + del success_count_threshold # consumed in __init__ (used by reset()) + asset: RigidObject = env.scene[object_cfg.name] + command = env.command_manager.get_command(command_name) + object_pos = asset.data.root_pos_w.torch - env.scene.env_origins + actions = ( + env.action_manager.action if action_name is None else env.action_manager.get_term(action_name).raw_actions + ) + # single per-step success evaluation: the recorder and the reward reuse it + goal_reached, orientation_error = evaluate_reorient_success( + asset.data.root_quat_w.torch, command[:, 3:7], success_tolerance + ) + self._orientation_error.update(orientation_error) + reward, _, self._successes, self._consecutive_successes = reorient_reward( + env.reset_buf, + torch.zeros_like(env.reset_buf), + self._successes, + self._consecutive_successes, + object_pos, + command[:, :3], + goal_reached, + orientation_error, + actions, + distance_scale, + rotation_scale, + rotation_epsilon, + action_penalty_scale, + success_bonus, + fall_distance, + fall_penalty, + averaging_factor, + ) + env.extras.setdefault("log", {})["consecutive_successes"] = self._consecutive_successes.mean() + return reward / env.step_dt diff --git a/source/isaaclab_tasks/isaaclab_tasks/core/reorient/mdp/terminations.py b/source/isaaclab_tasks/isaaclab_tasks/core/reorient/mdp/terminations.py index 32d6df6aa18a..052e5cf5e1d1 100644 --- a/source/isaaclab_tasks/isaaclab_tasks/core/reorient/mdp/terminations.py +++ b/source/isaaclab_tasks/isaaclab_tasks/core/reorient/mdp/terminations.py @@ -11,7 +11,9 @@ import torch -from isaaclab.managers import SceneEntityCfg +from isaaclab.managers import ManagerTermBase, SceneEntityCfg, TerminationTermCfg + +from .rewards import ReorientReward, evaluate_reorient_success if TYPE_CHECKING: from isaaclab.envs import ManagerBasedRLEnv @@ -19,69 +21,77 @@ from .commands import ReorientCommand -def max_consecutive_success(env: ManagerBasedRLEnv, num_success: int, command_name: str) -> torch.Tensor: - """Check if the task has been completed consecutively for a certain number of times. +class object_reorientation_out_of_reach(ManagerTermBase): + """Terminate when object-to-goal distance is at least the threshold [m]. - Args: - env: The environment object. - num_success: Threshold for the number of consecutive successes required. - command_name: The command term to be used for extracting the goal. + The scalar task parameters arrive as term params, set at the configuration + declaration site to match the Direct environment's values. """ - command_term: ReorientCommand = env.command_manager.get_term(command_name) - - return command_term.metrics["consecutive_success"] >= num_success - - -def object_away_from_goal( - env: ManagerBasedRLEnv, - threshold: float, - command_name: str, - object_cfg: SceneEntityCfg = SceneEntityCfg("object"), -) -> torch.Tensor: - """Check if object has gone far from the goal. - The object is considered to be out-of-reach if the distance between the goal and the object is greater - than the threshold. - - Args: - env: The environment object. - threshold: The threshold for the distance between the robot and the object. - command_name: The command term to be used for extracting the goal. - object_cfg: The configuration for the scene entity. Default is "object". + def __init__(self, cfg: TerminationTermCfg, env: ManagerBasedRLEnv): + super().__init__(cfg, env) + # resolved on first call: the command term does not exist yet during manager construction + self._command_term: ReorientCommand | None = None + + def __call__( + self, + env: ManagerBasedRLEnv, + command_name: str, + threshold: float, + object_cfg: SceneEntityCfg = SceneEntityCfg("object"), + ) -> torch.Tensor: + """Return per-environment termination flags.""" + asset = env.scene[object_cfg.name] + if self._command_term is None: + self._command_term = env.command_manager.get_term(command_name) + distance = torch.linalg.norm(asset.data.root_pos_w.torch - self._command_term.pos_command_w, ord=2, dim=-1) + return distance >= threshold + + +class ReorientTimeout(ManagerTermBase): + """Apply progress-reset and timeout semantics with a consecutive-success cap. + + Matches the OpenAI-variant timeout semantics of the Direct-workflow implementation + (:class:`~isaaclab_tasks.core.reorient.reorient_direct_env.ReorientDirectEnv`): + resets the episode timer whenever the goal is reached so episodes extend + across goal streaks (real Direct dynamics, not boundary cosmetics), and + terminates after the consecutive-success cap or the usual timeout. The + scalar task parameters arrive as term params, set at the configuration + declaration site to match the Direct environment's values. """ - # extract useful elements - command_term: ReorientCommand = env.command_manager.get_term(command_name) - asset = env.scene[object_cfg.name] - - # object pos - asset_pos_e = asset.data.root_pos_w.torch - env.scene.env_origins - goal_pos_e = command_term.command[:, :3] - - return torch.linalg.norm(asset_pos_e - goal_pos_e, ord=2, dim=1) > threshold - - -def object_away_from_robot( - env: ManagerBasedRLEnv, - threshold: float, - asset_cfg: SceneEntityCfg = SceneEntityCfg("robot"), - object_cfg: SceneEntityCfg = SceneEntityCfg("object"), -) -> torch.Tensor: - """Check if object has gone far from the robot. - - The object is considered to be out-of-reach if the distance between the robot and the object is greater - than the threshold. - - Args: - env: The environment object. - threshold: The threshold for the distance between the robot and the object. - asset_cfg: The configuration for the robot entity. Default is "robot". - object_cfg: The configuration for the object entity. Default is "object". - """ - # extract useful elements - robot = env.scene[asset_cfg.name] - object = env.scene[object_cfg.name] - - # compute distance - dist = torch.linalg.norm(robot.data.root_pos_w.torch - object.data.root_pos_w.torch, dim=1) - return dist > threshold + def __init__(self, cfg: TerminationTermCfg, env: ManagerBasedRLEnv): + super().__init__(cfg, env) + # resolved on first call: the command term does not exist yet during manager construction + self._command_term: ReorientCommand | None = None + + def __call__( + self, + env: ManagerBasedRLEnv, + command_name: str, + reward_name: str, + success_tolerance: float, + max_successes: int, + object_cfg: SceneEntityCfg = SceneEntityCfg("object"), + ) -> torch.Tensor: + """Return per-environment timeout flags. + + Args: + env: Environment containing the object, goal, and reward term. + command_name: Goal command term name. + reward_name: Reorientation reward term name. + success_tolerance: Goal orientation tolerance [rad]. + max_successes: Consecutive-success cap before forcing a reset. + object_cfg: Object scene entity. + """ + object_asset = env.scene[object_cfg.name] + if self._command_term is None: + self._command_term = env.command_manager.get_term(command_name) + goal_reached, _ = evaluate_reorient_success( + object_asset.data.root_quat_w.torch, self._command_term.quat_command_w, success_tolerance + ) + # in place: rebinding env.episode_length_buf would orphan references held elsewhere + env.episode_length_buf.masked_fill_(goal_reached, 0) + reward_term: ReorientReward = env.reward_manager.get_term_cfg(reward_name).func + max_success_reached = reward_term.successes >= max_successes + return (env.episode_length_buf >= env.max_episode_length - 1) | max_success_reached diff --git a/source/isaaclab_tasks/isaaclab_tasks/core/reorient/reorient_common.py b/source/isaaclab_tasks/isaaclab_tasks/core/reorient/reorient_common.py new file mode 100644 index 000000000000..946f1d2e2f08 --- /dev/null +++ b/source/isaaclab_tasks/isaaclab_tasks/core/reorient/reorient_common.py @@ -0,0 +1,32 @@ +# 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 + +"""Task-family identity shared by the Direct and manager-based reorientation tasks. + +Geometry constants only — task tunables live inline in the workflow +configuration files, and robot identity lives in the per-robot ``*_common`` +modules. +""" + +GOAL_MARKER_POSITION: tuple[float, float, float] = (-0.2, -0.45, 0.68) +"""Fixed goal-marker display position [m], environment frame (state-based tasks).""" + +CAMERA_GOAL_MARKER_POSITION: tuple[float, float, float] = (-0.2, 0.1, 0.6) +"""Goal-marker display position [m] for the camera tasks. + +Deviates from :data:`GOAL_MARKER_POSITION` so the goal cube sits inside the +tiled camera's frustum. +""" + +IN_HAND_POS_OFFSET: tuple[float, float, float] = (0.0, 0.0, -0.04) +"""Offset from the object's default position to the in-hand goal anchor [m]. + +Defines the Direct/manager goal-position parity: the Direct environments and +the manager command terms derive the same in-hand target point from the +object's default root position plus this offset. +""" + +CAMERA_PLAY_NUM_ENVS: int = 64 +"""Camera-task environment count for checkpoint playback.""" diff --git a/source/isaaclab_tasks/isaaclab_tasks/core/reorient/reorient_direct_env.py b/source/isaaclab_tasks/isaaclab_tasks/core/reorient/reorient_direct_env.py index 5a632dc52ed2..a250bda10ffb 100644 --- a/source/isaaclab_tasks/isaaclab_tasks/core/reorient/reorient_direct_env.py +++ b/source/isaaclab_tasks/isaaclab_tasks/core/reorient/reorient_direct_env.py @@ -9,7 +9,6 @@ from collections.abc import Sequence from typing import TYPE_CHECKING -import numpy as np import torch import isaaclab.sim as sim_utils @@ -19,19 +18,15 @@ from isaaclab.markers import VisualizationMarkers from isaaclab.sensors import JointWrenchSensor, JointWrenchSensorCfg from isaaclab.sim.spawners.from_files import GroundPlaneCfg, spawn_ground_plane -from isaaclab.utils.math import ( - quat_conjugate, - quat_from_angle_axis, - quat_mul, - sample_uniform, - saturate, - scale_transform, - unscale_transform, -) +from isaaclab.utils.math import quat_conjugate, quat_mul, sample_uniform, saturate, scale_transform, unscale_transform + +from isaaclab_tasks.core.reorient.mdp.events import randomize_rotation, sample_joint_positions_within_limits +from isaaclab_tasks.core.reorient.mdp.rewards import EpisodeErrorRecorder, evaluate_reorient_success, reorient_reward +from isaaclab_tasks.core.reorient.reorient_common import GOAL_MARKER_POSITION, IN_HAND_POS_OFFSET if TYPE_CHECKING: from isaaclab_tasks.core.reorient.config.allegro_hand.allegro_hand_direct_env_cfg import AllegroHandEnvCfg - from isaaclab_tasks.core.reorient.config.shadow_hand.shadow_hand_env_cfg import ShadowHandEnvCfg + from isaaclab_tasks.core.reorient.config.shadow_hand.shadow_hand_direct_env_cfg import ShadowHandEnvCfg class ReorientDirectEnv(DirectRLEnv): @@ -40,61 +35,58 @@ class ReorientDirectEnv(DirectRLEnv): def __init__(self, cfg: AllegroHandEnvCfg | ShadowHandEnvCfg, render_mode: str | None = None, **kwargs): super().__init__(cfg, render_mode, **kwargs) + # -- robot introspection: joints, bodies, limits -- self.num_hand_dofs = self.hand.num_joints - - # buffers for position targets - self.hand_dof_targets = torch.zeros((self.num_envs, self.num_hand_dofs), dtype=torch.float, device=self.device) - self.prev_targets = torch.zeros((self.num_envs, self.num_hand_dofs), dtype=torch.float, device=self.device) - self.cur_targets = torch.zeros((self.num_envs, self.num_hand_dofs), dtype=torch.float, device=self.device) - - # list of actuated joints - self.actuated_dof_indices = list() - for joint_name in cfg.actuated_joint_names: - self.actuated_dof_indices.append(self.hand.joint_names.index(joint_name)) - self.actuated_dof_indices.sort() - - # finger bodies - self.finger_bodies = list() - for body_name in self.cfg.fingertip_body_names: - self.finger_bodies.append(self.hand.body_names.index(body_name)) - self.finger_bodies.sort() + self.actuated_dof_indices, _ = self.hand.find_joints(cfg.actuated_joint_names) + if len(self.actuated_dof_indices) != len(cfg.actuated_joint_names): + raise ValueError( + f"Expected {len(cfg.actuated_joint_names)} actuated joints, found {len(self.actuated_dof_indices)}." + ) + self.finger_bodies, fingertip_body_names = self.hand.find_bodies(self.cfg.fingertip_body_names) + if len(self.finger_bodies) != len(self.cfg.fingertip_body_names): + raise ValueError( + f"Expected {len(self.cfg.fingertip_body_names)} fingertip bodies, found {len(self.finger_bodies)}." + ) self.num_fingertips = len(self.finger_bodies) - self.finger_wrench_bodies = [] if getattr(self, "_joint_wrench_sensor", None) is not None: - for body_name in self.cfg.fingertip_body_names: + for body_name in fingertip_body_names: self.finger_wrench_bodies.append(self._joint_wrench_sensor.body_names.index(body_name)) self.finger_wrench_bodies.sort() - - # joint limits joint_pos_limits = self.hand.data.joint_limits.torch.to(self.device) self.hand_dof_lower_limits = joint_pos_limits[..., 0] self.hand_dof_upper_limits = joint_pos_limits[..., 1] - # track goal resets - self.reset_goal_buf = torch.zeros(self.num_envs, dtype=torch.bool, device=self.device) - # used to compare object position + # -- actuation targets (EMA-smoothed joint position targets) -- + self.prev_targets = torch.zeros((self.num_envs, self.num_hand_dofs), dtype=torch.float, device=self.device) + self.cur_targets = torch.zeros((self.num_envs, self.num_hand_dofs), dtype=torch.float, device=self.device) + + # -- goal and success state -- + # in-hand target = object default position + shared offset (mirrors ReorientCommand) self.in_hand_pos = self.object.data.default_root_pose.torch[:, 0:3].clone() - self.in_hand_pos[:, 2] -= 0.04 - # default goal positions + self.in_hand_pos += torch.tensor(IN_HAND_POS_OFFSET, dtype=torch.float, device=self.device) self.goal_rot = torch.zeros((self.num_envs, 4), dtype=torch.float, device=self.device) - self.goal_rot[:, 0] = 1.0 + self.goal_rot[:, 3] = 1.0 # identity quaternion in (x, y, z, w) layout self.goal_pos = torch.zeros((self.num_envs, 3), dtype=torch.float, device=self.device) - self.goal_pos[:, :] = torch.tensor([-0.2, -0.45, 0.68], device=self.device) - # initialize goal marker - self.goal_markers = VisualizationMarkers(self.cfg.goal_object_cfg) - - # track successes + self.goal_pos[:, :] = torch.tensor(GOAL_MARKER_POSITION, device=self.device) + self.reset_goal_buf = torch.zeros(self.num_envs, dtype=torch.bool, device=self.device) self.successes = torch.zeros(self.num_envs, dtype=torch.float, device=self.device) self.consecutive_successes = torch.zeros(1, dtype=torch.float, device=self.device) self._last_episode_success = torch.zeros(self.num_envs, dtype=torch.bool, device=self.device) - # unit tensors + # -- per-step evaluation state and diagnostics -- + # written once per step in :meth:`_get_dones`; the reward and metrics reuse them + self._success_flags = torch.zeros(self.num_envs, dtype=torch.bool, device=self.device) + self._orientation_error_buf = torch.full((self.num_envs,), torch.inf, device=self.device) + self._orientation_error = EpisodeErrorRecorder(self.num_envs, self.device) + + # -- reset randomization constants -- self.x_unit_tensor = torch.tensor([1, 0, 0], dtype=torch.float, device=self.device).repeat((self.num_envs, 1)) self.y_unit_tensor = torch.tensor([0, 1, 0], dtype=torch.float, device=self.device).repeat((self.num_envs, 1)) self.z_unit_tensor = torch.tensor([0, 0, 1], dtype=torch.float, device=self.device).repeat((self.num_envs, 1)) - # bind write methods + # -- visualization and articulation write handles -- + self.goal_markers = VisualizationMarkers(self.cfg.goal_object_cfg) self._set_joint_pos_target = self.hand.set_joint_position_target_index self._write_obj_root_pose = self.object.write_root_pose_to_sim_index self._write_obj_root_vel = self.object.write_root_velocity_to_sim_index @@ -190,32 +182,31 @@ def _update_fingertip_force_sensors(self) -> None: self.fingertip_force_sensors = torch.cat((force, torque), dim=-1) def _get_rewards(self) -> torch.Tensor: - ( - total_reward, - self.reset_goal_buf, - self.successes[:], - self.consecutive_successes[:], - ) = compute_rewards( + # the success flags and orientation errors were computed this step by + # :meth:`_get_dones`; the recorder and the reward reuse them + self._orientation_error.update(self._orientation_error_buf) + total_reward, goal_resets, successes, consecutive_successes = reorient_reward( self.reset_buf, self.reset_goal_buf, self.successes, self.consecutive_successes, - self.max_episode_length, self.object_pos, - self.object_rot, self.in_hand_pos, - self.goal_rot, + self._success_flags, + self._orientation_error_buf, + self.actions, self.cfg.dist_reward_scale, self.cfg.rot_reward_scale, self.cfg.rot_eps, - self.actions, self.cfg.action_penalty_scale, - self.cfg.success_tolerance, self.cfg.reach_goal_bonus, self.cfg.fall_dist, self.cfg.fall_penalty, self.cfg.av_factor, ) + self.reset_goal_buf.copy_(goal_resets) + self.successes[:] = successes + self.consecutive_successes[:] = consecutive_successes if "log" not in self.extras: self.extras["log"] = dict() @@ -235,11 +226,15 @@ def _get_dones(self) -> tuple[torch.Tensor, torch.Tensor]: goal_dist = torch.linalg.norm(self.object_pos - self.in_hand_pos, ord=2, dim=-1) out_of_reach = goal_dist >= self.cfg.fall_dist + # single per-step success evaluation; the reward and metrics reuse these buffers + self._success_flags, self._orientation_error_buf = evaluate_reorient_success( + self.object_rot, self.goal_rot, self.cfg.success_tolerance + ) + if self.cfg.max_consecutive_success > 0: - # Reset progress (episode length buf) on goal envs if max_consecutive_success > 0 - rot_dist = rotation_distance(self.object_rot, self.goal_rot) + # reset progress (episode length buf) on goal environments self.episode_length_buf = torch.where( - torch.abs(rot_dist) <= self.cfg.success_tolerance, + self._success_flags, torch.zeros_like(self.episode_length_buf), self.episode_length_buf, ) @@ -253,9 +248,10 @@ def _get_dones(self) -> tuple[torch.Tensor, torch.Tensor]: def _reset_idx(self, env_ids: Sequence[int]): # Episode counts as successful when goals reached >= cfg.success_count_threshold. self._last_episode_success[env_ids] = self.successes[env_ids] >= self.cfg.success_count_threshold - self.extras.setdefault("log", {})["Metrics/success_rate"] = ( - self._last_episode_success[env_ids].float().mean().item() - ) + # 0-dim device tensor: avoids a host sync here; consumers read it at logging cadence + self.extras.setdefault("log", {})["Metrics/success_rate"] = self._last_episode_success[env_ids].float().mean() + for statistic, value in self._orientation_error.reset(env_ids).items(): + self.extras["log"][f"Diagnostics/episode_min_orientation_error_{statistic}"] = value super()._reset_idx(env_ids) @@ -281,19 +277,15 @@ def _reset_idx(self, env_ids: Sequence[int]): self._write_obj_root_vel(root_velocity=object_default_vel, env_ids=env_ids) # reset hand - delta_max = self.hand_dof_upper_limits[env_ids] - self.hand.data.default_joint_pos.torch[env_ids] - delta_min = self.hand_dof_lower_limits[env_ids] - self.hand.data.default_joint_pos.torch[env_ids] - - dof_pos_noise = sample_uniform(-1.0, 1.0, (len(env_ids), self.num_hand_dofs), device=self.device) - rand_delta = delta_min + (delta_max - delta_min) * 0.5 * dof_pos_noise - dof_pos = self.hand.data.default_joint_pos.torch[env_ids] + self.cfg.reset_dof_pos_noise * rand_delta + default_dof_pos = self.hand.data.default_joint_pos.torch[env_ids] + dof_limits = self.hand.data.joint_limits.torch[env_ids] + dof_pos = sample_joint_positions_within_limits(default_dof_pos, dof_limits, self.cfg.reset_dof_pos_noise) dof_vel_noise = sample_uniform(-1.0, 1.0, (len(env_ids), self.num_hand_dofs), device=self.device) dof_vel = self.hand.data.default_joint_vel.torch[env_ids] + self.cfg.reset_dof_vel_noise * dof_vel_noise self.prev_targets[env_ids] = dof_pos self.cur_targets[env_ids] = dof_pos - self.hand_dof_targets[env_ids] = dof_pos self._set_joint_pos_target(target=dof_pos, env_ids=env_ids) self._write_hand_joint_pos(position=dof_pos, env_ids=env_ids) @@ -317,6 +309,7 @@ def _reset_target_pose(self, env_ids): self.reset_goal_buf[env_ids] = 0 def _compute_intermediate_values(self): + """Refresh the torch-side state snapshots consumed by the observation and reward paths.""" # data for hand self.fingertip_pos = self.hand.data.body_pos_w.torch[:, self.finger_bodies] self.fingertip_rot = self.hand.data.body_quat_w.torch[:, self.finger_bodies] @@ -405,75 +398,3 @@ def compute_full_state(self): dim=-1, ) return states - - -@torch.jit.script -def randomize_rotation(rand0, rand1, x_unit_tensor, y_unit_tensor): - return quat_mul( - quat_from_angle_axis(rand0 * np.pi, x_unit_tensor), quat_from_angle_axis(rand1 * np.pi, y_unit_tensor) - ) - - -@torch.jit.script -def rotation_distance(object_rot, target_rot): - # Orientation alignment for the cube in hand and goal cube - quat_diff = quat_mul(object_rot, quat_conjugate(target_rot)) - return 2.0 * torch.asin(torch.clamp(torch.linalg.norm(quat_diff[:, 0:3], ord=2, dim=-1), max=1.0)) - - -@torch.jit.script -def compute_rewards( - reset_buf: torch.Tensor, - reset_goal_buf: torch.Tensor, - successes: torch.Tensor, - consecutive_successes: torch.Tensor, - max_episode_length: float, - object_pos: torch.Tensor, - object_rot: torch.Tensor, - target_pos: torch.Tensor, - target_rot: torch.Tensor, - dist_reward_scale: float, - rot_reward_scale: float, - rot_eps: float, - actions: torch.Tensor, - action_penalty_scale: float, - success_tolerance: float, - reach_goal_bonus: float, - fall_dist: float, - fall_penalty: float, - av_factor: float, -): - goal_dist = torch.linalg.norm(object_pos - target_pos, ord=2, dim=-1) - rot_dist = rotation_distance(object_rot, target_rot) - - dist_rew = goal_dist * dist_reward_scale - rot_rew = 1.0 / (torch.abs(rot_dist) + rot_eps) * rot_reward_scale - - action_penalty = torch.sum(actions**2, dim=-1) - - # Total reward is: position distance + orientation alignment + action regularization + success bonus + fall penalty - reward = dist_rew + rot_rew + action_penalty * action_penalty_scale - - # Find out which envs hit the goal and update successes count - goal_resets = torch.where(torch.abs(rot_dist) <= success_tolerance, torch.ones_like(reset_goal_buf), reset_goal_buf) - successes = successes + goal_resets - - # Success bonus: orientation is within `success_tolerance` of goal orientation - reward = torch.where(goal_resets == 1, reward + reach_goal_bonus, reward) - - # Fall penalty: distance to the goal is larger than a threshold - reward = torch.where(goal_dist >= fall_dist, reward + fall_penalty, reward) - - # Check env termination conditions, including maximum success number - resets = torch.where(goal_dist >= fall_dist, torch.ones_like(reset_buf), reset_buf) - - num_resets = torch.sum(resets) - finished_cons_successes = torch.sum(successes * resets.float()) - - cons_successes = torch.where( - num_resets > 0, - av_factor * finished_cons_successes / num_resets + (1.0 - av_factor) * consecutive_successes, - consecutive_successes, - ) - - return reward, goal_resets, successes, cons_successes diff --git a/source/isaaclab_tasks/isaaclab_tasks/core/reorient/reorient_manager_env_cfg.py b/source/isaaclab_tasks/isaaclab_tasks/core/reorient/reorient_manager_env_cfg.py deleted file mode 100644 index ef1639a61e32..000000000000 --- a/source/isaaclab_tasks/isaaclab_tasks/core/reorient/reorient_manager_env_cfg.py +++ /dev/null @@ -1,348 +0,0 @@ -# Copyright (c) 2022-2026, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md). -# All rights reserved. -# -# SPDX-License-Identifier: BSD-3-Clause - -from __future__ import annotations - -from dataclasses import MISSING - -from isaaclab_physx.physics import PhysxCfg - -import isaaclab.sim as sim_utils -from isaaclab.assets import ArticulationCfg, AssetBaseCfg, RigidObjectCfg -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.sim.simulation_cfg import SimulationCfg -from isaaclab.sim.spawners.materials import RigidBodyMaterialCfg -from isaaclab.utils.assets import ISAAC_NUCLEUS_DIR -from isaaclab.utils.configclass import configclass -from isaaclab.utils.noise import GaussianNoiseCfg as Gnoise - -import isaaclab_tasks.core.reorient.mdp as mdp - -## -# Scene definition -## - - -@configclass -class ReorientObjectSceneCfg(InteractiveSceneCfg): - """Configuration for a scene with an object and a dexterous hand.""" - - # robots - robot: ArticulationCfg = MISSING - - # objects - object: RigidObjectCfg = RigidObjectCfg( - prim_path="{ENV_REGEX_NS}/object", - spawn=sim_utils.UsdFileCfg( - usd_path=f"{ISAAC_NUCLEUS_DIR}/Props/Blocks/DexCube/dex_cube_instanceable.usd", - rigid_props=sim_utils.RigidBodyPropertiesCfg( - kinematic_enabled=False, - disable_gravity=False, - enable_gyroscopic_forces=True, - solver_position_iteration_count=8, - solver_velocity_iteration_count=0, - sleep_threshold=0.005, - stabilization_threshold=0.0025, - max_depenetration_velocity=1000.0, - ), - mass_props=sim_utils.MassPropertiesCfg(density=400.0), - ), - init_state=RigidObjectCfg.InitialStateCfg(pos=(0.0, -0.19, 0.56), rot=(0.0, 0.0, 0.0, 1.0)), - ) - - # lights - light = AssetBaseCfg( - prim_path="/World/light", - spawn=sim_utils.DistantLightCfg(color=(0.95, 0.95, 0.95), intensity=1000.0), - ) - - dome_light = AssetBaseCfg( - prim_path="/World/domeLight", - spawn=sim_utils.DomeLightCfg(color=(0.02, 0.02, 0.02), intensity=1000.0), - ) - - -## -# MDP settings -## - - -@configclass -class CommandsCfg: - """Command specifications for the MDP.""" - - object_pose = mdp.ReorientCommandCfg( - asset_name="object", - init_pos_offset=(0.0, 0.0, -0.04), - update_goal_on_success=True, - orientation_success_threshold=0.1, - make_quat_unique=False, - marker_pos_offset=(-0.2, -0.06, 0.08), - debug_vis=True, - ) - - -@configclass -class ActionsCfg: - """Action specifications for the MDP.""" - - joint_pos = mdp.EMAJointPositionToLimitsActionCfg( - asset_name="robot", - joint_names=[".*"], - alpha=0.95, - rescale_to_limits=True, - ) - - -@configclass -class ObservationsCfg: - """Observation specifications for the MDP.""" - - @configclass - class KinematicObsGroupCfg(ObsGroup): - """Observations with full-kinematic state information. - - This does not include acceleration or force information. - """ - - # observation terms (order preserved) - # -- robot terms - joint_pos = ObsTerm(func=mdp.joint_pos_limit_normalized, noise=Gnoise(std=0.005)) - joint_vel = ObsTerm(func=mdp.joint_vel_rel, scale=0.2, noise=Gnoise(std=0.01)) - - # -- object terms - object_pos = ObsTerm( - func=mdp.root_pos_w, noise=Gnoise(std=0.002), 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, noise=Gnoise(std=0.002), params={"asset_cfg": SceneEntityCfg("object")} - ) - object_ang_vel = ObsTerm( - func=mdp.root_ang_vel_w, - scale=0.2, - noise=Gnoise(std=0.002), - params={"asset_cfg": SceneEntityCfg("object")}, - ) - - # -- command terms - 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}, - ) - - # -- action terms - last_action = ObsTerm(func=mdp.last_action) - - def __post_init__(self): - self.enable_corruption = True - self.concatenate_terms = True - - @configclass - class NoVelocityKinematicObsGroupCfg(KinematicObsGroupCfg): - """Observations with partial kinematic state information. - - In contrast to the full-kinematic state group, this group does not include velocity information - about the robot joints and the object root frame. This is useful for tasks where velocity information - is not available or has a lot of noise. - """ - - def __post_init__(self): - # call parent post init - super().__post_init__() - # set unused terms to None - self.joint_vel = None - self.object_lin_vel = None - self.object_ang_vel = None - - # observation groups - policy: KinematicObsGroupCfg = KinematicObsGroupCfg() - - -@configclass -class EventCfg: - """Configuration for randomization.""" - - # startup - # -- robot - robot_physics_material = EventTerm( - func=mdp.randomize_rigid_body_material, - mode="startup", - params={ - "asset_cfg": SceneEntityCfg("robot", body_names=".*"), - "static_friction_range": (0.7, 1.3), - "dynamic_friction_range": (0.7, 1.3), - "restitution_range": (0.0, 0.0), - "num_buckets": 250, - }, - ) - robot_scale_mass = EventTerm( - func=mdp.randomize_rigid_body_mass, - mode="startup", - params={ - "asset_cfg": SceneEntityCfg("robot", body_names=".*"), - "mass_distribution_params": (0.95, 1.05), - "operation": "scale", - }, - ) - robot_joint_stiffness_and_damping = EventTerm( - func=mdp.randomize_actuator_gains, - mode="startup", - params={ - "asset_cfg": SceneEntityCfg("robot", joint_names=".*"), - "stiffness_distribution_params": (0.3, 3.0), # default: 3.0 - "damping_distribution_params": (0.75, 1.5), # default: 0.1 - "operation": "scale", - "distribution": "log_uniform", - }, - ) - - # -- object - object_physics_material = EventTerm( - func=mdp.randomize_rigid_body_material, - mode="startup", - params={ - "asset_cfg": SceneEntityCfg("object", body_names=".*"), - "static_friction_range": (0.7, 1.3), - "dynamic_friction_range": (0.7, 1.3), - "restitution_range": (0.0, 0.0), - "num_buckets": 250, - }, - ) - object_scale_mass = EventTerm( - func=mdp.randomize_rigid_body_mass, - mode="startup", - params={ - "asset_cfg": SceneEntityCfg("object"), - "mass_distribution_params": (0.4, 1.6), - "operation": "scale", - }, - ) - - # reset - reset_object = EventTerm( - func=mdp.reset_root_state_uniform, - mode="reset", - params={ - "pose_range": {"x": [-0.01, 0.01], "y": [-0.01, 0.01], "z": [-0.01, 0.01]}, - "velocity_range": {}, - "asset_cfg": SceneEntityCfg("object", body_names=".*"), - }, - ) - reset_robot_joints = EventTerm( - func=mdp.reset_joints_within_limits_range, - mode="reset", - params={ - "position_range": {".*": [0.2, 0.2]}, - "velocity_range": {".*": [0.0, 0.0]}, - "use_default_offset": True, - "operation": "scale", - }, - ) - - -@configclass -class RewardsCfg: - """Reward terms for the MDP.""" - - # -- task - # track_pos_l2 = RewTerm( - # func=mdp.track_pos_l2, - # weight=-10.0, - # params={"object_cfg": SceneEntityCfg("object"), "command_name": "object_pose"}, - # ) - track_orientation_inv_l2 = RewTerm( - func=mdp.track_orientation_inv_l2, - weight=1.0, - params={"object_cfg": SceneEntityCfg("object"), "rot_eps": 0.1, "command_name": "object_pose"}, - ) - success_bonus = RewTerm( - func=mdp.success_bonus, - weight=250.0, - params={"object_cfg": SceneEntityCfg("object"), "command_name": "object_pose"}, - ) - - # -- penalties - joint_vel_l2 = RewTerm(func=mdp.joint_vel_l2, weight=-2.5e-5) - action_l2 = RewTerm(func=mdp.action_l2, weight=-0.0001) - action_rate_l2 = RewTerm(func=mdp.action_rate_l2, weight=-0.01) - - # -- optional penalties (these are disabled by default) - # object_away_penalty = RewTerm( - # func=mdp.is_terminated_term, - # weight=-0.0, - # params={"term_keys": "object_out_of_reach"}, - # ) - - -@configclass -class TerminationsCfg: - """Termination terms for the MDP.""" - - time_out = DoneTerm(func=mdp.time_out, time_out=True) - - max_consecutive_success = DoneTerm( - func=mdp.max_consecutive_success, params={"num_success": 50, "command_name": "object_pose"} - ) - - object_out_of_reach = DoneTerm(func=mdp.object_away_from_robot, params={"threshold": 0.3}) - - # object_out_of_reach = DoneTerm( - # func=mdp.object_away_from_goal, params={"threshold": 0.24, "command_name": "object_pose"} - # ) - - -## -# Environment configuration -## - - -@configclass -class ReorientObjectEnvCfg(ManagerBasedRLEnvCfg): - """Configuration for the in hand reorientation environment.""" - - # Scene settings - scene: ReorientObjectSceneCfg = ReorientObjectSceneCfg(num_envs=8192, env_spacing=0.6) - # Simulation settings - sim: SimulationCfg = SimulationCfg( - physics_material=RigidBodyMaterialCfg( - static_friction=1.0, - dynamic_friction=1.0, - ), - physics=PhysxCfg( - bounce_threshold_velocity=0.2, - gpu_max_rigid_contact_count=2**20, - gpu_max_rigid_patch_count=2**23, - ), - ) - # Basic settings - observations: ObservationsCfg = ObservationsCfg() - actions: ActionsCfg = ActionsCfg() - commands: CommandsCfg = CommandsCfg() - # MDP settings - rewards: RewardsCfg = RewardsCfg() - terminations: TerminationsCfg = TerminationsCfg() - events: EventCfg = EventCfg() - - def __post_init__(self): - """Post initialization.""" - # general settings - self.decimation = 4 - self.episode_length_s = 20.0 - # simulation settings - self.sim.dt = 1.0 / 120.0 - self.sim.render_interval = self.decimation - # change viewer settings - self.viewer.eye = (2.0, 2.0, 2.0) diff --git a/source/isaaclab_tasks/test/benchmarking/configs.yaml b/source/isaaclab_tasks/test/benchmarking/configs.yaml index 62a652e995e3..9495a8196fe7 100644 --- a/source/isaaclab_tasks/test/benchmarking/configs.yaml +++ b/source/isaaclab_tasks/test/benchmarking/configs.yaml @@ -235,11 +235,18 @@ full: Isaac-Reorient-Cube-Allegro*: max_iterations: 500 lower_thresholds: - reward: 15 - episode_length: 300 + reward: 200 + episode_length: 150 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,11 +275,34 @@ 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: + # evidence-calibrated: full-budget Newton run crosses 1000 at iter ~2300 + # (PhysX at iter ~1000); tail success rate 0.83 at the 0.1 rad tolerance reward: 1000 episode_length: 400 + success_rate: + value: 0.3 + consecutive_samples: 20 + upper_thresholds: + duration: 40000 + Isaac-Reorient-Cube-Shadow-Camera: + max_iterations: 3000 + lower_thresholds: + # interim reward gate: catches pipeline breakage well below the truncated + # manager plateau (~174-290); tighten toward the Direct row's 1000 once a + # full-budget manager run calibrates it (success gates stay off until then + # — success is near-binary below convergence) + reward: 150 + episode_length: 400 upper_thresholds: duration: 40000 Isaac-Shadow-Handover-Direct: @@ -275,6 +312,19 @@ full: episode_length: 150 upper_thresholds: duration: 10000 + Isaac-Shadow-Handover: + max_iterations: 3000 + lower_thresholds: + # evidence-calibrated: the manager variant plateaus near 785 on PhysX + # (the Direct rows keep the shared 1000 gate) + reward: 500 + episode_length: 150 + # provisional gate matching the validation campaign's controller + success_rate: + value: 0.3 + consecutive_samples: 20 + upper_thresholds: + duration: 10000 # Dexterous manipulation - KukaAllegro Isaac-Lift-KukaAllegro: max_iterations: 500 diff --git a/source/isaaclab_tasks/test/benchmarking/env_benchmark_test_utils.py b/source/isaaclab_tasks/test/benchmarking/env_benchmark_test_utils.py index acfd84770ed9..02ad32285847 100644 --- a/source/isaaclab_tasks/test/benchmarking/env_benchmark_test_utils.py +++ b/source/isaaclab_tasks/test/benchmarking/env_benchmark_test_utils.py @@ -8,11 +8,20 @@ import math import os import re +from numbers import Real import numpy as np import yaml +def _is_training_task(task_id: str) -> bool: + """Return whether a registered task is intended for training benchmarks.""" + stem, separator, version = task_id.rpartition("-v") + if separator and version.isdigit(): + task_id = stem + return not {"Play", "Benchmark"}.intersection(task_id.split("-")) + + def _get_repo_path(): """Get the repository root by searching for marker files. @@ -81,14 +90,27 @@ def evaluate_job(workflow, task, env_config, duration): thresholds = {**env_config.get("lower_thresholds", {}), **env_config.get("upper_thresholds", {})} # evaluate all thresholds from the config - for threshold_name, threshold_val in thresholds.items(): + for threshold_name, threshold_spec in thresholds.items(): uses_lower_threshold = threshold_name in env_config.get("lower_thresholds", {}) + threshold_val, consecutive_samples = _parse_threshold_spec(threshold_spec) if threshold_name == "duration": val = duration else: - val = _extract_log_val(threshold_name, log_data, uses_lower_threshold, workflow) - # skip non-numeric values - if val is None or not isinstance(val, (int, float)) or (isinstance(val, float) and math.isnan(val)): + val = _extract_log_val( + threshold_name, + log_data, + uses_lower_threshold, + workflow, + consecutive_samples=consecutive_samples, + ) + if val is None or not isinstance(val, Real) or not math.isfinite(float(val)): + kpi_payload[threshold_name] = None + kpi_payload[f"{threshold_name}_threshold"] = threshold_val + if consecutive_samples is not None: + kpi_payload[f"{threshold_name}_consecutive_samples"] = consecutive_samples + kpi_payload["success"] = False + if not kpi_payload["msg"]: + kpi_payload["msg"] = f"{threshold_name} metric is missing or non-numeric" continue val = round(val, 4) threshold_val_rounded = round(threshold_val, 4) @@ -107,6 +129,8 @@ def evaluate_job(workflow, task, env_config, duration): normalized_reward = val / threshold_val kpi_payload[f"{threshold_name}_normalized"] = normalized_reward kpi_payload[f"{threshold_name}_threshold"] = threshold_val + if consecutive_samples is not None: + kpi_payload[f"{threshold_name}_consecutive_samples"] = consecutive_samples # add max iterations to the payload max_iterations = env_config.get("max_iterations") @@ -219,9 +243,26 @@ def _parse_tf_logs(log): return log_data -def _extract_log_val(name, log_data, uses_lower_threshold, workflow): +def _parse_threshold_spec(threshold_spec): + """Return a threshold value and optional sustained-sample requirement.""" + if isinstance(threshold_spec, Real): + return threshold_spec, None + if not isinstance(threshold_spec, dict): + raise TypeError(f"Threshold must be a number or mapping, got {type(threshold_spec).__name__}") + + threshold_val = threshold_spec.get("value") + consecutive_samples = threshold_spec.get("consecutive_samples") + if not isinstance(threshold_val, Real): + raise TypeError("Structured thresholds require a numeric 'value'") + if not isinstance(consecutive_samples, int) or isinstance(consecutive_samples, bool) or consecutive_samples < 1: + raise ValueError("Structured thresholds require 'consecutive_samples' to be a positive integer") + return threshold_val, consecutive_samples + + +def _extract_log_val(name, log_data, uses_lower_threshold, workflow, consecutive_samples=None): """Extract the value from the log data.""" try: + tag = None if name == "reward": reward_tags = { "rl_games": "rewards/iter", @@ -230,7 +271,7 @@ def _extract_log_val(name, log_data, uses_lower_threshold, workflow): "skrl": "Reward / Total reward (mean)", } tag = reward_tags.get(workflow) - if tag: + if tag and consecutive_samples is None: return _extract_reward(log_data, tag) elif name == "episode_length": @@ -241,8 +282,20 @@ def _extract_log_val(name, log_data, uses_lower_threshold, workflow): "skrl": "Episode / Total timesteps (mean)", } tag = episode_tags.get(workflow) - if tag: - return _extract_feature(log_data, tag, uses_lower_threshold) + elif name == "success_rate": + success_rate_tags = { + "rl_games": "Episode/Metrics/success_rate", + "rsl_rl": "Metrics/success_rate", + "skrl": "Metrics/success_rate", + } + tag = success_rate_tags.get(workflow) + + if tag: + if consecutive_samples is not None: + return _extract_sustained_feature(log_data, tag, uses_lower_threshold, consecutive_samples) + return _extract_feature(log_data, tag, uses_lower_threshold) + if name == "success_rate": + return None except KeyError as e: print(f"Warning: Metric '{name}' not found in logs for workflow '{workflow}': {e}") return None @@ -263,6 +316,21 @@ def _extract_feature(log_data, feature, uses_lower_threshold): return min(log_data) +def _extract_sustained_feature(log_data, feature, uses_lower_threshold, consecutive_samples): + """Extract the best threshold-facing value sustained over a sample window.""" + values = np.asarray(log_data[feature], dtype=float)[:, 1] + if len(values) < consecutive_samples: + return None + if not np.all(np.isfinite(values)): + return math.nan + + window_extrema = [] + for start in range(len(values) - consecutive_samples + 1): + window = values[start : start + consecutive_samples] + window_extrema.append(min(window) if uses_lower_threshold else max(window)) + return max(window_extrema) if uses_lower_threshold else min(window_extrema) + + def _extract_reward(log_data, feature, k=8): """Extract the averaged max reward from the log data.""" log_data = np.array(log_data[feature])[:, 1] diff --git a/source/isaaclab_tasks/test/benchmarking/test_env_benchmark_test_utils.py b/source/isaaclab_tasks/test/benchmarking/test_env_benchmark_test_utils.py new file mode 100644 index 000000000000..2aef0e5aba94 --- /dev/null +++ b/source/isaaclab_tasks/test/benchmarking/test_env_benchmark_test_utils.py @@ -0,0 +1,159 @@ +# Copyright (c) 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 + +"""Unit tests for training benchmark KPI evaluation.""" + +import math + +import env_benchmark_test_utils as utils +import pytest + +_SUSTAINED_SUCCESS_THRESHOLD = {"value": 0.3, "consecutive_samples": 20} + + +def _evaluate_success_rate(monkeypatch, workflow, tag, values): + """Evaluate a success-rate series with the sustained-success requirement.""" + log_data = {tag: list(enumerate(values))} + monkeypatch.setattr(utils, "_retrieve_logs", lambda workflow, task: log_data) + return utils.evaluate_job( + workflow, + "Isaac-Test", + {"lower_thresholds": {"success_rate": _SUSTAINED_SUCCESS_THRESHOLD}}, + duration=1.0, + ) + + +@pytest.mark.parametrize( + "task_id,expected", + [ + ("Isaac-Reorient-Cube-Shadow-Camera", True), + ("Isaac-Reorient-Cube-Shadow-Camera-v0", True), + ("Isaac-Reorient-Cube-Shadow-Camera-Play", False), + ("Isaac-Reorient-Cube-Shadow-Camera-Play-v0", False), + ("Isaac-Reorient-Cube-Shadow-Camera-Benchmark", False), + ("Isaac-Reorient-Cube-Shadow-Camera-Benchmark-v1", False), + ("Isaac-Reorient-Cube-Shadow-Camera-Benchmark-Direct", False), + ("Isaac-Reorient-Cube-Shadow-Camera-Benchmark-Direct-v0", False), + ], +) +def test_training_task_filter_excludes_play_and_benchmark(task_id: str, expected: bool): + """Verify inference-only variants never enter the training benchmark matrix.""" + assert utils._is_training_task(task_id) is expected + + +@pytest.mark.parametrize( + "log_data", + [ + {"Train/mean_episode_length": [(0, 100.0)]}, + {"Train/mean_reward": [(0, math.nan)]}, + ], + ids=["missing", "nan"], +) +def test_evaluate_job_fails_when_configured_reward_is_unavailable(monkeypatch, log_data): + """Verify missing or invalid configured metrics cannot produce a successful KPI payload.""" + monkeypatch.setattr(utils, "_retrieve_logs", lambda workflow, task: log_data) + + payload = utils.evaluate_job( + "rsl_rl", + "Isaac-Test", + {"lower_thresholds": {"reward": 1000.0}}, + duration=1.0, + ) + + assert payload["success"] is False + assert payload["msg"] == "reward metric is missing or non-numeric" + assert payload["reward"] is None + assert payload["reward_threshold"] == 1000.0 + + +@pytest.mark.parametrize( + "workflow_name,tag", + [ + ("rl_games", "Episode/Metrics/success_rate"), + ("rsl_rl", "Metrics/success_rate"), + ("skrl", "Metrics/success_rate"), + ], +) +def test_evaluate_job_passes_sustained_success_rate(monkeypatch, workflow_name, tag): + """Verify each supported workflow passes after 20 consecutive successful samples.""" + payload = _evaluate_success_rate(monkeypatch, workflow_name, tag, [0.2, *([0.3] * 20)]) + + assert payload["success"] is True + assert payload["success_rate"] == 0.3 + assert payload["success_rate_threshold"] == 0.3 + assert payload["success_rate_consecutive_samples"] == 20 + + +def test_evaluate_job_fails_success_rate_below_threshold(monkeypatch): + """Verify a sustained series below the configured success rate fails.""" + payload = _evaluate_success_rate(monkeypatch, "rsl_rl", "Metrics/success_rate", [0.29] * 20) + + assert payload["success"] is False + assert payload["msg"] == "success_rate below threshold: 0.29 < 0.3" + assert payload["success_rate"] == 0.29 + + +def test_evaluate_job_fails_when_success_rate_is_missing(monkeypatch): + """Verify a missing configured success-rate metric fails.""" + payload = _evaluate_success_rate(monkeypatch, "rsl_rl", "Train/mean_reward", [1000.0] * 20) + + assert payload["success"] is False + assert payload["msg"] == "success_rate metric is missing or non-numeric" + assert payload["success_rate"] is None + + +def test_evaluate_job_treats_success_rate_as_missing_for_unsupported_workflow(monkeypatch): + """Verify an unmapped workflow fails the KPI instead of raising an exception.""" + payload = _evaluate_success_rate(monkeypatch, "sb3", "Metrics/success_rate", [0.3] * 20) + + assert payload["success"] is False + assert payload["msg"] == "success_rate metric is missing or non-numeric" + assert payload["success_rate"] is None + + +def test_evaluate_job_fails_when_success_rate_is_nonfinite(monkeypatch): + """Verify a non-finite configured success-rate metric fails.""" + payload = _evaluate_success_rate( + monkeypatch, + "rsl_rl", + "Metrics/success_rate", + [0.3] * 19 + [math.nan], + ) + + assert payload["success"] is False + assert payload["msg"] == "success_rate metric is missing or non-numeric" + assert payload["success_rate"] is None + + +def test_evaluate_job_resets_sustained_success_streak(monkeypatch): + """Verify a below-threshold sample resets the consecutive-success streak.""" + payload = _evaluate_success_rate( + monkeypatch, + "rsl_rl", + "Metrics/success_rate", + [0.3] * 19 + [0.29] + [0.3] * 19, + ) + + assert payload["success"] is False + assert payload["msg"] == "success_rate below threshold: 0.29 < 0.3" + assert payload["success_rate"] == 0.29 + + +def test_scalar_threshold_behavior_is_preserved(monkeypatch): + """Verify numeric thresholds retain the existing reward aggregation behavior.""" + log_data = {"Train/mean_reward": list(enumerate([1.0, 2.0, 3.0]))} + monkeypatch.setattr(utils, "_retrieve_logs", lambda workflow, task: log_data) + + payload = utils.evaluate_job( + "rsl_rl", + "Isaac-Test", + {"lower_thresholds": {"reward": 2.0}}, + duration=1.0, + ) + + assert payload["success"] is True + assert payload["reward"] == 2.0 + assert payload["reward_threshold"] == 2.0 + assert "reward_consecutive_samples" not in payload diff --git a/source/isaaclab_tasks/test/benchmarking/test_environments_training.py b/source/isaaclab_tasks/test/benchmarking/test_environments_training.py index 1e82728e13bd..8df78522b06a 100644 --- a/source/isaaclab_tasks/test/benchmarking/test_environments_training.py +++ b/source/isaaclab_tasks/test/benchmarking/test_environments_training.py @@ -28,7 +28,7 @@ def setup_environment(): # Acquire all Isaac environments names registered_task_specs = [] for task_spec in gym.registry.values(): - if "Isaac" in task_spec.id and not task_spec.id.endswith("Play-v0"): + if "Isaac" in task_spec.id and utils._is_training_task(task_spec.id): registered_task_specs.append(task_spec) # Sort environments by name diff --git a/source/isaaclab_tasks/test/core/test_dexterous_task_math.py b/source/isaaclab_tasks/test/core/test_dexterous_task_math.py new file mode 100644 index 000000000000..b922111bde93 --- /dev/null +++ b/source/isaaclab_tasks/test/core/test_dexterous_task_math.py @@ -0,0 +1,142 @@ +# 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 + +"""Behavior tests for the dexterous-task math helpers. + +Each test asserts behavior on hand-computed cases (known rotations, distances, +and layouts) so quaternion-layout or formula regressions fail loudly without +re-implementing the tested math as a reference. +""" + +import math + +import pytest +import torch + +import isaaclab.utils.math as math_utils + +from isaaclab_tasks.core.handover.mdp.rewards import evaluate_handover_success, handover_reward +from isaaclab_tasks.core.reorient.mdp.observations import compute_cube_keypoints, cube_keypoints_from_quat +from isaaclab_tasks.core.reorient.mdp.rewards import evaluate_reorient_success + +_DEVICES = ["cpu"] + (["cuda:0"] if torch.cuda.is_available() else []) + +_SIN45 = math.sin(math.pi / 4) +_COS45 = math.cos(math.pi / 4) +# (x, y, z, w) storage everywhere +_IDENTITY = (0.0, 0.0, 0.0, 1.0) +_ROT90_X = (_SIN45, 0.0, 0.0, _COS45) +_ROT180_Z = (0.0, 0.0, 1.0, 0.0) + + +def _quats(device, *quats): + return torch.tensor(quats, dtype=torch.float32, device=device) + + +@pytest.mark.parametrize("device", _DEVICES) +def test_rotation_distance_recovers_known_angles(device): + _, distance = evaluate_reorient_success( + _quats(device, _IDENTITY, _ROT90_X, _ROT180_Z), _quats(device, _IDENTITY, _IDENTITY, _IDENTITY), 0.4 + ) + torch.testing.assert_close(distance, torch.tensor([0.0, math.pi / 2, math.pi], device=device), atol=1e-5, rtol=0.0) + + +@pytest.mark.parametrize("device", _DEVICES) +def test_reorient_success_thresholds_on_rotation_distance(device): + success, distance = evaluate_reorient_success( + _quats(device, _IDENTITY, _ROT90_X), _quats(device, _IDENTITY, _IDENTITY), 0.4 + ) + assert success.tolist() == [True, False] + torch.testing.assert_close(distance, torch.tensor([0.0, math.pi / 2], device=device), atol=1e-5, rtol=0.0) + + +@pytest.mark.parametrize("device", _DEVICES) +def test_goal_quat_error_composes_the_conjugate_product(device): + # error(q, identity) = q; error(q, q) = identity; the 180-degree case exercises w = 0. + # This is the composition the goal_quat_diff and object_goal observation terms rely on. + asset = _quats(device, _ROT90_X, _ROT90_X, _ROT180_Z) + goal = _quats(device, _IDENTITY, _ROT90_X, _IDENTITY) + out = math_utils.quat_mul(asset, math_utils.quat_conjugate(goal)) + torch.testing.assert_close(out[0], torch.tensor(_ROT90_X, device=device), atol=1e-6, rtol=0.0) + torch.testing.assert_close(out[1], torch.tensor(_IDENTITY, device=device), atol=1e-6, rtol=0.0) + # 180-degree rotations keep w = 0, so quat_unique must leave them unchanged + unique = math_utils.quat_unique(out) + torch.testing.assert_close(unique[2].abs(), torch.tensor(_ROT180_Z, device=device).abs(), atol=1e-6, rtol=0.0) + + +@pytest.mark.parametrize("device", _DEVICES) +def test_goal_quat_error_flips_sign_for_negative_real_part(device): + # error(270-degree X rotation, identity) has w = -cos45 < 0, exercising the sign flip + rot270_x = (_SIN45, 0.0, 0.0, -_COS45) + out = math_utils.quat_mul(_quats(device, rot270_x), math_utils.quat_conjugate(_quats(device, _IDENTITY))) + torch.testing.assert_close(out[0], torch.tensor(rot270_x, device=device), atol=1e-6, rtol=0.0) + unique = math_utils.quat_unique(out) + expected = (-_SIN45, 0.0, 0.0, _COS45) + torch.testing.assert_close(unique[0], torch.tensor(expected, device=device), atol=1e-6, rtol=0.0) + + +@pytest.mark.parametrize("device", _DEVICES) +def test_handover_success_measures_env_frame_distance(device): + object_pos = torch.tensor([[0.0, 0.0, 3.0], [0.0, 0.0, 0.0]], device=device) + goal_pos = torch.tensor([[0.0, 0.0, 3.05], [1.0, 0.0, 0.0]], device=device) + success, distance = evaluate_handover_success(object_pos, goal_pos, 0.1) + torch.testing.assert_close(distance, torch.tensor([0.05, 1.0], device=device), atol=1e-6, rtol=0.0) + assert success.tolist() == [True, False] + + +@pytest.mark.parametrize("device", _DEVICES) +def test_handover_reward_falls_off_exponentially(device): + scale = 20.0 + distance = torch.tensor([0.0, math.log(2.0) / scale], device=device) + reward = handover_reward(distance, scale) + # 2 * exp(-scale * d): d = 0 -> 2.0; d = ln(2)/scale -> 1.0 + torch.testing.assert_close(reward, torch.tensor([2.0, 1.0], device=device), atol=1e-6, rtol=1e-6) + + +@pytest.mark.parametrize("device", _DEVICES) +def test_cube_keypoints_identity_pose_gives_half_side_corners(device): + pose = torch.zeros(1, 7, device=device) + pose[0, 6] = 1.0 # identity orientation (x, y, z, w) + + keypoints = compute_cube_keypoints(pose, size=(0.4, 0.6, 0.8)) + + corners = {tuple(round(c, 3) for c in corner) for corner in keypoints[0].tolist()} + expected = {(sx * 0.2, sy * 0.3, sz * 0.4) for sx in (1, -1) for sy in (1, -1) for sz in (1, -1)} + assert corners == expected + + +@pytest.mark.parametrize("device", _DEVICES) +def test_cube_keypoints_write_into_optional_out_buffer(device): + pose = torch.zeros(2, 7, device=device) + pose[:, 6] = 1.0 + out = torch.full((2, 8, 3), torch.nan, dtype=torch.float32, device=device) + result = compute_cube_keypoints(pose, out=out) + assert result is out + assert not out.isnan().any() + torch.testing.assert_close(out, compute_cube_keypoints(pose)) + + +@pytest.mark.parametrize("device", _DEVICES) +def test_goal_keypoints_are_rotation_only_offsets(device): + quat = _quats(device, _IDENTITY) + flattened = cube_keypoints_from_quat(quat, half_size=(0.2, 0.3, 0.4)) + corners = {tuple(round(c, 3) for c in corner) for corner in flattened.view(8, 3).tolist()} + expected = {(sx * 0.2, sy * 0.3, sz * 0.4) for sx in (1, -1) for sy in (1, -1) for sz in (1, -1)} + assert corners == expected + + +@pytest.mark.parametrize("device", _DEVICES) +def test_cube_keypoints_deprecated_shim_delegates_and_warns(device): + import warnings + + from isaaclab_tasks.core.reorient.config.shadow_hand.shadow_hand_direct_camera_env import compute_keypoints + + pose = torch.zeros(2, 7, device=device) + pose[:, 6] = 1.0 + with warnings.catch_warnings(record=True) as caught: + warnings.simplefilter("always") + result = compute_keypoints(pose) + assert any(issubclass(w.category, DeprecationWarning) for w in caught) + torch.testing.assert_close(result, compute_cube_keypoints(pose)) diff --git a/source/isaaclab_tasks/test/core/test_dexterous_value_parity.py b/source/isaaclab_tasks/test/core/test_dexterous_value_parity.py new file mode 100644 index 000000000000..19e4e00ab6dc --- /dev/null +++ b/source/isaaclab_tasks/test/core/test_dexterous_value_parity.py @@ -0,0 +1,115 @@ +# 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 + +"""Direct-vs-manager scalar value parity for the dexterous task families. + +The Direct and manager configurations define their task scalars separately +(per repo convention); these checks catch one side being re-tuned without the +other. Each case maps a Direct cfg field to the manager cfg location that must +carry the same value. +""" + +import pytest + + +def _resolve(obj, path: str): + for part in path.split("."): + obj = obj[part] if isinstance(obj, dict) else getattr(obj, part) + return obj + + +def _reorient_cases(direct_cfg, manager_cfg): + """Field-to-term mapping shared by the Shadow, OpenAI, and Allegro pairs.""" + return [ + (direct_cfg.sim.dt, manager_cfg.sim.dt), + (direct_cfg.sim.render_interval, manager_cfg.sim.render_interval), + (direct_cfg.success_tolerance, manager_cfg.commands.object_pose.orientation_success_threshold), + (direct_cfg.success_tolerance, _resolve(manager_cfg, "rewards.reorient.params.success_tolerance")), + (direct_cfg.dist_reward_scale, _resolve(manager_cfg, "rewards.reorient.params.distance_scale")), + (direct_cfg.rot_reward_scale, _resolve(manager_cfg, "rewards.reorient.params.rotation_scale")), + (direct_cfg.rot_eps, _resolve(manager_cfg, "rewards.reorient.params.rotation_epsilon")), + (direct_cfg.action_penalty_scale, _resolve(manager_cfg, "rewards.reorient.params.action_penalty_scale")), + (direct_cfg.reach_goal_bonus, _resolve(manager_cfg, "rewards.reorient.params.success_bonus")), + (direct_cfg.fall_dist, _resolve(manager_cfg, "rewards.reorient.params.fall_distance")), + (direct_cfg.fall_penalty, _resolve(manager_cfg, "rewards.reorient.params.fall_penalty")), + (direct_cfg.av_factor, _resolve(manager_cfg, "rewards.reorient.params.averaging_factor")), + (direct_cfg.act_moving_average, manager_cfg.actions.joint_pos.alpha), + (direct_cfg.decimation, manager_cfg.decimation), + (direct_cfg.episode_length_s, manager_cfg.episode_length_s), + (direct_cfg.reset_position_noise, _reset_event_params(manager_cfg)["position_noise"]), + (direct_cfg.reset_dof_pos_noise, _reset_event_params(manager_cfg)["joint_position_noise"]), + (direct_cfg.reset_dof_vel_noise, _reset_event_params(manager_cfg)["joint_velocity_noise"]), + (direct_cfg.max_consecutive_success, _timeout_max_successes(manager_cfg)), + ] + + +def _reset_event_params(manager_cfg): + """Params of the reset event term (the OpenAI events are preset-wrapped).""" + events = manager_cfg.events + term = getattr(events, "reset_state", None) or events.default.reset_state + return term.params + + +def _timeout_max_successes(manager_cfg): + """Successes-based timeout threshold; 0 when the manager uses the plain timeout. + + Mirrors the Direct convention where ``max_consecutive_success = 0`` disables + the mechanism (state and Allegro), while the OpenAI variants enable it. + """ + params = getattr(manager_cfg.terminations.time_out, "params", None) or {} + return params.get("max_successes", 0) + + +def _pairs(): + from isaaclab_tasks.core.handover.handover_env_cfg import HandoverEnvCfg + from isaaclab_tasks.core.handover.handover_manager_env_cfg import HandoverManagerEnvCfg + from isaaclab_tasks.core.reorient.config.allegro_hand.allegro_hand_direct_env_cfg import AllegroHandEnvCfg + from isaaclab_tasks.core.reorient.config.allegro_hand.allegro_hand_manager_env_cfg import AllegroCubeEnvCfg + from isaaclab_tasks.core.reorient.config.shadow_hand.shadow_hand_direct_env_cfg import ( + ShadowHandEnvCfg, + ShadowHandOpenAIEnvCfg, + ) + from isaaclab_tasks.core.reorient.config.shadow_hand.shadow_hand_manager_env_cfg import ShadowHandManagerEnvCfg + from isaaclab_tasks.core.reorient.config.shadow_hand.shadow_hand_openai_manager_env_cfg import ( + ShadowHandOpenAIManagerEnvCfg, + ) + + return { + "shadow": (ShadowHandEnvCfg(), ShadowHandManagerEnvCfg()), + "openai": (ShadowHandOpenAIEnvCfg(), ShadowHandOpenAIManagerEnvCfg()), + "allegro": (AllegroHandEnvCfg(), AllegroCubeEnvCfg()), + "handover": (HandoverEnvCfg(), HandoverManagerEnvCfg()), + } + + +@pytest.mark.parametrize("family", ["shadow", "openai", "allegro"]) +def test_reorient_direct_manager_scalars_match(family): + """Direct cfg scalars equal the manager term params they mirror.""" + direct_cfg, manager_cfg = _pairs()[family] + for i, (direct_value, manager_value) in enumerate(_reorient_cases(direct_cfg, manager_cfg)): + assert direct_value == manager_value, f"{family} case {i}: direct={direct_value} manager={manager_value}" + + +def test_handover_direct_manager_scalars_match(): + """Handover Direct cfg scalars equal the manager term params they mirror.""" + direct_cfg, manager_cfg = _pairs()["handover"] + obs = manager_cfg.observations + cases = [ + (direct_cfg.sim.dt, manager_cfg.sim.dt), + (direct_cfg.sim.render_interval, manager_cfg.sim.render_interval), + (direct_cfg.dist_reward_scale, _resolve(manager_cfg, "rewards.handover.params.distance_scale")), + ( + direct_cfg.success_distance_threshold, + _resolve(manager_cfg, "rewards.handover.params.success_distance_threshold"), + ), + (direct_cfg.vel_obs_scale, _resolve(obs, "policy.right_object_goal.params.vel_obs_scale")), + (direct_cfg.vel_obs_scale, _resolve(obs, "policy.left_object_goal.params.vel_obs_scale")), + (direct_cfg.act_moving_average, manager_cfg.actions.right_hand.alpha), + (direct_cfg.act_moving_average, manager_cfg.actions.left_hand.alpha), + (direct_cfg.decimation, manager_cfg.decimation), + (direct_cfg.episode_length_s, manager_cfg.episode_length_s), + ] + for i, (direct_value, manager_value) in enumerate(cases): + assert direct_value == manager_value, f"handover case {i}: direct={direct_value} manager={manager_value}" diff --git a/source/isaaclab_tasks/test/core/test_rendering_registered_tasks.py b/source/isaaclab_tasks/test/core/test_rendering_registered_tasks.py index bb3f80d1eb2d..1eb4692c40a6 100644 --- a/source/isaaclab_tasks/test/core/test_rendering_registered_tasks.py +++ b/source/isaaclab_tasks/test/core/test_rendering_registered_tasks.py @@ -82,6 +82,12 @@ def _collect_camera_outputs(env: object) -> dict[str, dict[str, torch.Tensor]]: # require at least one pass while we tighten the validation tolerances for this scene. marks=pytest.mark.flaky(max_runs=3, min_passes=1), ), + pytest.param( + "Isaac-Reorient-Cube-Shadow-Camera", + None, + "shadow_hand", + marks=pytest.mark.flaky(max_runs=3, min_passes=1), + ), ] diff --git a/source/isaaclab_tasks/test/core/test_shadow_hand_camera_presets.py b/source/isaaclab_tasks/test/core/test_shadow_hand_camera_presets.py index 828c06fc0d1e..5700b9d86b99 100644 --- a/source/isaaclab_tasks/test/core/test_shadow_hand_camera_presets.py +++ b/source/isaaclab_tasks/test/core/test_shadow_hand_camera_presets.py @@ -36,7 +36,7 @@ from isaaclab.renderers import RendererCfg # noqa: E402 -from isaaclab_tasks.core.reorient.config.shadow_hand.shadow_hand_camera_env_cfg import ( # noqa: E402 +from isaaclab_tasks.core.reorient.config.shadow_hand.shadow_hand_direct_camera_env_cfg import ( # noqa: E402 ShadowHandCameraEnvCfg, ) from isaaclab_tasks.utils.hydra import collect_presets # noqa: E402 @@ -185,6 +185,7 @@ def shadow_hand_camera_presets(): ("default", ["rgb", "depth", "semantic_segmentation"]), ("full", ["rgb", "depth", "semantic_segmentation"]), ("rgb", ["rgb"]), + ("rgb_depth", ["rgb", "depth"]), ("albedo", ["albedo"]), ("simple_shading_constant_diffuse", ["simple_shading_constant_diffuse"]), ("simple_shading_diffuse_mdl", ["simple_shading_diffuse_mdl"]), @@ -277,7 +278,7 @@ def test_all_renderer_presets_present(shadow_hand_camera_presets): # when paired with the warp renderer preset # --------------------------------------------------------------------------- -_WARP_VALID_CAMERA_PRESETS = ["rgb", "depth"] +_WARP_VALID_CAMERA_PRESETS = ["rgb", "rgb_depth", "depth"] _WARP_INVALID_CAMERA_PRESETS = [ "default", "full", @@ -290,7 +291,7 @@ def test_all_renderer_presets_present(shadow_hand_camera_presets): @pytest.mark.parametrize("camera_preset", _WARP_VALID_CAMERA_PRESETS) def test_warp_with_valid_camera_preset(shadow_hand_camera_presets, camera_preset): - """Warp + {rgb, depth} camera presets must not raise (depth with CNN disabled).""" + """Warp + {rgb, rgb_depth, depth} camera presets must not raise (depth with CNN disabled).""" camera_cfg = shadow_hand_camera_presets["tiled_camera"][camera_preset] warp_cfg = shadow_hand_camera_presets["tiled_camera.renderer_cfg"]["newton_renderer"] enabled = camera_cfg.data_types != ["depth"] # disable CNN for depth-only diff --git a/source/isaaclab_tasks/test/golden_images/registered_tasks/Isaac-Reorient-Cube-Shadow-Camera-Direct/default_physics-default_renderer-depth.png b/source/isaaclab_tasks/test/golden_images/registered_tasks/Isaac-Reorient-Cube-Shadow-Camera-Direct/default_physics-default_renderer-depth.png new file mode 100644 index 000000000000..c229b583dfb7 --- /dev/null +++ b/source/isaaclab_tasks/test/golden_images/registered_tasks/Isaac-Reorient-Cube-Shadow-Camera-Direct/default_physics-default_renderer-depth.png @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:a37f6bca30bb2d093eb68186c601551d52aafe8ed19c6c090de149b3210d81a5 +size 3665 diff --git a/source/isaaclab_tasks/test/golden_images/registered_tasks/Isaac-Reorient-Cube-Shadow-Camera-Direct/default_physics-default_renderer-rgb.png b/source/isaaclab_tasks/test/golden_images/registered_tasks/Isaac-Reorient-Cube-Shadow-Camera-Direct/default_physics-default_renderer-rgb.png new file mode 100644 index 000000000000..472bbe6e9db0 --- /dev/null +++ b/source/isaaclab_tasks/test/golden_images/registered_tasks/Isaac-Reorient-Cube-Shadow-Camera-Direct/default_physics-default_renderer-rgb.png @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:bc17ca40a050eb357a607326ccb4fc553cb525abb3f3fa96d7b496f5dc51c93e +size 19878 diff --git a/source/isaaclab_tasks/test/golden_images/registered_tasks/Isaac-Reorient-Cube-Shadow-Camera-Direct/default_physics-default_renderer-rgba.png b/source/isaaclab_tasks/test/golden_images/registered_tasks/Isaac-Reorient-Cube-Shadow-Camera-Direct/default_physics-default_renderer-rgba.png new file mode 100644 index 000000000000..885109b2f728 --- /dev/null +++ b/source/isaaclab_tasks/test/golden_images/registered_tasks/Isaac-Reorient-Cube-Shadow-Camera-Direct/default_physics-default_renderer-rgba.png @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:5aa01e4ef3a1dbd299fbbc31d47c0db785c0db99252236376470d547fac4f509 +size 22056 diff --git a/source/isaaclab_tasks/test/golden_images/registered_tasks/Isaac-Reorient-Cube-Shadow-Camera-Direct/default_physics-default_renderer-semantic_segmentation.png b/source/isaaclab_tasks/test/golden_images/registered_tasks/Isaac-Reorient-Cube-Shadow-Camera-Direct/default_physics-default_renderer-semantic_segmentation.png new file mode 100644 index 000000000000..4bad29d72ce6 --- /dev/null +++ b/source/isaaclab_tasks/test/golden_images/registered_tasks/Isaac-Reorient-Cube-Shadow-Camera-Direct/default_physics-default_renderer-semantic_segmentation.png @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:75576f31118081f96b0cec2151ada7016794a933a035ea35665a752e8552b503 +size 1474 diff --git a/source/isaaclab_tasks/test/golden_images/registered_tasks/Isaac-Reorient-Cube-Shadow-Camera/default_physics-default_renderer-depth.png b/source/isaaclab_tasks/test/golden_images/registered_tasks/Isaac-Reorient-Cube-Shadow-Camera/default_physics-default_renderer-depth.png new file mode 100644 index 000000000000..c229b583dfb7 --- /dev/null +++ b/source/isaaclab_tasks/test/golden_images/registered_tasks/Isaac-Reorient-Cube-Shadow-Camera/default_physics-default_renderer-depth.png @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:a37f6bca30bb2d093eb68186c601551d52aafe8ed19c6c090de149b3210d81a5 +size 3665 diff --git a/source/isaaclab_tasks/test/golden_images/registered_tasks/Isaac-Reorient-Cube-Shadow-Camera/default_physics-default_renderer-rgb.png b/source/isaaclab_tasks/test/golden_images/registered_tasks/Isaac-Reorient-Cube-Shadow-Camera/default_physics-default_renderer-rgb.png new file mode 100644 index 000000000000..eace991f49eb --- /dev/null +++ b/source/isaaclab_tasks/test/golden_images/registered_tasks/Isaac-Reorient-Cube-Shadow-Camera/default_physics-default_renderer-rgb.png @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:b60263834743f0508a437281d36bd2f46296789b28ad930627fb72e406ab8700 +size 19962 diff --git a/source/isaaclab_tasks/test/golden_images/registered_tasks/Isaac-Reorient-Cube-Shadow-Camera/default_physics-default_renderer-rgba.png b/source/isaaclab_tasks/test/golden_images/registered_tasks/Isaac-Reorient-Cube-Shadow-Camera/default_physics-default_renderer-rgba.png new file mode 100644 index 000000000000..7d6fa735693d --- /dev/null +++ b/source/isaaclab_tasks/test/golden_images/registered_tasks/Isaac-Reorient-Cube-Shadow-Camera/default_physics-default_renderer-rgba.png @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:8da49db87ab2afe4f2a32c9354a72fd1b5a32aaf76484a21f787b01c0a47197a +size 22127 diff --git a/source/isaaclab_tasks/test/golden_images/registered_tasks/Isaac-Reorient-Cube-Shadow-Camera/default_physics-default_renderer-semantic_segmentation.png b/source/isaaclab_tasks/test/golden_images/registered_tasks/Isaac-Reorient-Cube-Shadow-Camera/default_physics-default_renderer-semantic_segmentation.png new file mode 100644 index 000000000000..4bad29d72ce6 --- /dev/null +++ b/source/isaaclab_tasks/test/golden_images/registered_tasks/Isaac-Reorient-Cube-Shadow-Camera/default_physics-default_renderer-semantic_segmentation.png @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:75576f31118081f96b0cec2151ada7016794a933a035ea35665a752e8552b503 +size 1474 diff --git a/source/isaaclab_tasks/test/rendering_test_utils.py b/source/isaaclab_tasks/test/rendering_test_utils.py index 452335b1e15f..a6a532297f88 100644 --- a/source/isaaclab_tasks/test/rendering_test_utils.py +++ b/source/isaaclab_tasks/test/rendering_test_utils.py @@ -1107,21 +1107,24 @@ def rendering_test_shadow_hand( from isaaclab.utils.configclass import configclass - from isaaclab_tasks.core.reorient.config.shadow_hand.shadow_hand_camera_env import ShadowHandCameraEnv - from isaaclab_tasks.core.reorient.config.shadow_hand.shadow_hand_camera_env_cfg import ( + from isaaclab_tasks.core.reorient.config.shadow_hand.shadow_hand_direct_camera_env import ShadowHandCameraEnv + from isaaclab_tasks.core.reorient.config.shadow_hand.shadow_hand_direct_camera_env_cfg import ( ShadowHandCameraEnvCfg, ShadowHandTiledCameraCfg, - _ShadowHandBaseTiledCameraCfg, ) @configclass class _ShadowHandTiledCameraTestCfg(ShadowHandTiledCameraCfg): - distance_to_camera = _ShadowHandBaseTiledCameraCfg(data_types=["distance_to_camera"]) - distance_to_image_plane = _ShadowHandBaseTiledCameraCfg(data_types=["distance_to_image_plane"]) - normals = _ShadowHandBaseTiledCameraCfg(data_types=["normals"]) - instance_segmentation_fast = _ShadowHandBaseTiledCameraCfg(data_types=["instance_segmentation_fast"]) - instance_id_segmentation_fast = _ShadowHandBaseTiledCameraCfg(data_types=["instance_id_segmentation_fast"]) - motion_vectors = _ShadowHandBaseTiledCameraCfg(data_types=["motion_vectors"]) + distance_to_camera = ShadowHandTiledCameraCfg.BaseTiledCameraCfg(data_types=["distance_to_camera"]) + distance_to_image_plane = ShadowHandTiledCameraCfg.BaseTiledCameraCfg(data_types=["distance_to_image_plane"]) + normals = ShadowHandTiledCameraCfg.BaseTiledCameraCfg(data_types=["normals"]) + instance_segmentation_fast = ShadowHandTiledCameraCfg.BaseTiledCameraCfg( + data_types=["instance_segmentation_fast"] + ) + instance_id_segmentation_fast = ShadowHandTiledCameraCfg.BaseTiledCameraCfg( + data_types=["instance_id_segmentation_fast"] + ) + motion_vectors = ShadowHandTiledCameraCfg.BaseTiledCameraCfg(data_types=["motion_vectors"]) @configclass class _ShadowHandCameraTestEnvCfg(ShadowHandCameraEnvCfg): diff --git a/source/isaaclab_tasks_experimental/changelog.d/task-cleanup-dex-warp-exp.minor.rst b/source/isaaclab_tasks_experimental/changelog.d/task-cleanup-dex-warp-exp.minor.rst new file mode 100644 index 000000000000..b1d6aeb124f3 --- /dev/null +++ b/source/isaaclab_tasks_experimental/changelog.d/task-cleanup-dex-warp-exp.minor.rst @@ -0,0 +1,10 @@ +Added +^^^^^ + +* Added warp-first Direct variants of the Shadow cube reorientation and handover + tasks (``Isaac-Reorient-Cube-Shadow-Direct-Warp-v0``, + ``Isaac-Reorient-Cube-Shadow-OpenAI-FF-Direct-Warp-v0``, + ``Isaac-Reorient-Cube-Shadow-OpenAI-LSTM-Direct-Warp-v0``, and + ``Isaac-Shadow-Handover-Direct-Warp-v0``). The environments compute rewards, + observations, and resets in Warp kernels and reuse the stable configurations + and agents from :mod:`isaaclab_tasks.core`. diff --git a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/direct/handover/__init__.py b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/direct/handover/__init__.py new file mode 100644 index 000000000000..a7dea11a8da1 --- /dev/null +++ b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/direct/handover/__init__.py @@ -0,0 +1,34 @@ +# 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 + +""" +Warp-first Direct variant of the Shadow handover task. + +The environment reuses the stable configuration and agent definitions from +:mod:`isaaclab_tasks.core.handover`. +""" + +import gymnasium as gym + +## +# Register Gym environments. +## + +stable_handover = "isaaclab_tasks.core.handover" +stable_handover_agents = f"{stable_handover}.agents" + +gym.register( + id="Isaac-Shadow-Handover-Direct-Warp-v0", + entry_point="isaaclab_tasks_experimental.direct.handover.handover_warp_env:HandoverWarpEnv", + disable_env_checker=True, + kwargs={ + "env_cfg_entry_point": f"{stable_handover}.handover_env_cfg:HandoverEnvCfg", + "rsl_rl_cfg_entry_point": f"{stable_handover_agents}.rsl_rl_ppo_cfg:HandoverPPORunnerCfg", + "rl_games_cfg_entry_point": f"{stable_handover_agents}:rl_games_ppo_cfg.yaml", + "skrl_cfg_entry_point": f"{stable_handover_agents}:skrl_ppo_cfg.yaml", + "skrl_ippo_cfg_entry_point": f"{stable_handover_agents}:skrl_ippo_cfg.yaml", + "skrl_mappo_cfg_entry_point": f"{stable_handover_agents}:skrl_mappo_cfg.yaml", + }, +) diff --git a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/direct/handover/handover_kernels.py b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/direct/handover/handover_kernels.py new file mode 100644 index 000000000000..db8c00d2621f --- /dev/null +++ b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/direct/handover/handover_kernels.py @@ -0,0 +1,154 @@ +# 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 + +"""Warp kernels for the handover task family. + +This is the family's math root: the experimental Direct environment launches +these kernels on the environment's Warp-native state (ProxyArray ``.warp`` +views) and its own pre-allocated output buffers. Shared hand idioms (fingertip +segment readers and the quaternion error) come from the reorientation family's +math root; the +reverse dependency never exists. +""" + +from __future__ import annotations + +import warp as wp + +from isaaclab_tasks_experimental.direct.reorient.reorient_kernels import ( + fingertip_pos_col, + fingertip_quat_col, + fingertip_vel_col, +) + + +@wp.kernel +def fall_kernel( + object_pos_w: wp.array(dtype=wp.vec3f), + env_origins: wp.array(dtype=wp.vec3f), + fall_height: float, + fell: wp.array(dtype=wp.bool), +): + i = wp.tid() + fell[i] = (object_pos_w[i][2] - env_origins[i][2]) <= fall_height + + +@wp.kernel +def hand_proprio_kernel( + joint_pos: wp.array2d(dtype=wp.float32), + joint_vel: wp.array2d(dtype=wp.float32), + lower: wp.array2d(dtype=wp.float32), + upper: wp.array2d(dtype=wp.float32), + vel_scale: float, + body_pos_w: wp.array2d(dtype=wp.vec3f), + body_quat_w: wp.array2d(dtype=wp.quatf), + body_vel_w: wp.array2d(dtype=wp.spatial_vectorf), + finger_ids: wp.array(dtype=wp.int32), + env_origins: wp.array(dtype=wp.vec3f), + actions: wp.array2d(dtype=wp.float32), + out: wp.array2d(dtype=wp.float32), +): + """Per-hand proprioceptive block, matching the torch concatenation order. + + Launched over ``(num_envs, obs_dim)``: each thread walks a branch ladder over the + segment boundaries and writes one output column, so warps (32 consecutive columns) + stay branch-uniform except at segment boundaries. + """ + i, j = wp.tid() + num_joints = joint_pos.shape[1] + num_fingers = finger_ids.shape[0] + # segment boundaries, in column order + end_joint = 2 * num_joints + end_tip_pos = end_joint + 3 * num_fingers + end_tip_quat = end_tip_pos + 4 * num_fingers + end_tip_vel = end_tip_quat + 6 * num_fingers + # hand: normalized DOF positions, scaled DOF velocities + if j < num_joints: + out[i, j] = 2.0 * (joint_pos[i, j] - lower[i, j]) / (upper[i, j] - lower[i, j]) - 1.0 + elif j < end_joint: + out[i, j] = vel_scale * joint_vel[i, j - num_joints] + # fingertips: environment-frame positions, rotations, spatial velocities + elif j < end_tip_pos: + out[i, j] = fingertip_pos_col(body_pos_w, env_origins, finger_ids, i, j - end_joint) + elif j < end_tip_quat: + out[i, j] = fingertip_quat_col(body_quat_w, finger_ids, i, j - end_tip_pos) + elif j < end_tip_vel: + out[i, j] = fingertip_vel_col(body_vel_w, finger_ids, i, j - end_tip_quat) + # actions + else: + out[i, j] = actions[i, j - end_tip_vel] + + +@wp.kernel +def handover_reward_kernel( + goal_distance: wp.array(dtype=wp.float32), + distance_scale: float, + reward: wp.array(dtype=wp.float32), +): + i = wp.tid() + # Direct handover reward: exponential falloff with object-goal distance + reward[i] = 2.0 * wp.exp(-distance_scale * goal_distance[i]) + + +@wp.kernel +def handover_success_kernel( + object_pos_w: wp.array(dtype=wp.vec3f), + env_origins: wp.array(dtype=wp.vec3f), + goal_pos_e: wp.array(dtype=wp.vec3f), + success_distance_threshold: float, + success: wp.array(dtype=wp.bool), + goal_distance: wp.array(dtype=wp.float32), +): + i = wp.tid() + # object position is world-frame; the goal is authored in the environment frame + distance = wp.length(object_pos_w[i] - env_origins[i] - goal_pos_e[i]) + goal_distance[i] = distance + success[i] = distance < success_distance_threshold + + +@wp.kernel +def object_goal_kernel( + object_pos: wp.array(dtype=wp.vec3f), + env_origins: wp.array(dtype=wp.vec3f), + object_quat: wp.array(dtype=wp.quatf), + lin_vel: wp.array(dtype=wp.vec3f), + ang_vel: wp.array(dtype=wp.vec3f), + goal_pos: wp.array(dtype=wp.vec3f), + goal_quat: wp.array(dtype=wp.quatf), + ang_vel_scale: float, + out: wp.array2d(dtype=wp.float32), +): + """Assemble one environment's 24-dimensional object/goal observation row.""" + i = wp.tid() + pos = object_pos[i] - env_origins[i] + q1 = object_quat[i] + q2 = goal_quat[i] + # quat_inverse == conjugate for these unit quaternions, matching + # isaaclab.utils.math.quat_mul/quat_conjugate semantics + qe = q1 * wp.quat_inverse(q2) + out[i, 0] = pos[0] + out[i, 1] = pos[1] + out[i, 2] = pos[2] + out[i, 3] = q1[0] + out[i, 4] = q1[1] + out[i, 5] = q1[2] + out[i, 6] = q1[3] + out[i, 7] = lin_vel[i][0] + out[i, 8] = lin_vel[i][1] + out[i, 9] = lin_vel[i][2] + out[i, 10] = ang_vel_scale * ang_vel[i][0] + out[i, 11] = ang_vel_scale * ang_vel[i][1] + out[i, 12] = ang_vel_scale * ang_vel[i][2] + out[i, 13] = goal_pos[i][0] + out[i, 14] = goal_pos[i][1] + out[i, 15] = goal_pos[i][2] + out[i, 16] = q2[0] + out[i, 17] = q2[1] + out[i, 18] = q2[2] + out[i, 19] = q2[3] + out[i, 20] = qe[0] + out[i, 21] = qe[1] + out[i, 22] = qe[2] + out[i, 23] = qe[3] diff --git a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/direct/handover/handover_warp_env.py b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/direct/handover/handover_warp_env.py new file mode 100644 index 000000000000..75b0b35be297 --- /dev/null +++ b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/direct/handover/handover_warp_env.py @@ -0,0 +1,381 @@ +# 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 + +"""Warp-first Direct environment for the Shadow handover task. + +Experimental counterpart of the torch-first mainline implementation in +:mod:`isaaclab_tasks.core.handover`; see the reorient variant for the conventions. +""" + +from __future__ import annotations + +from collections.abc import Sequence + +import torch +import warp as wp + +import isaaclab.sim as sim_utils +from isaaclab import cloner +from isaaclab.assets import Articulation, RigidObject +from isaaclab.envs import DirectMARLEnv +from isaaclab.markers import VisualizationMarkers +from isaaclab.sim.spawners.from_files import GroundPlaneCfg, spawn_ground_plane +from isaaclab.utils.math import sample_uniform + +from isaaclab_tasks.core.handover.handover_env_cfg import HandoverEnvCfg +from isaaclab_tasks.core.handover.handover_task_base import GOAL_POSITION_OFFSET +from isaaclab_tasks.core.utils import EpisodeErrorRecorder, randomize_rotation, sample_joint_positions_within_limits + +from isaaclab_tasks_experimental.direct.handover.handover_kernels import ( + fall_kernel, + hand_proprio_kernel, + handover_reward_kernel, + handover_success_kernel, + object_goal_kernel, +) +from isaaclab_tasks_experimental.direct.reorient.reorient_kernels import ema_actuation_kernel + + +class HandoverWarpEnv(DirectMARLEnv): + cfg: HandoverEnvCfg + + def __init__(self, cfg: HandoverEnvCfg, render_mode: str | None = None, **kwargs): + super().__init__(cfg, render_mode, **kwargs) + + self.num_hand_dofs = self.right_hand.num_joints + + # buffers for position targets + self.right_hand_prev_targets = torch.zeros( + (self.num_envs, self.num_hand_dofs), dtype=torch.float, device=self.device + ) + self.right_hand_curr_targets = torch.zeros( + (self.num_envs, self.num_hand_dofs), dtype=torch.float, device=self.device + ) + self.left_hand_prev_targets = torch.zeros( + (self.num_envs, self.num_hand_dofs), dtype=torch.float, device=self.device + ) + self.left_hand_curr_targets = torch.zeros( + (self.num_envs, self.num_hand_dofs), dtype=torch.float, device=self.device + ) + + # list of actuated joints + self.actuated_dof_indices, _ = self.right_hand.find_joints(cfg.actuated_joint_names) + if len(self.actuated_dof_indices) != len(cfg.actuated_joint_names): + raise ValueError( + f"Expected {len(cfg.actuated_joint_names)} actuated joints, found {len(self.actuated_dof_indices)}." + ) + + # finger bodies + self.finger_bodies, _ = self.right_hand.find_bodies(self.cfg.fingertip_body_names) + if len(self.finger_bodies) != len(self.cfg.fingertip_body_names): + raise ValueError( + f"Expected {len(self.cfg.fingertip_body_names)} fingertip bodies, found {len(self.finger_bodies)}." + ) + self.num_fingertips = len(self.finger_bodies) + + # joint limits + joint_pos_limits = self.right_hand.data.joint_limits.torch.to(self.device) + self.hand_dof_lower_limits = joint_pos_limits[..., 0] + self.hand_dof_upper_limits = joint_pos_limits[..., 1] + + # default goal positions + self.goal_rot = torch.zeros((self.num_envs, 4), dtype=torch.float, device=self.device) + self.goal_rot[:, 3] = 1.0 # identity quaternion in (x, y, z, w) layout + self.goal_pos = torch.zeros((self.num_envs, 3), dtype=torch.float, device=self.device) + # goal = object default position + shared offset (mirrors HandoverCommand.__init__) + self.goal_pos[:, :] = self.object.data.default_root_pose.torch[:, :3].to(self.device) + torch.tensor( + GOAL_POSITION_OFFSET, dtype=torch.float, device=self.device + ) + # initialize goal marker + self.goal_markers = VisualizationMarkers(self.cfg.goal_object_cfg) + + # Sticky per-env flag: True once the object reached the goal within threshold. + self._episode_succeeded = torch.zeros(self.num_envs, dtype=torch.bool, device=self.device) + self._goal_distance = EpisodeErrorRecorder(self.num_envs, self.device) + self._success_flags = torch.empty(self.num_envs, dtype=torch.bool, device=self.device) + self._goal_distance_buf = torch.empty(self.num_envs, dtype=torch.float32, device=self.device) + self._reward_buf = torch.empty(self.num_envs, dtype=torch.float32, device=self.device) + # cached Warp views of the buffers above; the hot loop launches kernels without conversions + self._success_flags_wp = wp.from_torch(self._success_flags) + self._goal_distance_buf_wp = wp.from_torch(self._goal_distance_buf) + self._reward_buf_wp = wp.from_torch(self._reward_buf) + self._env_origins_wp = wp.from_torch(self.scene.env_origins, dtype=wp.vec3f) + self._goal_pos_wp = wp.from_torch(self.goal_pos, dtype=wp.vec3f) + self._goal_rot_wp = wp.from_torch(self.goal_rot, dtype=wp.quatf) + self._object_goal_buf = torch.empty(self.num_envs, 24, dtype=torch.float32, device=self.device) + self._object_goal_buf_wp = wp.from_torch(self._object_goal_buf) + self._lower_limits_wp = wp.from_torch(self.hand_dof_lower_limits) + self._upper_limits_wp = wp.from_torch(self.hand_dof_upper_limits) + self._actuated_dof_ids_wp = wp.array(self.actuated_dof_indices, dtype=wp.int32, device=str(self.device)) + self._finger_ids_wp = wp.array(self.finger_bodies, dtype=wp.int32, device=str(self.device)) + self._hand_targets_wp = { + "right_hand": (wp.from_torch(self.right_hand_prev_targets), wp.from_torch(self.right_hand_curr_targets)), + "left_hand": (wp.from_torch(self.left_hand_prev_targets), wp.from_torch(self.left_hand_curr_targets)), + } + self._compact_targets = torch.zeros( + (self.num_envs, len(self.actuated_dof_indices)), dtype=torch.float32, device=self.device + ) + self._compact_targets_wp = wp.from_torch(self._compact_targets) + self._fell_flags = torch.empty(self.num_envs, dtype=torch.bool, device=self.device) + self._fell_flags_wp = wp.from_torch(self._fell_flags) + num_joints, num_fingers = self.num_hand_dofs, self.num_fingertips + proprio_dim = 2 * num_joints + 13 * num_fingers + self.cfg.action_spaces["right_hand"] + self._proprio_bufs = { + agent: torch.empty(self.num_envs, proprio_dim, device=self.device) for agent in ("right_hand", "left_hand") + } + self._proprio_bufs_wp = {agent: wp.from_torch(buf) for agent, buf in self._proprio_bufs.items()} + + # unit tensors for sampling goal/object rotations about the x and y axes + self.x_unit_tensor = torch.tensor([1, 0, 0], dtype=torch.float, device=self.device).repeat((self.num_envs, 1)) + self.y_unit_tensor = torch.tensor([0, 1, 0], dtype=torch.float, device=self.device).repeat((self.num_envs, 1)) + + def _setup_scene(self): + # add hand, in-hand object, and goal object + self.right_hand = Articulation(self.cfg.right_robot_cfg) + self.left_hand = Articulation(self.cfg.left_robot_cfg) + self.object = RigidObject(self.cfg.object_cfg) + # add ground plane + spawn_ground_plane(prim_path="/World/ground", cfg=GroundPlaneCfg()) + src, dest = "/World/envs/env_0", "/World/envs/env_{}" + pos = cloner.grid_transforms(self.scene.num_envs, self.scene.cfg.env_spacing, device=self.device)[0] + plan = cloner.ClonePlan.from_env_0(src, dest, self.scene.num_envs, self.device, pos) + cloner.replicate(plan, stage=self.scene.stage) + # add articulation to scene - we must register to scene to randomize with EventManager + self.scene.articulations["right_robot"] = self.right_hand + self.scene.articulations["left_robot"] = self.left_hand + self.scene.rigid_objects["object"] = self.object + # add lights + light_cfg = sim_utils.DomeLightCfg(intensity=2000.0, color=(0.75, 0.75, 0.75)) + light_cfg.func("/World/Light", light_cfg) + + def _pre_physics_step(self, actions: dict[str, torch.Tensor]) -> None: + self.actions = actions + + def _apply_action(self) -> None: + self._apply_hand_action(self.right_hand, "right_hand") + self._apply_hand_action(self.left_hand, "left_hand") + + def _apply_hand_action(self, hand: Articulation, agent: str) -> None: + """Map one agent's actions to joint position targets and write them to its hand. + + The raw ``[-1, 1]`` action is rescaled to the joint limits, blended with the previous + target via the exponential moving average, clamped to the limits, and set on the hand. + """ + prev_wp, curr_wp = self._hand_targets_wp[agent] + wp.launch( + ema_actuation_kernel, + dim=(self.num_envs, self._actuated_dof_ids_wp.shape[0]), + inputs=[ + wp.from_torch(self.actions[agent]), + self._lower_limits_wp, + self._upper_limits_wp, + self._actuated_dof_ids_wp, + self.cfg.act_moving_average, + prev_wp, + curr_wp, + ], + outputs=[self._compact_targets_wp], + device=self._compact_targets_wp.device, + ) + hand.set_joint_position_target_index(target=self._compact_targets, joint_ids=self.actuated_dof_indices) + + def _hand_proprio_obs(self, agent: str) -> torch.Tensor: + """Per-hand proprioceptive observation block for ``agent`` (133 dims). + + Layout: normalized DOF positions (24), scaled DOF velocities (24), fingertip positions + (5*3), rotations (5*4), linear+angular velocities (5*6), and the applied actions (20). + """ + hand = self.right_hand if agent == "right_hand" else self.left_hand + wp.launch( + hand_proprio_kernel, + dim=(self.num_envs, self._proprio_bufs_wp[agent].shape[1]), + inputs=[ + hand.data.joint_pos.warp, + hand.data.joint_vel.warp, + self._lower_limits_wp, + self._upper_limits_wp, + self.cfg.vel_obs_scale, + hand.data.body_pos_w.warp, + hand.data.body_quat_w.warp, + hand.data.body_vel_w.warp, + self._finger_ids_wp, + self._env_origins_wp, + wp.from_torch(self.actions[agent]), + ], + outputs=[self._proprio_bufs_wp[agent]], + device=self._proprio_bufs_wp[agent].device, + ) + return self._proprio_bufs[agent] + + def _object_goal_obs(self) -> torch.Tensor: + """Object and goal observation block shared by both agents and the critic state (24 dims). + + Layout: object position (3), rotation (4), linear velocity (3), scaled angular velocity (3), + goal position (3), goal rotation (4), and the goal-to-object rotation difference (4). + """ + wp.launch( + object_goal_kernel, + dim=self.num_envs, + inputs=[ + self.object.data.root_pos_w.warp, + self._env_origins_wp, + self.object.data.root_quat_w.warp, + self.object.data.root_lin_vel_w.warp, + self.object.data.root_ang_vel_w.warp, + self._goal_pos_wp, + self._goal_rot_wp, + self.cfg.vel_obs_scale, + ], + outputs=[self._object_goal_buf_wp], + device=self._object_goal_buf_wp.device, + ) + return self._object_goal_buf + + def _get_observations(self) -> dict[str, torch.Tensor]: + object_goal = self._object_goal_obs() + return { + "right_hand": torch.cat((self._hand_proprio_obs("right_hand"), object_goal), dim=-1), + "left_hand": torch.cat((self._hand_proprio_obs("left_hand"), object_goal), dim=-1), + } + + def _get_states(self) -> torch.Tensor: + # DirectMARLEnv.step() and reset() always run _get_observations() immediately before + # any state() call with no sim advance in between, so the proprio and object-goal + # buffers already hold this step's values; reuse them instead of relaunching kernels + return torch.cat( + (self._proprio_bufs["right_hand"], self._proprio_bufs["left_hand"], self._object_goal_buf), + dim=-1, + ) + + def _get_rewards(self) -> dict[str, torch.Tensor]: + # compute reward + wp.launch( + handover_success_kernel, + dim=self.num_envs, + inputs=[ + self.object.data.root_pos_w.warp, + self._env_origins_wp, + self._goal_pos_wp, + self.cfg.success_distance_threshold, + ], + outputs=[self._success_flags_wp, self._goal_distance_buf_wp], + device=self._success_flags_wp.device, + ) + succeeded, goal_dist = self._success_flags, self._goal_distance_buf + self._goal_distance.update(goal_dist) + wp.launch( + handover_reward_kernel, + dim=self.num_envs, + inputs=[self._goal_distance_buf_wp, self.cfg.dist_reward_scale], + outputs=[self._reward_buf_wp], + device=self._reward_buf_wp.device, + ) + rew_dist = self._reward_buf + + # log reward components + if "log" not in self.extras: + self.extras["log"] = dict() + self.extras["log"]["dist_reward"] = rew_dist.mean() + # tensors, not .item(): a host sync every step stalls the GPU at large env counts + goal_dist_mean = goal_dist.mean() + self.extras["log"]["dist_goal"] = goal_dist_mean + self.extras["log"]["Metrics/goal_distance"] = goal_dist_mean + # Sticky per-env success: True once the object reached the goal within threshold. + self._episode_succeeded |= succeeded + + return {"right_hand": rew_dist, "left_hand": rew_dist} + + def _get_dones(self) -> tuple[dict[str, torch.Tensor], dict[str, torch.Tensor]]: + # reset when object has fallen + wp.launch( + fall_kernel, + dim=self.num_envs, + inputs=[self.object.data.root_pos_w.warp, self._env_origins_wp, self.cfg.fall_dist], + outputs=[self._fell_flags_wp], + device=self._fell_flags_wp.device, + ) + out_of_reach = self._fell_flags + # reset when episode ends + time_out = self.episode_length_buf >= self.max_episode_length - 1 + + terminated = {agent: out_of_reach for agent in self.cfg.possible_agents} + time_outs = {agent: time_out for agent in self.cfg.possible_agents} + return terminated, time_outs + + def _reset_idx(self, env_ids: Sequence[int] | torch.Tensor | None): + if env_ids is None: + env_ids = self.right_hand._ALL_INDICES + # Flush per-episode success (sticky binary: object ever reached the goal within threshold). + self.extras.setdefault("log", {})["Metrics/success_rate"] = ( + self._episode_succeeded[env_ids].float().mean().item() + ) + for statistic, value in self._goal_distance.reset(env_ids).items(): + self.extras["log"][f"Diagnostics/episode_min_goal_distance_{statistic}"] = value + self._episode_succeeded[env_ids] = False + # reset articulation and rigid body attributes + super()._reset_idx(env_ids) + + # reset goals + self._reset_target_pose(env_ids) + + # reset object + object_default_pose = self.object.data.default_root_pose.torch.clone()[env_ids] + object_default_vel = self.object.data.default_root_vel.torch.clone()[env_ids] + pos_noise = sample_uniform(-1.0, 1.0, (len(env_ids), 3), device=self.device) + + object_default_pose[:, 0:3] = ( + object_default_pose[:, 0:3] + self.cfg.reset_position_noise * pos_noise + self.scene.env_origins[env_ids] + ) + + rot_noise = sample_uniform(-1.0, 1.0, (len(env_ids), 2), device=self.device) # noise for X and Y rotation + object_default_pose[:, 3:7] = randomize_rotation( + rot_noise[:, 0], rot_noise[:, 1], self.x_unit_tensor[env_ids], self.y_unit_tensor[env_ids] + ) + + object_default_vel[:] = 0.0 + self.object.write_root_pose_to_sim_index(root_pose=object_default_pose, env_ids=env_ids) + self.object.write_root_velocity_to_sim_index(root_velocity=object_default_vel, env_ids=env_ids) + + # reset right hand + default_dof_pos = self.right_hand.data.default_joint_pos.torch[env_ids] + dof_limits = self.right_hand.data.joint_limits.torch[env_ids] + dof_pos = sample_joint_positions_within_limits(default_dof_pos, dof_limits, self.cfg.reset_dof_pos_noise) + + dof_vel_noise = sample_uniform(-1.0, 1.0, (len(env_ids), self.num_hand_dofs), device=self.device) + dof_vel = self.right_hand.data.default_joint_vel.torch[env_ids] + self.cfg.reset_dof_vel_noise * dof_vel_noise + + self.right_hand_prev_targets[env_ids] = dof_pos + self.right_hand_curr_targets[env_ids] = dof_pos + + self.right_hand.set_joint_position_target_index(target=dof_pos, env_ids=env_ids) + self.right_hand.write_joint_position_to_sim_index(position=dof_pos, env_ids=env_ids) + self.right_hand.write_joint_velocity_to_sim_index(velocity=dof_vel, env_ids=env_ids) + + # reset left hand + default_dof_pos = self.left_hand.data.default_joint_pos.torch[env_ids] + dof_limits = self.left_hand.data.joint_limits.torch[env_ids] + dof_pos = sample_joint_positions_within_limits(default_dof_pos, dof_limits, self.cfg.reset_dof_pos_noise) + + dof_vel_noise = sample_uniform(-1.0, 1.0, (len(env_ids), self.num_hand_dofs), device=self.device) + dof_vel = self.left_hand.data.default_joint_vel.torch[env_ids] + self.cfg.reset_dof_vel_noise * dof_vel_noise + + self.left_hand_prev_targets[env_ids] = dof_pos + self.left_hand_curr_targets[env_ids] = dof_pos + + self.left_hand.set_joint_position_target_index(target=dof_pos, env_ids=env_ids) + self.left_hand.write_joint_position_to_sim_index(position=dof_pos, env_ids=env_ids) + self.left_hand.write_joint_velocity_to_sim_index(velocity=dof_vel, env_ids=env_ids) + + def _reset_target_pose(self, env_ids): + # reset goal rotation + rand_floats = sample_uniform(-1.0, 1.0, (len(env_ids), 2), device=self.device) + new_rot = randomize_rotation( + rand_floats[:, 0], rand_floats[:, 1], self.x_unit_tensor[env_ids], self.y_unit_tensor[env_ids] + ) + + # update goal pose and markers + self.goal_rot[env_ids] = new_rot + goal_pos = self.goal_pos + self.scene.env_origins + self.goal_markers.visualize(goal_pos, self.goal_rot) diff --git a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/direct/reorient/__init__.py b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/direct/reorient/__init__.py new file mode 100644 index 000000000000..bd7c2efd82db --- /dev/null +++ b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/direct/reorient/__init__.py @@ -0,0 +1,58 @@ +# 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 + +""" +Warp-first Direct variants of the cube reorientation tasks. + +The environments reuse the stable configuration and agent definitions from +:mod:`isaaclab_tasks.core.reorient`. The Allegro warp variant keeps its existing +registration in :mod:`isaaclab_tasks_experimental.direct.allegro_hand` (legacy +in-hand implementation), so only the Shadow variants are registered here. +""" + +import gymnasium as gym + +## +# Register Gym environments. +## + +reorient_warp_entry = "isaaclab_tasks_experimental.direct.reorient.reorient_warp_env:ReorientWarpEnv" +stable_shadow = "isaaclab_tasks.core.reorient.config.shadow_hand" +stable_shadow_agents = f"{stable_shadow}.agents" + +gym.register( + id="Isaac-Reorient-Cube-Shadow-Direct-Warp-v0", + entry_point=reorient_warp_entry, + disable_env_checker=True, + kwargs={ + "env_cfg_entry_point": f"{stable_shadow}.shadow_hand_env_cfg:ShadowHandEnvCfg", + "rl_games_cfg_entry_point": f"{stable_shadow_agents}:rl_games_ppo_cfg.yaml", + "rsl_rl_cfg_entry_point": f"{stable_shadow_agents}.rsl_rl_ppo_cfg:ShadowHandPPORunnerCfg", + "skrl_cfg_entry_point": f"{stable_shadow_agents}:skrl_ppo_cfg.yaml", + }, +) + +gym.register( + id="Isaac-Reorient-Cube-Shadow-OpenAI-FF-Direct-Warp-v0", + entry_point=reorient_warp_entry, + disable_env_checker=True, + kwargs={ + "env_cfg_entry_point": f"{stable_shadow}.shadow_hand_env_cfg:ShadowHandOpenAIEnvCfg", + "rl_games_cfg_entry_point": f"{stable_shadow_agents}:rl_games_ppo_ff_cfg.yaml", + "rsl_rl_cfg_entry_point": f"{stable_shadow_agents}.rsl_rl_ppo_cfg:ShadowHandAsymFFPPORunnerCfg", + "skrl_cfg_entry_point": f"{stable_shadow_agents}:skrl_ff_ppo_cfg.yaml", + }, +) + +gym.register( + id="Isaac-Reorient-Cube-Shadow-OpenAI-LSTM-Direct-Warp-v0", + entry_point=reorient_warp_entry, + disable_env_checker=True, + kwargs={ + "env_cfg_entry_point": f"{stable_shadow}.shadow_hand_env_cfg:ShadowHandOpenAIEnvCfg", + "rl_games_cfg_entry_point": f"{stable_shadow_agents}:rl_games_ppo_lstm_cfg.yaml", + "rsl_rl_cfg_entry_point": f"{stable_shadow_agents}.rsl_rl_ppo_cfg:ShadowHandAsymLSTMPPORunnerCfg", + }, +) diff --git a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/direct/reorient/reorient_kernels.py b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/direct/reorient/reorient_kernels.py new file mode 100644 index 000000000000..96a9fdc00191 --- /dev/null +++ b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/direct/reorient/reorient_kernels.py @@ -0,0 +1,765 @@ +# 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 + +"""Warp kernels and device math for the reorientation task family. + +This is the family's math root: the experimental Direct environments and the +handover family (which reuses the shared hand idioms) launch these kernels on +the environment's Warp-native state (ProxyArray ``.warp`` views) and their own +pre-allocated output buffers. The module deliberately imports no environment +or mdp code. + +A Direct environment step launches these kernels in the following order: + +1. ``_pre_physics_step`` / ``_apply_action``: :func:`ema_actuation_kernel` + turns the policy action into clamped joint position targets. +2. ``_get_dones``: :func:`out_of_reach_kernel` flags fallen objects, then + :func:`reorient_success_kernel` evaluates goal orientation success once per + step, and :func:`reorient_progress_kernel` folds it into the episode + bookkeeping (episode-counter reset, time-out flags, minimum-error tracking). +3. ``_get_rewards``: :func:`reorient_reward` reuses the success state + and launches :func:`reorient_reward_kernel` plus the two + ``consecutive_success_*`` reduction kernels, all writing in place. +4. Observations: :func:`full_obs_kernel` or :func:`reduced_obs_kernel` build + the flattened policy/critic vectors; the camera tasks derive cube corner + keypoints with :func:`compute_cube_keypoints` / + :func:`cube_keypoints_from_quat_kernel`. + +Kernel signatures group their parameters by role, marked with ``# input`` +(read only), ``# input/output`` (read and updated in place), and ``# output`` +(written) comments. +""" + +from __future__ import annotations + +import warp as wp + + +@wp.kernel +def ema_actuation_kernel( + # input + actions: wp.array2d(dtype=wp.float32), + lower: wp.array2d(dtype=wp.float32), + upper: wp.array2d(dtype=wp.float32), + dof_ids: wp.array(dtype=wp.int32), + moving_average: float, + # input/output + prev_targets: wp.array2d(dtype=wp.float32), + # output + cur_targets: wp.array2d(dtype=wp.float32), + compact_targets: wp.array2d(dtype=wp.float32), +): + """Blend actions into clamped joint position targets with an exponential moving average. + + Launched over ``(num_envs, num_actuated_joints)``. Each normalized action in + ``[-1, 1]`` is unscaled to its joint's position range [m or rad, depending on + joint type], blended with the previous target (``moving_average`` weighs the + new value), and clamped back to the limits. The result lands three ways: + ``cur_targets``/``prev_targets`` at the full-joint index ``dof_ids[j]`` (so + the smoothing state persists across steps), and ``compact_targets`` at the + actuated-joint index ``j`` for the articulation's indexed target write. + """ + i, j = wp.tid() + dj = dof_ids[j] + lo = lower[i, dj] + hi = upper[i, dj] + # unscale from [-1, 1] to the joint range, blend, and clamp + t = 0.5 * (actions[i, j] + 1.0) * (hi - lo) + lo + t = moving_average * t + (1.0 - moving_average) * prev_targets[i, dj] + t = wp.clamp(t, lo, hi) + cur_targets[i, dj] = t + prev_targets[i, dj] = t + compact_targets[i, j] = t + + +@wp.kernel +def fingertip_pos_kernel( + # input + body_pos_w: wp.array2d(dtype=wp.vec3f), + env_origins: wp.array(dtype=wp.vec3f), + body_ids: wp.array(dtype=wp.int32), + # output + out: wp.array2d(dtype=wp.float32), +): + """Gather flattened fingertip positions [m] in the environment frame. + + Launched over ``(num_envs, num_fingertips)``; fingertip ``j`` fills columns + ``[3j, 3j + 3)`` of ``out``. + """ + 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( + # input + body_quat_w: wp.array2d(dtype=wp.quatf), + body_ids: wp.array(dtype=wp.int32), + # output + out: wp.array2d(dtype=wp.float32), +): + """Gather flattened fingertip ``(x, y, z, w)`` orientations. + + Launched over ``(num_envs, num_fingertips)``; fingertip ``j`` fills columns + ``[4j, 4j + 4)`` of ``out``. + """ + 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_vel_kernel( + # input + body_vel_w: wp.array2d(dtype=wp.spatial_vectorf), + body_ids: wp.array(dtype=wp.int32), + # output + out: wp.array2d(dtype=wp.float32), +): + """Gather flattened fingertip spatial velocities [m/s, rad/s]. + + Launched over ``(num_envs, num_fingertips)``; fingertip ``j`` fills columns + ``[6j, 6j + 6)`` of ``out``. + """ + i, j = wp.tid() + v = body_vel_w[i, body_ids[j]] + for k in range(6): + out[i, 6 * j + k] = v[k] + + +# Column-parallel segment readers shared by the observation kernels launched over +# ``(num_envs, obs_dim)``: each maps a flattened segment column ``k`` to one value. + + +@wp.func +def fingertip_pos_col( + body_pos_w: wp.array2d(dtype=wp.vec3f), + env_origins: wp.array(dtype=wp.vec3f), + body_ids: wp.array(dtype=wp.int32), + i: int, + k: int, +) -> float: + """Column ``k`` of the flattened fingertip positions [m] in the environment frame.""" + p = body_pos_w[i, body_ids[k // 3]] - env_origins[i] + return p[k % 3] + + +@wp.func +def fingertip_quat_col( + body_quat_w: wp.array2d(dtype=wp.quatf), + body_ids: wp.array(dtype=wp.int32), + i: int, + k: int, +) -> float: + """Column ``k`` of the flattened fingertip ``(x, y, z, w)`` orientations.""" + return body_quat_w[i, body_ids[k // 4]][k % 4] + + +@wp.func +def fingertip_vel_col( + body_vel_w: wp.array2d(dtype=wp.spatial_vectorf), + body_ids: wp.array(dtype=wp.int32), + i: int, + k: int, +) -> float: + """Column ``k`` of the flattened fingertip spatial velocities [m/s, rad/s].""" + return body_vel_w[i, body_ids[k // 6]][k % 6] + + +# --------------------------------------------------------------------------- +# Success and reward math (Warp). +# +# Kernels shared by the Direct environments and the manager terms in +# :mod:`~isaaclab_tasks.core.reorient.mdp`; state arrays update in place. The +# quaternion arithmetic mirrors :func:`isaaclab.utils.math.quat_mul` +# (``(x, y, z, w)`` layout) term for term so results match the previously +# validated reference implementation. +# --------------------------------------------------------------------------- + + +@wp.func +def _rotation_distance(object_quat: wp.quatf, target_quat: wp.quatf) -> float: + """Orientation distance [rad] between two ``(x, y, z, w)`` quaternions.""" + q1 = object_quat + # quat_inverse == conjugate for these unit quaternions + q2 = wp.quat_inverse(target_quat) + ww = (q1[2] + q1[0]) * (q2[0] + q2[1]) + yy = (q1[3] - q1[1]) * (q2[3] + q2[2]) + zz = (q1[3] + q1[1]) * (q2[3] - q2[2]) + xx = ww + yy + zz + qq = 0.5 * (xx + (q1[2] - q1[0]) * (q2[0] - q2[1])) + x = qq - xx + (q1[0] + q1[3]) * (q2[0] + q2[3]) + y = qq - yy + (q1[3] - q1[0]) * (q2[1] + q2[2]) + z = qq - zz + (q1[2] + q1[1]) * (q2[3] - q2[0]) + return 2.0 * wp.asin(wp.min(wp.length(wp.vec3(x, y, z)), 1.0)) + + +@wp.kernel +def rotation_distance_kernel( + # input + object_quat: wp.array(dtype=wp.quatf), + target_quat: wp.array(dtype=wp.quatf), + # output + distance: wp.array(dtype=wp.float32), +): + """Per-environment orientation distance [rad] between object and target.""" + i = wp.tid() + distance[i] = _rotation_distance(object_quat[i], target_quat[i]) + + +@wp.kernel +def reorient_success_kernel( + # input + object_quat: wp.array(dtype=wp.quatf), + target_quat: wp.array(dtype=wp.quatf), + success_tolerance: float, + # output + success: wp.array(dtype=wp.bool), + distance: wp.array(dtype=wp.float32), +): + """Evaluate goal-orientation success while exposing the physical error. + + The single per-step success evaluation: writes the orientation distance + [rad] and whether it is within ``success_tolerance`` [rad]. Downstream + consumers — :func:`reorient_progress_kernel`, the manager command/reward/ + termination terms, and the environments' metrics — reuse these buffers + instead of recomputing the quaternion math. + """ + i = wp.tid() + d = _rotation_distance(object_quat[i], target_quat[i]) + distance[i] = d + success[i] = d <= success_tolerance + + +@wp.kernel +def reorient_reward_kernel( + # input + object_pos_w: wp.array(dtype=wp.vec3f), + env_origins: wp.array(dtype=wp.vec3f), + object_quat: wp.array(dtype=wp.quatf), + target_pos_e: wp.array(dtype=wp.vec3f), + target_quat: wp.array(dtype=wp.quatf), + actions: wp.array2d(dtype=wp.float32), + reset_buf: wp.array(dtype=wp.bool), + distance_scale: float, + rotation_scale: float, + rotation_epsilon: float, + action_penalty_scale: float, + success_tolerance: float, + success_bonus: float, + fall_distance: float, + fall_penalty: float, + # input/output + goal_resets: wp.array(dtype=wp.bool), + successes: wp.array(dtype=wp.float32), + # output + reward: wp.array(dtype=wp.float32), + resets: wp.array(dtype=wp.bool), +): + """Per-environment reorientation reward and success state transition. + + The reward sums four terms: goal distance [m] times ``distance_scale`` + (negative scale penalizes distance), the inverse orientation error + ``rotation_scale / (distance [rad] + rotation_epsilon)``, the squared-action + penalty times ``action_penalty_scale``, plus a one-time ``success_bonus`` + when the goal orientation is first within ``success_tolerance`` [rad] and a + ``fall_penalty`` when the object drifts beyond ``fall_distance`` [m]. + + State transition, in place: ``goal_resets`` marks environments due for goal + resampling (sticky until the caller resamples and clears it), each newly + reached goal increments ``successes``, and ``resets`` combines the incoming + episode resets with the fallen-object condition for the consecutive-success + reduction downstream. + """ + i = wp.tid() + # position error to the goal [m] + # object position is world-frame; the target is authored in the environment frame + goal_distance = wp.length(object_pos_w[i] - env_origins[i] - target_pos_e[i]) + rotation_distance = _rotation_distance(object_quat[i], target_quat[i]) + # a goal within tolerance marks this env for goal resampling (sticky until resampled) + goal_reset = rotation_distance <= success_tolerance or goal_resets[i] + action_penalty = float(0.0) + for j in range(actions.shape[1]): + action_penalty += actions[i, j] * actions[i, j] + value = ( + goal_distance * distance_scale + + rotation_scale / (rotation_distance + rotation_epsilon) + + action_penalty * action_penalty_scale + ) + # one-time bonus on reaching the goal; penalty and episode reset when the object falls + if goal_reset: + value += success_bonus + fell = goal_distance >= fall_distance + if fell: + value += fall_penalty + reward[i] = value + goal_resets[i] = goal_reset + # per-episode success count feeds the consecutive-success moving average + if goal_reset: + successes[i] = successes[i] + 1.0 + resets[i] = fell or reset_buf[i] + + +@wp.kernel +def consecutive_success_stats_kernel( + # input + successes: wp.array(dtype=wp.float32), + resets: wp.array(dtype=wp.bool), + # output + stats: wp.array(dtype=wp.float32), +): + """Accumulate the resetting-env count and their episode successes into ``stats``. + + First half of the consecutive-success moving-average reduction (atomic + adds). ``stats`` must be zeroed before the launch; ``[0]`` receives the + number of resetting environments and ``[1]`` the sum of their episode + success counts. + """ + i = wp.tid() + if resets[i]: + wp.atomic_add(stats, 0, 1.0) + wp.atomic_add(stats, 1, successes[i]) + + +@wp.kernel +def consecutive_success_update_kernel( + # input + stats: wp.array(dtype=wp.float32), + averaging_factor: float, + # input/output + consecutive_successes: wp.array(dtype=wp.float32), +): + """Fold the finished-episode statistics into the moving average. + + Second half of the reduction; launch with ``dim=1`` after + :func:`consecutive_success_stats_kernel` so the completed sums are visible. + Leaves the average untouched when no episode finished this step. + """ + if stats[0] > 0.0: + consecutive_successes[0] = ( + averaging_factor * stats[1] / stats[0] + (1.0 - averaging_factor) * consecutive_successes[0] + ) + + +class ReorientRewardBuffers: + """Caller-owned device buffers for :func:`reorient_reward`. + + Allocate once (Direct environments in ``__init__``, manager terms in their + term ``__init__``); every step writes in place. + """ + + def __init__(self, num_envs: int, device: str): + device = str(device) + self.reward = wp.zeros(num_envs, dtype=wp.float32, device=device) + self.resets = wp.empty(num_envs, dtype=wp.bool, device=device) + # scratch for the consecutive-success reduction: [0] resetting-env + # count, [1] summed episode successes of the resetting environments + self.stats = wp.zeros(2, dtype=wp.float32, device=device) + + +def reorient_reward( + reset_buf: wp.array, + goal_resets: wp.array, + successes: wp.array, + consecutive_successes: wp.array, + object_pos_w: wp.array, + env_origins: wp.array, + object_quat: wp.array, + target_pos_e: wp.array, + target_quat: wp.array, + actions: wp.array, + distance_scale: float, + rotation_scale: float, + rotation_epsilon: float, + action_penalty_scale: float, + success_tolerance: float, + success_bonus: float, + fall_distance: float, + fall_penalty: float, + averaging_factor: float, + buffers: ReorientRewardBuffers, +) -> wp.array: + """Compute the Direct reorientation reward and success state transition in place. + + Updates :paramref:`goal_resets`, :paramref:`successes`, and + :paramref:`consecutive_successes` in place and writes the per-environment + reward into ``buffers.reward``. + + Args: + reset_buf: Episode-reset flags, ``wp.bool``. The kernel ORs in the + fallen-object condition itself, so callers whose only other reset + source is the episode time-out may pass those flags directly. + goal_resets: Goal-reset flags, ``wp.bool``; read and updated in place. + successes: Goals reached in each episode, ``wp.float32``; updated in place. + consecutive_successes: One-element moving-average success count, + ``wp.float32``; updated in place. + object_pos_w: Object world positions [m], ``wp.vec3f``. + env_origins: Environment origins [m], ``wp.vec3f``. + object_quat: Object ``(x, y, z, w)`` orientations, ``wp.quatf``. + target_pos_e: Goal positions in the environment frame [m], ``wp.vec3f``. + target_quat: Goal ``(x, y, z, w)`` orientations, ``wp.quatf``. + actions: Normalized joint actions, 2D ``wp.float32``. + distance_scale: Position-distance reward scale [1/m]. + rotation_scale: Orientation reward scale [rad]. + rotation_epsilon: Orientation reward regularizer [rad]. + action_penalty_scale: Squared-action reward scale. + success_tolerance: Goal orientation tolerance [rad]. + success_bonus: Reward added when a goal is reached. + fall_distance: Object-to-goal termination distance [m]. + fall_penalty: Reward added when the object is out of reach. + averaging_factor: Consecutive-success moving-average factor. + buffers: Caller-owned output/scratch buffers. + + Returns: + The per-environment reward, a view of ``buffers.reward``. + """ + num_envs = buffers.reward.shape[0] + wp.launch( + reorient_reward_kernel, + dim=num_envs, + inputs=[ + object_pos_w, + env_origins, + object_quat, + target_pos_e, + target_quat, + actions, + reset_buf, + distance_scale, + rotation_scale, + rotation_epsilon, + action_penalty_scale, + success_tolerance, + success_bonus, + fall_distance, + fall_penalty, + ], + outputs=[ + goal_resets, + successes, + buffers.reward, + buffers.resets, + ], + device=buffers.reward.device, + ) + buffers.stats.zero_() + wp.launch( + consecutive_success_stats_kernel, + dim=num_envs, + inputs=[successes, buffers.resets], + outputs=[buffers.stats], + device=buffers.reward.device, + ) + wp.launch( + consecutive_success_update_kernel, + dim=1, + inputs=[buffers.stats, averaging_factor], + outputs=[consecutive_successes], + device=buffers.reward.device, + ) + return buffers.reward + + +# --------------------------------------------------------------------------- +# Cube-keypoint math shared by the Direct camera environment and the manager +# camera observation terms. +# --------------------------------------------------------------------------- + + +CUBE_HALF_SIZE = wp.vec3(0.03, 0.03, 0.03) +"""Half side lengths [m] of the reorientation cube.""" + + +@wp.kernel +def cube_keypoints_from_quat_kernel( + # input + quat: wp.array(dtype=wp.quatf), + half_size: wp.vec3, + # output + keypoints: wp.array2d(dtype=wp.float32), +): + """Rotation-only cube-corner offsets [m] from batched orientations. + + Launched over ``(num_envs, num_corners)``: corner ``c`` (its index bits + select the +/- half side per axis) fills the flattened columns + ``[3c, 3c + 3)``. Used where only the rotated corner offsets matter, e.g. + goal-orientation keypoints without a translation. + """ + env, corner = wp.tid() + # corner index bits select the +/- half-side per axis (bit set -> negative) + sign_x = wp.where(((corner >> 0) & 1) == 0, 1.0, -1.0) + sign_y = wp.where(((corner >> 1) & 1) == 0, 1.0, -1.0) + sign_z = wp.where(((corner >> 2) & 1) == 0, 1.0, -1.0) + offset = wp.vec3(sign_x * half_size[0], sign_y * half_size[1], sign_z * half_size[2]) + p = wp.quat_rotate(quat[env], offset) + keypoints[env, 3 * corner + 0] = p[0] + keypoints[env, 3 * corner + 1] = p[1] + keypoints[env, 3 * corner + 2] = p[2] + + +@wp.kernel +def _cube_keypoints_kernel( + # input + pose: wp.array(dtype=wp.float32, ndim=2), + half_size: wp.vec3, + # output + keypoints: wp.array2d(dtype=wp.vec3), +): + """Posed cube-corner positions [m]; launch via :func:`compute_cube_keypoints`.""" + env, corner = wp.tid() + # corner index bits select the +/- half-side per axis (bit set -> negative) + sign_x = wp.where(((corner >> 0) & 1) == 0, 1.0, -1.0) + sign_y = wp.where(((corner >> 1) & 1) == 0, 1.0, -1.0) + sign_z = wp.where(((corner >> 2) & 1) == 0, 1.0, -1.0) + offset = wp.vec3(sign_x * half_size[0], sign_y * half_size[1], sign_z * half_size[2]) + orientation = wp.quat(pose[env, 3], pose[env, 4], pose[env, 5], pose[env, 6]) + position = wp.vec3(pose[env, 0], pose[env, 1], pose[env, 2]) + keypoints[env, corner] = position + wp.quat_rotate(orientation, offset) + + +def compute_cube_keypoints( + pose: wp.array, + keypoints: wp.array, + size: tuple[float, float, float] = (2 * 0.03, 2 * 0.03, 2 * 0.03), +) -> None: + """Compute cube-corner positions for batched poses. + + Args: + pose: Cube center poses ``(x, y, z, qx, qy, qz, qw)`` [m, unit + quaternion], 2D ``wp.float32``. + keypoints: Output corner positions [m], 2D ``wp.vec3`` of shape + ``(num_envs, num_keypoints)``; the corner count is taken from it. + size: Cube side lengths along each axis [m]. + """ + wp.launch( + _cube_keypoints_kernel, + dim=(keypoints.shape[0], keypoints.shape[1]), + inputs=[ + pose, + wp.vec3(size[0] / 2.0, size[1] / 2.0, size[2] / 2.0), + ], + outputs=[keypoints], + device=keypoints.device, + ) + + +@wp.kernel +def out_of_reach_kernel( + # input + object_pos_w: wp.array(dtype=wp.vec3f), + env_origins: wp.array(dtype=wp.vec3f), + target_pos_e: wp.array(dtype=wp.vec3f), + fall_distance: float, + # output + out_of_reach: wp.array(dtype=wp.bool), +): + """Flag environments whose object drifted beyond ``fall_distance`` [m] from the in-hand target.""" + i = wp.tid() + out_of_reach[i] = wp.length(object_pos_w[i] - env_origins[i] - target_pos_e[i]) >= fall_distance + + +@wp.kernel +def full_obs_kernel( + # input + joint_pos: wp.array2d(dtype=wp.float32), + joint_vel: wp.array2d(dtype=wp.float32), + lower: wp.array2d(dtype=wp.float32), + upper: wp.array2d(dtype=wp.float32), + object_pos_w: wp.array(dtype=wp.vec3f), + env_origins: wp.array(dtype=wp.vec3f), + object_quat: wp.array(dtype=wp.quatf), + object_lin_vel: wp.array(dtype=wp.vec3f), + object_ang_vel: wp.array(dtype=wp.vec3f), + in_hand_pos_e: wp.array(dtype=wp.vec3f), + goal_quat: wp.array(dtype=wp.quatf), + body_pos_w: wp.array2d(dtype=wp.vec3f), + body_quat_w: wp.array2d(dtype=wp.quatf), + body_vel_w: wp.array2d(dtype=wp.spatial_vectorf), + finger_ids: wp.array(dtype=wp.int32), + force: wp.array2d(dtype=wp.vec3f), + torque: wp.array2d(dtype=wp.vec3f), + wrench_ids: wp.array(dtype=wp.int32), + actions: wp.array2d(dtype=wp.float32), + vel_scale: float, + force_scale: float, + with_forces: int, + # output + out: wp.array2d(dtype=wp.float32), +): + """Direct full observation / full state, matching the reference concatenation order. + + Column layout of ``out``, for ``J`` joints, ``F`` fingertips, and ``A`` + actions (segment widths in parentheses): + + 1. normalized joint positions (``J``), scaled joint velocities (``J``) + 2. object: environment-frame position (3), ``(x, y, z, w)`` rotation (4), + linear velocity (3), scaled angular velocity (3) + 3. goal: in-hand anchor position (3), goal rotation (4), object-to-goal + rotation difference (4) + 4. fingertips: positions (``3F``), rotations (``4F``), spatial + velocities (``6F``) + 5. fingertip wrenches (``6F``) — full state only; the segment exists when + ``with_forces != 0`` and zero-fills while the sensor has no data + (``with_forces == -1``) + 6. last actions (``A``) + + Launched over ``(num_envs, obs_dim)``: each thread walks a branch ladder over the + segment boundaries and writes one output column, so warps (32 consecutive columns) + stay branch-uniform except at segment boundaries. + """ + i, j = wp.tid() + num_joints = joint_pos.shape[1] + num_fingers = finger_ids.shape[0] + # segment boundaries, in column order + end_joint = 2 * num_joints + end_object = end_joint + 13 + end_goal = end_object + 11 + end_tip_pos = end_goal + 3 * num_fingers + end_tip_quat = end_tip_pos + 4 * num_fingers + end_tip_vel = end_tip_quat + 6 * num_fingers + end_wrench = end_tip_vel + if with_forces != 0: + end_wrench += 6 * num_fingers + # hand: normalized DOF positions, scaled DOF velocities + if j < num_joints: + out[i, j] = 2.0 * (joint_pos[i, j] - lower[i, j]) / (upper[i, j] - lower[i, j]) - 1.0 + elif j < end_joint: + out[i, j] = vel_scale * joint_vel[i, j - num_joints] + # object pose and velocities (environment frame position) + elif j < end_object: + k = j - end_joint + if k < 3: + p = object_pos_w[i] - env_origins[i] + out[i, j] = p[k] + elif k < 7: + out[i, j] = object_quat[i][k - 3] + elif k < 10: + out[i, j] = object_lin_vel[i][k - 7] + else: + out[i, j] = vel_scale * object_ang_vel[i][k - 10] + # goal: in-hand anchor, goal rotation, and the goal-to-object rotation difference + elif j < end_goal: + k = j - end_object + if k < 3: + out[i, j] = in_hand_pos_e[i][k] + elif k < 7: + out[i, j] = goal_quat[i][k - 3] + else: + # quat_inverse == conjugate for these unit quaternions, matching + # isaaclab.utils.math.quat_mul/quat_conjugate semantics + qe = object_quat[i] * wp.quat_inverse(goal_quat[i]) + out[i, j] = qe[k - 7] + # fingertips: environment-frame positions, rotations, spatial velocities + elif j < end_tip_pos: + out[i, j] = fingertip_pos_col(body_pos_w, env_origins, finger_ids, i, j - end_goal) + elif j < end_tip_quat: + out[i, j] = fingertip_quat_col(body_quat_w, finger_ids, i, j - end_tip_pos) + elif j < end_tip_vel: + out[i, j] = fingertip_vel_col(body_vel_w, finger_ids, i, j - end_tip_quat) + # fingertip force/torque sensors (full state only; absent when with_forces == 0) + elif j < end_wrench: + if with_forces == 1: + k = j - end_tip_vel + c = k % 6 + if c < 3: + out[i, j] = force_scale * force[i, wrench_ids[k // 6]][c] + else: + out[i, j] = force_scale * torque[i, wrench_ids[k // 6]][c - 3] + else: + # full state requested but the sensor has no data yet: zero block + out[i, j] = 0.0 + # actions + else: + out[i, j] = actions[i, j - end_wrench] + + +@wp.kernel +def reduced_obs_kernel( + # input + body_pos_w: wp.array2d(dtype=wp.vec3f), + env_origins: wp.array(dtype=wp.vec3f), + finger_ids: wp.array(dtype=wp.int32), + object_pos_w: wp.array(dtype=wp.vec3f), + object_quat: wp.array(dtype=wp.quatf), + goal_quat: wp.array(dtype=wp.quatf), + actions: wp.array2d(dtype=wp.float32), + # output + out: wp.array2d(dtype=wp.float32), +): + """Direct reduced (OpenAI) observation, matching the reference concatenation order. + + Column layout of ``out``, for ``F`` fingertips and ``A`` actions: + environment-frame fingertip positions (``3F``), then the object's + environment-frame position (3) and object-to-goal rotation difference (4), + then the last actions (``A``). One thread per environment writes the whole + row. + """ + i = wp.tid() + num_fingers = finger_ids.shape[0] + for f in range(num_fingers): + fp = body_pos_w[i, finger_ids[f]] - env_origins[i] + out[i, 3 * f + 0] = fp[0] + out[i, 3 * f + 1] = fp[1] + out[i, 3 * f + 2] = fp[2] + idx = 3 * num_fingers + p = object_pos_w[i] - env_origins[i] + # quat_inverse == conjugate for these unit quaternions, matching + # isaaclab.utils.math.quat_mul/quat_conjugate semantics + qe = object_quat[i] * wp.quat_inverse(goal_quat[i]) + out[i, idx + 0] = p[0] + out[i, idx + 1] = p[1] + out[i, idx + 2] = p[2] + out[i, idx + 3] = qe[0] + out[i, idx + 4] = qe[1] + out[i, idx + 5] = qe[2] + out[i, idx + 6] = qe[3] + idx += 7 + for a in range(actions.shape[1]): + out[i, idx + a] = actions[i, a] + + +@wp.kernel +def reorient_progress_kernel( + # input + success: wp.array(dtype=wp.bool), + distance: wp.array(dtype=wp.float32), + successes: wp.array(dtype=wp.float32), + max_consecutive_success: float, + max_episode_length: wp.int64, + # input/output + episode_length: wp.array(dtype=wp.int64), + minimum_error: wp.array(dtype=wp.float32), + # output + time_out: wp.array(dtype=wp.bool), + has_sample: wp.array(dtype=wp.bool), +): + """Per-step progress bookkeeping downstream of :func:`reorient_success_kernel`. + + Consumes the success flags and orientation distances [rad] and, in place: + zeroes the episode progress on goal-reached environments (only when + :paramref:`max_consecutive_success` > 0), flags episode time-out — by + elapsed episode length or by reaching the consecutive-success cap — and + tightens the episode-minimum error tracking (absorbing + ``EpisodeErrorRecorder.update``: non-finite distances leave + ``minimum_error`` and the sticky ``has_sample`` marker untouched). + """ + i = wp.tid() + max_success_reached = False + if max_consecutive_success > 0.0: + if success[i]: + episode_length[i] = wp.int64(0) + max_success_reached = successes[i] >= max_consecutive_success + time_out[i] = episode_length[i] >= max_episode_length - wp.int64(1) or max_success_reached + d = distance[i] + if wp.isfinite(d): + if d < minimum_error[i]: + minimum_error[i] = d + has_sample[i] = True diff --git a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/direct/reorient/reorient_warp_env.py b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/direct/reorient/reorient_warp_env.py new file mode 100644 index 000000000000..173e64c3ec6d --- /dev/null +++ b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/direct/reorient/reorient_warp_env.py @@ -0,0 +1,520 @@ +# 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 + +"""Warp-first Direct environment for the cube reorientation tasks. + +Experimental counterpart of the torch-first mainline implementation in +:mod:`isaaclab_tasks.core.reorient`: environment-owned state lives in Warp arrays +and rewards, observations, and resets are computed in Warp kernels. +""" + +from __future__ import annotations + +from collections.abc import Sequence +from typing import TYPE_CHECKING + +import torch +import warp as wp + +import isaaclab.sim as sim_utils +from isaaclab import cloner +from isaaclab.assets import Articulation, RigidObject +from isaaclab.envs import DirectRLEnv +from isaaclab.markers import VisualizationMarkers +from isaaclab.sensors import JointWrenchSensor, JointWrenchSensorCfg +from isaaclab.sim.spawners.from_files import GroundPlaneCfg, spawn_ground_plane +from isaaclab.utils.math import sample_uniform + +from isaaclab_tasks.core.reorient.reorient_task_base import GOAL_MARKER_POSITION, IN_HAND_POS_OFFSET +from isaaclab_tasks.core.utils import EpisodeErrorRecorder, randomize_rotation, sample_joint_positions_within_limits + +from isaaclab_tasks_experimental.direct.reorient.reorient_kernels import ( + ReorientRewardBuffers, + ema_actuation_kernel, + full_obs_kernel, + out_of_reach_kernel, + reduced_obs_kernel, + reorient_progress_kernel, + reorient_reward, + reorient_success_kernel, +) + +if TYPE_CHECKING: + from isaaclab_tasks.core.reorient.config.allegro_hand.allegro_hand_direct_env_cfg import AllegroHandEnvCfg + from isaaclab_tasks.core.reorient.config.shadow_hand.shadow_hand_env_cfg import ShadowHandEnvCfg + + +class ReorientWarpEnv(DirectRLEnv): + cfg: AllegroHandEnvCfg | ShadowHandEnvCfg + + def __init__(self, cfg: AllegroHandEnvCfg | ShadowHandEnvCfg, render_mode: str | None = None, **kwargs): + super().__init__(cfg, render_mode, **kwargs) + + # -- Robot introspection: joints, bodies, limits -- + self.num_hand_dofs = self.hand.num_joints + self.actuated_dof_indices, _ = self.hand.find_joints(cfg.actuated_joint_names) + if len(self.actuated_dof_indices) != len(cfg.actuated_joint_names): + raise ValueError( + f"Expected {len(cfg.actuated_joint_names)} actuated joints, found {len(self.actuated_dof_indices)}." + ) + self.finger_bodies, fingertip_body_names = self.hand.find_bodies(self.cfg.fingertip_body_names) + if len(self.finger_bodies) != len(self.cfg.fingertip_body_names): + raise ValueError( + f"Expected {len(self.cfg.fingertip_body_names)} fingertip bodies, found {len(self.finger_bodies)}." + ) + self.num_fingertips = len(self.finger_bodies) + self.finger_wrench_bodies = [] + if getattr(self, "_joint_wrench_sensor", None) is not None: + for body_name in fingertip_body_names: + self.finger_wrench_bodies.append(self._joint_wrench_sensor.body_names.index(body_name)) + self.finger_wrench_bodies.sort() + joint_pos_limits = self.hand.data.joint_limits.torch.to(self.device) + # warp-owned copies of the static joint limits (split into per-bound planes) + self.hand_dof_lower_limits = wp.clone(wp.from_torch(joint_pos_limits[..., 0].contiguous())) + self.hand_dof_upper_limits = wp.clone(wp.from_torch(joint_pos_limits[..., 1].contiguous())) + + # -- Observation layout, derived from the introspection -- + num_joints, num_fingers = self.num_hand_dofs, self.num_fingertips + num_actions = self.cfg.action_space + full_dim = 2 * num_joints + 13 + 11 + 13 * num_fingers + num_actions + reduced_dim = 3 * num_fingers + 7 + num_actions + state_dim = full_dim + 6 * num_fingers + + # -- Warp buffers: actuation -- + # joint/body index arrays consumed by the indexed kernels + self.actuated_dof_ids = wp.array(self.actuated_dof_indices, dtype=wp.int32, device=str(self.device)) + self.finger_ids = wp.array(self.finger_bodies, dtype=wp.int32, device=str(self.device)) + self.wrench_ids = wp.array( + self.finger_wrench_bodies if len(self.finger_wrench_bodies) else self.finger_bodies, + dtype=wp.int32, + device=str(self.device), + ) + # EMA-smoothed joint position targets: full joint set plus the actuated-subset write buffer + self.prev_targets = wp.zeros((self.num_envs, self.num_hand_dofs), dtype=wp.float32, device=str(self.device)) + self.cur_targets = wp.zeros((self.num_envs, self.num_hand_dofs), dtype=wp.float32, device=str(self.device)) + self.compact_targets = wp.zeros( + (self.num_envs, len(self.actuated_dof_indices)), dtype=wp.float32, device=str(self.device) + ) + # the policy action is copied in-place each step so the alias below stays bound + self.actions = wp.zeros((self.num_envs, num_actions), dtype=wp.float32, device=str(self.device)) + + # -- Warp buffers: goal and success state -- + # goal pose: the in-hand anchor and the sampled goal orientation / marker position + self.in_hand_pos = wp.zeros(self.num_envs, dtype=wp.vec3f, device=str(self.device)) + self.goal_rot = wp.zeros(self.num_envs, dtype=wp.quatf, device=str(self.device)) + self.goal_pos = wp.zeros(self.num_envs, dtype=wp.vec3f, device=str(self.device)) + # success tracking: per-episode counts, goal-resample flags, and logging state + self.successes = wp.zeros(self.num_envs, dtype=wp.float32, device=str(self.device)) + self.reset_goal_buf = wp.zeros(self.num_envs, dtype=wp.bool, device=str(self.device)) + self.consecutive_successes = wp.zeros(1, dtype=wp.float32, device=str(self.device)) + self._last_episode_success = wp.zeros(self.num_envs, dtype=wp.bool, device=str(self.device)) + + # -- Warp buffers: per-step evaluation outputs -- + # success evaluation, written once per step in :meth:`_get_dones` + self.success_flags = wp.empty(self.num_envs, dtype=wp.bool, device=str(self.device)) + self.orientation_error_buf = wp.empty(self.num_envs, dtype=wp.float32, device=str(self.device)) + # termination flags + self.out_of_reach_flags = wp.empty(self.num_envs, dtype=wp.bool, device=str(self.device)) + self.time_out_flags = wp.empty(self.num_envs, dtype=wp.bool, device=str(self.device)) + # reward and diagnostics machinery (Warp-native internally) + self._reward_buffers = ReorientRewardBuffers(self.num_envs, self.device) + self._orientation_error = EpisodeErrorRecorder(self.num_envs, self.device) + + # -- Warp buffers: observation outputs -- + self.policy_obs_buf = wp.empty( + (self.num_envs, reduced_dim if self.cfg.obs_type == "openai" else full_dim), + dtype=wp.float32, + device=str(self.device), + ) + self.state_obs_buf = wp.empty((self.num_envs, state_dim), dtype=wp.float32, device=str(self.device)) + # placeholder for kernel launches while the wrench sensor has no data + self.dummy_wrench = wp.zeros((1, 1), dtype=wp.vec3f, device=str(self.device)) + + # -- Warp buffers: reset randomization constants -- + # batched unit axes consumed (through their torch aliases) by the torch.jit + # reset helper; a Warp reset path would replace them with wp.vec3f constants + self.x_unit_vec = wp.zeros(self.num_envs, dtype=wp.vec3f, device=str(self.device)) + self.y_unit_vec = wp.zeros(self.num_envs, dtype=wp.vec3f, device=str(self.device)) + self.z_unit_vec = wp.zeros(self.num_envs, dtype=wp.vec3f, device=str(self.device)) + + # -- Visualization and articulation write handles -- + self.goal_markers = VisualizationMarkers(self.cfg.goal_object_cfg) + self._set_joint_pos_target = self.hand.set_joint_position_target_index + self._write_obj_root_pose = self.object.write_root_pose_to_sim_index + self._write_obj_root_vel = self.object.write_root_velocity_to_sim_index + self._write_hand_joint_pos = self.hand.write_joint_position_to_sim_index + self._write_hand_joint_vel = self.hand.write_joint_velocity_to_sim_index + + # -- Bind torch aliases to the warp-native buffers (zero-copy views) -- + # robot joint limits + self.hand_dof_lower_limits_torch = wp.to_torch(self.hand_dof_lower_limits) + self.hand_dof_upper_limits_torch = wp.to_torch(self.hand_dof_upper_limits) + # actuation + self.prev_targets_torch = wp.to_torch(self.prev_targets) + self.cur_targets_torch = wp.to_torch(self.cur_targets) + self.compact_targets_torch = wp.to_torch(self.compact_targets) + self.actions_torch = wp.to_torch(self.actions) + # goal and success state + self.in_hand_pos_torch = wp.to_torch(self.in_hand_pos) + self.goal_rot_torch = wp.to_torch(self.goal_rot) + self.goal_pos_torch = wp.to_torch(self.goal_pos) + self.successes_torch = wp.to_torch(self.successes) + self.reset_goal_buf_torch = wp.to_torch(self.reset_goal_buf) + self.consecutive_successes_torch = wp.to_torch(self.consecutive_successes) + self._last_episode_success_torch = wp.to_torch(self._last_episode_success) + # per-step evaluation and reward outputs + self.out_of_reach_flags_torch = wp.to_torch(self.out_of_reach_flags) + self.time_out_flags_torch = wp.to_torch(self.time_out_flags) + self.reward_torch = wp.to_torch(self._reward_buffers.reward) + # observation outputs + self.policy_obs_buf_torch = wp.to_torch(self.policy_obs_buf) + self.state_obs_buf_torch = wp.to_torch(self.state_obs_buf) + # reset randomization constants + self.x_unit_vec_torch = wp.to_torch(self.x_unit_vec) + self.y_unit_vec_torch = wp.to_torch(self.y_unit_vec) + self.z_unit_vec_torch = wp.to_torch(self.z_unit_vec) + + # -- Cached warp views of externally-owned torch tensors -- + self.episode_length_buf_warp = wp.from_torch(self.episode_length_buf) + self.env_origins_warp = wp.from_torch(self.scene.env_origins, dtype=wp.vec3f) + + # -- Initial buffer contents (written through the torch aliases) -- + # in-hand target = object default position + shared offset (mirrors ReorientCommand) + self.in_hand_pos_torch[:] = self.object.data.default_root_pose.torch[:, 0:3] + self.in_hand_pos_torch += torch.tensor( + IN_HAND_POS_OFFSET, dtype=self.in_hand_pos_torch.dtype, device=self.in_hand_pos_torch.device + ) + self.goal_rot_torch[:, 3] = 1.0 # identity quaternion in (x, y, z, w) layout + self.goal_pos_torch[:, :] = torch.tensor(GOAL_MARKER_POSITION, device=self.device) + self.x_unit_vec_torch[:] = torch.tensor([1.0, 0.0, 0.0], device=self.device) + self.y_unit_vec_torch[:] = torch.tensor([0.0, 1.0, 0.0], device=self.device) + self.z_unit_vec_torch[:] = torch.tensor([0.0, 0.0, 1.0], device=self.device) + + def _setup_scene(self): + # add hand, in-hand object, and goal object + self.hand = Articulation(self.cfg.robot_cfg) + self.object: Articulation | RigidObject = self.cfg.object_cfg.class_type(self.cfg.object_cfg) + self._joint_wrench_sensor = None + if self.cfg.asymmetric_obs: + self._joint_wrench_sensor = self._create_joint_wrench_sensor() + # add ground plane + spawn_ground_plane(prim_path="/World/ground", cfg=GroundPlaneCfg()) + src, dest = "/World/envs/env_0", "/World/envs/env_{}" + pos = cloner.grid_transforms(self.scene.num_envs, self.scene.cfg.env_spacing, device=self.device)[0] + plan = cloner.ClonePlan.from_env_0(src, dest, self.scene.num_envs, self.device, pos) + cloner.replicate(plan, stage=self.scene.stage) + # add articulation to scene - we must register to scene to randomize with EventManager + self.scene.articulations["robot"] = self.hand + self.scene.rigid_objects["object"] = self.object + if self._joint_wrench_sensor is not None: + self.scene.sensors["joint_wrench"] = self._joint_wrench_sensor + # add lights + light_cfg = sim_utils.DomeLightCfg(intensity=2000.0, color=(0.75, 0.75, 0.75)) + light_cfg.func("/World/Light", light_cfg) + + def _create_joint_wrench_sensor(self) -> JointWrenchSensor: + """Create the joint-wrench sensor used for fingertip force/torque observations.""" + return JointWrenchSensor(JointWrenchSensorCfg(prim_path=self.cfg.robot_cfg.prim_path)) + + def _pre_physics_step(self, actions: torch.Tensor) -> None: + # copy in-place: keeps the persistent warp-native action buffer (and its + # cached kernel view) bound to the same memory across steps + self.actions_torch.copy_(actions) + + def _apply_action(self) -> None: + wp.launch( + ema_actuation_kernel, + dim=(self.num_envs, self.actuated_dof_ids.shape[0]), + inputs=[ + self.actions, + self.hand_dof_lower_limits, + self.hand_dof_upper_limits, + self.actuated_dof_ids, + self.cfg.act_moving_average, + self.prev_targets, + self.cur_targets, + ], + outputs=[self.compact_targets], + device=self.compact_targets.device, + ) + self._set_joint_pos_target(target=self.compact_targets_torch, joint_ids=self.actuated_dof_indices) + + def _get_observations(self) -> dict: + if self.cfg.obs_type == "openai": + obs = self.compute_reduced_observations() + elif self.cfg.obs_type == "full": + obs = self.compute_full_observations() + else: + raise ValueError(f"Unknown observation type: {self.cfg.obs_type}. Should be 'full' or 'openai'.") + + # RSL-RL holds the observation reference across the next env.step, so hand out + # a per-step snapshot rather than the persistent kernel output buffer. + # TODO: drop the per-step clone by double-buffering the observation outputs + # (pre-allocated alternating Warp buffers with cached torch bindings). + observations = {"policy": obs.clone()} + if self.cfg.asymmetric_obs: + observations["critic"] = self.compute_full_state().clone() + return observations + + def _get_rewards(self) -> torch.Tensor: + # success flags, distances, and the episode-minimum error tracking were + # already computed this step by the kernel launches in :meth:`_get_dones`; + # the goal-reset flags, success counts, and moving average update in place. + # passing the time-out flags alone is exact here: the kernel ORs in the + # fallen-object condition itself, and that is this env's only other reset + reorient_reward( + self.time_out_flags, + self.reset_goal_buf, + self.successes, + self.consecutive_successes, + self.object.data.root_pos_w.warp, + self.env_origins_warp, + self.object.data.root_quat_w.warp, + self.in_hand_pos, + self.goal_rot, + self.actions, + self.cfg.dist_reward_scale, + self.cfg.rot_reward_scale, + self.cfg.rot_eps, + self.cfg.action_penalty_scale, + self.cfg.success_tolerance, + self.cfg.reach_goal_bonus, + self.cfg.fall_dist, + self.cfg.fall_penalty, + self.cfg.av_factor, + self._reward_buffers, + ) + + if "log" not in self.extras: + self.extras["log"] = dict() + self.extras["log"]["consecutive_successes"] = self.consecutive_successes_torch.mean() + + # reset goals if the goal has been reached + goal_env_ids = self.reset_goal_buf_torch.nonzero(as_tuple=False).squeeze(-1) + if len(goal_env_ids) > 0: + self._reset_target_pose(goal_env_ids) + + return self.reward_torch + + def _get_dones(self) -> tuple[torch.Tensor, torch.Tensor]: + # reset when cube has fallen + wp.launch( + out_of_reach_kernel, + dim=self.num_envs, + inputs=[ + self.object.data.root_pos_w.warp, + self.env_origins_warp, + self.in_hand_pos, + self.cfg.fall_dist, + ], + outputs=[self.out_of_reach_flags], + device=self.out_of_reach_flags.device, + ) + # single per-step success evaluation (rewards and metrics reuse the buffers), + # followed by the progress bookkeeping (episode-counter reset in place, + # time-out by length or consecutive-success cap, episode-minimum error tracking) + wp.launch( + reorient_success_kernel, + dim=self.num_envs, + inputs=[self.object.data.root_quat_w.warp, self.goal_rot, self.cfg.success_tolerance], + outputs=[self.success_flags, self.orientation_error_buf], + device=self.orientation_error_buf.device, + ) + wp.launch( + reorient_progress_kernel, + dim=self.num_envs, + inputs=[ + self.success_flags, + self.orientation_error_buf, + self.successes, + float(self.cfg.max_consecutive_success), + self.max_episode_length, + ], + outputs=[ + self.episode_length_buf_warp, + self._orientation_error.minimum_error, + self.time_out_flags, + self._orientation_error._has_sample, + ], + device=self.orientation_error_buf.device, + ) + return self.out_of_reach_flags_torch, self.time_out_flags_torch + + def _reset_idx(self, env_ids: Sequence[int]): + # Episode counts as successful when goals reached >= cfg.success_count_threshold. + self._last_episode_success_torch[env_ids] = self.successes_torch[env_ids] >= self.cfg.success_count_threshold + # 0-dim device tensor: avoids a host sync here; consumers read it at logging cadence + self.extras.setdefault("log", {})["Metrics/success_rate"] = ( + self._last_episode_success_torch[env_ids].float().mean() + ) + for statistic, value in self._orientation_error.reset(env_ids).items(): + self.extras["log"][f"Diagnostics/episode_min_orientation_error_{statistic}"] = value + + super()._reset_idx(env_ids) + + # reset goals + self._reset_target_pose(env_ids) + + # reset object + object_default_pose = self.object.data.default_root_pose.torch.clone()[env_ids] + object_default_vel = self.object.data.default_root_vel.torch.clone()[env_ids] + pos_noise = sample_uniform(-1.0, 1.0, (len(env_ids), 3), device=self.device) + # global object positions + object_default_pose[:, 0:3] = ( + object_default_pose[:, 0:3] + self.cfg.reset_position_noise * pos_noise + self.scene.env_origins[env_ids] + ) + + rot_noise = sample_uniform(-1.0, 1.0, (len(env_ids), 2), device=self.device) # noise for X and Y rotation + object_default_pose[:, 3:7] = randomize_rotation( + rot_noise[:, 0], rot_noise[:, 1], self.x_unit_vec_torch[env_ids], self.y_unit_vec_torch[env_ids] + ) + + object_default_vel[:] = 0.0 + self._write_obj_root_pose(root_pose=object_default_pose, env_ids=env_ids) + self._write_obj_root_vel(root_velocity=object_default_vel, env_ids=env_ids) + + # reset hand + default_dof_pos = self.hand.data.default_joint_pos.torch[env_ids] + dof_limits = self.hand.data.joint_limits.torch[env_ids] + dof_pos = sample_joint_positions_within_limits(default_dof_pos, dof_limits, self.cfg.reset_dof_pos_noise) + + dof_vel_noise = sample_uniform(-1.0, 1.0, (len(env_ids), self.num_hand_dofs), device=self.device) + dof_vel = self.hand.data.default_joint_vel.torch[env_ids] + self.cfg.reset_dof_vel_noise * dof_vel_noise + + self.prev_targets_torch[env_ids] = dof_pos + self.cur_targets_torch[env_ids] = dof_pos + + self._set_joint_pos_target(target=dof_pos, env_ids=env_ids) + self._write_hand_joint_pos(position=dof_pos, env_ids=env_ids) + self._write_hand_joint_vel(velocity=dof_vel, env_ids=env_ids) + + self.successes_torch[env_ids] = 0 + self._compute_intermediate_values() + + def _reset_target_pose(self, env_ids): + # reset goal rotation + rand_floats = sample_uniform(-1.0, 1.0, (len(env_ids), 2), device=self.device) + new_rot = randomize_rotation( + rand_floats[:, 0], rand_floats[:, 1], self.x_unit_vec_torch[env_ids], self.y_unit_vec_torch[env_ids] + ) + + # update goal pose and markers + self.goal_rot_torch[env_ids] = new_rot + goal_pos = self.goal_pos_torch + self.scene.env_origins + self.goal_markers.visualize(goal_pos, self.goal_rot_torch) + + self.reset_goal_buf_torch[env_ids] = False + + def _compute_intermediate_values(self): + """Refresh the torch-side state snapshots consumed by the camera environment's observations.""" + # data for hand + self.fingertip_pos = self.hand.data.body_pos_w.torch[:, self.finger_bodies] + self.fingertip_rot = self.hand.data.body_quat_w.torch[:, self.finger_bodies] + self.fingertip_pos -= self.scene.env_origins.repeat((1, self.num_fingertips)).reshape( + self.num_envs, self.num_fingertips, 3 + ) + self.fingertip_velocities = self.hand.data.body_vel_w.torch[:, self.finger_bodies] + + self.hand_dof_pos = self.hand.data.joint_pos.torch + self.hand_dof_vel = self.hand.data.joint_vel.torch + + # data for object + self.object_pos = self.object.data.root_pos_w.torch - self.scene.env_origins + self.object_rot = self.object.data.root_quat_w.torch + self.object_velocities = self.object.data.root_vel_w.torch + self.object_linvel = self.object.data.root_lin_vel_w.torch + self.object_angvel = self.object.data.root_ang_vel_w.torch + + def compute_reduced_observations(self): + # Per https://arxiv.org/pdf/1808.00177.pdf Table 2 + wp.launch( + reduced_obs_kernel, + dim=self.num_envs, + inputs=[ + self.hand.data.body_pos_w.warp, + self.env_origins_warp, + self.finger_ids, + self.object.data.root_pos_w.warp, + self.object.data.root_quat_w.warp, + self.goal_rot, + self.actions, + ], + outputs=[self.policy_obs_buf], + device=self.policy_obs_buf.device, + ) + return self.policy_obs_buf_torch + + def compute_full_observations(self): + force, torque = self.dummy_wrench, self.dummy_wrench + wp.launch( + full_obs_kernel, + dim=(self.num_envs, self.policy_obs_buf.shape[1]), + inputs=[ + self.hand.data.joint_pos.warp, + self.hand.data.joint_vel.warp, + self.hand_dof_lower_limits, + self.hand_dof_upper_limits, + self.object.data.root_pos_w.warp, + self.env_origins_warp, + self.object.data.root_quat_w.warp, + self.object.data.root_lin_vel_w.warp, + self.object.data.root_ang_vel_w.warp, + self.in_hand_pos, + self.goal_rot, + self.hand.data.body_pos_w.warp, + self.hand.data.body_quat_w.warp, + self.hand.data.body_vel_w.warp, + self.finger_ids, + force, + torque, + self.wrench_ids, + self.actions, + self.cfg.vel_obs_scale, + self.cfg.force_torque_obs_scale, + 0, + ], + outputs=[self.policy_obs_buf], + device=self.policy_obs_buf.device, + ) + return self.policy_obs_buf_torch + + def compute_full_state(self): + force, torque, with_forces = self.dummy_wrench, self.dummy_wrench, -1 + if self._joint_wrench_sensor is not None: + force_data = self._joint_wrench_sensor.data.force + torque_data = self._joint_wrench_sensor.data.torque + if force_data is not None and torque_data is not None: + force, torque, with_forces = force_data.warp, torque_data.warp, 1 + wp.launch( + full_obs_kernel, + dim=(self.num_envs, self.state_obs_buf.shape[1]), + inputs=[ + self.hand.data.joint_pos.warp, + self.hand.data.joint_vel.warp, + self.hand_dof_lower_limits, + self.hand_dof_upper_limits, + self.object.data.root_pos_w.warp, + self.env_origins_warp, + self.object.data.root_quat_w.warp, + self.object.data.root_lin_vel_w.warp, + self.object.data.root_ang_vel_w.warp, + self.in_hand_pos, + self.goal_rot, + self.hand.data.body_pos_w.warp, + self.hand.data.body_quat_w.warp, + self.hand.data.body_vel_w.warp, + self.finger_ids, + force, + torque, + self.wrench_ids, + self.actions, + self.cfg.vel_obs_scale, + self.cfg.force_torque_obs_scale, + with_forces, + ], + outputs=[self.state_obs_buf], + device=self.state_obs_buf.device, + ) + return self.state_obs_buf_torch diff --git a/source/isaaclab_tasks_experimental/test/direct/test_dexterous_warp_kernels.py b/source/isaaclab_tasks_experimental/test/direct/test_dexterous_warp_kernels.py new file mode 100644 index 000000000000..c5f2ef6aebe1 --- /dev/null +++ b/source/isaaclab_tasks_experimental/test/direct/test_dexterous_warp_kernels.py @@ -0,0 +1,220 @@ +# 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 + +"""Behavior tests for the dexterous-task Warp kernels. + +Each test asserts kernel behavior on hand-computed cases (known rotations, +distances, and layouts) so quaternion-layout or formula regressions fail +loudly without re-implementing the replaced math as a reference. +""" + +import math + +import pytest +import torch +import warp as wp +from isaaclab_tasks_experimental.direct.handover.handover_kernels import ( + handover_reward_kernel, + handover_success_kernel, + object_goal_kernel, +) +from isaaclab_tasks_experimental.direct.reorient.reorient_kernels import ( + compute_cube_keypoints, + reorient_success_kernel, + rotation_distance_kernel, +) + +# wp.from_torch below does not auto-initialize warp +wp.init() + + +def _cube_keypoints(pose: torch.Tensor, size: tuple[float, float, float] = (0.06, 0.06, 0.06)) -> torch.Tensor: + """Torch convenience wrapper over the Warp-signature keypoints launcher.""" + out = torch.empty(pose.shape[0], 8, 3, dtype=torch.float32, device=pose.device) + compute_cube_keypoints( + wp.from_torch(pose.contiguous(), dtype=wp.float32), + wp.from_torch(out, dtype=wp.vec3), + size=size, + ) + return out + + +_DEVICES = ["cpu"] + (["cuda:0"] if torch.cuda.is_available() else []) + +_SIN45 = math.sin(math.pi / 4) +_COS45 = math.cos(math.pi / 4) +# (x, y, z, w) storage everywhere +_IDENTITY = (0.0, 0.0, 0.0, 1.0) +_ROT90_X = (_SIN45, 0.0, 0.0, _COS45) +_ROT180_Z = (0.0, 0.0, 1.0, 0.0) + + +def _quats(device, *quats): + return wp.from_torch(torch.tensor(quats, dtype=torch.float32, device=device), dtype=wp.quatf) + + +def _vecs(device, *vecs): + return wp.from_torch(torch.tensor(vecs, dtype=torch.float32, device=device), dtype=wp.vec3f) + + +@pytest.mark.parametrize("device", _DEVICES) +def test_rotation_distance_recovers_known_angles(device): + distance = torch.empty(3, device=device) + wp.launch( + rotation_distance_kernel, + dim=3, + inputs=[_quats(device, _IDENTITY, _ROT90_X, _ROT180_Z), _quats(device, _IDENTITY, _IDENTITY, _IDENTITY)], + outputs=[wp.from_torch(distance)], + device=wp.device_from_torch(distance.device), + ) + torch.testing.assert_close(distance, torch.tensor([0.0, math.pi / 2, math.pi], device=device), atol=1e-5, rtol=0.0) + + +@pytest.mark.parametrize("device", _DEVICES) +def test_reorient_success_thresholds_on_rotation_distance(device): + success = torch.empty(2, dtype=torch.bool, device=device) + distance = torch.empty(2, device=device) + wp.launch( + reorient_success_kernel, + dim=2, + inputs=[_quats(device, _IDENTITY, _ROT90_X), _quats(device, _IDENTITY, _IDENTITY), 0.4], + outputs=[wp.from_torch(success), wp.from_torch(distance)], + device=wp.device_from_torch(distance.device), + ) + assert success.tolist() == [True, False] + torch.testing.assert_close(distance, torch.tensor([0.0, math.pi / 2], device=device), atol=1e-5, rtol=0.0) + + +@pytest.mark.parametrize("device", _DEVICES) +def test_handover_success_measures_env_frame_distance(device): + success = torch.empty(2, dtype=torch.bool, device=device) + distance = torch.empty(2, device=device) + wp.launch( + handover_success_kernel, + dim=2, + inputs=[ + _vecs(device, (1.0, 2.0, 3.0), (5.0, 5.0, 5.0)), # object, world frame + _vecs(device, (1.0, 2.0, 0.0), (5.0, 5.0, 5.0)), # environment origins + _vecs(device, (0.0, 0.0, 3.05), (1.0, 0.0, 0.0)), # goals, environment frame + 0.1, + ], + outputs=[wp.from_torch(success), wp.from_torch(distance)], + device=wp.device_from_torch(distance.device), + ) + torch.testing.assert_close(distance, torch.tensor([0.05, 1.0], device=device), atol=1e-6, rtol=0.0) + assert success.tolist() == [True, False] + + +@pytest.mark.parametrize("device", _DEVICES) +def test_handover_reward_falls_off_exponentially(device): + scale = 20.0 + distance = torch.tensor([0.0, math.log(2.0) / scale], device=device) + reward = torch.empty(2, device=device) + wp.launch( + handover_reward_kernel, + dim=2, + inputs=[wp.from_torch(distance), scale], + outputs=[wp.from_torch(reward)], + device=wp.device_from_torch(reward.device), + ) + # 2 * exp(-scale * d): d = 0 -> 2.0; d = ln(2)/scale -> 1.0 + torch.testing.assert_close(reward, torch.tensor([2.0, 1.0], device=device), atol=1e-6, rtol=1e-6) + + +@pytest.mark.parametrize("device", _DEVICES) +def test_object_goal_block_layout(device): + out = torch.empty(1, 24, dtype=torch.float32, device=device) + wp.launch( + object_goal_kernel, + dim=1, + inputs=[ + _vecs(device, (1.0, 2.0, 3.0)), # object position, world + _vecs(device, (1.0, 0.0, 0.0)), # environment origin + _quats(device, _ROT90_X), # object rotation + _vecs(device, (0.5, 0.0, 0.0)), # linear velocity + _vecs(device, (0.0, 10.0, 0.0)), # angular velocity + _vecs(device, (0.0, -0.6, 0.5)), # goal position, environment frame + _quats(device, _IDENTITY), # goal rotation + 0.2, # angular-velocity scale + ], + outputs=[wp.from_torch(out)], + device=wp.device_from_torch(out.device), + ) + expected = torch.tensor( + [[0.0, 2.0, 3.0, *_ROT90_X, 0.5, 0.0, 0.0, 0.0, 2.0, 0.0, 0.0, -0.6, 0.5, *_IDENTITY, *_ROT90_X]], + device=device, + ) + torch.testing.assert_close(out, expected, atol=1e-6, rtol=0.0) + + +@pytest.mark.parametrize("device", _DEVICES) +def test_cube_keypoints_identity_pose_gives_half_side_corners(device): + pose = torch.zeros(1, 7, device=device) + pose[0, 6] = 1.0 # identity orientation (x, y, z, w) + + keypoints = _cube_keypoints(pose, size=(0.4, 0.6, 0.8)) + + corners = {tuple(round(c, 3) for c in corner) for corner in keypoints[0].tolist()} + expected = {(sx * 0.2, sy * 0.3, sz * 0.4) for sx in (1, -1) for sy in (1, -1) for sz in (1, -1)} + assert corners == expected + + +@pytest.mark.parametrize("device", _DEVICES) +@pytest.mark.parametrize("with_forces", [1, 0, -1]) +def test_full_obs_kernel_force_block_modes(device, with_forces): + """The wrench segment is written (1), omitted (0), or zero-filled without reads (-1).""" + from isaaclab_tasks_experimental.direct.reorient.reorient_kernels import full_obs_kernel + + num_joints, num_fingers = 2, 1 + # layout: 2J | 13 object | 11 goal | 13F fingertip | (6F forces unless 0) | actions + force_offset = 2 * num_joints + 13 + 11 + 13 * num_fingers + obs_dim = force_offset + (6 * num_fingers if with_forces != 0 else 0) + 2 + zeros2 = torch.zeros(1, num_joints, device=device) + f32 = {"dtype": torch.float32, "device": device} + + out = torch.full((1, obs_dim), torch.nan, **f32) + actions = torch.tensor([[0.25, -0.75]], **f32) + # with_forces == -1 models the pre-init sensor: a (1, 1) dummy array whose payload + # must never be read; the sentinel 7.0 makes an accidental read fail the zero assert + force = torch.full((1, 1) if with_forces == -1 else (1, 2), 7.0, device=device).unsqueeze(-1).repeat(1, 1, 3) + torque = 2.0 * force + wrench_ids = torch.tensor([0 if with_forces == -1 else 1], dtype=torch.int32, device=device) + wp.launch( + full_obs_kernel, + dim=(1, obs_dim), + inputs=[ + wp.from_torch(zeros2), # joint_pos + wp.from_torch(zeros2), # joint_vel + wp.from_torch(zeros2 - 1.0), # lower + wp.from_torch(zeros2 + 1.0), # upper + _vecs(device, (0.0, 0.0, 0.0)), # object_pos_w + _vecs(device, (0.0, 0.0, 0.0)), # env_origins + _quats(device, _IDENTITY), # object_quat + _vecs(device, (0.0, 0.0, 0.0)), # object_lin_vel + _vecs(device, (0.0, 0.0, 0.0)), # object_ang_vel + _vecs(device, (0.0, 0.0, 0.0)), # in_hand_pos_e + _quats(device, _IDENTITY), # goal_quat + wp.from_torch(torch.zeros(1, 2, 3, **f32), dtype=wp.vec3f), # body_pos_w + wp.from_torch(torch.zeros(1, 2, 4, **f32), dtype=wp.quatf), # body_quat_w + wp.from_torch(torch.zeros(1, 2, 6, **f32), dtype=wp.spatial_vectorf), # body_vel_w + wp.from_torch(torch.tensor([1], dtype=torch.int32, device=device)), # finger_ids + wp.from_torch(force, dtype=wp.vec3f), + wp.from_torch(torque, dtype=wp.vec3f), + wp.from_torch(wrench_ids), + wp.from_torch(actions), + 1.0, # vel_scale + 0.5, # force_scale + with_forces, + ], + outputs=[wp.from_torch(out)], + device=wp.device_from_torch(out.device), + ) + if with_forces == 1: + expected_wrench = 0.5 * torch.tensor([7.0, 7.0, 7.0, 14.0, 14.0, 14.0], device=device) + torch.testing.assert_close(out[0, force_offset : force_offset + 6], expected_wrench) + elif with_forces == -1: + torch.testing.assert_close(out[0, force_offset : force_offset + 6], torch.zeros(6, device=device)) + torch.testing.assert_close(out[0, -2:], actions[0]) + assert not out.isnan().any() diff --git a/source/isaaclab_tasks_experimental/test/direct/test_reorient_kernels.py b/source/isaaclab_tasks_experimental/test/direct/test_reorient_kernels.py new file mode 100644 index 000000000000..6e8d32db399c --- /dev/null +++ b/source/isaaclab_tasks_experimental/test/direct/test_reorient_kernels.py @@ -0,0 +1,495 @@ +# 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 + +"""Parity unit tests for the reorient family Warp kernels. + +Kit-free: imports only torch/warp and the kernel module. Every kernel is +checked against an independent torch reference implementation on randomized +inputs, on CPU and CUDA. +""" + +from __future__ import annotations + +import pytest +import torch +import warp as wp +from isaaclab_tasks_experimental.direct.reorient.reorient_kernels import ( + ReorientRewardBuffers, + compute_cube_keypoints, + cube_keypoints_from_quat_kernel, + ema_actuation_kernel, + fingertip_pos_kernel, + fingertip_quat_kernel, + fingertip_vel_kernel, + full_obs_kernel, + out_of_reach_kernel, + reduced_obs_kernel, + reorient_progress_kernel, + reorient_reward, + reorient_success_kernel, + rotation_distance_kernel, +) + +wp.init() + +N_ENVS = 33 +N_JOINTS = 24 +N_BODIES = 26 +N_FINGERS = 5 +N_ACTIONS = 20 + +DEVICES = ["cpu"] + (["cuda:0"] if torch.cuda.is_available() else []) + + +# -- torch reference helpers (independent implementations, (x, y, z, w) layout) -- + + +def _rand_unit_quat(n: int, device: str) -> torch.Tensor: + q = torch.randn(n, 4, device=device) + return torch.nn.functional.normalize(q, dim=-1) + + +def _quat_conj(q: torch.Tensor) -> torch.Tensor: + return torch.cat([-q[..., :3], q[..., 3:4]], dim=-1) + + +def _quat_mul(q1: torch.Tensor, q2: torch.Tensor) -> torch.Tensor: + x1, y1, z1, w1 = q1.unbind(-1) + x2, y2, z2, w2 = q2.unbind(-1) + return torch.stack( + [ + w1 * x2 + x1 * w2 + y1 * z2 - z1 * y2, + w1 * y2 - x1 * z2 + y1 * w2 + z1 * x2, + w1 * z2 + x1 * y2 - y1 * x2 + z1 * w2, + w1 * w2 - x1 * x2 - y1 * y2 - z1 * z2, + ], + dim=-1, + ) + + +def _quat_rotate(q: torch.Tensor, v: torch.Tensor) -> torch.Tensor: + qv, qw = q[..., :3], q[..., 3:4] + t = 2.0 * torch.cross(qv, v, dim=-1) + return v + qw * t + torch.cross(qv, t, dim=-1) + + +def _rotation_distance_ref(q_obj: torch.Tensor, q_goal: torch.Tensor) -> torch.Tensor: + qe = _quat_mul(q_obj, _quat_conj(q_goal)) + return 2.0 * torch.asin(torch.clamp(qe[..., :3].norm(dim=-1), max=1.0)) + + +def _wp2d(t: torch.Tensor, dtype) -> wp.array: + return wp.from_torch(t.contiguous(), dtype=dtype) + + +# -- tests -- + + +@pytest.mark.parametrize("device", DEVICES) +def test_rotation_distance_kernel_matches_reference(device): + torch.manual_seed(0) + q1, q2 = _rand_unit_quat(N_ENVS, device), _rand_unit_quat(N_ENVS, device) + dist = torch.empty(N_ENVS, dtype=torch.float32, device=device) + wp.launch( + rotation_distance_kernel, + dim=N_ENVS, + inputs=[_wp2d(q1, wp.quatf), _wp2d(q2, wp.quatf)], + outputs=[wp.from_torch(dist)], + device=device, + ) + torch.testing.assert_close(dist, _rotation_distance_ref(q1, q2), atol=1e-4, rtol=1e-4) + # identical quaternions have zero distance + wp.launch( + rotation_distance_kernel, + dim=N_ENVS, + inputs=[_wp2d(q1, wp.quatf), _wp2d(q1.clone(), wp.quatf)], + outputs=[wp.from_torch(dist)], + device=device, + ) + assert dist.abs().max().item() < 1e-3 + + +@pytest.mark.parametrize("device", DEVICES) +def test_success_kernel_flags_and_distance(device): + torch.manual_seed(1) + q1, q2 = _rand_unit_quat(N_ENVS, device), _rand_unit_quat(N_ENVS, device) + tol = 0.4 + success = torch.empty(N_ENVS, dtype=torch.bool, device=device) + dist = torch.empty(N_ENVS, dtype=torch.float32, device=device) + wp.launch( + reorient_success_kernel, + dim=N_ENVS, + inputs=[_wp2d(q1, wp.quatf), _wp2d(q2, wp.quatf), tol], + outputs=[wp.from_torch(success), wp.from_torch(dist)], + device=device, + ) + ref_dist = _rotation_distance_ref(q1, q2) + torch.testing.assert_close(dist, ref_dist, atol=1e-4, rtol=1e-4) + assert torch.equal(success, dist <= tol) + + +@pytest.mark.parametrize("device", DEVICES) +def test_ema_actuation_kernel_matches_reference(device): + torch.manual_seed(2) + lower = -1.5 + 0.1 * torch.rand(N_ENVS, N_JOINTS, device=device) + upper = 1.5 + 0.1 * torch.rand(N_ENVS, N_JOINTS, device=device) + dof_ids = torch.randperm(N_JOINTS, device=device)[:N_ACTIONS].to(torch.int32) + actions = 2.0 * torch.rand(N_ENVS, N_ACTIONS, device=device) - 1.0 + prev = torch.zeros(N_ENVS, N_JOINTS, device=device) + prev_ref = prev.clone() + cur = torch.zeros_like(prev) + compact = torch.zeros(N_ENVS, N_ACTIONS, device=device) + m = 0.7 + wp.launch( + ema_actuation_kernel, + dim=(N_ENVS, N_ACTIONS), + inputs=[ + _wp2d(actions, wp.float32), + _wp2d(lower, wp.float32), + _wp2d(upper, wp.float32), + wp.from_torch(dof_ids), + m, + ], + outputs=[_wp2d(prev, wp.float32), _wp2d(cur, wp.float32), _wp2d(compact, wp.float32)], + device=device, + ) + ids = dof_ids.long() + lo, hi = lower[:, ids], upper[:, ids] + t = 0.5 * (actions + 1.0) * (hi - lo) + lo + t = m * t + (1.0 - m) * prev_ref[:, ids] + t = torch.clamp(t, lo, hi) + torch.testing.assert_close(compact, t, atol=1e-5, rtol=1e-5) + torch.testing.assert_close(cur[:, ids], t, atol=1e-5, rtol=1e-5) + torch.testing.assert_close(prev[:, ids], t, atol=1e-5, rtol=1e-5) + + +@pytest.mark.parametrize("device", DEVICES) +def test_fingertip_gather_kernels_match_reference(device): + torch.manual_seed(3) + body_pos = torch.randn(N_ENVS, N_BODIES, 3, device=device) + body_quat = _rand_unit_quat(N_ENVS * N_BODIES, device).view(N_ENVS, N_BODIES, 4) + body_vel = torch.randn(N_ENVS, N_BODIES, 6, device=device) + origins = torch.randn(N_ENVS, 3, device=device) + ids = torch.randperm(N_BODIES, device=device)[:N_FINGERS].to(torch.int32) + out_pos = torch.empty(N_ENVS, 3 * N_FINGERS, device=device) + out_quat = torch.empty(N_ENVS, 4 * N_FINGERS, device=device) + out_vel = torch.empty(N_ENVS, 6 * N_FINGERS, device=device) + wp.launch( + fingertip_pos_kernel, + dim=(N_ENVS, N_FINGERS), + inputs=[_wp2d(body_pos, wp.vec3f), _wp2d(origins, wp.vec3f), wp.from_torch(ids)], + outputs=[_wp2d(out_pos, wp.float32)], + device=device, + ) + wp.launch( + fingertip_quat_kernel, + dim=(N_ENVS, N_FINGERS), + inputs=[_wp2d(body_quat, wp.quatf), wp.from_torch(ids)], + outputs=[_wp2d(out_quat, wp.float32)], + device=device, + ) + wp.launch( + fingertip_vel_kernel, + dim=(N_ENVS, N_FINGERS), + inputs=[_wp2d(body_vel, wp.spatial_vectorf), wp.from_torch(ids)], + outputs=[_wp2d(out_vel, wp.float32)], + device=device, + ) + lids = ids.long() + torch.testing.assert_close(out_pos, (body_pos[:, lids] - origins[:, None]).reshape(N_ENVS, -1)) + torch.testing.assert_close(out_quat, body_quat[:, lids].reshape(N_ENVS, -1)) + torch.testing.assert_close(out_vel, body_vel[:, lids].reshape(N_ENVS, -1)) + + +@pytest.mark.parametrize("device", DEVICES) +def test_out_of_reach_kernel_matches_reference(device): + torch.manual_seed(4) + obj = torch.randn(N_ENVS, 3, device=device) + origins = torch.randn(N_ENVS, 3, device=device) + target = torch.randn(N_ENVS, 3, device=device) + fall = 1.0 + out = torch.empty(N_ENVS, dtype=torch.bool, device=device) + wp.launch( + out_of_reach_kernel, + dim=N_ENVS, + inputs=[_wp2d(obj, wp.vec3f), _wp2d(origins, wp.vec3f), _wp2d(target, wp.vec3f), fall], + outputs=[wp.from_torch(out)], + device=device, + ) + ref = (obj - origins - target).norm(dim=-1) >= fall + assert torch.equal(out, ref) + + +@pytest.mark.parametrize("device", DEVICES) +def test_cube_keypoints_match_reference(device): + torch.manual_seed(5) + pose = torch.cat([torch.randn(N_ENVS, 3, device=device), _rand_unit_quat(N_ENVS, device)], dim=-1) + out = torch.empty(N_ENVS, 8, 3, dtype=torch.float32, device=device) + compute_cube_keypoints(_wp2d(pose, wp.float32), wp.from_torch(out, dtype=wp.vec3)) + half = torch.tensor([0.03, 0.03, 0.03], device=device) + signs = torch.tensor( + [[1 - 2 * ((c >> b) & 1) for b in range(3)] for c in range(8)], dtype=torch.float32, device=device + ) + corners = signs * half # (8, 3) + ref = pose[:, None, :3] + _quat_rotate(pose[:, None, 3:7].expand(-1, 8, -1), corners.expand(N_ENVS, -1, -1)) + torch.testing.assert_close(out, ref, atol=1e-5, rtol=1e-5) + # quat-only variant: rotated offsets without the position term + out2 = torch.empty(N_ENVS, 24, dtype=torch.float32, device=device) + wp.launch( + cube_keypoints_from_quat_kernel, + dim=(N_ENVS, 8), + inputs=[_wp2d(pose[:, 3:7], wp.quatf), wp.vec3(0.03, 0.03, 0.03)], + outputs=[_wp2d(out2, wp.float32)], + device=device, + ) + ref2 = _quat_rotate(pose[:, None, 3:7].expand(-1, 8, -1), corners.expand(N_ENVS, -1, -1)) + torch.testing.assert_close(out2.view(N_ENVS, 8, 3), ref2, atol=1e-5, rtol=1e-5) + + +@pytest.mark.parametrize("device", DEVICES) +def test_reorient_reward_matches_reference(device): + torch.manual_seed(6) + obj = 0.2 * torch.randn(N_ENVS, 3, device=device) + origins = torch.randn(N_ENVS, 3, device=device) + target_pos = 0.2 * torch.randn(N_ENVS, 3, device=device) + q_obj, q_goal = _rand_unit_quat(N_ENVS, device), _rand_unit_quat(N_ENVS, device) + actions = 2.0 * torch.rand(N_ENVS, N_ACTIONS, device=device) - 1.0 + reset_buf = torch.rand(N_ENVS, device=device) < 0.2 + reset_goal_buf = torch.rand(N_ENVS, device=device) < 0.2 + successes_t = torch.randint(0, 5, (N_ENVS,), device=device).float() + goal_resets = wp.from_torch(reset_goal_buf.clone()) + successes = wp.from_torch(successes_t.clone()) + consecutive_t = torch.tensor([1.5], device=device) + consecutive = wp.array([1.5], dtype=wp.float32, device=device) + params = dict( + distance_scale=-10.0, + rotation_scale=1.0, + rotation_epsilon=0.1, + action_penalty_scale=-0.0002, + success_tolerance=0.6, + success_bonus=250.0, + fall_distance=0.24, + fall_penalty=-50.0, + averaging_factor=0.1, + ) + buffers = ReorientRewardBuffers(N_ENVS, device) + reward = reorient_reward( + wp.from_torch(reset_buf), + goal_resets, + successes, + consecutive, + _wp2d(obj, wp.vec3f), + _wp2d(origins, wp.vec3f), + _wp2d(q_obj, wp.quatf), + _wp2d(target_pos, wp.vec3f), + _wp2d(q_goal, wp.quatf), + _wp2d(actions, wp.float32), + buffers=buffers, + **params, + ) + # torch reference + goal_dist = (obj - origins - target_pos).norm(dim=-1) + rot_dist = _rotation_distance_ref(q_obj, q_goal) + goal_reset_ref = (rot_dist <= params["success_tolerance"]) | reset_goal_buf + penalty = (actions**2).sum(dim=-1) + value = ( + goal_dist * params["distance_scale"] + + params["rotation_scale"] / (rot_dist + params["rotation_epsilon"]) + + penalty * params["action_penalty_scale"] + ) + value = value + torch.where(goal_reset_ref, params["success_bonus"], 0.0) + fell = goal_dist >= params["fall_distance"] + value = value + torch.where(fell, params["fall_penalty"], 0.0) + resets_ref = (fell | reset_buf).float() + succ_ref = successes_t + goal_reset_ref.float() + num_resets = resets_ref.sum() + finished = (succ_ref * resets_ref).sum() + consec_ref = torch.where( + num_resets > 0, params["averaging_factor"] * finished / num_resets + 0.9 * consecutive_t, consecutive_t + ) + torch.testing.assert_close(wp.to_torch(reward), value, atol=1e-3, rtol=1e-4) + assert torch.equal(wp.to_torch(goal_resets), goal_reset_ref) + torch.testing.assert_close(wp.to_torch(successes), succ_ref) + torch.testing.assert_close(wp.to_torch(consecutive), consec_ref, atol=1e-4, rtol=1e-4) + + +def _obs_inputs(device): + torch.manual_seed(7) + data = dict( + joint_pos=torch.rand(N_ENVS, N_JOINTS, device=device) - 0.5, + joint_vel=torch.randn(N_ENVS, N_JOINTS, device=device), + lower=-1.0 - torch.rand(N_ENVS, N_JOINTS, device=device), + upper=1.0 + torch.rand(N_ENVS, N_JOINTS, device=device), + obj=torch.randn(N_ENVS, 3, device=device), + origins=torch.randn(N_ENVS, 3, device=device), + q_obj=_rand_unit_quat(N_ENVS, device), + lin_vel=torch.randn(N_ENVS, 3, device=device), + ang_vel=torch.randn(N_ENVS, 3, device=device), + in_hand=torch.randn(N_ENVS, 3, device=device), + q_goal=_rand_unit_quat(N_ENVS, device), + body_pos=torch.randn(N_ENVS, N_BODIES, 3, device=device), + body_quat=_rand_unit_quat(N_ENVS * N_BODIES, device).view(N_ENVS, N_BODIES, 4), + body_vel=torch.randn(N_ENVS, N_BODIES, 6, device=device), + force=torch.randn(N_ENVS, N_BODIES, 3, device=device), + torque=torch.randn(N_ENVS, N_BODIES, 3, device=device), + actions=2.0 * torch.rand(N_ENVS, N_ACTIONS, device=device) - 1.0, + ) + data["finger_ids"] = torch.randperm(N_BODIES, device=device)[:N_FINGERS].to(torch.int32) + data["wrench_ids"] = torch.randperm(N_BODIES, device=device)[:N_FINGERS].to(torch.int32) + return data + + +def _full_obs_ref(d, vel_scale, force_scale, with_forces): + fids, wids = d["finger_ids"].long(), d["wrench_ids"].long() + qe = _quat_mul(d["q_obj"], _quat_conj(d["q_goal"])) + segs = [ + 2.0 * (d["joint_pos"] - d["lower"]) / (d["upper"] - d["lower"]) - 1.0, + vel_scale * d["joint_vel"], + d["obj"] - d["origins"], + d["q_obj"], + d["lin_vel"], + vel_scale * d["ang_vel"], + d["in_hand"], + d["q_goal"], + qe, + (d["body_pos"][:, fids] - d["origins"][:, None]).reshape(N_ENVS, -1), + d["body_quat"][:, fids].reshape(N_ENVS, -1), + d["body_vel"][:, fids].reshape(N_ENVS, -1), + ] + # with_forces == 0: observation has no wrench segment at all; any other value + # reserves the segment — real data for 1, zero block otherwise (sensor not ready) + if with_forces != 0: + wrench = torch.cat([d["force"][:, wids, None, :], d["torque"][:, wids, None, :]], dim=2) + wrench = force_scale * wrench.reshape(N_ENVS, -1) + if with_forces != 1: + wrench = torch.zeros_like(wrench) + segs.append(wrench) + segs.append(d["actions"]) + return torch.cat(segs, dim=-1) + + +@pytest.mark.parametrize("device", DEVICES) +@pytest.mark.parametrize("with_forces", [1, 0, 2]) +def test_full_obs_kernel_matches_reference(device, with_forces): + d = _obs_inputs(device) + vel_scale, force_scale = 0.2, 0.05 + dim = 2 * N_JOINTS + 13 + 11 + 13 * N_FINGERS + N_ACTIONS + if with_forces != 0: + dim += 6 * N_FINGERS + out = torch.empty(N_ENVS, dim, dtype=torch.float32, device=device) + wp.launch( + full_obs_kernel, + dim=(N_ENVS, dim), + inputs=[ + _wp2d(d["joint_pos"], wp.float32), + _wp2d(d["joint_vel"], wp.float32), + _wp2d(d["lower"], wp.float32), + _wp2d(d["upper"], wp.float32), + _wp2d(d["obj"], wp.vec3f), + _wp2d(d["origins"], wp.vec3f), + _wp2d(d["q_obj"], wp.quatf), + _wp2d(d["lin_vel"], wp.vec3f), + _wp2d(d["ang_vel"], wp.vec3f), + _wp2d(d["in_hand"], wp.vec3f), + _wp2d(d["q_goal"], wp.quatf), + _wp2d(d["body_pos"], wp.vec3f), + _wp2d(d["body_quat"], wp.quatf), + _wp2d(d["body_vel"], wp.spatial_vectorf), + wp.from_torch(d["finger_ids"]), + _wp2d(d["force"], wp.vec3f), + _wp2d(d["torque"], wp.vec3f), + wp.from_torch(d["wrench_ids"]), + _wp2d(d["actions"], wp.float32), + vel_scale, + force_scale, + with_forces, + ], + outputs=[_wp2d(out, wp.float32)], + device=device, + ) + ref = _full_obs_ref(d, vel_scale, force_scale, with_forces) + torch.testing.assert_close(out, ref, atol=1e-5, rtol=1e-5) + + +@pytest.mark.parametrize("device", DEVICES) +def test_reduced_obs_kernel_matches_reference(device): + d = _obs_inputs(device) + fids = d["finger_ids"].long() + dim = 3 * N_FINGERS + 7 + N_ACTIONS + out = torch.empty(N_ENVS, dim, dtype=torch.float32, device=device) + wp.launch( + reduced_obs_kernel, + dim=N_ENVS, + inputs=[ + _wp2d(d["body_pos"], wp.vec3f), + _wp2d(d["origins"], wp.vec3f), + wp.from_torch(d["finger_ids"]), + _wp2d(d["obj"], wp.vec3f), + _wp2d(d["q_obj"], wp.quatf), + _wp2d(d["q_goal"], wp.quatf), + _wp2d(d["actions"], wp.float32), + ], + outputs=[_wp2d(out, wp.float32)], + device=device, + ) + qe = _quat_mul(d["q_obj"], _quat_conj(d["q_goal"])) + ref = torch.cat( + [ + (d["body_pos"][:, fids] - d["origins"][:, None]).reshape(N_ENVS, -1), + d["obj"] - d["origins"], + qe, + d["actions"], + ], + dim=-1, + ) + torch.testing.assert_close(out, ref, atol=1e-5, rtol=1e-5) + + +@pytest.mark.parametrize("device", DEVICES) +@pytest.mark.parametrize("max_consecutive", [50.0, 0.0]) +def test_progress_kernel_matches_reference(device, max_consecutive): + torch.manual_seed(8) + success = torch.rand(N_ENVS, device=device) < 0.3 + distance = torch.rand(N_ENVS, device=device) + distance[0] = float("inf") # non-finite samples must not touch the tracking state + successes = torch.randint(0, 100, (N_ENVS,), device=device).float() + ep_len = torch.randint(1, 400, (N_ENVS,), device=device, dtype=torch.long) + ep_len_ref = ep_len.clone() + max_ep_len = 300 # below the episode-length range so both time-out outcomes occur + time_out = torch.empty(N_ENVS, dtype=torch.bool, device=device) + minimum = torch.full((N_ENVS,), 0.5, device=device) + minimum_ref = minimum.clone() + has_sample = torch.zeros(N_ENVS, dtype=torch.bool, device=device) + wp.launch( + reorient_progress_kernel, + dim=N_ENVS, + inputs=[ + wp.from_torch(success), + wp.from_torch(distance), + wp.from_torch(successes), + max_consecutive, + max_ep_len, + ], + outputs=[ + wp.from_torch(ep_len), + wp.from_torch(minimum), + wp.from_torch(time_out), + wp.from_torch(has_sample), + ], + device=device, + ) + if max_consecutive > 0.0: + ep_ref = torch.where(success, torch.zeros_like(ep_len_ref), ep_len_ref) + reached_ref = successes >= max_consecutive + else: + ep_ref = ep_len_ref + reached_ref = torch.zeros(N_ENVS, dtype=torch.bool, device=device) + time_out_ref = (ep_ref >= max_ep_len - 1) | reached_ref + finite = torch.isfinite(distance) + min_ref = torch.where(finite & (distance < minimum_ref), distance, minimum_ref) + torch.testing.assert_close(ep_len, ep_ref) + assert torch.equal(time_out, time_out_ref) + torch.testing.assert_close(minimum, min_ref) + assert torch.equal(has_sample, finite)