diff --git a/source/isaaclab_newton/changelog.d/max-newton-ik-solver.rst b/source/isaaclab_newton/changelog.d/max-newton-ik-solver.rst new file mode 100644 index 000000000000..06fc1d6b38bb --- /dev/null +++ b/source/isaaclab_newton/changelog.d/max-newton-ik-solver.rst @@ -0,0 +1,9 @@ +Added +^^^^^ + +* 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. +* 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..04bde19f7a8c 100644 --- a/source/isaaclab_newton/isaaclab_newton/cloner/replicate.py +++ b/source/isaaclab_newton/isaaclab_newton/cloner/replicate.py @@ -43,8 +43,13 @@ 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. + + 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) if quaternions is None: @@ -98,7 +103,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 +213,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 +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._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/__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/__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__.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/__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/__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..5a91f113111a --- /dev/null +++ b/source/isaaclab_newton/isaaclab_newton/envs/mdp/actions/newton_ik_actions.py @@ -0,0 +1,380 @@ +# 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 dataclasses import dataclass +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.sim as sim_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_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__) + + +@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: the live body to read and its action slice offset.""" + + body_idx: int + action_offset: int + objective: NewtonIKPoseObjective + + +class NewtonInverseKinematicsAction(ActionTerm): + """Newton inverse-kinematics action term. + + 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 + _asset: BaseArticulation + + def __init__(self, cfg: NewtonInverseKinematicsActionCfg, env: ManagerBasedEnv): + super().__init__(cfg, env) + + 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) + self._joint_ids_warp = wp.array(self._joint_ids, dtype=wp.int32, device=self.device) + + 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) + prototype_view = ArticulationView( + prototype_model, + self._source_path, + verbose=False, + exclude_joint_types=[JointType.FREE, JointType.FIXED], + ) + 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( + self.cfg.controller, + model=prototype_model, + num_envs=self.num_envs, + device=self.device, + 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 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) + 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( + 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 bodies %s.", + self._joint_names, + self._joint_ids, + [(d.objective.name, d.body_idx) for d in self._drivers], + ) + + @property + def action_dim(self) -> int: + return self._action_dim + + @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.joint_names = self._joint_names + self._IO_descriptor.clip = self.cfg.clip + self._IO_descriptor.extras["controller_cfg"] = self.cfg.controller.__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 + if self._clip is not None: + self._processed_actions = torch.clamp( + self._processed_actions, min=self._clip[:, :, 0], max=self._clip[:, :, 1] + ) + # Each pose objective maps its action slice to a prototype-world target, + # written straight into its Warp target arrays. + self._validate_matching_root_orientations() + 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: + 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, + ) + + 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 _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_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] + ) -> 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) + } + 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, body_name: str) -> int: + layout = prototype_view.frequency_layouts[NewtonModel.AttributeFrequency.BODY] + selected_indices = self._layout_indices(layout) + local_link_index = prototype_view.link_names.index(body_name) + 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]: + 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 new file mode 100644 index 000000000000..e8330d62699e --- /dev/null +++ b/source/isaaclab_newton/isaaclab_newton/envs/mdp/actions/newton_ik_actions_cfg.py @@ -0,0 +1,54 @@ +# 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, 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: + from .newton_ik_actions import NewtonInverseKinematicsAction + + +@configclass +class NewtonInverseKinematicsActionCfg(ActionTermCfg): + """Configuration for a Newton inverse-kinematics action term. + + 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. + """ + + class_type: type[NewtonInverseKinematicsAction] | str = ( + "isaaclab_newton.envs.mdp.actions.newton_ik_actions:NewtonInverseKinematicsAction" + ) + + joint_names: list[str] = MISSING + """Joints actuated by the action. + + 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. + """ + + 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__.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..364300e05fcb --- /dev/null +++ b/source/isaaclab_newton/isaaclab_newton/ik/__init__.pyi @@ -0,0 +1,24 @@ +# 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__ = [ + "NewtonIKSolver", + "NewtonIKSolverCfg", + "NewtonIKObjective", + "NewtonIKPoseObjective", + "NewtonIKJointLimitObjective", + "NewtonIKObjectiveCfg", + "NewtonIKPoseObjectiveCfg", + "NewtonIKJointLimitObjectiveCfg", +] + +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..e18227d67327 --- /dev/null +++ b/source/isaaclab_newton/isaaclab_newton/ik/newton_ik_objectives.py @@ -0,0 +1,140 @@ +# 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. 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`. +""" + +from __future__ import annotations + +from collections.abc import Callable +from dataclasses import dataclass + +import newton.ik as ik +import warp as wp + +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: + """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`. Pose objectives also set a :attr:`name` and a + non-zero :attr:`action_dim`; 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. + + 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 + 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()) + + 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 {len(scale_values)} values." + ) + 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) + 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}") + + +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] + + +def _is_scalar(value) -> bool: + return isinstance(value, (int, float)) 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 new file mode 100644 index 000000000000..1dc35a4dbb01 --- /dev/null +++ b/source/isaaclab_newton/isaaclab_newton/ik/newton_ik_solver.py @@ -0,0 +1,104 @@ +# 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 Callable, Sequence + +import newton.ik as ik +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 + + +class NewtonIKSolver: + """Batched wrapper around Newton's inverse-kinematics solver. + + 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 + + def __init__( + self, + cfg: NewtonIKSolverCfg, + *, + model, + num_envs: int, + device: str, + objectives: Sequence[NewtonIKObjectiveCfg], + link_resolver: Callable[[str], int], + ): + if not objectives: + raise ValueError("NewtonIKSolver requires at least one objective cfg.") + + self.cfg = cfg + ctx = NewtonIKBuildContext(model=model, num_envs=num_envs, device=device, resolve_link=link_resolver) + + self.objectives: list[NewtonIKObjective] = [] + self.objectives_by_name: dict[str, NewtonIKObjective] = {} + solver_objectives: list[ik.IKObjective] = [] + 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, + 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 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: 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_pos, + self.joint_q_out, + iterations=self.cfg.iterations, + step_size=self.cfg.step_size, + ) + 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 new file mode 100644 index 000000000000..520433b856c3 --- /dev/null +++ b/source/isaaclab_newton/isaaclab_newton/ik/newton_ik_solver_cfg.py @@ -0,0 +1,66 @@ +# 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 NewtonIKSolverCfg: + """Configuration for the Newton inverse-kinematics solver. + + 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. + + :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"``.""" + + 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. + """ diff --git a/source/isaaclab_newton/isaaclab_newton/physics/newton_manager.py b/source/isaaclab_newton/isaaclab_newton/physics/newton_manager.py index bb5279e7d8be..67ed27f21d89 100644 --- a/source/isaaclab_newton/isaaclab_newton/physics/newton_manager.py +++ b/source/isaaclab_newton/isaaclab_newton/physics/newton_manager.py @@ -311,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]] = [] @@ -788,6 +792,7 @@ def clear(cls): 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 = [] @@ -1242,6 +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._cl_protos = source_builders global_site_indices, source_site_indices, env_root_sites = cls._cl_inject_sites(builder, source_builders) xform_cache = UsdGeom.XformCache() diff --git a/source/isaaclab_newton/test/ik/test_newton_ik_solver.py b/source/isaaclab_newton/test/ik/test_newton_ik_solver.py new file mode 100644 index 000000000000..a155c881c62c --- /dev/null +++ b/source/isaaclab_newton/test/ik/test_newton_ik_solver.py @@ -0,0 +1,187 @@ +# 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_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_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: + 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): + # ``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 _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=num_envs, + device="cpu", + objectives=[NewtonIKPoseObjectiveCfg(body_name="ee")] if objectives is None else objectives, + link_resolver=_resolver, + ) + + +def test_solve_writes_output_buffer(monkeypatch): + _patch_newton_ik(monkeypatch) + solver = _pose_solver() + seed = wp.from_torch(torch.tensor([[1.0, 2.0], [3.0, 4.0]]), dtype=wp.float32) + + result = solver.solve(seed) + + assert torch.allclose(wp.to_torch(result), torch.tensor([[2.0, 3.0], [4.0, 5.0]])) + + +def test_constraint_objectives_carry_no_target_or_action(monkeypatch): + _patch_newton_ik(monkeypatch) + solver = _pose_solver( + objectives=[ + NewtonIKPoseObjectiveCfg(body_name="ee"), + NewtonIKJointLimitObjectiveCfg(weight=0.1), + ] + ) + # 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), + ] + ) + 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_pose_objective_exposes_warp_command_data(monkeypatch): + _patch_newton_ik(monkeypatch) + 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) + ) + + # 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): + 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 = _pose_solver(objectives=[NewtonIKPoseObjectiveCfg(body_name="ee"), _CustomObjectiveCfg()]) + + assert "custom" in solver.objectives_by_name + assert _CustomObjective.SENTINEL in solver.solver.kwargs["objectives"] diff --git a/source/isaaclab_tasks/changelog.d/max-newton-ik-solver.rst b/source/isaaclab_tasks/changelog.d/max-newton-ik-solver.rst new file mode 100644 index 000000000000..a16171ca8db8 --- /dev/null +++ b/source/isaaclab_tasks/changelog.d/max-newton-ik-solver.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..018d0baa6ba0 --- /dev/null +++ b/source/isaaclab_tasks/isaaclab_tasks/core/reach/config/franka/ik_newton_env_cfg.py @@ -0,0 +1,53 @@ +# 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_objectives_cfg import NewtonIKJointLimitObjectiveCfg, NewtonIKPoseObjectiveCfg +from isaaclab_newton.ik.newton_ik_solver_cfg import NewtonIKSolverCfg + +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.*"], + 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), + ], + ) + + +@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