Adds example for gear assembly sim-to-real with UR10e

[ashwinvkNV]

rl-video-step-137600.mp4

Description

This PR introduces a new Gear Assembly manipulation task for sim-to-real training with the UR10e robot arm. This environment enables training policies for precise gear insertion tasks using reinforcement learning, with comprehensive sim-to-real transfer capabilities.

Summary of Changes

New Features

• Gear Assembly Environment: Complete environment implementation for gear insertion tasks

  • Environment configuration (gear_assembly_env_cfg.py)
  • UR10e-specific joint position control configuration (joint_pos_env_cfg.py)
  • RSL-RL PPO training configuration (rsl_rl_ppo_cfg.py)

• MDP Components: Task-specific observation, reward, termination, and event functions

  • mdp/events.py: Randomization and reset events for robust training
  • mdp/observations.py: State observation functions
  • mdp/rewards.py: Reward shaping for gear insertion
  • mdp/terminations.py: Episode termination conditions

• Noise Models: Enhanced noise simulation for domain randomization

  • Added configurable noise models (noise_model.py, noise_cfg.py)
  • Integration with observation and action spaces for realistic sim-to-real transfer

Documentation

• Sim-to-Real Training Walkthrough: Complete guide for training and deploying the gear assembly task
  • Step-by-step training instructions
  • Real robot deployment guidelines
  • Visual assets (GIFs and screenshots)

Core Enhancements

• Training Script: Enhanced train.py with additional logging and configuration options
• UR10e Robot Configuration: Updated universal_robots.py with gear assembly specific parameters
• Reward System: Extended core reward functions in isaaclab/envs/mdp/rewards.py
• RL Configuration: Updated RSL-RL integration (rl_cfg.py, setup.py)

Type of change

• New feature (non-breaking change which adds functionality)
• Documentation update

Checklist

• I have read and understood the contribution guidelines
• I have run the pre-commit checks with ./isaaclab.sh --format
• I have made corresponding changes to the documentation
• My changes generate no new warnings
• I have added tests that prove my fix is effective or that my feature works
• I have updated the changelog and the corresponding version in the extension's config/extension.toml file
• I have added my name to the CONTRIBUTORS.md or my name already exists there

Usage Example

# Train the gear assembly task
python scripts/reinforcement_learning/rsl_rl/train.py \
  --task Isaac-Deploy-GearAssembly-UR10e-2F140-ROS-Inference-v0 \
  --num_envs 256 \
  --headless

# Run inference with trained policy
python scripts/reinforcement_learning/rsl_rl/play.py \
  --task Isaac-Deploy-GearAssembly-UR10e-2F140-ROS-Inference-v0 \
  --num_envs 1 \
 --checkpoint <checkpoint_path>

[greptile-apps]

Greptile Summary

• Introduces complete gear assembly sim-to-real environment for UR10e with PPO/LSTM training supporting 2F-140 and 2F-85 Robotiq grippers
• Implements class-based MDP components with pre-cached tensors for efficient batch operations including dynamic gear type randomization, keypoint-based rewards, and IK-based grasp initialization
• Adds ResetSampledNoiseModel for domain randomization that samples noise once per episode reset rather than every step

Confidence Score: 4/5

• Safe to merge with minor style improvements recommended
• Well-structured implementation with comprehensive reward shaping, termination conditions, and domain randomization. Code follows IsaacLab patterns with class-based terms and proper tensor caching. Minor redundant operations in IK loop and temporary USD path workaround noted but non-critical.
• source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/deploy/mdp/events.py has redundant joint state reads in IK loop; source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/deploy/gear_assembly/config/ur_10e/joint_pos_env_cfg.py uses temporary USD path pending bug fix

Important Files Changed

Filename	Overview
source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/deploy/mdp/events.py	Implements gear type randomization and IK-based robot grasp pose initialization with pre-cached tensors for efficient batch operations; IK loop reads joint state redundantly on each iteration (line 232)
source/isaaclab/isaaclab/utils/noise/noise_model.py	Adds ResetSampledNoiseModel class that samples noise only during reset and applies it consistently throughout the episode
source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/deploy/gear_assembly/config/ur_10e/joint_pos_env_cfg.py	UR10e-specific configuration with 2F-140 and 2F-85 gripper support, gripper-specific joint setters, and IK-based grasp pose initialization; uses temporary USD path override (line 415)

Sequence Diagram

sequenceDiagram
    participant User
    participant TrainingScript
    participant Environment
    participant GearTypeManager
    participant RobotIK
    participant PPOAgent
    participant RewardManager

    User->>TrainingScript: run train.py with task config
    TrainingScript->>Environment: create env with UR10e gear assembly config
    Environment->>GearTypeManager: initialize RandomizeGearType event
    GearTypeManager->>Environment: register as _gear_type_manager
    Environment->>Environment: setup scene with robot and 3 gear types
    
    loop Training Episodes
        Environment->>GearTypeManager: reset - randomize gear type
        GearTypeManager->>Environment: set active gear per env
        Environment->>RobotIK: SetRobotToGraspPose event
        RobotIK->>RobotIK: run IK to compute grasp pose
        RobotIK->>Environment: update robot joint positions
        Environment->>Environment: RandomizeGearsAndBasePose event
        
        loop Episode Steps
            Environment->>PPOAgent: get observation (joint pos/vel, gear shaft pose)
            PPOAgent->>Environment: return action (delta joint positions)
            Environment->>Environment: apply action and step simulation
            Environment->>RewardManager: compute keypoint distance rewards
            RewardManager->>Environment: return reward signal
            Environment->>Environment: check terminations (gear dropped, orientation)
        end
        
        Environment->>PPOAgent: collect episode data
        PPOAgent->>PPOAgent: update policy with PPO
    end
    
    TrainingScript->>User: save trained model checkpoint

Loading

[greptile-apps]

_{22 files reviewed, 2 comments}

_{Edit Code Review Agent Settings | Greptile
React with 👍 or 👎 to share your feedback on this new summary format}

[kellyguo11]

could we try to avoid having large .gif files in the repo directly? we can upload them to the server if needed and referenced from docs.

[ooctipus]

Thanks for the edit and contribution : )),

I'd like to ask a high level questions why not put this PR in deploy folder we created for the reach eariler? Are we planing to add peg insert and nut thread as well? if that's the intention, it might be more benefitial to work a general structure with all of them, right now mdp seems just tailored to gearmesh.

[ooctipus]

Nice work :)

[Mayankm96]

Could there be an explanation to this choice? If this is something generally useful, maybe we should move it to the utils.noise module directly?

[ashwinvkNV]

I added a comment. It was decided after discussion with @ooctipus to not add it to the deafult noise_models. It was mainly because he felt it would not be used by other envs as it is and instead might confuse users.