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 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..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 @@ -734,7 +734,41 @@ 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. + + 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. .. figure:: ../../_static/policy_deployment/02_gear_assembly/sim_real_gear_assembly_train.jpg :align: center @@ -820,10 +854,59 @@ 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 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): + +.. 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..171d99c37f22 --- /dev/null +++ b/source/isaaclab_tasks/changelog.d/curiep-newton-gear-insertion.minor.rst @@ -0,0 +1,15 @@ +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 +^^^^^ + +* 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..122d4caf1666 --- /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:3bcba32e904a81cca277a8b2480634f34fcfd78ade2f261ad3ab5878b73f30f8 +size 240603 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..1af9c0812e05 --- /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:606b5a4b67b44f529468475c3e1a042996e6b2ef96651ebb378e7f1f7953ffa3 +size 690391 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..f7fab98c2dbe --- /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:c03bb09b10a9fa62353af29fbcd45a6a08e28b042a9bd459bf9f04b46950120c +size 495217 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..20d6569c74a4 --- /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:205aee2b77a1b308fa2828cb4bb9501cde6316aced6f78a0eb503855dfadfcca +size 332041 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..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 @@ -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,85 @@ ## +_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 + + +_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), + physx_sdf=(0.75, 0.75), + newton_mjwarp=(3.0, 3.0), + newton_sdf=(3.0, 3.0), + newton_hydroelastic=(3.0, 3.0), + ) - # 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 +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, + physx_sdf=physx_value, + newton_mjwarp=newton_value, + newton_sdf=newton_value, + newton_hydroelastic=newton_value, + ) + - # 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 _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, + ) ## @@ -89,8 +140,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 +152,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 +164,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 +221,22 @@ class EventCfg: "z": [0.0575, 0.0775], }, "velocity_range": {}, + "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, + ), + "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, + ), }, ) @@ -182,22 +249,32 @@ 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, + ) + @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 = _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" + 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 +283,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 +343,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 +372,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 +388,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 = _gear_asset_frame_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 = _gear_asset_frame_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 = _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.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 = _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}, ) - - # 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..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 @@ -6,7 +6,16 @@ import os from dataclasses import MISSING +from isaaclab_newton.physics import ( + HydroelasticSDFCfg, + MJWarpSolverCfg, + NewtonCfg, + NewtonCollisionPipelineCfg, + NewtonShapeCfg, +) +from isaaclab_newton.sim.schemas import NewtonSDFCollisionPropertiesCfg 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 +36,194 @@ 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 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], + "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 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 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, + 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 + ), + ) + ## # 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 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``) + 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, + ) + 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 + + @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, + physx_sdf=False, + newton_mjwarp=True, + newton_sdf=True, + newton_hydroelastic=True, + ) # world ground = AssetBaseCfg( @@ -55,9 +236,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 +255,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=_gear_collision_properties(), ), init_state=RigidObjectCfg.InitialStateCfg(pos=(-1.0200, 0.2100, -0.1), rot=(0.0, 0.0, 0.70711, 0.70711)), ) @@ -80,9 +264,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 +283,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=_gear_collision_properties(), ), init_state=RigidObjectCfg.InitialStateCfg(pos=(-1.0200, 0.2100, -0.1), rot=(0.0, 0.0, 0.70711, 0.70711)), ) @@ -105,9 +292,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 +311,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=_gear_collision_properties(), ), init_state=RigidObjectCfg.InitialStateCfg(pos=(-1.0200, 0.2100, -0.1), rot=(0.0, 0.0, 0.70711, 0.70711)), ) @@ -130,9 +320,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 +339,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=_gear_collision_properties(), ), init_state=RigidObjectCfg.InitialStateCfg(pos=(-1.0200, 0.2100, -0.1), rot=(0.0, 0.0, 0.70711, 0.70711)), ) @@ -302,11 +495,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 +505,53 @@ 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, + physx_sdf=4, + newton_mjwarp=3, + newton_sdf=3, + newton_hydroelastic=3, + ) + self.sim.render_interval = preset( + default=4, + physx=4, + physx_sdf=4, + newton_mjwarp=3, + newton_sdf=3, + newton_hydroelastic=3, + ) + 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, + ) + + physx_gear_offsets = _PHYSX_GEAR_OFFSETS + newton_gear_offsets = _NEWTON_GEAR_OFFSETS + 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, + ) + 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, + physx_sdf=newton_gear_offsets, + 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..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 @@ -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,144 @@ 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_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"}) + + 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 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 + + 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 == 4_194_304 + 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", + ): + 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 + 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(): + """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_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", + ): + 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"