# !/usr/bin/env python # Copyright 2025 The HuggingFace Inc. team. # All rights reserved. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # 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. This script implements the learner component of the distributed HILSerl architecture. It initializes the policy network, maintains replay buffers, and updates the policy based on transitions received from the actor server. Examples of usage: - Start a learner server for training: ```bash python -m lerobot.rl.learner --config_path src/lerobot/configs/train_config_hilserl_so100.json ``` **NOTE**: Start the learner server before launching the actor server. The learner opens a gRPC server to communicate with actors. **NOTE**: Training progress can be monitored through Weights & Biases if wandb.enable is set to true in your configuration. **WORKFLOW**: 1. Create training configuration with proper policy, dataset, and environment settings 2. Start this learner server with the configuration 3. Start an actor server with the same configuration 4. Monitor training progress through wandb dashboard For more details on the complete HILSerl training workflow, see: https://github.com/michel-aractingi/lerobot-hilserl-guide """ import logging import os import shutil import time from concurrent.futures import ThreadPoolExecutor from pathlib import Path from pprint import pformat from typing import TYPE_CHECKING, Any from lerobot.utils.import_utils import _grpc_available, require_package if TYPE_CHECKING or _grpc_available: import grpc from lerobot.transport import services_pb2_grpc else: grpc = None services_pb2_grpc = None import torch from huggingface_hub.constants import SAFETENSORS_SINGLE_FILE from safetensors.torch import load_file as load_safetensors from termcolor import colored from torch import nn from torch.multiprocessing import Queue from torch.optim.optimizer import Optimizer from lerobot.cameras import opencv # noqa: F401 from lerobot.common.train_utils import ( get_step_checkpoint_dir, load_training_state as utils_load_training_state, save_checkpoint, update_last_checkpoint, ) from lerobot.common.wandb_utils import WandBLogger from lerobot.configs import parser from lerobot.datasets import LeRobotDataset, make_dataset from lerobot.policies import make_policy, make_pre_post_processors from lerobot.robots import so_follower # noqa: F401 from lerobot.teleoperators import gamepad, so_leader # noqa: F401 from lerobot.teleoperators.utils import TeleopEvents from lerobot.transport.utils import ( MAX_MESSAGE_SIZE, bytes_to_python_object, bytes_to_transitions, state_to_bytes, ) from lerobot.utils.constants import ( ACTION, ALGORITHM_DIR, CHECKPOINTS_DIR, LAST_CHECKPOINT_LINK, PRETRAINED_MODEL_DIR, TRAINING_STATE_DIR, TRAINING_STEP, ) from lerobot.utils.device_utils import get_safe_torch_device from lerobot.utils.io_utils import load_json, write_json from lerobot.utils.process import ProcessSignalHandler from lerobot.utils.random_utils import set_seed from lerobot.utils.utils import ( format_big_number, init_logging, ) from .algorithms.base import RLAlgorithm from .algorithms.factory import make_algorithm from .buffer import ReplayBuffer from .data_sources import OnlineOfflineMixer from .learner_service import MAX_WORKERS, SHUTDOWN_TIMEOUT, LearnerService from .train_rl import TrainRLServerPipelineConfig from .trainer import RLTrainer @parser.wrap() def train_cli(cfg: TrainRLServerPipelineConfig): # 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): import torch.multiprocessing as mp mp.set_start_method("spawn") # Use the job_name from the config train( cfg, job_name=cfg.job_name, ) logging.info("[LEARNER] train_cli finished") def train(cfg: TrainRLServerPipelineConfig, job_name: str | None = None): """ 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.validate() if job_name is None: job_name = cfg.job_name if job_name is None: raise ValueError("Job name must be specified either in config or as a parameter") display_pid = False if not use_threads(cfg): display_pid = True # Create logs directory to ensure it exists log_dir = os.path.join(cfg.output_dir, "logs") os.makedirs(log_dir, exist_ok=True) log_file = os.path.join(log_dir, f"learner_{job_name}.log") # Initialize logging with explicit log file init_logging(log_file=log_file, display_pid=display_pid) logging.info(f"Learner logging initialized, writing to {log_file}") logging.info(pformat(cfg.to_dict())) # Setup WandB logging if enabled if cfg.wandb.enable and cfg.wandb.project: from lerobot.common.wandb_utils import WandBLogger wandb_logger = WandBLogger(cfg) else: wandb_logger = None logging.info(colored("Logs will be saved locally.", "yellow", attrs=["bold"])) # Handle resume logic cfg = handle_resume_logic(cfg) set_seed(seed=cfg.seed) torch.backends.cudnn.benchmark = True torch.backends.cuda.matmul.allow_tf32 = True is_threaded = use_threads(cfg) shutdown_event = ProcessSignalHandler(is_threaded, display_pid=display_pid).shutdown_event start_learner_threads( cfg=cfg, wandb_logger=wandb_logger, shutdown_event=shutdown_event, ) def start_learner_threads( cfg: TrainRLServerPipelineConfig, wandb_logger: WandBLogger | None, shutdown_event: Any, # Event ) -> None: """ Start the learner threads for training. Args: cfg (TrainRLServerPipelineConfig): Training configuration wandb_logger (WandBLogger | None): Logger for metrics shutdown_event: Event to signal shutdown """ # Create multiprocessing queues transition_queue = Queue() interaction_message_queue = Queue() parameters_queue = Queue() concurrency_entity = None if use_threads(cfg): from threading import Thread concurrency_entity = Thread else: from torch.multiprocessing import Process concurrency_entity = Process communication_process = concurrency_entity( target=start_learner, args=( parameters_queue, transition_queue, interaction_message_queue, shutdown_event, cfg, ), daemon=True, ) communication_process.start() try: add_actor_information_and_train( cfg=cfg, wandb_logger=wandb_logger, shutdown_event=shutdown_event, transition_queue=transition_queue, interaction_message_queue=interaction_message_queue, parameters_queue=parameters_queue, ) logging.info("[LEARNER] Training process stopped") except Exception: logging.exception("[LEARNER] Unhandled exception in training loop") shutdown_event.set() finally: logging.info("[LEARNER] Closing queues") transition_queue.close() interaction_message_queue.close() parameters_queue.close() communication_process.join() logging.info("[LEARNER] Communication process joined") transition_queue.cancel_join_thread() interaction_message_queue.cancel_join_thread() parameters_queue.cancel_join_thread() logging.info("[LEARNER] Cleanup complete") # Core algorithm functions def add_actor_information_and_train( cfg: TrainRLServerPipelineConfig, wandb_logger: WandBLogger | None, shutdown_event: Any, # Event transition_queue: Queue, 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. This function continuously: - Transfers transitions from the actor to the replay buffer. - Logs received interaction messages. - Ensures training begins only when the replay buffer has a sufficient number of transitions. - Delegates training updates to an ``RLAlgorithm``. - Periodically pushes updated weights to actors. - Logs training statistics, including loss values and optimization frequency. NOTE: This function doesn't have a single responsibility, it should be split into multiple functions in the future. The reason why we did that is the GIL in Python. It's super slow the performance are divided by 200. So we need to have a single thread that does all the work. Args: cfg (TrainRLServerPipelineConfig): Configuration object containing hyperparameters. wandb_logger (WandBLogger | None): Logger for tracking training progress. shutdown_event (Event): Event to signal shutdown. transition_queue (Queue): Queue for receiving transitions from the actor. interaction_message_queue (Queue): Queue for receiving interaction messages from the actor. parameters_queue (Queue): Queue for sending policy parameters to the actor. """ # Extract all configuration variables at the beginning, it improve the speed performance # of 7% device = get_safe_torch_device(try_device=cfg.policy.device, log=True) storage_device = get_safe_torch_device(try_device=cfg.policy.storage_device) online_step_before_learning = cfg.policy.online_step_before_learning fps = cfg.env.fps log_freq = cfg.log_freq save_freq = cfg.save_freq policy_parameters_push_frequency = cfg.policy.actor_learner_config.policy_parameters_push_frequency saving_checkpoint = cfg.save_checkpoint online_steps = cfg.policy.online_steps # Initialize logging for multiprocessing if not use_threads(cfg): log_dir = os.path.join(cfg.output_dir, "logs") os.makedirs(log_dir, exist_ok=True) log_file = os.path.join(log_dir, f"learner_train_process_{os.getpid()}.log") init_logging(log_file=log_file, display_pid=True) logging.info("Initialized logging for actor information and training process") logging.info("Initializing policy") policy = make_policy( cfg=cfg.policy, env_cfg=cfg.env, ) assert isinstance(policy, nn.Module) policy.train() algorithm = make_algorithm(cfg=cfg.algorithm, policy=policy) preprocessor, postprocessor = make_pre_post_processors( policy_cfg=cfg.policy, dataset_stats=cfg.policy.dataset_stats, ) # Push initial policy weights to actors push_actor_policy_to_queue(parameters_queue=parameters_queue, algorithm=algorithm) last_time_policy_pushed = time.time() log_training_info(cfg=cfg, policy=policy) replay_buffer = initialize_replay_buffer(cfg, device, storage_device) batch_size = cfg.batch_size offline_replay_buffer = None if cfg.dataset is not None: offline_replay_buffer = initialize_offline_replay_buffer( cfg=cfg, device=device, storage_device=storage_device, ) # DataMixer: online-only or online/offline 50-50 mix data_mixer = OnlineOfflineMixer( online_buffer=replay_buffer, offline_buffer=offline_replay_buffer, online_ratio=cfg.online_ratio, ) # RLTrainer owns the iterator, preprocessor, and creates optimizers. trainer = RLTrainer( algorithm=algorithm, data_mixer=data_mixer, batch_size=batch_size, preprocessor=preprocessor, ) # If we are resuming, we need to load the training state optimizers = algorithm.get_optimizers() resume_optimization_step, resume_interaction_step = load_training_state( cfg=cfg, optimizers=optimizers, algorithm=algorithm, device=device ) logging.info("Starting learner thread") interaction_message = None optimization_step = resume_optimization_step if resume_optimization_step is not None else 0 algorithm.optimization_step = optimization_step interaction_step_shift = resume_interaction_step if resume_interaction_step is not None else 0 dataset_repo_id = None if cfg.dataset is not None: dataset_repo_id = cfg.dataset.repo_id # NOTE: THIS IS THE MAIN LOOP OF THE LEARNER while True: # Exit the training loop if shutdown is requested if shutdown_event is not None and shutdown_event.is_set(): logging.info("[LEARNER] Shutdown signal received. Exiting...") break # Process all available transitions to the replay buffer, send by the actor server process_transitions( transition_queue=transition_queue, replay_buffer=replay_buffer, offline_replay_buffer=offline_replay_buffer, dataset_repo_id=dataset_repo_id, shutdown_event=shutdown_event, ) # Process all available interaction messages sent by the actor server interaction_message = process_interaction_messages( interaction_message_queue=interaction_message_queue, interaction_step_shift=interaction_step_shift, wandb_logger=wandb_logger, shutdown_event=shutdown_event, ) # Wait until the replay buffer has enough samples to start training if len(replay_buffer) < online_step_before_learning: continue time_for_one_optimization_step = time.time() # One training step (trainer owns data_mixer iterator; algorithm owns UTD loop) stats = trainer.training_step() # Push policy to actors if needed if time.time() - last_time_policy_pushed > policy_parameters_push_frequency: push_actor_policy_to_queue(parameters_queue=parameters_queue, algorithm=algorithm) last_time_policy_pushed = time.time() training_infos = stats.to_log_dict() # Log training metrics at specified intervals optimization_step = algorithm.optimization_step if optimization_step % log_freq == 0: training_infos["replay_buffer_size"] = len(replay_buffer) if offline_replay_buffer is not None: training_infos["offline_replay_buffer_size"] = len(offline_replay_buffer) training_infos["Optimization step"] = optimization_step # Log training metrics if wandb_logger: wandb_logger.log_dict(d=training_infos, mode="train", custom_step_key="Optimization step") # Calculate and log optimization frequency time_for_one_optimization_step = time.time() - time_for_one_optimization_step frequency_for_one_optimization_step = 1 / (time_for_one_optimization_step + 1e-9) logging.info(f"[LEARNER] Optimization frequency loop [Hz]: {frequency_for_one_optimization_step}") # Log optimization frequency if wandb_logger: wandb_logger.log_dict( { "Optimization frequency loop [Hz]": frequency_for_one_optimization_step, "Optimization step": optimization_step, }, mode="train", custom_step_key="Optimization step", ) if optimization_step % log_freq == 0: logging.info(f"[LEARNER] Number of optimization step: {optimization_step}") # Save checkpoint at specified intervals if saving_checkpoint and (optimization_step % save_freq == 0 or optimization_step == online_steps): save_training_checkpoint( cfg=cfg, optimization_step=optimization_step, online_steps=online_steps, interaction_message=interaction_message, policy=policy, optimizers=optimizers, replay_buffer=replay_buffer, algorithm=algorithm, offline_replay_buffer=offline_replay_buffer, dataset_repo_id=dataset_repo_id, fps=fps, preprocessor=preprocessor, postprocessor=postprocessor, ) def start_learner( parameters_queue: Queue, transition_queue: Queue, interaction_message_queue: Queue, 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. 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 """ if not use_threads(cfg): # Create a process-specific log file log_dir = os.path.join(cfg.output_dir, "logs") os.makedirs(log_dir, exist_ok=True) log_file = os.path.join(log_dir, f"learner_process_{os.getpid()}.log") # Initialize logging with explicit log file init_logging(log_file=log_file, display_pid=True) logging.info("Learner server process logging initialized") # Setup process handlers to handle shutdown signal # But use shutdown event from the main process # Return back for MP # TODO: Check if its useful _ = ProcessSignalHandler(False, display_pid=True) service = LearnerService( shutdown_event=shutdown_event, parameters_queue=parameters_queue, seconds_between_pushes=cfg.policy.actor_learner_config.policy_parameters_push_frequency, transition_queue=transition_queue, interaction_message_queue=interaction_message_queue, queue_get_timeout=cfg.policy.actor_learner_config.queue_get_timeout, ) server = grpc.server( ThreadPoolExecutor(max_workers=MAX_WORKERS), options=[ ("grpc.max_receive_message_length", MAX_MESSAGE_SIZE), ("grpc.max_send_message_length", MAX_MESSAGE_SIZE), ], ) services_pb2_grpc.add_LearnerServiceServicer_to_server( service, server, ) host = cfg.policy.actor_learner_config.learner_host port = cfg.policy.actor_learner_config.learner_port server.add_insecure_port(f"{host}:{port}") server.start() logging.info("[LEARNER] gRPC server started") shutdown_event.wait() logging.info("[LEARNER] Stopping gRPC server...") server.stop(SHUTDOWN_TIMEOUT) logging.info("[LEARNER] gRPC server stopped") def save_training_checkpoint( cfg: TrainRLServerPipelineConfig, optimization_step: int, online_steps: int, interaction_message: dict | None, policy: nn.Module, optimizers: dict[str, Optimizer], replay_buffer: ReplayBuffer, algorithm: RLAlgorithm | None = None, offline_replay_buffer: ReplayBuffer | None = None, dataset_repo_id: str | None = None, fps: int = 30, preprocessor=None, postprocessor=None, ) -> None: """ Save training checkpoint and associated data. This function performs the following steps: 1. Creates a checkpoint directory with the current optimization step 2. Saves the policy model, configuration, and optimizer states 3. Saves the current interaction step for resuming training 4. Updates the "last" checkpoint symlink to point to this checkpoint 5. Saves the replay buffer as a dataset for later use 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 """ logging.info(f"Checkpoint policy after step {optimization_step}") _num_digits = max(6, len(str(online_steps))) interaction_step = interaction_message["Interaction step"] if interaction_message is not None else 0 # Create checkpoint directory checkpoint_dir = get_step_checkpoint_dir(cfg.output_dir, online_steps, optimization_step) # Save policy artifacts (pretrained_model/) + Trainer scaffolding (training_state/). save_checkpoint( checkpoint_dir=checkpoint_dir, step=optimization_step, cfg=cfg, policy=policy, optimizer=optimizers, scheduler=None, preprocessor=preprocessor, postprocessor=postprocessor, ) # Algorithm-owned tensors live in their own component subfolder # so they can be `push_to_hub`'d independently and don't bloat the inference artifact. if algorithm is not None: algorithm.save_pretrained(checkpoint_dir / ALGORITHM_DIR) # Enrich training_step.json with the RL-specific interaction_step counter so # both can be restored from a single file. training_state_dir = checkpoint_dir / TRAINING_STATE_DIR write_json( {"step": optimization_step, "interaction_step": interaction_step}, training_state_dir / TRAINING_STEP, ) # Update the "last" symlink update_last_checkpoint(checkpoint_dir) # TODO : temporary save replay buffer here, remove later when on the robot # We want to control this with the keyboard inputs dataset_dir = os.path.join(cfg.output_dir, "dataset") if os.path.exists(dataset_dir) and os.path.isdir(dataset_dir): shutil.rmtree(dataset_dir) # Save dataset # NOTE: Handle the case where the dataset repo id is not specified in the config # eg. RL training without demonstrations data repo_id_buffer_save = cfg.env.task if dataset_repo_id is None else dataset_repo_id replay_buffer.to_lerobot_dataset(repo_id=repo_id_buffer_save, fps=fps, root=dataset_dir) if offline_replay_buffer is not None: dataset_offline_dir = os.path.join(cfg.output_dir, "dataset_offline") if os.path.exists(dataset_offline_dir) and os.path.isdir(dataset_offline_dir): shutil.rmtree(dataset_offline_dir) offline_replay_buffer.to_lerobot_dataset( cfg.dataset.repo_id, fps=fps, root=dataset_offline_dir, ) logging.info("Resume training") # Training setup functions def handle_resume_logic(cfg: TrainRLServerPipelineConfig) -> TrainRLServerPipelineConfig: """ Handle the resume logic for training. If resume is True: - Verifies that a checkpoint exists - Loads the checkpoint configuration - Logs resumption details - Returns the checkpoint configuration If resume is False: - Checks if an output directory exists (to prevent accidental overwriting) - Returns the original configuration Args: cfg (TrainRLServerPipelineConfig): The training configuration Returns: TrainRLServerPipelineConfig: The updated configuration Raises: RuntimeError: If resume is True but no checkpoint found, or if resume is False but directory exists """ out_dir = cfg.output_dir # Case 1: Not resuming, but need to check if directory exists to prevent overwrites if not cfg.resume: checkpoint_dir = os.path.join(out_dir, CHECKPOINTS_DIR, LAST_CHECKPOINT_LINK) if os.path.exists(checkpoint_dir): raise RuntimeError( f"Output directory {checkpoint_dir} already exists. Use `resume=true` to resume training." ) return cfg # Case 2: Resuming training checkpoint_dir = os.path.join(out_dir, CHECKPOINTS_DIR, LAST_CHECKPOINT_LINK) if not os.path.exists(checkpoint_dir): raise RuntimeError(f"No model checkpoint found in {checkpoint_dir} for resume=True") # Log that we found a valid checkpoint and are resuming logging.info( colored( "Valid checkpoint found: resume=True detected, resuming previous run", color="yellow", attrs=["bold"], ) ) # Load config using Draccus checkpoint_cfg_path = os.path.join(checkpoint_dir, PRETRAINED_MODEL_DIR, "train_config.json") checkpoint_cfg = TrainRLServerPipelineConfig.from_pretrained(checkpoint_cfg_path) # Ensure resume flag is set in returned config checkpoint_cfg.resume = True return checkpoint_cfg def load_training_state( cfg: TrainRLServerPipelineConfig, optimizers: Optimizer | dict[str, Optimizer], 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. Args: cfg: Training configuration. optimizers: Optimizers to load state into. algorithm: Algorithm whose state dict should be restored. Required for full main-equivalent resume; the policy itself is restored separately via ``make_policy``. device: Device on which to place loaded algorithm tensors. Returns: tuple: (optimization_step, interaction_step) or (None, None) if not resuming """ if not cfg.resume: return None, None # Construct path to the last checkpoint directory checkpoint_dir = Path(cfg.output_dir) / CHECKPOINTS_DIR / LAST_CHECKPOINT_LINK logging.info(f"Loading training state from {checkpoint_dir}") try: # Restore optimizers + RNG + step from the standard `training_state/` folder step, optimizers, _ = utils_load_training_state(checkpoint_dir, optimizers, None) # Restore algorithm-owned tensors if algorithm is not None: algo_dir = checkpoint_dir / ALGORITHM_DIR if algo_dir.is_dir(): tensors = load_safetensors(str(algo_dir / SAFETENSORS_SINGLE_FILE)) algorithm.load_state_dict(tensors, device=device) logging.info(f"Loaded algorithm state from {algo_dir}") else: logging.warning( f"No algorithm state found at {algo_dir}; " "will keep their freshly-initialised values. Adam moments restored from the " "old optimizer state may not match these reset parameters." ) # Read interaction_step from the enriched training_step.json training_step_path = checkpoint_dir / TRAINING_STATE_DIR / TRAINING_STEP interaction_step = int(load_json(training_step_path).get("interaction_step", 0)) logging.info(f"Resuming from step {step}, interaction step {interaction_step}") return step, interaction_step except Exception as e: logging.error(f"Failed to load training state: {e}") return None, None def log_training_info(cfg: TrainRLServerPipelineConfig, policy: nn.Module) -> None: """ Log information about the training process. Args: cfg (TrainRLServerPipelineConfig): Training configuration policy (nn.Module): Policy model """ num_learnable_params = sum(p.numel() for p in policy.parameters() if p.requires_grad) num_total_params = sum(p.numel() for p in policy.parameters()) logging.info(colored("Output dir:", "yellow", attrs=["bold"]) + f" {cfg.output_dir}") logging.info(f"{cfg.env.task=}") logging.info(f"{cfg.policy.online_steps=}") logging.info(f"{num_learnable_params=} ({format_big_number(num_learnable_params)})") logging.info(f"{num_total_params=} ({format_big_number(num_total_params)})") def initialize_replay_buffer( cfg: TrainRLServerPipelineConfig, device: str, storage_device: str ) -> ReplayBuffer: """ Initialize a replay buffer, either empty or from a dataset if resuming. Args: cfg (TrainRLServerPipelineConfig): Training configuration device (str): Device to store tensors on storage_device (str): Device for storage optimization Returns: ReplayBuffer: Initialized replay buffer """ if not cfg.resume: return ReplayBuffer( capacity=cfg.policy.online_buffer_capacity, device=device, state_keys=cfg.policy.input_features.keys(), storage_device=storage_device, optimize_memory=True, ) logging.info("Resume training load the online dataset") dataset_path = os.path.join(cfg.output_dir, "dataset") # NOTE: In RL is possible to not have a dataset. repo_id = None if cfg.dataset is not None: repo_id = cfg.dataset.repo_id dataset = LeRobotDataset( repo_id=repo_id, root=dataset_path, ) return ReplayBuffer.from_lerobot_dataset( lerobot_dataset=dataset, capacity=cfg.policy.online_buffer_capacity, device=device, state_keys=cfg.policy.input_features.keys(), optimize_memory=True, ) def initialize_offline_replay_buffer( cfg: TrainRLServerPipelineConfig, device: str, storage_device: str, ) -> ReplayBuffer: """ Initialize an offline replay buffer from a dataset. Args: cfg (TrainRLServerPipelineConfig): Training configuration device (str): Device to store tensors on storage_device (str): Device for storage optimization Returns: ReplayBuffer: Initialized offline replay buffer """ if not cfg.resume: logging.info("make_dataset offline buffer") offline_dataset = make_dataset(cfg) else: logging.info("load offline dataset") dataset_offline_path = os.path.join(cfg.output_dir, "dataset_offline") offline_dataset = LeRobotDataset( repo_id=cfg.dataset.repo_id, root=dataset_offline_path, ) logging.info("Convert to a offline replay buffer") offline_replay_buffer = ReplayBuffer.from_lerobot_dataset( offline_dataset, device=device, state_keys=cfg.policy.input_features.keys(), storage_device=storage_device, optimize_memory=True, capacity=cfg.policy.offline_buffer_capacity, ) return offline_replay_buffer # Utilities/Helpers functions def use_threads(cfg: TrainRLServerPipelineConfig) -> bool: return cfg.policy.concurrency.learner == "threads" def check_nan_in_transition( observations: torch.Tensor, actions: torch.Tensor, next_state: torch.Tensor, raise_error: bool = False, ) -> bool: """ 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 Returns: bool: True if NaN values were detected, False otherwise """ nan_detected = False # Check observations for key, tensor in observations.items(): if torch.isnan(tensor).any(): logging.error(f"observations[{key}] contains NaN values") nan_detected = True if raise_error: raise ValueError(f"NaN detected in observations[{key}]") # Check next state for key, tensor in next_state.items(): if torch.isnan(tensor).any(): logging.error(f"next_state[{key}] contains NaN values") nan_detected = True if raise_error: raise ValueError(f"NaN detected in next_state[{key}]") # Check actions if torch.isnan(actions).any(): logging.error("actions contains NaN values") nan_detected = True if raise_error: raise ValueError("NaN detected in actions") return nan_detected def push_actor_policy_to_queue(parameters_queue: Queue, algorithm: RLAlgorithm) -> None: logging.debug("[LEARNER] Pushing actor policy to the queue") # Create a dictionary to hold all the state dicts state_dicts = algorithm.get_weights() state_bytes = state_to_bytes(state_dicts) parameters_queue.put(state_bytes) def process_interaction_message( message, interaction_step_shift: int, wandb_logger: WandBLogger | None = None ): """Process a single interaction message with consistent handling.""" message = bytes_to_python_object(message) # Shift interaction step for consistency with checkpointed state message["Interaction step"] += interaction_step_shift # Log if logger available if wandb_logger: wandb_logger.log_dict(d=message, mode="train", custom_step_key="Interaction step") return message def process_transitions( transition_queue: Queue, replay_buffer: ReplayBuffer, offline_replay_buffer: ReplayBuffer, dataset_repo_id: str | None, shutdown_event: Any, # Event ): """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 """ while not transition_queue.empty() and not shutdown_event.is_set(): transition_list = transition_queue.get() transition_list = bytes_to_transitions(buffer=transition_list) for transition in transition_list: # Skip transitions with NaN values if check_nan_in_transition( observations=transition["state"], actions=transition[ACTION], next_state=transition["next_state"], ): logging.warning("[LEARNER] NaN detected in transition, skipping") continue replay_buffer.add(**transition) # Add to offline buffer if it's an intervention if dataset_repo_id is not None and transition.get("complementary_info", {}).get( TeleopEvents.IS_INTERVENTION.value ): offline_replay_buffer.add(**transition) def process_interaction_messages( interaction_message_queue: Queue, interaction_step_shift: int, wandb_logger: WandBLogger | None, shutdown_event: Any, # Event ) -> dict | None: """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 Returns: dict | None: The last interaction message processed, or None if none were processed """ last_message = None while not interaction_message_queue.empty() and not shutdown_event.is_set(): message = interaction_message_queue.get() last_message = process_interaction_message( message=message, interaction_step_shift=interaction_step_shift, wandb_logger=wandb_logger, ) return last_message if __name__ == "__main__": train_cli() logging.info("[LEARNER] main finished")