[RLlib] Support easy use_attention=True flag for using the GTrXL model. (#11698)

doc/source/rllib-models.rst

@@ -32,6 +32,7 @@ Thereby, always make sure that the last Conv2D output has an output shape of `[B
 X=last Conv2D layer's number of filters, so that RLlib can flatten it. An informative error will be thrown if this is not the case.
 
 In addition, if you set ``"model": {"use_lstm": true}``, the model output will be further processed by an LSTM cell (`TF <https://github.com/ray-project/ray/blob/master/rllib/models/tf/recurrent_net.py>`__ or `Torch <https://github.com/ray-project/ray/blob/master/rllib/models/torch/recurrent_net.py>`__).
+
 More generally, RLlib supports the use of recurrent models for its policy gradient algorithms (A3C, PPO, PG, IMPALA), and RNN support is built into its policy evaluation utilities.
 For custom RNN/LSTM setups, see the `Recurrent Models`_. section below.
 

rllib/evaluation/tests/test_trajectory_view_api.py

@@ -146,7 +146,7 @@ class TestTrajectoryViewAPI(unittest.TestCase):
         config["model"]["custom_model"] = GTrXLNet
         config["model"]["custom_model_config"] = {
             "num_transformer_units": 1,
-            "attn_dim": 64,
+            "attention_dim": 64,
             "num_heads": 2,
             "memory_inference": 50,
             "memory_training": 50,

rllib/examples/attention_net.py

@@ -3,8 +3,6 @@ import os
 
 import ray
 from ray import tune
-from ray.rllib.models.tf.attention_net import GTrXLNet
-from ray.rllib.models.torch.attention_net import GTrXLNet as TorchGTrXLNet
 from ray.rllib.examples.env.look_and_push import LookAndPush, OneHot
 from ray.rllib.examples.env.repeat_after_me_env import RepeatAfterMeEnv
 from ray.rllib.examples.env.repeat_initial_obs_env import RepeatInitialObsEnv
@@ -51,17 +49,15 @@ if __name__ == "__main__":
         "num_sgd_iter": 10,
         "vf_loss_coeff": 1e-5,
         "model": {
-            "custom_model": TorchGTrXLNet if args.torch else GTrXLNet,
+            "use_attention": True,
             "max_seq_len": 50,
-            "custom_model_config": {
-                "num_transformer_units": 1,
-                "attn_dim": 64,
-                "memory_inference": 100,
-                "memory_training": 50,
-                "head_dim": 32,
-                "num_heads": 2,
-                "ff_hidden_dim": 32,
-            },
+            "attention_num_transformer_units": 1,
+            "attention_dim": 64,
+            "attention_memory_inference": 100,
+            "attention_memory_training": 50,
+            "attention_num_heads": 2,
+            "attention_head_dim": 32,
+            "attention_position_wise_mlp_dim": 32,
         },
         "framework": "torch" if args.torch else "tf",
     }

rllib/examples/attention_net_supervised.py

@@ -33,10 +33,10 @@ def train_bit_shift(seq_length, num_iterations, print_every_n):
         model_config={"max_seq_len": seq_length},
         name="trxl",
         num_transformer_units=1,
-        attn_dim=10,
+        attention_dim=10,
         num_heads=5,
         head_dim=20,
-        ff_hidden_dim=20,
+        position_wise_mlp_dim=20,
     )
 
     shift = 10

rllib/models/catalog.py

@@ -11,7 +11,6 @@ from ray.rllib.models.action_dist import ActionDistribution
 from ray.rllib.models.jax.jax_action_dist import JAXCategorical
 from ray.rllib.models.modelv2 import ModelV2
 from ray.rllib.models.preprocessors import get_preprocessor, Preprocessor
-from ray.rllib.models.tf.recurrent_net import LSTMWrapper
 from ray.rllib.models.tf.tf_action_dist import Categorical, \
     Deterministic, DiagGaussian, Dirichlet, \
     MultiActionDistribution, MultiCategorical
@@ -68,6 +67,32 @@ MODEL_DEFAULTS: ModelConfigDict = {
     # Whether the LSTM is time-major (TxBx..) or batch-major (BxTx..).
     "_time_major": False,
 
+    # == Attention Nets (experimental: torch-version is untested) ==
+    # Whether to use a GTrXL ("Gru transformer XL"; attention net) as the
+    # wrapper Model around the default Model.
+    "use_attention": False,
+    # The number of transformer units within GTrXL.
+    # A transformer unit in GTrXL consists of a) MultiHeadAttention module and
+    # b) a position-wise MLP.
+    "attention_num_transformer_units": 1,
+    # The input and output size of each transformer unit.
+    "attention_dim": 64,
+    # The number of attention heads within the MultiHeadAttention units.
+    "attention_num_heads": 1,
+    # The dim of a single head (within the MultiHeadAttention units).
+    "attention_head_dim": 32,
+    # The memory sizes for inference and training.
+    "attention_memory_inference": 50,
+    "attention_memory_training": 50,
+    # The output dim of the position-wise MLP.
+    "attention_position_wise_mlp_dim": 32,
+    # The initial bias values for the 2 GRU gates within a transformer unit.
+    "attention_init_gru_gate_bias": 2.0,
+    # TODO: Whether to feed a_{t-n:t-1} to GTrXL (one-hot encoded if discrete).
+    # "attention_use_n_prev_actions": 0,
+    # Whether to feed r_{t-n:t-1} to GTrXL.
+    # "attention_use_n_prev_rewards": 0,
+
     # == Atari ==
     # Whether to enable framestack for Atari envs
     "framestack": True,
@@ -283,6 +308,8 @@ class ModelCatalog:
                 unflatten the tensor into a ragged tensor.
             action_space (Space): Action space of the target gym env.
             num_outputs (int): The size of the output vector of the model.
+            model_config (ModelConfigDict): The "model" sub-config dict
+                within the Trainer's config dict.
             framework (str): One of "tf2", "tf", "tfe", "torch", or "jax".
             name (str): Name (scope) for the model.
             model_interface (cls): Interface required for the model
@@ -294,6 +321,9 @@ class ModelCatalog:
             model (ModelV2): Model to use for the policy.
         """
 
+        # Validate the given config dict.
+        ModelCatalog._validate_config(config=model_config, framework=framework)
+
         if model_config.get("custom_model"):
             # Allow model kwargs to be overridden / augmented by
             # custom_model_config.
@@ -316,12 +346,18 @@ class ModelCatalog:
                                                      model_interface)
 
             if framework in ["tf2", "tf", "tfe"]:
-                # Try wrapping custom model with LSTM, if required.
-                if model_config.get("use_lstm"):
+                # Try wrapping custom model with LSTM/attention, if required.
+                if model_config.get("use_lstm") or \
+                        model_config.get("use_attention"):
+                    from ray.rllib.models.tf.attention_net import \
+                        AttentionWrapper
+                    from ray.rllib.models.tf.recurrent_net import LSTMWrapper
+
                     wrapped_cls = model_cls
                     forward = wrapped_cls.forward
                     model_cls = ModelCatalog._wrap_if_needed(
-                        wrapped_cls, LSTMWrapper)
+                        wrapped_cls, LSTMWrapper
+                        if model_config.get("use_lstm") else AttentionWrapper)
                     model_cls._wrapped_forward = forward
 
                 # Track and warn if vars were created but not registered.
@@ -367,6 +403,21 @@ class ModelCatalog:
                         "question?".format(not_registered, instance,
                                            registered))
             elif framework == "torch":
+                # Try wrapping custom model with LSTM/attention, if required.
+                if model_config.get("use_lstm") or \
+                        model_config.get("use_attention"):
+                    from ray.rllib.models.torch.attention_net import \
+                        AttentionWrapper
+                    from ray.rllib.models.torch.recurrent_net import \
+                        LSTMWrapper
+
+                    wrapped_cls = model_cls
+                    forward = wrapped_cls.forward
+                    model_cls = ModelCatalog._wrap_if_needed(
+                        wrapped_cls, LSTMWrapper
+                        if model_config.get("use_lstm") else AttentionWrapper)
+                    model_cls._wrapped_forward = forward
+
                 # PyTorch automatically tracks nn.Modules inside the parent
                 # nn.Module's constructor.
                 # Try calling with kwargs first (custom ModelV2 should
@@ -402,16 +453,27 @@ class ModelCatalog:
             # Try to get a default v2 model.
             if not model_config.get("custom_model"):
                 v2_class = default_model or ModelCatalog._get_v2_model_class(
-                    obs_space, model_config, framework=framework)
+                    obs_space, framework=framework)
 
             if not v2_class:
                 raise ValueError("ModelV2 class could not be determined!")
 
-            if model_config.get("use_lstm"):
+            if model_config.get("use_lstm") or \
+                    model_config.get("use_attention"):
+
+                from ray.rllib.models.tf.attention_net import \
+                    AttentionWrapper
+                from ray.rllib.models.tf.recurrent_net import LSTMWrapper
+
                 wrapped_cls = v2_class
                 forward = wrapped_cls.forward
-                v2_class = ModelCatalog._wrap_if_needed(
-                    wrapped_cls, LSTMWrapper)
+                if model_config.get("use_lstm"):
+                    v2_class = ModelCatalog._wrap_if_needed(
+                        wrapped_cls, LSTMWrapper)
+                else:
+                    v2_class = ModelCatalog._wrap_if_needed(
+                        wrapped_cls, AttentionWrapper)
+
                 v2_class._wrapped_forward = forward
 
             # Wrap in the requested interface.
@@ -421,17 +483,32 @@ class ModelCatalog:
 
         # Find a default TorchModelV2 and wrap with model_interface.
         elif framework == "torch":
-            v2_class = \
-                default_model or ModelCatalog._get_v2_model_class(
-                    obs_space, model_config, framework=framework)
-            if model_config.get("use_lstm"):
-                from ray.rllib.models.torch.recurrent_net import LSTMWrapper \
-                    as TorchLSTMWrapper
+            # Try to get a default v2 model.
+            if not model_config.get("custom_model"):
+                v2_class = default_model or ModelCatalog._get_v2_model_class(
+                    obs_space, framework=framework)
+
+            if not v2_class:
+                raise ValueError("ModelV2 class could not be determined!")
+
+            if model_config.get("use_lstm") or \
+                    model_config.get("use_attention"):
+
+                from ray.rllib.models.torch.attention_net import \
+                    AttentionWrapper
+                from ray.rllib.models.torch.recurrent_net import LSTMWrapper
+
                 wrapped_cls = v2_class
                 forward = wrapped_cls.forward
-                v2_class = ModelCatalog._wrap_if_needed(
-                    wrapped_cls, TorchLSTMWrapper)
+                if model_config.get("use_lstm"):
+                    v2_class = ModelCatalog._wrap_if_needed(
+                        wrapped_cls, LSTMWrapper)
+                else:
+                    v2_class = ModelCatalog._wrap_if_needed(
+                        wrapped_cls, AttentionWrapper)
+
                 v2_class._wrapped_forward = forward
+
             # Wrap in the requested interface.
             wrapper = ModelCatalog._wrap_if_needed(v2_class, model_interface)
             return wrapper(obs_space, action_space, num_outputs, model_config,
@@ -441,9 +518,7 @@ class ModelCatalog:
         elif framework == "jax":
             v2_class = \
                 default_model or ModelCatalog._get_v2_model_class(
-                    obs_space, model_config, framework=framework)
-            if model_config.get("use_lstm"):
-                raise NotImplementedError("JAXModel's not LSTM-wrappable yet!")
+                    obs_space, framework=framework)
             # Wrap in the requested interface.
             wrapper = ModelCatalog._wrap_if_needed(v2_class, model_interface)
             return wrapper(obs_space, action_space, num_outputs, model_config,
@@ -568,7 +643,6 @@ class ModelCatalog:
 
     @staticmethod
     def _get_v2_model_class(input_space: gym.Space,
-                            model_config: ModelConfigDict,
                             framework: str = "tf") -> ModelV2:
 
         VisionNet = None
@@ -622,3 +696,26 @@ class ModelCatalog:
                 child_distributions=child_dists,
                 input_lens=input_lens), int(sum(input_lens))
         return dist_class
+
+    @staticmethod
+    def _validate_config(config: ModelConfigDict, framework: str) -> None:
+        """Validates a given model config dict.
+
+        Args:
+            config (ModelConfigDict): The "model" sub-config dict
+                within the Trainer's config dict.
+            framework (str): One of "jax", "tf2", "tf", "tfe", or "torch".
+
+        Raises:
+            ValueError: If something is wrong with the given config.
+        """
+        if config.get("use_attention") and config.get("use_lstm"):
+            raise ValueError("Only one of `use_lstm` or `use_attention` may "
+                             "be set to True!")
+        if framework == "jax":
+            if config.get("use_attention"):
+                raise ValueError("`use_attention` not available for "
+                                 "framework=jax so far!")
+            elif config.get("use_lstm"):
+                raise ValueError("`use_lstm` not available for "
+                                 "framework=jax so far!")

rllib/models/tf/attention_net.py

@@ -8,14 +8,15 @@
       Z. Dai, Z. Yang, et al. - Carnegie Mellon U - 2019.
       https://www.aclweb.org/anthology/P19-1285.pdf
 """
-from gym.spaces import Box
+from gym.spaces import Box, Discrete, MultiDiscrete
 import numpy as np
 import gym
-from typing import Any, Optional
+from typing import Any, Dict, Optional, Union
 
 from ray.rllib.models.modelv2 import ModelV2
 from ray.rllib.models.tf.layers import GRUGate, RelativeMultiHeadAttention, \
     SkipConnection
+from ray.rllib.models.tf.tf_modelv2 import TFModelV2
 from ray.rllib.models.tf.recurrent_net import RecurrentNetwork
 from ray.rllib.policy.sample_batch import SampleBatch
 from ray.rllib.policy.view_requirement import ViewRequirement
@@ -27,7 +28,7 @@ tf1, tf, tfv = try_import_tf()
 
 
 # TODO(sven): Use RLlib's FCNet instead.
-class PositionwiseFeedforward(tf.keras.layers.Layer):
+class PositionwiseFeedforward(tf.keras.layers.Layer if tf else object):
     """A 2x linear layer with ReLU activation in between described in [1].
 
     Each timestep coming from the attention head will be passed through this
@@ -61,31 +62,31 @@ class TrXLNet(RecurrentNetwork):
     def __init__(self, observation_space: gym.spaces.Space,
                  action_space: gym.spaces.Space, num_outputs: int,
                  model_config: ModelConfigDict, name: str,
-                 num_transformer_units: int, attn_dim: int, num_heads: int,
-                 head_dim: int, ff_hidden_dim: int):
+                 num_transformer_units: int, attention_dim: int,
+                 num_heads: int, head_dim: int, position_wise_mlp_dim: int):
         """Initializes a TrXLNet object.
 
         Args:
             num_transformer_units (int): The number of Transformer repeats to
                 use (denoted L in [2]).
-            attn_dim (int): The input and output dimensions of one Transformer
-                unit.
+            attention_dim (int): The input and output dimensions of one
+                Transformer unit.
             num_heads (int): The number of attention heads to use in parallel.
                 Denoted as `H` in [3].
-            head_dim (int): The dimension of a single(!) head.
-                Denoted as `d` in [3].
-            ff_hidden_dim (int): The dimension of the hidden layer within
-                the position-wise MLP (after the multi-head attention block
-                within one Transformer unit). This is the size of the first
-                of the two layers within the PositionwiseFeedforward. The
-                second layer always has size=`attn_dim`.
+            head_dim (int): The dimension of a single(!) attention head within
+                a multi-head attention unit. Denoted as `d` in [3].
+            position_wise_mlp_dim (int): The dimension of the hidden layer
+                within the position-wise MLP (after the multi-head attention
+                block within one Transformer unit). This is the size of the
+                first of the two layers within the PositionwiseFeedforward. The
+                second layer always has size=`attention_dim`.
         """
 
         super().__init__(observation_space, action_space, num_outputs,
                          model_config, name)
 
         self.num_transformer_units = num_transformer_units
-        self.attn_dim = attn_dim
+        self.attention_dim = attention_dim
         self.num_heads = num_heads
         self.head_dim = head_dim
         self.max_seq_len = model_config["max_seq_len"]
@@ -93,19 +94,20 @@ class TrXLNet(RecurrentNetwork):
 
         inputs = tf.keras.layers.Input(
             shape=(self.max_seq_len, self.obs_dim), name="inputs")
-        E_out = tf.keras.layers.Dense(attn_dim)(inputs)
+        E_out = tf.keras.layers.Dense(attention_dim)(inputs)
 
         for _ in range(self.num_transformer_units):
             MHA_out = SkipConnection(
                 RelativeMultiHeadAttention(
-                    out_dim=attn_dim,
+                    out_dim=attention_dim,
                     num_heads=num_heads,
                     head_dim=head_dim,
                     input_layernorm=False,
                     output_activation=None),
                 fan_in_layer=None)(E_out)
             E_out = SkipConnection(
-                PositionwiseFeedforward(attn_dim, ff_hidden_dim))(MHA_out)
+                PositionwiseFeedforward(attention_dim,
+                                        position_wise_mlp_dim))(MHA_out)
             E_out = tf.keras.layers.LayerNormalization(axis=-1)(E_out)
 
         # Postprocess TrXL output with another hidden layer and compute values.
@@ -158,7 +160,7 @@ class GTrXLNet(RecurrentNetwork):
         >> config["model"]["max_seq_len"] = 10
         >> config["model"]["custom_model_config"] = {
         >>     num_transformer_units=1,
-        >>     attn_dim=32,
+        >>     attention_dim=32,
         >>     num_heads=2,
         >>     memory_inference=100,
         >>     memory_training=50,
@@ -169,24 +171,25 @@ class GTrXLNet(RecurrentNetwork):
     def __init__(self,
                  observation_space: gym.spaces.Space,
                  action_space: gym.spaces.Space,
-                 num_outputs: int,
+                 num_outputs: Optional[int],
                  model_config: ModelConfigDict,
                  name: str,
-                 num_transformer_units: int,
-                 attn_dim: int,
-                 num_heads: int,
-                 memory_inference: int,
-                 memory_training: int,
-                 head_dim: int,
-                 ff_hidden_dim: int,
-                 init_gate_bias: float = 2.0):
+                 *,
+                 num_transformer_units: int = 1,
+                 attention_dim: int = 64,
+                 num_heads: int = 2,
+                 memory_inference: int = 50,
+                 memory_training: int = 50,
+                 head_dim: int = 32,
+                 position_wise_mlp_dim: int = 32,
+                 init_gru_gate_bias: float = 2.0):
         """Initializes a GTrXLNet instance.
 
         Args:
             num_transformer_units (int): The number of Transformer repeats to
                 use (denoted L in [2]).
-            attn_dim (int): The input and output dimensions of one Transformer
-                unit.
+            attention_dim (int): The input and output dimensions of one
+                Transformer unit.
             num_heads (int): The number of attention heads to use in parallel.
                 Denoted as `H` in [3].
             memory_inference (int): The number of timesteps to concat (time
@@ -198,23 +201,23 @@ class GTrXLNet(RecurrentNetwork):
                 input (plus the actual input sequence of len=max_seq_len).
                 The first transformer unit will receive this number of
                 past observations (plus the input sequence), instead.
-            head_dim (int): The dimension of a single(!) head.
-                Denoted as `d` in [3].
-            ff_hidden_dim (int): The dimension of the hidden layer within
-                the position-wise MLP (after the multi-head attention block
-                within one Transformer unit). This is the size of the first
-                of the two layers within the PositionwiseFeedforward. The
-                second layer always has size=`attn_dim`.
-            init_gate_bias (float): Initial bias values for the GRU gates (two
-                GRUs per Transformer unit, one after the MHA, one after the
-                position-wise MLP).
+            head_dim (int): The dimension of a single(!) attention head within
+                a multi-head attention unit. Denoted as `d` in [3].
+            position_wise_mlp_dim (int): The dimension of the hidden layer
+                within the position-wise MLP (after the multi-head attention
+                block within one Transformer unit). This is the size of the
+                first of the two layers within the PositionwiseFeedforward. The
+                second layer always has size=`attention_dim`.
+            init_gru_gate_bias (float): Initial bias values for the GRU gates
+                (two GRUs per Transformer unit, one after the MHA, one after
+                the position-wise MLP).
         """
 
         super().__init__(observation_space, action_space, num_outputs,
                          model_config, name)
 
         self.num_transformer_units = num_transformer_units
-        self.attn_dim = attn_dim
+        self.attention_dim = attention_dim
         self.num_heads = num_heads
         self.memory_inference = memory_inference
         self.memory_training = memory_training
@@ -227,14 +230,14 @@ class GTrXLNet(RecurrentNetwork):
             shape=(None, self.obs_dim), name="inputs")
         memory_ins = [
             tf.keras.layers.Input(
-                shape=(None, self.attn_dim),
+                shape=(None, self.attention_dim),
                 dtype=tf.float32,
                 name="memory_in_{}".format(i))
             for i in range(self.num_transformer_units)
         ]
 
         # Map observation dim to input/output transformer (attention) dim.
-        E_out = tf.keras.layers.Dense(self.attn_dim)(input_layer)
+        E_out = tf.keras.layers.Dense(self.attention_dim)(input_layer)
         # Output, collected and concat'd to build the internal, tau-len
         # Memory units used for additional contextual information.
         memory_outs = [E_out]
@@ -244,12 +247,12 @@ class GTrXLNet(RecurrentNetwork):
             # RelativeMultiHeadAttention part.
             MHA_out = SkipConnection(
                 RelativeMultiHeadAttention(
-                    out_dim=self.attn_dim,
+                    out_dim=self.attention_dim,
                     num_heads=num_heads,
                     head_dim=head_dim,
                     input_layernorm=True,
                     output_activation=tf.nn.relu),
-                fan_in_layer=GRUGate(init_gate_bias),
+                fan_in_layer=GRUGate(init_gru_gate_bias),
                 name="mha_{}".format(i + 1))(
                     E_out, memory=memory_ins[i])
             # Position-wise MLP part.
@@ -257,29 +260,32 @@ class GTrXLNet(RecurrentNetwork):
                 tf.keras.Sequential(
                     (tf.keras.layers.LayerNormalization(axis=-1),
                      PositionwiseFeedforward(
-                         out_dim=self.attn_dim,
-                         hidden_dim=ff_hidden_dim,
+                         out_dim=self.attention_dim,
+                         hidden_dim=position_wise_mlp_dim,
                          output_activation=tf.nn.relu))),
-                fan_in_layer=GRUGate(init_gate_bias),
+                fan_in_layer=GRUGate(init_gru_gate_bias),
                 name="pos_wise_mlp_{}".format(i + 1))(MHA_out)
             # Output of position-wise MLP == E(l-1), which is concat'd
             # to the current Mem block (M(l-1)) to yield E~(l-1), which is then
             # used by the next transformer block.
             memory_outs.append(E_out)
 
-        # Postprocess TrXL output with another hidden layer and compute values.
-        logits = tf.keras.layers.Dense(
-            self.num_outputs,
-            activation=tf.keras.activations.linear,
-            name="logits")(E_out)
-
+        self._logits = None
         self._value_out = None
-        values_out = tf.keras.layers.Dense(
-            1, activation=None, name="values")(E_out)
+
+        # Postprocess TrXL output with another hidden layer and compute values.
+        if num_outputs is not None:
+            self._logits = tf.keras.layers.Dense(
+                self.num_outputs, activation=None, name="logits")(E_out)
+            values_out = tf.keras.layers.Dense(
+                1, activation=None, name="values")(E_out)
+            outs = [self._logits, values_out]
+        else:
+            outs = [E_out]
+            self.num_outputs = self.attention_dim
 
         self.trxl_model = tf.keras.Model(
-            inputs=[input_layer] + memory_ins,
-            outputs=[logits, values_out] + memory_outs[:-1])
+            inputs=[input_layer] + memory_ins, outputs=outs + memory_outs[:-1])
 
         self.register_variables(self.trxl_model.variables)
         self.trxl_model.summary()
@@ -287,7 +293,7 @@ class GTrXLNet(RecurrentNetwork):
         # __sphinx_doc_begin__
         # Setup trajectory views (`memory-inference` x past memory outs).
         for i in range(self.num_transformer_units):
-            space = Box(-1.0, 1.0, shape=(self.attn_dim, ))
+            space = Box(-1.0, 1.0, shape=(self.attention_dim, ))
             self.view_requirements["state_in_{}".format(i)] = \
                 ViewRequirement(
                     "state_out_{}".format(i),
@@ -317,12 +323,16 @@ class GTrXLNet(RecurrentNetwork):
 
         all_out = self.trxl_model([observations] + state)
 
-        logits = all_out[0]
-        self._value_out = all_out[1]
-        memory_outs = all_out[2:]
+        if self._logits is not None:
+            out = tf.reshape(all_out[0], [-1, self.num_outputs])
+            self._value_out = all_out[1]
+            memory_outs = all_out[2:]
+        else:
+            out = tf.reshape(all_out[0], [-1, self.attention_dim])
+            memory_outs = all_out[1:]
 
-        return tf.reshape(logits, [-1, self.num_outputs]), [
-            tf.reshape(m, [-1, self.attn_dim]) for m in memory_outs
+        return out, [
+            tf.reshape(m, [-1, self.attention_dim]) for m in memory_outs
         ]
 
     # TODO: (sven) Deprecate this once trajectory view API has fully matured.
@@ -333,3 +343,90 @@ class GTrXLNet(RecurrentNetwork):
     @override(ModelV2)
     def value_function(self) -> TensorType:
         return tf.reshape(self._value_out, [-1])
+
+
+class AttentionWrapper(TFModelV2):
+    """GTrXL wrapper serving as interface for ModelV2s that set use_attention.
+    """
+
+    def __init__(self, obs_space: gym.spaces.Space,
+                 action_space: gym.spaces.Space, num_outputs: int,
+                 model_config: ModelConfigDict, name: str):
+
+        super().__init__(obs_space, action_space, None, model_config, name)
+
+        if isinstance(action_space, Discrete):
+            self.action_dim = action_space.n
+        elif isinstance(action_space, MultiDiscrete):
+            self.action_dim = np.product(action_space.nvec)
+        elif action_space.shape is not None:
+            self.action_dim = int(np.product(action_space.shape))
+        else:
+            self.action_dim = int(len(action_space))
+
+        cfg = model_config
+
+        self.attention_dim = cfg["attention_dim"]
+
+        # Construct GTrXL sub-module w/ num_outputs=None (so it does not
+        # create a logits/value output; we'll do this ourselves in this wrapper
+        # here).
+        self.gtrxl = GTrXLNet(
+            obs_space,
+            action_space,
+            None,
+            model_config,
+            "gtrxl",
+            num_transformer_units=cfg["attention_num_transformer_units"],
+            attention_dim=self.attention_dim,
+            num_heads=cfg["attention_num_heads"],
+            head_dim=cfg["attention_head_dim"],
+            memory_inference=cfg["attention_memory_inference"],
+            memory_training=cfg["attention_memory_training"],
+            position_wise_mlp_dim=cfg["attention_position_wise_mlp_dim"],
+            init_gru_gate_bias=cfg["attention_init_gru_gate_bias"],
+        )
+        self.register_variables(self.gtrxl.variables())
+
+        # `self.num_outputs` right now is the number of nodes coming from the
+        # attention net.
+        input_ = tf.keras.layers.Input(shape=(self.gtrxl.num_outputs, ))
+
+        # Set final num_outputs to correct value (depending on action space).
+        self.num_outputs = num_outputs
+
+        # Postprocess GTrXL output with another hidden layer and compute
+        # values.
+        out = tf.keras.layers.Dense(self.num_outputs, activation=None)(input_)
+        self._logits_branch = tf.keras.models.Model([input_], [out])
+        self.register_variables(self._logits_branch.variables)
+
+        out = tf.keras.layers.Dense(1, activation=None)(input_)
+        self._value_branch = tf.keras.models.Model([input_], [out])
+        self.register_variables(self._value_branch.variables)
+
+        self.view_requirements = self.gtrxl.view_requirements
+
+    @override(RecurrentNetwork)
+    def forward(self, input_dict: Dict[str, TensorType],
+                state: List[TensorType],
+                seq_lens: TensorType) -> (TensorType, List[TensorType]):
+        assert seq_lens is not None
+        # Push obs through "unwrapped" net's `forward()` first.
+        wrapped_out, _ = self._wrapped_forward(input_dict, [], None)
+
+        # Then through our GTrXL.
+        input_dict["obs_flat"] = wrapped_out
+
+        self._features, memory_outs = self.gtrxl(input_dict, state, seq_lens)
+        model_out = self._logits_branch(self._features)
+        return model_out, memory_outs
+
+    @override(ModelV2)
+    def get_initial_state(self) -> Union[List[np.ndarray], List[TensorType]]:
+        return []
+
+    @override(ModelV2)
+    def value_function(self) -> TensorType:
+        assert self._features is not None, "Must call forward() first!"
+        return tf.reshape(self._value_branch(self._features), [-1])

rllib/models/tf/layers/relative_multi_head_attention.py

@@ -22,11 +22,12 @@ class RelativeMultiHeadAttention(tf.keras.layers.Layer if tf else object):
         """Initializes a RelativeMultiHeadAttention keras Layer object.
 
         Args:
-            out_dim (int):
+            out_dim (int): The output dimensions of the multi-head attention
+                unit.
             num_heads (int): The number of attention heads to use.
                 Denoted `H` in [2].
-            head_dim (int): The dimension of a single(!) attention head
-                Denoted `D` in [2].
+            head_dim (int): The dimension of a single(!) attention head within
+                a multi-head attention unit. Denoted as `d` in [3].
             input_layernorm (bool): Whether to prepend a LayerNorm before
                 everything else. Should be True for building a GTrXL.
             output_activation (Optional[tf.nn.activation]): Optional tf.nn

rllib/models/torch/attention_net.py

@@ -8,15 +8,17 @@
       Z. Dai, Z. Yang, et al. - Carnegie Mellon U - 2019.
       https://www.aclweb.org/anthology/P19-1285.pdf
 """
-import numpy as np
 import gym
-from gym.spaces import Box
+from gym.spaces import Box, Discrete, MultiDiscrete
+import numpy as np
+from typing import Dict, Optional, Union
 
 from ray.rllib.models.modelv2 import ModelV2
 from ray.rllib.models.torch.misc import SlimFC
 from ray.rllib.models.torch.modules import GRUGate, \
     RelativeMultiHeadAttention, SkipConnection
 from ray.rllib.models.torch.recurrent_net import RecurrentNetwork
+from ray.rllib.models.torch.torch_modelv2 import TorchModelV2
 from ray.rllib.policy.sample_batch import SampleBatch
 from ray.rllib.policy.view_requirement import ViewRequirement
 from ray.rllib.utils.annotations import override
@@ -41,7 +43,7 @@ class GTrXLNet(RecurrentNetwork, nn.Module):
         >> config["model"]["max_seq_len"] = 10
         >> config["model"]["custom_model_config"] = {
         >>     num_transformer_units=1,
-        >>     attn_dim=32,
+        >>     attention_dim=32,
         >>     num_heads=2,
         >>     memory_tau=50,
         >>     etc..
@@ -51,24 +53,25 @@ class GTrXLNet(RecurrentNetwork, nn.Module):
     def __init__(self,
                  observation_space: gym.spaces.Space,
                  action_space: gym.spaces.Space,
-                 num_outputs: int,
+                 num_outputs: Optional[int],
                  model_config: ModelConfigDict,
                  name: str,
-                 num_transformer_units: int,
-                 attn_dim: int,
-                 num_heads: int,
-                 memory_inference: int,
-                 memory_training: int,
-                 head_dim: int,
-                 ff_hidden_dim: int,
-                 init_gate_bias: float = 2.0):
+                 *,
+                 num_transformer_units: int = 1,
+                 attention_dim: int = 64,
+                 num_heads: int = 2,
+                 memory_inference: int = 50,
+                 memory_training: int = 50,
+                 head_dim: int = 32,
+                 position_wise_mlp_dim: int = 32,
+                 init_gru_gate_bias: float = 2.0):
         """Initializes a GTrXLNet.
 
         Args:
             num_transformer_units (int): The number of Transformer repeats to
                 use (denoted L in [2]).
-            attn_dim (int): The input and output dimensions of one Transformer
-                unit.
+            attention_dim (int): The input and output dimensions of one
+                Transformer unit.
             num_heads (int): The number of attention heads to use in parallel.
                 Denoted as `H` in [3].
             memory_inference (int): The number of timesteps to concat (time
@@ -80,16 +83,16 @@ class GTrXLNet(RecurrentNetwork, nn.Module):
                 input (plus the actual input sequence of len=max_seq_len).
                 The first transformer unit will receive this number of
                 past observations (plus the input sequence), instead.
-            head_dim (int): The dimension of a single(!) head.
-                Denoted as `d` in [3].
-            ff_hidden_dim (int): The dimension of the hidden layer within
-                the position-wise MLP (after the multi-head attention block
-                within one Transformer unit). This is the size of the first
-                of the two layers within the PositionwiseFeedforward. The
-                second layer always has size=`attn_dim`.
-            init_gate_bias (float): Initial bias values for the GRU gates (two
-                GRUs per Transformer unit, one after the MHA, one after the
-                position-wise MLP).
+            head_dim (int): The dimension of a single(!) attention head within
+                a multi-head attention unit. Denoted as `d` in [3].
+            position_wise_mlp_dim (int): The dimension of the hidden layer
+                within the position-wise MLP (after the multi-head attention
+                block within one Transformer unit). This is the size of the
+                first of the two layers within the PositionwiseFeedforward. The
+                second layer always has size=`attention_dim`.
+            init_gru_gate_bias (float): Initial bias values for the GRU gates
+                (two GRUs per Transformer unit, one after the MHA, one after
+                the position-wise MLP).
         """
 
         super().__init__(observation_space, action_space, num_outputs,
@@ -98,7 +101,7 @@ class GTrXLNet(RecurrentNetwork, nn.Module):
         nn.Module.__init__(self)
 
         self.num_transformer_units = num_transformer_units
-        self.attn_dim = attn_dim
+        self.attention_dim = attention_dim
         self.num_heads = num_heads
         self.memory_inference = memory_inference
         self.memory_training = memory_training
@@ -107,7 +110,7 @@ class GTrXLNet(RecurrentNetwork, nn.Module):
         self.obs_dim = observation_space.shape[0]
 
         self.linear_layer = SlimFC(
-            in_size=self.obs_dim, out_size=self.attn_dim)
+            in_size=self.obs_dim, out_size=self.attention_dim)
 
         self.layers = [self.linear_layer]
 
@@ -117,29 +120,29 @@ class GTrXLNet(RecurrentNetwork, nn.Module):
             # RelativeMultiHeadAttention part.
             MHA_layer = SkipConnection(
                 RelativeMultiHeadAttention(
-                    in_dim=self.attn_dim,
-                    out_dim=self.attn_dim,
+                    in_dim=self.attention_dim,
+                    out_dim=self.attention_dim,
                     num_heads=num_heads,
                     head_dim=head_dim,
                     input_layernorm=True,
                     output_activation=nn.ReLU),
-                fan_in_layer=GRUGate(self.attn_dim, init_gate_bias))
+                fan_in_layer=GRUGate(self.attention_dim, init_gru_gate_bias))
 
             # Position-wise MultiLayerPerceptron part.
             E_layer = SkipConnection(
                 nn.Sequential(
-                    torch.nn.LayerNorm(self.attn_dim),
+                    torch.nn.LayerNorm(self.attention_dim),
                     SlimFC(
-                        in_size=self.attn_dim,
-                        out_size=ff_hidden_dim,
+                        in_size=self.attention_dim,
+                        out_size=position_wise_mlp_dim,
                         use_bias=False,
                         activation_fn=nn.ReLU),
                     SlimFC(
-                        in_size=ff_hidden_dim,
-                        out_size=self.attn_dim,
+                        in_size=position_wise_mlp_dim,
+                        out_size=self.attention_dim,
                         use_bias=False,
                         activation_fn=nn.ReLU)),
-                fan_in_layer=GRUGate(self.attn_dim, init_gate_bias))
+                fan_in_layer=GRUGate(self.attention_dim, init_gru_gate_bias))
 
             # Build a list of all attanlayers in order.
             attention_layers.extend([MHA_layer, E_layer])
@@ -149,20 +152,27 @@ class GTrXLNet(RecurrentNetwork, nn.Module):
         self.attention_layers = nn.Sequential(*attention_layers)
         self.layers.extend(attention_layers)
 
-        # Postprocess GTrXL output with another hidden layer.
-        self.logits = SlimFC(
-            in_size=self.attn_dim,
-            out_size=self.num_outputs,
-            activation_fn=nn.ReLU)
-
-        # Value function used by all RLlib Torch RL implementations.
+        # Final layers if num_outputs not None.
+        self.logits = None
+        self.values_out = None
+        # Last value output.
         self._value_out = None
-        self.values_out = SlimFC(
-            in_size=self.attn_dim, out_size=1, activation_fn=None)
+        # Postprocess GTrXL output with another hidden layer.
+        if self.num_outputs is not None:
+            self.logits = SlimFC(
+                in_size=self.attention_dim,
+                out_size=self.num_outputs,
+                activation_fn=nn.ReLU)
+
+            # Value function used by all RLlib Torch RL implementations.
+            self.values_out = SlimFC(
+                in_size=self.attention_dim, out_size=1, activation_fn=None)
+        else:
+            self.num_outputs = self.attention_dim
 
         # Setup trajectory views (`memory-inference` x past memory outs).
         for i in range(self.num_transformer_units):
-            space = Box(-1.0, 1.0, shape=(self.attn_dim, ))
+            space = Box(-1.0, 1.0, shape=(self.attention_dim, ))
             self.view_requirements["state_in_{}".format(i)] = \
                 ViewRequirement(
                     "state_out_{}".format(i),
@@ -205,11 +215,16 @@ class GTrXLNet(RecurrentNetwork, nn.Module):
         # layer).
         memory_outs = memory_outs[:-1]
 
-        logits = self.logits(all_out)
-        self._value_out = self.values_out(all_out)
+        if self.logits is not None:
+            out = self.logits(all_out)
+            self._value_out = self.values_out(all_out)
+            out_dim = self.num_outputs
+        else:
+            out = all_out
+            out_dim = self.attention_dim
 
-        return torch.reshape(logits, [-1, self.num_outputs]), [
-            torch.reshape(m, [-1, self.attn_dim]) for m in memory_outs
+        return torch.reshape(out, [-1, out_dim]), [
+            torch.reshape(m, [-1, self.attention_dim]) for m in memory_outs
         ]
 
     # TODO: (sven) Deprecate this once trajectory view API has fully matured.
@@ -219,4 +234,92 @@ class GTrXLNet(RecurrentNetwork, nn.Module):
 
     @override(ModelV2)
     def value_function(self) -> TensorType:
+        assert self._value_out is not None,\
+            "Must call forward first AND must have value branch!"
         return torch.reshape(self._value_out, [-1])
+
+
+class AttentionWrapper(TorchModelV2, nn.Module):
+    """GTrXL wrapper serving as interface for ModelV2s that set use_attention.
+    """
+
+    def __init__(self, obs_space: gym.spaces.Space,
+                 action_space: gym.spaces.Space, num_outputs: int,
+                 model_config: ModelConfigDict, name: str):
+
+        nn.Module.__init__(self)
+        super().__init__(obs_space, action_space, None, model_config, name)
+
+        if isinstance(action_space, Discrete):
+            self.action_dim = action_space.n
+        elif isinstance(action_space, MultiDiscrete):
+            self.action_dim = np.product(action_space.nvec)
+        elif action_space.shape is not None:
+            self.action_dim = int(np.product(action_space.shape))
+        else:
+            self.action_dim = int(len(action_space))
+
+        cfg = model_config
+
+        self.attention_dim = cfg["attention_dim"]
+
+        # Construct GTrXL sub-module w/ num_outputs=None (so it does not
+        # create a logits/value output; we'll do this ourselves in this wrapper
+        # here).
+        self.gtrxl = GTrXLNet(
+            obs_space,
+            action_space,
+            None,
+            model_config,
+            "gtrxl",
+            num_transformer_units=cfg["attention_num_transformer_units"],
+            attention_dim=self.attention_dim,
+            num_heads=cfg["attention_num_heads"],
+            head_dim=cfg["attention_head_dim"],
+            memory_inference=cfg["attention_memory_inference"],
+            memory_training=cfg["attention_memory_training"],
+            position_wise_mlp_dim=cfg["attention_position_wise_mlp_dim"],
+            init_gru_gate_bias=cfg["attention_init_gru_gate_bias"],
+        )
+
+        # Set final num_outputs to correct value (depending on action space).
+        self.num_outputs = num_outputs
+
+        # Postprocess GTrXL output with another hidden layer and compute
+        # values.
+        self._logits_branch = SlimFC(
+            in_size=self.attention_dim,
+            out_size=self.num_outputs,
+            activation_fn=None,
+            initializer=torch.nn.init.xavier_uniform_)
+        self._value_branch = SlimFC(
+            in_size=self.attention_dim,
+            out_size=1,
+            activation_fn=None,
+            initializer=torch.nn.init.xavier_uniform_)
+
+        self.view_requirements = self.gtrxl.view_requirements
+
+    @override(RecurrentNetwork)
+    def forward(self, input_dict: Dict[str, TensorType],
+                state: List[TensorType],
+                seq_lens: TensorType) -> (TensorType, List[TensorType]):
+        assert seq_lens is not None
+        # Push obs through "unwrapped" net's `forward()` first.
+        wrapped_out, _ = self._wrapped_forward(input_dict, [], None)
+
+        # Then through our GTrXL.
+        input_dict["obs_flat"] = wrapped_out
+
+        self._features, memory_outs = self.gtrxl(input_dict, state, seq_lens)
+        model_out = self._logits_branch(self._features)
+        return model_out, [torch.squeeze(m, 0) for m in memory_outs]
+
+    @override(ModelV2)
+    def get_initial_state(self) -> Union[List[np.ndarray], List[TensorType]]:
+        return []
+
+    @override(ModelV2)
+    def value_function(self) -> TensorType:
+        assert self._features is not None, "Must call forward() first!"
+        return torch.reshape(self._value_branch(self._features), [-1])

rllib/models/torch/recurrent_net.py

@@ -116,7 +116,8 @@ class LSTMWrapper(RecurrentNetwork, nn.Module):
                  model_config: ModelConfigDict, name: str):
 
         nn.Module.__init__(self)
-        super().__init__(obs_space, action_space, None, model_config, name)
+        super(LSTMWrapper, self).__init__(obs_space, action_space, None,
+                                          model_config, name)
 
         # At this point, self.num_outputs is the number of nodes coming
         # from the wrapped (underlying) model. In other words, self.num_outputs

rllib/tests/test_attention_net_learning.py

@@ -42,12 +42,12 @@ class TestAttentionNetLearning(unittest.TestCase):
                     "max_seq_len": 10,
                     "custom_model_config": {
                         "num_transformer_units": 1,
-                        "attn_dim": 32,
+                        "attention_dim": 32,
                         "num_heads": 1,
                         "memory_inference": 5,
                         "memory_training": 5,
                         "head_dim": 32,
-                        "ff_hidden_dim": 32,
+                        "position_wise_mlp_dim": 32,
                     },
                 },
             })
@@ -70,12 +70,12 @@ class TestAttentionNetLearning(unittest.TestCase):
         #            "max_seq_len": 65,
         #            "custom_model_config": {
         #                "num_transformer_units": 1,
-        #                "attn_dim": 64,
+        #                "attention_dim": 64,
         #                "num_heads": 1,
         #                "memory_inference": 10,
         #                "memory_training": 10,
         #                "head_dim": 32,
-        #                "ff_hidden_dim": 32,
+        #                "position_wise_mlp_dim": 32,
         #            },
         #        },
         #    })