diff --git a/docs/source/policy_deployment/06_displayport_insertion/displayport_insertion_policy.rst b/docs/source/policy_deployment/06_displayport_insertion/displayport_insertion_policy.rst new file mode 100644 index 00000000000..38f5dcbab6b --- /dev/null +++ b/docs/source/policy_deployment/06_displayport_insertion/displayport_insertion_policy.rst @@ -0,0 +1,803 @@ +.. _walkthrough_dp_insertion: + +Training a DisplayPort Cable Insertion Policy and ROS Deployment +================================================================ + +This tutorial walks you through how to train a DisplayPort plug insertion reinforcement learning (RL) policy that transfers from simulation to a real Flexiv robot. The workflow consists of two main stages: + +1. **Simulation Training in Isaac Lab**: Train the policy in a high-fidelity physics simulation with domain randomization +2. **LEAPP Export and Real Robot Deployment**: Export the trained policy with LEAPP, then deploy on hardware with Isaac ROS / Isaac Manipulator + +This walkthrough covers the key principles and best practices for sim-to-real transfer using Isaac Lab. + +**Supported Robot:** + +- **Flexiv Rizon 4s**: 7-DOF collaborative robot arm with Grav parallel gripper + +**Task Details:** + +The DisplayPort insertion policy operates as follows: + +1. **Initial State**: The policy assumes the DisplayPort plug is already grasped by the gripper at the start of the episode +2. **Input Observations**: The policy receives the pose of the socket insertion point (position and orientation) from a separate perception pipeline +3. **Policy Output**: The policy outputs delta joint positions (incremental changes to arm joint angles) to control the robot and perform the insertion +4. **Task Goal**: Insert the right-angle DisplayPort plug into a fixed socket until the mate point aligns within the success threshold + +**Scope of This Tutorial:** + +This tutorial covers **training and LEAPP export** in Isaac Lab. For the complete on-robot workflow (vision pipeline, robot interface, ROS inference node), refer to the `Isaac ROS Documentation `_ after exporting your policy. + +**Code Layout:** + +The task follows the same structure as the gear assembly deploy environments: + +- ``isaaclab_tasks/contrib/deploy/cable_insertion/displayport_insertion_env_cfg.py`` — shared task MDP (scene, assets, observations, rewards) +- ``isaaclab_tasks/contrib/deploy/cable_insertion/insertion_env.py`` — environment class that logs insertion success metrics during training +- ``isaaclab_tasks/contrib/deploy/cable_insertion/config/displayport_rizon_4s/`` — Flexiv Rizon 4s + Grav robot-specific overrides and gym registrations + +Overview +-------- + +Successful sim-to-real transfer requires addressing three fundamental aspects: + +1. **Input Consistency**: Ensuring the observations your policy receives in simulation match those available on the real robot +2. **System Response Consistency**: Ensuring the robot and environment respond to actions in simulation the same way they do in reality +3. **Output Consistency**: Ensuring any post-processing applied to policy outputs in Isaac Lab is also applied during real-world inference + +When all three aspects are properly addressed, policies trained purely in simulation can achieve robust performance on real hardware without any real-world training data. + +**Debugging Tip**: When your policy fails on the real robot, set up the real robot with the same initial observations as in simulation, then compare how the controller responds. This isolates whether the problem is from observation mismatch (Input Consistency) or physics/controller mismatch (System Response Consistency). + +Asset Quality for Insertion Tasks +---------------------------------- + +For any contact-rich insertion task, **the quality of the plug and socket assets matters more than most other sim-to-real knobs**. DisplayPort insertion in particular operates at very small clearances between plug blades and the socket cavity. If the USD collision geometry, mass properties, or joint behavior are wrong, no amount of reward tuning or domain randomization will produce a policy that transfers well to hardware. + +The current DisplayPort assets in ``display_cable_insertion_assets/`` (``display_port_plug_fixed_sdf.usd`` and ``display_port_socket_fixed_sdf_noprotrusions.usd``) have been iterated extensively and work well for training policies that transfer sim-to-real. Expect significant upfront effort to reach this quality for a new connector or cable type. + +**What to validate before training:** + +1. **Static insertion pose stability**: Load the plug fully inserted into the socket at the goal pose (no robot, no gripper). The plug should remain seated without drifting, jittering, or being ejected by contact forces. Persistent separation or slow creep at the mated pose usually indicates incorrect collision meshes, rest offsets, or mass/inertia. +2. **Collision fidelity at clearance scale**: Blade-to-cavity gaps are sub-millimeter. Convex hulls or coarse meshes often produce false contacts, snagging, or penetration. SDF or carefully authored triangle meshes with tuned ``contact_offset`` / ``rest_offset`` are typically required. +3. **Engagement behavior**: Push the plug through the approach path by hand (or with scripted motion) and confirm contact feels plausible — no explosive pops, no tunneling through the socket wall, no sticky high-friction jamming unless that matches the real connector. +4. **Grasped plug stability**: With the gripper closed at the training grasp width, the plug should not spin or slip unrealistically when the arm moves. Cable mass and plug COM should be representative of the real assembly. +5. **Mate-point alignment**: Verify ``SOCKET_INSERTION_OFFSET``, ``PLUG_INSERTION_OFFSET``, and ``PLUG_GOAL_ROT`` in ``displayport_insertion_env_cfg.py`` match the intended physical mate frame. Reward and success metrics are computed from these offsets; a mismatch here looks like a perception error on the real robot. + +**Practical workflow:** + +1. Fix assets in isolation (drop-test or basic play env with fixed poses) before running full RL training. +2. Compare sim behavior to real hardware video at the same poses — look for drift, bounce, and penetration, not policy success rate. +3. Only after assets pass these checks, tune curriculum, rewards, and domain randomization. + +.. note:: + + Poor asset quality often shows up as policies that learn high training success but fail on hardware with inconsistent contact behavior, or as training that never achieves high ``Metrics/success_rate`` despite reward tuning. Fix the assets first. + +Part 1: Input Consistency +-------------------------- + +The observations your policy receives must be consistent between simulation and reality. This means: + +1. The observation space should only include information available from real sensors +2. Sensor noise and delays should be modeled appropriately + +Using Real-Robot-Available Observations +~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + +Your simulation environment should only use observations that are available on the real robot and not use "privileged" information that would not be available in deployment. The critic receives additional privileged observations (plug pose and joint velocities) to improve value estimation during training, but these are not passed to the actor at deployment time. + + +Observation Specification +^^^^^^^^^^^^^^^^^^^^^^^^^ + +The DisplayPort insertion environment uses proprioceptive and exteroceptive (vision) observations: + +.. list-table:: DisplayPort Insertion Environment Observations (Flexiv Rizon 4s) + :widths: 25 10 25 20 + :header-rows: 1 + + * - Observation + - Dim + - Real-World Source + - Noise + * - ``joint_pos`` (arm only) + - 7 + - Robot controller + - None + * - ``joint_vel`` (arm only, optional) + - 7 + - Robot controller + - None + * - ``socket_pos`` (insertion mate point) + - 3 + - Perception pipeline + - ±10mm + * - ``socket_quat`` + - 4 + - Perception pipeline + - None + +**Recommended shipping configuration** (``NoJointVel`` variants): **14** policy dimensions (7 joint positions + 3 socket position + 4 socket quaternion). + +**Training configuration with joint velocity** (``Grav`` variants without ``NoJointVel``): **21** policy dimensions. + +.. note:: + + **Sim-to-real recommendation: use the NoJointVel variant.** Policies trained with ``joint_vel`` in the actor observation can achieve slightly higher success rates in simulation, but we consistently observe less stable behavior on the real Flexiv robot (jittery motions, inconsistent contact during insertion). For deployment, train and ship with the ``NoJointVel`` environments. + + The ``NoJointVel`` configs remove ``joint_vel`` from the actor observation while keeping it in the critic observation group. This matches deployment setups where joint velocity is not exposed to the policy network but can still help the value function during training. + +**Implementation (base class):** + +.. code-block:: python + + @configclass + class PolicyCfg(ObsGroup): + """Observations for policy group.""" + + joint_pos = ObsTerm( + func=mdp.joint_pos, + params={"asset_cfg": SceneEntityCfg("robot", joint_names=[".*"])}, + ) + joint_vel = ObsTerm( + func=mdp.joint_vel, + params={"asset_cfg": SceneEntityCfg("robot", joint_names=[".*"])}, + ) + socket_pos = ObsTerm( + func=mdp.rigid_object_pos_w, + params={"asset_cfg": SceneEntityCfg("dp_socket"), "offset": SOCKET_INSERTION_OFFSET}, + noise=ResetSampledConstantNoiseModelCfg( + noise_cfg=UniformNoiseCfg(n_min=-0.01, n_max=0.01, operation="add") # ±10mm + ), + ) + socket_quat = ObsTerm( + func=mdp.rigid_object_quat_w, + params={"asset_cfg": SceneEntityCfg("dp_socket")}, + ) + + def __post_init__(self): + self.enable_corruption = True + self.concatenate_terms = True + +**Rizon 4s overrides** (in ``config/displayport_rizon_4s/joint_pos_env_cfg.py``): + +.. code-block:: python + + # Arm joints only — gripper joints are excluded from observations + self.observations.policy.joint_pos.params["asset_cfg"].joint_names = [ + "joint1", "joint2", "joint3", "joint4", "joint5", "joint6", "joint7", + ] + +**Why No Noise for Proprioceptive Observations?** + +As with the gear assembly task, policies trained without noise on proprioceptive observations (joint positions) transfer well to the Flexiv Rizon 4s. The controller provides sufficiently accurate joint state feedback that modeling sensor noise on joint states does not improve sim-to-real transfer for this task. + + +Part 2: System Response Consistency +------------------------------------ + +Once your observations are consistent, ensure the simulated robot and environment respond to actions the same way the real system does. For DisplayPort insertion this involves: + +1. Physics simulation parameters (friction, contact properties, plug/socket collision meshes) +2. Actuator modeling (PD controller gains, effort limits) +3. Domain randomization and curriculum + +Physics Parameter Tuning +~~~~~~~~~~~~~~~~~~~~~~~~~ + +Accurate physics simulation is critical for contact-rich insertion. The DisplayPort plug and socket use SDF collision meshes with high solver iteration counts on the rigid bodies: + +.. code-block:: python + + # From displayport_insertion_env_cfg.py — DisplayPortPlug / DisplayPortSocket + rigid_props=sim_utils.RigidBodyPropertiesCfg( + solver_position_iteration_count=128, + solver_velocity_iteration_count=1, + max_depenetration_velocity=0.5, # plug; socket uses 5.0 + ), + collision_props=sim_utils.CollisionPropertiesCfg( + contact_offset=0.00001, # plug + rest_offset=-0.00005, + ), + +The Flexiv Rizon 4s arm uses lower solver iteration counts for performance, matching the gear assembly Flexiv configuration: + +.. code-block:: python + + # From config/displayport_rizon_4s/joint_pos_env_cfg.py + rigid_props=sim_utils.RigidBodyPropertiesCfg( + disable_gravity=True, + max_depenetration_velocity=5.0, + solver_position_iteration_count=4, + solver_velocity_iteration_count=1, + max_contact_impulse=1e32, + ), + collision_props=sim_utils.CollisionPropertiesCfg( + contact_offset=0.005, + rest_offset=0.0, + ), + +**Friction randomization** (in ``config/displayport_rizon_4s/joint_pos_env_cfg.py``): + +.. code-block:: python + + plug_physics_material = EventTerm( + func=mdp.randomize_rigid_body_material, + mode="startup", + params={ + "asset_cfg": SceneEntityCfg("dp_plug", body_names=".*"), + "static_friction_range": (0.001, 0.001), + "dynamic_friction_range": (0.001, 0.001), + }, + ) + + robot_physics_material = EventTerm( + func=mdp.randomize_rigid_body_material, + mode="startup", + params={ + "asset_cfg": SceneEntityCfg("robot", body_names=".*finger.*"), + "static_friction_range": (0.75, 0.75), + "dynamic_friction_range": (0.75, 0.75), + }, + ), + +Low plug/socket friction reduces sticking during blade engagement. Gripper finger friction is set to match real grasp behavior. + +Actuator Modeling +~~~~~~~~~~~~~~~~~ + +The Rizon 4s uses ``ImplicitActuatorCfg`` with per-joint-group arm tuning from ``FLEXIV_RIZON4S_GRAV_GRIPPER_CFG``, plus dedicated Grav gripper actuators: + +.. code-block:: python + + # Grav gripper actuator configuration + self.scene.robot.actuators["gripper_drive"] = ImplicitActuatorCfg( + joint_names_expr=["finger_joint"], + effort_limit_sim=2.0, + velocity_limit_sim=1.0, + stiffness=2e3, + damping=1e1, + ) + self.scene.robot.actuators["gripper_passive"] = ImplicitActuatorCfg( + joint_names_expr=[".*_knuckle_joint"], + effort_limit_sim=1.0, + velocity_limit_sim=1.0, + stiffness=0.0, + damping=0.0, + ) + +.. note:: + + **Flexiv Rizon 4s**: Domain randomization for actuator gains and joint friction is not included in the Rizon 4s ``EventCfg``. The real-world Flexiv controller is stable and precise enough that the simulation policy transfers without these additional randomizations, consistent with the gear assembly Flexiv setup. + +Action Space Design +~~~~~~~~~~~~~~~~~~~ + +The policy controls only the 7 arm joints using **incremental joint position control**. The gripper is not in the action space — the plug is held at a fixed grasp width for the episode. + +.. code-block:: python + + self.joint_action_scale = 0.025 # ±1.4 degrees per step + + self.actions.arm_action = mdp.RelativeJointPositionActionCfg( + asset_name="robot", + joint_names=["joint1", "joint2", "joint3", "joint4", + "joint5", "joint6", "joint7"], + scale=self.joint_action_scale, + use_zero_offset=True, + ) + +**Action dimension:** 7 + +**Control frequency:** ``sim.dt = 1/240`` s with ``decimation = 8`` → 30 Hz policy rate. + +Domain Randomization Strategy +~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + +**Socket pose randomization** perturbs the fixed socket to cover perception and mounting variation: + +.. code-block:: python + + randomize_socket_pose = EventTerm( + func=mdp.reset_root_state_uniform, + mode="reset", + params={ + "pose_range": { + "x": [-0.01, 0.01], # ±1 cm + "y": [-0.01, 0.01], # ±1 cm + "z": [-0.02, 0.02], # ±2 cm + "roll": [-math.radians(2.0), math.radians(2.0)], + "pitch": [-math.radians(2.0), math.radians(2.0)], + "yaw": [-math.radians(2.0), math.radians(2.0)], + }, + "asset_cfg": SceneEntityCfg("dp_socket"), + }, + ) + +**Plug reset curriculum** starts episodes with the plug near the goal pose and anneals toward farther approach poses: + +.. code-block:: python + + reset_plug_curriculum = EventTerm( + func=mdp.reset_plug_at_goal_curriculum, + mode="reset", + params={ + "at_goal_prob": 0.8, + "at_goal_prob_final": 0.0, + "anneal_start_iter": 0.0, + "anneal_end_iter": 500.0, + "num_steps_per_env": 512, + "insertion_axis": [1.0, 0.0, 0.0], + "at_goal_depth_range": [0.0, 0.015], # 0–15 mm engaged + "approach_depth_range": [0.02, 0.06], # 20–60 mm approach + "normal_pose_range": { + "x": [-0.02, 0.02], + "y": [-0.02, 0.02], + "z": [0.0, 0.0], + }, + }, + ) + +At the start of training, 80% of resets place the plug near the inserted pose; this probability linearly anneals to 0% over 500 training iterations, forcing the policy to learn full approach and insertion. + +**Initial robot pose** is set via inverse kinematics to a grasp pose on the plug at each reset: + +.. code-block:: python + + set_robot_to_grasp_pose = EventTerm( + func=mdp.set_robot_to_object_grasp_pose, + mode="reset", + params={ + "target_object_name": "dp_plug", + "grasp_offset": [0.0025, 0.0, -0.1875], # plug local frame [m] + "end_effector_body_name": "flange", + "num_arm_joints": 7, + }, + ) + +Reward Shaping +~~~~~~~~~~~~~~ + +The environment uses keypoint-based rewards that measure alignment between the plug and socket insertion mate points. Reward terms are defined in ``displayport_insertion_env_cfg.py``: + +- **Keypoint tracking** (``plug_socket_keypoint_tracking``): Penalizes L2 keypoint distance between plug and socket mate frames +- **Exponential keypoint tracking** (``plug_socket_keypoint_tracking_exp``): Dense exponential reward for fine alignment +- **Action rate** (``action_rate_l2``): Penalizes large action changes for smooth motions + +The Rizon 4s config sets the linear and exponential keypoint reward weights to a **1:1 ratio**: + +.. code-block:: python + + self.rewards.plug_socket_keypoint_tracking_exp.weight = abs( + self.rewards.plug_socket_keypoint_tracking.weight + ) + +Terminations +~~~~~~~~~~~~ + +In addition to the episode timeout, the Rizon 4s config terminates early when: + +- **Plug dropped**: End-effector moves more than 15 cm away from the plug grasp point +- **Plug orientation exceeded**: Roll or pitch deviation exceeds 15° relative to the grasp frame + +Training Metrics +~~~~~~~~~~~~~~~~ + +Unlike gear assembly, this task uses a custom environment class (``DisplayportInsertionEnv``) to log insertion metrics to TensorBoard without changing the MDP: + +- ``Metrics/success_rate`` — fraction of environments within the 3 mm mate-point threshold +- ``Metrics/plug_socket_pos_error_m`` — mean mate-point distance +- ``Metrics/plug_socket_keypoint_dist_m`` — mean keypoint distance +- ``Metrics/terminal_success_rate`` — success rate at episode reset + + +Tuning Hyperparameters for Better Performance +---------------------------------------------- + +After asset quality and physics look correct, the following hyperparameters are the main levers for improving training speed, final success rate, and sim-to-real robustness. Defaults below are the shipped Flexiv Rizon 4s values; adjust one group at a time and monitor ``Metrics/success_rate`` in TensorBoard. + +Reward Weights +~~~~~~~~~~~~~~ + +Defined in ``displayport_insertion_env_cfg.py``; the Rizon 4s config overrides the exponential weight in ``joint_pos_env_cfg.py``. + +.. list-table:: Reward hyperparameters + :widths: 35 20 45 + :header-rows: 1 + + * - Parameter + - Default + - Effect + * - ``plug_socket_keypoint_tracking.weight`` + - ``-1.5`` + - Linear penalty on keypoint distance. More negative → stronger pull toward alignment. + * - ``plug_socket_keypoint_tracking_exp.weight`` + - ``1.5`` (matched to linear) + - Exponential bonus near the goal. Increase relative to linear for sharper fine-insertion behavior; decrease if policy is brittle or stalls short of full insertion. + * - ``kp_exp_coeffs`` + - ``[(50, 0.0001), (300, 0.0001), (600, 0.0001), (2000, 0.0001)]`` + - Per-keypoint exponential scales. Higher first values tighten the reward basin around the goal. + * - ``keypoint_scale`` + - ``0.15`` + - Spatial extent of keypoint offsets. Affects how rotation errors contribute relative to translation. + * - ``action_rate.weight`` + - ``-5e-6`` + - Smoothness penalty. More negative → slower, smoother motions; too strong can prevent final insertion force. + +**1:1 linear-to-exponential weighting** (current shipping default): + +.. code-block:: python + + self.rewards.plug_socket_keypoint_tracking_exp.weight = abs( + self.rewards.plug_socket_keypoint_tracking.weight + ) + +If the policy approaches but does not fully seat the plug, try increasing the exponential weight or tightening ``kp_exp_coeffs``. If it rushes and bounces off the socket, increase ``action_rate`` magnitude or reduce the exponential weight. + +Reset Curriculum +~~~~~~~~~~~~~~~~ + +Defined in ``config/displayport_rizon_4s/joint_pos_env_cfg.py`` → ``reset_plug_curriculum``. + +.. list-table:: Curriculum hyperparameters + :widths: 35 20 45 + :header-rows: 1 + + * - Parameter + - Default + - Effect + * - ``at_goal_prob`` / ``at_goal_prob_final`` + - ``0.8`` → ``0.0`` + - Fraction of resets with plug near full insertion. Higher start values make early learning easier; anneal to zero for full approach behavior. + * - ``anneal_end_iter`` + - ``500`` + - Training iterations over which at-goal probability anneals. Extend (e.g. 800–1000) if success rate drops when curriculum gets harder; shorten if training is too slow to reach approach poses. + * - ``at_goal_depth_range`` + - ``[0.0, 0.015]`` m + - How deep the plug starts when sampled "at goal" (0–15 mm engaged). Narrow for fine final-insertion practice; widen slightly if the policy never sees near-mated contacts. + * - ``approach_depth_range`` + - ``[0.02, 0.06]`` m + - Standoff distance when not at goal (20–60 mm). Increase upper bound for harder long-range approach; decrease if the policy struggles to reach the socket mouth. + * - ``normal_pose_range`` + - ±2 cm lateral + - Lateral misalignment when not at goal. Widen for more robustness to perception error; narrow if training fails to converge. + +If ``Metrics/success_rate`` is high early but collapses after iteration ~500, the curriculum may be annealing too aggressively — extend ``anneal_end_iter`` or raise ``at_goal_prob_final`` temporarily. + +Domain Randomization and Observations +~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + +.. list-table:: Randomization hyperparameters + :widths: 35 20 45 + :header-rows: 1 + + * - Parameter + - Default + - Effect + * - ``randomize_socket_pose`` ranges + - ±1 cm XY, ±2 cm Z, ±2° + - Socket pose DR. Widen to match real perception/mount error; narrow if the policy cannot learn a baseline insertion. + * - ``socket_pos`` observation noise + - ±10 mm + - Perception noise on mate point. Increase for more robust real-world pose error; decrease if sim policy is too conservative. + * - Plug/socket friction (startup) + - ``0.001`` + - Low friction reduces unrealistic jamming. Tune only after visual sim-vs-real comparison — wrong friction can dominate insertion feel. + * - Gripper finger friction + - ``0.75`` + - Affects grasp stability during insertion forces. + +Actions and Grasp +~~~~~~~~~~~~~~~~~~~ + +.. list-table:: Action / grasp hyperparameters + :widths: 35 20 45 + :header-rows: 1 + + * - Parameter + - Default + - Effect + * - ``joint_action_scale`` + - ``0.025`` + - Max joint delta per step (~±1.4°). Increase if real robot stiction prevents reaching targets; decrease for finer final alignment. + * - ``grasp_offset`` + - ``[0.0025, 0.0, -0.1875]`` m + - EE-to-plug transform for IK reset. Wrong values cause dropped-plug terminations or misaligned approach. + * - ``hand_hold_width`` / ``hand_close_width`` + - ``-0.05`` / ``-0.155`` rad + - Grav finger_joint grasp command. Adjust if plug slips or is over-compressed during insertion. + +Terminations and Success Metrics +~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + +.. list-table:: Termination / metric hyperparameters + :widths: 35 20 45 + :header-rows: 1 + + * - Parameter + - Default + - Effect + * - ``success_pos_threshold`` + - ``3`` mm + - Mate-point distance counted as success in ``Metrics/success_rate``. Tighten to match real acceptance criteria. + * - ``plug_dropped`` distance threshold + - ``15`` cm + - Early reset if EE leaves plug. Tighten to discourage release; loosen if false positives during large motions. + * - Orientation thresholds (roll/pitch) + - ``15°`` + - Reset if plug tilts excessively relative to grasp frame. + +RL Algorithm (PPO) +~~~~~~~~~~~~~~~~~~ + +Defined in ``config/displayport_rizon_4s/agents/rsl_rl_ppo_cfg.py``. + +.. list-table:: PPO hyperparameters + :widths: 35 20 45 + :header-rows: 1 + + * - Parameter + - Default + - Effect + * - ``max_iterations`` + - ``1500`` + - Total training iterations. Extend if success rate is still climbing at the end. + * - ``num_steps_per_env`` + - ``512`` + - Rollout length per iteration. Affects curriculum annealing rate (tied to ``anneal_end_iter``). + * - ``learning_rate`` + - ``5e-4`` + - PPO learning rate. Reduce if training is unstable; increase if learning is very slow. + * - ``desired_kl`` + - ``0.008`` + - Target KL for adaptive LR schedule. + * - ``init_noise_std`` + - ``1.0`` + - Exploration noise. Lower for fine-tuning a near-working policy. + +**Suggested tuning order:** (1) confirm asset/physics quality, (2) curriculum depth and anneal schedule, (3) linear vs exponential reward balance, (4) socket pose DR and observation noise, (5) action scale, (6) PPO training length. + + +Part 3: Training the Policy in Isaac Lab +----------------------------------------- + +Registered Gym Environments +~~~~~~~~~~~~~~~~~~~~~~~~~~~ + +.. list-table:: Flexiv Rizon 4s DisplayPort Insertion Environments + :widths: 55 45 + :header-rows: 1 + + * - Environment ID + - Purpose + * - ``Isaac-Deploy-DisplayportInsertion-Rizon4s-Grav-NoJointVel-v0`` + - **Training** (recommended for deployment) + * - ``Isaac-Deploy-DisplayportInsertion-Rizon4s-Grav-NoJointVel-Play-v0`` + - Evaluation / visualization + * - ``Isaac-Deploy-DisplayportInsertion-Rizon4s-Grav-NoJointVel-ROS-Inference-v0`` + - ROS / Isaac Manipulator inference metadata + * - ``Isaac-Deploy-DisplayportInsertion-Rizon4s-Grav-v0`` + - Training with joint velocity in actor obs (21-dim) + * - ``Isaac-Deploy-DisplayportInsertion-Rizon4s-Grav-ROS-Inference-v0`` + - ROS inference with joint velocity in actor obs + +Step 1: Visualize the Environment +~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + +Launch training with a small number of environments and visualization enabled to verify the setup: + +.. code-block:: bash + + ./isaaclab.sh train --rl_library rsl_rl \ + --task Isaac-Deploy-DisplayportInsertion-Rizon4s-Grav-NoJointVel-ROS-Inference-v0 \ + --num_envs 4 \ + --visualizer kit + +**What to Expect:** + +In early training, the robot moves the grasped plug toward the socket but will not insert reliably yet. Verify that: + +- The plug is grasped at reset and held throughout the episode +- The socket pose randomization and plug curriculum produce varied starting configurations +- Contact between plug blades and socket looks physically plausible +- With the plug placed in the fully inserted pose (no policy), it stays seated without drift or instability + +Stop training (Ctrl+C) once the environment looks correct, then proceed to full-scale training. + +Step 2: Full-Scale Training with Video Recording +~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + +Launch full training in headless mode with video recording: + +.. code-block:: bash + + ./isaaclab.sh train --rl_library rsl_rl \ + --task Isaac-Deploy-DisplayportInsertion-Rizon4s-Grav-NoJointVel-v0 \ + --num_envs 256 \ + --video --video_length 200 --video_interval 76800 + +**Command breakdown:** + +- ``--num_envs 256``: Runs 256 parallel environments +- ``--video_length 200``: One episode per video (``episode_length_s / (sim.dt * decimation)`` ≈ 200 steps) +- ``--video_interval 76800``: Records a video every 76,800 environment steps (~every 150 iterations with 512 steps/env) + +Training uses a recurrent PPO agent (LSTM, 1500 max iterations, 512 steps per environment). Videos are saved under ``logs/``. + +.. note:: + + **GPU Memory Considerations**: The default configuration uses 4096 environments in the base config but 256 is recommended for most GPUs. The plug and socket SDF collision meshes and high rigid-body solver counts increase GPU memory usage compared to primitive-shape tasks. Reduce ``num_envs`` or ``solver_position_iteration_count`` on the plug/socket assets if you encounter out-of-memory errors. + +**Monitoring Training Progress with TensorBoard:** + +.. code-block:: bash + + ./isaaclab.sh -p -m tensorboard.main --logdir logs/rsl_rl/displayport_insertion_rizon4s + +Monitor ``Metrics/success_rate`` and reward curves to confirm learning. The curriculum anneals over the first 500 iterations — expect success rate to rise as the at-goal reset probability decreases. + +Step 3: Export and Deploy on Real Robot +~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + +**Recommended workflow:** export the trained policy with **LEAPP**, validate the export in simulation, then deploy the LEAPP package with Isaac ROS / Isaac Manipulator on the Flexiv robot. + +Use the **NoJointVel** task for export and deployment so the observation space matches real hardware (14-dim actor input). + +Export with LEAPP (Recommended) +^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ + +`LEAPP `__ (Lightweight Export Annotations for Policy Pipelines) is the **default and recommended** path from a trained checkpoint to real-robot inference. It packages the policy together with input/output semantics (observation ordering, action scaling, recurrent LSTM state) so Isaac ROS deployment does not need to reimplement Isaac Lab preprocessing by hand. + +**Prerequisites:** ``leapp>=0.5.2`` and a trained NoJointVel checkpoint. + +.. code-block:: bash + + ./isaaclab.sh -p -m pip install leapp + +**Export the policy:** + +.. code-block:: bash + + ./isaaclab.sh -p scripts/reinforcement_learning/leapp/rsl_rl/export.py \ + --task Isaac-Deploy-DisplayportInsertion-Rizon4s-Grav-NoJointVel-ROS-Inference-v0 \ + --checkpoint logs/rsl_rl/displayport_insertion_rizon4s//model_.pt + +Use the ``...-ROS-Inference-v0`` task so the traced observation and action layout matches deployment. Replace ```` and ```` with your training log path. + +By default, export artifacts are written next to the checkpoint: + +- Exported model (``.onnx`` by default, or ``.pt`` depending on backend) +- LEAPP metadata YAML describing the policy I/O graph +- Initial recurrent hidden state (``.safetensors``) — this policy uses an LSTM actor +- Pipeline graph visualization (``.png``) + +Useful export flags: + +- ``--export_method onnx-dynamo`` — default ONNX export backend +- ``--validation_steps 5`` — replay traced rollout data to verify the export (recommended; set ``0`` only for debugging) +- ``--export_save_path `` — write artifacts to a custom directory + +See :doc:`Exporting Policies with LEAPP ` for full CLI options, backend choices, and troubleshooting. + +**Validate the LEAPP export in simulation** before real-robot deployment: + +.. code-block:: bash + + ./isaaclab.sh -p scripts/reinforcement_learning/leapp/deploy.py \ + --task Isaac-Deploy-DisplayportInsertion-Rizon4s-Grav-NoJointVel-ROS-Inference-v0 \ + --leapp_model logs/rsl_rl/displayport_insertion_rizon4s// \ + --viz kit + +This runs the packaged policy through the LEAPP deployment path in Isaac Lab and confirms that observation wiring and recurrent state handling survived export. + +**Deploy on hardware:** pass the LEAPP export directory and metadata to your Isaac ROS / Isaac Manipulator workflow. Refer to the `Isaac ROS manipulation DNN policy documentation `_ for on-robot setup. The on-robot pipeline typically includes: + +1. **Perception** — socket pose estimation +2. **Motion planning** — approach trajectory to the insertion station (if used) +3. **Policy inference** — LEAPP-exported policy at control frequency in the ROS inference node +4. **Robot control** — Flexiv low-level joint commands from policy actions + +The ROS inference environment (``Isaac-Deploy-DisplayportInsertion-Rizon4s-Grav-NoJointVel-ROS-Inference-v0``) defines the deployment metadata LEAPP traces during export: + +- ``obs_order``: ``["arm_dof_pos", "socket_pos", "socket_quat"]`` +- ``policy_action_space``: ``"joint"`` +- ``observation_space``: 14 +- ``action_space``: 7 +- ``joint_action_scale``: 0.025 + +Fixed deployment poses for the socket and plug are set in ``config/displayport_rizon_4s/ros_inference_env_cfg.py``. + +Alternative: Raw Checkpoint Deployment +^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ + +For development or legacy Isaac Manipulator setups, you can deploy the RSL-RL checkpoint directly without a LEAPP export step. This path uses the ``.pt`` checkpoint with ``agent.yaml`` and ``env.yaml`` from: + +.. code-block:: text + + logs/rsl_rl/displayport_insertion_rizon4s//model_.pt + +This is **not recommended for shipping** — you must manually ensure observation ordering, action scaling, and LSTM state handling match training. Prefer the LEAPP export path above for production deployment. + + +Troubleshooting +--------------- + +PhysX Collision Stack Overflow +~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + +**Error Message:** + +.. code-block:: text + + PhysX error: PxGpuDynamicsMemoryConfig::collisionStackSize buffer overflow detected + +**Cause:** GPU collision buffer is too small for contact-rich plug/socket interaction across many parallel environments. + +**Solution:** Increase ``gpu_collision_stack_size`` in ``displayport_insertion_env_cfg.py`` (default is ``2**30``): + +.. code-block:: python + + sim: SimulationCfg = SimulationCfg( + physics=PhysxCfg( + gpu_collision_stack_size=2**31, # Increase if overflow persists + gpu_max_rigid_contact_count=2**23, + gpu_max_rigid_patch_count=2**23, + ), + ) + +CUDA Out of Memory +~~~~~~~~~~~~~~~~~~ + +**Solutions (in order of preference):** + +1. Reduce parallel environments: + + .. code-block:: bash + + ./isaaclab.sh train --rl_library rsl_rl \ + --task Isaac-Deploy-DisplayportInsertion-Rizon4s-Grav-NoJointVel-v0 \ + --num_envs 128 + +2. Reduce plug/socket ``solver_position_iteration_count`` in ``displayport_insertion_env_cfg.py`` (trade-off: more penetration) + +3. Disable video recording during training + + +Deterministic Debugging (Play Environment) +------------------------------------------- + +The ``Isaac-Deploy-DisplayportInsertion-Rizon4s-Grav-NoJointVel-Play-v0`` environment disables observation corruption for repeatable evaluation: + +.. code-block:: bash + + ./isaaclab.sh play --rl_library rsl_rl \ + --task Isaac-Deploy-DisplayportInsertion-Rizon4s-Grav-NoJointVel-Play-v0 \ + --num_envs 1 \ + --checkpoint + +To match a specific real-world station layout, edit the workspace constants in ``config/displayport_rizon_4s/joint_pos_env_cfg.py`` (training layout) or ``config/displayport_rizon_4s/ros_inference_env_cfg.py`` (deployment layout): + +.. code-block:: python + + # Training station layout (joint_pos_env_cfg.py) + _GEOMETRY_POS = (0.475, 0.125, 0.06) + _SOCKET_ROT = (0.5, 0.5, 0.5, -0.5) + + # Deployment layout (ros_inference_env_cfg.py) + _DEPLOY_GEOMETRY_POS = (0.476, 0.127, 0.07) + _DEPLOY_SOCKET_ROT = (0.5, 0.5, 0.5, -0.5) + +This environment is useful for: + +- Comparing simulated and real-world policy behavior at a known socket pose +- Verifying plug grasp and approach trajectories before full perception integration +- Debugging insertion failures at a fixed station configuration + + +Further Resources +----------------- + +- Gear Assembly Sim-to-Real Tutorial: :ref:`walkthrough_sim_to_real` +- Exporting Policies with LEAPP: :doc:`/source/policy_deployment/05_leapp/exporting_policies_with_leapp` +- `Isaac ROS Manipulation Documentation `_ +- RL Training Tutorial: :ref:`tutorial-run-rl-training` diff --git a/docs/source/policy_deployment/index.rst b/docs/source/policy_deployment/index.rst index cd89350c177..d10e57910b8 100644 --- a/docs/source/policy_deployment/index.rst +++ b/docs/source/policy_deployment/index.rst @@ -10,6 +10,7 @@ Below, you'll find detailed examples of various policies for training and deploy 01_io_descriptors/io_descriptors_101 02_gear_assembly/gear_assembly_policy + 06_displayport_insertion/displayport_insertion_policy 03_compass_with_NuRec/compass_navigation_policy_with_NuRec 04_reach/reach_policy 05_leapp/exporting_policies_with_leapp diff --git a/source/isaaclab/isaaclab/envs/mdp/commands/velocity_command.py b/source/isaaclab/isaaclab/envs/mdp/commands/velocity_command.py index 930f663f65d..ae4b8766436 100644 --- a/source/isaaclab/isaaclab/envs/mdp/commands/velocity_command.py +++ b/source/isaaclab/isaaclab/envs/mdp/commands/velocity_command.py @@ -17,6 +17,7 @@ from isaaclab.assets import Articulation from isaaclab.managers import CommandTerm from isaaclab.markers import VisualizationMarkers +from isaaclab.utils.leapp.leapp_semantics import TWIST3_ELEMENT_NAMES if TYPE_CHECKING: from isaaclab.envs import ManagerBasedEnv @@ -95,7 +96,7 @@ def __init__(self, cfg: UniformVelocityCommandCfg, env: ManagerBasedEnv): # adds (optional) cmd kind and element names for leapp export # during export, semantic data about this command will be used to annotate the command input self.cfg.cmd_kind = self.cfg.cmd_kind or "command/body/velocity" - self.cfg.element_names = self.cfg.element_names or ["lin_vel_x", "lin_vel_y", "ang_vel_z"] + self.cfg.element_names = self.cfg.element_names or TWIST3_ELEMENT_NAMES def __str__(self) -> str: """Return a string representation of the command generator.""" diff --git a/source/isaaclab/isaaclab/utils/leapp/__init__.pyi b/source/isaaclab/isaaclab/utils/leapp/__init__.pyi index e2b8f497b5f..ad5f6c6ce74 100644 --- a/source/isaaclab/isaaclab/utils/leapp/__init__.pyi +++ b/source/isaaclab/isaaclab/utils/leapp/__init__.pyi @@ -11,6 +11,8 @@ __all__ = [ "POSE6_ELEMENT_NAMES", "POSE7_ELEMENT_NAMES", "QUAT_XYZW_ELEMENT_NAMES", + "TWIST3_ELEMENT_NAMES", + "TWIST6_ELEMENT_NAMES", "WRENCH6_ELEMENT_NAMES", "XYZ_ELEMENT_NAMES", "body_names_resolver", @@ -22,6 +24,7 @@ __all__ = [ "build_command_connection", "build_state_connection", "build_write_connection", + "canonicalize_command_element_names", "joint_names_resolver", "leapp_tensor_semantics", "patch_env_for_export", @@ -38,6 +41,8 @@ from .leapp_semantics import ( POSE6_ELEMENT_NAMES, POSE7_ELEMENT_NAMES, QUAT_XYZW_ELEMENT_NAMES, + TWIST3_ELEMENT_NAMES, + TWIST6_ELEMENT_NAMES, WRENCH6_ELEMENT_NAMES, XYZ_ELEMENT_NAMES, LeappTensorSemantics, @@ -47,6 +52,7 @@ from .leapp_semantics import ( body_quat_resolver, body_wrench_resolver, body_xyz_resolver, + canonicalize_command_element_names, joint_names_resolver, leapp_tensor_semantics, resolve_leapp_element_names, diff --git a/source/isaaclab/isaaclab/utils/leapp/export_annotator.py b/source/isaaclab/isaaclab/utils/leapp/export_annotator.py index 3ca82a66c08..8432aacc972 100644 --- a/source/isaaclab/isaaclab/utils/leapp/export_annotator.py +++ b/source/isaaclab/isaaclab/utils/leapp/export_annotator.py @@ -45,7 +45,7 @@ from isaaclab.assets.articulation.base_articulation import BaseArticulation from isaaclab.managers import ManagerTermBase -from .leapp_semantics import select_element_names +from .leapp_semantics import canonicalize_command_element_names, select_element_names from .proxy import _ArticulationWriteProxy, _DataProxy, _EnvProxy, _ManagerTermProxy from .utils import ( TracedProxyArray, @@ -540,11 +540,17 @@ def wrapped(env, command_name=None, **kwargs): command_cfg = None with suppress(AttributeError, KeyError): command_cfg = env.command_manager.get_term(leapp_input_name).cfg + kind = getattr(command_cfg, "cmd_kind", None) + element_names = canonicalize_command_element_names( + kind, + getattr(command_cfg, "element_names", None), + result, + ) sem = TensorSemantics( name=leapp_input_name, ref=result, - kind=getattr(command_cfg, "cmd_kind", None), - element_names=getattr(command_cfg, "element_names", None), + kind=kind, + element_names=element_names, extra=build_command_connection(leapp_input_name), ) return annotate.input_tensors(task_name, sem) diff --git a/source/isaaclab/isaaclab/utils/leapp/leapp_semantics.py b/source/isaaclab/isaaclab/utils/leapp/leapp_semantics.py index 340291de16a..9791ed2645b 100644 --- a/source/isaaclab/isaaclab/utils/leapp/leapp_semantics.py +++ b/source/isaaclab/isaaclab/utils/leapp/leapp_semantics.py @@ -48,6 +48,23 @@ class LeappTensorSemantics: POSE7_ELEMENT_NAMES: list[str] = ["x", "y", "z", "qx", "qy", "qz", "qw"] POSE6_ELEMENT_NAMES: list[str] = ["x", "y", "z", "angular_x", "angular_y", "angular_z"] WRENCH6_ELEMENT_NAMES: list[str] = ["fx", "fy", "fz", "tx", "ty", "tz"] +TWIST3_ELEMENT_NAMES: list[str] = ["lin_vel_x", "lin_vel_y", "ang_vel_z"] +TWIST6_ELEMENT_NAMES: list[str] = [ + "lin_vel_x", + "lin_vel_y", + "lin_vel_z", + "ang_vel_x", + "ang_vel_y", + "ang_vel_z", +] + +_COMMAND_BODY_VELOCITY_KIND = "command/body/velocity" +_TWIST_ELEMENT_NAME_ALIASES = { + ("lin_x", "lin_y", "ang_z"): TWIST3_ELEMENT_NAMES, + ("lin_x", "lin_y", "lin_z", "ang_x", "ang_y", "ang_z"): TWIST6_ELEMENT_NAMES, + ("linear_x", "linear_y", "angular_z"): TWIST3_ELEMENT_NAMES, + ("linear_x", "linear_y", "linear_z", "angular_x", "angular_y", "angular_z"): TWIST6_ELEMENT_NAMES, +} def select_element_names(names: list[str] | None, indices: Any = None) -> list[str] | None: @@ -67,6 +84,66 @@ def select_element_names(names: list[str] | None, indices: Any = None) -> list[s return None +def canonicalize_command_element_names( + kind: Any, + element_names: list[str] | list[list[str]] | None, + ref: Any | None = None, +) -> list[str] | list[list[str]] | None: + """Return Deploy-compatible element names for command tensors. + + Isaac ROS Deploy's Twist/TwistStamped converters publish body velocity + elements as ``lin_vel_*`` and ``ang_vel_*``. Older export configs used + shorter names such as ``lin_x`` / ``ang_x``. Canonicalize only the + command-body-velocity case so tensor ordering stays unchanged while the + exported metadata matches the runtime converter ``TensorSpec``. + """ + if getattr(kind, "value", kind) != _COMMAND_BODY_VELOCITY_KIND: + return element_names + + width = None + with suppress(AttributeError, IndexError, TypeError): + width = int(ref.shape[-1]) + if width is not None and width not in (3, 6): + raise ValueError(f"LEAPP command/body/velocity input must be 3D or 6D, but tensor width is {width}.") + + if element_names is None: + if width == 3: + return TWIST3_ELEMENT_NAMES + if width == 6: + return TWIST6_ELEMENT_NAMES + return None + + # Command tensors are flat. Nested names cannot match the current Deploy + # Twist converter TensorSpec. + if any(isinstance(name, (list, tuple)) for name in element_names): + raise ValueError("LEAPP command/body/velocity element names must be a flat list.") + + names_tuple = tuple(element_names) + canonical_names = _TWIST_ELEMENT_NAME_ALIASES.get(names_tuple, list(element_names)) + if len(canonical_names) == 3: + expected_names = TWIST3_ELEMENT_NAMES + elif len(canonical_names) == 6: + expected_names = TWIST6_ELEMENT_NAMES + else: + raise ValueError( + "LEAPP command/body/velocity input must have 3 or 6 element names, " + f"but got {len(canonical_names)}." + ) + + if width is not None and len(canonical_names) != width: + raise ValueError( + f"LEAPP command/body/velocity input has {len(canonical_names)} element names, " + f"but tensor width is {width}." + ) + if canonical_names != expected_names: + raise ValueError( + "LEAPP command/body/velocity element names must match Isaac ROS Deploy Twist converter names. " + f"Got {list(element_names)}; expected {expected_names}." + ) + + return list(canonical_names) + + def leapp_tensor_semantics( *, kind: Any = None, diff --git a/source/isaaclab_rl/test/export/test_leapp_proxy.py b/source/isaaclab_rl/test/export/test_leapp_proxy.py index eb91a66167f..0c7fe17dc5a 100644 --- a/source/isaaclab_rl/test/export/test_leapp_proxy.py +++ b/source/isaaclab_rl/test/export/test_leapp_proxy.py @@ -16,7 +16,7 @@ from isaaclab.utils import math as math_utils from isaaclab.utils.leapp import utils as leapp_utils from isaaclab.utils.leapp.export_annotator import ExportPatcher -from isaaclab.utils.leapp.leapp_semantics import InputKindEnum +from isaaclab.utils.leapp.leapp_semantics import InputKindEnum, TWIST6_ELEMENT_NAMES from isaaclab.utils.leapp.proxy import _DataProxy, _EnvProxy @@ -102,3 +102,57 @@ def test_projected_gravity_observation_exports_root_quat_w_input(monkeypatch: py assert semantics.name == "robot_root_quat_w" assert semantics.kind == InputKindEnum.BODY_ROTATION assert semantics.extra == {"isaaclab_connection": "state:robot:root_quat_w"} + + +def test_generated_body_velocity_command_exports_deploy_twist_names(monkeypatch: pytest.MonkeyPatch): + """Test legacy body-velocity command names are exported with Deploy Twist names.""" + annotated_inputs = _capture_leapp_inputs(monkeypatch) + command_tensor = torch.zeros(2, 6, dtype=torch.float32) + + def generated_commands(env, command_name=None, **kwargs): + return command_tensor + + command_cfg = SimpleNamespace( + cmd_kind="command/body/velocity", + element_names=["lin_x", "lin_y", "lin_z", "ang_x", "ang_y", "ang_z"], + ) + command_manager = SimpleNamespace(get_term=lambda name: SimpleNamespace(cfg=command_cfg)) + env = SimpleNamespace(command_manager=command_manager) + term_cfg = SimpleNamespace(params={"command_name": "target_twist"}) + + patcher = ExportPatcher(export_method="onnx-dynamo") + patcher.task_name = "Isaac-Test-Task" + + result = patcher._wrap_generated_commands(generated_commands, term_cfg)(env) + + assert result is command_tensor + assert len(annotated_inputs) == 1 + task_name, semantics = annotated_inputs[0] + assert task_name == "Isaac-Test-Task" + assert semantics.name == "target_twist" + assert semantics.kind == "command/body/velocity" + assert semantics.element_names == [TWIST6_ELEMENT_NAMES] + assert semantics.extra == {"isaaclab_connection": "command:target_twist"} + + +def test_generated_body_velocity_command_rejects_non_deploy_twist_names(monkeypatch: pytest.MonkeyPatch): + """Test unknown body-velocity command names fail before exporting unusable YAML.""" + _capture_leapp_inputs(monkeypatch) + command_tensor = torch.zeros(2, 6, dtype=torch.float32) + + def generated_commands(env, command_name=None, **kwargs): + return command_tensor + + command_cfg = SimpleNamespace( + cmd_kind="command/body/velocity", + element_names=["vx", "vy", "vz", "wx", "wy", "wz"], + ) + command_manager = SimpleNamespace(get_term=lambda name: SimpleNamespace(cfg=command_cfg)) + env = SimpleNamespace(command_manager=command_manager) + term_cfg = SimpleNamespace(params={"command_name": "target_twist"}) + + patcher = ExportPatcher(export_method="onnx-dynamo") + patcher.task_name = "Isaac-Test-Task" + + with pytest.raises(ValueError, match="command/body/velocity element names"): + patcher._wrap_generated_commands(generated_commands, term_cfg)(env) diff --git a/source/isaaclab_tasks/isaaclab_tasks/contrib/deploy/__init__.py b/source/isaaclab_tasks/isaaclab_tasks/contrib/deploy/__init__.py index 3de316ff330..5004a831f4a 100644 --- a/source/isaaclab_tasks/isaaclab_tasks/contrib/deploy/__init__.py +++ b/source/isaaclab_tasks/isaaclab_tasks/contrib/deploy/__init__.py @@ -11,5 +11,7 @@ The deploy module includes: - Reach environments for end-effector pose tracking +- Gear assembly environments for multi-gear insertion tasks +- Cable insertion environments for DisplayPort plug insertion tasks """ diff --git a/source/isaaclab_tasks/isaaclab_tasks/contrib/deploy/cable_insertion/__init__.py b/source/isaaclab_tasks/isaaclab_tasks/contrib/deploy/cable_insertion/__init__.py new file mode 100644 index 00000000000..b1d487e2b67 --- /dev/null +++ b/source/isaaclab_tasks/isaaclab_tasks/contrib/deploy/cable_insertion/__init__.py @@ -0,0 +1,6 @@ +# Copyright (c) 2022-2026, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +"""DisplayPort cable insertion environments.""" diff --git a/source/isaaclab_tasks/isaaclab_tasks/contrib/deploy/cable_insertion/config/__init__.py b/source/isaaclab_tasks/isaaclab_tasks/contrib/deploy/cable_insertion/config/__init__.py new file mode 100644 index 00000000000..1332dd34866 --- /dev/null +++ b/source/isaaclab_tasks/isaaclab_tasks/contrib/deploy/cable_insertion/config/__init__.py @@ -0,0 +1,6 @@ +# Copyright (c) 2022-2026, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +"""Configurations for arm-based cable insertion environments.""" diff --git a/source/isaaclab_tasks/isaaclab_tasks/contrib/deploy/cable_insertion/config/displayport_rizon_4s/__init__.py b/source/isaaclab_tasks/isaaclab_tasks/contrib/deploy/cable_insertion/config/displayport_rizon_4s/__init__.py new file mode 100644 index 00000000000..38f6be2c63c --- /dev/null +++ b/source/isaaclab_tasks/isaaclab_tasks/contrib/deploy/cable_insertion/config/displayport_rizon_4s/__init__.py @@ -0,0 +1,84 @@ +# 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 + +import gymnasium as gym + +from . import agents + +_INSERTION_ENV_ENTRY = ( + "isaaclab_tasks.contrib.deploy.cable_insertion.insertion_env:DisplayportInsertionEnv" +) + +## +# Register Gym environments. +## + +# Flexiv Rizon 4s - Joint space +gym.register( + id="Isaac-Deploy-DisplayportInsertion-Rizon4s-Grav-v0", + entry_point=_INSERTION_ENV_ENTRY, + disable_env_checker=True, + kwargs={ + "env_cfg_entry_point": f"{__name__}.joint_pos_env_cfg:Rizon4sGravDisplayportInsertionEnvCfg", + "rsl_rl_cfg_entry_point": f"{agents.__name__}.rsl_rl_ppo_cfg:Rizon4sGravDisplayportInsertionRNNPPORunnerCfg", + }, +) + +# Flexiv Rizon 4s - Joint space Play +gym.register( + id="Isaac-Deploy-DisplayportInsertion-Rizon4s-Grav-Play-v0", + entry_point=_INSERTION_ENV_ENTRY, + disable_env_checker=True, + kwargs={ + "env_cfg_entry_point": f"{__name__}.joint_pos_env_cfg:Rizon4sGravDisplayportInsertionEnvCfg_PLAY", + "rsl_rl_cfg_entry_point": f"{agents.__name__}.rsl_rl_ppo_cfg:Rizon4sGravDisplayportInsertionRNNPPORunnerCfg", + }, +) + +# Flexiv Rizon 4s - Joint space without joint velocity +gym.register( + id="Isaac-Deploy-DisplayportInsertion-Rizon4s-Grav-NoJointVel-v0", + entry_point=_INSERTION_ENV_ENTRY, + disable_env_checker=True, + kwargs={ + "env_cfg_entry_point": f"{__name__}.joint_pos_env_cfg:Rizon4sGravDisplayportInsertionNoJointVelEnvCfg", + "rsl_rl_cfg_entry_point": f"{agents.__name__}.rsl_rl_ppo_cfg:Rizon4sGravDisplayportInsertionRNNPPORunnerCfg", + }, +) + +# Flexiv Rizon 4s - Joint space without joint velocity Play +gym.register( + id="Isaac-Deploy-DisplayportInsertion-Rizon4s-Grav-NoJointVel-Play-v0", + entry_point=_INSERTION_ENV_ENTRY, + disable_env_checker=True, + kwargs={ + "env_cfg_entry_point": f"{__name__}.joint_pos_env_cfg:Rizon4sGravDisplayportInsertionNoJointVelEnvCfg_PLAY", + "rsl_rl_cfg_entry_point": f"{agents.__name__}.rsl_rl_ppo_cfg:Rizon4sGravDisplayportInsertionRNNPPORunnerCfg", + }, +) + +# Flexiv Rizon 4s - Joint space ROS Inference without joint velocity +gym.register( + id="Isaac-Deploy-DisplayportInsertion-Rizon4s-Grav-NoJointVel-ROS-Inference-v0", + entry_point=_INSERTION_ENV_ENTRY, + disable_env_checker=True, + kwargs={ + "env_cfg_entry_point": ( + f"{__name__}.ros_inference_env_cfg:Rizon4sGravDisplayportInsertionNoJointVelROSInferenceEnvCfg" + ), + "rsl_rl_cfg_entry_point": f"{agents.__name__}.rsl_rl_ppo_cfg:Rizon4sGravDisplayportInsertionRNNPPORunnerCfg", + }, +) + +# Flexiv Rizon 4s - Joint space ROS Inference +gym.register( + id="Isaac-Deploy-DisplayportInsertion-Rizon4s-Grav-ROS-Inference-v0", + entry_point=_INSERTION_ENV_ENTRY, + disable_env_checker=True, + kwargs={ + "env_cfg_entry_point": f"{__name__}.ros_inference_env_cfg:Rizon4sGravDisplayportInsertionROSInferenceEnvCfg", + "rsl_rl_cfg_entry_point": f"{agents.__name__}.rsl_rl_ppo_cfg:Rizon4sGravDisplayportInsertionRNNPPORunnerCfg", + }, +) diff --git a/source/isaaclab_tasks/isaaclab_tasks/contrib/deploy/cable_insertion/config/displayport_rizon_4s/agents/__init__.py b/source/isaaclab_tasks/isaaclab_tasks/contrib/deploy/cable_insertion/config/displayport_rizon_4s/agents/__init__.py new file mode 100644 index 00000000000..6fc9a6577ee --- /dev/null +++ b/source/isaaclab_tasks/isaaclab_tasks/contrib/deploy/cable_insertion/config/displayport_rizon_4s/agents/__init__.py @@ -0,0 +1,6 @@ +# Copyright (c) 2022-2026, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +"""Package containing agent configurations for RL with Rizon 4s for DisplayPort insertion.""" diff --git a/source/isaaclab_tasks/isaaclab_tasks/contrib/deploy/cable_insertion/config/displayport_rizon_4s/agents/rsl_rl_ppo_cfg.py b/source/isaaclab_tasks/isaaclab_tasks/contrib/deploy/cable_insertion/config/displayport_rizon_4s/agents/rsl_rl_ppo_cfg.py new file mode 100644 index 00000000000..7dcb69667ec --- /dev/null +++ b/source/isaaclab_tasks/isaaclab_tasks/contrib/deploy/cable_insertion/config/displayport_rizon_4s/agents/rsl_rl_ppo_cfg.py @@ -0,0 +1,49 @@ +# 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.utils.configclass import configclass + +from isaaclab_rl.rsl_rl import RslRlOnPolicyRunnerCfg, RslRlPpoActorCriticRecurrentCfg, RslRlPpoAlgorithmCfg + + +@configclass +class Rizon4sGravDisplayportInsertionRNNPPORunnerCfg(RslRlOnPolicyRunnerCfg): + num_steps_per_env = 512 + max_iterations = 1500 + save_interval = 50 + experiment_name = "displayport_insertion_rizon4s" + clip_actions = 1.0 + resume = False + obs_groups = { + "policy": ["policy"], + "critic": ["critic"], + } + policy = RslRlPpoActorCriticRecurrentCfg( + state_dependent_std=True, + init_noise_std=1.0, + actor_obs_normalization=True, + critic_obs_normalization=True, + actor_hidden_dims=[256, 128, 64], + critic_hidden_dims=[256, 128, 64], + noise_std_type="log", + activation="elu", + rnn_type="lstm", + rnn_hidden_dim=256, + rnn_num_layers=2, + ) + algorithm = RslRlPpoAlgorithmCfg( + value_loss_coef=1.0, + use_clipped_value_loss=True, + clip_param=0.2, + entropy_coef=0.0, + num_learning_epochs=8, + num_mini_batches=16, + learning_rate=5.0e-4, + schedule="adaptive", + gamma=0.99, + lam=0.95, + desired_kl=0.008, + max_grad_norm=1.0, + ) diff --git a/source/isaaclab_tasks/isaaclab_tasks/contrib/deploy/cable_insertion/config/displayport_rizon_4s/joint_pos_env_cfg.py b/source/isaaclab_tasks/isaaclab_tasks/contrib/deploy/cable_insertion/config/displayport_rizon_4s/joint_pos_env_cfg.py new file mode 100644 index 00000000000..cb16b33c10f --- /dev/null +++ b/source/isaaclab_tasks/isaaclab_tasks/contrib/deploy/cable_insertion/config/displayport_rizon_4s/joint_pos_env_cfg.py @@ -0,0 +1,428 @@ +# 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 + +"""Joint-space DisplayPort insertion environment for Flexiv Rizon 4S + Grav gripper.""" + +import math + +import torch + +import isaaclab.sim as sim_utils +from isaaclab.actuators import ImplicitActuatorCfg +from isaaclab.assets import ArticulationCfg, RigidObjectCfg +from isaaclab.managers import EventTermCfg as EventTerm +from isaaclab.managers import SceneEntityCfg +from isaaclab.managers import TerminationTermCfg as DoneTerm +from isaaclab.utils.configclass import configclass + +import isaaclab_tasks.contrib.deploy.mdp as mdp +import isaaclab_tasks.contrib.deploy.mdp.terminations as cable_terminations +from isaaclab_tasks.contrib.deploy.cable_insertion.displayport_insertion_env_cfg import ( + PLUG_GOAL_ROT, + PLUG_INSERTION_OFFSET, + SOCKET_INSERTION_OFFSET, + DisplayportInsertionEnvCfg, + compute_plug_pose, + compute_socket_root, +) + +# DisplayPort insertion station layout in the Flexiv workspace. +# _GEOMETRY_POS is the desired insertion (mate) point; _SOCKET_ROT orients the socket opening up. +_GEOMETRY_POS = (0.475, 0.125, 0.06) +_SOCKET_ROT = (0.5, 0.5, 0.5, -0.5) # opening faces +Z (top-down insertion) +_PLUG_CLEARANCE_Z = 0.068 # vertical clearance between plug and socket at reset + +_SOCKET_ROOT = compute_socket_root(_GEOMETRY_POS, _SOCKET_ROT) +_PLUG_ROOT, _PLUG_ROT = compute_plug_pose( + _GEOMETRY_POS, + _SOCKET_ROT, + z_clearance=_PLUG_CLEARANCE_Z, +) + +# Blade engagement along the insertion axis used by the at-goal curriculum [m]. +_INSERTION_LENGTH = 0.011 + +## +# Pre-defined configs +## +from isaaclab_assets import FLEXIV_RIZON4S_GRAV_GRIPPER_CFG # isort: skip + + +## +# Gripper-specific helper functions +## + + +def set_finger_joint_pos_grav( + joint_pos: torch.Tensor, + reset_ind_joint_pos: list[int], + finger_joints: list[int], + finger_joint_position: float, +): + """Set finger joint positions for Grav gripper. + + 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 + """ + 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)}") + + # 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 + + # 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 + + +## +# Environment configuration +## + + +@configclass +class EventCfg: + """Configuration for events.""" + + plug_physics_material = EventTerm( + func=mdp.randomize_rigid_body_material, + mode="startup", + params={ + "asset_cfg": SceneEntityCfg("dp_plug", body_names=".*"), + "static_friction_range": (0.001, 0.001), + "dynamic_friction_range": (0.001, 0.001), + "restitution_range": (0.0, 0.0), + "num_buckets": 16, + }, + ) + + socket_physics_material = EventTerm( + func=mdp.randomize_rigid_body_material, + mode="startup", + params={ + "asset_cfg": SceneEntityCfg("dp_socket", body_names=".*"), + "static_friction_range": (0.001, 0.001), + "dynamic_friction_range": (0.001, 0.001), + "restitution_range": (0.0, 0.0), + "num_buckets": 16, + }, + ) + + robot_physics_material = EventTerm( + func=mdp.randomize_rigid_body_material, + mode="startup", + params={ + "asset_cfg": SceneEntityCfg("robot", body_names=".*finger.*"), + "static_friction_range": (0.75, 0.75), + "dynamic_friction_range": (0.75, 0.75), + "restitution_range": (0.0, 0.0), + "num_buckets": 16, + }, + ) + + reset_all = EventTerm(func=mdp.reset_scene_to_default, mode="reset") + + randomize_socket_pose = EventTerm( + func=mdp.reset_root_state_uniform, + mode="reset", + params={ + "pose_range": { + "x": [-0.01, 0.01], + "y": [-0.01, 0.01], + "z": [-0.02, 0.02], + "roll": [-math.radians(2.0), math.radians(2.0)], # 2 degree + "pitch": [-math.radians(2.0), math.radians(2.0)], # 2 degree + "yaw": [-math.radians(2.0), math.radians(2.0)], # 2 degree + }, + "velocity_range": {}, + "asset_cfg": SceneEntityCfg("dp_socket"), + }, + ) + + reset_plug_curriculum = EventTerm( + func=mdp.reset_plug_at_goal_curriculum, + mode="reset", + params={ + "plug_cfg": SceneEntityCfg("dp_plug"), + "socket_cfg": SceneEntityCfg("dp_socket"), + "at_goal_prob": 0.8, + "at_goal_prob_final": 0.0, + "anneal_start_iter": 0.0, + "anneal_end_iter": 500.0, + "num_steps_per_env": 512, + "insertion_axis": [1.0, 0.0, 0.0], + "insertion_length": _INSERTION_LENGTH, + "at_goal_depth_range": [0.0, 0.015], + "approach_depth_range": [0.02, 0.06], + "socket_insertion_offset": SOCKET_INSERTION_OFFSET, + "plug_insertion_offset": PLUG_INSERTION_OFFSET, + "goal_rot": list(PLUG_GOAL_ROT), + "normal_pose_range": { + "x": [-0.02, 0.02], + "y": [-0.02, 0.02], + "z": [0.0, 0.0], + }, + }, + ) + + set_robot_to_grasp_pose = EventTerm( + func=mdp.set_robot_to_object_grasp_pose, + mode="reset", + params={ + "robot_asset_cfg": SceneEntityCfg("robot"), + "pos_randomization_range": {"x": [-0.0, 0.0], "y": [-0.0, 0.0], "z": [-0.0, 0.0]}, + "target_object_name": "dp_plug", + "grasp_offset": [0.0, 0.0, 0.0], + }, + ) + + +@configclass +class TerminationsCfg: + """Configuration for termination terms.""" + + time_out = DoneTerm(func=mdp.time_out, time_out=True) + + plug_dropped = DoneTerm( + func=cable_terminations.reset_when_plug_dropped, + params={ + "robot_asset_cfg": SceneEntityCfg("robot"), + "plug_asset_cfg": SceneEntityCfg("dp_plug"), + "distance_threshold": 0.15, + "end_effector_body_name": "link7", + "grasp_offset": [0.0, 0.0, 0.0], + "grasp_rot_offset": [0.0, 0.0, 0.0, 1.0], + }, + ) + + plug_orientation_exceeded = DoneTerm( + func=cable_terminations.reset_when_plug_orientation_exceeded, + params={ + "robot_asset_cfg": SceneEntityCfg("robot"), + "plug_asset_cfg": SceneEntityCfg("dp_plug"), + "roll_threshold_deg": 15.0, + "pitch_threshold_deg": 15.0, + "yaw_threshold_deg": 180.0, + "end_effector_body_name": "link7", + "grasp_rot_offset": [0.0, 0.0, 0.0, 1.0], + }, + ) + + +@configclass +class Rizon4sGravDisplayportInsertionEnvCfg(DisplayportInsertionEnvCfg): + """Configuration for Flexiv Rizon 4s with Grav Gripper DisplayPort insertion. + + The Flexiv Rizon 4s is a 7-DOF collaborative robot arm equipped with the + Flexiv Grav parallel gripper for DisplayPort plug insertion tasks. + """ + + def __post_init__(self): + # post init of parent + super().__post_init__() + + # Match exponential keypoint reward weight to the linear term (1:1 weighting) + self.rewards.plug_socket_keypoint_tracking_exp.weight = abs(self.rewards.plug_socket_keypoint_tracking.weight) + + # Robot-specific parameters for Flexiv Rizon 4s with Grav gripper + self.end_effector_body_name = "flange" # End effector body name for IK + self.num_arm_joints = 7 # Number of arm joints (Rizon 4s has 7 DOF) + # Grasp offset in the DisplayPort plug's local frame [m] + self.grasp_offset = [0.0025, 0.0, -0.1875] + # Rotation offset for grasp pose (quaternion [x, y, z, w]) + self.grasp_rot_offset = [0.0, 0.0, 0.0, 1.0] + self.gripper_joint_setter_func = set_finger_joint_pos_grav # Grav gripper joint setter function + + # Plug orientation termination thresholds (in degrees) + self.plug_orientation_roll_threshold_deg = 15.0 # Maximum allowed roll deviation + self.plug_orientation_pitch_threshold_deg = 15.0 # Maximum allowed pitch deviation + self.plug_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 = TerminationsCfg() + + # Update termination thresholds from config + self.terminations.plug_orientation_exceeded.params["roll_threshold_deg"] = ( + self.plug_orientation_roll_threshold_deg + ) + self.terminations.plug_orientation_exceeded.params["pitch_threshold_deg"] = ( + self.plug_orientation_pitch_threshold_deg + ) + self.terminations.plug_orientation_exceeded.params["yaw_threshold_deg"] = ( + self.plug_orientation_yaw_threshold_deg + ) + + # Action configuration for Rizon 4s arm + self.joint_action_scale = 0.025 + _arm_joint_names = ["joint1", "joint2", "joint3", "joint4", "joint5", "joint6", "joint7"] + self.actions.arm_action = mdp.DeployRelativeJointPositionActionCfg( + asset_name="robot", + joint_names=_arm_joint_names, + scale=self.joint_action_scale, + use_zero_offset=True, + ) + + # Switch robot to Flexiv Rizon 4s with Grav gripper + 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, + max_depenetration_velocity=5.0, + linear_damping=0.0, + angular_damping=0.0, + max_linear_velocity=1000.0, + max_angular_velocity=3666.0, + enable_gyroscopic_forces=True, + solver_position_iteration_count=4, + 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, + ), + collision_props=sim_utils.CollisionPropertiesCfg(contact_offset=0.005, rest_offset=0.0), + ), + # Joint positions for the DisplayPort insertion station home pose + init_state=ArticulationCfg.InitialStateCfg( + joint_pos={ + "joint1": math.radians(32.44), + "joint2": math.radians(-16.71), + "joint3": math.radians(-5.69), + "joint4": math.radians(128.38), + "joint5": math.radians(6.74), + "joint6": math.radians(55.95), + "joint7": math.radians(111.54), + }, + pos=(0.0, 0.0, 0.0), + rot=(0.0, 0.0, 0.0, 1.0), + ), + ) + + # Grav gripper actuator configuration + self.scene.robot.actuators["gripper_drive"] = ImplicitActuatorCfg( + joint_names_expr=["finger_joint"], + effort_limit_sim=2.0, + velocity_limit_sim=1.0, + stiffness=2e3, + damping=1e1, + ) + + # Passive/mimic joints in the gripper - set to zero stiffness/damping + self.scene.robot.actuators["gripper_passive"] = ImplicitActuatorCfg( + joint_names_expr=[".*_knuckle_joint"], + effort_limit_sim=1.0, + velocity_limit_sim=1.0, + stiffness=0.0, + damping=0.0, + ) + + # Override socket/plug initial states for the DisplayPort insertion station + self.scene.dp_socket.init_state = RigidObjectCfg.InitialStateCfg( + pos=_SOCKET_ROOT, + rot=_SOCKET_ROT, + ) + self.scene.dp_plug.init_state = RigidObjectCfg.InitialStateCfg( + pos=_PLUG_ROOT, + rot=_PLUG_ROT, + ) + + # Grasp widths for Grav gripper (raw radian values for finger_joint) + self.hand_grasp_width = 0.3 + self.hand_hold_width = -0.05 + self.hand_close_width = -0.155 + + # Populate event term parameters + 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["grasp_offset"] = self.grasp_offset + self.events.set_robot_to_grasp_pose.params["gripper_joint_setter_func"] = self.gripper_joint_setter_func + self.events.set_robot_to_grasp_pose.params["max_iterations"] = 150 + + # Populate termination term parameters + self.terminations.plug_dropped.params["end_effector_body_name"] = self.end_effector_body_name + self.terminations.plug_dropped.params["grasp_offset"] = self.grasp_offset + self.terminations.plug_dropped.params["grasp_rot_offset"] = self.grasp_rot_offset + + self.terminations.plug_orientation_exceeded.params["end_effector_body_name"] = self.end_effector_body_name + self.terminations.plug_orientation_exceeded.params["grasp_rot_offset"] = self.grasp_rot_offset + + +@configclass +class Rizon4sGravDisplayportInsertionEnvCfg_PLAY(Rizon4sGravDisplayportInsertionEnvCfg): + """Play configuration for Flexiv Rizon 4s DisplayPort insertion.""" + + def __post_init__(self): + super().__post_init__() + self.scene.num_envs = 50 + self.scene.env_spacing = 2.5 + self.observations.policy.enable_corruption = False + + +@configclass +class Rizon4sGravDisplayportInsertionNoJointVelEnvCfg(Rizon4sGravDisplayportInsertionEnvCfg): + """DisplayPort insertion without joint velocity in the policy observation. + + The critic retains joint velocity as privileged information for the value function. + """ + + def __post_init__(self): + super().__post_init__() + # Remove joint velocity from the actor observation group + self.observations.policy.joint_vel = None + + +@configclass +class Rizon4sGravDisplayportInsertionNoJointVelEnvCfg_PLAY(Rizon4sGravDisplayportInsertionNoJointVelEnvCfg): + """Play configuration for the no-joint-velocity joint-space variant.""" + + def __post_init__(self): + super().__post_init__() + self.scene.num_envs = 50 + self.scene.env_spacing = 2.5 + self.observations.policy.enable_corruption = False diff --git a/source/isaaclab_tasks/isaaclab_tasks/contrib/deploy/cable_insertion/config/displayport_rizon_4s/ros_inference_env_cfg.py b/source/isaaclab_tasks/isaaclab_tasks/contrib/deploy/cable_insertion/config/displayport_rizon_4s/ros_inference_env_cfg.py new file mode 100644 index 00000000000..8ed1d7ca5ef --- /dev/null +++ b/source/isaaclab_tasks/isaaclab_tasks/contrib/deploy/cable_insertion/config/displayport_rizon_4s/ros_inference_env_cfg.py @@ -0,0 +1,133 @@ +# 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 + +import math + +from isaaclab.assets import RigidObjectCfg +from isaaclab.utils.configclass import configclass + +from isaaclab_tasks.contrib.deploy.cable_insertion.displayport_insertion_env_cfg import ( + compute_plug_pose, + compute_socket_root, +) + +from .joint_pos_env_cfg import Rizon4sGravDisplayportInsertionEnvCfg + +# Deployment socket/plug poses for the physical DisplayPort insertion station. +_DEPLOY_GEOMETRY_POS = (0.476, 0.127, 0.07) +_DEPLOY_SOCKET_ROT = (0.5, 0.5, 0.5, -0.5) +_DEPLOY_PLUG_CLEARANCE_Z = 0.068 + +_DEPLOY_SOCKET_ROOT = compute_socket_root(_DEPLOY_GEOMETRY_POS, _DEPLOY_SOCKET_ROT) +_DEPLOY_PLUG_ROOT, _DEPLOY_PLUG_ROT = compute_plug_pose( + _DEPLOY_GEOMETRY_POS, + _DEPLOY_SOCKET_ROT, + z_clearance=_DEPLOY_PLUG_CLEARANCE_Z, +) + + +@configclass +class Rizon4sGravDisplayportInsertionROSInferenceEnvCfg(Rizon4sGravDisplayportInsertionEnvCfg): + """Configuration for ROS inference with Flexiv Rizon 4s and Grav gripper. + + This configuration: + - Exposes variables needed for ROS inference + - Overrides robot and plug/socket initial poses for fixed/deterministic setup + """ + + def __post_init__(self): + # post init of parent + super().__post_init__() + + # Variables used by Isaac Manipulator for on robot inference + # These parameters allow the ROS inference node to validate environment configuration, + # perform checks during inference, and correctly interpret observations and actions. + self.obs_order = ["arm_dof_pos", "arm_dof_vel", "socket_pos", "socket_quat"] + self.policy_action_space = "joint" + # Use inherited joint names from parent's observation configuration + self.arm_joint_names = self.observations.policy.joint_pos.params["asset_cfg"].joint_names + # Use inherited num_arm_joints from parent + self.action_space = self.num_arm_joints + # State: 7 joint pos + 7 joint vel + 3 socket pos + 4 socket quat + 3 plug pos + 4 plug quat = 28 + self.state_space = 28 + # Observation: 7 joint pos + 7 joint vel + 3 socket pos + 4 socket quat = 21 + self.observation_space = 21 + + # Set joint_action_scale from the existing arm_action.scale + self.joint_action_scale = self.actions.arm_action.scale + + # Dynamically generate action_scale_joint_space based on action_space + self.action_scale_joint_space = [self.joint_action_scale] * self.action_space + + # Override robot initial pose for ROS inference (fixed pose, no randomization) + self.scene.robot.init_state.pos = (0.0, 0.0, 0.0) + self.scene.robot.init_state.rot = (0.0, 0.0, 0.0, 1.0) # Identity quaternion (x, y, z, w) + self.scene.robot.init_state.joint_pos = { + "joint1": math.radians(32.44), + "joint2": math.radians(-16.71), + "joint3": math.radians(-5.69), + "joint4": math.radians(128.38), + "joint5": math.radians(6.74), + "joint6": math.radians(55.95), + "joint7": math.radians(111.54), + } + + # Override socket/plug initial poses (fixed poses for ROS inference) + self.scene.dp_socket.init_state = RigidObjectCfg.InitialStateCfg( + pos=_DEPLOY_SOCKET_ROOT, + rot=_DEPLOY_SOCKET_ROT, + ) + + self.scene.dp_plug.init_state = RigidObjectCfg.InitialStateCfg( + pos=_DEPLOY_PLUG_ROOT, + rot=_DEPLOY_PLUG_ROT, + ) + + self.events.set_robot_to_grasp_pose.params["max_iterations"] = 150 + + # Fixed asset parameters for ROS inference - derived from configuration + # These parameters are used by the ROS inference node to validate the environment setup + # and apply appropriate noise models for robust real-world deployment. + self.fixed_asset_init_pos_center = list(_DEPLOY_GEOMETRY_POS) + + pose_range = self.events.randomize_socket_pose.params["pose_range"] + self.fixed_asset_init_pos_range = [ + pose_range["x"][1], # max value + pose_range["y"][1], # max value + pose_range["z"][1], # max value + ] + # Orientation in degrees for the vertical table-top Flexiv mount + self.fixed_asset_init_orn_deg = [0.0, 0.0, 0.0] + # Derive orientation range from parent's pose_range (radians to degrees) + self.fixed_asset_init_orn_deg_range = [ + math.degrees(pose_range["roll"][1]), # convert radians to degrees + math.degrees(pose_range["pitch"][1]), + math.degrees(pose_range["yaw"][1]), + ] + # Derive observation noise level from parent's socket_pos noise configuration + socket_pos_noise = self.observations.policy.socket_pos.noise.noise_cfg.n_max + self.fixed_asset_pos_obs_noise_level = [ + socket_pos_noise, + socket_pos_noise, + socket_pos_noise, + ] + + +@configclass +class Rizon4sGravDisplayportInsertionNoJointVelROSInferenceEnvCfg(Rizon4sGravDisplayportInsertionROSInferenceEnvCfg): + """ROS inference configuration without joint velocity in the policy observation.""" + + def __post_init__(self): + super().__post_init__() + + # Remove joint velocity from the actor observation group + self.observations.policy.joint_vel = None + + # Update Isaac Manipulator metadata for the velocity-free actor + self.obs_order = ["arm_dof_pos", "socket_pos", "socket_quat"] + # Observation: 7 joint pos + 3 socket pos + 4 socket quat = 14 + self.observation_space = 14 + # State (critic) is unchanged: 7 jpos + 7 jvel + 3 socket pos + 4 socket quat + 3 plug pos + 4 plug quat = 28 + self.state_space = 28 diff --git a/source/isaaclab_tasks/isaaclab_tasks/contrib/deploy/cable_insertion/display_cable_insertion_assets/display_port_plug_fixed_sdf.usd b/source/isaaclab_tasks/isaaclab_tasks/contrib/deploy/cable_insertion/display_cable_insertion_assets/display_port_plug_fixed_sdf.usd new file mode 100644 index 00000000000..e3e65d95be9 --- /dev/null +++ b/source/isaaclab_tasks/isaaclab_tasks/contrib/deploy/cable_insertion/display_cable_insertion_assets/display_port_plug_fixed_sdf.usd @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:b9f1b7fc6f1ea314f9a43a5967836dcb3f35b27b163885b179efc45c7a1cd52d +size 1885641 diff --git a/source/isaaclab_tasks/isaaclab_tasks/contrib/deploy/cable_insertion/display_cable_insertion_assets/display_port_socket_fixed_sdf_noprotrusions.usd b/source/isaaclab_tasks/isaaclab_tasks/contrib/deploy/cable_insertion/display_cable_insertion_assets/display_port_socket_fixed_sdf_noprotrusions.usd new file mode 100644 index 00000000000..da586695e9d --- /dev/null +++ b/source/isaaclab_tasks/isaaclab_tasks/contrib/deploy/cable_insertion/display_cable_insertion_assets/display_port_socket_fixed_sdf_noprotrusions.usd @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:06f2c4d67bda068d9d5f61e0c2e2c35e7812d4cc574521c1ce34d9b280fdca9a +size 1460990 diff --git a/source/isaaclab_tasks/isaaclab_tasks/contrib/deploy/cable_insertion/displayport_insertion_env_cfg.py b/source/isaaclab_tasks/isaaclab_tasks/contrib/deploy/cable_insertion/displayport_insertion_env_cfg.py new file mode 100644 index 00000000000..48777026dea --- /dev/null +++ b/source/isaaclab_tasks/isaaclab_tasks/contrib/deploy/cable_insertion/displayport_insertion_env_cfg.py @@ -0,0 +1,364 @@ +# 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 + +"""Base RL environment for inserting a DisplayPort plug into a socket. + +Targets the right-angle DisplayPort plug/socket assets in +``display_cable_insertion_assets``. Assets load with plain +:class:`~isaaclab.sim.UsdFileCfg` at ``scale=(1,1,1)``. +""" + +import os +from dataclasses import MISSING + +from isaaclab_physx.physics import PhysxCfg + +import isaaclab.sim as sim_utils +from isaaclab.assets import ArticulationCfg, AssetBaseCfg, RigidObjectCfg +from isaaclab.envs import ManagerBasedRLEnvCfg +from isaaclab.managers import ActionTermCfg as ActionTerm +from isaaclab.managers import EventTermCfg as EventTerm +from isaaclab.managers import ObservationGroupCfg as ObsGroup +from isaaclab.managers import ObservationTermCfg as ObsTerm +from isaaclab.managers import RewardTermCfg as RewTerm +from isaaclab.managers import SceneEntityCfg +from isaaclab.managers import TerminationTermCfg as DoneTerm +from isaaclab.scene import InteractiveSceneCfg +from isaaclab.sim.simulation_cfg import SimulationCfg +from isaaclab.utils.configclass import configclass +from isaaclab.utils.noise import UniformNoiseCfg + +import isaaclab_tasks.contrib.deploy.mdp as mdp +from isaaclab_tasks.contrib.deploy.mdp.noise_models import ResetSampledConstantNoiseModelCfg + +CABLE_INSERTION_DIR = os.path.dirname(os.path.abspath(__file__)) +DISPLAY_ASSETS_DIR = os.path.join(CABLE_INSERTION_DIR, "display_cable_insertion_assets") + + +def _quat_rotate_vec(q_xyzw, v): + """Apply quaternion rotation to a 3D vector.""" + qx, qy, qz, qw = q_xyzw + vx, vy, vz = v + tx = 2.0 * (qy * vz - qz * vy) + ty = 2.0 * (qz * vx - qx * vz) + tz = 2.0 * (qx * vy - qy * vx) + return ( + vx + qw * tx + qy * tz - qz * ty, + vy + qw * ty + qz * tx - qx * tz, + vz + qw * tz + qx * ty - qy * tx, + ) + + +def _quat_mul(q1_xyzw, q2_xyzw): + """Multiply two quaternions in (x, y, z, w) format.""" + x1, y1, z1, w1 = q1_xyzw + x2, y2, z2, w2 = q2_xyzw + return ( + w1 * x2 + x1 * w2 + y1 * z2 - z1 * y2, + w1 * y2 - x1 * z2 + y1 * w2 + z1 * x2, + w1 * z2 + x1 * y2 - y1 * x2 + z1 * w2, + w1 * w2 - x1 * x2 - y1 * y2 - z1 * z2, + ) + + +# Asset geometry offsets expressed in each body's local frame. +SOCKET_INSERTION_OFFSET = [0.0375, 0.0, 0.0] +PLUG_INSERTION_OFFSET = [0.0, 0.0, 0.0221] +# Plug orientation relative to socket at the mated pose (x, y, z, w). +PLUG_GOAL_ROT = [0.0, -0.70711, 0.0, 0.70711] +PLUG_GOAL_ROT_INV = [0.0, 0.70711, 0.0, 0.70711] + + +def compute_socket_root(geometry_pos, socket_rot): + """Compute socket USD root position from a desired insertion-geometry world position. + + Inverts :data:`SOCKET_INSERTION_OFFSET` (expressed in the socket's local + frame) for a given world-frame socket rotation. + """ + rotated = _quat_rotate_vec(socket_rot, SOCKET_INSERTION_OFFSET) + return ( + geometry_pos[0] - rotated[0], + geometry_pos[1] - rotated[1], + geometry_pos[2] - rotated[2], + ) + + +def compute_plug_pose(geometry_pos, socket_rot, z_clearance=0.0): + """Compute plug USD root position and world-frame rotation. + + Returns ``(plug_root_pos, plug_rot)`` such that the plug insertion point + lands at ``geometry_pos`` (plus optional vertical clearance) with the + correct goal orientation relative to the socket. + """ + plug_rot = _quat_mul(socket_rot, tuple(PLUG_GOAL_ROT)) + plug_offset_world = _quat_rotate_vec(plug_rot, PLUG_INSERTION_OFFSET) + plug_root = ( + geometry_pos[0] - plug_offset_world[0], + geometry_pos[1] - plug_offset_world[1], + geometry_pos[2] - plug_offset_world[2] + z_clearance, + ) + return plug_root, plug_rot + + +_INSERTION_POINT = [0.0, 0.0, 0.1875] +_DEFAULT_SOCKET_ROT = (0.5, 0.5, 0.5, -0.5) # opening faces +Z + +_SOCKET_ROOT_POS = compute_socket_root(_INSERTION_POINT, _DEFAULT_SOCKET_ROT) +_PLUG_ROOT_POS, _DEFAULT_PLUG_ROT = compute_plug_pose( + _INSERTION_POINT, + _DEFAULT_SOCKET_ROT, + z_clearance=0.033, +) + +## +# Asset Configurations +## + + +@configclass +class DisplayPortPlug(RigidObjectCfg): + """DisplayPort right-angle plug (held asset) — dynamic.""" + + prim_path = "{ENV_REGEX_NS}/DisplayPortPlug" + spawn = sim_utils.UsdFileCfg( + usd_path=os.path.join(DISPLAY_ASSETS_DIR, "display_port_plug_fixed_sdf.usd"), + scale=(1.0, 1.0, 1.0), + activate_contact_sensors=True, + rigid_props=sim_utils.RigidBodyPropertiesCfg( + disable_gravity=False, + kinematic_enabled=False, + max_depenetration_velocity=0.5, + linear_damping=0.0, + angular_damping=0.0, + max_linear_velocity=1000.0, + max_angular_velocity=3666.0, + enable_gyroscopic_forces=True, + solver_position_iteration_count=128, + solver_velocity_iteration_count=1, + max_contact_impulse=None, + ), + mass_props=sim_utils.MassPropertiesCfg(mass=0.03), + collision_props=sim_utils.CollisionPropertiesCfg(contact_offset=0.00001, rest_offset=-0.00005), + ) + init_state = RigidObjectCfg.InitialStateCfg(pos=_PLUG_ROOT_POS, rot=_DEFAULT_PLUG_ROT) + + +@configclass +class DisplayPortSocket(RigidObjectCfg): + """DisplayPort socket (fixed asset) — kinematic.""" + + prim_path = "{ENV_REGEX_NS}/DisplayPortSocket" + spawn = sim_utils.UsdFileCfg( + usd_path=os.path.join(DISPLAY_ASSETS_DIR, "display_port_socket_fixed_sdf_noprotrusions.usd"), + scale=(1.0, 1.0, 1.0), + activate_contact_sensors=False, + rigid_props=sim_utils.RigidBodyPropertiesCfg( + disable_gravity=False, + kinematic_enabled=True, + max_depenetration_velocity=5.0, + linear_damping=0.0, + angular_damping=0.0, + max_linear_velocity=1000.0, + max_angular_velocity=3666.0, + enable_gyroscopic_forces=True, + solver_position_iteration_count=128, + solver_velocity_iteration_count=1, + max_contact_impulse=1e32, + ), + mass_props=sim_utils.MassPropertiesCfg(mass=None), + collision_props=sim_utils.CollisionPropertiesCfg(contact_offset=0.0001, rest_offset=-0.0001), + ) + init_state = RigidObjectCfg.InitialStateCfg(pos=_SOCKET_ROOT_POS, rot=_DEFAULT_SOCKET_ROT) + + +## +# Environment configuration +## + + +@configclass +class DisplayportInsertionSceneCfg(InteractiveSceneCfg): + """Configuration for the DisplayPort insertion scene.""" + + replicate_physics = True + + ground = AssetBaseCfg( + prim_path="/World/ground", + spawn=sim_utils.GroundPlaneCfg(), + init_state=AssetBaseCfg.InitialStateCfg(pos=(0.0, 0.0, -1.05)), + ) + + dp_plug = DisplayPortPlug() + dp_socket = DisplayPortSocket() + + robot: ArticulationCfg = MISSING + + light = AssetBaseCfg( + prim_path="/World/light", + spawn=sim_utils.DomeLightCfg(color=(0.75, 0.75, 0.75), intensity=2500.0), + ) + + +@configclass +class ActionsCfg: + """Action specifications for the MDP.""" + + arm_action: ActionTerm = MISSING + gripper_action: ActionTerm | None = None + + +@configclass +class ObservationsCfg: + """Observation specifications for the MDP.""" + + @configclass + class PolicyCfg(ObsGroup): + """Observations for policy group.""" + + joint_pos = ObsTerm(func=mdp.joint_pos, params={"asset_cfg": SceneEntityCfg("robot", joint_names=[".*"])}) + joint_vel = ObsTerm(func=mdp.joint_vel, params={"asset_cfg": SceneEntityCfg("robot", joint_names=[".*"])}) + socket_pos = ObsTerm( + func=mdp.rigid_object_pos_w, + params={"asset_cfg": SceneEntityCfg("dp_socket"), "offset": SOCKET_INSERTION_OFFSET}, + noise=ResetSampledConstantNoiseModelCfg( + noise_cfg=UniformNoiseCfg(n_min=-0.01, n_max=0.01, operation="add") + ), + ) + socket_quat = ObsTerm( + func=mdp.rigid_object_quat_w, + params={"asset_cfg": SceneEntityCfg("dp_socket")}, + ) + + def __post_init__(self): + self.enable_corruption = True + self.concatenate_terms = True + + @configclass + class CriticCfg(ObsGroup): + """Observations for critic group.""" + + joint_pos = ObsTerm(func=mdp.joint_pos, params={"asset_cfg": SceneEntityCfg("robot", joint_names=[".*"])}) + joint_vel = ObsTerm(func=mdp.joint_vel, params={"asset_cfg": SceneEntityCfg("robot", joint_names=[".*"])}) + socket_pos = ObsTerm( + func=mdp.rigid_object_pos_w, + params={"asset_cfg": SceneEntityCfg("dp_socket"), "offset": SOCKET_INSERTION_OFFSET}, + ) + socket_quat = ObsTerm( + func=mdp.rigid_object_quat_w, + params={"asset_cfg": SceneEntityCfg("dp_socket")}, + ) + plug_pos = ObsTerm( + func=mdp.rigid_object_pos_w, + params={"asset_cfg": SceneEntityCfg("dp_plug"), "offset": PLUG_INSERTION_OFFSET}, + ) + plug_quat = ObsTerm( + func=mdp.rigid_object_quat_w, + params={"asset_cfg": SceneEntityCfg("dp_plug")}, + ) + + policy: PolicyCfg = PolicyCfg() + critic: CriticCfg = CriticCfg() + + +@configclass +class EventCfg: + """Configuration for events.""" + + reset_all = EventTerm(func=mdp.reset_scene_to_default, mode="reset") + + +@configclass +class RewardsCfg: + """Reward terms for the MDP.""" + + plug_socket_keypoint_tracking = RewTerm( + func=mdp.keypoint_two_body_error, + weight=-1.5, + params={ + "asset_cfg_1": SceneEntityCfg("dp_socket"), + "asset_cfg_2": SceneEntityCfg("dp_plug"), + "offset_1": SOCKET_INSERTION_OFFSET, + "offset_2": PLUG_INSERTION_OFFSET, + "rot_offset_2": PLUG_GOAL_ROT_INV, + "keypoint_scale": 0.15, + }, + ) + + plug_socket_keypoint_tracking_exp = RewTerm( + func=mdp.keypoint_two_body_error_exp, + weight=3.0, + params={ + "asset_cfg_1": SceneEntityCfg("dp_socket"), + "asset_cfg_2": SceneEntityCfg("dp_plug"), + "offset_1": SOCKET_INSERTION_OFFSET, + "offset_2": PLUG_INSERTION_OFFSET, + "rot_offset_2": PLUG_GOAL_ROT_INV, + "kp_exp_coeffs": [(50, 0.0001), (300, 0.0001), (600, 0.0001), (2000, 0.0001)], + "kp_use_sum_of_exps": False, + "keypoint_scale": 0.15, + }, + ) + + action_rate = RewTerm(func=mdp.action_rate_l2, weight=-5.0e-06) + + +@configclass +class TerminationsCfg: + """Termination terms for the MDP.""" + + time_out = DoneTerm(func=mdp.time_out, time_out=True) + + +@configclass +class DisplayportInsertionEnvCfg(ManagerBasedRLEnvCfg): + """Base configuration for DisplayPort plug/socket insertion.""" + + scene: DisplayportInsertionSceneCfg = DisplayportInsertionSceneCfg(num_envs=4096, env_spacing=2.5) + observations: ObservationsCfg = ObservationsCfg() + actions: ActionsCfg = ActionsCfg() + rewards: RewardsCfg = RewardsCfg() + terminations: TerminationsCfg = TerminationsCfg() + events: EventCfg = EventCfg() + + # Task-success logging (consumed by DisplayportInsertionEnv) + log_success_metrics: bool = True + success_socket_asset: str = "dp_socket" + success_plug_asset: str = "dp_plug" + success_pos_threshold: float = 0.003 + success_keypoint_scale: float = 0.15 + success_socket_offset: list = MISSING + success_plug_offset: list = MISSING + success_plug_goal_rot_inv: list = MISSING + sim: SimulationCfg = SimulationCfg( + physics_material=sim_utils.RigidBodyMaterialCfg( + friction_combine_mode="multiply", + restitution_combine_mode="multiply", + static_friction=1.0, + dynamic_friction=1.0, + restitution=0.0, + ), + physics=PhysxCfg( + bounce_threshold_velocity=0.2, + friction_offset_threshold=0.01, + friction_correlation_distance=0.00625, + gpu_collision_stack_size=2**30, + gpu_max_rigid_contact_count=2**23, + gpu_max_rigid_patch_count=2**23, + ), + ) + + def __post_init__(self): + """Post initialization.""" + self.episode_length_s = 6.66 + self.viewer.eye = (0.5, -1.8, 1.2) + self.viewer.lookat = (0.5, 0.0, 0.5) + self.decimation = 8 + self.sim.render_interval = self.decimation + self.sim.dt = 1.0 / 240.0 + + # Success mate-point geometry mirrors the keypoint-tracking reward + self.success_socket_offset = list(SOCKET_INSERTION_OFFSET) + self.success_plug_offset = list(PLUG_INSERTION_OFFSET) + self.success_plug_goal_rot_inv = list(PLUG_GOAL_ROT_INV) diff --git a/source/isaaclab_tasks/isaaclab_tasks/contrib/deploy/cable_insertion/insertion_env.py b/source/isaaclab_tasks/isaaclab_tasks/contrib/deploy/cable_insertion/insertion_env.py new file mode 100644 index 00000000000..b1fc55cc5c6 --- /dev/null +++ b/source/isaaclab_tasks/isaaclab_tasks/contrib/deploy/cable_insertion/insertion_env.py @@ -0,0 +1,124 @@ +# 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 + +"""Manager-based insertion env with task-success detection and logging. + +Adds success metrics to ``extras["log"]`` for RSL-RL without changing the MDP +observation, action, reward, or termination logic. +""" + +from __future__ import annotations + +import torch +import warp as wp + +from isaaclab.envs import ManagerBasedRLEnv +from isaaclab.utils.math import combine_frame_transforms + + +def _keypoint_offsets_6d(device: torch.device) -> torch.Tensor: + """Return the 7 unit keypoint offsets used by the keypoint reward.""" + corners = torch.tensor([[0, 0, 0], [1, 0, 0], [0, 1, 0], [0, 0, 1]], device=device, dtype=torch.float32) + return torch.cat((corners, -corners[-3:]), dim=0) + + +class DisplayportInsertionEnv(ManagerBasedRLEnv): + """Manager-based RL env that logs insertion success metrics during training. + + The following scalars are added to ``extras["log"]``: + + - ``Metrics/success_rate``: fraction of envs within the success threshold + - ``Metrics/plug_socket_pos_error_m``: mean mate-point distance (m) + - ``Metrics/plug_socket_keypoint_dist_m``: mean keypoint distance (m) + - ``Metrics/terminal_success_rate``: success fraction at episode reset + """ + + def __init__(self, cfg, render_mode: str | None = None, **kwargs): + super().__init__(cfg, render_mode=render_mode, **kwargs) + + self._log_success_metrics: bool = bool(getattr(cfg, "log_success_metrics", True)) + self._success_socket_asset: str = getattr(cfg, "success_socket_asset", "dp_socket") + self._success_plug_asset: str = getattr(cfg, "success_plug_asset", "dp_plug") + self._success_pos_threshold: float = float(getattr(cfg, "success_pos_threshold", 0.003)) + self._success_keypoint_scale: float = float(getattr(cfg, "success_keypoint_scale", 0.15)) + + device = self.device + self._success_socket_offset = torch.tensor( + getattr(cfg, "success_socket_offset", [0.0, 0.0, 0.0]), device=device, dtype=torch.float32 + ) + self._success_plug_offset = torch.tensor( + getattr(cfg, "success_plug_offset", [0.0, 0.0, 0.0]), device=device, dtype=torch.float32 + ) + self._success_plug_goal_rot_inv = torch.tensor( + getattr(cfg, "success_plug_goal_rot_inv", [0.0, 0.0, 0.0, 1.0]), device=device, dtype=torch.float32 + ) + + self._success_identity_quat = torch.tensor([[0.0, 0.0, 0.0, 1.0]], device=device, dtype=torch.float32).repeat( + self.num_envs, 1 + ) + self._success_kp_offsets = _keypoint_offsets_6d(device) * self._success_keypoint_scale + + def _compute_success(self) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Compute per-env success mask, mate-point distance, and keypoint distance.""" + socket = self.scene[self._success_socket_asset] + plug = self.scene[self._success_plug_asset] + + socket_pos = wp.to_torch(socket.data.root_pos_w) + socket_quat = wp.to_torch(socket.data.root_quat_w) + plug_pos = wp.to_torch(plug.data.root_pos_w) + plug_quat = wp.to_torch(plug.data.root_quat_w) + + n = self.num_envs + socket_off = self._success_socket_offset.unsqueeze(0).expand(n, -1) + plug_off = self._success_plug_offset.unsqueeze(0).expand(n, -1) + plug_goal_rot_inv = self._success_plug_goal_rot_inv.unsqueeze(0).expand(n, -1) + + # Mate reference frames (same construction as keypoint_two_body_error) + kp_pos_s, kp_quat_s = combine_frame_transforms(socket_pos, socket_quat, socket_off, self._success_identity_quat) + kp_pos_p, kp_quat_p = combine_frame_transforms(plug_pos, plug_quat, plug_off, plug_goal_rot_inv) + + pos_error = torch.linalg.norm(kp_pos_p - kp_pos_s, dim=-1) + + k = self._success_kp_offsets.shape[0] + offs_flat = self._success_kp_offsets.unsqueeze(0).expand(n, -1, -1).reshape(-1, 3) + ident_flat = self._success_identity_quat.unsqueeze(1).expand(-1, k, -1).reshape(-1, 4) + + kp_s = combine_frame_transforms( + kp_pos_s.unsqueeze(1).expand(-1, k, -1).reshape(-1, 3), + kp_quat_s.unsqueeze(1).expand(-1, k, -1).reshape(-1, 4), + offs_flat, + ident_flat, + )[0].reshape(n, k, 3) + kp_p = combine_frame_transforms( + kp_pos_p.unsqueeze(1).expand(-1, k, -1).reshape(-1, 3), + kp_quat_p.unsqueeze(1).expand(-1, k, -1).reshape(-1, 4), + offs_flat, + ident_flat, + )[0].reshape(n, k, 3) + keypoint_dist = torch.linalg.norm(kp_p - kp_s, dim=-1).mean(dim=-1) + + is_success = pos_error < self._success_pos_threshold + return is_success, pos_error, keypoint_dist + + def step(self, action: torch.Tensor): + obs_buf, reward_buf, terminated, time_outs, extras = super().step(action) + if getattr(self, "_log_success_metrics", False): + is_success, pos_error, keypoint_dist = self._compute_success() + log = self.extras.setdefault("log", {}) + log["Metrics/success_rate"] = is_success.float().mean() + log["Metrics/plug_socket_pos_error_m"] = pos_error.mean() + log["Metrics/plug_socket_keypoint_dist_m"] = keypoint_dist.mean() + return obs_buf, reward_buf, terminated, time_outs, self.extras + + def _reset_idx(self, env_ids): + terminal_success = None + if getattr(self, "_log_success_metrics", False): + is_success, _, _ = self._compute_success() + terminal_success = is_success[env_ids].float().mean() + + super()._reset_idx(env_ids) + + if terminal_success is not None: + self.extras["log"]["Metrics/terminal_success_rate"] = terminal_success 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 2a200c888bc..f6dd621ad25 100644 --- a/source/isaaclab_tasks/isaaclab_tasks/contrib/deploy/mdp/__init__.pyi +++ b/source/isaaclab_tasks/isaaclab_tasks/contrib/deploy/mdp/__init__.pyi @@ -7,25 +7,60 @@ __all__ = [ "randomize_gear_type", "randomize_gears_and_base_pose", "set_robot_to_grasp_pose", + "set_robot_to_object_grasp_pose", + "reset_plug_at_goal_curriculum", + "DeployRelativeJointPositionAction", + "DeployRelativeJointPositionActionCfg", "ResetSampledConstantNoiseModel", "ResetSampledConstantNoiseModelCfg", + "joint_pos", + "joint_vel", "gear_pos_w", "gear_quat_w", "gear_shaft_pos_w", "gear_shaft_quat_w", + "rigid_object_pos_w", + "rigid_object_quat_w", + "rigid_object_rot_6d_w", + "eef_pos_w", + "eef_rot_6d_w", "keypoint_command_error", "keypoint_command_error_exp", "keypoint_entity_error", "keypoint_entity_error_exp", "keypoint_ee_grasp_error", "keypoint_ee_grasp_error_exp", + "keypoint_two_body_error", + "keypoint_two_body_error_exp", "reset_when_gear_dropped", "reset_when_gear_orientation_exceeds_threshold", + "reset_when_plug_dropped", + "reset_when_plug_orientation_exceeded", ] -from .events import randomize_gear_type, randomize_gears_and_base_pose, set_robot_to_grasp_pose +from .events import ( + randomize_gear_type, + randomize_gears_and_base_pose, + set_robot_to_grasp_pose, + set_robot_to_object_grasp_pose, + reset_plug_at_goal_curriculum, +) +from .actions import DeployRelativeJointPositionAction +from .actions_cfg import DeployRelativeJointPositionActionCfg from .noise_models import ResetSampledConstantNoiseModel, ResetSampledConstantNoiseModelCfg -from .observations import gear_pos_w, gear_quat_w, gear_shaft_pos_w, gear_shaft_quat_w +from .observations import ( + joint_pos, + joint_vel, + gear_pos_w, + gear_quat_w, + gear_shaft_pos_w, + gear_shaft_quat_w, + rigid_object_pos_w, + rigid_object_quat_w, + rigid_object_rot_6d_w, + eef_pos_w, + eef_rot_6d_w, +) from .rewards import ( keypoint_command_error, keypoint_command_error_exp, @@ -33,6 +68,13 @@ from .rewards import ( keypoint_entity_error_exp, keypoint_ee_grasp_error, keypoint_ee_grasp_error_exp, + keypoint_two_body_error, + keypoint_two_body_error_exp, +) +from .terminations import ( + reset_when_gear_dropped, + reset_when_gear_orientation_exceeds_threshold, + reset_when_plug_dropped, + reset_when_plug_orientation_exceeded, ) -from .terminations import reset_when_gear_dropped, reset_when_gear_orientation_exceeds_threshold from isaaclab.envs.mdp import * diff --git a/source/isaaclab_tasks/isaaclab_tasks/contrib/deploy/mdp/actions.py b/source/isaaclab_tasks/isaaclab_tasks/contrib/deploy/mdp/actions.py new file mode 100644 index 00000000000..87c5ec12900 --- /dev/null +++ b/source/isaaclab_tasks/isaaclab_tasks/contrib/deploy/mdp/actions.py @@ -0,0 +1,120 @@ +# 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 + +"""Deploy-specific action terms for LEAPP export workflows.""" + +from __future__ import annotations + +from typing import TYPE_CHECKING + +import warp as wp + +from isaaclab.envs.mdp.actions.joint_actions import RelativeJointPositionAction + +if TYPE_CHECKING: + from isaaclab.envs import ManagerBasedEnv + + from .actions_cfg import DeployRelativeJointPositionActionCfg + +_LEAPP_TRACED_OBSERVATION_INPUTS = "_leapp_traced_observation_inputs" +_LEAPP_CONSUMED_OBSERVATION_INPUTS = "_leapp_consumed_observation_inputs" + + +def _leapp_real_env(env): + real_env = object.__getattribute__(env, "_real_env") if type(env).__name__ == "_EnvProxy" else env + return real_env + + +def _tensor_data_to_torch(data): + """Return a torch tensor view for Isaac Lab data stored as torch or Warp-backed data.""" + return data.torch if hasattr(data, "torch") else wp.to_torch(data) + + +def _get_observation_term_from_buffer(env, group_name: str, term_name: str): + """Return a term slice from the cached observation buffer.""" + obs_buffer = getattr(env, "obs_buf", None) + if obs_buffer is None: + obs_buffer = getattr(getattr(env, "observation_manager", None), "_obs_buffer", None) + if not obs_buffer or group_name not in obs_buffer: + return None + + group_obs = obs_buffer[group_name] + if isinstance(group_obs, dict): + return group_obs.get(term_name) + + obs_manager = getattr(env, "observation_manager", None) + if obs_manager is None: + return None + + term_names = obs_manager.active_terms.get(group_name, []) + if term_name not in term_names: + return None + + term_index = term_names.index(term_name) + term_dims = obs_manager.group_obs_term_dim[group_name] + concat_dim = obs_manager._group_obs_concatenate_dim[group_name] + if concat_dim > 0: + concat_dim -= 1 + + start = sum(dim[concat_dim] for dim in term_dims[:term_index]) + length = term_dims[term_index][concat_dim] + return group_obs.narrow(dim=concat_dim, start=start, length=length) + + +def _pop_leapp_traced_observation_input(env, name: str, *, group_name: str, term_name: str): + """Consume one traced observation tensor for the current LEAPP action trace.""" + real_env = _leapp_real_env(env) + consumed_inputs = getattr(real_env, _LEAPP_CONSUMED_OBSERVATION_INPUTS, None) + if consumed_inputs is None: + consumed_inputs = set() + setattr(real_env, _LEAPP_CONSUMED_OBSERVATION_INPUTS, consumed_inputs) + + if name in consumed_inputs: + return None + + traced_inputs = getattr(real_env, _LEAPP_TRACED_OBSERVATION_INPUTS, {}) + traced_tensor = traced_inputs.pop(name, None) + if traced_tensor is None: + traced_tensor = _get_observation_term_from_buffer(real_env, group_name, term_name) + + if traced_tensor is not None: + consumed_inputs.add(name) + return traced_tensor + + +def _is_leapp_observation_input_consumed(env, name: str) -> bool: + real_env = _leapp_real_env(env) + return name in getattr(real_env, _LEAPP_CONSUMED_OBSERVATION_INPUTS, set()) + + +class DeployRelativeJointPositionAction(RelativeJointPositionAction): + """Relative joint action that reuses traced current joint observations during LEAPP export.""" + + def __init__(self, cfg: DeployRelativeJointPositionActionCfg, env: ManagerBasedEnv): + super().__init__(cfg, env) + + def apply_actions(self): + asset = self._asset + if type(asset).__name__ == "_ArticulationWriteProxy": + observation_input_name = f"{self.cfg.asset_name}_joint_pos" + current_joint_pos = _pop_leapp_traced_observation_input( + self._env, + observation_input_name, + group_name="policy", + term_name="joint_pos", + ) + if current_joint_pos is None: + if not _is_leapp_observation_input_consumed(self._env, observation_input_name): + raise RuntimeError( + "DeployRelativeJointPositionAction requires the traced " + f"'{self.cfg.asset_name}_joint_pos' observation during LEAPP export." + ) + real_asset = object.__getattribute__(asset, "_real_asset") + current_joint_pos = _tensor_data_to_torch(real_asset.data.joint_pos)[:, self._joint_ids] + else: + current_joint_pos = _tensor_data_to_torch(asset.data.joint_pos)[:, self._joint_ids] + + current_actions = self.processed_actions + current_joint_pos + self._asset.set_joint_position_target_index(target=current_actions, joint_ids=self._joint_ids) diff --git a/source/isaaclab_tasks/isaaclab_tasks/contrib/deploy/mdp/actions_cfg.py b/source/isaaclab_tasks/isaaclab_tasks/contrib/deploy/mdp/actions_cfg.py new file mode 100644 index 00000000000..faf8088b883 --- /dev/null +++ b/source/isaaclab_tasks/isaaclab_tasks/contrib/deploy/mdp/actions_cfg.py @@ -0,0 +1,18 @@ +# 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 + +"""Deploy-specific action configuration classes.""" + +from __future__ import annotations + +from isaaclab.envs.mdp.actions.actions_cfg import RelativeJointPositionActionCfg +from isaaclab.utils.configclass import configclass + + +@configclass +class DeployRelativeJointPositionActionCfg(RelativeJointPositionActionCfg): + """Configuration for deploy relative joint actions with explicit LEAPP current-joint input.""" + + class_type: type | str = "isaaclab_tasks.contrib.deploy.mdp.actions:DeployRelativeJointPositionAction" 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 8dde6b5c510..6e6f72618b4 100644 --- a/source/isaaclab_tasks/isaaclab_tasks/contrib/deploy/mdp/events.py +++ b/source/isaaclab_tasks/isaaclab_tasks/contrib/deploy/mdp/events.py @@ -3,7 +3,7 @@ # # SPDX-License-Identifier: BSD-3-Clause -"""Class-based event terms specific to the gear assembly manipulation environments.""" +"""Class-based event terms for manipulation deployment environments.""" from __future__ import annotations @@ -11,6 +11,7 @@ from typing import TYPE_CHECKING import torch +import warp as wp import isaaclab.utils.math as math_utils from isaaclab.managers import EventTermCfg, ManagerTermBase, SceneEntityCfg @@ -503,3 +504,450 @@ def __call__( 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) + + +class set_robot_to_object_grasp_pose(ManagerTermBase): + """Set robot to a grasp pose over a single named target object using IK. + + Generic single-object counterpart of :class:`set_robot_to_grasp_pose` (which + is keyed on the gear-type manager). This term targets a single named + :class:`~isaaclab.assets.RigidObject` with a fixed grasp offset, suitable + for cable insertion and other single-object manipulation tasks. + + Args: + target_object_name: Name of the rigid object in the scene to grasp. + end_effector_body_name: Name of the end-effector body on the robot. + num_arm_joints: Number of arm joints (the remaining joints are + assumed to be gripper/finger joints). + grasp_offset: Position offset ``[x, y, z]`` [m] applied in the + (rotated) object frame to define the IK target. Defaults to zero. + grasp_rot_offset: Quaternion offset ``(x, y, z, w)`` applied to the + object orientation to define the IK target. + gripper_joint_setter_func: Callable used to set finger joint positions + for the configured grasp/close widths. + robot_asset_cfg: Robot asset configuration. Defaults to + ``SceneEntityCfg("robot")``. + pos_threshold: IK position-error tolerance [m]. + rot_threshold: IK rotation-error tolerance [rad]. + max_iterations: Maximum IK iterations per env reset. + pos_randomization_range: Optional dict with keys ``"x"``, ``"y"``, + ``"z"`` mapping to ``(low, high)`` tuples [m] for per-reset + randomization of the grasp offset. + """ + + def __init__(self, cfg: EventTermCfg, env: ManagerBasedEnv): + super().__init__(cfg, env) + + self.robot_asset_cfg: SceneEntityCfg = cfg.params.get("robot_asset_cfg", SceneEntityCfg("robot")) + self.robot_asset: Articulation = env.scene[self.robot_asset_cfg.name] + + for required in ( + "end_effector_body_name", + "num_arm_joints", + "grasp_rot_offset", + "gripper_joint_setter_func", + "target_object_name", + ): + if required not in cfg.params: + raise ValueError(f"'{required}' is required in set_robot_to_object_grasp_pose configuration.") + + self.end_effector_body_name: str = cfg.params["end_effector_body_name"] + self.num_arm_joints: int = cfg.params["num_arm_joints"] + self.gripper_joint_setter_func = cfg.params["gripper_joint_setter_func"] + self.target_object_name: str = cfg.params["target_object_name"] + + grasp_offset = cfg.params.get("grasp_offset", [0.0, 0.0, 0.0]) + self.grasp_offset_tensor = torch.tensor(grasp_offset, device=env.device, dtype=torch.float32) + + grasp_rot_offset = cfg.params["grasp_rot_offset"] + self.grasp_rot_offset_tensor = ( + torch.tensor(grasp_rot_offset, device=env.device, dtype=torch.float32).unsqueeze(0).repeat(env.num_envs, 1) + ) + + self.grasp_offsets_buffer = torch.zeros(env.num_envs, 3, device=env.device, dtype=torch.float32) + + self.hand_grasp_width = env.cfg.hand_grasp_width + self.hand_close_width = env.cfg.hand_close_width + # hand_hold_width: joint angle where fingers just touch the held object + # surface. Written as the physical STATE so there is no mesh overlap. + # Falls back to hand_close_width when not set (original behaviour). + self.hand_hold_width = getattr(env.cfg, "hand_hold_width", self.hand_close_width) + + eef_indices, _ = self.robot_asset.find_bodies([self.end_effector_body_name]) + if len(eef_indices) == 0: + raise ValueError(f"End effector body '{self.end_effector_body_name}' not found in robot") + self.eef_idx = eef_indices[0] + self.jacobi_body_idx = self.eef_idx - 1 + + all_joints, _ = self.robot_asset.find_joints([".*"]) + self.all_joints = all_joints + self.finger_joints = all_joints[self.num_arm_joints :] + + def __call__( + self, + env: ManagerBasedEnv, + env_ids: torch.Tensor, + robot_asset_cfg: SceneEntityCfg = SceneEntityCfg("robot"), + pos_threshold: float = 1e-6, + rot_threshold: float = 1e-6, + max_iterations: int = 50, + pos_randomization_range: dict | None = None, + target_object_name: str | None = None, + grasp_offset: list | None = None, + 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, + ): + num_reset_envs = len(env_ids) + grasp_offsets = self.grasp_offsets_buffer[:num_reset_envs] + grasp_rot_offset_tensor = self.grasp_rot_offset_tensor[env_ids] + + # One-shot debug log to confirm the event fires and report IK convergence. + # Remove or guard once the grasp wiring is verified. + debug_first_call = not getattr(self, "_debug_printed", False) + if debug_first_call: + self._debug_printed = True + target_object_dbg: RigidObject = env.scene[self.target_object_name] + init_obj_pos = wp.to_torch(target_object_dbg.data.root_link_pos_w)[env_ids][0].tolist() + init_obj_quat = wp.to_torch(target_object_dbg.data.root_link_quat_w)[env_ids][0].tolist() + init_eef_pos = wp.to_torch(self.robot_asset.data.body_pos_w)[env_ids, self.eef_idx][0].tolist() + init_eef_quat = wp.to_torch(self.robot_asset.data.body_quat_w)[env_ids, self.eef_idx][0].tolist() + print( + f"[GRASP-DBG] set_robot_to_object_grasp_pose fired:" + f" target={self.target_object_name!r} num_reset_envs={num_reset_envs}" + f" eef_idx={self.eef_idx} num_arm_joints={self.num_arm_joints}\n" + f" grasp_offset={self.grasp_offset_tensor.tolist()}" + f" grasp_rot_offset(xyzw)={self.grasp_rot_offset_tensor[0].tolist()}" + f" hand_close_width={self.hand_close_width}\n" + f" init_obj_pos_w={init_obj_pos}" + f" init_obj_quat(xyzw)={init_obj_quat}\n" + f" init_eef_pos_w={init_eef_pos}" + f" init_eef_quat(xyzw)={init_eef_quat}" + ) + + last_pos_err = None + last_rot_err = None + converged_at = -1 + last_target_pos = None + last_target_quat = None + + for _iter in range(max_iterations): + joint_pos = wp.to_torch(self.robot_asset.data.joint_pos)[env_ids].clone() + joint_vel = wp.to_torch(self.robot_asset.data.joint_vel)[env_ids].clone() + + target_object: RigidObject = env.scene[self.target_object_name] + grasp_object_pos_world = wp.to_torch(target_object.data.root_link_pos_w)[env_ids] + grasp_object_quat = wp.to_torch(target_object.data.root_link_quat_w)[env_ids] + + grasp_object_quat = math_utils.quat_mul(grasp_object_quat, grasp_rot_offset_tensor) + + grasp_offsets[:] = self.grasp_offset_tensor + + if pos_randomization_range is not None: + pos_keys = ["x", "y", "z"] + range_list_pos = [pos_randomization_range.get(key, (0.0, 0.0)) for key in pos_keys] + ranges_pos = torch.tensor(range_list_pos, device=env.device) + rand_pos_offsets = math_utils.sample_uniform( + ranges_pos[:, 0], ranges_pos[:, 1], (len(env_ids), 3), device=env.device + ) + grasp_offsets = grasp_offsets + rand_pos_offsets + + grasp_object_pos_world = grasp_object_pos_world + math_utils.quat_apply(grasp_object_quat, grasp_offsets) + + eef_pos = wp.to_torch(self.robot_asset.data.body_pos_w)[env_ids, self.eef_idx] + eef_quat = wp.to_torch(self.robot_asset.data.body_quat_w)[env_ids, self.eef_idx] + + last_target_pos = grasp_object_pos_world.clone() + last_target_quat = grasp_object_quat.clone() + + pos_error, axis_angle_error = fc.get_pose_error( + fingertip_midpoint_pos=eef_pos, + fingertip_midpoint_quat=eef_quat, + ctrl_target_fingertip_midpoint_pos=grasp_object_pos_world, + ctrl_target_fingertip_midpoint_quat=grasp_object_quat, + jacobian_type="geometric", + rot_error_type="axis_angle", + ) + delta_hand_pose = torch.cat((pos_error, axis_angle_error), dim=-1) + + pos_error_norm = torch.linalg.norm(pos_error, dim=-1) + rot_error_norm = torch.linalg.norm(axis_angle_error, dim=-1) + last_pos_err = pos_error_norm + last_rot_err = rot_error_norm + + if torch.all(pos_error_norm < pos_threshold) and torch.all(rot_error_norm < rot_threshold): + converged_at = _iter + break + + jacobians = wp.to_torch(self.robot_asset.root_view.get_jacobians()).clone() + jacobian = jacobians[env_ids, self.jacobi_body_idx, :, :] + + delta_dof_pos = fc._get_delta_dof_pos( + delta_pose=delta_hand_pose, + ik_method="dls", + jacobian=jacobian, + device=env.device, + ) + + joint_pos = joint_pos + delta_dof_pos + + joint_pos_limits = wp.to_torch(self.robot_asset.data.joint_pos_limits)[env_ids, : self.num_arm_joints, :] + joint_min = joint_pos_limits[:, :, 0] + joint_max = joint_pos_limits[:, :, 1] + joint_range = joint_max - joint_min + + arm_joint_pos = joint_pos[:, : self.num_arm_joints] + arm_joint_pos = torch.where( + joint_range > 0, + joint_min + torch.remainder(arm_joint_pos - joint_min, joint_range), + arm_joint_pos, + ) + joint_pos[:, : self.num_arm_joints] = arm_joint_pos + + joint_vel = torch.zeros_like(joint_pos) + + self.robot_asset.set_joint_position_target_index(target=joint_pos, env_ids=env_ids) + self.robot_asset.set_joint_velocity_target_index(target=joint_vel, env_ids=env_ids) + self.robot_asset.write_joint_position_to_sim_index(position=joint_pos, env_ids=env_ids) + self.robot_asset.write_joint_velocity_to_sim_index(velocity=joint_vel, env_ids=env_ids) + + # Snap the held object to the achieved gripper pose so the gripper actually + # holds it after closing. Without this, any IK residual error or USD + # geometry offset leaves the object outside the finger gap and gravity drops + # Snap the held object to the achieved gripper pose so the gripper + # actually holds it after closing. + held_object = env.scene[self.target_object_name] + achieved_hand_pos = wp.to_torch(self.robot_asset.data.body_pos_w)[env_ids, self.eef_idx].clone() + achieved_hand_quat = wp.to_torch(self.robot_asset.data.body_quat_w)[env_ids, self.eef_idx].clone() + + # Object orientation: inverse of grasp_rot_offset applied to the achieved hand quat, + # because IK target was ``hand = obj * grasp_rot_offset`` => ``obj = hand * grasp_rot_offset^{-1}``. + inv_grasp_rot_offset = math_utils.quat_conjugate(grasp_rot_offset_tensor) + target_obj_quat = math_utils.quat_mul(achieved_hand_quat, inv_grasp_rot_offset) + + # Object position: IK target was + # ``hand_pos = obj_pos + R(obj * grasp_rot_offset) * grasp_offset`` + # = ``obj_pos + R(hand_quat) * grasp_offset`` + # so ``obj_pos = hand_pos - R(hand_quat) * grasp_offset``. + grasp_offset_in_world = math_utils.quat_apply(achieved_hand_quat, grasp_offsets) + target_obj_pos = achieved_hand_pos - grasp_offset_in_world + + new_root_pose = torch.cat([target_obj_pos, target_obj_quat], dim=-1) + zero_velocity = torch.zeros((len(env_ids), 6), device=env.device, dtype=torch.float32) + held_object.write_root_pose_to_sim(new_root_pose, env_ids=env_ids) + held_object.write_root_velocity_to_sim(zero_velocity, env_ids=env_ids) + + if debug_first_call: + pos_err_max = float(last_pos_err.max().item()) if last_pos_err is not None else float("nan") + rot_err_max = float(last_rot_err.max().item()) if last_rot_err is not None else float("nan") + tgt_pos0 = target_obj_pos[0].tolist() + tgt_quat0 = target_obj_quat[0].tolist() + eef_pos0 = achieved_hand_pos[0].tolist() + eef_quat0 = achieved_hand_quat[0].tolist() + ik_target_pos0 = last_target_pos[0].tolist() if last_target_pos is not None else None + ik_target_quat0 = last_target_quat[0].tolist() if last_target_quat is not None else None + print( + f"[GRASP-DBG] IK done: converged_at_iter={converged_at}\n" + f" max_pos_err={pos_err_max:.6f} m, max_rot_err={rot_err_max:.6f} rad\n" + f" IK_target_pos_w[0]={ik_target_pos0}\n" + f" IK_target_quat(xyzw)[0]={ik_target_quat0}\n" + f" achieved_eef_pos_w[0]={eef_pos0}\n" + f" achieved_eef_quat(xyzw)[0]={eef_quat0}\n" + f" snapped_obj_pos_w[0]={tgt_pos0}\n" + f" snapped_obj_quat(xyzw)[0]={tgt_quat0}" + ) + + joint_vel = torch.zeros_like(wp.to_torch(self.robot_asset.data.joint_vel)[env_ids]) + joint_pos = wp.to_torch(self.robot_asset.data.joint_pos)[env_ids].clone() + + # Write gripper STATE at ``hand_hold_width`` (fingers just touching the + # plug, no mesh overlap) and set the TARGET to ``hand_close_width`` + # (fully closed) so the actuator drive squeezes around the plug. + self.gripper_joint_setter_func(joint_pos, list(range(num_reset_envs)), self.finger_joints, self.hand_hold_width) + self.robot_asset.write_joint_position_to_sim_index(position=joint_pos, env_ids=env_ids) + self.robot_asset.write_joint_velocity_to_sim_index(velocity=joint_vel, env_ids=env_ids) + + self.gripper_joint_setter_func( + joint_pos, list(range(num_reset_envs)), self.finger_joints, self.hand_close_width + ) + self.robot_asset.set_joint_position_target_index(target=joint_pos, joint_ids=self.all_joints, env_ids=env_ids) + + +class reset_plug_at_goal_curriculum(ManagerTermBase): + """Reset a fraction of plugs at the goal position (at-goal curriculum). + + For each reset batch, a fraction ``at_goal_prob`` of environments have the + plug placed along the insertion axis at a random depth (from socket opening + to full insertion) with goal orientation. The remaining environments get + normal pose randomization. + + This replaces the simple ``reset_root_state_uniform`` for the plug when + curriculum-based training is desired. + """ + + def __init__(self, cfg: EventTermCfg, env: ManagerBasedEnv): + super().__init__(cfg, env) + + self.plug: RigidObject = env.scene[cfg.params["plug_cfg"].name] + self.socket: RigidObject = env.scene[cfg.params["socket_cfg"].name] + + self.at_goal_prob: float = cfg.params.get("at_goal_prob", 0.8) + + # Optional linear annealing of at_goal_prob over training iterations. + # `at_goal_prob` is the starting value; it decays linearly to + # `at_goal_prob_final` between `anneal_start_iter` and `anneal_end_iter`. + # Annealing is active only when both `at_goal_prob_final` and + # `anneal_end_iter` are provided (otherwise the probability is constant). + # Iterations are derived from the env step counter via `num_steps_per_env` + # (one RL iteration == `num_steps_per_env` env steps). + self.at_goal_prob_final = cfg.params.get("at_goal_prob_final", None) + self.anneal_start_iter: float = cfg.params.get("anneal_start_iter", 0.0) + self.anneal_end_iter = cfg.params.get("anneal_end_iter", None) + self.num_steps_per_env = cfg.params.get("num_steps_per_env", None) + + insertion_axis = cfg.params.get("insertion_axis", [0.0, 0.0, 1.0]) + self.insertion_axis = torch.tensor(insertion_axis, device=env.device, dtype=torch.float32) + self.insertion_axis = self.insertion_axis / self.insertion_axis.norm() + + self.insertion_length: float = cfg.params.get("insertion_length", 0.02) + + # Optional depth ranges along the insertion axis, measured from the socket keypoint origin. + self.at_goal_depth_range = cfg.params.get("at_goal_depth_range", None) + self.approach_depth_range = cfg.params.get("approach_depth_range", None) + + socket_offset = cfg.params.get("socket_insertion_offset", [0.0, 0.0, 0.0]) + self.socket_insertion_offset = torch.tensor(socket_offset, device=env.device, dtype=torch.float32) + + plug_offset = cfg.params.get("plug_insertion_offset", [0.0, 0.0, 0.0]) + self.plug_insertion_offset = torch.tensor(plug_offset, device=env.device, dtype=torch.float32) + + goal_rot = cfg.params.get("goal_rot", [0.0, 0.0, 0.0, 1.0]) + self.goal_rot = torch.tensor(goal_rot, device=env.device, dtype=torch.float32) + + self.normal_pose_range: dict = cfg.params.get("normal_pose_range", {}) + + self.identity_quat = torch.tensor([0.0, 0.0, 0.0, 1.0], device=env.device, dtype=torch.float32) + + def _current_at_goal_prob(self, env: ManagerBasedEnv) -> float: + """Return the at-goal probability for the current training progress. + + Linearly interpolates from ``at_goal_prob`` to ``at_goal_prob_final`` + between ``anneal_start_iter`` and ``anneal_end_iter``. Returns the + constant ``at_goal_prob`` when annealing is not fully configured. + """ + if self.at_goal_prob_final is None or self.anneal_end_iter is None or not self.num_steps_per_env: + return self.at_goal_prob + + current_iter = env.common_step_counter / float(self.num_steps_per_env) + span = max(float(self.anneal_end_iter) - float(self.anneal_start_iter), 1e-9) + frac = (current_iter - float(self.anneal_start_iter)) / span + frac = min(max(frac, 0.0), 1.0) + return self.at_goal_prob + frac * (float(self.at_goal_prob_final) - self.at_goal_prob) + + def __call__( + self, + env: ManagerBasedEnv, + env_ids: torch.Tensor, + plug_cfg: SceneEntityCfg | None = None, + socket_cfg: SceneEntityCfg | None = None, + at_goal_prob: float = 0.8, + insertion_axis: list | None = None, + insertion_length: float = 0.02, + socket_insertion_offset: list | None = None, + plug_insertion_offset: list | None = None, + goal_rot: list | None = None, + normal_pose_range: dict | None = None, + at_goal_prob_final: float | None = None, + anneal_start_iter: float = 0.0, + anneal_end_iter: float | None = None, + num_steps_per_env: int | None = None, + at_goal_depth_range: list | None = None, + approach_depth_range: list | None = None, + ): + num_envs = len(env_ids) + + socket_pos = wp.to_torch(self.socket.data.root_pos_w)[env_ids] + socket_quat = wp.to_torch(self.socket.data.root_quat_w)[env_ids] + + # Compute socket keypoint origin in world frame + socket_offset_batch = self.socket_insertion_offset.unsqueeze(0).expand(num_envs, -1) + id_quat_batch = self.identity_quat.unsqueeze(0).expand(num_envs, -1) + kp_origin_w, _ = math_utils.combine_frame_transforms( + socket_pos, + socket_quat, + socket_offset_batch, + id_quat_batch, + ) + + # Insertion axis in world frame (rotated by socket orientation) + insertion_axis_w = math_utils.quat_apply(socket_quat, self.insertion_axis.unsqueeze(0).expand(num_envs, -1)) + + # Goal plug orientation in world frame + goal_quat_w = math_utils.quat_mul(socket_quat, self.goal_rot.unsqueeze(0).expand(num_envs, -1)) + + # Plug keypoint offset rotated into world frame (for converting kp pos -> root pos) + plug_offset_batch = self.plug_insertion_offset.unsqueeze(0).expand(num_envs, -1) + plug_kp_in_world = math_utils.quat_apply(goal_quat_w, plug_offset_batch) + + # At-goal probability for the current training progress (may be annealed). + current_at_goal_prob = self._current_at_goal_prob(env) + + if self.approach_depth_range is None: + pose_range = self.normal_pose_range + rand_pos = torch.zeros(num_envs, 3, device=env.device) + for i, key in enumerate(["x", "y", "z"]): + rng = pose_range.get(key, [0.0, 0.0]) + rand_pos[:, i] = torch.empty(num_envs, device=env.device).uniform_(rng[0], rng[1]) + + default_plug_pos = ( + wp.to_torch(self.plug.data.default_root_state)[env_ids, :3] + env.scene.env_origins[env_ids] + ) + normal_plug_pos = default_plug_pos + rand_pos + normal_plug_quat = wp.to_torch(self.plug.data.default_root_state)[env_ids, 3:7] + + plug_pos = normal_plug_pos.clone() + plug_quat = normal_plug_quat.clone() + + if current_at_goal_prob > 0.0 and num_envs > 0: + at_goal_mask = torch.rand(num_envs, device=env.device) < current_at_goal_prob + at_goal_local = at_goal_mask.nonzero(as_tuple=False).squeeze(-1) + num_at_goal = int(at_goal_local.numel()) + if num_at_goal > 0: + depth_rand = torch.rand(num_at_goal, 1, device=env.device) + goal_kp_pos = ( + kp_origin_w[at_goal_local] + + depth_rand * insertion_axis_w[at_goal_local] * self.insertion_length + ) + + plug_pos[at_goal_local] = goal_kp_pos - plug_kp_in_world[at_goal_local] + plug_quat[at_goal_local] = goal_quat_w[at_goal_local] + else: + at_goal_mask = torch.rand(num_envs, device=env.device) < current_at_goal_prob + + at_goal_range = ( + self.at_goal_depth_range if self.at_goal_depth_range is not None else [0.0, self.insertion_length] + ) + depth_at_goal = torch.empty(num_envs, device=env.device).uniform_( + float(at_goal_range[0]), float(at_goal_range[1]) + ) + depth_approach = torch.empty(num_envs, device=env.device).uniform_( + float(self.approach_depth_range[0]), float(self.approach_depth_range[1]) + ) + depth = torch.where(at_goal_mask, depth_at_goal, depth_approach) + + pose_range = self.normal_pose_range + rand_pos = torch.zeros(num_envs, 3, device=env.device) + for i, key in enumerate(["x", "y", "z"]): + rng = pose_range.get(key, [0.0, 0.0]) + rand_pos[:, i] = torch.empty(num_envs, device=env.device).uniform_(rng[0], rng[1]) + rand_pos[at_goal_mask] = 0.0 + + goal_kp_pos = kp_origin_w + depth.unsqueeze(-1) * insertion_axis_w + plug_pos = goal_kp_pos - plug_kp_in_world + rand_pos + plug_quat = goal_quat_w.clone() + + new_root_pose = torch.cat([plug_pos, plug_quat], dim=-1) + zero_vel = torch.zeros(num_envs, 6, device=env.device, dtype=torch.float32) + self.plug.write_root_pose_to_sim(new_root_pose, env_ids=env_ids) + self.plug.write_root_velocity_to_sim(zero_vel, env_ids=env_ids) diff --git a/source/isaaclab_tasks/isaaclab_tasks/contrib/deploy/mdp/observations.py b/source/isaaclab_tasks/isaaclab_tasks/contrib/deploy/mdp/observations.py index ac12d8b22f7..bb22740898d 100644 --- a/source/isaaclab_tasks/isaaclab_tasks/contrib/deploy/mdp/observations.py +++ b/source/isaaclab_tasks/isaaclab_tasks/contrib/deploy/mdp/observations.py @@ -3,24 +3,133 @@ # # SPDX-License-Identifier: BSD-3-Clause -"""Class-based observation terms for the gear assembly manipulation environment.""" +"""Class-based observation terms for manipulation deployment environments.""" from __future__ import annotations from typing import TYPE_CHECKING import torch +import warp as wp from isaaclab.managers import ManagerTermBase, ObservationTermCfg, SceneEntityCfg -from isaaclab.utils.math import combine_frame_transforms +from isaaclab.utils.leapp import ( + QUAT_XYZW_ELEMENT_NAMES, + XYZ_ELEMENT_NAMES, + InputKindEnum, +) +from isaaclab.utils.math import combine_frame_transforms, matrix_from_quat if TYPE_CHECKING: - from isaaclab.assets import RigidObject + from isaaclab.assets import Articulation, RigidObject from isaaclab.envs import ManagerBasedRLEnv from .events import randomize_gear_type +_LEAPP_TRACED_OBSERVATION_INPUTS = "_leapp_traced_observation_inputs" +_LEAPP_CONSUMED_OBSERVATION_INPUTS = "_leapp_consumed_observation_inputs" + + +def _tensor_data_to_torch(data) -> torch.Tensor: + """Return a torch tensor view for Isaac Lab data stored as torch or Warp-backed data.""" + return data.torch if hasattr(data, "torch") else wp.to_torch(data) + + +def _selected_joint_names(asset, joint_ids) -> list[str] | None: + """Return joint names selected by the observation config.""" + joint_names = getattr(asset, "joint_names", None) + if joint_names is None: + return None + if joint_ids is None or joint_ids == slice(None): + return list(joint_names) + if isinstance(joint_ids, slice): + return list(joint_names[joint_ids]) + if hasattr(joint_ids, "tolist"): + joint_ids = joint_ids.tolist() + return [joint_names[int(joint_id)] for joint_id in joint_ids] + + +def _is_leapp_export_env(env) -> bool: + """Return whether the observation is running under the LEAPP export proxy.""" + return type(env).__name__ == "_EnvProxy" + + +def _leapp_real_env(env): + """Return the wrapped Isaac Lab env when LEAPP passes an export proxy.""" + if _is_leapp_export_env(env): + return object.__getattribute__(env, "_real_env") + return env + + +def _set_leapp_traced_observation_input(env, name: str, tensor: torch.Tensor) -> None: + """Store a traced observation tensor for later export-only reuse.""" + if not _is_leapp_export_env(env): + return + real_env = _leapp_real_env(env) + traced_inputs = getattr(real_env, _LEAPP_TRACED_OBSERVATION_INPUTS, None) + if traced_inputs is None: + traced_inputs = {} + setattr(real_env, _LEAPP_TRACED_OBSERVATION_INPUTS, traced_inputs) + traced_inputs[name] = tensor + getattr(real_env, _LEAPP_CONSUMED_OBSERVATION_INPUTS, set()).discard(name) + + +def _deploy_object_input_base_name(asset_name: str) -> str: + """Return a stable deploy input base name for common socket/plug assets.""" + for prefix in ("dp_", "gb300_", "factory_"): + if asset_name.startswith(prefix): + return asset_name[len(prefix) :] + return asset_name + + +def joint_pos(env: ManagerBasedRLEnv, asset_cfg: SceneEntityCfg = SceneEntityCfg("robot")) -> torch.Tensor: + """Joint positions for the configured joints, exposed as the LEAPP input boundary.""" + real_env = _leapp_real_env(env) + asset = real_env.scene[asset_cfg.name] + selected_joint_pos = _tensor_data_to_torch(asset.data.joint_pos)[:, asset_cfg.joint_ids] + joint_names = _selected_joint_names(asset, asset_cfg.joint_ids) + if _is_leapp_export_env(env): + from leapp import annotate + from leapp.utils.tensor_description import TensorSemantics + + selected_joint_pos = annotate.input_tensors( + env.unwrapped.spec.id, + TensorSemantics( + name=f"{asset_cfg.name}_joint_pos", + ref=selected_joint_pos, + kind=InputKindEnum.JOINT_POSITION, + element_names=joint_names, + extra={"isaaclab_connection": f"state:{asset_cfg.name}:joint_pos"}, + ), + ) + _set_leapp_traced_observation_input(env, f"{asset_cfg.name}_joint_pos", selected_joint_pos) + return selected_joint_pos + + +def joint_vel(env: ManagerBasedRLEnv, asset_cfg: SceneEntityCfg = SceneEntityCfg("robot")) -> torch.Tensor: + """Joint velocities for the configured joints, exposed as the LEAPP input boundary.""" + real_env = _leapp_real_env(env) + asset = real_env.scene[asset_cfg.name] + selected_joint_vel = _tensor_data_to_torch(asset.data.joint_vel)[:, asset_cfg.joint_ids] + joint_names = _selected_joint_names(asset, asset_cfg.joint_ids) + if _is_leapp_export_env(env): + from leapp import annotate + from leapp.utils.tensor_description import TensorSemantics + + selected_joint_vel = annotate.input_tensors( + env.unwrapped.spec.id, + TensorSemantics( + name=f"{asset_cfg.name}_joint_vel", + ref=selected_joint_vel, + kind=InputKindEnum.JOINT_VELOCITY, + element_names=joint_names, + extra={"isaaclab_connection": f"state:{asset_cfg.name}:joint_vel"}, + ), + ) + return selected_joint_vel + + class gear_shaft_pos_w(ManagerTermBase): """Gear shaft position in world frame with offset applied. @@ -338,3 +447,259 @@ def __call__(self, env: ManagerBasedRLEnv) -> torch.Tensor: gear_positive_quat[w_negative] = -gear_quat[w_negative] return gear_positive_quat + + +class rigid_object_pos_w(ManagerTermBase): + """Rigid object position in the environment frame, with optional local-frame offset. + + Generic observation term that returns the position of any + :class:`~isaaclab.assets.RigidObject` in the environment frame. An optional + 3D offset can be applied in the object's local frame before subtracting the + environment origin. + + Args: + asset_cfg: The asset configuration. Required. + offset: A 3D offset ``[x, y, z]`` [m] applied in the object's local frame. + Defaults to ``[0, 0, 0]``. + + Returns: + Object position tensor in the environment frame, shape ``[num_envs, 3]`` [m]. + """ + + def __init__(self, cfg: ObservationTermCfg, env: ManagerBasedRLEnv): + super().__init__(cfg, env) + + if "asset_cfg" not in cfg.params: + raise ValueError("'asset_cfg' parameter is required in rigid_object_pos_w configuration.") + self.asset_cfg: SceneEntityCfg = cfg.params["asset_cfg"] + self.asset: RigidObject = env.scene[self.asset_cfg.name] + + offset = cfg.params.get("offset", [0.0, 0.0, 0.0]) + self.offset_tensor = torch.tensor(offset, device=env.device, dtype=torch.float32) + + 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 __call__( + self, + env: ManagerBasedRLEnv, + asset_cfg: SceneEntityCfg | None = None, + offset: list | None = None, + ) -> torch.Tensor: + real_env = _leapp_real_env(env) + asset = real_env.scene[self.asset_cfg.name] + obj_pos = _tensor_data_to_torch(asset.data.root_pos_w) + obj_quat = _tensor_data_to_torch(asset.data.root_quat_w) + + if torch.any(self.offset_tensor != 0): + offset_repeated = self.offset_tensor.unsqueeze(0).repeat(real_env.num_envs, 1) + obj_pos, _ = combine_frame_transforms(obj_pos, obj_quat, offset_repeated, self.identity_quat) + + obj_pos = obj_pos - real_env.scene.env_origins + if _is_leapp_export_env(env): + from leapp import annotate + from leapp.utils.tensor_description import TensorSemantics + + input_name = f"{_deploy_object_input_base_name(self.asset_cfg.name)}_pos" + obj_pos = annotate.input_tensors( + env.unwrapped.spec.id, + TensorSemantics( + name=input_name, + ref=obj_pos, + kind=InputKindEnum.BODY_POSITION, + element_names=XYZ_ELEMENT_NAMES, + extra={"isaaclab_connection": f"observation:policy:{input_name}"}, + ), + ) + return obj_pos + + +class rigid_object_quat_w(ManagerTermBase): + """Rigid object orientation in the world frame. + + Generic observation term that returns the orientation of any + :class:`~isaaclab.assets.RigidObject`. The quaternion is canonicalized so + that the ``w`` component is positive, reducing observation variation seen + by the policy. + + Args: + asset_cfg: The asset configuration. Required. + + Returns: + Object orientation as a quaternion ``(x, y, z, w)``, shape ``[num_envs, 4]``. + """ + + def __init__(self, cfg: ObservationTermCfg, env: ManagerBasedRLEnv): + super().__init__(cfg, env) + + if "asset_cfg" not in cfg.params: + raise ValueError("'asset_cfg' parameter is required in rigid_object_quat_w configuration.") + self.asset_cfg: SceneEntityCfg = cfg.params["asset_cfg"] + self.asset: RigidObject = env.scene[self.asset_cfg.name] + + def __call__( + self, + env: ManagerBasedRLEnv, + asset_cfg: SceneEntityCfg | None = None, + ) -> torch.Tensor: + real_env = _leapp_real_env(env) + obj_quat = _tensor_data_to_torch(real_env.scene[self.asset_cfg.name].data.root_quat_w) + if _is_leapp_export_env(env): + from leapp import annotate + from leapp.utils.tensor_description import TensorSemantics + + input_name = f"{_deploy_object_input_base_name(self.asset_cfg.name)}_quat" + obj_quat = annotate.input_tensors( + env.unwrapped.spec.id, + TensorSemantics( + name=input_name, + ref=obj_quat, + kind=InputKindEnum.BODY_ROTATION, + element_names=QUAT_XYZW_ELEMENT_NAMES, + extra={"isaaclab_connection": f"observation:policy:{input_name}"}, + ), + ) + + sign = torch.where(obj_quat[:, 3:4] < 0, -1.0, 1.0) + return obj_quat * sign + + +def _quat_to_rot_6d(quat: torch.Tensor) -> torch.Tensor: + """Convert quaternion (x, y, z, w) to 6D rotation (Zhou et al.). + + Takes the first two rows of the 3x3 rotation matrix and flattens + them into a 6-element vector per sample. + + Args: + quat: Quaternion tensor of shape ``(..., 4)`` in ``(x, y, z, w)`` format. + + Returns: + 6D rotation tensor of shape ``(..., 6)``. + """ + rot_mat = matrix_from_quat(quat) + batch_shape = rot_mat.shape[:-2] + return rot_mat[..., :2, :].clone().reshape(batch_shape + (6,)) + + +class rigid_object_rot_6d_w(ManagerTermBase): + """Rigid object 6D rotation in the world frame (Zhou et al.). + + Returns the first two rows of the 3x3 rotation matrix derived from + the object's root quaternion, giving a continuous 6D rotation + representation that avoids quaternion discontinuities. + + Args: + asset_cfg: The asset configuration. Required. + + Returns: + 6D rotation tensor, shape ``[num_envs, 6]``. + """ + + def __init__(self, cfg: ObservationTermCfg, env: ManagerBasedRLEnv): + super().__init__(cfg, env) + if "asset_cfg" not in cfg.params: + raise ValueError("'asset_cfg' parameter is required in rigid_object_rot_6d_w configuration.") + self.asset_cfg: SceneEntityCfg = cfg.params["asset_cfg"] + self.asset: RigidObject = env.scene[self.asset_cfg.name] + + def __call__( + self, + env: ManagerBasedRLEnv, + asset_cfg: SceneEntityCfg | None = None, + ) -> torch.Tensor: + obj_quat = wp.to_torch(self.asset.data.root_quat_w) + return _quat_to_rot_6d(obj_quat) + + +class eef_pos_w(ManagerTermBase): + """End-effector position in the environment frame. + + Gets the position of a specified body on a robot articulation and + returns it relative to the environment origin. An optional 3D offset can be + applied in the body's local frame, e.g. to report the gripper tool-center + point (TCP) rather than the raw flange. + + Args: + asset_cfg: The robot articulation configuration. Required. + body_name: Name of the end-effector body link. Required. + offset: A 3D offset ``[x, y, z]`` [m] applied in the body's local frame. + Defaults to ``[0, 0, 0]``. + + Returns: + EEF position tensor, shape ``[num_envs, 3]`` [m]. + """ + + def __init__(self, cfg: ObservationTermCfg, env: ManagerBasedRLEnv): + super().__init__(cfg, env) + if "asset_cfg" not in cfg.params: + raise ValueError("'asset_cfg' parameter is required in eef_pos_w configuration.") + if "body_name" not in cfg.params: + raise ValueError("'body_name' parameter is required in eef_pos_w configuration.") + + self.asset_cfg: SceneEntityCfg = cfg.params["asset_cfg"] + self.robot: Articulation = env.scene[self.asset_cfg.name] + self.body_name: str = cfg.params["body_name"] + self.body_idx = self.robot.find_bodies(self.body_name)[0][0] + + offset = cfg.params.get("offset", [0.0, 0.0, 0.0]) + self.offset_tensor = torch.tensor(offset, device=env.device, dtype=torch.float32) + 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 __call__( + self, + env: ManagerBasedRLEnv, + asset_cfg: SceneEntityCfg | None = None, + body_name: str | None = None, + offset: list | None = None, + ) -> torch.Tensor: + body_pos = wp.to_torch(self.robot.data.body_pos_w)[:, self.body_idx, :] + + if torch.any(self.offset_tensor != 0): + body_quat = wp.to_torch(self.robot.data.body_quat_w)[:, self.body_idx, :] + offset_repeated = self.offset_tensor.unsqueeze(0).repeat(env.num_envs, 1) + body_pos, _ = combine_frame_transforms(body_pos, body_quat, offset_repeated, self.identity_quat) + + return body_pos - env.scene.env_origins + + +class eef_rot_6d_w(ManagerTermBase): + """End-effector 6D rotation in the world frame (Zhou et al.). + + Gets the quaternion of a specified body on a robot articulation and + converts it to a continuous 6D rotation representation. + + Args: + asset_cfg: The robot articulation configuration. Required. + body_name: Name of the end-effector body link. Required. + + Returns: + 6D rotation tensor, shape ``[num_envs, 6]``. + """ + + def __init__(self, cfg: ObservationTermCfg, env: ManagerBasedRLEnv): + super().__init__(cfg, env) + if "asset_cfg" not in cfg.params: + raise ValueError("'asset_cfg' parameter is required in eef_rot_6d_w configuration.") + if "body_name" not in cfg.params: + raise ValueError("'body_name' parameter is required in eef_rot_6d_w configuration.") + + self.asset_cfg: SceneEntityCfg = cfg.params["asset_cfg"] + self.robot: Articulation = env.scene[self.asset_cfg.name] + self.body_name: str = cfg.params["body_name"] + self.body_idx = self.robot.find_bodies(self.body_name)[0][0] + + def __call__( + self, + env: ManagerBasedRLEnv, + asset_cfg: SceneEntityCfg | None = None, + body_name: str | None = None, + ) -> torch.Tensor: + body_quat = wp.to_torch(self.robot.data.body_quat_w)[:, self.body_idx, :] + return _quat_to_rot_6d(body_quat) 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 c776168e5b4..1c1d763a894 100644 --- a/source/isaaclab_tasks/isaaclab_tasks/contrib/deploy/mdp/rewards.py +++ b/source/isaaclab_tasks/isaaclab_tasks/contrib/deploy/mdp/rewards.py @@ -568,6 +568,139 @@ def __call__( return scaled_reward +class keypoint_two_body_error(ManagerTermBase): + """Keypoint distance between two rigid objects with body-frame offsets. + + Computes keypoint frames for each object by applying a local-frame offset + (translation and optional rotation) to each object's root pose, then + measures the mean keypoint distance between the two frames. + + This handles USD assets whose root frame doesn't coincide with the + functionally relevant geometry (e.g., the GB300 socket whose root frame + is far from the actual insertion slot). + + Args: + asset_cfg_1: Scene entity config for the first rigid object (e.g., socket). + asset_cfg_2: Scene entity config for the second rigid object (e.g., plug). + offset_1: 3D offset ``[x, y, z]`` in asset 1's local frame to its keypoint + origin. Defaults to ``[0, 0, 0]``. + offset_2: 3D offset ``[x, y, z]`` in asset 2's local frame to its keypoint + origin. Defaults to ``[0, 0, 0]``. + rot_offset_2: Quaternion ``(x, y, z, w)`` rotation applied to asset 2's + keypoint frame orientation. Typically the inverse of the goal rotation + so that keypoint frames align at the insertion goal. Defaults to identity. + keypoint_scale: Scale factor for keypoint offsets. Defaults to 0.15. + """ + + def __init__(self, cfg: RewardTermCfg, env: ManagerBasedRLEnv): + super().__init__(cfg, env) + + self.asset_1 = env.scene[cfg.params["asset_cfg_1"].name] + self.asset_2 = env.scene[cfg.params["asset_cfg_2"].name] + + offset_1 = cfg.params.get("offset_1", [0.0, 0.0, 0.0]) + self.offset_1 = torch.tensor(offset_1, device=env.device, dtype=torch.float32) + + offset_2 = cfg.params.get("offset_2", [0.0, 0.0, 0.0]) + self.offset_2 = torch.tensor(offset_2, device=env.device, dtype=torch.float32) + + rot_offset_2 = cfg.params.get("rot_offset_2", [0.0, 0.0, 0.0, 1.0]) + self.rot_offset_2 = ( + torch.tensor(rot_offset_2, device=env.device, dtype=torch.float32).unsqueeze(0).repeat(env.num_envs, 1) + ) + + self.identity_quat = torch.tensor([[0.0, 0.0, 0.0, 1.0]], device=env.device, dtype=torch.float32).repeat( + env.num_envs, 1 + ) + + self.keypoint_computer = _compute_keypoint_distance(cfg, env) + + def _get_kp_frames(self, env: ManagerBasedRLEnv) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + """Compute keypoint frames for both assets with offsets applied.""" + import warp as wp + + pos_1 = wp.to_torch(self.asset_1.data.root_pos_w) + quat_1 = wp.to_torch(self.asset_1.data.root_quat_w) + pos_2 = wp.to_torch(self.asset_2.data.root_pos_w) + quat_2 = wp.to_torch(self.asset_2.data.root_quat_w) + + offset_1_batch = self.offset_1.unsqueeze(0).expand(env.num_envs, -1) + kp_pos_1, kp_quat_1 = combine_frame_transforms(pos_1, quat_1, offset_1_batch, self.identity_quat) + + offset_2_batch = self.offset_2.unsqueeze(0).expand(env.num_envs, -1) + kp_pos_2, kp_quat_2 = combine_frame_transforms(pos_2, quat_2, offset_2_batch, self.rot_offset_2) + + return kp_pos_1, kp_quat_1, kp_pos_2, kp_quat_2 + + def __call__( + self, + env: ManagerBasedRLEnv, + asset_cfg_1: SceneEntityCfg = SceneEntityCfg("factory_gear_base"), + asset_cfg_2: SceneEntityCfg = SceneEntityCfg("factory_gear_small"), + keypoint_scale: float = 0.15, + offset_1: list | None = None, + offset_2: list | None = None, + rot_offset_2: list | None = None, + ) -> torch.Tensor: + kp_pos_1, kp_quat_1, kp_pos_2, kp_quat_2 = self._get_kp_frames(env) + + keypoint_dist_sep = self.keypoint_computer.compute( + current_pos=kp_pos_1, + current_quat=kp_quat_1, + target_pos=kp_pos_2, + target_quat=kp_quat_2, + keypoint_scale=keypoint_scale, + ) + + return keypoint_dist_sep.mean(-1) + + +class keypoint_two_body_error_exp(keypoint_two_body_error): + """Exponential keypoint reward between two rigid objects with body-frame offsets. + + Same offset logic as :class:`keypoint_two_body_error`, but applies an + exponential reward transformation for sharper shaping near the goal. + """ + + def __call__( + self, + env: ManagerBasedRLEnv, + asset_cfg_1: SceneEntityCfg = SceneEntityCfg("factory_gear_base"), + asset_cfg_2: SceneEntityCfg = SceneEntityCfg("factory_gear_small"), + kp_exp_coeffs: list[tuple[float, float]] = [(1.0, 0.1)], + kp_use_sum_of_exps: bool = True, + keypoint_scale: float = 0.15, + offset_1: list | None = None, + offset_2: list | None = None, + rot_offset_2: list | None = None, + ) -> torch.Tensor: + kp_pos_1, kp_quat_1, kp_pos_2, kp_quat_2 = self._get_kp_frames(env) + + keypoint_dist_sep = self.keypoint_computer.compute( + current_pos=kp_pos_1, + current_quat=kp_quat_1, + target_pos=kp_pos_2, + target_quat=kp_quat_2, + keypoint_scale=keypoint_scale, + ) + + keypoint_reward_exp = torch.zeros_like(keypoint_dist_sep[:, 0]) + + if kp_use_sum_of_exps: + for coeff in kp_exp_coeffs: + a, b = coeff + keypoint_reward_exp += ( + 1.0 / (torch.exp(a * keypoint_dist_sep) + b + torch.exp(-a * keypoint_dist_sep)) + ).mean(-1) + else: + keypoint_dist = keypoint_dist_sep.mean(-1) + for coeff in kp_exp_coeffs: + a, b = coeff + keypoint_reward_exp += 1.0 / (torch.exp(a * keypoint_dist) + b + torch.exp(-a * keypoint_dist)) + + return keypoint_reward_exp + + ## # Helper functions and classes ## 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 9f217c31637..983541859f2 100644 --- a/source/isaaclab_tasks/isaaclab_tasks/contrib/deploy/mdp/terminations.py +++ b/source/isaaclab_tasks/isaaclab_tasks/contrib/deploy/mdp/terminations.py @@ -11,6 +11,7 @@ from typing import TYPE_CHECKING import torch +import warp as wp import isaaclab.utils.math as math_utils from isaaclab.managers import ManagerTermBase, SceneEntityCfg, TerminationTermCfg @@ -18,7 +19,7 @@ logger = logging.getLogger(__name__) if TYPE_CHECKING: - from isaaclab.assets import Articulation + from isaaclab.assets import Articulation, RigidObject from isaaclab.envs import ManagerBasedEnv from .events import randomize_gear_type @@ -330,3 +331,184 @@ def __call__( ) return self.reset_flags + + +class reset_when_plug_dropped(ManagerTermBase): + """Check if a held plug/connector has fallen out of the gripper. + + Generic single-object counterpart of :class:`reset_when_gear_dropped`. + Computes the distance between the end effector and the plug's expected + grasp position (derived from the plug's pose and the grasp offset). + If the distance exceeds ``distance_threshold``, the environment is reset. + + Args: + robot_asset_cfg: Robot asset configuration. Defaults to + ``SceneEntityCfg("robot")``. + plug_asset_cfg: Plug rigid-object asset configuration. + distance_threshold: Distance [m] above which the plug is considered + dropped. + end_effector_body_name: End-effector body name on the robot. + grasp_offset: Position offset ``[x, y, z]`` [m] from the plug origin + to the expected grasp point in the (rotated) plug frame. + grasp_rot_offset: Quaternion offset ``(x, y, z, w)`` applied to the + plug orientation to define the grasp frame. + """ + + def __init__(self, cfg: TerminationTermCfg, env: ManagerBasedEnv): + super().__init__(cfg, env) + + self.robot_asset_cfg: SceneEntityCfg = cfg.params.get("robot_asset_cfg", SceneEntityCfg("robot")) + self.robot_asset: Articulation = env.scene[self.robot_asset_cfg.name] + + if "plug_asset_cfg" not in cfg.params: + raise ValueError("'plug_asset_cfg' is required in reset_when_plug_dropped configuration.") + self.plug_asset_cfg: SceneEntityCfg = cfg.params["plug_asset_cfg"] + self.plug_asset: RigidObject = env.scene[self.plug_asset_cfg.name] + + for required in ("end_effector_body_name", "grasp_rot_offset", "grasp_offset"): + if required not in cfg.params: + raise ValueError(f"'{required}' is required in reset_when_plug_dropped configuration.") + + self.end_effector_body_name: str = cfg.params["end_effector_body_name"] + + grasp_offset = cfg.params["grasp_offset"] + self.grasp_offset_tensor = torch.tensor(grasp_offset, device=env.device, dtype=torch.float32) + + grasp_rot_offset = cfg.params["grasp_rot_offset"] + self.grasp_rot_offset_tensor = ( + torch.tensor(grasp_rot_offset, device=env.device, dtype=torch.float32).unsqueeze(0).repeat(env.num_envs, 1) + ) + + self.reset_flags = torch.zeros(env.num_envs, dtype=torch.bool, device=env.device) + + eef_indices, _ = self.robot_asset.find_bodies([self.end_effector_body_name]) + if len(eef_indices) == 0: + logger.warning( + f"{self.end_effector_body_name} not found in robot body names. Cannot check plug drop condition." + ) + self.eef_idx = None + else: + self.eef_idx = eef_indices[0] + + def __call__( + self, + env: ManagerBasedEnv, + distance_threshold: float = 0.1, + robot_asset_cfg: SceneEntityCfg = SceneEntityCfg("robot"), + plug_asset_cfg: SceneEntityCfg | None = None, + grasp_offset: list | None = None, + end_effector_body_name: str | None = None, + grasp_rot_offset: list | None = None, + ) -> torch.Tensor: + self.reset_flags.fill_(False) + + if self.eef_idx is None: + return self.reset_flags + + eef_pos_world = wp.to_torch(self.robot_asset.data.body_link_pos_w)[:, self.eef_idx] + + plug_pos_world = wp.to_torch(self.plug_asset.data.root_link_pos_w) + plug_quat_world = wp.to_torch(self.plug_asset.data.root_link_quat_w) + + plug_quat_world = math_utils.quat_mul(plug_quat_world, self.grasp_rot_offset_tensor) + + grasp_offset_batch = self.grasp_offset_tensor.unsqueeze(0).expand(plug_pos_world.shape[0], -1) + plug_grasp_pos_world = plug_pos_world + math_utils.quat_apply(plug_quat_world, grasp_offset_batch) + + distances = torch.linalg.norm(plug_grasp_pos_world - eef_pos_world, dim=-1) + self.reset_flags[:] = distances > distance_threshold + + return self.reset_flags + + +class reset_when_plug_orientation_exceeded(ManagerTermBase): + """Check if a held plug/connector has rotated too much relative to the end effector. + + Generic single-object counterpart of + :class:`reset_when_gear_orientation_exceeds_threshold`. Computes the + relative orientation between the plug (with grasp rotation offset + applied) and the end effector. If the roll, pitch, or yaw angles exceed + the configured thresholds, the environment is reset. + + Args: + robot_asset_cfg: Robot asset configuration. Defaults to + ``SceneEntityCfg("robot")``. + plug_asset_cfg: Plug rigid-object asset configuration. + roll_threshold_deg: Roll threshold [deg]. + pitch_threshold_deg: Pitch threshold [deg]. + yaw_threshold_deg: Yaw threshold [deg]. + end_effector_body_name: End-effector body name on the robot. + grasp_rot_offset: Quaternion offset ``(x, y, z, w)`` applied to the + plug orientation to define the grasp frame. + """ + + def __init__(self, cfg: TerminationTermCfg, env: ManagerBasedEnv): + super().__init__(cfg, env) + + self.robot_asset_cfg: SceneEntityCfg = cfg.params.get("robot_asset_cfg", SceneEntityCfg("robot")) + self.robot_asset: Articulation = env.scene[self.robot_asset_cfg.name] + + if "plug_asset_cfg" not in cfg.params: + raise ValueError("'plug_asset_cfg' is required in reset_when_plug_orientation_exceeded configuration.") + self.plug_asset_cfg: SceneEntityCfg = cfg.params["plug_asset_cfg"] + self.plug_asset: RigidObject = env.scene[self.plug_asset_cfg.name] + + for required in ("end_effector_body_name", "grasp_rot_offset"): + if required not in cfg.params: + raise ValueError(f"'{required}' is required in reset_when_plug_orientation_exceeded configuration.") + + self.end_effector_body_name: str = cfg.params["end_effector_body_name"] + + grasp_rot_offset = cfg.params["grasp_rot_offset"] + self.grasp_rot_offset_tensor = ( + torch.tensor(grasp_rot_offset, device=env.device, dtype=torch.float32).unsqueeze(0).repeat(env.num_envs, 1) + ) + + self.reset_flags = torch.zeros(env.num_envs, dtype=torch.bool, device=env.device) + + eef_indices, _ = self.robot_asset.find_bodies([self.end_effector_body_name]) + if len(eef_indices) == 0: + logger.warning( + f"{self.end_effector_body_name} not found in robot. Cannot check plug orientation condition." + ) + self.eef_idx = None + else: + self.eef_idx = eef_indices[0] + + def __call__( + self, + env: ManagerBasedEnv, + roll_threshold_deg: float = 30.0, + pitch_threshold_deg: float = 30.0, + yaw_threshold_deg: float = 180.0, + robot_asset_cfg: SceneEntityCfg = SceneEntityCfg("robot"), + plug_asset_cfg: SceneEntityCfg | None = None, + end_effector_body_name: str | None = None, + grasp_rot_offset: list | None = None, + ) -> torch.Tensor: + self.reset_flags.fill_(False) + + if self.eef_idx is None: + return self.reset_flags + + roll_threshold_rad = torch.deg2rad(torch.tensor(roll_threshold_deg, device=env.device)) + pitch_threshold_rad = torch.deg2rad(torch.tensor(pitch_threshold_deg, device=env.device)) + yaw_threshold_rad = torch.deg2rad(torch.tensor(yaw_threshold_deg, device=env.device)) + + eef_quat_world = wp.to_torch(self.robot_asset.data.body_link_quat_w)[:, self.eef_idx] + + plug_quat_world = wp.to_torch(self.plug_asset.data.root_link_quat_w) + plug_quat_world = math_utils.quat_mul(plug_quat_world, self.grasp_rot_offset_tensor) + + eef_quat_inv = math_utils.quat_conjugate(eef_quat_world) + relative_quat = math_utils.quat_mul(plug_quat_world, eef_quat_inv) + + roll, pitch, yaw = math_utils.euler_xyz_from_quat(relative_quat) + + self.reset_flags[:] = ( + (torch.abs(roll) > roll_threshold_rad) + | (torch.abs(pitch) > pitch_threshold_rad) + | (torch.abs(yaw) > yaw_threshold_rad) + ) + + return self.reset_flags