rollout() runs `while not np.all(done)` with `done` latched, so a sub-env that
terminates early keeps being driven -- physics and offscreen rendering included
-- until the slowest sub-env in the batch finishes. A batch of N runs for
max(episode_lengths) iterations to complete work that only needs
mean(episode_lengths), and all of the surplus is discarded.
Adds FreezeAfterEpisodeEnd, applied to each sub-env of the eval vector env. It
caches the terminal transition and replays it for any further step(), and also
absorbs Gymnasium's autoreset -- under AutoresetMode.NEXT_STEP the vector env
otherwise rebuilds a finished sub-env and runs it through an entire extra episode
the rollout throws away. Reward is zeroed on replay so a frozen sub-env cannot
inflate a return if a caller sums over the padded tail.
Thawing is signalled explicitly: rollout() passes NEW_ROLLOUT_OPTION in
reset(options=...). Gymnasium's autoreset calls reset() with no arguments, but so
would a caller passing seeds=None, and inferring from that would strand an env
frozen for a whole rollout.
AutoresetMode.DISABLED is not an alternative -- Gymnasium asserts that no
terminated env is ever stepped in that mode, so the wrapper is never reached. An
earlier version of this patch used it and the vector-env test caught the assert.
Co-authored-by: Claude Opus 5 <noreply@anthropic.com>
Dimitar Dimitrov