wehub-resource-sync
655 lines
27 KiB
Python
655 lines
27 KiB
Python
import dataclasses
|
|
import enum
|
|
import functools
|
|
from typing import Optional
|
|
|
|
import gymnasium as gym
|
|
import numpy as np
|
|
import tree
|
|
from gymnasium.spaces import Box, Dict, Discrete, MultiDiscrete, Tuple
|
|
|
|
from ray._common.deprecation import DEPRECATED_VALUE, deprecation_warning
|
|
from ray.rllib.core.distribution.distribution import Distribution
|
|
from ray.rllib.core.models.base import Encoder
|
|
from ray.rllib.core.models.configs import (
|
|
CNNEncoderConfig,
|
|
MLPEncoderConfig,
|
|
ModelConfig,
|
|
RecurrentEncoderConfig,
|
|
)
|
|
from ray.rllib.core.rl_module.default_model_config import DefaultModelConfig
|
|
from ray.rllib.models.preprocessors import Preprocessor, get_preprocessor
|
|
from ray.rllib.models.utils import get_filter_config
|
|
from ray.rllib.utils.annotations import (
|
|
OverrideToImplementCustomLogic,
|
|
OverrideToImplementCustomLogic_CallToSuperRecommended,
|
|
)
|
|
from ray.rllib.utils.error import UnsupportedSpaceException
|
|
from ray.rllib.utils.spaces.simplex import Simplex
|
|
from ray.rllib.utils.spaces.space_utils import flatten_space, get_base_struct_from_space
|
|
|
|
|
|
class Catalog:
|
|
"""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
|
|
|
|
.. testcode::
|
|
|
|
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]]))
|
|
"""
|
|
|
|
# TODO (Sven): Add `framework` arg to c'tor and remove this arg from `build`
|
|
# methods. This way, we can already know in the c'tor of Catalog, what the exact
|
|
# action distibution objects are and thus what the output dims for e.g. a pi-head
|
|
# will be.
|
|
def __init__(
|
|
self,
|
|
observation_space: gym.Space,
|
|
action_space: gym.Space,
|
|
model_config_dict: dict,
|
|
# deprecated args.
|
|
view_requirements=DEPRECATED_VALUE,
|
|
):
|
|
"""Initializes a Catalog with a default encoder config.
|
|
|
|
Args:
|
|
observation_space: The observation space of the environment.
|
|
action_space: The action space of the environment.
|
|
model_config_dict: The model config that specifies things like hidden
|
|
dimensions and activations functions to use in this Catalog.
|
|
"""
|
|
if view_requirements != DEPRECATED_VALUE:
|
|
deprecation_warning(old="Catalog(view_requirements=..)", error=True)
|
|
|
|
# TODO (sven): The following logic won't be needed anymore, once we get rid of
|
|
# Catalogs entirely. We will assert directly inside the algo's DefaultRLModule
|
|
# class that the `model_config` is a DefaultModelConfig. Thus users won't be
|
|
# able to pass in partial config dicts into a default model (alternatively, we
|
|
# could automatically augment the user provided dict by the default config
|
|
# dataclass object only(!) for default modules).
|
|
if dataclasses.is_dataclass(model_config_dict):
|
|
model_config_dict = dataclasses.asdict(model_config_dict)
|
|
default_config = dataclasses.asdict(DefaultModelConfig())
|
|
# end: TODO
|
|
|
|
self.observation_space = observation_space
|
|
self.action_space = action_space
|
|
|
|
self._model_config_dict = default_config | model_config_dict
|
|
self._latent_dims = None
|
|
|
|
self._determine_components_hook()
|
|
|
|
@OverrideToImplementCustomLogic_CallToSuperRecommended
|
|
def _determine_components_hook(self):
|
|
"""Decision tree hook for subclasses to override.
|
|
|
|
By default, this method executes the decision tree that determines the
|
|
components that a Catalog builds. You can extend the components by overriding
|
|
this or by adding to the constructor of your subclass.
|
|
|
|
Override this method if you don't want to use the default components
|
|
determined here. If you want to use them but add additional components, you
|
|
should call `super()._determine_components()` at the beginning of your
|
|
implementation.
|
|
|
|
This makes it so that subclasses are not forced to create an encoder config
|
|
if the rest of their catalog is not dependent on it or if it breaks.
|
|
At the end of this method, an attribute `Catalog.latent_dims`
|
|
should be set so that heads can be built using that information.
|
|
"""
|
|
self._encoder_config = self._get_encoder_config(
|
|
observation_space=self.observation_space,
|
|
action_space=self.action_space,
|
|
model_config_dict=self._model_config_dict,
|
|
)
|
|
|
|
# Create a function that can be called when framework is known to retrieve the
|
|
# class type for action distributions
|
|
self._action_dist_class_fn = functools.partial(
|
|
self._get_dist_cls_from_action_space, action_space=self.action_space
|
|
)
|
|
|
|
# The dimensions of the latent vector that is output by the encoder and fed
|
|
# to the heads.
|
|
self.latent_dims = self._encoder_config.output_dims
|
|
|
|
@property
|
|
def latent_dims(self):
|
|
"""Returns the latent dimensions of the encoder.
|
|
|
|
This establishes an agreement between encoder and heads about the latent
|
|
dimensions. Encoders can be built to output a latent tensor with
|
|
`latent_dims` dimensions, and heads can be built with tensors of
|
|
`latent_dims` dimensions as inputs. This can be safely ignored if this
|
|
agreement is not needed in case of modifications to the Catalog.
|
|
|
|
Returns:
|
|
The latent dimensions of the encoder.
|
|
"""
|
|
return self._latent_dims
|
|
|
|
@latent_dims.setter
|
|
def latent_dims(self, value):
|
|
self._latent_dims = value
|
|
|
|
@OverrideToImplementCustomLogic
|
|
def build_encoder(self, framework: str) -> Encoder:
|
|
"""Builds the encoder.
|
|
|
|
By default, this method builds an encoder instance from Catalog._encoder_config.
|
|
|
|
You should override this if you want to use RLlib's default RL Modules but
|
|
only want to change the encoder. For example, if you want to use a custom
|
|
encoder, but want to use RLlib's default heads, action distribution and how
|
|
tensors are routed between them. If you want to have full control over the
|
|
RL Module, we recommend writing your own RL Module by inheriting from one of
|
|
RLlib's RL Modules instead.
|
|
|
|
Args:
|
|
framework: The framework to use. Either "torch" or "tf2".
|
|
|
|
Returns:
|
|
The encoder.
|
|
"""
|
|
assert hasattr(self, "_encoder_config"), (
|
|
"You must define a `Catalog._encoder_config` attribute in your Catalog "
|
|
"subclass or override the `Catalog.build_encoder` method. By default, "
|
|
"an encoder_config is created in the __post_init__ method."
|
|
)
|
|
return self._encoder_config.build(framework=framework)
|
|
|
|
@OverrideToImplementCustomLogic
|
|
def get_action_dist_cls(self, framework: str):
|
|
"""Get the action distribution class.
|
|
|
|
The default behavior is to get the action distribution from the
|
|
`Catalog._action_dist_class_fn`.
|
|
|
|
You should override this to have RLlib build your custom action
|
|
distribution instead of the default one. For example, if you don't want to
|
|
use RLlib's default RLModules with their default models, but only want to
|
|
change the distribution that Catalog returns.
|
|
|
|
Args:
|
|
framework: The framework to use. Either "torch" or "tf2".
|
|
|
|
Returns:
|
|
The action distribution.
|
|
"""
|
|
assert hasattr(self, "_action_dist_class_fn"), (
|
|
"You must define a `Catalog._action_dist_class_fn` attribute in your "
|
|
"Catalog subclass or override the `Catalog.action_dist_class_fn` method. "
|
|
"By default, an action_dist_class_fn is created in the __post_init__ "
|
|
"method."
|
|
)
|
|
return self._action_dist_class_fn(framework=framework)
|
|
|
|
@classmethod
|
|
def _get_encoder_config(
|
|
cls,
|
|
observation_space: gym.Space,
|
|
model_config_dict: dict,
|
|
action_space: gym.Space = None,
|
|
) -> ModelConfig:
|
|
"""Returns an EncoderConfig for the given input_space and model_config_dict.
|
|
|
|
Encoders are usually used in RLModules to transform the input space into a
|
|
latent space that is then fed to the heads. The returned EncoderConfig
|
|
objects correspond to the built-in Encoder classes in RLlib.
|
|
For example, for a simple 1D-Box input_space, RLlib offers an
|
|
MLPEncoder, hence this method returns the MLPEncoderConfig. You can overwrite
|
|
this method to produce specific EncoderConfigs for your custom Models.
|
|
|
|
The following input spaces lead to the following configs:
|
|
- 1D-Box: MLPEncoderConfig
|
|
- 3D-Box: CNNEncoderConfig
|
|
# TODO (Artur): Support more spaces here
|
|
# ...
|
|
|
|
Args:
|
|
observation_space: The observation space to use.
|
|
model_config_dict: The model config to use.
|
|
action_space: The action space to use if actions are to be encoded. This
|
|
is commonly the case for LSTM models.
|
|
|
|
Returns:
|
|
The encoder config.
|
|
"""
|
|
activation = model_config_dict["fcnet_activation"]
|
|
output_activation = model_config_dict["fcnet_activation"]
|
|
use_lstm = model_config_dict["use_lstm"]
|
|
|
|
if use_lstm:
|
|
encoder_config = RecurrentEncoderConfig(
|
|
input_dims=observation_space.shape,
|
|
recurrent_layer_type="lstm",
|
|
hidden_dim=model_config_dict["lstm_cell_size"],
|
|
hidden_weights_initializer=model_config_dict["lstm_kernel_initializer"],
|
|
hidden_weights_initializer_config=model_config_dict[
|
|
"lstm_kernel_initializer_kwargs"
|
|
],
|
|
hidden_bias_initializer=model_config_dict["lstm_bias_initializer"],
|
|
hidden_bias_initializer_config=model_config_dict[
|
|
"lstm_bias_initializer_kwargs"
|
|
],
|
|
batch_major=True,
|
|
num_layers=1,
|
|
tokenizer_config=cls.get_tokenizer_config(
|
|
observation_space,
|
|
model_config_dict,
|
|
),
|
|
)
|
|
else:
|
|
# TODO (Artur): Maybe check for original spaces here
|
|
# input_space is a 1D Box
|
|
if isinstance(observation_space, Box) and len(observation_space.shape) == 1:
|
|
# In order to guarantee backward compatability with old configs,
|
|
# we need to check if no latent dim was set and simply reuse the last
|
|
# fcnet hidden dim for that purpose.
|
|
hidden_layer_dims = model_config_dict["fcnet_hiddens"][:-1]
|
|
encoder_latent_dim = model_config_dict["fcnet_hiddens"][-1]
|
|
encoder_config = MLPEncoderConfig(
|
|
input_dims=observation_space.shape,
|
|
hidden_layer_dims=hidden_layer_dims,
|
|
hidden_layer_activation=activation,
|
|
hidden_layer_use_layernorm=model_config_dict.get(
|
|
"fcnet_use_layernorm", False
|
|
),
|
|
output_layer_use_layernorm=model_config_dict.get(
|
|
"fcnet_use_layernorm", False
|
|
),
|
|
hidden_layer_weights_initializer=model_config_dict[
|
|
"fcnet_kernel_initializer"
|
|
],
|
|
hidden_layer_weights_initializer_config=model_config_dict[
|
|
"fcnet_kernel_initializer_kwargs"
|
|
],
|
|
hidden_layer_bias_initializer=model_config_dict[
|
|
"fcnet_bias_initializer"
|
|
],
|
|
hidden_layer_bias_initializer_config=model_config_dict[
|
|
"fcnet_bias_initializer_kwargs"
|
|
],
|
|
output_layer_dim=encoder_latent_dim,
|
|
output_layer_activation=output_activation,
|
|
output_layer_weights_initializer=model_config_dict[
|
|
"fcnet_kernel_initializer"
|
|
],
|
|
output_layer_weights_initializer_config=model_config_dict[
|
|
"fcnet_kernel_initializer_kwargs"
|
|
],
|
|
output_layer_bias_initializer=model_config_dict[
|
|
"fcnet_bias_initializer"
|
|
],
|
|
output_layer_bias_initializer_config=model_config_dict[
|
|
"fcnet_bias_initializer_kwargs"
|
|
],
|
|
)
|
|
|
|
# input_space is a 3D Box
|
|
elif (
|
|
isinstance(observation_space, Box) and len(observation_space.shape) == 3
|
|
):
|
|
if not model_config_dict.get("conv_filters"):
|
|
model_config_dict["conv_filters"] = get_filter_config(
|
|
observation_space.shape
|
|
)
|
|
|
|
encoder_config = CNNEncoderConfig(
|
|
input_dims=observation_space.shape,
|
|
cnn_filter_specifiers=model_config_dict["conv_filters"],
|
|
cnn_activation=model_config_dict["conv_activation"],
|
|
cnn_kernel_initializer=model_config_dict["conv_kernel_initializer"],
|
|
cnn_kernel_initializer_config=model_config_dict[
|
|
"conv_kernel_initializer_kwargs"
|
|
],
|
|
cnn_bias_initializer=model_config_dict["conv_bias_initializer"],
|
|
cnn_bias_initializer_config=model_config_dict[
|
|
"conv_bias_initializer_kwargs"
|
|
],
|
|
)
|
|
# input_space is a 2D Box
|
|
elif (
|
|
isinstance(observation_space, Box) and len(observation_space.shape) == 2
|
|
):
|
|
# RLlib used to support 2D Box spaces by silently flattening them
|
|
raise ValueError(
|
|
f"No default encoder config for obs space={observation_space},"
|
|
f" lstm={use_lstm} found. 2D Box "
|
|
f"spaces are not supported. They should be either flattened to a "
|
|
f"1D Box space or enhanced to be a 3D box space."
|
|
)
|
|
# input_space is a possibly nested structure of spaces.
|
|
else:
|
|
# NestedModelConfig
|
|
raise ValueError(
|
|
f"No default encoder config for obs space={observation_space},"
|
|
f" lstm={use_lstm} found."
|
|
)
|
|
|
|
return encoder_config
|
|
|
|
@classmethod
|
|
@OverrideToImplementCustomLogic
|
|
def get_tokenizer_config(
|
|
cls,
|
|
observation_space: gym.Space,
|
|
model_config_dict: dict,
|
|
# deprecated args.
|
|
view_requirements=DEPRECATED_VALUE,
|
|
) -> ModelConfig:
|
|
"""Returns a tokenizer config for the given space.
|
|
|
|
This is useful for recurrent / transformer models that need to tokenize their
|
|
inputs. By default, RLlib uses the models supported by Catalog out of the box to
|
|
tokenize.
|
|
|
|
You should override this method if you want to change the custom tokenizer
|
|
inside current encoders that Catalog returns without providing the recurrent
|
|
network as a whole. For example, if you want to define some custom CNN layers
|
|
as a tokenizer for a recurrent encoder that already includes the recurrent
|
|
layers and handles the state.
|
|
|
|
Args:
|
|
observation_space: The observation space to use.
|
|
model_config_dict: The model config to use.
|
|
"""
|
|
if view_requirements != DEPRECATED_VALUE:
|
|
deprecation_warning(old="Catalog(view_requirements=..)", error=True)
|
|
|
|
return cls._get_encoder_config(
|
|
observation_space=observation_space,
|
|
# Use model_config_dict without flags that would end up in complex models
|
|
model_config_dict={
|
|
**model_config_dict,
|
|
**{"use_lstm": False, "use_attention": False},
|
|
},
|
|
)
|
|
|
|
@classmethod
|
|
def _get_dist_cls_from_action_space(
|
|
cls,
|
|
action_space: gym.Space,
|
|
*,
|
|
framework: Optional[str] = None,
|
|
) -> Distribution:
|
|
"""Returns a distribution class for the given action space.
|
|
|
|
You can get the required input dimension for the distribution by calling
|
|
`action_dict_cls.required_input_dim(action_space)`
|
|
on the retrieved class. This is useful, because the Catalog needs to find out
|
|
about the required input dimension for the distribution before the model that
|
|
outputs these inputs is configured.
|
|
|
|
Args:
|
|
action_space: Action space of the target gym env.
|
|
framework: The framework to use.
|
|
|
|
Returns:
|
|
The distribution class for the given action space.
|
|
"""
|
|
# If no framework provided, return no action distribution class (None).
|
|
if framework is None:
|
|
return None
|
|
# This method is structured in two steps:
|
|
# Firstly, construct a dictionary containing the available distribution classes.
|
|
# Secondly, return the correct distribution class for the given action space.
|
|
|
|
# Step 1: Construct the dictionary.
|
|
|
|
class DistEnum(enum.Enum):
|
|
Categorical = "Categorical"
|
|
DiagGaussian = "Gaussian"
|
|
Deterministic = "Deterministic"
|
|
MultiDistribution = "MultiDistribution"
|
|
MultiCategorical = "MultiCategorical"
|
|
|
|
if framework == "torch":
|
|
from ray.rllib.core.distribution.torch.torch_distribution import (
|
|
TorchCategorical,
|
|
TorchDeterministic,
|
|
TorchDiagGaussian,
|
|
)
|
|
|
|
distribution_dicts = {
|
|
DistEnum.Deterministic: TorchDeterministic,
|
|
DistEnum.DiagGaussian: TorchDiagGaussian,
|
|
DistEnum.Categorical: TorchCategorical,
|
|
}
|
|
else:
|
|
raise ValueError(
|
|
f"Unknown framework: {framework}. Only 'torch' and 'tf2' are "
|
|
"supported for RLModule Catalogs."
|
|
)
|
|
|
|
# Only add a MultiAction distribution class to the dict if we can compute its
|
|
# components (we need a Tuple/Dict space for this).
|
|
if isinstance(action_space, (Tuple, Dict)):
|
|
partial_multi_action_distribution_cls = _multi_action_dist_partial_helper(
|
|
catalog_cls=cls,
|
|
action_space=action_space,
|
|
framework=framework,
|
|
)
|
|
|
|
distribution_dicts[
|
|
DistEnum.MultiDistribution
|
|
] = partial_multi_action_distribution_cls
|
|
|
|
# Only add a MultiCategorical distribution class to the dict if we can compute
|
|
# its components (we need a MultiDiscrete space for this).
|
|
if isinstance(action_space, MultiDiscrete):
|
|
partial_multi_categorical_distribution_cls = (
|
|
_multi_categorical_dist_partial_helper(
|
|
action_space=action_space,
|
|
framework=framework,
|
|
)
|
|
)
|
|
|
|
distribution_dicts[
|
|
DistEnum.MultiCategorical
|
|
] = partial_multi_categorical_distribution_cls
|
|
|
|
# Step 2: Return the correct distribution class for the given action space.
|
|
|
|
# Box space -> DiagGaussian OR Deterministic.
|
|
if isinstance(action_space, Box):
|
|
if action_space.dtype.char in np.typecodes["AllInteger"]:
|
|
raise ValueError(
|
|
"Box(..., `int`) action spaces are not supported. "
|
|
"Use MultiDiscrete or Box(..., `float`)."
|
|
)
|
|
else:
|
|
if len(action_space.shape) > 1:
|
|
raise UnsupportedSpaceException(
|
|
f"Action space has multiple dimensions {action_space.shape}. "
|
|
f"Consider reshaping this into a single dimension, using a "
|
|
f"custom action distribution, using a Tuple action space, "
|
|
f"or the multi-agent API."
|
|
)
|
|
return distribution_dicts[DistEnum.DiagGaussian]
|
|
|
|
# Discrete Space -> Categorical.
|
|
elif isinstance(action_space, Discrete):
|
|
return distribution_dicts[DistEnum.Categorical]
|
|
|
|
# Tuple/Dict Spaces -> MultiAction.
|
|
elif isinstance(action_space, (Tuple, Dict)):
|
|
return distribution_dicts[DistEnum.MultiDistribution]
|
|
|
|
# Simplex -> Dirichlet.
|
|
elif isinstance(action_space, Simplex):
|
|
# TODO(Artur): Supported Simplex (in torch).
|
|
raise NotImplementedError("Simplex action space not yet supported.")
|
|
|
|
# MultiDiscrete -> MultiCategorical.
|
|
elif isinstance(action_space, MultiDiscrete):
|
|
return distribution_dicts[DistEnum.MultiCategorical]
|
|
|
|
# Unknown type -> Error.
|
|
else:
|
|
raise NotImplementedError(f"Unsupported action space: `{action_space}`")
|
|
|
|
@staticmethod
|
|
def get_preprocessor(observation_space: gym.Space, **kwargs) -> Preprocessor:
|
|
"""Returns a suitable preprocessor for the given observation space.
|
|
|
|
Args:
|
|
observation_space: The input observation space.
|
|
**kwargs: Forward-compatible kwargs.
|
|
|
|
Returns:
|
|
preprocessor: Preprocessor for the observations.
|
|
"""
|
|
# TODO(Artur): Since preprocessors have long been @PublicAPI with the options
|
|
# kwarg as part of their constructor, we fade out support for this,
|
|
# beginning with this entrypoint.
|
|
# Next, we should deprecate the `options` kwarg from the Preprocessor itself,
|
|
# after deprecating the old catalog and other components that still pass this.
|
|
options = kwargs.get("options", {})
|
|
if options:
|
|
deprecation_warning(
|
|
old="get_preprocessor_for_space(..., options={...})",
|
|
help="Override `Catalog.get_preprocessor()` "
|
|
"in order to implement custom behaviour.",
|
|
error=False,
|
|
)
|
|
|
|
if options.get("custom_preprocessor"):
|
|
deprecation_warning(
|
|
old="model_config['custom_preprocessor']",
|
|
help="Custom preprocessors are deprecated, "
|
|
"since they sometimes conflict with the built-in "
|
|
"preprocessors for handling complex observation spaces. "
|
|
"Please use wrapper classes around your environment "
|
|
"instead.",
|
|
error=True,
|
|
)
|
|
else:
|
|
# TODO(Artur): Inline the get_preprocessor() call here once we have
|
|
# deprecated the old model catalog.
|
|
cls = get_preprocessor(observation_space)
|
|
prep = cls(observation_space, options)
|
|
return prep
|
|
|
|
|
|
def _multi_action_dist_partial_helper(
|
|
catalog_cls: "Catalog", action_space: gym.Space, framework: str
|
|
) -> Distribution:
|
|
"""Helper method to get a partial of a MultiActionDistribution.
|
|
|
|
This is useful for when we want to create MultiActionDistributions from
|
|
logits only (!) later, but know the action space now already.
|
|
|
|
Args:
|
|
catalog_cls: The ModelCatalog class to use.
|
|
action_space: The action space to get the child distribution classes for.
|
|
framework: The framework to use.
|
|
|
|
Returns:
|
|
A partial of the TorchMultiActionDistribution class.
|
|
"""
|
|
action_space_struct = get_base_struct_from_space(action_space)
|
|
flat_action_space = flatten_space(action_space)
|
|
child_distribution_cls_struct = tree.map_structure(
|
|
lambda s: catalog_cls._get_dist_cls_from_action_space(
|
|
action_space=s,
|
|
framework=framework,
|
|
),
|
|
action_space_struct,
|
|
)
|
|
flat_distribution_clses = tree.flatten(child_distribution_cls_struct)
|
|
|
|
logit_lens = [
|
|
int(dist_cls.required_input_dim(space))
|
|
for dist_cls, space in zip(flat_distribution_clses, flat_action_space)
|
|
]
|
|
|
|
if framework == "torch":
|
|
from ray.rllib.core.distribution.torch.torch_distribution import (
|
|
TorchMultiDistribution,
|
|
)
|
|
|
|
multi_action_dist_cls = TorchMultiDistribution
|
|
else:
|
|
raise ValueError(f"Unsupported framework: {framework}")
|
|
|
|
partial_dist_cls = multi_action_dist_cls.get_partial_dist_cls(
|
|
space=action_space,
|
|
child_distribution_cls_struct=child_distribution_cls_struct,
|
|
input_lens=logit_lens,
|
|
)
|
|
return partial_dist_cls
|
|
|
|
|
|
def _multi_categorical_dist_partial_helper(
|
|
action_space: gym.Space, framework: str
|
|
) -> Distribution:
|
|
"""Helper method to get a partial of a MultiCategorical Distribution.
|
|
|
|
This is useful for when we want to create MultiCategorical Distribution from
|
|
logits only (!) later, but know the action space now already.
|
|
|
|
Args:
|
|
action_space: The action space to get the child distribution classes for.
|
|
framework: The framework to use.
|
|
|
|
Returns:
|
|
A partial of the MultiCategorical class.
|
|
"""
|
|
|
|
if framework == "torch":
|
|
from ray.rllib.core.distribution.torch.torch_distribution import (
|
|
TorchMultiCategorical,
|
|
)
|
|
|
|
multi_categorical_dist_cls = TorchMultiCategorical
|
|
else:
|
|
raise ValueError(f"Unsupported framework: {framework}")
|
|
|
|
partial_dist_cls = multi_categorical_dist_cls.get_partial_dist_cls(
|
|
space=action_space, input_lens=list(action_space.nvec)
|
|
)
|
|
|
|
return partial_dist_cls
|