83 lines
2.7 KiB
Python
83 lines
2.7 KiB
Python
import hilserl_pb2 # type: ignore
|
|
import hilserl_pb2_grpc # type: ignore
|
|
import logging
|
|
from multiprocessing import Event, Queue
|
|
|
|
from lerobot.scripts.server.network_utils import receive_bytes_in_chunks
|
|
from lerobot.scripts.server.network_utils import send_bytes_in_chunks
|
|
|
|
MAX_MESSAGE_SIZE = 4 * 1024 * 1024 # 4 MB
|
|
MAX_WORKERS = 3 # Stream parameters, send transitions and interactions
|
|
STUTDOWN_TIMEOUT = 10
|
|
|
|
|
|
class LearnerService(hilserl_pb2_grpc.LearnerServiceServicer):
|
|
def __init__(
|
|
self,
|
|
shutdown_event: Event,
|
|
parameters_queue: Queue,
|
|
seconds_between_pushes: float,
|
|
transition_queue: Queue,
|
|
interaction_message_queue: Queue,
|
|
):
|
|
self.shutdown_event = shutdown_event
|
|
self.parameters_queue = parameters_queue
|
|
self.seconds_between_pushes = seconds_between_pushes
|
|
self.transition_queue = transition_queue
|
|
self.interaction_message_queue = interaction_message_queue
|
|
|
|
def StreamParameters(self, request, context):
|
|
# TODO: authorize the request
|
|
logging.info("[LEARNER] Received request to stream parameters from the Actor")
|
|
|
|
while not self.shutdown_event.is_set():
|
|
logging.info("[LEARNER] Push parameters to the Actor")
|
|
buffer = self.parameters_queue.get()
|
|
|
|
yield from send_bytes_in_chunks(
|
|
buffer,
|
|
hilserl_pb2.Parameters,
|
|
log_prefix="[LEARNER] Sending parameters",
|
|
silent=True,
|
|
)
|
|
|
|
logging.info("[LEARNER] Parameters sent")
|
|
|
|
self.shutdown_event.wait(self.seconds_between_pushes)
|
|
|
|
logging.info("[LEARNER] Stream parameters finished")
|
|
return hilserl_pb2.Empty()
|
|
|
|
def SendTransitions(self, request_iterator, _context):
|
|
# TODO: authorize the request
|
|
logging.info("[LEARNER] Received request to receive transitions from the Actor")
|
|
|
|
receive_bytes_in_chunks(
|
|
request_iterator,
|
|
self.transition_queue,
|
|
self.shutdown_event,
|
|
log_prefix="[LEARNER] transitions",
|
|
)
|
|
|
|
logging.debug("[LEARNER] Finished receiving transitions")
|
|
return hilserl_pb2.Empty()
|
|
|
|
def SendInteractions(self, request_iterator, _context):
|
|
# TODO: authorize the request
|
|
logging.info(
|
|
"[LEARNER] Received request to receive interactions from the Actor"
|
|
)
|
|
|
|
receive_bytes_in_chunks(
|
|
request_iterator,
|
|
self.interaction_message_queue,
|
|
self.shutdown_event,
|
|
log_prefix="[LEARNER] interactions",
|
|
)
|
|
|
|
logging.debug("[LEARNER] Finished receiving interactions")
|
|
return hilserl_pb2.Empty()
|
|
|
|
def Ready(self, request, context):
|
|
return hilserl_pb2.Empty()
|