Реализуйте собственный Robot Processor

В этом руководстве вы узнаете, как реализовать собственный Robot Processor. Сначала мы рассмотрим необходимость в кастомном процессоре, а затем на примере NormalizerProcessorStep объясним, как реализовать, настроить и сериализовать процессор. С шагами процессора, входящими в состав LeRobot, можно ознакомиться в разделе Common Processor Steps.

Зачем нужен собственный процессор?

В большинстве случаев при чтении сырых данных с датчиков или при получении от моделей действий необходимо обработать эти данные, чтобы сделать их совместимыми с целевой системой. Например, часто требуется нормализовать диапазоны данных, чтобы они стали пригодными для нейронных сетей.

NormalizerProcessorStep в LeRobot выполняет эту важную задачу:

# 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, чтобы добавить переименованную характеристику
• Всегда возвращайте изменённый словарь признаков
• Чётко описывайте преобразования данных в docstring

Использование переопределений

Вы можете переопределить параметры шага при загрузке с помощью overrides. Это удобно для объектов, которые нельзя сериализовать, или для настроек, специфичных для конкретного окружения. Это работает как в фабриках политик, так и с DataProcessorPipeline.from_pretrained(...).

Адаптация базовой модели: это особенно полезно при работе с базовыми предобученными политиками, когда у вас редко есть доступ к исходным статистикам обучения. Вы можете внедрить статистику собственного набора данных, чтобы адаптировать нормализатор под данные конкретного робота или среды.

Пример: при оценке политики на роботе переопределите устройство и карту переименований. Используйте это для запуска политики, обученной на CUDA, на роботе с только CPU, или для переназначения ключей камер, если робот использует другие имена, чем в наборе данных.

Прямое использование с 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) и внутреннее состояние (тензоры, буферы). Для внутреннего состояния, которое не должно отображаться в конструкторе или строковом представлении, используйте поля dataclass с init=False, repr=False.

Реализуйте метод reset() для очистки внутреннего состояния между эпизодами. Это критически важно для процессоров с состоянием, которые накапливают данные со временем, например скользящие средние или временные фильтры.

Помните, что get_config() должен возвращать только конфигурацию, сериализуемую в JSON, тогда как state_dict() отдельно обрабатывает состояние тензоров.

5. Валидация входных данных и обработка ошибок

Проверяйте типы и формы входных данных перед обработкой. Убедитесь в совместимости с вашими алгоритмами, проверив свойства тензоров, такие как dtype и размерности. Для действий робота убедитесь, что необходимые компоненты позы или значения суставов присутствуют и находятся в ожидаемых пределах.

Используйте ранние возвраты для граничных случаев, когда обработка не требуется. Предоставляйте ясные и описательные сообщения об ошибках, указывающие ожидаемые и фактические типы данных или формы. Это значительно упрощает отладку для пользователей.

6. Осведомленность о устройстве и типе данных

Проектируйте процессоры так, чтобы они автоматически адаптировались к устройству и типу данных входных тензоров. Внутренние тензоры (например, статистика нормализации) должны соответствовать устройству и типу данных входного тензора, чтобы обеспечить совместимость с обучением на нескольких GPU, смешанной точностью и распределёнными конфигурациями.

Реализуйте метод to(), который перемещает внутреннее состояние вашего процессора на указанное устройство. Проверяйте совместимость устройства и типа данных во время выполнения и автоматически мигрируйте внутреннее состояние при необходимости. Этот подход обеспечивает бесшовную работу на различных аппаратных конфигурациях без ручного вмешательства.

Заключение

Теперь у вас есть все инструменты для реализации собственных процессоров в LeRobot! Ключевые шаги:

1. Определите ваш процессор в виде dataclass с необходимыми методами (__call__, get_config, state_dict, load_state_dict, reset, transform_features)
2. Зарегистрируйте его с помощью @ProcessorStepRegistry.register("name") для возможности обнаружения
3. Интегрируйте его в DataProcessorPipeline вместе с другими шагами обработки
4. Используйте базовые классы, такие как ObservationProcessorStep, по возможности, чтобы сократить шаблонный код
5. Реализовать учёт устройства/типа данных для поддержки конфигураций с несколькими GPU и смешанной точностью

Система процессоров спроектирована как модульная и компонуемая, что позволяет строить сложные конвейеры обработки данных из простых, сфокусированных компонентов. Независимо от того, предобрабатываете ли вы данные датчиков для обучения или постобрабатываете выходные данные модели для выполнения роботом, настраиваемые процессоры предоставляют гибкость для обработки любых преобразований данных, необходимых вашему приложению робототехники.

Ключевые принципы для надёжных процессоров:

• Адаптация устройства/типа данных: внутренние тензоры должны соответствовать входным тензорам
• Ясные сообщения об ошибках: помогайте пользователям понять, что пошло не так
• Использование базовых классов: Используйте специализированные базовые классы для сокращения шаблонного кода
• Контракты признаков: объявляйте изменения структуры данных с помощью transform_features()

Начните с простого, тщательно тестируйте и убедитесь, что ваши процессоры работают безупречно на разных конфигурациях оборудования!