Refactor for HRA

network/shallow_network.py

@@ -61,4 +61,55 @@ class ActorCriticFCNet(nn.Module, ActorCriticNet):
         phi = self.fc2(x)
         return phi
 
+class FruitHRFCNet(nn.Module, VanillaNet):
+    def __init__(self, state_dim, action_dim, head_weights, optimizer_fn=None, gpu=True):
+        super(FruitHRFCNet, self).__init__()
+        hidden_size = 250
+        self.fc1 = nn.Linear(state_dim, hidden_size)
+        self.fc2 = nn.ModuleList([nn.Linear(hidden_size, action_dim) for _ in head_weights])
+        self.criterion = nn.MSELoss()
+        self.head_weights = head_weights
+        BasicNet.__init__(self, optimizer_fn, gpu)
 
+    def forward(self, x, heads_only):
+        x = self.to_torch_variable(x)
+        x = x.view(x.size(0), -1)
+        x = F.relu(self.fc1(x))
+        head_q = [fc(x) for fc in self.fc2]
+        if not heads_only:
+            q = [h * w for h, w in zip(head_q, self.head_weights)]
+            q = torch.stack(q, dim=0)
+            q = q.sum(0).squeeze(0)
+            return q
+        else:
+            return head_q
+
+    def predict(self, x, heads_only):
+        return self.forward(x, heads_only)
+
+class FruitMultiStatesFCNet(nn.Module, BasicNet):
+    def __init__(self, state_dim, action_dim, head_weights, optimizer_fn=None, gpu=True):
+        super(FruitMultiStatesFCNet, self).__init__()
+        hidden_size = 250
+        self.fc1 = nn.ModuleList([nn.Linear(state_dim, hidden_size) for _ in head_weights])
+        self.fc2 = nn.ModuleList([nn.Linear(hidden_size, action_dim) for _ in head_weights])
+        self.criterion = nn.MSELoss()
+        self.head_weights = head_weights
+        self.state_dim = state_dim
+        self.n_heads = head_weights.shape[0]
+        BasicNet.__init__(self, optimizer_fn, gpu)
+
+    def predict(self, x, merge):
+        head_q = []
+        for i in range(self.n_heads):
+            q = self.to_torch_variable(x[:, i, :])
+            q = self.fc1[i](q)
+            q = F.relu(q)
+            q = self.fc2[i](q)
+            head_q.append(q)
+        if merge:
+            q = [q * w for q, w in zip(head_q, self.head_weights)]
+            q = torch.stack(q, dim=0)
+            q = q.sum(0).squeeze(0)
+            return q
+        return head_q