diff --git a/rllib/models/catalog.py b/rllib/models/catalog.py index 1a2650d66..8c241a3fa 100644 --- a/rllib/models/catalog.py +++ b/rllib/models/catalog.py @@ -532,7 +532,9 @@ class ModelCatalog: return wrapper @staticmethod - def _get_v2_model_class(input_space, model_config, framework="tf"): + def _get_v2_model_class(input_space: gym.Space, + model_config: ModelConfigDict, + framework: str = "tf") -> ModelV2: if framework == "torch": from ray.rllib.models.torch.fcnet import (FullyConnectedNetwork as FCNet) diff --git a/rllib/models/preprocessors.py b/rllib/models/preprocessors.py index c31cc533f..51a6f5442 100644 --- a/rllib/models/preprocessors.py +++ b/rllib/models/preprocessors.py @@ -7,6 +7,7 @@ from typing import Any, List from ray.rllib.utils.annotations import override, PublicAPI from ray.rllib.utils.spaces.repeated import Repeated +from ray.rllib.utils.typing import TensorType ATARI_OBS_SHAPE = (210, 160, 3) ATARI_RAM_OBS_SHAPE = (128, ) @@ -42,11 +43,12 @@ class Preprocessor: raise NotImplementedError @PublicAPI - def transform(self, observation: Any) -> np.ndarray: + def transform(self, observation: TensorType) -> np.ndarray: """Returns the preprocessed observation.""" raise NotImplementedError - def write(self, observation: Any, array: np.ndarray, offset: int) -> None: + def write(self, observation: TensorType, array: np.ndarray, + offset: int) -> None: """Alternative to transform for more efficient flattening.""" array[offset:offset + self._size] = self.transform(observation) @@ -93,7 +95,7 @@ class GenericPixelPreprocessor(Preprocessor): """ @override(Preprocessor) - def _init_shape(self, obs_space, options): + def _init_shape(self, obs_space: gym.Space, options: dict) -> List[int]: self._grayscale = options.get("grayscale") self._zero_mean = options.get("zero_mean") self._dim = options.get("dim") @@ -105,7 +107,7 @@ class GenericPixelPreprocessor(Preprocessor): return shape @override(Preprocessor) - def transform(self, observation): + def transform(self, observation: TensorType) -> np.ndarray: """Downsamples images from (210, 160, 3) by the configured factor.""" self.check_shape(observation) scaled = observation[25:-25, :, :] @@ -129,50 +131,52 @@ class GenericPixelPreprocessor(Preprocessor): class AtariRamPreprocessor(Preprocessor): @override(Preprocessor) - def _init_shape(self, obs_space, options): + def _init_shape(self, obs_space: gym.Space, options: dict) -> List[int]: return (128, ) @override(Preprocessor) - def transform(self, observation): + def transform(self, observation: TensorType) -> np.ndarray: self.check_shape(observation) return (observation - 128) / 128 class OneHotPreprocessor(Preprocessor): @override(Preprocessor) - def _init_shape(self, obs_space, options): + def _init_shape(self, obs_space: gym.Space, options: dict) -> List[int]: return (self._obs_space.n, ) @override(Preprocessor) - def transform(self, observation): + def transform(self, observation: TensorType) -> np.ndarray: self.check_shape(observation) arr = np.zeros(self._obs_space.n, dtype=np.float32) arr[observation] = 1 return arr @override(Preprocessor) - def write(self, observation, array, offset): + def write(self, observation: TensorType, array: np.ndarray, + offset: int) -> None: array[offset + observation] = 1 class NoPreprocessor(Preprocessor): @override(Preprocessor) - def _init_shape(self, obs_space, options): + def _init_shape(self, obs_space: gym.Space, options: dict) -> List[int]: return self._obs_space.shape @override(Preprocessor) - def transform(self, observation): + def transform(self, observation: TensorType) -> np.ndarray: self.check_shape(observation) return observation @override(Preprocessor) - def write(self, observation, array, offset): + def write(self, observation: TensorType, array: np.ndarray, + offset: int) -> None: array[offset:offset + self._size] = np.array( observation, copy=False).ravel() @property @override(Preprocessor) - def observation_space(self): + def observation_space(self) -> gym.Space: return self._obs_space @@ -183,7 +187,7 @@ class TupleFlatteningPreprocessor(Preprocessor): """ @override(Preprocessor) - def _init_shape(self, obs_space, options): + def _init_shape(self, obs_space: gym.Space, options: dict) -> List[int]: assert isinstance(self._obs_space, gym.spaces.Tuple) size = 0 self.preprocessors = [] @@ -196,14 +200,15 @@ class TupleFlatteningPreprocessor(Preprocessor): return (size, ) @override(Preprocessor) - def transform(self, observation): + def transform(self, observation: TensorType) -> np.ndarray: self.check_shape(observation) array = np.zeros(self.shape) self.write(observation, array, 0) return array @override(Preprocessor) - def write(self, observation, array, offset): + def write(self, observation: TensorType, array: np.ndarray, + offset: int) -> None: assert len(observation) == len(self.preprocessors), observation for o, p in zip(observation, self.preprocessors): p.write(o, array, offset) @@ -217,7 +222,7 @@ class DictFlatteningPreprocessor(Preprocessor): """ @override(Preprocessor) - def _init_shape(self, obs_space, options): + def _init_shape(self, obs_space: gym.Space, options: dict) -> List[int]: assert isinstance(self._obs_space, gym.spaces.Dict) size = 0 self.preprocessors = [] @@ -229,14 +234,15 @@ class DictFlatteningPreprocessor(Preprocessor): return (size, ) @override(Preprocessor) - def transform(self, observation): + def transform(self, observation: TensorType) -> np.ndarray: self.check_shape(observation) array = np.zeros(self.shape) self.write(observation, array, 0) return array @override(Preprocessor) - def write(self, observation, array, offset): + def write(self, observation: TensorType, array: np.ndarray, + offset: int) -> None: if not isinstance(observation, OrderedDict): observation = OrderedDict(sorted(observation.items())) assert len(observation) == len(self.preprocessors), \ @@ -250,7 +256,7 @@ class RepeatedValuesPreprocessor(Preprocessor): """Pads and batches the variable-length list value.""" @override(Preprocessor) - def _init_shape(self, obs_space, options): + def _init_shape(self, obs_space: gym.Space, options: dict) -> List[int]: assert isinstance(self._obs_space, Repeated) child_space = obs_space.child_space self.child_preprocessor = get_preprocessor(child_space)(child_space, @@ -260,7 +266,7 @@ class RepeatedValuesPreprocessor(Preprocessor): return (size, ) @override(Preprocessor) - def transform(self, observation): + def transform(self, observation: TensorType) -> np.ndarray: array = np.zeros(self.shape) if isinstance(observation, list): for elem in observation: @@ -271,7 +277,8 @@ class RepeatedValuesPreprocessor(Preprocessor): return array @override(Preprocessor) - def write(self, observation, array, offset): + def write(self, observation: TensorType, array: np.ndarray, + offset: int) -> None: if not isinstance(observation, list): raise ValueError("Input for {} must be list type, got {}".format( self, observation)) diff --git a/rllib/models/torch/attention_net.py b/rllib/models/torch/attention_net.py index dc1453eb6..58480a64b 100644 --- a/rllib/models/torch/attention_net.py +++ b/rllib/models/torch/attention_net.py @@ -9,6 +9,7 @@ https://www.aclweb.org/anthology/P19-1285.pdf """ import numpy as np +import gym from ray.rllib.models.modelv2 import ModelV2 from ray.rllib.models.torch.misc import SlimFC @@ -17,11 +18,12 @@ from ray.rllib.models.torch.modules import GRUGate, \ from ray.rllib.models.torch.recurrent_net import RecurrentNetwork from ray.rllib.utils.annotations import override from ray.rllib.utils.framework import try_import_torch +from ray.rllib.utils.typing import ModelConfigDict, TensorType, List torch, nn = try_import_torch() -def relative_position_embedding(seq_length, out_dim): +def relative_position_embedding(seq_length: int, out_dim: int) -> TensorType: """Creates a [seq_length x seq_length] matrix for rel. pos encoding. Denoted as Phi in [2] and [3]. Phi is the standard sinusoid encoding @@ -64,18 +66,18 @@ class GTrXLNet(RecurrentNetwork, nn.Module): """ def __init__(self, - observation_space, - action_space, - num_outputs, - model_config, - name, - num_transformer_units, - attn_dim, - num_heads, - memory_tau, - head_dim, - ff_hidden_dim, - init_gate_bias=2.0): + observation_space: gym.spaces.Space, + action_space: gym.spaces.Space, + num_outputs: int, + model_config: ModelConfigDict, + name: str, + num_transformer_units: int, + attn_dim: int, + num_heads: int, + memory_tau: int, + head_dim: int, + ff_hidden_dim: int, + init_gate_bias: float = 2.0): """Initializes a GTrXLNet. Args: @@ -167,7 +169,8 @@ class GTrXLNet(RecurrentNetwork, nn.Module): in_size=self.attn_dim, out_size=1, activation_fn=None) @override(RecurrentNetwork) - def forward_rnn(self, inputs, state, seq_lens): + def forward_rnn(self, inputs: TensorType, state: List[TensorType], + seq_lens: TensorType) -> (TensorType, List[TensorType]): # To make Attention work with current RLlib's ModelV2 API: # We assume `state` is the history of L recent observations (all # concatenated into one tensor) and append the current inputs to the @@ -214,7 +217,7 @@ class GTrXLNet(RecurrentNetwork, nn.Module): return logits, [observations] + memory_outs @override(RecurrentNetwork) - def get_initial_state(self): + def get_initial_state(self) -> List[np.ndarray]: # State is the T last observations concat'd together into one Tensor. # Plus all Transformer blocks' E(l) outputs concat'd together (up to # tau timesteps). @@ -223,5 +226,5 @@ class GTrXLNet(RecurrentNetwork, nn.Module): for _ in range(self.num_transformer_units)] @override(ModelV2) - def value_function(self): + def value_function(self) -> TensorType: return torch.reshape(self._value_out, [-1]) diff --git a/rllib/models/torch/fcnet.py b/rllib/models/torch/fcnet.py index 4080c9318..e9d779c30 100644 --- a/rllib/models/torch/fcnet.py +++ b/rllib/models/torch/fcnet.py @@ -1,11 +1,13 @@ import logging import numpy as np +import gym from ray.rllib.models.torch.torch_modelv2 import TorchModelV2 from ray.rllib.models.torch.misc import SlimFC, AppendBiasLayer, \ normc_initializer from ray.rllib.utils.annotations import override from ray.rllib.utils.framework import try_import_torch +from ray.rllib.utils.typing import Dict, TensorType, List, ModelConfigDict torch, nn = try_import_torch() @@ -15,8 +17,9 @@ logger = logging.getLogger(__name__) class FullyConnectedNetwork(TorchModelV2, nn.Module): """Generic fully connected network.""" - def __init__(self, obs_space, action_space, num_outputs, model_config, - name): + def __init__(self, obs_space: gym.spaces.Space, + action_space: gym.spaces.Space, num_outputs: int, + model_config: ModelConfigDict, name: str): TorchModelV2.__init__(self, obs_space, action_space, num_outputs, model_config, name) nn.Module.__init__(self) @@ -111,7 +114,9 @@ class FullyConnectedNetwork(TorchModelV2, nn.Module): self._last_flat_in = None @override(TorchModelV2) - def forward(self, input_dict, state, seq_lens): + def forward(self, input_dict: Dict[str, TensorType], + state: List[TensorType], + seq_lens: TensorType) -> (TensorType, List[TensorType]): obs = input_dict["obs_flat"].float() self._last_flat_in = obs.reshape(obs.shape[0], -1) self._features = self._hidden_layers(self._last_flat_in) @@ -122,7 +127,7 @@ class FullyConnectedNetwork(TorchModelV2, nn.Module): return logits, state @override(TorchModelV2) - def value_function(self): + def value_function(self) -> TensorType: assert self._features is not None, "must call forward() first" if self._value_branch_separate: return self._value_branch( diff --git a/rllib/models/torch/misc.py b/rllib/models/torch/misc.py index 977a5d907..e68fac3a9 100644 --- a/rllib/models/torch/misc.py +++ b/rllib/models/torch/misc.py @@ -1,13 +1,14 @@ """ Code adapted from https://github.com/ikostrikov/pytorch-a3c""" import numpy as np -from typing import List +from typing import Union, Tuple, Any, List from ray.rllib.utils.framework import get_activation_fn, try_import_torch +from ray.rllib.utils.typing import TensorType torch, nn = try_import_torch() -def normc_initializer(std=1.0): +def normc_initializer(std: float = 1.0) -> Any: def initializer(tensor): tensor.data.normal_(0, 1) tensor.data *= std / torch.sqrt( @@ -16,7 +17,9 @@ def normc_initializer(std=1.0): return initializer -def same_padding(in_size, filter_size, stride_size): +def same_padding(in_size: Tuple[int, int], filter_size: Tuple[int, int], + stride_size: Union[int, Tuple[int, int]] + ) -> (Union[int, Tuple[int, int]], Tuple[int, int]): """Note: Padding is added to match TF conv2d `same` padding. See www.tensorflow.org/versions/r0.12/api_docs/python/nn/convolution @@ -58,15 +61,31 @@ class SlimConv2d(nn.Module): def __init__( self, - in_channels, - out_channels, - kernel, - stride, - padding, + in_channels: int, + out_channels: int, + kernel: Union[int, Tuple[int, int]], + stride: Union[int, Tuple[int, int]], + padding: Union[int, Tuple[int, int]], # Defaulting these to nn.[..] will break soft torch import. - initializer="default", - activation_fn="default", - bias_init=0): + initializer: Any = "default", + activation_fn: Any = "default", + bias_init: float = 0): + """Creates a standard Conv2d layer, similar to torch.nn.Conv2d + + Args: + in_channels(int): Number of input channels + out_channels (int): Number of output channels + kernel (Union[int, Tuple[int, int]]): If int, the kernel is + a tuple(x,x). Elsewise, the tuple can be specified + stride (Union[int, Tuple[int, int]]): Controls the stride + for the cross-correlation. If int, the stride is a + tuple(x,x). Elsewise, the tuple can be specified + padding (Union[int, Tuple[int, int]]): Controls the amount + of implicit zero-paddings during the conv operation + initializer (Any): Initializer function for kernel weights + activation_fn (Any): Activation function at the end of layer + bias_init (float): Initalize bias weights to bias_init const + """ super(SlimConv2d, self).__init__() layers = [] # Padding layer. @@ -91,7 +110,7 @@ class SlimConv2d(nn.Module): # Put everything in sequence. self._model = nn.Sequential(*layers) - def forward(self, x): + def forward(self, x: TensorType) -> TensorType: return self._model(x) @@ -99,12 +118,22 @@ class SlimFC(nn.Module): """Simple PyTorch version of `linear` function""" def __init__(self, - in_size, - out_size, - initializer=None, - activation_fn=None, - use_bias=True, - bias_init=0.0): + in_size: int, + out_size: int, + initializer: Any = None, + activation_fn: Any = None, + use_bias: bool = True, + bias_init: float = 0.0): + """Creates a standard FC layer, similar to torch.nn.Linear + + Args: + in_size(int): Input size for FC Layer + out_size (int): Output size for FC Layer + initializer (Any): Initializer function for FC layer weights + activation_fn (Any): Activation function at the end of layer + use_bias (bool): Whether to add bias weights or not + bias_init (float): Initalize bias weights to bias_init const + """ super(SlimFC, self).__init__() layers = [] # Actual Conv2D layer (including correct initialization logic). @@ -122,20 +151,20 @@ class SlimFC(nn.Module): # Put everything in sequence. self._model = nn.Sequential(*layers) - def forward(self, x): + def forward(self, x: TensorType) -> TensorType: return self._model(x) class AppendBiasLayer(nn.Module): """Simple bias appending layer for free_log_std.""" - def __init__(self, num_bias_vars): + def __init__(self, num_bias_vars: int): super().__init__() self.log_std = torch.nn.Parameter( torch.as_tensor([0.0] * num_bias_vars)) self.register_parameter("log_std", self.log_std) - def forward(self, x): + def forward(self, x: TensorType) -> TensorType: out = torch.cat( [x, self.log_std.unsqueeze(0).repeat([len(x), 1])], axis=1) return out diff --git a/rllib/models/torch/modules/gru_gate.py b/rllib/models/torch/modules/gru_gate.py index 76d0f9eb3..4cabc5eb2 100644 --- a/rllib/models/torch/modules/gru_gate.py +++ b/rllib/models/torch/modules/gru_gate.py @@ -1,4 +1,5 @@ from ray.rllib.utils.framework import try_import_torch +from ray.rllib.utils.framework import TensorType torch, nn = try_import_torch() @@ -6,7 +7,7 @@ torch, nn = try_import_torch() class GRUGate(nn.Module): """Implements a gated recurrent unit for use in AttentionNet""" - def __init__(self, dim, init_bias=0., **kwargs): + def __init__(self, dim: int, init_bias: int = 0., **kwargs): """ input_shape (torch.Tensor): dimension of the input init_bias (int): Bias added to every input to stabilize training @@ -33,7 +34,7 @@ class GRUGate(nn.Module): self._bias_z = torch.zeros(dim, ).fill_(self._init_bias) - def forward(self, inputs, **kwargs): + def forward(self, inputs: TensorType, **kwargs) -> TensorType: # Pass in internal state first. h, X = inputs diff --git a/rllib/models/torch/modules/multi_head_attention.py b/rllib/models/torch/modules/multi_head_attention.py index ea27ad650..09ce5b795 100644 --- a/rllib/models/torch/modules/multi_head_attention.py +++ b/rllib/models/torch/modules/multi_head_attention.py @@ -6,6 +6,7 @@ from ray.rllib.utils.framework import try_import_torch from ray.rllib.models.torch.misc import SlimFC from ray.rllib.utils.torch_ops import sequence_mask +from ray.rllib.utils.framework import TensorType torch, nn = try_import_torch() @@ -13,7 +14,8 @@ torch, nn = try_import_torch() class MultiHeadAttention(nn.Module): """A multi-head attention layer described in [1].""" - def __init__(self, in_dim, out_dim, num_heads, head_dim, **kwargs): + def __init__(self, in_dim: int, out_dim: int, num_heads: int, + head_dim: int, **kwargs): """ in_dim (int): Dimension of input out_dim (int): Dimension of output @@ -31,7 +33,7 @@ class MultiHeadAttention(nn.Module): self._linear_layer = SlimFC( in_size=num_heads * head_dim, out_size=out_dim, use_bias=False) - def forward(self, inputs): + def forward(self, inputs: TensorType) -> TensorType: L = list(inputs.size())[1] # length of segment H = self._num_heads # number of attention heads D = self._head_dim # attention head dimension diff --git a/rllib/models/torch/modules/noisy_layer.py b/rllib/models/torch/modules/noisy_layer.py index 547ac74ae..ee553c73c 100644 --- a/rllib/models/torch/modules/noisy_layer.py +++ b/rllib/models/torch/modules/noisy_layer.py @@ -1,6 +1,7 @@ import numpy as np from ray.rllib.utils.framework import get_activation_fn, try_import_torch +from ray.rllib.utils.framework import TensorType torch, nn = try_import_torch() @@ -17,14 +18,18 @@ class NoisyLayer(nn.Module): vanish along the training procedure. """ - def __init__(self, in_size, out_size, sigma0, activation="relu"): + def __init__(self, + in_size: int, + out_size: int, + sigma0: float, + activation: str = "relu"): """Initializes a NoisyLayer object. Args: - in_size: - out_size: - sigma0: - non_linear: + in_size: Input size for Noisy Layer + out_size: Output size for Noisy Layer + sigma0: Initialization value for sigma_b (bias noise) + activation: Non-linear activation for Noisy Layer """ super().__init__() @@ -59,7 +64,7 @@ class NoisyLayer(nn.Module): b = nn.Parameter(torch.from_numpy(np.zeros([out_size])).float()) self.register_parameter("b", b) - def forward(self, inputs): + def forward(self, inputs: TensorType) -> TensorType: epsilon_in = self._f_epsilon( torch.normal( mean=torch.zeros([self.in_size]), @@ -81,5 +86,5 @@ class NoisyLayer(nn.Module): action_activation = self.activation(action_activation) return action_activation - def _f_epsilon(self, x): + def _f_epsilon(self, x: TensorType) -> TensorType: return torch.sign(x) * torch.pow(torch.abs(x), 0.5) diff --git a/rllib/models/torch/modules/relative_multi_head_attention.py b/rllib/models/torch/modules/relative_multi_head_attention.py index 204495517..fe28d6f73 100644 --- a/rllib/models/torch/modules/relative_multi_head_attention.py +++ b/rllib/models/torch/modules/relative_multi_head_attention.py @@ -1,6 +1,7 @@ from ray.rllib.utils.framework import try_import_torch from ray.rllib.models.torch.misc import SlimFC from ray.rllib.utils.torch_ops import sequence_mask +from ray.rllib.utils.typing import TensorType, Any torch, nn = try_import_torch() @@ -12,13 +13,13 @@ class RelativeMultiHeadAttention(nn.Module): """ def __init__(self, - in_dim, - out_dim, - num_heads, - head_dim, - rel_pos_encoder, - input_layernorm=False, - output_activation=None, + in_dim: int, + out_dim: int, + num_heads: int, + head_dim: int, + rel_pos_encoder: Any, + input_layernorm: bool = False, + output_activation: Any = None, **kwargs): """Initializes a RelativeMultiHeadAttention nn.Module object. @@ -66,7 +67,8 @@ class RelativeMultiHeadAttention(nn.Module): if input_layernorm: self._input_layernorm = torch.nn.LayerNorm(in_dim) - def forward(self, inputs, memory=None): + def forward(self, inputs: TensorType, + memory: TensorType = None) -> TensorType: T = list(inputs.size())[1] # length of segment (time) H = self._num_heads # number of attention heads d = self._head_dim # attention head dimension @@ -119,7 +121,7 @@ class RelativeMultiHeadAttention(nn.Module): return self._linear_layer(out) @staticmethod - def rel_shift(x): + def rel_shift(x: TensorType) -> TensorType: # Transposed version of the shift approach described in [3]. # https://github.com/kimiyoung/transformer-xl/blob/ # 44781ed21dbaec88b280f74d9ae2877f52b492a5/tf/model.py#L31 diff --git a/rllib/models/torch/modules/skip_connection.py b/rllib/models/torch/modules/skip_connection.py index 9c85eea85..126274b1d 100644 --- a/rllib/models/torch/modules/skip_connection.py +++ b/rllib/models/torch/modules/skip_connection.py @@ -1,4 +1,6 @@ from ray.rllib.utils.framework import try_import_torch +from ray.rllib.utils.typing import TensorType +from typing import Optional torch, nn = try_import_torch() @@ -10,7 +12,11 @@ class SkipConnection(nn.Module): input as hidden state input to a given fan_in_layer. """ - def __init__(self, layer, fan_in_layer=None, add_memory=False, **kwargs): + def __init__(self, + layer: nn.Module, + fan_in_layer: Optional[nn.Module] = None, + add_memory: bool = False, + **kwargs): """Initializes a SkipConnection nn Module object. Args: @@ -23,7 +29,7 @@ class SkipConnection(nn.Module): self._layer = layer self._fan_in_layer = fan_in_layer - def forward(self, inputs, **kwargs): + def forward(self, inputs: TensorType, **kwargs) -> TensorType: # del kwargs outputs = self._layer(inputs, **kwargs) # Residual case, just add inputs to outputs. diff --git a/rllib/models/torch/recurrent_net.py b/rllib/models/torch/recurrent_net.py index c500450d0..d48fa737f 100644 --- a/rllib/models/torch/recurrent_net.py +++ b/rllib/models/torch/recurrent_net.py @@ -1,5 +1,7 @@ from gym.spaces import Box import numpy as np +import gym +from typing import Dict, List, Union from ray.rllib.models.modelv2 import ModelV2 from ray.rllib.models.torch.misc import SlimFC @@ -9,6 +11,7 @@ from ray.rllib.policy.sample_batch import SampleBatch from ray.rllib.policy.view_requirement import ViewRequirement from ray.rllib.utils.annotations import override, DeveloperAPI from ray.rllib.utils.framework import try_import_torch +from ray.rllib.utils.typing import ModelConfigDict, TensorType torch, nn = try_import_torch() @@ -59,7 +62,9 @@ class RecurrentNetwork(TorchModelV2): """ @override(ModelV2) - def forward(self, input_dict, state, seq_lens): + def forward(self, input_dict: Dict[str, TensorType], + state: List[TensorType], + seq_lens: TensorType) -> (TensorType, List[TensorType]): """Adds time dimension to batch before sending inputs to forward_rnn(). You should implement forward_rnn() in your subclass.""" @@ -78,7 +83,8 @@ class RecurrentNetwork(TorchModelV2): output = torch.reshape(output, [-1, self.num_outputs]) return output, new_state - def forward_rnn(self, inputs, state, seq_lens): + def forward_rnn(self, inputs: TensorType, state: List[TensorType], + seq_lens: TensorType) -> (TensorType, List[TensorType]): """Call the model with the given input tensors and state. Args: @@ -104,8 +110,9 @@ class LSTMWrapper(RecurrentNetwork, nn.Module): """An LSTM wrapper serving as an interface for ModelV2s that set use_lstm. """ - def __init__(self, obs_space, action_space, num_outputs, model_config, - name): + def __init__(self, obs_space: gym.spaces.Space, + action_space: gym.spaces.Space, num_outputs: int, + model_config: ModelConfigDict, name: str): nn.Module.__init__(self) super().__init__(obs_space, action_space, None, model_config, name) @@ -154,7 +161,9 @@ class LSTMWrapper(RecurrentNetwork, nn.Module): space=Box(-1.0, 1.0, shape=(self.cell_size,))) @override(RecurrentNetwork) - def forward(self, input_dict, state, seq_lens): + def forward(self, input_dict: Dict[str, TensorType], + state: List[TensorType], + seq_lens: TensorType) -> (TensorType, List[TensorType]): assert seq_lens is not None # Push obs through "unwrapped" net's `forward()` first. wrapped_out, _ = self._wrapped_forward(input_dict, [], None) @@ -176,7 +185,8 @@ class LSTMWrapper(RecurrentNetwork, nn.Module): return super().forward(input_dict, state, seq_lens) @override(RecurrentNetwork) - def forward_rnn(self, inputs, state, seq_lens): + def forward_rnn(self, inputs: TensorType, state: List[TensorType], + seq_lens: TensorType) -> (TensorType, List[TensorType]): # Don't show paddings to RNN(?) # TODO: (sven) For now, only allow, iff time_major=True to not break # anything retrospectively (time_major not supported previously). @@ -199,7 +209,7 @@ class LSTMWrapper(RecurrentNetwork, nn.Module): return model_out, [torch.squeeze(h, 0), torch.squeeze(c, 0)] @override(ModelV2) - def get_initial_state(self): + def get_initial_state(self) -> Union[List[np.ndarray], List[TensorType]]: # Place hidden states on same device as model. linear = next(self._logits_branch._model.children()) h = [ @@ -209,6 +219,6 @@ class LSTMWrapper(RecurrentNetwork, nn.Module): return h @override(ModelV2) - def value_function(self): + def value_function(self) -> TensorType: assert self._features is not None, "must call forward() first" return torch.reshape(self._value_branch(self._features), [-1]) diff --git a/rllib/models/torch/torch_action_dist.py b/rllib/models/torch/torch_action_dist.py index 57b197b80..2f0f7ed4c 100644 --- a/rllib/models/torch/torch_action_dist.py +++ b/rllib/models/torch/torch_action_dist.py @@ -2,6 +2,7 @@ import functools from math import log import numpy as np import tree +import gym from ray.rllib.models.action_dist import ActionDistribution from ray.rllib.models.torch.torch_modelv2 import TorchModelV2 @@ -11,7 +12,8 @@ from ray.rllib.utils.numpy import SMALL_NUMBER, MIN_LOG_NN_OUTPUT, \ MAX_LOG_NN_OUTPUT from ray.rllib.utils.spaces.space_utils import get_base_struct_from_space from ray.rllib.utils.torch_ops import atanh -from ray.rllib.utils.typing import TensorType, List +from ray.rllib.utils.typing import TensorType, List, Union, \ + Tuple, ModelConfigDict torch, nn = try_import_torch() @@ -58,7 +60,10 @@ class TorchCategorical(TorchDistributionWrapper): """Wrapper class for PyTorch Categorical distribution.""" @override(ActionDistribution) - def __init__(self, inputs, model=None, temperature=1.0): + def __init__(self, + inputs: List[TensorType], + model: TorchModelV2 = None, + temperature: float = 1.0): if temperature != 1.0: assert temperature > 0.0, \ "Categorical `temperature` must be > 0.0!" @@ -68,13 +73,15 @@ class TorchCategorical(TorchDistributionWrapper): logits=self.inputs) @override(ActionDistribution) - def deterministic_sample(self): + def deterministic_sample(self) -> TensorType: self.last_sample = self.dist.probs.argmax(dim=1) return self.last_sample @staticmethod @override(ActionDistribution) - def required_model_output_shape(action_space, model_config): + def required_model_output_shape( + action_space: gym.Space, + model_config: ModelConfigDict) -> Union[int, np.ndarray]: return action_space.n @@ -82,7 +89,8 @@ class TorchMultiCategorical(TorchDistributionWrapper): """MultiCategorical distribution for MultiDiscrete action spaces.""" @override(TorchDistributionWrapper) - def __init__(self, inputs, model, input_lens): + def __init__(self, inputs: List[TensorType], model: TorchModelV2, + input_lens: Union[List[int], np.ndarray, Tuple[int, ...]]): super().__init__(inputs, model) # If input_lens is np.ndarray or list, force-make it a tuple. inputs_split = self.inputs.split(tuple(input_lens), dim=1) @@ -92,19 +100,19 @@ class TorchMultiCategorical(TorchDistributionWrapper): ] @override(TorchDistributionWrapper) - def sample(self): + def sample(self) -> TensorType: arr = [cat.sample() for cat in self.cats] self.last_sample = torch.stack(arr, dim=1) return self.last_sample @override(ActionDistribution) - def deterministic_sample(self): + def deterministic_sample(self) -> TensorType: arr = [torch.argmax(cat.probs, -1) for cat in self.cats] self.last_sample = torch.stack(arr, dim=1) return self.last_sample @override(TorchDistributionWrapper) - def logp(self, actions): + def logp(self, actions: TensorType) -> TensorType: # # If tensor is provided, unstack it into list. if isinstance(actions, torch.Tensor): actions = torch.unbind(actions, dim=1) @@ -113,15 +121,15 @@ class TorchMultiCategorical(TorchDistributionWrapper): return torch.sum(logps, dim=0) @override(ActionDistribution) - def multi_entropy(self): + def multi_entropy(self) -> TensorType: return torch.stack([cat.entropy() for cat in self.cats], dim=1) @override(TorchDistributionWrapper) - def entropy(self): + def entropy(self) -> TensorType: return torch.sum(self.multi_entropy(), dim=1) @override(ActionDistribution) - def multi_kl(self, other): + def multi_kl(self, other: ActionDistribution) -> TensorType: return torch.stack( [ torch.distributions.kl.kl_divergence(cat, oth_cat) @@ -131,12 +139,14 @@ class TorchMultiCategorical(TorchDistributionWrapper): ) @override(TorchDistributionWrapper) - def kl(self, other): + def kl(self, other: ActionDistribution) -> TensorType: return torch.sum(self.multi_kl(other), dim=1) @staticmethod @override(ActionDistribution) - def required_model_output_shape(action_space, model_config): + def required_model_output_shape( + action_space: gym.Space, + model_config: ModelConfigDict) -> Union[int, np.ndarray]: return np.sum(action_space.nvec) @@ -144,31 +154,33 @@ class TorchDiagGaussian(TorchDistributionWrapper): """Wrapper class for PyTorch Normal distribution.""" @override(ActionDistribution) - def __init__(self, inputs, model): + def __init__(self, inputs: List[TensorType], model: TorchModelV2): super().__init__(inputs, model) mean, log_std = torch.chunk(self.inputs, 2, dim=1) self.dist = torch.distributions.normal.Normal(mean, torch.exp(log_std)) @override(ActionDistribution) - def deterministic_sample(self): + def deterministic_sample(self) -> TensorType: self.last_sample = self.dist.mean return self.last_sample @override(TorchDistributionWrapper) - def logp(self, actions): + def logp(self, actions: TensorType) -> TensorType: return super().logp(actions).sum(-1) @override(TorchDistributionWrapper) - def entropy(self): + def entropy(self) -> TensorType: return super().entropy().sum(-1) @override(TorchDistributionWrapper) - def kl(self, other): + def kl(self, other: ActionDistribution) -> TensorType: return super().kl(other).sum(-1) @staticmethod @override(ActionDistribution) - def required_model_output_shape(action_space, model_config): + def required_model_output_shape( + action_space: gym.Space, + model_config: ModelConfigDict) -> Union[int, np.ndarray]: return np.prod(action_space.shape) * 2 @@ -179,7 +191,11 @@ class TorchSquashedGaussian(TorchDistributionWrapper): `low`+SMALL_NUMBER or `high`-SMALL_NUMBER respectively. """ - def __init__(self, inputs, model, low=-1.0, high=1.0): + def __init__(self, + inputs: List[TensorType], + model: TorchModelV2, + low: float = -1.0, + high: float = 1.0): """Parameterizes the distribution via `inputs`. Args: @@ -200,12 +216,12 @@ class TorchSquashedGaussian(TorchDistributionWrapper): self.high = high @override(ActionDistribution) - def deterministic_sample(self): + def deterministic_sample(self) -> TensorType: self.last_sample = self._squash(self.dist.mean) return self.last_sample @override(TorchDistributionWrapper) - def sample(self): + def sample(self) -> TensorType: # Use the reparameterization version of `dist.sample` to allow for # the results to be backprop'able e.g. in a loss term. normal_sample = self.dist.rsample() @@ -213,7 +229,7 @@ class TorchSquashedGaussian(TorchDistributionWrapper): return self.last_sample @override(ActionDistribution) - def logp(self, x): + def logp(self, x: TensorType) -> TensorType: # Unsquash values (from [low,high] to ]-inf,inf[) unsquashed_values = self._unsquash(x) # Get log prob of unsquashed values from our Normal. @@ -227,13 +243,13 @@ class TorchSquashedGaussian(TorchDistributionWrapper): torch.log(1 - unsquashed_values_tanhd**2 + SMALL_NUMBER), dim=-1) return log_prob - def _squash(self, raw_values): + def _squash(self, raw_values: TensorType) -> TensorType: # Returned values are within [low, high] (including `low` and `high`). squashed = ((torch.tanh(raw_values) + 1.0) / 2.0) * \ (self.high - self.low) + self.low return torch.clamp(squashed, self.low, self.high) - def _unsquash(self, values): + def _unsquash(self, values: TensorType) -> TensorType: normed_values = (values - self.low) / (self.high - self.low) * 2.0 - \ 1.0 # Stabilize input to atanh. @@ -244,7 +260,9 @@ class TorchSquashedGaussian(TorchDistributionWrapper): @staticmethod @override(ActionDistribution) - def required_model_output_shape(action_space, model_config): + def required_model_output_shape( + action_space: gym.Space, + model_config: ModelConfigDict) -> Union[int, np.ndarray]: return np.prod(action_space.shape) * 2 @@ -258,7 +276,11 @@ class TorchBeta(TorchDistributionWrapper): and Gamma(n) = (n - 1)! """ - def __init__(self, inputs, model, low=0.0, high=1.0): + def __init__(self, + inputs: List[TensorType], + model: TorchModelV2, + low: float = 0.0, + high: float = 1.0): super().__init__(inputs, model) # Stabilize input parameters (possibly coming from a linear layer). self.inputs = torch.clamp(self.inputs, log(SMALL_NUMBER), @@ -272,12 +294,12 @@ class TorchBeta(TorchDistributionWrapper): concentration1=alpha, concentration0=beta) @override(ActionDistribution) - def deterministic_sample(self): + def deterministic_sample(self) -> TensorType: self.last_sample = self._squash(self.dist.mean) return self.last_sample @override(TorchDistributionWrapper) - def sample(self): + def sample(self) -> TensorType: # Use the reparameterization version of `dist.sample` to allow for # the results to be backprop'able e.g. in a loss term. normal_sample = self.dist.rsample() @@ -285,19 +307,21 @@ class TorchBeta(TorchDistributionWrapper): return self.last_sample @override(ActionDistribution) - def logp(self, x): + def logp(self, x: TensorType) -> TensorType: unsquashed_values = self._unsquash(x) return torch.sum(self.dist.log_prob(unsquashed_values), dim=-1) - def _squash(self, raw_values): + def _squash(self, raw_values: TensorType) -> TensorType: return raw_values * (self.high - self.low) + self.low - def _unsquash(self, values): + def _unsquash(self, values: TensorType) -> TensorType: return (values - self.low) / (self.high - self.low) @staticmethod @override(ActionDistribution) - def required_model_output_shape(action_space, model_config): + def required_model_output_shape( + action_space: gym.Space, + model_config: ModelConfigDict) -> Union[int, np.ndarray]: return np.prod(action_space.shape) * 2 @@ -309,20 +333,22 @@ class TorchDeterministic(TorchDistributionWrapper): """ @override(ActionDistribution) - def deterministic_sample(self): + def deterministic_sample(self) -> TensorType: return self.inputs @override(TorchDistributionWrapper) - def sampled_action_logp(self): + def sampled_action_logp(self) -> TensorType: return torch.zeros((self.inputs.size()[0], ), dtype=torch.float32) @override(TorchDistributionWrapper) - def sample(self): + def sample(self) -> TensorType: return self.deterministic_sample() @staticmethod @override(ActionDistribution) - def required_model_output_shape(action_space, model_config): + def required_model_output_shape( + action_space: gym.Space, + model_config: ModelConfigDict) -> Union[int, np.ndarray]: return np.prod(action_space.shape) diff --git a/rllib/models/torch/visionnet.py b/rllib/models/torch/visionnet.py index a283f0378..1de30d722 100644 --- a/rllib/models/torch/visionnet.py +++ b/rllib/models/torch/visionnet.py @@ -1,4 +1,6 @@ import numpy as np +from typing import Dict, List +import gym from ray.rllib.models.torch.torch_modelv2 import TorchModelV2 from ray.rllib.models.torch.misc import normc_initializer, same_padding, \ @@ -6,6 +8,7 @@ from ray.rllib.models.torch.misc import normc_initializer, same_padding, \ from ray.rllib.models.utils import get_filter_config from ray.rllib.utils.annotations import override from ray.rllib.utils.framework import try_import_torch +from ray.rllib.utils.typing import ModelConfigDict, TensorType _, nn = try_import_torch() @@ -13,11 +16,11 @@ _, nn = try_import_torch() class VisionNetwork(TorchModelV2, nn.Module): """Generic vision network.""" - def __init__(self, obs_space, action_space, num_outputs, model_config, - name): + def __init__(self, obs_space: gym.spaces.Space, + action_space: gym.spaces.Space, num_outputs: int, + model_config: ModelConfigDict, name: str): if not model_config.get("conv_filters"): model_config["conv_filters"] = get_filter_config(obs_space.shape) - TorchModelV2.__init__(self, obs_space, action_space, num_outputs, model_config, name) nn.Module.__init__(self) @@ -148,7 +151,9 @@ class VisionNetwork(TorchModelV2, nn.Module): self._features = None @override(TorchModelV2) - def forward(self, input_dict, state, seq_lens): + def forward(self, input_dict: Dict[str, TensorType], + state: List[TensorType], + seq_lens: TensorType) -> (TensorType, List[TensorType]): self._features = input_dict["obs"].float().permute(0, 3, 1, 2) conv_out = self._convs(self._features) # Store features to save forward pass when getting value_function out. @@ -173,7 +178,7 @@ class VisionNetwork(TorchModelV2, nn.Module): return conv_out, state @override(TorchModelV2) - def value_function(self): + def value_function(self) -> TensorType: assert self._features is not None, "must call forward() first" if self._value_branch_separate: value = self._value_branch_separate(self._features) @@ -188,7 +193,7 @@ class VisionNetwork(TorchModelV2, nn.Module): features = self._features return self._value_branch(features).squeeze(1) - def _hidden_layers(self, obs): + def _hidden_layers(self, obs: TensorType) -> TensorType: res = self._convs(obs.permute(0, 3, 1, 2)) # switch to channel-major res = res.squeeze(3) res = res.squeeze(2)