from typing import Dict
import numpy as np
import torch
import torch.nn as nn
from navsim.agents.transfuser.transfuser_backbone import TransfuserBackbone
from navsim.agents.transfuser.transfuser_backbone_conv import TransfuserBackboneConv
from navsim.agents.transfuser.transfuser_backbone_moe import TransfuserBackboneMoe
from navsim.agents.transfuser.transfuser_backbone_moe_ult32 import TransfuserBackboneMoeUlt32
from navsim.agents.transfuser.transfuser_backbone_vit import TransfuserBackboneViT
from navsim.agents.transfuser.transfuser_model import AgentHead
from navsim.agents.utils.attn import MemoryEffTransformer
from navsim.agents.utils.nerf import nerf_positional_encoding
from navsim.agents.vadv2.vadv2_config import Vadv2Config
class Vadv2ModelPDMProgress(nn.Module):
def __init__(self, config: Vadv2Config):
super().__init__()
self._query_splits = [
config.num_bounding_boxes,
]
self._config = config
assert config.backbone_type in ['vit', 'intern', 'vov', 'resnet', 'eva', 'moe', 'moe_ult32', 'swin']
if config.backbone_type == 'vit' or config.backbone_type == 'eva':
self._backbone = TransfuserBackboneViT(config)
elif config.backbone_type == 'intern' or config.backbone_type == 'vov' or config.backbone_type == 'swin':
self._backbone = TransfuserBackboneConv(config)
elif config.backbone_type == 'moe':
self._backbone = TransfuserBackboneMoe(config)
elif config.backbone_type == 'moe_ult32':
self._backbone = TransfuserBackboneMoeUlt32(config)
else:
self._backbone = TransfuserBackbone(config)
bev_size = config.lidar_vert_anchors * config.lidar_horz_anchors
bev_c = self._backbone.lidar_encoder.feature_info.info[4]['num_chs']
self._keyval_embedding = nn.Embedding(
bev_size, config.tf_d_model
) # 8x8 feature grid + trajectory
self._query_embedding = nn.Embedding(sum(self._query_splits), config.tf_d_model)
# usually, the BEV features are variable in size.
self._bev_downscale = nn.Conv2d(bev_c, config.tf_d_model, kernel_size=1)
# todo drop ego status like plantf
# assert config.num_ego_status == 1
# assert not config.use_nerf
self._status_encoding = nn.Linear((4 + 2 + 2) * config.num_ego_status, config.tf_d_model)
self._bev_semantic_head = nn.Sequential(
nn.Conv2d(
config.bev_features_channels,
config.bev_features_channels,
kernel_size=(3, 3),
stride=1,
padding=(1, 1),
bias=True,
),
nn.ReLU(inplace=True),
nn.Conv2d(
config.bev_features_channels,
config.num_bev_classes,
kernel_size=(1, 1),
stride=1,
padding=0,
bias=True,
),
nn.Upsample(
size=(config.lidar_resolution_height // 2, config.lidar_resolution_width),
mode="bilinear",
align_corners=False,
),
)
tf_decoder_layer = nn.TransformerDecoderLayer(
d_model=config.tf_d_model,
nhead=config.tf_num_head,
dim_feedforward=config.tf_d_ffn,
dropout=config.tf_dropout,
batch_first=True,
)
self._tf_decoder = nn.TransformerDecoder(tf_decoder_layer, config.tf_num_layers)
self._agent_head = AgentHead(
num_agents=config.num_bounding_boxes,
d_ffn=config.tf_d_ffn,
d_model=config.tf_d_model,
)
self._trajectory_head = Vadv2HeadPDMProgress(
num_poses=config.trajectory_sampling.num_poses,
d_ffn=config.tf_d_ffn,
d_model=config.tf_d_model,
nhead=config.vadv2_head_nhead,
nlayers=config.vadv2_head_nlayers,
vocab_path=config.vocab_path,
config=config
)
def forward(self, features: Dict[str, torch.Tensor],
interpolated_traj=None) -> Dict[str, torch.Tensor]:
# Todo egostatus
camera_feature: torch.Tensor = features["camera_feature"]
lidar_feature: torch.Tensor = features["lidar_feature"]
status_feature: torch.Tensor = features["status_feature"]
batch_size = status_feature.shape[0]
bev_feature_upscale, bev_feature, _ = self._backbone(camera_feature, lidar_feature)
bev_feature = self._bev_downscale(bev_feature).flatten(-2, -1)
bev_feature = bev_feature.permute(0, 2, 1)
if self._config.num_ego_status == 1 and status_feature.shape[1] == 32:
status_encoding = self._status_encoding(status_feature[:, :8])
else:
status_encoding = self._status_encoding(status_feature)
keyval = bev_feature
keyval += self._keyval_embedding.weight[None, ...]
query = self._query_embedding.weight[None, ...].repeat(batch_size, 1, 1)
agents_query = self._tf_decoder(query, keyval)
bev_semantic_map = self._bev_semantic_head(bev_feature_upscale)
output: Dict[str, torch.Tensor] = {"bev_semantic_map": bev_semantic_map}
# 轨迹预测head
trajectory = self._trajectory_head(keyval, status_encoding, interpolated_traj)
output.update(trajectory)
agents = self._agent_head(agents_query)
output.update(agents)
return output
class Vadv2HeadPDMProgress(nn.Module):
def __init__(self, num_poses: int, d_ffn: int, d_model: int, vocab_path: str,
nhead: int, nlayers: int, config: Vadv2Config = None
):
super().__init__()
self._num_poses = num_poses
self.transformer = nn.TransformerDecoder(
nn.TransformerDecoderLayer(
d_model, nhead, d_ffn,
dropout=0.0, batch_first=True
), nlayers
)
self.vocab = nn.Parameter(
torch.from_numpy(np.load(vocab_path)),
requires_grad=False
)
self.heads = nn.ModuleDict({
'noc': nn.Sequential(
nn.Linear(d_model, d_ffn),
nn.ReLU(),
nn.Linear(d_ffn, 1),
),
'da':
nn.Sequential(
nn.Linear(d_model, d_ffn),
nn.ReLU(),
nn.Linear(d_ffn, 1),
),
'ttc': nn.Sequential(
nn.Linear(d_model, d_ffn),
nn.ReLU(),
nn.Linear(d_ffn, 1),
),
'comfort': nn.Sequential(
nn.Linear(d_model, d_ffn),
nn.ReLU(),
nn.Linear(d_ffn, 1),
),
'progress': nn.Sequential(
nn.Linear(d_model, d_ffn),
nn.ReLU(),
nn.Linear(d_ffn, 1),
),
'imi': nn.Sequential(
nn.Linear(d_model, d_ffn),
nn.ReLU(),
nn.Linear(d_ffn, d_ffn),
nn.ReLU(),
nn.Linear(d_ffn, 1),
)
})
self.inference_imi_weight = config.inference_imi_weight
self.inference_da_weight = config.inference_da_weight
self.normalize_vocab_pos = config.normalize_vocab_pos
if self.normalize_vocab_pos:
self.encoder = MemoryEffTransformer(
d_model=d_model,
nhead=nhead,
dim_feedforward=d_model * 4,
dropout=0.0
)
self.use_nerf = config.use_nerf
if self.use_nerf:
self.pos_embed = nn.Sequential(
nn.Linear(1040, d_ffn),
nn.ReLU(),
nn.Linear(d_ffn, d_model),
)
else:
self.pos_embed = nn.Sequential(
nn.Linear(num_poses * 3, d_ffn),
nn.ReLU(),
nn.Linear(d_ffn, d_model),
)
def forward(self, bev_feature, status_encoding, interpolated_traj) -> Dict[str, torch.Tensor]:
# todo sinusoidal embedding
# vocab: 4096, 40, 3
# bev_feature: B, 32, C
# embedded_vocab: B, 4096, C
vocab = self.vocab.data
L, HORIZON, _ = vocab.shape
B = bev_feature.shape[0]
if self.use_nerf:
vocab = torch.cat(
[
nerf_positional_encoding(vocab[..., :2]),
torch.cos(vocab[..., -1])[..., None],
torch.sin(vocab[..., -1])[..., None],
], dim=-1
)
if self.normalize_vocab_pos:
embedded_vocab = self.pos_embed(vocab.view(L, -1))[None]
embedded_vocab = self.encoder(embedded_vocab).repeat(B, 1, 1)
else:
embedded_vocab = self.pos_embed(vocab.view(L, -1))[None].repeat(B, 1, 1)
tr_out = self.transformer(embedded_vocab, bev_feature)
dist_status = tr_out + status_encoding.unsqueeze(1)
result = {}
# selected_indices: B,
for k, head in self.heads.items():
if k == 'imi':
result[k] = head(dist_status).squeeze(-1)
else:
result[k] = head(dist_status).squeeze(-1).sigmoid()
# how
scores = (
self.inference_imi_weight * result['imi'].softmax(-1).log() +
result['noc'].log() +
self.inference_da_weight * result['da'].log() +
(5 * result['ttc'] + 2 * result['comfort'] + 5 * result['progress']).log()
)
selected_indices = scores.argmax(1)
result["trajectory"] = self.vocab.data[selected_indices]
result["trajectory_vocab"] = self.vocab.data
result["selected_indices"] = selected_indices
return result