64 lines
1.9 KiB
Python
64 lines
1.9 KiB
Python
|
"""Public data config dataclasses.
|
||
|
|
|
||
|
|
Only data backends shipped in the public package should define config classes
|
||
|
|
here. Internal backends register their dataclasses from their own packages via
|
||
|
|
``wall_x.config.registry.register_data_config``.
|
||
|
|
"""
|
||
|
|
|
||
|
|
from dataclasses import dataclass, field
|
||
|
|
from typing import Any, Dict, Optional
|
||
|
|
|
||
|
|
from .registry import register_data_config
|
||
|
|
|
||
|
|
|
||
|
|
@dataclass
|
||
|
|
class DataConfig:
|
||
|
|
"""Base fields shared by data backends.
|
||
|
|
|
||
|
|
``normalizer_config`` may contain:
|
||
|
|
- ``min_key``: stats key for the minimum value.
|
||
|
|
- ``delta_key``: stats key for the value range.
|
||
|
|
- ``customized_action_statistic_dof``: explicit action-stats JSON path.
|
||
|
|
"""
|
||
|
|
|
||
|
|
dataset_type: str = "lerobot"
|
||
|
|
resolution: Dict[str, int] = field(
|
||
|
|
default_factory=lambda: {
|
||
|
|
"face_view": 256,
|
||
|
|
"left_wrist_view": 256,
|
||
|
|
"right_wrist_view": 256,
|
||
|
|
}
|
||
|
|
)
|
||
|
|
train_test_split: float = 0.95
|
||
|
|
normalizer_config: Optional[Dict[str, Any]] = None
|
||
|
|
|
||
|
|
|
||
|
|
@register_data_config("lerobot")
|
||
|
|
@dataclass
|
||
|
|
class LeRobotDataConfig(DataConfig):
|
||
|
|
"""LeRobot data config.
|
||
|
|
|
||
|
|
``lerobot_config`` is expected to contain fields such as ``repo_id`` and
|
||
|
|
``root`` for a HuggingFace LeRobot dataset. ``norm_stats_path`` points to
|
||
|
|
explicit action normalizer stats; the core package does not bundle private
|
||
|
|
defaults.
|
||
|
|
"""
|
||
|
|
|
||
|
|
dataset_type: str = "lerobot"
|
||
|
|
lerobot_config: Optional[Dict[str, Any]] = None
|
||
|
|
key_mappings: Optional[Dict[str, Any]] = None
|
||
|
|
norm_stats_path: Optional[str] = None
|
||
|
|
priority_order: Optional[Dict[str, float]] = None
|
||
|
|
camera_name_mapping: Optional[Dict[str, str]] = None
|
||
|
|
num_workers: int = 4
|
||
|
|
action_tokenizer_path: Optional[str] = None
|
||
|
|
use_fast_tokenizer: bool = False
|
||
|
|
padding_side: str = "left"
|
||
|
|
noise_scheduler: Optional[Dict[str, Any]] = None
|
||
|
|
|
||
|
|
|
||
|
|
__all__ = [
|
||
|
|
"DataConfig",
|
||
|
|
"LeRobotDataConfig",
|
||
|
|
]
|