LeRobot 文件

實作您自己的機器人處理器

Hugging Face's logo
加入 Hugging Face 社群

並獲得增強的文件體驗

開始使用

實作您自己的機器人處理器

在本教學中,您將學習如何實作自己的機器人處理器 (Robot Processor)。首先會探討為何需要自訂處理器,接著以 NormalizerProcessorStep 為例,說明如何實作、設定及序列化處理器。最後,會列出 LeRobot 隨附的所有輔助處理器。

為什麼您會需要自訂處理器?

在大多數情況下,當從感測器讀取原始數據或當模型輸出動作 (Action) 時,您需要處理這些數據使其與目標系統相容。例如,常見的需求是正規化 (Normalizing) 數據範圍,使其適用於神經網路。

LeRobot 的 NormalizerProcessorStep 負責處理這項關鍵任務。

# Input: raw joint positions in [0, 180] degrees
raw_action = torch.tensor([90.0, 45.0, 135.0])

# After processing: normalized to [-1, 1] range for model training
normalizer = NormalizerProcessorStep(features=features, norm_map=norm_map, stats=dataset_stats)
normalized_result = normalizer(transition)
# ...

其他常見的處理需求包括:

  • 裝置放置 (Device placement):在 CPU/GPU 之間移動張量 (Tensor) 並轉換數據類型
  • 格式轉換:在不同的數據結構之間進行轉換
  • 批次處理 (Batching):增加/移除批次維度以滿足模型相容性
  • 安全約束:對機器人指令套用限制
# Example pipeline combining multiple processors
pipeline = PolicyProcessorPipeline([
    RenameObservationsProcessorStep(rename_map={}),
    AddBatchDimensionProcessorStep(),
    NormalizerProcessorStep(features=features, stats=stats),
    DeviceProcessorStep(device="cuda"),
    # ...
])

LeRobot 提供了一種管道 (Pipeline) 機制,用於實作輸入數據和輸出動作的處理步驟序列,讓您可以輕鬆地按正確順序組合這些轉換,以獲得最佳效能。

如何實作您自己的處理器?

我們將以 NormalizerProcessorStep 作為主要範例,因為它展示了必要的處理器模式,包括您通常會需要的狀態管理、設定序列化和張量處理。

準備針對您的問題所需的處理步驟序列。處理步驟是一個實作以下方法的類別:

  • __call__:實作輸入轉換 (transition) 的處理步驟。
  • get_config:取得處理步驟的設定 (Configuration)。
  • state_dict:取得處理步驟的狀態 (State)。
  • load_state_dict:載入處理步驟的狀態。
  • reset:重設處理步驟的狀態。
  • feature_contract:顯示處理步驟期間對特徵空間 (Feature Space) 的修改。

實作 __call__ 方法

__call__ 方法是處理步驟的核心。它接收一個 EnvTransition 並返回一個修改後的 EnvTransition。以下是 NormalizerProcessorStep 的運作方式:

@dataclass
@ProcessorStepRegistry.register("normalizer_processor")
class NormalizerProcessorStep(ProcessorStep):
    """Normalize observations/actions using dataset statistics."""

    features: dict[str, PolicyFeature]
    norm_map: dict[FeatureType, NormalizationMode]
    stats: dict[str, dict[str, Any]] | None = None
    eps: float = 1e-8
    _tensor_stats: dict = field(default_factory=dict, init=False, repr=False)

    def __post_init__(self):
        """Convert stats to tensors for efficient computation."""
        self.stats = self.stats or {}
        self._tensor_stats = to_tensor(self.stats, device=self.device, dtype=torch.float32)

    def __call__(self, transition: EnvTransition) -> EnvTransition:
        new_transition = transition.copy()
        # Normalize observations
        # ...
        # Normalize action
        # ...
        return new_transition

如需完整細節,請參閱 src/lerobot/processor/normalize_processor.py 中的完整實作。

關鍵原則:

  • 始終使用 transition.copy() 以避免副作用
  • 一致地處理觀察 (Observations) 和動作 (Actions)
  • 區分設定與狀態get_config() 返回可 JSON 序列化的參數,state_dict() 返回張量
  • __post_init__() 中將統計數據轉換為張量,以實現高效計算

設定與狀態管理

處理器透過將設定與張量狀態分離的三個方法支援序列化。NormalizerProcessorStep 完美展示了這一點——它在狀態中攜帶數據集統計信息(張量),而在設定中攜帶超參數。

# Continuing the NormalizerProcessorStep example...

def get_config(self) -> dict[str, Any]:
    """JSON-serializable configuration (no tensors)."""
    return {
        "eps": self.eps,
        "features": {k: {"type": v.type.value, "shape": v.shape} for k, v in self.features.items()},
        "norm_map": {ft.value: nm.value for ft, nm in self.norm_map.items()},
        # ...
    }

def state_dict(self) -> dict[str, torch.Tensor]:
    """Tensor state only (e.g., dataset statistics)."""
    flat: dict[str, torch.Tensor] = {}
    for key, sub in self._tensor_stats.items():
        for stat_name, tensor in sub.items():
            flat[f"{key}.{stat_name}"] = tensor.cpu()  # Always save to CPU
    return flat

def load_state_dict(self, state: dict[str, torch.Tensor]) -> None:
    """Restore tensor state at runtime."""
    self._tensor_stats.clear()
    for flat_key, tensor in state.items():
        key, stat_name = flat_key.rsplit(".", 1)
        # Load to processor's configured device
        self._tensor_stats.setdefault(key, {})[stat_name] = tensor.to(
            dtype=torch.float32, device=self.device
        )
        # ...

用法

# Save (e.g., inside a policy)
config = normalizer.get_config()
tensors = normalizer.state_dict()

# Restore (e.g., loading a pretrained policy)
new_normalizer = NormalizerProcessorStep(**config)
new_normalizer.load_state_dict(tensors)
# Now new_normalizer has the same stats and configuration

轉換特徵

transform_features 方法定義了您的處理器如何轉換特徵名稱和形狀。這對於策略 (Policy) 設定和除錯至關重要。

對於 NormalizerProcessorStep,特徵通常保持不變,因為正規化不會改變鍵值 (Key) 或形狀。

def transform_features(self, features: dict[PipelineFeatureType, dict[str, PolicyFeature]]) -> dict[PipelineFeatureType, dict[str, PolicyFeature]]:
    """Normalization preserves all feature definitions."""
    return features  # No changes to feature structure
    # ...

當您的處理器重新命名或重塑數據時,請實作此方法以反映下游組件的對應關係。例如,一個簡單的重新命名處理器:

def transform_features(self, features: dict[str, PolicyFeature]) -> dict[str, PolicyFeature]:
    # Simple renaming
    if "pixels" in features:
        features["observation.image"] = features.pop("pixels")

    # Pattern-based renaming
    for key in list(features.keys()):
        if key.startswith("env_state."):
            suffix = key[len("env_state."):]
            features[f"observation.{suffix}"] = features.pop(key)
            # ...

    return features

關鍵原則:

  • 使用 features.pop(old_key) 來移除並取得舊特徵
  • 使用 features[new_key] = old_feature 來新增重新命名後的特徵
  • 始終返回修改後的特徵字典
  • 在 docstring 中清楚記錄轉換過程

使用覆寫 (Overrides)

您可以在載入時使用 overrides 覆寫步驟參數。這對於不可序列化的物件或特定場域的設定非常方便。它在策略工廠和 DataProcessorPipeline.from_pretrained(...) 中均可運作。

基礎模型調適 (Foundational model adaptation):這在使用基礎預訓練策略時特別有用,因為您很少能獲得原始訓練統計數據。您可以注入自己的數據集統計信息,使正規化器適配您特定的機器人或環境數據。

範例:在機器人上進行策略評估期間,覆寫裝置 (Device) 和重新命名地圖 (Rename map)。使用此功能可以在僅有 CPU 的機器人上執行在 CUDA 上訓練的策略,或在機器人使用的名稱與數據集不同時重新映射相機鍵值。

直接與 from_pretrained 搭配使用:

from lerobot.processor import RobotProcessorPipeline

# Load a foundational policy trained on diverse robot data
# but adapt normalization to your specific robot/environment
new_stats = LeRobotDataset(repo_id="username/my-dataset").meta.stats
processor = RobotProcessorPipeline.from_pretrained(
    "huggingface/foundational-robot-policy",  # Pretrained foundation model
    overrides={
        "normalizer_processor": {"stats": new_stats},     # Inject your robot's statistics
        "device_processor": {"device": "cuda:0"},         # registry name for registered steps
        "rename_processor": {"rename_map": robot_key_map}, # Map your robot's observation keys
        # ...
    },
)

最佳實踐

基於對所有 LeRobot 處理器實作的分析,以下是關鍵模式和實作:

1. 數據安全處理

始終建立輸入數據的副本,以避免意外的副作用。使用 transition.copy()observation.copy(),而不是直接就地 (In-place) 修改數據。這可防止您的處理器意外影響管道中的其他組件。

在處理前檢查必要的數據,並優雅地處理缺失數據。如果您的處理器需要特定的鍵值(例如圖像處理需要的 "pixels"),請先驗證它們是否存在。對於選填數據,請使用 transition.get() 等安全存取模式,並適當處理 None 值。

當數據驗證失敗時,提供清晰且可操作的錯誤訊息,幫助使用者了解出了什麼問題以及如何修正。

2. 選擇合適的基底類別

LeRobot 提供了專門的基底類別 (Base classes),可減少重複程式碼並確保一致性。當您只需要修改觀察時使用 ObservationProcessorStep;對於僅限動作的處理使用 ActionProcessorStep;針對基於字典的機器人動作則使用 RobotActionProcessorStep

只有在需要完全控制整個 transition 或同時處理多個 transition 組件時,才直接繼承 ProcessorStep。專門的基底類別會為您處理 transition 管理並提供類型安全。

3. 註冊與命名

使用 @ProcessorStepRegistry.register() 並賦予具描述性、具命名空間的名稱來註冊您的處理器。使用組織前綴,如 "robotics_lab/safety_clipper""acme_corp/vision_enhancer",以避免命名衝突。避免使用 "processor""step" 等可能與其他實作產生衝突的通用名稱。

良好的註冊能讓您的處理器易於被發現,並在儲存與載入管道時實現乾淨的序列化/反序列化。

4. 狀態管理模式

區分設定參數(可 JSON 序列化的值)和內部狀態(張量、緩衝區)。對於不應出現在建構子或字串表示中的內部狀態,請使用帶有 init=False, repr=False 的 dataclass 欄位。

實作 reset() 方法,以在回合 (Episodes) 之間清除內部狀態。這對於隨時間累積數據的有狀態 (Stateful) 處理器(如移動平均值或時間濾波器)至關重要。

請記住,get_config() 應僅返回可 JSON 序列化的設定,而 state_dict() 則分開處理張量狀態。

5. 輸入驗證與錯誤處理

在處理前驗證輸入類型與形狀。檢查張量屬性(如 dtype 和維度),以確保與您的演算法相容。對於機器人動作,請驗證所需的姿態 (Pose) 組件或關節值是否存在且位於預期範圍內。

對於不需要處理的邊緣情況 (Edge cases),請使用提早返回 (Early returns)。提供清晰且具描述性的錯誤訊息,包含預期與實際的數據類型或形狀。這能大幅降低使用者的除錯難度。

6. 裝置與數據類型感知 (Device and Dtype Awareness)

設計您的處理器以自動適應輸入張量的裝置與數據類型。內部張量(如正規化統計數據)應與輸入張量的裝置和 dtype 相匹配,以確保與多 GPU 訓練、混合精度 (Mixed precision) 和分散式設置相容。

實作一個 to() 方法,將處理器的內部狀態移動到指定裝置。在運行時檢查裝置/dtype 相容性,並在需要時自動遷移內部狀態。這種模式能在不同硬體配置之間實現無縫運作,而無需手動干預。

結論

您現在已擁有在 LeRobot 中實作自訂處理器的所有工具!關鍵步驟包括:

  1. 將處理器定義為一個 dataclass 並包含所需的方法 (__call__, get_config, state_dict, load_state_dict, reset, transform_features)
  2. 使用 @ProcessorStepRegistry.register("name") 註冊它 以便於發現
  3. 將其與其他處理步驟一同整合DataProcessorPipeline
  4. 儘可能使用 ObservationProcessorStep 等基底類別以減少重複程式碼
  5. 實作裝置/dtype 感知,以支援多 GPU 和混合精度設置

處理器系統旨在模組化與可組合化,讓您可以從簡單、專注的組件建構複雜的數據處理管道。無論您是為了訓練而預處理感測器數據,還是為了機器人執行而後處理模型輸出,自訂處理器都能提供靈活性,以處理機器人應用程序所需的任何數據轉換。

強韌處理器的關鍵原則:

  • 裝置/dtype 適配:內部張量應匹配輸入張量
  • 清晰的錯誤訊息:幫助使用者了解出了什麼問題
  • 基底類別使用:利用專門的基底類別減少重複程式碼
  • 特徵契約 (Feature contracts):使用 transform_features() 聲明數據結構變化

從簡單開始,徹底測試,並確保您的處理器在不同硬體配置上都能無縫運作!

在 GitHub 上更新

© . This site is unofficial and not affiliated with Hugging Face, Inc.