From ac23116326d1db1fa5a882fabbf8f68f7b91bd73 Mon Sep 17 00:00:00 2001 From: curiep Date: Thu, 30 Jul 2026 12:54:28 -0700 Subject: [PATCH 1/4] Add Newton gear insertion training Add Newton point-SDF and hydroelastic presets, local concave collision assets, and backend-specific reset, reward, and gripper configuration for Rizon 4s gear assembly. Keep PhysX behavior unchanged, document Newton training, and add configuration regression coverage plus a flange IK variant. --- .../02_gear_assembly/gear_assembly_policy.rst | 80 +++- .../curiep-newton-gear-insertion.minor.rst | 13 + .../factory_gear_base/factory_gear_base.usda | 3 + .../factory_gear_base_hydroelastic.usda | 3 + .../factory_gear_large.usda | 3 + .../factory_gear_large_hydroelastic.usda | 3 + .../factory_gear_medium.usda | 3 + .../factory_gear_medium_hydroelastic.usda | 3 + .../factory_gear_small.usda | 3 + .../factory_gear_small_hydroelastic.usda | 3 + .../gear_assembly/config/rizon_4s/__init__.py | 11 + .../config/rizon_4s/ik_newton_env_cfg.py | 48 ++ .../config/rizon_4s/joint_pos_env_cfg.py | 423 ++++++++++-------- .../config/ur_10e/joint_pos_env_cfg.py | 4 + .../gear_assembly/gear_assembly_env_cfg.py | 262 +++++++++-- .../contrib/deploy/mdp/__init__.pyi | 8 +- .../contrib/deploy/mdp/events.py | 303 +++++++++++-- .../contrib/deploy/mdp/rewards.py | 125 ++++-- .../contrib/deploy/mdp/terminations.py | 20 +- .../contrib/test_deploy_gear_assembly_cfg.py | 127 ++++++ 20 files changed, 1127 insertions(+), 321 deletions(-) create mode 100644 source/isaaclab_tasks/changelog.d/curiep-newton-gear-insertion.minor.rst create mode 100644 source/isaaclab_tasks/isaaclab_tasks/contrib/deploy/gear_assembly/assets/newton/factory_gear_base/factory_gear_base.usda create mode 100644 source/isaaclab_tasks/isaaclab_tasks/contrib/deploy/gear_assembly/assets/newton/factory_gear_base/factory_gear_base_hydroelastic.usda create mode 100644 source/isaaclab_tasks/isaaclab_tasks/contrib/deploy/gear_assembly/assets/newton/factory_gear_large/factory_gear_large.usda create mode 100644 source/isaaclab_tasks/isaaclab_tasks/contrib/deploy/gear_assembly/assets/newton/factory_gear_large/factory_gear_large_hydroelastic.usda create mode 100644 source/isaaclab_tasks/isaaclab_tasks/contrib/deploy/gear_assembly/assets/newton/factory_gear_medium/factory_gear_medium.usda create mode 100644 source/isaaclab_tasks/isaaclab_tasks/contrib/deploy/gear_assembly/assets/newton/factory_gear_medium/factory_gear_medium_hydroelastic.usda create mode 100644 source/isaaclab_tasks/isaaclab_tasks/contrib/deploy/gear_assembly/assets/newton/factory_gear_small/factory_gear_small.usda create mode 100644 source/isaaclab_tasks/isaaclab_tasks/contrib/deploy/gear_assembly/assets/newton/factory_gear_small/factory_gear_small_hydroelastic.usda create mode 100644 source/isaaclab_tasks/isaaclab_tasks/contrib/deploy/gear_assembly/config/rizon_4s/ik_newton_env_cfg.py diff --git a/docs/source/policy_deployment/02_gear_assembly/gear_assembly_policy.rst b/docs/source/policy_deployment/02_gear_assembly/gear_assembly_policy.rst index 5b39572eb2d2..ff30123d67a7 100644 --- a/docs/source/policy_deployment/02_gear_assembly/gear_assembly_policy.rst +++ b/docs/source/policy_deployment/02_gear_assembly/gear_assembly_policy.rst @@ -734,7 +734,36 @@ First, launch the training with a small number of environments and visualization --num_envs 4 \ --visualizer kit -This will open the Isaac Sim viewer where you can observe the training process in real-time. + .. tab-item:: Flexiv Rizon 4s + Grav (Newton) + + .. tab-set:: + + .. tab-item:: uv (Recommended) + + .. code-block:: bash + + uv run isaaclab train --rl_library rsl_rl \ + --task IsaacContrib-Deploy-GearAssembly-Rizon4s-Grav \ + --num_envs 2 \ + --visualizer newton \ + presets=newton_hydroelastic + + .. tab-item:: isaaclab.sh / isaaclab.bat + + .. code-block:: bash + + ./isaaclab.sh train --rl_library rsl_rl \ + --task IsaacContrib-Deploy-GearAssembly-Rizon4s-Grav \ + --num_envs 2 \ + --visualizer newton \ + presets=newton_hydroelastic + + Use ``presets=newton_sdf`` to compare point-SDF contacts. To train a task-space + policy, use the ``IsaacContrib-Deploy-GearAssembly-Rizon4s-Grav-Newton-IK`` task + with the same Newton preset. + +This opens the selected viewer (Isaac Sim for PhysX or Newton Viewer for Newton) so you can +observe the training process in real time. .. figure:: ../../_static/policy_deployment/02_gear_assembly/sim_real_gear_assembly_train.jpg :align: center @@ -820,10 +849,55 @@ Now launch the full training run with more parallel environments in headless mod --num_envs 256 \ --video --video_length 200 --video_interval 76800 + .. tab-item:: Flexiv Rizon 4s + Grav (Newton) + + .. tab-set:: + + .. tab-item:: uv (Recommended) + + .. code-block:: bash + + uv run isaaclab train --rl_library rsl_rl \ + --task IsaacContrib-Deploy-GearAssembly-Rizon4s-Grav \ + --num_envs 256 \ + --visualizer none \ + presets=newton_hydroelastic + + .. tab-item:: isaaclab.sh / isaaclab.bat + + .. code-block:: bash + + ./isaaclab.sh train --rl_library rsl_rl \ + --task IsaacContrib-Deploy-GearAssembly-Rizon4s-Grav \ + --num_envs 256 \ + --visualizer none \ + presets=newton_hydroelastic + +The Newton hydroelastic preset uses package-local concave SDF collision assets. Its defaults +separate the expensive collision rate from contact integration: the outer physics and collision +tick is 100 Hz, while 20 solver substeps produce a 2 kHz solver rate. A policy decimation of three +gives a 33.3 Hz control rate. The preset also uses a 5 mm shape gap and a bounded triangle-pair +buffer. The arm uses solver-native actuator gravity compensation so gravity remains enabled for +the gears, and the Grav linkage is held by physical PD actuators; no action term rewrites the +selected gear pose. + +For four-GPU training, launch 256 environments per rank (1,024 total): + +.. code-block:: bash + + uv run python scripts/reinforcement_learning/train_multigpu.py \ + --rl_library rsl_rl --num_gpus 4 \ + --task IsaacContrib-Deploy-GearAssembly-Rizon4s-Grav \ + --num_envs 256 --max_iterations 1000 \ + --visualizer none \ + presets=newton_hydroelastic + **Command breakdown:** -- ``--num_envs 256``: Runs 256 parallel environments for efficient training -- ``--video_length 200``: Each video captures approximately one full episode (``episode_length_s / (sim.dt * decimation)`` = ``6.66 / (1/1000 * 33)`` ≈ 200 steps) +- ``--num_envs 256``: Runs 256 parallel environments per trainer process +- ``--video_length 200``: Captures approximately one PhysX episode + (``6.66 / (1/120 * 4)`` ≈ 200 policy steps); a Newton episode is approximately + 222 policy steps (``6.66 / (0.01 * 3)``) - ``--video_interval 76800``: Records a video every 76,800 environment steps (~every 150 iterations), producing ~10 videos over full training Training typically takes ~12-24 hours for a robust insertion policy. The videos will be saved in the ``logs`` directory and can be reviewed to assess policy performance during training. diff --git a/source/isaaclab_tasks/changelog.d/curiep-newton-gear-insertion.minor.rst b/source/isaaclab_tasks/changelog.d/curiep-newton-gear-insertion.minor.rst new file mode 100644 index 000000000000..670536a1b2f8 --- /dev/null +++ b/source/isaaclab_tasks/changelog.d/curiep-newton-gear-insertion.minor.rst @@ -0,0 +1,13 @@ +Added +^^^^^ + +* Added Newton point-SDF and hydroelastic-SDF physics presets and collision + assets to the Flexiv Rizon 4s gear assembly task. +* Added a Newton inverse-kinematics task variant for task-space policies. + +Fixed +^^^^^ + +* Fixed Newton gear reset, reward, and termination frames so selected gears + target their shafts and physical grasps are measured at the fingertip + midpoint without changing the PhysX defaults. diff --git a/source/isaaclab_tasks/isaaclab_tasks/contrib/deploy/gear_assembly/assets/newton/factory_gear_base/factory_gear_base.usda b/source/isaaclab_tasks/isaaclab_tasks/contrib/deploy/gear_assembly/assets/newton/factory_gear_base/factory_gear_base.usda new file mode 100644 index 000000000000..f66f09bd8392 --- /dev/null +++ b/source/isaaclab_tasks/isaaclab_tasks/contrib/deploy/gear_assembly/assets/newton/factory_gear_base/factory_gear_base.usda @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:2810d8d2faecd2a7d8e08ab78aa59dbfc513f77f047b3f49a8506ba99bc82771 +size 240517 diff --git a/source/isaaclab_tasks/isaaclab_tasks/contrib/deploy/gear_assembly/assets/newton/factory_gear_base/factory_gear_base_hydroelastic.usda b/source/isaaclab_tasks/isaaclab_tasks/contrib/deploy/gear_assembly/assets/newton/factory_gear_base/factory_gear_base_hydroelastic.usda new file mode 100644 index 000000000000..9462f1626f28 --- /dev/null +++ b/source/isaaclab_tasks/isaaclab_tasks/contrib/deploy/gear_assembly/assets/newton/factory_gear_base/factory_gear_base_hydroelastic.usda @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:e2cf23310834fff8d049797a9e47a383478f501cd326fd59a800b95cb51523dd +size 473 diff --git a/source/isaaclab_tasks/isaaclab_tasks/contrib/deploy/gear_assembly/assets/newton/factory_gear_large/factory_gear_large.usda b/source/isaaclab_tasks/isaaclab_tasks/contrib/deploy/gear_assembly/assets/newton/factory_gear_large/factory_gear_large.usda new file mode 100644 index 000000000000..1e03caa179c0 --- /dev/null +++ b/source/isaaclab_tasks/isaaclab_tasks/contrib/deploy/gear_assembly/assets/newton/factory_gear_large/factory_gear_large.usda @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:f72d166f8978491cc5ccc826ba1683d8ae817722065f34779ffdafb4c8e4ec6a +size 690305 diff --git a/source/isaaclab_tasks/isaaclab_tasks/contrib/deploy/gear_assembly/assets/newton/factory_gear_large/factory_gear_large_hydroelastic.usda b/source/isaaclab_tasks/isaaclab_tasks/contrib/deploy/gear_assembly/assets/newton/factory_gear_large/factory_gear_large_hydroelastic.usda new file mode 100644 index 000000000000..9fb03f560b7e --- /dev/null +++ b/source/isaaclab_tasks/isaaclab_tasks/contrib/deploy/gear_assembly/assets/newton/factory_gear_large/factory_gear_large_hydroelastic.usda @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:12dfce05c39ec5927d777c59961999eec99ba6fa6c57498715a148d5a8039317 +size 477 diff --git a/source/isaaclab_tasks/isaaclab_tasks/contrib/deploy/gear_assembly/assets/newton/factory_gear_medium/factory_gear_medium.usda b/source/isaaclab_tasks/isaaclab_tasks/contrib/deploy/gear_assembly/assets/newton/factory_gear_medium/factory_gear_medium.usda new file mode 100644 index 000000000000..8638cdad4fd2 --- /dev/null +++ b/source/isaaclab_tasks/isaaclab_tasks/contrib/deploy/gear_assembly/assets/newton/factory_gear_medium/factory_gear_medium.usda @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:3736a62c7c48b3e8ce9c8773e9983c6e9d2ecfd52d8e39a7a1c50e0360f56868 +size 495131 diff --git a/source/isaaclab_tasks/isaaclab_tasks/contrib/deploy/gear_assembly/assets/newton/factory_gear_medium/factory_gear_medium_hydroelastic.usda b/source/isaaclab_tasks/isaaclab_tasks/contrib/deploy/gear_assembly/assets/newton/factory_gear_medium/factory_gear_medium_hydroelastic.usda new file mode 100644 index 000000000000..e59584596bc0 --- /dev/null +++ b/source/isaaclab_tasks/isaaclab_tasks/contrib/deploy/gear_assembly/assets/newton/factory_gear_medium/factory_gear_medium_hydroelastic.usda @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:72e61fea6b51a6bf364680b9b2eb289efa2acf8bfc74d83d1faef04961fe0b02 +size 481 diff --git a/source/isaaclab_tasks/isaaclab_tasks/contrib/deploy/gear_assembly/assets/newton/factory_gear_small/factory_gear_small.usda b/source/isaaclab_tasks/isaaclab_tasks/contrib/deploy/gear_assembly/assets/newton/factory_gear_small/factory_gear_small.usda new file mode 100644 index 000000000000..98607558332b --- /dev/null +++ b/source/isaaclab_tasks/isaaclab_tasks/contrib/deploy/gear_assembly/assets/newton/factory_gear_small/factory_gear_small.usda @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:e2122445d00107c7698992ade203c72553988298e8a5c01ea58900e7bdf140c7 +size 331955 diff --git a/source/isaaclab_tasks/isaaclab_tasks/contrib/deploy/gear_assembly/assets/newton/factory_gear_small/factory_gear_small_hydroelastic.usda b/source/isaaclab_tasks/isaaclab_tasks/contrib/deploy/gear_assembly/assets/newton/factory_gear_small/factory_gear_small_hydroelastic.usda new file mode 100644 index 000000000000..d17b73df9a23 --- /dev/null +++ b/source/isaaclab_tasks/isaaclab_tasks/contrib/deploy/gear_assembly/assets/newton/factory_gear_small/factory_gear_small_hydroelastic.usda @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:43d91b4a9701c040c6df4136e2404540f3caa04d6a9becbe07b3f526a1fa7af7 +size 477 diff --git a/source/isaaclab_tasks/isaaclab_tasks/contrib/deploy/gear_assembly/config/rizon_4s/__init__.py b/source/isaaclab_tasks/isaaclab_tasks/contrib/deploy/gear_assembly/config/rizon_4s/__init__.py index 3885057c2596..2ff1b032e44b 100644 --- a/source/isaaclab_tasks/isaaclab_tasks/contrib/deploy/gear_assembly/config/rizon_4s/__init__.py +++ b/source/isaaclab_tasks/isaaclab_tasks/contrib/deploy/gear_assembly/config/rizon_4s/__init__.py @@ -23,6 +23,17 @@ }, ) +# Flexiv Rizon 4s with Newton inverse-kinematics actions +gym.register( + id="IsaacContrib-Deploy-GearAssembly-Rizon4s-Grav-Newton-IK", + entry_point="isaaclab.envs:ManagerBasedRLEnv", + disable_env_checker=True, + kwargs={ + "env_cfg_entry_point": f"{__name__}.ik_newton_env_cfg:Rizon4sGearAssemblyIKNewtonEnvCfg", + "rsl_rl_cfg_entry_point": f"{agents.__name__}.rsl_rl_ppo_cfg:Rizon4sGearAssemblyRNNPPORunnerCfg", + }, +) + # Flexiv Rizon 4s - ROS Inference gym.register( id="IsaacContrib-Deploy-GearAssembly-Rizon4s-Grav-ROS-Inference", diff --git a/source/isaaclab_tasks/isaaclab_tasks/contrib/deploy/gear_assembly/config/rizon_4s/ik_newton_env_cfg.py b/source/isaaclab_tasks/isaaclab_tasks/contrib/deploy/gear_assembly/config/rizon_4s/ik_newton_env_cfg.py new file mode 100644 index 000000000000..389cf28e8342 --- /dev/null +++ b/source/isaaclab_tasks/isaaclab_tasks/contrib/deploy/gear_assembly/config/rizon_4s/ik_newton_env_cfg.py @@ -0,0 +1,48 @@ +# Copyright (c) 2022-2026, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +from isaaclab_newton.envs.mdp.actions.newton_ik_actions_cfg import NewtonInverseKinematicsActionCfg +from isaaclab_newton.ik.newton_ik_objectives_cfg import NewtonIKJointLimitObjectiveCfg, NewtonIKPoseObjectiveCfg +from isaaclab_newton.ik.newton_ik_solver_cfg import NewtonIKSolverCfg + +from isaaclab.utils.configclass import configclass + +from . import joint_pos_env_cfg + + +@configclass +class Rizon4sGearAssemblyIKNewtonEnvCfg(joint_pos_env_cfg.Rizon4sGearAssemblyEnvCfg): + """Gear-assembly with a Newton inverse-kinematics (task-space) action for the Rizon 4s arm. + + Replaces the joint-space :class:`~isaaclab.envs.mdp.RelativeJointPositionActionCfg` with a + Newton-solved relative end-effector pose action, so the policy commands end-effector motion + directly. This provides a six-dimensional task-space alternative to the seven-dimensional + joint-space action. The gripper remains a fixed mimic mechanism (not commanded by the policy). + + Note: + Newton IK consumes the replicated robot prototype, so this variant requires a Newton preset + (``presets=newton_mjwarp``, ``presets=newton_sdf``, or + ``presets=newton_hydroelastic``); it is not compatible with ``presets=physx``. + """ + + def __post_init__(self): + super().__post_init__() + + # Command the physical flange frame used by the real Flexiv Cartesian controller. + self.actions.arm_action = NewtonInverseKinematicsActionCfg( + asset_name="robot", + joint_names=["joint1", "joint2", "joint3", "joint4", "joint5", "joint6", "joint7"], + controller=NewtonIKSolverCfg(optimizer="lm", jacobian_mode="analytic", iterations=24), + clip={".*": (-0.5, 0.5)}, + objectives=[ + NewtonIKPoseObjectiveCfg( + body_name="flange", + command_type="pose", + use_relative_mode=True, + scale=0.025, + ), + NewtonIKJointLimitObjectiveCfg(weight=0.1), + ], + ) diff --git a/source/isaaclab_tasks/isaaclab_tasks/contrib/deploy/gear_assembly/config/rizon_4s/joint_pos_env_cfg.py b/source/isaaclab_tasks/isaaclab_tasks/contrib/deploy/gear_assembly/config/rizon_4s/joint_pos_env_cfg.py index 885e9b8b0d27..38e0730405c0 100644 --- a/source/isaaclab_tasks/isaaclab_tasks/contrib/deploy/gear_assembly/config/rizon_4s/joint_pos_env_cfg.py +++ b/source/isaaclab_tasks/isaaclab_tasks/contrib/deploy/gear_assembly/config/rizon_4s/joint_pos_env_cfg.py @@ -6,6 +6,11 @@ import math import torch +from isaaclab_physx.sim.schemas import ( + PhysxArticulationRootPropertiesCfg, + PhysxCollisionPropertiesCfg, + PhysxRigidBodyPropertiesCfg, +) import isaaclab.sim as sim_utils from isaaclab.actuators import ImplicitActuatorCfg @@ -18,11 +23,16 @@ import isaaclab_tasks.contrib.deploy.mdp as mdp import isaaclab_tasks.contrib.deploy.mdp.events as gear_assembly_events -from isaaclab_tasks.contrib.deploy.gear_assembly.gear_assembly_env_cfg import GearAssemblyEnvCfg +from isaaclab_tasks.contrib.deploy.gear_assembly.gear_assembly_env_cfg import ( + _NEWTON_GEAR_OFFSETS, + _PHYSX_GEAR_OFFSETS, + GearAssemblyEnvCfg, +) from isaaclab_tasks.contrib.deploy.mdp.noise_models import ( ResetSampledConstantNoiseModelCfg, ResetSampledQuaternionNoiseModelCfg, ) +from isaaclab_tasks.utils import PresetCfg, preset ## # Pre-defined configs @@ -35,44 +45,71 @@ ## +_GRAV_GRIPPER_MIMIC_GEARING = { + "finger_joint": 1.0, + "left_inner_knuckle_joint": 1.0, + "right_inner_knuckle_joint": 1.0, + "right_outer_knuckle_joint": 1.0, + "left_outer_finger_joint": -1.0, + "right_outer_finger_joint": -1.0, +} + + def set_finger_joint_pos_grav( joint_pos: torch.Tensor, reset_ind_joint_pos: list[int], finger_joints: list[int], finger_joint_position: float, + joint_name_to_idx: dict[str, int] | None = None, ): - """Set finger joint positions for Grav gripper. + """Set Grav gripper joint positions by joint name. Args: - joint_pos: Joint positions tensor - reset_ind_joint_pos: Row indices into the sliced joint_pos tensor - finger_joints: List of all gripper joint indices (6 joints total) - finger_joint_position: Target position for main finger joint (in radians) - - Note: - Grav gripper joint structure (indices from finger_joints list): - [0] finger_joint - main controllable joint - [1] left_inner_knuckle_joint - mimic with -1 gearing - [2] right_inner_knuckle_joint - mimic with -1 gearing - [3] right_outer_knuckle_joint - mimic with -1 gearing - [4] left_outer_finger_joint - mimic with +1 gearing - [5] right_outer_finger_joint - mimic with +1 gearing + joint_pos: Joint positions [rad]. + reset_ind_joint_pos: Row indices into ``joint_pos``. + finger_joints: Gripper joint indices retained for setter compatibility. + finger_joint_position: Main finger-joint target [rad]. + joint_name_to_idx: Mapping from joint names to simulation indices. """ + if joint_name_to_idx is None: + raise ValueError("set_finger_joint_pos_grav requires 'joint_name_to_idx'") + + missing = [name for name in _GRAV_GRIPPER_MIMIC_GEARING if name not in joint_name_to_idx] + if missing: + raise ValueError(f"Grav gripper joints not found on robot: {missing}") + for idx in reset_ind_joint_pos: - if len(finger_joints) < 6: - raise ValueError(f"Grav gripper requires at least 6 finger joints, got {len(finger_joints)}") + for joint_name, gearing in _GRAV_GRIPPER_MIMIC_GEARING.items(): + joint_pos[idx, joint_name_to_idx[joint_name]] = gearing * finger_joint_position - # Main controllable joint - joint_pos[idx, finger_joints[0]] = finger_joint_position - # Mimic joints with -1 gearing - joint_pos[idx, finger_joints[1]] = finger_joint_position # left_inner_knuckle_joint - joint_pos[idx, finger_joints[2]] = finger_joint_position # right_inner_knuckle_joint - joint_pos[idx, finger_joints[3]] = finger_joint_position # right_outer_knuckle_joint +_NEWTON_PIN_UNSELECTED_GEARS_EVENT = EventTerm( + func=gear_assembly_events.pin_unselected_gears_to_shafts, + mode="interval", + interval_range_s=(0.0, 0.0), + params={"gear_offsets": _NEWTON_GEAR_OFFSETS, "seated_gear_z_offset": 0.0075}, +) + + +def _gear_friction_range() -> PresetCfg: + return preset( + default=(0.75, 0.75), + physx=(0.75, 0.75), + newton_mjwarp=(3.0, 3.0), + newton_sdf=(3.0, 3.0), + newton_hydroelastic=(3.0, 3.0), + ) + - # Mimic joints with +1 gearing - joint_pos[idx, finger_joints[4]] = -finger_joint_position # left_outer_finger_joint - joint_pos[idx, finger_joints[5]] = -finger_joint_position # right_outer_finger_joint +def _backend_preset(physx_value: object, newton_value: object) -> PresetCfg: + """Create a PhysX-preserving preset with one value for all Newton variants.""" + return preset( + default=physx_value, + physx=physx_value, + newton_mjwarp=newton_value, + newton_sdf=newton_value, + newton_hydroelastic=newton_value, + ) ## @@ -89,8 +126,8 @@ class EventCfg: mode="startup", params={ "asset_cfg": SceneEntityCfg("factory_gear_small", body_names=".*"), - "static_friction_range": (0.75, 0.75), - "dynamic_friction_range": (0.75, 0.75), + "static_friction_range": _gear_friction_range(), + "dynamic_friction_range": _gear_friction_range(), "restitution_range": (0.0, 0.0), "num_buckets": 16, }, @@ -101,8 +138,8 @@ class EventCfg: mode="startup", params={ "asset_cfg": SceneEntityCfg("factory_gear_medium", body_names=".*"), - "static_friction_range": (0.75, 0.75), - "dynamic_friction_range": (0.75, 0.75), + "static_friction_range": _gear_friction_range(), + "dynamic_friction_range": _gear_friction_range(), "restitution_range": (0.0, 0.0), "num_buckets": 16, }, @@ -113,8 +150,8 @@ class EventCfg: mode="startup", params={ "asset_cfg": SceneEntityCfg("factory_gear_large", body_names=".*"), - "static_friction_range": (0.75, 0.75), - "dynamic_friction_range": (0.75, 0.75), + "static_friction_range": _gear_friction_range(), + "dynamic_friction_range": _gear_friction_range(), "restitution_range": (0.0, 0.0), "num_buckets": 16, }, @@ -170,6 +207,20 @@ class EventCfg: "z": [0.0575, 0.0775], }, "velocity_range": {}, + "gear_offsets": preset( + default=None, + physx=None, + newton_mjwarp=_NEWTON_GEAR_OFFSETS, + newton_sdf=_NEWTON_GEAR_OFFSETS, + newton_hydroelastic=_NEWTON_GEAR_OFFSETS, + ), + "seated_gear_z_offset": preset( + default=0.0, + physx=0.0, + newton_mjwarp=0.0075, + newton_sdf=0.0075, + newton_hydroelastic=0.0075, + ), }, ) @@ -182,22 +233,31 @@ class EventCfg: }, ) + pin_unselected_gears_to_shafts = preset( + default=None, + physx=None, + newton_mjwarp=_NEWTON_PIN_UNSELECTED_GEARS_EVENT, + newton_sdf=_NEWTON_PIN_UNSELECTED_GEARS_EVENT, + newton_hydroelastic=_NEWTON_PIN_UNSELECTED_GEARS_EVENT, + ) + @configclass class Rizon4sGearAssemblyEnvCfg(GearAssemblyEnvCfg): - """Configuration for Flexiv Rizon 4s with Grav Gripper Gear Assembly Environment. - - The Flexiv Rizon 4s is a 7-DOF collaborative robot arm equipped with the - Flexiv Grav parallel gripper for gear manipulation tasks. - """ + """Configure Flexiv Rizon 4s with the Grav gripper for gear assembly.""" ee_grasp_weight_ramp_steps: int = 512_000 def __post_init__(self): - # post init of parent super().__post_init__() + self.gear_offsets = _backend_preset(_PHYSX_GEAR_OFFSETS, _NEWTON_GEAR_OFFSETS) + + arm_joint_names = ["joint1", "joint2", "joint3", "joint4", "joint5", "joint6", "joint7"] + self.end_effector_body_name = "link7" + self.num_arm_joints = len(arm_joint_names) + self.grasp_rot_offset = [-0.707, 0.707, 0.0, 0.0] + self.gripper_joint_setter_func = set_finger_joint_pos_grav - # Flexiv-specific observation noise overrides self.observations.policy.gear_shaft_pos.noise = ResetSampledConstantNoiseModelCfg( noise_cfg=UniformNoiseCfg(n_min=-0.01, n_max=0.01, operation="add") ) @@ -206,83 +266,56 @@ def __post_init__(self): pitch_range=(-0.03491, 0.03491), yaw_range=(-0.03491, 0.03491), ) + self.observations.policy.joint_pos.params["asset_cfg"].joint_names = arm_joint_names + self.observations.policy.joint_vel.params["asset_cfg"].joint_names = arm_joint_names - # Robot-specific parameters for Flexiv Rizon 4s with Grav gripper - self.end_effector_body_name = "link7" # End effector body name for IK - self.num_arm_joints = 7 # Number of arm joints (Rizon 4s has 7 DOF) - # Rotation offset for grasp pose (quaternion [x, y, z, w]) - # Computed from IK convergence for downward-facing end effector - self.grasp_rot_offset = [ - -0.707, - 0.707, - 0.0, - 0.0, - ] - self.gripper_joint_setter_func = set_finger_joint_pos_grav # Grav gripper joint setter function - - # Gear orientation termination thresholds (in degrees) - self.gear_orientation_roll_threshold_deg = 15.0 # Maximum allowed roll deviation - self.gear_orientation_pitch_threshold_deg = 15.0 # Maximum allowed pitch deviation - self.gear_orientation_yaw_threshold_deg = 180.0 # Maximum allowed yaw deviation - - # Common observation configuration for Rizon 4s joints (arm only, not gripper) - self.observations.policy.joint_pos.params["asset_cfg"].joint_names = [ - "joint1", - "joint2", - "joint3", - "joint4", - "joint5", - "joint6", - "joint7", - ] - self.observations.policy.joint_vel.params["asset_cfg"].joint_names = [ - "joint1", - "joint2", - "joint3", - "joint4", - "joint5", - "joint6", - "joint7", - ] - - # override events self.events = EventCfg() + self.terminations.gear_orientation_exceeded.params["roll_threshold_deg"] = 15.0 + self.terminations.gear_orientation_exceeded.params["pitch_threshold_deg"] = 15.0 + self.terminations.gear_orientation_exceeded.params["yaw_threshold_deg"] = 180.0 - # Update termination thresholds from config - self.terminations.gear_orientation_exceeded.params["roll_threshold_deg"] = ( - self.gear_orientation_roll_threshold_deg - ) - self.terminations.gear_orientation_exceeded.params["pitch_threshold_deg"] = ( - self.gear_orientation_pitch_threshold_deg - ) - self.terminations.gear_orientation_exceeded.params["yaw_threshold_deg"] = ( - self.gear_orientation_yaw_threshold_deg - ) - - # Action configuration for Rizon 4s arm - # Using smaller action scale for stability self.joint_action_scale = 0.025 self.actions.arm_action = mdp.RelativeJointPositionActionCfg( asset_name="robot", - joint_names=[ - "joint1", - "joint2", - "joint3", - "joint4", - "joint5", - "joint6", - "joint7", - ], + joint_names=arm_joint_names, scale=self.joint_action_scale, use_zero_offset=True, ) - # Switch robot to Flexiv Rizon 4s with Grav gripper + physx_robot_init = ArticulationCfg.InitialStateCfg( + joint_pos={ + "joint1": 0.0, + "joint2": -0.698, + "joint3": 0.0, + "joint4": 1.571, + "joint5": 0.0, + "joint6": 0.698, + "joint7": 0.0, + }, + pos=(0.0, 0.0, 0.0), + rot=(0.0, 0.0, 0.0, 1.0), + ) + newton_robot_init = ArticulationCfg.InitialStateCfg( + joint_pos={ + "joint1": 0.050265, + "joint2": -0.372105, + "joint3": 0.111177, + "joint4": 2.276781, + "joint5": -0.083078, + "joint6": 1.074427, + "joint7": 0.230907, + }, + pos=(0.0, 0.0, 0.0), + rot=(0.0, 0.0, 0.0, 1.0), + ) self.scene.robot = FLEXIV_RIZON4S_GRAV_GRIPPER_CFG.replace( prim_path="{ENV_REGEX_NS}/Robot", spawn=FLEXIV_RIZON4S_GRAV_GRIPPER_CFG.spawn.replace( - rigid_props=sim_utils.RigidBodyPropertiesCfg( - disable_gravity=True, + # Newton currently cannot exclude only the robot from gravity. Actuator gravity + # compensation below holds the arm while preserving gravity for the gears. + joint_drive_props=_backend_preset(None, sim_utils.MujocoJointDrivePropertiesCfg(actuatorgravcomp=True)), + rigid_props=PhysxRigidBodyPropertiesCfg( + disable_gravity=_backend_preset(True, False), max_depenetration_velocity=5.0, linear_damping=0.0, angular_damping=0.0, @@ -293,29 +326,27 @@ def __post_init__(self): solver_velocity_iteration_count=1, max_contact_impulse=1e32, ), - articulation_props=sim_utils.ArticulationRootPropertiesCfg( - enabled_self_collisions=False, solver_position_iteration_count=4, solver_velocity_iteration_count=1 + articulation_props=PhysxArticulationRootPropertiesCfg( + enabled_self_collisions=False, + solver_position_iteration_count=4, + solver_velocity_iteration_count=1, ), - collision_props=sim_utils.CollisionPropertiesCfg(contact_offset=0.005, rest_offset=0.0), - ), - # Joint positions based on IK from center of distribution for randomized gear positions - init_state=ArticulationCfg.InitialStateCfg( - joint_pos={ - "joint1": 0.0, - "joint2": -0.698, - "joint3": 0.0, - "joint4": 1.571, - "joint5": 0.0, - "joint6": 0.698, - "joint7": 0.0, - }, - pos=(0.0, 0.0, 0.0), - rot=(0.0, 0.0, 0.0, 1.0), + collision_props=PhysxCollisionPropertiesCfg(contact_offset=0.005, rest_offset=0.0), ), + init_state=_backend_preset(physx_robot_init, newton_robot_init), ) - # Grav gripper actuator configuration for gear manipulation - self.scene.robot.actuators["gripper_drive"] = ImplicitActuatorCfg( + # Use the validated bare-arm gains for Newton while preserving the existing PhysX gains. + for actuator_name, physx_stiffness, physx_damping, newton_stiffness, newton_damping in ( + ("shoulder", 1320.0, 72.0, 6000.0, 108.5), + ("elbow", 600.0, 35.0, 4200.0, 90.7), + ("wrist", 216.0, 29.0, 1500.0, 54.2), + ): + actuator = self.scene.robot.actuators[actuator_name] + actuator.stiffness = _backend_preset(physx_stiffness, newton_stiffness) + actuator.damping = _backend_preset(physx_damping, newton_damping) + + physx_gripper_drive = ImplicitActuatorCfg( joint_names_expr=["finger_joint"], effort_limit_sim=2.0, velocity_limit_sim=1.0, @@ -324,9 +355,14 @@ def __post_init__(self): friction=0.0, armature=0.0, ) + newton_gripper_drive = physx_gripper_drive.replace( + effort_limit_sim=200.0, + velocity_limit_sim=2.0, + armature=0.1, + ) + self.scene.robot.actuators["gripper_drive"] = _backend_preset(physx_gripper_drive, newton_gripper_drive) - # Passive/mimic joints in the gripper - set to zero stiffness/damping - self.scene.robot.actuators["gripper_passive"] = ImplicitActuatorCfg( + physx_gripper_passive = ImplicitActuatorCfg( joint_names_expr=[".*_knuckle_joint"], effort_limit_sim=1.0, velocity_limit_sim=1.0, @@ -335,91 +371,94 @@ def __post_init__(self): friction=0.0, armature=0.0, ) - - # Override gear initial states for Rizon (closer to robot, centered) - self.scene.factory_gear_base.init_state = RigidObjectCfg.InitialStateCfg( - pos=(0.481, -0.073, 0.071), - rot=(0.0, 0.0, 0.70711, -0.70711), + newton_gripper_passive = physx_gripper_passive.replace( + joint_names_expr=[".*_knuckle_joint", ".*_outer_finger_joint"], + effort_limit_sim=20.0, + stiffness=2e3, + damping=10.0, + armature=0.05, ) - self.scene.factory_gear_small.init_state = RigidObjectCfg.InitialStateCfg( - pos=(0.481, -0.073, 0.071), - rot=(0.0, 0.0, 0.70711, -0.70711), + self.scene.robot.actuators["gripper_passive"] = _backend_preset(physx_gripper_passive, newton_gripper_passive) + + base_rot = (0.0, 0.0, 0.70711, -0.70711) + physx_asset_state = RigidObjectCfg.InitialStateCfg(pos=(0.481, -0.073, 0.071), rot=base_rot) + newton_base_pos = (0.481, -0.073, -0.005) + self.scene.factory_gear_base.init_state = _backend_preset( + physx_asset_state, + RigidObjectCfg.InitialStateCfg(pos=newton_base_pos, rot=base_rot), ) - self.scene.factory_gear_medium.init_state = RigidObjectCfg.InitialStateCfg( - pos=(0.481, -0.073, 0.071), - rot=(0.0, 0.0, 0.70711, -0.70711), + for gear_name, asset in ( + ("gear_small", self.scene.factory_gear_small), + ("gear_medium", self.scene.factory_gear_medium), + ("gear_large", self.scene.factory_gear_large), + ): + newton_gear_pos = ( + newton_base_pos[0] - _NEWTON_GEAR_OFFSETS[gear_name][0], + newton_base_pos[1], + newton_base_pos[2], + ) + asset.init_state = _backend_preset( + physx_asset_state, + RigidObjectCfg.InitialStateCfg(pos=newton_gear_pos, rot=base_rot), + ) + + physx_grasp_offsets = {name: [0.0, -offset[0], -0.35] for name, offset in _PHYSX_GEAR_OFFSETS.items()} + newton_grasp_offsets = { + "gear_small": [0.0, 0.0, -0.026], + "gear_medium": [0.0, 0.0, -0.026], + "gear_large": [0.0, 0.0, -0.026], + } + self.gear_offsets_grasp = _backend_preset(physx_grasp_offsets, newton_grasp_offsets) + self.grasp_center_body_names = _backend_preset(None, ("left_finger_tip", "right_finger_tip")) + self.hand_grasp_width = _backend_preset( + {"gear_small": 0.05, "gear_medium": 0.2, "gear_large": 0.28}, + {"gear_small": 0.01, "gear_medium": 0.2, "gear_large": 0.28}, ) - self.scene.factory_gear_large.init_state = RigidObjectCfg.InitialStateCfg( - pos=(0.481, -0.073, 0.071), - rot=(0.0, 0.0, 0.70711, -0.70711), + self.hand_close_width = _backend_preset( + {"gear_small": 0.0, "gear_medium": 0.139626, "gear_large": 0.139626}, + {"gear_small": 0.01, "gear_medium": 0.139626, "gear_large": 0.139626}, ) - - # Gear offsets and grasp positions for Rizon 4s with Grav gripper - # These offsets are relative to the end effector frame (link7) - # Z offset accounts for the gripper length from link7 to finger tip - self.gear_offsets_grasp = { - "gear_small": [0.0, -self.gear_offsets["gear_small"][0], -0.35], - "gear_medium": [0.0, -self.gear_offsets["gear_medium"][0], -0.35], - "gear_large": [0.0, -self.gear_offsets["gear_large"][0], -0.35], - } - - # Grasp widths for Grav gripper (raw radian values for finger_joint) - self.hand_grasp_width = { - "gear_small": 0.05, - "gear_medium": 0.2, - "gear_large": 0.28, - } - - # Close widths for Grav gripper (raw radian values for finger_joint) - self.hand_close_width = { - "gear_small": 0.0, - "gear_medium": 0.139626, - "gear_large": 0.139626, + self.ee_grasp_weight_ramp_start = _backend_preset(0.0, 0.2) + self.ee_grasp_weight_ramp_steps = _backend_preset(512_000, 250_000) + + grasp_event_params = self.events.set_robot_to_grasp_pose.params + grasp_event_params["gear_offsets_grasp"] = self.gear_offsets_grasp + grasp_event_params["end_effector_body_name"] = self.end_effector_body_name + grasp_event_params["num_arm_joints"] = self.num_arm_joints + grasp_event_params["grasp_rot_offset"] = self.grasp_rot_offset + grasp_event_params["gripper_joint_setter_func"] = self.gripper_joint_setter_func + grasp_event_params["grasp_center_body_names"] = self.grasp_center_body_names + + grasp_reward_params = { + "robot_asset_cfg": SceneEntityCfg("robot"), + "keypoint_scale": 0.15, + "ee_grasp_threshold": 0.0, + "weight_ramp_start": self.ee_grasp_weight_ramp_start, + "weight_ramp_steps": self.ee_grasp_weight_ramp_steps, + "end_effector_body_name": self.end_effector_body_name, + "grasp_rot_offset": self.grasp_rot_offset, + "gear_offsets_grasp": self.gear_offsets_grasp, + "grasp_center_body_names": self.grasp_center_body_names, } - - # Populate event term parameters - self.events.set_robot_to_grasp_pose.params["gear_offsets_grasp"] = self.gear_offsets_grasp - self.events.set_robot_to_grasp_pose.params["end_effector_body_name"] = self.end_effector_body_name - self.events.set_robot_to_grasp_pose.params["num_arm_joints"] = self.num_arm_joints - self.events.set_robot_to_grasp_pose.params["grasp_rot_offset"] = self.grasp_rot_offset - self.events.set_robot_to_grasp_pose.params["gripper_joint_setter_func"] = self.gripper_joint_setter_func - - # Flexiv-specific reward terms for EE-grasp keypoint tracking self.rewards.end_effector_grasp_keypoint_tracking = RewTerm( func=mdp.keypoint_ee_grasp_error, weight=-0.5, - params={ - "robot_asset_cfg": SceneEntityCfg("robot"), - "keypoint_scale": 0.15, - "ee_grasp_threshold": 0.00, - "weight_ramp_start": 0.0, - "weight_ramp_steps": self.ee_grasp_weight_ramp_steps, - "end_effector_body_name": self.end_effector_body_name, - "grasp_rot_offset": self.grasp_rot_offset, - "gear_offsets_grasp": self.gear_offsets_grasp, - }, + params=grasp_reward_params, ) self.rewards.end_effector_grasp_keypoint_tracking_exp = RewTerm( func=mdp.keypoint_ee_grasp_error_exp, weight=0.5, params={ - "robot_asset_cfg": SceneEntityCfg("robot"), + **grasp_reward_params, "kp_exp_coeffs": [(50, 0.0001), (300, 0.0001)], "kp_use_sum_of_exps": False, - "keypoint_scale": 0.15, - "ee_grasp_threshold": 0.00, - "weight_ramp_start": 0.0, - "weight_ramp_steps": self.ee_grasp_weight_ramp_steps, - "end_effector_body_name": self.end_effector_body_name, - "grasp_rot_offset": self.grasp_rot_offset, - "gear_offsets_grasp": self.gear_offsets_grasp, }, ) - # Populate termination term parameters - self.terminations.gear_dropped.params["gear_offsets_grasp"] = self.gear_offsets_grasp - self.terminations.gear_dropped.params["end_effector_body_name"] = self.end_effector_body_name - self.terminations.gear_dropped.params["grasp_rot_offset"] = self.grasp_rot_offset - + drop_params = self.terminations.gear_dropped.params + drop_params["gear_offsets_grasp"] = self.gear_offsets_grasp + drop_params["grasp_center_body_names"] = self.grasp_center_body_names + drop_params["end_effector_body_name"] = self.end_effector_body_name + drop_params["grasp_rot_offset"] = self.grasp_rot_offset self.terminations.gear_orientation_exceeded.params["end_effector_body_name"] = self.end_effector_body_name self.terminations.gear_orientation_exceeded.params["grasp_rot_offset"] = self.grasp_rot_offset diff --git a/source/isaaclab_tasks/isaaclab_tasks/contrib/deploy/gear_assembly/config/ur_10e/joint_pos_env_cfg.py b/source/isaaclab_tasks/isaaclab_tasks/contrib/deploy/gear_assembly/config/ur_10e/joint_pos_env_cfg.py index 626abb2bbccf..78729955504e 100644 --- a/source/isaaclab_tasks/isaaclab_tasks/contrib/deploy/gear_assembly/config/ur_10e/joint_pos_env_cfg.py +++ b/source/isaaclab_tasks/isaaclab_tasks/contrib/deploy/gear_assembly/config/ur_10e/joint_pos_env_cfg.py @@ -34,6 +34,7 @@ def set_finger_joint_pos_robotiq_2f140( reset_ind_joint_pos: list[int], finger_joints: list[int], finger_joint_position: float, + joint_name_to_idx: dict[str, int] | None = None, ): """Set finger joint positions for Robotiq 2F-140 gripper. @@ -42,6 +43,7 @@ def set_finger_joint_pos_robotiq_2f140( reset_ind_joint_pos: Row indices into the sliced joint_pos tensor finger_joints: List of finger joint indices finger_joint_position: Target position for finger joints + joint_name_to_idx: Optional joint-name map, unused by this index-based gripper. """ for idx in reset_ind_joint_pos: # For 2F-140 gripper (8 joints expected) @@ -69,6 +71,7 @@ def set_finger_joint_pos_robotiq_2f85( reset_ind_joint_pos: list[int], finger_joints: list[int], finger_joint_position: float, + joint_name_to_idx: dict[str, int] | None = None, ): """Set finger joint positions for Robotiq 2F-85 gripper. @@ -77,6 +80,7 @@ def set_finger_joint_pos_robotiq_2f85( reset_ind_joint_pos: Row indices into the sliced joint_pos tensor finger_joints: List of finger joint indices finger_joint_position: Target position for finger joints + joint_name_to_idx: Optional joint-name map, unused by this index-based gripper. """ for idx in reset_ind_joint_pos: # For 2F-85 gripper (6 joints expected) diff --git a/source/isaaclab_tasks/isaaclab_tasks/contrib/deploy/gear_assembly/gear_assembly_env_cfg.py b/source/isaaclab_tasks/isaaclab_tasks/contrib/deploy/gear_assembly/gear_assembly_env_cfg.py index 8ac943fdd9bf..ec86be7721e0 100644 --- a/source/isaaclab_tasks/isaaclab_tasks/contrib/deploy/gear_assembly/gear_assembly_env_cfg.py +++ b/source/isaaclab_tasks/isaaclab_tasks/contrib/deploy/gear_assembly/gear_assembly_env_cfg.py @@ -6,7 +6,15 @@ import os from dataclasses import MISSING +from isaaclab_newton.physics import ( + HydroelasticSDFCfg, + MJWarpSolverCfg, + NewtonCfg, + NewtonCollisionPipelineCfg, + NewtonShapeCfg, +) from isaaclab_physx.physics import PhysxCfg +from isaaclab_physx.sim.schemas import PhysxCollisionPropertiesCfg, PhysxRigidBodyPropertiesCfg import isaaclab.sim as sim_utils from isaaclab.assets import ArticulationCfg, AssetBaseCfg, RigidObjectCfg @@ -27,22 +35,170 @@ import isaaclab_tasks.contrib.deploy.mdp as mdp import isaaclab_tasks.contrib.deploy.mdp.terminations as gear_assembly_terminations from isaaclab_tasks.contrib.deploy.mdp.noise_models import ResetSampledConstantNoiseModelCfg +from isaaclab_tasks.utils import PresetCfg, preset # Get the directory where this configuration file is located CONFIG_DIR = os.path.dirname(os.path.abspath(__file__)) ASSETS_DIR = os.path.join(CONFIG_DIR, "assets") +NEWTON_GEAR_ASSETS_DIR = os.path.join(ASSETS_DIR, "newton") + +# A 256-world GPU shard peaks near 1.1M broad-phase pairs during randomized resets. +_GEAR_MAX_TRIANGLE_PAIRS = 1_500_000 +_PHYSX_GEAR_OFFSETS = { + "gear_small": [0.076125, 0.0, 0.0], + "gear_medium": [0.030375, 0.0, 0.0], + "gear_large": [-0.045375, 0.0, 0.0], +} +_NEWTON_GEAR_OFFSETS = { + "gear_small": [0.0823685, 0.0, 0.0], + "gear_medium": [0.0366185, 0.0, 0.0], + "gear_large": [-0.0391315, 0.0, 0.0], +} + + +def _gear_usd_path(default_usd_path: str, asset_name: str) -> PresetCfg: + """Create a gear USD path preset with Newton-specific collision assets. + + Args: + default_usd_path: Factory asset USD path used by the default and PhysX presets. + asset_name: Gear asset directory and USD stem. + + Returns: + Preset that resolves to package-local point-SDF or hydroelastic assets for Newton collision presets. + """ + return preset( + default=default_usd_path, + newton_mjwarp=os.path.join(NEWTON_GEAR_ASSETS_DIR, asset_name, f"{asset_name}.usda"), + newton_sdf=os.path.join(NEWTON_GEAR_ASSETS_DIR, asset_name, f"{asset_name}.usda"), + newton_hydroelastic=os.path.join(NEWTON_GEAR_ASSETS_DIR, asset_name, f"{asset_name}_hydroelastic.usda"), + ) + ## # Environment configuration ## +@configclass +class GearAssemblyPhysicsCfg(PresetCfg): + """Physics backend presets for gear assembly. + + Gear insertion is contact-rich (gear teeth, shaft walls, gripper fingers), so the + Newton (MuJoCo) solver limits are set conservatively. Select a preset at runtime + with the ``presets=`` CLI override: + + * ``default`` and ``physx`` -- PhysX with contact buffers sized for assembly. + * ``newton_mjwarp`` -- Newton with MuJoCo's internal contact solver. + * ``newton_sdf`` -- Newton's collision pipeline with reduced point-SDF contacts. + * ``newton_hydroelastic`` -- Newton's own collision pipeline (``use_mujoco_contacts=False``) + with SDF-based hydroelastic contacts. Produces distributed contact areas instead of + point contacts, which can improve fidelity for the gear-teeth/shaft-wall interaction. + More expensive; A/B test against ``newton_sdf`` before committing to it for training. + + Note: + ``collision_cfg`` (and therefore hydroelastic contacts) is only valid when the Newton + collision pipeline is active, i.e. ``use_mujoco_contacts=False``. Setting it alongside + ``use_mujoco_contacts=True`` raises ``ValueError``, which is why it lives in a separate + preset rather than ``newton_mjwarp``. + """ + + newton_mjwarp: NewtonCfg = NewtonCfg( + solver_cfg=MJWarpSolverCfg( + solver="newton", + integrator="implicitfast", + njmax=200, + nconmax=100, + impratio=10.0, + cone="elliptic", + iterations=100, + ls_iterations=50, + use_mujoco_contacts=True, + update_data_interval=10, + ), + num_substeps=20, + default_shape_cfg=NewtonShapeCfg(gap=0.005), + collision_decimation=0, + debug_mode=False, + ) + newton_sdf: NewtonCfg = NewtonCfg( + solver_cfg=MJWarpSolverCfg( + solver="newton", + integrator="implicitfast", + njmax=4096, + nconmax=4096, + impratio=10.0, + cone="elliptic", + iterations=100, + ls_iterations=50, + use_mujoco_contacts=False, + ccd_iterations=35, + update_data_interval=10, + ), + collision_cfg=NewtonCollisionPipelineCfg( + reduce_contacts=True, + max_triangle_pairs=_GEAR_MAX_TRIANGLE_PAIRS, + ), + num_substeps=20, + default_shape_cfg=NewtonShapeCfg(gap=0.005), + collision_decimation=0, + debug_mode=False, + ) + newton_hydroelastic: NewtonCfg = NewtonCfg( + solver_cfg=MJWarpSolverCfg( + solver="newton", + integrator="implicitfast", + # The hydroelastic SDF pipeline produces distributed contact areas (thousands of points + # for a gripped concave gear), so the per-world contact/constraint buffers must be far + # larger than the MuJoCo ``newton_mjwarp`` preset's. Sized for ~4k constraints/contacts. + njmax=4096, + nconmax=4096, + impratio=10.0, + cone="elliptic", + iterations=100, + ls_iterations=50, + # Hand collision detection to Newton's pipeline so hydroelastic SDF contacts apply. + use_mujoco_contacts=False, + ccd_iterations=35, + update_data_interval=10, + ), + collision_cfg=NewtonCollisionPipelineCfg( + max_triangle_pairs=_GEAR_MAX_TRIANGLE_PAIRS, + sdf_hydroelastic_config=HydroelasticSDFCfg( + reduce_contacts=True, + normal_matching=True, + ), + ), + num_substeps=20, + default_shape_cfg=NewtonShapeCfg(gap=0.005), + collision_decimation=0, + debug_mode=False, + ) + physx: PhysxCfg = PhysxCfg( + # Important to prevent collisionStackSize buffer overflow in contact-rich environments. + gpu_collision_stack_size=2**30, + gpu_max_rigid_contact_count=2**23, + gpu_max_rigid_patch_count=2**23, + ) + default = physx + + @configclass class GearAssemblySceneCfg(InteractiveSceneCfg): """Configuration for the scene with a robotic arm.""" - # Disable scene replication to allow USD-level randomization - replicate_physics = False + # Replicate physics so each environment gets its own physics instance. The Newton backend + # only creates per-environment bodies through the physics-replication path; with + # ``replicate_physics=False`` every environment collapses onto a single physics instance + # (root states come back shaped ``(1, ...)`` instead of ``(num_envs, ...)``). Per-environment + # gear/base variation is applied at reset via the randomization events + # (``write_root_pose_to_sim``), so it does not rely on USD-level authoring and is preserved. + replicate_physics = preset( + default=False, + physx=False, + newton_mjwarp=True, + newton_sdf=True, + newton_hydroelastic=True, + ) # world ground = AssetBaseCfg( @@ -55,9 +211,12 @@ class GearAssemblySceneCfg(InteractiveSceneCfg): prim_path="{ENV_REGEX_NS}/FactoryGearBase", # TODO: change to common isaac sim directory spawn=sim_utils.UsdFileCfg( - usd_path=f"{ISAAC_NUCLEUS_DIR}/Props/Factory/gear_assets/factory_gear_base/factory_gear_base.usd", + usd_path=_gear_usd_path( + f"{ISAAC_NUCLEUS_DIR}/Props/Factory/gear_assets/factory_gear_base/factory_gear_base.usd", + "factory_gear_base", + ), activate_contact_sensors=False, - rigid_props=sim_utils.RigidBodyPropertiesCfg( + rigid_props=PhysxRigidBodyPropertiesCfg( disable_gravity=False, kinematic_enabled=True, max_depenetration_velocity=5.0, @@ -71,7 +230,7 @@ class GearAssemblySceneCfg(InteractiveSceneCfg): max_contact_impulse=1e32, ), mass_props=sim_utils.MassPropertiesCfg(mass=None), - collision_props=sim_utils.CollisionPropertiesCfg(contact_offset=0.02, rest_offset=0.0), + collision_props=PhysxCollisionPropertiesCfg(contact_offset=0.02, rest_offset=0.0), ), init_state=RigidObjectCfg.InitialStateCfg(pos=(-1.0200, 0.2100, -0.1), rot=(0.0, 0.0, 0.70711, 0.70711)), ) @@ -80,9 +239,12 @@ class GearAssemblySceneCfg(InteractiveSceneCfg): prim_path="{ENV_REGEX_NS}/FactoryGearSmall", # TODO: change to common isaac sim directory spawn=sim_utils.UsdFileCfg( - usd_path=f"{ISAAC_NUCLEUS_DIR}/Props/Factory/gear_assets/factory_gear_small/factory_gear_small.usd", + usd_path=_gear_usd_path( + f"{ISAAC_NUCLEUS_DIR}/Props/Factory/gear_assets/factory_gear_small/factory_gear_small.usd", + "factory_gear_small", + ), activate_contact_sensors=False, - rigid_props=sim_utils.RigidBodyPropertiesCfg( + rigid_props=PhysxRigidBodyPropertiesCfg( disable_gravity=False, kinematic_enabled=False, max_depenetration_velocity=5.0, @@ -96,7 +258,7 @@ class GearAssemblySceneCfg(InteractiveSceneCfg): max_contact_impulse=1e32, ), mass_props=sim_utils.MassPropertiesCfg(mass=None), - collision_props=sim_utils.CollisionPropertiesCfg(contact_offset=0.02, rest_offset=0.0), + collision_props=PhysxCollisionPropertiesCfg(contact_offset=0.02, rest_offset=0.0), ), init_state=RigidObjectCfg.InitialStateCfg(pos=(-1.0200, 0.2100, -0.1), rot=(0.0, 0.0, 0.70711, 0.70711)), ) @@ -105,9 +267,12 @@ class GearAssemblySceneCfg(InteractiveSceneCfg): prim_path="{ENV_REGEX_NS}/FactoryGearMedium", # TODO: change to common isaac sim directory spawn=sim_utils.UsdFileCfg( - usd_path=f"{ISAAC_NUCLEUS_DIR}/Props/Factory/gear_assets/factory_gear_medium/factory_gear_medium.usd", + usd_path=_gear_usd_path( + f"{ISAAC_NUCLEUS_DIR}/Props/Factory/gear_assets/factory_gear_medium/factory_gear_medium.usd", + "factory_gear_medium", + ), activate_contact_sensors=False, - rigid_props=sim_utils.RigidBodyPropertiesCfg( + rigid_props=PhysxRigidBodyPropertiesCfg( disable_gravity=False, kinematic_enabled=False, max_depenetration_velocity=5.0, @@ -121,7 +286,7 @@ class GearAssemblySceneCfg(InteractiveSceneCfg): max_contact_impulse=1e32, ), mass_props=sim_utils.MassPropertiesCfg(mass=None), - collision_props=sim_utils.CollisionPropertiesCfg(contact_offset=0.02, rest_offset=0.0), + collision_props=PhysxCollisionPropertiesCfg(contact_offset=0.02, rest_offset=0.0), ), init_state=RigidObjectCfg.InitialStateCfg(pos=(-1.0200, 0.2100, -0.1), rot=(0.0, 0.0, 0.70711, 0.70711)), ) @@ -130,9 +295,12 @@ class GearAssemblySceneCfg(InteractiveSceneCfg): prim_path="{ENV_REGEX_NS}/FactoryGearLarge", # TODO: change to common isaac sim directory spawn=sim_utils.UsdFileCfg( - usd_path=f"{ISAAC_NUCLEUS_DIR}/Props/Factory/gear_assets/factory_gear_large/factory_gear_large.usd", + usd_path=_gear_usd_path( + f"{ISAAC_NUCLEUS_DIR}/Props/Factory/gear_assets/factory_gear_large/factory_gear_large.usd", + "factory_gear_large", + ), activate_contact_sensors=False, - rigid_props=sim_utils.RigidBodyPropertiesCfg( + rigid_props=PhysxRigidBodyPropertiesCfg( disable_gravity=False, kinematic_enabled=False, max_depenetration_velocity=5.0, @@ -146,7 +314,7 @@ class GearAssemblySceneCfg(InteractiveSceneCfg): max_contact_impulse=1e32, ), mass_props=sim_utils.MassPropertiesCfg(mass=None), - collision_props=sim_utils.CollisionPropertiesCfg(contact_offset=0.02, rest_offset=0.0), + collision_props=PhysxCollisionPropertiesCfg(contact_offset=0.02, rest_offset=0.0), ), init_state=RigidObjectCfg.InitialStateCfg(pos=(-1.0200, 0.2100, -0.1), rot=(0.0, 0.0, 0.70711, 0.70711)), ) @@ -302,11 +470,9 @@ class GearAssemblyEnvCfg(ManagerBasedRLEnvCfg): rewards: RewardsCfg = RewardsCfg() terminations: TerminationsCfg = TerminationsCfg() events: EventCfg = EventCfg() - sim: SimulationCfg = SimulationCfg( - physics=PhysxCfg( # Important to prevent collisionStackSize buffer overflow in contact-rich environments. - gpu_collision_stack_size=2**30, gpu_max_rigid_contact_count=2**23, gpu_max_rigid_patch_count=2**23 - ), - ) + # PhysX remains the backwards-compatible default; explicit Newton presets select + # package-local SDF collision assets. See :class:`GearAssemblyPhysicsCfg`. + sim: SimulationCfg = SimulationCfg(physics=GearAssemblyPhysicsCfg()) def __post_init__(self): """Post initialization.""" @@ -314,16 +480,48 @@ def __post_init__(self): self.episode_length_s = 6.66 self.viewer.eye = (3.5, 3.5, 3.5) # simulation settings - self.decimation = 4 - self.sim.render_interval = self.decimation - self.sim.dt = 1.0 / 120.0 - - self.gear_offsets = { - "gear_small": [0.076125, 0.0, 0.0], - "gear_medium": [0.030375, 0.0, 0.0], - "gear_large": [-0.045375, 0.0, 0.0], - } - - # Populate observation term parameters with gear offsets - self.observations.policy.gear_shaft_pos.params["gear_offsets"] = self.gear_offsets - self.observations.critic.gear_shaft_pos.params["gear_offsets"] = self.gear_offsets + self.decimation = preset( + default=4, + physx=4, + newton_mjwarp=3, + newton_sdf=3, + newton_hydroelastic=3, + ) + self.sim.render_interval = preset( + default=4, + physx=4, + newton_mjwarp=3, + newton_sdf=3, + newton_hydroelastic=3, + ) + self.sim.dt = preset( + default=1.0 / 120.0, + physx=1.0 / 120.0, + newton_mjwarp=0.01, + newton_sdf=0.01, + newton_hydroelastic=0.01, + ) + + physx_gear_offsets = _PHYSX_GEAR_OFFSETS + newton_gear_offsets = _NEWTON_GEAR_OFFSETS + gear_offsets = preset( + default=physx_gear_offsets, + physx=physx_gear_offsets, + newton_mjwarp=newton_gear_offsets, + newton_sdf=newton_gear_offsets, + newton_hydroelastic=newton_gear_offsets, + ) + self.gear_offsets = physx_gear_offsets + + # Populate observation and reward term parameters with backend-specific shaft offsets. + self.observations.policy.gear_shaft_pos.params["gear_offsets"] = gear_offsets + self.observations.critic.gear_shaft_pos.params["gear_offsets"] = gear_offsets + reward_gear_offsets = preset( + default=None, + physx=None, + newton_mjwarp=newton_gear_offsets, + newton_sdf=newton_gear_offsets, + newton_hydroelastic=newton_gear_offsets, + ) + self.rewards.end_effector_gear_keypoint_tracking.params["gear_offsets"] = reward_gear_offsets + self.rewards.end_effector_gear_keypoint_tracking_exp.params["gear_offsets"] = reward_gear_offsets diff --git a/source/isaaclab_tasks/isaaclab_tasks/contrib/deploy/mdp/__init__.pyi b/source/isaaclab_tasks/isaaclab_tasks/contrib/deploy/mdp/__init__.pyi index 2a200c888bc0..6178634ef3cd 100644 --- a/source/isaaclab_tasks/isaaclab_tasks/contrib/deploy/mdp/__init__.pyi +++ b/source/isaaclab_tasks/isaaclab_tasks/contrib/deploy/mdp/__init__.pyi @@ -6,6 +6,7 @@ __all__ = [ "randomize_gear_type", "randomize_gears_and_base_pose", + "pin_unselected_gears_to_shafts", "set_robot_to_grasp_pose", "ResetSampledConstantNoiseModel", "ResetSampledConstantNoiseModelCfg", @@ -23,7 +24,12 @@ __all__ = [ "reset_when_gear_orientation_exceeds_threshold", ] -from .events import randomize_gear_type, randomize_gears_and_base_pose, set_robot_to_grasp_pose +from .events import ( + pin_unselected_gears_to_shafts, + randomize_gear_type, + randomize_gears_and_base_pose, + set_robot_to_grasp_pose, +) from .noise_models import ResetSampledConstantNoiseModel, ResetSampledConstantNoiseModelCfg from .observations import gear_pos_w, gear_quat_w, gear_shaft_pos_w, gear_shaft_quat_w from .rewards import ( diff --git a/source/isaaclab_tasks/isaaclab_tasks/contrib/deploy/mdp/events.py b/source/isaaclab_tasks/isaaclab_tasks/contrib/deploy/mdp/events.py index 8dde6b5c5105..67328a41053b 100644 --- a/source/isaaclab_tasks/isaaclab_tasks/contrib/deploy/mdp/events.py +++ b/source/isaaclab_tasks/isaaclab_tasks/contrib/deploy/mdp/events.py @@ -8,6 +8,7 @@ from __future__ import annotations import random +from collections.abc import Callable, Sequence from typing import TYPE_CHECKING import torch @@ -74,8 +75,7 @@ def __call__( env_ids: Environment IDs to randomize gear_types: List of available gear types to choose from """ - # Randomly select gear type for each environment - # Use the parameter passed to __call__ (not self.gear_types) to allow runtime overrides + # Randomly select gear type for each environment. for env_id in env_ids.tolist(): chosen_gear = random.choice(gear_types) self._current_gear_type[env_id] = chosen_gear @@ -203,10 +203,26 @@ def __init__(self, cfg: EventTermCfg, env: ManagerBasedEnv): # Find jacobian body index (for fixed-base robots, subtract 1) self.jacobi_body_idx = self.eef_idx - 1 + # Optionally derive the end-effector -> grasp-center transform from named bodies. + self.grasp_center_body_indices: list[int] | None = None + grasp_center_body_names = cfg.params.get("grasp_center_body_names") + if grasp_center_body_names is not None: + grasp_center_body_names = list(grasp_center_body_names) + grasp_center_body_indices, _ = self.robot_asset.find_bodies(grasp_center_body_names) + if len(grasp_center_body_indices) != len(grasp_center_body_names): + raise ValueError(f"Grasp center bodies not found on robot: {grasp_center_body_names}") + self.grasp_center_body_indices = grasp_center_body_indices + # The existing PhysX task leaves this disabled; Newton targets the fingertip midpoint. + self._grasp_center_offset_in_eef = None + # Find all joints once all_joints, all_joints_names = self.robot_asset.find_joints([".*"]) self.all_joints = all_joints self.finger_joints = all_joints[self.num_arm_joints :] + # Map joint name -> simulation joint index. The order returned by ``find_joints`` is the + # backend DOF order, which differs between PhysX and Newton; gripper joint setters must + # therefore resolve joints by name (not by positional slice) to remain backend-agnostic. + self.joint_name_to_idx = {name: idx for idx, name in zip(all_joints, all_joints_names)} def __call__( self, @@ -221,7 +237,8 @@ def __call__( end_effector_body_name: str | None = None, num_arm_joints: int | None = None, grasp_rot_offset: list | None = None, - gripper_joint_setter_func: callable | None = None, + gripper_joint_setter_func: Callable[..., None] | None = None, + grasp_center_body_names: tuple[str, ...] | None = None, ): """Set robot to grasp pose using IK. @@ -250,6 +267,17 @@ def __call__( gear_grasp_offsets = self.gear_grasp_offsets_buffer[:num_reset_envs] grasp_rot_offset_tensor = self.grasp_rot_offset_tensor[env_ids] + # Measure the rigid end-effector -> grasp-center vector in the end-effector frame. + if self.grasp_center_body_indices is not None: + eef_pos_now = self.robot_asset.data.body_pos_w.torch[env_ids, self.eef_idx] + eef_quat_now = self.robot_asset.data.body_quat_w.torch[env_ids, self.eef_idx] + grasp_center_now = self.robot_asset.data.body_pos_w.torch[ + env_ids[:, None], self.grasp_center_body_indices + ].mean(dim=1) + self._grasp_center_offset_in_eef = math_utils.quat_apply( + math_utils.quat_conjugate(eef_quat_now), grasp_center_now - eef_pos_now + ) + # IK loop for i in range(max_iterations): # Get current joint state @@ -289,6 +317,11 @@ def __call__( # Get grasp offsets (vectorized) gear_grasp_offsets[:] = self.gear_grasp_offsets_stacked[gear_type_indices] + if self._grasp_center_offset_in_eef is not None: + # Target the configured grasp-center bodies instead of the end-effector link origin. + gear_grasp_offsets = gear_grasp_offsets.clone() + gear_grasp_offsets = gear_grasp_offsets - self._grasp_center_offset_in_eef + # Add position randomization if specified if pos_randomization_range is not None: pos_keys = ["x", "y", "z"] @@ -376,7 +409,9 @@ def __call__( for row_idx, env_id in enumerate(env_ids.tolist()): gear_key = all_gear_types[env_id] hand_grasp_width = self.hand_grasp_width[gear_key] - self.gripper_joint_setter_func(joint_pos, [row_idx], self.finger_joints, hand_grasp_width) + self.gripper_joint_setter_func( + joint_pos, [row_idx], self.finger_joints, hand_grasp_width, self.joint_name_to_idx + ) self.robot_asset.set_joint_position_target_index(target=joint_pos, joint_ids=self.all_joints, env_ids=env_ids) self.robot_asset.write_joint_position_to_sim_index(position=joint_pos, env_ids=env_ids) @@ -386,7 +421,9 @@ def __call__( for row_idx, env_id in enumerate(env_ids.tolist()): gear_key = all_gear_types[env_id] hand_close_width = self.hand_close_width[gear_key] - self.gripper_joint_setter_func(joint_pos, [row_idx], self.finger_joints, hand_close_width) + self.gripper_joint_setter_func( + joint_pos, [row_idx], self.finger_joints, hand_close_width, self.joint_name_to_idx + ) self.robot_asset.set_joint_position_target_index(target=joint_pos, joint_ids=self.all_joints, env_ids=env_ids) @@ -414,6 +451,71 @@ def __init__(self, cfg: EventTermCfg, env: ManagerBasedEnv): self.gear_asset_names = ["factory_gear_small", "factory_gear_medium", "factory_gear_large"] self.base_asset_name = "factory_gear_base" + def _randomize_legacy( + self, + env: ManagerBasedEnv, + env_ids: torch.Tensor, + pose_range: dict, + velocity_range: dict, + gear_pos_range: dict, + ) -> None: + """Apply the existing PhysX reset behavior when shaft offsets are not configured.""" + if not hasattr(env, "_gear_type_manager"): + raise RuntimeError( + "Gear type manager not initialized. Ensure randomize_gear_type event is configured " + "in your environment's event configuration before this event term is used." + ) + + gear_type_manager: randomize_gear_type = env._gear_type_manager + device = env.device + pose_keys = ["x", "y", "z", "roll", "pitch", "yaw"] + range_list_pose = [pose_range.get(key, (0.0, 0.0)) for key in pose_keys] + ranges_pose = torch.tensor(range_list_pose, device=device) + rand_pose_samples = math_utils.sample_uniform( + ranges_pose[:, 0], ranges_pose[:, 1], (len(env_ids), 6), device=device + ) + orientations_delta = math_utils.quat_from_euler_xyz( + rand_pose_samples[:, 3], rand_pose_samples[:, 4], rand_pose_samples[:, 5] + ) + + range_list_vel = [velocity_range.get(key, (0.0, 0.0)) for key in pose_keys] + ranges_vel = torch.tensor(range_list_vel, device=device) + rand_vel_samples = math_utils.sample_uniform( + ranges_vel[:, 0], ranges_vel[:, 1], (len(env_ids), 6), device=device + ) + + positions_by_asset = {} + orientations_by_asset = {} + velocities_by_asset = {} + for asset_name in [self.base_asset_name] + self.gear_asset_names: + asset: RigidObject | Articulation = env.scene[asset_name] + default_root_pose = asset.data.default_root_pose.torch[env_ids].clone() + default_root_vel = asset.data.default_root_vel.torch[env_ids].clone() + positions_by_asset[asset_name] = ( + default_root_pose[:, 0:3] + env.scene.env_origins[env_ids] + rand_pose_samples[:, 0:3] + ) + orientations_by_asset[asset_name] = math_utils.quat_mul(default_root_pose[:, 3:7], orientations_delta) + velocities_by_asset[asset_name] = default_root_vel + rand_vel_samples + + range_list_gear = [gear_pos_range.get(key, (0.0, 0.0)) for key in ["x", "y", "z"]] + ranges_gear = torch.tensor(range_list_gear, device=device) + rand_gear_offsets = math_utils.sample_uniform( + ranges_gear[:, 0], ranges_gear[:, 1], (len(env_ids), 3), device=device + ) + gear_type_indices = self.gear_type_indices[: len(env_ids)] + gear_type_indices[:] = gear_type_manager.get_all_gear_type_indices()[env_ids] + + for gear_idx, asset_name in enumerate(self.gear_asset_names): + mask = gear_type_indices == gear_idx + positions_by_asset[asset_name][mask] += rand_gear_offsets[mask] + + for asset_name, positions in positions_by_asset.items(): + asset = env.scene[asset_name] + asset.write_root_pose_to_sim_index( + root_pose=torch.cat([positions, orientations_by_asset[asset_name]], dim=-1), env_ids=env_ids + ) + asset.write_root_velocity_to_sim_index(root_velocity=velocities_by_asset[asset_name], env_ids=env_ids) + def __call__( self, env: ManagerBasedEnv, @@ -421,16 +523,30 @@ def __call__( pose_range: dict = {}, velocity_range: dict = {}, gear_pos_range: dict = {}, + gear_offsets: dict | None = None, + seated_gear_z_offset: float = 0.0, ): """Randomize gear base and gear poses. + Every gear is seated on its shaft (at the shaft target the gear-shaft observation and reward + use), expressed relative to the randomized base pose. The single gear selected for the trial + is then lifted off its shaft by ``gear_pos_range`` so the gripper can grasp it for insertion. + Args: env: Environment instance env_ids: Environment IDs to randomize - pose_range: Pose randomization range for base and all gears - velocity_range: Velocity randomization range - gear_pos_range: Additional position randomization for selected gear only + pose_range: Pose randomization range for the base (and, rigidly, all gears) + velocity_range: Unused; gears and base are reset at rest + gear_pos_range: Lift applied to the selected gear only, in world frame + gear_offsets: Per-gear shaft offset in the base frame, mapping ``gear_small`` / + ``gear_medium`` / ``gear_large`` to ``[x, y, z]`` (the insertion target) + seated_gear_z_offset: Rest-height offset [m] applied in the base frame for the two + non-selected gears, so they remain centered on their shafts under base rotation. """ + if gear_offsets is None: + self._randomize_legacy(env, env_ids, pose_range, velocity_range, gear_pos_range) + return + if not hasattr(env, "_gear_type_manager"): raise RuntimeError( "Gear type manager not initialized. Ensure randomize_gear_type event is configured " @@ -452,54 +568,145 @@ def __call__( rand_pose_samples[:, 3], rand_pose_samples[:, 4], rand_pose_samples[:, 5] ) - # Shared velocity samples - range_list_vel = [velocity_range.get(key, (0.0, 0.0)) for key in pose_keys] - ranges_vel = torch.tensor(range_list_vel, device=device) - rand_vel_samples = math_utils.sample_uniform( - ranges_vel[:, 0], ranges_vel[:, 1], (len(env_ids), 6), device=device + # Per-gear shaft offset (base frame), ordered to match self.gear_asset_names. + key_map = { + "factory_gear_small": "gear_small", + "factory_gear_medium": "gear_medium", + "factory_gear_large": "gear_large", + } + gear_offsets_t = torch.tensor( + [gear_offsets[key_map[name]] for name in self.gear_asset_names], device=device, dtype=torch.float32 ) - # Prepare poses for all assets - positions_by_asset = {} - orientations_by_asset = {} - velocities_by_asset = {} + # Perturbed base world pose. Each gear is seated at its shaft target (combined with the base + # pose) so it stays on its shaft under any base rotation; only the gear selected for this + # trial is lifted off. Placing gears at the linear spawn position instead drifts them off the + # shafts because the base orientation is not a pure z-rotation. + base: RigidObject | Articulation = env.scene[self.base_asset_name] + base_default = base.data.default_root_pose.torch[env_ids].clone() + base_world_pos = base_default[:, 0:3] + env.scene.env_origins[env_ids] + rand_pose_samples[:, 0:3] + base_world_quat = math_utils.quat_mul(base_default[:, 3:7], orientations_delta) + zero_vel = torch.zeros((len(env_ids), 6), device=device) + base.write_root_pose_to_sim_index( + root_pose=torch.cat([base_world_pos, base_world_quat], dim=-1), env_ids=env_ids + ) + base.write_root_velocity_to_sim_index(root_velocity=zero_vel, env_ids=env_ids) - asset_names_to_process = [self.base_asset_name] + self.gear_asset_names - for asset_name in asset_names_to_process: - asset: RigidObject | Articulation = env.scene[asset_name] - default_root_pose = asset.data.default_root_pose.torch[env_ids].clone() - default_root_vel = asset.data.default_root_vel.torch[env_ids].clone() - positions = default_root_pose[:, 0:3] + env.scene.env_origins[env_ids] + rand_pose_samples[:, 0:3] - orientations = math_utils.quat_mul(default_root_pose[:, 3:7], orientations_delta) - velocities = default_root_vel + rand_vel_samples - positions_by_asset[asset_name] = positions - orientations_by_asset[asset_name] = orientations - velocities_by_asset[asset_name] = velocities - - # Per-env gear offset (gear_pos_range) applied only to selected gear + # Per-env selected gear index. + num_reset_envs = len(env_ids) + gear_type_indices = self.gear_type_indices[:num_reset_envs] + gear_type_indices[:] = gear_type_manager.get_all_gear_type_indices()[env_ids] + + # Per-env lift (gear_pos_range), applied in world frame to the selected gear only. range_list_gear = [gear_pos_range.get(key, (0.0, 0.0)) for key in ["x", "y", "z"]] ranges_gear = torch.tensor(range_list_gear, device=device) rand_gear_offsets = math_utils.sample_uniform( ranges_gear[:, 0], ranges_gear[:, 1], (len(env_ids), 3), device=device ) - # Get gear type indices directly as tensor - num_reset_envs = len(env_ids) - gear_type_indices = self.gear_type_indices[:num_reset_envs] - all_gear_type_indices = gear_type_manager.get_all_gear_type_indices() - gear_type_indices[:] = all_gear_type_indices[env_ids] - - # Apply offsets using vectorized operations with masks for gear_idx, asset_name in enumerate(self.gear_asset_names): - if asset_name in positions_by_asset: - mask = gear_type_indices == gear_idx - positions_by_asset[asset_name][mask] = positions_by_asset[asset_name][mask] + rand_gear_offsets[mask] + gear: RigidObject = env.scene[asset_name] + gear_default = gear.data.default_root_pose.torch[env_ids].clone() + # Orientation of the gear relative to the base in its default (seated) configuration. + _, rel_quat = math_utils.subtract_frame_transforms( + base_default[:, 0:3], base_default[:, 3:7], gear_default[:, 0:3], gear_default[:, 3:7] + ) + # Seat non-selected gears at the rest height in the base frame. Applying this offset before + # the base transform keeps each gear centered on its shaft even when base roll/pitch are + # randomized. + off = gear_offsets_t[gear_idx].unsqueeze(0).expand(len(env_ids), 3) + seated_off = off.clone() + seated_off[:, 2] += seated_gear_z_offset + gear_world_pos, gear_world_quat = math_utils.combine_frame_transforms( + base_world_pos, base_world_quat, seated_off, rel_quat + ) - # Write to sim - for asset_name in positions_by_asset.keys(): - asset = env.scene[asset_name] - positions = positions_by_asset[asset_name] - orientations = orientations_by_asset[asset_name] - velocities = velocities_by_asset[asset_name] - asset.write_root_pose_to_sim_index(root_pose=torch.cat([positions, orientations], dim=-1), env_ids=env_ids) - asset.write_root_velocity_to_sim_index(root_velocity=velocities, env_ids=env_ids) + # Lift only the selected gear from the insertion target so the gripper can grasp it. + mask = gear_type_indices == gear_idx + if torch.any(mask): + grasp_world_pos, grasp_world_quat = math_utils.combine_frame_transforms( + base_world_pos, base_world_quat, off, rel_quat + ) + gear_world_pos[mask] = grasp_world_pos[mask] + rand_gear_offsets[mask] + gear_world_quat[mask] = grasp_world_quat[mask] + gear.write_root_pose_to_sim_index( + root_pose=torch.cat([gear_world_pos, gear_world_quat], dim=-1), env_ids=env_ids + ) + gear.write_root_velocity_to_sim_index(root_velocity=zero_vel, env_ids=env_ids) + + +def pin_unselected_gears_to_shafts( + env: ManagerBasedEnv, + env_ids: torch.Tensor | Sequence[int] | slice | None, + gear_offsets: dict[str, Sequence[float]], + seated_gear_z_offset: float = 0.0, +) -> None: + """Keep non-selected gears seated on their base shafts. + + The gear assembly task manipulates one selected gear per environment; the other gears are + scene context and shaft obstacles. Under Newton hydroelastic contacts those free, light gears + can settle into small SDF penetrations and tip off their shafts. This event rewrites only the + non-selected gears to their shaft rest poses after each environment step. + + Args: + env: Environment instance. + env_ids: Environment IDs to update. If ``None``, all environments are updated. + gear_offsets: Per-gear shaft offsets in the base frame [m]. + seated_gear_z_offset: Rest-height offset in the base frame [m]. + """ + if not hasattr(env, "_gear_type_manager"): + raise RuntimeError( + "Gear type manager not initialized. Ensure randomize_gear_type event is configured " + "before pin_unselected_gears_to_shafts is used." + ) + + device = env.device + if env_ids is None: + env_ids = torch.arange(env.num_envs, device=device, dtype=torch.long) + elif isinstance(env_ids, slice): + env_ids = torch.arange(env.num_envs, device=device, dtype=torch.long)[env_ids] + elif not isinstance(env_ids, torch.Tensor): + env_ids = torch.tensor(env_ids, device=device, dtype=torch.long) + else: + env_ids = env_ids.to(device=device, dtype=torch.long) + + if env_ids.numel() == 0: + return + + gear_asset_names = ["factory_gear_small", "factory_gear_medium", "factory_gear_large"] + key_map = { + "factory_gear_small": "gear_small", + "factory_gear_medium": "gear_medium", + "factory_gear_large": "gear_large", + } + gear_offsets_t = torch.tensor( + [gear_offsets[key_map[name]] for name in gear_asset_names], device=device, dtype=torch.float32 + ) + + base: RigidObject | Articulation = env.scene["factory_gear_base"] + base_world_pos = base.data.root_link_pos_w.torch[env_ids] + base_world_quat = base.data.root_link_quat_w.torch[env_ids] + base_default = base.data.default_root_pose.torch[env_ids] + selected_gear_indices = env._gear_type_manager.get_all_gear_type_indices()[env_ids] + zero_vel = torch.zeros((len(env_ids), 6), device=device) + + for gear_idx, asset_name in enumerate(gear_asset_names): + mask = selected_gear_indices != gear_idx + if not torch.any(mask): + continue + + gear: RigidObject = env.scene[asset_name] + gear_default = gear.data.default_root_pose.torch[env_ids] + _, rel_quat = math_utils.subtract_frame_transforms( + base_default[:, 0:3], base_default[:, 3:7], gear_default[:, 0:3], gear_default[:, 3:7] + ) + seated_off = gear_offsets_t[gear_idx].unsqueeze(0).expand(len(env_ids), 3).clone() + seated_off[:, 2] += seated_gear_z_offset + gear_world_pos, gear_world_quat = math_utils.combine_frame_transforms( + base_world_pos, base_world_quat, seated_off, rel_quat + ) + gear_env_ids = env_ids[mask] + gear.write_root_pose_to_sim_index( + root_pose=torch.cat([gear_world_pos[mask], gear_world_quat[mask]], dim=-1), env_ids=gear_env_ids + ) + gear.write_root_velocity_to_sim_index(root_velocity=zero_vel[mask], env_ids=gear_env_ids) diff --git a/source/isaaclab_tasks/isaaclab_tasks/contrib/deploy/mdp/rewards.py b/source/isaaclab_tasks/isaaclab_tasks/contrib/deploy/mdp/rewards.py index c776168e5b48..26443b54f5d3 100644 --- a/source/isaaclab_tasks/isaaclab_tasks/contrib/deploy/mdp/rewards.py +++ b/source/isaaclab_tasks/isaaclab_tasks/contrib/deploy/mdp/rewards.py @@ -175,9 +175,12 @@ def __call__( class keypoint_entity_error(ManagerTermBase): - """Compute keypoint distance between a RigidObject and the dynamically selected gear. + """Compute keypoint distance between a RigidObject target and the selected gear. - This class-based term pre-caches gear type mapping and asset references. + This class-based term pre-caches gear type mapping and asset references. If ``gear_offsets`` + is provided, the target pose is shifted from ``asset_cfg_1`` by the selected gear's offset in + that asset's frame. This is used by gear assembly to reward insertion at the selected shaft + instead of at the gear-base origin. """ def __init__(self, cfg: RewardTermCfg, env: ManagerBasedRLEnv): @@ -194,6 +197,7 @@ def __init__(self, cfg: RewardTermCfg, env: ManagerBasedRLEnv): self.asset_1 = env.scene[self.asset_cfg_1.name] self._init_gear_selection(env) + self._init_target_offsets(cfg, env) # Create keypoint distance computer self.keypoint_computer = _compute_keypoint_distance(cfg, env) @@ -210,6 +214,46 @@ def _init_gear_selection(self, env: ManagerBasedRLEnv) -> None: "gear_large": env.scene["factory_gear_large"], } + def _init_target_offsets(self, cfg: RewardTermCfg, env: ManagerBasedRLEnv) -> None: + """Pre-cache optional selected-gear target offsets.""" + gear_offsets = cfg.params.get("gear_offsets") + self.gear_offsets_stacked = None + if gear_offsets is None: + return + if not isinstance(gear_offsets, dict): + raise TypeError( + f"'gear_offsets' parameter must be a dict, got {type(gear_offsets).__name__}. " + "It should have keys 'gear_small', 'gear_medium', 'gear_large' mapping to [x, y, z] offsets." + ) + + offset_tensors = [] + for gear_type in ["gear_small", "gear_medium", "gear_large"]: + if gear_type not in gear_offsets: + raise ValueError( + f"'{gear_type}' offset is required in 'gear_offsets' parameter. " + f"Found keys: {list(gear_offsets.keys())}" + ) + offset_tensors.append(torch.tensor(gear_offsets[gear_type], device=env.device, dtype=torch.float32)) + + self.gear_offsets_stacked = torch.stack(offset_tensors, dim=0) + self.identity_quat = ( + torch.tensor([[0.0, 0.0, 0.0, 1.0]], device=env.device, dtype=torch.float32) + .repeat(env.num_envs, 1) + .contiguous() + ) + + def _get_target_asset_pose(self) -> tuple[torch.Tensor, torch.Tensor]: + """Retrieve the target asset pose, optionally shifted to the selected gear shaft.""" + target_pos = self.asset_1.data.body_pos_w.torch[:, 0] + target_quat = self.asset_1.data.body_quat_w.torch[:, 0] + + if self.gear_offsets_stacked is None: + return target_pos, target_quat + + target_offsets = self.gear_offsets_stacked[self.gear_type_indices] + target_pos, _ = combine_frame_transforms(target_pos, target_quat, target_offsets, self.identity_quat) + return target_pos, target_quat + def _get_selected_gear_poses(self, env: ManagerBasedRLEnv) -> tuple[torch.Tensor, torch.Tensor]: """Retrieve world-frame position and quaternion of the active gear per environment. @@ -254,6 +298,7 @@ def __call__( asset_cfg_1: SceneEntityCfg, keypoint_scale: float = 1.0, add_cube_center_kp: bool = True, + gear_offsets: dict | None = None, ) -> torch.Tensor: """Compute keypoint distance error. @@ -262,23 +307,23 @@ def __call__( asset_cfg_1: Configuration of the first asset (RigidObject) keypoint_scale: Scale factor for keypoint offsets add_cube_center_kp: Whether to include center keypoint + gear_offsets: Optional selected-gear target offsets [m], cached during initialization. Returns: Mean keypoint distance tensor of shape (num_envs,) """ - # Get current pose of asset_1 (RigidObject) - curr_pos_1 = self.asset_1.data.body_pos_w.torch[:, 0] - curr_quat_1 = self.asset_1.data.body_quat_w.torch[:, 0] - # Get selected gear pose - curr_pos_2, curr_quat_2 = self._get_selected_gear_poses(env) + gear_pos, gear_quat = self._get_selected_gear_poses(env) + + # Get target pose of asset_1 (RigidObject), optionally shifted to selected shaft + target_pos, target_quat = self._get_target_asset_pose() # Compute keypoint distance keypoint_dist_sep = self.keypoint_computer.compute( - current_pos=curr_pos_1, - current_quat=curr_quat_1, - target_pos=curr_pos_2, - target_quat=curr_quat_2, + current_pos=gear_pos, + current_quat=gear_quat, + target_pos=target_pos, + target_quat=target_quat, keypoint_scale=keypoint_scale, ) @@ -300,6 +345,7 @@ def __call__( kp_use_sum_of_exps: bool = True, keypoint_scale: float = 1.0, add_cube_center_kp: bool = True, + gear_offsets: dict | None = None, ) -> torch.Tensor: """Compute exponential keypoint reward. @@ -310,23 +356,23 @@ def __call__( kp_use_sum_of_exps: Whether to use sum of exponentials keypoint_scale: Scale factor for keypoint offsets add_cube_center_kp: Whether to include center keypoint + gear_offsets: Optional selected-gear target offsets [m], cached during initialization. Returns: Exponential keypoint reward tensor of shape (num_envs,) """ - # Get current pose of asset_1 (RigidObject) - curr_pos_1 = self.asset_1.data.body_pos_w.torch[:, 0] - curr_quat_1 = self.asset_1.data.body_quat_w.torch[:, 0] - # Get selected gear pose - curr_pos_2, curr_quat_2 = self._get_selected_gear_poses(env) + gear_pos, gear_quat = self._get_selected_gear_poses(env) + + # Get target pose of asset_1 (RigidObject), optionally shifted to selected shaft + target_pos, target_quat = self._get_target_asset_pose() # Compute keypoint distance keypoint_dist_sep = self.keypoint_computer.compute( - current_pos=curr_pos_1, - current_quat=curr_quat_1, - target_pos=curr_pos_2, - target_quat=curr_quat_2, + current_pos=gear_pos, + current_quat=gear_quat, + target_pos=target_pos, + target_quat=target_quat, keypoint_scale=keypoint_scale, ) @@ -395,7 +441,15 @@ def __init__(self, cfg: RewardTermCfg, env: ManagerBasedRLEnv): eef_indices, _ = self.robot_asset.find_bodies([self.end_effector_body_name]) self.eef_idx = eef_indices[0] if len(eef_indices) > 0 else None - self._step_count = 0 + + self.grasp_center_body_indices: list[int] | None = None + grasp_center_body_names = cfg.params.get("grasp_center_body_names") + if grasp_center_body_names is not None: + grasp_center_body_names = list(grasp_center_body_names) + grasp_center_body_indices, _ = self.robot_asset.find_bodies(grasp_center_body_names) + if len(grasp_center_body_indices) != len(grasp_center_body_names): + raise ValueError(f"Grasp center bodies not found on robot: {grasp_center_body_names}") + self.grasp_center_body_indices = grasp_center_body_indices def _get_weight_scale(self, env: ManagerBasedRLEnv) -> float: progress = min(env.common_step_counter / max(self.weight_ramp_steps, 1), 1.0) @@ -409,7 +463,11 @@ def _get_grasp_corrected_target( Returns: Tuple of (eef_pos, eef_quat, gear_grasp_pos, gear_quat_grasp). """ - eef_pos = self.robot_asset.data.body_link_pos_w.torch[:, self.eef_idx] + if self.grasp_center_body_indices is None: + eef_pos = self.robot_asset.data.body_link_pos_w.torch[:, self.eef_idx] + else: + grasp_center_pos = self.robot_asset.data.body_link_pos_w.torch[:, self.grasp_center_body_indices] + eef_pos = grasp_center_pos.mean(dim=1) eef_quat = self.robot_asset.data.body_link_quat_w.torch[:, self.eef_idx] gear_pos, gear_quat = self._get_selected_gear_poses(env) @@ -432,6 +490,7 @@ def __call__( weight_ramp_start: float = 0.0, weight_ramp_steps: int = 1, ee_grasp_threshold: float = 0.0, + grasp_center_body_names: tuple[str, ...] | None = None, ) -> torch.Tensor: if self.eef_idx is None: return torch.zeros(env.num_envs, device=env.device) @@ -464,16 +523,6 @@ def __call__( env.extras["log"]["ee_grasp_kp_error/pct_envs_active"] = pct_active env.extras["log"]["ee_grasp_kp_error/weight_scale"] = weight_scale - self._step_count += 1 - import carb - - carb.log_info( - f"[ee_grasp_kp_error] step={self._step_count}" - f" | mean_kp_error={mean_error_scalar:.5f}" - f" | pct_active={pct_active:.3f}" - f" | weight_scale={weight_scale:.4f}" - ) - return scaled_reward @@ -507,6 +556,7 @@ def __call__( weight_ramp_start: float = 0.0, weight_ramp_steps: int = 1, ee_grasp_threshold: float = 0.0, + grasp_center_body_names: tuple[str, ...] | None = None, ) -> torch.Tensor: if self.eef_idx is None: return torch.zeros(env.num_envs, device=env.device) @@ -554,17 +604,6 @@ def __call__( env.extras["log"]["ee_grasp_kp_error_exp/pct_envs_active"] = pct_active env.extras["log"]["ee_grasp_kp_error_exp/weight_scale"] = weight_scale - self._step_count += 1 - import carb - - carb.log_info( - f"[ee_grasp_kp_error_exp] step={self._step_count}" - f" | mean_kp_error={mean_error_scalar:.5f}" - f" | pct_active={pct_active:.3f}" - f" | weight_scale={weight_scale:.4f}" - f" | mean_exp_reward={mean_reward_scalar:.5f}" - ) - return scaled_reward diff --git a/source/isaaclab_tasks/isaaclab_tasks/contrib/deploy/mdp/terminations.py b/source/isaaclab_tasks/isaaclab_tasks/contrib/deploy/mdp/terminations.py index 9f217c316372..78b187e2e6dc 100644 --- a/source/isaaclab_tasks/isaaclab_tasks/contrib/deploy/mdp/terminations.py +++ b/source/isaaclab_tasks/isaaclab_tasks/contrib/deploy/mdp/terminations.py @@ -124,6 +124,15 @@ def __init__(self, cfg: TerminationTermCfg, env: ManagerBasedEnv): else: self.eef_idx = eef_indices[0] + self.grasp_center_body_indices: list[int] | None = None + grasp_center_body_names = cfg.params.get("grasp_center_body_names") + if grasp_center_body_names is not None: + grasp_center_body_names = list(grasp_center_body_names) + grasp_center_body_indices, _ = self.robot_asset.find_bodies(grasp_center_body_names) + if len(grasp_center_body_indices) != len(grasp_center_body_names): + raise ValueError(f"Grasp center bodies not found on robot: {grasp_center_body_names}") + self.grasp_center_body_indices = grasp_center_body_indices + def __call__( self, env: ManagerBasedEnv, @@ -132,6 +141,7 @@ def __call__( gear_offsets_grasp: dict | None = None, end_effector_body_name: str | None = None, grasp_rot_offset: list | None = None, + grasp_center_body_names: tuple[str, ...] | None = None, ) -> torch.Tensor: """Check if gear has dropped and return reset flags. @@ -160,8 +170,14 @@ def __call__( # Get gear type indices directly as tensor (no Python loops!) self.gear_type_indices = gear_type_manager.get_all_gear_type_indices() - # Get end effector position - eef_pos_world = self.robot_asset.data.body_link_pos_w.torch[:, self.eef_idx] + # Get grasp-center position. For grippers with explicit fingertip bodies, use the + # fingertip midpoint instead of a distant wrist/link frame so small gear rotations do not + # amplify into false drop resets. + if self.grasp_center_body_indices is None: + eef_pos_world = self.robot_asset.data.body_link_pos_w.torch[:, self.eef_idx] + else: + grasp_center_pos = self.robot_asset.data.body_link_pos_w.torch[:, self.grasp_center_body_indices] + eef_pos_world = grasp_center_pos.mean(dim=1) # Update gear positions and quaternions in buffers self.all_gear_pos_buffer[:, 0, :] = self.gear_assets["gear_small"].data.root_link_pos_w.torch diff --git a/source/isaaclab_tasks/test/contrib/test_deploy_gear_assembly_cfg.py b/source/isaaclab_tasks/test/contrib/test_deploy_gear_assembly_cfg.py index b86487c18abf..69ef8ac3a2e4 100644 --- a/source/isaaclab_tasks/test/contrib/test_deploy_gear_assembly_cfg.py +++ b/source/isaaclab_tasks/test/contrib/test_deploy_gear_assembly_cfg.py @@ -6,8 +6,23 @@ """Tests for Deploy GearAssembly environment configuration defaults.""" import pytest +from isaaclab_newton.envs.mdp.actions.newton_ik_actions_cfg import NewtonInverseKinematicsActionCfg + +from pxr import Usd import isaaclab_tasks # noqa: F401 +from isaaclab_tasks.contrib.deploy.gear_assembly.config.rizon_4s.ik_newton_env_cfg import ( + Rizon4sGearAssemblyIKNewtonEnvCfg, +) +from isaaclab_tasks.contrib.deploy.gear_assembly.config.rizon_4s.joint_pos_env_cfg import ( + Rizon4sGearAssemblyEnvCfg, +) +from isaaclab_tasks.contrib.deploy.gear_assembly.gear_assembly_env_cfg import ( + _NEWTON_GEAR_OFFSETS, + _PHYSX_GEAR_OFFSETS, + NEWTON_GEAR_ASSETS_DIR, +) +from isaaclab_tasks.utils.hydra import resolve_presets from isaaclab_tasks.utils.parse_cfg import parse_env_cfg @@ -23,3 +38,115 @@ def test_ur10e_gear_assembly_default_num_envs(task_name: str): env_cfg = parse_env_cfg(task_name) assert env_cfg.scene.num_envs == 2048 + + +def test_rizon_gear_newton_collision_presets_use_local_sdf_assets(): + """Newton point and hydro presets should preserve concave package-local gear collision.""" + default_cfg = resolve_presets(Rizon4sGearAssemblyEnvCfg(), {"default"}) + mujoco_cfg = resolve_presets(Rizon4sGearAssemblyEnvCfg(), {"newton_mjwarp"}) + point_cfg = resolve_presets(Rizon4sGearAssemblyEnvCfg(), {"newton_sdf"}) + hydro_cfg = resolve_presets(Rizon4sGearAssemblyEnvCfg(), {"newton_hydroelastic"}) + + assert type(default_cfg.sim.physics).__name__ == "PhysxCfg" + assert default_cfg.sim.dt == pytest.approx(1.0 / 120.0) + assert default_cfg.decimation == 4 + assert default_cfg.scene.replicate_physics is False + assert default_cfg.scene.robot.spawn.joint_drive_props is None + assert default_cfg.scene.robot.spawn.rigid_props.disable_gravity is True + assert default_cfg.scene.factory_gear_base.init_state.pos[2] == 0.071 + assert "/assets/newton/" not in default_cfg.scene.factory_gear_base.spawn.usd_path + + assert mujoco_cfg.sim.physics.solver_cfg.use_mujoco_contacts is True + assert "/assets/newton/" in mujoco_cfg.scene.factory_gear_base.spawn.usd_path + + for env_cfg in (mujoco_cfg, point_cfg, hydro_cfg): + assert env_cfg.sim.dt == 0.01 + assert env_cfg.decimation == 3 + assert env_cfg.scene.replicate_physics is True + assert env_cfg.sim.physics.num_substeps == 20 + assert env_cfg.sim.physics.collision_decimation == 0 + assert env_cfg.sim.physics.solver_cfg.update_data_interval == 10 + assert env_cfg.scene.robot.spawn.joint_drive_props.actuatorgravcomp is True + assert env_cfg.scene.robot.spawn.rigid_props.disable_gravity is False + + for env_cfg in (point_cfg, hydro_cfg): + assert env_cfg.sim.physics.solver_cfg.use_mujoco_contacts is False + assert env_cfg.sim.physics.collision_cfg.max_triangle_pairs == 1_500_000 + assert env_cfg.sim.physics.default_shape_cfg.gap == 0.005 + for asset_name in ( + "factory_gear_base", + "factory_gear_small", + "factory_gear_medium", + "factory_gear_large", + ): + usd_path = getattr(env_cfg.scene, asset_name).spawn.usd_path + assert "/assets/newton/" in usd_path + assert usd_path.endswith(".usda") + + assert point_cfg.sim.physics.collision_cfg.sdf_hydroelastic_config is None + assert hydro_cfg.sim.physics.collision_cfg.sdf_hydroelastic_config is not None + assert point_cfg.scene.factory_gear_small.spawn.usd_path.endswith("factory_gear_small.usda") + assert hydro_cfg.scene.factory_gear_small.spawn.usd_path.endswith("factory_gear_small_hydroelastic.usda") + + +def test_rizon_gear_uses_shaft_targets_relative_actions_and_physical_gripper(): + """Newton should use shaft targets and a physical grasp while preserving PhysX defaults.""" + default_cfg = resolve_presets(Rizon4sGearAssemblyEnvCfg(), {"default"}) + hydro_cfg = resolve_presets(Rizon4sGearAssemblyEnvCfg(), {"newton_hydroelastic"}) + + for env_cfg in (default_cfg, hydro_cfg): + assert type(env_cfg.actions.arm_action).__name__ == "RelativeJointPositionActionCfg" + assert env_cfg.actions.gripper_action is None + + assert default_cfg.gear_offsets == _PHYSX_GEAR_OFFSETS + assert default_cfg.rewards.end_effector_gear_keypoint_tracking.params["gear_offsets"] is None + assert default_cfg.events.pin_unselected_gears_to_shafts is None + assert default_cfg.grasp_center_body_names is None + assert default_cfg.scene.robot.actuators["gripper_passive"].stiffness == 0.0 + + assert hydro_cfg.scene.factory_gear_base.init_state.pos[2] == -0.005 + assert hydro_cfg.gear_offsets == _NEWTON_GEAR_OFFSETS + assert hydro_cfg.observations.policy.gear_shaft_pos.params["gear_offsets"] == _NEWTON_GEAR_OFFSETS + assert hydro_cfg.rewards.end_effector_gear_keypoint_tracking.params["gear_offsets"] == _NEWTON_GEAR_OFFSETS + assert hydro_cfg.rewards.end_effector_gear_keypoint_tracking_exp.params["gear_offsets"] == (_NEWTON_GEAR_OFFSETS) + assert hydro_cfg.events.pin_unselected_gears_to_shafts is not None + assert hydro_cfg.grasp_center_body_names == ("left_finger_tip", "right_finger_tip") + assert all(offset == [0.0, 0.0, -0.026] for offset in hydro_cfg.gear_offsets_grasp.values()) + + passive_gripper = hydro_cfg.scene.robot.actuators["gripper_passive"] + assert passive_gripper.stiffness == 2_000.0 + assert passive_gripper.damping == 10.0 + assert passive_gripper.effort_limit_sim == 20.0 + assert hydro_cfg.scene.robot.actuators["shoulder"].stiffness == 6_000.0 + + +def test_rizon_gear_newton_ik_commands_physical_flange(): + """Newton IK should expose a relative flange-pose action with normal safety limits.""" + env_cfg = resolve_presets(Rizon4sGearAssemblyIKNewtonEnvCfg(), {"newton_hydroelastic"}) + pose_objective = env_cfg.actions.arm_action.objectives[0] + + assert isinstance(env_cfg.actions.arm_action, NewtonInverseKinematicsActionCfg) + assert pose_objective.body_name == "flange" + assert pose_objective.use_relative_mode is True + assert pose_objective.scale == 0.025 + assert env_cfg.terminations.gear_orientation_exceeded.params["roll_threshold_deg"] == 15.0 + assert env_cfg.terminations.gear_orientation_exceeded.params["pitch_threshold_deg"] == 15.0 + + +def test_rizon_gear_assets_author_newton_sdf_per_collider(): + """Every package-local gear mesh should preserve concavity with authored Newton SDF metadata.""" + for asset_name in ( + "factory_gear_base", + "factory_gear_small", + "factory_gear_medium", + "factory_gear_large", + ): + for suffix, hydroelastic_enabled in (("", False), ("_hydroelastic", True)): + usd_path = f"{NEWTON_GEAR_ASSETS_DIR}/{asset_name}/{asset_name}{suffix}.usda" + stage = Usd.Stage.Open(usd_path) + sdf_prims = [prim for prim in stage.Traverse() if prim.HasAttribute("newton:sdfMaxResolution")] + + assert len(sdf_prims) == 1 + assert sdf_prims[0].GetAttribute("newton:sdfMaxResolution").Get() == 128 + assert sdf_prims[0].GetAttribute("newton:hydroelasticEnabled").Get() is hydroelastic_enabled + assert sdf_prims[0].GetAttribute("physics:approximation").Get() == "sdf" From 571f25593dc9be180227ec82a512944a0cfb4940 Mon Sep 17 00:00:00 2001 From: curiep Date: Fri, 31 Jul 2026 10:10:36 -0700 Subject: [PATCH 2/4] Increase Newton gear collision capacity Raise the per-rank triangle-pair capacity above the measured reset peak so 256-environment shards do not discard candidate contacts. Document the sizing assumption and cover it in the task configuration test. --- .../02_gear_assembly/gear_assembly_policy.rst | 10 +++++++--- .../deploy/gear_assembly/gear_assembly_env_cfg.py | 5 +++-- .../test/contrib/test_deploy_gear_assembly_cfg.py | 2 +- 3 files changed, 11 insertions(+), 6 deletions(-) diff --git a/docs/source/policy_deployment/02_gear_assembly/gear_assembly_policy.rst b/docs/source/policy_deployment/02_gear_assembly/gear_assembly_policy.rst index ff30123d67a7..4d153f5d46fa 100644 --- a/docs/source/policy_deployment/02_gear_assembly/gear_assembly_policy.rst +++ b/docs/source/policy_deployment/02_gear_assembly/gear_assembly_policy.rst @@ -877,9 +877,13 @@ The Newton hydroelastic preset uses package-local concave SDF collision assets. separate the expensive collision rate from contact integration: the outer physics and collision tick is 100 Hz, while 20 solver substeps produce a 2 kHz solver rate. A policy decimation of three gives a 33.3 Hz control rate. The preset also uses a 5 mm shape gap and a bounded triangle-pair -buffer. The arm uses solver-native actuator gravity compensation so gravity remains enabled for -the gears, and the Grav linkage is held by physical PD actuators; no action term rewrites the -selected gear pose. +buffer sized for 256 environments per GPU rank. The arm uses solver-native actuator gravity +compensation so gravity remains enabled for the gears, and the Grav linkage is held by physical +PD actuators; no action term rewrites the selected gear pose. + +The default triangle-pair capacity is 4,194,304 per GPU rank. If Newton reports a triangle-pair +buffer overflow, reduce the number of environments per rank or increase the capacity before using +the resulting policy; overflowed candidate contacts are discarded. For four-GPU training, launch 256 environments per rank (1,024 total): diff --git a/source/isaaclab_tasks/isaaclab_tasks/contrib/deploy/gear_assembly/gear_assembly_env_cfg.py b/source/isaaclab_tasks/isaaclab_tasks/contrib/deploy/gear_assembly/gear_assembly_env_cfg.py index ec86be7721e0..4381e50dad0f 100644 --- a/source/isaaclab_tasks/isaaclab_tasks/contrib/deploy/gear_assembly/gear_assembly_env_cfg.py +++ b/source/isaaclab_tasks/isaaclab_tasks/contrib/deploy/gear_assembly/gear_assembly_env_cfg.py @@ -42,8 +42,9 @@ ASSETS_DIR = os.path.join(CONFIG_DIR, "assets") NEWTON_GEAR_ASSETS_DIR = os.path.join(ASSETS_DIR, "newton") -# A 256-world GPU shard peaks near 1.1M broad-phase pairs during randomized resets. -_GEAR_MAX_TRIANGLE_PAIRS = 1_500_000 +# A 256-world GPU shard reached 1.54M broad-phase pairs during randomized resets. +# Keep power-of-two headroom so reset spikes do not discard candidate contacts. +_GEAR_MAX_TRIANGLE_PAIRS = 4_194_304 _PHYSX_GEAR_OFFSETS = { "gear_small": [0.076125, 0.0, 0.0], "gear_medium": [0.030375, 0.0, 0.0], diff --git a/source/isaaclab_tasks/test/contrib/test_deploy_gear_assembly_cfg.py b/source/isaaclab_tasks/test/contrib/test_deploy_gear_assembly_cfg.py index 69ef8ac3a2e4..60a8e98e067a 100644 --- a/source/isaaclab_tasks/test/contrib/test_deploy_gear_assembly_cfg.py +++ b/source/isaaclab_tasks/test/contrib/test_deploy_gear_assembly_cfg.py @@ -71,7 +71,7 @@ def test_rizon_gear_newton_collision_presets_use_local_sdf_assets(): for env_cfg in (point_cfg, hydro_cfg): assert env_cfg.sim.physics.solver_cfg.use_mujoco_contacts is False - assert env_cfg.sim.physics.collision_cfg.max_triangle_pairs == 1_500_000 + assert env_cfg.sim.physics.collision_cfg.max_triangle_pairs == 4_194_304 assert env_cfg.sim.physics.default_shape_cfg.gap == 0.005 for asset_name in ( "factory_gear_base", From 1ea7a1eee60ac77e687ec5aa8c8cc3a0bb28d7f9 Mon Sep 17 00:00:00 2001 From: curiep Date: Fri, 31 Jul 2026 13:05:30 -0700 Subject: [PATCH 3/4] Add Curie Park to contributors --- CONTRIBUTORS.md | 1 + 1 file changed, 1 insertion(+) diff --git a/CONTRIBUTORS.md b/CONTRIBUTORS.md index 4fd3a69de84a..603aa52fdfcf 100644 --- a/CONTRIBUTORS.md +++ b/CONTRIBUTORS.md @@ -75,6 +75,7 @@ Guidelines for modifications: * Chenyu Cao * Chenyu Yang * Connor Smith +* Curie Park * CY (Chien-Ying) Chen * David Cao-Mueller * David Leon From c70ca458710f018fab9431db96d33fd1db2b6544 Mon Sep 17 00:00:00 2001 From: curiep Date: Fri, 31 Jul 2026 15:49:26 -0700 Subject: [PATCH 4/4] Share gear SDF assets across backends Use one centered collision mesh set for PhysX and Newton while preserving the legacy PhysX asset frames for existing checkpoints. Select Newton hydroelastic behavior at spawn time and expose an explicit PhysX SDF validation preset. --- .../02_gear_assembly/gear_assembly_policy.rst | 5 ++ .../curiep-newton-gear-insertion.minor.rst | 2 + .../factory_gear_base/factory_gear_base.usda | 4 +- .../factory_gear_base_hydroelastic.usda | 3 - .../factory_gear_large.usda | 4 +- .../factory_gear_large_hydroelastic.usda | 3 - .../factory_gear_medium.usda | 4 +- .../factory_gear_medium_hydroelastic.usda | 3 - .../factory_gear_small.usda | 4 +- .../factory_gear_small_hydroelastic.usda | 3 - .../config/rizon_4s/joint_pos_env_cfg.py | 31 +++++++-- .../gear_assembly/gear_assembly_env_cfg.py | 49 +++++++++++--- .../contrib/test_deploy_gear_assembly_cfg.py | 65 ++++++++++++++----- 13 files changed, 125 insertions(+), 55 deletions(-) delete mode 100644 source/isaaclab_tasks/isaaclab_tasks/contrib/deploy/gear_assembly/assets/newton/factory_gear_base/factory_gear_base_hydroelastic.usda delete mode 100644 source/isaaclab_tasks/isaaclab_tasks/contrib/deploy/gear_assembly/assets/newton/factory_gear_large/factory_gear_large_hydroelastic.usda delete mode 100644 source/isaaclab_tasks/isaaclab_tasks/contrib/deploy/gear_assembly/assets/newton/factory_gear_medium/factory_gear_medium_hydroelastic.usda delete mode 100644 source/isaaclab_tasks/isaaclab_tasks/contrib/deploy/gear_assembly/assets/newton/factory_gear_small/factory_gear_small_hydroelastic.usda diff --git a/docs/source/policy_deployment/02_gear_assembly/gear_assembly_policy.rst b/docs/source/policy_deployment/02_gear_assembly/gear_assembly_policy.rst index 4d153f5d46fa..cab1a8b2f529 100644 --- a/docs/source/policy_deployment/02_gear_assembly/gear_assembly_policy.rst +++ b/docs/source/policy_deployment/02_gear_assembly/gear_assembly_policy.rst @@ -762,6 +762,11 @@ First, launch the training with a small number of environments and visualization policy, use the ``IsaacContrib-Deploy-GearAssembly-Rizon4s-Grav-Newton-IK`` task with the same Newton preset. + The centered SDF assets can also be tested with PhysX by using + ``presets=physx_sdf`` on the ``IsaacContrib-Deploy-GearAssembly-Rizon4s-Grav`` + task. The default PhysX preset retains the legacy Factory assets and frame conventions + for existing checkpoints. + This opens the selected viewer (Isaac Sim for PhysX or Newton Viewer for Newton) so you can observe the training process in real time. diff --git a/source/isaaclab_tasks/changelog.d/curiep-newton-gear-insertion.minor.rst b/source/isaaclab_tasks/changelog.d/curiep-newton-gear-insertion.minor.rst index 670536a1b2f8..171d99c37f22 100644 --- a/source/isaaclab_tasks/changelog.d/curiep-newton-gear-insertion.minor.rst +++ b/source/isaaclab_tasks/changelog.d/curiep-newton-gear-insertion.minor.rst @@ -3,6 +3,8 @@ Added * Added Newton point-SDF and hydroelastic-SDF physics presets and collision assets to the Flexiv Rizon 4s gear assembly task. +* Added a PhysX SDF preset that uses the same centered collision assets for + cross-backend validation while preserving legacy PhysX defaults. * Added a Newton inverse-kinematics task variant for task-space policies. Fixed diff --git a/source/isaaclab_tasks/isaaclab_tasks/contrib/deploy/gear_assembly/assets/newton/factory_gear_base/factory_gear_base.usda b/source/isaaclab_tasks/isaaclab_tasks/contrib/deploy/gear_assembly/assets/newton/factory_gear_base/factory_gear_base.usda index f66f09bd8392..122d4caf1666 100644 --- a/source/isaaclab_tasks/isaaclab_tasks/contrib/deploy/gear_assembly/assets/newton/factory_gear_base/factory_gear_base.usda +++ b/source/isaaclab_tasks/isaaclab_tasks/contrib/deploy/gear_assembly/assets/newton/factory_gear_base/factory_gear_base.usda @@ -1,3 +1,3 @@ version https://git-lfs.github.com/spec/v1 -oid sha256:2810d8d2faecd2a7d8e08ab78aa59dbfc513f77f047b3f49a8506ba99bc82771 -size 240517 +oid sha256:3bcba32e904a81cca277a8b2480634f34fcfd78ade2f261ad3ab5878b73f30f8 +size 240603 diff --git a/source/isaaclab_tasks/isaaclab_tasks/contrib/deploy/gear_assembly/assets/newton/factory_gear_base/factory_gear_base_hydroelastic.usda b/source/isaaclab_tasks/isaaclab_tasks/contrib/deploy/gear_assembly/assets/newton/factory_gear_base/factory_gear_base_hydroelastic.usda deleted file mode 100644 index 9462f1626f28..000000000000 --- a/source/isaaclab_tasks/isaaclab_tasks/contrib/deploy/gear_assembly/assets/newton/factory_gear_base/factory_gear_base_hydroelastic.usda +++ /dev/null @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:e2cf23310834fff8d049797a9e47a383478f501cd326fd59a800b95cb51523dd -size 473 diff --git a/source/isaaclab_tasks/isaaclab_tasks/contrib/deploy/gear_assembly/assets/newton/factory_gear_large/factory_gear_large.usda b/source/isaaclab_tasks/isaaclab_tasks/contrib/deploy/gear_assembly/assets/newton/factory_gear_large/factory_gear_large.usda index 1e03caa179c0..1af9c0812e05 100644 --- a/source/isaaclab_tasks/isaaclab_tasks/contrib/deploy/gear_assembly/assets/newton/factory_gear_large/factory_gear_large.usda +++ b/source/isaaclab_tasks/isaaclab_tasks/contrib/deploy/gear_assembly/assets/newton/factory_gear_large/factory_gear_large.usda @@ -1,3 +1,3 @@ version https://git-lfs.github.com/spec/v1 -oid sha256:f72d166f8978491cc5ccc826ba1683d8ae817722065f34779ffdafb4c8e4ec6a -size 690305 +oid sha256:606b5a4b67b44f529468475c3e1a042996e6b2ef96651ebb378e7f1f7953ffa3 +size 690391 diff --git a/source/isaaclab_tasks/isaaclab_tasks/contrib/deploy/gear_assembly/assets/newton/factory_gear_large/factory_gear_large_hydroelastic.usda b/source/isaaclab_tasks/isaaclab_tasks/contrib/deploy/gear_assembly/assets/newton/factory_gear_large/factory_gear_large_hydroelastic.usda deleted file mode 100644 index 9fb03f560b7e..000000000000 --- a/source/isaaclab_tasks/isaaclab_tasks/contrib/deploy/gear_assembly/assets/newton/factory_gear_large/factory_gear_large_hydroelastic.usda +++ /dev/null @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:12dfce05c39ec5927d777c59961999eec99ba6fa6c57498715a148d5a8039317 -size 477 diff --git a/source/isaaclab_tasks/isaaclab_tasks/contrib/deploy/gear_assembly/assets/newton/factory_gear_medium/factory_gear_medium.usda b/source/isaaclab_tasks/isaaclab_tasks/contrib/deploy/gear_assembly/assets/newton/factory_gear_medium/factory_gear_medium.usda index 8638cdad4fd2..f7fab98c2dbe 100644 --- a/source/isaaclab_tasks/isaaclab_tasks/contrib/deploy/gear_assembly/assets/newton/factory_gear_medium/factory_gear_medium.usda +++ b/source/isaaclab_tasks/isaaclab_tasks/contrib/deploy/gear_assembly/assets/newton/factory_gear_medium/factory_gear_medium.usda @@ -1,3 +1,3 @@ version https://git-lfs.github.com/spec/v1 -oid sha256:3736a62c7c48b3e8ce9c8773e9983c6e9d2ecfd52d8e39a7a1c50e0360f56868 -size 495131 +oid sha256:c03bb09b10a9fa62353af29fbcd45a6a08e28b042a9bd459bf9f04b46950120c +size 495217 diff --git a/source/isaaclab_tasks/isaaclab_tasks/contrib/deploy/gear_assembly/assets/newton/factory_gear_medium/factory_gear_medium_hydroelastic.usda b/source/isaaclab_tasks/isaaclab_tasks/contrib/deploy/gear_assembly/assets/newton/factory_gear_medium/factory_gear_medium_hydroelastic.usda deleted file mode 100644 index e59584596bc0..000000000000 --- a/source/isaaclab_tasks/isaaclab_tasks/contrib/deploy/gear_assembly/assets/newton/factory_gear_medium/factory_gear_medium_hydroelastic.usda +++ /dev/null @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:72e61fea6b51a6bf364680b9b2eb289efa2acf8bfc74d83d1faef04961fe0b02 -size 481 diff --git a/source/isaaclab_tasks/isaaclab_tasks/contrib/deploy/gear_assembly/assets/newton/factory_gear_small/factory_gear_small.usda b/source/isaaclab_tasks/isaaclab_tasks/contrib/deploy/gear_assembly/assets/newton/factory_gear_small/factory_gear_small.usda index 98607558332b..20d6569c74a4 100644 --- a/source/isaaclab_tasks/isaaclab_tasks/contrib/deploy/gear_assembly/assets/newton/factory_gear_small/factory_gear_small.usda +++ b/source/isaaclab_tasks/isaaclab_tasks/contrib/deploy/gear_assembly/assets/newton/factory_gear_small/factory_gear_small.usda @@ -1,3 +1,3 @@ version https://git-lfs.github.com/spec/v1 -oid sha256:e2122445d00107c7698992ade203c72553988298e8a5c01ea58900e7bdf140c7 -size 331955 +oid sha256:205aee2b77a1b308fa2828cb4bb9501cde6316aced6f78a0eb503855dfadfcca +size 332041 diff --git a/source/isaaclab_tasks/isaaclab_tasks/contrib/deploy/gear_assembly/assets/newton/factory_gear_small/factory_gear_small_hydroelastic.usda b/source/isaaclab_tasks/isaaclab_tasks/contrib/deploy/gear_assembly/assets/newton/factory_gear_small/factory_gear_small_hydroelastic.usda deleted file mode 100644 index d17b73df9a23..000000000000 --- a/source/isaaclab_tasks/isaaclab_tasks/contrib/deploy/gear_assembly/assets/newton/factory_gear_small/factory_gear_small_hydroelastic.usda +++ /dev/null @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:43d91b4a9701c040c6df4136e2404540f3caa04d6a9becbe07b3f526a1fa7af7 -size 477 diff --git a/source/isaaclab_tasks/isaaclab_tasks/contrib/deploy/gear_assembly/config/rizon_4s/joint_pos_env_cfg.py b/source/isaaclab_tasks/isaaclab_tasks/contrib/deploy/gear_assembly/config/rizon_4s/joint_pos_env_cfg.py index 38e0730405c0..fcc7387f3dc6 100644 --- a/source/isaaclab_tasks/isaaclab_tasks/contrib/deploy/gear_assembly/config/rizon_4s/joint_pos_env_cfg.py +++ b/source/isaaclab_tasks/isaaclab_tasks/contrib/deploy/gear_assembly/config/rizon_4s/joint_pos_env_cfg.py @@ -95,6 +95,7 @@ def _gear_friction_range() -> PresetCfg: return preset( default=(0.75, 0.75), physx=(0.75, 0.75), + physx_sdf=(0.75, 0.75), newton_mjwarp=(3.0, 3.0), newton_sdf=(3.0, 3.0), newton_hydroelastic=(3.0, 3.0), @@ -106,12 +107,25 @@ def _backend_preset(physx_value: object, newton_value: object) -> PresetCfg: return preset( default=physx_value, physx=physx_value, + physx_sdf=physx_value, newton_mjwarp=newton_value, newton_sdf=newton_value, newton_hydroelastic=newton_value, ) +def _gear_asset_frame_preset(legacy_value: object, centered_value: object) -> PresetCfg: + """Create a preset that selects centered gear frames across physics backends.""" + return preset( + default=legacy_value, + physx=legacy_value, + physx_sdf=centered_value, + newton_mjwarp=centered_value, + newton_sdf=centered_value, + newton_hydroelastic=centered_value, + ) + + ## # Environment configuration ## @@ -210,6 +224,7 @@ class EventCfg: "gear_offsets": preset( default=None, physx=None, + physx_sdf=_NEWTON_GEAR_OFFSETS, newton_mjwarp=_NEWTON_GEAR_OFFSETS, newton_sdf=_NEWTON_GEAR_OFFSETS, newton_hydroelastic=_NEWTON_GEAR_OFFSETS, @@ -217,6 +232,7 @@ class EventCfg: "seated_gear_z_offset": preset( default=0.0, physx=0.0, + physx_sdf=0.0075, newton_mjwarp=0.0075, newton_sdf=0.0075, newton_hydroelastic=0.0075, @@ -236,6 +252,7 @@ class EventCfg: pin_unselected_gears_to_shafts = preset( default=None, physx=None, + physx_sdf=None, newton_mjwarp=_NEWTON_PIN_UNSELECTED_GEARS_EVENT, newton_sdf=_NEWTON_PIN_UNSELECTED_GEARS_EVENT, newton_hydroelastic=_NEWTON_PIN_UNSELECTED_GEARS_EVENT, @@ -250,7 +267,7 @@ class Rizon4sGearAssemblyEnvCfg(GearAssemblyEnvCfg): def __post_init__(self): super().__post_init__() - self.gear_offsets = _backend_preset(_PHYSX_GEAR_OFFSETS, _NEWTON_GEAR_OFFSETS) + self.gear_offsets = _gear_asset_frame_preset(_PHYSX_GEAR_OFFSETS, _NEWTON_GEAR_OFFSETS) arm_joint_names = ["joint1", "joint2", "joint3", "joint4", "joint5", "joint6", "joint7"] self.end_effector_body_name = "link7" @@ -383,7 +400,7 @@ def __post_init__(self): base_rot = (0.0, 0.0, 0.70711, -0.70711) physx_asset_state = RigidObjectCfg.InitialStateCfg(pos=(0.481, -0.073, 0.071), rot=base_rot) newton_base_pos = (0.481, -0.073, -0.005) - self.scene.factory_gear_base.init_state = _backend_preset( + self.scene.factory_gear_base.init_state = _gear_asset_frame_preset( physx_asset_state, RigidObjectCfg.InitialStateCfg(pos=newton_base_pos, rot=base_rot), ) @@ -397,7 +414,7 @@ def __post_init__(self): newton_base_pos[1], newton_base_pos[2], ) - asset.init_state = _backend_preset( + asset.init_state = _gear_asset_frame_preset( physx_asset_state, RigidObjectCfg.InitialStateCfg(pos=newton_gear_pos, rot=base_rot), ) @@ -408,13 +425,13 @@ def __post_init__(self): "gear_medium": [0.0, 0.0, -0.026], "gear_large": [0.0, 0.0, -0.026], } - self.gear_offsets_grasp = _backend_preset(physx_grasp_offsets, newton_grasp_offsets) - self.grasp_center_body_names = _backend_preset(None, ("left_finger_tip", "right_finger_tip")) - self.hand_grasp_width = _backend_preset( + self.gear_offsets_grasp = _gear_asset_frame_preset(physx_grasp_offsets, newton_grasp_offsets) + self.grasp_center_body_names = _gear_asset_frame_preset(None, ("left_finger_tip", "right_finger_tip")) + self.hand_grasp_width = _gear_asset_frame_preset( {"gear_small": 0.05, "gear_medium": 0.2, "gear_large": 0.28}, {"gear_small": 0.01, "gear_medium": 0.2, "gear_large": 0.28}, ) - self.hand_close_width = _backend_preset( + self.hand_close_width = _gear_asset_frame_preset( {"gear_small": 0.0, "gear_medium": 0.139626, "gear_large": 0.139626}, {"gear_small": 0.01, "gear_medium": 0.139626, "gear_large": 0.139626}, ) diff --git a/source/isaaclab_tasks/isaaclab_tasks/contrib/deploy/gear_assembly/gear_assembly_env_cfg.py b/source/isaaclab_tasks/isaaclab_tasks/contrib/deploy/gear_assembly/gear_assembly_env_cfg.py index 4381e50dad0f..0e882dee2a41 100644 --- a/source/isaaclab_tasks/isaaclab_tasks/contrib/deploy/gear_assembly/gear_assembly_env_cfg.py +++ b/source/isaaclab_tasks/isaaclab_tasks/contrib/deploy/gear_assembly/gear_assembly_env_cfg.py @@ -13,6 +13,7 @@ NewtonCollisionPipelineCfg, NewtonShapeCfg, ) +from isaaclab_newton.sim.schemas import NewtonSDFCollisionPropertiesCfg from isaaclab_physx.physics import PhysxCfg from isaaclab_physx.sim.schemas import PhysxCollisionPropertiesCfg, PhysxRigidBodyPropertiesCfg @@ -58,20 +59,36 @@ def _gear_usd_path(default_usd_path: str, asset_name: str) -> PresetCfg: - """Create a gear USD path preset with Newton-specific collision assets. + """Create a gear USD path preset with package-local SDF collision assets. Args: default_usd_path: Factory asset USD path used by the default and PhysX presets. asset_name: Gear asset directory and USD stem. Returns: - Preset that resolves to package-local point-SDF or hydroelastic assets for Newton collision presets. + Preset that resolves to the package-local SDF asset for SDF collision presets. """ + newton_usd_path = os.path.join(NEWTON_GEAR_ASSETS_DIR, asset_name, f"{asset_name}.usda") return preset( default=default_usd_path, - newton_mjwarp=os.path.join(NEWTON_GEAR_ASSETS_DIR, asset_name, f"{asset_name}.usda"), - newton_sdf=os.path.join(NEWTON_GEAR_ASSETS_DIR, asset_name, f"{asset_name}.usda"), - newton_hydroelastic=os.path.join(NEWTON_GEAR_ASSETS_DIR, asset_name, f"{asset_name}_hydroelastic.usda"), + physx_sdf=newton_usd_path, + newton_mjwarp=newton_usd_path, + newton_sdf=newton_usd_path, + newton_hydroelastic=newton_usd_path, + ) + + +def _gear_collision_properties() -> PresetCfg: + """Create backend-specific collision properties for the gear assets.""" + return preset( + default=PhysxCollisionPropertiesCfg(contact_offset=0.02, rest_offset=0.0), + physx=PhysxCollisionPropertiesCfg(contact_offset=0.02, rest_offset=0.0), + physx_sdf=PhysxCollisionPropertiesCfg(contact_offset=0.02, rest_offset=0.0), + newton_mjwarp=NewtonSDFCollisionPropertiesCfg(contact_offset=0.02, rest_offset=0.0, hydroelastic_enabled=False), + newton_sdf=NewtonSDFCollisionPropertiesCfg(contact_offset=0.02, rest_offset=0.0, hydroelastic_enabled=False), + newton_hydroelastic=NewtonSDFCollisionPropertiesCfg( + contact_offset=0.02, rest_offset=0.0, hydroelastic_enabled=True + ), ) @@ -88,7 +105,8 @@ class GearAssemblyPhysicsCfg(PresetCfg): Newton (MuJoCo) solver limits are set conservatively. Select a preset at runtime with the ``presets=`` CLI override: - * ``default`` and ``physx`` -- PhysX with contact buffers sized for assembly. + * ``default`` and ``physx`` -- PhysX with the legacy Factory assets. + * ``physx_sdf`` -- PhysX with the same centered, package-local SDF assets as Newton. * ``newton_mjwarp`` -- Newton with MuJoCo's internal contact solver. * ``newton_sdf`` -- Newton's collision pipeline with reduced point-SDF contacts. * ``newton_hydroelastic`` -- Newton's own collision pipeline (``use_mujoco_contacts=False``) @@ -180,6 +198,11 @@ class GearAssemblyPhysicsCfg(PresetCfg): gpu_max_rigid_contact_count=2**23, gpu_max_rigid_patch_count=2**23, ) + physx_sdf: PhysxCfg = PhysxCfg( + gpu_collision_stack_size=2**30, + gpu_max_rigid_contact_count=2**23, + gpu_max_rigid_patch_count=2**23, + ) default = physx @@ -196,6 +219,7 @@ class GearAssemblySceneCfg(InteractiveSceneCfg): replicate_physics = preset( default=False, physx=False, + physx_sdf=False, newton_mjwarp=True, newton_sdf=True, newton_hydroelastic=True, @@ -231,7 +255,7 @@ class GearAssemblySceneCfg(InteractiveSceneCfg): max_contact_impulse=1e32, ), mass_props=sim_utils.MassPropertiesCfg(mass=None), - collision_props=PhysxCollisionPropertiesCfg(contact_offset=0.02, rest_offset=0.0), + collision_props=_gear_collision_properties(), ), init_state=RigidObjectCfg.InitialStateCfg(pos=(-1.0200, 0.2100, -0.1), rot=(0.0, 0.0, 0.70711, 0.70711)), ) @@ -259,7 +283,7 @@ class GearAssemblySceneCfg(InteractiveSceneCfg): max_contact_impulse=1e32, ), mass_props=sim_utils.MassPropertiesCfg(mass=None), - collision_props=PhysxCollisionPropertiesCfg(contact_offset=0.02, rest_offset=0.0), + collision_props=_gear_collision_properties(), ), init_state=RigidObjectCfg.InitialStateCfg(pos=(-1.0200, 0.2100, -0.1), rot=(0.0, 0.0, 0.70711, 0.70711)), ) @@ -287,7 +311,7 @@ class GearAssemblySceneCfg(InteractiveSceneCfg): max_contact_impulse=1e32, ), mass_props=sim_utils.MassPropertiesCfg(mass=None), - collision_props=PhysxCollisionPropertiesCfg(contact_offset=0.02, rest_offset=0.0), + collision_props=_gear_collision_properties(), ), init_state=RigidObjectCfg.InitialStateCfg(pos=(-1.0200, 0.2100, -0.1), rot=(0.0, 0.0, 0.70711, 0.70711)), ) @@ -315,7 +339,7 @@ class GearAssemblySceneCfg(InteractiveSceneCfg): max_contact_impulse=1e32, ), mass_props=sim_utils.MassPropertiesCfg(mass=None), - collision_props=PhysxCollisionPropertiesCfg(contact_offset=0.02, rest_offset=0.0), + collision_props=_gear_collision_properties(), ), init_state=RigidObjectCfg.InitialStateCfg(pos=(-1.0200, 0.2100, -0.1), rot=(0.0, 0.0, 0.70711, 0.70711)), ) @@ -484,6 +508,7 @@ def __post_init__(self): self.decimation = preset( default=4, physx=4, + physx_sdf=4, newton_mjwarp=3, newton_sdf=3, newton_hydroelastic=3, @@ -491,6 +516,7 @@ def __post_init__(self): self.sim.render_interval = preset( default=4, physx=4, + physx_sdf=4, newton_mjwarp=3, newton_sdf=3, newton_hydroelastic=3, @@ -498,6 +524,7 @@ def __post_init__(self): self.sim.dt = preset( default=1.0 / 120.0, physx=1.0 / 120.0, + physx_sdf=1.0 / 120.0, newton_mjwarp=0.01, newton_sdf=0.01, newton_hydroelastic=0.01, @@ -508,6 +535,7 @@ def __post_init__(self): gear_offsets = preset( default=physx_gear_offsets, physx=physx_gear_offsets, + physx_sdf=newton_gear_offsets, newton_mjwarp=newton_gear_offsets, newton_sdf=newton_gear_offsets, newton_hydroelastic=newton_gear_offsets, @@ -520,6 +548,7 @@ def __post_init__(self): reward_gear_offsets = preset( default=None, physx=None, + physx_sdf=newton_gear_offsets, newton_mjwarp=newton_gear_offsets, newton_sdf=newton_gear_offsets, newton_hydroelastic=newton_gear_offsets, diff --git a/source/isaaclab_tasks/test/contrib/test_deploy_gear_assembly_cfg.py b/source/isaaclab_tasks/test/contrib/test_deploy_gear_assembly_cfg.py index 60a8e98e067a..28884f863360 100644 --- a/source/isaaclab_tasks/test/contrib/test_deploy_gear_assembly_cfg.py +++ b/source/isaaclab_tasks/test/contrib/test_deploy_gear_assembly_cfg.py @@ -40,9 +40,10 @@ def test_ur10e_gear_assembly_default_num_envs(task_name: str): assert env_cfg.scene.num_envs == 2048 -def test_rizon_gear_newton_collision_presets_use_local_sdf_assets(): - """Newton point and hydro presets should preserve concave package-local gear collision.""" +def test_rizon_gear_sdf_presets_use_local_assets(): + """PhysX and Newton SDF presets should share concave package-local collision assets.""" default_cfg = resolve_presets(Rizon4sGearAssemblyEnvCfg(), {"default"}) + physx_sdf_cfg = resolve_presets(Rizon4sGearAssemblyEnvCfg(), {"physx_sdf"}) mujoco_cfg = resolve_presets(Rizon4sGearAssemblyEnvCfg(), {"newton_mjwarp"}) point_cfg = resolve_presets(Rizon4sGearAssemblyEnvCfg(), {"newton_sdf"}) hydro_cfg = resolve_presets(Rizon4sGearAssemblyEnvCfg(), {"newton_hydroelastic"}) @@ -56,6 +57,23 @@ def test_rizon_gear_newton_collision_presets_use_local_sdf_assets(): assert default_cfg.scene.factory_gear_base.init_state.pos[2] == 0.071 assert "/assets/newton/" not in default_cfg.scene.factory_gear_base.spawn.usd_path + assert type(physx_sdf_cfg.sim.physics).__name__ == "PhysxCfg" + assert physx_sdf_cfg.sim.dt == pytest.approx(1.0 / 120.0) + assert physx_sdf_cfg.decimation == 4 + assert physx_sdf_cfg.scene.replicate_physics is False + assert physx_sdf_cfg.scene.factory_gear_base.init_state.pos[2] == -0.005 + assert physx_sdf_cfg.gear_offsets == _NEWTON_GEAR_OFFSETS + assert physx_sdf_cfg.rewards.end_effector_gear_keypoint_tracking.params["gear_offsets"] == (_NEWTON_GEAR_OFFSETS) + assert physx_sdf_cfg.events.pin_unselected_gears_to_shafts is None + assert physx_sdf_cfg.grasp_center_body_names == ("left_finger_tip", "right_finger_tip") + for asset_name in ( + "factory_gear_base", + "factory_gear_small", + "factory_gear_medium", + "factory_gear_large", + ): + assert "/assets/newton/" in getattr(physx_sdf_cfg.scene, asset_name).spawn.usd_path + assert mujoco_cfg.sim.physics.solver_cfg.use_mujoco_contacts is True assert "/assets/newton/" in mujoco_cfg.scene.factory_gear_base.spawn.usd_path @@ -79,14 +97,23 @@ def test_rizon_gear_newton_collision_presets_use_local_sdf_assets(): "factory_gear_medium", "factory_gear_large", ): - usd_path = getattr(env_cfg.scene, asset_name).spawn.usd_path - assert "/assets/newton/" in usd_path - assert usd_path.endswith(".usda") + spawn = getattr(env_cfg.scene, asset_name).spawn + assert "/assets/newton/" in spawn.usd_path + assert spawn.usd_path.endswith(".usda") assert point_cfg.sim.physics.collision_cfg.sdf_hydroelastic_config is None assert hydro_cfg.sim.physics.collision_cfg.sdf_hydroelastic_config is not None - assert point_cfg.scene.factory_gear_small.spawn.usd_path.endswith("factory_gear_small.usda") - assert hydro_cfg.scene.factory_gear_small.spawn.usd_path.endswith("factory_gear_small_hydroelastic.usda") + for asset_name in ( + "factory_gear_base", + "factory_gear_small", + "factory_gear_medium", + "factory_gear_large", + ): + point_spawn = getattr(point_cfg.scene, asset_name).spawn + hydro_spawn = getattr(hydro_cfg.scene, asset_name).spawn + assert hydro_spawn.usd_path == point_spawn.usd_path + assert point_spawn.collision_props.hydroelastic_enabled is False + assert hydro_spawn.collision_props.hydroelastic_enabled is True def test_rizon_gear_uses_shaft_targets_relative_actions_and_physical_gripper(): @@ -133,20 +160,22 @@ def test_rizon_gear_newton_ik_commands_physical_flange(): assert env_cfg.terminations.gear_orientation_exceeded.params["pitch_threshold_deg"] == 15.0 -def test_rizon_gear_assets_author_newton_sdf_per_collider(): - """Every package-local gear mesh should preserve concavity with authored Newton SDF metadata.""" +def test_rizon_gear_assets_author_sdf_per_collider(): + """Every package-local gear mesh should author both PhysX and Newton SDF metadata.""" for asset_name in ( "factory_gear_base", "factory_gear_small", "factory_gear_medium", "factory_gear_large", ): - for suffix, hydroelastic_enabled in (("", False), ("_hydroelastic", True)): - usd_path = f"{NEWTON_GEAR_ASSETS_DIR}/{asset_name}/{asset_name}{suffix}.usda" - stage = Usd.Stage.Open(usd_path) - sdf_prims = [prim for prim in stage.Traverse() if prim.HasAttribute("newton:sdfMaxResolution")] - - assert len(sdf_prims) == 1 - assert sdf_prims[0].GetAttribute("newton:sdfMaxResolution").Get() == 128 - assert sdf_prims[0].GetAttribute("newton:hydroelasticEnabled").Get() is hydroelastic_enabled - assert sdf_prims[0].GetAttribute("physics:approximation").Get() == "sdf" + usd_path = f"{NEWTON_GEAR_ASSETS_DIR}/{asset_name}/{asset_name}.usda" + stage = Usd.Stage.Open(usd_path) + sdf_prims = [prim for prim in stage.Traverse() if prim.HasAttribute("newton:sdfMaxResolution")] + + assert len(sdf_prims) == 1 + authored_schemas = sdf_prims[0].GetPrimStack()[0].GetInfo("apiSchemas").prependedItems + assert "PhysxSDFMeshCollisionAPI" in authored_schemas + assert sdf_prims[0].GetAttribute("physxSDFMeshCollision:sdfResolution").Get() == 256 + assert sdf_prims[0].GetAttribute("newton:sdfMaxResolution").Get() == 128 + assert sdf_prims[0].GetAttribute("newton:hydroelasticEnabled").Get() is False + assert sdf_prims[0].GetAttribute("physics:approximation").Get() == "sdf"