chore: import upstream snapshot with attribution

rllib/algorithms/__init__.py

@@ -0,0 +1,38 @@
+from ray.rllib.algorithms.algorithm import Algorithm
+from ray.rllib.algorithms.algorithm_config import AlgorithmConfig
+from ray.rllib.algorithms.appo.appo import APPO, APPOConfig
+from ray.rllib.algorithms.bc.bc import BC, BCConfig
+from ray.rllib.algorithms.cql.cql import CQL, CQLConfig
+from ray.rllib.algorithms.dqn.dqn import DQN, DQNConfig
+from ray.rllib.algorithms.impala.impala import (
+    IMPALA,
+    Impala,
+    IMPALAConfig,
+    ImpalaConfig,
+)
+from ray.rllib.algorithms.marwil.marwil import MARWIL, MARWILConfig
+from ray.rllib.algorithms.ppo.ppo import PPO, PPOConfig
+from ray.rllib.algorithms.sac.sac import SAC, SACConfig
+
+__all__ = [
+    "Algorithm",
+    "AlgorithmConfig",
+    "APPO",
+    "APPOConfig",
+    "BC",
+    "BCConfig",
+    "CQL",
+    "CQLConfig",
+    "DQN",
+    "DQNConfig",
+    "IMPALA",
+    "IMPALAConfig",
+    "Impala",
+    "ImpalaConfig",
+    "MARWIL",
+    "MARWILConfig",
+    "PPO",
+    "PPOConfig",
+    "SAC",
+    "SACConfig",
+]

rllib/algorithms/appo/README.md

@@ -0,0 +1,34 @@
+# Asynchronous Proximal Policy Optimization (APPO)
+
+## Overview
+
+[PPO](https://arxiv.org/abs/1707.06347) is a model-free on-policy RL algorithm that works
+well for both discrete and continuous action space environments. PPO utilizes an
+actor-critic framework, where there are two networks, an actor (policy network) and
+critic network (value function).
+
+## Distributed PPO Algorithms
+
+### Distributed baseline PPO
+[See implementation here](https://github.com/ray-project/ray/blob/master/rllib/algorithms/ppo/ppo.py)
+
+### Asychronous PPO (APPO) ..
+
+.. opts to imitate IMPALA as its distributed execution plan.
+Data collection nodes gather data asynchronously, which are collected in a circular replay
+buffer. A target network and doubly-importance sampled surrogate objective is introduced
+to enforce training stability in the asynchronous data-collection setting.
+[See implementation here](https://github.com/ray-project/ray/blob/master/rllib/algorithms/appo/appo.py)
+
+### Decentralized Distributed PPO (DDPPO)
+
+[See implementation here](https://github.com/ray-project/ray/blob/master/rllib/algorithms/ddppo/ddppo.py)
+
+
+## Documentation & Implementation:
+
+### [Asynchronous Proximal Policy Optimization (APPO)](https://arxiv.org/abs/1912.00167).
+
+**[Detailed Documentation](https://docs.ray.io/en/master/rllib-algorithms.html#appo)**
+
+**[Implementation](https://github.com/ray-project/ray/blob/master/rllib/agents/ppo/appo.py)**

rllib/algorithms/appo/__init__.py

@@ -0,0 +1,12 @@
+from ray.rllib.algorithms.appo.appo import APPO, APPOConfig
+from ray.rllib.algorithms.appo.appo_tf_policy import APPOTF1Policy, APPOTF2Policy
+from ray.rllib.algorithms.appo.appo_torch_policy import APPOTorchPolicy
+
+__all__ = [
+    "APPO",
+    "APPOConfig",
+    # @OldAPIStack
+    "APPOTF1Policy",
+    "APPOTF2Policy",
+    "APPOTorchPolicy",
+]

rllib/algorithms/appo/appo.py

@@ -0,0 +1,456 @@
+"""Asynchronous Proximal Policy Optimization (APPO)
+
+The algorithm is described in [1] (under the name of "IMPACT"):
+
+Detailed documentation:
+https://docs.ray.io/en/master/rllib-algorithms.html#appo
+
+[1] IMPACT: Importance Weighted Asynchronous Architectures with Clipped Target Networks.
+Luo et al. 2020
+https://arxiv.org/pdf/1912.00167
+"""
+
+import logging
+from typing import Optional, Type
+
+from typing_extensions import Self
+
+from ray._common.deprecation import DEPRECATED_VALUE, deprecation_warning
+from ray.rllib.algorithms.algorithm_config import AlgorithmConfig, NotProvided
+from ray.rllib.algorithms.impala.impala import IMPALA, IMPALAConfig
+from ray.rllib.core.rl_module.rl_module import RLModuleSpec
+from ray.rllib.policy.policy import Policy
+from ray.rllib.utils.annotations import override
+from ray.rllib.utils.metrics import (
+    LAST_TARGET_UPDATE_TS,
+    LEARNER_STATS_KEY,
+    NUM_AGENT_STEPS_SAMPLED,
+    NUM_ENV_STEPS_SAMPLED,
+    NUM_TARGET_UPDATES,
+)
+
+logger = logging.getLogger(__name__)
+
+
+LEARNER_RESULTS_KL_KEY = "mean_kl_loss"
+LEARNER_RESULTS_CURR_KL_COEFF_KEY = "curr_kl_coeff"
+OLD_ACTION_DIST_KEY = "old_action_dist"
+# Mean and variance of the IMPACT clipped IS ratio
+# (`clip(pi_behaviour / pi_old_target, 0, 2)`)
+LEARNER_RESULTS_MEAN_IS_KEY = "mean_IS"
+LEARNER_RESULTS_VAR_IS_KEY = "var_IS"
+
+
+class APPOConfig(IMPALAConfig):
+    """Defines a configuration class from which an APPO Algorithm can be built.
+
+    .. testcode::
+
+        from ray.rllib.algorithms.appo import APPOConfig
+        config = (
+            APPOConfig()
+            .training(lr=0.01, grad_clip=30.0, train_batch_size_per_learner=50)
+        )
+        config = config.learners(num_learners=1)
+        config = config.env_runners(num_env_runners=1)
+        config = config.environment("CartPole-v1")
+
+        # Build an Algorithm object from the config and run 1 training iteration.
+        algo = config.build()
+        algo.train()
+        del algo
+
+    .. testcode::
+
+        from ray.rllib.algorithms.appo import APPOConfig
+        from ray import tune
+
+        config = APPOConfig()
+        # Update the config object.
+        config = config.training(lr=tune.grid_search([0.001,]))
+        # Set the config object's env.
+        config = config.environment(env="CartPole-v1")
+        # Use to_dict() to get the old-style python config dict when running with tune.
+        tune.Tuner(
+            "APPO",
+            run_config=tune.RunConfig(
+                stop={"training_iteration": 1},
+                verbose=0,
+            ),
+            param_space=config.to_dict(),
+
+        ).fit()
+
+    .. testoutput::
+        :hide:
+
+        ...
+    """
+
+    def __init__(self, algo_class=None):
+        """Initializes a APPOConfig instance."""
+        self.exploration_config = {
+            # The Exploration class to use. In the simplest case, this is the name
+            # (str) of any class present in the `rllib.utils.exploration` package.
+            # You can also provide the python class directly or the full location
+            # of your class (e.g. "ray.rllib.utils.exploration.epsilon_greedy.
+            # EpsilonGreedy").
+            "type": "StochasticSampling",
+            # Add constructor kwargs here (if any).
+        }
+
+        super().__init__(algo_class=algo_class or APPO)
+
+        # fmt: off
+        # __sphinx_doc_begin__
+        # APPO specific settings:
+        self.vtrace = True
+        self.use_gae = True
+        self.lambda_ = 1.0
+        self.clip_param = 0.4
+        self.use_kl_loss = False
+        self.kl_coeff = 1.0
+        self.kl_target = 0.01
+        self.target_worker_clipping = 2.0
+
+        # If a circular buffer should be used to store training batches. The
+        # alternative is a simple `Queue`.
+        self.use_circular_buffer = True
+        # Circular replay buffer settings.
+        # Used in [1] for discrete action tasks:
+        # `circular_buffer_num_batches=4` and `circular_buffer_iterations_per_batch=2`
+        # For cont. action tasks:
+        # `circular_buffer_num_batches=16` and `circular_buffer_iterations_per_batch=20`
+        self.circular_buffer_num_batches = 8
+        self.circular_buffer_iterations_per_batch = 2
+
+        # Size of the simple queue (if `use_circular_buffer` is False).
+        self.simple_queue_size = 32
+
+        # Override some of IMPALAConfig's default values with APPO-specific values.
+        self.num_env_runners = 2
+        self.target_network_update_freq = 2
+        self.broadcast_interval = 1
+        self.grad_clip = 40.0
+        # Note: Only when using enable_rl_module_and_learner=True can the clipping mode
+        # be configured by the user. On the old API stack, RLlib will always clip by
+        # global_norm, no matter the value of `grad_clip_by`.
+        self.grad_clip_by = "global_norm"
+
+        self.opt_type = "adam"
+        self.lr = 0.0005
+        self.decay = 0.99
+        self.momentum = 0.0
+        self.epsilon = 0.1
+        self.vf_loss_coeff = 0.5
+        self.entropy_coeff = 0.01
+        self.tau = 1.0
+        # __sphinx_doc_end__
+        # fmt: on
+
+        self.lr_schedule = None  # @OldAPIStack
+        self.entropy_coeff_schedule = None  # @OldAPIStack
+        self.num_gpus = 0  # @OldAPIStack
+        self.num_multi_gpu_tower_stacks = 1  # @OldAPIStack
+        self.minibatch_buffer_size = 1  # @OldAPIStack
+        self.replay_proportion = 0.0  # @OldAPIStack
+        self.replay_buffer_num_slots = 100  # @OldAPIStack
+        self.learner_queue_size = 16  # @OldAPIStack
+        self.learner_queue_timeout = 300  # @OldAPIStack
+
+        # Deprecated keys.
+        self.target_update_frequency = DEPRECATED_VALUE
+        self.use_critic = DEPRECATED_VALUE
+
+    @override(IMPALAConfig)
+    def training(
+        self,
+        *,
+        vtrace: Optional[bool] = NotProvided,
+        use_gae: Optional[bool] = NotProvided,
+        lambda_: Optional[float] = NotProvided,
+        clip_param: Optional[float] = NotProvided,
+        use_kl_loss: Optional[bool] = NotProvided,
+        kl_coeff: Optional[float] = NotProvided,
+        kl_target: Optional[float] = NotProvided,
+        target_network_update_freq: Optional[int] = NotProvided,
+        tau: Optional[float] = NotProvided,
+        target_worker_clipping: Optional[float] = NotProvided,
+        use_circular_buffer: Optional[bool] = NotProvided,
+        circular_buffer_num_batches: Optional[int] = NotProvided,
+        circular_buffer_iterations_per_batch: Optional[int] = NotProvided,
+        simple_queue_size: Optional[int] = NotProvided,
+        # Deprecated keys.
+        target_update_frequency=DEPRECATED_VALUE,
+        use_critic=DEPRECATED_VALUE,
+        **kwargs,
+    ) -> Self:
+        """Sets the training related configuration.
+
+        Args:
+            vtrace: Whether to use V-trace weighted advantages. If false, PPO GAE
+                advantages will be used instead.
+            use_gae: If true, use the Generalized Advantage Estimator (GAE)
+                with a value function, see https://arxiv.org/pdf/1506.02438.pdf.
+                Only applies if vtrace=False.
+            lambda_: GAE (lambda) parameter.
+            clip_param: PPO surrogate slipping parameter.
+            use_kl_loss: Whether to use the KL-term in the loss function.
+            kl_coeff: Coefficient for weighting the KL-loss term.
+            kl_target: Target term for the KL-term to reach (via adjusting the
+                `kl_coeff` automatically).
+            target_network_update_freq: NOTE: This parameter is only applicable on
+                the new API stack. The frequency with which to update the target
+                policy network from the main trained policy network. The metric
+                used is `NUM_ENV_STEPS_TRAINED_LIFETIME` and the unit is `n` (see [1]
+                4.1.1), where: `n = [circular_buffer_num_batches (N)] *
+                [circular_buffer_iterations_per_batch (K)] * [train batch size]`
+                For example, if you set `target_network_update_freq=2`, and N=4, K=2,
+                and `train_batch_size_per_learner=500`, then the target net is updated
+                every 2*4*2*500=8000 trained env steps (every 16 batch updates on each
+                learner).
+                The authors in [1] suggests that this setting is robust to a range of
+                choices (try values between 0.125 and 4).
+            target_network_update_freq: The frequency to update the target policy and
+                tune the kl loss coefficients that are used during training. After
+                setting this parameter, the algorithm waits for at least
+                `target_network_update_freq` number of environment samples to be trained
+                on before updating the target networks and tune the kl loss
+                coefficients. NOTE: This parameter is only applicable when using the
+                Learner API (enable_rl_module_and_learner=True).
+            tau: The factor by which to update the target policy network towards
+                the current policy network. Can range between 0 and 1.
+                e.g. updated_param = tau * current_param + (1 - tau) * target_param
+            target_worker_clipping: The maximum value for the target-worker-clipping
+                used for computing the IS ratio, described in [1]
+                IS = min(π(i) / π(target), ρ) * (π / π(i))
+            use_circular_buffer: Whether to use a circular buffer for storing
+                training batches. If false, a simple Queue will be used. Defaults to
+                True.
+            circular_buffer_num_batches: The number of train batches that fit
+                into the circular buffer. Each such train batch can be sampled for
+                training max. `circular_buffer_iterations_per_batch` times.
+            circular_buffer_iterations_per_batch: The number of times any train
+                batch in the circular buffer can be sampled for training. A batch gets
+                evicted from the buffer either if it's the oldest batch in the buffer
+                and a new batch is added OR if the batch reaches this max. number of
+                being sampled.
+            simple_queue_size: The size of the simple queue (if `use_circular_buffer`
+                is False) for storing training batches.
+
+        Returns:
+            This updated AlgorithmConfig object.
+        """
+        if target_update_frequency != DEPRECATED_VALUE:
+            deprecation_warning(
+                old="target_update_frequency",
+                new="target_network_update_freq",
+                error=True,
+            )
+        if use_critic != DEPRECATED_VALUE:
+            deprecation_warning(
+                old="use_critic",
+                help="`use_critic` no longer supported! APPO always uses a value "
+                "function (critic).",
+                error=True,
+            )
+
+        # Pass kwargs onto super's `training()` method.
+        super().training(**kwargs)
+
+        if vtrace is not NotProvided:
+            self.vtrace = vtrace
+        if use_gae is not NotProvided:
+            self.use_gae = use_gae
+        if lambda_ is not NotProvided:
+            self.lambda_ = lambda_
+        if clip_param is not NotProvided:
+            self.clip_param = clip_param
+        if use_kl_loss is not NotProvided:
+            self.use_kl_loss = use_kl_loss
+        if kl_coeff is not NotProvided:
+            self.kl_coeff = kl_coeff
+        if kl_target is not NotProvided:
+            self.kl_target = kl_target
+        if target_network_update_freq is not NotProvided:
+            self.target_network_update_freq = target_network_update_freq
+        if tau is not NotProvided:
+            self.tau = tau
+        if target_worker_clipping is not NotProvided:
+            self.target_worker_clipping = target_worker_clipping
+        if use_circular_buffer is not NotProvided:
+            self.use_circular_buffer = use_circular_buffer
+        if circular_buffer_num_batches is not NotProvided:
+            self.circular_buffer_num_batches = circular_buffer_num_batches
+        if circular_buffer_iterations_per_batch is not NotProvided:
+            self.circular_buffer_iterations_per_batch = (
+                circular_buffer_iterations_per_batch
+            )
+        if simple_queue_size is not NotProvided:
+            self.simple_queue_size = simple_queue_size
+
+        return self
+
+    @override(IMPALAConfig)
+    def validate(self) -> None:
+        super().validate()
+
+        # On new API stack, circular buffer should be used, not `minibatch_buffer_size`.
+        if self.enable_rl_module_and_learner:
+            if self.minibatch_buffer_size != 1 or self.replay_proportion != 0.0:
+                self._value_error(
+                    "`minibatch_buffer_size/replay_proportion` not valid on new API "
+                    "stack with APPO! "
+                    "Use `circular_buffer_num_batches` for the number of train batches "
+                    "in the circular buffer. To change the maximum number of times "
+                    "any batch may be sampled, set "
+                    "`circular_buffer_iterations_per_batch`."
+                )
+            if self.num_multi_gpu_tower_stacks != 1:
+                self._value_error(
+                    "`num_multi_gpu_tower_stacks` not supported on new API stack with "
+                    "APPO! In order to train on multi-GPU, use "
+                    "`config.learners(num_learners=[number of GPUs], "
+                    "num_gpus_per_learner=1)`. To scale the throughput of batch-to-GPU-"
+                    "pre-loading on each of your `Learners`, set "
+                    "`num_gpu_loader_threads` to a higher number (recommended values: "
+                    "1-8)."
+                )
+            if self.learner_queue_size != 16:
+                self._value_error(
+                    "`learner_queue_size` not supported on new API stack with "
+                    "APPO! In order set the size of the circular buffer (which acts as "
+                    "a 'learner queue'), use "
+                    "`config.training(circular_buffer_num_batches=..)`. To change the "
+                    "maximum number of times any batch may be sampled, set "
+                    "`config.training(circular_buffer_iterations_per_batch=..)`."
+                )
+
+    @override(IMPALAConfig)
+    def get_default_learner_class(self):
+        if self.framework_str == "torch":
+            from ray.rllib.algorithms.appo.torch.appo_torch_learner import (
+                APPOTorchLearner,
+            )
+
+            return APPOTorchLearner
+        elif self.framework_str in ["tf2", "tf"]:
+            raise ValueError(
+                "TensorFlow is no longer supported on the new API stack! "
+                "Use `framework='torch'`."
+            )
+        else:
+            raise ValueError(
+                f"The framework {self.framework_str} is not supported. "
+                "Use `framework='torch'`."
+            )
+
+    @override(IMPALAConfig)
+    def get_default_rl_module_spec(self) -> RLModuleSpec:
+        if self.framework_str == "torch":
+            from ray.rllib.algorithms.appo.torch.appo_torch_rl_module import (
+                APPOTorchRLModule as RLModule,
+            )
+        else:
+            raise ValueError(
+                f"The framework {self.framework_str} is not supported. "
+                "Use either 'torch' or 'tf2'."
+            )
+
+        return RLModuleSpec(module_class=RLModule)
+
+    @property
+    @override(AlgorithmConfig)
+    def _model_config_auto_includes(self):
+        return super()._model_config_auto_includes | {"vf_share_layers": False}
+
+
+class APPO(IMPALA):
+    def __init__(self, config, *args, **kwargs):
+        """Initializes an APPO instance."""
+        super().__init__(config, *args, **kwargs)
+
+        # After init: Initialize target net.
+
+        # TODO(avnishn): Does this need to happen in __init__? I think we can move it
+        #  to setup()
+        if not self.config.enable_rl_module_and_learner:
+            self.env_runner.foreach_policy_to_train(lambda p, _: p.update_target())
+
+    @override(IMPALA)
+    def training_step(self) -> None:
+        if self.config.enable_rl_module_and_learner:
+            return super().training_step()
+
+        train_results = super().training_step()
+        # Update the target network and the KL coefficient for the APPO-loss.
+        # The target network update frequency is calculated automatically by the product
+        # of `num_epochs` setting (usually 1 for APPO) and `minibatch_buffer_size`.
+        last_update = self._counters[LAST_TARGET_UPDATE_TS]
+        cur_ts = self._counters[
+            (
+                NUM_AGENT_STEPS_SAMPLED
+                if self.config.count_steps_by == "agent_steps"
+                else NUM_ENV_STEPS_SAMPLED
+            )
+        ]
+        target_update_freq = self.config.num_epochs * self.config.minibatch_buffer_size
+        if cur_ts - last_update > target_update_freq:
+            self._counters[NUM_TARGET_UPDATES] += 1
+            self._counters[LAST_TARGET_UPDATE_TS] = cur_ts
+
+            # Update our target network.
+            self.env_runner.foreach_policy_to_train(lambda p, _: p.update_target())
+
+            # Also update the KL-coefficient for the APPO loss, if necessary.
+            if self.config.use_kl_loss:
+
+                def update(pi, pi_id):
+                    assert LEARNER_STATS_KEY not in train_results, (
+                        "{} should be nested under policy id key".format(
+                            LEARNER_STATS_KEY
+                        ),
+                        train_results,
+                    )
+                    if pi_id in train_results:
+                        kl = train_results[pi_id][LEARNER_STATS_KEY].get("kl")
+                        assert kl is not None, (train_results, pi_id)
+                        # Make the actual `Policy.update_kl()` call.
+                        pi.update_kl(kl)
+                    else:
+                        logger.warning("No data for {}, not updating kl".format(pi_id))
+
+                # Update KL on all trainable policies within the local (trainer)
+                # Worker.
+                self.env_runner.foreach_policy_to_train(update)
+
+        return train_results
+
+    @classmethod
+    @override(IMPALA)
+    def get_default_config(cls) -> APPOConfig:
+        return APPOConfig()
+
+    @classmethod
+    @override(IMPALA)
+    def get_default_policy_class(
+        cls, config: AlgorithmConfig
+    ) -> Optional[Type[Policy]]:
+        if config["framework"] == "torch":
+            from ray.rllib.algorithms.appo.appo_torch_policy import APPOTorchPolicy
+
+            return APPOTorchPolicy
+        elif config["framework"] == "tf":
+            if config.enable_rl_module_and_learner:
+                raise ValueError(
+                    "RLlib's RLModule and Learner API is not supported for"
+                    " tf1. Use "
+                    "framework='tf2' instead."
+                )
+            from ray.rllib.algorithms.appo.appo_tf_policy import APPOTF1Policy
+
+            return APPOTF1Policy
+        else:
+            from ray.rllib.algorithms.appo.appo_tf_policy import APPOTF2Policy
+
+            return APPOTF2Policy

rllib/algorithms/appo/appo_learner.py

@@ -0,0 +1,163 @@
+import abc
+from collections import defaultdict
+from queue import Queue
+from typing import Any, Dict, Optional
+
+from ray.rllib.algorithms.appo.appo import APPOConfig
+from ray.rllib.algorithms.appo.utils import CircularBuffer
+from ray.rllib.algorithms.impala.impala_learner import IMPALALearner
+from ray.rllib.core.learner.learner import Learner
+from ray.rllib.core.learner.utils import update_target_network
+from ray.rllib.core.rl_module.apis import TargetNetworkAPI, ValueFunctionAPI
+from ray.rllib.core.rl_module.multi_rl_module import MultiRLModuleSpec
+from ray.rllib.core.rl_module.rl_module import RLModuleSpec
+from ray.rllib.utils.annotations import override
+from ray.rllib.utils.lambda_defaultdict import LambdaDefaultDict
+from ray.rllib.utils.metrics import (
+    LAST_TARGET_UPDATE_TS,
+    NUM_ENV_STEPS_TRAINED_LIFETIME,
+    NUM_MODULE_STEPS_TRAINED,
+    NUM_TARGET_UPDATES,
+)
+from ray.rllib.utils.schedules.scheduler import Scheduler
+from ray.rllib.utils.typing import ModuleID, ShouldModuleBeUpdatedFn
+
+
+class APPOLearner(IMPALALearner):
+    """Adds KL coeff updates via `after_gradient_based_update()` to IMPALA logic.
+
+    Framework-specific subclasses must override `_update_module_kl_coeff()`.
+    """
+
+    @override(IMPALALearner)
+    def build(self):
+        self._last_update_ts_by_mid = defaultdict(int)
+
+        # Use a CircularBuffer as learner-in-queue if configured to do so.
+        if self.config.use_circular_buffer:
+            self._learner_thread_in_queue = CircularBuffer(
+                num_batches=self.config.circular_buffer_num_batches,
+                iterations_per_batch=self.config.circular_buffer_iterations_per_batch,
+            )
+        # Otherwise, use a simple Queue.
+        else:
+            # For APPO use a large queue.
+            self._learner_thread_in_queue = Queue(maxsize=self.config.simple_queue_size)
+
+        # Now build the super class. Otherwise the learner-queue would be overridden.
+        super().build()
+
+        # Make target networks.
+        self.module.foreach_module(
+            lambda mid, mod: (
+                mod.make_target_networks()
+                if isinstance(mod, TargetNetworkAPI)
+                else None
+            )
+        )
+
+        # The current kl coefficients per module as (framework specific) tensor
+        # variables.
+        self.curr_kl_coeffs_per_module: LambdaDefaultDict[
+            ModuleID, Scheduler
+        ] = LambdaDefaultDict(
+            lambda module_id: self._get_tensor_variable(
+                self.config.get_config_for_module(module_id).kl_coeff
+            )
+        )
+
+    @override(Learner)
+    def add_module(
+        self,
+        *,
+        module_id: ModuleID,
+        module_spec: RLModuleSpec,
+        config_overrides: Optional[Dict] = None,
+        new_should_module_be_updated: Optional[ShouldModuleBeUpdatedFn] = None,
+    ) -> MultiRLModuleSpec:
+        marl_spec = super().add_module(
+            module_id=module_id,
+            module_spec=module_spec,
+            config_overrides=config_overrides,
+            new_should_module_be_updated=new_should_module_be_updated,
+        )
+        # Create target networks for added Module, if applicable.
+        if isinstance(self.module[module_id].unwrapped(), TargetNetworkAPI):
+            self.module[module_id].unwrapped().make_target_networks()
+        return marl_spec
+
+    @override(IMPALALearner)
+    def remove_module(self, module_id: str) -> MultiRLModuleSpec:
+        marl_spec = super().remove_module(module_id)
+        self.curr_kl_coeffs_per_module.pop(module_id)
+        return marl_spec
+
+    @override(Learner)
+    def after_gradient_based_update(self, *, timesteps: Dict[str, Any]) -> None:
+        """Updates the target Q Networks."""
+        super().after_gradient_based_update(timesteps=timesteps)
+
+        # TODO (sven): Maybe we should have a `after_gradient_based_update`
+        #  method per module?
+        curr_timestep = timesteps.get(NUM_ENV_STEPS_TRAINED_LIFETIME, 0)
+        for module_id, module in self.module._rl_modules.items():
+            config = self.config.get_config_for_module(module_id)
+
+            if isinstance(module.unwrapped(), TargetNetworkAPI) and (
+                curr_timestep - self._last_update_ts_by_mid[module_id]
+                >= (
+                    config.target_network_update_freq
+                    * config.circular_buffer_num_batches
+                    * config.circular_buffer_iterations_per_batch
+                    * config.train_batch_size_per_learner
+                )
+            ):
+                for (
+                    main_net,
+                    target_net,
+                ) in module.unwrapped().get_target_network_pairs():
+                    update_target_network(
+                        main_net=main_net,
+                        target_net=target_net,
+                        tau=config.tau,
+                    )
+                # Increase lifetime target network update counter by one.
+                self.metrics.log_value(
+                    (module_id, NUM_TARGET_UPDATES), 1, reduce="lifetime_sum"
+                )
+
+                # Update the (single-value -> window=1) last updated timestep metric.
+                self._last_update_ts_by_mid[module_id] = curr_timestep
+                self.metrics.log_value(
+                    (module_id, LAST_TARGET_UPDATE_TS), curr_timestep, reduce="max"
+                )
+
+            if (
+                config.use_kl_loss
+                and self.metrics.peek((module_id, NUM_MODULE_STEPS_TRAINED), default=0)
+                > 0
+            ):
+                self._update_module_kl_coeff(module_id=module_id, config=config)
+
+    @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 [TargetNetworkAPI, ValueFunctionAPI]
+
+    @abc.abstractmethod
+    def _update_module_kl_coeff(self, module_id: ModuleID, config: APPOConfig) -> None:
+        """Dynamically update the KL loss coefficients of each module.
+
+        The update is completed using the mean KL divergence between the action
+        distributions current policy and old policy of each module. That action
+        distribution is computed during the most recent update/call to `compute_loss`.
+
+        Args:
+            module_id: The module whose KL loss coefficient to update.
+            config: The AlgorithmConfig specific to the given `module_id`.
+        """
+
+
+AppoLearner = APPOLearner

rllib/algorithms/appo/appo_rl_module.py

@@ -0,0 +1,11 @@
+# Backward compat import.
+from ray.rllib.algorithms.appo.default_appo_rl_module import (  # noqa
+    DefaultAPPORLModule as APPORLModule,
+)
+from ray._common.deprecation import deprecation_warning
+
+deprecation_warning(
+    old="ray.rllib.algorithms.appo.appo_rl_module.APPORLModule",
+    new="ray.rllib.algorithms.appo.default_appo_rl_module.DefaultAPPORLModule",
+    error=False,
+)

rllib/algorithms/appo/appo_tf_policy.py

@@ -0,0 +1,393 @@
+"""
+TensorFlow policy class used for APPO.
+
+Adapted from VTraceTFPolicy to use the PPO surrogate loss.
+Keep in sync with changes to VTraceTFPolicy.
+"""
+
+import logging
+from typing import Dict, List, Optional, Type, Union
+
+import gymnasium as gym
+import numpy as np
+
+from ray.rllib.algorithms.appo.utils import make_appo_models
+from ray.rllib.algorithms.impala import vtrace_tf as vtrace
+from ray.rllib.algorithms.impala.impala_tf_policy import (
+    VTraceClipGradients,
+    VTraceOptimizer,
+    _make_time_major,
+)
+from ray.rllib.evaluation.postprocessing import (
+    Postprocessing,
+    compute_bootstrap_value,
+    compute_gae_for_sample_batch,
+)
+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,
+    KLCoeffMixin,
+    LearningRateSchedule,
+    TargetNetworkMixin,
+    ValueNetworkMixin,
+)
+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 TensorType
+
+tf1, tf, tfv = try_import_tf()
+
+logger = logging.getLogger(__name__)
+
+
+# TODO (sven): Deprecate once APPO and IMPALA fully on RLModules/Learner APIs.
+def get_appo_tf_policy(name: str, base: type) -> type:
+    """Construct an APPOTFPolicy 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.
+    """
+
+    class APPOTFPolicy(
+        VTraceClipGradients,
+        VTraceOptimizer,
+        LearningRateSchedule,
+        KLCoeffMixin,
+        EntropyCoeffSchedule,
+        ValueNetworkMixin,
+        TargetNetworkMixin,
+        GradStatsMixin,
+        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()
+
+            # Although this is a no-op, we call __init__ here to make it clear
+            # that base.__init__ will use the make_model() call.
+            VTraceClipGradients.__init__(self)
+            VTraceOptimizer.__init__(self)
+
+            # Initialize base class.
+            base.__init__(
+                self,
+                observation_space,
+                action_space,
+                config,
+                existing_inputs=existing_inputs,
+                existing_model=existing_model,
+            )
+
+            # TF LearningRateSchedule depends on self.framework, so initialize
+            # after base.__init__() is called.
+            LearningRateSchedule.__init__(self, config["lr"], config["lr_schedule"])
+            EntropyCoeffSchedule.__init__(
+                self, config["entropy_coeff"], config["entropy_coeff_schedule"]
+            )
+            ValueNetworkMixin.__init__(self, config)
+            KLCoeffMixin.__init__(self, config)
+
+            GradStatsMixin.__init__(self)
+
+            # 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()
+
+            # Initiate TargetNetwork ops after loss initialization.
+            TargetNetworkMixin.__init__(self)
+
+        @override(base)
+        def make_model(self) -> ModelV2:
+            return make_appo_models(self)
+
+        @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.multi_discrete.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_logits = train_batch[SampleBatch.ACTION_DIST_INPUTS]
+
+            target_model_out, _ = self.target_model(train_batch)
+            prev_action_dist = dist_class(behaviour_logits, self.model)
+            values = self.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])
+                mask = make_time_major(mask)
+
+                def reduce_mean_valid(t):
+                    return tf.reduce_mean(tf.boolean_mask(t, mask))
+
+            else:
+                reduce_mean_valid = tf.reduce_mean
+
+            if self.config["vtrace"]:
+                logger.debug("Using V-Trace surrogate loss (vtrace=True)")
+
+                # Prepare actions for loss.
+                loss_actions = (
+                    actions if is_multidiscrete else tf.expand_dims(actions, axis=1)
+                )
+
+                old_policy_behaviour_logits = tf.stop_gradient(target_model_out)
+                old_policy_action_dist = dist_class(old_policy_behaviour_logits, model)
+
+                # Prepare KL for Loss
+                mean_kl = make_time_major(old_policy_action_dist.multi_kl(action_dist))
+
+                unpacked_behaviour_logits = tf.split(
+                    behaviour_logits, output_hidden_shape, axis=1
+                )
+                unpacked_old_policy_behaviour_logits = tf.split(
+                    old_policy_behaviour_logits, output_hidden_shape, axis=1
+                )
+
+                # Compute vtrace on the CPU for better perf.
+                with tf.device("/cpu:0"):
+                    vtrace_returns = vtrace.multi_from_logits(
+                        behaviour_policy_logits=make_time_major(
+                            unpacked_behaviour_logits
+                        ),
+                        target_policy_logits=make_time_major(
+                            unpacked_old_policy_behaviour_logits
+                        ),
+                        actions=tf.unstack(make_time_major(loss_actions), axis=2),
+                        discounts=tf.cast(
+                            ~make_time_major(tf.cast(dones, tf.bool)),
+                            tf.float32,
+                        )
+                        * 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,
+                        clip_rho_threshold=tf.cast(
+                            self.config["vtrace_clip_rho_threshold"], tf.float32
+                        ),
+                        clip_pg_rho_threshold=tf.cast(
+                            self.config["vtrace_clip_pg_rho_threshold"], tf.float32
+                        ),
+                    )
+
+                actions_logp = make_time_major(action_dist.logp(actions))
+                prev_actions_logp = make_time_major(prev_action_dist.logp(actions))
+                old_policy_actions_logp = make_time_major(
+                    old_policy_action_dist.logp(actions)
+                )
+
+                is_ratio = tf.clip_by_value(
+                    tf.math.exp(prev_actions_logp - old_policy_actions_logp), 0.0, 2.0
+                )
+                logp_ratio = is_ratio * tf.exp(actions_logp - prev_actions_logp)
+                self._is_ratio = is_ratio
+
+                advantages = vtrace_returns.pg_advantages
+                surrogate_loss = tf.minimum(
+                    advantages * logp_ratio,
+                    advantages
+                    * tf.clip_by_value(
+                        logp_ratio,
+                        1 - self.config["clip_param"],
+                        1 + self.config["clip_param"],
+                    ),
+                )
+
+                action_kl = (
+                    tf.reduce_mean(mean_kl, axis=0) if is_multidiscrete else mean_kl
+                )
+                mean_kl_loss = reduce_mean_valid(action_kl)
+                mean_policy_loss = -reduce_mean_valid(surrogate_loss)
+
+                # The value function loss.
+                value_targets = vtrace_returns.vs
+                delta = values_time_major - value_targets
+                mean_vf_loss = 0.5 * reduce_mean_valid(tf.math.square(delta))
+
+                # The entropy loss.
+                actions_entropy = make_time_major(action_dist.multi_entropy())
+                mean_entropy = reduce_mean_valid(actions_entropy)
+
+            else:
+                logger.debug("Using PPO surrogate loss (vtrace=False)")
+
+                # Prepare KL for Loss
+                mean_kl = make_time_major(prev_action_dist.multi_kl(action_dist))
+
+                logp_ratio = tf.math.exp(
+                    make_time_major(action_dist.logp(actions))
+                    - make_time_major(prev_action_dist.logp(actions))
+                )
+
+                advantages = make_time_major(train_batch[Postprocessing.ADVANTAGES])
+                surrogate_loss = tf.minimum(
+                    advantages * logp_ratio,
+                    advantages
+                    * tf.clip_by_value(
+                        logp_ratio,
+                        1 - self.config["clip_param"],
+                        1 + self.config["clip_param"],
+                    ),
+                )
+
+                action_kl = (
+                    tf.reduce_mean(mean_kl, axis=0) if is_multidiscrete else mean_kl
+                )
+                mean_kl_loss = reduce_mean_valid(action_kl)
+                mean_policy_loss = -reduce_mean_valid(surrogate_loss)
+
+                # The value function loss.
+                value_targets = make_time_major(
+                    train_batch[Postprocessing.VALUE_TARGETS]
+                )
+                delta = values_time_major - value_targets
+                mean_vf_loss = 0.5 * reduce_mean_valid(tf.math.square(delta))
+
+                # The entropy loss.
+                mean_entropy = reduce_mean_valid(
+                    make_time_major(action_dist.multi_entropy())
+                )
+
+            # The summed weighted loss.
+            total_loss = mean_policy_loss - mean_entropy * self.entropy_coeff
+            # Optional KL loss.
+            if self.config["use_kl_loss"]:
+                total_loss += self.kl_coeff * mean_kl_loss
+            # Optional vf loss (or in a separate term due to separate
+            # optimizers/networks).
+            loss_wo_vf = total_loss
+            if not self.config["_separate_vf_optimizer"]:
+                total_loss += mean_vf_loss * self.config["vf_loss_coeff"]
+
+            # Store stats in policy for stats_fn.
+            self._total_loss = total_loss
+            self._loss_wo_vf = loss_wo_vf
+            self._mean_policy_loss = mean_policy_loss
+            # Backward compatibility: Deprecate policy._mean_kl.
+            self._mean_kl_loss = self._mean_kl = mean_kl_loss
+            self._mean_vf_loss = mean_vf_loss
+            self._mean_entropy = mean_entropy
+            self._value_targets = value_targets
+
+            # Return one total loss or two losses: vf vs rest (policy + kl).
+            if self.config["_separate_vf_optimizer"]:
+                return loss_wo_vf, mean_vf_loss
+            else:
+                return 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(),
+            )
+
+            stats_dict = {
+                "cur_lr": tf.cast(self.cur_lr, tf.float64),
+                "total_loss": self._total_loss,
+                "policy_loss": self._mean_policy_loss,
+                "entropy": self._mean_entropy,
+                "var_gnorm": tf.linalg.global_norm(self.model.trainable_variables()),
+                "vf_loss": self._mean_vf_loss,
+                "vf_explained_var": explained_variance(
+                    tf.reshape(self._value_targets, [-1]),
+                    tf.reshape(values_batched, [-1]),
+                ),
+                "entropy_coeff": tf.cast(self.entropy_coeff, tf.float64),
+            }
+
+            if self.config["vtrace"]:
+                is_stat_mean, is_stat_var = tf.nn.moments(self._is_ratio, [0, 1])
+                stats_dict["mean_IS"] = is_stat_mean
+                stats_dict["var_IS"] = is_stat_var
+
+            if self.config["use_kl_loss"]:
+                stats_dict["kl"] = self._mean_kl_loss
+                stats_dict["KL_Coeff"] = self.kl_coeff
+
+            return stats_dict
+
+        @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 not self.config["vtrace"]:
+                sample_batch = compute_gae_for_sample_batch(
+                    self, sample_batch, other_agent_batches, episode
+                )
+            else:
+                # Add the Columns.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"]
+
+    APPOTFPolicy.__name__ = name
+    APPOTFPolicy.__qualname__ = name
+
+    return APPOTFPolicy
+
+
+APPOTF1Policy = get_appo_tf_policy("APPOTF1Policy", DynamicTFPolicyV2)
+APPOTF2Policy = get_appo_tf_policy("APPOTF2Policy", EagerTFPolicyV2)

rllib/algorithms/appo/appo_torch_policy.py

@@ -0,0 +1,413 @@
+"""
+PyTorch policy class used for APPO.
+
+Adapted from VTraceTFPolicy to use the PPO surrogate loss.
+Keep in sync with changes to VTraceTFPolicy.
+"""
+
+import logging
+from typing import Any, Dict, List, Optional, Type, Union
+
+import gymnasium as gym
+import numpy as np
+
+import ray
+import ray.rllib.algorithms.impala.vtrace_torch as vtrace
+from ray.rllib.algorithms.appo.utils import make_appo_models
+from ray.rllib.algorithms.impala.impala_torch_policy import (
+    VTraceOptimizer,
+    make_time_major,
+)
+from ray.rllib.evaluation.postprocessing import (
+    Postprocessing,
+    compute_bootstrap_value,
+    compute_gae_for_sample_batch,
+)
+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,
+    TorchDistributionWrapper,
+)
+from ray.rllib.models.torch.torch_modelv2 import TorchModelV2
+from ray.rllib.policy.sample_batch import SampleBatch
+from ray.rllib.policy.torch_mixins import (
+    EntropyCoeffSchedule,
+    KLCoeffMixin,
+    LearningRateSchedule,
+    TargetNetworkMixin,
+    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__)
+
+
+# TODO (sven): Deprecate once APPO and IMPALA fully on RLModules/Learner APIs.
+class APPOTorchPolicy(
+    VTraceOptimizer,
+    LearningRateSchedule,
+    EntropyCoeffSchedule,
+    KLCoeffMixin,
+    ValueNetworkMixin,
+    TargetNetworkMixin,
+    TorchPolicyV2,
+):
+    """PyTorch policy class used with APPO."""
+
+    def __init__(self, observation_space, action_space, config):
+        config = dict(ray.rllib.algorithms.appo.appo.APPOConfig().to_dict(), **config)
+        config["enable_rl_module_and_learner"] = False
+        config["enable_env_runner_and_connector_v2"] = False
+
+        # Although this is a no-op, we call __init__ here to make it clear
+        # that base.__init__ will use the make_model() call.
+        VTraceOptimizer.__init__(self)
+
+        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
+        )
+
+        TorchPolicyV2.__init__(
+            self,
+            observation_space,
+            action_space,
+            config,
+            max_seq_len=config["model"]["max_seq_len"],
+        )
+
+        EntropyCoeffSchedule.__init__(
+            self, config["entropy_coeff"], config["entropy_coeff_schedule"]
+        )
+        ValueNetworkMixin.__init__(self, config)
+        KLCoeffMixin.__init__(self, config)
+
+        self._initialize_loss_from_dummy_batch()
+
+        # Initiate TargetNetwork ops after loss initialization.
+        TargetNetworkMixin.__init__(self)
+
+    @override(TorchPolicyV2)
+    def init_view_requirements(self):
+        self.view_requirements = self._get_default_view_requirements()
+
+    @override(TorchPolicyV2)
+    def make_model(self) -> ModelV2:
+        return make_appo_models(self)
+
+    @override(TorchPolicyV2)
+    def loss(
+        self,
+        model: ModelV2,
+        dist_class: Type[ActionDistribution],
+        train_batch: SampleBatch,
+    ) -> Union[TensorType, List[TensorType]]:
+        """Constructs the loss for APPO.
+
+        With IS modifications and V-trace for Advantage Estimation.
+
+        Args:
+            model (ModelV2): The Model to calculate the loss for.
+            dist_class (Type[ActionDistribution]): The action distr. class.
+            train_batch: The training data.
+
+        Returns:
+            Union[TensorType, List[TensorType]]: A single loss tensor or a list
+                of loss tensors.
+        """
+        target_model = self.target_models[model]
+
+        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.multi_discrete.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, **kwargs):
+            return make_time_major(
+                self, train_batch.get(SampleBatch.SEQ_LENS), *args, **kwargs
+            )
+
+        actions = train_batch[SampleBatch.ACTIONS]
+        dones = train_batch[SampleBatch.TERMINATEDS]
+        rewards = train_batch[SampleBatch.REWARDS]
+        behaviour_logits = train_batch[SampleBatch.ACTION_DIST_INPUTS]
+
+        target_model_out, _ = target_model(train_batch)
+
+        prev_action_dist = dist_class(behaviour_logits, model)
+        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 = sequence_mask(train_batch[SampleBatch.SEQ_LENS], max_seq_len)
+            mask = torch.reshape(mask, [-1])
+            mask = _make_time_major(mask)
+            num_valid = torch.sum(mask)
+
+            def reduce_mean_valid(t):
+                return torch.sum(t[mask]) / num_valid
+
+        else:
+            reduce_mean_valid = torch.mean
+
+        if self.config["vtrace"]:
+            logger.debug("Using V-Trace surrogate loss (vtrace=True)")
+
+            old_policy_behaviour_logits = target_model_out.detach()
+            old_policy_action_dist = dist_class(old_policy_behaviour_logits, model)
+
+            if isinstance(output_hidden_shape, (list, tuple, np.ndarray)):
+                unpacked_behaviour_logits = torch.split(
+                    behaviour_logits, list(output_hidden_shape), dim=1
+                )
+                unpacked_old_policy_behaviour_logits = torch.split(
+                    old_policy_behaviour_logits, list(output_hidden_shape), dim=1
+                )
+            else:
+                unpacked_behaviour_logits = torch.chunk(
+                    behaviour_logits, output_hidden_shape, dim=1
+                )
+                unpacked_old_policy_behaviour_logits = torch.chunk(
+                    old_policy_behaviour_logits, output_hidden_shape, dim=1
+                )
+
+            # Prepare actions for loss.
+            loss_actions = (
+                actions if is_multidiscrete else torch.unsqueeze(actions, dim=1)
+            )
+
+            # Prepare KL for loss.
+            action_kl = _make_time_major(old_policy_action_dist.kl(action_dist))
+
+            # Compute vtrace on the CPU for better perf.
+            vtrace_returns = vtrace.multi_from_logits(
+                behaviour_policy_logits=_make_time_major(unpacked_behaviour_logits),
+                target_policy_logits=_make_time_major(
+                    unpacked_old_policy_behaviour_logits
+                ),
+                actions=torch.unbind(_make_time_major(loss_actions), dim=2),
+                discounts=(1.0 - _make_time_major(dones).float())
+                * 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,
+                clip_rho_threshold=self.config["vtrace_clip_rho_threshold"],
+                clip_pg_rho_threshold=self.config["vtrace_clip_pg_rho_threshold"],
+            )
+
+            actions_logp = _make_time_major(action_dist.logp(actions))
+            prev_actions_logp = _make_time_major(prev_action_dist.logp(actions))
+            old_policy_actions_logp = _make_time_major(
+                old_policy_action_dist.logp(actions)
+            )
+            is_ratio = torch.clamp(
+                torch.exp(prev_actions_logp - old_policy_actions_logp), 0.0, 2.0
+            )
+            logp_ratio = is_ratio * torch.exp(actions_logp - prev_actions_logp)
+            self._is_ratio = is_ratio
+
+            advantages = vtrace_returns.pg_advantages.to(logp_ratio.device)
+            surrogate_loss = torch.min(
+                advantages * logp_ratio,
+                advantages
+                * torch.clamp(
+                    logp_ratio,
+                    1 - self.config["clip_param"],
+                    1 + self.config["clip_param"],
+                ),
+            )
+
+            mean_kl_loss = reduce_mean_valid(action_kl)
+            mean_policy_loss = -reduce_mean_valid(surrogate_loss)
+
+            # The value function loss.
+            value_targets = vtrace_returns.vs.to(values_time_major.device)
+            delta = values_time_major - value_targets
+            mean_vf_loss = 0.5 * reduce_mean_valid(torch.pow(delta, 2.0))
+
+            # The entropy loss.
+            mean_entropy = reduce_mean_valid(_make_time_major(action_dist.entropy()))
+
+        else:
+            logger.debug("Using PPO surrogate loss (vtrace=False)")
+
+            # Prepare KL for Loss
+            action_kl = _make_time_major(prev_action_dist.kl(action_dist))
+
+            actions_logp = _make_time_major(action_dist.logp(actions))
+            prev_actions_logp = _make_time_major(prev_action_dist.logp(actions))
+            logp_ratio = torch.exp(actions_logp - prev_actions_logp)
+
+            advantages = _make_time_major(train_batch[Postprocessing.ADVANTAGES])
+            surrogate_loss = torch.min(
+                advantages * logp_ratio,
+                advantages
+                * torch.clamp(
+                    logp_ratio,
+                    1 - self.config["clip_param"],
+                    1 + self.config["clip_param"],
+                ),
+            )
+
+            mean_kl_loss = reduce_mean_valid(action_kl)
+            mean_policy_loss = -reduce_mean_valid(surrogate_loss)
+
+            # The value function loss.
+            value_targets = _make_time_major(train_batch[Postprocessing.VALUE_TARGETS])
+            delta = values_time_major - value_targets
+            mean_vf_loss = 0.5 * reduce_mean_valid(torch.pow(delta, 2.0))
+
+            # The entropy loss.
+            mean_entropy = reduce_mean_valid(_make_time_major(action_dist.entropy()))
+
+        # The summed weighted loss.
+        total_loss = mean_policy_loss - mean_entropy * self.entropy_coeff
+        # Optional additional KL Loss
+        if self.config["use_kl_loss"]:
+            total_loss += self.kl_coeff * mean_kl_loss
+
+        # Optional vf loss (or in a separate term due to separate
+        # optimizers/networks).
+        loss_wo_vf = total_loss
+        if not self.config["_separate_vf_optimizer"]:
+            total_loss += mean_vf_loss * self.config["vf_loss_coeff"]
+
+        # 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["total_loss"] = total_loss
+        model.tower_stats["mean_policy_loss"] = mean_policy_loss
+        model.tower_stats["mean_kl_loss"] = mean_kl_loss
+        model.tower_stats["mean_vf_loss"] = mean_vf_loss
+        model.tower_stats["mean_entropy"] = mean_entropy
+        model.tower_stats["value_targets"] = value_targets
+        model.tower_stats["vf_explained_var"] = explained_variance(
+            torch.reshape(value_targets, [-1]),
+            torch.reshape(values_time_major, [-1]),
+        )
+
+        # Return one total loss or two losses: vf vs rest (policy + kl).
+        if self.config["_separate_vf_optimizer"]:
+            return loss_wo_vf, mean_vf_loss
+        else:
+            return total_loss
+
+    @override(TorchPolicyV2)
+    def stats_fn(self, train_batch: SampleBatch) -> Dict[str, TensorType]:
+        """Stats function for APPO. Returns a dict with important loss stats.
+
+        Args:
+            policy: The Policy to generate stats for.
+            train_batch: The SampleBatch (already) used for training.
+
+        Returns:
+            Dict[str, TensorType]: The stats dict.
+        """
+        stats_dict = {
+            "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("mean_policy_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("mean_vf_loss"))),
+            "vf_explained_var": torch.mean(
+                torch.stack(self.get_tower_stats("vf_explained_var"))
+            ),
+        }
+
+        if self.config["vtrace"]:
+            is_stat_mean = torch.mean(self._is_ratio, [0, 1])
+            is_stat_var = torch.var(self._is_ratio, [0, 1])
+            stats_dict["mean_IS"] = is_stat_mean
+            stats_dict["var_IS"] = is_stat_var
+
+        if self.config["use_kl_loss"]:
+            stats_dict["kl"] = torch.mean(
+                torch.stack(self.get_tower_stats("mean_kl_loss"))
+            )
+            stats_dict["KL_Coeff"] = self.kl_coeff
+
+        return convert_to_numpy(stats_dict)
+
+    @override(TorchPolicyV2)
+    def extra_action_out(
+        self,
+        input_dict: Dict[str, TensorType],
+        state_batches: List[TensorType],
+        model: TorchModelV2,
+        action_dist: TorchDistributionWrapper,
+    ) -> Dict[str, TensorType]:
+        return {SampleBatch.VF_PREDS: model.value_function()}
+
+    @override(TorchPolicyV2)
+    def postprocess_trajectory(
+        self,
+        sample_batch: SampleBatch,
+        other_agent_batches: Optional[Dict[Any, SampleBatch]] = None,
+        episode=None,
+    ):
+        # Call super's postprocess_trajectory first.
+        # sample_batch = super().postprocess_trajectory(
+        #    sample_batch, other_agent_batches, episode
+        # )
+
+        # Do all post-processing always with no_grad().
+        # Not using this here will introduce a memory leak
+        # in torch (issue #6962).
+        with torch.no_grad():
+            if not self.config["vtrace"]:
+                sample_batch = compute_gae_for_sample_batch(
+                    self, sample_batch, other_agent_batches, episode
+                )
+            else:
+                # 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/appo/default_appo_rl_module.py

@@ -0,0 +1,56 @@
+import abc
+from typing import Any, Dict, List, Tuple
+
+from ray.rllib.algorithms.ppo.default_ppo_rl_module import DefaultPPORLModule
+from ray.rllib.core.learner.utils import make_target_network
+from ray.rllib.core.models.base import ACTOR, ENCODER_OUT
+from ray.rllib.core.rl_module.apis import (
+    TARGET_NETWORK_ACTION_DIST_INPUTS,
+    TargetNetworkAPI,
+)
+from ray.rllib.utils.annotations import (
+    OverrideToImplementCustomLogic_CallToSuperRecommended,
+    override,
+)
+from ray.rllib.utils.typing import NetworkType
+from ray.util.annotations import DeveloperAPI
+
+
+@DeveloperAPI
+class DefaultAPPORLModule(DefaultPPORLModule, TargetNetworkAPI, abc.ABC):
+    """Default RLModule used by APPO, if user does not specify a custom RLModule.
+
+    Users who want to train their RLModules with APPO may implement any RLModule
+    (or TorchRLModule) subclass as long as the custom class also implements the
+    `ValueFunctionAPI` (see ray.rllib.core.rl_module.apis.value_function_api.py)
+    and the `TargetNetworkAPI` (see
+    ray.rllib.core.rl_module.apis.target_network_api.py).
+    """
+
+    @override(TargetNetworkAPI)
+    def make_target_networks(self):
+        self._old_encoder = make_target_network(self.encoder)
+        self._old_pi = make_target_network(self.pi)
+
+    @override(TargetNetworkAPI)
+    def get_target_network_pairs(self) -> List[Tuple[NetworkType, NetworkType]]:
+        return [
+            (self.encoder, self._old_encoder),
+            (self.pi, self._old_pi),
+        ]
+
+    @override(TargetNetworkAPI)
+    def forward_target(self, batch: Dict[str, Any]) -> Dict[str, Any]:
+        old_pi_inputs_encoded = self._old_encoder(batch)[ENCODER_OUT][ACTOR]
+        old_action_dist_logits = self._old_pi(old_pi_inputs_encoded)
+        return {TARGET_NETWORK_ACTION_DIST_INPUTS: old_action_dist_logits}
+
+    @OverrideToImplementCustomLogic_CallToSuperRecommended
+    @override(DefaultPPORLModule)
+    def get_non_inference_attributes(self) -> List[str]:
+        # Get the NON inference-only attributes from the parent class.
+        ret = super().get_non_inference_attributes()
+        # Add the two (APPO) target networks to it (NOT needed in
+        # inference-only mode).
+        ret += ["_old_encoder", "_old_pi"]
+        return ret

rllib/algorithms/appo/tests/test_appo.py

@@ -0,0 +1,274 @@
+import unittest
+
+import ray
+import ray.rllib.algorithms.appo as appo
+from ray.rllib.algorithms.impala.impala import LEARNER_RESULTS_CURR_ENTROPY_COEFF_KEY
+from ray.rllib.core.rl_module.default_model_config import DefaultModelConfig
+from ray.rllib.policy.sample_batch import DEFAULT_POLICY_ID
+from ray.rllib.utils.metrics import (
+    ENV_RUNNER_RESULTS,
+    LEARNER_RESULTS,
+    NUM_ENV_STEPS_SAMPLED_LIFETIME,
+)
+from ray.rllib.utils.test_utils import (
+    check_compute_single_action,
+    check_train_results,
+    check_train_results_new_api_stack,
+)
+
+
+class TestAPPO(unittest.TestCase):
+    @classmethod
+    def setUpClass(cls):
+        ray.init()
+
+    @classmethod
+    def tearDownClass(cls):
+        ray.shutdown()
+
+    def test_appo_compilation(self):
+        """Test whether APPO can be built with both frameworks."""
+        config = (
+            appo.APPOConfig().environment("CartPole-v1").env_runners(num_env_runners=1)
+        )
+        algo = config.build()
+
+        num_iterations = 2
+        for i in range(num_iterations):
+            results = algo.train()
+            print(results)
+        check_train_results_new_api_stack(results)
+
+        algo.stop()
+
+    def test_appo_compilation_use_kl_loss(self):
+        """Test whether APPO can be built with kl_loss enabled."""
+        config = (
+            appo.APPOConfig().env_runners(num_env_runners=1).training(use_kl_loss=True)
+        )
+        num_iterations = 2
+
+        algo = config.build(env="CartPole-v1")
+        for i in range(num_iterations):
+            results = algo.train()
+            print(results)
+        check_train_results_new_api_stack(results)
+        algo.stop()
+
+    def test_appo_two_optimizers_two_lrs(self):
+        # Not explicitly setting this should cause a warning, but not fail.
+        # config["_tf_policy_handles_more_than_one_loss"] = True
+        config = (
+            appo.APPOConfig()
+            .api_stack(
+                enable_rl_module_and_learner=False,
+                enable_env_runner_and_connector_v2=False,
+            )
+            .env_runners(num_env_runners=1)
+            .training(
+                _separate_vf_optimizer=True,
+                _lr_vf=0.002,
+                # Make sure we have two completely separate models for policy and
+                # value function.
+                model={
+                    "vf_share_layers": False,
+                },
+            )
+        )
+
+        num_iterations = 2
+
+        # Only supported for tf so far.
+        algo = config.build(env="CartPole-v1")
+        for i in range(num_iterations):
+            results = algo.train()
+            check_train_results(results)
+            print(results)
+        check_compute_single_action(algo)
+        algo.stop()
+
+    def test_appo_entropy_coeff_schedule(self):
+        config = (
+            appo.APPOConfig()
+            .environment("CartPole-v1")
+            .env_runners(
+                num_env_runners=1,
+                rollout_fragment_length=10,
+            )
+            .training(
+                train_batch_size_per_learner=20,
+                entropy_coeff=[
+                    [0, 0.1],
+                    [50000, 0.01],
+                ],
+            )
+            .reporting(
+                min_train_timesteps_per_iteration=20,
+                # 0 metrics reporting delay, this makes sure timestep,
+                # which entropy coeff depends on, is updated after each worker rollout.
+                min_time_s_per_iteration=0,
+            )
+        )
+
+        def _step_n_times(algo, n: int):
+            """Step Algorithm n times.
+
+            Returns:
+                learning rate at the end of the execution.
+            """
+            for _ in range(n):
+                results = algo.train()
+                print(results)
+            return results[LEARNER_RESULTS][DEFAULT_POLICY_ID][
+                LEARNER_RESULTS_CURR_ENTROPY_COEFF_KEY
+            ]
+
+        algo = config.build()
+
+        coeff = _step_n_times(algo, 10)
+        # Should be close to the starting coeff of 0.01.
+        ts_sampled = algo.metrics.peek(
+            (ENV_RUNNER_RESULTS, NUM_ENV_STEPS_SAMPLED_LIFETIME)
+        )
+        expected_coeff = 0.1 - ((0.1 - 0.01) / 50000 * ts_sampled)
+        self.assertLessEqual(coeff, expected_coeff + 0.005)
+        self.assertGreaterEqual(coeff, expected_coeff - 0.005)
+
+        coeff = _step_n_times(algo, 20)
+        ts_sampled = algo.metrics.peek(
+            (ENV_RUNNER_RESULTS, NUM_ENV_STEPS_SAMPLED_LIFETIME)
+        )
+        expected_coeff = 0.1 - ((0.1 - 0.01) / 50000 * ts_sampled)
+        self.assertLessEqual(coeff, expected_coeff + 0.005)
+        self.assertGreaterEqual(coeff, expected_coeff - 0.005)
+
+        algo.stop()
+
+    def test_appo_learning_rate_schedule(self):
+        config = (
+            appo.APPOConfig()
+            .env_runners(
+                num_env_runners=1,
+                batch_mode="truncate_episodes",
+                rollout_fragment_length=10,
+            )
+            .training(
+                train_batch_size_per_learner=20,
+                entropy_coeff=0.01,
+                # Setup lr schedule for testing.
+                lr=[[0, 5e-2], [50000, 0.0]],
+            )
+            .reporting(
+                min_train_timesteps_per_iteration=20,
+                # 0 metrics reporting delay, this makes sure timestep,
+                # which entropy coeff depends on, is updated after each worker rollout.
+                min_time_s_per_iteration=0,
+            )
+        )
+
+        def _step_n_times(algo, n: int):
+            """Step Algorithm n times.
+
+            Returns:
+                learning rate at the end of the execution.
+            """
+            for _ in range(n):
+                results = algo.train()
+                print(results)
+            return results[LEARNER_RESULTS][DEFAULT_POLICY_ID][
+                "default_optimizer_learning_rate"
+            ]
+
+        algo = config.build(env="CartPole-v1")
+
+        lr1 = _step_n_times(algo, 10)
+        lr2 = _step_n_times(algo, 10)
+
+        self.assertGreater(lr1, lr2)
+
+        algo.stop()
+
+    def test_appo_model_variables(self):
+        config = (
+            appo.APPOConfig()
+            .environment("CartPole-v1")
+            .env_runners(
+                num_env_runners=1,
+                rollout_fragment_length=10,
+            )
+            .training(
+                train_batch_size_per_learner=20,
+            )
+            .rl_module(
+                model_config=DefaultModelConfig(
+                    fcnet_hiddens=[16],
+                    vf_share_layers=True,
+                ),
+            )
+        )
+
+        algo = config.build()
+        state = algo.get_module(DEFAULT_POLICY_ID).get_state()
+        # Weights and biases for the encoder hidden layer (2) and the output layer
+        # of the policy (2), plus the `log_std_clip` param (1), makes 5 altogether.
+        # We should not get the tensors from the target model here or any value function
+        # related parameters (inference-only).
+        self.assertEqual(len(state), 5)
+
+    def test_env_runner_state_server_on_vs_off(self):
+        """PULL-based EnvRunnerStateServer: APPO learns with the flag ON and OFF.
+
+        Also checks the global server actor is created only when the flag is enabled.
+        """
+        for use_server in [False, True]:
+            print(f"Testing with use_server={use_server}")
+            config = (
+                appo.APPOConfig()
+                .environment("CartPole-v1")
+                .env_runners(
+                    num_env_runners=2,
+                    use_env_runner_state_server=use_server,
+                )
+            )
+            algo = config.build()
+            # The global server actor exists iff the flag is enabled.
+            self.assertEqual(algo._env_runner_state_server is not None, use_server)
+
+            results = algo.train()
+            check_train_results_new_api_stack(results)
+            algo.stop()
+
+    def test_env_runner_state_server_kill_and_recover(self):
+        """Killing the EnvRunnerStateServer must not stop training; it recovers."""
+        config = (
+            appo.APPOConfig()
+            .environment("CartPole-v1")
+            .env_runners(num_env_runners=2, use_env_runner_state_server=True)
+        )
+        algo = config.build()
+        self.assertIsNotNone(algo._env_runner_state_server)
+
+        for _ in range(3):
+            algo.train()
+        version_before = ray.get(algo._env_runner_state_server.get_version.remote())
+        self.assertGreater(version_before, 0)
+
+        # Kill the server. `max_restarts=-1` makes Ray restart it (with empty state).
+        ray.kill(algo._env_runner_state_server, no_restart=False)
+
+        # Training continues through the gap and the next push re-seeds the server.
+        for _ in range(3):
+            results = algo.train()
+        check_train_results_new_api_stack(results)
+        version_after = ray.get(algo._env_runner_state_server.get_version.remote())
+        self.assertGreaterEqual(version_after, version_before)
+
+        algo.stop()
+
+
+if __name__ == "__main__":
+    import sys
+
+    import pytest
+
+    sys.exit(pytest.main(["-v", __file__]))

rllib/algorithms/appo/tests/test_appo_learner.py

@@ -0,0 +1,125 @@
+import unittest
+
+import numpy as np
+import tree  # pip install dm_tree
+
+import ray
+import ray.rllib.algorithms.appo as appo
+from ray.rllib.algorithms.appo.appo import LEARNER_RESULTS_CURR_KL_COEFF_KEY
+from ray.rllib.core import DEFAULT_MODULE_ID
+from ray.rllib.core.columns import Columns
+from ray.rllib.core.rl_module.default_model_config import DefaultModelConfig
+from ray.rllib.policy.sample_batch import SampleBatch
+from ray.rllib.utils.metrics import LEARNER_RESULTS
+from ray.rllib.utils.torch_utils import convert_to_torch_tensor
+
+frag_length = 50
+
+FAKE_BATCH = {
+    Columns.OBS: np.random.uniform(low=0, high=1, size=(frag_length, 4)).astype(
+        np.float32
+    ),
+    Columns.ACTIONS: np.random.choice(2, frag_length).astype(np.float32),
+    Columns.REWARDS: np.random.uniform(low=-1, high=1, size=(frag_length,)).astype(
+        np.float32
+    ),
+    Columns.TERMINATEDS: np.array(
+        [False for _ in range(frag_length - 1)] + [True]
+    ).astype(np.float32),
+    Columns.VF_PREDS: np.array(list(reversed(range(frag_length))), dtype=np.float32),
+    Columns.ACTION_LOGP: np.log(
+        np.random.uniform(low=0, high=1, size=(frag_length,))
+    ).astype(np.float32),
+    Columns.LOSS_MASK: np.ones(shape=(frag_length,)),
+}
+
+
+class TestAPPOLearner(unittest.TestCase):
+    @classmethod
+    def setUpClass(cls):
+        ray.init()
+
+    @classmethod
+    def tearDownClass(cls):
+        ray.shutdown()
+
+    def test_appo_loss(self):
+        """Test that appo_policy_rlm loss matches the appo learner loss."""
+        config = (
+            appo.APPOConfig()
+            .environment("CartPole-v1")
+            .env_runners(
+                num_env_runners=0,
+                rollout_fragment_length=frag_length,
+            )
+            .training(
+                gamma=0.99,
+                model=dict(
+                    fcnet_hiddens=[10, 10],
+                    fcnet_activation="linear",
+                    vf_share_layers=False,
+                ),
+            )
+        )
+        # We have to set exploration_config here manually because setting it through
+        # config.env_runners() only deep-updates it
+        config.exploration_config = {}
+
+        algo = config.build()
+
+        train_batch = SampleBatch(
+            tree.map_structure(lambda x: convert_to_torch_tensor(x), FAKE_BATCH)
+        )
+
+        algo_config = config.copy(copy_frozen=False)
+        algo_config.learners(num_learners=0).experimental(_validate_config=False)
+        algo_config.validate()
+
+        learner_group = algo_config.build_learner_group(env=algo.env_runner.env)
+        learner_group.update(batch=train_batch.as_multi_agent())
+
+        algo.stop()
+
+    def test_kl_coeff_changes(self):
+        initial_kl_coeff = 0.01
+        config = (
+            appo.APPOConfig()
+            .environment("CartPole-v1")
+            .env_runners(
+                num_env_runners=0,
+                rollout_fragment_length=frag_length,
+                exploration_config={},
+            )
+            .learners(num_learners=0)
+            .experimental(_validate_config=False)
+            .training(
+                use_kl_loss=True,
+                kl_coeff=initial_kl_coeff,
+            )
+            .rl_module(
+                model_config=DefaultModelConfig(
+                    fcnet_hiddens=[10, 10],
+                    fcnet_activation="linear",
+                    vf_share_layers=False,
+                ),
+            )
+        )
+        algo = config.build()
+        # Call train while results aren't returned because this is
+        # a asynchronous algorithm and results are returned asynchronously.
+        curr_kl_coeff = None
+        while curr_kl_coeff is None:
+            results = algo.train()
+            print(results)
+            results = results.get(LEARNER_RESULTS, {})
+            results = results.get(DEFAULT_MODULE_ID, {})
+            curr_kl_coeff = results.get(LEARNER_RESULTS_CURR_KL_COEFF_KEY)
+        self.assertNotEqual(curr_kl_coeff, initial_kl_coeff)
+
+
+if __name__ == "__main__":
+    import sys
+
+    import pytest
+
+    sys.exit(pytest.main(["-v", __file__]))

rllib/algorithms/appo/tests/test_appo_multi_agent_data_balance.py

@@ -0,0 +1,74 @@
+"""Test how APPO handles per-policy data imbalance in multi-agent setups.
+
+Note: PPO will always use "equalize" data across policies. So each policy will train on the same amount of data.
+APPO, in contrast to PPO, will train on varying amounts of data per policy.
+
+When a policy_mapping_fn maps more agents to one policy than another, the resulting
+MultiAgentBatch has unequal per-policy data. This test verifies:
+1. Default APPO (no minibatch_size): policies train on unequal amounts of data.
+2. With minibatch_size set: MiniBatchCyclicIterator equalizes per-policy batch sizes.
+"""
+
+from ray.rllib.algorithms.appo import APPOConfig
+from ray.rllib.examples.envs.classes.multi_agent import MultiAgentCartPole
+from ray.rllib.utils.metrics import NUM_MODULE_STEPS_TRAINED
+
+NUM_AGENTS = 5
+
+
+def policy_mapping_fn(agent_id, episode, **kw):
+    return "policy_a" if agent_id in (0, 1, 2, 3) else "policy_b"
+
+
+def _build_config(*, minibatch_size=None, num_epochs=1):
+    config = (
+        APPOConfig()
+        .environment(MultiAgentCartPole, env_config={"num_agents": NUM_AGENTS})
+        .multi_agent(
+            policies={"policy_a", "policy_b"},
+            policy_mapping_fn=policy_mapping_fn,
+        )
+        .training(
+            train_batch_size_per_learner=500,
+        )
+        .learners(num_learners=0)
+        .env_runners(num_env_runners=1)
+    )
+    if minibatch_size is not None:
+        config.training(minibatch_size=minibatch_size, num_epochs=num_epochs)
+    return config
+
+
+def test_default_appo_unequal_data():
+    """Without minibatch_size, policy_a trains on more data than policy_b."""
+    algo = _build_config().build()
+
+    learners = algo.train()["learners"]
+    steps_a = learners["policy_a"][NUM_MODULE_STEPS_TRAINED]
+    steps_b = learners["policy_b"][NUM_MODULE_STEPS_TRAINED]
+    # steps_a should be 4x more data than steps_b
+    assert steps_a / steps_b > 2.5, (
+        "Expected policy_a to train on more data than policy_b "
+        "with biased policy mapping and no minibatch_size."
+    )
+
+
+def test_minibatch_equalizes_data():
+    """With minibatch_size, both policies train on equal amounts of data."""
+    algo = _build_config(minibatch_size=50, num_epochs=4).build()
+
+    learners = algo.train()["learners"]
+    steps_a = learners["policy_a"][NUM_MODULE_STEPS_TRAINED]
+    steps_b = learners["policy_b"][NUM_MODULE_STEPS_TRAINED]
+    assert steps_a == steps_b, (
+        "Expected equal per-policy training steps when "
+        "minibatch_size is set (MiniBatchCyclicIterator)."
+    )
+
+
+if __name__ == "__main__":
+    import sys
+
+    import pytest
+
+    sys.exit(pytest.main(["-v", __file__]))

rllib/algorithms/appo/torch/appo_torch_learner.py

@@ -0,0 +1,257 @@
+"""Asynchronous Proximal Policy Optimization (APPO)
+
+The algorithm is described in [1] (under the name of "IMPACT"):
+
+Detailed documentation:
+https://docs.ray.io/en/master/rllib-algorithms.html#appo
+
+[1] IMPACT: Importance Weighted Asynchronous Architectures with Clipped Target Networks.
+Luo et al. 2020
+https://arxiv.org/pdf/1912.00167
+"""
+from typing import Dict
+
+from ray.rllib.algorithms.appo.appo import (
+    LEARNER_RESULTS_CURR_KL_COEFF_KEY,
+    LEARNER_RESULTS_KL_KEY,
+    LEARNER_RESULTS_MEAN_IS_KEY,
+    LEARNER_RESULTS_VAR_IS_KEY,
+    APPOConfig,
+)
+from ray.rllib.algorithms.appo.appo_learner import APPOLearner
+from ray.rllib.algorithms.impala.torch.impala_torch_learner import IMPALATorchLearner
+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, POLICY_LOSS_KEY, VF_LOSS_KEY
+from ray.rllib.core.rl_module.apis import (
+    TARGET_NETWORK_ACTION_DIST_INPUTS,
+    TargetNetworkAPI,
+    ValueFunctionAPI,
+)
+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.typing import ModuleID, TensorType
+
+torch, nn = try_import_torch()
+
+
+class APPOTorchLearner(APPOLearner, IMPALATorchLearner):
+    """Implements APPO loss / update logic on top of IMPALATorchLearner."""
+
+    @override(IMPALATorchLearner)
+    def compute_loss_for_module(
+        self,
+        *,
+        module_id: ModuleID,
+        config: APPOConfig,
+        batch: Dict,
+        fwd_out: Dict[str, TensorType],
+    ) -> TensorType:
+        module = self.module[module_id].unwrapped()
+        assert isinstance(module, TargetNetworkAPI)
+        assert isinstance(module, ValueFunctionAPI)
+
+        # 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)
+
+        values = module.compute_values(
+            batch, embeddings=fwd_out.get(Columns.EMBEDDINGS)
+        )
+
+        action_dist_cls_train = module.get_train_action_dist_cls()
+        target_policy_dist = action_dist_cls_train.from_logits(
+            fwd_out[Columns.ACTION_DIST_INPUTS]
+        )
+
+        old_target_policy_dist = action_dist_cls_train.from_logits(
+            module.forward_target(batch)[TARGET_NETWORK_ACTION_DIST_INPUTS]
+        )
+        old_target_policy_actions_logp = old_target_policy_dist.logp(
+            batch[Columns.ACTIONS]
+        )
+        behaviour_actions_logp = batch[Columns.ACTION_LOGP]
+        target_actions_logp = target_policy_dist.logp(batch[Columns.ACTIONS])
+
+        behaviour_actions_logp_time_major = make_time_major(
+            behaviour_actions_logp,
+            trajectory_len=rollout_frag_or_episode_len,
+            recurrent_seq_len=recurrent_seq_len,
+        )
+        target_actions_logp_time_major = make_time_major(
+            target_actions_logp,
+            trajectory_len=rollout_frag_or_episode_len,
+            recurrent_seq_len=recurrent_seq_len,
+        )
+        old_actions_logp_time_major = make_time_major(
+            old_target_policy_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,
+        )
+        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,
+        )
+
+        # 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,
+                ).float()
+            )
+            * 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=old_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_pg_rho_threshold=config.vtrace_clip_pg_rho_threshold,
+            clip_rho_threshold=config.vtrace_clip_rho_threshold,
+        )
+        pg_advantages = pg_advantages * loss_mask_time_major
+
+        # The policy gradients loss.
+        is_ratio = torch.clip(
+            torch.exp(behaviour_actions_logp_time_major - old_actions_logp_time_major),
+            0.0,
+            2.0,
+        )
+        logp_ratio = is_ratio * torch.exp(
+            target_actions_logp_time_major - behaviour_actions_logp_time_major
+        )
+
+        surrogate_loss = torch.minimum(
+            pg_advantages * logp_ratio,
+            pg_advantages
+            * torch.clip(logp_ratio, 1 - config.clip_param, 1 + config.clip_param),
+        )
+
+        # IS-ratio diagnostics; restricted to valid timesteps via loss mask.
+        mean_is_ratio = torch.sum(is_ratio * loss_mask_time_major) / size_loss_mask
+        var_is_ratio = (
+            torch.sum(torch.square(is_ratio - mean_is_ratio) * loss_mask_time_major)
+            / size_loss_mask
+        )
+
+        if config.use_kl_loss:
+            action_kl = old_target_policy_dist.kl(target_policy_dist) * loss_mask
+            mean_kl_loss = torch.sum(action_kl) / size_loss_mask
+        else:
+            mean_kl_loss = 0.0
+        mean_pi_loss = -(torch.sum(surrogate_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.
+        mean_entropy_loss = (
+            -torch.sum(target_policy_dist.entropy() * loss_mask) / 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()
+            )
+            + (mean_kl_loss * self.curr_kl_coeffs_per_module[module_id])
+        )
+
+        # Log important loss stats.
+        self.metrics.log_dict(
+            {
+                POLICY_LOSS_KEY: mean_pi_loss,
+                VF_LOSS_KEY: mean_vf_loss,
+                ENTROPY_KEY: -mean_entropy_loss,
+                LEARNER_RESULTS_KL_KEY: mean_kl_loss,
+                LEARNER_RESULTS_CURR_KL_COEFF_KEY: (
+                    self.curr_kl_coeffs_per_module[module_id]
+                ),
+                LEARNER_RESULTS_MEAN_IS_KEY: mean_is_ratio,
+                LEARNER_RESULTS_VAR_IS_KEY: var_is_ratio,
+            },
+            key=module_id,
+            window=1,  # <- single items (should not be mean/ema-reduced over time).
+        )
+        # Return the total loss.
+        return total_loss
+
+    @override(APPOLearner)
+    def _update_module_kl_coeff(self, module_id: ModuleID, config: APPOConfig) -> None:
+        # Update the current KL value based on the recently measured value.
+        # Increase.
+        kl = convert_to_numpy(self.metrics.peek((module_id, LEARNER_RESULTS_KL_KEY)))
+        kl_coeff_var = self.curr_kl_coeffs_per_module[module_id]
+
+        if kl > 2.0 * config.kl_target:
+            # TODO (Kourosh) why not *2.0?
+            kl_coeff_var.data *= 1.5
+        # Decrease.
+        elif kl < 0.5 * config.kl_target:
+            kl_coeff_var.data *= 0.5
+
+        self.metrics.log_value(
+            (module_id, LEARNER_RESULTS_CURR_KL_COEFF_KEY),
+            kl_coeff_var.item(),
+            window=1,
+        )

rllib/algorithms/appo/torch/appo_torch_rl_module.py

@@ -0,0 +1,13 @@
+# Backward compat import.
+from ray.rllib.algorithms.appo.torch.default_appo_torch_rl_module import (  # noqa
+    DefaultAPPOTorchRLModule as APPOTorchRLModule,
+)
+from ray._common.deprecation import deprecation_warning
+
+
+deprecation_warning(
+    old="ray.rllib.algorithms.appo.torch.appo_torch_rl_module.APPOTorchRLModule",
+    new="ray.rllib.algorithms.appo.torch.default_appo_torch_rl_module."
+    "DefaultAPPOTorchRLModule",
+    error=False,
+)

rllib/algorithms/appo/torch/default_appo_torch_rl_module.py

@@ -0,0 +1,10 @@
+from ray.rllib.algorithms.appo.default_appo_rl_module import DefaultAPPORLModule
+from ray.rllib.algorithms.ppo.torch.default_ppo_torch_rl_module import (
+    DefaultPPOTorchRLModule,
+)
+from ray.util.annotations import DeveloperAPI
+
+
+@DeveloperAPI
+class DefaultAPPOTorchRLModule(DefaultPPOTorchRLModule, DefaultAPPORLModule):
+    pass

rllib/algorithms/appo/utils.py

@@ -0,0 +1,264 @@
+"""
+[1] IMPACT: Importance Weighted Asynchronous Architectures with Clipped Target Networks.
+Luo et al. 2020
+https://arxiv.org/pdf/1912.00167
+"""
+import threading
+import time
+from collections import deque
+from typing import Any, Optional
+
+import numpy as np
+
+from ray.rllib.models.catalog import ModelCatalog
+from ray.rllib.models.modelv2 import ModelV2
+from ray.rllib.utils.annotations import OldAPIStack
+from ray.rllib.utils.metrics.ray_metrics import (
+    DEFAULT_HISTOGRAM_BOUNDARIES_SHORT_EVENTS,
+    TimerAndPrometheusLogger,
+)
+from ray.util.metrics import Counter, Histogram
+
+POLICY_SCOPE = "func"
+TARGET_POLICY_SCOPE = "target_func"
+
+
+class CircularBuffer:
+    """A circular batch-wise buffer with Queue-like interface.
+
+    The buffer holds at most N batches, which are sampled at random (uniformly).
+    If full and a new batch is added, the oldest batch is discarded. Each batch
+    can be sampled at most K times (after which it is also discarded).
+
+    This version implements Queue-like put/get methods with blocking support.
+    """
+
+    def __init__(self, num_batches: int, iterations_per_batch: int):
+        """
+        Args:
+            num_batches: N from the paper (queue buffer size).
+            iterations_per_batch: K ("replay coefficient") from the paper. Defines
+                how often a single batch can sampled before being discarded. If a
+                new batch is added when the buffer is full, the oldest batch is
+                discarded entirely (regardless of how often it has been sampled).
+        """
+        self.num_batches = num_batches
+        self.iterations_per_batch = iterations_per_batch
+
+        self._NxK = self.num_batches * self.iterations_per_batch
+        self._num_added = 0
+
+        self._buffer = deque([None for _ in range(self._NxK)], maxlen=self._NxK)
+        self._indices = set()
+        self._offset = self._NxK
+        self._lock = threading.Lock()
+
+        # Semaphore tracks the number of *available* samples.
+        self._items_available = threading.Semaphore(0)
+
+        self._rng = np.random.default_rng()
+
+        # Statistics
+        self._total_puts = 0
+        self._total_gets = 0
+        self._total_dropped = 0
+
+        # Ray metrics
+        self._metrics_circular_buffer_put_time = Histogram(
+            name="rllib_utils_circular_buffer_put_time",
+            description="Time spent in CircularBuffer.put()",
+            boundaries=DEFAULT_HISTOGRAM_BOUNDARIES_SHORT_EVENTS,
+            tag_keys=("rllib",),
+        )
+        self._metrics_circular_buffer_put_time.set_default_tags(
+            {"rllib": self.__class__.__name__}
+        )
+
+        self._metrics_circular_buffer_put_ts_dropped = Counter(
+            name="rllib_utils_circular_buffer_put_ts_dropped_counter",
+            description="Total number of env steps dropped by the CircularBuffer.",
+            tag_keys=("rllib",),
+        )
+        self._metrics_circular_buffer_put_ts_dropped.set_default_tags(
+            {"rllib": self.__class__.__name__}
+        )
+
+        self._metrics_circular_buffer_get_time = Histogram(
+            name="rllib_utils_circular_buffer_get_time",
+            description="Time spent in CircularBuffer.get()",
+            boundaries=DEFAULT_HISTOGRAM_BOUNDARIES_SHORT_EVENTS,
+            tag_keys=("rllib",),
+        )
+        self._metrics_circular_buffer_get_time.set_default_tags(
+            {"rllib": self.__class__.__name__}
+        )
+
+    def put(
+        self, item: Any, block: bool = True, timeout: Optional[float] = None
+    ) -> int:
+        """Add a new batch to the buffer.
+
+        The batch is added K times (iterations_per_batch) to allow for K samples.
+        If full, the oldest batch entries are dropped.
+
+        Args:
+            item: The batch to add
+            block: Not used (always non-blocking for puts)
+            timeout: Not used
+
+        Returns:
+            Number of dropped entries (0 or iterations_per_batch)
+        """
+        with TimerAndPrometheusLogger(self._metrics_circular_buffer_put_time):
+            with self._lock:
+                self._total_puts += 1
+
+                # Check if we'll drop old entries
+                dropped_entry = self._buffer[0]
+
+                # Add buffer K times with new indices
+                for _ in range(self.iterations_per_batch):
+                    self._buffer.append(item)
+                    self._indices.add(self._offset)
+                    self._indices.discard(self._offset - self._NxK)
+                    self._offset += 1
+
+                    # Release semaphore for each available sample
+                    self._items_available.release()
+
+                self._num_added += 1
+
+                # A valid entry (w/ a batch whose k has not been reach K yet) was dropped.
+                dropped_ts = 0
+                if dropped_entry is not None:
+                    dropped_ts = (
+                        dropped_entry[0].env_steps()
+                        if isinstance(dropped_entry, tuple)
+                        else dropped_entry.env_steps()
+                    )
+                    if dropped_ts > 0:
+                        self._metrics_circular_buffer_put_ts_dropped.inc(
+                            value=dropped_ts
+                        )
+
+        return dropped_ts
+
+    def put_nowait(self, item: Any) -> int:
+        """Equivalent to self.put(block=False)."""
+        return self.put(item, block=False)
+
+    def get(self, block: bool = True, timeout: Optional[float] = None) -> Any:
+        """Sample a random batch from the buffer.
+
+        The sampled entry is removed and won't be sampled again.
+        Blocks if the buffer is empty (when block=True).
+
+        Args:
+            block: If True, block until an item is available
+            timeout: Maximum time to wait (only used when block=True)
+
+        Returns:
+            A randomly sampled batch
+
+        Raises:
+            TimeoutError: If timeout expires while blocking
+            IndexError: If buffer is empty and block=False
+        """
+        # Only initially, the buffer may be empty -> Just wait for some time.
+        with TimerAndPrometheusLogger(self._metrics_circular_buffer_get_time):
+            while len(self) == 0:
+                time.sleep(0.0001)
+
+            # Sample a random buffer index.
+            with self._lock:
+                idx = self._rng.choice(list(self._indices))
+                actual_buffer_idx = idx - self._offset + self._NxK
+                batch = self._buffer[actual_buffer_idx]
+                assert batch is not None, (
+                    idx,
+                    actual_buffer_idx,
+                    self._offset,
+                    self._indices,
+                    [b is None for b in self._buffer],
+                )
+                self._buffer[actual_buffer_idx] = None
+                self._indices.discard(idx)
+
+        # Return the sampled batch.
+        return batch
+
+    def get_nowait(self) -> Any:
+        """Equivalent to self.get(block=False)."""
+        return self.get(block=False)
+
+    @property
+    def filled(self) -> bool:
+        """Whether the buffer has been filled once with at least `self.num_batches`."""
+        with self._lock:
+            return self._num_added >= self.num_batches
+
+    def qsize(self) -> int:
+        """Returns the number of actually valid (non-expired) batches in the buffer."""
+        with self._lock:
+            return len(self._indices)
+
+    def __len__(self) -> int:
+        return self.qsize()
+
+    def task_done(self):
+        """No-op for Queue compatibility."""
+        pass
+
+    def get_stats(self) -> dict:
+        """Get buffer statistics for monitoring."""
+        with self._lock:
+            return {
+                "size": len(self._indices),
+                "capacity": self._NxK,
+                "num_batches": self.num_batches,
+                "iterations_per_batch": self.iterations_per_batch,
+                "total_puts": self._total_puts,
+                "total_gets": self._total_gets,
+                "total_dropped": self._total_dropped,
+                "filled": self._num_added >= self.num_batches,
+            }
+
+
+@OldAPIStack
+def make_appo_models(policy) -> ModelV2:
+    """Builds model and target model for APPO.
+
+    Returns:
+        ModelV2: The Model for the Policy to use.
+            Note: The target model will not be returned, just assigned to
+            `policy.target_model`.
+    """
+    # Get the num_outputs for the following model construction calls.
+    _, logit_dim = ModelCatalog.get_action_dist(
+        policy.action_space, policy.config["model"]
+    )
+
+    # Construct the (main) model.
+    policy.model = ModelCatalog.get_model_v2(
+        policy.observation_space,
+        policy.action_space,
+        logit_dim,
+        policy.config["model"],
+        name=POLICY_SCOPE,
+        framework=policy.framework,
+    )
+    policy.model_variables = policy.model.variables()
+
+    # Construct the target model.
+    policy.target_model = ModelCatalog.get_model_v2(
+        policy.observation_space,
+        policy.action_space,
+        logit_dim,
+        policy.config["model"],
+        name=TARGET_POLICY_SCOPE,
+        framework=policy.framework,
+    )
+    policy.target_model_variables = policy.target_model.variables()
+
+    # Return only the model (not the target model).
+    return policy.model

rllib/algorithms/bc/__init__.py

@@ -0,0 +1,6 @@
+from ray.rllib.algorithms.bc.bc import BC, BCConfig
+
+__all__ = [
+    "BC",
+    "BCConfig",
+]

rllib/algorithms/bc/bc.py

@@ -0,0 +1,120 @@
+from ray.rllib.algorithms.algorithm_config import AlgorithmConfig
+from ray.rllib.algorithms.marwil.marwil import MARWIL, MARWILConfig
+from ray.rllib.core.rl_module.rl_module import RLModuleSpec
+from ray.rllib.utils.annotations import override
+from ray.rllib.utils.typing import RLModuleSpecType
+
+
+class BCConfig(MARWILConfig):
+    """Defines a configuration class from which a new BC Algorithm can be built
+
+    .. testcode::
+        :skipif: True
+
+        from ray.rllib.algorithms.bc import BCConfig
+        # Run this from the ray directory root.
+        config = BCConfig().training(lr=0.00001, gamma=0.99)
+        config = config.offline_data(
+            input_="./rllib/offline/tests/data/cartpole/large.json")
+
+        # Build an Algorithm object from the config and run 1 training iteration.
+        algo = config.build()
+        algo.train()
+
+    .. testcode::
+        :skipif: True
+
+        from ray.rllib.algorithms.bc import BCConfig
+        from ray import tune
+        config = BCConfig()
+        # Print out some default values.
+        print(config.beta)
+        # Update the config object.
+        config.training(
+            lr=tune.grid_search([0.001, 0.0001]), beta=0.75
+        )
+        # Set the config object's data path.
+        # Run this from the ray directory root.
+        config.offline_data(
+            input_="./rllib/offline/tests/data/cartpole/large.json"
+        )
+        # Set the config object's env, used for evaluation.
+        config.environment(env="CartPole-v1")
+        # Use to_dict() to get the old-style python config dict
+        # when running with tune.
+        tune.Tuner(
+            "BC",
+            param_space=config.to_dict(),
+        ).fit()
+    """
+
+    def __init__(self, algo_class=None):
+        super().__init__(algo_class=algo_class or BC)
+
+        # fmt: off
+        # __sphinx_doc_begin__
+        # No need to calculate advantages (or do anything else with the rewards).
+        self.beta = 0.0
+        # Advantages (calculated during postprocessing)
+        # not important for behavioral cloning.
+        self.postprocess_inputs = False
+
+        # Materialize only the mapped data. This is optimal as long
+        # as no connector in the connector pipeline holds a state.
+        self.materialize_data = False
+        self.materialize_mapped_data = True
+        # __sphinx_doc_end__
+        # fmt: on
+
+    @override(AlgorithmConfig)
+    def get_default_rl_module_spec(self) -> RLModuleSpecType:
+        if self.framework_str == "torch":
+            from ray.rllib.algorithms.bc.torch.default_bc_torch_rl_module import (
+                DefaultBCTorchRLModule,
+            )
+
+            return RLModuleSpec(module_class=DefaultBCTorchRLModule)
+        else:
+            raise ValueError(
+                f"The framework {self.framework_str} is not supported. "
+                "Use `torch` instead."
+            )
+
+    @override(AlgorithmConfig)
+    def build_learner_connector(
+        self,
+        input_observation_space,
+        input_action_space,
+        device=None,
+    ):
+        pipeline = super().build_learner_connector(
+            input_observation_space=input_observation_space,
+            input_action_space=input_action_space,
+            device=device,
+        )
+
+        # Remove unneeded connectors from the MARWIL connector pipeline.
+        pipeline.remove("AddOneTsToEpisodesAndTruncate")
+        pipeline.remove("GeneralAdvantageEstimation")
+
+        return pipeline
+
+    @override(MARWILConfig)
+    def validate(self) -> None:
+        # Call super's validation method.
+        super().validate()
+
+        if self.beta != 0.0:
+            self._value_error("For behavioral cloning, `beta` parameter must be 0.0!")
+
+
+class BC(MARWIL):
+    """Behavioral Cloning (derived from MARWIL).
+
+    Uses MARWIL with beta force-set to 0.0.
+    """
+
+    @classmethod
+    @override(MARWIL)
+    def get_default_config(cls) -> BCConfig:
+        return BCConfig()

rllib/algorithms/bc/bc_catalog.py

@@ -0,0 +1,112 @@
+# __sphinx_doc_begin__
+import gymnasium as gym
+
+from ray.rllib.algorithms.ppo.ppo_catalog import _check_if_diag_gaussian
+from ray.rllib.core.models.base import Model
+from ray.rllib.core.models.catalog import Catalog
+from ray.rllib.core.models.configs import FreeLogStdMLPHeadConfig, MLPHeadConfig
+from ray.rllib.utils.annotations import OverrideToImplementCustomLogic
+
+
+class BCCatalog(Catalog):
+    """The Catalog class used to build models for BC.
+
+    BCCatalog provides the following models:
+        - Encoder: The encoder used to encode the observations.
+        - Pi Head: The head used for the policy logits.
+
+    The default encoder is chosen by RLlib dependent on the observation space.
+    See `ray.rllib.core.models.encoders::Encoder` for details. To define the
+    network architecture use the `model_config_dict[fcnet_hiddens]` and
+    `model_config_dict[fcnet_activation]`.
+
+    To implement custom logic, override `BCCatalog.build_encoder()` or modify the
+    `EncoderConfig` at `BCCatalog.encoder_config`.
+
+    Any custom head can be built by overriding the `build_pi_head()` method.
+    Alternatively, the `PiHeadConfig` can be overridden to build a custom
+    policy head during runtime. To change solely the network architecture,
+    `model_config_dict["head_fcnet_hiddens"]` and
+    `model_config_dict["head_fcnet_activation"]` can be used.
+    """
+
+    def __init__(
+        self,
+        observation_space: gym.Space,
+        action_space: gym.Space,
+        model_config_dict: dict,
+    ):
+        """Initializes the BCCatalog.
+
+        Args:
+            observation_space: The observation space if the Encoder.
+            action_space: The action space for the Pi Head.
+            model_cnfig_dict: The model config to use..
+        """
+        super().__init__(
+            observation_space=observation_space,
+            action_space=action_space,
+            model_config_dict=model_config_dict,
+        )
+
+        self.pi_head_hiddens = self._model_config_dict["head_fcnet_hiddens"]
+        self.pi_head_activation = self._model_config_dict["head_fcnet_activation"]
+
+        # At this time we do not have the precise (framework-specific) action
+        # distribution class, i.e. we do  not know the output dimension of the
+        # policy head. The config for the policy head is therefore build in the
+        # `self.build_pi_head()` method.
+        self.pi_head_config = None
+
+    @OverrideToImplementCustomLogic
+    def build_pi_head(self, framework: str) -> Model:
+        """Builds the policy head.
+
+        The default behavior is to build the head from the pi_head_config.
+        This can be overridden to build a custom policy head as a means of configuring
+        the behavior of a BC specific RLModule implementation.
+
+        Args:
+            framework: The framework to use. Either "torch" or "tf2".
+
+        Returns:
+            The policy head.
+        """
+
+        # Define the output dimension via the action distribution.
+        action_distribution_cls = self.get_action_dist_cls(framework=framework)
+        if self._model_config_dict["free_log_std"]:
+            _check_if_diag_gaussian(
+                action_distribution_cls=action_distribution_cls, framework=framework
+            )
+            is_diag_gaussian = True
+        else:
+            is_diag_gaussian = _check_if_diag_gaussian(
+                action_distribution_cls=action_distribution_cls,
+                framework=framework,
+                no_error=True,
+            )
+        required_output_dim = action_distribution_cls.required_input_dim(
+            space=self.action_space, model_config=self._model_config_dict
+        )
+        # With the action distribution class and the number of outputs defined,
+        # we can build the config for the policy head.
+        pi_head_config_cls = (
+            FreeLogStdMLPHeadConfig
+            if self._model_config_dict["free_log_std"]
+            else MLPHeadConfig
+        )
+        self.pi_head_config = pi_head_config_cls(
+            input_dims=self._latent_dims,
+            hidden_layer_dims=self.pi_head_hiddens,
+            hidden_layer_activation=self.pi_head_activation,
+            output_layer_dim=required_output_dim,
+            output_layer_activation="linear",
+            clip_log_std=is_diag_gaussian,
+            log_std_clip_param=self._model_config_dict.get("log_std_clip_param", 20),
+        )
+
+        return self.pi_head_config.build(framework=framework)
+
+
+# __sphinx_doc_end__

rllib/algorithms/bc/tests/test_bc.py

@@ -0,0 +1,146 @@
+import unittest
+from pathlib import Path
+
+import ray
+from ray.rllib.algorithms.bc import BCConfig
+from ray.rllib.policy.sample_batch import DEFAULT_POLICY_ID
+from ray.rllib.utils.metrics import (
+    ENV_RUNNER_RESULTS,
+    EPISODE_RETURN_MEAN,
+    EVALUATION_RESULTS,
+    LEARNER_RESULTS,
+    NUM_ENV_STEPS_SAMPLED_LIFETIME,
+)
+
+
+class TestBC(unittest.TestCase):
+    @classmethod
+    def setUpClass(cls) -> None:
+        ray.init()
+
+    @classmethod
+    def tearDownClass(cls) -> None:
+        ray.shutdown()
+
+    def test_bc_compilation_and_learning_from_offline_file(self):
+        # Define the data paths.
+        data_path = "offline/tests/data/cartpole/cartpole-v1_large"
+        base_path = Path(__file__).parents[3]
+        print(f"base_path={base_path}")
+        data_path = "local://" / base_path / data_path
+
+        print(f"data_path={data_path}")
+
+        # Define the BC config.
+        config = (
+            BCConfig()
+            .environment(env="CartPole-v1")
+            .learners(
+                num_learners=0,
+            )
+            .evaluation(
+                evaluation_interval=3,
+                evaluation_num_env_runners=1,
+                evaluation_duration=5,
+                evaluation_parallel_to_training=True,
+            )
+            # Note, the `input_` argument is the major argument for the
+            # new offline API.
+            .offline_data(
+                input_=[data_path.as_posix()],
+                dataset_num_iters_per_learner=1,
+            )
+            .training(
+                lr=0.0008,
+                train_batch_size_per_learner=2000,
+            )
+        )
+
+        num_iterations = 350
+        min_return_to_reach = 120.0
+
+        # TODO (simon): Add support for recurrent modules.
+        algo = config.build()
+        learnt = False
+        for i in range(num_iterations):
+            results = algo.train()
+            print(results)
+
+            eval_results = results.get(EVALUATION_RESULTS, {})
+            if eval_results:
+                episode_return_mean = eval_results[ENV_RUNNER_RESULTS][
+                    EPISODE_RETURN_MEAN
+                ]
+                print(f"iter={i}, R={episode_return_mean}")
+                if episode_return_mean > min_return_to_reach:
+                    print("BC has learnt the task!")
+                    learnt = True
+                    break
+
+        if not learnt:
+            raise ValueError(
+                f"`BC` did not reach {min_return_to_reach} reward from "
+                "expert offline data!"
+            )
+
+        algo.stop()
+
+    def test_bc_lr_schedule(self):
+        # Define the data paths.
+        data_path = "offline/tests/data/cartpole/cartpole-v1_large"
+        base_path = Path(__file__).parents[3]
+        data_path = "local://" / base_path / data_path
+
+        config = (
+            BCConfig()
+            .environment(env="CartPole-v1")
+            .learners(
+                num_learners=0,
+            )
+            .evaluation(
+                evaluation_interval=3,
+                evaluation_num_env_runners=1,
+                evaluation_duration=5,
+                evaluation_parallel_to_training=True,
+            )
+            # Note, the `input_` argument is the major argument for the
+            # new offline API.
+            .offline_data(
+                input_=[data_path.as_posix()],
+                dataset_num_iters_per_learner=1,
+            )
+            .training(
+                lr=[
+                    [0, 0.001],
+                    [3000, 0.01],
+                ],
+                train_batch_size_per_learner=2000,
+            )
+        )
+        algo = config.build()
+
+        done = False
+        while not done:
+            results = algo.train()
+            ts = results[NUM_ENV_STEPS_SAMPLED_LIFETIME]
+            assert ts > 0
+            lr = results[LEARNER_RESULTS][DEFAULT_POLICY_ID][
+                "default_optimizer_learning_rate"
+            ]
+            if ts < 3000:
+                # The learning rate should be linearly interpolated.
+                expected_lr = 0.001 + (ts / 3000) * (0.01 - 0.001)
+                self.assertAlmostEqual(lr, expected_lr, places=6)
+            else:
+                self.assertEqual(lr, 0.01)
+                done = True
+
+        algo.stop()
+
+
+if __name__ == "__main__":
+    import sys
+
+    import pytest
+
+    sys.exit(pytest.main(["-v", __file__]))

rllib/algorithms/bc/torch/default_bc_torch_rl_module.py

@@ -0,0 +1,56 @@
+import abc
+from typing import Any, Dict
+
+from ray.rllib.algorithms.bc.bc_catalog import BCCatalog
+from ray.rllib.core.columns import Columns
+from ray.rllib.core.models.base import ENCODER_OUT
+from ray.rllib.core.rl_module.rl_module import RLModule
+from ray.rllib.core.rl_module.torch.torch_rl_module import TorchRLModule
+from ray.rllib.utils.annotations import override
+from ray.util.annotations import DeveloperAPI
+
+
+@DeveloperAPI
+class DefaultBCTorchRLModule(TorchRLModule, abc.ABC):
+    """The default TorchRLModule used, if no custom RLModule is provided.
+
+    Builds an encoder net based on the observation space.
+    Builds a pi head based on the action space.
+
+    Passes observations from the input batch through the encoder, then the pi head to
+    compute action logits.
+    """
+
+    def __init__(self, *args, **kwargs):
+        catalog_class = kwargs.pop("catalog_class", None)
+        if catalog_class is None:
+            catalog_class = BCCatalog
+        super().__init__(*args, **kwargs, catalog_class=catalog_class)
+
+    @override(RLModule)
+    def setup(self):
+        # Build model components (encoder and pi head) from catalog.
+        super().setup()
+        self._encoder = self.catalog.build_encoder(framework=self.framework)
+        self._pi_head = self.catalog.build_pi_head(framework=self.framework)
+
+    @override(TorchRLModule)
+    def _forward(self, batch: Dict, **kwargs) -> Dict[str, Any]:
+        """Generic BC forward pass (for all phases of training/evaluation)."""
+        # Encoder embeddings.
+        encoder_outs = self._encoder(batch)
+        # Action dist inputs.
+        outputs = {Columns.ACTION_DIST_INPUTS: self._pi_head(encoder_outs[ENCODER_OUT])}
+
+        # Add the state if the encoder is stateful.
+        if Columns.STATE_OUT in encoder_outs:
+            outputs[Columns.STATE_OUT] = encoder_outs[Columns.STATE_OUT]
+        # Return the outputs.
+        return outputs
+
+    @override(RLModule)
+    def get_initial_state(self) -> dict:
+        if hasattr(self._encoder, "get_initial_state"):
+            return self._encoder.get_initial_state()
+        else:
+            return {}

rllib/algorithms/callbacks.py

@@ -0,0 +1,7 @@
+# @OldAPIStack
+from ray.rllib.callbacks.callbacks import RLlibCallback
+from ray.rllib.callbacks.utils import _make_multi_callbacks
+
+# Backward compatibility
+DefaultCallbacks = RLlibCallback
+make_multi_callbacks = _make_multi_callbacks

rllib/algorithms/cql/README.md

@@ -0,0 +1,13 @@
+# Conservative Q-Learning (CQL)
+
+## Overview
+
+[CQL](https://arxiv.org/abs/2006.04779) is an offline RL algorithm that mitigates the overestimation of Q-values outside the dataset distribution via convservative critic estimates. CQL does this by adding a simple Q regularizer loss to the standard Belman update loss. This ensures that the critic does not output overly-optimistic Q-values and can be added on top of any off-policy Q-learning algorithm (in this case, we use SAC).
+
+## Documentation & Implementation:
+
+Conservative Q-Learning (CQL).
+
+   **[Detailed Documentation](https://docs.ray.io/en/master/rllib-algorithms.html#cql)**
+
+   **[Implementation](https://github.com/ray-project/ray/blob/master/rllib/algorithms/cql/cql.py)**

rllib/algorithms/cql/__init__.py

@@ -0,0 +1,9 @@
+from ray.rllib.algorithms.cql.cql import CQL, CQLConfig
+from ray.rllib.algorithms.cql.cql_torch_policy import CQLTorchPolicy
+
+__all__ = [
+    "CQL",
+    "CQLConfig",
+    # @OldAPIStack
+    "CQLTorchPolicy",
+]

rllib/algorithms/cql/cql.py

@@ -0,0 +1,378 @@
+import logging
+from typing import Optional, Type, Union
+
+from typing_extensions import Self
+
+from ray._common.deprecation import (
+    DEPRECATED_VALUE,
+    deprecation_warning,
+)
+from ray.rllib.algorithms.algorithm_config import AlgorithmConfig, NotProvided
+from ray.rllib.algorithms.cql.cql_tf_policy import CQLTFPolicy
+from ray.rllib.algorithms.cql.cql_torch_policy import CQLTorchPolicy
+from ray.rllib.algorithms.sac.sac import (
+    SAC,
+    SACConfig,
+)
+from ray.rllib.connectors.common.add_observations_from_episodes_to_batch import (
+    AddObservationsFromEpisodesToBatch,
+)
+from ray.rllib.connectors.learner.add_next_observations_from_episodes_to_train_batch import (  # noqa
+    AddNextObservationsFromEpisodesToTrainBatch,
+)
+from ray.rllib.core.learner.learner import Learner
+from ray.rllib.core.rl_module.rl_module import RLModuleSpec
+from ray.rllib.execution.rollout_ops import (
+    synchronous_parallel_sample,
+)
+from ray.rllib.execution.train_ops import (
+    multi_gpu_train_one_step,
+    train_one_step,
+)
+from ray.rllib.policy.policy import Policy
+from ray.rllib.utils.annotations import OldAPIStack, override
+from ray.rllib.utils.framework import try_import_tf, try_import_tfp
+from ray.rllib.utils.metrics import (
+    LAST_TARGET_UPDATE_TS,
+    LEARNER_RESULTS,
+    LEARNER_UPDATE_TIMER,
+    NUM_AGENT_STEPS_SAMPLED,
+    NUM_AGENT_STEPS_TRAINED,
+    NUM_ENV_STEPS_SAMPLED,
+    NUM_ENV_STEPS_TRAINED,
+    NUM_TARGET_UPDATES,
+    OFFLINE_SAMPLING_TIMER,
+    SAMPLE_TIMER,
+    SYNCH_WORKER_WEIGHTS_TIMER,
+    TARGET_NET_UPDATE_TIMER,
+    TIMERS,
+)
+from ray.rllib.utils.typing import ResultDict, RLModuleSpecType
+
+tf1, tf, tfv = try_import_tf()
+tfp = try_import_tfp()
+logger = logging.getLogger(__name__)
+
+
+class CQLConfig(SACConfig):
+    """Defines a configuration class from which a CQL can be built.
+
+    .. testcode::
+        :skipif: True
+
+        from ray.rllib.algorithms.cql import CQLConfig
+        config = CQLConfig().training(gamma=0.9, lr=0.01)
+        config = config.resources(num_gpus=0)
+        config = config.env_runners(num_env_runners=4)
+        print(config.to_dict())
+        # Build a Algorithm object from the config and run 1 training iteration.
+        algo = config.build(env="CartPole-v1")
+        algo.train()
+    """
+
+    def __init__(self, algo_class=None):
+        super().__init__(algo_class=algo_class or CQL)
+
+        # fmt: off
+        # __sphinx_doc_begin__
+        # CQL-specific config settings:
+        self.bc_iters = 20000
+        self.temperature = 1.0
+        self.num_actions = 10
+        self.lagrangian = False
+        self.lagrangian_thresh = 5.0
+        self.min_q_weight = 5.0
+        self.deterministic_backup = True
+        self.lr = 3e-4
+        # Note, the new stack defines learning rates for each component.
+        # The base learning rate `lr` has to be set to `None`, if using
+        # the new stack.
+        self.actor_lr = 1e-4
+        self.critic_lr = 1e-3
+        self.alpha_lr = 1e-3
+
+        self.replay_buffer_config = {
+            "_enable_replay_buffer_api": True,
+            "type": "MultiAgentPrioritizedReplayBuffer",
+            "capacity": int(1e6),
+            # If True prioritized replay buffer will be used.
+            "prioritized_replay": False,
+            "prioritized_replay_alpha": 0.6,
+            "prioritized_replay_beta": 0.4,
+            "prioritized_replay_eps": 1e-6,
+            # Whether to compute priorities already on the remote worker side.
+            "worker_side_prioritization": False,
+        }
+
+        # Changes to Algorithm's/SACConfig's default:
+
+        # .reporting()
+        self.min_sample_timesteps_per_iteration = 0
+        self.min_train_timesteps_per_iteration = 100
+        # fmt: on
+        # __sphinx_doc_end__
+
+        self.timesteps_per_iteration = DEPRECATED_VALUE
+
+    @override(SACConfig)
+    def training(
+        self,
+        *,
+        bc_iters: Optional[int] = NotProvided,
+        temperature: Optional[float] = NotProvided,
+        num_actions: Optional[int] = NotProvided,
+        lagrangian: Optional[bool] = NotProvided,
+        lagrangian_thresh: Optional[float] = NotProvided,
+        min_q_weight: Optional[float] = NotProvided,
+        deterministic_backup: Optional[bool] = NotProvided,
+        **kwargs,
+    ) -> Self:
+        """Sets the training-related configuration.
+
+        Args:
+            bc_iters: Number of iterations with Behavior Cloning pretraining.
+            temperature: CQL loss temperature.
+            num_actions: Number of actions to sample for CQL loss
+            lagrangian: Whether to use the Lagrangian for Alpha Prime (in CQL loss).
+            lagrangian_thresh: Lagrangian threshold.
+            min_q_weight: in Q weight multiplier.
+            deterministic_backup: If the target in the Bellman update should have an
+                entropy backup. Defaults to `True`.
+
+        Returns:
+            This updated AlgorithmConfig object.
+        """
+        # Pass kwargs onto super's `training()` method.
+        super().training(**kwargs)
+
+        if bc_iters is not NotProvided:
+            self.bc_iters = bc_iters
+        if temperature is not NotProvided:
+            self.temperature = temperature
+        if num_actions is not NotProvided:
+            self.num_actions = num_actions
+        if lagrangian is not NotProvided:
+            self.lagrangian = lagrangian
+        if lagrangian_thresh is not NotProvided:
+            self.lagrangian_thresh = lagrangian_thresh
+        if min_q_weight is not NotProvided:
+            self.min_q_weight = min_q_weight
+        if deterministic_backup is not NotProvided:
+            self.deterministic_backup = deterministic_backup
+
+        return self
+
+    @override(AlgorithmConfig)
+    def offline_data(self, **kwargs) -> Self:
+
+        super().offline_data(**kwargs)
+
+        # Check, if the passed in class incorporates the `OfflinePreLearner`
+        # interface.
+        if "prelearner_class" in kwargs:
+            from ray.rllib.offline.offline_data import OfflinePreLearner
+
+            if not issubclass(kwargs.get("prelearner_class"), OfflinePreLearner):
+                raise ValueError(
+                    f"`prelearner_class` {kwargs.get('prelearner_class')} is not a "
+                    "subclass of `OfflinePreLearner`. Any class passed to "
+                    "`prelearner_class` needs to implement the interface given by "
+                    "`OfflinePreLearner`."
+                )
+
+        return self
+
+    @override(SACConfig)
+    def get_default_learner_class(self) -> Union[Type["Learner"], str]:
+        if self.framework_str == "torch":
+            from ray.rllib.algorithms.cql.torch.cql_torch_learner import CQLTorchLearner
+
+            return CQLTorchLearner
+        else:
+            raise ValueError(
+                f"The framework {self.framework_str} is not supported. "
+                "Use `'torch'` instead."
+            )
+
+    @override(AlgorithmConfig)
+    def build_learner_connector(
+        self,
+        input_observation_space,
+        input_action_space,
+        device=None,
+    ):
+        pipeline = super().build_learner_connector(
+            input_observation_space=input_observation_space,
+            input_action_space=input_action_space,
+            device=device,
+        )
+
+        # Prepend the "add-NEXT_OBS-from-episodes-to-train-batch" connector piece (right
+        # after the corresponding "add-OBS-..." default piece).
+        pipeline.insert_after(
+            AddObservationsFromEpisodesToBatch,
+            AddNextObservationsFromEpisodesToTrainBatch(),
+        )
+
+        return pipeline
+
+    @override(SACConfig)
+    def validate(self) -> None:
+        # First check, whether old `timesteps_per_iteration` is used.
+        if self.timesteps_per_iteration != DEPRECATED_VALUE:
+            deprecation_warning(
+                old="timesteps_per_iteration",
+                new="min_train_timesteps_per_iteration",
+                error=True,
+            )
+
+        # Call super's validation method.
+        super().validate()
+
+        # CQL-torch performs the optimizer steps inside the loss function.
+        # Using the multi-GPU optimizer will therefore not work (see multi-GPU
+        # check above) and we must use the simple optimizer for now.
+        if self.simple_optimizer is not True and self.framework_str == "torch":
+            self.simple_optimizer = True
+
+        if self.framework_str in ["tf", "tf2"] and tfp is None:
+            logger.warning(
+                "You need `tensorflow_probability` in order to run CQL! "
+                "Install it via `pip install tensorflow_probability`. Your "
+                f"tf.__version__={tf.__version__ if tf else None}."
+                "Trying to import tfp results in the following error:"
+            )
+            try_import_tfp(error=True)
+
+        # Assert that for a local learner the number of iterations is 1. Note,
+        # this is needed because we have no iterators, but instead a single
+        # batch returned directly from the `OfflineData.sample` method.
+        if (
+            self.num_learners == 0
+            and not self.dataset_num_iters_per_learner
+            and self.enable_rl_module_and_learner
+        ):
+            self._value_error(
+                "When using a single local learner the number of iterations "
+                "per learner, `dataset_num_iters_per_learner` has to be defined. "
+                "Set this hyperparameter in the `AlgorithmConfig.offline_data`."
+            )
+
+    @override(SACConfig)
+    def get_default_rl_module_spec(self) -> RLModuleSpecType:
+        if self.framework_str == "torch":
+            from ray.rllib.algorithms.cql.torch.default_cql_torch_rl_module import (
+                DefaultCQLTorchRLModule,
+            )
+
+            return RLModuleSpec(module_class=DefaultCQLTorchRLModule)
+        else:
+            raise ValueError(
+                f"The framework {self.framework_str} is not supported. Use `torch`."
+            )
+
+    @property
+    def _model_config_auto_includes(self):
+        return super()._model_config_auto_includes | {
+            "num_actions": self.num_actions,
+        }
+
+
+class CQL(SAC):
+    """CQL (derived from SAC)."""
+
+    @classmethod
+    @override(SAC)
+    def get_default_config(cls) -> CQLConfig:
+        return CQLConfig()
+
+    @classmethod
+    @override(SAC)
+    def get_default_policy_class(
+        cls, config: AlgorithmConfig
+    ) -> Optional[Type[Policy]]:
+        if config["framework"] == "torch":
+            return CQLTorchPolicy
+        else:
+            return CQLTFPolicy
+
+    @override(SAC)
+    def training_step(self) -> None:
+        # Old API stack (Policy, RolloutWorker, Connector).
+        if not self.config.enable_env_runner_and_connector_v2:
+            return self._training_step_old_api_stack()
+
+        # Sampling from offline data.
+        with self.metrics.log_time((TIMERS, OFFLINE_SAMPLING_TIMER)):
+            # If we should use an iterator in the learner(s). Note, in case of
+            # multiple learners we must always return a list of iterators.
+            return_iterator = return_iterator = (
+                self.config.num_learners > 0
+                or self.config.dataset_num_iters_per_learner != 1
+            )
+
+            # Return an iterator in case we are using remote learners.
+            batch_or_iterator = self.offline_data.sample(
+                num_samples=self.config.train_batch_size_per_learner,
+                num_shards=self.config.num_learners,
+                # Return an iterator, if a `Learner` should update
+                # multiple times per RLlib iteration.
+                return_iterator=return_iterator,
+            )
+
+        # Updating the policy.
+        with self.metrics.log_time((TIMERS, LEARNER_UPDATE_TIMER)):
+            learner_results = self.learner_group.update(
+                data_iterators=batch_or_iterator,
+                minibatch_size=self.config.train_batch_size_per_learner,
+                num_iters=self.config.dataset_num_iters_per_learner,
+            )
+
+            # Log training results.
+            self.metrics.aggregate(learner_results, key=LEARNER_RESULTS)
+
+    @OldAPIStack
+    def _training_step_old_api_stack(self) -> ResultDict:
+        # Collect SampleBatches from sample workers.
+        with self._timers[SAMPLE_TIMER]:
+            train_batch = synchronous_parallel_sample(worker_set=self.env_runner_group)
+        train_batch = train_batch.as_multi_agent()
+        self._counters[NUM_AGENT_STEPS_SAMPLED] += train_batch.agent_steps()
+        self._counters[NUM_ENV_STEPS_SAMPLED] += train_batch.env_steps()
+
+        # Postprocess batch before we learn on it.
+        post_fn = self.config.get("before_learn_on_batch") or (lambda b, *a: b)
+        train_batch = post_fn(train_batch, self.env_runner_group, self.config)
+
+        # Learn on training batch.
+        # Use simple optimizer (only for multi-agent or tf-eager; all other
+        # cases should use the multi-GPU optimizer, even if only using 1 GPU)
+        if self.config.get("simple_optimizer") is True:
+            train_results = train_one_step(self, train_batch)
+        else:
+            train_results = multi_gpu_train_one_step(self, train_batch)
+
+        # Update target network every `target_network_update_freq` training steps.
+        cur_ts = self._counters[
+            NUM_AGENT_STEPS_TRAINED
+            if self.config.count_steps_by == "agent_steps"
+            else NUM_ENV_STEPS_TRAINED
+        ]
+        last_update = self._counters[LAST_TARGET_UPDATE_TS]
+        if cur_ts - last_update >= self.config.target_network_update_freq:
+            with self._timers[TARGET_NET_UPDATE_TIMER]:
+                to_update = self.env_runner.get_policies_to_train()
+                self.env_runner.foreach_policy_to_train(
+                    lambda p, pid: pid in to_update and p.update_target()
+                )
+            self._counters[NUM_TARGET_UPDATES] += 1
+            self._counters[LAST_TARGET_UPDATE_TS] = cur_ts
+
+        # Update remote workers's weights after learning on local worker
+        # (only those policies that were actually trained).
+        if self.env_runner_group.num_remote_workers() > 0:
+            with self._timers[SYNCH_WORKER_WEIGHTS_TIMER]:
+                self.env_runner_group.sync_weights(policies=list(train_results.keys()))
+
+        # Return all collected metrics for the iteration.
+        return train_results

rllib/algorithms/cql/cql_tf_policy.py

@@ -0,0 +1,426 @@
+"""
+TensorFlow policy class used for CQL.
+"""
+import logging
+from functools import partial
+from typing import Dict, List, Type, Union
+
+import gymnasium as gym
+import numpy as np
+import tree
+
+import ray
+from ray.rllib.algorithms.sac.sac_tf_policy import (
+    ActorCriticOptimizerMixin as SACActorCriticOptimizerMixin,
+    ComputeTDErrorMixin,
+    _get_dist_class,
+    apply_gradients as sac_apply_gradients,
+    build_sac_model,
+    compute_and_clip_gradients as sac_compute_and_clip_gradients,
+    get_distribution_inputs_and_class,
+    postprocess_trajectory,
+    setup_late_mixins,
+    stats,
+    validate_spaces,
+)
+from ray.rllib.models.modelv2 import ModelV2
+from ray.rllib.models.tf.tf_action_dist import TFActionDistribution
+from ray.rllib.policy.policy import Policy
+from ray.rllib.policy.sample_batch import SampleBatch
+from ray.rllib.policy.tf_mixins import TargetNetworkMixin
+from ray.rllib.policy.tf_policy_template import build_tf_policy
+from ray.rllib.utils.exploration.random import Random
+from ray.rllib.utils.framework import get_variable, try_import_tf, try_import_tfp
+from ray.rllib.utils.typing import (
+    AlgorithmConfigDict,
+    LocalOptimizer,
+    ModelGradients,
+    TensorType,
+)
+
+tf1, tf, tfv = try_import_tf()
+tfp = try_import_tfp()
+
+logger = logging.getLogger(__name__)
+
+MEAN_MIN = -9.0
+MEAN_MAX = 9.0
+
+
+def _repeat_tensor(t: TensorType, n: int):
+    # Insert new axis at position 1 into tensor t
+    t_rep = tf.expand_dims(t, 1)
+    # Repeat tensor t_rep along new axis n times
+    multiples = tf.concat([[1, n], tf.tile([1], tf.expand_dims(tf.rank(t) - 1, 0))], 0)
+    t_rep = tf.tile(t_rep, multiples)
+    # Merge new axis into batch axis
+    t_rep = tf.reshape(t_rep, tf.concat([[-1], tf.shape(t)[1:]], 0))
+    return t_rep
+
+
+# Returns policy tiled actions and log probabilities for CQL Loss
+def policy_actions_repeat(model, action_dist, obs, num_repeat=1):
+    batch_size = tf.shape(tree.flatten(obs)[0])[0]
+    obs_temp = tree.map_structure(lambda t: _repeat_tensor(t, num_repeat), obs)
+    logits, _ = model.get_action_model_outputs(obs_temp)
+    policy_dist = action_dist(logits, model)
+    actions, logp_ = policy_dist.sample_logp()
+    logp = tf.expand_dims(logp_, -1)
+    return actions, tf.reshape(logp, [batch_size, num_repeat, 1])
+
+
+def q_values_repeat(model, obs, actions, twin=False):
+    action_shape = tf.shape(actions)[0]
+    obs_shape = tf.shape(tree.flatten(obs)[0])[0]
+    num_repeat = action_shape // obs_shape
+    obs_temp = tree.map_structure(lambda t: _repeat_tensor(t, num_repeat), obs)
+    if not twin:
+        preds_, _ = model.get_q_values(obs_temp, actions)
+    else:
+        preds_, _ = model.get_twin_q_values(obs_temp, actions)
+    preds = tf.reshape(preds_, [obs_shape, num_repeat, 1])
+    return preds
+
+
+def cql_loss(
+    policy: Policy,
+    model: ModelV2,
+    dist_class: Type[TFActionDistribution],
+    train_batch: SampleBatch,
+) -> Union[TensorType, List[TensorType]]:
+    logger.info(f"Current iteration = {policy.cur_iter}")
+    policy.cur_iter += 1
+
+    # For best performance, turn deterministic off
+    deterministic = policy.config["_deterministic_loss"]
+    assert not deterministic
+    twin_q = policy.config["twin_q"]
+    discount = policy.config["gamma"]
+
+    # CQL Parameters
+    bc_iters = policy.config["bc_iters"]
+    cql_temp = policy.config["temperature"]
+    num_actions = policy.config["num_actions"]
+    min_q_weight = policy.config["min_q_weight"]
+    use_lagrange = policy.config["lagrangian"]
+    target_action_gap = policy.config["lagrangian_thresh"]
+
+    obs = train_batch[SampleBatch.CUR_OBS]
+    actions = tf.cast(train_batch[SampleBatch.ACTIONS], tf.float32)
+    rewards = tf.cast(train_batch[SampleBatch.REWARDS], tf.float32)
+    next_obs = train_batch[SampleBatch.NEXT_OBS]
+    terminals = train_batch[SampleBatch.TERMINATEDS]
+
+    model_out_t, _ = model(SampleBatch(obs=obs, _is_training=True), [], None)
+
+    model_out_tp1, _ = model(SampleBatch(obs=next_obs, _is_training=True), [], None)
+
+    target_model_out_tp1, _ = policy.target_model(
+        SampleBatch(obs=next_obs, _is_training=True), [], None
+    )
+
+    action_dist_class = _get_dist_class(policy, policy.config, policy.action_space)
+    action_dist_inputs_t, _ = model.get_action_model_outputs(model_out_t)
+    action_dist_t = action_dist_class(action_dist_inputs_t, model)
+    policy_t, log_pis_t = action_dist_t.sample_logp()
+    log_pis_t = tf.expand_dims(log_pis_t, -1)
+
+    # Unlike original SAC, Alpha and Actor Loss are computed first.
+    # Alpha Loss
+    alpha_loss = -tf.reduce_mean(
+        model.log_alpha * tf.stop_gradient(log_pis_t + model.target_entropy)
+    )
+
+    # Policy Loss (Either Behavior Clone Loss or SAC Loss)
+    alpha = tf.math.exp(model.log_alpha)
+    if policy.cur_iter >= bc_iters:
+        min_q, _ = model.get_q_values(model_out_t, policy_t)
+        if twin_q:
+            twin_q_, _ = model.get_twin_q_values(model_out_t, policy_t)
+            min_q = tf.math.minimum(min_q, twin_q_)
+        actor_loss = tf.reduce_mean(tf.stop_gradient(alpha) * log_pis_t - min_q)
+    else:
+        bc_logp = action_dist_t.logp(actions)
+        actor_loss = tf.reduce_mean(tf.stop_gradient(alpha) * log_pis_t - bc_logp)
+        # actor_loss = -tf.reduce_mean(bc_logp)
+
+    # Critic Loss (Standard SAC Critic L2 Loss + CQL Entropy Loss)
+    # SAC Loss:
+    # Q-values for the batched actions.
+    action_dist_inputs_tp1, _ = model.get_action_model_outputs(model_out_tp1)
+    action_dist_tp1 = action_dist_class(action_dist_inputs_tp1, model)
+    policy_tp1, _ = action_dist_tp1.sample_logp()
+
+    q_t, _ = model.get_q_values(model_out_t, actions)
+    q_t_selected = tf.squeeze(q_t, axis=-1)
+    if twin_q:
+        twin_q_t, _ = model.get_twin_q_values(model_out_t, actions)
+        twin_q_t_selected = tf.squeeze(twin_q_t, axis=-1)
+
+    # Target q network evaluation.
+    q_tp1, _ = policy.target_model.get_q_values(target_model_out_tp1, policy_tp1)
+    if twin_q:
+        twin_q_tp1, _ = policy.target_model.get_twin_q_values(
+            target_model_out_tp1, policy_tp1
+        )
+        # Take min over both twin-NNs.
+        q_tp1 = tf.math.minimum(q_tp1, twin_q_tp1)
+
+    q_tp1_best = tf.squeeze(input=q_tp1, axis=-1)
+    q_tp1_best_masked = (1.0 - tf.cast(terminals, tf.float32)) * q_tp1_best
+
+    # compute RHS of bellman equation
+    q_t_target = tf.stop_gradient(
+        rewards + (discount ** policy.config["n_step"]) * q_tp1_best_masked
+    )
+
+    # Compute the TD-error (potentially clipped), for priority replay buffer
+    base_td_error = tf.math.abs(q_t_selected - q_t_target)
+    if twin_q:
+        twin_td_error = tf.math.abs(twin_q_t_selected - q_t_target)
+        td_error = 0.5 * (base_td_error + twin_td_error)
+    else:
+        td_error = base_td_error
+
+    critic_loss_1 = tf.keras.losses.MSE(q_t_selected, q_t_target)
+    if twin_q:
+        critic_loss_2 = tf.keras.losses.MSE(twin_q_t_selected, q_t_target)
+
+    # CQL Loss (We are using Entropy version of CQL (the best version))
+    rand_actions, _ = policy._random_action_generator.get_exploration_action(
+        action_distribution=action_dist_class(
+            tf.tile(action_dist_tp1.inputs, (num_actions, 1)), model
+        ),
+        timestep=0,
+        explore=True,
+    )
+    curr_actions, curr_logp = policy_actions_repeat(
+        model, action_dist_class, model_out_t, num_actions
+    )
+    next_actions, next_logp = policy_actions_repeat(
+        model, action_dist_class, model_out_tp1, num_actions
+    )
+
+    q1_rand = q_values_repeat(model, model_out_t, rand_actions)
+    q1_curr_actions = q_values_repeat(model, model_out_t, curr_actions)
+    q1_next_actions = q_values_repeat(model, model_out_t, next_actions)
+
+    if twin_q:
+        q2_rand = q_values_repeat(model, model_out_t, rand_actions, twin=True)
+        q2_curr_actions = q_values_repeat(model, model_out_t, curr_actions, twin=True)
+        q2_next_actions = q_values_repeat(model, model_out_t, next_actions, twin=True)
+
+    random_density = np.log(0.5 ** int(curr_actions.shape[-1]))
+    cat_q1 = tf.concat(
+        [
+            q1_rand - random_density,
+            q1_next_actions - tf.stop_gradient(next_logp),
+            q1_curr_actions - tf.stop_gradient(curr_logp),
+        ],
+        1,
+    )
+    if twin_q:
+        cat_q2 = tf.concat(
+            [
+                q2_rand - random_density,
+                q2_next_actions - tf.stop_gradient(next_logp),
+                q2_curr_actions - tf.stop_gradient(curr_logp),
+            ],
+            1,
+        )
+
+    min_qf1_loss_ = (
+        tf.reduce_mean(tf.reduce_logsumexp(cat_q1 / cql_temp, axis=1))
+        * min_q_weight
+        * cql_temp
+    )
+    min_qf1_loss = min_qf1_loss_ - (tf.reduce_mean(q_t) * min_q_weight)
+    if twin_q:
+        min_qf2_loss_ = (
+            tf.reduce_mean(tf.reduce_logsumexp(cat_q2 / cql_temp, axis=1))
+            * min_q_weight
+            * cql_temp
+        )
+        min_qf2_loss = min_qf2_loss_ - (tf.reduce_mean(twin_q_t) * min_q_weight)
+
+    if use_lagrange:
+        alpha_prime = tf.clip_by_value(model.log_alpha_prime.exp(), 0.0, 1000000.0)[0]
+        min_qf1_loss = alpha_prime * (min_qf1_loss - target_action_gap)
+        if twin_q:
+            min_qf2_loss = alpha_prime * (min_qf2_loss - target_action_gap)
+            alpha_prime_loss = 0.5 * (-min_qf1_loss - min_qf2_loss)
+        else:
+            alpha_prime_loss = -min_qf1_loss
+
+    cql_loss = [min_qf1_loss]
+    if twin_q:
+        cql_loss.append(min_qf2_loss)
+
+    critic_loss = [critic_loss_1 + min_qf1_loss]
+    if twin_q:
+        critic_loss.append(critic_loss_2 + min_qf2_loss)
+
+    # Save for stats function.
+    policy.q_t = q_t_selected
+    policy.policy_t = policy_t
+    policy.log_pis_t = log_pis_t
+    policy.td_error = td_error
+    policy.actor_loss = actor_loss
+    policy.critic_loss = critic_loss
+    policy.alpha_loss = alpha_loss
+    policy.log_alpha_value = model.log_alpha
+    policy.alpha_value = alpha
+    policy.target_entropy = model.target_entropy
+    # CQL Stats
+    policy.cql_loss = cql_loss
+    if use_lagrange:
+        policy.log_alpha_prime_value = model.log_alpha_prime[0]
+        policy.alpha_prime_value = alpha_prime
+        policy.alpha_prime_loss = alpha_prime_loss
+
+    # Return all loss terms corresponding to our optimizers.
+    if use_lagrange:
+        return actor_loss + tf.math.add_n(critic_loss) + alpha_loss + alpha_prime_loss
+    return actor_loss + tf.math.add_n(critic_loss) + alpha_loss
+
+
+def cql_stats(policy: Policy, train_batch: SampleBatch) -> Dict[str, TensorType]:
+    sac_dict = stats(policy, train_batch)
+    sac_dict["cql_loss"] = tf.reduce_mean(tf.stack(policy.cql_loss))
+    if policy.config["lagrangian"]:
+        sac_dict["log_alpha_prime_value"] = policy.log_alpha_prime_value
+        sac_dict["alpha_prime_value"] = policy.alpha_prime_value
+        sac_dict["alpha_prime_loss"] = policy.alpha_prime_loss
+    return sac_dict
+
+
+class ActorCriticOptimizerMixin(SACActorCriticOptimizerMixin):
+    def __init__(self, config):
+        super().__init__(config)
+        if config["lagrangian"]:
+            # Eager mode.
+            if config["framework"] == "tf2":
+                self._alpha_prime_optimizer = tf.keras.optimizers.Adam(
+                    learning_rate=config["optimization"]["critic_learning_rate"]
+                )
+            # Static graph mode.
+            else:
+                self._alpha_prime_optimizer = tf1.train.AdamOptimizer(
+                    learning_rate=config["optimization"]["critic_learning_rate"]
+                )
+
+
+def setup_early_mixins(
+    policy: Policy,
+    obs_space: gym.spaces.Space,
+    action_space: gym.spaces.Space,
+    config: AlgorithmConfigDict,
+) -> None:
+    """Call mixin classes' constructors before Policy's initialization.
+
+    Adds the necessary optimizers to the given Policy.
+
+    Args:
+        policy: The Policy object.
+        obs_space (gym.spaces.Space): The Policy's observation space.
+        action_space (gym.spaces.Space): The Policy's action space.
+        config: The Policy's config.
+    """
+    policy.cur_iter = 0
+    ActorCriticOptimizerMixin.__init__(policy, config)
+    if config["lagrangian"]:
+        policy.model.log_alpha_prime = get_variable(
+            0.0, framework="tf", trainable=True, tf_name="log_alpha_prime"
+        )
+        policy.alpha_prime_optim = tf.keras.optimizers.Adam(
+            learning_rate=config["optimization"]["critic_learning_rate"],
+        )
+    # Generic random action generator for calculating CQL-loss.
+    policy._random_action_generator = Random(
+        action_space,
+        model=None,
+        framework="tf2",
+        policy_config=config,
+        num_workers=0,
+        worker_index=0,
+    )
+
+
+def compute_gradients_fn(
+    policy: Policy, optimizer: LocalOptimizer, loss: TensorType
+) -> ModelGradients:
+    grads_and_vars = sac_compute_and_clip_gradients(policy, optimizer, loss)
+
+    if policy.config["lagrangian"]:
+        # Eager: Use GradientTape (which is a property of the `optimizer`
+        # object (an OptimizerWrapper): see rllib/policy/eager_tf_policy.py).
+        if policy.config["framework"] == "tf2":
+            tape = optimizer.tape
+            log_alpha_prime = [policy.model.log_alpha_prime]
+            alpha_prime_grads_and_vars = list(
+                zip(
+                    tape.gradient(policy.alpha_prime_loss, log_alpha_prime),
+                    log_alpha_prime,
+                )
+            )
+        # Tf1.x: Use optimizer.compute_gradients()
+        else:
+            alpha_prime_grads_and_vars = (
+                policy._alpha_prime_optimizer.compute_gradients(
+                    policy.alpha_prime_loss, var_list=[policy.model.log_alpha_prime]
+                )
+            )
+
+        # Clip if necessary.
+        if policy.config["grad_clip"]:
+            clip_func = partial(tf.clip_by_norm, clip_norm=policy.config["grad_clip"])
+        else:
+            clip_func = tf.identity
+
+        # Save grads and vars for later use in `build_apply_op`.
+        policy._alpha_prime_grads_and_vars = [
+            (clip_func(g), v) for (g, v) in alpha_prime_grads_and_vars if g is not None
+        ]
+
+        grads_and_vars += policy._alpha_prime_grads_and_vars
+    return grads_and_vars
+
+
+def apply_gradients_fn(policy, optimizer, grads_and_vars):
+    sac_results = sac_apply_gradients(policy, optimizer, grads_and_vars)
+
+    if policy.config["lagrangian"]:
+        # Eager mode -> Just apply and return None.
+        if policy.config["framework"] == "tf2":
+            policy._alpha_prime_optimizer.apply_gradients(
+                policy._alpha_prime_grads_and_vars
+            )
+            return
+        # Tf static graph -> Return grouped op.
+        else:
+            alpha_prime_apply_op = policy._alpha_prime_optimizer.apply_gradients(
+                policy._alpha_prime_grads_and_vars,
+                global_step=tf1.train.get_or_create_global_step(),
+            )
+            return tf.group([sac_results, alpha_prime_apply_op])
+    return sac_results
+
+
+# Build a child class of `TFPolicy`, given the custom functions defined
+# above.
+CQLTFPolicy = build_tf_policy(
+    name="CQLTFPolicy",
+    loss_fn=cql_loss,
+    get_default_config=lambda: ray.rllib.algorithms.cql.cql.CQLConfig(),
+    validate_spaces=validate_spaces,
+    stats_fn=cql_stats,
+    postprocess_fn=postprocess_trajectory,
+    before_init=setup_early_mixins,
+    after_init=setup_late_mixins,
+    make_model=build_sac_model,
+    extra_learn_fetches_fn=lambda policy: {"td_error": policy.td_error},
+    mixins=[ActorCriticOptimizerMixin, TargetNetworkMixin, ComputeTDErrorMixin],
+    action_distribution_fn=get_distribution_inputs_and_class,
+    compute_gradients_fn=compute_gradients_fn,
+    apply_gradients_fn=apply_gradients_fn,
+)

rllib/algorithms/cql/cql_torch_policy.py

@@ -0,0 +1,406 @@
+"""
+PyTorch policy class used for CQL.
+"""
+import logging
+from typing import Dict, List, Tuple, Type, Union
+
+import gymnasium as gym
+import numpy as np
+import tree
+
+import ray
+from ray.rllib.algorithms.sac.sac_tf_policy import (
+    postprocess_trajectory,
+    validate_spaces,
+)
+from ray.rllib.algorithms.sac.sac_torch_policy import (
+    ComputeTDErrorMixin,
+    _get_dist_class,
+    action_distribution_fn,
+    build_sac_model_and_action_dist,
+    optimizer_fn,
+    setup_late_mixins,
+    stats,
+)
+from ray.rllib.models.modelv2 import ModelV2
+from ray.rllib.models.torch.torch_action_dist import TorchDistributionWrapper
+from ray.rllib.policy.policy import Policy
+from ray.rllib.policy.policy_template import build_policy_class
+from ray.rllib.policy.sample_batch import SampleBatch
+from ray.rllib.policy.torch_mixins import TargetNetworkMixin
+from ray.rllib.utils.framework import try_import_torch
+from ray.rllib.utils.metrics.learner_info import LEARNER_STATS_KEY
+from ray.rllib.utils.torch_utils import (
+    apply_grad_clipping,
+    concat_multi_gpu_td_errors,
+    convert_to_torch_tensor,
+)
+from ray.rllib.utils.typing import AlgorithmConfigDict, LocalOptimizer, TensorType
+
+torch, nn = try_import_torch()
+F = nn.functional
+
+logger = logging.getLogger(__name__)
+
+MEAN_MIN = -9.0
+MEAN_MAX = 9.0
+
+
+def _repeat_tensor(t: TensorType, n: int):
+    # Insert new dimension at posotion 1 into tensor t
+    t_rep = t.unsqueeze(1)
+    # Repeat tensor t_rep along new dimension n times
+    t_rep = torch.repeat_interleave(t_rep, n, dim=1)
+    # Merge new dimension into batch dimension
+    t_rep = t_rep.view(-1, *t.shape[1:])
+    return t_rep
+
+
+# Returns policy tiled actions and log probabilities for CQL Loss
+def policy_actions_repeat(model, action_dist, obs, num_repeat=1):
+    batch_size = tree.flatten(obs)[0].shape[0]
+    obs_temp = tree.map_structure(lambda t: _repeat_tensor(t, num_repeat), obs)
+    logits, _ = model.get_action_model_outputs(obs_temp)
+    policy_dist = action_dist(logits, model)
+    actions, logp_ = policy_dist.sample_logp()
+    logp = logp_.unsqueeze(-1)
+    return actions, logp.view(batch_size, num_repeat, 1)
+
+
+def q_values_repeat(model, obs, actions, twin=False):
+    action_shape = actions.shape[0]
+    obs_shape = tree.flatten(obs)[0].shape[0]
+    num_repeat = int(action_shape / obs_shape)
+    obs_temp = tree.map_structure(lambda t: _repeat_tensor(t, num_repeat), obs)
+    if not twin:
+        preds_, _ = model.get_q_values(obs_temp, actions)
+    else:
+        preds_, _ = model.get_twin_q_values(obs_temp, actions)
+    preds = preds_.view(obs_shape, num_repeat, 1)
+    return preds
+
+
+def cql_loss(
+    policy: Policy,
+    model: ModelV2,
+    dist_class: Type[TorchDistributionWrapper],
+    train_batch: SampleBatch,
+) -> Union[TensorType, List[TensorType]]:
+    logger.info(f"Current iteration = {policy.cur_iter}")
+    policy.cur_iter += 1
+
+    # Look up the target model (tower) using the model tower.
+    target_model = policy.target_models[model]
+
+    # For best performance, turn deterministic off
+    deterministic = policy.config["_deterministic_loss"]
+    assert not deterministic
+    twin_q = policy.config["twin_q"]
+    discount = policy.config["gamma"]
+    action_low = model.action_space.low[0]
+    action_high = model.action_space.high[0]
+
+    # CQL Parameters
+    bc_iters = policy.config["bc_iters"]
+    cql_temp = policy.config["temperature"]
+    num_actions = policy.config["num_actions"]
+    min_q_weight = policy.config["min_q_weight"]
+    use_lagrange = policy.config["lagrangian"]
+    target_action_gap = policy.config["lagrangian_thresh"]
+
+    obs = train_batch[SampleBatch.CUR_OBS]
+    actions = train_batch[SampleBatch.ACTIONS]
+    rewards = train_batch[SampleBatch.REWARDS].float()
+    next_obs = train_batch[SampleBatch.NEXT_OBS]
+    terminals = train_batch[SampleBatch.TERMINATEDS]
+
+    model_out_t, _ = model(SampleBatch(obs=obs, _is_training=True), [], None)
+
+    model_out_tp1, _ = model(SampleBatch(obs=next_obs, _is_training=True), [], None)
+
+    target_model_out_tp1, _ = target_model(
+        SampleBatch(obs=next_obs, _is_training=True), [], None
+    )
+
+    action_dist_class = _get_dist_class(policy, policy.config, policy.action_space)
+    action_dist_inputs_t, _ = model.get_action_model_outputs(model_out_t)
+    action_dist_t = action_dist_class(action_dist_inputs_t, model)
+    policy_t, log_pis_t = action_dist_t.sample_logp()
+    log_pis_t = torch.unsqueeze(log_pis_t, -1)
+
+    # Unlike original SAC, Alpha and Actor Loss are computed first.
+    # Alpha Loss
+    alpha_loss = -(model.log_alpha * (log_pis_t + model.target_entropy).detach()).mean()
+
+    batch_size = tree.flatten(obs)[0].shape[0]
+    if batch_size == policy.config["train_batch_size"]:
+        policy.alpha_optim.zero_grad()
+        alpha_loss.backward()
+        policy.alpha_optim.step()
+
+    # Policy Loss (Either Behavior Clone Loss or SAC Loss)
+    alpha = torch.exp(model.log_alpha)
+    if policy.cur_iter >= bc_iters:
+        min_q, _ = model.get_q_values(model_out_t, policy_t)
+        if twin_q:
+            twin_q_, _ = model.get_twin_q_values(model_out_t, policy_t)
+            min_q = torch.min(min_q, twin_q_)
+        actor_loss = (alpha.detach() * log_pis_t - min_q).mean()
+    else:
+        bc_logp = action_dist_t.logp(actions)
+        actor_loss = (alpha.detach() * log_pis_t - bc_logp).mean()
+        # actor_loss = -bc_logp.mean()
+
+    if batch_size == policy.config["train_batch_size"]:
+        policy.actor_optim.zero_grad()
+        actor_loss.backward(retain_graph=True)
+        policy.actor_optim.step()
+
+    # Critic Loss (Standard SAC Critic L2 Loss + CQL Entropy Loss)
+    # SAC Loss:
+    # Q-values for the batched actions.
+    action_dist_inputs_tp1, _ = model.get_action_model_outputs(model_out_tp1)
+    action_dist_tp1 = action_dist_class(action_dist_inputs_tp1, model)
+    policy_tp1, _ = action_dist_tp1.sample_logp()
+
+    q_t, _ = model.get_q_values(model_out_t, train_batch[SampleBatch.ACTIONS])
+    q_t_selected = torch.squeeze(q_t, dim=-1)
+    if twin_q:
+        twin_q_t, _ = model.get_twin_q_values(
+            model_out_t, train_batch[SampleBatch.ACTIONS]
+        )
+        twin_q_t_selected = torch.squeeze(twin_q_t, dim=-1)
+
+    # Target q network evaluation.
+    q_tp1, _ = target_model.get_q_values(target_model_out_tp1, policy_tp1)
+    if twin_q:
+        twin_q_tp1, _ = target_model.get_twin_q_values(target_model_out_tp1, policy_tp1)
+        # Take min over both twin-NNs.
+        q_tp1 = torch.min(q_tp1, twin_q_tp1)
+
+    q_tp1_best = torch.squeeze(input=q_tp1, dim=-1)
+    q_tp1_best_masked = (1.0 - terminals.float()) * q_tp1_best
+
+    # compute RHS of bellman equation
+    q_t_target = (
+        rewards + (discount ** policy.config["n_step"]) * q_tp1_best_masked
+    ).detach()
+
+    # Compute the TD-error (potentially clipped), for priority replay buffer
+    base_td_error = torch.abs(q_t_selected - q_t_target)
+    if twin_q:
+        twin_td_error = torch.abs(twin_q_t_selected - q_t_target)
+        td_error = 0.5 * (base_td_error + twin_td_error)
+    else:
+        td_error = base_td_error
+
+    critic_loss_1 = nn.functional.mse_loss(q_t_selected, q_t_target)
+    if twin_q:
+        critic_loss_2 = nn.functional.mse_loss(twin_q_t_selected, q_t_target)
+
+    # CQL Loss (We are using Entropy version of CQL (the best version))
+    rand_actions = convert_to_torch_tensor(
+        torch.FloatTensor(actions.shape[0] * num_actions, actions.shape[-1]).uniform_(
+            action_low, action_high
+        ),
+        policy.device,
+    )
+    curr_actions, curr_logp = policy_actions_repeat(
+        model, action_dist_class, model_out_t, num_actions
+    )
+    next_actions, next_logp = policy_actions_repeat(
+        model, action_dist_class, model_out_tp1, num_actions
+    )
+
+    q1_rand = q_values_repeat(model, model_out_t, rand_actions)
+    q1_curr_actions = q_values_repeat(model, model_out_t, curr_actions)
+    q1_next_actions = q_values_repeat(model, model_out_t, next_actions)
+
+    if twin_q:
+        q2_rand = q_values_repeat(model, model_out_t, rand_actions, twin=True)
+        q2_curr_actions = q_values_repeat(model, model_out_t, curr_actions, twin=True)
+        q2_next_actions = q_values_repeat(model, model_out_t, next_actions, twin=True)
+
+    random_density = np.log(0.5 ** curr_actions.shape[-1])
+    cat_q1 = torch.cat(
+        [
+            q1_rand - random_density,
+            q1_next_actions - next_logp.detach(),
+            q1_curr_actions - curr_logp.detach(),
+        ],
+        1,
+    )
+    if twin_q:
+        cat_q2 = torch.cat(
+            [
+                q2_rand - random_density,
+                q2_next_actions - next_logp.detach(),
+                q2_curr_actions - curr_logp.detach(),
+            ],
+            1,
+        )
+
+    min_qf1_loss_ = (
+        torch.logsumexp(cat_q1 / cql_temp, dim=1).mean() * min_q_weight * cql_temp
+    )
+    min_qf1_loss = min_qf1_loss_ - (q_t.mean() * min_q_weight)
+    if twin_q:
+        min_qf2_loss_ = (
+            torch.logsumexp(cat_q2 / cql_temp, dim=1).mean() * min_q_weight * cql_temp
+        )
+        min_qf2_loss = min_qf2_loss_ - (twin_q_t.mean() * min_q_weight)
+
+    if use_lagrange:
+        alpha_prime = torch.clamp(model.log_alpha_prime.exp(), min=0.0, max=1000000.0)[
+            0
+        ]
+        min_qf1_loss = alpha_prime * (min_qf1_loss - target_action_gap)
+        if twin_q:
+            min_qf2_loss = alpha_prime * (min_qf2_loss - target_action_gap)
+            alpha_prime_loss = 0.5 * (-min_qf1_loss - min_qf2_loss)
+        else:
+            alpha_prime_loss = -min_qf1_loss
+
+    cql_loss = [min_qf1_loss]
+    if twin_q:
+        cql_loss.append(min_qf2_loss)
+
+    critic_loss = [critic_loss_1 + min_qf1_loss]
+    if twin_q:
+        critic_loss.append(critic_loss_2 + min_qf2_loss)
+
+    if batch_size == policy.config["train_batch_size"]:
+        policy.critic_optims[0].zero_grad()
+        critic_loss[0].backward(retain_graph=True)
+        policy.critic_optims[0].step()
+
+        if twin_q:
+            policy.critic_optims[1].zero_grad()
+            critic_loss[1].backward(retain_graph=False)
+            policy.critic_optims[1].step()
+
+    # Store values for stats function in model (tower), such that for
+    # multi-GPU, we do not override them during the parallel loss phase.
+    # SAC stats.
+    model.tower_stats["q_t"] = q_t_selected
+    model.tower_stats["policy_t"] = policy_t
+    model.tower_stats["log_pis_t"] = log_pis_t
+    model.tower_stats["actor_loss"] = actor_loss
+    model.tower_stats["critic_loss"] = critic_loss
+    model.tower_stats["alpha_loss"] = alpha_loss
+    model.tower_stats["log_alpha_value"] = model.log_alpha
+    model.tower_stats["alpha_value"] = alpha
+    model.tower_stats["target_entropy"] = model.target_entropy
+    # CQL stats.
+    model.tower_stats["cql_loss"] = cql_loss
+
+    # TD-error tensor in final stats
+    # will be concatenated and retrieved for each individual batch item.
+    model.tower_stats["td_error"] = td_error
+
+    if use_lagrange:
+        model.tower_stats["log_alpha_prime_value"] = model.log_alpha_prime[0]
+        model.tower_stats["alpha_prime_value"] = alpha_prime
+        model.tower_stats["alpha_prime_loss"] = alpha_prime_loss
+
+        if batch_size == policy.config["train_batch_size"]:
+            policy.alpha_prime_optim.zero_grad()
+            alpha_prime_loss.backward()
+            policy.alpha_prime_optim.step()
+
+    # Return all loss terms corresponding to our optimizers.
+    return tuple(
+        [actor_loss]
+        + critic_loss
+        + [alpha_loss]
+        + ([alpha_prime_loss] if use_lagrange else [])
+    )
+
+
+def cql_stats(policy: Policy, train_batch: SampleBatch) -> Dict[str, TensorType]:
+    # Get SAC loss stats.
+    stats_dict = stats(policy, train_batch)
+
+    # Add CQL loss stats to the dict.
+    stats_dict["cql_loss"] = torch.mean(
+        torch.stack(*policy.get_tower_stats("cql_loss"))
+    )
+
+    if policy.config["lagrangian"]:
+        stats_dict["log_alpha_prime_value"] = torch.mean(
+            torch.stack(policy.get_tower_stats("log_alpha_prime_value"))
+        )
+        stats_dict["alpha_prime_value"] = torch.mean(
+            torch.stack(policy.get_tower_stats("alpha_prime_value"))
+        )
+        stats_dict["alpha_prime_loss"] = torch.mean(
+            torch.stack(policy.get_tower_stats("alpha_prime_loss"))
+        )
+    return stats_dict
+
+
+def cql_optimizer_fn(
+    policy: Policy, config: AlgorithmConfigDict
+) -> Tuple[LocalOptimizer]:
+    policy.cur_iter = 0
+    opt_list = optimizer_fn(policy, config)
+    if config["lagrangian"]:
+        log_alpha_prime = nn.Parameter(torch.zeros(1, requires_grad=True).float())
+        policy.model.register_parameter("log_alpha_prime", log_alpha_prime)
+        policy.alpha_prime_optim = torch.optim.Adam(
+            params=[policy.model.log_alpha_prime],
+            lr=config["optimization"]["critic_learning_rate"],
+            eps=1e-7,  # to match tf.keras.optimizers.Adam's epsilon default
+        )
+        return tuple(
+            [policy.actor_optim]
+            + policy.critic_optims
+            + [policy.alpha_optim]
+            + [policy.alpha_prime_optim]
+        )
+    return opt_list
+
+
+def cql_setup_late_mixins(
+    policy: Policy,
+    obs_space: gym.spaces.Space,
+    action_space: gym.spaces.Space,
+    config: AlgorithmConfigDict,
+) -> None:
+    setup_late_mixins(policy, obs_space, action_space, config)
+    if config["lagrangian"]:
+        policy.model.log_alpha_prime = policy.model.log_alpha_prime.to(policy.device)
+
+
+def compute_gradients_fn(policy, postprocessed_batch):
+    batches = [policy._lazy_tensor_dict(postprocessed_batch)]
+    model = policy.model
+    policy._loss(policy, model, policy.dist_class, batches[0])
+    stats = {LEARNER_STATS_KEY: policy._convert_to_numpy(cql_stats(policy, batches[0]))}
+    return [None, stats]
+
+
+def apply_gradients_fn(policy, gradients):
+    return
+
+
+# Build a child class of `TorchPolicy`, given the custom functions defined
+# above.
+CQLTorchPolicy = build_policy_class(
+    name="CQLTorchPolicy",
+    framework="torch",
+    loss_fn=cql_loss,
+    get_default_config=lambda: ray.rllib.algorithms.cql.cql.CQLConfig(),
+    stats_fn=cql_stats,
+    postprocess_fn=postprocess_trajectory,
+    extra_grad_process_fn=apply_grad_clipping,
+    optimizer_fn=cql_optimizer_fn,
+    validate_spaces=validate_spaces,
+    before_loss_init=cql_setup_late_mixins,
+    make_model_and_action_dist=build_sac_model_and_action_dist,
+    extra_learn_fetches_fn=concat_multi_gpu_td_errors,
+    mixins=[TargetNetworkMixin, ComputeTDErrorMixin],
+    action_distribution_fn=action_distribution_fn,
+    compute_gradients_fn=compute_gradients_fn,
+    apply_gradients_fn=apply_gradients_fn,
+)

rllib/algorithms/cql/tests/test_cql_old_api_stack.py

@@ -0,0 +1,128 @@
+import os
+import unittest
+from pathlib import Path
+
+import ray
+from ray.rllib.algorithms import cql
+from ray.rllib.utils.framework import try_import_torch
+from ray.rllib.utils.metrics import (
+    ENV_RUNNER_RESULTS,
+    EPISODE_RETURN_MEAN,
+    EVALUATION_RESULTS,
+)
+from ray.rllib.utils.test_utils import check_compute_single_action, check_train_results
+
+torch, _ = try_import_torch()
+
+
+class TestCQL(unittest.TestCase):
+    @classmethod
+    def setUpClass(cls):
+        ray.init()
+
+    @classmethod
+    def tearDownClass(cls):
+        ray.shutdown()
+
+    def test_cql_compilation(self):
+        """Test whether CQL can be built with all frameworks."""
+
+        # Learns from a historic-data file.
+        # To generate this data, first run:
+        # $ ./train.py --run=SAC --env=Pendulum-v1 \
+        #   --stop='{"timesteps_total": 50000}' \
+        #   --config='{"output": "/tmp/out"}'
+        rllib_dir = Path(__file__).parent.parent.parent.parent
+        print("rllib dir={}".format(rllib_dir))
+        data_file = os.path.join(rllib_dir, "offline/tests/data/pendulum/small.json")
+        print("data_file={} exists={}".format(data_file, os.path.isfile(data_file)))
+
+        config = (
+            cql.CQLConfig()
+            .api_stack(
+                enable_rl_module_and_learner=False,
+                enable_env_runner_and_connector_v2=False,
+            )
+            .environment(
+                env="Pendulum-v1",
+            )
+            .offline_data(
+                input_=data_file,
+                # In the files, we use here for testing, actions have already
+                # been normalized.
+                # This is usually the case when the file was generated by another
+                # RLlib algorithm (e.g. PPO or SAC).
+                actions_in_input_normalized=False,
+            )
+            .training(
+                clip_actions=False,
+                train_batch_size=2000,
+                twin_q=True,
+                num_steps_sampled_before_learning_starts=0,
+                bc_iters=2,
+            )
+            .evaluation(
+                evaluation_interval=2,
+                evaluation_duration=10,
+                evaluation_config=cql.CQLConfig.overrides(input_="sampler"),
+                evaluation_parallel_to_training=False,
+                evaluation_num_env_runners=2,
+            )
+            .env_runners(num_env_runners=0)
+            .reporting(min_time_s_per_iteration=0)
+        )
+        num_iterations = 4
+
+        algo = config.build()
+        for i in range(num_iterations):
+            results = algo.train()
+            check_train_results(results)
+            print(results)
+            eval_results = results.get(EVALUATION_RESULTS)
+            if eval_results:
+                print(
+                    f"iter={algo.iteration} "
+                    f"R={eval_results[ENV_RUNNER_RESULTS][EPISODE_RETURN_MEAN]}"
+                )
+        check_compute_single_action(algo)
+
+        # Get policy and model.
+        pol = algo.get_policy()
+        cql_model = pol.model
+
+        # Example on how to do evaluation on the trained Algorithm
+        # using the data from CQL's global replay buffer.
+        # Get a sample (MultiAgentBatch).
+
+        batch = algo.env_runner.input_reader.next()
+        multi_agent_batch = batch.as_multi_agent()
+        # All experiences have been buffered for `default_policy`
+        batch = multi_agent_batch.policy_batches["default_policy"]
+
+        obs = torch.from_numpy(batch["obs"])
+
+        # Pass the observations through our model to get the
+        # features, which then to pass through the Q-head.
+        model_out, _ = cql_model({"obs": obs})
+        # The estimated Q-values from the (historic) actions in the batch.
+        q_values_old = cql_model.get_q_values(
+            model_out, torch.from_numpy(batch["actions"])
+        )
+
+        # The estimated Q-values for the new actions computed
+        # by our policy.
+        actions_new = pol.compute_actions_from_input_dict({"obs": obs})[0]
+        q_values_new = cql_model.get_q_values(model_out, torch.from_numpy(actions_new))
+
+        print(f"Q-val batch={q_values_old}")
+        print(f"Q-val policy={q_values_new}")
+
+        algo.stop()
+
+
+if __name__ == "__main__":
+    import sys
+
+    import pytest
+
+    sys.exit(pytest.main(["-v", __file__]))

rllib/algorithms/cql/torch/cql_torch_learner.py

@@ -0,0 +1,280 @@
+from typing import Dict
+
+from ray.rllib.algorithms.cql.cql import CQLConfig
+from ray.rllib.algorithms.sac.sac_learner import (
+    LOGPS_KEY,
+    QF_LOSS_KEY,
+    QF_MAX_KEY,
+    QF_MEAN_KEY,
+    QF_MIN_KEY,
+    QF_PREDS,
+    QF_TWIN_LOSS_KEY,
+    QF_TWIN_PREDS,
+    TD_ERROR_MEAN_KEY,
+)
+from ray.rllib.algorithms.sac.torch.sac_torch_learner import SACTorchLearner
+from ray.rllib.core.columns import Columns
+from ray.rllib.core.learner.learner import (
+    POLICY_LOSS_KEY,
+)
+from ray.rllib.utils.annotations import override
+from ray.rllib.utils.framework import try_import_torch
+from ray.rllib.utils.metrics import ALL_MODULES
+from ray.rllib.utils.typing import ModuleID, ParamDict, TensorType
+from ray.tune.result import TRAINING_ITERATION
+
+torch, nn = try_import_torch()
+
+
+class CQLTorchLearner(SACTorchLearner):
+    @override(SACTorchLearner)
+    def compute_loss_for_module(
+        self,
+        *,
+        module_id: ModuleID,
+        config: CQLConfig,
+        batch: Dict,
+        fwd_out: Dict[str, TensorType],
+    ) -> TensorType:
+
+        # TODO (simon, sven): Add upstream information pieces into this timesteps
+        #  call arg to Learner.update_...().
+        self.metrics.log_value(
+            (ALL_MODULES, TRAINING_ITERATION),
+            1,
+            reduce="sum",
+        )
+        # Get the train action distribution for the current policy and current state.
+        # This is needed for the policy (actor) loss and the `alpha`` loss.
+        action_dist_class = self.module[module_id].get_train_action_dist_cls()
+        action_dist_curr = action_dist_class.from_logits(
+            fwd_out[Columns.ACTION_DIST_INPUTS]
+        )
+
+        # Optimize also the hyperparameter `alpha` by using the current policy
+        # evaluated at the current state (from offline data). Note, in contrast
+        # to the original SAC loss, here the `alpha` and actor losses are
+        # calculated first.
+        # TODO (simon): Check, why log(alpha) is used, prob. just better
+        # to optimize and monotonic function. Original equation uses alpha.
+        alpha_loss = -torch.mean(
+            self.curr_log_alpha[module_id]
+            * (fwd_out["logp_resampled"].detach() + self.target_entropy[module_id])
+        )
+
+        # Get the current alpha.
+        alpha = torch.exp(self.curr_log_alpha[module_id])
+        # Start training with behavior cloning and turn to the classic Soft-Actor Critic
+        # after `bc_iters` of training iterations.
+        if (
+            self.metrics.peek((ALL_MODULES, TRAINING_ITERATION), default=0)
+            >= config.bc_iters
+        ):
+            actor_loss = torch.mean(
+                alpha.detach() * fwd_out["logp_resampled"] - fwd_out["q_curr"]
+            )
+        else:
+            # Use log-probabilities of the current action distribution to clone
+            # the behavior policy (selected actions in data) in the first `bc_iters`
+            # training iterations.
+            bc_logps_curr = action_dist_curr.logp(batch[Columns.ACTIONS])
+            actor_loss = torch.mean(
+                alpha.detach() * fwd_out["logp_resampled"] - bc_logps_curr
+            )
+
+        # The critic loss is composed of the standard SAC Critic L2 loss and the
+        # CQL entropy loss.
+
+        # Get the Q-values for the actually selected actions in the offline data.
+        # In the critic loss we use these as predictions.
+        q_selected = fwd_out[QF_PREDS]
+        if config.twin_q:
+            q_twin_selected = fwd_out[QF_TWIN_PREDS]
+
+        if not config.deterministic_backup:
+            q_next = (
+                fwd_out["q_target_next"]
+                - alpha.detach() * fwd_out["logp_next_resampled"]
+            )
+        else:
+            q_next = fwd_out["q_target_next"]
+
+        # Now mask all Q-values with terminating next states in the targets.
+        q_next_masked = (1.0 - batch[Columns.TERMINATEDS].float()) * q_next
+
+        # Compute the right hand side of the Bellman equation. Detach this node
+        # from the computation graph as we do not want to backpropagate through
+        # the target network when optimizing the Q loss.
+        # TODO (simon, sven): Kumar et al. (2020) use here also a reward scaler.
+        q_selected_target = (
+            # TODO (simon): Add an `n_step` option to the `AddNextObsToBatch` connector.
+            batch[Columns.REWARDS]
+            # TODO (simon): Implement n_step.
+            + (config.gamma) * q_next_masked
+        ).detach()
+
+        # Calculate the TD error.
+        td_error = torch.abs(q_selected - q_selected_target)
+        # Calculate a TD-error for twin-Q values, if needed.
+        if config.twin_q:
+            td_error += torch.abs(q_twin_selected - q_selected_target)
+            # Rescale the TD error
+            td_error *= 0.5
+
+        # MSBE loss for the critic(s) (i.e. Q, see eqs. (7-8) Haarnoja et al. (2018)).
+        # Note, this needs a sample from the current policy given the next state.
+        # Note further, we could also use here the Huber loss instead of the MSE.
+        # TODO (simon): Add the huber loss as an alternative (SAC uses it).
+        sac_critic_loss = torch.nn.MSELoss(reduction="mean")(
+            q_selected,
+            q_selected_target,
+        )
+        if config.twin_q:
+            sac_critic_twin_loss = torch.nn.MSELoss(reduction="mean")(
+                q_twin_selected,
+                q_selected_target,
+            )
+
+        # Now calculate the CQL loss (we use the entropy version of the CQL algorithm).
+        # Note, the entropy version performs best in shown experiments.
+
+        # Compute the log-probabilities for the random actions (note, we generate random
+        # actions (from the mu distribution as named in Kumar et al. (2020))).
+        # Note, all actions, action log-probabilities and Q-values are already computed
+        # by the module's `_forward_train` method.
+        # TODO (simon): This is the density for a discrete uniform, however, actions
+        # come from a continuous one. So actually this density should use (1/(high-low))
+        # instead of (1/2).
+        random_density = torch.log(
+            torch.pow(
+                0.5,
+                torch.tensor(
+                    fwd_out["actions_curr_repeat"].shape[-1],
+                    device=fwd_out["actions_curr_repeat"].device,
+                ),
+            )
+        )
+        # Merge all Q-values and subtract the log-probabilities (note, we use the
+        # entropy version of CQL).
+        q_repeat = torch.cat(
+            [
+                fwd_out["q_rand_repeat"] - random_density,
+                fwd_out["q_next_repeat"] - fwd_out["logps_next_repeat"].detach(),
+                fwd_out["q_curr_repeat"] - fwd_out["logps_curr_repeat"].detach(),
+            ],
+            dim=1,
+        )
+        cql_loss = (
+            torch.logsumexp(q_repeat / config.temperature, dim=1).mean()
+            * config.min_q_weight
+            * config.temperature
+        )
+        cql_loss -= q_selected.mean() * config.min_q_weight
+        # Add the CQL loss term to the SAC loss term.
+        critic_loss = sac_critic_loss + cql_loss
+
+        # If a twin Q-value function is implemented calculated its CQL loss.
+        if config.twin_q:
+            q_twin_repeat = torch.cat(
+                [
+                    fwd_out["q_twin_rand_repeat"] - random_density,
+                    fwd_out["q_twin_next_repeat"]
+                    - fwd_out["logps_next_repeat"].detach(),
+                    fwd_out["q_twin_curr_repeat"]
+                    - fwd_out["logps_curr_repeat"].detach(),
+                ],
+                dim=1,
+            )
+            cql_twin_loss = (
+                torch.logsumexp(q_twin_repeat / config.temperature, dim=1).mean()
+                * config.min_q_weight
+                * config.temperature
+            )
+            cql_twin_loss -= q_twin_selected.mean() * config.min_q_weight
+            # Add the CQL loss term to the SAC loss term.
+            critic_twin_loss = sac_critic_twin_loss + cql_twin_loss
+
+        # TODO (simon): Check, if we need to implement here also a Lagrangian
+        # loss.
+
+        total_loss = actor_loss + critic_loss + alpha_loss
+
+        # Add the twin critic loss to the total loss, if needed.
+        if config.twin_q:
+            # Reweigh the critic loss terms in the total loss.
+            total_loss += 0.5 * critic_twin_loss - 0.5 * critic_loss
+
+        # Log important loss stats (reduce=mean (default), but with window=1
+        # in order to keep them history free).
+        self.metrics.log_dict(
+            {
+                POLICY_LOSS_KEY: actor_loss,
+                QF_LOSS_KEY: critic_loss,
+                # TODO (simon): Add these keys to SAC Learner.
+                "cql_loss": cql_loss,
+                "alpha_loss": alpha_loss,
+                "alpha_value": alpha[0],
+                "log_alpha_value": torch.log(alpha)[0],
+                "target_entropy": self.target_entropy[module_id],
+                LOGPS_KEY: torch.mean(
+                    fwd_out["logp_resampled"]
+                ),  # torch.mean(logps_curr),
+                QF_MEAN_KEY: torch.mean(fwd_out["q_curr_repeat"]),
+                QF_MAX_KEY: torch.max(fwd_out["q_curr_repeat"]),
+                QF_MIN_KEY: torch.min(fwd_out["q_curr_repeat"]),
+                TD_ERROR_MEAN_KEY: torch.mean(td_error),
+            },
+            key=module_id,
+            window=1,  # <- single items (should not be mean/ema-reduced over time).
+        )
+        self._temp_losses[(module_id, POLICY_LOSS_KEY)] = actor_loss
+        self._temp_losses[(module_id, QF_LOSS_KEY)] = critic_loss
+        self._temp_losses[(module_id, "alpha_loss")] = alpha_loss
+
+        # TODO (simon): Add loss keys for langrangian, if needed.
+        # TODO (simon): Add only here then the Langrange parameter optimization.
+        if config.twin_q:
+            self.metrics.log_value(
+                key=(module_id, QF_TWIN_LOSS_KEY),
+                value=critic_twin_loss,
+                window=1,  # <- single items (should not be mean/ema-reduced over time).
+            )
+            self._temp_losses[(module_id, QF_TWIN_LOSS_KEY)] = critic_twin_loss
+
+        # Return the total loss.
+        return total_loss
+
+    @override(SACTorchLearner)
+    def compute_gradients(
+        self, loss_per_module: Dict[ModuleID, TensorType], **kwargs
+    ) -> ParamDict:
+
+        grads = {}
+        for module_id in set(loss_per_module.keys()) - {ALL_MODULES}:
+            # Loop through optimizers registered for this module.
+            for optim_name, optim in self.get_optimizers_for_module(module_id):
+                # Zero the gradients. Note, we need to reset the gradients b/c
+                # each component for a module operates on the same graph.
+                optim.zero_grad(set_to_none=True)
+
+                # Compute the gradients for the component and module.
+                loss_tensor = self._temp_losses.pop((module_id, optim_name + "_loss"))
+                loss_tensor.backward(
+                    retain_graph=False if optim_name in ["policy", "alpha"] else True
+                )
+                # Store the gradients for the component and module.
+                # TODO (simon): Check another time the graph for overlapping
+                # gradients.
+                grads.update(
+                    {
+                        pid: grads[pid] + p.grad.clone()
+                        if pid in grads
+                        else p.grad.clone()
+                        for pid, p in self.filter_param_dict_for_optimizer(
+                            self._params, optim
+                        ).items()
+                    }
+                )
+
+        assert not self._temp_losses
+        return grads

rllib/algorithms/cql/torch/default_cql_torch_rl_module.py

@@ -0,0 +1,207 @@
+from typing import Any, Dict, Optional
+
+import tree
+
+from ray.rllib.algorithms.sac.sac_catalog import SACCatalog
+from ray.rllib.algorithms.sac.sac_learner import (
+    QF_PREDS,
+    QF_TWIN_PREDS,
+)
+from ray.rllib.algorithms.sac.torch.default_sac_torch_rl_module import (
+    DefaultSACTorchRLModule,
+)
+from ray.rllib.core.columns import Columns
+from ray.rllib.core.models.base import ENCODER_OUT
+from ray.rllib.utils.annotations import override
+from ray.rllib.utils.framework import try_import_torch
+from ray.rllib.utils.typing import TensorType
+
+torch, nn = try_import_torch()
+
+
+class DefaultCQLTorchRLModule(DefaultSACTorchRLModule):
+    def __init__(self, *args, **kwargs):
+        catalog_class = kwargs.pop("catalog_class", None)
+        if catalog_class is None:
+            catalog_class = SACCatalog
+        super().__init__(*args, **kwargs, catalog_class=catalog_class)
+
+    @override(DefaultSACTorchRLModule)
+    def _forward_train(self, batch: Dict) -> Dict[str, Any]:
+        # Call the super method.
+        fwd_out = super()._forward_train(batch)
+
+        # Make sure we perform a "straight-through gradient" pass here,
+        # ignoring the gradients of the q-net, however, still recording
+        # the gradients of the policy net (which was used to rsample the actions used
+        # here). This is different from doing `.detach()` or `with torch.no_grads()`,
+        # as these two methds would fully block all gradient recordings, including
+        # the needed policy ones.
+        all_params = list(self.pi_encoder.parameters()) + list(self.pi.parameters())
+        # if self.twin_q:
+        #     all_params += list(self.qf_twin.parameters()) + list(
+        #         self.qf_twin_encoder.parameters()
+        #     )
+
+        for param in all_params:
+            param.requires_grad = False
+
+        # Compute the repeated actions, action log-probabilites and Q-values for all
+        # observations.
+        # First for the random actions (from the mu-distribution as named by Kumar et
+        # al. (2020)).
+        low = torch.tensor(
+            self.action_space.low,
+            device=fwd_out[QF_PREDS].device,
+        )
+        high = torch.tensor(
+            self.action_space.high,
+            device=fwd_out[QF_PREDS].device,
+        )
+        num_samples = batch[Columns.ACTIONS].shape[0] * self.model_config["num_actions"]
+        actions_rand_repeat = low + (high - low) * torch.rand(
+            (num_samples, low.shape[0]), device=fwd_out[QF_PREDS].device
+        )
+
+        # First for the random actions (from the mu-distribution as named in Kumar
+        # et al. (2020)) using repeated observations.
+        rand_repeat_out = self._repeat_actions(batch[Columns.OBS], actions_rand_repeat)
+        (fwd_out["actions_rand_repeat"], fwd_out["q_rand_repeat"]) = (
+            rand_repeat_out[Columns.ACTIONS],
+            rand_repeat_out[QF_PREDS],
+        )
+        # Sample current and next actions (from the pi distribution as named in Kumar
+        # et al. (2020)) using repeated observations
+        # Second for the current observations and the current action distribution.
+        curr_repeat_out = self._repeat_actions(batch[Columns.OBS])
+        (
+            fwd_out["actions_curr_repeat"],
+            fwd_out["logps_curr_repeat"],
+            fwd_out["q_curr_repeat"],
+        ) = (
+            curr_repeat_out[Columns.ACTIONS],
+            curr_repeat_out[Columns.ACTION_LOGP],
+            curr_repeat_out[QF_PREDS],
+        )
+        # Then, for the next observations and the current action distribution.
+        next_repeat_out = self._repeat_actions(batch[Columns.NEXT_OBS])
+        (
+            fwd_out["actions_next_repeat"],
+            fwd_out["logps_next_repeat"],
+            fwd_out["q_next_repeat"],
+        ) = (
+            next_repeat_out[Columns.ACTIONS],
+            next_repeat_out[Columns.ACTION_LOGP],
+            next_repeat_out[QF_PREDS],
+        )
+        if self.twin_q:
+            # First for the random actions from the mu-distribution.
+            fwd_out["q_twin_rand_repeat"] = rand_repeat_out[QF_TWIN_PREDS]
+            # Second for the current observations and the current action distribution.
+            fwd_out["q_twin_curr_repeat"] = curr_repeat_out[QF_TWIN_PREDS]
+            # Then, for the next observations and the current action distribution.
+            fwd_out["q_twin_next_repeat"] = next_repeat_out[QF_TWIN_PREDS]
+        # Reset the gradient requirements for all Q-function parameters.
+        for param in all_params:
+            param.requires_grad = True
+
+        return fwd_out
+
+    def _repeat_tensor(self, tensor: TensorType, repeat: int) -> TensorType:
+        """Generates a repeated version of a tensor.
+
+        The repetition is done similar `np.repeat` and repeats each value
+        instead of the complete vector.
+
+        Args:
+            tensor: The tensor to be repeated.
+            repeat: How often each value in the tensor should be repeated.
+
+        Returns:
+            A tensor holding `repeat`  repeated values of the input `tensor`
+        """
+        # Insert the new dimension at axis 1 into the tensor.
+        t_repeat = tensor.unsqueeze(1)
+        # Repeat the tensor along the new dimension.
+        t_repeat = torch.repeat_interleave(t_repeat, repeat, dim=1)
+        # Stack the repeated values into the batch dimension.
+        t_repeat = t_repeat.view(-1, *tensor.shape[1:])
+        # Return the repeated tensor.
+        return t_repeat
+
+    def _repeat_actions(
+        self, obs: TensorType, actions: Optional[TensorType] = None
+    ) -> Dict[str, TensorType]:
+        """Generated actions and Q-values for repeated observations.
+
+        The `self.model_config["num_actions"]` define a multiplier
+        used for generating `num_actions` as many actions as the batch size.
+        Observations are repeated and then a model forward pass is made.
+
+        Args:
+            obs: A batched observation tensor.
+            actions: An optional batched actions tensor.
+
+        Returns:
+            A dictionary holding the (sampled or passed-in actions), the log
+            probabilities (of sampled actions), the Q-values and if available
+            the twin-Q values.
+        """
+        output = {}
+        # Receive the batch size.
+        batch_size = obs.shape[0]
+        # Receive the number of action to sample.
+        num_actions = self.model_config["num_actions"]
+        # Repeat the observations `num_actions` times.
+        obs_repeat = tree.map_structure(
+            lambda t: self._repeat_tensor(t, num_actions), obs
+        )
+        # Generate a batch for the forward pass.
+        temp_batch = {Columns.OBS: obs_repeat}
+        if actions is None:
+            # TODO (simon): Run the forward pass in inference mode.
+            # Compute the action logits.
+            pi_encoder_outs = self.pi_encoder(temp_batch)
+            action_logits = self.pi(pi_encoder_outs[ENCODER_OUT])
+            # Generate the squashed Gaussian from the model's logits.
+            action_dist = self.get_train_action_dist_cls().from_logits(action_logits)
+            # Sample the actions. Note, we want to make a backward pass through
+            # these actions.
+            output[Columns.ACTIONS] = action_dist.rsample()
+            # Compute the action log-probabilities.
+            output[Columns.ACTION_LOGP] = action_dist.logp(
+                output[Columns.ACTIONS]
+            ).view(batch_size, num_actions, 1)
+        else:
+            output[Columns.ACTIONS] = actions
+
+        # Compute all Q-values.
+        temp_batch.update(
+            {
+                Columns.ACTIONS: output[Columns.ACTIONS],
+            }
+        )
+        output.update(
+            {
+                QF_PREDS: self._qf_forward_train_helper(
+                    temp_batch,
+                    self.qf_encoder,
+                    self.qf,
+                ).view(batch_size, num_actions, 1)
+            }
+        )
+        # If we have a twin-Q network, compute its Q-values, too.
+        if self.twin_q:
+            output.update(
+                {
+                    QF_TWIN_PREDS: self._qf_forward_train_helper(
+                        temp_batch,
+                        self.qf_twin_encoder,
+                        self.qf_twin,
+                    ).view(batch_size, num_actions, 1)
+                }
+            )
+        del temp_batch
+
+        # Return
+        return output

rllib/algorithms/dqn/README.md

@@ -0,0 +1,69 @@
+# Deep Q Networks (DQN)
+
+Code in this package is adapted from https://github.com/openai/baselines/tree/master/baselines/deepq.
+
+
+## Overview
+
+[DQN](https://www.cs.toronto.edu/~vmnih/docs/dqn.pdf) is a model-free off-policy RL
+algorithm and one of the first deep RL algorithms developed. DQN proposes using a
+neural network as a function approximator for the Q-function in Q-learning.
+The algorithm aims to minimize the L2 norm between the Q-value predictions and the
+Q-value targets, which is computed as 1-step TD. The paper proposes two important concepts,
+a target network and an experience replay buffer. The target network is a copy of the
+main Q network and is used to compute Q-value targets for loss-function calculations.
+To stabilize training, the target network lags slightly behind the main Q-network.
+Meanwhile, the experience replay stores all data encountered by the agent during
+training and is uniformly sampled from to generate gradient updates for the Q-value network.
+
+
+## Supported DQN Algorithms
+
+[Double DQN](https://arxiv.org/pdf/1509.06461.pdf) - As opposed to learning one Q network in vanilla DQN, Double DQN proposes learning two Q networks akin to double Q-learning. As a solution, Double DQN aims to solve the issue of vanilla DQN's overly-optimistic Q-values, which limits performance.
+
+[Dueling DQN](https://arxiv.org/pdf/1511.06581.pdf) - Dueling DQN proposes splitting learning a Q-value function approximator into learning two networks: a value and advantage approximator.
+
+[Distributional DQN](https://arxiv.org/pdf/1707.06887.pdf) - Usually, the Q network outputs the predicted Q-value of a state-action pair. Distributional DQN takes this further by predicting the distribution of Q-values (e.g. mean and std of a normal distribution) of a state-action pair. Doing this captures uncertainty of the Q-value and can improve the performance of DQN algorithms.
+
+[APEX-DQN](https://arxiv.org/pdf/1803.00933.pdf) - Standard DQN algorithms propose using a experience replay buffer to sample data uniformly and compute gradients from the sampled data. APEX introduces the notion of weighted replay data, where elements in the replay buffer are more or less likely to be sampled depending on the TD-error.
+
+[Rainbow](https://arxiv.org/pdf/1710.02298.pdf) - Rainbow DQN, as the word Rainbow suggests, aggregates the many improvements discovered in research to improve DQN performance. This includes a multi-step distributional loss (extended from Distributional DQN), prioritized replay (inspired from APEX-DQN), double Q-networks (inspired from Double DQN), and dueling networks (inspired from Dueling DQN).
+
+
+## Documentation & Implementation:
+
+1) Vanilla DQN (DQN).
+
+    **[Detailed Documentation](https://docs.ray.io/en/master/rllib-algorithms.html#dqn)**
+
+    **[Implementation](https://github.com/ray-project/ray/blob/master/rllib/algorithms/dqn/simple_q.py)**
+
+2) Double DQN.
+
+    **[Detailed Documentation](https://docs.ray.io/en/master/rllib-algorithms.html#dqn)**
+
+    **[Implementation](https://github.com/ray-project/ray/blob/master/rllib/algorithms/dqn/dqn.py)**
+
+3) Dueling DQN
+
+    **[Detailed Documentation](https://docs.ray.io/en/master/rllib-algorithms.html#dqn)**
+
+    **[Implementation](https://github.com/ray-project/ray/blob/master/rllib/algorithms/dqn/dqn.py)**
+
+3) Distributional DQN
+
+    **[Detailed Documentation](https://docs.ray.io/en/master/rllib-algorithms.html#dqn)**
+
+    **[Implementation](https://github.com/ray-project/ray/blob/master/rllib/algorithms/dqn/dqn.py)**
+
+4) APEX DQN
+
+    **[Detailed Documentation](https://docs.ray.io/en/master/rllib-algorithms.html#dqn)**
+
+    **[Implementation](https://github.com/ray-project/ray/blob/master/rllib/agents/dqn/apex.py)**
+
+5) Rainbow DQN
+
+    **[Detailed Documentation](https://docs.ray.io/en/master/rllib-algorithms.html#dqn)**
+
+    **[Implementation](https://github.com/ray-project/ray/blob/master/rllib/algorithms/dqn/dqn.py)**

rllib/algorithms/dqn/__init__.py

@@ -0,0 +1,10 @@
+from ray.rllib.algorithms.dqn.dqn import DQN, DQNConfig
+from ray.rllib.algorithms.dqn.dqn_tf_policy import DQNTFPolicy
+from ray.rllib.algorithms.dqn.dqn_torch_policy import DQNTorchPolicy
+
+__all__ = [
+    "DQN",
+    "DQNConfig",
+    "DQNTFPolicy",
+    "DQNTorchPolicy",
+]

rllib/algorithms/dqn/default_dqn_rl_module.py

@@ -0,0 +1,166 @@
+import abc
+from typing import Any, Dict, List, Tuple, Union
+
+from ray.rllib.core.learner.utils import make_target_network
+from ray.rllib.core.models.base import Encoder, Model
+from ray.rllib.core.rl_module.apis import InferenceOnlyAPI, QNetAPI, TargetNetworkAPI
+from ray.rllib.core.rl_module.rl_module import RLModule
+from ray.rllib.utils.annotations import (
+    OverrideToImplementCustomLogic,
+    override,
+)
+from ray.rllib.utils.schedules.scheduler import Scheduler
+from ray.rllib.utils.typing import NetworkType, TensorType
+from ray.util.annotations import DeveloperAPI
+
+QF_PREDS = "qf_preds"
+ATOMS = "atoms"
+QF_LOGITS = "qf_logits"
+QF_NEXT_PREDS = "qf_next_preds"
+QF_PROBS = "qf_probs"
+QF_TARGET_NEXT_PREDS = "qf_target_next_preds"
+QF_TARGET_NEXT_PROBS = "qf_target_next_probs"
+
+
+@DeveloperAPI
+class DefaultDQNRLModule(RLModule, InferenceOnlyAPI, TargetNetworkAPI, QNetAPI):
+    @override(RLModule)
+    def setup(self):
+        # If a dueling architecture is used.
+        self.uses_dueling: bool = self.model_config.get("dueling")
+        # If double Q learning is used.
+        self.uses_double_q: bool = self.model_config.get("double_q")
+        # The number of atoms for a distribution support.
+        self.num_atoms: int = self.model_config.get("num_atoms")
+        # If distributional learning is requested configure the support.
+        if self.num_atoms > 1:
+            self.v_min: float = self.model_config.get("v_min")
+            self.v_max: float = self.model_config.get("v_max")
+        # The epsilon scheduler for epsilon greedy exploration.
+        self.epsilon_schedule = Scheduler(
+            fixed_value_or_schedule=self.model_config["epsilon"],
+            framework=self.framework,
+        )
+
+        # Build the encoder for the advantage and value streams. Note,
+        # the same encoder is used.
+        # Note further, by using the base encoder the correct encoder
+        # is chosen for the observation space used.
+        self.encoder = self.catalog.build_encoder(framework=self.framework)
+
+        # Build heads.
+        self.af = self.catalog.build_af_head(framework=self.framework)
+        if self.uses_dueling:
+            # If in a dueling setting setup the value function head.
+            self.vf = self.catalog.build_vf_head(framework=self.framework)
+
+    @override(InferenceOnlyAPI)
+    def get_non_inference_attributes(self) -> List[str]:
+        return ["_target_encoder", "_target_af"] + (
+            ["_target_vf"] if self.uses_dueling else []
+        )
+
+    @override(TargetNetworkAPI)
+    def make_target_networks(self) -> None:
+        self._target_encoder = make_target_network(self.encoder)
+        self._target_af = make_target_network(self.af)
+        if self.uses_dueling:
+            self._target_vf = make_target_network(self.vf)
+
+    @override(TargetNetworkAPI)
+    def get_target_network_pairs(self) -> List[Tuple[NetworkType, NetworkType]]:
+        return [(self.encoder, self._target_encoder), (self.af, self._target_af)] + (
+            # If we have a dueling architecture we need to update the value stream
+            # target, too.
+            [
+                (self.vf, self._target_vf),
+            ]
+            if self.uses_dueling
+            else []
+        )
+
+    @override(TargetNetworkAPI)
+    def forward_target(self, batch: Dict[str, Any]) -> Dict[str, Any]:
+        """Computes Q-values from the target network.
+
+        Note, these can be accompanied by logits and probabilities
+        in case of distributional Q-learning, i.e. `self.num_atoms > 1`.
+
+        Args:
+            batch: The batch received in the forward pass.
+
+        Results:
+            A dictionary containing the target Q-value predictions ("qf_preds")
+            and in case of distributional Q-learning in addition to the target
+            Q-value predictions ("qf_preds") the support atoms ("atoms"), the target
+            Q-logits  ("qf_logits"), and the probabilities ("qf_probs").
+        """
+        # If we have a dueling architecture we have to add the value stream.
+        return self._qf_forward_helper(
+            batch,
+            self._target_encoder,
+            (
+                {"af": self._target_af, "vf": self._target_vf}
+                if self.uses_dueling
+                else self._target_af
+            ),
+        )
+
+    @override(QNetAPI)
+    def compute_q_values(self, batch: Dict[str, TensorType]) -> Dict[str, TensorType]:
+        """Computes Q-values, given encoder, q-net and (optionally), advantage net.
+
+        Note, these can be accompanied by logits and probabilities
+        in case of distributional Q-learning, i.e. `self.num_atoms > 1`.
+
+        Args:
+            batch: The batch received in the forward pass.
+
+        Results:
+            A dictionary containing the Q-value predictions ("qf_preds")
+            and in case of distributional Q-learning - in addition to the Q-value
+            predictions ("qf_preds") - the support atoms ("atoms"), the Q-logits
+            ("qf_logits"), and the probabilities ("qf_probs").
+        """
+        # If we have a dueling architecture we have to add the value stream.
+        return self._qf_forward_helper(
+            batch,
+            self.encoder,
+            {"af": self.af, "vf": self.vf} if self.uses_dueling else self.af,
+        )
+
+    @override(RLModule)
+    def get_initial_state(self) -> dict:
+        if hasattr(self.encoder, "get_initial_state"):
+            return self.encoder.get_initial_state()
+        else:
+            return {}
+
+    @abc.abstractmethod
+    @OverrideToImplementCustomLogic
+    def _qf_forward_helper(
+        self,
+        batch: Dict[str, TensorType],
+        encoder: Encoder,
+        head: Union[Model, Dict[str, Model]],
+    ) -> Dict[str, TensorType]:
+        """Computes Q-values.
+
+        This is a helper function that takes care of all different cases,
+        i.e. if we use a dueling architecture or not and if we use distributional
+        Q-learning or not.
+
+        Args:
+            batch: The batch received in the forward pass.
+            encoder: The encoder network to use. Here we have a single encoder
+                for all heads (Q or advantages and value in case of a dueling
+                architecture).
+            head: Either a head model or a dictionary of head model (dueling
+                architecture) containing advantage and value stream heads.
+
+        Returns:
+            In case of expectation learning the Q-value predictions ("qf_preds")
+            and in case of distributional Q-learning in addition to the predictions
+            the atoms ("atoms"), the Q-value predictions ("qf_preds"), the Q-logits
+            ("qf_logits") and the probabilities for the support atoms ("qf_probs").
+        """

rllib/algorithms/dqn/distributional_q_tf_model.py

@@ -0,0 +1,191 @@
+"""Tensorflow model for DQN"""
+
+from typing import List
+
+import gymnasium as gym
+
+from ray.rllib.models.tf.layers import NoisyLayer
+from ray.rllib.models.tf.tf_modelv2 import TFModelV2
+from ray.rllib.utils.annotations import OldAPIStack
+from ray.rllib.utils.framework import try_import_tf
+from ray.rllib.utils.typing import ModelConfigDict, TensorType
+
+tf1, tf, tfv = try_import_tf()
+
+
+@OldAPIStack
+class DistributionalQTFModel(TFModelV2):
+    """Extension of standard TFModel to provide distributional Q values.
+
+    It also supports options for noisy nets and parameter space noise.
+
+    Data flow:
+        obs -> forward() -> model_out
+        model_out -> get_q_value_distributions() -> Q(s, a) atoms
+        model_out -> get_state_value() -> V(s)
+
+    Note that this class by itself is not a valid model unless you
+    implement forward() in a subclass."""
+
+    def __init__(
+        self,
+        obs_space: gym.spaces.Space,
+        action_space: gym.spaces.Space,
+        num_outputs: int,
+        model_config: ModelConfigDict,
+        name: str,
+        q_hiddens=(256,),
+        dueling: bool = False,
+        num_atoms: int = 1,
+        use_noisy: bool = False,
+        v_min: float = -10.0,
+        v_max: float = 10.0,
+        sigma0: float = 0.5,
+        # TODO(sven): Move `add_layer_norm` into ModelCatalog as
+        #  generic option, then error if we use ParameterNoise as
+        #  Exploration type and do not have any LayerNorm layers in
+        #  the net.
+        add_layer_norm: bool = False,
+    ):
+        """Initialize variables of this model.
+
+        Extra model kwargs:
+            q_hiddens (List[int]): List of layer-sizes after(!) the
+                Advantages(A)/Value(V)-split. Hence, each of the A- and V-
+                branches will have this structure of Dense layers. To define
+                the NN before this A/V-split, use - as always -
+                config["model"]["fcnet_hiddens"].
+            dueling: Whether to build the advantage(A)/value(V) heads
+                for DDQN. If True, Q-values are calculated as:
+                Q = (A - mean[A]) + V. If False, raw NN output is interpreted
+                as Q-values.
+            num_atoms: If >1, enables distributional DQN.
+            use_noisy: Use noisy nets.
+            v_min: Min value support for distributional DQN.
+            v_max: Max value support for distributional DQN.
+            sigma0 (float): Initial value of noisy layers.
+            add_layer_norm: Enable layer norm (for param noise).
+
+        Note that the core layers for forward() are not defined here, this
+        only defines the layers for the Q head. Those layers for forward()
+        should be defined in subclasses of DistributionalQModel.
+        """
+        super(DistributionalQTFModel, self).__init__(
+            obs_space, action_space, num_outputs, model_config, name
+        )
+
+        # setup the Q head output (i.e., model for get_q_values)
+        self.model_out = tf.keras.layers.Input(shape=(num_outputs,), name="model_out")
+
+        def build_action_value(prefix: str, model_out: TensorType) -> List[TensorType]:
+            if q_hiddens:
+                action_out = model_out
+                for i in range(len(q_hiddens)):
+                    if use_noisy:
+                        action_out = NoisyLayer(
+                            "{}hidden_{}".format(prefix, i), q_hiddens[i], sigma0
+                        )(action_out)
+                    elif add_layer_norm:
+                        action_out = tf.keras.layers.Dense(
+                            units=q_hiddens[i], activation=tf.nn.relu
+                        )(action_out)
+                        action_out = tf.keras.layers.LayerNormalization()(action_out)
+                    else:
+                        action_out = tf.keras.layers.Dense(
+                            units=q_hiddens[i],
+                            activation=tf.nn.relu,
+                            name="hidden_%d" % i,
+                        )(action_out)
+            else:
+                # Avoid postprocessing the outputs. This enables custom models
+                # to be used for parametric action DQN.
+                action_out = model_out
+
+            if use_noisy:
+                action_scores = NoisyLayer(
+                    "{}output".format(prefix),
+                    self.action_space.n * num_atoms,
+                    sigma0,
+                    activation=None,
+                )(action_out)
+            elif q_hiddens:
+                action_scores = tf.keras.layers.Dense(
+                    units=self.action_space.n * num_atoms, activation=None
+                )(action_out)
+            else:
+                action_scores = model_out
+
+            if num_atoms > 1:
+                # Distributional Q-learning uses a discrete support z
+                # to represent the action value distribution
+                z = tf.range(num_atoms, dtype=tf.float32)
+                z = v_min + z * (v_max - v_min) / float(num_atoms - 1)
+
+                def _layer(x):
+                    support_logits_per_action = tf.reshape(
+                        tensor=x, shape=(-1, self.action_space.n, num_atoms)
+                    )
+                    support_prob_per_action = tf.nn.softmax(
+                        logits=support_logits_per_action
+                    )
+                    x = tf.reduce_sum(input_tensor=z * support_prob_per_action, axis=-1)
+                    logits = support_logits_per_action
+                    dist = support_prob_per_action
+                    return [x, z, support_logits_per_action, logits, dist]
+
+                return tf.keras.layers.Lambda(_layer)(action_scores)
+            else:
+                logits = tf.expand_dims(tf.ones_like(action_scores), -1)
+                dist = tf.expand_dims(tf.ones_like(action_scores), -1)
+                return [action_scores, logits, dist]
+
+        def build_state_score(prefix: str, model_out: TensorType) -> TensorType:
+            state_out = model_out
+            for i in range(len(q_hiddens)):
+                if use_noisy:
+                    state_out = NoisyLayer(
+                        "{}dueling_hidden_{}".format(prefix, i), q_hiddens[i], sigma0
+                    )(state_out)
+                else:
+                    state_out = tf.keras.layers.Dense(
+                        units=q_hiddens[i], activation=tf.nn.relu
+                    )(state_out)
+                    if add_layer_norm:
+                        state_out = tf.keras.layers.LayerNormalization()(state_out)
+            if use_noisy:
+                state_score = NoisyLayer(
+                    "{}dueling_output".format(prefix),
+                    num_atoms,
+                    sigma0,
+                    activation=None,
+                )(state_out)
+            else:
+                state_score = tf.keras.layers.Dense(units=num_atoms, activation=None)(
+                    state_out
+                )
+            return state_score
+
+        q_out = build_action_value(name + "/action_value/", self.model_out)
+        self.q_value_head = tf.keras.Model(self.model_out, q_out)
+
+        if dueling:
+            state_out = build_state_score(name + "/state_value/", self.model_out)
+            self.state_value_head = tf.keras.Model(self.model_out, state_out)
+
+    def get_q_value_distributions(self, model_out: TensorType) -> List[TensorType]:
+        """Returns distributional values for Q(s, a) given a state embedding.
+
+        Override this in your custom model to customize the Q output head.
+
+        Args:
+            model_out: embedding from the model layers
+
+        Returns:
+            (action_scores, logits, dist) if num_atoms == 1, otherwise
+            (action_scores, z, support_logits_per_action, logits, dist)
+        """
+        return self.q_value_head(model_out)
+
+    def get_state_value(self, model_out: TensorType) -> TensorType:
+        """Returns the state value prediction for the given state embedding."""
+        return self.state_value_head(model_out)

rllib/algorithms/dqn/dqn.py

@@ -0,0 +1,859 @@
+"""
+Deep Q-Networks (DQN, Rainbow, Parametric DQN)
+==============================================
+
+This file defines the distributed Algorithm class for the Deep Q-Networks
+algorithm. See `dqn_[tf|torch]_policy.py` for the definition of the policies.
+
+Detailed documentation:
+https://docs.ray.io/en/master/rllib-algorithms.html#deep-q-networks-dqn-rainbow-parametric-dqn
+"""  # noqa: E501
+
+import logging
+from collections import defaultdict
+from typing import Any, Callable, Dict, List, Optional, Tuple, Type, Union
+
+import numpy as np
+from typing_extensions import Self
+
+from ray._common.deprecation import DEPRECATED_VALUE
+from ray.rllib.algorithms.algorithm import Algorithm
+from ray.rllib.algorithms.algorithm_config import AlgorithmConfig, NotProvided
+from ray.rllib.algorithms.dqn.dqn_tf_policy import DQNTFPolicy
+from ray.rllib.algorithms.dqn.dqn_torch_policy import DQNTorchPolicy
+from ray.rllib.core.learner import Learner
+from ray.rllib.core.rl_module.rl_module import RLModuleSpec
+from ray.rllib.execution.rollout_ops import (
+    synchronous_parallel_sample,
+)
+from ray.rllib.execution.train_ops import (
+    multi_gpu_train_one_step,
+    train_one_step,
+)
+from ray.rllib.policy.policy import Policy
+from ray.rllib.policy.sample_batch import MultiAgentBatch
+from ray.rllib.utils import deep_update
+from ray.rllib.utils.annotations import override
+from ray.rllib.utils.metrics import (
+    ALL_MODULES,
+    ENV_RUNNER_RESULTS,
+    ENV_RUNNER_SAMPLING_TIMER,
+    LAST_TARGET_UPDATE_TS,
+    LEARNER_RESULTS,
+    LEARNER_UPDATE_TIMER,
+    NUM_AGENT_STEPS_SAMPLED,
+    NUM_AGENT_STEPS_SAMPLED_LIFETIME,
+    NUM_ENV_STEPS_SAMPLED,
+    NUM_ENV_STEPS_SAMPLED_LIFETIME,
+    NUM_TARGET_UPDATES,
+    REPLAY_BUFFER_ADD_DATA_TIMER,
+    REPLAY_BUFFER_RESULTS,
+    REPLAY_BUFFER_SAMPLE_TIMER,
+    REPLAY_BUFFER_UPDATE_PRIOS_TIMER,
+    SAMPLE_TIMER,
+    SYNCH_WORKER_WEIGHTS_TIMER,
+    TD_ERROR_KEY,
+    TIMERS,
+)
+from ray.rllib.utils.numpy import convert_to_numpy
+from ray.rllib.utils.replay_buffers.utils import (
+    sample_min_n_steps_from_buffer,
+    update_priorities_in_episode_replay_buffer,
+    update_priorities_in_replay_buffer,
+    validate_buffer_config,
+)
+from ray.rllib.utils.typing import (
+    LearningRateOrSchedule,
+    ResultDict,
+    RLModuleSpecType,
+    SampleBatchType,
+)
+
+logger = logging.getLogger(__name__)
+
+
+class DQNConfig(AlgorithmConfig):
+    r"""Defines a configuration class from which a DQN Algorithm can be built.
+
+    .. testcode::
+
+        from ray.rllib.algorithms.dqn.dqn import DQNConfig
+
+        config = (
+            DQNConfig()
+            .environment("CartPole-v1")
+            .training(replay_buffer_config={
+                "type": "PrioritizedEpisodeReplayBuffer",
+                "capacity": 60000,
+                "alpha": 0.5,
+                "beta": 0.5,
+            })
+            .env_runners(num_env_runners=1)
+        )
+        algo = config.build()
+        algo.train()
+        algo.stop()
+
+    .. testcode::
+
+        from ray.rllib.algorithms.dqn.dqn import DQNConfig
+        from ray import tune
+
+        config = (
+            DQNConfig()
+            .environment("CartPole-v1")
+            .training(
+                num_atoms=tune.grid_search([1,])
+            )
+        )
+        tune.Tuner(
+            "DQN",
+            run_config=tune.RunConfig(stop={"training_iteration":1}),
+            param_space=config,
+        ).fit()
+
+    .. testoutput::
+        :hide:
+
+        ...
+
+
+    """
+
+    def __init__(self, algo_class=None):
+        """Initializes a DQNConfig instance."""
+        self.exploration_config = {
+            "type": "EpsilonGreedy",
+            "initial_epsilon": 1.0,
+            "final_epsilon": 0.02,
+            "epsilon_timesteps": 10000,
+        }
+
+        super().__init__(algo_class=algo_class or DQN)
+
+        # Overrides of AlgorithmConfig defaults
+        # `env_runners()`
+        # Set to `self.n_step`, if 'auto'.
+        self.rollout_fragment_length: Union[int, str] = "auto"
+        # New stack uses `epsilon` as either a constant value or a scheduler
+        # defined like this.
+        # TODO (simon): Ensure that users can understand how to provide epsilon.
+        #  (sven): Should we add this to `self.env_runners(epsilon=..)`?
+        self.epsilon = [(0, 1.0), (10000, 0.05)]
+
+        # `training()`
+        self.grad_clip = 40.0
+        # Note: Only when using enable_rl_module_and_learner=True can the clipping mode
+        # be configured by the user. On the old API stack, RLlib will always clip by
+        # global_norm, no matter the value of `grad_clip_by`.
+        self.grad_clip_by = "global_norm"
+        self.lr = 5e-4
+        self.train_batch_size = 32
+
+        # `evaluation()`
+        self.evaluation(evaluation_config=AlgorithmConfig.overrides(explore=False))
+
+        # `reporting()`
+        self.min_time_s_per_iteration = None
+        self.min_sample_timesteps_per_iteration = 1000
+
+        # DQN specific config settings.
+        # fmt: off
+        # __sphinx_doc_begin__
+        self.target_network_update_freq = 500
+        self.num_steps_sampled_before_learning_starts = 1000
+        self.store_buffer_in_checkpoints = False
+        self.adam_epsilon = 1e-8
+
+        self.tau = 1.0
+
+        self.num_atoms = 1
+        self.v_min = -10.0
+        self.v_max = 10.0
+        self.noisy = False
+        self.sigma0 = 0.5
+        self.dueling = True
+        self.hiddens = [256]
+        self.double_q = True
+        self.n_step = 1
+        self.before_learn_on_batch = None
+        self.training_intensity = None
+        self.td_error_loss_fn = "huber"
+        self.categorical_distribution_temperature = 1.0
+        # The burn-in for stateful `RLModule`s.
+        self.burn_in_len = 0
+
+        # Replay buffer configuration.
+        self.replay_buffer_config = {
+            "type": "PrioritizedEpisodeReplayBuffer",
+            # Size of the replay buffer. Note that if async_updates is set,
+            # then each worker will have a replay buffer of this size.
+            "capacity": 50000,
+            "alpha": 0.6,
+            # Beta parameter for sampling from prioritized replay buffer.
+            "beta": 0.4,
+        }
+        # fmt: on
+        # __sphinx_doc_end__
+
+        self.lr_schedule = None  # @OldAPIStack
+
+        # Deprecated
+        self.buffer_size = DEPRECATED_VALUE
+        self.prioritized_replay = DEPRECATED_VALUE
+        self.learning_starts = DEPRECATED_VALUE
+        self.replay_batch_size = DEPRECATED_VALUE
+        # Can not use DEPRECATED_VALUE here because -1 is a common config value
+        self.replay_sequence_length = None
+        self.prioritized_replay_alpha = DEPRECATED_VALUE
+        self.prioritized_replay_beta = DEPRECATED_VALUE
+        self.prioritized_replay_eps = DEPRECATED_VALUE
+
+    @override(AlgorithmConfig)
+    def training(
+        self,
+        *,
+        target_network_update_freq: Optional[int] = NotProvided,
+        replay_buffer_config: Optional[dict] = NotProvided,
+        store_buffer_in_checkpoints: Optional[bool] = NotProvided,
+        lr_schedule: Optional[List[List[Union[int, float]]]] = NotProvided,
+        epsilon: Optional[LearningRateOrSchedule] = NotProvided,
+        adam_epsilon: Optional[float] = NotProvided,
+        grad_clip: Optional[int] = NotProvided,
+        num_steps_sampled_before_learning_starts: Optional[int] = NotProvided,
+        tau: Optional[float] = NotProvided,
+        num_atoms: Optional[int] = NotProvided,
+        v_min: Optional[float] = NotProvided,
+        v_max: Optional[float] = NotProvided,
+        noisy: Optional[bool] = NotProvided,
+        sigma0: Optional[float] = NotProvided,
+        dueling: Optional[bool] = NotProvided,
+        hiddens: Optional[int] = NotProvided,
+        double_q: Optional[bool] = NotProvided,
+        n_step: Optional[Union[int, Tuple[int, int]]] = NotProvided,
+        before_learn_on_batch: Callable[
+            [Type[MultiAgentBatch], List[Type[Policy]], Type[int]],
+            Type[MultiAgentBatch],
+        ] = NotProvided,
+        training_intensity: Optional[float] = NotProvided,
+        td_error_loss_fn: Optional[str] = NotProvided,
+        categorical_distribution_temperature: Optional[float] = NotProvided,
+        burn_in_len: Optional[int] = NotProvided,
+        **kwargs,
+    ) -> Self:
+        """Sets the training related configuration.
+
+        Args:
+            target_network_update_freq: Update the target network every
+                `target_network_update_freq` sample steps.
+            replay_buffer_config: Replay buffer config.
+                Examples:
+                {
+                "_enable_replay_buffer_api": True,
+                "type": "MultiAgentReplayBuffer",
+                "capacity": 50000,
+                "replay_sequence_length": 1,
+                }
+                - OR -
+                {
+                "_enable_replay_buffer_api": True,
+                "type": "MultiAgentPrioritizedReplayBuffer",
+                "capacity": 50000,
+                "prioritized_replay_alpha": 0.6,
+                "prioritized_replay_beta": 0.4,
+                "prioritized_replay_eps": 1e-6,
+                "replay_sequence_length": 1,
+                }
+                - Where -
+                prioritized_replay_alpha: Alpha parameter controls the degree of
+                prioritization in the buffer. In other words, when a buffer sample has
+                a higher temporal-difference error, with how much more probability
+                should it drawn to use to update the parametrized Q-network. 0.0
+                corresponds to uniform probability. Setting much above 1.0 may quickly
+                result as the sampling distribution could become heavily “pointy” with
+                low entropy.
+                prioritized_replay_beta: Beta parameter controls the degree of
+                importance sampling which suppresses the influence of gradient updates
+                from samples that have higher probability of being sampled via alpha
+                parameter and the temporal-difference error.
+                prioritized_replay_eps: Epsilon parameter sets the baseline probability
+                for sampling so that when the temporal-difference error of a sample is
+                zero, there is still a chance of drawing the sample.
+            store_buffer_in_checkpoints: Set this to True, if you want the contents of
+                your buffer(s) to be stored in any saved checkpoints as well.
+                Warnings will be created if:
+                - This is True AND restoring from a checkpoint that contains no buffer
+                data.
+                - This is False AND restoring from a checkpoint that does contain
+                buffer data.
+            epsilon: Epsilon exploration schedule. In the format of [[timestep, value],
+                [timestep, value], ...]. A schedule must start from
+                timestep 0.
+            adam_epsilon: Adam optimizer's epsilon hyper parameter.
+            grad_clip: If not None, clip gradients during optimization at this value.
+            num_steps_sampled_before_learning_starts: Number of timesteps to collect
+                from rollout workers before we start sampling from replay buffers for
+                learning. Whether we count this in agent steps or environment steps
+                depends on config.multi_agent(count_steps_by=..).
+            tau: Update the target by \tau * policy + (1-\tau) * target_policy.
+            num_atoms: Number of atoms for representing the distribution of return.
+                When this is greater than 1, distributional Q-learning is used.
+            v_min: Minimum value estimation
+            v_max: Maximum value estimation
+            noisy: Whether to use noisy network to aid exploration. This adds parametric
+                noise to the model weights.
+            sigma0: Control the initial parameter noise for noisy nets.
+            dueling: Whether to use dueling DQN.
+            hiddens: Dense-layer setup for each the advantage branch and the value
+                branch in a dueling architecture.
+            double_q: Whether to use double DQN.
+            n_step: N-step target updates. If >1, sars' tuples in trajectories will be
+                postprocessed to become sa[discounted sum of R][s t+n] tuples. An
+                integer will be interpreted as a fixed n-step value. If a tuple of 2
+                ints is provided here, the n-step value will be drawn for each sample(!)
+                in the train batch from a uniform distribution over the closed interval
+                defined by `[n_step[0], n_step[1]]`.
+            before_learn_on_batch: Callback to run before learning on a multi-agent
+                batch of experiences.
+            training_intensity: The intensity with which to update the model (vs
+                collecting samples from the env).
+                If None, uses "natural" values of:
+                `train_batch_size` / (`rollout_fragment_length` x `num_env_runners` x
+                `num_envs_per_env_runner`).
+                If not None, will make sure that the ratio between timesteps inserted
+                into and sampled from the buffer matches the given values.
+                Example:
+                training_intensity=1000.0
+                train_batch_size=250
+                rollout_fragment_length=1
+                num_env_runners=1 (or 0)
+                num_envs_per_env_runner=1
+                -> natural value = 250 / 1 = 250.0
+                -> will make sure that replay+train op will be executed 4x asoften as
+                rollout+insert op (4 * 250 = 1000).
+                See: rllib/algorithms/dqn/dqn.py::calculate_rr_weights for further
+                details.
+            td_error_loss_fn: "huber" or "mse". loss function for calculating TD error
+                when num_atoms is 1. Note that if num_atoms is > 1, this parameter
+                is simply ignored, and softmax cross entropy loss will be used.
+            categorical_distribution_temperature: Set the temperature parameter used
+                by Categorical action distribution. A valid temperature is in the range
+                of [0, 1]. Note that this mostly affects evaluation since TD error uses
+                argmax for return calculation.
+            burn_in_len: The burn-in period for a stateful RLModule. It allows the
+                Learner to utilize the initial `burn_in_len` steps in a replay sequence
+                solely for unrolling the network and establishing a typical starting
+                state. The network is then updated on the remaining steps of the
+                sequence. This process helps mitigate issues stemming from a poor
+                initial state - zero or an outdated recorded state. Consider setting
+                this parameter to a positive integer if your stateful RLModule faces
+                convergence challenges or exhibits signs of catastrophic forgetting.
+
+        Returns:
+            This updated AlgorithmConfig object.
+        """
+        # Pass kwargs onto super's `training()` method.
+        super().training(**kwargs)
+
+        if target_network_update_freq is not NotProvided:
+            self.target_network_update_freq = target_network_update_freq
+        if replay_buffer_config is not NotProvided:
+            # Override entire `replay_buffer_config` if `type` key changes.
+            # Update, if `type` key remains the same or is not specified.
+            new_replay_buffer_config = deep_update(
+                {"replay_buffer_config": self.replay_buffer_config},
+                {"replay_buffer_config": replay_buffer_config},
+                False,
+                ["replay_buffer_config"],
+                ["replay_buffer_config"],
+            )
+            self.replay_buffer_config = new_replay_buffer_config["replay_buffer_config"]
+        if store_buffer_in_checkpoints is not NotProvided:
+            self.store_buffer_in_checkpoints = store_buffer_in_checkpoints
+        if lr_schedule is not NotProvided:
+            self.lr_schedule = lr_schedule
+        if epsilon is not NotProvided:
+            self.epsilon = epsilon
+        if adam_epsilon is not NotProvided:
+            self.adam_epsilon = adam_epsilon
+        if grad_clip is not NotProvided:
+            self.grad_clip = grad_clip
+        if num_steps_sampled_before_learning_starts is not NotProvided:
+            self.num_steps_sampled_before_learning_starts = (
+                num_steps_sampled_before_learning_starts
+            )
+        if tau is not NotProvided:
+            self.tau = tau
+        if num_atoms is not NotProvided:
+            self.num_atoms = num_atoms
+        if v_min is not NotProvided:
+            self.v_min = v_min
+        if v_max is not NotProvided:
+            self.v_max = v_max
+        if noisy is not NotProvided:
+            self.noisy = noisy
+        if sigma0 is not NotProvided:
+            self.sigma0 = sigma0
+        if dueling is not NotProvided:
+            self.dueling = dueling
+        if hiddens is not NotProvided:
+            self.hiddens = hiddens
+        if double_q is not NotProvided:
+            self.double_q = double_q
+        if n_step is not NotProvided:
+            self.n_step = n_step
+        if before_learn_on_batch is not NotProvided:
+            self.before_learn_on_batch = before_learn_on_batch
+        if training_intensity is not NotProvided:
+            self.training_intensity = training_intensity
+        if td_error_loss_fn is not NotProvided:
+            self.td_error_loss_fn = td_error_loss_fn
+        if categorical_distribution_temperature is not NotProvided:
+            self.categorical_distribution_temperature = (
+                categorical_distribution_temperature
+            )
+        if burn_in_len is not NotProvided:
+            self.burn_in_len = burn_in_len
+
+        return self
+
+    @override(AlgorithmConfig)
+    def validate(self) -> None:
+        # Call super's validation method.
+        super().validate()
+
+        if self.enable_rl_module_and_learner:
+            # `lr_schedule` checking.
+            if self.lr_schedule is not None:
+                self._value_error(
+                    "`lr_schedule` is deprecated and must be None! Use the "
+                    "`lr` setting to setup a schedule."
+                )
+        else:
+            if not self.in_evaluation:
+                validate_buffer_config(self)
+
+            # TODO (simon): Find a clean solution to deal with configuration configs
+            #  when using the new API stack.
+            if self.exploration_config["type"] == "ParameterNoise":
+                if self.batch_mode != "complete_episodes":
+                    self._value_error(
+                        "ParameterNoise Exploration requires `batch_mode` to be "
+                        "'complete_episodes'. Try setting `config.env_runners("
+                        "batch_mode='complete_episodes')`."
+                    )
+                if self.noisy:
+                    self._value_error(
+                        "ParameterNoise Exploration and `noisy` network cannot be"
+                        " used at the same time!"
+                    )
+
+        if self.td_error_loss_fn not in ["huber", "mse"]:
+            self._value_error("`td_error_loss_fn` must be 'huber' or 'mse'!")
+
+        # Check rollout_fragment_length to be compatible with n_step.
+        if (
+            not self.in_evaluation
+            and self.rollout_fragment_length != "auto"
+            and self.rollout_fragment_length < self.n_step
+        ):
+            self._value_error(
+                f"Your `rollout_fragment_length` ({self.rollout_fragment_length}) is "
+                f"smaller than `n_step` ({self.n_step})! "
+                "Try setting config.env_runners(rollout_fragment_length="
+                f"{self.n_step})."
+            )
+
+        # Check, if the `max_seq_len` is longer then the burn-in.
+        if (
+            "max_seq_len" in self.model_config
+            and 0 < self.model_config["max_seq_len"] <= self.burn_in_len
+        ):
+            raise ValueError(
+                f"Your defined `burn_in_len`={self.burn_in_len} is larger or equal "
+                f"`max_seq_len`={self.model_config['max_seq_len']}! Either decrease "
+                "the `burn_in_len` or increase your `max_seq_len`."
+            )
+
+        # Validate that we use the corresponding `EpisodeReplayBuffer` when using
+        # episodes.
+        # TODO (sven, simon): Implement the multi-agent case for replay buffers.
+        from ray.rllib.utils.replay_buffers.episode_replay_buffer import (
+            EpisodeReplayBuffer,
+        )
+
+        if (
+            self.enable_env_runner_and_connector_v2
+            and not isinstance(self.replay_buffer_config["type"], str)
+            and not issubclass(self.replay_buffer_config["type"], EpisodeReplayBuffer)
+        ):
+            self._value_error(
+                "When using the new `EnvRunner API` the replay buffer must be of type "
+                "`EpisodeReplayBuffer`."
+            )
+        elif not self.enable_env_runner_and_connector_v2 and (
+            (
+                isinstance(self.replay_buffer_config["type"], str)
+                and "Episode" in self.replay_buffer_config["type"]
+            )
+            or issubclass(self.replay_buffer_config["type"], EpisodeReplayBuffer)
+        ):
+            self._value_error(
+                "When using the old API stack the replay buffer must not be of type "
+                "`EpisodeReplayBuffer`! We suggest you use the following config to run "
+                "DQN on the old API stack: `config.training(replay_buffer_config={"
+                "'type': 'MultiAgentPrioritizedReplayBuffer', "
+                "'prioritized_replay_alpha': [alpha], "
+                "'prioritized_replay_beta': [beta], "
+                "'prioritized_replay_eps': [eps], "
+                "})`."
+            )
+
+    @override(AlgorithmConfig)
+    def get_rollout_fragment_length(self, worker_index: int = 0) -> int:
+        if self.rollout_fragment_length == "auto":
+            return (
+                self.n_step[1]
+                if isinstance(self.n_step, (tuple, list))
+                else self.n_step
+            )
+        else:
+            return self.rollout_fragment_length
+
+    @override(AlgorithmConfig)
+    def get_default_rl_module_spec(self) -> RLModuleSpecType:
+        if self.framework_str == "torch":
+            from ray.rllib.algorithms.dqn.torch.default_dqn_torch_rl_module import (
+                DefaultDQNTorchRLModule,
+            )
+
+            return RLModuleSpec(
+                module_class=DefaultDQNTorchRLModule,
+                model_config=self.model_config,
+            )
+        else:
+            raise ValueError(
+                f"The framework {self.framework_str} is not supported! "
+                "Use `config.framework('torch')` instead."
+            )
+
+    @property
+    @override(AlgorithmConfig)
+    def _model_config_auto_includes(self) -> Dict[str, Any]:
+        return super()._model_config_auto_includes | {
+            "double_q": self.double_q,
+            "dueling": self.dueling,
+            "epsilon": self.epsilon,
+            "num_atoms": self.num_atoms,
+            "std_init": self.sigma0,
+            "v_max": self.v_max,
+            "v_min": self.v_min,
+        }
+
+    @override(AlgorithmConfig)
+    def get_default_learner_class(self) -> Union[Type["Learner"], str]:
+        if self.framework_str == "torch":
+            from ray.rllib.algorithms.dqn.torch.dqn_torch_learner import (
+                DQNTorchLearner,
+            )
+
+            return DQNTorchLearner
+        else:
+            raise ValueError(
+                f"The framework {self.framework_str} is not supported! "
+                "Use `config.framework('torch')` instead."
+            )
+
+
+def calculate_rr_weights(config: AlgorithmConfig) -> List[float]:
+    """Calculate the round robin weights for the rollout and train steps"""
+    if not config.training_intensity:
+        return [1, 1]
+
+    # Calculate the "native ratio" as:
+    # [train-batch-size] / [size of env-rolled-out sampled data]
+    # This is to set freshly rollout-collected data in relation to
+    # the data we pull from the replay buffer (which also contains old
+    # samples).
+    native_ratio = config.total_train_batch_size / (
+        config.get_rollout_fragment_length()
+        * config.num_envs_per_env_runner
+        # Add one to workers because the local
+        # worker usually collects experiences as well, and we avoid division by zero.
+        * max(config.num_env_runners + 1, 1)
+    )
+
+    # Training intensity is specified in terms of
+    # (steps_replayed / steps_sampled), so adjust for the native ratio.
+    sample_and_train_weight = config.training_intensity / native_ratio
+    if sample_and_train_weight < 1:
+        return [int(np.round(1 / sample_and_train_weight)), 1]
+    else:
+        return [1, int(np.round(sample_and_train_weight))]
+
+
+class DQN(Algorithm):
+    @classmethod
+    @override(Algorithm)
+    def get_default_config(cls) -> DQNConfig:
+        return DQNConfig()
+
+    @classmethod
+    @override(Algorithm)
+    def get_default_policy_class(
+        cls, config: AlgorithmConfig
+    ) -> Optional[Type[Policy]]:
+        if config["framework"] == "torch":
+            return DQNTorchPolicy
+        else:
+            return DQNTFPolicy
+
+    @override(Algorithm)
+    def setup(self, config: AlgorithmConfig) -> None:
+        super().setup(config)
+
+        if self.config.enable_env_runner_and_connector_v2 and self.env_runner_group:
+            if self.env_runner is None:
+                self._module_is_stateful = self.env_runner_group.foreach_env_runner(
+                    lambda er: er.module.is_stateful(),
+                    remote_worker_ids=[1],
+                    local_env_runner=False,
+                )[0]
+            else:
+                self._module_is_stateful = self.env_runner.module.is_stateful()
+
+    @override(Algorithm)
+    def training_step(self) -> None:
+        """DQN training iteration function.
+
+        Each training iteration, we:
+        - Sample (MultiAgentBatch) from workers.
+        - Store new samples in replay buffer.
+        - Sample training batch (MultiAgentBatch) from replay buffer.
+        - Learn on training batch.
+        - Update remote workers' new policy weights.
+        - Update target network every `target_network_update_freq` sample steps.
+        - Return all collected metrics for the iteration.
+
+        Returns:
+            The results dict from executing the training iteration.
+        """
+        # Old API stack (Policy, RolloutWorker, Connector).
+        if not self.config.enable_env_runner_and_connector_v2:
+            return self._training_step_old_api_stack()
+
+        # New API stack (RLModule, Learner, EnvRunner, ConnectorV2).
+        return self._training_step_new_api_stack()
+
+    def _training_step_new_api_stack(self):
+        # Alternate between storing and sampling and training.
+        store_weight, sample_and_train_weight = calculate_rr_weights(self.config)
+
+        # Run multiple sampling + storing to buffer iterations.
+        for _ in range(store_weight):
+            with self.metrics.log_time((TIMERS, ENV_RUNNER_SAMPLING_TIMER)):
+                # Sample in parallel from workers.
+                episodes, env_runner_results = synchronous_parallel_sample(
+                    worker_set=self.env_runner_group,
+                    concat=True,
+                    sample_timeout_s=self.config.sample_timeout_s,
+                    _uses_new_env_runners=True,
+                    _return_metrics=True,
+                )
+            # Reduce EnvRunner metrics over the n EnvRunners.
+            self.metrics.aggregate(env_runner_results, key=ENV_RUNNER_RESULTS)
+
+            # Add the sampled experiences to the replay buffer.
+            with self.metrics.log_time((TIMERS, REPLAY_BUFFER_ADD_DATA_TIMER)):
+                self.local_replay_buffer.add(episodes)
+
+        if self.config.count_steps_by == "agent_steps":
+            current_ts = sum(
+                self.metrics.peek(
+                    (ENV_RUNNER_RESULTS, NUM_AGENT_STEPS_SAMPLED_LIFETIME), default={}
+                ).values()
+            )
+        else:
+            current_ts = self.metrics.peek(
+                (ENV_RUNNER_RESULTS, NUM_ENV_STEPS_SAMPLED_LIFETIME), default=0
+            )
+
+        # If enough experiences have been sampled start training.
+        if current_ts >= self.config.num_steps_sampled_before_learning_starts:
+            # Run multiple sample-from-buffer and update iterations.
+            for _ in range(sample_and_train_weight):
+                # Sample a list of episodes used for learning from the replay buffer.
+                with self.metrics.log_time((TIMERS, REPLAY_BUFFER_SAMPLE_TIMER)):
+
+                    episodes = self.local_replay_buffer.sample(
+                        num_items=self.config.total_train_batch_size,
+                        n_step=self.config.n_step,
+                        # In case an `EpisodeReplayBuffer` is used we need to provide
+                        # the sequence length.
+                        batch_length_T=(
+                            self._module_is_stateful
+                            * self.config.model_config.get("max_seq_len", 0)
+                        ),
+                        lookback=int(self._module_is_stateful),
+                        # TODO (simon): Implement `burn_in_len` in SAC and remove this
+                        # if-else clause.
+                        min_batch_length_T=self.config.burn_in_len
+                        if hasattr(self.config, "burn_in_len")
+                        else 0,
+                        gamma=self.config.gamma,
+                        beta=self.config.replay_buffer_config.get("beta"),
+                        sample_episodes=True,
+                    )
+
+                    # Get the replay buffer metrics.
+                    replay_buffer_results = self.local_replay_buffer.get_metrics()
+                    self.metrics.aggregate(
+                        [replay_buffer_results], key=REPLAY_BUFFER_RESULTS
+                    )
+
+                # Perform an update on the buffer-sampled train batch.
+                with self.metrics.log_time((TIMERS, LEARNER_UPDATE_TIMER)):
+                    learner_results = self.learner_group.update(
+                        episodes=episodes,
+                        timesteps={
+                            NUM_ENV_STEPS_SAMPLED_LIFETIME: (
+                                self.metrics.peek(
+                                    (ENV_RUNNER_RESULTS, NUM_ENV_STEPS_SAMPLED_LIFETIME)
+                                )
+                            ),
+                            NUM_AGENT_STEPS_SAMPLED_LIFETIME: (
+                                self.metrics.peek(
+                                    (
+                                        ENV_RUNNER_RESULTS,
+                                        NUM_AGENT_STEPS_SAMPLED_LIFETIME,
+                                    )
+                                )
+                            ),
+                        },
+                    )
+                    # Isolate TD-errors from result dicts (we should not log these to
+                    # disk or WandB, they might be very large).
+                    td_errors = defaultdict(list)
+                    for res in learner_results:
+                        for module_id, module_results in res.items():
+                            if TD_ERROR_KEY in module_results:
+                                td_errors[module_id].extend(
+                                    convert_to_numpy(
+                                        module_results.pop(TD_ERROR_KEY).peek()
+                                    )
+                                )
+                    td_errors = {
+                        module_id: {TD_ERROR_KEY: np.concatenate(s, axis=0)}
+                        for module_id, s in td_errors.items()
+                    }
+                    self.metrics.aggregate(learner_results, key=LEARNER_RESULTS)
+
+                # Update replay buffer priorities.
+                with self.metrics.log_time((TIMERS, REPLAY_BUFFER_UPDATE_PRIOS_TIMER)):
+                    update_priorities_in_episode_replay_buffer(
+                        replay_buffer=self.local_replay_buffer,
+                        td_errors=td_errors,
+                    )
+
+            # Update weights and global_vars - after learning on the local worker -
+            # on all remote workers.
+            with self.metrics.log_time((TIMERS, SYNCH_WORKER_WEIGHTS_TIMER)):
+                modules_to_update = set(learner_results[0].keys()) - {ALL_MODULES}
+                # NOTE: the new API stack does not use global vars.
+                self.env_runner_group.sync_weights(
+                    from_worker_or_learner_group=self.learner_group,
+                    policies=modules_to_update,
+                    global_vars=None,
+                    inference_only=True,
+                )
+
+    def _training_step_old_api_stack(self) -> ResultDict:
+        """Training step for the old API stack.
+
+        More specifically this training step relies on `RolloutWorker`.
+        """
+        train_results = {}
+
+        # We alternate between storing new samples and sampling and training
+        store_weight, sample_and_train_weight = calculate_rr_weights(self.config)
+
+        for _ in range(store_weight):
+            # Sample (MultiAgentBatch) from workers.
+            with self._timers[SAMPLE_TIMER]:
+                new_sample_batch: SampleBatchType = synchronous_parallel_sample(
+                    worker_set=self.env_runner_group,
+                    concat=True,
+                    sample_timeout_s=self.config.sample_timeout_s,
+                )
+
+            # Return early if all our workers failed.
+            if not new_sample_batch:
+                return {}
+
+            # Update counters
+            self._counters[NUM_AGENT_STEPS_SAMPLED] += new_sample_batch.agent_steps()
+            self._counters[NUM_ENV_STEPS_SAMPLED] += new_sample_batch.env_steps()
+
+            # Store new samples in replay buffer.
+            self.local_replay_buffer.add(new_sample_batch)
+
+        global_vars = {
+            "timestep": self._counters[NUM_ENV_STEPS_SAMPLED],
+        }
+
+        # Update target network every `target_network_update_freq` sample steps.
+        cur_ts = self._counters[
+            (
+                NUM_AGENT_STEPS_SAMPLED
+                if self.config.count_steps_by == "agent_steps"
+                else NUM_ENV_STEPS_SAMPLED
+            )
+        ]
+
+        if cur_ts > self.config.num_steps_sampled_before_learning_starts:
+            for _ in range(sample_and_train_weight):
+                # Sample training batch (MultiAgentBatch) from replay buffer.
+                train_batch = sample_min_n_steps_from_buffer(
+                    self.local_replay_buffer,
+                    self.config.total_train_batch_size,
+                    count_by_agent_steps=self.config.count_steps_by == "agent_steps",
+                )
+
+                # Postprocess batch before we learn on it
+                post_fn = self.config.get("before_learn_on_batch") or (lambda b, *a: b)
+                train_batch = post_fn(train_batch, self.env_runner_group, self.config)
+
+                # Learn on training batch.
+                # Use simple optimizer (only for multi-agent or tf-eager; all other
+                # cases should use the multi-GPU optimizer, even if only using 1 GPU)
+                if self.config.get("simple_optimizer") is True:
+                    train_results = train_one_step(self, train_batch)
+                else:
+                    train_results = multi_gpu_train_one_step(self, train_batch)
+
+                # Update replay buffer priorities.
+                update_priorities_in_replay_buffer(
+                    self.local_replay_buffer,
+                    self.config,
+                    train_batch,
+                    train_results,
+                )
+
+                last_update = self._counters[LAST_TARGET_UPDATE_TS]
+                if cur_ts - last_update >= self.config.target_network_update_freq:
+                    to_update = self.env_runner.get_policies_to_train()
+                    self.env_runner.foreach_policy_to_train(
+                        lambda p, pid, to_update=to_update: (
+                            pid in to_update and p.update_target()
+                        )
+                    )
+                    self._counters[NUM_TARGET_UPDATES] += 1
+                    self._counters[LAST_TARGET_UPDATE_TS] = cur_ts
+
+                # Update weights and global_vars - after learning on the local worker -
+                # on all remote workers.
+                with self._timers[SYNCH_WORKER_WEIGHTS_TIMER]:
+                    self.env_runner_group.sync_weights(global_vars=global_vars)
+
+        # Return all collected metrics for the iteration.
+        return train_results

rllib/algorithms/dqn/dqn_catalog.py

@@ -0,0 +1,179 @@
+import gymnasium as gym
+
+from ray.rllib.core.distribution.torch.torch_distribution import TorchCategorical
+from ray.rllib.core.models.base import Model
+from ray.rllib.core.models.catalog import Catalog
+from ray.rllib.core.models.configs import MLPHeadConfig
+from ray.rllib.utils.annotations import (
+    ExperimentalAPI,
+    OverrideToImplementCustomLogic,
+    override,
+)
+
+
+@ExperimentalAPI
+class DQNCatalog(Catalog):
+    """The catalog class used to build models for DQN Rainbow.
+
+    `DQNCatalog` provides the following models:
+        - Encoder: The encoder used to encode the observations.
+        - Target_Encoder: The encoder used to encode the observations
+            for the target network.
+        - Af Head: Either the head of the advantage stream, if a dueling
+            architecture is used or the head of the Q-function. This is
+            a multi-node head with `action_space.n` many nodes in case
+            of expectation learning and `action_space.n` times the number
+            of atoms (`num_atoms`) in case of distributional Q-learning.
+        - Vf Head (optional): The head of the value function in case a
+            dueling architecture is chosen. This is a single node head.
+            If no dueling architecture is used, this head does not exist.
+
+    Any custom head can be built by overridng the `build_af_head()` and
+    `build_vf_head()`. Alternatively, the `AfHeadConfig` or `VfHeadConfig`
+    can be overridden to build custom logic during `RLModule` runtime.
+
+    All heads can optionally use distributional learning. In this case the
+    number of output neurons corresponds to the number of actions times the
+    number of support atoms of the discrete distribution.
+
+    Any module built for exploration or inference is built with the flag
+    `ìnference_only=True` and does not contain any target networks. This flag can
+    be set in a `SingleAgentModuleSpec` through the `inference_only` boolean flag.
+    """
+
+    @override(Catalog)
+    def __init__(
+        self,
+        observation_space: gym.Space,
+        action_space: gym.Space,
+        model_config_dict: dict,
+        view_requirements: dict = None,
+    ):
+        """Initializes the DQNCatalog.
+
+        Args:
+            observation_space: The observation space of the Encoder.
+            action_space: The action space for the Af Head.
+            model_config_dict: The model config to use.
+        """
+        assert view_requirements is None, (
+            "Instead, use the new ConnectorV2 API to pick whatever information "
+            "you need from the running episodes"
+        )
+
+        super().__init__(
+            observation_space=observation_space,
+            action_space=action_space,
+            model_config_dict=model_config_dict,
+        )
+
+        # The number of atoms to be used for distributional Q-learning.
+        self.num_atoms: bool = self._model_config_dict["num_atoms"]
+
+        # Advantage and value streams have MLP heads. Note, the advantage
+        # stream will has an output dimension that is the product of the
+        # action space dimension and the number of atoms to approximate the
+        # return distribution in distributional reinforcement learning.
+        self.af_head_config = self._get_head_config(
+            output_layer_dim=int(self.action_space.n * self.num_atoms)
+        )
+        self.vf_head_config = self._get_head_config(output_layer_dim=1)
+
+    @OverrideToImplementCustomLogic
+    def build_af_head(self, framework: str) -> Model:
+        """Build the A/Q-function head.
+
+        Note, if no dueling architecture is chosen, this will
+        be the Q-function head.
+
+        The default behavior is to build the head from the `af_head_config`.
+        This can be overridden to build a custom policy head as a means to
+        configure the behavior of a `DQNRLModule` implementation.
+
+        Args:
+            framework: The framework to use. Either "torch" or "tf2".
+
+        Returns:
+            The advantage head in case a dueling architecutre is chosen or
+            the Q-function head in the other case.
+        """
+        return self.af_head_config.build(framework=framework)
+
+    @OverrideToImplementCustomLogic
+    def build_vf_head(self, framework: str) -> Model:
+        """Build the value function head.
+
+        Note, this function is only called in case of a dueling architecture.
+
+        The default behavior is to build the head from the `vf_head_config`.
+        This can be overridden to build a custom policy head as a means to
+        configure the behavior of a `DQNRLModule` implementation.
+
+        Args:
+            framework: The framework to use. Either "torch" or "tf2".
+
+        Returns:
+            The value function head.
+        """
+
+        return self.vf_head_config.build(framework=framework)
+
+    @override(Catalog)
+    def get_action_dist_cls(self, framework: str) -> "TorchCategorical":
+        # We only implement DQN Rainbow for Torch.
+        if framework != "torch":
+            raise ValueError("DQN Rainbow is only supported for framework `torch`.")
+        else:
+            return TorchCategorical
+
+    def _get_head_config(self, output_layer_dim: int):
+        """Returns a head config.
+
+        Args:
+            output_layer_dim: Integer defining the output layer dimension.
+                This is 1 for the Vf-head and `action_space.n * num_atoms`
+                for the Af(Qf)-head.
+
+        Returns:
+            A `MLPHeadConfig`.
+        """
+        # Return the appropriate config.
+        return MLPHeadConfig(
+            input_dims=self.latent_dims,
+            hidden_layer_dims=self._model_config_dict["head_fcnet_hiddens"],
+            # Note, `"post_fcnet_activation"` is `"relu"` by definition.
+            hidden_layer_activation=self._model_config_dict["head_fcnet_activation"],
+            # TODO (simon): Not yet available.
+            # hidden_layer_use_layernorm=self._model_config_dict[
+            #     "hidden_layer_use_layernorm"
+            # ],
+            # hidden_layer_use_bias=self._model_config_dict["hidden_layer_use_bias"],
+            hidden_layer_weights_initializer=self._model_config_dict[
+                "head_fcnet_kernel_initializer"
+            ],
+            hidden_layer_weights_initializer_config=self._model_config_dict[
+                "head_fcnet_kernel_initializer_kwargs"
+            ],
+            hidden_layer_bias_initializer=self._model_config_dict[
+                "head_fcnet_bias_initializer"
+            ],
+            hidden_layer_bias_initializer_config=self._model_config_dict[
+                "head_fcnet_bias_initializer_kwargs"
+            ],
+            output_layer_activation="linear",
+            output_layer_dim=output_layer_dim,
+            # TODO (simon): Not yet available.
+            # output_layer_use_bias=self._model_config_dict["output_layer_use_bias"],
+            output_layer_weights_initializer=self._model_config_dict[
+                "head_fcnet_kernel_initializer"
+            ],
+            output_layer_weights_initializer_config=self._model_config_dict[
+                "head_fcnet_kernel_initializer_kwargs"
+            ],
+            output_layer_bias_initializer=self._model_config_dict[
+                "head_fcnet_bias_initializer"
+            ],
+            output_layer_bias_initializer_config=self._model_config_dict[
+                "head_fcnet_bias_initializer_kwargs"
+            ],
+        )

rllib/algorithms/dqn/dqn_learner.py

@@ -0,0 +1,126 @@
+from collections import defaultdict
+from typing import Any, Dict, Optional
+
+from ray.rllib.connectors.common.add_observations_from_episodes_to_batch import (
+    AddObservationsFromEpisodesToBatch,
+)
+from ray.rllib.connectors.learner.add_next_observations_from_episodes_to_train_batch import (  # noqa
+    AddNextObservationsFromEpisodesToTrainBatch,
+)
+from ray.rllib.core.learner.learner import Learner
+from ray.rllib.core.learner.utils import update_target_network
+from ray.rllib.core.rl_module.apis import QNetAPI, TargetNetworkAPI
+from ray.rllib.core.rl_module.multi_rl_module import MultiRLModuleSpec
+from ray.rllib.core.rl_module.rl_module import RLModuleSpec
+from ray.rllib.utils.annotations import (
+    OverrideToImplementCustomLogic_CallToSuperRecommended,
+    override,
+)
+from ray.rllib.utils.metrics import (
+    LAST_TARGET_UPDATE_TS,
+    NUM_ENV_STEPS_SAMPLED_LIFETIME,
+    NUM_TARGET_UPDATES,
+)
+from ray.rllib.utils.typing import ModuleID, ShouldModuleBeUpdatedFn
+
+# Now, this is double defined: In `SACRLModule` and here. I would keep it here
+# or push it into the `Learner` as these are recurring keys in RL.
+ATOMS = "atoms"
+QF_LOSS_KEY = "qf_loss"
+QF_LOGITS = "qf_logits"
+QF_MEAN_KEY = "qf_mean"
+QF_MAX_KEY = "qf_max"
+QF_MIN_KEY = "qf_min"
+QF_NEXT_PREDS = "qf_next_preds"
+QF_TARGET_NEXT_PREDS = "qf_target_next_preds"
+QF_TARGET_NEXT_PROBS = "qf_target_next_probs"
+QF_PREDS = "qf_preds"
+QF_PROBS = "qf_probs"
+TD_ERROR_MEAN_KEY = "td_error_mean"
+
+
+class DQNLearner(Learner):
+    @OverrideToImplementCustomLogic_CallToSuperRecommended
+    @override(Learner)
+    def build(self) -> None:
+        super().build()
+
+        self.last_update_ts_by_mid = defaultdict(int)  # Returns 0 for missing keys
+
+        # Make target networks.
+        self.module.foreach_module(
+            lambda mid, mod: (
+                mod.make_target_networks()
+                if isinstance(mod, TargetNetworkAPI)
+                else None
+            )
+        )
+
+        # Prepend the "add-NEXT_OBS-from-episodes-to-train-batch" connector piece (right
+        # after the corresponding "add-OBS-..." default piece).
+        self._learner_connector.insert_after(
+            AddObservationsFromEpisodesToBatch,
+            AddNextObservationsFromEpisodesToTrainBatch(),
+        )
+
+    @override(Learner)
+    def add_module(
+        self,
+        *,
+        module_id: ModuleID,
+        module_spec: RLModuleSpec,
+        config_overrides: Optional[Dict] = None,
+        new_should_module_be_updated: Optional[ShouldModuleBeUpdatedFn] = None,
+    ) -> MultiRLModuleSpec:
+        marl_spec = super().add_module(
+            module_id=module_id,
+            module_spec=module_spec,
+            config_overrides=config_overrides,
+            new_should_module_be_updated=new_should_module_be_updated,
+        )
+        # Create target networks for added Module, if applicable.
+        if isinstance(self.module[module_id].unwrapped(), TargetNetworkAPI):
+            self.module[module_id].unwrapped().make_target_networks()
+        return marl_spec
+
+    @override(Learner)
+    def after_gradient_based_update(self, *, timesteps: Dict[str, Any]) -> None:
+        """Updates the target Q Networks."""
+        super().after_gradient_based_update(timesteps=timesteps)
+
+        timestep = timesteps.get(NUM_ENV_STEPS_SAMPLED_LIFETIME, 0)
+
+        # TODO (sven): Maybe we should have a `after_gradient_based_update`
+        #  method per module?
+        for module_id, module in self.module._rl_modules.items():
+            config = self.config.get_config_for_module(module_id)
+            if timestep - self.last_update_ts_by_mid[
+                module_id
+            ] >= config.target_network_update_freq and isinstance(
+                module.unwrapped(), TargetNetworkAPI
+            ):
+                for (
+                    main_net,
+                    target_net,
+                ) in module.unwrapped().get_target_network_pairs():
+                    update_target_network(
+                        main_net=main_net,
+                        target_net=target_net,
+                        tau=config.tau,
+                    )
+                # Increase lifetime target network update counter by one.
+                self.metrics.log_value(
+                    (module_id, NUM_TARGET_UPDATES), 1, reduce="lifetime_sum"
+                )
+                # Update the (single-value -> window=1) last updated timestep metric.
+                self.last_update_ts_by_mid[module_id] = timestep
+                self.metrics.log_value(
+                    (module_id, LAST_TARGET_UPDATE_TS), timestep, reduce="max"
+                )
+
+    @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 [QNetAPI, TargetNetworkAPI]

rllib/algorithms/dqn/dqn_tf_policy.py

@@ -0,0 +1,511 @@
+from typing import Dict
+
+import gymnasium as gym
+import numpy as np
+
+import ray
+from ray.rllib.algorithms.dqn.distributional_q_tf_model import DistributionalQTFModel
+from ray.rllib.evaluation.postprocessing import adjust_nstep
+from ray.rllib.models import ModelCatalog
+from ray.rllib.models.modelv2 import ModelV2
+from ray.rllib.models.tf.tf_action_dist import get_categorical_class_with_temperature
+from ray.rllib.policy.policy import Policy
+from ray.rllib.policy.sample_batch import SampleBatch
+from ray.rllib.policy.tf_mixins import LearningRateSchedule, TargetNetworkMixin
+from ray.rllib.policy.tf_policy_template import build_tf_policy
+from ray.rllib.utils.annotations import OldAPIStack
+from ray.rllib.utils.error import UnsupportedSpaceException
+from ray.rllib.utils.exploration import ParameterNoise
+from ray.rllib.utils.framework import try_import_tf
+from ray.rllib.utils.numpy import convert_to_numpy
+from ray.rllib.utils.tf_utils import (
+    huber_loss,
+    l2_loss,
+    make_tf_callable,
+    minimize_and_clip,
+    reduce_mean_ignore_inf,
+)
+from ray.rllib.utils.typing import AlgorithmConfigDict, ModelGradients, TensorType
+
+tf1, tf, tfv = try_import_tf()
+
+# Importance sampling weights for prioritized replay
+PRIO_WEIGHTS = "weights"
+Q_SCOPE = "q_func"
+Q_TARGET_SCOPE = "target_q_func"
+
+
+@OldAPIStack
+class QLoss:
+    def __init__(
+        self,
+        q_t_selected: TensorType,
+        q_logits_t_selected: TensorType,
+        q_tp1_best: TensorType,
+        q_dist_tp1_best: TensorType,
+        importance_weights: TensorType,
+        rewards: TensorType,
+        done_mask: TensorType,
+        gamma: float = 0.99,
+        n_step: int = 1,
+        num_atoms: int = 1,
+        v_min: float = -10.0,
+        v_max: float = 10.0,
+        loss_fn=huber_loss,
+    ):
+
+        if num_atoms > 1:
+            # Distributional Q-learning which corresponds to an entropy loss
+
+            z = tf.range(num_atoms, dtype=tf.float32)
+            z = v_min + z * (v_max - v_min) / float(num_atoms - 1)
+
+            # (batch_size, 1) * (1, num_atoms) = (batch_size, num_atoms)
+            r_tau = tf.expand_dims(rewards, -1) + gamma**n_step * tf.expand_dims(
+                1.0 - done_mask, -1
+            ) * tf.expand_dims(z, 0)
+            r_tau = tf.clip_by_value(r_tau, v_min, v_max)
+            b = (r_tau - v_min) / ((v_max - v_min) / float(num_atoms - 1))
+            lb = tf.floor(b)
+            ub = tf.math.ceil(b)
+            # indispensable judgement which is missed in most implementations
+            # when b happens to be an integer, lb == ub, so pr_j(s', a*) will
+            # be discarded because (ub-b) == (b-lb) == 0
+            floor_equal_ceil = tf.cast(tf.less(ub - lb, 0.5), tf.float32)
+
+            l_project = tf.one_hot(
+                tf.cast(lb, dtype=tf.int32), num_atoms
+            )  # (batch_size, num_atoms, num_atoms)
+            u_project = tf.one_hot(
+                tf.cast(ub, dtype=tf.int32), num_atoms
+            )  # (batch_size, num_atoms, num_atoms)
+            ml_delta = q_dist_tp1_best * (ub - b + floor_equal_ceil)
+            mu_delta = q_dist_tp1_best * (b - lb)
+            ml_delta = tf.reduce_sum(l_project * tf.expand_dims(ml_delta, -1), axis=1)
+            mu_delta = tf.reduce_sum(u_project * tf.expand_dims(mu_delta, -1), axis=1)
+            m = ml_delta + mu_delta
+
+            # Rainbow paper claims that using this cross entropy loss for
+            # priority is robust and insensitive to `prioritized_replay_alpha`
+            self.td_error = tf.nn.softmax_cross_entropy_with_logits(
+                labels=m, logits=q_logits_t_selected
+            )
+            self.loss = tf.reduce_mean(
+                self.td_error * tf.cast(importance_weights, tf.float32)
+            )
+            self.stats = {
+                # TODO: better Q stats for dist dqn
+                "mean_td_error": tf.reduce_mean(self.td_error),
+            }
+        else:
+            q_tp1_best_masked = (1.0 - done_mask) * q_tp1_best
+
+            # compute RHS of bellman equation
+            q_t_selected_target = rewards + gamma**n_step * q_tp1_best_masked
+
+            # compute the error (potentially clipped)
+            self.td_error = q_t_selected - tf.stop_gradient(q_t_selected_target)
+            self.loss = tf.reduce_mean(
+                tf.cast(importance_weights, tf.float32) * loss_fn(self.td_error)
+            )
+            self.stats = {
+                "mean_q": tf.reduce_mean(q_t_selected),
+                "min_q": tf.reduce_min(q_t_selected),
+                "max_q": tf.reduce_max(q_t_selected),
+                "mean_td_error": tf.reduce_mean(self.td_error),
+            }
+
+
+@OldAPIStack
+class ComputeTDErrorMixin:
+    """Assign the `compute_td_error` method to the DQNTFPolicy
+
+    This allows us to prioritize on the worker side.
+    """
+
+    def __init__(self):
+        @make_tf_callable(self.get_session(), dynamic_shape=True)
+        def compute_td_error(
+            obs_t, act_t, rew_t, obs_tp1, terminateds_mask, importance_weights
+        ):
+            # Do forward pass on loss to update td error attribute
+            build_q_losses(
+                self,
+                self.model,
+                None,
+                {
+                    SampleBatch.CUR_OBS: tf.convert_to_tensor(obs_t),
+                    SampleBatch.ACTIONS: tf.convert_to_tensor(act_t),
+                    SampleBatch.REWARDS: tf.convert_to_tensor(rew_t),
+                    SampleBatch.NEXT_OBS: tf.convert_to_tensor(obs_tp1),
+                    SampleBatch.TERMINATEDS: tf.convert_to_tensor(terminateds_mask),
+                    PRIO_WEIGHTS: tf.convert_to_tensor(importance_weights),
+                },
+            )
+
+            return self.q_loss.td_error
+
+        self.compute_td_error = compute_td_error
+
+
+@OldAPIStack
+def build_q_model(
+    policy: Policy,
+    obs_space: gym.spaces.Space,
+    action_space: gym.spaces.Space,
+    config: AlgorithmConfigDict,
+) -> ModelV2:
+    """Build q_model and target_model for DQN
+
+    Args:
+        policy: The Policy, which will use the model for optimization.
+        obs_space (gym.spaces.Space): The policy's observation space.
+        action_space (gym.spaces.Space): The policy's action space.
+        config (AlgorithmConfigDict):
+
+    Returns:
+        ModelV2: The Model for the Policy to use.
+            Note: The target q model will not be returned, just assigned to
+            `policy.target_model`.
+    """
+    if not isinstance(action_space, gym.spaces.Discrete):
+        raise UnsupportedSpaceException(
+            "Action space {} is not supported for DQN.".format(action_space)
+        )
+
+    if config["hiddens"]:
+        # try to infer the last layer size, otherwise fall back to 256
+        num_outputs = ([256] + list(config["model"]["fcnet_hiddens"]))[-1]
+        config["model"]["no_final_linear"] = True
+    else:
+        num_outputs = action_space.n
+
+    q_model = ModelCatalog.get_model_v2(
+        obs_space=obs_space,
+        action_space=action_space,
+        num_outputs=num_outputs,
+        model_config=config["model"],
+        framework="tf",
+        model_interface=DistributionalQTFModel,
+        name=Q_SCOPE,
+        num_atoms=config["num_atoms"],
+        dueling=config["dueling"],
+        q_hiddens=config["hiddens"],
+        use_noisy=config["noisy"],
+        v_min=config["v_min"],
+        v_max=config["v_max"],
+        sigma0=config["sigma0"],
+        # TODO(sven): Move option to add LayerNorm after each Dense
+        #  generically into ModelCatalog.
+        add_layer_norm=isinstance(getattr(policy, "exploration", None), ParameterNoise)
+        or config["exploration_config"]["type"] == "ParameterNoise",
+    )
+
+    policy.target_model = ModelCatalog.get_model_v2(
+        obs_space=obs_space,
+        action_space=action_space,
+        num_outputs=num_outputs,
+        model_config=config["model"],
+        framework="tf",
+        model_interface=DistributionalQTFModel,
+        name=Q_TARGET_SCOPE,
+        num_atoms=config["num_atoms"],
+        dueling=config["dueling"],
+        q_hiddens=config["hiddens"],
+        use_noisy=config["noisy"],
+        v_min=config["v_min"],
+        v_max=config["v_max"],
+        sigma0=config["sigma0"],
+        # TODO(sven): Move option to add LayerNorm after each Dense
+        #  generically into ModelCatalog.
+        add_layer_norm=isinstance(getattr(policy, "exploration", None), ParameterNoise)
+        or config["exploration_config"]["type"] == "ParameterNoise",
+    )
+
+    return q_model
+
+
+@OldAPIStack
+def get_distribution_inputs_and_class(
+    policy: Policy, model: ModelV2, input_dict: SampleBatch, *, explore=True, **kwargs
+):
+    q_vals = compute_q_values(
+        policy, model, input_dict, state_batches=None, explore=explore
+    )
+    q_vals = q_vals[0] if isinstance(q_vals, tuple) else q_vals
+
+    policy.q_values = q_vals
+
+    # Return a Torch TorchCategorical distribution where the temperature
+    # parameter is partially binded to the configured value.
+    temperature = policy.config["categorical_distribution_temperature"]
+
+    return (
+        policy.q_values,
+        get_categorical_class_with_temperature(temperature),
+        [],
+    )  # state-out
+
+
+@OldAPIStack
+def build_q_losses(policy: Policy, model, _, train_batch: SampleBatch) -> TensorType:
+    """Constructs the loss for DQNTFPolicy.
+
+    Args:
+        policy: The Policy to calculate the loss for.
+        model (ModelV2): The Model to calculate the loss for.
+        train_batch: The training data.
+
+    Returns:
+        TensorType: A single loss tensor.
+    """
+    config = policy.config
+    # q network evaluation
+    q_t, q_logits_t, q_dist_t, _ = compute_q_values(
+        policy,
+        model,
+        SampleBatch({"obs": train_batch[SampleBatch.CUR_OBS]}),
+        state_batches=None,
+        explore=False,
+    )
+
+    # target q network evalution
+    q_tp1, q_logits_tp1, q_dist_tp1, _ = compute_q_values(
+        policy,
+        policy.target_model,
+        SampleBatch({"obs": train_batch[SampleBatch.NEXT_OBS]}),
+        state_batches=None,
+        explore=False,
+    )
+    if not hasattr(policy, "target_q_func_vars"):
+        policy.target_q_func_vars = policy.target_model.variables()
+
+    # q scores for actions which we know were selected in the given state.
+    one_hot_selection = tf.one_hot(
+        tf.cast(train_batch[SampleBatch.ACTIONS], tf.int32), policy.action_space.n
+    )
+    q_t_selected = tf.reduce_sum(q_t * one_hot_selection, 1)
+    q_logits_t_selected = tf.reduce_sum(
+        q_logits_t * tf.expand_dims(one_hot_selection, -1), 1
+    )
+
+    # compute estimate of best possible value starting from state at t + 1
+    if config["double_q"]:
+        (
+            q_tp1_using_online_net,
+            q_logits_tp1_using_online_net,
+            q_dist_tp1_using_online_net,
+            _,
+        ) = compute_q_values(
+            policy,
+            model,
+            SampleBatch({"obs": train_batch[SampleBatch.NEXT_OBS]}),
+            state_batches=None,
+            explore=False,
+        )
+        q_tp1_best_using_online_net = tf.argmax(q_tp1_using_online_net, 1)
+        q_tp1_best_one_hot_selection = tf.one_hot(
+            q_tp1_best_using_online_net, policy.action_space.n
+        )
+        q_tp1_best = tf.reduce_sum(q_tp1 * q_tp1_best_one_hot_selection, 1)
+        q_dist_tp1_best = tf.reduce_sum(
+            q_dist_tp1 * tf.expand_dims(q_tp1_best_one_hot_selection, -1), 1
+        )
+    else:
+        q_tp1_best_one_hot_selection = tf.one_hot(
+            tf.argmax(q_tp1, 1), policy.action_space.n
+        )
+        q_tp1_best = tf.reduce_sum(q_tp1 * q_tp1_best_one_hot_selection, 1)
+        q_dist_tp1_best = tf.reduce_sum(
+            q_dist_tp1 * tf.expand_dims(q_tp1_best_one_hot_selection, -1), 1
+        )
+
+    loss_fn = huber_loss if policy.config["td_error_loss_fn"] == "huber" else l2_loss
+
+    policy.q_loss = QLoss(
+        q_t_selected,
+        q_logits_t_selected,
+        q_tp1_best,
+        q_dist_tp1_best,
+        train_batch[PRIO_WEIGHTS],
+        tf.cast(train_batch[SampleBatch.REWARDS], tf.float32),
+        tf.cast(train_batch[SampleBatch.TERMINATEDS], tf.float32),
+        config["gamma"],
+        config["n_step"],
+        config["num_atoms"],
+        config["v_min"],
+        config["v_max"],
+        loss_fn,
+    )
+
+    return policy.q_loss.loss
+
+
+@OldAPIStack
+def adam_optimizer(
+    policy: Policy, config: AlgorithmConfigDict
+) -> "tf.keras.optimizers.Optimizer":
+    if policy.config["framework"] == "tf2":
+        return tf.keras.optimizers.Adam(
+            learning_rate=policy.cur_lr, epsilon=config["adam_epsilon"]
+        )
+    else:
+        return tf1.train.AdamOptimizer(
+            learning_rate=policy.cur_lr, epsilon=config["adam_epsilon"]
+        )
+
+
+@OldAPIStack
+def clip_gradients(
+    policy: Policy, optimizer: "tf.keras.optimizers.Optimizer", loss: TensorType
+) -> ModelGradients:
+    if not hasattr(policy, "q_func_vars"):
+        policy.q_func_vars = policy.model.variables()
+
+    return minimize_and_clip(
+        optimizer,
+        loss,
+        var_list=policy.q_func_vars,
+        clip_val=policy.config["grad_clip"],
+    )
+
+
+@OldAPIStack
+def build_q_stats(policy: Policy, batch) -> Dict[str, TensorType]:
+    return dict(
+        {
+            "cur_lr": tf.cast(policy.cur_lr, tf.float64),
+        },
+        **policy.q_loss.stats
+    )
+
+
+@OldAPIStack
+def setup_mid_mixins(policy: Policy, obs_space, action_space, config) -> None:
+    LearningRateSchedule.__init__(policy, config["lr"], config["lr_schedule"])
+    ComputeTDErrorMixin.__init__(policy)
+
+
+@OldAPIStack
+def setup_late_mixins(
+    policy: Policy,
+    obs_space: gym.spaces.Space,
+    action_space: gym.spaces.Space,
+    config: AlgorithmConfigDict,
+) -> None:
+    TargetNetworkMixin.__init__(policy)
+
+
+@OldAPIStack
+def compute_q_values(
+    policy: Policy,
+    model: ModelV2,
+    input_batch: SampleBatch,
+    state_batches=None,
+    seq_lens=None,
+    explore=None,
+    is_training: bool = False,
+):
+
+    config = policy.config
+
+    model_out, state = model(input_batch, state_batches or [], seq_lens)
+
+    if config["num_atoms"] > 1:
+        (
+            action_scores,
+            z,
+            support_logits_per_action,
+            logits,
+            dist,
+        ) = model.get_q_value_distributions(model_out)
+    else:
+        (action_scores, logits, dist) = model.get_q_value_distributions(model_out)
+
+    if config["dueling"]:
+        state_score = model.get_state_value(model_out)
+        if config["num_atoms"] > 1:
+            support_logits_per_action_mean = tf.reduce_mean(
+                support_logits_per_action, 1
+            )
+            support_logits_per_action_centered = (
+                support_logits_per_action
+                - tf.expand_dims(support_logits_per_action_mean, 1)
+            )
+            support_logits_per_action = (
+                tf.expand_dims(state_score, 1) + support_logits_per_action_centered
+            )
+            support_prob_per_action = tf.nn.softmax(logits=support_logits_per_action)
+            value = tf.reduce_sum(input_tensor=z * support_prob_per_action, axis=-1)
+            logits = support_logits_per_action
+            dist = support_prob_per_action
+        else:
+            action_scores_mean = reduce_mean_ignore_inf(action_scores, 1)
+            action_scores_centered = action_scores - tf.expand_dims(
+                action_scores_mean, 1
+            )
+            value = state_score + action_scores_centered
+    else:
+        value = action_scores
+
+    return value, logits, dist, state
+
+
+@OldAPIStack
+def postprocess_nstep_and_prio(
+    policy: Policy, batch: SampleBatch, other_agent=None, episode=None
+) -> SampleBatch:
+    # N-step Q adjustments.
+    if policy.config["n_step"] > 1:
+        adjust_nstep(policy.config["n_step"], policy.config["gamma"], batch)
+
+    # Create dummy prio-weights (1.0) in case we don't have any in
+    # the batch.
+    if PRIO_WEIGHTS not in batch:
+        batch[PRIO_WEIGHTS] = np.ones_like(batch[SampleBatch.REWARDS])
+
+    # Prioritize on the worker side.
+    if batch.count > 0 and policy.config["replay_buffer_config"].get(
+        "worker_side_prioritization", False
+    ):
+        td_errors = policy.compute_td_error(
+            batch[SampleBatch.OBS],
+            batch[SampleBatch.ACTIONS],
+            batch[SampleBatch.REWARDS],
+            batch[SampleBatch.NEXT_OBS],
+            batch[SampleBatch.TERMINATEDS],
+            batch[PRIO_WEIGHTS],
+        )
+        # Retain compatibility with old-style Replay args
+        epsilon = policy.config.get("replay_buffer_config", {}).get(
+            "prioritized_replay_eps"
+        ) or policy.config.get("prioritized_replay_eps")
+        if epsilon is None:
+            raise ValueError("prioritized_replay_eps not defined in config.")
+
+        new_priorities = np.abs(convert_to_numpy(td_errors)) + epsilon
+        batch[PRIO_WEIGHTS] = new_priorities
+
+    return batch
+
+
+DQNTFPolicy = build_tf_policy(
+    name="DQNTFPolicy",
+    get_default_config=lambda: ray.rllib.algorithms.dqn.dqn.DQNConfig(),
+    make_model=build_q_model,
+    action_distribution_fn=get_distribution_inputs_and_class,
+    loss_fn=build_q_losses,
+    stats_fn=build_q_stats,
+    postprocess_fn=postprocess_nstep_and_prio,
+    optimizer_fn=adam_optimizer,
+    compute_gradients_fn=clip_gradients,
+    extra_action_out_fn=lambda policy: {"q_values": policy.q_values},
+    extra_learn_fetches_fn=lambda policy: {"td_error": policy.q_loss.td_error},
+    before_loss_init=setup_mid_mixins,
+    after_init=setup_late_mixins,
+    mixins=[
+        TargetNetworkMixin,
+        ComputeTDErrorMixin,
+        LearningRateSchedule,
+    ],
+)

rllib/algorithms/dqn/dqn_torch_model.py

@@ -0,0 +1,177 @@
+"""PyTorch model for DQN"""
+
+from typing import Sequence
+
+import gymnasium as gym
+
+from ray.rllib.models.torch.misc import SlimFC
+from ray.rllib.models.torch.modules.noisy_layer import NoisyLayer
+from ray.rllib.models.torch.torch_modelv2 import TorchModelV2
+from ray.rllib.utils.annotations import OldAPIStack
+from ray.rllib.utils.framework import try_import_torch
+from ray.rllib.utils.typing import ModelConfigDict
+
+torch, nn = try_import_torch()
+
+
+@OldAPIStack
+class DQNTorchModel(TorchModelV2, nn.Module):
+    """Extension of standard TorchModelV2 to provide dueling-Q functionality."""
+
+    def __init__(
+        self,
+        obs_space: gym.spaces.Space,
+        action_space: gym.spaces.Space,
+        num_outputs: int,
+        model_config: ModelConfigDict,
+        name: str,
+        *,
+        q_hiddens: Sequence[int] = (256,),
+        dueling: bool = False,
+        dueling_activation: str = "relu",
+        num_atoms: int = 1,
+        use_noisy: bool = False,
+        v_min: float = -10.0,
+        v_max: float = 10.0,
+        sigma0: float = 0.5,
+        # TODO(sven): Move `add_layer_norm` into ModelCatalog as
+        #  generic option, then error if we use ParameterNoise as
+        #  Exploration type and do not have any LayerNorm layers in
+        #  the net.
+        add_layer_norm: bool = False
+    ):
+        """Initialize variables of this model.
+
+        Extra model kwargs:
+            q_hiddens (Sequence[int]): List of layer-sizes after(!) the
+                Advantages(A)/Value(V)-split. Hence, each of the A- and V-
+                branches will have this structure of Dense layers. To define
+                the NN before this A/V-split, use - as always -
+                config["model"]["fcnet_hiddens"].
+            dueling: Whether to build the advantage(A)/value(V) heads
+                for DDQN. If True, Q-values are calculated as:
+                Q = (A - mean[A]) + V. If False, raw NN output is interpreted
+                as Q-values.
+            dueling_activation: The activation to use for all dueling
+                layers (A- and V-branch). One of "relu", "tanh", "linear".
+            num_atoms: If >1, enables distributional DQN.
+            use_noisy: Use noisy layers.
+            v_min: Min value support for distributional DQN.
+            v_max: Max value support for distributional DQN.
+            sigma0 (float): Initial value of noisy layers.
+            add_layer_norm: Enable layer norm (for param noise).
+        """
+        nn.Module.__init__(self)
+        super(DQNTorchModel, self).__init__(
+            obs_space, action_space, num_outputs, model_config, name
+        )
+
+        self.dueling = dueling
+        self.num_atoms = num_atoms
+        self.v_min = v_min
+        self.v_max = v_max
+        self.sigma0 = sigma0
+        ins = num_outputs
+
+        advantage_module = nn.Sequential()
+        value_module = nn.Sequential()
+
+        # Dueling case: Build the shared (advantages and value) fc-network.
+        for i, n in enumerate(q_hiddens):
+            if use_noisy:
+                advantage_module.add_module(
+                    "dueling_A_{}".format(i),
+                    NoisyLayer(
+                        ins, n, sigma0=self.sigma0, activation=dueling_activation
+                    ),
+                )
+                value_module.add_module(
+                    "dueling_V_{}".format(i),
+                    NoisyLayer(
+                        ins, n, sigma0=self.sigma0, activation=dueling_activation
+                    ),
+                )
+            else:
+                advantage_module.add_module(
+                    "dueling_A_{}".format(i),
+                    SlimFC(ins, n, activation_fn=dueling_activation),
+                )
+                value_module.add_module(
+                    "dueling_V_{}".format(i),
+                    SlimFC(ins, n, activation_fn=dueling_activation),
+                )
+                # Add LayerNorm after each Dense.
+                if add_layer_norm:
+                    advantage_module.add_module(
+                        "LayerNorm_A_{}".format(i), nn.LayerNorm(n)
+                    )
+                    value_module.add_module("LayerNorm_V_{}".format(i), nn.LayerNorm(n))
+            ins = n
+
+        # Actual Advantages layer (nodes=num-actions).
+        if use_noisy:
+            advantage_module.add_module(
+                "A",
+                NoisyLayer(
+                    ins, self.action_space.n * self.num_atoms, sigma0, activation=None
+                ),
+            )
+        elif q_hiddens:
+            advantage_module.add_module(
+                "A", SlimFC(ins, action_space.n * self.num_atoms, activation_fn=None)
+            )
+
+        self.advantage_module = advantage_module
+
+        # Value layer (nodes=1).
+        if self.dueling:
+            if use_noisy:
+                value_module.add_module(
+                    "V", NoisyLayer(ins, self.num_atoms, sigma0, activation=None)
+                )
+            elif q_hiddens:
+                value_module.add_module(
+                    "V", SlimFC(ins, self.num_atoms, activation_fn=None)
+                )
+            self.value_module = value_module
+
+    def get_q_value_distributions(self, model_out):
+        """Returns distributional values for Q(s, a) given a state embedding.
+
+        Override this in your custom model to customize the Q output head.
+
+        Args:
+            model_out: Embedding from the model layers.
+
+        Returns:
+            (action_scores, logits, dist) if num_atoms == 1, otherwise
+            (action_scores, z, support_logits_per_action, logits, dist)
+        """
+        action_scores = self.advantage_module(model_out)
+
+        if self.num_atoms > 1:
+            # Distributional Q-learning uses a discrete support z
+            # to represent the action value distribution
+            z = torch.arange(0.0, self.num_atoms, dtype=torch.float32).to(
+                action_scores.device
+            )
+            z = self.v_min + z * (self.v_max - self.v_min) / float(self.num_atoms - 1)
+
+            support_logits_per_action = torch.reshape(
+                action_scores, shape=(-1, self.action_space.n, self.num_atoms)
+            )
+            support_prob_per_action = nn.functional.softmax(
+                support_logits_per_action, dim=-1
+            )
+            action_scores = torch.sum(z * support_prob_per_action, dim=-1)
+            logits = support_logits_per_action
+            probs = support_prob_per_action
+            return action_scores, z, support_logits_per_action, logits, probs
+        else:
+            logits = torch.unsqueeze(torch.ones_like(action_scores), -1)
+            return action_scores, logits, logits
+
+    def get_state_value(self, model_out):
+        """Returns the state value prediction for the given state embedding."""
+
+        return self.value_module(model_out)

rllib/algorithms/dqn/dqn_torch_policy.py

@@ -0,0 +1,519 @@
+"""PyTorch policy class used for DQN"""
+
+from typing import Dict, List, Tuple
+
+import gymnasium as gym
+
+import ray
+from ray.rllib.algorithms.dqn.dqn_tf_policy import (
+    PRIO_WEIGHTS,
+    Q_SCOPE,
+    Q_TARGET_SCOPE,
+    postprocess_nstep_and_prio,
+)
+from ray.rllib.algorithms.dqn.dqn_torch_model import DQNTorchModel
+from ray.rllib.models.catalog import ModelCatalog
+from ray.rllib.models.modelv2 import ModelV2
+from ray.rllib.models.torch.torch_action_dist import (
+    TorchDistributionWrapper,
+    get_torch_categorical_class_with_temperature,
+)
+from ray.rllib.policy.policy import Policy
+from ray.rllib.policy.policy_template import build_policy_class
+from ray.rllib.policy.sample_batch import SampleBatch
+from ray.rllib.policy.torch_mixins import (
+    LearningRateSchedule,
+    TargetNetworkMixin,
+)
+from ray.rllib.utils.annotations import OldAPIStack
+from ray.rllib.utils.error import UnsupportedSpaceException
+from ray.rllib.utils.exploration.parameter_noise import ParameterNoise
+from ray.rllib.utils.framework import try_import_torch
+from ray.rllib.utils.torch_utils import (
+    FLOAT_MIN,
+    apply_grad_clipping,
+    concat_multi_gpu_td_errors,
+    huber_loss,
+    l2_loss,
+    reduce_mean_ignore_inf,
+    softmax_cross_entropy_with_logits,
+)
+from ray.rllib.utils.typing import AlgorithmConfigDict, TensorType
+
+torch, nn = try_import_torch()
+F = None
+if nn:
+    F = nn.functional
+
+
+@OldAPIStack
+class QLoss:
+    def __init__(
+        self,
+        q_t_selected: TensorType,
+        q_logits_t_selected: TensorType,
+        q_tp1_best: TensorType,
+        q_probs_tp1_best: TensorType,
+        importance_weights: TensorType,
+        rewards: TensorType,
+        done_mask: TensorType,
+        gamma=0.99,
+        n_step=1,
+        num_atoms=1,
+        v_min=-10.0,
+        v_max=10.0,
+        loss_fn=huber_loss,
+    ):
+
+        if num_atoms > 1:
+            # Distributional Q-learning which corresponds to an entropy loss
+            z = torch.arange(0.0, num_atoms, dtype=torch.float32).to(rewards.device)
+            z = v_min + z * (v_max - v_min) / float(num_atoms - 1)
+
+            # (batch_size, 1) * (1, num_atoms) = (batch_size, num_atoms)
+            r_tau = torch.unsqueeze(rewards, -1) + gamma**n_step * torch.unsqueeze(
+                1.0 - done_mask, -1
+            ) * torch.unsqueeze(z, 0)
+            r_tau = torch.clamp(r_tau, v_min, v_max)
+            b = (r_tau - v_min) / ((v_max - v_min) / float(num_atoms - 1))
+            lb = torch.floor(b)
+            ub = torch.ceil(b)
+
+            # Indispensable judgement which is missed in most implementations
+            # when b happens to be an integer, lb == ub, so pr_j(s', a*) will
+            # be discarded because (ub-b) == (b-lb) == 0.
+            floor_equal_ceil = ((ub - lb) < 0.5).float()
+
+            # (batch_size, num_atoms, num_atoms)
+            l_project = F.one_hot(lb.long(), num_atoms)
+            # (batch_size, num_atoms, num_atoms)
+            u_project = F.one_hot(ub.long(), num_atoms)
+            ml_delta = q_probs_tp1_best * (ub - b + floor_equal_ceil)
+            mu_delta = q_probs_tp1_best * (b - lb)
+            ml_delta = torch.sum(l_project * torch.unsqueeze(ml_delta, -1), dim=1)
+            mu_delta = torch.sum(u_project * torch.unsqueeze(mu_delta, -1), dim=1)
+            m = ml_delta + mu_delta
+
+            # Rainbow paper claims that using this cross entropy loss for
+            # priority is robust and insensitive to `prioritized_replay_alpha`
+            self.td_error = softmax_cross_entropy_with_logits(
+                logits=q_logits_t_selected, labels=m.detach()
+            )
+            self.loss = torch.mean(self.td_error * importance_weights)
+            self.stats = {
+                # TODO: better Q stats for dist dqn
+            }
+        else:
+            q_tp1_best_masked = (1.0 - done_mask) * q_tp1_best
+
+            # compute RHS of bellman equation
+            q_t_selected_target = rewards + gamma**n_step * q_tp1_best_masked
+
+            # compute the error (potentially clipped)
+            self.td_error = q_t_selected - q_t_selected_target.detach()
+            self.loss = torch.mean(importance_weights.float() * loss_fn(self.td_error))
+            self.stats = {
+                "mean_q": torch.mean(q_t_selected),
+                "min_q": torch.min(q_t_selected),
+                "max_q": torch.max(q_t_selected),
+            }
+
+
+@OldAPIStack
+class ComputeTDErrorMixin:
+    """Assign the `compute_td_error` method to the DQNTorchPolicy
+
+    This allows us to prioritize on the worker side.
+    """
+
+    def __init__(self):
+        def compute_td_error(
+            obs_t, act_t, rew_t, obs_tp1, terminateds_mask, importance_weights
+        ):
+            input_dict = self._lazy_tensor_dict({SampleBatch.CUR_OBS: obs_t})
+            input_dict[SampleBatch.ACTIONS] = act_t
+            input_dict[SampleBatch.REWARDS] = rew_t
+            input_dict[SampleBatch.NEXT_OBS] = obs_tp1
+            input_dict[SampleBatch.TERMINATEDS] = terminateds_mask
+            input_dict[PRIO_WEIGHTS] = importance_weights
+
+            # Do forward pass on loss to update td error attribute
+            build_q_losses(self, self.model, None, input_dict)
+
+            return self.model.tower_stats["q_loss"].td_error
+
+        self.compute_td_error = compute_td_error
+
+
+@OldAPIStack
+def build_q_model_and_distribution(
+    policy: Policy,
+    obs_space: gym.spaces.Space,
+    action_space: gym.spaces.Space,
+    config: AlgorithmConfigDict,
+) -> Tuple[ModelV2, TorchDistributionWrapper]:
+    """Build q_model and target_model for DQN
+
+    Args:
+        policy: The policy, which will use the model for optimization.
+        obs_space (gym.spaces.Space): The policy's observation space.
+        action_space (gym.spaces.Space): The policy's action space.
+        config (AlgorithmConfigDict):
+
+    Returns:
+        (q_model, TorchCategorical)
+            Note: The target q model will not be returned, just assigned to
+            `policy.target_model`.
+    """
+    if not isinstance(action_space, gym.spaces.Discrete):
+        raise UnsupportedSpaceException(
+            "Action space {} is not supported for DQN.".format(action_space)
+        )
+
+    if config["hiddens"]:
+        # try to infer the last layer size, otherwise fall back to 256
+        num_outputs = ([256] + list(config["model"]["fcnet_hiddens"]))[-1]
+        config["model"]["no_final_linear"] = True
+    else:
+        num_outputs = action_space.n
+
+    # TODO(sven): Move option to add LayerNorm after each Dense
+    #  generically into ModelCatalog.
+    add_layer_norm = (
+        isinstance(getattr(policy, "exploration", None), ParameterNoise)
+        or config["exploration_config"]["type"] == "ParameterNoise"
+    )
+
+    model = ModelCatalog.get_model_v2(
+        obs_space=obs_space,
+        action_space=action_space,
+        num_outputs=num_outputs,
+        model_config=config["model"],
+        framework="torch",
+        model_interface=DQNTorchModel,
+        name=Q_SCOPE,
+        q_hiddens=config["hiddens"],
+        dueling=config["dueling"],
+        num_atoms=config["num_atoms"],
+        use_noisy=config["noisy"],
+        v_min=config["v_min"],
+        v_max=config["v_max"],
+        sigma0=config["sigma0"],
+        # TODO(sven): Move option to add LayerNorm after each Dense
+        #  generically into ModelCatalog.
+        add_layer_norm=add_layer_norm,
+    )
+
+    policy.target_model = ModelCatalog.get_model_v2(
+        obs_space=obs_space,
+        action_space=action_space,
+        num_outputs=num_outputs,
+        model_config=config["model"],
+        framework="torch",
+        model_interface=DQNTorchModel,
+        name=Q_TARGET_SCOPE,
+        q_hiddens=config["hiddens"],
+        dueling=config["dueling"],
+        num_atoms=config["num_atoms"],
+        use_noisy=config["noisy"],
+        v_min=config["v_min"],
+        v_max=config["v_max"],
+        sigma0=config["sigma0"],
+        # TODO(sven): Move option to add LayerNorm after each Dense
+        #  generically into ModelCatalog.
+        add_layer_norm=add_layer_norm,
+    )
+
+    # Return a Torch TorchCategorical distribution where the temperature
+    # parameter is partially binded to the configured value.
+    temperature = config["categorical_distribution_temperature"]
+
+    return model, get_torch_categorical_class_with_temperature(temperature)
+
+
+@OldAPIStack
+def get_distribution_inputs_and_class(
+    policy: Policy,
+    model: ModelV2,
+    input_dict: SampleBatch,
+    *,
+    explore: bool = True,
+    is_training: bool = False,
+    **kwargs
+) -> Tuple[TensorType, type, List[TensorType]]:
+    q_vals = compute_q_values(
+        policy, model, input_dict, explore=explore, is_training=is_training
+    )
+    q_vals = q_vals[0] if isinstance(q_vals, tuple) else q_vals
+
+    model.tower_stats["q_values"] = q_vals
+
+    # Return a Torch TorchCategorical distribution where the temperature
+    # parameter is partially binded to the configured value.
+    temperature = policy.config["categorical_distribution_temperature"]
+
+    return (
+        q_vals,
+        get_torch_categorical_class_with_temperature(temperature),
+        [],  # state-out
+    )
+
+
+@OldAPIStack
+def build_q_losses(policy: Policy, model, _, train_batch: SampleBatch) -> TensorType:
+    """Constructs the loss for DQNTorchPolicy.
+
+    Args:
+        policy: The Policy to calculate the loss for.
+        model (ModelV2): The Model to calculate the loss for.
+        train_batch: The training data.
+
+    Returns:
+        TensorType: A single loss tensor.
+    """
+
+    config = policy.config
+    # Q-network evaluation.
+    q_t, q_logits_t, q_probs_t, _ = compute_q_values(
+        policy,
+        model,
+        {"obs": train_batch[SampleBatch.CUR_OBS]},
+        explore=False,
+        is_training=True,
+    )
+
+    # Target Q-network evaluation.
+    q_tp1, q_logits_tp1, q_probs_tp1, _ = compute_q_values(
+        policy,
+        policy.target_models[model],
+        {"obs": train_batch[SampleBatch.NEXT_OBS]},
+        explore=False,
+        is_training=True,
+    )
+
+    # Q scores for actions which we know were selected in the given state.
+    one_hot_selection = F.one_hot(
+        train_batch[SampleBatch.ACTIONS].long(), policy.action_space.n
+    )
+    q_t_selected = torch.sum(
+        torch.where(q_t > FLOAT_MIN, q_t, torch.tensor(0.0, device=q_t.device))
+        * one_hot_selection,
+        1,
+    )
+    q_logits_t_selected = torch.sum(
+        q_logits_t * torch.unsqueeze(one_hot_selection, -1), 1
+    )
+
+    # compute estimate of best possible value starting from state at t + 1
+    if config["double_q"]:
+        (
+            q_tp1_using_online_net,
+            q_logits_tp1_using_online_net,
+            q_dist_tp1_using_online_net,
+            _,
+        ) = compute_q_values(
+            policy,
+            model,
+            {"obs": train_batch[SampleBatch.NEXT_OBS]},
+            explore=False,
+            is_training=True,
+        )
+        q_tp1_best_using_online_net = torch.argmax(q_tp1_using_online_net, 1)
+        q_tp1_best_one_hot_selection = F.one_hot(
+            q_tp1_best_using_online_net, policy.action_space.n
+        )
+        q_tp1_best = torch.sum(
+            torch.where(
+                q_tp1 > FLOAT_MIN, q_tp1, torch.tensor(0.0, device=q_tp1.device)
+            )
+            * q_tp1_best_one_hot_selection,
+            1,
+        )
+        q_probs_tp1_best = torch.sum(
+            q_probs_tp1 * torch.unsqueeze(q_tp1_best_one_hot_selection, -1), 1
+        )
+    else:
+        q_tp1_best_one_hot_selection = F.one_hot(
+            torch.argmax(q_tp1, 1), policy.action_space.n
+        )
+        q_tp1_best = torch.sum(
+            torch.where(
+                q_tp1 > FLOAT_MIN, q_tp1, torch.tensor(0.0, device=q_tp1.device)
+            )
+            * q_tp1_best_one_hot_selection,
+            1,
+        )
+        q_probs_tp1_best = torch.sum(
+            q_probs_tp1 * torch.unsqueeze(q_tp1_best_one_hot_selection, -1), 1
+        )
+
+    loss_fn = huber_loss if policy.config["td_error_loss_fn"] == "huber" else l2_loss
+
+    q_loss = QLoss(
+        q_t_selected,
+        q_logits_t_selected,
+        q_tp1_best,
+        q_probs_tp1_best,
+        train_batch[PRIO_WEIGHTS],
+        train_batch[SampleBatch.REWARDS],
+        train_batch[SampleBatch.TERMINATEDS].float(),
+        config["gamma"],
+        config["n_step"],
+        config["num_atoms"],
+        config["v_min"],
+        config["v_max"],
+        loss_fn,
+    )
+
+    # 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["td_error"] = q_loss.td_error
+    # TD-error tensor in final stats
+    # will be concatenated and retrieved for each individual batch item.
+    model.tower_stats["q_loss"] = q_loss
+
+    return q_loss.loss
+
+
+@OldAPIStack
+def adam_optimizer(
+    policy: Policy, config: AlgorithmConfigDict
+) -> "torch.optim.Optimizer":
+
+    # By this time, the models have been moved to the GPU - if any - and we
+    # can define our optimizers using the correct CUDA variables.
+    if not hasattr(policy, "q_func_vars"):
+        policy.q_func_vars = policy.model.variables()
+
+    return torch.optim.Adam(
+        policy.q_func_vars, lr=policy.cur_lr, eps=config["adam_epsilon"]
+    )
+
+
+@OldAPIStack
+def build_q_stats(policy: Policy, batch) -> Dict[str, TensorType]:
+    stats = {}
+    for stats_key in policy.model_gpu_towers[0].tower_stats["q_loss"].stats.keys():
+        stats[stats_key] = torch.mean(
+            torch.stack(
+                [
+                    t.tower_stats["q_loss"].stats[stats_key].to(policy.device)
+                    for t in policy.model_gpu_towers
+                    if "q_loss" in t.tower_stats
+                ]
+            )
+        )
+    stats["cur_lr"] = policy.cur_lr
+    return stats
+
+
+@OldAPIStack
+def setup_early_mixins(
+    policy: Policy, obs_space, action_space, config: AlgorithmConfigDict
+) -> None:
+    LearningRateSchedule.__init__(policy, config["lr"], config["lr_schedule"])
+
+
+@OldAPIStack
+def before_loss_init(
+    policy: Policy,
+    obs_space: gym.spaces.Space,
+    action_space: gym.spaces.Space,
+    config: AlgorithmConfigDict,
+) -> None:
+    ComputeTDErrorMixin.__init__(policy)
+    TargetNetworkMixin.__init__(policy)
+
+
+@OldAPIStack
+def compute_q_values(
+    policy: Policy,
+    model: ModelV2,
+    input_dict,
+    state_batches=None,
+    seq_lens=None,
+    explore=None,
+    is_training: bool = False,
+):
+    config = policy.config
+
+    model_out, state = model(input_dict, state_batches or [], seq_lens)
+
+    if config["num_atoms"] > 1:
+        (
+            action_scores,
+            z,
+            support_logits_per_action,
+            logits,
+            probs_or_logits,
+        ) = model.get_q_value_distributions(model_out)
+    else:
+        (action_scores, logits, probs_or_logits) = model.get_q_value_distributions(
+            model_out
+        )
+
+    if config["dueling"]:
+        state_score = model.get_state_value(model_out)
+        if policy.config["num_atoms"] > 1:
+            support_logits_per_action_mean = torch.mean(
+                support_logits_per_action, dim=1
+            )
+            support_logits_per_action_centered = (
+                support_logits_per_action
+                - torch.unsqueeze(support_logits_per_action_mean, dim=1)
+            )
+            support_logits_per_action = (
+                torch.unsqueeze(state_score, dim=1) + support_logits_per_action_centered
+            )
+            support_prob_per_action = nn.functional.softmax(
+                support_logits_per_action, dim=-1
+            )
+            value = torch.sum(z * support_prob_per_action, dim=-1)
+            logits = support_logits_per_action
+            probs_or_logits = support_prob_per_action
+        else:
+            advantages_mean = reduce_mean_ignore_inf(action_scores, 1)
+            advantages_centered = action_scores - torch.unsqueeze(advantages_mean, 1)
+            value = state_score + advantages_centered
+    else:
+        value = action_scores
+
+    return value, logits, probs_or_logits, state
+
+
+@OldAPIStack
+def grad_process_and_td_error_fn(
+    policy: Policy, optimizer: "torch.optim.Optimizer", loss: TensorType
+) -> Dict[str, TensorType]:
+    # Clip grads if configured.
+    return apply_grad_clipping(policy, optimizer, loss)
+
+
+@OldAPIStack
+def extra_action_out_fn(
+    policy: Policy, input_dict, state_batches, model, action_dist
+) -> Dict[str, TensorType]:
+    return {"q_values": model.tower_stats["q_values"]}
+
+
+DQNTorchPolicy = build_policy_class(
+    name="DQNTorchPolicy",
+    framework="torch",
+    loss_fn=build_q_losses,
+    get_default_config=lambda: ray.rllib.algorithms.dqn.dqn.DQNConfig(),
+    make_model_and_action_dist=build_q_model_and_distribution,
+    action_distribution_fn=get_distribution_inputs_and_class,
+    stats_fn=build_q_stats,
+    postprocess_fn=postprocess_nstep_and_prio,
+    optimizer_fn=adam_optimizer,
+    extra_grad_process_fn=grad_process_and_td_error_fn,
+    extra_learn_fetches_fn=concat_multi_gpu_td_errors,
+    extra_action_out_fn=extra_action_out_fn,
+    before_init=setup_early_mixins,
+    before_loss_init=before_loss_init,
+    mixins=[
+        TargetNetworkMixin,
+        ComputeTDErrorMixin,
+        LearningRateSchedule,
+    ],
+)

rllib/algorithms/dqn/tests/test_dqn.py

@@ -0,0 +1,54 @@
+import unittest
+
+import ray
+import ray.rllib.algorithms.dqn as dqn
+from ray.rllib.utils.test_utils import check_train_results_new_api_stack
+
+
+class TestDQN(unittest.TestCase):
+    @classmethod
+    def setUpClass(cls) -> None:
+        ray.init()
+
+    @classmethod
+    def tearDownClass(cls) -> None:
+        ray.shutdown()
+
+    def test_dqn_compilation(self):
+        """Test whether DQN can be built and trained."""
+        num_iterations = 2
+        config = (
+            dqn.dqn.DQNConfig()
+            .environment("CartPole-v1")
+            .env_runners(num_env_runners=2)
+            .training(num_steps_sampled_before_learning_starts=0)
+        )
+
+        # Double-dueling DQN.
+        print("Double-dueling")
+        algo = config.build()
+        for i in range(num_iterations):
+            results = algo.train()
+            check_train_results_new_api_stack(results)
+            print(results)
+
+        algo.stop()
+
+        # Rainbow.
+        print("Rainbow")
+        config.training(num_atoms=10, double_q=True, dueling=True, n_step=5)
+        algo = config.build()
+        for i in range(num_iterations):
+            results = algo.train()
+            check_train_results_new_api_stack(results)
+            print(results)
+
+        algo.stop()
+
+
+if __name__ == "__main__":
+    import sys
+
+    import pytest
+
+    sys.exit(pytest.main(["-v", __file__]))

rllib/algorithms/dqn/tests/test_dqn_rl_module.py

@@ -0,0 +1,190 @@
+import dataclasses
+
+import numpy as np
+import pytest
+import tree
+from gymnasium.spaces import Box, Dict, Discrete
+
+from ray.rllib.algorithms.dqn.dqn_catalog import DQNCatalog
+from ray.rllib.algorithms.dqn.torch.default_dqn_torch_rl_module import (
+    DefaultDQNTorchRLModule,
+)
+from ray.rllib.core.columns import Columns
+from ray.rllib.core.models.base import ENCODER_OUT
+from ray.rllib.core.rl_module.default_model_config import DefaultModelConfig
+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()
+
+
+# Custom encoder, config and catalog to test Dict observation spaces.
+# RLlib does not build encoders for Dict observation spaces out of the box so we define our own.
+class DictFlattenEncoder(nn.Module):
+    def __init__(self, obs_space, output_dim=64):
+        super().__init__()
+        total_dim = sum(
+            int(np.prod(space.shape)) for space in obs_space.spaces.values()
+        )
+        self.net = nn.Sequential(
+            nn.Linear(total_dim, output_dim),
+            nn.ReLU(),
+        )
+
+    def forward(self, inputs):
+        obs = inputs[Columns.OBS]
+        flat_obs = torch.cat(
+            [obs[k].reshape(obs[k].shape[0], -1) for k in sorted(obs.keys())],
+            dim=-1,
+        )
+        return {ENCODER_OUT: self.net(flat_obs)}
+
+
+class DictEncoderConfig:
+    def __init__(self, obs_space, output_dim=64):
+        self.obs_space = obs_space
+        self.output_dims = (output_dim,)
+
+    def build(self, framework):
+        return DictFlattenEncoder(self.obs_space, output_dim=self.output_dims[0])
+
+
+class DictObsDQNCatalog(DQNCatalog):
+    @classmethod
+    def _get_encoder_config(
+        cls, observation_space, model_config_dict, action_space=None
+    ):
+        return DictEncoderConfig(observation_space, output_dim=64)
+
+
+# Observation space definitions.
+OBS_SPACES = {
+    "box": Box(low=-1.0, high=1.0, shape=(8,), dtype=np.float32),
+    "image": Box(low=0, high=255, shape=(64, 64, 3), dtype=np.uint8),
+    "dict": Dict(
+        {
+            "sensors": Box(low=-1.0, high=1.0, shape=(4,), dtype=np.float32),
+            "position": Box(low=-10.0, high=10.0, shape=(3,), dtype=np.float32),
+            "mode": Discrete(4),
+        }
+    ),
+}
+
+
+def _get_dqn_module(observation_space, action_space, **config_overrides):
+    model_config = dataclasses.asdict(DefaultModelConfig())
+    model_config.update(
+        {
+            "double_q": True,
+            "dueling": True,
+            "epsilon": [(0, 1.0), (10000, 0.05)],
+            "num_atoms": 1,
+            "v_min": -10.0,
+            "v_max": 10.0,
+        }
+    )
+    model_config.update(config_overrides)
+
+    # Use custom catalog for Dict observation spaces.
+    catalog_class = (
+        DictObsDQNCatalog if isinstance(observation_space, Dict) else DQNCatalog
+    )
+
+    module = DefaultDQNTorchRLModule(
+        observation_space=observation_space,
+        action_space=action_space,
+        model_config=model_config,
+        catalog_class=catalog_class,
+        inference_only=False,
+    )
+
+    # Create target networks (normally done by the learner).
+    module.make_target_networks()
+    return module
+
+
+class TestDQNRLModule:
+    @pytest.mark.parametrize("obs_space_name", ["box", "image", "dict"])
+    @pytest.mark.parametrize("forward_method", ["train", "exploration", "inference"])
+    @pytest.mark.parametrize("double_q", [True, False])
+    @pytest.mark.parametrize("dueling", [True, False])
+    def test_forward(self, obs_space_name, forward_method, double_q, dueling):
+        """Test forward methods with different obs spaces and config settings."""
+        obs_space = OBS_SPACES[obs_space_name]
+        action_space = Discrete(4)
+
+        module = _get_dqn_module(
+            obs_space, action_space, double_q=double_q, dueling=dueling
+        )
+
+        if (
+            forward_method == "train"
+        ):  # forward train needs batching, exploration and inference don't
+            module.train()
+            # Create a batch first
+            batch_size = 4
+            obs_list = [obs_space.sample() for _ in range(batch_size)]
+            next_obs_list = [obs_space.sample() for _ in range(batch_size)]
+
+            obs_batch = tree.map_structure(
+                lambda *x: np.stack(x, axis=0, dtype=np.float32), *obs_list
+            )
+            next_obs_batch = tree.map_structure(
+                lambda *x: np.stack(x, axis=0, dtype=np.float32), *next_obs_list
+            )
+
+            batch = {
+                Columns.OBS: convert_to_torch_tensor(obs_batch),
+                Columns.NEXT_OBS: convert_to_torch_tensor(next_obs_batch),
+                Columns.ACTIONS: convert_to_torch_tensor(
+                    np.array([0] * batch_size, dtype=np.int64)
+                ),
+                Columns.REWARDS: convert_to_torch_tensor(
+                    np.array([1.0] * batch_size, dtype=np.float32)
+                ),
+                Columns.TERMINATEDS: convert_to_torch_tensor(
+                    np.array([False] * batch_size, dtype=np.bool_)
+                ),
+                Columns.TRUNCATEDS: convert_to_torch_tensor(
+                    np.array([False] * batch_size, dtype=np.bool_)
+                ),
+            }
+
+            # Forward pass and check outputs
+            output = module.forward_train(batch)
+
+            assert "qf_preds" in output
+            assert output["qf_preds"].shape == (4, action_space.n)
+
+            if double_q:
+                assert "qf_next_preds" in output
+                assert output["qf_next_preds"].shape == (4, action_space.n)
+            else:
+                assert "qf_next_preds" not in output
+        else:
+            module.eval()
+            # Create a single observation batch
+            obs = obs_space.sample()
+            if isinstance(obs_space, Dict):
+                obs_tensor = tree.map_structure(
+                    lambda x: convert_to_torch_tensor(x.astype(np.float32)[None]),
+                    obs,
+                )
+            else:
+                obs_tensor = convert_to_torch_tensor(obs.astype(np.float32)[None])
+            batch = {Columns.OBS: obs_tensor}
+
+            # Forward pass and check outputs
+            if forward_method == "exploration":
+                output = module.forward_exploration(batch, t=0)
+            else:
+                output = module.forward_inference(batch)
+
+            assert Columns.ACTIONS in output
+            assert output[Columns.ACTIONS].shape == (1,)
+
+
+if __name__ == "__main__":
+    import sys
+
+    sys.exit(pytest.main(["-v", __file__]))

rllib/algorithms/dqn/torch/default_dqn_torch_rl_module.py

@@ -0,0 +1,330 @@
+from typing import Dict, Union
+
+import tree
+
+from ray.rllib.algorithms.dqn.default_dqn_rl_module import (
+    ATOMS,
+    QF_LOGITS,
+    QF_NEXT_PREDS,
+    QF_PREDS,
+    QF_PROBS,
+    QF_TARGET_NEXT_PREDS,
+    QF_TARGET_NEXT_PROBS,
+    DefaultDQNRLModule,
+)
+from ray.rllib.algorithms.dqn.dqn_catalog import DQNCatalog
+from ray.rllib.core.columns import Columns
+from ray.rllib.core.models.base import ENCODER_OUT, Encoder, Model
+from ray.rllib.core.rl_module.apis.q_net_api import QNetAPI
+from ray.rllib.core.rl_module.rl_module import RLModule
+from ray.rllib.core.rl_module.torch.torch_rl_module import TorchRLModule
+from ray.rllib.utils.annotations import override
+from ray.rllib.utils.framework import try_import_torch
+from ray.rllib.utils.typing import TensorStructType, TensorType
+from ray.util.annotations import DeveloperAPI
+
+torch, nn = try_import_torch()
+
+
+@DeveloperAPI
+class DefaultDQNTorchRLModule(TorchRLModule, DefaultDQNRLModule):
+    framework: str = "torch"
+
+    def __init__(self, *args, **kwargs):
+        catalog_class = kwargs.pop("catalog_class", None)
+        if catalog_class is None:
+            catalog_class = DQNCatalog
+        super().__init__(*args, **kwargs, catalog_class=catalog_class)
+
+    @override(RLModule)
+    def _forward_inference(self, batch: Dict[str, TensorType]) -> Dict[str, TensorType]:
+        # Q-network forward pass.
+        qf_outs = self.compute_q_values(batch)
+
+        # Get action distribution.
+        action_dist_cls = self.get_exploration_action_dist_cls()
+        action_dist = action_dist_cls.from_logits(qf_outs[QF_PREDS])
+        # Note, the deterministic version of the categorical distribution
+        # outputs directly the `argmax` of the logits.
+        exploit_actions = action_dist.to_deterministic().sample()
+
+        output = {Columns.ACTIONS: exploit_actions}
+        if Columns.STATE_OUT in qf_outs:
+            output[Columns.STATE_OUT] = qf_outs[Columns.STATE_OUT]
+
+        # In inference, we only need the exploitation actions.
+        return output
+
+    @override(RLModule)
+    def _forward_exploration(
+        self, batch: Dict[str, TensorType], t: int
+    ) -> Dict[str, TensorType]:
+        # Define the return dictionary.
+        output = {}
+
+        # Q-network forward pass.
+        qf_outs = self.compute_q_values(batch)
+
+        # Get action distribution.
+        action_dist_cls = self.get_exploration_action_dist_cls()
+        action_dist = action_dist_cls.from_logits(qf_outs[QF_PREDS])
+        # Note, the deterministic version of the categorical distribution
+        # outputs directly the `argmax` of the logits.
+        exploit_actions = action_dist.to_deterministic().sample()
+
+        # We need epsilon greedy to support exploration.
+        # TODO (simon): Implement sampling for nested spaces.
+        # Update scheduler.
+        self.epsilon_schedule.update(t)
+        # Get the actual epsilon,
+        epsilon = self.epsilon_schedule.get_current_value()
+        # Apply epsilon-greedy exploration.
+        B = qf_outs[QF_PREDS].shape[0]
+        random_actions = torch.squeeze(
+            torch.multinomial(
+                (
+                    torch.nan_to_num(
+                        qf_outs[QF_PREDS].reshape(-1, qf_outs[QF_PREDS].size(-1)),
+                        neginf=0.0,
+                    )
+                    != 0.0
+                ).float(),
+                num_samples=1,
+            ),
+            dim=1,
+        )
+
+        actions = torch.where(
+            torch.rand((B,)) < epsilon,
+            random_actions,
+            exploit_actions,
+        )
+
+        # Add the actions to the return dictionary.
+        output[Columns.ACTIONS] = actions
+
+        # If this is a stateful module, add output states.
+        if Columns.STATE_OUT in qf_outs:
+            output[Columns.STATE_OUT] = qf_outs[Columns.STATE_OUT]
+
+        return output
+
+    @override(RLModule)
+    def _forward_train(
+        self, batch: Dict[str, TensorType]
+    ) -> Dict[str, TensorStructType]:
+        if self.inference_only:
+            raise RuntimeError(
+                "Trying to train a module that is not a learner module. Set the "
+                "flag `inference_only=False` when building the module."
+            )
+        output = {}
+
+        # If we use a double-Q setup.
+        if self.uses_double_q:
+            # Then we need to make a single forward pass with both,
+            # current and next observations.
+            batch_base = {
+                Columns.OBS: tree.map_structure(
+                    lambda x, y: torch.cat([x, y], dim=0),
+                    batch[Columns.OBS],
+                    batch[Columns.NEXT_OBS],
+                )
+            }
+            # If this is a stateful module add the input states.
+            if Columns.STATE_IN in batch:
+                # Add both, the input state for the actual observation and
+                # the one for the next observation.
+                batch_base.update(
+                    {
+                        Columns.STATE_IN: tree.map_structure(
+                            lambda t1, t2: torch.cat([t1, t2], dim=0),
+                            batch[Columns.STATE_IN],
+                            batch[Columns.NEXT_STATE_IN],
+                        )
+                    }
+                )
+        # Otherwise we can just use the current observations.
+        else:
+            batch_base = {Columns.OBS: batch[Columns.OBS]}
+            # If this is a stateful module add the input state.
+            if Columns.STATE_IN in batch:
+                batch_base.update({Columns.STATE_IN: batch[Columns.STATE_IN]})
+
+        batch_target = {Columns.OBS: batch[Columns.NEXT_OBS]}
+
+        # If we have a stateful encoder, add the states for the target forward
+        # pass.
+        if Columns.NEXT_STATE_IN in batch:
+            batch_target.update({Columns.STATE_IN: batch[Columns.NEXT_STATE_IN]})
+
+        # Q-network forward passes.
+        qf_outs = self.compute_q_values(batch_base)
+        if self.uses_double_q:
+            output[QF_PREDS], output[QF_NEXT_PREDS] = torch.chunk(
+                qf_outs[QF_PREDS], chunks=2, dim=0
+            )
+        else:
+            output[QF_PREDS] = qf_outs[QF_PREDS]
+        # The target Q-values for the next observations.
+        qf_target_next_outs = self.forward_target(batch_target)
+        output[QF_TARGET_NEXT_PREDS] = qf_target_next_outs[QF_PREDS]
+        # We are learning a Q-value distribution.
+        if self.num_atoms > 1:
+            # Add distribution artefacts to the output.
+            # Distribution support.
+            output[ATOMS] = qf_target_next_outs[ATOMS]
+            # Original logits from the Q-head.
+            output[QF_LOGITS] = qf_outs[QF_LOGITS]
+            # Probabilities of the Q-value distribution of the current state.
+            output[QF_PROBS] = qf_outs[QF_PROBS]
+            # Probabilities of the target Q-value distribution of the next state.
+            output[QF_TARGET_NEXT_PROBS] = qf_target_next_outs[QF_PROBS]
+
+        # Add the states to the output, if the module is stateful.
+        if Columns.STATE_OUT in qf_outs:
+            output[Columns.STATE_OUT] = qf_outs[Columns.STATE_OUT]
+        # For correctness, also add the output states from the target forward pass.
+        # Note, we do not backpropagate through this state.
+        if Columns.STATE_OUT in qf_target_next_outs:
+            output[Columns.NEXT_STATE_OUT] = qf_target_next_outs[Columns.STATE_OUT]
+
+        return output
+
+    @override(QNetAPI)
+    def compute_advantage_distribution(
+        self,
+        batch: Dict[str, TensorType],
+    ) -> Dict[str, TensorType]:
+        output = {}
+        # Distributional Q-learning uses a discrete support `z`
+        # to represent the action value distribution.
+        # TODO (simon): Check, if we still need here the device for torch.
+        z = torch.arange(0.0, self.num_atoms, dtype=torch.float32).to(
+            batch.device,
+        )
+        # Rescale the support.
+        z = self.v_min + z * (self.v_max - self.v_min) / float(self.num_atoms - 1)
+        # Reshape the action values.
+        # NOTE: Handcrafted action shape.
+        logits_per_action_per_atom = torch.reshape(
+            batch, shape=(*batch.shape[:-1], self.action_space.n, self.num_atoms)
+        )
+        # Calculate the probability for each action value atom. Note,
+        # the sum along action value atoms of a single action value
+        # must sum to one.
+        prob_per_action_per_atom = nn.functional.softmax(
+            logits_per_action_per_atom,
+            dim=-1,
+        )
+        # Compute expected action value by weighted sum.
+        output[ATOMS] = z
+        output["logits"] = logits_per_action_per_atom
+        output["probs"] = prob_per_action_per_atom
+
+        return output
+
+    # TODO (simon): Test, if providing the function with a `return_probs`
+    #  improves performance significantly.
+    @override(DefaultDQNRLModule)
+    def _qf_forward_helper(
+        self,
+        batch: Dict[str, TensorType],
+        encoder: Encoder,
+        head: Union[Model, Dict[str, Model]],
+    ) -> Dict[str, TensorType]:
+        """Computes Q-values.
+
+        This is a helper function that takes care of all different cases,
+        i.e. if we use a dueling architecture or not and if we use distributional
+        Q-learning or not.
+
+        Args:
+            batch: The batch received in the forward pass.
+            encoder: The encoder network to use. Here we have a single encoder
+                for all heads (Q or advantages and value in case of a dueling
+                architecture).
+            head: Either a head model or a dictionary of head model (dueling
+                architecture) containing advantage and value stream heads.
+
+        Returns:
+            In case of expectation learning the Q-value predictions ("qf_preds")
+            and in case of distributional Q-learning in addition to the predictions
+            the atoms ("atoms"), the Q-value predictions ("qf_preds"), the Q-logits
+            ("qf_logits") and the probabilities for the support atoms ("qf_probs").
+        """
+        output = {}
+
+        # Encoder forward pass.
+        encoder_outs = encoder(batch)
+
+        # Do we have a dueling architecture.
+        if self.uses_dueling:
+            # Head forward passes for advantage and value stream.
+            qf_outs = head["af"](encoder_outs[ENCODER_OUT])
+            vf_outs = head["vf"](encoder_outs[ENCODER_OUT])
+            # We learn a Q-value distribution.
+            if self.num_atoms > 1:
+                # Compute the advantage stream distribution.
+                af_dist_output = self.compute_advantage_distribution(qf_outs)
+                # Center the advantage stream distribution.
+                centered_af_logits = af_dist_output["logits"] - af_dist_output[
+                    "logits"
+                ].mean(dim=-1, keepdim=True)
+                # Calculate the Q-value distribution by adding advantage and
+                # value stream.
+                qf_logits = centered_af_logits + vf_outs.view(
+                    -1, *((1,) * (centered_af_logits.dim() - 1))
+                )
+                # Calculate probabilites for the Q-value distribution along
+                # the support given by the atoms.
+                qf_probs = nn.functional.softmax(qf_logits, dim=-1)
+                # Return also the support as we need it in the learner.
+                output[ATOMS] = af_dist_output[ATOMS]
+                # Calculate the Q-values by the weighted sum over the atoms.
+                output[QF_PREDS] = torch.sum(af_dist_output[ATOMS] * qf_probs, dim=-1)
+                output[QF_LOGITS] = qf_logits
+                output[QF_PROBS] = qf_probs
+            # Otherwise we learn an expectation.
+            else:
+                # Center advantages. Note, we cannot do an in-place operation here
+                # b/c we backpropagate through these values. See for a discussion
+                # https://discuss.pytorch.org/t/gradient-computation-issue-due-to-
+                # inplace-operation-unsure-how-to-debug-for-custom-model/170133
+                # Has to be a mean for each batch element.
+                af_outs_mean = torch.nan_to_num(qf_outs, neginf=torch.nan).nanmean(
+                    dim=-1, keepdim=True
+                )
+                qf_outs = qf_outs - af_outs_mean
+                # Add advantage and value stream. Note, we broadcast here.
+                output[QF_PREDS] = qf_outs + vf_outs
+        # No dueling architecture.
+        else:
+            # Note, in this case the advantage network is the Q-network.
+            # Forward pass through Q-head.
+            qf_outs = head(encoder_outs[ENCODER_OUT])
+            # We learn a Q-value distribution.
+            if self.num_atoms > 1:
+                # Note in a non-dueling architecture the advantage distribution is
+                # the Q-value distribution.
+                # Get the Q-value distribution.
+                qf_dist_outs = self.compute_advantage_distribution(qf_outs)
+                # Get the support of the Q-value distribution.
+                output[ATOMS] = qf_dist_outs[ATOMS]
+                # Calculate the Q-values by the weighted sum over the atoms.
+                output[QF_PREDS] = torch.sum(
+                    qf_dist_outs[ATOMS] * qf_dist_outs["probs"], dim=-1
+                )
+                output[QF_LOGITS] = qf_dist_outs["logits"]
+                output[QF_PROBS] = qf_dist_outs["probs"]
+            # Otherwise we learn an expectation.
+            else:
+                # In this case we have a Q-head of dimension (1, action_space.n).
+                output[QF_PREDS] = qf_outs
+
+        # If we have a stateful encoder add the output states to the return
+        # dictionary.
+        if Columns.STATE_OUT in encoder_outs:
+            output[Columns.STATE_OUT] = encoder_outs[Columns.STATE_OUT]
+
+        return output

rllib/algorithms/dqn/torch/dqn_torch_learner.py

@@ -0,0 +1,293 @@
+from typing import Dict
+
+from ray.rllib.algorithms.dqn.dqn import DQNConfig
+from ray.rllib.algorithms.dqn.dqn_learner import (
+    ATOMS,
+    QF_LOGITS,
+    QF_LOSS_KEY,
+    QF_MAX_KEY,
+    QF_MEAN_KEY,
+    QF_MIN_KEY,
+    QF_NEXT_PREDS,
+    QF_PREDS,
+    QF_PROBS,
+    QF_TARGET_NEXT_PREDS,
+    QF_TARGET_NEXT_PROBS,
+    TD_ERROR_MEAN_KEY,
+    DQNLearner,
+)
+from ray.rllib.core.columns import Columns
+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.metrics import TD_ERROR_KEY
+from ray.rllib.utils.typing import ModuleID, TensorType
+
+torch, nn = try_import_torch()
+
+
+class DQNTorchLearner(DQNLearner, TorchLearner):
+    """Implements `torch`-specific DQN Rainbow loss logic on top of `DQNLearner`
+
+    This ' Learner' class implements the loss in its
+    `self.compute_loss_for_module()` method.
+    """
+
+    @override(TorchLearner)
+    def compute_loss_for_module(
+        self,
+        *,
+        module_id: ModuleID,
+        config: DQNConfig,
+        batch: Dict,
+        fwd_out: Dict[str, TensorType]
+    ) -> TensorType:
+
+        # Possibly apply masking to some sub loss terms and to the total loss term
+        # at the end. Masking could be used for RNN-based model (zero padded `batch`)
+        # and for PPO's batched value function (and bootstrap value) computations,
+        # for which we add an (artificial) timestep to each episode to
+        # simplify the actual computation.
+        if Columns.LOSS_MASK in batch:
+            mask = batch[Columns.LOSS_MASK].clone()
+            # Check, if a burn-in should be used to recover from a poor state.
+            if self.config.burn_in_len > 0:
+                # Train only on the timesteps after the burn-in period.
+                mask[:, : self.config.burn_in_len] = False
+            num_valid = torch.sum(mask)
+
+            def possibly_masked_mean(data_):
+                return torch.sum(data_[mask]) / num_valid
+
+            def possibly_masked_min(data_):
+                # Prevent minimum over empty tensors, which can happened
+                # when all elements in the mask are `False`.
+                return (
+                    torch.tensor(float("nan"))
+                    if data_[mask].numel() == 0
+                    else torch.min(data_[mask])
+                )
+
+            def possibly_masked_max(data_):
+                # Prevent maximum over empty tensors, which can happened
+                # when all elements in the mask are `False`.
+                return (
+                    torch.tensor(float("nan"))
+                    if data_[mask].numel() == 0
+                    else torch.max(data_[mask])
+                )
+
+        else:
+            possibly_masked_mean = torch.mean
+            possibly_masked_min = torch.min
+            possibly_masked_max = torch.max
+
+        q_curr = fwd_out[QF_PREDS]
+        q_target_next = fwd_out[QF_TARGET_NEXT_PREDS]
+
+        # Get the Q-values for the selected actions in the rollout.
+        # TODO (simon, sven): Check, if we can use `gather` with a complex action
+        # space - we might need the one_hot_selection. Also test performance.
+        q_selected = torch.nan_to_num(
+            torch.gather(
+                q_curr,
+                dim=-1,
+                index=batch[Columns.ACTIONS]
+                .view(*batch[Columns.ACTIONS].shape, 1)
+                .long(),
+            ),
+            neginf=0.0,
+        ).squeeze(dim=-1)
+
+        # Use double Q learning.
+        if config.double_q:
+            # Then we evaluate the target Q-function at the best action (greedy action)
+            # over the online Q-function.
+            # Mark the best online Q-value of the next state.
+            q_next_best_idx = (
+                torch.argmax(fwd_out[QF_NEXT_PREDS], dim=-1).unsqueeze(dim=-1).long()
+            )
+            # Get the Q-value of the target network at maximum of the online network
+            # (bootstrap action).
+            q_next_best = torch.nan_to_num(
+                torch.gather(q_target_next, dim=-1, index=q_next_best_idx),
+                neginf=0.0,
+            ).squeeze()
+        else:
+            # Mark the maximum Q-value(s).
+            q_next_best_idx = (
+                torch.argmax(q_target_next, dim=-1).unsqueeze(dim=-1).long()
+            )
+            # Get the maximum Q-value(s).
+            q_next_best = torch.nan_to_num(
+                torch.gather(q_target_next, dim=-1, index=q_next_best_idx),
+                neginf=0.0,
+            ).squeeze()
+
+        # If we learn a Q-distribution.
+        if config.num_atoms > 1:
+            # Extract the Q-logits evaluated at the selected actions.
+            # (Note, `torch.gather` should be faster than multiplication
+            # with a one-hot tensor.)
+            # (32, 2, 10) -> (32, 10)
+            q_logits_selected = torch.gather(
+                fwd_out[QF_LOGITS],
+                dim=1,
+                # Note, the Q-logits are of shape (B, action_space.n, num_atoms)
+                # while the actions have shape (B, 1). We reshape actions to
+                # (B, 1, num_atoms).
+                index=batch[Columns.ACTIONS]
+                .view(-1, 1, 1)
+                .expand(-1, 1, config.num_atoms)
+                .long(),
+            ).squeeze(dim=1)
+            # Get the probabilies for the maximum Q-value(s).
+            q_probs_next_best = torch.gather(
+                fwd_out[QF_TARGET_NEXT_PROBS],
+                dim=1,
+                # Change the view and then expand to get to the dimensions
+                # of the probabilities (dims 0 and 2, 1 should be reduced
+                # from 2 -> 1).
+                index=q_next_best_idx.view(-1, 1, 1).expand(-1, 1, config.num_atoms),
+            ).squeeze(dim=1)
+
+            # For distributional Q-learning we use an entropy loss.
+
+            # Extract the support grid for the Q distribution.
+            z = fwd_out[ATOMS]
+            # TODO (simon): Enable computing on GPU.
+            # (batch_size, 1) * (1, num_atoms) = (batch_size, num_atoms)s
+            r_tau = torch.clamp(
+                batch[Columns.REWARDS].unsqueeze(dim=-1)
+                + (
+                    config.gamma ** batch["n_step"]
+                    * (1.0 - batch[Columns.TERMINATEDS].float())
+                ).unsqueeze(dim=-1)
+                * z,
+                config.v_min,
+                config.v_max,
+            ).squeeze(dim=1)
+            # (32, 10)
+            b = (r_tau - config.v_min) / (
+                (config.v_max - config.v_min) / float(config.num_atoms - 1.0)
+            )
+            lower_bound = torch.floor(b)
+            upper_bound = torch.ceil(b)
+
+            floor_equal_ceil = ((upper_bound - lower_bound) < 0.5).float()
+
+            # (B, num_atoms, num_atoms).
+            lower_projection = nn.functional.one_hot(
+                lower_bound.long(), config.num_atoms
+            )
+            upper_projection = nn.functional.one_hot(
+                upper_bound.long(), config.num_atoms
+            )
+            # (32, 10)
+            ml_delta = q_probs_next_best * (upper_bound - b + floor_equal_ceil)
+            mu_delta = q_probs_next_best * (b - lower_bound)
+            # (32, 10)
+            ml_delta = torch.sum(lower_projection * ml_delta.unsqueeze(dim=-1), dim=1)
+            mu_delta = torch.sum(upper_projection * mu_delta.unsqueeze(dim=-1), dim=1)
+            # We do not want to propagate through the distributional targets.
+            # (32, 10)
+            m = (ml_delta + mu_delta).detach()
+
+            # The Rainbow paper claims to use the KL-divergence loss. This is identical
+            # to using the cross-entropy (differs only by entropy which is constant)
+            # when optimizing by the gradient (the gradient is identical).
+            td_error = nn.CrossEntropyLoss(reduction="none")(q_logits_selected, m)
+            # Compute the weighted loss (importance sampling weights).
+            total_loss = torch.mean(batch["weights"] * td_error)
+        else:
+            # Masked all Q-values with terminated next states in the targets.
+            q_next_best_masked = (
+                1.0 - batch[Columns.TERMINATEDS].float()
+            ) * q_next_best
+
+            # Compute the RHS of the Bellman equation.
+            # Detach this node from the computation graph as we do not want to
+            # backpropagate through the target network when optimizing the Q loss.
+            q_selected_target = (
+                batch[Columns.REWARDS]
+                + (config.gamma ** batch["n_step"]) * q_next_best_masked
+            ).detach()
+
+            # Choose the requested loss function. Note, in case of the Huber loss
+            # we fall back to the default of `delta=1.0`.
+            loss_fn = nn.HuberLoss if config.td_error_loss_fn == "huber" else nn.MSELoss
+            # Compute the TD error.
+            td_error = torch.abs(q_selected - q_selected_target)
+            # Compute the weighted loss (importance sampling weights).
+            total_loss = possibly_masked_mean(
+                batch["weights"]
+                * loss_fn(reduction="none")(q_selected, q_selected_target)
+            )
+
+        # Log the TD-error with reduce="item_series", such that - in case we have n parallel
+        # Learners - we will re-concatenate the produced TD-error tensors to yield
+        # a 1:1 representation of the original batch.
+        self.metrics.log_value(
+            key=(module_id, TD_ERROR_KEY),
+            value=td_error,
+            reduce="item_series",
+        )
+        # Log other important loss stats (reduce=mean (default), but with window=1
+        # in order to keep them history free).
+        self.metrics.log_dict(
+            {
+                QF_LOSS_KEY: total_loss,
+                QF_MEAN_KEY: possibly_masked_mean(q_selected),
+                QF_MAX_KEY: possibly_masked_max(q_selected),
+                QF_MIN_KEY: possibly_masked_min(q_selected),
+                TD_ERROR_MEAN_KEY: possibly_masked_mean(td_error),
+            },
+            key=module_id,
+            window=1,  # <- single items (should not be mean/ema-reduced over time).
+        )
+        # If we learn a Q-value distribution log the support and average
+        # probabilities.
+        if config.num_atoms > 1:
+            # Log important loss stats.
+            self.metrics.log_dict(
+                {
+                    ATOMS: torch.mean(z),
+                    # The absolute difference in expectation between the actions
+                    # should (at least mildly) rise.
+                    "expectations_abs_diff": torch.mean(
+                        torch.abs(
+                            torch.diff(
+                                torch.sum(fwd_out[QF_PROBS].mean(dim=0) * z, dim=1)
+                            ).mean(dim=0)
+                        )
+                    ),
+                    # The total variation distance should measure the distance between
+                    # return distributions of different actions. This should (at least
+                    # mildly) increase during training when the agent differentiates
+                    # more between actions.
+                    "dist_total_variation_dist": torch.diff(
+                        fwd_out[QF_PROBS].mean(dim=0), dim=0
+                    )
+                    .abs()
+                    .sum()
+                    * 0.5,
+                    # The maximum distance between the action distributions. This metric
+                    # should increase over the course of training.
+                    "dist_max_abs_distance": torch.max(
+                        torch.diff(fwd_out[QF_PROBS].mean(dim=0), dim=0).abs()
+                    ),
+                    # Mean shannon entropy of action distributions. This should decrease
+                    # over the course of training.
+                    "action_dist_mean_entropy": torch.mean(
+                        (
+                            fwd_out[QF_PROBS].mean(dim=0)
+                            * torch.log(fwd_out[QF_PROBS].mean(dim=0))
+                        ).sum(dim=1),
+                        dim=0,
+                    ),
+                },
+                key=module_id,
+                window=1,  # <- single items (should not be mean/ema-reduced over time).
+            )
+
+        return total_loss

rllib/algorithms/dreamerv3/README.md

@@ -0,0 +1,200 @@
+# DreamerV3
+
+![DreamerV3](../../../doc/source/rllib/images/dreamerv3/dreamerv3.png)
+
+## Overview
+An RLlib-based implementation of the
+[DreamerV3 model-based reinforcement learning algorithm](https://arxiv.org/pdf/2301.04104v1.pdf)
+by D. Hafner et al. (Google DeepMind) 2023, in PyTorch.
+
+This implementation allows scaling up training by using multi-GPU machines for
+neural network updates (see below for tips and tricks, example configs, and command lines).
+
+DreamerV3 trains a world model in supervised fashion using real environment
+interactions. The world model's objective is to correctly predict all aspects
+of the transition dynamics of the RL environment, which includes predicting the
+correct next environment state, received rewards, as well as a boolean episode
+continuation flag.
+Just like in a standard policy gradient algorithm (e.g. REINFORCE), the critic tries to
+predict a correct value function and the actor tries to come up with good actions
+choices that maximize accumulated rewards over time.
+However, both actor and critic are never trained on real environment data, but solely on
+dreamed trajectories produced by the world model.
+
+For more specific details about DreamerV3 architecture and math refer to the
+[original paper](https://arxiv.org/pdf/2301.04104v1.pdf) (see below for all references).
+
+## Note on Hyperparameter Tuning for DreamerV3
+DreamerV3 is an extremely versatile and stable algorithm that not only works well on
+different action- and observation spaces (i.e. discrete and continuous actions, as well
+as image and vector observations) and reward functions (sparse or dense),
+but also has very little hyperparameters that require tuning.
+
+All you need is a simple "model size" setting (from "XS" to "XL") and a value for the training ratio, which
+specifies how many steps to replay from the buffer for a training update vs how many
+steps to take in the actual environment.
+
+Here are some examples on how to set these config settings within your `DreamerV3Config` objects:
+
+## Example Configs and Command Lines
+
+<b>Note:</b> For a quick setup guide on how to get started with RLlib, refer to this
+[documentation page here](https://docs.ray.io/en/latest/rllib/index.html#rllib-in-60-seconds).
+
+Use the config examples and templates in the
+[examples folder](../../examples/algorithms/dreamerv3)
+in combination with the following scripts and command lines in order to run RLlib's DreamerV3 algorithm in your experiments:
+
+### [Atari100k](../../examples/algorithms/dreamerv3/atari_100k_dreamerv3.py)
+```shell
+$ cd ray/rllib/examples/algorithms/dreamerv3/
+$ python atari_100k_dreamerv3.py --env ale_py:ALE/Pong-v5
+```
+
+### [DeepMind Control Suite (vision)](../../examples/algorithms/dreamerv3/dm_control_suite_vision_dreamerv3.py)
+```shell
+$ cd ray/rllib/examples/algorithms/dreamerv3/
+$ python dm_control_suite_vision_dreamerv3.py --env DMC/cartpole/swingup
+```
+Other `--env` options for the DM Control Suite would be `--env DMC/hopper/hop`, `--env DMC/walker/walk`, etc..
+Note that you can also switch on WandB logging with the above script via the options
+`--wandb-key=[your WandB API key] --wandb-project=[some project name] --wandb-run-name=[some run name]`
+
+## Running DreamerV3 with arbitrary Envs and Configs
+Can I run DreamerV3 with any gym or custom environments? Yes, you can!
+
+<img src="../../../doc/source/rllib/images/dreamerv3/flappy_bird_env.png" alt="Flappy Bird gymnasium env" width="300" height="300" />
+
+Let's try the Flappy Bird gymnasium env. It's image space is a cellphone-style
+288 x 512 RGB, very different from DreamerV3's Atari benchmark norm (which is 64x64 RGB).
+So we will have to custom-wrap observations to resize/normalize FlappyBird's ``Box(0, 255, (288, 512, 3), f32)``
+space into a new ``Box(-1, 1, (64, 64, 3), f32)``.
+
+First we quickly install ``flappy_bird_gymnasium`` in our dev environment:
+```shell
+$ pip install flappy_bird_gymnasium
+```
+
+Now, let's create a new python file for this RLlib experiment and call it ``flappy_bird.py``:
+
+```python
+from ray import tune
+from ray.rllib.algorithms.dreamerv3.dreamerv3 import DreamerV3Config
+
+
+def _env_creator(ctx):
+    import flappy_bird_gymnasium  # doctest: +SKIP
+    import gymnasium as gym
+    from supersuit.generic_wrappers import resize_v1
+    from ray.rllib.env.wrappers.atari_wrappers import NormalizedImageEnv
+
+    return NormalizedImageEnv(
+        resize_v1(  # resize to 64x64 and normalize images
+            gym.make("FlappyBird-rgb-v0", audio_on=False), x_size=64, y_size=64
+        )
+    )
+
+
+# Register the FlappyBird-rgb-v0 env including necessary wrappers via the
+# `tune.register_env()` API.
+tune.register_env("flappy-bird", _env_creator)
+
+# Define the `config` variable to use for training.
+config = (
+    DreamerV3Config()
+    # set the env to the pre-registered string
+    .environment("flappy-bird")
+    # play around with the insanely high number of hyperparameters for DreamerV3 ;)
+    .training(
+        model_size="S",
+        training_ratio=1024,
+    )
+)
+
+# Run the tuner job.
+results = tune.Tuner(trainable="DreamerV3", param_space=config).fit()
+```
+
+Great! Now, let's run this experiment:
+
+```shell
+$ python flappy_bird.py
+```
+
+This should be it. Feel free to try out running this on multiple GPUs using these
+more advanced config examples [here (Atari100k)](../../examples/algorithms/dreamerv3/atari_100k_dreamerv3.py) and
+[here (DM Control Suite)](../../examples/algorithms/dreamerv3/dm_control_suite_vision_dreamerv3.py).
+Also see the notes below on good recipes for running on multiple GPUs.
+
+<b>IMPORTANT:</b> DreamerV3 out-of-the-box only supports image observation spaces of
+shape 64x64x3 as well as any vector observations (1D float32 Box spaces).
+Should you require a special world model encoder- and decoder for other observation
+spaces (e.g. a text embedding or images of other dimensions), you will have to
+subclass [DreamerV3's catalog class](dreamerv3_catalog.py) and then configure this
+new catalog via your ``DreamerV3Config`` object as follows:
+
+```python
+from ray.rllib.algorithms.dreamerv3.torch.dreamerv3_torch_rl_module import DreamerV3TorchRLModule
+from ray.rllib.core.rl_module.rl_module import RLModuleSpec
+
+config.rl_module(
+    rl_module_spec=RLModuleSpec(
+        module_class=DreamerV3TorchRLModule,
+        catalog_class=[your DreamerV3Catalog subclass],
+    )
+)
+```
+
+
+## Note on multi-GPU Training with DreamerV3
+We found that when using multiple GPUs for DreamerV3 training, the following simple
+adjustments should be made on top of the default config.
+
+- Multiply the batch size (default `B=16`) by the number of GPUs you are using.
+  Use the `DreamerV3Config.training(batch_size_B=..)` API for this. For example, for 2 GPUs,
+  use a batch size of `B=32`.
+- Multiply the number of environments you sample from in parallel by the number of GPUs you are using.
+  Use the `DreamerV3Config.env_runners(num_envs_per_env_runner=..)` for this.
+  For example, for 4 GPUs and a default environment count of 8 (the single-GPU default for
+  this setting depends on the benchmark you are running), use 32
+  parallel environments instead.
+- Roughly use learning rates that are the default values multiplied by the square root of the number of GPUs.
+  For example, when using 4 GPUs, multiply all default learning rates (for world model, critic, and actor) by 2.
+  - Additionally, a "priming"-style warmup schedule might help. Thereby, increase the learning rates from 0.0
+    to the final value(s) over the first ~10% of total env steps needed for the experiment.
+  - For examples on how to set such schedules within your `DreamerV3Config`, see below.
+  - [See here](https://aws.amazon.com/blogs/machine-learning/the-importance-of-hyperparameter-tuning-for-scaling-deep-learning-training-to-multiple-gpus/) for more details on learning rate "priming".
+
+
+## Results
+Our results on the Atari 100k and (visual) DeepMind Control Suite benchmarks match those
+reported in the paper.
+
+### Pong-v5 (100k) 1GPU vs 2GPUs vs 4GPUs
+<img src="../../../doc/source/rllib/images/dreamerv3/pong_1_2_and_4gpus.svg" style="display:block;">
+
+### Atari 100k
+<img src="../../../doc/source/rllib/images/dreamerv3/atari100k_1_vs_4gpus.svg" style="display:block;">
+
+### DeepMind Control Suite (vision)
+<img src="../../../doc/source/rllib/images/dreamerv3/dmc_1_vs_4gpus.svg" style="display:block;">
+
+
+## Running Action Inference after Training
+
+To run action inference on a DreamerV3 Algorithm object, you can use
+[this simple environment loop script](https://github.com/ray-project/ray/tree/master/doc/source/rllib/doc_code/dreamerv3_inference.py).
+
+Note the slight complexity caused by the fact that DreamerV3 a) uses a recurrent model,
+b) uses the new RLModule-based API stack (no Policy class), and c) outputs actions in a one-hot
+fashion for discrete action spaces.
+
+
+## References
+For more algorithm details, see the original Dreamer-V3 paper:
+
+[1] [Mastering Diverse Domains through World Models - 2023 D. Hafner, J. Pasukonis, J. Ba, T. Lillicrap](https://arxiv.org/pdf/2301.04104v1.pdf)
+
+.. and the (predecessor) Dreamer-V2 paper:
+
+[2] [Mastering Atari with Discrete World Models - 2021 D. Hafner, T. Lillicrap, M. Norouzi, J. Ba](https://arxiv.org/pdf/2010.02193.pdf)

rllib/algorithms/dreamerv3/__init__.py

@@ -0,0 +1,15 @@
+"""
+[1] Mastering Diverse Domains through World Models - 2023
+D. Hafner, J. Pasukonis, J. Ba, T. Lillicrap
+https://arxiv.org/pdf/2301.04104v1.pdf
+
+[2] Mastering Atari with Discrete World Models - 2021
+D. Hafner, T. Lillicrap, M. Norouzi, J. Ba
+https://arxiv.org/pdf/2010.02193.pdf
+"""
+from ray.rllib.algorithms.dreamerv3.dreamerv3 import DreamerV3, DreamerV3Config
+
+__all__ = [
+    "DreamerV3",
+    "DreamerV3Config",
+]

rllib/algorithms/dreamerv3/dreamerv3.py

@@ -0,0 +1,732 @@
+"""
+[1] Mastering Diverse Domains through World Models - 2023
+D. Hafner, J. Pasukonis, J. Ba, T. Lillicrap
+https://arxiv.org/pdf/2301.04104v1.pdf
+
+[2] Mastering Atari with Discrete World Models - 2021
+D. Hafner, T. Lillicrap, M. Norouzi, J. Ba
+https://arxiv.org/pdf/2010.02193.pdf
+"""
+
+import logging
+from typing import Any, Dict, Optional, Union
+
+import gymnasium as gym
+from typing_extensions import Self
+
+from ray.rllib.algorithms.algorithm import Algorithm
+from ray.rllib.algorithms.algorithm_config import AlgorithmConfig, NotProvided
+from ray.rllib.algorithms.dreamerv3.dreamerv3_catalog import DreamerV3Catalog
+from ray.rllib.algorithms.dreamerv3.utils import do_symlog_obs
+from ray.rllib.algorithms.dreamerv3.utils.add_is_firsts_to_batch import (
+    AddIsFirstsToBatch,
+)
+from ray.rllib.algorithms.dreamerv3.utils.summaries import (
+    report_dreamed_eval_trajectory_vs_samples,
+    report_predicted_vs_sampled_obs,
+    report_sampling_and_replay_buffer,
+)
+from ray.rllib.connectors.common import AddStatesFromEpisodesToBatch
+from ray.rllib.core import DEFAULT_MODULE_ID
+from ray.rllib.core.columns import Columns
+from ray.rllib.core.rl_module.rl_module import RLModuleSpec
+from ray.rllib.env import INPUT_ENV_SINGLE_SPACES
+from ray.rllib.execution.rollout_ops import synchronous_parallel_sample
+from ray.rllib.policy.sample_batch import SampleBatch
+from ray.rllib.utils import deep_update
+from ray.rllib.utils.annotations import PublicAPI, override
+from ray.rllib.utils.metrics import (
+    ENV_RUNNER_RESULTS,
+    LEARN_ON_BATCH_TIMER,
+    LEARNER_RESULTS,
+    NUM_ENV_STEPS_SAMPLED_LIFETIME,
+    NUM_ENV_STEPS_TRAINED_LIFETIME,
+    NUM_GRAD_UPDATES_LIFETIME,
+    NUM_SYNCH_WORKER_WEIGHTS,
+    REPLAY_BUFFER_RESULTS,
+    SAMPLE_TIMER,
+    SYNCH_WORKER_WEIGHTS_TIMER,
+    TIMERS,
+)
+from ray.rllib.utils.numpy import one_hot
+from ray.rllib.utils.replay_buffers.episode_replay_buffer import EpisodeReplayBuffer
+from ray.rllib.utils.typing import LearningRateOrSchedule
+
+logger = logging.getLogger(__name__)
+
+
+class DreamerV3Config(AlgorithmConfig):
+    """Defines a configuration class from which a DreamerV3 can be built.
+
+    .. testcode::
+
+        from ray.rllib.algorithms.dreamerv3 import DreamerV3Config
+        config = (
+            DreamerV3Config()
+            .environment("CartPole-v1")
+            .training(
+                model_size="XS",
+                training_ratio=1,
+                # TODO
+                model={
+                    "batch_size_B": 1,
+                    "batch_length_T": 1,
+                    "horizon_H": 1,
+                    "gamma": 0.997,
+                    "model_size": "XS",
+                },
+            )
+        )
+
+        config = config.learners(num_learners=0)
+        # Build a Algorithm object from the config and run 1 training iteration.
+        algo = config.build()
+        # algo.train()
+        del algo
+    """
+
+    def __init__(self, algo_class=None):
+        """Initializes a DreamerV3Config instance."""
+        super().__init__(algo_class=algo_class or DreamerV3)
+
+        # fmt: off
+        # __sphinx_doc_begin__
+
+        # DreamerV3 specific settings:
+        self.model_size = "XS"
+        self.training_ratio = 1024
+
+        self.replay_buffer_config = {
+            "type": "EpisodeReplayBuffer",
+            "capacity": int(1e6),
+        }
+        self.world_model_lr = 1e-4
+        self.actor_lr = 3e-5
+        self.critic_lr = 3e-5
+        self.batch_size_B = 16
+        self.batch_length_T = 64
+        self.horizon_H = 15
+        self.gae_lambda = 0.95  # [1] eq. 7.
+        self.entropy_scale = 3e-4  # [1] eq. 11.
+        self.return_normalization_decay = 0.99  # [1] eq. 11 and 12.
+        self.train_critic = True
+        self.train_actor = True
+        self.intrinsic_rewards_scale = 0.1
+        self.world_model_grad_clip_by_global_norm = 1000.0
+        self.critic_grad_clip_by_global_norm = 100.0
+        self.actor_grad_clip_by_global_norm = 100.0
+        self.symlog_obs = "auto"
+        self.use_float16 = False
+        self.use_curiosity = False
+
+        # Reporting.
+        # DreamerV3 is super sample efficient and only needs very few episodes
+        # (normally) to learn. Leaving this at its default value would gravely
+        # underestimate the learning performance over the course of an experiment.
+        self.metrics_num_episodes_for_smoothing = 1
+        self.report_individual_batch_item_stats = False
+        self.report_dream_data = False
+        self.report_images_and_videos = False
+
+        # Override some of AlgorithmConfig's default values with DreamerV3-specific
+        # values.
+        self.lr = None
+        self.gamma = 0.997  # [1] eq. 7.
+        # Do not use! Set `batch_size_B` and `batch_length_T` instead.
+        self.train_batch_size = None
+        self.num_env_runners = 0
+        self.rollout_fragment_length = 1
+        # Dreamer only runs on the new API stack.
+        self.enable_rl_module_and_learner = True
+        self.enable_env_runner_and_connector_v2 = True
+        # TODO (sven): DreamerV3 still uses its own EnvRunner class. This env-runner
+        #  does not use connectors. We therefore should not attempt to merge/broadcast
+        #  the connector states between EnvRunners (if >0). Note that this is only
+        #  relevant if num_env_runners > 0, which is normally not the case when using
+        #  this algo.
+        self.use_worker_filter_stats = False
+        # __sphinx_doc_end__
+        # fmt: on
+
+    @override(AlgorithmConfig)
+    def build_env_to_module_connector(self, env, spaces, device):
+        connector = super().build_env_to_module_connector(env, spaces, device)
+
+        # Prepend the "is_first" connector such that the RSSM knows, when to insert
+        # its (learned) internal state into the batch.
+        # We have to do this before the `AddStatesFromEpisodesToBatch` piece
+        # such that the column is properly batched/time-ranked.
+        if self.add_default_connectors_to_learner_pipeline:
+            connector.insert_before(
+                AddStatesFromEpisodesToBatch,
+                AddIsFirstsToBatch(),
+            )
+        return connector
+
+    @property
+    def batch_size_B_per_learner(self):
+        """Returns the batch_size_B per Learner worker.
+
+        Needed by some of the DreamerV3 loss math."""
+        return self.batch_size_B // (self.num_learners or 1)
+
+    @override(AlgorithmConfig)
+    def training(
+        self,
+        *,
+        model_size: Optional[str] = NotProvided,
+        training_ratio: Optional[float] = NotProvided,
+        batch_size_B: Optional[int] = NotProvided,
+        batch_length_T: Optional[int] = NotProvided,
+        horizon_H: Optional[int] = NotProvided,
+        gae_lambda: Optional[float] = NotProvided,
+        entropy_scale: Optional[float] = NotProvided,
+        return_normalization_decay: Optional[float] = NotProvided,
+        train_critic: Optional[bool] = NotProvided,
+        train_actor: Optional[bool] = NotProvided,
+        intrinsic_rewards_scale: Optional[float] = NotProvided,
+        world_model_lr: Optional[LearningRateOrSchedule] = NotProvided,
+        actor_lr: Optional[LearningRateOrSchedule] = NotProvided,
+        critic_lr: Optional[LearningRateOrSchedule] = NotProvided,
+        world_model_grad_clip_by_global_norm: Optional[float] = NotProvided,
+        critic_grad_clip_by_global_norm: Optional[float] = NotProvided,
+        actor_grad_clip_by_global_norm: Optional[float] = NotProvided,
+        symlog_obs: Optional[Union[bool, str]] = NotProvided,
+        use_float16: Optional[bool] = NotProvided,
+        replay_buffer_config: Optional[dict] = NotProvided,
+        use_curiosity: Optional[bool] = NotProvided,
+        **kwargs,
+    ) -> Self:
+        """Sets the training related configuration.
+
+        Args:
+            model_size: The main switch for adjusting the overall model size. See [1]
+                (table B) for more information on the effects of this setting on the
+                model architecture.
+                Supported values are "XS", "S", "M", "L", "XL" (as per the paper), as
+                well as, "nano", "micro", "mini", and "XXS" (for RLlib's
+                implementation). See ray.rllib.algorithms.dreamerv3.utils.
+                __init__.py for the details on what exactly each size does to the layer
+                sizes, number of layers, etc..
+            training_ratio: The ratio of total steps trained (sum of the sizes of all
+                batches ever sampled from the replay buffer) over the total env steps
+                taken (in the actual environment, not the dreamed one). For example,
+                if the training_ratio is 1024 and the batch size is 1024, we would take
+                1 env step for every training update: 1024 / 1. If the training ratio
+                is 512 and the batch size is 1024, we would take 2 env steps and then
+                perform a single training update (on a 1024 batch): 1024 / 2.
+            batch_size_B: The batch size (B) interpreted as number of rows (each of
+                length `batch_length_T`) to sample from the replay buffer in each
+                iteration.
+            batch_length_T: The batch length (T) interpreted as the length of each row
+                sampled from the replay buffer in each iteration. Note that
+                `batch_size_B` rows will be sampled in each iteration. Rows normally
+                contain consecutive data (consecutive timesteps from the same episode),
+                but there might be episode boundaries in a row as well.
+            horizon_H: The horizon (in timesteps) used to create dreamed data from the
+                world model, which in turn is used to train/update both actor- and
+                critic networks.
+            gae_lambda: The lambda parameter used for computing the GAE-style
+                value targets for the actor- and critic losses.
+            entropy_scale: The factor with which to multiply the entropy loss term
+                inside the actor loss.
+            return_normalization_decay: The decay value to use when computing the
+                running EMA values for return normalization (used in the actor loss).
+            train_critic: Whether to train the critic network. If False, `train_actor`
+                must also be False (cannot train actor w/o training the critic).
+            train_actor: Whether to train the actor network. If True, `train_critic`
+                must also be True (cannot train actor w/o training the critic).
+            intrinsic_rewards_scale: The factor to multiply intrinsic rewards with
+                before adding them to the extrinsic (environment) rewards.
+            world_model_lr: The learning rate or schedule for the world model optimizer.
+            actor_lr: The learning rate or schedule for the actor optimizer.
+            critic_lr: The learning rate or schedule for the critic optimizer.
+            world_model_grad_clip_by_global_norm: World model grad clipping value
+                (by global norm).
+            critic_grad_clip_by_global_norm: Critic grad clipping value
+                (by global norm).
+            actor_grad_clip_by_global_norm: Actor grad clipping value (by global norm).
+            symlog_obs: Whether to symlog observations or not. If set to "auto"
+                (default), will check for the environment's observation space and then
+                only symlog if not an image space.
+            use_float16: Whether to train with mixed float16 precision. In this mode,
+                model parameters are stored as float32, but all computations are
+                performed in float16 space (except for losses and distribution params
+                and outputs).
+            replay_buffer_config: Replay buffer config.
+                Only serves in DreamerV3 to set the capacity of the replay buffer.
+                Note though that in the paper ([1]) a size of 1M is used for all
+                benchmarks and there doesn't seem to be a good reason to change this
+                parameter.
+                Examples:
+                {
+                "type": "EpisodeReplayBuffer",
+                "capacity": 100000,
+                }
+
+        Returns:
+            This updated AlgorithmConfig object.
+        """
+        # Not fully supported/tested yet.
+        if use_curiosity is not NotProvided:
+            raise ValueError(
+                "`DreamerV3Config.curiosity` is not fully supported and tested yet! "
+                "It thus remains disabled for now."
+            )
+
+        # Pass kwargs onto super's `training()` method.
+        super().training(**kwargs)
+
+        if model_size is not NotProvided:
+            self.model_size = model_size
+        if training_ratio is not NotProvided:
+            self.training_ratio = training_ratio
+        if batch_size_B is not NotProvided:
+            self.batch_size_B = batch_size_B
+        if batch_length_T is not NotProvided:
+            self.batch_length_T = batch_length_T
+        if horizon_H is not NotProvided:
+            self.horizon_H = horizon_H
+        if gae_lambda is not NotProvided:
+            self.gae_lambda = gae_lambda
+        if entropy_scale is not NotProvided:
+            self.entropy_scale = entropy_scale
+        if return_normalization_decay is not NotProvided:
+            self.return_normalization_decay = return_normalization_decay
+        if train_critic is not NotProvided:
+            self.train_critic = train_critic
+        if train_actor is not NotProvided:
+            self.train_actor = train_actor
+        if intrinsic_rewards_scale is not NotProvided:
+            self.intrinsic_rewards_scale = intrinsic_rewards_scale
+        if world_model_lr is not NotProvided:
+            self.world_model_lr = world_model_lr
+        if actor_lr is not NotProvided:
+            self.actor_lr = actor_lr
+        if critic_lr is not NotProvided:
+            self.critic_lr = critic_lr
+        if world_model_grad_clip_by_global_norm is not NotProvided:
+            self.world_model_grad_clip_by_global_norm = (
+                world_model_grad_clip_by_global_norm
+            )
+        if critic_grad_clip_by_global_norm is not NotProvided:
+            self.critic_grad_clip_by_global_norm = critic_grad_clip_by_global_norm
+        if actor_grad_clip_by_global_norm is not NotProvided:
+            self.actor_grad_clip_by_global_norm = actor_grad_clip_by_global_norm
+        if symlog_obs is not NotProvided:
+            self.symlog_obs = symlog_obs
+        if use_float16 is not NotProvided:
+            self.use_float16 = use_float16
+        if replay_buffer_config is not NotProvided:
+            # Override entire `replay_buffer_config` if `type` key changes.
+            # Update, if `type` key remains the same or is not specified.
+            new_replay_buffer_config = deep_update(
+                {"replay_buffer_config": self.replay_buffer_config},
+                {"replay_buffer_config": replay_buffer_config},
+                False,
+                ["replay_buffer_config"],
+                ["replay_buffer_config"],
+            )
+            self.replay_buffer_config = new_replay_buffer_config["replay_buffer_config"]
+
+        return self
+
+    @override(AlgorithmConfig)
+    def reporting(
+        self,
+        *,
+        report_individual_batch_item_stats: Optional[bool] = NotProvided,
+        report_dream_data: Optional[bool] = NotProvided,
+        report_images_and_videos: Optional[bool] = NotProvided,
+        **kwargs,
+    ):
+        """Sets the reporting related configuration.
+
+        Args:
+            report_individual_batch_item_stats: Whether to include loss and other stats
+                per individual timestep inside the training batch in the result dict
+                returned by `training_step()`. If True, besides the `CRITIC_L_total`,
+                the individual critic loss values per batch row and time axis step
+                in the train batch (CRITIC_L_total_B_T) will also be part of the
+                results.
+            report_dream_data:  Whether to include the dreamed trajectory data in the
+                result dict returned by `training_step()`. If True, however, will
+                slice each reported item in the dream data down to the shape.
+                (H, B, t=0, ...), where H is the horizon and B is the batch size. The
+                original time axis will only be represented by the first timestep
+                to not make this data too large to handle.
+            report_images_and_videos: Whether to include any image/video data in the
+                result dict returned by `training_step()`.
+            **kwargs:
+
+        Returns:
+            This updated AlgorithmConfig object.
+        """
+        super().reporting(**kwargs)
+
+        if report_individual_batch_item_stats is not NotProvided:
+            self.report_individual_batch_item_stats = report_individual_batch_item_stats
+        if report_dream_data is not NotProvided:
+            self.report_dream_data = report_dream_data
+        if report_images_and_videos is not NotProvided:
+            self.report_images_and_videos = report_images_and_videos
+
+        return self
+
+    @override(AlgorithmConfig)
+    def validate(self) -> None:
+        # Call the super class' validation method first.
+        super().validate()
+
+        # Make sure, users are not using DreamerV3 yet for multi-agent:
+        if self.is_multi_agent:
+            self._value_error("DreamerV3 does NOT support multi-agent setups yet!")
+
+        # Make sure, we are configure for the new API stack.
+        if not self.enable_rl_module_and_learner:
+            self._value_error(
+                "DreamerV3 must be run with `config.api_stack("
+                "enable_rl_module_and_learner=True)`!"
+            )
+
+        # If run on several Learners, the provided batch_size_B must be a multiple
+        # of `num_learners`.
+        if self.num_learners > 1 and (self.batch_size_B % self.num_learners != 0):
+            self._value_error(
+                f"Your `batch_size_B` ({self.batch_size_B}) must be a multiple of "
+                f"`num_learners` ({self.num_learners}) in order for "
+                "DreamerV3 to be able to split batches evenly across your Learner "
+                "processes."
+            )
+
+        # Cannot train actor w/o critic.
+        if self.train_actor and not self.train_critic:
+            self._value_error(
+                "Cannot train actor network (`train_actor=True`) w/o training critic! "
+                "Make sure you either set `train_critic=True` or `train_actor=False`."
+            )
+        # Use DreamerV3 specific batch size settings.
+        if self.train_batch_size is not None:
+            self._value_error(
+                "`train_batch_size` should NOT be set! Use `batch_size_B` and "
+                "`batch_length_T` instead."
+            )
+        # Must be run with `EpisodeReplayBuffer` type.
+        if self.replay_buffer_config.get("type") != "EpisodeReplayBuffer":
+            self._value_error(
+                "DreamerV3 must be run with the `EpisodeReplayBuffer` type! None "
+                "other supported."
+            )
+
+    @override(AlgorithmConfig)
+    def get_default_learner_class(self):
+        if self.framework_str == "torch":
+            from ray.rllib.algorithms.dreamerv3.torch.dreamerv3_torch_learner import (
+                DreamerV3TorchLearner,
+            )
+
+            return DreamerV3TorchLearner
+        else:
+            raise ValueError(f"The framework {self.framework_str} is not supported.")
+
+    @override(AlgorithmConfig)
+    def get_default_rl_module_spec(self) -> RLModuleSpec:
+        if self.framework_str == "torch":
+            from ray.rllib.algorithms.dreamerv3.torch.dreamerv3_torch_rl_module import (
+                DreamerV3TorchRLModule as module,
+            )
+
+        else:
+            raise ValueError(f"The framework {self.framework_str} is not supported.")
+
+        return RLModuleSpec(module_class=module, catalog_class=DreamerV3Catalog)
+
+    @property
+    @override(AlgorithmConfig)
+    def _model_config_auto_includes(self) -> Dict[str, Any]:
+        return super()._model_config_auto_includes | {
+            "gamma": self.gamma,
+            "horizon_H": self.horizon_H,
+            "model_size": self.model_size,
+            "symlog_obs": self.symlog_obs,
+            "use_float16": self.use_float16,
+            "batch_length_T": self.batch_length_T,
+        }
+
+
+class DreamerV3(Algorithm):
+    """Implementation of the model-based DreamerV3 RL algorithm described in [1]."""
+
+    # TODO (sven): Deprecate/do-over the Algorithm.compute_single_action() API.
+    @override(Algorithm)
+    def compute_single_action(self, *args, **kwargs):
+        raise NotImplementedError(
+            "DreamerV3 does not support the `compute_single_action()` API. Refer to the"
+            " README here (https://github.com/ray-project/ray/tree/master/rllib/"
+            "algorithms/dreamerv3) to find more information on how to run action "
+            "inference with this algorithm."
+        )
+
+    @classmethod
+    @override(Algorithm)
+    def get_default_config(cls) -> DreamerV3Config:
+        return DreamerV3Config()
+
+    @override(Algorithm)
+    def setup(self, config: AlgorithmConfig):
+        super().setup(config)
+
+        # Share RLModule between EnvRunner and single (local) Learner instance.
+        # To avoid possibly expensive weight synching step.
+        # if self.config.share_module_between_env_runner_and_learner:
+        #    assert self.env_runner.module is None
+        #    self.env_runner.module = self.learner_group._learner.module[
+        #        DEFAULT_MODULE_ID
+        #    ]
+
+        # Create a replay buffer for storing actual env samples.
+        self.replay_buffer = EpisodeReplayBuffer(
+            capacity=self.config.replay_buffer_config["capacity"],
+            batch_size_B=self.config.batch_size_B,
+            batch_length_T=self.config.batch_length_T,
+        )
+
+    @override(Algorithm)
+    def training_step(self) -> None:
+        # Push enough samples into buffer initially before we start training.
+        if self.training_iteration == 0:
+            logger.info(
+                "Filling replay buffer so it contains at least "
+                f"{self.config.batch_size_B * self.config.batch_length_T} timesteps "
+                "(required for a single train batch)."
+            )
+
+        # Have we sampled yet in this `training_step()` call?
+        have_sampled = False
+        with self.metrics.log_time((TIMERS, SAMPLE_TIMER)):
+            # Continue sampling from the actual environment (and add collected samples
+            # to our replay buffer) as long as we:
+            while (
+                # a) Don't have at least batch_size_B x batch_length_T timesteps stored
+                # in the buffer. This is the minimum needed to train.
+                self.replay_buffer.get_num_timesteps()
+                < (self.config.batch_size_B * self.config.batch_length_T)
+                # b) The computed `training_ratio` is >= the configured (desired)
+                # training ratio (meaning we should continue sampling).
+                or self.training_ratio >= self.config.training_ratio
+                # c) we have not sampled at all yet in this `training_step()` call.
+                or not have_sampled
+            ):
+                # Sample using the env runner's module.
+                episodes, env_runner_results = synchronous_parallel_sample(
+                    worker_set=self.env_runner_group,
+                    max_agent_steps=(
+                        self.config.rollout_fragment_length
+                        * self.config.num_envs_per_env_runner
+                    ),
+                    sample_timeout_s=self.config.sample_timeout_s,
+                    _uses_new_env_runners=True,
+                    _return_metrics=True,
+                )
+                self.metrics.aggregate(env_runner_results, key=ENV_RUNNER_RESULTS)
+                # Add ongoing and finished episodes into buffer. The buffer will
+                # automatically take care of properly concatenating (by episode IDs)
+                # the different chunks of the same episodes, even if they come in via
+                # separate `add()` calls.
+                self.replay_buffer.add(episodes=episodes)
+                have_sampled = True
+
+                # We took B x T env steps.
+                env_steps_last_regular_sample = sum(len(eps) for eps in episodes)
+                total_sampled = env_steps_last_regular_sample
+
+                # If we have never sampled before (just started the algo and not
+                # recovered from a checkpoint), sample B random actions first.
+                if (
+                    self.metrics.peek(
+                        (ENV_RUNNER_RESULTS, NUM_ENV_STEPS_SAMPLED_LIFETIME),
+                        default=0,
+                    )
+                    == 0
+                ):
+                    _episodes, _env_runner_results = synchronous_parallel_sample(
+                        worker_set=self.env_runner_group,
+                        max_agent_steps=(
+                            self.config.batch_size_B * self.config.batch_length_T
+                            - env_steps_last_regular_sample
+                        ),
+                        sample_timeout_s=self.config.sample_timeout_s,
+                        random_actions=True,
+                        _uses_new_env_runners=True,
+                        _return_metrics=True,
+                    )
+                    self.metrics.aggregate(_env_runner_results, key=ENV_RUNNER_RESULTS)
+                    self.replay_buffer.add(episodes=_episodes)
+                    total_sampled += sum(len(eps) for eps in _episodes)
+
+        # Summarize environment interaction and buffer data.
+        report_sampling_and_replay_buffer(
+            metrics=self.metrics, replay_buffer=self.replay_buffer
+        )
+        # Get the replay buffer metrics.
+        replay_buffer_results = self.local_replay_buffer.get_metrics()
+        self.metrics.aggregate([replay_buffer_results], key=REPLAY_BUFFER_RESULTS)
+
+        # Use self.spaces for the environment spaces of the env-runners
+        single_observation_space, single_action_space = self.spaces[
+            INPUT_ENV_SINGLE_SPACES
+        ]
+
+        # Continue sampling batch_size_B x batch_length_T sized batches from the buffer
+        # and using these to update our models (`LearnerGroup.update()`)
+        # until the computed `training_ratio` is larger than the configured one, meaning
+        # we should go back and collect more samples again from the actual environment.
+        # However, when calculating the `training_ratio` here, we use only the
+        # trained steps in this very `training_step()` call over the most recent sample
+        # amount (`env_steps_last_regular_sample`), not the global values. This is to
+        # avoid a heavy overtraining at the very beginning when we have just pre-filled
+        # the buffer with the minimum amount of samples.
+        replayed_steps_this_iter = sub_iter = 0
+        while (
+            replayed_steps_this_iter / env_steps_last_regular_sample
+        ) < self.config.training_ratio:
+            # Time individual batch updates.
+            with self.metrics.log_time((TIMERS, LEARN_ON_BATCH_TIMER)):
+                logger.info(f"\tSub-iteration {self.training_iteration}/{sub_iter})")
+
+                # Draw a new sample from the replay buffer.
+                sample = self.replay_buffer.sample(
+                    batch_size_B=self.config.batch_size_B,
+                    batch_length_T=self.config.batch_length_T,
+                )
+                replayed_steps = self.config.batch_size_B * self.config.batch_length_T
+                replayed_steps_this_iter += replayed_steps
+
+                if isinstance(single_action_space, gym.spaces.Discrete):
+                    sample["actions_ints"] = sample[Columns.ACTIONS]
+                    sample[Columns.ACTIONS] = one_hot(
+                        sample["actions_ints"],
+                        depth=single_action_space.n,
+                    )
+
+                # Perform the actual update via our learner group.
+                learner_results = self.learner_group.update(
+                    batch=SampleBatch(sample).as_multi_agent(),
+                    # TODO(sven): Maybe we should do this broadcase of global timesteps
+                    #  at the end, like for EnvRunner global env step counts. Maybe when
+                    #  we request the state from the Learners, we can - at the same
+                    #  time - send the current globally summed/reduced-timesteps.
+                    timesteps={
+                        NUM_ENV_STEPS_SAMPLED_LIFETIME: self.metrics.peek(
+                            (ENV_RUNNER_RESULTS, NUM_ENV_STEPS_SAMPLED_LIFETIME),
+                            default=0,
+                        )
+                    },
+                )
+                self.metrics.aggregate(learner_results, key=LEARNER_RESULTS)
+
+                sub_iter += 1
+                self.metrics.log_value(
+                    NUM_GRAD_UPDATES_LIFETIME, 1, reduce="lifetime_sum"
+                )
+
+        # Log videos showing how the decoder produces observation predictions
+        # from the posterior states.
+        # Only every n iterations and only for the first sampled batch row
+        # (videos are `config.batch_length_T` frames long).
+        report_predicted_vs_sampled_obs(
+            # TODO (sven): DreamerV3 is single-agent only.
+            metrics=self.metrics,
+            sample=sample,
+            batch_size_B=self.config.batch_size_B,
+            batch_length_T=self.config.batch_length_T,
+            symlog_obs=do_symlog_obs(
+                single_observation_space,
+                self.config.symlog_obs,
+            ),
+            do_report=(
+                self.config.report_images_and_videos
+                and self.training_iteration % 100 == 0
+            ),
+        )
+
+        # Log videos showing some of the dreamed trajectories and compare them with the
+        # actual trajectories from the train batch.
+        # Only every n iterations and only for the first sampled batch row AND first ts.
+        # (videos are `config.horizon_H` frames long originating from the observation
+        # at B=0 and T=0 in the train batch).
+        report_dreamed_eval_trajectory_vs_samples(
+            metrics=self.metrics,
+            sample=sample,
+            burn_in_T=0,
+            dreamed_T=self.config.horizon_H + 1,
+            dreamer_model=self.env_runner.module.dreamer_model,
+            symlog_obs=do_symlog_obs(
+                single_observation_space,
+                self.config.symlog_obs,
+            ),
+            do_report=(
+                self.config.report_dream_data and self.training_iteration % 100 == 0
+            ),
+            framework=self.config.framework_str,
+        )
+
+        # Update weights - after learning on the LearnerGroup - on all EnvRunner
+        # workers.
+        with self.metrics.log_time((TIMERS, SYNCH_WORKER_WEIGHTS_TIMER)):
+            # Only necessary if RLModule is not shared between (local) EnvRunner and
+            # (local) Learner.
+            # if not self.config.share_module_between_env_runner_and_learner:
+            self.metrics.log_value(NUM_SYNCH_WORKER_WEIGHTS, 1, reduce="sum")
+            self.env_runner_group.sync_weights(
+                from_worker_or_learner_group=self.learner_group,
+                inference_only=True,
+            )
+
+        # Add train results and the actual training ratio to stats. The latter should
+        # be close to the configured `training_ratio`.
+        self.metrics.log_value("actual_training_ratio", self.training_ratio, window=1)
+
+    @property
+    def training_ratio(self) -> float:
+        """Returns the actual training ratio of this Algorithm (not the configured one).
+
+        The training ratio is copmuted by dividing the total number of steps
+        trained thus far (replayed from the buffer) over the total number of actual
+        env steps taken thus far.
+        """
+        eps = 0.0001
+        return self.metrics.peek(NUM_ENV_STEPS_TRAINED_LIFETIME, default=0) / (
+            (
+                self.metrics.peek(
+                    (ENV_RUNNER_RESULTS, NUM_ENV_STEPS_SAMPLED_LIFETIME),
+                    default=eps,
+                )
+                or eps
+            )
+        )
+
+    # TODO (sven): Remove this once DreamerV3 is on the new SingleAgentEnvRunner.
+    @PublicAPI
+    def __setstate__(self, state) -> None:
+        """Sts the algorithm to the provided state
+
+        Args:
+            state: The state dictionary to restore this `DreamerV3` instance to.
+                `state` may have been returned by a call to an `Algorithm`'s
+                `__getstate__()` method.
+        """
+        # Call the `Algorithm`'s `__setstate__()` method.
+        super().__setstate__(state=state)
+
+        # Assign the module to the local `EnvRunner` if sharing is enabled.
+        # Note, in `Learner.restore_from_path()` the module is first deleted
+        # and then a new one is built - therefore the worker has no
+        # longer a copy of the learner.
+        if self.config.share_module_between_env_runner_and_learner:
+            assert id(self.env_runner.module) != id(
+                self.learner_group._learner.module[DEFAULT_MODULE_ID]
+            )
+            self.env_runner.module = self.learner_group._learner.module[
+                DEFAULT_MODULE_ID
+            ]

rllib/algorithms/dreamerv3/dreamerv3_catalog.py

@@ -0,0 +1,184 @@
+import gymnasium as gym
+import numpy as np
+
+from ray.rllib.algorithms.dreamerv3.utils import (
+    do_symlog_obs,
+    get_gru_units,
+    get_num_z_categoricals,
+    get_num_z_classes,
+)
+from ray.rllib.core.models.base import Encoder, Model
+from ray.rllib.core.models.catalog import Catalog
+from ray.rllib.utils import override
+
+
+class DreamerV3Catalog(Catalog):
+    """The Catalog class used to build all the models needed for DreamerV3 training."""
+
+    def __init__(
+        self,
+        observation_space: gym.Space,
+        action_space: gym.Space,
+        model_config_dict: dict,
+    ):
+        """Initializes a DreamerV3Catalog instance.
+
+        Args:
+            observation_space: The observation space of the environment.
+            action_space: The action space of the environment.
+            model_config_dict: The model config to use.
+        """
+        super().__init__(
+            observation_space=observation_space,
+            action_space=action_space,
+            model_config_dict=model_config_dict,
+        )
+
+        self.model_size = self._model_config_dict["model_size"]
+        self.is_img_space = len(self.observation_space.shape) in [2, 3]
+        self.is_gray_scale = (
+            self.is_img_space and len(self.observation_space.shape) == 2
+        )
+        # Compute the size of the vector coming out of the sequence model.
+        self.h_plus_z_flat = get_gru_units(self.model_size) + (
+            get_num_z_categoricals(self.model_size) * get_num_z_classes(self.model_size)
+        )
+
+        # TODO (sven): We should work with sub-component configurations here,
+        #  and even try replacing all current Dreamer model components with
+        #  our default primitives. But for now, we'll construct the DreamerV3Model
+        #  directly in our `build_...()` methods.
+
+    @override(Catalog)
+    def build_encoder(self, framework: str) -> Encoder:
+        """Builds the World-Model's encoder network depending on the obs space."""
+        if self.is_img_space:
+            if framework == "torch":
+                from ray.rllib.algorithms.dreamerv3.torch.models.components import (
+                    cnn_atari,
+                )
+
+                return cnn_atari.CNNAtari(
+                    gray_scaled=self.is_gray_scale,
+                    model_size=self.model_size,
+                )
+            else:
+                raise ValueError(f"`framework={framework}` not supported!")
+
+        else:
+            if framework == "torch":
+                from ray.rllib.algorithms.dreamerv3.torch.models.components import mlp
+
+                return mlp.MLP(
+                    input_size=int(np.prod(self.observation_space.shape)),
+                    model_size=self.model_size,
+                )
+            else:
+                raise ValueError(f"`framework={framework}` not supported!")
+
+    def build_decoder(self, framework: str) -> Model:
+        """Builds the World-Model's decoder network depending on the obs space."""
+
+        if self.is_img_space:
+            if framework == "torch":
+                from ray.rllib.algorithms.dreamerv3.torch.models.components import (
+                    conv_transpose_atari,
+                )
+
+                return conv_transpose_atari.ConvTransposeAtari(
+                    input_size=self.h_plus_z_flat,
+                    gray_scaled=self.is_gray_scale,
+                    model_size=self.model_size,
+                )
+            else:
+                raise ValueError(f"`framework={framework}` not supported!")
+
+        else:
+            if framework == "torch":
+                from ray.rllib.algorithms.dreamerv3.torch.models.components import (
+                    vector_decoder,
+                )
+
+                return vector_decoder.VectorDecoder(
+                    input_size=self.h_plus_z_flat,
+                    model_size=self.model_size,
+                    observation_space=self.observation_space,
+                )
+            else:
+                raise ValueError(f"`framework={framework}` not supported!")
+
+    def build_world_model(self, framework: str, *, encoder, decoder) -> Model:
+        symlog_obs = do_symlog_obs(
+            self.observation_space,
+            self._model_config_dict.get("symlog_obs", "auto"),
+        )
+
+        if framework == "torch":
+            from ray.rllib.algorithms.dreamerv3.torch.models.world_model import (
+                WorldModel,
+            )
+        else:
+            raise ValueError(f"`framework={framework}` not supported!")
+
+        return WorldModel(
+            model_size=self.model_size,
+            observation_space=self.observation_space,
+            action_space=self.action_space,
+            batch_length_T=self._model_config_dict["batch_length_T"],
+            encoder=encoder,
+            decoder=decoder,
+            symlog_obs=symlog_obs,
+        )
+
+    def build_actor(self, framework: str) -> Model:
+        if framework == "torch":
+            from ray.rllib.algorithms.dreamerv3.torch.models.actor_network import (
+                ActorNetwork,
+            )
+
+            return ActorNetwork(
+                input_size=self.h_plus_z_flat,
+                action_space=self.action_space,
+                model_size=self.model_size,
+            )
+        else:
+            raise ValueError(f"`framework={framework}` not supported!")
+
+    def build_critic(self, framework: str) -> Model:
+        if framework == "torch":
+            from ray.rllib.algorithms.dreamerv3.torch.models.critic_network import (
+                CriticNetwork,
+            )
+
+            return CriticNetwork(
+                input_size=self.h_plus_z_flat,
+                model_size=self.model_size,
+            )
+        else:
+            raise ValueError(f"`framework={framework}` not supported!")
+
+    def build_dreamer_model(
+        self, framework: str, *, world_model, actor, critic, horizon=None, gamma=None
+    ) -> Model:
+        if framework == "torch":
+            from ray.rllib.algorithms.dreamerv3.torch.models.dreamer_model import (
+                DreamerModel,
+            )
+        else:
+            raise ValueError(f"`framework={framework}` not supported!")
+
+        return DreamerModel(
+            model_size=self.model_size,
+            action_space=self.action_space,
+            world_model=world_model,
+            actor=actor,
+            critic=critic,
+            **(
+                {}
+                if framework == "torch"
+                else {
+                    "horizon": horizon,
+                    "gamma": gamma,
+                }
+            ),
+        )

rllib/algorithms/dreamerv3/dreamerv3_learner.py

@@ -0,0 +1,31 @@
+"""
+[1] Mastering Diverse Domains through World Models - 2023
+D. Hafner, J. Pasukonis, J. Ba, T. Lillicrap
+https://arxiv.org/pdf/2301.04104v1.pdf
+
+[2] Mastering Atari with Discrete World Models - 2021
+D. Hafner, T. Lillicrap, M. Norouzi, J. Ba
+https://arxiv.org/pdf/2010.02193.pdf
+"""
+from ray.rllib.core.learner.learner import Learner
+from ray.rllib.utils.annotations import (
+    OverrideToImplementCustomLogic_CallToSuperRecommended,
+    override,
+)
+
+
+class DreamerV3Learner(Learner):
+    """DreamerV3 specific Learner class.
+
+    Only implements the `after_gradient_based_update()` method to define the logic
+    for updating the critic EMA-copy after each training step.
+    """
+
+    @OverrideToImplementCustomLogic_CallToSuperRecommended
+    @override(Learner)
+    def after_gradient_based_update(self, *, timesteps):
+        super().after_gradient_based_update(timesteps=timesteps)
+
+        # Update EMA weights of the critic.
+        for module_id, module in self.module._rl_modules.items():
+            module.unwrapped().critic.update_ema()

rllib/algorithms/dreamerv3/dreamerv3_rl_module.py

@@ -0,0 +1,85 @@
+"""
+This file holds framework-agnostic components for DreamerV3's RLModule.
+"""
+
+import abc
+from typing import Dict
+
+from ray.rllib.algorithms.dreamerv3.torch.models.actor_network import ActorNetwork
+from ray.rllib.algorithms.dreamerv3.torch.models.critic_network import CriticNetwork
+from ray.rllib.algorithms.dreamerv3.torch.models.dreamer_model import DreamerModel
+from ray.rllib.algorithms.dreamerv3.torch.models.world_model import WorldModel
+from ray.rllib.algorithms.dreamerv3.utils import (
+    do_symlog_obs,
+    get_gru_units,
+    get_num_z_categoricals,
+    get_num_z_classes,
+)
+from ray.rllib.core.rl_module.rl_module import RLModule
+from ray.rllib.utils.annotations import override
+from ray.util.annotations import DeveloperAPI
+
+ACTIONS_ONE_HOT = "actions_one_hot"
+
+
+@DeveloperAPI(stability="alpha")
+class DreamerV3RLModule(RLModule, abc.ABC):
+    @override(RLModule)
+    def setup(self):
+        super().setup()
+
+        # Gather model-relevant settings.
+        T = self.model_config["batch_length_T"]
+        symlog_obs = do_symlog_obs(
+            self.observation_space,
+            self.model_config.get("symlog_obs", "auto"),
+        )
+        model_size = self.model_config["model_size"]
+
+        # Build encoder and decoder from catalog.
+        self.encoder = self.catalog.build_encoder(framework=self.framework)
+        self.decoder = self.catalog.build_decoder(framework=self.framework)
+
+        # Build the world model (containing encoder and decoder).
+        self.world_model = WorldModel(
+            model_size=model_size,
+            observation_space=self.observation_space,
+            action_space=self.action_space,
+            batch_length_T=T,
+            encoder=self.encoder,
+            decoder=self.decoder,
+            symlog_obs=symlog_obs,
+        )
+        input_size = get_gru_units(model_size) + get_num_z_classes(
+            model_size
+        ) * get_num_z_categoricals(model_size)
+        self.actor = ActorNetwork(
+            input_size=input_size,
+            action_space=self.action_space,
+            model_size=model_size,
+        )
+        self.critic = CriticNetwork(
+            input_size=input_size,
+            model_size=model_size,
+        )
+        # Build the final dreamer model (containing the world model).
+        self.dreamer_model = DreamerModel(
+            model_size=self.model_config["model_size"],
+            action_space=self.action_space,
+            world_model=self.world_model,
+            actor=self.actor,
+            critic=self.critic,
+            # horizon=horizon_H,
+            # gamma=gamma,
+        )
+        self.action_dist_cls = self.catalog.get_action_dist_cls(
+            framework=self.framework
+        )
+
+        # Initialize the critic EMA net:
+        self.critic.init_ema()
+
+    @override(RLModule)
+    def get_initial_state(self) -> Dict:
+        # Use `DreamerModel`'s `get_initial_state` method.
+        return self.dreamer_model.get_initial_state()

rllib/algorithms/dreamerv3/tests/test_dreamerv3.py

@@ -0,0 +1,324 @@
+"""
+[1] Mastering Diverse Domains through World Models - 2023
+D. Hafner, J. Pasukonis, J. Ba, T. Lillicrap
+https://arxiv.org/pdf/2301.04104v1.pdf
+
+[2] Mastering Atari with Discrete World Models - 2021
+D. Hafner, T. Lillicrap, M. Norouzi, J. Ba
+https://arxiv.org/pdf/2010.02193.pdf
+
+[3]
+D. Hafner's (author) original code repo (for JAX):
+https://github.com/danijar/dreamerv3
+"""
+
+import unittest
+
+import gymnasium as gym
+import numpy as np
+import tree  # pip install dm_tree
+
+import ray
+from ray import tune
+from ray.rllib.algorithms.dreamerv3 import dreamerv3
+from ray.rllib.connectors.env_to_module import FlattenObservations
+from ray.rllib.core import DEFAULT_MODULE_ID
+from ray.rllib.env.wrappers.atari_wrappers import wrap_atari_for_new_api_stack
+from ray.rllib.env.wrappers.dm_control_wrapper import ActionClip, DMCEnv
+from ray.rllib.utils.framework import try_import_torch
+from ray.rllib.utils.numpy import one_hot
+from ray.rllib.utils.test_utils import check
+
+torch, nn = try_import_torch()
+
+
+class TestDreamerV3(unittest.TestCase):
+    @classmethod
+    def setUpClass(cls):
+        ray.init()
+
+    @classmethod
+    def tearDownClass(cls):
+        ray.shutdown()
+
+    def test_dreamerv3_compilation(self):
+        """Test whether DreamerV3 can be built with all frameworks."""
+
+        # Build a DreamerV3Config object.
+        config = (
+            dreamerv3.DreamerV3Config()
+            .env_runners(num_env_runners=0)
+            .training(
+                # Keep things simple. Especially the long dream rollouts seem
+                # to take an enormous amount of time (initially).
+                batch_size_B=4,
+                horizon_H=5,
+                batch_length_T=16,
+                model_size="nano",  # Use a tiny model for testing
+                symlog_obs=True,
+                use_float16=False,
+            )
+            .learners(
+                num_learners=2,
+                num_cpus_per_learner=1,
+                num_gpus_per_learner=0,
+            )
+        )
+
+        num_iterations = 3
+
+        for env in [
+            # "DMC/cartpole/swingup",  # causes strange MuJoCo error(s) on CI
+            "FrozenLake-v1",
+            "CartPole-v1",
+            "ale_py:ALE/MsPacman-v5",
+            "Pendulum-v1",
+        ]:
+            print("Env={}".format(env))
+
+            # Add one-hot observations for FrozenLake env.
+            if env == "FrozenLake-v1":
+                config.env_runners(
+                    env_to_module_connector=(
+                        lambda env, spaces, device: FlattenObservations()
+                    )
+                )
+            else:
+                config.env_runners(env_to_module_connector=None)
+
+            # Add Atari preprocessing.
+            if env == "ale_py:ALE/MsPacman-v5":
+
+                def env_creator(cfg):
+                    return wrap_atari_for_new_api_stack(
+                        gym.make(env, **cfg, render_mode="rgb_array"),
+                        # No frame-stacking. DreamerV3 processes color images with a
+                        # GRU, so partial observability is ok.
+                        framestack=None,
+                        grayscale=False,
+                    )
+
+                tune.register_env("env", env_creator)
+                env = "env"
+
+            elif env.startswith("DMC"):
+                parts = env.split("/")
+                assert len(parts) == 3, (
+                    "ERROR: DMC env must be formatted as 'DMC/[task]/[domain]', e.g. "
+                    f"'DMC/cartpole/swingup'! You provided '{env}'."
+                )
+
+                def env_creator(cfg):
+                    return ActionClip(
+                        DMCEnv(
+                            parts[1],
+                            parts[2],
+                            from_pixels=True,
+                            channels_first=False,
+                        )
+                    )
+
+                tune.register_env("env", env_creator)
+                env = "env"
+
+            config.environment(env)
+            algo = config.build_algo()
+            obs_space = algo.env_runner._env_to_module.observation_space
+            act_space = algo.env_runner.env.single_action_space
+            rl_module = algo.env_runner.module
+
+            for i in range(num_iterations):
+                results = algo.train()
+                print(results)
+            # Test dream trajectory w/ recreated observations.
+            sample = algo.replay_buffer.sample()
+            start_states = rl_module.dreamer_model.get_initial_state()
+            start_states = tree.map_structure(
+                # Repeat only the batch dimension (B times).
+                lambda s: s.unsqueeze(0).repeat(1, *([1] * len(s.shape))),
+                start_states,
+            )
+            dream = rl_module.dreamer_model.dream_trajectory_with_burn_in(
+                start_states=start_states,
+                timesteps_burn_in=5,
+                timesteps_H=45,
+                observations=torch.from_numpy(sample["obs"][:1]),  # B=1
+                actions=torch.from_numpy(
+                    one_hot(
+                        sample["actions"],
+                        depth=act_space.n,
+                    )
+                    if isinstance(act_space, gym.spaces.Discrete)
+                    else sample["actions"]
+                )[
+                    :1
+                ],  # B=1
+            )
+            check(
+                dream["actions_dreamed_t0_to_H_BxT"].shape,
+                (46, 1)
+                + (
+                    (act_space.n,)
+                    if isinstance(act_space, gym.spaces.Discrete)
+                    else tuple(act_space.shape)
+                ),
+            )
+            check(dream["continues_dreamed_t0_to_H_BxT"].shape, (46, 1))
+            check(
+                dream["observations_dreamed_t0_to_H_BxT"].shape,
+                [46, 1] + list(obs_space.shape),
+            )
+            algo.stop()
+
+    def test_dreamerv3_dreamer_model_sizes(self):
+        """Tests, whether the different model sizes match the ones reported in [1]."""
+
+        # For Atari, these are the exact numbers from the repo ([3]).
+        # However, for CartPole + size "S" and "M", the author's original code will not
+        # match for the world model count. This is due to the fact that the author uses
+        # encoder/decoder nets with 5x1024 nodes (which corresponds to XL) regardless of
+        # the `model_size` settings (iff >="S").
+        expected_num_params_world_model = {
+            # XS encoder
+            # kernel=[4, 256], (no bias), layernorm=[256],[256]
+            # XS reward_predictor
+            # kernel=[1280, 256], (no bias), layernorm[256],[256]
+            # kernel=[256, 255] bias=[255]
+            # 1280=1024 (z-state) + 256 (h-state)
+            # XS continue_predictor
+            # kernel=[1280, 256], (no bias), layernorm=[256],[256]
+            # kernel=[256, 1] bias=[1]
+            # XS sequence_model
+            # [
+            # pre-MLP: kernel=[1026, 256], (no bias), layernorm=[256],[256], silu
+            # custom GRU: kernel=[512, 768], (no bias), layernorm=[768],[768]
+            # ]
+            # XS decoder
+            # kernel=[1280, 256], (no bias), layernorm=[256],[256]
+            # kernel=[256, 4] bias=[4]
+            # XS posterior_mlp
+            # kernel=[512, 256], (no bias), layernorm=[256],[256]
+            # XS posterior_representation_layer
+            # kernel=[256, 1024], bias=[1024]
+            "XS_cartpole": 2435076,
+            "S_cartpole": 7493380,
+            "M_cartpole": 16206084,
+            "L_cartpole": 37802244,
+            "XL_cartpole": 108353796,
+            # XS encoder (atari)
+            # cnn kernel=[4, 4, 3, 24], (no bias), layernorm=[24],[24],
+            # cnn kernel=[4, 4, 24, 48], (no bias), layernorm=[48],[48],
+            # cnn kernel=[4, 4, 48, 96], (no bias), layernorm=[96],[96],
+            # cnn kernel=[4, 4, 96, 192], (no bias), layernorm=[192],[192],
+            # XS decoder (atari)
+            # init dense kernel[1280, 3072] bias=[3072] -> reshape into image
+            # [4, 4, 96, 192], [96], [96]
+            # [4, 4, 48, 96], [48], [48],
+            # [4, 4, 24, 48], [24], [24],
+            # [4, 4, 3, 24], [3] <- no layernorm at end
+            "XS_atari": 7538979,
+            "S_atari": 15687811,
+            "M_atari": 32461635,
+            "L_atari": 68278275,
+            "XL_atari": 181558659,
+        }
+
+        # All values confirmed against [3] (100% match).
+        expected_num_params_actor = {
+            # hidden=[1280, 256]
+            # hidden_norm=[256], [256]
+            # pi (2 actions)=[256, 2], [2]
+            "XS_cartpole": 328706,
+            "S_cartpole": 1051650,
+            "M_cartpole": 2135042,
+            "L_cartpole": 4136450,
+            "XL_cartpole": 9449474,
+            "XS_atari": 329734,
+            "S_atari": 1053702,
+            "M_atari": 2137606,
+            "L_atari": 4139526,
+            "XL_atari": 9453574,
+        }
+
+        # All values confirmed against [3] (100% match).
+        expected_num_params_critic = {
+            # hidden=[1280, 256]
+            # hidden_norm=[256], [256]
+            # vf (buckets)=[256, 255], [255]
+            "XS_cartpole": 393727,
+            "S_cartpole": 1181439,
+            "M_cartpole": 2297215,
+            "L_cartpole": 4331007,
+            "XL_cartpole": 9708799,
+            "XS_atari": 393727,
+            "S_atari": 1181439,
+            "M_atari": 2297215,
+            "L_atari": 4331007,
+            "XL_atari": 9708799,
+        }
+
+        config = dreamerv3.DreamerV3Config().training(
+            batch_length_T=16,
+            horizon_H=5,
+            symlog_obs=True,
+        )
+
+        # Check all model_sizes described in the paper ([1]) on matching the number
+        # of parameters to RLlib's implementation.
+        for model_size in ["XS", "S", "M", "L", "XL"]:
+            config.model_size = model_size
+
+            # Atari and CartPole spaces.
+            for obs_space, num_actions, env_name in [
+                (gym.spaces.Box(-1.0, 0.0, (4,), np.float32), 2, "cartpole"),
+                (gym.spaces.Box(-1.0, 0.0, (64, 64, 3), np.float32), 6, "atari"),
+            ]:
+                print(f"Testing model_size={model_size} on env-type: {env_name} ..")
+                config.environment(
+                    observation_space=obs_space,
+                    action_space=gym.spaces.Discrete(num_actions),
+                )
+
+                # Create our RLModule to compute actions with.
+                policy_dict, _ = config.get_multi_agent_setup()
+                module_spec = config.get_multi_rl_module_spec(policy_dict=policy_dict)
+                rl_module = module_spec.build()[DEFAULT_MODULE_ID]
+
+                # Count the generated RLModule's parameters and compare to the
+                # paper's reported numbers ([1] and [3]).
+                num_params_world_model = sum(
+                    np.prod(v.shape)
+                    for v in rl_module.world_model.parameters()
+                    if v.requires_grad
+                )
+                self.assertEqual(
+                    num_params_world_model,
+                    expected_num_params_world_model[f"{model_size}_{env_name}"],
+                )
+                num_params_actor = sum(
+                    np.prod(v.shape)
+                    for v in rl_module.actor.parameters()
+                    if v.requires_grad
+                )
+                self.assertEqual(
+                    num_params_actor,
+                    expected_num_params_actor[f"{model_size}_{env_name}"],
+                )
+                num_params_critic = sum(
+                    np.prod(v.shape)
+                    for v in rl_module.critic.parameters()
+                    if v.requires_grad
+                )
+                self.assertEqual(
+                    num_params_critic,
+                    expected_num_params_critic[f"{model_size}_{env_name}"],
+                )
+                print("\tok")
+
+
+if __name__ == "__main__":
+    import sys
+
+    import pytest
+
+    sys.exit(pytest.main(["-v", __file__]))

rllib/algorithms/dreamerv3/torch/dreamerv3_torch_learner.py

@@ -0,0 +1,925 @@
+"""
+[1] Mastering Diverse Domains through World Models - 2023
+D. Hafner, J. Pasukonis, J. Ba, T. Lillicrap
+https://arxiv.org/pdf/2301.04104v1.pdf
+
+[2] Mastering Atari with Discrete World Models - 2021
+D. Hafner, T. Lillicrap, M. Norouzi, J. Ba
+https://arxiv.org/pdf/2010.02193.pdf
+"""
+from typing import Any, Dict, Tuple
+
+import gymnasium as gym
+
+from ray.rllib.algorithms.dreamerv3.dreamerv3 import DreamerV3Config
+from ray.rllib.algorithms.dreamerv3.dreamerv3_learner import DreamerV3Learner
+from ray.rllib.core import DEFAULT_MODULE_ID
+from ray.rllib.core.columns import Columns
+from ray.rllib.core.learner.learner import ParamDict
+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.torch_utils import clip_gradients, symlog, two_hot
+from ray.rllib.utils.typing import ModuleID, TensorType
+
+torch, nn = try_import_torch()
+
+
+class DreamerV3TorchLearner(DreamerV3Learner, TorchLearner):
+    """Implements DreamerV3 losses and gradient-based update logic in PyTorch.
+
+    The critic EMA-copy update step can be found in the `DreamerV3Learner` base class,
+    as it is framework independent.
+
+    We define 3 local PyTorch optimizers for the sub components "world_model",
+    "actor", and "critic". Each of these optimizers might use a different learning rate,
+    epsilon parameter, and gradient clipping thresholds and procedures.
+    """
+
+    def build(self) -> None:
+        super().build()
+
+        # Store loss tensors here temporarily inside the loss function for (exact)
+        # consumption later by the compute gradients function.
+        # Keys=(module_id, optimizer_name), values=loss tensors (in-graph).
+        self._temp_losses = {}
+
+    @override(TorchLearner)
+    def configure_optimizers_for_module(
+        self, module_id: ModuleID, config: DreamerV3Config = None
+    ):
+        """Create the 3 optimizers for Dreamer learning: world_model, actor, critic.
+
+        The learning rates used are described in [1] and the epsilon values used here
+        - albeit probably not that important - are used by the author's own
+        implementation.
+        """
+
+        dreamerv3_module = self._module[module_id]
+
+        # World Model optimizer.
+        optim_world_model = torch.optim.Adam(
+            dreamerv3_module.world_model.parameters(),
+            eps=1e-8,
+        )
+        self.register_optimizer(
+            module_id=module_id,
+            optimizer_name="world_model",
+            optimizer=optim_world_model,
+            params=list(dreamerv3_module.world_model.parameters()),
+            lr_or_lr_schedule=config.world_model_lr,
+        )
+
+        # Actor optimizer.
+        optim_actor = torch.optim.Adam(dreamerv3_module.actor.parameters(), eps=1e-5)
+        self.register_optimizer(
+            module_id=module_id,
+            optimizer_name="actor",
+            optimizer=optim_actor,
+            params=list(dreamerv3_module.actor.parameters()),
+            lr_or_lr_schedule=config.actor_lr,
+        )
+
+        # Critic optimizer.
+        optim_critic = torch.optim.Adam(dreamerv3_module.critic.parameters(), eps=1e-5)
+        self.register_optimizer(
+            module_id=module_id,
+            optimizer_name="critic",
+            optimizer=optim_critic,
+            params=list(dreamerv3_module.critic.parameters()),
+            lr_or_lr_schedule=config.critic_lr,
+        )
+
+    @override(TorchLearner)
+    def postprocess_gradients_for_module(
+        self,
+        *,
+        module_id: ModuleID,
+        config: DreamerV3Config,
+        module_gradients_dict: Dict[str, Any],
+    ) -> ParamDict:
+        """Performs gradient clipping on the 3 module components' computed grads.
+
+        Note that different grad global-norm clip values are used for the 3
+        module components: world model, actor, and critic.
+        """
+        for optimizer_name, optimizer in self.get_optimizers_for_module(
+            module_id=module_id
+        ):
+            grads_sub_dict = self.filter_param_dict_for_optimizer(
+                module_gradients_dict, optimizer
+            )
+            # Figure out which grad clip setting to use.
+            grad_clip = (
+                config.world_model_grad_clip_by_global_norm
+                if optimizer_name == "world_model"
+                else config.actor_grad_clip_by_global_norm
+                if optimizer_name == "actor"
+                else config.critic_grad_clip_by_global_norm
+            )
+            global_norm = clip_gradients(
+                grads_sub_dict, grad_clip=grad_clip, grad_clip_by="global_norm"
+            )
+            module_gradients_dict.update(grads_sub_dict)
+
+            # DreamerV3 stats have the format: [WORLD_MODEL|ACTOR|CRITIC]_[stats name].
+            self.metrics.log_dict(
+                {
+                    optimizer_name.upper()
+                    + "_gradients_global_norm": (global_norm.item()),
+                    optimizer_name.upper()
+                    + "_gradients_maxabs_after_clipping": (
+                        torch.max(
+                            torch.abs(
+                                torch.cat(
+                                    [g.flatten() for g in grads_sub_dict.values()]
+                                )
+                            )
+                        ).item()
+                    ),
+                },
+                key=module_id,
+                window=1,  # <- single items (should not be mean/ema-reduced over time).
+            )
+
+        return module_gradients_dict
+
+    @override(TorchLearner)
+    def compute_gradients(
+        self,
+        loss_per_module,
+        **kwargs,
+    ):
+        """Override of the default gradient computation method.
+
+        For DreamerV3, we need to compute gradients over the individual loss terms
+        as otherwise, the world model's parameters would have their gradients also
+        be influenced by the actor- and critic loss terms/gradient computations.
+        """
+        grads = {}
+
+        # Do actor and critic's grad computations first, such that after those two,
+        # we can zero out the gradients of the world model again (they will have values
+        # in them from the actor/critic backwards).
+        for component in ["actor", "critic", "world_model"]:
+            optim = self.get_optimizer(DEFAULT_MODULE_ID, component)
+            optim.zero_grad(set_to_none=True)
+            # Do the backward pass
+            loss = self._temp_losses.pop(component.upper())
+            loss.backward(retain_graph=component in ["actor", "critic"])
+            optim_grads = {
+                pid: p.grad
+                for pid, p in self.filter_param_dict_for_optimizer(
+                    self._params, optim
+                ).items()
+            }
+            for ref, grad in optim_grads.items():
+                assert ref not in grads
+                grads[ref] = grad
+
+        return grads
+
+    @override(TorchLearner)
+    def compute_loss_for_module(
+        self,
+        module_id: ModuleID,
+        config: DreamerV3Config,
+        batch: Dict[str, TensorType],
+        fwd_out: Dict[str, TensorType],
+    ) -> TensorType:
+        # World model losses.
+        prediction_losses = self._compute_world_model_prediction_losses(
+            config=config,
+            rewards_B_T=batch[Columns.REWARDS],
+            continues_B_T=(1.0 - batch["is_terminated"].float()),
+            fwd_out=fwd_out,
+        )
+
+        (
+            L_dyn_B_T,
+            L_rep_B_T,
+        ) = self._compute_world_model_dynamics_and_representation_loss(
+            config=config, fwd_out=fwd_out
+        )
+        L_dyn = torch.mean(L_dyn_B_T)
+        L_rep = torch.mean(L_rep_B_T)
+        # Make sure values for L_rep and L_dyn are the same (they only differ in their
+        # gradients).
+        assert torch.allclose(L_dyn, L_rep)
+
+        # Compute the actual total loss using fixed weights described in [1] eq. 4.
+        L_world_model_total_B_T = (
+            1.0 * prediction_losses["L_prediction_B_T"]
+            + 0.5 * L_dyn_B_T
+            + 0.1 * L_rep_B_T
+        )
+
+        # In the paper, it says to sum up timesteps, and average over
+        # batch (see eq. 4 in [1]). But Danijar's implementation only does
+        # averaging (over B and T), so we'll do this here as well. This is generally
+        # true for all other loss terms as well (we'll always just average, no summing
+        # over T axis!).
+        L_world_model_total = torch.mean(L_world_model_total_B_T)
+
+        # Log world model loss stats.
+        self.metrics.log_dict(
+            {
+                "WORLD_MODEL_learned_initial_h": self.module[module_id]
+                .unwrapped()
+                .world_model.initial_h.mean(),
+                # Prediction losses.
+                # Decoder (obs) loss.
+                "WORLD_MODEL_L_decoder": prediction_losses["L_decoder"],
+                # Reward loss.
+                "WORLD_MODEL_L_reward": prediction_losses["L_reward"],
+                # Continue loss.
+                "WORLD_MODEL_L_continue": prediction_losses["L_continue"],
+                # Total.
+                "WORLD_MODEL_L_prediction": prediction_losses["L_prediction"],
+                # Dynamics loss.
+                "WORLD_MODEL_L_dynamics": L_dyn,
+                # Representation loss.
+                "WORLD_MODEL_L_representation": L_rep,
+                # Total loss.
+                "WORLD_MODEL_L_total": L_world_model_total,
+            },
+            key=module_id,
+            window=1,  # <- single items (should not be mean/ema-reduced over time).
+        )
+
+        # Add the predicted obs distributions for possible (video) summarization.
+        if config.report_images_and_videos:
+            self.metrics.log_value(
+                (module_id, "WORLD_MODEL_fwd_out_obs_distribution_means_b0xT"),
+                fwd_out["obs_distribution_means_BxT"][: self.config.batch_length_T],
+                reduce="item_series",  # No reduction, we want the obs tensor to stay in-tact.
+                window=1,  # <- single items (should not be mean/ema-reduced over time).
+            )
+
+        if config.report_individual_batch_item_stats:
+            # Log important world-model loss stats.
+            self.metrics.log_dict(
+                {
+                    "WORLD_MODEL_L_decoder_B_T": prediction_losses["L_decoder_B_T"],
+                    "WORLD_MODEL_L_reward_B_T": prediction_losses["L_reward_B_T"],
+                    "WORLD_MODEL_L_continue_B_T": prediction_losses["L_continue_B_T"],
+                    "WORLD_MODEL_L_prediction_B_T": (
+                        prediction_losses["L_prediction_B_T"]
+                    ),
+                    "WORLD_MODEL_L_dynamics_B_T": L_dyn_B_T,
+                    "WORLD_MODEL_L_representation_B_T": L_rep_B_T,
+                    "WORLD_MODEL_L_total_B_T": L_world_model_total_B_T,
+                },
+                key=module_id,
+                window=1,  # <- single items (should not be mean/ema-reduced over time).
+            )
+
+        # Dream trajectories starting in all internal states (h + z_posterior) that were
+        # computed during world model training.
+        # Everything goes in as BxT: We are starting a new dream trajectory at every
+        # actually encountered timestep in the batch, so we are creating B*T
+        # trajectories of len `horizon_H`.
+        dream_data = (
+            self.module[module_id]
+            .unwrapped()
+            .dreamer_model.dream_trajectory(
+                start_states={
+                    "h": fwd_out["h_states_BxT"],
+                    "z": fwd_out["z_posterior_states_BxT"],
+                },
+                start_is_terminated=batch["is_terminated"].reshape(-1),  # -> BxT
+                timesteps_H=config.horizon_H,
+                gamma=config.gamma,
+            )
+        )
+        if config.report_dream_data:
+            # To reduce this massive amount of data a little, slice out a T=1 piece
+            # from each stats that has the shape (H, BxT), meaning convert e.g.
+            # `rewards_dreamed_t0_to_H_BxT` into `rewards_dreamed_t0_to_H_Bx1`.
+            # This will reduce the amount of data to be transferred and reported
+            # by the factor of `batch_length_T`.
+            self.metrics.log_dict(
+                {
+                    # Replace 'T' with '1'.
+                    key[:-1] + "1": value[:, :: config.batch_length_T]
+                    for key, value in dream_data.items()
+                    if key.endswith("H_BxT")
+                },
+                key=(module_id, "dream_data"),
+                reduce="item_series",
+                window=1,  # <- single items (should not be mean/ema-reduced over time).
+            )
+
+        value_targets_t0_to_Hm1_BxT = self._compute_value_targets(
+            config=config,
+            # Learn critic in symlog'd space.
+            rewards_t0_to_H_BxT=dream_data["rewards_dreamed_t0_to_H_BxT"],
+            intrinsic_rewards_t1_to_H_BxT=(
+                dream_data["rewards_intrinsic_t1_to_H_B"]
+                if config.use_curiosity
+                else None
+            ),
+            continues_t0_to_H_BxT=dream_data["continues_dreamed_t0_to_H_BxT"],
+            value_predictions_t0_to_H_BxT=dream_data["values_dreamed_t0_to_H_BxT"],
+        )
+        # self.metrics.log_value(
+        #    key=(module_id, "VALUE_TARGETS_H_BxT"),
+        #    value=value_targets_t0_to_Hm1_BxT,
+        #    window=1,  # <- single items (should not be mean/ema-reduced over time).
+        # )
+
+        CRITIC_L_total = self._compute_critic_loss(
+            module_id=module_id,
+            config=config,
+            dream_data=dream_data,
+            value_targets_t0_to_Hm1_BxT=value_targets_t0_to_Hm1_BxT,
+        )
+        if config.train_actor:
+            ACTOR_L_total = self._compute_actor_loss(
+                module_id=module_id,
+                config=config,
+                dream_data=dream_data,
+                value_targets_t0_to_Hm1_BxT=value_targets_t0_to_Hm1_BxT,
+            )
+        else:
+            ACTOR_L_total = 0.0
+
+        self._temp_losses["ACTOR"] = ACTOR_L_total
+        self._temp_losses["CRITIC"] = CRITIC_L_total
+        self._temp_losses["WORLD_MODEL"] = L_world_model_total
+
+        # Return the total loss as a sum of all individual losses.
+        return L_world_model_total + CRITIC_L_total + ACTOR_L_total
+
+    def _compute_world_model_prediction_losses(
+        self,
+        *,
+        config: DreamerV3Config,
+        rewards_B_T: TensorType,
+        continues_B_T: TensorType,
+        fwd_out: Dict[str, TensorType],
+    ) -> Dict[str, TensorType]:
+        """Helper method computing all world-model related prediction losses.
+
+        Prediction losses are used to train the predictors of the world model, which
+        are: Reward predictor, continue predictor, and the decoder (which predicts
+        observations).
+
+        Args:
+            config: The DreamerV3Config to use.
+            rewards_B_T: The rewards batch in the shape (B, T) and of type float32.
+            continues_B_T: The continues batch in the shape (B, T) and of type float32
+                (1.0 -> continue; 0.0 -> end of episode).
+            fwd_out: The `forward_train` outputs of the DreamerV3RLModule.
+        """
+
+        # Learn to produce symlog'd observation predictions.
+        # If symlog is disabled (e.g. for uint8 image inputs), `obs_symlog_BxT` is the
+        # same as `obs_BxT`.
+        obs_BxT = fwd_out["sampled_obs_symlog_BxT"]
+        obs_distr_means = fwd_out["obs_distribution_means_BxT"]
+
+        # Leave time dim folded (BxT) and flatten all other (e.g. image) dims.
+        obs_BxT = obs_BxT.reshape(obs_BxT.shape[0], -1)
+
+        # Squared diff loss w/ sum(!) over all (already folded) obs dims.
+        # decoder_loss_BxT = SUM[ (obs_distr.loc - observations)^2 ]
+        # Note: This is described strangely in the paper (stating a neglogp loss here),
+        # but the author's own implementation actually uses simple MSE with the loc
+        # of the Gaussian.
+        decoder_loss_BxT = torch.sum(torch.square(obs_distr_means - obs_BxT), dim=-1)
+
+        # Unfold time rank back in.
+        decoder_loss_B_T = decoder_loss_BxT.reshape(
+            config.batch_size_B_per_learner, config.batch_length_T
+        )
+        L_decoder = torch.mean(decoder_loss_B_T)
+
+        # The FiniteDiscrete reward bucket distribution computed by our reward
+        # predictor.
+        # [B x num_buckets].
+        reward_logits_BxT = fwd_out["reward_logits_BxT"]
+        # Learn to produce symlog'd reward predictions.
+        rewards_symlog_B_T = symlog(rewards_B_T)
+        # Fold time dim.
+        rewards_symlog_BxT = rewards_symlog_B_T.reshape(-1)
+
+        # Two-hot encode.
+        two_hot_rewards_symlog_BxT = two_hot(rewards_symlog_BxT, device=self._device)
+        # two_hot_rewards_symlog_BxT=[B*T, num_buckets]
+        reward_log_pred_BxT = reward_logits_BxT - torch.logsumexp(
+            reward_logits_BxT, dim=-1, keepdim=True
+        )
+        # Multiply with two-hot targets and neg.
+        reward_loss_two_hot_BxT = -torch.sum(
+            reward_log_pred_BxT * two_hot_rewards_symlog_BxT, dim=-1
+        )
+        # Unfold time rank back in.
+        reward_loss_two_hot_B_T = reward_loss_two_hot_BxT.reshape(
+            config.batch_size_B_per_learner, config.batch_length_T
+        )
+        L_reward_two_hot = torch.mean(reward_loss_two_hot_B_T)
+
+        # Probabilities that episode continues, computed by our continue predictor.
+        # [B]
+        continue_distr = fwd_out["continue_distribution_BxT"]
+        # -log(p) loss
+        # Fold time dim.
+        continues_BxT = continues_B_T.reshape(-1)
+        continue_loss_BxT = -continue_distr.log_prob(continues_BxT)
+        # Unfold time rank back in.
+        continue_loss_B_T = continue_loss_BxT.reshape(
+            config.batch_size_B_per_learner, config.batch_length_T
+        )
+        L_continue = torch.mean(continue_loss_B_T)
+
+        # Sum all losses together as the "prediction" loss.
+        L_pred_B_T = decoder_loss_B_T + reward_loss_two_hot_B_T + continue_loss_B_T
+        L_pred = torch.mean(L_pred_B_T)
+
+        return {
+            "L_decoder_B_T": decoder_loss_B_T,
+            "L_decoder": L_decoder,
+            "L_reward": L_reward_two_hot,
+            "L_reward_B_T": reward_loss_two_hot_B_T,
+            "L_continue": L_continue,
+            "L_continue_B_T": continue_loss_B_T,
+            "L_prediction": L_pred,
+            "L_prediction_B_T": L_pred_B_T,
+        }
+
+    def _compute_world_model_dynamics_and_representation_loss(
+        self, *, config: DreamerV3Config, fwd_out: Dict[str, Any]
+    ) -> Tuple[TensorType, TensorType]:
+        """Helper method computing the world-model's dynamics and representation losses.
+
+        Args:
+            config: The DreamerV3Config to use.
+            fwd_out: The `forward_train` outputs of the DreamerV3RLModule.
+
+        Returns:
+            Tuple consisting of a) dynamics loss: Trains the prior network, predicting
+            z^ prior states from h-states and b) representation loss: Trains posterior
+            network, predicting z posterior states from h-states and (encoded)
+            observations.
+        """
+
+        # Actual distribution over stochastic internal states (z) produced by the
+        # encoder.
+        z_posterior_probs_BxT = fwd_out["z_posterior_probs_BxT"]
+        z_posterior_distr_BxT = torch.distributions.Independent(
+            torch.distributions.OneHotCategorical(probs=z_posterior_probs_BxT),
+            reinterpreted_batch_ndims=1,
+        )
+
+        # Actual distribution over stochastic internal states (z) produced by the
+        # dynamics network.
+        z_prior_probs_BxT = fwd_out["z_prior_probs_BxT"]
+        z_prior_distr_BxT = torch.distributions.Independent(
+            torch.distributions.OneHotCategorical(probs=z_prior_probs_BxT),
+            reinterpreted_batch_ndims=1,
+        )
+
+        # Stop gradient for encoder's z-outputs:
+        sg_z_posterior_distr_BxT = torch.distributions.Independent(
+            torch.distributions.OneHotCategorical(probs=z_posterior_probs_BxT.detach()),
+            reinterpreted_batch_ndims=1,
+        )
+        # Stop gradient for dynamics model's z-outputs:
+        sg_z_prior_distr_BxT = torch.distributions.Independent(
+            torch.distributions.OneHotCategorical(probs=z_prior_probs_BxT.detach()),
+            reinterpreted_batch_ndims=1,
+        )
+
+        # Implement free bits. According to [1]:
+        # "To avoid a degenerate solution where the dynamics are trivial to predict but
+        # contain not enough information about the inputs, we employ free bits by
+        # clipping the dynamics and representation losses below the value of
+        # 1 nat ≈ 1.44 bits. This disables them while they are already minimized well to
+        # focus the world model on its prediction loss"
+        L_dyn_BxT = torch.clamp(
+            torch.distributions.kl.kl_divergence(
+                sg_z_posterior_distr_BxT, z_prior_distr_BxT
+            ),
+            min=1.0,
+        )
+        # Unfold time rank back in.
+        L_dyn_B_T = L_dyn_BxT.reshape(
+            config.batch_size_B_per_learner, config.batch_length_T
+        )
+
+        L_rep_BxT = torch.clamp(
+            torch.distributions.kl.kl_divergence(
+                z_posterior_distr_BxT, sg_z_prior_distr_BxT
+            ),
+            min=1.0,
+        )
+        # Unfold time rank back in.
+        L_rep_B_T = L_rep_BxT.reshape(
+            config.batch_size_B_per_learner, config.batch_length_T
+        )
+
+        return L_dyn_B_T, L_rep_B_T
+
+    def _compute_actor_loss(
+        self,
+        *,
+        module_id: ModuleID,
+        config: DreamerV3Config,
+        dream_data: Dict[str, TensorType],
+        value_targets_t0_to_Hm1_BxT: TensorType,
+    ) -> TensorType:
+        """Helper method computing the actor's loss terms.
+
+        Args:
+            module_id: The module_id for which to compute the actor loss.
+            config: The DreamerV3Config to use.
+            dream_data: The data generated by dreaming for H steps (horizon) starting
+                from any BxT state (sampled from the buffer for the train batch).
+            value_targets_t0_to_Hm1_BxT: The computed value function targets of the
+                shape (t0 to H-1, BxT).
+
+        Returns:
+            The total actor loss tensor.
+        """
+        actor = self.module[module_id].unwrapped().actor
+
+        # Note: `scaled_value_targets_t0_to_Hm1_B` are NOT stop_gradient'd yet.
+        scaled_value_targets_t0_to_Hm1_B = self._compute_scaled_value_targets(
+            module_id=module_id,
+            config=config,
+            value_targets_t0_to_Hm1_BxT=value_targets_t0_to_Hm1_BxT,
+            value_predictions_t0_to_Hm1_BxT=dream_data["values_dreamed_t0_to_H_BxT"][
+                :-1
+            ],
+        )
+
+        # Actions actually taken in the dream.
+        actions_dreamed = dream_data["actions_dreamed_t0_to_H_BxT"][:-1].detach()
+        actions_dreamed_dist_params_t0_to_Hm1_B = dream_data[
+            "actions_dreamed_dist_params_t0_to_H_BxT"
+        ][:-1]
+
+        dist_t0_to_Hm1_B = actor.get_action_dist_object(
+            actions_dreamed_dist_params_t0_to_Hm1_B
+        )
+
+        # Compute log(p)s of all possible actions in the dream.
+        if isinstance(
+            self.module[module_id].unwrapped().actor.action_space, gym.spaces.Discrete
+        ):
+            # Note that when we create the Categorical action distributions, we compute
+            # unimix probs, then math.log these and provide these log(p) as "logits" to
+            # the Categorical. So here, we'll continue to work with log(p)s (not
+            # really "logits")!
+            logp_actions_t0_to_Hm1_B = actions_dreamed_dist_params_t0_to_Hm1_B
+
+            # Log probs of actions actually taken in the dream.
+            logp_actions_dreamed_t0_to_Hm1_B = torch.sum(
+                actions_dreamed * logp_actions_t0_to_Hm1_B,
+                dim=-1,
+            )
+            # First term of loss function. [1] eq. 11.
+            logp_loss_H_B = (
+                logp_actions_dreamed_t0_to_Hm1_B
+                * scaled_value_targets_t0_to_Hm1_B.detach()
+            )
+        # Box space.
+        else:
+            logp_actions_dreamed_t0_to_Hm1_B = dist_t0_to_Hm1_B.log_prob(
+                actions_dreamed
+            )
+            # First term of loss function. [1] eq. 11.
+            logp_loss_H_B = scaled_value_targets_t0_to_Hm1_B
+
+        assert logp_loss_H_B.ndim == 2
+
+        # Add entropy loss term (second term [1] eq. 11).
+        entropy_H_B = dist_t0_to_Hm1_B.entropy()
+        assert entropy_H_B.ndim == 2
+        entropy = torch.mean(entropy_H_B)
+
+        L_actor_reinforce_term_H_B = -logp_loss_H_B
+        L_actor_action_entropy_term_H_B = -config.entropy_scale * entropy_H_B
+
+        L_actor_H_B = L_actor_reinforce_term_H_B + L_actor_action_entropy_term_H_B
+        # Mask out everything that goes beyond a predicted continue=False boundary.
+        L_actor_H_B *= dream_data["dream_loss_weights_t0_to_H_BxT"][:-1].detach()
+        L_actor = torch.mean(L_actor_H_B)
+
+        # Log important actor loss stats.
+        self.metrics.log_dict(
+            {
+                "ACTOR_L_total": L_actor,
+                "ACTOR_value_targets_pct95_ema": actor.ema_value_target_pct95,
+                "ACTOR_value_targets_pct5_ema": actor.ema_value_target_pct5,
+                "ACTOR_action_entropy": entropy,
+                # Individual loss terms.
+                "ACTOR_L_neglogp_reinforce_term": torch.mean(
+                    L_actor_reinforce_term_H_B
+                ),
+                "ACTOR_L_neg_entropy_term": torch.mean(L_actor_action_entropy_term_H_B),
+            },
+            key=module_id,
+            window=1,  # <- single items (should not be mean/ema-reduced over time).
+        )
+        if config.report_individual_batch_item_stats:
+            self.metrics.log_dict(
+                {
+                    "ACTOR_L_total_H_BxT": L_actor_H_B,
+                    "ACTOR_logp_actions_dreamed_H_BxT": (
+                        logp_actions_dreamed_t0_to_Hm1_B
+                    ),
+                    "ACTOR_scaled_value_targets_H_BxT": (
+                        scaled_value_targets_t0_to_Hm1_B
+                    ),
+                    "ACTOR_action_entropy_H_BxT": entropy_H_B,
+                    # Individual loss terms.
+                    "ACTOR_L_neglogp_reinforce_term_H_BxT": L_actor_reinforce_term_H_B,
+                    "ACTOR_L_neg_entropy_term_H_BxT": L_actor_action_entropy_term_H_B,
+                },
+                key=module_id,
+                window=1,  # <- single items (should not be mean/ema-reduced over time).
+            )
+
+        return L_actor
+
+    def _compute_critic_loss(
+        self,
+        *,
+        module_id: ModuleID,
+        config: DreamerV3Config,
+        dream_data: Dict[str, TensorType],
+        value_targets_t0_to_Hm1_BxT: TensorType,
+    ) -> TensorType:
+        """Helper method computing the critic's loss terms.
+
+        Args:
+            module_id: The ModuleID for which to compute the critic loss.
+            config: The DreamerV3Config to use.
+            dream_data: The data generated by dreaming for H steps (horizon) starting
+                from any BxT state (sampled from the buffer for the train batch).
+            value_targets_t0_to_Hm1_BxT: The computed value function targets of the
+                shape (t0 to H-1, BxT).
+
+        Returns:
+            The total critic loss tensor.
+        """
+        # B=BxT
+        H, B = dream_data["rewards_dreamed_t0_to_H_BxT"].shape[:2]
+        Hm1 = H - 1
+
+        # Note that value targets are NOT symlog'd and go from t0 to H-1, not H, like
+        # all the other dream data.
+
+        # From here on: B=BxT
+        value_targets_t0_to_Hm1_B = value_targets_t0_to_Hm1_BxT.detach()
+        value_symlog_targets_t0_to_Hm1_B = symlog(value_targets_t0_to_Hm1_B)
+        # Fold time rank (for two_hot'ing).
+        value_symlog_targets_HxB = value_symlog_targets_t0_to_Hm1_B.view(
+            -1,
+        )
+        value_symlog_targets_two_hot_HxB = two_hot(
+            value_symlog_targets_HxB, device=self._device
+        )
+        # Unfold time rank.
+        value_symlog_targets_two_hot_t0_to_Hm1_B = (
+            value_symlog_targets_two_hot_HxB.view(
+                [Hm1, B, value_symlog_targets_two_hot_HxB.shape[-1]]
+            )
+        )
+
+        # Get (B x T x probs) tensor from return distributions.
+        # Use the value function outputs that don't graph-trace back through the
+        # world model. The other corresponding value function outputs
+        # which do trace back through the world model are only used for cont. actions
+        # for the actor loss (to compute the scaled value targets).
+        value_symlog_logits_HxB = dream_data[
+            "values_symlog_dreamed_logits_t0_to_HxBxT_wm_detached"
+        ]
+        # Unfold time rank and cut last time index to match value targets.
+        value_symlog_logits_t0_to_Hm1_B = value_symlog_logits_HxB.view(
+            [H, B, value_symlog_logits_HxB.shape[-1]]
+        )[:-1]
+
+        values_log_pred_Hm1_B = value_symlog_logits_t0_to_Hm1_B - torch.logsumexp(
+            value_symlog_logits_t0_to_Hm1_B, dim=-1, keepdim=True
+        )
+        # Multiply with two-hot targets and neg.
+        value_loss_two_hot_H_B = -torch.sum(
+            values_log_pred_Hm1_B * value_symlog_targets_two_hot_t0_to_Hm1_B, dim=-1
+        )
+
+        # Compute EMA regularization loss.
+        # Expected values (dreamed) from the EMA (slow critic) net.
+        value_symlog_ema_t0_to_Hm1_B = dream_data[
+            "v_symlog_dreamed_ema_t0_to_H_BxT"
+        ].detach()[:-1]
+        # Fold time rank (for two_hot'ing).
+        value_symlog_ema_HxB = value_symlog_ema_t0_to_Hm1_B.view(
+            -1,
+        )
+        value_symlog_ema_two_hot_HxB = two_hot(
+            value_symlog_ema_HxB, device=self._device
+        )
+        # Unfold time rank.
+        value_symlog_ema_two_hot_t0_to_Hm1_B = value_symlog_ema_two_hot_HxB.view(
+            [Hm1, B, value_symlog_ema_two_hot_HxB.shape[-1]]
+        )
+
+        # Compute ema regularizer loss.
+        # In the paper, it is not described how exactly to form this regularizer term
+        # and how to weigh it.
+        # So we follow Danijar's repo here:
+        # `reg = -dist.log_prob(sg(self.slow(traj).mean()))`
+        # with a weight of 1.0, where dist is the bucket'ized distribution output by the
+        # fast critic. sg=stop gradient; mean() -> use the expected EMA values.
+        # Multiply with two-hot targets and neg.
+        ema_regularization_loss_H_B = -torch.sum(
+            values_log_pred_Hm1_B * value_symlog_ema_two_hot_t0_to_Hm1_B, dim=-1
+        )
+
+        L_critic_H_B = value_loss_two_hot_H_B + ema_regularization_loss_H_B
+
+        # Mask out everything that goes beyond a predicted continue=False boundary.
+        L_critic_H_B *= dream_data["dream_loss_weights_t0_to_H_BxT"].detach()[:-1]
+
+        # Reduce over both H- (time) axis and B-axis (mean).
+        L_critic = L_critic_H_B.mean()
+
+        # Log important critic loss stats.
+        self.metrics.log_dict(
+            {
+                "CRITIC_L_total": L_critic,
+                "CRITIC_L_neg_logp_of_value_targets": torch.mean(
+                    value_loss_two_hot_H_B
+                ),
+                "CRITIC_L_slow_critic_regularization": torch.mean(
+                    ema_regularization_loss_H_B
+                ),
+            },
+            key=module_id,
+            window=1,  # <- single items (should not be mean/ema-reduced over time).
+        )
+        if config.report_individual_batch_item_stats:
+            # Log important critic loss stats.
+            self.metrics.log_dict(
+                {
+                    # Symlog'd value targets. Critic learns to predict symlog'd values.
+                    "VALUE_TARGETS_symlog_H_BxT": value_symlog_targets_t0_to_Hm1_B,
+                    # Critic loss terms.
+                    "CRITIC_L_total_H_BxT": L_critic_H_B,
+                    "CRITIC_L_neg_logp_of_value_targets_H_BxT": value_loss_two_hot_H_B,
+                    "CRITIC_L_slow_critic_regularization_H_BxT": (
+                        ema_regularization_loss_H_B
+                    ),
+                },
+                key=module_id,
+                window=1,  # <- single items (should not be mean/ema-reduced over time).
+            )
+
+        return L_critic
+
+    def _compute_value_targets(
+        self,
+        *,
+        config: DreamerV3Config,
+        rewards_t0_to_H_BxT: torch.Tensor,
+        intrinsic_rewards_t1_to_H_BxT: torch.Tensor,
+        continues_t0_to_H_BxT: torch.Tensor,
+        value_predictions_t0_to_H_BxT: torch.Tensor,
+    ) -> torch.Tensor:
+        """Helper method computing the value targets.
+
+        All args are (H, BxT, ...) and in non-symlog'd (real) reward space.
+        Non-symlog is important b/c log(a+b) != log(a) + log(b).
+        See [1] eq. 8 and 10.
+        Thus, targets are always returned in real (non-symlog'd space).
+        They need to be re-symlog'd before computing the critic loss from them (b/c the
+        critic produces predictions in symlog space).
+        Note that the original B and T ranks together form the new batch dimension
+        (folded into BxT) and the new time rank is the dream horizon (hence: [H, BxT]).
+
+        Variable names nomenclature:
+        `H`=1+horizon_H (start state + H steps dreamed),
+        `BxT`=batch_size * batch_length (meaning the original trajectory time rank has
+        been folded).
+
+        Rewards, continues, and value predictions are all of shape [t0-H, BxT]
+        (time-major), whereas returned targets are [t0 to H-1, B] (last timestep missing
+        b/c the target value equals vf prediction in that location anyways.
+
+        Args:
+            config: The DreamerV3Config to use.
+            rewards_t0_to_H_BxT: The reward predictor's predictions over the
+                dreamed trajectory t0 to H (and for the batch BxT).
+            intrinsic_rewards_t1_to_H_BxT: The predicted intrinsic rewards over the
+                dreamed trajectory t0 to H (and for the batch BxT).
+            continues_t0_to_H_BxT: The continue predictor's predictions over the
+                dreamed trajectory t0 to H (and for the batch BxT).
+            value_predictions_t0_to_H_BxT: The critic's value predictions over the
+                dreamed trajectory t0 to H (and for the batch BxT).
+
+        Returns:
+            The value targets in the shape: [t0toH-1, BxT]. Note that the last step (H)
+            does not require a value target as it matches the critic's value prediction
+            anyways.
+        """
+        # The first reward is irrelevant (not used for any VF target).
+        rewards_t1_to_H_BxT = rewards_t0_to_H_BxT[1:]
+        if intrinsic_rewards_t1_to_H_BxT is not None:
+            rewards_t1_to_H_BxT += intrinsic_rewards_t1_to_H_BxT
+
+        # In all the following, when building value targets for t=1 to T=H,
+        # exclude rewards & continues for t=1 b/c we don't need r1 or c1.
+        # The target (R1) for V1 is built from r2, c2, and V2/R2.
+        discount = continues_t0_to_H_BxT[1:] * config.gamma  # shape=[2-16, BxT]
+        Rs = [value_predictions_t0_to_H_BxT[-1]]  # Rs indices=[16]
+        intermediates = (
+            rewards_t1_to_H_BxT
+            + discount * (1 - config.gae_lambda) * value_predictions_t0_to_H_BxT[1:]
+        )
+        # intermediates.shape=[2-16, BxT]
+
+        # Loop through reversed timesteps (axis=1) from T+1 to t=2.
+        for t in reversed(range(discount.shape[0])):
+            Rs.append(intermediates[t] + discount[t] * config.gae_lambda * Rs[-1])
+
+        # Reverse time axis and cut the last entry (value estimate at very end
+        # cannot be learnt from as it's the same as the ... well ... value estimate).
+        targets_t0toHm1_BxT = torch.stack(list(reversed(Rs))[:-1], dim=0)
+        # targets.shape=[t0 to H-1,BxT]
+
+        return targets_t0toHm1_BxT
+
+    def _compute_scaled_value_targets(
+        self,
+        *,
+        module_id: ModuleID,
+        config: DreamerV3Config,
+        value_targets_t0_to_Hm1_BxT: torch.Tensor,
+        value_predictions_t0_to_Hm1_BxT: torch.Tensor,
+    ) -> torch.Tensor:
+        """Helper method computing the scaled value targets.
+
+        Args:
+            module_id: The module_id to compute value targets for.
+            config: The DreamerV3Config to use.
+            value_targets_t0_to_Hm1_BxT: The value targets computed by
+                `self._compute_value_targets` in the shape of (t0 to H-1, BxT)
+                and of type float32.
+            value_predictions_t0_to_Hm1_BxT: The critic's value predictions over the
+                dreamed trajectories (w/o the last timestep). The shape of this
+                tensor is (t0 to H-1, BxT) and the type is float32.
+
+        Returns:
+            The scaled value targets used by the actor for REINFORCE policy updates
+            (using scaled advantages). See [1] eq. 12 for more details.
+        """
+        actor = self.module[module_id].unwrapped().actor
+
+        value_targets_H_B = value_targets_t0_to_Hm1_BxT
+        value_predictions_H_B = value_predictions_t0_to_Hm1_BxT
+
+        # Compute S: [1] eq. 12.
+        Per_R_5 = torch.quantile(value_targets_H_B, 0.05)
+        Per_R_95 = torch.quantile(value_targets_H_B, 0.95)
+
+        # Update EMA values for 5 and 95 percentile, stored as actor network's
+        # parameters.
+        # 5 percentile
+        new_val_pct5 = torch.where(
+            torch.isnan(actor.ema_value_target_pct5),
+            # is NaN: Initial values: Just set.
+            Per_R_5,
+            # Later update (something already stored in EMA variable): Update EMA.
+            (
+                config.return_normalization_decay * actor.ema_value_target_pct5
+                + (1.0 - config.return_normalization_decay) * Per_R_5
+            ),
+        )
+        actor.ema_value_target_pct5.data = new_val_pct5
+        # 95 percentile
+        new_val_pct95 = torch.where(
+            # is NaN: Initial values: Just set.
+            torch.isnan(actor.ema_value_target_pct95),
+            # Later update (something already stored in EMA variable): Update EMA.
+            Per_R_95,
+            (
+                config.return_normalization_decay * actor.ema_value_target_pct95
+                + (1.0 - config.return_normalization_decay) * Per_R_95
+            ),
+        )
+        actor.ema_value_target_pct95.data = new_val_pct95
+
+        # [1] eq. 11 (first term).
+        offset = actor.ema_value_target_pct5
+        invscale = torch.clamp(
+            (actor.ema_value_target_pct95 - actor.ema_value_target_pct5),
+            min=1e-8,
+        )
+        scaled_value_targets_H_B = (value_targets_H_B - offset) / invscale
+        scaled_value_predictions_H_B = (value_predictions_H_B - offset) / invscale
+
+        # Return advantages.
+        return scaled_value_targets_H_B - scaled_value_predictions_H_B

rllib/algorithms/dreamerv3/torch/dreamerv3_torch_rl_module.py

@@ -0,0 +1,78 @@
+"""
+[1] Mastering Diverse Domains through World Models - 2023
+D. Hafner, J. Pasukonis, J. Ba, T. Lillicrap
+https://arxiv.org/pdf/2301.04104v1.pdf
+
+[2] Mastering Atari with Discrete World Models - 2021
+D. Hafner, T. Lillicrap, M. Norouzi, J. Ba
+https://arxiv.org/pdf/2010.02193.pdf
+"""
+from typing import Any, Dict
+
+import gymnasium as gym
+import torch
+
+from ray.rllib.algorithms.dreamerv3.dreamerv3_rl_module import (
+    ACTIONS_ONE_HOT,
+    DreamerV3RLModule,
+)
+from ray.rllib.core.columns import Columns
+from ray.rllib.core.rl_module.rl_module import RLModule
+from ray.rllib.core.rl_module.torch.torch_rl_module import TorchRLModule
+from ray.rllib.utils.annotations import override
+
+
+class DreamerV3TorchRLModule(TorchRLModule, DreamerV3RLModule):
+    """The torch-specific RLModule class for DreamerV3.
+
+    Serves mainly as a thin-wrapper around the `DreamerModel` (a torch.nn.Module) class.
+    """
+
+    framework = "torch"
+
+    @override(TorchRLModule)
+    def _forward_inference(self, batch: Dict[str, Any], **kwargs) -> Dict[str, Any]:
+        # Call the Dreamer-Model's forward_inference method and return a dict.
+        with torch.no_grad():
+            actions, next_state = self.dreamer_model.forward_inference(
+                observations=batch[Columns.OBS],
+                previous_states=batch[Columns.STATE_IN],
+                is_first=batch["is_first"],
+            )
+        return self._forward_inference_or_exploration_helper(batch, actions, next_state)
+
+    @override(TorchRLModule)
+    def _forward_exploration(self, batch: Dict[str, Any], **kwargs) -> Dict[str, Any]:
+        # Call the Dreamer-Model's forward_exploration method and return a dict.
+        with torch.no_grad():
+            actions, next_state = self.dreamer_model.forward_exploration(
+                observations=batch[Columns.OBS],
+                previous_states=batch[Columns.STATE_IN],
+                is_first=batch["is_first"],
+            )
+        return self._forward_inference_or_exploration_helper(batch, actions, next_state)
+
+    @override(RLModule)
+    def _forward_train(self, batch: Dict[str, Any], **kwargs):
+        # Call the Dreamer-Model's forward_train method and return its outputs as-is.
+        return self.dreamer_model.forward_train(
+            observations=batch[Columns.OBS],
+            actions=batch[Columns.ACTIONS],
+            is_first=batch["is_first"],
+        )
+
+    def _forward_inference_or_exploration_helper(self, batch, actions, next_state):
+        # Unfold time dimension.
+        shape = batch[Columns.OBS].shape
+        B, T = shape[0], shape[1]
+        actions = actions.view((B, T) + actions.shape[1:])
+
+        output = {
+            Columns.ACTIONS: actions,
+            ACTIONS_ONE_HOT: actions,
+            Columns.STATE_OUT: next_state,
+        }
+        # Undo one-hot actions?
+        if isinstance(self.action_space, gym.spaces.Discrete):
+            output[Columns.ACTIONS] = torch.argmax(actions, dim=-1)
+        return output

rllib/algorithms/dreamerv3/torch/models/actor_network.py

@@ -0,0 +1,178 @@
+"""
+[1] Mastering Diverse Domains through World Models - 2023
+D. Hafner, J. Pasukonis, J. Ba, T. Lillicrap
+https://arxiv.org/pdf/2301.04104v1.pdf
+"""
+import gymnasium as gym
+import numpy as np
+
+from ray.rllib.algorithms.dreamerv3.torch.models.components.mlp import MLP
+from ray.rllib.utils.framework import try_import_torch
+
+torch, nn = try_import_torch()
+
+
+class ActorNetwork(nn.Module):
+    """The `actor` (policy net) of DreamerV3.
+
+    Consists of a simple MLP for Discrete actions and two MLPs for cont. actions (mean
+    and stddev).
+    Also contains two scalar variables to keep track of the percentile-5 and
+    percentile-95 values of the computed value targets within a batch. This is used to
+    compute the "scaled value targets" for actor learning. These two variables decay
+    over time exponentially (see [1] for more details).
+    """
+
+    def __init__(
+        self,
+        *,
+        input_size: int,
+        model_size: str = "XS",
+        action_space: gym.Space,
+    ):
+        """Initializes an ActorNetwork instance.
+
+        Args:
+            input_size: The input size of the actor network.
+            model_size: The "Model Size" used according to [1] Appendinx B.
+                Use None for manually setting the different network sizes.
+            action_space: The action space the our environment used.
+        """
+        super().__init__()
+
+        self.input_size = input_size
+        self.model_size = model_size
+        self.action_space = action_space
+
+        # The EMA decay variables used for the [Percentile(R, 95%) - Percentile(R, 5%)]
+        # diff to scale value targets for the actor loss.
+        self.ema_value_target_pct5 = nn.Parameter(
+            torch.tensor(float("nan")), requires_grad=False
+        )
+        self.ema_value_target_pct95 = nn.Parameter(
+            torch.tensor(float("nan")), requires_grad=False
+        )
+
+        # For discrete actions, use a single MLP that computes logits.
+        if isinstance(self.action_space, gym.spaces.Discrete):
+            self.mlp = MLP(
+                input_size=self.input_size,
+                model_size=self.model_size,
+                output_layer_size=self.action_space.n,
+            )
+        # For cont. actions, use separate MLPs for Gaussian mean and stddev.
+        # TODO (sven): In the author's original code repo, this is NOT the case,
+        #  inputs are pushed through a shared MLP, then only the two output linear
+        #  layers are separate for std- and mean logits.
+        elif isinstance(action_space, gym.spaces.Box):
+            output_layer_size = np.prod(action_space.shape)
+            self.mlp = MLP(
+                input_size=self.input_size,
+                model_size=self.model_size,
+                output_layer_size=output_layer_size,
+            )
+            self.std_mlp = MLP(
+                input_size=self.input_size,
+                model_size=self.model_size,
+                output_layer_size=output_layer_size,
+            )
+        else:
+            raise ValueError(f"Invalid action space: {action_space}")
+
+    def forward(self, h, z, return_distr_params=False):
+        """Performs a forward pass through this policy network.
+
+        Args:
+            h: The deterministic hidden state of the sequence model. [B, dim(h)].
+            z: The stochastic discrete representations of the original
+                observation input. [B, num_categoricals, num_classes].
+            return_distr_params: Whether to return (as a second tuple item) the action
+                distribution parameter tensor created by the policy.
+        """
+        # Flatten last two dims of z.
+        assert len(z.shape) == 3
+        z_shape = z.shape
+        z = z.view(z_shape[0], -1)
+        assert len(z.shape) == 2
+        out = torch.cat([h, z], dim=-1)
+        # Send h-cat-z through MLP.
+        action_logits = self.mlp(out)
+
+        if isinstance(self.action_space, gym.spaces.Discrete):
+            action_probs = nn.functional.softmax(action_logits, dim=-1)
+
+            # Add the unimix weighting (1% uniform) to the probs.
+            # See [1]: "Unimix categoricals: We parameterize the categorical
+            # distributions for the world model representations and dynamics, as well as
+            # for the actor network, as mixtures of 1% uniform and 99% neural network
+            # output to ensure a minimal amount of probability mass on every class and
+            # thus keep log probabilities and KL divergences well behaved."
+            action_probs = 0.99 * action_probs + 0.01 * (1.0 / self.action_space.n)
+
+            # Danijar's code does: distr = [Distr class](logits=torch.log(probs)).
+            # Not sure why we don't directly use the already available probs instead.
+            action_logits = torch.log(action_probs)
+
+            # Distribution parameters are the log(probs) directly.
+            distr_params = action_logits
+            distr = self.get_action_dist_object(distr_params)
+
+            action = distr.sample().float().detach() + (
+                action_probs - action_probs.detach()
+            )
+
+        elif isinstance(self.action_space, gym.spaces.Box):
+            # Send h-cat-z through MLP to compute stddev logits for Normal dist
+            std_logits = self.std_mlp(out)
+            # minstd, maxstd taken from [1] from configs.yaml
+            minstd = 0.1
+            maxstd = 1.0
+
+            # Distribution parameters are the squashed std_logits and the tanh'd
+            # mean logits.
+            # squash std_logits from (-inf, inf) to (minstd, maxstd)
+            std_logits = (maxstd - minstd) * torch.sigmoid(std_logits + 2.0) + minstd
+            mean_logits = torch.tanh(action_logits)
+
+            distr_params = torch.cat([mean_logits, std_logits], dim=-1)
+            distr = self.get_action_dist_object(distr_params)
+
+            action = distr.rsample()
+
+        if return_distr_params:
+            return action, distr_params
+        return action
+
+    def get_action_dist_object(self, action_dist_params_T_B):
+        """Helper method to create an action distribution object from (T, B, ..) params.
+
+        Args:
+            action_dist_params_T_B: The time-major action distribution parameters.
+                This could be simply the logits (discrete) or a to-be-split-in-2
+                tensor for mean and stddev (continuous).
+
+        Returns:
+            The torch action distribution object, from which one can sample, compute
+            log probs, entropy, etc..
+        """
+        if isinstance(self.action_space, gym.spaces.Discrete):
+            # Create the distribution object using the unimix'd logits.
+            distr = torch.distributions.OneHotCategorical(logits=action_dist_params_T_B)
+
+        elif isinstance(self.action_space, gym.spaces.Box):
+            # Compute Normal distribution from action_logits and std_logits
+            loc, scale = torch.split(
+                action_dist_params_T_B,
+                action_dist_params_T_B.shape[-1] // 2,
+                dim=-1,
+            )
+            distr = torch.distributions.Normal(loc=loc, scale=scale)
+
+            # If action_space is a box with multiple dims, make individual dims
+            # independent.
+            distr = torch.distributions.Independent(distr, len(self.action_space.shape))
+
+        else:
+            raise ValueError(f"Action space {self.action_space} not supported!")
+
+        return distr

rllib/algorithms/dreamerv3/torch/models/components/__init__.py

@@ -0,0 +1,37 @@
+import numpy as np
+
+from ray.rllib.utils import force_list
+from ray.rllib.utils.framework import try_import_torch
+
+torch, _ = try_import_torch()
+
+
+def dreamerv3_normal_initializer(parameters):
+    """From Danijar Hafner's DreamerV3 JAX repo.
+
+    Used on any layer whenever the config for that layer has `winit="normal"`.
+
+    Note: Not identical with Glorot normal. Differs in the std computation
+    glorot_std = sqrt(2/(fanin+fanout))
+    this_std = sqrt(1/AVG(fanin, fanout)) / [somemagicnumber=0.879...]
+    """
+    for param in force_list(parameters):
+        if param.dim() > 1:
+            fanin, fanout = _fans(param.shape)
+            scale = 1.0 / np.mean([fanin, fanout])
+            std = np.sqrt(scale) / 0.87962566103423978
+            with torch.no_grad():
+                param.normal_(0, std)
+                param.clamp_(-2, 2)
+
+
+def _fans(shape):
+    if len(shape) == 0:
+        return 1, 1
+    elif len(shape) == 1:
+        return shape[0], shape[0]
+    elif len(shape) == 2:
+        return shape
+    else:
+        space = int(np.prod(shape[:-2]))
+        return shape[-2] * space, shape[-1] * space

rllib/algorithms/dreamerv3/torch/models/components/cnn_atari.py

@@ -0,0 +1,70 @@
+"""
+[1] Mastering Diverse Domains through World Models - 2023
+D. Hafner, J. Pasukonis, J. Ba, T. Lillicrap
+https://arxiv.org/pdf/2301.04104v1.pdf
+"""
+from typing import Optional
+
+from ray.rllib.algorithms.dreamerv3.torch.models.components import (
+    dreamerv3_normal_initializer,
+)
+from ray.rllib.algorithms.dreamerv3.utils import get_cnn_multiplier
+from ray.rllib.core.models.base import ENCODER_OUT
+from ray.rllib.core.models.configs import CNNEncoderConfig
+from ray.rllib.policy.sample_batch import SampleBatch
+from ray.rllib.utils.framework import try_import_torch
+
+torch, nn = try_import_torch()
+
+
+class CNNAtari(nn.Module):
+    """An image encoder mapping 64x64 RGB images via 4 CNN layers into a 1D space."""
+
+    def __init__(
+        self,
+        *,
+        model_size: str = "XS",
+        cnn_multiplier: Optional[int] = None,
+        gray_scaled: bool,
+    ):
+        """Initializes a CNNAtari instance.
+
+        Args:
+            model_size: The "Model Size" used according to [1] Appendix B.
+                Use None for manually setting the `cnn_multiplier`.
+            cnn_multiplier: Optional override for the additional factor used to multiply
+                the number of filters with each CNN layer. Starting with
+                1 * `cnn_multiplier` filters in the first CNN layer, the number of
+                filters then increases via `2*cnn_multiplier`, `4*cnn_multiplier`, till
+                `8*cnn_multiplier`.
+            gray_scaled: Whether the input is a gray-scaled image (1 color channel) or
+                not (3 RGB channels).
+        """
+        super().__init__()
+
+        cnn_multiplier = get_cnn_multiplier(model_size, override=cnn_multiplier)
+
+        config = CNNEncoderConfig(
+            input_dims=[64, 64, 1 if gray_scaled else 3],
+            cnn_filter_specifiers=[
+                [1 * cnn_multiplier, 4, 2],
+                [2 * cnn_multiplier, 4, 2],
+                [4 * cnn_multiplier, 4, 2],
+                [8 * cnn_multiplier, 4, 2],
+            ],
+            cnn_use_bias=False,
+            cnn_use_layernorm=True,
+            cnn_activation="silu",
+            cnn_kernel_initializer=dreamerv3_normal_initializer,
+            flatten_at_end=True,
+        )
+        self.cnn_stack = config.build(framework="torch")
+        self.output_size = config.output_dims
+
+    def forward(self, inputs):
+        """Performs a forward pass through the CNN Atari encoder.
+
+        Args:
+            inputs: The image inputs of shape (B, 64, 64, 3).
+        """
+        return self.cnn_stack({SampleBatch.OBS: inputs})[ENCODER_OUT]

rllib/algorithms/dreamerv3/torch/models/components/continue_predictor.py

@@ -0,0 +1,62 @@
+"""
+[1] Mastering Diverse Domains through World Models - 2023
+D. Hafner, J. Pasukonis, J. Ba, T. Lillicrap
+https://arxiv.org/pdf/2301.04104v1.pdf
+"""
+from ray.rllib.algorithms.dreamerv3.torch.models.components.mlp import MLP
+from ray.rllib.utils.framework import try_import_torch
+
+torch, nn = try_import_torch()
+
+
+class ContinuePredictor(nn.Module):
+    """The world-model network sub-component used to predict the `continue` flags .
+
+    Predicted continue flags are used to produce "dream data" to learn the policy in.
+
+    The continue flags are predicted via a linear output used to parameterize a
+    Bernoulli distribution, from which simply the mode is used (no stochastic
+    sampling!). In other words, if the sigmoid of the output of the linear layer is
+    >0.5, we predict a continuation of the episode, otherwise we predict an episode
+    terminal.
+    """
+
+    def __init__(self, *, input_size: int, model_size: str = "XS"):
+        """Initializes a ContinuePredictor instance.
+
+        Args:
+            input_size: The input size of the continue predictor.
+            model_size: The "Model Size" used according to [1] Appendinx B.
+                Determines the exact size of the underlying MLP.
+        """
+        super().__init__()
+
+        self.mlp = MLP(
+            input_size=input_size,
+            model_size=model_size,
+            output_layer_size=1,
+        )
+
+    def forward(self, h, z, return_distribution=False):
+        """Performs a forward pass through the continue predictor.
+
+        Args:
+            h: The deterministic hidden state of the sequence model. [B, dim(h)].
+            z: The stochastic discrete representations of the original
+                observation input. [B, num_categoricals, num_classes].
+            return_distribution: Whether to return (as a second tuple item) the
+                Bernoulli distribution object created by the underlying MLP.
+        """
+        z_shape = z.size()
+        z = z.view(z_shape[0], -1)
+
+        out = torch.cat([h, z], dim=-1)
+        out = self.mlp(out)
+        logits = out.squeeze(dim=-1)
+        bernoulli = torch.distributions.Bernoulli(logits=logits)
+        # Use the mode of the Bernoulli distribution (greedy, deterministic "sample").
+        continue_ = bernoulli.probs > 0.5
+
+        if return_distribution:
+            return continue_, bernoulli
+        return continue_

rllib/algorithms/dreamerv3/torch/models/components/conv_transpose_atari.py

@@ -0,0 +1,95 @@
+"""
+[1] Mastering Diverse Domains through World Models - 2023
+D. Hafner, J. Pasukonis, J. Ba, T. Lillicrap
+https://arxiv.org/pdf/2301.04104v1.pdf
+"""
+from typing import Optional
+
+from ray.rllib.algorithms.dreamerv3.torch.models.components import (
+    dreamerv3_normal_initializer,
+)
+from ray.rllib.algorithms.dreamerv3.utils import get_cnn_multiplier
+from ray.rllib.core.models.configs import CNNTransposeHeadConfig
+from ray.rllib.utils.framework import try_import_torch
+
+torch, nn = try_import_torch()
+
+
+class ConvTransposeAtari(nn.Module):
+    """A Conv2DTranspose decoder to generate Atari images from a latent space.
+
+    Wraps an initial single linear layer with a stack of 4 Conv2DTranspose layers (with
+    layer normalization) and a diag Gaussian, from which we then sample the final image.
+    """
+
+    def __init__(
+        self,
+        *,
+        input_size: int,
+        model_size: str = "XS",
+        cnn_multiplier: Optional[int] = None,
+        gray_scaled: bool,
+    ):
+        """Initializes a ConvTransposeAtari instance.
+
+        Args:
+            input_size: The input size of the ConvTransposeAtari network.
+            model_size: The "Model Size" used according to [1] Appendinx B.
+                Use None for manually setting the `cnn_multiplier`.
+            cnn_multiplier: Optional override for the additional factor used to multiply
+                the number of filters with each CNN transpose layer. Starting with
+                8 * `cnn_multiplier` filters in the first CNN transpose layer, the
+                number of filters then decreases via `4*cnn_multiplier`,
+                `2*cnn_multiplier`, till `1*cnn_multiplier`.
+            gray_scaled: Whether the last Conv2DTranspose layer's output has only 1
+                color channel (gray_scaled=True) or 3 RGB channels (gray_scaled=False).
+        """
+        super().__init__()
+
+        cnn_multiplier = get_cnn_multiplier(model_size, override=cnn_multiplier)
+        self.gray_scaled = gray_scaled
+        config = CNNTransposeHeadConfig(
+            input_dims=[input_size],
+            initial_image_dims=(4, 4, 8 * cnn_multiplier),
+            initial_dense_weights_initializer=dreamerv3_normal_initializer,
+            cnn_transpose_filter_specifiers=[
+                [4 * cnn_multiplier, 4, 2],
+                [2 * cnn_multiplier, 4, 2],
+                [1 * cnn_multiplier, 4, 2],
+                [1 if self.gray_scaled else 3, 4, 2],
+            ],
+            cnn_transpose_use_bias=False,
+            cnn_transpose_use_layernorm=True,
+            cnn_transpose_activation="silu",
+            cnn_transpose_kernel_initializer=dreamerv3_normal_initializer,
+        )
+        # Make sure the output dims match Atari.
+        # assert config.output_dims == (64, 64, 1 if self.gray_scaled else 3)
+
+        self._transpose_2d_head = config.build(framework="torch")
+
+    def forward(self, h, z):
+        """Performs a forward pass through the Conv2D transpose decoder.
+
+        Args:
+            h: The deterministic hidden state of the sequence model.
+            z: The sequence of stochastic discrete representations of the original
+                observation input. Note: `z` is not used for the dynamics predictor
+                model (which predicts z from h).
+        """
+        z_shape = z.size()
+        z = z.view(z_shape[0], -1)
+
+        input_ = torch.cat([h, z], dim=-1)
+
+        out = self._transpose_2d_head(input_)
+
+        # Interpret output as means of a diag-Gaussian with std=1.0:
+        # From [2]:
+        # "Distributions: The image predictor outputs the mean of a diagonal Gaussian
+        # likelihood with unit variance, ..."
+
+        # Reshape `out` for the diagonal multi-variate Gaussian (each pixel is its own
+        # independent (b/c diagonal co-variance matrix) variable).
+        loc = torch.reshape(out, (z_shape[0], -1))
+        return loc

rllib/algorithms/dreamerv3/torch/models/components/dynamics_predictor.py

@@ -0,0 +1,74 @@
+"""
+[1] Mastering Diverse Domains through World Models - 2023
+D. Hafner, J. Pasukonis, J. Ba, T. Lillicrap
+https://arxiv.org/pdf/2301.04104v1.pdf
+"""
+from typing import Optional
+
+from ray.rllib.algorithms.dreamerv3.torch.models.components import (
+    representation_layer,
+)
+from ray.rllib.algorithms.dreamerv3.torch.models.components.mlp import MLP
+from ray.rllib.algorithms.dreamerv3.utils import get_dense_hidden_units
+from ray.rllib.utils.framework import try_import_torch
+
+torch, nn = try_import_torch()
+
+
+class DynamicsPredictor(nn.Module):
+    """The dynamics (or "prior") network described in [1], producing prior z-states.
+
+    The dynamics net is used to:
+    - compute the initial z-state (from the tanh'd initial h-state variable) at the
+    beginning of a sequence.
+    - compute prior-z-states during dream data generation. Note that during dreaming,
+    no actual observations are available and thus no posterior z-states can be computed.
+    """
+
+    def __init__(
+        self,
+        *,
+        input_size: int,
+        model_size: str = "XS",
+        num_categoricals: Optional[int] = None,
+        num_classes_per_categorical: Optional[int] = None,
+    ):
+        """Initializes a DynamicsPredictor instance.
+
+        Args:
+            input_size: The input size of the dynamics predictor.
+            model_size: The "Model Size" used according to [1] Appendinx B.
+                Use None for manually setting the different parameters.
+            num_categoricals: Overrides the number of categoricals used in the z-states.
+                In [1], 32 is used for any model size.
+            num_classes_per_categorical: Overrides the number of classes within each
+                categorical used for the z-states. In [1], 32 is used for any model
+                dimension.
+        """
+        super().__init__()
+
+        self.mlp = MLP(
+            input_size=input_size,
+            num_dense_layers=1,
+            model_size=model_size,
+            output_layer_size=None,
+        )
+        representation_layer_input_size = get_dense_hidden_units(model_size)
+        self.representation_layer = representation_layer.RepresentationLayer(
+            input_size=representation_layer_input_size,
+            model_size=model_size,
+            num_categoricals=num_categoricals,
+            num_classes_per_categorical=num_classes_per_categorical,
+        )
+
+    def forward(self, h, return_z_probs=False):
+        """Performs a forward pass through the dynamics (or "prior") network.
+
+        Args:
+            h: The deterministic hidden state of the sequence model.
+            return_z_probs: Whether to return the probabilities for the categorical
+                distribution (in the shape of [B, num_categoricals, num_classes])
+                as a second return value.
+        """
+        out = self.mlp(h)
+        return self.representation_layer(out, return_z_probs=return_z_probs)

rllib/algorithms/dreamerv3/torch/models/components/mlp.py

@@ -0,0 +1,93 @@
+"""
+[1] Mastering Diverse Domains through World Models - 2023
+D. Hafner, J. Pasukonis, J. Ba, T. Lillicrap
+https://arxiv.org/pdf/2301.04104v1.pdf
+
+[2] Mastering Atari with Discrete World Models - 2021
+D. Hafner, T. Lillicrap, M. Norouzi, J. Ba
+https://arxiv.org/pdf/2010.02193.pdf
+"""
+from typing import Optional
+
+from ray.rllib.algorithms.dreamerv3.torch.models.components import (
+    dreamerv3_normal_initializer,
+)
+from ray.rllib.algorithms.dreamerv3.utils import (
+    get_dense_hidden_units,
+    get_num_dense_layers,
+)
+from ray.rllib.utils.framework import try_import_torch
+
+torch, nn = try_import_torch()
+
+
+class MLP(nn.Module):
+    """An MLP primitive used by several DreamerV3 components and described in [1] Fig 5.
+
+    MLP=multi-layer perceptron.
+
+    See Appendix B in [1] for the MLP sizes depending on the given `model_size`.
+    """
+
+    def __init__(
+        self,
+        *,
+        input_size: int,
+        model_size: str = "XS",
+        num_dense_layers: Optional[int] = None,
+        dense_hidden_units: Optional[int] = None,
+        output_layer_size=None,
+    ):
+        """Initializes an MLP instance.
+
+        Args:
+            input_size: The input size of the MLP.
+            model_size: The "Model Size" used according to [1] Appendinx B.
+                Use None for manually setting the different network sizes.
+            num_dense_layers: The number of hidden layers in the MLP. If None,
+                will use `model_size` and appendix B to figure out this value.
+            dense_hidden_units: The number of nodes in each hidden layer. If None,
+                will use `model_size` and appendix B to figure out this value.
+            output_layer_size: The size of an optional linear (no activation) output
+                layer. If None, no output layer will be added on top of the MLP dense
+                stack.
+        """
+        super().__init__()
+
+        self.output_size = None
+
+        num_dense_layers = get_num_dense_layers(model_size, override=num_dense_layers)
+        dense_hidden_units = get_dense_hidden_units(
+            model_size, override=dense_hidden_units
+        )
+
+        layers = []
+        for _ in range(num_dense_layers):
+            # In this order: layer, normalization, activation.
+            linear = nn.Linear(input_size, dense_hidden_units, bias=False)
+            # Use same initializers as the Author in their JAX repo.
+            dreamerv3_normal_initializer(linear.weight)
+            layers.append(linear)
+            layers.append(nn.LayerNorm(dense_hidden_units, eps=0.001))
+            layers.append(nn.SiLU())
+            input_size = dense_hidden_units
+            self.output_size = (dense_hidden_units,)
+
+        self.output_layer = None
+        if output_layer_size:
+            linear = nn.Linear(input_size, output_layer_size, bias=True)
+            # Use same initializers as the Author in their JAX repo.
+            dreamerv3_normal_initializer(linear.weight)
+            nn.init.zeros_(linear.bias)
+            layers.append(linear)
+            self.output_size = (output_layer_size,)
+
+        self._net = nn.Sequential(*layers)
+
+    def forward(self, input_):
+        """Performs a forward pass through this MLP.
+
+        Args:
+            input_: The input tensor for the MLP dense stack.
+        """
+        return self._net(input_)

rllib/algorithms/dreamerv3/torch/models/components/representation_layer.py

@@ -0,0 +1,133 @@
+"""
+[1] Mastering Diverse Domains through World Models - 2023
+D. Hafner, J. Pasukonis, J. Ba, T. Lillicrap
+https://arxiv.org/pdf/2301.04104v1.pdf
+
+[2] Mastering Atari with Discrete World Models - 2021
+D. Hafner, T. Lillicrap, M. Norouzi, J. Ba
+https://arxiv.org/pdf/2010.02193.pdf
+"""
+from typing import Optional
+
+from ray.rllib.algorithms.dreamerv3.torch.models.components import (
+    dreamerv3_normal_initializer,
+)
+from ray.rllib.algorithms.dreamerv3.utils import (
+    get_num_z_categoricals,
+    get_num_z_classes,
+)
+from ray.rllib.utils.framework import try_import_torch
+
+torch, nn = try_import_torch()
+if torch:
+    F = nn.functional
+
+
+class RepresentationLayer(nn.Module):
+    """A representation (z-state) generating layer.
+
+    The value for z is the result of sampling from a categorical distribution with
+    shape B x `num_classes`. So a computed z-state consists of `num_categoricals`
+    one-hot vectors, each of size `num_classes_per_categorical`.
+    """
+
+    def __init__(
+        self,
+        *,
+        input_size: int,
+        model_size: str = "XS",
+        num_categoricals: Optional[int] = None,
+        num_classes_per_categorical: Optional[int] = None,
+    ):
+        """Initializes a RepresentationLayer instance.
+
+        Args:
+            input_size: The input size of the representation layer.
+            model_size: The "Model Size" used according to [1] Appendinx B.
+                Use None for manually setting the different parameters.
+            num_categoricals: Overrides the number of categoricals used in the z-states.
+                In [1], 32 is used for any model size.
+            num_classes_per_categorical: Overrides the number of classes within each
+                categorical used for the z-states. In [1], 32 is used for any model
+                dimension.
+        """
+        self.num_categoricals = get_num_z_categoricals(
+            model_size, override=num_categoricals
+        )
+        self.num_classes_per_categorical = get_num_z_classes(
+            model_size, override=num_classes_per_categorical
+        )
+
+        super().__init__()
+
+        self.z_generating_layer = nn.Linear(
+            input_size,
+            self.num_categoricals * self.num_classes_per_categorical,
+            bias=True,
+        )
+        # Use same initializers as the Author in their JAX repo.
+        dreamerv3_normal_initializer(self.z_generating_layer.weight)
+
+    def forward(self, inputs, return_z_probs=False):
+        """Produces a discrete, differentiable z-sample from some 1D input tensor.
+
+        Pushes the input_ tensor through our dense layer, which outputs
+        32(B=num categoricals)*32(c=num classes) logits. Logits are used to:
+
+        1) sample stochastically
+        2) compute probs (via softmax)
+        3) make sure the sampling step is differentiable (see [2] Algorithm 1):
+            sample=one_hot(draw(logits))
+            probs=softmax(logits)
+            sample=sample + probs - stop_grad(probs)
+            -> Now sample has the gradients of the probs.
+
+        Args:
+            inputs: The input to our z-generating layer. This might be a) the combined
+                (concatenated) outputs of the (image?) encoder + the last hidden
+                deterministic state, or b) the output of the dynamics predictor MLP
+                network.
+            return_z_probs: Whether to return the probabilities for the categorical
+                distribution (in the shape of [B, num_categoricals, num_classes])
+                as a second return value.
+        """
+        # Compute the logits (no activation) for our `num_categoricals` Categorical
+        # distributions (with `num_classes_per_categorical` classes each).
+        logits = self.z_generating_layer(inputs)
+        # Reshape the logits to [B, num_categoricals, num_classes]
+        logits = logits.reshape(
+            -1, self.num_categoricals, self.num_classes_per_categorical
+        )
+        # Compute the probs (based on logits) via softmax.
+        probs = F.softmax(logits, dim=-1)
+        # Add the unimix weighting (1% uniform) to the probs.
+        # See [1]: "Unimix categoricals: We parameterize the categorical distributions
+        # for the world model representations and dynamics, as well as for the actor
+        # network, as mixtures of 1% uniform and 99% neural network output to ensure
+        # a minimal amount of probability mass on every class and thus keep log
+        # probabilities and KL divergences well behaved."
+        probs = 0.99 * probs + 0.01 * (1.0 / self.num_classes_per_categorical)
+
+        # Danijar's code does: distr = [Distr class](logits=torch.log(probs)).
+        # Not sure why we don't directly use the already available probs instead.
+        logits = torch.log(probs)
+
+        # Create the distribution object using the unimix'd logits.
+        distribution = torch.distributions.Independent(
+            torch.distributions.OneHotCategorical(logits=logits),
+            reinterpreted_batch_ndims=1,
+        )
+
+        # Draw a one-hot sample (B, num_categoricals, num_classes).
+        sample = distribution.sample()
+        # Make sure we can take gradients "straight-through" the sampling step
+        # by adding the probs and subtracting the sg(probs). Note that `sample`
+        # does not have any gradients as it's the result of a Categorical sample step,
+        # which is non-differentiable (other than say a Gaussian sample step).
+        # [1] "The representations are sampled from a vector of softmax distributions
+        # and we take straight-through gradients through the sampling step."
+        # [2] Algorithm 1.
+        differentiable_sample = sample.detach() + probs - probs.detach()
+        if return_z_probs:
+            return differentiable_sample, probs
+        return differentiable_sample

rllib/algorithms/dreamerv3/torch/models/components/reward_predictor.py

@@ -0,0 +1,86 @@
+"""
+[1] Mastering Diverse Domains through World Models - 2023
+D. Hafner, J. Pasukonis, J. Ba, T. Lillicrap
+https://arxiv.org/pdf/2301.04104v1.pdf
+"""
+from ray.rllib.algorithms.dreamerv3.torch.models.components import (
+    reward_predictor_layer,
+)
+from ray.rllib.algorithms.dreamerv3.torch.models.components.mlp import MLP
+from ray.rllib.algorithms.dreamerv3.utils import get_dense_hidden_units
+from ray.rllib.utils.framework import try_import_torch
+
+torch, nn = try_import_torch()
+
+
+class RewardPredictor(nn.Module):
+    """Wrapper of MLP and RewardPredictorLayer to predict rewards for the world model.
+
+    Predicted rewards are used to produce "dream data" to learn the policy in.
+    """
+
+    def __init__(
+        self,
+        *,
+        input_size: int,
+        model_size: str = "XS",
+        num_buckets: int = 255,
+        lower_bound: float = -20.0,
+        upper_bound: float = 20.0,
+    ):
+        """Initializes a RewardPredictor instance.
+
+        Args:
+            input_size: The input size of the reward predictor.
+            model_size: The "Model Size" used according to [1] Appendinx B.
+                Determines the exact size of the underlying MLP.
+            num_buckets: The number of buckets to create. Note that the number of
+                possible symlog'd outcomes from the used distribution is
+                `num_buckets` + 1:
+                lower_bound --bucket-- o[1] --bucket-- o[2] ... --bucket-- upper_bound
+                o=outcomes
+                lower_bound=o[0]
+                upper_bound=o[num_buckets]
+            lower_bound: The symlog'd lower bound for a possible reward value.
+                Note that a value of -20.0 here already allows individual (actual env)
+                rewards to be as low as -400M. Buckets will be created between
+                `lower_bound` and `upper_bound`.
+            upper_bound: The symlog'd upper bound for a possible reward value.
+                Note that a value of +20.0 here already allows individual (actual env)
+                rewards to be as high as 400M. Buckets will be created between
+                `lower_bound` and `upper_bound`.
+        """
+        super().__init__()
+
+        self.mlp = MLP(
+            input_size=input_size,
+            model_size=model_size,
+            output_layer_size=None,
+        )
+        reward_predictor_input_size = get_dense_hidden_units(model_size)
+        self.reward_layer = reward_predictor_layer.RewardPredictorLayer(
+            input_size=reward_predictor_input_size,
+            num_buckets=num_buckets,
+            lower_bound=lower_bound,
+            upper_bound=upper_bound,
+        )
+
+    def forward(self, h, z, return_logits=False):
+        """Computes the expected reward using N equal sized buckets of possible values.
+
+        Args:
+            h: The deterministic hidden state of the sequence model. [B, dim(h)].
+            z: The stochastic discrete representations of the original
+                observation input. [B, num_categoricals, num_classes].
+            return_logits: Whether to return the logits over the reward buckets
+                as a second return value (besides the expected reward).
+        """
+        # Flatten last two dims of z.
+        z_shape = z.shape
+        z = z.view(z_shape[0], -1)
+        out = torch.cat([h, z], dim=-1)
+        # Send h-cat-z through MLP.
+        out = self.mlp(out)
+        # Return a) mean reward OR b) a tuple: (mean reward, logits over the reward
+        # buckets).
+        return self.reward_layer(out, return_logits=return_logits)

rllib/algorithms/dreamerv3/torch/models/components/reward_predictor_layer.py

@@ -0,0 +1,106 @@
+"""
+[1] Mastering Diverse Domains through World Models - 2023
+D. Hafner, J. Pasukonis, J. Ba, T. Lillicrap
+https://arxiv.org/pdf/2301.04104v1.pdf
+
+[2] Mastering Atari with Discrete World Models - 2021
+D. Hafner, T. Lillicrap, M. Norouzi, J. Ba
+https://arxiv.org/pdf/2010.02193.pdf
+"""
+
+from ray.rllib.utils.framework import try_import_torch
+
+torch, nn = try_import_torch()
+if torch:
+    F = nn.functional
+
+
+class RewardPredictorLayer(nn.Module):
+    """A layer outputting reward predictions using K bins and two-hot encoding.
+
+    This layer is used in two models in DreamerV3: The reward predictor of the world
+    model and the value function. K is 255 by default (see [1]) and doesn't change
+    with the model size.
+
+    Possible predicted reward/values range from symexp(-20.0) to symexp(20.0), which
+    should cover any possible environment. Outputs of this layer are generated by
+    generating logits/probs via a single linear layer, then interpreting the probs
+    as weights for a weighted average of the different possible reward (binned) values.
+    """
+
+    def __init__(
+        self,
+        *,
+        input_size: int,
+        num_buckets: int = 255,
+        lower_bound: float = -20.0,
+        upper_bound: float = 20.0,
+    ):
+        """Initializes a RewardPredictorLayer instance.
+
+        Args:
+            input_size: The input size of the reward predictor layer.
+            num_buckets: The number of buckets to create. Note that the number of
+                possible symlog'd outcomes from the used distribution is
+                `num_buckets` + 1:
+                lower_bound --bucket-- o[1] --bucket-- o[2] ... --bucket-- upper_bound
+                o=outcomes
+                lower_bound=o[0]
+                upper_bound=o[num_buckets]
+            lower_bound: The symlog'd lower bound for a possible reward value.
+                Note that a value of -20.0 here already allows individual (actual env)
+                rewards to be as low as -400M. Buckets will be created between
+                `lower_bound` and `upper_bound`.
+            upper_bound: The symlog'd upper bound for a possible reward value.
+                Note that a value of +20.0 here already allows individual (actual env)
+                rewards to be as high as 400M. Buckets will be created between
+                `lower_bound` and `upper_bound`.
+        """
+        self.num_buckets = num_buckets
+        super().__init__()
+
+        self.lower_bound = lower_bound
+        self.upper_bound = upper_bound
+        self.reward_buckets_layer = nn.Linear(
+            in_features=input_size, out_features=self.num_buckets, bias=True
+        )
+        nn.init.zeros_(self.reward_buckets_layer.weight)
+        nn.init.zeros_(self.reward_buckets_layer.bias)
+        # self.reward_buckets_layer.weight.data.fill_(0.0)
+        # self.reward_buckets_layer.bias.data.fill_(0.0)
+
+    def forward(self, inputs, return_logits=False):
+        """Computes the expected reward using N equal sized buckets of possible values.
+
+        Args:
+            inputs: The input tensor for the layer, which computes the reward bucket
+                weights (logits). [B, dim].
+            return_logits: Whether to return the logits over the reward buckets
+                as a second return value (besides the expected reward).
+
+        Returns:
+            The expected reward OR a tuple consisting of the expected reward and the
+            torch `FiniteDiscrete` distribution object. To get the individual bucket
+            probs, do `[FiniteDiscrete object].probs`.
+        """
+        # Compute the `num_buckets` weights.
+        logits = self.reward_buckets_layer(inputs)
+
+        # Compute the expected(!) reward using the formula:
+        # `softmax(Linear(x))` [vectordot] `possible_outcomes`, where
+        # `possible_outcomes` is the even-spaced (binned) encoding of all possible
+        # symexp'd reward/values.
+        probs = F.softmax(logits, dim=-1)
+        possible_outcomes = torch.linspace(
+            self.lower_bound, self.upper_bound, self.num_buckets, device=logits.device
+        )
+        # probs=possible_outcomes=[B, `num_buckets`]
+
+        # Simple vector dot product (over last dim) to get the mean reward
+        # weighted sum, where all weights sum to 1.0.
+        expected_rewards = torch.sum(probs * possible_outcomes, dim=-1)
+        # expected_rewards=[B]
+
+        if return_logits:
+            return expected_rewards, logits
+        return expected_rewards

rllib/algorithms/dreamerv3/torch/models/components/sequence_model.py

@@ -0,0 +1,132 @@
+"""
+[1] Mastering Diverse Domains through World Models - 2023
+D. Hafner, J. Pasukonis, J. Ba, T. Lillicrap
+https://arxiv.org/pdf/2301.04104v1.pdf
+"""
+from typing import Optional
+
+import gymnasium as gym
+
+from ray.rllib.algorithms.dreamerv3.torch.models.components import (
+    dreamerv3_normal_initializer,
+)
+from ray.rllib.algorithms.dreamerv3.torch.models.components.mlp import MLP
+from ray.rllib.algorithms.dreamerv3.utils import get_dense_hidden_units, get_gru_units
+from ray.rllib.utils.framework import try_import_torch
+
+torch, nn = try_import_torch()
+
+
+class SequenceModel(nn.Module):
+    """The "sequence model" of the RSSM, computing ht+1 given (ht, zt, at).
+
+    Note: The "internal state" always consists of:
+    The actions `a` (initially, this is a zeroed-out action), `h`-states (deterministic,
+    continuous), and `z`-states (stochastic, discrete).
+    There are two versions of z-states: "posterior" for world model training and "prior"
+    for creating the dream data.
+
+    Initial internal state values (`a`, `h`, and `z`) are used where ever a new episode
+    starts within a batch row OR at the beginning of each train batch's B rows,
+    regardless of whether there was an actual episode boundary or not. Thus, internal
+    states are not required to be stored in or retrieved from the replay buffer AND
+    retrieved batches from the buffer must not be zero padded.
+
+    Initial `a` is the zero "one hot" action, e.g. [0.0, 0.0] for Discrete(2), initial
+    `h` is a separate learned variable, and initial `z` are computed by the "dynamics"
+    (or "prior") net, using only the initial-h state as input.
+
+    The GRU in this SequenceModel always produces the next h-state, then.
+    """
+
+    def __init__(
+        self,
+        *,
+        input_size: int,
+        model_size: str = "XS",
+        action_space: gym.Space,
+        num_gru_units: Optional[int] = None,
+    ):
+        """Initializes a SequenceModel instance.
+
+        Args:
+            input_size: The input size of the pre-layer (Dense) of the sequence model.
+            model_size: The "Model Size" used according to [1] Appendinx B.
+                Use None for manually setting the number of GRU units used.
+            action_space: The action space the our environment used.
+            num_gru_units: Overrides the number of GRU units (dimension of the h-state).
+                If None, use the value given through `model_size`
+                (see [1] Appendix B).
+        """
+        super().__init__()
+
+        num_gru_units = get_gru_units(model_size, override=num_gru_units)
+        self.action_space = action_space
+
+        # In Danijar's code, there is an additional layer (units=[model_size])
+        # prior to the GRU (but always only with 1 layer), which is not mentioned in
+        # the paper.
+        # In Danijar's code, this layer is called: `img_in`.
+        self.pre_gru_layer = MLP(
+            input_size=input_size,
+            num_dense_layers=1,
+            model_size=model_size,
+            output_layer_size=None,
+        )
+        gru_input_size = get_dense_hidden_units(model_size)
+
+        # Use a custom GRU implementation w/ Normal init, layernorm, no bias
+        # (just like Danijar's GRU).
+        # In Danijar's code, this layer is called: `gru`.
+        self.gru_unit = DreamerV3GRU(input_size=gru_input_size, cell_size=num_gru_units)
+
+    def forward(self, a, h, z):
+        """
+
+        Args:
+            a: The previous action (already one-hot'd if applicable). (B, ...).
+            h: The previous deterministic hidden state of the sequence model.
+                (B, num_gru_units)
+            z: The previous stochastic discrete representations of the original
+                observation input. (B, num_categoricals, num_classes_per_categorical).
+        """
+        # Flatten last two dims of z.
+        z_shape = z.shape
+        z = z.view(z_shape[0], -1)
+        out = torch.cat([z, a], dim=-1)
+        # Pass through pre-GRU layer.
+        out = self.pre_gru_layer(out)
+        # Pass through GRU (add extra time axis at 0 to make time-major).
+        h_next, _ = self.gru_unit(out.unsqueeze(0), h.unsqueeze(0))
+        h_next = h_next.squeeze(0)  # Remove extra time dimension again.
+        # Return the GRU's output (the next h-state).
+        return h_next
+
+
+class DreamerV3GRU(nn.Module):
+    """Analogous to Danijar's JAX GRU unit code."""
+
+    def __init__(self, input_size, cell_size):
+        super().__init__()
+        self.cell_size = cell_size
+        self.output_size = 3 * self.cell_size
+
+        self.linear = nn.Linear(
+            input_size + self.cell_size,
+            self.output_size,
+            bias=False,
+        )
+        dreamerv3_normal_initializer(list(self.linear.parameters()))
+
+        self.layer_norm = nn.LayerNorm(self.output_size, eps=0.001)
+
+    def forward(self, x, h):
+        x = torch.cat([h, x], dim=-1)
+        x = self.linear(x)
+        x = self.layer_norm(x)
+        reset, cand, update = torch.split(x, self.cell_size, dim=-1)
+        reset = torch.sigmoid(reset)
+        cand = torch.tanh(reset * cand)
+        update = torch.sigmoid(update - 1)
+        h = update * cand + (1 - update) * h
+        return h, h

rllib/algorithms/dreamerv3/torch/models/components/vector_decoder.py

@@ -0,0 +1,68 @@
+"""
+[1] Mastering Diverse Domains through World Models - 2023
+D. Hafner, J. Pasukonis, J. Ba, T. Lillicrap
+https://arxiv.org/pdf/2301.04104v1.pdf
+"""
+import gymnasium as gym
+
+from ray.rllib.algorithms.dreamerv3.torch.models.components.mlp import MLP
+from ray.rllib.utils.framework import try_import_torch
+
+torch, nn = try_import_torch()
+
+
+class VectorDecoder(nn.Module):
+    """A simple vector decoder to reproduce non-image (1D vector) observations.
+
+    Wraps an MLP for mean parameter computations and a Gaussian distribution,
+    from which we then sample using these mean values and a fixed stddev of 1.0.
+    """
+
+    def __init__(
+        self,
+        *,
+        input_size: int,
+        model_size: str = "XS",
+        observation_space: gym.Space,
+    ):
+        """Initializes a VectorDecoder instance.
+
+        Args:
+            input_size: The input size of the vector decoder.
+            model_size: The "Model Size" used according to [1] Appendinx B.
+                Determines the exact size of the underlying MLP.
+            observation_space: The observation space to decode back into. This must
+                be a Box of shape (d,), where d >= 1.
+        """
+        super().__init__()
+
+        assert (
+            isinstance(observation_space, gym.spaces.Box)
+            and len(observation_space.shape) == 1
+        )
+
+        self.mlp = MLP(
+            input_size=input_size,
+            model_size=model_size,
+            output_layer_size=observation_space.shape[0],
+        )
+
+    def forward(self, h, z):
+        """Performs a forward pass through the vector encoder.
+
+        Args:
+            h: The deterministic hidden state of the sequence model. [B, dim(h)].
+            z: The stochastic discrete representations of the original
+                observation input. [B, num_categoricals, num_classes].
+        """
+        # Flatten last two dims of z.
+        assert len(z.shape) == 3
+        z_shape = z.shape
+        z = z.view(z_shape[0], -1)
+        assert len(z.shape) == 2
+        out = torch.cat([h, z], dim=-1)
+        # Send h-cat-z through MLP to get mean values of diag gaussian.
+        loc = self.mlp(out)
+
+        # Return only the predicted observations (mean, no sample).
+        return loc

rllib/algorithms/dreamerv3/torch/models/critic_network.py

@@ -0,0 +1,168 @@
+"""
+[1] Mastering Diverse Domains through World Models - 2023
+D. Hafner, J. Pasukonis, J. Ba, T. Lillicrap
+https://arxiv.org/pdf/2301.04104v1.pdf
+"""
+from ray.rllib.algorithms.dreamerv3.torch.models.components import (
+    reward_predictor_layer,
+)
+from ray.rllib.algorithms.dreamerv3.torch.models.components.mlp import MLP
+from ray.rllib.algorithms.dreamerv3.utils import get_dense_hidden_units
+from ray.rllib.utils.framework import try_import_torch
+
+torch, nn = try_import_torch()
+
+
+class CriticNetwork(nn.Module):
+    """The critic network described in [1], predicting values for policy learning.
+
+    Contains a copy of itself (EMA net) for weight regularization.
+    The EMA net is updated after each train step via EMA (using the `ema_decay`
+    parameter and the actual critic's weights). The EMA net is NOT used for target
+    computations (we use the actual critic for that), its only purpose is to compute a
+    weights regularizer term for the critic's loss such that the actual critic does not
+    move too quickly.
+    """
+
+    def __init__(
+        self,
+        *,
+        input_size: int,
+        model_size: str = "XS",
+        num_buckets: int = 255,
+        lower_bound: float = -20.0,
+        upper_bound: float = 20.0,
+        ema_decay: float = 0.98,
+    ):
+        """Initializes a CriticNetwork instance.
+
+        Args:
+            input_size: The input size of the critic network.
+            model_size: The "Model Size" used according to [1] Appendinx B.
+               Use None for manually setting the different network sizes.
+            num_buckets: The number of buckets to create. Note that the number of
+                possible symlog'd outcomes from the used distribution is
+                `num_buckets` + 1:
+                lower_bound --bucket-- o[1] --bucket-- o[2] ... --bucket-- upper_bound
+                o=outcomes
+                lower_bound=o[0]
+                upper_bound=o[num_buckets]
+            lower_bound: The symlog'd lower bound for a possible reward value.
+                Note that a value of -20.0 here already allows individual (actual env)
+                rewards to be as low as -400M. Buckets will be created between
+                `lower_bound` and `upper_bound`.
+            upper_bound: The symlog'd upper bound for a possible reward value.
+                Note that a value of +20.0 here already allows individual (actual env)
+                rewards to be as high as 400M. Buckets will be created between
+                `lower_bound` and `upper_bound`.
+            ema_decay: The weight to use for updating the weights of the critic's copy
+                vs the actual critic. After each training update, the EMA copy of the
+                critic gets updated according to:
+                ema_net=(`ema_decay`*ema_net) + (1.0-`ema_decay`)*critic_net
+                The EMA copy of the critic is used inside the critic loss function only
+                to produce a regularizer term against the current critic's weights, NOT
+                to compute any target values.
+        """
+        super().__init__()
+
+        self.input_size = input_size
+        self.model_size = model_size
+        self.ema_decay = ema_decay
+
+        # "Fast" critic network(s) (mlp + reward-pred-layer). This is the network
+        # we actually train with our critic loss.
+        # IMPORTANT: We also use this to compute the return-targets, BUT we regularize
+        # the critic loss term such that the weights of this fast critic stay close
+        # to the EMA weights (see below).
+        self.mlp = MLP(
+            input_size=self.input_size,
+            model_size=self.model_size,
+            output_layer_size=None,
+        )
+        reward_predictor_input_size = get_dense_hidden_units(self.model_size)
+        self.return_layer = reward_predictor_layer.RewardPredictorLayer(
+            input_size=reward_predictor_input_size,
+            num_buckets=num_buckets,
+            lower_bound=lower_bound,
+            upper_bound=upper_bound,
+        )
+
+        # Weights-EMA (EWMA) containing networks for critic loss (similar to a
+        # target net, BUT not used to compute anything, just for the
+        # weights regularizer term inside the critic loss).
+        self.mlp_ema = MLP(
+            input_size=self.input_size,
+            model_size=self.model_size,
+            output_layer_size=None,
+        )
+        self.return_layer_ema = reward_predictor_layer.RewardPredictorLayer(
+            input_size=reward_predictor_input_size,
+            num_buckets=num_buckets,
+            lower_bound=lower_bound,
+            upper_bound=upper_bound,
+        )
+
+    def forward(self, h, z, return_logits=False, use_ema=False):
+        """Performs a forward pass through the critic network.
+
+        Args:
+            h: The deterministic hidden state of the sequence model. [B, dim(h)].
+            z: The stochastic discrete representations of the original
+                observation input. [B, num_categoricals, num_classes].
+            return_logits: Whether also return (as a second tuple item) the logits
+                computed by the binned return layer (instead of only the value itself).
+            use_ema: Whether to use the EMA-copy of the critic instead of the actual
+                critic to perform this computation.
+        """
+        # Flatten last two dims of z.
+        assert len(z.shape) == 3
+        z_shape = z.shape
+        z = z.view(z_shape[0], -1)
+        assert len(z.shape) == 2
+        out = torch.cat([h, z], dim=-1)
+
+        if not use_ema:
+            # Send h-cat-z through MLP.
+            out = self.mlp(out)
+            # Return expected return OR (expected return, probs of bucket values).
+            return self.return_layer(out, return_logits=return_logits)
+        else:
+            out = self.mlp_ema(out)
+            return self.return_layer_ema(out, return_logits=return_logits)
+
+    def init_ema(self) -> None:
+        """Initializes the EMA-copy of the critic from the critic's weights.
+
+        After calling this method, the two networks have identical weights and the EMA
+        net will be non-trainable.
+        """
+        for param_ema, param in zip(self.mlp_ema.parameters(), self.mlp.parameters()):
+            param_ema.data.copy_(param.data)
+            # Make all EMA parameters non-trainable.
+            param_ema.requires_grad = False
+            assert param_ema.grad is None
+
+        for param_ema, param in zip(
+            self.return_layer_ema.parameters(), self.return_layer.parameters()
+        ):
+            param_ema.data.copy_(param.data)
+            # Make all EMA parameters non-trainable.
+            param_ema.requires_grad = False
+            assert param_ema.grad is None
+
+    def update_ema(self) -> None:
+        """Updates the EMA-copy of the critic according to the update formula:
+
+        ema_net=(`ema_decay`*ema_net) + (1.0-`ema_decay`)*critic_net
+        """
+        for param_ema, param in zip(self.mlp_ema.parameters(), self.mlp.parameters()):
+            param_ema.data.mul_(self.ema_decay).add_(
+                (1.0 - self.ema_decay) * param.data
+            )
+
+        for param_ema, param in zip(
+            self.return_layer_ema.parameters(), self.return_layer.parameters()
+        ):
+            param_ema.data.mul_(self.ema_decay).add_(
+                (1.0 - self.ema_decay) * param.data
+            )

rllib/algorithms/dreamerv3/torch/models/dreamer_model.py

@@ -0,0 +1,518 @@
+"""
+[1] Mastering Diverse Domains through World Models - 2023
+D. Hafner, J. Pasukonis, J. Ba, T. Lillicrap
+https://arxiv.org/pdf/2301.04104v1.pdf
+"""
+import re
+
+import gymnasium as gym
+import numpy as np
+
+from ray.rllib.algorithms.dreamerv3.torch.models.actor_network import ActorNetwork
+from ray.rllib.algorithms.dreamerv3.torch.models.critic_network import CriticNetwork
+from ray.rllib.algorithms.dreamerv3.torch.models.world_model import WorldModel
+from ray.rllib.utils.framework import try_import_torch
+from ray.rllib.utils.torch_utils import inverse_symlog
+
+torch, nn = try_import_torch()
+
+
+class DreamerModel(nn.Module):
+    """The main PyTorch model containing all necessary components for DreamerV3.
+
+    Includes:
+    - The world model with encoder, decoder, sequence-model (RSSM), dynamics
+    (generates prior z-state), and "posterior" model (generates posterior z-state).
+    Predicts env dynamics and produces dreamed trajectories for actor- and critic
+    learning.
+    - The actor network (policy).
+    - The critic network for value function prediction.
+    """
+
+    def __init__(
+        self,
+        *,
+        model_size: str = "XS",
+        action_space: gym.Space,
+        world_model: WorldModel,
+        actor: ActorNetwork,
+        critic: CriticNetwork,
+        use_curiosity: bool = False,
+        intrinsic_rewards_scale: float = 0.1,
+    ):
+        """Initializes a DreamerModel instance.
+
+        Args:
+             model_size: The "Model Size" used according to [1] Appendinx B.
+                Use None for manually setting the different network sizes.
+             action_space: The action space the our environment used.
+             world_model: The WorldModel component.
+             actor: The ActorNetwork component.
+             critic: The CriticNetwork component.
+        """
+        super().__init__()
+
+        self.model_size = model_size
+        self.action_space = action_space
+        self.use_curiosity = use_curiosity
+
+        self.world_model = world_model
+        self.actor = actor
+        self.critic = critic
+
+        self.disagree_nets = None
+        if self.use_curiosity:
+            raise NotImplementedError
+
+    def forward_inference(self, observations, previous_states, is_first):
+        """Performs a (non-exploring) action computation step given obs and states.
+
+        Note that all input data should not have a time rank (only a batch dimension).
+
+        Args:
+            observations: The current environment observation with shape (B, ...).
+            previous_states: Dict with keys `a`, `h`, and `z` used as input to the RSSM
+                to produce the next h-state, from which then to compute the action
+                using the actor network. All values in the dict should have shape
+                (B, ...) (no time rank).
+            is_first: Batch of is_first flags. These should be True if a new episode
+                has been started at the current timestep (meaning `observations` is the
+                reset observation from the environment).
+        """
+        # Perform one step in the world model (starting from `previous_state` and
+        # using the observations to yield a current (posterior) state).
+        states = self.world_model.forward_inference(
+            observations=observations,
+            previous_states=previous_states,
+            is_first=is_first,
+        )
+        # Compute action using our actor network and the current states.
+        _, distr_params = self.actor(
+            h=states["h"], z=states["z"], return_distr_params=True
+        )
+        # Use the mode of the distribution (Discrete=argmax, Normal=mean).
+        distr = self.actor.get_action_dist_object(distr_params)
+        actions = distr.mode
+        return actions, {"h": states["h"], "z": states["z"], "a": actions}
+
+    def forward_exploration(self, observations, previous_states, is_first):
+        """Performs an exploratory action computation step given obs and states.
+
+        Note that all input data should not have a time rank (only a batch dimension).
+
+        Args:
+            observations: The current environment observation with shape (B, ...).
+            previous_states: Dict with keys `a`, `h`, and `z` used as input to the RSSM
+                to produce the next h-state, from which then to compute the action
+                using the actor network. All values in the dict should have shape
+                (B, ...) (no time rank).
+            is_first: Batch of is_first flags. These should be True if a new episode
+                has been started at the current timestep (meaning `observations` is the
+                reset observation from the environment).
+        """
+        # Perform one step in the world model (starting from `previous_state` and
+        # using the observations to yield a current (posterior) state).
+        states = self.world_model.forward_inference(
+            observations=observations,
+            previous_states=previous_states,
+            is_first=is_first,
+        )
+        # Compute action using our actor network and the current states.
+        actions = self.actor(h=states["h"], z=states["z"])
+        return actions, {"h": states["h"], "z": states["z"], "a": actions}
+
+    def forward_train(self, observations, actions, is_first):
+        """Performs a training forward pass given observations and actions.
+
+        Note that all input data must have a time rank (batch-major: [B, T, ...]).
+
+        Args:
+            observations: The environment observations with shape (B, T, ...). Thus,
+                the batch has B rows of T timesteps each. Note that it's ok to have
+                episode boundaries (is_first=True) within a batch row. DreamerV3 will
+                simply insert an initial state before these locations and continue the
+                sequence modelling (with the RSSM). Hence, there will be no zero
+                padding.
+            actions: The actions actually taken in the environment with shape
+                (B, T, ...). See `observations` docstring for details on how B and T are
+                handled.
+            is_first: Batch of is_first flags. These should be True:
+                - if a new episode has been started at the current timestep (meaning
+                `observations` is the reset observation from the environment).
+                - in each batch row at T=0 (first timestep of each of the B batch
+                rows), regardless of whether the actual env had an episode boundary
+                there or not.
+        """
+        return self.world_model.forward_train(
+            observations=observations,
+            actions=actions,
+            is_first=is_first,
+        )
+
+    def get_initial_state(self):
+        """Returns the initial state of the dreamer model (a, h-, z-states).
+
+        An initial state is generated using the previous action, the tanh of the
+        (learned) h-state variable and the dynamics predictor (or "prior net") to
+        compute z^0 from h0. In this last step, it is important that we do NOT sample
+        the z^-state (as we would usually do during dreaming), but rather take the mode
+        (argmax, then one-hot again).
+
+        Note that the initial state is returned without batch dimension.
+        """
+        states = self.world_model.get_initial_state()
+
+        action_dim = (
+            self.action_space.n
+            if isinstance(self.action_space, gym.spaces.Discrete)
+            else np.prod(self.action_space.shape)
+        )
+        states["a"] = torch.zeros((action_dim,), dtype=torch.float32)
+        return states
+
+    def dream_trajectory(self, start_states, start_is_terminated, timesteps_H, gamma):
+        """Dreams trajectories of length H from batch of h- and z-states.
+
+        Note that incoming data will have the shapes (BxT, ...), where the original
+        batch- and time-dimensions are already folded together. Beginning from this
+        new batch dim (BxT), we will unroll `timesteps_H` timesteps in a time-major
+        fashion, such that the dreamed data will have shape (H, BxT, ...).
+
+        Args:
+            start_states: Dict of `h` and `z` states in the shape of (B, ...) and
+                (B, num_categoricals, num_classes), respectively, as
+                computed by a train forward pass. From each individual h-/z-state pair
+                in the given batch, we will branch off a dreamed trajectory of len
+                `timesteps_H`.
+            start_is_terminated: Float flags of shape (B,) indicating whether the
+                first timesteps of each batch row is already a terminated timestep
+                (given by the actual environment).
+            timesteps_H: The number of timesteps to dream for.
+            gamma: The discount factor gamma.
+        """
+        # Dreamed actions (one-hot encoded for discrete actions).
+        a_dreamed_t0_to_H = []
+        a_dreamed_dist_params_t0_to_H = []
+
+        h = start_states["h"].detach()
+        z = start_states["z"].detach()
+
+        # GRU outputs.
+        h_states_t0_to_H = [h]
+        # Dynamics model outputs.
+        z_states_prior_t0_to_H = [z]
+
+        # Compute `a` using actor network (already the first step uses a dreamed action,
+        # not a sampled one).
+        a, a_dist_params = self.actor(
+            # We have to stop the gradients through the states. B/c we are using a
+            # differentiable Discrete action distribution (straight through gradients
+            # with `a = stop_gradient(sample(probs)) + probs - stop_gradient(probs)`,
+            # we otherwise would add dependencies of the `-log(pi(a|s))` REINFORCE loss
+            # term on actions further back in the trajectory.
+            h=h.detach(),
+            z=z.detach(),
+            return_distr_params=True,
+        )
+        a_dreamed_t0_to_H.append(a)
+        a_dreamed_dist_params_t0_to_H.append(a_dist_params)
+
+        # Disable all gradients from the world model so they don't get backprop'd
+        # through twice when computing the actor loss (for cont. actions).
+        for p in self.world_model.parameters():
+            p.requires_grad_(False)
+
+        for i in range(timesteps_H):
+            # Move one step in the dream using the RSSM.
+            h = self.world_model.sequence_model(a=a, h=h, z=z)
+            h_states_t0_to_H.append(h)
+
+            # Compute prior z using dynamics model.
+            z = self.world_model.dynamics_predictor(h=h)
+            z_states_prior_t0_to_H.append(z)
+
+            # Compute `a` using actor network.
+            a, a_dist_params = self.actor(
+                h=h.detach(),
+                z=z.detach(),
+                return_distr_params=True,
+            )
+            a_dreamed_t0_to_H.append(a)
+            a_dreamed_dist_params_t0_to_H.append(a_dist_params)
+
+        h_states_H_B = torch.stack(h_states_t0_to_H, dim=0)  # (T, B, ...)
+        h_states_HxB = h_states_H_B.reshape([-1] + list(h_states_H_B.shape[2:]))
+
+        z_states_prior_H_B = torch.stack(z_states_prior_t0_to_H, dim=0)  # (T, B, ...)
+        z_states_prior_HxB = z_states_prior_H_B.reshape(
+            [-1] + list(z_states_prior_H_B.shape[2:])
+        )
+
+        a_dreamed_H_B = torch.stack(a_dreamed_t0_to_H, dim=0)  # (T, B, ...)
+        a_dreamed_dist_params_H_B = torch.stack(a_dreamed_dist_params_t0_to_H, dim=0)
+
+        # Compute r using reward predictor.
+        r_dreamed_H_B = inverse_symlog(
+            self.world_model.reward_predictor(h=h_states_HxB, z=z_states_prior_HxB)
+        )
+        r_dreamed_H_B = r_dreamed_H_B.reshape([timesteps_H + 1, -1])
+
+        # Compute intrinsic rewards.
+        if self.use_curiosity:
+            results_HxB = self.disagree_nets.compute_intrinsic_rewards(
+                h=h_states_HxB,
+                z=z_states_prior_HxB,
+                a=a_dreamed_H_B.reshape([-1] + a_dreamed_H_B.shape[2:]),
+            )
+            r_intrinsic_H_B = results_HxB["rewards_intrinsic"]
+            r_intrinsic_H_B = r_intrinsic_H_B.reshape([timesteps_H + 1, -1])[1:]
+            curiosity_forward_train_outs = results_HxB["forward_train_outs"]
+            del results_HxB
+
+        # Compute continues using continue predictor.
+        c_dreamed_HxB = self.world_model.continue_predictor(
+            h=h_states_HxB,
+            z=z_states_prior_HxB,
+        )
+        c_dreamed_H_B = c_dreamed_HxB.reshape([timesteps_H + 1, -1])
+        # Force-set first `continue` flags to False iff `start_is_terminated`.
+        # Note: This will cause the loss-weights for this row in the batch to be
+        # completely zero'd out. In general, we don't use dreamed data past any
+        # predicted (or actual first) continue=False flags.
+        c_dreamed_H_B = torch.cat(
+            [1.0 - start_is_terminated.unsqueeze(0).float(), c_dreamed_H_B[1:]], dim=0
+        )
+
+        # Loss weights for each individual dreamed timestep. Zero-out all timesteps
+        # that lie past continue=False flags. B/c our world model does NOT learn how
+        # to skip terminal/reset episode boundaries, dreamed data crossing such a
+        # boundary should not be used for critic/actor learning either.
+        dream_loss_weights_H_B = torch.cumprod(gamma * c_dreamed_H_B, dim=0) / gamma
+
+        # Reactivate world model gradients.
+        for p in self.world_model.parameters():
+            p.requires_grad_(True)
+
+        # Compute the symlog'd value logits (w/o world model gradients; used for the
+        # critic loss).
+        _, v_symlog_dreamed_logits_HxB_wm_detached = self.critic(
+            h=h_states_HxB.detach(),
+            z=z_states_prior_HxB.detach(),
+            use_ema=False,
+            return_logits=True,
+        )
+
+        # Compute the value estimates (including world model gradients -> 1 sequence
+        # model step after the action has been computed; used for the scaled value
+        # target used in the actor loss for cont. actions).
+        # Disable all gradients from the critic so they don't get backprop'd
+        # through twice when computing the actor loss (for cont. actions).
+        for p in self.critic.parameters():
+            p.requires_grad_(False)
+        v, _ = self.critic(
+            h=h_states_HxB,
+            z=z_states_prior_HxB,
+            use_ema=False,
+            return_logits=True,
+        )
+        # Reactivate critic gradients.
+        for p in self.critic.parameters():
+            p.requires_grad_(True)
+        v_dreamed_HxB = inverse_symlog(v)
+        v_dreamed_H_B = v_dreamed_HxB.reshape([timesteps_H + 1, -1])
+
+        # Compute the EMA net outputs w/o any gradients.
+        with torch.no_grad():
+            v_symlog_dreamed_ema_HxB = self.critic(
+                h=h_states_HxB.detach(),
+                z=z_states_prior_HxB.detach(),
+                return_logits=False,
+                use_ema=True,
+            )
+        v_symlog_dreamed_ema_H_B = v_symlog_dreamed_ema_HxB.reshape(
+            [timesteps_H + 1, -1]
+        )
+
+        ret = {
+            "h_states_t0_to_H_BxT": h_states_H_B,
+            "z_states_prior_t0_to_H_BxT": z_states_prior_H_B,
+            "rewards_dreamed_t0_to_H_BxT": r_dreamed_H_B,
+            "continues_dreamed_t0_to_H_BxT": c_dreamed_H_B,
+            "actions_dreamed_t0_to_H_BxT": a_dreamed_H_B,
+            "actions_dreamed_dist_params_t0_to_H_BxT": a_dreamed_dist_params_H_B,
+            # Critic (w/ world-model grads for actor loss).
+            "values_dreamed_t0_to_H_BxT": v_dreamed_H_B,
+            # Critic (world-model detached, for critic loss).
+            "values_symlog_dreamed_logits_t0_to_HxBxT_wm_detached": v_symlog_dreamed_logits_HxB_wm_detached,
+            # Critic EMA.
+            "v_symlog_dreamed_ema_t0_to_H_BxT": v_symlog_dreamed_ema_H_B,
+            # Loss weights for critic- and actor losses.
+            "dream_loss_weights_t0_to_H_BxT": dream_loss_weights_H_B,
+        }
+
+        if self.use_curiosity:
+            ret["rewards_intrinsic_t1_to_H_B"] = r_intrinsic_H_B
+            ret.update(curiosity_forward_train_outs)
+
+        if isinstance(self.action_space, gym.spaces.Discrete):
+            ret["actions_ints_dreamed_t0_to_H_B"] = torch.argmax(a_dreamed_H_B, dim=-1)
+
+        return ret
+
+    def dream_trajectory_with_burn_in(
+        self,
+        *,
+        start_states,
+        timesteps_burn_in: int,
+        timesteps_H: int,
+        observations,  # [B, >=timesteps_burn_in]
+        actions,  # [B, timesteps_burn_in (+timesteps_H)?]
+        use_sampled_actions_in_dream: bool = False,
+        use_random_actions_in_dream: bool = False,
+    ):
+        """Dreams trajectory from N initial observations and initial states.
+
+        Note: This is only used for reporting and debugging, not for actual world-model
+        or policy training.
+
+        Args:
+            start_states: The batch of start states (dicts with `a`, `h`, and `z` keys)
+                to begin dreaming with. These are used to compute the first h-state
+                using the sequence model.
+            timesteps_burn_in: For how many timesteps should be use the posterior
+                z-states (computed by the posterior net and actual observations from
+                the env)?
+            timesteps_H: For how many timesteps should we dream using the prior
+                z-states (computed by the dynamics (prior) net and h-states only)?
+                Note that the total length of the returned trajectories will
+                be `timesteps_burn_in` + `timesteps_H`.
+            observations: The batch (B, T, ...) of observations (to be used only during
+                burn-in over `timesteps_burn_in` timesteps).
+            actions: The batch (B, T, ...) of actions to use during a) burn-in over the
+                first `timesteps_burn_in` timesteps and - possibly - b) during
+                actual dreaming, iff use_sampled_actions_in_dream=True.
+            use_sampled_actions_in_dream: If True, instead of using our actor network
+                to compute fresh actions, we will use the one provided via the `actions`
+                argument. Note that in the latter case, the `actions` time dimension
+                must be at least `timesteps_burn_in` + `timesteps_H` long.
+            use_random_actions_in_dream: Whether to use randomly sampled actions in the
+                dream. Note that this does not apply to the burn-in phase, during which
+                we will always use the actions given in the `actions` argument.
+        """
+        assert not (use_sampled_actions_in_dream and use_random_actions_in_dream)
+
+        B = observations.shape[0]
+
+        # Produce initial N internal posterior states (burn-in) using the given
+        # observations:
+        states = start_states
+        for i in range(timesteps_burn_in):
+            states = self.world_model.forward_inference(
+                observations=observations[:, i : i + 1],
+                previous_states=states,
+                is_first=torch.full((B,), 1.0 if i == 0 else 0.0),
+            )
+            states["a"] = actions[:, i]
+
+        # Start producing the actual dream, using prior states and either the given
+        # actions, dreamed, or random ones.
+        h_states_t0_to_H = [states["h"]]
+        z_states_prior_t0_to_H = [states["z"]]
+        a_t0_to_H = [states["a"]]
+
+        for j in range(timesteps_H):
+            # Compute next h using sequence model.
+            h = self.world_model.sequence_model(
+                a=states["a"],
+                h=states["h"],
+                z=states["z"],
+            )
+            h_states_t0_to_H.append(h)
+            # Compute z from h, using the dynamics model (we don't have an actual
+            # observation at this timestep).
+            z = self.world_model.dynamics_predictor(h=h)
+            z_states_prior_t0_to_H.append(z)
+
+            # Compute next dreamed action or use sampled one or random one.
+            if use_sampled_actions_in_dream:
+                a = actions[:, timesteps_burn_in + j]
+            elif use_random_actions_in_dream:
+                if isinstance(self.action_space, gym.spaces.Discrete):
+                    a = torch.randint(self.action_space.n, (B,), dtype=torch.int64)
+                    a = torch.nn.functional.one_hot(a, num_classes=self.action_space.n)
+                else:
+                    a = torch.rand(
+                        (B,) + self.action_space.shape, dtype=self.action_space.dtype
+                    )
+            else:
+                a = self.actor(h=h, z=z)
+            a_t0_to_H.append(a)
+
+            states = {"h": h, "z": z, "a": a}
+
+        # Fold time-rank for upcoming batch-predictions (no sequences needed anymore).
+        h_states_t0_to_H_B = torch.stack(h_states_t0_to_H, dim=0)
+        h_states_t0_to_HxB = h_states_t0_to_H_B.reshape(
+            [-1] + list(h_states_t0_to_H_B.shape[2:])
+        )
+
+        z_states_prior_t0_to_H_B = torch.stack(z_states_prior_t0_to_H, dim=0)
+        z_states_prior_t0_to_HxB = z_states_prior_t0_to_H_B.reshape(
+            [-1] + list(z_states_prior_t0_to_H_B.shape[2:])
+        )
+
+        a_t0_to_H_B = torch.stack(a_t0_to_H, dim=0)
+
+        # Compute o using decoder.
+        o_dreamed_t0_to_HxB = self.world_model.decoder(
+            h=h_states_t0_to_HxB,
+            z=z_states_prior_t0_to_HxB,
+        )
+        if self.world_model.symlog_obs:
+            o_dreamed_t0_to_HxB = inverse_symlog(o_dreamed_t0_to_HxB)
+
+        # Compute r using reward predictor.
+        r_dreamed_t0_to_H_B = inverse_symlog(
+            self.world_model.reward_predictor(
+                h=h_states_t0_to_HxB,
+                z=z_states_prior_t0_to_HxB,
+            )
+        ).reshape([-1, B])
+
+        # Compute continues using continue predictor.
+        c_dreamed_t0_to_H_B = self.world_model.continue_predictor(
+            h=h_states_t0_to_HxB,
+            z=z_states_prior_t0_to_HxB,
+        ).reshape([-1, B])
+
+        # Return everything as time-major (H, B, ...), where H is the timesteps dreamed
+        # (NOT burn-in'd) and B is a batch dimension (this might or might not include
+        # an original time dimension from the real env, from all of which we then branch
+        # out our dream trajectories).
+        ret = {
+            "h_states_t0_to_H_BxT": h_states_t0_to_H_B,
+            "z_states_prior_t0_to_H_BxT": z_states_prior_t0_to_H_B,
+            # Unfold time-ranks in predictions.
+            "observations_dreamed_t0_to_H_BxT": torch.reshape(
+                o_dreamed_t0_to_HxB, [-1, B] + list(observations.shape)[2:]
+            ),
+            "rewards_dreamed_t0_to_H_BxT": r_dreamed_t0_to_H_B,
+            "continues_dreamed_t0_to_H_BxT": c_dreamed_t0_to_H_B,
+        }
+
+        # Figure out action key (random, sampled from env, dreamed?).
+        if use_sampled_actions_in_dream:
+            key = "actions_sampled_t0_to_H_BxT"
+        elif use_random_actions_in_dream:
+            key = "actions_random_t0_to_H_BxT"
+        else:
+            key = "actions_dreamed_t0_to_H_BxT"
+        ret[key] = a_t0_to_H_B
+
+        # Also provide int-actions, if discrete action space.
+        if isinstance(self.action_space, gym.spaces.Discrete):
+            ret[re.sub("^actions_", "actions_ints_", key)] = torch.argmax(
+                a_t0_to_H_B, dim=-1
+            )
+
+        return ret

rllib/algorithms/dreamerv3/torch/models/world_model.py

@@ -0,0 +1,431 @@
+"""
+[1] Mastering Diverse Domains through World Models - 2023
+D. Hafner, J. Pasukonis, J. Ba, T. Lillicrap
+https://arxiv.org/pdf/2301.04104v1.pdf
+"""
+from typing import Optional
+
+import gymnasium as gym
+import numpy as np
+import tree  # pip install dm_tree
+
+from ray.rllib.algorithms.dreamerv3.torch.models.components import (
+    representation_layer,
+)
+from ray.rllib.algorithms.dreamerv3.torch.models.components.continue_predictor import (
+    ContinuePredictor,
+)
+from ray.rllib.algorithms.dreamerv3.torch.models.components.dynamics_predictor import (
+    DynamicsPredictor,
+)
+from ray.rllib.algorithms.dreamerv3.torch.models.components.mlp import MLP
+from ray.rllib.algorithms.dreamerv3.torch.models.components.reward_predictor import (
+    RewardPredictor,
+)
+from ray.rllib.algorithms.dreamerv3.torch.models.components.sequence_model import (
+    SequenceModel,
+)
+from ray.rllib.algorithms.dreamerv3.utils import get_dense_hidden_units, get_gru_units
+from ray.rllib.utils.framework import try_import_torch
+from ray.rllib.utils.torch_utils import symlog
+
+torch, nn = try_import_torch()
+if torch:
+    F = nn.functional
+
+
+class WorldModel(nn.Module):
+    """WorldModel component of [1] w/ encoder, decoder, RSSM, reward/cont. predictors.
+
+    See eq. 3 of [1] for all components and their respective in- and outputs.
+    Note that in the paper, the "encoder" includes both the raw encoder plus the
+    "posterior net", which produces posterior z-states from observations and h-states.
+
+    Note: The "internal state" of the world model always consists of:
+    The actions `a` (initially, this is a zeroed-out action), `h`-states (deterministic,
+    continuous), and `z`-states (stochastic, discrete).
+    There are two versions of z-states: "posterior" for world model training and "prior"
+    for creating the dream data.
+
+    Initial internal state values (`a`, `h`, and `z`) are inserted where ever a new
+    episode starts within a batch row OR at the beginning of each train batch's B rows,
+    regardless of whether there was an actual episode boundary or not. Thus, internal
+    states are not required to be stored in or retrieved from the replay buffer AND
+    retrieved batches from the buffer must not be zero padded.
+
+    Initial `a` is the zero "one hot" action, e.g. [0.0, 0.0] for Discrete(2), initial
+    `h` is a separate learned variable, and initial `z` are computed by the "dynamics"
+    (or "prior") net, using only the initial-h state as input.
+    """
+
+    def __init__(
+        self,
+        *,
+        model_size: str = "XS",
+        observation_space: gym.Space,
+        action_space: gym.Space,
+        batch_length_T: int = 64,
+        encoder: nn.Module,
+        decoder: nn.Module,
+        num_gru_units: Optional[int] = None,
+        symlog_obs: bool = True,
+    ):
+        """Initializes a WorldModel instance.
+
+        Args:
+             model_size: The "Model Size" used according to [1] Appendinx B.
+                Use None for manually setting the different network sizes.
+             action_space: The action space the our environment used.
+             batch_length_T: The length (T) of the sequences used for training. The
+                actual shape of the input data (e.g. rewards) is then: [B, T, ...],
+                where B is the "batch size", T is the "batch length" (this arg) and
+                "..." is the dimension of the data (e.g. (64, 64, 3) for Atari image
+                observations). Note that a single row (within a batch) may contain data
+                from different episodes, but an already on-going episode is always
+                finished, before a new one starts within the same row.
+            encoder: The encoder Model taking observations as inputs and
+                outputting a 1D latent vector that will be used as input into the
+                posterior net (z-posterior state generating layer). Inputs are symlogged
+                if inputs are NOT images. For images, we use normalization between -1.0
+                and 1.0 (x / 128 - 1.0)
+            decoder: The decoder Model taking h- and z-states as inputs and generating
+                a (possibly symlogged) predicted observation. Note that for images,
+                the last decoder layer produces the exact, normalized pixel values
+                (not a Gaussian as described in [1]!).
+            num_gru_units: The number of GRU units to use. If None, use
+                `model_size` to figure out this parameter.
+            symlog_obs: Whether to predict decoded observations in symlog space.
+                This should be False for image based observations.
+                According to the paper [1] Appendix E: "NoObsSymlog: This ablation
+                removes the symlog encoding of inputs to the world model and also
+                changes the symlog MSE loss in the decoder to a simple MSE loss.
+                *Because symlog encoding is only used for vector observations*, this
+                ablation is equivalent to DreamerV3 on purely image-based environments".
+        """
+        super().__init__()
+
+        self.model_size = model_size
+        self.batch_length_T = batch_length_T
+        self.symlog_obs = symlog_obs
+        self.action_space = action_space
+        a_flat = (
+            action_space.n
+            if isinstance(action_space, gym.spaces.Discrete)
+            else (np.prod(action_space.shape))
+        )
+
+        # Encoder (latent 1D vector generator) (xt -> lt).
+        self.encoder = encoder
+
+        self.num_gru_units = get_gru_units(
+            model_size=self.model_size,
+            override=num_gru_units,
+        )
+
+        # Posterior predictor consisting of an MLP and a RepresentationLayer:
+        # [ht, lt] -> zt.
+        # In Danijar's code, this is called: `obs_out`.
+        self.posterior_mlp = MLP(
+            input_size=(self.num_gru_units + encoder.output_size[0]),
+            model_size=self.model_size,
+            output_layer_size=None,
+            # In Danijar's code, the posterior predictor only has a single layer,
+            # no matter the model size:
+            num_dense_layers=1,
+        )
+        # The (posterior) z-state generating layer.
+        # In Danijar's code, this is called: `obs_stats`.
+        self.posterior_representation_layer = representation_layer.RepresentationLayer(
+            input_size=get_dense_hidden_units(self.model_size),
+            model_size=self.model_size,
+        )
+
+        z_flat = (
+            self.posterior_representation_layer.num_categoricals
+            * self.posterior_representation_layer.num_classes_per_categorical
+        )
+        h_plus_z_flat = self.num_gru_units + z_flat
+
+        # Dynamics (prior z-state) predictor: ht -> z^t
+        # In Danijar's code, the layers in this network are called:
+        # `img_out` (1 Linear) and `img_stats` (representation layer).
+        self.dynamics_predictor = DynamicsPredictor(
+            input_size=self.num_gru_units, model_size=self.model_size
+        )
+
+        # GRU for the RSSM: [at, ht, zt] -> ht+1
+        # Initial h-state variable (learnt).
+        # -> tanh(self.initial_h) -> deterministic state
+        # Use our Dynamics predictor for initial stochastic state, BUT with greedy
+        # (mode) instead of sampling.
+        self.initial_h = nn.Parameter(
+            torch.zeros(self.num_gru_units), requires_grad=True
+        )
+        # The actual sequence model containing the GRU layer.
+        # In Danijar's code, the layers in this network are called:
+        # `img_in` (1 Linear) and `gru` (custom GRU implementation).
+        self.sequence_model = SequenceModel(
+            # Only z- and a-state go into pre-layer. The output of that goes then
+            # into GRU (together with h-state).
+            input_size=int(z_flat + a_flat),
+            model_size=self.model_size,
+            action_space=self.action_space,
+            num_gru_units=self.num_gru_units,
+        )
+
+        # Reward Predictor: [ht, zt] -> rt.
+        self.reward_predictor = RewardPredictor(
+            input_size=h_plus_z_flat,
+            model_size=self.model_size,
+        )
+        # Continue Predictor: [ht, zt] -> ct.
+        self.continue_predictor = ContinuePredictor(
+            input_size=h_plus_z_flat,
+            model_size=self.model_size,
+        )
+
+        # Decoder: [ht, zt] -> x^t.
+        self.decoder = decoder
+
+    def get_initial_state(self) -> dict:
+        """Returns the (current) initial state of the world model (h- and z-states).
+
+        An initial state is generated using the tanh of the (learned) h-state variable
+        and the dynamics predictor (or "prior net") to compute z^0 from h0. In this last
+        step, it is important that we do NOT sample the z^-state (as we would usually
+        do during dreaming), but rather take the mode (argmax, then one-hot again).
+        """
+        h = torch.tanh(self.initial_h)
+        # Use the mode, NOT a sample for the initial z-state.
+        _, z_probs = self.dynamics_predictor(h.unsqueeze(0), return_z_probs=True)
+        z = z_probs.squeeze(0).argmax(dim=-1)
+        z = F.one_hot(z, num_classes=z_probs.shape[-1])
+
+        return {"h": h, "z": z}
+
+    def forward_inference(
+        self,
+        observations: "torch.Tensor",
+        previous_states: dict,
+        is_first: "torch.Tensor",
+    ) -> dict:
+        """Performs a forward step for inference (e.g. environment stepping).
+
+        Works analogous to `forward_train`, except that all inputs are provided
+        for a single timestep in the shape of [B, ...] (no time dimension!).
+
+        Args:
+            observations: The batch (B, ...) of observations to be passed through
+                the encoder network to yield the inputs to the representation layer
+                (which then can compute the z-states).
+            previous_states: A dict with `h`, `z`, and `a` keys mapping to the
+                respective previous states/actions. All of the shape (B, ...), no time
+                rank.
+            is_first: The batch (B) of `is_first` flags.
+
+        Returns:
+            The next deterministic h-state (h(t+1)) as predicted by the sequence model.
+        """
+        B = observations.shape[0]
+        initial_states = tree.map_structure(
+            # Repeat only the batch dimension (B times).
+            lambda s: s.unsqueeze(0).repeat(B, *([1] * len(s.shape))),
+            self.get_initial_state(),
+        )
+
+        # If first, mask it with initial state/actions.
+        previous_h = self._mask(previous_states["h"], 1.0 - is_first)  # zero out
+        previous_h = previous_h + self._mask(initial_states["h"], is_first)  # add init
+
+        previous_z = self._mask(previous_states["z"], 1.0 - is_first)  # zero out
+        previous_z = previous_z + self._mask(initial_states["z"], is_first)  # add init
+
+        # Zero out actions (no special learnt initial state).
+        previous_a = self._mask(previous_states["a"], 1.0 - is_first)
+
+        # Compute new states.
+        h = self.sequence_model(a=previous_a, h=previous_h, z=previous_z)
+        z = self.compute_posterior_z(observations=observations, initial_h=h)
+
+        return {"h": h, "z": z}
+
+    def forward_train(
+        self,
+        observations: "torch.Tensor",
+        actions: "torch.Tensor",
+        is_first: "torch.Tensor",
+    ) -> dict:
+        """Performs a forward step for training.
+
+        1) Forwards all observations [B, T, ...] through the encoder network to yield
+        o_processed[B, T, ...].
+        2) Uses initial state (h0/z^0/a0[B, 0, ...]) and sequence model (RSSM) to
+        compute the first internal state (h1 and z^1).
+        3) Uses action a[B, 1, ...], z[B, 1, ...] and h[B, 1, ...] to compute the
+        next h-state (h[B, 2, ...]), etc..
+        4) Repeats 2) and 3) until t=T.
+        5) Uses all h[B, T, ...] and z[B, T, ...] to compute predicted/reconstructed
+        observations, rewards, and continue signals.
+        6) Returns predictions from 5) along with all z-states z[B, T, ...] and
+        the final h-state (h[B, ...] for t=T).
+
+        Should we encounter is_first=True flags in the middle of a batch row (somewhere
+        within an ongoing sequence of length T), we insert this world model's initial
+        state again (zero-action, learned init h-state, and prior-computed z^) and
+        simply continue (no zero-padding).
+
+        Args:
+            observations: The batch (B, T, ...) of observations to be passed through
+                the encoder network to yield the inputs to the representation layer
+                (which then can compute the posterior z-states).
+            actions: The batch (B, T, ...) of actions to be used in combination with
+                h-states and computed z-states to yield the next h-states.
+            is_first: The batch (B, T) of `is_first` flags.
+        """
+        if self.symlog_obs:
+            observations = symlog(observations)
+
+        # Compute bare encoder outs (not z; this is done later with involvement of the
+        # sequence model and the h-states).
+        # Fold time dimension for CNN pass.
+        shape = observations.shape
+        B, T = shape[0], shape[1]
+        observations = observations.view((-1,) + shape[2:])
+        encoder_out = self.encoder(observations)
+        # Unfold time dimension.
+        encoder_out = encoder_out.view(
+            (
+                B,
+                T,
+            )
+            + encoder_out.shape[1:]
+        )
+        # Make time major for faster upcoming loop.
+        encoder_out = encoder_out.transpose(0, 1)
+        # encoder_out=[T, B, ...]
+
+        initial_states = tree.map_structure(
+            # Repeat only the batch dimension (B times).
+            lambda s: s.unsqueeze(0).repeat(B, *([1] * len(s.shape))),
+            self.get_initial_state(),
+        )
+
+        # Make actions and `is_first` time-major.
+        actions = actions.transpose(0, 1)
+        is_first = is_first.transpose(0, 1).float()
+
+        # Loop through the T-axis of our samples and perform one computation step at
+        # a time. This is necessary because the sequence model's output (h(t+1)) depends
+        # on the current z(t), but z(t) depends on the current sequence model's output
+        # h(t).
+        z_t0_to_T = [initial_states["z"]]
+        z_posterior_probs = []
+        z_prior_probs = []
+        h_t0_to_T = [initial_states["h"]]
+        for t in range(self.batch_length_T):
+            # If first, mask it with initial state/actions.
+            h_tm1 = self._mask(h_t0_to_T[-1], 1.0 - is_first[t])  # zero out
+            h_tm1 = h_tm1 + self._mask(initial_states["h"], is_first[t])  # add init
+
+            z_tm1 = self._mask(z_t0_to_T[-1], 1.0 - is_first[t])  # zero out
+            z_tm1 = z_tm1 + self._mask(initial_states["z"], is_first[t])  # add init
+
+            # Zero out actions (no special learnt initial state).
+            a_tm1 = self._mask(actions[t - 1], 1.0 - is_first[t])
+
+            # Perform one RSSM (sequence model) step to get the current h.
+            h_t = self.sequence_model(a=a_tm1, h=h_tm1, z=z_tm1)
+            h_t0_to_T.append(h_t)
+
+            posterior_mlp_input = torch.cat([encoder_out[t], h_t], dim=-1)
+            repr_input = self.posterior_mlp(posterior_mlp_input)
+            # Draw one z-sample (z(t)) and also get the z-distribution for dynamics and
+            # representation loss computations.
+            z_t, z_probs = self.posterior_representation_layer(
+                repr_input,
+                return_z_probs=True,
+            )
+            # z_t=[B, num_categoricals, num_classes]
+            z_posterior_probs.append(z_probs)
+            z_t0_to_T.append(z_t)
+
+            # Compute the predicted z_t (z^) using the dynamics model.
+            _, z_probs = self.dynamics_predictor(h_t, return_z_probs=True)
+            z_prior_probs.append(z_probs)
+
+        # Stack at time dimension to yield: [B, T, ...].
+        h_t1_to_T = torch.stack(h_t0_to_T[1:], dim=1)
+        z_t1_to_T = torch.stack(z_t0_to_T[1:], dim=1)
+
+        # Fold time axis to retrieve the final (loss ready) Independent distribution
+        # (over `num_categoricals` Categoricals).
+        z_posterior_probs = torch.stack(z_posterior_probs, dim=1)
+        z_posterior_probs = z_posterior_probs.view(
+            (-1,) + z_posterior_probs.shape[2:],
+        )
+        # Fold time axis to retrieve the final (loss ready) Independent distribution
+        # (over `num_categoricals` Categoricals).
+        z_prior_probs = torch.stack(z_prior_probs, dim=1)
+        z_prior_probs = z_prior_probs.view((-1,) + z_prior_probs.shape[2:])
+
+        # Fold time dimension for parallelization of all dependent predictions:
+        # observations (reproduction via decoder), rewards, continues.
+        h_BxT = h_t1_to_T.view((-1,) + h_t1_to_T.shape[2:])
+        z_BxT = z_t1_to_T.view((-1,) + z_t1_to_T.shape[2:])
+
+        obs_distribution_means = self.decoder(h=h_BxT, z=z_BxT)
+
+        # Compute (predicted) reward distributions.
+        rewards, reward_logits = self.reward_predictor(
+            h=h_BxT, z=z_BxT, return_logits=True
+        )
+
+        # Compute (predicted) continue distributions.
+        continues, continue_distribution = self.continue_predictor(
+            h=h_BxT, z=z_BxT, return_distribution=True
+        )
+
+        # Return outputs for loss computation.
+        # Note that all shapes are [BxT, ...] (time axis already folded).
+        return {
+            # Obs.
+            "sampled_obs_symlog_BxT": observations,
+            "obs_distribution_means_BxT": obs_distribution_means,
+            # Rewards.
+            "reward_logits_BxT": reward_logits,
+            "rewards_BxT": rewards,
+            # Continues.
+            "continue_distribution_BxT": continue_distribution,
+            "continues_BxT": continues,
+            # Deterministic, continuous h-states (t1 to T).
+            "h_states_BxT": h_BxT,
+            # Sampled, discrete posterior z-states and their probs (t1 to T).
+            "z_posterior_states_BxT": z_BxT,
+            "z_posterior_probs_BxT": z_posterior_probs,
+            # Probs of the prior z-states (t1 to T).
+            "z_prior_probs_BxT": z_prior_probs,
+        }
+
+    def compute_posterior_z(
+        self, observations: "torch.Tensor", initial_h: "torch.Tensor"
+    ) -> "torch.Tensor":
+        # Fold time dimension for possible CNN pass.
+        shape = observations.shape
+        observations = observations.view((-1,) + shape[2:])
+        # Compute bare encoder outputs (not including z, which is computed in next step
+        # with involvement of the previous output (initial_h) of the sequence model).
+        # encoder_outs=[B, ...]
+        if self.symlog_obs:
+            observations = symlog(observations)
+        encoder_out = self.encoder(observations)
+        # Concat encoder outs with the h-states.
+        posterior_mlp_input = torch.cat([encoder_out, initial_h], dim=-1)
+        # Compute z.
+        repr_input = self.posterior_mlp(posterior_mlp_input)
+        # Draw one z-sample (no need to return the distribution here).
+        z_t = self.posterior_representation_layer(repr_input, return_z_probs=False)
+        return z_t
+
+    @staticmethod
+    def _mask(value: "torch.Tensor", mask: "torch.Tensor") -> "torch.Tensor":
+        return torch.einsum("b...,b->b...", value, mask)

rllib/algorithms/dreamerv3/utils/__init__.py

@@ -0,0 +1,168 @@
+"""
+Utility functions for the DreamerV3 ([1]) algorithm.
+
+[1] Mastering Diverse Domains through World Models - 2023
+D. Hafner, J. Pasukonis, J. Ba, T. Lillicrap
+https://arxiv.org/pdf/2301.04104v1.pdf
+"""
+
+_ALLOWED_MODEL_DIMS = [
+    # RLlib debug sizes (not mentioned in [1]).
+    "nano",
+    "micro",
+    "mini",
+    "XXS",
+    # Regular sizes (listed in table B in [1]).
+    "XS",
+    "S",
+    "M",
+    "L",
+    "XL",
+]
+
+
+def get_cnn_multiplier(model_size, override=None):
+    if override is not None:
+        return override
+
+    assert model_size in _ALLOWED_MODEL_DIMS
+    cnn_multipliers = {
+        "nano": 2,
+        "micro": 4,
+        "mini": 8,
+        "XXS": 16,
+        "XS": 24,
+        "S": 32,
+        "M": 48,
+        "L": 64,
+        "XL": 96,
+    }
+    return cnn_multipliers[model_size]
+
+
+def get_dense_hidden_units(model_size, override=None):
+    if override is not None:
+        return override
+
+    assert model_size in _ALLOWED_MODEL_DIMS
+    dense_units = {
+        "nano": 16,
+        "micro": 32,
+        "mini": 64,
+        "XXS": 128,
+        "XS": 256,
+        "S": 512,
+        "M": 640,
+        "L": 768,
+        "XL": 1024,
+    }
+    return dense_units[model_size]
+
+
+def get_gru_units(model_size, override=None):
+    if override is not None:
+        return override
+
+    assert model_size in _ALLOWED_MODEL_DIMS
+    gru_units = {
+        "nano": 16,
+        "micro": 32,
+        "mini": 64,
+        "XXS": 128,
+        "XS": 256,
+        "S": 512,
+        "M": 1024,
+        "L": 2048,
+        "XL": 4096,
+    }
+    return gru_units[model_size]
+
+
+def get_num_z_categoricals(model_size, override=None):
+    if override is not None:
+        return override
+
+    assert model_size in _ALLOWED_MODEL_DIMS
+    gru_units = {
+        "nano": 4,
+        "micro": 8,
+        "mini": 16,
+        "XXS": 32,
+        "XS": 32,
+        "S": 32,
+        "M": 32,
+        "L": 32,
+        "XL": 32,
+    }
+    return gru_units[model_size]
+
+
+def get_num_z_classes(model_size, override=None):
+    if override is not None:
+        return override
+
+    assert model_size in _ALLOWED_MODEL_DIMS
+    gru_units = {
+        "nano": 4,
+        "micro": 8,
+        "mini": 16,
+        "XXS": 32,
+        "XS": 32,
+        "S": 32,
+        "M": 32,
+        "L": 32,
+        "XL": 32,
+    }
+    return gru_units[model_size]
+
+
+def get_num_curiosity_nets(model_size, override=None):
+    if override is not None:
+        return override
+
+    assert model_size in _ALLOWED_MODEL_DIMS
+    num_curiosity_nets = {
+        "nano": 8,
+        "micro": 8,
+        "mini": 8,
+        "XXS": 8,
+        "XS": 8,
+        "S": 8,
+        "M": 8,
+        "L": 8,
+        "XL": 8,
+    }
+    return num_curiosity_nets[model_size]
+
+
+def get_num_dense_layers(model_size, override=None):
+    if override is not None:
+        return override
+
+    assert model_size in _ALLOWED_MODEL_DIMS
+    num_dense_layers = {
+        "nano": 1,
+        "micro": 1,
+        "mini": 1,
+        "XXS": 1,
+        "XS": 1,
+        "S": 2,
+        "M": 3,
+        "L": 4,
+        "XL": 5,
+    }
+    return num_dense_layers[model_size]
+
+
+def do_symlog_obs(observation_space, symlog_obs_user_setting):
+    # If our symlog_obs setting is NOT set specifically (it's set to "auto"), return
+    # True if we don't have an image observation space, otherwise return False.
+
+    # TODO (sven): Support mixed observation spaces.
+
+    is_image_space = len(observation_space.shape) in [2, 3]
+    return (
+        not is_image_space
+        if symlog_obs_user_setting == "auto"
+        else symlog_obs_user_setting
+    )

rllib/algorithms/dreamerv3/utils/add_is_firsts_to_batch.py

@@ -0,0 +1,36 @@
+from typing import Any, List, Optional
+
+from ray.rllib.connectors.connector_v2 import ConnectorV2
+from ray.rllib.core.rl_module.rl_module import RLModule
+from ray.rllib.utils.annotations import override
+from ray.rllib.utils.typing import EpisodeType
+
+
+class AddIsFirstsToBatch(ConnectorV2):
+    """Adds the "is_first" column to the batch."""
+
+    @override(ConnectorV2)
+    def __call__(
+        self,
+        *,
+        rl_module: RLModule,
+        batch: Optional[Any],
+        episodes: List[EpisodeType],
+        explore: Optional[bool] = None,
+        shared_data: Optional[dict] = None,
+        **kwargs,
+    ) -> Any:
+        # If "is_first" already in batch, early out.
+        if "is_first" in batch:
+            return batch
+
+        for sa_episode in self.single_agent_episode_iterator(episodes):
+            self.add_batch_item(
+                batch,
+                "is_first",
+                item_to_add=(
+                    1.0 if sa_episode.t_started == 0 and len(sa_episode) == 0 else 0.0
+                ),
+                single_agent_episode=sa_episode,
+            )
+        return batch

rllib/algorithms/dreamerv3/utils/debugging.py

@@ -0,0 +1,189 @@
+import gymnasium as gym
+import numpy as np
+from gymnasium.envs.classic_control.cartpole import CartPoleEnv
+from PIL import Image, ImageDraw
+
+from ray.rllib.utils.framework import try_import_torch
+
+torch, _ = try_import_torch()
+
+
+class CartPoleDebug(CartPoleEnv):
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+
+        low = np.concatenate([np.array([0.0]), self.observation_space.low])
+        high = np.concatenate([np.array([1000.0]), self.observation_space.high])
+
+        self.observation_space = gym.spaces.Box(low, high, shape=(5,), dtype=np.float32)
+
+        self.timesteps_ = 0
+        self._next_action = 0
+        self._seed = 1
+
+    def reset(self, *, seed=None, options=None):
+        ret = super().reset(seed=self._seed)
+        self._seed += 1
+        self.timesteps_ = 0
+        self._next_action = 0
+        obs = np.concatenate([np.array([self.timesteps_]), ret[0]])
+        return obs, ret[1]
+
+    def step(self, action):
+        ret = super().step(self._next_action)
+
+        self.timesteps_ += 1
+        self._next_action = 0 if self._next_action else 1
+
+        obs = np.concatenate([np.array([self.timesteps_]), ret[0]])
+        reward = 0.1 * self.timesteps_
+        return (obs, reward) + ret[2:]
+
+
+gym.register("CartPoleDebug-v0", CartPoleDebug)
+cartpole_env = gym.make("CartPoleDebug-v0", render_mode="rgb_array")
+cartpole_env.reset()
+
+frozenlake_env = gym.make(
+    "FrozenLake-v1", render_mode="rgb_array", is_slippery=False, map_name="4x4"
+)  # desc=["SF", "HG"])
+frozenlake_env.reset()
+
+
+def create_cartpole_dream_image(
+    dreamed_obs,  # real space (not symlog'd)
+    dreamed_V,  # real space (not symlog'd)
+    dreamed_a,
+    dreamed_r_tp1,  # real space (not symlog'd)
+    dreamed_ri_tp1,  # intrinsic reward
+    dreamed_c_tp1,  # continue flag
+    value_target,  # real space (not symlog'd)
+    initial_h,
+    as_tensor=False,
+):
+    # CartPoleDebug
+    if dreamed_obs.shape == (5,):
+        # Set the state of our env to the given observation.
+        cartpole_env.unwrapped.state = np.array(dreamed_obs[1:], dtype=np.float32)
+    # Normal CartPole-v1
+    else:
+        cartpole_env.unwrapped.state = np.array(dreamed_obs, dtype=np.float32)
+
+    # Produce an RGB-image of the current state.
+    rgb_array = cartpole_env.render()
+
+    # Add value-, action-, reward-, and continue-prediction information.
+    image = Image.fromarray(rgb_array)
+    draw_obj = ImageDraw.Draw(image)
+
+    # fnt = ImageFont.load_default(size=40)
+
+    draw_obj.text(
+        (5, 6), f"Vt={dreamed_V:.2f} (Rt={value_target:.2f})", fill=(0, 0, 0)
+    )  # , font=fnt.font, size=30)
+    draw_obj.text(
+        (5, 18),
+        f"at={'<--' if dreamed_a == 0 else '-->'} ({dreamed_a})",
+        fill=(0, 0, 0),
+    )
+    draw_obj.text((5, 30), f"rt+1={dreamed_r_tp1:.2f}", fill=(0, 0, 0))
+    if dreamed_ri_tp1 is not None:
+        draw_obj.text((5, 42), f"rit+1={dreamed_ri_tp1:.6f}", fill=(0, 0, 0))
+    draw_obj.text((5, 54), f"ct+1={dreamed_c_tp1}", fill=(0, 0, 0))
+    draw_obj.text((5, 66), f"|h|t={np.mean(np.abs(initial_h)):.5f}", fill=(0, 0, 0))
+
+    if dreamed_obs.shape == (5,):
+        draw_obj.text((20, 100), f"t={dreamed_obs[0]}", fill=(0, 0, 0))
+
+    # Return image.
+    np_img = np.asarray(image)
+    if as_tensor:
+        return torch.from_numpy(np_img, dtype=torch.uint8)
+    return np_img
+
+
+def create_frozenlake_dream_image(
+    dreamed_obs,  # real space (not symlog'd)
+    dreamed_V,  # real space (not symlog'd)
+    dreamed_a,
+    dreamed_r_tp1,  # real space (not symlog'd)
+    dreamed_ri_tp1,  # intrinsic reward
+    dreamed_c_tp1,  # continue flag
+    value_target,  # real space (not symlog'd)
+    initial_h,
+    as_tensor=False,
+):
+    frozenlake_env.unwrapped.s = np.argmax(dreamed_obs, axis=0)
+
+    # Produce an RGB-image of the current state.
+    rgb_array = frozenlake_env.render()
+
+    # Add value-, action-, reward-, and continue-prediction information.
+    image = Image.fromarray(rgb_array)
+    draw_obj = ImageDraw.Draw(image)
+
+    draw_obj.text((5, 6), f"Vt={dreamed_V:.2f} (Rt={value_target:.2f})", fill=(0, 0, 0))
+    action_arrow = (
+        "<--"
+        if dreamed_a == 0
+        else "v"
+        if dreamed_a == 1
+        else "-->"
+        if dreamed_a == 2
+        else "^"
+    )
+    draw_obj.text((5, 18), f"at={action_arrow} ({dreamed_a})", fill=(0, 0, 0))
+    draw_obj.text((5, 30), f"rt+1={dreamed_r_tp1:.2f}", fill=(0, 0, 0))
+    if dreamed_ri_tp1 is not None:
+        draw_obj.text((5, 42), f"rit+1={dreamed_ri_tp1:.6f}", fill=(0, 0, 0))
+    draw_obj.text((5, 54), f"ct+1={dreamed_c_tp1}", fill=(0, 0, 0))
+    draw_obj.text((5, 66), f"|h|t={np.mean(np.abs(initial_h)):.5f}", fill=(0, 0, 0))
+
+    # Return image.
+    np_img = np.asarray(image)
+    if as_tensor:
+        return torch.from_numpy(np_img, dtype=torch.uint8)
+    return np_img
+
+
+if __name__ == "__main__":
+    # CartPole debug.
+    rgb_array = create_cartpole_dream_image(
+        dreamed_obs=np.array([100.0, 1.0, -0.01, 1.5, 0.02]),
+        dreamed_V=4.3,
+        dreamed_a=1,
+        dreamed_r_tp1=1.0,
+        dreamed_c_tp1=True,
+        initial_h=0.0,
+        value_target=8.0,
+    )
+    # ImageFont.load("arial.pil")
+    image = Image.fromarray(rgb_array)
+    image.show()
+
+    # Normal CartPole.
+    rgb_array = create_cartpole_dream_image(
+        dreamed_obs=np.array([1.0, -0.01, 1.5, 0.02]),
+        dreamed_V=4.3,
+        dreamed_a=1,
+        dreamed_r_tp1=1.0,
+        dreamed_c_tp1=True,
+        initial_h=0.1,
+        value_target=8.0,
+    )
+    # ImageFont.load("arial.pil")
+    image = Image.fromarray(rgb_array)
+    image.show()
+
+    # Frozenlake
+    rgb_array = create_frozenlake_dream_image(
+        dreamed_obs=np.array([1.0] + [0.0] * (frozenlake_env.observation_space.n - 1)),
+        dreamed_V=4.3,
+        dreamed_a=1,
+        dreamed_r_tp1=1.0,
+        dreamed_c_tp1=True,
+        initial_h=0.1,
+        value_target=8.0,
+    )
+    image = Image.fromarray(rgb_array)
+    image.show()

rllib/algorithms/dreamerv3/utils/summaries.py

@@ -0,0 +1,407 @@
+"""
+[1] Mastering Diverse Domains through World Models - 2023
+D. Hafner, J. Pasukonis, J. Ba, T. Lillicrap
+https://arxiv.org/pdf/2301.04104v1.pdf
+
+[2] Mastering Atari with Discrete World Models - 2021
+D. Hafner, T. Lillicrap, M. Norouzi, J. Ba
+https://arxiv.org/pdf/2010.02193.pdf
+"""
+import numpy as np
+
+from ray.rllib.algorithms.dreamerv3.utils.debugging import (
+    create_cartpole_dream_image,
+    create_frozenlake_dream_image,
+)
+from ray.rllib.core import DEFAULT_MODULE_ID
+from ray.rllib.core.columns import Columns
+from ray.rllib.utils.framework import try_import_torch
+from ray.rllib.utils.metrics import (
+    LEARNER_RESULTS,
+    REPLAY_BUFFER_RESULTS,
+)
+from ray.rllib.utils.torch_utils import inverse_symlog
+
+torch, _ = try_import_torch()
+
+
+def reconstruct_obs_from_h_and_z(
+    h_t0_to_H,
+    z_t0_to_H,
+    dreamer_model,
+    obs_dims_shape,
+    framework="torch",
+):
+    """Returns"""
+    shape = h_t0_to_H.shape
+    T = shape[0]  # inputs are time-major
+    B = shape[1]
+    # Compute actual observations using h and z and the decoder net.
+    # Note that the last h-state (T+1) is NOT used here as it's already part of
+    # a new trajectory.
+    # Use mean() of the Gaussian, no sample! -> No need to construct dist object here.
+    if framework == "torch":
+        device = next(iter(dreamer_model.world_model.decoder.parameters())).device
+        reconstructed_obs_distr_means_TxB = (
+            dreamer_model.world_model.decoder(
+                # Fold time rank.
+                h=torch.from_numpy(h_t0_to_H).reshape((T * B, -1)).to(device),
+                z=torch.from_numpy(z_t0_to_H)
+                .reshape((T * B,) + z_t0_to_H.shape[2:])
+                .to(device),
+            )
+            .detach()
+            .cpu()
+            .numpy()
+        )
+    else:
+        reconstructed_obs_distr_means_TxB = dreamer_model.world_model.decoder(
+            # Fold time rank.
+            h=h_t0_to_H.reshape((T * B, -1)),
+            z=z_t0_to_H.reshape((T * B,) + z_t0_to_H.shape[2:]),
+        )
+
+    # Unfold time rank again.
+    reconstructed_obs_T_B = np.reshape(
+        reconstructed_obs_distr_means_TxB, (T, B) + obs_dims_shape
+    )
+    # Return inverse symlog'd (real env obs space) reconstructed observations.
+    return reconstructed_obs_T_B
+
+
+def report_dreamed_trajectory(
+    *,
+    results,
+    env,
+    dreamer_model,
+    obs_dims_shape,
+    batch_indices=(0,),
+    desc=None,
+    include_images=True,
+    framework="torch",
+):
+    if not include_images:
+        return
+
+    dream_data = results["dream_data"]
+    dreamed_obs_H_B = reconstruct_obs_from_h_and_z(
+        h_t0_to_H=dream_data["h_states_t0_to_H_BxT"],
+        z_t0_to_H=dream_data["z_states_prior_t0_to_H_BxT"],
+        dreamer_model=dreamer_model,
+        obs_dims_shape=obs_dims_shape,
+        framework=framework,
+    )
+    func = (
+        create_cartpole_dream_image
+        if env.startswith("CartPole")
+        else create_frozenlake_dream_image
+    )
+    # Take 0th dreamed trajectory and produce series of images.
+    for b in batch_indices:
+        images = []
+        for t in range(len(dreamed_obs_H_B) - 1):
+            images.append(
+                func(
+                    dreamed_obs=dreamed_obs_H_B[t][b],
+                    dreamed_V=dream_data["values_dreamed_t0_to_H_BxT"][t][b],
+                    dreamed_a=(dream_data["actions_ints_dreamed_t0_to_H_BxT"][t][b]),
+                    dreamed_r_tp1=(dream_data["rewards_dreamed_t0_to_H_BxT"][t + 1][b]),
+                    # `DISAGREE_intrinsic_rewards_H_B` are shifted by 1 already
+                    # (from t1 to H, not t0 to H like all other data here).
+                    dreamed_ri_tp1=(
+                        results["DISAGREE_intrinsic_rewards_H_BxT"][t][b]
+                        if "DISAGREE_intrinsic_rewards_H_BxT" in results
+                        else None
+                    ),
+                    dreamed_c_tp1=(
+                        dream_data["continues_dreamed_t0_to_H_BxT"][t + 1][b]
+                    ),
+                    value_target=results["VALUE_TARGETS_H_BxT"][t][b],
+                    initial_h=dream_data["h_states_t0_to_H_BxT"][t][b],
+                    as_tensor=True,
+                ).numpy()
+            )
+        # Concat images along width-axis (so they show as a "film sequence" next to each
+        # other).
+        results.update(
+            {
+                f"dreamed_trajectories{('_'+desc) if desc else ''}_B{b}": (
+                    np.concatenate(images, axis=1)
+                ),
+            }
+        )
+
+
+def report_predicted_vs_sampled_obs(
+    *,
+    metrics,
+    sample,
+    batch_size_B,
+    batch_length_T,
+    symlog_obs: bool = True,
+    do_report: bool = True,
+):
+    """Summarizes sampled data (from the replay buffer) vs world-model predictions.
+
+    World model predictions are based on the posterior states (z computed from actual
+    observation encoder input + the current h-states).
+
+    Observations: Computes MSE (sampled vs predicted/recreated) over all features.
+    For image observations, also creates direct image comparisons (sampled images
+    vs predicted (posterior) ones).
+    Rewards: Compute MSE (sampled vs predicted).
+    Continues: Compute MSE (sampled vs predicted).
+
+    Args:
+        metrics: The MetricsLogger object of the DreamerV3 algo.
+        sample: The sampled data (dict) from the replay buffer. Already torch-tensor
+            converted.
+        batch_size_B: The batch size (B). This is the number of trajectories sampled
+            from the buffer.
+        batch_length_T: The batch length (T). This is the length of an individual
+            trajectory sampled from the buffer.
+        do_report: Whether to actually log the report (default). If this is set to
+            False, this function serves as a clean-up on the given metrics, making sure
+            they do NOT contain anymore any (spacious) data relevant for producing
+            the report/videos.
+    """
+    fwd_output_key = (
+        LEARNER_RESULTS,
+        DEFAULT_MODULE_ID,
+        "WORLD_MODEL_fwd_out_obs_distribution_means_b0xT",
+    )
+    # logged as a non-reduced item (still a list)
+    predicted_observation_means_single_example = metrics.peek(
+        fwd_output_key, default=[None]
+    )[-1]
+    metrics.delete(fwd_output_key, key_error=False)
+
+    final_result_key = (
+        f"WORLD_MODEL_sampled_vs_predicted_posterior_b0x{batch_length_T}_videos"
+    )
+    if not do_report:
+        metrics.delete(final_result_key, key_error=False)
+        return
+
+    _report_obs(
+        metrics=metrics,
+        computed_float_obs_B_T_dims=np.reshape(
+            predicted_observation_means_single_example,
+            # WandB videos need to be channels first.
+            (1, batch_length_T) + sample[Columns.OBS].shape[2:],
+        ),
+        sampled_obs_B_T_dims=sample[Columns.OBS][0:1],
+        metrics_key=final_result_key,
+        symlog_obs=symlog_obs,
+    )
+
+
+def report_dreamed_eval_trajectory_vs_samples(
+    *,
+    metrics,
+    sample,
+    burn_in_T,
+    dreamed_T,
+    dreamer_model,
+    symlog_obs: bool = True,
+    do_report: bool = True,
+    framework="torch",
+) -> None:
+    """Logs dreamed observations, rewards, continues and compares them vs sampled data.
+
+    For obs, we'll try to create videos (side-by-side comparison) of the dreamed,
+    recreated-from-prior obs vs the sampled ones (over dreamed_T timesteps).
+
+    Args:
+        metrics: The MetricsLogger object of the DreamerV3 algo.
+        sample: The sampled data (dict) from the replay buffer. Already torch-tensor
+            converted.
+        burn_in_T: The number of burn-in timesteps (these will be skipped over in the
+            reported video comparisons and MSEs).
+        dreamed_T: The number of timesteps to produce dreamed data for.
+        dreamer_model: The DreamerModel to use to create observation vectors/images
+            from dreamed h- and (prior) z-states.
+        symlog_obs: Whether to inverse-symlog the computed observations or not. Set this
+            to True for environments, in which we should symlog the observations.
+        do_report: Whether to actually log the report (default). If this is set to
+            False, this function serves as a clean-up on the given metrics, making sure
+            they do NOT contain anymore any (spacious) data relevant for producing
+            the report/videos.
+    """
+    dream_data = metrics.peek(
+        (LEARNER_RESULTS, DEFAULT_MODULE_ID, "dream_data"),
+        default={},
+    )
+    metrics.delete(LEARNER_RESULTS, DEFAULT_MODULE_ID, "dream_data", key_error=False)
+
+    final_result_key_obs = f"EVALUATION_sampled_vs_dreamed_prior_H{dreamed_T}_obs"
+    final_result_key_rew = (
+        f"EVALUATION_sampled_vs_dreamed_prior_H{dreamed_T}_rewards_MSE"
+    )
+    final_result_key_cont = (
+        f"EVALUATION_sampled_vs_dreamed_prior_H{dreamed_T}_continues_MSE"
+    )
+    if not do_report:
+        metrics.delete(final_result_key_obs, key_error=False)
+        metrics.delete(final_result_key_rew, key_error=False)
+        metrics.delete(final_result_key_cont, key_error=False)
+        return
+
+    # Obs MSE.
+    dreamed_obs_H_B = reconstruct_obs_from_h_and_z(
+        h_t0_to_H=dream_data["h_states_t0_to_H_Bx1"][0],  # [0] b/c reduce=None (list)
+        z_t0_to_H=dream_data["z_states_prior_t0_to_H_Bx1"][0],
+        dreamer_model=dreamer_model,
+        obs_dims_shape=sample[Columns.OBS].shape[2:],
+        framework=framework,
+    )
+    t0 = burn_in_T
+    tH = t0 + dreamed_T
+    # Observation MSE and - if applicable - images comparisons.
+    _report_obs(
+        metrics=metrics,
+        # WandB videos need to be 5D (B, L, c, h, w) -> transpose/swap H and B axes.
+        computed_float_obs_B_T_dims=np.swapaxes(dreamed_obs_H_B, 0, 1)[
+            0:1
+        ],  # for now: only B=1
+        sampled_obs_B_T_dims=sample[Columns.OBS][0:1, t0:tH],
+        metrics_key=final_result_key_obs,
+        symlog_obs=symlog_obs,
+    )
+
+    # Reward MSE.
+    _report_rewards(
+        metrics=metrics,
+        computed_rewards=dream_data["rewards_dreamed_t0_to_H_Bx1"][0],
+        sampled_rewards=sample[Columns.REWARDS][:, t0:tH],
+        metrics_key=final_result_key_rew,
+    )
+
+    # Continues MSE.
+    _report_continues(
+        metrics=metrics,
+        computed_continues=dream_data["continues_dreamed_t0_to_H_Bx1"][0],
+        sampled_continues=(1.0 - sample["is_terminated"])[:, t0:tH],
+        metrics_key=final_result_key_cont,
+    )
+
+
+def report_sampling_and_replay_buffer(*, metrics, replay_buffer):
+    episodes_in_buffer = replay_buffer.get_num_episodes()
+    ts_in_buffer = replay_buffer.get_num_timesteps()
+    replayed_steps = replay_buffer.get_sampled_timesteps()
+    added_steps = replay_buffer.get_added_timesteps()
+
+    # Summarize buffer, sampling, and train ratio stats.
+    metrics.log_dict(
+        {
+            "capacity": replay_buffer.capacity,
+            "size_num_episodes": episodes_in_buffer,
+            "size_timesteps": ts_in_buffer,
+            "replayed_steps": replayed_steps,
+            "added_steps": added_steps,
+        },
+        key=REPLAY_BUFFER_RESULTS,
+        window=1,
+    )  # window=1 b/c these are current (total count/state) values.
+
+
+def _report_obs(
+    *,
+    metrics,
+    computed_float_obs_B_T_dims,
+    sampled_obs_B_T_dims,
+    metrics_key,
+    symlog_obs,
+):
+    """Summarizes computed- vs sampled observations: MSE and (if applicable) images.
+
+    Args:
+        metrics: The MetricsLogger object of the DreamerV3 algo.
+        computed_float_obs_B_T_dims: Computed float observations
+            (not clipped, not cast'd). Shape=(B, T, [dims ...]).
+        sampled_obs_B_T_dims: Sampled observations (as-is from the environment, meaning
+            this could be uint8, 0-255 clipped images). Shape=(B, T, [dims ...]).
+        metrics_key: The metrics key (or key sequence) under which to log ths resulting
+            video sequence.
+        symlog_obs: Whether to inverse-symlog the computed observations or not. Set this
+            to True for environments, in which we should symlog the observations.
+    """
+    # Videos: Create summary, comparing computed images with actual sampled ones.
+    # 4=[B, T, w, h] grayscale image; 5=[B, T, w, h, C] RGB image.
+    if len(sampled_obs_B_T_dims.shape) in [4, 5]:
+        # WandB videos need to be channels first.
+        transpose_axes = (
+            (0, 1, 4, 2, 3) if len(sampled_obs_B_T_dims.shape) == 5 else (0, 3, 1, 2)
+        )
+
+        if symlog_obs:
+            computed_float_obs_B_T_dims = inverse_symlog(computed_float_obs_B_T_dims)
+
+        # Restore image pixels from normalized (non-symlog'd) data.
+        if not symlog_obs:
+            computed_float_obs_B_T_dims = (computed_float_obs_B_T_dims + 1.0) * 128
+            sampled_obs_B_T_dims = (sampled_obs_B_T_dims + 1.0) * 128
+            sampled_obs_B_T_dims = np.clip(sampled_obs_B_T_dims, 0.0, 255.0).astype(
+                np.uint8
+            )
+            sampled_obs_B_T_dims = np.transpose(sampled_obs_B_T_dims, transpose_axes)
+        computed_images = np.clip(computed_float_obs_B_T_dims, 0.0, 255.0).astype(
+            np.uint8
+        )
+        computed_images = np.transpose(computed_images, transpose_axes)
+        # Concat sampled and computed images along the height axis (3) such that
+        # real images show below respective predicted ones.
+        # (B, T, C, h, w)
+        sampled_vs_computed_images = np.concatenate(
+            [computed_images, sampled_obs_B_T_dims],
+            axis=-1,  # concat on width axis (looks nicer)
+        )
+        # Add grayscale dim, if necessary.
+        if len(sampled_obs_B_T_dims.shape) == 2 + 2:
+            sampled_vs_computed_images = np.expand_dims(sampled_vs_computed_images, -1)
+
+        metrics.log_value(
+            metrics_key,
+            sampled_vs_computed_images,
+            reduce="item_series",  # No reduction, we want the obs tensor to stay in-tact.
+            window=1,
+        )
+
+
+def _report_rewards(
+    *,
+    metrics,
+    computed_rewards,
+    sampled_rewards,
+    metrics_key,
+):
+    mse_sampled_vs_computed_rewards = np.mean(
+        np.square(computed_rewards - sampled_rewards)
+    )
+    mse_sampled_vs_computed_rewards = np.mean(mse_sampled_vs_computed_rewards)
+    metrics.log_value(
+        metrics_key,
+        mse_sampled_vs_computed_rewards,
+        window=1,
+    )
+
+
+def _report_continues(
+    *,
+    metrics,
+    computed_continues,
+    sampled_continues,
+    metrics_key,
+):
+    # Continue MSE.
+    mse_sampled_vs_computed_continues = np.mean(
+        np.square(
+            computed_continues - sampled_continues.astype(computed_continues.dtype)
+        )
+    )
+    metrics.log_value(
+        metrics_key,
+        mse_sampled_vs_computed_continues,
+        window=1,
+    )

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

rllib/algorithms/iql/__init__.py

@@ -0,0 +1,6 @@
+from ray.rllib.algorithms.iql.iql import IQL, IQLConfig
+
+__all__ = [
+    "IQL",
+    "IQLConfig",
+]

rllib/algorithms/iql/default_iql_rl_module.py

@@ -0,0 +1,35 @@
+from ray.rllib.algorithms.sac.default_sac_rl_module import DefaultSACRLModule
+from ray.rllib.core.models.configs import MLPHeadConfig
+from ray.rllib.core.rl_module.apis.value_function_api import ValueFunctionAPI
+from ray.rllib.utils.annotations import (
+    OverrideToImplementCustomLogic_CallToSuperRecommended,
+    override,
+)
+
+
+class DefaultIQLRLModule(DefaultSACRLModule, ValueFunctionAPI):
+    @override(DefaultSACRLModule)
+    def setup(self):
+        # Setup the `DefaultSACRLModule` to get the catalog.
+        super().setup()
+
+        # Only, if the `RLModule` is used on a `Learner` we build the value network.
+        if not self.inference_only:
+            # Build the encoder for the value function.
+            self.vf_encoder = self.catalog.build_encoder(framework=self.framework)
+
+            # Build the vf head.
+            self.vf = MLPHeadConfig(
+                input_dims=self.catalog.latent_dims,
+                # Note, we use the same layers as for the policy and Q-network.
+                hidden_layer_dims=self.catalog.pi_and_qf_head_hiddens,
+                hidden_layer_activation=self.catalog.pi_and_qf_head_activation,
+                output_layer_activation="linear",
+                output_layer_dim=1,
+            ).build(framework=self.framework)
+
+    @override(DefaultSACRLModule)
+    @OverrideToImplementCustomLogic_CallToSuperRecommended
+    def get_non_inference_attributes(self):
+        # Use all of `super`'s attributes and add the value function attributes.
+        return super().get_non_inference_attributes() + ["vf_encoder", "vf"]

rllib/algorithms/iql/iql.py

@@ -0,0 +1,228 @@
+from typing import Optional, Type, Union
+
+from ray.rllib.algorithms.algorithm_config import AlgorithmConfig, NotProvided
+from ray.rllib.algorithms.marwil.marwil import MARWIL, MARWILConfig
+from ray.rllib.connectors.common.add_observations_from_episodes_to_batch import (
+    AddObservationsFromEpisodesToBatch,
+)
+from ray.rllib.connectors.learner.add_next_observations_from_episodes_to_train_batch import (  # noqa
+    AddNextObservationsFromEpisodesToTrainBatch,
+)
+from ray.rllib.core.learner.learner import Learner
+from ray.rllib.core.rl_module.rl_module import RLModuleSpec
+from ray.rllib.utils.annotations import override
+from ray.rllib.utils.typing import LearningRateOrSchedule, RLModuleSpecType
+
+
+class IQLConfig(MARWILConfig):
+    """Defines a configuration class from which a new IQL Algorithm can be built
+
+    .. testcode::
+        :skipif: True
+
+        from ray.rllib.algorithms.iql import IQLConfig
+        # Run this from the ray directory root.
+        config = IQLConfig().training(actor_lr=0.00001, gamma=0.99)
+        config = config.offline_data(
+            input_="./rllib/offline/tests/data/pendulum/pendulum-v1_enormous")
+
+        # Build an Algorithm object from the config and run 1 training iteration.
+        algo = config.build()
+        algo.train()
+
+    .. testcode::
+        :skipif: True
+
+        from ray.rllib.algorithms.iql import IQLConfig
+        from ray import tune
+        config = IQLConfig()
+        # Print out some default values.
+        print(config.beta)
+        # Update the config object.
+        config.training(
+            lr=tune.grid_search([0.001, 0.0001]), beta=0.75
+        )
+        # Set the config object's data path.
+        # Run this from the ray directory root.
+        config.offline_data(
+            input_="./rllib/offline/tests/data/pendulum/pendulum-v1_enormous"
+        )
+        # Set the config object's env, used for evaluation.
+        config.environment(env="Pendulum-v1")
+        # Use to_dict() to get the old-style python config dict
+        # when running with tune.
+        tune.Tuner(
+            "IQL",
+            param_space=config.to_dict(),
+        ).fit()
+    """
+
+    def __init__(self, algo_class=None):
+        super().__init__(algo_class=algo_class or IQL)
+
+        # fmt: off
+        # __sphinx_doc_begin__
+        # The temperature for the actor loss.
+        self.beta = 0.1
+
+        # The expectile to use in expectile regression.
+        self.expectile = 0.8
+
+        # The learning rates for the actor, critic and value network(s).
+        self.actor_lr = 3e-4
+        self.critic_lr = 3e-4
+        self.value_lr = 3e-4
+        # Set `lr` parameter to `None` and ensure it is not used.
+        self.lr = None
+
+        # If a twin-Q architecture should be used (advisable).
+        self.twin_q = True
+
+        # How often the target network should be updated.
+        self.target_network_update_freq = 0
+        # The weight for Polyak averaging.
+        self.tau = 1.0
+
+        # __sphinx_doc_end__
+        # fmt: on
+
+    @override(MARWILConfig)
+    def training(
+        self,
+        *,
+        twin_q: Optional[bool] = NotProvided,
+        expectile: Optional[float] = NotProvided,
+        actor_lr: Optional[LearningRateOrSchedule] = NotProvided,
+        critic_lr: Optional[LearningRateOrSchedule] = NotProvided,
+        value_lr: Optional[LearningRateOrSchedule] = NotProvided,
+        target_network_update_freq: Optional[int] = NotProvided,
+        tau: Optional[float] = NotProvided,
+        **kwargs,
+    ) -> "IQLConfig":
+        """Sets the training related configuration.
+
+        Args:
+            beta: The temperature to scaling advantages in exponential terms.
+                Must be >> 0.0. The higher this parameter the less greedy
+                (exploitative) the policy becomes. It also means that the policy
+                is fitting less to the best actions in the dataset.
+            twin_q: If a twin-Q architecture should be used (advisable).
+            expectile: The expectile to use in expectile regression for the value
+                function. For high expectiles the value function tries to match
+                the upper tail of the Q-value distribution.
+            actor_lr: The learning rate for the actor network. Actor learning rates
+                greater than critic learning rates work well in experiments.
+            critic_lr: The learning rate for the Q-network. Critic learning rates
+                greater than value function learning rates work well in experiments.
+            value_lr: The learning rate for the value function network.
+            target_network_update_freq: The number of timesteps in between the target
+                Q-network is fixed. Note, too high values here could harm convergence.
+                The target network is updated via Polyak-averaging.
+            tau: The update parameter for Polyak-averaging of the target Q-network.
+                The higher this value the faster the weights move towards the actual
+                Q-network.
+
+        Return:
+            This updated `AlgorithmConfig` object.
+        """
+        super().training(**kwargs)
+
+        if twin_q is not NotProvided:
+            self.twin_q = twin_q
+        if expectile is not NotProvided:
+            self.expectile = expectile
+        if actor_lr is not NotProvided:
+            self.actor_lr = actor_lr
+        if critic_lr is not NotProvided:
+            self.critic_lr = critic_lr
+        if value_lr is not NotProvided:
+            self.value_lr = value_lr
+        if target_network_update_freq is not NotProvided:
+            self.target_network_update_freq = target_network_update_freq
+        if tau is not NotProvided:
+            self.tau = tau
+
+        return self
+
+    @override(MARWILConfig)
+    def get_default_learner_class(self) -> Union[Type["Learner"], str]:
+        if self.framework_str == "torch":
+            from ray.rllib.algorithms.iql.torch.iql_torch_learner import IQLTorchLearner
+
+            return IQLTorchLearner
+        else:
+            raise ValueError(
+                f"The framework {self.framework_str} is not supported. "
+                "Use `'torch'` instead."
+            )
+
+    @override(MARWILConfig)
+    def get_default_rl_module_spec(self) -> RLModuleSpecType:
+        if self.framework_str == "torch":
+            from ray.rllib.algorithms.iql.torch.default_iql_torch_rl_module import (
+                DefaultIQLTorchRLModule,
+            )
+
+            return RLModuleSpec(module_class=DefaultIQLTorchRLModule)
+        else:
+            raise ValueError(
+                f"The framework {self.framework_str} is not supported. "
+                "Use `torch` instead."
+            )
+
+    @override(MARWILConfig)
+    def build_learner_connector(
+        self,
+        input_observation_space,
+        input_action_space,
+        device=None,
+    ):
+        pipeline = super().build_learner_connector(
+            input_observation_space=input_observation_space,
+            input_action_space=input_action_space,
+            device=device,
+        )
+
+        # Remove unneeded connectors from the MARWIL connector pipeline.
+        pipeline.remove("AddOneTsToEpisodesAndTruncate")
+        pipeline.remove("GeneralAdvantageEstimation")
+
+        # Prepend the "add-NEXT_OBS-from-episodes-to-train-batch" connector piece (right
+        # after the corresponding "add-OBS-..." default piece).
+        pipeline.insert_after(
+            AddObservationsFromEpisodesToBatch,
+            AddNextObservationsFromEpisodesToTrainBatch(),
+        )
+
+        return pipeline
+
+    @override(MARWILConfig)
+    def validate(self) -> None:
+        # Call super's validation method.
+        super().validate()
+
+        # Ensure hyperparameters are meaningful.
+        if self.beta <= 0.0:
+            self._value_error(
+                "For meaningful results, `beta` (temperature) parameter must be >> 0.0!"
+            )
+        if not 0.0 < self.expectile < 1.0:
+            self._value_error(
+                "For meaningful results, `expectile` parameter must be in (0, 1)."
+            )
+
+    @property
+    def _model_config_auto_includes(self):
+        return super()._model_config_auto_includes | {"twin_q": self.twin_q}
+
+
+class IQL(MARWIL):
+    """Implicit Q-learning (derived from MARWIL).
+
+    Uses MARWIL training step.
+    """
+
+    @classmethod
+    @override(MARWIL)
+    def get_default_config(cls) -> AlgorithmConfig:
+        return IQLConfig()

rllib/algorithms/iql/iql_learner.py

@@ -0,0 +1,84 @@
+from typing import Dict
+
+from ray.rllib.algorithms.dqn.dqn_learner import DQNLearner
+from ray.rllib.utils.annotations import (
+    OverrideToImplementCustomLogic_CallToSuperRecommended,
+    override,
+)
+from ray.rllib.utils.lambda_defaultdict import LambdaDefaultDict
+from ray.rllib.utils.typing import ModuleID, TensorType
+
+QF_TARGET_PREDS = "qf_target_preds"
+VF_PREDS_NEXT = "vf_preds_next"
+VF_LOSS = "value_loss"
+
+
+class IQLLearner(DQNLearner):
+    @OverrideToImplementCustomLogic_CallToSuperRecommended
+    @override(DQNLearner)
+    def build(self) -> None:
+        # Build the `DQNLearner` (builds the target network).
+        super().build()
+
+        # Define the expectile parameter(s).
+        self.expectile: Dict[ModuleID, TensorType] = LambdaDefaultDict(
+            lambda module_id: self._get_tensor_variable(
+                # Note, we want to train with a certain expectile.
+                [self.config.get_config_for_module(module_id).expectile],
+                trainable=False,
+            )
+        )
+
+        # Define the temperature for the actor advantage loss.
+        self.temperature: Dict[ModuleID, TensorType] = LambdaDefaultDict(
+            lambda module_id: self._get_tensor_variable(
+                # Note, we want to train with a certain expectile.
+                [self.config.get_config_for_module(module_id).beta],
+                trainable=False,
+            )
+        )
+
+        # Store loss tensors here temporarily inside the loss function for (exact)
+        # consumption later by the compute gradients function.
+        # Keys=(module_id, optimizer_name), values=loss tensors (in-graph).
+        self._temp_losses = {}
+
+    @override(DQNLearner)
+    def remove_module(self, module_id: ModuleID) -> None:
+        """Removes the expectile and temperature for removed modules."""
+        # First call `super`'s `remove_module` method.
+        super().remove_module(module_id)
+        # Remove the expectile from the mapping.
+        self.expectile.pop(module_id, None)
+        # Remove the temperature from the mapping.
+        self.temperature.pop(module_id, None)
+
+    @override(DQNLearner)
+    def add_module(
+        self,
+        *,
+        module_id,
+        module_spec,
+        config_overrides=None,
+        new_should_module_be_updated=None
+    ):
+        """Adds the expectile and temperature for new modules."""
+        # First call `super`'s `add_module` method.
+        super().add_module(
+            module_id=module_id,
+            module_spec=module_spec,
+            config_overrides=config_overrides,
+            new_should_module_be_updated=new_should_module_be_updated,
+        )
+        # Add the expectile to the mapping.
+        self.expectile[module_id] = self._get_tensor_variable(
+            # Note, we want to train with a certain expectile.
+            [self.config.get_config_for_module(module_id).beta],
+            trainable=False,
+        )
+        # Add the temperature to the mapping.
+        self.temperature[module_id] = self._get_tensor_variable(
+            # Note, we want to train with a certain expectile.
+            [self.config.get_config_for_module(module_id).beta],
+            trainable=False,
+        )