wassname/ray
1
0
Fork 0
mirror of https://github.com/wassname/ray.git synced 2026-08-03 13:10:57 +08:00
Code Issues Packages Projects Releases Wiki Activity

[RLlib] Model Annotations to Torch Models (#9749)

Browse Source
This commit is contained in:
Michael Luo
2020-11-12 12:16:12 +01:00
committed by GitHub
parent 3fbd8be851
commit b2984d1c34
13 changed files with 240 additions and 137 deletions
Download Patch File Download Diff File Expand all files Collapse all files
rllib/models/catalog.py
+3 -1
View File
@@ -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)
rllib/models/preprocessors.py
+29 -22
View File
@@ -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))
rllib/models/torch/attention_net.py
+19 -16
View File
@@ -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])
rllib/models/torch/fcnet.py
+9 -4
View File
@@ -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(
rllib/models/torch/misc.py
+50 -21
View File
@@ -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
rllib/models/torch/modules/gru_gate.py
+3 -2
View File
@@ -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
rllib/models/torch/modules/multi_head_attention.py
+4 -2
View File
@@ -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
rllib/models/torch/modules/noisy_layer.py
+12 -7
View File
@@ -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)
rllib/models/torch/modules/relative_multi_head_attention.py
+11 -9
View File
@@ -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
rllib/models/torch/modules/skip_connection.py
+8 -2
View File
@@ -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.
rllib/models/torch/recurrent_net.py
+18 -8
View File
@@ -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])
rllib/models/torch/torch_action_dist.py
+63 -37
View File
@@ -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)
rllib/models/torch/visionnet.py
+11 -6
View File
@@ -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)