added optimizer and sac to factory.py

This commit is contained in:
Michel Aractingi 2024-12-23 14:12:03 +01:00 committed by AdilZouitine
parent d96edbf5ac
commit 9dafad15e6
2 changed files with 7 additions and 0 deletions

View File

@ -54,6 +54,12 @@ def get_policy_class(name: str) -> PreTrainedPolicy:
from lerobot.common.policies.pi0.modeling_pi0 import PI0Policy
return PI0Policy
elif name == "sac":
from lerobot.common.policies.sac.configuration_sac import SACConfig
from lerobot.common.policies.sac.modeling_sac import SACPolicy
return SACPolicy, SACConfig
else:
raise NotImplementedError(f"Policy with name {name} is not implemented.")

View File

@ -26,6 +26,7 @@ class SACConfig:
num_subsample_critics = None
critic_lr = 3e-4
actor_lr = 3e-4
temperature_lr = 3e-4
critic_target_update_weight = 0.005
utd_ratio = 2
critic_network_kwargs = {