diff --git a/scripts/reinforcement_learning/leapp/deploy.py b/scripts/reinforcement_learning/leapp/deploy.py new file mode 100644 index 000000000000..5b032daa9dda --- /dev/null +++ b/scripts/reinforcement_learning/leapp/deploy.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 + +"""Deploy a LEAPP-exported policy in an Isaac Lab simulation.""" + +"""Launch Isaac Sim Simulator first.""" + +import argparse +import sys + +from isaaclab.app import AppLauncher + +parser = argparse.ArgumentParser(description="Deploy a LEAPP-exported policy in simulation.") +parser.add_argument("--task", type=str, default=None, help="Name of the registered Isaac Lab task.") +parser.add_argument("--leapp_model", type=str, default=None, help="Path to the LEAPP .yaml pipeline description.") +parser.add_argument("--seed", type=int, default=None, help="Seed for the environment.") +AppLauncher.add_app_launcher_args(parser) +args_cli, hydra_args = parser.parse_known_args() + +if args_cli.task is None or args_cli.leapp_model is None: + missing_args = [] + if args_cli.task is None: + missing_args.append("--task") + if args_cli.leapp_model is None: + missing_args.append("--leapp_model") + parser.error(f"the following arguments are required: {', '.join(missing_args)}") + +sys.argv = [sys.argv[0]] + hydra_args + +app_launcher = AppLauncher(args_cli) +simulation_app = app_launcher.app + +"""Rest everything follows.""" + +import torch + +from isaaclab.envs.leapp_deployment_env import LeappDeploymentEnv + +import isaaclab_tasks # noqa: F401 +from isaaclab_tasks.utils.parse_cfg import load_cfg_from_registry + + +def main(): + # ── Load env config from gym registry ───────────────────────── + task_name = args_cli.task.split(":")[-1] + env_cfg = load_cfg_from_registry(task_name, "env_cfg_entry_point") + + if args_cli.seed is not None: + env_cfg.seed = args_cli.seed + if args_cli.device is not None: + env_cfg.sim.device = args_cli.device + + # ── Create deploy env ───────────────────────────────────────── + env = LeappDeploymentEnv(env_cfg, args_cli.leapp_model) + + if getattr(args_cli, "headless", False): + print( + "[WARN]: Running deploy without a viewport. This happens when headless mode is active, " + "including the default case where no visualizer was selected. The policy may be " + "stepping normally, but no viewport will appear unless you specify the " + "`--visualizer` field." + ) + + print(f"[INFO]: Deploying task '{task_name}' with LEAPP model: {args_cli.leapp_model}") + print(f"[INFO]: Num envs: {env.num_envs}, decimation: {env.cfg.decimation}, step_dt: {env.step_dt:.4f}s") + + # ── Run loop ────────────────────────────────────────────────── + env.reset() + try: + with torch.inference_mode(): + while simulation_app.is_running(): + env.step() + env.close() + except KeyboardInterrupt: + pass + + +if __name__ == "__main__": + main() + simulation_app.close() diff --git a/scripts/reinforcement_learning/leapp/rsl_rl/export.py b/scripts/reinforcement_learning/leapp/rsl_rl/export.py new file mode 100644 index 000000000000..62f6d15c8a1e --- /dev/null +++ b/scripts/reinforcement_learning/leapp/rsl_rl/export.py @@ -0,0 +1,426 @@ +# 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 + +"""Script to export a checkpoint if an RL agent from RSL-RL.""" + +from __future__ import annotations + +import argparse +import contextlib +import importlib.metadata as metadata +import os +import sys +import time +from collections.abc import Mapping +from pathlib import Path + +from isaaclab.app import AppLauncher + +from isaaclab_tasks.utils import setup_preset_cli + +_RSL_RL_SCRIPTS_DIR = Path(__file__).resolve().parents[2] / "rsl_rl" +if str(_RSL_RL_SCRIPTS_DIR) not in sys.path: + sys.path.insert(0, str(_RSL_RL_SCRIPTS_DIR)) +import cli_args # isort: skip + + +RSL_RL_MIN_VERSION = "5.0.1" +_RUNTIME_IMPORTS_LOADED = False + +# Keep heavy/runtime-sensitive imports out of module import time. The CLI needs +# to parse launcher arguments and start Isaac Sim/Kit before importing torch, +# LEAPP, RSL-RL, and task modules; importing them earlier has caused launcher +# import-order failures. ``_load_runtime_dependencies()`` populates these +# globals immediately before export execution. +torch = None +leapp = None +annotate = None +gym = None +DistillationRunner = None +OnPolicyRunner = None +ManagerBasedRLEnv = None +RslRlVecEnvWrapper = None +handle_deprecated_rsl_rl_cfg = None +retrieve_file_path = None +patch_env_for_export = None +ensure_env_spec_id = None +get_published_pretrained_checkpoint = None +get_checkpoint_path = None +hydra_task_config = None +installed_version = None + + +def create_arg_parser() -> argparse.ArgumentParser: + """Create the command-line parser for RSL-RL policy export.""" + parser = argparse.ArgumentParser(description="Export an RL agent with RSL-RL.") + parser.add_argument( + "--disable_fabric", action="store_true", default=False, help="Disable fabric and use USD I/O operations." + ) + parser.add_argument("--task", type=str, default=None, help="Name of the task.") + parser.add_argument( + "--agent", type=str, default="rsl_rl_cfg_entry_point", help="Name of the RL agent configuration entry point." + ) + parser.add_argument("--seed", type=int, default=None, help="Seed used for the environment") + parser.add_argument( + "--use_pretrained_checkpoint", + action="store_true", + help="Use the pre-trained checkpoint from Nucleus.", + ) + + # LEAPP arguments + parser.add_argument( + "--export_task_name", + type=str, + default=None, + help="Name of the exported graph. Defaults to the task name.", + ) + parser.add_argument( + "--export_method", + type=str, + default="onnx-dynamo", + choices=["onnx-dynamo", "onnx-torchscript", "jit-script", "jit-trace"], + help="Method to export the policy", + ) + parser.add_argument( + "--export_save_path", + type=str, + default=None, + help="Path to save the exported model", + ) + parser.add_argument( + "--validation_steps", + type=int, + default=5, + help="Number of steps to validate the exported model", + ) + parser.add_argument( + "--disable_graph_visualization", + action="store_true", + default=False, + help="Disable LEAPP graph visualization during compile_graph().", + ) + + cli_args.add_rsl_rl_args(parser) + AppLauncher.add_app_launcher_args(parser) + return parser + + +def parse_export_args(argv: list[str] | None = None) -> tuple[argparse.Namespace, list[str]]: + """Parse export arguments and return remaining Hydra overrides.""" + parser = create_arg_parser() + # setup_preset_cli attaches the preset-selection help group then parses; + # remainder still carries typed selectors (physics=/renderer=/presets=) + # verbatim for run_export_with_hydra to fold before invoking Hydra. + args_cli, hydra_args = setup_preset_cli(parser, argv) + args_cli.headless = True + return args_cli, hydra_args + + +def _load_runtime_dependencies() -> None: + """Import runtime dependencies after Isaac Sim has been launched.""" + global _RUNTIME_IMPORTS_LOADED + global annotate, leapp, torch + global DistillationRunner, ManagerBasedRLEnv, OnPolicyRunner, RslRlVecEnvWrapper, get_checkpoint_path, gym + global ensure_env_spec_id, get_published_pretrained_checkpoint, handle_deprecated_rsl_rl_cfg, hydra_task_config + global installed_version + global patch_env_for_export, retrieve_file_path + + if _RUNTIME_IMPORTS_LOADED: + return + + try: + import leapp as leapp_module + except ImportError as e: + raise ImportError("LEAPP package is required for policy export. Install with: pip install leapp") from e + annotate_module = getattr(leapp_module, "annotate") + + import gymnasium as gym_module + import torch as torch_module + from packaging import version as packaging_version_module + from rsl_rl.runners import DistillationRunner as DistillationRunnerCls + from rsl_rl.runners import OnPolicyRunner as OnPolicyRunnerCls + + # Disable TorchScript before importing task/environment modules so any + # @torch.jit.script helpers resolve to plain Python functions during export. + torch_module.jit._state.disable() + + from isaaclab.envs import ManagerBasedRLEnv as ManagerBasedRLEnvCls + from isaaclab.utils.assets import retrieve_file_path as retrieve_file_path_fn + from isaaclab.utils.leapp import patch_env_for_export as patch_env_for_export_fn + from isaaclab.utils.leapp.utils import ensure_env_spec_id as ensure_env_spec_id_fn + + from isaaclab_rl.rsl_rl import RslRlVecEnvWrapper as RslRlVecEnvWrapperCls + from isaaclab_rl.rsl_rl import handle_deprecated_rsl_rl_cfg as handle_deprecated_rsl_rl_cfg_fn + from isaaclab_rl.utils.pretrained_checkpoint import ( + get_published_pretrained_checkpoint as get_published_pretrained_checkpoint_fn, + ) + + __import__("isaaclab_tasks") + from isaaclab_tasks.utils import get_checkpoint_path as get_checkpoint_path_fn + from isaaclab_tasks.utils.hydra import hydra_task_config as hydra_task_config_fn + + installed_version = metadata.version("rsl-rl-lib") + if packaging_version_module.parse(installed_version) < packaging_version_module.parse(RSL_RL_MIN_VERSION): + print( + f"[WARNING] LEAPP RSL-RL export is validated with rsl-rl-lib {RSL_RL_MIN_VERSION} or newer. " + f"Installed version is '{installed_version}'." + ) + + torch = torch_module + leapp = leapp_module + annotate = annotate_module + gym = gym_module + DistillationRunner = DistillationRunnerCls + OnPolicyRunner = OnPolicyRunnerCls + ManagerBasedRLEnv = ManagerBasedRLEnvCls + RslRlVecEnvWrapper = RslRlVecEnvWrapperCls + handle_deprecated_rsl_rl_cfg = handle_deprecated_rsl_rl_cfg_fn + retrieve_file_path = retrieve_file_path_fn + patch_env_for_export = patch_env_for_export_fn + ensure_env_spec_id = ensure_env_spec_id_fn + get_published_pretrained_checkpoint = get_published_pretrained_checkpoint_fn + get_checkpoint_path = get_checkpoint_path_fn + hydra_task_config = hydra_task_config_fn + _RUNTIME_IMPORTS_LOADED = True + + +def get_actor_memory_module(policy): + """Return the actor-side RNN module for supported RSL-RL recurrent policies.""" + if hasattr(policy, "rnn"): + return policy.rnn + return None + + +def is_actor_recurrent_policy(policy) -> bool: + """Return whether the actor policy has a supported recurrent state container.""" + return bool(getattr(policy, "is_recurrent", False) and get_actor_memory_module(policy) is not None) + + +def get_actor_hidden_state(policy): + """Return the actor-side recurrent hidden state for supported RSL-RL policy APIs.""" + if hasattr(policy, "get_hidden_state"): + return policy.get_hidden_state() + memory = get_actor_memory_module(policy) + return None if memory is None else getattr(memory, "hidden_state", None) + + +def set_actor_hidden_state(policy, actor_hidden) -> None: + """Assign the actor-side recurrent hidden state for supported RSL-RL policy APIs.""" + memory = get_actor_memory_module(policy) + if memory is not None: + memory.hidden_state = actor_hidden + + +def ensure_actor_hidden_state_initialized(policy, batch_size: int, device, dtype): + """Initialize and return the actor hidden state when a recurrent policy has not created it yet.""" + # ``torch`` is a lazy runtime global populated by ``_load_runtime_dependencies()`` + # after Isaac Sim launches and before export calls this helper. + assert torch is not None + actor_state = get_actor_hidden_state(policy) + if actor_state is not None: + return actor_state + + memory = get_actor_memory_module(policy) + if memory is None or not hasattr(memory, "rnn"): + return None + + num_layers = memory.rnn.num_layers + hidden_size = memory.rnn.hidden_size + zeros = torch.zeros(num_layers, batch_size, hidden_size, device=device, dtype=dtype) + if isinstance(memory.rnn, torch.nn.LSTM): + actor_state = (zeros.clone(), zeros.clone()) + else: + actor_state = zeros + set_actor_hidden_state(policy, actor_state) + return actor_state + + +def state_dict_from_actor_hidden(actor_hidden): + """Convert the actor hidden state into the named tensor mapping expected by LEAPP state APIs.""" + if actor_hidden is None: + return {} + if isinstance(actor_hidden, tuple): + return {f"actor_state_{idx}": tensor for idx, tensor in enumerate(actor_hidden)} + return {"actor_state": actor_hidden} + + +def actor_hidden_from_registered(registered_state, original_hidden): + """Restore the registered LEAPP state to the hidden-state structure expected by the actor memory module.""" + if isinstance(original_hidden, tuple): + if isinstance(registered_state, tuple): + return registered_state + return (registered_state,) + return registered_state + + +def export_rsl_rl_agent( + args_cli: argparse.Namespace, + env_cfg, + agent_cfg, + simulation_app=None, +) -> bool: + """Export a RSL-RL agent.""" + _load_runtime_dependencies() + + task_name = args_cli.task.split(":")[-1] + checkpoint_task_name = task_name.replace("-Play", "") + + agent_cfg = cli_args.update_rsl_rl_cfg(agent_cfg, args_cli) + env_cfg.scene.num_envs = 1 + + agent_cfg = handle_deprecated_rsl_rl_cfg(agent_cfg, installed_version) + + # note: certain randomizations occur in the environment initialization so we set the seed here + env_cfg.seed = agent_cfg.seed + env_cfg.sim.device = args_cli.device if args_cli.device is not None else env_cfg.sim.device + + log_root_path = os.path.join("logs", "rsl_rl", agent_cfg.experiment_name) + log_root_path = os.path.abspath(log_root_path) + print(f"[INFO] Loading checkpoint search path from directory: {log_root_path}") + if args_cli.use_pretrained_checkpoint: + resume_path = get_published_pretrained_checkpoint("rsl_rl", checkpoint_task_name) + if not resume_path: + print("[INFO] Unfortunately a pre-trained checkpoint is currently unavailable for this task.") + return False + elif args_cli.checkpoint: + resume_path = retrieve_file_path(args_cli.checkpoint) + else: + resume_path = get_checkpoint_path(log_root_path, agent_cfg.load_run, agent_cfg.load_checkpoint) + + if not resume_path: + print(f"[INFO] No checkpoint found for task: {checkpoint_task_name} in directory: {log_root_path}") + return False + + log_dir = os.path.dirname(resume_path) + + env_cfg.log_dir = log_dir + + env = None + leapp_started = False + + try: + env = gym.make(args_cli.task, cfg=env_cfg, render_mode=None) + policy_node_name = ensure_env_spec_id(env) + + graph_name = args_cli.export_task_name if args_cli.export_task_name is not None else task_name + + if isinstance(env.unwrapped, ManagerBasedRLEnv): + # Patch only the observation groups consumed by the actor policy. + # This filters out the critic and teacher observation groups. + obs_groups_cfg = getattr(agent_cfg, "obs_groups", None) + if isinstance(obs_groups_cfg, Mapping): + required_obs_groups = set(obs_groups_cfg.get("actor", ["policy"])) + else: + required_obs_groups = {"policy"} + patch_env_for_export( + env, + export_method=args_cli.export_method, + required_obs_groups=required_obs_groups, + ) + + env = RslRlVecEnvWrapper(env, clip_actions=agent_cfg.clip_actions) + + print(f"[INFO]: Loading model checkpoint from: {resume_path}") + if agent_cfg.class_name == "OnPolicyRunner": + runner = OnPolicyRunner(env, agent_cfg.to_dict(), log_dir=None, device=agent_cfg.device) + elif agent_cfg.class_name == "DistillationRunner": + runner = DistillationRunner(env, agent_cfg.to_dict(), log_dir=None, device=agent_cfg.device) + else: + raise ValueError(f"Unsupported runner class: {agent_cfg.class_name}") + runner.load(resume_path) + + policy = runner.get_inference_policy(device=env.unwrapped.device) + + if args_cli.export_save_path is not None: + save_path = args_cli.export_save_path + elif args_cli.use_pretrained_checkpoint: + # Use a predictable path independent of the Nucleus mirror directory structure. + save_path = os.path.join(".pretrained_checkpoints", "rsl_rl", checkpoint_task_name) + else: + save_path = log_dir + leapp.start(graph_name, save_path=save_path, max_cached_io=max(args_cli.validation_steps, 2)) + leapp_started = True + obs = env.reset()[0] + if simulation_app is not None: + while not simulation_app.is_running(): + time.sleep(0.5) + + for _ in range(max(args_cli.validation_steps, 2)): + with torch.inference_mode(): + if is_actor_recurrent_policy(policy): + actor_hidden = ensure_actor_hidden_state_initialized( + policy, + batch_size=env.num_envs, + device=env.unwrapped.device, + dtype=next(policy.parameters()).dtype, + ) + registered_state = annotate.state_tensors( + policy_node_name, + state_dict_from_actor_hidden(actor_hidden), + ) + set_actor_hidden_state(policy, actor_hidden_from_registered(registered_state, actor_hidden)) + + actions = policy(obs) + + if is_actor_recurrent_policy(policy): + actor_hidden_after = get_actor_hidden_state(policy) + annotate.update_state( + policy_node_name, + state_dict_from_actor_hidden(actor_hidden_after), + ) + + obs, _, _, _ = env.step(actions) + + leapp.stop() + leapp_started = False + validate = args_cli.validation_steps > 0 + leapp.compile_graph(visualize=not args_cli.disable_graph_visualization, validate=validate) + finally: + if leapp_started: + with contextlib.suppress(Exception): + leapp.stop() + if env is not None: + env.close() + + return True + + +def run_export_with_hydra(args_cli: argparse.Namespace, hydra_args: list[str]) -> bool: + """Resolve Hydra task configuration and export one RSL-RL policy.""" + try: + from isaaclab.app import launch_simulation + except ImportError: + from isaaclab_tasks.utils.sim_launcher import launch_simulation + + from isaaclab_tasks.utils.hydra import hydra_task_config + + original_argv = sys.argv + # Hydra reads the preset tokens (physics=/renderer=/presets=) from sys.argv directly. + sys.argv = [sys.argv[0]] + hydra_args + exported = False + + try: + + @hydra_task_config(args_cli.task, args_cli.agent) + def _main(env_cfg, agent_cfg) -> None: + nonlocal exported + with launch_simulation(env_cfg, args_cli): + exported = export_rsl_rl_agent(args_cli, env_cfg, agent_cfg) + + _main() + finally: + sys.argv = original_argv + + return exported + + +def main_cli(argv: list[str] | None = None) -> bool: + """Run the command-line export flow.""" + args_cli, hydra_args = parse_export_args(argv) + return run_export_with_hydra(args_cli, hydra_args) + + +if __name__ == "__main__": + main_cli() diff --git a/source/isaaclab/isaaclab/assets/articulation/base_articulation.py b/source/isaaclab/isaaclab/assets/articulation/base_articulation.py index 25ca2c4ceaf0..76e831bdad17 100644 --- a/source/isaaclab/isaaclab/assets/articulation/base_articulation.py +++ b/source/isaaclab/isaaclab/assets/articulation/base_articulation.py @@ -16,6 +16,8 @@ import torch import warp as wp +from ...sim import SimulationContext +from ...utils.leapp.leapp_semantics import OutputKindEnum, joint_names_resolver, leapp_tensor_semantics from ..asset_base import AssetBase if TYPE_CHECKING: @@ -94,6 +96,8 @@ def __init__(self, cfg: ArticulationCfg): cfg: A configuration instance. """ super().__init__(cfg) + sim_ctx = SimulationContext.instance() + self._sim_cfg = sim_ctx.cfg if sim_ctx is not None else None """ Properties @@ -173,6 +177,19 @@ def root_view(self): """ raise NotImplementedError() + @property + def num_base_dofs(self) -> int: + """Number of free DoFs of the floating base. + + A floating-base articulation can translate and rotate freely in space, so + its base contributes 6 DoFs (3 linear, 3 angular). A fixed-base articulation + is bolted to the world and contributes 0. + + Use this to map an actuated-joint index ``j`` to its column in the Jacobian + / mass matrix / gravity vector: ``column = j + num_base_dofs``. + """ + return 0 if self.is_fixed_base else 6 + @property @abstractmethod def instantaneous_wrench_composer(self) -> WrenchComposer: @@ -1066,11 +1083,12 @@ def write_joint_friction_coefficient_to_sim_index( joint_ids: Sequence[int] | torch.Tensor | wp.array | None = None, env_ids: Sequence[int] | torch.Tensor | wp.array | None = None, ) -> None: - r"""Write joint static friction coefficients into the simulation. + r"""Write backend-specific joint friction values into the simulation. - The joint static friction is a unitless quantity. It relates the magnitude of the spatial force transmitted - from the parent body to the child body to the maximal static friction force that may be applied by the solver - to resist the joint motion. + .. warning:: + The physical meaning and units of joint friction depend on the concrete backend and solver. Do not assume + values are comparable across backends; check the backend-specific implementation before interpreting or + reusing them. .. note:: This method expects partial data. @@ -1080,7 +1098,7 @@ def write_joint_friction_coefficient_to_sim_index( Some backends may provide optimized implementations for masks / indices. Args: - joint_friction_coeff: Joint static friction coefficient. Shape is (len(env_ids), len(joint_ids)). + joint_friction_coeff: Backend-specific joint friction values. Shape is (len(env_ids), len(joint_ids)). joint_ids: The joint indices to set the joint torque limits for. Defaults to None (all joints). env_ids: The environment indices to set the joint torque limits for. Defaults to None (all instances). """ @@ -1094,11 +1112,12 @@ def write_joint_friction_coefficient_to_sim_mask( joint_mask: wp.array | None = None, env_mask: wp.array | None = None, ) -> None: - r"""Write joint static friction coefficients into the simulation. + r"""Write backend-specific joint friction values into the simulation. - The joint static friction is a unitless quantity. It relates the magnitude of the spatial force transmitted - from the parent body to the child body to the maximal static friction force that may be applied by the solver - to resist the joint motion. + .. warning:: + The physical meaning and units of joint friction depend on the concrete backend and solver. Do not assume + values are comparable across backends; check the backend-specific implementation before interpreting or + reusing them. .. note:: This method expects full data. @@ -1108,7 +1127,7 @@ def write_joint_friction_coefficient_to_sim_mask( Some backends may provide optimized implementations for masks / indices. Args: - joint_friction_coeff: Joint static friction coefficient. Shape is (num_instances, num_joints). + joint_friction_coeff: Backend-specific joint friction values. Shape is (num_instances, num_joints). joint_mask: Joint mask. If None, then all the joints are updated. Shape is (num_joints,). env_mask: Environment mask. If None, then all the instances are updated. Shape is (num_instances,). """ @@ -1266,6 +1285,7 @@ def set_inertias_mask( raise NotImplementedError() @abstractmethod + @leapp_tensor_semantics(kind=OutputKindEnum.JOINT_POSITION, element_names_resolver=joint_names_resolver) def set_joint_position_target_index( self, *, @@ -1293,6 +1313,7 @@ def set_joint_position_target_index( raise NotImplementedError() @abstractmethod + @leapp_tensor_semantics(kind=OutputKindEnum.JOINT_POSITION, element_names_resolver=joint_names_resolver) def set_joint_position_target_mask( self, *, @@ -1320,6 +1341,7 @@ def set_joint_position_target_mask( raise NotImplementedError() @abstractmethod + @leapp_tensor_semantics(kind=OutputKindEnum.JOINT_VELOCITY, element_names_resolver=joint_names_resolver) def set_joint_velocity_target_index( self, *, @@ -1347,6 +1369,7 @@ def set_joint_velocity_target_index( raise NotImplementedError() @abstractmethod + @leapp_tensor_semantics(kind=OutputKindEnum.JOINT_VELOCITY, element_names_resolver=joint_names_resolver) def set_joint_velocity_target_mask( self, *, @@ -1374,6 +1397,7 @@ def set_joint_velocity_target_mask( raise NotImplementedError() @abstractmethod + @leapp_tensor_semantics(kind=OutputKindEnum.JOINT_EFFORT, element_names_resolver=joint_names_resolver) def set_joint_effort_target_index( self, *, @@ -1401,6 +1425,7 @@ def set_joint_effort_target_index( raise NotImplementedError() @abstractmethod + @leapp_tensor_semantics(kind=OutputKindEnum.JOINT_EFFORT, element_names_resolver=joint_names_resolver) def set_joint_effort_target_mask( self, *, @@ -2553,14 +2578,17 @@ def set_external_force_and_torque( env_ids: Sequence[int] | torch.Tensor | wp.array | None = None, is_global: bool = False, ) -> None: - """Deprecated, same as :meth:`permanent_wrench_composer.set_forces_and_torques`.""" + """Deprecated. Resets target environments, then adds forces and torques via the permanent wrench composer.""" warnings.warn( - "The function 'set_external_force_and_torque' will be deprecated in a future release. Please" - " use 'permanent_wrench_composer.set_forces_and_torques' instead.", + "The function 'set_external_force_and_torque' is deprecated. Please use" + " 'permanent_wrench_composer.reset' followed by 'permanent_wrench_composer.add_forces_and_torques'" + " instead.", DeprecationWarning, stacklevel=2, ) - self.permanent_wrench_composer.set_forces_and_torques( + # Reset only target env_ids then add (not set which clears all envs globally) + self.permanent_wrench_composer.reset(env_ids=env_ids) + self.permanent_wrench_composer.add_forces_and_torques( forces, torques, positions=positions, body_ids=body_ids, env_ids=env_ids, is_global=is_global ) diff --git a/source/isaaclab/isaaclab/assets/articulation/base_articulation_data.py b/source/isaaclab/isaaclab/assets/articulation/base_articulation_data.py index 61ee8f28e7a9..e73264bf509a 100644 --- a/source/isaaclab/isaaclab/assets/articulation/base_articulation_data.py +++ b/source/isaaclab/isaaclab/assets/articulation/base_articulation_data.py @@ -8,6 +8,14 @@ import warp as wp +from isaaclab.utils.leapp.leapp_semantics import ( + InputKindEnum, + QUAT_XYZW_ELEMENT_NAMES, + XYZ_ELEMENT_NAMES, + joint_names_resolver, + leapp_tensor_semantics, +) + class BaseArticulationData(ABC): """Data container for an articulation. @@ -624,6 +632,7 @@ def body_incoming_joint_wrench_b(self) -> wp.array: @property @abstractmethod + @leapp_tensor_semantics(kind=InputKindEnum.JOINT_POSITION, element_names_resolver=joint_names_resolver) def joint_pos(self) -> wp.array: """Joint positions of all joints. @@ -994,11 +1003,13 @@ def root_pose_w(self) -> wp.array: return self.root_link_pose_w @property + @leapp_tensor_semantics(kind=InputKindEnum.BODY_POSITION, element_names=XYZ_ELEMENT_NAMES) def root_pos_w(self) -> wp.array: """Shorthand for :attr:`root_link_pos_w`.""" return self.root_link_pos_w @property + @leapp_tensor_semantics(kind=InputKindEnum.BODY_ROTATION, element_names=QUAT_XYZW_ELEMENT_NAMES) def root_quat_w(self) -> wp.array: """Shorthand for :attr:`root_link_quat_w`.""" return self.root_link_quat_w diff --git a/source/isaaclab/isaaclab/assets/rigid_object/base_rigid_object_data.py b/source/isaaclab/isaaclab/assets/rigid_object/base_rigid_object_data.py index e04a1ffa8bc1..9cc8bb82b8eb 100644 --- a/source/isaaclab/isaaclab/assets/rigid_object/base_rigid_object_data.py +++ b/source/isaaclab/isaaclab/assets/rigid_object/base_rigid_object_data.py @@ -8,6 +8,13 @@ import warp as wp +from isaaclab.utils.leapp.leapp_semantics import ( + InputKindEnum, + QUAT_XYZW_ELEMENT_NAMES, + XYZ_ELEMENT_NAMES, + leapp_tensor_semantics, +) + class BaseRigidObjectData(ABC): """Data container for a rigid object. @@ -588,11 +595,13 @@ def root_pose_w(self) -> wp.array: return self.root_link_pose_w @property + @leapp_tensor_semantics(kind=InputKindEnum.BODY_POSITION, element_names=XYZ_ELEMENT_NAMES) def root_pos_w(self) -> wp.array: """Shorthand for :attr:`root_link_pos_w`.""" return self.root_link_pos_w @property + @leapp_tensor_semantics(kind=InputKindEnum.BODY_ROTATION, element_names=QUAT_XYZW_ELEMENT_NAMES) def root_quat_w(self) -> wp.array: """Shorthand for :attr:`root_link_quat_w`.""" return self.root_link_quat_w diff --git a/source/isaaclab/isaaclab/sensors/camera/camera_data.py b/source/isaaclab/isaaclab/sensors/camera/camera_data.py index 2cbe5006fecd..aa31ebb69e04 100644 --- a/source/isaaclab/isaaclab/sensors/camera/camera_data.py +++ b/source/isaaclab/isaaclab/sensors/camera/camera_data.py @@ -3,90 +3,237 @@ # # SPDX-License-Identifier: BSD-3-Clause -from dataclasses import dataclass +from __future__ import annotations + from typing import Any -import torch +import warp as wp + +# Re-exported as part of the public isaaclab.sensors.camera API +from isaaclab.renderers.output_contract import RenderBufferKind, RenderBufferSpec +from isaaclab.utils.warp import ProxyArray +from isaaclab.utils.warp.warp_math import convert_camera_frame_orientation_convention_wp -from isaaclab.utils.math import convert_camera_frame_orientation_convention +__all__ = ["CameraData", "RenderBufferKind", "RenderBufferSpec"] -@dataclass class CameraData: - """Data container for the camera sensor.""" + """Data container for the camera sensor. + + Public properties return :class:`~isaaclab.utils.warp.ProxyArray` wrappers. + Use ``.torch`` for a cached zero-copy :class:`torch.Tensor` view or + ``.warp`` for the underlying :class:`warp.array`. + """ + + def __init__(self): + # ProxyArray wrappers — created in create_buffers() + self._pos_w: ProxyArray | None = None + self._quat_w_world: ProxyArray | None = None + self._intrinsic_matrices: ProxyArray | None = None + self._quat_w_ros: ProxyArray | None = None + self._quat_w_opengl: ProxyArray | None = None + + # Output image buffers — allocated in allocate() + self._output: dict[str, ProxyArray] | None = None + + self.image_shape: tuple[int, int] | None = None + """A tuple containing (height, width) of the camera sensor.""" + + self.info: dict[str, Any] | None = None + """The retrieved sensor info with sensor types as key. + + This contains extra information provided by the sensor such as semantic segmentation label mapping, prim paths. + For semantic-based data, this corresponds to the ``"info"`` key in the output of the sensor. For other sensor + types, the info is empty. + """ ## # Frame state. ## - pos_w: torch.Tensor = None - """Position of the sensor origin in world frame, following ROS convention. + @property + def pos_w(self) -> ProxyArray: + """Position of the sensor origin in world frame [m], following ROS convention. - Shape is (N, 3) where N is the number of sensors. - """ + Shape is (N,), dtype ``wp.vec3f``. In torch this resolves to (N, 3), + where N is the number of sensors. Use ``.warp`` for the underlying + ``wp.array`` or ``.torch`` for a cached zero-copy ``torch.Tensor`` view. + """ + return self._pos_w - quat_w_world: torch.Tensor = None - """Quaternion orientation `(x, y, z, w)` of the sensor origin in world frame, following the world coordinate frame + @property + def quat_w_world(self) -> ProxyArray: + """Quaternion orientation ``(x, y, z, w)`` of the sensor origin in world frame, + following the world coordinate frame convention. - .. note:: - World frame convention follows the camera aligned with forward axis +X and up axis +Z. + .. note:: + World frame convention follows the camera aligned with forward axis +X and up axis +Z. - Shape is (N, 4) where N is the number of sensors. - """ + Shape is (N,), dtype ``wp.quatf``. In torch this resolves to (N, 4), + where N is the number of sensors. Use ``.warp`` for the underlying + ``wp.array`` or ``.torch`` for a cached zero-copy ``torch.Tensor`` view. + """ + return self._quat_w_world ## # Camera data ## - image_shape: tuple[int, int] = None - """A tuple containing (height, width) of the camera sensor.""" + @property + def intrinsic_matrices(self) -> ProxyArray: + """The intrinsic matrices for the camera. - intrinsic_matrices: torch.Tensor = None - """The intrinsic matrices for the camera. + Shape is (N,), dtype ``wp.mat33f``. In torch this resolves to (N, 3, 3), + where N is the number of sensors. Use ``.warp`` for the underlying + ``wp.array`` or ``.torch`` for a cached zero-copy ``torch.Tensor`` view. + """ + return self._intrinsic_matrices - Shape is (N, 3, 3) where N is the number of sensors. - """ + @property + def output(self) -> dict[str, ProxyArray] | None: + """The retrieved sensor data with sensor types as key. - output: dict[str, torch.Tensor] = None - """The retrieved sensor data with sensor types as key. + Each value is a :class:`~isaaclab.utils.warp.ProxyArray` of shape + ``(N, H, W, C)`` where N is the number of views, H/W are image dimensions, + and C is the number of channels. Use ``.torch`` for a ``torch.Tensor`` view + or ``.warp`` for the underlying ``wp.array``. - The format of the data is available in the `Replicator Documentation`_. For semantic-based data, - this corresponds to the ``"data"`` key in the output of the sensor. + The format of the data is available in the `Replicator Documentation`_. For semantic-based data, + this corresponds to the ``"data"`` key in the output of the sensor. - .. _Replicator Documentation: https://docs.omniverse.nvidia.com/prod_extensions/prod_extensions/ext_replicator/annotators_details.html#annotator-output - """ + .. _Replicator Documentation: https://docs.omniverse.nvidia.com/prod_extensions/prod_extensions/ext_replicator/annotators_details.html#annotator-output + """ + return self._output - info: list[dict[str, Any]] = None - """The retrieved sensor info with sensor types as key. + def create_buffers(self, num_views: int, device: str) -> None: + """Allocate warp arrays for pose and intrinsics and create their :class:`ProxyArray` wrappers. - This contains extra information provided by the sensor such as semantic segmentation label mapping, prim paths. - For semantic-based data, this corresponds to the ``"info"`` key in the output of the sensor. For other sensor - types, the info is empty. - """ + Called by :class:`~isaaclab.sensors.camera.Camera` after :meth:`allocate` to + populate the pose and intrinsics buffers. + + Args: + num_views: Number of camera views (batch dimension). + device: Device for tensor storage (e.g. ``"cuda:0"``). + """ + self._pos_w = ProxyArray(wp.zeros(num_views, dtype=wp.vec3f, device=device)) + self._quat_w_world = ProxyArray(wp.zeros(num_views, dtype=wp.quatf, device=device)) + self._intrinsic_matrices = ProxyArray(wp.zeros(num_views, dtype=wp.mat33f, device=device)) + self._quat_w_ros = ProxyArray(wp.zeros(num_views, dtype=wp.quatf, device=device)) + self._quat_w_opengl = ProxyArray(wp.zeros(num_views, dtype=wp.quatf, device=device)) + + @classmethod + def allocate( + cls, + data_types: list[str], + height: int, + width: int, + num_views: int, + device: str, + supported_specs: dict[RenderBufferKind, RenderBufferSpec], + ) -> CameraData: + """Build a :class:`CameraData` with output buffers pre-allocated as warp arrays. + + Allocates one ``(num_views, height, width, channels)`` warp array per kind + in the intersection of ``data_types`` and ``supported_specs``, using + the channels and dtype from each :class:`RenderBufferSpec`. Each buffer is + wrapped in a :class:`~isaaclab.utils.warp.ProxyArray`; call ``.torch`` on + the result to obtain a zero-copy :class:`torch.Tensor` view. + + Args: + data_types: Requested output names (typically :attr:`CameraCfg.data_types`). + Every name must be a member of :class:`RenderBufferKind`. + height: Image height in pixels. + width: Image width in pixels. + num_views: Number of camera views (batch dimension). + device: Device on which to allocate the buffers. + supported_specs: Per-buffer layout the active renderer can produce, + keyed by :class:`RenderBufferKind`. Names absent from this mapping + are not allocated. + + Returns: + A new :class:`CameraData` with :attr:`image_shape`, :attr:`output`, + and :attr:`info` populated; pose/intrinsic buffers must be created + separately via :meth:`create_buffers`. + + Raises: + ValueError: If ``data_types`` contains names that are not members of + :class:`RenderBufferKind`. + """ + valid_names = {kind.value for kind in RenderBufferKind} + unknown = [name for name in data_types if name not in valid_names] + if unknown: + raise ValueError(f"Unknown RenderBufferKind name(s): {unknown}. Expected members of RenderBufferKind.") + requested = {RenderBufferKind(name) for name in data_types} + + rgb_kinds = {RenderBufferKind.RGB, RenderBufferKind.RGBA} + rgb_alias = rgb_kinds <= supported_specs.keys() and not requested.isdisjoint(rgb_kinds) + if rgb_alias: + requested.update(rgb_kinds) + + allocated = requested.intersection(supported_specs) + if rgb_alias: + allocated.remove(RenderBufferKind.RGB) + + buffers: dict[str, ProxyArray] = {} + for name, spec in supported_specs.items(): + if name not in allocated: + continue + shape = (num_views, height, width, spec.channels) + buffers[str(name)] = ProxyArray(wp.zeros(shape, dtype=spec.dtype, device=device)) + + if rgb_alias: + # Zero-copy strided view into rgba: shape (N, H, W, 3), skipping the alpha channel. + # Byte strides for a contiguous (N, H, W, 4) uint8 array are (H*W*4, W*4, 4, 1). + # Using the same outer strides but limiting the last dim to 3 channels gives a + # non-contiguous view where each pixel reads RGB without the alpha byte. + rgba_wp = buffers[str(RenderBufferKind.RGBA)].warp + rgb_wp = wp.array( + ptr=rgba_wp.ptr, + shape=(num_views, height, width, 3), + strides=(height * width * 4, width * 4, 4, 1), + dtype=wp.uint8, + device=rgba_wp.device, + copy=False, + ) + buffers[str(RenderBufferKind.RGB)] = ProxyArray(rgb_wp) + + obj = cls() + obj.image_shape = (height, width) + obj._output = buffers + obj.info = {name: None for name in buffers} + return obj ## # Additional Frame orientation conventions ## @property - def quat_w_ros(self) -> torch.Tensor: - """Quaternion orientation `(x, y, z, w)` of the sensor origin in the world frame, following ROS convention. + def quat_w_ros(self) -> ProxyArray: + """Quaternion orientation ``(x, y, z, w)`` of the sensor origin in the world frame, following ROS convention. .. note:: ROS convention follows the camera aligned with forward axis +Z and up axis -Y. - Shape is (N, 4) where N is the number of sensors. + Shape is (N,), dtype ``wp.quatf``. In torch this resolves to (N, 4), + where N is the number of sensors. Use ``.warp`` for the underlying + ``wp.array`` or ``.torch`` for a cached zero-copy ``torch.Tensor`` view. """ - return convert_camera_frame_orientation_convention(self.quat_w_world, origin="world", target="ros") + convert_camera_frame_orientation_convention_wp(self._quat_w_world.warp, self._quat_w_ros.warp, "world", "ros") + return self._quat_w_ros @property - def quat_w_opengl(self) -> torch.Tensor: - """Quaternion orientation `(x, y, z, w)` of the sensor origin in the world frame, following + def quat_w_opengl(self) -> ProxyArray: + """Quaternion orientation ``(x, y, z, w)`` of the sensor origin in the world frame, following Opengl / USD Camera convention. .. note:: OpenGL convention follows the camera aligned with forward axis -Z and up axis +Y. - Shape is (N, 4) where N is the number of sensors. + Shape is (N,), dtype ``wp.quatf``. In torch this resolves to (N, 4), + where N is the number of sensors. Use ``.warp`` for the underlying + ``wp.array`` or ``.torch`` for a cached zero-copy ``torch.Tensor`` view. """ - return convert_camera_frame_orientation_convention(self.quat_w_world, origin="world", target="opengl") + convert_camera_frame_orientation_convention_wp( + self._quat_w_world.warp, self._quat_w_opengl.warp, "world", "opengl" + ) + return self._quat_w_opengl diff --git a/source/isaaclab/isaaclab/sensors/contact_sensor/base_contact_sensor_data.py b/source/isaaclab/isaaclab/sensors/contact_sensor/base_contact_sensor_data.py index 74acf5092761..2f5f69ef55db 100644 --- a/source/isaaclab/isaaclab/sensors/contact_sensor/base_contact_sensor_data.py +++ b/source/isaaclab/isaaclab/sensors/contact_sensor/base_contact_sensor_data.py @@ -9,7 +9,14 @@ from abc import ABC, abstractmethod -import warp as wp +from isaaclab.utils.leapp import ( + POSE7_ELEMENT_NAMES, + QUAT_XYZW_ELEMENT_NAMES, + XYZ_ELEMENT_NAMES, + InputKindEnum, + leapp_tensor_semantics, +) +from isaaclab.utils.warp import ProxyArray class BaseContactSensorData(ABC): @@ -21,7 +28,8 @@ class BaseContactSensorData(ABC): @property @abstractmethod - def pose_w(self) -> wp.array | None: + @leapp_tensor_semantics(kind=InputKindEnum.BODY_POSE, element_names=POSE7_ELEMENT_NAMES) + def pose_w(self) -> ProxyArray | None: """Pose of the sensor origin in world frame. None if :attr:`ContactSensorCfg.track_pose` is False. @@ -30,7 +38,8 @@ def pose_w(self) -> wp.array | None: @property @abstractmethod - def pos_w(self) -> wp.array | None: + @leapp_tensor_semantics(kind=InputKindEnum.BODY_POSITION, element_names=XYZ_ELEMENT_NAMES) + def pos_w(self) -> ProxyArray | None: """Position of the sensor origin in world frame. Shape is (num_instances, num_sensors), dtype = wp.vec3f. In torch this resolves to @@ -42,7 +51,8 @@ def pos_w(self) -> wp.array | None: @property @abstractmethod - def quat_w(self) -> wp.array | None: + @leapp_tensor_semantics(kind=InputKindEnum.BODY_ROTATION, element_names=QUAT_XYZW_ELEMENT_NAMES) + def quat_w(self) -> ProxyArray | None: """Orientation of the sensor origin in world frame. Shape is (num_instances, num_sensors), dtype = wp.quatf. In torch this resolves to @@ -54,7 +64,8 @@ def quat_w(self) -> wp.array | None: @property @abstractmethod - def net_forces_w(self) -> wp.array | None: + @leapp_tensor_semantics(kind=InputKindEnum.VECTOR3D, element_names=XYZ_ELEMENT_NAMES) + def net_forces_w(self) -> ProxyArray | None: """The net normal contact forces in world frame. Shape is (num_instances, num_sensors), dtype = wp.vec3f. In torch this resolves to @@ -64,7 +75,8 @@ def net_forces_w(self) -> wp.array | None: @property @abstractmethod - def net_forces_w_history(self) -> wp.array | None: + @leapp_tensor_semantics(kind=InputKindEnum.VECTOR3D, element_names=XYZ_ELEMENT_NAMES) + def net_forces_w_history(self) -> ProxyArray | None: """History of net normal contact forces. Shape is (num_instances, history_length, num_sensors), dtype = wp.vec3f. In torch this resolves to @@ -74,7 +86,8 @@ def net_forces_w_history(self) -> wp.array | None: @property @abstractmethod - def force_matrix_w(self) -> wp.array | None: + @leapp_tensor_semantics(kind=InputKindEnum.VECTOR3D, element_names=XYZ_ELEMENT_NAMES) + def force_matrix_w(self) -> ProxyArray | None: """Normal contact forces filtered between sensor and filtered bodies. Shape is (num_instances, num_sensors, num_filter_shapes), dtype = wp.vec3f. In torch this resolves to @@ -86,7 +99,8 @@ def force_matrix_w(self) -> wp.array | None: @property @abstractmethod - def force_matrix_w_history(self) -> wp.array | None: + @leapp_tensor_semantics(kind=InputKindEnum.VECTOR3D, element_names=XYZ_ELEMENT_NAMES) + def force_matrix_w_history(self) -> ProxyArray | None: """History of filtered contact forces. Shape is (num_instances, history_length, num_sensors, num_filter_shapes), dtype = wp.vec3f. @@ -98,7 +112,8 @@ def force_matrix_w_history(self) -> wp.array | None: @property @abstractmethod - def contact_pos_w(self) -> wp.array | None: + @leapp_tensor_semantics(kind=InputKindEnum.BODY_POSITION, element_names=XYZ_ELEMENT_NAMES) + def contact_pos_w(self) -> ProxyArray | None: """Average position of contact points. Shape is (num_instances, num_sensors, num_filter_shapes), dtype = wp.vec3f. In torch this resolves to @@ -110,7 +125,8 @@ def contact_pos_w(self) -> wp.array | None: @property @abstractmethod - def friction_forces_w(self) -> wp.array | None: + @leapp_tensor_semantics(kind=InputKindEnum.VECTOR3D, element_names=XYZ_ELEMENT_NAMES) + def friction_forces_w(self) -> ProxyArray | None: """Sum of friction forces. Shape is (num_instances, num_sensors, num_filter_shapes), dtype = wp.vec3f. In torch this resolves to @@ -122,7 +138,8 @@ def friction_forces_w(self) -> wp.array | None: @property @abstractmethod - def last_air_time(self) -> wp.array | None: + @leapp_tensor_semantics() + def last_air_time(self) -> ProxyArray | None: """Time spent in air before last contact. Shape is (num_instances, num_sensors), dtype = wp.float32. @@ -133,7 +150,8 @@ def last_air_time(self) -> wp.array | None: @property @abstractmethod - def current_air_time(self) -> wp.array | None: + @leapp_tensor_semantics() + def current_air_time(self) -> ProxyArray | None: """Time spent in air since last detach. Shape is (num_instances, num_sensors), dtype = wp.float32. @@ -144,7 +162,8 @@ def current_air_time(self) -> wp.array | None: @property @abstractmethod - def last_contact_time(self) -> wp.array | None: + @leapp_tensor_semantics() + def last_contact_time(self) -> ProxyArray | None: """Time spent in contact before last detach. Shape is (num_instances, num_sensors), dtype = wp.float32. @@ -155,7 +174,8 @@ def last_contact_time(self) -> wp.array | None: @property @abstractmethod - def current_contact_time(self) -> wp.array | None: + @leapp_tensor_semantics() + def current_contact_time(self) -> ProxyArray | None: """Time spent in contact since last contact. Shape is (num_instances, num_sensors), dtype = wp.float32. diff --git a/source/isaaclab/isaaclab/sensors/frame_transformer/base_frame_transformer_data.py b/source/isaaclab/isaaclab/sensors/frame_transformer/base_frame_transformer_data.py index 3b28a7b17d00..286af6e84ea3 100644 --- a/source/isaaclab/isaaclab/sensors/frame_transformer/base_frame_transformer_data.py +++ b/source/isaaclab/isaaclab/sensors/frame_transformer/base_frame_transformer_data.py @@ -9,7 +9,17 @@ from abc import ABC, abstractmethod -import warp as wp +from isaaclab.utils.leapp import ( + POSE7_ELEMENT_NAMES, + QUAT_XYZW_ELEMENT_NAMES, + XYZ_ELEMENT_NAMES, + InputKindEnum, + leapp_tensor_semantics, + target_frame_pose_resolver, + target_frame_quat_resolver, + target_frame_xyz_resolver, +) +from isaaclab.utils.warp import ProxyArray class BaseFrameTransformerData(ABC): @@ -30,7 +40,8 @@ def target_frame_names(self) -> list[str]: @property @abstractmethod - def target_pose_source(self) -> wp.array | None: + @leapp_tensor_semantics(kind=InputKindEnum.BODY_POSE, element_names_resolver=target_frame_pose_resolver) + def target_pose_source(self) -> ProxyArray | None: """Pose of the target frame(s) relative to source frame. Shape is (num_instances, num_target_frames), dtype = wp.transformf. In torch this resolves to @@ -40,7 +51,8 @@ def target_pose_source(self) -> wp.array | None: @property @abstractmethod - def target_pos_source(self) -> wp.array: + @leapp_tensor_semantics(kind=InputKindEnum.BODY_POSITION, element_names_resolver=target_frame_xyz_resolver) + def target_pos_source(self) -> ProxyArray: """Position of the target frame(s) relative to source frame. Shape is (num_instances, num_target_frames), dtype = wp.vec3f. In torch this resolves to @@ -50,7 +62,8 @@ def target_pos_source(self) -> wp.array: @property @abstractmethod - def target_quat_source(self) -> wp.array: + @leapp_tensor_semantics(kind=InputKindEnum.BODY_ROTATION, element_names_resolver=target_frame_quat_resolver) + def target_quat_source(self) -> ProxyArray: """Orientation of the target frame(s) relative to source frame. Shape is (num_instances, num_target_frames), dtype = wp.quatf. In torch this resolves to @@ -60,7 +73,8 @@ def target_quat_source(self) -> wp.array: @property @abstractmethod - def target_pose_w(self) -> wp.array | None: + @leapp_tensor_semantics(kind=InputKindEnum.BODY_POSE, element_names_resolver=target_frame_pose_resolver) + def target_pose_w(self) -> ProxyArray | None: """Pose of the target frame(s) after offset in world frame. Shape is (num_instances, num_target_frames), dtype = wp.transformf. In torch this resolves to @@ -70,7 +84,8 @@ def target_pose_w(self) -> wp.array | None: @property @abstractmethod - def target_pos_w(self) -> wp.array: + @leapp_tensor_semantics(kind=InputKindEnum.BODY_POSITION, element_names_resolver=target_frame_xyz_resolver) + def target_pos_w(self) -> ProxyArray: """Position of the target frame(s) after offset in world frame. Shape is (num_instances, num_target_frames), dtype = wp.vec3f. In torch this resolves to @@ -80,7 +95,8 @@ def target_pos_w(self) -> wp.array: @property @abstractmethod - def target_quat_w(self) -> wp.array: + @leapp_tensor_semantics(kind=InputKindEnum.BODY_ROTATION, element_names_resolver=target_frame_quat_resolver) + def target_quat_w(self) -> ProxyArray: """Orientation of the target frame(s) after offset in world frame. Shape is (num_instances, num_target_frames), dtype = wp.quatf. In torch this resolves to @@ -90,7 +106,8 @@ def target_quat_w(self) -> wp.array: @property @abstractmethod - def source_pose_w(self) -> wp.array | None: + @leapp_tensor_semantics(kind=InputKindEnum.BODY_POSE, element_names=POSE7_ELEMENT_NAMES) + def source_pose_w(self) -> ProxyArray | None: """Pose of the source frame after offset in world frame. Shape is (num_instances,), dtype = wp.transformf. In torch this resolves to (num_instances, 7). @@ -100,7 +117,8 @@ def source_pose_w(self) -> wp.array | None: @property @abstractmethod - def source_pos_w(self) -> wp.array: + @leapp_tensor_semantics(kind=InputKindEnum.BODY_POSITION, element_names=XYZ_ELEMENT_NAMES) + def source_pos_w(self) -> ProxyArray: """Position of the source frame after offset in world frame. Shape is (num_instances,), dtype = wp.vec3f. In torch this resolves to (num_instances, 3). @@ -109,7 +127,8 @@ def source_pos_w(self) -> wp.array: @property @abstractmethod - def source_quat_w(self) -> wp.array: + @leapp_tensor_semantics(kind=InputKindEnum.BODY_ROTATION, element_names=QUAT_XYZW_ELEMENT_NAMES) + def source_quat_w(self) -> ProxyArray: """Orientation of the source frame after offset in world frame. Shape is (num_instances,), dtype = wp.quatf. In torch this resolves to (num_instances, 4). diff --git a/source/isaaclab/isaaclab/sensors/frame_transformer/frame_transformer_data.py b/source/isaaclab/isaaclab/sensors/frame_transformer/frame_transformer_data.py index 8f202e15f40e..f6b28faea395 100644 --- a/source/isaaclab/isaaclab/sensors/frame_transformer/frame_transformer_data.py +++ b/source/isaaclab/isaaclab/sensors/frame_transformer/frame_transformer_data.py @@ -14,12 +14,17 @@ from .base_frame_transformer_data import BaseFrameTransformerData if TYPE_CHECKING: + from isaaclab_newton.sensors.frame_transformer.frame_transformer_data import ( + FrameTransformerData as NewtonFrameTransformerData, + ) from isaaclab_physx.sensors.frame_transformer import FrameTransformerData as PhysXFrameTransformerData class FrameTransformerData(FactoryBase, BaseFrameTransformerData): """Factory for creating frame transformer data instances.""" - def __new__(cls, *args, **kwargs) -> BaseFrameTransformerData | PhysXFrameTransformerData: + def __new__( + cls, *args, **kwargs + ) -> BaseFrameTransformerData | NewtonFrameTransformerData | PhysXFrameTransformerData: """Create a new instance of a frame transformer data based on the backend.""" return super().__new__(cls, *args, **kwargs) diff --git a/source/isaaclab/isaaclab/sensors/imu/base_imu_data.py b/source/isaaclab/isaaclab/sensors/imu/base_imu_data.py index a64eee79192f..039d5dd60f64 100644 --- a/source/isaaclab/isaaclab/sensors/imu/base_imu_data.py +++ b/source/isaaclab/isaaclab/sensors/imu/base_imu_data.py @@ -9,67 +9,29 @@ from abc import ABC, abstractmethod -import warp as wp +from isaaclab.utils.leapp import ( + XYZ_ELEMENT_NAMES, + InputKindEnum, + leapp_tensor_semantics, +) +from isaaclab.utils.warp import ProxyArray class BaseImuData(ABC): - """Data container for the Imu sensor. + """Data container for the IMU sensor. This base class defines the interface for IMU sensor data. Backend-specific implementations should inherit from this class and provide the actual data storage. - """ - - @property - @abstractmethod - def pose_w(self) -> wp.array | None: - """Pose of the sensor origin in world frame. - - Shape is (num_instances,), dtype = wp.transformf. In torch this resolves to (num_instances, 7). - The pose is provided in (x, y, z, qx, qy, qz, qw) format. - """ - raise NotImplementedError - - @property - @abstractmethod - def pos_w(self) -> wp.array: - """Position of the sensor origin in world frame. - - Shape is (num_instances,), dtype = wp.vec3f. In torch this resolves to (num_instances, 3). - """ - raise NotImplementedError - - @property - @abstractmethod - def quat_w(self) -> wp.array: - """Orientation of the sensor origin in world frame. - - Shape is (num_instances,), dtype = wp.quatf. In torch this resolves to (num_instances, 4). - The orientation is provided in (x, y, z, w) format. - """ - raise NotImplementedError - @property - @abstractmethod - def projected_gravity_b(self) -> wp.array: - """Gravity direction unit vector projected on the IMU frame. - - Shape is (num_instances,), dtype = wp.vec3f. In torch this resolves to (num_instances, 3). - """ - raise NotImplementedError - - @property - @abstractmethod - def lin_vel_b(self) -> wp.array: - """IMU frame linear velocity relative to the world expressed in IMU frame. - - Shape is (num_instances,), dtype = wp.vec3f. In torch this resolves to (num_instances, 3). - """ - raise NotImplementedError + Unlike the PVA sensor, the IMU only provides the two physical quantities that a + real inertial measurement unit measures: angular velocity and linear acceleration. + """ @property @abstractmethod - def ang_vel_b(self) -> wp.array: - """IMU frame angular velocity relative to the world expressed in IMU frame. + @leapp_tensor_semantics(kind=InputKindEnum.BODY_ANGULAR_VELOCITY, element_names=XYZ_ELEMENT_NAMES) + def ang_vel_b(self) -> ProxyArray: + """IMU frame angular velocity relative to the world expressed in IMU frame [rad/s]. Shape is (num_instances,), dtype = wp.vec3f. In torch this resolves to (num_instances, 3). """ @@ -77,17 +39,11 @@ def ang_vel_b(self) -> wp.array: @property @abstractmethod - def lin_acc_b(self) -> wp.array: - """IMU frame linear acceleration relative to the world expressed in IMU frame. - - Shape is (num_instances,), dtype = wp.vec3f. In torch this resolves to (num_instances, 3). - """ - raise NotImplementedError + @leapp_tensor_semantics(kind=InputKindEnum.BODY_LINEAR_ACCELERATION, element_names=XYZ_ELEMENT_NAMES) + def lin_acc_b(self) -> ProxyArray: + """Linear acceleration (proper) in the IMU frame [m/s^2]. - @property - @abstractmethod - def ang_acc_b(self) -> wp.array: - """IMU frame angular acceleration relative to the world expressed in IMU frame. + Zero in freefall, +g upward at rest. Shape is (num_instances,), dtype = wp.vec3f. In torch this resolves to (num_instances, 3). """ diff --git a/source/isaaclab/isaaclab/sensors/imu/imu_data.py b/source/isaaclab/isaaclab/sensors/imu/imu_data.py index 6f6ed268ad52..f23f2a3be6ca 100644 --- a/source/isaaclab/isaaclab/sensors/imu/imu_data.py +++ b/source/isaaclab/isaaclab/sensors/imu/imu_data.py @@ -3,7 +3,7 @@ # # SPDX-License-Identifier: BSD-3-Clause -"""Re-exports the base IMU data class for backwards compatibility.""" +"""Factory class for IMU data.""" from __future__ import annotations @@ -14,12 +14,13 @@ from .base_imu_data import BaseImuData if TYPE_CHECKING: + from isaaclab_newton.sensors.imu import ImuData as NewtonImuData from isaaclab_physx.sensors.imu import ImuData as PhysXImuData class ImuData(FactoryBase, BaseImuData): """Factory for creating IMU data instances.""" - def __new__(cls, *args, **kwargs) -> BaseImuData | PhysXImuData: - """Create a new instance of an IMU data based on the backend.""" + def __new__(cls, *args, **kwargs) -> BaseImuData | PhysXImuData | NewtonImuData: + """Create a new instance of IMU data based on the backend.""" return super().__new__(cls, *args, **kwargs) diff --git a/source/isaaclab/isaaclab/sensors/joint_wrench/base_joint_wrench_sensor_data.py b/source/isaaclab/isaaclab/sensors/joint_wrench/base_joint_wrench_sensor_data.py new file mode 100644 index 000000000000..282021a36784 --- /dev/null +++ b/source/isaaclab/isaaclab/sensors/joint_wrench/base_joint_wrench_sensor_data.py @@ -0,0 +1,42 @@ +# 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 + +"""Base class for joint-wrench sensor data containers.""" + +from __future__ import annotations + +from abc import ABC, abstractmethod + +from isaaclab.utils.warp import ProxyArray + + +class BaseJointWrenchSensorData(ABC): + """Data container for the joint reaction wrench sensor.""" + + @property + @abstractmethod + def force(self) -> ProxyArray | None: + """Linear component of the joint reaction wrench [N]. + + Expressed in the frame selected by + :attr:`~isaaclab.sensors.JointWrenchSensorCfg.convention`. Shape is + ``(num_envs, num_bodies)``, dtype ``wp.vec3f``. In torch this resolves + to ``(num_envs, num_bodies, 3)``. ``None`` before the simulation is + initialized. + """ + raise NotImplementedError + + @property + @abstractmethod + def torque(self) -> ProxyArray | None: + """Angular component of the joint reaction wrench [N·m]. + + Expressed in the frame selected by + :attr:`~isaaclab.sensors.JointWrenchSensorCfg.convention`. Shape is + ``(num_envs, num_bodies)``, dtype ``wp.vec3f``. In torch this resolves + to ``(num_envs, num_bodies, 3)``. ``None`` before the simulation is + initialized. + """ + raise NotImplementedError diff --git a/source/isaaclab/isaaclab/sensors/joint_wrench/joint_wrench_sensor_data.py b/source/isaaclab/isaaclab/sensors/joint_wrench/joint_wrench_sensor_data.py new file mode 100644 index 000000000000..5872640d143c --- /dev/null +++ b/source/isaaclab/isaaclab/sensors/joint_wrench/joint_wrench_sensor_data.py @@ -0,0 +1,28 @@ +# 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 + +"""Factory class for joint-wrench sensor data.""" + +from __future__ import annotations + +from typing import TYPE_CHECKING + +from isaaclab.utils.backend_utils import FactoryBase + +from .base_joint_wrench_sensor_data import BaseJointWrenchSensorData + +if TYPE_CHECKING: + from isaaclab_newton.sensors.joint_wrench import JointWrenchSensorData as NewtonJointWrenchSensorData + from isaaclab_physx.sensors.joint_wrench import JointWrenchSensorData as PhysXJointWrenchSensorData + + +class JointWrenchSensorData(FactoryBase, BaseJointWrenchSensorData): + """Factory for creating joint-wrench sensor data instances.""" + + def __new__( + cls, *args, **kwargs + ) -> BaseJointWrenchSensorData | PhysXJointWrenchSensorData | NewtonJointWrenchSensorData: + """Create a new instance of joint-wrench sensor data based on the backend.""" + return super().__new__(cls, *args, **kwargs) diff --git a/source/isaaclab/isaaclab/sensors/pva/base_pva_data.py b/source/isaaclab/isaaclab/sensors/pva/base_pva_data.py new file mode 100644 index 000000000000..07fcab62b0b6 --- /dev/null +++ b/source/isaaclab/isaaclab/sensors/pva/base_pva_data.py @@ -0,0 +1,111 @@ +# 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 + +"""Base class for PVA sensor data containers.""" + +from __future__ import annotations + +from abc import ABC, abstractmethod + +from isaaclab.utils.leapp import ( + POSE7_ELEMENT_NAMES, + QUAT_XYZW_ELEMENT_NAMES, + XYZ_ELEMENT_NAMES, + InputKindEnum, + leapp_tensor_semantics, +) +from isaaclab.utils.warp import ProxyArray + + +class BasePvaData(ABC): + """Data container for the PVA sensor. + + This base class defines the interface for PVA sensor data. Backend-specific + implementations should inherit from this class and provide the actual data storage. + """ + + @property + @abstractmethod + @leapp_tensor_semantics(kind=InputKindEnum.BODY_POSE, element_names=POSE7_ELEMENT_NAMES) + def pose_w(self) -> ProxyArray | None: + """Pose of the sensor origin in world frame [m, unitless]. + + Shape is (num_instances,), dtype = wp.transformf. In torch this resolves to (num_instances, 7). + The pose is provided in (x, y, z, qx, qy, qz, qw) format. + """ + raise NotImplementedError + + @property + @abstractmethod + @leapp_tensor_semantics(kind=InputKindEnum.BODY_POSITION, element_names=XYZ_ELEMENT_NAMES) + def pos_w(self) -> ProxyArray: + """Position of the sensor origin in world frame [m]. + + Shape is (num_instances,), dtype = wp.vec3f. In torch this resolves to (num_instances, 3). + """ + raise NotImplementedError + + @property + @abstractmethod + @leapp_tensor_semantics(kind=InputKindEnum.BODY_ROTATION, element_names=QUAT_XYZW_ELEMENT_NAMES) + def quat_w(self) -> ProxyArray: + """Orientation of the sensor origin in world frame. + + Shape is (num_instances,), dtype = wp.quatf. In torch this resolves to (num_instances, 4). + The orientation is provided in (x, y, z, w) format. + """ + raise NotImplementedError + + @property + @abstractmethod + @leapp_tensor_semantics(kind=InputKindEnum.VECTOR3D, element_names=XYZ_ELEMENT_NAMES) + def projected_gravity_b(self) -> ProxyArray: + """Gravity direction unit vector projected on the PVA frame. + + Shape is (num_instances,), dtype = wp.vec3f. In torch this resolves to (num_instances, 3). + """ + raise NotImplementedError + + @property + @abstractmethod + @leapp_tensor_semantics(kind=InputKindEnum.BODY_LINEAR_VELOCITY, element_names=XYZ_ELEMENT_NAMES) + def lin_vel_b(self) -> ProxyArray: + """PVA frame linear velocity relative to the world expressed in PVA frame [m/s]. + + Shape is (num_instances,), dtype = wp.vec3f. In torch this resolves to (num_instances, 3). + """ + raise NotImplementedError + + @property + @abstractmethod + @leapp_tensor_semantics(kind=InputKindEnum.BODY_ANGULAR_VELOCITY, element_names=XYZ_ELEMENT_NAMES) + def ang_vel_b(self) -> ProxyArray: + """PVA frame angular velocity relative to the world expressed in PVA frame [rad/s]. + + Shape is (num_instances,), dtype = wp.vec3f. In torch this resolves to (num_instances, 3). + """ + raise NotImplementedError + + @property + @abstractmethod + @leapp_tensor_semantics(kind=InputKindEnum.BODY_LINEAR_ACCELERATION, element_names=XYZ_ELEMENT_NAMES) + def lin_acc_b(self) -> ProxyArray: + """Linear acceleration (coordinate) in the PVA frame [m/s^2]. + + Equal to -g in freefall, zero at rest. + + Shape is (num_instances,), dtype = wp.vec3f. In torch this resolves to (num_instances, 3). + """ + raise NotImplementedError + + @property + @abstractmethod + @leapp_tensor_semantics(kind=InputKindEnum.BODY_ANGULAR_ACCELERATION, element_names=XYZ_ELEMENT_NAMES) + def ang_acc_b(self) -> ProxyArray: + """PVA frame angular acceleration relative to the world expressed in PVA frame [rad/s^2]. + + Shape is (num_instances,), dtype = wp.vec3f. In torch this resolves to (num_instances, 3). + """ + raise NotImplementedError diff --git a/source/isaaclab/isaaclab/sensors/pva/pva_data.py b/source/isaaclab/isaaclab/sensors/pva/pva_data.py new file mode 100644 index 000000000000..e1a614772a18 --- /dev/null +++ b/source/isaaclab/isaaclab/sensors/pva/pva_data.py @@ -0,0 +1,26 @@ +# 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 + +"""Re-exports the base PVA data class.""" + +from __future__ import annotations + +from typing import TYPE_CHECKING + +from isaaclab.utils.backend_utils import FactoryBase + +from .base_pva_data import BasePvaData + +if TYPE_CHECKING: + from isaaclab_newton.sensors.pva import PvaData as NewtonPvaData + from isaaclab_physx.sensors.pva import PvaData as PhysXPvaData + + +class PvaData(FactoryBase, BasePvaData): + """Factory for creating PVA data instances.""" + + def __new__(cls, *args, **kwargs) -> BasePvaData | NewtonPvaData | PhysXPvaData: + """Create a new instance of PVA data based on the backend.""" + return super().__new__(cls, *args, **kwargs) diff --git a/source/isaaclab/isaaclab/sensors/ray_caster/multi_mesh_ray_caster_camera_data.py b/source/isaaclab/isaaclab/sensors/ray_caster/multi_mesh_ray_caster_camera_data.py index d2f26abdbf47..687185c3e0bb 100644 --- a/source/isaaclab/isaaclab/sensors/ray_caster/multi_mesh_ray_caster_camera_data.py +++ b/source/isaaclab/isaaclab/sensors/ray_caster/multi_mesh_ray_caster_camera_data.py @@ -5,17 +5,20 @@ """Data container for the multi-mesh ray-cast camera sensor.""" -import torch - from isaaclab.sensors.camera import CameraData +from isaaclab.utils.warp import ProxyArray -from .ray_caster_data import RayCasterData +class MultiMeshRayCasterCameraData(CameraData): + """Data container for the multi-mesh ray-cast camera sensor. -class MultiMeshRayCasterCameraData(CameraData, RayCasterData): - """Data container for the multi-mesh ray-cast sensor.""" + This class extends :class:`CameraData` with additional mesh-id information. + It does not inherit from :class:`RayCasterData` because the camera variant + manages its own torch-based pose and hit buffers independently from the + warp-native :class:`RayCasterData`. + """ - image_mesh_ids: torch.Tensor = None + image_mesh_ids: ProxyArray = None """The mesh ids of the image pixels. Shape is (N, H, W, 1), where N is the number of sensors, H and W are the height and width of the image, diff --git a/source/isaaclab/isaaclab/sensors/ray_caster/multi_mesh_ray_caster_data.py b/source/isaaclab/isaaclab/sensors/ray_caster/multi_mesh_ray_caster_data.py index b9ae187591be..331dcb4af79f 100644 --- a/source/isaaclab/isaaclab/sensors/ray_caster/multi_mesh_ray_caster_data.py +++ b/source/isaaclab/isaaclab/sensors/ray_caster/multi_mesh_ray_caster_data.py @@ -6,7 +6,7 @@ """Data container for the multi-mesh ray-cast sensor.""" -import torch +from isaaclab.utils.warp import ProxyArray from .ray_caster_data import RayCasterData @@ -14,7 +14,7 @@ class MultiMeshRayCasterData(RayCasterData): """Data container for the multi-mesh ray-cast sensor.""" - ray_mesh_ids: torch.Tensor = None + ray_mesh_ids: ProxyArray = None """The mesh ids of the ray hits. Shape is (N, B, 1), where N is the number of sensors, B is the number of rays diff --git a/source/isaaclab/isaaclab/sensors/ray_caster/ray_caster_data.py b/source/isaaclab/isaaclab/sensors/ray_caster/ray_caster_data.py index 6103a2167d66..1265e1df0fc3 100644 --- a/source/isaaclab/isaaclab/sensors/ray_caster/ray_caster_data.py +++ b/source/isaaclab/isaaclab/sensors/ray_caster/ray_caster_data.py @@ -3,28 +3,82 @@ # # SPDX-License-Identifier: BSD-3-Clause -from dataclasses import dataclass +from __future__ import annotations -import torch +import warp as wp +from isaaclab.utils.leapp import ( + QUAT_XYZW_ELEMENT_NAMES, + XYZ_ELEMENT_NAMES, + leapp_tensor_semantics, +) +from isaaclab.utils.warp import ProxyArray -@dataclass -class RayCasterData: - """Data container for the ray-cast sensor.""" - pos_w: torch.Tensor = None - """Position of the sensor origin in world frame. +class RayCasterData: + """Data container for the ray-cast sensor. - Shape is (N, 3), where N is the number of sensors. + Public properties return :class:`~isaaclab.utils.warp.ProxyArray` wrappers. + Use ``.torch`` for a cached zero-copy :class:`torch.Tensor` view or + ``.warp`` for the underlying :class:`warp.array`. """ - quat_w: torch.Tensor = None - """Orientation of the sensor origin in quaternion (x, y, z, w) in world frame. - Shape is (N, 4), where N is the number of sensors. - """ - ray_hits_w: torch.Tensor = None - """The ray hit positions in the world frame. + def __init__(self): + self._pos_w: wp.array | None = None + self._quat_w: wp.array | None = None + self._ray_hits_w: wp.array | None = None - Shape is (N, B, 3), where N is the number of sensors, B is the number of rays - in the scan pattern per sensor. - """ + # _pos_w_ta / _quat_w_ta / _ray_hits_w_ta are created in create_buffers(). + # Accessing the public properties before create_buffers() raises AttributeError. + + @property + @leapp_tensor_semantics(kind="state/sensor/position", element_names=XYZ_ELEMENT_NAMES) + def pos_w(self) -> ProxyArray: + """Position of the sensor origin in world frame [m]. + + Shape is (N,), dtype ``wp.vec3f``. In torch this resolves to (N, 3), + where N is the number of sensors. Use ``.warp`` for the underlying + ``wp.array`` or ``.torch`` for a cached zero-copy ``torch.Tensor`` view. + """ + return self._pos_w_ta + + @property + @leapp_tensor_semantics(kind="state/sensor/rotation", element_names=QUAT_XYZW_ELEMENT_NAMES) + def quat_w(self) -> ProxyArray: + """Orientation of the sensor origin in quaternion (x, y, z, w) in world frame. + + Shape is (N,), dtype ``wp.quatf``. In torch this resolves to (N, 4), + where N is the number of sensors. Use ``.warp`` for the underlying + ``wp.array`` or ``.torch`` for a cached zero-copy ``torch.Tensor`` view. + """ + return self._quat_w_ta + + @property + @leapp_tensor_semantics(kind="state/sensor/ray_hit_position") + def ray_hits_w(self) -> ProxyArray: + """The ray hit positions in the world frame [m]. + + Shape is (N, B), dtype ``wp.vec3f``. In torch this resolves to (N, B, 3), + where N is the number of sensors and B is the number of rays per sensor. + Contains ``inf`` for missed hits. Use ``.warp`` for the underlying + ``wp.array`` or ``.torch`` for a cached zero-copy ``torch.Tensor`` view. + """ + return self._ray_hits_w_ta + + def create_buffers(self, num_envs: int, num_rays: int, device: str) -> None: + """Create internal warp buffers and their :class:`ProxyArray` wrappers. + + Args: + num_envs: Number of environments / sensors. + num_rays: Number of rays per sensor. + device: Device for tensor storage. + """ + self._device = device + + self._pos_w = wp.zeros(num_envs, dtype=wp.vec3f, device=device) + self._quat_w = wp.zeros(num_envs, dtype=wp.quatf, device=device) + self._ray_hits_w = wp.zeros((num_envs, num_rays), dtype=wp.vec3f, device=device) + + self._pos_w_ta = ProxyArray(self._pos_w) + self._quat_w_ta = ProxyArray(self._quat_w) + self._ray_hits_w_ta = ProxyArray(self._ray_hits_w) diff --git a/source/isaaclab/isaaclab/utils/leapp/__init__.py b/source/isaaclab/isaaclab/utils/leapp/__init__.py new file mode 100644 index 000000000000..f39b0e4d7eea --- /dev/null +++ b/source/isaaclab/isaaclab/utils/leapp/__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 + +"""Sub-module for LEAPP export annotation and proxy-based policy tracing.""" + +from isaaclab.utils.module import lazy_export + +lazy_export() diff --git a/source/isaaclab/isaaclab/utils/leapp/__init__.pyi b/source/isaaclab/isaaclab/utils/leapp/__init__.pyi new file mode 100644 index 000000000000..598395b99872 --- /dev/null +++ b/source/isaaclab/isaaclab/utils/leapp/__init__.pyi @@ -0,0 +1,64 @@ +# 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__ = [ + "ExportPatcher", + "InputKindEnum", + "LeappTensorSemantics", + "OutputKindEnum", + "POSE6_ELEMENT_NAMES", + "POSE7_ELEMENT_NAMES", + "QUAT_XYZW_ELEMENT_NAMES", + "WRENCH6_ELEMENT_NAMES", + "XYZ_ELEMENT_NAMES", + "body_names_resolver", + "body_pose6_resolver", + "body_pose_resolver", + "body_quat_resolver", + "body_wrench_resolver", + "body_xyz_resolver", + "build_command_connection", + "build_state_connection", + "build_write_connection", + "joint_names_resolver", + "leapp_input_tensor", + "leapp_real_env", + "leapp_tensor_semantics", + "patch_env_for_export", + "resolve_leapp_element_names", + "target_frame_pose_resolver", + "target_frame_quat_resolver", + "target_frame_xyz_resolver", +] + +from .export_annotator import ExportPatcher, patch_env_for_export +from .leapp_semantics import ( + InputKindEnum, + OutputKindEnum, + POSE6_ELEMENT_NAMES, + POSE7_ELEMENT_NAMES, + QUAT_XYZW_ELEMENT_NAMES, + WRENCH6_ELEMENT_NAMES, + XYZ_ELEMENT_NAMES, + LeappTensorSemantics, + body_names_resolver, + body_pose6_resolver, + body_pose_resolver, + body_quat_resolver, + body_wrench_resolver, + body_xyz_resolver, + joint_names_resolver, + leapp_tensor_semantics, + resolve_leapp_element_names, + target_frame_pose_resolver, + target_frame_quat_resolver, + target_frame_xyz_resolver, +) +from .manual_inputs import leapp_input_tensor, leapp_real_env +from .utils import ( + build_command_connection, + build_state_connection, + build_write_connection, +) diff --git a/source/isaaclab/isaaclab/utils/leapp/export_annotator.py b/source/isaaclab/isaaclab/utils/leapp/export_annotator.py new file mode 100644 index 000000000000..889a5c8a5da9 --- /dev/null +++ b/source/isaaclab/isaaclab/utils/leapp/export_annotator.py @@ -0,0 +1,744 @@ +# 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 + +"""Export annotations for Isaac Lab policies using proxy-based patching. + +Observation and action annotation share a unified dedup cache so that a +state property (e.g. ``joint_pos``) read by both an observation term and +an action term resolves to one LEAPP input edge. + +- Observation term functions see an ``_EnvProxy`` whose scene returns + ``_EntityProxy`` objects with annotating data proxies. + +- Action terms have their ``_asset`` attribute replaced with an + _ArticulationWriteProxy that intercepts ``_leapp_semantics``-decorated + write methods **and** routes ``.data`` reads through the same annotating + data proxy used by observations. + +Cache lifecycle (assuming single-env play-mode export): + + compute() clear cache → obs terms populate cache + policy inference TracedTensors propagate through NN + process_action() register_buffer for raw_actions + apply_action() [tracing] reuse cached TracedTensors for state reads, + capture write outputs, call output_tensors(), + then clear cache + apply_action() [decim.] clear cache → fresh reads for simulation + ... + compute() clear cache → fresh reads for next obs +""" + +from __future__ import annotations + +import inspect +import logging +from collections.abc import Callable +from contextlib import suppress +from typing import TYPE_CHECKING, Any + +import torch +import warp as wp +from leapp import annotate +from leapp.utils.tensor_description import TensorSemantics + +from isaaclab.assets.articulation.base_articulation import BaseArticulation +from isaaclab.managers import ManagerTermBase + +from .leapp_semantics import select_element_names +from .proxy import _ArticulationWriteProxy, _DataProxy, _EnvProxy, _ManagerTermProxy +from .utils import ( + TracedProxyArray, + build_command_connection, + build_write_connection, +) + +if TYPE_CHECKING: + from isaaclab.envs import ManagerBasedEnv + + +VARIABLE_IMPEDANCE_MODES = frozenset({"variable", "variable_kp"}) + + +def _to_torch_array(value): + """Return a torch tensor for ProxyArray or raw Warp array values.""" + return value.torch if hasattr(value, "torch") else wp.to_torch(value) + + +# ══════════════════════════════════════════════════════════════════ +# ExportPatcher +# ══════════════════════════════════════════════════════════════════ + + +class ExportPatcher: + """Unified patcher that annotates observation inputs and action outputs for LEAPP export. + + Observation-side property semantics are resolved lazily inside + ``_DataProxy`` by combining: + + - the concrete runtime getter from the backend data class + - the nearest ``_leapp_semantics`` metadata found while walking the MRO + + This lets backends override property implementations without duplicating + decorators from the abstract API. + + - The observation proxy chain (``_EnvProxy`` → ``_SceneProxy`` → + ``_EntityProxy`` → ``_DataProxy``) for state reads + by observation term functions. + - The ``_ArticulationWriteProxy`` on each action term, which intercepts + target writes **and** routes ``.data`` reads through the same + ``_DataProxy`` / cache. + + """ + + def __init__(self, export_method: str, required_obs_groups: set[str] | None = None): + """Initialize the export patcher. + + Args: + export_method: LEAPP export backend passed to + :func:`annotate.output_tensors`. + required_obs_groups: Observation groups that should be patched, or + ``None`` to patch all groups. + """ + self.task_name: str | None = None + self.export_method = export_method + self.required_obs_groups = required_obs_groups + self._annotated_tensor_cache: dict[tuple[int, str], TracedProxyArray] = {} + self._data_property_resolution_cache: dict[tuple[type, str], tuple[Callable, object] | None] = {} + self._write_method_resolution_cache: dict[ + tuple[type, str], tuple[Callable, object, inspect.Signature] | None + ] = {} + self._action_output_cache: list[TensorSemantics] = [] + self._captured_write_term_names: set[str] = set() + self._fallback_term_names: set[str] = set() + self._pending_action_output_export: bool = False + self._uses_last_action_state: bool = False + self._action_term_scene_keys: dict[str, str] = {} + + def setup(self, env): + """Patch the environment in place for LEAPP-aware export. + + Args: + env: Wrapped manager-based environment whose unwrapped instance + should be patched. + """ + unwrapped = env.env.unwrapped + task_name = str(unwrapped.spec.id) + self.task_name = task_name + + proxy_env = _EnvProxy( + unwrapped, + task_name, + self._data_property_resolution_cache, + self._annotated_tensor_cache, + ) + + self._disable_training_managers(unwrapped) + self._patch_observation_manager(unwrapped.observation_manager, proxy_env) + self._patch_history_buffers(unwrapped.observation_manager) + self._patch_action_manager( + unwrapped.action_manager, + self._annotated_tensor_cache, + ) + + # ── Disable training-only managers ───────────────────────────── + + @staticmethod + def _disable_training_managers(unwrapped): + """Replace training-only manager methods with no-ops. + + During export the curriculum, reward, termination, and recorder + managers serve no purpose. Disabling them avoids side-effect + crashes (e.g. ADR curriculum terms accessing nullified noise + configs) and removes unnecessary computation. + + Args: + unwrapped: Unwrapped environment whose training-only managers + should be disabled. + """ + num_envs = unwrapped.num_envs + device = unwrapped.device + _zero_reward = torch.zeros(num_envs, device=device) + _no_termination = torch.zeros(num_envs, dtype=torch.bool, device=device) + + def _noop_curriculum(env_ids=None): + return None + + def _zero_reward_compute(dt): + return _zero_reward + + def _no_termination_compute(): + return _no_termination + + def _noop(*args, **kwargs): + return None + + if hasattr(unwrapped, "curriculum_manager"): + unwrapped.curriculum_manager.compute = _noop_curriculum + + if hasattr(unwrapped, "reward_manager"): + unwrapped.reward_manager.compute = _zero_reward_compute + + if hasattr(unwrapped, "termination_manager"): + unwrapped.termination_manager.compute = _no_termination_compute + + if hasattr(unwrapped, "recorder_manager"): + rm = unwrapped.recorder_manager + + rm.record_pre_step = _noop + rm.record_post_step = _noop + rm.record_pre_reset = _noop + rm.record_post_reset = _noop + rm.record_post_physics_decimation_step = _noop + + @staticmethod + def _resolve_scene_entity_key(scene, entity: Any) -> str | None: + """Return the scene dictionary key for an entity. + + Args: + scene: Scene object that stores entity dictionaries. + entity: Entity instance to locate. + + Returns: + The scene key for ``entity`` if found, otherwise ``None``. + """ + for attr_value in vars(scene).values(): + if not isinstance(attr_value, dict): + continue + for key, candidate in attr_value.items(): + if candidate is entity: + return key + return None + + # ── Observation manager patches ─────────────────────────────── + + def _patch_history_buffers(self, obs_manager): + """Patch history-enabled observation buffers to export as LEAPP state. + + Args: + obs_manager: Observation manager whose history buffers should be + wrapped. + """ + history_buffers = getattr(obs_manager, "_group_obs_term_history_buffer", {}) + term_names_by_group = getattr(obs_manager, "_group_obs_term_names", {}) + + for group_name, term_cfgs in obs_manager._group_obs_term_cfgs.items(): + if self.required_obs_groups is not None and group_name not in self.required_obs_groups: + continue + group_buffers = history_buffers.get(group_name, {}) + group_term_names = term_names_by_group.get(group_name, []) + + for index, term_cfg in enumerate(term_cfgs): + history_length = getattr(term_cfg, "history_length", 0) or 0 + if history_length <= 0: + continue + + if index >= len(group_term_names): + continue + + term_name = group_term_names[index] + circular_buffer = group_buffers.get(term_name) + if circular_buffer is None: + continue + + state_name = f"h_{group_name}_{term_name}" + self._patch_history_buffer_append(circular_buffer, state_name) + + def _patch_history_buffer_append(self, circular_buffer, state_name: str): + """Wrap ``_append`` so history buffers become explicit LEAPP state. + + Args: + circular_buffer: Circular buffer instance to patch. + state_name: LEAPP state tensor name for the buffer contents. + """ + if hasattr(circular_buffer, "_leapp_original_append"): + return + + task_name = self.task_name + original_append = circular_buffer._append + + def patched_append(data: torch.Tensor): + """Annotate history buffer updates as LEAPP state transitions. + + Args: + data: New observation slice appended to the buffer. + + Returns: + ``None``. + """ + if circular_buffer._buffer is not None: + circular_buffer._buffer = annotate.state_tensors(task_name, {state_name: circular_buffer._buffer}) + + original_append(data) + + if circular_buffer._buffer is not None: + circular_buffer._buffer = annotate.update_state(task_name, {state_name: circular_buffer._buffer}) + + circular_buffer._leapp_original_append = original_append + circular_buffer._append = patched_append + + def _patch_observation_manager(self, obs_manager, proxy_env): + """Patch observation terms to use annotating proxies and disable noise. + + Args: + obs_manager: Observation manager instance to patch. + proxy_env: Proxy environment routed into observation terms. + """ + for group_name, term_cfgs in obs_manager._group_obs_term_cfgs.items(): + if self.required_obs_groups is not None and group_name not in self.required_obs_groups: + continue + for term_cfg in term_cfgs: + original_func = term_cfg.func + func_name = getattr(original_func, "__name__", None) + + if func_name == "last_action": + self._uses_last_action_state = True + term_cfg.func = self._wrap_last_action(original_func) + elif func_name == "generated_commands": + term_cfg.func = self._wrap_generated_commands(original_func, term_cfg) + elif func_name == "projected_gravity": + term_cfg.func = self._wrap_projected_gravity(original_func, proxy_env) + else: + term_cfg.func = self._wrap_with_proxy(original_func, proxy_env) + + term_cfg.noise = None + + original_compute = obs_manager.compute + cache = self._annotated_tensor_cache + + def patched_compute(*args, **kwargs): + """Clear the tensor dedup cache once per full observation pass.""" + cache.clear() + return original_compute(*args, **kwargs) + + obs_manager.compute = patched_compute + + # ── Action manager patches ──────────────────────────────────── + + def _patch_action_manager(self, action_manager, cache): + """Patch action terms with write/read proxies and manager hooks. + + Args: + action_manager: Action manager instance to patch. + cache: Shared tensor dedup cache for annotated state reads. + """ + assert self.task_name is not None + scene = action_manager._env.scene + for term_name, term in action_manager._terms.items(): + asset = getattr(term, "_asset", None) + if isinstance(asset, BaseArticulation): + real_asset: BaseArticulation = asset + scene_key = self._resolve_scene_entity_key(scene, real_asset) or "ego" + data_proxy = _DataProxy( + real_asset.data, + scene_key, + self.task_name, + self._data_property_resolution_cache, + cache, + input_name_resolver=lambda prop_name, k=scene_key: f"{k}_{prop_name}", + ) + term._asset = _ArticulationWriteProxy( + real_asset=real_asset, + entity_name=scene_key, + term_name=term_name, + output_cache=self._action_output_cache, + method_resolution_cache=self._write_method_resolution_cache, + captured_write_term_names=self._captured_write_term_names, + data_proxy=data_proxy, + ) + self._action_term_scene_keys[term_name] = scene_key + + self._patch_action_manager_methods(action_manager) + + def _patch_action_manager_methods(self, action_manager): + """Patch ``process_action`` and ``apply_action`` on the action manager instance. + + ``process_action`` registers raw_action buffers for LEAPP tracing and + preserves the action tensor clone. + + ``apply_action`` coordinates the cache and output lifecycle: + + - **Tracing pass** (first ``apply_action`` after ``process_action``): + The cache still holds TracedTensors populated by ``compute_group``. + Action terms that read state (e.g. ``RelativeJointPositionAction`` + reading ``joint_pos``) get those TracedTensors from the cache, + keeping the LEAPP graph connected. After ``output_tensors()`` the + cache is cleared so subsequent decimation sub-steps read fresh values. + + - **Non-tracing passes** (remaining decimation sub-steps and all + subsequent iterations): The cache is cleared **before** running + action terms so every ``.data`` read returns the current simulator + value, preserving simulation correctness. + + Args: + action_manager: Action manager whose instance methods should be + wrapped. + """ + original_process = action_manager.process_action + original_apply = action_manager.apply_action + task_name = self.task_name + cache = self._annotated_tensor_cache + + def patched_process_action(action: torch.Tensor): + """Register raw_action buffers, call real process_action, preserve action clone.""" + original_process(action) + action_manager._action = action.clone() + self._pending_action_output_export = True + + def patched_apply_action(): + """Coordinate cache lifecycle and LEAPP output annotation.""" + if not self._pending_action_output_export: + cache.clear() + return original_apply() + + # Tracing pass: cache still holds TracedTensors from compute_group. + self._action_output_cache.clear() + self._captured_write_term_names.clear() + original_apply() + + self._action_output_cache.extend(self._collect_action_outputs(action_manager)) + self._action_output_cache.extend(self._collect_processed_action_fallbacks(action_manager)) + if self._uses_last_action_state: + annotate.update_state(task_name, {"last_action": action_manager._action}) + fallback_terms = self._fallback_term_names + static_values = self._collect_action_static_outputs(action_manager, fallback_terms) + annotate.output_tensors( + task_name, + self._action_output_cache, + static_outputs=static_values or None, + export_with=self.export_method, + ) + self._pending_action_output_export = False + self._action_output_cache.clear() + cache.clear() + return None + + action_manager.process_action = patched_process_action + action_manager.apply_action = patched_apply_action + + # ── Observation term wrappers ───────────────────────────────── + + @staticmethod + def _wrap_with_proxy(original_func, proxy_env): + """Wrap a term function so it receives the proxy env. + + Args: + original_func: Original observation term function or manager term. + proxy_env: Proxy environment routed into the wrapped callable. + + Returns: + Wrapped callable that substitutes ``proxy_env`` for the real env. + """ + + if isinstance(original_func, ManagerTermBase): + return _ManagerTermProxy(original_func, proxy_env) + + def wrapped(*args, **kwargs): + """Invoke the original function with the proxy environment. + + Args: + *args: Original positional arguments. + **kwargs: Original keyword arguments. + + Returns: + Result of the wrapped observation term. + """ + if args: + args = (proxy_env, *args[1:]) + else: + args = (proxy_env,) + return original_func(*args, **kwargs) + + wrapped.__name__ = getattr(original_func, "__name__", "unknown") + return wrapped + + @staticmethod + def _wrap_projected_gravity(original_func, proxy_env): + """Wrap projected gravity as root-quaternion input plus fixed gravity projection. + + Deployment backends generally provide body orientation, not an already + projected gravity vector. During export, keep the policy observation as + projected gravity while exposing ``root_quat_w`` at the LEAPP graph + boundary. + """ + + def wrapped(*args, **kwargs): + """Compute projected gravity from an annotated root quaternion input.""" + kwargs.pop("inspect", None) + asset_cfg = kwargs.get("asset_cfg") + if asset_cfg is None and len(args) > 1: + asset_cfg = args[1] + asset_name = getattr(asset_cfg, "name", "robot") + root_quat_w = proxy_env.scene[asset_name].data.root_quat_w.torch + gravity_w = torch.zeros((*root_quat_w.shape[:-1], 3), dtype=root_quat_w.dtype, device=root_quat_w.device) + gravity_w[..., 2] = -1.0 + quat_xyz = root_quat_w[..., :3] + quat_w = root_quat_w[..., 3:4] + t = torch.cross(quat_xyz, gravity_w, dim=-1) * 2.0 + return gravity_w - quat_w * t + torch.cross(quat_xyz, t, dim=-1) + + wrapped.__name__ = getattr(original_func, "__name__", "unknown") + return wrapped + + def _wrap_last_action(self, original_func): + """Wrap ``last_action`` as a LEAPP state tensor. + + ``last_action`` is feedback state, not a regular dangling input. We + therefore register it through ``annotate.state_tensors(...)`` on the + observation side and update it through ``annotate.update_state(...)`` + after the traced action pass. + + Args: + original_func: Original ``last_action`` observation term. + + Returns: + Wrapped callable that exports ``last_action`` as LEAPP state. + """ + task_name = self.task_name + + def wrapped(env, action_name=None, **kwargs): + """Run the wrapped ``last_action`` term and annotate its output. + + Args: + env: Environment passed by the observation manager. + action_name: Optional action term name. + **kwargs: Additional keyword arguments for the term. + + Returns: + Annotated last-action tensor. + """ + result = original_func(env, action_name, **kwargs) + return annotate.state_tensors(task_name, {"last_action": result}) + + wrapped.__name__ = original_func.__name__ + return wrapped + + def _wrap_generated_commands(self, original_func, term_cfg): + """Wrap the ``generated_commands`` observation term to annotate its output as a LEAPP input. + + Resolves command semantics (kind, element_names) from the command manager + configuration when available. + + Args: + original_func: Original ``generated_commands`` observation term. + term_cfg: Observation term config used to resolve the command name. + + Returns: + Wrapped callable that exports generated commands as LEAPP inputs. + """ + task_name = self.task_name + command_name_from_cfg = term_cfg.params.get("command_name") + + def wrapped(env, command_name=None, **kwargs): + """Run the wrapped command term and annotate its output. + + Args: + env: Environment passed by the observation manager. + command_name: Optional command term name override. + **kwargs: Additional keyword arguments for the term. + + Returns: + Annotated command tensor. + """ + result = original_func(env, command_name, **kwargs) + leapp_input_name = command_name or command_name_from_cfg or "commands" + command_cfg = None + with suppress(AttributeError, KeyError): + command_cfg = env.command_manager.get_term(leapp_input_name).cfg + sem = TensorSemantics( + name=leapp_input_name, + ref=result, + kind=getattr(command_cfg, "cmd_kind", None), + element_names=getattr(command_cfg, "element_names", None), + extra=build_command_connection(leapp_input_name), + ) + return annotate.input_tensors(task_name, sem) + + wrapped.__name__ = original_func.__name__ + return wrapped + + # ── Output collection ───────────────────────────────────────── + + def _collect_action_outputs(self, action_manager) -> list[TensorSemantics]: + """Collect non-writer action tensors that should be exported. + + Args: + action_manager: Action manager whose terms should be inspected. + + Returns: + Exportable tensor semantics for dynamic action outputs such as OSC + gains. + """ + tensors: list[TensorSemantics] = [] + for term_name, term in action_manager._terms.items(): + osc = getattr(term, "_osc", None) + if osc and hasattr(osc, "cfg") and osc.cfg.impedance_mode in VARIABLE_IMPEDANCE_MODES: + asset = getattr(term, "_asset", None) + real_asset = getattr(asset, "_real_asset", asset) + joint_ids = getattr(term, "_joint_ids", None) + joint_names = getattr(real_asset, "joint_names", None) if real_asset else None + scene_key = self._action_term_scene_keys.get(term_name, "ego") + tensors.append( + TensorSemantics( + name=f"{term_name}_kp_gains", + ref=torch.diagonal(osc._motion_p_gains_task, dim1=-2, dim2=-1), + kind="kp", + element_names=select_element_names(joint_names, joint_ids), + extra=build_write_connection(scene_key, "write_joint_stiffness_to_sim_index"), + ) + ) + tensors.append( + TensorSemantics( + name=f"{term_name}_kd_gains", + ref=torch.diagonal(osc._motion_d_gains_task, dim1=-2, dim2=-1), + kind="kd", + element_names=select_element_names(joint_names, joint_ids), + extra=build_write_connection(scene_key, "write_joint_damping_to_sim_index"), + ) + ) + return tensors + + def _collect_processed_action_fallbacks(self, action_manager) -> list[TensorSemantics]: + """Fallback: use ``term.processed_actions`` for terms that produced no write outputs. + + When an action term does not call any ``_leapp_semantics``-decorated write method + (e.g. ``PreTrainedPolicyAction`` which delegates writes to a nested sub-policy), + we fall back to capturing ``term.processed_actions`` as the output tensor. + + Args: + action_manager: Action manager whose terms should be inspected. + + Returns: + Fallback tensor semantics built from ``processed_actions``. + """ + logger = logging.getLogger(__name__) + fallback_terms: set[str] = set() + tensors: list[TensorSemantics] = [] + for term_name, term in action_manager._terms.items(): + if term_name in self._captured_write_term_names: + continue + processed = getattr(term, "processed_actions", None) + if processed is None: + continue + if isinstance(processed, torch.Tensor): + logger.warning( + "Action term '%s' did not write to any asset directly. Falling back to processed_actions as the" + " export output.\nIf you wish to add semantic data to this policy, you need to manually annotate it" + " with output_tensors.", + term_name, + ) + tensors.append( + TensorSemantics( + name=term_name, + ref=processed.clone(), + kind=None, + element_names=None, + ) + ) + fallback_terms.add(term_name) + self._fallback_term_names = fallback_terms + return tensors + + def _collect_action_static_outputs( + self, action_manager, skip_terms: set[str] | None = None + ) -> list[TensorSemantics]: + """Collect static kp/kd gain values from action terms for export metadata. + + Terms in ``skip_terms`` are excluded — these are terms that fell back + to ``processed_actions`` and whose static gains (kp/kd) belong to a + lower abstraction level that is not part of the exported policy. + + Args: + action_manager: Action manager whose terms should be inspected. + skip_terms: Action term names whose static outputs should be + skipped. + + Returns: + Static tensor semantics for action gains exported as metadata. + """ + static_values: list[TensorSemantics] = [] + for term_name, term in action_manager._terms.items(): + if skip_terms and term_name in skip_terms: + continue + osc = getattr(term, "_osc", None) + if osc and hasattr(osc, "cfg") and osc.cfg.impedance_mode in VARIABLE_IMPEDANCE_MODES: + continue + asset = getattr(term, "_asset", None) + real_asset = getattr(asset, "_real_asset", asset) + if real_asset and hasattr(real_asset, "data"): + data = real_asset.data + joint_ids = getattr(term, "_joint_ids", None) + joint_names = getattr(real_asset, "joint_names", None) + scene_key = self._action_term_scene_keys.get(term_name, "ego") + if hasattr(data, "default_joint_stiffness") and data.default_joint_stiffness is not None: + gains = _to_torch_array(data.default_joint_stiffness) + static_values.append( + TensorSemantics( + name=f"{term_name}_kp_gains", + ref=gains[:, joint_ids] if joint_ids else gains, + kind="kp", + element_names=select_element_names(joint_names, joint_ids), + extra=build_write_connection(scene_key, "write_joint_stiffness_to_sim_index"), + ) + ) + if hasattr(data, "default_joint_damping") and data.default_joint_damping is not None: + gains = _to_torch_array(data.default_joint_damping) + static_values.append( + TensorSemantics( + name=f"{term_name}_kd_gains", + ref=gains[:, joint_ids] if joint_ids else gains, + kind="kd", + element_names=select_element_names(joint_names, joint_ids), + extra=build_write_connection(scene_key, "write_joint_damping_to_sim_index"), + ) + ) + return static_values + + +# ══════════════════════════════════════════════════════════════════ +# Public entry point +# ══════════════════════════════════════════════════════════════════ + + +def patch_env_for_export( + env: ManagerBasedEnv, + export_method: str, + required_obs_groups: set[str] | None = None, +) -> None: + """Patch the env's observation and action managers for LEAPP export. + + This is a thin public entry point around ``ExportPatcher``. It mutates + the provided env instance in-place so that: + + - Observation terms route through proxy objects that annotate tensor + reads from **any** scene entity data class (articulations, rigid + objects, sensors, etc.). + - Action terms route through proxy objects that annotate both data + reads **and** ``Articulation`` write methods. + + Data properties are resolved lazily through proxies — no hardcoded + class list is required. To produce LEAPP input annotations, the + accessed data property getter must carry ``_leapp_semantics``. + Likewise, action-side write methods must be annotated to produce + semantic LEAPP outputs. Undecorated reads and writes are forwarded + as normal runtime access, but they do not gain semantic annotation + metadata through this patching path. + + State reads are deduplicated across observation and action paths via a + shared cache, so a property like ``joint_pos`` that is read by both an + observation term and a relative-position action term appears as a single + LEAPP input edge. + + The underlying env, scene, assets, and tensors remain shared with the rest + of the pipeline; only the manager call paths are redirected. + + Args: + env: Manager-based environment to patch in place. + export_method: LEAPP export backend passed to + :func:`annotate.output_tensors`. + required_obs_groups: Observation groups that should be patched, or + ``None`` to patch all groups. + """ + patcher = ExportPatcher(export_method, required_obs_groups=required_obs_groups) + patcher.setup(env) diff --git a/source/isaaclab/isaaclab/utils/leapp/leapp_semantics.py b/source/isaaclab/isaaclab/utils/leapp/leapp_semantics.py new file mode 100644 index 000000000000..340291de16ad --- /dev/null +++ b/source/isaaclab/isaaclab/utils/leapp/leapp_semantics.py @@ -0,0 +1,145 @@ +# 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 + +"""LEAPP semantic metadata helpers for raw tensor-producing functions.""" + +from __future__ import annotations + +from collections.abc import Callable +from contextlib import suppress +from dataclasses import dataclass +from typing import Any + +try: + from leapp import InputKindEnum, OutputKindEnum +except ImportError: + + class _LeappEnumSentinel: + """Stand-in when leapp is not installed. + + Any attribute access returns ``None`` so that + ``@leapp_tensor_semantics(kind=InputKindEnum.BODY_POSE)`` + silently stores ``kind=None`` instead of crashing at import time. + The real enum values are only needed at export time, when leapp + *is* guaranteed to be available. + """ + + def __getattr__(self, name: str): + return None + + InputKindEnum = _LeappEnumSentinel() # type: ignore[assignment,misc] + OutputKindEnum = _LeappEnumSentinel() # type: ignore[assignment,misc] + + +@dataclass(frozen=True) +class LeappTensorSemantics: + """Semantic metadata attached directly to a raw tensor-producing function.""" + + kind: Any = None + element_names: list[str] | list[list[str]] | None = None + element_names_resolver: Callable | None = None + const: bool = False + + +XYZ_ELEMENT_NAMES: list[str] = ["x", "y", "z"] +QUAT_XYZW_ELEMENT_NAMES: list[str] = ["qx", "qy", "qz", "qw"] +POSE7_ELEMENT_NAMES: list[str] = ["x", "y", "z", "qx", "qy", "qz", "qw"] +POSE6_ELEMENT_NAMES: list[str] = ["x", "y", "z", "angular_x", "angular_y", "angular_z"] +WRENCH6_ELEMENT_NAMES: list[str] = ["fx", "fy", "fz", "tx", "ty", "tz"] + + +def select_element_names(names: list[str] | None, indices: Any = None) -> list[str] | None: + """Select element names using optional runtime indices.""" + if names is None: + return None + if indices is None or indices == slice(None): + return list(names) + if isinstance(indices, slice): + return list(names[indices]) + with suppress(AttributeError): + indices = indices.tolist() + if isinstance(indices, (list, tuple)): + return [names[int(index)] for index in indices] + if isinstance(indices, int): + return [names[indices]] + return None + + +def leapp_tensor_semantics( + *, + kind: Any = None, + element_names: list[str] | list[list[str]] | None = None, + element_names_resolver: Callable | None = None, + const: bool = False, +) -> Callable: + """Attach LEAPP semantic metadata to a raw tensor-producing function.""" + + semantics = LeappTensorSemantics( + kind=kind, + element_names=element_names, + element_names_resolver=element_names_resolver, + const=const, + ) + + def _apply(func: Callable) -> Callable: + func._leapp_semantics = semantics + return func + + return _apply + + +def resolve_leapp_element_names(semantics: LeappTensorSemantics | None, data_self) -> list | None: + """Resolve element names from attached semantics and a tensor-producing object.""" + if semantics is None: + return None + if semantics.element_names is not None: + return semantics.element_names + if semantics.element_names_resolver is not None: + return semantics.element_names_resolver(data_self) + return None + + +# ── Predefined element-name resolvers ───────────────────────────── + + +def joint_names_resolver(data_self) -> list[str] | None: + """Resolve joint element names from the data object at trace time.""" + return select_element_names( + getattr(data_self, "joint_names", getattr(data_self, "_joint_names", None)), + getattr(data_self, "_joint_ids", None), + ) + + +def body_names_resolver(data_self) -> list[str] | None: + """Resolve body element names from the data object at trace time.""" + return select_element_names( + getattr(data_self, "body_names", getattr(data_self, "_body_names", None)), + getattr(data_self, "_body_ids", None), + ) + + +def _compound_resolver(outer_fn: Callable, inner_names: list[str]) -> Callable: + """Build a 2D resolver: ``[outer_names, inner_constant_names]``.""" + + def resolver(data_self) -> list | None: + outer = outer_fn(data_self) + return [outer, inner_names] if outer else None + + return resolver + + +def _target_frame_names(data_self) -> list[str] | None: + names = getattr(data_self, "target_frame_names", None) + return list(names) if names is not None else None + + +body_xyz_resolver = _compound_resolver(body_names_resolver, XYZ_ELEMENT_NAMES) +body_pose_resolver = _compound_resolver(body_names_resolver, POSE7_ELEMENT_NAMES) +body_pose6_resolver = _compound_resolver(body_names_resolver, POSE6_ELEMENT_NAMES) +body_quat_resolver = _compound_resolver(body_names_resolver, QUAT_XYZW_ELEMENT_NAMES) +body_wrench_resolver = _compound_resolver(body_names_resolver, WRENCH6_ELEMENT_NAMES) +target_frame_xyz_resolver = _compound_resolver(_target_frame_names, XYZ_ELEMENT_NAMES) +target_frame_quat_resolver = _compound_resolver(_target_frame_names, QUAT_XYZW_ELEMENT_NAMES) +target_frame_pose_resolver = _compound_resolver(_target_frame_names, POSE7_ELEMENT_NAMES) diff --git a/source/isaaclab/isaaclab/utils/leapp/manual_inputs.py b/source/isaaclab/isaaclab/utils/leapp/manual_inputs.py new file mode 100644 index 000000000000..8d1edfcad7fa --- /dev/null +++ b/source/isaaclab/isaaclab/utils/leapp/manual_inputs.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 + +"""Helpers for manually exposing computed tensors as LEAPP inputs.""" + +from __future__ import annotations + +from typing import Any + +import torch + + +class _ManualTensorView: + """Small adapter matching the ``data..torch`` access pattern.""" + + def __init__(self, tensor: torch.Tensor): + self._tensor = tensor + + @property + def torch(self) -> torch.Tensor: + return self._tensor + + +def _cache_manual_data_proxy(data_proxy: Any, real_data: Any, property_name: str, tensor: torch.Tensor) -> None: + """Cache a manual input tensor on a LEAPP data proxy.""" + manual_cache = object.__getattribute__(data_proxy, "_manual_cache") + manual_cache[(id(real_data), property_name)] = _ManualTensorView(tensor) + + +def leapp_real_env(env: Any) -> Any: + """Return the wrapped real environment when LEAPP passes an export proxy.""" + if type(env).__name__ == "_EnvProxy": + return object.__getattribute__(env, "_real_env") + return env + + +def leapp_input_tensor( + env: Any, + name: str, + tensor: torch.Tensor, + *, + kind: Any | None = None, + element_names: list[str] | list[list[str]] | None = None, + connection: str | dict[str, str] | None = None, + cache: tuple[str, str] | None = None, +) -> torch.Tensor: + """Expose a tensor as a LEAPP input during export and return it unchanged otherwise. + + Args: + env: Environment passed into the observation term. + name: LEAPP input tensor name. + tensor: Tensor to expose. + kind: Optional LEAPP input kind used by deployment tooling. + element_names: Optional element names for tensor dimensions. + connection: Optional Isaac Lab connection string or metadata dict. + cache: Optional ``(entity_name, property_name)`` tuple. When provided, + the annotated tensor is cached as that LEAPP scene-data property read. + + Returns: + The annotated tensor during LEAPP export, otherwise ``tensor`` unchanged. + """ + if type(env).__name__ != "_EnvProxy": + return tensor + + from leapp import annotate + from leapp.utils.tensor_description import TensorSemantics + + extra = None + if isinstance(connection, str): + extra = {"isaaclab_connection": connection} + elif connection is not None: + extra = connection + + if kind is None and element_names is None and extra is None: + annotated = annotate.input_tensors(env.unwrapped.spec.id, {name: tensor}) + else: + semantics = TensorSemantics( + name=name, + ref=tensor, + kind=kind, + element_names=element_names, + extra=extra, + ) + annotated = annotate.input_tensors(env.unwrapped.spec.id, semantics) + + if cache is not None: + entity_name, property_name = cache + real_env = leapp_real_env(env) + real_data = real_env.scene[entity_name].data + scene_proxy = object.__getattribute__(env, "_scene_proxy") + proxy_cache = object.__getattribute__(scene_proxy, "_cache") + tensor_view = _ManualTensorView(annotated) + proxy_cache[(id(real_data), property_name)] = tensor_view + + proxied_entities = object.__getattribute__(scene_proxy, "_proxied") + entity_proxy = proxied_entities.get(entity_name) + if entity_proxy is not None: + data_proxy = object.__getattribute__(entity_proxy, "_data_proxy") + _cache_manual_data_proxy(data_proxy, real_data, property_name, annotated) + + action_manager = getattr(real_env, "action_manager", None) + if action_manager is not None: + for action_term in action_manager._terms.values(): + asset_proxy = getattr(action_term, "_asset", None) + if type(asset_proxy).__name__ != "_ArticulationWriteProxy": + continue + if object.__getattribute__(asset_proxy, "_entity_name") != entity_name: + continue + data_proxy = object.__getattribute__(asset_proxy, "_data_proxy") + _cache_manual_data_proxy(data_proxy, real_data, property_name, annotated) + + return annotated diff --git a/source/isaaclab/isaaclab/utils/leapp/proxy.py b/source/isaaclab/isaaclab/utils/leapp/proxy.py new file mode 100644 index 000000000000..02b77a753b00 --- /dev/null +++ b/source/isaaclab/isaaclab/utils/leapp/proxy.py @@ -0,0 +1,532 @@ +# 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 inspect +from collections.abc import Callable +from typing import Any, cast + +import torch +from leapp.utils.tensor_description import TensorSemantics + +from isaaclab.managers import ManagerTermBase +from isaaclab.utils.warp.proxy_array import ProxyArray + +from .leapp_semantics import resolve_leapp_element_names +from .utils import TracedProxyArray, build_write_connection + + +def _resolve_annotated_property( + property_resolution_cache: dict[tuple[type, str], tuple[Callable, Any] | None], + real_data: Any, + name: str, +) -> tuple[Callable, Any] | None: + """Resolve a concrete property getter and inherited semantics metadata. + + The execution getter always comes from the concrete runtime class. Semantic + metadata is resolved independently by walking the MRO until a property + definition with ``_leapp_semantics`` is found. This mirrors the output-side + export path, where semantics are authored on the base API while concrete + backends provide the runtime implementation. + """ + cache_key = (type(real_data), name) + if cache_key in property_resolution_cache: + return property_resolution_cache[cache_key] + + execution_prop = getattr(type(real_data), name, None) + if not isinstance(execution_prop, property) or execution_prop.fget is None: + property_resolution_cache[cache_key] = None + return None + + semantics_meta = None + for data_cls in type(real_data).__mro__: + prop = data_cls.__dict__.get(name) + if not isinstance(prop, property) or prop.fget is None: + continue + candidate = getattr(prop.fget, "_leapp_semantics", None) + if candidate is None: + continue + if getattr(candidate, "const", False): + property_resolution_cache[cache_key] = None + return None + semantics_meta = candidate + break + + if semantics_meta is None: + property_resolution_cache[cache_key] = None + return None + + resolution = (execution_prop.fget, semantics_meta) + property_resolution_cache[cache_key] = resolution + return resolution + + +def _resolve_annotated_method( + method_resolution_cache: dict[tuple[type, str], tuple[Callable, Any, inspect.Signature] | None], + real_asset: Any, + name: str, +) -> tuple[Callable, Any, inspect.Signature] | None: + """Resolve a concrete bound method and inherited semantics metadata.""" + cache_key = (type(real_asset), name) + if cache_key in method_resolution_cache: + return method_resolution_cache[cache_key] + + original_method = getattr(real_asset, name, None) + if not callable(original_method): + method_resolution_cache[cache_key] = None + return None + + for asset_cls in type(real_asset).__mro__: + candidate = asset_cls.__dict__.get(name) + if not callable(candidate): + continue + semantics_meta = getattr(candidate, "_leapp_semantics", None) + if semantics_meta is None: + continue + resolution = (original_method, semantics_meta, inspect.signature(candidate)) + method_resolution_cache[cache_key] = resolution + return resolution + + method_resolution_cache[cache_key] = None + return None + + +class _WriteJointNameContext: + """Resolve runtime joint-name subsets for lazy write interception.""" + + __slots__ = ("joint_names", "_joint_ids") + + def __init__(self, joint_names: list[str], joint_ids): + self.joint_names = joint_names + self._joint_ids = joint_ids + + +def _unique_output_name(term_name: str, method_name: str, output_cache: list[TensorSemantics]) -> str: + """Return a stable, unique output name for an action write entry.""" + existing = {t.name for t in output_cache} + candidate = term_name + if candidate in existing: + candidate = f"{term_name}_{method_name}" + suffix = 2 + while candidate in existing: + candidate = f"{term_name}_{method_name}_{suffix}" + suffix += 1 + return candidate + + +class _DataProxy: + """Proxy around a real data object that intercepts tensor-returning property reads. + + The real data object may be any scene entity data class (``ArticulationData``, + ``RigidObjectData``, sensor data classes, etc.). The proxy resolves property + semantics lazily on first access by walking the runtime class MRO. This lets + concrete backend overrides reuse semantic metadata authored on abstract base + properties without copying decorators onto every implementation. + + When a semantic property returns a :class:`~isaaclab.utils.warp.ProxyArray`, + the result is wrapped in a ``TracedProxyArray`` and cached for + deduplication. Non-proxy results and ordinary attributes are forwarded + transparently. + + All other attribute access is forwarded transparently to the real object. + """ + + def __init__( + self, + real_data: Any, + entity_name: str, + task_name: str, + property_resolution_cache: dict[tuple[type, str], tuple[Callable, Any] | None], + cache: dict, + input_name_resolver: Callable, + ): + object.__setattr__(self, "_real_data", real_data) + object.__setattr__(self, "_entity_name", entity_name) + object.__setattr__(self, "_task_name", task_name) + object.__setattr__(self, "_property_resolution_cache", property_resolution_cache) + object.__setattr__(self, "_cache", cache) + object.__setattr__(self, "_manual_cache", {}) + object.__setattr__(self, "_input_name_resolver", input_name_resolver) + + def __getattr__(self, name): + """Intercept semantic property reads; forward everything else.""" + real_data = object.__getattribute__(self, "_real_data") + resolution = _resolve_annotated_property( + object.__getattribute__(self, "_property_resolution_cache"), real_data, name + ) + if resolution is None: + return getattr(real_data, name) + + cache = object.__getattribute__(self, "_cache") + cache_key = (id(real_data), name) + manual_cache = object.__getattribute__(self, "_manual_cache") + if cache_key in manual_cache: + return manual_cache[cache_key] + if cache_key in cache: + return cache[cache_key] + + execution_fget, semantics_meta = resolution + result = execution_fget(real_data) + if isinstance(result, ProxyArray): + proxy_result = result + elif isinstance(result, torch.Tensor): + return result + else: + try: + proxy_result = ProxyArray(result) + except TypeError: + return result + + input_name = object.__getattribute__(self, "_input_name_resolver")(name) + traced = TracedProxyArray( + proxy_result, + input_name=input_name, + semantics_meta=semantics_meta, + real_data=real_data, + entity_name=object.__getattribute__(self, "_entity_name"), + property_name=name, + task_name=object.__getattribute__(self, "_task_name"), + ) + cache[cache_key] = traced + return traced + + +class _EntityProxy: + """Proxy around a real scene entity that returns a ``_DataProxy`` for ``.data``. + + All other attribute access is forwarded transparently to the real asset. + """ + + def __init__(self, real_entity: Any, data_proxy: _DataProxy): + object.__setattr__(self, "_real_entity", real_entity) + object.__setattr__(self, "_data_proxy", data_proxy) + + @property + def data(self): + """Return the annotating data proxy instead of the real data object.""" + return object.__getattribute__(self, "_data_proxy") + + def __getattr__(self, name): + """Forward all non-data attribute access to the real scene entity.""" + return getattr(object.__getattribute__(self, "_real_entity"), name) + + +class _EntityMappingProxy: + """Proxy around a mapping of scene entities that lazily wraps data-producing entries.""" + + def __init__( + self, + real_mapping, + task_name: str, + property_resolution_cache: dict[tuple[type, str], tuple[Callable, Any] | None], + cache: dict, + ): + object.__setattr__(self, "_real_mapping", real_mapping) + object.__setattr__(self, "_task_name", task_name) + object.__setattr__(self, "_property_resolution_cache", property_resolution_cache) + object.__setattr__(self, "_cache", cache) + object.__setattr__(self, "_proxied", {}) + + def __getitem__(self, key): + """Return a proxied entity when it has a ``.data`` attribute.""" + proxied = object.__getattribute__(self, "_proxied") + if key in proxied: + return proxied[key] + real_mapping = object.__getattribute__(self, "_real_mapping") + entity = real_mapping[key] + data = getattr(entity, "data", None) + if data is None: + return entity + data_proxy = _DataProxy( + data, + key, + object.__getattribute__(self, "_task_name"), + object.__getattribute__(self, "_property_resolution_cache"), + object.__getattribute__(self, "_cache"), + input_name_resolver=lambda prop_name: f"{key}_{prop_name}", + ) + proxy = _EntityProxy(entity, data_proxy) + proxied[key] = proxy + return proxy + + def get(self, key, default=None): + """Return a proxied entity when present, default otherwise.""" + real_mapping = object.__getattribute__(self, "_real_mapping") + if key not in real_mapping: + return default + return self[key] + + def __iter__(self): + return iter(object.__getattribute__(self, "_real_mapping")) + + def __len__(self): + return len(object.__getattribute__(self, "_real_mapping")) + + def __getattr__(self, name): + """Forward all other mapping access to the real mapping.""" + return getattr(object.__getattribute__(self, "_real_mapping"), name) + + +class _SceneProxy: + """Proxy around the real InteractiveScene. + + When an observation term looks up a scene entity by name, this proxy lazily + wraps any entity that has a ``.data`` attribute. All tensor-returning + properties on the data object are intercepted for LEAPP annotation. This + covers articulations, rigid objects, and sensors through both + ``scene["name"]`` and ``scene.sensors["name"]`` access paths. + """ + + def __init__( + self, + real_scene, + task_name: str, + property_resolution_cache: dict[tuple[type, str], tuple[Callable, Any] | None], + cache: dict, + ): + object.__setattr__(self, "_real_scene", real_scene) + object.__setattr__(self, "_task_name", task_name) + object.__setattr__(self, "_property_resolution_cache", property_resolution_cache) + object.__setattr__(self, "_cache", cache) + object.__setattr__(self, "_proxied", {}) + object.__setattr__(self, "_sensor_mapping_proxy", None) + + def _maybe_proxy_entity(self, key: str, entity: Any): + """Return a proxy for any entity that has a ``.data`` attribute.""" + proxied = object.__getattribute__(self, "_proxied") + if key in proxied: + return proxied[key] + + data = getattr(entity, "data", None) + if data is None: + return entity + + cache = object.__getattribute__(self, "_cache") + data_proxy = _DataProxy( + data, + key, + object.__getattribute__(self, "_task_name"), + object.__getattribute__(self, "_property_resolution_cache"), + cache, + input_name_resolver=lambda prop_name, k=key: f"{k}_{prop_name}", + ) + proxy = _EntityProxy(entity, data_proxy) + proxied[key] = proxy + return proxy + + def __getitem__(self, key): + """Return a proxied entity when it exposes annotated data getters.""" + real_scene = object.__getattribute__(self, "_real_scene") + entity = real_scene[key] + return self._maybe_proxy_entity(key, entity) + + @property + def sensors(self): + """Return a mapping proxy for scene sensors.""" + sensor_mapping_proxy = object.__getattribute__(self, "_sensor_mapping_proxy") + if sensor_mapping_proxy is None: + real_scene = object.__getattribute__(self, "_real_scene") + sensor_mapping_proxy = _EntityMappingProxy( + real_scene.sensors, + object.__getattribute__(self, "_task_name"), + object.__getattribute__(self, "_property_resolution_cache"), + object.__getattribute__(self, "_cache"), + ) + object.__setattr__(self, "_sensor_mapping_proxy", sensor_mapping_proxy) + return sensor_mapping_proxy + + def __getattr__(self, name): + """Forward all other scene access to the real scene.""" + return getattr(object.__getattribute__(self, "_real_scene"), name) + + +class _EnvProxy: + """Proxy around the real env that returns a _SceneProxy for ``.scene``. + + All other attribute access (``num_envs``, ``command_manager``, etc.) + is forwarded transparently to the real env. + """ + + def __init__( + self, + real_env, + task_name: str, + property_resolution_cache: dict[tuple[type, str], tuple[Callable, Any] | None], + cache: dict, + ): + object.__setattr__(self, "_real_env", real_env) + object.__setattr__( + self, + "_scene_proxy", + _SceneProxy(real_env.scene, task_name, property_resolution_cache, cache), + ) + + @property + def scene(self): + """Return the scene proxy instead of the real scene.""" + return object.__getattribute__(self, "_scene_proxy") + + def __getattr__(self, name): + """Forward all non-scene attribute access to the real env.""" + return getattr(object.__getattribute__(self, "_real_env"), name) + + +def _build_scene_entity_lookup(real_scene) -> dict[int, tuple[str, str]]: + """Map real scene entity object ids to their lookup path.""" + lookup: dict[int, tuple[str, str]] = {} + for attr_name, attr_value in vars(real_scene).items(): + if not isinstance(attr_value, dict): + continue + container_kind = "sensors" if attr_name == "sensors" else "scene" + for key, entity in attr_value.items(): + lookup[id(entity)] = (container_kind, key) + return lookup + + +class _ManagerTermProxy(ManagerTermBase): + """Proxy a class-based manager term while preserving its lifecycle methods. + + Observation manager terms can be stateful ``ManagerTermBase`` instances that + expose ``reset()`` and ``serialize()`` in addition to being callable. This + proxy preserves that interface while swapping the env argument passed into + ``__call__`` for the observation-side proxy env. + """ + + def __init__(self, target: ManagerTermBase, proxy_env: _EnvProxy): + super().__init__(target.cfg, target._env) + self._target = target + self._proxy_env = proxy_env + self._entity_lookup = _build_scene_entity_lookup(target._env.scene) + + @property + def __name__(self) -> str: + """Expose the wrapped term name for compatibility and debugging.""" + return getattr(self._target, "__name__", self._target.__class__.__name__) + + def reset(self, env_ids=None) -> None: + """Forward resets to the wrapped term instance.""" + self._target.reset(env_ids=env_ids) + + def serialize(self) -> dict: + """Forward serialization to the wrapped term instance.""" + return self._target.serialize() + + def __call__(self, *args, **kwargs): + """Call the wrapped term with the proxy env in place of the real env.""" + if args: + args = (self._proxy_env, *args[1:]) + else: + args = (self._proxy_env,) + swapped_attrs: list[tuple[str, Any]] = [] + for attr_name, attr_value in vars(self._target).items(): + lookup = self._entity_lookup.get(id(attr_value)) + if lookup is None: + continue + + container_kind, key = lookup + proxy_entity = ( + self._proxy_env.scene.sensors[key] if container_kind == "sensors" else self._proxy_env.scene[key] + ) + swapped_attrs.append((attr_name, attr_value)) + setattr(self._target, attr_name, proxy_entity) + + try: + return self._target(*args, **kwargs) + finally: + for attr_name, attr_value in swapped_attrs: + setattr(self._target, attr_name, attr_value) + + def __getattr__(self, name): + """Forward all other attribute access to the wrapped term instance.""" + return getattr(self._target, name) + + +# ══════════════════════════════════════════════════════════════════ +# Action-side proxy +# ══════════════════════════════════════════════════════════════════ + + +class _ArticulationWriteProxy: + """Proxy around a real articulation implementation for action terms. + + Intercepts ``_leapp_semantics``-decorated write methods **and** routes + ``.data`` reads through a shared ``_DataProxy`` so that + action-side state reads (e.g. ``self._asset.data.joint_pos`` inside + ``RelativeJointPositionAction``) participate in LEAPP annotation and + share the dedup cache with observation-side reads. + + All other attribute access is forwarded transparently to the real asset. + """ + + def __init__( + self, + real_asset: Any, + entity_name: str, + term_name: str, + output_cache: list[TensorSemantics], + method_resolution_cache: dict[tuple[type, str], tuple[Callable, Any, inspect.Signature] | None], + captured_write_term_names: set[str], + data_proxy: _DataProxy, + ): + object.__setattr__(self, "_real_asset", real_asset) + object.__setattr__(self, "_entity_name", entity_name) + object.__setattr__(self, "_term_name", term_name) + object.__setattr__(self, "_output_cache", output_cache) + object.__setattr__(self, "_method_resolution_cache", method_resolution_cache) + object.__setattr__(self, "_captured_write_term_names", captured_write_term_names) + object.__setattr__(self, "_data_proxy", data_proxy) + + @property + def data(self): + """Return the shared annotating data proxy.""" + return object.__getattribute__(self, "_data_proxy") + + def __getattr__(self, name): + """Return an annotating wrapper for semantic write methods; forward everything else.""" + real_asset = object.__getattribute__(self, "_real_asset") + resolution = _resolve_annotated_method( + object.__getattribute__(self, "_method_resolution_cache"), + real_asset, + name, + ) + if resolution is None: + return getattr(real_asset, name) + + original_method, semantics_meta, signature = resolution + term_name = object.__getattribute__(self, "_term_name") + output_cache = object.__getattribute__(self, "_output_cache") + captured_write_term_names = object.__getattribute__(self, "_captured_write_term_names") + + def interceptor(*args, **kwargs): + result = original_method(*args, **kwargs) + bound_args = signature.bind_partial(real_asset, *args, **kwargs) + target = bound_args.arguments.get("target") + + if not isinstance(target, torch.Tensor): + return result + + target_tensor = cast(torch.Tensor, target) + joint_ids = bound_args.arguments.get("joint_ids") + output_cache.append( + TensorSemantics( + name=_unique_output_name(term_name, name, output_cache), + ref=target_tensor.clone(), + kind=semantics_meta.kind, + element_names=resolve_leapp_element_names( + semantics_meta, + _WriteJointNameContext(real_asset.joint_names, joint_ids), + ), + extra=build_write_connection( + object.__getattribute__(self, "_entity_name"), + name, + ), + ) + ) + captured_write_term_names.add(term_name) + + return result + + return interceptor diff --git a/source/isaaclab/isaaclab/utils/leapp/utils.py b/source/isaaclab/isaaclab/utils/leapp/utils.py new file mode 100644 index 000000000000..2308f662ab80 --- /dev/null +++ b/source/isaaclab/isaaclab/utils/leapp/utils.py @@ -0,0 +1,87 @@ +# 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 types import SimpleNamespace +from typing import Any + +import torch +from leapp import annotate +from leapp.utils.tensor_description import TensorSemantics + +from isaaclab.utils.warp.proxy_array import ProxyArray + +from .leapp_semantics import LeappTensorSemantics, resolve_leapp_element_names + + +class TracedProxyArray(ProxyArray): + _traced_array: torch.Tensor + + def __init__( + self, + proxy_array: ProxyArray, + *, + input_name: str, + semantics_meta: LeappTensorSemantics, + real_data: Any, + entity_name: str, + property_name: str, + task_name: str, + ) -> None: + super().__init__(proxy_array.warp) + astorch = super().torch + sem = TensorSemantics( + name=input_name, + ref=astorch, + kind=semantics_meta.kind, + element_names=resolve_leapp_element_names(semantics_meta, real_data), + extra=build_state_connection(entity_name, property_name), + ) + annotated = annotate.input_tensors(task_name, sem) + object.__setattr__(self, "_traced_array", annotated) + + @property + def torch(self) -> torch.Tensor: + return self._traced_array + + @property + def warp(self) -> Any: + raise AttributeError("warp arrays are not supported for leapp export") + + +def ensure_env_spec_id(env, fallback_task_name: str = "policy") -> str: + """Return ``env.unwrapped.spec.id``, creating a fallback spec when needed.""" + spec = getattr(env.unwrapped, "spec", None) + if spec is None: + env.unwrapped.spec = SimpleNamespace(id=fallback_task_name) + return fallback_task_name + + task_name = getattr(spec, "id", None) + if task_name is None: + spec.id = fallback_task_name + return fallback_task_name + + return task_name + + +# ══════════════════════════════════════════════════════════════════ +# Connection Builders +# ══════════════════════════════════════════════════════════════════ + + +def build_state_connection(entity_name: str, property_name: str) -> dict[str, str]: + """Return a compact deployment connection string for a state property.""" + return {"isaaclab_connection": f"state:{entity_name}:{property_name}"} + + +def build_command_connection(command_name: str) -> dict[str, str]: + """Return a compact deployment connection string for a command term.""" + return {"isaaclab_connection": f"command:{command_name}"} + + +def build_write_connection(entity_name: str, method_name: str) -> dict[str, str]: + """Return a compact deployment connection string for an articulation write target.""" + return {"isaaclab_connection": f"write:{entity_name}:{method_name}"} diff --git a/source/isaaclab/isaaclab/utils/warp/__init__.py b/source/isaaclab/isaaclab/utils/warp/__init__.py index e5a4264c81c7..569b80c9b089 100644 --- a/source/isaaclab/isaaclab/utils/warp/__init__.py +++ b/source/isaaclab/isaaclab/utils/warp/__init__.py @@ -5,6 +5,8 @@ """Sub-module containing operations based on warp.""" +import warnings + import warp as wp wp.config.quiet = True @@ -13,3 +15,47 @@ from isaaclab.utils.module import lazy_export lazy_export() + +# Avoid a circular import at module load: `.proxy_array` imports warp, which is +# already loaded above. Importing it here ensures the class is available inside +# the shim. +from .proxy_array import ProxyArray # noqa: E402 + +_WP_TO_TORCH_ORIGINAL = wp.to_torch +_WP_TO_TORCH_WARNED = False + + +def _wp_to_torch_with_proxyarray(a, requires_grad=None): + """Shim for :func:`warp.to_torch` that gracefully handles :class:`ProxyArray`. + + Without this shim, ``wp.to_torch(proxy)`` would fail with + ``AttributeError: 'ProxyArray' object has no attribute 'requires_grad'`` + because :class:`ProxyArray` intentionally doesn't replicate the full + ``wp.array`` attribute surface. Users and third-party code that still use + ``wp.to_torch(asset.data.)`` from before the ProxyArray migration + would break hard. + + The shim routes :class:`ProxyArray` arguments to their cached ``.torch`` + view (a zero-copy :class:`torch.Tensor` of the same underlying memory) and + emits a one-shot :class:`DeprecationWarning`. For any other input type, + the original :func:`warp.to_torch` handles the call as before. + """ + global _WP_TO_TORCH_WARNED + if isinstance(a, ProxyArray): + if not _WP_TO_TORCH_WARNED: + _WP_TO_TORCH_WARNED = True + warnings.warn( + "wp.to_torch() is deprecated; use the `.torch` accessor on" + " the ProxyArray directly (e.g. `asset.data.joint_pos.torch`).", + DeprecationWarning, + stacklevel=2, + ) + return a.torch + return _WP_TO_TORCH_ORIGINAL(a, requires_grad=requires_grad) + + +# Patch at both the top-level ``warp`` namespace and the underlying module so +# callers using ``import warp as wp`` and rare ``from warp._src.torch import +# to_torch`` patterns both pick up the shim. +wp.to_torch = _wp_to_torch_with_proxyarray +wp._src.torch.to_torch = _wp_to_torch_with_proxyarray # noqa: SLF001 diff --git a/source/isaaclab/isaaclab/utils/warp/__init__.pyi b/source/isaaclab/isaaclab/utils/warp/__init__.pyi index ea25f04435d3..ab1de52f39f7 100644 --- a/source/isaaclab/isaaclab/utils/warp/__init__.pyi +++ b/source/isaaclab/isaaclab/utils/warp/__init__.pyi @@ -4,6 +4,7 @@ # SPDX-License-Identifier: BSD-3-Clause __all__ = [ + "ProxyArray", "convert_to_warp_mesh", "raycast_dynamic_meshes", "raycast_mesh", @@ -11,3 +12,4 @@ __all__ = [ ] from .ops import convert_to_warp_mesh, raycast_dynamic_meshes, raycast_mesh, raycast_single_mesh +from .proxy_array import ProxyArray diff --git a/source/isaaclab/isaaclab/utils/warp/math_ops.py b/source/isaaclab/isaaclab/utils/warp/math_ops.py new file mode 100644 index 000000000000..36a0f960a8b7 --- /dev/null +++ b/source/isaaclab/isaaclab/utils/warp/math_ops.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 + +import warp as wp + + +def transform_to_vec_quat( + t: wp.array, +) -> tuple[wp.array, wp.array]: + """Split a wp.transformf array into position (vec3f) and quaternion (quatf) arrays. + + Zero-copy: returns views into the same underlying memory. + + Args: + t: Array of transforms (dtype=wp.transformf). Shape ``(N,)``, ``(N, M)``, or ``(N, M, K)``. + + Returns: + Tuple of (positions, quaternions) as warp array views with matching dimensionality. + + Raises: + TypeError: If *t* does not have dtype ``wp.transformf``. + """ + if t.dtype != wp.transformf: + raise TypeError(f"Expected wp.transformf array, got dtype={t.dtype}") + floats = t.view(wp.float32) + if t.ndim == 1: + return floats[:, :3].view(wp.vec3f), floats[:, 3:].view(wp.quatf) + if t.ndim == 2: + return floats[:, :, :3].view(wp.vec3f), floats[:, :, 3:].view(wp.quatf) + if t.ndim == 3: + return floats[:, :, :, :3].view(wp.vec3f), floats[:, :, :, 3:].view(wp.quatf) + raise ValueError(f"Expected 1D, 2D, or 3D transform array, got ndim={t.ndim}") diff --git a/source/isaaclab/isaaclab/utils/warp/ops.py b/source/isaaclab/isaaclab/utils/warp/ops.py index 313a7fd43afb..3ecc875ae047 100644 --- a/source/isaaclab/isaaclab/utils/warp/ops.py +++ b/source/isaaclab/isaaclab/utils/warp/ops.py @@ -19,6 +19,64 @@ from . import kernels +# Cache of all-True env masks keyed by (n_envs, device) to avoid per-call allocations in +# raycast_dynamic_meshes. Populated lazily on first call with a given (n_envs, device) pair. +_all_env_mask_cache: dict[tuple[int, str], wp.array] = {} + + +# Tile size for the spatial-axis split in :func:`_uint8_spatial_mean`. Tuned on L40; +# robust across modern NVIDIA arches (Ampere/Ada/Hopper) at R256. +_UINT8_SUM_TILE_HW: int = 32 + +# Cache of int32 partials scratch tensors keyed by (src.shape, device, channel_dim). Avoids +# per-call allocation in :func:`_uint8_spatial_mean`. channel_dim is part of the key so +# BCHW and BHWC inputs with otherwise-identical shape get separate scratch slots. Typically +# holds one entry per training run (camera resolution, device, and layout are all fixed). +_uint8_sum_partials_cache: dict[tuple[tuple[int, ...], str, int], torch.Tensor] = {} + + +def _uint8_spatial_mean(src: torch.Tensor, scale: float, channel_dim: int = 3) -> torch.Tensor: + """Per-(batch, channel) mean of a uint8 image scaled by ``1 / scale``. + + Equivalent to ``src.sum(dim=spatial_dims, dtype=int64).float() / scale`` where + ``spatial_dims`` is the pair of non-batch, non-channel axes. The int64 + promotion is safe at any resolution; the per-tile Warp accumulator stays + int32 (overflow-safe up to ~16M values per tile). + + Args: + src: Input image. Shape is ``(B, H, W, C)`` (BHWC) or ``(B, C, H, W)`` (BCHW), + dtype ``torch.uint8``, contiguous. + scale: Multiplier for the per-channel sum. Pass ``H * W * 255`` to get the mean + of ``src / 255``. + channel_dim: Resolved positive position of the channel axis -- ``1`` (BCHW) or + ``3`` (BHWC). Defaults to ``3`` (BHWC) for back-compat with internal callers. + + Returns: + Per-(batch, channel) mean as float32. Shape is ``(B, C)``. + """ + if channel_dim == 1: + b, c, h, _ = src.shape + else: + b, h, _, c = src.shape + device_str = str(src.device) + cache_key = (src.shape, device_str, channel_dim) + partials = _uint8_sum_partials_cache.get(cache_key) + if partials is None: + num_tiles = (h + _UINT8_SUM_TILE_HW - 1) // _UINT8_SUM_TILE_HW + # C innermost: adjacent threads stride-1 along src's contiguous trailing dim (BHWC fast path). + partials = torch.empty((b, num_tiles, c), dtype=torch.int32, device=src.device) + _uint8_sum_partials_cache[cache_key] = partials + + src_wp = wp.from_torch(src, dtype=wp.uint8) + partials_wp = wp.from_torch(partials, dtype=wp.int32) + wp.launch( + kernel=kernels.spatial_sum_uint8_tiled, + dim=partials.shape, + inputs=[src_wp, partials_wp, _UINT8_SUM_TILE_HW, channel_dim], + device=device_str, + ) + return partials.sum(dim=1, dtype=torch.int64).float() / scale + def raycast_mesh( ray_starts: torch.Tensor, @@ -335,11 +393,19 @@ def raycast_dynamic_meshes( mesh_orientations_w = mesh_orientations_w.to(dtype=torch.float32, device=torch_device).contiguous() mesh_quat_wp_w = wp.from_torch(mesh_orientations_w, dtype=wp.quat) + # All environments active when called through this public API. + # Cache the mask by (n_envs, device) to avoid a per-call allocation. + cache_key = (n_envs, str(torch_device)) + if cache_key not in _all_env_mask_cache: + _all_env_mask_cache[cache_key] = wp.from_torch(torch.ones(n_envs, dtype=torch.bool, device=torch_device)) + all_env_mask = _all_env_mask_cache[cache_key] + # launch the warp kernel wp.launch( kernel=kernels.raycast_dynamic_meshes_kernel, dim=[n_meshes, n_envs, n_rays_per_env], inputs=[ + all_env_mask, mesh_ids_wp, ray_starts_wp, ray_directions_wp, @@ -392,3 +458,87 @@ def convert_to_warp_mesh(points: np.ndarray, indices: np.ndarray, device: str) - points=wp.array(points.astype(np.float32), dtype=wp.vec3, device=device), indices=wp.array(indices.astype(np.int32).flatten(), dtype=wp.int32, device=device), ) + + +def normalize_image_uint8( + src: torch.Tensor, + channel_dim: int = -1, + out: torch.Tensor | None = None, +) -> torch.Tensor: + """Compute ``(src / 255.0) - mean(src / 255.0, spatial_dims, keepdim=True)`` via a fused Warp kernel. + + Equivalent to the pure-PyTorch expression to within float32 precision. Pass an ``out`` + tensor to reuse storage across steps. + + Supports both image layouts via ``channel_dim``: + + - BHWC (``channel_dim=-1`` or ``3``, the default): mean is taken over axes (1, 2). + - BCHW (``channel_dim=-3`` or ``1``): mean is taken over axes (2, 3). + + Note: + Most callers should go through :func:`isaaclab.utils.images.normalize_camera_image`, + which dispatches non-uint8 / non-RGB inputs to a PyTorch fallback. + + Args: + src: Input uint8 image tensor. Shape is ``(B, H, W, C)`` or ``(B, C, H, W)``. + Must be contiguous. + channel_dim: Position of the channel axis. Must resolve to ``1`` (BCHW) or ``3`` (BHWC). + Negative values are supported (``-1`` == BHWC, ``-3`` == BCHW). Defaults to ``-1``. + out: Optional pre-allocated float32 output. Same shape as ``src``, contiguous, on the + same device. If omitted, a fresh tensor is allocated. Defaults to None. + + .. warning:: + + If you pass the same ``out`` tensor across calls that happen on either + side of an environment-step boundary (i.e., the result of one call is + still being read by the RL trainer when the next call is made), the + returned observation will alias the latest call's output and the + trainer will see overwritten data. Use a ping-pong of two ``out`` + buffers, or omit ``out`` entirely, when the result lifetime crosses + ``env.step()`` boundaries. + + Returns: + The normalized float32 tensor. Same object as ``out`` when provided. + + Raises: + ValueError: If ``src`` is not 4D uint8, not contiguous, ``channel_dim`` does not + resolve to 1 or 3, or ``out``'s shape / dtype / device does not match. + """ + if src.dtype != torch.uint8 or src.ndim != 4: + raise ValueError(f"src must be a 4D uint8 tensor; got dtype={src.dtype}, ndim={src.ndim}") + if not src.is_contiguous(): + raise ValueError("src must be contiguous (Warp kernel reads it as a 4D wp.array)") + + # Resolve negative channel_dim to its positive index in [1, src.ndim - 1]. + resolved_channel_dim = channel_dim + src.ndim if channel_dim < 0 else channel_dim + if resolved_channel_dim not in (1, 3): + raise ValueError( + f"channel_dim must resolve to 1 (BCHW) or 3 (BHWC) for 4D input;" + f" got channel_dim={channel_dim} -> {resolved_channel_dim}" + ) + + if out is None: + out = torch.empty(src.shape, dtype=torch.float32, device=src.device) + elif out.shape != src.shape or out.dtype != torch.float32 or out.device != src.device: + raise ValueError( + f"out shape/dtype/device mismatch: expected {tuple(src.shape)}/float32/{src.device}," + f" got {tuple(out.shape)}/{out.dtype}/{out.device}" + ) + elif not out.is_contiguous(): + raise ValueError("out must be contiguous") + + # Spatial dims = the two non-batch, non-channel axes; mean is shape (B, C) for both layouts. + spatial_dims = tuple(d for d in (1, 2, 3) if d != resolved_channel_dim) + spatial_size = src.shape[spatial_dims[0]] * src.shape[spatial_dims[1]] + mean = _uint8_spatial_mean(src, spatial_size * 255.0, channel_dim=resolved_channel_dim) + + src_wp = wp.from_torch(src, dtype=wp.uint8) + mean_wp = wp.from_torch(mean, dtype=wp.float32) + out_wp = wp.from_torch(out, dtype=wp.float32) + wp.launch( + kernel=kernels.normalize_image_uint8, + dim=src.shape, + inputs=[src_wp, mean_wp, out_wp, resolved_channel_dim], + device=str(src.device), + ) + return out diff --git a/source/isaaclab/isaaclab/utils/warp/proxy_array.py b/source/isaaclab/isaaclab/utils/warp/proxy_array.py new file mode 100644 index 000000000000..2a56c9bebc4b --- /dev/null +++ b/source/isaaclab/isaaclab/utils/warp/proxy_array.py @@ -0,0 +1,356 @@ +# 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 dual-access array wrapper with explicit ``.torch`` and ``.warp`` accessors. + +Inspired by ProxyArray from mujocolab/mjlab (BSD-3-Clause). +""" + +from __future__ import annotations + +import os +import warnings +from typing import ClassVar + +import torch +import warp as wp + +_QUATF_ACCESS_WARN_ENV = "WARN_ON_TORCH_QUATF_ACCESS" +"""Environment variable that, when set to ``"1"``, makes :attr:`ProxyArray.torch` +emit a :class:`UserWarning` on every read of a ``wp.quatf``-typed array. Used as a +runtime aid for tracking down call sites that may still assume Isaac Lab 2.x's +``(w, x, y, z)`` quaternion convention after the migration to Isaac Lab 3.x's +``(x, y, z, w)`` convention. See the Isaac Lab 3.0 migration guide for details.""" + + +class ProxyArray: + """Warp-first array wrapper providing cached zero-copy ``.torch`` and ``.warp`` accessors. + + This class wraps a :class:`warp.array` and provides: + + * A ``.warp`` property that returns the original warp array (for kernel interop). + * A ``.torch`` property that returns a cached, zero-copy :class:`torch.Tensor` view + (via :func:`warp.to_torch`). + * Convenience properties (``shape``, ``dtype``, ``device``) delegated to the warp array. + * A deprecation bridge for common torch functions, indexing, and arithmetic/comparison + operators while emitting a one-time :class:`DeprecationWarning`. Tensor instance methods + such as ``clone()`` are not forwarded; use explicit ``.torch`` access for those. + + Example: + + .. code-block:: python + + import warp as wp + from isaaclab.utils.warp.proxy_array import ProxyArray + + arr = wp.zeros(100, dtype=wp.vec3f, device="cuda:0") + ta = ProxyArray(arr) + + # Explicit access (preferred) + ta.warp # -> wp.array, shape (100,), dtype vec3f + ta.torch # -> torch.Tensor, shape (100, 3) + + # Deprecation bridge (warns once, then silent) + result = ta + 1.0 # works, emits DeprecationWarning + """ + + _deprecation_warned: ClassVar[bool] = False + """Class-level flag ensuring the deprecation warning is emitted at most once.""" + + def __init__(self, wp_array: wp.array) -> None: + """Initialize the ProxyArray wrapper. + + The instance is immutable after construction: the wrapped ``wp.array`` cannot + be reassigned. If the underlying simulation memory is re-allocated, construct + a new :class:`ProxyArray` instead of mutating an existing one. + + Args: + wp_array: The warp array to wrap. + + Raises: + TypeError: If ``wp_array`` is not a :class:`warp.array`. + """ + if not isinstance(wp_array, wp.array): + raise TypeError( + f"ProxyArray expects a warp.array, got {type(wp_array).__name__}." + " If you have a ProxyArray, use it directly instead of wrapping it again." + ) + # Bypass __setattr__ for the two internal fields — everything else raises. + object.__setattr__(self, "_warp", wp_array) + object.__setattr__(self, "_torch_cache", None) + # Cached once at construction so the .torch read path stays a constant-time + # check; only used when the WARN_ON_TORCH_QUATF_ACCESS env var is set. + object.__setattr__(self, "_is_quatf", wp_array.dtype is wp.quatf) + + def __setattr__(self, name: str, value) -> None: + """Forbid mutation of ProxyArray instances except for the internal torch cache. + + The torch view is populated lazily on first ``.torch`` access; that is the + only allowed post-init state change. Every other write raises + :class:`AttributeError` so callers don't accidentally re-point the wrapper. + """ + if name == "_torch_cache": + object.__setattr__(self, name, value) + return + raise AttributeError( + f"ProxyArray is immutable; cannot set attribute {name!r}." + " Construct a new ProxyArray instead of mutating an existing one." + ) + + @staticmethod + def _quatf_access_warning_enabled() -> bool: + """Return ``True`` when the ``WARN_ON_TORCH_QUATF_ACCESS`` env var is set to ``"1"``. + + Read on every :attr:`torch` access to keep the flag dynamic — a single + ``os.environ`` lookup is cheap relative to the warp/torch interop work + that follows. + """ + return os.environ.get(_QUATF_ACCESS_WARN_ENV, "0") == "1" + + # ------------------------------------------------------------------ + # Core accessors + # ------------------------------------------------------------------ + + @property + def warp(self) -> wp.array: + """The underlying warp array.""" + return self._warp + + @property + def torch(self) -> torch.Tensor: + """A cached, zero-copy :class:`torch.Tensor` view of the warp array. + + The tensor is created on first access via :func:`warp.to_torch` and cached + for subsequent calls. Since this is a zero-copy view, modifications to the + tensor are visible through the warp array and vice versa. + + When the underlying warp array has dtype ``wp.quatf`` and the + ``WARN_ON_TORCH_QUATF_ACCESS`` environment variable is set to ``"1"``, + each read emits a :class:`UserWarning` pointing at the call site. This + is a runtime aid for migrating Isaac Lab 2.x code (which used the + ``(w, x, y, z)`` quaternion convention) to Isaac Lab 3.x's + ``(x, y, z, w)`` convention. + """ + if self._is_quatf and self._quatf_access_warning_enabled(): + warnings.warn( + "Reading .torch on a wp.quatf-typed ProxyArray. The Isaac Lab" + " quaternion convention changed from (w, x, y, z) in 2.x to" + " (x, y, z, w) in 3.x. If your code assumes the old order," + " this is likely the source of incorrect rotations." + f" Unset {_QUATF_ACCESS_WARN_ENV} to silence this warning.", + UserWarning, + stacklevel=2, + ) + if self._torch_cache is None: + self._torch_cache = wp.to_torch(self._warp) + return self._torch_cache + + # ------------------------------------------------------------------ + # Convenience properties + # ------------------------------------------------------------------ + + @property + def shape(self) -> tuple[int, ...]: + """Shape of the underlying warp array.""" + return self._warp.shape + + @property + def dtype(self): + """Warp dtype of the underlying array.""" + return self._warp.dtype + + @property + def device(self) -> str: + """Device string of the underlying warp array.""" + return self._warp.device + + def __len__(self) -> int: + """Return the size of the first dimension.""" + return self._warp.shape[0] + + def __repr__(self) -> str: + """Return a string representation of the ProxyArray.""" + return f"ProxyArray(shape={self.shape}, dtype={self.dtype}, device={self.device})" + + # ------------------------------------------------------------------ + # Warp kernel interop + # ------------------------------------------------------------------ + + @property + def __cuda_array_interface__(self): + """Delegate the CUDA array interface to the underlying warp array. + + This allows a ``ProxyArray`` to be passed directly as an argument to + :func:`warp.launch` without explicitly accessing ``.warp``. + + Raises: + AttributeError: If the underlying warp array is not on a CUDA device. + """ + return self._warp.__cuda_array_interface__ + + @property + def __array_interface__(self): + """Delegate the NumPy array interface to the underlying warp array. + + This allows a ``ProxyArray`` to be passed directly as an argument to + :func:`warp.launch` on CPU without explicitly accessing ``.warp``. + + Raises: + AttributeError: If the underlying warp array is not on a CPU device. + """ + return self._warp.__array_interface__ + + # ------------------------------------------------------------------ + # Attribute forwarding (deprecation bridge — delegates to .torch) + # ------------------------------------------------------------------ + + def __getattr__(self, name: str): + """Forward unknown attribute access to the torch view (deprecation bridge). + + Called only when normal attribute lookup fails (i.e. the attribute is not + defined on :class:`ProxyArray` itself), so explicit properties such as + ``shape``, ``dtype``, ``device``, ``warp``, and ``torch`` are unaffected. + + This allows tensor instance methods (``float()``, ``clone()``, ``cpu()``, + ``permute()``, etc.) to be called on a :class:`ProxyArray` without an + explicit ``.torch`` accessor, emitting a one-time :class:`DeprecationWarning`. + """ + self._warn_implicit() + return getattr(self.torch, name) + + # ------------------------------------------------------------------ + # Indexing (deprecation bridge — delegates to .torch) + # ------------------------------------------------------------------ + + def __getitem__(self, key): + """Index into the torch view of this array. + + Supports all torch indexing: ``int``, ``slice``, ``tuple``, + boolean masks, and fancy indexing (multi-dimensional). + """ + self._warn_implicit() + return self.torch[key] + + def __setitem__(self, key, value): + """Write through the torch view into the shared warp memory. + + Supports all torch indexing: ``int``, ``slice``, ``tuple``, + boolean masks, and fancy indexing (multi-dimensional). + """ + self._warn_implicit() + self.torch[key] = value + + # ------------------------------------------------------------------ + # Deprecation bridge + # ------------------------------------------------------------------ + + @classmethod + def _warn_implicit(cls) -> None: + """Emit a one-time deprecation warning for implicit torch usage.""" + if not cls._deprecation_warned: + cls._deprecation_warned = True + warnings.warn( + "Implicit use of ProxyArray as a torch.Tensor is deprecated. " + "Use the explicit .torch property instead (e.g., array.torch).", + DeprecationWarning, + stacklevel=3, + ) + + @classmethod + def __torch_function__(cls, func, types, args=(), kwargs=None): + """Enable torch operations on ProxyArray by unwrapping to ``.torch``. + + This method is called by PyTorch when a torch function receives a + ``ProxyArray`` as an argument. It unwraps all ``ProxyArray`` instances + to their ``.torch`` tensors and delegates to the original function. + """ + if kwargs is None: + kwargs = {} + cls._warn_implicit() + + def unwrap(x): + if isinstance(x, ProxyArray): + return x.torch + if isinstance(x, (list, tuple)): + return type(x)(unwrap(i) for i in x) + return x + + args = unwrap(args) + kwargs = {k: unwrap(v) for k, v in kwargs.items()} + return func(*args, **kwargs) + + # ------------------------------------------------------------------ + # Arithmetic operators + # ------------------------------------------------------------------ + + def _binop(self, other, op: str) -> torch.Tensor: + """Helper for binary and reflected binary operations.""" + self._warn_implicit() + other_val = other.torch if isinstance(other, ProxyArray) else other + return getattr(self.torch, op)(other_val) + + def __add__(self, other) -> torch.Tensor: + return self._binop(other, "__add__") + + def __radd__(self, other) -> torch.Tensor: + return self._binop(other, "__radd__") + + def __sub__(self, other) -> torch.Tensor: + return self._binop(other, "__sub__") + + def __rsub__(self, other) -> torch.Tensor: + return self._binop(other, "__rsub__") + + def __mul__(self, other) -> torch.Tensor: + return self._binop(other, "__mul__") + + def __rmul__(self, other) -> torch.Tensor: + return self._binop(other, "__rmul__") + + def __truediv__(self, other) -> torch.Tensor: + return self._binop(other, "__truediv__") + + def __rtruediv__(self, other) -> torch.Tensor: + return self._binop(other, "__rtruediv__") + + def __pow__(self, other) -> torch.Tensor: + return self._binop(other, "__pow__") + + def __rpow__(self, other) -> torch.Tensor: + return self._binop(other, "__rpow__") + + def __neg__(self) -> torch.Tensor: + self._warn_implicit() + return -self.torch + + def __pos__(self) -> torch.Tensor: + self._warn_implicit() + return +self.torch + + def __abs__(self) -> torch.Tensor: + self._warn_implicit() + return abs(self.torch) + + # ------------------------------------------------------------------ + # Comparison operators + # ------------------------------------------------------------------ + + def __eq__(self, other) -> torch.Tensor: + return self._binop(other, "__eq__") + + def __ne__(self, other) -> torch.Tensor: + return self._binop(other, "__ne__") + + def __lt__(self, other) -> torch.Tensor: + return self._binop(other, "__lt__") + + def __le__(self, other) -> torch.Tensor: + return self._binop(other, "__le__") + + def __gt__(self, other) -> torch.Tensor: + return self._binop(other, "__gt__") + + def __ge__(self, other) -> torch.Tensor: + return self._binop(other, "__ge__") diff --git a/source/isaaclab/isaaclab/utils/warp/utils.py b/source/isaaclab/isaaclab/utils/warp/utils.py new file mode 100644 index 000000000000..2df288dfba2c --- /dev/null +++ b/source/isaaclab/isaaclab/utils/warp/utils.py @@ -0,0 +1,151 @@ +# 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 functools +from collections.abc import Sequence + +import torch +import warp as wp + +## +# Mask resolution - ids/mask to warp boolean mask. +## + + +@wp.kernel +def _populate_mask_from_ids( + mask: wp.array(dtype=wp.bool), + ids: wp.array(dtype=wp.int32), +): + i = wp.tid() + mask[ids[i]] = True + + +def resolve_1d_mask( + *, + ids: Sequence[int] | slice | torch.Tensor | wp.array | None = None, + mask: wp.array | torch.Tensor | None = None, + all_mask: wp.array, + scratch_mask: wp.array, + device: str, +) -> wp.array: + """Resolve ids/mask into a warp boolean mask. + + Callers provide pre-allocated ``all_mask`` (all-True) and ``scratch_mask`` (reusable + work buffer) so this function never allocates. + + Args: + ids: Index ids. Accepts ``Sequence[int]``, ``slice``, ``torch.Tensor``, + ``wp.array(dtype=wp.int32)``, or ``None`` (all elements). + mask: Direct boolean mask. ``wp.array`` is returned as-is; + ``torch.Tensor`` is converted. + all_mask: Pre-allocated all-True mask returned when both *ids* and *mask* + are ``None``. + scratch_mask: Pre-allocated scratch buffer populated in-place when *ids* + are provided. Not re-entrant (shared buffer). + device: Warp device string (e.g. ``"cuda:0"``). + + Returns: + A ``wp.array(dtype=wp.bool)`` mask. + """ + # Normalize slice(None) to None so the capture guard treats it identically to ids=None. + if isinstance(ids, slice) and ids == slice(None): + ids = None + + if wp.get_device().is_capturing: + if ids is not None or (mask is not None and not isinstance(mask, wp.array)): + raise RuntimeError( + "resolve_1d_mask is only capturable when mask is a wp.array or both ids and mask are None." + ) + + # --- Direct mask input --- + if mask is not None: + if isinstance(mask, wp.array): + return mask + if isinstance(mask, torch.Tensor): + if mask.dtype != torch.bool: + mask = mask.to(dtype=torch.bool) + if str(mask.device) != device: + mask = mask.to(device) + return wp.from_torch(mask, dtype=wp.bool) + raise TypeError(f"Unsupported mask type: {type(mask)}") + + # --- Fast path: all elements --- + if ids is None: + return all_mask + + # --- Normalize slice to list --- + if isinstance(ids, slice): + start, stop, step = ids.indices(scratch_mask.shape[0]) + ids = list(range(start, stop, step)) + + # --- Normalize to concrete type --- + if not isinstance(ids, (torch.Tensor, wp.array)): + ids = list(ids) + + # --- Populate scratch mask --- + scratch_mask.fill_(False) + + if isinstance(ids, torch.Tensor): + if ids.numel() == 0: + return scratch_mask + if str(ids.device) != device: + ids = ids.to(device) + if ids.dtype != torch.int32: + ids = ids.to(dtype=torch.int32) + if not ids.is_contiguous(): + ids = ids.contiguous() + ids_wp = wp.from_torch(ids, dtype=wp.int32) + elif isinstance(ids, wp.array): + if ids.shape[0] == 0: + return scratch_mask + if ids.dtype != wp.int32: + raise TypeError(f"Unsupported wp.array dtype for ids: {ids.dtype}. Expected wp.int32 index array.") + ids_wp = ids + else: + if len(ids) == 0: + return scratch_mask + ids_wp = wp.array(ids, dtype=wp.int32, device=device) + + wp.launch(_populate_mask_from_ids, dim=ids_wp.shape[0], inputs=[scratch_mask, ids_wp], device=device) + return scratch_mask + + +## +# Capture safety — property guard. +## + + +def capture_unsafe(reason: str | None = None): + """Mark a callable as not CUDA-graph-capture-safe. + + Raises ``RuntimeError`` if the decorated callable is invoked while + ``wp.get_device().is_capturing`` is ``True``. + + Args: + reason: Optional explanation appended to the error message. + + Usage:: + + @property + @capture_unsafe("Relies on a Python timestamp guard.") + def projected_gravity_b(self) -> wp.array: ... + """ + + def decorator(func): + @functools.wraps(func) + def wrapper(*args, **kwargs): + if wp.get_device().is_capturing: + msg = f"'{func.__qualname__}' cannot be called during CUDA graph capture." + if reason: + msg = f"{msg} {reason}" + raise RuntimeError(msg) + return func(*args, **kwargs) + + return wrapper + + return decorator diff --git a/source/isaaclab/isaaclab/utils/warp/warp_math.py b/source/isaaclab/isaaclab/utils/warp/warp_math.py new file mode 100644 index 000000000000..52d4e708f858 --- /dev/null +++ b/source/isaaclab/isaaclab/utils/warp/warp_math.py @@ -0,0 +1,186 @@ +# 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 helpers for camera-related math operations. + +These replace equivalent torch functions on the per-frame hot path, operating +directly on warp arrays without torch round-trips. +""" + +from __future__ import annotations + +from typing import Literal + +import warp as wp + +# Camera orientation convention conversion +# +# Every pair of (origin, target) conventions is equivalent to a single +# right-multiplication by a constant unit quaternion: +# +# q_out[i] = q_in[i] * q_const +# +# Derivations (xyzw): +# opengl ↔ ros : 180° around X → (1, 0, 0, 0) (self-inverse) +# world → opengl : Rx(+90°)·Ry(−90°) → (0.5, −0.5, −0.5, 0.5) +# opengl → world : inverse of above → (−0.5, 0.5, 0.5, 0.5) +# ros → world : compose ros→gl→world → (0.5, −0.5, 0.5, 0.5) +# world → ros : inverse of above → (−0.5, 0.5, −0.5, 0.5) + +_CAMERA_ORIENTATION_CONST: dict[tuple[str, str], wp.quatf] = { + ("opengl", "ros"): wp.quatf(1.0, 0.0, 0.0, 0.0), + ("ros", "opengl"): wp.quatf(1.0, 0.0, 0.0, 0.0), + ("world", "opengl"): wp.quatf(0.5, -0.5, -0.5, 0.5), + ("opengl", "world"): wp.quatf(-0.5, 0.5, 0.5, 0.5), + ("ros", "world"): wp.quatf(0.5, -0.5, 0.5, 0.5), + ("world", "ros"): wp.quatf(-0.5, 0.5, -0.5, 0.5), +} + + +# TODO: Optimize these kernels with tiled ops and use wp.static +@wp.kernel +def _convert_camera_orientation_all_kernel( + src: wp.array(dtype=wp.quatf), + dst: wp.array(dtype=wp.quatf), + q_const: wp.quatf, +): + """Apply constant-quaternion convention conversion to every element.""" + i = wp.tid() + dst[i] = src[i] * q_const + + +@wp.kernel +def _convert_camera_orientation_indexed_kernel( + src: wp.array(dtype=wp.quatf), + dst: wp.array(dtype=wp.quatf), + indices: wp.array(dtype=wp.int32), + q_const: wp.quatf, +): + """Apply constant-quaternion convention conversion to indexed elements. + + Reads ``src[i]`` and writes to ``dst[indices[i]]``. Use this for partial + camera updates (e.g. environment resets targeting a subset of cameras). + """ + i = wp.tid() + dst[indices[i]] = src[i] * q_const + + +def convert_camera_frame_orientation_convention_wp( + src: wp.array, + dst: wp.array, + origin: Literal["opengl", "ros", "world"], + target: Literal["opengl", "ros", "world"], + indices: wp.array | None = None, + device: str | None = None, +) -> None: + """Convert camera-frame quaternion orientations between conventions using a warp kernel. + + Replaces :func:`~isaaclab.utils.math.convert_camera_frame_orientation_convention` on + the per-frame hot path. All six convention pairs collapse to a single quaternion + right-multiplication by a pre-computed constant — no matrix round-trip, no torch. + + The operation is **in-place on** ``dst``: + - Without ``indices``: ``dst[i] = src[i] * q_const`` for all i. + - With ``indices``: ``dst[indices[i]] = src[i] * q_const`` for each i. + + Args: + src: Source quaternions ``(x, y, z, w)``. Shape ``(N,)``, dtype ``wp.quatf``. + dst: Destination quaternion array to write into. Shape ``(M,)``, dtype ``wp.quatf``. + ``M >= N`` when ``indices`` is provided; ``M == N`` otherwise. + origin: Source convention (``"opengl"``, ``"ros"``, or ``"world"``). + target: Target convention (``"opengl"``, ``"ros"``, or ``"world"``). + indices: Optional warp int32 array of shape ``(N,)`` selecting which slots of + ``dst`` to write. If ``None`` all N elements are written sequentially. + device: Warp device string. Defaults to ``src.device``. + """ + if origin == target: + if indices is None: + wp.copy(dst, src) + else: + # scatter copy: dst[indices[i]] = src[i] + wp.launch( + _convert_camera_orientation_indexed_kernel, + dim=indices.shape[0], + inputs=[src, dst, indices, wp.quatf(0.0, 0.0, 0.0, 1.0)], + device=device or src.device, + ) + return + + q_const = _CAMERA_ORIENTATION_CONST[(origin, target)] + dev = device or src.device + + if indices is None: + wp.launch( + _convert_camera_orientation_all_kernel, + dim=src.shape[0], + inputs=[src, dst, q_const], + device=dev, + ) + else: + wp.launch( + _convert_camera_orientation_indexed_kernel, + dim=indices.shape[0], + inputs=[src, dst, indices, q_const], + device=dev, + ) + + +@wp.kernel +def _clamp_depth_to_inf_kernel( + buf: wp.array(dtype=wp.float32, ndim=4), + max_range: float, +): + """Replace values above ``max_range`` with ``+inf``.""" + n, h, w, c = wp.tid() + v = buf[n, h, w, c] + if v > max_range: + buf[n, h, w, c] = wp.inf + + +@wp.kernel +def _replace_inf_kernel( + buf: wp.array(dtype=wp.float32, ndim=4), + replacement: float, +): + """Replace ``+inf`` values with ``replacement``.""" + n, h, w, c = wp.tid() + if wp.isinf(buf[n, h, w, c]): + buf[n, h, w, c] = replacement + + +def clamp_depth_to_inf_wp(buf: wp.array, max_range: float, device: str | None = None) -> None: + """Replace depth values above ``max_range`` with ``+inf`` using a warp kernel. + + Replaces ``t[t > max_range] = torch.inf`` on the hot path. + + Args: + buf: Depth buffer. Shape ``(N, H, W, C)``, dtype ``wp.float32``. + max_range: Depth values strictly greater than this are set to ``+inf``. + device: Warp device string. Defaults to ``buf.device``. + """ + wp.launch( + _clamp_depth_to_inf_kernel, + dim=buf.shape, + inputs=[buf, float(max_range)], + device=device or buf.device, + ) + + +def replace_inf_depth_wp(buf: wp.array, replacement: float, device: str | None = None) -> None: + """Replace ``+inf`` depth values with ``replacement`` using a warp kernel. + + Replaces ``t[torch.isinf(t)] = value`` on the hot path. + + Args: + buf: Depth buffer. Shape ``(N, H, W, C)``, dtype ``wp.float32``. + replacement: Value to write where ``+inf`` was found (e.g. ``0.0`` or ``max_range``). + device: Warp device string. Defaults to ``buf.device``. + """ + wp.launch( + _replace_inf_kernel, + dim=buf.shape, + inputs=[buf, float(replacement)], + device=device or buf.device, + ) diff --git a/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/deploy/cable_insertion/config/displayport_rizon_4s/joint_pos_env_cfg.py b/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/deploy/cable_insertion/config/displayport_rizon_4s/joint_pos_env_cfg.py index f2d509660719..293c05743494 100644 --- a/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/deploy/cable_insertion/config/displayport_rizon_4s/joint_pos_env_cfg.py +++ b/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/deploy/cable_insertion/config/displayport_rizon_4s/joint_pos_env_cfg.py @@ -11,6 +11,7 @@ """ import math +import os import torch @@ -63,7 +64,7 @@ # switches to a specific combination; see the experiment matrix / commit log. # The block between the START/END markers is what the experiment commits edit. # --- EXP TOGGLES START --- -EXP_SYSID = True # enable sim2real (sysid) action model + PhysX SysID gains +EXP_SYSID = os.getenv("DP_CABLE_EXP_SYSID", "1").lower() not in {"0", "false", "no", "off"} EXP_SOCKET_POS_RANGE = [0.01, 0.01, 0.02] # socket position randomization, +/- m per axis [x, y, z] EXP_SOCKET_ORN_DEG = 2.0 # socket orientation randomization, +/- deg on roll/pitch/yaw EXP_CURRICULUM = "anneal_80_0_500" # disabled|fixed80|anneal_80_0_1000|anneal_80_20_1000|anneal_80_20_500|anneal_80_0_500 @@ -483,7 +484,7 @@ def __post_init__(self): command_acceleration_limit=USE_SIM2REAL_COMMAND_ACCELERATION_LIMIT, ) else: - self.actions.arm_action = mdp.RelativeJointPositionActionCfg( + self.actions.arm_action = mdp.DeployRelativeJointPositionActionCfg( asset_name="robot", joint_names=_arm_joint_names, scale=self.joint_action_scale, diff --git a/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/deploy/mdp/__init__.pyi b/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/deploy/mdp/__init__.pyi index ccfb35c62f19..0fbfc330f8e5 100644 --- a/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/deploy/mdp/__init__.pyi +++ b/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/deploy/mdp/__init__.pyi @@ -11,6 +11,8 @@ __all__ = [ "reset_plug_at_goal_curriculum", "ResetSampledConstantNoiseModel", "ResetSampledConstantNoiseModelCfg", + "joint_pos", + "joint_vel", "gear_pos_w", "gear_quat_w", "gear_shaft_pos_w", @@ -38,8 +40,12 @@ __all__ = [ "ShapedDelayedRelativeJointPositionActionCfg", "FlexivDynamicsAwareRelativeJointPositionAction", "FlexivDynamicsAwareRelativeJointPositionActionCfg", + "DeployRelativeJointPositionAction", + "DeployRelativeJointPositionActionCfg", ] +from .actions import DeployRelativeJointPositionAction +from .actions_cfg import DeployRelativeJointPositionActionCfg from .delayed_joint_actions import ( DelayedRelativeJointPositionAction, ShapedDelayedRelativeJointPositionAction, @@ -59,6 +65,8 @@ from .events import ( ) from .noise_models import ResetSampledConstantNoiseModel, ResetSampledConstantNoiseModelCfg from .observations import ( + joint_pos, + joint_vel, gear_pos_w, gear_quat_w, gear_shaft_pos_w, diff --git a/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/deploy/mdp/actions.py b/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/deploy/mdp/actions.py new file mode 100644 index 000000000000..de98008a598d --- /dev/null +++ b/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/deploy/mdp/actions.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 + +"""Deploy-specific action terms.""" + +from __future__ import annotations + +from typing import TYPE_CHECKING + +import warp as wp + +from isaaclab.envs.mdp.actions.joint_actions import RelativeJointPositionAction + +if TYPE_CHECKING: + from isaaclab.envs import ManagerBasedEnv + + from .actions_cfg import DeployRelativeJointPositionActionCfg + + +_LEAPP_TRACED_OBSERVATION_INPUTS = "_leapp_traced_observation_inputs" + + +def _real_env(env): + return object.__getattribute__(env, "_real_env") if type(env).__name__ == "_EnvProxy" else env + + +def _to_torch(data): + return data.torch if hasattr(data, "torch") else wp.to_torch(data) + + +class DeployRelativeJointPositionAction(RelativeJointPositionAction): + """Relative joint action that exports absolute joint targets for deploy. + + During normal Isaac Lab execution this behaves like + :class:`RelativeJointPositionAction`. During LEAPP export it reuses the + traced policy joint-position observation so the exported graph computes: + + ``target_joint_pos = robot_joint_pos + raw_action * scale``. + """ + + cfg: DeployRelativeJointPositionActionCfg + + def __init__(self, cfg: DeployRelativeJointPositionActionCfg, env: ManagerBasedEnv): + super().__init__(cfg, env) + + def apply_actions(self): + current_joint_pos = None + if type(self._asset).__name__ == "_ArticulationWriteProxy": + traced_inputs = getattr(_real_env(self._env), _LEAPP_TRACED_OBSERVATION_INPUTS, {}) + current_joint_pos = traced_inputs.get(f"{self.cfg.asset_name}_joint_pos") + + if current_joint_pos is None: + current_joint_pos = _to_torch(self._asset.data.joint_pos)[:, self._joint_ids] + + current_actions = self.processed_actions + current_joint_pos + self._asset.set_joint_position_target_index(target=current_actions, joint_ids=self._joint_ids) diff --git a/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/deploy/mdp/actions_cfg.py b/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/deploy/mdp/actions_cfg.py new file mode 100644 index 000000000000..e7ad3d279434 --- /dev/null +++ b/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/deploy/mdp/actions_cfg.py @@ -0,0 +1,18 @@ +# 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 + +"""Deploy-specific action configuration classes.""" + +from __future__ import annotations + +from isaaclab.envs.mdp.actions.actions_cfg import RelativeJointPositionActionCfg +from isaaclab.utils import configclass + + +@configclass +class DeployRelativeJointPositionActionCfg(RelativeJointPositionActionCfg): + """Relative joint-position action that exports absolute joint targets.""" + + class_type: type | str = "{DIR}.actions:DeployRelativeJointPositionAction" diff --git a/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/deploy/mdp/observations.py b/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/deploy/mdp/observations.py index fbe42cf647f4..edab9e7bdff7 100644 --- a/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/deploy/mdp/observations.py +++ b/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/deploy/mdp/observations.py @@ -13,6 +13,7 @@ import warp as wp from isaaclab.managers import ManagerTermBase, ObservationTermCfg, SceneEntityCfg +from isaaclab.utils.leapp import QUAT_XYZW_ELEMENT_NAMES, XYZ_ELEMENT_NAMES, InputKindEnum from isaaclab.utils.math import combine_frame_transforms, matrix_from_quat if TYPE_CHECKING: @@ -22,6 +23,103 @@ from .events import randomize_gear_type +_LEAPP_TRACED_OBSERVATION_INPUTS = "_leapp_traced_observation_inputs" + + +def _is_leapp_export_env(env) -> bool: + return type(env).__name__ == "_EnvProxy" + + +def _real_env(env): + return object.__getattribute__(env, "_real_env") if _is_leapp_export_env(env) else env + + +def _to_torch(data): + return data.torch if hasattr(data, "torch") else wp.to_torch(data) + + +def _selected_joint_names(asset, joint_ids) -> list[str] | None: + joint_names = getattr(asset, "joint_names", None) + if joint_names is None: + return None + if joint_ids is None or joint_ids == slice(None): + return list(joint_names) + if isinstance(joint_ids, slice): + return list(joint_names[joint_ids]) + if hasattr(joint_ids, "tolist"): + joint_ids = joint_ids.tolist() + return [joint_names[int(joint_id)] for joint_id in joint_ids] + + +def _leapp_input(env, *, name: str, tensor: torch.Tensor, kind: str, element_names=None, connection: str): + if not _is_leapp_export_env(env): + return tensor + + from leapp import annotate + from leapp.utils.tensor_description import TensorSemantics + + return annotate.input_tensors( + env.unwrapped.spec.id, + TensorSemantics( + name=name, + ref=tensor, + kind=kind, + element_names=element_names, + extra={"isaaclab_connection": connection}, + ), + ) + + +def _set_traced_observation_input(env, name: str, tensor: torch.Tensor) -> None: + if not _is_leapp_export_env(env): + return + real_env = _real_env(env) + traced_inputs = getattr(real_env, _LEAPP_TRACED_OBSERVATION_INPUTS, None) + if traced_inputs is None: + traced_inputs = {} + setattr(real_env, _LEAPP_TRACED_OBSERVATION_INPUTS, traced_inputs) + traced_inputs[name] = tensor + + +def _deploy_object_input_base_name(asset_name: str) -> str: + for prefix in ("dp_", "gb300_", "factory_"): + if asset_name.startswith(prefix): + return asset_name[len(prefix) :] + return asset_name + + +def joint_pos(env: ManagerBasedRLEnv, asset_cfg: SceneEntityCfg = SceneEntityCfg("robot")) -> torch.Tensor: + """Joint positions for the configured joints, exposed as the LEAPP input boundary.""" + real_env = _real_env(env) + asset = real_env.scene[asset_cfg.name] + selected_joint_pos = _to_torch(asset.data.joint_pos)[:, asset_cfg.joint_ids] + selected_joint_pos = _leapp_input( + env, + name=f"{asset_cfg.name}_joint_pos", + tensor=selected_joint_pos, + kind=InputKindEnum.JOINT_POSITION, + element_names=_selected_joint_names(asset, asset_cfg.joint_ids), + connection=f"state:{asset_cfg.name}:joint_pos", + ) + _set_traced_observation_input(env, f"{asset_cfg.name}_joint_pos", selected_joint_pos) + return selected_joint_pos + + +def joint_vel(env: ManagerBasedRLEnv, asset_cfg: SceneEntityCfg = SceneEntityCfg("robot")) -> torch.Tensor: + """Joint velocities for the configured joints, exposed as the LEAPP input boundary.""" + real_env = _real_env(env) + asset = real_env.scene[asset_cfg.name] + selected_joint_vel = _to_torch(asset.data.joint_vel)[:, asset_cfg.joint_ids] + return _leapp_input( + env, + name=f"{asset_cfg.name}_joint_vel", + tensor=selected_joint_vel, + kind=InputKindEnum.JOINT_VELOCITY, + element_names=_selected_joint_names(asset, asset_cfg.joint_ids), + connection=f"state:{asset_cfg.name}:joint_vel", + ) + + class gear_shaft_pos_w(ManagerTermBase): """Gear shaft position in world frame with offset applied. @@ -382,14 +480,25 @@ def __call__( asset_cfg: SceneEntityCfg | None = None, offset: list | None = None, ) -> torch.Tensor: - obj_pos = wp.to_torch(self.asset.data.root_pos_w) - obj_quat = wp.to_torch(self.asset.data.root_quat_w) + real_env = _real_env(env) + asset = real_env.scene[self.asset_cfg.name] + obj_pos = _to_torch(asset.data.root_pos_w) + obj_quat = _to_torch(asset.data.root_quat_w) if torch.any(self.offset_tensor != 0): - offset_repeated = self.offset_tensor.unsqueeze(0).repeat(env.num_envs, 1) + offset_repeated = self.offset_tensor.unsqueeze(0).repeat(real_env.num_envs, 1) obj_pos, _ = combine_frame_transforms(obj_pos, obj_quat, offset_repeated, self.identity_quat) - return obj_pos - env.scene.env_origins + obj_pos = obj_pos - real_env.scene.env_origins + input_name = f"{_deploy_object_input_base_name(self.asset_cfg.name)}_pos" + return _leapp_input( + env, + name=input_name, + tensor=obj_pos, + kind=InputKindEnum.BODY_POSITION, + element_names=XYZ_ELEMENT_NAMES, + connection=f"observation:policy:{input_name}", + ) class rigid_object_quat_w(ManagerTermBase): @@ -420,13 +529,20 @@ def __call__( env: ManagerBasedRLEnv, asset_cfg: SceneEntityCfg | None = None, ) -> torch.Tensor: - obj_quat = wp.to_torch(self.asset.data.root_quat_w) - - w_negative = obj_quat[:, 3] < 0 - positive_quat = obj_quat.clone() - positive_quat[w_negative] = -obj_quat[w_negative] + real_env = _real_env(env) + obj_quat = _to_torch(real_env.scene[self.asset_cfg.name].data.root_quat_w) + input_name = f"{_deploy_object_input_base_name(self.asset_cfg.name)}_quat" + obj_quat = _leapp_input( + env, + name=input_name, + tensor=obj_quat, + kind=InputKindEnum.BODY_ROTATION, + element_names=QUAT_XYZW_ELEMENT_NAMES, + connection=f"observation:policy:{input_name}", + ) - return positive_quat + w_negative = obj_quat[:, 3:4] < 0 + return torch.where(w_negative, -obj_quat, obj_quat) def _quat_to_rot_6d(quat: torch.Tensor) -> torch.Tensor: diff --git a/source/isaaclab_tasks/isaaclab_tasks/utils/__init__.pyi b/source/isaaclab_tasks/isaaclab_tasks/utils/__init__.pyi index e960be9e290e..bfb7af42f762 100644 --- a/source/isaaclab_tasks/isaaclab_tasks/utils/__init__.pyi +++ b/source/isaaclab_tasks/isaaclab_tasks/utils/__init__.pyi @@ -12,14 +12,11 @@ __all__ = [ "preset", "resolve_task_config", "hydra_task_config", - "resolve_preset_defaults", - "add_launcher_args", - "launch_simulation", - "compute_kit_requirements", + "resolve_presets", + "setup_preset_cli", ] -from .hydra import PresetCfg, preset, hydra_task_config, resolve_task_config +from .hydra import PresetCfg, preset, hydra_task_config, resolve_task_config, resolve_presets from .importer import import_packages from .parse_cfg import get_checkpoint_path, load_cfg_from_registry, parse_env_cfg -from .hydra import resolve_task_config, hydra_task_config, resolve_preset_defaults -from .sim_launcher import add_launcher_args, launch_simulation, compute_kit_requirements +from .preset_cli import setup_preset_cli diff --git a/source/isaaclab_tasks/isaaclab_tasks/utils/preset_cli.py b/source/isaaclab_tasks/isaaclab_tasks/utils/preset_cli.py new file mode 100644 index 000000000000..312cb37a43f5 --- /dev/null +++ b/source/isaaclab_tasks/isaaclab_tasks/utils/preset_cli.py @@ -0,0 +1,321 @@ +# 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 + +"""Typed-preset selection via Hydra-style CLI tokens. + +Recognizes three ``key=value`` tokens (no leading dashes) on ``sys.argv``: + +* ``physics=NAME`` -- typed selector for ``PhysicsCfg`` variants. +* ``renderer=NAME`` -- typed selector for ``RendererCfg`` variants. +* ``presets=NAME[,NAME,...]`` -- broadcast applied to every matching ``PresetCfg``. + +:func:`setup_preset_cli` registers the preset-selection help description on the +parser and runs ``parse_known_args``, returning the verbatim remainder. The +tokens above are passed through unchanged; hydra's +:func:`~isaaclab_tasks.utils.hydra.register_task` parses them directly (applying +the names as presets and enforcing that ``physics=``/``renderer=`` resolve +against a config of that type). Callers simply assign the remainder to +``sys.argv``; no rewriting step is needed. + +No argparse arguments are registered for the typed selectors -- discoverability +lives in the ``argument_group`` description, so the parsed Namespace gains no +preset attributes and cannot shadow :class:`~isaaclab.app.AppLauncher` +SimulationApp config keys (``renderer`` notably). + +Typical script setup:: + + parser = argparse.ArgumentParser(...) + # ... script-specific args ... + add_launcher_args(parser) + args_cli, remaining = setup_preset_cli(parser) + sys.argv = [sys.argv[0]] + remaining + +Scripts that intersect the remainder with external-callback output (e.g. +``rsl_rl`` scripts' ``--external_callback`` hook) do the intersection on the +remainder before assigning ``sys.argv`` -- both sides share the same token +vocabulary:: + + args_cli, remaining = setup_preset_cli(parser) + if args_cli.external_callback: + remaining = list_intersection(remaining, external_callback_function()) + sys.argv = [sys.argv[0]] + remaining + +``setup_preset_cli`` does NOT add AppLauncher flags itself -- callers add them +explicitly via :func:`isaaclab.app.add_launcher_args` before calling. +""" + +from __future__ import annotations + +import argparse +import sys + +from .preset_target import PresetTarget + +# ============================================================================ +# Public entry point +# ============================================================================ + + +def setup_preset_cli( + parser: argparse.ArgumentParser, argv: list[str] | None = None +) -> tuple[argparse.Namespace, list[str]]: + """Register the preset-selection help description and parse argv. + + Must be called *after* AppLauncher flags and script-specific arguments are + registered on ``parser`` -- otherwise those unknown tokens land in + ``parse_known_args``'s remainder. + + The returned remainder contains the user-typed ``physics=`` / ``renderer=`` + / ``presets=`` tokens verbatim, alongside any Hydra path overrides and any + unknown argparse flags, ready to assign to ``sys.argv`` for hydra to parse. + + Does not mutate ``sys.argv``; the caller assigns + ``sys.argv = [sys.argv[0]] + remaining`` when ready, so any argv-aware logic + that re-reads ``sys.argv`` (e.g. an external callback) runs against the + user's original command line first. + + Args: + parser: Caller's argument parser. An ``argument_group`` is attached + for help-time variant discovery; no ``add_argument`` calls are + made, so the Namespace gains no preset attributes. + argv: Optional argument list to parse. When ``None`` (default), + ``parse_known_args`` reads from ``sys.argv``. Provided primarily + for in-process test paths that drive the parser with a synthetic + argv. Help-time variant enumeration always reads ``sys.argv`` -- + the user's interactive command line is the only argv that + triggers ``--help`` rendering. + + Returns: + ``(args, remaining)`` where ``remaining`` is the verbatim output of + ``parser.parse_known_args(argv)``, ready to hand to Hydra via + ``sys.argv``. + """ + # --help short-circuits parsing, so help text that depends on --task has to + # find it before argparse runs. Gate the env_cfg load on --help to keep + # normal training runs cheap. + argv_helper = _ArgvHelper(sys.argv) + actual_variants = ( + _enumerate_variants(argv_helper.task_name) if (argv_helper.task_name and argv_helper.help_requested) else None + ) + + # Argparse's default HelpFormatter reflows description text into one wrapped + # paragraph, which would collapse the per-variant bullets we emit. Use a + # formatter that wraps each blank-line-separated paragraph independently + # while preserving explicit newlines. Respect a caller-set custom formatter. + if parser.formatter_class is argparse.HelpFormatter: + parser.formatter_class = _PresetHelpFormatter + + # Help-only group: no add_argument() calls means no preset attributes on + # the Namespace, so AppLauncher can't accidentally forward one (notably + # ``renderer``) into SimulationApp config. + parser.add_argument_group("preset selection", description=_DescriptionBuilder.build(actual_variants)) + + return parser.parse_known_args(argv) + + +# ============================================================================ +# Public preset enumeration (for tooling, e.g. list_envs) +# ============================================================================ + + +def enumerate_task_presets(task_name: str) -> dict[PresetTarget, list[str]] | None: + """Return the available preset names for *task_name*, bucketed by selector type. + + Loads the env config registered under *task_name* and walks its preset tree + using the same logic that the CLI help-text renderer uses, so the returned + view matches what ``--task= --help`` shows at the command line. + + This function is safe to call after :class:`~isaaclab.app.AppLauncher` has + booted (i.e. inside a running Isaac Sim session). + + Args: + task_name: Gymnasium task ID (e.g. ``"Isaac-Cartpole"``). + + Returns: + A mapping ``{PresetTarget: sorted list of preset names}`` on success. + Returns ``None`` if the env config cannot be loaded (import error, + missing registration, etc.). The ``"default"`` fallback is excluded + from every list because it is implicit, not a user-selectable name. + """ + try: + result = _enumerate_variants(task_name) + return {target: sorted(names) for target, names in result.items()} + except Exception: + return None + + +# ============================================================================ +# Help-text rendering +# ============================================================================ + + +class _PresetHelpFormatter(argparse.HelpFormatter): + """Argparse help formatter that wraps each paragraph separately. + + Default :class:`argparse.HelpFormatter` reflows the entire description into + one paragraph, merging the variant listing into the surrounding prose, and + collapses ``\\n``-separated bullets onto one line. + :class:`~argparse.RawDescriptionHelpFormatter` preserves description + newlines but drops wrapping entirely. The ``_fill_text`` override below + splits the description on blank lines and wraps each paragraph indep- + endently, giving both readable paragraphs and per-line bullets. + """ + + def _fill_text(self, text: str, width: int, indent: str) -> str: + import textwrap + + paragraphs = text.split("\n\n") + rendered: list[str] = [] + for paragraph in paragraphs: + # A paragraph that already contains hard newlines (the bulleted + # variant listing) is rendered verbatim; otherwise word-wrap. + if "\n" in paragraph: + rendered.append("\n".join(f"{indent}{line}" for line in paragraph.splitlines())) + else: + rendered.append(textwrap.fill(paragraph, width, initial_indent=indent, subsequent_indent=indent)) + return "\n\n".join(rendered) + + +class _DescriptionBuilder: + """Renders the preset-selection ``argument_group`` description. + + Groups the column constants and per-row formatting that build the + selector table. Iterates :class:`PresetTarget` to produce one row per + selector; each row's syntax and description come from the enum, so + adding a new typed target needs no changes here. + """ + + # Column widths. ``SELECTOR_COL`` = width of the longest selector syntax + # (``presets=NAME[,NAME,...]`` = 23 chars); shorter selectors right-pad + # to this width. ``DESC_GAP`` is the gap between syntax and description. + SELECTOR_COL = 23 + DESC_GAP = 3 + ROW_PREFIX = " " + + INTRO = "Select named PresetCfg alternatives via Hydra-style overrides (key=value, no leading dashes):" + EPILOG = "Hydra also accepts path-targeted overrides like env.sim.physics=NAME." + HINT = "Pass `--task=X` along with `--help` to see preset variants available for that task." + + @classmethod + def build(cls, actual_variants: dict[PresetTarget, set[str]] | None) -> str: + """Build the description text. + + Args: + actual_variants: ``None`` when no ``--task=X --help`` is in argv; + otherwise a ``{target: set[name]}`` bucketed view from + :func:`_enumerate_variants`. + """ + with_available = actual_variants is not None + rows = [ + cls._row(t, with_available=with_available, variants=sorted((actual_variants or {}).get(t, set()))) + for t in PresetTarget + ] + middle = f"{cls.HINT}\n\n" if not with_available else "" + return f"{cls.INTRO}\n" + "\n".join(rows) + f"\n\n{middle}{cls.EPILOG}" + + @classmethod + def _row(cls, target: PresetTarget, *, with_available: bool, variants: list[str]) -> str: + syntax = cls._syntax(target).ljust(cls.SELECTOR_COL) + desc = cls._description(target) + suffix = ". Available:" if with_available else "" + header = f"{cls.ROW_PREFIX}{syntax}{' ' * cls.DESC_GAP}{desc}{suffix}" + if not with_available: + return header + # Bullet indent aligns with the description column once argparse + # prepends its 2-space group-description indent. + bullet_indent = " " * (len(cls.ROW_PREFIX) + cls.SELECTOR_COL + cls.DESC_GAP) + body = "\n".join(f"{bullet_indent}- {n}" for n in variants) if variants else f"{bullet_indent}(none)" + return f"{header}\n{body}" + + @staticmethod + def _syntax(target: PresetTarget) -> str: + """User-facing selector form: ``physics=NAME`` vs ``presets=NAME[,NAME,...]``.""" + if target.base_classes: # typed: single name + return f"{target.value}=NAME" + return f"{target.value}=NAME[,NAME,...]" # DOMAIN: comma-separated broadcast + + @staticmethod + def _description(target: PresetTarget) -> str: + """One-line description; for typed targets includes the cfg base class name.""" + if target.base_classes: + return f"(typed) selects a {target.base_classes[0].__name__} variant" + return "broadcast: applied to every matching PresetCfg" + + +# ============================================================================ +# argv inspection (pre-argparse peek for help-text rendering) +# ============================================================================ + + +class _ArgvHelper: + """Single-pass argv scan that exposes ``task_name`` and ``help_requested``. + + Needed because argparse's ``--help`` short-circuits parsing, so help text + that depends on ``--task`` has to find it before argparse runs. + + Attributes: + task_name: Last ``--task`` value (matching argparse's last-wins + semantics), or ``None`` if absent. + help_requested: ``True`` if ``--help`` or ``-h`` is present. + """ + + def __init__(self, argv: list[str]): + self.task_name: str | None = None + self.help_requested: bool = False + for i in range(1, len(argv)): + token = argv[i] + if token in ("--help", "-h"): + self.help_requested = True + elif token == "--task" and i + 1 < len(argv): + self.task_name = argv[i + 1] + elif token.startswith("--task="): + self.task_name = token[len("--task=") :] + + +# ============================================================================ +# Help-time variant enumeration (load env_cfg, walk, bucket by target) +# ============================================================================ + + +def _enumerate_variants(task_name: str) -> dict[PresetTarget, set[str]]: + """Load env_cfg for *task_name* and bucket its variants by target. + + Uses the same walker hydra's resolver runs so help and resolve see one + view of the cfg tree. The env_cfg load is safe before AppLauncher boots + because ``test_env_cfg_no_forbidden_imports`` blocks Kit-only imports at + the top level of cfg modules. Exceptions from the loader propagate + verbatim -- they surface as the natural error, not a buried help string. + """ + from isaaclab_tasks.utils.hydra import collect_presets + from isaaclab_tasks.utils.parse_cfg import load_cfg_from_registry + + env_cfg = load_cfg_from_registry(task_name, "env_cfg_entry_point") + return _bucket_variants_by_target(collect_presets(env_cfg)) + + +def _bucket_variants_by_target(walked: dict) -> dict[PresetTarget, set[str]]: + """Convert :func:`collect_presets` output into ``{target: set[name]}``. + + Routes each ``(name, cfg)`` by ``isinstance(cfg, target.base_classes)``; + cfgs matching no typed target fall into ``DOMAIN``. The implicit + ``default`` field is filtered -- it's the fallback, not a selectable name. + + Routing by class hierarchy means new backends subclassing + :class:`~isaaclab.physics.PhysicsCfg` / + :class:`~isaaclab.renderers.renderer_cfg.RendererCfg` bucket automatically + regardless of what name the env_cfg gives the field. + """ + typed_targets = [t for t in PresetTarget if t.base_classes] + result: dict[PresetTarget, set[str]] = {target: set() for target in PresetTarget} + for path_dict in walked.values(): + for name, cfg in path_dict.items(): + if name == "default": + continue + matched = next( + (t for t in typed_targets if isinstance(cfg, t.base_classes)), + PresetTarget.DOMAIN, + ) + result[matched].add(name) + return result diff --git a/source/isaaclab_tasks/isaaclab_tasks/utils/preset_target.py b/source/isaaclab_tasks/isaaclab_tasks/utils/preset_target.py new file mode 100644 index 000000000000..89878541cf2e --- /dev/null +++ b/source/isaaclab_tasks/isaaclab_tasks/utils/preset_target.py @@ -0,0 +1,128 @@ +# 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 + +"""Closed enum of typed preset categories with per-target metadata. + +Each :class:`PresetTarget` member carries everything the preset CLI layer +needs to know about that category in one place: + +* ``label`` -- the Hydra-style selector key (e.g. ``"physics"`` for + ``physics=NAME``) and ``self.value``. +* ``base_classes`` -- the cfg base classes whose subclass instances belong to + this bucket. Help-time bucketing in :mod:`isaaclab_tasks.utils.preset_cli` + routes variants by ``isinstance`` against these. Empty for + :attr:`PresetTarget.DOMAIN`, which is the catch-all whose membership is + "no typed target matched". +* ``legacy_aliases`` -- deprecated-name to canonical-name table for this + target, aggregated for hydra's resolver via :meth:`all_legacy_aliases`. + +Adding a new typed target = appending one enum member with its label, base +classes, and (optional) legacy alias map. The CLI layer needs no other wiring. +""" + +from __future__ import annotations + +import enum +import functools + +from isaaclab.physics import PhysicsCfg +from isaaclab.renderers.renderer_cfg import RendererCfg + + +class PresetTarget(enum.Enum): + """Typed preset categories. + + **Bucketing contract.** Help-time bucketing in + :mod:`isaaclab_tasks.utils.preset_cli` routes each preset variant to a + typed target by checking ``isinstance(cfg_value, target.base_classes)`` + against every typed target's bases. A variant whose cfg value does *not* + subclass any typed target's base falls into :attr:`DOMAIN` and shows up + under the ``presets:`` catch-all in ``--help``. + + To opt into the typed ``physics`` / ``renderer`` help-text listing, + a backend's cfg class must subclass :class:`~isaaclab.physics.PhysicsCfg` + or :class:`~isaaclab.renderers.renderer_cfg.RendererCfg` respectively. + A variant whose class does *not* subclass either base still **resolves + correctly at runtime** -- hydra applies the selected name across every + matching ``PresetCfg`` field regardless of class; the typed bucketing only + governs which header it appears under in ``--help``. + + Adding a new target = appending one enum member. + """ + + # Members. Tuple values are (label, base_classes, legacy_aliases); the + # enum metaclass collects the whole namespace before constructing members, + # so ``__new__`` below unpacks each tuple regardless of declaration order. + PHYSICS = ("physics", (PhysicsCfg,), {"newton": "newton_mjwarp", "kamino": "newton_kamino"}) + """Physics backends -- ``physics=NAME`` selector. + + Legacy aliases ``newton`` -> ``newton_mjwarp`` and ``kamino`` -> ``newton_kamino`` + exist because Newton-backend solver presets were renamed to use the + ``newton_`` prefix so they group together in autocomplete and read + distinctly from backend / package / visualizer names that also contain the + word ``newton``. Hydra's resolver (see + :func:`~isaaclab_tasks.utils.hydra._normalize_preset_name`) consults these + and emits a :class:`FutureWarning`; the aliases will be removed in a + future release. + """ + + RENDERER = ("renderer", (RendererCfg,)) + """Camera-sensor renderers -- ``renderer=NAME`` selector.""" + + DOMAIN = ("presets",) + """Free-form env-specific presets -- ``presets=NAME[,...]`` selector (catch-all). + + No ``base_classes`` -- any variant whose cfg class doesn't subclass a typed + target's base ends up here. The ``presets=`` token also acts as a + broadcast: hydra's resolver applies a DOMAIN-bucketed name to every + matching ``PresetCfg`` regardless of target. ``self.value`` matches the + CLI selector key (``"presets"``) so the CLI layer can dispatch by + enum value without a hardcoded constant. + """ + + def __new__( + cls, + label: str, + base_classes: tuple[type, ...] = (), + legacy_aliases: dict[str, str] | None = None, + ): + """Construct a member from its ``(label, base_classes, legacy_aliases)`` tuple. + + Args: + label: Hydra-style selector key (e.g. ``"physics"`` is recognized + as the ``physics=NAME`` token and becomes ``self.value``). + base_classes: Cfg base classes whose instances route to this + target via :func:`isinstance`. Defaults to ``()`` (no typed + routing). + legacy_aliases: Optional deprecated-to-canonical map for this + target; copied so members cannot alias each other's tables. + + Returns: + A new enum member with ``_value_`` set to *label*, plus + ``base_classes`` and ``legacy_aliases`` attributes. + """ + obj = object.__new__(cls) + obj._value_ = label + obj.base_classes = tuple(base_classes) + obj.legacy_aliases = dict(legacy_aliases) if legacy_aliases else {} + return obj + + @classmethod + @functools.cache + def all_legacy_aliases(cls) -> dict[str, str]: + """Flat ``{deprecated: canonical}`` view across every target. + + Resolver-layer code (in :mod:`isaaclab_tasks.utils.hydra`) needs a + target-agnostic lookup -- the ``presets=...`` token is target-agnostic + on the wire. Cached because per-member tables are immutable after + class construction, so the merged view never changes; this keeps + each lookup O(1) instead of rebuilding on every membership test or + ``[]`` access. Callers must not mutate the returned dict. + + Returns: + Mapping of every legacy alias to its canonical replacement, + aggregated across all members. + """ + return {name: rep for target in cls for name, rep in target.legacy_aliases.items()}