[RLlib] Fix inconsistency wrt batch size in SampleCollector (traj. view API). Makes DD-PPO work with traj. view API. (#12063)

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(

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,

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)
 

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,
 )

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,
 )

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

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.

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] = {}

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(

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
 

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):

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()

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")
 

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