# 实现你自己的 Robot Processor

在本教程中，您将学习如何实现自己的 Robot Processor。首先，我们将探讨为何需要自定义处理器，然后以 `NormalizerProcessorStep` 为例，说明如何实现、配置和序列化处理器。有关 LeRobot 自带的处理器步骤，请参阅 [常见处理器步骤](./introduction_processors#common-processor-steps)。

## 为什么需要自定义处理器？

在大多数情况下，当从传感器读取原始数据或模型输出动作时，你需要处理这些数据，使其与目标系统兼容。例如，一个常见的需求是归一化数据范围，使其适合神经网络使用。

LeRobot 的 `NormalizerProcessorStep` 负责处理这项关键任务：

```python
# 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 之间移动张量并转换数据类型
- **格式转换**：在不同数据结构之间进行转换
- **批处理**：添加/移除批处理维度以兼容模型
- **安全约束**：对机器人指令施加限制

```python
# 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` 的工作方式如下：

```python
@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` 完美地展示了这一点——它在其状态中携带数据集统计信息（张量），并在其配置中携带超参数：

```python
# 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
        )
        # ...
```

**用法：**

```python
# 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`，特征通常保持不变，因为归一化不会更改键或形状：

```python
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
    # ...
```

当您的处理器重命名或重塑数据时，请实现此方法以反映下游组件的映射关系。例如，一个简单的重命名处理器：

```python
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`：

```python
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 中实现自定义处理器的所有工具！关键步骤如下：

1. 将**处理器**定义为包含所需方法（`__call__`、`get_config`、`state_dict`、`load_state_dict`、`reset`、`transform_features`）的数据类
2. 使用 `@ProcessorStepRegistry.register("name")` **注册它** 以便发现
3. 将其**集成**到 `DataProcessorPipeline` 中，与其他处理步骤一起使用
4. 尽可能使用 **基类**，例如 `ObservationProcessorStep`，以减少样板代码
5. **实现设备/数据类型感知**以支持多 GPU 和混合精度配置

处理器系统采用模块化且可组合的设计，允许您从简单、专注的组件构建复杂的数据处理流水线。无论是为训练预处理传感器数据，还是为机器人执行后处理模型输出，自定义处理器都能为您提供所需的灵活性，以处理机器人应用中任何数据转换需求。

构建稳健处理器的关键原则：

- **设备/数据类型适配**：内部张量应与输入张量匹配
- **清晰的错误信息**：帮助用户理解出了什么问题
- **基类用法**：利用专门的基类来减少样板代码
- **功能契约**：使用 `transform_features()` 声明数据结构变更

从简单开始，充分测试，并确保你的处理器在不同的硬件配置上都能无缝运行！