chore: import upstream snapshot with attribution

rllib/algorithms/impala/__init__.py

@@ -0,0 +1,23 @@
+from ray.rllib.algorithms.impala.impala import (
+    IMPALA,
+    Impala,
+    IMPALAConfig,
+    ImpalaConfig,
+)
+from ray.rllib.algorithms.impala.impala_tf_policy import (
+    ImpalaTF1Policy,
+    ImpalaTF2Policy,
+)
+from ray.rllib.algorithms.impala.impala_torch_policy import ImpalaTorchPolicy
+
+__all__ = [
+    "IMPALA",
+    "IMPALAConfig",
+    # @OldAPIStack
+    "ImpalaTF1Policy",
+    "ImpalaTF2Policy",
+    "ImpalaTorchPolicy",
+    # Deprecated names (lowercase)
+    "ImpalaConfig",
+    "Impala",
+]

rllib/algorithms/impala/impala_learner.py

@@ -0,0 +1,632 @@
+import atexit
+import logging
+import queue
+import threading
+import weakref
+from queue import Queue
+from typing import Any, Dict, List
+
+import ray
+from ray.rllib.algorithms.impala.impala import LEARNER_RESULTS_CURR_ENTROPY_COEFF_KEY
+from ray.rllib.core import COMPONENT_RL_MODULE
+from ray.rllib.core.learner.learner import Learner
+from ray.rllib.core.learner.training_data import TrainingData
+from ray.rllib.core.rl_module.apis import ValueFunctionAPI
+from ray.rllib.utils.annotations import (
+    OverrideToImplementCustomLogic_CallToSuperRecommended,
+    override,
+)
+from ray.rllib.utils.framework import try_import_torch
+from ray.rllib.utils.lambda_defaultdict import LambdaDefaultDict
+from ray.rllib.utils.metrics import (
+    ALL_MODULES,
+    NUM_ENV_STEPS_SAMPLED_LIFETIME,
+)
+from ray.rllib.utils.metrics.metrics_logger import MetricsLogger
+from ray.rllib.utils.metrics.ray_metrics import (
+    DEFAULT_HISTOGRAM_BOUNDARIES_SHORT_EVENTS,
+    TimerAndPrometheusLogger,
+)
+from ray.rllib.utils.schedules.scheduler import Scheduler
+from ray.rllib.utils.typing import ModuleID, ResultDict
+from ray.util.metrics import Gauge, Histogram
+
+logger = logging.getLogger(__name__)
+
+torch, _ = try_import_torch()
+
+GPU_LOADER_QUEUE_WAIT_TIMER = "gpu_loader_queue_wait_timer"
+GPU_LOADER_LOAD_TO_GPU_TIMER = "gpu_loader_load_to_gpu_timer"
+LEARNER_THREAD_IN_QUEUE_WAIT_TIMER = "learner_thread_in_queue_wait_timer"
+LEARNER_THREAD_ENV_STEPS_DROPPED = "learner_thread_env_steps_dropped"
+LEARNER_THREAD_UPDATE_TIMER = "learner_thread_update_timer"
+RAY_GET_EPISODES_TIMER = "ray_get_episodes_timer"
+
+QUEUE_SIZE_GPU_LOADER_QUEUE = "queue_size_gpu_loader_queue"
+QUEUE_SIZE_LEARNER_THREAD_QUEUE = "queue_size_learner_thread_queue"
+QUEUE_SIZE_RESULTS_QUEUE = "queue_size_results_queue"
+
+# Aggregation cycle size.
+BATCHES_PER_AGGREGATION = 10
+
+# Stop sentinel for the `_LearnerThread`
+_STOP_SENTINEL = object()
+
+
+class IMPALALearner(Learner):
+    @override(Learner)
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+
+        # Ray metrics
+        self._metrics_learner_impala_update = Histogram(
+            name="rllib_learner_impala_update_time",
+            description="Time spent in the 'IMPALALearner.update()' method.",
+            boundaries=DEFAULT_HISTOGRAM_BOUNDARIES_SHORT_EVENTS,
+            tag_keys=("rllib",),
+        )
+        self._metrics_learner_impala_update.set_default_tags(
+            {"rllib": self.__class__.__name__}
+        )
+
+        self._metrics_learner_impala_update_solve_refs = Histogram(
+            name="rllib_learner_impala_update_solve_refs_time",
+            description="Time spent on resolving refs in the 'Learner.update()'",
+            boundaries=DEFAULT_HISTOGRAM_BOUNDARIES_SHORT_EVENTS,
+            tag_keys=("rllib",),
+        )
+        self._metrics_learner_impala_update_solve_refs.set_default_tags(
+            {"rllib": self.__class__.__name__}
+        )
+
+        self._metrics_learner_impala_update_make_batch_if_necessary = Histogram(
+            name="rllib_learner_impala_update_make_batch_if_necessary_time",
+            description="Time spent on making a batch in the 'Learner.update()'.",
+            boundaries=DEFAULT_HISTOGRAM_BOUNDARIES_SHORT_EVENTS,
+            tag_keys=("rllib",),
+        )
+        self._metrics_learner_impala_update_make_batch_if_necessary.set_default_tags(
+            {"rllib": self.__class__.__name__}
+        )
+
+        self._metrics_learner_impala_get_learner_state_time = Histogram(
+            name="rllib_learner_impala_get_learner_state_time",
+            description="Time spent on get_state() in IMPALALearner.update().",
+            boundaries=DEFAULT_HISTOGRAM_BOUNDARIES_SHORT_EVENTS,
+            tag_keys=("rllib",),
+        )
+        self._metrics_learner_impala_get_learner_state_time.set_default_tags(
+            {"rllib": self.__class__.__name__}
+        )
+
+        # Set the aggregation threshold to the broadcast interval. We return
+        # a state at the same time the metrics are aggregated.
+        global BATCHES_PER_AGGREGATION
+        BATCHES_PER_AGGREGATION = self.config.broadcast_interval
+
+    @override(Learner)
+    def build(self) -> None:
+        super().build()
+
+        # APPO/IMPALA require RLock for thread safety around metrics.
+        self.metrics._threading_lock = threading.RLock()
+
+        # Aggregation signaling (replaces condition-variable contention) ---
+        self._agg_event = threading.Event()
+        self._submitted_updates = 0  # producer-side counter (update thread(s))
+        self._num_updates = 0  # learner-side counter
+        self._num_updates_lock = threading.Lock()
+        # Set the update kwargs passed in the main thread for use in the learner thread.
+        self._update_kwargs = {}
+
+        self._model_io_lock = threading.RLock()
+
+        self._learner_state_queue = Queue(maxsize=1)
+        self._learner_state_lock = threading.Lock()
+        self._learner_state = None
+
+        # Dict mapping module IDs to entropy Scheduler instances.
+        self.entropy_coeff_schedulers_per_module: Dict[
+            ModuleID, Scheduler
+        ] = LambdaDefaultDict(
+            lambda module_id: Scheduler(
+                fixed_value_or_schedule=(
+                    self.config.get_config_for_module(module_id).entropy_coeff
+                ),
+                framework=self.framework,
+                device=self._device,
+            )
+        )
+
+        # Create queues as bounded queues to create real back-pressure & stabilize
+        # GPU memory usage.
+
+        # Small loader in-queue to keep threads busy without flooding.
+        # TODO (simon): Do extensive testing to find an optimal queue size.
+        loader_qsize = max(2, 10 * self.config.num_gpu_loader_threads)
+        # Note, we are passing now the timesteps dictionary through the queue.
+        self._gpu_loader_in_queue: "Queue[tuple[TrainingData, Dict[str, Any]]]" = Queue(
+            maxsize=loader_qsize
+        )
+
+        # Learner in-queue must be tiny. 1 strictly serializes GPU-resident batches.
+        # TODO (simon): Add a parameter to define queue size.
+        if not hasattr(self, "_learner_thread_in_queue"):
+            self._learner_thread_in_queue: "Queue[tuple[Any, Dict[str, Any]]]" = Queue(
+                maxsize=self.config.learner_queue_size
+            )
+
+        # Get the rank of this learner, if necessary.
+        self._rank: int = (
+            torch.distributed.get_rank() if torch.distributed.is_initialized() else 0
+        )
+
+        # Define the out-queue for the metrics from the `_LearnerThread`.
+        # TODO (simon): Add types for items.
+        self._learner_thread_out_queue: "Queue[Dict[str, Any]]" = Queue()
+
+        # Create and start `_GPULoaderThread`(s).
+        if self.config.num_gpus_per_learner > 0:
+            self._gpu_loader_threads: List[threading.Thread] = [
+                _GPULoaderThread(
+                    in_queue=self._gpu_loader_in_queue,
+                    out_queue=self._learner_thread_in_queue,
+                    device=self._device,
+                    metrics_logger=self.metrics,
+                )
+                for _ in range(self.config.num_gpu_loader_threads)
+            ]
+            for t in self._gpu_loader_threads:
+                t.start()
+
+        # Create and start the `_LearnerThread`.
+        self._learner_thread: threading.Thread = _LearnerThread(
+            update_method=Learner.update,
+            in_queue=self._learner_thread_in_queue,
+            out_queue=self._learner_thread_out_queue,
+            learner=self,
+        )
+        self._learner_thread.start()
+
+    @override(Learner)
+    def update(
+        self,
+        training_data: TrainingData,
+        *,
+        timesteps: Dict[str, Any],
+        return_state: bool = False,
+        **kwargs,
+    ) -> ResultDict:
+        """
+
+        Args:
+            batch:
+            timesteps:
+            return_state: Whether to include one of the Learner worker's state from
+                after the update step in the returned results dict (under the
+                `_rl_module_state_after_update` key). Note that after an update, all
+                Learner workers' states should be identical, so we use the first
+                Learner's state here. Useful for avoiding an extra `get_weights()` call,
+                e.g. for synchronizing EnvRunner weights.
+            **kwargs:
+
+        Returns:
+
+        """
+        # Set the update kwargs passed in the main thread for use in the learner thread.
+        self._update_kwargs = kwargs
+
+        with TimerAndPrometheusLogger(self._metrics_learner_impala_update):
+            # Get the train batch from the object store.
+            with TimerAndPrometheusLogger(
+                self._metrics_learner_impala_update_solve_refs
+            ):
+                # Resolve object refs and ensure we have a proper batch object.
+                # TODO (simon): Check, if we can resolve the object references and
+                # run the pipeline on the GPULoaderThreads.
+                training_data.solve_refs()
+
+            with TimerAndPrometheusLogger(
+                self._metrics_learner_impala_update_make_batch_if_necessary
+            ):
+                batch = self._make_batch_if_necessary(training_data=training_data)
+                assert batch is not None
+
+        # Enqeue the batch (bounded backpressure).
+        if self.config.num_gpus_per_learner > 0:
+            # Pass timesteps alongside batch (no globals).
+            self._gpu_loader_in_queue.put((batch, timesteps))
+            # Only occasionally log loader queue size.
+            if (self._submitted_updates & 0xFF) == 0:
+                self.metrics.log_value(
+                    (ALL_MODULES, QUEUE_SIZE_GPU_LOADER_QUEUE),
+                    self._gpu_loader_in_queue.qsize(),
+                    window=1,
+                )
+                # TODO (simon): Check, if we want to get here stats from the
+                # RingBuffer.
+        else:
+            # No GPU loader: directly enqueue to learner queue.
+            _LearnerThread.enqueue(
+                self._learner_thread_in_queue, (batch, timesteps), self.metrics
+            )
+
+        # Return the module state, if requested and available.
+        if return_state:
+            try:
+                with self._learner_state_lock:
+                    self._learner_state = self._learner_state_queue.get_nowait()
+            except queue.Empty:
+                logger.debug("No learner state available in the queue yet.")
+
+        # Every 20th block call we submit results. Otherwise we keep the
+        # thread running without interruption to avoid thread contention.
+        self._submitted_updates += 1
+        if (self._submitted_updates % BATCHES_PER_AGGREGATION) != 0:
+            result = {}
+            if return_state and self._learner_state:
+                result["_rl_module_state_after_update"] = self._learner_state
+            return result
+
+        # Result submission: wait until learner finished BATCHES_PER_AGGREGATION updates (blocking).
+        self._agg_event.wait()
+        # Reset the aggregation event to keep the `_LearnerThread` running.
+        self._agg_event.clear()
+
+        if self._learner_thread_out_queue:
+            try:
+                result = self._learner_thread_out_queue.get(timeout=0.001)
+            except queue.Empty:
+                result = {}
+
+        # Return the module state, if requested and existent.
+        if return_state and self._learner_state:
+            result["_rl_module_state_after_update"] = self._learner_state
+
+        return result
+
+    @OverrideToImplementCustomLogic_CallToSuperRecommended
+    def before_gradient_based_update(self, *, timesteps: Dict[str, Any]) -> None:
+        super().before_gradient_based_update(timesteps=timesteps)
+
+        for module_id in self.module.keys():
+            # Update entropy coefficient via our Scheduler.
+            new_entropy_coeff = self.entropy_coeff_schedulers_per_module[
+                module_id
+            ].update(timestep=timesteps.get(NUM_ENV_STEPS_SAMPLED_LIFETIME, 0))
+            self.metrics.log_value(
+                (module_id, LEARNER_RESULTS_CURR_ENTROPY_COEFF_KEY),
+                new_entropy_coeff,
+                window=1,
+            )
+
+    @override(Learner)
+    def remove_module(self, module_id: str):
+        super().remove_module(module_id)
+        self.entropy_coeff_schedulers_per_module.pop(module_id)
+
+    @override(Learner)
+    def shutdown(self) -> None:
+        # Stop the learner thread deterministically: setting the stop event
+        # and enqueuing a sentinel wakes the consumer if it's blocked on
+        # `_in_queue.get()`. Then `join` ensures it has fully exited before
+        # we return, so any subsequent `ray.shutdown()`/interpreter teardown
+        # can't race with the daemon thread.
+        thread = getattr(self, "_learner_thread", None)
+        if thread is not None and thread.is_alive():
+            thread.request_stop()
+            thread.join(timeout=5.0)
+
+    @classmethod
+    @override(Learner)
+    def rl_module_required_apis(cls) -> list[type]:
+        # In order for a PPOLearner to update an RLModule, it must implement the
+        # following APIs:
+        return [ValueFunctionAPI]
+
+
+ImpalaLearner = IMPALALearner
+
+
+class _GPULoaderThread(threading.Thread):
+    def __init__(
+        self,
+        *,
+        in_queue: "Queue[tuple[TrainingData, Dict[str, Any]]]",
+        out_queue: "Queue[tuple[Any, Dict[str, Any]]]",
+        device: "torch.device",
+        metrics_logger: MetricsLogger,
+    ):
+        super().__init__(name="_GPULoaderThread")
+        self.daemon = True
+
+        self._in_queue = in_queue
+        self._out_queue = out_queue
+        self._device = device
+        self.metrics = metrics_logger
+
+        # Use a single CUDA stream for each loader thread.
+        self._use_cuda_stream = (
+            torch is not None
+            and hasattr(torch, "cuda")
+            and device is not None
+            and getattr(device, "type", None) == "cuda"
+        )
+        self._stream = (
+            torch.cuda.Stream(device=self._device) if self._use_cuda_stream else None
+        )
+
+        self._metrics_impala_gpu_loader_thread_step_time = Histogram(
+            name="rllib_learner_impala_gpu_loader_thread_step_time",
+            description="Time taken in seconds for gpu loader thread _step.",
+            boundaries=DEFAULT_HISTOGRAM_BOUNDARIES_SHORT_EVENTS,
+            tag_keys=("rllib",),
+        )
+        self._metrics_impala_gpu_loader_thread_step_time.set_default_tags(
+            {"rllib": "IMPALA/GPULoaderThread"}
+        )
+
+        self._metrics_impala_gpu_loader_thread_step_in_queue_get_time = Histogram(
+            name="rllib_learner_impala_gpu_loader_thread_step_get_time",
+            description="Time taken in seconds for gpu loader thread _step _in_queue.get().",
+            boundaries=DEFAULT_HISTOGRAM_BOUNDARIES_SHORT_EVENTS,
+            tag_keys=("rllib",),
+        )
+        self._metrics_impala_gpu_loader_thread_step_in_queue_get_time.set_default_tags(
+            {"rllib": "IMPALA/GPULoaderThread"}
+        )
+
+        self._metrics_impala_gpu_loader_thread_step_load_to_gpu_time = Histogram(
+            name="rllib_learner_impala_gpu_loader_thread_step_load_to_gpu_time",
+            description="Time taken in seconds for GPU loader thread _step to load batch to GPU.",
+            boundaries=DEFAULT_HISTOGRAM_BOUNDARIES_SHORT_EVENTS,
+            tag_keys=("rllib",),
+        )
+        self._metrics_impala_gpu_loader_thread_step_load_to_gpu_time.set_default_tags(
+            {"rllib": "IMPALA/GPULoaderThread"}
+        )
+
+        self._metrics_impala_gpu_loader_thread_in_qsize_beginning_of_step = Gauge(
+            name="rllib_impala_gpu_loader_thread_in_qsize_beginning_of_step",
+            description="Size of the _GPULoaderThread in-queue size, at the beginning of the step.",
+            tag_keys=("rllib",),
+        )
+        self._metrics_impala_gpu_loader_thread_in_qsize_beginning_of_step.set_default_tags(
+            {"rllib": "IMPALA/GPULoaderThread"}
+        )
+
+    # Robust pinned-memory copy: fall back if batch contains CUDA tensors already.
+    # TODO (simon): Find a more compliant solution.
+    def _to_device_safe(self, batch):
+        try:
+            return batch.to_device(self._device, pin_memory=True)
+        except RuntimeError as e:
+            msg = str(e)
+            if "only dense CPU tensors can be pinned" in msg or "pin_memory" in msg:
+                return batch.to_device(self._device, pin_memory=False)
+            raise
+
+    def run(self) -> None:
+        while True:
+            with TimerAndPrometheusLogger(
+                self._metrics_impala_gpu_loader_thread_step_time
+            ):
+                self._step()
+
+    def _step(self) -> None:
+        self._metrics_impala_gpu_loader_thread_in_qsize_beginning_of_step.set(
+            value=self._in_queue.qsize()
+        )
+        # Get a new batch (CPU) and the global timesteps from the loader in--queue (blocking).
+        with self.metrics.log_time((ALL_MODULES, GPU_LOADER_QUEUE_WAIT_TIMER)):
+            with TimerAndPrometheusLogger(
+                self._metrics_impala_gpu_loader_thread_step_in_queue_get_time
+            ):
+                ma_batch_on_cpu, timesteps = self._in_queue.get()
+
+        # Load the batch onto the GPU device; enable pinned memory for async copies.
+        with self.metrics.log_time((ALL_MODULES, GPU_LOADER_LOAD_TO_GPU_TIMER)):
+            if self._use_cuda_stream and self._stream is not None:
+                # Issue copies on a non-default stream so they can overlap with compute.
+                with torch.cuda.stream(self._stream):
+                    ma_batch_on_gpu = self._to_device_safe(ma_batch_on_cpu)
+                    # TODO (simon): Maybe use the `use_stream` in `convert_to_tensor`.
+                # No explicit synching here. Consumer will naturally serialize when needed.
+            else:
+                ma_batch_on_gpu = self._to_device_safe(ma_batch_on_cpu)
+
+        # Enqueue to Learner thread’s in-queue (GPU-resident batch and global timesteps).
+        _LearnerThread.enqueue(
+            self._out_queue, (ma_batch_on_gpu, timesteps), self.metrics
+        )
+
+
+class _LearnerThread(threading.Thread):
+    def __init__(
+        self,
+        *,
+        update_method,
+        in_queue: "Queue[tuple[Any, Dict[str, Any]]]",
+        out_queue: "Queue[Dict[str, Any]]",
+        learner: IMPALALearner,
+    ):
+        super().__init__(name="_LearnerThread")
+        self.daemon = True
+        self.learner = learner
+        self._update_method = update_method
+        # Note, we pass now the timesteps dictionary through the queue.
+        self._in_queue: "queue.Queue[tuple[Any, Dict[str, Any]]]" = in_queue
+        # TODO (simon): Type hints.
+        self._out_queue = out_queue
+        self._stop_event = threading.Event()
+
+        # Ray metrics
+        self._metrics_learner_impala_thread_step = Histogram(
+            name="rllib_learner_impala_learner_thread_step_time",
+            description="Time taken in seconds for learner thread _step.",
+            boundaries=DEFAULT_HISTOGRAM_BOUNDARIES_SHORT_EVENTS,
+            tag_keys=("rllib",),
+        )
+        self._metrics_learner_impala_thread_step.set_default_tags(
+            {"rllib": "IMPALA/LearnerThread"}
+        )
+
+        self._metrics_learner_impala_thread_step_update = Histogram(
+            name="rllib_learner_impala_learner_thread_step_update_time",
+            description="Time taken in seconds for learner thread _step update.",
+            boundaries=DEFAULT_HISTOGRAM_BOUNDARIES_SHORT_EVENTS,
+            tag_keys=("rllib",),
+        )
+        self._metrics_learner_impala_thread_step_update.set_default_tags(
+            {"rllib": "IMPALA/LearnerThread"}
+        )
+
+        # Stop cleanly at interpreter shutdown so the daemon thread doesn't
+        # get killed mid-call inside an auto_init-wrapped Ray API (which
+        # would otherwise trigger e.g. `start_reaper` -> `preexec_fn not
+        # supported at interpreter shutdown`). Use a weakref so this hook
+        # doesn't pin the thread (and therefore the Learner) alive.
+        weak_self = weakref.ref(self)
+
+        def _request_stop_at_exit():
+            t = weak_self()
+            if t is not None:
+                t.request_stop()
+
+        atexit.register(_request_stop_at_exit)
+
+    # Keeps compatibility, but thread-safe.
+    @property
+    def stopped(self) -> bool:
+        return self._stop_event.is_set()
+
+    # Call this to stop the thread and wake it if it's blocked on .get()
+    def request_stop(self) -> None:
+        self._stop_event.set()
+        # Wake the consumer if it's blocked on an empty queue
+        try:
+            self._in_queue.put_nowait(_STOP_SENTINEL)
+        except queue.Full:
+            # If the queue is full, the consumer will wake soon anyway.
+            logger.warning(
+                "_LearnerThread.request_stop(): in_queue is full; cannot enqueue stop sentinel."
+            )
+
+    def run(self) -> None:
+        while True:
+            # Returns always `True` until stop-signal/sentinel is sent.
+            if not self.step():
+                break
+
+    def step(self) -> bool:
+        # Get a batch and wait, if the input queue is empty (blocking; no polling).
+        with self.learner.metrics.log_time(
+            (ALL_MODULES, LEARNER_THREAD_IN_QUEUE_WAIT_TIMER)
+        ):
+            item = self._in_queue.get()
+
+        # Handle the stop/sentinel signal(s).
+        # TODO (simon): Check, if we need `None` for belt-and-suspenders/comp.
+        if item is _STOP_SENTINEL or self.stopped:
+            try:
+                self._in_queue.task_done()
+            except Exception:
+                logger.warning(
+                    "_LearnerThread._in_queue.task_done() failed during stop handling."
+                )
+            # Signal `run` to exit.
+            return False
+
+        # Extract the multi-agent batch and the timesteps dictionary.
+        ma_batch_on_gpu, timesteps = item
+
+        # Update the `RLModule`, but do not reduce metrics.
+        with self.learner.metrics.log_time((ALL_MODULES, LEARNER_THREAD_UPDATE_TIMER)):
+            with TimerAndPrometheusLogger(
+                self._metrics_learner_impala_thread_step_update
+            ):
+                self._update_method(
+                    self=self.learner,
+                    training_data=TrainingData(batch=ma_batch_on_gpu),
+                    timesteps=timesteps,
+                    _no_metrics_reduce=True,
+                    # Include the learner update kwargs set in the main thread.
+                    **self.learner._update_kwargs,
+                )
+
+        # Signal queue done (unblocks producer’s put when bounded)
+        try:
+            self._in_queue.task_done()
+        finally:
+            # Set the Aggregation counter and signal this event (atomic).
+            with self.learner._num_updates_lock:
+                self.learner._num_updates += 1
+                # Check, if we need to aggregate.
+                do_agg = self.learner._num_updates == BATCHES_PER_AGGREGATION
+                if do_agg:
+                    # Reset the update counter inside the lock.
+                    self.learner._num_updates = 0
+
+            # If we need to aggregate, reduce metrics and queue them.
+            if do_agg:
+                # If in multi-learner setup, safeguard state retrieval within barriers.
+                if torch.distributed.is_initialized():
+                    torch.distributed.barrier()
+                # Only the first rank retrieves the state.
+                if self.learner._rank == 0:
+                    with self.learner._model_io_lock, torch.inference_mode():
+                        learner_state = self.learner.get_state(
+                            # Only return the state of those RLModules that are trainable.
+                            components=[
+                                COMPONENT_RL_MODULE + "/" + mid
+                                for mid in self.learner.module.keys()
+                                if self.learner.should_module_be_updated(mid)
+                            ],
+                            inference_only=True,
+                        )
+                        learner_state[COMPONENT_RL_MODULE] = ray.put(
+                            learner_state[COMPONENT_RL_MODULE]
+                        )
+                    try:
+                        if (self.learner._submitted_updates & ~0xFF) != (
+                            (self.learner._submitted_updates - BATCHES_PER_AGGREGATION)
+                            & ~0xFF
+                        ):
+                            with self.learner._learner_state_lock:
+                                self.learner.metrics.log_value(
+                                    (ALL_MODULES, "learner_thread_state_queue_size"),
+                                    self.learner._learner_state_queue.qsize(),
+                                    window=1,
+                                )
+                        # Remove any old learner state in the queue.
+                        self.learner._learner_state_queue.get_nowait()
+                    except queue.Empty:
+                        logger.debug("No old learner state to remove from the queue.")
+
+                    # Pass the learner state into the queue to the main process.
+                    self.learner._learner_state_queue.put_nowait(learner_state)
+                self.learner.metrics.log_value(
+                    (ALL_MODULES, "learner_thread_out_queue_size"),
+                    self._out_queue.qsize(),
+                    window=1,
+                )
+
+                # Reduce metrics and pass them into the queue for the main process.
+                self._out_queue.put(self.learner.metrics.reduce())
+                # Notify all listeners that aggregation is done and results can be
+                # retrieved.
+                self.learner._agg_event.set()
+                if torch.distributed.is_initialized():
+                    torch.distributed.barrier()
+
+        # Keep running (see `run` method).
+        return True
+
+    @staticmethod
+    def enqueue(
+        learner_queue: "queue.Queue[tuple[Any, Dict[str, Any]]]",
+        batch_with_ts,
+        metrics: MetricsLogger,
+    ):
+        # Put the batch into the queue (blocking if thread is updating).
+        learner_queue.put(batch_with_ts, block=True)

rllib/algorithms/impala/impala_tf_policy.py

@@ -0,0 +1,449 @@
+"""Adapted from A3CTFPolicy to add V-trace.
+
+Keep in sync with changes to A3CTFPolicy and VtraceSurrogatePolicy."""
+
+import logging
+from typing import Dict, List, Optional, Type, Union
+
+import gymnasium as gym
+import numpy as np
+
+from ray.rllib.algorithms.impala import vtrace_tf as vtrace
+from ray.rllib.evaluation.postprocessing import compute_bootstrap_value
+from ray.rllib.models.modelv2 import ModelV2
+from ray.rllib.models.tf.tf_action_dist import Categorical, TFActionDistribution
+from ray.rllib.policy.dynamic_tf_policy_v2 import DynamicTFPolicyV2
+from ray.rllib.policy.eager_tf_policy_v2 import EagerTFPolicyV2
+from ray.rllib.policy.sample_batch import SampleBatch
+from ray.rllib.policy.tf_mixins import (
+    EntropyCoeffSchedule,
+    GradStatsMixin,
+    LearningRateSchedule,
+    ValueNetworkMixin,
+)
+from ray.rllib.utils import force_list
+from ray.rllib.utils.annotations import override
+from ray.rllib.utils.framework import try_import_tf
+from ray.rllib.utils.tf_utils import explained_variance
+from ray.rllib.utils.typing import (
+    LocalOptimizer,
+    ModelGradients,
+    TensorType,
+    TFPolicyV2Type,
+)
+
+tf1, tf, tfv = try_import_tf()
+
+logger = logging.getLogger(__name__)
+
+
+class VTraceLoss:
+    def __init__(
+        self,
+        actions,
+        actions_logp,
+        actions_entropy,
+        dones,
+        behaviour_action_logp,
+        behaviour_logits,
+        target_logits,
+        discount,
+        rewards,
+        values,
+        bootstrap_value,
+        dist_class,
+        model,
+        valid_mask,
+        config,
+        vf_loss_coeff=0.5,
+        entropy_coeff=0.01,
+        clip_rho_threshold=1.0,
+        clip_pg_rho_threshold=1.0,
+    ):
+        """Policy gradient loss with vtrace importance weighting.
+
+        VTraceLoss takes tensors of shape [T, B, ...], where `B` is the
+        batch_size. The reason we need to know `B` is for V-trace to properly
+        handle episode cut boundaries.
+
+        Args:
+            actions: An int|float32 tensor of shape [T, B, ACTION_SPACE].
+            actions_logp: A float32 tensor of shape [T, B].
+            actions_entropy: A float32 tensor of shape [T, B].
+            dones: A bool tensor of shape [T, B].
+            behaviour_action_logp: Tensor of shape [T, B].
+            behaviour_logits: A list with length of ACTION_SPACE of float32
+                tensors of shapes
+                [T, B, ACTION_SPACE[0]],
+                ...,
+                [T, B, ACTION_SPACE[-1]]
+            target_logits: A list with length of ACTION_SPACE of float32
+                tensors of shapes
+                [T, B, ACTION_SPACE[0]],
+                ...,
+                [T, B, ACTION_SPACE[-1]]
+            discount: A float32 scalar.
+            rewards: A float32 tensor of shape [T, B].
+            values: A float32 tensor of shape [T, B].
+            bootstrap_value: A float32 tensor of shape [B].
+            dist_class: action distribution class for logits.
+            valid_mask: A bool tensor of valid RNN input elements (#2992).
+            config: Algorithm config dict.
+        """
+
+        # Compute vtrace on the CPU for better performance.
+        with tf.device("/cpu:0"):
+            self.vtrace_returns = vtrace.multi_from_logits(
+                behaviour_action_log_probs=behaviour_action_logp,
+                behaviour_policy_logits=behaviour_logits,
+                target_policy_logits=target_logits,
+                actions=tf.unstack(actions, axis=2),
+                discounts=tf.cast(~tf.cast(dones, tf.bool), tf.float32) * discount,
+                rewards=rewards,
+                values=values,
+                bootstrap_value=bootstrap_value,
+                dist_class=dist_class,
+                model=model,
+                clip_rho_threshold=tf.cast(clip_rho_threshold, tf.float32),
+                clip_pg_rho_threshold=tf.cast(clip_pg_rho_threshold, tf.float32),
+            )
+            self.value_targets = self.vtrace_returns.vs
+
+        # The policy gradients loss.
+        masked_pi_loss = tf.boolean_mask(
+            actions_logp * self.vtrace_returns.pg_advantages, valid_mask
+        )
+        self.pi_loss = -tf.reduce_sum(masked_pi_loss)
+        self.mean_pi_loss = -tf.reduce_mean(masked_pi_loss)
+
+        # The baseline loss.
+        delta = tf.boolean_mask(values - self.vtrace_returns.vs, valid_mask)
+        delta_squarred = tf.math.square(delta)
+        self.vf_loss = 0.5 * tf.reduce_sum(delta_squarred)
+        self.mean_vf_loss = 0.5 * tf.reduce_mean(delta_squarred)
+
+        # The entropy loss.
+        masked_entropy = tf.boolean_mask(actions_entropy, valid_mask)
+        self.entropy = tf.reduce_sum(masked_entropy)
+        self.mean_entropy = tf.reduce_mean(masked_entropy)
+
+        # The summed weighted loss.
+        self.total_loss = self.pi_loss - self.entropy * entropy_coeff
+
+        # Optional vf loss (or in a separate term due to separate
+        # optimizers/networks).
+        self.loss_wo_vf = self.total_loss
+        if not config["_separate_vf_optimizer"]:
+            self.total_loss += self.vf_loss * vf_loss_coeff
+
+
+def _make_time_major(policy, seq_lens, tensor):
+    """Swaps batch and trajectory axis.
+
+    Args:
+        policy: Policy reference
+        seq_lens: Sequence lengths if recurrent or None
+        tensor: A tensor or list of tensors to reshape.
+            trajectory item.
+
+    Returns:
+        res: A tensor with swapped axes or a list of tensors with
+        swapped axes.
+    """
+    if isinstance(tensor, list):
+        return [_make_time_major(policy, seq_lens, t) for t in tensor]
+
+    if policy.is_recurrent():
+        B = tf.shape(seq_lens)[0]
+        T = tf.shape(tensor)[0] // B
+    else:
+        # Important: chop the tensor into batches at known episode cut
+        # boundaries.
+        # TODO: (sven) this is kind of a hack and won't work for
+        #  batch_mode=complete_episodes.
+        T = policy.config["rollout_fragment_length"]
+        B = tf.shape(tensor)[0] // T
+    rs = tf.reshape(tensor, tf.concat([[B, T], tf.shape(tensor)[1:]], axis=0))
+
+    # swap B and T axes
+    res = tf.transpose(rs, [1, 0] + list(range(2, 1 + int(tf.shape(tensor).shape[0]))))
+
+    return res
+
+
+class VTraceClipGradients:
+    """VTrace version of gradient computation logic."""
+
+    def __init__(self):
+        """No special initialization required."""
+        pass
+
+    def compute_gradients_fn(
+        self, optimizer: LocalOptimizer, loss: TensorType
+    ) -> ModelGradients:
+        # Supporting more than one loss/optimizer.
+        trainable_variables = self.model.trainable_variables()
+        if self.config["_tf_policy_handles_more_than_one_loss"]:
+            optimizers = force_list(optimizer)
+            losses = force_list(loss)
+            assert len(optimizers) == len(losses)
+            clipped_grads_and_vars = []
+            for optim, loss_ in zip(optimizers, losses):
+                grads_and_vars = optim.compute_gradients(loss_, trainable_variables)
+                clipped_g_and_v = []
+                for g, v in grads_and_vars:
+                    if g is not None:
+                        clipped_g, _ = tf.clip_by_global_norm(
+                            [g], self.config["grad_clip"]
+                        )
+                        clipped_g_and_v.append((clipped_g[0], v))
+                clipped_grads_and_vars.append(clipped_g_and_v)
+
+            self.grads = [g for g_and_v in clipped_grads_and_vars for (g, v) in g_and_v]
+        # Only one optimizer and and loss term.
+        else:
+            grads_and_vars = optimizer.compute_gradients(
+                loss, self.model.trainable_variables()
+            )
+            grads = [g for (g, v) in grads_and_vars]
+            self.grads, _ = tf.clip_by_global_norm(grads, self.config["grad_clip"])
+            clipped_grads_and_vars = list(zip(self.grads, trainable_variables))
+
+        return clipped_grads_and_vars
+
+
+class VTraceOptimizer:
+    """Optimizer function for VTrace policies."""
+
+    def __init__(self):
+        pass
+
+    # TODO: maybe standardize this function, so the choice of optimizers are more
+    #  predictable for common algorithms.
+    def optimizer(
+        self,
+    ) -> Union["tf.keras.optimizers.Optimizer", List["tf.keras.optimizers.Optimizer"]]:
+        config = self.config
+        if config["opt_type"] == "adam":
+            if config["framework"] == "tf2":
+                optim = tf.keras.optimizers.Adam(self.cur_lr)
+                if config["_separate_vf_optimizer"]:
+                    return optim, tf.keras.optimizers.Adam(config["_lr_vf"])
+            else:
+                optim = tf1.train.AdamOptimizer(self.cur_lr)
+                if config["_separate_vf_optimizer"]:
+                    return optim, tf1.train.AdamOptimizer(config["_lr_vf"])
+        else:
+            if config["_separate_vf_optimizer"]:
+                raise ValueError(
+                    "RMSProp optimizer not supported for separate"
+                    "vf- and policy losses yet! Set `opt_type=adam`"
+                )
+
+            if tfv == 2:
+                optim = tf.keras.optimizers.RMSprop(
+                    self.cur_lr, config["decay"], config["momentum"], config["epsilon"]
+                )
+            else:
+                optim = tf1.train.RMSPropOptimizer(
+                    self.cur_lr, config["decay"], config["momentum"], config["epsilon"]
+                )
+
+        return optim
+
+
+# We need this builder function because we want to share the same
+# custom logics between TF1 dynamic and TF2 eager policies.
+def get_impala_tf_policy(name: str, base: TFPolicyV2Type) -> TFPolicyV2Type:
+    """Construct an ImpalaTFPolicy inheriting either dynamic or eager base policies.
+
+    Args:
+        base: Base class for this policy. DynamicTFPolicyV2 or EagerTFPolicyV2.
+
+    Returns:
+        A TF Policy to be used with Impala.
+    """
+
+    # VTrace mixins are placed in front of more general mixins to make sure
+    # their functions like optimizer() overrides all the other implementations
+    # (e.g., LearningRateSchedule.optimizer())
+    class ImpalaTFPolicy(
+        VTraceClipGradients,
+        VTraceOptimizer,
+        LearningRateSchedule,
+        EntropyCoeffSchedule,
+        GradStatsMixin,
+        ValueNetworkMixin,
+        base,
+    ):
+        def __init__(
+            self,
+            observation_space,
+            action_space,
+            config,
+            existing_model=None,
+            existing_inputs=None,
+        ):
+            # First thing first, enable eager execution if necessary.
+            base.enable_eager_execution_if_necessary()
+
+            # Initialize base class.
+            base.__init__(
+                self,
+                observation_space,
+                action_space,
+                config,
+                existing_inputs=existing_inputs,
+                existing_model=existing_model,
+            )
+            ValueNetworkMixin.__init__(self, config)
+
+            # If Learner API is used, we don't need any loss-specific mixins.
+            # However, we also would like to avoid creating special Policy-subclasses
+            # for this as the entire Policy concept will soon not be used anymore with
+            # the new Learner- and RLModule APIs.
+            GradStatsMixin.__init__(self)
+            VTraceClipGradients.__init__(self)
+            VTraceOptimizer.__init__(self)
+            LearningRateSchedule.__init__(self, config["lr"], config["lr_schedule"])
+            EntropyCoeffSchedule.__init__(
+                self, config["entropy_coeff"], config["entropy_coeff_schedule"]
+            )
+
+            # Note: this is a bit ugly, but loss and optimizer initialization must
+            # happen after all the MixIns are initialized.
+            self.maybe_initialize_optimizer_and_loss()
+
+        @override(base)
+        def loss(
+            self,
+            model: Union[ModelV2, "tf.keras.Model"],
+            dist_class: Type[TFActionDistribution],
+            train_batch: SampleBatch,
+        ) -> Union[TensorType, List[TensorType]]:
+            model_out, _ = model(train_batch)
+            action_dist = dist_class(model_out, model)
+
+            if isinstance(self.action_space, gym.spaces.Discrete):
+                is_multidiscrete = False
+                output_hidden_shape = [self.action_space.n]
+            elif isinstance(self.action_space, gym.spaces.MultiDiscrete):
+                is_multidiscrete = True
+                output_hidden_shape = self.action_space.nvec.astype(np.int32)
+            else:
+                is_multidiscrete = False
+                output_hidden_shape = 1
+
+            def make_time_major(*args, **kw):
+                return _make_time_major(
+                    self, train_batch.get(SampleBatch.SEQ_LENS), *args, **kw
+                )
+
+            actions = train_batch[SampleBatch.ACTIONS]
+            dones = train_batch[SampleBatch.TERMINATEDS]
+            rewards = train_batch[SampleBatch.REWARDS]
+            behaviour_action_logp = train_batch[SampleBatch.ACTION_LOGP]
+            behaviour_logits = train_batch[SampleBatch.ACTION_DIST_INPUTS]
+            unpacked_behaviour_logits = tf.split(
+                behaviour_logits, output_hidden_shape, axis=1
+            )
+            unpacked_outputs = tf.split(model_out, output_hidden_shape, axis=1)
+            values = model.value_function()
+            values_time_major = make_time_major(values)
+            bootstrap_values_time_major = make_time_major(
+                train_batch[SampleBatch.VALUES_BOOTSTRAPPED]
+            )
+            bootstrap_value = bootstrap_values_time_major[-1]
+
+            if self.is_recurrent():
+                max_seq_len = tf.reduce_max(train_batch[SampleBatch.SEQ_LENS])
+                mask = tf.sequence_mask(train_batch[SampleBatch.SEQ_LENS], max_seq_len)
+                mask = tf.reshape(mask, [-1])
+            else:
+                mask = tf.ones_like(rewards)
+
+            # Prepare actions for loss
+            loss_actions = (
+                actions if is_multidiscrete else tf.expand_dims(actions, axis=1)
+            )
+
+            # Inputs are reshaped from [B * T] => [(T|T-1), B] for V-trace calc.
+            self.vtrace_loss = VTraceLoss(
+                actions=make_time_major(loss_actions),
+                actions_logp=make_time_major(action_dist.logp(actions)),
+                actions_entropy=make_time_major(action_dist.multi_entropy()),
+                dones=make_time_major(dones),
+                behaviour_action_logp=make_time_major(behaviour_action_logp),
+                behaviour_logits=make_time_major(unpacked_behaviour_logits),
+                target_logits=make_time_major(unpacked_outputs),
+                discount=self.config["gamma"],
+                rewards=make_time_major(rewards),
+                values=values_time_major,
+                bootstrap_value=bootstrap_value,
+                dist_class=Categorical if is_multidiscrete else dist_class,
+                model=model,
+                valid_mask=make_time_major(mask),
+                config=self.config,
+                vf_loss_coeff=self.config["vf_loss_coeff"],
+                entropy_coeff=self.entropy_coeff,
+                clip_rho_threshold=self.config["vtrace_clip_rho_threshold"],
+                clip_pg_rho_threshold=self.config["vtrace_clip_pg_rho_threshold"],
+            )
+
+            if self.config.get("_separate_vf_optimizer"):
+                return self.vtrace_loss.loss_wo_vf, self.vtrace_loss.vf_loss
+            else:
+                return self.vtrace_loss.total_loss
+
+        @override(base)
+        def stats_fn(self, train_batch: SampleBatch) -> Dict[str, TensorType]:
+            values_batched = _make_time_major(
+                self,
+                train_batch.get(SampleBatch.SEQ_LENS),
+                self.model.value_function(),
+            )
+
+            return {
+                "cur_lr": tf.cast(self.cur_lr, tf.float64),
+                "policy_loss": self.vtrace_loss.mean_pi_loss,
+                "entropy": self.vtrace_loss.mean_entropy,
+                "entropy_coeff": tf.cast(self.entropy_coeff, tf.float64),
+                "var_gnorm": tf.linalg.global_norm(self.model.trainable_variables()),
+                "vf_loss": self.vtrace_loss.mean_vf_loss,
+                "vf_explained_var": explained_variance(
+                    tf.reshape(self.vtrace_loss.value_targets, [-1]),
+                    tf.reshape(values_batched, [-1]),
+                ),
+            }
+
+        @override(base)
+        def postprocess_trajectory(
+            self,
+            sample_batch: SampleBatch,
+            other_agent_batches: Optional[SampleBatch] = None,
+            episode=None,
+        ):
+            # Call super's postprocess_trajectory first.
+            # sample_batch = super().postprocess_trajectory(
+            #    sample_batch, other_agent_batches, episode
+            # )
+
+            if self.config["vtrace"]:
+                # Add the SampleBatch.VALUES_BOOTSTRAPPED column, which we'll need
+                # inside the loss for vtrace calculations.
+                sample_batch = compute_bootstrap_value(sample_batch, self)
+
+            return sample_batch
+
+        @override(base)
+        def get_batch_divisibility_req(self) -> int:
+            return self.config["rollout_fragment_length"]
+
+    ImpalaTFPolicy.__name__ = name
+    ImpalaTFPolicy.__qualname__ = name
+
+    return ImpalaTFPolicy
+
+
+ImpalaTF1Policy = get_impala_tf_policy("ImpalaTF1Policy", DynamicTFPolicyV2)
+ImpalaTF2Policy = get_impala_tf_policy("ImpalaTF2Policy", EagerTFPolicyV2)

rllib/algorithms/impala/impala_torch_policy.py

@@ -0,0 +1,425 @@
+import logging
+from typing import Dict, List, Optional, Type, Union
+
+import gymnasium as gym
+import numpy as np
+
+import ray
+from ray.rllib.evaluation.postprocessing import compute_bootstrap_value
+from ray.rllib.models.action_dist import ActionDistribution
+from ray.rllib.models.modelv2 import ModelV2
+from ray.rllib.models.torch.torch_action_dist import TorchCategorical
+from ray.rllib.policy.sample_batch import SampleBatch
+from ray.rllib.policy.torch_mixins import (
+    EntropyCoeffSchedule,
+    LearningRateSchedule,
+    ValueNetworkMixin,
+)
+from ray.rllib.policy.torch_policy_v2 import TorchPolicyV2
+from ray.rllib.utils.annotations import override
+from ray.rllib.utils.framework import try_import_torch
+from ray.rllib.utils.numpy import convert_to_numpy
+from ray.rllib.utils.torch_utils import (
+    apply_grad_clipping,
+    explained_variance,
+    global_norm,
+    sequence_mask,
+)
+from ray.rllib.utils.typing import TensorType
+
+torch, nn = try_import_torch()
+
+logger = logging.getLogger(__name__)
+
+
+class VTraceLoss:
+    def __init__(
+        self,
+        actions,
+        actions_logp,
+        actions_entropy,
+        dones,
+        behaviour_action_logp,
+        behaviour_logits,
+        target_logits,
+        discount,
+        rewards,
+        values,
+        bootstrap_value,
+        dist_class,
+        model,
+        valid_mask,
+        config,
+        vf_loss_coeff=0.5,
+        entropy_coeff=0.01,
+        clip_rho_threshold=1.0,
+        clip_pg_rho_threshold=1.0,
+    ):
+        """Policy gradient loss with vtrace importance weighting.
+
+        VTraceLoss takes tensors of shape [T, B, ...], where `B` is the
+        batch_size. The reason we need to know `B` is for V-trace to properly
+        handle episode cut boundaries.
+
+        Args:
+            actions: An int|float32 tensor of shape [T, B, ACTION_SPACE].
+            actions_logp: A float32 tensor of shape [T, B].
+            actions_entropy: A float32 tensor of shape [T, B].
+            dones: A bool tensor of shape [T, B].
+            behaviour_action_logp: Tensor of shape [T, B].
+            behaviour_logits: A list with length of ACTION_SPACE of float32
+                tensors of shapes
+                [T, B, ACTION_SPACE[0]],
+                ...,
+                [T, B, ACTION_SPACE[-1]]
+            target_logits: A list with length of ACTION_SPACE of float32
+                tensors of shapes
+                [T, B, ACTION_SPACE[0]],
+                ...,
+                [T, B, ACTION_SPACE[-1]]
+            discount: A float32 scalar.
+            rewards: A float32 tensor of shape [T, B].
+            values: A float32 tensor of shape [T, B].
+            bootstrap_value: A float32 tensor of shape [B].
+            dist_class: action distribution class for logits.
+            valid_mask: A bool tensor of valid RNN input elements (#2992).
+            config: Algorithm config dict.
+        """
+        import ray.rllib.algorithms.impala.vtrace_torch as vtrace
+
+        if valid_mask is None:
+            valid_mask = torch.ones_like(actions_logp)
+
+        # Compute vtrace on the CPU for better perf
+        # (devices handled inside `vtrace.multi_from_logits`).
+        device = behaviour_action_logp[0].device
+        self.vtrace_returns = vtrace.multi_from_logits(
+            behaviour_action_log_probs=behaviour_action_logp,
+            behaviour_policy_logits=behaviour_logits,
+            target_policy_logits=target_logits,
+            actions=torch.unbind(actions, dim=2),
+            discounts=(1.0 - dones.float()) * discount,
+            rewards=rewards,
+            values=values,
+            bootstrap_value=bootstrap_value,
+            dist_class=dist_class,
+            model=model,
+            clip_rho_threshold=clip_rho_threshold,
+            clip_pg_rho_threshold=clip_pg_rho_threshold,
+        )
+        # Move v-trace results back to GPU for actual loss computing.
+        self.value_targets = self.vtrace_returns.vs.to(device)
+
+        # The policy gradients loss.
+        self.pi_loss = -torch.sum(
+            actions_logp * self.vtrace_returns.pg_advantages.to(device) * valid_mask
+        )
+
+        # The baseline loss.
+        delta = (values - self.value_targets) * valid_mask
+        self.vf_loss = 0.5 * torch.sum(torch.pow(delta, 2.0))
+
+        # The entropy loss.
+        self.entropy = torch.sum(actions_entropy * valid_mask)
+        self.mean_entropy = self.entropy / torch.sum(valid_mask)
+
+        # The summed weighted loss.
+        self.total_loss = self.pi_loss - self.entropy * entropy_coeff
+
+        # Optional vf loss (or in a separate term due to separate
+        # optimizers/networks).
+        self.loss_wo_vf = self.total_loss
+        if not config["_separate_vf_optimizer"]:
+            self.total_loss += self.vf_loss * vf_loss_coeff
+
+
+def make_time_major(policy, seq_lens, tensor):
+    """Swaps batch and trajectory axis.
+
+    Args:
+        policy: Policy reference
+        seq_lens: Sequence lengths if recurrent or None
+        tensor: A tensor or list of tensors to reshape.
+
+    Returns:
+        res: A tensor with swapped axes or a list of tensors with
+        swapped axes.
+    """
+    if isinstance(tensor, (list, tuple)):
+        return [make_time_major(policy, seq_lens, t) for t in tensor]
+
+    if policy.is_recurrent():
+        B = seq_lens.shape[0]
+        T = tensor.shape[0] // B
+    else:
+        # Important: chop the tensor into batches at known episode cut
+        # boundaries.
+        # TODO: (sven) this is kind of a hack and won't work for
+        #  batch_mode=complete_episodes.
+        T = policy.config["rollout_fragment_length"]
+        B = tensor.shape[0] // T
+    rs = torch.reshape(tensor, [B, T] + list(tensor.shape[1:]))
+
+    # Swap B and T axes.
+    res = torch.transpose(rs, 1, 0)
+
+    return res
+
+
+class VTraceOptimizer:
+    """Optimizer function for VTrace torch policies."""
+
+    def __init__(self):
+        pass
+
+    def optimizer(
+        self,
+    ) -> Union[List["torch.optim.Optimizer"], "torch.optim.Optimizer"]:
+
+        if self.config["_separate_vf_optimizer"]:
+            # Figure out, which parameters of the model belong to the value
+            # function (and which to the policy net).
+            dummy_batch = self._lazy_tensor_dict(
+                self._get_dummy_batch_from_view_requirements()
+            )
+            # Zero out all gradients (set to None)
+            for param in self.model.parameters():
+                param.grad = None
+            # Perform a dummy forward pass (through the policy net, which should be
+            # separated from the value function in this particular user setup).
+            out = self.model(dummy_batch)
+            # Perform a (dummy) backward pass to be able to see, which params have
+            # gradients and are therefore used for the policy computations (vs vf
+            # computations).
+            torch.sum(out[0]).backward()  # [0] -> Model returns out and state-outs.
+            # Collect policy vs value function params separately.
+            policy_params = []
+            value_params = []
+            for param in self.model.parameters():
+                if param.grad is None:
+                    value_params.append(param)
+                else:
+                    policy_params.append(param)
+            if self.config["opt_type"] == "adam":
+                return (
+                    torch.optim.Adam(params=policy_params, lr=self.cur_lr),
+                    torch.optim.Adam(params=value_params, lr=self.cur_lr2),
+                )
+            else:
+                raise NotImplementedError
+
+        if self.config["opt_type"] == "adam":
+            return torch.optim.Adam(params=self.model.parameters(), lr=self.cur_lr)
+        else:
+            return torch.optim.RMSprop(
+                params=self.model.parameters(),
+                lr=self.cur_lr,
+                weight_decay=self.config["decay"],
+                momentum=self.config["momentum"],
+                eps=self.config["epsilon"],
+            )
+
+
+# VTrace mixins are placed in front of more general mixins to make sure
+# their functions like optimizer() overrides all the other implementations
+# (e.g., LearningRateSchedule.optimizer())
+class ImpalaTorchPolicy(
+    VTraceOptimizer,
+    LearningRateSchedule,
+    EntropyCoeffSchedule,
+    ValueNetworkMixin,
+    TorchPolicyV2,
+):
+    """PyTorch policy class used with IMPALA."""
+
+    def __init__(self, observation_space, action_space, config):
+        config = dict(
+            ray.rllib.algorithms.impala.impala.IMPALAConfig().to_dict(), **config
+        )
+        config["enable_rl_module_and_learner"] = False
+        config["enable_env_runner_and_connector_v2"] = False
+
+        # If Learner API is used, we don't need any loss-specific mixins.
+        # However, we also would like to avoid creating special Policy-subclasses
+        # for this as the entire Policy concept will soon not be used anymore with
+        # the new Learner- and RLModule APIs.
+        VTraceOptimizer.__init__(self)
+        # Need to initialize learning rate variable before calling
+        # TorchPolicyV2.__init__.
+        lr_schedule_additional_args = []
+        if config.get("_separate_vf_optimizer"):
+            lr_schedule_additional_args = (
+                [config["_lr_vf"][0][1], config["_lr_vf"]]
+                if isinstance(config["_lr_vf"], (list, tuple))
+                else [config["_lr_vf"], None]
+            )
+        LearningRateSchedule.__init__(
+            self, config["lr"], config["lr_schedule"], *lr_schedule_additional_args
+        )
+        EntropyCoeffSchedule.__init__(
+            self, config["entropy_coeff"], config["entropy_coeff_schedule"]
+        )
+
+        TorchPolicyV2.__init__(
+            self,
+            observation_space,
+            action_space,
+            config,
+            max_seq_len=config["model"]["max_seq_len"],
+        )
+
+        ValueNetworkMixin.__init__(self, config)
+
+        self._initialize_loss_from_dummy_batch()
+
+    @override(TorchPolicyV2)
+    def loss(
+        self,
+        model: ModelV2,
+        dist_class: Type[ActionDistribution],
+        train_batch: SampleBatch,
+    ) -> Union[TensorType, List[TensorType]]:
+        model_out, _ = model(train_batch)
+        action_dist = dist_class(model_out, model)
+
+        if isinstance(self.action_space, gym.spaces.Discrete):
+            is_multidiscrete = False
+            output_hidden_shape = [self.action_space.n]
+        elif isinstance(self.action_space, gym.spaces.MultiDiscrete):
+            is_multidiscrete = True
+            output_hidden_shape = self.action_space.nvec.astype(np.int32)
+        else:
+            is_multidiscrete = False
+            output_hidden_shape = 1
+
+        def _make_time_major(*args, **kw):
+            return make_time_major(
+                self, train_batch.get(SampleBatch.SEQ_LENS), *args, **kw
+            )
+
+        actions = train_batch[SampleBatch.ACTIONS]
+        dones = train_batch[SampleBatch.TERMINATEDS]
+        rewards = train_batch[SampleBatch.REWARDS]
+        behaviour_action_logp = train_batch[SampleBatch.ACTION_LOGP]
+        behaviour_logits = train_batch[SampleBatch.ACTION_DIST_INPUTS]
+        if isinstance(output_hidden_shape, (list, tuple, np.ndarray)):
+            unpacked_behaviour_logits = torch.split(
+                behaviour_logits, list(output_hidden_shape), dim=1
+            )
+            unpacked_outputs = torch.split(model_out, list(output_hidden_shape), dim=1)
+        else:
+            unpacked_behaviour_logits = torch.chunk(
+                behaviour_logits, output_hidden_shape, dim=1
+            )
+            unpacked_outputs = torch.chunk(model_out, output_hidden_shape, dim=1)
+        values = model.value_function()
+        values_time_major = _make_time_major(values)
+        bootstrap_values_time_major = _make_time_major(
+            train_batch[SampleBatch.VALUES_BOOTSTRAPPED]
+        )
+        bootstrap_value = bootstrap_values_time_major[-1]
+
+        if self.is_recurrent():
+            max_seq_len = torch.max(train_batch[SampleBatch.SEQ_LENS])
+            mask_orig = sequence_mask(train_batch[SampleBatch.SEQ_LENS], max_seq_len)
+            mask = torch.reshape(mask_orig, [-1])
+        else:
+            mask = torch.ones_like(rewards)
+
+        # Prepare actions for loss.
+        loss_actions = actions if is_multidiscrete else torch.unsqueeze(actions, dim=1)
+
+        # Inputs are reshaped from [B * T] => [(T|T-1), B] for V-trace calc.
+        loss = VTraceLoss(
+            actions=_make_time_major(loss_actions),
+            actions_logp=_make_time_major(action_dist.logp(actions)),
+            actions_entropy=_make_time_major(action_dist.entropy()),
+            dones=_make_time_major(dones),
+            behaviour_action_logp=_make_time_major(behaviour_action_logp),
+            behaviour_logits=_make_time_major(unpacked_behaviour_logits),
+            target_logits=_make_time_major(unpacked_outputs),
+            discount=self.config["gamma"],
+            rewards=_make_time_major(rewards),
+            values=values_time_major,
+            bootstrap_value=bootstrap_value,
+            dist_class=TorchCategorical if is_multidiscrete else dist_class,
+            model=model,
+            valid_mask=_make_time_major(mask),
+            config=self.config,
+            vf_loss_coeff=self.config["vf_loss_coeff"],
+            entropy_coeff=self.entropy_coeff,
+            clip_rho_threshold=self.config["vtrace_clip_rho_threshold"],
+            clip_pg_rho_threshold=self.config["vtrace_clip_pg_rho_threshold"],
+        )
+
+        # Store values for stats function in model (tower), such that for
+        # multi-GPU, we do not override them during the parallel loss phase.
+        model.tower_stats["pi_loss"] = loss.pi_loss
+        model.tower_stats["vf_loss"] = loss.vf_loss
+        model.tower_stats["entropy"] = loss.entropy
+        model.tower_stats["mean_entropy"] = loss.mean_entropy
+        model.tower_stats["total_loss"] = loss.total_loss
+
+        values_batched = make_time_major(
+            self,
+            train_batch.get(SampleBatch.SEQ_LENS),
+            values,
+        )
+        model.tower_stats["vf_explained_var"] = explained_variance(
+            torch.reshape(loss.value_targets, [-1]), torch.reshape(values_batched, [-1])
+        )
+
+        if self.config.get("_separate_vf_optimizer"):
+            return loss.loss_wo_vf, loss.vf_loss
+        else:
+            return loss.total_loss
+
+    @override(TorchPolicyV2)
+    def stats_fn(self, train_batch: SampleBatch) -> Dict[str, TensorType]:
+        return convert_to_numpy(
+            {
+                "cur_lr": self.cur_lr,
+                "total_loss": torch.mean(
+                    torch.stack(self.get_tower_stats("total_loss"))
+                ),
+                "policy_loss": torch.mean(torch.stack(self.get_tower_stats("pi_loss"))),
+                "entropy": torch.mean(
+                    torch.stack(self.get_tower_stats("mean_entropy"))
+                ),
+                "entropy_coeff": self.entropy_coeff,
+                "var_gnorm": global_norm(self.model.trainable_variables()),
+                "vf_loss": torch.mean(torch.stack(self.get_tower_stats("vf_loss"))),
+                "vf_explained_var": torch.mean(
+                    torch.stack(self.get_tower_stats("vf_explained_var"))
+                ),
+            }
+        )
+
+    @override(TorchPolicyV2)
+    def postprocess_trajectory(
+        self,
+        sample_batch: SampleBatch,
+        other_agent_batches: Optional[SampleBatch] = None,
+        episode=None,
+    ):
+        # Call super's postprocess_trajectory first.
+        # sample_batch = super().postprocess_trajectory(
+        #    sample_batch, other_agent_batches, episode
+        # )
+
+        if self.config["vtrace"]:
+            # Add the SampleBatch.VALUES_BOOTSTRAPPED column, which we'll need
+            # inside the loss for vtrace calculations.
+            sample_batch = compute_bootstrap_value(sample_batch, self)
+
+        return sample_batch
+
+    @override(TorchPolicyV2)
+    def extra_grad_process(
+        self, optimizer: "torch.optim.Optimizer", loss: TensorType
+    ) -> Dict[str, TensorType]:
+        return apply_grad_clipping(self, optimizer, loss)
+
+    @override(TorchPolicyV2)
+    def get_batch_divisibility_req(self) -> int:
+        return self.config["rollout_fragment_length"]

rllib/algorithms/impala/tests/test_impala.py

@@ -0,0 +1,109 @@
+import unittest
+
+import pytest
+
+import ray
+import ray.rllib.algorithms.impala as impala
+from ray.rllib.policy.sample_batch import DEFAULT_POLICY_ID
+from ray.rllib.utils.metrics import LEARNER_RESULTS
+from ray.rllib.utils.test_utils import check
+
+
+class TestIMPALA(unittest.TestCase):
+    @classmethod
+    def setUpClass(cls) -> None:
+        ray.init()
+
+    @classmethod
+    def tearDownClass(cls) -> None:
+        ray.shutdown()
+
+    def test_impala_minibatch_size_check(self):
+        config = (
+            impala.IMPALAConfig()
+            .environment("CartPole-v1")
+            .training(minibatch_size=100)
+            .env_runners(rollout_fragment_length=30)
+        )
+
+        with pytest.raises(
+            ValueError,
+            match=r"`minibatch_size` \(100\) must either be None or a multiple of `rollout_fragment_length` \(30\)",
+        ):
+            config.validate()
+
+    def test_impala_lr_schedule(self):
+        # Test whether we correctly ignore the "lr" setting.
+        # The first lr should be 0.05.
+        config = (
+            impala.IMPALAConfig()
+            .learners(num_learners=0)
+            .experimental(_validate_config=False)  #
+            .training(
+                lr=[
+                    [0, 0.05],
+                    [100000, 0.000001],
+                ],
+                train_batch_size=100,
+            )
+            .env_runners(num_envs_per_env_runner=2)
+            .environment(env="CartPole-v1")
+        )
+
+        def get_lr(result):
+            return result[LEARNER_RESULTS][DEFAULT_POLICY_ID][
+                "default_optimizer_learning_rate"
+            ]
+
+        algo = config.build()
+        optim = algo.learner_group._learner.get_optimizer()
+
+        try:
+            check(optim.param_groups[0]["lr"], 0.05)
+            for _ in range(1):
+                r1 = algo.train()
+            for _ in range(2):
+                r2 = algo.train()
+            for _ in range(2):
+                r3 = algo.train()
+            # Due to the asynch'ness of IMPALA, learner-stats metrics
+            # could be delayed by one iteration. Do 3 train() calls here
+            # and measure guaranteed decrease in lr between 1st and 3rd.
+            lr1 = get_lr(r1)
+            lr2 = get_lr(r2)
+            lr3 = get_lr(r3)
+            assert lr2 <= lr1, (lr1, lr2)
+            assert lr3 <= lr2, (lr2, lr3)
+            assert lr3 < lr1, (lr1, lr3)
+        finally:
+            algo.stop()
+
+    def test_local_learner_thread_stops_on_algo_stop(self):
+        # Regression test: `algo.stop()` -> `LearnerGroup.shutdown()` ->
+        # `IMPALALearner.shutdown()` must stop and join the local IMPALA
+        # `_LearnerThread`. Otherwise the daemon thread keeps spinning and
+        # can race against interpreter shutdown inside an auto_init-wrapped
+        # Ray API.
+        config = (
+            impala.IMPALAConfig()
+            .environment("CartPole-v1")
+            .learners(num_learners=0)
+            .env_runners(num_env_runners=0)
+        )
+        algo = config.build()
+        learner_thread = algo.learner_group._learner._learner_thread
+        self.assertTrue(learner_thread.is_alive())
+
+        algo.stop()
+
+        # `Learner.shutdown()` joins the thread, so it must be dead by the
+        # time `algo.stop()` returns — no extra `join()` needed here.
+        self.assertFalse(learner_thread.is_alive())
+
+
+if __name__ == "__main__":
+    import sys
+
+    import pytest
+
+    sys.exit(pytest.main(["-v", __file__]))

rllib/algorithms/impala/tests/test_vtrace_old_api_stack.py

@@ -0,0 +1,290 @@
+# Copyright 2018 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Tests for V-trace.
+
+For details and theory see:
+
+"IMPALA: Scalable Distributed Deep-RL with
+Importance Weighted Actor-Learner Architectures"
+by Espeholt, Soyer, Munos et al.
+"""
+
+import unittest
+
+import numpy as np
+from gymnasium.spaces import Box
+
+from ray.rllib.algorithms.impala import vtrace_torch as vtrace_torch
+from ray.rllib.utils.framework import try_import_torch
+from ray.rllib.utils.numpy import softmax
+from ray.rllib.utils.test_utils import check
+
+torch, nn = try_import_torch()
+
+
+def _ground_truth_vtrace_calculation(
+    discounts: np.ndarray,
+    log_rhos: np.ndarray,
+    rewards: np.ndarray,
+    values: np.ndarray,
+    bootstrap_value: np.ndarray,
+    clip_rho_threshold: float,
+    clip_pg_rho_threshold: float,
+):
+    """Calculates the ground truth for V-trace in Python/Numpy.
+
+    NOTE:
+    The discount, log_rhos, rewards, values, and bootstrap_values are all assumed to
+    come from trajectories of experience. Typically batches of trajectories could be
+    thought of as having the shape [B, T] where B is the batch dimension, and T is
+    the timestep dimension. Computing vtrace returns requires that the data is time
+    major, meaning that it has the shape [T, B]. One can use a function like
+    `make_time_major` to properly format their discount, log_rhos, rewards, values,
+    and bootstrap_values before calling _ground_truth_vtrace_calculation.
+
+    Args:
+        discounts: Array of shape [T*B] of discounts. T is the length of the trajectory
+            or sequence and B is the batch size.
+            NOTE: The discount will be equal to gamma, the discount factor when the
+            timestep that the discount is being applied to is not a terminal timestep.
+        log_rhos: Array of shape [T*B] of log likelihood ratios of target action
+            probabilities to behavior action probabilities.
+        rewards: Array of shape [T*B] of rewards.
+        values: Array of shape [T*B] of the value function estimated for every timestep
+            in a batch.
+        bootstrap_values: Array of shape [B] of the value function estimated at the last
+            timestep for each trajectory in the batch.
+        clip_rho_threshold: The threshold for clipping the importance weights.
+        clip_pg_rho_threshold: The threshold for clipping the importance weights for
+            the policy gradient loss.
+
+    Returns:
+        The v-trace adjusted values and the policy gradient advantages.
+
+    """
+    vs = []
+    seq_len = len(discounts)
+    rhos = np.exp(log_rhos)
+    cs = np.minimum(rhos, 1.0)
+    clipped_rhos = rhos
+    if clip_rho_threshold:
+        clipped_rhos = np.minimum(rhos, clip_rho_threshold)
+    clipped_pg_rhos = rhos
+    if clip_pg_rho_threshold:
+        clipped_pg_rhos = np.minimum(rhos, clip_pg_rho_threshold)
+
+    # This is a very inefficient way to calculate the V-trace ground truth.
+    # We calculate it this way because it is close to the mathematical notation
+    # of
+    # V-trace.
+    # v_s = V(x_s)
+    #       + \sum^{T-1}_{t=s} \gamma^{t-s}
+    #         * \prod_{i=s}^{t-1} c_i
+    #         * \rho_t (r_t + \gamma V(x_{t+1}) - V(x_t))
+    # Note that when we take the product over c_i, we write `s:t` as the
+    # notation
+    # of the paper is inclusive of the `t-1`, but Python is exclusive.
+    # Also note that np.prod([]) == 1.
+    values_t_plus_1 = np.concatenate([values[1:], bootstrap_value[None, :]], axis=0)
+    for s in range(seq_len):
+        v_s = np.copy(values[s])  # Very important copy.
+        for t in range(s, seq_len):
+            v_s += (
+                np.prod(discounts[s:t], axis=0)
+                * np.prod(cs[s:t], axis=0)
+                * clipped_rhos[t]
+                * (rewards[t] + discounts[t] * values_t_plus_1[t] - values[t])
+            )
+        vs.append(v_s)
+    vs = np.stack(vs, axis=0)
+    pg_advantages = clipped_pg_rhos * (
+        rewards
+        + discounts * np.concatenate([vs[1:], bootstrap_value[None, :]], axis=0)
+        - values
+    )
+    return vs, pg_advantages
+
+
+class LogProbsFromLogitsAndActionsTest(unittest.TestCase):
+    def test_log_probs_from_logits_and_actions(self):
+        """Tests log_probs_from_logits_and_actions."""
+        seq_len = 7
+        num_actions = 3
+        batch_size = 4
+
+        vtrace = vtrace_torch
+        policy_logits = Box(
+            -1.0, 1.0, (seq_len, batch_size, num_actions), np.float32
+        ).sample()
+        actions = np.random.randint(
+            0, num_actions - 1, size=(seq_len, batch_size), dtype=np.int32
+        )
+
+        action_log_probs_tensor = vtrace.log_probs_from_logits_and_actions(
+            torch.from_numpy(policy_logits), torch.from_numpy(actions)
+        )
+
+        # Ground Truth
+        # Using broadcasting to create a mask that indexes action logits
+        action_index_mask = actions[..., None] == np.arange(num_actions)
+
+        def index_with_mask(array, mask):
+            return array[mask].reshape(*array.shape[:-1])
+
+        # Note: Normally log(softmax) is not a good idea because it's not
+        # numerically stable. However, in this test we have well-behaved
+        # values.
+        ground_truth_v = index_with_mask(
+            np.log(softmax(policy_logits)), action_index_mask
+        )
+
+        check(action_log_probs_tensor, ground_truth_v)
+
+
+class VtraceTest(unittest.TestCase):
+    def test_vtrace(self):
+        """Tests V-trace against ground truth data calculated in python."""
+        seq_len = 5
+        batch_size = 10
+
+        # Create log_rhos such that rho will span from near-zero to above the
+        # clipping thresholds. In particular, calculate log_rhos in
+        # [-2.5, 2.5),
+        # so that rho is in approx [0.08, 12.2).
+        space_w_time = Box(-1.0, 1.0, (seq_len, batch_size), np.float32)
+        space_only_batch = Box(-1.0, 1.0, (batch_size,), np.float32)
+        log_rhos = space_w_time.sample() / (batch_size * seq_len)
+        log_rhos = 5 * (log_rhos - 0.5)  # [0.0, 1.0) -> [-2.5, 2.5).
+        values = {
+            "log_rhos": log_rhos,
+            # T, B where B_i: [0.9 / (i+1)] * T
+            "discounts": np.array(
+                [[0.9 / (b + 1) for b in range(batch_size)] for _ in range(seq_len)]
+            ),
+            "rewards": space_w_time.sample(),
+            "values": space_w_time.sample() / batch_size,
+            "bootstrap_value": space_only_batch.sample() + 1.0,
+            "clip_rho_threshold": 3.7,
+            "clip_pg_rho_threshold": 2.2,
+        }
+
+        vtrace = vtrace_torch
+        output = vtrace.from_importance_weights(**values)
+
+        gt_vs, gt_pg_advantags = _ground_truth_vtrace_calculation(**values)
+        check(output.vs, gt_vs)
+        check(output.pg_advantages, gt_pg_advantags)
+
+    def test_vtrace_from_logits(self):
+        """Tests V-trace calculated from logits."""
+        seq_len = 5
+        batch_size = 15
+        num_actions = 3
+        clip_rho_threshold = None  # No clipping.
+        clip_pg_rho_threshold = None  # No clipping.
+        space = Box(-1.0, 1.0, (seq_len, batch_size, num_actions))
+        action_space = Box(
+            0,
+            num_actions - 1,
+            (
+                seq_len,
+                batch_size,
+            ),
+            dtype=np.int32,
+        )
+        space_w_time = Box(
+            -1.0,
+            1.0,
+            (
+                seq_len,
+                batch_size,
+            ),
+        )
+        space_only_batch = Box(-1.0, 1.0, (batch_size,))
+
+        inputs_ = {
+            # T, B, NUM_ACTIONS
+            "behaviour_policy_logits": space.sample(),
+            # T, B, NUM_ACTIONS
+            "target_policy_logits": space.sample(),
+            "actions": action_space.sample(),
+            "discounts": space_w_time.sample(),
+            "rewards": space_w_time.sample(),
+            "values": space_w_time.sample(),
+            "bootstrap_value": space_only_batch.sample(),
+        }
+        from_logits_output = vtrace_torch.from_logits(
+            clip_rho_threshold=clip_rho_threshold,
+            clip_pg_rho_threshold=clip_pg_rho_threshold,
+            **inputs_
+        )
+
+        target_log_probs = vtrace_torch.log_probs_from_logits_and_actions(
+            torch.from_numpy(inputs_["target_policy_logits"]),
+            torch.from_numpy(inputs_["actions"]),
+        )
+        behaviour_log_probs = vtrace_torch.log_probs_from_logits_and_actions(
+            torch.from_numpy(inputs_["behaviour_policy_logits"]),
+            torch.from_numpy(inputs_["actions"]),
+        )
+        log_rhos = target_log_probs - behaviour_log_probs
+
+        from_iw = vtrace_torch.from_importance_weights(
+            log_rhos=log_rhos,
+            discounts=inputs_["discounts"],
+            rewards=inputs_["rewards"],
+            values=inputs_["values"],
+            bootstrap_value=inputs_["bootstrap_value"],
+            clip_rho_threshold=clip_rho_threshold,
+            clip_pg_rho_threshold=clip_pg_rho_threshold,
+        )
+
+        check(from_iw.vs, from_logits_output.vs)
+        check(from_iw.pg_advantages, from_logits_output.pg_advantages)
+        check(behaviour_log_probs, from_logits_output.behaviour_action_log_probs)
+        check(target_log_probs, from_logits_output.target_action_log_probs)
+        check(log_rhos, from_logits_output.log_rhos)
+
+    def test_higher_rank_inputs_for_importance_weights(self):
+        """Checks support for additional dimensions in inputs."""
+        inputs_ = {
+            "log_rhos": Box(-1.0, 1.0, (8, 10, 1)).sample(),
+            "discounts": Box(-1.0, 1.0, (8, 10, 1)).sample(),
+            "rewards": Box(-1.0, 1.0, (8, 10, 42)).sample(),
+            "values": Box(-1.0, 1.0, (8, 10, 42)).sample(),
+            "bootstrap_value": Box(-1.0, 1.0, (10, 42)).sample(),
+        }
+        output = vtrace_torch.from_importance_weights(**inputs_)
+        check(int(output.vs.shape[-1]), 42)
+
+    def test_inconsistent_rank_inputs_for_importance_weights(self):
+        """Test one of many possible errors in shape of inputs."""
+        inputs_ = {
+            "log_rhos": Box(-1.0, 1.0, (7, 15, 1)).sample(),
+            "discounts": Box(-1.0, 1.0, (7, 15, 1)).sample(),
+            "rewards": Box(-1.0, 1.0, (7, 15, 42)).sample(),
+            "values": Box(-1.0, 1.0, (7, 15, 42)).sample(),
+            # Should be [15, 42].
+            "bootstrap_value": Box(-1.0, 1.0, (7,)).sample(),
+        }
+        with self.assertRaisesRegex((ValueError, AssertionError), "must have rank 2"):
+            vtrace_torch.from_importance_weights(**inputs_)
+
+
+if __name__ == "__main__":
+    import sys
+
+    import pytest
+
+    sys.exit(pytest.main(["-v", __file__]))

rllib/algorithms/impala/tests/test_vtrace_v2.py

@@ -0,0 +1,154 @@
+import unittest
+
+import numpy as np
+from gymnasium.spaces import Box, Discrete
+
+from ray.rllib.algorithms.impala.tests.test_vtrace_old_api_stack import (
+    _ground_truth_vtrace_calculation,
+)
+from ray.rllib.algorithms.impala.torch.vtrace_torch_v2 import (
+    make_time_major,
+    vtrace_torch,
+)
+from ray.rllib.core.distribution.torch.torch_distribution import TorchCategorical
+from ray.rllib.utils.framework import try_import_torch
+from ray.rllib.utils.test_utils import check
+from ray.rllib.utils.torch_utils import convert_to_torch_tensor
+
+torch, _ = try_import_torch()
+
+
+def flatten_batch_and_time_dim(t):
+    if not torch.is_tensor(t):
+        t = torch.from_numpy(t)
+    new_shape = [-1] + list(t.shape[2:])
+    return torch.reshape(t, new_shape)
+
+
+class TestVtraceRLModule(unittest.TestCase):
+    """Tests V-trace-v2 against ground truth data calculated.
+
+    There is a ground truth implementation that we used to test our original
+    implementation against. This test checks that the new implementation still
+    matches the ground truth test from our first implementation of V-Trace.
+    """
+
+    @classmethod
+    def setUpClass(cls):
+        """Sets up inputs for V-Trace and calculate ground truth.
+
+        We use tf operations here to compile the inputs but convert to numpy arrays to
+        calculate the ground truth (and the other v-trace outputs in the
+        framework-specific tests).
+        """
+
+        # we can test against any trajectory length or batch size and it won't matter
+        trajectory_len = 5
+        batch_size = 10
+
+        action_space = Discrete(10)
+        action_logit_space = Box(-1.0, 1.0, (action_space.n,), np.float32)
+        behavior_action_logits = torch.from_numpy(
+            np.array([action_logit_space.sample()], dtype=np.float32)
+        )
+        target_action_logits = torch.from_numpy(
+            np.array([action_logit_space.sample()], dtype=np.float32)
+        )
+        behavior_dist = TorchCategorical(logits=behavior_action_logits)
+        target_dist = TorchCategorical(logits=target_action_logits)
+        dummy_action_batch = [
+            [action_space.sample() for _ in range(trajectory_len)]
+            for _ in range(batch_size)
+        ]
+
+        behavior_log_probs = torch.stack(
+            [
+                torch.squeeze(
+                    behavior_dist.logp(torch.from_numpy(np.array(v, np.uint8)))
+                )
+                for v in dummy_action_batch
+            ]
+        )
+        target_log_probs = torch.stack(
+            [
+                torch.squeeze(target_dist.logp(torch.from_numpy(np.array(v, np.uint8))))
+                for v in dummy_action_batch
+            ]
+        )
+        # target_log_probs = torch.stack(
+        #    tree.map_structure(
+        #        lambda v: torch.squeeze(target_dist.logp(v)), dummy_action_batch
+        #    )
+        # )
+
+        value_fn_space_w_time = Box(-1.0, 1.0, (batch_size, trajectory_len), np.float32)
+        value_fn_space = Box(-1.0, 1.0, (batch_size,), np.float32)
+
+        # using randomly sampled values in lieu of actual values sampled from a value fn
+        values = value_fn_space_w_time.sample()
+        # this is supposed to be the value function at the last timestep of each
+        # trajectory in the batch. In IMPALA its bootstrapped at training time
+        cls.bootstrap_values = np.array(value_fn_space.sample() + 1.0)
+
+        # discount factor used at all of the timesteps
+        discounts = torch.from_numpy(
+            np.array([0.9 for _ in range(trajectory_len * batch_size)])
+        )
+        rewards = value_fn_space_w_time.sample()
+        cls.clip_rho_threshold = 3.7
+        cls.clip_pg_rho_threshold = 2.2
+
+        # convert to time major dimension
+        cls.behavior_log_probs_time_major = make_time_major(
+            flatten_batch_and_time_dim(behavior_log_probs),
+            trajectory_len=trajectory_len,
+        ).numpy()
+        cls.target_log_probs_time_major = make_time_major(
+            flatten_batch_and_time_dim(target_log_probs), trajectory_len=trajectory_len
+        ).numpy()
+        cls.discounts_time_major = make_time_major(
+            flatten_batch_and_time_dim(discounts), trajectory_len=trajectory_len
+        ).numpy()
+        cls.rewards_time_major = make_time_major(
+            flatten_batch_and_time_dim(rewards), trajectory_len=trajectory_len
+        ).numpy()
+        cls.values_time_major = make_time_major(
+            flatten_batch_and_time_dim(values), trajectory_len=trajectory_len
+        ).numpy()
+
+        log_rhos = cls.target_log_probs_time_major - cls.behavior_log_probs_time_major
+
+        cls.ground_truth_v = _ground_truth_vtrace_calculation(
+            discounts=cls.discounts_time_major,
+            log_rhos=log_rhos,
+            rewards=cls.rewards_time_major,
+            values=cls.values_time_major,
+            bootstrap_value=cls.bootstrap_values,
+            clip_rho_threshold=cls.clip_rho_threshold,
+            clip_pg_rho_threshold=cls.clip_pg_rho_threshold,
+        )
+
+    def test_vtrace_torch(self):
+        output_torch_vtrace = vtrace_torch(
+            behaviour_action_log_probs=convert_to_torch_tensor(
+                self.behavior_log_probs_time_major
+            ),
+            target_action_log_probs=convert_to_torch_tensor(
+                self.target_log_probs_time_major
+            ),
+            discounts=convert_to_torch_tensor(self.discounts_time_major),
+            rewards=convert_to_torch_tensor(self.rewards_time_major),
+            values=convert_to_torch_tensor(self.values_time_major),
+            bootstrap_values=convert_to_torch_tensor(self.bootstrap_values),
+            clip_rho_threshold=self.clip_rho_threshold,
+            clip_pg_rho_threshold=self.clip_pg_rho_threshold,
+        )
+        check(output_torch_vtrace, self.ground_truth_v)
+
+
+if __name__ == "__main__":
+    import sys
+
+    import pytest
+
+    sys.exit(pytest.main(["-v", __file__]))

rllib/algorithms/impala/torch/impala_torch_learner.py

@@ -0,0 +1,353 @@
+import contextlib
+from typing import Dict
+
+from ray.rllib.algorithms.impala.impala import IMPALAConfig
+from ray.rllib.algorithms.impala.impala_learner import IMPALALearner
+from ray.rllib.algorithms.impala.torch.vtrace_torch_v2 import (
+    make_time_major,
+    vtrace_torch,
+)
+from ray.rllib.core.columns import Columns
+from ray.rllib.core.learner.learner import ENTROPY_KEY
+from ray.rllib.core.learner.torch.torch_learner import TorchLearner
+from ray.rllib.utils.annotations import override
+from ray.rllib.utils.framework import try_import_torch
+from ray.rllib.utils.typing import ModuleID, ParamDict, TensorType
+
+torch, nn = try_import_torch()
+
+
+class IMPALATorchLearner(IMPALALearner, TorchLearner):
+    """Implements the IMPALA loss function in torch."""
+
+    @override(TorchLearner)
+    def compute_loss_for_module(
+        self,
+        *,
+        module_id: ModuleID,
+        config: IMPALAConfig,
+        batch: Dict,
+        fwd_out: Dict[str, TensorType],
+    ) -> TensorType:
+        module = self.module[module_id].unwrapped()
+
+        # TODO (sven): Now that we do the +1ts trick to be less vulnerable about
+        #  bootstrap values at the end of rollouts in the new stack, we might make
+        #  this a more flexible, configurable parameter for users, e.g.
+        #  `v_trace_seq_len` (independent of `rollout_fragment_length`). Separation
+        #  of concerns (sampling vs learning).
+        rollout_frag_or_episode_len = config.get_rollout_fragment_length()
+        recurrent_seq_len = batch.get("seq_lens")
+
+        loss_mask = batch[Columns.LOSS_MASK].float()
+        loss_mask_time_major = make_time_major(
+            loss_mask,
+            trajectory_len=rollout_frag_or_episode_len,
+            recurrent_seq_len=recurrent_seq_len,
+        )
+        size_loss_mask = torch.sum(loss_mask)
+
+        # Behavior actions logp and target actions logp.
+        behaviour_actions_logp = batch[Columns.ACTION_LOGP]
+        target_policy_dist = module.get_train_action_dist_cls().from_logits(
+            fwd_out[Columns.ACTION_DIST_INPUTS]
+        )
+        target_actions_logp = target_policy_dist.logp(batch[Columns.ACTIONS])
+
+        # Values and bootstrap values.
+        values = module.compute_values(
+            batch, embeddings=fwd_out.get(Columns.EMBEDDINGS)
+        )
+        values_time_major = make_time_major(
+            values,
+            trajectory_len=rollout_frag_or_episode_len,
+            recurrent_seq_len=recurrent_seq_len,
+        )
+        assert Columns.VALUES_BOOTSTRAPPED not in batch
+        # Use as bootstrap values the vf-preds in the next "batch row", except
+        # for the very last row (which doesn't have a next row), for which the
+        # bootstrap value does not matter b/c it has a +1ts value at its end
+        # anyways. So we chose an arbitrary item (for simplicity of not having to
+        # move new data to the device).
+        bootstrap_values = torch.cat(
+            [
+                values_time_major[0][1:],  # 0th ts values from "next row"
+                values_time_major[0][0:1],  # <- can use any arbitrary value here
+            ],
+            dim=0,
+        )
+
+        # TODO(Artur): In the old impala code, actions were unsqueezed if they were
+        #  multi_discrete. Find out why and if we need to do the same here.
+        #  actions = actions if is_multidiscrete else torch.unsqueeze(actions, dim=1)
+        target_actions_logp_time_major = make_time_major(
+            target_actions_logp,
+            trajectory_len=rollout_frag_or_episode_len,
+            recurrent_seq_len=recurrent_seq_len,
+        )
+        behaviour_actions_logp_time_major = make_time_major(
+            behaviour_actions_logp,
+            trajectory_len=rollout_frag_or_episode_len,
+            recurrent_seq_len=recurrent_seq_len,
+        )
+        rewards_time_major = make_time_major(
+            batch[Columns.REWARDS],
+            trajectory_len=rollout_frag_or_episode_len,
+            recurrent_seq_len=recurrent_seq_len,
+        )
+
+        # Discount = gamma * (1 - terminated) * loss_mask.
+        # - The (1 - terminated) factor implements the Bellman gating: no
+        #   bootstrap from t -> t+1 across a terminal step.
+        # - The loss_mask factor zeros out the discount at the appended bootstrap
+        #   timestep (loss_mask=False there). Without it, the bootstrap-ts delta
+        #   (which references `bootstrap_values` from a neighbouring trajectory)
+        #   would leak into the V-trace recursion of the last real step. The
+        #   loss_mask gating is equivalent to the legacy convention of marking
+        #   the bootstrap ts as `terminated=True`, but keeps `terminateds`
+        #   meaning only "Gymnasium terminal state reached".
+        discounts_time_major = (
+            (
+                1.0
+                - make_time_major(
+                    batch[Columns.TERMINATEDS],
+                    trajectory_len=rollout_frag_or_episode_len,
+                    recurrent_seq_len=recurrent_seq_len,
+                ).type(dtype=torch.float32)
+            )
+            * config.gamma
+            * loss_mask_time_major
+        )
+
+        # Note that vtrace will compute the main loop on the CPU for better performance.
+        vtrace_adjusted_target_values, pg_advantages = vtrace_torch(
+            target_action_log_probs=target_actions_logp_time_major,
+            behaviour_action_log_probs=behaviour_actions_logp_time_major,
+            discounts=discounts_time_major,
+            rewards=rewards_time_major,
+            values=values_time_major,
+            bootstrap_values=bootstrap_values,
+            clip_rho_threshold=config.vtrace_clip_rho_threshold,
+            clip_pg_rho_threshold=config.vtrace_clip_pg_rho_threshold,
+        )
+
+        # The policy gradients loss.
+        pi_loss = -torch.sum(
+            target_actions_logp_time_major * pg_advantages * loss_mask_time_major
+        )
+        mean_pi_loss = pi_loss / size_loss_mask
+
+        # The baseline loss.
+        delta = values_time_major - vtrace_adjusted_target_values
+        vf_loss = 0.5 * torch.sum(torch.pow(delta, 2.0) * loss_mask_time_major)
+        mean_vf_loss = vf_loss / size_loss_mask
+
+        # The entropy loss.
+        entropy_loss = -torch.sum(target_policy_dist.entropy() * loss_mask)
+        mean_entropy_loss = entropy_loss / size_loss_mask
+
+        # The summed weighted loss.
+        total_loss = (
+            mean_pi_loss
+            + mean_vf_loss * config.vf_loss_coeff
+            + (
+                mean_entropy_loss
+                * self.entropy_coeff_schedulers_per_module[
+                    module_id
+                ].get_current_value()
+            )
+        )
+
+        # Log important loss stats.
+        self.metrics.log_dict(
+            {
+                "pi_loss": pi_loss,
+                "mean_pi_loss": mean_pi_loss,
+                "vf_loss": vf_loss,
+                "mean_vf_loss": mean_vf_loss,
+                ENTROPY_KEY: -mean_entropy_loss,
+            },
+            key=module_id,
+            window=1,  # <- single items (should not be mean/ema-reduced over time).
+        )
+        # Return the total loss.
+        return total_loss
+
+    @override(TorchLearner)
+    def _uncompiled_update(
+        self,
+        batch: Dict,
+        **kwargs,
+    ):
+        """Performs a single update given a batch of data.
+
+        This override is critical to prevent DDP (DistributedDataParallel)
+        deadlocks caused by the unique properties of APPO's asynchronous,
+        multi-agent data pipeline.
+
+        **The Problem: Asymmetric Graph Deadlock**
+        1.  APPO's asynchronous `EnvRunners` send data to each Learner (DDP rank)
+            independently.
+        2.  This means that ranks will receive sometimes **asymmetric batches**
+            (e.g., Rank 0 gets `{'p0', 'p1'}`, while Rank 1 gets just `{'p0'}`).
+        3.  The default DDP update path (using automatic, hook-based
+            synchronization) fails in this scenario. When `backward()` is
+            called, Rank 1 has no computation graph for `p1`'s parameters,
+            so its DDP hooks for `p1` never fire.
+        4.  Rank 0, which *does* have a graph for `p1`, waits forever for
+            `p1` gradients from Rank 1, causing a **permanent deadlock**.
+
+        The solution is manual synchronization. This function replaces DDP's
+        fragile, hook-based communication with a robust, manual, three-stage process:
+
+        1.  **Disable Hooks (The `no_sync` context):**
+            The *entire* `forward_train`, `compute_losses`, and
+            `compute_gradients` (which calls `backward()`) chain is wrapped
+            in a `mod.no_sync()` context. This is the most critical step, as
+            DDP hooks are attached during the *forward pass*. This correctly
+            prevents DDP's automatic communication from firing.
+
+        2.  **Synchronize `backward()` (avoid GPU race condition):**
+            A call to `torch.cuda.synchronize()` is added *after*
+            `total_loss.backward()` (inside `compute_gradients`). On GPU,
+            `backward()` is asynchronous. This `synchronize()` call forces
+            the CPU to wait for the GPU to *actually finish* computing the
+            gradients before we proceed to the next step. This prevents a
+            race condition where we try to `all_reduce` a `.grad` attribute
+            that is still `None` because the GPU is lagging.
+
+        3.  **Manual `all_reduce` (zero-padding):**
+            After the `no_sync` block, we manually `all_reduce` all
+            gradients. To prevent a deadlock here, we iterate over all
+            parameters. If `param.grad is None` (which happens on the
+            "incomplete" rank for policy `p1`), we zero-pad. This "zero-padding"
+            ensures that *all* ranks participate in the `all_reduce` call for
+            *all* parameters, making the manual synchronization robust.
+        """
+        # For single-learner setups, call the super's update.
+        if self.config.num_learners < 2 or not self.config.is_multi_agent:
+            return super()._uncompiled_update(batch=batch, **kwargs)
+
+        # Compute the off-policyness of the batch.
+        self._compute_off_policyness(batch)
+
+        # These must be defined outside the scope of the `with` block
+        fwd_out = {}
+        loss_per_module = {}
+        gradients = {}
+
+        # 1. Enter no_sync() for the forward and backward pass.
+        with contextlib.ExitStack() as stack:
+            for mod in self.module.values():
+                if isinstance(mod, torch.nn.Module):
+                    stack.enter_context(mod.no_sync())
+
+            # All Torch DDP-affected computation must be inside to avoid firing
+            # the hooks on policy gradients that are only trained on some ranks.
+            # 2. Forward pass (now inside no_sync)
+            fwd_out = self.module.forward_train(batch)
+
+            # 3. Loss computation (now inside no_sync)
+            loss_per_module = self.compute_losses(fwd_out=fwd_out, batch=batch)
+
+            # 4. Compute gradients LOCALLY (backward() is called inside)
+            gradients = self.compute_gradients(loss_per_module)
+
+        # 5. Manually All-Reduce gradients (outside no_sync).
+        # We iterate over all known parameters (`self._params`). This is important
+        # to ensure the `all_reduce`` calls are made on all ranks for all params.
+        for param in self._params.values():
+            # Is the parameter present on this rank?
+            present = 1
+            if param.grad is None:
+                # Parameter is not present on this rank. Keep track of that
+                # for averaging later.
+                present = 0
+                # This parameter was not used (e.g., p1 on Rank 0).
+                # Create a zero-gradient to participate in the all_reduce.
+                param.grad = torch.zeros_like(param)
+            # Now, all ranks have a valid `param.grad` tensor, i.e. all ranks will call
+            # all_reduce. No deadlock.
+            torch.distributed.all_reduce(param.grad, op=torch.distributed.ReduceOp.SUM)
+            # Scale the gradients accordingly.
+            denom = torch.tensor(present, device=param.device, dtype=param.dtype)
+            # Receive the number of participating ranks for this param.
+            torch.distributed.all_reduce(denom, op=torch.distributed.ReduceOp.SUM)
+            # Average the summed gradients.
+            param.grad.data.div_(denom.clamp(min=1.0))
+
+        # 6. Collect the gradients for all modules to update the modules synchronously.
+        gradients = {pid: p.grad for pid, p in self._params.items()}
+
+        # 7. Postprocess gradients, e.g. clipping.
+        postprocessed_gradients = self.postprocess_gradients(gradients)
+
+        # 8. Apply the post-processed gradients to the weigths.
+        self.apply_gradients(postprocessed_gradients)
+
+        return fwd_out, loss_per_module, {}
+
+    @override(TorchLearner)
+    def compute_gradients(
+        self, loss_per_module: Dict[ModuleID, TensorType], **kwargs
+    ) -> ParamDict:
+        """Computes the gradients by running `backward`.
+
+        This method is a core part of the manual synchronization logic
+        in `_uncompiled_update`. It performs the `backward()` pass locally
+        without DDP's automatic communication.
+
+        Key implementation details:
+        1.**GPU Synchronization:** It includes a `torch.cuda.synchronize()`
+            call after `total_loss.backward()`. This is critical for GPU
+            training, as `backward()` is asynchronous. This sync
+            prevents a race condition where the calling function
+            (`_uncompiled_update`) would try to access `param.grad`
+            before the GPU has finished computing it (mistakenly reading `None`).
+        2.  **Asymmetric Gradients:** On an incomplete batch, parameters for
+            missing modules will correctly have a `None` gradient. This is
+            expected and handled by the zero-padding logic in
+            `_uncompiled_update`'s `all_reduce` loop.
+        3. If the loss is zero (i.e., no modules were trained on this rank),
+            this method returns an empty dict.
+        """
+        # If a single learner is used, fall back to the super's method.
+        if self.config.num_learners < 2 or not self.config.is_multi_agent:
+            return super().compute_gradients(loss_per_module=loss_per_module, **kwargs)
+
+        for optim in self._optimizer_parameters:
+            # `set_to_none=True` is a faster way to zero out the gradients.
+            optim.zero_grad(set_to_none=True)
+
+        if self._grad_scalers is not None:
+            total_loss = sum(
+                self._grad_scalers[mid].scale(loss)
+                for mid, loss in loss_per_module.items()
+            )
+        else:
+            total_loss = sum(loss_per_module.values())
+
+        # If we don't have any loss computations, `sum` returns 0.
+        if isinstance(total_loss, int):
+            assert total_loss == 0
+            return {}
+
+        # This backward() call is inside no_sync(). It will be a clean,
+        # local-only operation.
+        total_loss.backward()
+
+        # We must force the CPU to wait for the async `backward()` call to
+        # finish on the GPU. Otherwise, the `param.grad is None` check in
+        # `_uncompiled_update` will race against the GPU and fail.
+        if torch.cuda.is_available():
+            torch.cuda.synchronize()
+
+        # This line is now safe. `grads` will have `None` for unused params
+        # (e.g., p1 on Rank 0). This is expected and handled by `_uncompiled_update`'s
+        # all_reduce loop.
+        grads = {pid: p.grad for pid, p in self._params.items()}
+
+        return grads
+
+
+ImpalaTorchLearner = IMPALATorchLearner

rllib/algorithms/impala/torch/vtrace_torch_v2.py

@@ -0,0 +1,169 @@
+from typing import List, Union
+
+from ray.rllib.utils.framework import try_import_torch
+
+torch, nn = try_import_torch()
+
+
+def make_time_major(
+    tensor: Union["torch.Tensor", List["torch.Tensor"]],
+    *,
+    trajectory_len: int = None,
+    recurrent_seq_len: int = None,
+):
+    """Swaps batch and trajectory axis.
+
+    Args:
+        tensor: A tensor or list of tensors to swap the axis of.
+            NOTE: Each tensor must have the shape [B * T] where B is the batch size and
+            T is the trajectory length.
+        trajectory_len: The length of each trajectory being transformed.
+            If None then `recurrent_seq_len` must be set.
+        recurrent_seq_len: Sequence lengths if recurrent.
+            If None then `trajectory_len` must be set.
+
+    Returns:
+        res: A tensor with swapped axes or a list of tensors with
+        swapped axes.
+    """
+    if isinstance(tensor, (list, tuple)):
+        return [
+            make_time_major(_tensor, trajectory_len, recurrent_seq_len)
+            for _tensor in tensor
+        ]
+
+    assert (
+        trajectory_len is not None or recurrent_seq_len is not None
+    ), "Either trajectory_len or recurrent_seq_len must be set."
+
+    # Figure out the sizes of the final B and T axes.
+    if recurrent_seq_len is not None:
+        assert len(tensor.shape) == 2
+        # Swap B and T axes.
+        tensor = torch.transpose(tensor, 1, 0)
+        return tensor
+    else:
+        T = trajectory_len
+        # Zero-pad, if necessary.
+        tensor_0 = tensor.shape[0]
+        B = tensor_0 // T
+        if B != (tensor_0 / T):
+            assert len(tensor.shape) == 1
+            tensor = torch.cat(
+                [
+                    tensor,
+                    torch.zeros(
+                        trajectory_len - tensor_0 % T,
+                        dtype=tensor.dtype,
+                        device=tensor.device,
+                    ),
+                ]
+            )
+            B += 1
+
+    # Reshape tensor (break up B axis into 2 axes: B and T).
+    tensor = torch.reshape(tensor, [B, T] + list(tensor.shape[1:]))
+
+    # Swap B and T axes.
+    tensor = torch.transpose(tensor, 1, 0)
+
+    return tensor
+
+
+def vtrace_torch(
+    *,
+    target_action_log_probs: "torch.Tensor",
+    behaviour_action_log_probs: "torch.Tensor",
+    discounts: "torch.Tensor",
+    rewards: "torch.Tensor",
+    values: "torch.Tensor",
+    bootstrap_values: "torch.Tensor",
+    clip_rho_threshold: Union[float, "torch.Tensor"] = 1.0,
+    clip_pg_rho_threshold: Union[float, "torch.Tensor"] = 1.0,
+):
+    r"""V-trace for softmax policies implemented with torch.
+
+    Calculates V-trace actor critic targets for softmax polices as described in
+    "IMPALA: Scalable Distributed Deep-RL with Importance Weighted Actor-Learner
+    Architectures" by Espeholt, Soyer, Munos et al. (https://arxiv.org/abs/1802.01561)
+
+    The V-trace implementation used here closely resembles the one found in the
+    scalable-agent repository by Google DeepMind, available at
+    https://github.com/deepmind/scalable_agent. This version has been optimized to
+    minimize the number of floating-point operations required per V-Trace
+    calculation, achieved through the use of dynamic programming techniques. It's
+    important to note that the mathematical expressions used in this implementation
+    may appear quite different from those presented in the IMPALA paper.
+
+    The following terminology applies:
+        - `target policy` refers to the policy we are interested in improving.
+        - `behaviour policy` refers to the policy that generated the given
+            rewards and actions.
+        - `T` refers to the time dimension. This is usually either the length of the
+            trajectory or the length of the sequence if recurrent.
+        - `B` refers to the batch size.
+
+    Args:
+        target_action_log_probs: Action log probs from the target policy. A float32
+            tensor of shape [T, B].
+        behaviour_action_log_probs: Action log probs from the behaviour policy. A
+            float32 tensor of shape [T, B].
+        discounts: A float32 tensor of shape [T, B] with the discount encountered when
+            following the behaviour policy. This will be 0 for terminal timesteps
+            (done=True) and gamma (the discount factor) otherwise.
+        rewards: A float32 tensor of shape [T, B] with the rewards generated by
+            following the behaviour policy.
+        values: A float32 tensor of shape [T, B] with the value function estimates
+            wrt. the target policy.
+        bootstrap_values: A float32 of shape [B] with the value function estimate at
+            time T.
+        clip_rho_threshold: A scalar float32 tensor with the clipping threshold for
+            importance weights (rho) when calculating the baseline targets (vs).
+            rho^bar in the paper.
+        clip_pg_rho_threshold: A scalar float32 tensor with the clipping threshold
+            on rho_s in \rho_s \delta log \pi(a|x) (r + \gamma v_{s+1} - V(x_s)).
+    """
+    log_rhos = target_action_log_probs - behaviour_action_log_probs
+
+    rhos = torch.exp(log_rhos)
+    if clip_rho_threshold is not None:
+        clipped_rhos = torch.clamp(rhos, max=clip_rho_threshold)
+    else:
+        clipped_rhos = rhos
+
+    cs = torch.clamp(rhos, max=1.0)
+    # Append bootstrapped value to get [v1, ..., v_t+1]
+    values_t_plus_1 = torch.cat(
+        [values[1:], torch.unsqueeze(bootstrap_values, 0)], axis=0
+    )
+
+    deltas = clipped_rhos * (rewards + discounts * values_t_plus_1 - values)
+
+    # Note: The original IMPALA code (and paper) suggested to perform the following
+    # v-trace for-loop on the CPU, due to its sequential nature. However, modern GPUs
+    # are quite optimized for these shorted for-loops, which is why it should be faster
+    # nowadays to leave these operations on the GPU to avoid the GPU<>CPU transfer
+    # penalty. This penalty can actually be quite massive on the LEarner actors, given
+    # all other code is already well optimized.
+    vs_minus_v_xs = [torch.zeros_like(bootstrap_values, device=deltas.device)]
+    for i in reversed(range(len(discounts))):
+        discount_t, c_t, delta_t = discounts[i], cs[i], deltas[i]
+        vs_minus_v_xs.append(delta_t + discount_t * c_t * vs_minus_v_xs[-1])
+    vs_minus_v_xs = torch.stack(vs_minus_v_xs[1:])
+
+    # Reverse the results back to original order.
+    vs_minus_v_xs = torch.flip(vs_minus_v_xs, dims=[0])
+
+    # Add V(x_s) to get v_s.
+    vs = torch.add(vs_minus_v_xs, values)
+
+    # Advantage for policy gradient.
+    vs_t_plus_1 = torch.cat([vs[1:], torch.unsqueeze(bootstrap_values, 0)], axis=0)
+    if clip_pg_rho_threshold is not None:
+        clipped_pg_rhos = torch.clamp(rhos, max=clip_pg_rho_threshold)
+    else:
+        clipped_pg_rhos = rhos
+    pg_advantages = clipped_pg_rhos * (rewards + discounts * vs_t_plus_1 - values)
+
+    # Make sure no gradients backpropagated through the returned values.
+    return torch.detach(vs), torch.detach(pg_advantages)

rllib/algorithms/impala/utils.py

@@ -0,0 +1,96 @@
+from collections import defaultdict, deque
+
+import numpy as np
+
+
+class _SleepTimeController:
+    def __init__(self):
+        self.L = 0.0
+        self.H = 0.4
+
+        self._recompute_candidates()
+
+        # Defaultdict mapping.
+        self.results = defaultdict(lambda: deque(maxlen=3))
+
+        self.iteration = 0
+
+    def _recompute_candidates(self):
+        self.center = (self.L + self.H) / 2
+        self.low = (self.L + self.center) / 2
+        self.high = (self.H + self.center) / 2
+
+        # Expand a little if range becomes too narrow to avoid
+        # overoptimization.
+        if self.H - self.L < 0.00001:
+            self.L = max(self.center - 0.1, 0.0)
+            self.H = min(self.center + 0.1, 1.0)
+            self._recompute_candidates()
+            # Reduce results, just in case it has grown too much.
+            c, l, h = (
+                self.results[self.center],
+                self.results[self.low],
+                self.results[self.high],
+            )
+            self.results = defaultdict(lambda: deque(maxlen=3))
+            self.results[self.center] = c
+            self.results[self.low] = l
+            self.results[self.high] = h
+
+    @property
+    def current(self):
+        if len(self.results[self.center]) < 3:
+            return self.center
+        elif len(self.results[self.low]) < 3:
+            return self.low
+        else:
+            return self.high
+
+    def log_result(self, performance):
+        self.iteration += 1
+
+        # Skip first 2 iterations for ignoring warm-up effect.
+        if self.iteration < 2:
+            return
+
+        self.results[self.current].append(performance)
+
+        # If all candidates have at least 3 results logged, re-evaluate
+        # and compute new L and H.
+        center, low, high = self.center, self.low, self.high
+        if (
+            len(self.results[center]) == 3
+            and len(self.results[low]) == 3
+            and len(self.results[high]) == 3
+        ):
+            perf_center = np.mean(self.results[center])
+            perf_low = np.mean(self.results[low])
+            perf_high = np.mean(self.results[high])
+            # Case: `center` is best.
+            if perf_center > perf_low and perf_center > perf_high:
+                self.L = low
+                self.H = high
+                # Erase low/high results: We'll not use these again.
+                self.results.pop(low, None)
+                self.results.pop(high, None)
+            # Case: `low` is best.
+            elif perf_low > perf_center and perf_low > perf_high:
+                self.H = center
+                # Erase center/high results: We'll not use these again.
+                self.results.pop(center, None)
+                self.results.pop(high, None)
+            # Case: `high` is best.
+            else:
+                self.L = center
+                # Erase center/low results: We'll not use these again.
+                self.results.pop(center, None)
+                self.results.pop(low, None)
+
+            self._recompute_candidates()
+
+
+if __name__ == "__main__":
+    controller = _SleepTimeController()
+    for _ in range(1000):
+        performance = np.random.random()
+        controller.log_result(performance)

rllib/algorithms/impala/vtrace_tf.py

@@ -0,0 +1,425 @@
+# Copyright 2018 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Functions to compute V-trace off-policy actor critic targets.
+
+For details and theory see:
+
+"IMPALA: Scalable Distributed Deep-RL with
+Importance Weighted Actor-Learner Architectures"
+by Espeholt, Soyer, Munos et al.
+
+See https://arxiv.org/abs/1802.01561 for the full paper.
+
+In addition to the original paper's code, changes have been made
+to support MultiDiscrete action spaces. behaviour_policy_logits,
+target_policy_logits and actions parameters in the entry point
+multi_from_logits method accepts lists of tensors instead of just
+tensors.
+"""
+
+import collections
+
+from ray.rllib.models.tf.tf_action_dist import Categorical
+from ray.rllib.utils.framework import try_import_tf
+
+tf1, tf, tfv = try_import_tf()
+
+VTraceFromLogitsReturns = collections.namedtuple(
+    "VTraceFromLogitsReturns",
+    [
+        "vs",
+        "pg_advantages",
+        "log_rhos",
+        "behaviour_action_log_probs",
+        "target_action_log_probs",
+    ],
+)
+
+VTraceReturns = collections.namedtuple("VTraceReturns", "vs pg_advantages")
+
+
+def log_probs_from_logits_and_actions(
+    policy_logits, actions, dist_class=Categorical, model=None
+):
+    return multi_log_probs_from_logits_and_actions(
+        [policy_logits], [actions], dist_class, model
+    )[0]
+
+
+def multi_log_probs_from_logits_and_actions(policy_logits, actions, dist_class, model):
+    """Computes action log-probs from policy logits and actions.
+
+    In the notation used throughout documentation and comments, T refers to the
+    time dimension ranging from 0 to T-1. B refers to the batch size and
+    ACTION_SPACE refers to the list of numbers each representing a number of
+    actions.
+
+    Args:
+        policy_logits: A list with length of ACTION_SPACE of float32
+            tensors of shapes [T, B, ACTION_SPACE[0]], ...,
+            [T, B, ACTION_SPACE[-1]] with un-normalized log-probabilities
+            parameterizing a softmax policy.
+        actions: A list with length of ACTION_SPACE of tensors of shapes
+            [T, B, ...], ..., [T, B, ...]
+            with actions.
+        dist_class: Python class of the action distribution.
+
+    Returns:
+        A list with length of ACTION_SPACE of float32 tensors of shapes
+            [T, B], ..., [T, B] corresponding to the sampling log probability
+            of the chosen action w.r.t. the policy.
+    """
+    log_probs = []
+    for i in range(len(policy_logits)):
+        p_shape = tf.shape(policy_logits[i])
+        a_shape = tf.shape(actions[i])
+        policy_logits_flat = tf.reshape(
+            policy_logits[i], tf.concat([[-1], p_shape[2:]], axis=0)
+        )
+        actions_flat = tf.reshape(actions[i], tf.concat([[-1], a_shape[2:]], axis=0))
+        log_probs.append(
+            tf.reshape(
+                dist_class(policy_logits_flat, model).logp(actions_flat), a_shape[:2]
+            )
+        )
+
+    return log_probs
+
+
+def from_logits(
+    behaviour_policy_logits,
+    target_policy_logits,
+    actions,
+    discounts,
+    rewards,
+    values,
+    bootstrap_value,
+    dist_class=Categorical,
+    model=None,
+    clip_rho_threshold=1.0,
+    clip_pg_rho_threshold=1.0,
+    name="vtrace_from_logits",
+):
+    """multi_from_logits wrapper used only for tests"""
+
+    res = multi_from_logits(
+        [behaviour_policy_logits],
+        [target_policy_logits],
+        [actions],
+        discounts,
+        rewards,
+        values,
+        bootstrap_value,
+        dist_class,
+        model,
+        clip_rho_threshold=clip_rho_threshold,
+        clip_pg_rho_threshold=clip_pg_rho_threshold,
+        name=name,
+    )
+
+    return VTraceFromLogitsReturns(
+        vs=res.vs,
+        pg_advantages=res.pg_advantages,
+        log_rhos=res.log_rhos,
+        behaviour_action_log_probs=tf.squeeze(res.behaviour_action_log_probs, axis=0),
+        target_action_log_probs=tf.squeeze(res.target_action_log_probs, axis=0),
+    )
+
+
+def multi_from_logits(
+    behaviour_policy_logits,
+    target_policy_logits,
+    actions,
+    discounts,
+    rewards,
+    values,
+    bootstrap_value,
+    dist_class,
+    model,
+    behaviour_action_log_probs=None,
+    clip_rho_threshold=1.0,
+    clip_pg_rho_threshold=1.0,
+    name="vtrace_from_logits",
+):
+    r"""V-trace for softmax policies.
+
+    Calculates V-trace actor critic targets for softmax polices as described in
+
+    "IMPALA: Scalable Distributed Deep-RL with
+    Importance Weighted Actor-Learner Architectures"
+    by Espeholt, Soyer, Munos et al.
+
+    Target policy refers to the policy we are interested in improving and
+    behaviour policy refers to the policy that generated the given
+    rewards and actions.
+
+    In the notation used throughout documentation and comments, T refers to the
+    time dimension ranging from 0 to T-1. B refers to the batch size and
+    ACTION_SPACE refers to the list of numbers each representing a number of
+    actions.
+
+    Args:
+      behaviour_policy_logits: A list with length of ACTION_SPACE of float32
+        tensors of shapes
+        [T, B, ACTION_SPACE[0]],
+        ...,
+        [T, B, ACTION_SPACE[-1]]
+        with un-normalized log-probabilities parameterizing the softmax behaviour
+        policy.
+      target_policy_logits: A list with length of ACTION_SPACE of float32
+        tensors of shapes
+        [T, B, ACTION_SPACE[0]],
+        ...,
+        [T, B, ACTION_SPACE[-1]]
+        with un-normalized log-probabilities parameterizing the softmax target
+        policy.
+      actions: A list with length of ACTION_SPACE of
+        tensors of shapes
+        [T, B, ...],
+        ...,
+        [T, B, ...]
+        with actions sampled from the behaviour policy.
+      discounts: A float32 tensor of shape [T, B] with the discount encountered
+        when following the behaviour policy.
+      rewards: A float32 tensor of shape [T, B] with the rewards generated by
+        following the behaviour policy.
+      values: A float32 tensor of shape [T, B] with the value function estimates
+        wrt. the target policy.
+      bootstrap_value: A float32 of shape [B] with the value function estimate at
+        time T.
+      dist_class: action distribution class for the logits.
+      model: backing ModelV2 instance
+      behaviour_action_log_probs: precalculated values of the behaviour actions
+      clip_rho_threshold: A scalar float32 tensor with the clipping threshold for
+        importance weights (rho) when calculating the baseline targets (vs).
+        rho^bar in the paper.
+      clip_pg_rho_threshold: A scalar float32 tensor with the clipping threshold
+        on rho_s in \rho_s \delta log \pi(a|x) (r + \gamma v_{s+1} - V(x_s)).
+      name: The name scope that all V-trace operations will be created in.
+
+    Returns:
+      A `VTraceFromLogitsReturns` namedtuple with the following fields:
+        vs: A float32 tensor of shape [T, B]. Can be used as target to train a
+            baseline (V(x_t) - vs_t)^2.
+        pg_advantages: A float 32 tensor of shape [T, B]. Can be used as an
+          estimate of the advantage in the calculation of policy gradients.
+        log_rhos: A float32 tensor of shape [T, B] containing the log importance
+          sampling weights (log rhos).
+        behaviour_action_log_probs: A float32 tensor of shape [T, B] containing
+          behaviour policy action log probabilities (log \mu(a_t)).
+        target_action_log_probs: A float32 tensor of shape [T, B] containing
+          target policy action probabilities (log \pi(a_t)).
+    """
+
+    for i in range(len(behaviour_policy_logits)):
+        behaviour_policy_logits[i] = tf.convert_to_tensor(
+            behaviour_policy_logits[i], dtype=tf.float32
+        )
+        target_policy_logits[i] = tf.convert_to_tensor(
+            target_policy_logits[i], dtype=tf.float32
+        )
+
+        # Make sure tensor ranks are as expected.
+        # The rest will be checked by from_action_log_probs.
+        behaviour_policy_logits[i].shape.assert_has_rank(3)
+        target_policy_logits[i].shape.assert_has_rank(3)
+
+    with tf1.name_scope(
+        name,
+        values=[
+            behaviour_policy_logits,
+            target_policy_logits,
+            actions,
+            discounts,
+            rewards,
+            values,
+            bootstrap_value,
+        ],
+    ):
+        target_action_log_probs = multi_log_probs_from_logits_and_actions(
+            target_policy_logits, actions, dist_class, model
+        )
+
+        if len(behaviour_policy_logits) > 1 or behaviour_action_log_probs is None:
+            # can't use precalculated values, recompute them. Note that
+            # recomputing won't work well for autoregressive action dists
+            # which may have variables not captured by 'logits'
+            behaviour_action_log_probs = multi_log_probs_from_logits_and_actions(
+                behaviour_policy_logits, actions, dist_class, model
+            )
+
+        log_rhos = get_log_rhos(target_action_log_probs, behaviour_action_log_probs)
+
+        vtrace_returns = from_importance_weights(
+            log_rhos=log_rhos,
+            discounts=discounts,
+            rewards=rewards,
+            values=values,
+            bootstrap_value=bootstrap_value,
+            clip_rho_threshold=clip_rho_threshold,
+            clip_pg_rho_threshold=clip_pg_rho_threshold,
+        )
+
+        return VTraceFromLogitsReturns(
+            log_rhos=log_rhos,
+            behaviour_action_log_probs=behaviour_action_log_probs,
+            target_action_log_probs=target_action_log_probs,
+            **vtrace_returns._asdict()
+        )
+
+
+def from_importance_weights(
+    log_rhos,
+    discounts,
+    rewards,
+    values,
+    bootstrap_value,
+    clip_rho_threshold=1.0,
+    clip_pg_rho_threshold=1.0,
+    name="vtrace_from_importance_weights",
+):
+    r"""V-trace from log importance weights.
+
+    Calculates V-trace actor critic targets as described in
+
+    "IMPALA: Scalable Distributed Deep-RL with
+    Importance Weighted Actor-Learner Architectures"
+    by Espeholt, Soyer, Munos et al.
+
+    In the notation used throughout documentation and comments, T refers to the
+    time dimension ranging from 0 to T-1. B refers to the batch size. This code
+    also supports the case where all tensors have the same number of additional
+    dimensions, e.g., `rewards` is [T, B, C], `values` is [T, B, C],
+    `bootstrap_value` is [B, C].
+
+    Args:
+      log_rhos: A float32 tensor of shape [T, B] representing the
+        log importance sampling weights, i.e.
+        log(target_policy(a) / behaviour_policy(a)). V-trace performs operations
+        on rhos in log-space for numerical stability.
+      discounts: A float32 tensor of shape [T, B] with discounts encountered when
+        following the behaviour policy.
+      rewards: A float32 tensor of shape [T, B] containing rewards generated by
+        following the behaviour policy.
+      values: A float32 tensor of shape [T, B] with the value function estimates
+        wrt. the target policy.
+      bootstrap_value: A float32 of shape [B] with the value function estimate at
+        time T.
+      clip_rho_threshold: A scalar float32 tensor with the clipping threshold for
+        importance weights (rho) when calculating the baseline targets (vs).
+        rho^bar in the paper. If None, no clipping is applied.
+      clip_pg_rho_threshold: A scalar float32 tensor with the clipping threshold
+        on rho_s in \rho_s \delta log \pi(a|x) (r + \gamma v_{s+1} - V(x_s)). If
+        None, no clipping is applied.
+      name: The name scope that all V-trace operations will be created in.
+
+    Returns:
+      A VTraceReturns namedtuple (vs, pg_advantages) where:
+        vs: A float32 tensor of shape [T, B]. Can be used as target to
+          train a baseline (V(x_t) - vs_t)^2.
+        pg_advantages: A float32 tensor of shape [T, B]. Can be used as the
+          advantage in the calculation of policy gradients.
+    """
+    log_rhos = tf.convert_to_tensor(log_rhos, dtype=tf.float32)
+    discounts = tf.convert_to_tensor(discounts, dtype=tf.float32)
+    rewards = tf.convert_to_tensor(rewards, dtype=tf.float32)
+    values = tf.convert_to_tensor(values, dtype=tf.float32)
+    bootstrap_value = tf.convert_to_tensor(bootstrap_value, dtype=tf.float32)
+    if clip_rho_threshold is not None:
+        clip_rho_threshold = tf.convert_to_tensor(clip_rho_threshold, dtype=tf.float32)
+    if clip_pg_rho_threshold is not None:
+        clip_pg_rho_threshold = tf.convert_to_tensor(
+            clip_pg_rho_threshold, dtype=tf.float32
+        )
+
+    # Make sure tensor ranks are consistent.
+    rho_rank = log_rhos.shape.ndims  # Usually 2.
+    values.shape.assert_has_rank(rho_rank)
+    bootstrap_value.shape.assert_has_rank(rho_rank - 1)
+    discounts.shape.assert_has_rank(rho_rank)
+    rewards.shape.assert_has_rank(rho_rank)
+    if clip_rho_threshold is not None:
+        clip_rho_threshold.shape.assert_has_rank(0)
+    if clip_pg_rho_threshold is not None:
+        clip_pg_rho_threshold.shape.assert_has_rank(0)
+
+    with tf1.name_scope(
+        name, values=[log_rhos, discounts, rewards, values, bootstrap_value]
+    ):
+        rhos = tf.math.exp(log_rhos)
+        if clip_rho_threshold is not None:
+            clipped_rhos = tf.minimum(clip_rho_threshold, rhos, name="clipped_rhos")
+        else:
+            clipped_rhos = rhos
+
+        cs = tf.minimum(1.0, rhos, name="cs")
+        # Append bootstrapped value to get [v1, ..., v_t+1]
+        values_t_plus_1 = tf.concat(
+            [values[1:], tf.expand_dims(bootstrap_value, 0)], axis=0
+        )
+        deltas = clipped_rhos * (rewards + discounts * values_t_plus_1 - values)
+
+        # All sequences are reversed, computation starts from the back.
+        sequences = (
+            tf.reverse(discounts, axis=[0]),
+            tf.reverse(cs, axis=[0]),
+            tf.reverse(deltas, axis=[0]),
+        )
+
+        # V-trace vs are calculated through a scan from the back to the
+        # beginning of the given trajectory.
+        def scanfunc(acc, sequence_item):
+            discount_t, c_t, delta_t = sequence_item
+            return delta_t + discount_t * c_t * acc
+
+        initial_values = tf.zeros_like(bootstrap_value)
+        vs_minus_v_xs = tf.nest.map_structure(
+            tf.stop_gradient,
+            tf.scan(
+                fn=scanfunc,
+                elems=sequences,
+                initializer=initial_values,
+                parallel_iterations=1,
+                name="scan",
+            ),
+        )
+        # Reverse the results back to original order.
+        vs_minus_v_xs = tf.reverse(vs_minus_v_xs, [0], name="vs_minus_v_xs")
+
+        # Add V(x_s) to get v_s.
+        vs = tf.add(vs_minus_v_xs, values, name="vs")
+
+        # Advantage for policy gradient.
+        vs_t_plus_1 = tf.concat([vs[1:], tf.expand_dims(bootstrap_value, 0)], axis=0)
+        if clip_pg_rho_threshold is not None:
+            clipped_pg_rhos = tf.minimum(
+                clip_pg_rho_threshold, rhos, name="clipped_pg_rhos"
+            )
+        else:
+            clipped_pg_rhos = rhos
+        pg_advantages = clipped_pg_rhos * (rewards + discounts * vs_t_plus_1 - values)
+
+        # Make sure no gradients backpropagated through the returned values.
+        return VTraceReturns(
+            vs=tf.stop_gradient(vs), pg_advantages=tf.stop_gradient(pg_advantages)
+        )
+
+
+def get_log_rhos(target_action_log_probs, behaviour_action_log_probs):
+    """With the selected log_probs for multi-discrete actions of behaviour
+    and target policies we compute the log_rhos for calculating the vtrace."""
+    t = tf.stack(target_action_log_probs)
+    b = tf.stack(behaviour_action_log_probs)
+    log_rhos = tf.reduce_sum(t - b, axis=0)
+    return log_rhos

rllib/algorithms/impala/vtrace_torch.py

@@ -0,0 +1,359 @@
+# Copyright 2018 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""PyTorch version of the functions to compute V-trace off-policy actor critic
+targets.
+
+For details and theory see:
+
+"IMPALA: Scalable Distributed Deep-RL with
+Importance Weighted Actor-Learner Architectures"
+by Espeholt, Soyer, Munos et al.
+
+See https://arxiv.org/abs/1802.01561 for the full paper.
+
+In addition to the original paper's code, changes have been made
+to support MultiDiscrete action spaces. behaviour_policy_logits,
+target_policy_logits and actions parameters in the entry point
+multi_from_logits method accepts lists of tensors instead of just
+tensors.
+"""
+
+from ray.rllib.algorithms.impala.vtrace_tf import VTraceFromLogitsReturns, VTraceReturns
+from ray.rllib.models.torch.torch_action_dist import TorchCategorical
+from ray.rllib.utils import force_list
+from ray.rllib.utils.framework import try_import_torch
+from ray.rllib.utils.torch_utils import convert_to_torch_tensor
+
+torch, nn = try_import_torch()
+
+
+def log_probs_from_logits_and_actions(
+    policy_logits, actions, dist_class=TorchCategorical, model=None
+):
+    return multi_log_probs_from_logits_and_actions(
+        [policy_logits], [actions], dist_class, model
+    )[0]
+
+
+def multi_log_probs_from_logits_and_actions(policy_logits, actions, dist_class, model):
+    """Computes action log-probs from policy logits and actions.
+
+    In the notation used throughout documentation and comments, T refers to the
+    time dimension ranging from 0 to T-1. B refers to the batch size and
+    ACTION_SPACE refers to the list of numbers each representing a number of
+    actions.
+
+    Args:
+        policy_logits: A list with length of ACTION_SPACE of float32
+            tensors of shapes [T, B, ACTION_SPACE[0]], ...,
+            [T, B, ACTION_SPACE[-1]] with un-normalized log-probabilities
+            parameterizing a softmax policy.
+        actions: A list with length of ACTION_SPACE of tensors of shapes
+            [T, B, ...], ..., [T, B, ...]
+            with actions.
+        dist_class: Python class of the action distribution.
+
+    Returns:
+        A list with length of ACTION_SPACE of float32 tensors of shapes
+            [T, B], ..., [T, B] corresponding to the sampling log probability
+            of the chosen action w.r.t. the policy.
+    """
+    log_probs = []
+    for i in range(len(policy_logits)):
+        p_shape = policy_logits[i].shape
+        a_shape = actions[i].shape
+        policy_logits_flat = torch.reshape(policy_logits[i], (-1,) + tuple(p_shape[2:]))
+        actions_flat = torch.reshape(actions[i], (-1,) + tuple(a_shape[2:]))
+        log_probs.append(
+            torch.reshape(
+                dist_class(policy_logits_flat, model).logp(actions_flat), a_shape[:2]
+            )
+        )
+
+    return log_probs
+
+
+def from_logits(
+    behaviour_policy_logits,
+    target_policy_logits,
+    actions,
+    discounts,
+    rewards,
+    values,
+    bootstrap_value,
+    dist_class=TorchCategorical,
+    model=None,
+    clip_rho_threshold=1.0,
+    clip_pg_rho_threshold=1.0,
+):
+    """multi_from_logits wrapper used only for tests"""
+
+    res = multi_from_logits(
+        [behaviour_policy_logits],
+        [target_policy_logits],
+        [actions],
+        discounts,
+        rewards,
+        values,
+        bootstrap_value,
+        dist_class,
+        model,
+        clip_rho_threshold=clip_rho_threshold,
+        clip_pg_rho_threshold=clip_pg_rho_threshold,
+    )
+
+    assert len(res.behaviour_action_log_probs) == 1
+    assert len(res.target_action_log_probs) == 1
+    return VTraceFromLogitsReturns(
+        vs=res.vs,
+        pg_advantages=res.pg_advantages,
+        log_rhos=res.log_rhos,
+        behaviour_action_log_probs=res.behaviour_action_log_probs[0],
+        target_action_log_probs=res.target_action_log_probs[0],
+    )
+
+
+def multi_from_logits(
+    behaviour_policy_logits,
+    target_policy_logits,
+    actions,
+    discounts,
+    rewards,
+    values,
+    bootstrap_value,
+    dist_class,
+    model,
+    behaviour_action_log_probs=None,
+    clip_rho_threshold=1.0,
+    clip_pg_rho_threshold=1.0,
+):
+    r"""V-trace for softmax policies.
+
+    Calculates V-trace actor critic targets for softmax polices as described in
+
+    "IMPALA: Scalable Distributed Deep-RL with
+    Importance Weighted Actor-Learner Architectures"
+    by Espeholt, Soyer, Munos et al.
+
+    Target policy refers to the policy we are interested in improving and
+    behaviour policy refers to the policy that generated the given
+    rewards and actions.
+
+    In the notation used throughout documentation and comments, T refers to the
+    time dimension ranging from 0 to T-1. B refers to the batch size and
+    ACTION_SPACE refers to the list of numbers each representing a number of
+    actions.
+
+    Args:
+        behaviour_policy_logits: A list with length of ACTION_SPACE of float32
+            tensors of shapes [T, B, ACTION_SPACE[0]], ...,
+            [T, B, ACTION_SPACE[-1]] with un-normalized log-probabilities
+            parameterizing the softmax behavior policy.
+        target_policy_logits: A list with length of ACTION_SPACE of float32
+            tensors of shapes [T, B, ACTION_SPACE[0]], ...,
+            [T, B, ACTION_SPACE[-1]] with un-normalized log-probabilities
+            parameterizing the softmax target policy.
+        actions: A list with length of ACTION_SPACE of tensors of shapes
+            [T, B, ...], ..., [T, B, ...]
+            with actions sampled from the behavior policy.
+        discounts: A float32 tensor of shape [T, B] with the discount
+            encountered when following the behavior policy.
+        rewards: A float32 tensor of shape [T, B] with the rewards generated by
+            following the behavior policy.
+        values: A float32 tensor of shape [T, B] with the value function
+            estimates wrt. the target policy.
+        bootstrap_value: A float32 of shape [B] with the value function
+            estimate at time T.
+        dist_class: action distribution class for the logits.
+        model: backing ModelV2 instance
+        behaviour_action_log_probs: Precalculated values of the behavior
+            actions.
+        clip_rho_threshold: A scalar float32 tensor with the clipping threshold
+            for importance weights (rho) when calculating the baseline targets
+            (vs). rho^bar in the paper.
+        clip_pg_rho_threshold: A scalar float32 tensor with the clipping
+            threshold on rho_s in:
+            \rho_s \delta log \pi(a|x) (r + \gamma v_{s+1} - V(x_s)).
+
+    Returns:
+        A `VTraceFromLogitsReturns` namedtuple with the following fields:
+        vs: A float32 tensor of shape [T, B]. Can be used as target to train a
+            baseline (V(x_t) - vs_t)^2.
+        pg_advantages: A float 32 tensor of shape [T, B]. Can be used as an
+            estimate of the advantage in the calculation of policy gradients.
+        log_rhos: A float32 tensor of shape [T, B] containing the log
+            importance sampling weights (log rhos).
+        behaviour_action_log_probs: A float32 tensor of shape [T, B] containing
+            behaviour policy action log probabilities (log \mu(a_t)).
+        target_action_log_probs: A float32 tensor of shape [T, B] containing
+            target policy action probabilities (log \pi(a_t)).
+    """
+
+    behaviour_policy_logits = convert_to_torch_tensor(
+        behaviour_policy_logits, device="cpu"
+    )
+    target_policy_logits = convert_to_torch_tensor(target_policy_logits, device="cpu")
+    actions = convert_to_torch_tensor(actions, device="cpu")
+
+    # Make sure tensor ranks are as expected.
+    # The rest will be checked by from_action_log_probs.
+    for i in range(len(behaviour_policy_logits)):
+        assert len(behaviour_policy_logits[i].size()) == 3
+        assert len(target_policy_logits[i].size()) == 3
+
+    target_action_log_probs = multi_log_probs_from_logits_and_actions(
+        target_policy_logits, actions, dist_class, model
+    )
+
+    if len(behaviour_policy_logits) > 1 or behaviour_action_log_probs is None:
+        # can't use precalculated values, recompute them. Note that
+        # recomputing won't work well for autoregressive action dists
+        # which may have variables not captured by 'logits'
+        behaviour_action_log_probs = multi_log_probs_from_logits_and_actions(
+            behaviour_policy_logits, actions, dist_class, model
+        )
+
+    behaviour_action_log_probs = convert_to_torch_tensor(
+        behaviour_action_log_probs, device="cpu"
+    )
+    behaviour_action_log_probs = force_list(behaviour_action_log_probs)
+    # log_rhos = target_logp - behavior_logp
+    log_rhos = get_log_rhos(target_action_log_probs, behaviour_action_log_probs)
+
+    vtrace_returns = from_importance_weights(
+        log_rhos=log_rhos,
+        discounts=discounts,
+        rewards=rewards,
+        values=values,
+        bootstrap_value=bootstrap_value,
+        clip_rho_threshold=clip_rho_threshold,
+        clip_pg_rho_threshold=clip_pg_rho_threshold,
+    )
+
+    return VTraceFromLogitsReturns(
+        log_rhos=log_rhos,
+        behaviour_action_log_probs=behaviour_action_log_probs,
+        target_action_log_probs=target_action_log_probs,
+        **vtrace_returns._asdict()
+    )
+
+
+def from_importance_weights(
+    log_rhos,
+    discounts,
+    rewards,
+    values,
+    bootstrap_value,
+    clip_rho_threshold=1.0,
+    clip_pg_rho_threshold=1.0,
+):
+    r"""V-trace from log importance weights.
+
+    Calculates V-trace actor critic targets as described in
+
+    "IMPALA: Scalable Distributed Deep-RL with
+    Importance Weighted Actor-Learner Architectures"
+    by Espeholt, Soyer, Munos et al.
+
+    In the notation used throughout documentation and comments, T refers to the
+    time dimension ranging from 0 to T-1. B refers to the batch size. This code
+    also supports the case where all tensors have the same number of additional
+    dimensions, e.g., `rewards` is [T, B, C], `values` is [T, B, C],
+    `bootstrap_value` is [B, C].
+
+    Args:
+        log_rhos: A float32 tensor of shape [T, B] representing the log
+            importance sampling weights, i.e.
+            log(target_policy(a) / behaviour_policy(a)). V-trace performs
+            operations on rhos in log-space for numerical stability.
+        discounts: A float32 tensor of shape [T, B] with discounts encountered
+            when following the behaviour policy.
+        rewards: A float32 tensor of shape [T, B] containing rewards generated
+            by following the behaviour policy.
+        values: A float32 tensor of shape [T, B] with the value function
+            estimates wrt. the target policy.
+        bootstrap_value: A float32 of shape [B] with the value function
+            estimate at time T.
+        clip_rho_threshold: A scalar float32 tensor with the clipping threshold
+            for importance weights (rho) when calculating the baseline targets
+            (vs). rho^bar in the paper. If None, no clipping is applied.
+        clip_pg_rho_threshold: A scalar float32 tensor with the clipping
+            threshold on rho_s in
+            \rho_s \delta log \pi(a|x) (r + \gamma v_{s+1} - V(x_s)).
+            If None, no clipping is applied.
+
+    Returns:
+        A VTraceReturns namedtuple (vs, pg_advantages) where:
+        vs: A float32 tensor of shape [T, B]. Can be used as target to
+            train a baseline (V(x_t) - vs_t)^2.
+        pg_advantages: A float32 tensor of shape [T, B]. Can be used as the
+            advantage in the calculation of policy gradients.
+    """
+    log_rhos = convert_to_torch_tensor(log_rhos, device="cpu")
+    discounts = convert_to_torch_tensor(discounts, device="cpu")
+    rewards = convert_to_torch_tensor(rewards, device="cpu")
+    values = convert_to_torch_tensor(values, device="cpu")
+    bootstrap_value = convert_to_torch_tensor(bootstrap_value, device="cpu")
+
+    # Make sure tensor ranks are consistent.
+    rho_rank = len(log_rhos.size())  # Usually 2.
+    assert rho_rank == len(values.size())
+    assert rho_rank - 1 == len(bootstrap_value.size()), "must have rank {}".format(
+        rho_rank - 1
+    )
+    assert rho_rank == len(discounts.size())
+    assert rho_rank == len(rewards.size())
+
+    rhos = torch.exp(log_rhos)
+    if clip_rho_threshold is not None:
+        clipped_rhos = torch.clamp_max(rhos, clip_rho_threshold)
+    else:
+        clipped_rhos = rhos
+
+    cs = torch.clamp_max(rhos, 1.0)
+    # Append bootstrapped value to get [v1, ..., v_t+1]
+    values_t_plus_1 = torch.cat(
+        [values[1:], torch.unsqueeze(bootstrap_value, 0)], dim=0
+    )
+    deltas = clipped_rhos * (rewards + discounts * values_t_plus_1 - values)
+
+    vs_minus_v_xs = [torch.zeros_like(bootstrap_value)]
+    for i in reversed(range(len(discounts))):
+        discount_t, c_t, delta_t = discounts[i], cs[i], deltas[i]
+        vs_minus_v_xs.append(delta_t + discount_t * c_t * vs_minus_v_xs[-1])
+    vs_minus_v_xs = torch.stack(vs_minus_v_xs[1:])
+    # Reverse the results back to original order.
+    vs_minus_v_xs = torch.flip(vs_minus_v_xs, dims=[0])
+    # Add V(x_s) to get v_s.
+    vs = vs_minus_v_xs + values
+
+    # Advantage for policy gradient.
+    vs_t_plus_1 = torch.cat([vs[1:], torch.unsqueeze(bootstrap_value, 0)], dim=0)
+    if clip_pg_rho_threshold is not None:
+        clipped_pg_rhos = torch.clamp_max(rhos, clip_pg_rho_threshold)
+    else:
+        clipped_pg_rhos = rhos
+    pg_advantages = clipped_pg_rhos * (rewards + discounts * vs_t_plus_1 - values)
+
+    # Make sure no gradients backpropagated through the returned values.
+    return VTraceReturns(vs=vs.detach(), pg_advantages=pg_advantages.detach())
+
+
+def get_log_rhos(target_action_log_probs, behaviour_action_log_probs):
+    """With the selected log_probs for multi-discrete actions of behavior
+    and target policies we compute the log_rhos for calculating the vtrace."""
+    t = torch.stack(target_action_log_probs)
+    b = torch.stack(behaviour_action_log_probs)
+    log_rhos = torch.sum(t - b, dim=0)
+    return log_rhos