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.get_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)
# 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