diff --git a/doc/source/rllib-training.rst b/doc/source/rllib-training.rst index f42a1cde7..57b3bb3c0 100644 --- a/doc/source/rllib-training.rst +++ b/doc/source/rllib-training.rst @@ -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 ` or :ref:`SAC `. These algorithms default to ``num_workers: 0`` for single-process operation. Consider also batch RL training with the `offline data `__ 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 ` or :ref:`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 `__ 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 `, :ref:`IMPALA `, or :ref:`APEX `. These can be scaled by increasing ``num_workers`` to add rollout workers. It may also make sense to enable `vectorization `__ 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 `, :ref:`IMPALA `, or :ref:`APEX `. These can be scaled by increasing ``num_workers`` to add rollout workers. It may also make sense to enable `vectorization `__ 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 `__. diff --git a/rllib/agents/dqn/apex.py b/rllib/agents/dqn/apex.py index 0caf90f05..463a1cd80 100644 --- a/rllib/agents/dqn/apex.py +++ b/rllib/agents/dqn/apex.py @@ -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. diff --git a/rllib/agents/dqn/dqn.py b/rllib/agents/dqn/dqn.py index 76bc21817..e08f0c169 100644 --- a/rllib/agents/dqn/dqn.py +++ b/rllib/agents/dqn/dqn.py @@ -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") diff --git a/rllib/agents/dqn/simple_q.py b/rllib/agents/dqn/simple_q.py index 152d05d96..d24bb786a 100644 --- a/rllib/agents/dqn/simple_q.py +++ b/rllib/agents/dqn/simple_q.py @@ -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") diff --git a/rllib/agents/dqn/tests/test_apex_dqn.py b/rllib/agents/dqn/tests/test_apex_dqn.py index af208ff38..04acbe595 100644 --- a/rllib/agents/dqn/tests/test_apex_dqn.py +++ b/rllib/agents/dqn/tests/test_apex_dqn.py @@ -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 diff --git a/rllib/agents/trainer.py b/rllib/agents/trainer.py index b578f3ec4..07b0ecddc 100644 --- a/rllib/agents/trainer.py +++ b/rllib/agents/trainer.py @@ -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`. diff --git a/rllib/contrib/maddpg/maddpg.py b/rllib/contrib/maddpg/maddpg.py index bdede3875..ad87212ba 100644 --- a/rllib/contrib/maddpg/maddpg.py +++ b/rllib/contrib/maddpg/maddpg.py @@ -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. diff --git a/rllib/evaluation/observation_function.py b/rllib/evaluation/observation_function.py index b1d50fadd..21c29a46f 100644 --- a/rllib/evaluation/observation_function.py +++ b/rllib/evaluation/observation_function.py @@ -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. diff --git a/rllib/evaluation/sample_batch_builder.py b/rllib/evaluation/sample_batch_builder.py index cf1409da8..e99844883 100644 --- a/rllib/evaluation/sample_batch_builder.py +++ b/rllib/evaluation/sample_batch_builder.py @@ -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): diff --git a/rllib/execution/replay_buffer.py b/rllib/execution/replay_buffer.py index a1ecab230..655190e7c 100644 --- a/rllib/execution/replay_buffer.py +++ b/rllib/execution/replay_buffer.py @@ -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: diff --git a/rllib/execution/replay_ops.py b/rllib/execution/replay_ops.py index 4cca52f6e..6a5bf96dd 100644 --- a/rllib/execution/replay_ops.py +++ b/rllib/execution/replay_ops.py @@ -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.""" diff --git a/rllib/execution/tests/test_prioritized_replay_buffer.py b/rllib/execution/tests/test_prioritized_replay_buffer.py index 38441b896..5f3f3084c 100644 --- a/rllib/execution/tests/test_prioritized_replay_buffer.py +++ b/rllib/execution/tests/test_prioritized_replay_buffer.py @@ -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 diff --git a/rllib/models/modelv2.py b/rllib/models/modelv2.py index 927baff9c..af270c0b2 100644 --- a/rllib/models/modelv2.py +++ b/rllib/models/modelv2.py @@ -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 diff --git a/rllib/policy/rnn_sequencing.py b/rllib/policy/rnn_sequencing.py index a45c84845..d38ec9158 100644 --- a/rllib/policy/rnn_sequencing.py +++ b/rllib/policy/rnn_sequencing.py @@ -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: diff --git a/rllib/policy/sample_batch.py b/rllib/policy/sample_batch.py index da5ec8334..2036a4503 100644 --- a/rllib/policy/sample_batch.py +++ b/rllib/policy/sample_batch.py @@ -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() diff --git a/rllib/tests/test_rollout_worker.py b/rllib/tests/test_rollout_worker.py index 03c97976d..09957e5ba 100644 --- a/rllib/tests/test_rollout_worker.py +++ b/rllib/tests/test_rollout_worker.py @@ -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) diff --git a/rllib/utils/debug.py b/rllib/utils/debug.py index 24a09302d..c302a3008 100644 --- a/rllib/utils/debug.py +++ b/rllib/utils/debug.py @@ -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: