def __init__(self,
capacity: int = 10000,
storage_unit: str = "timesteps",
alpha: float = 1.0,
**kwargs):
"""Initializes a PrioritizedReplayBuffer instance.
Args:
capacity: Max number of timesteps to store in the FIFO
buffer. After reaching this number, older samples will be
dropped to make space for new ones.
storage_unit: Either 'timesteps', 'sequences' or
'episodes'. Specifies how experiences are stored.
alpha: How much prioritization is used
(0.0=no prioritization, 1.0=full prioritization).
**kwargs: Forward compatibility kwargs.
"""
ReplayBuffer.__init__(self, capacity, storage_unit, **kwargs)
assert alpha > 0
self._alpha = alpha
# Segment tree must have capacity that is a power of 2
it_capacity = 1
while it_capacity < self.capacity:
it_capacity *= 2
self._it_sum = SumSegmentTree(it_capacity)
self._it_min = MinSegmentTree(it_capacity)
self._max_priority = 1.0
self._prio_change_stats = WindowStat("reprio", 1000)
def __init__(self, capacity: int = 10000, storage_unit: str = "timesteps"):
"""Initializes a ReservoirBuffer instance.
Args:
capacity: Max number of timesteps to store in the FIFO
buffer. After reaching this number, older samples will be
dropped to make space for new ones.
storage_unit: Either 'sequences' or 'timesteps'. Specifies how
experiences are stored.
"""
ReplayBuffer.__init__(self, capacity, storage_unit)
self._num_add_calls = 0
self._num_evicted = 0
def __init__(self,
capacity: int = 10000,
storage_unit: str = "timesteps",
num_shards: int = 1,
learning_starts: int = 1000,
replay_mode: str = "independent",
replay_sequence_length: int = 1,
replay_burn_in: int = 0,
replay_zero_init_states: bool = True,
underlying_buffer_config: dict = None,
**kwargs):
"""Initializes a MultiAgentReplayBuffer instance.
Args:
num_shards: The number of buffer shards that exist in total
(including this one).
storage_unit: Either 'timesteps', 'sequences' or
'episodes'. Specifies how experiences are stored. If they
are stored in episodes, replay_sequence_length is ignored.
learning_starts: Number of timesteps after which a call to
`sample()` will yield samples (before that, `sample()` will
return None).
capacity: The capacity of the buffer, measured in `storage_unit`.
replay_mode: One of "independent" or "lockstep". Determines,
whether batches are sampled independently or to an equal
amount.
replay_sequence_length: The sequence length (T) of a single
sample. If > 1, we will sample B x T from this buffer. This
only has an effect if storage_unit is 'timesteps'.
replay_burn_in: This is the number of timesteps
each sequence overlaps with the previous one to generate a
better internal state (=state after the burn-in), instead of
starting from 0.0 each RNN rollout. This only has an effect
if storage_unit is `sequences`.
replay_zero_init_states: Whether the initial states in the
buffer (if replay_sequence_length > 0) are alwayas 0.0 or
should be updated with the previous train_batch state outputs.
underlying_buffer_config: A config that contains all necessary
constructor arguments and arguments for methods to call on
the underlying buffers.
**kwargs: Forward compatibility kwargs.
"""
shard_capacity = capacity // num_shards
ReplayBuffer.__init__(self, capacity, storage_unit)
# If the user provides an underlying buffer config, we use to
# instantiate and interact with underlying buffers
self.underlying_buffer_config = underlying_buffer_config
if self.underlying_buffer_config is not None:
self.underlying_buffer_call_args = self.underlying_buffer_config
else:
self.underlying_buffer_call_args = {}
self.replay_starts = learning_starts // num_shards
self.replay_mode = replay_mode
self.replay_sequence_length = replay_sequence_length
self.replay_burn_in = replay_burn_in
self.replay_zero_init_states = replay_zero_init_states
if (replay_sequence_length > 1
and self._storage_unit is not StorageUnit.SEQUENCES):
logger.warning(
"MultiAgentReplayBuffer configured with "
"`replay_sequence_length={}`, but `storage_unit={}`. "
"replay_sequence_length will be ignored and set to 1.".format(
replay_sequence_length, storage_unit))
self.replay_sequence_length = 1
if replay_sequence_length == 1 and self._storage_unit is StorageUnit.SEQUENCES:
logger.warning(
"MultiAgentReplayBuffer configured with "
"`replay_sequence_length={}`, but `storage_unit={}`. "
"This will result in sequences equal to timesteps.".format(
replay_sequence_length, storage_unit))
if replay_mode in ["lockstep", ReplayMode.LOCKSTEP]:
self.replay_mode = ReplayMode.LOCKSTEP
if self._storage_unit in [
StorageUnit.EPISODES, StorageUnit.SEQUENCES
]:
raise ValueError("MultiAgentReplayBuffer does not support "
"lockstep mode with storage unit `episodes`"
"or `sequences`.")
elif replay_mode in ["independent", ReplayMode.INDEPENDENT]:
self.replay_mode = ReplayMode.INDEPENDENT
else:
raise ValueError("Unsupported replay mode: {}".format(replay_mode))
if self.underlying_buffer_config:
ctor_args = {
**{
"capacity": shard_capacity,
"storage_unit": StorageUnit.FRAGMENTS
},
**self.underlying_buffer_config,
}
def new_buffer():
return from_config(self.underlying_buffer_config["type"],
ctor_args)
else:
# Default case
def new_buffer():
self.underlying_buffer_call_args = {}
return ReplayBuffer(
self.capacity,
storage_unit=StorageUnit.FRAGMENTS,
)
self.replay_buffers = collections.defaultdict(new_buffer)
# Metrics.
self.add_batch_timer = TimerStat()
self.replay_timer = TimerStat()
self._num_added = 0
def __init__(self,
capacity: int = 10000,
storage_unit: str = "timesteps",
num_shards: int = 1,
learning_starts: int = 1000,
replay_batch_size: int = 1,
prioritized_replay_alpha: float = 0.6,
prioritized_replay_beta: float = 0.4,
prioritized_replay_eps: float = 1e-6,
replay_mode: str = "independent",
replay_sequence_length: int = 1,
replay_burn_in: int = 0,
replay_zero_init_states: bool = True,
**kwargs):
"""Initializes a MultiAgentReplayBuffer instance.
Args:
num_shards: The number of buffer shards that exist in total
(including this one).
learning_starts: Number of timesteps after which a call to
`replay()` will yield samples (before that, `replay()` will
return None).
capacity: The capacity of the buffer. Note that when
`replay_sequence_length` > 1, this is the number of sequences
(not single timesteps) stored.
replay_batch_size: The batch size to be sampled (in timesteps).
Note that if `replay_sequence_length` > 1,
`self.replay_batch_size` will be set to the number of
sequences sampled (B).
prioritized_replay_alpha: Alpha parameter for a prioritized
replay buffer. Use 0.0 for no prioritization.
prioritized_replay_beta: Beta parameter for a prioritized
replay buffer.
prioritized_replay_eps: Epsilon parameter for a prioritized
replay buffer.
replay_mode: One of "independent" or "lockstep". Determined,
whether in the multiagent case, sampling is done across all
agents/policies equally.
replay_sequence_length: The sequence length (T) of a single
sample. If > 1, we will sample B x T from this buffer.
replay_burn_in: The burn-in length in case
`replay_sequence_length` > 0. This is the number of timesteps
each sequence overlaps with the previous one to generate a
better internal state (=state after the burn-in), instead of
starting from 0.0 each RNN rollout.
replay_zero_init_states: Whether the initial states in the
buffer (if replay_sequence_length > 0) are alwayas 0.0 or
should be updated with the previous train_batch state outputs.
**kwargs: Forward compatibility kwargs.
"""
shard_capacity = capacity // num_shards
ReplayBuffer.__init__(self, shard_capacity, storage_unit)
self.replay_starts = learning_starts // num_shards
self.replay_batch_size = replay_batch_size
self.prioritized_replay_beta = prioritized_replay_beta
self.prioritized_replay_eps = prioritized_replay_eps
self.replay_mode = replay_mode
self.replay_sequence_length = replay_sequence_length
self.replay_burn_in = replay_burn_in
self.replay_zero_init_states = replay_zero_init_states
if replay_sequence_length > 1:
self.replay_batch_size = int(
max(1, replay_batch_size // replay_sequence_length))
logger.info(
"Since replay_sequence_length={} and replay_batch_size={}, "
"we will replay {} sequences at a time.".format(
replay_sequence_length, replay_batch_size,
self.replay_batch_size))
if replay_mode not in ["lockstep", "independent"]:
raise ValueError("Unsupported replay mode: {}".format(replay_mode))
def new_buffer():
if prioritized_replay_alpha == 0.0:
return ReplayBuffer(self.capacity)
else:
return PrioritizedReplayBuffer(self.capacity,
alpha=prioritized_replay_alpha)
self.replay_buffers = collections.defaultdict(new_buffer)
# Metrics.
self.add_batch_timer = TimerStat()
self.replay_timer = TimerStat()
self.update_priorities_timer = TimerStat()
self._num_added = 0
# Make externally accessible for testing.
global _local_replay_buffer
_local_replay_buffer = self
# If set, return this instead of the usual data for testing.
self._fake_batch = None
