From dab241dcc6c10243f7427b64abe34a71865ca931 Mon Sep 17 00:00:00 2001 From: Sven Mika Date: Thu, 19 Nov 2020 19:01:14 +0100 Subject: [PATCH] [RLlib] Fix inconsistency wrt batch size in SampleCollector (traj. view API). Makes DD-PPO work with traj. view API. (#12063) --- rllib/BUILD | 4 +- rllib/agents/callbacks.py | 45 +++-- rllib/agents/dqn/tests/test_simple_q.py | 5 +- rllib/agents/ppo/appo.py | 2 - rllib/agents/ppo/ddppo.py | 2 - .../evaluation/collectors/sample_collector.py | 42 +++-- .../collectors/simple_list_collector.py | 158 ++++++++++-------- rllib/evaluation/episode.py | 6 +- rllib/evaluation/rollout_worker.py | 3 + rllib/evaluation/sampler.py | 46 ++--- .../tests/test_trajectory_view_api.py | 11 ++ rllib/examples/custom_eval.py | 15 +- rllib/utils/test_utils.py | 7 +- rllib/utils/typing.py | 6 +- 14 files changed, 217 insertions(+), 135 deletions(-) diff --git a/rllib/BUILD b/rllib/BUILD index e6ab78e6d..94f15cef7 100644 --- a/rllib/BUILD +++ b/rllib/BUILD @@ -1667,7 +1667,7 @@ py_test( tags = ["examples", "examples_C"], size = "small", srcs = ["examples/custom_eval.py"], - args = ["--num-cpus=4"] + args = ["--num-cpus=4", "--as-test"] ) py_test( @@ -1676,7 +1676,7 @@ py_test( tags = ["examples", "examples_C"], size = "small", srcs = ["examples/custom_eval.py"], - args = ["--torch", "--num-cpus=4"] + args = ["--num-cpus=4", "--as-test", "--torch"] ) py_test( diff --git a/rllib/agents/callbacks.py b/rllib/agents/callbacks.py index f0460c17d..bf5284740 100644 --- a/rllib/agents/callbacks.py +++ b/rllib/agents/callbacks.py @@ -1,4 +1,4 @@ -from typing import Dict, TYPE_CHECKING +from typing import Dict, Optional, TYPE_CHECKING from ray.rllib.env import BaseEnv from ray.rllib.policy import Policy @@ -7,6 +7,7 @@ from ray.rllib.evaluation import MultiAgentEpisode from ray.rllib.utils.annotations import PublicAPI from ray.rllib.utils.deprecation import deprecation_warning from ray.rllib.utils.typing import AgentID, PolicyID +from ray.util.debug import log_once if TYPE_CHECKING: from ray.rllib.evaluation import RolloutWorker @@ -30,9 +31,13 @@ class DefaultCallbacks: "a class extending rllib.agents.callbacks.DefaultCallbacks") self.legacy_callbacks = legacy_callbacks_dict or {} - def on_episode_start(self, *, worker: "RolloutWorker", base_env: BaseEnv, + def on_episode_start(self, + *, + worker: "RolloutWorker", + base_env: BaseEnv, policies: Dict[PolicyID, Policy], - episode: MultiAgentEpisode, env_index: int, + episode: MultiAgentEpisode, + env_index: Optional[int] = None, **kwargs) -> None: """Callback run on the rollout worker before each episode starts. @@ -46,11 +51,15 @@ class DefaultCallbacks: state. You can use the `episode.user_data` dict to store temporary data, and `episode.custom_metrics` to store custom metrics for the episode. - env_index (int): The index of the (vectorized) env, which the + env_index (EnvID): Obsoleted: The ID of the environment, which the episode belongs to. kwargs: Forward compatibility placeholder. """ + if env_index is not None: + if log_once("callbacks_env_index_deprecated"): + deprecation_warning("env_index", "episode.env_id", error=False) + if self.legacy_callbacks.get("on_episode_start"): self.legacy_callbacks["on_episode_start"]({ "env": base_env, @@ -58,8 +67,12 @@ class DefaultCallbacks: "episode": episode, }) - def on_episode_step(self, *, worker: "RolloutWorker", base_env: BaseEnv, - episode: MultiAgentEpisode, env_index: int, + def on_episode_step(self, + *, + worker: "RolloutWorker", + base_env: BaseEnv, + episode: MultiAgentEpisode, + env_index: Optional[int] = None, **kwargs) -> None: """Runs on each episode step. @@ -71,20 +84,28 @@ class DefaultCallbacks: state. You can use the `episode.user_data` dict to store temporary data, and `episode.custom_metrics` to store custom metrics for the episode. - env_index (int): The index of the (vectorized) env, which the + env_index (EnvID): Obsoleted: The ID of the environment, which the episode belongs to. kwargs: Forward compatibility placeholder. """ + if env_index is not None: + if log_once("callbacks_env_index_deprecated"): + deprecation_warning("env_index", "episode.env_id", error=False) + if self.legacy_callbacks.get("on_episode_step"): self.legacy_callbacks["on_episode_step"]({ "env": base_env, "episode": episode }) - def on_episode_end(self, *, worker: "RolloutWorker", base_env: BaseEnv, + def on_episode_end(self, + *, + worker: "RolloutWorker", + base_env: BaseEnv, policies: Dict[PolicyID, Policy], - episode: MultiAgentEpisode, env_index: int, + episode: MultiAgentEpisode, + env_index: Optional[int] = None, **kwargs) -> None: """Runs when an episode is done. @@ -98,11 +119,15 @@ class DefaultCallbacks: state. You can use the `episode.user_data` dict to store temporary data, and `episode.custom_metrics` to store custom metrics for the episode. - env_index (int): The index of the (vectorized) env, which the + env_index (EnvID): Obsoleted: The ID of the environment, which the episode belongs to. kwargs: Forward compatibility placeholder. """ + if env_index is not None: + if log_once("callbacks_env_index_deprecated"): + deprecation_warning("env_index", "episode.env_id", error=False) + if self.legacy_callbacks.get("on_episode_end"): self.legacy_callbacks["on_episode_end"]({ "env": base_env, diff --git a/rllib/agents/dqn/tests/test_simple_q.py b/rllib/agents/dqn/tests/test_simple_q.py index a87314e80..175c4bc51 100644 --- a/rllib/agents/dqn/tests/test_simple_q.py +++ b/rllib/agents/dqn/tests/test_simple_q.py @@ -19,11 +19,14 @@ class TestSimpleQ(unittest.TestCase): """Test whether a SimpleQTrainer can be built on all frameworks.""" config = dqn.SIMPLE_Q_DEFAULT_CONFIG.copy() config["num_workers"] = 0 # Run locally. + num_iterations = 2 for _ in framework_iterator(config): trainer = dqn.SimpleQTrainer(config=config, env="CartPole-v0") - num_iterations = 2 + rw = trainer.workers.local_worker() for i in range(num_iterations): + sb = rw.sample() + assert sb.count == config["rollout_fragment_length"] results = trainer.train() print(results) diff --git a/rllib/agents/ppo/appo.py b/rllib/agents/ppo/appo.py index f1bbe3361..61510eb04 100644 --- a/rllib/agents/ppo/appo.py +++ b/rllib/agents/ppo/appo.py @@ -76,8 +76,6 @@ DEFAULT_CONFIG = impala.ImpalaTrainer.merge_trainer_configs( "vf_loss_coeff": 0.5, "entropy_coeff": 0.01, "entropy_coeff_schedule": None, - # Trajectory View API not supported for DD-PPO yet. - "_use_trajectory_view_api": False, }, _allow_unknown_configs=True, ) diff --git a/rllib/agents/ppo/ddppo.py b/rllib/agents/ppo/ddppo.py index 59ab2c4a4..9c14a7647 100644 --- a/rllib/agents/ppo/ddppo.py +++ b/rllib/agents/ppo/ddppo.py @@ -74,8 +74,6 @@ DEFAULT_CONFIG = ppo.PPOTrainer.merge_trainer_configs( "truncate_episodes": True, # This is auto set based on sample batch size. "train_batch_size": -1, - # Trajectory View API not supported for DD-PPO yet. - "_use_trajectory_view_api": False, }, _allow_unknown_configs=True, ) diff --git a/rllib/evaluation/collectors/sample_collector.py b/rllib/evaluation/collectors/sample_collector.py index 7d6f00b52..4689c9261 100644 --- a/rllib/evaluation/collectors/sample_collector.py +++ b/rllib/evaluation/collectors/sample_collector.py @@ -1,6 +1,6 @@ from abc import abstractmethod, ABCMeta import logging -from typing import Dict, Union +from typing import Dict, List, Optional, Union from ray.rllib.evaluation.episode import MultiAgentEpisode from ray.rllib.policy.sample_batch import MultiAgentBatch, SampleBatch @@ -145,7 +145,8 @@ class _SampleCollector(metaclass=ABCMeta): def postprocess_episode(self, episode: MultiAgentEpisode, is_done: bool = False, - check_dones: bool = False) -> None: + check_dones: bool = False, + build: bool = False) -> Optional[MultiAgentBatch]: """Postprocesses all agents' trajectories in a given episode. Generates (single-trajectory) SampleBatches for all Policies/Agents and @@ -159,31 +160,27 @@ class _SampleCollector(metaclass=ABCMeta): episode (MultiAgentEpisode): The Episode object for which to post-process data. is_done (bool): Whether the given episode is actually terminated - (all agents are done). + (all agents are done OR we hit a hard horizon). If True, the + episode will no longer be used/continued and we may need to + recycle/erase it internally. If a soft-horizon is hit, the + episode will continue to be used and `is_done` should be set + to False here. check_dones (bool): Whether we need to check that all agents' trajectories have dones=True at the end. - """ - raise NotImplementedError - - @abstractmethod - def build_multi_agent_batch(self, env_steps: int) -> \ - Union[MultiAgentBatch, SampleBatch]: - """Builds a MultiAgentBatch of size=env_steps from the collected data. - - Args: - env_steps (int): The sum of all env-steps (across all agents) taken - so far. + build (bool): Whether to build a MultiAgentBatch from the given + episode (and only that episode!) and return that + MultiAgentBatch. Used for batch_mode=`complete_episodes`. Returns: - Union[MultiAgentBatch, SampleBatch]: Returns the accumulated - sample batches for each policy inside one MultiAgentBatch - object (or a simple SampleBatch if only one policy). + Any: An ID that can be used in `build_multi_agent_batch` to + retrieve the samples that have been postprocessed as a + ready-built MultiAgentBatch. """ raise NotImplementedError @abstractmethod def try_build_truncated_episode_multi_agent_batch(self) -> \ - Union[MultiAgentBatch, SampleBatch, None]: + List[Union[MultiAgentBatch, SampleBatch]]: """Tries to build an MA-batch, if `rollout_fragment_length` is reached. Any unprocessed data will be first postprocessed with a policy @@ -193,9 +190,10 @@ class _SampleCollector(metaclass=ABCMeta): returns None. Returns: - Union[MultiAgentBatch, SampleBatch, None]: Returns the accumulated - sample batches for each policy inside one MultiAgentBatch - object (or a simple SampleBatch if only one policy) or None - if `self.rollout_fragment_length` has not been reached yet. + List[Union[MultiAgentBatch, SampleBatch]]: Returns a (possibly + empty) list of MultiAgentBatches (containing the accumulated + SampleBatches for each policy or a simple SampleBatch if only + one policy). The list will be empty if + `self.rollout_fragment_length` has not been reached yet. """ raise NotImplementedError diff --git a/rllib/evaluation/collectors/simple_list_collector.py b/rllib/evaluation/collectors/simple_list_collector.py index fa426d40f..97623b473 100644 --- a/rllib/evaluation/collectors/simple_list_collector.py +++ b/rllib/evaluation/collectors/simple_list_collector.py @@ -1,7 +1,7 @@ import collections import logging import numpy as np -from typing import List, Any, Dict, Tuple, TYPE_CHECKING, Union +from typing import Any, List, Dict, Tuple, TYPE_CHECKING, Union from ray.rllib.env.base_env import _DUMMY_AGENT_ID from ray.rllib.evaluation.collectors.sample_collector import _SampleCollector @@ -251,6 +251,15 @@ class _PolicyCollector: return batch +class _PolicyCollectorGroup: + def __init__(self, policy_map): + self.policy_collectors = { + pid: _PolicyCollector() + for pid in policy_map.keys() + } + self.count = 0 + + class _SimpleListCollector(_SampleCollector): """Util to build SampleBatches for each policy in a multi-agent env. @@ -285,38 +294,41 @@ class _SimpleListCollector(_SampleCollector): 1000, rollout_fragment_length * 10) if rollout_fragment_length != float("inf") else 5000 - # Build each Policies' single collector. - self.policy_collectors = { - pid: _PolicyCollector() - for pid in policy_map.keys() - } - self.policy_collectors_env_steps = 0 # Whenever we observe a new episode+agent, add a new # _SingleTrajectoryCollector. self.agent_collectors: Dict[Tuple[EpisodeID, AgentID], _AgentCollector] = {} - # Internal agent-key-to-policy map. - self.agent_key_to_policy = {} + # Internal agent-key-to-policy-id map. + self.agent_key_to_policy_id = {} + # Pool of used/unused PolicyCollectorGroups (attached to episodes for + # across-episode multi-agent sample collection). + self.policy_collector_groups = [] # Agents to collect data from for the next forward pass (per policy). self.forward_pass_agent_keys = {pid: [] for pid in policy_map.keys()} self.forward_pass_size = {pid: 0 for pid in policy_map.keys()} - # Maps episode ID to _EpisodeRecord objects. - self.episode_steps: Dict[EpisodeID, int] = collections.defaultdict(int) + # Maps episode ID to the (non-built) env steps taken in this episode. + self.episode_steps: Dict[EpisodeID, int] = \ + collections.defaultdict(int) + # Maps episode ID to MultiAgentEpisode. self.episodes: Dict[EpisodeID, MultiAgentEpisode] = {} @override(_SampleCollector) def episode_step(self, episode_id: EpisodeID) -> None: + episode = self.episodes[episode_id] self.episode_steps[episode_id] += 1 + episode.length += 1 + assert episode.batch_builder is not None + env_steps = episode.batch_builder.count + num_observations = sum( + c.count for c in episode.batch_builder.policy_collectors.values()) - env_steps = \ - self.policy_collectors_env_steps + self.episode_steps[episode_id] - if (env_steps > self.large_batch_threshold - and log_once("large_batch_warning")): + if num_observations > self.large_batch_threshold and \ + log_once("large_batch_warning"): logger.warning( - "More than {} observations for {} env steps ".format( - env_steps, env_steps) + + "More than {} observations in {} env steps for " + "episode {} ".format(num_observations, env_steps, episode_id) + "are buffered in the sampler. If this is more than you " "expected, check that that you set a horizon on your " "environment correctly and that it terminates at some point. " @@ -324,7 +336,7 @@ class _SimpleListCollector(_SampleCollector): "sets the batch size based on (across-agents) environment " "steps, not the steps of individual agents, which can result " "in unexpectedly large batches." + - ("Also, you may be in evaluation waiting for your Env to " + ("Also, you may be waiting for your Env to " "terminate (batch_mode=`complete_episodes`). Make sure it " "does at some point." if not self.multiple_episodes_in_batch else "")) @@ -335,10 +347,10 @@ class _SimpleListCollector(_SampleCollector): init_obs: TensorType) -> None: # Make sure our mappings are up to date. agent_key = (episode.episode_id, agent_id) - if agent_key not in self.agent_key_to_policy: - self.agent_key_to_policy[agent_key] = policy_id + if agent_key not in self.agent_key_to_policy_id: + self.agent_key_to_policy_id[agent_key] = policy_id else: - assert self.agent_key_to_policy[agent_key] == policy_id + assert self.agent_key_to_policy_id[agent_key] == policy_id policy = self.policy_map[policy_id] view_reqs = policy.model.inference_view_requirements if \ getattr(policy, "model", None) else policy.view_requirements @@ -355,8 +367,12 @@ class _SimpleListCollector(_SampleCollector): view_requirements=view_reqs) self.episodes[episode.episode_id] = episode + if episode.batch_builder is None: + episode.batch_builder = self.policy_collector_groups.pop() if \ + self.policy_collector_groups else _PolicyCollectorGroup( + self.policy_map) - self._add_to_next_inference_call(agent_key, env_id) + self._add_to_next_inference_call(agent_key) @override(_SampleCollector) def add_action_reward_next_obs(self, episode_id: EpisodeID, @@ -365,7 +381,7 @@ class _SimpleListCollector(_SampleCollector): values: Dict[str, TensorType]) -> None: # Make sure, episode/agent already has some (at least init) data. agent_key = (episode_id, agent_id) - assert self.agent_key_to_policy[agent_key] == policy_id + assert self.agent_key_to_policy_id[agent_key] == policy_id assert agent_key in self.agent_collectors # Include the current agent id for multi-agent algorithms. @@ -376,7 +392,7 @@ class _SimpleListCollector(_SampleCollector): self.agent_collectors[agent_key].add_action_reward_next_obs(values) if not agent_done: - self._add_to_next_inference_call(agent_key, env_id) + self._add_to_next_inference_call(agent_key) @override(_SampleCollector) def total_env_steps(self) -> int: @@ -417,8 +433,10 @@ class _SimpleListCollector(_SampleCollector): def postprocess_episode(self, episode: MultiAgentEpisode, is_done: bool = False, - check_dones: bool = False) -> None: + check_dones: bool = False, + build: bool = False) -> None: episode_id = episode.episode_id + policy_collector_group = episode.batch_builder # TODO: (sven) Once we implement multi-agent communication channels, # we have to resolve the restriction of only sending other agent @@ -429,8 +447,8 @@ class _SimpleListCollector(_SampleCollector): # Build only if there is data and agent is part of given episode. if collector.count == 0 or eps_id != episode_id: continue - policy = self.policy_map[self.agent_key_to_policy[(eps_id, - agent_id)]] + pid = self.agent_key_to_policy_id[(eps_id, agent_id)] + policy = self.policy_map[pid] pre_batch = collector.build(policy.view_requirements) pre_batches[agent_id] = (policy, pre_batch) @@ -455,7 +473,7 @@ class _SimpleListCollector(_SampleCollector): "Episode {} terminated for all agents, but we still don't " "don't have a last observation for agent {} (policy " "{}). ".format( - episode_id, agent_id, self.agent_key_to_policy[( + episode_id, agent_id, self.agent_key_to_policy_id[( episode_id, agent_id)]) + "Please ensure that you include the last observations " "of all live agents when setting done[__all__] to " @@ -467,8 +485,8 @@ class _SimpleListCollector(_SampleCollector): other_batches = pre_batches.copy() del other_batches[agent_id] - policy = self.policy_map[self.agent_key_to_policy[(episode_id, - agent_id)]] + pid = self.agent_key_to_policy_id[(episode_id, agent_id)] + policy = self.policy_map[pid] if any(pre_batch["dones"][:-1]) or len(set( pre_batch["eps_id"])) > 1: raise ValueError( @@ -491,7 +509,7 @@ class _SimpleListCollector(_SampleCollector): # Append into policy batches and reset. from ray.rllib.evaluation.rollout_worker import get_global_worker for agent_id, post_batch in sorted(post_batches.items()): - pid = self.agent_key_to_policy[(episode_id, agent_id)] + pid = self.agent_key_to_policy_id[(episode_id, agent_id)] policy = self.policy_map[pid] self.callbacks.on_postprocess_trajectory( worker=get_global_worker(), @@ -503,60 +521,65 @@ class _SimpleListCollector(_SampleCollector): original_batches=pre_batches) # Add the postprocessed SampleBatch to the policy collectors for # training. - self.policy_collectors[pid].add_postprocessed_batch_for_training( - post_batch, policy.view_requirements) + policy_collector_group.policy_collectors[ + pid].add_postprocessed_batch_for_training( + post_batch, policy.view_requirements) env_steps = self.episode_steps[episode_id] - self.policy_collectors_env_steps += env_steps + policy_collector_group.count += env_steps if is_done: del self.episode_steps[episode_id] del self.episodes[episode_id] + # Make PolicyCollectorGroup available for more agent batches in + # other episodes. Do not reset count to 0. + self.policy_collector_groups.append(policy_collector_group) else: self.episode_steps[episode_id] = 0 - @override(_SampleCollector) - def build_multi_agent_batch(self, env_steps: int) -> \ + # Build a MultiAgentBatch from the episode and return. + if build: + return self._build_multi_agent_batch(episode) + + def _build_multi_agent_batch(self, episode: MultiAgentEpisode) -> \ Union[MultiAgentBatch, SampleBatch]: + + ma_batch = {} + for pid, collector in episode.batch_builder.policy_collectors.items(): + if collector.count > 0: + ma_batch[pid] = collector.build() + # Create the batch. ma_batch = MultiAgentBatch.wrap_as_needed( - { - pid: collector.build() - for pid, collector in self.policy_collectors.items() - if collector.count > 0 - }, - env_steps=env_steps) - self.policy_collectors_env_steps = 0 + ma_batch, env_steps=episode.batch_builder.count) + + # PolicyCollectorGroup is empty. + episode.batch_builder.count = 0 + return ma_batch @override(_SampleCollector) def try_build_truncated_episode_multi_agent_batch(self) -> \ - Union[MultiAgentBatch, SampleBatch, None]: - # Have something to loop through, even if there are currently no - # ongoing episodes. - episode_steps = self.episode_steps or {"_fake_id": 0} + List[Union[MultiAgentBatch, SampleBatch]]: + batches = [] # Loop through ongoing episodes and see whether their length plus # what's already in the policy collectors reaches the fragment-len. - for episode_id, count in episode_steps.items(): - env_steps = self.policy_collectors_env_steps + count + for episode_id, episode in self.episodes.items(): + env_steps = episode.batch_builder.count + \ + self.episode_steps[episode_id] # Reached the fragment-len -> We should build an MA-Batch. if env_steps >= self.rollout_fragment_length: + assert env_steps == self.rollout_fragment_length # If we reached the fragment-len only because of `episode_id` # (still ongoing) -> postprocess `episode_id` first. - if self.policy_collectors_env_steps < \ - self.rollout_fragment_length: - self.postprocess_episode( - self.episodes[episode_id], is_done=False) - # Otherwise, create MA-batch only from what's already in our - # policy buffers (do not include `episode_id`'s data). - else: - env_steps = self.policy_collectors_env_steps + if episode.batch_builder.count < self.rollout_fragment_length: + self.postprocess_episode(episode, is_done=False) # Build the MA-batch and return. - ma_batch = self.build_multi_agent_batch(env_steps=env_steps) - return ma_batch - return None + batch = self._build_multi_agent_batch(episode=episode) + batches.append(batch) + return batches - def _add_to_next_inference_call(self, agent_key: Tuple[EpisodeID, AgentID], - env_id: EnvID) -> None: + def _add_to_next_inference_call( + self, agent_key: Tuple[EpisodeID, AgentID]) -> None: """Adds an Agent key (episode+agent IDs) to the next inference call. This makes sure that the agent's current data (in the trajectory) is @@ -566,14 +589,13 @@ class _SimpleListCollector(_SampleCollector): Args: agent_key (Tuple[EpisodeID, AgentID]: A unique agent key (across vectorized environments). - env_id (EnvID): The environment index (in a vectorized setup). """ - policy_id = self.agent_key_to_policy[agent_key] - idx = self.forward_pass_size[policy_id] + pid = self.agent_key_to_policy_id[agent_key] + idx = self.forward_pass_size[pid] if idx == 0: - self.forward_pass_agent_keys[policy_id].clear() - self.forward_pass_agent_keys[policy_id].append(agent_key) - self.forward_pass_size[policy_id] += 1 + self.forward_pass_agent_keys[pid].clear() + self.forward_pass_agent_keys[pid].append(agent_key) + self.forward_pass_size[pid] += 1 def _reset_inference_calls(self, policy_id: PolicyID) -> None: """Resets internal inference input-dict registries. diff --git a/rllib/evaluation/episode.py b/rllib/evaluation/episode.py index dbb33ac2c..4bf5a7172 100644 --- a/rllib/evaluation/episode.py +++ b/rllib/evaluation/episode.py @@ -8,7 +8,7 @@ from ray.rllib.policy.policy import Policy from ray.rllib.utils.annotations import DeveloperAPI from ray.rllib.utils.spaces.space_utils import flatten_to_single_ndarray from ray.rllib.utils.typing import SampleBatchType, AgentID, PolicyID, \ - EnvObsType, EnvInfoDict, EnvActionType + EnvActionType, EnvID, EnvInfoDict, EnvObsType if TYPE_CHECKING: from ray.rllib.evaluation.sample_batch_builder import \ @@ -48,7 +48,8 @@ class MultiAgentEpisode: policy_mapping_fn: Callable[[AgentID], PolicyID], batch_builder_factory: Callable[ [], "MultiAgentSampleBatchBuilder"], - extra_batch_callback: Callable[[SampleBatchType], None]): + extra_batch_callback: Callable[[SampleBatchType], None], + env_id: EnvID): self.new_batch_builder: Callable[ [], "MultiAgentSampleBatchBuilder"] = batch_builder_factory self.add_extra_batch: Callable[[SampleBatchType], @@ -58,6 +59,7 @@ class MultiAgentEpisode: self.total_reward: float = 0.0 self.length: int = 0 self.episode_id: int = random.randrange(2e9) + self.env_id = env_id self.agent_rewards: Dict[AgentID, float] = defaultdict(float) self.custom_metrics: Dict[str, float] = {} self.user_data: Dict[str, Any] = {} diff --git a/rllib/evaluation/rollout_worker.py b/rllib/evaluation/rollout_worker.py index faee68024..f98d29e32 100644 --- a/rllib/evaluation/rollout_worker.py +++ b/rllib/evaluation/rollout_worker.py @@ -613,6 +613,9 @@ class RolloutWorker(ParallelIteratorWorker): if self.fake_sampler and self.last_batch is not None: return self.last_batch + elif self.input_reader is None: + raise ValueError("RolloutWorker has no `input_reader` object! " + "Cannot call `sample()`.") if log_once("sample_start"): logger.info("Generating sample batch of size {}".format( diff --git a/rllib/evaluation/sampler.py b/rllib/evaluation/sampler.py index f7a7dfd46..7c098dd65 100644 --- a/rllib/evaluation/sampler.py +++ b/rllib/evaluation/sampler.py @@ -51,11 +51,11 @@ StateBatch = List[List[Any]] class NewEpisodeDefaultDict(defaultdict): - def __missing__(self, env_index): + def __missing__(self, env_id): if self.default_factory is None: - raise KeyError(env_index) + raise KeyError(env_id) else: - ret = self[env_index] = self.default_factory(env_index) + ret = self[env_id] = self.default_factory(env_id) return ret @@ -517,9 +517,13 @@ def _env_runner( return MultiAgentSampleBatchBuilder(policies, clip_rewards, callbacks) - def new_episode(env_index): - episode = MultiAgentEpisode(policies, policy_mapping_fn, - get_batch_builder, extra_batch_callback) + def new_episode(env_id): + episode = MultiAgentEpisode( + policies, + policy_mapping_fn, + get_batch_builder, + extra_batch_callback, + env_id=env_id) # Call each policy's Exploration.on_episode_start method. # type: Policy for p in policies.values(): @@ -534,7 +538,7 @@ def _env_runner( base_env=base_env, policies=policies, episode=episode, - env_index=env_index, + env_index=env_id, ) return episode @@ -972,7 +976,6 @@ def _process_observations_w_trajectory_view_api( if not is_new_episode: _sample_collector.episode_step(episode.episode_id) - episode.length += 1 episode._add_agent_rewards(rewards[env_id]) # Check episode termination conditions. @@ -1077,20 +1080,23 @@ def _process_observations_w_trajectory_view_api( episode=episode, env_index=env_id) - # Episode is done for all agents - # (dones[__all__] == True or hit horizon). - # Make sure postprocessor stays within one episode. + # Episode is done for all agents (dones[__all__] == True) + # or we hit the horizon. if all_agents_done: is_done = dones[env_id]["__all__"] check_dones = is_done and not no_done_at_end - _sample_collector.postprocess_episode( - episode, is_done=is_done, check_dones=check_dones) - # We are not allowed to pack the next episode into the same + + # If, we are not allowed to pack the next episode into the same # SampleBatch (batch_mode=complete_episodes) -> Build the # MultiAgentBatch from a single episode and add it to "outputs". - if not multiple_episodes_in_batch: - ma_sample_batch = \ - _sample_collector.build_multi_agent_batch(episode.length) + # Otherwise, just postprocess and continue collecting across + # episodes. + ma_sample_batch = _sample_collector.postprocess_episode( + episode, + is_done=is_done or (hit_horizon and not soft_horizon), + check_dones=check_dones, + build=not multiple_episodes_in_batch) + if ma_sample_batch: outputs.append(ma_sample_batch) # Call each policy's Exploration.on_episode_end method. @@ -1155,10 +1161,10 @@ def _process_observations_w_trajectory_view_api( # Try to build something. if multiple_episodes_in_batch: - sample_batch = \ + sample_batches = \ _sample_collector.try_build_truncated_episode_multi_agent_batch() - if sample_batch is not None: - outputs.append(sample_batch) + if sample_batches: + outputs.extend(sample_batches) return active_envs, to_eval, outputs diff --git a/rllib/evaluation/tests/test_trajectory_view_api.py b/rllib/evaluation/tests/test_trajectory_view_api.py index a21161bb8..d92d2b499 100644 --- a/rllib/evaluation/tests/test_trajectory_view_api.py +++ b/rllib/evaluation/tests/test_trajectory_view_api.py @@ -28,6 +28,9 @@ class TestTrajectoryViewAPI(unittest.TestCase): """Tests, whether Model and Policy return the correct ViewRequirements. """ config = dqn.DEFAULT_CONFIG.copy() + config["num_envs_per_worker"] = 10 + config["rollout_fragment_length"] = 4 + for _ in framework_iterator(config): trainer = dqn.DQNTrainer( config, @@ -55,6 +58,14 @@ class TestTrajectoryViewAPI(unittest.TestCase): else: assert view_req_policy[key].data_col == SampleBatch.OBS assert view_req_policy[key].shift == 1 + rollout_worker = trainer.workers.local_worker() + sample_batch = rollout_worker.sample() + expected_count = \ + config["num_envs_per_worker"] * \ + config["rollout_fragment_length"] + assert sample_batch.count == expected_count + for v in sample_batch.data.values(): + assert len(v) == expected_count trainer.stop() def test_traj_view_lstm_prev_actions_and_rewards(self): diff --git a/rllib/examples/custom_eval.py b/rllib/examples/custom_eval.py index 459c4c627..1c5bd7bd6 100644 --- a/rllib/examples/custom_eval.py +++ b/rllib/examples/custom_eval.py @@ -73,10 +73,15 @@ import ray from ray import tune from ray.rllib.evaluation.metrics import collect_episodes, summarize_episodes from ray.rllib.examples.env.simple_corridor import SimpleCorridor +from ray.rllib.utils.test_utils import check_learning_achieved parser = argparse.ArgumentParser() parser.add_argument("--num-cpus", type=int, default=0) parser.add_argument("--torch", action="store_true") +parser.add_argument("--as-test", action="store_true") +parser.add_argument("--stop-iters", type=int, default=50) +parser.add_argument("--stop-timesteps", type=int, default=20000) +parser.add_argument("--stop-reward", type=float, default=0.7) parser.add_argument("--no-custom-eval", action="store_true") @@ -169,9 +174,15 @@ if __name__ == "__main__": } stop = { - "training_iteration": 10, + "training_iteration": args.stop_iters, + "timesteps_total": args.stop_timesteps, + "episode_reward_mean": args.stop_reward, } - tune.run("PG", config=config, stop=stop) + results = tune.run("PG", config=config, stop=stop, verbose=1) + # Check eval results (from eval workers using the custom function), + # not results from the regular workers. + if args.as_test: + check_learning_achieved(results, args.stop_reward, evaluation=True) ray.shutdown() diff --git a/rllib/utils/test_utils.py b/rllib/utils/test_utils.py index 01138001f..f9a779ce0 100644 --- a/rllib/utils/test_utils.py +++ b/rllib/utils/test_utils.py @@ -243,7 +243,7 @@ def check(x, y, decimals=5, atol=None, rtol=None, false=False): "ERROR: x ({}) is the same as y ({})!".format(x, y) -def check_learning_achieved(tune_results, min_reward): +def check_learning_achieved(tune_results, min_reward, evaluation=False): """Throws an error if `min_reward` is not reached within tune_results. Checks the last iteration found in tune_results for its @@ -256,7 +256,10 @@ def check_learning_achieved(tune_results, min_reward): Raises: ValueError: If `min_reward` not reached. """ - if tune_results.trials[0].last_result["episode_reward_mean"] < min_reward: + last_result = tune_results.trials[0].last_result + avg_reward = last_result["episode_reward_mean"] if not evaluation else \ + last_result["evaluation"]["episode_reward_mean"] + if avg_reward < min_reward: raise ValueError("`stop-reward` of {} not reached!".format(min_reward)) print("ok") diff --git a/rllib/utils/typing.py b/rllib/utils/typing.py index eaf3e4e63..6684ae53a 100644 --- a/rllib/utils/typing.py +++ b/rllib/utils/typing.py @@ -37,8 +37,10 @@ PolicyID = str MultiAgentPolicyConfigDict = Dict[PolicyID, Tuple[Union[ type, None], gym.Space, gym.Space, PartialTrainerConfigDict]] -# Represents an environment id. -EnvID = int +# Represents an environment id. These could be: +# - An int index for a sub-env within a vectorized env. +# - An external env ID (str), which changes(!) each episode. +EnvID = Union[int, str] # Represents an episode id. EpisodeID = int