File size: 1,808 Bytes
68e55d1 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 | """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):
# None skips the legacy all-dataset scorer cache and Ray initialization.
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):
# Legacy zip(proposal_list, scores_index) truncated stages at small batch
# sizes. B2D has no anchor pseudo labels: supervise every stage with GT.
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
|