mirror of
https://github.com/huggingface/lerobot.git
synced 2026-07-23 17:56:07 +00:00
Michel Aractingi
d0b7690bc0
Added support for hil_serl classifier to be trained with train.py run classifier training by python lerobot/scripts/train.py --policy.type=hilserl_classifier fixes in find_joint_limits, control_robot, end_effector_control_utils
55 lines
1.6 KiB
Python
55 lines
1.6 KiB
Python
from dataclasses import dataclass, field
|
|
from typing import Dict, List
|
|
|
|
from lerobot.common.optim.optimizers import AdamWConfig, OptimizerConfig
|
|
from lerobot.common.optim.schedulers import LRSchedulerConfig
|
|
from lerobot.configs.policies import PreTrainedConfig
|
|
from lerobot.configs.types import FeatureType, PolicyFeature
|
|
|
|
|
|
@PreTrainedConfig.register_subclass(name="hilserl_classifier")
|
|
@dataclass
|
|
class ClassifierConfig(PreTrainedConfig):
|
|
"""Configuration for the Classifier model."""
|
|
|
|
name: str = "hilserl_classifier"
|
|
num_classes: int = 2
|
|
hidden_dim: int = 256
|
|
dropout_rate: float = 0.1
|
|
model_name: str = "helper2424/resnet10"
|
|
device: str = "cpu"
|
|
model_type: str = "cnn" # "transformer" or "cnn"
|
|
num_cameras: int = 2
|
|
learning_rate: float = 1e-4
|
|
normalization_mode = None
|
|
# output_features: Dict[str, PolicyFeature] = field(
|
|
# default_factory=lambda: {"next.reward": PolicyFeature(type=FeatureType.REWARD, shape=(1,))}
|
|
# )
|
|
|
|
@property
|
|
def observation_delta_indices(self) -> List | None:
|
|
return None
|
|
|
|
@property
|
|
def action_delta_indices(self) -> List | None:
|
|
return None
|
|
|
|
@property
|
|
def reward_delta_indices(self) -> List | None:
|
|
return None
|
|
|
|
def get_optimizer_preset(self) -> OptimizerConfig:
|
|
return AdamWConfig(
|
|
lr=self.learning_rate,
|
|
weight_decay=0.01,
|
|
grad_clip_norm=1.0,
|
|
)
|
|
|
|
def get_scheduler_preset(self) -> LRSchedulerConfig | None:
|
|
return None
|
|
|
|
def validate_features(self) -> None:
|
|
"""Validate feature configurations."""
|
|
# Classifier doesn't need specific feature validation
|
|
pass
|