chore: import upstream snapshot with attribution

rllib/examples/rl_modules/classes/__init__.py

@@ -0,0 +1,9 @@
+from ray.rllib.examples.rl_modules.classes.rock_paper_scissors_heuristic_rlm import (
+    AlwaysSameHeuristicRLM,
+    BeatLastHeuristicRLM,
+)
+
+__all__ = [
+    "AlwaysSameHeuristicRLM",
+    "BeatLastHeuristicRLM",
+]

rllib/examples/rl_modules/classes/action_masking_rlm.py

@@ -0,0 +1,211 @@
+from typing import Dict, Optional, Tuple, Union
+
+import gymnasium as gym
+
+from ray.rllib.algorithms.ppo.torch.ppo_torch_rl_module import PPOTorchRLModule
+from ray.rllib.core.columns import Columns
+from ray.rllib.core.rl_module.apis.value_function_api import ValueFunctionAPI
+from ray.rllib.core.rl_module.default_model_config import DefaultModelConfig
+from ray.rllib.core.rl_module.rl_module import RLModule
+from ray.rllib.utils.annotations import override
+from ray.rllib.utils.framework import try_import_torch
+from ray.rllib.utils.torch_utils import FLOAT_MIN
+from ray.rllib.utils.typing import TensorType
+
+torch, nn = try_import_torch()
+
+
+class ActionMaskingRLModule(RLModule):
+    """An RLModule that implements an action masking for safe RL.
+
+    This RLModule implements action masking to avoid unsafe/unwanted actions
+    dependent on the current state (observations). It does so by using an
+    environment generated action mask defining which actions are allowed and
+    which should be avoided. The action mask is extracted from the
+    environment's `gymnasium.spaces.dict.Dict` observation and applied after
+    the module's `forward`-pass to the action logits. The resulting action
+    logits prevent unsafe/unwanted actions to be sampled from the corresponding
+    action distribution.
+
+    Note, this RLModule is implemented for the `PPO` algorithm only. It is not
+    guaranteed to work with other algorithms. Furthermore, not that for this
+    module to work it requires an environment with a `gymnasium.spaces.dict.Dict`
+    observation space containing tow key, `"action_mask"` and `"observations"`.
+    """
+
+    @override(RLModule)
+    def __init__(
+        self,
+        *,
+        observation_space: Optional[gym.Space] = None,
+        action_space: Optional[gym.Space] = None,
+        inference_only: Optional[bool] = None,
+        learner_only: bool = False,
+        model_config: Optional[Union[dict, DefaultModelConfig]] = None,
+        catalog_class=None,
+        **kwargs,
+    ):
+        # If observation space is not of type `Dict` raise an error.
+        if not isinstance(observation_space, gym.spaces.dict.Dict):
+            raise ValueError(
+                "This RLModule requires the environment to provide a "
+                "`gym.spaces.Dict` observation space of the form: \n"
+                " {'action_mask': Box(0.0, 1.0, shape=(self.action_space.n,)),"
+                "  'observation_space': self.observation_space}"
+            )
+
+        # While the environment holds an observation space that contains, both,
+        # the action mask and the original observation space, the 'RLModule'
+        # receives only the `"observation"` element of the space, but not the
+        # action mask.
+        self.observation_space_with_mask = observation_space
+        self.observation_space = observation_space["observations"]
+
+        # Keeps track if observation specs have been checked already.
+        self._checked_observations = False
+
+        # The DefaultPPORLModule, in its constructor will build networks for the
+        # original observation space (i.e. without the action mask).
+        super().__init__(
+            observation_space=self.observation_space,
+            action_space=action_space,
+            inference_only=inference_only,
+            learner_only=learner_only,
+            model_config=model_config,
+            catalog_class=catalog_class,
+            **kwargs,
+        )
+
+
+class ActionMaskingTorchRLModule(ActionMaskingRLModule, PPOTorchRLModule):
+    @override(PPOTorchRLModule)
+    def setup(self):
+        super().setup()
+        # We need to reset here the observation space such that the
+        # super`s (`PPOTorchRLModule`) observation space is the
+        # original space (i.e. without the action mask) and `self`'s
+        # observation space contains the action mask.
+        self.observation_space = self.observation_space_with_mask
+
+    @override(PPOTorchRLModule)
+    def _forward_inference(
+        self, batch: Dict[str, TensorType], **kwargs
+    ) -> Dict[str, TensorType]:
+        # Preprocess the original batch to extract the action mask.
+        action_mask, batch = self._preprocess_batch(batch)
+        # Run the forward pass.
+        outs = super()._forward_inference(batch, **kwargs)
+        # Mask the action logits and return.
+        return self._mask_action_logits(outs, action_mask)
+
+    @override(PPOTorchRLModule)
+    def _forward_exploration(
+        self, batch: Dict[str, TensorType], **kwargs
+    ) -> Dict[str, TensorType]:
+        # Preprocess the original batch to extract the action mask.
+        action_mask, batch = self._preprocess_batch(batch)
+        # Run the forward pass.
+        outs = super()._forward_exploration(batch, **kwargs)
+        # Mask the action logits and return.
+        return self._mask_action_logits(outs, action_mask)
+
+    @override(PPOTorchRLModule)
+    def _forward_train(
+        self, batch: Dict[str, TensorType], **kwargs
+    ) -> Dict[str, TensorType]:
+        # Run the forward pass.
+        outs = super()._forward_train(batch, **kwargs)
+        # Mask the action logits and return.
+        return self._mask_action_logits(outs, batch["action_mask"])
+
+    @override(ValueFunctionAPI)
+    def compute_values(self, batch: Dict[str, TensorType], embeddings=None):
+        # Check, if the observations are still in `dict` form.
+        if isinstance(batch[Columns.OBS], dict):
+            # Preprocess the batch to extract the `observations` to `Columns.OBS`.
+            action_mask, batch = self._preprocess_batch(batch)
+            # NOTE: Because we manipulate the batch we need to add the `action_mask`
+            # to the batch to access them in `_forward_train`.
+            batch["action_mask"] = action_mask
+        # Call the super's method to compute values for GAE.
+        return super().compute_values(batch, embeddings)
+
+    def _preprocess_batch(
+        self, batch: Dict[str, TensorType], **kwargs
+    ) -> Tuple[TensorType, Dict[str, TensorType]]:
+        """Extracts observations and action mask from the batch
+
+        Args:
+            batch: A dictionary containing tensors (at least `Columns.OBS`)
+
+        Returns:
+            A tuple with the action mask tensor and the modified batch containing
+                the original observations.
+        """
+        # Check observation specs for action mask and observation keys.
+        self._check_batch(batch)
+
+        # Extract the available actions tensor from the observation.
+        action_mask = batch[Columns.OBS].pop("action_mask")
+
+        # Modify the batch for the `DefaultPPORLModule`'s `forward` method, i.e.
+        # pass only `"obs"` into the `forward` method.
+        batch[Columns.OBS] = batch[Columns.OBS].pop("observations")
+
+        # Return the extracted action mask and the modified batch.
+        return action_mask, batch
+
+    def _mask_action_logits(
+        self, batch: Dict[str, TensorType], action_mask: TensorType
+    ) -> Dict[str, TensorType]:
+        """Masks the action logits for the output of `forward` methods
+
+        Args:
+            batch: A dictionary containing tensors (at least action logits).
+            action_mask: A tensor containing the action mask for the current
+                observations.
+
+        Returns:
+            A modified batch with masked action logits for the action distribution
+            inputs.
+        """
+        # Convert action mask into an `[0.0][-inf]`-type mask.
+        inf_mask = torch.clamp(torch.log(action_mask), min=FLOAT_MIN)
+
+        # Mask the logits.
+        batch[Columns.ACTION_DIST_INPUTS] += inf_mask
+
+        # Return the batch with the masked action logits.
+        return batch
+
+    def _check_batch(self, batch: Dict[str, TensorType]) -> Optional[ValueError]:
+        """Assert that the batch includes action mask and observations.
+
+        Args:
+            batch: A dicitonary containing tensors (at least `Columns.OBS`) to be
+                checked.
+
+        Raises:
+            `ValueError` if the column `Columns.OBS`  does not contain observations
+                and action mask.
+        """
+        if not self._checked_observations:
+            if "action_mask" not in batch[Columns.OBS]:
+                raise ValueError(
+                    "No action mask found in observation. This `RLModule` requires "
+                    "the environment to provide observations that include an "
+                    "action mask (i.e. an observation space of the Dict space "
+                    "type that looks as follows: \n"
+                    "{'action_mask': Box(0.0, 1.0, shape=(self.action_space.n,)),"
+                    "'observations': self.observation_space}"
+                )
+            if "observations" not in batch[Columns.OBS]:
+                raise ValueError(
+                    "No observations found in observation. This 'RLModule` requires "
+                    "the environment to provide observations that include the original "
+                    "observations under a key `'observations'` in a dict (i.e. an "
+                    "observation space of the Dict space type that looks as follows: \n"
+                    "{'action_mask': Box(0.0, 1.0, shape=(self.action_space.n,)),"
+                    "'observations': <observation_space>}"
+                )
+            self._checked_observations = True

rllib/examples/rl_modules/classes/autoregressive_actions_rlm.py

@@ -0,0 +1,129 @@
+from typing import Dict
+
+import gymnasium as gym
+
+from ray.rllib.core import Columns
+from ray.rllib.core.distribution.torch.torch_distribution import (
+    TorchCategorical,
+    TorchDiagGaussian,
+    TorchMultiDistribution,
+)
+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.torch_rl_module import TorchRLModule
+from ray.rllib.utils.annotations import override
+from ray.rllib.utils.framework import try_import_torch
+from ray.rllib.utils.torch_utils import one_hot
+from ray.rllib.utils.typing import TensorType
+
+torch, nn = try_import_torch()
+
+
+class AutoregressiveActionsRLM(TorchRLModule, ValueFunctionAPI):
+    """An RLModule that uses an autoregressive action distribution.
+
+    Actions are sampled in two steps. The first (prior) action component is sampled from
+    a categorical distribution. Then, the second (posterior) action component is sampled
+    from a posterior distribution that depends on the first action component and the
+    other input data (observations).
+
+    Note, this RLModule works in combination with any algorithm, whose Learners require
+    the `ValueFunctionAPI`.
+    """
+
+    @override(RLModule)
+    def setup(self):
+        super().setup()
+
+        # Assert the action space is correct.
+        assert isinstance(self.action_space, gym.spaces.Tuple)
+        assert isinstance(self.action_space[0], gym.spaces.Discrete)
+        assert self.action_space[0].n == 3
+        assert isinstance(self.action_space[1], gym.spaces.Box)
+
+        self._prior_net = nn.Sequential(
+            nn.Linear(
+                in_features=self.observation_space.shape[0],
+                out_features=256,
+            ),
+            nn.Tanh(),
+            nn.Linear(in_features=256, out_features=self.action_space[0].n),
+        )
+
+        self._posterior_net = nn.Sequential(
+            nn.Linear(
+                in_features=self.observation_space.shape[0] + self.action_space[0].n,
+                out_features=256,
+            ),
+            nn.Tanh(),
+            nn.Linear(in_features=256, out_features=self.action_space[1].shape[0] * 2),
+        )
+
+        # Build the value function head.
+        self._value_net = nn.Sequential(
+            nn.Linear(
+                in_features=self.observation_space.shape[0],
+                out_features=256,
+            ),
+            nn.Tanh(),
+            nn.Linear(in_features=256, out_features=1),
+        )
+
+    @override(TorchRLModule)
+    def _forward_inference(self, batch: Dict[str, TensorType]) -> Dict[str, TensorType]:
+        return self._pi(batch[Columns.OBS], inference=True)
+
+    @override(TorchRLModule)
+    def _forward_exploration(
+        self, batch: Dict[str, TensorType], **kwargs
+    ) -> Dict[str, TensorType]:
+        return self._pi(batch[Columns.OBS], inference=False)
+
+    @override(TorchRLModule)
+    def _forward_train(self, batch: Dict[str, TensorType]) -> Dict[str, TensorType]:
+        return self._forward_exploration(batch)
+
+    @override(ValueFunctionAPI)
+    def compute_values(self, batch: Dict[str, TensorType], embeddings=None):
+        # Value function forward pass.
+        vf_out = self._value_net(batch[Columns.OBS])
+        # Squeeze out last dimension (single node value head).
+        return vf_out.squeeze(-1)
+
+    # __sphinx_begin__
+    def _pi(self, obs, inference: bool):
+        # Prior forward pass and sample a1.
+        prior_out = self._prior_net(obs)
+        dist_a1 = TorchCategorical.from_logits(prior_out)
+        if inference:
+            dist_a1 = dist_a1.to_deterministic()
+        a1 = dist_a1.sample()
+
+        # Posterior forward pass and sample a2.
+        posterior_batch = torch.cat(
+            [obs, one_hot(a1, self.action_space[0])],
+            dim=-1,
+        )
+        posterior_out = self._posterior_net(posterior_batch)
+        dist_a2 = TorchDiagGaussian.from_logits(posterior_out)
+        if inference:
+            dist_a2 = dist_a2.to_deterministic()
+        a2 = dist_a2.sample()
+        actions = (a1, a2)
+
+        # We need logp and distribution parameters for the loss.
+        return {
+            Columns.ACTION_LOGP: (
+                TorchMultiDistribution((dist_a1, dist_a2)).logp(actions)
+            ),
+            Columns.ACTION_DIST_INPUTS: torch.cat([prior_out, posterior_out], dim=-1),
+            Columns.ACTIONS: actions,
+        }
+        # __sphinx_end__
+
+    @override(TorchRLModule)
+    def get_inference_action_dist_cls(self):
+        return TorchMultiDistribution.get_partial_dist_cls(
+            child_distribution_cls_struct=(TorchCategorical, TorchDiagGaussian),
+            input_lens=(3, 2),
+        )

rllib/examples/rl_modules/classes/custom_action_distribution_rlm.py

@@ -0,0 +1,153 @@
+from typing import Any, Dict, Optional
+
+from ray.rllib.core.columns import Columns
+from ray.rllib.core.distribution.torch.torch_distribution import TorchCategorical
+from ray.rllib.core.rl_module.apis import ValueFunctionAPI
+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
+
+torch, nn = try_import_torch()
+
+
+def _make_categorical_with_temperature(temp):
+    """Helper function to create a new action distribution class.
+
+    The returned class takes a temperature parameter in its constructor with the default
+    value `temp`.
+
+    Args:
+        temp: The default temperature to use for the generated distribution class.
+    """
+
+    class TorchCategoricalWithTemp(TorchCategorical):
+        def __init__(self, logits=None, probs=None, temperature: float = temp):
+            """Initializes a TorchCategoricalWithTemp instance.
+
+            Args:
+                logits: Event log probabilities (non-normalized).
+                probs: The probabilities of each event.
+                temperature: In case of using logits, this parameter can be used to
+                    determine the sharpness of the distribution. i.e.
+                    ``probs = softmax(logits / temperature)``. The temperature must be
+                    strictly positive. A low value (e.g. 1e-10) will result in argmax
+                    sampling while a larger value will result in uniform sampling.
+            """
+            # Either divide logits or probs by the temperature.
+            assert (
+                temperature > 0.0
+            ), f"Temperature ({temperature}) must be strictly positive!"
+            if logits is not None:
+                logits /= temperature
+            else:
+                probs = torch.nn.functional.softmax(probs / temperature)
+            super().__init__(logits, probs)
+
+    return TorchCategoricalWithTemp
+
+
+class CustomActionDistributionRLModule(TorchRLModule, ValueFunctionAPI):
+    """A simple TorchRLModule with its own custom action distribution.
+
+    The distribution differs from the default one by an additional temperature
+    parameter applied on top of the Categorical base distribution. See the above
+    `TorchCategoricalWithTemp` class for details.
+
+    .. testcode::
+
+        import numpy as np
+        import gymnasium as gym
+
+        my_net = CustomActionDistributionRLModule(
+            observation_space=gym.spaces.Box(-1.0, 1.0, (4,), np.float32),
+            action_space=gym.spaces.Discrete(4),
+            model_config={"action_dist_temperature": 5.0},
+        )
+
+        B = 10
+        data = torch.from_numpy(
+            np.random.random_sample(size=(B, 4)).astype(np.float32)
+        )
+        # Expect a relatively high-temperature distribution.
+        # Set "action_dist_temperature" to small values << 1.0 to approximate greedy
+        # behavior (even when stochastically sampling from the distribution).
+        print(my_net.forward_exploration({"obs": data}))
+    """
+
+    @override(TorchRLModule)
+    def setup(self):
+        """Use this method to create all the model components that you require.
+
+        Feel free to access the following useful properties in this class:
+        - `self.model_config`: The config dict for this RLModule class,
+        which should contain flexible settings, for example: {"hiddens": [256, 256]}.
+        - `self.observation|action_space`: The observation and action space that
+        this RLModule is subject to. Note that the observation space might not be the
+        exact space from your env, but that it might have already gone through
+        preprocessing through a connector pipeline (for example, flattening,
+        frame-stacking, mean/std-filtering, etc..).
+        - `self.inference_only`: If True, this model should be built only for inference
+        purposes, in which case you may want to exclude any components that are not used
+        for computing actions, for example a value function branch.
+        """
+        input_dim = self.observation_space.shape[0]
+        hidden_dim = self.model_config.get("hidden_dim", 256)
+        output_dim = self.action_space.n
+
+        # Define simple encoder, and policy- and vf heads.
+        self._encoder = torch.nn.Sequential(
+            torch.nn.Linear(input_dim, hidden_dim),
+            torch.nn.ReLU(),
+        )
+        self._policy_net = torch.nn.Linear(hidden_dim, output_dim)
+        self._vf = nn.Linear(hidden_dim, 1)
+
+        # Plug in a custom action dist class.
+        # NOTE: If you need more granularity as to which distribution class is used by
+        # which forward method (`forward_inference`, `forward_exploration`,
+        # `forward_train`), override the RLModule methods
+        # `get_inference_action_dist_cls`, `get_exploration_action_dist_cls`, and
+        # `get_train_action_dist_cls`, and return
+        # your custom class(es) from these. In this case, leave `self.action_dist_cls`
+        # set to None, its default value.
+        self.action_dist_cls = _make_categorical_with_temperature(
+            self.model_config["action_dist_temperature"]
+        )
+
+    @override(TorchRLModule)
+    def _forward(self, batch, **kwargs):
+        # Compute the basic 1D feature tensor (inputs to policy- and value-heads).
+        _, logits = self._compute_embeddings_and_logits(batch)
+        # Return features and logits as ACTION_DIST_INPUTS (categorical distribution).
+        return {
+            Columns.ACTION_DIST_INPUTS: logits,
+        }
+
+    @override(TorchRLModule)
+    def _forward_train(self, batch, **kwargs):
+        # Compute the basic 1D feature tensor (inputs to policy- and value-heads).
+        embeddings, logits = self._compute_embeddings_and_logits(batch)
+        # Return features and logits as ACTION_DIST_INPUTS (categorical distribution).
+        return {
+            Columns.ACTION_DIST_INPUTS: logits,
+            Columns.EMBEDDINGS: embeddings,
+        }
+
+    # We implement this RLModule as a ValueFunctionAPI RLModule, so it can be used
+    # by value-based methods like PPO or IMPALA.
+    @override(ValueFunctionAPI)
+    def compute_values(
+        self,
+        batch: Dict[str, Any],
+        embeddings: Optional[Any] = None,
+    ) -> TensorType:
+        # Features not provided -> We need to compute them first.
+        if embeddings is None:
+            embeddings = self._encoder(batch[Columns.OBS])
+        return self._vf(embeddings).squeeze(-1)
+
+    def _compute_embeddings_and_logits(self, batch):
+        embeddings = self._encoder(batch[Columns.OBS])
+        logits = self._policy_net(embeddings)
+        return embeddings, logits

rllib/examples/rl_modules/classes/intrinsic_curiosity_model_rlm.py

@@ -0,0 +1,240 @@
+from typing import TYPE_CHECKING, Any, Dict
+
+import tree  # pip install dm_tree
+
+from ray.rllib.core.columns import Columns
+from ray.rllib.core.rl_module.apis import SelfSupervisedLossAPI
+from ray.rllib.core.rl_module.torch import TorchRLModule
+from ray.rllib.models.utils import get_activation_fn
+from ray.rllib.utils.annotations import override
+from ray.rllib.utils.framework import try_import_torch
+from ray.rllib.utils.torch_utils import one_hot
+from ray.rllib.utils.typing import ModuleID
+
+if TYPE_CHECKING:
+    from ray.rllib.algorithms.algorithm_config import AlgorithmConfig
+    from ray.rllib.core.learner.torch.torch_learner import TorchLearner
+
+torch, nn = try_import_torch()
+
+
+class IntrinsicCuriosityModel(TorchRLModule, SelfSupervisedLossAPI):
+    """An intrinsic curiosity model (ICM) as TorchRLModule for better exploration.
+
+    For more details, see:
+    [1] Curiosity-driven Exploration by Self-supervised Prediction
+    Pathak, Agrawal, Efros, and Darrell - UC Berkeley - ICML 2017.
+    https://arxiv.org/pdf/1705.05363.pdf
+
+    Learns a simplified model of the environment based on three networks:
+    1) Embedding observations into latent space ("feature" network).
+    2) Predicting the action, given two consecutive embedded observations
+    ("inverse" network).
+    3) Predicting the next embedded obs, given an obs and action
+    ("forward" network).
+
+    The less the agent is able to predict the actually observed next feature
+    vector, given obs and action (through the forwards network), the larger the
+    "intrinsic reward", which will be added to the extrinsic reward.
+    Therefore, if a state transition was unexpected, the agent becomes
+    "curious" and will further explore this transition leading to better
+    exploration in sparse rewards environments.
+
+    .. testcode::
+
+            import numpy as np
+            import gymnasium as gym
+            import torch
+
+            from ray.rllib.core import Columns
+            from ray.rllib.examples.rl_modules.classes.intrinsic_curiosity_model_rlm import (  # noqa
+                IntrinsicCuriosityModel
+            )
+
+            B = 10  # batch size
+            O = 4  # obs (1D) dim
+            A = 2  # num actions
+            f = 25  # feature dim
+
+            # Construct the RLModule.
+            icm_net = IntrinsicCuriosityModel(
+                observation_space=gym.spaces.Box(-1.0, 1.0, (O,), np.float32),
+                action_space=gym.spaces.Discrete(A),
+            )
+
+            # Create some dummy input.
+            obs = torch.from_numpy(
+                np.random.random_sample(size=(B, O)).astype(np.float32)
+            )
+            next_obs = torch.from_numpy(
+                np.random.random_sample(size=(B, O)).astype(np.float32)
+            )
+            actions = torch.from_numpy(
+                np.random.random_integers(0, A - 1, size=(B,))
+            )
+            input_dict = {
+                Columns.OBS: obs,
+                Columns.NEXT_OBS: next_obs,
+                Columns.ACTIONS: actions,
+            }
+
+            # Call `forward_train()` to get phi (feature vector from obs), next-phi
+            # (feature vector from next obs), and the intrinsic rewards (individual, per
+            # batch-item forward loss values).
+            print(icm_net.forward_train(input_dict))
+
+            # Print out the number of parameters.
+            num_all_params = sum(int(np.prod(p.size())) for p in icm_net.parameters())
+            print(f"num params = {num_all_params}")
+    """
+
+    @override(TorchRLModule)
+    def setup(self):
+        # Get the ICM achitecture settings from the `model_config` attribute:
+        cfg = self.model_config
+
+        feature_dim = cfg.get("feature_dim", 288)
+
+        # Build the feature model (encoder of observations to feature space).
+        layers = []
+        dense_layers = cfg.get("feature_net_hiddens", (256, 256))
+        # `in_size` is the observation space (assume a simple Box(1D)).
+        in_size = self.observation_space.shape[0]
+        for out_size in dense_layers:
+            layers.append(nn.Linear(in_size, out_size))
+            if cfg.get("feature_net_activation") not in [None, "linear"]:
+                layers.append(
+                    get_activation_fn(cfg["feature_net_activation"], "torch")()
+                )
+            in_size = out_size
+        # Last feature layer of n nodes (feature dimension).
+        layers.append(nn.Linear(in_size, feature_dim))
+        self._feature_net = nn.Sequential(*layers)
+
+        # Build the inverse model (predicting the action between two observations).
+        layers = []
+        dense_layers = cfg.get("inverse_net_hiddens", (256,))
+        # `in_size` is 2x the feature dim.
+        in_size = feature_dim * 2
+        for out_size in dense_layers:
+            layers.append(nn.Linear(in_size, out_size))
+            if cfg.get("inverse_net_activation") not in [None, "linear"]:
+                layers.append(
+                    get_activation_fn(cfg["inverse_net_activation"], "torch")()
+                )
+            in_size = out_size
+        # Last feature layer of n nodes (action space).
+        layers.append(nn.Linear(in_size, self.action_space.n))
+        self._inverse_net = nn.Sequential(*layers)
+
+        # Build the forward model (predicting the next observation from current one and
+        # action).
+        layers = []
+        dense_layers = cfg.get("forward_net_hiddens", (256,))
+        # `in_size` is the feature dim + action space (one-hot).
+        in_size = feature_dim + self.action_space.n
+        for out_size in dense_layers:
+            layers.append(nn.Linear(in_size, out_size))
+            if cfg.get("forward_net_activation") not in [None, "linear"]:
+                layers.append(
+                    get_activation_fn(cfg["forward_net_activation"], "torch")()
+                )
+            in_size = out_size
+        # Last feature layer of n nodes (feature dimension).
+        layers.append(nn.Linear(in_size, feature_dim))
+        self._forward_net = nn.Sequential(*layers)
+
+    @override(TorchRLModule)
+    def _forward_train(self, batch, **kwargs):
+        # Push both observations through feature net to get feature vectors (phis).
+        # We cat/batch them here for efficiency reasons (save one forward pass).
+        obs = tree.map_structure(
+            lambda obs, next_obs: torch.cat([obs, next_obs], dim=0),
+            batch[Columns.OBS],
+            batch[Columns.NEXT_OBS],
+        )
+        phis = self._feature_net(obs)
+        # Split again to yield 2 individual phi tensors.
+        phi, next_phi = torch.chunk(phis, 2)
+
+        # Predict next feature vector (next_phi) with forward model (using obs and
+        # actions).
+        predicted_next_phi = self._forward_net(
+            torch.cat(
+                [
+                    phi,
+                    one_hot(batch[Columns.ACTIONS].long(), self.action_space).float(),
+                ],
+                dim=-1,
+            )
+        )
+
+        # Forward loss term: Predicted phi - given phi and action - vs actually observed
+        # phi (square-root of L2 norm). Note that this is the intrinsic reward that
+        # will be used and the mean of this is the forward net loss.
+        forward_l2_norm_sqrt = 0.5 * torch.sum(
+            torch.pow(predicted_next_phi - next_phi, 2.0), dim=-1
+        )
+
+        output = {
+            Columns.INTRINSIC_REWARDS: forward_l2_norm_sqrt,
+            # Computed feature vectors (used to compute the losses later).
+            "phi": phi,
+            "next_phi": next_phi,
+        }
+
+        return output
+
+    @override(SelfSupervisedLossAPI)
+    def compute_self_supervised_loss(
+        self,
+        *,
+        learner: "TorchLearner",
+        module_id: ModuleID,
+        config: "AlgorithmConfig",
+        batch: Dict[str, Any],
+        fwd_out: Dict[str, Any],
+    ) -> Dict[str, Any]:
+        module = learner.module[module_id].unwrapped()
+
+        # Forward net loss.
+        forward_loss = torch.mean(fwd_out[Columns.INTRINSIC_REWARDS])
+
+        # Inverse loss term (predicted action that led from phi to phi' vs
+        # actual action taken).
+        dist_inputs = module._inverse_net(
+            torch.cat([fwd_out["phi"], fwd_out["next_phi"]], dim=-1)
+        )
+        action_dist = module.get_train_action_dist_cls().from_logits(dist_inputs)
+
+        # Neg log(p); p=probability of observed action given the inverse-NN
+        # predicted action distribution.
+        inverse_loss = -action_dist.logp(batch[Columns.ACTIONS])
+        inverse_loss = torch.mean(inverse_loss)
+
+        # Calculate the ICM loss.
+        total_loss = (
+            config.learner_config_dict["forward_loss_weight"] * forward_loss
+            + (1.0 - config.learner_config_dict["forward_loss_weight"]) * inverse_loss
+        )
+
+        learner.metrics.log_dict(
+            {
+                "mean_intrinsic_rewards": forward_loss,
+                "forward_loss": forward_loss,
+                "inverse_loss": inverse_loss,
+            },
+            key=module_id,
+            window=1,
+        )
+
+        return total_loss
+
+    # Inference and exploration not supported (this is a world-model that should only
+    # be used for training).
+    @override(TorchRLModule)
+    def _forward(self, batch, **kwargs):
+        raise NotImplementedError(
+            "`IntrinsicCuriosityModel` should only be used for training! "
+            "Only calls to `forward_train()` supported."
+        )

rllib/examples/rl_modules/classes/lstm_containing_rlm.py

@@ -0,0 +1,281 @@
+import abc
+from typing import Any, Dict, List, Optional, Tuple
+
+import numpy as np
+
+from ray.rllib.core.columns import Columns
+from ray.rllib.core.learner.utils import make_target_network
+from ray.rllib.core.rl_module.apis import (
+    TARGET_NETWORK_ACTION_DIST_INPUTS,
+    InferenceOnlyAPI,
+    TargetNetworkAPI,
+)
+from ray.rllib.core.rl_module.apis.value_function_api import ValueFunctionAPI
+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 NetworkType, TensorType
+
+torch, nn = try_import_torch()
+
+
+class LSTMContainingRLModule(TorchRLModule, ValueFunctionAPI):
+    """An example TorchRLModule that contains an LSTM layer.
+
+    .. testcode::
+
+        import numpy as np
+        import gymnasium as gym
+
+        B = 10  # batch size
+        T = 5  # seq len
+        e = 25  # embedding dim
+        CELL = 32  # LSTM cell size
+
+        # Construct the RLModule.
+        my_net = LSTMContainingRLModule(
+            observation_space=gym.spaces.Box(-1.0, 1.0, (e,), np.float32),
+            action_space=gym.spaces.Discrete(4),
+            model_config={"lstm_cell_size": CELL}
+        )
+
+        # Create some dummy input.
+        obs = torch.from_numpy(
+            np.random.random_sample(size=(B, T, e)
+        ).astype(np.float32))
+        state_in = my_net.get_initial_state()
+        # Repeat state_in across batch.
+        state_in = tree.map_structure(
+            lambda s: torch.from_numpy(s).unsqueeze(0).repeat(B, 1), state_in
+        )
+        input_dict = {
+            Columns.OBS: obs,
+            Columns.STATE_IN: state_in,
+        }
+
+        # Run through all 3 forward passes.
+        print(my_net.forward_inference(input_dict))
+        print(my_net.forward_exploration(input_dict))
+        print(my_net.forward_train(input_dict))
+
+        # Print out the number of parameters.
+        num_all_params = sum(int(np.prod(p.size())) for p in my_net.parameters())
+        print(f"num params = {num_all_params}")
+    """
+
+    @override(TorchRLModule)
+    def setup(self):
+        """Use this method to create all the model components that you require.
+
+        Feel free to access the following useful properties in this class:
+        - `self.model_config`: The config dict for this RLModule class,
+        which should contain flxeible settings, for example: {"hiddens": [256, 256]}.
+        - `self.observation|action_space`: The observation and action space that
+        this RLModule is subject to. Note that the observation space might not be the
+        exact space from your env, but that it might have already gone through
+        preprocessing through a connector pipeline (for example, flattening,
+        frame-stacking, mean/std-filtering, etc..).
+        """
+        # Assume a simple Box(1D) tensor as input shape.
+        in_size = self.observation_space.shape[0]
+
+        # Get the LSTM cell size from the `model_config` attribute:
+        self._lstm_cell_size = self.model_config.get("lstm_cell_size", 256)
+        self._lstm = nn.LSTM(in_size, self._lstm_cell_size, batch_first=True)
+        in_size = self._lstm_cell_size
+
+        # Build a sequential stack.
+        layers = []
+        # Get the dense layer pre-stack configuration from the same config dict.
+        dense_layers = self.model_config.get("dense_layers", [128, 128])
+        for out_size in dense_layers:
+            # Dense layer.
+            layers.append(nn.Linear(in_size, out_size))
+            # ReLU activation.
+            layers.append(nn.ReLU())
+            in_size = out_size
+
+        self._fc_net = nn.Sequential(*layers)
+
+        # Logits layer (no bias, no activation).
+        self._pi_head = nn.Linear(in_size, self.action_space.n)
+        # Single-node value layer.
+        self._values = nn.Linear(in_size, 1)
+
+    @override(TorchRLModule)
+    def get_initial_state(self) -> Any:
+        return {
+            "h": np.zeros(shape=(self._lstm_cell_size,), dtype=np.float32),
+            "c": np.zeros(shape=(self._lstm_cell_size,), dtype=np.float32),
+        }
+
+    @override(TorchRLModule)
+    def _forward(self, batch, **kwargs):
+        # Compute the basic 1D embedding tensor (inputs to policy- and value-heads).
+        embeddings, state_outs = self._compute_embeddings_and_state_outs(batch)
+        logits = self._pi_head(embeddings)
+
+        # Return logits as ACTION_DIST_INPUTS (categorical distribution).
+        # Note that the default `GetActions` connector piece (in the EnvRunner) will
+        # take care of argmax-"sampling" from the logits to yield the inference (greedy)
+        # action.
+        return {
+            Columns.ACTION_DIST_INPUTS: logits,
+            Columns.STATE_OUT: state_outs,
+        }
+
+    @override(TorchRLModule)
+    def _forward_train(self, batch, **kwargs):
+        # Same logic as _forward, but also return embeddings to be used by value
+        # function branch during training.
+        embeddings, state_outs = self._compute_embeddings_and_state_outs(batch)
+        logits = self._pi_head(embeddings)
+        return {
+            Columns.ACTION_DIST_INPUTS: logits,
+            Columns.STATE_OUT: state_outs,
+            Columns.EMBEDDINGS: embeddings,
+        }
+
+    # We implement this RLModule as a ValueFunctionAPI RLModule, so it can be used
+    # by value-based methods like PPO or IMPALA.
+    @override(ValueFunctionAPI)
+    def compute_values(
+        self, batch: Dict[str, Any], embeddings: Optional[Any] = None
+    ) -> TensorType:
+        if embeddings is None:
+            embeddings, _ = self._compute_embeddings_and_state_outs(batch)
+        values = self._values(embeddings).squeeze(-1)
+        return values
+
+    def _compute_embeddings_and_state_outs(self, batch):
+        obs = batch[Columns.OBS]
+        state_in = batch[Columns.STATE_IN]
+        h, c = state_in["h"], state_in["c"]
+        # Unsqueeze the layer dim (we only have 1 LSTM layer).
+        embeddings, (h, c) = self._lstm(obs, (h.unsqueeze(0), c.unsqueeze(0)))
+        # Push through our FC net.
+        embeddings = self._fc_net(embeddings)
+        # Squeeze the layer dim (we only have 1 LSTM layer).
+        return embeddings, {"h": h.squeeze(0), "c": c.squeeze(0)}
+
+
+class LSTMContainingRLModuleWithTargetNetwork(
+    LSTMContainingRLModule, TargetNetworkAPI, InferenceOnlyAPI, abc.ABC
+):
+    """LSTMContainingRLModule with TargetNetworkAPI support for use with APPO.
+
+    This class extends LSTMContainingRLModule to add target network functionality,
+    which is required by algorithms like APPO that use target networks for
+    importance sampling and policy updates.
+
+    .. testcode::
+
+        import numpy as np
+        import gymnasium as gym
+        import tree
+        import torch
+        from ray.rllib.core.columns import Columns
+
+        B = 10  # batch size
+        T = 5  # seq len
+        e = 25  # embedding dim
+        CELL = 32  # LSTM cell size
+
+        # Construct the RLModule with target network support.
+        my_net = LSTMContainingRLModuleWithTargetNetwork(
+            observation_space=gym.spaces.Box(-1.0, 1.0, (e,), np.float32),
+            action_space=gym.spaces.Discrete(4),
+            model_config={"lstm_cell_size": CELL}
+        )
+
+        # Create target networks (required for TargetNetworkAPI).
+        my_net.make_target_networks()
+
+        # Create some dummy input.
+        obs = torch.from_numpy(
+            np.random.random_sample(size=(B, T, e)
+        ).astype(np.float32))
+        state_in = my_net.get_initial_state()
+        # Repeat state_in across batch.
+        state_in = tree.map_structure(
+            lambda s: torch.from_numpy(s).unsqueeze(0).repeat(B, 1), state_in
+        )
+        input_dict = {
+            Columns.OBS: obs,
+            Columns.STATE_IN: state_in,
+        }
+
+        # Run through all forward passes including target network forward.
+        print("Forward inference:", my_net.forward_inference(input_dict))
+        print("Forward exploration:", my_net.forward_exploration(input_dict))
+        print("Forward train:", my_net.forward_train(input_dict))
+        print("Forward target:", my_net.forward_target(input_dict))
+
+        # Get target network pairs for synchronization.
+        target_pairs = my_net.get_target_network_pairs()
+        print(f"Number of target network pairs: {len(target_pairs)}")
+
+        # Print out the number of parameters.
+        num_all_params = sum(int(np.prod(p.size())) for p in my_net.parameters())
+        print(f"num params = {num_all_params}")
+
+    Example usage with APPO:
+
+    .. testcode::
+
+        from ray.rllib.algorithms.appo import APPOConfig
+        from ray.rllib.core.rl_module.rl_module import RLModuleSpec
+        from ray.rllib.examples.rl_modules.classes.lstm_containing_rlm import (
+            LSTMContainingRLModuleWithTargetNetwork,
+        )
+
+        config = (
+            APPOConfig()
+            .environment("CartPole-v1")
+            .rl_module(
+                rl_module_spec=RLModuleSpec(
+                    module_class=LSTMContainingRLModuleWithTargetNetwork,
+                    model_config={"lstm_cell_size": 256, "dense_layers": [128, 128]},
+                )
+            )
+        )
+    """
+
+    @override(TargetNetworkAPI)
+    def make_target_networks(self):
+        """Creates target networks for LSTM, FC net, and policy head."""
+        self._old_lstm = make_target_network(self._lstm)
+        self._old_fc_net = make_target_network(self._fc_net)
+        self._old_pi_head = make_target_network(self._pi_head)
+
+    @override(TargetNetworkAPI)
+    def get_target_network_pairs(self) -> List[Tuple[NetworkType, NetworkType]]:
+        """Returns pairs of (main_net, target_net) for target network updates."""
+        return [
+            (self._lstm, self._old_lstm),
+            (self._fc_net, self._old_fc_net),
+            (self._pi_head, self._old_pi_head),
+        ]
+
+    @override(TargetNetworkAPI)
+    def forward_target(self, batch: Dict[str, Any]) -> Dict[str, Any]:
+        """Forward pass through target networks to get action distribution inputs."""
+        # Compute embeddings using target networks (similar to _compute_embeddings_and_state_outs)
+        obs = batch[Columns.OBS]
+        state_in = batch[Columns.STATE_IN]
+        h, c = state_in["h"], state_in["c"]
+        # Unsqueeze the layer dim (we only have 1 LSTM layer) and forward through target LSTM
+        embeddings, (h, c) = self._old_lstm(obs, (h.unsqueeze(0), c.unsqueeze(0)))
+        # Push through target FC net
+        embeddings = self._old_fc_net(embeddings)
+        # Forward through target policy head to get action distribution inputs
+        old_action_dist_logits = self._old_pi_head(embeddings)
+
+        return {TARGET_NETWORK_ACTION_DIST_INPUTS: old_action_dist_logits}
+
+    @override(InferenceOnlyAPI)
+    def get_non_inference_attributes(self) -> List[str]:
+        """Returns attributes that should not be included in inference-only mode."""
+        return ["_old_lstm", "_old_fc_net", "_old_pi_head", "_values"]

rllib/examples/rl_modules/classes/mobilenet_rlm.py

@@ -0,0 +1,130 @@
+"""
+This example shows how to take full control over what models and action distribution
+are being built inside an RL Module. With this pattern, we can bypass a Catalog and
+explicitly define our own models within a given RL Module.
+
+Here, we plug a pre-trained MobileNet V3 (small) image encoder from `torchvision` into
+a PPO RLModule and use it to encode image observations before the policy- and
+value-heads. You can modify this example to accommodate your own encoder network or
+other pre-trained networks.
+"""
+# __sphinx_doc_begin__
+from dataclasses import dataclass
+
+import gymnasium as gym
+import numpy as np
+
+from ray.rllib.algorithms.ppo.ppo import PPOConfig
+from ray.rllib.algorithms.ppo.torch.default_ppo_torch_rl_module import (
+    DefaultPPOTorchRLModule,
+)
+from ray.rllib.core.models.base import ENCODER_OUT, Encoder
+from ray.rllib.core.models.configs import (
+    ActorCriticEncoderConfig,
+    MLPHeadConfig,
+    ModelConfig,
+)
+from ray.rllib.core.models.torch.base import TorchModel
+from ray.rllib.core.rl_module.rl_module import RLModuleSpec
+from ray.rllib.examples.envs.classes.random_env import RandomEnv
+from ray.rllib.utils.framework import try_import_torch
+
+torch, nn = try_import_torch()
+
+# torchvision's pre-trained image classifiers expect (3, 224, 224) inputs.
+MOBILENET_INPUT_SHAPE = (3, 224, 224)
+
+
+@dataclass
+class MobileNetV3EncoderConfig(ModelConfig):
+    freeze: bool = True
+
+    @property
+    def output_dims(self):
+        # MobileNet V3 (small) has a flat output of length 1000 (its ImageNet logits).
+        return (1000,)
+
+    def build(self, framework):
+        assert framework == "torch", "Unsupported framework `{}`!".format(framework)
+        return MobileNetV3Encoder(self)
+
+
+class MobileNetV3Encoder(TorchModel, Encoder):
+    """A MobileNet V3 (small) encoder for RLlib."""
+
+    def __init__(self, config):
+        super().__init__(config)
+        # Load MobileNet V3 (small) with its default pre-trained ImageNet weights from
+        # the installed torchvision. We use torchvision directly (rather than
+        # `torch.hub.load`) so the model code always matches the installed torch.
+        from torchvision.models import MobileNet_V3_Small_Weights, mobilenet_v3_small
+
+        self.net = mobilenet_v3_small(weights=MobileNet_V3_Small_Weights.DEFAULT)
+        if config.freeze:
+            # We don't want to train this encoder, so freeze its parameters!
+            for p in self.net.parameters():
+                p.requires_grad = False
+
+    def _forward(self, input_dict, **kwargs):
+        return {ENCODER_OUT: (self.net(input_dict["obs"]))}
+
+
+class MobileNetTorchPPORLModule(DefaultPPOTorchRLModule):
+    """A DefaultPPORLModule with MobileNet V3 (small) as an encoder.
+
+    The idea behind this model is to demonstrate how we can bypass catalog to
+    take full control over what models and action distribution are being built.
+    In this example, we do this to modify an existing RLModule with a custom encoder.
+    """
+
+    def setup(self):
+        mobilenet_config = MobileNetV3EncoderConfig()
+        # Since we want to use PPO, which is an actor-critic algorithm, we need to
+        # use an ActorCriticEncoderConfig to wrap the base encoder config.
+        actor_critic_encoder_config = ActorCriticEncoderConfig(
+            base_encoder_config=mobilenet_config
+        )
+
+        self.encoder = actor_critic_encoder_config.build(framework="torch")
+        mobilenet_output_dims = mobilenet_config.output_dims
+
+        pi_config = MLPHeadConfig(
+            input_dims=mobilenet_output_dims,
+            output_layer_dim=2,
+        )
+
+        vf_config = MLPHeadConfig(input_dims=mobilenet_output_dims, output_layer_dim=1)
+
+        self.pi = pi_config.build(framework="torch")
+        self.vf = vf_config.build(framework="torch")
+
+
+config = (
+    PPOConfig()
+    .environment(
+        RandomEnv,
+        env_config={
+            "action_space": gym.spaces.Discrete(2),
+            # Test a simple Image observation space.
+            "observation_space": gym.spaces.Box(
+                0.0,
+                1.0,
+                shape=MOBILENET_INPUT_SHAPE,
+                dtype=np.float32,
+            ),
+        },
+    )
+    .env_runners(num_env_runners=0)
+    # The following training settings make it so that a training iteration is very
+    # quick. This is just for the sake of this example. PPO will not learn properly
+    # with these settings!
+    .training(
+        train_batch_size_per_learner=32,
+        minibatch_size=16,
+        num_epochs=1,
+    )
+    .rl_module(rl_module_spec=RLModuleSpec(module_class=MobileNetTorchPPORLModule))
+)
+
+config.build().train()
+# __sphinx_doc_end__

rllib/examples/rl_modules/classes/modelv2_to_rlm.py

@@ -0,0 +1,207 @@
+import pathlib
+from typing import Any, Dict, Optional
+
+import tree
+
+from ray.rllib.core import DEFAULT_POLICY_ID, Columns
+from ray.rllib.core.distribution.torch.torch_distribution import (
+    TorchCategorical,
+    TorchDiagGaussian,
+    TorchMultiCategorical,
+    TorchMultiDistribution,
+    TorchSquashedGaussian,
+)
+from ray.rllib.core.rl_module.apis import ValueFunctionAPI
+from ray.rllib.core.rl_module.torch import TorchRLModule
+from ray.rllib.models.torch.torch_action_dist import (
+    TorchCategorical as OldTorchCategorical,
+    TorchDiagGaussian as OldTorchDiagGaussian,
+    TorchMultiActionDistribution as OldTorchMultiActionDistribution,
+    TorchMultiCategorical as OldTorchMultiCategorical,
+    TorchSquashedGaussian as OldTorchSquashedGaussian,
+)
+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
+
+torch, _ = try_import_torch()
+
+
+class ModelV2ToRLModule(TorchRLModule, ValueFunctionAPI):
+    """An RLModule containing a (old stack) ModelV2.
+
+    The `ModelV2` may be define either through
+    - an existing Policy checkpoint
+    - an existing Algorithm checkpoint (and a policy ID or "default_policy")
+    - or through an AlgorithmConfig object
+
+    The ModelV2 is created in the `setup` and contines to live through the lifetime
+    of the RLModule.
+    """
+
+    @override(TorchRLModule)
+    def setup(self):
+        # Try extracting the policy ID from this RLModule's config dict.
+        policy_id = self.model_config.get("policy_id", DEFAULT_POLICY_ID)
+
+        # Try getting the algorithm checkpoint from the `model_config`.
+        algo_checkpoint_dir = self.model_config.get("algo_checkpoint_dir")
+        if algo_checkpoint_dir:
+            algo_checkpoint_dir = pathlib.Path(algo_checkpoint_dir)
+            if not algo_checkpoint_dir.is_dir():
+                raise ValueError(
+                    "The `model_config` of your RLModule must contain a "
+                    "`algo_checkpoint_dir` key pointing to the algo checkpoint "
+                    "directory! You can find this dir inside the results dir of your "
+                    "experiment. You can then add this path "
+                    "through `config.rl_module(model_config={"
+                    "'algo_checkpoint_dir': [your algo checkpoint dir]})`."
+                )
+            policy_checkpoint_dir = algo_checkpoint_dir / "policies" / policy_id
+        # Try getting the policy checkpoint from the `model_config`.
+        else:
+            policy_checkpoint_dir = self.model_config.get("policy_checkpoint_dir")
+
+        # Create the ModelV2 from the Policy.
+        if policy_checkpoint_dir:
+            policy_checkpoint_dir = pathlib.Path(policy_checkpoint_dir)
+            if not policy_checkpoint_dir.is_dir():
+                raise ValueError(
+                    "The `model_config` of your RLModule must contain a "
+                    "`policy_checkpoint_dir` key pointing to the policy checkpoint "
+                    "directory! You can find this dir under the Algorithm's checkpoint "
+                    "dir in subdirectory: [algo checkpoint dir]/policies/[policy ID "
+                    "ex. `default_policy`]. You can then add this path through `config"
+                    ".rl_module(model_config={'policy_checkpoint_dir': "
+                    "[your policy checkpoint dir]})`."
+                )
+            # Create a temporary policy object.
+            policy = TorchPolicyV2.from_checkpoint(policy_checkpoint_dir)
+        # Create the ModelV2 from scratch using the config.
+        else:
+            config = self.model_config.get("old_api_stack_algo_config")
+            if not config:
+                raise ValueError(
+                    "The `model_config` of your RLModule must contain a "
+                    "`algo_config` key with a AlgorithmConfig object in it that "
+                    "contains all the settings that would be necessary to construct a "
+                    "old API stack Algorithm/Policy/ModelV2! You can add this setting "
+                    "through `config.rl_module(model_config={'algo_config': "
+                    "[your old config]})`."
+                )
+            # Get the multi-agent policies dict.
+            policy_dict, _ = config.get_multi_agent_setup(
+                spaces={
+                    policy_id: (self.observation_space, self.action_space),
+                },
+                default_policy_class=config.algo_class.get_default_policy_class(config),
+            )
+            config = config.to_dict()
+            config["__policy_id"] = policy_id
+            policy = policy_dict[policy_id].policy_class(
+                self.observation_space,
+                self.action_space,
+                config,
+            )
+
+        self._model_v2 = policy.model
+
+        # Translate the action dist classes from the old API stack to the new.
+        self.action_dist_class = self._translate_dist_class(policy.dist_class)
+
+        # Erase the torch policy from memory, so it can be garbage collected.
+        del policy
+
+    @override(TorchRLModule)
+    def _forward_inference(self, batch: Dict[str, Any], **kwargs) -> Dict[str, Any]:
+        return self._forward_pass(batch, inference=True)
+
+    @override(TorchRLModule)
+    def _forward_exploration(self, batch: Dict[str, Any], **kwargs) -> Dict[str, Any]:
+        return self._forward_inference(batch, **kwargs)
+
+    @override(TorchRLModule)
+    def _forward_train(self, batch: Dict[str, Any], **kwargs) -> Dict[str, Any]:
+        out = self._forward_pass(batch, inference=False)
+        out[Columns.ACTION_LOGP] = self.get_train_action_dist_cls()(
+            out[Columns.ACTION_DIST_INPUTS]
+        ).logp(batch[Columns.ACTIONS])
+        out[Columns.VF_PREDS] = self._model_v2.value_function()
+        if Columns.STATE_IN in batch and Columns.SEQ_LENS in batch:
+            out[Columns.VF_PREDS] = torch.reshape(
+                out[Columns.VF_PREDS], [len(batch[Columns.SEQ_LENS]), -1]
+            )
+        return out
+
+    def _forward_pass(self, batch, inference=True):
+        # Translate states and seq_lens into old API stack formats.
+        batch = batch.copy()
+        state_in = batch.pop(Columns.STATE_IN, {})
+        state_in = [s for i, s in sorted(state_in.items())]
+        seq_lens = batch.pop(Columns.SEQ_LENS, None)
+
+        if state_in:
+            if inference and seq_lens is None:
+                seq_lens = torch.tensor(
+                    [1.0] * state_in[0].shape[0], device=state_in[0].device
+                )
+            elif not inference:
+                assert seq_lens is not None
+            # Perform the actual ModelV2 forward pass.
+            # A recurrent ModelV2 adds and removes the time-rank itself (whereas in the
+            # new API stack, the connector pipelines are responsible for doing this) ->
+            # We have to remove, then re-add the time rank here to make ModelV2 work.
+            batch = tree.map_structure(
+                lambda s: torch.reshape(s, [-1] + list(s.shape[2:])), batch
+            )
+        nn_output, state_out = self._model_v2(batch, state_in, seq_lens)
+        # Put back 1ts time rank into nn-output (inference).
+        if state_in:
+            if inference:
+                nn_output = tree.map_structure(
+                    lambda s: torch.unsqueeze(s, axis=1), nn_output
+                )
+            else:
+                nn_output = tree.map_structure(
+                    lambda s: torch.reshape(s, [len(seq_lens), -1] + list(s.shape[1:])),
+                    nn_output,
+                )
+        # Interpret the NN output as action logits.
+        output = {Columns.ACTION_DIST_INPUTS: nn_output}
+        # Add the `state_out` to the `output`, new API stack style.
+        if state_out:
+            output[Columns.STATE_OUT] = {}
+        for i, o in enumerate(state_out):
+            output[Columns.STATE_OUT][i] = o
+
+        return output
+
+    @override(ValueFunctionAPI)
+    def compute_values(self, batch: Dict[str, Any], embeddings: Optional[Any] = None):
+        self._forward_pass(batch, inference=False)
+        v_preds = self._model_v2.value_function()
+        if Columns.STATE_IN in batch and Columns.SEQ_LENS in batch:
+            v_preds = torch.reshape(v_preds, [len(batch[Columns.SEQ_LENS]), -1])
+        return v_preds
+
+    @override(TorchRLModule)
+    def get_initial_state(self):
+        """Converts the initial state list of ModelV2 into a dict (new API stack)."""
+        init_state_list = self._model_v2.get_initial_state()
+        return dict(enumerate(init_state_list))
+
+    def _translate_dist_class(self, old_dist_class):
+        map_ = {
+            OldTorchCategorical: TorchCategorical,
+            OldTorchDiagGaussian: TorchDiagGaussian,
+            OldTorchMultiActionDistribution: TorchMultiDistribution,
+            OldTorchMultiCategorical: TorchMultiCategorical,
+            OldTorchSquashedGaussian: TorchSquashedGaussian,
+        }
+        if old_dist_class not in map_:
+            raise ValueError(
+                f"ModelV2ToRLModule does NOT support {old_dist_class} action "
+                f"distributions yet!"
+            )
+
+        return map_[old_dist_class]

rllib/examples/rl_modules/classes/random_rlm.py

@@ -0,0 +1,63 @@
+import gymnasium as gym
+import numpy as np
+import tree  # pip install dm_tree
+
+from ray.rllib.core.columns import Columns
+from ray.rllib.core.rl_module import RLModule
+from ray.rllib.policy.sample_batch import SampleBatch
+from ray.rllib.utils.annotations import override
+from ray.rllib.utils.spaces.space_utils import batch as batch_func
+
+
+class RandomRLModule(RLModule):
+    @override(RLModule)
+    def _forward(self, batch, **kwargs):
+        obs_batch_size = len(tree.flatten(batch[SampleBatch.OBS])[0])
+        actions = batch_func(
+            [self.action_space.sample() for _ in range(obs_batch_size)]
+        )
+        return {SampleBatch.ACTIONS: actions}
+
+    @override(RLModule)
+    def _forward_train(self, *args, **kwargs):
+        # RandomRLModule should always be configured as non-trainable.
+        # To do so, set in your config:
+        # `config.multi_agent(policies_to_train=[list of ModuleIDs to be trained,
+        # NOT including the ModuleID of this RLModule])`
+        raise NotImplementedError("Random RLModule: Should not be trained!")
+
+    def compile(self, *args, **kwargs):
+        """Dummy method for compatibility with TorchRLModule.
+
+        This is hit when RolloutWorker tries to compile TorchRLModule."""
+        pass
+
+
+class StatefulRandomRLModule(RandomRLModule):
+    """A stateful RLModule that returns STATE_OUT from its forward methods.
+
+    - Implements the `get_initial_state` method (returning a all-zeros dummy state).
+    - Returns a dummy state under the `Columns.STATE_OUT` from its forward methods.
+    """
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self._internal_state_space = gym.spaces.Box(-1.0, 1.0, (1,))
+
+    @override(RLModule)
+    def get_initial_state(self):
+        return {
+            "state": np.zeros_like([self._internal_state_space.sample()]),
+        }
+
+    def _random_forward(self, batch, **kwargs):
+        batch = super()._random_forward(batch, **kwargs)
+        batch[Columns.STATE_OUT] = {
+            "state": batch_func(
+                [
+                    self._internal_state_space.sample()
+                    for _ in range(len(batch[Columns.ACTIONS]))
+                ]
+            ),
+        }
+        return batch

rllib/examples/rl_modules/classes/rock_paper_scissors_heuristic_rlm.py

@@ -0,0 +1,92 @@
+from collections import defaultdict
+
+import numpy as np
+
+from ray.rllib.core.columns import Columns
+from ray.rllib.core.rl_module.rl_module import RLModule
+from ray.rllib.utils.annotations import override
+
+
+class AlwaysSameHeuristicRLM(RLModule):
+    """In rock-paper-scissors, always chooses the same action within an episode.
+
+    The first move is random, all the following moves are the same as the first one.
+    """
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self._actions_per_vector_idx = defaultdict(int)
+
+    @override(RLModule)
+    def _forward_inference(self, batch, **kwargs):
+        ret = []
+        # Note that the obs is the previous move of the opponens (0-2). If it's 3, it
+        # means that there was no previous move and thus, the episode just started.
+        for i, obs in enumerate(batch[Columns.OBS]):
+            if obs == 3:
+                self._actions_per_vector_idx[i] = np.random.choice([0, 1, 2])
+            ret.append(self._actions_per_vector_idx[i])
+        return {Columns.ACTIONS: np.array(ret)}
+
+    @override(RLModule)
+    def _forward_exploration(self, batch, **kwargs):
+        return self._forward_inference(batch, **kwargs)
+
+    @override(RLModule)
+    def _forward_train(self, batch, **kwargs):
+        raise NotImplementedError(
+            "AlwaysSameHeuristicRLM is not trainable! Make sure you do NOT include it "
+            "in your `config.multi_agent(policies_to_train={...})` set."
+        )
+
+
+class BeatLastHeuristicRLM(RLModule):
+    """In rock-paper-scissors, always acts such that it beats prev. move of opponent.
+
+    The first move is random.
+
+    For example, after opponent played `rock` (and this policy made a random
+    move), the next move would be `paper`(to beat `rock`).
+    """
+
+    @override(RLModule)
+    def _forward_inference(self, batch, **kwargs):
+        """Returns the exact action that would beat the previous action of the opponent.
+
+        The opponent's previous action is the current observation for this agent.
+
+        Both action- and observation spaces are discrete. There are 3 actions available.
+        (0-2) and 4 observations (0-2 plus 3, where 3 is the observation after the env
+        reset, when no action has been taken yet). Thereby:
+        0=Rock
+        1=Paper
+        2=Scissors
+        3=[after reset] (observation space only)
+        """
+        return {
+            Columns.ACTIONS: np.array(
+                [self._pick_single_action(obs) for obs in batch[Columns.OBS]]
+            ),
+        }
+
+    @override(RLModule)
+    def _forward_exploration(self, batch, **kwargs):
+        return self._forward_inference(batch, **kwargs)
+
+    @override(RLModule)
+    def _forward_train(self, batch, **kwargs):
+        raise NotImplementedError(
+            "BeatLastHeuristicRLM is not trainable! Make sure you do NOT include it in "
+            "your `config.multi_agent(policies_to_train={...})` set."
+        )
+
+    @staticmethod
+    def _pick_single_action(prev_opponent_obs):
+        if prev_opponent_obs == 0:
+            return 1
+        elif prev_opponent_obs == 1:
+            return 2
+        elif prev_opponent_obs == 2:
+            return 0
+        else:
+            return np.random.choice([0, 1, 2])

rllib/examples/rl_modules/classes/tiny_atari_cnn_rlm.py

@@ -0,0 +1,194 @@
+from typing import Any, Dict, Optional
+
+from ray.rllib.core.columns import Columns
+from ray.rllib.core.learner.utils import make_target_network
+from ray.rllib.core.rl_module.apis import (
+    TARGET_NETWORK_ACTION_DIST_INPUTS,
+    TargetNetworkAPI,
+    ValueFunctionAPI,
+)
+from ray.rllib.core.rl_module.torch import TorchRLModule
+from ray.rllib.models.torch.misc import (
+    normc_initializer,
+    same_padding,
+    valid_padding,
+)
+from ray.rllib.utils.annotations import override
+from ray.rllib.utils.framework import try_import_torch
+from ray.rllib.utils.typing import TensorType
+
+torch, nn = try_import_torch()
+
+
+class TinyAtariCNN(TorchRLModule, ValueFunctionAPI, TargetNetworkAPI):
+    """A tiny CNN stack for fast-learning of Atari envs.
+
+    The architecture here is the exact same as the one used by the old API stack as
+    CNN default ModelV2.
+
+    We stack 3 CNN layers based on the config, then a 4th one with linear activation
+    and n 1x1 filters, where n is the number of actions in the (discrete) action space.
+    Simple reshaping (no flattening or extra linear layers necessary) lead to the
+    action logits, which can directly be used inside a distribution or loss.
+
+    .. testcode::
+
+        import numpy as np
+        import gymnasium as gym
+
+        my_net = TinyAtariCNN(
+            observation_space=gym.spaces.Box(-1.0, 1.0, (42, 42, 4), np.float32),
+            action_space=gym.spaces.Discrete(4),
+        )
+
+        B = 10
+        w = 42
+        h = 42
+        c = 4
+        data = torch.from_numpy(
+            np.random.random_sample(size=(B, w, h, c)).astype(np.float32)
+        )
+        print(my_net.forward_inference({"obs": data}))
+        print(my_net.forward_exploration({"obs": data}))
+        print(my_net.forward_train({"obs": data}))
+
+        num_all_params = sum(int(np.prod(p.size())) for p in my_net.parameters())
+        print(f"num params = {num_all_params}")
+    """
+
+    @override(TorchRLModule)
+    def setup(self):
+        """Use this method to create all the model components that you require.
+
+        Feel free to access the following useful properties in this class:
+        - `self.model_config`: The config dict for this RLModule class,
+        which should contain flxeible settings, for example: {"hiddens": [256, 256]}.
+        - `self.observation|action_space`: The observation and action space that
+        this RLModule is subject to. Note that the observation space might not be the
+        exact space from your env, but that it might have already gone through
+        preprocessing through a connector pipeline (for example, flattening,
+        frame-stacking, mean/std-filtering, etc..).
+        """
+        # Get the CNN stack config from our RLModuleConfig's (self.config)
+        # `model_config` property:
+        conv_filters = self.model_config.get("conv_filters")
+        # Default CNN stack with 3 layers:
+        if conv_filters is None:
+            conv_filters = [
+                [16, 4, 2, "same"],  # num filters, kernel wxh, stride wxh, padding type
+                [32, 4, 2, "same"],
+                [256, 11, 1, "valid"],
+            ]
+
+        # Build the CNN layers.
+        layers = []
+
+        # Add user-specified hidden convolutional layers first
+        width, height, in_depth = self.observation_space.shape
+        in_size = [width, height]
+        for filter_specs in conv_filters:
+            if len(filter_specs) == 4:
+                out_depth, kernel_size, strides, padding = filter_specs
+            else:
+                out_depth, kernel_size, strides = filter_specs
+                padding = "same"
+
+            # Pad like in tensorflow's SAME mode.
+            if padding == "same":
+                padding_size, out_size = same_padding(in_size, kernel_size, strides)
+                layers.append(nn.ZeroPad2d(padding_size))
+            # No actual padding is performed for "valid" mode, but we will still
+            # compute the output size (input for the next layer).
+            else:
+                out_size = valid_padding(in_size, kernel_size, strides)
+
+            layer = nn.Conv2d(in_depth, out_depth, kernel_size, strides, bias=True)
+            # Initialize CNN layer kernel and bias.
+            nn.init.xavier_uniform_(layer.weight)
+            nn.init.zeros_(layer.bias)
+            layers.append(layer)
+            # Activation.
+            layers.append(nn.ReLU())
+
+            in_size = out_size
+            in_depth = out_depth
+
+        self._base_cnn_stack = nn.Sequential(*layers)
+
+        # Add the final CNN 1x1 layer with num_filters == num_actions to be reshaped to
+        # yield the logits (no flattening, no additional linear layers required).
+        _final_conv = nn.Conv2d(in_depth, self.action_space.n, 1, 1, bias=True)
+        nn.init.xavier_uniform_(_final_conv.weight)
+        nn.init.zeros_(_final_conv.bias)
+        self._logits = nn.Sequential(
+            nn.ZeroPad2d(same_padding(in_size, 1, 1)[0]), _final_conv
+        )
+
+        self._values = nn.Linear(in_depth, 1)
+        # Mimick old API stack behavior of initializing the value function with `normc`
+        # std=0.01.
+        normc_initializer(0.01)(self._values.weight)
+
+    @override(TorchRLModule)
+    def _forward(self, batch, **kwargs):
+        # Compute the basic 1D feature tensor (inputs to policy- and value-heads).
+        _, logits = self._compute_embeddings_and_logits(batch)
+        # Return features and logits as ACTION_DIST_INPUTS (categorical distribution).
+        return {
+            Columns.ACTION_DIST_INPUTS: logits,
+        }
+
+    @override(TorchRLModule)
+    def _forward_train(self, batch, **kwargs):
+        # Compute the basic 1D feature tensor (inputs to policy- and value-heads).
+        embeddings, logits = self._compute_embeddings_and_logits(batch)
+        # Return features and logits as ACTION_DIST_INPUTS (categorical distribution).
+        return {
+            Columns.ACTION_DIST_INPUTS: logits,
+            Columns.EMBEDDINGS: embeddings,
+        }
+
+    # We implement this RLModule as a TargetNetworkAPI RLModule, so it can be used
+    # by the APPO algorithm.
+    @override(TargetNetworkAPI)
+    def make_target_networks(self) -> None:
+        self._target_base_cnn_stack = make_target_network(self._base_cnn_stack)
+        self._target_logits = make_target_network(self._logits)
+
+    @override(TargetNetworkAPI)
+    def get_target_network_pairs(self):
+        return [
+            (self._base_cnn_stack, self._target_base_cnn_stack),
+            (self._logits, self._target_logits),
+        ]
+
+    @override(TargetNetworkAPI)
+    def forward_target(self, batch, **kw):
+        obs = batch[Columns.OBS].permute(0, 3, 1, 2)
+        embeddings = self._target_base_cnn_stack(obs)
+        logits = self._target_logits(embeddings)
+        return {TARGET_NETWORK_ACTION_DIST_INPUTS: torch.squeeze(logits, dim=[-1, -2])}
+
+    # We implement this RLModule as a ValueFunctionAPI RLModule, so it can be used
+    # by value-based methods like PPO or IMPALA.
+    @override(ValueFunctionAPI)
+    def compute_values(
+        self,
+        batch: Dict[str, Any],
+        embeddings: Optional[Any] = None,
+    ) -> TensorType:
+        # Features not provided -> We need to compute them first.
+        if embeddings is None:
+            obs = batch[Columns.OBS]
+            embeddings = self._base_cnn_stack(obs.permute(0, 3, 1, 2))
+            embeddings = torch.squeeze(embeddings, dim=[-1, -2])
+        return self._values(embeddings).squeeze(-1)
+
+    def _compute_embeddings_and_logits(self, batch):
+        obs = batch[Columns.OBS].permute(0, 3, 1, 2)
+        embeddings = self._base_cnn_stack(obs)
+        logits = self._logits(embeddings)
+        return (
+            torch.squeeze(embeddings, dim=[-1, -2]),
+            torch.squeeze(logits, dim=[-1, -2]),
+        )

rllib/examples/rl_modules/classes/vpg_torch_rlm.py

@@ -0,0 +1,59 @@
+import torch
+
+from ray.rllib.core.columns import Columns
+from ray.rllib.core.rl_module.torch import TorchRLModule
+
+
+class VPGTorchRLModule(TorchRLModule):
+    """A simple VPG (vanilla policy gradient)-style RLModule for testing purposes.
+
+    Use this as a minimum, bare-bones example implementation of a custom TorchRLModule.
+    """
+
+    def setup(self):
+        """Use this method to create all the model components that you require.
+
+        Feel free to access the following useful properties in this class:
+        - `self.model_config`: The config dict for this RLModule class,
+        which should contain flexible settings, for example: {"hiddens": [256, 256]}.
+        - `self.observation|action_space`: The observation and action space that
+        this RLModule is subject to. Note that the observation space might not be the
+        exact space from your env, but that it might have already gone through
+        preprocessing through a connector pipeline (for example, flattening,
+        frame-stacking, mean/std-filtering, etc..).
+        - `self.inference_only`: If True, this model should be built only for inference
+        purposes, in which case you may want to exclude any components that are not used
+        for computing actions, for example a value function branch.
+        """
+        input_dim = self.observation_space.shape[0]
+        hidden_dim = self.model_config["hidden_dim"]
+        output_dim = self.action_space.n
+
+        self._policy_net = torch.nn.Sequential(
+            torch.nn.Linear(input_dim, hidden_dim),
+            torch.nn.ReLU(),
+            torch.nn.Linear(hidden_dim, output_dim),
+        )
+
+    def _forward(self, batch, **kwargs):
+        # Push the observations from the batch through our `self._policy_net`.
+        action_logits = self._policy_net(batch[Columns.OBS])
+        # Return parameters for the (default) action distribution, which is
+        # `TorchCategorical` (due to our action space being `gym.spaces.Discrete`).
+        return {Columns.ACTION_DIST_INPUTS: action_logits}
+
+        # If you need more granularity between the different forward behaviors during
+        # the different phases of the module's lifecycle, implement three different
+        # forward methods. Thereby, it is recommended to put the inference and
+        # exploration versions inside a `with torch.no_grad()` context for better
+        # performance.
+        # def _forward_train(self, batch):
+        #    ...
+        #
+        # def _forward_inference(self, batch):
+        #    with torch.no_grad():
+        #        return self._forward_train(batch)
+        #
+        # def _forward_exploration(self, batch):
+        #    with torch.no_grad():
+        #        return self._forward_train(batch)

rllib/examples/rl_modules/classes/vpg_using_shared_encoder_rlm.py

@@ -0,0 +1,254 @@
+from typing import (
+    Any,
+    Dict,
+    Union,
+)
+
+import torch
+
+from ray.rllib.core import Columns
+from ray.rllib.core.models.base import ENCODER_OUT
+from ray.rllib.core.rl_module.multi_rl_module import MultiRLModule
+from ray.rllib.core.rl_module.torch.torch_rl_module import TorchRLModule
+from ray.rllib.utils.annotations import override
+from ray.rllib.utils.typing import ModuleID
+
+SHARED_ENCODER_ID = "shared_encoder"
+
+
+# __sphinx_doc_policy_begin__
+class VPGPolicyAfterSharedEncoder(TorchRLModule):
+    """A VPG (vanilla pol. gradient)-style RLModule using a shared encoder.
+    # __sphinx_doc_policy_end__
+
+        The shared encoder RLModule must be held by the same MultiRLModule, under which
+        this RLModule resides. The shared encoder's forward is called before this
+        RLModule's forward and returns the embeddings under the "encoder_embeddings"
+        key.
+    """
+
+    # __sphinx_doc_policy_2_begin__
+
+    def setup(self):
+        super().setup()
+
+        # Incoming feature dim from the shared encoder.
+        embedding_dim = self.model_config["embedding_dim"]
+        hidden_dim = self.model_config["hidden_dim"]
+
+        self._pi_head = torch.nn.Sequential(
+            torch.nn.Linear(embedding_dim, hidden_dim),
+            torch.nn.ReLU(),
+            torch.nn.Linear(hidden_dim, self.action_space.n),
+        )
+
+    def _forward(self, batch, **kwargs):
+        embeddings = batch[ENCODER_OUT]  # Get the output of the encoder
+        logits = self._pi_head(embeddings)
+        return {Columns.ACTION_DIST_INPUTS: logits}
+
+
+# __sphinx_doc_policy_2_end__
+
+
+# __sphinx_doc_mrlm_begin__
+class VPGMultiRLModuleWithSharedEncoder(MultiRLModule):
+    """VPG (vanilla pol. gradient)-style MultiRLModule handling a shared encoder.
+    # __sphinx_doc_mrlm_end__
+
+        This MultiRLModule needs to be configured appropriately as below.
+
+        .. testcode::
+
+            # __sphinx_doc_how_to_run_begin__
+            import gymnasium as gym
+            from ray.rllib.core.rl_module.rl_module import RLModuleSpec
+            from ray.rllib.core.rl_module.multi_rl_module import MultiRLModuleSpec
+
+            from ray.rllib.examples.algorithms.classes.vpg import VPGConfig
+            from ray.rllib.examples.learners.classes.vpg_torch_learner_shared_optimizer import VPGTorchLearnerSharedOptimizer
+            from ray.rllib.examples.envs.classes.multi_agent import MultiAgentCartPole
+            from ray.rllib.examples.rl_modules.classes.vpg_using_shared_encoder_rlm import (
+                SHARED_ENCODER_ID,
+                SharedEncoder,
+                VPGPolicyAfterSharedEncoder,
+                VPGMultiRLModuleWithSharedEncoder,
+            )
+
+            single_agent_env = gym.make("CartPole-v1")
+
+            EMBEDDING_DIM = 64  # encoder output dim
+
+            config = (
+                VPGConfig()
+                .environment(MultiAgentCartPole, env_config={"num_agents": 2})
+                .training(
+                    learner_class=VPGTorchLearnerSharedOptimizer,
+                )
+                .multi_agent(
+                    # Declare the two policies trained.
+                    policies={"p0", "p1"},
+                    # Agent IDs of `MultiAgentCartPole` are 0 and 1. They are mapped to
+                    # the two policies with ModuleIDs "p0" and "p1", respectively.
+                    policy_mapping_fn=lambda agent_id, episode, **kw: f"p{agent_id}"
+                )
+                .rl_module(
+                    rl_module_spec=MultiRLModuleSpec(
+                        multi_rl_module_class=VPGMultiRLModuleWithSharedEncoder,
+                        rl_module_specs={
+                            # Shared encoder.
+                            SHARED_ENCODER_ID: RLModuleSpec(
+                                module_class=SharedEncoder,
+                                model_config={"embedding_dim": EMBEDDING_DIM},
+                                observation_space=single_agent_env.observation_space,
+                                action_space=single_agent_env.action_space,
+                            ),
+                            # Large policy net.
+                            "p0": RLModuleSpec(
+                                module_class=VPGPolicyAfterSharedEncoder,
+                                model_config={
+                                    "embedding_dim": EMBEDDING_DIM,
+                                    "hidden_dim": 1024,
+                                },
+                            ),
+                            # Small policy net.
+                            "p1": RLModuleSpec(
+                                module_class=VPGPolicyAfterSharedEncoder,
+                                model_config={
+                                    "embedding_dim": EMBEDDING_DIM,
+                                    "hidden_dim": 64,
+                                },
+                            ),
+                        },
+                    ),
+                )
+            )
+            algo = config.build_algo()
+            print(algo.train())
+            # __sphinx_doc_how_to_run_end__
+    """
+
+    # __sphinx_doc_mrlm_2_begin__
+
+    def setup(self):
+        # Call the super's setup().
+        super().setup()
+        # Assert, we have the shared encoder submodule.
+        assert SHARED_ENCODER_ID in self._rl_modules and len(self._rl_modules) > 1
+        # Assign the encoder to a convenience attribute.
+        self.encoder = self._rl_modules[SHARED_ENCODER_ID]
+
+    def _forward(self, batch, forward_type, **kwargs):
+        # Collect our policies' outputs in this dict.
+        fwd_out = {}
+        # Loop through the policy nets (through the given batch's keys).
+        for policy_id, policy_batch in batch.items():
+            # Feed this policy's observation into the shared encoder
+            encoder_output = self.encoder._forward(batch[policy_id])
+            policy_batch[ENCODER_OUT] = encoder_output[ENCODER_OUT]
+            # Get the desired module
+            m = getattr(self._rl_modules[policy_id], forward_type)
+            # Pass the policy's embeddings through the policy net.
+            fwd_out[policy_id] = m(batch[policy_id], **kwargs)
+        return fwd_out
+
+    # These methods could probably stand to be adjusted in MultiRLModule using something like this, so that subclasses that tweak _forward don't need to rewrite all of them. The prior implementation errored out because of this issue.
+    @override(MultiRLModule)
+    def _forward_inference(
+        self, batch: Dict[str, Any], **kwargs
+    ) -> Union[Dict[str, Any], Dict[ModuleID, Dict[str, Any]]]:
+        return self._forward(batch, "_forward_inference", **kwargs)
+
+    @override(MultiRLModule)
+    def _forward_exploration(
+        self, batch: Dict[str, Any], **kwargs
+    ) -> Union[Dict[str, Any], Dict[ModuleID, Dict[str, Any]]]:
+        return self._forward(batch, "_forward_exploration", **kwargs)
+
+    @override(MultiRLModule)
+    def _forward_train(
+        self, batch: Dict[str, Any], **kwargs
+    ) -> Union[Dict[str, Any], Dict[ModuleID, Dict[str, Any]]]:
+        return self._forward(batch, "_forward_train", **kwargs)
+
+
+# __sphinx_doc_mrlm_2_end__
+
+
+# __sphinx_doc_encoder_begin__
+class SharedEncoder(TorchRLModule):
+    """A shared encoder that can be used with `VPGMultiRLModuleWithSharedEncoder`."""
+
+    def setup(self):
+        super().setup()
+
+        input_dim = self.observation_space.shape[0]
+        embedding_dim = self.model_config["embedding_dim"]
+
+        # A very simple encoder network.
+        self._net = torch.nn.Sequential(
+            torch.nn.Linear(input_dim, embedding_dim),
+        )
+
+    def _forward(self, batch, **kwargs):
+        # Pass observations through the net and return outputs.
+        return {ENCODER_OUT: self._net(batch[Columns.OBS])}
+
+
+# __sphinx_doc_encoder_end__
+
+
+# __sphinx_doc_ns_encoder_begin__
+class VPGIndividualEncoder(torch.nn.Module):
+    def __init__(self, observation_space, embedding_dim):
+        """
+        An individual version of SharedEncoder, supporting direct comparison between
+        the two architectures.
+        """
+        super().__init__()
+
+        input_dim = observation_space.shape[0]
+
+        # A very simple encoder network.
+        self._net = torch.nn.Sequential(
+            torch.nn.Linear(input_dim, embedding_dim),
+        )
+
+    def forward(self, batch, **kwargs):
+        # Pass observations through the net and return outputs.
+        return {ENCODER_OUT: self._net(batch[Columns.OBS])}
+
+
+# __sphinx_doc_ns_encoder_end__
+
+
+# __sphinx_doc_ns_policy_begin__
+class VPGPolicyNoSharedEncoder(TorchRLModule):
+    """
+    A VPG (vanilla pol. gradient)-style RLModule that doesn't use a shared encoder.
+    Facilitates experiments comparing shared and individual encoder architectures.
+    """
+
+    def setup(self):
+        super().setup()
+
+        # Incoming feature dim from the encoder.
+        embedding_dim = self.model_config["embedding_dim"]
+        hidden_dim = self.model_config["hidden_dim"]
+
+        self._pi_head = torch.nn.Sequential(
+            torch.nn.Linear(embedding_dim, hidden_dim),
+            torch.nn.ReLU(),
+            torch.nn.Linear(hidden_dim, self.action_space.n),
+        )
+        self.encoder = VPGIndividualEncoder(self.observation_space, embedding_dim)
+
+    def _forward(self, batch, **kwargs):
+        if ENCODER_OUT not in batch:
+            batch = self.encoder(batch)
+        embeddings = batch[ENCODER_OUT]
+        logits = self._pi_head(embeddings)
+        return {Columns.ACTION_DIST_INPUTS: logits}
+
+
+# __sphinx_doc_ns_policy_end__