| 136 | |
| 137 | |
| 138 | def forward(self, img, state, target_point): |
| 139 | feature_emb, cnn_feature = self.perception(img) |
| 140 | outputs = {} |
| 141 | outputs['pred_speed'] = self.speed_branch(feature_emb) |
| 142 | measurement_feature = self.measurements(state) |
| 143 | |
| 144 | j_traj = self.join_traj(torch.cat([feature_emb, measurement_feature], 1)) |
| 145 | outputs['pred_value_traj'] = self.value_branch_traj(j_traj) |
| 146 | outputs['pred_features_traj'] = j_traj |
| 147 | z = j_traj |
| 148 | output_wp = list() |
| 149 | traj_hidden_state = list() |
| 150 | |
| 151 | # initial input variable to GRU |
| 152 | x = torch.zeros(size=(z.shape[0], 2), dtype=z.dtype).type_as(z) |
| 153 | |
| 154 | # autoregressive generation of output waypoints |
| 155 | for _ in range(self.config.pred_len): |
| 156 | x_in = torch.cat([x, target_point], dim=1) |
| 157 | z = self.decoder_traj(x_in, z) |
| 158 | traj_hidden_state.append(z) |
| 159 | dx = self.output_traj(z) |
| 160 | x = dx + x |
| 161 | output_wp.append(x) |
| 162 | |
| 163 | pred_wp = torch.stack(output_wp, dim=1) |
| 164 | outputs['pred_wp'] = pred_wp |
| 165 | |
| 166 | traj_hidden_state = torch.stack(traj_hidden_state, dim=1) |
| 167 | init_att = self.init_att(measurement_feature).view(-1, 1, 8, 29) |
| 168 | feature_emb = torch.sum(cnn_feature*init_att, dim=(2, 3)) |
| 169 | j_ctrl = self.join_ctrl(torch.cat([feature_emb, measurement_feature], 1)) |
| 170 | outputs['pred_value_ctrl'] = self.value_branch_ctrl(j_ctrl) |
| 171 | outputs['pred_features_ctrl'] = j_ctrl |
| 172 | policy = self.policy_head(j_ctrl) |
| 173 | outputs['mu_branches'] = self.dist_mu(policy) |
| 174 | outputs['sigma_branches'] = self.dist_sigma(policy) |
| 175 | |
| 176 | x = j_ctrl |
| 177 | mu = outputs['mu_branches'] |
| 178 | sigma = outputs['sigma_branches'] |
| 179 | future_feature, future_mu, future_sigma = [], [], [] |
| 180 | |
| 181 | # initial hidden variable to GRU |
| 182 | h = torch.zeros(size=(x.shape[0], 256), dtype=x.dtype).type_as(x) |
| 183 | |
| 184 | for _ in range(self.config.pred_len): |
| 185 | x_in = torch.cat([x, mu, sigma], dim=1) |
| 186 | h = self.decoder_ctrl(x_in, h) |
| 187 | wp_att = self.wp_att(torch.cat([h, traj_hidden_state[:, _]], 1)).view(-1, 1, 8, 29) |
| 188 | new_feature_emb = torch.sum(cnn_feature*wp_att, dim=(2, 3)) |
| 189 | merged_feature = self.merge(torch.cat([h, new_feature_emb], 1)) |
| 190 | dx = self.output_ctrl(merged_feature) |
| 191 | x = dx + x |
| 192 | |
| 193 | policy = self.policy_head(x) |
| 194 | mu = self.dist_mu(policy) |
| 195 | sigma = self.dist_sigma(policy) |