Pytorch AttentionNet (#9088)

rllib/models/tests/test_torch_modules.py

@@ -0,0 +1,262 @@
+import gym
+import numpy as np
+import unittest
+
+from ray.rllib.models.tf.attention_net import relative_position_embedding, \
+    GTrXLNet
+from ray.rllib.models.tf.layers import MultiHeadAttention
+from ray.rllib.models.torch.attention_net import relative_position_embedding \
+    as relative_position_embedding_torch, GTrXLNet as TorchGTrXLNet
+from ray.rllib.models.torch.modules.multi_head_attention import \
+    MultiHeadAttention as TorchMultiHeadAttention
+from ray.rllib.utils.framework import try_import_torch, try_import_tf
+from ray.rllib.utils.test_utils import framework_iterator
+
+torch, nn = try_import_torch()
+tf = try_import_tf()
+
+
+class TestModules(unittest.TestCase):
+    """Tests various torch/modules and tf/layers required for AttentionNet"""
+
+    def train_torch_full_model(self,
+                               model,
+                               inputs,
+                               outputs,
+                               num_epochs=250,
+                               state=None,
+                               seq_lens=None):
+        """Convenience method that trains a Torch model for num_epochs epochs
+            and tests whether loss decreased, as expected.
+
+        Args:
+            model (nn.Module): Torch model to be trained.
+            inputs (torch.Tensor): Training data
+            outputs (torch.Tensor): Training labels
+            num_epochs (int): Number of epochs to train for
+            state (torch.Tensor): Internal state of module
+            seq_lens (torch.Tensor): Tensor of sequence lengths
+        """
+
+        criterion = torch.nn.MSELoss(reduction="sum")
+        optimizer = torch.optim.Adam(model.parameters(), lr=3e-4)
+
+        # Check that the layer trains correctly
+        for t in range(num_epochs):
+            y_pred = model(inputs, state, seq_lens)
+            loss = criterion(y_pred[0], torch.squeeze(outputs[0]))
+
+            if t % 10 == 1:
+                print(t, loss.item())
+
+            if t == 1:
+                init_loss = loss.item()
+
+            optimizer.zero_grad()
+            loss.backward()
+            optimizer.step()
+
+        final_loss = loss.item()
+
+        # The final loss has decreased, which tests
+        # that the model is learning from the training data.
+        self.assertLess(final_loss / init_loss, 0.99)
+
+    def train_torch_layer(self, model, inputs, outputs, num_epochs=250):
+        """Convenience method that trains a Torch model for num_epochs epochs
+            and tests whether loss decreased, as expected.
+
+        Args:
+            model (nn.Module): Torch model to be trained.
+            inputs (torch.Tensor): Training data
+            outputs (torch.Tensor): Training labels
+            num_epochs (int): Number of epochs to train for
+        """
+        criterion = torch.nn.MSELoss(reduction="sum")
+        optimizer = torch.optim.SGD(model.parameters(), lr=1e-4)
+
+        # Check that the layer trains correctly
+        for t in range(num_epochs):
+            y_pred = model(inputs)
+            loss = criterion(y_pred, outputs)
+
+            if t == 1:
+                init_loss = loss.item()
+
+            optimizer.zero_grad()
+            loss.backward()
+            optimizer.step()
+
+        final_loss = loss.item()
+
+        # The final loss has decreased by a factor of 2, which tests
+        # that the model is learning from the training data.
+        self.assertLess(final_loss / init_loss, 0.5)
+
+    def train_tf_model(self,
+                       model,
+                       inputs,
+                       outputs,
+                       num_epochs=250,
+                       minibatch_size=32):
+        """Convenience method that trains a Tensorflow model for num_epochs
+            epochs and tests whether loss decreased, as expected.
+
+        Args:
+            model (tf.Model): Torch model to be trained.
+            inputs (np.array): Training data
+            outputs (np.array): Training labels
+            num_epochs (int): Number of training epochs
+            batch_size (int): Number of samples in each minibatch
+        """
+
+        # Configure a model for mean-squared error loss.
+        model.compile(optimizer="SGD", loss="mse", metrics=["mae"])
+
+        hist = model.fit(
+            inputs,
+            outputs,
+            verbose=0,
+            epochs=num_epochs,
+            batch_size=minibatch_size).history
+        init_loss = hist["loss"][0]
+        final_loss = hist["loss"][-1]
+
+        self.assertLess(final_loss / init_loss, 0.5)
+
+    def test_multi_head_attention(self):
+        """Tests the MultiHeadAttention mechanism of Vaswani et al."""
+        # B is batch size
+        B = 1
+        # D_in is attention dim, L is memory_tau
+        L, D_in, D_out = 2, 32, 10
+
+        for fw, sess in framework_iterator(
+                frameworks=("tfe", "torch", "tf"), session=True):
+
+            # Create a single attention layer with 2 heads
+            if fw == "torch":
+
+                # Create random Tensors to hold inputs and outputs
+                x = torch.randn(B, L, D_in)
+                y = torch.randn(B, L, D_out)
+
+                model = TorchMultiHeadAttention(
+                    in_dim=D_in, out_dim=D_out, num_heads=2, head_dim=32)
+
+                self.train_torch_layer(model, x, y)
+
+            else:  # framework is tensorflow or tensorflow-eager
+
+                x = np.random.random((B, L, D_in))
+                y = np.random.random((B, L, D_out))
+
+                inputs = tf.keras.layers.Input(shape=(L, D_in))
+
+                model = tf.keras.Sequential([
+                    inputs,
+                    MultiHeadAttention(
+                        out_dim=D_out, num_heads=2, head_dim=32)
+                ])
+                self.train_tf_model(model, x, y)
+
+    def test_attention_net(self):
+        """Tests the GTrXL. Builds a full AttentionNet and checks
+            that it trains in a supervised setting."""
+
+        # Checks that torch and tf embedding matrices are the same
+        with tf.Session().as_default() as sess:
+            assert np.allclose(
+                relative_position_embedding(20, 15).eval(session=sess),
+                relative_position_embedding_torch(20, 15).numpy())
+
+        # B is batch size
+        B = 32
+        # D_in is attention dim, L is memory_tau
+        L, D_in, D_out = 2, 16, 2
+
+        for fw, sess in framework_iterator(
+                frameworks=("tfe", "torch", "tf"), session=True):
+
+            # Create a single attention layer with 2 heads
+            if fw == "torch":
+                # Create random Tensors to hold inputs and outputs
+                x = torch.randn(B, L, D_in)
+                y = torch.randn(B, L, D_out)
+
+                value_labels = torch.randn(B, L, D_in)
+                memory_labels = torch.randn(B, L, D_out)
+
+                attention_net = TorchGTrXLNet(
+                    observation_space=gym.spaces.Box(
+                        low=float("-inf"), high=float("inf"), shape=(D_in, )),
+                    action_space=gym.spaces.Discrete(D_out),
+                    num_outputs=D_out,
+                    model_config={"max_seq_len": 2},
+                    name="TestTorchAttentionNet",
+                    num_transformer_units=2,
+                    attn_dim=D_in,
+                    num_heads=2,
+                    memory_tau=L,
+                    head_dim=D_out,
+                    ff_hidden_dim=16,
+                    init_gate_bias=2.0)
+
+                init_state = attention_net.get_initial_state()
+
+                # Get initial state and add a batch dimension.
+                init_state = [np.expand_dims(s, 0) for s in init_state]
+                seq_lens_init = torch.full(size=(B, ), fill_value=L)
+
+                # Torch implementation expects a formatted input_dict instead
+                # of a numpy array as input.
+                input_dict = {"obs": x}
+                self.train_torch_full_model(
+                    attention_net,
+                    input_dict, [y, value_labels, memory_labels],
+                    num_epochs=250,
+                    state=init_state,
+                    seq_lens=seq_lens_init)
+
+            else:  # Framework is tensorflow or tensorflow-eager.
+                x = np.random.random((B, L, D_in))
+                y = np.random.random((B, L, D_out))
+
+                value_labels = np.random.random((B, L, 1))
+                memory_labels = np.random.random((B, L, D_in))
+
+                # We need to create (N-1) MLP labels for N transformer units
+                mlp_labels = np.random.random((B, L, D_in))
+
+                attention_net = GTrXLNet(
+                    observation_space=gym.spaces.Box(
+                        low=float("-inf"), high=float("inf"), shape=(D_in, )),
+                    action_space=gym.spaces.Discrete(D_out),
+                    num_outputs=D_out,
+                    model_config={"max_seq_len": 2},
+                    name="TestTFAttentionNet",
+                    num_transformer_units=2,
+                    attn_dim=D_in,
+                    num_heads=2,
+                    memory_tau=L,
+                    head_dim=D_out,
+                    ff_hidden_dim=16,
+                    init_gate_bias=2.0)
+                model = attention_net.trxl_model
+
+                # Get initial state and add a batch dimension.
+                init_state = attention_net.get_initial_state()
+                init_state = [np.tile(s, (B, 1, 1)) for s in init_state]
+
+                self.train_tf_model(
+                    model, [x] + init_state,
+                    [y, value_labels, memory_labels, mlp_labels],
+                    num_epochs=50,
+                    minibatch_size=B)
+
+
+if __name__ == "__main__":
+    import pytest
+    import sys
+
+    sys.exit(pytest.main(["-v", __file__]))

rllib/models/tf/layers/__init__.py

@@ -2,5 +2,9 @@ from ray.rllib.models.tf.layers.gru_gate import GRUGate
 from ray.rllib.models.tf.layers.relative_multi_head_attention import \
     RelativeMultiHeadAttention
 from ray.rllib.models.tf.layers.skip_connection import SkipConnection
+from ray.rllib.models.tf.layers.multi_head_attention import MultiHeadAttention
 
-__all__ = ["GRUGate", "RelativeMultiHeadAttention", "SkipConnection"]
+__all__ = [
+    "GRUGate", "RelativeMultiHeadAttention", "SkipConnection",
+    "MultiHeadAttention"
+]

rllib/models/tf/layers/multi_head_attention.py

@@ -47,5 +47,6 @@ class MultiHeadAttention(tf.keras.layers.Layer):
         wmat = tf.nn.softmax(masked_score, axis=2)
 
         out = tf.einsum("bijh,bjhd->bihd", wmat, values)
-        out = tf.reshape(out, tf.concat((tf.shape(out)[:2], [H * D]), axis=0))
+        shape = tf.concat([tf.shape(out)[:2], [H * D]], axis=0)
+        out = tf.reshape(out, shape)
         return self._linear_layer(out)

rllib/models/tf/layers/relative_multi_head_attention.py

@@ -113,7 +113,7 @@ class RelativeMultiHeadAttention(tf.keras.layers.Layer):
 
         x = tf.pad(x, [[0, 0], [0, 0], [1, 0], [0, 0]])
         x = tf.reshape(x, [x_size[0], x_size[2] + 1, x_size[1], x_size[3]])
-        x = tf.slice(x, [0, 1, 0, 0], [-1, -1, -1, -1])
+        x = x[:, 1:, :, :]
         x = tf.reshape(x, x_size)
 
         return x

rllib/models/torch/attention_net.py

@@ -0,0 +1,227 @@
+"""
+[1] - Attention Is All You Need - Vaswani, Jones, Shazeer, Parmar,
+      Uszkoreit, Gomez, Kaiser - Google Brain/Research, U Toronto - 2017.
+      https://arxiv.org/pdf/1706.03762.pdf
+[2] - Stabilizing Transformers for Reinforcement Learning - E. Parisotto
+      et al. - DeepMind - 2019. https://arxiv.org/pdf/1910.06764.pdf
+[3] - Transformer-XL: Attentive Language Models Beyond a Fixed-Length Context.
+      Z. Dai, Z. Yang, et al. - Carnegie Mellon U - 2019.
+      https://www.aclweb.org/anthology/P19-1285.pdf
+"""
+import numpy as np
+
+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.utils.annotations import override
+from ray.rllib.utils.framework import try_import_torch
+
+torch, nn = try_import_torch()
+
+
+def relative_position_embedding(seq_length, out_dim):
+    """Creates a [seq_length x seq_length] matrix for rel. pos encoding.
+
+    Denoted as Phi in [2] and [3]. Phi is the standard sinusoid encoding
+    matrix.
+
+    Args:
+        seq_length (int): The max. sequence length (time axis).
+        out_dim (int): The number of nodes to go into the first Tranformer
+            layer with.
+
+    Returns:
+        torch.Tensor: The encoding matrix Phi.
+    """
+    inverse_freq = 1 / (10000**(torch.arange(0, out_dim, 2.0) / out_dim))
+    pos_offsets = torch.arange(seq_length - 1, -1, -1)
+    inputs = pos_offsets[:, None] * inverse_freq[None, :]
+    return torch.cat((torch.sin(inputs), torch.cos(inputs)), dim=-1)
+
+
+class GTrXLNet(RecurrentNetwork, nn.Module):
+    """A GTrXL net Model described in [2].
+
+    This is still in an experimental phase.
+    Can be used as a drop-in replacement for LSTMs in PPO and IMPALA.
+    For an example script, see: `ray/rllib/examples/attention_net.py`.
+
+    To use this network as a replacement for an RNN, configure your Trainer
+    as follows:
+
+    Examples:
+        >> config["model"]["custom_model"] = GTrXLNet
+        >> config["model"]["max_seq_len"] = 10
+        >> config["model"]["custom_model_config"] = {
+        >>     num_transformer_units=1,
+        >>     attn_dim=32,
+        >>     num_heads=2,
+        >>     memory_tau=50,
+        >>     etc..
+        >> }
+    """
+
+    def __init__(self,
+                 observation_space,
+                 action_space,
+                 num_outputs,
+                 model_config,
+                 name,
+                 num_transformer_units,
+                 attn_dim,
+                 num_heads,
+                 memory_tau,
+                 head_dim,
+                 ff_hidden_dim,
+                 init_gate_bias=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.
+            num_heads (int): The number of attention heads to use in parallel.
+                Denoted as `H` in [3].
+            memory_tau (int): The number of timesteps to store in each
+                transformer block's memory M (concat'd over time and fed into
+                next transformer block as input).
+            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).
+        """
+
+        super().__init__(observation_space, action_space, num_outputs,
+                         model_config, name)
+
+        nn.Module.__init__(self)
+
+        self.num_transformer_units = num_transformer_units
+        self.attn_dim = attn_dim
+        self.num_heads = num_heads
+        self.memory_tau = memory_tau
+        self.head_dim = head_dim
+        self.max_seq_len = model_config["max_seq_len"]
+        self.obs_dim = observation_space.shape[0]
+
+        # Constant (non-trainable) sinusoid rel pos encoding matrix.
+        Phi = relative_position_embedding(self.max_seq_len + self.memory_tau,
+                                          self.attn_dim)
+
+        self.linear_layer = SlimFC(
+            in_size=self.obs_dim, out_size=self.attn_dim)
+
+        self.layers = [self.linear_layer]
+
+        # 2) Create L Transformer blocks according to [2].
+        for i in range(self.num_transformer_units):
+            # RelativeMultiHeadAttention part.
+            MHA_layer = SkipConnection(
+                RelativeMultiHeadAttention(
+                    in_dim=self.attn_dim,
+                    out_dim=self.attn_dim,
+                    num_heads=num_heads,
+                    head_dim=head_dim,
+                    rel_pos_encoder=Phi,
+                    input_layernorm=True,
+                    output_activation=nn.ReLU),
+                fan_in_layer=GRUGate(self.attn_dim, init_gate_bias))
+
+            # Position-wise MultiLayerPerceptron part.
+            E_layer = SkipConnection(
+                nn.Sequential(
+                    torch.nn.LayerNorm(self.attn_dim),
+                    SlimFC(
+                        in_size=self.attn_dim,
+                        out_size=ff_hidden_dim,
+                        use_bias=False,
+                        activation_fn=nn.ReLU),
+                    SlimFC(
+                        in_size=ff_hidden_dim,
+                        out_size=self.attn_dim,
+                        use_bias=False,
+                        activation_fn=nn.ReLU)),
+                fan_in_layer=GRUGate(self.attn_dim, init_gate_bias))
+
+            # Build a list of all layers in order.
+            self.layers.extend([MHA_layer, E_layer])
+
+        # 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.
+        self._value_out = None
+        self.values_out = SlimFC(
+            in_size=self.attn_dim, out_size=1, activation_fn=None)
+
+    @override(RecurrentNetwork)
+    def forward_rnn(self, inputs, state, seq_lens):
+        # To make Attention work with current RLlib's ModelV2 API:
+        # We assume `state` is the history of L recent observations (all
+        # concatenated into one tensor) and append the current inputs to the
+        # end and only keep the most recent (up to `max_seq_len`). This allows
+        # us to deal with timestep-wise inference and full sequence training
+        # within the same logic.
+        state = [torch.from_numpy(item) for item in state]
+        observations = state[0]
+        memory = state[1:]
+
+        inputs = torch.reshape(inputs, [1, -1, observations.shape[-1]])
+        observations = torch.cat(
+            (observations, inputs), axis=1)[:, -self.max_seq_len:]
+
+        all_out = observations
+        for i in range(len(self.layers)):
+            # MHA layers which need memory passed in.
+            if i % 2 == 1:
+                all_out = self.layers[i](all_out, memory=memory[i // 2])
+            # Either linear layers or MultiLayerPerceptrons.
+            else:
+                all_out = self.layers[i](all_out)
+
+        logits = self.logits(all_out)
+        self._value_out = self.values_out(all_out)
+
+        memory_outs = all_out[2:]
+        # If memory_tau > max_seq_len -> overlap w/ previous `memory` input.
+        if self.memory_tau > self.max_seq_len:
+            memory_outs = [
+                torch.cat(
+                    [memory[i][:, -(self.memory_tau - self.max_seq_len):], m],
+                    axis=1) for i, m in enumerate(memory_outs)
+            ]
+        else:
+            memory_outs = [m[:, -self.memory_tau:] for m in memory_outs]
+
+        T = list(inputs.size())[1]  # Length of input segment (time).
+
+        # Postprocessing final output.
+        logits = logits[:, -T:]
+        self._value_out = self._value_out[:, -T:]
+
+        return logits, [observations] + memory_outs
+
+    @override(RecurrentNetwork)
+    def get_initial_state(self):
+        # State is the T last observations concat'd together into one Tensor.
+        # Plus all Transformer blocks' E(l) outputs concat'd together (up to
+        # tau timesteps).
+        return [np.zeros((self.max_seq_len, self.obs_dim), np.float32)] + \
+               [np.zeros((self.memory_tau, self.attn_dim), np.float32)
+                for _ in range(self.num_transformer_units)]
+
+    @override(ModelV2)
+    def value_function(self):
+        return torch.reshape(self._value_out, [-1])

rllib/models/torch/misc.py

@@ -92,13 +92,15 @@ class SlimFC(nn.Module):
                  out_size,
                  initializer=None,
                  activation_fn=None,
+                 use_bias=True,
                  bias_init=0.0):
         super(SlimFC, self).__init__()
         layers = []
-        linear = nn.Linear(in_size, out_size)
+        linear = nn.Linear(in_size, out_size, bias=use_bias)
         if initializer:
             initializer(linear.weight)
-        nn.init.constant_(linear.bias, bias_init)
+        if use_bias is True:
+            nn.init.constant_(linear.bias, bias_init)
         layers.append(linear)
         if activation_fn:
             layers.append(activation_fn())

rllib/models/torch/modules/__init__.py

@@ -0,0 +1,11 @@
+from ray.rllib.models.torch.modules.gru_gate import GRUGate
+from ray.rllib.models.torch.modules.multi_head_attention import \
+    MultiHeadAttention
+from ray.rllib.models.torch.modules.relative_multi_head_attention import \
+    RelativeMultiHeadAttention
+from ray.rllib.models.torch.modules.skip_connection import SkipConnection
+
+__all__ = [
+    "GRUGate", "RelativeMultiHeadAttention", "SkipConnection",
+    "MultiHeadAttention"
+]

rllib/models/torch/modules/gru_gate.py

@@ -0,0 +1,52 @@
+from ray.rllib.utils.framework import try_import_torch
+
+torch, nn = try_import_torch()
+
+
+class GRUGate(nn.Module):
+    """Implements a gated recurrent unit for use in AttentionNet"""
+
+    def __init__(self, dim, init_bias=0., **kwargs):
+        """
+        input_shape (torch.Tensor): dimension of the input
+        init_bias (int): Bias added to every input to stabilize training
+        """
+        super().__init__(**kwargs)
+        self._init_bias = init_bias
+
+        # Xavier initialization of torch tensors
+        self._w_r = torch.zeros(dim, dim)
+        self._w_z = torch.zeros(dim, dim)
+        self._w_h = torch.zeros(dim, dim)
+
+        self._u_r = torch.zeros(dim, dim)
+        self._u_z = torch.zeros(dim, dim)
+        self._u_h = torch.zeros(dim, dim)
+
+        nn.init.xavier_uniform_(self._w_r)
+        nn.init.xavier_uniform_(self._w_z)
+        nn.init.xavier_uniform_(self._w_h)
+
+        nn.init.xavier_uniform_(self._u_r)
+        nn.init.xavier_uniform_(self._u_z)
+        nn.init.xavier_uniform_(self._u_h)
+
+        self._bias_z = torch.zeros(dim, ).fill_(self._init_bias)
+
+    def forward(self, inputs, **kwargs):
+        # Pass in internal state first.
+        h, X = inputs
+
+        r = torch.tensordot(X, self._w_r, dims=1) + \
+            torch.tensordot(h, self._u_r, dims=1)
+        r = torch.sigmoid(r)
+
+        z = torch.tensordot(X, self._w_z, dims=1) + \
+            torch.tensordot(h, self._u_z, dims=1) - self._bias_z
+        z = torch.sigmoid(z)
+
+        h_next = torch.tensordot(X, self._w_h, dims=1) + \
+            torch.tensordot((h * r), self._u_h, dims=1)
+        h_next = torch.tanh(h_next)
+
+        return (1 - z) * h + z * h_next

rllib/models/torch/modules/multi_head_attention.py

@@ -0,0 +1,63 @@
+"""
+[1] - Attention Is All You Need - Vaswani, Jones, Shazeer, Parmar,
+      Uszkoreit, Gomez, Kaiser - Google Brain/Research, U Toronto - 2017.
+      https://arxiv.org/pdf/1706.03762.pdf
+"""
+from ray.rllib.utils.framework import try_import_torch
+from ray.rllib.models.torch.misc import SlimFC
+from ray.rllib.utils.torch_ops import sequence_mask
+
+torch, nn = try_import_torch()
+
+
+class MultiHeadAttention(nn.Module):
+    """A multi-head attention layer described in [1]."""
+
+    def __init__(self, in_dim, out_dim, num_heads, head_dim, **kwargs):
+        """
+        in_dim (int): Dimension of input
+        out_dim (int): Dimension of output
+        num_heads (int): Number of attention heads
+        head_dim (int): Output dimension of each attention head
+        """
+        super().__init__(**kwargs)
+
+        # No bias or non-linearity.
+        self._num_heads = num_heads
+        self._head_dim = head_dim
+        self._qkv_layer = SlimFC(
+            in_size=in_dim, out_size=3 * num_heads * head_dim, use_bias=False)
+
+        self._linear_layer = SlimFC(
+            in_size=num_heads * head_dim, out_size=out_dim, use_bias=False)
+
+    def forward(self, inputs):
+        L = list(inputs.size())[1]  # length of segment
+        H = self._num_heads  # number of attention heads
+        D = self._head_dim  # attention head dimension
+
+        qkv = self._qkv_layer(inputs)
+
+        queries, keys, values = torch.chunk(input=qkv, chunks=3, dim=-1)
+        queries = queries[:, -L:]  # only query based on the segment
+
+        queries = torch.reshape(queries, [-1, L, H, D])
+        keys = torch.reshape(keys, [-1, L, H, D])
+        values = torch.reshape(values, [-1, L, H, D])
+
+        score = torch.einsum("bihd,bjhd->bijh", queries, keys)
+        score = score / D**0.5
+
+        # causal mask of the same length as the sequence
+        mask = sequence_mask(torch.arange(1, L + 1), dtype=score.dtype)
+        mask = mask[None, :, :, None]
+        mask = mask.float()
+
+        masked_score = score * mask + 1e30 * (mask - 1.)
+        wmat = nn.functional.softmax(masked_score, dim=2)
+
+        out = torch.einsum("bijh,bjhd->bihd", wmat, values)
+        shape = list(out.size())[:2] + [H * D]
+        #        temp = torch.cat(temp2, [H * D], dim=0)
+        out = torch.reshape(out, shape)
+        return self._linear_layer(out)

rllib/models/torch/modules/relative_multi_head_attention.py

@@ -0,0 +1,133 @@
+from ray.rllib.utils.framework import try_import_torch
+from ray.rllib.models.torch.misc import SlimFC
+from ray.rllib.utils.torch_ops import sequence_mask
+
+torch, nn = try_import_torch()
+
+
+class RelativeMultiHeadAttention(nn.Module):
+    """A RelativeMultiHeadAttention layer as described in [3].
+
+    Uses segment level recurrence with state reuse.
+    """
+
+    def __init__(self,
+                 in_dim,
+                 out_dim,
+                 num_heads,
+                 head_dim,
+                 rel_pos_encoder,
+                 input_layernorm=False,
+                 output_activation=None,
+                 **kwargs):
+        """Initializes a RelativeMultiHeadAttention nn.Module object.
+
+        Args:
+            in_dim (int):
+            out_dim (int):
+            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].
+            rel_pos_encoder (:
+            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
+                activation function. Should be relu for GTrXL.
+            **kwargs:
+        """
+        super().__init__(**kwargs)
+
+        # No bias or non-linearity.
+        self._num_heads = num_heads
+        self._head_dim = head_dim
+
+        # 3=Query, key, and value inputs.
+        self._qkv_layer = SlimFC(
+            in_size=in_dim, out_size=3 * num_heads * head_dim, use_bias=False)
+
+        self._linear_layer = SlimFC(
+            in_size=num_heads * head_dim,
+            out_size=out_dim,
+            use_bias=False,
+            activation_fn=output_activation)
+
+        self._pos_proj = SlimFC(
+            in_size=in_dim, out_size=num_heads * head_dim, use_bias=False)
+
+        self._uvar = torch.zeros(num_heads, head_dim)
+        self._vvar = torch.zeros(num_heads, head_dim)
+        nn.init.xavier_uniform_(self._uvar)
+        nn.init.xavier_uniform_(self._vvar)
+
+        self._rel_pos_encoder = rel_pos_encoder
+        self._input_layernorm = None
+
+        if input_layernorm:
+            self._input_layernorm = torch.nn.LayerNorm(in_dim)
+
+    def forward(self, inputs, memory=None):
+        T = list(inputs.size())[1]  # length of segment (time)
+        H = self._num_heads  # number of attention heads
+        d = self._head_dim  # attention head dimension
+
+        # Add previous memory chunk (as const, w/o gradient) to input.
+        # Tau (number of (prev) time slices in each memory chunk).
+        Tau = list(memory.shape)[1] if memory is not None else 0
+        if memory is not None:
+            memory.requires_grad_(False)
+            inputs = torch.cat((memory, inputs), dim=1)
+
+        # Apply the Layer-Norm.
+        if self._input_layernorm is not None:
+            inputs = self._input_layernorm(inputs)
+
+        qkv = self._qkv_layer(inputs)
+
+        queries, keys, values = torch.chunk(input=qkv, chunks=3, dim=-1)
+        # Cut out Tau memory timesteps from query.
+        queries = queries[:, -T:]
+
+        queries = torch.reshape(queries, [-1, T, H, d])
+        keys = torch.reshape(keys, [-1, T + Tau, H, d])
+        values = torch.reshape(values, [-1, T + Tau, H, d])
+
+        R = self._pos_proj(self._rel_pos_encoder)
+        R = torch.reshape(R, [T + Tau, H, d])
+
+        # b=batch
+        # i and j=time indices (i=max-timesteps (inputs); j=Tau memory space)
+        # h=head
+        # d=head-dim (over which we will reduce-sum)
+        score = torch.einsum("bihd,bjhd->bijh", queries + self._uvar, keys)
+        pos_score = torch.einsum("bihd,jhd->bijh", queries + self._vvar, R)
+        score = score + self.rel_shift(pos_score)
+        score = score / d**0.5
+
+        # causal mask of the same length as the sequence
+        mask = sequence_mask(
+            torch.arange(Tau + 1, T + Tau + 1), dtype=score.dtype)
+        mask = mask[None, :, :, None]
+
+        masked_score = score * mask + 1e30 * (mask.to(torch.float32) - 1.)
+        wmat = nn.functional.softmax(masked_score, dim=2)
+
+        out = torch.einsum("bijh,bjhd->bihd", wmat, values)
+        shape = list(out.shape)[:2] + [H * d]
+        out = torch.reshape(out, shape)
+
+        return self._linear_layer(out)
+
+    @staticmethod
+    def rel_shift(x):
+        # Transposed version of the shift approach described in [3].
+        # https://github.com/kimiyoung/transformer-xl/blob/
+        # 44781ed21dbaec88b280f74d9ae2877f52b492a5/tf/model.py#L31
+        x_size = list(x.shape)
+
+        x = torch.nn.functional.pad(x, (0, 0, 1, 0, 0, 0, 0, 0))
+        x = torch.reshape(x, [x_size[0], x_size[2] + 1, x_size[1], x_size[3]])
+        x = x[:, 1:, :, :]
+        x = torch.reshape(x, x_size)
+
+        return x

rllib/models/torch/modules/skip_connection.py

@@ -0,0 +1,37 @@
+from ray.rllib.utils.framework import try_import_torch
+
+torch, nn = try_import_torch()
+
+
+class SkipConnection(nn.Module):
+    """Skip connection layer.
+
+    Adds the original input to the output (regular residual layer) OR uses
+    input as hidden state input to a given fan_in_layer.
+    """
+
+    def __init__(self, layer, fan_in_layer=None, add_memory=False, **kwargs):
+        """Initializes a SkipConnection nn Module object.
+
+        Args:
+            layer (nn.Module): Any layer processing inputs.
+            fan_in_layer (Optional[nn.Module]): An optional
+                layer taking two inputs: The original input and the output
+                of `layer`.
+        """
+        super().__init__(**kwargs)
+        self._layer = layer
+        self._fan_in_layer = fan_in_layer
+
+    def forward(self, inputs, **kwargs):
+        # del kwargs
+        outputs = self._layer(inputs, **kwargs)
+        # Residual case, just add inputs to outputs.
+        if self._fan_in_layer is None:
+            outputs = outputs + inputs
+        # Fan-in e.g. RNN: Call fan-in with `inputs` and `outputs`.
+        else:
+            # NOTE: In the GRU case, `inputs` is the state input.
+            outputs = self._fan_in_layer((inputs, outputs))
+
+        return outputs

rllib/utils/torch_ops.py

@@ -77,7 +77,7 @@ def sequence_mask(lengths, maxlen=None, dtype=None):
     39036).
     """
     if maxlen is None:
-        maxlen = lengths.max()
+        maxlen = int(lengths.max())
 
     mask = ~(torch.ones((len(lengths), maxlen)).to(
         lengths.device).cumsum(dim=1).t() > lengths).t()