05_dataloading.ipynb: bridge dataset loading

[AasherH]

I am attempting to run the example 05_dataloading.ipynb with the bridge dataset. However, I cannot seem to get the code to run. I have tried manually changing the builder from tfds.builder(name, data_dir=data_dir) to builder = tfds.builder_from_directory(builder_dir='gs://gresearch/robotics/bridge/0.1.0/') as in the 01_example, but I still run into issues, presumably due to how the trajectories (actions) are stored.

Does anyone know how I can use the dataloader for the bridge dataset? The error in 05_dataloading.ipynb occurs at full_dataset = full_dataset.traj_map(restructure).filter(is_nonzero_length) in dataset.py.