[Bug] MBPO - Epoch loop missing steps_epoch reset

[matthiaskiller]

Observed

mbpo.py - line 199

in the for loop that is incrementing the steps of the current epoch, the steps_epoch iteration variable is not being reset after we observe a termination (done=True).
This will mean the next epoch will start from the steps_epoch+1 in which the previous epoch ended.
And the next epoch will be shorter than the actual epoch length (i.e. epoch_length - steps_epoch)

 for steps_epoch in range(cfg.overrides.epoch_length):
            if steps_epoch == 0 or done:
                obs, done = env.reset(), False

Is this behaviour desired?

Expected

If it's not desired we would propose something like this to set steps_epoch=0 :

 for steps_epoch in range(cfg.overrides.epoch_length):
            if steps_epoch == 0 or done:
                steps_epoch = 0
                obs, done = env.reset(), False

Thanks!:)

[luisenp]

This wasn't intentional, good catch! Would you be interested in submitting a PR with the fix? If not, that's OK, thanks a lot for the bug report!

[matthiaskiller]

Yes, sure!