独自のRobot Processorを実装する

このチュートリアルでは、独自のRobot Processorを実装する方法を学びます。 まず、カスタムプロセッサが必要とされる背景を調べ、その後、NormalizerProcessorStepを例として、プロセッサの実装、設定、シリアライズの方法を説明します。LeRobotに同梱されているプロセッサステップについては、共通のプロセッサステップを参照してください。

なぜカスタムプロセッサが必要になるのでしょうか？

ほとんどの場合、センサーから生データを読み込む際や、モデルがアクションを出力する際、ターゲットシステムとの互換性を確保するためにこのデータを処理する必要があります。例えば、ニューラルネットワークに適した形にするためにデータ範囲を正規化することは、よくあるニーズです。

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__: 入力遷移の処理ステップを実装します。
• 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__() 内で統計情報をテンソルに変換し、効率的な計算を実現します

設定と状態管理

プロセッサは、設定とテンソル状態を分離する3つの方法によるシリアライズをサポートしています。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 を使用して、名前変更された特徴量を追加します
• 変更されたfeatures辞書を常に返します
• 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 シリアライズ可能な値）と内部状態（テンソル、バッファ）を区別してください。コンストラクターや文字列表現に表示すべきでない内部状態には、init=False, repr=False を指定した dataclass フィールドを使用してください。

エピソード間で内部状態をクリアするためにreset()メソッドを実装します。これは、移動平均や時間フィルタのように時間とともにデータを蓄積する状態付きプロセッサにとって重要です。

get_config()はJSONシリアライズ可能な設定のみを返すべきであり、テンソル状態はstate_dict()が別途処理することを覚えておいてください。

5. 入力検証とエラー処理

処理前に、入力データの型と形状を検証してください。dtypeや次元などのテンソルプロパティを確認し、アルゴリズムとの互換性を確保します。ロボットのアクションについては、必要な姿勢成分や関節値が存在し、想定された範囲内にあることを確認してください。

処理が不要なエッジケースには、早期リターンを使用してください。期待されるデータ型や形状と実際のデータ型や形状を含む、明確で説明的なエラーメッセージを提供してください。これにより、ユーザーのデバッグが大幅に容易になります。

6. デバイスと Dtype の認識

プロセッサは、入力テンソルのデバイスとdtypeに自動的に適応するように設計してください。内部テンソル（正規化統計量など）は、マルチGPUトレーニング、混合精度、分散環境との互換性を確保するために、入力テンソルのデバイスとdtypeと一致させる必要があります。

プロセッサの内部状態を指定されたデバイスに移動するto()メソッドを実装します。実行時にデバイス/dtypeの互換性を確認し、必要に応じて内部状態を自動的に移行します。このパターンにより、手動での介入なしに、異なるハードウェア構成間でシームレスに動作させることができます。

結論

これで、LeRobotでカスタムプロセッサを実装するためのすべてのツールが揃いました。主な手順は以下の通りです。

1. プロセッサを、必要なメソッド（__call__、get_config、state_dict、load_state_dict、reset、transform_features）を含むデータクラスとして定義します。
2. 登録するには、検索性のために@ProcessorStepRegistry.register("name")を使用します
3. DataProcessorPipelineに他の処理ステップと組み合わせて統合する
4. ObservationProcessorStep などのベースクラスを可能な限り使用し、ボイラープレートコードを削減してください
5. マルチGPUおよび混合精度構成をサポートするために、デバイス/データ型への対応を実装する

プロセッサシステムはモジュール化と組み合わせやすさを前提に設計されており、シンプルで目的に特化したコンポーネントから複雑なデータ処理パイプラインを構築できます。センサーデータの事前処理や、ロボット実行のためのモデル出力の事後処理など、ロボティクスアプリケーションに必要なあらゆるデータ変換に対応する柔軟性を、カスタムプロセッサが提供します。

堅牢なプロセッサのための主要な原則:

• デバイス/dtypeの適応: 内部テンソルは入力テンソルと一致させる必要があります
• 明確なエラーメッセージ: ユーザーが何が問題だったかを理解できるようにします
• ベースクラスの使用方法: 専用のベースクラスを活用して、ボイラープレートコードを削減します
• 機能契約: transform_features() を使用してデータ構造の変更を宣言します

シンプルに始め、徹底的にテストし、プロセッサが異なるハードウェア構成でシームレスに動作することを確認してください！