def __init__(self, action_size, layer_fn, activation_fn=nn.ReLU, bias=True, parity="odd"): """ Conditional affine coupling layer used in Real NVP Bijector. Original paper: https://arxiv.org/abs/1605.08803 Adaptation to RL: https://arxiv.org/abs/1804.02808 Important notes --------------- 1. State embeddings are supposed to have size (action_size * 2). 2. Scale and translation networks used in the Real NVP Bijector both have one hidden layer of (action_size) (activation_fn) units. 3. Parity ("odd" or "even") determines which part of the input is being copied and which is being transformed. """ super().__init__() layer_fn = MODULES.get_if_str(layer_fn) activation_fn = MODULES.get_if_str(activation_fn) self.parity = parity if self.parity == "odd": self.copy_size = action_size // 2 else: self.copy_size = action_size - action_size // 2 self.scale_prenet = SequentialNet( hiddens=[action_size * 2 + self.copy_size, action_size], layer_fn=layer_fn, activation_fn=activation_fn, norm_fn=None, bias=bias) self.scale_net = SequentialNet( hiddens=[action_size, action_size - self.copy_size], layer_fn=layer_fn, activation_fn=None, norm_fn=None, bias=True) self.translation_prenet = SequentialNet( hiddens=[action_size * 2 + self.copy_size, action_size], layer_fn=layer_fn, activation_fn=activation_fn, norm_fn=None, bias=bias) self.translation_net = SequentialNet( hiddens=[action_size, action_size - self.copy_size], layer_fn=layer_fn, activation_fn=None, norm_fn=None, bias=True) inner_init = create_optimal_inner_init(nonlinearity=activation_fn) self.scale_prenet.apply(inner_init) self.scale_net.apply(outer_init) self.translation_prenet.apply(inner_init) self.translation_net.apply(outer_init)
def __init__(self, hiddens, layer_fn=nn.Linear, bias=True, norm_fn=None, activation_fn=nn.ReLU, dropout=None, layer_order=None, residual=False): super().__init__() assert len(hiddens) > 1, "No sequence found" layer_fn = MODULES.get_if_str(layer_fn) activation_fn = MODULES.get_if_str(activation_fn) norm_fn = MODULES.get_if_str(norm_fn) dropout = MODULES.get_if_str(dropout) inner_init = create_optimal_inner_init(nonlinearity=activation_fn) layer_order = layer_order or ["layer", "norm", "drop", "act"] if isinstance(dropout, float): dropout_fn = lambda: nn.Dropout(dropout) else: dropout_fn = dropout def _layer_fn(f_in, f_out, bias): return layer_fn(f_in, f_out, bias=bias) def _normalize_fn(f_in, f_out, bias): return norm_fn(f_out) if norm_fn is not None else None def _dropout_fn(f_in, f_out, bias): return dropout_fn() if dropout_fn is not None else None def _activation_fn(f_in, f_out, bias): return activation_fn() if activation_fn is not None else None name2fn = { "layer": _layer_fn, "norm": _normalize_fn, "drop": _dropout_fn, "act": _activation_fn, } net = [] for i, (f_in, f_out) in enumerate(pairwise(hiddens)): block = [] for key in layer_order: fn = name2fn[key](f_in, f_out, bias) if fn is not None: block.append((f"{key}", fn)) block = torch.nn.Sequential(OrderedDict(block)) if residual: block = ResidualWrapper(net=block) net.append((f"block_{i}", block)) self.net = torch.nn.Sequential(OrderedDict(net)) self.net.apply(inner_init)
def _normalization_fn(normalization_fn, f_in, f_out, **kwargs): normalization_fn = MODULES.get_if_str(normalization_fn) normalization_fn = \ normalization_fn(f_out, **kwargs) \ if normalization_fn is not None \ else None return normalization_fn
def __init__(self, squashing_fn=nn.Tanh): """ Layer that squashes samples from some distribution to be bounded. """ super().__init__() self.squashing_fn = MODULES.get_if_str(squashing_fn)()
def _get_block(**conv_params): layers = [nn.Conv2d(**conv_params)] if normalization is not None: normalization_fn = MODULES.get_if_str(normalization) layers.append(normalization_fn(conv_params["out_channels"])) if dropout_rate is not None: layers.append(nn.Dropout2d(p=dropout_rate)) layers.append(activation_fn(inplace=True)) return layers
def __init__(self, in_features, activation_fn="Sigmoid"): super().__init__() activation_fn = MODULES.get_if_str(activation_fn) self.attn = nn.Sequential( nn.Conv2d( in_features, 1, kernel_size=1, stride=1, padding=0, bias=False ), activation_fn() )
def __init__(self, action_size, layer_fn, activation_fn=nn.ReLU, squashing_fn=nn.Tanh, bias=False): super().__init__() activation_fn = MODULES.get_if_str(activation_fn) self.action_size = action_size self.coupling1 = CouplingLayer(action_size=action_size, layer_fn=layer_fn, activation_fn=activation_fn, bias=bias, parity="odd") self.coupling2 = CouplingLayer(action_size=action_size, layer_fn=layer_fn, activation_fn=activation_fn, bias=bias, parity="even") self.squashing_layer = SquashingLayer(squashing_fn)
def _activation_fn(activation_fn, f_in, f_out, **kwargs): activation_fn = MODULES.get_if_str(activation_fn) activation_fn = activation_fn(**kwargs) \ if activation_fn is not None \ else None return activation_fn
def _dropout_fn(dropout_fn, f_in, f_out, **kwargs): dropout_fn = MODULES.get_if_str(dropout_fn) dropout_fn = dropout_fn(**kwargs) \ if dropout_fn is not None \ else None return dropout_fn
def _layer_fn(layer_fn, f_in, f_out, **kwargs): layer_fn = MODULES.get_if_str(layer_fn) layer_fn = layer_fn(f_in, f_out, **kwargs) return layer_fn