feat: add gradient accumulation support

[irisTa56]

What this does

Adds gradient accumulation support to the training script.
This allows simulating larger batch sizes without increasing GPU memory usage, which is useful when training large models or when memory is limited.

• Added gradient_accumulation_steps configuration parameter to TrainPipelineConfig (default: 1)
• Updated update_policy() to use accelerator.accumulate() context manager
• Updated effective batch size calculation in logging and MetricsTracker
• Training loop now skips evaluation and checkpointing during gradient accumulation steps (i.e., treats step as an optimizer step)

How it was tested

• Added tests in tests/training/test_update_policy.py.
  • test_update_policy_sync_gradients: Verifies gradient sync behavior
  • test_update_policy_gradient_accumulation_equivalence: Validates mathematical equivalence
• Added test_metrics_tracker_step_with_accelerator in tests/utils/test_logging_utils.py

P.S. I have removed those tests since we’re using the Accelerate API, and testing it is out of the scope of this PR. But they can be seen at d5df208 and d365eb8.

How to checkout & try? (for the reviewer)

# Effective batch size: 8 × 1 × 4 = 32
lerobot-train \
  --dataset.repo_id=lerobot/svla_so101_pickplace \
  --policy.type=act \
  --policy.repo_id=foo/my_policy \
  --policy.push_to_hub=false \
  --batch_size=8 \
  --gradient_accumulation_steps=4 \
  --log_freq=1 \
  --steps=50

[Copilot]

Pull request overview

This PR adds gradient accumulation support to the training pipeline, enabling simulation of larger batch sizes without increasing GPU memory usage. The implementation leverages Accelerator's built-in gradient accumulation features.

Key changes:

• Added gradient_accumulation_steps configuration parameter (default: 1) to control how many batches to accumulate before performing an optimizer step
• Updated training loop to properly handle gradient synchronization, skipping evaluation/checkpointing during accumulation steps
• Modified effective batch size calculations throughout the codebase to account for gradient accumulation

Reviewed changes

Copilot reviewed 5 out of 5 changed files in this pull request and generated 2 comments.

Show a summary per file

File	Description
src/lerobot/configs/train.py	Adds gradient_accumulation_steps configuration parameter with documentation
src/lerobot/scripts/lerobot_train.py	Updates training loop and update_policy() to use Accelerator's accumulation context, removes unused lock parameter, and adds gradient sync checks
src/lerobot/utils/logging_utils.py	Updates MetricsTracker.step() to calculate effective batch size including gradient accumulation
tests/training/test_update_policy.py	Adds comprehensive tests for gradient sync behavior and mathematical equivalence
tests/utils/test_logging_utils.py	Adds test for MetricsTracker with gradient accumulation

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

[jadechoghari]

Hey @irisTa56 that's a very useful feature! I did it locally to train bigger models. Could you maybe resolve the conflicts with main, so we kickoff a review?
Thanks!

[irisTa56]

@jadechoghari
Thanks for the feedback! I've resolved the merge conflicts with main and also added code comments for the newly introduced RA-BC feature. The PR should be ready for review now.

[jadechoghari]

Thanks! @irisTa56, the tests in this PR seem unnecessary since we’re using the accelerate api, and testing this functionality is the responsibility of accelerate itself. Removing the accelerate-related tests would make the PR cleaner and shorter 😄

[irisTa56]

Thanks! @irisTa56, the tests in this PR seem unnecessary since we’re using the accelerate api, and testing this functionality is the responsibility of accelerate itself. Removing the accelerate-related tests would make the PR cleaner and shorter 😄

Thanks for the feedback! I agree and have removed the tests.