wehub-resource-sync
457 lines
16 KiB
Python
457 lines
16 KiB
Python
import abc
|
|
from typing import List, Optional, Tuple, Union
|
|
|
|
from ray.rllib.core.columns import Columns
|
|
from ray.rllib.core.models.configs import ModelConfig
|
|
from ray.rllib.core.models.specs.specs_base import Spec
|
|
from ray.rllib.policy.rnn_sequencing import get_fold_unfold_fns
|
|
from ray.rllib.utils.annotations import ExperimentalAPI, override
|
|
from ray.rllib.utils.typing import TensorType
|
|
from ray.util.annotations import DeveloperAPI
|
|
|
|
# Top level keys that unify model i/o.
|
|
ENCODER_OUT: str = "encoder_out"
|
|
# For Actor-Critic algorithms, these signify data related to the actor and critic
|
|
ACTOR: str = "actor"
|
|
CRITIC: str = "critic"
|
|
|
|
|
|
@ExperimentalAPI
|
|
class Model(abc.ABC):
|
|
"""Framework-agnostic base class for RLlib models.
|
|
|
|
Models are low-level neural network components that offer input- and
|
|
output-specification, a forward method, and a get_initial_state method. Models
|
|
are composed in RLModules.
|
|
|
|
Usage Example together with ModelConfig:
|
|
|
|
.. testcode::
|
|
|
|
from ray.rllib.core.models.base import Model
|
|
from ray.rllib.core.models.configs import ModelConfig
|
|
from ray.rllib.core.models.configs import ModelConfig
|
|
from dataclasses import dataclass
|
|
|
|
class MyModel(Model):
|
|
def __init__(self, config):
|
|
super().__init__(config)
|
|
self.my_param = config.my_param * 2
|
|
|
|
def _forward(self, input_dict):
|
|
return input_dict["obs"] * self.my_param
|
|
|
|
def get_num_parameters(self):
|
|
return (0, 0)
|
|
|
|
def _set_to_dummy_weights(self, value_sequence=(-0.02, -0.01, 0.01, 0.02)):
|
|
pass
|
|
|
|
|
|
@dataclass
|
|
class MyModelConfig(ModelConfig):
|
|
my_param: int = 42
|
|
|
|
def build(self, framework: str):
|
|
if framework == "bork":
|
|
return MyModel(self)
|
|
|
|
|
|
config = MyModelConfig(my_param=3)
|
|
model = config.build(framework="bork")
|
|
print(model._forward({"obs": 1}))
|
|
|
|
.. testoutput::
|
|
|
|
6
|
|
|
|
"""
|
|
|
|
def __init__(self, config: ModelConfig):
|
|
self.config = config
|
|
|
|
def __init_subclass__(cls, **kwargs):
|
|
# Automatically add a __post_init__ method to all subclasses of Model.
|
|
# This method is called after the __init__ method of the subclass.
|
|
def init_decorator(previous_init):
|
|
def new_init(self, *args, **kwargs):
|
|
previous_init(self, *args, **kwargs)
|
|
if type(self) is cls:
|
|
self.__post_init__()
|
|
|
|
return new_init
|
|
|
|
cls.__init__ = init_decorator(cls.__init__)
|
|
|
|
def __post_init__(self):
|
|
"""Called automatically after the __init__ method of the subclasses.
|
|
|
|
The module first calls the __init__ method of the subclass, With in the
|
|
__init__ you should call the super().__init__ method. Then after the __init__
|
|
method of the subclass is called, the __post_init__ method is called.
|
|
|
|
This is a good place to do any initialization that requires access to the
|
|
subclass's attributes.
|
|
"""
|
|
self._input_specs = self.get_input_specs()
|
|
self._output_specs = self.get_output_specs()
|
|
|
|
def get_input_specs(self) -> Optional[Spec]:
|
|
"""Returns the input specs of this model.
|
|
|
|
Override `get_input_specs` to define your own input specs.
|
|
This method should not be called often, e.g. every forward pass.
|
|
Instead, it should be called once at instantiation to define Model.input_specs.
|
|
|
|
Returns:
|
|
Spec: The input specs.
|
|
"""
|
|
return None
|
|
|
|
def get_output_specs(self) -> Optional[Spec]:
|
|
"""Returns the output specs of this model.
|
|
|
|
Override `get_output_specs` to define your own output specs.
|
|
This method should not be called often, e.g. every forward pass.
|
|
Instead, it should be called once at instantiation to define Model.output_specs.
|
|
|
|
Returns:
|
|
Spec: The output specs.
|
|
"""
|
|
return None
|
|
|
|
@property
|
|
def input_specs(self) -> Spec:
|
|
"""Returns the input spec of this model."""
|
|
return self._input_specs
|
|
|
|
@input_specs.setter
|
|
def input_specs(self, spec: Spec) -> None:
|
|
raise ValueError(
|
|
"`input_specs` cannot be set directly. Override "
|
|
"Model.get_input_specs() instead. Set Model._input_specs if "
|
|
"you want to override this behavior."
|
|
)
|
|
|
|
@property
|
|
def output_specs(self) -> Spec:
|
|
"""Returns the output specs of this model."""
|
|
return self._output_specs
|
|
|
|
@output_specs.setter
|
|
def output_specs(self, spec: Spec) -> None:
|
|
raise ValueError(
|
|
"`output_specs` cannot be set directly. Override "
|
|
"Model.get_output_specs() instead. Set Model._output_specs if "
|
|
"you want to override this behavior."
|
|
)
|
|
|
|
def get_initial_state(self) -> Union[dict, List[TensorType]]:
|
|
"""Returns the initial state of the Model.
|
|
|
|
It can be left empty if this Model is not stateful.
|
|
"""
|
|
return dict()
|
|
|
|
@abc.abstractmethod
|
|
def _forward(self, input_dict: dict, **kwargs) -> dict:
|
|
"""Returns the output of this model for the given input.
|
|
|
|
This method is called by the forwarding method of the respective framework
|
|
that is itself wrapped by RLlib in order to check model inputs and outputs.
|
|
|
|
Args:
|
|
input_dict: The input tensors.
|
|
**kwargs: Forward compatibility kwargs.
|
|
|
|
Returns:
|
|
dict: The output tensors.
|
|
"""
|
|
|
|
@abc.abstractmethod
|
|
def get_num_parameters(self) -> Tuple[int, int]:
|
|
"""Returns a tuple of (num trainable params, num non-trainable params)."""
|
|
|
|
@abc.abstractmethod
|
|
def _set_to_dummy_weights(self, value_sequence=(-0.02, -0.01, 0.01, 0.02)) -> None:
|
|
"""Helper method to set all weights to deterministic dummy values.
|
|
|
|
Calling this method on two `Models` that have the same architecture using
|
|
the exact same `value_sequence` arg should make both models output the exact
|
|
same values on arbitrary inputs. This will work, even if the two `Models`
|
|
are of different DL frameworks.
|
|
|
|
Args:
|
|
value_sequence: Looping through the list of all parameters (weight matrices,
|
|
bias tensors, etc..) of this model, in each iteration i, we set all
|
|
values in this parameter to `value_sequence[i % len(value_sequence)]`
|
|
(round robin).
|
|
|
|
Example:
|
|
TODO:
|
|
"""
|
|
|
|
|
|
@ExperimentalAPI
|
|
class Encoder(Model, abc.ABC):
|
|
"""The framework-agnostic base class for all RLlib encoders.
|
|
|
|
Encoders are used to transform observations to a latent space.
|
|
Therefore, their `input_specs` contains the observation space dimensions.
|
|
Similarly, their `output_specs` contains the latent space dimensions.
|
|
Encoders can be recurrent, in which case the state should be part of input- and
|
|
output_specs. The latent vectors produced by an encoder are fed into subsequent
|
|
"heads". Any implementation of Encoder should also be callable. This should be done
|
|
by also inheriting from a framework-specific model base-class, s.a. TorchModel or
|
|
TfModel.
|
|
|
|
Abstract illustration of typical flow of tensors:
|
|
|
|
Inputs
|
|
|
|
|
Encoder
|
|
| \
|
|
Head1 Head2
|
|
| /
|
|
Outputs
|
|
|
|
Outputs of encoders are generally of shape (B, latent_dim) or (B, T, latent_dim).
|
|
That is, for time-series data, we encode into the latent space for each time step.
|
|
This should be reflected in the `output_specs`.
|
|
|
|
Usage example together with a ModelConfig:
|
|
|
|
.. testcode::
|
|
|
|
from dataclasses import dataclass
|
|
import numpy as np
|
|
|
|
from ray.rllib.core.columns import Columns
|
|
from ray.rllib.core.models.base import Encoder, ENCODER_OUT
|
|
from ray.rllib.core.models.configs import ModelConfig
|
|
from ray.rllib.policy.sample_batch import SampleBatch
|
|
|
|
class NumpyEncoder(Encoder):
|
|
def __init__(self, config):
|
|
super().__init__(config)
|
|
self.factor = config.factor
|
|
|
|
def __call__(self, *args, **kwargs):
|
|
# This is a dummy method to do checked forward passes.
|
|
return self._forward(*args, **kwargs)
|
|
|
|
def _forward(self, input_dict, **kwargs):
|
|
obs = input_dict[Columns.OBS]
|
|
return {
|
|
ENCODER_OUT: int(np.array(obs) * self.factor),
|
|
Columns.STATE_OUT: int(
|
|
np.array(input_dict[Columns.STATE_IN])
|
|
* self.factor
|
|
),
|
|
}
|
|
|
|
def get_num_parameters(self):
|
|
return (0, 0)
|
|
|
|
def _set_to_dummy_weights(self, value_sequence=(-0.02, -0.01, 0.01, 0.02)):
|
|
pass
|
|
|
|
@dataclass
|
|
class NumpyEncoderConfig(ModelConfig):
|
|
factor: int = None
|
|
|
|
def build(self, framework: str):
|
|
return NumpyEncoder(self)
|
|
|
|
config = NumpyEncoderConfig(factor=2)
|
|
encoder = NumpyEncoder(config)
|
|
print(encoder({Columns.OBS: 1, Columns.STATE_IN: 2}))
|
|
|
|
.. testoutput::
|
|
|
|
{'encoder_out': 2, 'state_out': 4}
|
|
|
|
"""
|
|
|
|
@abc.abstractmethod
|
|
def _forward(self, input_dict: dict, **kwargs) -> dict:
|
|
"""Returns the latent of the encoder for the given inputs.
|
|
|
|
This method is called by the forwarding method of the respective framework
|
|
that is itself wrapped by RLlib in order to check model inputs and outputs.
|
|
|
|
The input dict contains at minimum the observation and the state of the encoder
|
|
(None for stateless encoders).
|
|
The output dict contains at minimum the latent and the state of the encoder
|
|
(None for stateless encoders).
|
|
To establish an agreement between the encoder and RLModules, these values
|
|
have the fixed keys `Columns.OBS` for the `input_dict`,
|
|
and `ACTOR` and `CRITIC` for the returned dict.
|
|
|
|
Args:
|
|
input_dict: The input tensors. Must contain at a minimum the keys
|
|
Columns.OBS and Columns.STATE_IN (which might be None for stateless
|
|
encoders).
|
|
**kwargs: Forward compatibility kwargs.
|
|
|
|
Returns:
|
|
The output tensors. Must contain at a minimum the key ENCODER_OUT.
|
|
"""
|
|
|
|
|
|
@ExperimentalAPI
|
|
class ActorCriticEncoder(Encoder):
|
|
"""An encoder that potentially holds two stateless encoders.
|
|
|
|
This is a special case of Encoder that can either enclose a single,
|
|
shared encoder or two separate encoders: One for the actor and one for the
|
|
critic. The two encoders are of the same type, and we can therefore make the
|
|
assumption that they have the same input and output specs.
|
|
"""
|
|
|
|
framework = None
|
|
|
|
def __init__(self, config: ModelConfig) -> None:
|
|
super().__init__(config)
|
|
|
|
if config.shared:
|
|
self.encoder = config.base_encoder_config.build(framework=self.framework)
|
|
else:
|
|
self.actor_encoder = config.base_encoder_config.build(
|
|
framework=self.framework
|
|
)
|
|
self.critic_encoder = None
|
|
if not config.inference_only:
|
|
self.critic_encoder = config.base_encoder_config.build(
|
|
framework=self.framework
|
|
)
|
|
|
|
@override(Model)
|
|
def _forward(self, inputs: dict, **kwargs) -> dict:
|
|
if self.config.shared:
|
|
encoder_outs = self.encoder(inputs, **kwargs)
|
|
return {
|
|
ENCODER_OUT: {
|
|
ACTOR: encoder_outs[ENCODER_OUT],
|
|
**(
|
|
{}
|
|
if self.config.inference_only
|
|
else {CRITIC: encoder_outs[ENCODER_OUT]}
|
|
),
|
|
}
|
|
}
|
|
else:
|
|
# Encoders should not modify inputs, so we can pass the same inputs
|
|
actor_out = self.actor_encoder(inputs, **kwargs)
|
|
if self.critic_encoder:
|
|
critic_out = self.critic_encoder(inputs, **kwargs)
|
|
|
|
return {
|
|
ENCODER_OUT: {
|
|
ACTOR: actor_out[ENCODER_OUT],
|
|
**(
|
|
{}
|
|
if self.config.inference_only
|
|
else {CRITIC: critic_out[ENCODER_OUT]}
|
|
),
|
|
}
|
|
}
|
|
|
|
|
|
@ExperimentalAPI
|
|
class StatefulActorCriticEncoder(Encoder):
|
|
"""An encoder that potentially holds two potentially stateful encoders.
|
|
|
|
This is a special case of Encoder that can either enclose a single,
|
|
shared encoder or two separate encoders: One for the actor and one for the
|
|
critic. The two encoders are of the same type, and we can therefore make the
|
|
assumption that they have the same input and output specs.
|
|
|
|
If this encoder wraps a single encoder, state in input- and output dicts
|
|
is simply stored under the key `STATE_IN` and `STATE_OUT`, respectively.
|
|
If this encoder wraps two encoders, state in input- and output dicts is
|
|
stored under the keys `(STATE_IN, ACTOR)` and `(STATE_IN, CRITIC)` and
|
|
`(STATE_OUT, ACTOR)` and `(STATE_OUT, CRITIC)`, respectively.
|
|
"""
|
|
|
|
framework = None
|
|
|
|
def __init__(self, config: ModelConfig) -> None:
|
|
super().__init__(config)
|
|
|
|
if config.shared:
|
|
self.encoder = config.base_encoder_config.build(framework=self.framework)
|
|
else:
|
|
self.actor_encoder = config.base_encoder_config.build(
|
|
framework=self.framework
|
|
)
|
|
self.critic_encoder = config.base_encoder_config.build(
|
|
framework=self.framework
|
|
)
|
|
|
|
@override(Model)
|
|
def get_initial_state(self):
|
|
if self.config.shared:
|
|
return self.encoder.get_initial_state()
|
|
else:
|
|
return {
|
|
ACTOR: self.actor_encoder.get_initial_state(),
|
|
CRITIC: self.critic_encoder.get_initial_state(),
|
|
}
|
|
|
|
@override(Model)
|
|
def _forward(self, inputs: dict, **kwargs) -> dict:
|
|
outputs = {}
|
|
|
|
if self.config.shared:
|
|
outs = self.encoder(inputs, **kwargs)
|
|
encoder_out = outs.pop(ENCODER_OUT)
|
|
outputs[ENCODER_OUT] = {ACTOR: encoder_out, CRITIC: encoder_out}
|
|
outputs[Columns.STATE_OUT] = outs[Columns.STATE_OUT]
|
|
else:
|
|
# Shallow copy inputs so that we can add states without modifying
|
|
# original dict.
|
|
actor_inputs = inputs.copy()
|
|
critic_inputs = inputs.copy()
|
|
actor_inputs[Columns.STATE_IN] = inputs[Columns.STATE_IN][ACTOR]
|
|
critic_inputs[Columns.STATE_IN] = inputs[Columns.STATE_IN][CRITIC]
|
|
|
|
actor_out = self.actor_encoder(actor_inputs, **kwargs)
|
|
critic_out = self.critic_encoder(critic_inputs, **kwargs)
|
|
|
|
outputs[ENCODER_OUT] = {
|
|
ACTOR: actor_out[ENCODER_OUT],
|
|
CRITIC: critic_out[ENCODER_OUT],
|
|
}
|
|
|
|
outputs[Columns.STATE_OUT] = {
|
|
ACTOR: actor_out[Columns.STATE_OUT],
|
|
CRITIC: critic_out[Columns.STATE_OUT],
|
|
}
|
|
|
|
return outputs
|
|
|
|
|
|
@DeveloperAPI
|
|
def tokenize(tokenizer: Encoder, inputs: dict, framework: str) -> dict:
|
|
"""Tokenizes the observations from the input dict.
|
|
|
|
Args:
|
|
tokenizer: The tokenizer to use.
|
|
inputs: The input dict.
|
|
|
|
Returns:
|
|
The output dict.
|
|
"""
|
|
# Tokenizer may depend solely on observations.
|
|
obs = inputs[Columns.OBS]
|
|
tokenizer_inputs = {Columns.OBS: obs}
|
|
size = list(obs.size() if framework == "torch" else obs.shape)
|
|
b_dim, t_dim = size[:2]
|
|
fold, unfold = get_fold_unfold_fns(b_dim, t_dim, framework=framework)
|
|
# Push through the tokenizer encoder.
|
|
out = tokenizer(fold(tokenizer_inputs))
|
|
out = out[ENCODER_OUT]
|
|
# Then unfold batch- and time-dimensions again.
|
|
return unfold(out)
|