mirror of
https://github.com/huggingface/lerobot.git
synced 2026-06-03 04:11:24 +00:00
fd4ae3466b
* chore(pipeline): remove unused features + device torch + envtransition keys * refactor(pipeline): ImageProcessor & StateProcessor are both implemented directly in VanillaObservationPRocessor * refactor(pipeline): RenameProcessor now inherits from ObservationProcessor + remove unused code * test(pipeline): fix broken test after refactors * docs(pipeline): update docstrings VanillaObservationProcessor * chore(pipeline): move None check to base pipeline classes
1265 lines
48 KiB
Python
1265 lines
48 KiB
Python
#!/usr/bin/env python
|
||
|
||
# Copyright 2025 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.
|
||
from __future__ import annotations
|
||
|
||
import importlib
|
||
import json
|
||
import os
|
||
from collections.abc import Callable, Iterable, Sequence
|
||
from copy import deepcopy
|
||
from dataclasses import dataclass, field
|
||
from enum import Enum
|
||
from pathlib import Path
|
||
from typing import Any, Protocol, TypedDict
|
||
|
||
import torch
|
||
from huggingface_hub import ModelHubMixin, hf_hub_download
|
||
from huggingface_hub.errors import HfHubHTTPError
|
||
from safetensors.torch import load_file, save_file
|
||
|
||
from lerobot.configs.types import PolicyFeature
|
||
|
||
|
||
class TransitionKey(str, Enum):
|
||
"""Keys for accessing EnvTransition dictionary components."""
|
||
|
||
# TODO(Steven): Use consts
|
||
OBSERVATION = "observation"
|
||
ACTION = "action"
|
||
REWARD = "reward"
|
||
DONE = "done"
|
||
TRUNCATED = "truncated"
|
||
INFO = "info"
|
||
COMPLEMENTARY_DATA = "complementary_data"
|
||
|
||
|
||
EnvTransition = TypedDict(
|
||
"EnvTransition",
|
||
{
|
||
TransitionKey.OBSERVATION.value: dict[str, Any] | None,
|
||
TransitionKey.ACTION.value: Any | torch.Tensor | None,
|
||
TransitionKey.REWARD.value: float | torch.Tensor | None,
|
||
TransitionKey.DONE.value: bool | torch.Tensor | None,
|
||
TransitionKey.TRUNCATED.value: bool | torch.Tensor | None,
|
||
TransitionKey.INFO.value: dict[str, Any] | None,
|
||
TransitionKey.COMPLEMENTARY_DATA.value: dict[str, Any] | None,
|
||
},
|
||
)
|
||
|
||
|
||
class ProcessorStepRegistry:
|
||
"""Registry for processor steps that enables saving/loading by name instead of module path."""
|
||
|
||
_registry: dict[str, type] = {}
|
||
|
||
@classmethod
|
||
def register(cls, name: str = None):
|
||
"""Decorator to register a processor step class.
|
||
|
||
Args:
|
||
name: Optional registration name. If not provided, uses class name.
|
||
|
||
Example:
|
||
@ProcessorStepRegistry.register("adaptive_normalizer")
|
||
class AdaptiveObservationNormalizer:
|
||
...
|
||
"""
|
||
|
||
def decorator(step_class: type) -> type:
|
||
registration_name = name if name is not None else step_class.__name__
|
||
|
||
if registration_name in cls._registry:
|
||
raise ValueError(
|
||
f"Processor step '{registration_name}' is already registered. "
|
||
f"Use a different name or unregister the existing one first."
|
||
)
|
||
|
||
cls._registry[registration_name] = step_class
|
||
# Store the registration name on the class for later reference
|
||
step_class._registry_name = registration_name
|
||
return step_class
|
||
|
||
return decorator
|
||
|
||
@classmethod
|
||
def get(cls, name: str) -> type:
|
||
"""Get a registered processor step class by name.
|
||
|
||
Args:
|
||
name: The registration name of the step.
|
||
|
||
Returns:
|
||
The registered step class.
|
||
|
||
Raises:
|
||
KeyError: If the step is not registered.
|
||
"""
|
||
if name not in cls._registry:
|
||
available = list(cls._registry.keys())
|
||
raise KeyError(
|
||
f"Processor step '{name}' not found in registry. "
|
||
f"Available steps: {available}. "
|
||
f"Make sure the step is registered using @ProcessorStepRegistry.register()"
|
||
)
|
||
return cls._registry[name]
|
||
|
||
@classmethod
|
||
def unregister(cls, name: str) -> None:
|
||
"""Remove a step from the registry."""
|
||
cls._registry.pop(name, None)
|
||
|
||
@classmethod
|
||
def list(cls) -> list[str]:
|
||
"""List all registered step names."""
|
||
return list(cls._registry.keys())
|
||
|
||
@classmethod
|
||
def clear(cls) -> None:
|
||
"""Clear all registrations."""
|
||
cls._registry.clear()
|
||
|
||
|
||
class ProcessorStep(Protocol):
|
||
"""Structural typing interface for a single processor step.
|
||
|
||
A step is any callable accepting a full `EnvTransition` dict and
|
||
returning a (possibly modified) dict of the same structure. Implementers
|
||
are encouraged—but not required—to expose the optional helper methods
|
||
listed below. When present, these hooks let `RobotProcessor`
|
||
automatically serialise the step's configuration and learnable state using
|
||
a safe-to-share JSON + SafeTensors format.
|
||
|
||
|
||
**Required**:
|
||
- ``__call__(transition: EnvTransition) -> EnvTransition``
|
||
- ``feature_contract(features: dict[str, PolicyFeature]) -> dict[str, PolicyFeature]``
|
||
|
||
Optional helper protocol:
|
||
* ``get_config() -> dict[str, Any]`` – User-defined JSON-serializable
|
||
configuration and state. YOU decide what to save here. This is where all
|
||
non-tensor state goes (e.g., name, counter, threshold, window_size).
|
||
The config dict will be passed to your class constructor when loading.
|
||
* ``state_dict() -> dict[str, torch.Tensor]`` – PyTorch tensor state ONLY.
|
||
This is exclusively for torch.Tensor objects (e.g., learned weights,
|
||
running statistics as tensors). Never put simple Python types here.
|
||
* ``load_state_dict(state)`` – Inverse of ``state_dict``. Receives a dict
|
||
containing torch tensors only.
|
||
* ``reset()`` – Clear internal buffers at episode boundaries.
|
||
|
||
Example separation:
|
||
- get_config(): {"name": "my_step", "learning_rate": 0.01, "window_size": 10}
|
||
- state_dict(): {"weights": torch.tensor(...), "running_mean": torch.tensor(...)}
|
||
"""
|
||
|
||
def __call__(self, transition: EnvTransition) -> EnvTransition: ...
|
||
|
||
def get_config(self) -> dict[str, Any]: ...
|
||
|
||
def state_dict(self) -> dict[str, torch.Tensor]: ...
|
||
|
||
def load_state_dict(self, state: dict[str, torch.Tensor]) -> None: ...
|
||
|
||
def reset(self) -> None: ...
|
||
|
||
def feature_contract(self, features: dict[str, PolicyFeature]) -> dict[str, PolicyFeature]: ...
|
||
|
||
|
||
def _default_batch_to_transition(batch: dict[str, Any]) -> EnvTransition: # noqa: D401
|
||
"""Convert a *batch* dict coming from Learobot replay/dataset code into an
|
||
``EnvTransition`` dictionary.
|
||
|
||
The function maps well known keys to the EnvTransition structure. Missing keys are
|
||
filled with sane defaults (``None`` or ``0.0``/``False``).
|
||
|
||
Keys recognised (case-sensitive):
|
||
|
||
* "observation.*" (keys starting with "observation." are grouped into observation dict)
|
||
* "action"
|
||
* "next.reward"
|
||
* "next.done"
|
||
* "next.truncated"
|
||
* "info"
|
||
|
||
Additional keys are ignored so that existing dataloaders can carry extra
|
||
metadata without breaking the processor.
|
||
"""
|
||
|
||
# Extract observation keys
|
||
observation_keys = {k: v for k, v in batch.items() if k.startswith("observation.")}
|
||
observation = observation_keys if observation_keys else None
|
||
|
||
# Extract padding and task keys for complementary data
|
||
pad_keys = {k: v for k, v in batch.items() if "_is_pad" in k}
|
||
task_key = {"task": batch["task"]} if "task" in batch else {}
|
||
complementary_data = {**pad_keys, **task_key} if pad_keys or task_key else {}
|
||
|
||
transition: EnvTransition = {
|
||
TransitionKey.OBSERVATION: observation,
|
||
TransitionKey.ACTION: batch.get("action"),
|
||
TransitionKey.REWARD: batch.get("next.reward", 0.0),
|
||
TransitionKey.DONE: batch.get("next.done", False),
|
||
TransitionKey.TRUNCATED: batch.get("next.truncated", False),
|
||
TransitionKey.INFO: batch.get("info", {}),
|
||
TransitionKey.COMPLEMENTARY_DATA: complementary_data,
|
||
}
|
||
return transition
|
||
|
||
|
||
def _default_transition_to_batch(transition: EnvTransition) -> dict[str, Any]: # noqa: D401
|
||
"""Inverse of :pyfunc:`_default_batch_to_transition`. Returns a dict with
|
||
the canonical field names used throughout *LeRobot*.
|
||
"""
|
||
|
||
batch = {
|
||
"action": transition.get(TransitionKey.ACTION),
|
||
"next.reward": transition.get(TransitionKey.REWARD, 0.0),
|
||
"next.done": transition.get(TransitionKey.DONE, False),
|
||
"next.truncated": transition.get(TransitionKey.TRUNCATED, False),
|
||
"info": transition.get(TransitionKey.INFO, {}),
|
||
}
|
||
|
||
# Add padding and task data from complementary_data
|
||
complementary_data = transition.get(TransitionKey.COMPLEMENTARY_DATA)
|
||
if complementary_data:
|
||
pad_data = {k: v for k, v in complementary_data.items() if "_is_pad" in k}
|
||
batch.update(pad_data)
|
||
|
||
if "task" in complementary_data:
|
||
batch["task"] = complementary_data["task"]
|
||
|
||
# Handle observation - flatten dict to observation.* keys if it's a dict
|
||
observation = transition.get(TransitionKey.OBSERVATION)
|
||
if isinstance(observation, dict):
|
||
batch.update(observation)
|
||
|
||
return batch
|
||
|
||
|
||
@dataclass
|
||
class RobotProcessor(ModelHubMixin):
|
||
"""
|
||
Composable, debuggable post-processing processor for robot transitions.
|
||
|
||
The class orchestrates an ordered collection of small, functional transforms—steps—executed
|
||
left-to-right on each incoming `EnvTransition`. It can process both `EnvTransition` dicts
|
||
and batch dictionaries, automatically converting between formats as needed.
|
||
|
||
Args:
|
||
steps: Ordered list of processing steps executed on every call. Defaults to empty list.
|
||
name: Human-readable identifier that is persisted inside the JSON config.
|
||
Defaults to "RobotProcessor".
|
||
to_transition: Function to convert batch dict to EnvTransition dict.
|
||
Defaults to _default_batch_to_transition.
|
||
to_output: Function to convert EnvTransition dict to the desired output format.
|
||
Usually it is a batch dict or EnvTransition dict.
|
||
Defaults to _default_transition_to_batch.
|
||
before_step_hooks: List of hooks called before each step. Each hook receives the step
|
||
index and transition, and can optionally return a modified transition.
|
||
after_step_hooks: List of hooks called after each step. Each hook receives the step
|
||
index and transition, and can optionally return a modified transition.
|
||
|
||
Hook Semantics:
|
||
- Hooks are executed sequentially in the order they were registered. There is no way to
|
||
reorder hooks after registration without creating a new pipeline.
|
||
- Hooks are for observation/monitoring only and DO NOT modify transitions. They are called
|
||
with the step index and current transition for logging, debugging, or monitoring purposes.
|
||
- All hooks for a given type (before/after) are executed for every step, or none at all if
|
||
an error occurs. There is no partial execution of hooks.
|
||
- Hooks should generally be stateless to maintain predictable behavior. If you need stateful
|
||
processing, consider implementing a proper ProcessorStep instead.
|
||
- To remove hooks, use the unregister methods. To remove steps, you must create a new pipeline.
|
||
- Hooks ALWAYS receive transitions in EnvTransition format, regardless of the input format
|
||
passed to __call__. This ensures consistent hook behavior whether processing batch dicts
|
||
or EnvTransition objects.
|
||
"""
|
||
|
||
steps: Sequence[ProcessorStep] = field(default_factory=list)
|
||
name: str = "RobotProcessor"
|
||
|
||
to_transition: Callable[[dict[str, Any]], EnvTransition] = field(
|
||
default_factory=lambda: _default_batch_to_transition, repr=False
|
||
)
|
||
to_output: Callable[[EnvTransition], dict[str, Any] | EnvTransition] = field(
|
||
default_factory=lambda: _default_transition_to_batch, repr=False
|
||
)
|
||
|
||
# Processor-level hooks for observation/monitoring
|
||
# Hooks do not modify transitions - they are called for logging, debugging, or monitoring purposes
|
||
before_step_hooks: list[Callable[[int, EnvTransition], None]] = field(default_factory=list, repr=False)
|
||
after_step_hooks: list[Callable[[int, EnvTransition], None]] = field(default_factory=list, repr=False)
|
||
|
||
def __call__(self, data: EnvTransition | dict[str, Any]):
|
||
"""Process data through all steps.
|
||
|
||
The method accepts either the classic EnvTransition dict or a batch dictionary
|
||
(like the ones returned by ReplayBuffer or LeRobotDataset). If a dict is supplied
|
||
it is first converted to the internal dict format using to_transition; after all
|
||
steps are executed the dict is transformed back into a batch dict with to_batch and the
|
||
result is returned – thereby preserving the caller's original data type.
|
||
|
||
Args:
|
||
data: Either an EnvTransition dict or a batch dictionary to process.
|
||
|
||
Returns:
|
||
The processed data in the same format as the input (EnvTransition or batch dict).
|
||
|
||
Raises:
|
||
ValueError: If the transition is not a valid EnvTransition format.
|
||
"""
|
||
# Check if we need to convert back to batch format at the end
|
||
_, called_with_batch = self._prepare_transition(data)
|
||
|
||
# Use step_through to get the iterator
|
||
step_iterator = self.step_through(data)
|
||
|
||
# Get initial state (before any steps)
|
||
current_transition = next(step_iterator)
|
||
|
||
# Process each step with hooks
|
||
for idx, next_transition in enumerate(step_iterator):
|
||
# Apply before hooks with current state (before step execution)
|
||
for hook in self.before_step_hooks:
|
||
hook(idx, current_transition)
|
||
|
||
# Move to next state (after step execution)
|
||
current_transition = next_transition
|
||
|
||
# Apply after hooks with updated state
|
||
for hook in self.after_step_hooks:
|
||
hook(idx, current_transition)
|
||
|
||
# Convert back to original format if needed
|
||
return self.to_output(current_transition) if called_with_batch else current_transition
|
||
|
||
def _prepare_transition(self, data: EnvTransition | dict[str, Any]) -> tuple[EnvTransition, bool]:
|
||
"""Prepare and validate transition data for processing.
|
||
|
||
Args:
|
||
data: Either an EnvTransition dict or a batch dictionary to process.
|
||
|
||
Returns:
|
||
A tuple of (prepared_transition, called_with_batch_flag)
|
||
|
||
Raises:
|
||
ValueError: If the transition is not a valid EnvTransition format.
|
||
"""
|
||
# Check if data is already an EnvTransition or needs conversion
|
||
if isinstance(data, dict) and not all(isinstance(k, TransitionKey) for k in data.keys()):
|
||
# It's a batch dict, convert it
|
||
called_with_batch = True
|
||
transition = self.to_transition(data)
|
||
else:
|
||
# It's already an EnvTransition
|
||
called_with_batch = False
|
||
transition = data
|
||
|
||
# Basic validation
|
||
if not isinstance(transition, dict):
|
||
raise ValueError(f"EnvTransition must be a dictionary. Got {type(transition).__name__}")
|
||
|
||
return transition, called_with_batch
|
||
|
||
def step_through(self, data: EnvTransition | dict[str, Any]) -> Iterable[EnvTransition]:
|
||
"""Yield the intermediate results after each processor step.
|
||
|
||
This is a low-level method that does NOT apply hooks. It simply executes each step
|
||
and yields the intermediate results. This allows users to debug the pipeline or
|
||
apply custom logic between steps if needed.
|
||
|
||
Note: This method always yields EnvTransition objects regardless of input format.
|
||
If you need the results in the original input format, you'll need to convert them
|
||
using `to_output()`.
|
||
|
||
Args:
|
||
data: Either an EnvTransition dict or a batch dictionary to process.
|
||
|
||
Yields:
|
||
The intermediate EnvTransition results after each step.
|
||
"""
|
||
transition, _ = self._prepare_transition(data)
|
||
|
||
# Yield initial state
|
||
yield transition
|
||
|
||
# Process each step WITHOUT hooks (low-level method)
|
||
for processor_step in self.steps:
|
||
transition = processor_step(transition)
|
||
yield transition
|
||
|
||
def _save_pretrained(self, save_directory: Path, **kwargs):
|
||
"""Internal save method for ModelHubMixin compatibility."""
|
||
# Extract config_filename from kwargs if provided
|
||
config_filename = kwargs.pop("config_filename", None)
|
||
self.save_pretrained(save_directory, config_filename=config_filename)
|
||
|
||
def save_pretrained(self, save_directory: str | Path, config_filename: str | None = None, **kwargs):
|
||
"""Serialize the processor definition and parameters to *save_directory*.
|
||
|
||
Args:
|
||
save_directory: Directory where the processor will be saved.
|
||
config_filename: Optional custom config filename. If not provided, defaults to
|
||
"{self.name}.json" where self.name is sanitized for filesystem compatibility.
|
||
"""
|
||
os.makedirs(str(save_directory), exist_ok=True)
|
||
|
||
# Sanitize processor name for use in filenames
|
||
import re
|
||
|
||
# The huggingface hub does not allow special characters in the repo name, so we sanitize the name
|
||
sanitized_name = re.sub(r"[^a-zA-Z0-9_]", "_", self.name.lower())
|
||
|
||
# Use sanitized name for config if not provided
|
||
if config_filename is None:
|
||
config_filename = f"{sanitized_name}.json"
|
||
|
||
config: dict[str, Any] = {
|
||
"name": self.name,
|
||
"steps": [],
|
||
}
|
||
|
||
for step_index, processor_step in enumerate(self.steps):
|
||
# Check if step was registered
|
||
registry_name = getattr(processor_step.__class__, "_registry_name", None)
|
||
|
||
step_entry: dict[str, Any] = {}
|
||
if registry_name:
|
||
# Use registry name for registered steps
|
||
step_entry["registry_name"] = registry_name
|
||
else:
|
||
# Fall back to full module path for unregistered steps
|
||
step_entry["class"] = (
|
||
f"{processor_step.__class__.__module__}.{processor_step.__class__.__name__}"
|
||
)
|
||
|
||
if hasattr(processor_step, "get_config"):
|
||
step_entry["config"] = processor_step.get_config()
|
||
|
||
if hasattr(processor_step, "state_dict"):
|
||
state = processor_step.state_dict()
|
||
if state:
|
||
# Clone tensors to avoid shared memory issues
|
||
# This ensures each tensor has its own memory allocation
|
||
# The reason is to avoid the following error:
|
||
# RuntimeError: Some tensors share memory, this will lead to duplicate memory on disk
|
||
# and potential differences when loading them again
|
||
# ------------------------------------------------------------------------------
|
||
# Since the state_dict of processor will be light, we can just clone the tensors
|
||
# and save them to the disk.
|
||
cloned_state = {}
|
||
for key, tensor in state.items():
|
||
cloned_state[key] = tensor.clone()
|
||
|
||
# Include pipeline name and step index to ensure unique filenames
|
||
# This prevents conflicts when multiple processors are saved in the same directory
|
||
if registry_name:
|
||
state_filename = f"{sanitized_name}_step_{step_index}_{registry_name}.safetensors"
|
||
else:
|
||
state_filename = f"{sanitized_name}_step_{step_index}.safetensors"
|
||
|
||
save_file(cloned_state, os.path.join(str(save_directory), state_filename))
|
||
step_entry["state_file"] = state_filename
|
||
|
||
config["steps"].append(step_entry)
|
||
|
||
with open(os.path.join(str(save_directory), config_filename), "w") as file_pointer:
|
||
json.dump(config, file_pointer, indent=2)
|
||
|
||
@classmethod
|
||
def from_pretrained(
|
||
cls,
|
||
pretrained_model_name_or_path: str | Path,
|
||
*,
|
||
force_download: bool = False,
|
||
resume_download: bool | None = None,
|
||
proxies: dict[str, str] | None = None,
|
||
token: str | bool | None = None,
|
||
cache_dir: str | Path | None = None,
|
||
local_files_only: bool = False,
|
||
revision: str | None = None,
|
||
config_filename: str | None = None,
|
||
overrides: dict[str, Any] | None = None,
|
||
**kwargs,
|
||
) -> RobotProcessor:
|
||
"""Load a serialized processor from source (local path or Hugging Face Hub identifier).
|
||
|
||
Args:
|
||
pretrained_model_name_or_path: Local path to a saved processor directory or Hugging Face Hub identifier
|
||
(e.g., "username/processor-name").
|
||
config_filename: Optional specific config filename to load. If not provided, will:
|
||
- For local paths: look for any .json file in the directory (error if multiple found)
|
||
- For HF Hub: try common names ("processor.json", "preprocessor.json", "postprocessor.json")
|
||
overrides: Optional dictionary mapping step names to configuration overrides.
|
||
Keys must match exact step class names (for unregistered steps) or registry names
|
||
(for registered steps). Values are dictionaries containing parameter overrides
|
||
that will be merged with the saved configuration. This is useful for providing
|
||
non-serializable objects like environment instances.
|
||
|
||
Returns:
|
||
A RobotProcessor instance loaded from the saved configuration.
|
||
|
||
Raises:
|
||
ImportError: If a processor step class cannot be loaded or imported.
|
||
ValueError: If a step cannot be instantiated with the provided configuration.
|
||
KeyError: If an override key doesn't match any step in the saved configuration.
|
||
|
||
Examples:
|
||
Basic loading:
|
||
```python
|
||
processor = RobotProcessor.from_pretrained("path/to/processor")
|
||
```
|
||
|
||
Loading specific config file:
|
||
```python
|
||
processor = RobotProcessor.from_pretrained(
|
||
"username/multi-processor-repo", config_filename="preprocessor.json"
|
||
)
|
||
```
|
||
|
||
Loading with overrides for non-serializable objects:
|
||
```python
|
||
import gym
|
||
|
||
env = gym.make("CartPole-v1")
|
||
processor = RobotProcessor.from_pretrained(
|
||
"username/cartpole-processor", overrides={"ActionRepeatStep": {"env": env}}
|
||
)
|
||
```
|
||
|
||
Multiple overrides:
|
||
```python
|
||
processor = RobotProcessor.from_pretrained(
|
||
"path/to/processor",
|
||
overrides={
|
||
"CustomStep": {"param1": "new_value"},
|
||
"device_processor": {"device": "cuda:1"}, # For registered steps
|
||
},
|
||
)
|
||
```
|
||
"""
|
||
# Use the local variable name 'source' for clarity
|
||
source = str(pretrained_model_name_or_path)
|
||
|
||
if Path(source).is_dir():
|
||
# Local path - use it directly
|
||
base_path = Path(source)
|
||
|
||
if config_filename is None:
|
||
# Look for any .json file in the directory
|
||
json_files = list(base_path.glob("*.json"))
|
||
if len(json_files) == 0:
|
||
raise FileNotFoundError(f"No .json configuration files found in {source}")
|
||
elif len(json_files) > 1:
|
||
raise ValueError(
|
||
f"Multiple .json files found in {source}: {[f.name for f in json_files]}. "
|
||
f"Please specify which one to load using the config_filename parameter."
|
||
)
|
||
config_filename = json_files[0].name
|
||
|
||
with open(base_path / config_filename) as file_pointer:
|
||
loaded_config: dict[str, Any] = json.load(file_pointer)
|
||
else:
|
||
# Hugging Face Hub - download all required files
|
||
if config_filename is None:
|
||
# Try common config names
|
||
common_names = [
|
||
"processor.json",
|
||
"preprocessor.json",
|
||
"postprocessor.json",
|
||
"robotprocessor.json",
|
||
]
|
||
config_path = None
|
||
for name in common_names:
|
||
try:
|
||
config_path = hf_hub_download(
|
||
source,
|
||
name,
|
||
repo_type="model",
|
||
force_download=force_download,
|
||
resume_download=resume_download,
|
||
proxies=proxies,
|
||
token=token,
|
||
cache_dir=cache_dir,
|
||
local_files_only=local_files_only,
|
||
revision=revision,
|
||
)
|
||
config_filename = name
|
||
break
|
||
except (FileNotFoundError, OSError, HfHubHTTPError):
|
||
# FileNotFoundError: local file issues
|
||
# OSError: network/system errors
|
||
# HfHubHTTPError: file not found on Hub (404) or other HTTP errors
|
||
continue
|
||
|
||
if config_path is None:
|
||
raise FileNotFoundError(
|
||
f"No processor configuration file found in {source}. "
|
||
f"Tried: {common_names}. Please specify the config_filename parameter."
|
||
)
|
||
else:
|
||
# Download specific config file
|
||
config_path = hf_hub_download(
|
||
source,
|
||
config_filename,
|
||
repo_type="model",
|
||
force_download=force_download,
|
||
resume_download=resume_download,
|
||
proxies=proxies,
|
||
token=token,
|
||
cache_dir=cache_dir,
|
||
local_files_only=local_files_only,
|
||
revision=revision,
|
||
)
|
||
|
||
with open(config_path) as file_pointer:
|
||
loaded_config = json.load(file_pointer)
|
||
|
||
# Store downloaded files in the same directory as the config
|
||
base_path = Path(config_path).parent
|
||
|
||
# Handle None overrides
|
||
if overrides is None:
|
||
overrides = {}
|
||
|
||
# Validate that all override keys will be matched
|
||
override_keys = set(overrides.keys())
|
||
|
||
steps: list[ProcessorStep] = []
|
||
for step_entry in loaded_config["steps"]:
|
||
# Check if step uses registry name or module path
|
||
if "registry_name" in step_entry:
|
||
# Load from registry
|
||
try:
|
||
step_class = ProcessorStepRegistry.get(step_entry["registry_name"])
|
||
step_key = step_entry["registry_name"]
|
||
except KeyError as e:
|
||
raise ImportError(f"Failed to load processor step from registry. {str(e)}") from e
|
||
else:
|
||
# Fall back to module path loading for backward compatibility
|
||
full_class_path = step_entry["class"]
|
||
module_path, class_name = full_class_path.rsplit(".", 1)
|
||
|
||
# Import the module containing the step class
|
||
try:
|
||
module = importlib.import_module(module_path)
|
||
step_class = getattr(module, class_name)
|
||
step_key = class_name
|
||
except (ImportError, AttributeError) as e:
|
||
raise ImportError(
|
||
f"Failed to load processor step '{full_class_path}'. "
|
||
f"Make sure the module '{module_path}' is installed and contains class '{class_name}'. "
|
||
f"Consider registering the step using @ProcessorStepRegistry.register() for better portability. "
|
||
f"Error: {str(e)}"
|
||
) from e
|
||
|
||
# Instantiate the step with its config
|
||
try:
|
||
saved_cfg = step_entry.get("config", {})
|
||
step_overrides = overrides.get(step_key, {})
|
||
merged_cfg = {**saved_cfg, **step_overrides}
|
||
step_instance: ProcessorStep = step_class(**merged_cfg)
|
||
|
||
# Track which override keys were used
|
||
if step_key in override_keys:
|
||
override_keys.discard(step_key)
|
||
|
||
except Exception as e:
|
||
step_name = step_entry.get("registry_name", step_entry.get("class", "Unknown"))
|
||
raise ValueError(
|
||
f"Failed to instantiate processor step '{step_name}' with config: {step_entry.get('config', {})}. "
|
||
f"Error: {str(e)}"
|
||
) from e
|
||
|
||
# Load state if available
|
||
if "state_file" in step_entry and hasattr(step_instance, "load_state_dict"):
|
||
if Path(source).is_dir():
|
||
# Local path - read directly
|
||
state_path = str(base_path / step_entry["state_file"])
|
||
else:
|
||
# Hugging Face Hub - download the state file
|
||
state_path = hf_hub_download(
|
||
source,
|
||
step_entry["state_file"],
|
||
repo_type="model",
|
||
force_download=force_download,
|
||
resume_download=resume_download,
|
||
proxies=proxies,
|
||
token=token,
|
||
cache_dir=cache_dir,
|
||
local_files_only=local_files_only,
|
||
revision=revision,
|
||
)
|
||
|
||
step_instance.load_state_dict(load_file(state_path))
|
||
|
||
steps.append(step_instance)
|
||
|
||
# Check for unused override keys
|
||
if override_keys:
|
||
available_keys = []
|
||
for step_entry in loaded_config["steps"]:
|
||
if "registry_name" in step_entry:
|
||
available_keys.append(step_entry["registry_name"])
|
||
else:
|
||
full_class_path = step_entry["class"]
|
||
class_name = full_class_path.rsplit(".", 1)[1]
|
||
available_keys.append(class_name)
|
||
|
||
raise KeyError(
|
||
f"Override keys {list(override_keys)} do not match any step in the saved configuration. "
|
||
f"Available step keys: {available_keys}. "
|
||
f"Make sure override keys match exact step class names or registry names."
|
||
)
|
||
|
||
return cls(steps, loaded_config.get("name", "RobotProcessor"))
|
||
|
||
def __len__(self) -> int:
|
||
"""Return the number of steps in the processor."""
|
||
return len(self.steps)
|
||
|
||
def __getitem__(self, idx: int | slice) -> ProcessorStep | RobotProcessor:
|
||
"""Indexing helper exposing underlying steps.
|
||
* ``int`` – returns the idx-th ProcessorStep.
|
||
* ``slice`` – returns a new RobotProcessor with the sliced steps.
|
||
"""
|
||
if isinstance(idx, slice):
|
||
return RobotProcessor(self.steps[idx], self.name)
|
||
return self.steps[idx]
|
||
|
||
def register_before_step_hook(self, fn: Callable[[int, EnvTransition], None]):
|
||
"""Attach fn to be executed before every processor step."""
|
||
self.before_step_hooks.append(fn)
|
||
|
||
def unregister_before_step_hook(self, fn: Callable[[int, EnvTransition], None]):
|
||
"""Remove a previously registered before_step hook.
|
||
|
||
Args:
|
||
fn: The exact function reference that was registered. Must be the same object.
|
||
|
||
Raises:
|
||
ValueError: If the hook is not found in the registered hooks.
|
||
"""
|
||
try:
|
||
self.before_step_hooks.remove(fn)
|
||
except ValueError:
|
||
raise ValueError(
|
||
f"Hook {fn} not found in before_step_hooks. Make sure to pass the exact same function reference."
|
||
) from None
|
||
|
||
def register_after_step_hook(self, fn: Callable[[int, EnvTransition], None]):
|
||
"""Attach fn to be executed after every processor step."""
|
||
self.after_step_hooks.append(fn)
|
||
|
||
def unregister_after_step_hook(self, fn: Callable[[int, EnvTransition], None]):
|
||
"""Remove a previously registered after_step hook.
|
||
|
||
Args:
|
||
fn: The exact function reference that was registered. Must be the same object.
|
||
|
||
Raises:
|
||
ValueError: If the hook is not found in the registered hooks.
|
||
"""
|
||
try:
|
||
self.after_step_hooks.remove(fn)
|
||
except ValueError:
|
||
raise ValueError(
|
||
f"Hook {fn} not found in after_step_hooks. Make sure to pass the exact same function reference."
|
||
) from None
|
||
|
||
def reset(self):
|
||
"""Clear state in every step that implements ``reset()`` and fire registered hooks."""
|
||
for step in self.steps:
|
||
if hasattr(step, "reset"):
|
||
step.reset() # type: ignore[attr-defined]
|
||
|
||
def __repr__(self) -> str:
|
||
"""Return a readable string representation of the processor."""
|
||
step_names = [step.__class__.__name__ for step in self.steps]
|
||
|
||
if not step_names:
|
||
steps_repr = "steps=0: []"
|
||
elif len(step_names) <= 3:
|
||
steps_repr = f"steps={len(step_names)}: [{', '.join(step_names)}]"
|
||
else:
|
||
# Show first 2 and last 1 with ellipsis for long lists
|
||
displayed = f"{step_names[0]}, {step_names[1]}, ..., {step_names[-1]}"
|
||
steps_repr = f"steps={len(step_names)}: [{displayed}]"
|
||
|
||
parts = [f"name='{self.name}'", steps_repr]
|
||
|
||
return f"RobotProcessor({', '.join(parts)})"
|
||
|
||
def __post_init__(self):
|
||
for i, step in enumerate(self.steps):
|
||
if not callable(step):
|
||
raise TypeError(
|
||
f"Step {i} ({type(step).__name__}) must define __call__(transition) -> EnvTransition"
|
||
)
|
||
|
||
fc = getattr(step, "feature_contract", None)
|
||
if not callable(fc):
|
||
raise TypeError(
|
||
f"Step {i} ({type(step).__name__}) must define feature_contract(features) -> dict[str, Any]"
|
||
)
|
||
|
||
def feature_contract(self, initial_features: dict[str, PolicyFeature]) -> dict[str, PolicyFeature]:
|
||
"""
|
||
Apply ALL steps in order. Each step must implement
|
||
feature_contract(features) and return a dict (full or incremental schema).
|
||
"""
|
||
features: dict[str, PolicyFeature] = deepcopy(initial_features)
|
||
|
||
for _, step in enumerate(self.steps):
|
||
out = step.feature_contract(features)
|
||
if not isinstance(out, dict):
|
||
raise TypeError(f"{step.__class__.__name__}.feature_contract must return dict[str, Any]")
|
||
features = out
|
||
return features
|
||
|
||
|
||
class ObservationProcessor:
|
||
"""Base class for processors that modify only the observation component of a transition.
|
||
|
||
Subclasses should override the `observation` method to implement custom observation processing.
|
||
This class handles the boilerplate of extracting and reinserting the processed observation
|
||
into the transition dict, eliminating the need to implement the `__call__` method in subclasses.
|
||
|
||
Example:
|
||
```python
|
||
class MyObservationScaler(ObservationProcessor):
|
||
def __init__(self, scale_factor):
|
||
self.scale_factor = scale_factor
|
||
|
||
def observation(self, observation):
|
||
return observation * self.scale_factor
|
||
```
|
||
|
||
By inheriting from this class, you avoid writing repetitive code to handle transition dict
|
||
manipulation, focusing only on the specific observation processing logic.
|
||
"""
|
||
|
||
def observation(self, observation):
|
||
"""Process the observation component.
|
||
|
||
Args:
|
||
observation: The observation to process
|
||
|
||
Returns:
|
||
The processed observation
|
||
"""
|
||
return observation
|
||
|
||
def __call__(self, transition: EnvTransition) -> EnvTransition:
|
||
observation = transition.get(TransitionKey.OBSERVATION)
|
||
if observation is None:
|
||
return transition
|
||
|
||
processed_observation = self.observation(observation)
|
||
# Create a new transition dict with the processed observation
|
||
new_transition = transition.copy()
|
||
new_transition[TransitionKey.OBSERVATION] = processed_observation
|
||
return new_transition
|
||
|
||
def get_config(self) -> dict[str, Any]:
|
||
return {}
|
||
|
||
def state_dict(self) -> dict[str, torch.Tensor]:
|
||
return {}
|
||
|
||
def load_state_dict(self, state: dict[str, torch.Tensor]) -> None:
|
||
pass
|
||
|
||
def reset(self) -> None:
|
||
pass
|
||
|
||
def feature_contract(self, features: dict[str, PolicyFeature]) -> dict[str, PolicyFeature]:
|
||
return features
|
||
|
||
|
||
class ActionProcessor:
|
||
"""Base class for processors that modify only the action component of a transition.
|
||
|
||
Subclasses should override the `action` method to implement custom action processing.
|
||
This class handles the boilerplate of extracting and reinserting the processed action
|
||
into the transition dict, eliminating the need to implement the `__call__` method in subclasses.
|
||
|
||
Example:
|
||
```python
|
||
class ActionClipping(ActionProcessor):
|
||
def __init__(self, min_val, max_val):
|
||
self.min_val = min_val
|
||
self.max_val = max_val
|
||
|
||
def action(self, action):
|
||
return np.clip(action, self.min_val, self.max_val)
|
||
```
|
||
|
||
By inheriting from this class, you avoid writing repetitive code to handle transition dict
|
||
manipulation, focusing only on the specific action processing logic.
|
||
"""
|
||
|
||
def action(self, action):
|
||
"""Process the action component.
|
||
|
||
Args:
|
||
action: The action to process
|
||
|
||
Returns:
|
||
The processed action
|
||
"""
|
||
return action
|
||
|
||
def __call__(self, transition: EnvTransition) -> EnvTransition:
|
||
action = transition.get(TransitionKey.ACTION)
|
||
if action is None:
|
||
return transition
|
||
|
||
processed_action = self.action(action)
|
||
# Create a new transition dict with the processed action
|
||
new_transition = transition.copy()
|
||
new_transition[TransitionKey.ACTION] = processed_action
|
||
return new_transition
|
||
|
||
def get_config(self) -> dict[str, Any]:
|
||
return {}
|
||
|
||
def state_dict(self) -> dict[str, torch.Tensor]:
|
||
return {}
|
||
|
||
def load_state_dict(self, state: dict[str, torch.Tensor]) -> None:
|
||
pass
|
||
|
||
def reset(self) -> None:
|
||
pass
|
||
|
||
def feature_contract(self, features: dict[str, PolicyFeature]) -> dict[str, PolicyFeature]:
|
||
return features
|
||
|
||
|
||
class RewardProcessor:
|
||
"""Base class for processors that modify only the reward component of a transition.
|
||
|
||
Subclasses should override the `reward` method to implement custom reward processing.
|
||
This class handles the boilerplate of extracting and reinserting the processed reward
|
||
into the transition dict, eliminating the need to implement the `__call__` method in subclasses.
|
||
|
||
Example:
|
||
```python
|
||
class RewardScaler(RewardProcessor):
|
||
def __init__(self, scale_factor):
|
||
self.scale_factor = scale_factor
|
||
|
||
def reward(self, reward):
|
||
return reward * self.scale_factor
|
||
```
|
||
|
||
By inheriting from this class, you avoid writing repetitive code to handle transition dict
|
||
manipulation, focusing only on the specific reward processing logic.
|
||
"""
|
||
|
||
def reward(self, reward):
|
||
"""Process the reward component.
|
||
|
||
Args:
|
||
reward: The reward to process
|
||
|
||
Returns:
|
||
The processed reward
|
||
"""
|
||
return reward
|
||
|
||
def __call__(self, transition: EnvTransition) -> EnvTransition:
|
||
reward = transition.get(TransitionKey.REWARD)
|
||
if reward is None:
|
||
return transition
|
||
|
||
processed_reward = self.reward(reward)
|
||
# Create a new transition dict with the processed reward
|
||
new_transition = transition.copy()
|
||
new_transition[TransitionKey.REWARD] = processed_reward
|
||
return new_transition
|
||
|
||
def get_config(self) -> dict[str, Any]:
|
||
return {}
|
||
|
||
def state_dict(self) -> dict[str, torch.Tensor]:
|
||
return {}
|
||
|
||
def load_state_dict(self, state: dict[str, torch.Tensor]) -> None:
|
||
pass
|
||
|
||
def reset(self) -> None:
|
||
pass
|
||
|
||
def feature_contract(self, features: dict[str, PolicyFeature]) -> dict[str, PolicyFeature]:
|
||
return features
|
||
|
||
|
||
class DoneProcessor:
|
||
"""Base class for processors that modify only the done flag of a transition.
|
||
|
||
Subclasses should override the `done` method to implement custom done flag processing.
|
||
This class handles the boilerplate of extracting and reinserting the processed done flag
|
||
into the transition dict, eliminating the need to implement the `__call__` method in subclasses.
|
||
|
||
Example:
|
||
```python
|
||
class TimeoutDone(DoneProcessor):
|
||
def __init__(self, max_steps):
|
||
self.steps = 0
|
||
self.max_steps = max_steps
|
||
|
||
def done(self, done):
|
||
self.steps += 1
|
||
return done or self.steps >= self.max_steps
|
||
|
||
def reset(self):
|
||
self.steps = 0
|
||
```
|
||
|
||
By inheriting from this class, you avoid writing repetitive code to handle transition dict
|
||
manipulation, focusing only on the specific done flag processing logic.
|
||
"""
|
||
|
||
def done(self, done):
|
||
"""Process the done flag.
|
||
|
||
Args:
|
||
done: The done flag to process
|
||
|
||
Returns:
|
||
The processed done flag
|
||
"""
|
||
return done
|
||
|
||
def __call__(self, transition: EnvTransition) -> EnvTransition:
|
||
done = transition.get(TransitionKey.DONE)
|
||
if done is None:
|
||
return transition
|
||
|
||
processed_done = self.done(done)
|
||
# Create a new transition dict with the processed done flag
|
||
new_transition = transition.copy()
|
||
new_transition[TransitionKey.DONE] = processed_done
|
||
return new_transition
|
||
|
||
def get_config(self) -> dict[str, Any]:
|
||
return {}
|
||
|
||
def state_dict(self) -> dict[str, torch.Tensor]:
|
||
return {}
|
||
|
||
def load_state_dict(self, state: dict[str, torch.Tensor]) -> None:
|
||
pass
|
||
|
||
def reset(self) -> None:
|
||
pass
|
||
|
||
def feature_contract(self, features: dict[str, PolicyFeature]) -> dict[str, PolicyFeature]:
|
||
return features
|
||
|
||
|
||
class TruncatedProcessor:
|
||
"""Base class for processors that modify only the truncated flag of a transition.
|
||
|
||
Subclasses should override the `truncated` method to implement custom truncated flag processing.
|
||
This class handles the boilerplate of extracting and reinserting the processed truncated flag
|
||
into the transition dict, eliminating the need to implement the `__call__` method in subclasses.
|
||
|
||
Example:
|
||
```python
|
||
class EarlyTruncation(TruncatedProcessor):
|
||
def __init__(self, threshold):
|
||
self.threshold = threshold
|
||
|
||
def truncated(self, truncated):
|
||
# Additional truncation condition
|
||
return truncated or some_condition > self.threshold
|
||
```
|
||
|
||
By inheriting from this class, you avoid writing repetitive code to handle transition dict
|
||
manipulation, focusing only on the specific truncated flag processing logic.
|
||
"""
|
||
|
||
def truncated(self, truncated):
|
||
"""Process the truncated flag.
|
||
|
||
Args:
|
||
truncated: The truncated flag to process
|
||
|
||
Returns:
|
||
The processed truncated flag
|
||
"""
|
||
return truncated
|
||
|
||
def __call__(self, transition: EnvTransition) -> EnvTransition:
|
||
truncated = transition.get(TransitionKey.TRUNCATED)
|
||
if truncated is None:
|
||
return transition
|
||
|
||
processed_truncated = self.truncated(truncated)
|
||
# Create a new transition dict with the processed truncated flag
|
||
new_transition = transition.copy()
|
||
new_transition[TransitionKey.TRUNCATED] = processed_truncated
|
||
return new_transition
|
||
|
||
def get_config(self) -> dict[str, Any]:
|
||
return {}
|
||
|
||
def state_dict(self) -> dict[str, torch.Tensor]:
|
||
return {}
|
||
|
||
def load_state_dict(self, state: dict[str, torch.Tensor]) -> None:
|
||
pass
|
||
|
||
def reset(self) -> None:
|
||
pass
|
||
|
||
def feature_contract(self, features: dict[str, PolicyFeature]) -> dict[str, PolicyFeature]:
|
||
return features
|
||
|
||
|
||
class InfoProcessor:
|
||
"""Base class for processors that modify only the info dictionary of a transition.
|
||
|
||
Subclasses should override the `info` method to implement custom info processing.
|
||
This class handles the boilerplate of extracting and reinserting the processed info
|
||
into the transition dict, eliminating the need to implement the `__call__` method in subclasses.
|
||
|
||
Example:
|
||
```python
|
||
class InfoAugmenter(InfoProcessor):
|
||
def __init__(self):
|
||
self.step_count = 0
|
||
|
||
def info(self, info):
|
||
info = info.copy() # Create a copy to avoid modifying the original
|
||
info["steps"] = self.step_count
|
||
self.step_count += 1
|
||
return info
|
||
|
||
def reset(self):
|
||
self.step_count = 0
|
||
```
|
||
|
||
By inheriting from this class, you avoid writing repetitive code to handle transition dict
|
||
manipulation, focusing only on the specific info dictionary processing logic.
|
||
"""
|
||
|
||
def info(self, info):
|
||
"""Process the info dictionary.
|
||
|
||
Args:
|
||
info: The info dictionary to process
|
||
|
||
Returns:
|
||
The processed info dictionary
|
||
"""
|
||
return info
|
||
|
||
def __call__(self, transition: EnvTransition) -> EnvTransition:
|
||
info = transition.get(TransitionKey.INFO)
|
||
if info is None:
|
||
return transition
|
||
|
||
processed_info = self.info(info)
|
||
# Create a new transition dict with the processed info
|
||
new_transition = transition.copy()
|
||
new_transition[TransitionKey.INFO] = processed_info
|
||
return new_transition
|
||
|
||
def get_config(self) -> dict[str, Any]:
|
||
return {}
|
||
|
||
def state_dict(self) -> dict[str, torch.Tensor]:
|
||
return {}
|
||
|
||
def load_state_dict(self, state: dict[str, torch.Tensor]) -> None:
|
||
pass
|
||
|
||
def reset(self) -> None:
|
||
pass
|
||
|
||
def feature_contract(self, features: dict[str, PolicyFeature]) -> dict[str, PolicyFeature]:
|
||
return features
|
||
|
||
|
||
class ComplementaryDataProcessor:
|
||
"""Base class for processors that modify only the complementary data of a transition.
|
||
|
||
Subclasses should override the `complementary_data` method to implement custom complementary data processing.
|
||
This class handles the boilerplate of extracting and reinserting the processed complementary data
|
||
into the transition dict, eliminating the need to implement the `__call__` method in subclasses.
|
||
"""
|
||
|
||
def complementary_data(self, complementary_data):
|
||
"""Process the complementary data.
|
||
|
||
Args:
|
||
complementary_data: The complementary data to process
|
||
|
||
Returns:
|
||
The processed complementary data
|
||
"""
|
||
return complementary_data
|
||
|
||
def __call__(self, transition: EnvTransition) -> EnvTransition:
|
||
complementary_data = transition.get(TransitionKey.COMPLEMENTARY_DATA)
|
||
if complementary_data is None:
|
||
return transition
|
||
|
||
processed_complementary_data = self.complementary_data(complementary_data)
|
||
# Create a new transition dict with the processed complementary data
|
||
new_transition = transition.copy()
|
||
new_transition[TransitionKey.COMPLEMENTARY_DATA] = processed_complementary_data
|
||
return new_transition
|
||
|
||
def get_config(self) -> dict[str, Any]:
|
||
return {}
|
||
|
||
def state_dict(self) -> dict[str, torch.Tensor]:
|
||
return {}
|
||
|
||
def load_state_dict(self, state: dict[str, torch.Tensor]) -> None:
|
||
pass
|
||
|
||
def reset(self) -> None:
|
||
pass
|
||
|
||
def feature_contract(self, features: dict[str, PolicyFeature]) -> dict[str, PolicyFeature]:
|
||
return features
|
||
|
||
|
||
class IdentityProcessor:
|
||
"""Identity processor that does nothing."""
|
||
|
||
def __call__(self, transition: EnvTransition) -> EnvTransition:
|
||
return transition
|
||
|
||
def get_config(self) -> dict[str, Any]:
|
||
return {}
|
||
|
||
def state_dict(self) -> dict[str, torch.Tensor]:
|
||
return {}
|
||
|
||
def load_state_dict(self, state: dict[str, torch.Tensor]) -> None:
|
||
pass
|
||
|
||
def reset(self) -> None:
|
||
pass
|
||
|
||
def feature_contract(self, features: dict[str, PolicyFeature]) -> dict[str, PolicyFeature]:
|
||
return features
|