From 94c14dc644674eb44556c27b7779628f204186b7 Mon Sep 17 00:00:00 2001 From: Francesco Capuano Date: Mon, 7 Jul 2025 19:37:49 +0200 Subject: [PATCH 1/2] fix: decorator calls with parentheses --- src/lerobot/policies/act/modeling_act.py | 4 ++-- src/lerobot/policies/diffusion/modeling_diffusion.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 ++-- 8 files changed, 14 insertions(+), 13 deletions(-) diff --git a/src/lerobot/policies/act/modeling_act.py b/src/lerobot/policies/act/modeling_act.py index ed911e9be5..f66c8ae82d 100644 --- a/src/lerobot/policies/act/modeling_act.py +++ b/src/lerobot/policies/act/modeling_act.py @@ -107,7 +107,7 @@ def reset(self): 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 @@ def select_action(self, batch: dict[str, Tensor]) -> Tensor: 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 af40f7a86c..6dad8fb895 100644 --- a/src/lerobot/policies/diffusion/modeling_diffusion.py +++ b/src/lerobot/policies/diffusion/modeling_diffusion.py @@ -99,7 +99,7 @@ def reset(self): 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 @@ def predict_action_chunk(self, batch: dict[str, Tensor]) -> Tensor: 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/pi0/modeling_pi0.py b/src/lerobot/policies/pi0/modeling_pi0.py index 241509d0be..badfb4b8c1 100644 --- a/src/lerobot/policies/pi0/modeling_pi0.py +++ b/src/lerobot/policies/pi0/modeling_pi0.py @@ -260,12 +260,12 @@ def reset(self): 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 d3e576d1c0..0e53bd3497 100644 --- a/src/lerobot/policies/pi0fast/modeling_pi0fast.py +++ b/src/lerobot/policies/pi0fast/modeling_pi0fast.py @@ -192,12 +192,12 @@ def _pi_aloha_encode_actions_inv(self, actions): 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 54ea122a8a..93cfe6c93c 100644 --- a/src/lerobot/policies/sac/modeling_sac.py +++ b/src/lerobot/policies/sac/modeling_sac.py @@ -76,7 +76,7 @@ def reset(self): """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 11bb8bf527..a31e1b0783 100644 --- a/src/lerobot/policies/smolvla/modeling_smolvla.py +++ b/src/lerobot/policies/smolvla/modeling_smolvla.py @@ -413,6 +413,7 @@ def _prepare_batch(self, batch: dict[str, Tensor]) -> dict[str, Tensor]: 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 @@ def predict_action_chunk(self, batch: dict[str, Tensor], noise: Tensor | None = 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 8b70b265df..c276893874 100644 --- a/src/lerobot/policies/tdmpc/modeling_tdmpc.py +++ b/src/lerobot/policies/tdmpc/modeling_tdmpc.py @@ -110,7 +110,7 @@ def reset(self): # 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 c045ccbd2f..59c820a961 100644 --- a/src/lerobot/policies/vqbet/modeling_vqbet.py +++ b/src/lerobot/policies/vqbet/modeling_vqbet.py @@ -124,14 +124,14 @@ def reset(self): 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. From 553f2a9fc2ea8b19335f49f967cbc890247f4e9f Mon Sep 17 00:00:00 2001 From: Francesco Capuano <74058581+fracapuano@users.noreply.github.com> Date: Tue, 8 Jul 2025 12:56:31 +0200 Subject: [PATCH 2/2] fix no grad for normalize too Signed-off-by: Francesco Capuano <74058581+fracapuano@users.noreply.github.com> --- src/lerobot/policies/normalize.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/lerobot/policies/normalize.py b/src/lerobot/policies/normalize.py index 9cc94b9298..119055873c 100644 --- a/src/lerobot/policies/normalize.py +++ b/src/lerobot/policies/normalize.py @@ -149,7 +149,7 @@ def __init__( 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 @@ def __init__( 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():