|
import os |
|
import pickle |
|
from typing import Any, Union |
|
from typing import Dict, List |
|
|
|
import numpy as np |
|
import pytorch_lightning as pl |
|
import torch |
|
from pytorch_lightning.callbacks import ModelCheckpoint |
|
from torch.optim import Optimizer |
|
from torch.optim.lr_scheduler import LRScheduler |
|
|
|
from navsim.agents.abstract_agent import AbstractAgent |
|
from navsim.agents.dm.dm_config import DMConfig |
|
from navsim.agents.dm.dm_features import DMTargetBuilder, DMFeatureBuilder |
|
from navsim.agents.dm.dm_loss_fn import dm_imi_loss |
|
from navsim.agents.dm.dm_model import DMModel |
|
from navsim.common.dataclasses import SensorConfig |
|
from navsim.planning.training.abstract_feature_target_builder import ( |
|
AbstractFeatureBuilder, |
|
AbstractTargetBuilder, |
|
) |
|
|
|
|
|
class DMAgent(AbstractAgent): |
|
def __init__( |
|
self, |
|
config: DMConfig, |
|
lr: float, |
|
checkpoint_path: str = None, |
|
pdm_split=None, |
|
metrics=None, |
|
): |
|
super().__init__() |
|
config.trajectory_pdm_weight = { |
|
'noc': 3.0, |
|
'da': 3.0, |
|
'ttc': 2.0, |
|
'progress': config.progress_weight, |
|
'comfort': 1.0, |
|
} |
|
self._config = config |
|
self._lr = lr |
|
self.metrics = metrics |
|
self._checkpoint_path = checkpoint_path |
|
self.vadv2_model:DMModel = DMModel(config) |
|
self.vocab_size = config.vocab_size |
|
self.backbone_wd = config.backbone_wd |
|
new_pkl_dir = f'vocab_score_full_{self.vocab_size}_navtrain' |
|
self.vocab_pdm_score_full = pickle.load( |
|
open(f'{os.getenv("NAVSIM_TRAJPDM_ROOT")}/{new_pkl_dir}/{pdm_split}.pkl', 'rb')) |
|
|
|
def name(self) -> str: |
|
"""Inherited, see superclass.""" |
|
|
|
return self.__class__.__name__ |
|
|
|
def initialize(self) -> None: |
|
"""Inherited, see superclass.""" |
|
state_dict: Dict[str, Any] = torch.load(self._checkpoint_path, map_location=torch.device("cpu"))["state_dict"] |
|
self.load_state_dict({k.replace("agent.", ""): v for k, v in state_dict.items()}) |
|
|
|
def get_sensor_config(self) -> SensorConfig: |
|
"""Inherited, see superclass.""" |
|
return SensorConfig( |
|
cam_f0=[3], |
|
cam_l0=[3], |
|
cam_l1=[3], |
|
cam_l2=[3], |
|
cam_r0=[3], |
|
cam_r1=[3], |
|
cam_r2=[3], |
|
cam_b0=[3], |
|
lidar_pc=[], |
|
) |
|
|
|
def get_target_builders(self) -> List[AbstractTargetBuilder]: |
|
return [DMTargetBuilder(config=self._config)] |
|
|
|
def get_feature_builders(self) -> List[AbstractFeatureBuilder]: |
|
return [DMFeatureBuilder(config=self._config)] |
|
|
|
def forward(self, features: Dict[str, torch.Tensor]) -> Dict[str, torch.Tensor]: |
|
return self.vadv2_model(features) |
|
|
|
def forward_train(self, features, interpolated_traj): |
|
return self.vadv2_model(features, interpolated_traj) |
|
|
|
def compute_loss( |
|
self, |
|
features: Dict[str, torch.Tensor], |
|
targets: Dict[str, torch.Tensor], |
|
predictions: Dict[str, torch.Tensor], |
|
tokens=None |
|
) -> Union[torch.Tensor, Dict[str, torch.Tensor]]: |
|
return dm_imi_loss(targets, predictions, self._config, self.vadv2_model._trajectory_head) |
|
|
|
def get_optimizers(self) -> Union[Optimizer, Dict[str, Union[Optimizer, LRScheduler]]]: |
|
backbone_params_name = '_backbone.image_encoder' |
|
img_backbone_params = list( |
|
filter(lambda kv: backbone_params_name in kv[0], self.vadv2_model.named_parameters())) |
|
default_params = list(filter(lambda kv: backbone_params_name not in kv[0], self.vadv2_model.named_parameters())) |
|
params_lr_dict = [ |
|
{'params': [tmp[1] for tmp in default_params]}, |
|
{ |
|
'params': [tmp[1] for tmp in img_backbone_params], |
|
'lr': self._lr * self._config.lr_mult_backbone, |
|
'weight_decay': self.backbone_wd |
|
} |
|
] |
|
return torch.optim.Adam(params_lr_dict, lr=self._lr) |
|
|
|
def get_training_callbacks(self) -> List[pl.Callback]: |
|
return [ |
|
|
|
ModelCheckpoint( |
|
save_top_k=30, |
|
monitor="val/loss_epoch", |
|
mode="min", |
|
dirpath=f"{os.environ.get('NAVSIM_EXP_ROOT')}/{self._config.ckpt_path}/", |
|
filename="{epoch:02d}-{step:04d}", |
|
) |
|
] |
|
|
