Parse Sketchy dataset as episodes rather than individual timesteps?

[mbchang]

This issue concerns the Sketchy dataset project.

The provided parsing script 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?