mirror of
https://github.com/huggingface/lerobot.git
synced 2026-07-29 04:36:04 +00:00
Maxime Ellerbach
99 lines
3.4 KiB
Python
99 lines
3.4 KiB
Python
# Copyright 2025 The HuggingFace Inc. team. All rights reserved.
|
|
#
|
|
# Licensed under the Apache License, Version 2.0 (the "License");
|
|
# you may not use this file except in compliance with the License.
|
|
# You may obtain a copy of the License at
|
|
#
|
|
# http://www.apache.org/licenses/LICENSE-2.0
|
|
#
|
|
# Unless required by applicable law or agreed to in writing, software
|
|
# distributed under the License is distributed on an "AS IS" BASIS,
|
|
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
# See the License for the specific language governing permissions and
|
|
# limitations under the License.
|
|
|
|
"""Inference engine ABC.
|
|
|
|
Rollout strategies consume actions through this small interface so they
|
|
do not need to know whether inference happens inline on the control thread
|
|
or asynchronously in a background thread (RTC).
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import abc
|
|
|
|
import torch
|
|
|
|
|
|
class InferenceEngine(abc.ABC):
|
|
"""Abstract backend for producing actions during rollout.
|
|
|
|
Subclasses decide whether inference happens inline on the control
|
|
thread or asynchronously in a background thread. The contract is
|
|
minimal so additional backends can be plugged in without touching
|
|
rollout strategies.
|
|
|
|
Lifecycle
|
|
---------
|
|
``start`` — prepare the backend (e.g. launch a background thread).
|
|
``stop`` — shut the backend down cleanly.
|
|
``reset`` — clear episode-scoped state (policy hidden state, queues…).
|
|
|
|
Action production
|
|
-----------------
|
|
``get_action(obs_frame)`` — return the next action tensor, or
|
|
``None`` if none is available (e.g. async queue empty). Sync
|
|
backends always compute from ``obs_frame``; async backends ignore
|
|
it (they receive observations via ``notify_observation``).
|
|
|
|
Optional hooks
|
|
--------------
|
|
``notify_observation`` / ``pause`` / ``resume`` have a no-op default
|
|
so rollout strategies can invoke them unconditionally.
|
|
"""
|
|
|
|
@abc.abstractmethod
|
|
def start(self) -> None:
|
|
"""Initialise the backend."""
|
|
|
|
@abc.abstractmethod
|
|
def stop(self) -> None:
|
|
"""Tear the backend down."""
|
|
|
|
@abc.abstractmethod
|
|
def reset(self) -> None:
|
|
"""Clear episode-scoped state."""
|
|
|
|
@abc.abstractmethod
|
|
def get_action(self, obs_frame: dict | None) -> torch.Tensor | None:
|
|
"""Return the next action tensor, or ``None`` if unavailable."""
|
|
|
|
def get_intermediate_predictions(self) -> dict | None:
|
|
"""Extra display-ready model outputs to visualize this tick, or ``None``.
|
|
|
|
Lets a backend surface a world model's intermediate predictions (e.g. imagined video
|
|
frames) into the rollout visualization path, keyed by ``"<datatype>.<name>"`` (mirroring
|
|
observation feature keys). Default: nothing extra.
|
|
"""
|
|
return None
|
|
|
|
def notify_observation(self, obs: dict) -> None: # noqa: B027
|
|
"""Publish the latest processed observation. Default: no-op."""
|
|
|
|
def pause(self) -> None: # noqa: B027
|
|
"""Pause background inference. Default: no-op."""
|
|
|
|
def resume(self) -> None: # noqa: B027
|
|
"""Resume background inference. Default: no-op."""
|
|
|
|
@property
|
|
def ready(self) -> bool:
|
|
"""True once the backend can produce actions (e.g. warmup done)."""
|
|
return True
|
|
|
|
@property
|
|
def failed(self) -> bool:
|
|
"""True if an unrecoverable error occurred in the backend."""
|
|
return False
|