Implemente seu próprio Robot Processor

Neste tutorial, você aprenderá a implementar seu próprio Robot Processor. Ele começa explorando a necessidade de um processador personalizado e, em seguida, usa o NormalizerProcessorStep como exemplo prático para explicar como implementar, configurar e serializar um processador. Para as etapas de processador que acompanham o LeRobot, consulte Common Processor Steps.

Por que você precisaria de um processador personalizado?

Na maioria dos casos, ao ler dados brutos de sensores ou quando os modelos produzem ações, é necessário processar esses dados para torná-los compatíveis com o sistema de destino. Por exemplo, uma necessidade comum é normalizar os intervalos de dados para que sejam adequados para redes neurais.

O NormalizerProcessorStep do LeRobot trata dessa tarefa 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)
# ...

Outras necessidades comuns de processamento incluem:

• Posicionamento de dispositivos: mover tensores entre CPU/GPU e converter tipos de dados
• Conversão de formato: Transformação entre diferentes estruturas de dados
• Loteamento: Adicionar/remover dimensões de lote para compatibilidade do modelo
• Restrições de segurança: Aplicando limites a comandos de robô

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

LeRobot fornece um mecanismo de pipeline para implementar sequências de etapas de processamento para dados de entrada e ações de saída, facilitando a composição dessas transformações na ordem correta para desempenho ideal.

Como implementar seu próprio processador?

Usaremos o NormalizerProcessorStep como nosso exemplo principal, pois ele demonstra padrões essenciais de processadores, incluindo gerenciamento de estado, serialização de configuração e manipulação de tensores, que você precisará com frequência.

Prepare a sequência de etapas de processamento necessárias para o seu problema. Uma etapa de processador é uma classe que implementa os seguintes métodos:

• __call__: implementa a etapa de processamento para a transição de entrada.
• get_config: obtém a configuração da etapa do processador.
• state_dict: obtém o estado da etapa do processador.
• load_state_dict: carrega o estado da etapa do processador.
• reset: redefine o estado da etapa do processador.
• transform_features: descreve como a etapa do processador modifica o espaço de características.

Implemente o método __call__

O método __call__ é o núcleo da sua etapa de processador. Ele recebe um EnvTransition e retorna um EnvTransition modificado. Veja como o NormalizerProcessorStep funciona:

@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

Veja a implementação completa em src/lerobot/processor/normalize_processor.py para todos os detalhes.

Princípios fundamentais:

• Use sempre transition.copy() para evitar efeitos colaterais
• Tratar observações e ações de forma consistente
• Separar configuração de estado: get_config() retorna parâmetros serializáveis em JSON, state_dict() retorna tensores
• Converter estatísticas para tensores em __post_init__() para computação eficiente

Configuração e Gerenciamento de Estado

Os processadores suportam serialização por meio de três métodos que separam a configuração do estado dos tensores. O NormalizerProcessorStep demonstra isso perfeitamente - ele carrega as estatísticas do conjunto de dados (tensores) em seu estado e os hiperparâmetros em sua configuração:

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

O método transform_features define como o seu processador transforma os nomes e as formas dos recursos. Isso é crucial para a configuração da política e para a depuração.

Para NormalizerProcessorStep, os recursos são tipicamente preservados sem alterações, pois a normalização não altera chaves ou 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
    # ...

Quando o seu processor renomeia ou reorganiza dados, implemente este método para refletir o mapeamento para os componentes a jusante. Por exemplo, um processor de renomeação simples:

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

Princípios fundamentais:

• Use features.pop(old_key) para remover e obter o recurso antigo
• Use features[new_key] = old_feature para adicionar o recurso renomeado
• Sempre retorne o dicionário de recursos modificado
• Documente as transformações de dados claramente na docstring

Usando substituições

Você pode sobrescrever os parâmetros de etapa no momento do carregamento usando overrides. Isso é útil para objetos não serializáveis ou configurações específicas do site. Funciona tanto em fábricas de política quanto com DataProcessorPipeline.from_pretrained(...).

Adaptação de modelos fundacionais: isso é particularmente útil ao trabalhar com políticas pré-treinadas fundacionais, onde raramente se tem acesso às estatísticas de treinamento originais. Você pode injetar as estatísticas do seu próprio conjunto de dados para adaptar o normalizador aos dados específicos do seu robô ou ambiente.

Exemplo: durante a avaliação da política no robô, substitua o dispositivo e o mapa de renomeação. Use isso para executar uma política treinada em CUDA em um robô somente com CPU, ou para remapear as chaves das câmeras quando o robô usa nomes diferentes dos do conjunto de dados.

Uso direto com 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
        # ...
    },
)

Boas Práticas

Com base na análise de todas as implementações de processadores do LeRobot, aqui estão os padrões e práticas principais:

1. Manipulação segura de dados

Sempre crie cópias dos dados de entrada para evitar efeitos colaterais indesejados. Use transition.copy() e observation.copy() em vez de modificar os dados no local. Isso impede que seu processador afete acidentalmente outros componentes no pipeline.

Verifique os dados obrigatórios antes do processamento e trate dados ausentes de forma adequada. Se o seu processador espera certas chaves (como "pixels" para processamento de imagens), valide primeiro a presença delas. Para dados opcionais, use padrões de acesso seguros como transition.get() e trate valores None de forma apropriada.

Quando a validação de dados falha, forneça mensagens de erro claras e acionáveis que ajudem os usuários a entender o que deu errado e como corrigir.

2. Escolha as classes base apropriadas

LeRobot fornece classes base especializadas que reduzem o código repetitivo e garantem consistência. Use ObservationProcessorStep quando você precisar apenas modificar observações, ActionProcessorStep para processamento apenas de ações e RobotActionProcessorStep especificamente para ações de robô baseadas em dicionários.

Herdite diretamente de ProcessorStep apenas quando precisar de controle total sobre a transição inteira ou ao processar múltiplos componentes de transição simultaneamente. As classes base especializadas gerenciam a transição para você e fornecem segurança de tipos.

3. Registro e Nomenclatura

Registre seus processadores com nomes descritivos e namespaced usando @ProcessorStepRegistry.register(). Use prefixos de organização como "robotics_lab/safety_clipper" ou "acme_corp/vision_enhancer" para evitar conflitos de nomenclatura. Evite nomes genéricos como "processor" ou "step" que podem conflitar com outras implementações.

Uma boa inscrição torna seus processadores descobríveis e permite uma serialização/desserialização limpa ao salvar e carregar pipelines.

4. Padrões de Gerenciamento de Estado

Distingua entre parâmetros de configuração (valores serializáveis em JSON) e estado interno (tensors, buffers). Use campos de dataclass com init=False, repr=False para o estado interno que não deve aparecer no construtor ou na representação em string.

Implemente o método reset() para limpar o estado interno entre episódios. Isso é crucial para processadores com estado que acumulam dados ao longo do tempo, como médias móveis ou filtros temporais.

Lembre-se de que get_config() deve retornar apenas configuração serializável em JSON, enquanto state_dict() trata o estado dos tensores separadamente.

5. Validação de entrada e tratamento de erros

Valide os tipos e as formas de entrada antes do processamento. Verifique as propriedades dos tensores, como dtype e dimensões, para garantir a compatibilidade com seus algoritmos. Para ações de robôs, verifique se os componentes de pose necessários ou os valores das juntas estão presentes e dentro dos intervalos esperados.

Use retornos antecipados para casos extremos em que nenhum processamento é necessário. Forneça mensagens de erro claras e descritivas que incluam os tipos ou formatos de dados esperados versus os reais. Isso torna o depuramento muito mais fácil para os usuários.

6. Consciência de Dispositivo e Dtype

Projete seus processadores para se adaptarem automaticamente ao dispositivo e ao dtype dos tensores de entrada. Os tensores internos (como estatísticas de normalização) devem corresponder ao dispositivo e ao dtype do tensor de entrada para garantir compatibilidade com treinamento multi-GPU, precisão mista e configurações distribuídas.

Implemente um método to() que move o estado interno do seu processador para o dispositivo especificado. Verifique a compatibilidade de dispositivo/dtype em tempo de execução e migre automaticamente o estado interno quando necessário. Esse padrão permite uma operação sem problemas em diferentes configurações de hardware, sem intervenção manual.

Conclusão

Agora você tem todas as ferramentas para implementar processadores personalizados no LeRobot! As etapas principais são:

1. Defina seu processador como um dataclass com os métodos necessários (__call__, get_config, state_dict, load_state_dict, reset, transform_features)
2. Registre-o usando @ProcessorStepRegistry.register("name") para facilitar a descoberta
3. Integre-o em um DataProcessorPipeline com outras etapas de processamento
4. Use classes base como ObservationProcessorStep sempre que possível para reduzir código repetitivo
5. Implementar consciência de dispositivo/dtype para suportar configurações multi-GPU e de precisão mista

O sistema de processadores é projetado para ser modular e composível, permitindo que você construa pipelines de processamento de dados complexos a partir de componentes simples e focados. Seja no pré-processamento de dados de sensores para treinamento ou no pós-processamento das saídas do modelo para a execução do robô, os processadores personalizados oferecem a flexibilidade necessária para lidar com qualquer transformação de dados exigida pela sua aplicação de robótica.

Princípios-chave para processadores robustos:

• Adaptação de dispositivo/dtype: os tensores internos devem corresponder aos tensores de entrada
• Mensagens de erro claras: Ajude os usuários a entender o que deu errado
• Uso de classes base: Aproveite classes base especializadas para reduzir código repetitivo
• Contratos de recursos: Declare mudanças na estrutura de dados com transform_features()

Comece de forma simples, teste minuciosamente e garanta que seus processadores funcionem perfeitamente em diferentes configurações de hardware!