fix epoch/episode iteration bug in mbpo.py (facebookresearch#160)

* fix epoch/episode iteration bug in mbpo.py * fixed trailing whitespace Co-authored-by: Matthias <ga26qen@mytum.de>
T-will19 · Nov 9, 2022 · f90a297 · f90a297
1 parent 25cffd5
commit f90a297
Showing 1 changed file with 1 addition and 0 deletions.
diff --git a/mbrl/algorithms/mbpo.py b/mbrl/algorithms/mbpo.py
@@ -198,6 +198,7 @@ def train(
         obs, done = None, False
         for steps_epoch in range(cfg.overrides.epoch_length):
             if steps_epoch == 0 or done:
+                steps_epoch = 0
                 obs, done = env.reset(), False
             # --- Doing env step and adding to model dataset ---
             next_obs, reward, done, _ = mbrl.util.common.step_env_and_add_to_buffer(