Fix encoder training

Eclectic-Sheep · Jul 17, 2023 · 5c28ae9 · 5c28ae9
1 parent 3f4be4d
commit 5c28ae9
Showing 1 changed file with 3 additions and 1 deletion.
diff --git a/sheeprl/algos/sac_pixel/sac_pixel_continuous.py b/sheeprl/algos/sac_pixel/sac_pixel_continuous.py
@@ -48,6 +48,7 @@
 def train(
     fabric: Fabric,
     agent: SACPixelAgent,
+    encoder: Union[Encoder, _FabricModule],
     decoder: Union[Decoder, _FabricModule],
     actor_optimizer: Optimizer,
     qf_optimizer: Optimizer,
@@ -101,7 +102,7 @@ def train(
 
     # Update the decoder
     if global_step % args.decoder_update_freq == 0:
-        hidden = agent.critic.encoder(normalized_obs)
+        hidden = encoder(normalized_obs)
         reconstruction = decoder(hidden)
         reconstruction_loss = (
             F.mse_loss(preprocess_obs(data["observations"], bits=5), reconstruction)  # Reconstruction
@@ -359,6 +360,7 @@ def main():
                     train(
                         fabric,
                         agent,
+                        encoder,
                         decoder,
                         actor_optimizer,
                         qf_optimizer,