class SimpleStochasticPolicy(Module): def __init__(self, **kwargs): super(SimpleStochasticPolicy, self).__init__() hidden_size = kwargs['linear_layers_size'] # actor self.bn = BatchNorm1d(kwargs['input_dim']) self.linears = ModuleList( [Linear(kwargs['input_dim'], hidden_size[0])]) self.linears.extend([ Linear(hidden_size[i - 1], hidden_size[i]) for i in range(1, len(hidden_size)) ]) self.mu = Linear(hidden_size[-1], kwargs['action_dim']) self.log_var = Linear(hidden_size[-1], kwargs['action_dim']) # self.log_var = torch.nn.Parameter(torch.zeros(kwargs['action_dim'])) self.relu = ReLU() self.tanh = Tanh() self.apply(init_weights) # xavier uniform init def forward(self, input, action=None): x = input x = self.bn(x) for l in self.linears: x = l(x) x = self.relu(x) mu = self.tanh(self.mu(x)) log_var = -self.relu(self.log_var(x)) sigmas = log_var.exp().sqrt() dists = Normal(mu, sigmas + 1.0e-4) if action is None: action = dists.rsample() log_prob = dists.log_prob(action).sum(dim=-1, keepdim=True) return action, log_prob, dists.entropy()
class SimpleDeterministicPolicy(Module): def __init__(self, **kwargs): super(SimpleDeterministicPolicy, self).__init__() hidden_size = kwargs['linear_layers_size'] # actor self.bn = BatchNorm1d(kwargs['input_dim']) self.linears = ModuleList( [Linear(kwargs['input_dim'], hidden_size[0])]) self.linears.extend([ Linear(hidden_size[i - 1], hidden_size[i]) for i in range(1, len(hidden_size)) ]) self.out = Linear(hidden_size[-1], kwargs['action_dim']) self.relu = ReLU() self.tanh = Tanh() self.apply(init_weights) # xavier uniform init def forward(self, input, action=None): x = input x = self.bn(x) for l in self.linears: x = l(x) x = self.relu(x) action = self.tanh(self.out(x)) return action
class DeterministicLstmModelDynamics(Module): def __init__(self, **kwargs): super(DeterministicLstmModelDynamics, self).__init__() self.__acoustic_state_dim = kwargs['goal_dim'] self.__action_dim = kwargs['action_dim'] self.__state_dim = kwargs['state_dim'] self.__lstm_sizes = kwargs['lstm_layers_size'] self.__linears_size = kwargs['linear_layers_size'] input_size = self.__acoustic_state_dim + self.__state_dim + self.__action_dim self.__bn1 = torch.nn.BatchNorm1d(input_size) self.lstms = ModuleList( [LSTM(input_size, self.__lstm_sizes[0], batch_first=True)]) self.lstms.extend([ LSTM(self.__lstm_sizes[i - 1], self.__lstm_sizes[i], batch_first=True) for i in range(1, len(self.__lstm_sizes)) ]) self.hiddens = [None] * len(self.__lstm_sizes) self.linears = ModuleList( [Linear(self.__lstm_sizes[-1], self.__linears_size[0])]) self.linears.extend([ Linear(self.__linears_size[i - 1], self.__linears_size[i]) for i in range(1, len(self.__linears_size)) ]) self.goal = Linear(self.__linears_size[-1], kwargs['goal_dim']) self.state = Linear(self.__linears_size[-1], kwargs['state_dim']) self.relu = ReLU() self.tanh = Tanh() self.apply(init_weights) # xavier uniform init def forward(self, states, actions, hidden=None): x = torch.cat((states, actions), -1) original_dim = x.shape x = self.__bn1(x.view(-1, original_dim[-1])) x = x.view(original_dim) for i, lstm in enumerate(self.lstms): x, self.hiddens[i] = lstm(x, self.hiddens[i]) for linear in self.linears: x = self.relu(linear(x)) # predict state states_out = self.state(x) + states[:, :self.__state_dim] # predict goal goals_out = self.goal(x) + states[:, self.__state_dim:] return states_out, goals_out def reset_hidden_state(self): self.hiddens = [None] * len(self.__lstm_sizes)
class SimpleStochasticModelDynamics(Module): def __init__(self, **kwargs): super(SimpleStochasticModelDynamics, self).__init__() self.__acoustic_state_dim = kwargs['goal_dim'] self.__action_dim = kwargs['action_dim'] self.__state_dim = kwargs['state_dim'] self.__linears_size = kwargs['linear_layers_size'] input_size = self.__acoustic_state_dim + self.__state_dim + self.__action_dim self.__bn1 = torch.nn.BatchNorm1d(input_size) self.linears = ModuleList([Linear(input_size, self.__linears_size[0])]) self.linears.extend([ Linear(self.__linears_size[i - 1], self.__linears_size[i]) for i in range(1, len(self.__linears_size)) ]) self.goal_mu = Linear(self.__linears_size[-1], kwargs['goal_dim']) self.goal_log_var = Linear(self.__linears_size[-1], kwargs['goal_dim']) self.state_mu = Linear(self.__linears_size[-1], kwargs['state_dim']) self.state_log_var = Linear(self.__linears_size[-1], kwargs['state_dim']) self.state_log_var_const = torch.zeros(kwargs['state_dim']) self.relu = ReLU() self.tanh = Tanh() self.apply(init_weights) # xavier uniform init def forward(self, states, actions): x = torch.cat((states, actions), -1) original_dim = x.shape x = self.__bn1(x.view(-1, original_dim[-1])) x = x.view(original_dim) for linear in self.linears: x = self.relu(linear(x)) # predict state state_mu = self.tanh(self.state_mu(x)) # state_log_var = -self.relu(self.state_log_var(x)) state_log_var = -5. - self.relu(self.state_log_var_const) state_sigmas = state_log_var.exp().sqrt() state_dists = Normal(state_mu, state_sigmas + 1.0e-4) states = state_dists.rsample() state_log_prob = state_dists.log_prob(states).sum(dim=-1, keepdim=True) # predict goal goal_mu = self.tanh(self.goal_mu(x)) goal_log_var = -self.relu(self.goal_log_var(x)) goal_sigmas = goal_log_var.exp().sqrt() goal_dists = Normal(goal_mu, goal_sigmas + 1.0e-4) goals = goal_dists.rsample() goal_log_prob = goal_dists.log_prob(goals).sum(dim=-1, keepdim=True) return states, goals, state_log_prob, goal_log_prob, state_dists, goal_dists
class SimpleDDPGAgent(Module): def __init__(self, **kwargs): super(SimpleDDPGAgent, self).__init__() hidden_size = kwargs['hidden_size'] # actor self.actor_linears = ModuleList( [Linear(kwargs['state_dim'], hidden_size[0])]) self.actor_linears.extend([ Linear(hidden_size[i - 1], hidden_size[i]) for i in range(1, len(hidden_size)) ]) self.action = Linear(hidden_size[-1], kwargs['action_dim']) # critic self.critic_linears = ModuleList([ Linear(kwargs['state_dim'] + kwargs['action_dim'], hidden_size[0]) ]) self.critic_linears.extend([ Linear(hidden_size[i - 1], hidden_size[i]) for i in range(1, len(hidden_size)) ]) self.q = Linear(hidden_size[-1], 1) self.relu = ReLU() self.sigmoid = Sigmoid() self.tanh = Tanh() self.apply(init_weights) # xavier uniform init def act(self, state): x = state for l in self.actor_linears: x = l(x) x = self.relu(x) action = self.tanh(self.action(x)) return action def Q(self, state, action): x = torch.cat([state, action], dim=1) for l in self.critic_linears: x = l(x) x = self.relu(x) q = self.q(x) return q def get_actor_parameters(self): return list(self.actor_linears.parameters()) + list( self.action.parameters()) def get_critic_parameters(self): return list(self.critic_linears.parameters()) + list( self.q.parameters())
class SimpleDeterministicModelDynamics(Module): def __init__(self, **kwargs): super(SimpleDeterministicModelDynamics, self).__init__() self.__acoustic_state_dim = kwargs['goal_dim'] self.__action_dim = kwargs['action_dim'] self.__state_dim = kwargs['state_dim'] self.__linears_size = kwargs['linear_layers_size'] input_size = self.__acoustic_state_dim + self.__state_dim + self.__action_dim self.__bn1 = torch.nn.BatchNorm1d(input_size) self.linears = ModuleList([Linear(input_size, self.__linears_size[0])]) self.linears.extend([ Linear(self.__linears_size[i - 1], self.__linears_size[i]) for i in range(1, len(self.__linears_size)) ]) self.goal = Linear(self.__linears_size[-1], kwargs['goal_dim']) self.state = Linear(self.__linears_size[-1], kwargs['state_dim']) self.relu = ReLU() self.tanh = Tanh() self.apply(init_weights) # xavier uniform init def forward(self, states, actions): x = torch.cat((states, actions), -1) original_dim = x.shape x = self.__bn1(x.view(-1, original_dim[-1])) x = x.view(original_dim) for linear in self.linears: x = self.relu(linear(x)) # predict state states = self.tanh(self.state(x)) # predict goal goals = self.tanh(self.goal(x)) return states, goals
class MLPFeatureExtractor(FeatureExtractor): def __init__(self, num_inputs: int, layers: Sequence[LinearSpec], dropout=0.0): super(MLPFeatureExtractor, self).__init__() self._num_inputs = num_inputs self._extractors = ModuleList() last_units = num_inputs for units, activation in layers: self._extractors.extend([ Linear(last_units, units), activations[activation] if isinstance(activation, str) else activation, Dropout(dropout) ]) last_units = units self._num_features = last_units def forward(self, X): for extractor in self._extractors: X = extractor(X) return X def features(self, X): with torch.no_grad(): features = [extractor(X) for extractor in self._extractors if isinstance(extractor, Linear)] return features @property def num_features(self): return self._num_features @property def input_shape(self): return (self._num_inputs, )
class Block(Module): """ TCSConv, Batch Normalisation, Activation, Dropout """ def __init__(self, in_channels, out_channels, activation, repeat=5, kernel_size=1, stride=1, dilation=1, dropout=0.0, residual=False, separable=False): super(Block, self).__init__() self.use_res = residual self.conv = ModuleList() _in_channels = in_channels padding = self.get_padding(kernel_size[0], stride[0], dilation[0]) # add the first n - 1 convolutions + activation for _ in range(repeat - 1): self.conv.extend( self.get_tcs(_in_channels, out_channels, kernel_size=kernel_size, stride=stride, dilation=dilation, padding=padding, separable=separable)) self.conv.extend(self.get_activation(activation, dropout)) _in_channels = out_channels # add the last conv and batch norm self.conv.extend( self.get_tcs(_in_channels, out_channels, kernel_size=kernel_size, stride=stride, dilation=dilation, padding=padding, separable=separable)) # add the residual connection if self.use_res: self.residual = Sequential( *self.get_tcs(in_channels, out_channels)) # add the activation and dropout self.activation = Sequential(*self.get_activation(activation, dropout)) def get_activation(self, activation, dropout): return activation, Dropout(p=dropout) def get_padding(self, kernel_size, stride, dilation): if stride > 1 and dilation > 1: raise ValueError( "Dilation and stride can not both be greater than 1") return (kernel_size // 2) * dilation def get_tcs(self, in_channels, out_channels, kernel_size=1, stride=1, dilation=1, padding=0, bias=False, separable=False): return [ TCSConv1d(in_channels, out_channels, kernel_size, stride=stride, dilation=dilation, padding=padding, bias=bias, separable=separable), BatchNorm1d(out_channels, eps=1e-3, momentum=0.1) ] def forward(self, x): _x = x for layer in self.conv: _x = layer(_x) if self.use_res: _x = _x + self.residual(x) return self.activation(_x)
class SimplePPOAgent(Module): def __init__(self, **kwargs): super(SimplePPOAgent, self).__init__() hidden_size = kwargs['hidden_size'] # actor self.actor_bn = BatchNorm1d(kwargs['state_dim']) self.actor_linears = ModuleList( [Linear(kwargs['state_dim'], hidden_size[0])]) self.actor_linears.extend([ Linear(hidden_size[i - 1], hidden_size[i]) for i in range(1, len(hidden_size)) ]) self.mu = Linear(hidden_size[-1], kwargs['action_dim']) self.log_var = Linear(hidden_size[-1], kwargs['action_dim']) # self.log_var = torch.nn.Parameter(torch.zeros(kwargs['action_dim'])) # critic self.critic_bn = BatchNorm1d(kwargs['state_dim']) self.critic_linears = ModuleList( [Linear(kwargs['state_dim'], hidden_size[0])]) self.critic_linears.extend([ Linear(hidden_size[i - 1], hidden_size[i]) for i in range(1, len(hidden_size)) ]) self.v = Linear(hidden_size[-1], 1) self.relu = ReLU() self.tanh = Tanh() self.apply(init_weights) # xavier uniform init def forward(self, state, action=None): x = state x = self.actor_bn(x) for l in self.actor_linears: x = l(x) x = self.relu(x) mu = self.tanh(self.mu(x)) log_var = -self.relu(self.log_var(x)) sigmas = log_var.exp().sqrt() dists = Normal(mu, sigmas) if action is None: action = dists.sample() log_prob = dists.log_prob(action).sum(dim=-1, keepdim=True) x = state x = self.critic_bn(x) for l in self.critic_linears: x = l(x) x = self.relu(x) v = self.v(x) return action, log_prob, dists.entropy(), v def get_actor_parameters(self): return [ *self.actor_bn.parameters(), *self.actor_linears.parameters(), *self.mu.parameters(), *self.log_var.parameters() ] def get_critic_parameters(self): return [ *self.critic_bn.parameters(), *self.critic_linears.parameters(), *self.v.parameters() ]
class BlockDP(Module): def __init__(self, in_channels, out_channels, activation, kernel_size, norm=None, prediction_size=32, dropout=0.05): super(BlockDP, self).__init__() self.conv = ModuleList() self.conv.extend( self.get_tcs(in_channels, out_channels, kernel_size=kernel_size, padding=kernel_size[0] // 2, separable=False)) self.predictor = torch.nn.Sequential( torch.nn.Conv1d(2, prediction_size, 31, stride=1, padding=15), torch.nn.BatchNorm1d(prediction_size), torch.nn.SiLU(), torch.nn.Conv1d(prediction_size, prediction_size, 15, stride=1, padding=7), torch.nn.BatchNorm1d(prediction_size), torch.nn.SiLU(), torch.nn.Conv1d(prediction_size, 2, 15, stride=1, padding=7)) self.predictor[-1].weight.data *= 0.01 self.predictor[-1].bias.data *= 0 self.predictor[-1].bias.data[0] -= np.log(2) self.predictor[-1].bias.data[1] -= np.log(2) self.activation = Sequential(*self.get_activation(activation, dropout)) self.norm_target = norm self.register_buffer('norm_mean', torch.ones(1)) def get_activation(self, activation, dropout): return nn.GLU(dim=1), Dropout(p=dropout) def row_pool(self, features, moves, weights): fw = features * weights.to(features.dtype).unsqueeze(1) poses = torch.cumsum(moves.detach(), 0) poses = poses.unsqueeze(1) floors = torch.floor(poses) ceils = floors + 1 w1 = (1 - (poses - floors)).to(features.dtype) w2 = (1 - (ceils - poses)).to(features.dtype) out = torch.zeros((int(ceils[-1].item()) + 1, features.shape[1]), device=features.device, dtype=features.dtype) out.index_add_(0, floors.to(torch.long).squeeze(1), w1 * fw) out.index_add_(0, ceils.to(torch.long).squeeze(1), w2 * fw) return out def get_tcs(self, in_channels, out_channels, kernel_size=1, stride=1, dilation=1, padding=0, bias=False, separable=False): return [ TCSConv1d(in_channels, out_channels * 2, kernel_size, stride=stride, dilation=dilation, padding=padding, bias=bias, separable=separable), BatchNorm1d(out_channels * 2, eps=1e-3, momentum=0.1) ] def forward(self, x): _x = x for layer in self.conv: _x = layer(_x) _x = self.activation(_x) features = _x jumps_mat = self.predictor(torch.cat([x, x * x], dim=1)) weights = torch.sigmoid(jumps_mat[:, 0, :]) moves = torch.sigmoid(jumps_mat[:, 1, :]) bmoves = moves if self.training: renorm = (1 / self.norm_target / moves.mean().detach()).detach() self.norm_mean.copy_(0.99 * self.norm_mean + 0.01 * renorm) else: renorm = self.norm_mean moves = moves * renorm features = features.permute((0, 2, 1)) lens = [] x_evs = [] for f, m, w in zip(features.unbind(0), moves.unbind(0), weights.unbind(0)): pooled = self.row_pool(f, m, w) x_evs.append(pooled) lens.append(pooled.shape[0]) x_evs = torch.nn.utils.rnn.pad_sequence(x_evs, True) x_evs = x_evs.permute(0, 2, 1) x_evs = F.pad(x_evs, (0, 3 - (x_evs.shape[2] % 3))) #x_evs = self.activation(x_evs) return x_evs, lens, bmoves, weights
class Merger(Component): def protocol_average(self, x) -> torch.Tensor: for m in self.merged_modules: if self.output is None: self.output = m(x) else: self.output.add_(m(x)) self.output.div_(len(self.merged_modules)) return self.output def protocol_voting(self, x) -> torch.Tensor: warn( "Some tensor operations in the voting merging do not support CUDA") n_classes = len(self.merged_modules[0](x[0])) tensors = tuple( m(x).argmax(dim=1)[:, None] for m in self.merged_modules) # Assuming m(x) predictions 2-D tensor class_predictions = torch.cat(tensors, dim=1) voting = class_predictions.mode().values self.output = one_hot(voting, num_classes=n_classes).float( ) # Returns a probability distribution return self.output def protocol_max(self, x) -> torch.Tensor: warn("Max protocol is uneficient, and does not support CUDA") tuple_pred = tuple(m(x) for m in self.merged_modules) stack_red = torch.stack(tuple_pred, dim=2) max_prob_pred = torch.max(stack_red, dim=1).values valid_pd = max_prob_pred.argmax(dim=1) # Fill the output with the valid distribution self.output = torch.empty_like(tuple_pred[0]) for i in torch.unique(valid_pd): idx = torch.where(valid_pd == i)[0] self.output[idx, :] = tuple_pred[i][idx, :] return self.output def __init__(self, merged_modules: List[Component], protocol=MergeProtocol.AVERAGE): if len(merged_modules) == 0: raise ValueError("Number of merged components has to be > 0") super().__init__(p=1, t=0) self.protocol = protocol self.merged_modules = ModuleList(merged_modules) self.register_buffer("output", None) def forward(self, x: torch.Tensor) -> torch.Tensor: if self.protocol == MergeProtocol.AVERAGE: return self.protocol_average(x) elif self.protocol == MergeProtocol.VOTING: return self.protocol_voting(x) elif self.protocol == MergeProtocol.MAX: return self.protocol_max(x) else: raise ValueError("Merging protocol not supported") def get_merge_protocol(self): return self.protocol def get_merged_modules(self): return self.merged_modules def update_merge_protocol(self, p: MergeProtocol): self.protocol = p def add_classifier(self, c: Component): self.merged_modules.extend([c])
class SimpleDeterministicModelDynamicsDeltaPredict(Module): def __init__(self, **kwargs): super(SimpleDeterministicModelDynamicsDeltaPredict, self).__init__() self.__acoustic_state_dim = kwargs['goal_dim'] self.__action_dim = kwargs['action_dim'] self.__state_dim = kwargs['state_dim'] self.__acoustic_dim = 26 self.__linears_size = kwargs['linear_layers_size'] # input_size = self.__acoustic_state_dim + self.__state_dim + self.__action_dim input_size = self.__state_dim + self.__action_dim self.__bn1 = torch.nn.BatchNorm1d(input_size) self.drop = torch.nn.modules.Dropout(p=0.1) # self.artic_state_0 = Linear(self.__state_dim + self.__action_dim - self.__acoustic_dim, 64) # self.artic_state_1 = Linear(64, self.__state_dim - self.__acoustic_dim) self.artic_state_0 = Linear( self.__state_dim + self.__action_dim - self.__acoustic_dim, self.__state_dim - self.__acoustic_dim) # self.artic_state_1 = Linear(64, ) self.linears = ModuleList([Linear(input_size, self.__linears_size[0])]) self.linears.extend([ Linear(self.__linears_size[i - 1], self.__linears_size[i]) for i in range(1, len(self.__linears_size)) ]) self.goal = Linear(self.__linears_size[-1], kwargs['goal_dim']) self.acoustic_state = Linear(self.__linears_size[-1], self.__acoustic_dim) self.state = Linear(self.__state_dim, self.__state_dim) self.relu = ReLU() self.tanh = Tanh() self.apply(init_weights) # xavier uniform init def forward(self, states, actions): x = torch.cat((states[:, :self.__state_dim], actions), -1) original_dim = x.shape x = self.__bn1(x.view(-1, original_dim[-1])) x = x.view(original_dim) x = self.drop(x) # artic artic_x = x[:, :self.__state_dim - self.__acoustic_dim] actions_x = x[:, -self.__action_dim:] # artic_state_delta = self.artic_state_1(self.relu(self.artic_state_0(torch.cat((artic_x, actions_x), -1)))) artic_state_delta = self.artic_state_0( torch.cat((artic_x, actions_x), -1)) # acoustic for linear in self.linears: x = self.relu(linear(x)) # predict state acoustic_state_delta = self.acoustic_state(x) states_delta = torch.cat((artic_state_delta, acoustic_state_delta), -1) # states_delta = self.tanh(torch.cat((artic_state_delta, acoustic_state_delta), -1)) out_states = self.tanh(states[:, :self.__state_dim] + states_delta) # predict goal goals_delta = self.tanh(self.goal(x)) out_goals = self.tanh(states[:, self.__state_dim:] + goals_delta) return out_states, out_goals
class Conv2dFeatureExtractor(FeatureExtractor): def __init__(self, num_channels: int, width: int, height: int, layers: Sequence[Conv2dSpec], dropout=0.0): super(Conv2dFeatureExtractor, self).__init__() self._num_channels = num_channels self._width = width self._height = height self._extractors = ModuleList() last_kernels = num_channels for num_kernels, kernel_size, stride, activation in layers: self._extractors.extend([ Conv2d(last_kernels, num_kernels, kernel_size, stride), activations[activation] if isinstance(activation, str) else activation, Dropout2d(dropout) ]) last_kernels = num_kernels width, height = compute_output((width, height), kernel_size, stride) self._num_features = last_kernels * width * height def forward(self, X): for extractor in self._extractors: X = extractor(X) F = flatten(X) return F @property def num_features(self): return self._num_features @property def input_shape(self): return self._num_channels, self._width, self._height def feature_maps(self, X): feature_maps = [] with torch.no_grad(): for extractor in self._extractors: X = extractor(X) if isinstance(extractor, Conv2d): w, h = X.shape[-2], X.shape[-1] fmaps = X.view(-1, w, h).cpu() [feature_maps.append(fmap) for fmap in fmaps] return feature_maps def plot_feature_maps(self, X): for feature_map in self.feature_maps(X): plt.imshow(feature_map) plt.show() @property def filters(self): filters = [] for extractor in self._extractors: if isinstance(extractor, Conv2d): weights = extractor.weight.cpu().detach() w, h = weights.shape[-2], weights.shape[-1] weights = weights.view(-1, w, h).cpu() [filters.append(weight) for weight in weights] return filters def plot_filters(self): for filter in self.filters: plt.imshow(filter) plt.show()