def sample(self, batchsize=None, chunk_length=None):
if batchsize == 0:
return self.all_data()
B = batchsize or self.batch_size
T = chunk_length or self._chunk_length
assert T <= self._horizon_length
if self._horizon_length == self.max_horizon:
start = self._pointer - self.max_horizon
else:
start = 0
end = self._pointer - T + 1
x = np.random.randint(start, end, B) # [B, ]
y = np.random.randint(0, self.n_copies, B) # (B, )
# (T, B) + (B, ) = (T, B)
xs = (np.tile(np.arange(T)[:, np.newaxis],
B) + x) % self._horizon_length
sample_idxs = (xs, y)
samples = {}
for k, v in self._buffer.items():
samples[k] = Data.from_nested_dict(
{_k: _v[sample_idxs] for _k, _v in v.items()}
)
return samples # [T, B, *]
def all_data(self):
samples = {}
for k, v in self._buffer.items():
samples[k] = Data.from_nested_dict(v)
return samples