chore: import upstream snapshot with attribution

rllib/execution/__init__.py

@@ -0,0 +1,23 @@
+from ray.rllib.execution.learner_thread import LearnerThread
+from ray.rllib.execution.minibatch_buffer import MinibatchBuffer
+from ray.rllib.execution.multi_gpu_learner_thread import MultiGPULearnerThread
+from ray.rllib.execution.replay_ops import SimpleReplayBuffer
+from ray.rllib.execution.rollout_ops import (
+    standardize_fields,
+    synchronous_parallel_sample,
+)
+from ray.rllib.execution.train_ops import (
+    multi_gpu_train_one_step,
+    train_one_step,
+)
+
+__all__ = [
+    "multi_gpu_train_one_step",
+    "standardize_fields",
+    "synchronous_parallel_sample",
+    "train_one_step",
+    "LearnerThread",
+    "MultiGPULearnerThread",
+    "SimpleReplayBuffer",
+    "MinibatchBuffer",
+]

rllib/execution/buffers/mixin_replay_buffer.py

@@ -0,0 +1,173 @@
+import collections
+import platform
+import random
+from typing import Optional
+
+from ray.rllib.execution.replay_ops import SimpleReplayBuffer
+from ray.rllib.policy.sample_batch import DEFAULT_POLICY_ID, concat_samples
+from ray.rllib.utils.annotations import OldAPIStack
+from ray.rllib.utils.replay_buffers.multi_agent_replay_buffer import ReplayMode
+from ray.rllib.utils.replay_buffers.replay_buffer import _ALL_POLICIES
+from ray.rllib.utils.typing import PolicyID, SampleBatchType
+from ray.util.timer import _Timer
+
+
+@OldAPIStack
+class MixInMultiAgentReplayBuffer:
+    """This buffer adds replayed samples to a stream of new experiences.
+
+    - Any newly added batch (`add()`) is immediately returned upon
+    the next `replay` call (close to on-policy) as well as being moved
+    into the buffer.
+    - Additionally, a certain number of old samples is mixed into the
+    returned sample according to a given "replay ratio".
+    - If >1 calls to `add()` are made without any `replay()` calls
+    in between, all newly added batches are returned (plus some older samples
+    according to the "replay ratio").
+
+    .. testcode::
+
+        from ray.rllib.execution.buffers.mixin_replay_buffer import (
+            MixInMultiAgentReplayBuffer)
+        from ray.rllib.policy.sample_batch import SampleBatch
+        # replay ratio 0.66 (2/3 replayed, 1/3 new samples):
+        buffer = MixInMultiAgentReplayBuffer(capacity=100,
+                                             replay_ratio=0.66)
+        A, B, C = (SampleBatch({"obs": [1]}), SampleBatch({"obs": [2]}),
+            SampleBatch({"obs": [3]}))
+        buffer.add(A)
+        buffer.add(B)
+        buffer.add(B)
+        print(buffer.replay()["obs"])
+
+    .. testoutput::
+        :hide:
+
+        ...
+    """
+
+    def __init__(
+        self,
+        capacity: int,
+        replay_ratio: float,
+        replay_mode: ReplayMode = ReplayMode.INDEPENDENT,
+    ):
+        """Initializes MixInReplay instance.
+
+        Args:
+            capacity: Number of batches to store in total.
+            replay_ratio: Ratio of replayed samples in the returned
+                batches. E.g. a ratio of 0.0 means only return new samples
+                (no replay), a ratio of 0.5 means always return newest sample
+                plus one old one (1:1), a ratio of 0.66 means always return
+                the newest sample plus 2 old (replayed) ones (1:2), etc...
+        """
+        self.capacity = capacity
+        self.replay_ratio = replay_ratio
+        self.replay_proportion = None
+        if self.replay_ratio != 1.0:
+            self.replay_proportion = self.replay_ratio / (1.0 - self.replay_ratio)
+
+        if replay_mode in ["lockstep", ReplayMode.LOCKSTEP]:
+            self.replay_mode = ReplayMode.LOCKSTEP
+        elif replay_mode in ["independent", ReplayMode.INDEPENDENT]:
+            self.replay_mode = ReplayMode.INDEPENDENT
+        else:
+            raise ValueError("Unsupported replay mode: {}".format(replay_mode))
+
+        def new_buffer():
+            return SimpleReplayBuffer(num_slots=capacity)
+
+        self.replay_buffers = collections.defaultdict(new_buffer)
+
+        # Metrics.
+        self.add_batch_timer = _Timer()
+        self.replay_timer = _Timer()
+        self.update_priorities_timer = _Timer()
+
+        # Added timesteps over lifetime.
+        self.num_added = 0
+
+        # Last added batch(es).
+        self.last_added_batches = collections.defaultdict(list)
+
+    def add(self, batch: SampleBatchType) -> None:
+        """Adds a batch to the appropriate policy's replay buffer.
+
+        Turns the batch into a MultiAgentBatch of the DEFAULT_POLICY_ID if
+        it is not a MultiAgentBatch. Subsequently adds the individual policy
+        batches to the storage.
+
+        Args:
+            batch: The batch to be added.
+        """
+        # Make a copy so the replay buffer doesn't pin plasma memory.
+        batch = batch.copy()
+        batch = batch.as_multi_agent()
+
+        with self.add_batch_timer:
+            if self.replay_mode == ReplayMode.LOCKSTEP:
+                # Lockstep mode: Store under _ALL_POLICIES key (we will always
+                # only sample from all policies at the same time).
+                # This means storing a MultiAgentBatch to the underlying buffer
+                self.replay_buffers[_ALL_POLICIES].add_batch(batch)
+                self.last_added_batches[_ALL_POLICIES].append(batch)
+            else:
+                # Store independent SampleBatches
+                for policy_id, sample_batch in batch.policy_batches.items():
+                    self.replay_buffers[policy_id].add_batch(sample_batch)
+                    self.last_added_batches[policy_id].append(sample_batch)
+
+        self.num_added += batch.count
+
+    def replay(
+        self, policy_id: PolicyID = DEFAULT_POLICY_ID
+    ) -> Optional[SampleBatchType]:
+        if self.replay_mode == ReplayMode.LOCKSTEP and policy_id != _ALL_POLICIES:
+            raise ValueError(
+                "Trying to sample from single policy's buffer in lockstep "
+                "mode. In lockstep mode, all policies' experiences are "
+                "sampled from a single replay buffer which is accessed "
+                "with the policy id `{}`".format(_ALL_POLICIES)
+            )
+
+        buffer = self.replay_buffers[policy_id]
+        # Return None, if:
+        # - Buffer empty or
+        # - `replay_ratio` < 1.0 (new samples required in returned batch)
+        #   and no new samples to mix with replayed ones.
+        if len(buffer) == 0 or (
+            len(self.last_added_batches[policy_id]) == 0 and self.replay_ratio < 1.0
+        ):
+            return None
+
+        # Mix buffer's last added batches with older replayed batches.
+        with self.replay_timer:
+            output_batches = self.last_added_batches[policy_id]
+            self.last_added_batches[policy_id] = []
+
+            # No replay desired -> Return here.
+            if self.replay_ratio == 0.0:
+                return concat_samples(output_batches)
+            # Only replay desired -> Return a (replayed) sample from the
+            # buffer.
+            elif self.replay_ratio == 1.0:
+                return buffer.replay()
+
+            # Replay ratio = old / [old + new]
+            # Replay proportion: old / new
+            num_new = len(output_batches)
+            replay_proportion = self.replay_proportion
+            while random.random() < num_new * replay_proportion:
+                replay_proportion -= 1
+                output_batches.append(buffer.replay())
+            return concat_samples(output_batches)
+
+    def get_host(self) -> str:
+        """Returns the computer's network name.
+
+        Returns:
+            The computer's networks name or an empty string, if the network
+            name could not be determined.
+        """
+        return platform.node()

rllib/execution/learner_thread.py

@@ -0,0 +1,137 @@
+import copy
+import queue
+import threading
+from typing import Dict, Optional
+
+from ray.rllib.evaluation.rollout_worker import RolloutWorker
+from ray.rllib.execution.minibatch_buffer import MinibatchBuffer
+from ray.rllib.utils.annotations import OldAPIStack
+from ray.rllib.utils.framework import try_import_tf
+from ray.rllib.utils.metrics.learner_info import LEARNER_INFO, LearnerInfoBuilder
+from ray.rllib.utils.metrics.window_stat import WindowStat
+from ray.util.iter import _NextValueNotReady
+from ray.util.timer import _Timer
+
+tf1, tf, tfv = try_import_tf()
+
+
+@OldAPIStack
+class LearnerThread(threading.Thread):
+    """Background thread that updates the local model from sample trajectories.
+
+    The learner thread communicates with the main thread through Queues. This
+    is needed since Ray operations can only be run on the main thread. In
+    addition, moving heavyweight gradient ops session runs off the main thread
+    improves overall throughput.
+    """
+
+    def __init__(
+        self,
+        local_worker: RolloutWorker,
+        minibatch_buffer_size: int,
+        num_sgd_iter: int,
+        learner_queue_size: int,
+        learner_queue_timeout: int,
+    ):
+        """Initialize the learner thread.
+
+        Args:
+            local_worker: process local rollout worker holding
+                policies this thread will call learn_on_batch() on
+            minibatch_buffer_size: max number of train batches to store
+                in the minibatching buffer
+            num_sgd_iter: number of passes to learn on per train batch
+            learner_queue_size: max size of queue of inbound
+                train batches to this thread
+            learner_queue_timeout: raise an exception if the queue has
+                been empty for this long in seconds
+        """
+        threading.Thread.__init__(self)
+        self.learner_queue_size = WindowStat("size", 50)
+        self.local_worker = local_worker
+        self.inqueue = queue.Queue(maxsize=learner_queue_size)
+        self.outqueue = queue.Queue()
+        self.minibatch_buffer = MinibatchBuffer(
+            inqueue=self.inqueue,
+            size=minibatch_buffer_size,
+            timeout=learner_queue_timeout,
+            num_passes=num_sgd_iter,
+            init_num_passes=num_sgd_iter,
+        )
+        self.queue_timer = _Timer()
+        self.grad_timer = _Timer()
+        self.load_timer = _Timer()
+        self.load_wait_timer = _Timer()
+        self.daemon = True
+        self.policy_ids_updated = []
+        self.learner_info = {}
+        self.stopped = False
+        self.num_steps = 0
+
+    def run(self) -> None:
+        # Switch on eager mode if configured.
+        if self.local_worker.config.framework_str == "tf2":
+            tf1.enable_eager_execution()
+        while not self.stopped:
+            self.step()
+
+    def step(self) -> Optional[_NextValueNotReady]:
+        with self.queue_timer:
+            try:
+                batch, _ = self.minibatch_buffer.get()
+            except queue.Empty:
+                return _NextValueNotReady()
+        with self.grad_timer:
+            # Use LearnerInfoBuilder as a unified way to build the final
+            # results dict from `learn_on_loaded_batch` call(s).
+            # This makes sure results dicts always have the same structure
+            # no matter the setup (multi-GPU, multi-agent, minibatch SGD,
+            # tf vs torch).
+            learner_info_builder = LearnerInfoBuilder(num_devices=1)
+            if self.local_worker.config.policy_states_are_swappable:
+                self.local_worker.lock()
+            multi_agent_results = self.local_worker.learn_on_batch(batch)
+            if self.local_worker.config.policy_states_are_swappable:
+                self.local_worker.unlock()
+            self.policy_ids_updated.extend(list(multi_agent_results.keys()))
+            for pid, results in multi_agent_results.items():
+                learner_info_builder.add_learn_on_batch_results(results, pid)
+            self.learner_info = learner_info_builder.finalize()
+
+        self.num_steps += 1
+        # Put tuple: env-steps, agent-steps, and learner info into the queue.
+        self.outqueue.put((batch.count, batch.agent_steps(), self.learner_info))
+        self.learner_queue_size.push(self.inqueue.qsize())
+
+    def add_learner_metrics(self, result: Dict, overwrite_learner_info=True) -> Dict:
+        """Add internal metrics to a result dict."""
+
+        def timer_to_ms(timer):
+            return round(1000 * timer.mean, 3)
+
+        if overwrite_learner_info:
+            result["info"].update(
+                {
+                    "learner_queue": self.learner_queue_size.stats(),
+                    LEARNER_INFO: copy.deepcopy(self.learner_info),
+                    "timing_breakdown": {
+                        "learner_grad_time_ms": timer_to_ms(self.grad_timer),
+                        "learner_load_time_ms": timer_to_ms(self.load_timer),
+                        "learner_load_wait_time_ms": timer_to_ms(self.load_wait_timer),
+                        "learner_dequeue_time_ms": timer_to_ms(self.queue_timer),
+                    },
+                }
+            )
+        else:
+            result["info"].update(
+                {
+                    "learner_queue": self.learner_queue_size.stats(),
+                    "timing_breakdown": {
+                        "learner_grad_time_ms": timer_to_ms(self.grad_timer),
+                        "learner_load_time_ms": timer_to_ms(self.load_timer),
+                        "learner_load_wait_time_ms": timer_to_ms(self.load_wait_timer),
+                        "learner_dequeue_time_ms": timer_to_ms(self.queue_timer),
+                    },
+                }
+            )
+        return result

rllib/execution/minibatch_buffer.py

@@ -0,0 +1,61 @@
+import queue
+from typing import Any, Tuple
+
+from ray.rllib.utils.annotations import OldAPIStack
+
+
+@OldAPIStack
+class MinibatchBuffer:
+    """Ring buffer of recent data batches for minibatch SGD.
+
+    This is for use with AsyncSamplesOptimizer.
+    """
+
+    def __init__(
+        self,
+        inqueue: queue.Queue,
+        size: int,
+        timeout: float,
+        num_passes: int,
+        init_num_passes: int = 1,
+    ):
+        """Initialize a minibatch buffer.
+
+        Args:
+           inqueue (queue.Queue): Queue to populate the internal ring buffer
+              from.
+           size: Max number of data items to buffer.
+           timeout: Queue timeout
+           num_passes: Max num times each data item should be emitted.
+           init_num_passes: Initial passes for each data item.
+              Maxiumum number of passes per item are increased to num_passes over
+              time.
+        """
+        self.inqueue = inqueue
+        self.size = size
+        self.timeout = timeout
+        self.max_initial_ttl = num_passes
+        self.cur_initial_ttl = init_num_passes
+        self.buffers = [None] * size
+        self.ttl = [0] * size
+        self.idx = 0
+
+    def get(self) -> Tuple[Any, bool]:
+        """Get a new batch from the internal ring buffer.
+
+        Returns:
+           buf: Data item saved from inqueue.
+           released: True if the item is now removed from the ring buffer.
+        """
+        if self.ttl[self.idx] <= 0:
+            self.buffers[self.idx] = self.inqueue.get(timeout=self.timeout)
+            self.ttl[self.idx] = self.cur_initial_ttl
+            if self.cur_initial_ttl < self.max_initial_ttl:
+                self.cur_initial_ttl += 1
+        buf = self.buffers[self.idx]
+        self.ttl[self.idx] -= 1
+        released = self.ttl[self.idx] <= 0
+        if released:
+            self.buffers[self.idx] = None
+        self.idx = (self.idx + 1) % len(self.buffers)
+        return buf, released

rllib/execution/multi_gpu_learner_thread.py

@@ -0,0 +1,245 @@
+import logging
+import queue
+import threading
+
+from ray._common.deprecation import deprecation_warning
+from ray.rllib.evaluation.rollout_worker import RolloutWorker
+from ray.rllib.execution.learner_thread import LearnerThread
+from ray.rllib.execution.minibatch_buffer import MinibatchBuffer
+from ray.rllib.policy.sample_batch import SampleBatch
+from ray.rllib.utils.annotations import OldAPIStack, override
+from ray.rllib.utils.framework import try_import_tf
+from ray.rllib.utils.metrics.learner_info import LearnerInfoBuilder
+from ray.util.timer import _Timer
+
+tf1, tf, tfv = try_import_tf()
+
+logger = logging.getLogger(__name__)
+
+
+@OldAPIStack
+class MultiGPULearnerThread(LearnerThread):
+    """Learner that can use multiple GPUs and parallel loading.
+
+    This class is used for async sampling algorithms.
+
+    Example workflow: 2 GPUs and 3 multi-GPU tower stacks.
+    -> On each GPU, there are 3 slots for batches, indexed 0, 1, and 2.
+
+    Workers collect data from env and push it into inqueue:
+    Workers -> (data) -> self.inqueue
+
+    We also have two queues, indicating, which stacks are loaded and which
+    are not.
+    - idle_tower_stacks = [0, 1, 2]  <- all 3 stacks are free at first.
+    - ready_tower_stacks = []  <- None of the 3 stacks is loaded with data.
+
+    `ready_tower_stacks` is managed by `ready_tower_stacks_buffer` for
+    possible minibatch-SGD iterations per loaded batch (this avoids a reload
+    from CPU to GPU for each SGD iter).
+
+    n _MultiGPULoaderThreads: self.inqueue -get()->
+    policy.load_batch_into_buffer() -> ready_stacks = [0 ...]
+
+    This thread: self.ready_tower_stacks_buffer -get()->
+    policy.learn_on_loaded_batch() -> if SGD-iters done,
+    put stack index back in idle_tower_stacks queue.
+    """
+
+    def __init__(
+        self,
+        local_worker: RolloutWorker,
+        num_gpus: int = 1,
+        lr=None,  # deprecated.
+        train_batch_size: int = 500,
+        num_multi_gpu_tower_stacks: int = 1,
+        num_sgd_iter: int = 1,
+        learner_queue_size: int = 16,
+        learner_queue_timeout: int = 300,
+        num_data_load_threads: int = 16,
+        _fake_gpus: bool = False,
+        # Deprecated arg, use
+        minibatch_buffer_size=None,
+    ):
+        """Initializes a MultiGPULearnerThread instance.
+
+        Args:
+            local_worker: Local RolloutWorker holding
+                policies this thread will call `load_batch_into_buffer` and
+                `learn_on_loaded_batch` on.
+            num_gpus: Number of GPUs to use for data-parallel SGD.
+            train_batch_size: Size of batches (minibatches if
+                `num_sgd_iter` > 1) to learn on.
+            num_multi_gpu_tower_stacks: Number of buffers to parallelly
+                load data into on one device. Each buffer is of size of
+                `train_batch_size` and hence increases GPU memory usage
+                accordingly.
+            num_sgd_iter: Number of passes to learn on per train batch
+                (minibatch if `num_sgd_iter` > 1).
+            learner_queue_size: Max size of queue of inbound
+                train batches to this thread.
+            num_data_load_threads: Number of threads to use to load
+                data into GPU memory in parallel.
+        """
+        # Deprecated: No need to specify as we don't need the actual
+        # minibatch-buffer anyways.
+        if minibatch_buffer_size:
+            deprecation_warning(
+                old="MultiGPULearnerThread.minibatch_buffer_size",
+                error=True,
+            )
+        super().__init__(
+            local_worker=local_worker,
+            minibatch_buffer_size=0,
+            num_sgd_iter=num_sgd_iter,
+            learner_queue_size=learner_queue_size,
+            learner_queue_timeout=learner_queue_timeout,
+        )
+        # Delete reference to parent's minibatch_buffer, which is not needed.
+        # Instead, in multi-GPU mode, we pull tower stack indices from the
+        # `self.ready_tower_stacks_buffer` buffer, whose size is exactly
+        # `num_multi_gpu_tower_stacks`.
+        self.minibatch_buffer = None
+
+        self.train_batch_size = train_batch_size
+
+        self.policy_map = self.local_worker.policy_map
+        self.devices = next(iter(self.policy_map.values())).devices
+
+        logger.info("MultiGPULearnerThread devices {}".format(self.devices))
+        assert self.train_batch_size % len(self.devices) == 0
+        assert self.train_batch_size >= len(self.devices), "batch too small"
+
+        self.tower_stack_indices = list(range(num_multi_gpu_tower_stacks))
+
+        # Two queues for tower stacks:
+        # a) Those that are loaded with data ("ready")
+        # b) Those that are ready to be loaded with new data ("idle").
+        self.idle_tower_stacks = queue.Queue()
+        self.ready_tower_stacks = queue.Queue()
+        # In the beginning, all stacks are idle (no loading has taken place
+        # yet).
+        for idx in self.tower_stack_indices:
+            self.idle_tower_stacks.put(idx)
+        # Start n threads that are responsible for loading data into the
+        # different (idle) stacks.
+        for i in range(num_data_load_threads):
+            self.loader_thread = _MultiGPULoaderThread(self, share_stats=(i == 0))
+            self.loader_thread.start()
+
+        # Create a buffer that holds stack indices that are "ready"
+        # (loaded with data). Those are stacks that we can call
+        # "learn_on_loaded_batch" on.
+        self.ready_tower_stacks_buffer = MinibatchBuffer(
+            self.ready_tower_stacks,
+            num_multi_gpu_tower_stacks,
+            learner_queue_timeout,
+            num_sgd_iter,
+        )
+
+    @override(LearnerThread)
+    def step(self) -> None:
+        if not self.loader_thread.is_alive():
+            raise RuntimeError(
+                "The `_MultiGPULoaderThread` has died! Will therefore also terminate "
+                "the `MultiGPULearnerThread`."
+            )
+
+        with self.load_wait_timer:
+            buffer_idx, released = self.ready_tower_stacks_buffer.get()
+
+        get_num_samples_loaded_into_buffer = 0
+        with self.grad_timer:
+            # Use LearnerInfoBuilder as a unified way to build the final
+            # results dict from `learn_on_loaded_batch` call(s).
+            # This makes sure results dicts always have the same structure
+            # no matter the setup (multi-GPU, multi-agent, minibatch SGD,
+            # tf vs torch).
+            learner_info_builder = LearnerInfoBuilder(num_devices=len(self.devices))
+
+            for pid in self.policy_map.keys():
+                # Not a policy-to-train.
+                if (
+                    self.local_worker.is_policy_to_train is not None
+                    and not self.local_worker.is_policy_to_train(pid)
+                ):
+                    continue
+                policy = self.policy_map[pid]
+                default_policy_results = policy.learn_on_loaded_batch(
+                    offset=0, buffer_index=buffer_idx
+                )
+                learner_info_builder.add_learn_on_batch_results(
+                    default_policy_results, policy_id=pid
+                )
+                self.policy_ids_updated.append(pid)
+                get_num_samples_loaded_into_buffer += (
+                    policy.get_num_samples_loaded_into_buffer(buffer_idx)
+                )
+
+            self.learner_info = learner_info_builder.finalize()
+
+        if released:
+            self.idle_tower_stacks.put(buffer_idx)
+
+        # Put tuple: env-steps, agent-steps, and learner info into the queue.
+        self.outqueue.put(
+            (
+                get_num_samples_loaded_into_buffer,
+                get_num_samples_loaded_into_buffer,
+                self.learner_info,
+            )
+        )
+        self.learner_queue_size.push(self.inqueue.qsize())
+
+
+class _MultiGPULoaderThread(threading.Thread):
+    def __init__(
+        self, multi_gpu_learner_thread: MultiGPULearnerThread, share_stats: bool
+    ):
+        threading.Thread.__init__(self)
+        self.multi_gpu_learner_thread = multi_gpu_learner_thread
+        self.daemon = True
+        if share_stats:
+            self.queue_timer = multi_gpu_learner_thread.queue_timer
+            self.load_timer = multi_gpu_learner_thread.load_timer
+        else:
+            self.queue_timer = _Timer()
+            self.load_timer = _Timer()
+
+    def run(self) -> None:
+        while True:
+            self._step()
+
+    def _step(self) -> None:
+        s = self.multi_gpu_learner_thread
+        policy_map = s.policy_map
+
+        # Get a new batch from the data (inqueue).
+        with self.queue_timer:
+            batch = s.inqueue.get()
+
+        # Get next idle stack for loading.
+        buffer_idx = s.idle_tower_stacks.get()
+
+        # Load the batch into the idle stack.
+        with self.load_timer:
+            for pid in policy_map.keys():
+                if (
+                    s.local_worker.is_policy_to_train is not None
+                    and not s.local_worker.is_policy_to_train(pid, batch)
+                ):
+                    continue
+                policy = policy_map[pid]
+                if isinstance(batch, SampleBatch):
+                    policy.load_batch_into_buffer(
+                        batch=batch,
+                        buffer_index=buffer_idx,
+                    )
+                elif pid in batch.policy_batches:
+                    policy.load_batch_into_buffer(
+                        batch=batch.policy_batches[pid],
+                        buffer_index=buffer_idx,
+                    )
+
+        # Tag just-loaded stack as "ready".
+        s.ready_tower_stacks.put(buffer_idx)

rllib/execution/replay_ops.py

@@ -0,0 +1,37 @@
+import random
+from typing import Optional
+
+from ray.rllib.utils.annotations import OldAPIStack
+from ray.rllib.utils.replay_buffers.replay_buffer import warn_replay_capacity
+from ray.rllib.utils.typing import SampleBatchType
+
+
+@OldAPIStack
+class SimpleReplayBuffer:
+    """Simple replay buffer that operates over batches."""
+
+    def __init__(self, num_slots: int, replay_proportion: Optional[float] = None):
+        """Initialize SimpleReplayBuffer.
+
+        Args:
+            num_slots: Number of batches to store in total.
+        """
+        self.num_slots = num_slots
+        self.replay_batches = []
+        self.replay_index = 0
+
+    def add_batch(self, sample_batch: SampleBatchType) -> None:
+        warn_replay_capacity(item=sample_batch, num_items=self.num_slots)
+        if self.num_slots > 0:
+            if len(self.replay_batches) < self.num_slots:
+                self.replay_batches.append(sample_batch)
+            else:
+                self.replay_batches[self.replay_index] = sample_batch
+                self.replay_index += 1
+                self.replay_index %= self.num_slots
+
+    def replay(self) -> SampleBatchType:
+        return random.choice(self.replay_batches)
+
+    def __len__(self):
+        return len(self.replay_batches)

rllib/execution/rollout_ops.py

@@ -0,0 +1,209 @@
+import logging
+from typing import List, Optional, Union
+
+import tree
+
+from ray.rllib.env.env_runner_group import EnvRunnerGroup
+from ray.rllib.policy.sample_batch import (
+    DEFAULT_POLICY_ID,
+    SampleBatch,
+    concat_samples,
+)
+from ray.rllib.utils.annotations import ExperimentalAPI, OldAPIStack
+from ray.rllib.utils.metrics import NUM_AGENT_STEPS_SAMPLED, NUM_ENV_STEPS_SAMPLED
+from ray.rllib.utils.sgd import standardized
+from ray.rllib.utils.typing import EpisodeType, SampleBatchType
+
+logger = logging.getLogger(__name__)
+
+
+@ExperimentalAPI
+def synchronous_parallel_sample(
+    *,
+    worker_set: EnvRunnerGroup,
+    max_agent_steps: Optional[int] = None,
+    max_env_steps: Optional[int] = None,
+    concat: bool = True,
+    sample_timeout_s: Optional[float] = None,
+    random_actions: bool = False,
+    _uses_new_env_runners: bool = False,
+    _return_metrics: bool = False,
+) -> Union[List[SampleBatchType], SampleBatchType, List[EpisodeType], EpisodeType]:
+    """Runs parallel and synchronous rollouts on all remote workers.
+
+    Waits for all workers to return from the remote calls.
+
+    If no remote workers exist (num_workers == 0), use the local worker
+    for sampling.
+
+    Alternatively to calling `worker.sample.remote()`, the user can provide a
+    `remote_fn()`, which will be applied to the worker(s) instead.
+
+    Args:
+        worker_set: The EnvRunnerGroup to use for sampling.
+        remote_fn: If provided, use `worker.apply.remote(remote_fn)` instead
+            of `worker.sample.remote()` to generate the requests.
+        max_agent_steps: Optional number of agent steps to be included in the
+            final batch or list of episodes.
+        max_env_steps: Optional number of environment steps to be included in the
+            final batch or list of episodes.
+        concat: Whether to aggregate all resulting batches or episodes. in case of
+            batches the list of batches is concatinated at the end. in case of
+            episodes all episode lists from workers are flattened into a single list.
+        sample_timeout_s: The timeout in sec to use on the `foreach_env_runner` call.
+            After this time, the call will return with a result (or not if all
+            EnvRunners are stalling). If None, will block indefinitely and not timeout.
+        _uses_new_env_runners: Whether the new `EnvRunner API` is used. In this case
+            episodes instead of `SampleBatch` objects are returned.
+
+    Returns:
+        The list of collected sample batch types or episode types (one for each parallel
+        rollout worker in the given `worker_set`).
+
+    .. testcode::
+
+        # Define an RLlib Algorithm.
+        from ray.rllib.algorithms.ppo import PPO, PPOConfig
+        config = (
+            PPOConfig()
+            .environment("CartPole-v1")
+        )
+        algorithm = config.build()
+        # 2 remote EnvRunners (num_env_runners=2):
+        episodes = synchronous_parallel_sample(
+            worker_set=algorithm.env_runner_group,
+            _uses_new_env_runners=True,
+            concat=False,
+        )
+        print(len(episodes))
+
+    .. testoutput::
+
+        2
+    """
+    # Only allow one of `max_agent_steps` or `max_env_steps` to be defined.
+    assert not (max_agent_steps is not None and max_env_steps is not None)
+
+    agent_or_env_steps = 0
+    max_agent_or_env_steps = max_agent_steps or max_env_steps or None
+    sample_batches_or_episodes = []
+    all_stats_dicts = []
+
+    random_action_kwargs = {} if not random_actions else {"random_actions": True}
+
+    # Stop collecting batches as soon as one criterium is met.
+    while (max_agent_or_env_steps is None and agent_or_env_steps == 0) or (
+        max_agent_or_env_steps is not None
+        and agent_or_env_steps < max_agent_or_env_steps
+    ):
+        # No remote workers in the set -> Use local worker for collecting
+        # samples.
+        if worker_set.num_remote_workers() <= 0:
+            sampled_data = [worker_set.local_env_runner.sample(**random_action_kwargs)]
+            if _return_metrics:
+                stats_dicts = [worker_set.local_env_runner.get_metrics()]
+        # Loop over remote workers' `sample()` method in parallel.
+        else:
+            sampled_data = worker_set.foreach_env_runner(
+                (
+                    (lambda w: w.sample(**random_action_kwargs))
+                    if not _return_metrics
+                    else (lambda w: (w.sample(**random_action_kwargs), w.get_metrics()))
+                ),
+                local_env_runner=False,
+                timeout_seconds=sample_timeout_s,
+            )
+            # Nothing was returned (maybe all workers are stalling) or no healthy
+            # remote workers left: Break.
+            # There is no point staying in this loop, since we will not be able to
+            # get any new samples if we don't have any healthy remote workers left.
+            if not sampled_data or worker_set.num_healthy_remote_workers() <= 0:
+                if not sampled_data:
+                    logger.warning(
+                        "No samples returned from remote workers. If you have a "
+                        "slow environment or model, consider increasing the "
+                        "`sample_timeout_s` or decreasing the "
+                        "`rollout_fragment_length` in `AlgorithmConfig.env_runners()."
+                    )
+                elif worker_set.num_healthy_remote_workers() <= 0:
+                    logger.warning(
+                        "No healthy remote workers left. Trying to restore workers ..."
+                    )
+                break
+
+            if _return_metrics:
+                stats_dicts = [s[1] for s in sampled_data]
+                sampled_data = [s[0] for s in sampled_data]
+
+        # Update our counters for the stopping criterion of the while loop.
+        if _return_metrics:
+            if max_agent_steps:
+                agent_or_env_steps += sum(
+                    int(agent_stat)
+                    for stat_dict in stats_dicts
+                    for agent_stat in stat_dict[NUM_AGENT_STEPS_SAMPLED].values()
+                )
+            else:
+                agent_or_env_steps += sum(
+                    int(stat_dict.get(NUM_ENV_STEPS_SAMPLED, 0))
+                    for stat_dict in stats_dicts
+                )
+            sample_batches_or_episodes.extend(sampled_data)
+            all_stats_dicts.extend(stats_dicts)
+        else:
+            for batch_or_episode in sampled_data:
+                if max_agent_steps:
+                    agent_or_env_steps += (
+                        sum(e.agent_steps() for e in batch_or_episode)
+                        if _uses_new_env_runners
+                        else batch_or_episode.agent_steps()
+                    )
+                else:
+                    agent_or_env_steps += (
+                        sum(e.env_steps() for e in batch_or_episode)
+                        if _uses_new_env_runners
+                        else batch_or_episode.env_steps()
+                    )
+                sample_batches_or_episodes.append(batch_or_episode)
+                # Break out (and ignore the remaining samples) if max timesteps (batch
+                # size) reached. We want to avoid collecting batches that are too large
+                # only because of a failed/restarted worker causing a second iteration
+                # of the main loop.
+                if (
+                    max_agent_or_env_steps is not None
+                    and agent_or_env_steps >= max_agent_or_env_steps
+                ):
+                    break
+
+    if concat is True:
+        # If we have episodes flatten the episode list.
+        if _uses_new_env_runners:
+            sample_batches_or_episodes = tree.flatten(sample_batches_or_episodes)
+        # Otherwise we concatenate the `SampleBatch` objects
+        else:
+            sample_batches_or_episodes = concat_samples(sample_batches_or_episodes)
+
+    if _return_metrics:
+        return sample_batches_or_episodes, all_stats_dicts
+    return sample_batches_or_episodes
+
+
+@OldAPIStack
+def standardize_fields(samples: SampleBatchType, fields: List[str]) -> SampleBatchType:
+    """Standardize fields of the given SampleBatch"""
+    wrapped = False
+
+    if isinstance(samples, SampleBatch):
+        samples = samples.as_multi_agent()
+        wrapped = True
+
+    for policy_id in samples.policy_batches:
+        batch = samples.policy_batches[policy_id]
+        for field in fields:
+            if field in batch:
+                batch[field] = standardized(batch[field])
+
+    if wrapped:
+        samples = samples.policy_batches[DEFAULT_POLICY_ID]
+
+    return samples

rllib/execution/segment_tree.py

@@ -0,0 +1,217 @@
+import operator
+from typing import Any, Optional
+
+
+class SegmentTree:
+    """A Segment Tree data structure.
+
+    https://en.wikipedia.org/wiki/Segment_tree
+
+    Can be used as regular array, but with two important differences:
+
+      a) Setting an item's value is slightly slower. It is O(lg capacity),
+         instead of O(1).
+      b) Offers efficient `reduce` operation which reduces the tree's values
+         over some specified contiguous subsequence of items in the array.
+         Operation could be e.g. min/max/sum.
+
+    The data is stored in a list, where the length is 2 * capacity.
+    The second half of the list stores the actual values for each index, so if
+    capacity=8, values are stored at indices 8 to 15. The first half of the
+    array contains the reduced-values of the different (binary divided)
+    segments, e.g. (capacity=4):
+    0=not used
+    1=reduced-value over all elements (array indices 4 to 7).
+    2=reduced-value over array indices (4 and 5).
+    3=reduced-value over array indices (6 and 7).
+    4-7: values of the tree.
+    NOTE that the values of the tree are accessed by indices starting at 0, so
+    `tree[0]` accesses `internal_array[4]` in the above example.
+    """
+
+    def __init__(
+        self, capacity: int, operation: Any, neutral_element: Optional[Any] = None
+    ):
+        """Initializes a Segment Tree object.
+
+        Args:
+            capacity: Total size of the array - must be a power of two.
+            operation: Lambda obj, obj -> obj
+                The operation for combining elements (eg. sum, max).
+                Must be a mathematical group together with the set of
+                possible values for array elements.
+            neutral_element (Optional[obj]): The neutral element for
+                `operation`. Use None for automatically finding a value:
+                max: float("-inf"), min: float("inf"), sum: 0.0.
+        """
+
+        assert (
+            capacity > 0 and capacity & (capacity - 1) == 0
+        ), "Capacity must be positive and a power of 2!"
+        self.capacity = capacity
+        if neutral_element is None:
+            neutral_element = (
+                0.0
+                if operation is operator.add
+                else float("-inf")
+                if operation is max
+                else float("inf")
+            )
+        self.neutral_element = neutral_element
+        self.value = [self.neutral_element for _ in range(2 * capacity)]
+        self.operation = operation
+
+    def reduce(self, start: int = 0, end: Optional[int] = None) -> Any:
+        """Applies `self.operation` to subsequence of our values.
+
+        Subsequence is contiguous, includes `start` and excludes `end`.
+
+          self.operation(
+              arr[start], operation(arr[start+1], operation(... arr[end])))
+
+        Args:
+            start: Start index to apply reduction to.
+            end (Optional[int]): End index to apply reduction to (excluded).
+
+        Returns:
+            any: The result of reducing self.operation over the specified
+                range of `self._value` elements.
+        """
+        if end is None:
+            end = self.capacity
+        elif end < 0:
+            end += self.capacity
+
+        # Init result with neutral element.
+        result = self.neutral_element
+        # Map start/end to our actual index space (second half of array).
+        start += self.capacity
+        end += self.capacity
+
+        # Example:
+        # internal-array (first half=sums, second half=actual values):
+        # 0 1 2 3 | 4 5 6 7
+        # - 6 1 5 | 1 0 2 3
+
+        # tree.sum(0, 3) = 3
+        # internally: start=4, end=7 -> sum values 1 0 2 = 3.
+
+        # Iterate over tree starting in the actual-values (second half)
+        # section.
+        # 1) start=4 is even -> do nothing.
+        # 2) end=7 is odd -> end-- -> end=6 -> add value to result: result=2
+        # 3) int-divide start and end by 2: start=2, end=3
+        # 4) start still smaller end -> iterate once more.
+        # 5) start=2 is even -> do nothing.
+        # 6) end=3 is odd -> end-- -> end=2 -> add value to result: result=1
+        #    NOTE: This adds the sum of indices 4 and 5 to the result.
+
+        # Iterate as long as start != end.
+        while start < end:
+
+            # If start is odd: Add its value to result and move start to
+            # next even value.
+            if start & 1:
+                result = self.operation(result, self.value[start])
+                start += 1
+
+            # If end is odd: Move end to previous even value, then add its
+            # value to result. NOTE: This takes care of excluding `end` in any
+            # situation.
+            if end & 1:
+                end -= 1
+                result = self.operation(result, self.value[end])
+
+            # Divide both start and end by 2 to make them "jump" into the
+            # next upper level reduce-index space.
+            start //= 2
+            end //= 2
+
+            # Then repeat till start == end.
+
+        return result
+
+    def __setitem__(self, idx: int, val: float) -> None:
+        """
+        Inserts/overwrites a value in/into the tree.
+
+        Args:
+            idx: The index to insert to. Must be in [0, `self.capacity`)
+            val: The value to insert.
+        """
+        assert 0 <= idx < self.capacity, f"idx={idx} capacity={self.capacity}"
+
+        # Index of the leaf to insert into (always insert in "second half"
+        # of the tree, the first half is reserved for already calculated
+        # reduction-values).
+        idx += self.capacity
+        self.value[idx] = val
+
+        # Recalculate all affected reduction values (in "first half" of tree).
+        idx = idx >> 1  # Divide by 2 (faster than division).
+        while idx >= 1:
+            update_idx = 2 * idx  # calculate only once
+            # Update the reduction value at the correct "first half" idx.
+            self.value[idx] = self.operation(
+                self.value[update_idx], self.value[update_idx + 1]
+            )
+            idx = idx >> 1  # Divide by 2 (faster than division).
+
+    def __getitem__(self, idx: int) -> Any:
+        assert 0 <= idx < self.capacity
+        return self.value[idx + self.capacity]
+
+    def get_state(self):
+        return self.value
+
+    def set_state(self, state):
+        assert len(state) == self.capacity * 2
+        self.value = state
+
+
+class SumSegmentTree(SegmentTree):
+    """A SegmentTree with the reduction `operation`=operator.add."""
+
+    def __init__(self, capacity: int):
+        super(SumSegmentTree, self).__init__(capacity=capacity, operation=operator.add)
+
+    def sum(self, start: int = 0, end: Optional[Any] = None) -> Any:
+        """Returns the sum over a sub-segment of the tree."""
+        return self.reduce(start, end)
+
+    def find_prefixsum_idx(self, prefixsum: float) -> int:
+        """Finds highest i, for which: sum(arr[0]+..+arr[i - i]) <= prefixsum.
+
+        Args:
+            prefixsum: `prefixsum` upper bound in above constraint.
+
+        Returns:
+            int: Largest possible index (i) satisfying above constraint.
+        """
+        assert 0 <= prefixsum <= self.sum() + 1e-5
+        # Global sum node.
+        idx = 1
+
+        # Edge case when prefixsum can clip into the invalid regions
+        #   https://github.com/ray-project/ray/issues/54284
+        if prefixsum >= self.value[idx]:
+            prefixsum = self.value[idx] - 1e-5
+
+        # While non-leaf (first half of tree).
+        while idx < self.capacity:
+            update_idx = 2 * idx
+            if self.value[update_idx] > prefixsum:
+                idx = update_idx
+            else:
+                prefixsum -= self.value[update_idx]
+                idx = update_idx + 1
+        return idx - self.capacity
+
+
+class MinSegmentTree(SegmentTree):
+    def __init__(self, capacity: int):
+        super(MinSegmentTree, self).__init__(capacity=capacity, operation=min)
+
+    def min(self, start: int = 0, end: Optional[Any] = None) -> Any:
+        """Returns min(arr[start], ...,  arr[end])"""
+        return self.reduce(start, end)

rllib/execution/train_ops.py

@@ -0,0 +1,205 @@
+import logging
+import math
+from typing import Dict
+
+import numpy as np
+
+from ray._common.deprecation import deprecation_warning
+from ray.rllib.policy.sample_batch import DEFAULT_POLICY_ID
+from ray.rllib.utils.annotations import OldAPIStack
+from ray.rllib.utils.framework import try_import_tf
+from ray.rllib.utils.metrics import (
+    LEARN_ON_BATCH_TIMER,
+    LOAD_BATCH_TIMER,
+    NUM_AGENT_STEPS_TRAINED,
+    NUM_ENV_STEPS_TRAINED,
+)
+from ray.rllib.utils.metrics.learner_info import LearnerInfoBuilder
+from ray.rllib.utils.sgd import do_minibatch_sgd
+from ray.util import log_once
+
+tf1, tf, tfv = try_import_tf()
+
+logger = logging.getLogger(__name__)
+
+
+@OldAPIStack
+def train_one_step(algorithm, train_batch, policies_to_train=None) -> Dict:
+    """Function that improves the all policies in `train_batch` on the local worker.
+
+    .. testcode::
+        :skipif: True
+
+        from ray.rllib.execution.rollout_ops import synchronous_parallel_sample
+        algo = [...]
+        train_batch = synchronous_parallel_sample(algo.env_runner_group)
+        # This trains the policy on one batch.
+        print(train_one_step(algo, train_batch)))
+
+    .. testoutput::
+
+        {"default_policy": ...}
+
+    Updates the NUM_ENV_STEPS_TRAINED and NUM_AGENT_STEPS_TRAINED counters as well as
+    the LEARN_ON_BATCH_TIMER timer of the `algorithm` object.
+    """
+    config = algorithm.config
+    workers = algorithm.env_runner_group
+    local_worker = workers.local_env_runner
+    num_sgd_iter = config.get("num_epochs", config.get("num_sgd_iter", 1))
+    minibatch_size = config.get("minibatch_size")
+    if minibatch_size is None:
+        minibatch_size = config.get("sgd_minibatch_size", 0)
+
+    learn_timer = algorithm._timers[LEARN_ON_BATCH_TIMER]
+    with learn_timer:
+        # Subsample minibatches (size=`minibatch_size`) from the
+        # train batch and loop through train batch `num_sgd_iter` times.
+        if num_sgd_iter > 1 or minibatch_size > 0:
+            info = do_minibatch_sgd(
+                train_batch,
+                {
+                    pid: local_worker.get_policy(pid)
+                    for pid in policies_to_train
+                    or local_worker.get_policies_to_train(train_batch)
+                },
+                local_worker,
+                num_sgd_iter,
+                minibatch_size,
+                [],
+            )
+        # Single update step using train batch.
+        else:
+            info = local_worker.learn_on_batch(train_batch)
+
+    learn_timer.push_units_processed(train_batch.count)
+    algorithm._counters[NUM_ENV_STEPS_TRAINED] += train_batch.count
+    algorithm._counters[NUM_AGENT_STEPS_TRAINED] += train_batch.agent_steps()
+
+    if algorithm.reward_estimators:
+        info[DEFAULT_POLICY_ID]["off_policy_estimation"] = {}
+        for name, estimator in algorithm.reward_estimators.items():
+            info[DEFAULT_POLICY_ID]["off_policy_estimation"][name] = estimator.train(
+                train_batch
+            )
+    return info
+
+
+@OldAPIStack
+def multi_gpu_train_one_step(algorithm, train_batch) -> Dict:
+    """Multi-GPU version of train_one_step.
+
+    Uses the policies' `load_batch_into_buffer` and `learn_on_loaded_batch` methods
+    to be more efficient wrt CPU/GPU data transfers. For example, when doing multiple
+    passes through a train batch (e.g. for PPO) using `config.num_sgd_iter`, the
+    actual train batch is only split once and loaded once into the GPU(s).
+
+    .. testcode::
+        :skipif: True
+
+        from ray.rllib.execution.rollout_ops import synchronous_parallel_sample
+        algo = [...]
+        train_batch = synchronous_parallel_sample(algo.env_runner_group)
+        # This trains the policy on one batch.
+        print(multi_gpu_train_one_step(algo, train_batch)))
+
+    .. testoutput::
+
+        {"default_policy": ...}
+
+    Updates the NUM_ENV_STEPS_TRAINED and NUM_AGENT_STEPS_TRAINED counters as well as
+    the LOAD_BATCH_TIMER and LEARN_ON_BATCH_TIMER timers of the Algorithm instance.
+    """
+    if log_once("mulit_gpu_train_one_step_deprecation_warning"):
+        deprecation_warning(
+            old=("ray.rllib.execution.train_ops.multi_gpu_train_one_step")
+        )
+    config = algorithm.config
+    workers = algorithm.env_runner_group
+    local_worker = workers.local_env_runner
+    num_sgd_iter = config.get("num_epochs", config.get("num_sgd_iter", 1))
+    minibatch_size = config.get("minibatch_size")
+    if minibatch_size is None:
+        minibatch_size = config["train_batch_size"]
+
+    # Determine the number of devices (GPUs or 1 CPU) we use.
+    num_devices = int(math.ceil(config["num_gpus"] or 1))
+
+    # Make sure total batch size is dividable by the number of devices.
+    # Batch size per tower.
+    per_device_batch_size = minibatch_size // num_devices
+    # Total batch size.
+    batch_size = per_device_batch_size * num_devices
+    assert batch_size % num_devices == 0
+    assert batch_size >= num_devices, "Batch size too small!"
+
+    # Handle everything as if multi-agent.
+    train_batch = train_batch.as_multi_agent()
+
+    # Load data into GPUs.
+    load_timer = algorithm._timers[LOAD_BATCH_TIMER]
+    with load_timer:
+        num_loaded_samples = {}
+        for policy_id, batch in train_batch.policy_batches.items():
+            # Not a policy-to-train.
+            if (
+                local_worker.is_policy_to_train is not None
+                and not local_worker.is_policy_to_train(policy_id, train_batch)
+            ):
+                continue
+
+            # Decompress SampleBatch, in case some columns are compressed.
+            batch.decompress_if_needed()
+
+            # Load the entire train batch into the Policy's only buffer
+            # (idx=0). Policies only have >1 buffers, if we are training
+            # asynchronously.
+            num_loaded_samples[policy_id] = local_worker.policy_map[
+                policy_id
+            ].load_batch_into_buffer(batch, buffer_index=0)
+
+    # Execute minibatch SGD on loaded data.
+    learn_timer = algorithm._timers[LEARN_ON_BATCH_TIMER]
+    with learn_timer:
+        # Use LearnerInfoBuilder as a unified way to build the final
+        # results dict from `learn_on_loaded_batch` call(s).
+        # This makes sure results dicts always have the same structure
+        # no matter the setup (multi-GPU, multi-agent, minibatch SGD,
+        # tf vs torch).
+        learner_info_builder = LearnerInfoBuilder(num_devices=num_devices)
+
+        for policy_id, samples_per_device in num_loaded_samples.items():
+            policy = local_worker.policy_map[policy_id]
+            num_batches = max(1, int(samples_per_device) // int(per_device_batch_size))
+            logger.debug("== sgd epochs for {} ==".format(policy_id))
+            for _ in range(num_sgd_iter):
+                permutation = np.random.permutation(num_batches)
+                for batch_index in range(num_batches):
+                    # Learn on the pre-loaded data in the buffer.
+                    # Note: For minibatch SGD, the data is an offset into
+                    # the pre-loaded entire train batch.
+                    results = policy.learn_on_loaded_batch(
+                        permutation[batch_index] * per_device_batch_size, buffer_index=0
+                    )
+
+                    learner_info_builder.add_learn_on_batch_results(results, policy_id)
+
+        # Tower reduce and finalize results.
+        learner_info = learner_info_builder.finalize()
+
+    load_timer.push_units_processed(train_batch.count)
+    learn_timer.push_units_processed(train_batch.count)
+
+    # TODO: Move this into Algorithm's `training_step` method for
+    #  better transparency.
+    algorithm._counters[NUM_ENV_STEPS_TRAINED] += train_batch.count
+    algorithm._counters[NUM_AGENT_STEPS_TRAINED] += train_batch.agent_steps()
+
+    if algorithm.reward_estimators:
+        learner_info[DEFAULT_POLICY_ID]["off_policy_estimation"] = {}
+        for name, estimator in algorithm.reward_estimators.items():
+            learner_info[DEFAULT_POLICY_ID]["off_policy_estimation"][
+                name
+            ] = estimator.train(train_batch)
+
+    return learner_info