|
from typing import Dict, Tuple, List
|
|
|
|
import pytorch_lightning as pl
|
|
import torch
|
|
from nuplan.planning.simulation.trajectory.trajectory_sampling import TrajectorySampling
|
|
from torch import Tensor
|
|
|
|
from navsim.agents.abstract_agent import AbstractAgent
|
|
from navsim.agents.vadv2.vadv2_agent import Vadv2Agent
|
|
from navsim.common.dataclasses import Trajectory
|
|
|
|
|
|
class AgentLightningModuleMap(pl.LightningModule):
|
|
def __init__(
|
|
self,
|
|
agent: AbstractAgent,
|
|
):
|
|
super().__init__()
|
|
self.agent = agent
|
|
|
|
def _step(
|
|
self,
|
|
batch: Tuple[Dict[str, Tensor], Dict[str, Tensor], List[str]],
|
|
logging_prefix: str,
|
|
):
|
|
features, targets = batch
|
|
if logging_prefix in ['train', 'val'] and isinstance(self.agent, Vadv2Agent):
|
|
prediction = self.agent.forward_train(features, targets['interpolated_traj'])
|
|
else:
|
|
prediction = self.agent.forward(features)
|
|
|
|
loss, loss_dict = self.agent.compute_loss(features, targets, prediction)
|
|
|
|
for k, v in loss_dict.items():
|
|
self.log(f"{logging_prefix}/{k}", v, on_step=True, on_epoch=True, prog_bar=True, sync_dist=True)
|
|
self.log(f"{logging_prefix}/loss", loss, on_step=True, on_epoch=True, prog_bar=True, sync_dist=True)
|
|
return loss
|
|
|
|
def training_step(
|
|
self,
|
|
batch: Tuple[Dict[str, Tensor], Dict[str, Tensor]],
|
|
batch_idx: int
|
|
):
|
|
return self._step(batch, "train")
|
|
|
|
def validation_step(
|
|
self,
|
|
batch: Tuple[Dict[str, Tensor], Dict[str, Tensor]],
|
|
batch_idx: int
|
|
):
|
|
return self._step(batch, "val")
|
|
|
|
def configure_optimizers(self):
|
|
return self.agent.get_optimizers()
|
|
|
|
def predict_step(
|
|
self,
|
|
batch: Tuple[Dict[str, Tensor], Dict[str, Tensor]],
|
|
batch_idx: int
|
|
):
|
|
features, targets, tokens = batch
|
|
self.agent.eval()
|
|
with torch.no_grad():
|
|
predictions = self.agent.forward(features)
|
|
poses = predictions["trajectory"].cpu().numpy()
|
|
|
|
imis = predictions["imi"].softmax(-1).log().cpu().numpy()
|
|
nocs = predictions["noc"].log().cpu().numpy()
|
|
das = predictions["da"].log().cpu().numpy()
|
|
ttcs = predictions["ttc"].log().cpu().numpy()
|
|
comforts = predictions["comfort"].log().cpu().numpy()
|
|
progresses = predictions["progress"].log().cpu().numpy()
|
|
if poses.shape[1] == 40:
|
|
interval_length = 0.1
|
|
else:
|
|
interval_length = 0.5
|
|
|
|
return {token: {
|
|
'trajectory': Trajectory(pose, TrajectorySampling(time_horizon=4, interval_length=interval_length)),
|
|
'imi': imi,
|
|
'noc': noc,
|
|
'da': da,
|
|
'ttc': ttc,
|
|
'comfort': comfort,
|
|
'progress': progress
|
|
} for pose, imi, noc, da, ttc, comfort, progress, token in zip(poses, imis, nocs, das, ttcs, comforts, progresses,
|
|
tokens)}
|
|
|
|
|
|
|
|
|
|
|
|
|
