diff --git a/docs/source/api/index.rst b/docs/source/api/index.rst index 46343eb6419a..15c8c1974345 100644 --- a/docs/source/api/index.rst +++ b/docs/source/api/index.rst @@ -152,6 +152,7 @@ The following modules are available in the ``isaaclab_newton`` extension: assets cloner + controllers physics renderers scene_data_providers @@ -162,6 +163,7 @@ The following modules are available in the ``isaaclab_newton`` extension: lab_newton/isaaclab_newton.assets lab_newton/isaaclab_newton.cloner + lab_newton/isaaclab_newton.controllers lab_newton/isaaclab_newton.physics lab_newton/isaaclab_newton.renderers lab_newton/isaaclab_newton.scene_data_providers diff --git a/docs/source/api/lab_newton/isaaclab_newton.controllers.rst b/docs/source/api/lab_newton/isaaclab_newton.controllers.rst new file mode 100644 index 000000000000..e502f5a03bc5 --- /dev/null +++ b/docs/source/api/lab_newton/isaaclab_newton.controllers.rst @@ -0,0 +1,4 @@ +isaaclab\_newton.controllers +============================ + +.. automodule:: isaaclab_newton.controllers diff --git a/source/isaaclab/isaaclab/envs/mdp/__init__.pyi b/source/isaaclab/isaaclab/envs/mdp/__init__.pyi index 399035f2c09c..adf609106188 100644 --- a/source/isaaclab/isaaclab/envs/mdp/__init__.pyi +++ b/source/isaaclab/isaaclab/envs/mdp/__init__.pyi @@ -55,6 +55,7 @@ __all__ = [ "randomize_physics_scene_gravity", "randomize_rigid_body_collider_offsets", "randomize_rigid_body_com", + "randomize_rigid_body_inertia", "randomize_rigid_body_mass", "randomize_rigid_body_material", "randomize_rigid_body_scale", @@ -195,6 +196,7 @@ from .events import ( randomize_physics_scene_gravity, randomize_rigid_body_collider_offsets, randomize_rigid_body_com, + randomize_rigid_body_inertia, randomize_rigid_body_mass, randomize_rigid_body_material, randomize_rigid_body_scale, diff --git a/source/isaaclab/isaaclab/envs/mdp/events.py b/source/isaaclab/isaaclab/envs/mdp/events.py index 18c9b5eff63a..3781df73cfdc 100644 --- a/source/isaaclab/isaaclab/envs/mdp/events.py +++ b/source/isaaclab/isaaclab/envs/mdp/events.py @@ -26,7 +26,7 @@ import isaaclab.utils.math as math_utils from isaaclab.actuators import ImplicitActuator from isaaclab.managers import EventTermCfg, ManagerTermBase, SceneEntityCfg -from isaaclab.utils.version import compare_versions, get_isaac_sim_version +from isaaclab.utils.version import compare_versions if TYPE_CHECKING: from isaaclab_physx.assets import DeformableObject @@ -151,69 +151,50 @@ def randomize_rigid_body_scale( op_order_spec.default = Vt.TokenArray(["xformOp:translate", "xformOp:orient", "xformOp:scale"]) -class randomize_rigid_body_material(ManagerTermBase): - """Randomize the physics materials on all geometries of the asset. - - This function creates a set of physics materials with random static friction, dynamic friction, and restitution - values. The number of materials is specified by ``num_buckets``. The materials are generated by sampling - uniform random values from the given ranges. - - The material properties are then assigned to the geometries of the asset. The assignment is done by - creating a random integer tensor of shape (num_instances, max_num_shapes) where ``num_instances`` - is the number of assets spawned and ``max_num_shapes`` is the maximum number of shapes in the asset (over - all bodies). The integer values are used as indices to select the material properties from the - material buckets. - - If the flag ``make_consistent`` is set to ``True``, the dynamic friction is set to be less than or equal to - the static friction. This obeys the physics constraint on friction values. However, it may not always be - essential for the application. Thus, the flag is set to ``False`` by default. +class _RandomizeRigidBodyMaterialPhysx: + """PhysX backend implementation for material randomization. - .. attention:: - This function uses CPU tensors to assign the material properties. It is recommended to use this function - only during the initialization of the environment. Otherwise, it may lead to a significant performance - overhead. - - .. note:: - PhysX only allows 64000 unique physics materials in the scene. If the number of materials exceeds this - limit, the simulation will crash. Due to this reason, we sample the materials only once during initialization. - Afterwards, these materials are randomly assigned to the geometries of the asset. + Uses the bucket-based approach required by PhysX's 64000 unique material limit. + Materials are pre-sampled into buckets and randomly assigned to shapes. """ - def __init__(self, cfg: EventTermCfg, env: ManagerBasedEnv): - """Initialize the term. - - Args: - cfg: The configuration of the event term. - env: The environment instance. + def __init__( + self, cfg: EventTermCfg, env: ManagerBasedEnv, asset: RigidObject | Articulation, asset_cfg: SceneEntityCfg + ): + from isaaclab.assets import BaseArticulation - Raises: - ValueError: If the asset is not a RigidObject or an Articulation. - """ - from isaaclab.assets import BaseArticulation, BaseRigidObject + # obtain parameters for sampling friction and restitution values + static_friction_range = cfg.params.get("static_friction_range", (1.0, 1.0)) + dynamic_friction_range = cfg.params.get("dynamic_friction_range", (1.0, 1.0)) + restitution_range = cfg.params.get("restitution_range", (0.0, 0.0)) + num_buckets = int(cfg.params.get("num_buckets", 1)) - super().__init__(cfg, env) + # sample material properties from the given ranges + # note: we only sample the materials once during initialization + # afterwards these are randomly assigned to the geometries of the asset + range_list = [static_friction_range, dynamic_friction_range, restitution_range] + ranges = torch.tensor(range_list, device="cpu") + self.material_buckets = math_utils.sample_uniform(ranges[:, 0], ranges[:, 1], (num_buckets, 3), device="cpu") - # extract the used quantities (to enable type-hinting) - self.asset_cfg: SceneEntityCfg = cfg.params["asset_cfg"] - self.asset: RigidObject | Articulation = env.scene[self.asset_cfg.name] + # ensure dynamic friction is always less than static friction + make_consistent = cfg.params.get("make_consistent", False) + if make_consistent: + self.material_buckets[:, 1] = torch.min(self.material_buckets[:, 0], self.material_buckets[:, 1]) - if not isinstance(self.asset, (BaseRigidObject, BaseArticulation)): - raise ValueError( - f"Randomization term 'randomize_rigid_body_material' not supported for asset: '{self.asset_cfg.name}'" - f" with type: '{type(self.asset)}'." - ) + self.asset = asset + self.asset_cfg = asset_cfg # obtain number of shapes per body (needed for indexing the material properties correctly) # note: this is a workaround since the Articulation does not provide a direct way to obtain the number of shapes # per body. We use the physics simulation view to obtain the number of shapes per body. - if isinstance(self.asset, BaseArticulation) and self.asset_cfg.body_ids != slice(None): + if isinstance(asset, BaseArticulation) and asset_cfg.body_ids != slice(None): self.num_shapes_per_body = [] - for link_path in self.asset.root_view.link_paths[0]: - link_physx_view = self.asset._physics_sim_view.create_rigid_body_view(link_path) # type: ignore + for link_path in asset.root_view.link_paths[0]: + link_physx_view = asset._physics_sim_view.create_rigid_body_view(link_path) # type: ignore self.num_shapes_per_body.append(link_physx_view.max_shapes) # ensure the parsing is correct num_shapes = sum(self.num_shapes_per_body) - expected_shapes = self.asset.root_view.max_shapes + expected_shapes = asset.root_view.max_shapes if num_shapes != expected_shapes: raise ValueError( "Randomization term 'randomize_rigid_body_material' failed to parse the number of shapes per body." @@ -223,24 +204,6 @@ def __init__(self, cfg: EventTermCfg, env: ManagerBasedEnv): # in this case, we don't need to do special indexing self.num_shapes_per_body = None - # obtain parameters for sampling friction and restitution values - static_friction_range = cfg.params.get("static_friction_range", (1.0, 1.0)) - dynamic_friction_range = cfg.params.get("dynamic_friction_range", (1.0, 1.0)) - restitution_range = cfg.params.get("restitution_range", (0.0, 0.0)) - num_buckets = int(cfg.params.get("num_buckets", 1)) - - # sample material properties from the given ranges - # note: we only sample the materials once during initialization - # afterwards these are randomly assigned to the geometries of the asset - range_list = [static_friction_range, dynamic_friction_range, restitution_range] - ranges = torch.tensor(range_list, device="cpu") - self.material_buckets = math_utils.sample_uniform(ranges[:, 0], ranges[:, 1], (num_buckets, 3), device="cpu") - - # ensure dynamic friction is always less than static friction - make_consistent = cfg.params.get("make_consistent", False) - if make_consistent: - self.material_buckets[:, 1] = torch.min(self.material_buckets[:, 0], self.material_buckets[:, 1]) - def __call__( self, env: ManagerBasedEnv, @@ -286,6 +249,174 @@ def __call__( ) +class _RandomizeRigidBodyMaterialNewton: + """Newton backend implementation for material randomization. + + Newton can assign arbitrary friction/restitution per shape (no bucket limitation). + Samples friction (mu) and restitution continuously from the given ranges. + Newton uses a single friction coefficient (mu), so ``dynamic_friction_range`` + and ``num_buckets`` are ignored. + """ + + def __init__( + self, cfg: EventTermCfg, env: ManagerBasedEnv, asset: RigidObject | Articulation, asset_cfg: SceneEntityCfg + ): + import isaaclab_newton.physics.newton_manager as newton_manager_module # noqa: PLC0415 + from newton.solvers import SolverNotifyFlags # noqa: PLC0415 + + from isaaclab.assets import BaseArticulation + + self.asset = asset + self.asset_cfg = asset_cfg + self._newton_manager = newton_manager_module.NewtonManager + self._notify_shape_properties = SolverNotifyFlags.SHAPE_PROPERTIES + + # cache friction/restitution ranges for continuous per-shape sampling + self._static_friction_range = cfg.params.get("static_friction_range", (1.0, 1.0)) + self._restitution_range = cfg.params.get("restitution_range", (0.0, 0.0)) + + # compute shape indices for body-specific randomization + if isinstance(asset, BaseArticulation) and asset_cfg.body_ids != slice(None): + num_shapes_per_body = asset.num_shapes_per_body + shape_indices_list = [] + for body_id in asset_cfg.body_ids: + start_idx = sum(num_shapes_per_body[:body_id]) + end_idx = start_idx + num_shapes_per_body[body_id] + shape_indices_list.extend(range(start_idx, end_idx)) + self._shape_indices = torch.tensor(shape_indices_list, dtype=torch.long) + else: + total_shapes = sum(asset.num_shapes_per_body) + self._shape_indices = torch.arange(total_shapes, dtype=torch.long) + + # get friction/restitution view-level bindings + model = self._newton_manager.get_model() + self._friction_binding = asset._root_view.get_attribute("shape_material_mu", model)[:, 0] # type: ignore + self._restitution_binding = asset._root_view.get_attribute("shape_material_restitution", model)[:, 0] # type: ignore + + def __call__( + self, + env: ManagerBasedEnv, + env_ids: torch.Tensor | None, + static_friction_range: tuple[float, float], + dynamic_friction_range: tuple[float, float], + restitution_range: tuple[float, float], + num_buckets: int, + asset_cfg: SceneEntityCfg, + make_consistent: bool = False, + ): + device = env.device + # resolve environment ids + if env_ids is None: + env_ids = torch.arange(env.scene.num_envs, device=device, dtype=torch.int32) + else: + env_ids = env_ids.to(device) + + num_shapes = len(self._shape_indices) + shape_idx = self._shape_indices.to(device) + + # sample friction (mu) and restitution continuously per shape + friction_range = torch.tensor(self._static_friction_range, device=device) + restitution_range_t = torch.tensor(self._restitution_range, device=device) + friction_samples = math_utils.sample_uniform( + friction_range[0], friction_range[1], (len(env_ids), num_shapes), device=device + ) + restitution_samples = math_utils.sample_uniform( + restitution_range_t[0], restitution_range_t[1], (len(env_ids), num_shapes), device=device + ) + + # write only the affected env_ids to the warp binding + friction_view = wp.to_torch(self._friction_binding) + restitution_view = wp.to_torch(self._restitution_binding) + friction_view[env_ids[:, None], shape_idx] = friction_samples + restitution_view[env_ids[:, None], shape_idx] = restitution_samples + + # notify the physics engine + self._newton_manager.add_model_change(self._notify_shape_properties) + + +class randomize_rigid_body_material(ManagerTermBase): + """Randomize the physics materials on all geometries of the asset. + + This function creates a set of physics materials with random static friction, dynamic friction, and restitution + values and assigns them to the geometries of the asset. + + Automatically detects the active physics backend (PhysX or Newton) and delegates to + the appropriate backend-specific implementation: + + - **PhysX**: Uses the 3-tuple material format (static_friction, dynamic_friction, restitution) + with bucket-based assignment (limited to 64000 unique materials). Applied via the PhysX + tensor API (``root_view.set_material_properties``). + - **Newton**: Samples friction (mu) and restitution continuously per shape (no bucket + limitation). Newton uses a single friction coefficient, so ``dynamic_friction_range`` + and ``num_buckets`` are ignored. Applied directly to Newton's view-level bindings. + + If the flag ``make_consistent`` is set to ``True``, the dynamic friction is set to be less than or equal to + the static friction (PhysX only). This obeys the physics constraint on friction values. + + .. attention:: + On PhysX, this function uses CPU tensors to assign the material properties. It is recommended to + use this function only during the initialization of the environment. + + .. note:: + PhysX only allows 64000 unique physics materials in the scene. If the number of materials exceeds this + limit, the simulation will crash. Due to this reason, we sample the materials only once during initialization. + Afterwards, these materials are randomly assigned to the geometries of the asset. + """ + + def __init__(self, cfg: EventTermCfg, env: ManagerBasedEnv): + """Initialize the term. + + Args: + cfg: The configuration of the event term. + env: The environment instance. + + Raises: + ValueError: If the asset is not a RigidObject or an Articulation. + """ + from isaaclab.assets import BaseArticulation, BaseRigidObject + + super().__init__(cfg, env) + + # extract the used quantities (to enable type-hinting) + self.asset_cfg: SceneEntityCfg = cfg.params["asset_cfg"] + self.asset: RigidObject | Articulation = env.scene[self.asset_cfg.name] + + if not isinstance(self.asset, (BaseRigidObject, BaseArticulation)): + raise ValueError( + f"Randomization term 'randomize_rigid_body_material' not supported for asset: '{self.asset_cfg.name}'" + f" with type: '{type(self.asset)}'." + ) + + # detect physics backend and instantiate the appropriate implementation + manager_name = env.sim.physics_manager.__name__.lower() + if "newton" in manager_name: + self._impl = _RandomizeRigidBodyMaterialNewton(cfg, env, self.asset, self.asset_cfg) + else: + self._impl = _RandomizeRigidBodyMaterialPhysx(cfg, env, self.asset, self.asset_cfg) + + def __call__( + self, + env: ManagerBasedEnv, + env_ids: torch.Tensor | None, + static_friction_range: tuple[float, float], + dynamic_friction_range: tuple[float, float], + restitution_range: tuple[float, float], + num_buckets: int, + asset_cfg: SceneEntityCfg, + make_consistent: bool = False, + ): + self._impl( + env, + env_ids, + static_friction_range, + dynamic_friction_range, + restitution_range, + num_buckets, + asset_cfg, + make_consistent, + ) + + class randomize_rigid_body_mass(ManagerTermBase): """Randomize the mass of the bodies by adding, scaling, or setting random values. @@ -405,104 +536,419 @@ def __call__( ) -def randomize_rigid_body_com( - env: ManagerBasedEnv, - env_ids: torch.Tensor | None, - com_range: dict[str, tuple[float, float]], - asset_cfg: SceneEntityCfg, -): - """Randomize the center of mass (CoM) of rigid bodies by adding a random value sampled from the given ranges. +class randomize_rigid_body_inertia(ManagerTermBase): + """Randomize the inertia tensor of rigid bodies by adding, scaling, or setting values. + + This function modifies body inertia tensors independently of mass. The inertia tensor + is a 3x3 symmetric matrix stored as 9 elements: ``[Ixx, Ixy, Ixz, Iyx, Iyy, Iyz, Izx, Izy, Izz]``. + + Two modes are supported via the :attr:`diagonal_only` parameter: + + - **diagonal_only=True** (default): Only modifies diagonal elements (Ixx, Iyy, Izz at + indices 0, 4, 8). This is useful for adding numerical stability (armature/regularization) + without changing rotational coupling between axes. The diagonal elements represent + resistance to rotation about each principal axis. + + - **diagonal_only=False**: Modifies all 9 elements of the inertia tensor. This can + simulate manufacturing variations or asymmetric mass distributions. Off-diagonal + elements represent coupling between rotations about different axes. .. note:: - This function does not properly track the original CoM values. It is recommended to use this function - only once, during the initialization of the environment. + Unlike :class:`randomize_rigid_body_mass` which recomputes inertia based on mass + ratios, this function modifies inertia directly without affecting mass. """ - # extract the used quantities (to enable type-hinting) - asset: Articulation = env.scene[asset_cfg.name] - # resolve environment ids - if env_ids is None: - env_ids = torch.arange(env.scene.num_envs, device=asset.device) - else: - env_ids = env_ids.to(asset.device) - # resolve body indices - if asset_cfg.body_ids == slice(None): - body_ids = torch.arange(asset.num_bodies, dtype=torch.int, device=asset.device) - else: - body_ids = torch.tensor(asset_cfg.body_ids, dtype=torch.int, device=asset.device) + def __init__(self, cfg: EventTermCfg, env: ManagerBasedEnv): + """Initialize the term. - # sample random CoM values - range_list = [com_range.get(key, (0.0, 0.0)) for key in ["x", "y", "z"]] - ranges = torch.tensor(range_list, device=asset.device) - rand_samples = math_utils.sample_uniform( - ranges[:, 0], ranges[:, 1], (len(env_ids), 3), device=asset.device - ).unsqueeze(1) + Args: + cfg: The configuration of the event term. + env: The environment instance. - # get the current com of the bodies (num_assets, num_bodies) - coms = wp.to_torch(asset.data.body_com_pose_b).clone() + Raises: + ValueError: If the operation is not supported. + ValueError: If the lower bound is negative or zero when not allowed for scale operation. + ValueError: If the upper bound is less than the lower bound. + """ + from isaaclab.assets import BaseArticulation, BaseRigidObject, BaseRigidObjectCollection - # Randomize the com in range - coms[env_ids[:, None], body_ids, :3] += rand_samples + super().__init__(cfg, env) - # Set the new coms - asset.set_coms_index(coms=coms, env_ids=env_ids) + # extract the used quantities (to enable type-hinting) + self.asset_cfg: SceneEntityCfg = cfg.params["asset_cfg"] + self.asset: RigidObject | Articulation = env.scene[self.asset_cfg.name] + if not isinstance(self.asset, (BaseRigidObject, BaseArticulation, BaseRigidObjectCollection)): + raise ValueError( + f"Randomization term 'randomize_rigid_body_inertia' not supported for asset: '{self.asset_cfg.name}'" + f" with type: '{type(self.asset)}'." + ) -def randomize_rigid_body_collider_offsets( - env: ManagerBasedEnv, - env_ids: torch.Tensor | None, - asset_cfg: SceneEntityCfg, - rest_offset_distribution_params: tuple[float, float] | None = None, - contact_offset_distribution_params: tuple[float, float] | None = None, - distribution: Literal["uniform", "log_uniform", "gaussian"] = "uniform", -): - """Randomize the collider parameters of rigid bodies in an asset by adding, scaling, or setting random values. + # check for valid operation + if cfg.params["operation"] == "scale": + if "inertia_distribution_params" in cfg.params: + _validate_scale_range( + cfg.params["inertia_distribution_params"], "inertia_distribution_params", allow_zero=False + ) + elif cfg.params["operation"] not in ("abs", "add"): + raise ValueError( + f"Randomization term 'randomize_rigid_body_inertia' does not support operation:" + f" '{cfg.params['operation']}'." + ) + + self.default_inertia = None + # cache inertia indices: diagonal (0, 4, 8) for regularization, or all elements + diagonal_only = cfg.params.get("diagonal_only", True) + self._inertia_idx = torch.tensor([0, 4, 8], device=self.asset.device) if diagonal_only else slice(None) + + def __call__( + self, + env: ManagerBasedEnv, + env_ids: torch.Tensor | None, + asset_cfg: SceneEntityCfg, + inertia_distribution_params: tuple[float, float], + operation: Literal["add", "scale", "abs"] = "add", + distribution: Literal["uniform", "log_uniform", "gaussian"] = "uniform", + diagonal_only: bool = True, + ): + """Randomize body inertia tensors. + + Args: + env: The environment instance. + env_ids: The environment indices to randomize. If None, all environments are randomized. + asset_cfg: The asset configuration specifying the asset and body names. + inertia_distribution_params: Distribution parameters as a tuple of two floats + ``(low, high)`` for sampling inertia modification values. + operation: The operation to apply. Options: ``"add"``, ``"scale"``, ``"abs"``. + Defaults to ``"add"`` which is typical for regularization/armature. + distribution: The distribution to sample from. Options: ``"uniform"``, + ``"log_uniform"``, ``"gaussian"``. Defaults to ``"uniform"``. + diagonal_only: If True, only modify diagonal elements (Ixx, Iyy, Izz) for + numerical stability. If False, modify all 9 elements. Defaults to True. + """ + # store default inertia on first call for repeatable randomization + if self.default_inertia is None: + self.default_inertia = wp.to_torch(self.asset.data.body_inertia).clone() + + # resolve environment ids + if env_ids is None: + env_ids = torch.arange(env.scene.num_envs, device=self.asset.device, dtype=torch.int32) + else: + env_ids = env_ids.to(self.asset.device) + + # resolve body indices + if self.asset_cfg.body_ids == slice(None): + body_ids = torch.arange(self.asset.num_bodies, dtype=torch.int32, device=self.asset.device) + else: + body_ids = torch.tensor(self.asset_cfg.body_ids, dtype=torch.int32, device=self.asset.device) + + # get default inertias for affected envs/bodies (advanced indexing creates a copy) + # shape: (len(env_ids), len(body_ids), 9) + inertias = self.default_inertia[env_ids[:, None], body_ids] + + # resolve the distribution function + if distribution == "uniform": + dist_fn = math_utils.sample_uniform + elif distribution == "log_uniform": + dist_fn = math_utils.sample_log_uniform + elif distribution == "gaussian": + dist_fn = math_utils.sample_gaussian + else: + raise NotImplementedError( + f"Unknown distribution: '{distribution}' for inertia randomization." + " Please use 'uniform', 'log_uniform', or 'gaussian'." + ) + + # sample random values once per (env, body) - shape: (len(env_ids), len(body_ids)) + random_values = dist_fn(*inertia_distribution_params, (len(env_ids), len(body_ids)), device=self.asset.device) + + # apply the operation with the SAME random value per body + if operation == "add": + inertias[:, :, self._inertia_idx] += random_values[..., None] + elif operation == "scale": + inertias[:, :, self._inertia_idx] *= random_values[..., None] + elif operation == "abs": + inertias[:, :, self._inertia_idx] = random_values[..., None] + + # set the inertia tensors into the physics simulation + self.asset.set_inertias_index(inertias=inertias, body_ids=body_ids, env_ids=env_ids) + + +class randomize_rigid_body_com(ManagerTermBase): + """Randomize the center of mass (CoM) of rigid bodies by adding a random value sampled from the given ranges. + + This class tracks the original CoM values and randomizes from those defaults on each call, + ensuring repeatable randomization across resets. + + Automatically detects the active physics backend: + + - **PhysX**: Passes the full CoM pose (position + quaternion) to ``set_coms_index``. + - **Newton**: Passes position-only (vec3) to ``set_coms_index``. Note that on Newton + (MuJoCo Warp), runtime CoM changes may cause simulation instability because + ``notify_model_changed(BODY_INERTIAL_PROPERTIES)`` does not fully recompute the + mass matrix after ``body_ipos`` changes. Use with caution until this is fixed upstream. + """ + + def __init__(self, cfg: EventTermCfg, env: ManagerBasedEnv): + """Initialize the term. + + Args: + cfg: The configuration of the event term. + env: The environment instance. + """ + super().__init__(cfg, env) + + self.asset_cfg: SceneEntityCfg = cfg.params["asset_cfg"] + self.asset: RigidObject | Articulation = env.scene[self.asset_cfg.name] + + # detect physics backend + manager_name = env.sim.physics_manager.__name__.lower() + self._is_newton = "newton" in manager_name + + self.default_com = None + + def __call__( + self, + env: ManagerBasedEnv, + env_ids: torch.Tensor | None, + com_range: dict[str, tuple[float, float]], + asset_cfg: SceneEntityCfg, + ): + # store default CoM on first call for repeatable randomization + if self.default_com is None: + self.default_com = wp.to_torch(self.asset.data.body_com_pose_b).clone() + + # resolve environment ids + if env_ids is None: + env_ids = torch.arange(env.scene.num_envs, device=self.asset.device) + else: + env_ids = env_ids.to(self.asset.device) + + # resolve body indices + if self.asset_cfg.body_ids == slice(None): + body_ids = torch.arange(self.asset.num_bodies, dtype=torch.int, device=self.asset.device) + else: + body_ids = torch.tensor(self.asset_cfg.body_ids, dtype=torch.int, device=self.asset.device) + + # sample random CoM values + range_list = [com_range.get(key, (0.0, 0.0)) for key in ["x", "y", "z"]] + ranges = torch.tensor(range_list, device=self.asset.device) + rand_samples = math_utils.sample_uniform( + ranges[:, 0], ranges[:, 1], (len(env_ids), 3), device=self.asset.device + ).unsqueeze(1) + + # start from defaults and add random offsets + coms = self.default_com.clone() + coms[env_ids[:, None], body_ids, :3] += rand_samples + + # Newton expects position-only (vec3f), PhysX expects the full pose (pos + quat) + # note: pass partial data of shape (len(env_ids), len(body_ids), ...) to match the API + if self._is_newton: + self.asset.set_coms_index(coms=coms[env_ids[:, None], body_ids, :3], body_ids=body_ids, env_ids=env_ids) + else: + self.asset.set_coms_index(coms=coms[env_ids[:, None], body_ids], body_ids=body_ids, env_ids=env_ids) + + +class _RandomizeRigidBodyColliderOffsetsPhysx: + """PhysX backend implementation for collider offset randomization. + + Uses rest offset and contact offset directly via the PhysX tensor API. + """ + + def __init__(self, asset: RigidObject | Articulation): + self.asset = asset + self.default_rest_offsets = wp.to_torch(asset.root_view.get_rest_offsets()).clone() + self.default_contact_offsets = wp.to_torch(asset.root_view.get_contact_offsets()).clone() + + def __call__( + self, + env: ManagerBasedEnv, + env_ids: torch.Tensor | None, + asset_cfg: SceneEntityCfg, + rest_offset_distribution_params: tuple[float, float] | None = None, + contact_offset_distribution_params: tuple[float, float] | None = None, + distribution: Literal["uniform", "log_uniform", "gaussian"] = "uniform", + ): + if env_ids is None: + env_ids = torch.arange(env.scene.num_envs, device="cpu") + else: + env_ids = env_ids.cpu() + + if rest_offset_distribution_params is not None: + rest_offset = self.default_rest_offsets.clone() + rest_offset = _randomize_prop_by_op( + rest_offset, + rest_offset_distribution_params, + None, + slice(None), + operation="abs", + distribution=distribution, + ) + self.asset.root_view.set_rest_offsets(rest_offset, env_ids) + + if contact_offset_distribution_params is not None: + contact_offset = self.default_contact_offsets.clone() + contact_offset = _randomize_prop_by_op( + contact_offset, + contact_offset_distribution_params, + None, + slice(None), + operation="abs", + distribution=distribution, + ) + self.asset.root_view.set_contact_offsets(contact_offset, env_ids) + + +class _RandomizeRigidBodyColliderOffsetsNewton: + """Newton backend implementation for collider offset randomization. + + Maps PhysX concepts to Newton's geometry properties: + + - ``rest_offset`` -> ``shape_margin`` (Newton margin) + - ``contact_offset`` -> ``shape_gap`` (Newton gap = contact_offset - margin) + + See the `Newton collision schema`_ for details. + + .. _Newton collision schema: https://newton-physics.github.io/newton/latest/concepts/collisions.html + """ + + def __init__(self, asset: RigidObject | Articulation): + import isaaclab_newton.physics.newton_manager as newton_manager_module # noqa: PLC0415 + from newton.solvers import SolverNotifyFlags # noqa: PLC0415 + + self.asset = asset + self._newton_manager = newton_manager_module.NewtonManager + self._notify_shape_properties = SolverNotifyFlags.SHAPE_PROPERTIES + + model = self._newton_manager.get_model() + self._sim_bind_shape_margin = asset._root_view.get_attribute("shape_margin", model)[:, 0] # type: ignore + self._sim_bind_shape_gap = asset._root_view.get_attribute("shape_gap", model)[:, 0] # type: ignore + + self.default_margin = wp.to_torch(self._sim_bind_shape_margin).clone() + self.default_gap = wp.to_torch(self._sim_bind_shape_gap).clone() + + def __call__( + self, + env: ManagerBasedEnv, + env_ids: torch.Tensor | None, + asset_cfg: SceneEntityCfg, + rest_offset_distribution_params: tuple[float, float] | None = None, + contact_offset_distribution_params: tuple[float, float] | None = None, + distribution: Literal["uniform", "log_uniform", "gaussian"] = "uniform", + ): + device = env.device + if env_ids is None: + env_ids = torch.arange(env.scene.num_envs, device=device, dtype=torch.int32) + else: + env_ids = env_ids.to(device) + + margin_view = wp.to_torch(self._sim_bind_shape_margin) + + if rest_offset_distribution_params is not None: + margin = self.default_margin.clone() + margin = _randomize_prop_by_op( + margin, + rest_offset_distribution_params, + None, + slice(None), + operation="abs", + distribution=distribution, + ) + self.default_margin[env_ids] = margin[env_ids] + margin_view[env_ids] = margin[env_ids] + + if contact_offset_distribution_params is not None: + current_margin = self.default_margin + contact_offset = torch.zeros_like(self.default_gap) + contact_offset = _randomize_prop_by_op( + contact_offset, + contact_offset_distribution_params, + None, + slice(None), + operation="abs", + distribution=distribution, + ) + gap = torch.clamp(contact_offset - current_margin, min=0.0) + self.default_gap[env_ids] = gap[env_ids] + gap_view = wp.to_torch(self._sim_bind_shape_gap) + gap_view[env_ids] = gap[env_ids] + + if rest_offset_distribution_params is not None or contact_offset_distribution_params is not None: + self._newton_manager.add_model_change(self._notify_shape_properties) + + +class randomize_rigid_body_collider_offsets(ManagerTermBase): + """Randomize the collider parameters of rigid bodies by setting random values. This function allows randomizing the collider parameters of the asset, such as rest and contact offsets. - These correspond to the physics engine collider properties that affect the collision checking. + These correspond to the physics engine collider properties that affect collision checking. - The function samples random values from the given distribution parameters and applies the operation to - the collider properties. It then sets the values into the physics simulation. If the distribution parameters - are not provided for a particular property, the function does not modify the property. + Automatically detects the active physics backend (PhysX or Newton) and delegates to + the appropriate backend-specific implementation: + + - **PhysX**: Uses rest offset and contact offset directly via the PhysX tensor API + (``root_view.set_rest_offsets`` / ``root_view.set_contact_offsets``). + - **Newton**: Maps PhysX concepts to Newton's geometry properties. PhysX ``rest_offset`` + maps to Newton ``shape_margin``, and PhysX ``contact_offset`` is converted to Newton + ``shape_gap`` via ``gap = contact_offset - margin``. - Currently, the distribution parameters are applied as absolute values. + The function samples random values from the given distribution parameters and applies them + as absolute values to the collider properties. If the distribution parameters are not + provided for a particular property, the function does not modify it. .. tip:: - This function uses CPU tensors to assign the collision properties. It is recommended to use this function - only during the initialization of the environment. + This function uses CPU tensors (PhysX) or GPU tensors (Newton) to assign the collision + properties. It is recommended to use this function only during the initialization of + the environment. """ - # extract the used quantities (to enable type-hinting) - asset: RigidObject | Articulation = env.scene[asset_cfg.name] - # resolve environment ids - if env_ids is None: - env_ids = torch.arange(env.scene.num_envs, device="cpu") + def __init__(self, cfg: EventTermCfg, env: ManagerBasedEnv): + """Initialize the term. - # sample collider properties from the given ranges and set into the physics simulation - # -- rest offsets - if rest_offset_distribution_params is not None: - rest_offset = wp.to_torch(asset.root_view.get_rest_offsets()).clone() - rest_offset = _randomize_prop_by_op( - rest_offset, + Args: + cfg: The configuration of the event term. + env: The environment instance. + + Raises: + ValueError: If the asset is not a RigidObject or an Articulation. + """ + from isaaclab.assets import BaseArticulation, BaseRigidObject + + super().__init__(cfg, env) + + self.asset_cfg: SceneEntityCfg = cfg.params["asset_cfg"] + self.asset: RigidObject | Articulation = env.scene[self.asset_cfg.name] + + if not isinstance(self.asset, (BaseRigidObject, BaseArticulation)): + raise ValueError( + f"Randomization term 'randomize_rigid_body_collider_offsets' not supported for asset:" + f" '{self.asset_cfg.name}' with type: '{type(self.asset)}'." + ) + + # detect physics backend and instantiate the appropriate implementation + manager_name = env.sim.physics_manager.__name__.lower() + if "newton" in manager_name: + self._impl = _RandomizeRigidBodyColliderOffsetsNewton(self.asset) + else: + self._impl = _RandomizeRigidBodyColliderOffsetsPhysx(self.asset) + + def __call__( + self, + env: ManagerBasedEnv, + env_ids: torch.Tensor | None, + asset_cfg: SceneEntityCfg, + rest_offset_distribution_params: tuple[float, float] | None = None, + contact_offset_distribution_params: tuple[float, float] | None = None, + distribution: Literal["uniform", "log_uniform", "gaussian"] = "uniform", + ): + self._impl( + env, + env_ids, + asset_cfg, rest_offset_distribution_params, - None, - slice(None), - operation="abs", - distribution=distribution, - ) - asset.root_view.set_rest_offsets(rest_offset, env_ids.cpu()) - # -- contact offsets - if contact_offset_distribution_params is not None: - contact_offset = wp.to_torch(asset.root_view.get_contact_offsets()).clone() - contact_offset = _randomize_prop_by_op( - contact_offset, contact_offset_distribution_params, - None, - slice(None), - operation="abs", - distribution=distribution, + distribution, ) - asset.root_view.set_contact_offsets(contact_offset, env_ids.cpu()) class randomize_physics_scene_gravity(ManagerTermBase): @@ -802,14 +1248,26 @@ def __init__(self, cfg: EventTermCfg, env: ManagerBasedEnv): # extract the used quantities (to enable type-hinting) self.asset_cfg: SceneEntityCfg = cfg.params["asset_cfg"] - self.asset: RigidObject | Articulation = env.scene[self.asset_cfg.name] + self.asset: Articulation = env.scene[self.asset_cfg.name] + # detect physics backend + manager_name = env.sim.physics_manager.__name__.lower() + self._backend = "newton" if "newton" in manager_name else "physx" + + # cache default values (common to both backends) self.default_joint_friction_coeff = wp.to_torch(self.asset.data.joint_friction_coeff).clone() - self.default_dynamic_joint_friction_coeff = wp.to_torch(self.asset.data.joint_dynamic_friction_coeff).clone() - self.default_viscous_joint_friction_coeff = wp.to_torch(self.asset.data.joint_viscous_friction_coeff).clone() self.default_joint_armature = wp.to_torch(self.asset.data.joint_armature).clone() self.default_joint_pos_limits = wp.to_torch(self.asset.data.joint_pos_limits).clone() + # cache dynamic/viscous friction (PhysX only - Newton only has static friction) + if self._backend == "physx": + self.default_dynamic_joint_friction_coeff = wp.to_torch( + self.asset.data.joint_dynamic_friction_coeff + ).clone() + self.default_viscous_joint_friction_coeff = wp.to_torch( + self.asset.data.joint_viscous_friction_coeff + ).clone() + # check for valid operation if cfg.params["operation"] == "scale": if "friction_distribution_params" in cfg.params: @@ -867,8 +1325,14 @@ def __call__( # Always set static friction (indexed once) static_friction_coeff = friction_coeff[env_ids_for_slice, joint_ids] - # if isaacsim version is lower than 5.0.0 we can set only the static friction coefficient - if get_isaac_sim_version().major >= 5: + if self._backend == "newton": + # Newton only supports static friction coefficient + self.asset.write_joint_friction_coefficient_to_sim_index( + joint_friction_coeff=static_friction_coeff, + joint_ids=joint_ids, + env_ids=env_ids, + ) + else: # Randomize raw tensors dynamic_friction_coeff = _randomize_prop_by_op( self.default_dynamic_joint_friction_coeff.clone(), @@ -897,19 +1361,15 @@ def __call__( # Index once at the end dynamic_friction_coeff = dynamic_friction_coeff[env_ids_for_slice, joint_ids] viscous_friction_coeff = viscous_friction_coeff[env_ids_for_slice, joint_ids] - else: - # For versions < 5.0.0, we do not set these values - dynamic_friction_coeff = None - viscous_friction_coeff = None - - # Single write call for all versions - self.asset.write_joint_friction_coefficient_to_sim_index( - joint_friction_coeff=static_friction_coeff, - joint_dynamic_friction_coeff=dynamic_friction_coeff, - joint_viscous_friction_coeff=viscous_friction_coeff, - joint_ids=joint_ids, - env_ids=env_ids, - ) + + # Single write call for all versions + self.asset.write_joint_friction_coefficient_to_sim_index( + joint_friction_coeff=static_friction_coeff, + joint_dynamic_friction_coeff=dynamic_friction_coeff, + joint_viscous_friction_coeff=viscous_friction_coeff, + joint_ids=joint_ids, + env_ids=env_ids, + ) # joint armature if armature_distribution_params is not None: diff --git a/source/isaaclab_newton/config/extension.toml b/source/isaaclab_newton/config/extension.toml index 154968319a0c..9cfd4c5fc105 100644 --- a/source/isaaclab_newton/config/extension.toml +++ b/source/isaaclab_newton/config/extension.toml @@ -1,7 +1,7 @@ [package] # Note: Semantic Versioning is used: https://semver.org/ -version = "0.5.9" +version = "0.5.11" # Description title = "Newton simulation interfaces for IsaacLab core package" diff --git a/source/isaaclab_newton/docs/CHANGELOG.rst b/source/isaaclab_newton/docs/CHANGELOG.rst index 59a80364d19f..1b9ee9788e04 100644 --- a/source/isaaclab_newton/docs/CHANGELOG.rst +++ b/source/isaaclab_newton/docs/CHANGELOG.rst @@ -1,6 +1,51 @@ Changelog --------- +0.5.11 (2026-04-21) +~~~~~~~~~~~~~~~~~~~ + +Added +^^^^^ + +* Initial scaffolding for Newton inverse-kinematics integration. +* Added :class:`~isaaclab_newton.controllers.IKModelInfo` and + :func:`~isaaclab_newton.controllers.build_single_arm_ik_model` for + constructing a single-arm Newton model from an IsaacLab asset cfg. +* Added :class:`~isaaclab_newton.controllers.NewtonIKControllerCfg` + configuration dataclass for the Newton inverse-kinematics controller. +* Added :class:`~isaaclab_newton.controllers.NewtonIKController` skeleton + with batched IK solver construction across all envs. +* Added :class:`~isaaclab_newton.envs.mdp.actions.NewtonInverseKinematicsActionCfg` + configuration dataclass for the Newton IK action term. +* Added :class:`~isaaclab_newton.envs.mdp.actions.NewtonInverseKinematicsAction` + action term that drives Newton's batched IK across all envs in one fused + pass per env step. +* Added :meth:`~isaaclab_newton.controllers.NewtonIKController.compute` + for batched IK solving across all envs in one fused pass. +* Added Gym task ``Isaac-Reach-Franka-Newton-IK-Rel-v0`` (and ``Isaac-Reach-Franka-Newton-IK-Rel-Play-v0``) + using Newton physics with the Newton IK action term for relative-pose EE + control. + + +0.5.10 (2026-04-02) +~~~~~~~~~~~~~~~~~~~ + +Added +^^^^^ + +* Added :class:`~isaaclab_newton.physics.NewtonCollisionPipelineCfg` for configuring + Newton's collision pipeline parameters (broad phase, contact limits, triangle pairs). + Set via :attr:`~isaaclab_newton.physics.NewtonCfg.collision_cfg`. +* Added :class:`~isaaclab_newton.physics.SDFCfg` for configuring SDF-based mesh + collisions via Newton's ``mesh.build_sdf()`` API. Supports per-body and per-shape + regex pattern matching, per-pattern resolution overrides, and optional creation of + collision shapes from visual meshes. +* Added :class:`~isaaclab_newton.physics.HydroelasticCfg` for distributed surface + contacts computed via SDF overlap with marching cubes. +* Added SDF pattern skip in the Newton cloner to preserve original triangle + meshes for shapes that will use SDF collision. + + 0.5.9 (2026-03-16) ~~~~~~~~~~~~~~~~~~ diff --git a/source/isaaclab_newton/isaaclab_newton/assets/rigid_object/rigid_object.py b/source/isaaclab_newton/isaaclab_newton/assets/rigid_object/rigid_object.py index 5e02e3622985..4df2f8923776 100644 --- a/source/isaaclab_newton/isaaclab_newton/assets/rigid_object/rigid_object.py +++ b/source/isaaclab_newton/isaaclab_newton/assets/rigid_object/rigid_object.py @@ -94,6 +94,19 @@ def root_view(self) -> ArticulationView: """ return self._root_view + @property + def num_shapes_per_body(self) -> list[int]: + """Number of collision shapes per body in the rigid object. + + Returns a list where each element is the number of shapes for + the corresponding body. Cached after first access. + """ + if not hasattr(self, "_num_shapes_per_body"): + self._num_shapes_per_body = [] + for shapes in self._root_view.body_shapes: + self._num_shapes_per_body.append(len(shapes)) + return self._num_shapes_per_body + @property def instantaneous_wrench_composer(self) -> WrenchComposer: """Instantaneous wrench composer. diff --git a/source/isaaclab_newton/isaaclab_newton/cloner/newton_replicate.py b/source/isaaclab_newton/isaaclab_newton/cloner/newton_replicate.py index 59e5d2476ccc..69f255587cf9 100644 --- a/source/isaaclab_newton/isaaclab_newton/cloner/newton_replicate.py +++ b/source/isaaclab_newton/isaaclab_newton/cloner/newton_replicate.py @@ -5,15 +5,17 @@ from __future__ import annotations +import re from collections.abc import Callable import torch import warp as wp -from newton import ModelBuilder, solvers +from newton import GeoType, ModelBuilder, solvers from newton._src.usd.schemas import SchemaResolverNewton, SchemaResolverPhysx from pxr import Usd, UsdGeom +from isaaclab.physics import PhysicsManager from isaaclab.physics.scene_data_requirements import VisualizerPrebuiltArtifacts from isaaclab_newton.physics import NewtonManager @@ -62,6 +64,17 @@ def _build_newton_builder_from_mapping( # The prototype is built from env_0 in absolute world coordinates. # add_builder xforms are deltas from env_0 so positions don't get double-counted. env0_pos = positions[0] + + # SDF collision requires original triangle meshes for mesh.build_sdf(). + # Convex hull approximation destroys the source geometry, so shapes + # matching SDF patterns must be excluded from approximation here. + # _apply_sdf_config() builds the SDF on each prototype after approximation. + cfg = PhysicsManager._cfg + sdf_cfg = getattr(cfg, "sdf_cfg", None) if cfg is not None else None + body_pats = [re.compile(x) for x in sdf_cfg.body_patterns] if sdf_cfg and sdf_cfg.body_patterns else None + shape_pats = [re.compile(x) for x in sdf_cfg.shape_patterns] if sdf_cfg and sdf_cfg.shape_patterns else None + has_sdf_patterns = body_pats is not None or shape_pats is not None + protos: dict[str, ModelBuilder] = {} for src_path in sources: p = ModelBuilder(up_axis=up_axis) @@ -74,7 +87,33 @@ def _build_newton_builder_from_mapping( schema_resolvers=schema_resolvers, ) if simplify_meshes: - p.approximate_meshes("convex_hull", keep_visual_shapes=True) + if has_sdf_patterns: + sdf_bodies: set[int] = set() + if body_pats is not None: + for bi in range(len(p.body_label)): + if any(pat.search(p.body_label[bi]) for pat in body_pats): + sdf_bodies.add(bi) + + approx_indices = [] + for i in range(len(p.shape_type)): + if p.shape_type[i] != GeoType.MESH: + continue + # Skip shapes that will use SDF (matched by body or shape pattern) + if p.shape_body[i] in sdf_bodies: + continue + if shape_pats is not None: + lbl = p.shape_label[i] if i < len(p.shape_label) else "" + if any(pat.search(lbl) for pat in shape_pats): + continue + approx_indices.append(i) + if approx_indices: + p.approximate_meshes("convex_hull", shape_indices=approx_indices, keep_visual_shapes=True) + else: + p.approximate_meshes("convex_hull", keep_visual_shapes=True) + # Build SDF on prototype before add_builder copies it N times. + # Mesh objects are shared by reference, so SDF is built once and + # all environments inherit it. + NewtonManager._apply_sdf_config(p) protos[src_path] = p # create a separate world for each environment (heterogeneous spawning) diff --git a/source/isaaclab_newton/isaaclab_newton/controllers/__init__.py b/source/isaaclab_newton/isaaclab_newton/controllers/__init__.py new file mode 100644 index 000000000000..1b4a8fed74c6 --- /dev/null +++ b/source/isaaclab_newton/isaaclab_newton/controllers/__init__.py @@ -0,0 +1,17 @@ +# 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-native controllers for IsaacLab (IK, etc.).""" + +from .newton_ik import NewtonIKController +from .newton_ik_cfg import NewtonIKControllerCfg +from .newton_ik_model_builder import IKModelInfo, build_single_arm_ik_model + +__all__ = [ + "IKModelInfo", + "NewtonIKController", + "NewtonIKControllerCfg", + "build_single_arm_ik_model", +] diff --git a/source/isaaclab_newton/isaaclab_newton/controllers/newton_ik.py b/source/isaaclab_newton/isaaclab_newton/controllers/newton_ik.py new file mode 100644 index 000000000000..fc2f1451c133 --- /dev/null +++ b/source/isaaclab_newton/isaaclab_newton/controllers/newton_ik.py @@ -0,0 +1,243 @@ +# 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 + +"""Batched Newton-native inverse-kinematics controller for IsaacLab.""" + +from __future__ import annotations + +import newton +import newton.ik as ik +import numpy as np +import torch +import warp as wp + +from .newton_ik_cfg import NewtonIKControllerCfg + + +class NewtonIKController: + """Batched optimization-based IK controller backed by :mod:`newton.ik`. + + Each :meth:`compute` call runs ``cfg.iterations`` optimizer steps for all + ``num_envs`` problems in one fused Warp kernel launch. Targets are + end-effector pose(s) in the **robot-base frame**; the IK model is built + with its base at the world origin so no per-env frame correction is needed + inside this class — the action term is responsible for converting world + poses into base-frame targets before calling :meth:`set_command`. + + Args: + cfg: Controller configuration. + ik_model: Single-articulation Newton model (built via + :func:`~isaaclab_newton.controllers.build_single_arm_ik_model`). + num_envs: Number of environments solved together per :meth:`compute`. + ee_link_index: Index of the EE body in ``ik_model.body_q``. + arm_dof_count: Number of arm-controlled DOFs (IK-relevant subset). + When provided, :attr:`num_arm_dofs` returns this value. When + ``None``, falls back to ``ik_model.joint_coord_count``. + ee_link_offset: Translation [m] of the constrained point in the EE + link's local frame. + ee_link_offset_rot: Quaternion ``(x, y, z, w)`` of the constrained + frame in the EE link's local frame. + device: Warp device string. + """ + + def __init__( + self, + cfg: NewtonIKControllerCfg, + ik_model: newton.Model, + num_envs: int, + ee_link_index: int, + arm_dof_count: int | None = None, + ee_link_offset: tuple[float, float, float] = (0.0, 0.0, 0.0), + ee_link_offset_rot: tuple[float, float, float, float] = (0.0, 0.0, 0.0, 1.0), + device: str = "cuda:0", + ) -> None: + # Validate (calls validate_config hook via @configclass). + cfg.validate_config() + self.cfg = cfg + self.model = ik_model + self.num_envs = num_envs + self.device = device + + self._n_coords = ik_model.joint_coord_count + self._arm_dof_count = arm_dof_count if arm_dof_count is not None else self._n_coords + + # Per-env target buffers (allocated empty; updated by set_command). + self._target_pos_wp = wp.zeros(num_envs, dtype=wp.vec3, device=device) + # Initialize quat targets to identity so quat residual is well-defined on first solve. + ident_np = np.tile(np.array([0.0, 0.0, 0.0, 1.0], dtype=np.float32), (num_envs, 1)) + self._target_rot_wp = wp.array(ident_np, dtype=wp.vec4, device=device) + + # Objectives. + self._pos_obj = ik.IKObjectivePosition( + link_index=ee_link_index, + link_offset=wp.vec3(*ee_link_offset), + target_positions=self._target_pos_wp, + weight=cfg.position_weight, + ) + self._rot_obj = ik.IKObjectiveRotation( + link_index=ee_link_index, + link_offset_rotation=wp.quat(*ee_link_offset_rot), + target_rotations=self._target_rot_wp, + weight=cfg.rotation_weight, + ) + self._limit_obj = ik.IKObjectiveJointLimit( + joint_limit_lower=ik_model.joint_limit_lower, + joint_limit_upper=ik_model.joint_limit_upper, + weight=cfg.joint_limit_weight, + ) + + objectives: list[ik.IKObjective] = [self._pos_obj, self._rot_obj, self._limit_obj] + if cfg.command_type == "position": + objectives = [self._pos_obj, self._limit_obj] + + self._solver = ik.IKSolver( + model=ik_model, + n_problems=num_envs, + objectives=objectives, + optimizer=cfg.optimizer, + jacobian_mode=cfg.jacobian_mode, + sampler=cfg.sampler, + n_seeds=cfg.n_seeds, + noise_std=cfg.noise_std, + rng_seed=cfg.rng_seed, + lambda_initial=cfg.lambda_initial, + lambda_factor=cfg.lambda_factor, + lambda_min=cfg.lambda_min, + lambda_max=cfg.lambda_max, + rho_min=cfg.rho_min, + history_len=cfg.history_len, + h0_scale=cfg.h0_scale, + ) + + self._joint_q_in_wp = wp.zeros((num_envs, self._n_coords), dtype=wp.float32, device=device) + self._joint_q_out_wp = wp.zeros((num_envs, self._n_coords), dtype=wp.float32, device=device) + self._previous_solution_wp: wp.array | None = None + if cfg.seed_source == "previous_solution": + self._previous_solution_wp = wp.zeros_like(self._joint_q_in_wp) + + @property + def action_dim(self) -> int: + """Action dimension implied by ``cfg`` (3, 6, or 7).""" + return self.cfg.action_dim() + + @property + def num_arm_dofs(self) -> int: + """Number of arm-controlled joint DOFs. + + When ``arm_dof_count`` is passed to the constructor, this returns that + value (the IK-controlled subset). Otherwise falls back to the full + ``joint_coord_count`` of the IK model. + """ + return self._arm_dof_count + + def set_command( + self, + command: torch.Tensor, + ee_pos: torch.Tensor | None = None, + ee_quat: torch.Tensor | None = None, + ) -> None: + """Set the EE pose target(s) for the next :meth:`compute` call. + + Args: + command: Target command tensor, shape ``[num_envs, action_dim]``. + Interpretation depends on :attr:`cfg.command_type` and + :attr:`cfg.use_relative_mode`. For ``pose`` + relative this is + ``(dx, dy, dz, drx, dry, drz)``. + ee_pos: Current EE position [m], shape ``[num_envs, 3]``. Required + for any relative mode. + ee_quat: Current EE orientation ``(x, y, z, w)``, shape + ``[num_envs, 4]``. Required for any pose mode with relative. + """ + from isaaclab.utils.math import apply_delta_pose + + if self.cfg.command_type == "pose": + if self.cfg.use_relative_mode: + if ee_pos is None or ee_quat is None: + raise ValueError("ee_pos and ee_quat are required for relative pose mode") + target_pos, target_quat = apply_delta_pose(ee_pos, ee_quat, command) + else: + target_pos = command[:, 0:3] + target_quat = command[:, 3:7] + self._write_pos_target(target_pos) + self._write_quat_target(target_quat) + else: # position-only + if self.cfg.use_relative_mode: + if ee_pos is None: + raise ValueError("ee_pos is required for position_rel mode") + target_pos = ee_pos + command + else: + target_pos = command + self._write_pos_target(target_pos) + + def _write_pos_target(self, target_pos: torch.Tensor) -> None: + """Copy a torch ``(N, 3)`` position tensor into the Warp target_positions buffer.""" + target_pos_wp = wp.from_torch(target_pos.contiguous(), dtype=wp.vec3) + wp.copy(self._target_pos_wp, target_pos_wp) + + def _write_quat_target(self, target_quat: torch.Tensor) -> None: + """Copy a torch ``(N, 4)`` quaternion tensor into the Warp target_rotations buffer.""" + target_quat_wp = wp.from_torch(target_quat.contiguous(), dtype=wp.vec4) + wp.copy(self._target_rot_wp, target_quat_wp) + + def compute(self, joint_q_in: torch.Tensor) -> torch.Tensor: + """Solve IK for all envs in one batched pass. + + Args: + joint_q_in: Current joint positions [m or rad] for the full IK + model, shape ``[num_envs, joint_coord_count]``. The caller + (typically the action term) is responsible for providing + complete joint coords — including non-arm joints such as + gripper fingers, which the solver leaves unchanged because + they have zero position/rotation Jacobian contribution. + + Returns: + Target joint positions [m or rad], shape + ``[num_envs, joint_coord_count]``. + """ + if joint_q_in.shape != (self.num_envs, self._n_coords): + raise ValueError( + f"joint_q_in shape {tuple(joint_q_in.shape)} does not match " + f"(num_envs={self.num_envs}, joint_coord_count={self._n_coords})" + ) + + if self.cfg.seed_source == "sim_joint_pos": + seed_wp = wp.from_torch(joint_q_in.contiguous(), dtype=wp.float32).reshape((self.num_envs, self._n_coords)) + wp.copy(self._joint_q_in_wp, seed_wp) + elif self.cfg.seed_source == "previous_solution": + assert self._previous_solution_wp is not None + wp.copy(self._joint_q_in_wp, self._previous_solution_wp) + elif self.cfg.seed_source == "default_pose": + default_q = self.model.joint_q.numpy() # full n_coords + tile = torch.tensor(default_q, device=joint_q_in.device).repeat(self.num_envs, 1) + wp.copy( + self._joint_q_in_wp, + wp.from_torch(tile.contiguous(), dtype=wp.float32).reshape((self.num_envs, self._n_coords)), + ) + else: + raise ValueError(f"unknown seed_source: {self.cfg.seed_source}") + + self._solver.step( + self._joint_q_in_wp, + self._joint_q_out_wp, + iterations=self.cfg.iterations, + ) + + if self._previous_solution_wp is not None: + wp.copy(self._previous_solution_wp, self._joint_q_out_wp) + + return wp.to_torch(self._joint_q_out_wp) + + def reset(self, env_ids: torch.Tensor | None = None) -> None: + """Reset internal solver/seed state. + + Args: + env_ids: Currently unused (the underlying Newton IK solver has no + per-env state worth preserving across resets). Present for + parity with the IsaacLab :class:`ActionTerm` reset signature. + """ + del env_ids # unused for now + self._solver.reset() + if self._previous_solution_wp is not None: + self._previous_solution_wp.zero_() diff --git a/source/isaaclab_newton/isaaclab_newton/controllers/newton_ik_cfg.py b/source/isaaclab_newton/isaaclab_newton/controllers/newton_ik_cfg.py new file mode 100644 index 000000000000..74520fea0636 --- /dev/null +++ b/source/isaaclab_newton/isaaclab_newton/controllers/newton_ik_cfg.py @@ -0,0 +1,111 @@ +# Copyright (c) 2022-2026, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +"""Configuration for :class:`~isaaclab_newton.controllers.NewtonIKController`.""" + +from __future__ import annotations + +from typing import Literal + +from isaaclab.utils import configclass + + +@configclass +class NewtonIKControllerCfg: + """Configuration for the Newton optimization-based IK controller. + + The defaults match the Standard tracking setup: relative-pose 6-D action, + Levenberg-Marquardt optimizer with analytic Jacobians, 8 iterations per + step, no multi-seed sampling, and warm-starting from the current sim + joint positions. + """ + + command_type: Literal["position", "pose"] = "pose" + """Whether the action command is position-only or full pose.""" + + use_relative_mode: bool = True + """If ``True``, the command is interpreted as a delta applied to the + current end-effector pose.""" + + # --- Solver --- + optimizer: Literal["lm", "lbfgs"] = "lm" + """Newton IK optimizer backend.""" + + jacobian_mode: Literal["analytic", "autodiff", "mixed"] = "analytic" + """Jacobian computation backend.""" + + iterations: int = 8 + """Number of optimizer iterations per :meth:`compute` call.""" + + # --- Sampling (only used when n_seeds > 1) --- + sampler: Literal["none", "gauss", "roberts", "uniform"] = "none" + n_seeds: int = 1 + noise_std: float = 0.1 + rng_seed: int = 12345 + + # --- Objective weights --- + # NOTE on unit scales: position residuals are in meters, rotation residuals + # are in radians. With equal weights, LM over-prioritizes rotation because + # a 1-rad error costs ~10,000× more than a 1-cm error in ``sum(r²)``. The + # default ``position_weight=100`` rescales meters to cm-scale so that a + # ~1 mm position error and a ~0.01 rad (~0.57°) rotation error contribute + # equally to the cost. Lower ``position_weight`` loosens position tracking; + # raise it for sub-mm precision. + position_weight: float = 100.0 + rotation_weight: float = 1.0 + joint_limit_weight: float = 0.1 + + # --- LM-specific (ignored when ``optimizer == "lbfgs"``) --- + lambda_initial: float = 0.1 + lambda_factor: float = 2.0 + lambda_min: float = 1e-5 + lambda_max: float = 1e10 + rho_min: float = 1e-3 + + # --- L-BFGS-specific (ignored when ``optimizer == "lm"``) --- + history_len: int = 10 + h0_scale: float = 1.0 + + # --- Warm-start --- + seed_source: Literal["sim_joint_pos", "previous_solution", "default_pose"] = "sim_joint_pos" + """Source of the seed joint positions for each :meth:`compute` call.""" + + def action_dim(self) -> int: + """Return the action dimension implied by this configuration. + + Returns: + ``3`` for position-only, ``6`` for pose-relative, or ``7`` for + absolute pose. + """ + if self.command_type == "position": + return 3 + if self.use_relative_mode: + return 6 + return 7 + + def validate_config(self) -> None: + """Validate enum-like fields beyond what :func:`Literal` enforces. + + This method is automatically called by :func:`~isaaclab.utils.configclass` + after the standard missing-field check. + + Raises: + ValueError: If any enum-like field has an unsupported value or a + numeric field is out of range. + """ + if self.optimizer not in ("lm", "lbfgs"): + raise ValueError(f"optimizer must be 'lm' or 'lbfgs', got {self.optimizer!r}") + if self.jacobian_mode not in ("analytic", "autodiff", "mixed"): + raise ValueError(f"jacobian_mode invalid: {self.jacobian_mode!r}") + if self.sampler not in ("none", "gauss", "roberts", "uniform"): + raise ValueError(f"sampler invalid: {self.sampler!r}") + if self.seed_source not in ("sim_joint_pos", "previous_solution", "default_pose"): + raise ValueError(f"seed_source invalid: {self.seed_source!r}") + if self.iterations < 1: + raise ValueError(f"iterations must be >= 1, got {self.iterations}") + if self.n_seeds < 1: + raise ValueError(f"n_seeds must be >= 1, got {self.n_seeds}") + if self.sampler == "none" and self.n_seeds != 1: + raise ValueError("sampler='none' requires n_seeds == 1") diff --git a/source/isaaclab_newton/isaaclab_newton/controllers/newton_ik_model_builder.py b/source/isaaclab_newton/isaaclab_newton/controllers/newton_ik_model_builder.py new file mode 100644 index 000000000000..977a5f4d09d3 --- /dev/null +++ b/source/isaaclab_newton/isaaclab_newton/controllers/newton_ik_model_builder.py @@ -0,0 +1,154 @@ +# Copyright (c) 2022-2026, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +"""Builder for a single-arm Newton :class:`~newton.Model` used by Newton IK. + +The Newton IK solver expects each of its ``n_problems`` to be a replication of +a single articulation topology. The IsaacLab simulation Newton model instead +holds N articulations concatenated, so we build a small dedicated IK model +from the same USD asset that the sim uses. Topology must match the per-env +arm slice of the sim model; we validate this at construction time. +""" + +from __future__ import annotations + +import re +from dataclasses import dataclass + +import newton +import warp as wp +from newton import ModelBuilder + +from pxr import Usd + +from isaaclab.assets import ArticulationCfg +from isaaclab.utils.assets import retrieve_file_path + + +@dataclass +class IKModelInfo: + """Metadata for a single-arm IK Newton model bound to an IsaacLab asset. + + Attributes: + ee_link_index: Index into ``ik_model.body_q`` for the end-effector body. + arm_dof_count: Number of arm joint DOFs controlled by IK. + ik_joint_names: Names of the IK-controlled joints, in IK-model order. + sim_to_ik_joint_perm: Permutation mapping IsaacLab sim arm joint + indices to IK-model joint indices, length ``arm_dof_count``. + """ + + ee_link_index: int + arm_dof_count: int + ik_joint_names: list[str] + sim_to_ik_joint_perm: list[int] + + +def _label_short_name(label: str) -> str: + """Return the final path component of a USD-path label. + + Newton body/joint labels may be full USD paths such as + ``'/panda/panda_link0/panda_joint1'``. This helper strips the prefix so + that caller-facing names stay short (e.g. ``'panda_joint1'``). + """ + return label.rsplit("/", 1)[-1] + + +def _resolve_joint_names(all_labels: list[str], patterns: list[str]) -> list[str]: + """Resolve IsaacLab-style regex patterns against a list of joint labels. + + Matching is performed against the short name (last path component) of each + label so callers can write plain names like ``panda_joint.*`` even when the + Newton model stores full USD paths. + """ + resolved: list[str] = [] + for label in all_labels: + short = _label_short_name(label) + if any(re.fullmatch(pat, short) for pat in patterns): + resolved.append(label) + return resolved + + +def build_single_arm_ik_model( + asset_cfg: ArticulationCfg, + body_name: str, + joint_names: list[str], + device: str, +) -> tuple[newton.Model, IKModelInfo]: + """Build a Newton :class:`~newton.Model` containing one copy of the robot. + + Loads the same USD asset used by the sim into a fresh + :class:`~newton.ModelBuilder`, finalizes a single-articulation + :class:`~newton.Model` with the base at the world origin, and returns + metadata needed by :class:`~isaaclab_newton.controllers.NewtonIKController`. + + Newton body and joint labels may be full USD paths (e.g. + ``'/panda/panda_hand'``). The ``body_name`` and ``joint_names`` patterns + are matched against the short name (last path component), so callers should + use simple names such as ``'panda_hand'`` and ``'panda_joint.*'``. + + For instanceable USDs the Newton builder may produce a small extra + articulation for the world-root fixed joint. The function automatically + selects the largest articulation (by joint count) as the arm. + + Args: + asset_cfg: IsaacLab articulation configuration whose ``spawn.usd_path`` + is loaded. + body_name: Short name of the end-effector body in the asset (last USD + path component, e.g. ``'panda_hand'``). + joint_names: Joint name regex patterns controlled by IK, matched + against short names (e.g. ``['panda_joint.*']``). + device: Warp device string (e.g. ``"cuda:0"`` or ``"cpu"``). + + Returns: + Tuple of the finalized single-articulation :class:`~newton.Model` and + the :class:`IKModelInfo` describing the EE body and joint mapping. + + Raises: + ValueError: If the EE body or any IK-controlled joints cannot be + resolved from the loaded model. + ValueError: If the loaded USD does not produce any articulations. + """ + usd_path = retrieve_file_path(asset_cfg.spawn.usd_path) + builder = ModelBuilder() + stage = Usd.Stage.Open(usd_path) + # Force fixed base at world origin so world-frame targets equal base-frame targets. + builder.add_usd( + stage, + floating=False, + xform=wp.transform_identity(), + override_root_xform=True, + ) + + model = builder.finalize(device=device, requires_grad=True) + + if model.articulation_count < 1: + raise ValueError( + f"Expected USD '{usd_path}' to produce at least 1 articulation, got {model.articulation_count}." + ) + + # Resolve body names (matching by short name / last USD path component). + body_labels = list(model.body_label) + body_short_names = [_label_short_name(lbl) for lbl in body_labels] + if body_name not in body_short_names: + raise ValueError(f"Body name '{body_name}' not found in IK model. Available (short names): {body_short_names}") + ee_link_index = body_short_names.index(body_name) + + # Resolve joint names (matching by short name / last USD path component). + all_joint_labels = list(model.joint_label) + ik_joint_labels = _resolve_joint_names(all_joint_labels, joint_names) + if not ik_joint_labels: + all_short = [_label_short_name(lbl) for lbl in all_joint_labels] + raise ValueError(f"No joints in IK model match {joint_names}. Available (short names): {all_short}") + ik_joint_names = [_label_short_name(lbl) for lbl in ik_joint_labels] + arm_dof_count = len(ik_joint_names) + sim_to_ik_joint_perm = [all_joint_labels.index(lbl) for lbl in ik_joint_labels] + + info = IKModelInfo( + ee_link_index=ee_link_index, + arm_dof_count=arm_dof_count, + ik_joint_names=ik_joint_names, + sim_to_ik_joint_perm=sim_to_ik_joint_perm, + ) + return model, info 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..e46a0d432fee --- /dev/null +++ b/source/isaaclab_newton/isaaclab_newton/envs/__init__.py @@ -0,0 +1,6 @@ +# 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-native env helpers for IsaacLab.""" 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..64e22f104075 --- /dev/null +++ b/source/isaaclab_newton/isaaclab_newton/envs/mdp/__init__.py @@ -0,0 +1,6 @@ +# 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-native MDP terms.""" 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..d728b87484a1 --- /dev/null +++ b/source/isaaclab_newton/isaaclab_newton/envs/mdp/actions/__init__.py @@ -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 + +"""Newton-native MDP action terms for IsaacLab manager-based envs.""" + +from .newton_ik_actions import NewtonInverseKinematicsAction +from .newton_ik_actions_cfg import NewtonInverseKinematicsActionCfg + +__all__ = ["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 new file mode 100644 index 000000000000..ab9d02d3daa1 --- /dev/null +++ b/source/isaaclab_newton/isaaclab_newton/envs/mdp/actions/newton_ik_actions.py @@ -0,0 +1,173 @@ +# 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 action term for IsaacLab manager-based envs.""" + +from __future__ import annotations + +import logging +from collections.abc import Sequence +from typing import TYPE_CHECKING + +import torch +import warp as wp + +import isaaclab.utils.math as math_utils +from isaaclab.managers.action_manager import ActionTerm + +from isaaclab_newton.assets.articulation import Articulation +from isaaclab_newton.controllers import NewtonIKController, build_single_arm_ik_model + +if TYPE_CHECKING: + from isaaclab.envs import ManagerBasedEnv + + from .newton_ik_actions_cfg import NewtonInverseKinematicsActionCfg + +logger = logging.getLogger(__name__) + + +class NewtonInverseKinematicsAction(ActionTerm): + """Newton optimization-based IK action term for manager-based envs. + + Owns a :class:`~isaaclab_newton.controllers.NewtonIKController` and a + dedicated single-arm Newton :class:`~newton.Model` built from the asset's + USD. Each :meth:`process_actions` call computes EE pose targets in the + robot-base frame and pushes them into the controller's IK objectives; each + :meth:`apply_actions` call solves IK across all envs in one batched pass + and writes the resulting joint targets via + :meth:`Articulation.set_joint_position_target_index`. + """ + + cfg: NewtonInverseKinematicsActionCfg + _asset: Articulation + + def __init__(self, cfg: NewtonInverseKinematicsActionCfg, env: ManagerBasedEnv) -> None: + super().__init__(cfg, env) + + # Resolve joint and body indices on the sim asset. + self._joint_ids, self._joint_names = self._asset.find_joints(self.cfg.joint_names) + body_ids, body_names = self._asset.find_bodies(self.cfg.body_name) + if len(body_ids) != 1: + raise ValueError(f"Expected exactly 1 body matching '{self.cfg.body_name}', got {body_names}") + self._body_idx = body_ids[0] + + # Build a single-arm IK Newton model (and validate topology match). + ik_model, ik_info = build_single_arm_ik_model( + asset_cfg=self._asset.cfg, + body_name=self.cfg.body_name, + joint_names=self.cfg.joint_names, + device=self.device, + ) + if ik_info.arm_dof_count != len(self._joint_ids): + raise ValueError( + f"IK model arm DOF count ({ik_info.arm_dof_count}) does not match " + f"sim asset joint count ({len(self._joint_ids)})." + ) + sim_joint_count = self._asset.data.joint_pos.shape[1] + if sim_joint_count != ik_model.joint_coord_count: + raise ValueError( + f"Sim asset joint count ({sim_joint_count}) does not match IK " + f"model joint_coord_count ({ik_model.joint_coord_count}). The " + f"two must be built from the same USD asset with matching " + f"topology." + ) + self._ik_joint_coord_count = ik_model.joint_coord_count + self._ik_info = ik_info + + # EE offset from cfg. + if self.cfg.body_offset is not None: + ee_link_offset = tuple(self.cfg.body_offset.pos) + ee_link_offset_rot = tuple(self.cfg.body_offset.rot) + else: + ee_link_offset = (0.0, 0.0, 0.0) + ee_link_offset_rot = (0.0, 0.0, 0.0, 1.0) + + self._controller = NewtonIKController( + cfg=self.cfg.controller, + ik_model=ik_model, + num_envs=self.num_envs, + ee_link_index=ik_info.ee_link_index, + arm_dof_count=ik_info.arm_dof_count, + ee_link_offset=ee_link_offset, + ee_link_offset_rot=ee_link_offset_rot, + device=self.device, + ) + + # Action tensors. + self._raw_actions = torch.zeros(self.num_envs, self.action_dim, device=self.device) + self._processed_actions = torch.zeros_like(self._raw_actions) + + # Scale: broadcast to (1, action_dim). + if isinstance(self.cfg.scale, (int, float)): + self._scale = torch.full((1, self.action_dim), float(self.cfg.scale), device=self.device) + else: + self._scale = torch.tensor(self.cfg.scale, device=self.device).reshape(1, -1) + if self._scale.shape[1] != self.action_dim: + raise ValueError( + f"scale tuple length ({self._scale.shape[1]}) must match action_dim ({self.action_dim})" + ) + + # Body offset tensors for EE pose composition (if set). + 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) + else: + self._offset_pos = None + self._offset_rot = None + + # ---- properties ---- + + @property + def action_dim(self) -> int: + return self._controller.action_dim + + @property + def raw_actions(self) -> torch.Tensor: + return self._raw_actions + + @property + def processed_actions(self) -> torch.Tensor: + return self._processed_actions + + # ---- operations ---- + + def process_actions(self, actions: torch.Tensor) -> None: + """Scale/clip actions and push EE pose targets into the IK objectives.""" + self._raw_actions[:] = actions + self._processed_actions[:] = self._raw_actions * self._scale + ee_pos_b, ee_quat_b = self._compute_ee_pose_in_base_frame() + self._controller.set_command(self._processed_actions, ee_pos=ee_pos_b, ee_quat=ee_quat_b) + + def apply_actions(self) -> None: + """Solve IK across all envs in one pass and write joint position targets.""" + # Read the full per-env joint_pos (all sim joints); shape (num_envs, joint_coord_count). + joint_pos_full = wp.to_torch(self._asset.data.joint_pos) + # Batched IK solve returns the same shape. + joint_targets_full = self._controller.compute(joint_pos_full) + # Extract only the arm-controlled subset and push to the articulation. + arm_targets = joint_targets_full[:, self._joint_ids] + self._asset.set_joint_position_target_index(target=arm_targets, joint_ids=self._joint_ids) + + def reset(self, env_ids: Sequence[int] | None = None) -> None: + if env_ids is None: + self._raw_actions[:] = 0.0 + else: + self._raw_actions[env_ids] = 0.0 + self._controller.reset(env_ids=None) + + # ---- helpers ---- + + def _compute_ee_pose_in_base_frame(self) -> tuple[torch.Tensor, torch.Tensor]: + """Compute EE pose (with optional offset) in the robot-base frame.""" + ee_pos_w = wp.to_torch(self._asset.data.body_pos_w)[:, self._body_idx] + ee_quat_w = wp.to_torch(self._asset.data.body_quat_w)[:, self._body_idx] + root_pos_w = wp.to_torch(self._asset.data.root_pos_w) + root_quat_w = wp.to_torch(self._asset.data.root_quat_w) + 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._offset_pos 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 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..914553f6368c --- /dev/null +++ b/source/isaaclab_newton/isaaclab_newton/envs/mdp/actions/newton_ik_actions_cfg.py @@ -0,0 +1,50 @@ +# 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 + +"""Configuration for the Newton inverse-kinematics action term.""" + +from __future__ import annotations + +from dataclasses import MISSING + +from isaaclab.managers.action_manager import ActionTermCfg +from isaaclab.utils import configclass + +from isaaclab_newton.controllers import NewtonIKControllerCfg + + +@configclass +class NewtonInverseKinematicsActionCfg(ActionTermCfg): + """Configuration for the Newton IK action term.""" + + @configclass + class OffsetCfg: + """Pose offset from the parent body frame to the constrained EE frame.""" + + pos: tuple[float, float, float] = (0.0, 0.0, 0.0) + """Translation w.r.t. the parent body [m].""" + + rot: tuple[float, float, float, float] = (0.0, 0.0, 0.0, 1.0) + """Quaternion ``(x, y, z, w)`` w.r.t. the parent body.""" + + class_type: type | str = "{DIR}.newton_ik_actions:NewtonInverseKinematicsAction" + + joint_names: list[str] = MISSING + """Joint name regex patterns controlled by IK.""" + + body_name: str = MISSING + """Name of the EE body for which IK is performed.""" + + body_offset: OffsetCfg | None = None + """Optional pose offset from the EE body frame to the constrained frame.""" + + scale: float | tuple[float, ...] = 1.0 + """Scale factor applied to the input action.""" + + clip: dict[str, tuple[float, float]] | None = None + """Optional per-joint clip dict.""" + + controller: NewtonIKControllerCfg = MISSING + """Configuration for the underlying :class:`NewtonIKController`.""" diff --git a/source/isaaclab_newton/isaaclab_newton/physics/__init__.pyi b/source/isaaclab_newton/isaaclab_newton/physics/__init__.pyi index 4e589e6f69d1..da6ae54f7c47 100644 --- a/source/isaaclab_newton/isaaclab_newton/physics/__init__.pyi +++ b/source/isaaclab_newton/isaaclab_newton/physics/__init__.pyi @@ -5,12 +5,24 @@ __all__ = [ "FeatherstoneSolverCfg", + "HydroelasticCfg", "MJWarpSolverCfg", "NewtonCfg", + "NewtonCollisionPipelineCfg", "NewtonManager", "NewtonSolverCfg", + "SDFCfg", "XPBDSolverCfg", ] from .newton_manager import NewtonManager -from .newton_manager_cfg import FeatherstoneSolverCfg, MJWarpSolverCfg, NewtonCfg, NewtonSolverCfg, XPBDSolverCfg +from .newton_manager_cfg import ( + FeatherstoneSolverCfg, + HydroelasticCfg, + MJWarpSolverCfg, + NewtonCfg, + NewtonCollisionPipelineCfg, + NewtonSolverCfg, + SDFCfg, + XPBDSolverCfg, +) diff --git a/source/isaaclab_newton/isaaclab_newton/physics/newton_manager.py b/source/isaaclab_newton/isaaclab_newton/physics/newton_manager.py index a96749d7b3f3..6b12923b49c7 100644 --- a/source/isaaclab_newton/isaaclab_newton/physics/newton_manager.py +++ b/source/isaaclab_newton/isaaclab_newton/physics/newton_manager.py @@ -11,6 +11,7 @@ import ctypes import inspect import logging +import re from typing import TYPE_CHECKING import numpy as np @@ -28,6 +29,7 @@ _cudart = None from newton import Axis, CollisionPipeline, Contacts, Control, Model, ModelBuilder, State, eval_fk from newton._src.usd.schemas import SchemaResolverNewton, SchemaResolverPhysx +from newton.geometry import HydroelasticSDF from newton.sensors import SensorContact as NewtonContactSensor from newton.solvers import SolverBase, SolverFeatherstone, SolverMuJoCo, SolverNotifyFlags, SolverXPBD @@ -377,6 +379,242 @@ def add_model_change(cls, change: SolverNotifyFlags) -> None: """Register a model change to notify the solver.""" cls._model_changes.add(change) + @staticmethod + def _build_sdf_on_mesh(mesh, sdf_cfg, res_overrides, label: str): + """Build SDF on a mesh, resolving per-pattern resolution overrides. + + Args: + mesh: Newton mesh object to build SDF on. + sdf_cfg: The active :class:`SDFCfg` instance. + res_overrides: Compiled (pattern, resolution) pairs, or None. + label: Shape label used for pattern resolution matching. + """ + if mesh is None: + return + if mesh.sdf is not None: + return # SDF already built (shared by reference from prototype) + resolution = sdf_cfg.max_resolution + if res_overrides is not None: + for pat, res in res_overrides: + if pat.search(label): + resolution = res + break + sdf_kwargs = dict(narrow_band_range=sdf_cfg.narrow_band_range) + if resolution is not None: + sdf_kwargs["max_resolution"] = resolution + if sdf_cfg.target_voxel_size is not None: + sdf_kwargs["target_voxel_size"] = sdf_cfg.target_voxel_size + mesh.build_sdf(**sdf_kwargs) + + @classmethod + def _create_sdf_collision_from_visual( + cls, builder: ModelBuilder, sdf_shape_indices: set[int], sdf_cfg, res_overrides + ): + """Create collision shapes from visual meshes for matched bodies that lack collision geometry. + + Args: + builder: Newton model builder to modify. + sdf_shape_indices: Shape indices that matched SDF patterns. + sdf_cfg: The active :class:`SDFCfg` instance. + res_overrides: Compiled (pattern, resolution) pairs, or None. + + Returns: + Tuple of (num_added, num_hydro) counts. + """ + from newton import ShapeFlags + + hydro_cfg = sdf_cfg.hydroelastic_cfg + + matched_bodies: set[int] = {builder.shape_body[si] for si in sdf_shape_indices} + bodies_with_collision: set[int] = set() + for si in range(builder.shape_count): + if builder.shape_flags[si] & ShapeFlags.COLLIDE_SHAPES and builder.shape_body[si] in matched_bodies: + bodies_with_collision.add(builder.shape_body[si]) + + hydro_patterns = None + if hydro_cfg is not None and hydro_cfg.shape_patterns is not None: + hydro_patterns = [re.compile(p) for p in hydro_cfg.shape_patterns] + + num_added = 0 + num_hydro = 0 + for body_idx in matched_bodies - bodies_with_collision: + visual_si = None + for si in sdf_shape_indices: + if builder.shape_body[si] == body_idx and builder.shape_source[si] is not None: + visual_si = si + break + if visual_si is None: + body_lbl = builder.body_label[body_idx] + logger.warning(f"SDF: body '{body_lbl}' matched but has no visual mesh to create collision from.") + continue + + mesh = builder.shape_source[visual_si] + cls._build_sdf_on_mesh(mesh, sdf_cfg, res_overrides, builder.shape_label[visual_si]) + + body_lbl = builder.body_label[body_idx] + shape_label = f"{body_lbl}/sdf_collision" + + # Check if this shape should get hydroelastic (respecting hydro shape_patterns) + apply_hydro = hydro_cfg is not None and ( + hydro_patterns is None or any(p.search(shape_label) for p in hydro_patterns) + ) + + shape_cfg_kwargs = dict( + density=0.0, + has_shape_collision=True, + has_particle_collision=True, + is_visible=False, + ) + if sdf_cfg.margin is not None: + shape_cfg_kwargs["margin"] = sdf_cfg.margin + if apply_hydro: + shape_cfg_kwargs["is_hydroelastic"] = True + shape_cfg_kwargs["kh"] = hydro_cfg.k_hydro + + builder.add_shape_mesh( + body=body_idx, + xform=builder.shape_transform[visual_si], + mesh=mesh, + scale=builder.shape_scale[visual_si], + cfg=ModelBuilder.ShapeConfig(**shape_cfg_kwargs), + label=shape_label, + ) + num_added += 1 + if apply_hydro: + num_hydro += 1 + + return num_added, num_hydro + + @classmethod + def _apply_sdf_config(cls, builder: ModelBuilder): + """Apply SDF collision and optional hydroelastic flags to matching mesh shapes.""" + from newton import GeoType, ShapeFlags + + cfg = PhysicsManager._cfg + if cfg is None: + return + sdf_cfg = cfg.sdf_cfg + if sdf_cfg is None: + return + + if sdf_cfg.max_resolution is None and sdf_cfg.target_voxel_size is None: + logger.warning("SDFCfg provided but neither max_resolution nor target_voxel_size is set. SDF disabled.") + return + + hydro_cfg = sdf_cfg.hydroelastic_cfg + + # Compile patterns + body_patterns = [re.compile(p) for p in sdf_cfg.body_patterns] if sdf_cfg.body_patterns else None + shape_patterns = [re.compile(p) for p in sdf_cfg.shape_patterns] if sdf_cfg.shape_patterns else None + res_overrides = ( + [(re.compile(p), r) for p, r in sdf_cfg.pattern_resolutions.items()] + if sdf_cfg.pattern_resolutions + else None + ) + hydro_patterns = None + if hydro_cfg is not None and hydro_cfg.shape_patterns is not None: + hydro_patterns = [re.compile(p) for p in hydro_cfg.shape_patterns] + + # --- Collect shape indices that should get SDF --- + # Build reverse map once: body_idx -> [mesh shape indices] + body_to_shapes: dict[int, list[int]] = {} + for si in range(builder.shape_count): + if builder.shape_type[si] == GeoType.MESH: + body_to_shapes.setdefault(builder.shape_body[si], []).append(si) + + sdf_shape_indices: set[int] = set() + + if body_patterns is not None: + for body_idx in range(len(builder.body_label)): + if any(p.search(builder.body_label[body_idx]) for p in body_patterns): + sdf_shape_indices.update(body_to_shapes.get(body_idx, [])) + + if shape_patterns is not None: + for shape_indices in body_to_shapes.values(): + for si in shape_indices: + if any(p.search(builder.shape_label[si]) for p in shape_patterns): + sdf_shape_indices.add(si) + + if body_patterns is None and shape_patterns is None: + logger.warning("SDFCfg has no body_patterns or shape_patterns set. No shapes will receive SDF.") + return + + # --- Patch existing collision meshes --- + num_patched = 0 + num_hydro = 0 + for si in sdf_shape_indices: + if not (builder.shape_flags[si] & ShapeFlags.COLLIDE_SHAPES): + continue + cls._build_sdf_on_mesh(builder.shape_source[si], sdf_cfg, res_overrides, builder.shape_label[si]) + if sdf_cfg.margin is not None: + builder.shape_margin[si] = sdf_cfg.margin + if hydro_cfg is not None: + apply_hydro = hydro_patterns is None or any(p.search(builder.shape_label[si]) for p in hydro_patterns) + if apply_hydro: + builder.shape_flags[si] |= ShapeFlags.HYDROELASTIC + builder.shape_material_kh[si] = hydro_cfg.k_hydro + num_hydro += 1 + num_patched += 1 + + # --- Optionally create collision shapes from visual meshes --- + num_added = 0 + if sdf_cfg.use_visual_meshes: + num_added, hydro_from_visual = cls._create_sdf_collision_from_visual( + builder, sdf_shape_indices, sdf_cfg, res_overrides + ) + num_hydro += hydro_from_visual + + hydro_msg = f", {num_hydro} hydroelastic shape(s)" if hydro_cfg is not None else "" + logger.info( + f"SDF config: {num_added} collision shape(s) added, {num_patched} existing shape(s) patched{hydro_msg}. " + f"(max_resolution={sdf_cfg.max_resolution}, narrow_band={sdf_cfg.narrow_band_range})" + ) + + @classmethod + def _create_collision_pipeline(cls, model: Model): + """Create a collision pipeline from :attr:`collision_cfg` and optional hydroelastic config. + + Pipeline parameters come from :class:`NewtonCollisionPipelineCfg` on ``NewtonCfg.collision_cfg``. + Hydroelastic config comes from :class:`HydroelasticCfg` on ``SDFCfg.hydroelastic_cfg``. + Falls back to ``broad_phase="explicit"`` when no ``collision_cfg`` is provided. + """ + cfg = PhysicsManager._cfg + + # Pipeline parameters from collision_cfg (or defaults) + collision_cfg = cfg.collision_cfg + if collision_cfg is not None: + if not hasattr(collision_cfg, "to_dict"): + raise TypeError( + f"collision_cfg {type(collision_cfg)} does not implement to_dict(). " + "Only @configclass instances are supported." + ) + pipeline_kwargs = {k: v for k, v in collision_cfg.to_dict().items() if v is not None} + else: + pipeline_kwargs = {"broad_phase": "explicit"} + + # Hydroelastic config from sdf_cfg + sdf_cfg = cfg.sdf_cfg + hydro_cfg = sdf_cfg.hydroelastic_cfg if sdf_cfg is not None else None + + if hydro_cfg is not None: + pipeline_kwargs["sdf_hydroelastic_config"] = HydroelasticSDF.Config( + reduce_contacts=hydro_cfg.reduce_contacts, + output_contact_surface=hydro_cfg.output_contact_surface, + normal_matching=hydro_cfg.normal_matching, + moment_matching=hydro_cfg.moment_matching, + margin_contact_area=hydro_cfg.margin_contact_area, + buffer_mult_broad=hydro_cfg.buffer_mult_broad, + buffer_mult_iso=hydro_cfg.buffer_mult_iso, + buffer_mult_contact=hydro_cfg.buffer_mult_contact, + grid_size=hydro_cfg.grid_size, + ) + logger.info( + f"Hydroelastic contacts enabled (k_hydro={hydro_cfg.k_hydro}, " + f"reduce_contacts={hydro_cfg.reduce_contacts})" + ) + + return CollisionPipeline(model, **pipeline_kwargs) + @classmethod def start_simulation(cls) -> None: """Start simulation by finalizing model and initializing state. @@ -387,7 +625,9 @@ def start_simulation(cls) -> None: """ logger.debug(f"Builder: {cls._builder}") - # Create builder from USD stage if not provided + # Create builder from USD stage if not provided. + # When the builder was set externally (e.g. by newton_replicate), + # SDF was already applied on the prototype before replication. if cls._builder is None: cls.instantiate_builder_from_stage() @@ -473,14 +713,49 @@ def instantiate_builder_from_stage(cls): builder.add_usd(stage, ignore_paths=ignore_paths, schema_resolvers=schema_resolvers) # Build a prototype from the first env (all envs assumed identical) + # Use skip_mesh_approximation + load_visual_shapes to preserve + # triangle meshes needed for SDF (same as newton_replicate path). _, proto_path = env_paths[0] proto = ModelBuilder(up_axis=up_axis) proto.add_usd( stage, root_path=proto_path, + load_visual_shapes=True, + skip_mesh_approximation=True, schema_resolvers=schema_resolvers, ) + # Approximate non-SDF meshes with convex hulls (preserves SDF meshes) + cfg = PhysicsManager._cfg + sdf_cfg = getattr(cfg, "sdf_cfg", None) if cfg is not None else None + body_pats = [re.compile(x) for x in sdf_cfg.body_patterns] if sdf_cfg and sdf_cfg.body_patterns else None + shape_pats = [re.compile(x) for x in sdf_cfg.shape_patterns] if sdf_cfg and sdf_cfg.shape_patterns else None + has_sdf_patterns = body_pats is not None or shape_pats is not None + + if has_sdf_patterns: + from newton import GeoType + + sdf_bodies: set[int] = set() + if body_pats is not None: + for bi in range(len(proto.body_label)): + if any(pat.search(proto.body_label[bi]) for pat in body_pats): + sdf_bodies.add(bi) + approx_indices = [] + for i in range(len(proto.shape_type)): + if proto.shape_type[i] != GeoType.MESH: + continue + if proto.shape_body[i] in sdf_bodies: + continue + if shape_pats is not None: + lbl = proto.shape_label[i] if i < len(proto.shape_label) else "" + if any(pat.search(lbl) for pat in shape_pats): + continue + approx_indices.append(i) + if approx_indices: + proto.approximate_meshes("convex_hull", shape_indices=approx_indices, keep_visual_shapes=True) + else: + proto.approximate_meshes("convex_hull", keep_visual_shapes=True) + # Add each env as a separate Newton world xform_cache = UsdGeom.XformCache() for _, env_path in env_paths: @@ -500,6 +775,7 @@ def instantiate_builder_from_stage(cls): cls._num_envs = len(env_paths) + cls._apply_sdf_config(builder) cls.set_builder(builder) @classmethod @@ -512,7 +788,7 @@ def _initialize_contacts(cls) -> None: if cls._needs_collision_pipeline: # Newton collision pipeline: create pipeline and generate contacts if cls._collision_pipeline is None: - cls._collision_pipeline = CollisionPipeline(cls._model, broad_phase="explicit") + cls._collision_pipeline = cls._create_collision_pipeline(cls._model) if cls._contacts is None: cls._contacts = cls._collision_pipeline.contacts() @@ -588,6 +864,21 @@ def initialize_solver(cls) -> None: else: cls._needs_collision_pipeline = True + # Force Newton pipeline when collision_cfg or SDF is configured + if cfg.collision_cfg is not None and not cls._needs_collision_pipeline: + logger.warning("collision_cfg set — enabling Newton collision pipeline.") + cls._needs_collision_pipeline = True + + sdf_cfg = cfg.sdf_cfg + has_sdf = ( + sdf_cfg is not None + and (sdf_cfg.body_patterns is not None or sdf_cfg.shape_patterns is not None) + and (sdf_cfg.max_resolution is not None or sdf_cfg.target_voxel_size is not None) + ) + if has_sdf and not cls._needs_collision_pipeline: + logger.warning("SDF collision requires Newton collision pipeline. Overriding use_mujoco_contacts.") + cls._needs_collision_pipeline = True + # Initialize contacts and collision pipeline cls._initialize_contacts() diff --git a/source/isaaclab_newton/isaaclab_newton/physics/newton_manager_cfg.py b/source/isaaclab_newton/isaaclab_newton/physics/newton_manager_cfg.py index afbf7f54ba0c..7605f64b309c 100644 --- a/source/isaaclab_newton/isaaclab_newton/physics/newton_manager_cfg.py +++ b/source/isaaclab_newton/isaaclab_newton/physics/newton_manager_cfg.py @@ -16,6 +16,191 @@ from .newton_manager import NewtonManager +@configclass +class HydroelasticCfg: + """Configuration for hydroelastic contact handling. + + Hydroelastic contacts use SDF overlap between shape pairs to compute distributed + contact surfaces with area-weighted forces via marching cubes. This requires SDF + to be enabled on the parent :class:`SDFCfg`. + + Both shapes in a colliding pair must have the hydroelastic flag for hydroelastic + contacts to be generated. Shapes that only have SDF (but not hydroelastic) will + fall back to standard point contacts. + + .. note:: + Hydroelastic contacts require the unified collision pipeline, which is used + when ``use_mujoco_contacts=False`` on :class:`MJWarpSolverCfg`, or when using + a non-MuJoCo solver (XPBD, Featherstone). + """ + + k_hydro: float = 1e10 + """Hydroelastic stiffness coefficient [Pa] applied to shapes. + + Controls the compliance of the hydroelastic contact surface. Higher values produce + stiffer contacts. + """ + + shape_patterns: list[str] | None = None + """Regex patterns to select which shapes get hydroelastic contacts. + + If None, all shapes that have SDF enabled will also get hydroelastic contacts. + If provided, only shapes whose label (USD prim path) matches at least one pattern + will have the ``HYDROELASTIC`` flag set. + + Example: ``[".*Gear.*", ".*gear.*"]`` + """ + + reduce_contacts: bool = True + """Whether to reduce contacts to a smaller representative set per shape pair.""" + + output_contact_surface: bool = False + """Whether to output hydroelastic contact surface vertices for visualization.""" + + normal_matching: bool = True + """Whether to adjust reduced contact normals so their net force direction matches + the unreduced reference. Only active when :attr:`reduce_contacts` is True.""" + + moment_matching: bool = False + """Whether to adjust reduced contact friction coefficients so their net maximum + moment matches the unreduced reference. Only active when :attr:`reduce_contacts` + is True.""" + + margin_contact_area: float = 1e-2 + """Contact area [m^2] used for non-penetrating contacts at the margin.""" + + buffer_mult_broad: int = 1 + """Multiplier for the preallocated broadphase buffer. Increase if a broadphase + overflow warning is issued.""" + + buffer_mult_iso: int = 1 + """Multiplier for preallocated iso-surface extraction buffers. Increase if an + iso buffer overflow warning is issued.""" + + buffer_mult_contact: int = 1 + """Multiplier for the preallocated face contact buffer. Increase if a face + contact overflow warning is issued.""" + + grid_size: int = 256 * 8 * 128 + """Grid size for contact handling. Can be tuned for performance.""" + + +@configclass +class SDFCfg: + """Configuration for SDF (Signed Distance Field) collision on Newton meshes. + + When provided as ``sdf_cfg`` on :class:`NewtonCfg`, mesh collision shapes loaded + from USD will have SDF-based collision enabled via Newton's ``mesh.build_sdf()`` + API at simulation start. At least one of ``max_resolution`` or ``target_voxel_size`` + must be set. + + Regex patterns in ``shape_patterns`` and ``pattern_resolutions`` allow selective + SDF application and per-shape resolution tuning. + """ + + max_resolution: int | None = None + """Maximum dimension [voxels] for sparse SDF grid (must be divisible by 8). + + If set, mesh collision shapes loaded from USD will have SDF-based collision enabled. + Typical values: 128, 256, 512. + """ + + narrow_band_range: tuple[float, float] = (-0.1, 0.1) + """Narrow band distance range (inner, outer) [m] for SDF computation.""" + + target_voxel_size: float | None = None + """Target voxel size [m] for sparse SDF grid. + + If provided, takes precedence over :attr:`max_resolution`. + """ + + margin: float | None = None + """Collision margin [m] for SDF shapes. If None, uses the builder's default.""" + + body_patterns: list[str] | None = None + """List of regex patterns to match body labels (USD prim paths) for SDF. + + If None, no body-level filtering is applied. If provided, bodies whose label + matches at least one pattern will have SDF applied to all their mesh shapes. + Example: ``[".*elbow.*", ".*wrist.*"]`` + """ + + shape_patterns: list[str] | None = None + """List of regex patterns to match shape labels (USD prim paths) for SDF. + + If None, no shape-level filtering is applied. If provided, only shapes whose + label matches at least one pattern get SDF. + Example: ``[".*Gear.*", ".*gear.*"]`` + + .. note:: + At least one of :attr:`body_patterns` or :attr:`shape_patterns` must be set + for SDF to be applied. If both are None, no shapes will receive SDF. + """ + + use_visual_meshes: bool = False + """Whether to create collision shapes from visual meshes for bodies that have + no collision geometry. When False (default), only existing collision meshes + are patched with SDF. When True, bodies that match the configured patterns + but lack collision shapes will get a new collision shape created from their + first visual mesh. + """ + + pattern_resolutions: dict[str, int] | None = None + """Per-pattern SDF resolution overrides. + + Maps regex pattern to max_resolution for matching shapes. Shapes not matching any + pattern here use the global ``max_resolution``. First matching pattern wins. + Example: ``{".*elbow.*": 128, ".*power_supply.*": 512}`` + """ + + hydroelastic_cfg: HydroelasticCfg | None = None + """Hydroelastic contact configuration. + + If None (default), hydroelastic contacts are disabled and standard point contacts + are used. When set, shapes matching the SDF patterns (or the hydroelastic-specific + ``shape_patterns``) will have the ``HYDROELASTIC`` flag enabled and use distributed + surface contacts computed via SDF overlap. + """ + + +@configclass +class NewtonCollisionPipelineCfg: + """Configuration for Newton's collision pipeline. + + Full-featured collision pipeline with GJK/MPR narrow phase and pluggable broad phase. + When set on :attr:`NewtonCfg.collision_cfg`, Newton's collision pipeline replaces the + solver's native collision handling. + + For more details, see the `Newton Collisions Guide`_. + + .. _Newton Collisions Guide: https://newton-physics.github.io/newton/latest/concepts/collisions.html + """ + + broad_phase: str = "explicit" + """Broad phase algorithm for collision detection. + + Options: ``"explicit"`` (precomputed pairs), ``"nxn"`` (all-pairs), ``"sap"`` (sweep-and-prune). + """ + + reduce_contacts: bool = True + """Whether to reduce contacts for mesh-mesh collisions.""" + + rigid_contact_max: int | None = None + """Maximum number of rigid contacts to allocate. If None, auto-estimated from model.""" + + max_triangle_pairs: int = 1_000_000 + """Maximum triangle pairs allocated by narrow phase for mesh/heightfield collisions.""" + + soft_contact_max: int | None = None + """Maximum number of soft contacts. If None, computed as ``shape_count * particle_count``.""" + + soft_contact_margin: float = 0.01 + """Margin [m] for soft contact generation.""" + + requires_grad: bool | None = None + """Whether to enable gradient computation. If None, uses ``model.requires_grad``.""" + + @configclass class NewtonSolverCfg: """Configuration for Newton solver-related parameters. @@ -218,3 +403,14 @@ class NewtonCfg(PhysicsCfg): solver_cfg: NewtonSolverCfg = MJWarpSolverCfg() """Solver configuration. Default is MJWarpSolverCfg().""" + + collision_cfg: NewtonCollisionPipelineCfg | None = None + """Newton collision pipeline configuration. + + If None (default), the solver's native collision is used (e.g., MuJoCo's internal + contact solver). If set, Newton's :class:`CollisionPipeline` is used with the + specified parameters. + """ + + sdf_cfg: SDFCfg | None = None + """SDF collision configuration. If None (default), SDF is disabled.""" diff --git a/source/isaaclab_newton/test/controllers/__init__.py b/source/isaaclab_newton/test/controllers/__init__.py new file mode 100644 index 000000000000..460a30569089 --- /dev/null +++ b/source/isaaclab_newton/test/controllers/__init__.py @@ -0,0 +1,4 @@ +# 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 diff --git a/source/isaaclab_newton/test/controllers/test_imports.py b/source/isaaclab_newton/test/controllers/test_imports.py new file mode 100644 index 000000000000..9469ae2fb633 --- /dev/null +++ b/source/isaaclab_newton/test/controllers/test_imports.py @@ -0,0 +1,16 @@ +# 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 + +"""Smoke tests for the isaaclab_newton.controllers package surface.""" + +from isaaclab.app import AppLauncher + +simulation_app = AppLauncher(headless=True).app + +import isaaclab_newton.controllers as controllers + + +def test_package_imports(): + assert controllers is not None diff --git a/source/isaaclab_newton/test/controllers/test_newton_ik.py b/source/isaaclab_newton/test/controllers/test_newton_ik.py new file mode 100644 index 000000000000..910fbb54eb26 --- /dev/null +++ b/source/isaaclab_newton/test/controllers/test_newton_ik.py @@ -0,0 +1,295 @@ +# 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 + +"""Tests for NewtonIKController construction, properties, and solve behavior.""" + +from isaaclab.app import AppLauncher + +simulation_app = AppLauncher(headless=True).app + +import pytest +import warp as wp +from isaaclab_newton.controllers import ( + NewtonIKController, + NewtonIKControllerCfg, + build_single_arm_ik_model, +) + +from isaaclab_assets.robots.franka import FRANKA_PANDA_CFG + + +@pytest.fixture(scope="module") +def franka_ik_setup(): + device = "cuda:0" if wp.is_cuda_available() else "cpu" + model, info = build_single_arm_ik_model( + asset_cfg=FRANKA_PANDA_CFG, + body_name="panda_hand", + joint_names=["panda_joint.*"], + device=device, + ) + return model, info, device + + +def test_controller_constructs_and_exposes_action_dim(franka_ik_setup): + model, info, device = franka_ik_setup + cfg = NewtonIKControllerCfg(command_type="pose", use_relative_mode=True) + ctrl = NewtonIKController( + cfg=cfg, + ik_model=model, + num_envs=4, + ee_link_index=info.ee_link_index, + arm_dof_count=info.arm_dof_count, + device=device, + ) + assert ctrl.action_dim == 6 + assert ctrl.num_arm_dofs == info.arm_dof_count + assert ctrl.num_envs == 4 + + +import torch + + +def test_set_command_pose_relative_pushes_to_objective(franka_ik_setup): + """Pose-relative delta should land in pos target buffer after set_command.""" + model, info, device = franka_ik_setup + cfg = NewtonIKControllerCfg(command_type="pose", use_relative_mode=True) + ctrl = NewtonIKController( + cfg=cfg, + ik_model=model, + num_envs=2, + ee_link_index=info.ee_link_index, + arm_dof_count=info.arm_dof_count, + device=device, + ) + # Current EE pose for both envs (identity rotation, origin position). + torch_device = device if device == "cpu" else device + ee_pos = torch.zeros(2, 3, device=torch_device) + ee_quat = torch.tensor([[0.0, 0.0, 0.0, 1.0]] * 2, device=torch_device) + # Pose-relative action: +X 10 cm translation, no rotation. + delta = torch.tensor([[0.1, 0.0, 0.0, 0.0, 0.0, 0.0]] * 2, device=torch_device) + ctrl.set_command(delta, ee_pos=ee_pos, ee_quat=ee_quat) + target_pos_np = ctrl._target_pos_wp.numpy() + assert abs(target_pos_np[0][0] - 0.1) < 1e-6, f"expected 0.1, got {target_pos_np[0][0]}" + assert abs(target_pos_np[1][0] - 0.1) < 1e-6 + + +def test_set_command_pose_absolute_pushes_to_objective(franka_ik_setup): + """Absolute pose action should land both position and quaternion in target buffers.""" + model, info, device = franka_ik_setup + cfg = NewtonIKControllerCfg(command_type="pose", use_relative_mode=False) + ctrl = NewtonIKController( + cfg=cfg, + ik_model=model, + num_envs=2, + ee_link_index=info.ee_link_index, + arm_dof_count=info.arm_dof_count, + device=device, + ) + torch_device = device if device == "cpu" else device + target_pos = torch.tensor([[0.3, 0.2, 0.5], [0.4, 0.1, 0.6]], device=torch_device) + target_quat = torch.tensor([[0.0, 0.0, 0.0, 1.0]] * 2, device=torch_device) + action = torch.cat([target_pos, target_quat], dim=1) + ctrl.set_command(action) + assert abs(ctrl._target_pos_wp.numpy()[0][0] - 0.3) < 1e-6 + assert abs(ctrl._target_pos_wp.numpy()[1][2] - 0.6) < 1e-6 + + +def test_set_command_position_relative(franka_ik_setup): + model, info, device = franka_ik_setup + cfg = NewtonIKControllerCfg(command_type="position", use_relative_mode=True) + ctrl = NewtonIKController( + cfg=cfg, + ik_model=model, + num_envs=1, + ee_link_index=info.ee_link_index, + arm_dof_count=info.arm_dof_count, + device=device, + ) + torch_device = device if device == "cpu" else device + ee_pos = torch.tensor([[0.5, 0.0, 0.3]], device=torch_device) + delta = torch.tensor([[0.05, 0.0, 0.0]], device=torch_device) + ctrl.set_command(delta, ee_pos=ee_pos) + assert abs(ctrl._target_pos_wp.numpy()[0][0] - 0.55) < 1e-6 + + +def test_set_command_position_absolute(franka_ik_setup): + model, info, device = franka_ik_setup + cfg = NewtonIKControllerCfg(command_type="position", use_relative_mode=False) + ctrl = NewtonIKController( + cfg=cfg, + ik_model=model, + num_envs=1, + ee_link_index=info.ee_link_index, + arm_dof_count=info.arm_dof_count, + device=device, + ) + torch_device = device if device == "cpu" else device + target_pos = torch.tensor([[0.6, 0.1, 0.4]], device=torch_device) + ctrl.set_command(target_pos) + assert abs(ctrl._target_pos_wp.numpy()[0][0] - 0.6) < 1e-6 + + +import newton +import numpy as np + + +def test_compute_pose_accuracy_single_env(franka_ik_setup): + """Home-pose target → solve → measured EE pose error < 1mm. + + The joint_limit_weight is set to 0.0 so the only objectives are position + and rotation; with the exact home pose as seed and target the solver + converges to machine precision (< 1 µm in practice). + """ + model, info, device = franka_ik_setup + cfg = NewtonIKControllerCfg( + command_type="pose", + use_relative_mode=False, + iterations=24, + joint_limit_weight=0.0, + ) + ctrl = NewtonIKController( + cfg=cfg, + ik_model=model, + num_envs=1, + ee_link_index=info.ee_link_index, + arm_dof_count=info.arm_dof_count, + device=device, + ) + + # Resolve the home EE pose via Newton FK. + state = model.state() + newton.eval_fk(model, model.joint_q, model.joint_qd, state) + body_q_np = state.body_q.numpy() + ee_tf = body_q_np[info.ee_link_index] # (7,) as [tx,ty,tz, qx,qy,qz,qw] + ee_pos_home = torch.tensor([[float(ee_tf[0]), float(ee_tf[1]), float(ee_tf[2])]], device=device) + ee_quat_home = torch.tensor([[float(ee_tf[3]), float(ee_tf[4]), float(ee_tf[5]), float(ee_tf[6])]], device=device) + + # Seed the IK with the current (home) joints (full n_coords). + home_q = torch.tensor(model.joint_q.numpy(), device=device).reshape(1, -1) + + target_action = torch.cat([ee_pos_home, ee_quat_home], dim=1) + ctrl.set_command(target_action, ee_pos=ee_pos_home, ee_quat=ee_quat_home) + out = ctrl.compute(home_q) + assert out.shape == (1, model.joint_coord_count), f"unexpected out shape: {out.shape}" + + # Resolve the EE pose at the solved joint config using a one-shot FK. + q_wp = wp.from_torch(out.contiguous().reshape(-1), dtype=wp.float32) + state2 = model.state() + newton.eval_fk(model, q_wp, model.joint_qd, state2) + body_q2 = state2.body_q.numpy()[info.ee_link_index] + pos_err = float(np.linalg.norm(np.array(body_q2[:3]) - np.array(ee_tf[:3]))) + assert pos_err < 1e-3, f"position error {pos_err:.4e} > 1mm" + + +def test_compute_batched_solve_multi_env(franka_ik_setup): + """Single solver.step solves 16 distinct envs in one pass (no per-env loop).""" + model, info, device = franka_ik_setup + cfg = NewtonIKControllerCfg( + command_type="pose", + use_relative_mode=False, + iterations=16, + ) + n = 16 + ctrl = NewtonIKController( + cfg=cfg, + ik_model=model, + num_envs=n, + ee_link_index=info.ee_link_index, + arm_dof_count=info.arm_dof_count, + device=device, + ) + state = model.state() + newton.eval_fk(model, model.joint_q, model.joint_qd, state) + body_q = state.body_q.numpy()[info.ee_link_index] + ee_pos = torch.tensor(body_q[:3], device=device).repeat(n, 1).float() + ee_quat = torch.tensor( + [float(body_q[3]), float(body_q[4]), float(body_q[5]), float(body_q[6])], device=device + ).repeat(n, 1) + action = torch.cat([ee_pos, ee_quat], dim=1) + home_q = torch.tensor(model.joint_q.numpy(), device=device).repeat(n, 1) + ctrl.set_command(action, ee_pos=ee_pos, ee_quat=ee_quat) + out = ctrl.compute(home_q) + assert out.shape == (n, model.joint_coord_count) + assert not torch.isnan(out).any() + + +def test_compute_pose_accuracy_perturbed_targets(franka_ik_setup): + """Perturb the home pose, solve, expect pos error < 2mm in ≥95% of trials. + + Uses ``joint_limit_weight=0.01`` so the regularizer is weak enough that + position error dominates the cost near the home pose. Targets are small + (<=2.5 cm) perturbations which keep the joints well within limits. + """ + model, info, device = franka_ik_setup + cfg = NewtonIKControllerCfg( + command_type="pose", + use_relative_mode=False, + iterations=100, + joint_limit_weight=0.0, + ) + n_trials = 50 + ctrl = NewtonIKController( + cfg=cfg, + ik_model=model, + num_envs=n_trials, + ee_link_index=info.ee_link_index, + arm_dof_count=info.arm_dof_count, + device=device, + ) + state = model.state() + newton.eval_fk(model, model.joint_q, model.joint_qd, state) + body_q = state.body_q.numpy()[info.ee_link_index] + home_pos = torch.tensor(body_q[:3], device=device).float() + home_quat = torch.tensor( + [float(body_q[3]), float(body_q[4]), float(body_q[5]), float(body_q[6])], + device=device, + ) + torch.manual_seed(0) + deltas = (torch.rand(n_trials, 3, device=device) - 0.5) * 0.04 # ±2 cm + target_pos = home_pos.unsqueeze(0) + deltas + target_quat = home_quat.unsqueeze(0).repeat(n_trials, 1) + action = torch.cat([target_pos, target_quat], dim=1) + ee_pos_curr = home_pos.unsqueeze(0).repeat(n_trials, 1) + ee_quat_curr = home_quat.unsqueeze(0).repeat(n_trials, 1) + home_q = torch.tensor(model.joint_q.numpy(), device=device).repeat(n_trials, 1) + ctrl.set_command(action, ee_pos=ee_pos_curr, ee_quat=ee_quat_curr) + out = ctrl.compute(home_q) + + # Resolve the EE pose at each solved joint config and measure pos error. + pos_errs = [] + for i in range(n_trials): + q_wp = wp.from_torch(out[i : i + 1].contiguous().reshape(-1), dtype=wp.float32) + st = model.state() + newton.eval_fk(model, q_wp, model.joint_qd, st) + ee = st.body_q.numpy()[info.ee_link_index][:3] + pos_errs.append(float(np.linalg.norm(np.array(ee) - target_pos[i].cpu().numpy()))) + success = sum(1 for e in pos_errs if e < 2e-3) + assert success >= int(0.95 * n_trials), ( + f"only {success}/{n_trials} solves under 2mm; max err {max(pos_errs):.4e}, " + f"mean err {sum(pos_errs) / len(pos_errs):.4e}" + ) + + +def test_reset_clears_previous_solution(franka_ik_setup): + """reset() must zero the previous_solution buffer when seed_source='previous_solution'.""" + model, info, device = franka_ik_setup + cfg = NewtonIKControllerCfg(seed_source="previous_solution") + ctrl = NewtonIKController( + cfg=cfg, + ik_model=model, + num_envs=2, + ee_link_index=info.ee_link_index, + arm_dof_count=info.arm_dof_count, + device=device, + ) + # Seed with non-zero values, then reset and check buffer is zeroed. + n_coords = model.joint_coord_count + nonzero = torch.full((2, n_coords), 0.5, device=device) + wp.copy( + ctrl._previous_solution_wp, + wp.from_torch(nonzero.contiguous(), dtype=wp.float32).reshape((2, n_coords)), + ) + ctrl.reset() + assert (ctrl._previous_solution_wp.numpy() == 0).all() diff --git a/source/isaaclab_newton/test/controllers/test_newton_ik_cfg.py b/source/isaaclab_newton/test/controllers/test_newton_ik_cfg.py new file mode 100644 index 000000000000..91eddaea4122 --- /dev/null +++ b/source/isaaclab_newton/test/controllers/test_newton_ik_cfg.py @@ -0,0 +1,42 @@ +# Copyright (c) 2022-2026, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +"""Tests for NewtonIKControllerCfg.""" + +from isaaclab.app import AppLauncher + +simulation_app = AppLauncher(headless=True).app + +import pytest +from isaaclab_newton.controllers import NewtonIKControllerCfg + + +def test_default_cfg_is_pose_relative_lm_analytic_warm(): + cfg = NewtonIKControllerCfg() + assert cfg.command_type == "pose" + assert cfg.use_relative_mode is True + assert cfg.optimizer == "lm" + assert cfg.jacobian_mode == "analytic" + assert cfg.iterations == 8 + assert cfg.sampler == "none" + assert cfg.n_seeds == 1 + assert cfg.seed_source == "sim_joint_pos" + + +def test_action_dim_helper(): + # 3-D for position-only + cfg = NewtonIKControllerCfg(command_type="position", use_relative_mode=False) + assert cfg.action_dim() == 3 + # 6-D for pose-relative + cfg = NewtonIKControllerCfg(command_type="pose", use_relative_mode=True) + assert cfg.action_dim() == 6 + # 7-D for absolute pose + cfg = NewtonIKControllerCfg(command_type="pose", use_relative_mode=False) + assert cfg.action_dim() == 7 + + +def test_invalid_optimizer_raises(): + with pytest.raises(ValueError): + NewtonIKControllerCfg(optimizer="bogus").validate_config() diff --git a/source/isaaclab_newton/test/controllers/test_newton_ik_model_builder.py b/source/isaaclab_newton/test/controllers/test_newton_ik_model_builder.py new file mode 100644 index 000000000000..e2551562ac6f --- /dev/null +++ b/source/isaaclab_newton/test/controllers/test_newton_ik_model_builder.py @@ -0,0 +1,47 @@ +# 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 + +"""Tests for the single-arm IK Newton Model builder.""" + +from isaaclab.app import AppLauncher + +simulation_app = AppLauncher(headless=True).app + +import pytest +import warp as wp +from isaaclab_newton.controllers.newton_ik_model_builder import ( + IKModelInfo, + build_single_arm_ik_model, +) + +from isaaclab_assets.robots.franka import FRANKA_PANDA_CFG + + +def test_build_franka_ik_model_returns_single_articulation(): + model, info = build_single_arm_ik_model( + asset_cfg=FRANKA_PANDA_CFG, + body_name="panda_hand", + joint_names=["panda_joint.*"], + device="cuda:0" if wp.is_cuda_available() else "cpu", + ) + # Newton may produce > 1 articulation for instanceable USDs (e.g. the + # rootJoint fixed joint creates a separate one-joint articulation). + assert model.articulation_count >= 1, f"expected at least 1 articulation, got {model.articulation_count}" + assert isinstance(info, IKModelInfo) + assert info.arm_dof_count == 7 # Franka has 7 arm DOFs (panda_joint1..panda_joint7) + assert info.ee_link_index >= 0 + assert len(info.ik_joint_names) == 7 + assert all(n.startswith("panda_joint") for n in info.ik_joint_names) + assert len(info.sim_to_ik_joint_perm) == 7 + + +def test_build_raises_on_unknown_body(): + with pytest.raises(ValueError, match="not found"): + build_single_arm_ik_model( + asset_cfg=FRANKA_PANDA_CFG, + body_name="this_body_does_not_exist", + joint_names=["panda_joint.*"], + device="cuda:0" if wp.is_cuda_available() else "cpu", + ) diff --git a/source/isaaclab_newton/test/envs/__init__.py b/source/isaaclab_newton/test/envs/__init__.py new file mode 100644 index 000000000000..460a30569089 --- /dev/null +++ b/source/isaaclab_newton/test/envs/__init__.py @@ -0,0 +1,4 @@ +# 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 diff --git a/source/isaaclab_newton/test/envs/mdp/__init__.py b/source/isaaclab_newton/test/envs/mdp/__init__.py new file mode 100644 index 000000000000..460a30569089 --- /dev/null +++ b/source/isaaclab_newton/test/envs/mdp/__init__.py @@ -0,0 +1,4 @@ +# 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 diff --git a/source/isaaclab_newton/test/envs/mdp/actions/__init__.py b/source/isaaclab_newton/test/envs/mdp/actions/__init__.py new file mode 100644 index 000000000000..460a30569089 --- /dev/null +++ b/source/isaaclab_newton/test/envs/mdp/actions/__init__.py @@ -0,0 +1,4 @@ +# 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 diff --git a/source/isaaclab_newton/test/envs/mdp/actions/test_newton_ik_action.py b/source/isaaclab_newton/test/envs/mdp/actions/test_newton_ik_action.py new file mode 100644 index 000000000000..53c570985ebc --- /dev/null +++ b/source/isaaclab_newton/test/envs/mdp/actions/test_newton_ik_action.py @@ -0,0 +1,29 @@ +# 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 + +"""Unit tests for the Newton IK action term class wiring. + +Runtime pipeline behavior (reset/step/no-NaN) is validated in the gym smoke +test under ``test/envs/`` (T12), which exercises the action term inside a +real ManagerBasedRLEnv. This file only verifies the class-level contract. +""" + +from isaaclab.app import AppLauncher + +simulation_app = AppLauncher(headless=True).app + +from isaaclab_newton.envs.mdp.actions import NewtonInverseKinematicsAction + + +def test_action_class_is_exported(): + """The action class must be importable and named correctly.""" + assert NewtonInverseKinematicsAction.__name__ == "NewtonInverseKinematicsAction" + + +def test_action_module_path(): + """The module containing the action class must exist.""" + import isaaclab_newton.envs.mdp.actions.newton_ik_actions as mod + + assert hasattr(mod, "NewtonInverseKinematicsAction") diff --git a/source/isaaclab_newton/test/envs/mdp/actions/test_newton_ik_action_cfg.py b/source/isaaclab_newton/test/envs/mdp/actions/test_newton_ik_action_cfg.py new file mode 100644 index 000000000000..bae27540d4d7 --- /dev/null +++ b/source/isaaclab_newton/test/envs/mdp/actions/test_newton_ik_action_cfg.py @@ -0,0 +1,25 @@ +# 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.app import AppLauncher + +simulation_app = AppLauncher(headless=True).app + +from isaaclab_newton.controllers import NewtonIKControllerCfg +from isaaclab_newton.envs.mdp.actions import NewtonInverseKinematicsActionCfg + + +def test_action_cfg_defaults_are_sane(): + cfg = NewtonInverseKinematicsActionCfg( + asset_name="robot", + joint_names=["panda_joint.*"], + body_name="panda_hand", + controller=NewtonIKControllerCfg(), + ) + assert cfg.asset_name == "robot" + assert cfg.body_name == "panda_hand" + assert cfg.scale == 1.0 + assert cfg.body_offset is None + assert cfg.controller.iterations == 8 diff --git a/source/isaaclab_newton/test/envs/test_gym_franka_newton_ik_rel.py b/source/isaaclab_newton/test/envs/test_gym_franka_newton_ik_rel.py new file mode 100644 index 000000000000..ac4bdd8a1f7c --- /dev/null +++ b/source/isaaclab_newton/test/envs/test_gym_franka_newton_ik_rel.py @@ -0,0 +1,34 @@ +# Copyright (c) 2022-2026, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +"""Gym smoke test for the Franka reach Newton-IK relative-pose env.""" + +from isaaclab.app import AppLauncher + +simulation_app = AppLauncher(headless=True).app + +import gymnasium as gym +import torch + +import isaaclab_tasks # noqa: F401 registers gym envs +from isaaclab_tasks.utils.parse_cfg import parse_env_cfg + + +def test_gym_make_franka_newton_ik_rel(): + """Create env via gym.make, reset, step 20 times with small random actions, assert no NaN.""" + env_cfg = parse_env_cfg("Isaac-Reach-Franka-Newton-IK-Rel-v0", num_envs=4) + env = gym.make("Isaac-Reach-Franka-Newton-IK-Rel-v0", cfg=env_cfg) + try: + obs, _ = env.reset() + assert "policy" in obs, f"expected 'policy' key in obs, got {list(obs.keys())}" + torch.manual_seed(0) + for _ in range(20): + action = (torch.rand(4, 6, device=env.unwrapped.device) - 0.5) * 0.02 # ±1cm + obs, _, _, _, _ = env.step(action) + for key, val in obs.items(): + if isinstance(val, torch.Tensor): + assert not torch.isnan(val).any(), f"NaN in observation '{key}'" + finally: + env.close() diff --git a/source/isaaclab_newton/test/physics/test_sdf_config.py b/source/isaaclab_newton/test/physics/test_sdf_config.py new file mode 100644 index 000000000000..ec99c256b922 --- /dev/null +++ b/source/isaaclab_newton/test/physics/test_sdf_config.py @@ -0,0 +1,440 @@ +# 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 + +"""Tests for SDF collision configuration and application logic.""" + +import re +from unittest.mock import MagicMock, patch + +from newton import GeoType, ModelBuilder, ShapeFlags + + +class TestBuildSdfOnMesh: + """Tests for NewtonManager._build_sdf_on_mesh.""" + + @staticmethod + def _make_sdf_cfg(max_resolution=256, narrow_band_range=(-0.1, 0.1), target_voxel_size=None): + cfg = MagicMock() + cfg.max_resolution = max_resolution + cfg.narrow_band_range = narrow_band_range + cfg.target_voxel_size = target_voxel_size + return cfg + + def test_none_mesh_is_noop(self): + """Passing None as mesh should not raise.""" + from isaaclab_newton.physics.newton_manager import NewtonManager + + sdf_cfg = self._make_sdf_cfg() + # Should not raise + NewtonManager._build_sdf_on_mesh(None, sdf_cfg, None, "test_label") + + def test_builds_sdf_with_max_resolution(self): + """SDF is built on mesh with max_resolution and narrow_band_range.""" + from isaaclab_newton.physics.newton_manager import NewtonManager + + mesh = MagicMock() + mesh.sdf = None + sdf_cfg = self._make_sdf_cfg(max_resolution=128) + + NewtonManager._build_sdf_on_mesh(mesh, sdf_cfg, None, "test_label") + + mesh.build_sdf.assert_called_once_with(narrow_band_range=(-0.1, 0.1), max_resolution=128) + + def test_skips_rebuild_when_sdf_already_exists(self): + """Existing SDF on mesh is preserved (shared by reference from prototype).""" + from isaaclab_newton.physics.newton_manager import NewtonManager + + mesh = MagicMock() + mesh.sdf = "existing_sdf" + sdf_cfg = self._make_sdf_cfg() + + NewtonManager._build_sdf_on_mesh(mesh, sdf_cfg, None, "test_label") + + mesh.clear_sdf.assert_not_called() + mesh.build_sdf.assert_not_called() + + def test_target_voxel_size_takes_precedence(self): + """When target_voxel_size is set, it is passed alongside max_resolution.""" + from isaaclab_newton.physics.newton_manager import NewtonManager + + mesh = MagicMock() + mesh.sdf = None + sdf_cfg = self._make_sdf_cfg(max_resolution=256, target_voxel_size=0.005) + + NewtonManager._build_sdf_on_mesh(mesh, sdf_cfg, None, "test_label") + + call_kwargs = mesh.build_sdf.call_args[1] + assert call_kwargs["target_voxel_size"] == 0.005 + assert call_kwargs["max_resolution"] == 256 + + def test_resolution_override_by_pattern(self): + """Per-pattern resolution override is applied when label matches.""" + from isaaclab_newton.physics.newton_manager import NewtonManager + + mesh = MagicMock() + mesh.sdf = None + sdf_cfg = self._make_sdf_cfg(max_resolution=256) + res_overrides = [(re.compile(".*elbow.*"), 128)] + + NewtonManager._build_sdf_on_mesh(mesh, sdf_cfg, res_overrides, "/World/Robot/elbow_link/collision") + + call_kwargs = mesh.build_sdf.call_args[1] + assert call_kwargs["max_resolution"] == 128 + + def test_resolution_override_no_match_uses_global(self): + """When label doesn't match any override, global max_resolution is used.""" + from isaaclab_newton.physics.newton_manager import NewtonManager + + mesh = MagicMock() + mesh.sdf = None + sdf_cfg = self._make_sdf_cfg(max_resolution=256) + res_overrides = [(re.compile(".*elbow.*"), 128)] + + NewtonManager._build_sdf_on_mesh(mesh, sdf_cfg, res_overrides, "/World/Robot/wrist_link/collision") + + call_kwargs = mesh.build_sdf.call_args[1] + assert call_kwargs["max_resolution"] == 256 + + def test_resolution_override_first_match_wins(self): + """First matching pattern in res_overrides determines resolution.""" + from isaaclab_newton.physics.newton_manager import NewtonManager + + mesh = MagicMock() + mesh.sdf = None + sdf_cfg = self._make_sdf_cfg(max_resolution=256) + res_overrides = [ + (re.compile(".*link.*"), 64), + (re.compile(".*elbow.*"), 128), + ] + + NewtonManager._build_sdf_on_mesh(mesh, sdf_cfg, res_overrides, "/World/Robot/elbow_link/collision") + + call_kwargs = mesh.build_sdf.call_args[1] + assert call_kwargs["max_resolution"] == 64 # ".*link.*" matches first + + +class TestApplySdfConfig: + """Tests for NewtonManager._apply_sdf_config shape index collection and patching.""" + + @staticmethod + def _make_builder(bodies, shapes): + """Create a minimal ModelBuilder-like mock. + + Args: + bodies: List of body label strings. + shapes: List of dicts with keys: body_idx, label, geo_type, flags, source. + """ + builder = MagicMock(spec=ModelBuilder) + builder.body_label = bodies + builder.shape_count = len(shapes) + builder.shape_body = [s["body_idx"] for s in shapes] + builder.shape_label = [s["label"] for s in shapes] + builder.shape_type = [s["geo_type"] for s in shapes] + builder.shape_flags = [s["flags"] for s in shapes] + builder.shape_source = [s.get("source", MagicMock()) for s in shapes] + builder.shape_transform = [None] * len(shapes) + builder.shape_scale = [None] * len(shapes) + builder.shape_margin = [0.0] * len(shapes) + builder.shape_material_kh = [0.0] * len(shapes) + return builder + + def test_sdf_cfg_none_is_noop(self): + """When sdf_cfg is None, _apply_sdf_config does nothing.""" + from isaaclab_newton.physics.newton_manager import NewtonManager + + builder = MagicMock() + with patch.object(NewtonManager, "_build_sdf_on_mesh") as mock_build: + with patch("isaaclab_newton.physics.newton_manager.PhysicsManager") as mock_pm: + mock_pm._cfg = MagicMock() + mock_pm._cfg.sdf_cfg = None + NewtonManager._apply_sdf_config(builder) + mock_build.assert_not_called() + + def test_no_patterns_warns_and_returns(self): + """When both body_patterns and shape_patterns are None, warns and returns.""" + from isaaclab_newton.physics.newton_manager import NewtonManager + + builder = MagicMock() + with patch("isaaclab_newton.physics.newton_manager.PhysicsManager") as mock_pm: + sdf_cfg = MagicMock() + sdf_cfg.max_resolution = 256 + sdf_cfg.target_voxel_size = None + sdf_cfg.body_patterns = None + sdf_cfg.shape_patterns = None + sdf_cfg.hydroelastic_cfg = None + mock_pm._cfg = MagicMock() + mock_pm._cfg.sdf_cfg = sdf_cfg + + with patch("isaaclab_newton.physics.newton_manager.logger") as mock_logger: + NewtonManager._apply_sdf_config(builder) + mock_logger.warning.assert_called_once() + assert "no body_patterns or shape_patterns" in mock_logger.warning.call_args[0][0] + + def test_no_resolution_warns_and_returns(self): + """When neither max_resolution nor target_voxel_size is set, warns and returns.""" + from isaaclab_newton.physics.newton_manager import NewtonManager + + builder = MagicMock() + with patch("isaaclab_newton.physics.newton_manager.PhysicsManager") as mock_pm: + sdf_cfg = MagicMock() + sdf_cfg.max_resolution = None + sdf_cfg.target_voxel_size = None + mock_pm._cfg = MagicMock() + mock_pm._cfg.sdf_cfg = sdf_cfg + + with patch("isaaclab_newton.physics.newton_manager.logger") as mock_logger: + NewtonManager._apply_sdf_config(builder) + mock_logger.warning.assert_called_once() + assert "neither max_resolution nor target_voxel_size" in mock_logger.warning.call_args[0][0] + + def test_shape_patterns_match_shape_label(self): + """shape_patterns should match against shape_label, not body_label.""" + from isaaclab_newton.physics.newton_manager import NewtonManager + + builder = self._make_builder( + bodies=["/World/Robot/base", "/World/Robot/arm"], + shapes=[ + { + "body_idx": 0, + "label": "/World/Robot/base/gear_mesh", + "geo_type": GeoType.MESH, + "flags": ShapeFlags.COLLIDE_SHAPES, + }, + { + "body_idx": 1, + "label": "/World/Robot/arm/plain_mesh", + "geo_type": GeoType.MESH, + "flags": ShapeFlags.COLLIDE_SHAPES, + }, + ], + ) + # Mock the mesh sources + for s in builder.shape_source: + s.sdf = None + + with patch("isaaclab_newton.physics.newton_manager.PhysicsManager") as mock_pm: + sdf_cfg = MagicMock() + sdf_cfg.max_resolution = 256 + sdf_cfg.target_voxel_size = None + sdf_cfg.body_patterns = None + sdf_cfg.shape_patterns = [".*gear.*"] + sdf_cfg.pattern_resolutions = None + sdf_cfg.margin = None + sdf_cfg.use_visual_meshes = False + sdf_cfg.hydroelastic_cfg = None + sdf_cfg.narrow_band_range = (-0.1, 0.1) + mock_pm._cfg = MagicMock() + mock_pm._cfg.sdf_cfg = sdf_cfg + + NewtonManager._apply_sdf_config(builder) + + # Only gear_mesh should have had SDF built + builder.shape_source[0].build_sdf.assert_called_once() + builder.shape_source[1].build_sdf.assert_not_called() + + def test_body_patterns_match_body_label(self): + """body_patterns should match against body_label and apply SDF to all shapes under that body.""" + from isaaclab_newton.physics.newton_manager import NewtonManager + + builder = self._make_builder( + bodies=["/World/Robot/arm", "/World/Robot/base"], + shapes=[ + { + "body_idx": 0, + "label": "/World/Robot/arm/mesh_a", + "geo_type": GeoType.MESH, + "flags": ShapeFlags.COLLIDE_SHAPES, + }, + { + "body_idx": 0, + "label": "/World/Robot/arm/mesh_b", + "geo_type": GeoType.MESH, + "flags": ShapeFlags.COLLIDE_SHAPES, + }, + { + "body_idx": 1, + "label": "/World/Robot/base/mesh_c", + "geo_type": GeoType.MESH, + "flags": ShapeFlags.COLLIDE_SHAPES, + }, + ], + ) + for s in builder.shape_source: + s.sdf = None + + with patch("isaaclab_newton.physics.newton_manager.PhysicsManager") as mock_pm: + sdf_cfg = MagicMock() + sdf_cfg.max_resolution = 256 + sdf_cfg.target_voxel_size = None + sdf_cfg.body_patterns = [".*arm.*"] + sdf_cfg.shape_patterns = None + sdf_cfg.pattern_resolutions = None + sdf_cfg.margin = None + sdf_cfg.use_visual_meshes = False + sdf_cfg.hydroelastic_cfg = None + sdf_cfg.narrow_band_range = (-0.1, 0.1) + mock_pm._cfg = MagicMock() + mock_pm._cfg.sdf_cfg = sdf_cfg + + NewtonManager._apply_sdf_config(builder) + + # Both arm shapes should get SDF, base should not + builder.shape_source[0].build_sdf.assert_called_once() + builder.shape_source[1].build_sdf.assert_called_once() + builder.shape_source[2].build_sdf.assert_not_called() + + def test_non_mesh_shapes_are_skipped(self): + """Non-MESH shapes (e.g., BOX, SPHERE) should never get SDF even if patterns match.""" + from isaaclab_newton.physics.newton_manager import NewtonManager + + builder = self._make_builder( + bodies=["/World/Robot/base"], + shapes=[ + { + "body_idx": 0, + "label": "/World/Robot/base/box", + "geo_type": GeoType.BOX, + "flags": ShapeFlags.COLLIDE_SHAPES, + }, + { + "body_idx": 0, + "label": "/World/Robot/base/mesh", + "geo_type": GeoType.MESH, + "flags": ShapeFlags.COLLIDE_SHAPES, + }, + ], + ) + for s in builder.shape_source: + s.sdf = None + + with patch("isaaclab_newton.physics.newton_manager.PhysicsManager") as mock_pm: + sdf_cfg = MagicMock() + sdf_cfg.max_resolution = 256 + sdf_cfg.target_voxel_size = None + sdf_cfg.body_patterns = [".*base.*"] + sdf_cfg.shape_patterns = None + sdf_cfg.pattern_resolutions = None + sdf_cfg.margin = None + sdf_cfg.use_visual_meshes = False + sdf_cfg.hydroelastic_cfg = None + sdf_cfg.narrow_band_range = (-0.1, 0.1) + mock_pm._cfg = MagicMock() + mock_pm._cfg.sdf_cfg = sdf_cfg + + NewtonManager._apply_sdf_config(builder) + + # Box should not get SDF, mesh should + builder.shape_source[0].build_sdf.assert_not_called() + builder.shape_source[1].build_sdf.assert_called_once() + + def test_visual_shapes_skipped_when_use_visual_meshes_false(self): + """Shapes without COLLIDE_SHAPES flag should not be patched when use_visual_meshes=False.""" + from isaaclab_newton.physics.newton_manager import NewtonManager + + builder = self._make_builder( + bodies=["/World/Robot/arm"], + shapes=[ + { + "body_idx": 0, + "label": "/World/Robot/arm/visual", + "geo_type": GeoType.MESH, + "flags": ShapeFlags.VISIBLE, + }, + ], + ) + for s in builder.shape_source: + s.sdf = None + + with patch("isaaclab_newton.physics.newton_manager.PhysicsManager") as mock_pm: + sdf_cfg = MagicMock() + sdf_cfg.max_resolution = 256 + sdf_cfg.target_voxel_size = None + sdf_cfg.body_patterns = [".*arm.*"] + sdf_cfg.shape_patterns = None + sdf_cfg.pattern_resolutions = None + sdf_cfg.margin = None + sdf_cfg.use_visual_meshes = False + sdf_cfg.hydroelastic_cfg = None + sdf_cfg.narrow_band_range = (-0.1, 0.1) + mock_pm._cfg = MagicMock() + mock_pm._cfg.sdf_cfg = sdf_cfg + + NewtonManager._apply_sdf_config(builder) + + # Visual-only shape should not get SDF built + builder.shape_source[0].build_sdf.assert_not_called() + + def test_margin_applied_when_set(self): + """When sdf_cfg.margin is set, it should be applied to matched collision shapes.""" + from isaaclab_newton.physics.newton_manager import NewtonManager + + builder = self._make_builder( + bodies=["/World/Robot/arm"], + shapes=[ + { + "body_idx": 0, + "label": "/World/Robot/arm/mesh", + "geo_type": GeoType.MESH, + "flags": ShapeFlags.COLLIDE_SHAPES, + }, + ], + ) + for s in builder.shape_source: + s.sdf = None + + with patch("isaaclab_newton.physics.newton_manager.PhysicsManager") as mock_pm: + sdf_cfg = MagicMock() + sdf_cfg.max_resolution = 256 + sdf_cfg.target_voxel_size = None + sdf_cfg.body_patterns = [".*arm.*"] + sdf_cfg.shape_patterns = None + sdf_cfg.pattern_resolutions = None + sdf_cfg.margin = 0.02 + sdf_cfg.use_visual_meshes = False + sdf_cfg.hydroelastic_cfg = None + sdf_cfg.narrow_band_range = (-0.1, 0.1) + mock_pm._cfg = MagicMock() + mock_pm._cfg.sdf_cfg = sdf_cfg + + NewtonManager._apply_sdf_config(builder) + + assert builder.shape_margin[0] == 0.02 + + def test_margin_none_leaves_default(self): + """When sdf_cfg.margin is None, shape_margin should not be modified.""" + from isaaclab_newton.physics.newton_manager import NewtonManager + + builder = self._make_builder( + bodies=["/World/Robot/arm"], + shapes=[ + { + "body_idx": 0, + "label": "/World/Robot/arm/mesh", + "geo_type": GeoType.MESH, + "flags": ShapeFlags.COLLIDE_SHAPES, + }, + ], + ) + for s in builder.shape_source: + s.sdf = None + original_margin = builder.shape_margin[0] + + with patch("isaaclab_newton.physics.newton_manager.PhysicsManager") as mock_pm: + sdf_cfg = MagicMock() + sdf_cfg.max_resolution = 256 + sdf_cfg.target_voxel_size = None + sdf_cfg.body_patterns = [".*arm.*"] + sdf_cfg.shape_patterns = None + sdf_cfg.pattern_resolutions = None + sdf_cfg.margin = None + sdf_cfg.use_visual_meshes = False + sdf_cfg.hydroelastic_cfg = None + sdf_cfg.narrow_band_range = (-0.1, 0.1) + mock_pm._cfg = MagicMock() + mock_pm._cfg.sdf_cfg = sdf_cfg + + NewtonManager._apply_sdf_config(builder) + + assert builder.shape_margin[0] == original_margin diff --git a/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/reach/config/franka/__init__.py b/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/reach/config/franka/__init__.py index 47158f64a650..3ccffdc4c6d9 100644 --- a/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/reach/config/franka/__init__.py +++ b/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/reach/config/franka/__init__.py @@ -49,6 +49,7 @@ entry_point="isaaclab.envs:ManagerBasedRLEnv", kwargs={ "env_cfg_entry_point": f"{__name__}.ik_abs_env_cfg:FrankaReachEnvCfg", + "rsl_rl_cfg_entry_point": f"{agents.__name__}.rsl_rl_ppo_cfg:FrankaReachPPORunnerCfg", }, disable_env_checker=True, ) @@ -62,6 +63,7 @@ entry_point="isaaclab.envs:ManagerBasedRLEnv", kwargs={ "env_cfg_entry_point": f"{__name__}.ik_rel_env_cfg:FrankaReachEnvCfg", + "rsl_rl_cfg_entry_point": f"{agents.__name__}.rsl_rl_ppo_cfg:FrankaReachPPORunnerCfg", }, disable_env_checker=True, ) @@ -89,3 +91,27 @@ "rsl_rl_cfg_entry_point": f"{agents.__name__}.rsl_rl_ppo_cfg:FrankaReachPPORunnerCfg", }, ) + +## +# 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__}.newton_ik_rel_env_cfg:FrankaNewtonIKReachEnvCfg", + "rsl_rl_cfg_entry_point": f"{agents.__name__}.rsl_rl_ppo_cfg:FrankaReachPPORunnerCfg", + }, + 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__}.newton_ik_rel_env_cfg:FrankaNewtonIKReachEnvCfg_PLAY", + "rsl_rl_cfg_entry_point": f"{agents.__name__}.rsl_rl_ppo_cfg:FrankaReachPPORunnerCfg", + }, + disable_env_checker=True, +) diff --git a/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/reach/config/franka/newton_ik_rel_env_cfg.py b/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/reach/config/franka/newton_ik_rel_env_cfg.py new file mode 100644 index 000000000000..bd064350374a --- /dev/null +++ b/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/reach/config/franka/newton_ik_rel_env_cfg.py @@ -0,0 +1,100 @@ +# 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 + +"""Franka reach env using Newton physics + Newton IK relative-pose control. + +Train command:: + + ./isaaclab.sh -p scripts/reinforcement_learning/rsl_rl/train.py \\ + --task Isaac-Reach-Franka-Newton-IK-Rel-v0 --num_envs 4096 --headless + +Play command:: + + ./isaaclab.sh -p scripts/reinforcement_learning/rsl_rl/play.py \\ + --task Isaac-Reach-Franka-Newton-IK-Rel-Play-v0 +""" + +from isaaclab_newton.controllers import NewtonIKControllerCfg +from isaaclab_newton.envs.mdp.actions import NewtonInverseKinematicsActionCfg +from isaaclab_newton.physics import MJWarpSolverCfg, NewtonCfg + +from isaaclab.utils 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 FrankaNewtonIKReachEnvCfg(joint_pos_env_cfg.FrankaReachEnvCfg): + def __post_init__(self): + super().__post_init__() + + # Switch to Newton physics (parent uses PhysX preset by default). + self.sim.physics = NewtonCfg( + solver_cfg=MJWarpSolverCfg( + njmax=50, + nconmax=20, + cone="pyramidal", + integrator="implicitfast", + impratio=1, + ), + num_substeps=1, + debug_mode=False, + ) + + # Use the high-PD Franka so IK joint targets track well. + self.scene.robot = FRANKA_PANDA_HIGH_PD_CFG.replace(prim_path="{ENV_REGEX_NS}/Robot") + + # Newton IK action term. + # NOTE: ``scale=0.5`` matches the PhysX ``DifferentialInverseKinematicsActionCfg`` + # baseline so the PPO input distribution is identical across backends. + # + # IK-specific knobs (no PhysX equivalent, empirically tuned for RL + # tracking against the PhysX DLS baseline on 1024-env Franka reach): + # - ``iterations=1``: a single Gauss-Newton (LM) step per control + # step. More iterations (e.g. 2) make LM converge more tightly, + # commanding joint deltas the PD cannot track in one 33 ms control + # window, which regresses position tracking significantly. + # - ``lambda_initial=1.0``: higher damping shrinks per-step joint + # commands toward what the PD can track. + # - ``joint_limit_weight=0.02``: the default 0.1 competes too hard + # with position/rotation objectives inside the reach workspace. + # - ``rotation_weight=60.0``: unit-normalizes rotation (rad) against + # position (m) — 1 cm pos ≈ 1° (0.017 rad) rot contribute equally + # to the LM cost. Higher rw (e.g. 100) makes rotation cost + # dominate and position tracking regresses. + # + # Measured vs PhysX DLS baseline at 3000 PPO iters: + # - EE position: 22 mm (PhysX 18 mm) + # - EE orientation: 0.27 rad (PhysX 0.17 rad) + # - Throughput: ~1.5× faster than PhysX per step + self.actions.arm_action = NewtonInverseKinematicsActionCfg( + asset_name="robot", + joint_names=["panda_joint.*"], + body_name="panda_hand", + controller=NewtonIKControllerCfg( + command_type="pose", + use_relative_mode=True, + iterations=1, + joint_limit_weight=0.02, + lambda_initial=1.0, + rotation_weight=60.0, + ), + scale=0.5, + body_offset=NewtonInverseKinematicsActionCfg.OffsetCfg(pos=[0.0, 0.0, 0.107]), + ) + + +@configclass +class FrankaNewtonIKReachEnvCfg_PLAY(FrankaNewtonIKReachEnvCfg): + def __post_init__(self): + super().__post_init__() + self.scene.num_envs = 50 + self.scene.env_spacing = 2.5 + self.observations.policy.enable_corruption = False diff --git a/source/isaaclab_visualizers/isaaclab_visualizers/newton/newton_visualizer.py b/source/isaaclab_visualizers/isaaclab_visualizers/newton/newton_visualizer.py index 70d44d1956f7..d1a62f3bbfe3 100644 --- a/source/isaaclab_visualizers/isaaclab_visualizers/newton/newton_visualizer.py +++ b/source/isaaclab_visualizers/isaaclab_visualizers/newton/newton_visualizer.py @@ -190,9 +190,18 @@ def _render_left_panel(self): show_contacts = self.show_contacts changed, self.show_contacts = imgui.checkbox("Show Contacts", show_contacts) + show_collision = self.show_collision + changed, self.show_collision = imgui.checkbox("Show Collision", show_collision) + + show_visual = self.show_visual + changed, self.show_visual = imgui.checkbox("Show Visual Geometries", show_visual) + show_springs = self.show_springs changed, self.show_springs = imgui.checkbox("Show Springs", show_springs) + show_inertia_boxes = self.show_inertia_boxes + changed, self.show_inertia_boxes = imgui.checkbox("Show Inertia Boxes", show_inertia_boxes) + show_com = self.show_com changed, self.show_com = imgui.checkbox("Show Center of Mass", show_com) @@ -316,7 +325,10 @@ def initialize(self, scene_data_provider: BaseSceneDataProvider) -> None: self._viewer.show_joints = self.cfg.show_joints self._viewer.show_contacts = self.cfg.show_contacts + self._viewer.show_collision = self.cfg.show_collision + self._viewer.show_visual = self.cfg.show_visual self._viewer.show_springs = self.cfg.show_springs + self._viewer.show_inertia_boxes = self.cfg.show_inertia_boxes self._viewer.show_com = self.cfg.show_com self._viewer.renderer.draw_shadows = self.cfg.enable_shadows diff --git a/source/isaaclab_visualizers/isaaclab_visualizers/newton/newton_visualizer_cfg.py b/source/isaaclab_visualizers/isaaclab_visualizers/newton/newton_visualizer_cfg.py index 2449db73b5e3..0d3a93897e3e 100644 --- a/source/isaaclab_visualizers/isaaclab_visualizers/newton/newton_visualizer_cfg.py +++ b/source/isaaclab_visualizers/isaaclab_visualizers/newton/newton_visualizer_cfg.py @@ -40,9 +40,18 @@ class NewtonVisualizerCfg(VisualizerCfg): show_contacts: bool = False """Show contact visualization.""" + show_collision: bool = False + """Show collision visualization.""" + + show_visual: bool = True + """Show visual geometry.""" + show_springs: bool = False """Show spring visualization.""" + show_inertia_boxes: bool = False + """Show inertia box visualization.""" + show_com: bool = False """Show center of mass visualization."""