This is the Recurrent Policy Gradient model I’m gonna use. It has two GRU layers and one dense for policy output (just like the one above).
The network takes a sequence of 5 observations (.i.e. input size (batch, 5, features)) and pads it to the desired sequence number. It takes into account episode end
For example:
At the beginning of the episode the input is only one observation of shape (batch, 1, features). The functions adds zeros before the observation. But at episode end it adds them after the last observation.
class RNNPG(nn.Module):
    r"""Recurrent plain vanila policy gradient model."""

    def __init__(self, obs_size, act_size, sequence=5, hid_size=128, num_layers=2):
        super().__init__()
        self.sequence = sequence

        self.gru = nn.GRU(obs_size, hid_size, num_layers, batch_first=True)
        self.fc = nn.Linear(hid_size*sequence, act_size)

    @staticmethod
    def pad_sequence(states, episode_end=False, sequence=5, padding_value=0.0):
        r"""Pad a sequence of observations upto sequence size by adding `padding_value` above or below the input.

        It pads above states by default unless episode end is `True` then it pads after (see example).

        Args:
            states: tensor of size (batch, sequence_0, features)

            episode_end: bool   

            sequence = 5: int. desired sequence size    

            padding_value = 0.0: float. value to pad tensor with     

        Example:
            >>> a = a = torch.ones(1,2,4)
            >>> print(a) 
            tensor([[[1., 1., 1., 1.],
                 [1., 1., 1., 1.]]])
            >>> a = pad_sequence(a,episode_end=False, sequence=5)
            >>> print(a)
            tensor([[[0., 0., 0., 0.],
                     [0., 0., 0., 0.],
                     [0., 0., 0., 0.],
                     [1., 1., 1., 1.],
                     [1., 1., 1., 1.]]])
        """
        # assuming batch is always first
        out_dims = (len(states), sequence, states.size(2))
        out_tensor = states[0].data.new(*out_dims).fill_(padding_value)
        length = states.size(1)
        if episode_end:
            out_tensor[:, :length, :] = states
        else:
            out_tensor[:, -length:, :] = states
        return out_tensor

    def forward(self, x, hx=None, episode_end=False):
        r"""Return output and pad the input if observation is less than sequence."""
        if x.size(1) < self.sequence:
            x = self.pad_sequence(x, episode_end, self.sequence)
        out, hx = self.gru(x, hx)
        out = F.relu(out.flatten(1))
        return self.fc(out), hx