Rectifying Simba Implementation (#4)

leeyngdo · web-flow · commit faa844a4957f · 2025-03-17T14:12:48.000+09:00
diff --git a/scale_rl/agents/simba/simba_network.py b/scale_rl/agents/simba/simba_network.py
@@ -23,8 +23,9 @@ def setup(self):
         )
         self.encoder = nn.Sequential(
             [
-                PreLNResidualBlock(hidden_dim=self.hidden_dim)
-                for _ in range(self.num_blocks)
+                *[PreLNResidualBlock(hidden_dim=self.hidden_dim) 
+                  for _ in range(self.num_blocks)], 
+                nn.LayerNorm(),
             ]
         )
         self.predictor = NormalTanhPolicy(self.action_dim)
@@ -51,8 +52,9 @@ def setup(self):
         )
         self.encoder = nn.Sequential(
             [
-                PreLNResidualBlock(hidden_dim=self.hidden_dim)
-                for _ in range(self.num_blocks)
+                *[PreLNResidualBlock(hidden_dim=self.hidden_dim) 
+                  for _ in range(self.num_blocks)], 
+                nn.LayerNorm(),
             ]
         )
         self.predictor = LinearCritic()

Original file line number	Diff line number	Diff line change
`@@ -23,8 +23,9 @@ def setup(self):`
`23`	`23`	`)`
`24`	`24`	`self.encoder = nn.Sequential(`
`25`	`25`	`[`
`26`		`- PreLNResidualBlock(hidden_dim=self.hidden_dim)`
`27`		`- for _ in range(self.num_blocks)`
	`26`	`+ *[PreLNResidualBlock(hidden_dim=self.hidden_dim)`
	`27`	`+ for _ in range(self.num_blocks)],`
	`28`	`+ nn.LayerNorm(),`
`28`	`29`	`]`
`29`	`30`	`)`
`30`	`31`	`self.predictor = NormalTanhPolicy(self.action_dim)`
`@@ -51,8 +52,9 @@ def setup(self):`
`51`	`52`	`)`
`52`	`53`	`self.encoder = nn.Sequential(`
`53`	`54`	`[`
`54`		`- PreLNResidualBlock(hidden_dim=self.hidden_dim)`
`55`		`- for _ in range(self.num_blocks)`
	`55`	`+ *[PreLNResidualBlock(hidden_dim=self.hidden_dim)`
	`56`	`+ for _ in range(self.num_blocks)],`
	`57`	`+ nn.LayerNorm(),`
`56`	`58`	`]`
`57`	`59`	`)`
`58`	`60`	`self.predictor = LinearCritic()`