From 7bd0d62ce5cf8cc63812eabde687b3e6539dd37b Mon Sep 17 00:00:00 2001 From: Adil Zouitine Date: Mon, 7 Jul 2025 17:56:34 +0200 Subject: [PATCH] refactor(policies): Remove unnormalization step from action predictions - Eliminated the unnormalization of actions in both `TDMPCPolicy` and `VQBeTPolicy` classes to streamline action prediction. - This change improves code clarity and aligns with recent updates to policy processors. --- src/lerobot/policies/tdmpc/modeling_tdmpc.py | 1 - src/lerobot/policies/vqbet/modeling_vqbet.py | 1 - 2 files changed, 2 deletions(-) diff --git a/src/lerobot/policies/tdmpc/modeling_tdmpc.py b/src/lerobot/policies/tdmpc/modeling_tdmpc.py index 92518d4ad..0ee2e3a43 100644 --- a/src/lerobot/policies/tdmpc/modeling_tdmpc.py +++ b/src/lerobot/policies/tdmpc/modeling_tdmpc.py @@ -129,7 +129,6 @@ class TDMPCPolicy(PreTrainedPolicy): actions = torch.clamp(actions, -1, +1) - actions = self.unnormalize_outputs({ACTION: actions})[ACTION] return actions @torch.no_grad() diff --git a/src/lerobot/policies/vqbet/modeling_vqbet.py b/src/lerobot/policies/vqbet/modeling_vqbet.py index 06c9bcfb9..41e021d8f 100644 --- a/src/lerobot/policies/vqbet/modeling_vqbet.py +++ b/src/lerobot/policies/vqbet/modeling_vqbet.py @@ -118,7 +118,6 @@ class VQBeTPolicy(PreTrainedPolicy): 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()