Update

[ghstack-poisoned]
pytorch · Feb 7, 2025 · 87a342c · 87a342c
1 parent c6069fe
commit 87a342c
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/torchrl/envs/gym_like.py b/torchrl/envs/gym_like.py
@@ -293,7 +293,8 @@ def read_obs(
 
     def _step(self, tensordict: TensorDictBase) -> TensorDictBase:
         if len(self.action_keys) == 1:
-            action = tensordict.get(self.action_key)
+            # Use brackets to get non-tensor data
+            action = tensordict[self.action_key]
         else:
             action = tensordict.select(*self.action_keys).to_dict()
         if self._convert_actions_to_numpy: