[RLlib] Trajectory View API (part 2.5): Actual implementations (not used yet) of a SampleCollector. (#10112)

rllib/evaluation/episode.py

@@ -92,11 +92,17 @@ class MultiAgentEpisode:
         self._agent_reward_history = defaultdict(list)
 
     @DeveloperAPI
-    def policy_for(self, agent_id: AgentID = _DUMMY_AGENT_ID) -> Policy:
-        """Returns the policy for the specified agent.
+    def policy_for(self, agent_id: AgentID = _DUMMY_AGENT_ID) -> PolicyID:
+        """Returns and stores the policy ID for the specified agent.
 
         If the agent is new, the policy mapping fn will be called to bind the
         agent to a policy for the duration of the episode.
+
+        Args:
+            agent_id (AgentID): The agent ID to lookup the policy ID for.
+
+        Returns:
+            PolicyID: The policy ID for the specified agent.
         """
 
         if agent_id not in self._agent_to_policy:

rllib/evaluation/multi_agent_sample_collector.py

@@ -0,0 +1,247 @@
+import logging
+from typing import Dict, Optional
+
+from ray.rllib.agents.callbacks import DefaultCallbacks
+from ray.rllib.env.base_env import _DUMMY_AGENT_ID
+from ray.rllib.evaluation.episode import MultiAgentEpisode
+from ray.rllib.evaluation.per_policy_sample_collector import \
+    _PerPolicySampleCollector
+from ray.rllib.evaluation.sample_collector import _SampleCollector
+from ray.rllib.policy.policy import Policy
+from ray.rllib.policy.sample_batch import MultiAgentBatch
+from ray.rllib.utils import force_list
+from ray.rllib.utils.annotations import override
+from ray.rllib.utils.debug import summarize
+from ray.rllib.utils.types import AgentID, EnvID, EpisodeID, PolicyID, \
+    TensorType
+from ray.util.debug import log_once
+
+logger = logging.getLogger(__name__)
+
+
+class _MultiAgentSampleCollector(_SampleCollector):
+    """Builds SampleBatches for each policy (and agent) in a multi-agent env.
+
+    Note: This is an experimental class only used when
+    `config._use_trajectory_view_api` = True.
+    Once `_use_trajectory_view_api` becomes the default in configs:
+    This class will deprecate the `SampleBatchBuilder` class.
+
+    Input data is collected in central per-policy buffers, which
+    efficiently pre-allocate memory (over n timesteps) and re-use the same
+    memory even for succeeding agents and episodes.
+    Input_dicts for action computations, SampleBatches for postprocessing, and
+    train_batch dicts are - if possible - created from the central per-policy
+    buffers via views to avoid copying of data).
+    """
+
+    def __init__(
+            self,
+            policy_map: Dict[PolicyID, Policy],
+            callbacks: DefaultCallbacks,
+            # TODO: (sven) make `num_agents` flexibly grow in size.
+            num_agents: int = 100,
+            num_timesteps=None,
+            time_major: Optional[bool] = False):
+        """Initializes a _MultiAgentSampleCollector object.
+
+        Args:
+            policy_map (Dict[PolicyID,Policy]): Maps policy ids to policy
+                instances.
+            callbacks (DefaultCallbacks): RLlib callbacks (configured in the
+                Trainer config dict). Used for trajectory postprocessing event.
+            num_agents (int): The max number of agent slots to pre-allocate
+                in the buffer.
+            num_timesteps (int): The max number of timesteps to pre-allocate
+                in the buffer.
+            time_major (Optional[bool]): Whether to preallocate buffers and
+                collect samples in time-major fashion (TxBx...).
+        """
+
+        self.policy_map = policy_map
+        self.callbacks = callbacks
+        if num_agents == float("inf") or num_agents is None:
+            num_agents = 1000
+        self.num_agents = int(num_agents)
+
+        # Collect SampleBatches per-policy in PolicyTrajectories objects.
+        self.rollout_sample_collectors = {}
+        for pid, policy in policy_map.items():
+            # Figure out max-shifts (before and after).
+            view_reqs = policy.training_view_requirements
+            max_shift_before = 0
+            max_shift_after = 0
+            for vr in view_reqs.values():
+                shift = force_list(vr.shift)
+                if max_shift_before > shift[0]:
+                    max_shift_before = shift[0]
+                if max_shift_after < shift[-1]:
+                    max_shift_after = shift[-1]
+            # Figure out num_timesteps and num_agents.
+            kwargs = {"time_major": time_major}
+            if policy.is_recurrent():
+                kwargs["num_timesteps"] = \
+                    policy.config["model"]["max_seq_len"]
+                kwargs["time_major"] = True
+            elif num_timesteps is not None:
+                kwargs["num_timesteps"] = num_timesteps
+
+            self.rollout_sample_collectors[pid] = _PerPolicySampleCollector(
+                num_agents=self.num_agents,
+                shift_before=-max_shift_before,
+                shift_after=max_shift_after,
+                **kwargs)
+
+        # Internal agent-to-policy map.
+        self.agent_to_policy = {}
+        # Number of "inference" steps taken in the environment.
+        # Regardless of the number of agents involved in each of these steps.
+        self.count = 0
+
+    @override(_SampleCollector)
+    def add_init_obs(self, episode_id: EpisodeID, agent_id: AgentID,
+                     env_id: EnvID, policy_id: PolicyID,
+                     obs: TensorType) -> None:
+        # Make sure our mappings are up to date.
+        if agent_id not in self.agent_to_policy:
+            self.agent_to_policy[agent_id] = policy_id
+        else:
+            assert self.agent_to_policy[agent_id] == policy_id
+
+        # Add initial obs to Trajectory.
+        self.rollout_sample_collectors[policy_id].add_init_obs(
+            episode_id, agent_id, env_id, chunk_num=0, init_obs=obs)
+
+    @override(_SampleCollector)
+    def add_action_reward_next_obs(self, episode_id: EpisodeID,
+                                   agent_id: AgentID, env_id: EnvID,
+                                   policy_id: PolicyID, agent_done: bool,
+                                   values: Dict[str, TensorType]) -> None:
+        assert policy_id in self.rollout_sample_collectors
+
+        # Make sure our mappings are up to date.
+        if agent_id not in self.agent_to_policy:
+            self.agent_to_policy[agent_id] = policy_id
+        else:
+            assert 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
+
+        # Add action/reward/next-obs (and other data) to Trajectory.
+        self.rollout_sample_collectors[policy_id].add_action_reward_next_obs(
+            episode_id, agent_id, env_id, agent_done, values)
+
+    @override(_SampleCollector)
+    def total_env_steps(self) -> int:
+        return sum(a.timesteps_since_last_reset
+                   for a in self.rollout_sample_collectors.values())
+
+    def total(self):
+        # TODO: (sven) deprecate; use `self.total_env_steps`, instead.
+        #  Sampler is currently still using `total()`.
+        return self.total_env_steps()
+
+    @override(_SampleCollector)
+    def get_inference_input_dict(self, policy_id: PolicyID) -> \
+            Dict[str, TensorType]:
+        policy = self.policy_map[policy_id]
+        view_reqs = policy.model.inference_view_requirements
+        return self.rollout_sample_collectors[
+            policy_id].get_inference_input_dict(view_reqs)
+
+    @override(_SampleCollector)
+    def has_non_postprocessed_data(self) -> bool:
+        return self.total_env_steps() > 0
+
+    @override(_SampleCollector)
+    def postprocess_trajectories_so_far(
+            self, episode: Optional[MultiAgentEpisode] = None) -> None:
+        # Loop through each per-policy collector and create a view (for each
+        # agent as SampleBatch) from its buffers for post-processing
+        all_agent_batches = {}
+        for pid, rc in self.rollout_sample_collectors.items():
+            policy = self.policy_map[pid]
+            view_reqs = policy.training_view_requirements
+            agent_batches = rc.get_postprocessing_sample_batches(
+                episode, view_reqs)
+
+            for agent_key, batch in agent_batches.items():
+                other_batches = None
+                if len(agent_batches) > 1:
+                    other_batches = agent_batches.copy()
+                    del other_batches[agent_key]
+
+                agent_batches[agent_key] = policy.postprocess_trajectory(
+                    batch, other_batches, episode)
+                # Call the Policy's Exploration's postprocess method.
+                if getattr(policy, "exploration", None) is not None:
+                    agent_batches[
+                        agent_key] = policy.exploration.postprocess_trajectory(
+                            policy, agent_batches[agent_key],
+                            getattr(policy, "_sess", None))
+
+                # Add new columns' data to buffers.
+                for col in agent_batches[agent_key].new_columns:
+                    data = agent_batches[agent_key].data[col]
+                    rc._build_buffers({col: data[0]})
+                    timesteps = data.shape[0]
+                    rc.buffers[col][rc.shift_before:rc.shift_before +
+                                    timesteps, rc.agent_key_to_slot[
+                                        agent_key]] = data
+
+            all_agent_batches.update(agent_batches)
+
+        if log_once("after_post"):
+            logger.info("Trajectory fragment after postprocess_trajectory():"
+                        "\n\n{}\n".format(summarize(all_agent_batches)))
+
+        # Append into policy batches and reset
+        from ray.rllib.evaluation.rollout_worker import get_global_worker
+        for agent_key, batch in sorted(all_agent_batches.items()):
+            self.callbacks.on_postprocess_trajectory(
+                worker=get_global_worker(),
+                episode=episode,
+                agent_id=agent_key[0],
+                policy_id=self.agent_to_policy[agent_key[0]],
+                policies=self.policy_map,
+                postprocessed_batch=batch,
+                original_batches=None)  # TODO: (sven) do we really need this?
+
+    @override(_SampleCollector)
+    def check_missing_dones(self, episode_id: EpisodeID) -> None:
+        for pid, rc in self.rollout_sample_collectors.items():
+            for agent_key in rc.agent_key_to_slot.keys():
+                # Only check for given episode and only for last chunk
+                # (all previous chunks for that agent in the episode are
+                # non-terminal).
+                if (agent_key[1] == episode_id
+                        and rc.agent_key_to_chunk_num[agent_key[:2]] ==
+                        agent_key[2]):
+                    t = rc.agent_key_to_timestep[agent_key] - 1
+                    b = rc.agent_key_to_slot[agent_key]
+                    if not rc.buffers["dones"][t][b]:
+                        raise ValueError(
+                            "Episode {} terminated for all agents, but we "
+                            "still don't have a last observation for "
+                            "agent {} (policy {}). ".format(agent_key[0], pid)
+                            + "Please ensure that you include the last "
+                            "observations of all live agents when setting "
+                            "'__all__' done to True. Alternatively, set "
+                            "no_done_at_end=True to allow this.")
+
+    @override(_SampleCollector)
+    def get_multi_agent_batch_and_reset(self):
+        self.postprocess_trajectories_so_far()
+        policy_batches = {}
+        for pid, rc in self.rollout_sample_collectors.items():
+            policy = self.policy_map[pid]
+            view_reqs = policy.training_view_requirements
+            policy_batches[pid] = rc.get_train_sample_batch_and_reset(
+                view_reqs)
+
+        ma_batch = MultiAgentBatch.wrap_as_needed(policy_batches, self.count)
+        # Reset our across-all-agents env step count.
+        self.count = 0
+        return ma_batch

rllib/evaluation/per_policy_sample_collector.py

@@ -0,0 +1,487 @@
+import logging
+import numpy as np
+from typing import Dict, Optional
+
+from ray.rllib.evaluation.episode import MultiAgentEpisode
+from ray.rllib.policy.sample_batch import SampleBatch
+from ray.rllib.policy.view_requirement import ViewRequirement
+from ray.rllib.utils.framework import try_import_tf, try_import_torch
+from ray.rllib.utils.types import AgentID, EnvID, EpisodeID, TensorType
+
+tf1, tf, tfv = try_import_tf()
+torch, _ = try_import_torch()
+
+logger = logging.getLogger(__name__)
+
+
+class _PerPolicySampleCollector:
+    """A class for efficiently collecting samples for a single (fixed) policy.
+
+    Can be used by a _MultiAgentSampleCollector for its different policies.
+    """
+
+    def __init__(self,
+                 num_agents: Optional[int] = None,
+                 num_timesteps: Optional[int] = None,
+                 time_major: bool = True,
+                 shift_before: int = 0,
+                 shift_after: int = 0):
+        """Initializes a _PerPolicySampleCollector object.
+
+        Args:
+            num_agents (int): The max number of agent slots to pre-allocate
+                in the buffer.
+            num_timesteps (int): The max number of timesteps to pre-allocate
+                in the buffer.
+            time_major (Optional[bool]): Whether to preallocate buffers and
+                collect samples in time-major fashion (TxBx...).
+            shift_before (int): The additional number of time slots to
+                pre-allocate at the beginning of a time window (for possible
+                underlying data column shifts, e.g. PREV_ACTIONS).
+            shift_after (int): The additional number of time slots to
+                pre-allocate at the end of a time window (for possible
+                underlying data column shifts, e.g. NEXT_OBS).
+        """
+
+        self.num_agents = num_agents or 100
+        self.num_timesteps = num_timesteps
+        self.time_major = time_major
+        # `shift_before must at least be 1 for the init obs timestep.
+        self.shift_before = max(shift_before, 1)
+        self.shift_after = shift_after
+
+        # The offset on the agent dim to start the next SampleBatch build from.
+        self.sample_batch_offset = 0
+
+        # The actual underlying data-buffers.
+        self.buffers = {}
+        self.postprocessed_agents = [False] * self.num_agents
+
+        # Next agent-slot to be used by a new agent/env combination.
+        self.agent_slot_cursor = 0
+        # Maps agent/episode ID/chunk-num to an agent slot.
+        self.agent_key_to_slot = {}
+        # Maps agent/episode ID to the last chunk-num.
+        self.agent_key_to_chunk_num = {}
+        # Maps agent slot number to agent keys.
+        self.slot_to_agent_key = [None] * self.num_agents
+        # Maps agent/episode ID/chunk-num to a time step cursor.
+        self.agent_key_to_timestep = {}
+
+        # Total timesteps taken in the env over all agents since last reset.
+        self.timesteps_since_last_reset = 0
+
+        # Indices (T,B) to pick from the buffers for the next forward pass.
+        self.forward_pass_indices = [[], []]
+        self.forward_pass_size = 0
+        # Maps index from the forward pass batch to (agent_id, episode_id,
+        # env_id) tuple.
+        self.forward_pass_index_to_agent_info = {}
+        self.agent_key_to_forward_pass_index = {}
+
+    def add_init_obs(self, episode_id: EpisodeID, agent_id: AgentID,
+                     env_id: EnvID, chunk_num: int,
+                     init_obs: TensorType) -> None:
+        """Adds a single initial observation (after env.reset()) to the buffer.
+
+        Args:
+            episode_id (EpisodeID): Unique ID for the episode we are adding the
+                initial observation for.
+            agent_id (AgentID): Unique ID for the agent we are adding the
+                initial observation for.
+            env_id (EnvID): The env ID to which `init_obs` belongs.
+            chunk_num (int): The time-chunk number (0-based). Some episodes
+                may last for longer than self.num_timesteps and therefore
+                have to be chopped into chunks.
+            init_obs (TensorType): Initial observation (after env.reset()).
+        """
+        agent_key = (agent_id, episode_id, chunk_num)
+        agent_slot = self.agent_slot_cursor
+        self.agent_key_to_slot[agent_key] = agent_slot
+        self.agent_key_to_chunk_num[agent_key[:2]] = chunk_num
+        self.slot_to_agent_key[agent_slot] = agent_key
+        self._next_agent_slot()
+
+        if SampleBatch.OBS not in self.buffers:
+            self._build_buffers(single_row={SampleBatch.OBS: init_obs})
+        if self.time_major:
+            self.buffers[SampleBatch.OBS][self.shift_before-1, agent_slot] = \
+                init_obs
+        else:
+            self.buffers[SampleBatch.OBS][agent_slot, self.shift_before-1] = \
+                init_obs
+        self.agent_key_to_timestep[agent_key] = self.shift_before
+
+        self._add_to_next_inference_call(agent_key, env_id, agent_slot,
+                                         self.shift_before - 1)
+
+    def add_action_reward_next_obs(
+            self, episode_id: EpisodeID, agent_id: AgentID, env_id: EnvID,
+            agent_done: bool, values: Dict[str, TensorType]) -> None:
+        """Add the given dictionary (row) of values to this batch.
+
+        Args:
+            episode_id (EpisodeID): Unique ID for the episode we are adding the
+                values for.
+            agent_id (AgentID): Unique ID for the agent we are adding the
+                values for.
+            env_id (EnvID): The env ID to which the given data belongs.
+            agent_done (bool): Whether next obs should not be used for an
+                upcoming inference call. Default: False = next-obs should be
+                used for upcoming inference.
+            values (Dict[str, TensorType]): Data dict (interpreted as a single
+                row) to be added to buffer. Must contain keys:
+                SampleBatch.ACTIONS, REWARDS, DONES, and NEXT_OBS.
+        """
+        assert (SampleBatch.ACTIONS in values and SampleBatch.REWARDS in values
+                and SampleBatch.NEXT_OBS in values
+                and SampleBatch.DONES in values)
+
+        assert SampleBatch.OBS not in values
+        values[SampleBatch.OBS] = values[SampleBatch.NEXT_OBS]
+        del values[SampleBatch.NEXT_OBS]
+
+        chunk_num = self.agent_key_to_chunk_num[(agent_id, episode_id)]
+        agent_key = (agent_id, episode_id, chunk_num)
+        agent_slot = self.agent_key_to_slot[agent_key]
+        ts = self.agent_key_to_timestep[agent_key]
+        for k, v in values.items():
+            if k not in self.buffers:
+                self._build_buffers(single_row=values)
+            if self.time_major:
+                self.buffers[k][ts, agent_slot] = v
+            else:
+                self.buffers[k][agent_slot, ts] = v
+        self.agent_key_to_timestep[agent_key] += 1
+
+        # Time-axis is "full" -> Cut-over to new chunk (only if not DONE).
+        if self.agent_key_to_timestep[
+            agent_key] - self.shift_before == self.num_timesteps and \
+                not values[SampleBatch.DONES]:
+            self._new_chunk_from(agent_slot, agent_key,
+                                 self.agent_key_to_timestep[agent_key])
+
+        self.timesteps_since_last_reset += 1
+
+        if not agent_done:
+            self._add_to_next_inference_call(agent_key, env_id, agent_slot, ts)
+
+    def get_inference_input_dict(self, view_reqs: Dict[str, ViewRequirement]
+                                 ) -> Dict[str, TensorType]:
+        """Returns an input_dict for an (inference) forward pass.
+
+        The input_dict can then be used for action computations inside a
+        Policy via `Policy.compute_actions_from_input_dict()`.
+
+        Args:
+            view_reqs (Dict[str, ViewRequirement]): The view requirements
+                dict to use.
+
+        Returns:
+            Dict[str, TensorType]: The input_dict to be passed into the ModelV2
+                for inference/training.
+
+        Examples:
+            >>> obs, r, done, info = env.step(action)
+            >>> collector.add_action_reward_next_obs(12345, 0, "pol0", {
+            ...     "action": action, "obs": obs, "reward": r, "done": done
+            ... })
+            >>> input_dict = collector.get_inference_input_dict(policy.model)
+            >>> action = policy.compute_actions_from_input_dict(input_dict)
+            >>> # repeat
+        """
+        input_dict = {}
+        for view_col, view_req in view_reqs.items():
+            # Create the batch of data from the different buffers.
+            data_col = view_req.data_col or view_col
+            if data_col not in self.buffers:
+                self._build_buffers({data_col: view_req.space.sample()})
+
+            indices = self.forward_pass_indices
+            if self.time_major:
+                input_dict[view_col] = self.buffers[data_col][indices]
+            else:
+                if isinstance(view_req.shift, (list, tuple)):
+                    time_indices = \
+                        np.array(view_req.shift) + np.array(indices[0])
+                    input_dict[view_col] = self.buffers[data_col][indices[1],
+                                                                  time_indices]
+                else:
+                    input_dict[view_col] = \
+                        self.buffers[data_col][indices[1], indices[0]]
+
+        self._reset_inference_call()
+
+        return input_dict
+
+    def get_postprocessing_sample_batches(
+            self,
+            episode: MultiAgentEpisode,
+            view_reqs: Dict[str, ViewRequirement]) -> \
+            Dict[AgentID, SampleBatch]:
+        """Returns a SampleBatch object ready for postprocessing.
+
+        Args:
+            episode (MultiAgentEpisode): The MultiAgentEpisode object to
+                get the to-be-postprocessed SampleBatches for.
+            view_reqs (Dict[str, ViewRequirement]): The view requirements dict
+                to use for creating the SampleBatch from our buffers.
+
+        Returns:
+            Dict[AgentID, SampleBatch]: The sample batch objects to be passed
+                to `Policy.postprocess_trajectory()`.
+        """
+        # Loop through all agents and create a SampleBatch
+        # (as "view"; no copying).
+
+        # Construct the SampleBatch-dict.
+        sample_batch_data = {}
+
+        range_ = self.agent_slot_cursor - self.sample_batch_offset
+        if range_ < 0:
+            range_ = self.num_agents + range_
+        for i in range(range_):
+            agent_slot = self.sample_batch_offset + i
+            if agent_slot >= self.num_agents:
+                agent_slot = agent_slot % self.num_agents
+            # Do not postprocess the same slot twice.
+            if self.postprocessed_agents[agent_slot]:
+                continue
+            agent_key = self.slot_to_agent_key[agent_slot]
+            # Skip other episodes (if episode provided).
+            if episode and agent_key[1] != episode.episode_id:
+                continue
+            end = self.agent_key_to_timestep[agent_key]
+            # Do not build any empty SampleBatches.
+            if end == self.shift_before:
+                continue
+            self.postprocessed_agents[agent_slot] = True
+
+            assert agent_key not in sample_batch_data
+            sample_batch_data[agent_key] = {}
+            batch = sample_batch_data[agent_key]
+
+            for view_col, view_req in view_reqs.items():
+                # Skip columns that will only get added through postprocessing
+                # (these may not even exist yet).
+                if view_req.created_during_postprocessing:
+                    continue
+
+                data_col = view_req.data_col or view_col
+                shift = view_req.shift
+                if data_col == SampleBatch.OBS:
+                    shift -= 1
+
+                batch[view_col] = self.buffers[data_col][
+                    self.shift_before + shift:end + shift, agent_slot]
+
+        batches = {}
+        for agent_key, data in sample_batch_data.items():
+            batches[agent_key] = SampleBatch(data)
+        return batches
+
+    def get_train_sample_batch_and_reset(self, view_reqs) -> SampleBatch:
+        """Returns the accumulated sample batche for this policy.
+
+        This is usually called to collect samples for policy training.
+
+        Returns:
+            SampleBatch: Returns the accumulated sample batch for this
+                policy.
+        """
+        seq_lens = [
+            self.agent_key_to_timestep[k] - self.shift_before
+            for k in self.slot_to_agent_key if k is not None
+        ]
+        first_zero_len = len(seq_lens)
+        if seq_lens[-1] == 0:
+            first_zero_len = seq_lens.index(0)
+            # Assert that all zeros lie at the end of the seq_lens array.
+            try:
+                assert all(seq_lens[i] == 0
+                           for i in range(first_zero_len, len(seq_lens)))
+            except AssertionError as e:
+                print()
+                raise e
+
+        t_start = self.shift_before
+        t_end = t_start + self.num_timesteps
+
+        # The agent_slot cursor that points to the newest agent-slot that
+        # actually already has at least 1 timestep of data (thus it excludes
+        # just-rolled over chunks (which only have the initial obs in them)).
+        valid_agent_cursor = \
+            (self.agent_slot_cursor - (len(seq_lens) - first_zero_len)) % \
+            self.num_agents
+
+        # Construct the view dict.
+        view = {}
+        for view_col, view_req in view_reqs.items():
+            data_col = view_req.data_col or view_col
+            assert data_col in self.buffers
+            # For OBS, indices must be shifted by -1.
+            extra_shift = 0 if data_col != SampleBatch.OBS else -1
+            # If agent_slot has been rolled-over to beginning, we have to copy
+            # here.
+            if valid_agent_cursor < self.sample_batch_offset:
+                time_slice = self.buffers[data_col][t_start + extra_shift:
+                                                    t_end + extra_shift]
+                one_ = time_slice[:, self.sample_batch_offset:]
+                two_ = time_slice[:, :valid_agent_cursor]
+                if torch and isinstance(time_slice, torch.Tensor):
+                    view[view_col] = torch.cat([one_, two_], dim=1)
+                else:
+                    view[view_col] = np.concatenate([one_, two_], axis=1)
+            else:
+                view[view_col] = \
+                    self.buffers[data_col][
+                    t_start + extra_shift:t_end + extra_shift,
+                    self.sample_batch_offset:valid_agent_cursor]
+
+        # Copy all still ongoing trajectories to new agent slots
+        # (including the ones that just started (are seq_len=0)).
+        new_chunk_args = []
+        for i, seq_len in enumerate(seq_lens):
+            if seq_len < self.num_timesteps:
+                agent_slot = self.sample_batch_offset + i
+                if agent_slot >= self.num_agents:
+                    agent_slot = agent_slot % self.num_agents
+                if not self.buffers[SampleBatch.
+                                    DONES][seq_len - 1 +
+                                           self.shift_before][agent_slot]:
+                    agent_key = self.slot_to_agent_key[agent_slot]
+                    new_chunk_args.append(
+                        (agent_slot, agent_key,
+                         self.agent_key_to_timestep[agent_key]))
+        # Cut out all 0 seq-lens.
+        seq_lens = seq_lens[:first_zero_len]
+        batch = SampleBatch(
+            view, _seq_lens=np.array(seq_lens), _time_major=True)
+
+        # Reset everything for new data.
+        self.postprocessed_agents = [False] * self.num_agents
+        self.agent_key_to_slot.clear()
+        self.agent_key_to_chunk_num.clear()
+        self.slot_to_agent_key = [None] * self.num_agents
+        self.agent_key_to_timestep.clear()
+        self.timesteps_since_last_reset = 0
+        self.forward_pass_size = 0
+        self.sample_batch_offset = self.agent_slot_cursor
+
+        for args in new_chunk_args:
+            self._new_chunk_from(*args)
+
+        return batch
+
+    def _build_buffers(self, single_row: Dict[str, TensorType]) -> None:
+        """Builds the internal data buffers based on a single given row.
+
+        Args:
+            single_row (Dict[str, TensorType]): A single datarow with one or
+                more columns (str as key, np.ndarray|tensor as data).
+        """
+        time_size = self.num_timesteps + self.shift_before + self.shift_after
+        for col, data in single_row.items():
+            if col in self.buffers:
+                continue
+            base_shape = (time_size, self.num_agents) if self.time_major else \
+                (self.num_agents, time_size)
+            # Python primitive -> np.array.
+            if isinstance(data, (int, float, bool)):
+                t_ = type(data)
+                dtype = np.float32 if t_ == float else \
+                    np.int32 if type(data) == int else np.bool_
+                self.buffers[col] = np.zeros(shape=base_shape, dtype=dtype)
+            # np.ndarray, torch.Tensor, or tf.Tensor.
+            else:
+                shape = base_shape + data.shape
+                dtype = data.dtype
+                if torch and isinstance(data, torch.Tensor):
+                    self.buffers[col] = torch.zeros(
+                        *shape, dtype=dtype, device=data.device)
+                elif tf and isinstance(data, tf.Tensor):
+                    self.buffers[col] = tf.zeros(shape=shape, dtype=dtype)
+                else:
+                    self.buffers[col] = np.zeros(shape=shape, dtype=dtype)
+
+    def _next_agent_slot(self):
+        """Starts a new agent slot at the end of the agent-axis.
+
+        Also makes sure, the new slot is not taken yet.
+        """
+        self.agent_slot_cursor += 1
+        if self.agent_slot_cursor >= self.num_agents:
+            self.agent_slot_cursor = 0
+        # Just make sure, there is space in our buffer.
+        assert self.slot_to_agent_key[self.agent_slot_cursor] is None
+
+    def _new_chunk_from(self, agent_slot, agent_key, timestep):
+        """Creates a new time-window (chunk) given an agent.
+
+        The agent may already have an unfinished episode going on (in a
+        previous chunk). The end of that previous chunk will be copied to the
+        beginning of the new one for proper data-shift handling (e.g.
+        PREV_ACTIONS/REWARDS).
+
+        Args:
+            agent_slot (int): The agent to start a new chunk for (from an
+                ongoing episode (chunk)).
+            agent_key (Tuple[AgentID, EpisodeID, int]): The internal key to
+                identify an active agent in some episode.
+            timestep (int): The timestep in the old chunk being continued.
+        """
+        new_agent_slot = self.agent_slot_cursor
+        # Increase chunk num by 1.
+        new_agent_key = agent_key[:2] + (agent_key[2] + 1, )
+        # Copy relevant timesteps at end of old chunk into new one.
+        if self.time_major:
+            for k in self.buffers.keys():
+                self.buffers[k][0:self.shift_before, new_agent_slot] = \
+                    self.buffers[k][
+                    timestep - self.shift_before:timestep, agent_slot]
+        else:
+            for k in self.buffers.keys():
+                self.buffers[k][new_agent_slot, 0:self.shift_before] = \
+                    self.buffers[k][
+                    agent_slot, timestep - self.shift_before:timestep]
+
+        self.agent_key_to_slot[new_agent_key] = new_agent_slot
+        self.agent_key_to_chunk_num[new_agent_key[:2]] = new_agent_key[2]
+        self.slot_to_agent_key[new_agent_slot] = new_agent_key
+        self._next_agent_slot()
+        self.agent_key_to_timestep[new_agent_key] = self.shift_before
+
+    def _add_to_next_inference_call(self, agent_key, env_id, agent_slot,
+                                    timestep):
+        """Registers given T and B (agent_slot) for get_inference_input_dict.
+
+        Calling `get_inference_input_dict` will produce an input_dict (for
+        Policy.compute_actions_from_input_dict) with all registered agent/time
+        indices and then automatically reset the registry.
+
+        Args:
+            agent_key (Tuple[AgentID, EpisodeID, int]): The internal key to
+                identify an active agent in some episode.
+            env_id (EnvID): The env ID of the given agent.
+            agent_slot (int): The agent_slot to register (B axis).
+            timestep (int): The timestep to register (T axis).
+        """
+        idx = self.forward_pass_size
+        self.forward_pass_index_to_agent_info[idx] = (agent_key[0],
+                                                      agent_key[1], env_id)
+        self.agent_key_to_forward_pass_index[agent_key[:2]] = idx
+        if self.forward_pass_size == 0:
+            self.forward_pass_indices[0].clear()
+            self.forward_pass_indices[1].clear()
+        self.forward_pass_indices[0].append(timestep)
+        self.forward_pass_indices[1].append(agent_slot)
+        self.forward_pass_size += 1
+
+    def _reset_inference_call(self):
+        """Resets indices for the next inference call.
+
+        After calling this, new calls to `add_init_obs()` and
+        `add_action_reward_next_obs()` will count for the next input_dict
+        returned by `get_inference_input_dict()`.
+        """
+        self.forward_pass_size = 0

rllib/evaluation/rollout_worker.py

@@ -6,7 +6,7 @@ import pickle
 import platform
 import os
 from typing import Callable, Any, List, Dict, Tuple, Union, Optional, \
-    TYPE_CHECKING, TypeVar
+    TYPE_CHECKING, Type, TypeVar
 
 import ray
 from ray.rllib.env.atari_wrappers import wrap_deepmind, is_atari
@@ -36,15 +36,15 @@ from ray.rllib.utils.filter import get_filter, Filter
 from ray.rllib.utils.framework import try_import_tf, try_import_torch
 from ray.rllib.utils.sgd import do_minibatch_sgd
 from ray.rllib.utils.tf_run_builder import TFRunBuilder
-from ray.rllib.utils.typing import EnvType, AgentID, PolicyID, EnvConfigDict, \
-    ModelConfigDict, TrainerConfigDict, SampleBatchType, ModelWeights, \
-    ModelGradients, MultiAgentPolicyConfigDict
+from ray.rllib.utils.typing import AgentID, EnvConfigDict, EnvType, \
+    ModelConfigDict, ModelGradients, ModelWeights, \
+    MultiAgentPolicyConfigDict, PartialTrainerConfigDict, PolicyID, \
+    SampleBatchType, TrainerConfigDict
 from ray.util.debug import log_once, disable_log_once_globally, \
     enable_periodic_logging
 from ray.util.iter import ParallelIteratorWorker
 
 if TYPE_CHECKING:
-    from ray.rllib.agents.callbacks import DefaultCallbacks
     from ray.rllib.evaluation.observation_function import ObservationFunction
 
 # Generic type var for foreach_* methods.
@@ -129,63 +129,67 @@ class RolloutWorker(ParallelIteratorWorker):
             resources=resources)(cls)
 
     @DeveloperAPI
-    def __init__(self,
-                 env_creator: Callable[[EnvContext], EnvType],
-                 policy: type,
-                 policy_mapping_fn: Callable[[AgentID], PolicyID] = None,
-                 policies_to_train: List[PolicyID] = None,
-                 tf_session_creator: Callable[[], Any] = None,
-                 rollout_fragment_length: int = 100,
-                 batch_mode: str = "truncate_episodes",
-                 episode_horizon: int = None,
-                 preprocessor_pref: str = "deepmind",
-                 sample_async: bool = False,
-                 compress_observations: bool = False,
-                 num_envs: int = 1,
-                 observation_fn: "ObservationFunction" = None,
-                 observation_filter: str = "NoFilter",
-                 clip_rewards: bool = None,
-                 clip_actions: bool = True,
-                 env_config: EnvConfigDict = None,
-                 model_config: ModelConfigDict = None,
-                 policy_config: TrainerConfigDict = None,
-                 worker_index: int = 0,
-                 num_workers: int = 0,
-                 monitor_path: str = None,
-                 log_dir: str = None,
-                 log_level: str = None,
-                 callbacks: "DefaultCallbacks" = None,
-                 input_creator: Callable[[
-                     IOContext
-                 ], InputReader] = lambda ioctx: ioctx.default_sampler_input(),
-                 input_evaluation: List[str] = frozenset([]),
-                 output_creator: Callable[
-                     [IOContext], OutputWriter] = lambda ioctx: NoopOutput(),
-                 remote_worker_envs: bool = False,
-                 remote_env_batch_wait_ms: int = 0,
-                 soft_horizon: bool = False,
-                 no_done_at_end: bool = False,
-                 seed: int = None,
-                 extra_python_environs: dict = None,
-                 fake_sampler: bool = False):
+    def __init__(
+            self,
+            env_creator: Callable[[EnvContext], EnvType],
+            policy: Union[type, Dict[str, Tuple[Optional[
+                type], gym.Space, gym.Space, PartialTrainerConfigDict]]],
+            policy_mapping_fn: Callable[[AgentID], PolicyID] = None,
+            policies_to_train: Optional[List[PolicyID]] = None,
+            tf_session_creator: Optional[Callable[[], "tf1.Session"]] = None,
+            rollout_fragment_length: int = 100,
+            batch_mode: str = "truncate_episodes",
+            episode_horizon: int = None,
+            preprocessor_pref: str = "deepmind",
+            sample_async: bool = False,
+            compress_observations: bool = False,
+            num_envs: int = 1,
+            observation_fn: "ObservationFunction" = None,
+            observation_filter: str = "NoFilter",
+            clip_rewards: bool = None,
+            clip_actions: bool = True,
+            env_config: EnvConfigDict = None,
+            model_config: ModelConfigDict = None,
+            policy_config: TrainerConfigDict = None,
+            worker_index: int = 0,
+            num_workers: int = 0,
+            monitor_path: str = None,
+            log_dir: str = None,
+            log_level: str = None,
+            callbacks: Type["DefaultCallbacks"] = None,
+            input_creator: Callable[[
+                IOContext
+            ], InputReader] = lambda ioctx: ioctx.default_sampler_input(),
+            input_evaluation: List[str] = frozenset([]),
+            output_creator: Callable[
+                [IOContext], OutputWriter] = lambda ioctx: NoopOutput(),
+            remote_worker_envs: bool = False,
+            remote_env_batch_wait_ms: int = 0,
+            soft_horizon: bool = False,
+            no_done_at_end: bool = False,
+            seed: int = None,
+            extra_python_environs: dict = None,
+            fake_sampler: bool = False):
         """Initialize a rollout worker.
 
-        Arguments:
-            env_creator (func): Function that returns a gym.Env given an
-                EnvContext wrapped configuration.
-            policy (class|dict): Either a class implementing
-                Policy, or a dictionary of policy id strings to
-                (Policy, obs_space, action_space, config) tuples. If a
-                dict is specified, then we are in multi-agent mode and a
-                policy_mapping_fn should also be set.
-            policy_mapping_fn (func): A function that maps agent ids to
-                policy ids in multi-agent mode. This function will be called
-                each time a new agent appears in an episode, to bind that agent
-                to a policy for the duration of the episode.
-            policies_to_train (list): Optional list of policies to train,
-                or None for all policies.
-            tf_session_creator (func): A function that returns a TF session.
-                This is optional and only useful with TFPolicy.
+        Args:
+            env_creator (Callable[[EnvContext], EnvType]): Function that
+                returns a gym.Env given an EnvContext wrapped configuration.
+            policy (Union[type, Dict[str, Tuple[Optional[type], gym.Space,
+                gym.Space, PartialTrainerConfigDict]]]): Either a Policy class
+                or a dict of policy id strings to
+                (Policy (None for default), obs_space, action_space,
+                config)-tuples. If a dict is specified, then we are in
+                multi-agent mode and a policy_mapping_fn should also be set.
+            policy_mapping_fn (Callable[[AgentID], PolicyID]): A function that
+                maps agent ids to policy ids in multi-agent mode. This function
+                will be called each time a new agent appears in an episode, to
+                bind that agent to a policy for the duration of the episode.
+            policies_to_train (Optional[List[PolicyID]]): Optional list of
+                policies to train, or None for all policies.
+            tf_session_creator (Optional[Callable[[], tf1.Session]]): A
+                function that returns a TF session. This is optional and only
+                useful with TFPolicy.
             rollout_fragment_length (int): The target number of env transitions
                 to include in each sample batch returned from this worker.
             batch_mode (str): One of the following batch modes:
@@ -221,10 +225,11 @@ class RolloutWorker(ParallelIteratorWorker):
                 only.
             clip_actions (bool): Whether to clip action values to the range
                 specified by the policy action space.
-            env_config (dict): Config to pass to the env creator.
-            model_config (dict): Config to use when creating the policy model.
-            policy_config (dict): Config to pass to the policy. In the
-                multi-agent case, this config will be merged with the
+            env_config (EnvConfigDict): Config to pass to the env creator.
+            model_config (ModelConfigDict): Config to use when creating the
+                policy model.
+            policy_config (TrainerConfigDict): Config to pass to the policy.
+                In the multi-agent case, this config will be merged with the
                 per-policy configs specified by `policy`.
             worker_index (int): For remote workers, this should be set to a
                 non-zero and unique value. This index is passed to created envs
@@ -236,17 +241,19 @@ class RolloutWorker(ParallelIteratorWorker):
             log_dir (str): Directory where logs can be placed.
             log_level (str): Set the root log level on creation.
             callbacks (DefaultCallbacks): Custom training callbacks.
-            input_creator (func): Function that returns an InputReader object
-                for loading previous generated experiences.
-            input_evaluation (list): How to evaluate the policy performance.
-                This only makes sense to set when the input is reading offline
-                data. The possible values include:
+            input_creator (Callable[[IOContext], InputReader]): Function that
+                returns an InputReader object for loading previous generated
+                experiences.
+            input_evaluation (List[str]): How to evaluate the policy
+                performance. This only makes sense to set when the input is
+                reading offline data. The possible values include:
                   - "is": the step-wise importance sampling estimator.
                   - "wis": the weighted step-wise is estimator.
                   - "simulation": run the environment in the background, but
                     use this data for evaluation only and never for learning.
-            output_creator (func): Function that returns an OutputWriter object
-                for saving generated experiences.
+            output_creator (Callable[[IOContext], OutputWriter]): Function that
+                returns an OutputWriter object for saving generated
+                experiences.
             remote_worker_envs (bool): If using num_envs > 1, whether to create
                 those new envs in remote processes instead of in the current
                 process. This adds overheads, but can make sense if your envs

rllib/evaluation/sample_batch_builder.py

@@ -44,7 +44,7 @@ class SampleBatchBuilder:
         self.count = 0
 
     @PublicAPI
-    def add_values(self, **values: Dict[str, Any]) -> None:
+    def add_values(self, **values: Any) -> None:
         """Add the given dictionary (row) of values to this batch."""
 
         for k, v in values.items():
@@ -138,7 +138,7 @@ class MultiAgentSampleBatchBuilder:
 
     @DeveloperAPI
     def add_values(self, agent_id: AgentID, policy_id: AgentID,
-                   **values: Dict[str, Any]) -> None:
+                   **values: Any) -> None:
         """Add the given dictionary (row) of values to this batch.
 
         Arguments:

rllib/evaluation/sample_collector.py

@@ -3,7 +3,6 @@ import logging
 from typing import Dict, Optional
 
 from ray.rllib.evaluation.episode import MultiAgentEpisode
-from ray.rllib.models.modelv2 import ModelV2
 from ray.rllib.utils.typing import AgentID, EpisodeID, PolicyID, \
     TensorType
 
@@ -101,15 +100,15 @@ class _SampleCollector(metaclass=ABCMeta):
         raise NotImplementedError
 
     @abstractmethod
-    def get_inference_input_dict(self, model: ModelV2) -> \
+    def get_inference_input_dict(self, policy_id: PolicyID) -> \
             Dict[str, TensorType]:
-        """Returns input_dict for an inference forward pass from our data.
+        """Returns an input_dict for an (inference) forward pass from our data.
 
-        The input_dict can then be used for action computations.
+        The input_dict can then be used for action computations inside a
+        Policy via `Policy.compute_actions_from_input_dict()`.
 
         Args:
-            model (ModelV2): The ModelV2 object for which to generate the view
-                (input_dict) from `data`.
+            policy_id (PolicyID): The Policy ID to get the input dict for.
 
         Returns:
             Dict[str, TensorType]: The input_dict to be passed into the ModelV2
@@ -155,23 +154,32 @@ class _SampleCollector(metaclass=ABCMeta):
         """
         raise NotImplementedError
 
+    @abstractmethod
+    def check_missing_dones(self, episode_id: EpisodeID) -> None:
+        """Checks whether given episode is properly terminated with done=True.
+
+        This applies to all agents in the episode.
+
+        Args:
+            episode_id (EpisodeID): The episode ID to check for proper
+                termination.
+
+        Raises:
+            ValueError: If `episode` has no done=True at the end.
+        """
+        raise NotImplementedError
+
     @abstractmethod
     def get_multi_agent_batch_and_reset(self):
         """Returns the accumulated sample batches for each policy.
 
         Any unprocessed rows will be first postprocessed with a policy
-        postprocessor. The internal state of this builder will be reset.
-
-        Args:
-            episode (Optional[MultiAgentEpisode]): The Episode object that
-                holds this MultiAgentBatchBuilder object or None.
+        postprocessor. The internal state of this builder will be reset to
+        start the next batch.
+        This is usually called to collect samples for policy training.
 
         Returns:
             MultiAgentBatch: Returns the accumulated sample batches for each
                 policy inside one MultiAgentBatch object.
         """
         raise NotImplementedError
-
-    @abstractmethod
-    def check_missing_dones(self, episode_id: EpisodeID) -> None:
-        raise NotImplementedError

rllib/evaluation/tests/test_trajectories.py

@@ -1,70 +0,0 @@
-from gym.spaces import Box, Discrete
-import numpy as np
-import unittest
-
-from ray.rllib.evaluation.trajectory import Trajectory
-
-
-class TestTrajectories(unittest.TestCase):
-    """Tests Trajectory classes."""
-
-    def test_trajectory(self):
-        """Tests the Trajectory class."""
-
-        buffer_size = 5
-
-        # Small trajecory object for testing purposes.
-        trajectory = Trajectory(buffer_size=buffer_size)
-        self.assertEqual(trajectory.cursor, 0)
-        self.assertEqual(trajectory.timestep, 0)
-        self.assertEqual(trajectory.sample_batch_offset, 0)
-        assert not trajectory.buffers
-        observation_space = Box(-1.0, 1.0, shape=(3, ))
-        action_space = Discrete(2)
-        trajectory.add_init_obs(
-            env_id=0,
-            agent_id="agent",
-            policy_id="policy",
-            init_obs=observation_space.sample())
-        self.assertEqual(trajectory.cursor, 0)
-        self.assertEqual(trajectory.initial_obs.shape, observation_space.shape)
-
-        # Fill up the buffer and make it extend if it hits the limit.
-        cur_buffer_size = buffer_size
-        for i in range(buffer_size + 1):
-            trajectory.add_action_reward_next_obs(
-                env_id=0,
-                agent_id="agent",
-                policy_id="policy",
-                values=dict(
-                    t=i,
-                    actions=action_space.sample(),
-                    rewards=1.0,
-                    dones=i == buffer_size,
-                    new_obs=observation_space.sample(),
-                    action_logp=-0.5,
-                    action_dist_inputs=np.array([[0.5, 0.5]]),
-                ))
-            self.assertEqual(trajectory.cursor, i + 1)
-            self.assertEqual(trajectory.timestep, i + 1)
-            self.assertEqual(trajectory.sample_batch_offset, 0)
-            if i == buffer_size - 1:
-                cur_buffer_size *= 2
-            self.assertEqual(
-                len(trajectory.buffers["new_obs"]), cur_buffer_size)
-            self.assertEqual(
-                len(trajectory.buffers["rewards"]), cur_buffer_size)
-
-        # Create a SampleBatch from the Trajectory and reset it.
-        batch = trajectory.get_sample_batch_and_reset()
-        self.assertEqual(batch.count, buffer_size + 1)
-        # Make sure, Trajectory was reset properly.
-        self.assertEqual(trajectory.cursor, buffer_size + 1)
-        self.assertEqual(trajectory.timestep, 0)
-        self.assertEqual(trajectory.sample_batch_offset, buffer_size + 1)
-
-
-if __name__ == "__main__":
-    import pytest
-    import sys
-    sys.exit(pytest.main(["-v", __file__]))

rllib/policy/view_requirement.py

@@ -1,5 +1,5 @@
 import gym
-from typing import Optional
+from typing import List, Optional, Union
 
 from ray.rllib.utils.framework import try_import_torch
 
@@ -21,7 +21,7 @@ class ViewRequirement:
 
     Examples:
         >>> # The default ViewRequirement for a Model is:
-        >>> req = [ModelV2].inference_view_requirements()
+        >>> req = [ModelV2].inference_view_requirements
         >>> print(req)
         {"obs": ViewRequirement(shift=0)}
     """
@@ -29,7 +29,8 @@ class ViewRequirement:
     def __init__(self,
                  data_col: Optional[str] = None,
                  space: gym.Space = None,
-                 shift: int = 0):
+                 shift: Union[int, List[int]] = 0,
+                 created_during_postprocessing: bool = False):
         """Initializes a ViewRequirement object.
 
         Args:
@@ -39,15 +40,18 @@ class ViewRequirement:
             space (gym.Space): The gym Space used in case we need to pad data
                 in inaccessible areas of the trajectory (t<0 or t>H).
                 Default: Simple box space, e.g. rewards.
-            shift (Union[List[int], int]): Single shift value of list of
+            shift (Union[int, List[int]]): Single shift value of list of
                 shift values to use relative to the underlying `data_col`.
                 Example: For a view column "prev_actions", you can set
                 `data_col="actions"` and `shift=-1`.
                 Example: For a view column "obs" in an Atari framestacking
                 fashion, you can set `data_col="obs"` and
                 `shift=[-3, -2, -1, 0]`.
+            created_during_postprocessing (bool): Whether this column only gets
+                created during postprocessing.
         """
         self.data_col = data_col
         self.space = space or gym.spaces.Box(
             float("-inf"), float("inf"), shape=())
         self.shift = shift
+        self.created_during_postprocessing = created_during_postprocessing

rllib/rollout.py

@@ -3,13 +3,13 @@
 import argparse
 import collections
 import copy
+import gym
 import json
 import os
 from pathlib import Path
 import pickle
 import shelve
 
-import gym
 import ray
 from ray.rllib.env import MultiAgentEnv
 from ray.rllib.env.base_env import _DUMMY_AGENT_ID