DSP: Diffusion Stabilizer Policy for Automated Surgical Robot Manipulations

The official implementation of the paper "Diffusion Stabilizer Policy for Automated Surgical Robot Manipulations".

this repo implemented the training and evaluation of DSP in SurRoL, the surgical environment.

Preparation

1. clone the repository

git clone --recursive https://github.com/Rappy325/DSP.git
cd DSP

2. Create a virtual environment

conda create -n dsp python=3.10 -y
conda activate dsp

3. install packages

pip install -e SurRoL/	# install surrol environments
pip install -e CleanDiffuser/
pip install -r requirements.txt

4. install torch

Install PyTorch 2.2.2 Following the Official Guideline using Conda.

For default, use the command:

pip3 install torch torchvision torchaudio

5. Then add one line of code at the top of gym/gym/envs/__init__.py to register SurRoL tasks:

# directory: anaconda3/envs/dsp/lib/python3.10/site-packages/
import surrol.gym

Usage

training

DSP (noise-aware training)

This repo provides pipelines for the two category of noise.

Dataset is generated by script, not manually collected. We use shell scripts to loop over SurRoL tasks and imperfect trajectory types, then calls data_generation.py to roll out oracle-policy demonstrations. The generated files are saved under datasets/demo/ (or datasets/video/ when video recording is enabled).

1. Action-level noise pipeline

Action-level noise is injected during data generation/training in the corresponding training script. You can directly edit the script and run it.

To test the main result of DSP, run the following script:

bash scripts/training/train_action-level.sh

Scripts for ablation studies are provided for reference as well.

• train_action-level_noise_setups.sh is for testing DSP on different noise strength and density
• train_thresholds_action-level.sh is for testing different thresholding method.

bash scripts/training/ablation/train_action-level_noise_setups.sh
bash scripts/training/ablation/train_thresholds_action-level.sh

2. Trajectory-level noise pipeline

Trajectory-level noise training requires you to run two scripts:

1. Generate trajectory-level noisy datasets first.
2. Run training using the generated datasets.

please edit the shell if you want to tweak the configurations.

Stage 1: dataset generation

To generate data, please run the following command.

bash scripts/training/traj-level_data_generate.sh

• if you want to design and try out your own traj-level perturbation setup, the relevant files are being defined in the wrapper files in dsp/surrol_wrappers/tasks. You might have to define your own trajectory-level imperfect type in there and do some testing.

Stage 2: training

bash scripts/training/train_traj-level.sh

Simularly, scripts for ablation studies for Trajectory-level noises are provided for reference as well.

bash scripts/training/ablation/train_thresholds_traj-level.sh

Notes:

• In trajectory-level generation, task-specific imperfect types are configured in scripts/training/traj-level_data_generate.sh.
• The imperfect type is passed from shell to SurRoL via environment variable (TRAJ_IMPERFECT_TYPE), and SurRoL selects the corresponding waypoint pattern for each task.
• Some imperfect types are configured in SurRoL to enforce success-only data collection automatically.

original DP

if you are only testing for the original diffusion policy, use train_clean.sh instead.

edit the parameters in scripts/training/train_clean.sh and run in terminal:

bash scripts/training/train_clean.sh

evaluation

Evaluates the model's performance on the provided dataset and logs the evaluation metrics. The evaluation process includes calculating the success rates with five different seeds, which returns 5 success rates for calculating IQM. When you run the shell scripts below, log files containing the evaluation results will be saved in the 'logs/eval/' directory by default.

DSP

Evaluation is provided with the same two pipelines as training.

1. Action-level noise evaluation

Edit the script and run:

bash scripts/evaluation/eval_action-level.sh

Ablation scripts for action-level evaluation:

• eval_action-level_noise_setups.sh evaluates DSP on different action-level noise strength/density setups (e.g., different (sigma, noisy_steps) combinations).
• eval_thresholds_action-level.sh evaluates different adaptive thresholding methods (otsu, k-means, logGMM, strict, loose, and normal-online baseline) under action-level noise.

bash scripts/evaluation/ablation/eval_action-level_noise_setups.sh
bash scripts/evaluation/ablation/eval_thresholds_action-level.sh

2. Trajectory-level noise evaluation

Edit the script and run:

bash scripts/evaluation/eval_traj-level.sh

Ablation script for trajectory-level thresholding:

bash scripts/evaluation/ablation/eval_thresholds_traj-level.sh

original DP

if you are only testing for the original diffusion policy, use eval_clean.sh instead.

bash scripts/evaluation/eval_clean.sh

Others

online filter test

script for figure 3 is provided as testFilter.sh in scripts/training/ablation. This is used to evaluate how the noisy-sample filter behaves during stage-2 online training. It loads stage-1 checkpoints and runs diffusion_trainer.py with --test_filter and --filter_noisy so you can measure filter-related metrics while training progresses, instead of only checking final task success. By looping over all six SurRoL tasks and multiple clean/noisy dataset combinations, it provides a consistent way to compare filter quality and stability across different environments and data scales.

Before running:

• make sure stage-1 checkpoints already exist for each (env, clean_num, noisy_num) combination
• make sure the corresponding datasets are available (the script uses noisy dataset names with noise_02_10)

Run command:

bash scripts/training/ablation/testFilter.sh

The script will call diffusion_trainer.py with --test_filter and save logs/checkpoints using the configured --msg naming pattern.