From ac83f4797c7a013b5a2c12590237e5cf61aa5a1a Mon Sep 17 00:00:00 2001
From: Khalil Meftah <khalil.meftah@huggingface.co>
Date: Thu, 7 May 2026 10:43:59 +0200
Subject: [PATCH] chore: address reviewer comments

---
 .../configuration_gaussian_actor.py           | 26 ++++++++++----
 .../gaussian_actor/modeling_gaussian_actor.py |  3 +-
 src/lerobot/rl/actor.py                       |  4 +--
 src/lerobot/rl/algorithms/base.py             |  8 ++---
 src/lerobot/rl/algorithms/configs.py          |  2 +-
 src/lerobot/rl/algorithms/factory.py          |  4 +--
 .../rl/algorithms/sac/configuration_sac.py    | 34 ++++++++++++++++---
 .../rl/algorithms/sac/sac_algorithm.py        | 31 ++++++++++++++---
 src/lerobot/rl/data_sources/__init__.py       |  4 ++-
 src/lerobot/rl/data_sources/data_mixer.py     |  5 +--
 src/lerobot/rl/eval_policy.py                 |  2 +-
 src/lerobot/rl/learner.py                     | 12 +++----
 src/lerobot/rl/train_rl.py                    |  7 ++--
 src/lerobot/rl/trainer.py                     |  8 +++--
 src/lerobot/types.py                          |  1 +
 15 files changed, 110 insertions(+), 41 deletions(-)

diff --git a/src/lerobot/policies/gaussian_actor/configuration_gaussian_actor.py b/src/lerobot/policies/gaussian_actor/configuration_gaussian_actor.py
index e6f9ebea1..e51653992 100644
--- a/src/lerobot/policies/gaussian_actor/configuration_gaussian_actor.py
+++ b/src/lerobot/policies/gaussian_actor/configuration_gaussian_actor.py
@@ -143,34 +143,48 @@ class GaussianActorConfig(PreTrainedConfig):
     latent_dim: int = 256
 
     # Online training (TODO(Khalil): relocate to TrainRLServerPipelineConfig)
+    # Number of steps for online training
     online_steps: int = 1000000
+    # Capacity of the online replay buffer
     online_buffer_capacity: int = 100000
+    # Capacity of the offline replay buffer
     offline_buffer_capacity: int = 100000
+    # Whether to use asynchronous prefetching for the buffers
     async_prefetch: bool = False
+    # Number of steps before learning starts
     online_step_before_learning: int = 100
 
     # Actor-learner transport (TODO(Khalil): relocate to TrainRLServerPipelineConfig).
+    # Configuration for actor-learner architecture
     actor_learner_config: ActorLearnerConfig = field(default_factory=ActorLearnerConfig)
+    # Configuration for concurrency settings (you can use threads or processes for the actor and learner)
     concurrency: ConcurrencyConfig = field(default_factory=ConcurrencyConfig)
 
     # Network architecture
-    # Actor network
+    # Configuration for the actor network architecture
     actor_network_kwargs: ActorNetworkConfig = field(default_factory=ActorNetworkConfig)
-    # Gaussian head parameters
+    # Configuration for the policy parameters (Gaussian head)
     policy_kwargs: PolicyConfig = field(default_factory=PolicyConfig)
-    # Discrete critic
+    # Configuration for the discrete critic network
     discrete_critic_network_kwargs: CriticNetworkConfig = field(default_factory=CriticNetworkConfig)
 
     def __post_init__(self):
         super().__post_init__()
+        # Any validation specific to GaussianActor configuration
 
     def get_optimizer_preset(self) -> MultiAdamConfig:
+        # Default learning rate used to satisfy the abstract ``get_optimizer_preset()``
+        # contract from ``PreTrainedConfig``. The actual optimizers used during RL
+        # training are built by ``SACAlgorithm.make_optimizers_and_scheduler()`` from
+        # ``SACAlgorithmConfig.{actor_lr,critic_lr,temperature_lr}`` and fully bypass
+        # this preset.
+        default_lr = 3e-4
         return MultiAdamConfig(
             weight_decay=0.0,
             optimizer_groups={
-                "actor": {"lr": 3e-4},
-                "critic": {"lr": 3e-4},
-                "temperature": {"lr": 3e-4},
+                "actor": {"lr": default_lr},
+                "critic": {"lr": default_lr},
+                "temperature": {"lr": default_lr},
             },
         )
 
diff --git a/src/lerobot/policies/gaussian_actor/modeling_gaussian_actor.py b/src/lerobot/policies/gaussian_actor/modeling_gaussian_actor.py
index 21828ba5d..9a7bcf1bc 100644
--- a/src/lerobot/policies/gaussian_actor/modeling_gaussian_actor.py
+++ b/src/lerobot/policies/gaussian_actor/modeling_gaussian_actor.py
@@ -25,6 +25,7 @@ from torch import Tensor
 from torch.distributions import MultivariateNormal, TanhTransform, Transform, TransformedDistribution
 
 from lerobot.utils.constants import ACTION, OBS_ENV_STATE, OBS_STATE
+from lerobot.utils.transition import move_state_dict_to_device
 
 from ..pretrained import PreTrainedPolicy
 from ..utils import get_device_from_parameters
@@ -113,8 +114,6 @@ class GaussianActorPolicy(
         return {"action": actions, "log_prob": log_probs, "action_mean": means}
 
     def load_actor_weights(self, state_dicts: dict[str, Any], device: str | torch.device = "cpu") -> None:
-        from lerobot.utils.transition import move_state_dict_to_device
-
         actor_state_dict = move_state_dict_to_device(state_dicts["policy"], device=device)
         self.actor.load_state_dict(actor_state_dict)
 
diff --git a/src/lerobot/rl/actor.py b/src/lerobot/rl/actor.py
index 2da2fba6f..ebb746343 100644
--- a/src/lerobot/rl/actor.py
+++ b/src/lerobot/rl/actor.py
@@ -62,8 +62,6 @@ from lerobot.cameras import opencv  # noqa: F401
 from lerobot.configs import parser
 from lerobot.policies import PreTrainedPolicy, make_policy, make_pre_post_processors
 from lerobot.processor import TransitionKey
-from lerobot.rl.queue import get_last_item_from_queue
-from lerobot.rl.train_rl import TrainRLServerPipelineConfig
 from lerobot.robots import so_follower  # noqa: F401
 from lerobot.teleoperators import gamepad, so_leader  # noqa: F401
 from lerobot.teleoperators.utils import TeleopEvents
@@ -95,6 +93,8 @@ from .gym_manipulator import (
     reset_and_build_transition,
     step_env_and_process_transition,
 )
+from .queue import get_last_item_from_queue
+from .train_rl import TrainRLServerPipelineConfig
 
 # Main entry point
 
diff --git a/src/lerobot/rl/algorithms/base.py b/src/lerobot/rl/algorithms/base.py
index b9f2c908c..69beff36b 100644
--- a/src/lerobot/rl/algorithms/base.py
+++ b/src/lerobot/rl/algorithms/base.py
@@ -21,12 +21,12 @@ from typing import TYPE_CHECKING, Any
 import torch
 from torch.optim import Optimizer
 
-from lerobot.rl.algorithms.configs import RLAlgorithmConfig, TrainingStats
+from lerobot.types import BatchType
+
+from .configs import RLAlgorithmConfig, TrainingStats
 
 if TYPE_CHECKING:
-    from lerobot.rl.data_sources.data_mixer import DataMixer
-
-BatchType = dict[str, Any]
+    from ..data_sources.data_mixer import DataMixer
 
 
 class RLAlgorithm(abc.ABC):
diff --git a/src/lerobot/rl/algorithms/configs.py b/src/lerobot/rl/algorithms/configs.py
index 58818f9db..c87042bc3 100644
--- a/src/lerobot/rl/algorithms/configs.py
+++ b/src/lerobot/rl/algorithms/configs.py
@@ -22,7 +22,7 @@ import draccus
 import torch
 
 if TYPE_CHECKING:
-    from lerobot.rl.algorithms.base import RLAlgorithm
+    from .base import RLAlgorithm
 
 
 @dataclass
diff --git a/src/lerobot/rl/algorithms/factory.py b/src/lerobot/rl/algorithms/factory.py
index ea89e598c..8adc9883d 100644
--- a/src/lerobot/rl/algorithms/factory.py
+++ b/src/lerobot/rl/algorithms/factory.py
@@ -16,8 +16,8 @@ from __future__ import annotations
 
 import torch
 
-from lerobot.rl.algorithms.base import RLAlgorithm
-from lerobot.rl.algorithms.configs import RLAlgorithmConfig
+from .base import RLAlgorithm
+from .configs import RLAlgorithmConfig
 
 
 def make_algorithm_config(algorithm_type: str, **kwargs) -> RLAlgorithmConfig:
diff --git a/src/lerobot/rl/algorithms/sac/configuration_sac.py b/src/lerobot/rl/algorithms/sac/configuration_sac.py
index 0ccb740f3..d126925ff 100644
--- a/src/lerobot/rl/algorithms/sac/configuration_sac.py
+++ b/src/lerobot/rl/algorithms/sac/configuration_sac.py
@@ -23,34 +23,57 @@ from lerobot.policies.gaussian_actor.configuration_gaussian_actor import (
     CriticNetworkConfig,
     GaussianActorConfig,
 )
-from lerobot.rl.algorithms.configs import RLAlgorithmConfig
+
+from ..configs import RLAlgorithmConfig
 
 if TYPE_CHECKING:
-    from lerobot.rl.algorithms.sac.sac_algorithm import SACAlgorithm
+    from .sac_algorithm import SACAlgorithm
 
 
 @RLAlgorithmConfig.register_subclass("sac")
 @dataclass
 class SACAlgorithmConfig(RLAlgorithmConfig):
-    """SAC algorithm hyperparameters."""
+    """Soft Actor-Critic (SAC) algorithm configuration.
+
+    SAC is an off-policy actor-critic deep RL algorithm based on the maximum
+    entropy reinforcement learning framework. It learns a policy and a Q-function
+    simultaneously using experience collected from the environment.
+
+    This configuration class contains the algorithm-side hyperparameters: critic
+    ensemble, target networks, temperature / entropy tuning, and the Bellman
+    update loop. The policy-side (actor + observation encoder) lives in
+    :class:`~lerobot.policies.gaussian_actor.GaussianActorConfig` and is
+    referenced via :attr:`policy_config`.
+    """
 
     # Optimizer learning rates
+    # Learning rate for the actor network
     actor_lr: float = 3e-4
+    # Learning rate for the critic network
     critic_lr: float = 3e-4
+    # Learning rate for the temperature parameter
     temperature_lr: float = 3e-4
 
     # Bellman update
+    # Discount factor for the SAC algorithm
     discount: float = 0.99
+    # Whether to use backup entropy for the SAC algorithm
     use_backup_entropy: bool = True
+    # Weight for the critic target update
     critic_target_update_weight: float = 0.005
 
     # Critic ensemble
+    # Number of critics in the ensemble
     num_critics: int = 2
+    # Number of subsampled critics for training
     num_subsample_critics: int | None = None
+    # Configuration for the critic network architecture
     critic_network_kwargs: CriticNetworkConfig = field(default_factory=CriticNetworkConfig)
+    # Configuration for the discrete critic network
     discrete_critic_network_kwargs: CriticNetworkConfig = field(default_factory=CriticNetworkConfig)
 
     # Temperature / entropy
+    # Initial temperature value
     temperature_init: float = 1.0
     # Target entropy for automatic temperature tuning. If ``None``, defaults to
     # ``-|A|/2`` where ``|A|`` is the total action dimension (continuous + 1 if
@@ -58,8 +81,11 @@ class SACAlgorithmConfig(RLAlgorithmConfig):
     target_entropy: float | None = None
 
     # Update loop
+    # Update-to-data ratio. Set to >1 to enable extra critic updates per env step.
     utd_ratio: int = 1
+    # Frequency of policy updates
     policy_update_freq: int = 1
+    # Gradient clipping norm for the SAC algorithm
     grad_clip_norm: float = 40.0
 
     # Optimizations
@@ -85,6 +111,6 @@ class SACAlgorithmConfig(RLAlgorithmConfig):
                 "before calling build_algorithm()."
             )
 
-        from lerobot.rl.algorithms.sac.sac_algorithm import SACAlgorithm
+        from .sac_algorithm import SACAlgorithm
 
         return SACAlgorithm(policy=policy, config=self)
diff --git a/src/lerobot/rl/algorithms/sac/sac_algorithm.py b/src/lerobot/rl/algorithms/sac/sac_algorithm.py
index 78cb40536..49640eb1a 100644
--- a/src/lerobot/rl/algorithms/sac/sac_algorithm.py
+++ b/src/lerobot/rl/algorithms/sac/sac_algorithm.py
@@ -35,12 +35,14 @@ from lerobot.policies.gaussian_actor.modeling_gaussian_actor import (
     orthogonal_init,
 )
 from lerobot.policies.utils import get_device_from_parameters
-from lerobot.rl.algorithms.base import BatchType, RLAlgorithm
-from lerobot.rl.algorithms.configs import TrainingStats
-from lerobot.rl.algorithms.sac.configuration_sac import SACAlgorithmConfig
+from lerobot.types import BatchType
 from lerobot.utils.constants import ACTION
 from lerobot.utils.transition import move_state_dict_to_device
 
+from ..base import RLAlgorithm
+from ..configs import TrainingStats
+from .configuration_sac import SACAlgorithmConfig
+
 
 class SACAlgorithm(RLAlgorithm):
     """Soft Actor-Critic. Owns critics, targets, temperature, and loss computation."""
@@ -175,6 +177,25 @@ class SACAlgorithm(RLAlgorithm):
         return q_values
 
     def update(self, batch_iterator: Iterator[BatchType]) -> TrainingStats:
+        """Run one SAC training step (critic / discrete-critic / actor / temperature).
+
+        Pulls ``utd_ratio`` batches from ``batch_iterator``, computes the relevant
+        losses, backpropagates each, and updates target networks.
+
+        Args:
+            batch_iterator: yields batches each containing
+                - ``action``: Action tensor
+                - ``reward``: Reward tensor
+                - ``state``: Observations tensor dict
+                - ``next_state``: Next observations tensor dict
+                - ``done``: Done mask tensor
+                - ``observation_feature``: Optional pre-computed observation features
+                - ``next_observation_feature``: Optional pre-computed next observation features
+                - ``complementary_info`` (optional): per-step extras like discrete penalties
+
+        Returns:
+            TrainingStats with per-component losses and grad norms.
+        """
         clip = self.config.grad_clip_norm
 
         for _ in range(self.config.utd_ratio - 1):
@@ -248,12 +269,14 @@ class SACAlgorithm(RLAlgorithm):
         return stats
 
     def _compute_loss_critic(self, batch: dict[str, Any]) -> Tensor:
+        # Extract common components from batch
         observations = batch["state"]
         actions = batch[ACTION]
+        observation_features = batch.get("observation_feature")
+        # Extract critic-specific components
         rewards = batch["reward"]
         next_observations = batch["next_state"]
         done = batch["done"]
-        observation_features = batch.get("observation_feature")
         next_observation_features = batch.get("next_observation_feature")
 
         with torch.no_grad():
diff --git a/src/lerobot/rl/data_sources/__init__.py b/src/lerobot/rl/data_sources/__init__.py
index b4c0bcf3d..97cfe5001 100644
--- a/src/lerobot/rl/data_sources/__init__.py
+++ b/src/lerobot/rl/data_sources/__init__.py
@@ -12,6 +12,8 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from .data_mixer import BatchType, DataMixer, OnlineOfflineMixer
+from lerobot.types import BatchType
+
+from .data_mixer import DataMixer, OnlineOfflineMixer
 
 __all__ = ["BatchType", "DataMixer", "OnlineOfflineMixer"]
diff --git a/src/lerobot/rl/data_sources/data_mixer.py b/src/lerobot/rl/data_sources/data_mixer.py
index 7546bde0c..69929b10a 100644
--- a/src/lerobot/rl/data_sources/data_mixer.py
+++ b/src/lerobot/rl/data_sources/data_mixer.py
@@ -16,8 +16,9 @@ from __future__ import annotations
 
 import abc
 
-from lerobot.rl.algorithms.base import BatchType
-from lerobot.rl.buffer import ReplayBuffer, concatenate_batch_transitions
+from lerobot.types import BatchType
+
+from ..buffer import ReplayBuffer, concatenate_batch_transitions
 
 
 class DataMixer(abc.ABC):
diff --git a/src/lerobot/rl/eval_policy.py b/src/lerobot/rl/eval_policy.py
index b7eb25e95..0f42d7573 100644
--- a/src/lerobot/rl/eval_policy.py
+++ b/src/lerobot/rl/eval_policy.py
@@ -19,7 +19,6 @@ from lerobot.cameras import opencv  # noqa: F401
 from lerobot.configs import parser
 from lerobot.datasets import LeRobotDataset
 from lerobot.policies import make_policy
-from lerobot.rl.train_rl import TrainRLServerPipelineConfig
 from lerobot.robots import (  # noqa: F401
     RobotConfig,
     make_robot_from_config,
@@ -31,6 +30,7 @@ from lerobot.teleoperators import (
 )
 
 from .gym_manipulator import make_robot_env
+from .train_rl import TrainRLServerPipelineConfig
 
 logging.basicConfig(level=logging.INFO)
 
diff --git a/src/lerobot/rl/learner.py b/src/lerobot/rl/learner.py
index 7f0632237..7cfee9cba 100644
--- a/src/lerobot/rl/learner.py
+++ b/src/lerobot/rl/learner.py
@@ -71,12 +71,6 @@ from lerobot.common.wandb_utils import WandBLogger
 from lerobot.configs import parser
 from lerobot.datasets import LeRobotDataset, make_dataset
 from lerobot.policies import make_policy, make_pre_post_processors
-from lerobot.rl.algorithms.base import RLAlgorithm
-from lerobot.rl.algorithms.factory import make_algorithm
-from lerobot.rl.buffer import ReplayBuffer
-from lerobot.rl.data_sources import OnlineOfflineMixer
-from lerobot.rl.train_rl import TrainRLServerPipelineConfig
-from lerobot.rl.trainer import RLTrainer
 from lerobot.robots import so_follower  # noqa: F401
 from lerobot.teleoperators import gamepad, so_leader  # noqa: F401
 from lerobot.teleoperators.utils import TeleopEvents
@@ -102,7 +96,13 @@ from lerobot.utils.utils import (
     init_logging,
 )
 
+from .algorithms.base import RLAlgorithm
+from .algorithms.factory import make_algorithm
+from .buffer import ReplayBuffer
+from .data_sources import OnlineOfflineMixer
 from .learner_service import MAX_WORKERS, SHUTDOWN_TIMEOUT, LearnerService
+from .train_rl import TrainRLServerPipelineConfig
+from .trainer import RLTrainer
 
 
 @parser.wrap()
diff --git a/src/lerobot/rl/train_rl.py b/src/lerobot/rl/train_rl.py
index fe8078858..e5ae0f9f5 100644
--- a/src/lerobot/rl/train_rl.py
+++ b/src/lerobot/rl/train_rl.py
@@ -20,9 +20,10 @@ from dataclasses import dataclass
 
 from lerobot.configs.default import DatasetConfig
 from lerobot.configs.train import TrainPipelineConfig
-from lerobot.rl.algorithms.configs import RLAlgorithmConfig
-from lerobot.rl.algorithms.factory import make_algorithm_config
-from lerobot.rl.algorithms.sac import SACAlgorithmConfig  # noqa: F401
+
+from .algorithms.configs import RLAlgorithmConfig
+from .algorithms.factory import make_algorithm_config
+from .algorithms.sac import SACAlgorithmConfig  # noqa: F401
 
 
 @dataclass(kw_only=True)
diff --git a/src/lerobot/rl/trainer.py b/src/lerobot/rl/trainer.py
index 9332ccd15..65f00568e 100644
--- a/src/lerobot/rl/trainer.py
+++ b/src/lerobot/rl/trainer.py
@@ -17,9 +17,11 @@ from __future__ import annotations
 from collections.abc import Iterator
 from typing import Any
 
-from lerobot.rl.algorithms.base import BatchType, RLAlgorithm
-from lerobot.rl.algorithms.configs import TrainingStats
-from lerobot.rl.data_sources.data_mixer import DataMixer
+from lerobot.types import BatchType
+
+from .algorithms.base import RLAlgorithm
+from .algorithms.configs import TrainingStats
+from .data_sources.data_mixer import DataMixer
 
 
 class RLTrainer:
diff --git a/src/lerobot/types.py b/src/lerobot/types.py
index d9b8166c5..9de504870 100644
--- a/src/lerobot/types.py
+++ b/src/lerobot/types.py
@@ -40,6 +40,7 @@ PolicyAction = torch.Tensor
 RobotAction = dict[str, Any]
 EnvAction = np.ndarray
 RobotObservation = dict[str, Any]
+BatchType = dict[str, Any]
 
 
 EnvTransition = TypedDict(