Compare commits

...

2 Commits

Author SHA1 Message Date
Steven Palma 21ba0ffd8e feat(rl): add multiprocessing_context option + check to rl training pipeline 2026-07-24 13:34:58 +02:00
Jash Shah e66da375c8 fix(rl): only set multiprocessing start method if not already set
mp.set_start_method() raises RuntimeError if called more than once
(e.g., when train_cli is invoked in a process that already set the
start method via pytest, Jupyter, or another entrypoint). Guard with
get_start_method(allow_none=True) check.
2026-07-24 13:28:08 +02:00
5 changed files with 40 additions and 8 deletions
@@ -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
+2 -4
View File
@@ -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
+2 -4
View File
@@ -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,9 +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
mp.set_start_method("spawn")
ensure_multiprocessing_start_method(cfg.policy.concurrency.multiprocessing_context)
# Use the job_name from the config
train(
+28
View File
@@ -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.
@@ -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():