def __iter__(self):
        '''
        Sample batches of episodes for an epoch.
        '''
        while True:
            dataset_name = np.random.choice(self.dataset_names, p=self.sample_weights)
            episode = next(self.name2dataset[dataset_name])
            if dataset_name == "agilex":
                episode_steps = flatten_episode_agilex(episode)
            else:
                episode_steps = flatten_episode(episode)
            # Filter too short
            if len(episode_steps) < self.epsd_len_thresh_low:
                continue
            # Randomly sample too long
            if len(episode_steps) > self.epsd_len_thresh_high:
                episode_steps = random.sample(episode_steps, self.epsd_len_thresh_high)
                
            yield episode_steps

        Sample batches of episodes for an epoch.
        
policy/RDT/data/vla_dataset.py