feat(diffusion): add gradient checkpointing for memory optimization

[jashshah999]

Summary

Adds gradient checkpointing to DiffusionPolicy's UNet (encoder, mid, and decoder residual blocks). This trades compute for memory, allowing training with larger batch sizes or higher-resolution inputs on memory-constrained GPUs.

Currently only Pi0 and XVLA have gradient checkpointing. This extends it to DiffusionPolicy as called for in the 0.6.0 roadmap (item 3.3).

Usage

lerobot-train \
  --policy.type=diffusion \
  --policy.gradient_checkpointing=true \
  --dataset.repo_id=lerobot/pusht

What changed

• Added gradient_checkpointing: bool = False to DiffusionConfig
• Wrapped UNet encoder/mid/decoder residual blocks with torch.utils.checkpoint.checkpoint when enabled and training
• Uses use_reentrant=False for compatibility with torch.compile

Test plan

• Training without gradient_checkpointing works as before
• Training with gradient_checkpointing=true runs without error
• Peak GPU memory is reduced with gradient_checkpointing enabled
• Training loss converges similarly with and without checkpointing

0.6.0 roadmap item 3.3.