LeRobot 文件
實作您自己的機器人處理器
並獲得增強的文件體驗
開始使用
實作您自己的機器人處理器
在本教學中,您將學習如何實作自己的機器人處理器 (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 中實作自訂處理器的所有工具!關鍵步驟包括:
- 將處理器定義為一個 dataclass 並包含所需的方法 (
__call__,get_config,state_dict,load_state_dict,reset,transform_features) - 使用
@ProcessorStepRegistry.register("name")註冊它 以便於發現 - 將其與其他處理步驟一同整合到
DataProcessorPipeline中 - 儘可能使用
ObservationProcessorStep等基底類別以減少重複程式碼 - 實作裝置/dtype 感知,以支援多 GPU 和混合精度設置
處理器系統旨在模組化與可組合化,讓您可以從簡單、專注的組件建構複雜的數據處理管道。無論您是為了訓練而預處理感測器數據,還是為了機器人執行而後處理模型輸出,自訂處理器都能提供靈活性,以處理機器人應用程序所需的任何數據轉換。
強韌處理器的關鍵原則:
- 裝置/dtype 適配:內部張量應匹配輸入張量
- 清晰的錯誤訊息:幫助使用者了解出了什麼問題
- 基底類別使用:利用專門的基底類別減少重複程式碼
- 特徵契約 (Feature contracts):使用
transform_features()聲明數據結構變化
從簡單開始,徹底測試,並確保您的處理器在不同硬體配置上都能無縫運作!
在 GitHub 上更新