def __init__(self, cfg, obs_space, timing, self_obs_dim=18, neighbor_obs_dim=6, neighbor_hidden_size=32): super().__init__(cfg, timing) self.self_obs_dim = self_obs_dim self.neighbor_obs_dim = neighbor_obs_dim self.neighbor_hidden_size = neighbor_hidden_size fc_encoder_layer = cfg.hidden_size # encode the current drone's observations self.self_encoder = nn.Sequential( nn.Linear(self.self_obs_dim, fc_encoder_layer), nonlinearity(cfg), nn.Linear(fc_encoder_layer, fc_encoder_layer), nonlinearity(cfg)) # encode the neighboring drone's observations self.neighbor_encoder = nn.Sequential( nn.Linear(self.neighbor_obs_dim, self.neighbor_hidden_size), nonlinearity(cfg), ) self.self_encoder_out_size = calc_num_elements(self.self_encoder, (self.self_obs_dim, )) self.neighbor_encoder_out_size = calc_num_elements( self.neighbor_encoder, (self.neighbor_obs_dim, )) # Feed forward self obs and neighbor obs after concatenation self.feed_forward = nn.Linear( self.self_encoder_out_size + self.neighbor_encoder_out_size, cfg.hidden_size) self.init_fc_blocks(cfg.hidden_size)
def __init__(self, cfg, obs_space, timing): super().__init__(cfg, timing) obs_shape = get_obs_shape(obs_space) input_ch = obs_shape.obs[0] log.debug('Num input channels: %d', input_ch) if cfg.encoder_subtype == 'convnet_simple': conv_filters = [[input_ch, 32, 8, 4], [32, 64, 4, 2], [64, 128, 3, 2]] elif cfg.encoder_subtype == 'convnet_impala': conv_filters = [[input_ch, 16, 8, 4], [16, 32, 4, 2]] elif cfg.encoder_subtype == 'minigrid_convnet_tiny': conv_filters = [[3, 16, 3, 1], [16, 32, 2, 1], [32, 64, 2, 1]] else: raise NotImplementedError(f'Unknown encoder {cfg.encoder_subtype}') activation = nonlinearity(self.cfg) fc_layer_size = fc_after_encoder_size(self.cfg) encoder_extra_fc_layers = self.cfg.encoder_extra_fc_layers enc = self.ConvEncoderImpl(activation, conv_filters, fc_layer_size, encoder_extra_fc_layers, obs_shape) self.enc = torch.jit.script(enc) self.encoder_out_size = calc_num_elements(self.enc, obs_shape.obs) log.debug('Encoder output size: %r', self.encoder_out_size)
def __init__(self, activation, conv_filters, fc_layer_size, encoder_extra_fc_layers, obs_shape): super(ConvEncoder.ConvEncoderImpl, self).__init__() conv_layers = [] for layer in conv_filters: if layer == 'maxpool_2x2': conv_layers.append(nn.MaxPool2d((2, 2))) elif isinstance(layer, (list, tuple)): inp_ch, out_ch, filter_size, stride = layer conv_layers.append( nn.Conv2d(inp_ch, out_ch, filter_size, stride=stride)) conv_layers.append(activation) else: raise NotImplementedError(f'Layer {layer} not supported!') self.conv_head = nn.Sequential(*conv_layers) self.conv_head_out_size = calc_num_elements( self.conv_head, obs_shape.obs) fc_layers = [] for i in range(encoder_extra_fc_layers): size = self.conv_head_out_size if i == 0 else fc_layer_size fc_layers.extend([nn.Linear(size, fc_layer_size), activation]) self.fc_layers = nn.Sequential(*fc_layers)
def __init__(self, cfg, obs_space, timing): super().__init__(cfg, timing) obs_shape = get_obs_shape(obs_space) conv_layers = [ nn.Conv2d(1, 8, 3, stride=2), nonlinearity(cfg), nn.Conv2d(8, 16, 2, stride=1), nonlinearity(cfg), ] self.conv_head = nn.Sequential(*conv_layers) self.conv_head_out_size = calc_num_elements(self.conv_head, obs_shape.obs) self.init_fc_blocks(self.conv_head_out_size)
def __init__(self, cfg, obs_space, timing): super().__init__(cfg, timing) self.basic_encoder = create_standard_encoder(cfg, obs_space, timing) self.encoder_out_size = self.basic_encoder.encoder_out_size obs_shape = get_obs_shape(obs_space) self.measurements_head = None if 'measurements' in obs_shape: self.measurements_head = nn.Sequential( nn.Linear(obs_shape.measurements[0], 128), nonlinearity(cfg), nn.Linear(128, 128), nonlinearity(cfg), ) measurements_out_size = calc_num_elements(self.measurements_head, obs_shape.measurements) self.encoder_out_size += measurements_out_size log.debug('Policy head output size: %r', self.encoder_out_size)
def __init__(self, cfg, obs_space, timing): super().__init__(cfg, timing) obs_shape = get_obs_shape(obs_space) input_ch = obs_shape.obs[0] log.debug('Num input channels: %d', input_ch) if cfg.encoder_subtype == 'resnet_impala': # configuration from the IMPALA paper resnet_conf = [[16, 2], [32, 2], [32, 2]] else: raise NotImplementedError( f'Unknown resnet subtype {cfg.encoder_subtype}') curr_input_channels = input_ch layers = [] for i, (out_channels, res_blocks) in enumerate(resnet_conf): layers.extend([ nn.Conv2d(curr_input_channels, out_channels, kernel_size=3, stride=1, padding=1), # padding SAME nn.MaxPool2d(kernel_size=3, stride=2, padding=1), # padding SAME ]) for j in range(res_blocks): layers.append( ResBlock(cfg, out_channels, out_channels, self.timing)) curr_input_channels = out_channels layers.append(nonlinearity(cfg)) self.conv_head = nn.Sequential(*layers) self.conv_head_out_size = calc_num_elements(self.conv_head, obs_shape.obs) log.debug('Convolutional layer output size: %r', self.conv_head_out_size) self.init_fc_blocks(self.conv_head_out_size)
def __init__(self, cfg, obs_space, timing): super().__init__(cfg, timing) obs_shape = get_obs_shape(obs_space) input_ch = obs_shape.obs[0] log.debug('Num input channels: %d', input_ch) if cfg.encoder_subtype == 'convnet_simple': conv_filters = [[input_ch, 32, 8, 4], [32, 64, 4, 2], [64, 128, 3, 2]] elif cfg.encoder_subtype == 'convnet_impala': conv_filters = [[input_ch, 16, 8, 4], [16, 32, 4, 2]] elif cfg.encoder_subtype == 'minigrid_convnet_tiny': conv_filters = [[3, 16, 3, 1], [16, 32, 2, 1], [32, 64, 2, 1]] else: raise NotImplementedError(f'Unknown encoder {cfg.encoder_subtype}') conv_layers = [] for layer in conv_filters: if layer == 'maxpool_2x2': conv_layers.append(nn.MaxPool2d((2, 2))) elif isinstance(layer, (list, tuple)): inp_ch, out_ch, filter_size, stride = layer conv_layers.append( nn.Conv2d(inp_ch, out_ch, filter_size, stride=stride)) conv_layers.append(nonlinearity(cfg)) else: raise NotImplementedError(f'Layer {layer} not supported!') self.conv_head = nn.Sequential(*conv_layers) self.conv_head_out_size = calc_num_elements(self.conv_head, obs_shape.obs) log.debug('Convolutional layer output size: %r', self.conv_head_out_size) self.init_fc_blocks(self.conv_head_out_size)
def __init__(self, cfg, obs_space, timing): super().__init__(cfg, timing) # internal params -- cannot change from cmd line if cfg.quads_obs_repr == 'xyz_vxyz_R_omega': self.self_obs_dim = 18 elif cfg.quads_obs_repr == 'xyz_vxyz_R_omega_wall': self.self_obs_dim = 24 else: raise NotImplementedError( f'Layer {cfg.quads_obs_repr} not supported!') self.neighbor_hidden_size = cfg.quads_neighbor_hidden_size self.neighbor_obs_type = cfg.neighbor_obs_type self.use_spectral_norm = cfg.use_spectral_norm self.obstacle_mode = cfg.quads_obstacle_mode if cfg.quads_local_obs == -1: self.num_use_neighbor_obs = cfg.quads_num_agents - 1 else: self.num_use_neighbor_obs = cfg.quads_local_obs if self.neighbor_obs_type == 'pos_vel_goals': self.neighbor_obs_dim = 9 # include goal pos info elif self.neighbor_obs_type == 'pos_vel': self.neighbor_obs_dim = 6 elif self.neighbor_obs_type == 'pos_vel_goals_ndist_gdist': self.neighbor_obs_dim = 11 elif self.neighbor_obs_type == 'none': # override these params so that neighbor encoder is a no-op during inference self.neighbor_obs_dim = 0 self.num_use_neighbor_obs = 0 else: raise NotImplementedError( f'Unknown value {cfg.neighbor_obs_type} passed to --neighbor_obs_type' ) # encode the neighboring drone's observations neighbor_encoder_out_size = 0 self.neighbor_encoder = None if self.num_use_neighbor_obs > 0: neighbor_encoder_type = cfg.quads_neighbor_encoder_type if neighbor_encoder_type == 'mean_embed': self.neighbor_encoder = QuadNeighborhoodEncoderDeepsets( cfg, self.neighbor_obs_dim, self.neighbor_hidden_size, self.use_spectral_norm, self.self_obs_dim, self.num_use_neighbor_obs) elif neighbor_encoder_type == 'attention': self.neighbor_encoder = QuadNeighborhoodEncoderAttention( cfg, self.neighbor_obs_dim, self.neighbor_hidden_size, self.use_spectral_norm, self.self_obs_dim, self.num_use_neighbor_obs) elif neighbor_encoder_type == 'mlp': self.neighbor_encoder = QuadNeighborhoodEncoderMlp( cfg, self.neighbor_obs_dim, self.neighbor_hidden_size, self.use_spectral_norm, self.self_obs_dim, self.num_use_neighbor_obs) elif neighbor_encoder_type == 'no_encoder': self.neighbor_encoder = None # blind agent else: raise NotImplementedError if self.neighbor_encoder: neighbor_encoder_out_size = self.neighbor_hidden_size fc_encoder_layer = cfg.hidden_size # encode the current drone's observations self.self_encoder = nn.Sequential( fc_layer(self.self_obs_dim, fc_encoder_layer, spec_norm=self.use_spectral_norm), nonlinearity(cfg), fc_layer(fc_encoder_layer, fc_encoder_layer, spec_norm=self.use_spectral_norm), nonlinearity(cfg)) self_encoder_out_size = calc_num_elements(self.self_encoder, (self.self_obs_dim, )) # encode the obstacle observations obstacle_encoder_out_size = 0 if self.obstacle_mode != 'no_obstacles': self.obstacle_obs_dim = 10 # internal param, pos_vel_size_type, 3 * 3 + 1, note: for size, we should consider it's length in xyz direction self.obstacle_hidden_size = cfg.quads_obstacle_hidden_size # internal param self.obstacle_encoder = nn.Sequential( fc_layer(self.obstacle_obs_dim, self.obstacle_hidden_size, spec_norm=self.use_spectral_norm), nonlinearity(cfg), fc_layer(self.obstacle_hidden_size, self.obstacle_hidden_size, spec_norm=self.use_spectral_norm), nonlinearity(cfg), ) obstacle_encoder_out_size = calc_num_elements( self.obstacle_encoder, (self.obstacle_obs_dim, )) total_encoder_out_size = self_encoder_out_size + neighbor_encoder_out_size + obstacle_encoder_out_size # this is followed by another fully connected layer in the action parameterization, so we add a nonlinearity here self.feed_forward = nn.Sequential( fc_layer(total_encoder_out_size, 2 * cfg.hidden_size, spec_norm=self.use_spectral_norm), nn.Tanh(), ) self.encoder_out_size = 2 * cfg.hidden_size