def add(
self,
batch: Batch,
buffer_ids: Optional[Union[np.ndarray, List[int]]] = None
) -> Tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray]:
"""Add a batch of data into replay buffer.
:param Batch batch: the input data batch. Its keys must belong to the 7
reserved keys, and "obs", "act", "rew", "done" is required.
:param buffer_ids: to make consistent with other buffer's add function; if it
is not None, we assume the input batch's first dimension is always 1.
Return (current_index, episode_reward, episode_length, episode_start_index). If
the episode is not finished, the return value of episode_length and
episode_reward is 0.
"""
# preprocess batch
b = Batch()
for key in set(self._reserved_keys).intersection(batch.keys()):
b.__dict__[key] = batch[key]
batch = b
assert set(["obs", "act", "rew", "done"]).issubset(batch.keys())
stacked_batch = buffer_ids is not None
if stacked_batch:
assert len(batch) == 1
if self._save_only_last_obs:
batch.obs = batch.obs[:, -1] if stacked_batch else batch.obs[-1]
if not self._save_obs_next:
batch.pop("obs_next", None)
elif self._save_only_last_obs:
batch.obs_next = (
batch.obs_next[:, -1] if stacked_batch else batch.obs_next[-1]
)
# get ptr
if stacked_batch:
rew, done = batch.rew[0], batch.done[0]
else:
rew, done = batch.rew, batch.done
ptr, ep_rew, ep_len, ep_idx = list(
map(lambda x: np.array([x]), self._add_index(rew, done))
)
try:
self._meta[ptr] = batch
except ValueError:
stack = not stacked_batch
batch.rew = batch.rew.astype(float)
batch.done = batch.done.astype(bool)
if self._meta.is_empty():
self._meta = _create_value( # type: ignore
batch, self.maxsize, stack)
else: # dynamic key pops up in batch
_alloc_by_keys_diff(self._meta, batch, self.maxsize, stack)
self._meta[ptr] = batch
return ptr, ep_rew, ep_len, ep_idx
def add(
self,
batch: Batch,
buffer_ids: Optional[Union[np.ndarray, List[int]]] = None
) -> Tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray]:
"""Add a batch of data into ReplayBufferManager.
Each of the data's length (first dimension) must equal to the length of
buffer_ids. By default buffer_ids is [0, 1, ..., buffer_num - 1].
Return (current_index, episode_reward, episode_length, episode_start_index). If
the episode is not finished, the return value of episode_length and
episode_reward is 0.
"""
# preprocess batch
new_batch = Batch()
for key in set(self._reserved_keys).intersection(batch.keys()):
new_batch.__dict__[key] = batch[key]
batch = new_batch
assert set(["obs", "act", "rew", "done"]).issubset(batch.keys())
if self._save_only_last_obs:
batch.obs = batch.obs[:, -1]
if not self._save_obs_next:
batch.pop("obs_next", None)
elif self._save_only_last_obs:
batch.obs_next = batch.obs_next[:, -1]
# get index
if buffer_ids is None:
buffer_ids = np.arange(self.buffer_num)
ptrs, ep_lens, ep_rews, ep_idxs = [], [], [], []
for batch_idx, buffer_id in enumerate(buffer_ids):
ptr, ep_rew, ep_len, ep_idx = self.buffers[buffer_id]._add_index(
batch.rew[batch_idx], batch.done[batch_idx]
)
ptrs.append(ptr + self._offset[buffer_id])
ep_lens.append(ep_len)
ep_rews.append(ep_rew)
ep_idxs.append(ep_idx + self._offset[buffer_id])
self.last_index[buffer_id] = ptr + self._offset[buffer_id]
self._lengths[buffer_id] = len(self.buffers[buffer_id])
ptrs = np.array(ptrs)
try:
self._meta[ptrs] = batch
except ValueError:
batch.rew = batch.rew.astype(float)
batch.done = batch.done.astype(bool)
if self._meta.is_empty():
self._meta = _create_value( # type: ignore
batch, self.maxsize, stack=False)
else: # dynamic key pops up in batch
_alloc_by_keys_diff(self._meta, batch, self.maxsize, False)
self._set_batch_for_children()
self._meta[ptrs] = batch
return ptrs, np.array(ep_rews), np.array(ep_lens), np.array(ep_idxs)
def set_batch(self, batch: Batch) -> None:
"""Manually choose the batch you want the ReplayBuffer to manage."""
assert len(batch) == self.maxsize and set(batch.keys()).issubset(
self._reserved_keys
), "Input batch doesn't meet ReplayBuffer's data form requirement."
self._meta = batch
