admin/lerobot
1
0
Fork 0
mirror of https://github.com/huggingface/lerobot.git synced 2026-07-21 17:01:53 +00:00
Code Issues Packages Projects Releases Wiki Activity

refactor: update SACAlgorithm to pass action_dim to _init_critics and fix encoder reference

Browse Source
This commit is contained in:
Khalil Meftah
2026-04-13 18:31:17 +02:00
parent 7b0bdf2a98
commit ee0814ef60
1 changed files with 5 additions and 5 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/lerobot/rl/algorithms/sac/sac_algorithm.py
+5 -5
View File
@@ -58,16 +58,16 @@ class SACAlgorithm(RLAlgorithm):
self.optimizers: dict[str, Optimizer] = {} self.optimizers: dict[str, Optimizer] = {}
self._optimization_step: int = 0 self._optimization_step: int = 0
self._init_critics() action_dim = self.policy.config.output_features[ACTION].shape[0]
self._init_critics(action_dim)
self._init_temperature() self._init_temperature()
self._device = torch.device(self.policy.config.device) self._device = torch.device(self.policy.config.device)
self._move_to_device() self._move_to_device()
def _init_critics(self) -> None: def _init_critics(self, action_dim) -> None:
"""Build critic ensemble, targets.""" """Build critic ensemble, targets."""
encoder = self.policy.encoder_critic encoder = self.policy.encoder_critic
action_dim = self.policy.config.output_features[ACTION].shape[0]
heads = [ heads = [
CriticHead( CriticHead(
@@ -76,7 +76,7 @@ class SACAlgorithm(RLAlgorithm):
) )
for _ in range(self.config.num_critics) for _ in range(self.config.num_critics)
] ]
self.critic_ensemble = CriticEnsemble(encoder=self.encoder_critic, ensemble=heads) self.critic_ensemble = CriticEnsemble(encoder=encoder, ensemble=heads)
target_heads = [ target_heads = [
CriticHead( CriticHead(
input_dim=encoder.output_dim + action_dim, input_dim=encoder.output_dim + action_dim,
@@ -84,7 +84,7 @@ class SACAlgorithm(RLAlgorithm):
) )
for _ in range(self.config.num_critics) for _ in range(self.config.num_critics)
] ]
self.critic_target = CriticEnsemble(encoder=self.encoder_critic, ensemble=target_heads) self.critic_target = CriticEnsemble(encoder=encoder, ensemble=target_heads)
self.critic_target.load_state_dict(self.critic_ensemble.state_dict()) self.critic_target.load_state_dict(self.critic_ensemble.state_dict())
# TODO(Khalil): Investigate and fix torch.compile # TODO(Khalil): Investigate and fix torch.compile