Compare commits

...

1 Commits

Author SHA1 Message Date
CarolinePascal 1834f819a6 docs(rl): write the API reference docstrings
Second module of Wave 4 (training & eval): brings src/lerobot/rl/ to 100%
public docstring coverage, following the standard in
docs/source/writing_docstrings.mdx.

- Documents the remaining gaps across the SAC algorithm (SACAlgorithm.__init__,
  CriticHead, CriticEnsemble.__init__/forward, get_optimizers), the
  RLAlgorithm/RLAlgorithmConfig base contract (optimization_step setter,
  from_pretrained), SACAlgorithmConfig (converts inline `#` field comments to
  a proper Args: block), ReplayBuffer/BatchTransition, OnlineOfflineMixer,
  TrainRLServerPipelineConfig (documents every inherited TrainPipelineConfig
  field, since the base class itself is undocumented and out of scope), and
  the actor/learner gRPC entry points (actor_cli, train_cli,
  transitions_stream/interactions_stream, LearnerService's 5 servicer
  methods) and their smaller helpers (queue.get_last_item_from_queue,
  crop_dataset_roi.mouse_callback, eval_policy).
- Also documents 3 dunder methods (RLTrainer's _PreprocessedIterator.__iter__/
  __next__, ReplayBuffer.__len__) that a naive "skip all underscore-prefixed
  names" gap scan misses but interrogate's ignore-magic=false requires.
- Removing the D-ignore surfaced ~30 pre-existing docstrings with D205/D415/
  D417 issues (missing blank line after summary, missing punctuation, stale
  Args entries that didn't match the real signature) across actor.py,
  crop_dataset_roi.py, gym_manipulator.py, learner.py, and
  learner_service.py — all fixed as part of this PR.
- Removes "src/lerobot/rl/**" = ["D"] from pyproject.toml's ruff ignore list;
  the whole module is now checked (no per-family split to narrow, unlike
  policies).
- Adds lerobot.rl to check_docstrings.py's MODULES_TO_CHECK ratchet.
- Creates docs/source/api/rl.mdx from scratch (algorithm base contract, SAC,
  replay buffer, data mixers, trainer, actor/learner CLIs and gRPC service)
  and wires it into _toctree.yml, cross-linked from the existing hilserl.mdx/
  hilserl_sim.mdx guides. Verified via a full doc-builder build — no dead
  cross-references, no leftover placeholder text.
- Ratchets interrogate's fail-under from 55 to 55.5 (measured 55.9% with this
  PR).

Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
2026-08-07 15:55:37 +02:00
21 changed files with 665 additions and 225 deletions
+2
View File
@@ -215,4 +215,6 @@
title: Environments
- local: api/configs
title: Configuration
- local: api/rl
title: Reinforcement Learning
title: "API Reference"
+87
View File
@@ -0,0 +1,87 @@
# Reinforcement Learning
`lerobot.rl` is the distributed actor/learner reinforcement-learning stack behind
[Train a Robot with RL](../hilserl) (HIL-SERL) and [Train RL in Simulation](../hilserl_sim). Algorithms,
the replay buffer, data sources, and the trainer are gRPC-free and usable standalone; the actor/learner
entry points (`actor`, `learner`, `learner_service`) additionally require `pip install 'lerobot[hilserl]'`.
## TrainRLServerPipelineConfig
Top-level configuration for both the `lerobot-actor` and `lerobot-learner` CLIs.
[[autodoc]] lerobot.rl.train_rl.TrainRLServerPipelineConfig
## RLAlgorithm
Abstract base every RL algorithm subclasses.
[[autodoc]] lerobot.rl.algorithms.base.RLAlgorithm
## RLAlgorithmConfig
[[autodoc]] lerobot.rl.algorithms.configs.RLAlgorithmConfig
## TrainingStats
[[autodoc]] lerobot.rl.algorithms.configs.TrainingStats
## make_algorithm
[[autodoc]] lerobot.rl.algorithms.factory.make_algorithm
## make_algorithm_config
[[autodoc]] lerobot.rl.algorithms.factory.make_algorithm_config
## get_algorithm_class
[[autodoc]] lerobot.rl.algorithms.factory.get_algorithm_class
## SAC
[[autodoc]] lerobot.rl.algorithms.sac.sac_algorithm.SACAlgorithm
- all
[[autodoc]] lerobot.rl.algorithms.sac.configuration_sac.SACAlgorithmConfig
## ReplayBuffer
In-memory replay buffer of transitions, sampled in batches for off-policy training.
[[autodoc]] lerobot.rl.buffer.ReplayBuffer
- all
[[autodoc]] lerobot.rl.buffer.BatchTransition
## DataMixer
Abstract interface for combining online and offline data sources into training batches.
[[autodoc]] lerobot.rl.data_sources.DataMixer
- all
[[autodoc]] lerobot.rl.data_sources.OnlineOfflineMixer
- all
## RLTrainer
Unified training-step orchestrator: holds the algorithm, a `DataMixer`, and an optional preprocessor.
[[autodoc]] lerobot.rl.trainer.RLTrainer
- all
## Actor / learner CLIs
The distributed actor and learner processes communicate over gRPC; see [Train a Robot with
RL](../hilserl) for the full workflow.
[[autodoc]] lerobot.rl.actor.actor_cli
[[autodoc]] lerobot.rl.learner.train_cli
[[autodoc]] lerobot.rl.learner_service.LearnerService
- all
## eval_policy
[[autodoc]] lerobot.rl.eval_policy.eval_policy
+1 -2
View File
@@ -451,7 +451,6 @@ ignore = [
"src/lerobot/policies/**" = ["D"]
"src/lerobot/processor/**" = ["D"]
"src/lerobot/rewards/**" = ["D"]
"src/lerobot/rl/**" = ["D"]
"src/lerobot/rollout/**" = ["D"]
"src/lerobot/scripts/**" = ["D"]
"src/lerobot/teleoperators/**" = ["D"]
@@ -515,7 +514,7 @@ ignore-private = false
ignore-property-decorators = false
ignore-module = false
ignore-setters = false
fail-under = 55
fail-under = 55.5
output-format = "term-missing"
color = true
paths = ["src/lerobot"]
+59 -23
View File
@@ -13,8 +13,7 @@
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""
Actor server runner for distributed HILSerl robot policy training.
"""Actor server runner for distributed HILSerl robot policy training.
This script implements the actor component of the distributed HILSerl architecture.
It executes the policy in the robot environment, collects experience,
@@ -119,6 +118,16 @@ from .train_rl import TrainRLServerPipelineConfig
@parser.wrap()
def actor_cli(cfg: TrainRLServerPipelineConfig):
"""CLI entry point for the HILSerl actor server.
Connects to the learner server over gRPC, then launches (as threads or processes, depending on
`cfg.policy.concurrency.multiprocessing_context`) the background workers that receive updated
policy parameters and stream transitions/interactions back to the learner, while running the
policy-environment interaction loop (`act_with_policy`) on the main thread/process.
Args:
cfg (`TrainRLServerPipelineConfig`): Parsed from the CLI.
"""
# Fail fast with a friendly error if the optional ``hilserl`` extra is missing.
require_package("grpcio", extra="hilserl", import_name="grpc")
cfg.validate()
@@ -234,18 +243,19 @@ def act_with_policy(
transitions_queue: Queue,
interactions_queue: Queue,
):
"""
Executes policy interaction within the environment.
"""Executes policy interaction within the environment.
This function rolls out the policy in the environment, collecting interaction data and pushing it to a queue for streaming to the learner.
Once an episode is completed, updated network parameters received from the learner are retrieved from a queue and loaded into the network.
Args:
cfg: Configuration settings for the interaction process.
shutdown_event: Event to check if the process should shutdown.
parameters_queue: Queue to receive updated network parameters from the learner.
transitions_queue: Queue to send transitions to the learner.
interactions_queue: Queue to send interactions to the learner.
cfg (`TrainRLServerPipelineConfig`): Training configuration.
shutdown_event (`Event`): Set to stop the policy loop.
parameters_queue (`Queue`): Queue of serialized learner weights, drained via
`update_policy_parameters`.
transitions_queue (`Queue`): Queue transitions are pushed to for streaming to the learner.
interactions_queue (`Queue`): Queue interaction messages are pushed to for streaming to the
learner.
"""
# Initialize logging for multiprocessing
if not use_threads(cfg):
@@ -440,7 +450,8 @@ def establish_learner_connection(
Args:
stub (services_pb2_grpc.LearnerServiceStub): The stub to use for the connection.
shutdown_event (Event): The event to check if the connection should be established.
attempts (int): The number of attempts to establish the connection.
attempts (int, *optional*, defaults to 30): The number of attempts to establish the connection.
Returns:
bool: True if the connection is established, False otherwise.
"""
@@ -473,7 +484,6 @@ def learner_service_client(
Returns:
tuple[services_pb2_grpc.LearnerServiceStub, grpc.Channel]: The stub and the channel.
"""
channel = grpc.insecure_channel(
f"{host}:{port}",
grpc_channel_options(),
@@ -496,8 +506,8 @@ def receive_policy(
cfg (TrainRLServerPipelineConfig): The configuration for the actor.
parameters_queue (Queue): The queue to receive the parameters.
shutdown_event (Event): The event to check if the process should shutdown.
learner_client (services_pb2_grpc.LearnerServiceStub | None): Optional pre-created stub.
grpc_channel (grpc.Channel | None): Optional pre-created channel.
learner_client (services_pb2_grpc.LearnerServiceStub | None, *optional*): Optional pre-created stub.
grpc_channel (grpc.Channel | None, *optional*): Optional pre-created channel.
"""
logging.info("[ACTOR] Start receiving parameters from the Learner")
if not use_threads(cfg):
@@ -557,10 +567,9 @@ def send_transitions(
cfg (TrainRLServerPipelineConfig): The configuration for the actor.
transitions_queue (Queue): The queue to receive the transitions.
shutdown_event (Event): The event to check if the process should shutdown.
learner_client (services_pb2_grpc.LearnerServiceStub | None): Optional pre-created stub.
grpc_channel (grpc.Channel | None): Optional pre-created channel.
learner_client (services_pb2_grpc.LearnerServiceStub | None, *optional*): Optional pre-created stub.
grpc_channel (grpc.Channel | None, *optional*): Optional pre-created channel.
"""
if not use_threads(cfg):
# Create a process-specific log file
log_dir = os.path.join(cfg.output_dir, "logs")
@@ -612,10 +621,9 @@ def send_interactions(
cfg (TrainRLServerPipelineConfig): The configuration for the actor.
interactions_queue (Queue): The queue to receive the interactions.
shutdown_event (Event): The event to check if the process should shutdown.
learner_client (services_pb2_grpc.LearnerServiceStub | None): Optional pre-created stub.
grpc_channel (grpc.Channel | None): Optional pre-created channel.
learner_client (services_pb2_grpc.LearnerServiceStub | None, *optional*): Optional pre-created stub.
grpc_channel (grpc.Channel | None, *optional*): Optional pre-created channel.
"""
if not use_threads(cfg):
# Create a process-specific log file
log_dir = os.path.join(cfg.output_dir, "logs")
@@ -657,6 +665,17 @@ def transitions_stream(
transitions_queue: Queue,
timeout: float,
) -> "Generator[Any, None, services_pb2.Empty]":
"""GRPC client-streaming generator that forwards queued transitions to the learner.
Args:
shutdown_event (`Event`): Set to stop streaming and return.
transitions_queue (`Queue`): Queue of serialized transition batches, filled by
`push_transitions_to_transport_queue`.
timeout (`float`): Seconds to wait for a queue item before checking `shutdown_event` again.
Yields:
Chunks of a `services_pb2.Transition` message, produced by `send_bytes_in_chunks`.
"""
while not shutdown_event.is_set():
try:
message = transitions_queue.get(block=True, timeout=timeout)
@@ -676,6 +695,16 @@ def interactions_stream(
interactions_queue: Queue,
timeout: float,
) -> "Generator[Any, None, services_pb2.Empty]":
"""GRPC client-streaming generator that forwards queued interaction messages to the learner.
Args:
shutdown_event (`Event`): Set to stop streaming and return.
interactions_queue (`Queue`): Queue of serialized interaction messages.
timeout (`float`): Seconds to wait for a queue item before checking `shutdown_event` again.
Yields:
Chunks of a `services_pb2.InteractionMessage`, produced by `send_bytes_in_chunks`.
"""
while not shutdown_event.is_set():
try:
message = interactions_queue.get(block=True, timeout=timeout)
@@ -718,12 +747,11 @@ def update_policy_parameters(algorithm: RLAlgorithm, parameters_queue: Queue, de
def push_transitions_to_transport_queue(transitions: list, transitions_queue):
"""Send transitions to learner in smaller chunks to avoid network issues.
"""Move `transitions` to CPU, check for NaNs, and enqueue them for the learner.
Args:
transitions: List of transitions to send
message_queue: Queue to send messages to learner
chunk_size: Size of each chunk to send
transitions (`list`): Transitions to send, as produced by the actor's rollout loop.
transitions_queue (`Queue`): Queue drained by `transitions_stream`.
"""
transition_to_send_to_learner = []
for transition in transitions:
@@ -760,6 +788,13 @@ def get_frequency_stats(timer: TimerManager) -> dict[str, float]:
def log_policy_frequency_issue(policy_fps: float, cfg: TrainRLServerPipelineConfig, interaction_step: int):
"""Log a warning if `policy_fps` is below the environment's target `cfg.env.fps`.
Args:
policy_fps (`float`): Measured policy loop frequency.
cfg (`TrainRLServerPipelineConfig`): Provides the target `cfg.env.fps` to compare against.
interaction_step (`int`): Current interaction step, included in the warning message.
"""
if policy_fps < cfg.env.fps:
logging.warning(
f"[ACTOR] Policy FPS {policy_fps:.1f} below required {cfg.env.fps} at step {interaction_step}"
@@ -767,6 +802,7 @@ def log_policy_frequency_issue(policy_fps: float, cfg: TrainRLServerPipelineConf
def use_threads(cfg: TrainRLServerPipelineConfig) -> bool:
"""Whether the actor's background workers should run as threads instead of processes."""
return cfg.policy.concurrency.actor == "threads"
+1
View File
@@ -101,6 +101,7 @@ class RLAlgorithm(HubMixin, abc.ABC):
@optimization_step.setter
def optimization_step(self, value: int) -> None:
"""Set the current learner optimization step."""
self._optimization_step = int(value)
def get_weights(self) -> dict[str, Any]:
+29 -1
View File
@@ -45,7 +45,6 @@ class TrainingStats:
def to_log_dict(self) -> dict[str, float]:
"""Flatten all stats into a single dict for logging."""
d: dict[str, float] = {}
for name, val in self.losses.items():
d[name] = val
@@ -98,6 +97,35 @@ class RLAlgorithmConfig(draccus.ChoiceRegistry, HubMixin, abc.ABC):
revision: str | None = None,
**algo_kwargs: Any,
) -> T:
"""Load an algorithm config from a local directory or the Hugging Face Hub.
Args:
pretrained_name_or_path (`str | Path`):
Local directory containing `config.json`, or a Hub repo id.
force_download (`bool`, *optional*, defaults to `False`):
Whether to force re-download the config even if it's cached.
resume_download (`bool | None`, *optional*):
Whether to resume an interrupted download.
proxies (`dict[Any, Any] | None`, *optional*):
Proxies to use for the download request.
token (`str | bool | None`, *optional*):
Hugging Face Hub authentication token.
cache_dir (`str | Path | None`, *optional*):
Directory to cache the downloaded config in.
local_files_only (`bool`, *optional*, defaults to `False`):
Whether to only look for files locally, without querying the Hub.
revision (`str | None`, *optional*):
Hub revision (branch, tag, or commit hash) to load from.
**algo_kwargs: Attribute overrides applied to the loaded config instance.
Returns:
RLAlgorithmConfig: The loaded config, as the concrete registered subclass.
Raises:
FileNotFoundError: If no `config.json` is found locally or on the Hub.
TypeError: If loaded via a specific subclass but the config's registered type doesn't
match it.
"""
model_id = str(pretrained_name_or_path)
config_file: str | None = None
if Path(model_id).is_dir():
+7 -9
View File
@@ -24,8 +24,8 @@ def make_algorithm_config(algorithm_type: str, **kwargs) -> RLAlgorithmConfig:
"""Instantiate an `RLAlgorithmConfig` from its registered type name.
Args:
algorithm_type: Registry key of the algorithm (e.g. ``"sac"``).
**kwargs: Keyword arguments forwarded to the config class constructor.
algorithm_type (`str`): Registry key of the algorithm (e.g. `"sac"`).
kwargs (`Any`, *optional*): Keyword arguments forwarded to the config class constructor.
Returns:
An instance of the matching ``RLAlgorithmConfig`` subclass.
@@ -44,14 +44,13 @@ def make_algorithm_config(algorithm_type: str, **kwargs) -> RLAlgorithmConfig:
def get_algorithm_class(name: str) -> type[RLAlgorithm]:
"""
Retrieves an RL algorithm class by its registered name.
"""Retrieves an RL algorithm class by its registered name.
This function uses dynamic imports to avoid loading all algorithm classes into
memory at once, improving startup time and reducing dependencies.
Args:
name: The name of the algorithm. Supported names are "sac".
name (`str`): The name of the algorithm. Supported names are "sac".
Returns:
The algorithm class corresponding to the given name.
@@ -70,8 +69,7 @@ def get_algorithm_class(name: str) -> type[RLAlgorithm]:
def make_algorithm(cfg: RLAlgorithmConfig, policy: torch.nn.Module) -> RLAlgorithm:
"""
Instantiate an RL algorithm.
"""Instantiate an RL algorithm.
This factory function looks up the :class:`RLAlgorithm` subclass that matches
``cfg.type`` and instantiates it with the provided policy. It also enforces
@@ -79,8 +77,8 @@ def make_algorithm(cfg: RLAlgorithmConfig, policy: torch.nn.Module) -> RLAlgorit
normally handled by :meth:`TrainRLServerPipelineConfig.validate`).
Args:
cfg: The algorithm configuration. Must have ``policy_config`` set.
policy: The policy module the algorithm will train.
cfg (`RLAlgorithmConfig`): The algorithm configuration. Must have `policy_config` set.
policy (`torch.nn.Module`): The policy module the algorithm will train.
Returns:
An instantiated :class:`RLAlgorithm`.
@@ -39,52 +39,73 @@ class SACAlgorithmConfig(RLAlgorithmConfig):
update loop. The policy-side (actor + observation encoder) lives in
:class:`~lerobot.policies.gaussian_actor.GaussianActorConfig` and is
referenced via :attr:`policy_config`.
Args:
actor_lr (`float`, *optional*, defaults to 0.0003):
Learning rate for the actor network.
critic_lr (`float`, *optional*, defaults to 0.0003):
Learning rate for the critic network.
temperature_lr (`float`, *optional*, defaults to 0.0003):
Learning rate for the temperature parameter.
discount (`float`, *optional*, defaults to 0.99):
Discount factor for the Bellman update.
use_backup_entropy (`bool`, *optional*, defaults to `True`):
Whether to use backup entropy in the Bellman target.
critic_target_update_weight (`float`, *optional*, defaults to 0.005):
Polyak-averaging weight for the critic target update.
num_critics (`int`, *optional*, defaults to 2):
Number of critics in the ensemble.
num_subsample_critics (`int | None`, *optional*):
Number of critics to subsample from the ensemble for each Bellman target computation.
`None` uses the full ensemble.
critic_network_kwargs (`CriticNetworkConfig`, *optional*):
Configuration for the (continuous-action) critic network architecture.
discrete_critic_network_kwargs (`CriticNetworkConfig`, *optional*):
Configuration for the discrete-action critic network architecture.
temperature_init (`float`, *optional*, defaults to 1.0):
Initial value of the entropy temperature.
target_entropy (`float | None`, *optional*):
Target entropy for automatic temperature tuning. If `None`, defaults to `-|A|/2` where
`|A|` is the total action dimension (continuous + 1 if there is a discrete action head).
utd_ratio (`int`, *optional*, defaults to 1):
Update-to-data ratio. Set to `>1` to enable extra critic updates per env step.
policy_update_freq (`int`, *optional*, defaults to 1):
Frequency of policy updates, in units of critic updates.
grad_clip_norm (`float`, *optional*, defaults to 40.0):
Gradient-clipping norm applied during optimization.
use_torch_compile (`bool`, *optional*, defaults to `False`):
Whether to `torch.compile` the algorithm's forward passes. Currently disabled by default.
policy_config (`PreTrainedConfig | None`, *optional*):
The policy (actor) config this algorithm trains. Populated via `from_policy_config` or by
`TrainRLServerPipelineConfig.validate` before the algorithm is constructed.
"""
# 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
# there is a discrete action head).
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
# torch.compile is currently disabled by default
use_torch_compile: bool = False
# Policy config
+65 -17
View File
@@ -55,6 +55,15 @@ class SACAlgorithm(RLAlgorithm):
policy: GaussianActorPolicy,
config: SACAlgorithmConfig,
):
"""Build the critic ensemble, target networks, and temperature from `config`.
Args:
policy (`GaussianActorPolicy`):
The actor policy this algorithm trains. Its observation encoder is shared with the
critics.
config (`SACAlgorithmConfig`):
Algorithm configuration.
"""
self.config = config
self.policy_config = config.policy_config
self.policy = policy
@@ -144,17 +153,18 @@ class SACAlgorithm(RLAlgorithm):
use_target: bool = False,
observation_features: Tensor | None = None,
) -> Tensor:
"""Forward pass through a critic network ensemble
"""Forward pass through a critic network ensemble.
Args:
observations: Dictionary of observations
actions: Action tensor
use_target: If True, use target critics, otherwise use ensemble critics
observation_features: Optional pre-computed observation features to avoid recomputing
encoder output
Returns:
Tensor of Q-values from all critics
"""
critics = self.critic_target if use_target else self.critic_ensemble
q_values = critics(observations, actions, observation_features)
return q_values
@@ -162,7 +172,7 @@ class SACAlgorithm(RLAlgorithm):
def _discrete_critic_forward(
self, observations, use_target=False, observation_features=None
) -> torch.Tensor:
"""Forward pass through a discrete critic network
"""Forward pass through a discrete critic network.
Args:
observations: Dictionary of observations
@@ -408,7 +418,7 @@ class SACAlgorithm(RLAlgorithm):
return actor_loss
def _compute_loss_temperature(self, batch: dict[str, Any]) -> Tensor:
"""Compute the temperature loss"""
"""Compute the temperature loss."""
observations = batch["state"]
observation_features = batch.get("observation_feature")
@@ -420,7 +430,7 @@ class SACAlgorithm(RLAlgorithm):
return temperature_loss
def _update_target_networks(self) -> None:
"""Update target networks with exponential moving average"""
"""Update target networks with exponential moving average."""
for target_p, p in zip(
self.critic_target.parameters(), self.critic_ensemble.parameters(), strict=True
):
@@ -461,8 +471,7 @@ class SACAlgorithm(RLAlgorithm):
return forward_batch
def make_optimizers_and_scheduler(self) -> dict[str, Optimizer]:
"""
Creates and returns optimizers for the actor, critic, and temperature components of a reinforcement learning policy.
"""Creates and returns optimizers for the actor, critic, and temperature components of a reinforcement learning policy.
This function sets up Adam optimizers for:
- The **actor network**, ensuring that only relevant parameters are optimized.
@@ -471,7 +480,7 @@ class SACAlgorithm(RLAlgorithm):
It also initializes a learning rate scheduler, though currently, it is set to `None`.
NOTE:
Note:
- If the encoder is shared, its parameters are excluded from the actor's optimization process.
- The policy's log temperature (`log_alpha`) is wrapped in a list to ensure proper optimization as a standalone tensor.
@@ -496,6 +505,7 @@ class SACAlgorithm(RLAlgorithm):
return self.optimizers
def get_optimizers(self) -> dict[str, Optimizer]:
"""See [`~rl.algorithms.RLAlgorithm.get_optimizers`]."""
return self.optimizers
def get_weights(self) -> dict[str, Any]:
@@ -560,20 +570,18 @@ class SACAlgorithm(RLAlgorithm):
def get_observation_features(
self, observations: Tensor, next_observations: Tensor
) -> tuple[Tensor | None, Tensor | None]:
"""
Get observation features from the policy encoder. It act as cache for the observation features.
when the encoder is frozen, the observation features are not updated.
We can save compute by caching the observation features.
"""Get observation features from the policy encoder, acting as a cache.
When the encoder is frozen, the observation features are not updated, so we can save compute
by caching them here instead of recomputing on every critic/actor forward pass.
Args:
policy: The policy model
observations: The current observations
next_observations: The next observations
Returns:
tuple: observation_features, next_observation_features
"""
if self.policy.config.vision_encoder_name is None or not self.policy.config.freeze_vision_encoder:
return None, None
@@ -595,6 +603,8 @@ def _split_prefix(state: dict[str, torch.Tensor], prefix: str) -> dict[str, torc
class CriticHead(nn.Module):
"""A single Q-value head: an MLP followed by a scalar linear output layer."""
def __init__(
self,
input_dim: int,
@@ -605,6 +615,23 @@ class CriticHead(nn.Module):
init_final: float | None = None,
final_activation: Callable[[torch.Tensor], torch.Tensor] | str | None = None,
):
"""Build the MLP trunk and scalar output layer.
Args:
input_dim (`int`): Dimension of the concatenated observation-encoding + action input.
hidden_dims (`list[int]`): Hidden layer widths of the MLP trunk.
activations (`Callable[[torch.Tensor], torch.Tensor] | str`, *optional*, defaults to `nn.SiLU()`):
Activation used between hidden layers.
activate_final (`bool`, *optional*, defaults to `False`): Whether to apply `activations`
after the last hidden layer.
dropout_rate (`float | None`, *optional*): Dropout probability applied between hidden
layers. `None` disables dropout.
init_final (`float | None`, *optional*): When set, the output layer's weight and bias are
initialized uniformly in `[-init_final, init_final]` instead of the default
orthogonal initialization.
final_activation (`Callable[[torch.Tensor], torch.Tensor] | str | None`, *optional*):
Activation applied after the MLP trunk's last hidden layer, before the output layer.
"""
super().__init__()
self.net = MLP(
input_dim=input_dim,
@@ -622,17 +649,17 @@ class CriticHead(nn.Module):
orthogonal_init()(self.output_layer.weight)
def forward(self, x: torch.Tensor) -> torch.Tensor:
"""Compute the scalar Q-value for `x` (a concatenated observation-encoding + action tensor)."""
return self.output_layer(self.net(x))
class CriticEnsemble(nn.Module):
"""
CriticEnsemble wraps multiple CriticHead modules into an ensemble.
"""CriticEnsemble wraps multiple CriticHead modules into an ensemble.
Args:
encoder (GaussianActorObservationEncoder): encoder for observations.
ensemble (List[CriticHead]): list of critic heads.
init_final (float | None): optional initializer scale for final layers.
init_final (float | None, *optional*): optional initializer scale for final layers.
Forward returns a tensor of shape (num_critics, batch_size) containing Q-values.
"""
@@ -643,6 +670,14 @@ class CriticEnsemble(nn.Module):
ensemble: list[CriticHead],
init_final: float | None = None,
):
"""Wrap `ensemble` behind the shared `encoder`.
Args:
encoder (`GaussianActorObservationEncoder`): Shared observation encoder for all critics.
ensemble (`list[CriticHead]`): The critic heads making up the ensemble.
init_final (`float | None`, *optional*): Stored for introspection; each `CriticHead` is
already initialized with it before being passed in here.
"""
super().__init__()
self.encoder = encoder
self.init_final = init_final
@@ -654,6 +689,19 @@ class CriticEnsemble(nn.Module):
actions: torch.Tensor,
observation_features: torch.Tensor | None = None,
) -> torch.Tensor:
"""Encode `observations` and return each ensemble member's Q-value for `actions`.
Args:
observations (`dict[str, torch.Tensor]`): Raw observation tensors, moved to the module's
device.
actions (`torch.Tensor`): Action tensor to evaluate.
observation_features (`torch.Tensor | None`, *optional*): Pre-computed encoder output,
e.g. from `SACAlgorithm.get_observation_features`. Bypasses re-encoding when the
vision encoder is frozen.
Returns:
torch.Tensor: Q-values of shape `(num_critics, batch_size)`.
"""
device = get_device_from_parameters(self)
# Move each tensor in observations to device
observations = {k: v.to(device) for k, v in observations.items()}
+34 -27
View File
@@ -30,6 +30,19 @@ from lerobot.utils.transition import Transition
class BatchTransition(TypedDict):
"""A batch of transitions sampled from a `ReplayBuffer`.
**Attributes**:
- **state** (`dict[str, torch.Tensor]`) -- Batched observation tensors at time `t`.
- **action** (`torch.Tensor`) -- Batched actions taken at time `t`.
- **reward** (`torch.Tensor`) -- Batched rewards received after `action`.
- **next_state** (`dict[str, torch.Tensor]`) -- Batched observation tensors at time `t+1`.
- **done** (`torch.Tensor`) -- Batched episode-termination flags.
- **truncated** (`torch.Tensor`) -- Batched episode-truncation flags.
- **complementary_info** (`dict[str, torch.Tensor | float | int] | None`) -- Optional extra
per-transition data (e.g. intervention flags), when present in the underlying dataset.
"""
state: dict[str, torch.Tensor]
action: torch.Tensor
reward: torch.Tensor
@@ -40,10 +53,7 @@ class BatchTransition(TypedDict):
def random_crop_vectorized(images: torch.Tensor, output_size: tuple) -> torch.Tensor:
"""
Perform a per-image random crop over a batch of images in a vectorized way.
(Same as shown previously.)
"""
"""Perform a per-image random crop over a batch of images in a vectorized way."""
B, C, H, W = images.shape # noqa: N806
crop_h, crop_w = output_size
@@ -72,13 +82,15 @@ def random_crop_vectorized(images: torch.Tensor, output_size: tuple) -> torch.Te
def random_shift(images: torch.Tensor, pad: int = 4):
"""Vectorized random shift, imgs: (B,C,H,W), pad: #pixels"""
"""Vectorized random shift. `images` has shape `(B, C, H, W)`; `pad` is the shift range in pixels."""
_, _, h, w = images.shape
images = F.pad(input=images, pad=(pad, pad, pad, pad), mode="replicate")
return random_crop_vectorized(images=images, output_size=(h, w))
class ReplayBuffer:
"""In-memory replay buffer of `Transition`s, sampled in batches for off-policy RL training."""
def __init__(
self,
capacity: int,
@@ -89,11 +101,12 @@ class ReplayBuffer:
storage_device: str = "cpu",
optimize_memory: bool = False,
):
"""
Replay buffer for storing transitions.
"""Replay buffer for storing transitions.
It will allocate tensors on the specified device, when the first transition is added.
NOTE: If you encounter memory issues, you can try to use the `optimize_memory` flag to save memory or
and use the `storage_device` flag to store the buffer on a different device.
Args:
capacity (int): Maximum number of transitions to store in the buffer.
device (str): The device where the tensors will be moved when sampling ("cuda:0" or "cpu").
@@ -187,6 +200,7 @@ class ReplayBuffer:
self.initialized = True
def __len__(self):
"""Number of transitions currently stored in the buffer."""
return self.size
def add(
@@ -305,8 +319,8 @@ class ReplayBuffer:
async_prefetch: bool = True,
queue_size: int = 2,
):
"""
Creates an infinite iterator that yields batches of transitions.
"""Creates an infinite iterator that yields batches of transitions.
Will automatically restart when internal iterator is exhausted.
Args:
@@ -329,10 +343,9 @@ class ReplayBuffer:
yield from iterator
def _get_async_iterator(self, batch_size: int, queue_size: int = 2):
"""
Create an iterator that continuously yields prefetched batches in a
background thread. The design is intentionally simple and avoids busy
waiting / complex state management.
"""Create an iterator that continuously yields prefetched batches in a background thread.
The design is intentionally simple and avoids busy waiting / complex state management.
Args:
batch_size (int): Size of batches to sample.
@@ -383,8 +396,7 @@ class ReplayBuffer:
producer_thread.join(timeout=1.0)
def _get_naive_iterator(self, batch_size: int, queue_size: int = 2):
"""
Creates a simple non-threaded iterator that yields batches.
"""Creates a simple non-threaded iterator that yields batches.
Args:
batch_size (int): Size of batches to sample
@@ -398,6 +410,7 @@ class ReplayBuffer:
queue = collections.deque()
def enqueue(n):
"""Sample `n` more batches and append them to `queue`."""
for _ in range(n):
data = self.sample(batch_size)
queue.append(data)
@@ -419,8 +432,7 @@ class ReplayBuffer:
storage_device: str = "cpu",
optimize_memory: bool = False,
) -> "ReplayBuffer":
"""
Convert a LeRobotDataset into a ReplayBuffer.
"""Convert a LeRobotDataset into a ReplayBuffer.
Args:
lerobot_dataset (LeRobotDataset): The dataset to convert.
@@ -509,9 +521,7 @@ class ReplayBuffer:
root=None,
task_name="from_replay_buffer",
) -> LeRobotDataset:
"""
Converts all transitions in this ReplayBuffer into a single LeRobotDataset object.
"""
"""Converts all transitions in this ReplayBuffer into a single LeRobotDataset object."""
if self.size == 0:
raise ValueError("The replay buffer is empty. Cannot convert to a dataset.")
@@ -612,8 +622,7 @@ class ReplayBuffer:
dataset: LeRobotDataset,
state_keys: Sequence[str] | None = None,
) -> list[Transition]:
"""
Convert a LeRobotDataset into a list of RL (s, a, r, s', done) transitions.
"""Convert a LeRobotDataset into a list of RL (s, a, r, s', done) transitions.
Args:
dataset (LeRobotDataset):
@@ -733,12 +742,11 @@ class ReplayBuffer:
# Utility function to guess shapes/dtypes from a tensor
def guess_feature_info(t, name: str):
"""
Return a dictionary with the 'dtype' and 'shape' for a given tensor or scalar value.
"""Return a dictionary with the 'dtype' and 'shape' for a given tensor or scalar value.
If it looks like a 3D (C,H,W) shape, we might consider it an 'image'.
Otherwise default to appropriate dtype for numeric.
"""
shape = tuple(t.shape)
# Basic guess: if we have exactly 3 dims and shape[0] in {1, 3}, guess 'image'
if len(shape) == 3 and shape[0] in [1, 3]:
@@ -757,8 +765,7 @@ def guess_feature_info(t, name: str):
def concatenate_batch_transitions(
left_batch_transitions: BatchTransition, right_batch_transition: BatchTransition
) -> BatchTransition:
"""
Concatenates two BatchTransition objects into one.
"""Concatenates two BatchTransition objects into one.
This function merges the right BatchTransition into the left one by concatenating
all corresponding tensors along dimension 0. The operation modifies the left_batch_transitions
+21 -20
View File
@@ -29,8 +29,7 @@ from lerobot.utils.constants import DONE, REWARD
def select_rect_roi(img):
"""
Allows the user to draw a rectangular ROI on the image.
"""Allows the user to draw a rectangular ROI on the image.
The user must click and drag to draw the rectangle.
- While dragging, the rectangle is dynamically drawn.
@@ -52,6 +51,7 @@ def select_rect_roi(img):
index_x, index_y = -1, -1 # Initial click coordinates
def mouse_callback(event, x, y, flags, param):
"""`cv2.setMouseCallback` handler that drives the click-and-drag ROI selection."""
nonlocal index_x, index_y, drawing, roi, working_img
if event == cv2.EVENT_LBUTTONDOWN:
@@ -118,12 +118,11 @@ def select_rect_roi(img):
def select_square_roi_for_images(images: dict) -> dict:
"""
For each image in the provided dictionary, open a window to allow the user
to select a rectangular ROI. Returns a dictionary mapping each key to a tuple
(top, left, height, width) representing the ROI.
"""For each image in the provided dictionary, open a window to allow the user to select a ROI.
Parameters:
Returns a dictionary mapping each key to a tuple (top, left, height, width) representing the ROI.
Args:
images (dict): Dictionary where keys are identifiers and values are OpenCV images.
Returns:
@@ -149,9 +148,7 @@ def select_square_roi_for_images(images: dict) -> dict:
def get_image_from_lerobot_dataset(dataset: LeRobotDataset):
"""
Find the first row in the dataset and extract the image in order to be used for the crop.
"""
"""Find the first row in the dataset and extract the image in order to be used for the crop."""
row = dataset[0]
image_dict = {}
for k in row:
@@ -169,19 +166,23 @@ def convert_lerobot_dataset_to_cropped_lerobot_dataset(
push_to_hub: bool = False,
task: str = "",
) -> LeRobotDataset:
"""
Converts an existing LeRobotDataset by iterating over its episodes and frames,
applying cropping and resizing to image observations, and saving a new dataset
with the transformed data.
"""Converts an existing LeRobotDataset to a new one with cropped/resized image observations.
Iterates over the source dataset's episodes and frames, applying cropping and resizing to image
observations, and saves a new dataset with the transformed data.
Args:
original_dataset (LeRobotDataset): The source dataset.
crop_params_dict (dict[str, Tuple[int, int, int, int]]):
original_dataset (`LeRobotDataset`): The source dataset.
crop_params_dict (`dict[str, tuple[int, int, int, int]]`):
A dictionary mapping observation keys to crop parameters (top, left, height, width).
new_repo_id (str): Repository id for the new dataset.
new_dataset_root (str): The root directory where the new dataset will be written.
resize_size (tuple[int, int], optional): The target size (height, width) after cropping.
Defaults to (128, 128).
new_repo_id (`str`): Repository id for the new dataset.
new_dataset_root (`str`): The root directory where the new dataset will be written.
resize_size (`tuple[int, int]`, *optional*, defaults to `(128, 128)`): The target size
(height, width) after cropping.
push_to_hub (`bool`, *optional*, defaults to `False`): Whether to push the new dataset to the
Hugging Face Hub.
task (`str`, *optional*, defaults to `""`): Task description recorded on every frame of the
new dataset.
Returns:
LeRobotDataset: A new LeRobotDataset where the specified image observations have been cropped
+13 -1
View File
@@ -49,6 +49,18 @@ class OnlineOfflineMixer(DataMixer):
offline_buffer: ReplayBuffer | None = None,
online_ratio: float = 1.0,
):
"""Create the mixer.
Args:
online_buffer (`ReplayBuffer`): Buffer of transitions collected online during training.
offline_buffer (`ReplayBuffer | None`, *optional*): Buffer of pre-collected offline
transitions. When `None`, every batch is drawn from `online_buffer` alone.
online_ratio (`float`, *optional*, defaults to 1.0): Fraction of each batch drawn from
`online_buffer`; the remainder comes from `offline_buffer`. Must be in `[0, 1]`.
Raises:
ValueError: If `online_ratio` is not in `[0, 1]`.
"""
if not 0.0 <= online_ratio <= 1.0:
raise ValueError(f"online_ratio must be in [0, 1], got {online_ratio}")
self.online_buffer = online_buffer
@@ -56,6 +68,7 @@ class OnlineOfflineMixer(DataMixer):
self.online_ratio = online_ratio
def sample(self, batch_size: int) -> BatchType:
"""See [`~rl.data_sources.DataMixer.sample`]."""
if self.offline_buffer is None:
return self.online_buffer.sample(batch_size)
@@ -73,7 +86,6 @@ class OnlineOfflineMixer(DataMixer):
queue_size: int = 2,
):
"""Yield batches by composing buffer async iterators."""
n_online = max(1, int(batch_size * self.online_ratio))
online_iter = self.online_buffer.get_iterator(
+13
View File
@@ -36,6 +36,13 @@ logging.basicConfig(level=logging.INFO)
def eval_policy(env, policy, n_episodes):
"""Roll out `policy` in `env` for `n_episodes` and log the per-episode and average reward.
Args:
env (`gymnasium.Env`): A robot environment, built via `make_robot_env`.
policy (`PreTrainedPolicy`): A policy exposing `select_action(obs) -> action`.
n_episodes (`int`): Number of episodes to run.
"""
sum_reward_episode = []
for _ in range(n_episodes):
obs, _ = env.reset()
@@ -54,6 +61,12 @@ def eval_policy(env, policy, n_episodes):
@parser.wrap()
def main(cfg: TrainRLServerPipelineConfig):
"""CLI entry point: load a pretrained policy and evaluate it for 10 episodes.
Args:
cfg (`TrainRLServerPipelineConfig`): Parsed from the CLI. `cfg.env.pretrained_policy_name_or_path`
selects the checkpoint to load; `cfg.dataset.repo_id` provides normalization stats.
"""
env_cfg = cfg.env
env = make_robot_env(env_cfg)
dataset_cfg = cfg.dataset
+27 -19
View File
@@ -305,7 +305,9 @@ def make_robot_env(cfg: HILSerlRobotEnvConfig) -> tuple[gym.Env, Any]:
"""Create robot environment from configuration.
Args:
cfg: Environment configuration.
cfg (`HILSerlRobotEnvConfig`): Environment configuration. `cfg.name == "gym_hil"` selects the
GymHIL simulation environment; otherwise a real-robot `RobotEnv` is built from
`cfg.robot`/`cfg.teleop`.
Returns:
Tuple of (gym environment, teleoperator device).
@@ -363,10 +365,13 @@ def make_processors(
"""Create environment and action processors.
Args:
env: Robot environment instance.
teleop_device: Teleoperator device for intervention.
cfg: Processor configuration.
device: Target device for computations.
env (`Env`): The environment returned by `make_robot_env`.
teleop_device (`lerobot.teleoperators.teleoperator.Teleoperator | None`): The teleoperator
device returned by `make_robot_env`, used to configure intervention-related processor
steps. `None` for simulation environments.
cfg (`HILSerlRobotEnvConfig`): Environment configuration; provides the reward classifier,
gripper, and reset-behavior settings for the built processor steps.
device (`str`, *optional*, defaults to `"cpu"`): Torch device the processors run on.
Returns:
Tuple of (environment processor, action processor).
@@ -536,20 +541,21 @@ def step_env_and_process_transition(
env_processor: DataProcessorPipeline[EnvTransition, EnvTransition],
action_processor: DataProcessorPipeline[EnvTransition, EnvTransition],
) -> EnvTransition:
"""
Execute one step with processor pipeline.
"""Execute one step with processor pipeline.
Args:
env: The robot environment
transition: Current transition state
action: Action to execute
env_processor: Environment processor
action_processor: Action processor
env (`Env`): The environment to step.
transition (`EnvTransition`): The current transition; its observation is overwritten with the
action processor's input before dispatch, then discarded.
action (`Tensor`): The raw action to process and send to `env`.
env_processor (`DataProcessorPipeline`): Post-processes the environment-produced transition
(e.g. reward shaping, termination overrides).
action_processor (`DataProcessorPipeline`): Pre-processes `action` before it reaches `env`
(e.g. intervention overrides, gripper handling).
Returns:
Processed transition with updated state.
"""
# Create action transition
transition[TransitionKey.ACTION] = action
transition[TransitionKey.OBSERVATION] = (
@@ -618,14 +624,16 @@ def control_loop(
cfg: GymManipulatorConfig,
) -> None:
"""Main control loop for robot environment interaction.
if cfg.mode == "record": then a dataset will be created and recorded
When `cfg.mode == "record"`, a dataset is created and recorded.
Args:
env: The robot environment
env_processor: Environment processor
action_processor: Action processor
teleop_device: Teleoperator device
cfg: gym_manipulator configuration
env (`Env`): The environment to control, built via `make_robot_env`.
env_processor (`DataProcessorPipeline`): Post-processes environment-produced transitions.
action_processor (`DataProcessorPipeline`): Pre-processes teleoperator actions before they
reach `env`.
teleop_device (`Teleoperator`): Teleoperator device driving the robot.
cfg (`GymManipulatorConfig`): Control-loop configuration (mode, fps, episode/dataset settings).
"""
dt = 1.0 / cfg.env.fps
+7 -14
View File
@@ -31,8 +31,7 @@ from lerobot.utils.constants import OBS_STATE
@dataclass
@ProcessorStepRegistry.register("joint_velocity_processor")
class JointVelocityProcessorStep(ObservationProcessorStep):
"""
Calculates and appends joint velocity information to the observation state.
"""Calculates and appends joint velocity information to the observation state.
This step computes the velocity of each joint by calculating the finite
difference between the current and the last observed joint positions. The
@@ -50,8 +49,7 @@ class JointVelocityProcessorStep(ObservationProcessorStep):
last_joint_positions: torch.Tensor | None = None
def observation(self, observation: dict) -> dict:
"""
Computes joint velocities and adds them to the observation state.
"""Computes joint velocities and adds them to the observation state.
Args:
observation: The input observation dictionary, expected to contain
@@ -89,8 +87,7 @@ class JointVelocityProcessorStep(ObservationProcessorStep):
return new_observation
def get_config(self) -> dict[str, Any]:
"""
Returns the configuration of the step for serialization.
"""Returns the configuration of the step for serialization.
Returns:
A dictionary containing the time step `dt`.
@@ -106,8 +103,7 @@ class JointVelocityProcessorStep(ObservationProcessorStep):
def transform_features(
self, features: dict[PipelineFeatureType, dict[str, PolicyFeature]]
) -> dict[PipelineFeatureType, dict[str, PolicyFeature]]:
"""
Updates the `observation.state` feature to reflect the added velocities.
"""Updates the `observation.state` feature to reflect the added velocities.
This method doubles the size of the first dimension of the `observation.state`
shape to account for the concatenation of position and velocity vectors.
@@ -132,8 +128,7 @@ class JointVelocityProcessorStep(ObservationProcessorStep):
@dataclass
@ProcessorStepRegistry.register("current_processor")
class MotorCurrentProcessorStep(ObservationProcessorStep):
"""
Reads motor currents from a robot and appends them to the observation state.
"""Reads motor currents from a robot and appends them to the observation state.
This step queries the robot's hardware interface to get the present current
for each motor and concatenates this information to the existing state vector.
@@ -146,8 +141,7 @@ class MotorCurrentProcessorStep(ObservationProcessorStep):
robot: Robot | None = None
def observation(self, observation: dict) -> dict:
"""
Fetches motor currents and adds them to the observation state.
"""Fetches motor currents and adds them to the observation state.
Args:
observation: The input observation dictionary.
@@ -184,8 +178,7 @@ class MotorCurrentProcessorStep(ObservationProcessorStep):
def transform_features(
self, features: dict[PipelineFeatureType, dict[str, PolicyFeature]]
) -> dict[PipelineFeatureType, dict[str, PolicyFeature]]:
"""
Updates the `observation.state` feature to reflect the added motor currents.
"""Updates the `observation.state` feature to reflect the added motor currents.
This method increases the size of the first dimension of the `observation.state`
shape by the number of motors in the robot.
+85 -70
View File
@@ -14,8 +14,7 @@
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""
Learner server runner for distributed HILSerl robot policy training.
"""Learner server runner for distributed HILSerl robot policy training.
This script implements the learner component of the distributed HILSerl architecture.
It initializes the policy network, maintains replay buffers, and updates
@@ -121,6 +120,11 @@ from .trainer import RLTrainer
@parser.wrap()
def train_cli(cfg: TrainRLServerPipelineConfig):
"""CLI entry point for the HILSerl learner server.
Args:
cfg (`TrainRLServerPipelineConfig`): Parsed from the CLI, forwarded to `train`.
"""
# Fail fast with a friendly error if the optional ``hilserl`` extra is missing.
require_package("grpcio", extra="hilserl", import_name="grpc")
if not use_threads(cfg):
@@ -136,14 +140,13 @@ def train_cli(cfg: TrainRLServerPipelineConfig):
def train(cfg: TrainRLServerPipelineConfig, job_name: str | None = None):
"""
Main training function that initializes and runs the training process.
"""Main training function that initializes and runs the training process.
Args:
cfg (TrainRLServerPipelineConfig): The training configuration
job_name (str | None, optional): Job name for logging. Defaults to None.
cfg (`TrainRLServerPipelineConfig`): The training configuration.
job_name (`str | None`, *optional*): Job name for logging. Defaults to `cfg.job_name` when
unset.
"""
cfg.validate()
if job_name is None:
@@ -198,13 +201,12 @@ def start_learner_threads(
wandb_logger: WandBLogger | None,
shutdown_event: Any, # Event
) -> None:
"""
Start the learner threads for training.
"""Start the learner threads for training.
Args:
cfg (TrainRLServerPipelineConfig): Training configuration
wandb_logger (WandBLogger | None): Logger for metrics
shutdown_event: Event to signal shutdown
cfg (`TrainRLServerPipelineConfig`): Training configuration.
wandb_logger (`WandBLogger | None`): Logger for metrics.
shutdown_event (`Event`): Event signaling the learner and its background workers to stop.
"""
# Create multiprocessing queues
transition_queue = Queue()
@@ -275,9 +277,7 @@ def add_actor_information_and_train(
interaction_message_queue: Queue,
parameters_queue: Queue,
):
"""
Handles data transfer from the actor to the learner, manages training updates,
and logs training progress in an online reinforcement learning setup.
"""Handles data transfer from the actor to the learner, manages training updates, and logs progress.
This function continuously:
- Transfers transitions from the actor to the replay buffer.
@@ -482,17 +482,18 @@ def start_learner(
shutdown_event: Any, # Event
cfg: TrainRLServerPipelineConfig,
):
"""
Start the learner server for training.
It will receive transitions and interaction messages from the actor server,
and send policy parameters to the actor server.
"""Start the learner server for training.
Receives transitions and interaction messages from the actor server, and sends policy parameters
to the actor server.
Args:
parameters_queue: Queue for sending policy parameters to the actor
transition_queue: Queue for receiving transitions from the actor
interaction_message_queue: Queue for receiving interaction messages from the actor
shutdown_event: Event to signal shutdown
cfg: Training configuration
parameters_queue (`Queue`): Queue of serialized policy weights, drained and streamed to the
actor by `LearnerService.StreamParameters`.
transition_queue (`Queue`): Queue filled by `LearnerService.SendTransitions`.
interaction_message_queue (`Queue`): Queue filled by `LearnerService.SendInteractions`.
shutdown_event (`Event`): Event signaling this process/thread to stop.
cfg (`TrainRLServerPipelineConfig`): Training configuration.
"""
if not use_threads(cfg):
# Create a process-specific log file
@@ -560,8 +561,7 @@ def save_training_checkpoint(
preprocessor=None,
postprocessor=None,
) -> None:
"""
Save training checkpoint and associated data.
"""Save training checkpoint and associated data.
This function performs the following steps:
1. Creates a checkpoint directory with the current optimization step
@@ -572,18 +572,26 @@ def save_training_checkpoint(
6. If an offline replay buffer exists, saves it as a separate dataset
Args:
cfg: Training configuration
optimization_step: Current optimization step
online_steps: Total number of online steps
interaction_message: Dictionary containing interaction information
policy: Policy model to save
optimizers: Dictionary of optimizers
replay_buffer: Replay buffer to save as dataset
offline_replay_buffer: Optional offline replay buffer to save
dataset_repo_id: Repository ID for dataset
fps: Frames per second for dataset
preprocessor: Optional preprocessor pipeline to save
postprocessor: Optional postprocessor pipeline to save
cfg (`TrainRLServerPipelineConfig`): Training configuration, saved alongside the checkpoint.
optimization_step (`int`): Current optimization step; used to name the checkpoint directory.
online_steps (`int`): Total number of online steps; used to size the checkpoint directory's
zero-padded step number.
interaction_message (`dict | None`): Latest interaction message; its `"Interaction step"`
entry is saved for resuming training.
policy (`Module`): Policy model to save.
optimizers (`dict`): Dictionary of optimizers whose states are saved.
replay_buffer (`ReplayBuffer`): Replay buffer to save as a dataset.
algorithm (`lerobot.rl.algorithms.base.RLAlgorithm | None`, *optional*): Algorithm whose state
dict (critic ensembles, temperature, etc.) should also be saved.
offline_replay_buffer (`lerobot.rl.buffer.ReplayBuffer | None`, *optional*): Optional offline
replay buffer, saved as a separate dataset when provided.
dataset_repo_id (`str | None`, *optional*): Repository id used when converting the replay
buffer(s) to a dataset.
fps (`int`, *optional*, defaults to 30): Frames per second recorded on the saved dataset(s).
preprocessor (`PolicyProcessorPipeline | None`, *optional*): Optional preprocessor pipeline to
save alongside the policy.
postprocessor (`PolicyProcessorPipeline | None`, *optional*): Optional postprocessor pipeline
to save alongside the policy.
"""
logging.info(f"Checkpoint policy after step {optimization_step}")
_num_digits = max(6, len(str(online_steps)))
@@ -650,8 +658,7 @@ def save_training_checkpoint(
def handle_resume_logic(cfg: TrainRLServerPipelineConfig) -> TrainRLServerPipelineConfig:
"""
Handle the resume logic for training.
"""Handle the resume logic for training.
If resume is True:
- Verifies that a checkpoint exists
@@ -712,19 +719,19 @@ def load_training_state(
algorithm: RLAlgorithm | None = None,
device: str | torch.device = "cpu",
):
"""
Loads the training state (optimizers, RNG, step + interaction step, and
algorithm-owned tensors) from the most recent checkpoint.
"""Loads the training state from the most recent checkpoint.
Restores optimizers, RNG state, the optimization/interaction step, and algorithm-owned tensors.
Args:
cfg (TrainRLServerPipelineConfig): Training configuration; `cfg.resume` gates the load and
cfg (`TrainRLServerPipelineConfig`): Training configuration; `cfg.resume` gates the load and
`cfg.output_dir` locates the last checkpoint.
optimizers (Optimizer | dict[str, Optimizer]): Optimizers to load state into.
algorithm (RLAlgorithm | None, optional): Algorithm whose state dict should be restored.
optimizers (`Optimizer | dict[str, Optimizer]`): Optimizers to load state into.
algorithm (`RLAlgorithm | None`, *optional*): Algorithm whose state dict should be restored.
Required for full main-equivalent resume; the policy itself is restored separately via
`make_policy`. Defaults to None.
device (str | torch.device, optional): Device on which to place loaded algorithm tensors.
Defaults to "cpu".
`make_policy`.
device (`str | torch.device`, *optional*, defaults to `"cpu"`): Device on which to place
loaded algorithm tensors.
Returns:
tuple[int | None, int | None]: `(optimization_step, interaction_step)`, or `(None, None)`
@@ -772,8 +779,7 @@ def load_training_state(
def log_training_info(cfg: TrainRLServerPipelineConfig, policy: nn.Module) -> None:
"""
Log information about the training process.
"""Log information about the training process.
Args:
cfg (TrainRLServerPipelineConfig): Training configuration
@@ -792,8 +798,7 @@ def log_training_info(cfg: TrainRLServerPipelineConfig, policy: nn.Module) -> No
def initialize_replay_buffer(
cfg: TrainRLServerPipelineConfig, device: str, storage_device: str
) -> ReplayBuffer:
"""
Initialize a replay buffer, either empty or from a dataset if resuming.
"""Initialize a replay buffer, either empty or from a dataset if resuming.
Args:
cfg (TrainRLServerPipelineConfig): Training configuration
@@ -837,8 +842,7 @@ def initialize_offline_replay_buffer(
device: str,
storage_device: str,
) -> ReplayBuffer:
"""
Initialize an offline replay buffer from a dataset.
"""Initialize an offline replay buffer from a dataset.
Args:
cfg (TrainRLServerPipelineConfig): Training configuration
@@ -875,6 +879,7 @@ def initialize_offline_replay_buffer(
def use_threads(cfg: TrainRLServerPipelineConfig) -> bool:
"""Whether the learner's background workers should run as threads instead of processes."""
return cfg.policy.concurrency.learner == "threads"
@@ -884,14 +889,14 @@ def check_nan_in_transition(
next_state: torch.Tensor,
raise_error: bool = False,
) -> bool:
"""
Check for NaN values in transition data.
"""Check for NaN values in transition data.
Args:
observations: Dictionary of observation tensors
actions: Action tensor
next_state: Dictionary of next state tensors
raise_error: If True, raises ValueError when NaN is detected
observations (`Tensor`): Dictionary of observation tensors.
actions (`Tensor`): Action tensor.
next_state (`Tensor`): Dictionary of next-observation tensors.
raise_error (`bool`, *optional*, defaults to `False`): Whether to raise a `ValueError` instead
of just logging when a NaN is found.
Returns:
bool: True if NaN values were detected, False otherwise
@@ -925,6 +930,12 @@ def check_nan_in_transition(
def push_actor_policy_to_queue(parameters_queue: Queue, algorithm: RLAlgorithm) -> None:
"""Serialize `algorithm`'s current weights and enqueue them for the actor-facing gRPC stream.
Args:
parameters_queue (`Queue`): Queue drained by `LearnerService.StreamParameters`.
algorithm (`RLAlgorithm`): Source of the weights, via `get_weights`.
"""
logging.debug("[LEARNER] Pushing actor policy to the queue")
# Create a dictionary to hold all the state dicts
@@ -958,11 +969,13 @@ def process_transitions(
"""Process all available transitions from the queue.
Args:
transition_queue: Queue for receiving transitions from the actor
replay_buffer: Replay buffer to add transitions to
offline_replay_buffer: Offline replay buffer to add transitions to
dataset_repo_id: Repository ID for dataset
shutdown_event: Event to signal shutdown
transition_queue (`Queue`): Queue filled by `LearnerService.SendTransitions`.
replay_buffer (`ReplayBuffer`): Buffer every non-NaN transition is added to.
offline_replay_buffer (`ReplayBuffer`): Buffer intervention transitions are additionally added
to, when `dataset_repo_id` is set.
dataset_repo_id (`str | None`): When set, transitions tagged as interventions are also added
to `offline_replay_buffer`.
shutdown_event (`Event`): Event that stops the loop when set.
"""
while not transition_queue.empty() and not shutdown_event.is_set():
transition_list = transition_queue.get()
@@ -996,10 +1009,12 @@ def process_interaction_messages(
"""Process all available interaction messages from the queue.
Args:
interaction_message_queue: Queue for receiving interaction messages
interaction_step_shift: Amount to shift interaction step by
wandb_logger: Logger for tracking progress
shutdown_event: Event to signal shutdown
interaction_message_queue (`Queue`): Queue filled by `LearnerService.SendInteractions`.
interaction_step_shift (`int`): Offset added to each message's `"Interaction step"` so it
stays consistent with checkpointed state after a resume.
wandb_logger (`lerobot.common.wandb_utils.WandBLogger | None`): Logger the message is
forwarded to, when set.
shutdown_event (`Event`): Event that stops the loop when set.
Returns:
dict | None: The last interaction message processed, or None if none were processed
+49 -4
View File
@@ -44,10 +44,10 @@ SHUTDOWN_TIMEOUT = 10
class LearnerService(_ServicerBase):
"""
Implementation of the LearnerService gRPC service
This service is used to send parameters to the Actor and receive transitions and interactions from the Actor
check transport.proto for the gRPC service definition
"""Implementation of the LearnerService gRPC service.
Sends policy parameters to the actor and receives transitions and interactions from it; see
`transport.proto` for the gRPC service definition.
"""
def __init__(
@@ -59,6 +59,19 @@ class LearnerService(_ServicerBase):
interaction_message_queue: Queue,
queue_get_timeout: float = 0.001,
):
"""Create the servicer.
Args:
shutdown_event (`Event`): Set to stop `StreamParameters`'s push loop.
parameters_queue (`Queue`): Queue of serialized policy weights, drained and streamed to
the actor by `StreamParameters`.
seconds_between_pushes (`float`): Minimum interval between successive parameter pushes.
transition_queue (`Queue`): Queue filled by `SendTransitions` with received transitions.
interaction_message_queue (`Queue`): Queue filled by `SendInteractions` with received
interaction messages.
queue_get_timeout (`float`, *optional*, defaults to 0.001): Timeout used when polling
`parameters_queue`.
"""
self.shutdown_event = shutdown_event
self.parameters_queue = parameters_queue
self.seconds_between_pushes = seconds_between_pushes
@@ -69,6 +82,17 @@ class LearnerService(_ServicerBase):
def StreamParameters( # noqa: N802
self, request: "services_pb2.Empty", context: "grpc.ServicerContext"
):
"""GRPC server-streaming RPC: push the latest policy parameters to the actor.
Runs until `shutdown_event` is set, pushing at most once every `seconds_between_pushes`.
Args:
request (`services_pb2.Empty`): Unused; required by the gRPC service signature.
context (`grpc.ServicerContext`): gRPC call context.
Yields:
Chunks of a `services_pb2.Parameters` message, produced by `send_bytes_in_chunks`.
"""
# TODO: authorize the request
logging.info("[LEARNER] Received request to stream parameters from the Actor")
@@ -104,6 +128,16 @@ class LearnerService(_ServicerBase):
return services_pb2.Empty()
def SendTransitions(self, request_iterator, _context: "grpc.ServicerContext"): # noqa: N802
"""GRPC client-streaming RPC: receive transition chunks from the actor into `transition_queue`.
Args:
request_iterator: Stream of `services_pb2.Transition` chunks sent by the actor's
`transitions_stream`.
_context (`grpc.ServicerContext`): gRPC call context.
Returns:
services_pb2.Empty: Acknowledgement sent once the actor closes the stream.
"""
# TODO: authorize the request
logging.info("[LEARNER] Received request to receive transitions from the Actor")
@@ -118,6 +152,16 @@ class LearnerService(_ServicerBase):
return services_pb2.Empty()
def SendInteractions(self, request_iterator, _context: "grpc.ServicerContext"): # noqa: N802
"""GRPC client-streaming RPC: receive interaction-message chunks into `interaction_message_queue`.
Args:
request_iterator: Stream of `services_pb2.InteractionMessage` chunks sent by the actor's
`interactions_stream`.
_context (`grpc.ServicerContext`): gRPC call context.
Returns:
services_pb2.Empty: Acknowledgement sent once the actor closes the stream.
"""
# TODO: authorize the request
logging.info("[LEARNER] Received request to receive interactions from the Actor")
@@ -132,4 +176,5 @@ class LearnerService(_ServicerBase):
return services_pb2.Empty()
def Ready(self, request: "services_pb2.Empty", context: "grpc.ServicerContext"): # noqa: N802
"""GRPC health check: returns immediately, confirming the learner server is up."""
return services_pb2.Empty()
+13
View File
@@ -23,6 +23,19 @@ from torch.multiprocessing import Queue
def get_last_item_from_queue(queue: Queue, block=True, timeout: float = 0.1) -> Any:
"""Drain `queue` and return only the most recently enqueued item.
Args:
queue (`Queue`): A `torch.multiprocessing.Queue` to drain.
block (`bool`, *optional*, defaults to `True`): Whether to block for up to `timeout` seconds
waiting for a first item before draining. When `False`, returns `None` if the queue is
currently empty.
timeout (`float`, *optional*, defaults to 0.1): Seconds to wait for a first item when `block`
is `True`.
Returns:
Any: The most recent item, or `None` if the queue was (and stayed) empty.
"""
if block:
try:
item = queue.get(timeout=timeout)
+99
View File
@@ -28,6 +28,100 @@ from .algorithms.sac import SACAlgorithmConfig # noqa: F401
@dataclass(kw_only=True)
class TrainRLServerPipelineConfig(TrainPipelineConfig):
"""Top-level config for the actor/learner distributed RL training server.
Extends [`~configs.train.TrainPipelineConfig`] with an optional (rather than required) offline
`dataset` and the RL-specific algorithm/data-mixing fields below.
Args:
env (`lerobot.envs.configs.EnvConfig | None`, *optional*):
Simulation environment configuration, used for `env_eval_freq` evaluation rollouts.
policy (`lerobot.configs.policies.PreTrainedConfig | None`, *optional*):
The actor policy configuration.
reward_model (`lerobot.configs.rewards.RewardModelConfig | None`, *optional*):
Reward model configuration, when training a reward model instead of a policy.
output_dir (`pathlib.Path | None`, *optional*):
Directory to save run outputs to. Reusing the same value across runs overwrites its
contents unless `resume` is `True`.
job_name (`str | None`, *optional*):
Name used for logging and checkpoint directory naming.
resume (`bool`, *optional*, defaults to `False`):
Whether to resume a previous run from `--config_path`'s checkpoint.
seed (`int | None`, *optional*, defaults to 1000):
Random seed for model initialization, dataset shuffling, and evaluation environments.
cudnn_deterministic (`bool`, *optional*, defaults to `False`):
Whether to use deterministic cuDNN algorithms for reproducibility. Disables
`cudnn.benchmark`, which may reduce training speed.
num_workers (`int`, *optional*, defaults to 4):
Number of dataloader worker processes.
batch_size (`int`, *optional*, defaults to 8):
Offline-dataset dataloader batch size.
prefetch_factor (`int`, *optional*, defaults to 4):
Number of batches prefetched per dataloader worker.
persistent_workers (`bool`, *optional*, defaults to `True`):
Whether dataloader workers stay alive between epochs.
dataloader_multiprocessing_context (`str | None`, *optional*, defaults to `"spawn"`):
DataLoader worker start method. `None` uses Python's platform default.
steps (`int`, *optional*, defaults to 100000):
Total number of training steps.
env_eval_freq (`int`, *optional*, defaults to 20000):
Run the policy in the simulation environment every N steps to measure reward/success.
`0` disables environment evaluation.
log_freq (`int`, *optional*, defaults to 200):
Log training metrics every N steps.
eval_steps (`int`, *optional*, defaults to 0):
Compute eval loss on held-out episodes every N steps. `0` disables it.
max_eval_samples (`int`, *optional*, defaults to 0):
Cap on total eval samples, split uniformly across tasks. `0` uses all held-out data.
tolerance_s (`float`, *optional*, defaults to 0.0001):
Maximum timestamp tolerance, in seconds, when loading dataset frames.
save_checkpoint (`bool`, *optional*, defaults to `True`):
Whether to save training checkpoints at all.
save_freq (`int`, *optional*, defaults to 20000):
Save a checkpoint every N training steps, and after the last step. A non-positive value
disables periodic saving, keeping only the final checkpoint.
checkpoint_format (`CheckpointFormat`, *optional*, defaults to `CheckpointFormat.SAFETENSORS`):
Model-artifact format inside checkpoints.
use_policy_training_preset (`bool`, *optional*, defaults to `True`):
Whether to use the policy's own recommended optimizer/scheduler preset when `optimizer`/
`scheduler` are unset.
optimizer (`lerobot.optim.optimizers.OptimizerConfig | None`, *optional*):
Optimizer configuration override.
scheduler (`lerobot.optim.schedulers.LRSchedulerConfig | None`, *optional*):
Learning-rate scheduler configuration override.
parallelism (`ParallelismConfig`, *optional*):
Process topology: `dp_replicate`/`dp_shard` (HSDP) and context-parallel degree.
accelerator (`AcceleratorConfig`, *optional*):
Execution runtime handed to the Accelerator: mixed precision, gradient accumulation,
FSDP/DDP tuning knobs, compile & activation-checkpointing.
eval (`EvalConfig`, *optional*):
Simulation-environment evaluation configuration (number of episodes, batch size).
wandb (`WandBConfig`, *optional*):
Weights & Biases logging configuration.
peft (`lerobot.configs.default.PeftConfig | None`, *optional*):
PEFT (e.g. LoRA) adapter configuration for parameter-efficient fine-tuning.
job (`JobConfig`, *optional*):
Where to run training: local (default) or an HF Jobs flavor.
save_checkpoint_to_hub (`bool`, *optional*, defaults to `False`):
Whether to push each saved checkpoint to the Hub as it is written, not just the final
model.
sample_weighting (`lerobot.utils.sample_weighting.SampleWeightingConfig | None`, *optional*):
Sample weighting configuration (e.g. for RA-BC training).
rename_map (`dict`, *optional*):
Mapping to override observation image/state key names.
dataset (`DatasetConfig | None`, *optional*):
Optional offline dataset config. Unlike imitation-learning training, RL doesn't require an
offline dataset data comes from the online replay buffer.
algorithm (`RLAlgorithmConfig | None`, *optional*):
RL algorithm configuration. Defaults to a SAC config (with `policy_config` populated from
`self.policy`) in `validate` when unset.
mixer (`str`, *optional*, defaults to `"online_offline"`):
Data mixer strategy name. Currently only `"online_offline"` is supported.
online_ratio (`float`, *optional*, defaults to 0.5):
Fraction of each training batch sampled from the online replay buffer when using
`OnlineOfflineMixer`; the remainder comes from the offline dataset.
"""
# NOTE: In RL, we don't need an offline dataset
# TODO: Make `TrainPipelineConfig.dataset` optional
dataset: DatasetConfig | None = None # type: ignore[assignment] # because the parent class has made it's type non-optional
@@ -41,6 +135,11 @@ class TrainRLServerPipelineConfig(TrainPipelineConfig):
online_ratio: float = 0.5
def validate(self) -> None:
"""See [`~configs.train.TrainPipelineConfig.validate`].
Additionally defaults `algorithm` to a SAC config and populates its `policy_config` from
`self.policy` when unset.
"""
super().validate()
if self.algorithm is None:
+13
View File
@@ -38,6 +38,16 @@ class RLTrainer:
*,
preprocessor: Any | None = None,
):
"""Build the trainer and its optimizers.
Args:
algorithm (`RLAlgorithm`): The RL algorithm to train. `make_optimizers_and_scheduler` is
called on it immediately.
data_mixer (`DataMixer`): Data source the training-batch iterator is built from.
batch_size (`int`): Batch size requested from `data_mixer` on each training step.
preprocessor (`Any | None`, *optional*): When set, each sampled batch is passed through
`preprocess_rl_batch` before reaching the algorithm.
"""
self.algorithm = algorithm
self.data_mixer = data_mixer
self.batch_size = batch_size
@@ -90,12 +100,15 @@ class _PreprocessedIterator:
__slots__ = ("_raw", "_preprocessor")
def __init__(self, raw_iterator: Iterator[BatchType], preprocessor: Any) -> None:
"""Wrap `raw_iterator`, applying `preprocessor` to each yielded batch."""
self._raw = raw_iterator
self._preprocessor = preprocessor
def __iter__(self) -> _PreprocessedIterator:
"""Return `self` (this object is its own iterator)."""
return self
def __next__(self) -> BatchType:
"""Return the next preprocessed batch from the wrapped iterator."""
batch = next(self._raw)
return preprocess_rl_batch(self._preprocessor, batch)
+1
View File
@@ -60,6 +60,7 @@ PATH_TO_LEROBOT = PATH_TO_REPO / "src" / "lerobot"
# Modules whose public objects are checked. Add a module here once its docstrings follow the standard.
MODULES_TO_CHECK = [
"lerobot.robots",
"lerobot.rl",
]
# Objects that do not yet follow the standard, so the check can be green from day one. Removing an entry