amend

Signed-off-by: Matteo Bettini <matbet@meta.com>
pytorch · vmoens · Jul 6, 2023 · Jun 21, 2023 · Jun 21, 2023 · Jun 21, 2023
commit 15e6036adda0fca86c97fb53635e54007e18ac78
diff --git a/torchrl/modules/tensordict_module/exploration.py b/torchrl/modules/tensordict_module/exploration.py
@@ -574,7 +574,7 @@ def add_sample(
 
         # Get the nested tensordict where the action lives
         if isinstance(self.key, tuple) and len(self.key) > 1:
-            action_tensordict = tensordict.get(self.key[-1])
+            action_tensordict = tensordict.get(self.key[:-1])
         else:
             action_tensordict = tensordict