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, neighbor_obs_dim, neighbor_hidden_size,
use_spectral_norm, self_obs_dim, num_use_neighbor_obs):
super().__init__(cfg, self_obs_dim, neighbor_obs_dim,
neighbor_hidden_size, num_use_neighbor_obs)
self.embedding_mlp = nn.Sequential(
fc_layer(neighbor_obs_dim,
neighbor_hidden_size,
spec_norm=use_spectral_norm), nonlinearity(cfg),
fc_layer(neighbor_hidden_size,
neighbor_hidden_size,
spec_norm=use_spectral_norm), nonlinearity(cfg))
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, neighbor_obs_dim, neighbor_hidden_size,
use_spectral_norm, self_obs_dim, num_use_neighbor_obs):
super().__init__(cfg, self_obs_dim, neighbor_obs_dim,
neighbor_hidden_size, num_use_neighbor_obs)
self.self_obs_dim = self_obs_dim
# outputs e_i from the paper
self.embedding_mlp = nn.Sequential(
fc_layer(self_obs_dim + neighbor_obs_dim,
neighbor_hidden_size,
spec_norm=use_spectral_norm), nonlinearity(cfg),
fc_layer(neighbor_hidden_size,
neighbor_hidden_size,
spec_norm=use_spectral_norm), nonlinearity(cfg))
# outputs h_i from the paper
self.neighbor_value_mlp = nn.Sequential(
fc_layer(neighbor_hidden_size,
neighbor_hidden_size,
spec_norm=use_spectral_norm),
nonlinearity(cfg),
fc_layer(neighbor_hidden_size,
neighbor_hidden_size,
spec_norm=use_spectral_norm),
nonlinearity(cfg),
)
# outputs scalar score alpha_i for each neighbor i
self.attention_mlp = nn.Sequential(
fc_layer(
neighbor_hidden_size * 2,
neighbor_hidden_size,
spec_norm=use_spectral_norm
), # neighbor_hidden_size * 2 because we concat e_i and e_m
nonlinearity(cfg),
fc_layer(neighbor_hidden_size,
neighbor_hidden_size,
spec_norm=use_spectral_norm),
nonlinearity(cfg),
fc_layer(neighbor_hidden_size, 1),
)
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