| """Original EMA-JQTF Formula model with batch-local measured geometry.""" |
| import gzip,pickle |
| import torch |
| from navsim.agents.drive_jepa_perception_based.drive_jepa_agent import DriveJEPAAgent |
| from navsim.agents.drive_jepa_perception_based.score_module.compute_b2d_score import b2d_before_score,get_scores |
|
|
| class DenseAgent(DriveJEPAAgent): |
| def __init__(self,config,lr): |
| |
| super().__init__(config,lr,checkpoint_path=None) |
| self.bce_logit_loss=torch.nn.BCEWithLogitsLoss();self.b2d=True;self.ray=False |
| with gzip.open(config.b2d_map_cache_path,'rb') as f:self.b2d_map_infos=pickle.load(f) |
| self.b2d_before_score=b2d_before_score;self.get_scores=get_scores |
| def compute_score(self,targets,proposals,test=True): |
| local={sample['token']:sample for sample in targets['_scorer']} |
| self.train_metric_cache_paths=local;self.test_metric_cache_paths=local |
| return super().compute_score(targets,proposals,test) |
| def trajectory_loss_anchors(self,proposal_list,target_trajectory,config,scores_index): |
| |
| |
| if any(len(x) for x in scores_index):raise RuntimeError('Unexpected trajectory teacher/anchor targets') |
| total=0;mins=[];divs=[] |
| for proposals in proposal_list: |
| matched=(proposals-target_trajectory[:,None]).abs().sum(-1).mean(-1).amin(1).mean() |
| diversity=self.diversity_loss(proposals) |
| total=config.prev_weight*total+matched+config.inter_weight*diversity |
| mins.append(matched);divs.append(diversity) |
| return total,mins[-1],divs[-1],mins,divs |
|
|
