Skip to content

YuweiWang2002/PoseInsert

Repository files navigation

📈 Exploring Pose-Guided Imitation Learning for Robotic Precise Insertion

[Paper]

video.mp4

🛫 Getting Started

💻 Installation

Please follow the instructions to install the conda environments, as well as the real robot environments. We recommend using CUDA 11.8 during installations to avoid compatibility issues. Also, remember to adjust the Hand-Eye-Calibration in eval_ros_Pose.py and eval_ros_RPDP.py according to your own environment.

  1. Create a new conda environment and activate the environment.

    conda create -n PoseInsert python=3.8
    conda activate PoseInsert
  2. Manually install cudatoolkit, then install necessary dependencies.

    pip install -r requirements.txt

📷 Calibration

First of all, we use the Cobot Mobile ALOHA, manufactured by agilex.ai. Please calibrate the camera with the robot before data collection and evaluation to ensure correct spatial transformations between camera and the robot.

  1. Place the collect_data/ros_pose_gripper.py to FoundationPose. Place the gripper in front of camera.
    python ros_pose_gripper.py # publish the gripper pose
    python calibation/calibation_fk.py # get the camera_in_base

📷 Collect

Human demonstrations is collected.

  1. Place the collect_data/ros_pose_source2.py and collect_data/ros_pose_target2.pyto FoundationPose. And get the source/target object pose.

     python ros_pose_source2.py # publish the source object pose
     python ros_pose_target2.py # publish the target object pose
  2. Collect the train data. Remember to adjust the save_dir in collect_pose.py .

     python collect_pose.py --idx 0
  3. Show the train data.

     python replay_with_workspace.py 
  4. Show the train trajectories.

     python vis_traj.py 

📷 Train and Test

  1. Get the workspace for normalization.

     python get_workspace.py 
  2. Train the model. Remember to adjust the data_path and ckpt_dir.

     python train_pose.py 
     python train_RPDP.py 
  3. Test. Remember to adjust the ckpt .

     python eval_pose.py 
     python eval_RPDP.py 
  4. Test in the real-world. Before start the policy, the robot should grasp the source object.

     python device/robot_bringup.py # start up the robot
     python eval_ros_Pose.py 
     python eval_ros_RPDP.py 

📈 PoseInsert Policy

The PoseInsert policy consists of (1) a pose encoder (policy/cnn.py), (2) a RGBD encoder (policy/cnn.py), (3) a Pose-Guided Residual Gated Fusion (policy/cnn.py) and (4) a diffusion action decoder (policy/diffusion.py).

🙏 Acknowledgement

  • Our diffusion module is adapted from RISE.
  • Our RGBD encoder is adapted from FoundationPose. .

✍️ Citation

@article{sun2025exploringposeguidedimitationlearning,
    title   = {Exploring Pose-Guided Imitation Learning for Robotic Precise Insertion},
    author  = {Han Sun and Yizhao Wang and Zhenning Zhou and Shuai Wang and Haibo Yang and Jingyuan Sun and Qixin Cao},
    journal = {arXiv preprint arXiv:2404.12281},
    year    = {2025}
}

📃 License

PoseInsert (including data and codebase) by is licensed under CC BY-NC-SA 4.0

About

No description, website, or topics provided.

Resources

Stars

Watchers

Forks

Releases

No releases published

Packages

 
 
 

Contributors

Languages