Rnn support: DRQN agent + recurrent buffer

hive/agents/drqn.py

@@ -1,6 +1,7 @@
 import copy
 import os
 
+import gym
 import numpy as np
 import torch
 

hive/agents/qnets/rnn.py

@@ -5,6 +5,7 @@
 from hive.agents.qnets.mlp import MLPNetwork
 from hive.agents.qnets.conv import ConvNetwork
 from hive.agents.qnets.utils import calculate_output_dim
+from hive.agents.qnets.sequence_models import SequenceModule
 
 
 class ConvRNNNetwork(nn.Module):
@@ -23,22 +24,23 @@ class ConvRNNNetwork(nn.Module):
     def __init__(
         self,
         in_dim,
+        sequence_fn: SequenceModule,
         channels=None,
         mlp_layers=None,
         kernel_sizes=1,
         strides=1,
         paddings=0,
         normalization_factor=255,
-        rnn_type="lstm",
-        rnn_hidden_size=128,
-        num_rnn_layers=1,
         noisy=False,
         std_init=0.5,
     ):
         """
         Args:
             in_dim (tuple): The tuple of observations dimension (channels, width,
                 height).
+            sequence_fn (SequenceModule): A sequence neural network that learns
+                recurrent representation. Usually placed between the convolutional
+                layers and mlp layers.
             channels (list): The size of output channel for each convolutional layer.
             mlp_layers (list): The number of neurons for each mlp layer after the
                 convolutional layers.
@@ -48,9 +50,6 @@ def __init__(
                 layer.
             normalization_factor (float | int): What the input is divided by before
                 the forward pass of the network.
-            rnn_type (str): Type of the recurrent layer. For now, we support lstm and gru.
-            rnn_hidden_size (int):  The number of features in the hidden state h.
-            num_rnn_layers (int): Number of recurrent layers.
             noisy (bool): Whether the MLP part of the network will use
                 :py:class:`~hive.agents.qnets.noisy_linear.NoisyLinear` layers or
                 :py:class:`torch.nn.Linear` layers.
@@ -59,9 +58,6 @@ def __init__(
                 :py:class:`~hive.agents.qnets.noisy_linear.NoisyLinear`.
         """
         super().__init__()
-        self._rnn_type = rnn_type
-        self._rnn_hidden_size = rnn_hidden_size
-        self._num_rnn_layers = num_rnn_layers
         self._normalization_factor = normalization_factor
         if channels is not None:
             if isinstance(kernel_sizes, int):
@@ -96,27 +92,14 @@ def __init__(
 
         # RNN Layers
         conv_output_size = calculate_output_dim(self.conv, in_dim)
-        if self._rnn_type == "lstm":
-            self.rnn = nn.LSTM(
-                np.prod(conv_output_size),
-                rnn_hidden_size,
-                num_rnn_layers,
-                batch_first=True,
-            )
-        elif self._rnn_type == "gru":
-            self.rnn = nn.GRU(
-                np.prod(conv_output_size),
-                rnn_hidden_size,
-                num_rnn_layers,
-                batch_first=True,
-            )
-        else:
-            raise ValueError("Invalid rnn type: {}".format(self._rnn_type))
+        self.rnn = sequence_fn(
+            rnn_input_size=np.prod(conv_output_size),
+        )
 
         if mlp_layers is not None:
             # MLP Layers
             self.mlp = MLPNetwork(
-                rnn_hidden_size, mlp_layers, noisy=noisy, std_init=std_init
+                self.rnn.hidden_size, mlp_layers, noisy=noisy, std_init=std_init
             )
         else:
             self.mlp = nn.Identity()
@@ -145,24 +128,9 @@ def forward(self, x, hidden_state=None):
         return x, hidden_state
 
     def init_hidden(self, batch_size, device="cpu"):
-        if self._rnn_type == "lstm":
-            hidden_state = (
-                torch.zeros(
-                    (self._num_rnn_layers, batch_size, self._rnn_hidden_size),
-                    dtype=torch.float32,
-                    device=device,
-                ),
-                torch.zeros(
-                    (self._num_rnn_layers, batch_size, self._rnn_hidden_size),
-                    dtype=torch.float32,
-                    device=device,
-                ),
-            )
-        elif self._rnn_type == "gru":
-            hidden_state = torch.zeros(
-                (self._num_rnn_layers, batch_size, self._rnn_hidden_size),
-                dtype=torch.float32,
-                device=device,
-            )
+        hidden_state = self.rnn.init_hidden(
+            batch_size=batch_size,
+            device=device,
+        )
 
         return hidden_state

hive/agents/qnets/sequence_models.py

@@ -0,0 +1,145 @@
+import torch
+from torch import nn
+
+from hive.utils.registry import registry, Registrable
+
+
+class SequenceModule(nn.Module, Registrable):
+    """
+    Base sequence neural network architecture.
+    """
+
+    def __init__(
+        self,
+        rnn_input_size=256,
+        rnn_hidden_size=128,
+        num_rnn_layers=1,
+        batch_first=True,
+    ):
+        """
+        Args:
+            rnn_input_size (int): The number of expected features in the input x.
+            rnn_hidden_size (int):  The number of features in the hidden state h.
+            num_rnn_layers (int): Number of recurrent layers.
+            batch_first (bool): If True, then the input and output tensors are provided as (batch, seq, feature) instead of (seq, batch, feature).
+        """
+        super().__init__()
+        self._rnn_input_size = rnn_input_size
+        self._rnn_hidden_size = rnn_hidden_size
+        self._num_rnn_layers = num_rnn_layers
+        self._batch_first = batch_first
+        self.core = None
+
+    def forward(self, x, hidden_state=None):
+        x, hidden_state = self.core(x, hidden_state)
+        return x, hidden_state
+
+    @property
+    def hidden_size(self):
+        return self._rnn_hidden_size
+
+    @classmethod
+    def type_name(cls):
+        return "sequence_fn"
+
+
+class LSTMModule(SequenceModule):
+    """
+    A multi-layer long short-term memory (LSTM) RNN.
+    """
+
+    def __init__(
+        self,
+        rnn_input_size=256,
+        rnn_hidden_size=128,
+        num_rnn_layers=1,
+        batch_first=True,
+    ):
+        """
+        Args:
+            rnn_input_size (int): The number of expected features in the input x.
+            rnn_hidden_size (int):  The number of features in the hidden state h.
+            num_rnn_layers (int): Number of recurrent layers.
+            batch_first (bool): If True, then the input and output tensors are provided as (batch, seq, feature) instead of (seq, batch, feature).
+        """
+        super().__init__(
+            rnn_input_size=rnn_input_size,
+            rnn_hidden_size=rnn_hidden_size,
+            num_rnn_layers=num_rnn_layers,
+            batch_first=batch_first,
+        )
+        self.core = nn.LSTM(
+            input_size=self._rnn_input_size,
+            hidden_size=self._rnn_hidden_size,
+            num_layers=self._num_rnn_layers,
+            batch_first=self._batch_first,
+        )
+
+    def init_hidden(self, batch_size, device="cpu"):
+        hidden_state = (
+            torch.zeros(
+                (self._num_rnn_layers, batch_size, self._rnn_hidden_size),
+                dtype=torch.float32,
+                device=device,
+            ),
+            torch.zeros(
+                (self._num_rnn_layers, batch_size, self._rnn_hidden_size),
+                dtype=torch.float32,
+                device=device,
+            ),
+        )
+
+        return hidden_state
+
+
+class GRUModule(SequenceModule):
+    """
+    A multi-layer gated recurrent unit (GRU) RNN.
+    """
+
+    def __init__(
+        self,
+        rnn_input_size=256,
+        rnn_hidden_size=128,
+        num_rnn_layers=1,
+        batch_first=True,
+    ):
+        """
+        Args:
+            rnn_input_size (int): The number of expected features in the input x.
+            rnn_hidden_size (int):  The number of features in the hidden state h.
+            num_rnn_layers (int): Number of recurrent layers.
+            batch_first (bool): If True, then the input and output tensors are provided as (batch, seq, feature) instead of (seq, batch, feature).
+        """
+        super().__init__(
+            rnn_input_size=rnn_input_size,
+            rnn_hidden_size=rnn_hidden_size,
+            num_rnn_layers=num_rnn_layers,
+            batch_first=batch_first,
+        )
+        self.core = nn.GRU(
+            input_size=self._rnn_input_size,
+            hidden_size=self._rnn_hidden_size,
+            num_layers=self._num_rnn_layers,
+            batch_first=self._batch_first,
+        )
+
+    def init_hidden(self, batch_size, device="cpu"):
+        hidden_state = torch.zeros(
+            (self._num_rnn_layers, batch_size, self._rnn_hidden_size),
+            dtype=torch.float32,
+            device=device,
+        )
+
+        return hidden_state
+
+
+registry.register_all(
+    SequenceModule,
+    {
+        "LSTM": LSTMModule,
+        "GRU": GRUModule,
+    },
+)
+
+get_sequence_fn = getattr(registry, f"get_{SequenceModule.type_name()}")

hive/configs/atari/drqn.yml

@@ -28,6 +28,11 @@ agent:
         strides: [4, 2, 1]
         paddings: [2, 2, 1]
         mlp_layers: [512]
+        sequence_fn:
+          name: 'LSTM'
+          kwargs:
+            rnn_hidden_size: 128
+            num_rnn_layers: 1
     optimizer_fn:
       name: 'RMSpropTF'
       kwargs: