navsim/agents/hydra/hydra_model_pe.py

from typing import Dict

import numpy as np
import torch
import torch.nn as nn

from navsim.agents.hydra.hydra_backbone_pe import HydraBackbonePE
from navsim.agents.hydra.hydra_config import HydraConfig
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
from mmcv.cnn.bricks.transformer import FFN, build_positional_encoding
from navsim.agents.utils.positional_encoding import SinePositionalEncoding3D
from mmcv.cnn import Conv2d
class HydraModelPE(nn.Module):
    def __init__(self, config: HydraConfig):
        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':
            raise ValueError(f'{config.backbone_type} not supported')
        elif config.backbone_type == 'intern' or config.backbone_type == 'vov' or config.backbone_type == 'swin' \
                or config.backbone_type == 'resnet':
            self._backbone = HydraBackbonePE(config)

        self._keyval_embedding = nn.Embedding(
            config.img_vert_anchors * config.img_horz_anchors, 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.downscale_layer = nn.Conv2d(self._backbone.img_feat_c, config.tf_d_model, kernel_size=1)
        self._status_encoding = nn.Linear((4 + 2 + 2) * config.num_ego_status, config.tf_d_model)

        self.depth_num = 64
        self.depth_start = 1
        self.position_range = [-32.0, -32.0, -10.0, 32.0, 32.0, 10.0]
        self.position_dim = 3 * self.depth_num
        self.embed_dims = 256
        self.sin_positional_encoding = dict(
            type='SinePositionalEncoding3D', num_feats=128, normalize=True)
        self.positional_encoding = build_positional_encoding(
            self.sin_positional_encoding)
        self.adapt_pos3d = nn.Sequential(
            nn.Conv2d(self.embed_dims*3//2, self.embed_dims*4, kernel_size=1, stride=1, padding=0),
            nn.ReLU(),
            nn.Conv2d(self.embed_dims*4, self.embed_dims, kernel_size=1, stride=1, padding=0),
        )
        self.position_encoder = nn.Sequential(
            nn.Conv2d(self.position_dim, self.embed_dims * 4, kernel_size=1, stride=1, padding=0),
            nn.ReLU(),
            nn.Conv2d(self.embed_dims * 4, self.embed_dims, kernel_size=1, stride=1, padding=0),
        )
        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 = HydraTrajHead(
            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 inverse_sigmoid(self, x, eps=1e-6):
        """Inverse sigmoid function.

        Args:
            x (Tensor): The input tensor.
            eps (float): A small value to avoid numerical issues.

        Returns:
            Tensor: The logit value of the input.
        """
        x = x.clamp(min=eps, max=1 - eps)  # Ensure the input is within the valid range
        return torch.log(x / (1 - x))

    def position_embedding(self, features, img_features):
        eps = 1e-5
        img_features = img_features.unsqueeze(1)
        B, N, C, tar_H, tar_W = img_features.shape
        device = img_features.device
        crop_top = 28
        crop_left = 416
        H = [self._config.img_vert_anchors for _ in range(3)]
        W = [
            self._config.img_horz_anchors * 1088 // (1088 * 2 + 1920),
            self._config.img_horz_anchors * 1920 // (1088 * 2 + 1920),
            self._config.img_horz_anchors * 1088 // (1088 * 2 + 1920)
        ]

        # 左视图(16,17)
        coords_h_l = torch.arange(H[0], device=device).float() * 1080 / H[0] + crop_top / H[0]
        coords_w_l = torch.arange(W[0], device=device).float() * 1920 / W[0] + crop_left / W[0]
        # 前视图(16,30)
        coords_h_f = torch.arange(H[1], device=device).float() * 1080 / H[1] + crop_top / H[1]
        coords_w_f = torch.arange(W[1], device=device).float() * 1920 / W[1]
        # 右视图(16,17)
        coords_h_r = torch.arange(H[2], device=device).float() * 1080 / H[2] + crop_top / H[2]
        coords_w_r = torch.arange(W[2], device=device).float() * 1920 / W[2] + crop_left / W[2]

        index = torch.arange(start=0, end=self.depth_num, step=1, device=img_features.device).float()
        index_1 = index + 1
        bin_size = (self.position_range[3] - self.depth_start) / (self.depth_num * (1 + self.depth_num))
        coords_d = self.depth_start + bin_size * index * index_1

        D = coords_d.shape[0]
        coords = [1] * 3  # 0,1,2 -> front, left, right
        coords[0] = torch.stack(torch.meshgrid([coords_w_l, coords_h_l, coords_d])).permute(1, 2, 3, 0)  # W, H, D, 3
        coords[1] = torch.stack(torch.meshgrid([coords_w_f, coords_h_f, coords_d])).permute(1, 2, 3, 0)  # W, H, D, 3
        coords[2] = torch.stack(torch.meshgrid([coords_w_r, coords_h_r, coords_d])).permute(1, 2, 3, 0)  # W, H, D, 3
        # coords = torch.cat((coords, torch.ones_like(coords[..., :1])), -1)
        coords[0][..., :2] = coords[0][..., :2] * torch.max(coords[0][..., 2:3],
                                                            torch.ones_like(coords[0][..., 2:3]) * eps)
        coords[1][..., :2] = coords[1][..., :2] * torch.max(coords[1][..., 2:3],
                                                            torch.ones_like(coords[1][..., 2:3]) * eps)
        coords[2][..., :2] = coords[2][..., :2] * torch.max(coords[2][..., 2:3],
                                                            torch.ones_like(coords[2][..., 2:3]) * eps)

        # img_meta
        # img2lidars = ?
        pos_3d_embed = None
        for i in range(3):
            sensor2lidar_rotation = features["sensor2lidar_rotation"][i]
            sensor2lidar_translation = features["sensor2lidar_translation"][i]
            intrinsics = features["intrinsics"][i]
            combine = torch.matmul(sensor2lidar_rotation, torch.inverse(intrinsics)).float()  # (B, 1, 3, 3) ?
            # print(combine.shape)

            # coords_front,coords_fleft,coords_fright (W, H, D, 3)
            # coords3d = torch.stack((coords_front, coords_fleft, coords_fright), dim=0) # (N, W, H, D, 3) -> (B, N, W, H, D, 3, 1)
            # coords = coords.view(1, H, W, D, 1, 3).repeat(B, 1, 1, 1, 1, 1)
            coords3d = coords[i].view(1, N, W[i], H[i], D, 3, 1).repeat(B, 1, 1, 1, 1, 1,
                                                                        1)  # (B, N, W, H, D, 3, 1) -> (B, N, W, H, D, 3, 3)
            combine = combine.view(B, N, 1, 1, 1, 3, 3).repeat(1, 1, W[i], H[i], D, 1, 1)
            coords3d = torch.matmul(combine, coords3d).squeeze(-1)  # (B, N, W, H, D, 3)
            sensor2lidar_translation = sensor2lidar_translation.view(B, N, 1, 1, 1, 3)
            coords3d += sensor2lidar_translation

            coords3d[..., 0:1] = (coords3d[..., 0:1] - self.position_range[0]) / (
                    self.position_range[3] - self.position_range[0])
            coords3d[..., 1:2] = (coords3d[..., 1:2] - self.position_range[1]) / (
                    self.position_range[4] - self.position_range[1])
            coords3d[..., 2:3] = (coords3d[..., 2:3] - self.position_range[2]) / (
                    self.position_range[5] - self.position_range[2])
            # coords_mask = (coords3d > 1.0) | (coords3d < 0.0)
            # coords_mask = coords_mask.flatten(-2).sum(-1) > (D * 0.5)
            # coords_mask = coords_mask.permute(0, 1, 3, 2)
            #     for j in range(1000000):
            #         print(coords3d.shape)
            # (2, 1, 17, 16, 64, 3) -> (B, N, W, H, D, 3)
            # (2, 1, 30, 16, 64, 3)
            # -> (2, 1, 17+30+17, 16, 64, 3)
            # coords3d = coords3d.permute(0, 1, 4, 5, 3, 2).contiguous().view(B * N, -1, H[i], W[i])
            if pos_3d_embed is None:
                pos_3d_embed = coords3d
            else:
                pos_3d_embed = torch.cat((pos_3d_embed, coords3d), dim=2)
        # for i in range(100000):
        # print(img_features.shape)
        pos_3d_embed = pos_3d_embed.permute(0, 1, 4, 5, 3, 2).contiguous().view(B * N, -1, tar_H, tar_W)
        coords3d = self.inverse_sigmoid(pos_3d_embed)
        coords_position_embeding = self.position_encoder(coords3d)
        return coords_position_embeding.view(B, N, self.embed_dims, tar_H, tar_W)

    def forward(self, features: Dict[str, torch.Tensor],
                interpolated_traj=None) -> Dict[str, torch.Tensor]:
        # Todo egostatus
        camera_feature: torch.Tensor = features["camera_feature"][0]
        # lidar_feature: torch.Tensor = features["lidar_feature"]
        status_feature: torch.Tensor = features["status_feature"]

        batch_size = status_feature.shape[0]
        assert (camera_feature.shape[0] == batch_size)
        img_features = self._backbone(camera_feature)
        img_features = self.downscale_layer(img_features)
        input_img_h, input_img_w = img_features.size(-2), img_features.size(-1)
        masks = img_features.new_ones(
            (img_features.shape[0], 1, input_img_h, input_img_w))

        coords_position_embeding = self.position_embedding(features, img_features)
        sin_embed = self.positional_encoding(masks)
        sin_embed = self.adapt_pos3d(sin_embed.flatten(0, 1)).view(img_features.size())
        pos_embed = coords_position_embeding.squeeze(1) + sin_embed
        # img_features = img_features.copy()
        img_features = img_features + pos_embed  # (B, N, self.embed_dims, H, W)
        img_features = img_features.flatten(-2, -1)
        img_features = img_features.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 = img_features

        keyval += self._keyval_embedding.weight[None, ...]

        query = self._query_embedding.weight[None, ...].repeat(batch_size, 1, 1)
        agents_query = self._tf_decoder(query, keyval)

        output: Dict[str, torch.Tensor] = {}
        trajectory = self._trajectory_head(keyval, status_encoding, interpolated_traj)
        output.update(trajectory)
        agents = self._agent_head(agents_query)
        output.update(agents)

        return output


class HydraTrajHead(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 = (
                0.05 * result['imi'].softmax(-1).log() +
                0.5 * result['noc'].log() +
                0.5 * result['da'].log() +
                8.0 * (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