def sample(self, count):
result = [
self.samples[random.randrange(len(self.samples))]
for idx in range(count)
]
result = default_collate(result)
return result
def __getitem__(self, index):
result = ...
if isinstance(index, slice):
result = default_collate(self.experiences[index])
else:
result = self.experiences[index]
return result
def out_of_sample(self):
result = []
for idx in range(self.batch_size):
sample = self.bad_prepare()
result.append(sample)
data, *args = to_device(default_collate(result), self.device)
#data = self.integrator.integrate(self.score, data, *args).detach()
detached = to_device(data.detach(), "cpu")
return to_device((data, *args), self.device)
def get_batch(self, size=64, length=0):
items = []
for idx in range(size):
index = random.randrange(0, len(self))
traj = self[index]
offset = random.randrange(0, len(traj) - length)
window = ...
if length == 0:
window = traj[offset]
else:
window = traj[offset:offset + length]
items.append(window)
return default_collate(items)
def collate(cls, inputs):
initial_state = []
final_state = []
action = []
reward = []
logits = []
outputs = []
terminal = []
for experience in inputs:
initial_state.append(experience.initial_state)
final_state.append(experience.final_state)
action.append(experience.action)
reward.append(experience.reward)
logits.append(experience.logits)
outputs.append(experience.outputs)
terminal.append(experience.terminal)
if logits[0] is None:
logits = None
else:
logits = default_collate(logits)
if outputs[0] is None:
outputs = None
else:
outputs = default_collate(outputs)
initial_state = default_collate(initial_state)
final_state = default_collate(final_state)
action = default_collate(action)
reward = default_collate(reward)
terminal = default_collate(terminal)
result = Experience(initial_state=initial_state,
final_state=final_state,
action=action,
reward=reward.float(),
terminal=terminal,
logits=logits,
outputs=outputs)
return result
def prepare_sample(self):
results = []
for idx in range(self.batch_size):
results.append(self.prepare())
return default_collate(results)