[rllib] example and docs on how to use parametric actions with DQN / PG algorithms

doc/source/rllib-env.rst

@@ -7,20 +7,22 @@ RLlib works with several different types of environments, including `OpenAI Gym
 
 **Compatibility matrix**:
 
-=============  ================  ==================  ===========  ==================
-Algorithm      Discrete Actions  Continuous Actions  Multi-Agent  Recurrent Policies
-=============  ================  ==================  ===========  ==================
-A2C, A3C        **Yes**           **Yes**             **Yes**      **Yes**
-PPO             **Yes**           **Yes**             **Yes**      **Yes**
-PG              **Yes**           **Yes**             **Yes**      **Yes**
-IMPALA          **Yes**           No                  **Yes**      **Yes**
-DQN, Rainbow    **Yes**           No                  **Yes**      No
-DDPG            No                **Yes**             **Yes**      No
-APEX-DQN        **Yes**           No                  **Yes**      No
-APEX-DDPG       No                **Yes**             **Yes**      No
-ES              **Yes**           **Yes**             No           No
-ARS             **Yes**           **Yes**             No           No
-=============  ================  ==================  ===========  ==================
+=============  =======================  ==================  ===========  ==================
+Algorithm      Discrete Actions         Continuous Actions  Multi-Agent  Recurrent Policies
+=============  =======================  ==================  ===========  ==================
+A2C, A3C        **Yes** `+parametric`_  **Yes**             **Yes**      **Yes**
+PPO             **Yes** `+parametric`_  **Yes**             **Yes**      **Yes**
+PG              **Yes** `+parametric`_  **Yes**             **Yes**      **Yes**
+IMPALA          **Yes** `+parametric`_  No                  **Yes**      **Yes**
+DQN, Rainbow    **Yes**                 No                  **Yes**      No
+DDPG            No                      **Yes**             **Yes**      No
+APEX-DQN        **Yes**                 No                  **Yes**      No
+APEX-DDPG       No                      **Yes**             **Yes**      No
+ES              **Yes**                 **Yes**             No           No
+ARS             **Yes**                 **Yes**             No           No
+=============  =======================  ==================  ===========  ==================
+
+.. _`+parametric`: rllib-models.html#variable-length-parametric-action-spaces
 
 In the high-level agent APIs, environments are identified with string names. By default, the string will be interpreted as a gym `environment name <https://gym.openai.com/envs>`__, however you can also register custom environments by name:
 

doc/source/rllib-models.rst

@@ -188,6 +188,53 @@ Then, you can create an agent with your custom policy graph by:
 
 In this example we overrode existing methods of the existing DDPG policy graph, i.e., `_build_q_network`, `_build_p_network`, `_build_action_network`, `_build_actor_critic_loss`, but you can also replace the entire graph class entirely.
 
+Variable-length / Parametric Action Spaces
+------------------------------------------
+
+Custom models can be used to work with environments where (1) the set of valid actions varies per step, and/or (2) the number of valid actions is `very large <https://neuro.cs.ut.ee/the-use-of-embeddings-in-openai-five/>`__ or `potentially infinite <https://arxiv.org/abs/1811.00260>`__. This applies to algorithms in the `policy-gradient family <rllib-env.html>`__ and works as follows:
+
+1. The environment should return a mask and/or list of valid action embeddings as part of the observation for each step. To enable batching, the number of actions can be allowed to vary from 1 to some max num available:
+
+.. code-block:: python
+
+   class MyParamActionEnv(
+       def __init__(self, max_avail_actions):
+           self.action_space = Discrete(max_avail_actions)
+           self.observation_space = Dict({
+               "action_mask": Box(0, 1, shape=(max_avail_actions, )),
+               "avail_actions": Box(-1, 1, shape=(max_avail_actions, action_embedding_sz)),
+               "real_obs": ...,
+           })
+
+2. A custom model can be defined that can interpret the ``action_mask`` and ``avail_actions`` portions of the observation. Here the model computes the action logits via the dot product of some network output and each action embedding. Invalid actions can be masked out of the softmax by scaling the probability to zero:
+
+.. code-block:: python
+
+    class MyParamActionModel(Model):
+        def _build_layers_v2(self, input_dict, num_outputs, options):
+            avail_actions = input_dict["obs"]["avail_actions"]
+            action_mask = input_dict["obs"]["action_mask"]
+
+            output = FullyConnectedNetwork(
+                input_dict["obs"]["real_obs"], num_outputs=action_embedding_sz)
+
+            # Expand the model output to [BATCH, 1, EMBED_SIZE]. Note that the
+            # avail actions tensor is of shape [BATCH, MAX_ACTIONS, EMBED_SIZE].
+            intent_vector = tf.expand_dims(output, 1)
+
+            # Shape of logits is [BATCH, MAX_ACTIONS].
+            action_logits = tf.reduce_sum(avail_actions * intent_vector, axis=2)
+
+            # Mask out invalid actions (use tf.float32.min for stability)
+            inf_mask = tf.maximum(tf.log(action_mask), tf.float32.min)
+            masked_logits = inf_mask + action_logits
+
+            return masked_logits, last_layer
+
+
+Depending on your use case it may make sense to use just the masking, just action embeddings, or both. For a runnable example of this in code, check out `parametric_action_cartpole.py <https://github.com/ray-project/ray/blob/master/python/ray/rllib/examples/parametric_action_cartpole.py>`__.
+
+
 Model-Based Rollouts
 --------------------
 

doc/source/rllib.rst

@@ -75,6 +75,7 @@ Models and Preprocessors
 * `Custom Models <rllib-models.html#custom-models>`__
 * `Custom Preprocessors <rllib-models.html#custom-preprocessors>`__
 * `Customizing Policy Graphs <rllib-models.html#customizing-policy-graphs>`__
+* `Variable-length / Parametric Action Spaces <rllib-models.html#variable-length-parametric-action-spaces>`__
 * `Model-Based Rollouts <rllib-models.html#model-based-rollouts>`__
 
 RLlib Concepts

python/ray/rllib/agents/ppo/ppo.py

@@ -110,6 +110,11 @@ def _validate_config(self):
                 and not self.config["simple_optimizer"]):
             logger.warn("forcing simple_optimizer=True in multi-agent mode")
             self.config["simple_optimizer"] = True
+        if self.config["observation_filter"] != "NoFilter":
+            # TODO(ekl): consider setting the default to be NoFilter
+            logger.warn(
+                "By default, observations will be normalized with {}".format(
+                    self.config["observation_filter"]))
 
     def _train(self):
         prev_steps = self.optimizer.num_steps_sampled

python/ray/rllib/examples/parametric_action_cartpole.py

@@ -0,0 +1,183 @@
+"""Example of handling variable length and/or parametric action spaces.
+
+This is a toy example of the action-embedding based approach for handling large
+discrete action spaces (potentially infinite in size), similar to how
+OpenAI Five works:
+
+    https://neuro.cs.ut.ee/the-use-of-embeddings-in-openai-five/
+
+Note: this currently only works with RLlib's policy gradient style algorithms
+i.e., PG / PPO and not DQN / DDPG.
+"""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import argparse
+import random
+import numpy as np
+import gym
+from gym.spaces import Box, Discrete, Dict
+import tensorflow as tf
+import tensorflow.contrib.slim as slim
+
+import ray
+from ray.rllib.models import Model, ModelCatalog
+from ray.rllib.models.misc import normc_initializer
+from ray.tune import run_experiments
+from ray.tune.registry import register_env
+
+parser = argparse.ArgumentParser()
+parser.add_argument("--num-iters", type=int, default=200)
+parser.add_argument("--run", type=str, default="PPO")
+
+
+class ParametricActionCartpole(gym.Env):
+    """Parametric action version of CartPole.
+
+    In this env there are only ever two valid actions, but we pretend there are
+    actually up to `max_avail_actions` actions that can be taken, and the two
+    valid actions are randomly hidden among this set.
+
+    At each step, we emit a dict of:
+        - the actual cart observation
+        - a mask of valid actions (e.g., [0, 1, 1, 0, 0, 1] for 3 / 6 active)
+        - the list of action embeddings (w/ zeroes for invalid actions) (e.g.,
+            [[0, 0],
+             [-1.4414, 0.9071],
+             [-0.2322, -0.2569],
+             [0, 0],
+             [0, 0],
+             [0.7878, 1.2297]] for 3 active of 6 max)
+
+    In a real environment, the actions embeddings would be larger than two
+    units of course, and also there would be a variable number of valid actions
+    per step instead of always [LEFT, RIGHT].
+    """
+
+    def __init__(self, max_avail_actions):
+        # Use simple random 2-unit action embeddings for [LEFT, RIGHT]
+        self.left_action_embed = np.random.randn(2)
+        self.right_action_embed = np.random.randn(2)
+        self.action_space = Discrete(max_avail_actions)
+        self.wrapped = gym.make("CartPole-v0")
+        self.observation_space = Dict({
+            "action_mask": Box(0, 1, shape=(max_avail_actions, )),
+            "avail_actions": Box(-1, 1, shape=(max_avail_actions, 2)),
+            "cart": self.wrapped.observation_space,
+        })
+
+    def update_avail_actions(self):
+        self.action_assignments = [[0, 0]] * self.action_space.n
+        self.action_mask = [0] * self.action_space.n
+        self.left_idx, self.right_idx = random.sample(
+            range(self.action_space.n), 2)
+        self.action_assignments[self.left_idx] = self.left_action_embed
+        self.action_assignments[self.right_idx] = self.right_action_embed
+        self.action_mask[self.left_idx] = 1
+        self.action_mask[self.right_idx] = 1
+
+    def reset(self):
+        self.update_avail_actions()
+        return {
+            "action_mask": self.action_mask,
+            "avail_actions": self.action_assignments,
+            "cart": self.wrapped.reset(),
+        }
+
+    def step(self, action):
+        if action == self.left_idx:
+            actual_action = 0
+        elif action == self.right_idx:
+            actual_action = 1
+        else:
+            raise ValueError(
+                "Chosen action was not one of the non-zero action embeddings",
+                action, self.action_assignments, self.action_mask,
+                self.left_idx, self.right_idx)
+        orig_obs, rew, done, info = self.wrapped.step(actual_action)
+        self.update_avail_actions()
+        obs = {
+            "action_mask": self.action_mask,
+            "avail_actions": self.action_assignments,
+            "cart": orig_obs,
+        }
+        return obs, rew, done, info
+
+
+class ParametricActionsModel(Model):
+    """Parametric action model that handles the dot product and masking.
+
+    This assumes the outputs are logits for a Categorical action dist."""
+
+    def _build_layers_v2(self, input_dict, num_outputs, options):
+        # Extract the available actions tensor from the observation.
+        avail_actions = input_dict["obs"]["avail_actions"]
+        action_mask = input_dict["obs"]["action_mask"]
+        action_embed_size = avail_actions.shape[2].value
+        if num_outputs != avail_actions.shape[1].value:
+            raise ValueError(
+                "This model assumes num outputs is equal to max avail actions",
+                num_outputs, avail_actions)
+
+        # Standard FC net component.
+        last_layer = input_dict["obs"]["cart"]
+        hiddens = [256, 256]
+        for i, size in enumerate(hiddens):
+            label = "fc{}".format(i)
+            last_layer = slim.fully_connected(
+                last_layer,
+                size,
+                weights_initializer=normc_initializer(1.0),
+                activation_fn=tf.nn.tanh,
+                scope=label)
+        output = slim.fully_connected(
+            last_layer,
+            action_embed_size,
+            weights_initializer=normc_initializer(0.01),
+            activation_fn=None,
+            scope="fc_out")
+
+        # Expand the model output to [BATCH, 1, EMBED_SIZE]. Note that the
+        # avail actions tensor is of shape [BATCH, MAX_ACTIONS, EMBED_SIZE].
+        intent_vector = tf.expand_dims(output, 1)
+
+        # Shape of logits is [BATCH, MAX_ACTIONS].
+        action_logits = tf.reduce_sum(avail_actions * intent_vector, axis=2)
+
+        # Mask out invalid actions (use tf.float32.min for stability)
+        inf_mask = tf.maximum(tf.log(action_mask), tf.float32.min)
+        masked_logits = inf_mask + action_logits
+
+        return masked_logits, last_layer
+
+
+if __name__ == "__main__":
+    args = parser.parse_args()
+    ray.init()
+
+    ModelCatalog.register_custom_model("pa_model", ParametricActionsModel)
+    register_env("pa_cartpole", lambda _: ParametricActionCartpole(10))
+    if args.run == "PPO":
+        cfg = {
+            "observation_filter": "NoFilter",  # don't filter the action list
+            "vf_share_layers": True,  # don't create duplicate value model
+        }
+    else:
+        cfg = {}
+    run_experiments({
+        "parametric_cartpole": {
+            "run": args.run,
+            "env": "pa_cartpole",
+            "stop": {
+                "training_iteration": args.num_iters
+            },
+            "config": dict({
+                "model": {
+                    "custom_model": "pa_model",
+                },
+                "num_workers": 0,
+            }, **cfg),
+        },
+    })

python/ray/rllib/models/catalog.py

@@ -215,7 +215,7 @@ def _get_model(input_dict, obs_space, num_outputs, options, state_in,
                    seq_lens):
         if options.get("custom_model"):
             model = options["custom_model"]
-            logger.info("Using custom model {}".format(model))
+            logger.debug("Using custom model {}".format(model))
             return _global_registry.get(RLLIB_MODEL, model)(
                 input_dict,
                 obs_space,

python/ray/rllib/models/preprocessors.py

@@ -2,6 +2,7 @@
 from __future__ import division
 from __future__ import print_function
 
+from collections import OrderedDict
 import cv2
 import logging
 import numpy as np
@@ -164,6 +165,8 @@ def _init_shape(self, obs_space, options):
         return (size, )
 
     def transform(self, observation):
+        if not isinstance(observation, OrderedDict):
+            observation = OrderedDict(sorted(list(observation.items())))
         assert len(observation) == len(self.preprocessors), \
             (len(observation), len(self.preprocessors))
         return np.concatenate([

test/jenkins_tests/run_multi_node_tests.sh

@@ -252,6 +252,9 @@ docker run --rm --shm-size=10G --memory=10G $DOCKER_SHA \
 docker run --rm --shm-size=10G --memory=10G $DOCKER_SHA \
     python /ray/python/ray/rllib/test/test_external_env.py
 
+docker run --rm --shm-size=10G --memory=10G $DOCKER_SHA \
+    python /ray/python/ray/rllib/examples/parametric_action_cartpole.py --num-iters=1
+
 docker run --rm --shm-size=10G --memory=10G $DOCKER_SHA \
     python /ray/python/ray/rllib/test/test_lstm.py