def test_data_storage_in_raw_data_with_data_size_limit(self):
     config_dict = {
         'output_path': self.output_dir,
         'max_size': 25,
         'separate_raw_data_runs': True
     }
     config = DataSaverConfig().create(config_dict=config_dict)
     self.data_saver = DataSaver(config=config)
     first_info = generate_dummy_dataset(self.data_saver, num_runs=2)
     self.assertEqual(sum(first_info['episode_lengths']),
                      self.data_saver._frame_counter)
     self.data_saver.update_saving_directory()
     second_info = generate_dummy_dataset(self.data_saver, num_runs=2)
     self.assertTrue(
         (sum(first_info['episode_lengths']) +
          sum(second_info['episode_lengths'])) > config_dict['max_size'])
     self.assertTrue(
         self.data_saver._frame_counter <= config_dict['max_size'])
     raw_data_dir = os.path.dirname(self.data_saver.get_saving_directory())
     count_actual_frames = sum([
         len(
             os.listdir(
                 os.path.join(raw_data_dir, episode_dir, 'observation')))
         for episode_dir in os.listdir(raw_data_dir)
     ])
     self.assertEqual(count_actual_frames, self.data_saver._frame_counter)
def generate_random_dataset_in_raw_data(
        output_dir: str,
        num_runs: int = 20,
        input_size: tuple = (100, 100, 3),
        output_size: tuple = (1, ),
        continuous: bool = True,
        fixed_input_value: Union[float, np.ndarray] = None,
        fixed_output_value: Union[float, np.ndarray] = None,
        store_hdf5: bool = False) -> dict:
    """Generate data, stored in raw_data directory of output_dir"""
    data_saver = DataSaver(config=DataSaverConfig().create(
        config_dict={
            'output_path': output_dir,
            'store_hdf5': store_hdf5,
            'separate_raw_data_runs': True
        }))
    info = generate_dummy_dataset(data_saver,
                                  num_runs=num_runs,
                                  input_size=input_size,
                                  output_size=output_size,
                                  continuous=continuous,
                                  fixed_input_value=fixed_input_value,
                                  fixed_output_value=fixed_output_value,
                                  store_hdf5=store_hdf5)
    return info
    def test_create_hdf5_files_subsampled_in_time(self):
        num_runs = 10
        split = 1.0
        subsample = 3
        config_dict = {
            'output_path': self.output_dir,
            'training_validation_split': split,
            'store_hdf5': True,
            'subsample_hdf5': subsample,
            'separate_raw_data_runs': True
        }
        config = DataSaverConfig().create(config_dict=config_dict)
        self.data_saver = DataSaver(config=config)
        info = generate_dummy_dataset(self.data_saver, num_runs=num_runs)
        self.data_saver.create_train_validation_hdf5_files()

        config = DataLoaderConfig().create(
            config_dict={
                'output_path': self.output_dir,
                'hdf5_files': [os.path.join(self.output_dir, 'train.hdf5')]
            })
        training_data_loader = DataLoader(config=config)
        training_data_loader.load_dataset()
        training_data = training_data_loader.get_dataset()

        self.assertEqual(
            len(training_data),
            sum([
                np.ceil((el - 1) / subsample) + 1
                for el in info['episode_lengths']
            ]))
    def test_big_data_hdf5_loop(self):
        # create 3 datasets as hdf5 files
        hdf5_files = []
        infos = []
        for index in range(3):
            output_path = os.path.join(self.output_dir, f'ds{index}')
            os.makedirs(output_path, exist_ok=True)
            config_dict = {
                'output_path': output_path,
                'store_hdf5': True,
                'training_validation_split': 1.0
            }
            config = DataSaverConfig().create(config_dict=config_dict)
            self.data_saver = DataSaver(config=config)
            infos.append(
                generate_dummy_dataset(self.data_saver,
                                       num_runs=2,
                                       input_size=(3, 10, 10),
                                       fixed_input_value=(0.3 * index) *
                                       np.ones((3, 10, 10)),
                                       store_hdf5=True))
            self.assertTrue(
                os.path.isfile(os.path.join(output_path, 'train.hdf5')))
            hdf5_files.append(os.path.join(output_path, 'train.hdf5'))
            hdf5_files.append(os.path.join(output_path, 'wrong.hdf5'))

        # create data loader with big data tag and three hdf5 training sets
        conf = {
            'output_path': self.output_dir,
            'hdf5_files': hdf5_files,
            'batch_size': 15,
            'loop_over_hdf5_files': True
        }
        loader = DataLoader(DataLoaderConfig().create(config_dict=conf))

        # sample data batches and see that index increases every two batches sampled
        for batch in loader.get_data_batch():
            self.assertAlmostEqual(batch.observations[0][0, 0, 0].item(), 0)
        for batch in loader.get_data_batch():
            self.assertAlmostEqual(batch.observations[0][0, 0, 0].item(), 0.3,
                                   2)
        for batch in loader.get_data_batch():
            self.assertAlmostEqual(batch.observations[0][0, 0, 0].item(), 0.6,
                                   2)
        for batch in loader.get_data_batch():
            self.assertAlmostEqual(batch.observations[0][0, 0, 0].item(), 0, 2)
        for batch in loader.sample_shuffled_batch():
            self.assertAlmostEqual(batch.observations[0][0, 0, 0].item(), 0.3,
                                   2)
        for batch in loader.sample_shuffled_batch():
            self.assertAlmostEqual(batch.observations[0][0, 0, 0].item(), 0.6,
                                   2)
        for batch in loader.sample_shuffled_batch():
            self.assertAlmostEqual(batch.observations[0][0, 0, 0].item(), 0, 2)
 def setUp(self) -> None:
     self.output_dir = f'{os.environ["PWD"]}/test_dir/{get_filename_without_extension(__file__)}'
     if not os.path.isdir(self.output_dir):
         os.makedirs(self.output_dir)
     config_dict = {'output_path': self.output_dir, 'store_hdf5': True}
     config = DataSaverConfig().create(config_dict=config_dict)
     self.data_saver = DataSaver(config=config)
     self.info = generate_dummy_dataset(self.data_saver,
                                        num_runs=20,
                                        input_size=(100, 100, 3),
                                        output_size=(3, ),
                                        continuous=False)
    def test_sample_batch(self):
        self.info = generate_dummy_dataset(self.data_saver,
                                           num_runs=20,
                                           input_size=(100, 100, 3),
                                           output_size=(3, ),
                                           continuous=False)
        max_num_batches = 2
        config_dict = {
            'data_directories': self.info['episode_directories'],
            'output_path': self.output_dir,
            'random_seed': 1,
            'batch_size': 3
        }
        data_loader = DataLoader(config=DataLoaderConfig().create(
            config_dict=config_dict))
        data_loader.load_dataset()
        first_batch = []
        index = 0
        for index, batch in enumerate(
                data_loader.sample_shuffled_batch(
                    max_number_of_batches=max_num_batches)):
            if index == 0:
                first_batch = deepcopy(batch)
            self.assertEqual(len(batch), config_dict['batch_size'])
        self.assertEqual(index, max_num_batches - 1)

        # test sampling seed for reproduction
        config_dict['random_seed'] = 2
        data_loader = DataLoader(config=DataLoaderConfig().create(
            config_dict=config_dict))
        data_loader.load_dataset()
        second_batch = []
        for index, batch in enumerate(
                data_loader.sample_shuffled_batch(
                    max_number_of_batches=max_num_batches)):
            second_batch = deepcopy(batch)
            break
        self.assertNotEqual(np.sum(np.asarray(first_batch.observations[0])),
                            np.sum(np.asarray(second_batch.observations[0])))
        config_dict['random_seed'] = 1
        data_loader = DataLoader(config=DataLoaderConfig().create(
            config_dict=config_dict))
        data_loader.load_dataset()
        third_batch = []
        for index, batch in enumerate(
                data_loader.sample_shuffled_batch(
                    max_number_of_batches=max_num_batches)):
            third_batch = deepcopy(batch)
            break
        self.assertEqual(np.sum(np.asarray(first_batch.observations[0])),
                         np.sum(np.asarray(third_batch.observations[0])))
 def test_store_in_ram(self):
     config_dict = {
         'output_path': self.output_dir,
         'store_on_ram_only': True,
         'max_size': 10
     }
     number_of_runs = 10
     config = DataSaverConfig().create(config_dict=config_dict)
     self.data_saver = DataSaver(config=config)
     info = generate_dummy_dataset(self.data_saver, num_runs=number_of_runs)
     data = self.data_saver.get_dataset()
     self.assertEqual(len(data), config_dict['max_size'])
     for lst in [data.observations, data.actions, data.rewards, data.done]:
         self.assertEqual(len(lst), config_dict['max_size'])
         self.assertTrue(isinstance(lst[0], torch.Tensor))
 def test_empty_saving_directory(self):
     config_dict = {
         'output_path': self.output_dir,
         'separate_raw_data_runs': True
     }
     number_of_runs = 5
     config = DataSaverConfig().create(config_dict=config_dict)
     self.data_saver = DataSaver(config=config)
     info = generate_dummy_dataset(self.data_saver, num_runs=number_of_runs)
     self.assertEqual(
         len(os.listdir(os.path.join(self.output_dir, 'raw_data'))),
         number_of_runs)
     self.data_saver.empty_raw_data_in_output_directory()
     self.assertEqual(
         len(os.listdir(os.path.join(self.output_dir, 'raw_data'))), 0)
    def test_data_loader_from_raw_path_dirs(self):
        self.info = generate_dummy_dataset(self.data_saver,
                                           num_runs=20,
                                           input_size=(100, 100, 3),
                                           output_size=(3, ),
                                           continuous=False)
        config_dict = {
            'data_directories': [self.output_dir],
            'output_path': self.output_dir,
        }
        config = DataLoaderConfig().create(config_dict=config_dict)
        data_loader = DataLoader(config=config)
        data_loader.load_dataset()

        config = DataLoaderConfig().create(config_dict=config_dict)
        for d in config.data_directories:
            self.assertTrue(os.path.isdir(d))
    def test_data_batch(self):
        self.info = generate_dummy_dataset(self.data_saver,
                                           num_runs=20,
                                           input_size=(100, 100, 3),
                                           output_size=(3, ),
                                           continuous=False)
        config_dict = {
            'data_directories': self.info['episode_directories'],
            'output_path': self.output_dir,
            'random_seed': 1,
            'batch_size': 3
        }
        data_loader = DataLoader(config=DataLoaderConfig().create(
            config_dict=config_dict))
        data_loader.load_dataset()

        for batch in data_loader.get_data_batch():
            self.assertEqual(len(batch), config_dict['batch_size'])
            break
    def test_data_loading(self):
        self.info = generate_dummy_dataset(self.data_saver,
                                           num_runs=20,
                                           input_size=(100, 100, 3),
                                           output_size=(3, ),
                                           continuous=False)
        config_dict = {
            'data_directories': self.info['episode_directories'],
            'output_path': self.output_dir,
        }
        config = DataLoaderConfig().create(config_dict=config_dict)
        data_loader = DataLoader(config=config)
        data_loader.load_dataset()

        # assert nothing is empty
        for k in ['observations', 'actions', 'rewards', 'done']:
            data = eval(f'data_loader.get_dataset().{k}')
            self.assertTrue(len(data) > 0)
            self.assertTrue(sum(data[0].shape) > 0)
 def test_data_storage_in_raw_data(self):
     config_dict = {
         'output_path': self.output_dir,
         'separate_raw_data_runs': True
     }
     config = DataSaverConfig().create(config_dict=config_dict)
     self.data_saver = DataSaver(config=config)
     info = generate_dummy_dataset(self.data_saver, num_runs=2)
     for total, episode_dir in zip(info['episode_lengths'],
                                   info['episode_directories']):
         self.assertEqual(
             len(
                 os.listdir(
                     os.path.join(self.output_dir, 'raw_data', episode_dir,
                                  'observation'))), total)
         with open(
                 os.path.join(self.output_dir, 'raw_data', episode_dir,
                              'action.data')) as f:
             expert_controls = f.readlines()
             self.assertEqual(len(expert_controls), total)
 def test_data_subsample(self):
     self.info = generate_dummy_dataset(self.data_saver,
                                        num_runs=20,
                                        input_size=(100, 100, 3),
                                        output_size=(3, ),
                                        continuous=False)
     subsample = 4
     config_dict = {
         'data_directories': self.info['episode_directories'],
         'output_path': self.output_dir,
         'random_seed': 1,
         'batch_size': 3,
         'subsample': subsample
     }
     data_loader = DataLoader(config=DataLoaderConfig().create(
         config_dict=config_dict))
     data_loader.load_dataset()
     self.assertTrue(
         sum([
             np.ceil((el - 1) / subsample) + 1
             for el in self.info['episode_lengths']
         ]), len(data_loader.get_dataset()))
    def test_create_train_validation_hdf5_files(self):
        num_runs = 10
        split = 0.7
        config_dict = {
            'output_path': self.output_dir,
            'training_validation_split': split,
            'store_hdf5': True,
            'separate_raw_data_runs': True
        }
        config = DataSaverConfig().create(config_dict=config_dict)
        self.data_saver = DataSaver(config=config)
        info = generate_dummy_dataset(self.data_saver, num_runs=num_runs)
        self.data_saver.create_train_validation_hdf5_files()

        config = DataLoaderConfig().create(
            config_dict={
                'output_path': self.output_dir,
                'hdf5_files': [os.path.join(self.output_dir, 'train.hdf5')]
            })
        training_data_loader = DataLoader(config=config)
        training_data_loader.load_dataset()
        training_data = training_data_loader.get_dataset()

        config = DataLoaderConfig().create(
            config_dict={
                'output_path': self.output_dir,
                'hdf5_files':
                [os.path.join(self.output_dir, 'validation.hdf5')]
            })
        validation_data_loader = DataLoader(config=config)
        validation_data_loader.load_dataset()
        validation_data = validation_data_loader.get_dataset()

        self.assertEqual(len(training_data),
                         sum(info['episode_lengths'][:int(split * num_runs)]))
        self.assertEqual(len(validation_data),
                         sum(info['episode_lengths'][int(split * num_runs):]))