Implementa tu propio Robot Processor

En este tutorial aprenderás a implementar tu propio Robot Processor. Comienza explorando la necesidad de un procesador personalizado y luego utiliza NormalizerProcessorStep como ejemplo práctico para explicar cómo implementar, configurar y serializar un procesador. Para los pasos de procesador que se incluyen con LeRobot, consulta Common Processor Steps.

¿Por qué necesitarías un procesor personalizado?

En la mayoría de los casos, al leer datos crudos de sensores o cuando los modelos generan acciones, es necesario procesar estos datos para hacerlos compatibles con el sistema objetivo. Por ejemplo, una necesidad común es normalizar los rangos de datos para que sean adecuados para las redes neuronales.

NormalizerProcessorStep de LeRobot se encarga de esta tarea crucial:

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

Otras necesidades de procesamiento habituales incluyen:

• Colocación de dispositivos: mover tensores entre CPU/GPU y convertir tipos de datos
• Conversión de formato: Transformación entre diferentes estructuras de datos
• Batching: Añadir/quitar dimensiones de lote para la compatibilidad del modelo
• Restricciones de seguridad: Aplicar límites a los comandos del robot

# Example pipeline combining multiple processors
pipeline = PolicyProcessorPipeline([
    RenameObservationsProcessorStep(rename_map={}),
    AddBatchDimensionProcessorStep(),
    NormalizerProcessorStep(features=features, norm_map=norm_map, stats=stats),
    DeviceProcessorStep(device="cuda"),
    # ...
])

LeRobot proporciona un mecanismo de pipeline para implementar secuencias de pasos de procesamiento tanto para los datos de entrada como para las acciones de salida, lo que facilita componer estas transformaciones en el orden correcto para un rendimiento óptimo.

Cómo implementar su propio procesador?

Utilizaremos NormalizerProcessorStep como ejemplo principal, ya que demuestra patrones esenciales de procesador, como la gestión de estado, la serialización de configuración y el manejo de tensores, que necesitarás con frecuencia.

Prepara la secuencia de pasos de procesamiento necesarios para tu problema. Un paso de procesador es una clase que implementa los siguientes métodos:

• __call__: implementa el paso de procesamiento para la transición de entrada.
• get_config: obtiene la configuración del paso del procesador.
• state_dict: obtiene el estado del paso del procesador.
• load_state_dict: carga el estado del paso del procesador.
• reset: restablece el estado del paso del procesador.
• transform_features: describe cómo el paso del procesador modifica el espacio de características.

Implementar el método __call__

El método __call__ es el núcleo de su paso de procesador. Toma un EnvTransition y devuelve un EnvTransition modificado. Así funciona el 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

Vea la implementación completa en src/lerobot/processor/normalize_processor.py para obtener todos los detalles.

Principios clave:

• Usa siempre transition.copy() para evitar efectos secundarios
• Manejar tanto las observaciones como las acciones de forma coherente
• Separar la configuración del estado: get_config() devuelve parámetros serializables en JSON, state_dict() devuelve tensores
• Convertir las estadísticas en tensores en __post_init__() para un cálculo eficiente

Configuración y gestión del estado

Los procesadores admiten la serialización a través de tres métodos que separan la configuración del estado de los tensores. El NormalizerProcessorStep lo demuestra a la perfección: lleva las estadísticas del conjunto de datos (tensores) en su estado y los hiperparámetros en su configuración:

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

Uso:

# 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

Transformar características

El método transform_features define cómo su procesador transforma los nombres y las formas de las características. Esto es crucial para la configuración de la política y la depuración.

Para NormalizerProcessorStep, las características suelen conservarse sin cambios, ya que la normalización no altera las claves ni las formas:

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

Cuando su procesador cambia el nombre o la forma de los datos, implemente este método para reflejar el mapeo para los componentes posteriores. Por ejemplo, un procesador de cambio de nombre simple:

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

Principios clave:

• Use features.pop(old_key) para eliminar y obtener la característica anterior
• Use features[new_key] = old_feature para añadir la característica renombrada
• Devuelve siempre el diccionario de características modificado
• Documenta las transformaciones de forma clara en la docstring

Uso de reemplazos

Puede anular los parámetros de paso en el momento de la carga mediante overrides. Esto resulta útil para objetos no serializables o configuraciones específicas del sitio. Funciona tanto en las fábricas de políticas como con DataProcessorPipeline.from_pretrained(...).

Adaptación de modelos fundamentales: Esto resulta especialmente útil al trabajar con políticas preentrenadas fundamentales, ya que rara vez se tiene acceso a las estadísticas de entrenamiento originales. Puedes inyectar las estadísticas de tu propio conjunto de datos para adaptar el normalizador a los datos específicos de tu robot o entorno.

Ejemplo: durante la evaluación de la política en el robot, sobrescribe el dispositivo y el mapa de nombres. Usa esto para ejecutar una política entrenada en CUDA en un robot solo con CPU, o para remapear las claves de las cámaras cuando el robot usa nombres diferentes a los del conjunto de datos.

Uso directo con 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
        # ...
    },
)

Buenas prácticas

Basado en el análisis de todas las implementaciones de procesadores de LeRobot, estos son los patrones y prácticas clave:

1. Manejo seguro de datos

Cree siempre copias de los datos de entrada para evitar efectos secundarios no deseados. Utilice transition.copy() y observation.copy() en lugar de modificar los datos in situ. Esto impide que su procesador afecte accidentalmente a otros componentes de la tubería.

Comprueba los datos obligatorios antes de procesarlos y gestiona los datos faltantes de forma adecuada. Si tu procesador espera ciertas claves (como "pixels" para el procesamiento de imágenes), valida primero su presencia. Para los datos opcionales, utiliza patrones de acceso seguros como transition.get() y gestiona los valores None de manera apropiada.

Cuando la validación de datos falla, proporciona mensajes de error claros y accionables que ayuden a los usuarios a comprender qué salió mal y cómo solucionarlo.

2. Elegir las clases base adecuadas

LeRobot proporciona clases base especializadas que reducen el código repetitivo y garantizan la consistencia. Utilice ObservationProcessorStep cuando solo necesite modificar observaciones, ActionProcessorStep para el procesamiento exclusivo de acciones y RobotActionProcessorStep específicamente para acciones de robot basadas en diccionarios.

Hereda directamente de ProcessorStep solo cuando necesites un control total sobre la transición completa o cuando proceses varios componentes de transición simultáneamente. Las clases base especializadas gestionan la transición por ti y proporcionan seguridad de tipos.

3. Registro y Nomenclatura

Registre sus procesadores con nombres descriptivos y con espacio de nombres usando @ProcessorStepRegistry.register(). Utilice prefijos de organización como "robotics_lab/safety_clipper" o "acme_corp/vision_enhancer" para evitar conflictos de nombres. Evite nombres genéricos como "processor" o "step" que podrían entrar en conflicto con otras implementaciones.

Un registro correcto hace que tus procesadores sean detectables y permite una serialización/deserialización limpia al guardar y cargar pipelines.

4. Patrones de gestión de estado

Distinga entre los parámetros de configuración (valores serializables en JSON) y el estado interno (tensores, búferes). Utilice campos de dataclass con init=False, repr=False para el estado interno que no debe aparecer en el constructor ni en la representación en cadena.

Implemente el método reset() para limpiar el estado interno entre episodios. Esto es crucial para los procesadores con estado que acumulan datos con el tiempo, como las medias móviles o los filtros temporales.

Recuerde que get_config() solo debe devolver una configuración serializable en JSON, mientras que state_dict() maneja el estado de los tensores por separado.

5. Validación de entrada y manejo de errores

Valide los tipos y formas de entrada antes de procesarlos. Compruebe las propiedades de los tensores, como dtype y las dimensiones, para garantizar la compatibilidad con sus algoritmos. Para las acciones del robot, verifique que los componentes de pose requeridos o los valores de las articulaciones estén presentes y dentro de los rangos esperados.

Utilice retornos anticipados para casos límite en los que no se requiere procesamiento. Proporcione mensajes de error claros y descriptivos que incluyan los tipos o formas de datos esperados frente a los reales. Esto facilita mucho la depuración para los usuarios.

6. Conciencia del dispositivo y el Dtype

Diseñen sus procesadores para que se adapten automáticamente al dispositivo y al dtype de los tensores de entrada. Los tensores internos (como las estadísticas de normalización) deben coincidir con el dispositivo y el dtype del tensor de entrada para garantizar la compatibilidad con el entrenamiento multi-GPU, la precisión mixta y las configuraciones distribuidas.

Implementa un método to() que mueva el estado interno de tu procesador al dispositivo especificado. Verifica la compatibilidad de dispositivo/dtype en tiempo de ejecución y migra automáticamente el estado interno cuando sea necesario. Este patrón permite una operación sin problemas en diferentes configuraciones de hardware sin intervención manual.

Conclusión

Ahora tienes todas las herramientas para implementar procesadores personalizados en LeRobot. Los pasos clave son:

1. Defina su procesador como un dataclass con los métodos requeridos (__call__, get_config, state_dict, load_state_dict, reset, transform_features)
2. Regístrelo usando @ProcessorStepRegistry.register("name") para que sea descubreble
3. Integrarlo en un DataProcessorPipeline con otras etapas de procesamiento
4. Use clases base como ObservationProcessorStep cuando sea posible para reducir el código repetitivo
5. Implementar la toma de conciencia del dispositivo/dtype para soportar configuraciones de GPU múltiple y precisión mixta

El sistema de procesadores está diseñado para ser modular y componible, lo que permite construir pipelines de procesamiento de datos complejos a partir de componentes simples y enfocados. Ya sea que esté preprocesando datos de sensores para el entrenamiento o postprocesando las salidas del modelo para la ejecución del robot, los procesadores personalizados le brindan la flexibilidad para manejar cualquier transformación de datos que su aplicación de robótica requiera.

Principios clave para procesadores robustos:

• Adaptación de dispositivo/dtype: Los tensores internos deben coincidir con los tensores de entrada
• Mensajes de error claros: Ayuda a los usuarios a entender qué salió mal
• Uso de clases base: Aprovecha las clases base especializadas para reducir el código repetitivo
• Contratos de características: Declara cambios en la estructura de los datos con transform_features()

Comienza de forma sencilla, realiza pruebas exhaustivas y asegúrate de que tus procesadores funcionen sin problemas en diferentes configuraciones de hardware.