wassname/ray
1
0
Fork 0
mirror of https://github.com/wassname/ray.git synced 2026-07-17 11:32:33 +08:00
Code Issues Packages Projects Releases Wiki Activity

[rllib] Flexible multi-agent replay modes and replay_sequence_length (#8893)

Browse Source
This commit is contained in:
Eric Liang
2020-06-12 20:17:27 -07:00
committed by GitHub
parent 5c56760fac
commit 34bae27ac7
17 changed files with 343 additions and 254 deletions
Download Patch File Download Diff File Expand all files Collapse all files
doc/source/rllib-training.rst
+2 -2
View File
@@ -99,10 +99,10 @@ Scaling Guide
Here are some rules of thumb for scaling training with RLlib.
1. If the environment is slow and cannot be replicated (e.g., since it requires interaction with physical systems), then you should use a sample-efficient off-policy algorithm such as :ref:`DQN <dqn>` or :ref:`SAC <sac>`. These algorithms default to ``num_workers: 0`` for single-process operation. Consider also batch RL training with the `offline data <rllib-offline.html>`__ API.
1. If the environment is slow and cannot be replicated (e.g., since it requires interaction with physical systems), then you should use a sample-efficient off-policy algorithm such as :ref:`DQN <dqn>` or :ref:`SAC <sac>`. These algorithms default to ``num_workers: 0`` for single-process operation. Make sure to set ``num_gpus: 1`` if you want to use a GPU. Consider also batch RL training with the `offline data <rllib-offline.html>`__ API.
2. If the environment is fast and the model is small (most models for RL are), use time-efficient algorithms such as :ref:`PPO <ppo>`, :ref:`IMPALA <impala>`, or :ref:`APEX <apex>`. These can be scaled by increasing ``num_workers`` to add rollout workers. It may also make sense to enable `vectorization <rllib-env.html#vectorized>`__ for inference. If the learner becomes a bottleneck, multiple GPUs can be used for learning by setting ``num_gpus > 1``.
2. If the environment is fast and the model is small (most models for RL are), use time-efficient algorithms such as :ref:`PPO <ppo>`, :ref:`IMPALA <impala>`, or :ref:`APEX <apex>`. These can be scaled by increasing ``num_workers`` to add rollout workers. It may also make sense to enable `vectorization <rllib-env.html#vectorized>`__ for inference. Make sure to set ``num_gpus: 1`` if you want to use a GPU. If the learner becomes a bottleneck, multiple GPUs can be used for learning by setting ``num_gpus > 1``.
3. If the model is compute intensive (e.g., a large deep residual network) and inference is the bottleneck, consider allocating GPUs to workers by setting ``num_gpus_per_worker: 1``. If you only have a single GPU, consider ``num_workers: 0`` to use the learner GPU for inference. For efficient use of GPU time, use a small number of GPU workers and a large number of `envs per worker <rllib-env.html#vectorized>`__.
rllib/agents/dqn/apex.py
+2
View File
@@ -86,6 +86,8 @@ def apex_execution_plan(workers: WorkerSet, config: dict):
config["prioritized_replay_alpha"],
config["prioritized_replay_beta"],
config["prioritized_replay_eps"],
config["multiagent"]["replay_mode"],
config["replay_sequence_length"],
], num_replay_buffer_shards)
# Start the learner thread.
rllib/agents/dqn/dqn.py
+10 -4
View File
@@ -83,9 +83,6 @@ DEFAULT_CONFIG = with_common_config({
"prioritized_replay_eps": 1e-6,
# Whether to LZ4 compress observations
"compress_observations": False,
# In multi-agent mode, whether to replay experiences from the same time
# step for all policies. This is required for MADDPG.
"multiagent_sync_replay": False,
# Callback to run before learning on a multi-agent batch of experiences.
"before_learn_on_batch": None,
# If set, this will fix the ratio of sampled to replayed timesteps.
@@ -227,6 +224,14 @@ def validate_config(config):
config.get("n_step", 1))
config["rollout_fragment_length"] = adjusted_batch_size
if config.get("prioritized_replay"):
if config["multiagent"]["replay_mode"] == "lockstep":
raise ValueError("Prioritized replay is not supported when "
"replay_mode=lockstep.")
elif config["replay_sequence_length"] > 1:
raise ValueError("Prioritized replay is not supported when "
"replay_sequence_length > 1.")
def execution_plan(workers, config):
if config.get("prioritized_replay"):
@@ -243,7 +248,8 @@ def execution_plan(workers, config):
learning_starts=config["learning_starts"],
buffer_size=config["buffer_size"],
replay_batch_size=config["train_batch_size"],
multiagent_sync_replay=config.get("multiagent_sync_replay"),
replay_mode=config["multiagent"]["replay_mode"],
replay_sequence_length=config["replay_sequence_length"],
**prio_args)
rollouts = ParallelRollouts(workers, mode="bulk_sync")
rllib/agents/dqn/simple_q.py
+3 -1
View File
@@ -92,7 +92,9 @@ def execution_plan(workers, config):
num_shards=1,
learning_starts=config["learning_starts"],
buffer_size=config["buffer_size"],
replay_batch_size=config["train_batch_size"])
replay_batch_size=config["train_batch_size"],
replay_mode=config["multiagent"]["replay_mode"],
replay_sequence_length=config["replay_sequence_length"])
rollouts = ParallelRollouts(workers, mode="bulk_sync")
rllib/agents/dqn/tests/test_apex_dqn.py
+2
View File
@@ -17,6 +17,7 @@ class TestApexDQN(unittest.TestCase):
def test_apex_zero_workers(self):
config = apex.APEX_DEFAULT_CONFIG.copy()
config["num_workers"] = 0
config["learning_starts"] = 1000
config["prioritized_replay"] = True
config["timesteps_per_iteration"] = 100
config["min_iter_time_s"] = 1
@@ -30,6 +31,7 @@ class TestApexDQN(unittest.TestCase):
"""Test whether an APEX-DQNTrainer can be built on all frameworks."""
config = apex.APEX_DEFAULT_CONFIG.copy()
config["num_workers"] = 3
config["learning_starts"] = 1000
config["prioritized_replay"] = True
config["timesteps_per_iteration"] = 100
config["min_iter_time_s"] = 1
rllib/agents/trainer.py
+11
View File
@@ -347,8 +347,19 @@ COMMON_CONFIG = {
# observations to include more state.
# See rllib/evaluation/observation_function.py for more info.
"observation_fn": None,
# When replay_mode=lockstep, RLlib will replay all the agent
# transitions at a particular timestep together in a batch. This allows
# the policy to implement differentiable shared computations between
# agents it controls at that timestep. When replay_mode=independent,
# transitions are replayed independently per policy.
"replay_mode": "independent",
},
# === Replay Settings ===
# The number of contiguous environment steps to replay at once. This may
# be set to greater than 1 to support recurrent models.
"replay_sequence_length": 1,
# Deprecated keys:
"use_pytorch": DEPRECATED_VALUE, # Replaced by `framework=torch`.
"eager": DEPRECATED_VALUE, # Replaced by `framework=tfe`.
rllib/contrib/maddpg/maddpg.py
+6 -4
View File
@@ -11,10 +11,11 @@ with the multi-agent particle envs.
import logging
from ray.rllib.agents.trainer import with_common_config
from ray.rllib.agents.trainer import COMMON_CONFIG, with_common_config
from ray.rllib.agents.dqn.dqn import GenericOffPolicyTrainer
from ray.rllib.contrib.maddpg.maddpg_policy import MADDPGTFPolicy
from ray.rllib.policy.sample_batch import SampleBatch, MultiAgentBatch
from ray.rllib.utils import merge_dicts
logger = logging.getLogger(__name__)
logger.setLevel(logging.INFO)
@@ -66,13 +67,14 @@ DEFAULT_CONFIG = with_common_config({
# Observation compression. Note that compression makes simulation slow in
# MPE.
"compress_observations": False,
# In multi-agent mode, whether to replay experiences from the same time
# step for all policies. This is required for MADDPG.
"multiagent_sync_replay": True,
# If set, this will fix the ratio of sampled to replayed timesteps.
# Otherwise, replay will proceed at the native ratio determined by
# (train_batch_size / rollout_fragment_length).
"training_intensity": None,
# Force lockstep replay mode for MADDPG.
"multiagent": merge_dicts(COMMON_CONFIG["multiagent"], {
"replay_mode": "lockstep",
}),
# === Optimization ===
# Learning rate for the critic (Q-function) optimizer.
rllib/evaluation/observation_function.py
+1 -1
View File
@@ -12,7 +12,7 @@ class ObservationFunction:
These callbacks can be used for preprocessing of observations, especially
in multi-agent scenarios.
Observations functions can be specified in the multi-agent config by
Observation functions can be specified in the multi-agent config by
specifying ``{"observation_function": your_obs_func}``. Note that
``your_obs_func`` can be a plain Python function.
rllib/evaluation/sample_batch_builder.py
+12 -4
View File
@@ -5,6 +5,7 @@ import numpy as np
from ray.rllib.policy.sample_batch import SampleBatch, MultiAgentBatch
from ray.rllib.utils.annotations import PublicAPI, DeveloperAPI
from ray.rllib.utils.debug import summarize
from ray.rllib.env.base_env import _DUMMY_AGENT_ID
from ray.util.debug import log_once
logger = logging.getLogger(__name__)
@@ -25,11 +26,12 @@ class SampleBatchBuilder:
However, it is useful to add data one row (dict) at a time.
"""
_next_unroll_id = 0 # disambiguates unrolls within a single episode
@PublicAPI
def __init__(self):
self.buffers = collections.defaultdict(list)
self.count = 0
self.unroll_id = 0 # disambiguates unrolls within a single episode
@PublicAPI
def add_values(self, **values):
@@ -54,11 +56,12 @@ class SampleBatchBuilder:
batch = SampleBatch(
{k: to_float_array(v)
for k, v in self.buffers.items()})
batch.data[SampleBatch.UNROLL_ID] = np.repeat(self.unroll_id,
batch.count)
if SampleBatch.UNROLL_ID not in batch.data:
batch.data[SampleBatch.UNROLL_ID] = np.repeat(
SampleBatchBuilder._next_unroll_id, batch.count)
SampleBatchBuilder._next_unroll_id += 1
self.buffers.clear()
self.count = 0
self.unroll_id += 1
return batch
@@ -132,6 +135,11 @@ class MultiAgentSampleBatchBuilder:
if agent_id not in self.agent_builders:
self.agent_builders[agent_id] = SampleBatchBuilder()
self.agent_to_policy[agent_id] = policy_id
# Include the current agent id for multi-agent algorithms.
if agent_id != _DUMMY_AGENT_ID:
values["agent_id"] = agent_id
self.agent_builders[agent_id].add_values(**values)
def postprocess_batch_so_far(self, episode=None):
rllib/execution/replay_buffer.py
+106 -169
View File
@@ -1,31 +1,34 @@
import numpy as np
import random
import collections
import logging
import numpy as np
import platform
import sys
import random
from typing import List
import ray
from ray.rllib.execution.common import SampleBatchType
from ray.rllib.execution.segment_tree import SumSegmentTree, MinSegmentTree
from ray.rllib.policy.sample_batch import SampleBatch, DEFAULT_POLICY_ID, \
MultiAgentBatch
from ray.rllib.policy.sample_batch import SampleBatch, MultiAgentBatch, \
DEFAULT_POLICY_ID
from ray.rllib.utils.annotations import DeveloperAPI
from ray.rllib.utils.compression import unpack_if_needed
from ray.util.iter import ParallelIteratorWorker
from ray.rllib.utils.timer import TimerStat
from ray.rllib.utils.window_stat import WindowStat
# Constant that represents all policies in lockstep replay mode.
_ALL_POLICIES = "__all__"
logger = logging.getLogger(__name__)
@DeveloperAPI
class ReplayBuffer:
@DeveloperAPI
def __init__(self, size):
def __init__(self, size: int):
"""Create Prioritized Replay buffer.
Parameters
----------
size: int
Max number of transitions to store in the buffer. When the buffer
overflows the old memories are dropped.
Args:
size (int): Max number of items to store in the FIFO buffer.
"""
self._storage = []
self._maxsize = size
@@ -41,15 +44,15 @@ class ReplayBuffer:
return len(self._storage)
@DeveloperAPI
def add(self, obs_t, action, reward, obs_tp1, done, weight):
data = (obs_t, action, reward, obs_tp1, done)
def add(self, item: SampleBatchType, weight: float):
assert item.count > 0, item
self._num_added += 1
if self._next_idx >= len(self._storage):
self._storage.append(data)
self._est_size_bytes += sum(sys.getsizeof(d) for d in data)
self._storage.append(item)
self._est_size_bytes += item.size_bytes()
else:
self._storage[self._next_idx] = data
self._storage[self._next_idx] = item
if self._next_idx + 1 >= self._maxsize:
self._eviction_started = True
self._next_idx = (self._next_idx + 1) % self._maxsize
@@ -57,57 +60,26 @@ class ReplayBuffer:
self._evicted_hit_stats.push(self._hit_count[self._next_idx])
self._hit_count[self._next_idx] = 0
def _encode_sample(self, idxes):
obses_t, actions, rewards, obses_tp1, dones = [], [], [], [], []
for i in idxes:
data = self._storage[i]
obs_t, action, reward, obs_tp1, done = data
obses_t.append(np.array(unpack_if_needed(obs_t), copy=False))
actions.append(np.array(action, copy=False))
rewards.append(reward)
obses_tp1.append(np.array(unpack_if_needed(obs_tp1), copy=False))
dones.append(done)
self._hit_count[i] += 1
return (np.array(obses_t), np.array(actions), np.array(rewards),
np.array(obses_tp1), np.array(dones))
def _encode_sample(self, idxes: List[int]) -> SampleBatchType:
out = SampleBatch.concat_samples([self._storage[i] for i in idxes])
out.decompress_if_needed()
return out
@DeveloperAPI
def sample_idxes(self, batch_size):
return np.random.randint(0, len(self._storage), batch_size)
@DeveloperAPI
def sample_with_idxes(self, idxes):
self._num_sampled += len(idxes)
return self._encode_sample(idxes)
@DeveloperAPI
def sample(self, batch_size):
def sample(self, num_items: int) -> SampleBatchType:
"""Sample a batch of experiences.
Parameters
----------
batch_size: int
How many transitions to sample.
Args:
num_items (int): Number of items to sample from this buffer.
Returns
-------
obs_batch: np.array
batch of observations
act_batch: np.array
batch of actions executed given obs_batch
rew_batch: np.array
rewards received as results of executing act_batch
next_obs_batch: np.array
next set of observations seen after executing act_batch
done_mask: np.array
done_mask[i] = 1 if executing act_batch[i] resulted in
the end of an episode and 0 otherwise.
Returns:
SampleBatchType: concatenated batch of items.
"""
idxes = [
random.randint(0,
len(self._storage) - 1) for _ in range(batch_size)
len(self._storage) - 1) for _ in range(num_items)
]
self._num_sampled += batch_size
self._num_sampled += num_items
return self._encode_sample(idxes)
@DeveloperAPI
@@ -126,21 +98,16 @@ class ReplayBuffer:
@DeveloperAPI
class PrioritizedReplayBuffer(ReplayBuffer):
@DeveloperAPI
def __init__(self, size, alpha):
def __init__(self, size: int, alpha: float):
"""Create Prioritized Replay buffer.
Parameters
----------
size: int
Max number of transitions to store in the buffer. When the buffer
overflows the old memories are dropped.
alpha: float
how much prioritization is used
(0 - no prioritization, 1 - full prioritization)
Args:
size (int): Max number of items to store in the FIFO buffer.
alpha (float): how much prioritization is used
(0 - no prioritization, 1 - full prioritization).
See Also
--------
ReplayBuffer.__init__
See also:
ReplayBuffer.__init__()
"""
super(PrioritizedReplayBuffer, self).__init__(size)
assert alpha > 0
@@ -156,20 +123,17 @@ class PrioritizedReplayBuffer(ReplayBuffer):
self._prio_change_stats = WindowStat("reprio", 1000)
@DeveloperAPI
def add(self, obs_t, action, reward, obs_tp1, done, weight):
"""See ReplayBuffer.store_effect"""
def add(self, item: SampleBatchType, weight: float):
idx = self._next_idx
super(PrioritizedReplayBuffer, self).add(obs_t, action, reward,
obs_tp1, done, weight)
super(PrioritizedReplayBuffer, self).add(item, weight)
if weight is None:
weight = self._max_priority
self._it_sum[idx] = weight**self._alpha
self._it_min[idx] = weight**self._alpha
def _sample_proportional(self, batch_size):
def _sample_proportional(self, num_items: int):
res = []
for _ in range(batch_size):
for _ in range(num_items):
# TODO(szymon): should we ensure no repeats?
mass = random.random() * self._it_sum.sum(0, len(self._storage))
idx = self._it_sum.find_prefixsum_idx(mass)
@@ -177,79 +141,45 @@ class PrioritizedReplayBuffer(ReplayBuffer):
return res
@DeveloperAPI
def sample_idxes(self, batch_size):
return self._sample_proportional(batch_size)
def sample(self, num_items: int, beta: float) -> SampleBatchType:
"""Sample a batch of experiences and return priority weights, indices.
@DeveloperAPI
def sample_with_idxes(self, idxes, beta):
assert beta > 0
self._num_sampled += len(idxes)
Args:
num_items (int): Number of items to sample from this buffer.
beta (float): To what degree to use importance weights
(0 - no corrections, 1 - full correction).
weights = []
p_min = self._it_min.min() / self._it_sum.sum()
max_weight = (p_min * len(self._storage))**(-beta)
for idx in idxes:
p_sample = self._it_sum[idx] / self._it_sum.sum()
weight = (p_sample * len(self._storage))**(-beta)
weights.append(weight / max_weight)
weights = np.array(weights)
encoded_sample = self._encode_sample(idxes)
return tuple(list(encoded_sample) + [weights, idxes])
@DeveloperAPI
def sample(self, batch_size, beta):
"""Sample a batch of experiences.
compared to ReplayBuffer.sample
it also returns importance weights and idxes
of sampled experiences.
Parameters
----------
batch_size: int
How many transitions to sample.
beta: float
To what degree to use importance weights
(0 - no corrections, 1 - full correction)
Returns
-------
obs_batch: np.array
batch of observations
act_batch: np.array
batch of actions executed given obs_batch
rew_batch: np.array
rewards received as results of executing act_batch
next_obs_batch: np.array
next set of observations seen after executing act_batch
done_mask: np.array
done_mask[i] = 1 if executing act_batch[i] resulted in
the end of an episode and 0 otherwise.
weights: np.array
Array of shape (batch_size,) and dtype np.float32
denoting importance weight of each sampled transition
idxes: np.array
Array of shape (batch_size,) and dtype np.int32
idexes in buffer of sampled experiences
Returns:
SampleBatchType: Concatenated batch of items including "weights"
and "batch_indexes" fields denoting IS of each sampled
transition and original idxes in buffer of sampled experiences.
"""
assert beta >= 0.0
self._num_sampled += batch_size
self._num_sampled += num_items
idxes = self._sample_proportional(batch_size)
idxes = self._sample_proportional(num_items)
weights = []
batch_indexes = []
p_min = self._it_min.min() / self._it_sum.sum()
max_weight = (p_min * len(self._storage))**(-beta)
for idx in idxes:
p_sample = self._it_sum[idx] / self._it_sum.sum()
weight = (p_sample * len(self._storage))**(-beta)
weights.append(weight / max_weight)
weights = np.array(weights)
encoded_sample = self._encode_sample(idxes)
return tuple(list(encoded_sample) + [weights, idxes])
count = self._storage[idx].count
weights.extend([weight / max_weight] * count)
batch_indexes.extend([idx] * count)
batch = self._encode_sample(idxes)
# Note: prioritization is not supported in lockstep replay mode.
if isinstance(batch, SampleBatch):
assert len(weights) == batch.count
assert len(batch_indexes) == batch.count
batch["weights"] = np.array(weights)
batch["batch_indexes"] = np.array(batch_indexes)
return batch
@DeveloperAPI
def update_priorities(self, idxes, priorities):
@@ -290,7 +220,6 @@ class PrioritizedReplayBuffer(ReplayBuffer):
_local_replay_buffer = None
# TODO(ekl) move this class to common
class LocalReplayBuffer(ParallelIteratorWorker):
"""A replay buffer shard.
@@ -305,13 +234,27 @@ class LocalReplayBuffer(ParallelIteratorWorker):
prioritized_replay_alpha=0.6,
prioritized_replay_beta=0.4,
prioritized_replay_eps=1e-6,
multiagent_sync_replay=False):
replay_mode="independent",
replay_sequence_length=1):
self.replay_starts = learning_starts // num_shards
self.buffer_size = buffer_size // num_shards
self.replay_batch_size = replay_batch_size
self.prioritized_replay_beta = prioritized_replay_beta
self.prioritized_replay_eps = prioritized_replay_eps
self.multiagent_sync_replay = multiagent_sync_replay
self.replay_mode = replay_mode
self.replay_sequence_length = replay_sequence_length
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 gen_replay():
while True:
@@ -352,12 +295,18 @@ class LocalReplayBuffer(ParallelIteratorWorker):
if isinstance(batch, SampleBatch):
batch = MultiAgentBatch({DEFAULT_POLICY_ID: batch}, batch.count)
with self.add_batch_timer:
for policy_id, s in batch.policy_batches.items():
for row in s.rows():
self.replay_buffers[policy_id].add(
row["obs"], row["actions"], row["rewards"],
row["new_obs"], row["dones"], row["weights"]
if "weights" in row else None)
if self.replay_mode == "lockstep":
# Note that prioritization is not supported in this mode.
for s in batch.timeslices(self.replay_sequence_length):
self.replay_buffers[_ALL_POLICIES].add(s, weight=None)
else:
for policy_id, b in batch.policy_batches.items():
for s in b.timeslices(self.replay_sequence_length):
if "weights" in s:
weight = np.mean(s["weights"])
else:
weight = None
self.replay_buffers[policy_id].add(s, weight=weight)
self.num_added += batch.count
def replay(self):
@@ -371,28 +320,16 @@ class LocalReplayBuffer(ParallelIteratorWorker):
return None
with self.replay_timer:
samples = {}
idxes = None
for policy_id, replay_buffer in self.replay_buffers.items():
if self.multiagent_sync_replay:
if idxes is None:
idxes = replay_buffer.sample_idxes(
self.replay_batch_size)
else:
idxes = replay_buffer.sample_idxes(self.replay_batch_size)
(obses_t, actions, rewards, obses_tp1, dones, weights,
batch_indexes) = replay_buffer.sample_with_idxes(
idxes, beta=self.prioritized_replay_beta)
samples[policy_id] = SampleBatch({
"obs": obses_t,
"actions": actions,
"rewards": rewards,
"new_obs": obses_tp1,
"dones": dones,
"weights": weights,
"batch_indexes": batch_indexes
})
return MultiAgentBatch(samples, self.replay_batch_size)
if self.replay_mode == "lockstep":
return self.replay_buffers[_ALL_POLICIES].sample(
self.replay_batch_size, beta=self.prioritized_replay_beta)
else:
samples = {}
for policy_id, replay_buffer in self.replay_buffers.items():
samples[policy_id] = replay_buffer.sample(
self.replay_batch_size,
beta=self.prioritized_replay_beta)
return MultiAgentBatch(samples, self.replay_batch_size)
def update_priorities(self, prio_dict):
with self.update_priorities_timer:
rllib/execution/replay_ops.py
+1
View File
@@ -106,6 +106,7 @@ class WaitUntilTimestepsElapsed:
return ts > self.target_num_timesteps
# TODO(ekl) deprecate this in favor of the replay_sequence_length option.
class SimpleReplayBuffer:
"""Simple replay buffer that operates over batches."""
rllib/execution/tests/test_prioritized_replay_buffer.py
+31 -25
View File
@@ -3,6 +3,7 @@ import numpy as np
import unittest
from ray.rllib.execution.replay_buffer import PrioritizedReplayBuffer
from ray.rllib.policy.sample_batch import SampleBatch
from ray.rllib.utils.test_utils import check
@@ -17,13 +18,13 @@ class TestPrioritizedReplayBuffer(unittest.TestCase):
max_priority = 1.0
def _generate_data(self):
return (
np.random.random((4, )), # obs_t
np.random.choice([0, 1]), # action
np.random.rand(), # reward
np.random.random((4, )), # obs_tp1
np.random.choice([False, True]), # done
)
return SampleBatch({
"obs_t": [np.random.random((4, ))],
"action": [np.random.choice([0, 1])],
"reward": [np.random.rand()],
"obs_tp1": [np.random.random((4, ))],
"done": [np.random.choice([False, True])],
})
def test_add(self):
memory = PrioritizedReplayBuffer(
@@ -37,19 +38,19 @@ class TestPrioritizedReplayBuffer(unittest.TestCase):
# Insert single record.
data = self._generate_data()
memory.add(*data, weight=0.5)
memory.add(data, weight=0.5)
self.assertTrue(len(memory) == 1)
self.assertTrue(memory._next_idx == 1)
# Insert single record.
data = self._generate_data()
memory.add(*data, weight=0.1)
memory.add(data, weight=0.1)
self.assertTrue(len(memory) == 2)
self.assertTrue(memory._next_idx == 0)
# Insert over capacity.
data = self._generate_data()
memory.add(*data, weight=1.0)
memory.add(data, weight=1.0)
self.assertTrue(len(memory) == 2)
self.assertTrue(memory._next_idx == 1)
@@ -60,13 +61,14 @@ class TestPrioritizedReplayBuffer(unittest.TestCase):
num_records = 5
for i in range(num_records):
data = self._generate_data()
memory.add(*data, weight=1.0)
memory.add(data, weight=1.0)
self.assertTrue(len(memory) == i + 1)
self.assertTrue(memory._next_idx == i + 1)
# Fetch records, their indices and weights.
_, _, _, _, _, weights, indices = \
memory.sample(3, beta=self.beta)
batch = memory.sample(3, beta=self.beta)
weights = batch["weights"]
indices = batch["batch_indexes"]
check(weights, np.ones(shape=(3, )))
self.assertEqual(3, len(indices))
self.assertTrue(len(memory) == num_records)
@@ -78,8 +80,8 @@ class TestPrioritizedReplayBuffer(unittest.TestCase):
# Expect to sample almost only index 1
# (which still has a weight of 1.0).
for _ in range(10):
_, _, _, _, _, weights, indices = memory.sample(
1000, beta=self.beta)
batch = memory.sample(1000, beta=self.beta)
indices = batch["batch_indexes"]
self.assertTrue(970 < np.sum(indices) < 1100)
# Update weight of indices 0 and 1 to >> 0.01.
@@ -87,7 +89,8 @@ class TestPrioritizedReplayBuffer(unittest.TestCase):
for _ in range(10):
rand = np.random.random() + 0.2
memory.update_priorities(np.array([0, 1]), np.array([rand, rand]))
_, _, _, _, _, _, indices = memory.sample(1000, beta=self.beta)
batch = memory.sample(1000, beta=self.beta)
indices = batch["batch_indexes"]
# Expect biased to higher values due to some 2s, 3s, and 4s.
# print(np.sum(indices))
self.assertTrue(400 < np.sum(indices) < 800)
@@ -99,7 +102,8 @@ class TestPrioritizedReplayBuffer(unittest.TestCase):
rand = np.random.random() + 0.2
memory.update_priorities(
np.array([0, 1]), np.array([rand, rand * 2]))
_, _, _, _, _, _, indices = memory.sample(1000, beta=self.beta)
batch = memory.sample(1000, beta=self.beta)
indices = batch["batch_indexes"]
# print(np.sum(indices))
self.assertTrue(600 < np.sum(indices) < 850)
@@ -110,7 +114,8 @@ class TestPrioritizedReplayBuffer(unittest.TestCase):
rand = np.random.random() + 0.2
memory.update_priorities(
np.array([0, 1]), np.array([rand, rand * 4]))
_, _, _, _, _, _, indices = memory.sample(1000, beta=self.beta)
batch = memory.sample(1000, beta=self.beta)
indices = batch["batch_indexes"]
# print(np.sum(indices))
self.assertTrue(750 < np.sum(indices) < 950)
@@ -121,7 +126,8 @@ class TestPrioritizedReplayBuffer(unittest.TestCase):
rand = np.random.random() + 0.2
memory.update_priorities(
np.array([0, 1]), np.array([rand, rand * 9]))
_, _, _, _, _, _, indices = memory.sample(1000, beta=self.beta)
batch = memory.sample(1000, beta=self.beta)
indices = batch["batch_indexes"]
# print(np.sum(indices))
self.assertTrue(850 < np.sum(indices) < 1100)
@@ -129,7 +135,7 @@ class TestPrioritizedReplayBuffer(unittest.TestCase):
num_records = 5
for i in range(num_records):
data = self._generate_data()
memory.add(*data, weight=1.0)
memory.add(data, weight=1.0)
self.assertTrue(len(memory) == i + 6)
self.assertTrue(memory._next_idx == (i + 6) % self.capacity)
@@ -139,8 +145,8 @@ class TestPrioritizedReplayBuffer(unittest.TestCase):
np.array([0.001, 0.1, 2., 8., 16., 32., 64., 128., 256., 512.]))
counts = Counter()
for _ in range(10):
_, _, _, _, _, _, indices = memory.sample(
np.random.randint(100, 600), beta=self.beta)
batch = memory.sample(np.random.randint(100, 600), beta=self.beta)
indices = batch["batch_indexes"]
for i in indices:
counts[i] += 1
print(counts)
@@ -158,13 +164,13 @@ class TestPrioritizedReplayBuffer(unittest.TestCase):
num_records = 5
for i in range(num_records):
data = self._generate_data()
memory.add(*data, weight=np.random.rand())
memory.add(data, weight=np.random.rand())
self.assertTrue(len(memory) == i + 1)
self.assertTrue(memory._next_idx == i + 1)
# Fetch records, their indices and weights.
_, _, _, _, _, weights, indices = \
memory.sample(1000, beta=self.beta)
batch = memory.sample(1000, beta=self.beta)
indices = batch["batch_indexes"]
counts = Counter()
for i in indices:
counts[i] += 1
rllib/models/modelv2.py
+2 -1
View File
@@ -82,7 +82,8 @@ class ModelV2:
Args:
input_dict (dict): dictionary of input tensors, including "obs",
"obs_flat", "prev_action", "prev_reward", "is_training"
"obs_flat", "prev_action", "prev_reward", "is_training",
"eps_id", "agent_id", "infos", and "t".
state (list): list of state tensors with sizes matching those
returned by get_initial_state + the batch dimension
seq_lens (Tensor): 1d tensor holding input sequence lengths
rllib/policy/rnn_sequencing.py
+5 -1
View File
@@ -223,7 +223,11 @@ def chop_into_sequences(episode_ids,
feature_sequences = []
for f in feature_columns:
f = np.array(f)
f_pad = np.zeros((len(seq_lens) * max_seq_len, ) + np.shape(f)[1:])
length = len(seq_lens) * max_seq_len
if f.dtype == np.object or f.dtype.type is np.str_:
f_pad = [None] * length
else:
f_pad = np.zeros((length, ) + np.shape(f)[1:])
seq_base = 0
i = 0
for l in seq_lens:
rllib/policy/sample_batch.py
+145 -40
View File
@@ -1,13 +1,20 @@
import collections
import numpy as np
import sys
import itertools
from typing import Dict, List, Any
from ray.rllib.utils.annotations import PublicAPI, DeveloperAPI
from ray.rllib.utils.compression import pack, unpack, is_compressed
from ray.rllib.utils.memory import concat_aligned
from ray.rllib.utils.deprecation import deprecation_warning
# Default policy id for single agent environments
DEFAULT_POLICY_ID = "default_policy"
# TODO(ekl) reuse the other id def once we fix imports
PolicyID = Any
@PublicAPI
class SampleBatch:
@@ -195,6 +202,15 @@ class SampleBatch:
return SampleBatch({k: v[start:end] for k, v in self.data.items()})
@PublicAPI
def timeslices(self, k: int) -> List["SampleBatch"]:
out = []
i = 0
while i < self.count:
out.append(self.slice(i, i + k))
i += k
return out
@PublicAPI
def keys(self):
return self.data.keys()
@@ -207,6 +223,10 @@ class SampleBatch:
def get(self, key):
return self.data.get(key)
@PublicAPI
def size_bytes(self) -> int:
return sum(sys.getsizeof(d) for d in self.data)
@PublicAPI
def __getitem__(self, key):
return self.data[key]
@@ -252,45 +272,133 @@ class SampleBatch:
@PublicAPI
class MultiAgentBatch:
"""A batch of experiences from multiple policies in the environment.
"""
"""A batch of experiences from multiple agents in the environment."""
@PublicAPI
def __init__(self, policy_batches, count):
"""Initializes a MultiAgentBatch object.
def __init__(self, policy_batches: Dict[PolicyID, SampleBatch],
env_steps: int):
"""Initialize a MultiAgentBatch object.
Args:
policy_batches (Dict[str,SampleBatch]): Mapping from policy id
(str) to a SampleBatch of experiences. Note that these batches
may be of different length.
count (int): The number of timesteps in the environment this batch
contains. This will be less than the number of transitions this
batch contains across all policies in total.
policy_batches (Dict[PolicyID, SampleBatch]): Mapping from policy
ids to SampleBatches of experiences.
env_steps (int): The number of timesteps in the environment this
batch contains. This will be less than the number of
transitions this batch contains across all policies in total.
Attributes:
policy_batches (Dict[PolicyID, SampleBatch]): Mapping from policy
ids to SampleBatches of experiences.
count (int): the number of env steps in this batch.
"""
for v in policy_batches.values():
assert isinstance(v, SampleBatch)
self.policy_batches = policy_batches
self.count = count
# Called count for uniformity with SampleBatch. Prefer to access this
# via the env_steps() method when possible for clarity.
self.count = env_steps
@PublicAPI
def env_steps(self) -> int:
"""The number of env steps (there are >= 1 agent steps per env step).
Returns:
int: the number of environment steps contained in this batch.
"""
return self.count
@PublicAPI
def agent_steps(self) -> int:
"""The number of agent steps (there are >= 1 agent steps per env step).
Returns:
int: the number of agent steps total in this batch.
"""
ct = 0
for batch in self.policy_batches.values():
ct += batch.count
return ct
@PublicAPI
def timeslices(self, k: int) -> List["MultiAgentBatch"]:
"""Returns k-step batches holding data for each agent at those steps.
For examples, suppose we have agent1 observations [a1t1, a1t2, a1t3],
for agent2, [a2t1, a2t3], and for agent3, [a3t3] only.
Calling timeslices(1) would return three MultiAgentBatches containing
[a1t1, a2t1], [a1t2], and [a1t3, a2t3, a3t3].
Calling timeslices(2) would return two MultiAgentBatches containing
[a1t1, a1t2, a2t1], and [a1t3, a2t3, a3t3].
This method is used to implement "lockstep" replay mode. Note that this
method does not guarantee each batch contains only data from a single
unroll. Batches might contain data from multiple different envs.
"""
from ray.rllib.evaluation.sample_batch_builder import \
SampleBatchBuilder
# Build a sorted set of (eps_id, t, policy_id, data...)
steps = []
for policy_id, batch in self.policy_batches.items():
for row in batch.rows():
steps.append((row[SampleBatch.EPS_ID], row["t"], policy_id,
row))
steps.sort()
finished_slices = []
cur_slice = collections.defaultdict(SampleBatchBuilder)
cur_slice_size = 0
def finish_slice():
nonlocal cur_slice_size
assert cur_slice_size > 0
batch = MultiAgentBatch(
{k: v.build_and_reset()
for k, v in cur_slice.items()}, cur_slice_size)
cur_slice_size = 0
finished_slices.append(batch)
# For each unique env timestep.
for _, group in itertools.groupby(steps, lambda x: x[:2]):
# Accumulate into the current slice.
for _, _, policy_id, row in group:
cur_slice[policy_id].add_values(**row)
cur_slice_size += 1
# Slice has reached target number of env steps.
if cur_slice_size >= k:
finish_slice()
assert cur_slice_size == 0
if cur_slice_size > 0:
finish_slice()
assert len(finished_slices) > 0, finished_slices
return finished_slices
@staticmethod
@PublicAPI
def wrap_as_needed(batches, count):
def wrap_as_needed(policy_batches: Dict[PolicyID, SampleBatch],
env_steps: int) -> Any:
"""Returns SampleBatch or MultiAgentBatch, depending on given policies.
Args:
batches (Dict[str,SampleBatch]): Mapping from policy ID to
SampleBatch.
count (int): A count to use, when returning a MultiAgentBatch.
policy_batches (Dict[PolicyID, SampleBatch]): Mapping from policy
ids to SampleBatch.
env_steps (int): Number of env steps in the batch.
Returns:
Union[SampleBatch,MultiAgentBatch]: The single default policy's
Union[SampleBatch, MultiAgentBatch]: The single default policy's
SampleBatch or a MultiAgentBatch (more than one policy).
"""
if len(batches) == 1 and DEFAULT_POLICY_ID in batches:
return batches[DEFAULT_POLICY_ID]
return MultiAgentBatch(batches, count)
if len(policy_batches) == 1 and DEFAULT_POLICY_ID in policy_batches:
return policy_batches[DEFAULT_POLICY_ID]
return MultiAgentBatch(policy_batches, env_steps)
@staticmethod
@PublicAPI
def concat_samples(samples):
def concat_samples(samples: List["MultiAgentBatch"]) -> "MultiAgentBatch":
"""Concatenates a list of MultiAgentBatches into a new MultiAgentBatch.
Args:
@@ -302,22 +410,22 @@ class MultiAgentBatch:
concatenated inputs.
"""
policy_batches = collections.defaultdict(list)
total_count = 0
env_steps = 0
for s in samples:
if not isinstance(s, MultiAgentBatch):
raise ValueError(
"`MultiAgentBatch.concat_samples()` can only concat "
"MultiAgentBatch types, not {}!".format(type(s).__name__))
for policy_id, batch in s.policy_batches.items():
policy_batches[policy_id].append(batch)
total_count += s.count
for key, batch in s.policy_batches.items():
policy_batches[key].append(batch)
env_steps += s.env_steps()
out = {}
for policy_id, batches in policy_batches.items():
out[policy_id] = SampleBatch.concat_samples(batches)
return MultiAgentBatch(out, total_count)
for key, batches in policy_batches.items():
out[key] = SampleBatch.concat_samples(batches)
return MultiAgentBatch(out, env_steps)
@PublicAPI
def copy(self):
def copy(self) -> "MultiAgentBatch":
"""Deep-copies self into a new MultiAgentBatch.
Returns:
@@ -328,16 +436,8 @@ class MultiAgentBatch:
for (k, v) in self.policy_batches.items()}, self.count)
@PublicAPI
def total(self):
"""Calculates the sum of all step-counts over all policy batches.
Returns:
int: The sum of counts over all policy batches.
"""
ct = 0
for batch in self.policy_batches.values():
ct += batch.count
return ct
def size_bytes(self) -> int:
return sum(b.size_bytes() for b in self.policy_batches.values())
@DeveloperAPI
def compress(self, bulk=False, columns=frozenset(["obs", "new_obs"])):
@@ -364,9 +464,14 @@ class MultiAgentBatch:
return self
def __str__(self):
return "MultiAgentBatch({}, count={})".format(
return "MultiAgentBatch({}, env_steps={})".format(
str(self.policy_batches), self.count)
def __repr__(self):
return "MultiAgentBatch({}, count={})".format(
return "MultiAgentBatch({}, env_steps={})".format(
str(self.policy_batches), self.count)
# Deprecated.
def total(self):
deprecation_warning("batch.total()", "batch.agent_steps()")
return self.agent_steps()
rllib/tests/test_rollout_worker.py
+3 -1
View File
@@ -161,7 +161,9 @@ class TestRolloutWorker(unittest.TestCase):
def test_batch_ids(self):
ev = RolloutWorker(
env_creator=lambda _: gym.make("CartPole-v0"), policy=MockPolicy)
env_creator=lambda _: gym.make("CartPole-v0"),
policy=MockPolicy,
rollout_fragment_length=1)
batch1 = ev.sample()
batch2 = ev.sample()
self.assertEqual(len(set(batch1["unroll_id"])), 1)
rllib/utils/debug.py
+1 -1
View File
@@ -33,7 +33,7 @@ def _summarize(obj):
if obj.size == 0:
return _StringValue("np.ndarray({}, dtype={})".format(
obj.shape, obj.dtype))
elif obj.dtype == np.object:
elif obj.dtype == np.object or obj.dtype.type is np.str_:
return _StringValue("np.ndarray({}, dtype={}, head={})".format(
obj.shape, obj.dtype, _summarize(obj[0])))
else: