correct for state representation

AOS55 · Oct 26, 2022 · c2bfb72 · c2bfb72
1 parent a29bf56
commit c2bfb72
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/libraries/latentsafesets/rl_trainers/mpc_trainer.py b/libraries/latentsafesets/rl_trainers/mpc_trainer.py
@@ -16,7 +16,6 @@ def __init__(self, env, cfg, modules):
         self.logdir = cfg.log_dir
 
         loss_plotter = LossPlotter(os.path.join(self.logdir, 'loss_plots'))
-        self.encoder_data_loader = EncoderDataLoader(cfg.env, frame_stack=cfg.frame_stack)
 
         self.trainers = []
 
@@ -33,6 +32,8 @@ def initial_train(self, replay_buffer):
         os.makedirs(update_dir, exist_ok=True)
         for trainer in self.trainers:
             if type(trainer) == VAETrainer:
+
+                self.encoder_data_loader = EncoderDataLoader(self.cfg.env, frame_stack=self.cfg.frame_stack)
                 trainer.initial_train(self.encoder_data_loader, update_dir)
             else:
                 trainer.initial_train(replay_buffer, update_dir)