Support for DDIMScheduler in Diffusion Policy

[kashyapakshay]

What this does

This PR adds DDIMScheduler as a choice for Diffusion Policy. DDPM can be too slow for real-world rollouts (takes ~0.7s on a 3080 with 100 timesteps, n_obs=2 and 216x288 images). This aims to get closer to real-time control (as is also mentioned in the DP paper).

How it was tested

I ran the standard test suite + diffusion-pushT training for some steps (screenshot attached).
I have not done a full training run and do not have performance numbers for this.

How to checkout & try? (for the reviewer)

Add/update the noise_scheduler_type field in your diffusion config yaml.

# Noise scheduler.
noise_scheduler_type: DDIM
num_train_timesteps: 100
num_inference_timesteps: 10 # optional but recommended

[alexander-soare]

Hi! Thanks so much for this PR. I have one request. A couple of things to discuss:

1. It's not passing style requirements. There's a section in CONTRIBUTING.md "6. Follow our style." that you can follow to set up ruff and a pre-commit hook.
2. This is great. There's just one slight hiccup that makes it less accessible. You can't use it with pretrained models without patching the code. That's because the make_policy calls PytorchModelHubMixin.from_pretrained therefore ignoring any command line overrides you add. For example I did python lerobot/scripts/eval.py -p lerobot/diffusion_pusht eval.n_episodes=10 eval.batch_size=10 policy.num_inference_steps=10 +policy.noise_scheduler_type=DDIM and my policy arguments were ignored. We can handle this in another PR. (btw just left this with huggingface_hub as it should help Don't override 'config' in model_kwargs huggingface_hub#2274) (btw I have a temporary fix in the works! Override pretrained model config #147).
3. Just FYI, it's possible we will remove the dependency on diffusers in which case we'll port the noise scheduler code into modeling_diffusion.py. We'll keep the DDIM feature in place though :)

[alexander-soare]

Could you please add a dosctring argument for this (preserving ordering)?

[kashyapakshay]

Done

[kashyapakshay]

Thanks @alexander-soare!

1. Oops, fixed, passing pre-commit checks now.
2. I merged in your change. anything else I need to do there?
3. Good to know. I added a little scheduler factory which may make that switch a little easier.

[alexander-soare]

@kashyapakshay There are still style issues somehow. Curious, did you follow the instructions in CONTRIBUTING? If so, we may need to fix the instructions! cc @aliberts I noticed that the style issue triggers on a yaml config file. Does the pre-commit handle that?

Since we are doing another round I have a nitty nit for you please: can you please move the noise scheduler factory so that the main policy class stays as the first definition in the file?

[kashyapakshay]

@alexander-soare I did follow the instructions and seems like ruff passed?

Anyway, fixed that style issue and also moved the factory func (I like that placement too but wasnt sure if you had were trying to ordering methods before classes)

[aliberts]

Since pre-commits hooks only check/affect staged files, if commits were pushed before installing the pre-commits then they can stay there since subsequent changes may not stage the same files.

To solve this, run pre-commit run -a to format every files (not just the staged ones), then re-stage, commit & push.

[kashyapakshay]

ah right, sorry. Will keep that in mind. All should pass now.

[alexander-soare]

Approved!