class RandomStateDataset(data.Dataset): def __init__(self, env, cur_length, full_dataset_length, max_length, num_samples, epsilon): self.env = env self.cur_length = cur_length self.max_length = max_length self.num_samples = num_samples self.epsilon = epsilon self.full_dataset_length = full_dataset_length self.generator = DataGenerator(env, epsilon) self.states_full, self.actions, self.next_states_full,\ self.masks_full = self.generator.calc_data_full(self.full_dataset_length) n = self.cur_length - self.full_dataset_length if n > 0: self.states_rand, self.next_states_rand, self.masks_rand\ = self.generator.calc_data_rand(self.states_full[-1].view(-1, 2, 2, 2), self.actions, n) def reinitialize(self): self.states_full, self.actions, self.next_states_full,\ self.masks_full = self.generator.calc_data_full(self.full_dataset_length) n = self.cur_length - self.full_dataset_length if n > 0: self.states_rand, self.next_states_rand, self.masks_rand\ = self.generator.calc_data_rand(self.states_full[-1].view(-1, 2, 2, 2), self.actions, n) def __len__(self): return self.num_samples def __getitem__(self, _): length = torch.randint(0, self.cur_length, ()) if length < self.full_dataset_length: idx = torch.randint(0, len(self.states_full[length]), ()) state = self.states_full[length][idx] next_states = self.next_states_full[length][idx] mask = self.masks_full[length][idx] else: length = length - self.full_dataset_length idx = torch.randint(0, len(self.states_rand[length]), ()) state = self.states_rand[length][idx] next_states = self.next_states_rand[length][idx] mask = self.masks_rand[length][idx] return {'state': state, 'next_states': next_states, 'mask': mask}
class TestDataset(data.Dataset): def __init__(self, env, cur_length, full_dataset_length, max_length, num_samples, epsilon): self.env = env self.cur_length = cur_length self.max_length = max_length self.num_samples = num_samples self.epsilon = epsilon self.full_dataset_length = full_dataset_length self.generator = DataGenerator(env, epsilon) self.states_full, self.actions, _, _ = self.generator.calc_data_full( self.full_dataset_length) self.reinitialize() def reinitialize(self): n = self.cur_length - self.full_dataset_length if n > 0: self.states_rand, _, _ = self.generator.calc_data_rand( self.states_full[-1].view(-1, 2, 2, 2), self.actions, n) def __len__(self): return self.num_samples def __getitem__(self, _): length = torch.randint(0, self.cur_length, ()) if length < self.full_dataset_length: idx = torch.randint(0, len(self.states_full[length]), ()) state = self.states_full[length][idx] return {'state': state, 'length': length + 1} else: length = length - self.full_dataset_length idx = torch.randint(0, len(self.states_rand[length]), ()) state = self.states_rand[length][idx] return { 'state': state, 'length': length + self.full_dataset_length + 1 }