diff --git a/rllib/models/tf/tf_action_dist.py b/rllib/models/tf/tf_action_dist.py index d3a56f92d..585d8eb3f 100644 --- a/rllib/models/tf/tf_action_dist.py +++ b/rllib/models/tf/tf_action_dist.py @@ -230,7 +230,7 @@ class DiagGaussian(TFActionDistribution): @override(ActionDistribution) def logp(self, x): return -0.5 * tf.reduce_sum( - tf.square((x - self.mean) / self.std), axis=1) - \ + tf.square((tf.to_float(x) - self.mean) / self.std), axis=1) - \ 0.5 * np.log(2.0 * np.pi) * tf.to_float(tf.shape(x)[1]) - \ tf.reduce_sum(self.log_std, axis=1) diff --git a/rllib/policy/eager_tf_policy.py b/rllib/policy/eager_tf_policy.py index 5760d2eb6..162731de9 100644 --- a/rllib/policy/eager_tf_policy.py +++ b/rllib/policy/eager_tf_policy.py @@ -5,6 +5,7 @@ It supports both traced and non-traced eager execution modes.""" import functools import logging import numpy as np +from gym.spaces import Tuple, Dict from ray.util.debug import log_once from ray.rllib.models.catalog import ModelCatalog @@ -586,10 +587,17 @@ def build_eager_tf_policy(name, SampleBatch.NEXT_OBS: np.array( [self.observation_space.sample()]), SampleBatch.DONES: np.array([False], dtype=np.bool), - SampleBatch.ACTIONS: tf.nest.map_structure( - lambda c: np.array([c]), self.action_space.sample()), SampleBatch.REWARDS: np.array([0], dtype=np.float32), } + if isinstance(self.action_space, Tuple) or isinstance( + self.action_space, Dict): + dummy_batch[SampleBatch.ACTIONS] = [ + flatten_to_single_ndarray(self.action_space.sample()) + ] + else: + dummy_batch[SampleBatch.ACTIONS] = tf.nest.map_structure( + lambda c: np.array([c]), self.action_space.sample()) + if obs_include_prev_action_reward: dummy_batch.update({ SampleBatch.PREV_ACTIONS: dummy_batch[SampleBatch.ACTIONS], diff --git a/rllib/tests/test_supported_spaces.py b/rllib/tests/test_supported_spaces.py index 488576b92..59d30b79a 100644 --- a/rllib/tests/test_supported_spaces.py +++ b/rllib/tests/test_supported_spaces.py @@ -48,7 +48,7 @@ OBSERVATION_SPACES_TO_TEST = { } -def check_support(alg, config, check_bounds=False): +def check_support(alg, config, check_bounds=False, tfe=False): config["log_level"] = "ERROR" def _do_check(alg, config, a_name, o_name): @@ -95,7 +95,10 @@ def check_support(alg, config, check_bounds=False): pass print(stat) - for _ in framework_iterator(config, frameworks=("tf", "torch")): + frameworks = ("tf", "torch") + if tfe: + frameworks += ("tfe", ) + for _ in framework_iterator(config, frameworks=frameworks): # Check all action spaces. for a_name, action_space in ACTION_SPACES_TO_TEST.items(): _do_check(alg, config, a_name, "discrete") @@ -141,7 +144,7 @@ class ModelSupportedSpaces(unittest.TestCase): def test_dqn(self): config = {"timesteps_per_iteration": 1} - check_support("DQN", config) + check_support("DQN", config, tfe=True) def test_es(self): check_support( @@ -163,11 +166,11 @@ class ModelSupportedSpaces(unittest.TestCase): "rollout_fragment_length": 10, "sgd_minibatch_size": 1, } - check_support("PPO", config, check_bounds=True) + check_support("PPO", config, check_bounds=True, tfe=True) def test_pg(self): config = {"num_workers": 1, "optimizer": {}} - check_support("PG", config, check_bounds=True) + check_support("PG", config, check_bounds=True, tfe=True) def test_sac(self): check_support("SAC", {}, check_bounds=True)