bair_robot_pushing dataset: S3 version (issue #737).

PiperOrigin-RevId: 257603054

Author: pierrot0

tensorflow_datasets/video/bair_robot_pushing.py

@@ -54,6 +54,11 @@ class BairRobotPushingSmall(tfds.core.GeneratorBasedBuilder):
 
   VERSION = tfds.core.Version("1.0.0",
                               experiments={tfds.core.Experiment.S3: False})
+  SUPPORTED_VERSIONS = [
+      tfds.core.Version("2.0.0"),
+  ]
+  # Versions history:
+  # 2.0.0: S3 (new shuffling, sharding and slicing mechanism).
 
   def _info(self):
     # The Bair dataset consist of a sequence of frames (video) with associated
@@ -103,7 +108,8 @@ def _generate_examples(self, filedir):
       filepath = os.path.join(filedir, filename)
 
       # For each video inside the file
-      for example_str in tf.compat.v1.io.tf_record_iterator(filepath):
+      for i, example_str in enumerate(
+          tf.compat.v1.io.tf_record_iterator(filepath)):
         example = tf.train.SequenceExample.FromString(example_str)
 
         # Merge all frames together
@@ -140,4 +146,7 @@ def _generate_examples(self, filedir):
         #     {'action': [...], 'image_main': img_frame1, ...},  # Frame 1
         #     ...,
         # ]
-        yield all_frames
+        if self.version.implements(tfds.core.Experiment.S3):
+          yield "%s_%s" % (filepath, i), all_frames
+        else:
+          yield all_frames

tensorflow_datasets/video/bair_robot_pushing_test.py

@@ -32,5 +32,9 @@ class BairRobotPushingTest(testing.DatasetBuilderTestCase):
   }
 
 
+class BairRobotPushingS3Test(BairRobotPushingTest):
+  VERSION = "experimental_latest"
+
+
 if __name__ == "__main__":
   testing.test_main()