chore: import upstream snapshot with attribution

rllib/examples/rl_modules/action_masking_rl_module.py

@@ -0,0 +1,125 @@
+"""An example script showing how to define and load an `RLModule` that applies
+action masking
+
+This example:
+    - Defines an `RLModule` that applies action masking.
+    - It does so by using a `gymnasium.spaces.dict.Dict` observation space
+        with two keys, namely `"observations"`, holding the original observations
+        and `"action_mask"` defining the action mask for the current environment
+        state. Note, by this definition you can wrap any `gymnasium` environment
+        and use it for this module.
+    - Furthermore, it derives its `TorchRLModule` from the `PPOTorchRLModule` and
+        can therefore be easily plugged into our `PPO` algorithm.
+    - It overrides the `forward` methods of the `PPOTorchRLModule` to apply the
+        action masking and it overrides the `_compute_values` method for GAE
+        computation to extract the `"observations"` from the batch `Columns.OBS`
+        key.
+    - It uses the custom `ActionMaskEnv` that defines for each step a new action
+        mask that defines actions that are allowed (1.0) and others that are not
+        (0.0).
+    - It runs 10 iterations with PPO and finishes.
+
+How to run this script
+----------------------
+`python [script file name].py --num-env-runners 2`
+
+Control the number of `EnvRunner`s with the `--num-env-runners` flag. This
+will increase the sampling speed.
+
+For debugging, use the following additional command line options
+`--no-tune --num-env-runners=0`
+which should allow you to set breakpoints anywhere in the RLlib code and
+have the execution stop there for inspection and debugging.
+
+For logging to your WandB account, use:
+`--wandb-key=[your WandB API key] --wandb-project=[some project name]
+--wandb-run-name=[optional: WandB run name (within the defined project)]`
+
+
+Results to expect
+-----------------
+You should expect a mean episode reward of around 0.35. The environment is a random
+environment paying out random rewards - so the agent cannot learn, but it can obey the
+action mask and should do so (no `AssertionError` should happen).
+After 40,000 environment steps and 10 training iterations the run should stop
+successfully:
+
++-------------------------------+------------+----------------------+--------+
+| Trial name                    | status     | loc                  |   iter |
+|                               |            |                      |        |
+|-------------------------------+------------+----------------------+--------+
+| PPO_ActionMaskEnv_dedc8_00000 | TERMINATED | 192.168.1.178:103298 |     10 |
++-------------------------------+------------+----------------------+--------+
++------------------+------------------------+------------------------+
+|   total time (s) |   num_env_steps_sample |   num_env_steps_traine |
+|                  |             d_lifetime |             d_lifetime |
++------------------+------------------------+------------------------+
+|          57.9207 |                  40000 |                  40000 |
++------------------+------------------------+------------------------+
+*------------------------+
+|   num_episodes_lifetim |
+|                      e |
++------------------------|
+|                   3898 |
++------------------------+
+"""
+from gymnasium.spaces import Box, Discrete
+
+from ray.rllib.algorithms.ppo import PPOConfig
+from ray.rllib.core.rl_module.rl_module import RLModuleSpec
+from ray.rllib.examples.envs.classes.action_mask_env import ActionMaskEnv
+from ray.rllib.examples.rl_modules.classes.action_masking_rlm import (
+    ActionMaskingTorchRLModule,
+)
+from ray.rllib.examples.utils import (
+    add_rllib_example_script_args,
+    run_rllib_example_script_experiment,
+)
+
+parser = add_rllib_example_script_args(
+    default_iters=10,
+    default_timesteps=100000,
+    default_reward=150.0,
+)
+
+if __name__ == "__main__":
+    args = parser.parse_args()
+
+    if args.algo != "PPO":
+        raise ValueError("This example only supports PPO. Please use --algo=PPO.")
+
+    base_config = (
+        PPOConfig()
+        .environment(
+            env=ActionMaskEnv,
+            env_config={
+                "action_space": Discrete(100),
+                # This defines the 'original' observation space that is used in the
+                # `RLModule`. The environment will wrap this space into a
+                # `gym.spaces.Dict` together with an 'action_mask' that signals the
+                # `RLModule` to adapt the action distribution inputs for the underlying
+                # `DefaultPPORLModule`.
+                "observation_space": Box(-1.0, 1.0, (5,)),
+            },
+        )
+        .rl_module(
+            # We need to explicitly specify here RLModule to use and
+            # the catalog needed to build it.
+            rl_module_spec=RLModuleSpec(
+                module_class=ActionMaskingTorchRLModule,
+                model_config={
+                    "head_fcnet_hiddens": [64, 64],
+                    "head_fcnet_activation": "relu",
+                },
+            ),
+        )
+        .evaluation(
+            evaluation_num_env_runners=1,
+            evaluation_interval=1,
+            # Run evaluation parallel to training to speed up the example.
+            evaluation_parallel_to_training=True,
+        )
+    )
+
+    # Run the example (with Tune).
+    run_rllib_example_script_experiment(base_config, args)

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__

rllib/examples/rl_modules/custom_cnn_rl_module.py

@@ -0,0 +1,115 @@
+"""Example of implementing and configuring a custom (torch) CNN containing RLModule.
+
+This example:
+    - demonstrates how you can subclass the TorchRLModule base class and set up your
+    own CNN-stack architecture by overriding the `setup()` method.
+    - shows how to override the 3 forward methods: `_forward_inference()`,
+    `_forward_exploration()`, and `forward_train()` to implement your own custom forward
+    logic(s). You will also learn, when each of these 3 methods is called by RLlib or
+    the users of your RLModule.
+    - shows how you then configure an RLlib Algorithm such that it uses your custom
+    RLModule (instead of a default RLModule).
+
+We implement a tiny CNN stack here, the exact same one that is used by the old API
+stack as default CNN net. It comprises 4 convolutional layers, the last of which
+ends in a 1x1 filter size and the number of filters exactly matches the number of
+discrete actions (logits). This way, the (non-activated) output of the last layer only
+needs to be reshaped in order to receive the policy's logit outputs. No flattening
+or additional dense layer required.
+
+The network is then used in a fast ALE/Pong-v5 experiment.
+
+
+How to run this script
+----------------------
+`python [script file name].py`
+
+For debugging, use the following additional command line options
+`--no-tune --num-env-runners=0`
+which should allow you to set breakpoints anywhere in the RLlib code and
+have the execution stop there for inspection and debugging.
+
+For logging to your WandB account, use:
+`--wandb-key=[your WandB API key] --wandb-project=[some project name]
+--wandb-run-name=[optional: WandB run name (within the defined project)]`
+
+
+Results to expect
+-----------------
+You should see the following output (during the experiment) in your console:
+
+Number of trials: 1/1 (1 RUNNING)
++---------------------+----------+----------------+--------+------------------+
+| Trial name          | status   | loc            |   iter |   total time (s) |
+|                     |          |                |        |                  |
+|---------------------+----------+----------------+--------+------------------+
+| PPO_env_82b44_00000 | RUNNING  | 127.0.0.1:9718 |      1 |          98.3585 |
++---------------------+----------+----------------+--------+------------------+
++------------------------+------------------------+------------------------+
+|   num_env_steps_sample |   num_env_steps_traine |   num_episodes_lifetim |
+|             d_lifetime |             d_lifetime |                      e |
+|------------------------+------------------------+------------------------|
+|                   4000 |                   4000 |                      4 |
++------------------------+------------------------+------------------------+
+"""
+import gymnasium as gym
+
+from ray.rllib.core.rl_module.rl_module import RLModuleSpec
+from ray.rllib.env.wrappers.atari_wrappers import wrap_atari_for_new_api_stack
+from ray.rllib.examples.rl_modules.classes.tiny_atari_cnn_rlm import TinyAtariCNN
+from ray.rllib.examples.utils import (
+    add_rllib_example_script_args,
+    run_rllib_example_script_experiment,
+)
+from ray.tune.registry import get_trainable_cls, register_env
+
+parser = add_rllib_example_script_args(default_iters=100, default_timesteps=600000)
+parser.set_defaults(
+    env="ale_py:ALE/Pong-v5",
+)
+
+
+if __name__ == "__main__":
+    args = parser.parse_args()
+
+    register_env(
+        "env",
+        lambda cfg: wrap_atari_for_new_api_stack(
+            gym.make(args.env, **cfg),
+            dim=42,  # <- need images to be "tiny" for our custom model
+            framestack=4,
+        ),
+    )
+
+    base_config = (
+        get_trainable_cls(args.algo)
+        .get_default_config()
+        .environment(
+            env="env",
+            env_config=dict(
+                frameskip=1,
+                full_action_space=False,
+                repeat_action_probability=0.0,
+            ),
+        )
+        .rl_module(
+            # Plug-in our custom RLModule class.
+            rl_module_spec=RLModuleSpec(
+                module_class=TinyAtariCNN,
+                # Feel free to specify your own `model_config` settings below.
+                # The `model_config` defined here will be available inside your
+                # custom RLModule class through the `self.model_config`
+                # property.
+                model_config={
+                    "conv_filters": [
+                        # num filters, kernel wxh, stride wxh, padding type
+                        [16, 4, 2, "same"],
+                        [32, 4, 2, "same"],
+                        [256, 11, 1, "valid"],
+                    ],
+                },
+            ),
+        )
+    )
+
+    run_rllib_example_script_experiment(base_config, args)

rllib/examples/rl_modules/custom_lstm_rl_module.py

@@ -0,0 +1,101 @@
+"""Example of implementing and configuring a custom (torch) LSTM containing RLModule.
+
+This example:
+    - demonstrates how you can subclass the TorchRLModule base class and set up your
+    own LSTM-containing NN architecture by overriding the `setup()` method.
+    - shows how to override the 3 forward methods: `_forward_inference()`,
+    `_forward_exploration()`, and `forward_train()` to implement your own custom forward
+    logic(s), including how to handle STATE in- and outputs to and from these calls.
+    - explains when each of these 3 methods is called by RLlib or the users of your
+    RLModule.
+    - shows how you then configure an RLlib Algorithm such that it uses your custom
+    RLModule (instead of a default RLModule).
+
+We implement a simple LSTM layer here, followed by a series of Linear layers.
+After the last Linear layer, we add fork of 2 Linear (non-activated) layers, one for the
+action logits and one for the value function output.
+
+We test the LSTM containing RLModule on the StatelessCartPole environment, a variant
+of CartPole that is non-Markovian (partially observable). Only an RNN-network can learn
+a decent policy in this environment due to the lack of any velocity information. By
+looking at one observation, one cannot know whether the cart is currently moving left or
+right and whether the pole is currently moving up or down).
+
+
+How to run this script
+----------------------
+`python [script file name].py`
+
+For debugging, use the following additional command line options
+`--no-tune --num-env-runners=0`
+which should allow you to set breakpoints anywhere in the RLlib code and
+have the execution stop there for inspection and debugging.
+
+For logging to your WandB account, use:
+`--wandb-key=[your WandB API key] --wandb-project=[some project name]
+--wandb-run-name=[optional: WandB run name (within the defined project)]`
+
+
+Results to expect
+-----------------
+You should see the following output (during the experiment) in your console:
+
+"""
+from ray.rllib.core.rl_module.rl_module import RLModuleSpec
+from ray.rllib.examples.envs.classes.multi_agent import MultiAgentStatelessCartPole
+from ray.rllib.examples.envs.classes.stateless_cartpole import StatelessCartPole
+from ray.rllib.examples.rl_modules.classes.lstm_containing_rlm import (
+    LSTMContainingRLModule,
+)
+from ray.rllib.examples.utils import (
+    add_rllib_example_script_args,
+    run_rllib_example_script_experiment,
+)
+from ray.tune.registry import get_trainable_cls, register_env
+
+parser = add_rllib_example_script_args(
+    default_reward=300.0,
+    default_timesteps=2000000,
+)
+
+
+if __name__ == "__main__":
+    args = parser.parse_args()
+
+    if args.num_agents == 0:
+        register_env("env", lambda cfg: StatelessCartPole())
+    else:
+        register_env("env", lambda cfg: MultiAgentStatelessCartPole(cfg))
+
+    base_config = (
+        get_trainable_cls(args.algo)
+        .get_default_config()
+        .environment(
+            env="env",
+            env_config={"num_agents": args.num_agents},
+        )
+        .training(
+            train_batch_size_per_learner=1024,
+            num_epochs=6,
+            lr=0.0009,
+            vf_loss_coeff=0.001,
+            entropy_coeff=0.0,
+        )
+        .rl_module(
+            # Plug-in our custom RLModule class.
+            rl_module_spec=RLModuleSpec(
+                module_class=LSTMContainingRLModule,
+                # Feel free to specify your own `model_config` settings below.
+                # The `model_config` defined here will be available inside your
+                # custom RLModule class through the `self.model_config`
+                # property.
+                model_config={
+                    "lstm_cell_size": 256,
+                    "dense_layers": [256, 256],
+                    "max_seq_len": 20,
+                },
+            ),
+        )
+    )
+
+    run_rllib_example_script_experiment(base_config, args)

rllib/examples/rl_modules/migrate_modelv2_to_new_api_stack_by_config.py

@@ -0,0 +1,76 @@
+from ray.rllib.algorithms.ppo import PPOConfig
+from ray.rllib.core import DEFAULT_POLICY_ID
+from ray.rllib.core.rl_module.rl_module import RLModuleSpec
+from ray.rllib.examples.rl_modules.classes.modelv2_to_rlm import ModelV2ToRLModule
+from ray.rllib.utils.metrics import (
+    ENV_RUNNER_RESULTS,
+    EPISODE_RETURN_MEAN,
+)
+
+if __name__ == "__main__":
+    # Configure an old stack default ModelV2.
+    config_old_stack = (
+        PPOConfig()
+        .api_stack(
+            enable_env_runner_and_connector_v2=False,
+            enable_rl_module_and_learner=False,
+        )
+        .environment("CartPole-v1")
+        .training(
+            lr=0.0003,
+            num_sgd_iter=6,
+            vf_loss_coeff=0.01,
+            # Change the ModelV2 settings a bit.
+            model={
+                "fcnet_hiddens": [32],
+                "fcnet_activation": "linear",
+                "use_lstm": True,
+                "vf_share_layers": True,
+            },
+        )
+    )
+
+    # Training with the (configured and wrapped) ModelV2.
+
+    # We change the original (old API stack) `config` into a new API stack one:
+    config_new_stack = (
+        config_old_stack.copy(copy_frozen=False)
+        .api_stack(
+            enable_rl_module_and_learner=True,
+            enable_env_runner_and_connector_v2=True,
+        )
+        .rl_module(
+            rl_module_spec=RLModuleSpec(
+                module_class=ModelV2ToRLModule,
+                model_config={
+                    "policy_id": DEFAULT_POLICY_ID,
+                    "old_api_stack_algo_config": config_old_stack,
+                    "max_seq_len": 20,
+                },
+            ),
+        )
+    )
+
+    # Build the new stack algo.
+    algo_new_stack = config_new_stack.build()
+
+    # Train until a higher return.
+    min_return_new_stack = 350.0
+    results = None
+    passed = False
+    for i in range(100):
+        results = algo_new_stack.train()
+        print(results)
+        if results[ENV_RUNNER_RESULTS][EPISODE_RETURN_MEAN] >= min_return_new_stack:
+            print(
+                f"Reached episode return of {min_return_new_stack} -> stopping "
+                "new API stack training."
+            )
+            passed = True
+            break
+
+    if not passed:
+        raise ValueError(
+            "Continuing training on the new stack did not succeed! Last return: "
+            f"{results[ENV_RUNNER_RESULTS][EPISODE_RETURN_MEAN]}"
+        )

rllib/examples/rl_modules/migrate_modelv2_to_new_api_stack_by_policy_checkpoint.py

@@ -0,0 +1,120 @@
+import os
+import pathlib
+
+import gymnasium as gym
+import numpy as np
+import torch
+
+from ray.rllib.algorithms.ppo import PPOConfig
+from ray.rllib.core.rl_module.rl_module import RLModuleSpec
+from ray.rllib.examples.rl_modules.classes.modelv2_to_rlm import ModelV2ToRLModule
+from ray.rllib.utils.metrics import (
+    ENV_RUNNER_RESULTS,
+    EPISODE_RETURN_MEAN,
+)
+from ray.rllib.utils.spaces.space_utils import batch
+
+if __name__ == "__main__":
+    # Configure and train an old stack default ModelV2.
+    config = (
+        PPOConfig()
+        # Old API stack.
+        .api_stack(
+            enable_env_runner_and_connector_v2=False,
+            enable_rl_module_and_learner=False,
+        )
+        .environment("CartPole-v1")
+        .training(
+            lr=0.0003,
+            num_epochs=6,
+            vf_loss_coeff=0.01,
+        )
+    )
+    algo_old_stack = config.build()
+
+    min_return_old_stack = 100.0
+    while True:
+        results = algo_old_stack.train()
+        print(results)
+        if results[ENV_RUNNER_RESULTS][EPISODE_RETURN_MEAN] >= min_return_old_stack:
+            print(
+                f"Reached episode return of {min_return_old_stack} -> stopping "
+                "old API stack training."
+            )
+            break
+
+    checkpoint = algo_old_stack.save()
+    policy_path = (
+        pathlib.Path(checkpoint.checkpoint.path) / "policies" / "default_policy"
+    )
+    assert policy_path.is_dir()
+    algo_old_stack.stop()
+
+    print("done")
+
+    # Move the old API stack (trained) ModelV2 into the new API stack's RLModule.
+    # Run a simple CartPole inference experiment. Render only if a display is
+    # available (local runs); skip rendering in headless CI to avoid pygame/X11 errors.
+    render_mode = "human" if os.environ.get("DISPLAY") else None
+    env = gym.make("CartPole-v1", render_mode=render_mode)
+    rl_module = ModelV2ToRLModule(
+        observation_space=env.observation_space,
+        action_space=env.action_space,
+        model_config={"policy_checkpoint_dir": policy_path},
+    )
+
+    obs, _ = env.reset()
+    env.render()
+    done = False
+    episode_return = 0.0
+    while not done:
+        output = rl_module.forward_inference({"obs": torch.from_numpy(batch([obs]))})
+        action_logits = output["action_dist_inputs"].detach().numpy()[0]
+        action = np.argmax(action_logits)
+        obs, reward, terminated, truncated, _ = env.step(action)
+        done = terminated or truncated
+        episode_return += reward
+        env.render()
+
+    print(f"Ran episode with trained ModelV2: return={episode_return}")
+
+    # Continue training with the (checkpointed) ModelV2.
+
+    # We change the original (old API stack) `config` into a new API stack one:
+    config = config.api_stack(
+        enable_rl_module_and_learner=True,
+        enable_env_runner_and_connector_v2=True,
+    ).rl_module(
+        rl_module_spec=RLModuleSpec(
+            module_class=ModelV2ToRLModule,
+            model_config={"policy_checkpoint_dir": policy_path},
+        ),
+    )
+
+    # Build the new stack algo.
+    algo_new_stack = config.build()
+
+    # Train until a higher return.
+    min_return_new_stack = 450.0
+    passed = False
+    for i in range(50):
+        results = algo_new_stack.train()
+        print(results)
+        # Make sure that the model's weights from the old API stack training
+        # were properly transferred to the new API RLModule wrapper. Thus, even
+        # after only one iteration of new stack training, we already expect the
+        # return to be higher than it was at the end of the old stack training.
+        assert results[ENV_RUNNER_RESULTS][EPISODE_RETURN_MEAN] >= min_return_old_stack
+        if results[ENV_RUNNER_RESULTS][EPISODE_RETURN_MEAN] >= min_return_new_stack:
+            print(
+                f"Reached episode return of {min_return_new_stack} -> stopping "
+                "new API stack training."
+            )
+            passed = True
+            break
+
+    if not passed:
+        raise ValueError(
+            "Continuing training on the new stack did not succeed! Last return: "
+            f"{results[ENV_RUNNER_RESULTS][EPISODE_RETURN_MEAN]}"
+        )

rllib/examples/rl_modules/pretraining_single_agent_training_multi_agent.py

@@ -0,0 +1,194 @@
+"""Example of running a single-agent pre-training followed with a multi-agent training.
+
+This examples `num_agents` agents each of them with its own `RLModule` that defines its
+policy. The first agent is pre-trained using a single-agent PPO algorithm. All agents
+are trained together in the main training run using a multi-agent PPO algorithm where
+the pre-trained module is used for the first agent.
+
+The environment is MultiAgentCartPole, in which there are n agents both policies.
+
+How to run this script
+----------------------
+`python [script file name].py --num-agents=2`
+
+For debugging, use the following additional command line options
+`--no-tune --num-env-runners=0`
+which should allow you to set breakpoints anywhere in the RLlib code and
+have the execution stop there for inspection and debugging.
+
+For logging to your WandB account, use:
+`--wandb-key=[your WandB API key] --wandb-project=[some project name]
+--wandb-run-name=[optional: WandB run name (within the defined project)]`
+
+
+Results to expect
+-----------------
+In the console output, you can see that the single-agent policy is first trained until
+the specified `--stop-reward-pretraining` value. For example, with the command line:
+`--num-agents=2 --stop-reward-pretraining=250.0 --stop-reward=250.0
+--stop-iters=3 --as-test`, you should get something like:
++-----------------------+------------+------+----------------+---------------------+
+| Trial name            | status     | iter | total time (s) | episode_return_mean |
+|                       |            |      |                |                     |
+|-----------------------+------------+------+----------------+---------------------+
+| PPO_CartPole-v1_00000 | TERMINATED |   16 |        25.6009 |               256.2 |
++-----------------------+------------+------+----------------+---------------------+
+
+Then, in the second experiment, where we run in a multi-agent setup with two policies
+("p0" from the single-agent checkpoint and "p1" randomly initialized), you can see that
+only "p0" immediately (after 1 iteration) reaches the same episode return as at the end
+of pretraining:
++----------------------------+------------+--------+------------------+------+
+| Trial name                 | status     |   iter |   total time (s) |   ts |
+|----------------------------+------------+--------+------------------+------+
+| PPO_multi-cart_6274d_00000 | TERMINATED |      1 |          2.71681 | 4000 |
++----------------------------+------------+--------+------------------+------+
++-------------------+-------------+-------------+
+|   combined return |   return p0 |   return p1 |
+|-------------------+-------------|-------------+
+|           451.625 |     433.125 |        18.5 |
++-------------------+-------------+-------------+
+"""
+from pathlib import Path
+
+import gymnasium as gym
+
+from ray.rllib.algorithms.ppo import PPOConfig
+from ray.rllib.algorithms.ppo.ppo_catalog import PPOCatalog
+from ray.rllib.algorithms.ppo.torch.ppo_torch_rl_module import PPOTorchRLModule
+from ray.rllib.core import (
+    COMPONENT_LEARNER,
+    COMPONENT_LEARNER_GROUP,
+    COMPONENT_RL_MODULE,
+    DEFAULT_MODULE_ID,
+)
+from ray.rllib.core.rl_module.default_model_config import DefaultModelConfig
+from ray.rllib.core.rl_module.multi_rl_module import MultiRLModuleSpec
+from ray.rllib.core.rl_module.rl_module import RLModuleSpec
+from ray.rllib.examples.envs.classes.multi_agent import MultiAgentCartPole
+from ray.rllib.examples.utils import (
+    add_rllib_example_script_args,
+    run_rllib_example_script_experiment,
+)
+from ray.tune import register_env
+
+# Read in common example script command line arguments.
+parser = add_rllib_example_script_args(
+    # Use less training steps for the main training run.
+    default_timesteps=50000,
+    default_reward=200.0,
+    default_iters=20,
+)
+parser.set_defaults(
+    checkpoint_freq=1,
+    checkpoint_at_end=True,
+)
+parser.add_argument(
+    "--stop-reward-pretraining",
+    type=float,
+    default=250.0,
+    help="The min. episode return to reach during pre-training.",
+)
+
+
+if __name__ == "__main__":
+
+    # Parse the command line arguments.
+    args = parser.parse_args()
+
+    # Ensure that the user has set the number of agents.
+    if args.num_agents == 0:
+        raise ValueError(
+            "This pre-training example script requires at least 1 agent. "
+            "Try setting the command line argument `--num-agents` to the "
+            "number of agents you want to use."
+        )
+
+    # Store the user's stopping criteria for the later training run.
+    stop_iters = args.stop_iters
+    stop_timesteps = args.stop_timesteps
+    stop_reward = args.stop_reward
+    num_agents = args.num_agents
+    as_test = args.as_test
+
+    # Override these criteria for the pre-training run.
+    args.stop_iters = 10000
+    args.stop_timesteps = 100000000
+    args.stop_reward = args.stop_reward_pretraining
+    args.num_agents = 0
+    args.as_test = False
+
+    # Define out pre-training single-agent algorithm. We will use the same module
+    # configuration for the pre-training and the training.
+    config = (
+        PPOConfig()
+        .environment("CartPole-v1")
+        .rl_module(
+            # Use a different number of hidden units for the pre-trained module.
+            model_config=DefaultModelConfig(fcnet_hiddens=[64]),
+        )
+    )
+
+    # Run the pre-training.
+    results = run_rllib_example_script_experiment(config, args, keep_ray_up=True)
+    # Get the checkpoint path.
+    module_chkpt_path = (
+        Path(results.get_best_result().checkpoint.path)
+        / COMPONENT_LEARNER_GROUP
+        / COMPONENT_LEARNER
+        / COMPONENT_RL_MODULE
+        / DEFAULT_MODULE_ID
+    )
+    assert module_chkpt_path.is_dir()
+
+    # Restore the user's stopping criteria for the training run.
+    args.stop_iters = stop_iters
+    args.stop_timesteps = stop_timesteps
+    args.stop_reward = stop_reward
+    args.num_agents = num_agents
+    args.as_test = as_test
+
+    # Create a new MultiRLModule using the pre-trained module for policy 0.
+    env = gym.make("CartPole-v1")
+    module_specs = {}
+    module_class = PPOTorchRLModule
+    for i in range(args.num_agents):
+        module_specs[f"p{i}"] = RLModuleSpec(
+            module_class=PPOTorchRLModule,
+            observation_space=env.observation_space,
+            action_space=env.action_space,
+            model_config=DefaultModelConfig(fcnet_hiddens=[32]),
+            catalog_class=PPOCatalog,
+        )
+
+    # Swap in the pre-trained module for policy 0.
+    module_specs["p0"] = RLModuleSpec(
+        module_class=PPOTorchRLModule,
+        observation_space=env.observation_space,
+        action_space=env.action_space,
+        model_config=DefaultModelConfig(fcnet_hiddens=[64]),
+        catalog_class=PPOCatalog,
+        # Note, we load here the module directly from the checkpoint.
+        load_state_path=module_chkpt_path,
+    )
+    multi_rl_module_spec = MultiRLModuleSpec(rl_module_specs=module_specs)
+
+    # Register our environment with tune if we use multiple agents.
+    register_env(
+        "multi-cart",
+        lambda _: MultiAgentCartPole(config={"num_agents": args.num_agents}),
+    )
+
+    # Configure the main (multi-agent) training run.
+    config = (
+        PPOConfig()
+        .environment("multi-cart")
+        .multi_agent(
+            policies={f"p{i}" for i in range(args.num_agents)},
+            policy_mapping_fn=lambda aid, eps, **kw: f"p{aid}",
+        )
+        .rl_module(rl_module_spec=multi_rl_module_spec)
+    )
+
+    # Run the main training run.
+    run_rllib_example_script_experiment(config, args)