diff --git a/src/lerobot/robots/lekiwi/lekiwi_host.py b/src/lerobot/robots/lekiwi/lekiwi_host.py index 1155cf71c..ae7ac58c4 100644 --- a/src/lerobot/robots/lekiwi/lekiwi_host.py +++ b/src/lerobot/robots/lekiwi/lekiwi_host.py @@ -18,14 +18,24 @@ import base64 import json import logging import time +from dataclasses import dataclass, field import cv2 +import draccus import zmq from .config_lekiwi import LeKiwiConfig, LeKiwiHostConfig from .lekiwi import LeKiwi +@dataclass +class LeKiwiServerConfig: + """Configuration for the LeKiwi host script.""" + + robot: LeKiwiConfig = field(default_factory=LeKiwiConfig) + host: LeKiwiHostConfig = field(default_factory=LeKiwiHostConfig) + + class LeKiwiHost: def __init__(self, config: LeKiwiHostConfig): self.zmq_context = zmq.Context() @@ -47,17 +57,16 @@ class LeKiwiHost: self.zmq_context.term() -def main(): +@draccus.wrap() +def main(cfg: LeKiwiServerConfig): logging.info("Configuring LeKiwi") - robot_config = LeKiwiConfig() - robot = LeKiwi(robot_config) + robot = LeKiwi(cfg.robot) logging.info("Connecting LeKiwi") robot.connect() logging.info("Starting HostAgent") - host_config = LeKiwiHostConfig() - host = LeKiwiHost(host_config) + host = LeKiwiHost(cfg.host) last_cmd_time = time.time() watchdog_active = False