From 21ba0ffd8ebd6942dde077bf6d4ea4a4f7b5226a Mon Sep 17 00:00:00 2001 From: Steven Palma Date: Fri, 24 Jul 2026 13:34:58 +0200 Subject: [PATCH] feat(rl): add multiprocessing_context option + check to rl training pipeline --- .../configuration_gaussian_actor.py | 6 ++++ src/lerobot/rl/actor.py | 6 ++-- src/lerobot/rl/learner.py | 7 ++--- src/lerobot/utils/process.py | 28 +++++++++++++++++++ tests/policies/test_gaussian_actor_config.py | 2 ++ 5 files changed, 40 insertions(+), 9 deletions(-) diff --git a/src/lerobot/policies/gaussian_actor/configuration_gaussian_actor.py b/src/lerobot/policies/gaussian_actor/configuration_gaussian_actor.py index e51653992..cfc926007 100644 --- a/src/lerobot/policies/gaussian_actor/configuration_gaussian_actor.py +++ b/src/lerobot/policies/gaussian_actor/configuration_gaussian_actor.py @@ -37,13 +37,19 @@ def is_image_feature(key: str) -> bool: @dataclass class ConcurrencyConfig: """Configuration for the concurrency of the actor and learner. + Possible values are: - "threads": Use threads for the actor and learner. - "processes": Use processes for the actor and learner. + + ``multiprocessing_context`` selects the process-wide start method when + processes are used. Set it to ``None`` to preserve Python's default or a + method already selected by the embedding application. """ actor: str = "threads" learner: str = "threads" + multiprocessing_context: str | None = "spawn" @dataclass diff --git a/src/lerobot/rl/actor.py b/src/lerobot/rl/actor.py index bfc7f1882..df703cc7b 100644 --- a/src/lerobot/rl/actor.py +++ b/src/lerobot/rl/actor.py @@ -91,7 +91,7 @@ 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.utils.device_utils import get_safe_torch_device -from lerobot.utils.process import ProcessSignalHandler +from lerobot.utils.process import ProcessSignalHandler, ensure_multiprocessing_start_method from lerobot.utils.random_utils import set_seed from lerobot.utils.robot_utils import precise_sleep from lerobot.utils.transition import ( @@ -124,9 +124,7 @@ def actor_cli(cfg: TrainRLServerPipelineConfig): cfg.validate() display_pid = False if not use_threads(cfg): - import torch.multiprocessing as mp - - mp.set_start_method("spawn") + ensure_multiprocessing_start_method(cfg.policy.concurrency.multiprocessing_context) display_pid = True # Create logs directory to ensure it exists diff --git a/src/lerobot/rl/learner.py b/src/lerobot/rl/learner.py index 0440d569f..ee95d831e 100644 --- a/src/lerobot/rl/learner.py +++ b/src/lerobot/rl/learner.py @@ -102,7 +102,7 @@ from lerobot.utils.constants import ( ) 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.process import ProcessSignalHandler, ensure_multiprocessing_start_method from lerobot.utils.random_utils import set_seed from lerobot.utils.utils import ( format_big_number, @@ -123,10 +123,7 @@ 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 - - if mp.get_start_method(allow_none=True) is None: - mp.set_start_method("spawn") + ensure_multiprocessing_start_method(cfg.policy.concurrency.multiprocessing_context) # Use the job_name from the config train( diff --git a/src/lerobot/utils/process.py b/src/lerobot/utils/process.py index 72438b6f9..d855ad8ab 100644 --- a/src/lerobot/utils/process.py +++ b/src/lerobot/utils/process.py @@ -16,11 +16,39 @@ # limitations under the License. import logging +import multiprocessing import os import signal import sys +def ensure_multiprocessing_start_method(start_method: str | None) -> None: + """Set a multiprocessing start method once, or verify the existing method matches. + + Passing ``None`` leaves Python's process-wide default untouched. This is useful + when LeRobot is embedded in an application that owns multiprocessing setup. + """ + if start_method is None: + return + + available_methods = multiprocessing.get_all_start_methods() + if start_method not in available_methods: + raise ValueError( + f"Multiprocessing start method must be one of {available_methods} on this platform, " + f"got {start_method!r}." + ) + + current_method = multiprocessing.get_start_method(allow_none=True) + if current_method is None: + multiprocessing.set_start_method(start_method) + elif current_method != start_method: + raise RuntimeError( + f"Multiprocessing start method is already {current_method!r}; cannot change it to " + f"{start_method!r}. Set the configured multiprocessing context to null to keep the " + "application's existing method, or launch LeRobot in a fresh process." + ) + + class ProcessSignalHandler: """Utility class to attach graceful shutdown signal handlers. diff --git a/tests/policies/test_gaussian_actor_config.py b/tests/policies/test_gaussian_actor_config.py index 004612374..511432441 100644 --- a/tests/policies/test_gaussian_actor_config.py +++ b/tests/policies/test_gaussian_actor_config.py @@ -113,6 +113,7 @@ def test_gaussian_actor_config_default_initialization(): # Concurrency configuration assert config.concurrency.actor == "threads" assert config.concurrency.learner == "threads" + assert config.concurrency.multiprocessing_context == "spawn" assert isinstance(config.actor_network_kwargs, ActorNetworkConfig) assert isinstance(config.policy_kwargs, PolicyConfig) @@ -152,6 +153,7 @@ def test_concurrency_config(): config = ConcurrencyConfig() assert config.actor == "threads" assert config.learner == "threads" + assert config.multiprocessing_context == "spawn" def test_gaussian_actor_config_custom_initialization():