From 170c09e7f63660934d8e784df202d04bea2e77bc Mon Sep 17 00:00:00 2001
From: Steven Palma <imstevenpmwork@ieee.org>
Date: Wed, 24 Sep 2025 17:10:52 +0200
Subject: [PATCH] chore(utils): move queue utils and wandb_utils to their
 respective modules (#2030)

* chore(utils): move queue utils and wandb_utils to their respective modules

* fix(rl): remove double imports

---------

Signed-off-by: Steven Palma <imstevenpmwork@ieee.org>
---
 src/lerobot/rl/actor.py                  | 2 +-
 src/lerobot/rl/learner.py                | 4 ++--
 src/lerobot/rl/learner_service.py        | 2 +-
 src/lerobot/{utils => rl}/queue.py       | 0
 src/lerobot/{utils => rl}/wandb_utils.py | 0
 src/lerobot/scripts/train.py             | 2 +-
 tests/{utils => rl}/test_queue.py        | 2 +-
 7 files changed, 6 insertions(+), 6 deletions(-)
 rename src/lerobot/{utils => rl}/queue.py (100%)
 rename src/lerobot/{utils => rl}/wandb_utils.py (100%)
 rename tests/{utils => rl}/test_queue.py (98%)

diff --git a/src/lerobot/rl/actor.py b/src/lerobot/rl/actor.py
index 2606481d..3c025a05 100644
--- a/src/lerobot/rl/actor.py
+++ b/src/lerobot/rl/actor.py
@@ -64,6 +64,7 @@ from lerobot.policies.factory import make_policy
 from lerobot.policies.sac.modeling_sac import SACPolicy
 from lerobot.processor import TransitionKey
 from lerobot.rl.process import ProcessSignalHandler
+from lerobot.rl.queue import get_last_item_from_queue
 from lerobot.robots import so100_follower  # noqa: F401
 from lerobot.teleoperators import gamepad, so101_leader  # noqa: F401
 from lerobot.teleoperators.utils import TeleopEvents
@@ -76,7 +77,6 @@ from lerobot.transport.utils import (
     send_bytes_in_chunks,
     transitions_to_bytes,
 )
-from lerobot.utils.queue import get_last_item_from_queue
 from lerobot.utils.random_utils import set_seed
 from lerobot.utils.robot_utils import busy_wait
 from lerobot.utils.transition import (
diff --git a/src/lerobot/rl/learner.py b/src/lerobot/rl/learner.py
index 1ff34376..0faa460e 100644
--- a/src/lerobot/rl/learner.py
+++ b/src/lerobot/rl/learner.py
@@ -68,6 +68,7 @@ from lerobot.policies.factory import make_policy
 from lerobot.policies.sac.modeling_sac import SACPolicy
 from lerobot.rl.buffer import ReplayBuffer, concatenate_batch_transitions
 from lerobot.rl.process import ProcessSignalHandler
+from lerobot.rl.wandb_utils import WandBLogger
 from lerobot.robots import so100_follower  # noqa: F401
 from lerobot.teleoperators import gamepad, so101_leader  # noqa: F401
 from lerobot.teleoperators.utils import TeleopEvents
@@ -97,7 +98,6 @@ from lerobot.utils.utils import (
     get_safe_torch_device,
     init_logging,
 )
-from lerobot.utils.wandb_utils import WandBLogger
 
 from .learner_service import MAX_WORKERS, SHUTDOWN_TIMEOUT, LearnerService
 
@@ -153,7 +153,7 @@ def train(cfg: TrainRLServerPipelineConfig, job_name: str | None = None):
 
     # Setup WandB logging if enabled
     if cfg.wandb.enable and cfg.wandb.project:
-        from lerobot.utils.wandb_utils import WandBLogger
+        from lerobot.rl.wandb_utils import WandBLogger
 
         wandb_logger = WandBLogger(cfg)
     else:
diff --git a/src/lerobot/rl/learner_service.py b/src/lerobot/rl/learner_service.py
index b07c296e..7ef38119 100644
--- a/src/lerobot/rl/learner_service.py
+++ b/src/lerobot/rl/learner_service.py
@@ -19,9 +19,9 @@ import logging
 import time
 from multiprocessing import Event, Queue
 
+from lerobot.rl.queue import get_last_item_from_queue
 from lerobot.transport import services_pb2, services_pb2_grpc
 from lerobot.transport.utils import receive_bytes_in_chunks, send_bytes_in_chunks
-from lerobot.utils.queue import get_last_item_from_queue
 
 MAX_WORKERS = 3  # Stream parameters, send transitions and interactions
 SHUTDOWN_TIMEOUT = 10
diff --git a/src/lerobot/utils/queue.py b/src/lerobot/rl/queue.py
similarity index 100%
rename from src/lerobot/utils/queue.py
rename to src/lerobot/rl/queue.py
diff --git a/src/lerobot/utils/wandb_utils.py b/src/lerobot/rl/wandb_utils.py
similarity index 100%
rename from src/lerobot/utils/wandb_utils.py
rename to src/lerobot/rl/wandb_utils.py
diff --git a/src/lerobot/scripts/train.py b/src/lerobot/scripts/train.py
index 21da62bb..df33f1db 100644
--- a/src/lerobot/scripts/train.py
+++ b/src/lerobot/scripts/train.py
@@ -35,6 +35,7 @@ from lerobot.optim.factory import make_optimizer_and_scheduler
 from lerobot.policies.factory import make_policy, make_pre_post_processors
 from lerobot.policies.pretrained import PreTrainedPolicy
 from lerobot.policies.utils import get_device_from_parameters
+from lerobot.rl.wandb_utils import WandBLogger
 from lerobot.scripts.eval import eval_policy_all
 from lerobot.utils.logging_utils import AverageMeter, MetricsTracker
 from lerobot.utils.random_utils import set_seed
@@ -51,7 +52,6 @@ from lerobot.utils.utils import (
     has_method,
     init_logging,
 )
-from lerobot.utils.wandb_utils import WandBLogger
 
 
 def update_policy(
diff --git a/tests/utils/test_queue.py b/tests/rl/test_queue.py
similarity index 98%
rename from tests/utils/test_queue.py
rename to tests/rl/test_queue.py
index 6e42acdb..b6716fbd 100644
--- a/tests/utils/test_queue.py
+++ b/tests/rl/test_queue.py
@@ -20,7 +20,7 @@ from queue import Queue
 
 from torch.multiprocessing import Queue as TorchMPQueue
 
-from lerobot.utils.queue import get_last_item_from_queue
+from lerobot.rl.queue import get_last_item_from_queue
 
 
 def test_get_last_item_single_item():