Towards Effective Utilization of Mixed-Quality Demonstrations in Robotic Manipulation via Segment-Level Selection and Optimization

[Paper] [Project Page]

🧑🏻‍💻 Run

For the representation model training stage, run the command bash command_train.sh to execute the data training script, which will preprocess the dataset and training the model.

Here are the argument explanations in the training process:

• --dataset : Specifies the entire dataset used for the representation model training.
• --aug_path : The path where the results of the augmented dataset will be stored.
• --save_mode ：Indicates the format or type of the dataset.
• --size : Specifies the size to which the images will be resized.
• --numbers ：The index or specific identifier used for data augmentation within the dataset.

For the eval stage, run the command bash command_eval.sh to complete the segment selection and trajectory optimization processes.

Here are the argument explanations in the evaluation process:

• --train_data_folder : The dataset used for distance-weighted voting during the segment selection process.
• --val_data_folder : The folder containing the full mixed-quality demonstration dataset for validation.
• --size : Specifies the size to which the images will be resized.

🤖 Training Manipulation Policy

After Select Segments to Imitate (S2I), the dataset can be directly used for downstream manipulation policy training as a plug-and-play solution.

For simulation experiments, we use the state-based BC-RNN and the Diffusion Policy (DP) that can be applied to both state and image data as robot manipulation policies. For real-world experiments, we choose DP and ACT as our image-based policies, as well as RISE as our point-cloud-based policy. Some minor modifications have been made to the sampler and rollout functions. The modified Python file is available in ./policy. Refer to the documentation for more details.

🙏 Acknowledgement

Our code is built upon: Diffusion Policy, RoboMimic, SupContrast, RISE and ACT. We thank all the authors for the contributions to the community.

✍️ Citation

If you find S2I useful in your research, please consider citing the following paper:

@article{
  chen2024towards,
  title = {Towards Effective Utilization of Mixed-Quality Demonstrations in Robotic Manipulation via Segment-Level Selection and Optimization},
  author = {Chen, Jingjing and Fang, Hongjie and Fang, Hao-Shu and Lu, Cewu},
  journal = {arXiv preprint arXiv:2409.19917},
  year = {2024}
}      

📃 License

S2I by Jingjing Chen, Hongjie Fang, Hao-Shu Fang, Cewu Lu is licensed under MIT License.