modified network structures to match the paper

2023-02-18 10:13:02 +09:00
parent fb5c21557a
commit f96ad071d1
2 changed files with 6 additions and 5 deletions
--- a/networks.py
+++ b/networks.py
@@ -60,6 +60,7 @@ class RSSM(nn.Module):
            inp_dim += self._embed
        for i in range(self._layers_input):
            inp_layers.append(nn.Linear(inp_dim, self._hidden))
+            inp_layers.append(self._norm(self._hidden))
            inp_layers.append(self._act())
            if i == 0:
                inp_dim = self._hidden