Repository: ray-project/ray

Issue:
[rllib] Make QMix support complex observation spaces
<!--Please include [tune], [rllib], [autoscaler] etc. in the issue title if relevant-->

### Describe your feature request

The QMix example and the default model only uses a simple MultiDiscrete observation space. It would be nice if more abstract observation spaces and custom models would be supported. The following example using a nested dict observation space ([based on the two step game](https://github.com/ray-project/ray/blob/master/rllib/examples/env/two_step_game.py)) fails [here](https://github.com/ray-project/ray/blob/master/rllib/agents/qmix/qmix_policy.py#L482) (I am aware that the model would have to be adapted as well, but that does not seem to be the cause of this error):

```
import argparse
from gym.spaces import Tuple, MultiDiscrete, Dict, Discrete

import ray
from ray import tune
from ray.tune import register_env, grid_search
from ray.rllib.env.multi_agent_env import ENV_STATE
from ray.rllib.examples.env.two_step_game import TwoStepGame
from ray.rllib.utils.test_utils import check_learning_achieved

from gym.spaces import MultiDiscrete, Dict, Discrete
import numpy as np

from ray.rllib.env.multi_agent_env import MultiAgentEnv, ENV_STATE


class TwoStepGame(MultiAgentEnv):
    action_space = Discrete(2)

    def __init__(self, env_config):
        self.state = None
        self.agent_1 = 0
        self.agent_2 = 1
        # MADDPG emits action logits instead of actual discrete actions
        self.actions_are_logits = env_config.get("actions_are_logits", False)
        self.one_hot_state_encoding = env_config.get("one_hot_state_encoding",
                                                     False)
        self.with_state = env_config.get("separate_state_space", False)

        if not self.one_hot_state_encoding:
            self.observation_space = Discrete(6)
            self.with_state = False
        else:
            # Each agent gets the full state (one-hot encoding of which of the
            # three states are active) as input with the receiving agent's
            # ID (1 or 2) concatenated onto the end.
            if self.with_state:
                self.observation_space = Dict({
                    "obs": MultiDiscrete([2, 2, 2, 3]),
                    ENV_STATE: MultiDiscrete([2, 2, 2])
                })
            else:
                self.observation_space = MultiDiscrete([2, 2, 2, 3])

    def reset(self):
        self.state = np.array([1, 0, 0])
        return self._obs()

    def step(self, action_dict):
        if self.actions_are_logits:
            action_dict = {
                k: np.random.choice([0, 1], p=v)
                for k, v in action_dict.items()
            }

        state_index = np.flatnonzero(self.state)
        if state_index == 0:
            action = action_dict[self.agent_1]
            assert action in [0, 1], action
            if action == 0:
                self.state = np.array([0, 1, 0])
            else:
                self.state = np.array([0, 0, 1])
            global_rew = 0
            done = False
        elif state_index == 1:
            global_rew = 7
            done = True
        else:
            if action_dict[self.agent_1] == 0 and action_dict[self.
                                                              agent_2] == 0:
                global_rew = 0
            elif action_dict[self.agent_1] == 1 and action_dict[self.
                                                                agent_2] == 1:
                global_rew = 8
            else:
                global_rew = 1
            done = True

        rewards = {
            self.agent_1: global_rew / 2.0,
            self.agent_2: global_rew / 2.0
        }
        obs = self._obs()
        dones = {"__all__": done}
        infos = {}
        return obs, rewards, dones, infos

    def _obs(self):
        if self.with_state:
            return {
                self.agent_1: {
                    "obs": {
                        "data": self.agent_1_obs(),
                        "test": 1
                    },
                    ENV_STATE: self.state
                },
                self.agent_2: {
                    "obs": {
                        "data": self.agent_2_obs(),
                        "test": 2
                    },
                    ENV_STATE: self.state
                }
            }
        else:
            return {
                self.agent_1: self.agent_1_obs(),
                self.agent_2: self.agent_2_obs()
            }

    def agent_1_obs(self):
        if self.one_hot_state_encoding:
            return np.concatenate([self.state, [1]])
        else:
            return np.flatnonzero(self.state)[0]

    def agent_2_obs(self):
        if self.one_hot_state_encoding:
            return np.concatenate([self.state, [2]])
        else:
            return np.flatnonzero(self.state)[0] + 3

if __name__ == "__main__":
    grouping = {
        "group_1": [0, 1],
    }
    obs_space = Tuple([
        Dict({
            "obs": Dict({
                'data': MultiDiscrete([2, 2, 2, 3]),
                'test': Discrete(5)
            }),
            ENV_STATE: MultiDiscrete([2, 2, 2])
        }),
        Dict({
            "obs": Dict({
                'data': MultiDiscrete([2, 2, 2, 3]),
                'test': Discrete(5)
            }),
            ENV_STATE: MultiDiscrete([2, 2, 2])
        }),
    ])
    act_space = Tuple([
        TwoStepGame.action_space,
        TwoStepGame.action_space,
    ])
    register_env(
        "grouped_twostep",
        lambda config: TwoStepGame(config).with_agent_groups(
            grouping, obs_space=obs_space, act_space=act_space))

    ray.init()

    results = tune.run("QMIX", config={
        "rollout_fragment_length": 4,
        "train_batch_size": 32,
        "exploration_fraction": .4,
        "exploration_final_eps": 0.0,
        "num_workers": 0,
        "mixer": "qmix",
        "env_config": {
            "separate_state_space": True,
            "one_hot_state_encoding": True
        },
        "use_pytorch": True,
        "env": "grouped_twostep"
    })

    ray.shutdown()
```
```
2020-05-20 15:40:03,477 ERROR trial_runner.py:519 -- Trial QMIX_grouped_twostep_c960c_00002: Error processing event.             
Traceback (most recent call last):
  File "[...]/ray/tune/trial_runner.py", line 467, in _process_trial        
    result = self.trial_executor.fetch_result(trial)
  File "[...]/ray/tune/ray_trial_executor.py", line 430, in fetch_result         
    result = ray.get(trial_future[0], DEFAULT_GET_TIMEOUT)         
  File "[...]/ray/worker.py", line 1516, in get                                      
    raise value.as_instanceof_cause()                              
ray.exceptions.RayTaskError(ValueError): ray::QMIX.train() (pid=54834, ip=128.232.69.20)
  File "python/ray/_raylet.pyx", line 460, in ray._raylet.execute_task
  File "python/ray/_raylet.pyx", line 414, in ray._raylet.execute_task.function_executor
  File "[...]/ray/rllib/agents/trainer.py", line 504, in train
    raise e                     
  File "[...]/ray/rllib/agents/trainer.py", line 490, in train
    result = Trainable.train(self)          
  File "[...]/ray/tune/trainable.py", line 260, in train
    result = self._train()                                       
  File "[...]/ray/rllib/agents/trainer_template.py", line 138, in _train
    return self._train_exec_impl()                               
  File "[...]/ray/rllib/agents/trainer_template.py", line 173, in _train_exec_impl
    res = next(self.train_exec_impl)                             
  File "[...]/ray/util/iter.py", line 689, in __next__
    return next(self.built_iterator)                             
  File "[...]/ray/util/iter.py", line 702, in apply_foreach
    for item in it:                                                                                                                                        
  File "[...]/ray/util/iter.py", line 772, in apply_filter                                        
    for item in it:                                                                                                                                        
  File "[...]/ray/util/iter.py", line 772, in apply_filter                                        
    for item in it:                                                                                                                                        
  File "[...]/ray/util/iter.py", line 702, in apply_foreach                                       
    for item in it:                                                                                                                                        
  File "[...]/ray/util/iter.py", line 772, in apply_filter
    for item in it:               
  File "[...]/ray/util/iter.py", line 977, in build_union
    item = next(it)         
  File "[...]/ray/util/iter.py", line 689, in __next__
    return next(self.built_iterator)                             
  File "[...]/ray/util/iter.py", line 702, in apply_foreach                                         
    for item in it:                                                       
  File "[...]/ray/util/iter.py", line 702, in apply_foreach
    for item in it:                                                                            
  File "[...]/ray/util/iter.py", line 702, in apply_foreach
    for item in it:
  File "[...]/ray/rllib/execution/rollout_ops.py", line 70, in sampler
    yield workers.local_worker().sample()
  File "[...]/ray/rllib/evaluation/rollout_worker.py", line 515, in sample
    batches = [self.input_reader.next()]
  File "[...]/ray/rllib/evaluation/sampler.py", line 56, in next
    batches = [self.get_data()]
  File "[...]/ray/rllib/evaluation/sampler.py", line 101, in get_data
    item = next(self.rollout_provider)
  File "[...]/ray/rllib/evaluation/sampler.py", line 367, in _env_runner
    active_episodes)
  File "[...]/ray/rllib/evaluation/sampler.py", line 637, in _do_policy_eval
    timestep=policy.global_timestep)
  File "[...]/ray/rllib/agents/qmix/qmix_policy.py", line 260, in compute_actions
    obs_batch, action_mask, _ = self._unpack_observation(obs_batch)
  File "[...]/ray/rllib/agents/qmix/qmix_policy.py", line 486, in _unpack_observation
    axis=1).reshape([len(obs_batch), self.n_agents, self.obs_size])
  File "<__array_function__ internals>", line 6, in concatenate
ValueError: zero-dimensional arrays cannot be concatenated
```
[rllib] Make QMix support complex observation spaces
<!--Please include [tune], [rllib], [autoscaler] etc. in the issue title if relevant-->

### Describe your feature request

The QMix example and the default model only uses a simple MultiDiscrete observation space. It would be nice if more abstract observation spaces and custom models would be supported. The following example using a nested dict observation space ([based on the two step game](https://github.com/ray-project/ray/blob/master/rllib/examples/env/two_step_game.py)) fails [here](https://github.com/ray-project/ray/blob/master/rllib/agents/qmix/qmix_policy.py#L482) (I am aware that the model would have to be adapted as well, but that does not seem to be the cause of this error):

```
import argparse
from gym.spaces import Tuple, MultiDiscrete, Dict, Discrete

import ray
from ray import tune
from ray.tune import register_env, grid_search
from ray.rllib.env.multi_agent_env import ENV_STATE
from ray.rllib.examples.env.two_step_game import TwoStepGame
from ray.rllib.utils.test_utils import check_learning_achieved

from gym.spaces import MultiDiscrete, Dict, Discrete
import numpy as np

from ray.rllib.env.multi_agent_env import MultiAgentEnv, ENV_STATE


class TwoStepGame(MultiAgentEnv):
    action_space = Discrete(2)

    def __init__(self, env_config):
        self.state = None
        self.agent_1 = 0
        self.agent_2 = 1
        # MADDPG emits action logits instead of actual discrete actions
        self.actions_are_logits = env_config.get("actions_are_logits", False)
        self.one_hot_state_encoding = env_config.get("one_hot_state_encoding",
                                                     False)
        self.with_state = env_config.get("separate_state_space", False)

        if not self.one_hot_state_encoding:
            self.observation_space = Discrete(6)
            self.with_state = False
        else:
            # Each agent gets the full state (one-hot encoding of which of the
            # three states are active) as input with the receiving agent's
            # ID (1 or 2) concatenated onto the end.
            if self.with_state:
                self.observation_space = Dict({
                    "obs": MultiDiscrete([2, 2, 2, 3]),
                    ENV_STATE: MultiDiscrete([2, 2, 2])
                })
            else:
                self.observation_space = MultiDiscrete([2, 2, 2, 3])

    def reset(self):
        self.state = np.array([1, 0, 0])
        return self._obs()

    def step(self, action_dict):
        if self.actions_are_logits:
            action_dict = {
                k: np.random.choice([0, 1], p=v)
                for k, v in action_dict.items()
            }

        state_index = np.flatnonzero(self.state)
        if state_index == 0:
            action = action_dict[self.agent_1]
            assert action in [0, 1], action
            if action == 0:
                self.state = np.array([0, 1, 0])
            else:
                self.state = np.array([0, 0, 1])
            global_rew = 0
            done = False
        elif state_index == 1:
            global_rew = 7
            done = True
        else:
            if action_dict[self.agent_1] == 0 and action_dict[self.
                                                              agent_2] == 0:
                global_rew = 0
            elif action_dict[self.agent_1] == 1 and action_dict[self.
                                                                agent_2] == 1:
                global_rew = 8
            else:
                global_rew = 1
            done = True

        rewards = {
            self.agent_1: global_rew / 2.0,
            self.agent_2: global_rew / 2.0
        }
        obs = self._obs()
        dones = {"__all__": done}
        infos = {}
        return obs, rewards, dones, infos

    def _obs(self):
        if self.with_state:
            return {
                self.agent_1: {
                    "obs": {
                        "data": self.agent_1_obs(),
                        "test": 1
                    },
                    ENV_STATE: self.state
                },
                self.agent_2: {
                    "obs": {
                        "data": self.agent_2_obs(),
                        "test": 2
                    },
                    ENV_STATE: self.state
                }
            }
        else:
            return {
                self.agent_1: self.agent_1_obs(),
                self.agent_2: self.agent_2_obs()
            }

    def agent_1_obs(self):
        if self.one_hot_state_encoding:
            return np.concatenate([self.state, [1]])
        else:
            return np.flatnonzero(self.state)[0]

    def agent_2_obs(self):
        if self.one_hot_state_encoding:
            return np.concatenate([self.state, [2]])
        else:
            return np.flatnonzero(self.state)[0] + 3

if __name__ == "__main__":
    grouping = {
        "group_1": [0, 1],
    }
    obs_space = Tuple([
        Dict({
            "obs": Dict({
                'data': MultiDiscrete([2, 2, 2, 3]),
                'test': Discrete(5)
            }),
            ENV_STATE: MultiDiscrete([2, 2, 2])
        }),
        Dict({
            "obs": Dict({
                'data': MultiDiscrete([2, 2, 2, 3]),
                'test': Discrete(5)
            }),
            ENV_STATE: MultiDiscrete([2, 2, 2])
        }),
    ])
    act_space = Tuple([
        TwoStepGame.action_space,
        TwoStepGame.action_space,
    ])
    register_env(
        "grouped_twostep",
        lambda config: TwoStepGame(config).with_agent_groups(
            grouping, obs_space=obs_space, act_space=act_space))

    ray.init()

    results = tune.run("QMIX", config={
        "rollout_fragment_length": 4,
        "train_batch_size": 32,
        "exploration_fraction": .4,
        "exploration_final_eps": 0.0,
        "num_workers": 0,
        "mixer": "qmix",
        "env_config": {
            "separate_state_space": True,
            "one_hot_state_encoding": True
        },
        "use_pytorch": True,
        "env": "grouped_twostep"
    })

    ray.shutdown()
```
```
2020-05-20 15:40:03,477 ERROR trial_runner.py:519 -- Trial QMIX_grouped_twostep_c960c_00002: Error processing event.             
Traceback (most recent call last):
  File "[...]/ray/tune/trial_runner.py", line 467, in _process_trial        
    result = self.trial_executor.fetch_result(trial)
  File "[...]/ray/tune/ray_trial_executor.py", line 430, in fetch_result         
    result = ray.get(trial_future[0], DEFAULT_GET_TIMEOUT)         
  File "[...]/ray/worker.py", line 1516, in get                                      
    raise value.as_instanceof_cause()                              
ray.exceptions.RayTaskError(ValueError): ray::QMIX.train() (pid=54834, ip=128.232.69.20)
  File "python/ray/_raylet.pyx", line 460, in ray._raylet.execute_task
  File "python/ray/_raylet.pyx", line 414, in ray._raylet.execute_task.function_executor
  File "[...]/ray/rllib/agents/trainer.py", line 504, in train
    raise e                     
  File "[...]/ray/rllib/agents/trainer.py", line 490, in train
    result = Trainable.train(self)          
  File "[...]/ray/tune/trainable.py", line 260, in train
    result = self._train()                                       
  File "[...]/ray/rllib/agents/trainer_template.py", line 138, in _train
    return self._train_exec_impl()                               
  File "[...]/ray/rllib/agents/trainer_template.py", line 173, in _train_exec_impl
    res = next(self.train_exec_impl)                             
  File "[...]/ray/util/iter.py", line 689, in __next__
    return next(self.built_iterator)                             
  File "[...]/ray/util/iter.py", line 702, in apply_foreach
    for item in it:                                                                                                                                        
  File "[...]/ray/util/iter.py", line 772, in apply_filter                                        
    for item in it:                                                                                                                                        
  File "[...]/ray/util/iter.py", line 772, in apply_filter                                        
    for item in it:                                                                                                                                        
  File "[...]/ray/util/iter.py", line 702, in apply_foreach                                       
    for item in it:                                                                                                                                        
  File "[...]/ray/util/iter.py", line 772, in apply_filter
    for item in it:               
  File "[...]/ray/util/iter.py", line 977, in build_union
    item = next(it)         
  File "[...]/ray/util/iter.py", line 689, in __next__
    return next(self.built_iterator)                             
  File "[...]/ray/util/iter.py", line 702, in apply_foreach                                         
    for item in it:                                                       
  File "[...]/ray/util/iter.py", line 702, in apply_foreach
    for item in it:                                                                            
  File "[...]/ray/util/iter.py", line 702, in apply_foreach
    for item in it:
  File "[...]/ray/rllib/execution/rollout_ops.py", line 70, in sampler
    yield workers.local_worker().sample()
  File "[...]/ray/rllib/evaluation/rollout_worker.py", line 515, in sample
    batches = [self.input_reader.next()]
  File "[...]/ray/rllib/evaluation/sampler.py", line 56, in next
    batches = [self.get_data()]
  File "[...]/ray/rllib/evaluation/sampler.py", line 101, in get_data
    item = next(self.rollout_provider)
  File "[...]/ray/rllib/evaluation/sampler.py", line 367, in _env_runner
    active_episodes)
  File "[...]/ray/rllib/evaluation/sampler.py", line 637, in _do_policy_eval
    timestep=policy.global_timestep)
  File "[...]/ray/rllib/agents/qmix/qmix_policy.py", line 260, in compute_actions
    obs_batch, action_mask, _ = self._unpack_observation(obs_batch)
  File "[...]/ray/rllib/agents/qmix/qmix_policy.py", line 486, in _unpack_observation
    axis=1).reshape([len(obs_batch), self.n_agents, self.obs_size])
  File "<__array_function__ internals>", line 6, in concatenate
ValueError: zero-dimensional arrays cannot be concatenated
```


Discussion:

Resolving patch:
diff --git a/rllib/agents/qmix/qmix_policy.py b/rllib/agents/qmix/qmix_policy.py
--- a/rllib/agents/qmix/qmix_policy.py
+++ b/rllib/agents/qmix/qmix_policy.py
@@ -13,11 +13,13 @@
 from ray.rllib.models.catalog import ModelCatalog
 from ray.rllib.models.modelv2 import _unpack_obs
 from ray.rllib.env.constants import GROUP_REWARDS
+from ray.rllib.utils import try_import_tree
 from ray.rllib.utils.framework import try_import_torch
 from ray.rllib.utils.annotations import override
 
 # Torch must be installed.
 torch, nn = try_import_torch(error=True)
+tree = try_import_tree()
 
 logger = logging.getLogger(__name__)
 
@@ -463,25 +465,28 @@ def _unpack_observation(self, obs_batch):
             state (np.ndarray or None): state tensor of shape [B, state_size]
                 or None if it is not in the batch
         """
+
         unpacked = _unpack_obs(
             np.array(obs_batch, dtype=np.float32),
             self.observation_space.original_space,
             tensorlib=np)
+
+        if isinstance(unpacked[0], dict):
+            unpacked_obs = [
+                np.concatenate(tree.flatten(u["obs"]), 1) for u in unpacked
+            ]
+        else:
+            unpacked_obs = unpacked
+
+        obs = np.concatenate(
+            unpacked_obs,
+            axis=1).reshape([len(obs_batch), self.n_agents, self.obs_size])
+
         if self.has_action_mask:
-            obs = np.concatenate(
-                [o["obs"] for o in unpacked],
-                axis=1).reshape([len(obs_batch), self.n_agents, self.obs_size])
             action_mask = np.concatenate(
                 [o["action_mask"] for o in unpacked], axis=1).reshape(
                     [len(obs_batch), self.n_agents, self.n_actions])
         else:
-            if isinstance(unpacked[0], dict):
-                unpacked_obs = [u["obs"] for u in unpacked]
-            else:
-                unpacked_obs = unpacked
-            obs = np.concatenate(
-                unpacked_obs,
-                axis=1).reshape([len(obs_batch), self.n_agents, self.obs_size])
             action_mask = np.ones(
                 [len(obs_batch), self.n_agents, self.n_actions],
                 dtype=np.float32)


Test patch:
diff --git a/rllib/tests/test_avail_actions_qmix.py b/rllib/tests/test_avail_actions_qmix.py
--- a/rllib/tests/test_avail_actions_qmix.py
+++ b/rllib/tests/test_avail_actions_qmix.py
@@ -1,4 +1,4 @@
-from gym.spaces import Tuple, Discrete, Dict, Box
+from gym.spaces import Box, Dict, Discrete, MultiDiscrete, Tuple
 import numpy as np
 import unittest
 
@@ -9,10 +9,17 @@
 
 
 class AvailActionsTestEnv(MultiAgentEnv):
-    action_space = Discrete(10)
+    num_actions = 10
+    action_space = Discrete(num_actions)
     observation_space = Dict({
-        "obs": Discrete(3),
-        "action_mask": Box(0, 1, (10, )),
+        "obs": Dict({
+            "test": Dict({
+                "a": Discrete(2),
+                "b": MultiDiscrete([2, 3, 4])
+            }),
+            "state": MultiDiscrete([2, 2, 2])
+        }),
+        "action_mask": Box(0, 1, (num_actions, )),
     })
 
     def __init__(self, env_config):
@@ -25,7 +32,7 @@ def reset(self):
         self.state = 0
         return {
             "agent_1": {
-                "obs": self.state,
+                "obs": self.observation_space["obs"].sample(),
                 "action_mask": self.action_mask
             }
         }
@@ -36,7 +43,12 @@ def step(self, action_dict):
                 "Failed to obey available actions mask!"
         self.state += 1
         rewards = {"agent_1": 1}
-        obs = {"agent_1": {"obs": 0, "action_mask": self.action_mask}}
+        obs = {
+            "agent_1": {
+                "obs": self.observation_space["obs"].sample(),
+                "action_mask": self.action_mask
+            }
+        }
         dones = {"__all__": self.state > 20}
         return obs, rewards, dones, {}