def __init__(self, seed, dataset_type, repeat=True):
        '''
        seed: the random seed
        dataset_type: 'pretrain' or 'finetune', which dataset to load
        repeat: whether to repeat to infinite length
        '''
        dataset_names_cfg = 'configs/pretrain_datasets.json' \
            if dataset_type == "pretrain" else 'configs/finetune_datasets.json'
        with open(dataset_names_cfg, 'r') as file:
            DATASET_NAMES = json.load(file)
        self.dataset_names = DATASET_NAMES
        sample_weights_cfg = 'configs/pretrain_sample_weights.json' \
            if dataset_type == "pretrain" else 'configs/finetune_sample_weights.json'
        # Load the sample weights
        with open(sample_weights_cfg, 'r') as file:
            SAMPLE_WEIGHTS = json.load(file)
        self.openx_dir = OPENX_EMBOD_DIR
        self.epsd_len_thresh_low = EPSD_LEN_THRESH_LOW
        self.epsd_len_thresh_high = EPSD_LEN_THRESH_HIGH
        self.repeat = repeat

        # Set the random seed
        tf.random.set_seed(seed)
        np.random.seed(seed)

        # Weights of the each dataset in the collection to sample from
        sample_weights = []

        self.name2dataset = {}
        for dataset_name in self.dataset_names:
            if dataset_name in DATASET_NAMES_NOOPENX:
                dataset = globals()[dataset_name].load_dataset(seed)
            else:
                dataset_path = dataset_to_path(dataset_name, self.openx_dir)
                dataset = tfds.builder_from_directory(builder_dir=dataset_path)
                dataset = dataset.as_dataset(split='all', shuffle_files=True)
                
                # You can add filter for other datasets
                if dataset_name == 'kuka':
                    dataset = dataset.filter(
                        lambda x: x['success'])
                elif dataset_name == 'bc_z':
                    dataset = dataset.filter(
                        lambda x: tf.math.greater(
                            next(iter(x['steps']))['observation']['episode_success'], 0.5))
                elif dataset_name == 'ucsd_pick_and_place_dataset_converted_externally_to_rlds':
                    dataset = dataset.filter(
                        lambda x: x['episode_metadata']['success'])
                elif dataset_name == 'utokyo_xarm_bimanual_converted_externally_to_rlds':
                    # Only preserve the meaningful episodes
                    dataset = dataset.filter(
                        lambda x: tf.math.equal(
                            next(iter(x['steps']))['language_instruction'],
                            tf.constant('Unfold a wrinkled towel.')))

            # Note: use cache() will cause the unexpected crash
            # dataset = dataset.map().cache().shuffle().repeat()
            dataset = dataset\
                .map(
                    lambda x: process_episode(x, dataset_name, 
                        IMAGE_KEYS[dataset_name]['image_keys'],
                        IMAGE_KEYS[dataset_name]['image_mask'])
                )
            
            # Change BGR to RGB if needed
            if dataset_name == 'fmb':
                dataset = dataset.map(bgr_to_rgb)
            
            if self.repeat:
                dataset = dataset.repeat()
            self.name2dataset[dataset_name] = iter(dataset)
            sample_weights.append(SAMPLE_WEIGHTS[dataset_name])
        # Normalize the sample weights
        sample_weights = np.array(sample_weights)
        self.sample_weights = sample_weights / np.sum(sample_weights)

        seed: the random seed
        dataset_type: 'pretrain' or 'finetune', which dataset to load
        repeat: whether to repeat to infinite length
        
policy/RDT/data/vla_dataset.py