ray.rllib.core.models.catalog.Catalog#

class ray.rllib.core.models.catalog.Catalog(observation_space: gymnasium.Space, action_space: gymnasium.Space, model_config_dict: dict, view_requirements=-1)[source]#

Describes the sub-module-architectures to be used in RLModules.

RLlib’s native RLModules get their Models from a Catalog object. By default, that Catalog builds the configs it has as attributes. This component was build to be hackable and extensible. You can inject custom components into RL Modules by overriding the build_xxx methods of this class. Note that it is recommended to write a custom RL Module for a single use-case. Modifications to Catalogs mostly make sense if you want to reuse the same Catalog for different RL Modules. For example if you have written a custom encoder and want to inject it into different RL Modules (e.g. for PPO, DQN, etc.). You can influence the decision tree that determines the sub-components by modifying Catalog._determine_components_hook.

Usage example:

# Define a custom catalog

import torch
import gymnasium as gym
from ray.rllib.core.models.configs import MLPHeadConfig
from ray.rllib.core.models.catalog import Catalog

class MyCatalog(Catalog):
    def __init__(
        self,
        observation_space: gym.Space,
        action_space: gym.Space,
        model_config_dict: dict,
    ):
        super().__init__(observation_space, action_space, model_config_dict)
        self.my_model_config = MLPHeadConfig(
            hidden_layer_dims=[64, 32],
            input_dims=[self.observation_space.shape[0]],
        )

    def build_my_head(self, framework: str):
        return self.my_model_config.build(framework=framework)

# With that, RLlib can build and use models from this catalog like this:
catalog = MyCatalog(gym.spaces.Box(0, 1), gym.spaces.Box(0, 1), {})
my_head = catalog.build_my_head(framework="torch")

# Make a call to the built model.
out = my_head(torch.Tensor([[1]]))

Methods

__init__

Initializes a Catalog with a default encoder config.

build_encoder

Builds the encoder.

get_action_dist_cls

Get the action distribution class.

get_preprocessor

Returns a suitable preprocessor for the given observation space.

get_tokenizer_config

Returns a tokenizer config for the given space.

Attributes

latent_dims

Returns the latent dimensions of the encoder.