Add files using upload-large-folder tool

lib/python3.12/site-packages/ray/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",
+]

lib/python3.12/site-packages/ray/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",
+]

lib/python3.12/site-packages/ray/rllib/algorithms/appo/appo.py

@@ -0,0 +1,435 @@
+"""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"
+
+
+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
+
+        # 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 = 4
+        self.circular_buffer_iterations_per_batch = 2
+
+        # 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,
+        circular_buffer_num_batches: Optional[int] = NotProvided,
+        circular_buffer_iterations_per_batch: 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))
+            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.
+
+        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 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
+            )
+
+        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

lib/python3.12/site-packages/ray/rllib/algorithms/appo/appo_learner.py

@@ -0,0 +1,147 @@
+import abc
+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._learner_thread_in_queue = CircularBuffer(
+            num_batches=self.config.circular_buffer_num_batches,
+            iterations_per_batch=self.config.circular_buffer_iterations_per_batch,
+        )
+
+        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)
+
+            last_update_ts_key = (module_id, LAST_TARGET_UPDATE_TS)
+            if isinstance(module.unwrapped(), TargetNetworkAPI) and (
+                curr_timestep - self.metrics.peek(last_update_ts_key, default=0)
+                >= (
+                    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="sum")
+                # Update the (single-value -> window=1) last updated timestep metric.
+                self.metrics.log_value(last_update_ts_key, curr_timestep, window=1)
+
+            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

lib/python3.12/site-packages/ray/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,
+)

lib/python3.12/site-packages/ray/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)

lib/python3.12/site-packages/ray/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"]

lib/python3.12/site-packages/ray/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

lib/python3.12/site-packages/ray/rllib/algorithms/appo/torch/appo_torch_learner.py

@@ -0,0 +1,234 @@
+"""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,
+    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,
+        )
+
+        # The discount factor that is used should be gamma except for timesteps where
+        # the episode is terminated. In that case, the discount factor should be 0.
+        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
+
+        # 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),
+        )
+
+        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]
+                ),
+            },
+            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,
+        )

lib/python3.12/site-packages/ray/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,
+)

lib/python3.12/site-packages/ray/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

lib/python3.12/site-packages/ray/rllib/algorithms/appo/utils.py

@@ -0,0 +1,174 @@
+"""
+[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
+
+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 as described in [1] for APPO.
+
+    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. Also, each batch
+    currently in the buffer can be sampled at most K times (after which it is also
+    discarded).
+    """
+
+    def __init__(self, num_batches: int, iterations_per_batch: int):
+        # N from the paper (buffer size).
+        self.num_batches = num_batches
+        # K ("replay coefficient") from the paper.
+        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()
+
+        self._rng = np.random.default_rng()
+
+        # Ray metrics
+        self._metrics_circular_buffer_add_time = Histogram(
+            name="rllib_utils_circular_buffer_add_time",
+            description="Time spent in CircularBuffer.add()",
+            boundaries=DEFAULT_HISTOGRAM_BOUNDARIES_SHORT_EVENTS,
+            tag_keys=("rllib",),
+        )
+        self._metrics_circular_buffer_add_time.set_default_tags(
+            {"rllib": self.__class__.__name__}
+        )
+
+        self._metrics_circular_buffer_add_ts_dropped = Counter(
+            name="rllib_utils_circular_buffer_add_ts_dropped_counter",
+            description="Total number of env steps dropped by the CircularBuffer.",
+            tag_keys=("rllib",),
+        )
+        self._metrics_circular_buffer_add_ts_dropped.set_default_tags(
+            {"rllib": self.__class__.__name__}
+        )
+
+        self._metrics_circular_buffer_sample_time = Histogram(
+            name="rllib_utils_circular_buffer_sample_time",
+            description="Time spent in CircularBuffer.sample()",
+            boundaries=DEFAULT_HISTOGRAM_BOUNDARIES_SHORT_EVENTS,
+            tag_keys=("rllib",),
+        )
+        self._metrics_circular_buffer_sample_time.set_default_tags(
+            {"rllib": self.__class__.__name__}
+        )
+
+    def add(self, batch):
+        # Add buffer and k=0 information to the deque.
+        with TimerAndPrometheusLogger(self._metrics_circular_buffer_add_time):
+            with self._lock:
+                dropped_entry = self._buffer[0]
+                for _ in range(self.iterations_per_batch):
+                    self._buffer.append(batch)
+                    self._indices.add(self._offset)
+                    self._indices.discard(self._offset - self._NxK)
+                    self._offset += 1
+                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.env_steps()
+                if dropped_ts > 0:
+                    self._metrics_circular_buffer_add_ts_dropped.inc(value=dropped_ts)
+
+        return dropped_ts
+
+    def sample(self):
+        # Only initially, the buffer may be empty -> Just wait for some time.
+        with TimerAndPrometheusLogger(self._metrics_circular_buffer_sample_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
+
+    @property
+    def filled(self):
+        """Whether the buffer has been filled once with at least `self.num_batches`."""
+        with self._lock:
+            return self._num_added >= self.num_batches
+
+    def __len__(self) -> int:
+        """Returns the number of actually valid (non-expired) batches in the buffer."""
+        with self._lock:
+            return len(self._indices)
+
+
+@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

lib/python3.12/site-packages/ray/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

lib/python3.12/site-packages/ray/rllib/algorithms/marwil/__init__.py

@@ -0,0 +1,18 @@
+from ray.rllib.algorithms.marwil.marwil import (
+    MARWIL,
+    MARWILConfig,
+)
+from ray.rllib.algorithms.marwil.marwil_tf_policy import (
+    MARWILTF1Policy,
+    MARWILTF2Policy,
+)
+from ray.rllib.algorithms.marwil.marwil_torch_policy import MARWILTorchPolicy
+
+__all__ = [
+    "MARWIL",
+    "MARWILConfig",
+    # @OldAPIStack
+    "MARWILTF1Policy",
+    "MARWILTF2Policy",
+    "MARWILTorchPolicy",
+]

lib/python3.12/site-packages/ray/rllib/algorithms/marwil/marwil.py

@@ -0,0 +1,536 @@
+from typing import Callable, Optional, Type, Union
+
+from typing_extensions import Self
+
+from ray._common.deprecation import deprecation_warning
+from ray.rllib.algorithms.algorithm import Algorithm
+from ray.rllib.algorithms.algorithm_config import AlgorithmConfig, NotProvided
+from ray.rllib.connectors.learner import (
+    AddNextObservationsFromEpisodesToTrainBatch,
+    AddObservationsFromEpisodesToBatch,
+    AddOneTsToEpisodesAndTruncate,
+    GeneralAdvantageEstimation,
+)
+from ray.rllib.core.learner.learner import Learner
+from ray.rllib.core.learner.training_data import TrainingData
+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.metrics import (
+    LEARNER_RESULTS,
+    LEARNER_UPDATE_TIMER,
+    NUM_AGENT_STEPS_SAMPLED,
+    NUM_ENV_STEPS_SAMPLED,
+    OFFLINE_SAMPLING_TIMER,
+    SAMPLE_TIMER,
+    SYNCH_WORKER_WEIGHTS_TIMER,
+    TIMERS,
+)
+from ray.rllib.utils.typing import (
+    EnvType,
+    ResultDict,
+    RLModuleSpecType,
+)
+from ray.tune.logger import Logger
+
+
+class MARWILConfig(AlgorithmConfig):
+    """Defines a configuration class from which a MARWIL Algorithm can be built.
+
+    .. testcode::
+
+        import gymnasium as gym
+        import numpy as np
+
+        from pathlib import Path
+        from ray.rllib.algorithms.marwil import MARWILConfig
+
+        # Get the base path (to ray/rllib)
+        base_path = Path(__file__).parents[2]
+        # Get the path to the data in rllib folder.
+        data_path = base_path / "tests/data/cartpole/cartpole-v1_large"
+
+        config = MARWILConfig()
+        # Enable the new API stack.
+        config.api_stack(
+            enable_rl_module_and_learner=True,
+            enable_env_runner_and_connector_v2=True,
+        )
+        # Define the environment for which to learn a policy
+        # from offline data.
+        config.environment(
+            observation_space=gym.spaces.Box(
+                np.array([-4.8, -np.inf, -0.41887903, -np.inf]),
+                np.array([4.8, np.inf, 0.41887903, np.inf]),
+                shape=(4,),
+                dtype=np.float32,
+            ),
+            action_space=gym.spaces.Discrete(2),
+        )
+        # Set the training parameters.
+        config.training(
+            beta=1.0,
+            lr=1e-5,
+            gamma=0.99,
+            # We must define a train batch size for each
+            # learner (here 1 local learner).
+            train_batch_size_per_learner=2000,
+        )
+        # Define the data source for offline data.
+        config.offline_data(
+            input_=[data_path.as_posix()],
+            # Run exactly one update per training iteration.
+            dataset_num_iters_per_learner=1,
+        )
+
+        # Build an `Algorithm` object from the config and run 1 training
+        # iteration.
+        algo = config.build()
+        algo.train()
+
+    .. testcode::
+
+        import gymnasium as gym
+        import numpy as np
+
+        from pathlib import Path
+        from ray.rllib.algorithms.marwil import MARWILConfig
+        from ray import tune
+
+        # Get the base path (to ray/rllib)
+        base_path = Path(__file__).parents[2]
+        # Get the path to the data in rllib folder.
+        data_path = base_path / "tests/data/cartpole/cartpole-v1_large"
+
+        config = MARWILConfig()
+        # Enable the new API stack.
+        config.api_stack(
+            enable_rl_module_and_learner=True,
+            enable_env_runner_and_connector_v2=True,
+        )
+        # Print out some default values
+        print(f"beta: {config.beta}")
+        # Update the config object.
+        config.training(
+            lr=tune.grid_search([1e-3, 1e-4]),
+            beta=0.75,
+            # We must define a train batch size for each
+            # learner (here 1 local learner).
+            train_batch_size_per_learner=2000,
+        )
+        # Set the config's data path.
+        config.offline_data(
+            input_=[data_path.as_posix()],
+            # Set the number of updates to be run per learner
+            # per training step.
+            dataset_num_iters_per_learner=1,
+        )
+        # Set the config's environment for evalaution.
+        config.environment(
+            observation_space=gym.spaces.Box(
+                np.array([-4.8, -np.inf, -0.41887903, -np.inf]),
+                np.array([4.8, np.inf, 0.41887903, np.inf]),
+                shape=(4,),
+                dtype=np.float32,
+            ),
+            action_space=gym.spaces.Discrete(2),
+        )
+        # Set up a tuner to run the experiment.
+        tuner = tune.Tuner(
+            "MARWIL",
+            param_space=config,
+            run_config=tune.RunConfig(
+                stop={"training_iteration": 1},
+            ),
+        )
+        # Run the experiment.
+        tuner.fit()
+    """
+
+    def __init__(self, algo_class=None):
+        """Initializes a MARWILConfig 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 MARWIL)
+        self._is_online = False
+        # fmt: off
+        # __sphinx_doc_begin__
+        # MARWIL specific settings:
+        self.beta = 1.0
+        self.bc_logstd_coeff = 0.0
+        self.moving_average_sqd_adv_norm_update_rate = 1e-8
+        self.moving_average_sqd_adv_norm_start = 100.0
+        self.vf_coeff = 1.0
+        self.model["vf_share_layers"] = False
+        self.grad_clip = None
+
+        # Override some of AlgorithmConfig's default values with MARWIL-specific values.
+
+        # You should override input_ to point to an offline dataset
+        # (see algorithm.py and algorithm_config.py).
+        # The dataset may have an arbitrary number of timesteps
+        # (and even episodes) per line.
+        # However, each line must only contain consecutive timesteps in
+        # order for MARWIL to be able to calculate accumulated
+        # discounted returns. It is ok, though, to have multiple episodes in
+        # the same line.
+        self.input_ = "sampler"
+        self.postprocess_inputs = True
+        self.lr = 1e-4
+        self.lambda_ = 1.0
+        self.train_batch_size = 2000
+
+        # Materialize only the data in raw format, but not the mapped data b/c
+        # MARWIL uses a connector to calculate values and therefore the module
+        # needs to be updated frequently. This updating would not work if we
+        # map the data once at the beginning.
+        # TODO (simon, sven): The module is only updated when the OfflinePreLearner
+        #   gets reinitiated, i.e. when the iterator gets reinitiated. This happens
+        #   frequently enough with a small dataset, but with a big one this does not
+        #   update often enough. We might need to put model weigths every couple of
+        #   iterations into the object storage (maybe also connector states).
+        self.materialize_data = True
+        self.materialize_mapped_data = False
+        # __sphinx_doc_end__
+        # fmt: on
+        self._set_off_policy_estimation_methods = False
+
+    @override(AlgorithmConfig)
+    def training(
+        self,
+        *,
+        beta: Optional[float] = NotProvided,
+        bc_logstd_coeff: Optional[float] = NotProvided,
+        moving_average_sqd_adv_norm_update_rate: Optional[float] = NotProvided,
+        moving_average_sqd_adv_norm_start: Optional[float] = NotProvided,
+        vf_coeff: Optional[float] = NotProvided,
+        grad_clip: Optional[float] = NotProvided,
+        **kwargs,
+    ) -> Self:
+        """Sets the training related configuration.
+
+        Args:
+            beta: Scaling  of advantages in exponential terms. When beta is 0.0,
+                MARWIL is reduced to behavior cloning (imitation learning);
+                see bc.py algorithm in this same directory.
+            bc_logstd_coeff: A coefficient to encourage higher action distribution
+                entropy for exploration.
+            moving_average_sqd_adv_norm_update_rate: The rate for updating the
+                squared moving average advantage norm (c^2). A higher rate leads
+                to faster updates of this moving avergage.
+            moving_average_sqd_adv_norm_start: Starting value for the
+                squared moving average advantage norm (c^2).
+            vf_coeff: Balancing value estimation loss and policy optimization loss.
+            grad_clip: If specified, clip the global norm of gradients by this amount.
+
+        Returns:
+            This updated AlgorithmConfig object.
+        """
+        # Pass kwargs onto super's `training()` method.
+        super().training(**kwargs)
+        if beta is not NotProvided:
+            self.beta = beta
+        if bc_logstd_coeff is not NotProvided:
+            self.bc_logstd_coeff = bc_logstd_coeff
+        if moving_average_sqd_adv_norm_update_rate is not NotProvided:
+            self.moving_average_sqd_adv_norm_update_rate = (
+                moving_average_sqd_adv_norm_update_rate
+            )
+        if moving_average_sqd_adv_norm_start is not NotProvided:
+            self.moving_average_sqd_adv_norm_start = moving_average_sqd_adv_norm_start
+        if vf_coeff is not NotProvided:
+            self.vf_coeff = vf_coeff
+        if grad_clip is not NotProvided:
+            self.grad_clip = grad_clip
+        return self
+
+    @override(AlgorithmConfig)
+    def get_default_rl_module_spec(self) -> RLModuleSpecType:
+        if self.framework_str == "torch":
+            from ray.rllib.algorithms.ppo.torch.default_ppo_torch_rl_module import (
+                DefaultPPOTorchRLModule,
+            )
+
+            return RLModuleSpec(module_class=DefaultPPOTorchRLModule)
+        else:
+            raise ValueError(
+                f"The framework {self.framework_str} is not supported. "
+                "Use 'torch' instead."
+            )
+
+    @override(AlgorithmConfig)
+    def get_default_learner_class(self) -> Union[Type["Learner"], str]:
+        if self.framework_str == "torch":
+            from ray.rllib.algorithms.marwil.torch.marwil_torch_learner import (
+                MARWILTorchLearner,
+            )
+
+            return MARWILTorchLearner
+        else:
+            raise ValueError(
+                f"The framework {self.framework_str} is not supported. "
+                "Use 'torch' instead."
+            )
+
+    @override(AlgorithmConfig)
+    def evaluation(
+        self,
+        **kwargs,
+    ) -> Self:
+        """Sets the evaluation related configuration.
+        Returns:
+            This updated AlgorithmConfig object.
+        """
+        # Pass kwargs onto super's `evaluation()` method.
+        super().evaluation(**kwargs)
+
+        if "off_policy_estimation_methods" in kwargs:
+            # User specified their OPE methods.
+            self._set_off_policy_estimation_methods = True
+
+        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(AlgorithmConfig)
+    def build(
+        self,
+        env: Optional[Union[str, EnvType]] = None,
+        logger_creator: Optional[Callable[[], Logger]] = None,
+    ) -> "Algorithm":
+        if not self._set_off_policy_estimation_methods:
+            deprecation_warning(
+                old=r"MARWIL used to have off_policy_estimation_methods "
+                "is and wis by default. This has"
+                r"changed to off_policy_estimation_methods: \{\}."
+                "If you want to use an off-policy estimator, specify it in"
+                ".evaluation(off_policy_estimation_methods=...)",
+                error=False,
+            )
+        return super().build(env, logger_creator)
+
+    @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,
+        )
+
+        # Before anything, add one ts to each episode (and record this in the loss
+        # mask, so that the computations at this extra ts are not used to compute
+        # the loss).
+        pipeline.prepend(AddOneTsToEpisodesAndTruncate())
+
+        # 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(),
+        )
+
+        # At the end of the pipeline (when the batch is already completed), add the
+        # GAE connector, which performs a vf forward pass, then computes the GAE
+        # computations, and puts the results of this (advantages, value targets)
+        # directly back in the batch. This is then the batch used for
+        # `forward_train` and `compute_losses`.
+        pipeline.append(
+            GeneralAdvantageEstimation(gamma=self.gamma, lambda_=self.lambda_)
+        )
+
+        return pipeline
+
+    @override(AlgorithmConfig)
+    def validate(self) -> None:
+        # Call super's validation method.
+        super().validate()
+
+        if self.beta < 0.0 or self.beta > 1.0:
+            self._value_error("`beta` must be within 0.0 and 1.0!")
+
+        if self.postprocess_inputs is False and self.beta > 0.0:
+            self._value_error(
+                "`postprocess_inputs` must be True for MARWIL (to "
+                "calculate accum., discounted returns)! Try setting "
+                "`config.offline_data(postprocess_inputs=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 local Learner (`config.num_learners=0`), the number of "
+                "iterations per learner (`dataset_num_iters_per_learner`) has to be "
+                "defined! Set this hyperparameter through `config.offline_data("
+                "dataset_num_iters_per_learner=...)`."
+            )
+
+    @property
+    def _model_auto_keys(self):
+        return super()._model_auto_keys | {"beta": self.beta, "vf_share_layers": False}
+
+
+class MARWIL(Algorithm):
+    @classmethod
+    @override(Algorithm)
+    def get_default_config(cls) -> MARWILConfig:
+        return MARWILConfig()
+
+    @classmethod
+    @override(Algorithm)
+    def get_default_policy_class(
+        cls, config: AlgorithmConfig
+    ) -> Optional[Type[Policy]]:
+        if config["framework"] == "torch":
+            from ray.rllib.algorithms.marwil.marwil_torch_policy import (
+                MARWILTorchPolicy,
+            )
+
+            return MARWILTorchPolicy
+        elif config["framework"] == "tf":
+            from ray.rllib.algorithms.marwil.marwil_tf_policy import (
+                MARWILTF1Policy,
+            )
+
+            return MARWILTF1Policy
+        else:
+            from ray.rllib.algorithms.marwil.marwil_tf_policy import MARWILTF2Policy
+
+            return MARWILTF2Policy
+
+    @override(Algorithm)
+    def training_step(self) -> None:
+        """Implements training logic for the new stack
+
+        Note, this includes so far training with the `OfflineData`
+        class (multi-/single-learner setup) and evaluation on
+        `EnvRunner`s. Note further, evaluation on the dataset itself
+        using estimators is not implemented, yet.
+        """
+        # Old API stack (Policy, RolloutWorker, Connector).
+        if not self.config.enable_env_runner_and_connector_v2:
+            return self._training_step_old_api_stack()
+
+        # TODO (simon): Take care of sampler metrics: right
+        #  now all rewards are `nan`, which possibly confuses
+        #  the user that sth. is not right, although it is as
+        #  we do not step the env.
+        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 = (
+                self.config.num_learners > 0
+                or self.config.dataset_num_iters_per_learner != 1
+            )
+
+            # Sampling from offline data.
+            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,
+            )
+            if return_iterator:
+                training_data = TrainingData(data_iterators=batch_or_iterator)
+            else:
+                training_data = TrainingData(batch=batch_or_iterator)
+
+        with self.metrics.log_time((TIMERS, LEARNER_UPDATE_TIMER)):
+            # Updating the policy.
+            learner_results = self.learner_group.update(
+                training_data=training_data,
+                minibatch_size=self.config.train_batch_size_per_learner,
+                num_iters=self.config.dataset_num_iters_per_learner,
+                **self.offline_data.iter_batches_kwargs,
+            )
+
+            # Log training results.
+            self.metrics.aggregate(learner_results, key=LEARNER_RESULTS)
+
+    @OldAPIStack
+    def _training_step_old_api_stack(self) -> ResultDict:
+        """Implements training step for the old stack.
+
+        Note, there is no hybrid stack anymore. If you need to use `RLModule`s,
+        use the new api stack.
+        """
+        # 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(
+            module_id=list(self.config.policies)[0]
+        )
+        self._counters[NUM_AGENT_STEPS_SAMPLED] += train_batch.agent_steps()
+        self._counters[NUM_ENV_STEPS_SAMPLED] += train_batch.env_steps()
+
+        # Train.
+        if self.config.simple_optimizer:
+            train_results = train_one_step(self, train_batch)
+        else:
+            train_results = multi_gpu_train_one_step(self, train_batch)
+
+        # TODO: Move training steps counter update outside of `train_one_step()` method.
+        # # Update train step counters.
+        # self._counters[NUM_ENV_STEPS_TRAINED] += train_batch.env_steps()
+        # self._counters[NUM_AGENT_STEPS_TRAINED] += train_batch.agent_steps()
+
+        global_vars = {
+            "timestep": self._counters[NUM_AGENT_STEPS_SAMPLED],
+        }
+
+        # Update weights - after learning on the local worker - on all remote
+        # workers (only those policies that were actually trained).
+        if self.env_runner_group.num_remote_env_runners() > 0:
+            with self._timers[SYNCH_WORKER_WEIGHTS_TIMER]:
+                self.env_runner_group.sync_weights(
+                    policies=list(train_results.keys()), global_vars=global_vars
+                )
+
+        # Update global vars on local worker as well.
+        self.env_runner.set_global_vars(global_vars)
+
+        return train_results

lib/python3.12/site-packages/ray/rllib/algorithms/marwil/marwil_learner.py

@@ -0,0 +1,51 @@
+from typing import Dict, Optional
+
+from ray.rllib.core.learner.learner import Learner
+from ray.rllib.core.rl_module.apis import ValueFunctionAPI
+from ray.rllib.utils.annotations import override
+from ray.rllib.utils.lambda_defaultdict import LambdaDefaultDict
+from ray.rllib.utils.typing import ModuleID, ShouldModuleBeUpdatedFn, TensorType
+
+LEARNER_RESULTS_MOVING_AVG_SQD_ADV_NORM_KEY = "moving_avg_sqd_adv_norm"
+LEARNER_RESULTS_VF_EXPLAINED_VAR_KEY = "vf_explained_variance"
+
+
+# TODO (simon): Check, if the norm update should be done inside
+# the Learner.
+class MARWILLearner(Learner):
+    @override(Learner)
+    def build(self) -> None:
+        super().build()
+
+        # Dict mapping module IDs to the respective moving averages of squared
+        # advantages.
+        self.moving_avg_sqd_adv_norms_per_module: Dict[
+            ModuleID, TensorType
+        ] = LambdaDefaultDict(
+            lambda module_id: self._get_tensor_variable(
+                self.config.get_config_for_module(
+                    module_id
+                ).moving_average_sqd_adv_norm_start
+            )
+        )
+
+    @override(Learner)
+    def remove_module(
+        self,
+        module_id: ModuleID,
+        *,
+        new_should_module_be_updated: Optional[ShouldModuleBeUpdatedFn] = None,
+    ) -> None:
+        super().remove_module(
+            module_id,
+            new_should_module_be_updated=new_should_module_be_updated,
+        )
+        # In case of BC (beta==0.0 and this property never being used),
+        self.moving_avg_sqd_adv_norms_per_module.pop(module_id, None)
+
+    @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]

lib/python3.12/site-packages/ray/rllib/algorithms/marwil/marwil_tf_policy.py

@@ -0,0 +1,251 @@
+import logging
+from typing import Any, Dict, List, Optional, Type, Union
+
+from ray.rllib.evaluation.postprocessing import Postprocessing, compute_advantages
+from ray.rllib.models.action_dist import ActionDistribution
+from ray.rllib.models.modelv2 import ModelV2
+from ray.rllib.models.tf.tf_action_dist import 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.policy import Policy
+from ray.rllib.policy.sample_batch import SampleBatch
+from ray.rllib.policy.tf_mixins import (
+    ValueNetworkMixin,
+    compute_gradients,
+)
+from ray.rllib.utils.annotations import override
+from ray.rllib.utils.framework import get_variable, try_import_tf
+from ray.rllib.utils.tf_utils import explained_variance
+from ray.rllib.utils.typing import (
+    LocalOptimizer,
+    ModelGradients,
+    TensorType,
+)
+
+tf1, tf, tfv = try_import_tf()
+
+logger = logging.getLogger(__name__)
+
+
+class PostprocessAdvantages:
+    """Marwil's custom trajectory post-processing mixin."""
+
+    def __init__(self):
+        pass
+
+    def postprocess_trajectory(
+        self,
+        sample_batch: SampleBatch,
+        other_agent_batches: Optional[Dict[Any, SampleBatch]] = None,
+        episode=None,
+    ):
+        sample_batch = super().postprocess_trajectory(
+            sample_batch, other_agent_batches, episode
+        )
+
+        # Trajectory is actually complete -> last r=0.0.
+        if sample_batch[SampleBatch.TERMINATEDS][-1]:
+            last_r = 0.0
+        # Trajectory has been truncated -> last r=VF estimate of last obs.
+        else:
+            # Input dict is provided to us automatically via the Model's
+            # requirements. It's a single-timestep (last one in trajectory)
+            # input_dict.
+            # Create an input dict according to the Model's requirements.
+            index = "last" if SampleBatch.NEXT_OBS in sample_batch else -1
+            input_dict = sample_batch.get_single_step_input_dict(
+                self.view_requirements, index=index
+            )
+            last_r = self._value(**input_dict)
+
+        # Adds the "advantages" (which in the case of MARWIL are simply the
+        # discounted cumulative rewards) to the SampleBatch.
+        return compute_advantages(
+            sample_batch,
+            last_r,
+            self.config["gamma"],
+            # We just want the discounted cumulative rewards, so we won't need
+            # GAE nor critic (use_critic=True: Subtract vf-estimates from returns).
+            use_gae=False,
+            use_critic=False,
+        )
+
+
+class MARWILLoss:
+    def __init__(
+        self,
+        policy: Policy,
+        value_estimates: TensorType,
+        action_dist: ActionDistribution,
+        train_batch: SampleBatch,
+        vf_loss_coeff: float,
+        beta: float,
+    ):
+        # L = - A * log\pi_\theta(a|s)
+        logprobs = action_dist.logp(train_batch[SampleBatch.ACTIONS])
+        if beta != 0.0:
+            cumulative_rewards = train_batch[Postprocessing.ADVANTAGES]
+            # Advantage Estimation.
+            adv = cumulative_rewards - value_estimates
+            adv_squared = tf.reduce_mean(tf.math.square(adv))
+            # Value function's loss term (MSE).
+            self.v_loss = 0.5 * adv_squared
+
+            # Perform moving averaging of advantage^2.
+            rate = policy.config["moving_average_sqd_adv_norm_update_rate"]
+
+            # Update averaged advantage norm.
+            # Eager.
+            if policy.config["framework"] == "tf2":
+                update_term = adv_squared - policy._moving_average_sqd_adv_norm
+                policy._moving_average_sqd_adv_norm.assign_add(rate * update_term)
+
+                # Exponentially weighted advantages.
+                c = tf.math.sqrt(policy._moving_average_sqd_adv_norm)
+                exp_advs = tf.math.exp(beta * (adv / (1e-8 + c)))
+            # Static graph.
+            else:
+                update_adv_norm = tf1.assign_add(
+                    ref=policy._moving_average_sqd_adv_norm,
+                    value=rate * (adv_squared - policy._moving_average_sqd_adv_norm),
+                )
+
+                # Exponentially weighted advantages.
+                with tf1.control_dependencies([update_adv_norm]):
+                    exp_advs = tf.math.exp(
+                        beta
+                        * tf.math.divide(
+                            adv,
+                            1e-8 + tf.math.sqrt(policy._moving_average_sqd_adv_norm),
+                        )
+                    )
+            exp_advs = tf.stop_gradient(exp_advs)
+
+            self.explained_variance = tf.reduce_mean(
+                explained_variance(cumulative_rewards, value_estimates)
+            )
+
+        else:
+            # Value function's loss term (MSE).
+            self.v_loss = tf.constant(0.0)
+            exp_advs = 1.0
+
+        # logprob loss alone tends to push action distributions to
+        # have very low entropy, resulting in worse performance for
+        # unfamiliar situations.
+        # A scaled logstd loss term encourages stochasticity, thus
+        # alleviate the problem to some extent.
+        logstd_coeff = policy.config["bc_logstd_coeff"]
+        if logstd_coeff > 0.0:
+            logstds = tf.reduce_sum(action_dist.log_std, axis=1)
+        else:
+            logstds = 0.0
+
+        self.p_loss = -1.0 * tf.reduce_mean(
+            exp_advs * (logprobs + logstd_coeff * logstds)
+        )
+
+        self.total_loss = self.p_loss + vf_loss_coeff * self.v_loss
+
+
+# We need this builder function because we want to share the same
+# custom logics between TF1 dynamic and TF2 eager policies.
+def get_marwil_tf_policy(name: str, base: type) -> type:
+    """Construct a MARWILTFPolicy 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 MAML.
+    """
+
+    class MARWILTFPolicy(ValueNetworkMixin, PostprocessAdvantages, 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)
+            PostprocessAdvantages.__init__(self)
+
+            # Not needed for pure BC.
+            if config["beta"] != 0.0:
+                # Set up a tf-var for the moving avg (do this here to make it work
+                # with eager mode); "c^2" in the paper.
+                self._moving_average_sqd_adv_norm = get_variable(
+                    config["moving_average_sqd_adv_norm_start"],
+                    framework="tf",
+                    tf_name="moving_average_of_advantage_norm",
+                    trainable=False,
+                )
+
+            # 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)
+            value_estimates = model.value_function()
+
+            self._marwil_loss = MARWILLoss(
+                self,
+                value_estimates,
+                action_dist,
+                train_batch,
+                self.config["vf_coeff"],
+                self.config["beta"],
+            )
+
+            return self._marwil_loss.total_loss
+
+        @override(base)
+        def stats_fn(self, train_batch: SampleBatch) -> Dict[str, TensorType]:
+            stats = {
+                "policy_loss": self._marwil_loss.p_loss,
+                "total_loss": self._marwil_loss.total_loss,
+            }
+            if self.config["beta"] != 0.0:
+                stats["moving_average_sqd_adv_norm"] = self._moving_average_sqd_adv_norm
+                stats["vf_explained_var"] = self._marwil_loss.explained_variance
+                stats["vf_loss"] = self._marwil_loss.v_loss
+
+            return stats
+
+        @override(base)
+        def compute_gradients_fn(
+            self, optimizer: LocalOptimizer, loss: TensorType
+        ) -> ModelGradients:
+            return compute_gradients(self, optimizer, loss)
+
+    MARWILTFPolicy.__name__ = name
+    MARWILTFPolicy.__qualname__ = name
+
+    return MARWILTFPolicy
+
+
+MARWILTF1Policy = get_marwil_tf_policy("MARWILTF1Policy", DynamicTFPolicyV2)
+MARWILTF2Policy = get_marwil_tf_policy("MARWILTF2Policy", EagerTFPolicyV2)

lib/python3.12/site-packages/ray/rllib/algorithms/marwil/marwil_torch_policy.py

@@ -0,0 +1,132 @@
+from typing import Dict, List, Type, Union
+
+from ray.rllib.algorithms.marwil.marwil_tf_policy import PostprocessAdvantages
+from ray.rllib.evaluation.postprocessing import Postprocessing
+from ray.rllib.models.modelv2 import ModelV2
+from ray.rllib.models.torch.torch_action_dist import TorchDistributionWrapper
+from ray.rllib.policy.sample_batch import SampleBatch
+from ray.rllib.policy.torch_mixins import 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
+from ray.rllib.utils.typing import TensorType
+
+torch, _ = try_import_torch()
+
+
+class MARWILTorchPolicy(ValueNetworkMixin, PostprocessAdvantages, TorchPolicyV2):
+    """PyTorch policy class used with Marwil."""
+
+    def __init__(self, observation_space, action_space, config):
+        TorchPolicyV2.__init__(
+            self,
+            observation_space,
+            action_space,
+            config,
+            max_seq_len=config["model"]["max_seq_len"],
+        )
+
+        ValueNetworkMixin.__init__(self, config)
+        PostprocessAdvantages.__init__(self)
+
+        # Not needed for pure BC.
+        if config["beta"] != 0.0:
+            # Set up a torch-var for the squared moving avg. advantage norm.
+            self._moving_average_sqd_adv_norm = torch.tensor(
+                [config["moving_average_sqd_adv_norm_start"]],
+                dtype=torch.float32,
+                requires_grad=False,
+            ).to(self.device)
+
+        # TODO: Don't require users to call this manually.
+        self._initialize_loss_from_dummy_batch()
+
+    @override(TorchPolicyV2)
+    def loss(
+        self,
+        model: ModelV2,
+        dist_class: Type[TorchDistributionWrapper],
+        train_batch: SampleBatch,
+    ) -> Union[TensorType, List[TensorType]]:
+        model_out, _ = model(train_batch)
+        action_dist = dist_class(model_out, model)
+        actions = train_batch[SampleBatch.ACTIONS]
+        # log\pi_\theta(a|s)
+        logprobs = action_dist.logp(actions)
+
+        # Advantage estimation.
+        if self.config["beta"] != 0.0:
+            cumulative_rewards = train_batch[Postprocessing.ADVANTAGES]
+            state_values = model.value_function()
+            adv = cumulative_rewards - state_values
+            adv_squared_mean = torch.mean(torch.pow(adv, 2.0))
+
+            explained_var = explained_variance(cumulative_rewards, state_values)
+            ev = torch.mean(explained_var)
+            model.tower_stats["explained_variance"] = ev
+
+            # Policy loss.
+            # Update averaged advantage norm.
+            rate = self.config["moving_average_sqd_adv_norm_update_rate"]
+            self._moving_average_sqd_adv_norm = (
+                rate * (adv_squared_mean.detach() - self._moving_average_sqd_adv_norm)
+                + self._moving_average_sqd_adv_norm
+            )
+            model.tower_stats[
+                "_moving_average_sqd_adv_norm"
+            ] = self._moving_average_sqd_adv_norm
+            # Exponentially weighted advantages.
+            exp_advs = torch.exp(
+                self.config["beta"]
+                * (adv / (1e-8 + torch.pow(self._moving_average_sqd_adv_norm, 0.5)))
+            ).detach()
+            # Value loss.
+            v_loss = 0.5 * adv_squared_mean
+        else:
+            # Policy loss (simple BC loss term).
+            exp_advs = 1.0
+            # Value loss.
+            v_loss = 0.0
+        model.tower_stats["v_loss"] = v_loss
+        # logprob loss alone tends to push action distributions to
+        # have very low entropy, resulting in worse performance for
+        # unfamiliar situations.
+        # A scaled logstd loss term encourages stochasticity, thus
+        # alleviate the problem to some extent.
+        logstd_coeff = self.config["bc_logstd_coeff"]
+        if logstd_coeff > 0.0:
+            logstds = torch.mean(action_dist.log_std, dim=1)
+        else:
+            logstds = 0.0
+
+        p_loss = -torch.mean(exp_advs * (logprobs + logstd_coeff * logstds))
+        model.tower_stats["p_loss"] = p_loss
+        # Combine both losses.
+        self.v_loss = v_loss
+        self.p_loss = p_loss
+        total_loss = p_loss + self.config["vf_coeff"] * v_loss
+        model.tower_stats["total_loss"] = total_loss
+        return total_loss
+
+    @override(TorchPolicyV2)
+    def stats_fn(self, train_batch: SampleBatch) -> Dict[str, TensorType]:
+        stats = {
+            "policy_loss": self.get_tower_stats("p_loss")[0].item(),
+            "total_loss": self.get_tower_stats("total_loss")[0].item(),
+        }
+        if self.config["beta"] != 0.0:
+            stats["moving_average_sqd_adv_norm"] = self.get_tower_stats(
+                "_moving_average_sqd_adv_norm"
+            )[0].item()
+            stats["vf_explained_var"] = self.get_tower_stats("explained_variance")[
+                0
+            ].item()
+            stats["vf_loss"] = self.get_tower_stats("v_loss")[0].item()
+        return convert_to_numpy(stats)
+
+    def extra_grad_process(
+        self, optimizer: "torch.optim.Optimizer", loss: TensorType
+    ) -> Dict[str, TensorType]:
+        return apply_grad_clipping(self, optimizer, loss)

lib/python3.12/site-packages/ray/rllib/algorithms/marwil/torch/marwil_torch_learner.py

@@ -0,0 +1,138 @@
+from typing import Any, Dict, Optional
+
+from ray.rllib.algorithms.marwil.marwil import MARWILConfig
+from ray.rllib.algorithms.marwil.marwil_learner import (
+    LEARNER_RESULTS_MOVING_AVG_SQD_ADV_NORM_KEY,
+    LEARNER_RESULTS_VF_EXPLAINED_VAR_KEY,
+    MARWILLearner,
+)
+from ray.rllib.core.columns import Columns
+from ray.rllib.core.learner.learner import POLICY_LOSS_KEY, VF_LOSS_KEY
+from ray.rllib.core.learner.torch.torch_learner import TorchLearner
+from ray.rllib.evaluation.postprocessing import Postprocessing
+from ray.rllib.utils.framework import try_import_torch
+from ray.rllib.utils.torch_utils import explained_variance
+from ray.rllib.utils.typing import TensorType
+
+torch, nn = try_import_torch()
+
+
+class MARWILTorchLearner(MARWILLearner, TorchLearner):
+    """Implements torch-specific MARWIL loss on top of MARWILLearner.
+
+    This class implements the MARWIL loss under `self.compute_loss_for_module()`.
+    """
+
+    def compute_loss_for_module(
+        self,
+        *,
+        module_id: str,
+        config: Optional[MARWILConfig] = None,
+        batch: Dict[str, Any],
+        fwd_out: Dict[str, TensorType]
+    ) -> TensorType:
+        module = self.module[module_id].unwrapped()
+
+        # 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 additional (artificial) timestep to each episode to
+        # simplify the actual computation.
+        if Columns.LOSS_MASK in batch:
+            num_valid = torch.sum(batch[Columns.LOSS_MASK])
+
+            def possibly_masked_mean(data_):
+                return torch.sum(data_[batch[Columns.LOSS_MASK]]) / num_valid
+
+        else:
+            possibly_masked_mean = torch.mean
+
+        action_dist_class_train = module.get_train_action_dist_cls()
+        curr_action_dist = action_dist_class_train.from_logits(
+            fwd_out[Columns.ACTION_DIST_INPUTS]
+        )
+
+        log_probs = curr_action_dist.logp(batch[Columns.ACTIONS])
+
+        # If beta is zero, we fall back to BC.
+        if config.beta == 0.0:
+            # Value function's loss term.
+            mean_vf_loss = 0.0
+            # Policy's loss term.
+            exp_weighted_advantages = 1.0
+        # Otherwise, compute advantages.
+        else:
+            # cumulative_rewards = batch[Columns.ADVANTAGES]
+            value_fn_out = module.compute_values(
+                batch, embeddings=fwd_out.get(Columns.EMBEDDINGS)
+            )
+            advantages = batch[Columns.VALUE_TARGETS] - value_fn_out
+            advantages_squared_mean = possibly_masked_mean(torch.pow(advantages, 2.0))
+
+            # Compute the value loss.
+            mean_vf_loss = 0.5 * advantages_squared_mean
+
+            # Compute the policy loss.
+            self.moving_avg_sqd_adv_norms_per_module[module_id] = (
+                config.moving_average_sqd_adv_norm_update_rate
+                * (
+                    advantages_squared_mean.detach()
+                    - self.moving_avg_sqd_adv_norms_per_module[module_id]
+                )
+                + self.moving_avg_sqd_adv_norms_per_module[module_id]
+            )
+            # Exponentially weighted advantages.
+            # TODO (simon): Check, if we need the mask here.
+            exp_weighted_advantages = torch.exp(
+                config.beta
+                * (
+                    advantages
+                    / (
+                        1e-8
+                        + torch.pow(
+                            self.moving_avg_sqd_adv_norms_per_module[module_id], 0.5
+                        )
+                    )
+                )
+            ).detach()
+
+        # Note, using solely a log-probability loss term tends to push the action
+        # distributions to have very low entropy, which results in worse performance
+        # specifically in unknown situations.
+        # Scaling the loss term with the logarithm of the action distribution's
+        # standard deviation encourages stochasticity in the policy.
+        if config.bc_logstd_coeff > 0.0:
+            log_stds = possibly_masked_mean(curr_action_dist.log_std, dim=1)
+        else:
+            log_stds = 0.0
+
+        # Compute the policy loss.
+        policy_loss = -possibly_masked_mean(
+            exp_weighted_advantages * (log_probs + config.bc_logstd_coeff * log_stds)
+        )
+
+        # Compute the total loss.
+        total_loss = policy_loss + config.vf_coeff * mean_vf_loss
+
+        # Log import loss stats. In case of the BC loss this is simply
+        # the policy loss.
+        if config.beta == 0.0:
+            self.metrics.log_value((module_id, POLICY_LOSS_KEY), policy_loss, window=1)
+        # Log more stats, if using the MARWIL loss.
+        else:
+            ma_sqd_adv_norms = self.moving_avg_sqd_adv_norms_per_module[module_id]
+            self.metrics.log_dict(
+                {
+                    POLICY_LOSS_KEY: policy_loss,
+                    VF_LOSS_KEY: mean_vf_loss,
+                    LEARNER_RESULTS_VF_EXPLAINED_VAR_KEY: explained_variance(
+                        batch[Postprocessing.VALUE_TARGETS], value_fn_out
+                    ),
+                    LEARNER_RESULTS_MOVING_AVG_SQD_ADV_NORM_KEY: ma_sqd_adv_norms,
+                },
+                key=module_id,
+                window=1,  # <- single items (should not be mean/ema-reduced over time).
+            )
+
+        # Return the total loss.
+        return total_loss

lib/python3.12/site-packages/ray/rllib/algorithms/mock.py

@@ -0,0 +1,154 @@
+import os
+import pickle
+import time
+
+import numpy as np
+
+from ray.rllib.algorithms.algorithm import Algorithm, AlgorithmConfig
+from ray.rllib.utils.annotations import override
+from ray.tune import result as tune_result
+
+
+class _MockTrainer(Algorithm):
+    """Mock Algorithm for use in tests."""
+
+    @classmethod
+    @override(Algorithm)
+    def get_default_config(cls) -> AlgorithmConfig:
+        return (
+            AlgorithmConfig()
+            .framework("tf")
+            .update_from_dict(
+                {
+                    "mock_error": False,
+                    "persistent_error": False,
+                    "test_variable": 1,
+                    "user_checkpoint_freq": 0,
+                    "sleep": 0,
+                }
+            )
+        )
+
+    @classmethod
+    def default_resource_request(cls, config: AlgorithmConfig):
+        return None
+
+    @override(Algorithm)
+    def setup(self, config):
+        self.callbacks = self.config.callbacks_class()
+
+        # Add needed properties.
+        self.info = None
+        self.restored = False
+
+    @override(Algorithm)
+    def step(self):
+        if (
+            self.config.mock_error
+            and self.iteration == 1
+            and (self.config.persistent_error or not self.restored)
+        ):
+            raise Exception("mock error")
+        if self.config.sleep:
+            time.sleep(self.config.sleep)
+        result = dict(
+            episode_reward_mean=10, episode_len_mean=10, timesteps_this_iter=10, info={}
+        )
+        if self.config.user_checkpoint_freq > 0 and self.iteration > 0:
+            if self.iteration % self.config.user_checkpoint_freq == 0:
+                result.update({tune_result.SHOULD_CHECKPOINT: True})
+        return result
+
+    @override(Algorithm)
+    def save_checkpoint(self, checkpoint_dir):
+        path = os.path.join(checkpoint_dir, "mock_agent.pkl")
+        with open(path, "wb") as f:
+            pickle.dump(self.info, f)
+
+    @override(Algorithm)
+    def load_checkpoint(self, checkpoint_dir):
+        path = os.path.join(checkpoint_dir, "mock_agent.pkl")
+        with open(path, "rb") as f:
+            info = pickle.load(f)
+        self.info = info
+        self.restored = True
+
+    @staticmethod
+    @override(Algorithm)
+    def _get_env_id_and_creator(env_specifier, config):
+        # No env to register.
+        return None, None
+
+    def set_info(self, info):
+        self.info = info
+        return info
+
+    def get_info(self, sess=None):
+        return self.info
+
+
+class _SigmoidFakeData(_MockTrainer):
+    """Algorithm that returns sigmoid learning curves.
+
+    This can be helpful for evaluating early stopping algorithms."""
+
+    @classmethod
+    @override(Algorithm)
+    def get_default_config(cls) -> AlgorithmConfig:
+        return AlgorithmConfig().update_from_dict(
+            {
+                "width": 100,
+                "height": 100,
+                "offset": 0,
+                "iter_time": 10,
+                "iter_timesteps": 1,
+            }
+        )
+
+    def step(self):
+        i = max(0, self.iteration - self.config.offset)
+        v = np.tanh(float(i) / self.config.width)
+        v *= self.config.height
+        return dict(
+            episode_reward_mean=v,
+            episode_len_mean=v,
+            timesteps_this_iter=self.config.iter_timesteps,
+            time_this_iter_s=self.config.iter_time,
+            info={},
+        )
+
+
+class _ParameterTuningTrainer(_MockTrainer):
+    @classmethod
+    @override(Algorithm)
+    def get_default_config(cls) -> AlgorithmConfig:
+        return AlgorithmConfig().update_from_dict(
+            {
+                "reward_amt": 10,
+                "dummy_param": 10,
+                "dummy_param2": 15,
+                "iter_time": 10,
+                "iter_timesteps": 1,
+            }
+        )
+
+    def step(self):
+        return dict(
+            episode_reward_mean=self.config.reward_amt * self.iteration,
+            episode_len_mean=self.config.reward_amt,
+            timesteps_this_iter=self.config.iter_timesteps,
+            time_this_iter_s=self.config.iter_time,
+            info={},
+        )
+
+
+def _algorithm_import_failed(trace):
+    """Returns dummy Algorithm class for if PyTorch etc. is not installed."""
+
+    class _AlgorithmImportFailed(Algorithm):
+        _name = "AlgorithmImportFailed"
+
+        def setup(self, config):
+            raise ImportError(trace)
+
+    return _AlgorithmImportFailed

lib/python3.12/site-packages/ray/rllib/algorithms/ppo/__init__.py

@@ -0,0 +1,12 @@
+from ray.rllib.algorithms.ppo.ppo import PPO, PPOConfig
+from ray.rllib.algorithms.ppo.ppo_tf_policy import PPOTF1Policy, PPOTF2Policy
+from ray.rllib.algorithms.ppo.ppo_torch_policy import PPOTorchPolicy
+
+__all__ = [
+    "PPO",
+    "PPOConfig",
+    # @OldAPIStack
+    "PPOTF1Policy",
+    "PPOTF2Policy",
+    "PPOTorchPolicy",
+]