From 6439faa39e1ceb3af761f006ff9f5dd58c52713c Mon Sep 17 00:00:00 2001
From: AdilZouitine <adilzouitinegm@gmail.com>
Date: Wed, 19 Mar 2025 13:16:31 +0000
Subject: [PATCH] Enhance training information logging in learner server

- Added tracking for replay buffer size and offline replay buffer size during training steps.
---
 lerobot/scripts/server/learner_server.py | 5 +++++
 1 file changed, 5 insertions(+)

diff --git a/lerobot/scripts/server/learner_server.py b/lerobot/scripts/server/learner_server.py
index 713fc2a8..2b19fea2 100644
--- a/lerobot/scripts/server/learner_server.py
+++ b/lerobot/scripts/server/learner_server.py
@@ -681,6 +681,11 @@ def add_actor_information_and_train(
 
         policy.update_target_networks()
         if optimization_step % cfg.training.log_freq == 0:
+            training_infos["replay_buffer_size"] = len(replay_buffer)
+            if offline_replay_buffer is not None:
+                training_infos["offline_replay_buffer_size"] = len(
+                    offline_replay_buffer
+                )
             training_infos["Optimization step"] = optimization_step
             logger.log_dict(
                 d=training_infos, mode="train", custom_step_key="Optimization step"