forked from tangger/lerobot
Co-authored-by: Remi <remi.cadene@huggingface.co> Co-authored-by: HUANG TZU-CHUN <137322177+tc-huang@users.noreply.github.com>
29 lines
611 B
Python
29 lines
611 B
Python
# Note: We subclass str so that serialization is straightforward
|
|
# https://stackoverflow.com/questions/24481852/serialising-an-enum-member-to-json
|
|
from dataclasses import dataclass
|
|
from enum import Enum
|
|
from typing import Any, Protocol
|
|
|
|
|
|
class FeatureType(str, Enum):
|
|
STATE = "STATE"
|
|
VISUAL = "VISUAL"
|
|
ENV = "ENV"
|
|
ACTION = "ACTION"
|
|
|
|
|
|
class NormalizationMode(str, Enum):
|
|
MIN_MAX = "MIN_MAX"
|
|
MEAN_STD = "MEAN_STD"
|
|
IDENTITY = "IDENTITY"
|
|
|
|
|
|
class DictLike(Protocol):
|
|
def __getitem__(self, key: Any) -> Any: ...
|
|
|
|
|
|
@dataclass
|
|
class PolicyFeature:
|
|
type: FeatureType
|
|
shape: tuple
|