From 87a342cab1f4a73283b8c988efdc56a66e6dadf5 Mon Sep 17 00:00:00 2001
From: Vincent Moens <vmoens@meta.com>
Date: Fri, 7 Feb 2025 11:24:04 +0000
Subject: [PATCH] Update

[ghstack-poisoned]
---
 torchrl/envs/gym_like.py | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/torchrl/envs/gym_like.py b/torchrl/envs/gym_like.py
index 5e6dd55be8b..aa1883958ca 100644
--- 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: