chore: import upstream snapshot with attribution

rllib/algorithms/ppo/README.md

@@ -0,0 +1,49 @@
+# Proximal Policy Optimization (PPO)
+
+## Overview
+
+[PPO](https://arxiv.org/abs/1707.06347) is a model-free on-policy RL algorithm that works
+well for both discrete and continuous action space environments. PPO utilizes an
+actor-critic framework, where there are two networks, an actor (policy network) and
+critic network (value function).
+
+There are two formulations of PPO, which are both implemented in RLlib. The first
+formulation of PPO imitates the prior paper [TRPO](https://arxiv.org/abs/1502.05477)
+without the complexity of second-order optimization. In this formulation, for every
+iteration, an old version of an actor-network is saved and the agent seeks to optimize
+the RL objective while staying close to the old policy. This makes sure that the agent
+does not destabilize during training. In the second formulation, To mitigate destructive
+large policy updates, an issue discovered for vanilla policy gradient methods, PPO
+introduces the surrogate objective, which clips large action probability ratios between
+the current and old policy. Clipping has been shown in the paper to significantly
+improve training stability and speed.
+
+## Distributed PPO Algorithms
+
+PPO is a core algorithm in RLlib due to its ability to scale well with the number of nodes.
+
+In RLlib, we provide various implementations of distributed PPO, with different underlying
+execution plans, as shown below:
+
+### Distributed baseline PPO ..
+.. is a synchronous distributed RL algorithm (this algo here).
+Data collection nodes, which represent the old policy, gather data synchronously to
+create a large pool of on-policy data from which the agent performs minibatch
+gradient descent on.
+
+### Asychronous PPO (APPO)
+
+[See implementation here](https://github.com/ray-project/ray/blob/master/rllib/algorithms/appo/appo.py)
+
+### Decentralized Distributed PPO (DDPPO)
+
+[See implementation here](https://github.com/ray-project/ray/blob/master/rllib/algorithms/ddppo/ddppo.py)
+
+
+## Documentation & Implementation:
+
+### Proximal Policy Optimization (PPO).
+
+**[Detailed Documentation](https://docs.ray.io/en/master/rllib-algorithms.html#ppo)**
+
+**[Implementation](https://github.com/ray-project/ray/blob/master/rllib/algorithms/ppo/ppo.py)**

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",
+]

rllib/algorithms/ppo/default_ppo_rl_module.py

@@ -0,0 +1,62 @@
+import abc
+from typing import List
+
+from ray.rllib.core.models.configs import RecurrentEncoderConfig
+from ray.rllib.core.rl_module.apis import InferenceOnlyAPI, ValueFunctionAPI
+from ray.rllib.core.rl_module.rl_module import RLModule
+from ray.rllib.utils.annotations import (
+    OverrideToImplementCustomLogic_CallToSuperRecommended,
+    override,
+)
+from ray.util.annotations import DeveloperAPI
+
+
+@DeveloperAPI
+class DefaultPPORLModule(RLModule, InferenceOnlyAPI, ValueFunctionAPI, abc.ABC):
+    """Default RLModule used by PPO, if user does not specify a custom RLModule.
+
+    Users who want to train their RLModules with PPO 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)
+    """
+
+    @override(RLModule)
+    def setup(self):
+        # __sphinx_doc_begin__
+        # If we have a stateful model, states for the critic need to be collected
+        # during sampling and `inference-only` needs to be `False`. Note, at this
+        # point the encoder is not built, yet and therefore `is_stateful()` does
+        # not work.
+        is_stateful = isinstance(
+            self.catalog.actor_critic_encoder_config.base_encoder_config,
+            RecurrentEncoderConfig,
+        )
+        if is_stateful:
+            self.inference_only = False
+        # If this is an `inference_only` Module, we'll have to pass this information
+        # to the encoder config as well.
+        if self.inference_only and self.framework == "torch":
+            self.catalog.actor_critic_encoder_config.inference_only = True
+
+        # Build models from catalog.
+        self.encoder = self.catalog.build_actor_critic_encoder(framework=self.framework)
+        self.pi = self.catalog.build_pi_head(framework=self.framework)
+        self.vf = self.catalog.build_vf_head(framework=self.framework)
+        # __sphinx_doc_end__
+
+    @override(RLModule)
+    def get_initial_state(self) -> dict:
+        if hasattr(self.encoder, "get_initial_state"):
+            return self.encoder.get_initial_state()
+        else:
+            return {}
+
+    @OverrideToImplementCustomLogic_CallToSuperRecommended
+    @override(InferenceOnlyAPI)
+    def get_non_inference_attributes(self) -> List[str]:
+        """Return attributes, which are NOT inference-only (only used for training)."""
+        return ["vf"] + (
+            []
+            if self.model_config.get("vf_share_layers")
+            else ["encoder.critic_encoder"]
+        )

rllib/algorithms/ppo/ppo.py

@@ -0,0 +1,570 @@
+"""
+Proximal Policy Optimization (PPO)
+==================================
+
+This file defines the distributed Algorithm class for proximal policy
+optimization.
+See `ppo_[tf|torch]_policy.py` for the definition of the policy loss.
+
+Detailed documentation: https://docs.ray.io/en/master/rllib-algorithms.html#ppo
+"""
+
+import logging
+from typing import TYPE_CHECKING, Any, Dict, List, Optional, Type, Union
+
+from typing_extensions import Self
+
+from ray._common.deprecation import DEPRECATED_VALUE
+from ray.rllib.algorithms.algorithm import Algorithm
+from ray.rllib.algorithms.algorithm_config import AlgorithmConfig, NotProvided
+from ray.rllib.core.rl_module.rl_module import RLModuleSpec
+from ray.rllib.execution.rollout_ops import (
+    standardize_fields,
+    synchronous_parallel_sample,
+)
+from ray.rllib.execution.train_ops import (
+    multi_gpu_train_one_step,
+    train_one_step,
+)
+from ray.rllib.policy.policy import Policy
+from ray.rllib.utils.annotations import OldAPIStack, override
+from ray.rllib.utils.metrics import (
+    ALL_MODULES,
+    ENV_RUNNER_RESULTS,
+    ENV_RUNNER_SAMPLING_TIMER,
+    LEARNER_RESULTS,
+    LEARNER_UPDATE_TIMER,
+    NUM_AGENT_STEPS_SAMPLED,
+    NUM_ENV_STEPS_SAMPLED,
+    NUM_ENV_STEPS_SAMPLED_LIFETIME,
+    NUM_MODULE_STEPS_TRAINED_LIFETIME,
+    SAMPLE_TIMER,
+    SYNCH_WORKER_WEIGHTS_TIMER,
+    TIMERS,
+)
+from ray.rllib.utils.metrics.learner_info import LEARNER_STATS_KEY
+from ray.rllib.utils.schedules.scheduler import Scheduler
+from ray.rllib.utils.typing import ResultDict
+from ray.util.debug import log_once
+
+if TYPE_CHECKING:
+    from ray.rllib.core.learner.learner import Learner
+
+
+logger = logging.getLogger(__name__)
+
+LEARNER_RESULTS_VF_LOSS_UNCLIPPED_KEY = "vf_loss_unclipped"
+LEARNER_RESULTS_VF_EXPLAINED_VAR_KEY = "vf_explained_var"
+LEARNER_RESULTS_KL_KEY = "mean_kl_loss"
+LEARNER_RESULTS_CURR_KL_COEFF_KEY = "curr_kl_coeff"
+LEARNER_RESULTS_CURR_ENTROPY_COEFF_KEY = "curr_entropy_coeff"
+
+
+class PPOConfig(AlgorithmConfig):
+    """Defines a configuration class from which a PPO Algorithm can be built.
+
+    .. testcode::
+
+        from ray.rllib.algorithms.ppo import PPOConfig
+
+        config = PPOConfig()
+        config.environment("CartPole-v1")
+        config.env_runners(num_env_runners=1)
+        config.training(
+            gamma=0.9, lr=0.01, kl_coeff=0.3, train_batch_size_per_learner=256
+        )
+
+        # Build a Algorithm object from the config and run 1 training iteration.
+        algo = config.build()
+        algo.train()
+
+    .. testcode::
+
+        from ray.rllib.algorithms.ppo import PPOConfig
+        from ray import tune
+
+        config = (
+            PPOConfig()
+            # Set the config object's env.
+            .environment(env="CartPole-v1")
+            # Update the config object's training parameters.
+            .training(
+                lr=0.001, clip_param=0.2
+            )
+        )
+
+        tune.Tuner(
+            "PPO",
+            run_config=tune.RunConfig(stop={"training_iteration": 1}),
+            param_space=config,
+        ).fit()
+
+    .. testoutput::
+        :hide:
+
+        ...
+    """
+
+    def __init__(self, algo_class=None):
+        """Initializes a PPOConfig 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 PPO)
+
+        # fmt: off
+        # __sphinx_doc_begin__
+        self.lr = 5e-5
+        self.rollout_fragment_length = "auto"
+        self.train_batch_size = 4000
+
+        # PPO specific settings:
+        self.use_critic = True
+        self.use_gae = True
+        self.num_epochs = 30
+        self.minibatch_size = 128
+        self.shuffle_batch_per_epoch = True
+        self.lambda_ = 1.0
+        self.use_kl_loss = True
+        self.kl_coeff = 0.2
+        self.kl_target = 0.01
+        self.vf_loss_coeff = 1.0
+        self.entropy_coeff = 0.0
+        self.clip_param = 0.3
+        self.vf_clip_param = 10.0
+        self.grad_clip = None
+
+        # Override some of AlgorithmConfig's default values with PPO-specific values.
+        self.num_env_runners = 2
+        # __sphinx_doc_end__
+        # fmt: on
+
+        self.model["vf_share_layers"] = False  # @OldAPIStack
+        self.entropy_coeff_schedule = None  # @OldAPIStack
+        self.lr_schedule = None  # @OldAPIStack
+
+        # Deprecated keys.
+        self.sgd_minibatch_size = DEPRECATED_VALUE
+        self.vf_share_layers = DEPRECATED_VALUE
+
+    @override(AlgorithmConfig)
+    def get_default_rl_module_spec(self) -> RLModuleSpec:
+        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 either 'torch' or 'tf2'."
+            )
+
+    @override(AlgorithmConfig)
+    def get_default_learner_class(self) -> Union[Type["Learner"], str]:
+        if self.framework_str == "torch":
+            from ray.rllib.algorithms.ppo.torch.ppo_torch_learner import (
+                PPOTorchLearner,
+            )
+
+            return PPOTorchLearner
+        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(AlgorithmConfig)
+    def training(
+        self,
+        *,
+        use_critic: Optional[bool] = NotProvided,
+        use_gae: Optional[bool] = NotProvided,
+        lambda_: Optional[float] = NotProvided,
+        use_kl_loss: Optional[bool] = NotProvided,
+        kl_coeff: Optional[float] = NotProvided,
+        kl_target: Optional[float] = NotProvided,
+        vf_loss_coeff: Optional[float] = NotProvided,
+        entropy_coeff: Optional[float] = NotProvided,
+        entropy_coeff_schedule: Optional[List[List[Union[int, float]]]] = NotProvided,
+        clip_param: Optional[float] = NotProvided,
+        vf_clip_param: Optional[float] = NotProvided,
+        grad_clip: Optional[float] = NotProvided,
+        # @OldAPIStack
+        lr_schedule: Optional[List[List[Union[int, float]]]] = NotProvided,
+        # Deprecated.
+        vf_share_layers=DEPRECATED_VALUE,
+        **kwargs,
+    ) -> Self:
+        """Sets the training related configuration.
+
+        Args:
+            use_critic: Should use a critic as a baseline (otherwise don't use value
+                baseline; required for using GAE).
+            use_gae: If true, use the Generalized Advantage Estimator (GAE)
+                with a value function, see https://arxiv.org/pdf/1506.02438.pdf.
+            lambda_: The lambda parameter for General Advantage Estimation (GAE).
+                Defines the exponential weight used between actually measured rewards
+                vs value function estimates over multiple time steps. Specifically,
+                `lambda_` balances short-term, low-variance estimates against long-term,
+                high-variance returns. A `lambda_` of 0.0 makes the GAE rely only on
+                immediate rewards (and vf predictions from there on, reducing variance,
+                but increasing bias), while a `lambda_` of 1.0 only incorporates vf
+                predictions at the truncation points of the given episodes or episode
+                chunks (reducing bias but increasing variance).
+            use_kl_loss: Whether to use the KL-term in the loss function.
+            kl_coeff: Initial coefficient for KL divergence.
+            kl_target: Target value for KL divergence.
+            vf_loss_coeff: Coefficient of the value function loss. IMPORTANT: you must
+                tune this if you set vf_share_layers=True inside your model's config.
+            entropy_coeff: The entropy coefficient (float) or entropy coefficient
+                schedule in the format of
+                [[timestep, coeff-value], [timestep, coeff-value], ...]
+                In case of a schedule, intermediary timesteps will be assigned to
+                linearly interpolated coefficient values. A schedule config's first
+                entry must start with timestep 0, i.e.: [[0, initial_value], [...]].
+            clip_param: The PPO clip parameter.
+            vf_clip_param: Clip param for the value function. Note that this is
+                sensitive to the scale of the rewards. If your expected V is large,
+                increase this.
+            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 use_critic is not NotProvided:
+            self.use_critic = use_critic
+            # TODO (Kourosh) This is experimental.
+            #  Don't forget to remove .use_critic from algorithm config.
+        if use_gae is not NotProvided:
+            self.use_gae = use_gae
+        if lambda_ is not NotProvided:
+            self.lambda_ = lambda_
+        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 vf_loss_coeff is not NotProvided:
+            self.vf_loss_coeff = vf_loss_coeff
+        if entropy_coeff is not NotProvided:
+            self.entropy_coeff = entropy_coeff
+        if clip_param is not NotProvided:
+            self.clip_param = clip_param
+        if vf_clip_param is not NotProvided:
+            self.vf_clip_param = vf_clip_param
+        if grad_clip is not NotProvided:
+            self.grad_clip = grad_clip
+
+        # TODO (sven): Remove these once new API stack is only option for PPO.
+        if lr_schedule is not NotProvided:
+            self.lr_schedule = lr_schedule
+        if entropy_coeff_schedule is not NotProvided:
+            self.entropy_coeff_schedule = entropy_coeff_schedule
+
+        return self
+
+    @override(AlgorithmConfig)
+    def validate(self) -> None:
+        # Call super's validation method.
+        super().validate()
+
+        # Synchronous sampling, on-policy/PPO algos -> Check mismatches between
+        # `rollout_fragment_length` and `train_batch_size_per_learner` to avoid user
+        # confusion.
+        # TODO (sven): Make rollout_fragment_length a property and create a private
+        #  attribute to store (possibly) user provided value (or "auto") in. Deprecate
+        #  `self.get_rollout_fragment_length()`.
+        self.validate_train_batch_size_vs_rollout_fragment_length()
+
+        # SGD minibatch size must be smaller than train_batch_size (b/c
+        # we subsample a batch of `minibatch_size` from the train-batch for
+        # each `num_epochs`).
+        if (
+            not self.enable_rl_module_and_learner
+            and self.minibatch_size > self.train_batch_size
+        ):
+            self._value_error(
+                f"`minibatch_size` ({self.minibatch_size}) must be <= "
+                f"`train_batch_size` ({self.train_batch_size}). In PPO, the train batch"
+                f" will be split into {self.minibatch_size} chunks, each of which "
+                f"is iterated over (used for updating the policy) {self.num_epochs} "
+                "times."
+            )
+        elif self.enable_rl_module_and_learner:
+            mbs = self.minibatch_size
+            tbs = self.train_batch_size_per_learner or self.train_batch_size
+            if isinstance(mbs, int) and isinstance(tbs, int) and mbs > tbs:
+                self._value_error(
+                    f"`minibatch_size` ({mbs}) must be <= "
+                    f"`train_batch_size_per_learner` ({tbs}). In PPO, the train batch"
+                    f" will be split into {mbs} chunks, each of which is iterated over "
+                    f"(used for updating the policy) {self.num_epochs} times."
+                )
+
+        # Episodes may only be truncated (and passed into PPO's
+        # `postprocessing_fn`), iff generalized advantage estimation is used
+        # (value function estimate at end of truncated episode to estimate
+        # remaining value).
+        if (
+            not self.in_evaluation
+            and self.batch_mode == "truncate_episodes"
+            and not self.use_gae
+        ):
+            self._value_error(
+                "Episode truncation is not supported without a value "
+                "function (to estimate the return at the end of the truncated"
+                " trajectory). Consider setting "
+                "batch_mode=complete_episodes."
+            )
+
+        # New API stack checks.
+        if self.enable_rl_module_and_learner:
+            # `lr_schedule` checking.
+            if self.lr_schedule is not None:
+                self._value_error(
+                    "`lr_schedule` is deprecated and must be None! Use the "
+                    "`lr` setting to setup a schedule."
+                )
+            if self.entropy_coeff_schedule is not None:
+                self._value_error(
+                    "`entropy_coeff_schedule` is deprecated and must be None! Use the "
+                    "`entropy_coeff` setting to setup a schedule."
+                )
+            Scheduler.validate(
+                fixed_value_or_schedule=self.entropy_coeff,
+                setting_name="entropy_coeff",
+                description="entropy coefficient",
+            )
+        if isinstance(self.entropy_coeff, float) and self.entropy_coeff < 0.0:
+            self._value_error("`entropy_coeff` must be >= 0.0")
+
+    @property
+    @override(AlgorithmConfig)
+    def _model_config_auto_includes(self) -> Dict[str, Any]:
+        return super()._model_config_auto_includes | {"vf_share_layers": False}
+
+
+class PPO(Algorithm):
+    @classmethod
+    @override(Algorithm)
+    def get_default_config(cls) -> PPOConfig:
+        return PPOConfig()
+
+    @classmethod
+    @override(Algorithm)
+    def get_default_policy_class(
+        cls, config: AlgorithmConfig
+    ) -> Optional[Type[Policy]]:
+        if config["framework"] == "torch":
+
+            from ray.rllib.algorithms.ppo.ppo_torch_policy import PPOTorchPolicy
+
+            return PPOTorchPolicy
+        elif config["framework"] == "tf":
+            from ray.rllib.algorithms.ppo.ppo_tf_policy import PPOTF1Policy
+
+            return PPOTF1Policy
+        else:
+            from ray.rllib.algorithms.ppo.ppo_tf_policy import PPOTF2Policy
+
+            return PPOTF2Policy
+
+    @override(Algorithm)
+    def training_step(self) -> None:
+        # Old API stack (Policy, RolloutWorker, Connector).
+        if not self.config.enable_env_runner_and_connector_v2:
+            return self._training_step_old_api_stack()
+
+        # Collect batches from sample workers until we have a full batch.
+        with self.metrics.log_time((TIMERS, ENV_RUNNER_SAMPLING_TIMER)):
+            # Sample in parallel from the workers.
+            if self.config.count_steps_by == "agent_steps":
+                episodes, env_runner_results = synchronous_parallel_sample(
+                    worker_set=self.env_runner_group,
+                    max_agent_steps=self.config.total_train_batch_size,
+                    sample_timeout_s=self.config.sample_timeout_s,
+                    _uses_new_env_runners=(
+                        self.config.enable_env_runner_and_connector_v2
+                    ),
+                    _return_metrics=True,
+                )
+            else:
+                episodes, env_runner_results = synchronous_parallel_sample(
+                    worker_set=self.env_runner_group,
+                    max_env_steps=self.config.total_train_batch_size,
+                    sample_timeout_s=self.config.sample_timeout_s,
+                    _uses_new_env_runners=(
+                        self.config.enable_env_runner_and_connector_v2
+                    ),
+                    _return_metrics=True,
+                )
+            # Return early if all our workers failed.
+            if not episodes:
+                return
+
+            # Reduce EnvRunner metrics over the n EnvRunners.
+            self.metrics.aggregate(env_runner_results, key=ENV_RUNNER_RESULTS)
+
+        # Perform a learner update step on the collected episodes.
+        with self.metrics.log_time((TIMERS, LEARNER_UPDATE_TIMER)):
+            learner_results = self.learner_group.update(
+                episodes=episodes,
+                timesteps={
+                    NUM_ENV_STEPS_SAMPLED_LIFETIME: (
+                        self.metrics.peek(
+                            (ENV_RUNNER_RESULTS, NUM_ENV_STEPS_SAMPLED_LIFETIME)
+                        )
+                    ),
+                    NUM_MODULE_STEPS_TRAINED_LIFETIME: (
+                        self.metrics.peek(
+                            (
+                                LEARNER_RESULTS,
+                                ALL_MODULES,
+                                NUM_MODULE_STEPS_TRAINED_LIFETIME,
+                            ),
+                            default=0,
+                        )
+                    ),
+                },
+                num_epochs=self.config.num_epochs,
+                minibatch_size=self.config.minibatch_size,
+                shuffle_batch_per_epoch=self.config.shuffle_batch_per_epoch,
+            )
+            self.metrics.aggregate(learner_results, key=LEARNER_RESULTS)
+
+        # Update weights - after learning on the local worker - on all remote
+        # workers.
+        with self.metrics.log_time((TIMERS, SYNCH_WORKER_WEIGHTS_TIMER)):
+            # The train results's loss keys are ModuleIDs to their loss values.
+            # But we also return a total_loss key at the same level as the ModuleID
+            # keys. So we need to subtract that to get the correct set of ModuleIDs to
+            # update.
+            # TODO (sven): We should not be using `learner_results` as a messenger
+            #  to infer which modules to update. `policies_to_train` might also NOT work
+            #  as it might be a very large set (100s of Modules) vs a smaller Modules
+            #  set that's present in the current train batch.
+            modules_to_update = set(learner_results[0].keys()) - {ALL_MODULES}
+            self.env_runner_group.sync_weights(
+                # Sync weights from learner_group to all EnvRunners.
+                from_worker_or_learner_group=self.learner_group,
+                policies=modules_to_update,
+                inference_only=True,
+            )
+
+    @OldAPIStack
+    def _training_step_old_api_stack(self) -> ResultDict:
+        # Collect batches from sample workers until we have a full batch.
+        with self._timers[SAMPLE_TIMER]:
+            if self.config.count_steps_by == "agent_steps":
+                train_batch = synchronous_parallel_sample(
+                    worker_set=self.env_runner_group,
+                    max_agent_steps=self.config.total_train_batch_size,
+                    sample_timeout_s=self.config.sample_timeout_s,
+                )
+            else:
+                train_batch = synchronous_parallel_sample(
+                    worker_set=self.env_runner_group,
+                    max_env_steps=self.config.total_train_batch_size,
+                    sample_timeout_s=self.config.sample_timeout_s,
+                )
+            # Return early if all our workers failed.
+            if not train_batch:
+                return {}
+            train_batch = train_batch.as_multi_agent()
+            self._counters[NUM_AGENT_STEPS_SAMPLED] += train_batch.agent_steps()
+            self._counters[NUM_ENV_STEPS_SAMPLED] += train_batch.env_steps()
+            # Standardize advantages.
+            train_batch = standardize_fields(train_batch, ["advantages"])
+
+        if self.config.simple_optimizer:
+            train_results = train_one_step(self, train_batch)
+        else:
+            train_results = multi_gpu_train_one_step(self, train_batch)
+
+        policies_to_update = list(train_results.keys())
+
+        global_vars = {
+            "timestep": self._counters[NUM_AGENT_STEPS_SAMPLED],
+            # TODO (sven): num_grad_updates per each policy should be
+            #  accessible via `train_results` (and get rid of global_vars).
+            "num_grad_updates_per_policy": {
+                pid: self.env_runner.policy_map[pid].num_grad_updates
+                for pid in policies_to_update
+            },
+        }
+
+        # Update weights - after learning on the local worker - on all remote
+        # workers.
+        with self._timers[SYNCH_WORKER_WEIGHTS_TIMER]:
+            if self.env_runner_group.num_remote_workers() > 0:
+                from_worker_or_learner_group = None
+                self.env_runner_group.sync_weights(
+                    from_worker_or_learner_group=from_worker_or_learner_group,
+                    policies=policies_to_update,
+                    global_vars=global_vars,
+                )
+
+        # For each policy: Update KL scale and warn about possible issues
+        for policy_id, policy_info in train_results.items():
+            # Update KL loss with dynamic scaling
+            # for each (possibly multiagent) policy we are training
+            kl_divergence = policy_info[LEARNER_STATS_KEY].get("kl")
+            self.get_policy(policy_id).update_kl(kl_divergence)
+
+            # Warn about excessively high value function loss
+            scaled_vf_loss = (
+                self.config.vf_loss_coeff * policy_info[LEARNER_STATS_KEY]["vf_loss"]
+            )
+            policy_loss = policy_info[LEARNER_STATS_KEY]["policy_loss"]
+            if (
+                log_once("ppo_warned_lr_ratio")
+                and self.config.get("model", {}).get("vf_share_layers")
+                and scaled_vf_loss > 100
+            ):
+                logger.warning(
+                    "The magnitude of your value function loss for policy: {} is "
+                    "extremely large ({}) compared to the policy loss ({}). This "
+                    "can prevent the policy from learning. Consider scaling down "
+                    "the VF loss by reducing vf_loss_coeff, or disabling "
+                    "vf_share_layers.".format(policy_id, scaled_vf_loss, policy_loss)
+                )
+            # Warn about bad clipping configs.
+            train_batch.policy_batches[policy_id].set_get_interceptor(None)
+            mean_reward = train_batch.policy_batches[policy_id]["rewards"].mean()
+            if (
+                log_once("ppo_warned_vf_clip")
+                and mean_reward > self.config.vf_clip_param
+            ):
+                self.warned_vf_clip = True
+                logger.warning(
+                    f"The mean reward returned from the environment is {mean_reward}"
+                    f" but the vf_clip_param is set to {self.config['vf_clip_param']}."
+                    f" Consider increasing it for policy: {policy_id} to improve"
+                    " value function convergence."
+                )
+
+        # Update global vars on local worker as well.
+        # TODO (simon): At least in RolloutWorker obsolete I guess as called in
+        #  `sync_weights()` called above if remote workers. Can we call this
+        #  where `set_weights()` is called on the local_worker?
+        self.env_runner.set_global_vars(global_vars)
+
+        return train_results

rllib/algorithms/ppo/ppo_catalog.py

@@ -0,0 +1,201 @@
+# __sphinx_doc_begin__
+import gymnasium as gym
+
+from ray.rllib.core.models.base import ActorCriticEncoder, Encoder, Model
+from ray.rllib.core.models.catalog import Catalog
+from ray.rllib.core.models.configs import (
+    ActorCriticEncoderConfig,
+    FreeLogStdMLPHeadConfig,
+    MLPHeadConfig,
+)
+from ray.rllib.utils import override
+from ray.rllib.utils.annotations import OverrideToImplementCustomLogic
+
+
+def _check_if_diag_gaussian(action_distribution_cls, framework, no_error=False):
+    if framework == "torch":
+        from ray.rllib.core.distribution.torch.torch_distribution import (
+            TorchDiagGaussian,
+        )
+
+        is_diag_gaussian = issubclass(action_distribution_cls, TorchDiagGaussian)
+        if no_error:
+            return is_diag_gaussian
+        else:
+            assert is_diag_gaussian, (
+                f"free_log_std is only supported for DiagGaussian action "
+                f"distributions. Found action distribution: {action_distribution_cls}."
+            )
+    else:
+        raise ValueError(f"Framework {framework} not supported for free_log_std.")
+
+
+class PPOCatalog(Catalog):
+    """The Catalog class used to build models for PPO.
+
+    PPOCatalog provides the following models:
+        - ActorCriticEncoder: The encoder used to encode the observations.
+        - Pi Head: The head used to compute the policy logits.
+        - Value Function Head: The head used to compute the value function.
+
+    The ActorCriticEncoder is a wrapper around Encoders to produce separate outputs
+    for the policy and value function. See implementations of DefaultPPORLModule for
+    more details.
+
+    Any custom ActorCriticEncoder can be built by overriding the
+    build_actor_critic_encoder() method. Alternatively, the ActorCriticEncoderConfig
+    at PPOCatalog.actor_critic_encoder_config can be overridden to build a custom
+    ActorCriticEncoder during RLModule runtime.
+
+    Any custom head can be built by overriding the build_pi_head() and build_vf_head()
+    methods. Alternatively, the PiHeadConfig and VfHeadConfig can be overridden to
+    build custom heads during RLModule runtime.
+
+    Any module built for exploration or inference is built with the flag
+    `ìnference_only=True` and does not contain a value network. This flag can be set
+    in the `SingleAgentModuleSpec` through the `inference_only` boolean flag.
+    In case that the actor-critic-encoder is not shared between the policy and value
+    function, the inference-only module will contain only the actor encoder network.
+    """
+
+    def __init__(
+        self,
+        observation_space: gym.Space,
+        action_space: gym.Space,
+        model_config_dict: dict,
+    ):
+        """Initializes the PPOCatalog.
+
+        Args:
+            observation_space: The observation space of the Encoder.
+            action_space: The action space for the Pi Head.
+            model_config_dict: The model config to use.
+        """
+        super().__init__(
+            observation_space=observation_space,
+            action_space=action_space,
+            model_config_dict=model_config_dict,
+        )
+        # Replace EncoderConfig by ActorCriticEncoderConfig
+        self.actor_critic_encoder_config = ActorCriticEncoderConfig(
+            base_encoder_config=self._encoder_config,
+            shared=self._model_config_dict["vf_share_layers"],
+        )
+
+        self.pi_and_vf_head_hiddens = self._model_config_dict["head_fcnet_hiddens"]
+        self.pi_and_vf_head_activation = self._model_config_dict[
+            "head_fcnet_activation"
+        ]
+
+        # We don't have the exact (framework specific) action dist class yet and thus
+        # cannot determine the exact number of output nodes (action space) required.
+        # -> Build pi config only in the `self.build_pi_head` method.
+        self.pi_head_config = None
+
+        self.vf_head_config = MLPHeadConfig(
+            input_dims=self.latent_dims,
+            hidden_layer_dims=self.pi_and_vf_head_hiddens,
+            hidden_layer_activation=self.pi_and_vf_head_activation,
+            hidden_layer_use_layernorm=self._model_config_dict.get(
+                "head_fcnet_use_layernorm", False
+            ),
+            output_layer_activation="linear",
+            output_layer_dim=1,
+        )
+
+    @OverrideToImplementCustomLogic
+    def build_actor_critic_encoder(self, framework: str) -> ActorCriticEncoder:
+        """Builds the ActorCriticEncoder.
+
+        The default behavior is to build the encoder from the encoder_config.
+        This can be overridden to build a custom ActorCriticEncoder as a means of
+        configuring the behavior of a PPORLModule implementation.
+
+        Args:
+            framework: The framework to use. Either "torch" or "tf2".
+
+        Returns:
+            The ActorCriticEncoder.
+        """
+        return self.actor_critic_encoder_config.build(framework=framework)
+
+    @override(Catalog)
+    def build_encoder(self, framework: str) -> Encoder:
+        """Builds the encoder.
+
+        Since PPO uses an ActorCriticEncoder, this method should not be implemented.
+        """
+        raise NotImplementedError(
+            "Use PPOCatalog.build_actor_critic_encoder() instead for PPO."
+        )
+
+    @OverrideToImplementCustomLogic
+    def build_pi_head(self, framework: str) -> Model:
+        """Builds the policy head.
+
+        The default behavior is to build the head from the pi_head_config.
+        This can be overridden to build a custom policy head as a means of configuring
+        the behavior of a PPORLModule implementation.
+
+        Args:
+            framework: The framework to use. Either "torch" or "tf2".
+
+        Returns:
+            The policy head.
+        """
+        # Get action_distribution_cls to find out about the output dimension for pi_head
+        action_distribution_cls = self.get_action_dist_cls(framework=framework)
+        if self._model_config_dict["free_log_std"]:
+            _check_if_diag_gaussian(
+                action_distribution_cls=action_distribution_cls, framework=framework
+            )
+            is_diag_gaussian = True
+        else:
+            is_diag_gaussian = _check_if_diag_gaussian(
+                action_distribution_cls=action_distribution_cls,
+                framework=framework,
+                no_error=True,
+            )
+        required_output_dim = action_distribution_cls.required_input_dim(
+            space=self.action_space, model_config=self._model_config_dict
+        )
+        # Now that we have the action dist class and number of outputs, we can define
+        # our pi-config and build the pi head.
+        pi_head_config_class = (
+            FreeLogStdMLPHeadConfig
+            if self._model_config_dict["free_log_std"]
+            else MLPHeadConfig
+        )
+        self.pi_head_config = pi_head_config_class(
+            input_dims=self.latent_dims,
+            hidden_layer_dims=self.pi_and_vf_head_hiddens,
+            hidden_layer_activation=self.pi_and_vf_head_activation,
+            hidden_layer_use_layernorm=self._model_config_dict.get(
+                "head_fcnet_use_layernorm", False
+            ),
+            output_layer_dim=required_output_dim,
+            output_layer_activation="linear",
+            clip_log_std=is_diag_gaussian,
+            log_std_clip_param=self._model_config_dict.get("log_std_clip_param", 20),
+        )
+
+        return self.pi_head_config.build(framework=framework)
+
+    @OverrideToImplementCustomLogic
+    def build_vf_head(self, framework: str) -> Model:
+        """Builds the value function head.
+
+        The default behavior is to build the head from the vf_head_config.
+        This can be overridden to build a custom value function head as a means of
+        configuring the behavior of a PPORLModule implementation.
+
+        Args:
+            framework: The framework to use. Either "torch" or "tf2".
+
+        Returns:
+            The value function head.
+        """
+        return self.vf_head_config.build(framework=framework)
+
+
+# __sphinx_doc_end__

rllib/algorithms/ppo/ppo_learner.py

@@ -0,0 +1,146 @@
+import abc
+from typing import Any, Dict
+
+from ray.rllib.algorithms.ppo.ppo import (
+    LEARNER_RESULTS_CURR_ENTROPY_COEFF_KEY,
+    LEARNER_RESULTS_KL_KEY,
+    PPOConfig,
+)
+from ray.rllib.connectors.learner import (
+    AddOneTsToEpisodesAndTruncate,
+    GeneralAdvantageEstimation,
+)
+from ray.rllib.core.learner.learner import Learner
+from ray.rllib.core.rl_module.apis.value_function_api import ValueFunctionAPI
+from ray.rllib.utils.annotations import (
+    OverrideToImplementCustomLogic_CallToSuperRecommended,
+    override,
+)
+from ray.rllib.utils.lambda_defaultdict import LambdaDefaultDict
+from ray.rllib.utils.metrics import (
+    NUM_ENV_STEPS_SAMPLED_LIFETIME,
+)
+from ray.rllib.utils.numpy import convert_to_numpy
+from ray.rllib.utils.schedules.scheduler import Scheduler
+from ray.rllib.utils.typing import ModuleID, TensorType
+
+
+class PPOLearner(Learner):
+    @override(Learner)
+    def build(self) -> None:
+        super().build()
+
+        # Dict mapping module IDs to the respective entropy Scheduler instance.
+        self.entropy_coeff_schedulers_per_module: Dict[
+            ModuleID, Scheduler
+        ] = LambdaDefaultDict(
+            lambda module_id: Scheduler(
+                fixed_value_or_schedule=(
+                    self.config.get_config_for_module(module_id).entropy_coeff
+                ),
+                framework=self.framework,
+                device=self._device,
+            )
+        )
+
+        # Set up KL coefficient variables (per module).
+        # Note that the KL coeff is not controlled by a Scheduler, but seeks
+        # to stay close to a given kl_target value.
+        self.curr_kl_coeffs_per_module: Dict[ModuleID, TensorType] = LambdaDefaultDict(
+            lambda module_id: self._get_tensor_variable(
+                self.config.get_config_for_module(module_id).kl_coeff
+            )
+        )
+
+        # Extend all episodes by one artificial timestep to allow the value function net
+        # to compute the bootstrap values (and add a mask to the batch to know, which
+        # slots to mask out).
+        if (
+            self._learner_connector is not None
+            and self.config.add_default_connectors_to_learner_pipeline
+        ):
+            # 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).
+            self._learner_connector.prepend(AddOneTsToEpisodesAndTruncate())
+            # 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`.
+            self._learner_connector.append(
+                GeneralAdvantageEstimation(
+                    gamma=self.config.gamma, lambda_=self.config.lambda_
+                )
+            )
+
+    @override(Learner)
+    def remove_module(self, module_id: ModuleID, **kwargs):
+        marl_spec = super().remove_module(module_id, **kwargs)
+
+        self.entropy_coeff_schedulers_per_module.pop(module_id, None)
+        self.curr_kl_coeffs_per_module.pop(module_id, None)
+
+        return marl_spec
+
+    @OverrideToImplementCustomLogic_CallToSuperRecommended
+    @override(Learner)
+    def after_gradient_based_update(
+        self,
+        *,
+        timesteps: Dict[str, Any],
+    ) -> None:
+        super().after_gradient_based_update(timesteps=timesteps)
+
+        for module_id, module in self.module._rl_modules.items():
+            config = self.config.get_config_for_module(module_id)
+
+            # Update entropy coefficient via our Scheduler.
+            new_entropy_coeff = self.entropy_coeff_schedulers_per_module[
+                module_id
+            ].update(timestep=timesteps.get(NUM_ENV_STEPS_SAMPLED_LIFETIME, 0))
+            self.metrics.log_value(
+                (module_id, LEARNER_RESULTS_CURR_ENTROPY_COEFF_KEY),
+                new_entropy_coeff,
+                window=1,
+            )
+            if (
+                config.use_kl_loss
+                and (module_id, LEARNER_RESULTS_KL_KEY) in self.metrics
+            ):
+                kl_loss = convert_to_numpy(
+                    self.metrics.peek((module_id, LEARNER_RESULTS_KL_KEY))
+                )
+                self._update_module_kl_coeff(
+                    module_id=module_id,
+                    config=config,
+                    kl_loss=kl_loss,
+                )
+
+    @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]
+
+    @abc.abstractmethod
+    def _update_module_kl_coeff(
+        self,
+        *,
+        module_id: ModuleID,
+        config: PPOConfig,
+        kl_loss: float,
+    ) -> 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`.
+            kl_loss: The mean KL loss of the module, computed inside
+                `compute_loss_for_module()`.
+        """

rllib/algorithms/ppo/ppo_rl_module.py

@@ -0,0 +1,11 @@
+# Backward compat import.
+from ray.rllib.algorithms.ppo.default_ppo_rl_module import (  # noqa
+    DefaultPPORLModule as PPORLModule,
+)
+from ray._common.deprecation import deprecation_warning
+
+deprecation_warning(
+    old="ray.rllib.algorithms.ppo.ppo_rl_module.PPORLModule",
+    new="ray.rllib.algorithms.ppo.default_ppo_rl_module.DefaultPPORLModule",
+    error=False,
+)

rllib/algorithms/ppo/ppo_tf_policy.py

@@ -0,0 +1,235 @@
+"""
+TensorFlow policy class used for PPO.
+"""
+
+import logging
+from typing import Dict, List, Type, Union
+
+from ray.rllib.evaluation.postprocessing import (
+    Postprocessing,
+    compute_gae_for_sample_batch,
+)
+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.sample_batch import SampleBatch
+from ray.rllib.policy.tf_mixins import (
+    EntropyCoeffSchedule,
+    KLCoeffMixin,
+    LearningRateSchedule,
+    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, warn_if_infinite_kl_divergence
+from ray.rllib.utils.typing import AlgorithmConfigDict, TensorType, TFPolicyV2Type
+
+tf1, tf, tfv = try_import_tf()
+
+logger = logging.getLogger(__name__)
+
+
+def validate_config(config: AlgorithmConfigDict) -> None:
+    """Executed before Policy is "initialized" (at beginning of constructor).
+    Args:
+        config: The Policy's config.
+    """
+    # If vf_share_layers is True, inform about the need to tune vf_loss_coeff.
+    if config.get("model", {}).get("vf_share_layers") is True:
+        logger.info(
+            "`vf_share_layers=True` in your model. "
+            "Therefore, remember to tune the value of `vf_loss_coeff`!"
+        )
+
+
+# We need this builder function because we want to share the same
+# custom logics between TF1 dynamic and TF2 eager policies.
+def get_ppo_tf_policy(name: str, base: TFPolicyV2Type) -> TFPolicyV2Type:
+    """Construct a PPOTFPolicy 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 PPO.
+    """
+
+    class PPOTFPolicy(
+        EntropyCoeffSchedule,
+        LearningRateSchedule,
+        KLCoeffMixin,
+        ValueNetworkMixin,
+        base,
+    ):
+        def __init__(
+            self,
+            observation_space,
+            action_space,
+            config,
+            existing_model=None,
+            existing_inputs=None,
+        ):
+            # First thing first, enable eager execution if necessary.
+            base.enable_eager_execution_if_necessary()
+
+            # TODO: Move into Policy API, if needed at all here. Why not move this into
+            #  `PPOConfig`?.
+            validate_config(config)
+
+            # Initialize base class.
+            base.__init__(
+                self,
+                observation_space,
+                action_space,
+                config,
+                existing_inputs=existing_inputs,
+                existing_model=existing_model,
+            )
+
+            # Initialize MixIns.
+            ValueNetworkMixin.__init__(self, config)
+            EntropyCoeffSchedule.__init__(
+                self, config["entropy_coeff"], config["entropy_coeff_schedule"]
+            )
+            LearningRateSchedule.__init__(self, config["lr"], config["lr_schedule"])
+            KLCoeffMixin.__init__(self, config)
+
+            # 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]]:
+            if isinstance(model, tf.keras.Model):
+                logits, state, extra_outs = model(train_batch)
+                value_fn_out = extra_outs[SampleBatch.VF_PREDS]
+            else:
+                logits, state = model(train_batch)
+                value_fn_out = model.value_function()
+
+            curr_action_dist = dist_class(logits, model)
+
+            # RNN case: Mask away 0-padded chunks at end of time axis.
+            if state:
+                # Derive max_seq_len from the data itself, not from the seq_lens
+                # tensor. This is in case e.g. seq_lens=[2, 3], but the data is still
+                # 0-padded up to T=5 (as it's the case for attention nets).
+                B = tf.shape(train_batch[SampleBatch.SEQ_LENS])[0]
+                max_seq_len = tf.shape(logits)[0] // B
+
+                mask = tf.sequence_mask(train_batch[SampleBatch.SEQ_LENS], max_seq_len)
+                mask = tf.reshape(mask, [-1])
+
+                def reduce_mean_valid(t):
+                    return tf.reduce_mean(tf.boolean_mask(t, mask))
+
+            # non-RNN case: No masking.
+            else:
+                mask = None
+                reduce_mean_valid = tf.reduce_mean
+
+            prev_action_dist = dist_class(
+                train_batch[SampleBatch.ACTION_DIST_INPUTS], model
+            )
+
+            logp_ratio = tf.exp(
+                curr_action_dist.logp(train_batch[SampleBatch.ACTIONS])
+                - train_batch[SampleBatch.ACTION_LOGP]
+            )
+
+            # Only calculate kl loss if necessary (kl-coeff > 0.0).
+            if self.config["kl_coeff"] > 0.0:
+                action_kl = prev_action_dist.kl(curr_action_dist)
+                mean_kl_loss = reduce_mean_valid(action_kl)
+                warn_if_infinite_kl_divergence(self, mean_kl_loss)
+            else:
+                mean_kl_loss = tf.constant(0.0)
+
+            curr_entropy = curr_action_dist.entropy()
+            mean_entropy = reduce_mean_valid(curr_entropy)
+
+            surrogate_loss = tf.minimum(
+                train_batch[Postprocessing.ADVANTAGES] * logp_ratio,
+                train_batch[Postprocessing.ADVANTAGES]
+                * tf.clip_by_value(
+                    logp_ratio,
+                    1 - self.config["clip_param"],
+                    1 + self.config["clip_param"],
+                ),
+            )
+
+            # Compute a value function loss.
+            if self.config["use_critic"]:
+                vf_loss = tf.math.square(
+                    value_fn_out - train_batch[Postprocessing.VALUE_TARGETS]
+                )
+                vf_loss_clipped = tf.clip_by_value(
+                    vf_loss,
+                    0,
+                    self.config["vf_clip_param"],
+                )
+                mean_vf_loss = reduce_mean_valid(vf_loss_clipped)
+            # Ignore the value function.
+            else:
+                vf_loss_clipped = mean_vf_loss = tf.constant(0.0)
+
+            total_loss = reduce_mean_valid(
+                -surrogate_loss
+                + self.config["vf_loss_coeff"] * vf_loss_clipped
+                - self.entropy_coeff * curr_entropy
+            )
+            # Add mean_kl_loss (already processed through `reduce_mean_valid`),
+            # if necessary.
+            if self.config["kl_coeff"] > 0.0:
+                total_loss += self.kl_coeff * mean_kl_loss
+
+            # Store stats in policy for stats_fn.
+            self._total_loss = total_loss
+            self._mean_policy_loss = reduce_mean_valid(-surrogate_loss)
+            self._mean_vf_loss = mean_vf_loss
+            self._mean_entropy = mean_entropy
+            # Backward compatibility: Deprecate self._mean_kl.
+            self._mean_kl_loss = self._mean_kl = mean_kl_loss
+            self._value_fn_out = value_fn_out
+
+            return total_loss
+
+        @override(base)
+        def stats_fn(self, train_batch: SampleBatch) -> Dict[str, TensorType]:
+            return {
+                "cur_kl_coeff": tf.cast(self.kl_coeff, tf.float64),
+                "cur_lr": tf.cast(self.cur_lr, tf.float64),
+                "total_loss": self._total_loss,
+                "policy_loss": self._mean_policy_loss,
+                "vf_loss": self._mean_vf_loss,
+                "vf_explained_var": explained_variance(
+                    train_batch[Postprocessing.VALUE_TARGETS], self._value_fn_out
+                ),
+                "kl": self._mean_kl_loss,
+                "entropy": self._mean_entropy,
+                "entropy_coeff": tf.cast(self.entropy_coeff, tf.float64),
+            }
+
+        @override(base)
+        def postprocess_trajectory(
+            self, sample_batch, other_agent_batches=None, episode=None
+        ):
+            sample_batch = super().postprocess_trajectory(sample_batch)
+            return compute_gae_for_sample_batch(
+                self, sample_batch, other_agent_batches, episode
+            )
+
+    PPOTFPolicy.__name__ = name
+    PPOTFPolicy.__qualname__ = name
+
+    return PPOTFPolicy
+
+
+PPOTF1Policy = get_ppo_tf_policy("PPOTF1Policy", DynamicTFPolicyV2)
+PPOTF2Policy = get_ppo_tf_policy("PPOTF2Policy", EagerTFPolicyV2)

rllib/algorithms/ppo/ppo_torch_policy.py

@@ -0,0 +1,217 @@
+import logging
+from typing import Dict, List, Type, Union
+
+import ray
+from ray.rllib.algorithms.ppo.ppo_tf_policy import validate_config
+from ray.rllib.evaluation.postprocessing import (
+    Postprocessing,
+    compute_gae_for_sample_batch,
+)
+from ray.rllib.models.action_dist import ActionDistribution
+from ray.rllib.models.modelv2 import ModelV2
+from ray.rllib.policy.sample_batch import SampleBatch
+from ray.rllib.policy.torch_mixins import (
+    EntropyCoeffSchedule,
+    KLCoeffMixin,
+    LearningRateSchedule,
+    ValueNetworkMixin,
+)
+from ray.rllib.policy.torch_policy_v2 import TorchPolicyV2
+from ray.rllib.utils.annotations import override
+from ray.rllib.utils.framework import try_import_torch
+from ray.rllib.utils.numpy import convert_to_numpy
+from ray.rllib.utils.torch_utils import (
+    apply_grad_clipping,
+    explained_variance,
+    sequence_mask,
+    warn_if_infinite_kl_divergence,
+)
+from ray.rllib.utils.typing import TensorType
+
+torch, nn = try_import_torch()
+
+logger = logging.getLogger(__name__)
+
+
+class PPOTorchPolicy(
+    ValueNetworkMixin,
+    LearningRateSchedule,
+    EntropyCoeffSchedule,
+    KLCoeffMixin,
+    TorchPolicyV2,
+):
+    """PyTorch policy class used with PPO."""
+
+    def __init__(self, observation_space, action_space, config):
+        config = dict(ray.rllib.algorithms.ppo.ppo.PPOConfig().to_dict(), **config)
+        validate_config(config)
+
+        TorchPolicyV2.__init__(
+            self,
+            observation_space,
+            action_space,
+            config,
+            max_seq_len=config["model"]["max_seq_len"],
+        )
+
+        ValueNetworkMixin.__init__(self, config)
+        LearningRateSchedule.__init__(self, config["lr"], config["lr_schedule"])
+        EntropyCoeffSchedule.__init__(
+            self, config["entropy_coeff"], config["entropy_coeff_schedule"]
+        )
+        KLCoeffMixin.__init__(self, config)
+
+        self._initialize_loss_from_dummy_batch()
+
+    @override(TorchPolicyV2)
+    def loss(
+        self,
+        model: ModelV2,
+        dist_class: Type[ActionDistribution],
+        train_batch: SampleBatch,
+    ) -> Union[TensorType, List[TensorType]]:
+        """Compute loss for Proximal Policy Objective.
+
+        Args:
+            model: The Model to calculate the loss for.
+            dist_class: The action distr. class.
+            train_batch: The training data.
+
+        Returns:
+            The PPO loss tensor given the input batch.
+        """
+
+        logits, state = model(train_batch)
+        curr_action_dist = dist_class(logits, model)
+
+        # RNN case: Mask away 0-padded chunks at end of time axis.
+        if state:
+            B = len(train_batch[SampleBatch.SEQ_LENS])
+            max_seq_len = logits.shape[0] // B
+            mask = sequence_mask(
+                train_batch[SampleBatch.SEQ_LENS],
+                max_seq_len,
+                time_major=model.is_time_major(),
+            )
+            mask = torch.reshape(mask, [-1])
+            num_valid = torch.sum(mask)
+
+            def reduce_mean_valid(t):
+                return torch.sum(t[mask]) / num_valid
+
+        # non-RNN case: No masking.
+        else:
+            mask = None
+            reduce_mean_valid = torch.mean
+
+        prev_action_dist = dist_class(
+            train_batch[SampleBatch.ACTION_DIST_INPUTS], model
+        )
+
+        logp_ratio = torch.exp(
+            curr_action_dist.logp(train_batch[SampleBatch.ACTIONS])
+            - train_batch[SampleBatch.ACTION_LOGP]
+        )
+
+        # Only calculate kl loss if necessary (kl-coeff > 0.0).
+        if self.config["kl_coeff"] > 0.0:
+            action_kl = prev_action_dist.kl(curr_action_dist)
+            mean_kl_loss = reduce_mean_valid(action_kl)
+            # TODO smorad: should we do anything besides warn? Could discard KL term
+            # for this update
+            warn_if_infinite_kl_divergence(self, mean_kl_loss)
+        else:
+            mean_kl_loss = torch.tensor(0.0, device=logp_ratio.device)
+
+        curr_entropy = curr_action_dist.entropy()
+        mean_entropy = reduce_mean_valid(curr_entropy)
+
+        surrogate_loss = torch.min(
+            train_batch[Postprocessing.ADVANTAGES] * logp_ratio,
+            train_batch[Postprocessing.ADVANTAGES]
+            * torch.clamp(
+                logp_ratio, 1 - self.config["clip_param"], 1 + self.config["clip_param"]
+            ),
+        )
+
+        # Compute a value function loss.
+        if self.config["use_critic"]:
+            value_fn_out = model.value_function()
+            vf_loss = torch.pow(
+                value_fn_out - train_batch[Postprocessing.VALUE_TARGETS], 2.0
+            )
+            vf_loss_clipped = torch.clamp(vf_loss, 0, self.config["vf_clip_param"])
+            mean_vf_loss = reduce_mean_valid(vf_loss_clipped)
+        # Ignore the value function.
+        else:
+            value_fn_out = torch.tensor(0.0).to(surrogate_loss.device)
+            vf_loss_clipped = mean_vf_loss = torch.tensor(0.0).to(surrogate_loss.device)
+
+        total_loss = reduce_mean_valid(
+            -surrogate_loss
+            + self.config["vf_loss_coeff"] * vf_loss_clipped
+            - self.entropy_coeff * curr_entropy
+        )
+
+        # Add mean_kl_loss (already processed through `reduce_mean_valid`),
+        # if necessary.
+        if self.config["kl_coeff"] > 0.0:
+            total_loss += self.kl_coeff * mean_kl_loss
+
+        # 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"] = reduce_mean_valid(-surrogate_loss)
+        model.tower_stats["mean_vf_loss"] = mean_vf_loss
+        model.tower_stats["vf_explained_var"] = explained_variance(
+            train_batch[Postprocessing.VALUE_TARGETS], value_fn_out
+        )
+        model.tower_stats["mean_entropy"] = mean_entropy
+        model.tower_stats["mean_kl_loss"] = mean_kl_loss
+
+        return total_loss
+
+    # TODO: Make this an event-style subscription (e.g.:
+    #  "after_gradients_computed").
+    @override(TorchPolicyV2)
+    def extra_grad_process(self, local_optimizer, loss):
+        return apply_grad_clipping(self, local_optimizer, loss)
+
+    @override(TorchPolicyV2)
+    def stats_fn(self, train_batch: SampleBatch) -> Dict[str, TensorType]:
+        return convert_to_numpy(
+            {
+                "cur_kl_coeff": self.kl_coeff,
+                "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"))
+                ),
+                "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"))
+                ),
+                "kl": torch.mean(torch.stack(self.get_tower_stats("mean_kl_loss"))),
+                "entropy": torch.mean(
+                    torch.stack(self.get_tower_stats("mean_entropy"))
+                ),
+                "entropy_coeff": self.entropy_coeff,
+            }
+        )
+
+    @override(TorchPolicyV2)
+    def postprocess_trajectory(
+        self, sample_batch, other_agent_batches=None, episode=None
+    ):
+        # Do all post-processing always with no_grad().
+        # Not using this here will introduce a memory leak
+        # in torch (issue #6962).
+        # TODO: no_grad still necessary?
+        with torch.no_grad():
+            return compute_gae_for_sample_batch(
+                self, sample_batch, other_agent_batches, episode
+            )

rllib/algorithms/ppo/tests/test_ppo.py

@@ -0,0 +1,168 @@
+import unittest
+
+import ray
+import ray.rllib.algorithms.ppo as ppo
+from ray.rllib.algorithms.ppo.ppo_learner import LEARNER_RESULTS_CURR_ENTROPY_COEFF_KEY
+from ray.rllib.core import DEFAULT_MODULE_ID
+from ray.rllib.core.learner.learner import DEFAULT_OPTIMIZER, LR_KEY
+from ray.rllib.core.rl_module.default_model_config import DefaultModelConfig
+from ray.rllib.utils.metrics import LEARNER_RESULTS
+from ray.rllib.utils.test_utils import check, check_train_results_new_api_stack
+
+
+def get_model_config(lstm=False):
+    return (
+        dict(
+            use_lstm=True,
+            lstm_use_prev_action=True,
+            lstm_use_prev_reward=True,
+            lstm_cell_size=10,
+            max_seq_len=20,
+        )
+        if lstm
+        else {"use_lstm": False}
+    )
+
+
+def on_train_result(algorithm, result: dict, **kwargs):
+    stats = result[LEARNER_RESULTS][DEFAULT_MODULE_ID]
+    # Entropy coeff goes to 0.05, then 0.0 (per iter).
+    check(
+        stats[LEARNER_RESULTS_CURR_ENTROPY_COEFF_KEY],
+        0.05 if algorithm.iteration == 1 else 0.0,
+    )
+
+    # Learning rate should decrease by 0.0001/4 per iteration.
+    check(
+        stats[DEFAULT_OPTIMIZER + "_" + LR_KEY],
+        0.0000075 if algorithm.iteration == 1 else 0.000005,
+    )
+    # Compare reported curr lr vs the actual lr found in the optimizer object.
+    optim = algorithm.learner_group._learner.get_optimizer()
+    actual_optimizer_lr = (
+        optim.param_groups[0]["lr"]
+        if algorithm.config.framework_str == "torch"
+        else optim.lr
+    )
+    check(stats[DEFAULT_OPTIMIZER + "_" + LR_KEY], actual_optimizer_lr)
+
+
+class TestPPO(unittest.TestCase):
+    @classmethod
+    def setUpClass(cls):
+        ray.init()
+
+    @classmethod
+    def tearDownClass(cls):
+        ray.shutdown()
+
+    def test_ppo_compilation_and_schedule_mixins(self):
+        """Test whether PPO can be built with all frameworks."""
+
+        # Build a PPOConfig object with the `SingleAgentEnvRunner` class.
+        config = (
+            ppo.PPOConfig()
+            .env_runners(num_env_runners=0)
+            .training(
+                num_epochs=2,
+                # Setup lr schedule for testing lr-scheduling correctness.
+                lr=[[0, 0.00001], [512, 0.0]],  # 512=4x128
+                # Setup `entropy_coeff` schedule for testing whether it's scheduled
+                # correctly.
+                entropy_coeff=[[0, 0.1], [256, 0.0]],  # 256=2x128,
+                train_batch_size=128,
+            )
+            .callbacks(on_train_result=on_train_result)
+            .evaluation(
+                # Also test evaluation with remote workers.
+                evaluation_num_env_runners=2,
+                evaluation_duration=3,
+                evaluation_duration_unit="episodes",
+                evaluation_parallel_to_training=True,
+            )
+        )
+
+        num_iterations = 2
+
+        for env in [
+            "CartPole-v1",
+            "Pendulum-v1",
+        ]:
+            print("Env={}".format(env))
+            for lstm in [False]:
+                print("LSTM={}".format(lstm))
+                config.rl_module(model_config=get_model_config(lstm=lstm))
+
+                algo = config.build(env=env)
+                # TODO: Maybe add an API to get the Learner(s) instances within
+                #  a learner group, remote or not.
+                learner = algo.learner_group._learner
+                optim = learner.get_optimizer()
+                # Check initial LR directly set in optimizer vs the first (ts=0)
+                # value from the schedule.
+                lr = optim.param_groups[0]["lr"]
+                check(lr, config.lr[0][1])
+
+                # Check current entropy coeff value using the respective Scheduler.
+                entropy_coeff = learner.entropy_coeff_schedulers_per_module[
+                    DEFAULT_MODULE_ID
+                ].get_current_value()
+                check(entropy_coeff, 0.1)
+
+                for i in range(num_iterations):
+                    results = algo.train()
+                    check_train_results_new_api_stack(results)
+                    print(results)
+
+                # algo.evaluate()
+                algo.stop()
+
+    def test_ppo_free_log_std(self):
+        """Tests the free log std option works."""
+        config = (
+            ppo.PPOConfig()
+            .environment("Pendulum-v1")
+            .env_runners(
+                num_env_runners=1,
+            )
+            .rl_module(
+                model_config=DefaultModelConfig(
+                    fcnet_hiddens=[10],
+                    fcnet_activation="linear",
+                    free_log_std=True,
+                    vf_share_layers=True,
+                ),
+            )
+            .training(
+                gamma=0.99,
+            )
+        )
+
+        algo = config.build()
+        module = algo.get_module(DEFAULT_MODULE_ID)
+
+        # Check the free log std var is created.
+        matching = [v for (n, v) in module.named_parameters() if "log_std" in n]
+        assert len(matching) == 1, matching
+        log_std_var = matching[0]
+
+        def get_value(log_std_var=log_std_var):
+            return log_std_var.detach().cpu().numpy()[0]
+
+        # Check the variable is initially zero.
+        init_std = get_value()
+        assert init_std == 0.0, init_std
+        algo.train()
+
+        # Check the variable is updated.
+        post_std = get_value()
+        assert post_std != 0.0, post_std
+        algo.stop()
+
+
+if __name__ == "__main__":
+    import sys
+
+    import pytest
+
+    sys.exit(pytest.main(["-v", __file__]))

rllib/algorithms/ppo/tests/test_ppo_learner.py

@@ -0,0 +1,142 @@
+import tempfile
+import unittest
+
+import gymnasium as gym
+import numpy as np
+
+import ray
+import ray.rllib.algorithms.ppo as ppo
+from ray.rllib.algorithms.ppo.ppo import LEARNER_RESULTS_CURR_KL_COEFF_KEY
+from ray.rllib.core.columns import Columns
+from ray.rllib.examples.envs.classes.multi_agent import MultiAgentCartPole
+from ray.rllib.utils.metrics import LEARNER_RESULTS
+from ray.rllib.utils.test_utils import check
+from ray.tune.registry import register_env
+
+# Fake CartPole episode of n time steps.
+FAKE_BATCH = {
+    Columns.OBS: np.array(
+        [[0.1, 0.2, 0.3, 0.4], [0.5, 0.6, 0.7, 0.8], [0.9, 1.0, 1.1, 1.2]],
+        dtype=np.float32,
+    ),
+    Columns.NEXT_OBS: np.array(
+        [[0.1, 0.2, 0.3, 0.4], [0.5, 0.6, 0.7, 0.8], [0.9, 1.0, 1.1, 1.2]],
+        dtype=np.float32,
+    ),
+    Columns.ACTIONS: np.array([0, 1, 1]),
+    Columns.REWARDS: np.array([1.0, -1.0, 0.5], dtype=np.float32),
+    Columns.TERMINATEDS: np.array([False, False, True]),
+    Columns.TRUNCATEDS: np.array([False, False, False]),
+    Columns.VF_PREDS: np.array([0.5, 0.6, 0.7], dtype=np.float32),
+    Columns.ACTION_DIST_INPUTS: np.array(
+        [[-2.0, 0.5], [-3.0, -0.3], [-0.1, 2.5]], dtype=np.float32
+    ),
+    Columns.ACTION_LOGP: np.array([-0.5, -0.1, -0.2], dtype=np.float32),
+    Columns.EPS_ID: np.array([0, 0, 0]),
+}
+
+
+class TestPPO(unittest.TestCase):
+    ENV = gym.make("CartPole-v1")
+
+    @classmethod
+    def setUpClass(cls):
+        ray.init()
+
+    @classmethod
+    def tearDownClass(cls):
+        ray.shutdown()
+
+    def test_save_to_path_and_restore_from_path(self):
+        """Tests saving and loading the state of the PPO Learner Group."""
+        config = (
+            ppo.PPOConfig()
+            .environment("CartPole-v1")
+            .env_runners(
+                num_env_runners=0,
+            )
+            .training(
+                gamma=0.99,
+                model=dict(
+                    fcnet_hiddens=[10, 10],
+                    fcnet_activation="linear",
+                    vf_share_layers=False,
+                ),
+            )
+        )
+
+        algo_config = config.copy(copy_frozen=False)
+        algo_config.validate()
+        algo_config.freeze()
+        learner_group1 = algo_config.build_learner_group(env=self.ENV)
+        learner_group2 = algo_config.build_learner_group(env=self.ENV)
+        with tempfile.TemporaryDirectory() as tmpdir:
+            learner_group1.save_to_path(tmpdir)
+            learner_group2.restore_from_path(tmpdir)
+            # Remove functions from state b/c they are not comparable via `check`.
+            s1 = learner_group1.get_state()
+            s2 = learner_group2.get_state()
+            check(s1, s2)
+
+    def test_kl_coeff_changes(self):
+        # Simple environment with 4 independent cartpole entities
+        register_env(
+            "multi_agent_cartpole", lambda _: MultiAgentCartPole({"num_agents": 2})
+        )
+
+        initial_kl_coeff = 0.01
+        config = (
+            ppo.PPOConfig()
+            .environment("CartPole-v1")
+            .env_runners(
+                num_env_runners=0,
+                rollout_fragment_length=50,
+                exploration_config={},
+            )
+            .training(
+                gamma=0.99,
+                model=dict(
+                    fcnet_hiddens=[10, 10],
+                    fcnet_activation="linear",
+                    vf_share_layers=False,
+                ),
+                kl_coeff=initial_kl_coeff,
+            )
+            .environment("multi_agent_cartpole")
+            .multi_agent(
+                policies={"p0", "p1"},
+                policy_mapping_fn=lambda agent_id, episode, **kwargs: (
+                    "p{}".format(agent_id % 2)
+                ),
+            )
+        )
+
+        algo = config.build()
+        # Call train while results aren't returned because this is
+        # a asynchronous Algorithm and results are returned asynchronously.
+        curr_kl_coeff_1 = None
+        curr_kl_coeff_2 = None
+        while not curr_kl_coeff_1 or not curr_kl_coeff_2:
+            results = algo.train()
+
+            # Attempt to get the current KL coefficient from the learner.
+            # Iterate until we have found both coefficients at least once.
+            if "p0" in results[LEARNER_RESULTS]:
+                curr_kl_coeff_1 = results[LEARNER_RESULTS]["p0"][
+                    LEARNER_RESULTS_CURR_KL_COEFF_KEY
+                ]
+            if "p1" in results[LEARNER_RESULTS]:
+                curr_kl_coeff_2 = results[LEARNER_RESULTS]["p1"][
+                    LEARNER_RESULTS_CURR_KL_COEFF_KEY
+                ]
+
+        self.assertNotEqual(curr_kl_coeff_1, initial_kl_coeff)
+        self.assertNotEqual(curr_kl_coeff_2, initial_kl_coeff)
+
+
+if __name__ == "__main__":
+    import sys
+
+    import pytest
+
+    sys.exit(pytest.main(["-v", __file__]))

rllib/algorithms/ppo/tests/test_ppo_rl_module.py

@@ -0,0 +1,192 @@
+import itertools
+import unittest
+
+import gymnasium as gym
+import numpy as np
+import tree
+
+import ray
+from ray.rllib.algorithms.ppo.ppo_catalog import PPOCatalog
+from ray.rllib.algorithms.ppo.torch.default_ppo_torch_rl_module import (
+    DefaultPPOTorchRLModule,
+)
+from ray.rllib.core.columns import Columns
+from ray.rllib.core.rl_module.default_model_config import DefaultModelConfig
+from ray.rllib.models.preprocessors import get_preprocessor
+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 convert_to_torch_tensor
+
+torch, nn = try_import_torch()
+
+
+def dummy_torch_ppo_loss(module, batch, fwd_out):
+    adv = batch[Columns.REWARDS] - module.compute_values(batch)
+    action_dist_class = module.get_train_action_dist_cls()
+    action_probs = action_dist_class.from_logits(
+        fwd_out[Columns.ACTION_DIST_INPUTS]
+    ).logp(batch[Columns.ACTIONS])
+    actor_loss = -(action_probs * adv).mean()
+    critic_loss = (adv**2).mean()
+    loss = actor_loss + critic_loss
+
+    return loss
+
+
+def _get_input_batch_from_obs(obs, lstm):
+    batch = {
+        Columns.OBS: convert_to_torch_tensor(obs)[None],
+    }
+    if lstm:
+        batch[Columns.OBS] = batch[Columns.OBS][None]
+    return batch
+
+
+class TestPPO(unittest.TestCase):
+    @classmethod
+    def setUpClass(cls):
+        ray.init()
+
+    @classmethod
+    def tearDownClass(cls):
+        ray.shutdown()
+
+    def test_rollouts(self):
+        # TODO: Add FrozenLake-v1 to cover LSTM case.
+        env_names = ["CartPole-v1", "Pendulum-v1", "ale_py:ALE/Breakout-v5"]
+        fwd_fns = ["forward_exploration", "forward_inference"]
+        lstm = [True, False]
+        config_combinations = [env_names, fwd_fns, lstm]
+        for config in itertools.product(*config_combinations):
+            env_name, fwd_fn, lstm = config
+            print(f"ENV={env_name}; FWD={fwd_fn}; LSTM={lstm}")
+            env = gym.make(env_name)
+
+            preprocessor_cls = get_preprocessor(env.observation_space)
+            preprocessor = preprocessor_cls(env.observation_space)
+
+            module = DefaultPPOTorchRLModule(
+                observation_space=preprocessor.observation_space,
+                action_space=env.action_space,
+                model_config=DefaultModelConfig(use_lstm=lstm),
+                catalog_class=PPOCatalog,
+            )
+
+            obs, _ = env.reset()
+            obs = preprocessor.transform(obs)
+
+            batch = _get_input_batch_from_obs(obs, lstm)
+
+            if lstm:
+                state_in = module.get_initial_state()
+                state_in = convert_to_torch_tensor(state_in)
+                state_in = tree.map_structure(lambda x: x[None], state_in)
+                batch[Columns.STATE_IN] = state_in
+
+            if fwd_fn == "forward_exploration":
+                module.forward_exploration(batch)
+            else:
+                module.forward_inference(batch)
+
+    def test_forward_train(self):
+        # TODO: Add FrozenLake-v1 to cover LSTM case.
+        env_names = ["CartPole-v1", "Pendulum-v1", "ale_py:ALE/Breakout-v5"]
+        lstm = [False, True]
+        config_combinations = [env_names, lstm]
+        for config in itertools.product(*config_combinations):
+            env_name, lstm = config
+            print(f"ENV={env_name}; LSTM={lstm}")
+            env = gym.make(env_name)
+
+            preprocessor_cls = get_preprocessor(env.observation_space)
+            preprocessor = preprocessor_cls(env.observation_space)
+
+            module = DefaultPPOTorchRLModule(
+                observation_space=preprocessor.observation_space,
+                action_space=env.action_space,
+                model_config=DefaultModelConfig(use_lstm=lstm),
+                catalog_class=PPOCatalog,
+            )
+
+            # collect a batch of data
+            batches = []
+            obs, _ = env.reset()
+            obs = preprocessor.transform(obs)
+            tstep = 0
+
+            if lstm:
+                state_in = module.get_initial_state()
+                state_in = tree.map_structure(
+                    lambda x: x[None], convert_to_torch_tensor(state_in)
+                )
+                initial_state = state_in
+
+            while tstep < 10:
+                input_batch = _get_input_batch_from_obs(obs, lstm=lstm)
+                if lstm:
+                    input_batch[Columns.STATE_IN] = state_in
+
+                fwd_out = module.forward_exploration(input_batch)
+                action_dist_cls = module.get_exploration_action_dist_cls()
+                action_dist = action_dist_cls.from_logits(
+                    fwd_out[Columns.ACTION_DIST_INPUTS]
+                )
+                _action = action_dist.sample()
+                action = convert_to_numpy(_action[0])
+                action_logp = convert_to_numpy(action_dist.logp(_action)[0])
+                if lstm:
+                    # Since this is inference, fwd out should only contain one action
+                    assert len(action) == 1
+                    action = action[0]
+                new_obs, reward, terminated, truncated, _ = env.step(action)
+                new_obs = preprocessor.transform(new_obs)
+                output_batch = {
+                    Columns.OBS: obs,
+                    Columns.NEXT_OBS: new_obs,
+                    Columns.ACTIONS: action,
+                    Columns.ACTION_LOGP: action_logp,
+                    Columns.REWARDS: np.array(reward),
+                    Columns.TERMINATEDS: np.array(terminated),
+                    Columns.TRUNCATEDS: np.array(truncated),
+                    Columns.STATE_IN: None,
+                }
+
+                if lstm:
+                    assert Columns.STATE_OUT in fwd_out
+                    state_in = fwd_out[Columns.STATE_OUT]
+                batches.append(output_batch)
+                obs = new_obs
+                tstep += 1
+
+            # convert the list of dicts to dict of lists
+            batch = tree.map_structure(lambda *x: np.array(x), *batches)
+            # convert dict of lists to dict of tensors
+            fwd_in = {k: convert_to_torch_tensor(np.array(v)) for k, v in batch.items()}
+            if lstm:
+                fwd_in[Columns.STATE_IN] = initial_state
+                # If we test lstm, the collected timesteps make up only one batch
+                fwd_in = {
+                    k: torch.unsqueeze(v, 0) if k != Columns.STATE_IN else v
+                    for k, v in fwd_in.items()
+                }
+
+            # forward train
+            # before training make sure module is on the right device
+            # and in training mode
+            module.to("cpu")
+            module.train()
+            fwd_out = module.forward_train(fwd_in)
+            loss = dummy_torch_ppo_loss(module, fwd_in, fwd_out)
+            loss.backward()
+
+            # check that all neural net parameters have gradients
+            for param in module.parameters():
+                self.assertIsNotNone(param.grad)
+
+
+if __name__ == "__main__":
+    import sys
+
+    import pytest
+
+    sys.exit(pytest.main(["-v", __file__]))

rllib/algorithms/ppo/tests/test_ppo_value_bootstrapping.py

@@ -0,0 +1,212 @@
+"""Unit tests for PPO's value-bootstrapping wiring.
+
+Exercises the connector pipeline (``AddOneTsToEpisodesAndTruncate`` +
+``AddColumnsFromEpisodesToTrainBatch`` + ``BatchIndividualItems``) feeding into
+``compute_value_targets``. Targets are pinned to closed-form GAE answers so a
+regression in either the connector layout or the GAE recursion is caught.
+"""
+import numpy as np
+import pytest
+
+from ray.rllib.algorithms.ppo import PPOConfig
+from ray.rllib.connectors.env_to_module import FlattenObservations
+from ray.rllib.connectors.learner import (
+    AddColumnsFromEpisodesToTrainBatch,
+    AddOneTsToEpisodesAndTruncate,
+    BatchIndividualItems,
+    LearnerConnectorPipeline,
+)
+from ray.rllib.core import DEFAULT_MODULE_ID
+from ray.rllib.core.columns import Columns
+from ray.rllib.core.rl_module.default_model_config import DefaultModelConfig
+from ray.rllib.env.single_agent_episode import SingleAgentEpisode
+from ray.rllib.utils.framework import try_import_torch
+from ray.rllib.utils.postprocessing.value_predictions import compute_value_targets
+from ray.rllib.utils.postprocessing.zero_padding import unpad_data_if_necessary
+from ray.rllib.utils.torch_utils import convert_to_torch_tensor
+
+torch, _ = try_import_torch()
+
+
+def _targets(per_ep_values, per_ep_rewards, terminated, truncated, gamma, lambda_):
+    """Run the real learner pipeline, then ``compute_value_targets``."""
+    episodes = [
+        SingleAgentEpisode(
+            observations=[0] * len(v),
+            actions=[0] * len(r),
+            rewards=r,
+            terminated=t,
+            truncated=u,
+            len_lookback_buffer=0,
+        )
+        for v, r, t, u in zip(per_ep_values, per_ep_rewards, terminated, truncated)
+    ]
+    pipe = LearnerConnectorPipeline(
+        connectors=[
+            AddOneTsToEpisodesAndTruncate(),
+            AddColumnsFromEpisodesToTrainBatch(),
+            BatchIndividualItems(),
+        ]
+    )
+    batch = pipe(
+        episodes=episodes, batch={}, rl_module=None, explore=False, shared_data={}
+    )
+    lens = [len(e) for e in episodes]
+    flat_values = np.array([v for vs in per_ep_values for v in vs], dtype=np.float32)
+    return compute_value_targets(
+        values=flat_values,
+        rewards=unpad_data_if_necessary(lens, np.array(batch[Columns.REWARDS])),
+        terminateds=unpad_data_if_necessary(lens, np.array(batch[Columns.TERMINATEDS])),
+        truncateds=unpad_data_if_necessary(lens, np.array(batch[Columns.TRUNCATEDS])),
+        gamma=gamma,
+        lambda_=lambda_,
+    )
+
+
+# Length-2 episode, values=[0, 0.95, 0.95] (last entry is the duplicated
+# bootstrap slot), rewards=[0, 1, 0], gamma=0.99.
+#   terminated: target[0] = gamma*v1 + gamma*lambda*(r1 - v1) = 0.9405 + 0.0495*lambda
+#               target[1] = r1 = 1.0
+#   truncated:  target[0] = 0.9405 + gamma*lambda*delta_1 = 0.9405 + 0.99*lambda*0.9905
+#               target[1] = r1 + gamma*v_extra = 1.9405
+@pytest.mark.parametrize(
+    "lambda_,is_terminated,expected",
+    [
+        (0.0, True, [0.9405, 1.0]),
+        (0.5, True, [0.9405 + 0.99 * 0.5 * 0.05, 1.0]),
+        (1.0, True, [0.99, 1.0]),
+        (0.0, False, [0.9405, 1.9405]),
+        (0.5, False, [0.9405 + 0.99 * 0.5 * 0.9905, 1.9405]),
+        (1.0, False, [0.9405 + 0.99 * 0.9905, 1.9405]),
+    ],
+)
+def test_single_episode_targets(lambda_, is_terminated, expected):
+    """Single episode: terminal reward propagates; truncation keeps the bootstrap."""
+    out = _targets(
+        per_ep_values=[[0.0, 0.95, 0.95]],
+        per_ep_rewards=[[0.0, 1.0]],
+        terminated=[is_terminated],
+        truncated=[not is_terminated],
+        gamma=0.99,
+        lambda_=lambda_,
+    )
+    np.testing.assert_allclose(out[:2], expected, atol=1e-4)
+
+
+@pytest.mark.parametrize(
+    "ep1_term,ep2_term",
+    [(True, True), (True, False), (False, True), (False, False)],
+)
+def test_no_cross_episode_leak(ep1_term, ep2_term):
+    """At lambda=1, episode 1's targets must not depend on episode 2."""
+    pair = _targets(
+        per_ep_values=[[0.0, 0.95, 0.95], [0.0, 0.95, 0.95]],
+        per_ep_rewards=[[0.0, 1.0], [0.0, 1.0]],
+        terminated=[ep1_term, ep2_term],
+        truncated=[not ep1_term, not ep2_term],
+        gamma=0.99,
+        lambda_=1.0,
+    )
+    solo = _targets(
+        per_ep_values=[[0.0, 0.95, 0.95]],
+        per_ep_rewards=[[0.0, 1.0]],
+        terminated=[ep1_term],
+        truncated=[not ep1_term],
+        gamma=0.99,
+        lambda_=1.0,
+    )
+    np.testing.assert_allclose(pair[:2], solo[:2], atol=1e-4)
+
+
+# 2x2 deterministic FrozenLake used for the end-to-end convergence check below:
+#   row 0: S F   states 0, 1
+#   row 1: H G   states 2, 3
+# Reward 1.0 at G; episodes terminate at H or G.
+# Optimal policy from S: right (to F=1), then down (to G=3, reward=1).
+# Bellman closed form with gamma=0.99 on the non-terminal states:
+#   V(F) = 1 + gamma * V(G_terminal) = 1.0
+#   V(S) = 0 + gamma * V(F)          = 0.99
+# V on the terminal states (H, G) is never targeted during training and is
+# therefore left out of the comparison.
+_FROZEN_LAKE_2X2_CFG = {"desc": ["SF", "HG"], "is_slippery": False}
+_TRUE_V_NON_TERMINAL = np.array([0.99, 1.0], dtype=np.float32)
+
+
+def _train_and_get_state_values(gae_lambda: float, num_iters: int, seed: int):
+    """Train PPO on 2x2 FrozenLake and return V for all 4 states."""
+    config = (
+        PPOConfig()
+        .environment("FrozenLake-v1", env_config=_FROZEN_LAKE_2X2_CFG)
+        .env_runners(
+            num_env_runners=0,
+            num_envs_per_env_runner=4,
+            # Discrete obs -> one-hot for the FC encoder.
+            env_to_module_connector=(lambda env, spaces, device: FlattenObservations()),
+        )
+        .training(
+            gamma=0.99,
+            lambda_=gae_lambda,
+            lr=3e-3,
+            train_batch_size=256,
+            num_epochs=10,
+            minibatch_size=64,
+            # Up-weight the value loss and disable entropy so V converges
+            # quickly and the test stays short.
+            vf_loss_coeff=1.0,
+            entropy_coeff=0.0,
+        )
+        .rl_module(
+            model_config=DefaultModelConfig(
+                fcnet_hiddens=[32],
+                fcnet_activation="tanh",
+            ),
+        )
+        .debugging(seed=seed)
+    )
+    algo = config.build_algo()
+
+    for _ in range(num_iters):
+        algo.train()
+    # `algo.get_module(...)` returns the EnvRunner's inference-only
+    # module (no critic). Reach into the Learner's module to call
+    # compute_values.
+    learner_module = algo.learner_group._learner.module[DEFAULT_MODULE_ID]
+    obs = np.eye(4, dtype=np.float32)  # one-hot for each of the 4 states
+    with torch.no_grad():
+        return (
+            learner_module.compute_values({Columns.OBS: convert_to_torch_tensor(obs)})
+            .detach()
+            .cpu()
+            .numpy()
+        )
+
+
+def test_value_function_converges_across_gae_lambda():
+    """
+    End-to-end check that PPO trains a consistent V across `gae_lambda`.
+    Different lambda values should converge to the same fixed-point V.
+    """
+    v_by_lambda = {
+        lam: _train_and_get_state_values(gae_lambda=lam, num_iters=40, seed=42)
+        for lam in [0.0, 0.9, 1.0]
+    }
+    # 1) V on the visited (non-terminal) states matches the analytic V.
+    for lam, v in v_by_lambda.items():
+        np.testing.assert_allclose(
+            v[:2],
+            _TRUE_V_NON_TERMINAL,
+            atol=0.05,
+            err_msg=(
+                f"V on non-terminal states diverged from analytic V "
+                f"for lambda={lam}: got {v[:2]}, expected "
+                f"{_TRUE_V_NON_TERMINAL}"
+            ),
+        )
+    # 2) V across the three lambdas must converge together.
+    assert np.ptp([v[:2] for v in v_by_lambda.values()], axis=0).max() < 0.05
+
+
+if __name__ == "__main__":
+    import sys
+
+    sys.exit(pytest.main(["-v", __file__]))

rllib/algorithms/ppo/torch/default_ppo_torch_rl_module.py

@@ -0,0 +1,73 @@
+from typing import Any, Dict, Optional
+
+from ray.rllib.algorithms.ppo.default_ppo_rl_module import DefaultPPORLModule
+from ray.rllib.algorithms.ppo.ppo_catalog import PPOCatalog
+from ray.rllib.core.columns import Columns
+from ray.rllib.core.models.base import ACTOR, CRITIC, ENCODER_OUT
+from ray.rllib.core.rl_module.apis.value_function_api import ValueFunctionAPI
+from ray.rllib.core.rl_module.rl_module import RLModule
+from ray.rllib.core.rl_module.torch import TorchRLModule
+from ray.rllib.utils.annotations import override
+from ray.rllib.utils.framework import try_import_torch
+from ray.rllib.utils.typing import TensorType
+from ray.util.annotations import DeveloperAPI
+
+torch, nn = try_import_torch()
+
+
+@DeveloperAPI
+class DefaultPPOTorchRLModule(TorchRLModule, DefaultPPORLModule):
+    def __init__(self, *args, **kwargs):
+        catalog_class = kwargs.pop("catalog_class", None)
+        if catalog_class is None:
+            catalog_class = PPOCatalog
+        super().__init__(*args, **kwargs, catalog_class=catalog_class)
+
+    @override(RLModule)
+    def _forward(self, batch: Dict[str, Any], **kwargs) -> Dict[str, Any]:
+        """Default forward pass (used for inference and exploration)."""
+        output = {}
+        # Encoder forward pass.
+        encoder_outs = self.encoder(batch)
+        # Stateful encoder?
+        if Columns.STATE_OUT in encoder_outs:
+            output[Columns.STATE_OUT] = encoder_outs[Columns.STATE_OUT]
+        # Pi head.
+        output[Columns.ACTION_DIST_INPUTS] = self.pi(encoder_outs[ENCODER_OUT][ACTOR])
+        return output
+
+    @override(RLModule)
+    def _forward_train(self, batch: Dict[str, Any], **kwargs) -> Dict[str, Any]:
+        """Train forward pass (keep embeddings for possible shared value func. call)."""
+        output = {}
+        encoder_outs = self.encoder(batch)
+        output[Columns.EMBEDDINGS] = encoder_outs[ENCODER_OUT][CRITIC]
+        if Columns.STATE_OUT in encoder_outs:
+            output[Columns.STATE_OUT] = encoder_outs[Columns.STATE_OUT]
+        output[Columns.ACTION_DIST_INPUTS] = self.pi(encoder_outs[ENCODER_OUT][ACTOR])
+        return output
+
+    @override(ValueFunctionAPI)
+    def compute_values(
+        self,
+        batch: Dict[str, Any],
+        embeddings: Optional[Any] = None,
+    ) -> TensorType:
+        if embeddings is None:
+            # Separate vf-encoder.
+            if hasattr(self.encoder, "critic_encoder"):
+                batch_ = batch
+                if self.is_stateful():
+                    # The recurrent encoders expect a `(state_in, h)`  key in the
+                    # input dict while the key returned is `(state_in, critic, h)`.
+                    batch_ = batch.copy()
+                    batch_[Columns.STATE_IN] = batch[Columns.STATE_IN][CRITIC]
+                embeddings = self.encoder.critic_encoder(batch_)[ENCODER_OUT]
+            # Shared encoder.
+            else:
+                embeddings = self.encoder(batch)[ENCODER_OUT][CRITIC]
+
+        # Value head.
+        vf_out = self.vf(embeddings)
+        # Squeeze out last dimension (single node value head).
+        return vf_out.squeeze(-1)

rllib/algorithms/ppo/torch/ppo_torch_learner.py

@@ -0,0 +1,173 @@
+import logging
+from typing import Any, Dict
+
+import numpy as np
+
+from ray.rllib.algorithms.ppo.ppo import (
+    LEARNER_RESULTS_CURR_KL_COEFF_KEY,
+    LEARNER_RESULTS_KL_KEY,
+    LEARNER_RESULTS_VF_EXPLAINED_VAR_KEY,
+    LEARNER_RESULTS_VF_LOSS_UNCLIPPED_KEY,
+    PPOConfig,
+)
+from ray.rllib.algorithms.ppo.ppo_learner import PPOLearner
+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.learner.torch.torch_learner import TorchLearner
+from ray.rllib.evaluation.postprocessing import Postprocessing
+from ray.rllib.utils.annotations import override
+from ray.rllib.utils.framework import try_import_torch
+from ray.rllib.utils.torch_utils import explained_variance
+from ray.rllib.utils.typing import ModuleID, TensorType
+
+torch, nn = try_import_torch()
+
+logger = logging.getLogger(__name__)
+
+
+class PPOTorchLearner(PPOLearner, TorchLearner):
+    """Implements torch-specific PPO loss logic on top of PPOLearner.
+
+    This class implements the ppo loss under `self.compute_loss_for_module()`.
+    """
+
+    @override(TorchLearner)
+    def compute_loss_for_module(
+        self,
+        *,
+        module_id: ModuleID,
+        config: PPOConfig,
+        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 (artificial) timestep to each episode to
+        # simplify the actual computation.
+        if Columns.LOSS_MASK in batch:
+            mask = batch[Columns.LOSS_MASK]
+            num_valid = torch.sum(mask)
+
+            def possibly_masked_mean(data_):
+                return torch.sum(data_[mask]) / num_valid
+
+        else:
+            possibly_masked_mean = torch.mean
+
+        action_dist_class_train = module.get_train_action_dist_cls()
+        action_dist_class_exploration = module.get_exploration_action_dist_cls()
+
+        curr_action_dist = action_dist_class_train.from_logits(
+            fwd_out[Columns.ACTION_DIST_INPUTS]
+        )
+        # TODO (sven): We should ideally do this in the LearnerConnector (separation of
+        #  concerns: Only do things on the EnvRunners that are required for computing
+        #  actions, do NOT do anything on the EnvRunners that's only required for a
+        #   training update).
+        prev_action_dist = action_dist_class_exploration.from_logits(
+            batch[Columns.ACTION_DIST_INPUTS]
+        )
+
+        logp_ratio = torch.exp(
+            curr_action_dist.logp(batch[Columns.ACTIONS]) - batch[Columns.ACTION_LOGP]
+        )
+
+        # Only calculate kl loss if necessary (kl-coeff > 0.0).
+        if config.use_kl_loss:
+            action_kl = prev_action_dist.kl(curr_action_dist)
+            mean_kl_loss = possibly_masked_mean(action_kl)
+        else:
+            mean_kl_loss = torch.tensor(0.0, device=logp_ratio.device)
+
+        curr_entropy = curr_action_dist.entropy()
+        mean_entropy = possibly_masked_mean(curr_entropy)
+
+        surrogate_loss = torch.min(
+            batch[Postprocessing.ADVANTAGES] * logp_ratio,
+            batch[Postprocessing.ADVANTAGES]
+            * torch.clamp(logp_ratio, 1 - config.clip_param, 1 + config.clip_param),
+        )
+
+        # Compute a value function loss.
+        if config.use_critic:
+            value_fn_out = module.compute_values(
+                batch, embeddings=fwd_out.get(Columns.EMBEDDINGS)
+            )
+            vf_loss = torch.pow(value_fn_out - batch[Postprocessing.VALUE_TARGETS], 2.0)
+            vf_loss_clipped = torch.clamp(vf_loss, 0, config.vf_clip_param)
+            mean_vf_loss = possibly_masked_mean(vf_loss_clipped)
+            mean_vf_unclipped_loss = possibly_masked_mean(vf_loss)
+        # Ignore the value function -> Set all to 0.0.
+        else:
+            z = torch.tensor(0.0, device=surrogate_loss.device)
+            value_fn_out = mean_vf_unclipped_loss = vf_loss_clipped = mean_vf_loss = z
+
+        total_loss = possibly_masked_mean(
+            -surrogate_loss
+            + config.vf_loss_coeff * vf_loss_clipped
+            - (
+                self.entropy_coeff_schedulers_per_module[module_id].get_current_value()
+                * curr_entropy
+            )
+        )
+
+        # Add mean_kl_loss (already processed through `possibly_masked_mean`),
+        # if necessary.
+        if config.use_kl_loss:
+            total_loss += self.curr_kl_coeffs_per_module[module_id] * mean_kl_loss
+
+        # Log important loss stats.
+        self.metrics.log_dict(
+            {
+                POLICY_LOSS_KEY: -possibly_masked_mean(surrogate_loss),
+                VF_LOSS_KEY: mean_vf_loss,
+                LEARNER_RESULTS_VF_LOSS_UNCLIPPED_KEY: mean_vf_unclipped_loss,
+                LEARNER_RESULTS_VF_EXPLAINED_VAR_KEY: explained_variance(
+                    batch[Postprocessing.VALUE_TARGETS], value_fn_out
+                ),
+                ENTROPY_KEY: mean_entropy,
+                LEARNER_RESULTS_KL_KEY: mean_kl_loss,
+            },
+            key=module_id,
+            window=1,  # <- single items (should not be mean/ema-reduced over time).
+        )
+        # Return the total loss.
+        return total_loss
+
+    @override(PPOLearner)
+    def _update_module_kl_coeff(
+        self,
+        *,
+        module_id: ModuleID,
+        config: PPOConfig,
+        kl_loss: float,
+    ) -> None:
+        if np.isnan(kl_loss):
+            logger.warning(
+                f"KL divergence for Module {module_id} is non-finite, this "
+                "will likely destabilize your model and the training "
+                "process. Action(s) in a specific state have near-zero "
+                "probability. This can happen naturally in deterministic "
+                "environments where the optimal policy has zero mass for a "
+                "specific action. To fix this issue, consider setting "
+                "`kl_coeff` to 0.0 or increasing `entropy_coeff` in your "
+                "config."
+            )
+
+        # Update the KL coefficient.
+        curr_var = self.curr_kl_coeffs_per_module[module_id]
+        if kl_loss > 2.0 * config.kl_target:
+            # TODO (Kourosh) why not 2?
+            curr_var.data *= 1.5
+        elif kl_loss < 0.5 * config.kl_target:
+            curr_var.data *= 0.5
+
+        # Log the updated KL-coeff value.
+        self.metrics.log_value(
+            (module_id, LEARNER_RESULTS_CURR_KL_COEFF_KEY),
+            curr_var.item(),
+            window=1,
+        )

rllib/algorithms/ppo/torch/ppo_torch_rl_module.py

@@ -0,0 +1,13 @@
+# Backward compat import.
+from ray.rllib.algorithms.ppo.torch.default_ppo_torch_rl_module import (  # noqa
+    DefaultPPOTorchRLModule as PPOTorchRLModule,
+)
+from ray._common.deprecation import deprecation_warning
+
+
+deprecation_warning(
+    old="ray.rllib.algorithms.ppo.torch.ppo_torch_rl_module.PPOTorchRLModule",
+    new="ray.rllib.algorithms.ppo.torch.default_ppo_torch_rl_module."
+    "DefaultPPOTorchRLModule",
+    error=False,
+)