LeRobot documentation
实现你自己的 Robot Processor
实现你自己的 Robot Processor
在本教程中，您将学习如何实现自己的 Robot Processor。
首先，我们将探讨为何需要自定义处理器，然后以 NormalizerProcessorStep 为例，说明如何实现、配置和序列化处理器。有关 LeRobot 自带的处理器步骤，请参阅 常见处理器步骤。
为什么需要自定义处理器？
在大多数情况下，当从传感器读取原始数据或模型输出动作时，你需要处理这些数据，使其与目标系统兼容。例如，一个常见的需求是归一化数据范围，使其适合神经网络使用。
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)
# ...其他常见的处理需求包括：
- 设备放置：在 CPU/GPU 之间移动张量并转换数据类型
- 格式转换：在不同数据结构之间进行转换
- 批处理：添加/移除批处理维度以兼容模型
- 安全约束：对机器人指令施加限制
# Example pipeline combining multiple processors
pipeline = PolicyProcessorPipeline([
RenameObservationsProcessorStep(rename_map={}),
AddBatchDimensionProcessorStep(),
NormalizerProcessorStep(features=features, norm_map=norm_map, stats=stats),
DeviceProcessorStep(device="cuda"),
# ...
])LeRobot 提供了一套流水线机制，用于为输入数据和输出动作实现一系列处理步骤，从而可以轻松地将这些转换按正确的顺序组合起来，以获得最佳性能。
如何实现你自己的处理器？
我们将以 NormalizerProcessorStep 作为主要示例，因为它展示了处理器中常见的几种基本模式，包括状态管理、配置序列化和张量处理。
准备解决您的问题所需的处理步骤序列。处理步骤是一个实现了以下方法的类：
__call__：实现输入转换的处理步骤。get_config：获取处理器步骤的配置。state_dict：获取处理器步骤的状态。load_state_dict：加载处理器步骤的状态。reset：重置处理器步骤的状态。transform_features：描述处理器步骤如何修改特征空间。
实现 __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()以避免副作用 - 一致地处理观测和动作
- 将配置与状态分离：
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 方法定义了处理器如何转换特征名称和形状。这对于策略配置和调试至关重要。
对于 NormalizerProcessorStep，特征通常保持不变，因为归一化不会更改键或形状：
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添加重命名后的特征 - 始终返回修改后的 features 字典
- 在文档字符串中清晰地说明数据转换
使用覆盖
您可以在加载时使用 overrides 来覆盖步骤参数。这对于不可序列化的对象或特定于站点的设置非常有用。它既适用于策略工厂，也适用于 DataProcessorPipeline.from_pretrained(...)。
基础模型适配：当处理基础预训练策略时，这尤其有用，因为您通常无法获取原始训练统计数据。您可以注入自己的数据集统计信息，以将归一化器适配到您特定的机器人或环境数据。
示例：在机器人上进行策略评估时，覆盖设备并重命名映射。 使用此功能可在仅配备 CPU 的机器人上运行在 CUDA 上训练的策略，或在机器人使用的名称与数据集不同时重新映射相机键。
直接使用 from_pretrained：
from lerobot.datasets import LeRobotDataset
from lerobot.processor import PolicyProcessorPipeline
# 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 = PolicyProcessorPipeline.from_pretrained(
"huggingface/foundational-robot-policy", # Pretrained foundation model
config_filename="policy_preprocessor.json", # The policy's saved preprocessor
overrides={
"normalizer_processor": {"stats": new_stats}, # Inject your robot's statistics
"device_processor": {"device": "cuda:0"}, # registry name for registered steps
"rename_observations_processor": {"rename_map": robot_key_map}, # Map your robot's observation keys
# ...
},
)最佳实践
基于对所有 LeRobot 处理器实现的分析，以下是关键的模式和实践：
1. 安全的数据处理
始终创建输入数据的副本，以避免产生意外的副作用。请使用 transition.copy() 和 observation.copy()，而不是就地修改数据。这可以防止您的处理器意外影响流水线中的其他组件。
在处理之前检查所需的数据，并妥善处理缺失的数据。如果你的处理器期望某些键（例如用于图像处理的 "pixels"），请先验证它们是否存在。对于可选数据，请使用 transition.get() 等安全访问模式，并适当处理 None 值。
当数据验证失败时，请提供清晰且可操作的错误信息，帮助用户了解问题所在以及如何修复。
2. 选择合适的基类
LeRobot 提供了专用的基类，以减少样板代码并确保一致性。当仅需修改观测值时，请使用 ObservationProcessorStep；仅处理动作时，请使用 ActionProcessorStep；而对于基于字典的机器人动作，请专门使用 RobotActionProcessorStep。
仅在需要完全控制整个过渡过程，或需要同时处理多个过渡组件时，才直接继承自 ProcessorStep。专用基类会为您处理过渡管理，并提供类型安全。
3. 注册和命名
使用 @ProcessorStepRegistry.register() 以具有描述性且带命名空间的名称注册您的处理器。使用诸如 "robotics_lab/safety_clipper" 或 "acme_corp/vision_enhancer" 的组织前缀来避免命名冲突。避免使用诸如 "processor" 或 "step" 等通用名称，因为它们可能会与其他实现发生冲突。
良好的注册机制可以让您的处理器被发现，并在保存和加载流水线时实现干净的序列化/反序列化。
4. 状态管理模式
区分配置参数（可 JSON 序列化的值）和内部状态（张量、缓冲区）。对于不应出现在构造函数或字符串表示中的内部状态，请使用带有 init=False, repr=False 的 dataclass 字段。
实现 reset() 方法以在回合之间清除内部状态。这对于随时间累积数据的有状态处理器（例如移动平均或时间滤波器）至关重要。
请记住，get_config() 应仅返回可 JSON 序列化的配置，而 state_dict() 则单独处理张量状态。
5. 输入验证与错误处理
在处理之前验证输入类型和形状。检查张量属性（如 dtype 和维度），以确保其与您的算法兼容。对于机器人动作，请验证所需的位姿分量或关节值是否存在，并且处于预期范围内。
对于无需处理的边缘情况，请使用提前返回。提供清晰且具有描述性的错误信息，其中应包含预期与实际的数据类型或形状。这将使用户调试变得更加容易。
6. 设备和数据类型感知
设计你的处理器，使其能够自动适配输入张量的设备和数据类型。内部张量（例如归一化统计量）应与输入张量的设备和数据类型保持一致，以确保与多 GPU 训练、混合精度以及分布式设置的兼容性。
实现一个 to() 方法，将处理器的内部状态移动到指定设备。在运行时检查设备/数据类型兼容性，并在需要时自动迁移内部状态。此模式可实现在不同硬件配置之间无缝运行，无需手动干预。
结论
现在，你已经拥有了在 LeRobot 中实现自定义处理器的所有工具！关键步骤如下：
- 将处理器定义为包含所需方法（
__call__、get_config、state_dict、load_state_dict、reset、transform_features）的数据类 - 使用
@ProcessorStepRegistry.register("name")注册它 以便发现 - 将其集成到
DataProcessorPipeline中，与其他处理步骤一起使用 - 尽可能使用 基类，例如
ObservationProcessorStep，以减少样板代码 - 实现设备/数据类型感知以支持多 GPU 和混合精度配置
处理器系统采用模块化且可组合的设计，允许您从简单、专注的组件构建复杂的数据处理流水线。无论是为训练预处理传感器数据，还是为机器人执行后处理模型输出，自定义处理器都能为您提供所需的灵活性，以处理机器人应用中任何数据转换需求。
构建稳健处理器的关键原则：
- 设备/数据类型适配：内部张量应与输入张量匹配
- 清晰的错误信息：帮助用户理解出了什么问题
- 基类用法：利用专门的基类来减少样板代码
- 功能契约：使用
transform_features()声明数据结构变更
从简单开始，充分测试，并确保你的处理器在不同的硬件配置上都能无缝运行！
Update on GitHub