직접 Robot Processor 구현하기

이 튜토리얼에서는 자체 Robot Processor를 구현하는 방법을 배웁니다. 먼저 커스텀 프로세서가 필요한 이유를 살펴본 후, NormalizerProcessorStep을 실행 예시로 사용하여 프로세서의 구현, 구성 및 직렬화 방법을 설명합니다. LeRobot에 포함된 프로세서 스텝에 대해서는 Common Processor Steps를 참고하세요.

왜 커스텀 프로세서가 필요할까요?

대부분의 경우, 센서에서 원시 데이터를 읽거나 모델이 동작을 출력할 때, 대상 시스템과 호환되도록 이 데이터를 처리해야 합니다. 예를 들어, 신경망에 적합하도록 데이터 범위를 정규화하는 것은 일반적인 요구 사항입니다.

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__: 입력 전이(input transition)에 대한 처리 단계를 구현합니다.
• 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__()에서 통계를 텐서로 변환하여 효율적인 계산 수행

구성 및 상태 관리

Processor는 구성과 텐서 상태를 분리하는 세 가지 방법을 통해 직렬화를 지원합니다. 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)를 사용하여 기존 feature를 제거하고 가져옵니다
• features[new_key] = old_feature를 사용하여 이름이 변경된 기능을 추가합니다
• features 사전이 수정된 경우 항상 반환
• 문서 문자열에 변환을 명확하게 문서화하세요

오버라이드 사용

overrides를 사용하여 로드 시 step 매개변수를 재정의할 수 있습니다. 이는 직렬화할 수 없는 객체나 사이트별 설정에 유용합니다. 정책 팩토리와 DataProcessorPipeline.from_pretrained(...) 모두에서 작동합니다.

기초 모델 적응: 이는 주로 기초 사전 학습된 정책을 다룰 때 특히 유용합니다. 원래의 학습 통계에 접근할 수 있는 경우가 드물기 때문입니다. 자체 데이터셋의 통계를 주입하여 정규화기를 특정 로봇이나 환경 데이터에 맞게 적응시킬 수 있습니다.

예: 로봇에서 정책 평가 중에는 device와 rename map을 오버라이드할 수 있습니다. 이를 통해 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. 안전한 데이터 처리

의도하지 않은 부수 효과를 피하려면 항상 입력 데이터의 복사본을 생성하세요. 데이터를 인플레이스(in-place)로 수정하는 대신 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() 메서드를 구현하세요. 이동 평균이나 시간 필터처럼 시간이 지남에 따라 데이터를 축적하는 상태ful 프로세서의 경우 이것이 매우 중요합니다.

get_config()은 JSON 직렬화 가능한 구성만 반환해야 하며, 텐서 상태는 state_dict()가 별도로 처리한다는 점을 기억하세요.

5. 입력 검증 및 오류 처리

처리 전에 입력의 유형과 형태를 검증하세요. 알고리즘과의 호환성을 보장하기 위해 dtype 및 차원과 같은 텐서 속성을 확인하세요. 로봇 액션의 경우, 필요한 자세 구성 요소나 관절 값이 존재하고 예상 범위 내에 있는지 확인하세요.

처리할 필요가 있는 엣지 케이스에는 조기 반환을 사용하세요. 예상 데이터 유형 또는 형태와 실제 데이터 유형 또는 형태를 포함하는 명확하고 설명적인 오류 메시지를 제공하세요. 이를 통해 사용자의 디버깅이 훨씬 쉬워집니다.

6. Device 및 Dtype 인식

프로세서를 입력 텐서의 디바이스와 dtype에 자동으로 적응하도록 설계하세요. 내부 텐서(정규화 통계값 등)는 입력 텐서의 디바이스와 dtype과 일치해야 멀티-GPU 학습, 혼합 정밀도, 분산 환경과의 호환성을 보장할 수 있습니다.

프로세서의 내부 상태를 지정된 장치로 이동시키는 to() 메서드를 구현하세요. 런타임에 장치/dtype 호환성을 확인하고 필요 시 내부 상태를 자동으로 마이그레이션합니다. 이 패턴을 사용하면 수동 개입 없이 다양한 하드웨어 구성에서 원활하게 작동할 수 있습니다.

결론

이제 LeRobot에서 커스텀 프로세서(Processor)를 구현하는 데 필요한 모든 도구를 갖추셨습니다. 주요 단계는 다음과 같습니다:

1. 필수 메서드(__call__, get_config, state_dict, load_state_dict, reset, transform_features)를 가진 dataclass로 프로세서 정의
2. 발견 가능성을 높이기 위해 @ProcessorStepRegistry.register("name")를 사용하여 등록합니다
3. DataProcessorPipeline에 다른 처리 단계와 함께 통합
4. ObservationProcessorStep와 같은 기본 클래스를 가능한 한 사용하여 불필요한 코드를 줄이세요
5. 멀티 GPU 및 혼합 정밀도 구성을 지원하기 위해 디바이스/dtype 인식 구현

프로세서 시스템은 모듈형 및 구성 가능한 구조로 설계되어, 단순하고 목적에 집중된 구성 요소로 복잡한 데이터 처리 파이프라인을 구축할 수 있습니다. 학습을 위한 센서 데이터 전처리부터 로봇 실행을 위한 모델 출력 후처리까지, 커스텀 프로세서를 사용하면 로봇 애플리케이션에 필요한 모든 데이터 변환을 유연하게 처리할 수 있습니다.

견고한 프로세서의 핵심 원칙:

• 디바이스/dtype 적응: 내부 텐서는 입력 텐서와 일치해야 합니다
• 명확한 오류 메시지: 사용자가 무엇이 잘못되었는지 이해할 수 있도록 돕습니다
• 기본 클래스 사용: 전문화된 기본 클래스를 활용하여 불필요한 코드를 줄입니다
• 기능 계약: transform_features()로 데이터 구조 변경을 선언

간단하게 시작하고, 철저히 테스트하며, 프로세서가 다양한 하드웨어 구성에서 원활하게 작동하도록 확인하세요!