| 19 | |
| 20 | |
| 21 | class TCP_planner(pl.LightningModule): |
| 22 | def __init__(self, config, lr): |
| 23 | super().__init__() |
| 24 | self.lr = lr |
| 25 | self.config = config |
| 26 | self.model = TCP(config) |
| 27 | self._load_weight() |
| 28 | |
| 29 | def _load_weight(self): |
| 30 | rl_state_dict = torch.load(self.config.rl_ckpt, map_location='cpu')['policy_state_dict'] |
| 31 | self._load_state_dict(self.model.value_branch_traj, rl_state_dict, 'value_head') |
| 32 | self._load_state_dict(self.model.value_branch_ctrl, rl_state_dict, 'value_head') |
| 33 | self._load_state_dict(self.model.dist_mu, rl_state_dict, 'dist_mu') |
| 34 | self._load_state_dict(self.model.dist_sigma, rl_state_dict, 'dist_sigma') |
| 35 | |
| 36 | def _load_state_dict(self, il_net, rl_state_dict, key_word): |
| 37 | rl_keys = [k for k in rl_state_dict.keys() if key_word in k] |
| 38 | il_keys = il_net.state_dict().keys() |
| 39 | assert len(rl_keys) == len(il_net.state_dict().keys()), f'mismatch number of layers loading {key_word}' |
| 40 | new_state_dict = OrderedDict() |
| 41 | for k_il, k_rl in zip(il_keys, rl_keys): |
| 42 | new_state_dict[k_il] = rl_state_dict[k_rl] |
| 43 | il_net.load_state_dict(new_state_dict) |
| 44 | |
| 45 | def forward(self, batch): |
| 46 | pass |
| 47 | |
| 48 | def training_step(self, batch, batch_idx): |
| 49 | front_img = batch['front_img'] |
| 50 | speed = batch['speed'].to(dtype=torch.float32).view(-1,1) / 12. |
| 51 | target_point = batch['target_point'].to(dtype=torch.float32) |
| 52 | command = batch['target_command'] |
| 53 | |
| 54 | state = torch.cat([speed, target_point, command], 1) |
| 55 | value = batch['value'].view(-1,1) |
| 56 | feature = batch['feature'] |
| 57 | |
| 58 | gt_waypoints = batch['waypoints'] |
| 59 | |
| 60 | pred = self.model(front_img, state, target_point) |
| 61 | |
| 62 | dist_sup = Beta(batch['action_mu'], batch['action_sigma']) |
| 63 | dist_pred = Beta(pred['mu_branches'], pred['sigma_branches']) |
| 64 | kl_div = torch.distributions.kl_divergence(dist_sup, dist_pred) |
| 65 | action_loss = torch.mean(kl_div[:, 0]) *0.5 + torch.mean(kl_div[:, 1]) *0.5 |
| 66 | speed_loss = F.l1_loss(pred['pred_speed'], speed) * self.config.speed_weight |
| 67 | value_loss = (F.mse_loss(pred['pred_value_traj'], value) + F.mse_loss(pred['pred_value_ctrl'], value)) * self.config.value_weight |
| 68 | feature_loss = (F.mse_loss(pred['pred_features_traj'], feature) +F.mse_loss(pred['pred_features_ctrl'], feature))* self.config.features_weight |
| 69 | |
| 70 | future_feature_loss = 0 |
| 71 | future_action_loss = 0 |
| 72 | for i in range(self.config.pred_len): |
| 73 | dist_sup = Beta(batch['future_action_mu'][i], batch['future_action_sigma'][i]) |
| 74 | dist_pred = Beta(pred['future_mu'][i], pred['future_sigma'][i]) |
| 75 | kl_div = torch.distributions.kl_divergence(dist_sup, dist_pred) |
| 76 | future_action_loss += torch.mean(kl_div[:, 0]) *0.5 + torch.mean(kl_div[:, 1]) *0.5 |
| 77 | future_feature_loss += F.mse_loss(pred['future_feature'][i], batch['future_feature'][i]) * self.config.features_weight |
| 78 | future_feature_loss /= self.config.pred_len |