chore: import upstream snapshot with attribution

rllib/examples/learners/classes/custom_ppo_loss_fn_learner.py

@@ -0,0 +1,54 @@
+from typing import Any, Dict
+
+from ray.rllib.algorithms.ppo import PPOConfig
+from ray.rllib.algorithms.ppo.torch.ppo_torch_learner import PPOTorchLearner
+from ray.rllib.utils.annotations import override
+from ray.rllib.utils.framework import try_import_torch
+from ray.rllib.utils.typing import ModuleID, TensorType
+
+torch, _ = try_import_torch()
+
+
+class PPOTorchLearnerWithWeightRegularizerLoss(PPOTorchLearner):
+    """A custom PPO torch learner adding a weight regularizer term to the loss.
+
+    We compute a naive regularizer term averaging over all parameters of the RLModule
+    and add this mean value (multiplied by the regularizer coefficient) to the base PPO
+    loss.
+    The experiment shows that even with a large learning rate, our custom Learner is
+    still able to learn properly as it's forced to keep the weights small.
+    """
+
+    @override(PPOTorchLearner)
+    def compute_loss_for_module(
+        self,
+        *,
+        module_id: ModuleID,
+        config: PPOConfig,
+        batch: Dict[str, Any],
+        fwd_out: Dict[str, TensorType],
+    ) -> TensorType:
+
+        base_total_loss = super().compute_loss_for_module(
+            module_id=module_id,
+            config=config,
+            batch=batch,
+            fwd_out=fwd_out,
+        )
+
+        # Compute the mean of all the RLModule's weights.
+        parameters = self.get_parameters(self.module[module_id])
+        mean_weight = torch.mean(torch.stack([w.mean() for w in parameters]))
+
+        self.metrics.log_value(
+            key=(module_id, "mean_weight"),
+            value=mean_weight,
+            window=1,
+        )
+
+        total_loss = (
+            base_total_loss
+            + config.learner_config_dict["regularizer_coeff"] * mean_weight
+        )
+
+        return total_loss

rllib/examples/learners/classes/intrinsic_curiosity_learners.py

@@ -0,0 +1,162 @@
+from typing import Any, List, Optional
+
+import gymnasium as gym
+import torch
+
+from ray.rllib.algorithms.dqn.torch.dqn_torch_learner import DQNTorchLearner
+from ray.rllib.algorithms.ppo.torch.ppo_torch_learner import PPOTorchLearner
+from ray.rllib.connectors.common.add_observations_from_episodes_to_batch import (
+    AddObservationsFromEpisodesToBatch,
+)
+from ray.rllib.connectors.common.numpy_to_tensor import NumpyToTensor
+from ray.rllib.connectors.connector_v2 import ConnectorV2
+from ray.rllib.connectors.learner.add_next_observations_from_episodes_to_train_batch import (  # noqa
+    AddNextObservationsFromEpisodesToTrainBatch,
+)
+from ray.rllib.core import DEFAULT_MODULE_ID, Columns
+from ray.rllib.core.learner.torch.torch_learner import TorchLearner
+from ray.rllib.core.rl_module.rl_module import RLModule
+from ray.rllib.utils.typing import EpisodeType
+
+ICM_MODULE_ID = "_intrinsic_curiosity_model"
+
+
+class DQNTorchLearnerWithCuriosity(DQNTorchLearner):
+    def build(self) -> None:
+        super().build()
+        add_intrinsic_curiosity_connectors(self)
+
+
+class PPOTorchLearnerWithCuriosity(PPOTorchLearner):
+    def build(self) -> None:
+        super().build()
+        add_intrinsic_curiosity_connectors(self)
+
+
+def add_intrinsic_curiosity_connectors(torch_learner: TorchLearner) -> None:
+    """Adds two connector pieces to the Learner pipeline, needed for ICM training.
+
+    - The `AddNextObservationsFromEpisodesToTrainBatch` connector makes sure the train
+    batch contains the NEXT_OBS for ICM's forward- and inverse dynamics net training.
+    - The `IntrinsicCuriosityModelConnector` piece computes intrinsic rewards from the
+    ICM and adds the results to the extrinsic reward of the main module's train batch.
+
+    Args:
+        torch_learner: The TorchLearner, to whose Learner pipeline the two ICM connector
+            pieces should be added.
+    """
+    learner_config_dict = torch_learner.config.learner_config_dict
+
+    # Assert, we are only training one policy (RLModule) and we have the ICM
+    # in our MultiRLModule.
+    assert (
+        len(torch_learner.module) == 2
+        and DEFAULT_MODULE_ID in torch_learner.module
+        and ICM_MODULE_ID in torch_learner.module
+    )
+
+    # Make sure both curiosity loss settings are explicitly set in the
+    # `learner_config_dict`.
+    if (
+        "forward_loss_weight" not in learner_config_dict
+        or "intrinsic_reward_coeff" not in learner_config_dict
+    ):
+        raise KeyError(
+            "When using the IntrinsicCuriosityTorchLearner, both `forward_loss_weight` "
+            " and `intrinsic_reward_coeff` must be part of your config's "
+            "`learner_config_dict`! Add these values through: `config.training("
+            "learner_config_dict={'forward_loss_weight': .., 'intrinsic_reward_coeff': "
+            "..})`."
+        )
+
+    if torch_learner.config.add_default_connectors_to_learner_pipeline:
+        # Prepend a "add-NEXT_OBS-from-episodes-to-train-batch" connector piece
+        # (right after the corresponding "add-OBS-..." default piece).
+        torch_learner._learner_connector.insert_after(
+            AddObservationsFromEpisodesToBatch,
+            AddNextObservationsFromEpisodesToTrainBatch(),
+        )
+        # Append the ICM connector, computing intrinsic rewards and adding these to
+        # the main model's extrinsic rewards.
+        torch_learner._learner_connector.insert_after(
+            NumpyToTensor,
+            IntrinsicCuriosityModelConnector(
+                intrinsic_reward_coeff=(
+                    torch_learner.config.learner_config_dict["intrinsic_reward_coeff"]
+                )
+            ),
+        )
+
+
+class IntrinsicCuriosityModelConnector(ConnectorV2):
+    """Learner ConnectorV2 piece to compute intrinsic rewards based on an ICM.
+
+    For more details, see here:
+    [1] Curiosity-driven Exploration by Self-supervised Prediction
+    Pathak, Agrawal, Efros, and Darrell - UC Berkeley - ICML 2017.
+    https://arxiv.org/pdf/1705.05363.pdf
+
+    This connector piece:
+    - requires two RLModules to be present in the MultiRLModule:
+    DEFAULT_MODULE_ID (the policy model to be trained) and ICM_MODULE_ID (the instrinsic
+    curiosity architecture).
+    - must be located toward the end of to your Learner pipeline (after the
+    `NumpyToTensor` piece) in order to perform a forward pass on the ICM model with the
+    readily compiled batch and a following forward-loss computation to get the intrinsi
+    rewards.
+    - these intrinsic rewards will then be added to the (extrinsic) rewards in the main
+    model's train batch.
+    """
+
+    def __init__(
+        self,
+        input_observation_space: Optional[gym.Space] = None,
+        input_action_space: Optional[gym.Space] = None,
+        *,
+        intrinsic_reward_coeff: float,
+        **kwargs,
+    ):
+        """Initializes a CountBasedCuriosity instance.
+
+        Args:
+            intrinsic_reward_coeff: The weight with which to multiply the intrinsic
+                reward before adding it to the extrinsic rewards of the main model.
+        """
+        super().__init__(input_observation_space, input_action_space)
+
+        self.intrinsic_reward_coeff = intrinsic_reward_coeff
+
+    def __call__(
+        self,
+        *,
+        rl_module: RLModule,
+        batch: Any,
+        episodes: List[EpisodeType],
+        explore: Optional[bool] = None,
+        shared_data: Optional[dict] = None,
+        **kwargs,
+    ) -> Any:
+        # Assert that the batch is ready.
+        assert DEFAULT_MODULE_ID in batch and ICM_MODULE_ID not in batch
+        assert (
+            Columns.OBS in batch[DEFAULT_MODULE_ID]
+            and Columns.NEXT_OBS in batch[DEFAULT_MODULE_ID]
+        )
+        # TODO (sven): We are performing two forward passes per update right now.
+        #  Once here in the connector (w/o grad) to just get the intrinsic rewards
+        #  and once in the learner to actually compute the ICM loss and update the ICM.
+        #  Maybe we can save one of these, but this would currently harm the DDP-setup
+        #  for multi-GPU training.
+        with torch.no_grad():
+            # Perform ICM forward pass.
+            fwd_out = rl_module[ICM_MODULE_ID].forward_train(batch[DEFAULT_MODULE_ID])
+
+        # Add the intrinsic rewards to the main module's extrinsic rewards.
+        batch[DEFAULT_MODULE_ID][Columns.REWARDS] += (
+            self.intrinsic_reward_coeff * fwd_out[Columns.INTRINSIC_REWARDS]
+        )
+
+        # Duplicate the batch such that the ICM also has data to learn on.
+        batch[ICM_MODULE_ID] = batch[DEFAULT_MODULE_ID]
+
+        return batch

rllib/examples/learners/classes/separate_vf_lr_and_optimizer_learner.py

@@ -0,0 +1,83 @@
+from ray.rllib.algorithms.algorithm_config import AlgorithmConfig
+from ray.rllib.algorithms.ppo.torch.ppo_torch_learner import PPOTorchLearner
+from ray.rllib.core.learner.torch.torch_learner import TorchLearner
+from ray.rllib.utils.annotations import override
+from ray.rllib.utils.framework import try_import_torch
+from ray.rllib.utils.typing import ModuleID
+
+torch, _ = try_import_torch()
+
+
+class PPOTorchLearnerWithSeparateVfOptimizer(PPOTorchLearner):
+    """A custom PPO torch learner with 2 optimizers, for policy and value function.
+
+    Overrides the Learner's standard `configure_optimizers_for_module()` method to
+    register the additional vf optimizer.
+
+    The standard PPOLearner only uses a single optimizer (and single learning rate) to
+    update the model, regardless of whether the value function network
+    is separate from the policy network or whether they have shared components.
+
+    We may leave the loss function of PPO completely untouched. It already returns a
+    sum of policy loss and vf loss (and entropy loss), and thus - given that the neural
+    networks used to compute each of these terms are separate and don't share any
+    components - gradients are computed separately per neural network (policy vs vf)
+    and applied separately through the two optimizers.
+    """
+
+    @override(TorchLearner)
+    def configure_optimizers_for_module(
+        self,
+        module_id: ModuleID,
+        config: "AlgorithmConfig" = None,
+    ) -> None:
+        """Registers 2 optimizers for the given ModuleID with this Learner."""
+
+        # Make sure the RLModule has the correct properties.
+        module = self.module[module_id]
+        # TODO (sven): We should move this into a new `ValueFunction` API, which
+        #  should have a `get_value_function_params` method. This way, any custom
+        #  RLModule that implements this API can be used here, not just the standard
+        #  PPO one.
+        assert (
+            hasattr(module, "pi")
+            and hasattr(module, "vf")
+            and hasattr(module, "encoder")
+            and hasattr(module.encoder, "actor_encoder")
+            and hasattr(module.encoder, "critic_encoder")
+        )
+        assert config.model_config["vf_share_layers"] is False
+
+        # Get all policy-related parameters from the RLModule.
+        pi_params = (
+            # Actor encoder and policy head.
+            self.get_parameters(self.module[module_id].encoder.actor_encoder)
+            + self.get_parameters(self.module[module_id].pi)
+        )
+        # Register the policy optimizer.
+        self.register_optimizer(
+            module_id=module_id,
+            optimizer_name="optim_for_pi",
+            optimizer=torch.optim.Adam(params=pi_params),
+            params=pi_params,
+            # For the policy learning rate, we use the "main" lr in the AlgorithmConfig.
+            lr_or_lr_schedule=config.lr,
+        )
+
+        # Get all value function-related parameters from the RLModule.
+        vf_params = (
+            # Critic encoder and value head.
+            self.get_parameters(self.module[module_id].encoder.critic_encoder)
+            + self.get_parameters(self.module[module_id].vf)
+        )
+        # Register the value function optimizer.
+        self.register_optimizer(
+            module_id=module_id,
+            optimizer_name="optim_for_vf",
+            optimizer=torch.optim.Adam(params=vf_params),
+            params=vf_params,
+            # For the value function learning rate, we use a user-provided custom
+            # setting in the `learner_config_dict` in the AlgorithmConfig. If this
+            # is not provided, use the same lr as for the policy optimizer.
+            lr_or_lr_schedule=config.learner_config_dict.get("lr_vf", config.lr),
+        )

rllib/examples/learners/classes/vpg_torch_learner.py

@@ -0,0 +1,71 @@
+from typing import TYPE_CHECKING, Any, Dict
+
+import numpy as np
+import torch
+
+from ray.rllib.connectors.learner import ComputeReturnsToGo
+from ray.rllib.core.columns import Columns
+from ray.rllib.core.learner.learner import Learner
+from ray.rllib.core.learner.torch.torch_learner import TorchLearner
+from ray.rllib.utils.annotations import override
+from ray.rllib.utils.numpy import convert_to_numpy
+from ray.rllib.utils.typing import ModuleID, TensorType
+
+if TYPE_CHECKING:
+    from ray.rllib.algorithms.algorithm_config import AlgorithmConfig
+
+
+class VPGTorchLearner(TorchLearner):
+    @override(TorchLearner)
+    def build(self) -> None:
+        super().build()
+
+        # Prepend the returns-to-go connector piece to have that information
+        # available in the train batch.
+        if self.config.add_default_connectors_to_learner_pipeline:
+            self._learner_connector.prepend(ComputeReturnsToGo(gamma=self.config.gamma))
+
+    @override(TorchLearner)
+    def compute_loss_for_module(
+        self,
+        *,
+        module_id: ModuleID,
+        config: "AlgorithmConfig",
+        batch: Dict[str, Any],
+        fwd_out: Dict[str, TensorType],
+    ) -> TensorType:
+        rl_module = self.module[module_id]
+
+        # Create the action distribution from the parameters output by the RLModule.
+        action_dist_inputs = fwd_out[Columns.ACTION_DIST_INPUTS]
+        action_dist_class = rl_module.get_train_action_dist_cls()
+        action_dist = action_dist_class.from_logits(action_dist_inputs)
+
+        # Compute log probabilities of the actions taken during sampling.
+        log_probs = action_dist.logp(batch[Columns.ACTIONS])
+
+        # Compute the policy gradient loss.
+        # Since we're not using a baseline, we use returns to go directly.
+        loss = -torch.mean(log_probs * batch[Columns.RETURNS_TO_GO])
+
+        # Just for exercise, log the average return to go per discrete action.
+        for act, ret_to_go in zip(batch[Columns.ACTIONS], batch[Columns.RETURNS_TO_GO]):
+            self.metrics.log_value(
+                key=(module_id, f"action_{act}_return_to_go_mean"),
+                value=ret_to_go,
+                reduce="mean",
+            )
+
+        return loss
+
+    @override(Learner)
+    def after_gradient_based_update(self, *, timesteps):
+        # This is to check if in the multi-gpu case, the weights across workers are
+        # the same. Only for testing purposes.
+        if self.config.report_mean_weights:
+            for module_id in self.module.keys():
+                parameters = convert_to_numpy(
+                    self.get_parameters(self.module[module_id])
+                )
+                mean_ws = np.mean([w.mean() for w in parameters])
+                self.metrics.log_value((module_id, "mean_weight"), mean_ws, window=1)

rllib/examples/learners/classes/vpg_torch_learner_shared_optimizer.py

@@ -0,0 +1,32 @@
+from ray.rllib.core.learner.torch.torch_learner import TorchLearner
+from ray.rllib.examples.learners.classes.vpg_torch_learner import VPGTorchLearner
+from ray.rllib.utils.annotations import override
+from ray.rllib.utils.framework import try_import_torch
+
+torch, _ = try_import_torch()
+
+
+class VPGTorchLearnerSharedOptimizer(VPGTorchLearner):
+    """
+    In order for a shared module to learn properly, a special, multi-agent Learner
+    has been set up. There is only one optimizer (used to train all submodules, e.g.
+    a shared encoder and n policy nets), in order to not destabilize learning. The
+    latter may happen if more than one optimizer would try to alternatingly optimize
+    the same shared submodule.
+    """
+
+    @override(TorchLearner)
+    def configure_optimizers(self) -> None:
+        # Get and aggregate parameters for every module
+        param_list = []
+        for m in self.module.values():
+            if self.rl_module_is_compatible(m):
+                param_list.extend(m.parameters())
+
+        self.register_optimizer(
+            optimizer_name="shared_optimizer",
+            optimizer=torch.optim.Adam(params=param_list),
+            params=param_list,
+            # For the policy learning rate, we use the "main" lr in the AlgorithmConfig.
+            lr_or_lr_schedule=self.config.lr,
+        )