mirror of
https://github.com/huggingface/lerobot.git
synced 2026-08-05 16:09:47 +00:00
Dimitar Dimitrov
8135a8a8d1
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>