| 24 | |
| 25 | |
| 26 | class Model(nn.Module): |
| 27 | def __init__(self, obs_shape, num_of_uav, device, trainable=True, hidden_size=params.temporal_hidden_size): |
| 28 | # todo 1: add parameter trainable=True |
| 29 | super(Model, self).__init__() |
| 30 | # feature extract |
| 31 | self.base = NNBase(obs_shape[0], device, trainable) |
| 32 | # actor |
| 33 | self.dist_dia = DiagGaussian(hidden_size + num_of_uav * 4, params.uav_action_dim * num_of_uav, |
| 34 | device) # continuous, (dx, dy,v,coll_t) |
| 35 | init_ = lambda m: init(m, |
| 36 | nn.init.orthogonal_, |
| 37 | lambda x: nn.init.constant_(x, 0)) |
| 38 | # critic |
| 39 | self.critic = nn.Sequential( |
| 40 | init_(nn.Linear(hidden_size + num_of_uav * 4, 1)) |
| 41 | ) |
| 42 | self.device = device |
| 43 | # todo 2: distinguish train and eval |
| 44 | if trainable: |
| 45 | self.train() |
| 46 | else: |
| 47 | self.eval() |
| 48 | |
| 49 | def act(self, inputs, uav_aoi, uav_snr, uav_compl, uav_tc_compl, temporal_hidden_state=None, mask=None, |
| 50 | spatial_hidden_state=None): |
| 51 | if params.use_rnn: |
| 52 | if params.use_spatial_att: |
| 53 | obs_feature, temporal_hidden_state, spatial_hidden_state, = self.base(inputs, temporal_hidden_state, |
| 54 | mask, spatial_hidden_state) |
| 55 | else: |
| 56 | obs_feature, temporal_hidden_state = self.base(inputs, temporal_hidden_state, mask) |
| 57 | |
| 58 | else: |
| 59 | obs_feature = self.base(inputs) |
| 60 | full_obs_feature = torch.cat( |
| 61 | [obs_feature, uav_aoi.float(), uav_snr.float(), uav_compl.float(), uav_tc_compl.float()], dim=1) |
| 62 | value = self.critic(full_obs_feature) |
| 63 | dist_dia = self.dist_dia(full_obs_feature) |
| 64 | action_dia = dist_dia.sample() |
| 65 | |
| 66 | action_log_probs_dia = dist_dia.log_probs(action_dia) |
| 67 | |
| 68 | # print('action log probs dia', action_log_probs_dia.mean()) |
| 69 | if params.use_rnn: |
| 70 | if params.use_spatial_att: |
| 71 | return value, action_dia, action_log_probs_dia, temporal_hidden_state, spatial_hidden_state |
| 72 | else: |
| 73 | return value, action_dia, action_log_probs_dia, temporal_hidden_state |
| 74 | |
| 75 | else: |
| 76 | return value, action_dia, action_log_probs_dia |
| 77 | |
| 78 | def get_value(self, inputs, uav_aoi, uav_snr, uav_compl, uav_tc_compl, temporal_hidden_state=None, mask=None, |
| 79 | spatial_hidden_state=None): |
| 80 | if params.use_rnn: |
| 81 | if params.use_spatial_att: |
| 82 | obs_feature, _, _ = self.base(inputs, temporal_hidden_state, mask, spatial_hidden_state) |
| 83 | else: |