基于Implicit Quantile Network和Normalization Flow的深度强化学习算法.zip

''' copied and modified from https://github.com/CW-Huang/torchkit/blob/33f61b914bf8e79faebab3d3d64c17ea921ce6d2/torchkit/nn.py copied and modified from https://github.com/lim0606/pytorch-flows-dev/blob/master/flows.py ''' import math import torch import torch.nn as nn import torch.nn.functional as F from utils import get_nonlinear_func from torch.nn.modules.utils import _pair ''' miscellanious layers ''' class Identity(nn.Module): def __init__(self,): super().__init__() def forward(self, input): return input ''' copied and modified from https://github.com/CW-Huang/torchkit/blob/master/nn.py ''' class WeightNormalizedLinear(nn.Module): def __init__(self, in_features, out_features, bias=True, mask=None, norm=True): super().__init__() self.in_features = in_features self.out_features = out_features self.register_buffer('mask',mask) self.norm = norm self.direction = nn.Parameter(torch.Tensor(out_features, in_features)) self.scale = nn.Parameter(torch.Tensor(out_features)) if bias: self.bias = nn.Parameter(torch.Tensor(out_features)) else: self.register_parameter('bias', None) self.reset_parameters() def reset_parameters(self): stdv = 1. / math.sqrt(self.direction.size(1)) self.direction.data.uniform_(-stdv, stdv) self.scale.data.uniform_(1, 1) if self.bias is not None: self.bias.data.uniform_(-stdv, stdv) def forward(self, input): if self.norm: dir_ = self.direction direction = dir_.div(dir_.pow(2).sum(1).sqrt()[:,None]) weight = self.scale[:,None].mul(direction) else: weight = self.scale[:,None].mul(self.direction) if self.mask is not None: #weight = weight * getattr(self.mask,⋅ # ('cpu', 'cuda')[weight.is_cuda])() weight = weight * self.mask return F.linear(input, weight, self.bias) def __repr__(self): return self.__class__.__name__ + '(' \ + 'in_features=' + str(self.in_features) \ + ', out_features=' + str(self.out_features) + ')' class ResLinear(nn.Module): def __init__(self, in_features, out_features, bias=True, same_dim=False, activation=nn.ReLU(), oper=WeightNormalizedLinear, oper_kwargs={'norm': False}): super().__init__() self.same_dim = same_dim self.dot_0h = oper(in_features, out_features, bias, **oper_kwargs) self.dot_h1 = oper(out_features, out_features, bias, **oper_kwargs) if not same_dim: self.dot_01 = oper(in_features, out_features, bias, **oper_kwargs) self.activation = activation def forward(self, input): h = self.activation(self.dot_0h(input)) out_nonlinear = self.dot_h1(h) out_skip = input if self.same_dim else self.dot_01(input) return out_nonlinear + out_skip class ContextResLinear(nn.Module): def __init__(self, in_features, out_features, context_features, bias=True, same_dim=False, activation=nn.ReLU(), oper=WeightNormalizedLinear, oper_kwargs={'norm': False}): super().__init__() self.same_dim = same_dim self.dot_0h = oper(in_features, out_features, bias, **oper_kwargs) self.dot_h1 = oper(out_features, out_features, bias, **oper_kwargs) if not same_dim: self.dot_01 = oper(in_features, out_features, bias, **oper_kwargs) self.dot_0c = oper(context_features, out_features, bias, **oper_kwargs) self.dot_c1 = oper(out_features, out_features, bias, **oper_kwargs) self.activation = activation def forward(self, input, context): h = self.activation(self.dot_0h(input)) outi_nonlinear = self.dot_h1(h) c = self.activation(self.dot_0c(context)) outc_nonlinear = self.dot_c1(c) out_skip = input if self.same_dim else self.dot_01(input) return outi_nonlinear + outc_nonlinear + out_skip ''' context ''' class ContextLinear(nn.Module): __constants__ = ['bias', 'in_features', 'out_features'] def __init__(self, in_features, out_features, context_features, bias=True): super().__init__() self.in_features = in_features self.out_features = out_features self.context_features = context_features self.direction = nn.Parameter(torch.Tensor(out_features, in_features)) self.cscale = nn.Linear(context_features, out_features, bias=False) self.cbias = nn.Linear(context_features, out_features, bias=bias) #self.cbias = nn.Linear(in_features+context_features, out_features, bias=bias) self.reset_parameters() def reset_parameters(self): torch.nn.init.kaiming_uniform_(self.direction, a=math.sqrt(5)) self.cscale.weight.data.normal_(0, 0.005) #torch.nn.init.constant_(self.cscale.bias, 1) #self.cbias.weight.data.normal_(0, 0.001) #torch.nn.init.constant_(self.cbias.bias, 0) def forward(self, input, context): scale = 1.+self.cscale(context) bias = self.cbias(context) return scale * F.linear(input, self.direction, None) + bias #return scale * self.cbias(torch.cat([input, context], dim=1)) def extra_repr(self): return 'in_features={}, out_features={}, context_features={}'.format( self.in_features, self.out_features, self.context_features, ) class ContextConv2d(nn.Module): def __init__(self, in_channels, out_channels, context_channels, kernel_size, stride=1, padding=0, dilation=1, groups=1, bias=True, padding_mode='zeros'): super().__init__() self.in_channels = in_channels self.out_channels = out_channels self.context_channels = context_channels self.direction = nn.Conv2d(in_channels, out_channels, bias=False, kernel_size=kernel_size, stride=stride, padding=padding, dilation=dilation, groups=groups)#, padding_mode=padding_mode) self.cscale = nn.Conv2d(context_channels, out_channels, bias=False, kernel_size=kernel_size, stride=stride, padding=padding, dilation=dilation, groups=groups)#, padding_mode=padding_mode) self.cbias = nn.Conv2d(context_channels, out_channels, bias=bias, kernel_size=kernel_size, stride=stride, padding=padding, dilation=dilation, groups=groups)#, padding_mode=padding_mode) self.reset_parameters() def reset_parameters(self): #torch.nn.init.kaiming_uniform_(self.direction, a=math.sqrt(5)) self.cscale.weight.data.normal_(0, 0.005) #torch.nn.init.constant_(self.cscale.bias, 1) #self.cbias.weight.data.normal_(0, 0.001) #torch.nn.init.constant_(self.cbias.bias, 0) def forward(self, input, context): scale = 1.+self.cscale(context) bias = self.cbias(context) return scale * self.direction(input) + bias def extra_repr(self): return 'in_channels={}, out_channels={}, context_channels={}'.format( self.in_channels, self.out_channels, self.context_channels, ) class ContextWeightNormalizedLinear(nn.Module): def __init__(self, in_features, out_features, context_features, bias=True, in_norm=False, ctx_norm=True, ctx_scale=0.1): super().__init__() self.in_features = in_features self.out_features = out_features self.context_features = context_features self.in_norm = in_norm self.ctx_norm = ctx_norm self.ctx_scale = ctx_scale self.direction = nn.Parameter(torch.Tensor(out_features, in_features)) self.cscale = nn.Parameter(torch.Tensor(out_features, context_features)) self.cbias = nn.Linear(context_features, out_features, bias=bias) self.reset_parameters() def reset_parameters(self): torch.nn.init.kaiming_un