feat: Enable torch.compile for DiffusionPolicy inference

[jashshah999]

What this does

(⚡️ Performance) Adds torch.compile support to DiffusionPolicy to optimize inference latency. This allows users to optionally enable PyTorch 2.0's compilation feature for the U-Net component, which is the computational bottleneck during diffusion inference.

Changes:

• Added compile_model: bool = False and compile_mode: str = "reduce-overhead" configuration flags to DiffusionConfig
• Applied torch.compile to the U-Net in DiffusionModel.__init__ when compile_model=True
• Defaults maintain backward compatibility (compilation disabled by default)

How it was tested

• Created benchmark script to measure inference latency with and without compilation
• Tested on CUDA with 50 iterations after proper warmup
• Verified backward compatibility (default behavior unchanged)
• Confirmed code compiles without errors

Performance Results:

• Baseline (no compile): 1796.57 ms average latency
• With torch.compile: 1468.84 ms average latency
• Speedup: 1.22x (18.2% improvement)

How to checkout & try? (for the reviewer)

from lerobot.policies.diffusion.configuration_diffusion import DiffusionConfig
from lerobot.configs.types import FeatureType, PolicyFeature

# Enable compilation
cfg = DiffusionConfig(
    input_features={
        "observation.state": PolicyFeature(type=FeatureType.STATE, shape=(14,)),
        "observation.image.camera": PolicyFeature(type=FeatureType.VISUAL, shape=(3, 96, 96)),
    },
    output_features={
        "action": PolicyFeature(type=FeatureType.ACTION, shape=(2,)),
    },
    compile_model=True,  # Enable torch.compile
    compile_mode="reduce-overhead",  # Optional: "default", "reduce-overhead", or "max-autotune"
)

policy = DiffusionPolicy(cfg)
# Use policy as normal - first inference will be slower due to compilation overhead
# Subsequent inferences benefit from the optimization

[imstevenpmwork]

LGTM, thanks for this contribution !