search

real-stanford / diffusion_policy

[RSS 2023] Diffusion Policy Visuomotor Policy Learning via Action Diffusion
https://diffusion-policy.cs.columbia.edu/
MIT License
1.1k stars 206 forks source link

Any suggestions on speeding up training? #23

Closed WangYixuan12 closed 9 months ago

WangYixuan12 commented 9 months ago

Hi Cheng, this work is incredible and elegant! I tried the code, and I found that the training time for each task and method is around 12 hrs. I am also testing on the robomimic can task, and it seems to take a while to get a performant policy. I am wondering whether you have any suggestions for speeding up the training process.

cheng-chi commented 9 months ago

Hi @WangYixuan12, 12 hrs training time sounds about right! If you need futher training acceleration, you can try to:

  1. reduce the number of training diffusion iterations (e.g. 50 or 25)
  2. reduce the frequency of simluation rollouts evaluations (they are very slow and can sometimes take more time than training)
  3. try improving policy runtime using torch.compile
WangYixuan12 commented 9 months ago

Thank you for your helpful tips! I will try them.