From a5e0aae13a3efd0080ac7ab6b461980d644014ab Mon Sep 17 00:00:00 2001 From: Francesco Capuano <74058581+fracapuano@users.noreply.github.com> Date: Tue, 8 Jul 2025 13:08:32 +0200 Subject: [PATCH] Fixes `@torch.no_grad()` usage (#1455) * fix: decorator calls with parentheses * fix no grad for normalize too Signed-off-by: Francesco Capuano <74058581+fracapuano@users.noreply.github.com> --------- Signed-off-by: Francesco Capuano <74058581+fracapuano@users.noreply.github.com> --- src/lerobot/policies/act/modeling_act.py | 4 ++-- src/lerobot/policies/diffusion/modeling_diffusion.py | 4 ++-- src/lerobot/policies/normalize.py | 4 ++-- src/lerobot/policies/pi0/modeling_pi0.py | 4 ++-- src/lerobot/policies/pi0fast/modeling_pi0fast.py | 4 ++-- src/lerobot/policies/sac/modeling_sac.py | 2 +- src/lerobot/policies/smolvla/modeling_smolvla.py | 3 ++- src/lerobot/policies/tdmpc/modeling_tdmpc.py | 2 +- src/lerobot/policies/vqbet/modeling_vqbet.py | 4 ++-- 9 files changed, 16 insertions(+), 15 deletions(-) diff --git a/src/lerobot/policies/act/modeling_act.py b/src/lerobot/policies/act/modeling_act.py index ed911e9b..f66c8ae8 100644 --- a/src/lerobot/policies/act/modeling_act.py +++ b/src/lerobot/policies/act/modeling_act.py @@ -107,7 +107,7 @@ class ACTPolicy(PreTrainedPolicy): else: self._action_queue = deque([], maxlen=self.config.n_action_steps) - @torch.no_grad + @torch.no_grad() def select_action(self, batch: dict[str, Tensor]) -> Tensor: """Select a single action given environment observations. @@ -132,7 +132,7 @@ class ACTPolicy(PreTrainedPolicy): self._action_queue.extend(actions.transpose(0, 1)) return self._action_queue.popleft() - @torch.no_grad + @torch.no_grad() def predict_action_chunk(self, batch: dict[str, Tensor]) -> Tensor: """Predict a chunk of actions given environment observations.""" self.eval() diff --git a/src/lerobot/policies/diffusion/modeling_diffusion.py b/src/lerobot/policies/diffusion/modeling_diffusion.py index af40f7a8..6dad8fb8 100644 --- a/src/lerobot/policies/diffusion/modeling_diffusion.py +++ b/src/lerobot/policies/diffusion/modeling_diffusion.py @@ -99,7 +99,7 @@ class DiffusionPolicy(PreTrainedPolicy): if self.config.env_state_feature: self._queues["observation.environment_state"] = deque(maxlen=self.config.n_obs_steps) - @torch.no_grad + @torch.no_grad() def predict_action_chunk(self, batch: dict[str, Tensor]) -> Tensor: """Predict a chunk of actions given environment observations.""" # stack n latest observations from the queue @@ -111,7 +111,7 @@ class DiffusionPolicy(PreTrainedPolicy): return actions - @torch.no_grad + @torch.no_grad() def select_action(self, batch: dict[str, Tensor]) -> Tensor: """Select a single action given environment observations. diff --git a/src/lerobot/policies/normalize.py b/src/lerobot/policies/normalize.py index 9cc94b92..11905587 100644 --- a/src/lerobot/policies/normalize.py +++ b/src/lerobot/policies/normalize.py @@ -149,7 +149,7 @@ class Normalize(nn.Module): setattr(self, "buffer_" + key.replace(".", "_"), buffer) # TODO(rcadene): should we remove torch.no_grad? - @torch.no_grad + @torch.no_grad() def forward(self, batch: dict[str, Tensor]) -> dict[str, Tensor]: # TODO: Remove this shallow copy batch = dict(batch) # shallow copy avoids mutating the input batch @@ -224,7 +224,7 @@ class Unnormalize(nn.Module): setattr(self, "buffer_" + key.replace(".", "_"), buffer) # TODO(rcadene): should we remove torch.no_grad? - @torch.no_grad + @torch.no_grad() def forward(self, batch: dict[str, Tensor]) -> dict[str, Tensor]: batch = dict(batch) # shallow copy avoids mutating the input batch for key, ft in self.features.items(): diff --git a/src/lerobot/policies/pi0/modeling_pi0.py b/src/lerobot/policies/pi0/modeling_pi0.py index 241509d0..badfb4b8 100644 --- a/src/lerobot/policies/pi0/modeling_pi0.py +++ b/src/lerobot/policies/pi0/modeling_pi0.py @@ -260,12 +260,12 @@ class PI0Policy(PreTrainedPolicy): def get_optim_params(self) -> dict: return self.parameters() - @torch.no_grad + @torch.no_grad() def predict_action_chunk(self, batch: dict[str, Tensor]) -> Tensor: """Predict a chunk of actions given environment observations.""" raise NotImplementedError("Currently not implemented for PI0") - @torch.no_grad + @torch.no_grad() def select_action(self, batch: dict[str, Tensor], noise: Tensor | None = None) -> Tensor: """Select a single action given environment observations. diff --git a/src/lerobot/policies/pi0fast/modeling_pi0fast.py b/src/lerobot/policies/pi0fast/modeling_pi0fast.py index d3e576d1..0e53bd34 100644 --- a/src/lerobot/policies/pi0fast/modeling_pi0fast.py +++ b/src/lerobot/policies/pi0fast/modeling_pi0fast.py @@ -192,12 +192,12 @@ class PI0FASTPolicy(PreTrainedPolicy): actions[:, :, motor_idx] = aloha_gripper_from_angular_inv(actions[:, :, motor_idx]) return actions - @torch.no_grad + @torch.no_grad() def predict_action_chunk(self, batch: dict[str, Tensor]) -> Tensor: """Predict a chunk of actions given environment observations.""" raise NotImplementedError("Currently not implemented for PI0FAST") - @torch.no_grad + @torch.no_grad() def select_action(self, batch: dict[str, Tensor]) -> Tensor: """Select a single action given environment observations. diff --git a/src/lerobot/policies/sac/modeling_sac.py b/src/lerobot/policies/sac/modeling_sac.py index 54ea122a..93cfe6c9 100644 --- a/src/lerobot/policies/sac/modeling_sac.py +++ b/src/lerobot/policies/sac/modeling_sac.py @@ -76,7 +76,7 @@ class SACPolicy( """Reset the policy""" pass - @torch.no_grad + @torch.no_grad() def predict_action_chunk(self, batch: dict[str, Tensor]) -> Tensor: """Predict a chunk of actions given environment observations.""" raise NotImplementedError("SACPolicy does not support action chunking. It returns single actions!") diff --git a/src/lerobot/policies/smolvla/modeling_smolvla.py b/src/lerobot/policies/smolvla/modeling_smolvla.py index 11bb8bf5..a31e1b07 100644 --- a/src/lerobot/policies/smolvla/modeling_smolvla.py +++ b/src/lerobot/policies/smolvla/modeling_smolvla.py @@ -413,6 +413,7 @@ class SmolVLAPolicy(PreTrainedPolicy): return batch + @torch.no_grad() def predict_action_chunk(self, batch: dict[str, Tensor], noise: Tensor | None = None) -> Tensor: self.eval() @@ -422,7 +423,7 @@ class SmolVLAPolicy(PreTrainedPolicy): actions = self._get_action_chunk(batch, noise) return actions - @torch.no_grad + @torch.no_grad() def select_action(self, batch: dict[str, Tensor], noise: Tensor | None = None) -> Tensor: """Select a single action given environment observations. diff --git a/src/lerobot/policies/tdmpc/modeling_tdmpc.py b/src/lerobot/policies/tdmpc/modeling_tdmpc.py index 8b70b265..c2768938 100644 --- a/src/lerobot/policies/tdmpc/modeling_tdmpc.py +++ b/src/lerobot/policies/tdmpc/modeling_tdmpc.py @@ -110,7 +110,7 @@ class TDMPCPolicy(PreTrainedPolicy): # CEM for the next step. self._prev_mean: torch.Tensor | None = None - @torch.no_grad + @torch.no_grad() def predict_action_chunk(self, batch: dict[str, Tensor]) -> Tensor: """Predict a chunk of actions given environment observations.""" batch = {key: torch.stack(list(self._queues[key]), dim=1) for key in batch if key in self._queues} diff --git a/src/lerobot/policies/vqbet/modeling_vqbet.py b/src/lerobot/policies/vqbet/modeling_vqbet.py index c045ccbd..59c820a9 100644 --- a/src/lerobot/policies/vqbet/modeling_vqbet.py +++ b/src/lerobot/policies/vqbet/modeling_vqbet.py @@ -124,14 +124,14 @@ class VQBeTPolicy(PreTrainedPolicy): ACTION: deque(maxlen=self.config.action_chunk_size), } - @torch.no_grad + @torch.no_grad() def predict_action_chunk(self, batch: dict[str, Tensor]) -> Tensor: batch = {k: torch.stack(list(self._queues[k]), dim=1) for k in batch if k in self._queues} actions = self.vqbet(batch, rollout=True)[:, : self.config.action_chunk_size] actions = self.unnormalize_outputs({ACTION: actions})[ACTION] return actions - @torch.no_grad + @torch.no_grad() def select_action(self, batch: dict[str, Tensor]) -> Tensor: """Select a single action given environment observations.