From 0b07af374a275f50d8b3ac644f5c5d31ca77ae0f Mon Sep 17 00:00:00 2001 From: Jiajie Xiao <6828820+jiajiexiao@users.noreply.github.com> Date: Thu, 29 Oct 2020 08:05:38 -0700 Subject: [PATCH] allow tuple action space (#11429) Co-authored-by: Jiajie Xiao --- rllib/models/tf/recurrent_net.py | 5 ++++- rllib/models/torch/recurrent_net.py | 5 ++++- 2 files changed, 8 insertions(+), 2 deletions(-) diff --git a/rllib/models/tf/recurrent_net.py b/rllib/models/tf/recurrent_net.py index 80a8f56d4..9a0ee7438 100644 --- a/rllib/models/tf/recurrent_net.py +++ b/rllib/models/tf/recurrent_net.py @@ -113,7 +113,10 @@ class LSTMWrapper(RecurrentNetwork): self.cell_size = model_config["lstm_cell_size"] self.use_prev_action_reward = model_config[ "lstm_use_prev_action_reward"] - self.action_dim = int(np.product(action_space.shape)) + if action_space.shape is not None: + self.action_dim = int(np.product(action_space.shape)) + else: + self.action_dim = int(len(action_space)) # Add prev-action/reward nodes to input to LSTM. if self.use_prev_action_reward: self.num_outputs += 1 + self.action_dim diff --git a/rllib/models/torch/recurrent_net.py b/rllib/models/torch/recurrent_net.py index f1c962c4b..c500450d0 100644 --- a/rllib/models/torch/recurrent_net.py +++ b/rllib/models/torch/recurrent_net.py @@ -114,7 +114,10 @@ class LSTMWrapper(RecurrentNetwork, nn.Module): self.time_major = model_config.get("_time_major", False) self.use_prev_action_reward = model_config[ "lstm_use_prev_action_reward"] - self.action_dim = int(np.product(action_space.shape)) + if action_space.shape is not None: + self.action_dim = int(np.product(action_space.shape)) + else: + self.action_dim = int(len(action_space)) # Add prev-action/reward nodes to input to LSTM. if self.use_prev_action_reward: self.num_outputs += 1 + self.action_dim