[RLlib] Examples folder restructuring (models) part 1 (#8353)

rllib/agents/dqn/__init__.py

@@ -1,5 +1,7 @@
 from ray.rllib.agents.dqn.apex import ApexTrainer
 from ray.rllib.agents.dqn.dqn import DQNTrainer, DEFAULT_CONFIG
+from ray.rllib.agents.dqn.dqn_tf_policy import DQNTFPolicy
+from ray.rllib.agents.dqn.dqn_torch_policy import DQNTorchPolicy
 from ray.rllib.agents.dqn.simple_q import SimpleQTrainer, \
     DEFAULT_CONFIG as SIMPLE_Q_DEFAULT_CONFIG
 from ray.rllib.agents.dqn.simple_q_tf_policy import SimpleQTFPolicy
@@ -7,6 +9,8 @@ from ray.rllib.agents.dqn.simple_q_torch_policy import SimpleQTorchPolicy
 
 __all__ = [
     "ApexTrainer",
+    "DQNTFPolicy",
+    "DQNTorchPolicy",
     "DQNTrainer",
     "DEFAULT_CONFIG",
     "SIMPLE_Q_DEFAULT_CONFIG",

rllib/agents/dqn/dqn_torch_model.py

@@ -87,9 +87,9 @@ class DQNTorchModel(TorchModelV2):
             advantage_module.add_module("A", nn.Linear(ins, action_space.n))
             value_module.add_module("V", nn.Linear(ins, 1))
         # Non-dueling:
-        # Q-value layer (use Advantage module's outputs as Q-values).
+        # Q-value layer (use main module's outputs as Q-values).
         else:
-            advantage_module.add_module("Q", nn.Linear(ins, action_space.n))
+            pass
 
         self.advantage_module = advantage_module
         self.value_module = value_module

rllib/agents/dqn/dqn_torch_policy.py

@@ -19,7 +19,7 @@ from ray.rllib.utils import try_import_torch
 torch, nn = try_import_torch()
 F = None
 if nn:
-    F = torch.nn.functional
+    F = nn.functional
 
 
 class QLoss:

rllib/agents/pg/__init__.py

@@ -1,9 +1,14 @@
 from ray.rllib.agents.pg.pg import PGTrainer, DEFAULT_CONFIG
 from ray.rllib.agents.pg.pg_tf_policy import pg_tf_loss, \
-    post_process_advantages
-from ray.rllib.agents.pg.pg_torch_policy import pg_torch_loss
+    post_process_advantages, PGTFPolicy
+from ray.rllib.agents.pg.pg_torch_policy import pg_torch_loss, PGTorchPolicy
 
 __all__ = [
-    "PGTrainer", "pg_tf_loss", "pg_torch_loss", "post_process_advantages",
-    "DEFAULT_CONFIG"
+    "pg_tf_loss",
+    "pg_torch_loss",
+    "post_process_advantages",
+    "DEFAULT_CONFIG",
+    "PGTFPolicy",
+    "PGTorchPolicy",
+    "PGTrainer",
 ]

rllib/agents/ppo/__init__.py

@@ -1,5 +1,14 @@
 from ray.rllib.agents.ppo.ppo import PPOTrainer, DEFAULT_CONFIG
+from ray.rllib.agents.ppo.ppo_tf_policy import PPOTFPolicy
+from ray.rllib.agents.ppo.ppo_torch_policy import PPOTorchPolicy
 from ray.rllib.agents.ppo.appo import APPOTrainer
 from ray.rllib.agents.ppo.ddppo import DDPPOTrainer
 
-__all__ = ["APPOTrainer", "DDPPOTrainer", "PPOTrainer", "DEFAULT_CONFIG"]
+__all__ = [
+    "APPOTrainer",
+    "DDPPOTrainer",
+    "DEFAULT_CONFIG",
+    "PPOTFPolicy",
+    "PPOTorchPolicy",
+    "PPOTrainer",
+]

rllib/agents/qmix/model.py

@@ -23,7 +23,7 @@ class RNNModel(TorchModelV2, nn.Module):
 
     @override(ModelV2)
     def get_initial_state(self):
-        # make hidden states on same device as model
+        # Place hidden states on same device as model.
         return [self.fc1.weight.new(1, self.rnn_hidden_dim).zero_().squeeze(0)]
 
     @override(ModelV2)

rllib/examples/env/correlated_actions_env.py

@@ -11,21 +11,23 @@ class CorrelatedActionsEnv(gym.Env):
     def __init__(self, _):
         self.observation_space = Discrete(2)
         self.action_space = Tuple([Discrete(2), Discrete(2)])
+        self.last_observation = None
 
     def reset(self):
         self.t = 0
-        self.last = random.choice([0, 1])
-        return self.last
+        self.last_observation = random.choice([0, 1])
+        return self.last_observation
 
     def step(self, action):
         self.t += 1
         a1, a2 = action
         reward = 0
-        if a1 == self.last:
+        # Encourage correlation between most recent observation and a1.
+        if a1 == self.last_observation:
             reward += 5
-        # encourage correlation between a1 and a2
+        # Encourage correlation between a1 and a2.
         if a1 == a2:
             reward += 5
         done = self.t > 20
-        self.last = random.choice([0, 1])
-        return self.last, reward, done, {}
+        self.last_observation = random.choice([0, 1])
+        return self.last_observation, reward, done, {}

rllib/examples/env/rock_paper_scissors.py

@@ -23,6 +23,9 @@ class RockPaperScissors(MultiAgentEnv):
         self.last_move = None
         self.num_moves = 0
 
+        # For test-case inspections (compare both players' scores).
+        self.player1_score = self.player2_score = 0
+
     def reset(self):
         self.last_move = (0, 0)
         self.num_moves = 0
@@ -79,4 +82,10 @@ class RockPaperScissors(MultiAgentEnv):
         done = {
             "__all__": self.num_moves >= 10,
         }
+
+        if rew["player1"] > rew["player2"]:
+            self.player1_score += 1
+        elif rew["player2"] > rew["player1"]:
+            self.player2_score += 1
+
         return obs, rew, done, {}

rllib/examples/models/autoregressive_action_dist.py

@@ -0,0 +1,147 @@
+from ray.rllib.models.tf.tf_action_dist import Categorical, ActionDistribution
+from ray.rllib.models.torch.torch_action_dist import TorchCategorical, \
+    TorchDistributionWrapper
+from ray.rllib.utils.framework import try_import_tf, try_import_torch
+
+tf = try_import_tf()
+torch, nn = try_import_torch()
+
+
+class BinaryAutoregressiveDistribution(ActionDistribution):
+    """Action distribution P(a1, a2) = P(a1) * P(a2 | a1)"""
+
+    def deterministic_sample(self):
+        # first, sample a1
+        a1_dist = self._a1_distribution()
+        a1 = a1_dist.deterministic_sample()
+
+        # sample a2 conditioned on a1
+        a2_dist = self._a2_distribution(a1)
+        a2 = a2_dist.deterministic_sample()
+        self._action_logp = a1_dist.logp(a1) + a2_dist.logp(a2)
+
+        # return the action tuple
+        return (a1, a2)
+
+    def sample(self):
+        # first, sample a1
+        a1_dist = self._a1_distribution()
+        a1 = a1_dist.sample()
+
+        # sample a2 conditioned on a1
+        a2_dist = self._a2_distribution(a1)
+        a2 = a2_dist.sample()
+        self._action_logp = a1_dist.logp(a1) + a2_dist.logp(a2)
+
+        # return the action tuple
+        return (a1, a2)
+
+    def logp(self, actions):
+        a1, a2 = actions[:, 0], actions[:, 1]
+        a1_vec = tf.expand_dims(tf.cast(a1, tf.float32), 1)
+        a1_logits, a2_logits = self.model.action_model([self.inputs, a1_vec])
+        return (
+            Categorical(a1_logits).logp(a1) + Categorical(a2_logits).logp(a2))
+
+    def sampled_action_logp(self):
+        return tf.exp(self._action_logp)
+
+    def entropy(self):
+        a1_dist = self._a1_distribution()
+        a2_dist = self._a2_distribution(a1_dist.sample())
+        return a1_dist.entropy() + a2_dist.entropy()
+
+    def kl(self, other):
+        a1_dist = self._a1_distribution()
+        a1_terms = a1_dist.kl(other._a1_distribution())
+
+        a1 = a1_dist.sample()
+        a2_terms = self._a2_distribution(a1).kl(other._a2_distribution(a1))
+        return a1_terms + a2_terms
+
+    def _a1_distribution(self):
+        BATCH = tf.shape(self.inputs)[0]
+        a1_logits, _ = self.model.action_model(
+            [self.inputs, tf.zeros((BATCH, 1))])
+        a1_dist = Categorical(a1_logits)
+        return a1_dist
+
+    def _a2_distribution(self, a1):
+        a1_vec = tf.expand_dims(tf.cast(a1, tf.float32), 1)
+        _, a2_logits = self.model.action_model([self.inputs, a1_vec])
+        a2_dist = Categorical(a2_logits)
+        return a2_dist
+
+    @staticmethod
+    def required_model_output_shape(self, model_config):
+        return 16  # controls model output feature vector size
+
+
+class TorchBinaryAutoregressiveDistribution(TorchDistributionWrapper):
+    """Action distribution P(a1, a2) = P(a1) * P(a2 | a1)"""
+
+    def deterministic_sample(self):
+        # first, sample a1
+        a1_dist = self._a1_distribution()
+        a1 = a1_dist.deterministic_sample()
+
+        # sample a2 conditioned on a1
+        a2_dist = self._a2_distribution(a1)
+        a2 = a2_dist.deterministic_sample()
+        self._action_logp = a1_dist.logp(a1) + a2_dist.logp(a2)
+
+        # return the action tuple
+        return (a1, a2)
+
+    def sample(self):
+        # first, sample a1
+        a1_dist = self._a1_distribution()
+        a1 = a1_dist.sample()
+
+        # sample a2 conditioned on a1
+        a2_dist = self._a2_distribution(a1)
+        a2 = a2_dist.sample()
+        self._action_logp = a1_dist.logp(a1) + a2_dist.logp(a2)
+
+        # return the action tuple
+        return (a1, a2)
+
+    def logp(self, actions):
+        a1, a2 = actions[:, 0], actions[:, 1]
+        a1_vec = torch.unsqueeze(a1.float(), 1)
+        a1_logits, a2_logits = self.model.action_module(self.inputs, a1_vec)
+        return (TorchCategorical(a1_logits).logp(a1) +
+                TorchCategorical(a2_logits).logp(a2))
+
+    def sampled_action_logp(self):
+        return torch.exp(self._action_logp)
+
+    def entropy(self):
+        a1_dist = self._a1_distribution()
+        a2_dist = self._a2_distribution(a1_dist.sample())
+        return a1_dist.entropy() + a2_dist.entropy()
+
+    def kl(self, other):
+        a1_dist = self._a1_distribution()
+        a1_terms = a1_dist.kl(other._a1_distribution())
+
+        a1 = a1_dist.sample()
+        a2_terms = self._a2_distribution(a1).kl(other._a2_distribution(a1))
+        return a1_terms + a2_terms
+
+    def _a1_distribution(self):
+        BATCH = self.inputs.shape[0]
+        a1_logits, _ = self.model.action_module(self.inputs,
+                                                torch.zeros((BATCH, 1)))
+        a1_dist = TorchCategorical(a1_logits)
+        return a1_dist
+
+    def _a2_distribution(self, a1):
+        a1_vec = torch.unsqueeze(a1.float(), 1)
+        _, a2_logits = self.model.action_module(self.inputs, a1_vec)
+        a2_dist = TorchCategorical(a2_logits)
+        return a2_dist
+
+    @staticmethod
+    def required_model_output_shape(self, model_config):
+        return 16  # controls model output feature vector size

rllib/examples/models/autoregressive_action_model.py

@@ -0,0 +1,157 @@
+from gym.spaces import Discrete, Tuple
+
+from ray.rllib.models.tf.misc import normc_initializer
+from ray.rllib.models.tf.tf_modelv2 import TFModelV2
+from ray.rllib.models.torch.misc import normc_initializer as normc_init_torch
+from ray.rllib.models.torch.misc import SlimFC
+from ray.rllib.models.torch.torch_modelv2 import TorchModelV2
+from ray.rllib.utils.framework import try_import_tf, try_import_torch
+
+tf = try_import_tf()
+torch, nn = try_import_torch()
+
+
+class AutoregressiveActionModel(TFModelV2):
+    """Implements the `.action_model` branch required above."""
+
+    def __init__(self, obs_space, action_space, num_outputs, model_config,
+                 name):
+        super(AutoregressiveActionModel, self).__init__(
+            obs_space, action_space, num_outputs, model_config, name)
+        if action_space != Tuple([Discrete(2), Discrete(2)]):
+            raise ValueError(
+                "This model only supports the [2, 2] action space")
+
+        # Inputs
+        obs_input = tf.keras.layers.Input(
+            shape=obs_space.shape, name="obs_input")
+        a1_input = tf.keras.layers.Input(shape=(1, ), name="a1_input")
+        ctx_input = tf.keras.layers.Input(
+            shape=(num_outputs, ), name="ctx_input")
+
+        # Output of the model (normally 'logits', but for an autoregressive
+        # dist this is more like a context/feature layer encoding the obs)
+        context = tf.keras.layers.Dense(
+            num_outputs,
+            name="hidden",
+            activation=tf.nn.tanh,
+            kernel_initializer=normc_initializer(1.0))(obs_input)
+
+        # V(s)
+        value_out = tf.keras.layers.Dense(
+            1,
+            name="value_out",
+            activation=None,
+            kernel_initializer=normc_initializer(0.01))(context)
+
+        # P(a1 | obs)
+        a1_logits = tf.keras.layers.Dense(
+            2,
+            name="a1_logits",
+            activation=None,
+            kernel_initializer=normc_initializer(0.01))(ctx_input)
+
+        # P(a2 | a1)
+        # --note: typically you'd want to implement P(a2 | a1, obs) as follows:
+        # a2_context = tf.keras.layers.Concatenate(axis=1)(
+        #     [ctx_input, a1_input])
+        a2_context = a1_input
+        a2_hidden = tf.keras.layers.Dense(
+            16,
+            name="a2_hidden",
+            activation=tf.nn.tanh,
+            kernel_initializer=normc_initializer(1.0))(a2_context)
+        a2_logits = tf.keras.layers.Dense(
+            2,
+            name="a2_logits",
+            activation=None,
+            kernel_initializer=normc_initializer(0.01))(a2_hidden)
+
+        # Base layers
+        self.base_model = tf.keras.Model(obs_input, [context, value_out])
+        self.register_variables(self.base_model.variables)
+        self.base_model.summary()
+
+        # Autoregressive action sampler
+        self.action_model = tf.keras.Model([ctx_input, a1_input],
+                                           [a1_logits, a2_logits])
+        self.action_model.summary()
+        self.register_variables(self.action_model.variables)
+
+    def forward(self, input_dict, state, seq_lens):
+        context, self._value_out = self.base_model(input_dict["obs"])
+        return context, state
+
+    def value_function(self):
+        return tf.reshape(self._value_out, [-1])
+
+
+class TorchAutoregressiveActionModel(TorchModelV2, nn.Module):
+    """PyTorch version of the AutoregressiveActionModel above."""
+
+    def __init__(self, obs_space, action_space, num_outputs, model_config,
+                 name):
+        TorchModelV2.__init__(self, obs_space, action_space, num_outputs,
+                              model_config, name)
+        nn.Module.__init__(self)
+
+        if action_space != Tuple([Discrete(2), Discrete(2)]):
+            raise ValueError(
+                "This model only supports the [2, 2] action space")
+
+        # Output of the model (normally 'logits', but for an autoregressive
+        # dist this is more like a context/feature layer encoding the obs)
+        self.context_layer = SlimFC(
+            in_size=obs_space.shape[0],
+            out_size=num_outputs,
+            initializer=normc_init_torch(1.0),
+            activation_fn=nn.Tanh,
+        )
+
+        # V(s)
+        self.value_branch = SlimFC(
+            in_size=num_outputs,
+            out_size=1,
+            initializer=normc_init_torch(0.01),
+            activation_fn=None,
+        )
+
+        # P(a1 | obs)
+        self.a1_logits = SlimFC(
+            in_size=num_outputs,
+            out_size=2,
+            activation_fn=None,
+            initializer=normc_init_torch(0.01))
+
+        class _ActionModel(nn.Module):
+            def __init__(self):
+                nn.Module.__init__(self)
+                self.a2_hidden = SlimFC(
+                    in_size=1,
+                    out_size=16,
+                    activation_fn=nn.Tanh,
+                    initializer=normc_init_torch(1.0))
+                self.a2_logits = SlimFC(
+                    in_size=16,
+                    out_size=2,
+                    activation_fn=None,
+                    initializer=normc_init_torch(0.01))
+
+            def forward(self_, ctx_input, a1_input):
+                a1_logits = self.a1_logits(ctx_input)
+                a2_logits = self_.a2_logits(self_.a2_hidden(a1_input))
+                return a1_logits, a2_logits
+
+        # P(a2 | a1)
+        # --note: typically you'd want to implement P(a2 | a1, obs) as follows:
+        # a2_context = tf.keras.layers.Concatenate(axis=1)(
+        #     [ctx_input, a1_input])
+        self.action_module = _ActionModel()
+
+    def forward(self, input_dict, state, seq_lens):
+        context = self.context_layer(input_dict["obs"])
+        self._value_out = self.value_branch(context)
+        return context, state
+
+    def value_function(self):
+        return torch.reshape(self._value_out, [-1])

rllib/examples/models/batch_norm_model.py

@@ -0,0 +1,190 @@
+import numpy as np
+
+from ray.rllib.models.modelv2 import ModelV2
+from ray.rllib.models.tf.misc import normc_initializer
+from ray.rllib.models.tf.tf_modelv2 import TFModelV2
+from ray.rllib.models.torch.misc import SlimFC, normc_initializer as \
+    torch_normc_initializer
+from ray.rllib.models.torch.torch_modelv2 import TorchModelV2
+from ray.rllib.utils import try_import_tf, try_import_torch
+from ray.rllib.utils.annotations import override
+
+tf = try_import_tf()
+torch, nn = try_import_torch()
+
+
+class BatchNormModel(TFModelV2):
+    """Example of a TFModelV2 that is built w/o using tf.keras.
+
+    NOTE: This example does not work when using a keras-based TFModelV2 due
+    to a bug in keras related to missing values for input placeholders, even
+    though these input values have been provided in a forward pass through the
+    actual keras Model.
+
+    All Model logic (layers) is defined in the `forward` method (incl.
+    the batch_normalization layers). Also, all variables are registered
+    (only once) at the end of `forward`, so an optimizer knows which tensors
+    to train on. A standard `value_function` override is used.
+    """
+    capture_index = 0
+
+    def __init__(self, obs_space, action_space, num_outputs, model_config,
+                 name):
+        super().__init__(obs_space, action_space, num_outputs, model_config,
+                         name)
+        # Have we registered our vars yet (see `forward`)?
+        self._registered = False
+
+    @override(ModelV2)
+    def forward(self, input_dict, state, seq_lens):
+        last_layer = input_dict["obs"]
+        hiddens = [256, 256]
+        with tf.variable_scope("model", reuse=tf.AUTO_REUSE):
+            for i, size in enumerate(hiddens):
+                last_layer = tf.layers.dense(
+                    last_layer,
+                    size,
+                    kernel_initializer=normc_initializer(1.0),
+                    activation=tf.nn.tanh,
+                    name="fc{}".format(i))
+                # Add a batch norm layer
+                last_layer = tf.layers.batch_normalization(
+                    last_layer,
+                    training=input_dict["is_training"],
+                    name="bn_{}".format(i))
+
+            output = tf.layers.dense(
+                last_layer,
+                self.num_outputs,
+                kernel_initializer=normc_initializer(0.01),
+                activation=None,
+                name="out")
+            self._value_out = tf.layers.dense(
+                last_layer,
+                1,
+                kernel_initializer=normc_initializer(1.0),
+                activation=None,
+                name="vf")
+        if not self._registered:
+            self.register_variables(
+                tf.get_collection(
+                    tf.GraphKeys.TRAINABLE_VARIABLES, scope=".+/model/.+"))
+            self._registered = True
+
+        return output, []
+
+    @override(ModelV2)
+    def value_function(self):
+        return tf.reshape(self._value_out, [-1])
+
+
+class KerasBatchNormModel(TFModelV2):
+    """Keras version of above BatchNormModel with exactly the same structure.
+
+    IMORTANT NOTE: This model will not work with PPO due to a bug in keras
+    that surfaces when having more than one input placeholder (here: `inputs`
+    and `is_training`) AND using the `make_tf_callable` helper (e.g. used by
+    PPO), in which auto-placeholders are generated, then passed through the
+    tf.keras. models.Model. In this last step, the connection between 1) the
+    provided value in the auto-placeholder and 2) the keras `is_training`
+    Input is broken and keras complains.
+    Use the above `BatchNormModel` (a non-keras based TFModelV2), instead.
+    """
+
+    def __init__(self, obs_space, action_space, num_outputs, model_config,
+                 name):
+        super().__init__(obs_space, action_space, num_outputs, model_config,
+                         name)
+        inputs = tf.keras.layers.Input(shape=obs_space.shape, name="inputs")
+        is_training = tf.keras.layers.Input(
+            shape=(), dtype=tf.bool, batch_size=1, name="is_training")
+        last_layer = inputs
+        hiddens = [256, 256]
+        for i, size in enumerate(hiddens):
+            label = "fc{}".format(i)
+            last_layer = tf.keras.layers.Dense(
+                units=size,
+                kernel_initializer=normc_initializer(1.0),
+                activation=tf.nn.tanh,
+                name=label)(last_layer)
+            # Add a batch norm layer
+            last_layer = tf.keras.layers.BatchNormalization()(
+                last_layer, training=is_training[0])
+        output = tf.keras.layers.Dense(
+            units=self.num_outputs,
+            kernel_initializer=normc_initializer(0.01),
+            activation=None,
+            name="fc_out")(last_layer)
+        value_out = tf.keras.layers.Dense(
+            units=1,
+            kernel_initializer=normc_initializer(0.01),
+            activation=None,
+            name="value_out")(last_layer)
+
+        self.base_model = tf.keras.models.Model(
+            inputs=[inputs, is_training], outputs=[output, value_out])
+        self.register_variables(self.base_model.variables)
+
+    @override(ModelV2)
+    def forward(self, input_dict, state, seq_lens):
+        out, self._value_out = self.base_model(
+            [input_dict["obs"], input_dict["is_training"]])
+        return out, []
+
+    @override(ModelV2)
+    def value_function(self):
+        return tf.reshape(self._value_out, [-1])
+
+
+class TorchBatchNormModel(TorchModelV2, nn.Module):
+    """Example of a TorchModelV2 using batch normalization."""
+    capture_index = 0
+
+    def __init__(self, obs_space, action_space, num_outputs, model_config,
+                 name, **kwargs):
+        TorchModelV2.__init__(self, obs_space, action_space, num_outputs,
+                              model_config, name)
+        nn.Module.__init__(self)
+        layers = []
+        prev_layer_size = int(np.product(obs_space.shape))
+        self._logits = None
+
+        # Create layers 0 to second-last.
+        for size in [256, 256]:
+            layers.append(
+                SlimFC(
+                    in_size=prev_layer_size,
+                    out_size=size,
+                    initializer=torch_normc_initializer(1.0),
+                    activation_fn=nn.ReLU))
+            prev_layer_size = size
+            # Add a batch norm layer.
+            layers.append(nn.BatchNorm1d(prev_layer_size))
+
+        self._logits = SlimFC(
+            in_size=prev_layer_size,
+            out_size=self.num_outputs,
+            initializer=torch_normc_initializer(0.01),
+            activation_fn=None)
+
+        self._value_branch = SlimFC(
+            in_size=prev_layer_size,
+            out_size=1,
+            initializer=torch_normc_initializer(1.0),
+            activation_fn=None)
+
+        self._hidden_layers = nn.Sequential(*layers)
+
+    @override(ModelV2)
+    def forward(self, input_dict, state, seq_lens):
+        # Set the correct train-mode for our hidden module (only important
+        # b/c we have some batch-norm layers).
+        self._hidden_layers.train(mode=input_dict["is_training"])
+        hidden_out = self._hidden_layers(input_dict["obs"])
+        logits = self._logits(hidden_out)
+        self._value_out = self._value_branch(hidden_out)
+        return logits, []
+
+    @override(ModelV2)
+    def value_function(self):
+        return torch.reshape(self._value_out, [-1])

rllib/examples/models/centralized_critic_models.py

@@ -0,0 +1,173 @@
+from gym.spaces import Box
+
+from ray.rllib.models.modelv2 import ModelV2
+from ray.rllib.models.tf.tf_modelv2 import TFModelV2
+from ray.rllib.models.tf.fcnet_v2 import FullyConnectedNetwork
+from ray.rllib.models.torch.misc import SlimFC
+from ray.rllib.models.torch.torch_modelv2 import TorchModelV2
+from ray.rllib.models.torch.fcnet import FullyConnectedNetwork as TorchFC
+from ray.rllib.utils.annotations import override
+from ray.rllib.utils.framework import try_import_tf, try_import_torch
+
+tf = try_import_tf()
+torch, nn = try_import_torch()
+
+
+class CentralizedCriticModel(TFModelV2):
+    """Multi-agent model that implements a centralized value function."""
+
+    def __init__(self, obs_space, action_space, num_outputs, model_config,
+                 name):
+        super(CentralizedCriticModel, self).__init__(
+            obs_space, action_space, num_outputs, model_config, name)
+        # Base of the model
+        self.model = FullyConnectedNetwork(obs_space, action_space,
+                                           num_outputs, model_config, name)
+        self.register_variables(self.model.variables())
+
+        # Central VF maps (obs, opp_obs, opp_act) -> vf_pred
+        obs = tf.keras.layers.Input(shape=(6, ), name="obs")
+        opp_obs = tf.keras.layers.Input(shape=(6, ), name="opp_obs")
+        opp_act = tf.keras.layers.Input(shape=(2, ), name="opp_act")
+        concat_obs = tf.keras.layers.Concatenate(axis=1)(
+            [obs, opp_obs, opp_act])
+        central_vf_dense = tf.keras.layers.Dense(
+            16, activation=tf.nn.tanh, name="c_vf_dense")(concat_obs)
+        central_vf_out = tf.keras.layers.Dense(
+            1, activation=None, name="c_vf_out")(central_vf_dense)
+        self.central_vf = tf.keras.Model(
+            inputs=[obs, opp_obs, opp_act], outputs=central_vf_out)
+        self.register_variables(self.central_vf.variables)
+
+    @override(ModelV2)
+    def forward(self, input_dict, state, seq_lens):
+        return self.model.forward(input_dict, state, seq_lens)
+
+    def central_value_function(self, obs, opponent_obs, opponent_actions):
+        return tf.reshape(
+            self.central_vf(
+                [obs, opponent_obs,
+                 tf.one_hot(opponent_actions, 2)]), [-1])
+
+    @override(ModelV2)
+    def value_function(self):
+        return self.model.value_function()  # not used
+
+
+class YetAnotherCentralizedCriticModel(TFModelV2):
+    """Multi-agent model that implements a centralized value function.
+
+    It assumes the observation is a dict with 'own_obs' and 'opponent_obs', the
+    former of which can be used for computing actions (i.e., decentralized
+    execution), and the latter for optimization (i.e., centralized learning).
+
+    This model has two parts:
+    - An action model that looks at just 'own_obs' to compute actions
+    - A value model that also looks at the 'opponent_obs' / 'opponent_action'
+      to compute the value (it does this by using the 'obs_flat' tensor).
+    """
+
+    def __init__(self, obs_space, action_space, num_outputs, model_config,
+                 name):
+        super(YetAnotherCentralizedCriticModel, self).__init__(
+            obs_space, action_space, num_outputs, model_config, name)
+
+        self.action_model = FullyConnectedNetwork(
+            Box(low=0, high=1, shape=(6, )),  # one-hot encoded Discrete(6)
+            action_space,
+            num_outputs,
+            model_config,
+            name + "_action")
+        self.register_variables(self.action_model.variables())
+
+        self.value_model = FullyConnectedNetwork(obs_space, action_space, 1,
+                                                 model_config, name + "_vf")
+        self.register_variables(self.value_model.variables())
+
+    def forward(self, input_dict, state, seq_lens):
+        self._value_out, _ = self.value_model({
+            "obs": input_dict["obs_flat"]
+        }, state, seq_lens)
+        return self.action_model({
+            "obs": input_dict["obs"]["own_obs"]
+        }, state, seq_lens)
+
+    def value_function(self):
+        return tf.reshape(self._value_out, [-1])
+
+
+class TorchCentralizedCriticModel(TorchModelV2, nn.Module):
+    """Multi-agent model that implements a centralized VF."""
+
+    def __init__(self, obs_space, action_space, num_outputs, model_config,
+                 name):
+        TorchModelV2.__init__(self, obs_space, action_space, num_outputs,
+                              model_config, name)
+        nn.Module.__init__(self)
+
+        # Base of the model
+        self.model = TorchFC(obs_space, action_space, num_outputs,
+                             model_config, name)
+
+        # Central VF maps (obs, opp_obs, opp_act) -> vf_pred
+        input_size = 6 + 6 + 2  # obs + opp_obs + opp_act
+        self.central_vf_dense = SlimFC(input_size, 16, activation_fn=nn.Tanh)
+        self.central_vf_out = SlimFC(16, 1)
+
+    @override(ModelV2)
+    def forward(self, input_dict, state, seq_lens):
+        model_out, _ = self.model(input_dict, state, seq_lens)
+        return model_out, []
+
+    def central_value_function(self, obs, opponent_obs, opponent_actions):
+        input_ = torch.cat([
+            obs, opponent_obs,
+            torch.nn.functional.one_hot(opponent_actions, 2)
+        ], 1)
+        return torch.reshape(
+            self.central_vf_out(self.central_vf_dense(input_)), [-1])
+
+    @override(ModelV2)
+    def value_function(self):
+        return self.model.value_function()  # not used
+
+
+class YetAnotherTorchCentralizedCriticModel(TorchModelV2, nn.Module):
+    """Multi-agent model that implements a centralized value function.
+
+    It assumes the observation is a dict with 'own_obs' and 'opponent_obs', the
+    former of which can be used for computing actions (i.e., decentralized
+    execution), and the latter for optimization (i.e., centralized learning).
+
+    This model has two parts:
+    - An action model that looks at just 'own_obs' to compute actions
+    - A value model that also looks at the 'opponent_obs' / 'opponent_action'
+      to compute the value (it does this by using the 'obs_flat' tensor).
+    """
+
+    def __init__(self, obs_space, action_space, num_outputs, model_config,
+                 name):
+        TorchModelV2.__init__(self, obs_space, action_space, num_outputs,
+                              model_config, name)
+        nn.Module.__init__(self)
+
+        self.action_model = TorchFC(
+            Box(low=0, high=1, shape=(6, )),  # one-hot encoded Discrete(6)
+            action_space,
+            num_outputs,
+            model_config,
+            name + "_action")
+
+        self.value_model = TorchFC(obs_space, action_space, 1, model_config,
+                                   name + "_vf")
+
+    def forward(self, input_dict, state, seq_lens):
+        self._value_out, _ = self.value_model({
+            "obs": input_dict["obs_flat"]
+        }, state, seq_lens)
+        return self.action_model({
+            "obs": input_dict["obs"]["own_obs"]
+        }, state, seq_lens)
+
+    def value_function(self):
+        return torch.reshape(self._value_out, [-1])

rllib/examples/models/eager_model.py

@@ -0,0 +1,60 @@
+import random
+
+from ray.rllib.models.modelv2 import ModelV2
+from ray.rllib.models.tf.fcnet_v2 import FullyConnectedNetwork
+from ray.rllib.models.tf.tf_modelv2 import TFModelV2
+from ray.rllib.utils.annotations import override
+from ray.rllib.utils.framework import try_import_tf
+
+tf = try_import_tf()
+
+
+class EagerModel(TFModelV2):
+    """Example of using embedded eager execution in a custom model.
+
+    This shows how to use tf.py_function() to execute a snippet of TF code
+    in eager mode. Here the `self.forward_eager` method just prints out
+    the intermediate tensor for debug purposes, but you can in general
+    perform any TF eager operation in tf.py_function().
+    """
+
+    def __init__(self, observation_space, action_space, num_outputs,
+                 model_config, name):
+        super().__init__(observation_space, action_space, num_outputs,
+                         model_config, name)
+
+        inputs = tf.keras.layers.Input(shape=observation_space.shape)
+        self.fcnet = FullyConnectedNetwork(
+            obs_space=self.obs_space,
+            action_space=self.action_space,
+            num_outputs=self.num_outputs,
+            model_config=self.model_config,
+            name="fc1")
+        out, value_out = self.fcnet.base_model(inputs)
+
+        def lambda_(x):
+            eager_out = tf.py_function(self.forward_eager, [x], tf.float32)
+            with tf.control_dependencies([eager_out]):
+                eager_out.set_shape(x.shape)
+                return eager_out
+
+        out = tf.keras.layers.Lambda(lambda_)(out)
+        self.base_model = tf.keras.models.Model(inputs, [out, value_out])
+        self.register_variables(self.base_model.variables)
+
+    @override(ModelV2)
+    def forward(self, input_dict, state, seq_lens):
+        out, self._value_out = self.base_model(input_dict["obs"], state,
+                                               seq_lens)
+        return out, []
+
+    @override(ModelV2)
+    def value_function(self):
+        return tf.reshape(self._value_out, [-1])
+
+    def forward_eager(self, feature_layer):
+        assert tf.executing_eagerly()
+        if random.random() > 0.99:
+            print("Eagerly printing the feature layer mean value",
+                  tf.reduce_mean(feature_layer))
+        return feature_layer

rllib/examples/models/fast_model.py

@@ -0,0 +1,76 @@
+from ray.rllib.models.modelv2 import ModelV2
+from ray.rllib.models.tf.tf_modelv2 import TFModelV2
+from ray.rllib.models.torch.misc import SlimFC
+from ray.rllib.models.torch.torch_modelv2 import TorchModelV2
+from ray.rllib.utils.annotations import override
+from ray.rllib.utils.framework import try_import_tf, try_import_torch
+
+tf = try_import_tf()
+torch, nn = try_import_torch()
+
+
+class FastModel(TFModelV2):
+    """An example for a non-Keras ModelV2 in tf that learns a single weight.
+
+    Defines all network architecture in `forward` (not `__init__` as it's
+    usually done for Keras-style TFModelV2s).
+    """
+
+    def __init__(self, obs_space, action_space, num_outputs, model_config,
+                 name):
+        super().__init__(obs_space, action_space, num_outputs, model_config,
+                         name)
+        # Have we registered our vars yet (see `forward`)?
+        self._registered = False
+
+    @override(ModelV2)
+    def forward(self, input_dict, state, seq_lens):
+        with tf.variable_scope("model", reuse=tf.AUTO_REUSE):
+            bias = tf.get_variable(
+                dtype=tf.float32,
+                name="bias",
+                initializer=tf.zeros_initializer,
+                shape=())
+            output = bias + \
+                tf.zeros([tf.shape(input_dict["obs"])[0], self.num_outputs])
+            self._value_out = tf.reduce_mean(output, -1)  # fake value
+
+        if not self._registered:
+            self.register_variables(
+                tf.get_collection(
+                    tf.GraphKeys.TRAINABLE_VARIABLES, scope=".+/model/.+"))
+            self._registered = True
+
+        return output, []
+
+    @override(ModelV2)
+    def value_function(self):
+        return tf.reshape(self._value_out, [-1])
+
+
+class TorchFastModel(TorchModelV2, nn.Module):
+    """Torch version of FastModel (tf)."""
+
+    def __init__(self, obs_space, action_space, num_outputs, model_config,
+                 name):
+        TorchModelV2.__init__(self, obs_space, action_space, num_outputs,
+                              model_config, name)
+        nn.Module.__init__(self)
+
+        self.bias = torch.tensor(
+            [0.0], dtype=torch.float32, requires_grad=True)
+
+        # Only needed to give some params to the optimizer (even though,
+        # they are never used anywhere).
+        self.dummy_layer = SlimFC(1, 1)
+
+    @override(ModelV2)
+    def forward(self, input_dict, state, seq_lens):
+        output = self.bias + \
+            torch.zeros(size=(input_dict["obs"].shape[0], self.num_outputs))
+        self._value_out = torch.mean(output, -1)  # fake value
+        return output, []
+
+    @override(ModelV2)
+    def value_function(self):
+        return torch.reshape(self._value_out, [-1])

rllib/examples/models/mobilenet_v2_with_lstm_models.py

@@ -0,0 +1,150 @@
+import numpy as np
+
+from ray.rllib.models.modelv2 import ModelV2
+from ray.rllib.models.tf.recurrent_tf_modelv2 import RecurrentTFModelV2
+from ray.rllib.models.torch.misc import SlimFC
+from ray.rllib.models.torch.recurrent_torch_model import RecurrentTorchModel
+from ray.rllib.utils.annotations import override
+from ray.rllib.utils.framework import try_import_tf, try_import_torch
+
+tf = try_import_tf()
+torch, nn = try_import_torch()
+
+
+class MobileV2PlusRNNModel(RecurrentTFModelV2):
+    """A conv. + recurrent keras net example using a pre-trained MobileNet."""
+
+    def __init__(self, obs_space, action_space, num_outputs, model_config,
+                 name, cnn_shape):
+
+        super(MobileV2PlusRNNModel, self).__init__(
+            obs_space, action_space, num_outputs, model_config, name)
+
+        self.cell_size = 16
+        visual_size = cnn_shape[0] * cnn_shape[1] * cnn_shape[2]
+
+        state_in_h = tf.keras.layers.Input(shape=(self.cell_size, ), name="h")
+        state_in_c = tf.keras.layers.Input(shape=(self.cell_size, ), name="c")
+        seq_in = tf.keras.layers.Input(shape=(), name="seq_in", dtype=tf.int32)
+
+        inputs = tf.keras.layers.Input(
+            shape=(None, visual_size), name="visual_inputs")
+
+        input_visual = inputs
+        input_visual = tf.reshape(
+            input_visual, [-1, cnn_shape[0], cnn_shape[1], cnn_shape[2]])
+        cnn_input = tf.keras.layers.Input(shape=cnn_shape, name="cnn_input")
+
+        cnn_model = tf.keras.applications.mobilenet_v2.MobileNetV2(
+            alpha=1.0,
+            include_top=True,
+            weights=None,
+            input_tensor=cnn_input,
+            pooling=None)
+        vision_out = cnn_model(input_visual)
+        vision_out = tf.reshape(
+            vision_out,
+            [-1, tf.shape(inputs)[1],
+             vision_out.shape.as_list()[-1]])
+
+        lstm_out, state_h, state_c = tf.keras.layers.LSTM(
+            self.cell_size,
+            return_sequences=True,
+            return_state=True,
+            name="lstm")(
+                inputs=vision_out,
+                mask=tf.sequence_mask(seq_in),
+                initial_state=[state_in_h, state_in_c])
+
+        # Postprocess LSTM output with another hidden layer and compute values.
+        logits = tf.keras.layers.Dense(
+            self.num_outputs,
+            activation=tf.keras.activations.linear,
+            name="logits")(lstm_out)
+        values = tf.keras.layers.Dense(
+            1, activation=None, name="values")(lstm_out)
+
+        # Create the RNN model
+        self.rnn_model = tf.keras.Model(
+            inputs=[inputs, seq_in, state_in_h, state_in_c],
+            outputs=[logits, values, state_h, state_c])
+        self.register_variables(self.rnn_model.variables)
+        self.rnn_model.summary()
+
+    @override(RecurrentTFModelV2)
+    def forward_rnn(self, inputs, state, seq_lens):
+        model_out, self._value_out, h, c = self.rnn_model([inputs, seq_lens] +
+                                                          state)
+        return model_out, [h, c]
+
+    @override(ModelV2)
+    def get_initial_state(self):
+        return [
+            np.zeros(self.cell_size, np.float32),
+            np.zeros(self.cell_size, np.float32),
+        ]
+
+    @override(ModelV2)
+    def value_function(self):
+        return tf.reshape(self._value_out, [-1])
+
+
+class TorchMobileV2PlusRNNModel(RecurrentTorchModel):
+    """A conv. + recurrent torch net example using a pre-trained MobileNet."""
+
+    def __init__(self, obs_space, action_space, num_outputs, model_config,
+                 name, cnn_shape):
+
+        super().__init__(obs_space, action_space, num_outputs, model_config,
+                         name)
+
+        self.lstm_state_size = 16
+        self.cnn_shape = list(cnn_shape)
+        self.visual_size_in = cnn_shape[0] * cnn_shape[1] * cnn_shape[2]
+        # MobileNetV2 has a flat output of (1000,).
+        self.visual_size_out = 1000
+
+        # Load the MobileNetV2 from torch.hub.
+        self.cnn_model = torch.hub.load(
+            "pytorch/vision:v0.6.0", "mobilenet_v2", pretrained=True)
+
+        self.lstm = nn.LSTM(
+            self.visual_size_out, self.lstm_state_size, batch_first=True)
+
+        # Postprocess LSTM output with another hidden layer and compute values.
+        self.logits = SlimFC(self.lstm_state_size, self.num_outputs)
+        self.value_branch = SlimFC(self.lstm_state_size, 1)
+
+    @override(RecurrentTFModelV2)
+    def forward_rnn(self, inputs, state, seq_lens):
+        # Create image dims.
+        vision_in = torch.reshape(inputs, [-1] + self.cnn_shape)
+        vision_out = self.cnn_model(vision_in)
+        # Flatten.
+        vision_out_time_ranked = torch.reshape(
+            vision_out,
+            [inputs.shape[0], inputs.shape[1], vision_out.shape[-1]])
+        if len(state[0].shape) == 2:
+            state[0] = state[0].unsqueeze(0)
+            state[1] = state[1].unsqueeze(0)
+        # Forward through LSTM.
+        lstm_out, [h, c] = self.lstm(vision_out_time_ranked, state)
+        # Forward LSTM out through logits layer and value layer.
+        logits = self.logits(lstm_out)
+        self._value_out = self.value_branch(lstm_out)
+        return logits, [h.squeeze(0), c.squeeze(0)]
+
+    @override(ModelV2)
+    def get_initial_state(self):
+        # Place hidden states on same device as model.
+        h = [
+            list(self.cnn_model.modules())[-1].weight.new(
+                1, self.lstm_state_size).zero_().squeeze(0),
+            list(self.cnn_model.modules())[-1].weight.new(
+                1, self.lstm_state_size).zero_().squeeze(0),
+        ]
+        return h
+
+    @override(ModelV2)
+    def value_function(self):
+        return torch.reshape(self._value_out, [-1])

rllib/examples/models/parametric_actions_model.py

@@ -0,0 +1,111 @@
+from gym.spaces import Box
+
+from ray.rllib.agents.dqn.distributional_q_tf_model import \
+    DistributionalQTFModel
+from ray.rllib.agents.dqn.dqn_torch_model import \
+    DQNTorchModel
+from ray.rllib.models.tf.fcnet_v2 import FullyConnectedNetwork
+from ray.rllib.models.torch.fcnet import FullyConnectedNetwork as TorchFC
+from ray.rllib.utils.framework import try_import_tf, try_import_torch
+from ray.rllib.utils.numpy import LARGE_INTEGER
+
+tf = try_import_tf()
+torch, nn = try_import_torch()
+
+
+class ParametricActionsModel(DistributionalQTFModel):
+    """Parametric action model that handles the dot product and masking.
+
+    This assumes the outputs are logits for a single Categorical action dist.
+    Getting this to work with a more complex output (e.g., if the action space
+    is a tuple of several distributions) is also possible but left as an
+    exercise to the reader.
+    """
+
+    def __init__(self,
+                 obs_space,
+                 action_space,
+                 num_outputs,
+                 model_config,
+                 name,
+                 true_obs_shape=(4, ),
+                 action_embed_size=2,
+                 **kw):
+        super(ParametricActionsModel, self).__init__(
+            obs_space, action_space, num_outputs, model_config, name, **kw)
+        self.action_embed_model = FullyConnectedNetwork(
+            Box(-1, 1, shape=true_obs_shape), action_space, action_embed_size,
+            model_config, name + "_action_embed")
+        self.register_variables(self.action_embed_model.variables())
+
+    def forward(self, input_dict, state, seq_lens):
+        # Extract the available actions tensor from the observation.
+        avail_actions = input_dict["obs"]["avail_actions"]
+        action_mask = input_dict["obs"]["action_mask"]
+
+        # Compute the predicted action embedding
+        action_embed, _ = self.action_embed_model({
+            "obs": input_dict["obs"]["cart"]
+        })
+
+        # 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(action_embed, 1)
+
+        # Batch dot product => 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)
+        return action_logits + inf_mask, state
+
+    def value_function(self):
+        return self.action_embed_model.value_function()
+
+
+class TorchParametricActionsModel(DQNTorchModel, nn.Module):
+    """PyTorch version of above ParametricActionsModel."""
+
+    def __init__(self,
+                 obs_space,
+                 action_space,
+                 num_outputs,
+                 model_config,
+                 name,
+                 true_obs_shape=(4, ),
+                 action_embed_size=2,
+                 **kw):
+        nn.Module.__init__(self)
+        DQNTorchModel.__init__(self, obs_space, action_space, num_outputs,
+                               model_config, name, **kw)
+
+        self.action_embed_model = TorchFC(
+            Box(-1, 1, shape=true_obs_shape), action_space, action_embed_size,
+            model_config, name + "_action_embed")
+
+    def forward(self, input_dict, state, seq_lens):
+        # Extract the available actions tensor from the observation.
+        avail_actions = input_dict["obs"]["avail_actions"]
+        action_mask = input_dict["obs"]["action_mask"]
+
+        # Compute the predicted action embedding
+        action_embed, _ = self.action_embed_model({
+            "obs": input_dict["obs"]["cart"]
+        })
+
+        # 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 = torch.unsqueeze(action_embed, 1)
+
+        # Batch dot product => shape of logits is [BATCH, MAX_ACTIONS].
+        action_logits = torch.sum(avail_actions * intent_vector, dim=2)
+
+        # Mask out invalid actions (use -LARGE_INTEGER to tag invalid).
+        # These are then recognized by the EpsilonGreedy exploration component
+        # as invalid actions that are not to be chosen.
+        inf_mask = torch.clamp(
+            torch.log(action_mask), -float(LARGE_INTEGER), float("inf"))
+        return action_logits + inf_mask, state
+
+    def value_function(self):
+        return self.action_embed_model.value_function()

rllib/examples/models/rnn_model.py

@@ -0,0 +1,138 @@
+import numpy as np
+
+from ray.rllib.models.modelv2 import ModelV2
+from ray.rllib.models.preprocessors import get_preprocessor
+from ray.rllib.models.tf.recurrent_tf_modelv2 import RecurrentTFModelV2
+from ray.rllib.models.torch.recurrent_torch_model import RecurrentTorchModel
+from ray.rllib.utils.annotations import override
+from ray.rllib.utils.framework import try_import_tf, try_import_torch
+
+tf = try_import_tf()
+torch, nn = try_import_torch()
+
+
+class RNNModel(RecurrentTFModelV2):
+    """Example of using the Keras functional API to define a RNN model."""
+
+    def __init__(self,
+                 obs_space,
+                 action_space,
+                 num_outputs,
+                 model_config,
+                 name,
+                 hiddens_size=256,
+                 cell_size=64):
+        super(RNNModel, self).__init__(obs_space, action_space, num_outputs,
+                                       model_config, name)
+        self.cell_size = cell_size
+
+        # Define input layers
+        input_layer = tf.keras.layers.Input(
+            shape=(None, obs_space.shape[0]), name="inputs")
+        state_in_h = tf.keras.layers.Input(shape=(cell_size, ), name="h")
+        state_in_c = tf.keras.layers.Input(shape=(cell_size, ), name="c")
+        seq_in = tf.keras.layers.Input(shape=(), name="seq_in", dtype=tf.int32)
+
+        # Preprocess observation with a hidden layer and send to LSTM cell
+        dense1 = tf.keras.layers.Dense(
+            hiddens_size, activation=tf.nn.relu, name="dense1")(input_layer)
+        lstm_out, state_h, state_c = tf.keras.layers.LSTM(
+            cell_size, return_sequences=True, return_state=True, name="lstm")(
+                inputs=dense1,
+                mask=tf.sequence_mask(seq_in),
+                initial_state=[state_in_h, state_in_c])
+
+        # Postprocess LSTM output with another hidden layer and compute values
+        logits = tf.keras.layers.Dense(
+            self.num_outputs,
+            activation=tf.keras.activations.linear,
+            name="logits")(lstm_out)
+        values = tf.keras.layers.Dense(
+            1, activation=None, name="values")(lstm_out)
+
+        # Create the RNN model
+        self.rnn_model = tf.keras.Model(
+            inputs=[input_layer, seq_in, state_in_h, state_in_c],
+            outputs=[logits, values, state_h, state_c])
+        self.register_variables(self.rnn_model.variables)
+        self.rnn_model.summary()
+
+    @override(RecurrentTFModelV2)
+    def forward_rnn(self, inputs, state, seq_lens):
+        model_out, self._value_out, h, c = self.rnn_model([inputs, seq_lens] +
+                                                          state)
+        return model_out, [h, c]
+
+    @override(ModelV2)
+    def get_initial_state(self):
+        return [
+            np.zeros(self.cell_size, np.float32),
+            np.zeros(self.cell_size, np.float32),
+        ]
+
+    @override(ModelV2)
+    def value_function(self):
+        return tf.reshape(self._value_out, [-1])
+
+
+class TorchRNNModel(RecurrentTorchModel):
+    def __init__(self,
+                 obs_space,
+                 action_space,
+                 num_outputs,
+                 model_config,
+                 name,
+                 fc_size=64,
+                 lstm_state_size=256):
+        super().__init__(obs_space, action_space, num_outputs, model_config,
+                         name)
+
+        self.obs_size = get_preprocessor(obs_space)(obs_space).size
+        self.fc_size = fc_size
+        self.lstm_state_size = lstm_state_size
+
+        # Build the Module from fc + LSTM + 2xfc (action + value outs).
+        self.fc1 = nn.Linear(self.obs_size, self.fc_size)
+        self.lstm = nn.LSTM(
+            self.fc_size, self.lstm_state_size, batch_first=True)
+        self.action_branch = nn.Linear(self.lstm_state_size, num_outputs)
+        self.value_branch = nn.Linear(self.lstm_state_size, 1)
+        # Store the value output to save an extra forward pass.
+        self._cur_value = None
+
+    @override(ModelV2)
+    def get_initial_state(self):
+        # Place hidden states on same device as model.
+        h = [
+            self.fc1.weight.new(1, self.lstm_state_size).zero_().squeeze(0),
+            self.fc1.weight.new(1, self.lstm_state_size).zero_().squeeze(0)
+        ]
+        return h
+
+    @override(ModelV2)
+    def value_function(self):
+        assert self._cur_value is not None, "must call forward() first"
+        return self._cur_value
+
+    @override(RecurrentTorchModel)
+    def forward_rnn(self, inputs, state, seq_lens):
+        """Feeds `inputs` (B x T x ..) through the Gru Unit.
+
+        Returns the resulting outputs as a sequence (B x T x ...).
+        Values are stored in self._cur_value in simple (B) shape (where B
+        contains both the B and T dims!).
+
+        Returns:
+            NN Outputs (B x T x ...) as sequence.
+            The state batches as a List of two items (c- and h-states).
+        """
+        x = nn.functional.relu(self.fc1(inputs))
+        lstm_out = self.lstm(
+            x, [torch.unsqueeze(state[0], 0),
+                torch.unsqueeze(state[1], 0)])
+        action_out = self.action_branch(lstm_out[0])
+        self._cur_value = torch.reshape(self.value_branch(lstm_out[0]), [-1])
+        return action_out, [
+            torch.squeeze(lstm_out[1][0], 0),
+            torch.squeeze(lstm_out[1][1], 0)
+        ]

rllib/examples/models/shared_weights_model.py

@@ -0,0 +1,124 @@
+import numpy as np
+
+from ray.rllib.models.modelv2 import ModelV2
+from ray.rllib.models.tf.tf_modelv2 import TFModelV2
+from ray.rllib.models.torch.misc import SlimFC
+from ray.rllib.models.torch.torch_modelv2 import TorchModelV2
+from ray.rllib.utils.annotations import override
+from ray.rllib.utils import try_import_tf, try_import_torch
+
+tf = try_import_tf()
+torch, nn = try_import_torch()
+
+
+class SharedWeightsModel1(TFModelV2):
+    """Example of weight sharing between two different TFModelV2s.
+
+    Here, we share the variables defined in the 'shared' variable scope
+    by entering it explicitly with tf.AUTO_REUSE. This creates the
+    variables for the 'fc1' layer in a global scope called 'shared'
+    (outside of the Policy's normal variable scope).
+    """
+
+    def __init__(self, observation_space, action_space, num_outputs,
+                 model_config, name):
+        super().__init__(observation_space, action_space, num_outputs,
+                         model_config, name)
+
+        inputs = tf.keras.layers.Input(observation_space.shape)
+        with tf.variable_scope(
+                tf.VariableScope(tf.AUTO_REUSE, "shared"),
+                reuse=tf.AUTO_REUSE,
+                auxiliary_name_scope=False):
+            last_layer = tf.keras.layers.Dense(
+                units=64, activation=tf.nn.relu, name="fc1")(inputs)
+        output = tf.keras.layers.Dense(
+            units=num_outputs, activation=None, name="fc_out")(last_layer)
+        vf = tf.keras.layers.Dense(
+            units=1, activation=None, name="value_out")(last_layer)
+        self.base_model = tf.keras.models.Model(inputs, [output, vf])
+        self.register_variables(self.base_model.variables)
+
+    @override(ModelV2)
+    def forward(self, input_dict, state, seq_lens):
+        out, self._value_out = self.base_model(input_dict["obs"])
+        return out, []
+
+    @override(ModelV2)
+    def value_function(self):
+        return tf.reshape(self._value_out, [-1])
+
+
+class SharedWeightsModel2(TFModelV2):
+    """The "other" TFModelV2 using the same shared space as the one above."""
+
+    def __init__(self, observation_space, action_space, num_outputs,
+                 model_config, name):
+        super().__init__(observation_space, action_space, num_outputs,
+                         model_config, name)
+
+        inputs = tf.keras.layers.Input(observation_space.shape)
+
+        # Weights shared with SharedWeightsModel1.
+        with tf.variable_scope(
+                tf.VariableScope(tf.AUTO_REUSE, "shared"),
+                reuse=tf.AUTO_REUSE,
+                auxiliary_name_scope=False):
+            last_layer = tf.keras.layers.Dense(
+                units=64, activation=tf.nn.relu, name="fc1")(inputs)
+        output = tf.keras.layers.Dense(
+            units=num_outputs, activation=None, name="fc_out")(last_layer)
+        vf = tf.keras.layers.Dense(
+            units=1, activation=None, name="value_out")(last_layer)
+        self.base_model = tf.keras.models.Model(inputs, [output, vf])
+        self.register_variables(self.base_model.variables)
+
+    @override(ModelV2)
+    def forward(self, input_dict, state, seq_lens):
+        out, self._value_out = self.base_model(input_dict["obs"])
+        return out, []
+
+    @override(ModelV2)
+    def value_function(self):
+        return tf.reshape(self._value_out, [-1])
+
+
+TORCH_GLOBAL_SHARED_LAYER = None
+if torch:
+    TORCH_GLOBAL_SHARED_LAYER = SlimFC(32, 32)
+
+
+class TorchSharedWeightsModel(TorchModelV2, nn.Module):
+    """Example of weight sharing between two different TorchModelV2s.
+
+    The shared (single) layer is simply defined outside of the two Models,
+    then used by both Models in their forward pass.
+    """
+
+    def __init__(self, observation_space, action_space, num_outputs,
+                 model_config, name):
+        TorchModelV2.__init__(self, observation_space, action_space,
+                              num_outputs, model_config, name)
+        nn.Module.__init__(self)
+
+        # Non-shared initial layer.
+        self.first_layer = SlimFC(
+            int(np.product(observation_space.shape)),
+            32,
+            activation_fn=nn.ReLU)
+
+        # Non-shared final layer.
+        self.last_layer = SlimFC(32, self.num_outputs, activation_fn=nn.ReLU)
+        self.vf = SlimFC(32, 1, activation_fn=None)
+
+    @override(ModelV2)
+    def forward(self, input_dict, state, seq_lens):
+        out = self.first_layer(input_dict["obs"])
+        out = TORCH_GLOBAL_SHARED_LAYER(out)
+        model_out = self.last_layer(out)
+        self._value_out = self.vf(out)
+        return model_out, []
+
+    @override(ModelV2)
+    def value_function(self):
+        return torch.reshape(self._value_out, [-1])

rllib/examples/policy/random_policy.py

@@ -0,0 +1,37 @@
+import numpy as np
+import random
+
+from ray.rllib.policy.policy import Policy
+from ray.rllib.utils.annotations import override
+
+
+class RandomPolicy(Policy):
+    """Hand-coded policy that returns random actions."""
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+
+    @override(Policy)
+    def compute_actions(self,
+                        obs_batch,
+                        state_batches=None,
+                        prev_action_batch=None,
+                        prev_reward_batch=None,
+                        **kwargs):
+        # Alternatively, a numpy array would work here as well.
+        # e.g.: np.array([random.choice([0, 1])] * len(obs_batch))
+        return [self.action_space.sample() for _ in obs_batch], [], {}
+
+    @override(Policy)
+    def learn_on_batch(self, samples):
+        """No learning."""
+        return {}
+
+    @override(Policy)
+    def compute_log_likelihoods(self,
+                                actions,
+                                obs_batch,
+                                state_batches=None,
+                                prev_action_batch=None,
+                                prev_reward_batch=None):
+        return np.array([random.random()] * len(obs_batch))

rllib/examples/policy/rock_paper_scissors_dummies.py

@@ -0,0 +1,65 @@
+import random
+
+from ray.rllib.examples.env.rock_paper_scissors import RockPaperScissors
+from ray.rllib.policy.policy import Policy
+
+
+class AlwaysSameHeuristic(Policy):
+    """Pick a random move and stick with it for the entire episode."""
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.exploration = self._create_exploration()
+
+    def get_initial_state(self):
+        return [
+            random.choice([
+                RockPaperScissors.ROCK, RockPaperScissors.PAPER,
+                RockPaperScissors.SCISSORS
+            ])
+        ]
+
+    def compute_actions(self,
+                        obs_batch,
+                        state_batches=None,
+                        prev_action_batch=None,
+                        prev_reward_batch=None,
+                        info_batch=None,
+                        episodes=None,
+                        **kwargs):
+        return state_batches[0], state_batches, {}
+
+
+class BeatLastHeuristic(Policy):
+    """Play the move that would beat the last move of the opponent."""
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.exploration = self._create_exploration()
+
+    def compute_actions(self,
+                        obs_batch,
+                        state_batches=None,
+                        prev_action_batch=None,
+                        prev_reward_batch=None,
+                        info_batch=None,
+                        episodes=None,
+                        **kwargs):
+        def successor(x):
+            if x[RockPaperScissors.ROCK] == 1:
+                return RockPaperScissors.PAPER
+            elif x[RockPaperScissors.PAPER] == 1:
+                return RockPaperScissors.SCISSORS
+            elif x[RockPaperScissors.SCISSORS] == 1:
+                return RockPaperScissors.ROCK
+
+        return [successor(x) for x in obs_batch], [], {}
+
+    def learn_on_batch(self, samples):
+        pass
+
+    def get_weights(self):
+        pass
+
+    def set_weights(self, weights):
+        pass

rllib/models/catalog.py

@@ -169,13 +169,13 @@ class ModelCatalog:
                 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]
+            input_lens = [int(e[1]) for e in child_dists_and_in_lens]
             return partial(
                 (TorchMultiActionDistribution
                  if framework == "torch" else MultiActionDistribution),
                 action_space=action_space,
                 child_distributions=child_dists,
-                input_lens=input_lens), sum(input_lens)
+                input_lens=input_lens), int(sum(input_lens))
         # Simplex -> Dirichlet.
         elif isinstance(action_space, Simplex):
             if framework == "torch":
@@ -281,14 +281,16 @@ class ModelCatalog:
             model_cls = _global_registry.get(RLLIB_MODEL,
                                              model_config["custom_model"])
             if issubclass(model_cls, ModelV2):
+
+                logger.info("Wrapping {} as {}".format(model_cls,
+                                                       model_interface))
+                model_cls = ModelCatalog._wrap_if_needed(
+                    model_cls, model_interface)
+
                 if framework == "tf":
-                    logger.info("Wrapping {} as {}".format(
-                        model_cls, model_interface))
-                    model_cls = ModelCatalog._wrap_if_needed(
-                        model_cls, model_interface)
+                    # Track and warn if vars were created but not registered.
                     created = set()
 
-                    # Track and warn if vars were created but not registered
                     def track_var_creation(next_creator, **kw):
                         v = next_creator(**kw)
                         created.add(v)
@@ -312,10 +314,13 @@ class ModelCatalog:
                             "question?".format(not_registered, instance,
                                                registered))
                 else:
-                    # no variable tracking
+                    # PyTorch automatically tracks nn.Modules inside the parent
+                    # nn.Module's constructor.
+                    # TODO(sven): Do this for TF as well.
                     instance = model_cls(obs_space, action_space, num_outputs,
                                          model_config, name, **model_kwargs)
                 return instance
+
             elif tf.executing_eagerly():
                 raise ValueError(
                     "Eager execution requires a TFModelV2 model to be "

rllib/models/torch/fcnet.py

@@ -88,7 +88,7 @@ class FullyConnectedNetwork(TorchModelV2, nn.Module):
 
     @override(TorchModelV2)
     def forward(self, input_dict, state, seq_lens):
-        obs = input_dict["obs_flat"]
+        obs = input_dict["obs_flat"].float()
         features = self._hidden_layers(obs.reshape(obs.shape[0], -1))
         logits = self._logits(features) if self._logits else features
         self._cur_value = self._value_branch(features).squeeze(1)

rllib/models/torch/recurrent_torch_model.py

@@ -31,7 +31,7 @@ class RecurrentTorchModel(TorchModelV2, nn.Module):
 
         @override(ModelV2)
         def get_initial_state(self):
-            # make hidden states on same device as model
+            # Place hidden states on same device as model.
             h = [self.fc1.weight.new(
                 1, self.rnn_hidden_dim).zero_().squeeze(0)]
             return h

rllib/policy/dynamic_tf_policy.py

@@ -156,11 +156,15 @@ class DynamicTFPolicy(TFPolicy):
             self.model = make_model(self, obs_space, action_space, config)
         else:
             self.model = ModelCatalog.get_model_v2(
-                obs_space,
-                action_space,
-                logit_dim,
-                self.config["model"],
+                obs_space=obs_space,
+                action_space=action_space,
+                num_outputs=logit_dim,
+                model_config=self.config["model"],
                 framework="tf")
+            # NOTE: Adding below line will break existing custom models
+            # that do not expect extra options in **kwargs but rather in
+            # model_config["custom_options"].
+            # **self.config["model"].get("custom_options", {}))
 
         # Create the Exploration object to use for this Policy.
         self.exploration = self._create_exploration()

rllib/policy/torch_policy.py

@@ -232,7 +232,7 @@ class TorchPolicy(Policy):
         loss_out = force_list(
             self._loss(self, self.model, self.dist_class, train_batch))
         assert len(loss_out) == len(self._optimizers)
-        assert not any(np.isnan(l.detach().numpy()) for l in loss_out)
+        # assert not any(np.isnan(l.detach().numpy()) for l in loss_out)
 
         # Loop through all optimizers.
         grad_info = {"allreduce_latency": 0.0}

rllib/rollout.py

@@ -32,6 +32,10 @@ Example Usage via executable:
 
 # Note: if you use any custom models or envs, register them here first, e.g.:
 #
+# from ray.rllib.examples.env.parametric_actions_cartpole import \
+#     ParametricActionsCartPole
+# from ray.rllib.examples.model.parametric_actions_model import \
+#     ParametricActionsModel
 # ModelCatalog.register_custom_model("pa_model", ParametricActionsModel)
 # register_env("pa_cartpole", lambda _: ParametricActionsCartPole(10))
 

rllib/utils/exploration/epsilon_greedy.py

@@ -6,6 +6,7 @@ from ray.rllib.utils.exploration.exploration import Exploration, TensorType
 from ray.rllib.utils.framework import try_import_tf, try_import_torch, \
     get_variable
 from ray.rllib.utils.from_config import from_config
+from ray.rllib.utils.numpy import LARGE_INTEGER
 from ray.rllib.utils.schedules import Schedule, PiecewiseSchedule
 
 tf = try_import_tf()
@@ -136,9 +137,8 @@ class EpsilonGreedy(Exploration):
             # Mask out actions, whose Q-values are -inf, so that we don't
             # even consider them for exploration.
             random_valid_action_logits = torch.where(
-                q_values == float("-inf"),
-                torch.ones_like(q_values) * float("-inf"),
-                torch.ones_like(q_values))
+                q_values == -float(LARGE_INTEGER),
+                torch.ones_like(q_values) * 0.0, torch.ones_like(q_values))
             # A random action.
             random_actions = torch.squeeze(
                 torch.multinomial(random_valid_action_logits, 1), axis=1)

rllib/utils/exploration/exploration.py

@@ -39,8 +39,11 @@ class Exploration:
         self.framework = check_framework(framework)
         # The device on which the Model has been placed.
         # This Exploration will be on the same device.
-        self.device = None if not isinstance(self.model, nn.Module) else \
-            next(self.model.parameters()).device
+        self.device = None
+        if isinstance(self.model, nn.Module):
+            params = list(self.model.parameters())
+            if params:
+                self.device = params[0].device
 
     @DeveloperAPI
     def before_compute_actions(self,

rllib/utils/test_utils.py

@@ -225,3 +225,21 @@ def check(x, y, decimals=5, atol=None, rtol=None, false=False):
                 if false is True:
                     assert False, \
                         "ERROR: x ({}) is the same as y ({})!".format(x, y)
+
+
+def check_learning_achieved(tune_results, min_reward):
+    """Throws an error if `min_reward` is not reached within tune_results.
+
+    Checks the last iteration found in tune_results for its
+    "episode_reward_mean" value and compares it to `min_reward`.
+
+    Args:
+        tune_results: The tune.run returned results object.
+        min_reward (float): The min reward that must be reached.
+
+    Throws:
+        ValueError: If `min_reward` not reached.
+    """
+    if tune_results.trials[0].last_result["episode_reward_mean"] < min_reward:
+        raise ValueError("`stop-reward` of {} not reached!".format(min_reward))
+    print("ok")