mirror of
https://github.com/huggingface/lerobot.git
synced 2026-07-19 07:51:43 +00:00
refactor(rl): update shutdown_event type hints from 'any' to 'Any' for consistency and clarity
This commit is contained in:
+14
-9
@@ -51,11 +51,12 @@ import os
|
|||||||
import time
|
import time
|
||||||
from functools import lru_cache
|
from functools import lru_cache
|
||||||
from queue import Empty
|
from queue import Empty
|
||||||
|
from typing import Any
|
||||||
|
|
||||||
import grpc
|
import grpc
|
||||||
import torch
|
import torch
|
||||||
from torch import nn
|
from torch import nn
|
||||||
from torch.multiprocessing import Event, Queue
|
from torch.multiprocessing import Queue
|
||||||
|
|
||||||
from lerobot.cameras import opencv # noqa: F401
|
from lerobot.cameras import opencv # noqa: F401
|
||||||
from lerobot.configs import parser
|
from lerobot.configs import parser
|
||||||
@@ -210,7 +211,7 @@ def actor_cli(cfg: TrainRLServerPipelineConfig):
|
|||||||
|
|
||||||
def act_with_policy(
|
def act_with_policy(
|
||||||
cfg: TrainRLServerPipelineConfig,
|
cfg: TrainRLServerPipelineConfig,
|
||||||
shutdown_event: any, # Event,
|
shutdown_event: Any, # Event
|
||||||
parameters_queue: Queue,
|
parameters_queue: Queue,
|
||||||
transitions_queue: Queue,
|
transitions_queue: Queue,
|
||||||
interactions_queue: Queue,
|
interactions_queue: Queue,
|
||||||
@@ -414,7 +415,7 @@ def act_with_policy(
|
|||||||
|
|
||||||
def establish_learner_connection(
|
def establish_learner_connection(
|
||||||
stub: services_pb2_grpc.LearnerServiceStub,
|
stub: services_pb2_grpc.LearnerServiceStub,
|
||||||
shutdown_event: Event, # type: ignore
|
shutdown_event: Any, # Event
|
||||||
attempts: int = 30,
|
attempts: int = 30,
|
||||||
):
|
):
|
||||||
"""Establish a connection with the learner.
|
"""Establish a connection with the learner.
|
||||||
@@ -466,7 +467,7 @@ def learner_service_client(
|
|||||||
def receive_policy(
|
def receive_policy(
|
||||||
cfg: TrainRLServerPipelineConfig,
|
cfg: TrainRLServerPipelineConfig,
|
||||||
parameters_queue: Queue,
|
parameters_queue: Queue,
|
||||||
shutdown_event: Event, # type: ignore
|
shutdown_event: Any, # Event
|
||||||
learner_client: services_pb2_grpc.LearnerServiceStub | None = None,
|
learner_client: services_pb2_grpc.LearnerServiceStub | None = None,
|
||||||
grpc_channel: grpc.Channel | None = None,
|
grpc_channel: grpc.Channel | None = None,
|
||||||
):
|
):
|
||||||
@@ -518,7 +519,7 @@ def receive_policy(
|
|||||||
def send_transitions(
|
def send_transitions(
|
||||||
cfg: TrainRLServerPipelineConfig,
|
cfg: TrainRLServerPipelineConfig,
|
||||||
transitions_queue: Queue,
|
transitions_queue: Queue,
|
||||||
shutdown_event: any, # Event,
|
shutdown_event: Any, # Event
|
||||||
learner_client: services_pb2_grpc.LearnerServiceStub | None = None,
|
learner_client: services_pb2_grpc.LearnerServiceStub | None = None,
|
||||||
grpc_channel: grpc.Channel | None = None,
|
grpc_channel: grpc.Channel | None = None,
|
||||||
) -> services_pb2.Empty:
|
) -> services_pb2.Empty:
|
||||||
@@ -568,7 +569,7 @@ def send_transitions(
|
|||||||
def send_interactions(
|
def send_interactions(
|
||||||
cfg: TrainRLServerPipelineConfig,
|
cfg: TrainRLServerPipelineConfig,
|
||||||
interactions_queue: Queue,
|
interactions_queue: Queue,
|
||||||
shutdown_event: Event, # type: ignore
|
shutdown_event: Any, # Event
|
||||||
learner_client: services_pb2_grpc.LearnerServiceStub | None = None,
|
learner_client: services_pb2_grpc.LearnerServiceStub | None = None,
|
||||||
grpc_channel: grpc.Channel | None = None,
|
grpc_channel: grpc.Channel | None = None,
|
||||||
) -> services_pb2.Empty:
|
) -> services_pb2.Empty:
|
||||||
@@ -618,7 +619,11 @@ def send_interactions(
|
|||||||
logging.info("[ACTOR] Interactions process stopped")
|
logging.info("[ACTOR] Interactions process stopped")
|
||||||
|
|
||||||
|
|
||||||
def transitions_stream(shutdown_event: Event, transitions_queue: Queue, timeout: float) -> services_pb2.Empty: # type: ignore
|
def transitions_stream(
|
||||||
|
shutdown_event: Any, # Event
|
||||||
|
transitions_queue: Queue,
|
||||||
|
timeout: float,
|
||||||
|
) -> services_pb2.Empty:
|
||||||
while not shutdown_event.is_set():
|
while not shutdown_event.is_set():
|
||||||
try:
|
try:
|
||||||
message = transitions_queue.get(block=True, timeout=timeout)
|
message = transitions_queue.get(block=True, timeout=timeout)
|
||||||
@@ -634,9 +639,9 @@ def transitions_stream(shutdown_event: Event, transitions_queue: Queue, timeout:
|
|||||||
|
|
||||||
|
|
||||||
def interactions_stream(
|
def interactions_stream(
|
||||||
shutdown_event: Event,
|
shutdown_event: Any, # Event
|
||||||
interactions_queue: Queue,
|
interactions_queue: Queue,
|
||||||
timeout: float, # type: ignore
|
timeout: float,
|
||||||
) -> services_pb2.Empty:
|
) -> services_pb2.Empty:
|
||||||
while not shutdown_event.is_set():
|
while not shutdown_event.is_set():
|
||||||
try:
|
try:
|
||||||
|
|||||||
@@ -51,6 +51,7 @@ import time
|
|||||||
from concurrent.futures import ThreadPoolExecutor
|
from concurrent.futures import ThreadPoolExecutor
|
||||||
from pathlib import Path
|
from pathlib import Path
|
||||||
from pprint import pformat
|
from pprint import pformat
|
||||||
|
from typing import Any
|
||||||
|
|
||||||
import grpc
|
import grpc
|
||||||
import torch
|
import torch
|
||||||
@@ -181,7 +182,7 @@ def train(cfg: TrainRLServerPipelineConfig, job_name: str | None = None):
|
|||||||
def start_learner_threads(
|
def start_learner_threads(
|
||||||
cfg: TrainRLServerPipelineConfig,
|
cfg: TrainRLServerPipelineConfig,
|
||||||
wandb_logger: WandBLogger | None,
|
wandb_logger: WandBLogger | None,
|
||||||
shutdown_event: any, # Event,
|
shutdown_event: Any, # Event
|
||||||
) -> None:
|
) -> None:
|
||||||
"""
|
"""
|
||||||
Start the learner threads for training.
|
Start the learner threads for training.
|
||||||
@@ -255,7 +256,7 @@ def start_learner_threads(
|
|||||||
def add_actor_information_and_train(
|
def add_actor_information_and_train(
|
||||||
cfg: TrainRLServerPipelineConfig,
|
cfg: TrainRLServerPipelineConfig,
|
||||||
wandb_logger: WandBLogger | None,
|
wandb_logger: WandBLogger | None,
|
||||||
shutdown_event: any, # Event,
|
shutdown_event: Any, # Event
|
||||||
transition_queue: Queue,
|
transition_queue: Queue,
|
||||||
interaction_message_queue: Queue,
|
interaction_message_queue: Queue,
|
||||||
parameters_queue: Queue,
|
parameters_queue: Queue,
|
||||||
@@ -465,7 +466,7 @@ def start_learner(
|
|||||||
parameters_queue: Queue,
|
parameters_queue: Queue,
|
||||||
transition_queue: Queue,
|
transition_queue: Queue,
|
||||||
interaction_message_queue: Queue,
|
interaction_message_queue: Queue,
|
||||||
shutdown_event: any, # Event,
|
shutdown_event: Any, # Event
|
||||||
cfg: TrainRLServerPipelineConfig,
|
cfg: TrainRLServerPipelineConfig,
|
||||||
):
|
):
|
||||||
"""
|
"""
|
||||||
@@ -907,7 +908,7 @@ def process_transitions(
|
|||||||
replay_buffer: ReplayBuffer,
|
replay_buffer: ReplayBuffer,
|
||||||
offline_replay_buffer: ReplayBuffer,
|
offline_replay_buffer: ReplayBuffer,
|
||||||
dataset_repo_id: str | None,
|
dataset_repo_id: str | None,
|
||||||
shutdown_event: any,
|
shutdown_event: Any, # Event
|
||||||
):
|
):
|
||||||
"""Process all available transitions from the queue.
|
"""Process all available transitions from the queue.
|
||||||
|
|
||||||
@@ -945,7 +946,7 @@ def process_interaction_messages(
|
|||||||
interaction_message_queue: Queue,
|
interaction_message_queue: Queue,
|
||||||
interaction_step_shift: int,
|
interaction_step_shift: int,
|
||||||
wandb_logger: WandBLogger | None,
|
wandb_logger: WandBLogger | None,
|
||||||
shutdown_event: any,
|
shutdown_event: Any, # Event
|
||||||
) -> dict | None:
|
) -> dict | None:
|
||||||
"""Process all available interaction messages from the queue.
|
"""Process all available interaction messages from the queue.
|
||||||
|
|
||||||
|
|||||||
Reference in New Issue
Block a user