Skip to content

Commit 3bad905

Browse files
author
Vincent Moens
committed
[BugFix] Fix single action pass to gym when action key is not "action"
ghstack-source-id: b91b9e1 Pull-Request-resolved: #2942
1 parent 46ccbae commit 3bad905

File tree

1 file changed

+5
-1
lines changed

1 file changed

+5
-1
lines changed

torchrl/envs/gym_like.py

Lines changed: 5 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -231,7 +231,11 @@ def read_action(self, action):
231231
Returns: an action in a format compatible with the contained environment.
232232
233233
"""
234-
return self.action_spec.to_numpy(action, safe=False)
234+
action_spec = self.full_action_spec
235+
action_keys = self.action_keys
236+
if len(action_keys) == 1:
237+
action_spec = action_spec[action_keys[0]]
238+
return action_spec.to_numpy(action, safe=False)
235239

236240
def read_done(
237241
self,

0 commit comments

Comments
 (0)