File size: 2,757 Bytes
7985e7e | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 | import numpy as np
def atleast_2d(x):
while x.ndim < 2:
x = np.expand_dims(x, axis=-1)
return x
class ReplayBuffer:
def __init__(self, max_n_episodes, max_path_length, termination_penalty):
self._dict = {
"path_lengths": np.zeros(max_n_episodes, dtype=np.int32),
}
self._count = 0
self.max_n_episodes = max_n_episodes
self.max_path_length = max_path_length
self.termination_penalty = termination_penalty
def __repr__(self):
return "[ datasets/buffer ] Fields:\n" + "\n".join(
f" {key}: {val.shape}" for key, val in self.items()
)
def __getitem__(self, key):
return self._dict[key]
def __setitem__(self, key, val):
self._dict[key] = val
self._add_attributes()
@property
def n_episodes(self):
return self._count
@property
def n_steps(self):
return sum(self["path_lengths"])
def _add_keys(self, path):
if hasattr(self, "keys"):
return
self.keys = list(path.keys())
def _add_attributes(self):
for key, val in self._dict.items():
setattr(self, key, val)
def items(self):
return {k: v for k, v in self._dict.items() if k != "path_lengths"}.items()
def _allocate(self, key, array):
assert key not in self._dict
dim = array.shape[-1]
shape = (self.max_n_episodes, self.max_path_length, dim)
self._dict[key] = np.zeros(shape, dtype=np.float32)
def add_path(self, path):
path_length = len(path["observations"])
assert path_length <= self.max_path_length
self._add_keys(path)
for key in self.keys:
array = atleast_2d(path[key])
if key not in self._dict:
self._allocate(key, array)
self._dict[key][self._count, :path_length] = array
if path["terminals"].any() and self.termination_penalty is not None:
assert not path[
"timeouts"
].any(), "Penalized a timeout episode for early termination"
self._dict["rewards"][
self._count, path_length - 1
] += self.termination_penalty
self._dict["path_lengths"][self._count] = path_length
self._count += 1
def truncate_path(self, path_ind, step):
old = self._dict["path_lengths"][path_ind]
new = min(step, old)
self._dict["path_lengths"][path_ind] = new
def finalize(self):
for key in self.keys + ["path_lengths"]:
self._dict[key] = self._dict[key][: self._count]
self._add_attributes()
print(f"[ datasets/buffer ] Finalized replay buffer | {self._count} episodes")
|