Small data loader changes we found helpful during DROID training:
add optional argument to pass dataset_statistics to make_interleaved_dataset that are used for normalization (we still compute individual stats for weight & thread balancing) -- useful if you eg want to do co-training and normalize all datasets with the same stats
add argument to make_dataset_from_rlds that allows to skip keys during normalization (eg in DROID we don't normalize proprio)
adds utility to compute aggregate dataset stats across a list of stats (eg for co-training setting)
@kvablack would be great to merge this ASAP if you have time to take a look!
Small data loader changes we found helpful during DROID training:
make_interleaved_datasetthat are used for normalization (we still compute individual stats for weight & thread balancing) -- useful if you eg want to do co-training and normalize all datasets with the same statsmake_dataset_from_rldsthat allows to skip keys during normalization (eg in DROID we don't normalize proprio)@kvablack would be great to merge this ASAP if you have time to take a look!