wehub-resource-sync
1194 lines
50 KiB
Python
1194 lines
50 KiB
Python
import abc
|
|
import functools
|
|
from dataclasses import dataclass, field
|
|
from typing import TYPE_CHECKING, Callable, Dict, List, Optional, Tuple, Union
|
|
|
|
import numpy as np
|
|
|
|
from ray.rllib.models.torch.misc import (
|
|
same_padding,
|
|
same_padding_transpose_after_stride,
|
|
valid_padding,
|
|
)
|
|
from ray.rllib.models.utils import get_activation_fn, get_initializer_fn
|
|
from ray.rllib.utils.annotations import DeveloperAPI
|
|
|
|
if TYPE_CHECKING:
|
|
from ray.rllib.core.models.base import Encoder, Model
|
|
|
|
|
|
@DeveloperAPI
|
|
def _framework_implemented(torch: bool = True, tf2: bool = True):
|
|
"""Decorator to check if a model was implemented in a framework.
|
|
|
|
Args:
|
|
torch: Whether we can build this model with torch.
|
|
tf2: Whether we can build this model with tf2.
|
|
|
|
Returns:
|
|
The decorated function.
|
|
|
|
Raises:
|
|
ValueError: If the framework is not available to build.
|
|
"""
|
|
accepted = []
|
|
if torch:
|
|
accepted.append("torch")
|
|
if tf2:
|
|
accepted.append("tf2")
|
|
|
|
def decorator(fn: Callable) -> Callable:
|
|
@functools.wraps(fn)
|
|
def checked_build(self, framework, **kwargs):
|
|
if framework not in accepted:
|
|
raise ValueError(
|
|
f"This config does not support framework "
|
|
f"{framework}. Only frameworks in {accepted} are "
|
|
f"supported."
|
|
)
|
|
return fn(self, framework, **kwargs)
|
|
|
|
return checked_build
|
|
|
|
return decorator
|
|
|
|
|
|
@DeveloperAPI
|
|
@dataclass
|
|
class ModelConfig(abc.ABC):
|
|
"""Base class for configuring a `Model` instance.
|
|
|
|
ModelConfigs are DL framework-agnostic.
|
|
A `Model` (as a sub-component of an `RLModule`) is built via calling the
|
|
respective ModelConfig's `build()` method.
|
|
RLModules build their sub-components this way after receiving one or more
|
|
`ModelConfig` instances from a Catalog object.
|
|
|
|
However, `ModelConfig` is not restricted to be used only with Catalog or RLModules.
|
|
Usage examples can be found in the individual Model classes', e.g.
|
|
see `ray.rllib.core.models.configs::MLPHeadConfig`.
|
|
|
|
Attributes:
|
|
input_dims: The input dimensions of the network
|
|
always_check_shapes: Whether to always check the inputs and outputs of the
|
|
model for the specifications. Input specifications are checked on failed
|
|
forward passes of the model regardless of this flag. If this flag is set
|
|
to `True`, inputs and outputs are checked on every call. This leads to
|
|
a slow-down and should only be used for debugging.
|
|
"""
|
|
|
|
input_dims: Union[List[int], Tuple[int, ...]] = None
|
|
always_check_shapes: bool = False
|
|
|
|
@abc.abstractmethod
|
|
def build(self, framework: str):
|
|
"""Builds the model.
|
|
|
|
Args:
|
|
framework: The framework to use for building the model.
|
|
"""
|
|
raise NotImplementedError
|
|
|
|
@property
|
|
def output_dims(self) -> Optional[Tuple[int, ...]]:
|
|
"""Read-only `output_dims` are inferred automatically from other settings."""
|
|
return None
|
|
|
|
|
|
@DeveloperAPI
|
|
@dataclass
|
|
class _MLPConfig(ModelConfig):
|
|
"""Generic configuration class for multi-layer-perceptron based Model classes.
|
|
|
|
`output_dims` is reached by either the provided `output_layer_dim` setting (int) OR
|
|
by the last entry of `hidden_layer_dims`. In the latter case, no special output
|
|
layer is added and all layers in the stack behave exactly the same. If
|
|
`output_layer_dim` is provided, users might also change this last layer's
|
|
activation (`output_layer_activation`) and its bias setting
|
|
(`output_layer_use_bias`).
|
|
|
|
This is a private class as users should not configure their models directly
|
|
through this class, but use one of the sub-classes, e.g. `MLPHeadConfig` or
|
|
`MLPEncoderConfig`.
|
|
|
|
Attributes:
|
|
input_dims: A 1D tensor indicating the input dimension, e.g. `[32]`.
|
|
hidden_layer_dims: The sizes of the hidden layers. If an empty list,
|
|
`output_layer_dim` must be provided (int) and only a single layer will be
|
|
built.
|
|
hidden_layer_use_bias: Whether to use bias on all dense layers in the network
|
|
(excluding a possible separate output layer defined by `output_layer_dim`).
|
|
hidden_layer_activation: The activation function to use after each layer (
|
|
except for the output). The default activation for hidden layers is "relu".
|
|
hidden_layer_use_layernorm: Whether to insert a LayerNorm functionality
|
|
in between each hidden layer's output and its activation.
|
|
hidden_layer_weights_initializer: The initializer function or class to use for
|
|
weight initialization in the hidden layers. If `None` the default
|
|
initializer of the respective dense layer of a framework (`"torch"` or
|
|
`"tf2"`) is used. Note, all initializers defined in the framework `"tf2`)
|
|
are allowed. For `"torch"` only the in-place initializers, i.e. ending with
|
|
an underscore "_" are allowed.
|
|
hidden_layer_weights_initializer_config: Configuration to pass into the
|
|
initializer defined in `hidden_layer_weights_initializer`.
|
|
hidden_layer_bias_initializer: The initializer function or class to use for
|
|
bias initialization in the hidden layers. If `None` the default initializer
|
|
of the respective dense layer of a framework (`"torch"` or `"tf2"`) is used.
|
|
Note, all initializers defined in the framework `"tf2`) are allowed. For
|
|
`"torch"` only the in-place initializers, i.e. ending with an underscore "_"
|
|
are allowed.
|
|
hidden_layer_bias_initializer_config: Configuration to pass into the
|
|
initializer defined in `hidden_layer_bias_initializer`.
|
|
output_layer_dim: An int indicating the size of the output layer. This may be
|
|
set to `None` in case no extra output layer should be built and only the
|
|
layers specified by `hidden_layer_dims` will be part of the network.
|
|
output_layer_use_layernorm: Whether to insert a LayerNorm after the output
|
|
layer (before its activation). Only applies when `output_layer_dim` is set.
|
|
output_layer_use_bias: Whether to use bias on the separate output layer, if any.
|
|
output_layer_activation: The activation function to use for the output layer,
|
|
if any. The default activation for the output layer, if any, is "linear",
|
|
meaning no activation.
|
|
output_layer_weights_initializer: The initializer function or class to use for
|
|
weight initialization in the output layers. If `None` the default
|
|
initializer of the respective dense layer of a framework (`"torch"` or `
|
|
"tf2"`) is used. Note, all initializers defined in the framework `"tf2`) are
|
|
allowed. For `"torch"` only the in-place initializers, i.e. ending with an
|
|
underscore "_" are allowed.
|
|
output_layer_weights_initializer_config: Configuration to pass into the
|
|
initializer defined in `output_layer_weights_initializer`.
|
|
output_layer_bias_initializer: The initializer function or class to use for
|
|
bias initialization in the output layers. If `None` the default initializer
|
|
of the respective dense layer of a framework (`"torch"` or `"tf2"`) is used.
|
|
For `"torch"` only the in-place initializers, i.e. ending with an underscore
|
|
"_" are allowed.
|
|
output_layer_bias_initializer_config: Configuration to pass into the
|
|
initializer defined in `output_layer_bias_initializer`.
|
|
clip_log_std: If log std should be clipped by `log_std_clip_param`. This applies
|
|
only to the action distribution parameters that encode the log standard
|
|
deviation of a `DiagGaussian` distribution.
|
|
log_std_clip_param: The clipping parameter for the log std, if clipping should
|
|
be applied - i.e. `clip_log_std=True`. The default value is 20, i.e. log
|
|
stds are clipped in between -20 and 20.
|
|
"""
|
|
|
|
hidden_layer_dims: Union[List[int], Tuple[int, ...]] = (256, 256)
|
|
hidden_layer_use_bias: bool = True
|
|
hidden_layer_activation: str = "relu"
|
|
hidden_layer_use_layernorm: bool = False
|
|
hidden_layer_weights_initializer: Optional[Union[str, Callable]] = None
|
|
hidden_layer_weights_initializer_config: Optional[Dict] = None
|
|
hidden_layer_bias_initializer: Optional[Union[str, Callable]] = None
|
|
hidden_layer_bias_initializer_config: Optional[Dict] = None
|
|
|
|
# Optional last output layer with - possibly - different activation and use_bias
|
|
# settings.
|
|
output_layer_dim: Optional[int] = None
|
|
output_layer_use_layernorm: bool = False
|
|
output_layer_use_bias: bool = True
|
|
output_layer_activation: str = "linear"
|
|
output_layer_weights_initializer: Optional[Union[str, Callable]] = None
|
|
output_layer_weights_initializer_config: Optional[Dict] = None
|
|
output_layer_bias_initializer: Optional[Union[str, Callable]] = None
|
|
output_layer_bias_initializer_config: Optional[Dict] = None
|
|
|
|
# Optional clipping of log standard deviation.
|
|
clip_log_std: bool = False
|
|
# Optional clip parameter for the log standard deviation.
|
|
log_std_clip_param: float = 20.0
|
|
|
|
@property
|
|
def output_dims(self):
|
|
if self.output_layer_dim is None and not self.hidden_layer_dims:
|
|
raise ValueError(
|
|
"If `output_layer_dim` is None, you must specify at least one hidden "
|
|
"layer dim, e.g. `hidden_layer_dims=[32]`!"
|
|
)
|
|
|
|
# Infer `output_dims` automatically.
|
|
return (int(self.output_layer_dim or self.hidden_layer_dims[-1]),)
|
|
|
|
def _validate(self, framework: str = "torch"):
|
|
"""Makes sure that settings are valid."""
|
|
if self.input_dims is not None and len(self.input_dims) != 1:
|
|
raise ValueError(
|
|
f"`input_dims` ({self.input_dims}) of MLPConfig must be 1D, "
|
|
"e.g. `[32]`!"
|
|
)
|
|
if len(self.output_dims) != 1:
|
|
raise ValueError(
|
|
f"`output_dims` ({self.output_dims}) of _MLPConfig must be "
|
|
"1D, e.g. `[32]`! This is an inferred value, hence other settings might"
|
|
" be wrong."
|
|
)
|
|
if self.log_std_clip_param is None:
|
|
raise ValueError(
|
|
"`log_std_clip_param` of _MLPConfig must be a float value, but is "
|
|
"`None`."
|
|
)
|
|
|
|
# Call these already here to catch errors early on.
|
|
get_activation_fn(self.hidden_layer_activation, framework=framework)
|
|
get_activation_fn(self.output_layer_activation, framework=framework)
|
|
get_initializer_fn(self.hidden_layer_weights_initializer, framework=framework)
|
|
get_initializer_fn(self.hidden_layer_bias_initializer, framework=framework)
|
|
get_initializer_fn(self.output_layer_weights_initializer, framework=framework)
|
|
get_initializer_fn(self.output_layer_bias_initializer, framework=framework)
|
|
|
|
|
|
@DeveloperAPI
|
|
@dataclass
|
|
class MLPHeadConfig(_MLPConfig):
|
|
"""Configuration for an MLP head.
|
|
|
|
See _MLPConfig for usage details.
|
|
|
|
Example:
|
|
|
|
.. testcode::
|
|
|
|
# Configuration:
|
|
config = MLPHeadConfig(
|
|
input_dims=[4], # must be 1D tensor
|
|
hidden_layer_dims=[8, 8],
|
|
hidden_layer_activation="relu",
|
|
hidden_layer_use_layernorm=False,
|
|
# final output layer with no activation (linear)
|
|
output_layer_dim=2,
|
|
output_layer_activation="linear",
|
|
)
|
|
model = config.build(framework="tf2")
|
|
|
|
# Resulting stack in pseudocode:
|
|
# Linear(4, 8, bias=True)
|
|
# ReLU()
|
|
# Linear(8, 8, bias=True)
|
|
# ReLU()
|
|
# Linear(8, 2, bias=True)
|
|
|
|
Example:
|
|
|
|
.. testcode::
|
|
|
|
# Configuration:
|
|
config = MLPHeadConfig(
|
|
input_dims=[2],
|
|
hidden_layer_dims=[10, 4],
|
|
hidden_layer_activation="silu",
|
|
hidden_layer_use_layernorm=True,
|
|
hidden_layer_use_bias=False,
|
|
# Initializer for `framework="torch"`.
|
|
hidden_layer_weights_initializer="xavier_normal_",
|
|
hidden_layer_weights_initializer_config={"gain": 0.8},
|
|
# No final output layer (use last dim in `hidden_layer_dims`
|
|
# as the size of the last layer in the stack).
|
|
output_layer_dim=None,
|
|
)
|
|
model = config.build(framework="torch")
|
|
|
|
# Resulting stack in pseudocode:
|
|
# Linear(2, 10, bias=False)
|
|
# LayerNorm((10,)) # layer norm always before activation
|
|
# SiLU()
|
|
# Linear(10, 4, bias=False)
|
|
# LayerNorm((4,)) # layer norm always before activation
|
|
# SiLU()
|
|
"""
|
|
|
|
@_framework_implemented()
|
|
def build(self, framework: str = "torch") -> "Model":
|
|
self._validate(framework=framework)
|
|
|
|
if framework == "torch":
|
|
from ray.rllib.core.models.torch.heads import TorchMLPHead
|
|
|
|
return TorchMLPHead(self)
|
|
|
|
|
|
@DeveloperAPI
|
|
@dataclass
|
|
class FreeLogStdMLPHeadConfig(_MLPConfig):
|
|
"""Configuration for an MLPHead with a floating second half of outputs.
|
|
|
|
This model can be useful together with Gaussian Distributions.
|
|
This gaussian distribution would be conditioned as follows:
|
|
- The first half of outputs from this model can be used as
|
|
state-dependent means when conditioning a gaussian distribution
|
|
- The second half are floating free biases that can be used as
|
|
state-independent standard deviations to condition a gaussian distribution.
|
|
The mean values are produced by an MLPHead, while the standard
|
|
deviations are added as floating free biases from a single 1D trainable variable
|
|
(not dependent on the net's inputs).
|
|
|
|
The output dimensions of the configured MLPHeadConfig must be even and are
|
|
divided by two to gain the output dimensions of each the mean-net and the
|
|
free std-variable.
|
|
|
|
Example:
|
|
.. testcode::
|
|
:skipif: True
|
|
|
|
# Configuration:
|
|
config = FreeLogStdMLPHeadConfig(
|
|
input_dims=[2],
|
|
hidden_layer_dims=[16],
|
|
hidden_layer_activation=None,
|
|
hidden_layer_use_layernorm=False,
|
|
hidden_layer_use_bias=True,
|
|
output_layer_dim=8, # <- this must be an even size
|
|
output_layer_use_bias=True,
|
|
)
|
|
model = config.build(framework="tf2")
|
|
|
|
# Resulting stack in pseudocode:
|
|
# Linear(2, 16, bias=True)
|
|
# Linear(8, 8, bias=True) # 16 / 2 = 8 -> 8 nodes for the mean
|
|
# Extra variable:
|
|
# Tensor((8,), float32) # for the free (observation independent) std outputs
|
|
|
|
Example:
|
|
.. testcode::
|
|
:skipif: True
|
|
|
|
# Configuration:
|
|
config = FreeLogStdMLPHeadConfig(
|
|
input_dims=[2],
|
|
hidden_layer_dims=[31, 100], # <- last idx must be an even size
|
|
hidden_layer_activation="relu",
|
|
hidden_layer_use_layernorm=False,
|
|
hidden_layer_use_bias=False,
|
|
output_layer_dim=None, # use the last hidden layer as output layer
|
|
)
|
|
model = config.build(framework="torch")
|
|
|
|
# Resulting stack in pseudocode:
|
|
# Linear(2, 31, bias=False)
|
|
# ReLu()
|
|
# Linear(31, 50, bias=False) # 100 / 2 = 50 -> 50 nodes for the mean
|
|
# ReLu()
|
|
# Extra variable:
|
|
# Tensor((50,), float32) # for the free (observation independent) std outputs
|
|
"""
|
|
|
|
def _validate(self, framework: str = "torch"):
|
|
if len(self.output_dims) > 1 or self.output_dims[0] % 2 == 1:
|
|
raise ValueError(
|
|
f"`output_layer_dim` ({self.ouput_layer_dim}) or the last value in "
|
|
f"`hidden_layer_dims` ({self.hidden_layer_dims}) of a "
|
|
"FreeLogStdMLPHeadConfig must be an even int (dividable by 2), "
|
|
"e.g. `output_layer_dim=8` or `hidden_layer_dims=[133, 128]`!"
|
|
)
|
|
|
|
@_framework_implemented()
|
|
def build(self, framework: str = "torch") -> "Model":
|
|
self._validate(framework=framework)
|
|
|
|
if framework == "torch":
|
|
from ray.rllib.core.models.torch.heads import TorchFreeLogStdMLPHead
|
|
|
|
return TorchFreeLogStdMLPHead(self)
|
|
|
|
|
|
@DeveloperAPI
|
|
@dataclass
|
|
class CNNTransposeHeadConfig(ModelConfig):
|
|
"""Configuration for a convolutional transpose head (decoder) network.
|
|
|
|
The configured Model transforms 1D-observations into an image space.
|
|
The stack of layers is composed of an initial Dense layer, followed by a sequence
|
|
of Conv2DTranspose layers.
|
|
`input_dims` describes the shape of the (1D) input tensor,
|
|
`initial_image_dims` describes the input into the first Conv2DTranspose
|
|
layer, where the translation from `input_dim` to `initial_image_dims` is done
|
|
via the initial Dense layer (w/o activation, w/o layer-norm, and w/ bias).
|
|
|
|
Beyond that, each layer specified by `cnn_transpose_filter_specifiers`
|
|
is followed by an activation function according to `cnn_transpose_activation`.
|
|
|
|
`output_dims` is reached after the final Conv2DTranspose layer.
|
|
Not that the last Conv2DTranspose layer is never activated and never layer-norm'd
|
|
regardless of the other settings.
|
|
|
|
An example for a single conv2d operation is as follows:
|
|
Input "image" is (4, 4, 24) (not yet strided), padding is "same", stride=2,
|
|
kernel=5.
|
|
|
|
First, the input "image" is strided (with stride=2):
|
|
|
|
Input image (4x4 (x24)):
|
|
A B C D
|
|
E F G H
|
|
I J K L
|
|
M N O P
|
|
|
|
Stride with stride=2 -> (7x7 (x24))
|
|
A 0 B 0 C 0 D
|
|
0 0 0 0 0 0 0
|
|
E 0 F 0 G 0 H
|
|
0 0 0 0 0 0 0
|
|
I 0 J 0 K 0 L
|
|
0 0 0 0 0 0 0
|
|
M 0 N 0 O 0 P
|
|
|
|
Then this strided "image" (strided_size=7x7) is padded (exact padding values will be
|
|
computed by the model):
|
|
|
|
Padding -> (left=3, right=2, top=3, bottom=2)
|
|
|
|
0 0 0 0 0 0 0 0 0 0 0 0
|
|
0 0 0 0 0 0 0 0 0 0 0 0
|
|
0 0 0 0 0 0 0 0 0 0 0 0
|
|
0 0 0 A 0 B 0 C 0 D 0 0
|
|
0 0 0 0 0 0 0 0 0 0 0 0
|
|
0 0 0 E 0 F 0 G 0 H 0 0
|
|
0 0 0 0 0 0 0 0 0 0 0 0
|
|
0 0 0 I 0 J 0 K 0 L 0 0
|
|
0 0 0 0 0 0 0 0 0 0 0 0
|
|
0 0 0 M 0 N 0 O 0 P 0 0
|
|
0 0 0 0 0 0 0 0 0 0 0 0
|
|
0 0 0 0 0 0 0 0 0 0 0 0
|
|
|
|
Then deconvolution with kernel=5 yields an output "image" of 8x8 (x num output
|
|
filters).
|
|
|
|
Attributes:
|
|
input_dims: The input dimensions of the network. This must be a 1D tensor.
|
|
initial_image_dims: The shape of the input to the first
|
|
Conv2DTranspose layer. We will make sure the input is transformed to
|
|
these dims via a preceding initial Dense layer, followed by a reshape,
|
|
before entering the Conv2DTranspose stack.
|
|
initial_dense_weights_initializer: The initializer function or class to use for
|
|
weight initialization in the initial dense layer. If `None` the default
|
|
initializer of the respective dense layer of a framework (`"torch"` or
|
|
`"tf2"`) is used. Note, all initializers defined in the framework `"tf2`)
|
|
are allowed. For `"torch"` only the in-place initializers, i.e. ending with
|
|
an underscore "_" are allowed.
|
|
initial_dense_weights_initializer_config: Configuration to pass into the
|
|
initializer defined in `initial_dense_weights_initializer`.
|
|
initial_dense_bias_initializer: The initializer function or class to use for
|
|
bias initialization in the initial dense layer. If `None` the default
|
|
initializer of the respective CNN layer of a framework (`"torch"` or `"tf2"`
|
|
) is used. For `"torch"` only the in-place initializers, i.e. ending with an
|
|
underscore "_" are allowed.
|
|
initial_dense_bias_initializer_config: Configuration to pass into the
|
|
initializer defined in `initial_dense_bias_initializer`.
|
|
cnn_transpose_filter_specifiers: A list of lists, where each element of an inner
|
|
list contains elements of the form
|
|
`[number of channels/filters, [kernel width, kernel height], stride]` to
|
|
specify a convolutional layer stacked in order of the outer list.
|
|
cnn_transpose_use_bias: Whether to use bias on all Conv2DTranspose layers.
|
|
cnn_transpose_activation: The activation function to use after each layer
|
|
(except for the output).
|
|
cnn_transpose_use_layernorm: Whether to insert a LayerNorm functionality
|
|
in between each Conv2DTranspose layer's output and its activation.
|
|
cnn_transpose_kernel_initializer: The initializer function or class to use for
|
|
kernel initialization in the CNN layers. If `None` the default initializer
|
|
of the respective CNN layer of a framework (`"torch"` or `"tf2"`) is used.
|
|
Note, all initializers defined in the framework `"tf2`) are allowed. For
|
|
`"torch"` only the in-place initializers, i.e. ending with an underscore "_"
|
|
are allowed.
|
|
cnn_transpose_kernel_initializer_config: Configuration to pass into the
|
|
initializer defined in `cnn_transpose_kernel_initializer`.
|
|
cnn_transpose_bias_initializer: The initializer function or class to use for
|
|
bias initialization in the CNN layers. If `None` the default initializer of
|
|
the respective CNN layer of a framework (`"torch"` or `"tf2"`) is used.
|
|
For `"torch"` only the in-place initializers, i.e. ending with an underscore
|
|
"_" are allowed.
|
|
cnn_transpose_bias_initializer_config: Configuration to pass into the
|
|
initializer defined in `cnn_transpose_bias_initializer`.
|
|
|
|
Example:
|
|
.. testcode::
|
|
:skipif: True
|
|
|
|
# Configuration:
|
|
config = CNNTransposeHeadConfig(
|
|
input_dims=[10], # 1D input vector (possibly coming from another NN)
|
|
initial_image_dims=[4, 4, 96], # first image input to deconv stack
|
|
# Initializer for TensorFlow.
|
|
initial_dense_weights_initializer="HeNormal",
|
|
initial_dense_weights_initializer={"seed": 334},
|
|
cnn_transpose_filter_specifiers=[
|
|
[48, [4, 4], 2],
|
|
[24, [4, 4], 2],
|
|
[3, [4, 4], 2],
|
|
],
|
|
cnn_transpose_activation="silu", # or "swish", which is the same
|
|
cnn_transpose_use_layernorm=False,
|
|
cnn_use_bias=True,
|
|
)
|
|
model = config.build(framework="torch)
|
|
|
|
# Resulting stack in pseudocode:
|
|
# Linear(10, 4*4*24)
|
|
# Conv2DTranspose(
|
|
# in_channels=96, out_channels=48,
|
|
# kernel_size=[4, 4], stride=2, bias=True,
|
|
# )
|
|
# Swish()
|
|
# Conv2DTranspose(
|
|
# in_channels=48, out_channels=24,
|
|
# kernel_size=[4, 4], stride=2, bias=True,
|
|
# )
|
|
# Swish()
|
|
# Conv2DTranspose(
|
|
# in_channels=24, out_channels=3,
|
|
# kernel_size=[4, 4], stride=2, bias=True,
|
|
# )
|
|
|
|
Example:
|
|
.. testcode::
|
|
:skipif: True
|
|
|
|
# Configuration:
|
|
config = CNNTransposeHeadConfig(
|
|
input_dims=[128], # 1D input vector (possibly coming from another NN)
|
|
initial_image_dims=[4, 4, 32], # first image input to deconv stack
|
|
cnn_transpose_filter_specifiers=[
|
|
[16, 4, 2],
|
|
[3, 4, 2],
|
|
],
|
|
cnn_transpose_activation="relu",
|
|
cnn_transpose_use_layernorm=True,
|
|
cnn_use_bias=False,
|
|
# Initializer for `framework="tf2"`.
|
|
# Note, for Torch only in-place initializers are allowed.
|
|
cnn_transpose_kernel_initializer="xavier_normal_",
|
|
cnn_transpose_kernel_initializer_config={"gain": 0.8},
|
|
)
|
|
model = config.build(framework="torch)
|
|
|
|
# Resulting stack in pseudocode:
|
|
# Linear(128, 4*4*32, bias=True) # bias always True for initial dense layer
|
|
# Conv2DTranspose(
|
|
# in_channels=32, out_channels=16,
|
|
# kernel_size=[4, 4], stride=2, bias=False,
|
|
# )
|
|
# LayerNorm((-3, -2, -1)) # layer normalize over last 3 axes
|
|
# ReLU()
|
|
# Conv2DTranspose(
|
|
# in_channels=16, out_channels=3,
|
|
# kernel_size=[4, 4], stride=2, bias=False,
|
|
# )
|
|
"""
|
|
|
|
input_dims: Union[List[int], Tuple[int, ...]] = None
|
|
initial_image_dims: Union[List[int], Tuple[int, ...]] = field(
|
|
default_factory=lambda: [4, 4, 96]
|
|
)
|
|
initial_dense_weights_initializer: Optional[Union[str, Callable]] = None
|
|
initial_dense_weights_initializer_config: Optional[Dict] = None
|
|
initial_dense_bias_initializer: Optional[Union[str, Callable]] = None
|
|
initial_dense_bias_initializer_config: Optional[Dict] = None
|
|
cnn_transpose_filter_specifiers: List[List[Union[int, List[int]]]] = field(
|
|
default_factory=lambda: [[48, [4, 4], 2], [24, [4, 4], 2], [3, [4, 4], 2]]
|
|
)
|
|
cnn_transpose_use_bias: bool = True
|
|
cnn_transpose_activation: str = "relu"
|
|
cnn_transpose_use_layernorm: bool = False
|
|
cnn_transpose_kernel_initializer: Optional[Union[str, Callable]] = None
|
|
cnn_transpose_kernel_initializer_config: Optional[Dict] = None
|
|
cnn_transpose_bias_initializer: Optional[Union[str, Callable]] = None
|
|
cnn_transpose_bias_initializer_config: Optional[Dict] = None
|
|
|
|
@property
|
|
def output_dims(self):
|
|
# Infer output dims, layer by layer.
|
|
dims = self.initial_image_dims
|
|
for filter_spec in self.cnn_transpose_filter_specifiers:
|
|
# Same padding.
|
|
num_filters, kernel, stride = filter_spec
|
|
# Compute stride output size first (striding is performed first in a
|
|
# conv transpose layer.
|
|
stride_w, stride_h = (stride, stride) if isinstance(stride, int) else stride
|
|
dims = [
|
|
dims[0] * stride_w - (stride_w - 1),
|
|
dims[1] * stride_h - (stride_h - 1),
|
|
num_filters,
|
|
]
|
|
# TODO (Sven): Support "valid" padding for Conv2DTranspose layers, too.
|
|
# Analogous to Conv2D Layers in a CNNEncoder.
|
|
# Apply the correct padding. Note that this might be asymetrical, meaning
|
|
# left padding might be != right padding, same for top/bottom.
|
|
_, padding_out_size = same_padding_transpose_after_stride(
|
|
(dims[0], dims[1]), kernel, stride
|
|
)
|
|
# Perform conv transpose operation with the kernel.
|
|
kernel_w, kernel_h = (kernel, kernel) if isinstance(kernel, int) else kernel
|
|
dims = [
|
|
padding_out_size[0] - (kernel_w - 1),
|
|
padding_out_size[1] - (kernel_h - 1),
|
|
num_filters,
|
|
]
|
|
return tuple(dims)
|
|
|
|
def _validate(self, framework: str = "torch"):
|
|
if len(self.input_dims) != 1:
|
|
raise ValueError(
|
|
f"`input_dims` ({self.input_dims}) of CNNTransposeHeadConfig must be a "
|
|
"3D tensor (image-like) with the dimensions meaning: width x height x "
|
|
"num_filters, e.g. `[4, 4, 92]`!"
|
|
)
|
|
|
|
@_framework_implemented()
|
|
def build(self, framework: str = "torch") -> "Model":
|
|
self._validate(framework)
|
|
|
|
if framework == "torch":
|
|
from ray.rllib.core.models.torch.heads import TorchCNNTransposeHead
|
|
|
|
return TorchCNNTransposeHead(self)
|
|
|
|
|
|
@DeveloperAPI
|
|
@dataclass
|
|
class CNNEncoderConfig(ModelConfig):
|
|
"""Configuration for a convolutional (encoder) network.
|
|
|
|
The configured CNN encodes 3D-observations into a latent space.
|
|
The stack of layers is composed of a sequence of convolutional layers.
|
|
`input_dims` describes the shape of the input tensor. Beyond that, each layer
|
|
specified by `filter_specifiers` is followed by an activation function according
|
|
to `filter_activation`.
|
|
|
|
`output_dims` is reached by either the final convolutional layer's output directly
|
|
OR by flatten this output.
|
|
|
|
See ModelConfig for usage details.
|
|
|
|
Example:
|
|
|
|
.. testcode::
|
|
|
|
# Configuration:
|
|
config = CNNEncoderConfig(
|
|
input_dims=[84, 84, 3], # must be 3D tensor (image: w x h x C)
|
|
cnn_filter_specifiers=[
|
|
[16, [8, 8], 4],
|
|
[32, [4, 4], 2],
|
|
],
|
|
cnn_activation="relu",
|
|
cnn_use_layernorm=False,
|
|
cnn_use_bias=True,
|
|
)
|
|
model = config.build(framework="torch")
|
|
|
|
# Resulting stack in pseudocode:
|
|
# Conv2D(
|
|
# in_channels=3, out_channels=16,
|
|
# kernel_size=[8, 8], stride=[4, 4], bias=True,
|
|
# )
|
|
# ReLU()
|
|
# Conv2D(
|
|
# in_channels=16, out_channels=32,
|
|
# kernel_size=[4, 4], stride=[2, 2], bias=True,
|
|
# )
|
|
# ReLU()
|
|
# Conv2D(
|
|
# in_channels=32, out_channels=1,
|
|
# kernel_size=[1, 1], stride=[1, 1], bias=True,
|
|
# )
|
|
# Flatten()
|
|
|
|
Attributes:
|
|
input_dims: The input dimension of the network. These must be given in the
|
|
form of `(width, height, channels)`.
|
|
cnn_filter_specifiers: A list in which each element is another (inner) list
|
|
of either the following forms:
|
|
`[number of channels/filters, kernel, stride]`
|
|
OR:
|
|
`[number of channels/filters, kernel, stride, padding]`, where `padding`
|
|
can either be "same" or "valid".
|
|
When using the first format w/o the `padding` specifier, `padding` is "same"
|
|
by default. Also, `kernel` and `stride` may be provided either as single
|
|
ints (square) or as a tuple/list of two ints (width- and height dimensions)
|
|
for non-squared kernel/stride shapes.
|
|
A good rule of thumb for constructing CNN stacks is:
|
|
When using padding="same", the input "image" will be reduced in size by
|
|
the factor `stride`, e.g. input=(84, 84, 3) stride=2 kernel=x padding="same"
|
|
filters=16 -> output=(42, 42, 16).
|
|
For example, if you would like to reduce an Atari image from its original
|
|
(84, 84, 3) dimensions down to (6, 6, F), you can construct the following
|
|
stack and reduce the w x h dimension of the image by 2 in each layer:
|
|
[[16, 4, 2], [32, 4, 2], [64, 4, 2], [128, 4, 2]] -> output=(6, 6, 128)
|
|
cnn_use_bias: Whether to use bias on all Conv2D layers.
|
|
cnn_activation: The activation function to use after each layer (
|
|
except for the output). The default activation for Conv2d layers is "relu".
|
|
cnn_use_layernorm: Whether to insert a LayerNorm functionality
|
|
in between each CNN layer's output and its activation. Note that
|
|
the output layer.
|
|
cnn_kernel_initializer: The initializer function or class to use for kernel
|
|
initialization in the CNN layers. If `None` the default initializer of the
|
|
respective CNN layer of a framework (`"torch"` or `"tf2"`) is used. Note,
|
|
all initializers defined in the framework `"tf2`) are allowed. For `"torch"`
|
|
only the in-place initializers, i.e. ending with an underscore "_" are
|
|
allowed.
|
|
cnn_kernel_initializer_config: Configuration to pass into the initializer
|
|
defined in `cnn_kernel_initializer`.
|
|
cnn_bias_initializer: The initializer function or class to use for bias
|
|
initialization in the CNN layers. If `None` the default initializer of
|
|
the respective CNN layer of a framework (`"torch"` or `"tf2"`) is used.
|
|
For `"torch"` only the in-place initializers, i.e. ending with an underscore
|
|
"_" are allowed.
|
|
cnn_bias_initializer_config: Configuration to pass into the initializer defined
|
|
in `cnn_bias_initializer`.
|
|
flatten_at_end: Whether to flatten the output of the last conv 2D layer into
|
|
a 1D tensor. By default, this is True. Note that if you set this to False,
|
|
you might simply stack another CNNEncoder on top of this one (maybe with
|
|
different activation and bias settings).
|
|
"""
|
|
|
|
input_dims: Union[List[int], Tuple[int, ...]] = None
|
|
cnn_filter_specifiers: List[List[Union[int, List[int]]]] = field(
|
|
default_factory=lambda: [[16, [4, 4], 2], [32, [4, 4], 2], [64, [8, 8], 2]]
|
|
)
|
|
cnn_use_bias: bool = True
|
|
cnn_activation: str = "relu"
|
|
cnn_use_layernorm: bool = False
|
|
cnn_kernel_initializer: Optional[Union[str, Callable]] = None
|
|
cnn_kernel_initializer_config: Optional[Dict] = None
|
|
cnn_bias_initializer: Optional[Union[str, Callable]] = None
|
|
cnn_bias_initializer_config: Optional[Dict] = None
|
|
flatten_at_end: bool = True
|
|
|
|
@property
|
|
def output_dims(self):
|
|
if not self.input_dims:
|
|
return None
|
|
|
|
# Infer output dims, layer by layer.
|
|
dims = self.input_dims # Creates a copy (works for tuple/list).
|
|
for filter_spec in self.cnn_filter_specifiers:
|
|
# Padding not provided, "same" by default.
|
|
if len(filter_spec) == 3:
|
|
num_filters, kernel, stride = filter_spec
|
|
padding = "same"
|
|
# Padding option provided, use given value.
|
|
else:
|
|
num_filters, kernel, stride, padding = filter_spec
|
|
|
|
# Same padding.
|
|
if padding == "same":
|
|
_, dims = same_padding(dims[:2], kernel, stride)
|
|
# Valid padding.
|
|
else:
|
|
dims = valid_padding(dims[:2], kernel, stride)
|
|
|
|
# Add depth (num_filters) to the end (our utility functions for same/valid
|
|
# only return the image width/height).
|
|
dims = [dims[0], dims[1], num_filters]
|
|
|
|
# Flatten everything.
|
|
if self.flatten_at_end:
|
|
return (int(np.prod(dims)),)
|
|
|
|
return tuple(dims)
|
|
|
|
def _validate(self, framework: str = "torch"):
|
|
if len(self.input_dims) != 3:
|
|
raise ValueError(
|
|
f"`input_dims` ({self.input_dims}) of CNNEncoderConfig must be a 3D "
|
|
"tensor (image) with the dimensions meaning: width x height x "
|
|
"channels, e.g. `[64, 64, 3]`!"
|
|
)
|
|
if not self.flatten_at_end and len(self.output_dims) != 3:
|
|
raise ValueError(
|
|
f"`output_dims` ({self.output_dims}) of CNNEncoderConfig must be "
|
|
"3D, e.g. `[4, 4, 128]`, b/c your `flatten_at_end` setting is False! "
|
|
"`output_dims` is an inferred value, hence other settings might be "
|
|
"wrong."
|
|
)
|
|
elif self.flatten_at_end and len(self.output_dims) != 1:
|
|
raise ValueError(
|
|
f"`output_dims` ({self.output_dims}) of CNNEncoderConfig must be "
|
|
"1D, e.g. `[32]`, b/c your `flatten_at_end` setting is True! "
|
|
"`output_dims` is an inferred value, hence other settings might be "
|
|
"wrong."
|
|
)
|
|
|
|
@_framework_implemented()
|
|
def build(self, framework: str = "torch") -> "Model":
|
|
self._validate(framework)
|
|
|
|
if framework == "torch":
|
|
from ray.rllib.core.models.torch.encoder import TorchCNNEncoder
|
|
|
|
return TorchCNNEncoder(self)
|
|
|
|
|
|
@DeveloperAPI
|
|
@dataclass
|
|
class MLPEncoderConfig(_MLPConfig):
|
|
"""Configuration for an MLP that acts as an encoder.
|
|
|
|
See _MLPConfig for usage details.
|
|
|
|
Example:
|
|
.. testcode::
|
|
|
|
# Configuration:
|
|
config = MLPEncoderConfig(
|
|
input_dims=[4], # must be 1D tensor
|
|
hidden_layer_dims=[16],
|
|
hidden_layer_activation="relu",
|
|
hidden_layer_use_layernorm=False,
|
|
output_layer_dim=None, # maybe None or an int
|
|
)
|
|
model = config.build(framework="torch")
|
|
|
|
# Resulting stack in pseudocode:
|
|
# Linear(4, 16, bias=True)
|
|
# ReLU()
|
|
|
|
Example:
|
|
.. testcode::
|
|
|
|
# Configuration:
|
|
config = MLPEncoderConfig(
|
|
input_dims=[2],
|
|
hidden_layer_dims=[8, 8],
|
|
hidden_layer_activation="silu",
|
|
hidden_layer_use_layernorm=True,
|
|
hidden_layer_use_bias=False,
|
|
output_layer_dim=4,
|
|
output_layer_activation="tanh",
|
|
output_layer_use_bias=False,
|
|
)
|
|
model = config.build(framework="tf2")
|
|
|
|
# Resulting stack in pseudocode:
|
|
# Linear(2, 8, bias=False)
|
|
# LayerNorm((8,)) # layernorm always before activation
|
|
# SiLU()
|
|
# Linear(8, 8, bias=False)
|
|
# LayerNorm((8,)) # layernorm always before activation
|
|
# SiLU()
|
|
# Linear(8, 4, bias=False)
|
|
# Tanh()
|
|
"""
|
|
|
|
@_framework_implemented()
|
|
def build(self, framework: str = "torch") -> "Encoder":
|
|
self._validate(framework)
|
|
|
|
if framework == "torch":
|
|
from ray.rllib.core.models.torch.encoder import TorchMLPEncoder
|
|
|
|
return TorchMLPEncoder(self)
|
|
|
|
|
|
@DeveloperAPI
|
|
@dataclass
|
|
class RecurrentEncoderConfig(ModelConfig):
|
|
"""Configuration for an LSTM-based or a GRU-based encoder.
|
|
|
|
The encoder consists of...
|
|
- Zero or one tokenizers
|
|
- N LSTM/GRU layers stacked on top of each other and feeding
|
|
their outputs as inputs to the respective next layer.
|
|
|
|
This makes for the following flow of tensors:
|
|
|
|
Inputs
|
|
|
|
|
[Tokenizer if present]
|
|
|
|
|
LSTM layer 1
|
|
|
|
|
(...)
|
|
|
|
|
LSTM layer n
|
|
|
|
|
Outputs
|
|
|
|
The internal state is structued as (num_layers, B, hidden-size) for all hidden
|
|
state components, e.g.
|
|
h- and c-states of the LSTM layer(s) or h-state of the GRU layer(s).
|
|
For example, the hidden states of an LSTMEncoder with num_layers=2 and hidden_dim=8
|
|
would be: {"h": (2, B, 8), "c": (2, B, 8)}.
|
|
|
|
`output_dims` is reached by the last recurrent layer's dimension, which is always
|
|
the `hidden_dims` value.
|
|
|
|
Example:
|
|
.. testcode::
|
|
|
|
# Configuration:
|
|
config = RecurrentEncoderConfig(
|
|
recurrent_layer_type="lstm",
|
|
input_dims=[16], # must be 1D tensor
|
|
hidden_dim=128,
|
|
num_layers=2,
|
|
use_bias=True,
|
|
)
|
|
model = config.build(framework="torch")
|
|
|
|
# Resulting stack in pseudocode:
|
|
# LSTM(16, 128, bias=True)
|
|
# LSTM(128, 128, bias=True)
|
|
|
|
# Resulting shape of the internal states (c- and h-states):
|
|
# (2, B, 128) for each c- and h-states.
|
|
|
|
Example:
|
|
.. testcode::
|
|
|
|
# Configuration:
|
|
config = RecurrentEncoderConfig(
|
|
recurrent_layer_type="gru",
|
|
input_dims=[32], # must be 1D tensor
|
|
hidden_dim=64,
|
|
num_layers=1,
|
|
use_bias=False,
|
|
)
|
|
model = config.build(framework="torch")
|
|
|
|
# Resulting stack in pseudocode:
|
|
# GRU(32, 64, bias=False)
|
|
|
|
# Resulting shape of the internal state:
|
|
# (1, B, 64)
|
|
|
|
Attributes:
|
|
input_dims: The input dimensions. Must be 1D. This is the 1D shape of the tensor
|
|
that goes into the first recurrent layer.
|
|
recurrent_layer_type: The type of the recurrent layer(s).
|
|
Either "lstm" or "gru".
|
|
hidden_dim: The size of the hidden internal state(s) of the recurrent layer(s).
|
|
For example, for an LSTM, this would be the size of the c- and h-tensors.
|
|
num_layers: The number of recurrent (LSTM or GRU) layers to stack.
|
|
batch_major: Whether the input is batch major (B, T, ..) or
|
|
time major (T, B, ..).
|
|
hidden_weights_initializer: The initializer function or class to use for
|
|
kernel initialization in the hidden layers. If `None` the default
|
|
initializer of the respective recurrent layer of a framework (`"torch"` or
|
|
`"tf2"`) is used. Note, all initializers defined in the frameworks (
|
|
`"torch"` or `"tf2`) are allowed. For `"torch"` only the in-place
|
|
initializers, i.e. ending with an underscore "_" are allowed.
|
|
hidden_weights_initializer_config: Configuration to pass into the
|
|
initializer defined in `hidden_weights_initializer`.
|
|
use_bias: Whether to use bias on the recurrent layers in the network.
|
|
hidden_bias_initializer: The initializer function or class to use for bias
|
|
initialization in the hidden layers. If `None` the default initializer of
|
|
the respective recurrent layer of a framework (`"torch"` or `"tf2"`) is
|
|
used. For `"torch"` only the in-place initializers, i.e. ending with an
|
|
underscore "_" are allowed.
|
|
hidden_bias_initializer_config: Configuration to pass into the initializer
|
|
defined in `hidden_bias_initializer`.
|
|
tokenizer_config: A ModelConfig to build tokenizers for observations,
|
|
actions and other spaces.
|
|
"""
|
|
|
|
recurrent_layer_type: str = "lstm"
|
|
hidden_dim: int = None
|
|
num_layers: int = None
|
|
batch_major: bool = True
|
|
hidden_weights_initializer: Optional[Union[str, Callable]] = None
|
|
hidden_weights_initializer_config: Optional[Dict] = None
|
|
use_bias: bool = True
|
|
hidden_bias_initializer: Optional[Union[str, Callable]] = None
|
|
hidden_bias_initializer_config: Optional[Dict] = None
|
|
tokenizer_config: ModelConfig = None
|
|
|
|
@property
|
|
def output_dims(self):
|
|
return (self.hidden_dim,)
|
|
|
|
def _validate(self, framework: str = "torch"):
|
|
"""Makes sure that settings are valid."""
|
|
if self.recurrent_layer_type not in ["gru", "lstm"]:
|
|
raise ValueError(
|
|
f"`recurrent_layer_type` ({self.recurrent_layer_type}) of "
|
|
"RecurrentEncoderConfig must be 'gru' or 'lstm'!"
|
|
)
|
|
if self.input_dims is not None and len(self.input_dims) != 1:
|
|
raise ValueError(
|
|
f"`input_dims` ({self.input_dims}) of RecurrentEncoderConfig must be "
|
|
"1D, e.g. `[32]`!"
|
|
)
|
|
if len(self.output_dims) != 1:
|
|
raise ValueError(
|
|
f"`output_dims` ({self.output_dims}) of RecurrentEncoderConfig must be "
|
|
"1D, e.g. `[32]`! This is an inferred value, hence other settings might"
|
|
" be wrong."
|
|
)
|
|
|
|
@_framework_implemented()
|
|
def build(self, framework: str = "torch") -> "Encoder":
|
|
if framework == "torch":
|
|
from ray.rllib.core.models.torch.encoder import (
|
|
TorchGRUEncoder as GRU,
|
|
TorchLSTMEncoder as LSTM,
|
|
)
|
|
|
|
if self.recurrent_layer_type == "lstm":
|
|
return LSTM(self)
|
|
else:
|
|
return GRU(self)
|
|
|
|
|
|
@DeveloperAPI
|
|
@dataclass
|
|
class ActorCriticEncoderConfig(ModelConfig):
|
|
"""Configuration for an ActorCriticEncoder.
|
|
|
|
The base encoder functions like other encoders in RLlib. It is wrapped by the
|
|
ActorCriticEncoder to provides a shared encoder Model to use in RLModules that
|
|
provides twofold outputs: one for the actor and one for the critic. See
|
|
ModelConfig for usage details.
|
|
|
|
Attributes:
|
|
base_encoder_config: The configuration for the wrapped encoder(s).
|
|
shared: Whether the base encoder is shared between the actor and critic.
|
|
inference_only: Whether the configured encoder will only ever be used as an
|
|
actor-encoder, never as a value-function encoder. Thus, if True and `shared`
|
|
is False, will only build the actor-related components.
|
|
"""
|
|
|
|
base_encoder_config: ModelConfig = None
|
|
shared: bool = True
|
|
inference_only: bool = False
|
|
|
|
@_framework_implemented()
|
|
def build(self, framework: str = "torch") -> "Encoder":
|
|
if framework == "torch":
|
|
from ray.rllib.core.models.torch.encoder import (
|
|
TorchActorCriticEncoder,
|
|
TorchStatefulActorCriticEncoder,
|
|
)
|
|
|
|
if isinstance(self.base_encoder_config, RecurrentEncoderConfig):
|
|
return TorchStatefulActorCriticEncoder(self)
|
|
else:
|
|
return TorchActorCriticEncoder(self)
|
|
|
|
|
|
@DeveloperAPI
|
|
@dataclass
|
|
class MultiStreamEncoderConfig(ModelConfig):
|
|
"""Configuration for a Multi-Stream Encoder.
|
|
|
|
The base encoders function like other encoders in RLlib. They are wrapped by the
|
|
Multi-Stream Encoder to provide a shared encoder Model to use in RLModules that
|
|
processes multiple input streams. The outputs of the individual encoders are
|
|
concatenated and further processed by shared fusion layers.
|
|
|
|
Attributes:
|
|
base_encoder_configs: A dictionary mapping stream names to their respective
|
|
encoder configurations.
|
|
hidden_layer_dims: Dimensions of the shared hidden layers after the
|
|
concatenation of the individual encoders' outputs.
|
|
hidden_layer_use_bias: Whether to use bias in the shared hidden layers.
|
|
hidden_layer_activation: Activation function to use in the shared hidden layers.
|
|
hidden_layer_weights_initializer: The initializer function or class to use for
|
|
weight initialization in the shared hidden layers. If `None` the default
|
|
initializer of the respective dense layer is used. Note, for `"torch"` only
|
|
the in-place initializers, i.e. ending with an underscore "_" are allowed.
|
|
hidden_layer_weights_initializer_config: Configuration to pass into the
|
|
initializer defined in `hidden_layer_weights_initializer`.
|
|
hidden_layer_bias_initializer: The initializer function or class to use for
|
|
bias initialization in the shared hidden layers. If `None` the default
|
|
initializer of the respective dense layer is used. Note, for `"torch"` only
|
|
the in-place initializers, i.e. ending with an underscore "_" are allowed.
|
|
hidden_layer_bias_initializer_config: Configuration to pass into the
|
|
initializer defined in `hidden_layer_bias_initializer`.
|
|
output_layer_dim: Optional dimension of the final output layer after the
|
|
shared hidden layers. If `None`, no output layer is added.
|
|
output_layer_use_bias: Whether to use bias in the output layer.
|
|
output_layer_activation: Activation function to use in the output layer.
|
|
output_layer_weights_initializer: The initializer function or class to use for
|
|
weight initialization in the output layer. If `None` the default
|
|
initializer of the respective dense layer is used. Note, for `"torch"` only
|
|
the in-place initializers, i.e. ending with an underscore "_" are allowed.
|
|
output_layer_weights_initializer_config: Configuration to pass into the
|
|
initializer defined in `output_layer_weights_initializer`.
|
|
output_layer_bias_initializer: The initializer function or class to use for
|
|
bias initialization in the output layer. If `None` the default
|
|
initializer of the respective dense layer is used. Note, for `"torch"` only
|
|
the in-place initializers, i.e. ending with an underscore "_" are allowed.
|
|
output_layer_bias_initializer_config: Configuration to pass into the
|
|
initializer defined in `output_layer_bias_initializer`.
|
|
"""
|
|
|
|
# Configuration for a Multi-Stream Encoder.s
|
|
base_encoder_configs: Dict[str, ModelConfig] = None
|
|
|
|
# Shared hidden layers after the individual encoders' outputs have been
|
|
# concatenated.
|
|
hidden_layer_dims: Union[List[int], Tuple[int, ...]] = (256, 256, 256)
|
|
hidden_layer_use_bias: bool = True
|
|
hidden_layer_activation: str = "relu"
|
|
# TODO (simon): enable layernorm for multi-stream encoder.
|
|
# hidden_layer_use_layernorm: bool = False
|
|
hidden_layer_weights_initializer: Optional[Union[str, Callable]] = None
|
|
hidden_layer_weights_initializer_config: Optional[Dict] = None
|
|
hidden_layer_bias_initializer: Optional[Union[str, Callable]] = None
|
|
hidden_layer_bias_initializer_config: Optional[Dict] = None
|
|
|
|
# Optional last output layer with - possibly - different activation and use_bias
|
|
# settings.
|
|
output_layer_dim: Optional[int] = None
|
|
output_layer_use_bias: bool = True
|
|
output_layer_activation: str = "linear"
|
|
output_layer_weights_initializer: Optional[Union[str, Callable]] = None
|
|
output_layer_weights_initializer_config: Optional[Dict] = None
|
|
output_layer_bias_initializer: Optional[Union[str, Callable]] = None
|
|
output_layer_bias_initializer_config: Optional[Dict] = None
|
|
|
|
@property
|
|
def output_dims(self):
|
|
# If output layer is present, its dim is the output dim.
|
|
if self.output_layer_dim is not None:
|
|
return (self.output_layer_dim,)
|
|
# If no output layer, the output dim is the last hidden layer's dim.
|
|
elif self.hidden_layer_dims:
|
|
return (self.hidden_layer_dims[-1],)
|
|
# No output layer and no hidden layers: output is concatenated embeddings.
|
|
else:
|
|
total_embed_dim = sum(
|
|
cfg.output_dims[0] for cfg in self.base_encoder_configs.values()
|
|
)
|
|
return (total_embed_dim,)
|
|
|
|
def _validate(self, framework: str = "torch"):
|
|
"""Makes sure that settings are valid."""
|
|
|
|
if self.base_encoder_configs is None or len(self.base_encoder_configs) == 0:
|
|
raise ValueError(
|
|
"`base_encoder_configs` of MultiStreamEncoderConfig must be a "
|
|
"non-empty dictionary mapping stream names to their respective "
|
|
"encoder configurations!"
|
|
)
|
|
|
|
if self.input_dims is not None and len(self.input_dims) != 1:
|
|
raise ValueError(
|
|
f"`input_dims` ({self.input_dims}) of MultiStreamEncoderConfig must be 1D, "
|
|
"e.g. `[32]`!"
|
|
)
|
|
if len(self.output_dims) != 1:
|
|
raise ValueError(
|
|
f"`output_dims` ({self.output_dims}) of MultiStreamEncoderConfig must be "
|
|
"1D, e.g. `[32]`! This is an inferred value, hence other settings might"
|
|
" be wrong."
|
|
)
|
|
|
|
# Call these already here to catch errors early on.
|
|
get_activation_fn(self.hidden_layer_activation, framework=framework)
|
|
get_activation_fn(self.output_layer_activation, framework=framework)
|
|
get_initializer_fn(self.hidden_layer_weights_initializer, framework=framework)
|
|
get_initializer_fn(self.hidden_layer_bias_initializer, framework=framework)
|
|
get_initializer_fn(self.output_layer_weights_initializer, framework=framework)
|
|
get_initializer_fn(self.output_layer_bias_initializer, framework=framework)
|
|
|
|
@_framework_implemented()
|
|
def build(self, framework: str = "torch") -> "Encoder":
|
|
# Validate config.
|
|
self._validate(framework)
|
|
|
|
if framework == "torch":
|
|
from ray.rllib.core.models.torch.encoder import TorchMultiStreamEncoder
|
|
|
|
return TorchMultiStreamEncoder(self)
|
|
else:
|
|
raise NotImplementedError(
|
|
f"MultiStreamEncoder is not implemented for framework {framework}"
|
|
)
|