code: def test_episode_boundaries_are_coherent(self):
      num_episodes = 0
      dataset = _build_dataset(self.dataset_dir)
      for episode in dataset:
          steps = episode["steps"]
          num_steps = len(steps["is_first"])
          for i in range(num_steps):
              if i == 0:
                  # Test first step:
                  self.assertEqual(steps["is_first"][i], True)
                  self.assertEqual(steps["is_last"][i], False)
                  self.assertEqual(steps["is_terminal"][i], False)
                  self.assertEqual(steps["discount"][i], 1.0)
              # Test middle steps:
              if i > 0 and i < num_steps - 1:
                  self.assertEqual(steps["is_first"][i], False)
                  self.assertEqual(steps["is_last"][i], False)
                  self.assertEqual(steps["discount"][i], 1.0)
              # Test final steps:
              if i == num_steps - 1:
                  self.assertEqual(steps["is_first"][i], False)
                  self.assertEqual(steps["is_last"][i], True)
                  self.assertEqual(steps["is_terminal"][i], True)
                  self.assertEqual(steps["discount"][i], 0.0)
  
          num_episodes += 1
      # The number of episodes should be consistent with the total available.
      self.assertEqual(num_episodes, 2)
  
file_path: datasets/tensorflow_datasets/robomimic/dataset_utils_test.py
imports: import inspect
  import os
  from absl import flags
  from absl.testing import absltest
  from etils import epath
  import h5py
  import numpy
  from tensorflow_datasets.robomimic import dataset_utils
  import tree
label_desc: Use image node_id 3 for calling the DatasetUtilsTest obj's underlying member method code with example usage: obj.test_episode_boundaries_are_coherent() without return types