mirror of
https://github.com/huggingface/lerobot.git
synced 2026-07-23 01:41:54 +00:00
update factory with dsrl
This commit is contained in:
@@ -249,7 +249,7 @@ class DSRLPolicy(PreTrainedPolicy):
|
|||||||
raise ValueError(f"Unknown model type: {model}")
|
raise ValueError(f"Unknown model type: {model}")
|
||||||
|
|
||||||
def update_target_networks(self):
|
def update_target_networks(self):
|
||||||
"""Update target networks with exponential moving average"""
|
"""Update target networks of the action critic with exponential moving average"""
|
||||||
for target_param, param in zip(
|
for target_param, param in zip(
|
||||||
self.action_critic_target.parameters(),
|
self.action_critic_target.parameters(),
|
||||||
self.action_critic_ensemble.parameters(),
|
self.action_critic_ensemble.parameters(),
|
||||||
|
|||||||
@@ -30,6 +30,7 @@ from lerobot.envs.configs import EnvConfig
|
|||||||
from lerobot.envs.utils import env_to_policy_features
|
from lerobot.envs.utils import env_to_policy_features
|
||||||
from lerobot.policies.act.configuration_act import ACTConfig
|
from lerobot.policies.act.configuration_act import ACTConfig
|
||||||
from lerobot.policies.diffusion.configuration_diffusion import DiffusionConfig
|
from lerobot.policies.diffusion.configuration_diffusion import DiffusionConfig
|
||||||
|
from lerobot.policies.dsrl.configuration_dsrl import DSRLConfig
|
||||||
from lerobot.policies.groot.configuration_groot import GrootConfig
|
from lerobot.policies.groot.configuration_groot import GrootConfig
|
||||||
from lerobot.policies.pi0.configuration_pi0 import PI0Config
|
from lerobot.policies.pi0.configuration_pi0 import PI0Config
|
||||||
from lerobot.policies.pi05.configuration_pi05 import PI05Config
|
from lerobot.policies.pi05.configuration_pi05 import PI05Config
|
||||||
@@ -59,7 +60,7 @@ def get_policy_class(name: str) -> type[PreTrainedPolicy]:
|
|||||||
|
|
||||||
Args:
|
Args:
|
||||||
name: The name of the policy. Supported names are "tdmpc", "diffusion", "act",
|
name: The name of the policy. Supported names are "tdmpc", "diffusion", "act",
|
||||||
"vqbet", "pi0", "pi05", "sac", "reward_classifier", "smolvla".
|
"vqbet", "pi0", "pi05", "sac", "reward_classifier", "smolvla", "dsrl".
|
||||||
|
|
||||||
Returns:
|
Returns:
|
||||||
The policy class corresponding to the given name.
|
The policy class corresponding to the given name.
|
||||||
@@ -103,6 +104,10 @@ def get_policy_class(name: str) -> type[PreTrainedPolicy]:
|
|||||||
from lerobot.policies.smolvla.modeling_smolvla import SmolVLAPolicy
|
from lerobot.policies.smolvla.modeling_smolvla import SmolVLAPolicy
|
||||||
|
|
||||||
return SmolVLAPolicy
|
return SmolVLAPolicy
|
||||||
|
elif name == "dsrl":
|
||||||
|
from lerobot.policies.dsrl.modeling_dsrl import DSRLPolicy
|
||||||
|
|
||||||
|
return DSRLPolicy
|
||||||
elif name == "groot":
|
elif name == "groot":
|
||||||
from lerobot.policies.groot.modeling_groot import GrootPolicy
|
from lerobot.policies.groot.modeling_groot import GrootPolicy
|
||||||
|
|
||||||
@@ -121,7 +126,7 @@ def make_policy_config(policy_type: str, **kwargs) -> PreTrainedConfig:
|
|||||||
Args:
|
Args:
|
||||||
policy_type: The type of the policy. Supported types include "tdmpc",
|
policy_type: The type of the policy. Supported types include "tdmpc",
|
||||||
"diffusion", "act", "vqbet", "pi0", "pi05", "sac", "smolvla",
|
"diffusion", "act", "vqbet", "pi0", "pi05", "sac", "smolvla",
|
||||||
"reward_classifier".
|
"reward_classifier", "dsrl".
|
||||||
**kwargs: Keyword arguments to be passed to the configuration class constructor.
|
**kwargs: Keyword arguments to be passed to the configuration class constructor.
|
||||||
|
|
||||||
Returns:
|
Returns:
|
||||||
@@ -148,6 +153,8 @@ def make_policy_config(policy_type: str, **kwargs) -> PreTrainedConfig:
|
|||||||
return SmolVLAConfig(**kwargs)
|
return SmolVLAConfig(**kwargs)
|
||||||
elif policy_type == "reward_classifier":
|
elif policy_type == "reward_classifier":
|
||||||
return RewardClassifierConfig(**kwargs)
|
return RewardClassifierConfig(**kwargs)
|
||||||
|
elif policy_type == "dsrl":
|
||||||
|
return DSRLConfig(**kwargs)
|
||||||
elif policy_type == "groot":
|
elif policy_type == "groot":
|
||||||
return GrootConfig(**kwargs)
|
return GrootConfig(**kwargs)
|
||||||
else:
|
else:
|
||||||
@@ -321,6 +328,21 @@ def make_pre_post_processors(
|
|||||||
config=policy_cfg,
|
config=policy_cfg,
|
||||||
dataset_stats=kwargs.get("dataset_stats"),
|
dataset_stats=kwargs.get("dataset_stats"),
|
||||||
)
|
)
|
||||||
|
elif isinstance(policy_cfg, DSRLConfig):
|
||||||
|
from lerobot.policies.dsrl.processor_dsrl import make_dsrl_pre_post_processors
|
||||||
|
|
||||||
|
processors = make_dsrl_pre_post_processors(
|
||||||
|
config=policy_cfg,
|
||||||
|
dataset_stats=kwargs.get("dataset_stats"),
|
||||||
|
)
|
||||||
|
|
||||||
|
elif isinstance(policy_cfg, GrootConfig):
|
||||||
|
from lerobot.policies.groot.processor_groot import make_groot_pre_post_processors
|
||||||
|
|
||||||
|
processors = make_groot_pre_post_processors(
|
||||||
|
config=policy_cfg,
|
||||||
|
dataset_stats=kwargs.get("dataset_stats"),
|
||||||
|
)
|
||||||
|
|
||||||
elif isinstance(policy_cfg, GrootConfig):
|
elif isinstance(policy_cfg, GrootConfig):
|
||||||
from lerobot.policies.groot.processor_groot import make_groot_pre_post_processors
|
from lerobot.policies.groot.processor_groot import make_groot_pre_post_processors
|
||||||
|
|||||||
Reference in New Issue
Block a user