EpisodeAwareSampler allows the dropping of first n or last n frames
from each episode.
Additionally, you can select which episodes to run on by passig in
their ids (complements functionality of the hf dataset slice syntax
to give extra flexibility to the user).
* modify train.py and diffusion.yaml to leverage new functionality
and train on pusht with last 7 steps per episode dropped
* add unit tests
Radek Osmulski