Repository: ray-project/ray

Issue:
[rllib] Crash when using hybrid action space env and eager mode.
<!--Please include [tune], [rllib], [autoscaler] etc. in the issue title if relevant-->

### What is the problem?

RLlib crashes on environment with **hybrid action space** (discrete and continuous) and **eager mode enabled**
It works well if eager mode is turned off.

_Python3.7
tensorflow==2.2.0
ray==0.9.0.dev0 (used the nightly version)
Ubuntu 18.04_

### Script to reproduce

Modified from `ray/rllib/examples/custom_env.py` adding `Box` to the action space to make it hybrid
```
import gym
import numpy as np
import ray
from gym.spaces import Discrete, Box
from ray import tune
from ray.rllib.utils.framework import try_import_tf

tf = try_import_tf()


class SimpleCorridor(gym.Env):
    """Example of a custom env in which you have to walk down a corridor.

    You can configure the length of the corridor via the env config."""

    def __init__(self, config):
        self.end_pos = config["corridor_length"]
        self.cur_pos = 0
        self.action_space = gym.spaces.Tuple([Discrete(2), Box(0.0, 1.0, shape=(1,))])  # Hybrid action space
        self.observation_space = Box(
            0.0, self.end_pos, shape=(1,), dtype=np.float32)

    def reset(self):
        self.cur_pos = 0
        return [self.cur_pos]

    def step(self, action):
        action = action[0]
        assert action in [0, 1], action
        if action == 0 and self.cur_pos > 0:
            self.cur_pos -= 1
        elif action == 1:
            self.cur_pos += 1
        done = self.cur_pos >= self.end_pos
        return [self.cur_pos], 1.0 if done else -0.1, done, {}


if __name__ == "__main__":
    ray.init()

    config = {
        "env": SimpleCorridor,  # or "corridor" if registered above
        "env_config": {
            "corridor_length": 5,
        },
        "model": {
        },
        "eager": True,  # If True crashes
        "vf_share_layers": True,
        "lr": 1e-2,
        "num_workers": 1,
        "framework": "tf",
    }

    results = tune.run("PPO", config=config)

    ray.shutdown()
```
Outputs following error
```
  File "/home/coac/anaconda3/envs/rllib/lib/python3.7/site-packages/ray/rllib/agents/trainer_template.py", line 90, in __init__
    Trainer.__init__(self, config, env, logger_creator)
  File "/home/coac/anaconda3/envs/rllib/lib/python3.7/site-packages/ray/rllib/agents/trainer.py", line 444, in __init__
    super().__init__(config, logger_creator)
  File "/home/coac/anaconda3/envs/rllib/lib/python3.7/site-packages/ray/tune/trainable.py", line 174, in __init__
    self._setup(copy.deepcopy(self.config))
  File "/home/coac/anaconda3/envs/rllib/lib/python3.7/site-packages/ray/rllib/agents/trainer.py", line 613, in _setup
    self._init(self.config, self.env_creator)
  File "/home/coac/anaconda3/envs/rllib/lib/python3.7/site-packages/ray/rllib/agents/trainer_template.py", line 115, in _init
    self.config["num_workers"])
  File "/home/coac/anaconda3/envs/rllib/lib/python3.7/site-packages/ray/rllib/agents/trainer.py", line 684, in _make_workers
    logdir=self.logdir)
  File "/home/coac/anaconda3/envs/rllib/lib/python3.7/site-packages/ray/rllib/evaluation/worker_set.py", line 59, in __init__
    RolloutWorker, env_creator, policy, 0, self._local_config)
  File "/home/coac/anaconda3/envs/rllib/lib/python3.7/site-packages/ray/rllib/evaluation/worker_set.py", line 282, in _make_worker
    extra_python_environs=extra_python_environs)
  File "/home/coac/anaconda3/envs/rllib/lib/python3.7/site-packages/ray/rllib/evaluation/rollout_worker.py", line 393, in __init__
    policy_dict, policy_config)
  File "/home/coac/anaconda3/envs/rllib/lib/python3.7/site-packages/ray/rllib/evaluation/rollout_worker.py", line 930, in _build_policy_map
    policy_map[name] = cls(obs_space, act_space, merged_conf)
  File "/home/coac/anaconda3/envs/rllib/lib/python3.7/site-packages/ray/rllib/policy/eager_tf_policy.py", line 258, in __init__
    self._initialize_loss_with_dummy_batch()
  File "/home/coac/anaconda3/envs/rllib/lib/python3.7/site-packages/ray/rllib/policy/eager_tf_policy.py", line 626, in _initialize_loss_with_dummy_batch
    explore=False)
  File "/home/coac/anaconda3/envs/rllib/lib/python3.7/site-packages/ray/rllib/policy/eager_tf_policy.py", line 60, in _func
    return func(*args, **kwargs)
  File "/home/coac/anaconda3/envs/rllib/lib/python3.7/site-packages/ray/rllib/policy/eager_tf_policy.py", line 68, in _func
    out = func(*args, **kwargs)
  File "/home/coac/anaconda3/envs/rllib/lib/python3.7/site-packages/ray/rllib/policy/eager_tf_policy.py", line 335, in compute_actions
    tf.convert_to_tensor(prev_action_batch)
  File "/home/coac/anaconda3/envs/rllib/lib/python3.7/site-packages/tensorflow/python/framework/ops.py", line 1217, in convert_to_tensor_v1
    return convert_to_tensor_v2(value, dtype, preferred_dtype, name)
  File "/home/coac/anaconda3/envs/rllib/lib/python3.7/site-packages/tensorflow/python/framework/ops.py", line 1283, in convert_to_tensor_v2
    as_ref=False)
  File "/home/coac/anaconda3/envs/rllib/lib/python3.7/site-packages/tensorflow/python/framework/ops.py", line 1341, in convert_to_tensor
    ret = conversion_func(value, dtype=dtype, name=name, as_ref=as_ref)
  File "/home/coac/anaconda3/envs/rllib/lib/python3.7/site-packages/tensorflow/python/ops/array_ops.py", line 1455, in _autopacking_conversion_function
    return _autopacking_helper(v, dtype, name or "packed")
  File "/home/coac/anaconda3/envs/rllib/lib/python3.7/site-packages/tensorflow/python/ops/array_ops.py", line 1361, in _autopacking_helper
    return gen_array_ops.pack(list_or_tuple, name=name)
  File "/home/coac/anaconda3/envs/rllib/lib/python3.7/site-packages/tensorflow/python/ops/gen_array_ops.py", line 6337, in pack
    _ops.raise_from_not_ok_status(e, name)
  File "/home/coac/anaconda3/envs/rllib/lib/python3.7/site-packages/tensorflow/python/framework/ops.py", line 6653, in raise_from_not_ok_status
    six.raise_from(core._status_to_exception(e.code, message), None)
  File "<string>", line 3, in raise_from
tensorflow.python.framework.errors_impl.InvalidArgumentError: cannot compute Pack as input #1(zero-based) was expected to be a int64 tensor but is a float tensor [Op:Pack] name: packed

```

If we cannot run your script, we cannot fix your issue.

- [x] I have verified my script runs in a clean environment and reproduces the issue.
- [x] I have verified the issue also occurs with the [latest wheels](https://docs.ray.io/en/latest/installation.html).



Discussion:
I think this is a matter of missing tf.cast() to float32 somewhere. Unfortunately, eager mode has slightly different type conversions than graph mode. Can you try finding the right place to add this conversion?
Sure, I will do a PR

Resolving patch:
diff --git a/rllib/models/tf/tf_action_dist.py b/rllib/models/tf/tf_action_dist.py
--- a/rllib/models/tf/tf_action_dist.py
+++ b/rllib/models/tf/tf_action_dist.py
@@ -230,7 +230,7 @@ def deterministic_sample(self):
     @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
--- a/rllib/policy/eager_tf_policy.py
+++ b/rllib/policy/eager_tf_policy.py
@@ -5,6 +5,7 @@
 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 _initialize_loss_with_dummy_batch(self):
                 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],


Test patch:
diff --git a/rllib/tests/test_supported_spaces.py b/rllib/tests/test_supported_spaces.py
--- a/rllib/tests/test_supported_spaces.py
+++ b/rllib/tests/test_supported_spaces.py
@@ -48,7 +48,7 @@
 }
 
 
-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 _do_check(alg, config, a_name, o_name):
                     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 @@ def test_ddpg(self):
 
     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 @@ def test_ppo(self):
             "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)