mirror of
https://github.com/huggingface/lerobot.git
synced 2026-06-01 03:11:29 +00:00
a6dbb65917
* feat(processor): introduce PolicyProcessorPipeline and RobotProcessorPipeline as type aliases for DataProcessorPipeline - Added PolicyProcessorPipeline and RobotProcessorPipeline type aliases to enhance clarity and maintainability in the processor module. - Updated the __all__ list to include the new pipelines for better module export consistency. * refactor(processor): replace DataProcessorPipeline with PolicyProcessorPipeline across multiple modules - Updated all instances of DataProcessorPipeline to PolicyProcessorPipeline in various processor files for consistency and clarity. - Adjusted function signatures to reflect the new pipeline type, enhancing maintainability and readability. * refactor(processor): update hotswap_stats function to use PolicyProcessorPipeline - Changed the parameter name from robot_processor to policy_processor for clarity. - Ensured consistency with recent updates to the processor module by reflecting the new pipeline type in the function signature. * refactor(processor): replace DataProcessorPipeline with PolicyProcessorPipeline in migrate_policy_normalization.py - Updated the preprocessor and postprocessor to use PolicyProcessorPipeline for consistency with recent changes in the processor module. - Enhanced clarity and maintainability by aligning with the new pipeline structure. * refactor(processor): update hotswap_stats to use PolicyProcessorPipeline - Changed the parameter type in hotswap_stats from DataProcessorPipeline to PolicyProcessorPipeline for consistency with recent updates. - Enhanced clarity by updating the function documentation to reflect the new pipeline type. * refactor(processor): replace DataProcessorPipeline with RobotProcessorPipeline across multiple files - Updated instances of DataProcessorPipeline to RobotProcessorPipeline in evaluate.py, record.py, replay.py, teleoperate.py, and other relevant files for consistency and clarity. - Adjusted function signatures and variable types to reflect the new pipeline structure, enhancing maintainability and readability.
71 lines
2.4 KiB
Python
71 lines
2.4 KiB
Python
#!/usr/bin/env python
|
|
|
|
# Copyright 2024 Columbia Artificial Intelligence, Robotics Lab,
|
|
# and The HuggingFace Inc. team. All rights reserved.
|
|
#
|
|
# Licensed under the Apache License, Version 2.0 (the "License");
|
|
# you may not use this file except in compliance with the License.
|
|
# You may obtain a copy of the License at
|
|
#
|
|
# http://www.apache.org/licenses/LICENSE-2.0
|
|
#
|
|
# Unless required by applicable law or agreed to in writing, software
|
|
# distributed under the License is distributed on an "AS IS" BASIS,
|
|
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
# See the License for the specific language governing permissions and
|
|
# limitations under the License.
|
|
import torch
|
|
|
|
from lerobot.constants import POSTPROCESSOR_DEFAULT_NAME, PREPROCESSOR_DEFAULT_NAME
|
|
from lerobot.policies.diffusion.configuration_diffusion import DiffusionConfig
|
|
from lerobot.processor import (
|
|
AddBatchDimensionProcessorStep,
|
|
DeviceProcessorStep,
|
|
NormalizerProcessorStep,
|
|
PolicyProcessorPipeline,
|
|
ProcessorKwargs,
|
|
RenameProcessorStep,
|
|
UnnormalizerProcessorStep,
|
|
)
|
|
|
|
|
|
def make_diffusion_pre_post_processors(
|
|
config: DiffusionConfig,
|
|
dataset_stats: dict[str, dict[str, torch.Tensor]] | None = None,
|
|
preprocessor_kwargs: ProcessorKwargs | None = None,
|
|
postprocessor_kwargs: ProcessorKwargs | None = None,
|
|
) -> tuple[PolicyProcessorPipeline, PolicyProcessorPipeline]:
|
|
if preprocessor_kwargs is None:
|
|
preprocessor_kwargs = {}
|
|
if postprocessor_kwargs is None:
|
|
postprocessor_kwargs = {}
|
|
|
|
input_steps = [
|
|
RenameProcessorStep(rename_map={}),
|
|
NormalizerProcessorStep(
|
|
features={**config.input_features, **config.output_features},
|
|
norm_map=config.normalization_mapping,
|
|
stats=dataset_stats,
|
|
),
|
|
AddBatchDimensionProcessorStep(),
|
|
DeviceProcessorStep(device=config.device),
|
|
]
|
|
output_steps = [
|
|
DeviceProcessorStep(device="cpu"),
|
|
UnnormalizerProcessorStep(
|
|
features=config.output_features, norm_map=config.normalization_mapping, stats=dataset_stats
|
|
),
|
|
]
|
|
return (
|
|
PolicyProcessorPipeline(
|
|
steps=input_steps,
|
|
name=PREPROCESSOR_DEFAULT_NAME,
|
|
**preprocessor_kwargs,
|
|
),
|
|
PolicyProcessorPipeline(
|
|
steps=output_steps,
|
|
name=POSTPROCESSOR_DEFAULT_NAME,
|
|
**postprocessor_kwargs,
|
|
),
|
|
)
|