feat(rl): add multiprocessing_context option + check to rl training pipeline

This commit is contained in:
Steven Palma
2026-07-24 13:34:58 +02:00
parent e66da375c8
commit 21ba0ffd8e
5 changed files with 40 additions and 9 deletions
@@ -37,13 +37,19 @@ def is_image_feature(key: str) -> bool:
@dataclass @dataclass
class ConcurrencyConfig: class ConcurrencyConfig:
"""Configuration for the concurrency of the actor and learner. """Configuration for the concurrency of the actor and learner.
Possible values are: Possible values are:
- "threads": Use threads for the actor and learner. - "threads": Use threads for the actor and learner.
- "processes": Use processes 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" actor: str = "threads"
learner: str = "threads" learner: str = "threads"
multiprocessing_context: str | None = "spawn"
@dataclass @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 import gamepad, so_leader # noqa: F401
from lerobot.teleoperators.utils import TeleopEvents from lerobot.teleoperators.utils import TeleopEvents
from lerobot.utils.device_utils import get_safe_torch_device 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.random_utils import set_seed
from lerobot.utils.robot_utils import precise_sleep from lerobot.utils.robot_utils import precise_sleep
from lerobot.utils.transition import ( from lerobot.utils.transition import (
@@ -124,9 +124,7 @@ def actor_cli(cfg: TrainRLServerPipelineConfig):
cfg.validate() cfg.validate()
display_pid = False display_pid = False
if not use_threads(cfg): if not use_threads(cfg):
import torch.multiprocessing as mp ensure_multiprocessing_start_method(cfg.policy.concurrency.multiprocessing_context)
mp.set_start_method("spawn")
display_pid = True display_pid = True
# Create logs directory to ensure it exists # Create logs directory to ensure it exists
+2 -5
View File
@@ -102,7 +102,7 @@ from lerobot.utils.constants import (
) )
from lerobot.utils.device_utils import get_safe_torch_device from lerobot.utils.device_utils import get_safe_torch_device
from lerobot.utils.io_utils import load_json, write_json 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.random_utils import set_seed
from lerobot.utils.utils import ( from lerobot.utils.utils import (
format_big_number, 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. # Fail fast with a friendly error if the optional ``hilserl`` extra is missing.
require_package("grpcio", extra="hilserl", import_name="grpc") require_package("grpcio", extra="hilserl", import_name="grpc")
if not use_threads(cfg): if not use_threads(cfg):
import torch.multiprocessing as mp ensure_multiprocessing_start_method(cfg.policy.concurrency.multiprocessing_context)
if mp.get_start_method(allow_none=True) is None:
mp.set_start_method("spawn")
# Use the job_name from the config # Use the job_name from the config
train( train(
+28
View File
@@ -16,11 +16,39 @@
# limitations under the License. # limitations under the License.
import logging import logging
import multiprocessing
import os import os
import signal import signal
import sys 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: class ProcessSignalHandler:
"""Utility class to attach graceful shutdown signal handlers. """Utility class to attach graceful shutdown signal handlers.
@@ -113,6 +113,7 @@ def test_gaussian_actor_config_default_initialization():
# Concurrency configuration # Concurrency configuration
assert config.concurrency.actor == "threads" assert config.concurrency.actor == "threads"
assert config.concurrency.learner == "threads" assert config.concurrency.learner == "threads"
assert config.concurrency.multiprocessing_context == "spawn"
assert isinstance(config.actor_network_kwargs, ActorNetworkConfig) assert isinstance(config.actor_network_kwargs, ActorNetworkConfig)
assert isinstance(config.policy_kwargs, PolicyConfig) assert isinstance(config.policy_kwargs, PolicyConfig)
@@ -152,6 +153,7 @@ def test_concurrency_config():
config = ConcurrencyConfig() config = ConcurrencyConfig()
assert config.actor == "threads" assert config.actor == "threads"
assert config.learner == "threads" assert config.learner == "threads"
assert config.multiprocessing_context == "spawn"
def test_gaussian_actor_config_custom_initialization(): def test_gaussian_actor_config_custom_initialization():