def loss_planning(self,
                      ego_fut_preds,
                      ego_fut_gt,
                      ego_fut_masks,
                      ego_fut_cmd,
                      lane_preds,
                      lane_score_preds,
                      agent_preds,
                      agent_fut_preds,
                      agent_score_preds,
                      agent_fut_cls_preds,
                      ego_cls_col_preds, 
                      ego_cls_bd_preds,
                      ego_cls_cl_preds,
                      ego_cls_expert_preds,
                      gt_agent_boxes,
                      gt_agent_feats,
                      gt_map_pts,
                      gt_map_labels,
                      ):
        """"Loss function for ego vehicle planning.
        Args:
            ego_fut_preds (Tensor): [B, num_cmd, fut_ts, 2]
            ego_fut_gt (Tensor): [B, fut_ts, 2]
            ego_fut_masks (Tensor): [B, fut_ts]
            ego_fut_cmd (Tensor): [B, num_cmd]
            lane_preds (Tensor): [B, num_vec, num_pts, 2]
            lane_score_preds (Tensor): [B, num_vec, 3]
            agent_preds (Tensor): [B, num_agent, 2]
            agent_fut_preds (Tensor): [B, num_agent, mot_fut_mode, fut_ts, 2]
            agent_score_preds (Tensor): [B, num_agent, 10]
            agent_fut_cls_scores (Tensor): [B, num_agent, mot_fut_mode]
            ego_cls_col_preds (Tensor): [B, num_plan_mode, 1]
            ego_cls_bd_preds (Tensor): [B, num_plan_mode, 1]
            ego_cls_cl_preds (Tensor): [B, num_plan_mode, 1]
            ego_cls_expert_preds (Tensor): [B, num_plan_mode, 1]

        Returns:
            loss_plan_cls_col (Tensor): cls col loss.
            loss_plan_cls_bd (Tensor): cls bd loss.
            loss_plan_cls_cl (Tensor): cls cl loss.
            loss_plan_cls_expert (Tensor): cls expert loss.
            loss_plan_reg (Tensor): planning l1 loss.
            loss_plan_bound (Tensor): planning map boundary loss.
            loss_plan_agent_dis (Tensor): planning agent distance loss.
            loss_plan_map_theta (Tensor): planning map theta loss.
        """

        batch = ego_fut_preds.shape[0]
        ego_fut_gt = ego_fut_gt.unsqueeze(1).repeat(1, self.plan_fut_mode, 1, 1)
        # loss_plan_l1_weight = ego_fut_cmd[..., None, None] * ego_fut_masks[:, None, :, None]
        # loss_plan_l1_weight = loss_plan_l1_weight.repeat(1, 1, 1, 2)

        # get plan cls target
        plan_col_labels, plan_bd_labels, plan_cl_labels = [], [], []
        plan_expert_labels, plan_expert_labels_weight = [], []
        for i in range(batch):
            plan_col_label = self.get_plan_col_target(
                ego_fut_preds[i].detach(),
                ego_fut_gt[i],
                gt_agent_boxes[i],
                gt_agent_feats[i])
            plan_bd_label = self.get_plan_bd_target(
                ego_fut_preds[i].detach(),
                gt_map_pts[i],
                gt_map_labels[i])
            plan_cl_label = self.get_plan_cl_target(
                ego_fut_preds[i].detach(),
                gt_map_pts[i],
                gt_map_labels[i])
            plan_expert_label, plan_expert_label_weight = self.get_plan_expert_target(
                ego_fut_preds[i].detach(),
                ego_fut_gt[i],
                ego_fut_masks[i],
                ego_cls_expert_preds[i],
                plan_col_label,
                plan_bd_label)

            plan_col_labels.append(plan_col_label)
            plan_bd_labels.append(plan_bd_label)
            plan_cl_labels.append(plan_cl_label)
            plan_expert_labels.append(plan_expert_label)
            plan_expert_labels_weight.append(plan_expert_label_weight)

        plan_col_labels = torch.stack(plan_col_labels, dim=0).to(ego_fut_preds.device)
        plan_bd_labels = torch.stack(plan_bd_labels, dim=0).to(ego_fut_preds.device)
        plan_cl_labels = torch.stack(plan_cl_labels, dim=0).to(ego_fut_preds.device)
        plan_expert_labels = torch.stack(plan_expert_labels, dim=0).to(ego_fut_preds.device)
        plan_expert_labels_weight = torch.stack(plan_expert_labels_weight,
                                                dim=0).to(ego_fut_preds.device)

        # plan collision classification loss
        loss_plan_cls_col = self.loss_plan_cls_col(
            ego_cls_col_preds.flatten(0, 1), plan_col_labels.flatten(),
            plan_col_labels.new_ones(batch*self.plan_fut_mode),
            avg_factor=batch*self.plan_fut_mode)

        # plan boundary overstepping classification loss
        loss_plan_cls_bd = self.loss_plan_cls_bd(
            ego_cls_bd_preds.flatten(0, 1), plan_bd_labels.flatten(),
            plan_bd_labels.new_ones(batch*self.plan_fut_mode),
            avg_factor=batch*self.plan_fut_mode)

        # plan centerline consistency classification loss
        plan_cl_weight = plan_cl_labels.flatten()
        plan_cl_labels_weight = (plan_cl_weight > -1.)  + (plan_cl_weight == -1.) * 0.01
        loss_plan_cls_cl = self.loss_plan_cls_cl(
            ego_cls_cl_preds.squeeze(-1).flatten(0, 1), plan_cl_labels.flatten(),
            plan_cl_labels_weight,
            # avg_factor=(plan_cl_labels > -1.).sum()
        )

        # plan expert driving behavior classification loss
        loss_plan_cls_expert = self.loss_plan_cls_expert(
            ego_cls_expert_preds.flatten(0, 1), plan_expert_labels.flatten(),
            plan_expert_labels_weight.flatten(),
            avg_factor=batch*self.plan_fut_mode)

        loss_plan_reg = (0. * ego_fut_preds).sum()
        # loss_plan_reg = self.loss_plan_reg(
        #     ego_fut_preds,
        #     ego_fut_gt,
        #     loss_plan_l1_weight
        # )

        loss_plan_bound = (0. * ego_fut_preds).sum()
        # loss_plan_bound = self.loss_plan_bound(
        #     ego_fut_preds[ego_fut_cmd==1],
        #     lane_preds,
        #     lane_score_preds,
        #     weight=ego_fut_masks
        # )

        loss_plan_agent_dis = (0. * ego_fut_preds).sum()
        # loss_plan_agent_dis = self.loss_plan_agent_dis(
        #     ego_fut_preds[ego_fut_cmd==1],
        #     agent_preds,
        #     agent_fut_preds,
        #     agent_score_preds,
        #     agent_fut_cls_preds,
        #     weight=ego_fut_masks[:, :, None].repeat(1, 1, 2)
        # )

        loss_plan_map_theta = (0. * ego_fut_preds).sum()
        # loss_plan_map_theta = self.loss_plan_map_theta(
        #     ego_fut_preds[ego_fut_cmd==1],
        #     lane_preds,
        #     lane_score_preds,
        #     weight=ego_fut_masks
        # )

        if digit_version(TORCH_VERSION) >= digit_version('1.8'):
            loss_plan_cls_col = torch.nan_to_num(loss_plan_cls_col)
            loss_plan_cls_bd = torch.nan_to_num(loss_plan_cls_bd)
            loss_plan_cls_cl = torch.nan_to_num(loss_plan_cls_cl)
            loss_plan_cls_expert = torch.nan_to_num(loss_plan_cls_expert)
            loss_plan_reg = torch.nan_to_num(loss_plan_reg)
            loss_plan_bound = torch.nan_to_num(loss_plan_bound)
            loss_plan_agent_dis = torch.nan_to_num(loss_plan_agent_dis)
            loss_plan_map_theta = torch.nan_to_num(loss_plan_map_theta)
        
        loss_plan_dict = dict()
        loss_plan_dict['loss_plan_cls_col'] = loss_plan_cls_col
        loss_plan_dict['loss_plan_cls_bd'] = loss_plan_cls_bd
        loss_plan_dict['loss_plan_cls_cl'] = loss_plan_cls_cl
        loss_plan_dict['loss_plan_cls_expert'] = loss_plan_cls_expert

        loss_plan_dict['loss_plan_reg'] = loss_plan_reg
        loss_plan_dict['loss_plan_bound'] = loss_plan_bound
        loss_plan_dict['loss_plan_agent_dis'] = loss_plan_agent_dis
        loss_plan_dict['loss_plan_map_theta'] = loss_plan_map_theta

        return loss_plan_dict
"Loss function for ego vehicle planning.
        Args:
            ego_fut_preds (Tensor): [B, num_cmd, fut_ts, 2]
            ego_fut_gt (Tensor): [B, fut_ts, 2]
            ego_fut_masks (Tensor): [B, fut_ts]
            ego_fut_cmd (Tensor): [B, num_cmd]
            lane_preds (Tensor): [B, num_vec, num_pts, 2]
            lane_score_preds (Tensor): [B, num_vec, 3]
            agent_preds (Tensor): [B, num_agent, 2]
            agent_fut_preds (Tensor): [B, num_agent, mot_fut_mode, fut_ts, 2]
            agent_score_preds (Tensor): [B, num_agent, 10]
            agent_fut_cls_scores (Tensor): [B, num_agent, mot_fut_mode]
            ego_cls_col_preds (Tensor): [B, num_plan_mode, 1]
            ego_cls_bd_preds (Tensor): [B, num_plan_mode, 1]
            ego_cls_cl_preds (Tensor): [B, num_plan_mode, 1]
            ego_cls_expert_preds (Tensor): [B, num_plan_mode, 1]

        Returns:
            loss_plan_cls_col (Tensor): cls col loss.
            loss_plan_cls_bd (Tensor): cls bd loss.
            loss_plan_cls_cl (Tensor): cls cl loss.
            loss_plan_cls_expert (Tensor): cls expert loss.
            loss_plan_reg (Tensor): planning l1 loss.
            loss_plan_bound (Tensor): planning map boundary loss.
            loss_plan_agent_dis (Tensor): planning agent distance loss.
            loss_plan_map_theta (Tensor): planning map theta loss.
        
VADv2/VADv2_head.py