diff --git a/docs/source/_toctree.yml b/docs/source/_toctree.yml index eb2400816..da1ae3966 100644 --- a/docs/source/_toctree.yml +++ b/docs/source/_toctree.yml @@ -215,4 +215,6 @@ title: Environments - local: api/configs title: Configuration + - local: api/rollout + title: Rollout title: "API Reference" diff --git a/docs/source/api/rollout.mdx b/docs/source/api/rollout.mdx new file mode 100644 index 000000000..aeca4d389 --- /dev/null +++ b/docs/source/api/rollout.mdx @@ -0,0 +1,118 @@ +# Rollout + +`lerobot.rollout` is the policy deployment engine behind [`lerobot-rollout`](../inference): it wires up a +policy, an inference backend, and a pluggable recording strategy, then runs the robot control loop. See +[Policy Deployment (lerobot-rollout)](../inference) for CLI usage and [Real-Time Chunking (RTC)](../rtc) +for the async inference backend. + +## RolloutConfig + +Top-level configuration for the `lerobot-rollout` CLI. + +[[autodoc]] lerobot.rollout.RolloutConfig + +## Strategies + +Strategies implement the control loop and are selected via `--strategy.type=`. + +[[autodoc]] lerobot.rollout.RolloutStrategy + - setup + - run + - teardown + +[[autodoc]] lerobot.rollout.RolloutStrategyConfig + +[[autodoc]] lerobot.rollout.BaseStrategy + - all + +[[autodoc]] lerobot.rollout.BaseStrategyConfig + +[[autodoc]] lerobot.rollout.SentryStrategy + - all + +[[autodoc]] lerobot.rollout.SentryStrategyConfig + +[[autodoc]] lerobot.rollout.HighlightStrategy + - all + +[[autodoc]] lerobot.rollout.HighlightStrategyConfig + +[[autodoc]] lerobot.rollout.EpisodicStrategy + - all + +[[autodoc]] lerobot.rollout.EpisodicStrategyConfig + +[[autodoc]] lerobot.rollout.DAggerStrategy + - all + +[[autodoc]] lerobot.rollout.DAggerStrategyConfig + +[[autodoc]] lerobot.rollout.DAggerKeyboardConfig + +[[autodoc]] lerobot.rollout.DAggerPedalConfig + +## create_strategy + +[[autodoc]] lerobot.rollout.create_strategy + +## Inference backends + +Inference backends produce actions during the control loop and are selected via +`--inference.type=`. + +[[autodoc]] lerobot.rollout.InferenceEngine + - start + - stop + - reset + - get_action + - notify_observation + - pause + - resume + +[[autodoc]] lerobot.rollout.InferenceEngineConfig + +[[autodoc]] lerobot.rollout.SyncInferenceEngine + - all + +[[autodoc]] lerobot.rollout.SyncInferenceConfig + +[[autodoc]] lerobot.rollout.RTCInferenceEngine + - all + +[[autodoc]] lerobot.rollout.RTCInferenceConfig + +## create_inference_engine + +[[autodoc]] lerobot.rollout.create_inference_engine + +## build_rollout_context + +Wires up policy, processors, hardware, dataset, and inference engine before strategy dispatch. + +[[autodoc]] lerobot.rollout.build_rollout_context + +## RolloutContext + +[[autodoc]] lerobot.rollout.RolloutContext + +[[autodoc]] lerobot.rollout.RuntimeContext + +[[autodoc]] lerobot.rollout.HardwareContext + +[[autodoc]] lerobot.rollout.PolicyContext + +[[autodoc]] lerobot.rollout.ProcessorContext + +[[autodoc]] lerobot.rollout.DatasetContext + +## ThreadSafeRobot + +[[autodoc]] lerobot.rollout.robot_wrapper.ThreadSafeRobot + - all + +## RolloutRingBuffer + +Memory-bounded ring buffer used by the Highlight strategy for on-demand recording. + +[[autodoc]] lerobot.rollout.ring_buffer.RolloutRingBuffer + - all diff --git a/pyproject.toml b/pyproject.toml index a350cd785..95483710a 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -452,7 +452,6 @@ ignore = [ "src/lerobot/processor/**" = ["D"] "src/lerobot/rewards/**" = ["D"] "src/lerobot/rl/**" = ["D"] -"src/lerobot/rollout/**" = ["D"] "src/lerobot/scripts/**" = ["D"] "src/lerobot/teleoperators/**" = ["D"] "src/lerobot/transforms/**" = ["D"] @@ -515,7 +514,7 @@ ignore-private = false ignore-property-decorators = false ignore-module = false ignore-setters = false -fail-under = 55 +fail-under = 55.5 output-format = "term-missing" color = true paths = ["src/lerobot"] diff --git a/src/lerobot/rollout/configs.py b/src/lerobot/rollout/configs.py index ad2c12524..210eafc9e 100644 --- a/src/lerobot/rollout/configs.py +++ b/src/lerobot/rollout/configs.py @@ -47,6 +47,7 @@ class RolloutStrategyConfig(draccus.ChoiceRegistry, abc.ABC): @property def type(self) -> str: + """The registered name of this strategy (e.g. `"base"`, `"sentry"`, `"dagger"`).""" return self.get_choice_name(self.__class__) @@ -69,12 +70,17 @@ class SentryStrategyConfig(RolloutStrategyConfig): boundaries with the dataset's video file chunking, so each ``push_to_hub`` call uploads complete video files rather than re-uploading a growing file that hasn't crossed the chunk boundary. + + Args: + upload_every_n_episodes (`int`, *optional*, defaults to 5): + Push the dataset to the Hub after every N saved episodes. + target_video_file_size_mb (`int | None`, *optional*): + Target video file size in MB for episode rotation. Episodes are saved once the estimated + video duration would exceed this limit. Defaults to `DEFAULT_VIDEO_FILE_SIZE_IN_MB` when + `None`. """ upload_every_n_episodes: int = 5 - # Target video file size in MB for episode rotation. Episodes are - # saved once the estimated video duration would exceed this limit. - # Defaults to DEFAULT_VIDEO_FILE_SIZE_IN_MB when set to None. target_video_file_size_mb: int | None = None @@ -87,6 +93,17 @@ class HighlightStrategyConfig(RolloutStrategyConfig): the user presses the save key, the buffer contents are flushed to the dataset and live recording continues until the key is pressed again. + + Args: + ring_buffer_seconds (`float`, *optional*, defaults to 10.0): + Duration, in seconds, of telemetry kept in the ring buffer before it's overwritten. + ring_buffer_max_memory_mb (`int`, *optional*, defaults to 1024): + Hard memory cap, in MiB, for the ring buffer. Frames are evicted early if this is reached + before `ring_buffer_seconds` of telemetry. + save_key (`str`, *optional*, defaults to `"s"`): + Keyboard key that flushes the ring buffer and starts (or ends) live recording. + push_key (`str`, *optional*, defaults to `"h"`): + Keyboard key that requests an on-demand push of the dataset to the Hub. """ ring_buffer_seconds: float = 10.0 @@ -101,6 +118,14 @@ class DAggerKeyboardConfig: Keys are specified as single characters (e.g. ``"c"``, ``"h"``) or special key names (``"space"``). + + Args: + pause_resume (`str`, *optional*, defaults to `"space"`): + Key that toggles policy execution on/off. + correction (`str`, *optional*, defaults to `"tab"`): + Key that toggles human correction recording. + upload (`str`, *optional*, defaults to `"enter"`): + Key that pushes the dataset to the Hub on demand (corrections-only mode). """ pause_resume: str = "space" @@ -113,6 +138,16 @@ class DAggerPedalConfig: """Foot pedal configuration for DAgger controls. Pedal codes are evdev key code strings (e.g. ``"KEY_A"``). + + Args: + device_path (`str`, *optional*, defaults to `"/dev/input/by-id/usb-PCsensor_FootSwitch-event-kbd"`): + evdev device path of the foot pedal. + pause_resume (`str`, *optional*, defaults to `"KEY_A"`): + evdev key code that toggles policy execution on/off. + correction (`str`, *optional*, defaults to `"KEY_B"`): + evdev key code that toggles human correction recording. + upload (`str`, *optional*, defaults to `"KEY_C"`): + evdev key code that pushes the dataset to the Hub on demand (corrections-only mode). """ device_path: str = "/dev/input/by-id/usb-PCsensor_FootSwitch-event-kbd" @@ -137,25 +172,27 @@ class EpisodicStrategyConfig(RolloutStrategyConfig): In between episodes: - if there is no teleop leader, the robot is held at its initial joint positions captured at startup. - else, the robot is moved smoothly to the position of the teleop leader. + + Args: + reset_to_initial_position (`bool`, *optional*, defaults to `True`): + Only applies when there is no teleop leader. When `True`, moves the robot back to the + joint positions captured at startup during the reset phase. Otherwise, leaves the robot in + its current position. + smooth_leader_to_follower_handover (`bool`, *optional*, defaults to `True`): + Whether to turn on or off the leader -> follower smooth handover behavior. When `False`, + falls back to follower -> leader handover. Leader -> follower handover is only supported + when the leader has `send_feedback` capability. + smooth_handover (`bool`, *optional*, defaults to `True`): + Whether to turn on or off the smooth handover behavior at the start of the reset phase: the + leader is driven to the follower position (actuated teleops, see + `smooth_leader_to_follower_handover`), or the follower is slid to the teleop pose + (non-actuated teleops). Disable for clutch-style teleoperators (e.g. VR controllers) that + re-reference at the current robot pose on engage: the handover is already continuous + there, and the blocking interpolation only delays the start of the reset phase. """ - # This only applies if there are no teleop leaders specified. - # When True (default), moves the robot back to the joint positions captured at startup. - # Otherwise, leave the robot in its current position. reset_to_initial_position: bool = True - - # Whether to turn on or off the leader -> follower smooth handover behavior. - # When False, fallback to follower -> leader handover. - # Note that leader -> follower handover is only supported when the leader has `send_feedback` capability. smooth_leader_to_follower_handover: bool = True - - # Whether to turn on or off the smooth handover behavior at the start of the - # reset phase: the leader is driven to the follower position (actuated - # teleops, see `smooth_leader_to_follower_handover`), or the follower is - # slid to the teleop pose (non-actuated teleops). Disable for clutch-style - # teleoperators (e.g. VR controllers) that re-reference at the current robot - # pose on engage: the handover is already continuous there, and the blocking - # interpolation only delays the start of the reset phase. smooth_handover: bool = True @@ -179,29 +216,54 @@ class DAggerStrategyConfig(RolloutStrategyConfig): to record both autonomous and correction frames with size-based episode rotation (same as Sentry) and background uploading. ``push_to_hub`` is blocked while a correction is in progress. + + Args: + num_episodes (`int | None`, *optional*): + Number of correction episodes to collect (corrections-only mode). When `None`, falls back + to `--dataset.num_episodes`. + record_autonomous (`bool`, *optional*, defaults to `False`): + When `False`, only human-correction windows are recorded, each becoming its own episode. + When `True`, both autonomous and correction frames are recorded with size-based episode + rotation (same as Sentry) and background uploading. + upload_every_n_episodes (`int`, *optional*, defaults to 5): + Push the dataset to the Hub after every N saved episodes (`record_autonomous=True` mode). + target_video_file_size_mb (`int | None`, *optional*): + Target video file size in MB for episode rotation (`record_autonomous=True` mode only). + Defaults to `DEFAULT_VIDEO_FILE_SIZE_IN_MB` when `None`. + smooth_handover (`bool`, *optional*, defaults to `True`): + Whether to turn on or off the smooth handover behavior at phase transitions: the leader is + driven to the follower position on pause (teleops with `send_feedback` capability), and + the follower is slid to the teleop pose when a correction starts (non-actuated teleops). + Disable for clutch-style teleoperators (e.g. VR controllers) that re-reference at the + current robot pose on engage: the handover is already continuous there, and the blocking + interpolation only delays the start of the correction. + input_device (`str`, *optional*, defaults to `"keyboard"`): + Input device used for the pause_resume/correction/upload controls. One of `"keyboard"` or + `"pedal"`. + keyboard (`DAggerKeyboardConfig`, *optional*): + Keyboard key bindings, used when `input_device="keyboard"`. + pedal (`DAggerPedalConfig`, *optional*): + Foot pedal configuration, used when `input_device="pedal"`. + + Raises: + ValueError: If `input_device` is not `"keyboard"` or `"pedal"`. """ - # Number of correction episodes to collect (corrections-only mode). - # When None, falls back to ``--dataset.num_episodes``. num_episodes: int | None = None record_autonomous: bool = False upload_every_n_episodes: int = 5 - # Target video file size in MB for episode rotation (record_autonomous - # mode only). Defaults to DEFAULT_VIDEO_FILE_SIZE_IN_MB when None. target_video_file_size_mb: int | None = None - # Whether to turn on or off the smooth handover behavior at phase transitions: - # the leader is driven to the follower position on pause (teleops with - # `send_feedback` capability), and the follower is slid to the teleop pose when - # a correction starts (non-actuated teleops). Disable for clutch-style - # teleoperators (e.g. VR controllers) that re-reference at the current robot - # pose on engage: the handover is already continuous there, and the blocking - # interpolation only delays the start of the correction. smooth_handover: bool = True input_device: str = "keyboard" keyboard: DAggerKeyboardConfig = field(default_factory=DAggerKeyboardConfig) pedal: DAggerPedalConfig = field(default_factory=DAggerPedalConfig) def __post_init__(self): + """Validate that `input_device` is a supported value. + + Raises: + ValueError: If `input_device` is not `"keyboard"` or `"pedal"`. + """ if self.input_device not in ("keyboard", "pedal"): raise ValueError(f"DAgger input_device must be 'keyboard' or 'pedal', got '{self.input_device}'") @@ -218,6 +280,68 @@ class RolloutConfig: Combines hardware, policy, strategy, and runtime settings. The ``__post_init__`` method performs fail-fast validation to reject invalid flag combinations early. + + Args: + robot (`RobotConfig | None`, *optional*): + Robot hardware configuration. Required — validated in `__post_init__`. + teleop (`TeleoperatorConfig | None`, *optional*): + Teleoperator hardware configuration. Required by the `dagger` strategy. + policy (`PreTrainedConfig | None`, *optional*): + Loaded automatically from `--policy.path` during `__post_init__`; do not set directly. + strategy (`RolloutStrategyConfig`, *optional*, defaults to `BaseStrategyConfig()`): + Polymorphic rollout strategy config, selected via `--strategy.type=base|sentry|highlight|dagger|episodic`. + inference (`InferenceEngineConfig`, *optional*, defaults to `SyncInferenceConfig()`): + Polymorphic inference backend config, selected via `--inference.type=sync|rtc`. + dataset (`DatasetRecordConfig | None`, *optional*): + Dataset recording configuration. Required for the `sentry`, `highlight`, `dagger`, and + `episodic` strategies; must be `None` for `base`. + fps (`float`, *optional*, defaults to 30.0): + Control loop frequency, in Hz. + duration (`float`, *optional*, defaults to 0.0): + Maximum rollout duration, in seconds. `0` means run indefinitely (24/7 mode). + interpolation_multiplier (`int`, *optional*, defaults to 1): + Number of interpolated control ticks generated per policy inference. + device (`str | None`, *optional*): + Torch device to run the policy on. Resolved from the policy config (or auto-selected) in + `__post_init__` when unset or unavailable. + task (`str`, *optional*, defaults to `""`): + Task description propagated to (or from) `dataset.single_task` in `__post_init__`. + display_data (`bool`, *optional*, defaults to `False`): + Whether to stream observation/action telemetry to a visualization backend. + display_mode (`str`, *optional*, defaults to `"rerun"`): + Visualization backend used when `display_data` is `True`: `"rerun"` or `"foxglove"`. + display_ip (`str | None`, *optional*): + For `"rerun"`: IP of a remote server to send to. For `"foxglove"`: interface to bind the + WebSocket server to (`127.0.0.1` for local only, `0.0.0.0` for all interfaces). + display_port (`int | None`, *optional*): + For `"rerun"`: port of the remote server. For `"foxglove"`: port to bind the WebSocket + server to. + display_compressed_images (`bool`, *optional*, defaults to `False`): + Whether to display compressed (JPEG) images instead of raw frames. + play_sounds (`bool`, *optional*, defaults to `True`): + Whether to use vocal synthesis to read out session events. + resume (`bool`, *optional*, defaults to `False`): + Whether to resume recording into an existing dataset instead of creating a new one. + rename_map (`dict[str, str]`, *optional*): + Mapping of robot/dataset observation keys to the policy's expected feature keys. + return_to_initial_position (`bool`, *optional*, defaults to `True`): + When `True`, smoothly interpolates the robot back to the joint positions captured at + startup before disconnecting. Set to `False` to leave the robot in its final achieved + pose at shutdown. + use_torch_compile (`bool`, *optional*, defaults to `False`): + Whether to wrap the policy's `predict_action_chunk` with `torch.compile`. + torch_compile_backend (`str`, *optional*, defaults to `"inductor"`): + Backend passed to `torch.compile`. + torch_compile_mode (`str`, *optional*, defaults to `"default"`): + Mode passed to `torch.compile`. + compile_warmup_inferences (`int`, *optional*, defaults to 2): + Number of warmup inferences run before `torch.compile`-backed inference is considered + ready. + + Raises: + ValueError: If a required flag combination is missing (e.g. `--robot.type`, `--policy.path`, + `--teleop.type` for DAgger, `--dataset.repo_id` for a recording strategy) or if the + strategy/dataset combination is invalid (e.g. a dataset passed to the `base` strategy). """ # Hardware @@ -243,25 +367,15 @@ class RolloutConfig: device: str | None = None task: str = "" display_data: bool = False - # Visualization backend used when display_data is True: "rerun" or "foxglove". display_mode: str = "rerun" - # For "rerun": IP of a remote server to send to. For "foxglove": interface to bind the WebSocket - # server to (127.0.0.1 for local only, 0.0.0.0 for all interfaces). display_ip: str | None = None - # For "rerun": port of the remote server. For "foxglove": port to bind the WebSocket server to. display_port: int | None = None - # Whether to display compressed (JPEG) images instead of raw frames display_compressed_images: bool = False - # Use vocal synthesis to read events play_sounds: bool = True resume: bool = False - # Rename map for mapping robot/dataset observation keys to policy keys rename_map: dict[str, str] = field(default_factory=dict) # Hardware teardown - # When True (default), smoothly interpolate the robot back to the joint - # positions captured at startup before disconnecting. Set to False to - # leave the robot in its final achieved pose at shutdown. return_to_initial_position: bool = True # Torch compile @@ -271,7 +385,12 @@ class RolloutConfig: compile_warmup_inferences: int = 2 def __post_init__(self): - """Validate config invariants and load the policy config from ``--policy.path``.""" + """Validate config invariants and load the policy config from ``--policy.path``. + + Raises: + ValueError: If a required flag combination is missing or the strategy/dataset combination + is invalid. + """ # --- Strategy-specific validation --- if isinstance(self.strategy, DAggerStrategyConfig) and self.teleop is None: raise ValueError("DAgger strategy requires --teleop.type to be set") @@ -384,4 +503,5 @@ class RolloutConfig: @classmethod def __get_path_fields__(cls) -> list[str]: + """Fields draccus resolves as pretrained-checkpoint paths (i.e. `--policy.path`).""" return ["policy"] diff --git a/src/lerobot/rollout/inference/factory.py b/src/lerobot/rollout/inference/factory.py index e600bed63..d57b104ce 100644 --- a/src/lerobot/rollout/inference/factory.py +++ b/src/lerobot/rollout/inference/factory.py @@ -54,6 +54,7 @@ class InferenceEngineConfig(draccus.ChoiceRegistry, abc.ABC): @property def type(self) -> str: + """The registered name of this backend (e.g. `"sync"`, `"rtc"`).""" return self.get_choice_name(self.__class__) @@ -66,10 +67,17 @@ class SyncInferenceConfig(InferenceEngineConfig): @InferenceEngineConfig.register_subclass("rtc") @dataclass class RTCInferenceConfig(InferenceEngineConfig): - """Real-Time Chunking: async policy inference in a background thread.""" + """Real-Time Chunking: async policy inference in a background thread. + + Args: + rtc (`RTCConfig`, *optional*): + RTC-specific configuration (e.g. prefix-attention schedule, execution horizon). Eagerly + constructed so draccus exposes nested fields directly on the CLI (e.g. + `--inference.rtc.execution_horizon=...`). + queue_threshold (`int`, *optional*, defaults to 30): + Action-queue size below which the background RTC thread starts producing a new chunk. + """ - # Eagerly constructed so draccus exposes nested fields directly on the CLI - # (e.g. ``--inference.rtc.execution_horizon=...``). rtc: RTCConfig = field(default_factory=RTCConfig) queue_threshold: int = 30 @@ -96,7 +104,44 @@ def create_inference_engine( compile_warmup_inferences: int = 2, shutdown_event: Event | None = None, ) -> InferenceEngine: - """Instantiate the appropriate inference engine from a config object.""" + """Instantiate the appropriate inference engine from a config object. + + Args: + config (`InferenceEngineConfig`): + Backend selector (`SyncInferenceConfig` or `RTCInferenceConfig`). + policy (`PreTrainedPolicy`): + The loaded policy to run inference with. + preprocessor (`PolicyProcessorPipeline`): + Observation pre-processor pipeline. + postprocessor (`PolicyProcessorPipeline`): + Action post-processor pipeline. + robot_wrapper (`ThreadSafeRobot`): + Thread-safe robot handle, used for RTC's background thread and to resolve `robot_type`. + hw_features (`dict`): + Raw hardware observation feature spec, used by RTC to rebuild dataset frames. + dataset_features (`dict`): + Dataset feature spec, used by sync inference to reorder policy outputs. + ordered_action_keys (`list[str]`): + Action key ordering the returned tensor should be mapped to. + task (`str`): + Task string passed through to the policy. + fps (`float`): + Control loop frequency, used by RTC to size its time-per-chunk estimate. + device (`str | None`): + Torch device to run inference on. + use_torch_compile (`bool`, *optional*, defaults to `False`): + Whether to `torch.compile` the policy's action-prediction call. + compile_warmup_inferences (`int`, *optional*, defaults to 2): + Number of warmup inferences before compiled inference is considered ready. + shutdown_event (`Event | None`, *optional*): + Global shutdown event RTC sets on an unrecoverable background-thread error. + + Returns: + InferenceEngine: The instantiated `SyncInferenceEngine` or `RTCInferenceEngine`. + + Raises: + ValueError: If `config` is not a recognized `InferenceEngineConfig` subclass. + """ logger.info("Creating inference engine: %s", config.type) if isinstance(config, SyncInferenceConfig): return SyncInferenceEngine( diff --git a/src/lerobot/rollout/inference/rtc.py b/src/lerobot/rollout/inference/rtc.py index 8ff44998a..195174ced 100644 --- a/src/lerobot/rollout/inference/rtc.py +++ b/src/lerobot/rollout/inference/rtc.py @@ -124,6 +124,37 @@ class RTCInferenceEngine(InferenceEngine): rtc_queue_threshold: int = 30, shutdown_event: Event | None = None, ) -> None: + """Build the engine (the background thread is started separately via `start`). + + Args: + policy (`PreTrainedPolicy`): + The RTC-capable policy to run inference with. + preprocessor (`PolicyProcessorPipeline`): + Observation pre-processor pipeline. + postprocessor (`PolicyProcessorPipeline`): + Action post-processor pipeline. + robot_wrapper (`ThreadSafeRobot`): + Thread-safe robot handle used to resolve `action_features` for relative-action + re-anchoring. + rtc_config (`RTCConfig`): + RTC configuration (execution horizon, prefix-attention schedule, etc.). + hw_features (`dict`): + Raw hardware observation feature spec used to rebuild dataset frames each tick. + task (`str`): + Task string passed through to the policy. + fps (`float`): + Control loop frequency, used to size the time-per-chunk estimate. + device (`str | None`): + Torch device to run inference on. Defaults to `"cpu"` when `None`. + use_torch_compile (`bool`, *optional*, defaults to `False`): + Whether to `torch.compile` the policy's action-prediction call. + compile_warmup_inferences (`int`, *optional*, defaults to 2): + Number of warmup inferences before `ready` reports `True`. + rtc_queue_threshold (`int`, *optional*, defaults to 30): + Action-queue size below which the background thread produces a new chunk. + shutdown_event (`Event | None`, *optional*): + Global shutdown event this engine sets on an unrecoverable background-thread error. + """ self._policy = policy self._preprocessor = preprocessor self._postprocessor = postprocessor diff --git a/src/lerobot/rollout/inference/sync.py b/src/lerobot/rollout/inference/sync.py index 2bb05b6ab..288d8a0cd 100644 --- a/src/lerobot/rollout/inference/sync.py +++ b/src/lerobot/rollout/inference/sync.py @@ -65,6 +65,26 @@ class SyncInferenceEngine(InferenceEngine): device: str | None, robot_type: str, ) -> None: + """Build the engine. + + Args: + policy (`PreTrainedPolicy`): + The policy to run inference with. + preprocessor (`PolicyProcessorPipeline`): + Observation pre-processor pipeline. + postprocessor (`PolicyProcessorPipeline`): + Action post-processor pipeline. + dataset_features (`dict`): + Dataset feature spec, used to reorder the policy's action output. + ordered_action_keys (`list[str]`): + Action key ordering the returned tensor is mapped to. + task (`str`): + Task string passed through to the policy. + device (`str | None`): + Torch device to run inference on. Defaults to `"cpu"` when `None`. + robot_type (`str`): + Robot type string, used for `prepare_observation_for_inference`'s per-robot handling. + """ self._policy = policy self._preprocessor = preprocessor self._postprocessor = postprocessor diff --git a/src/lerobot/rollout/ring_buffer.py b/src/lerobot/rollout/ring_buffer.py index 2c0a06301..534a83b91 100644 --- a/src/lerobot/rollout/ring_buffer.py +++ b/src/lerobot/rollout/ring_buffer.py @@ -34,20 +34,20 @@ class RolloutRingBuffer: must all be called from the same thread (the rollout main loop). Concurrent access from a background thread will corrupt ``_current_bytes`` accounting. - - Parameters - ---------- - max_seconds: - Maximum duration of buffered telemetry. - max_memory_mb: - Hard memory cap in MiB. Frames are evicted when the estimated - total size exceeds this. - fps: - Frames per second — used to convert ``max_seconds`` to a frame - count. """ def __init__(self, max_seconds: float = 30.0, max_memory_mb: int = 2048, fps: float = 30.0) -> None: + """Create an empty ring buffer. + + Args: + max_seconds (`float`, *optional*, defaults to 30.0): + Maximum duration of buffered telemetry. + max_memory_mb (`int`, *optional*, defaults to 2048): + Hard memory cap in MiB. Frames are evicted when the estimated total size exceeds + this. + fps (`float`, *optional*, defaults to 30.0): + Frames per second, used to convert `max_seconds` to a frame count. + """ self._max_frames = int(max_seconds * fps) self._max_bytes = int(max_memory_mb * 1024 * 1024) self._buffer: deque[dict] = deque(maxlen=self._max_frames) @@ -82,6 +82,7 @@ class RolloutRingBuffer: self._current_bytes = 0 def __len__(self) -> int: + """Number of frames currently buffered.""" return len(self._buffer) @property diff --git a/src/lerobot/rollout/robot_wrapper.py b/src/lerobot/rollout/robot_wrapper.py index 44f744812..a1b21448a 100644 --- a/src/lerobot/rollout/robot_wrapper.py +++ b/src/lerobot/rollout/robot_wrapper.py @@ -34,16 +34,23 @@ class ThreadSafeRobot: """ def __init__(self, robot: Robot) -> None: + """Wrap `robot` behind a lock. + + Args: + robot (`Robot`): The connected robot instance to protect. + """ self._robot = robot self._lock = Lock() # -- Lock-protected I/O -------------------------------------------------- def get_observation(self) -> dict[str, Any]: + """See [`~robots.Robot.get_observation`].""" with self._lock: return self._robot.get_observation() def send_action(self, action: dict[str, Any] | Any) -> Any: + """See [`~robots.Robot.send_action`].""" with self._lock: return self._robot.send_action(action) @@ -51,26 +58,32 @@ class ThreadSafeRobot: @property def observation_features(self) -> dict: + """See [`~robots.Robot.observation_features`].""" return self._robot.observation_features @property def action_features(self) -> dict: + """See [`~robots.Robot.action_features`].""" return self._robot.action_features @property def name(self) -> str: + """See [`~robots.Robot.name`].""" return self._robot.name @property def robot_type(self) -> str: + """See [`~robots.Robot.robot_type`].""" return self._robot.robot_type @property def cameras(self): + """The wrapped robot's cameras, or `{}` if it has none.""" return getattr(self._robot, "cameras", {}) @property def is_connected(self) -> bool: + """See [`~robots.Robot.is_connected`].""" return self._robot.is_connected @property diff --git a/src/lerobot/rollout/strategies/core.py b/src/lerobot/rollout/strategies/core.py index 460ad12e5..8087439b0 100644 --- a/src/lerobot/rollout/strategies/core.py +++ b/src/lerobot/rollout/strategies/core.py @@ -46,6 +46,11 @@ class RolloutStrategy(abc.ABC): """ def __init__(self, config: RolloutStrategyConfig) -> None: + """Store `config`; the inference engine is attached later via `_init_engine`. + + Args: + config (`RolloutStrategyConfig`): This strategy's configuration. + """ self.config = config self._engine: InferenceEngine | None = None self._interpolator: ActionInterpolator | None = None diff --git a/src/lerobot/rollout/strategies/dagger.py b/src/lerobot/rollout/strategies/dagger.py index dec403ea9..92f63a2fa 100644 --- a/src/lerobot/rollout/strategies/dagger.py +++ b/src/lerobot/rollout/strategies/dagger.py @@ -105,6 +105,7 @@ class DAggerEvents: """ def __init__(self) -> None: + """Create a fresh events container, starting in the `AUTONOMOUS` phase.""" self._lock = Lock() self._phase = DAggerPhase.AUTONOMOUS self._pending_transition: str | None = None @@ -123,6 +124,7 @@ class DAggerEvents: @phase.setter def phase(self, value: DAggerPhase) -> None: + """Set the current phase directly, bypassing `_DAGGER_TRANSITIONS` validation.""" with self._lock: self._phase = value @@ -207,6 +209,7 @@ def _init_dagger_pedal(events: DAggerEvents, cfg: DAggerPedalConfig): } def on_press(code: str) -> None: + """Apply a resolved pedal code to the DAgger events.""" if code in code_to_event: events.request_transition(code_to_event[code]) if code == cfg.upload: @@ -239,6 +242,7 @@ class DAggerStrategy(RolloutStrategy): config: DAggerStrategyConfig def __init__(self, config: DAggerStrategyConfig): + """See [`~rollout.RolloutStrategy.__init__`].""" super().__init__(config) self._listener = None self._pedal_thread = None diff --git a/src/lerobot/rollout/strategies/episodic.py b/src/lerobot/rollout/strategies/episodic.py index e4eb9a885..b87e97df4 100644 --- a/src/lerobot/rollout/strategies/episodic.py +++ b/src/lerobot/rollout/strategies/episodic.py @@ -74,6 +74,7 @@ class EpisodicStrategy(RolloutStrategy): config: EpisodicStrategyConfig def __init__(self, config: EpisodicStrategyConfig) -> None: + """See [`~rollout.RolloutStrategy.__init__`].""" super().__init__(config) self._listener = None self._events: dict | None = None diff --git a/src/lerobot/rollout/strategies/highlight.py b/src/lerobot/rollout/strategies/highlight.py index 385a9e2b6..3af1112dd 100644 --- a/src/lerobot/rollout/strategies/highlight.py +++ b/src/lerobot/rollout/strategies/highlight.py @@ -55,6 +55,7 @@ class HighlightStrategy(RolloutStrategy): config: HighlightStrategyConfig def __init__(self, config: HighlightStrategyConfig): + """See [`~rollout.RolloutStrategy.__init__`].""" super().__init__(config) self._ring: RolloutRingBuffer | None = None self._listener = None diff --git a/src/lerobot/rollout/strategies/sentry.py b/src/lerobot/rollout/strategies/sentry.py index 61e38aa68..797afefa9 100644 --- a/src/lerobot/rollout/strategies/sentry.py +++ b/src/lerobot/rollout/strategies/sentry.py @@ -60,6 +60,7 @@ class SentryStrategy(RolloutStrategy): config: SentryStrategyConfig def __init__(self, config: SentryStrategyConfig): + """See [`~rollout.RolloutStrategy.__init__`].""" super().__init__(config) self._push_executor: ThreadPoolExecutor | None = None self._pending_push: Future | None = None diff --git a/utils/check_docstrings.py b/utils/check_docstrings.py index 254c7c872..151f12f10 100644 --- a/utils/check_docstrings.py +++ b/utils/check_docstrings.py @@ -60,6 +60,7 @@ PATH_TO_LEROBOT = PATH_TO_REPO / "src" / "lerobot" # Modules whose public objects are checked. Add a module here once its docstrings follow the standard. MODULES_TO_CHECK = [ "lerobot.robots", + "lerobot.rollout", ] # Objects that do not yet follow the standard, so the check can be green from day one. Removing an entry