[RLlib] Remove TupleActions and support arbitrarily nested action spaces. (#8143)

Deprecate TupleActions and support arbitrarily nested action spaces.
Closes issue #8143.

.travis.yml

@@ -247,7 +247,7 @@ matrix:
         - . ./ci/travis/ci.sh preload
         - ./ci/keep_alive bazel test --config=ci --build_tests_only --test_tag_filters=examples_A,examples_B rllib/...
         - ./ci/keep_alive bazel test --config=ci --build_tests_only --test_tag_filters=examples_C rllib/...
-        - ./ci/keep_alive bazel test --config=ci --build_tests_only --test_tag_filters=examples_E,examples_L,examples_M,examples_P rllib/...
+        - ./ci/keep_alive bazel test --config=ci --build_tests_only --test_tag_filters=examples_E,examples_L,examples_M,examples_N,examples_P rllib/...
         - ./ci/keep_alive bazel test --config=ci --build_tests_only --test_tag_filters=examples_U,examples_R,examples_S,examples_T rllib/...
 
     # RLlib: tests_dir: Everything in rllib/tests/ directory (A-I).

rllib/BUILD

@@ -1136,10 +1136,11 @@ py_test(
 )
 
 py_test(
-    name = "tests/test_nested_spaces",
+    name = "tests/test_nested_observation_spaces",
+    main = "tests/test_nested_observation_spaces.py",
     tags = ["tests_dir", "tests_dir_N"],
     size = "small",
-    srcs = ["tests/test_nested_spaces.py"]
+    srcs = ["tests/test_nested_observation_spaces.py"]
 )
 
 py_test(
@@ -1401,6 +1402,15 @@ py_test(
     args = ["--num-iters=2"]
 )
 
+py_test(
+    name = "examples/nested_action_spaces_ppo",
+    main = "examples/nested_action_spaces.py",
+    tags = ["examples", "examples_N"],
+    size = "small",
+    srcs = ["examples/nested_action_spaces.py"],
+    args = ["--stop=-500", "--run=PPO"]
+)
+
 py_test(
     name = "examples/parametric_action_cartpole_pg", main="examples/parametric_action_cartpole.py",
     tags = ["examples", "examples_P"],

rllib/agents/ars/ars_tf_policy.py

@@ -7,7 +7,7 @@ import numpy as np
 import ray
 import ray.experimental.tf_utils
 from ray.rllib.agents.es.es_tf_policy import make_session
-from ray.rllib.evaluation.sampler import _unbatch_tuple_actions
+from ray.rllib.evaluation.sampler import unbatch_actions
 from ray.rllib.models import ModelCatalog
 from ray.rllib.policy.sample_batch import SampleBatch
 from ray.rllib.utils.filter import get_filter
@@ -56,7 +56,7 @@ class ARSTFPolicy:
         observation = self.observation_filter(observation[None], update=update)
         action = self.sess.run(
             self.sampler, feed_dict={self.inputs: observation})
-        action = _unbatch_tuple_actions(action)
+        action = unbatch_actions(action)
         if add_noise and isinstance(self.action_space, gym.spaces.Box):
             action += np.random.randn(*action.shape) * self.action_noise_std
         return action

rllib/agents/es/es_tf_policy.py

@@ -6,13 +6,16 @@ import numpy as np
 
 import ray
 import ray.experimental.tf_utils
-from ray.rllib.evaluation.sampler import _unbatch_tuple_actions
+from ray.rllib.evaluation.sampler import unbatch_actions
 from ray.rllib.models import ModelCatalog
 from ray.rllib.policy.sample_batch import SampleBatch
+from ray.rllib.utils import try_import_tree
 from ray.rllib.utils.filter import get_filter
 from ray.rllib.utils.framework import try_import_tf
+from ray.rllib.utils.space_utils import get_base_struct_from_space
 
 tf = try_import_tf()
+tree = try_import_tree()
 
 
 def rollout(policy, env, timestep_limit=None, add_noise=False, offset=0.0):
@@ -66,6 +69,7 @@ def make_session(single_threaded):
 class ESTFPolicy:
     def __init__(self, obs_space, action_space, config):
         self.action_space = action_space
+        self.action_space_struct = get_base_struct_from_space(action_space)
         self.action_noise_std = config["action_noise_std"]
         self.preprocessor = ModelCatalog.get_preprocessor_for_space(obs_space)
         self.observation_filter = get_filter(config["observation_filter"],
@@ -95,12 +99,22 @@ class ESTFPolicy:
     def compute_actions(self, observation, add_noise=False, update=True):
         observation = self.preprocessor.transform(observation)
         observation = self.observation_filter(observation[None], update=update)
-        action = self.sess.run(
+        # `actions` is a list of (component) batches.
+        actions = self.sess.run(
             self.sampler, feed_dict={self.inputs: observation})
-        action = _unbatch_tuple_actions(action)
-        if add_noise and isinstance(self.action_space, gym.spaces.Box):
-            action += np.random.randn(*action.shape) * self.action_noise_std
-        return action
+        if add_noise:
+            actions = tree.map_structure(self._add_noise, actions,
+                                         self.action_space_struct)
+        # Convert `flat_actions` to a list of lists of action components
+        # (list of single actions).
+        actions = unbatch_actions(actions)
+        return actions
+
+    def _add_noise(self, single_action, single_action_space):
+        if isinstance(single_action_space, gym.spaces.Box):
+            single_action += np.random.randn(*single_action.shape) * \
+                self.action_noise_std
+        return single_action
 
     def set_flat_weights(self, x):
         self.variables.set_flat(x)

rllib/agents/es/es_torch_policy.py

@@ -5,7 +5,7 @@ import gym
 import numpy as np
 
 import ray
-from ray.rllib.evaluation.sampler import _unbatch_tuple_actions
+from ray.rllib.evaluation.sampler import unbatch_actions
 from ray.rllib.models import ModelCatalog
 from ray.rllib.policy.sample_batch import SampleBatch
 from ray.rllib.policy.torch_policy_template import build_torch_policy
@@ -61,7 +61,7 @@ def before_init(policy, observation_space, action_space, config):
         }, [], None)
         dist = policy.dist_class(dist_inputs, policy.model)
         action = dist.sample().detach().numpy()
-        action = _unbatch_tuple_actions(action)
+        action = unbatch_actions(action)
         if add_noise and isinstance(policy.action_space, gym.spaces.Box):
             action += np.random.randn(*action.shape) * policy.action_noise_std
         return action

rllib/agents/qmix/qmix_policy.py

@@ -14,7 +14,6 @@ from ray.rllib.models.model import _unpack_obs
 from ray.rllib.env.constants import GROUP_REWARDS
 from ray.rllib.utils.framework import try_import_torch
 from ray.rllib.utils.annotations import override
-from ray.rllib.utils.tuple_actions import TupleActions
 
 # Torch must be installed.
 torch, nn = try_import_torch(error=True)
@@ -289,7 +288,7 @@ class QMixTorchPolicy(Policy):
             actions = actions.cpu().numpy()
             hiddens = [s.cpu().numpy() for s in hiddens]
 
-        return TupleActions(list(actions.transpose([1, 0]))), hiddens, {}
+        return tuple(actions.transpose([1, 0])), hiddens, {}
 
     @override(Policy)
     def compute_log_likelihoods(self,

rllib/evaluation/episode.py

@@ -1,12 +1,14 @@
 from collections import defaultdict
+import numpy as np
 import random
 
-import numpy as np
-
 from ray.rllib.env.base_env import _DUMMY_AGENT_ID
+from ray.rllib.utils import try_import_tree
 from ray.rllib.utils.annotations import DeveloperAPI
 from ray.rllib.utils.space_utils import flatten_to_single_ndarray
 
+tree = try_import_tree()
+
 
 @DeveloperAPI
 class MultiAgentEpisode:

rllib/evaluation/sampler.py

@@ -15,11 +15,13 @@ from ray.rllib.policy.tf_policy import TFPolicy
 from ray.rllib.env.base_env import BaseEnv, ASYNC_RESET_RETURN
 from ray.rllib.env.atari_wrappers import get_wrapper_by_cls, MonitorEnv
 from ray.rllib.offline import InputReader
+from ray.rllib.utils import try_import_tree
 from ray.rllib.utils.annotations import override
 from ray.rllib.utils.debug import summarize
-from ray.rllib.utils.tuple_actions import TupleActions
-from ray.rllib.utils.space_utils import flatten_to_single_ndarray
 from ray.rllib.utils.tf_run_builder import TFRunBuilder
+from ray.rllib.utils.space_utils import flatten_to_single_ndarray
+
+tree = try_import_tree()
 
 logger = logging.getLogger(__name__)
 
@@ -581,9 +583,7 @@ def _do_policy_eval(tf_sess, to_eval, policies, active_episodes):
 
             obs_batch = [t.obs for t in eval_data]
             state_batches = _to_column_format(rnn_in)
-
             # TODO(ekl): how can we make info batch available to TF code?
-            obs_batch = [t.obs for t in eval_data]
             prev_action_batch = [t.prev_action for t in eval_data]
             prev_reward_batch = [t.prev_reward for t in eval_data]
 
@@ -642,6 +642,11 @@ def _process_policy_eval_results(to_eval, eval_results, active_episodes,
         rnn_out_cols = eval_results[policy_id][1]
         pi_info_cols = eval_results[policy_id][2]
 
+        # In case actions is a list (representing the 0th dim of a batch of
+        # primitive actions), try to convert it first.
+        if isinstance(actions, list):
+            actions = np.array(actions)
+
         if len(rnn_in_cols) != len(rnn_out_cols):
             raise ValueError("Length of RNN in did not match RNN out, got: "
                              "{} vs {}".format(rnn_in_cols, rnn_out_cols))
@@ -650,17 +655,17 @@ def _process_policy_eval_results(to_eval, eval_results, active_episodes,
             pi_info_cols["state_in_{}".format(f_i)] = column
         for f_i, column in enumerate(rnn_out_cols):
             pi_info_cols["state_out_{}".format(f_i)] = column
-        # Save output rows
-        actions = _unbatch_tuple_actions(actions)
+
         policy = _get_or_raise(policies, policy_id)
+        # Clip if necessary (while action components are still batched).
+        if clip_actions:
+            actions = clip_action(actions, policy.action_space_struct)
+        # Split action-component batches into single action rows.
+        actions = unbatch_actions(actions)
         for i, action in enumerate(actions):
             env_id = eval_data[i].env_id
             agent_id = eval_data[i].agent_id
-            if clip_actions:
-                actions_to_send[env_id][agent_id] = clip_action(
-                    action, policy.action_space)
-            else:
-                actions_to_send[env_id][agent_id] = action
+            actions_to_send[env_id][agent_id] = action
             episode = active_episodes[env_id]
             episode._set_rnn_state(agent_id, [c[i] for c in rnn_out_cols])
             episode._set_last_pi_info(
@@ -694,17 +699,41 @@ def _fetch_atari_metrics(base_env):
     return atari_out
 
 
-def _unbatch_tuple_actions(action_batch):
-    # convert list of batches -> batch of lists
-    if isinstance(action_batch, TupleActions):
-        out = []
-        for j in range(len(action_batch.batches[0])):
-            out.append([
-                action_batch.batches[i][j]
-                for i in range(len(action_batch.batches))
-            ])
-        return out
-    return action_batch
+def unbatch_actions(action_batches):
+    """Converts action_batches from list of batches to batch of lists.
+
+    Input: Struct of batches:
+        {"a": [1, 2, 3], "b": ([4, 5, 6], [7.0, 8.0, 9.0])}
+    Output: Batch (list) of structs (each of these structs representing a
+        single action):
+        [
+            {"a": 1, "b": (4, 7.0)},  <- action 1
+            {"a": 2, "b": (5, 8.0)},  <- action 2
+            {"a": 3, "b": (6, 9.0)},  <- action 3
+        ]
+
+    Args:
+        action_batches (any): The list of action-component batches. Each item
+            in this list represents the batch for a single action component
+            (in case action is Tuple/Dict), meaning the list is already
+            flattened.
+            Alternatively, `action_batches` may also simply be a batch of
+            primitive actions (non Tuple/Dict).
+
+    Returns:
+        List[List[action-components]]: The list of action rows. Each item
+            in the returned list represents a single (maybe complex) action.
+    """
+    flat_action_batches = tree.flatten(action_batches)
+
+    out = []
+    for batch_pos in range(len(flat_action_batches[0])):
+        out.append(
+            tree.unflatten_as(action_batches, [
+                flat_action_batches[i][batch_pos]
+                for i in range(len(flat_action_batches))
+            ]))
+    return out
 
 
 def _to_column_format(rnn_state_rows):

rllib/examples/autoregressive_action_dist.py

@@ -21,8 +21,7 @@ from ray.rllib.models import ModelCatalog
 from ray.rllib.models.tf.tf_action_dist import Categorical, ActionDistribution
 from ray.rllib.models.tf.misc import normc_initializer
 from ray.rllib.models.tf.tf_modelv2 import TFModelV2
-from ray.rllib.utils.tuple_actions import TupleActions
-from ray.rllib.utils import try_import_tf
+from ray.rllib.utils.framework import try_import_tf
 
 tf = try_import_tf()
 
@@ -78,7 +77,7 @@ class BinaryAutoregressiveOutput(ActionDistribution):
         self._action_logp = a1_dist.logp(a1) + a2_dist.logp(a2)
 
         # return the action tuple
-        return TupleActions([a1, a2])
+        return (a1, a2)
 
     def sample(self):
         # first, sample a1
@@ -91,7 +90,7 @@ class BinaryAutoregressiveOutput(ActionDistribution):
         self._action_logp = a1_dist.logp(a1) + a2_dist.logp(a2)
 
         # return the action tuple
-        return TupleActions([a1, a2])
+        return (a1, a2)
 
     def logp(self, actions):
         a1, a2 = actions[:, 0], actions[:, 1]

rllib/examples/nested_action_spaces.py

@@ -0,0 +1,97 @@
+import argparse
+import gym
+from gym.spaces import Dict, Tuple, Box, Discrete
+import numpy as np
+import sys
+
+import ray
+from ray.tune.registry import register_env
+from ray.rllib.utils import try_import_tree
+from ray.rllib.utils.framework import try_import_tf
+from ray.rllib.utils.space_utils import flatten_space
+
+tf = try_import_tf()
+tree = try_import_tree()
+
+parser = argparse.ArgumentParser()
+parser.add_argument("--run", type=str, default="PPO")
+parser.add_argument("--stop", type=int, default=90)
+parser.add_argument("--max-trainstop", type=int, default=90)
+parser.add_argument("--num-cpus", type=int, default=0)
+
+
+class NestedSpaceRepeatAfterMeEnv(gym.Env):
+    """Env for which policy has to repeat the (possibly complex) observation.
+    """
+
+    def __init__(self, config):
+        self.observation_space = config.get(
+            "space", Tuple([Discrete(2),
+                            Dict({
+                                "a": Box(-1.0, 1.0, (2, ))
+                            })]))
+        self.action_space = self.observation_space
+        self.flattened_action_space = flatten_space(self.action_space)
+        self.episode_len = config.get("episode_len", 100)
+
+    def reset(self):
+        self.steps = 0
+        return self._next_obs()
+
+    def step(self, action):
+        self.steps += 1
+        action = tree.flatten(action)
+        reward = 0.0
+        for a, o, space in zip(action, self.current_obs_flattened,
+                               self.flattened_action_space):
+            # Box: -abs(diff).
+            if isinstance(space, gym.spaces.Box):
+                reward -= np.abs(np.sum(a - o))
+            # Discrete: +1.0 if exact match.
+            if isinstance(space, gym.spaces.Discrete):
+                reward += 1.0 if a == o else 0.0
+        done = self.steps >= self.episode_len
+        return self._next_obs(), reward, done, {}
+
+    def _next_obs(self):
+        self.current_obs = self.observation_space.sample()
+        self.current_obs_flattened = tree.flatten(self.current_obs)
+        return self.current_obs
+
+
+if __name__ == "__main__":
+    args = parser.parse_args()
+    ray.init(num_cpus=args.num_cpus or None)
+    register_env("NestedSpaceRepeatAfterMeEnv",
+                 lambda c: NestedSpaceRepeatAfterMeEnv(c))
+
+    config = {
+        "env": "NestedSpaceRepeatAfterMeEnv",
+        "env_config": {
+            "space": Dict({
+                "a": Tuple(
+                    [Dict({
+                        "d": Box(-10.0, 10.0, ()),
+                        "e": Discrete(2)
+                    })]),
+                "b": Box(-10.0, 10.0, (2, )),
+                "c": Discrete(4)
+            }),
+        },
+        "gamma": 0.0,  # No history in Env (bandit problem).
+        "num_workers": 0,
+        "num_envs_per_worker": 20,
+        "entropy_coeff": 0.00005,  # We don't want high entropy in this Env.
+        "num_sgd_iter": 20,
+        "vf_loss_coeff": 0.01,
+        "lr": 0.0003
+    }
+
+    import ray.rllib.agents.ppo as ppo
+    trainer = ppo.PPOTrainer(config=config)
+    for _ in range(100):
+        results = trainer.train()
+        print(results)
+        if results["episode_reward_mean"] > args.stop:
+            sys.exit(0)  # Learnt, exit gracefully.
+    sys.exit(1)  # Done, but did not learn, exit with error.

rllib/examples/rock_paper_scissors_multiagent.py

@@ -97,7 +97,7 @@ class AlwaysSameHeuristic(Policy):
                         info_batch=None,
                         episodes=None,
                         **kwargs):
-        return list(state_batches[0]), state_batches, {}
+        return state_batches[0], state_batches, {}
 
     def learn_on_batch(self, samples):
         pass
@@ -168,32 +168,29 @@ def run_heuristic_vs_learned(args, use_lstm=False, trainer="PG"):
         else:
             return random.choice(["always_same", "beat_last"])
 
-    tune.run(
-        trainer,
-        stop={"timesteps_total": args.stop},
-        config={
-            "env": RockPaperScissorsEnv,
-            "gamma": 0.9,
-            "num_workers": 0,
-            "num_envs_per_worker": 4,
-            "rollout_fragment_length": 10,
-            "train_batch_size": 200,
-            "multiagent": {
-                "policies_to_train": ["learned"],
-                "policies": {
-                    "always_same": (AlwaysSameHeuristic, Discrete(3),
-                                    Discrete(3), {}),
-                    "beat_last": (BeatLastHeuristic, Discrete(3), Discrete(3),
-                                  {}),
-                    "learned": (None, Discrete(3), Discrete(3), {
-                        "model": {
-                            "use_lstm": use_lstm
-                        }
-                    }),
-                },
-                "policy_mapping_fn": select_policy,
+    config = {
+        "env": RockPaperScissorsEnv,
+        "gamma": 0.9,
+        "num_workers": 0,
+        "num_envs_per_worker": 4,
+        "rollout_fragment_length": 10,
+        "train_batch_size": 200,
+        "multiagent": {
+            "policies_to_train": ["learned"],
+            "policies": {
+                "always_same": (AlwaysSameHeuristic, Discrete(3), Discrete(3),
+                                {}),
+                "beat_last": (BeatLastHeuristic, Discrete(3), Discrete(3), {}),
+                "learned": (None, Discrete(3), Discrete(3), {
+                    "model": {
+                        "use_lstm": use_lstm
+                    }
+                }),
             },
-        })
+            "policy_mapping_fn": select_policy,
+        },
+    }
+    tune.run(trainer, stop={"timesteps_total": args.stop}, config=config)
 
 
 def run_with_custom_entropy_loss(args):
@@ -217,11 +214,11 @@ def run_with_custom_entropy_loss(args):
 
 if __name__ == "__main__":
     args = parser.parse_args()
+    run_heuristic_vs_learned(args, use_lstm=False)
+    print("run_heuristic_vs_learned(w/o lstm): ok.")
     run_same_policy(args)
     print("run_same_policy: ok.")
     run_heuristic_vs_learned(args, use_lstm=True)
-    print("run_heuristic_vs_learned(w/ lstm): ok.")
-    run_heuristic_vs_learned(args, use_lstm=False)
-    print("run_heuristic_vs_learned (w/o lstm): ok.")
+    print("run_heuristic_vs_learned (w/ lstm): ok.")
     run_with_custom_entropy_loss(args)
     print("run_with_custom_entropy_loss: ok.")

rllib/examples/rollout_worker_custom_workflow.py

@@ -7,6 +7,7 @@ collection and policy optimization.
 
 import argparse
 import gym
+import numpy as np
 
 import ray
 from ray import tune
@@ -43,7 +44,8 @@ class CustomPolicy(TestPolicy):
                         episodes=None,
                         **kwargs):
         # return random actions
-        return [self.action_space.sample() for _ in obs_batch], [], {}
+        return np.array([self.action_space.sample()
+                         for _ in obs_batch]), [], {}
 
     def learn_on_batch(self, samples):
         # implement your learning code here

rllib/models/catalog.py

@@ -12,18 +12,22 @@ from ray.rllib.models.preprocessors import get_preprocessor
 from ray.rllib.models.tf.fcnet_v1 import FullyConnectedNetwork
 from ray.rllib.models.tf.lstm_v1 import LSTM
 from ray.rllib.models.tf.modelv1_compat import make_v1_wrapper
-from ray.rllib.models.tf.tf_action_dist import Categorical, MultiCategorical, \
-    Deterministic, DiagGaussian, MultiActionDistribution, Dirichlet
+from ray.rllib.models.tf.tf_action_dist import Categorical, \
+    Deterministic, DiagGaussian, Dirichlet, \
+    MultiActionDistribution, MultiCategorical
 from ray.rllib.models.tf.tf_modelv2 import TFModelV2
 from ray.rllib.models.tf.visionnet_v1 import VisionNetwork
 from ray.rllib.models.torch.torch_modelv2 import TorchModelV2
 from ray.rllib.models.torch.torch_action_dist import TorchCategorical, \
-    TorchMultiCategorical, TorchDeterministic, TorchDiagGaussian
-from ray.rllib.utils import try_import_tf
+    TorchDeterministic, TorchDiagGaussian, \
+    TorchMultiActionDistribution, TorchMultiCategorical
+from ray.rllib.utils import try_import_tf, try_import_tree
 from ray.rllib.utils.annotations import DeveloperAPI, PublicAPI
 from ray.rllib.utils.error import UnsupportedSpaceException
+from ray.rllib.utils.space_utils import flatten_space
 
 tf = try_import_tf()
+tree = try_import_tree()
 
 logger = logging.getLogger(__name__)
 
@@ -134,7 +138,8 @@ class ModelCatalog:
         # Dist_type is given directly as a class.
         elif type(dist_type) is type and \
                 issubclass(dist_type, ActionDistribution) and \
-                dist_type is not MultiActionDistribution:
+                dist_type not in (
+                MultiActionDistribution, TorchMultiActionDistribution):
             dist = dist_type
         # Box space -> DiagGaussian OR Deterministic.
         elif isinstance(action_space, gym.spaces.Box):
@@ -154,24 +159,21 @@ class ModelCatalog:
         # Discrete Space -> Categorical.
         elif isinstance(action_space, gym.spaces.Discrete):
             dist = Categorical if framework == "tf" else TorchCategorical
-        # Tuple Space -> MultiAction.
-        elif dist_type is MultiActionDistribution or \
-                isinstance(action_space, gym.spaces.Tuple):
-            if framework == "torch":
-                # TODO(sven): implement
-                raise NotImplementedError(
-                    "Tuple action spaces not supported for Pytorch.")
-            child_dist = []
-            input_lens = []
-            for action in action_space.spaces:
-                dist, action_size = ModelCatalog.get_action_dist(
-                    action, config)
-                child_dist.append(dist)
-                input_lens.append(action_size)
+        # Tuple/Dict Spaces -> MultiAction.
+        elif dist_type in (MultiActionDistribution,
+                           TorchMultiActionDistribution) or \
+                isinstance(action_space, (gym.spaces.Tuple, gym.spaces.Dict)):
+            flat_action_space = flatten_space(action_space)
+            child_dists_and_in_lens = tree.map_structure(
+                lambda s: ModelCatalog.get_action_dist(
+                    s, config, framework=framework), flat_action_space)
+            child_dists = [e[0] for e in child_dists_and_in_lens]
+            input_lens = [e[1] for e in child_dists_and_in_lens]
             return partial(
-                MultiActionDistribution,
+                (TorchMultiActionDistribution
+                 if framework == "torch" else MultiActionDistribution),
                 action_space=action_space,
-                child_distributions=child_dist,
+                child_distributions=child_dists,
                 input_lens=input_lens), sum(input_lens)
         # Simplex -> Dirichlet.
         elif isinstance(action_space, Simplex):
@@ -186,12 +188,6 @@ class ModelCatalog:
                 TorchMultiCategorical
             return partial(dist, input_lens=action_space.nvec), \
                 int(sum(action_space.nvec))
-        # Dict -> TODO(sven)
-        elif isinstance(action_space, gym.spaces.Dict):
-            # TODO(sven): implement
-            raise NotImplementedError(
-                "Dict action spaces are not supported, consider using "
-                "gym.spaces.Tuple instead")
         # Unknown type -> Error.
         else:
             raise NotImplementedError("Unsupported args: {} {}".format(
@@ -217,39 +213,38 @@ class ModelCatalog:
         elif isinstance(action_space, gym.spaces.MultiDiscrete):
             return (tf.as_dtype(action_space.dtype),
                     (None, ) + action_space.shape)
-        elif isinstance(action_space, gym.spaces.Tuple):
+        elif isinstance(action_space, (gym.spaces.Tuple, gym.spaces.Dict)):
+            flat_action_space = flatten_space(action_space)
             size = 0
             all_discrete = True
-            for i in range(len(action_space.spaces)):
-                if isinstance(action_space.spaces[i], gym.spaces.Discrete):
+            for i in range(len(flat_action_space)):
+                if isinstance(flat_action_space[i], gym.spaces.Discrete):
                     size += 1
                 else:
                     all_discrete = False
-                    size += np.product(action_space.spaces[i].shape)
+                    size += np.product(flat_action_space[i].shape)
             size = int(size)
             return (tf.int64 if all_discrete else tf.float32, (None, size))
-        elif isinstance(action_space, gym.spaces.Dict):
-            raise NotImplementedError(
-                "Dict action spaces are not supported, consider using "
-                "gym.spaces.Tuple instead")
         else:
-            raise NotImplementedError("action space {}"
-                                      " not supported".format(action_space))
+            raise NotImplementedError(
+                "Action space {} not supported".format(action_space))
 
     @staticmethod
     @DeveloperAPI
-    def get_action_placeholder(action_space, name=None):
+    def get_action_placeholder(action_space, name="action"):
         """Returns an action placeholder consistent with the action space
 
         Args:
             action_space (Space): Action space of the target gym env.
+            name (str): An optional string to name the placeholder by.
+                Default: "action".
         Returns:
             action_placeholder (Tensor): A placeholder for the actions
         """
 
         dtype, shape = ModelCatalog.get_action_shape(action_space)
 
-        return tf.placeholder(dtype, shape=shape, name=(name or "action"))
+        return tf.placeholder(dtype, shape=shape, name=name)
 
     @staticmethod
     @DeveloperAPI

rllib/models/tf/tf_action_dist.py

@@ -2,13 +2,15 @@ import numpy as np
 import functools
 
 from ray.rllib.models.action_dist import ActionDistribution
+from ray.rllib.utils import MIN_LOG_NN_OUTPUT, MAX_LOG_NN_OUTPUT, \
+    SMALL_NUMBER, try_import_tree
 from ray.rllib.utils.annotations import override, DeveloperAPI
-from ray.rllib.utils import try_import_tf, try_import_tfp, SMALL_NUMBER, \
-    MIN_LOG_NN_OUTPUT, MAX_LOG_NN_OUTPUT
-from ray.rllib.utils.tuple_actions import TupleActions
+from ray.rllib.utils.framework import try_import_tf, try_import_tfp
+from ray.rllib.utils.space_utils import get_base_struct_from_space
 
 tf = try_import_tf()
 tfp = try_import_tfp()
+tree = try_import_tree()
 
 
 @DeveloperAPI
@@ -350,70 +352,82 @@ class Deterministic(TFActionDistribution):
 
 
 class MultiActionDistribution(TFActionDistribution):
-    """Action distribution that operates for list of actions.
+    """Action distribution that operates on a set of actions.
 
     Args:
         inputs (Tensor list): A list of tensors from which to compute samples.
     """
 
-    def __init__(self, inputs, model, action_space, child_distributions,
-                 input_lens):
-        # skip TFActionDistribution init
+    def __init__(self, inputs, model, *, child_distributions, input_lens,
+                 action_space):
         ActionDistribution.__init__(self, inputs, model)
-        self.input_lens = input_lens
-        split_inputs = tf.split(inputs, self.input_lens, axis=1)
-        child_list = []
-        for i, distribution in enumerate(child_distributions):
-            child_list.append(distribution(split_inputs[i], model))
-        self.child_distributions = child_list
+
+        self.action_space_struct = get_base_struct_from_space(action_space)
+
+        input_lens = np.array(input_lens, dtype=np.int32)
+        split_inputs = tf.split(inputs, input_lens, axis=1)
+        self.flat_child_distributions = tree.map_structure(
+            lambda dist, input_: dist(input_, model), child_distributions,
+            split_inputs)
 
     @override(ActionDistribution)
     def logp(self, x):
-        split_indices = []
-        for dist in self.child_distributions:
+        # Single tensor input (all merged).
+        if isinstance(x, (tf.Tensor, np.ndarray)):
+            split_indices = []
+            for dist in self.flat_child_distributions:
+                if isinstance(dist, Categorical):
+                    split_indices.append(1)
+                else:
+                    split_indices.append(tf.shape(dist.sample())[1])
+            split_x = tf.split(x, split_indices, axis=1)
+        # Structured or flattened (by single action component) input.
+        else:
+            split_x = tree.flatten(x)
+
+        def map_(val, dist):
+            # Remove extra categorical dimension.
             if isinstance(dist, Categorical):
-                split_indices.append(1)
-            else:
-                split_indices.append(tf.shape(dist.sample())[1])
-        split_list = tf.split(x, split_indices, axis=1)
-        for i, distribution in enumerate(self.child_distributions):
-            # Remove extra categorical dimension
-            if isinstance(distribution, Categorical):
-                split_list[i] = tf.cast(
-                    tf.squeeze(split_list[i], axis=-1), tf.int32)
-        log_list = [
-            distribution.logp(split_x) for distribution, split_x in zip(
-                self.child_distributions, split_list)
-        ]
-        return functools.reduce(lambda a, b: a + b, log_list)
+                val = tf.cast(tf.squeeze(val, axis=-1), tf.int32)
+            return dist.logp(val)
+
+        # Remove extra categorical dimension and take the logp of each
+        # component.
+        flat_logps = tree.map_structure(map_, split_x,
+                                        self.flat_child_distributions)
+
+        return functools.reduce(lambda a, b: a + b, flat_logps)
 
     @override(ActionDistribution)
     def kl(self, other):
         kl_list = [
-            distribution.kl(other_distribution)
-            for distribution, other_distribution in zip(
-                self.child_distributions, other.child_distributions)
+            d.kl(o) for d, o in zip(self.flat_child_distributions,
+                                    other.flat_child_distributions)
         ]
         return functools.reduce(lambda a, b: a + b, kl_list)
 
     @override(ActionDistribution)
     def entropy(self):
-        entropy_list = [s.entropy() for s in self.child_distributions]
+        entropy_list = [d.entropy() for d in self.flat_child_distributions]
         return functools.reduce(lambda a, b: a + b, entropy_list)
 
     @override(ActionDistribution)
     def sample(self):
-        return TupleActions([s.sample() for s in self.child_distributions])
+        child_distributions = tree.unflatten_as(self.action_space_struct,
+                                                self.flat_child_distributions)
+        return tree.map_structure(lambda s: s.sample(), child_distributions)
 
     @override(ActionDistribution)
     def deterministic_sample(self):
-        return TupleActions(
-            [s.deterministic_sample() for s in self.child_distributions])
+        child_distributions = tree.unflatten_as(self.action_space_struct,
+                                                self.flat_child_distributions)
+        return tree.map_structure(lambda s: s.deterministic_sample(),
+                                  child_distributions)
 
     @override(TFActionDistribution)
     def sampled_action_logp(self):
-        p = self.child_distributions[0].sampled_action_logp()
-        for c in self.child_distributions[1:]:
+        p = self.flat_child_distributions[0].sampled_action_logp()
+        for c in self.flat_child_distributions[1:]:
             p += c.sampled_action_logp()
         return p
 

rllib/policy/eager_tf_policy.py

@@ -580,8 +580,6 @@ def build_eager_tf_policy(name,
 
         def _initialize_loss_with_dummy_batch(self):
             # Dummy forward pass to initialize any policy attributes, etc.
-            action_dtype, action_shape = ModelCatalog.get_action_shape(
-                self.action_space)
             dummy_batch = {
                 SampleBatch.CUR_OBS: np.array(
                     [self.observation_space.sample()]),

rllib/policy/policy.py

@@ -2,9 +2,13 @@ from abc import ABCMeta, abstractmethod
 import gym
 import numpy as np
 
+from ray.rllib.utils import try_import_tree
 from ray.rllib.utils.annotations import DeveloperAPI
 from ray.rllib.utils.exploration.exploration import Exploration
 from ray.rllib.utils.from_config import from_config
+from ray.rllib.utils.space_utils import get_base_struct_from_space
+
+tree = try_import_tree()
 
 # By convention, metrics from optimizing the loss can be reported in the
 # `grad_info` dict returned by learn_on_batch() / compute_grads() via this key.
@@ -47,6 +51,7 @@ class Policy(metaclass=ABCMeta):
         """
         self.observation_space = observation_space
         self.action_space = action_space
+        self.action_space_struct = get_base_struct_from_space(action_space)
         self.config = config
         # The global timestep, broadcast down from time to time from the
         # driver.
@@ -157,7 +162,7 @@ class Policy(metaclass=ABCMeta):
             timestep=timestep)
 
         if clip_actions:
-            action = clip_action(action, self.action_space)
+            action = clip_action(action, self.action_space_struct)
 
         # Return action, internal state(s), infos.
         return action, [s[0] for s in state_out], \
@@ -385,27 +390,22 @@ class Policy(metaclass=ABCMeta):
         return exploration
 
 
-def clip_action(action, space):
-    """
-    Called to clip actions to the specified range of this policy.
+def clip_action(action, action_space):
+    """Clips all actions in `flat_actions` according to the given Spaces.
 
-    Arguments:
-        action: Single action.
-        space: Action space the actions should be present in.
+    Args:
+        flat_actions (List[np.ndarray]): The (flattened) list of single action
+            components. List will have len=1 for "primitive" action Spaces.
+        flat_space (List[Space]): The (flattened) list of single action Space
+            objects. Has to be of same length as `flat_actions`.
 
     Returns:
-        Clipped batch of actions.
+        List[np.ndarray]: Flattened list of single clipped "primitive" actions.
     """
 
-    if isinstance(space, gym.spaces.Box):
-        return np.clip(action, space.low, space.high)
-    elif isinstance(space, gym.spaces.Tuple):
-        if type(action) not in (tuple, list):
-            raise ValueError("Expected tuple space for actions {}: {}".format(
-                action, space))
-        out = []
-        for a, s in zip(action, space.spaces):
-            out.append(clip_action(a, s))
-        return out
-    else:
-        return action
+    def map_(a, s):
+        if isinstance(s, gym.spaces.Box):
+            a = np.clip(a, s.low, s.high)
+        return a
+
+    return tree.map_structure(map_, action, action_space)

rllib/policy/tests/test_policy.py

@@ -1,3 +1,4 @@
+import numpy as np
 import random
 
 from ray.rllib.policy.policy import Policy
@@ -22,7 +23,7 @@ class TestPolicy(Policy):
                         explore=None,
                         timestep=None,
                         **kwargs):
-        return [random.choice([0, 1])] * len(obs_batch), [], {}
+        return np.array([random.choice([0, 1])] * len(obs_batch)), [], {}
 
     @override(Policy)
     def compute_log_likelihoods(self,
@@ -31,4 +32,4 @@ class TestPolicy(Policy):
                                 state_batches=None,
                                 prev_action_batch=None,
                                 prev_reward_batch=None):
-        return [random.random()] * len(obs_batch)
+        return np.array([random.random()] * len(obs_batch))

rllib/tests/test_multi_agent_pendulum.py

@@ -36,7 +36,7 @@ class TestMultiAgentPendulum(unittest.TestCase):
                     "sgd_minibatch_size": 64,
                     "num_sgd_iter": 10,
                     "model": {
-                        "fcnet_hiddens": [64, 64],
+                        "fcnet_hiddens": [128, 128],
                     },
                     "batch_mode": "complete_episodes",
                 },

rllib/tests/test_rollout_worker.py

@@ -29,7 +29,7 @@ class MockPolicy(TestPolicy):
                         explore=None,
                         timestep=None,
                         **kwargs):
-        return [random.choice([0, 1])] * len(obs_batch), [], {}
+        return np.array([random.choice([0, 1])] * len(obs_batch)), [], {}
 
     def postprocess_trajectory(self,
                                batch,

rllib/tests/test_supported_spaces.py

@@ -1,5 +1,5 @@
 import gym
-from gym.spaces import Box, Discrete, Tuple, Dict, MultiDiscrete
+from gym.spaces import Box, Dict, Discrete, Tuple, MultiDiscrete
 from gym.envs.registration import EnvSpec
 import numpy as np
 import unittest
@@ -31,6 +31,13 @@ ACTION_SPACES_TO_TEST = {
         [Discrete(2),
          Discrete(3),
          Box(-1.0, 1.0, (5, ), dtype=np.float32)]),
+    "dict": Dict({
+        "action_choice": Discrete(3),
+        "parameters": Box(-1.0, 1.0, (1, ), dtype=np.float32),
+        "yet_another_nested_dict": Dict({
+            "a": Tuple([Discrete(2), Discrete(3)])
+        })
+    }),
 }
 
 OBSERVATION_SPACES_TO_TEST = {
@@ -89,9 +96,6 @@ def check_support(alg, config, stats, check_bounds=False, name=None):
             try:
                 if a_name in covered_a and o_name in covered_o:
                     stat = "skip"  # speed up tests by avoiding full grid
-                # TODO(sven): Add necessary torch distributions.
-                elif torch and a_name in ["tuple", "multidiscrete"]:
-                    stat = "unsupported"
                 else:
                     a = get_agent_class(alg)(config=config, env="stub_env")
                     if alg not in ["DDPG", "ES", "ARS", "SAC"]:

rllib/utils/exploration/random.py

@@ -1,14 +1,14 @@
 from gym.spaces import Discrete, MultiDiscrete, Tuple
 import numpy as np
+import tree
 from typing import Union
 
 from ray.rllib.models.action_dist import ActionDistribution
 from ray.rllib.utils.annotations import override
 from ray.rllib.utils.exploration.exploration import Exploration
+from ray.rllib.utils import force_tuple
 from ray.rllib.utils.framework import try_import_tf, try_import_torch, \
     TensorType
-from ray.rllib.utils.tuple_actions import TupleActions
-from ray.rllib.utils import force_tuple
 
 tf = try_import_tf()
 torch, _ = try_import_torch()
@@ -75,10 +75,7 @@ class Random(Exploration):
             false_fn=false_fn)
 
         # TODO(sven): Move into (deterministic_)sample(logp=True|False)
-        if isinstance(action, TupleActions):
-            batch_size = tf.shape(action[0][0])[0]
-        else:
-            batch_size = tf.shape(action)[0]
+        batch_size = tf.shape(tree.flatten(action)[0])[0]
         logp = tf.zeros(shape=(batch_size, ), dtype=tf.float32)
         return action, logp
 

rllib/utils/exploration/stochastic_sampling.py

@@ -2,14 +2,15 @@ from typing import Union
 
 from ray.rllib.models.action_dist import ActionDistribution
 from ray.rllib.models.modelv2 import ModelV2
+from ray.rllib.utils import try_import_tree
 from ray.rllib.utils.annotations import override
 from ray.rllib.utils.exploration.exploration import Exploration
 from ray.rllib.utils.framework import try_import_tf, try_import_torch, \
     TensorType
-from ray.rllib.utils.tuple_actions import TupleActions
 
 tf = try_import_tf()
 torch, _ = try_import_torch()
+tree = try_import_tree()
 
 
 class StochasticSampling(Exploration):
@@ -55,11 +56,7 @@ class StochasticSampling(Exploration):
             false_fn=lambda: deterministic_sample)
 
         def logp_false_fn():
-            # TODO(sven): Move into (deterministic_)sample(logp=True|False)
-            if isinstance(sample, TupleActions):
-                batch_size = tf.shape(action[0])[0]
-            else:
-                batch_size = tf.shape(action)[0]
+            batch_size = tf.shape(tree.flatten(action)[0])[0]
             return tf.zeros(shape=(batch_size, ), dtype=tf.float32)
 
         logp = tf.cond(
@@ -67,8 +64,7 @@ class StochasticSampling(Exploration):
             true_fn=lambda: action_dist.sampled_action_logp(),
             false_fn=logp_false_fn)
 
-        return TupleActions(action) if isinstance(sample, TupleActions) \
-            else action, logp
+        return action, logp
 
     @staticmethod
     def _get_torch_exploration_action(action_dist, explore):

rllib/utils/space_utils.py

@@ -85,10 +85,10 @@ def flatten_to_single_ndarray(input_):
         >>> #     1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0, 9.0
         >>> # ])
     """
-    # Concatenate tuple actions
-    if isinstance(input_, (list, tuple)):
+    # Concatenate complex inputs.
+    if isinstance(input_, (list, tuple, dict)):
         expanded = []
-        for in_ in input_:
+        for in_ in tree.flatten(input_):
             expanded.append(np.reshape(in_, [-1]))
         input_ = np.concatenate(expanded, axis=0).flatten()
     return input_

rllib/utils/torch_ops.py

@@ -1,17 +1,9 @@
 import numpy as np
-import logging
 
-from ray.rllib.utils.framework import try_import_torch
+from ray.rllib.utils import try_import_torch, try_import_tree
 
 torch, _ = try_import_torch()
-
-logger = logging.getLogger(__name__)
-
-try:
-    import tree
-except (ImportError, ModuleNotFoundError) as e:
-    logger.warning("`dm-tree` is not installed! Run `pip install dm-tree`.")
-    raise e
+tree = try_import_tree()
 
 
 def global_norm(tensors):

rllib/utils/tuple_actions.py

@@ -1,10 +1,18 @@
 from collections import namedtuple
+from ray.rllib.utils.deprecation import deprecation_warning
 
 
+# NOTE: This is a deprecated class. Use native python tuples
+# or dicts (both arbitrarily nested) for multi-actions from here on.
 class TupleActions(namedtuple("TupleActions", ["batches"])):
     """Used to return tuple actions as a list of batches per tuple element."""
 
     def __new__(cls, batches):
+        # Throw an informative error if used.
+        deprecation_warning(
+            old="TupleActions",
+            new="`native python tuples (arbitrarily nested)`",
+            error=True)
         return super(TupleActions, cls).__new__(cls, batches)
 
     def numpy(self):