From 95a303c1da9e30956b394347709e27d92657dd34 Mon Sep 17 00:00:00 2001 From: Maximilian Krause Date: Wed, 3 Jun 2026 11:23:28 -0700 Subject: [PATCH 1/4] Add Newton inverse kinematics action --- .../changelog.d/max-newton-ik-manager.rst | 9 + .../isaaclab_newton/cloner/replicate.py | 13 +- .../isaaclab_newton/envs/__init__.py | 10 + .../isaaclab_newton/envs/mdp/__init__.py | 10 + .../envs/mdp/actions/__init__.py | 10 + .../envs/mdp/actions/__init__.pyi | 12 + .../envs/mdp/actions/newton_ik_actions.py | 309 ++++++++++++++++++ .../envs/mdp/actions/newton_ik_actions_cfg.py | 60 ++++ .../isaaclab_newton/ik/__init__.py | 10 + .../isaaclab_newton/ik/__init__.pyi | 13 + .../isaaclab_newton/ik/newton_ik_manager.py | 286 ++++++++++++++++ .../ik/newton_ik_manager_cfg.py | 67 ++++ .../isaaclab_newton/physics/newton_manager.py | 102 +++++- .../test/ik/test_newton_ik_manager.py | 153 +++++++++ .../physics/test_newton_prototype_models.py | 73 +++++ .../changelog.d/max-newton-ik-manager.rst | 4 + .../core/reach/config/franka/__init__.py | 24 ++ .../config/franka/agents/rsl_rl_ppo_cfg.py | 6 + .../reach/config/franka/ik_newton_env_cfg.py | 45 +++ 19 files changed, 1211 insertions(+), 5 deletions(-) create mode 100644 source/isaaclab_newton/changelog.d/max-newton-ik-manager.rst create mode 100644 source/isaaclab_newton/isaaclab_newton/envs/__init__.py create mode 100644 source/isaaclab_newton/isaaclab_newton/envs/mdp/__init__.py create mode 100644 source/isaaclab_newton/isaaclab_newton/envs/mdp/actions/__init__.py create mode 100644 source/isaaclab_newton/isaaclab_newton/envs/mdp/actions/__init__.pyi create mode 100644 source/isaaclab_newton/isaaclab_newton/envs/mdp/actions/newton_ik_actions.py create mode 100644 source/isaaclab_newton/isaaclab_newton/envs/mdp/actions/newton_ik_actions_cfg.py create mode 100644 source/isaaclab_newton/isaaclab_newton/ik/__init__.py create mode 100644 source/isaaclab_newton/isaaclab_newton/ik/__init__.pyi create mode 100644 source/isaaclab_newton/isaaclab_newton/ik/newton_ik_manager.py create mode 100644 source/isaaclab_newton/isaaclab_newton/ik/newton_ik_manager_cfg.py create mode 100644 source/isaaclab_newton/test/ik/test_newton_ik_manager.py create mode 100644 source/isaaclab_newton/test/physics/test_newton_prototype_models.py create mode 100644 source/isaaclab_tasks/changelog.d/max-newton-ik-manager.rst create mode 100644 source/isaaclab_tasks/isaaclab_tasks/core/reach/config/franka/ik_newton_env_cfg.py diff --git a/source/isaaclab_newton/changelog.d/max-newton-ik-manager.rst b/source/isaaclab_newton/changelog.d/max-newton-ik-manager.rst new file mode 100644 index 000000000000..48b597e676ba --- /dev/null +++ b/source/isaaclab_newton/changelog.d/max-newton-ik-manager.rst @@ -0,0 +1,9 @@ +Added +^^^^^ + +* Added :class:`~isaaclab_newton.ik.NewtonIKManager` and + :class:`~isaaclab_newton.envs.mdp.actions.NewtonInverseKinematicsAction` + for Newton-backed inverse kinematics, including named pose objectives and + custom Newton objective passthrough. +* Added persistent IK seeds and helpers to initialize pose-objective targets + from live Newton body transforms. diff --git a/source/isaaclab_newton/isaaclab_newton/cloner/replicate.py b/source/isaaclab_newton/isaaclab_newton/cloner/replicate.py index b19050da70d2..10506d06c3a2 100644 --- a/source/isaaclab_newton/isaaclab_newton/cloner/replicate.py +++ b/source/isaaclab_newton/isaaclab_newton/cloner/replicate.py @@ -43,8 +43,12 @@ def _build_newton_builder_from_mapping( quaternions: torch.Tensor | None = None, up_axis: str = "Z", simplify_meshes: bool = True, -) -> tuple[ModelBuilder, object, dict, list]: - """Build a Newton model builder from clone mapping inputs.""" +) -> tuple[ModelBuilder, object, dict, list, dict[str, ModelBuilder]]: + """Build a Newton model builder from clone mapping inputs. + + Returns the populated builder, stage metadata, the site index map, the + per-world transforms, and the per-source builders keyed by clone source path. + """ if positions is None: positions = torch.zeros((mapping.size(1), 3), device=mapping.device, dtype=torch.float32) if quaternions is None: @@ -98,7 +102,7 @@ def _build_newton_builder_from_mapping( site_index_map = {label: (idx, None) for label, idx in global_sites.items()} site_index_map.update((label, (None, per_world)) for label, per_world in local_site_map.items()) - return builder, stage_info, site_index_map, world_xforms + return builder, stage_info, site_index_map, world_xforms, source_builders class NewtonReplicateContext: @@ -208,7 +212,7 @@ def _merged_mapping(self) -> _MappingBatch: def replicate(self) -> tuple[ModelBuilder, object, dict]: """Build the Newton model builder from queued mappings and optionally publish it.""" sources, destinations, env_ids, mapping, positions, quaternions = self._merged_mapping() - builder, stage_info, site_index_map, world_xforms = _build_newton_builder_from_mapping( + builder, stage_info, site_index_map, world_xforms, source_builders = _build_newton_builder_from_mapping( stage=self.stage, sources=sources, destinations=destinations, @@ -224,6 +228,7 @@ def replicate(self) -> tuple[ModelBuilder, object, dict]: NewtonManager._cl_site_index_map = site_index_map NewtonManager._cl_fabric_body_bindings = fabric_body_bindings NewtonManager._world_xforms = world_xforms + NewtonManager.register_prototype_builders(sources, destinations, source_builders) NewtonManager.set_builder(builder) NewtonManager._num_envs = mapping.size(1) self._queue.clear() diff --git a/source/isaaclab_newton/isaaclab_newton/envs/__init__.py b/source/isaaclab_newton/isaaclab_newton/envs/__init__.py new file mode 100644 index 000000000000..ad7cada78cfa --- /dev/null +++ b/source/isaaclab_newton/isaaclab_newton/envs/__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 + +"""Newton-specific environment components.""" + +from isaaclab.utils.module import lazy_export + +lazy_export() diff --git a/source/isaaclab_newton/isaaclab_newton/envs/mdp/__init__.py b/source/isaaclab_newton/isaaclab_newton/envs/mdp/__init__.py new file mode 100644 index 000000000000..91e69bf91a81 --- /dev/null +++ b/source/isaaclab_newton/isaaclab_newton/envs/mdp/__init__.py @@ -0,0 +1,10 @@ +# Copyright (c) 2022-2026, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +"""Newton-specific MDP components.""" + +from isaaclab.utils.module import lazy_export + +lazy_export() diff --git a/source/isaaclab_newton/isaaclab_newton/envs/mdp/actions/__init__.py b/source/isaaclab_newton/isaaclab_newton/envs/mdp/actions/__init__.py new file mode 100644 index 000000000000..a3598c846846 --- /dev/null +++ b/source/isaaclab_newton/isaaclab_newton/envs/mdp/actions/__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 + +"""Newton-specific action terms.""" + +from isaaclab.utils.module import lazy_export + +lazy_export() diff --git a/source/isaaclab_newton/isaaclab_newton/envs/mdp/actions/__init__.pyi b/source/isaaclab_newton/isaaclab_newton/envs/mdp/actions/__init__.pyi new file mode 100644 index 000000000000..0858dc33a978 --- /dev/null +++ b/source/isaaclab_newton/isaaclab_newton/envs/mdp/actions/__init__.pyi @@ -0,0 +1,12 @@ +# 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__ = [ + "NewtonInverseKinematicsAction", + "NewtonInverseKinematicsActionCfg", +] + +from .newton_ik_actions import NewtonInverseKinematicsAction +from .newton_ik_actions_cfg import NewtonInverseKinematicsActionCfg diff --git a/source/isaaclab_newton/isaaclab_newton/envs/mdp/actions/newton_ik_actions.py b/source/isaaclab_newton/isaaclab_newton/envs/mdp/actions/newton_ik_actions.py new file mode 100644 index 000000000000..d78802c2cc3e --- /dev/null +++ b/source/isaaclab_newton/isaaclab_newton/envs/mdp/actions/newton_ik_actions.py @@ -0,0 +1,309 @@ +# 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 logging +from collections.abc import Sequence +from typing import TYPE_CHECKING + +import torch +import warp as wp +from newton import JointType +from newton import Model as NewtonModel +from newton.selection import ArticulationView + +import isaaclab.utils.math as math_utils +import isaaclab.utils.string as string_utils +from isaaclab.assets.articulation.base_articulation import BaseArticulation +from isaaclab.managers.action_manager import ActionTerm + +from isaaclab_newton.ik.newton_ik_manager import NewtonIKManager, NewtonIKPoseObjective +from isaaclab_newton.physics import NewtonManager + +if TYPE_CHECKING: + from isaaclab.envs import ManagerBasedEnv + from isaaclab.envs.utils.io_descriptors import GenericActionIODescriptor + + from .newton_ik_actions_cfg import NewtonInverseKinematicsActionCfg + + +logger = logging.getLogger(__name__) + + +class NewtonInverseKinematicsAction(ActionTerm): + """Newton inverse-kinematics action term. + + This action currently supports fixed-base articulations only. It solves IK + on the single-environment Newton prototype model registered by the cloner + and maps the resulting actuated joint coordinates back to the live batched + Isaac Lab articulation. + """ + + cfg: NewtonInverseKinematicsActionCfg + _asset: BaseArticulation + + def __init__(self, cfg: NewtonInverseKinematicsActionCfg, env: ManagerBasedEnv): + super().__init__(cfg, env) + + if not isinstance(self._asset, BaseArticulation): + raise TypeError( + f"NewtonInverseKinematicsAction expects a BaseArticulation asset, got {type(self._asset).__name__}." + ) + if not self._asset.is_fixed_base: + raise ValueError("NewtonInverseKinematicsAction currently supports fixed-base articulations only.") + + self._joint_ids, self._joint_names = self._asset.find_joints(self.cfg.joint_names) + if len(self._joint_ids) == 0: + raise ValueError(f"No joints matched Newton IK action joint_names={self.cfg.joint_names}.") + self._joint_ids_warp = wp.array(self._joint_ids, dtype=wp.int32, device=self.device) + + body_ids, body_names = self._asset.find_bodies(self.cfg.body_name) + if len(body_ids) != 1: + raise ValueError( + f"Expected one match for Newton IK body_name={self.cfg.body_name}. Found {len(body_ids)}: {body_names}." + ) + self._body_idx = body_ids[0] + self._body_name = body_names[0] + + self._raw_actions = torch.zeros(self.num_envs, self.action_dim, device=self.device) + self._processed_actions = torch.zeros_like(self._raw_actions) + self._target_pos_b = torch.zeros(self.num_envs, 3, device=self.device) + self._target_quat_b = torch.zeros(self.num_envs, 4, device=self.device) + self._target_quat_b[:, 3] = 1.0 + + self._scale = torch.zeros((self.num_envs, self.action_dim), device=self.device) + self._scale[:] = torch.tensor(self.cfg.scale, device=self.device) + + self._clip = None + if self.cfg.clip is not None: + self._clip = torch.tensor([[-float("inf"), float("inf")]], device=self.device).repeat( + self.num_envs, self.action_dim, 1 + ) + action_names = self._action_coordinate_names() + index_list, _, value_list = string_utils.resolve_matching_names_values(self.cfg.clip, action_names) + self._clip[:, index_list] = torch.tensor(value_list, device=self.device) + + if self.cfg.body_offset is not None: + self._offset_pos = torch.tensor(self.cfg.body_offset.pos, device=self.device).repeat(self.num_envs, 1) + self._offset_rot = torch.tensor(self.cfg.body_offset.rot, device=self.device).repeat(self.num_envs, 1) + link_offset_pos = tuple(self.cfg.body_offset.pos) + link_offset_rot = tuple(self.cfg.body_offset.rot) + else: + self._offset_pos, self._offset_rot = None, None + link_offset_pos = (0.0, 0.0, 0.0) + link_offset_rot = (0.0, 0.0, 0.0, 1.0) + + prototype_info = NewtonManager.get_prototype_model(self._asset.cfg.prim_path) + prototype_model = prototype_info.model + if prototype_model is None: + raise RuntimeError(f"Newton prototype model for '{self._asset.cfg.prim_path}' was not finalized.") + prototype_view = self._resolve_prototype_view(prototype_model) + self._prototype_joint_coord_ids = self._resolve_prototype_joint_coord_ids( + prototype_view, self._asset.joint_names + ) + self._prototype_controlled_coord_ids = self._resolve_prototype_joint_coord_ids( + prototype_view, self._joint_names + ) + self._prototype_link_index = self._resolve_prototype_link_index(prototype_view) + self._prototype_joint_seed = wp.to_torch(prototype_model.joint_q).to(device=self.device, dtype=torch.float32) + self._prototype_joint_seed = self._prototype_joint_seed.unsqueeze(0).repeat(self.num_envs, 1).contiguous() + self._ik_target_name = "target" + + self._ik_manager = NewtonIKManager( + self.cfg.controller, + model=prototype_model, + num_envs=self.num_envs, + device=self.device, + pose_objectives=[ + NewtonIKPoseObjective( + name=self._ik_target_name, + link_index=self._prototype_link_index, + link_offset_pos=link_offset_pos, + link_offset_rot=link_offset_rot, + ) + ], + ) + + logger.info( + "Resolved Newton IK action joints %s [%s] and body %s [%s].", + self._joint_names, + self._joint_ids, + self._body_name, + self._body_idx, + ) + + @property + def action_dim(self) -> int: + if self.cfg.controller.command_type == "position": + return 3 + if self.cfg.controller.command_type == "pose" and self.cfg.controller.use_relative_mode: + return 6 + if self.cfg.controller.command_type == "pose": + return 7 + raise ValueError(f"Unsupported Newton IK command type: {self.cfg.controller.command_type}") + + @property + def raw_actions(self) -> torch.Tensor: + return self._raw_actions + + @property + def processed_actions(self) -> torch.Tensor: + return self._processed_actions + + @property + def IO_descriptor(self) -> GenericActionIODescriptor: + super().IO_descriptor + self._IO_descriptor.shape = (self.action_dim,) + self._IO_descriptor.dtype = str(self.raw_actions.dtype) + self._IO_descriptor.action_type = "NewtonInverseKinematicsAction" + self._IO_descriptor.body_name = self._body_name + self._IO_descriptor.joint_names = self._joint_names + self._IO_descriptor.scale = self._scale + self._IO_descriptor.clip = self.cfg.clip + self._IO_descriptor.extras["controller_cfg"] = self.cfg.controller.__dict__ + self._IO_descriptor.extras["body_offset"] = ( + None if self.cfg.body_offset is None else self.cfg.body_offset.__dict__ + ) + return self._IO_descriptor + + def process_actions(self, actions: torch.Tensor) -> None: + self._raw_actions[:] = actions + self._processed_actions[:] = self.raw_actions * self._scale + if self._clip is not None: + self._processed_actions = torch.clamp( + self._processed_actions, min=self._clip[:, :, 0], max=self._clip[:, :, 1] + ) + + ee_pos_b, ee_quat_b = self._compute_frame_pose() + if self.cfg.controller.command_type == "position": + if self.cfg.controller.use_relative_mode: + self._target_pos_b[:] = ee_pos_b + self._processed_actions + else: + self._target_pos_b[:] = self._processed_actions + self._target_quat_b[:] = ee_quat_b + elif self.cfg.controller.use_relative_mode: + self._target_pos_b[:], self._target_quat_b[:] = math_utils.apply_delta_pose( + ee_pos_b, ee_quat_b, self._processed_actions + ) + else: + self._target_pos_b[:] = self._processed_actions[:, 0:3] + self._target_quat_b[:] = self._processed_actions[:, 3:7] + + def apply_actions(self) -> None: + # The IK solve runs on the single-env prototype model, so all batched + # root-frame targets are expressed in the prototype (env 0) world frame. + self._validate_matching_root_orientations() + root_pos_proto = self._asset.data.root_pos_w.torch[0:1].repeat(self.num_envs, 1) + root_quat_proto = self._asset.data.root_quat_w.torch[0:1].repeat(self.num_envs, 1) + target_pos_w, target_quat_w = math_utils.combine_frame_transforms( + root_pos_proto, root_quat_proto, self._target_pos_b, self._target_quat_b + ) + self._ik_manager.set_target_pose(self._ik_target_name, target_pos_w, target_quat_w) + + joint_seed = self._prototype_joint_seed.clone() + joint_seed[:, self._prototype_joint_coord_ids] = self._asset.data.joint_pos.torch + joint_pos_des_all = self._ik_manager.solve(joint_seed) + joint_pos_des = joint_pos_des_all[:, self._prototype_controlled_coord_ids].contiguous() + self._asset.set_joint_position_target_index(target=joint_pos_des, joint_ids=self._joint_ids_warp) + + def reset(self, env_ids: Sequence[int] | None = None) -> None: + env_ids = slice(None) if env_ids is None else env_ids + self._raw_actions[env_ids] = 0.0 + + def _compute_frame_pose(self) -> tuple[torch.Tensor, torch.Tensor]: + ee_pos_w = self._asset.data.body_pos_w.torch[:, self._body_idx] + ee_quat_w = self._asset.data.body_quat_w.torch[:, self._body_idx] + root_pos_w = self._asset.data.root_pos_w.torch + root_quat_w = self._asset.data.root_quat_w.torch + ee_pos_b, ee_quat_b = math_utils.subtract_frame_transforms(root_pos_w, root_quat_w, ee_pos_w, ee_quat_w) + if self.cfg.body_offset is not None: + ee_pos_b, ee_quat_b = math_utils.combine_frame_transforms( + ee_pos_b, ee_quat_b, self._offset_pos, self._offset_rot + ) + return ee_pos_b, ee_quat_b + + def _validate_matching_root_orientations(self) -> None: + """Guard the prototype-frame IK assumption for replicated fixed-base roots.""" + root_quat_w = self._asset.data.root_quat_w.torch + if root_quat_w.shape[0] <= 1: + return + # q and -q represent the same orientation, so compare absolute dot products. + same_orientation = torch.abs(torch.sum(root_quat_w * root_quat_w[0:1], dim=-1)) > 1.0 - 1e-5 + if not torch.all(same_orientation): + bad_env_ids = torch.nonzero(~same_orientation, as_tuple=False).flatten().tolist() + raise RuntimeError( + "NewtonInverseKinematicsAction solves against the env 0 prototype root orientation, but " + f"root orientations differ in env ids {bad_env_ids}. Use identical fixed-base root orientations " + "for this action." + ) + + def _resolve_prototype_view(self, model) -> ArticulationView: + requested_path = NewtonManager._to_first_env_path(self._asset.cfg.prim_path) + basename = requested_path.rsplit("/", 1)[-1] + patterns = [requested_path, requested_path.replace(".*", "*"), basename, f"*{basename}"] + last_error: Exception | None = None + for pattern in patterns: + try: + view = ArticulationView( + model, + pattern, + verbose=False, + exclude_joint_types=[JointType.FREE, JointType.FIXED], + ) + except (KeyError, ValueError) as exc: + last_error = exc + continue + if view.world_count != 1 or view.count_per_world != 1: + raise ValueError( + f"Newton IK expected one prototype articulation for '{self._asset.cfg.prim_path}', " + f"got world_count={view.world_count}, count_per_world={view.count_per_world}." + ) + if not view.is_fixed_base: + raise ValueError("Newton IK currently supports fixed-base prototype articulations only.") + return view + raise KeyError( + f"Failed to resolve Newton prototype articulation for '{self._asset.cfg.prim_path}'." + ) from last_error + + def _resolve_prototype_joint_coord_ids( + self, prototype_view: ArticulationView, joint_names: Sequence[str] + ) -> torch.Tensor: + layout = prototype_view.frequency_layouts[NewtonModel.AttributeFrequency.JOINT_COORD] + selected_indices = self._layout_indices(layout) + coord_indices_by_name = { + name: layout.offset + selected_indices[index] for index, name in enumerate(prototype_view.joint_coord_names) + } + try: + coord_ids = [coord_indices_by_name[name] for name in joint_names] + except KeyError as exc: + raise KeyError( + f"Joint '{exc.args[0]}' was resolved in Isaac Lab but not in the Newton prototype model." + ) from exc + return torch.tensor(coord_ids, device=self.device, dtype=torch.long) + + def _resolve_prototype_link_index(self, prototype_view: ArticulationView) -> int: + layout = prototype_view.frequency_layouts[NewtonModel.AttributeFrequency.BODY] + selected_indices = self._layout_indices(layout) + try: + local_link_index = prototype_view.link_names.index(self._body_name) + except ValueError as exc: + raise ValueError( + f"Body '{self._body_name}' was resolved in Isaac Lab but not in the Newton prototype model." + ) from exc + return layout.offset + selected_indices[local_link_index] + + @staticmethod + def _layout_indices(layout) -> list[int]: + if layout.slice is not None: + return list(range(layout.slice.start, layout.slice.stop)) + return [int(index) for index in layout.indices.numpy().tolist()] + + def _action_coordinate_names(self) -> list[str]: + if self.cfg.controller.command_type == "position": + return ["x", "y", "z"] + if self.cfg.controller.command_type == "pose" and self.cfg.controller.use_relative_mode: + return ["x", "y", "z", "roll", "pitch", "yaw"] + return ["x", "y", "z", "qx", "qy", "qz", "qw"] diff --git a/source/isaaclab_newton/isaaclab_newton/envs/mdp/actions/newton_ik_actions_cfg.py b/source/isaaclab_newton/isaaclab_newton/envs/mdp/actions/newton_ik_actions_cfg.py new file mode 100644 index 000000000000..31504333ac5b --- /dev/null +++ b/source/isaaclab_newton/isaaclab_newton/envs/mdp/actions/newton_ik_actions_cfg.py @@ -0,0 +1,60 @@ +# Copyright (c) 2022-2026, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +from __future__ import annotations + +from dataclasses import MISSING +from typing import TYPE_CHECKING + +from isaaclab.managers.action_manager import ActionTermCfg +from isaaclab.utils.configclass import configclass + +from isaaclab_newton.ik.newton_ik_manager_cfg import NewtonIKManagerCfg + +if TYPE_CHECKING: + from .newton_ik_actions import NewtonInverseKinematicsAction + + +@configclass +class NewtonInverseKinematicsActionCfg(ActionTermCfg): + """Configuration for a Newton inverse-kinematics action term. + + The action currently supports fixed-base articulations only. The configured + joints and body must resolve both in Isaac Lab and in the registered Newton + prototype model for the controlled asset. + """ + + @configclass + class OffsetCfg: + """Offset pose from the controlled body frame to the IK target frame.""" + + pos: tuple[float, float, float] = (0.0, 0.0, 0.0) + """Translation [m] w.r.t. the parent body frame.""" + + rot: tuple[float, float, float, float] = (0.0, 0.0, 0.0, 1.0) + """Quaternion rotation ``(x, y, z, w)`` w.r.t. the parent body frame.""" + + class_type: type[NewtonInverseKinematicsAction] | str = ( + "isaaclab_newton.envs.mdp.actions.newton_ik_actions:NewtonInverseKinematicsAction" + ) + + joint_names: list[str] = MISSING + """List of joint names or regex expressions controlled by the action.""" + + body_name: str = MISSING + """Name of the body for which IK is performed.""" + + body_offset: OffsetCfg | None = None + """Offset of the target frame w.r.t. the body frame.""" + + scale: float | tuple[float, ...] = 1.0 + """Scale factor applied to the raw action. + + For position coordinates this is in meters. For relative rotation coordinates + this is in radians. For quaternions this is dimensionless. + """ + + controller: NewtonIKManagerCfg = MISSING + """Configuration for the Newton IK manager.""" diff --git a/source/isaaclab_newton/isaaclab_newton/ik/__init__.py b/source/isaaclab_newton/isaaclab_newton/ik/__init__.py new file mode 100644 index 000000000000..ac24b2ccf6e8 --- /dev/null +++ b/source/isaaclab_newton/isaaclab_newton/ik/__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 + +"""Newton inverse-kinematics utilities.""" + +from isaaclab.utils.module import lazy_export + +lazy_export() diff --git a/source/isaaclab_newton/isaaclab_newton/ik/__init__.pyi b/source/isaaclab_newton/isaaclab_newton/ik/__init__.pyi new file mode 100644 index 000000000000..2d751bb38eb2 --- /dev/null +++ b/source/isaaclab_newton/isaaclab_newton/ik/__init__.pyi @@ -0,0 +1,13 @@ +# 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__ = [ + "NewtonIKManager", + "NewtonIKManagerCfg", + "NewtonIKPoseObjective", +] + +from .newton_ik_manager import NewtonIKManager, NewtonIKPoseObjective +from .newton_ik_manager_cfg import NewtonIKManagerCfg diff --git a/source/isaaclab_newton/isaaclab_newton/ik/newton_ik_manager.py b/source/isaaclab_newton/isaaclab_newton/ik/newton_ik_manager.py new file mode 100644 index 000000000000..f8b02119994e --- /dev/null +++ b/source/isaaclab_newton/isaaclab_newton/ik/newton_ik_manager.py @@ -0,0 +1,286 @@ +# 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 collections.abc import Sequence +from dataclasses import dataclass +from numbers import Integral +from typing import Any + +import newton.ik as ik +import torch +import warp as wp + +from .newton_ik_manager_cfg import NewtonIKManagerCfg + + +@dataclass(frozen=True) +class NewtonIKPoseObjective: + """Pose objective descriptor used to build Newton position/rotation objectives. + + Args: + name: Unique objective name used when updating targets. + link_index: Newton body index for the controlled link. + link_offset_pos: Target frame translation in meters relative to the link frame. + link_offset_rot: Target frame quaternion ``(x, y, z, w)`` relative to the link frame. + position_weight: Optional residual weight overriding the manager default. + rotation_weight: Optional residual weight overriding the manager default. + """ + + name: str + link_index: int + link_offset_pos: tuple[float, float, float] = (0.0, 0.0, 0.0) + link_offset_rot: tuple[float, float, float, float] = (0.0, 0.0, 0.0, 1.0) + position_weight: float | None = None + rotation_weight: float | None = None + + +class NewtonIKManager: + """Batched wrapper around Newton's inverse-kinematics solver. + + The manager mirrors Newton's objective-list design while adding torch/Warp + target updates convenient for Isaac Lab action terms. Pose objectives are + named so callers can update individual targets between solves. Additional + custom Newton objectives can be passed through ``extra_objectives``. + """ + + cfg: NewtonIKManagerCfg + + def __init__( + self, + cfg: NewtonIKManagerCfg, + *, + model, + num_envs: int, + device: str, + pose_objectives: Sequence[NewtonIKPoseObjective], + extra_objectives: Sequence[ik.IKObjective] | None = None, + ): + if not pose_objectives and not extra_objectives: + raise ValueError("NewtonIKManager requires at least one pose or custom objective.") + + self.cfg = cfg + self.model = model + self.num_envs = num_envs + self.device = device + self.num_coords = model.joint_coord_count + self.pose_objective_names = [objective.name for objective in pose_objectives] + self.pose_objective_cfgs = {objective.name: objective for objective in pose_objectives} + + if len(set(self.pose_objective_names)) != len(self.pose_objective_names): + raise ValueError(f"Newton IK pose objective names must be unique: {self.pose_objective_names}") + + self.position_objectives: dict[str, ik.IKObjectivePosition] = {} + self.rotation_objectives: dict[str, ik.IKObjectiveRotation] = {} + solver_objectives: list[ik.IKObjective] = [] + + for objective_cfg in pose_objectives: + target_positions = wp.zeros((num_envs,), dtype=wp.vec3, device=device) + target_rotations = wp.array( + [(0.0, 0.0, 0.0, 1.0)] * num_envs, + dtype=wp.vec4, + device=device, + ) + position_objective = ik.IKObjectivePosition( + link_index=objective_cfg.link_index, + link_offset=wp.vec3(*objective_cfg.link_offset_pos), + target_positions=target_positions, + weight=cfg.position_weight if objective_cfg.position_weight is None else objective_cfg.position_weight, + ) + rotation_objective = ik.IKObjectiveRotation( + link_index=objective_cfg.link_index, + link_offset_rotation=wp.quat(*objective_cfg.link_offset_rot), + target_rotations=target_rotations, + weight=cfg.rotation_weight if objective_cfg.rotation_weight is None else objective_cfg.rotation_weight, + ) + self.position_objectives[objective_cfg.name] = position_objective + self.rotation_objectives[objective_cfg.name] = rotation_objective + solver_objectives.extend((position_objective, rotation_objective)) + + if cfg.joint_limit_weight is not None: + self.joint_limit_objective = ik.IKObjectiveJointLimit( + joint_limit_lower=model.joint_limit_lower, + joint_limit_upper=model.joint_limit_upper, + weight=cfg.joint_limit_weight, + ) + solver_objectives.append(self.joint_limit_objective) + else: + self.joint_limit_objective = None + + solver_objectives.extend(extra_objectives or []) + + self.joint_q_out = wp.zeros((num_envs, self.num_coords), dtype=wp.float32, device=device) + self.joint_q_seed_t = torch.zeros((num_envs, self.num_coords), dtype=torch.float32, device=device) + self.joint_q_seed = wp.from_torch(self.joint_q_seed_t, dtype=wp.float32) + self._has_joint_q_seed = False + self.solver = ik.IKSolver( + model=model, + n_problems=num_envs, + objectives=solver_objectives, + optimizer=ik.IKOptimizer(cfg.optimizer), + jacobian_mode=ik.IKJacobianType(cfg.jacobian_mode), + sampler=ik.IKSampler(cfg.sampler), + n_seeds=cfg.n_seeds, + noise_std=cfg.noise_std, + rng_seed=cfg.rng_seed, + lambda_initial=cfg.lambda_initial, + ) + + @property + def action_dim(self) -> int: + """Dimension of the IK command expected by this manager.""" + if self.cfg.command_type == "position": + return 3 + if self.cfg.command_type == "pose" and self.cfg.use_relative_mode: + return 6 + if self.cfg.command_type == "pose": + return 7 + raise ValueError(f"Unsupported Newton IK command type: {self.cfg.command_type}") + + def set_target_pose(self, name: str, target_pos_w: torch.Tensor, target_quat_w: torch.Tensor) -> None: + """Update batched world-frame target poses for a named pose objective.""" + try: + position_objective = self.position_objectives[name] + rotation_objective = self.rotation_objectives[name] + except KeyError as exc: + raise KeyError( + f"Unknown Newton IK pose objective '{name}'. Available objectives: {self.pose_objective_names}." + ) from exc + position_objective.set_target_positions(wp.from_torch(target_pos_w.contiguous(), dtype=wp.vec3)) + rotation_objective.set_target_rotations(wp.from_torch(target_quat_w.contiguous(), dtype=wp.vec4)) + + def set_target_pose_from_body_q( + self, + name: str, + body_q: torch.Tensor | wp.array | Any, + *, + env_origins: torch.Tensor | None = None, + ) -> None: + """Set one pose objective target from body transforms.""" + objective_cfg = self.pose_objective_cfgs[name] + body_q_t = _as_torch(body_q, device=self.device) + if body_q_t.ndim == 2: + target_pos = body_q_t[objective_cfg.link_index, :3].reshape(1, 3).repeat(self.num_envs, 1) + target_quat = body_q_t[objective_cfg.link_index, 3:7].reshape(1, 4).repeat(self.num_envs, 1) + elif body_q_t.ndim == 3: + if body_q_t.shape[0] != self.num_envs: + raise ValueError(f"Expected body_q first dimension {self.num_envs}, got {body_q_t.shape[0]}.") + target_pos = body_q_t[:, objective_cfg.link_index, :3] + target_quat = body_q_t[:, objective_cfg.link_index, 3:7] + else: + raise ValueError( + f"Expected body_q shape (num_bodies, 7) or (num_envs, num_bodies, 7), got {body_q_t.shape}." + ) + if env_origins is not None: + target_pos = target_pos - env_origins.to(device=self.device, dtype=torch.float32) + self.set_target_pose(name, target_pos.to(torch.float32), target_quat.to(torch.float32)) + + def set_pose_targets_from_body_q( + self, + body_q: torch.Tensor | wp.array | Any, + names: Sequence[str] | None = None, + *, + env_origins: torch.Tensor | None = None, + ) -> None: + """Set multiple pose objective targets from body transforms.""" + for name in self.pose_objective_names if names is None else names: + self.set_target_pose_from_body_q(name, body_q, env_origins=env_origins) + + def set_joint_seed( + self, joint_pos: torch.Tensor, env_ids: Sequence[int] | torch.Tensor | slice | None = None + ) -> None: + """Set the persistent joint-coordinate seed used by ``solve()`` without an explicit seed.""" + joint_pos = joint_pos.to(device=self.device, dtype=torch.float32) + if env_ids is None: + if joint_pos.shape != (self.num_envs, self.num_coords): + raise ValueError( + f"Expected joint seed shape {(self.num_envs, self.num_coords)}, got {tuple(joint_pos.shape)}." + ) + self.joint_q_seed_t.copy_(joint_pos) + else: + ids = _env_ids_to_tensor(env_ids, self.num_envs, self.device) + if joint_pos.shape != (ids.numel(), self.num_coords): + raise ValueError( + f"Expected joint seed shape {(ids.numel(), self.num_coords)}, got {tuple(joint_pos.shape)}." + ) + self.joint_q_seed_t[ids] = joint_pos + self._has_joint_q_seed = True + + def reset( + self, + env_ids: Sequence[int] | torch.Tensor | slice | None = None, + joint_pos: torch.Tensor | None = None, + ) -> None: + """Reset Newton solver state, selected seeds, and sampler RNG.""" + self.solver.reset() + if joint_pos is not None: + self.set_joint_seed(joint_pos, env_ids=env_ids) + elif env_ids is None: + self._has_joint_q_seed = False + + @property + def costs(self) -> wp.array: + """Expanded per-seed costs from the most recent Newton solve.""" + return self.solver.costs + + @property + def joint_q(self) -> wp.array: + """Expanded joint-coordinate buffer storing all sampled seeds.""" + return self.solver.joint_q + + def solve(self, joint_pos: torch.Tensor | None = None) -> torch.Tensor: + """Solve IK from an explicit seed or the manager's persistent seed.""" + if joint_pos is None: + if not self._has_joint_q_seed: + raise RuntimeError("NewtonIKManager.solve() needs joint_pos or a seed set with set_joint_seed().") + joint_q_in = self.joint_q_seed + update_seed = True + else: + if joint_pos.shape != (self.num_envs, self.num_coords): + raise ValueError( + f"Expected joint seed shape {(self.num_envs, self.num_coords)}, got {tuple(joint_pos.shape)}." + ) + if self.cfg.use_persistent_seed: + self.set_joint_seed(joint_pos) + joint_q_in = self.joint_q_seed + update_seed = True + else: + joint_q_in = wp.from_torch(joint_pos.contiguous(), dtype=wp.float32) + update_seed = False + self.solver.step( + joint_q_in, + self.joint_q_out, + iterations=self.cfg.iterations, + step_size=self.cfg.step_size, + ) + result = wp.to_torch(self.joint_q_out).clone() + if update_seed: + self.joint_q_seed_t.copy_(result) + self._has_joint_q_seed = True + return result + + def step(self) -> torch.Tensor: + """Solve IK from the persistent seed and store the result as the next seed.""" + return self.solve() + + +def _as_torch(value, *, device: str) -> torch.Tensor: + if isinstance(value, torch.Tensor): + return value.to(device=device, dtype=torch.float32) + if hasattr(value, "numpy"): + return wp.to_torch(value).to(device=device, dtype=torch.float32) + return torch.as_tensor(value, device=device, dtype=torch.float32) + + +def _env_ids_to_tensor(env_ids: Sequence[int] | torch.Tensor | slice, num_envs: int, device: str) -> torch.Tensor: + if isinstance(env_ids, slice): + start, stop, step = env_ids.indices(num_envs) + return torch.arange(start, stop, step, device=device, dtype=torch.long) + if isinstance(env_ids, Integral): + return torch.as_tensor([env_ids], device=device, dtype=torch.long) + if isinstance(env_ids, torch.Tensor): + return env_ids.to(device=device, dtype=torch.long).flatten() + return torch.as_tensor(list(env_ids), device=device, dtype=torch.long) diff --git a/source/isaaclab_newton/isaaclab_newton/ik/newton_ik_manager_cfg.py b/source/isaaclab_newton/isaaclab_newton/ik/newton_ik_manager_cfg.py new file mode 100644 index 000000000000..10fb106ae7ef --- /dev/null +++ b/source/isaaclab_newton/isaaclab_newton/ik/newton_ik_manager_cfg.py @@ -0,0 +1,67 @@ +# 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 isaaclab.utils.configclass import configclass + + +@configclass +class NewtonIKManagerCfg: + """Configuration for the Newton inverse-kinematics manager.""" + + command_type: str = "pose" + """Action command type for manager-based actions: ``"position"`` or ``"pose"``.""" + + use_relative_mode: bool = True + """Whether manager-based action commands are relative to the current target pose.""" + + optimizer: str = "lm" + """Newton IK optimizer backend. Supported values are ``"lm"`` and ``"lbfgs"``.""" + + jacobian_mode: str = "analytic" + """Newton IK Jacobian backend. Supported values are ``"analytic"``, ``"autodiff"``, and ``"mixed"``.""" + + sampler: str = "none" + """Initial seed sampler. Supported values are ``"none"``, ``"gauss"``, ``"roberts"``, and ``"uniform"``.""" + + n_seeds: int = 1 + """Number of candidate seeds per IK problem. Must be ``1`` when ``sampler="none"``.""" + + noise_std: float = 0.1 + """Gaussian sampling standard deviation used when ``sampler="gauss"``.""" + + rng_seed: int = 12345 + """Random seed used by stochastic samplers.""" + + iterations: int = 24 + """Number of Newton IK solver iterations per action application. + + The default keeps manager-based action applications affordable while still + giving the Newton solver several refinement steps per control update. + Increase this for harder targets or tighter residual requirements. + """ + + step_size: float = 1.0 + """LM step scale passed to Newton ``IKSolver.step``. Ignored by L-BFGS.""" + + lambda_initial: float = 0.1 + """Initial damping value for the Newton Levenberg-Marquardt optimizer. + + This moderate default favors stable updates near singular configurations + over aggressive first-step motion. + """ + + position_weight: float = 1.0 + """Default residual weight for pose position objectives.""" + + rotation_weight: float = 1.0 + """Default residual weight for pose rotation objectives.""" + + joint_limit_weight: float | None = 0.1 + """Residual weight for the joint-limit objective. Set to ``None`` to disable it.""" + + use_persistent_seed: bool = False + """Whether solved joint coordinates should be reused as the next IK seed.""" diff --git a/source/isaaclab_newton/isaaclab_newton/physics/newton_manager.py b/source/isaaclab_newton/isaaclab_newton/physics/newton_manager.py index bb5279e7d8be..83164e5f27db 100644 --- a/source/isaaclab_newton/isaaclab_newton/physics/newton_manager.py +++ b/source/isaaclab_newton/isaaclab_newton/physics/newton_manager.py @@ -62,6 +62,18 @@ logger = logging.getLogger(__name__) + +@dataclass +class NewtonPrototypeModelInfo: + """Cached Newton prototype model built before physics replication.""" + + source_path: str + destination_path: str + builder: ModelBuilder + model: Model | None = None + model_device: object | None = None + + # Tagged union for entries in _cl_site_index_map. # _GlobalSite: (global_shape_idx, None) — body_pattern was None # _LocalSite: (None, [[env0_idx, ...], ...]) — per-world site indices @@ -227,6 +239,7 @@ class NewtonManager(PhysicsManager): # Newton model and state _builder: ModelBuilder = None _model: Model = None + _prototype_models: dict[str, NewtonPrototypeModelInfo] = {} _solver: SolverBase | None = None _use_single_state: bool | None = None """Use only one state for both input and output for solver stepping. Requires solver support.""" @@ -782,6 +795,7 @@ def clear(cls): NewtonManager._up_axis = "Z" NewtonManager._scene_data = None NewtonManager._scene_data_mapping = None + NewtonManager._prototype_models = {} NewtonManager._model_changes = set() NewtonManager._scene_data_backend = None NewtonManager._cl_pending_sites = {} @@ -798,6 +812,92 @@ def set_builder(cls, builder: ModelBuilder) -> None: """Set the Newton model builder.""" NewtonManager._builder = builder + @classmethod + def register_prototype_builders( + cls, + sources: list[str] | tuple[str, ...], + destinations: list[str] | tuple[str, ...], + proto_builders: dict[str, ModelBuilder], + ) -> None: + """Register prototype builders created before Newton physics replication. + + Prototype builders preserve the single-source model that is later added + into each replicated Newton world. Controllers that need a single-model + representation, such as batched Newton IK, should use this registry + instead of re-importing USD after the scene has already been built. + """ + for index, source_path in enumerate(sources): + builder = proto_builders.get(source_path) + if builder is None: + continue + destination_path = destinations[index] if index < len(destinations) else destinations[-1] + NewtonManager._prototype_models[source_path] = NewtonPrototypeModelInfo( + source_path=source_path, + destination_path=destination_path, + builder=builder, + ) + + @staticmethod + def _to_first_env_path(path: str) -> str: + """Convert common Isaac Lab env regex/template paths to env_0 paths.""" + return ( + path.replace("env_.*", "env_0") + .replace("env_\\d+", "env_0") + .replace("env_*", "env_0") + .replace("env_{}", "env_0") + ) + + @classmethod + def get_prototype_model(cls, prim_path: str) -> NewtonPrototypeModelInfo: + """Return the prototype model that owns ``prim_path``. + + Args: + prim_path: Live Isaac Lab prim path or env-regex path, such as + ``/World/envs/env_.*/Robot``. + + Returns: + Prototype model information with ``model`` finalized on the active + Newton device. + + Raises: + RuntimeError: If no prototype builders have been registered. + KeyError: If ``prim_path`` does not resolve to a registered prototype. + ValueError: If multiple registered prototypes match equally well. + """ + if not NewtonManager._prototype_models: + raise RuntimeError( + "No Newton prototype builders are registered. Prototype access is only available after " + "Newton physics replication has built the scene." + ) + + requested_path = cls._to_first_env_path(prim_path) + matches: list[NewtonPrototypeModelInfo] = [] + for info in NewtonManager._prototype_models.values(): + source_path = cls._to_first_env_path(info.source_path) + destination_path = cls._to_first_env_path(info.destination_path) + if requested_path in (source_path, destination_path) or requested_path.startswith( + (source_path + "/", destination_path + "/") + ): + matches.append(info) + + if not matches: + available = ", ".join(sorted(NewtonManager._prototype_models)) + raise KeyError(f"No Newton prototype model matches '{prim_path}'. Available prototypes: {available}") + + matches.sort(key=lambda item: len(cls._to_first_env_path(item.source_path)), reverse=True) + if len(matches) > 1: + best_len = len(cls._to_first_env_path(matches[0].source_path)) + if len(cls._to_first_env_path(matches[1].source_path)) == best_len: + paths = ", ".join(match.source_path for match in matches) + raise ValueError(f"Ambiguous Newton prototype model for '{prim_path}'. Matches: {paths}") + + info = matches[0] + device = cls._model.device if cls._model is not None else None + if info.model is None or info.model_device != device: + info.model = info.builder.finalize(device=device) + info.model_device = device + return info + @classmethod def create_builder(cls, up_axis: str | None = None, **kwargs) -> ModelBuilder: """Create a :class:`ModelBuilder` configured with default settings. @@ -1242,6 +1342,7 @@ def instantiate_builder_from_stage(cls): _, proto_path = env_paths[0] source_builders = {proto_path: cls.create_builder(up_axis=up_axis)} source_builders[proto_path].add_usd(stage, root_path=proto_path, schema_resolvers=schema_resolvers) + cls.register_prototype_builders((proto_path,), (proto_path,), source_builders) global_site_indices, source_site_indices, env_root_sites = cls._cl_inject_sites(builder, source_builders) xform_cache = UsdGeom.XformCache() @@ -1861,7 +1962,6 @@ def update_visualization_state(cls, scene_data_provider: SceneDataProvider | Non Invoked lazily from :meth:`get_state` so consumers do not need to coordinate the sync explicitly. """ - if scene_data_provider is None: scene_data_provider = cls.get_scene_data_provider() diff --git a/source/isaaclab_newton/test/ik/test_newton_ik_manager.py b/source/isaaclab_newton/test/ik/test_newton_ik_manager.py new file mode 100644 index 000000000000..865e279c6388 --- /dev/null +++ b/source/isaaclab_newton/test/ik/test_newton_ik_manager.py @@ -0,0 +1,153 @@ +# Copyright (c) 2022-2026, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +from __future__ import annotations + +import isaaclab_newton.ik.newton_ik_manager as ik_manager_module +import torch +import warp as wp +from isaaclab_newton.ik.newton_ik_manager import NewtonIKManager, NewtonIKPoseObjective +from isaaclab_newton.ik.newton_ik_manager_cfg import NewtonIKManagerCfg + + +class _Model: + joint_coord_count = 2 + joint_limit_lower = None + joint_limit_upper = None + + +class _PoseObjective: + def __init__(self, **kwargs): + self.kwargs = kwargs + self.target = kwargs.get("target_positions", kwargs.get("target_rotations")) + + def set_target_positions(self, target): + self.target = target + + def set_target_rotations(self, target): + self.target = target + + +class _JointLimitObjective: + def __init__(self, **kwargs): + self.kwargs = kwargs + + +class _Solver: + def __init__(self, **kwargs): + self.kwargs = kwargs + self.joint_q = wp.zeros((1, 2), dtype=wp.float32, device="cpu") + self.costs = wp.zeros((1,), dtype=wp.float32, device="cpu") + self.reset_count = 0 + + def reset(self): + self.reset_count += 1 + + def step(self, joint_q_in, joint_q_out, *, iterations, step_size): + del iterations, step_size + wp.to_torch(joint_q_out).copy_(wp.to_torch(joint_q_in) + 1.0) + + +def _patch_newton_ik(monkeypatch): + monkeypatch.setattr(ik_manager_module.ik, "IKObjectivePosition", _PoseObjective) + monkeypatch.setattr(ik_manager_module.ik, "IKObjectiveRotation", _PoseObjective) + monkeypatch.setattr(ik_manager_module.ik, "IKObjectiveJointLimit", _JointLimitObjective) + monkeypatch.setattr(ik_manager_module.ik, "IKSolver", _Solver) + monkeypatch.setattr(ik_manager_module.ik, "IKOptimizer", lambda value: value) + monkeypatch.setattr(ik_manager_module.ik, "IKJacobianType", lambda value: value) + monkeypatch.setattr(ik_manager_module.ik, "IKSampler", lambda value: value) + + +def _cfg() -> NewtonIKManagerCfg: + cfg = NewtonIKManagerCfg() + cfg.use_persistent_seed = True + cfg.joint_limit_weight = None + return cfg + + +def test_persistent_seed_is_reused_between_solves(monkeypatch): + _patch_newton_ik(monkeypatch) + manager = NewtonIKManager( + _cfg(), + model=_Model(), + num_envs=2, + device="cpu", + pose_objectives=[NewtonIKPoseObjective(name="ee", link_index=0)], + ) + + seed = torch.tensor([[1.0, 2.0], [3.0, 4.0]]) + manager.set_joint_seed(seed) + + first = manager.solve().clone() + second = manager.solve().clone() + + assert torch.allclose(first, seed + 1.0) + assert torch.allclose(second, seed + 2.0) + + +def test_solve_result_is_independent_from_next_solve(monkeypatch): + _patch_newton_ik(monkeypatch) + cfg = _cfg() + cfg.use_persistent_seed = False + manager = NewtonIKManager( + cfg, + model=_Model(), + num_envs=2, + device="cpu", + pose_objectives=[NewtonIKPoseObjective(name="ee", link_index=0)], + ) + seed = torch.tensor([[1.0, 2.0], [3.0, 4.0]]) + + first = manager.solve(seed) + expected_first = first.clone() + manager.solve(seed + 10.0) + + assert torch.allclose(first, expected_first) + + +def test_set_target_pose_updates_named_objective(monkeypatch): + _patch_newton_ik(monkeypatch) + manager = NewtonIKManager( + _cfg(), + model=_Model(), + num_envs=2, + device="cpu", + pose_objectives=[NewtonIKPoseObjective(name="ee", link_index=0)], + ) + target_pos = torch.tensor([[0.1, 0.2, 0.3], [1.0, 1.1, 1.2]]) + target_quat = torch.tensor([[0.0, 0.0, 0.0, 1.0], [0.0, 0.0, 1.0, 0.0]]) + + manager.set_target_pose("ee", target_pos, target_quat) + + assert torch.allclose(wp.to_torch(manager.position_objectives["ee"].target), target_pos) + assert torch.allclose(wp.to_torch(manager.rotation_objectives["ee"].target), target_quat) + + +def test_pose_targets_can_be_initialized_from_body_transforms(monkeypatch): + _patch_newton_ik(monkeypatch) + manager = NewtonIKManager( + _cfg(), + model=_Model(), + num_envs=2, + device="cpu", + pose_objectives=[ + NewtonIKPoseObjective(name="ee", link_index=0), + NewtonIKPoseObjective(name="torso", link_index=1, position_weight=50.0, rotation_weight=50.0), + ], + ) + body_q = torch.tensor( + [ + [1.0, 2.0, 3.0, 0.0, 0.0, 0.0, 1.0], + [4.0, 5.0, 6.0, 0.0, 0.0, 1.0, 0.0], + ] + ) + env_origins = torch.tensor([[0.0, 0.0, 0.0], [1.0, 1.0, 1.0]]) + + manager.set_pose_targets_from_body_q(body_q, names=["torso"], env_origins=env_origins) + + target_pos = wp.to_torch(manager.position_objectives["torso"].target) + target_quat = wp.to_torch(manager.rotation_objectives["torso"].target) + assert torch.allclose(target_pos, torch.tensor([[4.0, 5.0, 6.0], [3.0, 4.0, 5.0]])) + assert torch.allclose(target_quat, torch.tensor([[0.0, 0.0, 1.0, 0.0], [0.0, 0.0, 1.0, 0.0]])) diff --git a/source/isaaclab_newton/test/physics/test_newton_prototype_models.py b/source/isaaclab_newton/test/physics/test_newton_prototype_models.py new file mode 100644 index 000000000000..1a62eda1176b --- /dev/null +++ b/source/isaaclab_newton/test/physics/test_newton_prototype_models.py @@ -0,0 +1,73 @@ +# Copyright (c) 2022-2026, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +from types import SimpleNamespace + +import pytest +from isaaclab_newton.physics.newton_manager import NewtonManager + + +class _FakeBuilder: + def __init__(self): + self.finalize_calls = 0 + + def finalize(self, device=None): + self.finalize_calls += 1 + return SimpleNamespace(device=device) + + +@pytest.fixture(autouse=True) +def _restore_newton_manager_state(): + old_model = NewtonManager._model + old_prototypes = NewtonManager._prototype_models + try: + yield + finally: + NewtonManager._model = old_model + NewtonManager._prototype_models = old_prototypes + + +def test_get_prototype_model_matches_env_regex_path(): + builder = _FakeBuilder() + NewtonManager._model = SimpleNamespace(device="cuda:0") + NewtonManager.register_prototype_builders( + ("/World/envs/env_0",), ("/World/envs/env_{}",), {"/World/envs/env_0": builder} + ) + + info = NewtonManager.get_prototype_model("/World/envs/env_.*/Robot") + + assert info.source_path == "/World/envs/env_0" + assert info.model.device == "cuda:0" + assert builder.finalize_calls == 1 + + +def test_get_prototype_model_reuses_cached_model_on_same_device(): + builder = _FakeBuilder() + NewtonManager._model = SimpleNamespace(device="cpu") + NewtonManager.register_prototype_builders( + ("/World/envs/env_0/Robot",), ("/World/envs/env_{}/Robot",), {"/World/envs/env_0/Robot": builder} + ) + + info_0 = NewtonManager.get_prototype_model("/World/envs/env_.*/Robot") + info_1 = NewtonManager.get_prototype_model("/World/envs/env_.*/Robot") + + assert info_0.model is info_1.model + assert builder.finalize_calls == 1 + + +def test_register_prototype_builders_accumulates_across_calls(): + robot_builder = _FakeBuilder() + prop_builder = _FakeBuilder() + NewtonManager._model = SimpleNamespace(device="cpu") + + NewtonManager.register_prototype_builders( + ("/World/envs/env_0/Robot",), ("/World/envs/env_{}/Robot",), {"/World/envs/env_0/Robot": robot_builder} + ) + NewtonManager.register_prototype_builders( + ("/World/envs/env_0/Prop",), ("/World/envs/env_{}/Prop",), {"/World/envs/env_0/Prop": prop_builder} + ) + + assert NewtonManager.get_prototype_model("/World/envs/env_.*/Robot").builder is robot_builder + assert NewtonManager.get_prototype_model("/World/envs/env_.*/Prop").builder is prop_builder diff --git a/source/isaaclab_tasks/changelog.d/max-newton-ik-manager.rst b/source/isaaclab_tasks/changelog.d/max-newton-ik-manager.rst new file mode 100644 index 000000000000..a16171ca8db8 --- /dev/null +++ b/source/isaaclab_tasks/changelog.d/max-newton-ik-manager.rst @@ -0,0 +1,4 @@ +Added +^^^^^ + +* Added ``Isaac-Reach-Franka-Newton-IK-Rel-v0`` for Newton-backed Franka reach IK. diff --git a/source/isaaclab_tasks/isaaclab_tasks/core/reach/config/franka/__init__.py b/source/isaaclab_tasks/isaaclab_tasks/core/reach/config/franka/__init__.py index 5e28dc985b5c..f6f7a9cc8742 100644 --- a/source/isaaclab_tasks/isaaclab_tasks/core/reach/config/franka/__init__.py +++ b/source/isaaclab_tasks/isaaclab_tasks/core/reach/config/franka/__init__.py @@ -40,6 +40,30 @@ ) +## +# Newton Inverse Kinematics - Relative Pose Control +## + +gym.register( + id="Isaac-Reach-Franka-Newton-IK-Rel-v0", + entry_point="isaaclab.envs:ManagerBasedRLEnv", + kwargs={ + "env_cfg_entry_point": f"{__name__}.ik_newton_env_cfg:FrankaReachEnvCfg", + "rsl_rl_cfg_entry_point": f"{agents.__name__}.rsl_rl_ppo_cfg:FrankaReachNewtonIKPPORunnerCfg", + }, + disable_env_checker=True, +) + +gym.register( + id="Isaac-Reach-Franka-Newton-IK-Rel-Play-v0", + entry_point="isaaclab.envs:ManagerBasedRLEnv", + kwargs={ + "env_cfg_entry_point": f"{__name__}.ik_newton_env_cfg:FrankaReachEnvCfg_PLAY", + "rsl_rl_cfg_entry_point": f"{agents.__name__}.rsl_rl_ppo_cfg:FrankaReachNewtonIKPPORunnerCfg", + }, + disable_env_checker=True, +) + ## # Operational Space Control ## diff --git a/source/isaaclab_tasks/isaaclab_tasks/core/reach/config/franka/agents/rsl_rl_ppo_cfg.py b/source/isaaclab_tasks/isaaclab_tasks/core/reach/config/franka/agents/rsl_rl_ppo_cfg.py index ccce6ac2b32b..29a65e02ee74 100644 --- a/source/isaaclab_tasks/isaaclab_tasks/core/reach/config/franka/agents/rsl_rl_ppo_cfg.py +++ b/source/isaaclab_tasks/isaaclab_tasks/core/reach/config/franka/agents/rsl_rl_ppo_cfg.py @@ -40,3 +40,9 @@ class FrankaReachPPORunnerCfg(RslRlOnPolicyRunnerCfg): desired_kl=0.01, max_grad_norm=1.0, ) + + +@configclass +class FrankaReachNewtonIKPPORunnerCfg(FrankaReachPPORunnerCfg): + experiment_name = "franka_reach_newton_ik_rel" + clip_actions = 1.0 diff --git a/source/isaaclab_tasks/isaaclab_tasks/core/reach/config/franka/ik_newton_env_cfg.py b/source/isaaclab_tasks/isaaclab_tasks/core/reach/config/franka/ik_newton_env_cfg.py new file mode 100644 index 000000000000..40e72bc100fc --- /dev/null +++ b/source/isaaclab_tasks/isaaclab_tasks/core/reach/config/franka/ik_newton_env_cfg.py @@ -0,0 +1,45 @@ +# Copyright (c) 2022-2026, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +from isaaclab_newton.envs.mdp.actions.newton_ik_actions_cfg import NewtonInverseKinematicsActionCfg +from isaaclab_newton.ik.newton_ik_manager_cfg import NewtonIKManagerCfg + +from isaaclab.utils.configclass import configclass + +from . import joint_pos_env_cfg + +## +# Pre-defined configs +## +from isaaclab_assets.robots.franka import FRANKA_PANDA_HIGH_PD_CFG # isort: skip + + +@configclass +class FrankaReachEnvCfg(joint_pos_env_cfg.FrankaReachEnvCfg): + def __post_init__(self): + super().__post_init__() + + # Newton IK consumes the replicated robot prototype. + self.sim.physics = self.sim.physics.newton_mjwarp + self.scene.table = None + self.scene.robot = FRANKA_PANDA_HIGH_PD_CFG.replace(prim_path="{ENV_REGEX_NS}/Robot") + + self.actions.arm_action = NewtonInverseKinematicsActionCfg( + asset_name="robot", + joint_names=["panda_joint.*"], + body_name="panda_hand", + controller=NewtonIKManagerCfg(command_type="pose", use_relative_mode=True), + scale=0.2, + body_offset=NewtonInverseKinematicsActionCfg.OffsetCfg(pos=(0.0, 0.0, 0.107)), + ) + + +@configclass +class FrankaReachEnvCfg_PLAY(FrankaReachEnvCfg): + def __post_init__(self): + super().__post_init__() + self.scene.num_envs = 50 + self.scene.env_spacing = 2.5 + self.observations.policy.enable_corruption = False From 5211091a5a20105d85946a816345e7a34d0d4043 Mon Sep 17 00:00:00 2001 From: Maximilian Krause Date: Thu, 4 Jun 2026 15:59:14 -0700 Subject: [PATCH 2/4] Rename Newton IK manager to solver --- ...k-manager.rst => max-newton-ik-solver.rst} | 2 +- .../envs/mdp/actions/newton_ik_actions.py | 8 +-- .../envs/mdp/actions/newton_ik_actions_cfg.py | 6 +-- .../isaaclab_newton/ik/__init__.pyi | 8 +-- ...wton_ik_manager.py => newton_ik_solver.py} | 22 ++++---- ...manager_cfg.py => newton_ik_solver_cfg.py} | 4 +- ...ik_manager.py => test_newton_ik_solver.py} | 54 +++++++++---------- ...k-manager.rst => max-newton-ik-solver.rst} | 0 .../reach/config/franka/ik_newton_env_cfg.py | 4 +- 9 files changed, 54 insertions(+), 54 deletions(-) rename source/isaaclab_newton/changelog.d/{max-newton-ik-manager.rst => max-newton-ik-solver.rst} (84%) rename source/isaaclab_newton/isaaclab_newton/ik/{newton_ik_manager.py => newton_ik_solver.py} (95%) rename source/isaaclab_newton/isaaclab_newton/ik/{newton_ik_manager_cfg.py => newton_ik_solver_cfg.py} (96%) rename source/isaaclab_newton/test/ik/{test_newton_ik_manager.py => test_newton_ik_solver.py} (67%) rename source/isaaclab_tasks/changelog.d/{max-newton-ik-manager.rst => max-newton-ik-solver.rst} (100%) diff --git a/source/isaaclab_newton/changelog.d/max-newton-ik-manager.rst b/source/isaaclab_newton/changelog.d/max-newton-ik-solver.rst similarity index 84% rename from source/isaaclab_newton/changelog.d/max-newton-ik-manager.rst rename to source/isaaclab_newton/changelog.d/max-newton-ik-solver.rst index 48b597e676ba..06fc1d6b38bb 100644 --- a/source/isaaclab_newton/changelog.d/max-newton-ik-manager.rst +++ b/source/isaaclab_newton/changelog.d/max-newton-ik-solver.rst @@ -1,7 +1,7 @@ Added ^^^^^ -* Added :class:`~isaaclab_newton.ik.NewtonIKManager` and +* Added :class:`~isaaclab_newton.ik.NewtonIKSolver` and :class:`~isaaclab_newton.envs.mdp.actions.NewtonInverseKinematicsAction` for Newton-backed inverse kinematics, including named pose objectives and custom Newton objective passthrough. diff --git a/source/isaaclab_newton/isaaclab_newton/envs/mdp/actions/newton_ik_actions.py b/source/isaaclab_newton/isaaclab_newton/envs/mdp/actions/newton_ik_actions.py index d78802c2cc3e..08d737296ec4 100644 --- a/source/isaaclab_newton/isaaclab_newton/envs/mdp/actions/newton_ik_actions.py +++ b/source/isaaclab_newton/isaaclab_newton/envs/mdp/actions/newton_ik_actions.py @@ -20,7 +20,7 @@ from isaaclab.assets.articulation.base_articulation import BaseArticulation from isaaclab.managers.action_manager import ActionTerm -from isaaclab_newton.ik.newton_ik_manager import NewtonIKManager, NewtonIKPoseObjective +from isaaclab_newton.ik.newton_ik_solver import NewtonIKPoseObjective, NewtonIKSolver from isaaclab_newton.physics import NewtonManager if TYPE_CHECKING: @@ -112,7 +112,7 @@ def __init__(self, cfg: NewtonInverseKinematicsActionCfg, env: ManagerBasedEnv): self._prototype_joint_seed = self._prototype_joint_seed.unsqueeze(0).repeat(self.num_envs, 1).contiguous() self._ik_target_name = "target" - self._ik_manager = NewtonIKManager( + self._ik_solver = NewtonIKSolver( self.cfg.controller, model=prototype_model, num_envs=self.num_envs, @@ -201,11 +201,11 @@ def apply_actions(self) -> None: target_pos_w, target_quat_w = math_utils.combine_frame_transforms( root_pos_proto, root_quat_proto, self._target_pos_b, self._target_quat_b ) - self._ik_manager.set_target_pose(self._ik_target_name, target_pos_w, target_quat_w) + self._ik_solver.set_target_pose(self._ik_target_name, target_pos_w, target_quat_w) joint_seed = self._prototype_joint_seed.clone() joint_seed[:, self._prototype_joint_coord_ids] = self._asset.data.joint_pos.torch - joint_pos_des_all = self._ik_manager.solve(joint_seed) + joint_pos_des_all = self._ik_solver.solve(joint_seed) joint_pos_des = joint_pos_des_all[:, self._prototype_controlled_coord_ids].contiguous() self._asset.set_joint_position_target_index(target=joint_pos_des, joint_ids=self._joint_ids_warp) diff --git a/source/isaaclab_newton/isaaclab_newton/envs/mdp/actions/newton_ik_actions_cfg.py b/source/isaaclab_newton/isaaclab_newton/envs/mdp/actions/newton_ik_actions_cfg.py index 31504333ac5b..051b28ac50f2 100644 --- a/source/isaaclab_newton/isaaclab_newton/envs/mdp/actions/newton_ik_actions_cfg.py +++ b/source/isaaclab_newton/isaaclab_newton/envs/mdp/actions/newton_ik_actions_cfg.py @@ -11,7 +11,7 @@ from isaaclab.managers.action_manager import ActionTermCfg from isaaclab.utils.configclass import configclass -from isaaclab_newton.ik.newton_ik_manager_cfg import NewtonIKManagerCfg +from isaaclab_newton.ik.newton_ik_solver_cfg import NewtonIKSolverCfg if TYPE_CHECKING: from .newton_ik_actions import NewtonInverseKinematicsAction @@ -56,5 +56,5 @@ class OffsetCfg: this is in radians. For quaternions this is dimensionless. """ - controller: NewtonIKManagerCfg = MISSING - """Configuration for the Newton IK manager.""" + controller: NewtonIKSolverCfg = MISSING + """Configuration for the Newton IK solver.""" diff --git a/source/isaaclab_newton/isaaclab_newton/ik/__init__.pyi b/source/isaaclab_newton/isaaclab_newton/ik/__init__.pyi index 2d751bb38eb2..34ad7773f5ac 100644 --- a/source/isaaclab_newton/isaaclab_newton/ik/__init__.pyi +++ b/source/isaaclab_newton/isaaclab_newton/ik/__init__.pyi @@ -4,10 +4,10 @@ # SPDX-License-Identifier: BSD-3-Clause __all__ = [ - "NewtonIKManager", - "NewtonIKManagerCfg", + "NewtonIKSolver", + "NewtonIKSolverCfg", "NewtonIKPoseObjective", ] -from .newton_ik_manager import NewtonIKManager, NewtonIKPoseObjective -from .newton_ik_manager_cfg import NewtonIKManagerCfg +from .newton_ik_solver import NewtonIKPoseObjective, NewtonIKSolver +from .newton_ik_solver_cfg import NewtonIKSolverCfg diff --git a/source/isaaclab_newton/isaaclab_newton/ik/newton_ik_manager.py b/source/isaaclab_newton/isaaclab_newton/ik/newton_ik_solver.py similarity index 95% rename from source/isaaclab_newton/isaaclab_newton/ik/newton_ik_manager.py rename to source/isaaclab_newton/isaaclab_newton/ik/newton_ik_solver.py index f8b02119994e..8aff27f743c3 100644 --- a/source/isaaclab_newton/isaaclab_newton/ik/newton_ik_manager.py +++ b/source/isaaclab_newton/isaaclab_newton/ik/newton_ik_solver.py @@ -14,7 +14,7 @@ import torch import warp as wp -from .newton_ik_manager_cfg import NewtonIKManagerCfg +from .newton_ik_solver_cfg import NewtonIKSolverCfg @dataclass(frozen=True) @@ -26,8 +26,8 @@ class NewtonIKPoseObjective: link_index: Newton body index for the controlled link. link_offset_pos: Target frame translation in meters relative to the link frame. link_offset_rot: Target frame quaternion ``(x, y, z, w)`` relative to the link frame. - position_weight: Optional residual weight overriding the manager default. - rotation_weight: Optional residual weight overriding the manager default. + position_weight: Optional residual weight overriding the solver default. + rotation_weight: Optional residual weight overriding the solver default. """ name: str @@ -38,20 +38,20 @@ class NewtonIKPoseObjective: rotation_weight: float | None = None -class NewtonIKManager: +class NewtonIKSolver: """Batched wrapper around Newton's inverse-kinematics solver. - The manager mirrors Newton's objective-list design while adding torch/Warp + The solver mirrors Newton's objective-list design while adding torch/Warp target updates convenient for Isaac Lab action terms. Pose objectives are named so callers can update individual targets between solves. Additional custom Newton objectives can be passed through ``extra_objectives``. """ - cfg: NewtonIKManagerCfg + cfg: NewtonIKSolverCfg def __init__( self, - cfg: NewtonIKManagerCfg, + cfg: NewtonIKSolverCfg, *, model, num_envs: int, @@ -60,7 +60,7 @@ def __init__( extra_objectives: Sequence[ik.IKObjective] | None = None, ): if not pose_objectives and not extra_objectives: - raise ValueError("NewtonIKManager requires at least one pose or custom objective.") + raise ValueError("NewtonIKSolver requires at least one pose or custom objective.") self.cfg = cfg self.model = model @@ -131,7 +131,7 @@ def __init__( @property def action_dim(self) -> int: - """Dimension of the IK command expected by this manager.""" + """Dimension of the IK command expected by this solver.""" if self.cfg.command_type == "position": return 3 if self.cfg.command_type == "pose" and self.cfg.use_relative_mode: @@ -232,10 +232,10 @@ def joint_q(self) -> wp.array: return self.solver.joint_q def solve(self, joint_pos: torch.Tensor | None = None) -> torch.Tensor: - """Solve IK from an explicit seed or the manager's persistent seed.""" + """Solve IK from an explicit seed or the solver's persistent seed.""" if joint_pos is None: if not self._has_joint_q_seed: - raise RuntimeError("NewtonIKManager.solve() needs joint_pos or a seed set with set_joint_seed().") + raise RuntimeError("NewtonIKSolver.solve() needs joint_pos or a seed set with set_joint_seed().") joint_q_in = self.joint_q_seed update_seed = True else: diff --git a/source/isaaclab_newton/isaaclab_newton/ik/newton_ik_manager_cfg.py b/source/isaaclab_newton/isaaclab_newton/ik/newton_ik_solver_cfg.py similarity index 96% rename from source/isaaclab_newton/isaaclab_newton/ik/newton_ik_manager_cfg.py rename to source/isaaclab_newton/isaaclab_newton/ik/newton_ik_solver_cfg.py index 10fb106ae7ef..5c827a16754e 100644 --- a/source/isaaclab_newton/isaaclab_newton/ik/newton_ik_manager_cfg.py +++ b/source/isaaclab_newton/isaaclab_newton/ik/newton_ik_solver_cfg.py @@ -9,8 +9,8 @@ @configclass -class NewtonIKManagerCfg: - """Configuration for the Newton inverse-kinematics manager.""" +class NewtonIKSolverCfg: + """Configuration for the Newton inverse-kinematics solver.""" command_type: str = "pose" """Action command type for manager-based actions: ``"position"`` or ``"pose"``.""" diff --git a/source/isaaclab_newton/test/ik/test_newton_ik_manager.py b/source/isaaclab_newton/test/ik/test_newton_ik_solver.py similarity index 67% rename from source/isaaclab_newton/test/ik/test_newton_ik_manager.py rename to source/isaaclab_newton/test/ik/test_newton_ik_solver.py index 865e279c6388..8f2a30ab0b77 100644 --- a/source/isaaclab_newton/test/ik/test_newton_ik_manager.py +++ b/source/isaaclab_newton/test/ik/test_newton_ik_solver.py @@ -5,11 +5,11 @@ from __future__ import annotations -import isaaclab_newton.ik.newton_ik_manager as ik_manager_module +import isaaclab_newton.ik.newton_ik_solver as ik_solver_module import torch import warp as wp -from isaaclab_newton.ik.newton_ik_manager import NewtonIKManager, NewtonIKPoseObjective -from isaaclab_newton.ik.newton_ik_manager_cfg import NewtonIKManagerCfg +from isaaclab_newton.ik.newton_ik_solver import NewtonIKPoseObjective, NewtonIKSolver +from isaaclab_newton.ik.newton_ik_solver_cfg import NewtonIKSolverCfg class _Model: @@ -51,17 +51,17 @@ def step(self, joint_q_in, joint_q_out, *, iterations, step_size): def _patch_newton_ik(monkeypatch): - monkeypatch.setattr(ik_manager_module.ik, "IKObjectivePosition", _PoseObjective) - monkeypatch.setattr(ik_manager_module.ik, "IKObjectiveRotation", _PoseObjective) - monkeypatch.setattr(ik_manager_module.ik, "IKObjectiveJointLimit", _JointLimitObjective) - monkeypatch.setattr(ik_manager_module.ik, "IKSolver", _Solver) - monkeypatch.setattr(ik_manager_module.ik, "IKOptimizer", lambda value: value) - monkeypatch.setattr(ik_manager_module.ik, "IKJacobianType", lambda value: value) - monkeypatch.setattr(ik_manager_module.ik, "IKSampler", lambda value: value) + monkeypatch.setattr(ik_solver_module.ik, "IKObjectivePosition", _PoseObjective) + monkeypatch.setattr(ik_solver_module.ik, "IKObjectiveRotation", _PoseObjective) + monkeypatch.setattr(ik_solver_module.ik, "IKObjectiveJointLimit", _JointLimitObjective) + monkeypatch.setattr(ik_solver_module.ik, "IKSolver", _Solver) + monkeypatch.setattr(ik_solver_module.ik, "IKOptimizer", lambda value: value) + monkeypatch.setattr(ik_solver_module.ik, "IKJacobianType", lambda value: value) + monkeypatch.setattr(ik_solver_module.ik, "IKSampler", lambda value: value) -def _cfg() -> NewtonIKManagerCfg: - cfg = NewtonIKManagerCfg() +def _cfg() -> NewtonIKSolverCfg: + cfg = NewtonIKSolverCfg() cfg.use_persistent_seed = True cfg.joint_limit_weight = None return cfg @@ -69,7 +69,7 @@ def _cfg() -> NewtonIKManagerCfg: def test_persistent_seed_is_reused_between_solves(monkeypatch): _patch_newton_ik(monkeypatch) - manager = NewtonIKManager( + solver = NewtonIKSolver( _cfg(), model=_Model(), num_envs=2, @@ -78,10 +78,10 @@ def test_persistent_seed_is_reused_between_solves(monkeypatch): ) seed = torch.tensor([[1.0, 2.0], [3.0, 4.0]]) - manager.set_joint_seed(seed) + solver.set_joint_seed(seed) - first = manager.solve().clone() - second = manager.solve().clone() + first = solver.solve().clone() + second = solver.solve().clone() assert torch.allclose(first, seed + 1.0) assert torch.allclose(second, seed + 2.0) @@ -91,7 +91,7 @@ def test_solve_result_is_independent_from_next_solve(monkeypatch): _patch_newton_ik(monkeypatch) cfg = _cfg() cfg.use_persistent_seed = False - manager = NewtonIKManager( + solver = NewtonIKSolver( cfg, model=_Model(), num_envs=2, @@ -100,16 +100,16 @@ def test_solve_result_is_independent_from_next_solve(monkeypatch): ) seed = torch.tensor([[1.0, 2.0], [3.0, 4.0]]) - first = manager.solve(seed) + first = solver.solve(seed) expected_first = first.clone() - manager.solve(seed + 10.0) + solver.solve(seed + 10.0) assert torch.allclose(first, expected_first) def test_set_target_pose_updates_named_objective(monkeypatch): _patch_newton_ik(monkeypatch) - manager = NewtonIKManager( + solver = NewtonIKSolver( _cfg(), model=_Model(), num_envs=2, @@ -119,15 +119,15 @@ def test_set_target_pose_updates_named_objective(monkeypatch): target_pos = torch.tensor([[0.1, 0.2, 0.3], [1.0, 1.1, 1.2]]) target_quat = torch.tensor([[0.0, 0.0, 0.0, 1.0], [0.0, 0.0, 1.0, 0.0]]) - manager.set_target_pose("ee", target_pos, target_quat) + solver.set_target_pose("ee", target_pos, target_quat) - assert torch.allclose(wp.to_torch(manager.position_objectives["ee"].target), target_pos) - assert torch.allclose(wp.to_torch(manager.rotation_objectives["ee"].target), target_quat) + assert torch.allclose(wp.to_torch(solver.position_objectives["ee"].target), target_pos) + assert torch.allclose(wp.to_torch(solver.rotation_objectives["ee"].target), target_quat) def test_pose_targets_can_be_initialized_from_body_transforms(monkeypatch): _patch_newton_ik(monkeypatch) - manager = NewtonIKManager( + solver = NewtonIKSolver( _cfg(), model=_Model(), num_envs=2, @@ -145,9 +145,9 @@ def test_pose_targets_can_be_initialized_from_body_transforms(monkeypatch): ) env_origins = torch.tensor([[0.0, 0.0, 0.0], [1.0, 1.0, 1.0]]) - manager.set_pose_targets_from_body_q(body_q, names=["torso"], env_origins=env_origins) + solver.set_pose_targets_from_body_q(body_q, names=["torso"], env_origins=env_origins) - target_pos = wp.to_torch(manager.position_objectives["torso"].target) - target_quat = wp.to_torch(manager.rotation_objectives["torso"].target) + target_pos = wp.to_torch(solver.position_objectives["torso"].target) + target_quat = wp.to_torch(solver.rotation_objectives["torso"].target) assert torch.allclose(target_pos, torch.tensor([[4.0, 5.0, 6.0], [3.0, 4.0, 5.0]])) assert torch.allclose(target_quat, torch.tensor([[0.0, 0.0, 1.0, 0.0], [0.0, 0.0, 1.0, 0.0]])) diff --git a/source/isaaclab_tasks/changelog.d/max-newton-ik-manager.rst b/source/isaaclab_tasks/changelog.d/max-newton-ik-solver.rst similarity index 100% rename from source/isaaclab_tasks/changelog.d/max-newton-ik-manager.rst rename to source/isaaclab_tasks/changelog.d/max-newton-ik-solver.rst diff --git a/source/isaaclab_tasks/isaaclab_tasks/core/reach/config/franka/ik_newton_env_cfg.py b/source/isaaclab_tasks/isaaclab_tasks/core/reach/config/franka/ik_newton_env_cfg.py index 40e72bc100fc..fb465b30e970 100644 --- a/source/isaaclab_tasks/isaaclab_tasks/core/reach/config/franka/ik_newton_env_cfg.py +++ b/source/isaaclab_tasks/isaaclab_tasks/core/reach/config/franka/ik_newton_env_cfg.py @@ -4,7 +4,7 @@ # SPDX-License-Identifier: BSD-3-Clause from isaaclab_newton.envs.mdp.actions.newton_ik_actions_cfg import NewtonInverseKinematicsActionCfg -from isaaclab_newton.ik.newton_ik_manager_cfg import NewtonIKManagerCfg +from isaaclab_newton.ik.newton_ik_solver_cfg import NewtonIKSolverCfg from isaaclab.utils.configclass import configclass @@ -30,7 +30,7 @@ def __post_init__(self): asset_name="robot", joint_names=["panda_joint.*"], body_name="panda_hand", - controller=NewtonIKManagerCfg(command_type="pose", use_relative_mode=True), + controller=NewtonIKSolverCfg(command_type="pose", use_relative_mode=True), scale=0.2, body_offset=NewtonInverseKinematicsActionCfg.OffsetCfg(pos=(0.0, 0.0, 0.107)), ) From 023975a0768817fcfb84954b076a5eb4bf222b8d Mon Sep 17 00:00:00 2001 From: Octi Zhang Date: Thu, 11 Jun 2026 17:04:11 -0700 Subject: [PATCH 3/4] add improvements --- .../isaaclab_newton/cloner/replicate.py | 7 +- .../envs/mdp/actions/newton_ik_actions.py | 282 ++++++++---------- .../envs/mdp/actions/newton_ik_actions_cfg.py | 50 ++-- .../isaaclab_newton/ik/__init__.pyi | 13 +- .../ik/newton_ik_objectives.py | 145 +++++++++ .../ik/newton_ik_objectives_cfg.py | 103 +++++++ .../isaaclab_newton/ik/newton_ik_solver.py | 268 +++-------------- .../ik/newton_ik_solver_cfg.py | 33 +- .../isaaclab_newton/physics/newton_manager.py | 108 +------ .../test/ik/test_newton_ik_solver.py | 190 +++++++----- .../physics/test_newton_prototype_models.py | 73 ----- .../reach/config/franka/ik_newton_env_cfg.py | 16 +- 12 files changed, 614 insertions(+), 674 deletions(-) create mode 100644 source/isaaclab_newton/isaaclab_newton/ik/newton_ik_objectives.py create mode 100644 source/isaaclab_newton/isaaclab_newton/ik/newton_ik_objectives_cfg.py delete mode 100644 source/isaaclab_newton/test/physics/test_newton_prototype_models.py diff --git a/source/isaaclab_newton/isaaclab_newton/cloner/replicate.py b/source/isaaclab_newton/isaaclab_newton/cloner/replicate.py index 10506d06c3a2..04bde19f7a8c 100644 --- a/source/isaaclab_newton/isaaclab_newton/cloner/replicate.py +++ b/source/isaaclab_newton/isaaclab_newton/cloner/replicate.py @@ -46,8 +46,9 @@ def _build_newton_builder_from_mapping( ) -> tuple[ModelBuilder, object, dict, list, dict[str, ModelBuilder]]: """Build a Newton model builder from clone mapping inputs. - Returns the populated builder, stage metadata, the site index map, the - per-world transforms, and the per-source builders keyed by clone source path. + Also returns the per-source builders (``{source_path: ModelBuilder}``) so the + committing path can retain them for single-model consumers such as the + batched Newton IK action. """ if positions is None: positions = torch.zeros((mapping.size(1), 3), device=mapping.device, dtype=torch.float32) @@ -228,7 +229,7 @@ def replicate(self) -> tuple[ModelBuilder, object, dict]: NewtonManager._cl_site_index_map = site_index_map NewtonManager._cl_fabric_body_bindings = fabric_body_bindings NewtonManager._world_xforms = world_xforms - NewtonManager.register_prototype_builders(sources, destinations, source_builders) + NewtonManager._cl_protos = source_builders NewtonManager.set_builder(builder) NewtonManager._num_envs = mapping.size(1) self._queue.clear() diff --git a/source/isaaclab_newton/isaaclab_newton/envs/mdp/actions/newton_ik_actions.py b/source/isaaclab_newton/isaaclab_newton/envs/mdp/actions/newton_ik_actions.py index 08d737296ec4..8be7f052e061 100644 --- a/source/isaaclab_newton/isaaclab_newton/envs/mdp/actions/newton_ik_actions.py +++ b/source/isaaclab_newton/isaaclab_newton/envs/mdp/actions/newton_ik_actions.py @@ -7,6 +7,7 @@ import logging from collections.abc import Sequence +from dataclasses import dataclass, field from typing import TYPE_CHECKING import torch @@ -15,31 +16,63 @@ from newton import Model as NewtonModel from newton.selection import ArticulationView +import isaaclab.sim as sim_utils import isaaclab.utils.math as math_utils import isaaclab.utils.string as string_utils from isaaclab.assets.articulation.base_articulation import BaseArticulation +from isaaclab.cloner import resolve_clone_plan_source from isaaclab.managers.action_manager import ActionTerm -from isaaclab_newton.ik.newton_ik_solver import NewtonIKPoseObjective, NewtonIKSolver +from isaaclab_newton.ik.newton_ik_objectives_cfg import NewtonIKPoseObjectiveCfg from isaaclab_newton.physics import NewtonManager if TYPE_CHECKING: from isaaclab.envs import ManagerBasedEnv from isaaclab.envs.utils.io_descriptors import GenericActionIODescriptor + from isaaclab_newton.ik.newton_ik_objectives import NewtonIKPoseObjective + from .newton_ik_actions_cfg import NewtonInverseKinematicsActionCfg logger = logging.getLogger(__name__) +@dataclass +class _PoseDriver: + """Per-pose-objective binding to the live articulation and the action vector. + + Holds the Isaac Lab body index used to read the current end-effector pose, the + target-frame offset (batched to ``num_envs``), this objective's slice of the + action vector, and the built solver objective. The per-step root-frame target + buffers are allocated in :meth:`__post_init__` from the offset's shape/device. + """ + + body_idx: int + offset_pos: torch.Tensor + offset_rot: torch.Tensor + slice: slice + objective: NewtonIKPoseObjective + target_pos_b: torch.Tensor = field(init=False) + target_quat_b: torch.Tensor = field(init=False) + + def __post_init__(self): + num_envs, device = self.offset_pos.shape[0], self.offset_pos.device + self.target_pos_b = torch.zeros(num_envs, 3, device=device) + self.target_quat_b = torch.zeros(num_envs, 4, device=device) + self.target_quat_b[:, 3] = 1.0 + + class NewtonInverseKinematicsAction(ActionTerm): """Newton inverse-kinematics action term. - This action currently supports fixed-base articulations only. It solves IK - on the single-environment Newton prototype model registered by the cloner - and maps the resulting actuated joint coordinates back to the live batched - Isaac Lab articulation. + The action solves IK as a single list of objectives on the single-env Newton + prototype model registered by the cloner, then maps the resulting actuated + joint coordinates back to the live batched Isaac Lab articulation. Each pose + objective contributes its slice of the action vector and drives one + end-effector body; a single pose objective is single-body IK, several are + multi-body IK. Constraint objectives (e.g. joint limits) add no action + dimensions. Fixed-base articulations only. """ cfg: NewtonInverseKinematicsActionCfg @@ -48,102 +81,89 @@ class NewtonInverseKinematicsAction(ActionTerm): def __init__(self, cfg: NewtonInverseKinematicsActionCfg, env: ManagerBasedEnv): super().__init__(cfg, env) - if not isinstance(self._asset, BaseArticulation): - raise TypeError( - f"NewtonInverseKinematicsAction expects a BaseArticulation asset, got {type(self._asset).__name__}." - ) if not self._asset.is_fixed_base: raise ValueError("NewtonInverseKinematicsAction currently supports fixed-base articulations only.") self._joint_ids, self._joint_names = self._asset.find_joints(self.cfg.joint_names) - if len(self._joint_ids) == 0: - raise ValueError(f"No joints matched Newton IK action joint_names={self.cfg.joint_names}.") self._joint_ids_warp = wp.array(self._joint_ids, dtype=wp.int32, device=self.device) - body_ids, body_names = self._asset.find_bodies(self.cfg.body_name) - if len(body_ids) != 1: - raise ValueError( - f"Expected one match for Newton IK body_name={self.cfg.body_name}. Found {len(body_ids)}: {body_names}." - ) - self._body_idx = body_ids[0] - self._body_name = body_names[0] - - self._raw_actions = torch.zeros(self.num_envs, self.action_dim, device=self.device) - self._processed_actions = torch.zeros_like(self._raw_actions) - self._target_pos_b = torch.zeros(self.num_envs, 3, device=self.device) - self._target_quat_b = torch.zeros(self.num_envs, 4, device=self.device) - self._target_quat_b[:, 3] = 1.0 - - self._scale = torch.zeros((self.num_envs, self.action_dim), device=self.device) - self._scale[:] = torch.tensor(self.cfg.scale, device=self.device) - - self._clip = None - if self.cfg.clip is not None: - self._clip = torch.tensor([[-float("inf"), float("inf")]], device=self.device).repeat( - self.num_envs, self.action_dim, 1 - ) - action_names = self._action_coordinate_names() - index_list, _, value_list = string_utils.resolve_matching_names_values(self.cfg.clip, action_names) - self._clip[:, index_list] = torch.tensor(value_list, device=self.device) - - if self.cfg.body_offset is not None: - self._offset_pos = torch.tensor(self.cfg.body_offset.pos, device=self.device).repeat(self.num_envs, 1) - self._offset_rot = torch.tensor(self.cfg.body_offset.rot, device=self.device).repeat(self.num_envs, 1) - link_offset_pos = tuple(self.cfg.body_offset.pos) - link_offset_rot = tuple(self.cfg.body_offset.rot) - else: - self._offset_pos, self._offset_rot = None, None - link_offset_pos = (0.0, 0.0, 0.0) - link_offset_rot = (0.0, 0.0, 0.0, 1.0) - - prototype_info = NewtonManager.get_prototype_model(self._asset.cfg.prim_path) - prototype_model = prototype_info.model - if prototype_model is None: - raise RuntimeError(f"Newton prototype model for '{self._asset.cfg.prim_path}' was not finalized.") - prototype_view = self._resolve_prototype_view(prototype_model) + pose_cfgs = [obj for obj in self.cfg.objectives if isinstance(obj, NewtonIKPoseObjectiveCfg)] + if not pose_cfgs: + raise ValueError("NewtonInverseKinematicsAction requires at least one pose objective.") + + # Resolve the controlled asset to its clone-plan source and finalize the + # single-env prototype builder the cloner already retained -- the same + # source resolution other Newton consumers use, no bespoke registry. + plan = sim_utils.SimulationContext.instance().get_clone_plan() + source_path, _, asset_suffix = resolve_clone_plan_source(self._asset.cfg.prim_path, plan) + # The proto builder is keyed by the bare clone source; the articulation + # lives at the asset suffix below it (e.g. ".../env_0" + "/Robot"). + self._source_path = source_path + asset_suffix + prototype_model = NewtonManager._cl_protos[source_path].finalize(device=NewtonManager.get_model().device) + # The prototype is the cloner's env_0 source builder; its single articulation + # is addressed by the resolved source path. + prototype_view = ArticulationView( + prototype_model, + self._source_path, + verbose=False, + exclude_joint_types=[JointType.FREE, JointType.FIXED], + ) self._prototype_joint_coord_ids = self._resolve_prototype_joint_coord_ids( prototype_view, self._asset.joint_names ) self._prototype_controlled_coord_ids = self._resolve_prototype_joint_coord_ids( prototype_view, self._joint_names ) - self._prototype_link_index = self._resolve_prototype_link_index(prototype_view) self._prototype_joint_seed = wp.to_torch(prototype_model.joint_q).to(device=self.device, dtype=torch.float32) self._prototype_joint_seed = self._prototype_joint_seed.unsqueeze(0).repeat(self.num_envs, 1).contiguous() - self._ik_target_name = "target" - self._ik_solver = NewtonIKSolver( + # The solver resolves each pose objective's body via the prototype view. + self._ik_solver = self.cfg.controller.class_type( self.cfg.controller, model=prototype_model, num_envs=self.num_envs, device=self.device, - pose_objectives=[ - NewtonIKPoseObjective( - name=self._ik_target_name, - link_index=self._prototype_link_index, - link_offset_pos=link_offset_pos, - link_offset_rot=link_offset_rot, - ) - ], + objectives=self.cfg.objectives, + link_resolver=lambda body_name: self._resolve_prototype_link_index(prototype_view, body_name), ) + # Bind each pose objective to the live articulation and the action vector. + self._drivers: list[_PoseDriver] = [] + offset = 0 + for pose_cfg in pose_cfgs: + name = pose_cfg.name if pose_cfg.name is not None else pose_cfg.body_name + objective = self._ik_solver.objectives_by_name[name] + body_idx = self._resolve_isaac_body_index(pose_cfg.body_name) + offset_pos = torch.tensor(pose_cfg.body_offset_pos, device=self.device).repeat(self.num_envs, 1) + offset_rot = torch.tensor(pose_cfg.body_offset_rot, device=self.device).repeat(self.num_envs, 1) + self._drivers.append( + _PoseDriver(body_idx, offset_pos, offset_rot, slice(offset, offset + objective.action_dim), objective) + ) + offset += objective.action_dim + self._action_dim = offset + + self._raw_actions = torch.zeros(self.num_envs, self._action_dim, device=self.device) + self._processed_actions = torch.zeros_like(self._raw_actions) + + self._clip = None + if self.cfg.clip is not None: + self._clip = torch.tensor([[-float("inf"), float("inf")]], device=self.device).repeat( + self.num_envs, self._action_dim, 1 + ) + action_names = self._action_coordinate_names() + index_list, _, value_list = string_utils.resolve_matching_names_values(self.cfg.clip, action_names) + self._clip[:, index_list] = torch.tensor(value_list, device=self.device) + logger.info( - "Resolved Newton IK action joints %s [%s] and body %s [%s].", + "Resolved Newton IK action joints %s [%s] and bodies %s.", self._joint_names, self._joint_ids, - self._body_name, - self._body_idx, + [(d.objective.name, d.body_idx) for d in self._drivers], ) @property def action_dim(self) -> int: - if self.cfg.controller.command_type == "position": - return 3 - if self.cfg.controller.command_type == "pose" and self.cfg.controller.use_relative_mode: - return 6 - if self.cfg.controller.command_type == "pose": - return 7 - raise ValueError(f"Unsupported Newton IK command type: {self.cfg.controller.command_type}") + return self._action_dim @property def raw_actions(self) -> torch.Tensor: @@ -156,41 +176,32 @@ def processed_actions(self) -> torch.Tensor: @property def IO_descriptor(self) -> GenericActionIODescriptor: super().IO_descriptor - self._IO_descriptor.shape = (self.action_dim,) + self._IO_descriptor.shape = (self._action_dim,) self._IO_descriptor.dtype = str(self.raw_actions.dtype) self._IO_descriptor.action_type = "NewtonInverseKinematicsAction" - self._IO_descriptor.body_name = self._body_name self._IO_descriptor.joint_names = self._joint_names - self._IO_descriptor.scale = self._scale self._IO_descriptor.clip = self.cfg.clip self._IO_descriptor.extras["controller_cfg"] = self.cfg.controller.__dict__ - self._IO_descriptor.extras["body_offset"] = ( - None if self.cfg.body_offset is None else self.cfg.body_offset.__dict__ - ) + self._IO_descriptor.extras["objective_names"] = [d.objective.name for d in self._drivers] + self._IO_descriptor.extras["coordinate_names"] = self._action_coordinate_names() return self._IO_descriptor def process_actions(self, actions: torch.Tensor) -> None: self._raw_actions[:] = actions - self._processed_actions[:] = self.raw_actions * self._scale + self._processed_actions[:] = self._raw_actions if self._clip is not None: self._processed_actions = torch.clamp( self._processed_actions, min=self._clip[:, :, 0], max=self._clip[:, :, 1] ) - - ee_pos_b, ee_quat_b = self._compute_frame_pose() - if self.cfg.controller.command_type == "position": - if self.cfg.controller.use_relative_mode: - self._target_pos_b[:] = ee_pos_b + self._processed_actions - else: - self._target_pos_b[:] = self._processed_actions - self._target_quat_b[:] = ee_quat_b - elif self.cfg.controller.use_relative_mode: - self._target_pos_b[:], self._target_quat_b[:] = math_utils.apply_delta_pose( - ee_pos_b, ee_quat_b, self._processed_actions + # Each pose objective maps its own action slice (scaled internally) onto a + # root-frame target from the body's current pose. + for driver in self._drivers: + ee_pos_b, ee_quat_b = self._compute_frame_pose(driver) + target_pos_b, target_quat_b = driver.objective.compute_target_b( + self._processed_actions[:, driver.slice], ee_pos_b, ee_quat_b ) - else: - self._target_pos_b[:] = self._processed_actions[:, 0:3] - self._target_quat_b[:] = self._processed_actions[:, 3:7] + driver.target_pos_b[:] = target_pos_b + driver.target_quat_b[:] = target_quat_b def apply_actions(self) -> None: # The IK solve runs on the single-env prototype model, so all batched @@ -198,14 +209,15 @@ def apply_actions(self) -> None: self._validate_matching_root_orientations() root_pos_proto = self._asset.data.root_pos_w.torch[0:1].repeat(self.num_envs, 1) root_quat_proto = self._asset.data.root_quat_w.torch[0:1].repeat(self.num_envs, 1) - target_pos_w, target_quat_w = math_utils.combine_frame_transforms( - root_pos_proto, root_quat_proto, self._target_pos_b, self._target_quat_b - ) - self._ik_solver.set_target_pose(self._ik_target_name, target_pos_w, target_quat_w) + for driver in self._drivers: + target_pos_w, target_quat_w = math_utils.combine_frame_transforms( + root_pos_proto, root_quat_proto, driver.target_pos_b, driver.target_quat_b + ) + driver.objective.set_target_pose(target_pos_w, target_quat_w) joint_seed = self._prototype_joint_seed.clone() joint_seed[:, self._prototype_joint_coord_ids] = self._asset.data.joint_pos.torch - joint_pos_des_all = self._ik_solver.solve(joint_seed) + joint_pos_des_all = wp.to_torch(self._ik_solver.solve(wp.from_torch(joint_seed.contiguous(), dtype=wp.float32))) joint_pos_des = joint_pos_des_all[:, self._prototype_controlled_coord_ids].contiguous() self._asset.set_joint_position_target_index(target=joint_pos_des, joint_ids=self._joint_ids_warp) @@ -213,16 +225,15 @@ def reset(self, env_ids: Sequence[int] | None = None) -> None: env_ids = slice(None) if env_ids is None else env_ids self._raw_actions[env_ids] = 0.0 - def _compute_frame_pose(self) -> tuple[torch.Tensor, torch.Tensor]: - ee_pos_w = self._asset.data.body_pos_w.torch[:, self._body_idx] - ee_quat_w = self._asset.data.body_quat_w.torch[:, self._body_idx] + def _compute_frame_pose(self, driver: _PoseDriver) -> tuple[torch.Tensor, torch.Tensor]: + ee_pos_w = self._asset.data.body_pos_w.torch[:, driver.body_idx] + ee_quat_w = self._asset.data.body_quat_w.torch[:, driver.body_idx] root_pos_w = self._asset.data.root_pos_w.torch root_quat_w = self._asset.data.root_quat_w.torch ee_pos_b, ee_quat_b = math_utils.subtract_frame_transforms(root_pos_w, root_quat_w, ee_pos_w, ee_quat_w) - if self.cfg.body_offset is not None: - ee_pos_b, ee_quat_b = math_utils.combine_frame_transforms( - ee_pos_b, ee_quat_b, self._offset_pos, self._offset_rot - ) + ee_pos_b, ee_quat_b = math_utils.combine_frame_transforms( + ee_pos_b, ee_quat_b, driver.offset_pos, driver.offset_rot + ) return ee_pos_b, ee_quat_b def _validate_matching_root_orientations(self) -> None: @@ -240,33 +251,13 @@ def _validate_matching_root_orientations(self) -> None: "for this action." ) - def _resolve_prototype_view(self, model) -> ArticulationView: - requested_path = NewtonManager._to_first_env_path(self._asset.cfg.prim_path) - basename = requested_path.rsplit("/", 1)[-1] - patterns = [requested_path, requested_path.replace(".*", "*"), basename, f"*{basename}"] - last_error: Exception | None = None - for pattern in patterns: - try: - view = ArticulationView( - model, - pattern, - verbose=False, - exclude_joint_types=[JointType.FREE, JointType.FIXED], - ) - except (KeyError, ValueError) as exc: - last_error = exc - continue - if view.world_count != 1 or view.count_per_world != 1: - raise ValueError( - f"Newton IK expected one prototype articulation for '{self._asset.cfg.prim_path}', " - f"got world_count={view.world_count}, count_per_world={view.count_per_world}." - ) - if not view.is_fixed_base: - raise ValueError("Newton IK currently supports fixed-base prototype articulations only.") - return view - raise KeyError( - f"Failed to resolve Newton prototype articulation for '{self._asset.cfg.prim_path}'." - ) from last_error + def _resolve_isaac_body_index(self, body_name: str) -> int: + body_ids, body_names = self._asset.find_bodies(body_name) + if len(body_ids) != 1: + raise ValueError( + f"Expected one match for Newton IK body_name={body_name}. Found {len(body_ids)}: {body_names}." + ) + return body_ids[0] def _resolve_prototype_joint_coord_ids( self, prototype_view: ArticulationView, joint_names: Sequence[str] @@ -276,23 +267,13 @@ def _resolve_prototype_joint_coord_ids( coord_indices_by_name = { name: layout.offset + selected_indices[index] for index, name in enumerate(prototype_view.joint_coord_names) } - try: - coord_ids = [coord_indices_by_name[name] for name in joint_names] - except KeyError as exc: - raise KeyError( - f"Joint '{exc.args[0]}' was resolved in Isaac Lab but not in the Newton prototype model." - ) from exc + coord_ids = [coord_indices_by_name[name] for name in joint_names] return torch.tensor(coord_ids, device=self.device, dtype=torch.long) - def _resolve_prototype_link_index(self, prototype_view: ArticulationView) -> int: + def _resolve_prototype_link_index(self, prototype_view: ArticulationView, body_name: str) -> int: layout = prototype_view.frequency_layouts[NewtonModel.AttributeFrequency.BODY] selected_indices = self._layout_indices(layout) - try: - local_link_index = prototype_view.link_names.index(self._body_name) - except ValueError as exc: - raise ValueError( - f"Body '{self._body_name}' was resolved in Isaac Lab but not in the Newton prototype model." - ) from exc + local_link_index = prototype_view.link_names.index(body_name) return layout.offset + selected_indices[local_link_index] @staticmethod @@ -302,8 +283,7 @@ def _layout_indices(layout) -> list[int]: return [int(index) for index in layout.indices.numpy().tolist()] def _action_coordinate_names(self) -> list[str]: - if self.cfg.controller.command_type == "position": - return ["x", "y", "z"] - if self.cfg.controller.command_type == "pose" and self.cfg.controller.use_relative_mode: - return ["x", "y", "z", "roll", "pitch", "yaw"] - return ["x", "y", "z", "qx", "qy", "qz", "qw"] + names: list[str] = [] + for driver in self._drivers: + names.extend(f"{driver.objective.name}/{coord}" for coord in driver.objective.command_coordinate_names()) + return names diff --git a/source/isaaclab_newton/isaaclab_newton/envs/mdp/actions/newton_ik_actions_cfg.py b/source/isaaclab_newton/isaaclab_newton/envs/mdp/actions/newton_ik_actions_cfg.py index 051b28ac50f2..e8330d62699e 100644 --- a/source/isaaclab_newton/isaaclab_newton/envs/mdp/actions/newton_ik_actions_cfg.py +++ b/source/isaaclab_newton/isaaclab_newton/envs/mdp/actions/newton_ik_actions_cfg.py @@ -5,12 +5,13 @@ from __future__ import annotations -from dataclasses import MISSING +from dataclasses import MISSING, field from typing import TYPE_CHECKING from isaaclab.managers.action_manager import ActionTermCfg from isaaclab.utils.configclass import configclass +from isaaclab_newton.ik.newton_ik_objectives_cfg import NewtonIKObjectiveCfg from isaaclab_newton.ik.newton_ik_solver_cfg import NewtonIKSolverCfg if TYPE_CHECKING: @@ -21,40 +22,33 @@ class NewtonInverseKinematicsActionCfg(ActionTermCfg): """Configuration for a Newton inverse-kinematics action term. - The action currently supports fixed-base articulations only. The configured - joints and body must resolve both in Isaac Lab and in the registered Newton - prototype model for the controlled asset. + The action solves IK as a single list of objectives. Pose objectives + (:class:`~isaaclab_newton.ik.NewtonIKPoseObjectiveCfg`) are command-driven + and contribute action dimensions -- one drives a single-body solve, several + drive a multi-body solve. Constraint objectives such as + :class:`~isaaclab_newton.ik.NewtonIKJointLimitObjectiveCfg` add residuals + but no action dimensions. The action vector is the concatenation of every + pose objective's slice, in list order. + + The action currently supports fixed-base articulations only. Each pose + objective's body and the configured joints must resolve both in Isaac Lab + and in the registered Newton prototype model for the controlled asset. """ - @configclass - class OffsetCfg: - """Offset pose from the controlled body frame to the IK target frame.""" - - pos: tuple[float, float, float] = (0.0, 0.0, 0.0) - """Translation [m] w.r.t. the parent body frame.""" - - rot: tuple[float, float, float, float] = (0.0, 0.0, 0.0, 1.0) - """Quaternion rotation ``(x, y, z, w)`` w.r.t. the parent body frame.""" - class_type: type[NewtonInverseKinematicsAction] | str = ( "isaaclab_newton.envs.mdp.actions.newton_ik_actions:NewtonInverseKinematicsAction" ) joint_names: list[str] = MISSING - """List of joint names or regex expressions controlled by the action.""" - - body_name: str = MISSING - """Name of the body for which IK is performed.""" + """Joints actuated by the action. - body_offset: OffsetCfg | None = None - """Offset of the target frame w.r.t. the body frame.""" - - scale: float | tuple[float, ...] = 1.0 - """Scale factor applied to the raw action. - - For position coordinates this is in meters. For relative rotation coordinates - this is in radians. For quaternions this is dimensionless. + The Newton solve resolves the whole prototype joint configuration jointly + against all objectives, so this is the single set of joints written back to + the articulation -- not a per-objective property. """ - controller: NewtonIKSolverCfg = MISSING - """Configuration for the Newton IK solver.""" + objectives: list[NewtonIKObjectiveCfg] = MISSING + """Ordered IK objectives. Must contain at least one pose objective.""" + + controller: NewtonIKSolverCfg = field(default_factory=NewtonIKSolverCfg) + """Configuration for the Newton IK solver (solver hyperparameters only).""" diff --git a/source/isaaclab_newton/isaaclab_newton/ik/__init__.pyi b/source/isaaclab_newton/isaaclab_newton/ik/__init__.pyi index 34ad7773f5ac..364300e05fcb 100644 --- a/source/isaaclab_newton/isaaclab_newton/ik/__init__.pyi +++ b/source/isaaclab_newton/isaaclab_newton/ik/__init__.pyi @@ -6,8 +6,19 @@ __all__ = [ "NewtonIKSolver", "NewtonIKSolverCfg", + "NewtonIKObjective", "NewtonIKPoseObjective", + "NewtonIKJointLimitObjective", + "NewtonIKObjectiveCfg", + "NewtonIKPoseObjectiveCfg", + "NewtonIKJointLimitObjectiveCfg", ] -from .newton_ik_solver import NewtonIKPoseObjective, NewtonIKSolver +from .newton_ik_objectives import NewtonIKJointLimitObjective, NewtonIKObjective, NewtonIKPoseObjective +from .newton_ik_objectives_cfg import ( + NewtonIKJointLimitObjectiveCfg, + NewtonIKObjectiveCfg, + NewtonIKPoseObjectiveCfg, +) +from .newton_ik_solver import NewtonIKSolver from .newton_ik_solver_cfg import NewtonIKSolverCfg diff --git a/source/isaaclab_newton/isaaclab_newton/ik/newton_ik_objectives.py b/source/isaaclab_newton/isaaclab_newton/ik/newton_ik_objectives.py new file mode 100644 index 000000000000..11262ddb3172 --- /dev/null +++ b/source/isaaclab_newton/isaaclab_newton/ik/newton_ik_objectives.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 + +"""Runtime Newton IK objective implementations. + +Each class is built by :class:`~isaaclab_newton.ik.NewtonIKSolver` from the +matching :class:`~isaaclab_newton.ik.newton_ik_objectives_cfg.NewtonIKObjectiveCfg` +and owns the concrete :class:`newton.ik.IKObjective` instances appended to the +solver. Command-driven objectives (currently pose) additionally own their +action contribution: an :attr:`~NewtonIKObjective.action_dim`, the coordinate +names for that slice, and the mapping from a raw action slice to a target pose. +The action term is therefore generic -- it sums :attr:`action_dim` across the +objective list and dispatches each slice without branching on command type. + +Importing this module pulls ``newton`` (and ``pxr``), so it is loaded lazily via +the package ``lazy_export`` only after Kit has launched. Custom objectives +integrate by subclassing :class:`NewtonIKObjective`, taking +``(cfg, ctx)`` in ``__init__`` -- pulling only the :class:`NewtonIKBuildContext` +fields they need -- and populating :attr:`NewtonIKObjective.solver_objectives`. +""" + +from __future__ import annotations + +from collections.abc import Callable +from dataclasses import dataclass + +import newton.ik as ik +import torch +import warp as wp + +import isaaclab.utils.math as math_utils + +from .newton_ik_objectives_cfg import NewtonIKJointLimitObjectiveCfg, NewtonIKPoseObjectiveCfg + + +@dataclass(frozen=True) +class NewtonIKBuildContext: + """Build-time inputs shared with every objective; each pulls what it needs.""" + + model: object + """Finalized Newton prototype model (e.g. for joint limits).""" + + num_envs: int + """Number of parallel IK problems (target-array batch size).""" + + device: str + """Warp device string for objective-owned arrays.""" + + resolve_link: Callable[[str], int] + """Maps a body name to its Newton link index in the prototype model.""" + + +class NewtonIKObjective: + """Base built IK objective. + + Owns the concrete :class:`newton.ik.IKObjective` instances in + :attr:`solver_objectives`. Command-driven objectives (pose) also set a + :attr:`name` and a non-zero :attr:`action_dim` and add ``compute_target_b`` + / ``command_coordinate_names``; constraint objectives leave the defaults. + """ + + name: str | None = None + """Unique objective name, or ``None`` when the objective has no runtime target.""" + + action_dim: int = 0 + """Number of action coordinates this objective consumes (0 for constraints).""" + + solver_objectives: list[ik.IKObjective] + """Concrete Newton objectives appended to the solver's objective list.""" + + +class NewtonIKPoseObjective(NewtonIKObjective): + """Command-driven position + rotation objective tracking one end-effector body.""" + + def __init__(self, cfg: NewtonIKPoseObjectiveCfg, ctx: NewtonIKBuildContext): + self.name = cfg.name if cfg.name is not None else cfg.body_name + self.command_type = cfg.command_type + self.use_relative_mode = cfg.use_relative_mode + self.link_index = ctx.resolve_link(cfg.body_name) + self.action_dim = len(self.command_coordinate_names()) + + scale = torch.as_tensor(cfg.scale, dtype=torch.float32, device=ctx.device) + if scale.ndim == 0: + scale = scale.repeat(self.action_dim) + if scale.shape != (self.action_dim,): + raise ValueError( + f"Newton IK pose objective '{self.name}' scale must be a float or length-{self.action_dim} " + f"sequence, got shape {tuple(scale.shape)}." + ) + self._scale = scale + + target_positions = wp.zeros((ctx.num_envs,), dtype=wp.vec3, device=ctx.device) + target_rotations = wp.array([(0.0, 0.0, 0.0, 1.0)] * ctx.num_envs, dtype=wp.vec4, device=ctx.device) + self.position_objective = ik.IKObjectivePosition( + link_index=self.link_index, + link_offset=wp.vec3(*cfg.body_offset_pos), + target_positions=target_positions, + weight=cfg.position_weight, + ) + self.rotation_objective = ik.IKObjectiveRotation( + link_index=self.link_index, + link_offset_rotation=wp.quat(*cfg.body_offset_rot), + target_rotations=target_rotations, + weight=cfg.rotation_weight, + ) + self.solver_objectives = [self.position_objective, self.rotation_objective] + + def command_coordinate_names(self) -> list[str]: + if self.command_type == "position": + return ["x", "y", "z"] + if self.command_type == "pose": + if self.use_relative_mode: + return ["x", "y", "z", "roll", "pitch", "yaw"] + return ["x", "y", "z", "qx", "qy", "qz", "qw"] + raise ValueError(f"Unsupported Newton IK command type: {self.command_type}") + + def compute_target_b( + self, action: torch.Tensor, ee_pos_b: torch.Tensor, ee_quat_b: torch.Tensor + ) -> tuple[torch.Tensor, torch.Tensor]: + processed = action * self._scale + if self.command_type == "position": + target_pos_b = ee_pos_b + processed if self.use_relative_mode else processed + return target_pos_b, ee_quat_b + if self.use_relative_mode: + return math_utils.apply_delta_pose(ee_pos_b, ee_quat_b, processed) + return processed[:, 0:3], processed[:, 3:7] + + def set_target_pose(self, target_pos_w: torch.Tensor, target_quat_w: torch.Tensor) -> None: + """Update the batched world-frame target pose, shape ``[num_envs, 3]`` and ``[num_envs, 4]``.""" + self.position_objective.set_target_positions(wp.from_torch(target_pos_w.contiguous(), dtype=wp.vec3)) + self.rotation_objective.set_target_rotations(wp.from_torch(target_quat_w.contiguous(), dtype=wp.vec4)) + + +class NewtonIKJointLimitObjective(NewtonIKObjective): + """Soft joint-limit constraint reading the model's coordinate limits.""" + + def __init__(self, cfg: NewtonIKJointLimitObjectiveCfg, ctx: NewtonIKBuildContext): + self.objective = ik.IKObjectiveJointLimit( + joint_limit_lower=ctx.model.joint_limit_lower, + joint_limit_upper=ctx.model.joint_limit_upper, + weight=cfg.weight, + ) + self.solver_objectives = [self.objective] diff --git a/source/isaaclab_newton/isaaclab_newton/ik/newton_ik_objectives_cfg.py b/source/isaaclab_newton/isaaclab_newton/ik/newton_ik_objectives_cfg.py new file mode 100644 index 000000000000..244113606db5 --- /dev/null +++ b/source/isaaclab_newton/isaaclab_newton/ik/newton_ik_objectives_cfg.py @@ -0,0 +1,103 @@ +# 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 + +"""Pure-dataclass configuration for Newton IK objectives. + +Inverse kinematics is expressed as a single ordered list of objectives. Each +entry is a :class:`NewtonIKObjectiveCfg` subclass describing one constraint: +a tracked end-effector pose, a joint limit, a collision penalty, and so on. +There is no separate notion of a "target" versus an "extra" objective -- a +single-body solve is a one-element list, a multi-body solve simply has more +pose entries, and constraints like joint limits are further entries that add +no action dimensions. + +The solver builds each objective via +``string_to_callable(cfg.class_type)(cfg, model=..., num_envs=..., device=..., +link_resolver=...)`` and the action term reads each built objective's action +contribution, so an objective owns both its Newton residual and how a policy +command (if any) maps onto its target. + +This module imports only the standard library and Isaac Lab's config utilities; +it must stay free of ``import newton`` so action/env configs remain importable +before Kit has launched. Matching runtime implementations live in +:mod:`isaaclab_newton.ik.newton_ik_objectives`. +""" + +from __future__ import annotations + +from dataclasses import MISSING + +from isaaclab.utils.configclass import configclass + + +@configclass +class NewtonIKObjectiveCfg: + """Base configuration for a Newton IK objective. + + Subclasses set :attr:`class_type` to the runtime implementation, which the + solver instantiates as + ``class_type(cfg, model=..., num_envs=..., device=..., link_resolver=...)``. + The implementation exposes the concrete :class:`newton.ik.IKObjective` + instances appended to the solver and, for command-driven objectives, an + action-dimension contribution. + """ + + class_type: type | str = MISSING # type: ignore[assignment] + """Runtime objective implementation, as a type or a ``"module:Class"`` string.""" + + +@configclass +class NewtonIKPoseObjectiveCfg(NewtonIKObjectiveCfg): + """A pose objective tracking one end-effector body. + + This is the command-driven objective: it contributes action dimensions + (3 for ``"position"``, 6 for relative ``"pose"``, 7 for absolute ``"pose"``) + and maps its slice of the policy action onto a target pose for + :attr:`body_name`. Multiple pose objectives drive a multi-body solve, each + with its own body, command convention, weights and scale. + """ + + class_type: type | str = "isaaclab_newton.ik.newton_ik_objectives:NewtonIKPoseObjective" + + body_name: str = MISSING # type: ignore[assignment] + """Name of the controlled end-effector body.""" + + name: str | None = None + """Unique objective name used to update its target. Defaults to :attr:`body_name`.""" + + body_offset_pos: tuple[float, float, float] = (0.0, 0.0, 0.0) + """Target-frame translation [m] relative to the body frame.""" + + body_offset_rot: tuple[float, float, float, float] = (0.0, 0.0, 0.0, 1.0) + """Target-frame quaternion ``(x, y, z, w)`` relative to the body frame.""" + + command_type: str = "pose" + """How the policy action is interpreted: ``"position"`` or ``"pose"``.""" + + use_relative_mode: bool = True + """Whether the command is a delta from the current end-effector pose.""" + + scale: float | tuple[float, ...] = 1.0 + """Scale applied to this objective's raw action slice [m for position, rad for relative rotation].""" + + position_weight: float = 1.0 + """Residual weight [unitless] for the position component.""" + + rotation_weight: float = 1.0 + """Residual weight [unitless] for the rotation component.""" + + +@configclass +class NewtonIKJointLimitObjectiveCfg(NewtonIKObjectiveCfg): + """Soft joint-limit constraint penalizing coordinates outside the model limits. + + A constraint-only objective: it adds a Newton residual but no action + dimensions. + """ + + class_type: type | str = "isaaclab_newton.ik.newton_ik_objectives:NewtonIKJointLimitObjective" + + weight: float = 0.1 + """Residual weight [unitless] applied to limit violations.""" diff --git a/source/isaaclab_newton/isaaclab_newton/ik/newton_ik_solver.py b/source/isaaclab_newton/isaaclab_newton/ik/newton_ik_solver.py index 8aff27f743c3..1dc35a4dbb01 100644 --- a/source/isaaclab_newton/isaaclab_newton/ik/newton_ik_solver.py +++ b/source/isaaclab_newton/isaaclab_newton/ik/newton_ik_solver.py @@ -5,46 +5,34 @@ from __future__ import annotations -from collections.abc import Sequence -from dataclasses import dataclass -from numbers import Integral -from typing import Any +from collections.abc import Callable, Sequence import newton.ik as ik -import torch import warp as wp +from .newton_ik_objectives import NewtonIKBuildContext, NewtonIKObjective +from .newton_ik_objectives_cfg import NewtonIKObjectiveCfg from .newton_ik_solver_cfg import NewtonIKSolverCfg -@dataclass(frozen=True) -class NewtonIKPoseObjective: - """Pose objective descriptor used to build Newton position/rotation objectives. - - Args: - name: Unique objective name used when updating targets. - link_index: Newton body index for the controlled link. - link_offset_pos: Target frame translation in meters relative to the link frame. - link_offset_rot: Target frame quaternion ``(x, y, z, w)`` relative to the link frame. - position_weight: Optional residual weight overriding the solver default. - rotation_weight: Optional residual weight overriding the solver default. - """ - - name: str - link_index: int - link_offset_pos: tuple[float, float, float] = (0.0, 0.0, 0.0) - link_offset_rot: tuple[float, float, float, float] = (0.0, 0.0, 0.0, 1.0) - position_weight: float | None = None - rotation_weight: float | None = None - - class NewtonIKSolver: """Batched wrapper around Newton's inverse-kinematics solver. - The solver mirrors Newton's objective-list design while adding torch/Warp - target updates convenient for Isaac Lab action terms. Pose objectives are - named so callers can update individual targets between solves. Additional - custom Newton objectives can be passed through ``extra_objectives``. + The solver is configured by an ordered list of + :class:`~isaaclab_newton.ik.newton_ik_objectives_cfg.NewtonIKObjectiveCfg`. + Each cfg is resolved to its runtime + :class:`~isaaclab_newton.ik.newton_ik_objectives.NewtonIKObjective` and its + concrete Newton objectives are appended to the underlying + :class:`newton.ik.IKSolver`. The built objectives are exposed via + :attr:`objectives` / :attr:`objectives_by_name`; callers update a pose + objective's target by calling :meth:`~..newton_ik_objectives.NewtonIKPoseObjective.set_target_pose` + on it directly. + + The solver solves ``num_envs`` independent problems and is agnostic to how + targets are produced -- the prototype-broadcast policy used by the Newton IK + action term lives in the action, not here. ``link_resolver`` maps an + objective's body name to a Newton link index; the caller owns it because the + name-to-index mapping depends on the model layout (e.g. cloned env prefixes). """ cfg: NewtonIKSolverCfg @@ -56,66 +44,28 @@ def __init__( model, num_envs: int, device: str, - pose_objectives: Sequence[NewtonIKPoseObjective], - extra_objectives: Sequence[ik.IKObjective] | None = None, + objectives: Sequence[NewtonIKObjectiveCfg], + link_resolver: Callable[[str], int], ): - if not pose_objectives and not extra_objectives: - raise ValueError("NewtonIKSolver requires at least one pose or custom objective.") + if not objectives: + raise ValueError("NewtonIKSolver requires at least one objective cfg.") self.cfg = cfg - self.model = model - self.num_envs = num_envs - self.device = device - self.num_coords = model.joint_coord_count - self.pose_objective_names = [objective.name for objective in pose_objectives] - self.pose_objective_cfgs = {objective.name: objective for objective in pose_objectives} - - if len(set(self.pose_objective_names)) != len(self.pose_objective_names): - raise ValueError(f"Newton IK pose objective names must be unique: {self.pose_objective_names}") + ctx = NewtonIKBuildContext(model=model, num_envs=num_envs, device=device, resolve_link=link_resolver) - self.position_objectives: dict[str, ik.IKObjectivePosition] = {} - self.rotation_objectives: dict[str, ik.IKObjectiveRotation] = {} + self.objectives: list[NewtonIKObjective] = [] + self.objectives_by_name: dict[str, NewtonIKObjective] = {} solver_objectives: list[ik.IKObjective] = [] - - for objective_cfg in pose_objectives: - target_positions = wp.zeros((num_envs,), dtype=wp.vec3, device=device) - target_rotations = wp.array( - [(0.0, 0.0, 0.0, 1.0)] * num_envs, - dtype=wp.vec4, - device=device, - ) - position_objective = ik.IKObjectivePosition( - link_index=objective_cfg.link_index, - link_offset=wp.vec3(*objective_cfg.link_offset_pos), - target_positions=target_positions, - weight=cfg.position_weight if objective_cfg.position_weight is None else objective_cfg.position_weight, - ) - rotation_objective = ik.IKObjectiveRotation( - link_index=objective_cfg.link_index, - link_offset_rotation=wp.quat(*objective_cfg.link_offset_rot), - target_rotations=target_rotations, - weight=cfg.rotation_weight if objective_cfg.rotation_weight is None else objective_cfg.rotation_weight, - ) - self.position_objectives[objective_cfg.name] = position_objective - self.rotation_objectives[objective_cfg.name] = rotation_objective - solver_objectives.extend((position_objective, rotation_objective)) - - if cfg.joint_limit_weight is not None: - self.joint_limit_objective = ik.IKObjectiveJointLimit( - joint_limit_lower=model.joint_limit_lower, - joint_limit_upper=model.joint_limit_upper, - weight=cfg.joint_limit_weight, - ) - solver_objectives.append(self.joint_limit_objective) - else: - self.joint_limit_objective = None - - solver_objectives.extend(extra_objectives or []) - - self.joint_q_out = wp.zeros((num_envs, self.num_coords), dtype=wp.float32, device=device) - self.joint_q_seed_t = torch.zeros((num_envs, self.num_coords), dtype=torch.float32, device=device) - self.joint_q_seed = wp.from_torch(self.joint_q_seed_t, dtype=wp.float32) - self._has_joint_q_seed = False + for objective_cfg in objectives: + objective = objective_cfg.class_type(objective_cfg, ctx) + if objective.name is not None: + if objective.name in self.objectives_by_name: + raise ValueError(f"Newton IK objective names must be unique: duplicate '{objective.name}'.") + self.objectives_by_name[objective.name] = objective + self.objectives.append(objective) + solver_objectives.extend(objective.solver_objectives) + + self.joint_q_out = wp.zeros((num_envs, model.joint_coord_count), dtype=wp.float32, device=device) self.solver = ik.IKSolver( model=model, n_problems=num_envs, @@ -129,98 +79,6 @@ def __init__( lambda_initial=cfg.lambda_initial, ) - @property - def action_dim(self) -> int: - """Dimension of the IK command expected by this solver.""" - if self.cfg.command_type == "position": - return 3 - if self.cfg.command_type == "pose" and self.cfg.use_relative_mode: - return 6 - if self.cfg.command_type == "pose": - return 7 - raise ValueError(f"Unsupported Newton IK command type: {self.cfg.command_type}") - - def set_target_pose(self, name: str, target_pos_w: torch.Tensor, target_quat_w: torch.Tensor) -> None: - """Update batched world-frame target poses for a named pose objective.""" - try: - position_objective = self.position_objectives[name] - rotation_objective = self.rotation_objectives[name] - except KeyError as exc: - raise KeyError( - f"Unknown Newton IK pose objective '{name}'. Available objectives: {self.pose_objective_names}." - ) from exc - position_objective.set_target_positions(wp.from_torch(target_pos_w.contiguous(), dtype=wp.vec3)) - rotation_objective.set_target_rotations(wp.from_torch(target_quat_w.contiguous(), dtype=wp.vec4)) - - def set_target_pose_from_body_q( - self, - name: str, - body_q: torch.Tensor | wp.array | Any, - *, - env_origins: torch.Tensor | None = None, - ) -> None: - """Set one pose objective target from body transforms.""" - objective_cfg = self.pose_objective_cfgs[name] - body_q_t = _as_torch(body_q, device=self.device) - if body_q_t.ndim == 2: - target_pos = body_q_t[objective_cfg.link_index, :3].reshape(1, 3).repeat(self.num_envs, 1) - target_quat = body_q_t[objective_cfg.link_index, 3:7].reshape(1, 4).repeat(self.num_envs, 1) - elif body_q_t.ndim == 3: - if body_q_t.shape[0] != self.num_envs: - raise ValueError(f"Expected body_q first dimension {self.num_envs}, got {body_q_t.shape[0]}.") - target_pos = body_q_t[:, objective_cfg.link_index, :3] - target_quat = body_q_t[:, objective_cfg.link_index, 3:7] - else: - raise ValueError( - f"Expected body_q shape (num_bodies, 7) or (num_envs, num_bodies, 7), got {body_q_t.shape}." - ) - if env_origins is not None: - target_pos = target_pos - env_origins.to(device=self.device, dtype=torch.float32) - self.set_target_pose(name, target_pos.to(torch.float32), target_quat.to(torch.float32)) - - def set_pose_targets_from_body_q( - self, - body_q: torch.Tensor | wp.array | Any, - names: Sequence[str] | None = None, - *, - env_origins: torch.Tensor | None = None, - ) -> None: - """Set multiple pose objective targets from body transforms.""" - for name in self.pose_objective_names if names is None else names: - self.set_target_pose_from_body_q(name, body_q, env_origins=env_origins) - - def set_joint_seed( - self, joint_pos: torch.Tensor, env_ids: Sequence[int] | torch.Tensor | slice | None = None - ) -> None: - """Set the persistent joint-coordinate seed used by ``solve()`` without an explicit seed.""" - joint_pos = joint_pos.to(device=self.device, dtype=torch.float32) - if env_ids is None: - if joint_pos.shape != (self.num_envs, self.num_coords): - raise ValueError( - f"Expected joint seed shape {(self.num_envs, self.num_coords)}, got {tuple(joint_pos.shape)}." - ) - self.joint_q_seed_t.copy_(joint_pos) - else: - ids = _env_ids_to_tensor(env_ids, self.num_envs, self.device) - if joint_pos.shape != (ids.numel(), self.num_coords): - raise ValueError( - f"Expected joint seed shape {(ids.numel(), self.num_coords)}, got {tuple(joint_pos.shape)}." - ) - self.joint_q_seed_t[ids] = joint_pos - self._has_joint_q_seed = True - - def reset( - self, - env_ids: Sequence[int] | torch.Tensor | slice | None = None, - joint_pos: torch.Tensor | None = None, - ) -> None: - """Reset Newton solver state, selected seeds, and sampler RNG.""" - self.solver.reset() - if joint_pos is not None: - self.set_joint_seed(joint_pos, env_ids=env_ids) - elif env_ids is None: - self._has_joint_q_seed = False - @property def costs(self) -> wp.array: """Expanded per-seed costs from the most recent Newton solve.""" @@ -231,56 +89,16 @@ def joint_q(self) -> wp.array: """Expanded joint-coordinate buffer storing all sampled seeds.""" return self.solver.joint_q - def solve(self, joint_pos: torch.Tensor | None = None) -> torch.Tensor: - """Solve IK from an explicit seed or the solver's persistent seed.""" - if joint_pos is None: - if not self._has_joint_q_seed: - raise RuntimeError("NewtonIKSolver.solve() needs joint_pos or a seed set with set_joint_seed().") - joint_q_in = self.joint_q_seed - update_seed = True - else: - if joint_pos.shape != (self.num_envs, self.num_coords): - raise ValueError( - f"Expected joint seed shape {(self.num_envs, self.num_coords)}, got {tuple(joint_pos.shape)}." - ) - if self.cfg.use_persistent_seed: - self.set_joint_seed(joint_pos) - joint_q_in = self.joint_q_seed - update_seed = True - else: - joint_q_in = wp.from_torch(joint_pos.contiguous(), dtype=wp.float32) - update_seed = False + def solve(self, joint_pos: wp.array) -> wp.array: + """Solve IK from the Warp seed ``joint_pos``, shape ``[num_envs, joint_coord_count]``. + + Returns the solver's output buffer, overwritten on the next solve -- consume + or copy it before solving again. + """ self.solver.step( - joint_q_in, + joint_pos, self.joint_q_out, iterations=self.cfg.iterations, step_size=self.cfg.step_size, ) - result = wp.to_torch(self.joint_q_out).clone() - if update_seed: - self.joint_q_seed_t.copy_(result) - self._has_joint_q_seed = True - return result - - def step(self) -> torch.Tensor: - """Solve IK from the persistent seed and store the result as the next seed.""" - return self.solve() - - -def _as_torch(value, *, device: str) -> torch.Tensor: - if isinstance(value, torch.Tensor): - return value.to(device=device, dtype=torch.float32) - if hasattr(value, "numpy"): - return wp.to_torch(value).to(device=device, dtype=torch.float32) - return torch.as_tensor(value, device=device, dtype=torch.float32) - - -def _env_ids_to_tensor(env_ids: Sequence[int] | torch.Tensor | slice, num_envs: int, device: str) -> torch.Tensor: - if isinstance(env_ids, slice): - start, stop, step = env_ids.indices(num_envs) - return torch.arange(start, stop, step, device=device, dtype=torch.long) - if isinstance(env_ids, Integral): - return torch.as_tensor([env_ids], device=device, dtype=torch.long) - if isinstance(env_ids, torch.Tensor): - return env_ids.to(device=device, dtype=torch.long).flatten() - return torch.as_tensor(list(env_ids), device=device, dtype=torch.long) + return self.joint_q_out diff --git a/source/isaaclab_newton/isaaclab_newton/ik/newton_ik_solver_cfg.py b/source/isaaclab_newton/isaaclab_newton/ik/newton_ik_solver_cfg.py index 5c827a16754e..520433b856c3 100644 --- a/source/isaaclab_newton/isaaclab_newton/ik/newton_ik_solver_cfg.py +++ b/source/isaaclab_newton/isaaclab_newton/ik/newton_ik_solver_cfg.py @@ -10,13 +10,24 @@ @configclass class NewtonIKSolverCfg: - """Configuration for the Newton inverse-kinematics solver.""" + """Configuration for the Newton inverse-kinematics solver. - command_type: str = "pose" - """Action command type for manager-based actions: ``"position"`` or ``"pose"``.""" + Holds solver hyperparameters only. Objectives (and their residual weights) + are configured separately as a list of + :class:`~isaaclab_newton.ik.newton_ik_objectives_cfg.NewtonIKObjectiveCfg` + passed to the solver. Command semantics for manager-based actions + (``command_type``, ``use_relative_mode``) live on the action cfg. - use_relative_mode: bool = True - """Whether manager-based action commands are relative to the current target pose.""" + :attr:`class_type` selects the solver implementation, so an alternative + solver can be dropped in via config without changing callers. + """ + + class_type: type | str = "isaaclab_newton.ik.newton_ik_solver:NewtonIKSolver" + """Solver implementation, as a type or a ``"module:Class"`` string. + + Instantiated as ``class_type(cfg, model=..., num_envs=..., device=..., + objectives=..., link_resolver=...)``. + """ optimizer: str = "lm" """Newton IK optimizer backend. Supported values are ``"lm"`` and ``"lbfgs"``.""" @@ -53,15 +64,3 @@ class NewtonIKSolverCfg: This moderate default favors stable updates near singular configurations over aggressive first-step motion. """ - - position_weight: float = 1.0 - """Default residual weight for pose position objectives.""" - - rotation_weight: float = 1.0 - """Default residual weight for pose rotation objectives.""" - - joint_limit_weight: float | None = 0.1 - """Residual weight for the joint-limit objective. Set to ``None`` to disable it.""" - - use_persistent_seed: bool = False - """Whether solved joint coordinates should be reused as the next IK seed.""" diff --git a/source/isaaclab_newton/isaaclab_newton/physics/newton_manager.py b/source/isaaclab_newton/isaaclab_newton/physics/newton_manager.py index 83164e5f27db..67ed27f21d89 100644 --- a/source/isaaclab_newton/isaaclab_newton/physics/newton_manager.py +++ b/source/isaaclab_newton/isaaclab_newton/physics/newton_manager.py @@ -62,18 +62,6 @@ logger = logging.getLogger(__name__) - -@dataclass -class NewtonPrototypeModelInfo: - """Cached Newton prototype model built before physics replication.""" - - source_path: str - destination_path: str - builder: ModelBuilder - model: Model | None = None - model_device: object | None = None - - # Tagged union for entries in _cl_site_index_map. # _GlobalSite: (global_shape_idx, None) — body_pattern was None # _LocalSite: (None, [[env0_idx, ...], ...]) — per-world site indices @@ -239,7 +227,6 @@ class NewtonManager(PhysicsManager): # Newton model and state _builder: ModelBuilder = None _model: Model = None - _prototype_models: dict[str, NewtonPrototypeModelInfo] = {} _solver: SolverBase | None = None _use_single_state: bool | None = None """Use only one state for both input and output for solver stepping. Requires solver support.""" @@ -324,6 +311,10 @@ class NewtonManager(PhysicsManager): _cl_site_index_map: dict[str, _SiteEntry] = {} _cl_fabric_body_bindings: list[tuple[str, int]] | None = None _world_xforms: list[wp.transform] | None = None + # Per-source builders retained from replication, keyed by clone-plan source + # path. Single-model consumers (e.g. batched Newton IK) finalize a single-env + # model from these and resolve it via ``resolve_clone_plan_source``. + _cl_protos: dict[str, ModelBuilder] = {} _deformable_registry: list = [] _per_world_builder_hooks: list[Callable[[ModelBuilder, int, list[float], list[float]], None]] = [] _post_replicate_hooks: list[Callable[[ModelBuilder], None]] = [] @@ -795,13 +786,13 @@ def clear(cls): NewtonManager._up_axis = "Z" NewtonManager._scene_data = None NewtonManager._scene_data_mapping = None - NewtonManager._prototype_models = {} NewtonManager._model_changes = set() NewtonManager._scene_data_backend = None NewtonManager._cl_pending_sites = {} NewtonManager._cl_site_index_map = {} NewtonManager._cl_fabric_body_bindings = None NewtonManager._world_xforms = None + NewtonManager._cl_protos = {} NewtonManager._pending_extended_state_attributes = set() NewtonManager._pending_extended_contact_attributes = set() NewtonManager._views = [] @@ -812,92 +803,6 @@ def set_builder(cls, builder: ModelBuilder) -> None: """Set the Newton model builder.""" NewtonManager._builder = builder - @classmethod - def register_prototype_builders( - cls, - sources: list[str] | tuple[str, ...], - destinations: list[str] | tuple[str, ...], - proto_builders: dict[str, ModelBuilder], - ) -> None: - """Register prototype builders created before Newton physics replication. - - Prototype builders preserve the single-source model that is later added - into each replicated Newton world. Controllers that need a single-model - representation, such as batched Newton IK, should use this registry - instead of re-importing USD after the scene has already been built. - """ - for index, source_path in enumerate(sources): - builder = proto_builders.get(source_path) - if builder is None: - continue - destination_path = destinations[index] if index < len(destinations) else destinations[-1] - NewtonManager._prototype_models[source_path] = NewtonPrototypeModelInfo( - source_path=source_path, - destination_path=destination_path, - builder=builder, - ) - - @staticmethod - def _to_first_env_path(path: str) -> str: - """Convert common Isaac Lab env regex/template paths to env_0 paths.""" - return ( - path.replace("env_.*", "env_0") - .replace("env_\\d+", "env_0") - .replace("env_*", "env_0") - .replace("env_{}", "env_0") - ) - - @classmethod - def get_prototype_model(cls, prim_path: str) -> NewtonPrototypeModelInfo: - """Return the prototype model that owns ``prim_path``. - - Args: - prim_path: Live Isaac Lab prim path or env-regex path, such as - ``/World/envs/env_.*/Robot``. - - Returns: - Prototype model information with ``model`` finalized on the active - Newton device. - - Raises: - RuntimeError: If no prototype builders have been registered. - KeyError: If ``prim_path`` does not resolve to a registered prototype. - ValueError: If multiple registered prototypes match equally well. - """ - if not NewtonManager._prototype_models: - raise RuntimeError( - "No Newton prototype builders are registered. Prototype access is only available after " - "Newton physics replication has built the scene." - ) - - requested_path = cls._to_first_env_path(prim_path) - matches: list[NewtonPrototypeModelInfo] = [] - for info in NewtonManager._prototype_models.values(): - source_path = cls._to_first_env_path(info.source_path) - destination_path = cls._to_first_env_path(info.destination_path) - if requested_path in (source_path, destination_path) or requested_path.startswith( - (source_path + "/", destination_path + "/") - ): - matches.append(info) - - if not matches: - available = ", ".join(sorted(NewtonManager._prototype_models)) - raise KeyError(f"No Newton prototype model matches '{prim_path}'. Available prototypes: {available}") - - matches.sort(key=lambda item: len(cls._to_first_env_path(item.source_path)), reverse=True) - if len(matches) > 1: - best_len = len(cls._to_first_env_path(matches[0].source_path)) - if len(cls._to_first_env_path(matches[1].source_path)) == best_len: - paths = ", ".join(match.source_path for match in matches) - raise ValueError(f"Ambiguous Newton prototype model for '{prim_path}'. Matches: {paths}") - - info = matches[0] - device = cls._model.device if cls._model is not None else None - if info.model is None or info.model_device != device: - info.model = info.builder.finalize(device=device) - info.model_device = device - return info - @classmethod def create_builder(cls, up_axis: str | None = None, **kwargs) -> ModelBuilder: """Create a :class:`ModelBuilder` configured with default settings. @@ -1342,7 +1247,7 @@ def instantiate_builder_from_stage(cls): _, proto_path = env_paths[0] source_builders = {proto_path: cls.create_builder(up_axis=up_axis)} source_builders[proto_path].add_usd(stage, root_path=proto_path, schema_resolvers=schema_resolvers) - cls.register_prototype_builders((proto_path,), (proto_path,), source_builders) + cls._cl_protos = source_builders global_site_indices, source_site_indices, env_root_sites = cls._cl_inject_sites(builder, source_builders) xform_cache = UsdGeom.XformCache() @@ -1962,6 +1867,7 @@ def update_visualization_state(cls, scene_data_provider: SceneDataProvider | Non Invoked lazily from :meth:`get_state` so consumers do not need to coordinate the sync explicitly. """ + if scene_data_provider is None: scene_data_provider = cls.get_scene_data_provider() diff --git a/source/isaaclab_newton/test/ik/test_newton_ik_solver.py b/source/isaaclab_newton/test/ik/test_newton_ik_solver.py index 8f2a30ab0b77..041cc8b46ee1 100644 --- a/source/isaaclab_newton/test/ik/test_newton_ik_solver.py +++ b/source/isaaclab_newton/test/ik/test_newton_ik_solver.py @@ -5,17 +5,34 @@ from __future__ import annotations +import isaaclab_newton.ik.newton_ik_objectives as objectives_module import isaaclab_newton.ik.newton_ik_solver as ik_solver_module import torch import warp as wp -from isaaclab_newton.ik.newton_ik_solver import NewtonIKPoseObjective, NewtonIKSolver +from isaaclab_newton.ik.newton_ik_objectives import NewtonIKBuildContext, NewtonIKObjective, NewtonIKPoseObjective +from isaaclab_newton.ik.newton_ik_objectives_cfg import ( + NewtonIKJointLimitObjectiveCfg, + NewtonIKObjectiveCfg, + NewtonIKPoseObjectiveCfg, +) +from isaaclab_newton.ik.newton_ik_solver import NewtonIKSolver from isaaclab_newton.ik.newton_ik_solver_cfg import NewtonIKSolverCfg +from isaaclab.utils.configclass import configclass + +# Maps the stub body names used across these tests to Newton link indices. +_LINKS = {"ee": 0, "torso": 1, "custom": 0} + + +def _resolver(body_name: str) -> int: + return _LINKS[body_name] + class _Model: joint_coord_count = 2 joint_limit_lower = None joint_limit_upper = None + body_label = None class _PoseObjective: @@ -51,103 +68,134 @@ def step(self, joint_q_in, joint_q_out, *, iterations, step_size): def _patch_newton_ik(monkeypatch): - monkeypatch.setattr(ik_solver_module.ik, "IKObjectivePosition", _PoseObjective) - monkeypatch.setattr(ik_solver_module.ik, "IKObjectiveRotation", _PoseObjective) - monkeypatch.setattr(ik_solver_module.ik, "IKObjectiveJointLimit", _JointLimitObjective) + # ``newton.ik`` is the same module object in both the solver and objective + # modules, so patching its attributes affects objective construction too. + monkeypatch.setattr(objectives_module.ik, "IKObjectivePosition", _PoseObjective) + monkeypatch.setattr(objectives_module.ik, "IKObjectiveRotation", _PoseObjective) + monkeypatch.setattr(objectives_module.ik, "IKObjectiveJointLimit", _JointLimitObjective) monkeypatch.setattr(ik_solver_module.ik, "IKSolver", _Solver) monkeypatch.setattr(ik_solver_module.ik, "IKOptimizer", lambda value: value) monkeypatch.setattr(ik_solver_module.ik, "IKJacobianType", lambda value: value) monkeypatch.setattr(ik_solver_module.ik, "IKSampler", lambda value: value) -def _cfg() -> NewtonIKSolverCfg: - cfg = NewtonIKSolverCfg() - cfg.use_persistent_seed = True - cfg.joint_limit_weight = None - return cfg - - -def test_persistent_seed_is_reused_between_solves(monkeypatch): - _patch_newton_ik(monkeypatch) - solver = NewtonIKSolver( - _cfg(), +def _pose_solver(cfg: NewtonIKSolverCfg | None = None, num_envs: int = 2, objectives=None) -> NewtonIKSolver: + return NewtonIKSolver( + NewtonIKSolverCfg() if cfg is None else cfg, model=_Model(), - num_envs=2, + num_envs=num_envs, device="cpu", - pose_objectives=[NewtonIKPoseObjective(name="ee", link_index=0)], + objectives=[NewtonIKPoseObjectiveCfg(body_name="ee")] if objectives is None else objectives, + link_resolver=_resolver, ) - seed = torch.tensor([[1.0, 2.0], [3.0, 4.0]]) - solver.set_joint_seed(seed) - - first = solver.solve().clone() - second = solver.solve().clone() - - assert torch.allclose(first, seed + 1.0) - assert torch.allclose(second, seed + 2.0) - -def test_solve_result_is_independent_from_next_solve(monkeypatch): +def test_solve_writes_output_buffer(monkeypatch): _patch_newton_ik(monkeypatch) - cfg = _cfg() - cfg.use_persistent_seed = False - solver = NewtonIKSolver( - cfg, - model=_Model(), - num_envs=2, - device="cpu", - pose_objectives=[NewtonIKPoseObjective(name="ee", link_index=0)], - ) - seed = torch.tensor([[1.0, 2.0], [3.0, 4.0]]) + solver = _pose_solver() + seed = wp.from_torch(torch.tensor([[1.0, 2.0], [3.0, 4.0]]), dtype=wp.float32) - first = solver.solve(seed) - expected_first = first.clone() - solver.solve(seed + 10.0) + result = solver.solve(seed) - assert torch.allclose(first, expected_first) + assert torch.allclose(wp.to_torch(result), torch.tensor([[2.0, 3.0], [4.0, 5.0]])) def test_set_target_pose_updates_named_objective(monkeypatch): _patch_newton_ik(monkeypatch) - solver = NewtonIKSolver( - _cfg(), - model=_Model(), - num_envs=2, - device="cpu", - pose_objectives=[NewtonIKPoseObjective(name="ee", link_index=0)], - ) + solver = _pose_solver() target_pos = torch.tensor([[0.1, 0.2, 0.3], [1.0, 1.1, 1.2]]) target_quat = torch.tensor([[0.0, 0.0, 0.0, 1.0], [0.0, 0.0, 1.0, 0.0]]) - solver.set_target_pose("ee", target_pos, target_quat) + pose = solver.objectives_by_name["ee"] + pose.set_target_pose(target_pos, target_quat) - assert torch.allclose(wp.to_torch(solver.position_objectives["ee"].target), target_pos) - assert torch.allclose(wp.to_torch(solver.rotation_objectives["ee"].target), target_quat) + assert torch.allclose(wp.to_torch(pose.position_objective.target), target_pos) + assert torch.allclose(wp.to_torch(pose.rotation_objective.target), target_quat) -def test_pose_targets_can_be_initialized_from_body_transforms(monkeypatch): +def test_constraint_objectives_carry_no_target_or_action(monkeypatch): _patch_newton_ik(monkeypatch) - solver = NewtonIKSolver( - _cfg(), - model=_Model(), - num_envs=2, - device="cpu", - pose_objectives=[ - NewtonIKPoseObjective(name="ee", link_index=0), - NewtonIKPoseObjective(name="torso", link_index=1, position_weight=50.0, rotation_weight=50.0), - ], + solver = _pose_solver( + objectives=[ + NewtonIKPoseObjectiveCfg(body_name="ee"), + NewtonIKJointLimitObjectiveCfg(weight=0.1), + ] ) - body_q = torch.tensor( - [ - [1.0, 2.0, 3.0, 0.0, 0.0, 0.0, 1.0], - [4.0, 5.0, 6.0, 0.0, 0.0, 1.0, 0.0], + # Only the pose objective is named and command-driven; the joint limit is a + # pure constraint (no name, no action dimensions). + assert list(solver.objectives_by_name) == ["ee"] + assert [obj.action_dim for obj in solver.objectives] == [6, 0] + + +def test_multiple_pose_objectives_register_distinct_targets(monkeypatch): + _patch_newton_ik(monkeypatch) + solver = _pose_solver( + objectives=[ + NewtonIKPoseObjectiveCfg(body_name="ee"), + NewtonIKPoseObjectiveCfg(body_name="torso"), + NewtonIKJointLimitObjectiveCfg(weight=0.1), ] ) - env_origins = torch.tensor([[0.0, 0.0, 0.0], [1.0, 1.0, 1.0]]) + assert list(solver.objectives_by_name) == ["ee", "torso"] + assert solver.objectives_by_name["ee"].link_index == 0 + assert solver.objectives_by_name["torso"].link_index == 1 + + +def _build_pose_objective(cfg: NewtonIKPoseObjectiveCfg, num_envs: int = 2) -> NewtonIKPoseObjective: + ctx = NewtonIKBuildContext(model=_Model(), num_envs=num_envs, device="cpu", resolve_link=_resolver) + return NewtonIKPoseObjective(cfg, ctx) + + +def test_pose_objective_action_dim_and_coordinate_names(monkeypatch): + _patch_newton_ik(monkeypatch) + rel_pose = _build_pose_objective( + NewtonIKPoseObjectiveCfg(body_name="ee", command_type="pose", use_relative_mode=True) + ) + abs_pose = _build_pose_objective( + NewtonIKPoseObjectiveCfg(body_name="ee", command_type="pose", use_relative_mode=False) + ) + position = _build_pose_objective(NewtonIKPoseObjectiveCfg(body_name="ee", command_type="position")) + + assert rel_pose.action_dim == 6 + assert abs_pose.action_dim == 7 + assert position.action_dim == 3 + assert position.command_coordinate_names() == ["x", "y", "z"] + assert rel_pose.command_coordinate_names() == ["x", "y", "z", "roll", "pitch", "yaw"] + + +def test_relative_position_command_offsets_current_pose(monkeypatch): + _patch_newton_ik(monkeypatch) + obj = _build_pose_objective( + NewtonIKPoseObjectiveCfg(body_name="ee", command_type="position", use_relative_mode=True, scale=0.5) + ) + ee_pos_b = torch.tensor([[1.0, 1.0, 1.0], [2.0, 2.0, 2.0]]) + ee_quat_b = torch.tensor([[0.0, 0.0, 0.0, 1.0], [0.0, 0.0, 0.0, 1.0]]) + action = torch.tensor([[2.0, 0.0, -2.0], [0.0, 4.0, 0.0]]) + + target_pos_b, target_quat_b = obj.compute_target_b(action, ee_pos_b, ee_quat_b) + + # scale 0.5 applied internally, then added to the current position. + assert torch.allclose(target_pos_b, torch.tensor([[2.0, 1.0, 0.0], [2.0, 4.0, 2.0]])) + assert torch.allclose(target_quat_b, ee_quat_b) + + +class _CustomObjective(NewtonIKObjective): + SENTINEL = object() + + def __init__(self, cfg, ctx): + del cfg, ctx + self.name = "custom" + self.solver_objectives = [_CustomObjective.SENTINEL] + + +def test_custom_objective_cfg_is_built_and_wired(monkeypatch): + _patch_newton_ik(monkeypatch) + + @configclass + class _CustomObjectiveCfg(NewtonIKObjectiveCfg): + class_type: type | str = _CustomObjective - solver.set_pose_targets_from_body_q(body_q, names=["torso"], env_origins=env_origins) + solver = _pose_solver(objectives=[NewtonIKPoseObjectiveCfg(body_name="ee"), _CustomObjectiveCfg()]) - target_pos = wp.to_torch(solver.position_objectives["torso"].target) - target_quat = wp.to_torch(solver.rotation_objectives["torso"].target) - assert torch.allclose(target_pos, torch.tensor([[4.0, 5.0, 6.0], [3.0, 4.0, 5.0]])) - assert torch.allclose(target_quat, torch.tensor([[0.0, 0.0, 1.0, 0.0], [0.0, 0.0, 1.0, 0.0]])) + assert "custom" in solver.objectives_by_name + assert _CustomObjective.SENTINEL in solver.solver.kwargs["objectives"] diff --git a/source/isaaclab_newton/test/physics/test_newton_prototype_models.py b/source/isaaclab_newton/test/physics/test_newton_prototype_models.py deleted file mode 100644 index 1a62eda1176b..000000000000 --- a/source/isaaclab_newton/test/physics/test_newton_prototype_models.py +++ /dev/null @@ -1,73 +0,0 @@ -# Copyright (c) 2022-2026, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md). -# All rights reserved. -# -# SPDX-License-Identifier: BSD-3-Clause - -from types import SimpleNamespace - -import pytest -from isaaclab_newton.physics.newton_manager import NewtonManager - - -class _FakeBuilder: - def __init__(self): - self.finalize_calls = 0 - - def finalize(self, device=None): - self.finalize_calls += 1 - return SimpleNamespace(device=device) - - -@pytest.fixture(autouse=True) -def _restore_newton_manager_state(): - old_model = NewtonManager._model - old_prototypes = NewtonManager._prototype_models - try: - yield - finally: - NewtonManager._model = old_model - NewtonManager._prototype_models = old_prototypes - - -def test_get_prototype_model_matches_env_regex_path(): - builder = _FakeBuilder() - NewtonManager._model = SimpleNamespace(device="cuda:0") - NewtonManager.register_prototype_builders( - ("/World/envs/env_0",), ("/World/envs/env_{}",), {"/World/envs/env_0": builder} - ) - - info = NewtonManager.get_prototype_model("/World/envs/env_.*/Robot") - - assert info.source_path == "/World/envs/env_0" - assert info.model.device == "cuda:0" - assert builder.finalize_calls == 1 - - -def test_get_prototype_model_reuses_cached_model_on_same_device(): - builder = _FakeBuilder() - NewtonManager._model = SimpleNamespace(device="cpu") - NewtonManager.register_prototype_builders( - ("/World/envs/env_0/Robot",), ("/World/envs/env_{}/Robot",), {"/World/envs/env_0/Robot": builder} - ) - - info_0 = NewtonManager.get_prototype_model("/World/envs/env_.*/Robot") - info_1 = NewtonManager.get_prototype_model("/World/envs/env_.*/Robot") - - assert info_0.model is info_1.model - assert builder.finalize_calls == 1 - - -def test_register_prototype_builders_accumulates_across_calls(): - robot_builder = _FakeBuilder() - prop_builder = _FakeBuilder() - NewtonManager._model = SimpleNamespace(device="cpu") - - NewtonManager.register_prototype_builders( - ("/World/envs/env_0/Robot",), ("/World/envs/env_{}/Robot",), {"/World/envs/env_0/Robot": robot_builder} - ) - NewtonManager.register_prototype_builders( - ("/World/envs/env_0/Prop",), ("/World/envs/env_{}/Prop",), {"/World/envs/env_0/Prop": prop_builder} - ) - - assert NewtonManager.get_prototype_model("/World/envs/env_.*/Robot").builder is robot_builder - assert NewtonManager.get_prototype_model("/World/envs/env_.*/Prop").builder is prop_builder diff --git a/source/isaaclab_tasks/isaaclab_tasks/core/reach/config/franka/ik_newton_env_cfg.py b/source/isaaclab_tasks/isaaclab_tasks/core/reach/config/franka/ik_newton_env_cfg.py index fb465b30e970..018d0baa6ba0 100644 --- a/source/isaaclab_tasks/isaaclab_tasks/core/reach/config/franka/ik_newton_env_cfg.py +++ b/source/isaaclab_tasks/isaaclab_tasks/core/reach/config/franka/ik_newton_env_cfg.py @@ -4,6 +4,7 @@ # SPDX-License-Identifier: BSD-3-Clause from isaaclab_newton.envs.mdp.actions.newton_ik_actions_cfg import NewtonInverseKinematicsActionCfg +from isaaclab_newton.ik.newton_ik_objectives_cfg import NewtonIKJointLimitObjectiveCfg, NewtonIKPoseObjectiveCfg from isaaclab_newton.ik.newton_ik_solver_cfg import NewtonIKSolverCfg from isaaclab.utils.configclass import configclass @@ -29,10 +30,17 @@ def __post_init__(self): self.actions.arm_action = NewtonInverseKinematicsActionCfg( asset_name="robot", joint_names=["panda_joint.*"], - body_name="panda_hand", - controller=NewtonIKSolverCfg(command_type="pose", use_relative_mode=True), - scale=0.2, - body_offset=NewtonInverseKinematicsActionCfg.OffsetCfg(pos=(0.0, 0.0, 0.107)), + controller=NewtonIKSolverCfg(optimizer="lm", jacobian_mode="analytic", iterations=24), + objectives=[ + NewtonIKPoseObjectiveCfg( + body_name="panda_hand", + body_offset_pos=(0.0, 0.0, 0.107), + command_type="pose", + use_relative_mode=True, + scale=0.2, + ), + NewtonIKJointLimitObjectiveCfg(weight=0.1), + ], ) From a85a217eca16ab7d84f5a1fabab87cdb73e89327 Mon Sep 17 00:00:00 2001 From: Octi Zhang Date: Thu, 11 Jun 2026 22:55:29 -0700 Subject: [PATCH 4/4] Make Newton IK objective and action pure Warp Move the per-step IK compute off Torch: the pose objective exposes its command convention as Warp data (command code, relative flag, scale, offset), and the action term computes targets, assembles the seed, solves and gathers through Warp kernels. Torch now appears only as the policy action at the boundary, viewed zero-copy into Warp. --- .../isaaclab_newton/envs/__init__.pyi | 10 + .../isaaclab_newton/envs/mdp/__init__.pyi | 11 + .../envs/mdp/actions/newton_ik_actions.py | 245 ++++++++++++------ .../ik/newton_ik_objectives.py | 69 +++-- .../test/ik/test_newton_ik_solver.py | 34 +-- 5 files changed, 231 insertions(+), 138 deletions(-) create mode 100644 source/isaaclab_newton/isaaclab_newton/envs/__init__.pyi create mode 100644 source/isaaclab_newton/isaaclab_newton/envs/mdp/__init__.pyi diff --git a/source/isaaclab_newton/isaaclab_newton/envs/__init__.pyi b/source/isaaclab_newton/isaaclab_newton/envs/__init__.pyi new file mode 100644 index 000000000000..3955a1e9ba9f --- /dev/null +++ b/source/isaaclab_newton/isaaclab_newton/envs/__init__.pyi @@ -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 + +__all__ = [ + "mdp", +] + +from . import mdp diff --git a/source/isaaclab_newton/isaaclab_newton/envs/mdp/__init__.pyi b/source/isaaclab_newton/isaaclab_newton/envs/mdp/__init__.pyi new file mode 100644 index 000000000000..f3b63db0ae76 --- /dev/null +++ b/source/isaaclab_newton/isaaclab_newton/envs/mdp/__init__.pyi @@ -0,0 +1,11 @@ +# 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__ = [ + "NewtonInverseKinematicsAction", + "NewtonInverseKinematicsActionCfg", +] + +from .actions import NewtonInverseKinematicsAction, NewtonInverseKinematicsActionCfg diff --git a/source/isaaclab_newton/isaaclab_newton/envs/mdp/actions/newton_ik_actions.py b/source/isaaclab_newton/isaaclab_newton/envs/mdp/actions/newton_ik_actions.py index 8be7f052e061..5a91f113111a 100644 --- a/source/isaaclab_newton/isaaclab_newton/envs/mdp/actions/newton_ik_actions.py +++ b/source/isaaclab_newton/isaaclab_newton/envs/mdp/actions/newton_ik_actions.py @@ -7,7 +7,7 @@ import logging from collections.abc import Sequence -from dataclasses import dataclass, field +from dataclasses import dataclass from typing import TYPE_CHECKING import torch @@ -17,7 +17,6 @@ from newton.selection import ArticulationView import isaaclab.sim as sim_utils -import isaaclab.utils.math as math_utils import isaaclab.utils.string as string_utils from isaaclab.assets.articulation.base_articulation import BaseArticulation from isaaclab.cloner import resolve_clone_plan_source @@ -38,41 +37,125 @@ logger = logging.getLogger(__name__) +@wp.kernel(enable_backward=False) +def _ik_world_target_kernel( + body_pos_w: wp.array2d(dtype=wp.vec3f), + body_quat_w: wp.array2d(dtype=wp.quatf), + root_pos_w: wp.array(dtype=wp.vec3f), + root_quat_w: wp.array(dtype=wp.quatf), + body_idx: int, + offset: wp.transformf, + action: wp.array2d(dtype=wp.float32), + action_offset: int, + scale: wp.array(dtype=wp.float32), + command_code: int, + use_relative: int, + out_pos: wp.array(dtype=wp.vec3f), + out_rot: wp.array(dtype=wp.vec4f), +): + """Map one pose objective's action slice to a prototype-world target pose. + + Mirrors ``subtract_frame_transforms`` -> body offset -> command (position / + relative-pose / absolute-pose) -> ``combine_frame_transforms`` against the + env-0 root, writing the target straight into the objective's Warp arrays. + """ + i = wp.tid() + # End-effector (offset) pose in the env's root frame. + root_t = wp.transformf(root_pos_w[i], root_quat_w[i]) + ee_t = wp.transform_multiply( + wp.transform_inverse(root_t), wp.transformf(body_pos_w[i, body_idx], body_quat_w[i, body_idx]) + ) + ee_t = wp.transform_multiply(ee_t, offset) + ee_pos = wp.transform_get_translation(ee_t) + ee_rot = wp.transform_get_rotation(ee_t) + + target_pos = ee_pos + target_rot = ee_rot + if command_code == 0: # COMMAND_POSITION + disp = wp.vec3f( + action[i, action_offset + 0] * scale[0], + action[i, action_offset + 1] * scale[1], + action[i, action_offset + 2] * scale[2], + ) + target_pos = ee_pos + disp if use_relative == 1 else disp + else: + if use_relative == 1: + target_pos = ee_pos + wp.vec3f( + action[i, action_offset + 0] * scale[0], + action[i, action_offset + 1] * scale[1], + action[i, action_offset + 2] * scale[2], + ) + rot_vec = wp.vec3f( + action[i, action_offset + 3] * scale[3], + action[i, action_offset + 4] * scale[4], + action[i, action_offset + 5] * scale[5], + ) + angle = wp.length(rot_vec) + delta_rot = wp.quat_identity() + if angle > 1.0e-6: + delta_rot = wp.quat_from_axis_angle(rot_vec / angle, angle) + target_rot = delta_rot * ee_rot + else: + target_pos = wp.vec3f( + action[i, action_offset + 0] * scale[0], + action[i, action_offset + 1] * scale[1], + action[i, action_offset + 2] * scale[2], + ) + target_rot = wp.quatf( + action[i, action_offset + 3] * scale[3], + action[i, action_offset + 4] * scale[4], + action[i, action_offset + 5] * scale[5], + action[i, action_offset + 6] * scale[6], + ) + + # Broadcast against the env-0 prototype root (all roots identical, validated). + world_t = wp.transform_multiply(wp.transformf(root_pos_w[0], root_quat_w[0]), wp.transformf(target_pos, target_rot)) + out_pos[i] = wp.transform_get_translation(world_t) + q = wp.transform_get_rotation(world_t) + out_rot[i] = wp.vec4f(q[0], q[1], q[2], q[3]) + + +@wp.kernel(enable_backward=False) +def _ik_seed_scatter_kernel( + joint_pos: wp.array2d(dtype=wp.float32), + coord_ids: wp.array(dtype=wp.int32), + seed: wp.array2d(dtype=wp.float32), +): + """Overwrite the seed's actuated coordinates with the live joint positions.""" + i, j = wp.tid() + seed[i, coord_ids[j]] = joint_pos[i, j] + + +@wp.kernel(enable_backward=False) +def _ik_gather_kernel( + solved: wp.array2d(dtype=wp.float32), + coord_ids: wp.array(dtype=wp.int32), + out: wp.array2d(dtype=wp.float32), +): + """Gather the controlled coordinates from the full solved joint vector.""" + i, k = wp.tid() + out[i, k] = solved[i, coord_ids[k]] + + @dataclass class _PoseDriver: - """Per-pose-objective binding to the live articulation and the action vector. - - Holds the Isaac Lab body index used to read the current end-effector pose, the - target-frame offset (batched to ``num_envs``), this objective's slice of the - action vector, and the built solver objective. The per-step root-frame target - buffers are allocated in :meth:`__post_init__` from the offset's shape/device. - """ + """Per-pose-objective binding: the live body to read and its action slice offset.""" body_idx: int - offset_pos: torch.Tensor - offset_rot: torch.Tensor - slice: slice + action_offset: int objective: NewtonIKPoseObjective - target_pos_b: torch.Tensor = field(init=False) - target_quat_b: torch.Tensor = field(init=False) - - def __post_init__(self): - num_envs, device = self.offset_pos.shape[0], self.offset_pos.device - self.target_pos_b = torch.zeros(num_envs, 3, device=device) - self.target_quat_b = torch.zeros(num_envs, 4, device=device) - self.target_quat_b[:, 3] = 1.0 class NewtonInverseKinematicsAction(ActionTerm): """Newton inverse-kinematics action term. - The action solves IK as a single list of objectives on the single-env Newton - prototype model registered by the cloner, then maps the resulting actuated - joint coordinates back to the live batched Isaac Lab articulation. Each pose - objective contributes its slice of the action vector and drives one - end-effector body; a single pose objective is single-body IK, several are - multi-body IK. Constraint objectives (e.g. joint limits) add no action - dimensions. Fixed-base articulations only. + Solves IK as a single list of objectives on the cloner's single-env Newton + prototype model, then maps the actuated joint coordinates back to the live + batched articulation. Each pose objective drives one end-effector body (one is + single-body IK, several are multi-body); constraint objectives add no action + dimensions. The per-step target computation, seed assembly, solve and gather + run entirely in Warp -- Torch appears only as the policy action at the + boundary, viewed zero-copy into Warp. Fixed-base articulations only. """ cfg: NewtonInverseKinematicsActionCfg @@ -100,22 +183,14 @@ def __init__(self, cfg: NewtonInverseKinematicsActionCfg, env: ManagerBasedEnv): # lives at the asset suffix below it (e.g. ".../env_0" + "/Robot"). self._source_path = source_path + asset_suffix prototype_model = NewtonManager._cl_protos[source_path].finalize(device=NewtonManager.get_model().device) - # The prototype is the cloner's env_0 source builder; its single articulation - # is addressed by the resolved source path. prototype_view = ArticulationView( prototype_model, self._source_path, verbose=False, exclude_joint_types=[JointType.FREE, JointType.FIXED], ) - self._prototype_joint_coord_ids = self._resolve_prototype_joint_coord_ids( - prototype_view, self._asset.joint_names - ) - self._prototype_controlled_coord_ids = self._resolve_prototype_joint_coord_ids( - prototype_view, self._joint_names - ) - self._prototype_joint_seed = wp.to_torch(prototype_model.joint_q).to(device=self.device, dtype=torch.float32) - self._prototype_joint_seed = self._prototype_joint_seed.unsqueeze(0).repeat(self.num_envs, 1).contiguous() + coord_ids = self._resolve_prototype_joint_coord_ids(prototype_view, self._asset.joint_names) + controlled_ids = self._resolve_prototype_joint_coord_ids(prototype_view, self._joint_names) # The solver resolves each pose objective's body via the prototype view. self._ik_solver = self.cfg.controller.class_type( @@ -127,24 +202,29 @@ def __init__(self, cfg: NewtonInverseKinematicsActionCfg, env: ManagerBasedEnv): link_resolver=lambda body_name: self._resolve_prototype_link_index(prototype_view, body_name), ) - # Bind each pose objective to the live articulation and the action vector. + # Bind each pose objective to the live body it reads and its action slice. self._drivers: list[_PoseDriver] = [] offset = 0 for pose_cfg in pose_cfgs: name = pose_cfg.name if pose_cfg.name is not None else pose_cfg.body_name objective = self._ik_solver.objectives_by_name[name] body_idx = self._resolve_isaac_body_index(pose_cfg.body_name) - offset_pos = torch.tensor(pose_cfg.body_offset_pos, device=self.device).repeat(self.num_envs, 1) - offset_rot = torch.tensor(pose_cfg.body_offset_rot, device=self.device).repeat(self.num_envs, 1) - self._drivers.append( - _PoseDriver(body_idx, offset_pos, offset_rot, slice(offset, offset + objective.action_dim), objective) - ) + self._drivers.append(_PoseDriver(body_idx, offset, objective)) offset += objective.action_dim self._action_dim = offset self._raw_actions = torch.zeros(self.num_envs, self._action_dim, device=self.device) self._processed_actions = torch.zeros_like(self._raw_actions) + # Warp scratch for the seed -> solve -> gather pipeline. + num_coords = prototype_model.joint_coord_count + default_seed = wp.to_torch(prototype_model.joint_q).to(device=self.device, dtype=torch.float32) + self._default_seed = wp.from_torch(default_seed.unsqueeze(0).repeat(self.num_envs, 1).contiguous()) + self._seed = wp.zeros((self.num_envs, num_coords), dtype=wp.float32, device=self.device) + self._joint_pos_des = wp.zeros((self.num_envs, len(self._joint_ids)), dtype=wp.float32, device=self.device) + self._coord_ids = wp.from_torch(coord_ids.to(torch.int32).contiguous()) + self._controlled_ids = wp.from_torch(controlled_ids.to(torch.int32).contiguous()) + self._clip = None if self.cfg.clip is not None: self._clip = torch.tensor([[-float("inf"), float("inf")]], device=self.device).repeat( @@ -193,49 +273,60 @@ def process_actions(self, actions: torch.Tensor) -> None: self._processed_actions = torch.clamp( self._processed_actions, min=self._clip[:, :, 0], max=self._clip[:, :, 1] ) - # Each pose objective maps its own action slice (scaled internally) onto a - # root-frame target from the body's current pose. - for driver in self._drivers: - ee_pos_b, ee_quat_b = self._compute_frame_pose(driver) - target_pos_b, target_quat_b = driver.objective.compute_target_b( - self._processed_actions[:, driver.slice], ee_pos_b, ee_quat_b - ) - driver.target_pos_b[:] = target_pos_b - driver.target_quat_b[:] = target_quat_b - - def apply_actions(self) -> None: - # The IK solve runs on the single-env prototype model, so all batched - # root-frame targets are expressed in the prototype (env 0) world frame. + # Each pose objective maps its action slice to a prototype-world target, + # written straight into its Warp target arrays. self._validate_matching_root_orientations() - root_pos_proto = self._asset.data.root_pos_w.torch[0:1].repeat(self.num_envs, 1) - root_quat_proto = self._asset.data.root_quat_w.torch[0:1].repeat(self.num_envs, 1) + action_wp = wp.from_torch(self._processed_actions.contiguous(), dtype=wp.float32) + body_pos_w = self._asset.data.body_pos_w.warp + body_quat_w = self._asset.data.body_quat_w.warp + root_pos_w = self._asset.data.root_pos_w.warp + root_quat_w = self._asset.data.root_quat_w.warp for driver in self._drivers: - target_pos_w, target_quat_w = math_utils.combine_frame_transforms( - root_pos_proto, root_quat_proto, driver.target_pos_b, driver.target_quat_b + obj = driver.objective + wp.launch( + _ik_world_target_kernel, + dim=self.num_envs, + inputs=[ + body_pos_w, + body_quat_w, + root_pos_w, + root_quat_w, + driver.body_idx, + obj.offset, + action_wp, + driver.action_offset, + obj.scale, + obj.command_code, + obj.use_relative, + obj.position_objective.target_positions, + obj.rotation_objective.target_rotations, + ], + device=self.device, ) - driver.objective.set_target_pose(target_pos_w, target_quat_w) - joint_seed = self._prototype_joint_seed.clone() - joint_seed[:, self._prototype_joint_coord_ids] = self._asset.data.joint_pos.torch - joint_pos_des_all = wp.to_torch(self._ik_solver.solve(wp.from_torch(joint_seed.contiguous(), dtype=wp.float32))) - joint_pos_des = joint_pos_des_all[:, self._prototype_controlled_coord_ids].contiguous() - self._asset.set_joint_position_target_index(target=joint_pos_des, joint_ids=self._joint_ids_warp) + def apply_actions(self) -> None: + # Seed the solver from the live joint positions on top of the prototype + # default, solve, and write the controlled coordinates back -- all in Warp. + wp.copy(self._seed, self._default_seed) + wp.launch( + _ik_seed_scatter_kernel, + dim=(self.num_envs, len(self._asset.joint_names)), + inputs=[self._asset.data.joint_pos.warp, self._coord_ids, self._seed], + device=self.device, + ) + solved = self._ik_solver.solve(self._seed) + wp.launch( + _ik_gather_kernel, + dim=(self.num_envs, len(self._joint_ids)), + inputs=[solved, self._controlled_ids, self._joint_pos_des], + device=self.device, + ) + self._asset.set_joint_position_target_index(target=self._joint_pos_des, joint_ids=self._joint_ids_warp) def reset(self, env_ids: Sequence[int] | None = None) -> None: env_ids = slice(None) if env_ids is None else env_ids self._raw_actions[env_ids] = 0.0 - def _compute_frame_pose(self, driver: _PoseDriver) -> tuple[torch.Tensor, torch.Tensor]: - ee_pos_w = self._asset.data.body_pos_w.torch[:, driver.body_idx] - ee_quat_w = self._asset.data.body_quat_w.torch[:, driver.body_idx] - root_pos_w = self._asset.data.root_pos_w.torch - root_quat_w = self._asset.data.root_quat_w.torch - ee_pos_b, ee_quat_b = math_utils.subtract_frame_transforms(root_pos_w, root_quat_w, ee_pos_w, ee_quat_w) - ee_pos_b, ee_quat_b = math_utils.combine_frame_transforms( - ee_pos_b, ee_quat_b, driver.offset_pos, driver.offset_rot - ) - return ee_pos_b, ee_quat_b - def _validate_matching_root_orientations(self) -> None: """Guard the prototype-frame IK assumption for replicated fixed-base roots.""" root_quat_w = self._asset.data.root_quat_w.torch diff --git a/source/isaaclab_newton/isaaclab_newton/ik/newton_ik_objectives.py b/source/isaaclab_newton/isaaclab_newton/ik/newton_ik_objectives.py index 11262ddb3172..e18227d67327 100644 --- a/source/isaaclab_newton/isaaclab_newton/ik/newton_ik_objectives.py +++ b/source/isaaclab_newton/isaaclab_newton/ik/newton_ik_objectives.py @@ -8,17 +8,17 @@ Each class is built by :class:`~isaaclab_newton.ik.NewtonIKSolver` from the matching :class:`~isaaclab_newton.ik.newton_ik_objectives_cfg.NewtonIKObjectiveCfg` and owns the concrete :class:`newton.ik.IKObjective` instances appended to the -solver. Command-driven objectives (currently pose) additionally own their -action contribution: an :attr:`~NewtonIKObjective.action_dim`, the coordinate -names for that slice, and the mapping from a raw action slice to a target pose. -The action term is therefore generic -- it sums :attr:`action_dim` across the -objective list and dispatches each slice without branching on command type. +solver. Pose objectives also describe their action contribution as Warp data: +an :attr:`~NewtonIKObjective.action_dim`, the coordinate names for that slice, +a numeric :attr:`~NewtonIKPoseObjective.command_code` / relative flag, a Warp +``scale`` array and a target-frame ``offset`` transform. The action term reads +these directly into a Warp kernel; nothing here touches Torch. Importing this module pulls ``newton`` (and ``pxr``), so it is loaded lazily via the package ``lazy_export`` only after Kit has launched. Custom objectives -integrate by subclassing :class:`NewtonIKObjective`, taking -``(cfg, ctx)`` in ``__init__`` -- pulling only the :class:`NewtonIKBuildContext` -fields they need -- and populating :attr:`NewtonIKObjective.solver_objectives`. +integrate by subclassing :class:`NewtonIKObjective`, taking ``(cfg, ctx)`` in +``__init__`` -- pulling only the :class:`NewtonIKBuildContext` fields they need -- +and populating :attr:`NewtonIKObjective.solver_objectives`. """ from __future__ import annotations @@ -27,13 +27,14 @@ from dataclasses import dataclass import newton.ik as ik -import torch import warp as wp -import isaaclab.utils.math as math_utils - from .newton_ik_objectives_cfg import NewtonIKJointLimitObjectiveCfg, NewtonIKPoseObjectiveCfg +# Numeric command codes consumed by the action's Warp kernel. +COMMAND_POSITION = 0 +COMMAND_POSE = 1 + @dataclass(frozen=True) class NewtonIKBuildContext: @@ -56,9 +57,8 @@ class NewtonIKObjective: """Base built IK objective. Owns the concrete :class:`newton.ik.IKObjective` instances in - :attr:`solver_objectives`. Command-driven objectives (pose) also set a - :attr:`name` and a non-zero :attr:`action_dim` and add ``compute_target_b`` - / ``command_coordinate_names``; constraint objectives leave the defaults. + :attr:`solver_objectives`. Pose objectives also set a :attr:`name` and a + non-zero :attr:`action_dim`; constraint objectives leave the defaults. """ name: str | None = None @@ -72,7 +72,13 @@ class NewtonIKObjective: class NewtonIKPoseObjective(NewtonIKObjective): - """Command-driven position + rotation objective tracking one end-effector body.""" + """Command-driven position + rotation objective tracking one end-effector body. + + Exposes its command convention to the action's Warp kernel as data: + :attr:`command_code` / :attr:`use_relative`, the per-coordinate :attr:`scale` + (``wp.float32``), the target-frame :attr:`offset` (``wp.transformf``), and the + position/rotation target arrays the kernel writes into. + """ def __init__(self, cfg: NewtonIKPoseObjectiveCfg, ctx: NewtonIKBuildContext): self.name = cfg.name if cfg.name is not None else cfg.body_name @@ -81,15 +87,16 @@ def __init__(self, cfg: NewtonIKPoseObjectiveCfg, ctx: NewtonIKBuildContext): self.link_index = ctx.resolve_link(cfg.body_name) self.action_dim = len(self.command_coordinate_names()) - scale = torch.as_tensor(cfg.scale, dtype=torch.float32, device=ctx.device) - if scale.ndim == 0: - scale = scale.repeat(self.action_dim) - if scale.shape != (self.action_dim,): + self.command_code = COMMAND_POSITION if cfg.command_type == "position" else COMMAND_POSE + self.use_relative = int(cfg.use_relative_mode) + scale_values = [float(cfg.scale)] * self.action_dim if _is_scalar(cfg.scale) else [float(s) for s in cfg.scale] + if len(scale_values) != self.action_dim: raise ValueError( f"Newton IK pose objective '{self.name}' scale must be a float or length-{self.action_dim} " - f"sequence, got shape {tuple(scale.shape)}." + f"sequence, got {len(scale_values)} values." ) - self._scale = scale + self.scale = wp.array(scale_values, dtype=wp.float32, device=ctx.device) + self.offset = wp.transformf(wp.vec3f(*cfg.body_offset_pos), wp.quatf(*cfg.body_offset_rot)) target_positions = wp.zeros((ctx.num_envs,), dtype=wp.vec3, device=ctx.device) target_rotations = wp.array([(0.0, 0.0, 0.0, 1.0)] * ctx.num_envs, dtype=wp.vec4, device=ctx.device) @@ -116,22 +123,6 @@ def command_coordinate_names(self) -> list[str]: return ["x", "y", "z", "qx", "qy", "qz", "qw"] raise ValueError(f"Unsupported Newton IK command type: {self.command_type}") - def compute_target_b( - self, action: torch.Tensor, ee_pos_b: torch.Tensor, ee_quat_b: torch.Tensor - ) -> tuple[torch.Tensor, torch.Tensor]: - processed = action * self._scale - if self.command_type == "position": - target_pos_b = ee_pos_b + processed if self.use_relative_mode else processed - return target_pos_b, ee_quat_b - if self.use_relative_mode: - return math_utils.apply_delta_pose(ee_pos_b, ee_quat_b, processed) - return processed[:, 0:3], processed[:, 3:7] - - def set_target_pose(self, target_pos_w: torch.Tensor, target_quat_w: torch.Tensor) -> None: - """Update the batched world-frame target pose, shape ``[num_envs, 3]`` and ``[num_envs, 4]``.""" - self.position_objective.set_target_positions(wp.from_torch(target_pos_w.contiguous(), dtype=wp.vec3)) - self.rotation_objective.set_target_rotations(wp.from_torch(target_quat_w.contiguous(), dtype=wp.vec4)) - class NewtonIKJointLimitObjective(NewtonIKObjective): """Soft joint-limit constraint reading the model's coordinate limits.""" @@ -143,3 +134,7 @@ def __init__(self, cfg: NewtonIKJointLimitObjectiveCfg, ctx: NewtonIKBuildContex weight=cfg.weight, ) self.solver_objectives = [self.objective] + + +def _is_scalar(value) -> bool: + return isinstance(value, (int, float)) diff --git a/source/isaaclab_newton/test/ik/test_newton_ik_solver.py b/source/isaaclab_newton/test/ik/test_newton_ik_solver.py index 041cc8b46ee1..a155c881c62c 100644 --- a/source/isaaclab_newton/test/ik/test_newton_ik_solver.py +++ b/source/isaaclab_newton/test/ik/test_newton_ik_solver.py @@ -100,19 +100,6 @@ def test_solve_writes_output_buffer(monkeypatch): assert torch.allclose(wp.to_torch(result), torch.tensor([[2.0, 3.0], [4.0, 5.0]])) -def test_set_target_pose_updates_named_objective(monkeypatch): - _patch_newton_ik(monkeypatch) - solver = _pose_solver() - target_pos = torch.tensor([[0.1, 0.2, 0.3], [1.0, 1.1, 1.2]]) - target_quat = torch.tensor([[0.0, 0.0, 0.0, 1.0], [0.0, 0.0, 1.0, 0.0]]) - - pose = solver.objectives_by_name["ee"] - pose.set_target_pose(target_pos, target_quat) - - assert torch.allclose(wp.to_torch(pose.position_objective.target), target_pos) - assert torch.allclose(wp.to_torch(pose.rotation_objective.target), target_quat) - - def test_constraint_objectives_carry_no_target_or_action(monkeypatch): _patch_newton_ik(monkeypatch) solver = _pose_solver( @@ -163,20 +150,19 @@ def test_pose_objective_action_dim_and_coordinate_names(monkeypatch): assert rel_pose.command_coordinate_names() == ["x", "y", "z", "roll", "pitch", "yaw"] -def test_relative_position_command_offsets_current_pose(monkeypatch): +def test_pose_objective_exposes_warp_command_data(monkeypatch): _patch_newton_ik(monkeypatch) - obj = _build_pose_objective( - NewtonIKPoseObjectiveCfg(body_name="ee", command_type="position", use_relative_mode=True, scale=0.5) + rel_pose = _build_pose_objective( + NewtonIKPoseObjectiveCfg(body_name="ee", command_type="pose", use_relative_mode=True, scale=0.5) + ) + position = _build_pose_objective( + NewtonIKPoseObjectiveCfg(body_name="ee", command_type="position", use_relative_mode=False) ) - ee_pos_b = torch.tensor([[1.0, 1.0, 1.0], [2.0, 2.0, 2.0]]) - ee_quat_b = torch.tensor([[0.0, 0.0, 0.0, 1.0], [0.0, 0.0, 0.0, 1.0]]) - action = torch.tensor([[2.0, 0.0, -2.0], [0.0, 4.0, 0.0]]) - - target_pos_b, target_quat_b = obj.compute_target_b(action, ee_pos_b, ee_quat_b) - # scale 0.5 applied internally, then added to the current position. - assert torch.allclose(target_pos_b, torch.tensor([[2.0, 1.0, 0.0], [2.0, 4.0, 2.0]])) - assert torch.allclose(target_quat_b, ee_quat_b) + # Command convention is exposed to the action's Warp kernel as plain data. + assert (rel_pose.command_code, rel_pose.use_relative) == (1, 1) + assert (position.command_code, position.use_relative) == (0, 0) + assert torch.allclose(wp.to_torch(rel_pose.scale), torch.full((6,), 0.5)) class _CustomObjective(NewtonIKObjective):