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

fix: move algorithm-owned modules to the policy device

Browse Source
This commit is contained in:
Khalil Meftah
2026-03-18 15:27:41 +01:00
parent 1f5487eea8
commit d3e6f14d4f
1 changed files with 9 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/lerobot/rl/algorithms/sac/sac_algorithm.py
+9
View File
@@ -74,6 +74,7 @@ class SACAlgorithm(RLAlgorithm):
self._init_critic_encoder() self._init_critic_encoder()
self._init_critics() self._init_critics()
self._init_temperature() self._init_temperature()
self._move_to_device()
def _init_critic_encoder(self) -> None: def _init_critic_encoder(self) -> None:
"""Build or share the encoder used by critics.""" """Build or share the encoder used by critics."""
@@ -131,6 +132,14 @@ class SACAlgorithm(RLAlgorithm):
dim = action_dim + (1 if self.config.num_discrete_actions is not None else 0) dim = action_dim + (1 if self.config.num_discrete_actions is not None else 0)
self.target_entropy = -np.prod(dim) / 2 self.target_entropy = -np.prod(dim) / 2
def _move_to_device(self) -> None:
"""Move algorithm-owned modules to the policy device."""
self.critic_ensemble.to(self._device)
self.critic_target.to(self._device)
self.log_alpha = nn.Parameter(self.log_alpha.data.to(self._device))
if hasattr(self, "discrete_critic_target"):
self.discrete_critic_target.to(self._device)
@property @property
def temperature(self) -> float: def temperature(self) -> float:
return self.log_alpha.exp().item() return self.log_alpha.exp().item()