def _build_replay_buffer(self, use_staging):
"""Creates the replay buffer used by the agent.
Args:
use_staging: bool, if True, uses a staging area to prefetch data for
faster training.
Returns:
A `WrappedPrioritizedReplayBuffer` object.
Raises:
ValueError: if given an invalid replay scheme.
"""
if self._replay_scheme not in ['uniform', 'prioritized']:
raise ValueError('Invalid replay scheme: {}'.format(self._replay_scheme))
# Both replay schemes use the same data structure, but the 'uniform' scheme
# sets all priorities to the same value (which yields uniform sampling).
extra_elements = [ReplayElement('return', (), np.float32)]
return prioritized_replay_buffer.WrappedPrioritizedReplayBuffer(
observation_shape=self.observation_shape,
stack_size=self.stack_size,
use_staging=use_staging,
update_horizon=self.update_horizon,
gamma=self.gamma,
observation_dtype=self.observation_dtype.as_numpy_dtype,
extra_storage_types=extra_elements,
replay_forgetting=self._replay_forgetting,
sample_newest_immediately=self._sample_newest_immediately)
def get_transition_elements(self, batch_size=None):
"""Returns a 'type signature' for sample_transition_batch.
Args:
batch_size: int, number of transitions returned. If None, the default
batch_size will be used.
Returns:
signature: A namedtuple describing the method's return type signature.
"""
parent_transition_type = (
super(OutOfGraphPrioritizedReplayBuffer,
self).get_transition_elements(batch_size))
probablilities_type = [
ReplayElement('sampling_probabilities', (batch_size,), np.float32)
]
return parent_transition_type + probablilities_type
def get_add_args_signature(self):
"""The signature of the add function.
The signature is the same as the one for OutOfGraphReplayBuffer, with an
added priority.
Returns:
list of ReplayElements defining the type of the argument signature needed
by the add function.
"""
parent_add_signature = super(OutOfGraphPrioritizedReplayBuffer,
self).get_add_args_signature()
add_signature = parent_add_signature + [
ReplayElement('priority', (), np.float32)
]
return add_signature