mirror of
https://github.com/huggingface/lerobot.git
synced 2026-08-08 17:39:44 +00:00
Compare commits
1 Commits
| Author | SHA1 | Date | |
|---|---|---|---|
| 1834f819a6 |
@@ -215,4 +215,6 @@
|
||||
title: Environments
|
||||
- local: api/configs
|
||||
title: Configuration
|
||||
- local: api/rl
|
||||
title: Reinforcement Learning
|
||||
title: "API Reference"
|
||||
|
||||
@@ -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
@@ -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
@@ -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"
|
||||
|
||||
|
||||
|
||||
@@ -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]:
|
||||
|
||||
@@ -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():
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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
@@ -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
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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(
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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
|
||||
|
||||
|
||||
@@ -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
@@ -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
|
||||
|
||||
@@ -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()
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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:
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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
|
||||
|
||||
Reference in New Issue
Block a user