From 803d0d781f38a84dde7dd24d7ad7213bef45ef5c Mon Sep 17 00:00:00 2001 From: jichuanh Date: Wed, 8 Jul 2026 11:45:27 +0000 Subject: [PATCH 1/2] Add renderer presets to the Direct camera task The Shadow Hand camera environment gains an RGB-depth camera preset so it can train with the Newton Warp renderer in addition to RTX, and the feature extractor exposes compute_cube_keypoints (the module-level compute_keypoints shim is deprecated in its favor). Validated by full camera training runs and replay video generation. --- .../task-cleanup-dex-part06.minor.rst | 11 ++ .../config/shadow_hand/feature_extractor.py | 147 ++++++++++++++++++ .../shadow_hand/shadow_hand_camera_env.py | 48 +++--- .../shadow_hand/shadow_hand_camera_env_cfg.py | 62 +++++--- .../core/test_shadow_hand_camera_presets.py | 5 +- 5 files changed, 224 insertions(+), 49 deletions(-) create mode 100644 source/isaaclab_tasks/changelog.d/task-cleanup-dex-part06.minor.rst diff --git a/source/isaaclab_tasks/changelog.d/task-cleanup-dex-part06.minor.rst b/source/isaaclab_tasks/changelog.d/task-cleanup-dex-part06.minor.rst new file mode 100644 index 000000000000..cdbe9e4a16d2 --- /dev/null +++ b/source/isaaclab_tasks/changelog.d/task-cleanup-dex-part06.minor.rst @@ -0,0 +1,11 @@ +Added +^^^^^ + +* Added an RGB-depth camera preset for training the Shadow Hand camera + environment with the Newton Warp renderer. + +Deprecated +^^^^^^^^^^ + +* Deprecated ``shadow_hand_camera_env.compute_keypoints`` in favor of + ``feature_extractor.compute_cube_keypoints``. 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..7f5063bf474d 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,15 +3,26 @@ # # SPDX-License-Identifier: BSD-3-Clause +from __future__ import annotations + import glob import os +from collections.abc import Sequence +from typing import TYPE_CHECKING import torch import torch.nn as nn import torchvision +from isaaclab.managers import ManagerTermBase, ObservationTermCfg, SceneEntityCfg from isaaclab.sensors import save_images_to_file from isaaclab.utils.configclass import configclass +from isaaclab.utils.math import quat_apply + +if TYPE_CHECKING: + from isaaclab.assets import RigidObject + from isaaclab.envs import ManagerBasedRLEnv + from isaaclab.sensors import Camera # Number of output channels for each supported camera data type. _DATA_TYPE_CHANNELS: dict[str, int] = { @@ -139,6 +150,37 @@ class FeatureExtractorCfg: """ +@torch.jit.script +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] + 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): + positive_axes = [((i >> axis) & 1) == 0 for axis in range(3)] + corner_values = ([(1 if positive_axes[axis] else -1) * side / 2 for axis, side in enumerate(size)],) + corner = torch.tensor(corner_values, dtype=torch.float32, device=pose.device) * out[:, i, :] + out[:, i, :] = pose[:, :3] + quat_apply(pose[:, 3:7], corner) + return out + + class FeatureExtractor: """Class for extracting features from image data. @@ -323,3 +365,108 @@ def step( else: predicted_pose = self.feature_extractor(img_input) return None, predicted_pose + + +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] + 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 + + 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) + 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: + """Return zero-origin cube keypoints for the current goal orientation. + + Args: + env: Environment containing the goal command. + command_name: Goal command term name. + + Returns: + Flattened zero-origin cube keypoints [m], shape ``(num_envs, 24)``. + """ + command = env.command_manager.get_command(command_name) + goal_pose = torch.cat((torch.zeros_like(command[:, :3]), command[:, 3:7]), dim=-1) + return compute_cube_keypoints(goal_pose).flatten(start_dim=1) 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_camera_env.py index 90de75ba3ee7..1d9b1e74b19b 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_camera_env.py @@ -6,6 +6,7 @@ from __future__ import annotations +import warnings from typing import TYPE_CHECKING import torch @@ -14,9 +15,9 @@ 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.config.shadow_hand.feature_extractor import FeatureExtractor, compute_cube_keypoints from isaaclab_tasks.core.reorient.reorient_direct_env import ReorientDirectEnv if TYPE_CHECKING: @@ -66,7 +67,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 +79,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 +139,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_camera_env_cfg.py b/source/isaaclab_tasks/isaaclab_tasks/core/reorient/config/shadow_hand/shadow_hand_camera_env_cfg.py index 3ecb37b503d7..0a37bc5c7362 100644 --- 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 @@ -16,6 +16,41 @@ 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", "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 _ShadowHandBaseTiledCameraCfg(CameraCfg): """Base camera configuration for the shadow hand vision environment. @@ -49,6 +84,7 @@ class ShadowHandTiledCameraCfg(PresetCfg): 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) @@ -70,6 +106,9 @@ class ShadowHandTiledCameraCfg(PresetCfg): rgb: _ShadowHandBaseTiledCameraCfg = _ShadowHandBaseTiledCameraCfg(data_types=["rgb"]) """RGB only (3 CNN input channels).""" + rgb_depth: _ShadowHandBaseTiledCameraCfg = _ShadowHandBaseTiledCameraCfg(data_types=["rgb", "depth"]) + """RGB and depth (4 CNN input channels).""" + albedo: _ShadowHandBaseTiledCameraCfg = _ShadowHandBaseTiledCameraCfg(data_types=["albedo"]) """Albedo (3 CNN input channels).""" @@ -123,28 +162,7 @@ class ShadowHandCameraEnvCfg(ShadowHandEnvCfg): 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", "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." - ) + validate_shadow_hand_camera_settings(self.tiled_camera, self.feature_extractor) @configclass 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..f0faacc1dd76 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 @@ -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 From 75cd2a5adb122e2ee12cd92f6d2366f1229fc92d Mon Sep 17 00:00:00 2001 From: jichuanh Date: Thu, 9 Jul 2026 04:18:21 +0000 Subject: [PATCH 2/2] Restrict this part to the renderer preset Strict review scope: the camera manager enablement (validation helper, keypoint relocation, manager observation terms) moves to the camera manager part; this part keeps the rgb_depth preset, its tests, and the Benchmark registration deprecation. --- .../task-cleanup-dex-part06.minor.rst | 5 +- .../config/shadow_hand/feature_extractor.py | 147 ------------------ .../shadow_hand/shadow_hand_camera_env.py | 48 +++--- .../shadow_hand/shadow_hand_camera_env_cfg.py | 63 ++++---- 4 files changed, 55 insertions(+), 208 deletions(-) diff --git a/source/isaaclab_tasks/changelog.d/task-cleanup-dex-part06.minor.rst b/source/isaaclab_tasks/changelog.d/task-cleanup-dex-part06.minor.rst index cdbe9e4a16d2..e33743b348de 100644 --- a/source/isaaclab_tasks/changelog.d/task-cleanup-dex-part06.minor.rst +++ b/source/isaaclab_tasks/changelog.d/task-cleanup-dex-part06.minor.rst @@ -7,5 +7,6 @@ Added Deprecated ^^^^^^^^^^ -* Deprecated ``shadow_hand_camera_env.compute_keypoints`` in favor of - ``feature_extractor.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. 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 7f5063bf474d..56c159a1446e 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,26 +3,15 @@ # # SPDX-License-Identifier: BSD-3-Clause -from __future__ import annotations - import glob import os -from collections.abc import Sequence -from typing import TYPE_CHECKING import torch import torch.nn as nn import torchvision -from isaaclab.managers import ManagerTermBase, ObservationTermCfg, SceneEntityCfg from isaaclab.sensors import save_images_to_file from isaaclab.utils.configclass import configclass -from isaaclab.utils.math import quat_apply - -if TYPE_CHECKING: - from isaaclab.assets import RigidObject - from isaaclab.envs import ManagerBasedRLEnv - from isaaclab.sensors import Camera # Number of output channels for each supported camera data type. _DATA_TYPE_CHANNELS: dict[str, int] = { @@ -150,37 +139,6 @@ class FeatureExtractorCfg: """ -@torch.jit.script -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] - 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): - positive_axes = [((i >> axis) & 1) == 0 for axis in range(3)] - corner_values = ([(1 if positive_axes[axis] else -1) * side / 2 for axis, side in enumerate(size)],) - corner = torch.tensor(corner_values, dtype=torch.float32, device=pose.device) * out[:, i, :] - out[:, i, :] = pose[:, :3] + quat_apply(pose[:, 3:7], corner) - return out - - class FeatureExtractor: """Class for extracting features from image data. @@ -365,108 +323,3 @@ def step( else: predicted_pose = self.feature_extractor(img_input) return None, predicted_pose - - -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] - 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 - - 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) - 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: - """Return zero-origin cube keypoints for the current goal orientation. - - Args: - env: Environment containing the goal command. - command_name: Goal command term name. - - Returns: - Flattened zero-origin cube keypoints [m], shape ``(num_envs, 24)``. - """ - command = env.command_manager.get_command(command_name) - goal_pose = torch.cat((torch.zeros_like(command[:, :3]), command[:, 3:7]), dim=-1) - return compute_cube_keypoints(goal_pose).flatten(start_dim=1) 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_camera_env.py index 1d9b1e74b19b..90de75ba3ee7 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_camera_env.py @@ -6,7 +6,6 @@ from __future__ import annotations -import warnings from typing import TYPE_CHECKING import torch @@ -15,9 +14,9 @@ from isaaclab import cloner from isaaclab.assets import Articulation, RigidObject from isaaclab.sensors import Camera -from isaaclab.utils.math import scale_transform +from isaaclab.utils.math import quat_apply, scale_transform -from isaaclab_tasks.core.reorient.config.shadow_hand.feature_extractor import FeatureExtractor, compute_cube_keypoints +from isaaclab_tasks.core.reorient.config.shadow_hand.feature_extractor import FeatureExtractor from isaaclab_tasks.core.reorient.reorient_direct_env import ReorientDirectEnv if TYPE_CHECKING: @@ -67,7 +66,7 @@ def _setup_scene(self): def _compute_image_observations(self): # generate ground truth keypoints for in-hand cube - compute_cube_keypoints(pose=torch.cat((self.object_pos, self.object_rot), dim=1), out=self.gt_keypoints) + compute_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) @@ -79,7 +78,7 @@ def _compute_image_observations(self): self.embeddings = embeddings.clone().detach() # compute keypoints for goal cube - compute_cube_keypoints( + compute_keypoints( pose=torch.cat((torch.zeros_like(self.goal_pos), self.goal_rot), dim=-1), out=self.goal_keypoints ) @@ -139,29 +138,32 @@ 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, -) -> torch.Tensor: - """Compute cube keypoints using the shared implementation. - - .. deprecated:: 9.0.0 - Use :func:`compute_cube_keypoints` instead. +): + """Computes positions of 8 corner keypoints of a cube. 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)``. + 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. """ - 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) + 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 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 index 0a37bc5c7362..38c36c4e4a29 100644 --- 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 @@ -16,41 +16,6 @@ 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", "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 _ShadowHandBaseTiledCameraCfg(CameraCfg): """Base camera configuration for the shadow hand vision environment. @@ -162,7 +127,28 @@ class ShadowHandCameraEnvCfg(ShadowHandEnvCfg): def validate_config(self): """Check renderer/data-type and feature-extractor compatibility.""" - validate_shadow_hand_camera_settings(self.tiled_camera, self.feature_extractor) + renderer_type = getattr(self.tiled_camera.renderer_cfg, "renderer_type", None) + warp_supported = {"rgb", "depth", "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 @@ -177,6 +163,11 @@ class ShadowHandCameraEnvPlayCfg(ShadowHandCameraEnvCfg): 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.