wehub-resource-sync
426 lines
16 KiB
Python
426 lines
16 KiB
Python
from ray.rllib.core.columns import Columns
|
|
from ray.rllib.core.models.base import (
|
|
ENCODER_OUT,
|
|
ActorCriticEncoder,
|
|
Encoder,
|
|
Model,
|
|
StatefulActorCriticEncoder,
|
|
tokenize,
|
|
)
|
|
from ray.rllib.core.models.configs import (
|
|
ActorCriticEncoderConfig,
|
|
CNNEncoderConfig,
|
|
MLPEncoderConfig,
|
|
MultiStreamEncoderConfig,
|
|
RecurrentEncoderConfig,
|
|
)
|
|
from ray.rllib.core.models.torch.base import TorchModel
|
|
from ray.rllib.core.models.torch.primitives import TorchCNN, TorchMLP
|
|
from ray.rllib.models.utils import get_activation_fn, get_initializer_fn
|
|
from ray.rllib.utils.annotations import override
|
|
from ray.rllib.utils.framework import try_import_torch
|
|
|
|
torch, nn = try_import_torch()
|
|
|
|
|
|
class TorchActorCriticEncoder(TorchModel, ActorCriticEncoder):
|
|
"""An actor-critic encoder for torch."""
|
|
|
|
framework = "torch"
|
|
|
|
def __init__(self, config: ActorCriticEncoderConfig) -> None:
|
|
TorchModel.__init__(self, config)
|
|
ActorCriticEncoder.__init__(self, config)
|
|
|
|
|
|
class TorchStatefulActorCriticEncoder(TorchModel, StatefulActorCriticEncoder):
|
|
"""A stateful actor-critic encoder for torch."""
|
|
|
|
framework = "torch"
|
|
|
|
def __init__(self, config: ActorCriticEncoderConfig) -> None:
|
|
TorchModel.__init__(self, config)
|
|
StatefulActorCriticEncoder.__init__(self, config)
|
|
|
|
|
|
class TorchMLPEncoder(TorchModel, Encoder):
|
|
def __init__(self, config: MLPEncoderConfig) -> None:
|
|
TorchModel.__init__(self, config)
|
|
Encoder.__init__(self, config)
|
|
|
|
# Create the neural network.
|
|
self.net = TorchMLP(
|
|
input_dim=config.input_dims[0],
|
|
hidden_layer_dims=config.hidden_layer_dims,
|
|
hidden_layer_activation=config.hidden_layer_activation,
|
|
hidden_layer_use_layernorm=config.hidden_layer_use_layernorm,
|
|
hidden_layer_use_bias=config.hidden_layer_use_bias,
|
|
hidden_layer_weights_initializer=config.hidden_layer_weights_initializer,
|
|
hidden_layer_weights_initializer_config=(
|
|
config.hidden_layer_weights_initializer_config
|
|
),
|
|
hidden_layer_bias_initializer=config.hidden_layer_bias_initializer,
|
|
hidden_layer_bias_initializer_config=(
|
|
config.hidden_layer_bias_initializer_config
|
|
),
|
|
output_dim=config.output_layer_dim,
|
|
output_activation=config.output_layer_activation,
|
|
output_use_bias=config.output_layer_use_bias,
|
|
output_layer_use_layernorm=config.output_layer_use_layernorm,
|
|
output_weights_initializer=config.output_layer_weights_initializer,
|
|
output_weights_initializer_config=(
|
|
config.output_layer_weights_initializer_config
|
|
),
|
|
output_bias_initializer=config.output_layer_bias_initializer,
|
|
output_bias_initializer_config=config.output_layer_bias_initializer_config,
|
|
)
|
|
|
|
@override(Model)
|
|
def _forward(self, inputs: dict, **kwargs) -> dict:
|
|
return {ENCODER_OUT: self.net(inputs[Columns.OBS])}
|
|
|
|
|
|
class TorchCNNEncoder(TorchModel, Encoder):
|
|
def __init__(self, config: CNNEncoderConfig) -> None:
|
|
TorchModel.__init__(self, config)
|
|
Encoder.__init__(self, config)
|
|
|
|
layers = []
|
|
# The bare-bones CNN (no flatten, no succeeding dense).
|
|
cnn = TorchCNN(
|
|
input_dims=config.input_dims,
|
|
cnn_filter_specifiers=config.cnn_filter_specifiers,
|
|
cnn_activation=config.cnn_activation,
|
|
cnn_use_layernorm=config.cnn_use_layernorm,
|
|
cnn_use_bias=config.cnn_use_bias,
|
|
cnn_kernel_initializer=config.cnn_kernel_initializer,
|
|
cnn_kernel_initializer_config=config.cnn_kernel_initializer_config,
|
|
cnn_bias_initializer=config.cnn_bias_initializer,
|
|
cnn_bias_initializer_config=config.cnn_bias_initializer_config,
|
|
)
|
|
layers.append(cnn)
|
|
|
|
# Add a flatten operation to move from 2/3D into 1D space.
|
|
if config.flatten_at_end:
|
|
layers.append(nn.Flatten())
|
|
|
|
# Create the network from gathered layers.
|
|
self.net = nn.Sequential(*layers)
|
|
|
|
@override(Model)
|
|
def _forward(self, inputs: dict, **kwargs) -> dict:
|
|
return {ENCODER_OUT: self.net(inputs[Columns.OBS])}
|
|
|
|
|
|
class TorchGRUEncoder(TorchModel, Encoder):
|
|
"""A recurrent GRU encoder.
|
|
|
|
On the usage of `torch.utils._pytree`: Unlike `dm_tree`/`tree`, it is traceable by
|
|
`torch.export`/TorchDynamo (used by the modern `torch.onnx.export(dynamo=True)`
|
|
path), so we use it for the state (un)mapping inside the recurrent encoders.
|
|
"""
|
|
|
|
def __init__(self, config: RecurrentEncoderConfig) -> None:
|
|
TorchModel.__init__(self, config)
|
|
|
|
# Maybe create a tokenizer
|
|
if config.tokenizer_config is not None:
|
|
self.tokenizer = config.tokenizer_config.build(framework="torch")
|
|
gru_input_dims = config.tokenizer_config.output_dims
|
|
else:
|
|
self.tokenizer = None
|
|
gru_input_dims = config.input_dims
|
|
|
|
# We only support 1D spaces right now.
|
|
assert len(gru_input_dims) == 1
|
|
gru_input_dim = gru_input_dims[0]
|
|
|
|
gru_weights_initializer = get_initializer_fn(
|
|
config.hidden_weights_initializer, framework="torch"
|
|
)
|
|
gru_bias_initializer = get_initializer_fn(
|
|
config.hidden_bias_initializer, framework="torch"
|
|
)
|
|
|
|
# Create the torch GRU layer.
|
|
self.gru = nn.GRU(
|
|
gru_input_dim,
|
|
config.hidden_dim,
|
|
config.num_layers,
|
|
batch_first=config.batch_major,
|
|
bias=config.use_bias,
|
|
)
|
|
|
|
# Initialize, GRU weights, if necessary.
|
|
if gru_weights_initializer:
|
|
gru_weights_initializer(
|
|
self.gru.weight, **config.hidden_weights_initializer_config or {}
|
|
)
|
|
# Initialize GRU bias, if necessary.
|
|
if gru_bias_initializer:
|
|
gru_bias_initializer(
|
|
self.gru.weight, **config.hidden_bias_initializer_config or {}
|
|
)
|
|
|
|
@override(Model)
|
|
def get_initial_state(self):
|
|
return {
|
|
"h": torch.zeros(self.config.num_layers, self.config.hidden_dim),
|
|
}
|
|
|
|
@override(Model)
|
|
def _forward(self, inputs: dict, **kwargs) -> dict:
|
|
outputs = {}
|
|
|
|
if self.tokenizer is not None:
|
|
# Push observations through the tokenizer encoder if we built one.
|
|
out = tokenize(self.tokenizer, inputs, framework="torch")
|
|
else:
|
|
# Otherwise, just use the raw observations.
|
|
out = inputs[Columns.OBS].float()
|
|
|
|
# States are batch-first when coming in. Make them layers-first.
|
|
states_in = torch.utils._pytree.tree_map(
|
|
lambda s: s.transpose(0, 1), inputs[Columns.STATE_IN]
|
|
)
|
|
|
|
out, states_out = self.gru(out, states_in["h"])
|
|
states_out = {"h": states_out}
|
|
|
|
# Insert them into the output dict.
|
|
outputs[ENCODER_OUT] = out
|
|
outputs[Columns.STATE_OUT] = torch.utils._pytree.tree_map(
|
|
lambda s: s.transpose(0, 1), states_out
|
|
)
|
|
return outputs
|
|
|
|
|
|
class TorchLSTMEncoder(TorchModel, Encoder):
|
|
"""A recurrent LSTM encoder.
|
|
|
|
On the usage of `torch.utils._pytree`: Unlike `dm_tree`/`tree`, it is traceable by
|
|
`torch.export`/TorchDynamo (used by the modern `torch.onnx.export(dynamo=True)`
|
|
path), so we use it for the state (un)mapping inside the recurrent encoders.
|
|
"""
|
|
|
|
def __init__(self, config: RecurrentEncoderConfig) -> None:
|
|
TorchModel.__init__(self, config)
|
|
|
|
# Maybe create a tokenizer
|
|
if config.tokenizer_config is not None:
|
|
self.tokenizer = config.tokenizer_config.build(framework="torch")
|
|
lstm_input_dims = config.tokenizer_config.output_dims
|
|
else:
|
|
self.tokenizer = None
|
|
lstm_input_dims = config.input_dims
|
|
|
|
# We only support 1D spaces right now.
|
|
assert len(lstm_input_dims) == 1
|
|
lstm_input_dim = lstm_input_dims[0]
|
|
|
|
lstm_weights_initializer = get_initializer_fn(
|
|
config.hidden_weights_initializer, framework="torch"
|
|
)
|
|
lstm_bias_initializer = get_initializer_fn(
|
|
config.hidden_bias_initializer, framework="torch"
|
|
)
|
|
|
|
# Create the torch LSTM layer.
|
|
self.lstm = nn.LSTM(
|
|
lstm_input_dim,
|
|
config.hidden_dim,
|
|
config.num_layers,
|
|
batch_first=config.batch_major,
|
|
bias=config.use_bias,
|
|
)
|
|
|
|
# Initialize LSTM layer weigths and biases, if necessary.
|
|
for layer in self.lstm.all_weights:
|
|
if lstm_weights_initializer:
|
|
lstm_weights_initializer(
|
|
layer[0], **config.hidden_weights_initializer_config or {}
|
|
)
|
|
lstm_weights_initializer(
|
|
layer[1], **config.hidden_weights_initializer_config or {}
|
|
)
|
|
if lstm_bias_initializer:
|
|
lstm_bias_initializer(
|
|
layer[2], **config.hidden_bias_initializer_config or {}
|
|
)
|
|
lstm_bias_initializer(
|
|
layer[3], **config.hidden_bias_initializer_config or {}
|
|
)
|
|
|
|
@override(Model)
|
|
def get_initial_state(self):
|
|
return {
|
|
"h": torch.zeros(self.config.num_layers, self.config.hidden_dim),
|
|
"c": torch.zeros(self.config.num_layers, self.config.hidden_dim),
|
|
}
|
|
|
|
@override(Model)
|
|
def _forward(self, inputs: dict, **kwargs) -> dict:
|
|
outputs = {}
|
|
|
|
if self.tokenizer is not None:
|
|
# Push observations through the tokenizer encoder if we built one.
|
|
out = tokenize(self.tokenizer, inputs, framework="torch")
|
|
else:
|
|
# Otherwise, just use the raw observations.
|
|
out = inputs[Columns.OBS].float()
|
|
|
|
# States are batch-first when coming in. Make them layers-first.
|
|
states_in = torch.utils._pytree.tree_map(
|
|
lambda s: s.transpose(0, 1), inputs[Columns.STATE_IN]
|
|
)
|
|
|
|
out, states_out = self.lstm(out, (states_in["h"], states_in["c"]))
|
|
states_out = {"h": states_out[0], "c": states_out[1]}
|
|
|
|
# Insert them into the output dict.
|
|
outputs[ENCODER_OUT] = out
|
|
outputs[Columns.STATE_OUT] = torch.utils._pytree.tree_map(
|
|
lambda s: s.transpose(0, 1), states_out
|
|
)
|
|
return outputs
|
|
|
|
|
|
class TorchMultiStreamEncoder(TorchModel, Encoder):
|
|
"""An encoder that encodes multiple input streams separately.
|
|
|
|
Each input stream is encoded with its own encoder defined in
|
|
`base_encoder_configs`. The resulting embeddings are concatenated and
|
|
passed through a final fusion network to produce the output embedding.
|
|
"""
|
|
|
|
def __init__(self, config: MultiStreamEncoderConfig) -> None:
|
|
TorchModel.__init__(self, config)
|
|
Encoder.__init__(self, config)
|
|
|
|
# Create the neural network for observation stream.
|
|
self.base_encoders = nn.ModuleDict(
|
|
{
|
|
k: cfg.build(framework="torch")
|
|
for k, cfg in sorted(config.base_encoder_configs.items())
|
|
}
|
|
)
|
|
|
|
# Get activation functions.
|
|
self.hidden_activation = (
|
|
get_activation_fn(config.hidden_layer_activation, framework="torch")
|
|
or nn.Identity
|
|
)
|
|
|
|
# Calculate total embed dim.
|
|
self._total_embed_dim = sum(
|
|
cfg.output_dims[0] for cfg in config.base_encoder_configs.values()
|
|
)
|
|
|
|
# Build fusion layers (hidden layers with skip connections).
|
|
fusion_layers = []
|
|
if config.hidden_layer_dims:
|
|
# Input layer: total_embed_dim -> hidden_layer_dims[0]
|
|
input_layer = nn.Linear(
|
|
self._total_embed_dim,
|
|
config.hidden_layer_dims[0],
|
|
bias=config.hidden_layer_use_bias,
|
|
)
|
|
fusion_layers.append(input_layer)
|
|
|
|
# Intermediate fusion layers with skip connections.
|
|
for i in range(1, len(config.hidden_layer_dims)):
|
|
fusion_layers.append(
|
|
nn.Linear(
|
|
config.hidden_layer_dims[i - 1] + self._total_embed_dim,
|
|
config.hidden_layer_dims[i],
|
|
bias=config.hidden_layer_use_bias,
|
|
)
|
|
)
|
|
|
|
# Initialize hidden layer weights if necessary.
|
|
if config.hidden_layer_weights_initializer:
|
|
hidden_weights_initializer = get_initializer_fn(
|
|
config.hidden_layer_weights_initializer, framework="torch"
|
|
)
|
|
for layer in fusion_layers:
|
|
hidden_weights_initializer(
|
|
layer.weight,
|
|
**config.hidden_layer_weights_initializer_config or {}
|
|
)
|
|
# Initialize hidden layer bias if necessary.
|
|
if config.hidden_layer_bias_initializer:
|
|
hidden_bias_initializer = get_initializer_fn(
|
|
config.hidden_layer_bias_initializer, framework="torch"
|
|
)
|
|
for layer in fusion_layers:
|
|
hidden_bias_initializer(
|
|
layer.bias, **config.hidden_layer_bias_initializer_config or {}
|
|
)
|
|
|
|
self.fusion_layers = nn.ModuleList(fusion_layers)
|
|
|
|
# Build output layer only if output_layer_dim is defined.
|
|
self.output_layer = None
|
|
if config.output_layer_dim is not None:
|
|
# Get output activation function.
|
|
self.output_activation = (
|
|
get_activation_fn(config.output_layer_activation, framework="torch")
|
|
or nn.Identity
|
|
)
|
|
# Determine input dim for output layer.
|
|
if config.hidden_layer_dims:
|
|
output_input_dim = config.hidden_layer_dims[-1] + self._total_embed_dim
|
|
else:
|
|
output_input_dim = self._total_embed_dim
|
|
|
|
self.output_layer = nn.Linear(
|
|
output_input_dim,
|
|
config.output_layer_dim,
|
|
bias=config.output_layer_use_bias,
|
|
)
|
|
|
|
# Initialize output layer weights if necessary.
|
|
if config.output_layer_weights_initializer:
|
|
output_weights_initializer = get_initializer_fn(
|
|
config.output_layer_weights_initializer, framework="torch"
|
|
)
|
|
output_weights_initializer(
|
|
self.output_layer.weight,
|
|
**config.output_layer_weights_initializer_config or {}
|
|
)
|
|
# Initialize output layer bias if necessary.
|
|
if config.output_layer_bias_initializer:
|
|
output_bias_initializer = get_initializer_fn(
|
|
config.output_layer_bias_initializer, framework="torch"
|
|
)
|
|
output_bias_initializer(
|
|
self.output_layer.bias,
|
|
**config.output_layer_bias_initializer_config or {}
|
|
)
|
|
|
|
@override(Model)
|
|
def _forward(self, inputs, **kwargs):
|
|
# Run the inputs through the base encoders.
|
|
keys = sorted(self.config.base_encoder_configs.keys())
|
|
encoder_outs = [
|
|
self.base_encoders[k]({Columns.OBS: inputs[k]})[ENCODER_OUT] for k in keys
|
|
]
|
|
# Concatenate the embeddings.
|
|
embeds = torch.cat(encoder_outs, dim=-1)
|
|
|
|
# Pass through fusion layers (if any).
|
|
out = embeds
|
|
if self.fusion_layers:
|
|
out = self.hidden_activation()(self.fusion_layers[0](embeds))
|
|
for layer in self.fusion_layers[1:]:
|
|
out = self.hidden_activation()(layer(torch.cat([out, embeds], dim=-1)))
|
|
|
|
# Pass through output layer (if defined).
|
|
if self.output_layer is not None:
|
|
# Concatenate with skip connection if we have fusion layers.
|
|
if self.fusion_layers:
|
|
out = torch.cat([out, embeds], dim=-1)
|
|
out = self.output_activation()(self.output_layer(out))
|
|
|
|
return {ENCODER_OUT: out}
|