Thực hiện Robot Processor của riêng bạn

Trong hướng dẫn này, bạn sẽ học cách triển khai Robot Processor của riêng mình. Nội dung bắt đầu bằng việc khám phá nhu cầu về một processor tùy chỉnh, sau đó sử dụng NormalizerProcessorStep làm ví dụ minh họa để giải thích cách triển khai, cấu hình và serialize một processor. Đối với các bước processor đi kèm với LeRobot, xem Các bước Processor phổ biến.

Bạn cần một bộ xử lý tùy chỉnh vì lý do gì?

Trong hầu hết các trường hợp, khi đọc dữ liệu thô từ cảm biến hoặc khi mô hình xuất ra các hành động, bạn cần xử lý dữ liệu này để đảm bảo tương thích với hệ thống đích của bạn. Ví dụ, một nhu cầu phổ biến là chuẩn hóa phạm vi dữ liệu để phù hợp với mạng nơ-ron.

NormalizerProcessorStep của LeRobot xử lý nhiệm vụ quan trọng này:

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

Một số nhu cầu xử lý phổ biến khác bao gồm:

• Đặt thiết bị: Di chuyển tensor giữa CPU/GPU và chuyển đổi kiểu dữ liệu
• Chuyển đổi định dạng: Chuyển đổi giữa các cấu trúc dữ liệu khác nhau
• Batching: Thêm/bỏ các chiều batch để tương thích với mô hình
• Ràng buộc an toàn: Áp dụng các giới hạn cho lệnh điều khiển 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 cung cấp cơ chế pipeline để triển khai các chuỗi bước xử lý cho cả dữ liệu đầu vào và hành động đầu ra, giúp dễ dàng kết hợp các phép biến đổi này theo đúng thứ tự để đạt hiệu suất tối ưu.

Cách triển khai bộ xử lý của riêng bạn?

Chúng ta sẽ sử dụng NormalizerProcessorStep làm ví dụ chính vì nó minh họa các mẫu xử lý viên thiết yếu, bao gồm quản lý trạng thái, chuỗi hóa cấu hình và xử lý tensor mà bạn sẽ thường xuyên cần đến.

Chuẩn bị chuỗi các bước xử lý cần thiết cho bài toán của bạn. Một bước xử lý là một lớp cài đặt các phương thức sau:

• __call__: triển khai bước xử lý cho chuyển tiếp đầu vào.
• get_config: lấy cấu hình của bước xử lý.
• state_dict: lấy trạng thái của bước xử lý.
• load_state_dict: nạp trạng thái của bước xử lý.
• reset: đặt lại trạng thái của bước xử lý.
• transform_features: mô tả cách bước xử lý thay đổi không gian đặc trưng.

Thực hiện phương thức __call__

Phương thức __call__ là phần cốt lõi của bước xử lý của bạn. Nó nhận một EnvTransition và trả về một EnvTransition đã được chỉnh sửa. Dưới đây là cách NormalizerProcessorStep hoạt động:

@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

Xem cài đặt đầy đủ trong src/lerobot/processor/normalize_processor.py để biết chi tiết hoàn chỉnh.

Các nguyên tắc chính:

• Luôn sử dụng transition.copy() để tránh các tác dụng phụ
• Xử lý cả quan sát và hành động một cách nhất quán
• Tách cấu hình khỏi trạng thái: get_config() trả về các tham số có thể serialize thành JSON, state_dict() trả về các tensor
• Chuyển đổi thống kê sang tensor trong __post_init__() để tính toán hiệu quả

Cấu hình và Quản lý Trạng thái

Bộ xử lý hỗ trợ chuỗi hóa thông qua ba phương pháp tách biệt cấu hình khỏi trạng thái tensor. NormalizerProcessorStep minh họa điều này một cách hoàn hảo - nó mang theo thống kê tập dữ liệu (tensor) trong trạng thái của nó, và siêu tham số trong cấu hình của 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
        )
        # ...

Cách sử dụng:

# 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

Biến đổi các đặc trưng

Phương thức transform_features xác định cách bộ xử lý của bạn biến đổi tên và hình dạng của các đặc trưng. Điều này rất quan trọng cho việc cấu hình chính sách và gỡ lỗi.

Đối với NormalizerProcessorStep, các đặc trưng thường được giữ nguyên vì việc chuẩn hóa không làm thay đổi các khóa hay hình dạng:

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

Khi bộ xử lý của bạn đổi tên hoặc thay đổi hình dạng dữ liệu, hãy triển khai phương thức này để phản ánh ánh xạ cho các thành phần phía sau. Ví dụ, một bộ xử lý đổi tên đơn giản:

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

Các nguyên tắc chính:

• Sử dụng features.pop(old_key) để xóa và lấy tính năng cũ
• Sử dụng features[new_key] = old_feature để thêm tính năng đã đổi tên
• Luôn luôn trả về từ điển các đặc trưng đã được sửa đổi
• Mô tả rõ ràng các phép biến đổi tài liệu trong docstring

Sử dụng các ghi đè

Bạn có thể ghi đè các tham số bước khi tải bằng cách sử dụng overrides. Điều này rất hữu ích cho các đối tượng không thể phân chuỗi hoặc các cài đặt cụ thể cho từng vị trí. Nó hoạt động cả trong các nhà máy chính sách và với DataProcessorPipeline.from_pretrained(...).

Điều chỉnh mô hình nền tảng: Điều này đặc biệt hữu ích khi làm việc với các chính sách tiền huấn luyện nền tảng, nơi bạn hiếm khi có quyền truy cập vào các thống kê huấn luyện gốc. Bạn có thể chèn các thống kê tập dữ liệu của riêng mình để điều chỉnh bộ chuẩn hóa phù hợp với dữ liệu robot hoặc môi trường cụ thể của bạn.

Ví dụ: trong quá trình đánh giá chính sách trên robot, ghi đè thiết bị và bản đồ đổi tên. Sử dụng điều này để chạy một chính sách được huấn luyện trên CUDA trên một robot chỉ có CPU, hoặc để ánh xạ lại các khóa camera khi robot sử dụng tên khác với tập dữ liệu.

Sử dụng trực tiếp với 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
        # ...
    },
)

Thực hành tốt nhất

Dựa trên phân tích tất cả các cài đặt bộ xử lý LeRobot, dưới đây là các mẫu và thực hành chính:

1. Xử lý dữ liệu an toàn

Luôn tạo bản sao của dữ liệu đầu vào để tránh các tác dụng phụ không mong muốn. Hãy sử dụng transition.copy() và observation.copy() thay vì sửa đổi dữ liệu tại chỗ. Điều này ngăn chặn bộ xử lý của bạn vô tình ảnh hưởng đến các thành phần khác trong pipeline.

Kiểm tra dữ liệu bắt buộc trước khi xử lý và xử lý dữ liệu thiếu một cách hợp lý. Nếu bộ xử lý của bạn kỳ vọng một số khóa nhất định (chẳng hạn "pixels" cho xử lý ảnh), hãy xác nhận sự hiện diện của chúng trước. Đối với dữ liệu tùy chọn, hãy sử dụng các mẫu truy cập an toàn như transition.get() và xử lý các giá trị None một cách phù hợp.

Khi xác thực dữ liệu thất bại, hãy cung cấp các thông báo lỗi rõ ràng và có thể hành động, giúp người dùng hiểu điều gì đã xảy ra và cách khắc phục.

2. Chọn các lớp cơ sở phù hợp

LeRobot cung cấp các lớp cơ sở chuyên dụng giúp giảm mã lặp lại và đảm bảo tính nhất quán. Sử dụng ObservationProcessorStep khi bạn chỉ cần sửa đổi quan sát, ActionProcessorStep cho xử lý chỉ hành động, và RobotActionProcessorStep cụ thể cho các hành động robot dựa trên từ điển.

Chỉ kế thừa trực tiếp từ ProcessorStep khi bạn cần kiểm soát hoàn toàn toàn bộ quá trình chuyển đổi hoặc khi xử lý nhiều thành phần chuyển đổi đồng thời. Các lớp cơ sở chuyên dụng sẽ quản lý quá trình chuyển đổi cho bạn và cung cấp tính an toàn kiểu dữ liệu.

3. Đăng ký và Đặt tên

Đăng ký các processor của bạn với các tên mô tả, có namespace, bằng cách sử dụng @ProcessorStepRegistry.register(). Hãy dùng các tiền tố tổ chức như "robotics_lab/safety_clipper" hoặc "acme_corp/vision_enhancer" để tránh xung đột tên. Tránh các tên chung chung như "processor" hoặc "step" vì chúng có thể trùng với các triển khai khác.

Việc đăng ký tốt giúp các bộ xử lý của bạn dễ dàng được tìm thấy và cho phép thực hiện serialize/deserialize sạch sẽ khi lưu và tải các pipeline.

4. Các mẫu quản lý trạng thái

Phân biệt giữa các tham số cấu hình (giá trị có thể serialize bằng JSON) và trạng thái nội bộ (tensor, buffer). Sử dụng các trường dataclass với init=False, repr=False cho trạng thái nội bộ không nên xuất hiện trong constructor hoặc biểu diễn chuỗi.

Thực hiện phương thức reset() để xóa trạng thái nội bộ giữa các episode. Điều này rất quan trọng đối với các bộ xử lý có trạng thái (stateful) tích lũy dữ liệu theo thời gian, chẳng hạn như trung bình trượt hoặc bộ lọc thời gian.

Hãy nhớ rằng get_config() chỉ nên trả về cấu hình có thể serialize thành JSON, trong khi state_dict() xử lý trạng thái tensor riêng biệt.

5. Xác thực đầu vào và xử lý lỗi

Kiểm tra kiểu và hình dạng của dữ liệu đầu vào trước khi xử lý. Kiểm tra các thuộc tính của tensor như dtype và số chiều để đảm bảo tương thích với các thuật toán của bạn. Đối với các hành động của robot, xác minh rằng các thành phần tư thế bắt buộc hoặc các giá trị khớp có mặt và nằm trong khoảng giá trị mong đợi.

Sử dụng các lệnh trả về sớm cho các trường hợp biên không cần xử lý. Cung cấp các thông báo lỗi rõ ràng, mô tả chi tiết, bao gồm kiểu dữ liệu hoặc hình dạng mong đợi so với thực tế. Điều này giúp người dùng gỡ lỗi dễ dàng hơn rất nhiều.

6. Nhận thức về thiết bị và kiểu dữ liệu

Thiết kế các bộ xử lý của bạn để tự động thích ứng với thiết bị và dtype của các tensor đầu vào. Các tensor nội bộ (như thống kê chuẩn hóa) cần khớp với thiết bị và dtype của tensor đầu vào để đảm bảo tương thích với huấn luyện đa GPU, độ chính xác hỗn hợp và các cấu hình phân tán.

Thực hiện phương thức to() để di chuyển trạng thái nội bộ của bộ xử lý sang thiết bị được chỉ định. Kiểm tra tính tương thích của thiết bị/dtype khi chạy và tự động chuyển đổi trạng thái nội bộ khi cần. Mẫu thiết kế này cho phép hoạt động liền mạch trên các cấu hình phần cứng khác nhau mà không cần can thiệp thủ công.

Kết luận

Bạn đã có đủ các công cụ để triển khai các bộ xử lý tùy chỉnh trong LeRobot! Các bước chính là:

1. Định nghĩa bộ xử lý của bạn dưới dạng dataclass với các phương thức bắt buộc (__call__, get_config, state_dict, load_state_dict, reset, transform_features)
2. Đăng ký bằng @ProcessorStepRegistry.register("name") để có thể phát hiện
3. Tích hợp nó vào một DataProcessorPipeline cùng với các bước xử lý khác
4. Sử dụng các lớp cơ sở như ObservationProcessorStep khi có thể để giảm bớt mã lặp lại
5. Thực hiện nhận thức về thiết bị/dtype để hỗ trợ các cấu hình đa GPU và độ chính xác hỗn hợp

Hệ thống bộ xử lý được thiết kế theo hướng mô-đun và có thể kết hợp, cho phép bạn xây dựng các pipeline xử lý dữ liệu phức tạp từ các thành phần đơn giản, tập trung. Dù bạn đang tiền xử lý dữ liệu cảm biến để huấn luyện hay hậu xử lý đầu ra của mô hình để robot thực thi, các bộ xử lý tùy chỉnh mang lại cho bạn sự linh hoạt để xử lý bất kỳ phép biến đổi dữ liệu nào mà ứng dụng robotics của bạn yêu cầu.

Các nguyên tắc chính cho bộ xử lý bền vững:

• Điều chỉnh thiết bị/dtype: Các tensor nội bộ phải khớp với các tensor đầu vào
• Thông báo lỗi rõ ràng: Giúp người dùng hiểu điều gì đã xảy ra
• Cách sử dụng lớp cơ sở: Tận dụng các lớp cơ sở chuyên dụng để giảm bớt mã lặp lại
• Hợp đồng tính năng: Khai báo các thay đổi cấu trúc dữ liệu bằng transform_features()

Bắt đầu đơn giản, kiểm tra kỹ lưỡng và đảm bảo các bộ xử lý của bạn hoạt động liền mạch trên các cấu hình phần cứng khác nhau!