Parse Sketchy dataset as episodes rather than individual timesteps?

This issue concerns the [Sketchy dataset project](https://github.com/deepmind/deepmind-research/tree/master/sketchy).

The provided parsing [script](https://github.com/deepmind/deepmind-research/blob/master/sketchy/dataset_example.py) parses a single tfrecords file into a sequence of frames.

When I pass a sequence of filenames into `sketchy.load_frames`, it returns a dataset that that concatenates the frames from all of the files: each entry in the dataset is a single frame from one of the episodes, rather than a different epsiode.

How would I be able to create a dataset that parses the folder of tfrecord files as a dataset of episodes rather than a dataset of frames?