TrafficBots: Towards World Models for Autonomous Driving Simulation and Motion Prediction. ICRA 2023.
# Licensed under the CC BY-NC 4.0 license (https://creativecommons.org/licenses/by-nc/4.0/)
from typing import Optional
from pathlib import Path
import wandb
from pytorch_lightning import Callback, Trainer
from pytorch_lightning.callbacks import ModelCheckpoint
from pytorch_lightning.loggers import LoggerCollection, WandbLogger
from pytorch_lightning.utilities import rank_zero_only


def get_wandb_logger(trainer: Trainer) -> Optional[WandbLogger]:
    """Safely get Weights&Biases logger from Trainer."""

    if isinstance(trainer.logger, WandbLogger):
        return trainer.logger

    if isinstance(trainer.logger, LoggerCollection):
        for logger in trainer.logger:
            if isinstance(logger, WandbLogger):
                return logger

    return None
    # raise Exception("You are using wandb related callback, but WandbLogger was not found for some reason...")


class ModelCheckpointWB(ModelCheckpoint):
    def save_checkpoint(self, trainer) -> None:
        super().save_checkpoint(trainer)
        if not hasattr(self, "_logged_model_time"):
            self._logged_model_time = {}
        logger = get_wandb_logger(trainer)
        if self.current_score is None:
            self.current_score = trainer.callback_metrics.get(self.monitor)
        if logger is not None:
            self._scan_and_log_checkpoints(logger)

    @rank_zero_only
    def _scan_and_log_checkpoints(self, wb_logger: WandbLogger) -> None:
        # adapted from pytorch_lightning 1.4.0: loggers/wandb.py
        checkpoints = {
            self.last_model_path: self.current_score,
            self.best_model_path: self.best_model_score,
        }
        checkpoints = sorted((Path(p).stat().st_mtime, p, s) for p, s in checkpoints.items() if Path(p).is_file())
        checkpoints = [
            c for c in checkpoints if c[1] not in self._logged_model_time.keys() or self._logged_model_time[c[1]] < c[0]
        ]
        # log iteratively all new checkpoints
        for t, p, s in checkpoints:
            metadata = {
                "score": s.item(),
                "original_filename": Path(p).name,
                "ModelCheckpoint": {
                    k: getattr(self, k)
                    for k in [
                        "monitor",
                        "mode",
                        "save_last",
                        "save_top_k",
                        "save_weights_only",
                        "_every_n_train_steps",
                        "_every_n_val_epochs",
                    ]
                    # ensure it does not break if `ModelCheckpoint` args change
                    if hasattr(self, k)
                },
            }
            artifact = wandb.Artifact(name=wb_logger.experiment.id, type="model", metadata=metadata)
            artifact.add_file(p, name="model.ckpt")
            aliases = ["latest", "best"] if p == self.best_model_path else ["latest"]
            wb_logger.experiment.log_artifact(artifact, aliases=aliases)
            # remember logged models - timestamp needed in case filename didn't change (lastkckpt or custom name)
            self._logged_model_time[p] = t


class WatchModel(Callback):
    """Make wandb watch model at the beginning of the run."""

    def __init__(self, log: str = "gradients", log_freq: int = 100):
        self._log = log
        self._log_freq = log_freq

    def on_train_start(self, trainer, pl_module):
        logger = get_wandb_logger(trainer)
        logger.watch(model=trainer.model, log=self._log, log_freq=self._log_freq)

src/callbacks/wandb_callbacks.py
zhejz/TrafficBots
# Licensed under the CC BY-NC 4.0 license (https://creativecommons.org/licenses/by-nc/4.0/)
from typing import Optional, Dict, Any, Tuple
from pytorch_lightning import LightningDataModule
from torch.utils.data import DataLoader, Dataset
import numpy as np
import h5py


class DatasetBase(Dataset[Dict[str, np.ndarray]]):
    def __init__(self, filepath: str, tensor_size: Dict[str, Tuple]) -> None:
        super().__init__()
        self.tensor_size = tensor_size
        self.filepath = filepath
        with h5py.File(self.filepath, "r", libver="latest", swmr=True) as hf:
            self.dataset_len = int(hf.attrs["data_len"])

    def __len__(self) -> int:
        return self.dataset_len


class DatasetTrain(DatasetBase):
    """
    The waymo 9-sec trainging.h5 is repetitive, start at {0, 2, 4, 5, 6, 8, 10} seconds within the 20-sec episode.
    Always train with the whole training.h5 dataset.
    limit_train_batches just for controlling the validation frequency.
    """

    def __getitem__(self, idx: int) -> Dict[str, np.ndarray]:
        idx = np.random.randint(self.dataset_len)
        idx_key = str(idx)
        out_dict = {"episode_idx": idx}
        with h5py.File(self.filepath, "r", libver="latest", swmr=True) as hf:
            for k in self.tensor_size.keys():
                out_dict[k] = np.ascontiguousarray(hf[idx_key][k])
        return out_dict


class DatasetVal(DatasetBase):
    # for validation.h5 and testing.h5
    def __getitem__(self, idx: int) -> Dict[str, np.ndarray]:
        idx_key = str(idx)
        with h5py.File(self.filepath, "r", libver="latest", swmr=True) as hf:
            out_dict = {
                "episode_idx": idx,
                "scenario_id": hf[idx_key].attrs["scenario_id"],
                "scenario_center": hf[idx_key].attrs["scenario_center"],
                "scenario_yaw": hf[idx_key].attrs["scenario_yaw"],
                "with_map": hf[idx_key].attrs["with_map"],  # some epidosdes in the testing dataset do not have map.
            }
            for k, _size in self.tensor_size.items():
                out_dict[k] = np.ascontiguousarray(hf[idx_key][k])
                if out_dict[k].shape != _size:
                    assert "agent" in k
                    out_dict[k] = np.ones(_size, dtype=out_dict[k].dtype)
        return out_dict


class DataH5womd(LightningDataModule):
    def __init__(
        self,
        data_dir: str,
        filename_train: str = "training",
        filename_val: str = "validation",
        filename_test: str = "testing",
        batch_size: int = 3,
        num_workers: int = 4,
        n_agent: int = 64,  # if not the same as h5 dataset, use dummy agents, for scalability tests.
    ) -> None:
        super().__init__()
        self.interactive_challenge = "interactive" in filename_val or "interactive" in filename_test

        self.path_train_h5 = f"{data_dir}/{filename_train}.h5"
        self.path_val_h5 = f"{data_dir}/{filename_val}.h5"
        self.path_test_h5 = f"{data_dir}/{filename_test}.h5"
        self.batch_size = batch_size
        self.num_workers = num_workers

        n_step = 91
        n_step_history = 11
        n_agent_no_sim = 256
        n_pl = 1024
        n_tl = 100
        n_tl_stop = 40
        n_pl_node = 20
        self.tensor_size_train = {
            # agent states
            "agent/valid": (n_step, n_agent),  # bool,
            "agent/pos": (n_step, n_agent, 2),  # float32
            "agent/z": (n_step, n_agent, 1),  # float32
            # v[1] = p[1]-p[0]. if p[1] invalid, v[1] also invalid, v[2]=v[3]
            "agent/vel": (n_step, n_agent, 2),  # float32, v_x, v_y
            "agent/spd": (n_step, n_agent, 1),  # norm of vel, signed using yaw_bbox and vel_xy
            "agent/acc": (n_step, n_agent, 1),  # m/s2, acc[t] = (spd[t]-spd[t-1])/dt
            "agent/yaw_bbox": (n_step, n_agent, 1),  # float32, yaw of the bbox heading
            "agent/yaw_rate": (n_step, n_agent, 1),  # rad/s, yaw_rate[t] = (yaw[t]-yaw[t-1])/dt
            # agent attributes
            "agent/type": (n_agent, 3),  # bool one_hot [Vehicle=0, Pedestrian=1, Cyclist=2]
            "agent/cmd": (n_agent, 8),  # bool one_hot
            "agent/role": (n_agent, 3),  # bool [sdc=0, interest=1, predict=2]
            "agent/size": (n_agent, 3),  # float32: [length, width, height]
            "agent/goal": (n_agent, 4),  # float32: [x, y, theta, spd]
            "agent/dest": (n_agent,),  # int64: index to map n_pl
            # map polylines
            "map/valid": (n_pl, n_pl_node),  # bool
            "map/type": (n_pl, 11),  # bool one_hot
            "map/pos": (n_pl, n_pl_node, 2),  # float32
            "map/dir": (n_pl, n_pl_node, 2),  # float32
            "map/boundary": (4,),  # xmin, xmax, ymin, ymax
            # traffic lights
            "tl_lane/valid": (n_step, n_tl),  # bool
            "tl_lane/state": (n_step, n_tl, 5),  # bool one_hot
            "tl_lane/idx": (n_step, n_tl),  # int, -1 means not valid
            "tl_stop/valid": (n_step, n_tl_stop),  # bool
            "tl_stop/state": (n_step, n_tl_stop, 5),  # bool one_hot
            "tl_stop/pos": (n_step, n_tl_stop, 2),  # x,y
            "tl_stop/dir": (n_step, n_tl_stop, 2),  # x,y
        }

        self.tensor_size_test = {
            # object_id for waymo metrics
            "history/agent/object_id": (n_agent,),
            "history/agent_no_sim/object_id": (n_agent_no_sim,),
            # agent_sim
            "history/agent/valid": (n_step_history, n_agent),  # bool,
            "history/agent/pos": (n_step_history, n_agent, 2),  # float32
            "history/agent/z": (n_step_history, n_agent, 1),  # float32
            "history/agent/vel": (n_step_history, n_agent, 2),  # float32, v_x, v_y
            "history/agent/spd": (n_step_history, n_agent, 1),  # norm of vel, signed using yaw_bbox and vel_xy
            "history/agent/acc": (n_step_history, n_agent, 1),  # m/s2, acc[t] = (spd[t]-spd[t-1])/dt
            "history/agent/yaw_bbox": (n_step_history, n_agent, 1),  # float32, yaw of the bbox heading
            "history/agent/yaw_rate": (n_step_history, n_agent, 1),  # rad/s, yaw_rate[t] = (yaw[t]-yaw[t-1])/dt
            "history/agent/type": (n_agent, 3),  # bool one_hot [Vehicle=0, Pedestrian=1, Cyclist=2]
            "history/agent/role": (n_agent, 3),  # bool [sdc=0, interest=1, predict=2]
            "history/agent/size": (n_agent, 3),  # float32: [length, width, height]
            "history/agent_no_sim/valid": (n_step_history, n_agent_no_sim),
            "history/agent_no_sim/pos": (n_step_history, n_agent_no_sim, 2),
            "history/agent_no_sim/z": (n_step_history, n_agent_no_sim, 1),
            "history/agent_no_sim/vel": (n_step_history, n_agent_no_sim, 2),
            "history/agent_no_sim/spd": (n_step_history, n_agent_no_sim, 1),
            "history/agent_no_sim/yaw_bbox": (n_step_history, n_agent_no_sim, 1),
            "history/agent_no_sim/type": (n_agent_no_sim, 3),
            "history/agent_no_sim/size": (n_agent_no_sim, 3),
            # map
            "map/valid": (n_pl, n_pl_node),  # bool
            "map/type": (n_pl, 11),  # bool one_hot
            "map/pos": (n_pl, n_pl_node, 2),  # float32
            "map/dir": (n_pl, n_pl_node, 2),  # float32
            "map/boundary": (4,),  # xmin, xmax, ymin, ymax
            # traffic_light
            "history/tl_lane/valid": (n_step_history, n_tl),  # bool
            "history/tl_lane/state": (n_step_history, n_tl, 5),  # bool one_hot
            "history/tl_lane/idx": (n_step_history, n_tl),  # int, -1 means not valid
            "history/tl_stop/valid": (n_step_history, n_tl_stop),  # bool
            "history/tl_stop/state": (n_step_history, n_tl_stop, 5),  # bool one_hot
            "history/tl_stop/pos": (n_step_history, n_tl_stop, 2),  # x,y
            "history/tl_stop/dir": (n_step_history, n_tl_stop, 2),  # dx,dy
        }

        self.tensor_size_val = {
            "agent/object_id": (n_agent,),
            "agent_no_sim/object_id": (n_agent_no_sim,),
            # agent_no_sim
            "agent_no_sim/valid": (n_step, n_agent_no_sim),  # bool,
            "agent_no_sim/pos": (n_step, n_agent_no_sim, 2),  # float32
            "agent_no_sim/z": (n_step, n_agent_no_sim, 1),  # float32
            "agent_no_sim/vel": (n_step, n_agent_no_sim, 2),  # float32, v_x, v_y
            "agent_no_sim/spd": (n_step, n_agent_no_sim, 1),  # norm of vel, signed using yaw_bbox and vel_xy
            "agent_no_sim/yaw_bbox": (n_step, n_agent_no_sim, 1),  # float32, yaw of the bbox heading
            "agent_no_sim/type": (n_agent_no_sim, 3),  # bool one_hot [Vehicle=0, Pedestrian=1, Cyclist=2]
            "agent_no_sim/size": (n_agent_no_sim, 3),  # float32: [length, width, height]
        }

        self.tensor_size_val = self.tensor_size_val | self.tensor_size_train | self.tensor_size_test

    def setup(self, stage: Optional[str] = None) -> None:
        if stage == "fit" or stage is None:
            self.train_dataset = DatasetTrain(self.path_train_h5, self.tensor_size_train)
            self.val_dataset = DatasetVal(self.path_val_h5, self.tensor_size_val)
        elif stage == "validate":
            self.val_dataset = DatasetVal(self.path_val_h5, self.tensor_size_val)
        elif stage == "test":
            self.test_dataset = DatasetVal(self.path_test_h5, self.tensor_size_test)

    def train_dataloader(self) -> DataLoader[Any]:
        return self._get_dataloader(self.train_dataset, self.batch_size, self.num_workers)

    def val_dataloader(self) -> DataLoader[Any]:
        return self._get_dataloader(self.val_dataset, self.batch_size, self.num_workers)

    def test_dataloader(self) -> DataLoader[Any]:
        return self._get_dataloader(self.test_dataset, self.batch_size, self.num_workers)

    @staticmethod
    def _get_dataloader(ds: Dataset, batch_size: int, num_workers: int) -> DataLoader[Any]:
        return DataLoader(
            ds,
            batch_size=batch_size,
            num_workers=num_workers,
            pin_memory=True,
            shuffle=False,
            drop_last=False,
            persistent_workers=True,
        )

src/data_modules/data_h5_womd.py
zhejz/TrafficBots
# Licensed under the CC BY-NC 4.0 license (https://creativecommons.org/licenses/by-nc/4.0/)
from typing import Dict, Tuple
from omegaconf import DictConfig
import torch
from torch import nn, Tensor
from utils.pose_pe import PosePE


class SceneCentricInput(nn.Module):
    def __init__(
        self, time_step_current: int, data_size: DictConfig, dropout_p_history: float, pe_dim: int, pose_pe: DictConfig
    ) -> None:
        super().__init__()
        self.dropout_p_history = dropout_p_history  # [0, 1], turn off if set to negative
        self.n_step_hist = time_step_current + 1

        self.pose_pe_agent = PosePE(pose_pe["agent"], pe_dim=pe_dim)
        self.pose_pe_map = PosePE(pose_pe["map"], pe_dim=pe_dim)
        self.pose_pe_tl = PosePE(pose_pe["tl"], pe_dim=pe_dim)

        agent_attr_dim = (
            data_size["agent/vel"][-1]  # 2
            + data_size["agent/spd"][-1]  # 1
            + data_size["agent/yaw_rate"][-1]  # 1
            + data_size["agent/acc"][-1]  # 1
            + data_size["agent/size"][-1]  # 3
            + data_size["agent/type"][-1]  # 3
        )
        agent_pe_dim = self.pose_pe_agent.out_dim

        n_pl_node = data_size["map/valid"][-1]
        map_attr_dim = data_size["map/type"][-1] + n_pl_node
        map_pe_dim = self.pose_pe_map.out_dim
        self.register_buffer("pl_node_ohe", torch.eye(n_pl_node))

        tl_attr_dim = data_size["tl_stop/state"][-1]
        tl_pe_dim = self.pose_pe_tl.out_dim

        self.model_kwargs = {
            "agent_attr_dim": agent_attr_dim,
            "agent_pe_dim": agent_pe_dim,
            "map_attr_dim": map_attr_dim,
            "map_pe_dim": map_pe_dim,
            "tl_attr_dim": tl_attr_dim,
            "tl_pe_dim": tl_pe_dim,
            "n_step_hist": self.n_step_hist,
            "n_pl_node": n_pl_node,
        }

    def forward(self, batch: Dict[str, Tensor]) -> Dict[str, Tensor]:
        """
        Args: scene-centric Dict
            # (sc) scene-centric agents states
                "sc/agent_valid": [n_scene, n_step_hist, n_agent]
                "sc/agent_pos": [n_scene, n_step_hist, n_agent, 2]
                "sc/agent_z": [n_scene, n_step_hist, n_agent, 1]
                "sc/agent_vel": [n_scene, n_step_hist, n_agent, 2]
                "sc/agent_spd": [n_scene, n_step_hist, n_agent, 1]
                "sc/agent_acc": [n_scene, n_step_hist, n_agent, 1]
                "sc/agent_yaw_bbox": [n_scene, n_step_hist, n_agent, 1]
                "sc/agent_yaw_rate": [n_scene, n_step_hist, n_agent, 1]
            # agents attributes
                "sc/agent_type": [n_scene, n_agent, 3]
                "sc/agent_role": [n_scene, n_agent, 3]
                "sc/agent_size": [n_scene, n_agent, 3]
            # map polylines
                "sc/map_valid": [n_scene, n_pl, n_pl_node], bool
                "sc/map_type": [n_scene, n_pl, 11], bool one_hot
                "sc/map_pos": [n_scene, n_pl, n_pl_node, 2], float32
                "sc/map_dir": [n_scene, n_pl, n_pl_node, 2], float32
            # traffic lights
                "sc/tl_valid": [n_scene, n_step_hist, n_tl], bool
                "sc/tl_state": [n_scene, n_step_hist, n_tl, 5], bool one_hot
                "sc/tl_pos": [n_scene, n_step_hist, n_tl, 2], x,y
                "sc/tl_dir": [n_scene, n_step_hist, n_tl, 2], x,y

        Returns: add following keys to batch Dict
            # rollout
                # agent history
                "input/agent_valid": [n_scene, n_step_hist, n_agent], bool
                "input/agent_attr": [n_scene, n_step_hist, n_agent, agent_attr_dim], for input to MLP
                "input/agent_pe": [n_scene, n_step_hist, n_agent, hidden_dim], for input/cat/add to MLP
                "input/agent_pos": [n_scene, n_step_hist, n_agent, 2], for knn attention
                # map polylines
                "input/map_valid": [n_scene, n_pl, n_pl_node], bool
                "input/map_attr": [n_scene, n_pl, n_pl_node, map_attr_dim]
                "input/map_pe": [n_scene, n_pl, n_pl_node, hidden_dim], for input/cat/add to MLP
                "input/map_pos": [n_scene, n_pl, 2], for knn attention
                # traffic lights: stop point, detections are not tracked, singular node polyline.
                "input/tl_valid": [n_scene, n_step_hist, n_tl], bool
                "input/tl_attr": [n_scene, n_step_hist, n_tl, tl_attr_dim], for input to MLP
                "input/tl_pe": [n_scene, n_step_hist, n_tl, hidden_dim], for input/cat/add to MLP
                "input/tl_pos": [n_scene, n_step_hist, n_tl, 2], for knn attention
        """
        batch["input/agent_valid"] = batch["sc/agent_valid"]
        batch["input/tl_valid"] = batch["sc/tl_valid"]
        batch["input/map_valid"] = batch["sc/map_valid"]

        # ! randomly mask history agent/tl/map
        if self.training and (0 < self.dropout_p_history <= 1.0):
            prob_mask = torch.ones_like(batch["input/agent_valid"][:, :-1, :]) * (1 - self.dropout_p_history)
            batch["input/agent_valid"][:, :-1, :] &= torch.bernoulli(prob_mask).bool()
            prob_mask = torch.ones_like(batch["input/tl_valid"]) * (1 - self.dropout_p_history)
            batch["input/tl_valid"] &= torch.bernoulli(prob_mask).bool()
            prob_mask = torch.ones_like(batch["input/map_valid"]) * (1 - self.dropout_p_history)
            batch["input/map_valid"] &= torch.bernoulli(prob_mask).bool()

        # ! prepare "input/agent"
        batch["input/agent_pos"] = batch["sc/agent_pos"]
        batch["input/agent_attr"] = torch.cat(
            [
                batch["sc/agent_vel"],  # vel xy, 2
                batch["sc/agent_spd"],  # speed, 1
                batch["sc/agent_yaw_rate"],  # yaw rate, 1
                batch["sc/agent_acc"],  # acc, 1
                batch["sc/agent_size"].unsqueeze(1).expand(-1, self.n_step_hist, -1, -1),  # 3
                batch["sc/agent_type"].unsqueeze(1).expand(-1, self.n_step_hist, -1, -1),  # 3
            ],
            dim=-1,
        )
        # "input/agent_pe": [n_scene, n_agent, n_step_hist, hidden_dim], for input/cat/add to MLP
        batch["input/agent_pe"] = self.pose_pe_agent(batch["sc/agent_pos"], batch["sc/agent_yaw_bbox"])

        # ! prepare "input/map_attr": [n_scene, n_pl, n_pl_node, map_attr_dim]
        n_scene, n_pl, n_pl_node = batch["sc/map_valid"].shape
        batch["input/map_pos"] = batch["sc/map_pos"][:, :, 0].contiguous()
        batch["input/map_attr"] = torch.cat(
            [
                batch["sc/map_type"].unsqueeze(-2).expand(-1, -1, n_pl_node, -1),
                self.pl_node_ohe[None, None, :, :].expand(n_scene, n_pl, -1, -1),
            ],
            dim=-1,
        )
        batch["input/map_pe"] = self.pose_pe_map(batch["sc/map_pos"], batch["sc/map_dir"])

        # ! prepare "input/tl_attr": [n_scene, n_step_hist, n_tl, tl_attr_dim]
        batch["input/tl_pos"] = batch["sc/tl_pos"]
        batch["input/tl_attr"] = batch["sc/tl_state"].to(batch["input/tl_pos"].dtype)
        batch["input/tl_pe"] = self.pose_pe_tl(batch["sc/tl_pos"], batch["sc/tl_dir"])
        return batch

    def get_agent_attr_and_pe(
        self,
        agent_pos: Tensor,
        agent_yaw_bbox: Tensor,
        agent_vel: Tensor,
        agent_spd: Tensor,
        agent_yaw_rate: Tensor,
        agent_acc: Tensor,
        agent_size: Tensor,
        agent_type: Tensor,
    ) -> Tuple[Tensor, Tensor]:
        agent_attr = torch.cat(
            [
                agent_vel,  # vel xy, 2
                agent_spd,  # speed, 1
                agent_yaw_rate,  # yaw rate, 1
                agent_acc,  # acc, 1
                agent_size,  # 3
                agent_type,  # 3
            ],
            dim=-1,
        )  # [n_scene, n_agent, agent_attr_dim]
        agent_pe = self.pose_pe_agent(agent_pos, agent_yaw_bbox)  # for input/cat/add to MLP
        return agent_attr, agent_pe

src/data_modules/sc_input.py
zhejz/TrafficBots
# Licensed under the CC BY-NC 4.0 license (https://creativecommons.org/licenses/by-nc/4.0/)
from typing import Dict
from omegaconf import DictConfig
import torch
from torch import nn, Tensor
from utils.pose_pe import PosePE
from utils.transform_utils import torch_rad2rot, torch_pos2local, torch_dir2local, torch_rad2local


class SceneCentricLatent(nn.Module):
    def __init__(
        self,
        time_step_current: int,
        data_size: DictConfig,
        perturb_input_to_latent: bool,  # random transform input to latent encoder at training time
        dropout_p_history: float,
        pe_dim: int,
        pose_pe: DictConfig,
        max_meter: float = 50.0,
        max_rad: float = 3.14,
    ) -> None:
        super().__init__()
        self.perturb_input_to_latent = perturb_input_to_latent
        self.dropout_p_history = dropout_p_history  # [0, 1], turn off if set to negative
        self.n_step_hist = time_step_current + 1
        self.max_meter = max_meter
        self.max_rad = max_rad

        self.pose_pe_agent = PosePE(pose_pe["agent"], pe_dim=pe_dim)
        self.pose_pe_map = PosePE(pose_pe["map"], pe_dim=pe_dim)
        self.pose_pe_tl = PosePE(pose_pe["tl"], pe_dim=pe_dim)

        self.register_buffer("pl_node_ohe", torch.eye(data_size["map/valid"][-1]))

        self.model_kwargs = {}

    def forward(self, batch: Dict[str, Tensor]) -> Dict[str, Tensor]:
        """
        Args: scene-centric Dict
            # (sc) scene-centric agents states
                "sc/agent_valid": [n_scene, n_step_hist, n_agent]
                "sc/agent_pos": [n_scene, n_step_hist, n_agent, 2]
                "sc/agent_z": [n_scene, n_step_hist, n_agent, 1]
                "sc/agent_vel": [n_scene, n_step_hist, n_agent, 2]
                "sc/agent_spd": [n_scene, n_step_hist, n_agent, 1]
                "sc/agent_acc": [n_scene, n_step_hist, n_agent, 1]
                "sc/agent_yaw_bbox": [n_scene, n_step_hist, n_agent, 1]
                "sc/agent_yaw_rate": [n_scene, n_step_hist, n_agent, 1]
            # agents attributes
                "sc/agent_type": [n_scene, n_agent, 3]
                "sc/agent_role": [n_scene, n_agent, 3]
                "sc/agent_size": [n_scene, n_agent, 3]
            # map polylines
                "sc/map_valid": [n_scene, n_pl, n_pl_node], bool
                "sc/map_type": [n_scene, n_pl, 11], bool one_hot
                "sc/map_pos": [n_scene, n_pl, n_pl_node, 2], float32
                "sc/map_dir": [n_scene, n_pl, n_pl_node, 2], float32
            # traffic lights
                "sc/tl_valid": [n_scene, n_step_hist, n_tl], bool
                "sc/tl_state": [n_scene, n_step_hist, n_tl, 5], bool one_hot
                "sc/tl_pos": [n_scene, n_step_hist, n_tl, 2], x,y
                "sc/tl_dir": [n_scene, n_step_hist, n_tl, 2], x,y
            # input dict
                # agent history
                "input/agent_valid": [n_scene, n_step_hist, n_agent], bool
                "input/agent_attr": [n_scene, n_step_hist, n_agent, agent_attr_dim], for input to MLP
                "input/agent_pe": [n_scene, n_step_hist, n_agent, hidden_dim], for input/cat/add to MLP
                "input/agent_pos": [n_scene, n_step_hist, n_agent, 2], for knn attention
                # map polylines
                "input/map_valid": [n_scene, n_pl, n_pl_node], bool
                "input/map_attr": [n_scene, n_pl, n_pl_node, map_attr_dim]
                "input/map_pe": [n_scene, n_pl, n_pl_node, hidden_dim], for input/cat/add to MLP
                "input/map_pos": [n_scene, n_pl, 2], for knn attention
                # traffic lights: stop point, detections are not tracked, singular node polyline.
                "input/tl_valid": [n_scene, n_step_hist, n_tl], bool
                "input/tl_attr": [n_scene, n_step_hist, n_tl, tl_attr_dim], for input to MLP
                "input/tl_pe": [n_scene, n_step_hist, n_tl, hidden_dim], for input/cat/add to MLP
                "input/tl_pos": [n_scene, n_step_hist, n_tl, 2], for knn attention

        Returns: add following keys to batch Dict
            # latent_prior from "sc/" or "input/"
                # agent
                "latent_prior/agent_valid": [n_scene, n_agent, n_step_hist]
                "latent_prior/agent_attr": [n_scene, n_agent, n_step_hist, agent_attr_dim]
                "latent_prior/agent_pe": [n_scene, n_agent, n_step_hist, hidden_dim]
                "latent_prior/agent_pos": [n_scene, n_agent, n_step_hist, 2]
                # map polylines
                "latent_prior/map_valid": [n_scene, n_pl, n_pl_node]
                "latent_prior/map_attr": [n_scene, n_pl, n_pl_node, map_attr_dim]
                "latent_prior/map_pe": [n_scene, n_pl, n_pl_node, hidden_dim]
                "latent_prior/map_pos": [n_scene, n_pl, 2]
                # traffic lights:
                "latent_prior/tl_valid": [n_scene, n_step_hist, n_tl]
                "latent_prior/tl_attr": [n_scene, n_step_hist, n_tl, tl_attr_dim]
                "latent_prior/tl_pe": [n_scene, n_step_hist, n_tl, hidden_dim]
                "latent_prior/tl_pos": [n_scene, n_step_hist, n_tl, 2]
            # latent_post from original batch gt
                # agent
                "latent_post/agent_valid": [n_scene, n_agent, n_step]
                "latent_post/agent_attr": [n_scene, n_agent, n_step, agent_attr_dim]
                "latent_post/agent_pe": [n_scene, n_agent, n_step, hidden_dim]
                "latent_post/agent_pos": [n_scene, n_agent, n_step, 2]
                # map polylines
                "latent_post/map_valid": [n_scene, n_pl, n_pl_node]
                "latent_post/map_attr": [n_scene, n_pl, n_pl_node, map_attr_dim]
                "latent_post/map_pe": [n_scene, n_pl, n_pl_node, hidden_dim]
                "latent_post/map_pos": [n_scene, n_pl, 2]
                # traffic lights:
                "latent_post/tl_valid": [n_scene, n_step, n_tl]
                "latent_post/tl_attr": [n_scene, n_step, n_tl, tl_attr_dim]
                "latent_post/tl_pe": [n_scene, n_step, n_tl, hidden_dim]
                "latent_post/tl_pos": [n_scene, n_step, n_tl, 2]
        """
        gt_available = "agent/valid" in batch.keys()  # training/validation time, prepare latent_posterior
        if self.training and self.perturb_input_to_latent:
            # "sc/agent_pos": [n_scene, n_agent, n_step_hist, 2]
            n_scene = batch["sc/agent_pos"].shape[0]
            device = batch["sc/agent_pos"].device
            dtype = batch["sc/agent_pos"].dtype
            rand_yaw = torch.rand([n_scene], device=device, dtype=dtype) * 2 * self.max_rad - self.max_rad
            rand_rot = torch_rad2rot(rand_yaw)
            rand_pos = torch.rand([n_scene, 2], device=device, dtype=dtype) * 2 * self.max_meter - self.max_meter
            rand_yaw = rand_yaw[:, None, None]  # [n_scene, 1, 1]
            rand_rot = rand_rot.unsqueeze(1)  # [n_scene, 1, 2, 2]
            rand_pos = rand_pos[:, None, None, :]  # [n_scene, 1, 1, 2]

        # ! map
        if self.training and self.perturb_input_to_latent:
            n_scene, n_pl, n_pl_node = batch["sc/map_valid"].shape
            map_pos = torch_pos2local(batch["sc/map_pos"], rand_pos, rand_rot)
            map_dir = torch_dir2local(batch["sc/map_dir"], rand_rot)
            batch["latent_prior/map_pos"] = map_pos[:, :, 0].contiguous()
            batch["latent_prior/map_attr"] = torch.cat(
                [
                    batch["sc/map_type"].unsqueeze(-2).expand(-1, -1, n_pl_node, -1),
                    self.pl_node_ohe[None, None, :, :].expand(n_scene, n_pl, -1, -1),
                ],
                dim=-1,
            )
            batch["latent_prior/map_pe"] = self.pose_pe_map(map_pos, map_dir)
            batch["latent_prior/map_valid"] = batch["sc/map_valid"]
        else:
            for k in ["valid", "pos", "attr", "pe"]:
                batch[f"latent_prior/map_{k}"] = batch[f"input/map_{k}"]

        if gt_available:
            for k in ["valid", "pos", "attr", "pe"]:
                batch[f"latent_post/map_{k}"] = batch[f"latent_prior/map_{k}"]

        # ! traffic lights
        # prior
        if self.training and self.perturb_input_to_latent:
            batch["latent_prior/tl_valid"] = batch["sc/tl_valid"]
            if self.training and (0 < self.dropout_p_history <= 1.0):
                prob_mask = torch.ones_like(batch["latent_prior/tl_valid"]) * (1 - self.dropout_p_history)
                batch["latent_prior/tl_valid"] &= torch.bernoulli(prob_mask).bool()
            tl_pos = batch["sc/tl_pos"]
            tl_dir = batch["sc/tl_dir"]
            if self.training and self.perturb_input_to_latent:
                tl_pos = torch_pos2local(tl_pos, rand_pos, rand_rot)
                tl_dir = torch_dir2local(tl_dir, rand_rot)
            batch["latent_prior/tl_pos"] = tl_pos
            batch["latent_prior/tl_attr"] = batch["sc/tl_state"].to(batch["latent_prior/tl_pos"].dtype)
            batch["latent_prior/tl_pe"] = self.pose_pe_tl(tl_pos, tl_dir)
        else:
            for k in ["valid", "pos", "attr", "pe"]:
                batch[f"latent_prior/tl_{k}"] = batch[f"input/tl_{k}"]
        # post
        if gt_available:  # train, val
            batch["latent_post/tl_valid"] = batch["tl_stop/valid"]
            if self.training and (0 < self.dropout_p_history <= 1.0):
                prob_mask = torch.ones_like(batch["latent_post/tl_valid"]) * (1 - self.dropout_p_history)
                batch["latent_post/tl_valid"] &= torch.bernoulli(prob_mask).bool()
            tl_pos = batch["tl_stop/pos"]
            tl_dir = batch["tl_stop/dir"]
            if self.training and self.perturb_input_to_latent:
                tl_pos = torch_pos2local(tl_pos, rand_pos, rand_rot)
                tl_dir = torch_dir2local(tl_dir, rand_rot)
            batch["latent_post/tl_pos"] = tl_pos
            batch["latent_post/tl_attr"] = batch["tl_stop/state"].to(batch["latent_post/tl_pos"].dtype)
            batch["latent_post/tl_pe"] = self.pose_pe_tl(tl_pos, tl_dir)

        # ! agents
        if self.training and self.perturb_input_to_latent:
            batch["latent_prior/agent_valid"] = batch["sc/agent_valid"]
            if self.training and (0 < self.dropout_p_history <= 1.0):
                prob_mask = torch.ones_like(batch["latent_prior/agent_valid"][:, :-1, :]) * (1 - self.dropout_p_history)
                batch["latent_prior/agent_valid"][:, :-1, :] &= torch.bernoulli(prob_mask).bool()
            agent_pos = batch["sc/agent_pos"]
            agent_vel = batch["sc/agent_vel"]
            agent_yaw_bbox = batch["sc/agent_yaw_bbox"]
            if self.training and self.perturb_input_to_latent:
                agent_pos = torch_pos2local(agent_pos, rand_pos, rand_rot)
                agent_vel = torch_dir2local(agent_vel, rand_rot)
                agent_yaw_bbox = torch_rad2local(agent_yaw_bbox, rand_yaw, cast=False)
            batch["latent_prior/agent_pos"] = agent_pos
            n_step = batch["latent_prior/agent_valid"].shape[1]
            batch["latent_prior/agent_attr"] = torch.cat(
                [
                    agent_vel,  # vel xy, 2
                    batch["sc/agent_spd"],  # speed, 1
                    batch["sc/agent_yaw_rate"],  # yaw rate, 1
                    batch["sc/agent_acc"],  # acc, 1
                    batch["sc/agent_size"].unsqueeze(1).expand(-1, n_step, -1, -1),  # 3
                    batch["sc/agent_type"].unsqueeze(1).expand(-1, n_step, -1, -1),  # 3
                ],
                dim=-1,
            )
            batch["latent_prior/agent_pe"] = self.pose_pe_agent(agent_pos, agent_yaw_bbox)
        else:
            for k in ["valid", "pos", "attr", "pe"]:
                batch[f"latent_prior/agent_{k}"] = batch[f"input/agent_{k}"]

        if gt_available:  # train, val
            batch["latent_post/agent_valid"] = batch["agent/valid"]
            if self.training and (0 < self.dropout_p_history <= 1.0):
                prob_mask = torch.ones_like(batch["latent_post/agent_valid"]) * (1 - self.dropout_p_history)
                batch["latent_post/agent_valid"] &= torch.bernoulli(prob_mask).bool()
            agent_pos = batch["agent/pos"]
            agent_vel = batch["agent/vel"]
            agent_yaw_bbox = batch["agent/yaw_bbox"]
            if self.training and self.perturb_input_to_latent:
                agent_pos = torch_pos2local(agent_pos, rand_pos, rand_rot)
                agent_vel = torch_dir2local(agent_vel, rand_rot)
                agent_yaw_bbox = torch_rad2local(agent_yaw_bbox, rand_yaw, cast=False)
            batch["latent_post/agent_pos"] = agent_pos
            n_step = batch["latent_post/agent_valid"].shape[1]
            batch["latent_post/agent_attr"] = torch.cat(
                [
                    agent_vel,  # vel xy, 2
                    batch["agent/spd"],  # speed, 1
                    batch["agent/yaw_rate"],  # yaw rate, 1
                    batch["agent/acc"],  # acc, 1
                    batch["agent/size"].unsqueeze(1).expand(-1, n_step, -1, -1),  # 3
                    batch["agent/type"].unsqueeze(1).expand(-1, n_step, -1, -1),  # 3
                ],
                dim=-1,
            )
            batch["latent_post/agent_pe"] = self.pose_pe_agent(agent_pos, agent_yaw_bbox)

        return batch

src/data_modules/sc_latent.py
zhejz/TrafficBots
# Licensed under the CC BY-NC 4.0 license (https://creativecommons.org/licenses/by-nc/4.0/)
from typing import Dict
from omegaconf import DictConfig
from torch import nn, Tensor
import torch


class SceneCentricPreProcessing(nn.Module):
    def __init__(self, time_step_current: int, data_size: DictConfig) -> None:
        super().__init__()
        self.n_step_hist = time_step_current + 1
        self.model_kwargs = {}

    def forward(self, batch: Dict[str, Tensor]) -> Dict[str, Tensor]:
        """
        Args: scene-centric Dict
            # agent states
                "agent/valid": [n_scene, n_step, n_agent], bool,
                "agent/pos": [n_scene, n_step, n_agent, 2], float32
                "agent/z": [n_scene, n_step, n_agent, 1], float32
                "agent/vel": [n_scene, n_step, n_agent, 2], float32, v_x, v_y
                "agent/spd": [n_scene, n_step, n_agent, 1], norm of vel, signed using yaw_bbox and vel_xy
                "agent/acc": [n_scene, n_step, n_agent, 1], m/s2, acc[t] = (spd[t]-spd[t-1])/dt
                "agent/yaw_bbox": [n_scene, n_step, n_agent, 1], float32, yaw of the bbox heading
                "agent/yaw_rate": [n_scene, n_step, n_agent, 1], rad/s, yaw_rate[t] = (yaw[t]-yaw[t-1])/dt
            # agent attributes
                "agent/type": [n_scene, n_agent, 3], bool one_hot [Vehicle=0, Pedestrian=1, Cyclist=2]
                "agent/role": [n_scene, n_agent, 3], bool [sdc=0, interest=1, predict=2]
                "agent/size": [n_scene, n_agent, 3], float32: [length, width, height]
            # map polylines
                "map/valid": [n_scene, n_pl, n_pl_node], bool
                "map/type": [n_scene, n_pl, 11], bool one_hot
                "map/pos": [n_scene, n_pl, n_pl_node, 2], float32
                "map/dir": [n_scene, n_pl, n_pl_node, 2], float32
            # traffic lights
                "tl_stop/valid": [n_scene, n_step, n_tl_stop], bool
                "tl_stop/state": [n_scene, n_step, n_tl_stop, 5], bool one_hot
                "tl_stop/pos": [n_scene, n_step, n_tl_stop, 2], x,y
                "tl_stop/dir": [n_scene, n_step, n_tl_stop, 2], x,y

        Returns: scene-centric Dict, masked according to valid
            # (ref) reference information
                "ref/agent_type": [n_scene, n_agent, 3], bool one_hot [Vehicle=0, Pedestrian=1, Cyclist=2]
                "ref/agent_role": [n_scene, n_agent, 3], bool [sdc=0, interest=1, predict=2]
                "ref/agent_state": [n_scene, n_step_hist, n_agent, 4], (x,y,yaw,spd)
                "ref/map_type": [n_scene, n_pl, 11]
            # (gt) ground-truth agent future for training, not available for testing
                "gt/valid": [n_scene, n_step, n_agent], bool
                "gt/spd": [n_scene, n_step, n_agent, 1]
                "gt/pos": [n_scene, n_step, n_agent, 2]
                "gt/vel": [n_scene, n_step, n_agent, 2]
                "gt/yaw_bbox": [n_scene, n_step, n_agent, 1]
                "gt/state": [n_scene, n_step, n_agent, 4], (x,y,yaw,spd)
                "gt/cmd": [n_scene, n_agent, 8]
                "gt/goal": [n_scene, n_agent, 4], (x, y, theta, spd)
                "gt/dest": [n_scene, n_agent], int64: index to map n_pl
            # (sc) scene-centric agents states
                "sc/agent_valid": [n_scene, n_step_hist, n_agent]
                "sc/agent_pos": [n_scene, n_step_hist, n_agent, 2]
                "sc/agent_z": [n_scene, n_step_hist, n_agent, 1]
                "sc/agent_vel": [n_scene, n_step_hist, n_agent, 2]
                "sc/agent_spd": [n_scene, n_step_hist, n_agent, 1]
                "sc/agent_acc": [n_scene, n_step_hist, n_agent, 1]
                "sc/agent_yaw_bbox": [n_scene, n_step_hist, n_agent, 1]
                "sc/agent_yaw_rate": [n_scene, n_step_hist, n_agent, 1]
            # agents attributes
                "sc/agent_type": [n_scene, n_agent, 3]
                "sc/agent_role": [n_scene, n_agent, 3]
                "sc/agent_size": [n_scene, n_agent, 3]
            # map polylines
                "sc/map_valid": [n_scene, n_pl, n_pl_node], bool
                "sc/map_type": [n_scene, n_pl, 11], bool one_hot
                "sc/map_pos": [n_scene, n_pl, n_pl_node, 2], float32
                "sc/map_dir": [n_scene, n_pl, n_pl_node, 2], float32
            # traffic lights
                "sc/tl_valid": [n_scene, n_step_hist, n_tl], bool
                "sc/tl_state": [n_scene, n_step_hist, n_tl, 5], bool one_hot
                "sc/tl_pos": [n_scene, n_step_hist, n_tl, 2], x,y
                "sc/tl_dir": [n_scene, n_step_hist, n_tl, 2], x,y
            # agent_no_sim for validation and testing
                "sc/agent_no_sim_valid": [n_scene, n_step_hist, n_agent_no_sim]
                "sc/agent_no_sim_pos": [n_scene, n_step_hist, n_agent_no_sim, 2]
                "sc/agent_no_sim_z": [n_scene, n_step_hist, n_agent_no_sim, 1]
                "sc/agent_no_sim_vel": [n_scene, n_step_hist, n_agent_no_sim, 2]
                "sc/agent_no_sim_spd": [n_scene, n_step_hist, n_agent_no_sim, 1]
                "sc/agent_no_sim_acc": [n_scene, n_step_hist, n_agent_no_sim, 1]
                "sc/agent_no_sim_yaw_bbox": [n_scene, n_step_hist, n_agent_no_sim, 1]
                "sc/agent_no_sim_yaw_rate": [n_scene, n_step_hist, n_agent_no_sim, 1]
                "sc/agent_no_sim_type": [n_scene, n_agent_no_sim, 3]
                "sc/agent_no_sim_role": [n_scene, n_agent_no_sim, 3]
                "sc/agent_no_sim_size": [n_scene, n_agent_no_sim, 3]
        """
        prefix = "" if self.training else "history/"

        # ! prepare agents states
        # [n_scene, n_step, n_agent, ...] -> [n_scene, n_agent, n_step_hist, ...]
        for k in ("valid", "pos", "z", "vel", "spd", "acc", "yaw_bbox", "yaw_rate"):
            batch[f"sc/agent_{k}"] = batch[f"{prefix}agent/{k}"][:, : self.n_step_hist].contiguous()

        # ! prepare agents attributes
        for k in ("type", "role", "size"):
            batch[f"sc/agent_{k}"] = batch[f"{prefix}agent/{k}"]

        # ! training/validation time, prepare "gt/" for losses
        if "agent/valid" in batch.keys():
            for k in ("cmd", "goal", "dest"):
                batch[f"gt/{k}"] = batch[f"agent/{k}"]
            for k in ("valid", "spd", "pos", "vel", "yaw_bbox"):
                batch[f"gt/{k}"] = batch[f"agent/{k}"]
            batch["gt/state"] = torch.cat([batch["gt/pos"], batch["gt/yaw_bbox"], batch["gt/spd"]], dim=-1)

        # ! prepare map polylines
        for k in ("valid", "type", "pos", "dir"):
            batch[f"sc/map_{k}"] = batch[f"map/{k}"]

        # ! prepare traffic lights
        for k in ("valid", "state", "pos", "dir"):
            batch[f"sc/tl_{k}"] = batch[f"{prefix}tl_stop/{k}"][:, : self.n_step_hist].contiguous()

        # ! prepare agent_no_sim
        if not self.training:
            for k in ("valid", "pos", "z", "vel", "spd", "yaw_bbox"):
                batch[f"sc/agent_no_sim_{k}"] = batch[f"history/agent_no_sim/{k}"][:, : self.n_step_hist].contiguous()
            for k in ("type", "size"):
                batch[f"sc/agent_no_sim_{k}"] = batch[f"history/agent_no_sim/{k}"]

        # ! prepare "ref/"
        batch["ref/agent_type"] = batch[prefix + "agent/type"]
        batch["ref/agent_role"] = batch[prefix + "agent/role"]
        batch["ref/map_type"] = batch["map/type"]
        batch["ref/agent_state"] = torch.cat(
            [batch["sc/agent_pos"], batch["sc/agent_yaw_bbox"], batch["sc/agent_spd"]], dim=-1
        )  # x,y,yaw,spd

        return batch

src/data_modules/scene_centric.py
zhejz/TrafficBots
# Licensed under the CC BY-NC 4.0 license (https://creativecommons.org/licenses/by-nc/4.0/)
from typing import Dict, Tuple, List
from omegaconf import ListConfig
import torch
from torch import nn, Tensor


class WaymoPostProcessing(nn.Module):
    def __init__(
        self,
        k_pred: int,
        score_temperature: float,
        mpa_nms_thresh: ListConfig,
        mtr_nms_thresh: ListConfig,
        aggr_thresh: ListConfig,
        n_iter_em: int,
        use_ade: bool,
    ) -> None:
        """
        Args:
            score_temperature: set to > 0 to turn on
            mpa_nms_thresh, mtr_nms_thresh, aggr_thresh: list in meters
        """
        super().__init__()
        self.k_pred = k_pred
        self.score_temperature = score_temperature
        self.mpa_nms_thresh = list(mpa_nms_thresh)
        self.mtr_nms_thresh = list(mtr_nms_thresh)
        self.aggr_thresh = list(aggr_thresh)
        self.n_iter_em = n_iter_em
        self.use_ade = use_ade

    def forward(self, valid: Tensor, scores: Tensor, trajs: Tensor, agent_type: Tensor) -> Dict[str, Tensor]:
        """
        Args:
            valid: [n_scene, n_agent]
            scores: [n_scene, n_agent, n_pred], not normalized!
            trajs: [n_scene, n_agent, n_pred, n_step_future, 2/4] (x,y, yaw...)
            agent_type: [n_scene, n_agent, 3]

        Returns: pred_dict
            "waymo_valid": [n_scene, n_step, n_agent]
            "waymo_trajs": [n_scene, n_step, n_agent, k_pred, 2]
            "waymo_scores": [n_scene, n_agent, k_pred], normalized prob.
            "waymo_yaw_bbox": [n_scene, n_step, n_agent, k_pred, 1] or None
            "waymo_spd": [n_scene, n_step, n_agent, k_pred, 1] or None
        """
        scores = scores / scores.sum(-1, keepdim=True)  # normalized to prob
        n_scene, n_agent, n_pred, n_step, d_traj = trajs.shape

        if n_pred > self.k_pred:
            if len(self.aggr_thresh) > 0:
                trajs, scores = self.traj_aggr(
                    trajs, scores, self.k_pred, self.aggr_thresh, self.n_iter_em, self.use_ade
                )
            elif len(self.mtr_nms_thresh) > 0:
                trajs, scores = self.mtr_nms(trajs, scores, self.k_pred, self.mtr_nms_thresh, self.use_ade, agent_type)
            else:
                trajs, scores = self.traj_topk(trajs, scores, self.k_pred)

        # ! manually scale scores if necessary: [n_scene, n_agent, n_pred]
        if len(self.mpa_nms_thresh) > 0:
            scores = self.mpa_nms(valid, trajs, scores, self.mpa_nms_thresh, self.use_ade, agent_type)
        if self.score_temperature > 0:
            scores = torch.softmax(torch.log(scores) / self.score_temperature, dim=-1)

        trajs = trajs.movedim(3, 1)  # [n_scene, n_step, n_agent, k_pred, :]
        pred_dict = {"waymo_trajs": trajs[..., :2]}
        if d_traj >= 3:
            pred_dict["waymo_yaw_bbox"] = trajs[..., 2:3]
        else:
            pred_dict["waymo_yaw_bbox"] = None
        if d_traj >= 4:
            pred_dict["waymo_spd"] = trajs[..., 3:4]
        else:
            pred_dict["waymo_spd"] = None
        pred_dict["waymo_scores"] = scores  # [n_scene, n_agent, k_pred]
        # [n_scene, n_step, n_agent]
        pred_dict["waymo_valid"] = valid.unsqueeze(1).expand(-1, n_step, -1)
        return pred_dict

    @staticmethod
    def mpa_nms(
        valid: Tensor, trajs: Tensor, scores: Tensor, type_thresh: List[float], use_ade: bool, agent_type: Tensor
    ) -> Tensor:
        """
        Args:
            valid: [n_scene, n_agent]
            trajs: [n_scene, n_agent, k_pred, n_step, 2]
            scores: [n_scene, n_agent, k_pred], normalized prob
            type_thresh: in meters, list, len=3 [veh, ped, cyc].
            agent_type: [n_scene, n_agent, 3], [veh, ped, cyc].

        Returns:
            scores: [n_scene, n_agent, k_pred], normalized prob
        """
        # ! type dependent thresh
        thresh = 0
        for i in range(len(type_thresh)):
            thresh += agent_type[:, :, i] * type_thresh[i]
        thresh = thresh[:, :, None, None]  # [n_scene, n_agent, 1, 1]
        # within_dist: [n_scene, n_agent, n_pred, n_pred]
        xy = trajs[..., :2]
        if use_ade:
            within_dist = (torch.norm(xy.unsqueeze(2) - xy.unsqueeze(3), dim=-1).mean(-1)) < thresh
        else:
            within_dist = torch.norm(xy[:, :, :, -1].unsqueeze(2) - xy[:, :, :, -1].unsqueeze(3), dim=-1) < thresh

        for i in range(valid.shape[0]):
            for j in range(valid.shape[1]):
                if valid[i, j]:
                    for k in scores[i, j].argsort(descending=True):
                        # [k_pred]
                        mask = within_dist[i, j, k] & (scores[i, j] > scores[i, j, k])
                        if mask.any():
                            scores[i, j, k] = 1e-3
        scores = scores / scores.sum(-1, keepdim=True)
        return scores

    @staticmethod
    def mtr_nms(
        trajs: Tensor, scores: Tensor, k_pred: int, type_thresh: float, use_ade: bool, agent_type: Tensor
    ) -> Tuple[Tensor, Tensor]:
        """
        Args:
            trajs: [n_scene, n_agent, n_pred, n_step, 2]
            scores: [n_scene, n_agent, n_pred], normalized prob
            k_pred: int
            type_thresh: in meters, list, len=3 [veh, ped, cyc].
            agent_type: [n_scene, n_agent, 3], [veh, ped, cyc].

        Returns:
            trajs_k: [n_scene, n_agent, k_pred, n_step, 2]
            scores_k: [n_scene, n_agent, k_pred], normalized prob
        """
        # ! type dependent thresh
        thresh = 0
        for i in range(len(type_thresh)):
            thresh += agent_type[:, :, i] * type_thresh[i]
        thresh = thresh[:, :, None, None]  # [n_scene, n_agent, 1, 1]
        # within_dist: [n_scene, n_agent, n_pred, n_pred]
        xy = trajs[..., :2]
        if use_ade:
            within_dist = (torch.norm(xy.unsqueeze(2) - xy.unsqueeze(3), dim=-1).mean(-1)) < thresh
        else:
            within_dist = torch.norm(xy[:, :, :, -1].unsqueeze(2) - xy[:, :, :, -1].unsqueeze(3), dim=-1) < thresh

        # ! compute mode_idx: [n_scene, n_agent, k_pred]
        scene_idx = torch.arange(scores.shape[0]).unsqueeze(1)  # [n_scene, 1]
        agent_idx = torch.arange(scores.shape[1]).unsqueeze(0)  # [1, n_agent]
        mode_idx = []
        scores_clone = scores.clone()
        for _ in range(k_pred):
            # [n_scene, n_agent]
            _idx = scores_clone.max(-1)[1]
            # [n_scene, n_agent, n_pred], True entry has w=0.01, False entry has w=1.0
            w_mask = ~(within_dist[scene_idx, agent_idx, _idx]) * 0.99 + 0.01
            # [n_scene, n_agent, n_pred], suppress all preds close to the selected one, by multiplying the prob by 0.1
            scores_clone *= w_mask
            scores_clone[scene_idx, agent_idx, _idx] = -1
            # append to mode_idx
            mode_idx.append(_idx)

        mode_idx = torch.stack(mode_idx, dim=-1)  # [n_scene, n_agent, k_pred]
        scene_idx = scene_idx.unsqueeze(-1)  # [n_scene, 1, 1]
        agent_idx = agent_idx.unsqueeze(-1)  # [1, n_agent, 1]
        trajs_k = trajs[scene_idx, agent_idx, mode_idx]
        scores_k = scores[scene_idx, agent_idx, mode_idx]
        scores_k = scores_k / scores_k.sum(-1, keepdim=True)
        return trajs_k, scores_k

    @staticmethod
    def traj_topk(trajs: Tensor, scores: Tensor, k_pred: int) -> Tuple[Tensor, Tensor]:
        """
        Args:
            trajs: [n_scene, n_agent, n_pred, n_step, 2]
            scores: [n_scene, n_agent, n_pred], normalized prob
            k_pred: int

        Returns:
            trajs_k: [n_scene, n_agent, k_pred, n_step, 2]
            scores_k: [n_scene, n_agent, k_pred], normalized prob
        """
        scene_idx = torch.arange(scores.shape[0])[:, None, None]  # [n_scene, 1, 1]
        agent_idx = torch.arange(scores.shape[1])[None, :, None]  # [1, n_agent, 1]
        mode_idx = scores.topk(k_pred, dim=-1, sorted=False)[1]  # [n_scene, n_agent, k_pred]
        trajs_k = trajs[scene_idx, agent_idx, mode_idx]
        scores_k = scores[scene_idx, agent_idx, mode_idx]

        scores_k = scores_k / scores_k.sum(-1, keepdim=True)
        return trajs_k, scores_k

    @staticmethod
    def traj_aggr(
        trajs: Tensor, scores: Tensor, k_pred: int, thresh: float, n_iter_em: int, use_ade: bool,
    ) -> Tuple[Tensor, Tensor]:
        """
        Args:
            trajs: [n_scene, n_agent, n_pred, n_step, 2]
            scores: [n_scene, n_agent, n_pred], normalized prob
            k_pred: int
            thresh: in meters, if < 0, just topk scores
            n_iter_em: nubmer of iterations for k-means em
        Returns:
            trajs_k: [n_scene, n_agent, k_pred, n_step, 2]
            scores_k: [n_scene, n_agent, k_pred], normalized prob
        """
        scene_idx = torch.arange(scores.shape[0]).unsqueeze(1)  # [n_scene, 1]
        agent_idx = torch.arange(scores.shape[1]).unsqueeze(0)  # [1, n_agent]

        # ! compute mode_idx mpa: greedy k-means center
        # [n_scene, n_agent, n_pred, n_pred]
        # within_dist = torch.norm(trajs[:, :, :, -1].unsqueeze(2) - trajs[:, :, :, -1].unsqueeze(3), dim=-1) < thresh
        # mode_idx = []
        # for _ in range(k_pred):
        #     # [n_scene, n_agent]
        #     _idx = (scores.unsqueeze(2) * within_dist).sum(-1).max(-1)[1]
        #     # [n_scene, n_agent, n_pred], False at all nodes belonging to the selected center
        #     mask = ~(within_dist[scene_idx, agent_idx, _idx])
        #     # [n_scene, n_agent, n_pred, n_pred], remove selected nodes
        #     within_dist &= mask.unsqueeze(-1)
        #     # [n_scene, n_agent, n_pred, n_pred], set selected mode to 0 for remaning nodes
        #     within_dist &= mask.unsqueeze(-2)
        #     # append to mode_idx
        #     mode_idx.append(_idx)

        # ! compute mode_idx mtr: [n_scene, n_agent, k_pred]
        # [n_scene, n_agent, n_pred, n_pred]
        xy = trajs[..., :2]
        if use_ade:
            within_dist = (torch.norm(xy.unsqueeze(2) - xy.unsqueeze(3), dim=-1).mean(-1)) < thresh
        else:
            within_dist = torch.norm(xy[:, :, :, -1].unsqueeze(2) - xy[:, :, :, -1].unsqueeze(3), dim=-1) < thresh
        n_pred = scores.shape[-1]
        mode_idx = []
        scores_clone = scores.clone()
        for _ in range(k_pred):
            # [n_scene, n_agent]
            _idx = scores_clone.max(-1)[1]
            # [n_scene, n_agent, n_pred], True entry has w=0.1, False entry has w=1.0
            w_mask = ~(within_dist[scene_idx, agent_idx, _idx]) * 0.9 + 0.1
            # [n_scene, n_agent, n_pred], suppress all preds close to the selected one, by multiplying the prob by 0.1
            scores_clone *= w_mask
            # [n_scene, n_agent, n_pred], set prob to negative if selected
            w_mask_idx = nn.functional.one_hot(_idx, n_pred)
            scores_clone -= w_mask_idx
            # append to mode_idx
            mode_idx.append(_idx)

        mode_idx = torch.stack(mode_idx, dim=-1)  # [n_scene, n_agent, k_pred]
        scene_idx = scene_idx.unsqueeze(-1)  # [n_scene, 1, 1]
        agent_idx = agent_idx.unsqueeze(-1)  # [1, n_agent, 1]
        trajs_k = trajs[scene_idx, agent_idx, mode_idx]  # [n_scene, n_agent, k_pred, n_step, 2]
        scores_k = scores[scene_idx, agent_idx, mode_idx]  # [n_scene, n_agent, k_pred]

        # ! em for kmeans
        # trajs: [n_scene, n_agent, n_pred, n_step, 2]
        for _ in range(n_iter_em):
            # [n_scene, n_agent, n_pred, k_pred]
            xy_k = trajs_k[..., :2]
            xy = trajs[..., :2]
            if use_ade:
                dist = torch.norm(xy_k.unsqueeze(2) - xy.unsqueeze(3), dim=-1).mean(-1)
            else:
                dist = torch.norm(xy_k[:, :, :, -1].unsqueeze(2) - xy[:, :, :, -1].unsqueeze(3), dim=-1)
            # [n_scene, n_agent, n_pred]
            min_dist, assignment = dist.min(-1)
            # [n_scene, n_agent, n_pred, k_pred]
            assignment = nn.functional.one_hot(assignment, k_pred)
            empty_scene, empty_agent, empty_pred = torch.where(assignment.sum(2) == 0)

            if len(empty_scene) > 0:
                for i in range(len(empty_scene)):

                    # assign to furtherst
                    # max_i = min_dist[empty_scene[i], empty_agent[i]].argmax()
                    # min_dist[empty_scene[i], empty_agent[i], max_i] = -1
                    # assignment[empty_scene[i], empty_agent[i], max_i] = 0
                    # assignment[empty_scene[i], empty_agent[i], max_i, empty_pred[i]] = 1

                    # split
                    counter_n, max_i = assignment.sum(2)[empty_scene[i], empty_agent[i]].max(0)
                    _split = torch.where(assignment[empty_scene[i], empty_agent[i], :, max_i] == 1)[0][: counter_n // 2]
                    assignment[empty_scene[i], empty_agent[i], _split, max_i] = 0
                    assignment[empty_scene[i], empty_agent[i], _split, empty_pred[i]] = 1

            n_members = assignment.sum(2)  # [n_scene, n_agent, k_pred]
            # update trajs_k: [n_scene, n_agent, k_pred, n_step, 2]
            trajs_k = (trajs.unsqueeze(3) * assignment[:, :, :, :, None, None]).sum(2)  # sum over n_pred
            trajs_k /= n_members[:, :, :, None, None]
            # update scores_k: [n_scene, n_agent, k_pred]
            scores_k = (scores.unsqueeze(3) * assignment).sum(2) / n_members

        scores_k = scores_k / scores_k.sum(-1, keepdim=True)
        return trajs_k, scores_k

src/data_modules/waymo_post_processing.py
zhejz/TrafficBots
# Licensed under the CC BY-NC 4.0 license (https://creativecommons.org/licenses/by-nc/4.0/)
from typing import Optional, Tuple
from omegaconf import DictConfig
from torch import Tensor, nn
import torch
from .modules.mlp import MLP
from .modules.transformer import TransformerBlock
from .modules.agent_temporal import MultiAgentGRULoop, TemporalAggregate
from .modules.distributions import MyDist, DestCategorical, DiagGaussian
from .modules.attention import Attention
from utils.transform_utils import torch_rad2rot, torch_pos2local, torch_pos2global


class GoalManager(nn.Module):
    def __init__(
        self,
        tf_cfg: DictConfig,
        goal_predictor: DictConfig,
        goal_attr_mode: str,
        goal_in_local: bool,
        dest_detach_map_feature: bool,
        disable_if_reached: bool,
    ) -> None:
        super().__init__()
        self.goal_attr_mode = goal_attr_mode
        self.goal_in_local = goal_in_local
        self.dest_detach_map_feature = dest_detach_map_feature
        self.disable_if_reached = disable_if_reached
        hidden_dim = tf_cfg.d_model

        self.update_goal = False
        if self.goal_attr_mode == "dummy":
            self.dummy = True
            self.out_dim = -1
            self.goal_predictor = None
        elif self.goal_attr_mode == "dest":
            self.dummy = False
            self.out_dim = hidden_dim
            self.goal_predictor = DestPredictor(tf_cfg=tf_cfg, **goal_predictor)
        elif self.goal_attr_mode == "goal_xy":
            self.dummy = False
            if goal_in_local:
                self.update_goal = True
            self.out_dim = 2
            self.goal_predictor = GoalPredictor(tf_cfg=tf_cfg, goal_in_local=goal_in_local, **goal_predictor)
        else:
            raise NotImplementedError

    @torch.no_grad()
    def get_gt_goal(
        self, agent_valid: Tensor, gt_goal: Tensor, gt_dest: Tensor
    ) -> Tuple[Optional[Tensor], Optional[Tensor]]:
        """
        Args:
            agent_valid: [n_scene, n_step_hist, n_agent], bool
            gt_goal: [n_scene, n_agent, 4], (x, y, theta, spd)
            gt_dest: [n_scene, n_agent], int64: index to map n_pl

        Returns:
            gt: [n_scene, n_agent, (?)]
            valid: [n_scene, n_agent]
        """
        if self.goal_attr_mode == "dummy":
            gt = None
            valid = None
        elif self.goal_attr_mode == "dest":
            gt = gt_dest
            valid = agent_valid.any(1)
        elif self.goal_attr_mode == "goal_xy":
            gt = gt_goal[..., :2]
            valid = agent_valid.any(1)
        else:
            raise NotImplementedError

        return gt, valid

    def pred_goal(self, *args, **kwargs) -> Optional[MyDist]:
        if self.goal_predictor is None:
            return None
        else:
            return self.goal_predictor(*args, **kwargs)

    def get_goal_feature(self, goal: Tensor, as_state: Tensor, map_feature: Tensor) -> Tensor:
        """
        Args:
            goal: [n_batch, n_agent, 4] (x,y,yaw,spd) in global coordinate
            or goal: [n_batch, n_agent] int64 index to map n_pl
            as_state: [n_batch, (n_step), n_agent, 4] (x,y,yaw,spd) in global coordinate
            map_feature: [n_scene, n_pl, hidden_dim]

        Return:
            goal_feature: [n_batch, (n_step), n_agent, self.out_dim]
        """
        if as_state.dim() == 4:
            n_step = as_state.shape[1]
            if self.goal_attr_mode == "dest":
                # [n_batch, n_step, n_agent]
                goal = goal.unsqueeze(1).expand(-1, n_step, -1)
            else:
                # [n_batch, n_step, n_agent, 4]
                goal = goal.unsqueeze(1).expand(-1, n_step, -1, -1)

        if self.goal_attr_mode == "dest":
            if self.dest_detach_map_feature:
                map_feature = map_feature.detach()
            goal_feature = self._get_dest_feature(goal, map_feature)
        elif self.goal_attr_mode == "goal_xy":
            # [n_batch, (n_step), n_agent, 2]
            goal_feature = goal[..., :2]
            if self.goal_in_local:
                as_state = as_state.detach()
                ref_pos = as_state[..., :2].unsqueeze(-2)  # [n_batch, (n_step), n_agent, 1, 2]
                ref_rot = torch_rad2rot(as_state[..., 2])  # [n_batch, (n_step), n_agent, 2, 2]
                goal_feature = torch_pos2local(goal_feature.unsqueeze(-2), ref_pos, ref_rot).squeeze(-2)
        else:
            raise NotImplementedError

        return goal_feature

    @staticmethod
    def _get_dest_feature(dest: Tensor, map_feature: Tensor) -> Tensor:
        """
        Args:
            dest: [n_scene, (n_step), n_agent] int64(dest)
            map_feature: [n_scene, n_pl, hidden_dim]

        Return:
            dest_feature: [n_scene, (n_step), n_agent, hidden_dim]
        """
        # [n_scene, 1]
        batch_index = torch.arange(dest.shape[0]).unsqueeze(1)
        if dest.dim() == 3:
            dest_feature = []
            for k in range(dest.shape[1]):
                dest_feature.append(map_feature[batch_index, dest[:, k, :]])
            dest_feature = torch.stack(dest_feature, dim=1)
        else:
            dest_feature = map_feature[batch_index, dest]
        return dest_feature

    @torch.no_grad()
    def disable_goal_reached(
        self, goal_valid: Tensor, agent_valid: Tensor, dest_reached: Tensor, goal_reached: Tensor
    ) -> Tensor:
        """
        Args:
            goal_valid: [n_batch, n_agent] at t
            agent_valid: [n_batch, n_agent] at t
            dest_reached: [n_batch, n_agent] at t
            goal_reached: [n_batch, n_agent] at t

        Returns:
            goal_valid: [n_batch, n_agent] at t
        """
        if goal_valid is not None:
            goal_valid = goal_valid & agent_valid
        if self.disable_if_reached:
            if self.goal_attr_mode == "dest":
                goal_valid = goal_valid & (~dest_reached)
            elif self.goal_attr_mode == "goal_xy":
                goal_valid = goal_valid & (~goal_reached)
        return goal_valid


class DestPredictor(nn.Module):
    def __init__(
        self,
        tf_cfg: DictConfig,
        mode: str = "mlp",
        n_layer_gru: int = -1,
        use_layernorm: bool = True,
        res_add_gru: bool = True,
        detach_features: bool = True,
    ) -> None:
        super().__init__()
        assert mode in ["transformer", "transformer_aggr", "mlp", "attn"]
        self.mode = mode
        self.detach_features = detach_features
        self.res_add_gru = res_add_gru
        hidden_dim = tf_cfg.d_model

        if n_layer_gru > 0:
            self.gru_as = MultiAgentGRULoop(hidden_dim, num_layers=n_layer_gru, dropout=tf_cfg.dropout_p)
        else:
            self.gru_as = None
        self.agr_as = TemporalAggregate("last_valid")

        if self.mode == "transformer" or self.mode == "transformer_aggr":
            self.transformer_pl2as = TransformerBlock(hidden_dim, n_layer=1, **tf_cfg)
            self.mlp = MLP([hidden_dim, hidden_dim, 1], end_layer_activation=False, use_layernorm=use_layernorm)
        elif self.mode == "mlp":
            self.mlp = MLP(
                [hidden_dim * 2, hidden_dim, hidden_dim, 1], end_layer_activation=False, use_layernorm=use_layernorm
            )
        elif self.mode == "attn":
            self.attn = Attention(
                d_model=tf_cfg.d_model, n_head=tf_cfg.n_head, dropout_p=tf_cfg.dropout_p, bias=tf_cfg.bias
            )
        else:
            raise NotImplementedError

    def forward(
        self,
        agent_type: Tensor,
        map_type: Tensor,
        agent_state: Tensor,
        agent_feature: Tensor,
        agent_feature_valid: Tensor,
        map_feature: Tensor,
        map_feature_valid: Tensor,
        tl_feature: Optional[Tensor] = None,
        tl_feature_valid: Optional[Tensor] = None,
    ) -> MyDist:
        """
        Args:
            agent_type: [n_scene, n_agent, 3] [Vehicle=0, Pedestrian=1, Cyclist=2] one hot
            map_type: [n_scene, n_pl, 11], one_hot, n_pl_type=11
            agent_state: [n_scene, n_step_hist, n_agent, 4], (x,y,yaw,spd)
            agent_feature: [n_scene, n_step_hist, n_agent, hidden_dim]
            agent_feature_valid: [n_scene, n_step_hist, n_agent] bool
            map_feature: [n_scene, n_pl, hidden_dim]
            map_feature_valid: [n_scene, n_pl] bool
            tl_feature: [n_scene, n_step_hist, n_tl, hidden_dim]
            tl_feature_valid: [n_scene, n_step_hist, n_tl] bool

        Returns:
            dest_dist: DestCategorical
        """
        if self.detach_features:
            agent_feature = agent_feature.detach()
            map_feature = map_feature.detach()

        # [n_scene, n_pl]
        # WOMD: FREEWAY = 0, SURFACE_STREET = 1, STOP_SIGN = 2, BIKE_LANE = 3, TYPE_ROAD_EDGE_BOUNDARY = 4
        map_type_mask = ~(map_feature_valid & (map_type[:, :, :5].any(-1)))

        # exclude (3) for veh(0): [n_scene, n_agent, 1] & [n_scene, 1, n_pl]
        attn_mask_veh = agent_type[:, :, [0]] & map_type[:, :, 3].unsqueeze(1)
        # exclude (0,1,2,3) for ped(1): [n_scene, n_agent, 1] & [n_scene, 1, n_pl]
        attn_mask_ped = agent_type[:, :, [1]] & map_type[:, :, :4].any(-1).unsqueeze(1)
        # exclude (0,1,2) for cyc(2): [n_scene, n_agent, 1] & [n_scene, 1, n_pl]
        attn_mask_cyc = agent_type[:, :, [2]] & map_type[:, :, :3].any(-1).unsqueeze(1)
        # [n_scene, n_agent, n_pl]
        attn_mask = attn_mask_veh | attn_mask_ped | attn_mask_cyc

        n_scene, n_pl, hidden_dim = map_feature.shape
        n_agent = agent_feature_valid.shape[2]
        dist_valid = agent_feature_valid.any(1)
        if self.mode == "transformer_aggr":
            if self.gru_as is None:
                tgt = agent_feature
            else:
                tgt, _ = self.gru_as(agent_feature, agent_feature_valid)
                if self.res_add_gru:
                    tgt = tgt + agent_feature
            # [n_scene, n_agent, hidden_dim]
            tgt, tgt_valid = self.agr_as(tgt, agent_feature_valid)

            map_feature_repeated = map_feature.unsqueeze(1).expand(-1, n_agent, -1, -1).flatten(0, 1)
            map_valid_repeated = map_feature_valid.unsqueeze(1).expand(-1, n_agent, -1).flatten(0, 1)
            # [n_scene*n_agent, n_pl, hidden_dim]
            map_feature_repeated, _ = self.transformer_pl2as(
                src=map_feature_repeated,  # [n_scene*n_agent, n_pl, hidden_dim]
                src_padding_mask=~map_valid_repeated,  # [n_scene*n_agent, n_pl]
                tgt=tgt.flatten(0, 1).unsqueeze(1),  # [n_scene*n_agent, 1, hidden_dim]
                tgt_padding_mask=~tgt_valid.flatten(0, 1).unsqueeze(1),  # [n_scene*n_agent, 1]
            )
            # [n_scene, n_agent, n_pl]
            logits = self.mlp(map_feature_repeated.view(n_scene, n_agent, n_pl, hidden_dim)).squeeze(-1)

        elif self.mode == "transformer":
            # downsample
            k_skip = 2
            # [n_scene, n_step_hist, n_agent, hidden_dim] -> [n_scene*n_agent, n_step_hist/2, hidden_dim]
            tgt = agent_feature.transpose(1, 2)[:, :, 0::k_skip].flatten(0, 1)
            # [n_scene, n_step_hist, n_agent]  -> [n_scene*n_agent, n_step_hist/2]
            tgt_valid = agent_feature_valid.transpose(1, 2)[:, :, 0::k_skip].flatten(0, 1)

            # agent_feature: [n_scene, n_agent, n_step_hist, hidden_dim]
            # agent_feature_valid: [n_scene, n_agent, n_step_hist] bool

            map_feature_repeated = map_feature.unsqueeze(1).expand(-1, n_agent, -1, -1).flatten(0, 1)
            map_valid_repeated = map_feature_valid.unsqueeze(1).expand(-1, n_agent, -1).flatten(0, 1)
            # [n_scene*n_agent, n_pl, hidden_dim]
            map_feature_repeated, _ = self.transformer_pl2as(
                src=map_feature_repeated,  # [n_scene*n_agent, n_pl, hidden_dim]
                src_padding_mask=~map_valid_repeated,  # [n_scene*n_agent, n_pl]
                tgt=tgt,  # [n_scene*n_agent, n_step_hist, hidden_dim]
                tgt_padding_mask=~tgt_valid,  # [n_scene*n_agent, c_step]
            )
            # [n_scene, n_agent, n_pl]
            logits = self.mlp(map_feature_repeated.view(n_scene, n_agent, n_pl, hidden_dim)).squeeze(-1)

        elif self.mode == "mlp":
            if self.gru_as is None:
                tgt = agent_feature
            else:
                tgt, _ = self.gru_as(agent_feature, agent_feature_valid)
                if self.res_add_gru:
                    tgt = tgt + agent_feature
            # [n_scene, n_agent, hidden_dim]
            tgt, tgt_valid = self.agr_as(tgt, agent_feature_valid)
            # [n_scene, n_agent, n_pl, hidden_dim]
            tgt = tgt.unsqueeze(2).expand(-1, -1, n_pl, -1)
            src = map_feature.unsqueeze(1).expand(-1, n_agent, -1, -1)
            # [n_scene, n_agent, n_pl]
            logits = self.mlp(torch.cat([src, tgt], dim=-1)).squeeze(-1)

        elif self.mode == "attn":
            if self.gru_as is None:
                src = agent_feature
            else:
                src, _ = self.gru_as(agent_feature, agent_feature_valid)
                if self.res_add_gru:
                    src = src + agent_feature
            # [n_scene, n_agent, hidden_dim]
            src, _ = self.agr_as(src, agent_feature_valid)
            attn_mask = attn_mask.repeat_interleave(self.n_head, 0)
            # [n_scene, n_agent, n_pl]
            _, probs = self.attn(
                src=src,
                tgt=map_feature,
                tgt_padding_mask=map_feature,
                attn_mask=attn_mask,
                key_padding_mask=map_type_mask,
            )
            logits = torch.log(probs)

        logits = logits.masked_fill(map_type_mask.unsqueeze(1), float("-inf"))
        logits = logits.masked_fill(attn_mask, float("-inf"))
        logits = logits.masked_fill(~dist_valid.unsqueeze(-1), 0)
        logits = logits.masked_fill((logits == float("-inf")).all(-1).unsqueeze(-1), 0)
        return DestCategorical(logits=logits, valid=dist_valid)


class GoalPredictor(nn.Module):
    def __init__(
        self,
        tf_cfg: DictConfig,
        goal_in_local: bool,
        mode: str = "transformer",
        n_layer_gru: int = 3,
        use_layernorm: bool = True,
        res_add_gru: bool = True,
        detach_features: bool = True,
    ) -> None:
        super().__init__()
        self.mode = mode
        self.detach_features = detach_features
        self.res_add_gru = res_add_gru
        self.goal_in_local = goal_in_local
        hidden_dim = tf_cfg.d_model

        if n_layer_gru > 0:
            self.gru_as = MultiAgentGRULoop(hidden_dim, num_layers=n_layer_gru, dropout=tf_cfg.dropout_p)
        else:
            self.gru_as = None
        self.agr_as = TemporalAggregate("last_valid")

        self.transformer_as2pl = TransformerBlock(n_layer=1, **tf_cfg)

        self.mlp_mean = MLP([hidden_dim, hidden_dim, 2], end_layer_activation=False, use_layernorm=use_layernorm)
        self.log_std = nn.Parameter(2.0 * torch.ones(2), requires_grad=True)

    def forward(
        self,
        agent_type: Tensor,
        map_type: Tensor,
        agent_state: Tensor,
        agent_feature: Tensor,
        agent_feature_valid: Tensor,
        map_feature: Tensor,
        map_feature_valid: Tensor,
        tl_feature: Optional[Tensor] = None,
        tl_feature_valid: Optional[Tensor] = None,
    ) -> MyDist:
        """
        Args:
            agent_type: [n_scene, n_agent, 3] [Vehicle=0, Pedestrian=1, Cyclist=2] one hot
            map_type: [n_scene, n_pl, 11], one_hot, n_pl_type=11
            agent_state: [n_scene, n_step_hist, n_agent, 4], (x,y,yaw,spd)
            agent_feature: [n_scene, n_step_hist, n_agent, hidden_dim]
            agent_feature_valid: [n_scene, n_step_hist, n_agent] bool
            map_feature: [n_scene, n_pl, hidden_dim]
            map_feature_valid: [n_scene, n_pl] bool
            tl_feature: [n_scene, n_step_hist, n_tl, hidden_dim]
            tl_feature_valid: [n_scene, n_step_hist, n_tl] bool

        Returns:
            dest_dist: diag_gaus
        """
        if self.detach_features:
            agent_feature = agent_feature.detach()
            map_feature = map_feature.detach()

        if self.gru_as is None:
            src = agent_feature
        else:
            src, _ = self.gru_as(agent_feature, agent_feature_valid)
            if self.res_add_gru:
                src = src + agent_feature
        # [n_scene, n_agent, hidden_dim]
        src, src_valid = self.agr_as(src, agent_feature_valid)

        # [n_scene, n_agent, hidden_dim]
        goal_feature, _ = self.transformer_as2pl(
            src=src,  # [n_scene, n_agent, hidden_dim]
            src_padding_mask=~src_valid,  # [n_scene, n_agent]
            tgt=map_feature,  # [n_scene, n_pl, hidden_dim]
            tgt_padding_mask=~map_feature_valid,  # [n_scene, n_pl]
            need_weights=False,
        )

        goal_mean = self.mlp_mean(goal_feature)

        if self.goal_in_local:
            ref_pos = agent_state[:, -1, :, :2].unsqueeze(-2)  # [n_scene, n_agent, 1, 2]
            ref_rot = torch_rad2rot(agent_state[:, -1, :, 2])  # [n_scene, n_agent, 2, 2]
            goal_mean = torch_pos2global(goal_mean.unsqueeze(-2), ref_pos, ref_rot).squeeze(-2)

        goal_valid = agent_feature_valid.any(1)
        goal_mean = goal_mean.masked_fill(~goal_valid.unsqueeze(-1), 0)
        return DiagGaussian(mean=goal_mean, log_std=self.log_std, valid=goal_valid)

src/models/goal_manager.py
zhejz/TrafficBots
# Licensed under the CC BY-NC 4.0 license (https://creativecommons.org/licenses/by-nc/4.0/)
from typing import Optional
from torch import Tensor, nn
import torch
from omegaconf import DictConfig
import hydra
from .modules.mlp import MLP
from .modules.transformer import TransformerBlock
from .modules.distributions import MyDist, DummyLatent, DiagGaussian, MultiCategorical
from .modules.agent_interaction import MultiAgentTF
from .modules.agent_temporal import TemporalAggregate


class LatentEncoder(nn.Module):
    def __init__(
        self,
        latent_dim: int,
        temporal_down_sample_rate: int,
        shared_post_prior_net: bool,
        shared_transformer_as: bool,
        latent_prior: DictConfig,
        latent_post: DictConfig,
        tf_cfg: DictConfig,
        interaction_first: bool,
        transformer_as2pl: TransformerBlock,
        transformer_as2tl: TransformerBlock,
        agent_temporal: DictConfig,
        agent_interaction: DictConfig,
        temporal_aggregate: DictConfig,
    ):
        super().__init__()
        hidden_dim = tf_cfg.d_model
        self.out_dim = latent_dim
        self.dummy = latent_post.dist_type == "dummy"
        self.temporal_down_sample_rate = temporal_down_sample_rate
        self.interaction_first = interaction_first

        # agent state map aware
        if shared_transformer_as:
            self.transformer_as2pl = transformer_as2pl
            self.transformer_as2tl = transformer_as2tl
        else:
            self.transformer_as2pl = TransformerBlock(n_layer=len(transformer_as2pl.layers), **tf_cfg)
            self.transformer_as2tl = TransformerBlock(n_layer=len(transformer_as2tl.layers), **tf_cfg)

        # such that prior and posterior distribution have different standard deviation.
        self.latent_prior_dist = DistEncoder(**latent_prior, hidden_dim=hidden_dim, out_dim=latent_dim)
        self.latent_post_dist = DistEncoder(**latent_post, hidden_dim=hidden_dim, out_dim=latent_dim)

        if self.latent_post_dist.skip_forward:
            self.temporal_aggregate = None
            self.agent_temporal_post = None
            self.agent_interaction_post = None
            self.agent_temporal_prior = None
            self.agent_interaction_prior = None
        else:
            self.temporal_aggregate = TemporalAggregate(**temporal_aggregate)
            self.agent_temporal_post = hydra.utils.instantiate(agent_temporal, hidden_dim=hidden_dim)
            self.agent_interaction_post = MultiAgentTF(hidden_dim=hidden_dim, tf_cfg=tf_cfg, **agent_interaction)
            if self.latent_prior_dist.skip_forward:
                self.agent_temporal_prior = None
                self.agent_interaction_prior = None
            elif shared_post_prior_net:
                self.agent_temporal_prior = self.agent_temporal_post
                self.agent_interaction_prior = self.agent_interaction_post
            else:
                self.agent_temporal_prior = hydra.utils.instantiate(agent_temporal, hidden_dim=hidden_dim)
                self.agent_interaction_prior = MultiAgentTF(hidden_dim=hidden_dim, tf_cfg=tf_cfg, **agent_interaction)

    def forward(
        self,
        agent_feature: Tensor,
        agent_feature_valid: Tensor,
        map_feature: Tensor,
        map_feature_valid: Tensor,
        tl_feature: Optional[Tensor] = None,
        tl_feature_valid: Optional[Tensor] = None,
        posterior: bool = False,
    ) -> MyDist:
        """
        Args:
            agent_feature: [n_scene, n_step, n_agent, hidden_dim]
            agent_feature_valid: [n_scene, n_step, n_agent] bool
            map_feature: [n_scene, n_pl, hidden_dim]
            map_feature_valid: [n_scene, n_pl] bool
            tl_feature: [n_scene, n_step, n_tl, hidden_dim]
            tl_feature_valid: [n_scene, n_step, n_tl] bool

        Returns: for each agent a latent distribution that considers temporal relation and interaction between agents.
            latent: [n_batch, n_agent, hidden_dim]
        """
        if posterior and self.latent_post_dist.skip_forward:
            return self.latent_post_dist(agent_feature[:, 0], agent_feature_valid.any(1))
        elif (not posterior) and self.latent_prior_dist.skip_forward:
            return self.latent_prior_dist(agent_feature[:, 0], agent_feature_valid.any(1))
        else:
            # ! downsampling
            if self.temporal_down_sample_rate > 1:
                assert (agent_feature_valid.shape[1] - 1) % self.temporal_down_sample_rate == 0
                agent_feature_valid = agent_feature_valid[:, :: self.temporal_down_sample_rate]
                agent_feature = agent_feature[:, :: self.temporal_down_sample_rate]
                tl_feature_valid = tl_feature_valid[:, :: self.temporal_down_sample_rate]
                tl_feature = tl_feature[:, :: self.temporal_down_sample_rate]

            # [n_batch, n_step, n_agent, hidden_dim]
            as_feature_map_aware = agent_feature
            as_shape = as_feature_map_aware.shape
            # ! attention to map polyline
            as_feature_map_aware, _ = self.transformer_as2pl(
                src=as_feature_map_aware.flatten(1, 2),  # [n_batch, n_step * n_agent, hidden_dim]
                src_padding_mask=~agent_feature_valid.flatten(1, 2),  # [n_batch, n_step * n_agent]
                tgt=map_feature,  # [n_batch, n_pl, hidden_dim]
                tgt_padding_mask=~map_feature_valid,  # [n_batch, n_pl]
            )
            as_feature_map_aware = as_feature_map_aware.view(as_shape)
            # ! attention to traffic light
            as_feature_map_aware, _ = self.transformer_as2tl(
                src=as_feature_map_aware.flatten(0, 1),  # [n_batch * n_step, n_agent, hidden_dim]
                src_padding_mask=~agent_feature_valid.flatten(0, 1),  # [n_batch * n_step, n_agent]
                tgt=tl_feature.flatten(0, 1),  # [n_batch * n_step, n_tl, hidden_dim]
                tgt_padding_mask=~tl_feature_valid.flatten(0, 1),  # [n_batch * n_step, n_tl]
            )
            as_feature_map_aware = as_feature_map_aware.view(as_shape)

            # ! interaction and temporal
            if posterior:  # post
                if self.interaction_first:
                    latent_feature, _ = self.agent_interaction_post(
                        as_feature_map_aware, agent_feature, agent_feature_valid
                    )
                    latent_feature, _ = self.agent_temporal_post(latent_feature, agent_feature_valid)
                else:
                    latent_feature, _ = self.agent_temporal_post(as_feature_map_aware, agent_feature_valid)
                    latent_feature, _ = self.agent_interaction_post(latent_feature, agent_feature, agent_feature_valid)
                latent_feature, latent_valid = self.temporal_aggregate(latent_feature, agent_feature_valid)
                return self.latent_post_dist(latent_feature, latent_valid)
            else:  # prior
                if self.interaction_first:
                    latent_feature, _ = self.agent_interaction_prior(
                        as_feature_map_aware, agent_feature, agent_feature_valid
                    )
                    latent_feature, _ = self.agent_temporal_prior(latent_feature, agent_feature_valid)
                else:
                    latent_feature, _ = self.agent_temporal_prior(as_feature_map_aware, agent_feature_valid)
                    latent_feature, _ = self.agent_interaction_prior(latent_feature, agent_feature, agent_feature_valid)
                latent_feature, latent_valid = self.temporal_aggregate(latent_feature, agent_feature_valid)
                return self.latent_prior_dist(latent_feature, latent_valid)


class DistEncoder(nn.Module):
    def __init__(
        self, dist_type: str, hidden_dim: int, out_dim: int, use_layernorm: bool, log_std: float = 0.0, n_cat: int = 1
    ) -> None:
        """
        dist_type in {dummy, std_gaus, diag_gaus, gaus, cat}
        """
        super().__init__()
        self.dist_type = dist_type
        self.hidden_dim = hidden_dim
        self.out_dim = out_dim
        self.skip_forward = False

        if dist_type == "dummy":
            self.skip_forward = True
        elif dist_type == "std_gaus":
            self.log_std = nn.Parameter(log_std * torch.ones(out_dim), requires_grad=False)
            self.skip_forward = True
        elif dist_type == "diag_gaus":
            self.mlp_mean = MLP(
                [hidden_dim, hidden_dim, out_dim], end_layer_activation=False, use_layernorm=use_layernorm
            )
            if log_std is None:
                self.log_std = None
                self.mlp_log_std = MLP(
                    [hidden_dim, hidden_dim, out_dim], end_layer_activation=False, use_layernorm=use_layernorm,
                )
            else:
                self.log_std = nn.Parameter(log_std * torch.ones(out_dim), requires_grad=True)
        elif dist_type == "cat":
            assert out_dim % n_cat == 0
            self.n_cat = n_cat
            self.n_class = out_dim // self.n_cat
            self.mlp_logits = MLP(
                [hidden_dim, hidden_dim, out_dim], end_layer_activation=False, use_layernorm=use_layernorm
            )

    def forward(self, x: Tensor, valid: Tensor) -> MyDist:
        if self.dist_type == "dummy":
            out_dist = DummyLatent(x, valid)
        elif self.dist_type == "std_gaus":
            out_dist = DiagGaussian(
                torch.zeros([*valid.shape, self.out_dim], dtype=x.dtype, device=x.device), self.log_std, valid=valid
            )
        elif self.dist_type == "diag_gaus":
            if self.log_std is None:
                out_dist = DiagGaussian(self.mlp_mean(x, valid), self.mlp_log_std(x, valid), valid=valid)
            else:
                out_dist = DiagGaussian(self.mlp_mean(x, valid), self.log_std, valid=valid)
        elif self.dist_type == "cat":
            logits = self.mlp_logits(x, valid).view(*valid.shape, self.n_cat, self.n_class)
            out_dist = MultiCategorical(nn.functional.softmax(logits, -1), valid=valid)
        return out_dist

src/models/latent_encoder.py
zhejz/TrafficBots
# Licensed under the CC BY-NC 4.0 license (https://creativecommons.org/licenses/by-nc/4.0/)
from typing import Dict
import torch
from torch import Tensor, tensor
from torchmetrics.metric import Metric
from utils.transform_utils import cast_rad


class ErrorMetrics(Metric):
    def __init__(self, prefix: str, loss_for_teacher_forcing: bool = False) -> None:
        super().__init__(dist_sync_on_step=False)
        self.prefix = prefix
        self.loss_for_teacher_forcing = loss_for_teacher_forcing

        self.add_state("err_counter", default=tensor(0.0), dist_reduce_fx="sum")
        self.add_state("err_pos_meter", default=tensor(0.0), dist_reduce_fx="sum")
        self.add_state("err_rot_deg", default=tensor(0.0), dist_reduce_fx="sum")
        self.add_state("err_spd_m_per_s", default=tensor(0.0), dist_reduce_fx="sum")

    def update(
        self,
        pred_valid: Tensor,
        pred_states: Tensor,
        gt_valid: Tensor,
        gt_states: Tensor,
        override_masks: Tensor,
        agent_role: Tensor,
    ) -> None:
        """
        pred_valid: [n_batch, n_agent, (n_future), n_step_rollout]
        pred_states: [n_batch, n_agent, (n_future), n_step_rollout, 4]
        gt_valid: [n_batch, n_agent, n_step_rollout]
        gt_states: [n_batch, n_agent, n_step_rollout, 4]
        override_masks: [n_batch, n_agent, (n_future), n_step_rollout]
        agent_role: [n_batch, n_agent, 4] one_hot [sdc=0, interest=1, predict=2]
        """
        with torch.no_grad():
            mask_relevant = agent_role.any(-1).unsqueeze(-1).unsqueeze(-1)  # [n_batch, n_agent, 1, 1]
            gt_valid = gt_valid.unsqueeze(2)
            gt_states = gt_states.unsqueeze(2)

            pred_valid = pred_valid & mask_relevant
            if not self.loss_for_teacher_forcing:
                pred_valid = pred_valid & (~override_masks)

            # [n_batch, n_step, n_agent]
            err_valid = gt_valid & pred_valid
            # [n_batch, n_step, n_agent, 4]
            gt_states = gt_states.masked_fill(~err_valid.unsqueeze(-1), 0.0)
            states = pred_states.masked_fill(~err_valid.unsqueeze(-1), 0.0)

            self.err_counter += err_valid.sum()
            self.err_pos_meter += torch.norm(gt_states[..., :2] - states[..., :2], dim=-1).sum()
            self.err_rot_deg += torch.abs(torch.rad2deg(cast_rad(gt_states[..., 2] - states[..., 2]))).sum()
            self.err_spd_m_per_s += torch.abs(gt_states[..., 3] - states[..., 3]).sum()

    def compute(self) -> Dict[str, Tensor]:
        # logging
        out_dict = {
            f"{self.prefix}/err/pos_meter": self.err_pos_meter / self.err_counter,
            f"{self.prefix}/err/rot_deg": self.err_rot_deg / self.err_counter,
            f"{self.prefix}/err/spd_m_per_s": self.err_spd_m_per_s / self.err_counter,
        }

        return out_dict


class TrafficRuleMetrics(Metric):
    """
    Log traffic_rule_violations, Not based on ground truth trajectory.
    n_agent_collided / n_agent_valid, collided if collision happened at any time step.
    """

    def __init__(self, prefix: str, loss_for_teacher_forcing: bool = False) -> None:
        super().__init__(dist_sync_on_step=False)
        self.prefix = prefix
        self.loss_for_teacher_forcing = loss_for_teacher_forcing
        self.add_state("counter_agent", default=tensor(0.0), dist_reduce_fx="sum")
        self.add_state("counter_veh", default=tensor(0.0), dist_reduce_fx="sum")
        self.add_state("outside_map", default=tensor(0.0), dist_reduce_fx="sum")
        self.add_state("collided", default=tensor(0.0), dist_reduce_fx="sum")
        self.add_state("run_red_light", default=tensor(0.0), dist_reduce_fx="sum")
        self.add_state("goal_reached", default=tensor(0.0), dist_reduce_fx="sum")
        self.add_state("dest_reached", default=tensor(0.0), dist_reduce_fx="sum")
        self.add_state("run_road_edge", default=tensor(0.0), dist_reduce_fx="sum")
        self.add_state("passive", default=tensor(0.0), dist_reduce_fx="sum")

    def update(
        self,
        valid: Tensor,
        override_masks: Tensor,
        outside_map: Tensor,
        collided: Tensor,
        run_road_edge: Tensor,
        run_red_light: Tensor,
        passive: Tensor,
        goal_reached: Tensor,
        dest_reached: Tensor,
        agent_type: Tensor,
    ) -> None:
        with torch.no_grad():
            if self.loss_for_teacher_forcing:
                agent_valid = valid.any(-1)
            else:
                # exclude teacher_forcing agents
                # [n_batch, n_agent, n_step]
                agent_valid = valid & (~override_masks)
                invalid_mask = ~agent_valid
                outside_map = outside_map.masked_fill(invalid_mask, 0)
                collided = collided.masked_fill(invalid_mask, 0)
                run_road_edge = run_road_edge.masked_fill(invalid_mask, 0)
                run_red_light = run_red_light.masked_fill(invalid_mask, 0)
                passive = passive.masked_fill(invalid_mask, 0)
                goal_reached = goal_reached.masked_fill(invalid_mask, 0)
                dest_reached = dest_reached.masked_fill(invalid_mask, 0)
                # [n_batch, n_agent]
                agent_valid = agent_valid.any(-1)

        self.counter_agent += agent_valid.sum()
        mask_veh = agent_type[:, :, 0:1]
        self.counter_veh += (agent_valid & mask_veh).sum()
        self.outside_map += outside_map.any(-1).sum()
        self.collided += collided.any(-1).sum()
        self.run_road_edge += run_road_edge.any(-1).sum()
        self.run_red_light += run_red_light.any(-1).sum()
        self.passive += passive.any(-1).sum()
        self.goal_reached += goal_reached.any(-1).sum()
        self.dest_reached += dest_reached.any(-1).sum()

    def compute(self) -> Dict[str, Tensor]:
        out_dict = {
            f"{self.prefix}/traffic_rule/outside_map": self.outside_map / self.counter_agent,
            f"{self.prefix}/traffic_rule/collided": self.collided / self.counter_agent,
            f"{self.prefix}/traffic_rule/run_road_edge": self.run_road_edge / self.counter_veh,
            f"{self.prefix}/traffic_rule/run_red_light": self.run_red_light / self.counter_veh,
            f"{self.prefix}/traffic_rule/passive": self.passive / self.counter_veh,
            f"{self.prefix}/traffic_rule/goal_reached": self.goal_reached / self.counter_agent,
            f"{self.prefix}/traffic_rule/dest_reached": self.dest_reached / self.counter_agent,
        }
        return out_dict

src/models/metrics/logging.py
zhejz/TrafficBots
# Licensed under the CC BY-NC 4.0 license (https://creativecommons.org/licenses/by-nc/4.0/)
from typing import Optional
import torch
from torch import Tensor
from utils.transform_utils import cast_rad
from torch.distributions import kl_divergence, Independent, Normal, OneHotCategoricalStraightThrough


class AngularError:
    def __init__(self, criterion: str, angular_type: Optional[str]) -> None:
        assert angular_type in ["cast", "cosine", "vector", None]
        self.angular_type = angular_type
        if self.angular_type != "cosine":
            self.criterion = getattr(torch.nn, criterion)(reduction="none")

    def compute(self, preds: Tensor, target: Tensor) -> Tensor:
        """
        Args:
            preds: [n_batch, n_step, n_agent, (k_pred), 1]
            target: [n_batch, n_step, n_agent, (k_pred), 1]
        """
        if self.angular_type is None:
            error = self.criterion(preds, target)
        elif self.angular_type == "cast":
            diff_angle = cast_rad(preds - target)
            error = self.criterion(diff_angle, torch.zeros_like(diff_angle))
        elif self.angular_type == "cosine":
            error = 0.5 * (1 - torch.cos(preds - target))
        elif self.angular_type == "vector":
            error_cos = self.criterion(torch.cos(preds), torch.cos(target))
            error_sin = self.criterion(torch.sin(preds), torch.sin(target))
            error = error_cos + error_sin
        return error


class BalancedKL:
    """
    Mastering atari with discrete world models, Algorithm 2: KL Balancing with Automatic Differentiation
    """

    def __init__(self, kl_balance_scale: float, kl_free_nats: float) -> None:
        self.alpha = kl_balance_scale
        self.free_nats = kl_free_nats

    def compute(self, posterior: Normal, prior: Normal) -> Tensor:  # type: ignore
        """
        Args:
            posterior: [n_batch, n_agent, E], Normal
            prior: [n_batch, n_agent, E], Normal
        Return:
            error: [n_batch, n_agent, E]
        """
        if self.alpha > 0:
            # latent dist is either DiagGaussian or MultiCategorical, both are wrapped by Independent
            # assert type(posterior) == Independent
            if type(posterior.base_dist) == OneHotCategoricalStraightThrough:
                detach_post = Independent(OneHotCategoricalStraightThrough(probs=posterior.base_dist.probs.detach()), 1)
                detach_prior = Independent(OneHotCategoricalStraightThrough(probs=prior.base_dist.probs.detach()), 1)
            elif type(posterior.base_dist) == Normal:
                detach_post = Independent(
                    Normal(posterior.base_dist.loc.detach(), posterior.base_dist.scale.detach()), 1
                )
                detach_prior = Independent(Normal(prior.base_dist.loc.detach(), prior.base_dist.scale.detach()), 1)
            error_0 = kl_divergence(detach_post, prior)
            error_1 = kl_divergence(posterior, detach_prior)
            if self.free_nats > 0:
                error_0 = torch.max(error_0, error_0.new_full(error_0.size(), self.free_nats))
                error_1 = torch.max(error_1, error_1.new_full(error_1.size(), self.free_nats))
            error = self.alpha * error_0 + (1 - self.alpha) * error_1
        else:
            error = kl_divergence(posterior, prior)
            if self.free_nats > 0:
                error = torch.max(error, error.new_full(error.size(), self.free_nats))
        return error

src/models/metrics/loss.py
zhejz/TrafficBots
# Licensed under the CC BY-NC 4.0 license (https://creativecommons.org/licenses/by-nc/4.0/)
from typing import Dict, Optional
import torch
from torch import Tensor, tensor
from torchmetrics.metric import Metric
from .loss import BalancedKL
from models.modules.distributions import MyDist


class TrainingMetrics(Metric):
    def __init__(
        self,
        prefix: str,
        w_vae_kl: float = 0,
        kl_balance_scale: float = 0,
        kl_free_nats: float = -1,
        kl_for_unseen_agent: bool = True,
        w_diffbar_reward: float = 0,
        w_goal: float = 0,
        w_relevant_agent: float = 0,
        loss_for_teacher_forcing: bool = True,
        p_loss_for_irrelevant: float = -1.0,
        step_training_start: int = 0,
    ) -> None:
        super().__init__(dist_sync_on_step=False)
        self.prefix = prefix
        self.loss_for_teacher_forcing = loss_for_teacher_forcing
        self.p_loss_for_irrelevant = p_loss_for_irrelevant
        self.w_relevant_agent = w_relevant_agent  # set to greater than 0 to enable extra weights on relevant agents
        self.step_training_start = step_training_start

        # CVAE KL divergence
        self.w_vae_kl = w_vae_kl
        self.kl_for_unseen_agent = kl_for_unseen_agent
        if self.w_vae_kl > 0:
            self.use_vae_kl = True
            self.add_state("vae_kl_counter", default=tensor(0.0), dist_reduce_fx="sum")
            self.add_state("vae_kl", default=tensor(0.0), dist_reduce_fx="sum")
            self.l_vae_kl = BalancedKL(kl_balance_scale=kl_balance_scale, kl_free_nats=kl_free_nats)
        else:
            self.use_vae_kl = False

        # diffbar reward
        self.w_diffbar_reward = w_diffbar_reward
        if self.w_diffbar_reward > 0:
            self.use_diffbar_reward = True
            # for training
            self.add_state("diffbar_reward_counter", default=tensor(0.0), dist_reduce_fx="sum")
            self.add_state("diffbar_reward", default=tensor(0.0), dist_reduce_fx="sum")
        else:
            self.use_diffbar_reward = False

        self.w_goal = w_goal
        if self.w_goal > 0:
            self.add_state("goal_loss", default=tensor(0.0), dist_reduce_fx="sum")
            self.add_state("goal_counter", default=tensor(0.0), dist_reduce_fx="sum")
            self.train_goal_predictor = True
        else:
            self.train_goal_predictor = False

    def update(
        self,
        pred_valid: Tensor,
        diffbar_rewards_valid: Tensor,
        diffbar_rewards: Tensor,
        override_masks: Tensor,
        agent_role: Tensor,
        goal_valid: Optional[MyDist],
        goal_pred: Optional[MyDist],
        goal_gt: Optional[Tensor],
        latent_post: MyDist,
        latent_prior: MyDist,
    ) -> None:
        """
        pred_valid: [n_batch, n_agent, n_step_rollout]
        diffbar_rewards_valid: [n_batch, n_agent, n_step_rollout]
        diffbar_rewards: [n_batch, n_agent, n_step_rollout]
        override_masks: [n_batch, n_agent, n_step_rollout]
        agent_role: [n_batch, n_agent, 4] one_hot [sdc=0, interest=1, predict=2]
        goal_valid: [n_batch, n_agent]
        goal_gt: [n_batch, n_agent, 4] float32 [x, y, theta, v] or [n_batch, n_agent] int64: index to n_pl, or None
        goal_pred: MyDist or None  # distribution, gaussian or categorical
        """
        with torch.no_grad():
            if self.p_loss_for_irrelevant > 0:
                mask_relevant = agent_role.any(-1).unsqueeze(-1)  # [n_batch, n_agent, 1]
                pred_valid = pred_valid & mask_relevant
                mask_irrelevant = torch.bernoulli(torch.ones_like(mask_relevant) * self.p_loss_for_irrelevant).bool()
                pred_valid = pred_valid | mask_irrelevant
            if not self.loss_for_teacher_forcing:
                pred_valid = pred_valid & (~override_masks)
            if self.step_training_start > 0:
                pred_valid[:, :, : self.step_training_start] &= False

            if self.w_relevant_agent > 0:
                # [n_batch, n_agent]
                w_mask_rel = pred_valid.any(-1) + agent_role.any(-1) * self.w_relevant_agent
            else:
                w_mask_rel = None

        # ! CVAE KL divergence
        if self.use_vae_kl:
            # [n_batch, n_agent]
            if self.kl_for_unseen_agent:
                # posterior valid: unseen agent has unit prior.
                vae_kl_valid = latent_post.valid
            else:
                # prior valid: no kl loss for unseen agent.
                vae_kl_valid = latent_prior.valid
            vae_kl_valid = vae_kl_valid & (pred_valid.any(-1))
            error_vae = self.l_vae_kl.compute(latent_post.distribution, latent_prior.distribution)
            self.vae_kl_counter += vae_kl_valid.sum()
            if w_mask_rel is not None:
                error_vae *= w_mask_rel
            self.vae_kl += error_vae.masked_fill(~vae_kl_valid, 0.0).sum()

        # ! diffbar reward
        if self.use_diffbar_reward:
            reward_valid = pred_valid & diffbar_rewards_valid  # [n_batch, n_agent, n_step]
            # [n_batch, n_agent, n_step]
            error_rewards_dr = diffbar_rewards.masked_fill(~reward_valid, 0.0)
            if w_mask_rel is not None:
                error_rewards_dr *= w_mask_rel.unsqueeze(1)
            self.diffbar_reward -= error_rewards_dr.sum()
            self.diffbar_reward_counter += reward_valid.sum()

        # ! goal/dest prediction
        if self.train_goal_predictor:
            # [n_batch, n_agent]
            goal_valid = goal_pred.valid & (pred_valid.any(-1))
            # same as F.cross_entropy(self.distribution.logits.transpose(1, 2), goal_gt, reduction="none")
            goal_nll = -goal_pred.log_prob(goal_gt).masked_fill(~goal_valid, 0)
            if w_mask_rel is not None:
                goal_nll *= w_mask_rel
            self.goal_loss += goal_nll.sum()
            self.goal_counter += goal_valid.sum()

    def compute(self) -> Dict[str, Tensor]:
        out_dict = {f"{self.prefix}/loss": 0}

        # CVAE KL divergence
        if self.use_vae_kl:
            out_dict[f"{self.prefix}/vae_kl"] = self.w_vae_kl * self.vae_kl / self.vae_kl_counter
            out_dict[f"{self.prefix}/loss"] += out_dict[f"{self.prefix}/vae_kl"]

        # diffbar reward
        if self.use_diffbar_reward:
            # training
            out_dict[f"{self.prefix}/diffbar_reward"] = (
                self.w_diffbar_reward * self.diffbar_reward / self.diffbar_reward_counter
            )
            out_dict[f"{self.prefix}/loss"] += out_dict[f"{self.prefix}/diffbar_reward"]

        # goal/dest prediction
        if self.train_goal_predictor:
            out_dict[f"{self.prefix}/goal_loss"] = self.w_goal * self.goal_loss / self.goal_counter
            out_dict[f"{self.prefix}/loss"] += out_dict[f"{self.prefix}/goal_loss"]
        return out_dict

src/models/metrics/training.py
zhejz/TrafficBots
# Licensed under the CC BY-NC 4.0 license (https://creativecommons.org/licenses/by-nc/4.0/)
from typing import Dict, List, Optional, Tuple
import numpy as np
import torch
import tensorflow
from torch import Tensor
from torchmetrics.metric import Metric
from google.protobuf import text_format
from waymo_open_dataset.protos import motion_metrics_pb2
from waymo_open_dataset.metrics.python.config_util_py import get_breakdown_names_from_motion_config
from waymo_open_dataset.metrics.ops import py_metrics_ops


class WOMDMetrics(Metric):
    """
    validation metrics based on ground truth trajectory, using waymo_open_dataset api
    """

    def __init__(self, prefix: str, step_gt: int, step_current: int, interactive_challenge: bool) -> None:
        """
        submission_type: MOTION_PREDICTION = 1; INTERACTION_PREDICTION = 2; ARGOVERSE2 = 3
        """
        super().__init__(dist_sync_on_step=True)

        self.prefix = prefix
        self.step_gt = step_gt
        self.step_current = step_current
        self.interactive_challenge = interactive_challenge
        self.track_future_samples = step_gt - step_current
        assert self.track_future_samples == 80
        self.metrics_config, self.metrics_names = self._waymo_metrics_config_names(
            step_current, self.track_future_samples
        )
        self.metrics_config = self.metrics_config.SerializeToString()
        self.metrics_type = ["min_ade", "min_fde", "miss_rate", "overlap_rate", "mean_average_precision"]

        if self.interactive_challenge:
            self.m_joint = 1
            self.n_pred = 2
        else:
            self.m_joint = 8
            self.n_pred = 1

        self.add_state("prediction_trajectory_gpu", default=[], dist_reduce_fx="cat")
        self.add_state("prediction_score_gpu", default=[], dist_reduce_fx="cat")
        self.add_state("ground_truth_trajectory_gpu", default=[], dist_reduce_fx="cat")
        self.add_state("ground_truth_is_valid_gpu", default=[], dist_reduce_fx="cat")
        self.add_state("prediction_ground_truth_indices_mask_gpu", default=[], dist_reduce_fx="cat")
        self.add_state("object_type_gpu", default=[], dist_reduce_fx="cat")

        self.ops_inputs_cpu = {
            "prediction_trajectory": [],
            "prediction_score": [],
            "ground_truth_trajectory": [],
            "ground_truth_is_valid": [],
            "prediction_ground_truth_indices_mask": [],
            "object_type": [],
        }

    def update(self, batch: Dict[str, Tensor], pred_traj: Tensor, pred_score: Optional[Tensor] = None) -> None:
        """
        pytorch tensors on gpu/cpu
        pred_traj: [n_batch, step_start+1...step_end, n_agent, K, 2]
        pred_score: [n_batch, n_agent, K] normalized prob or None
        """
        # [n_batch, n_agent]
        mask_pred = batch["agent/role"][..., 2]
        # mask_other = (~mask_pred) & (batch["agent/valid"][:, self.step_current, :].cpu())
        mask_other = (~mask_pred) & batch["agent/valid"][:, : self.step_current + 1, :].all(1)

        # gt_traj: [n_batch, n_agent, step_gt, 7]
        gt_traj = torch.cat(
            [
                batch["agent/pos"],  # [n_batch, n_step, n_agent, 2]
                batch["agent/size"][..., :2].unsqueeze(1).expand(-1, batch["agent/pos"].shape[1], -1, -1),
                batch["agent/yaw_bbox"],  # [n_batch, n_step, n_agent, 1]
                batch["agent/vel"],  # [n_batch, n_step, n_agent, 2]
            ],
            axis=-1,
        ).transpose(1, 2)
        gt_traj = gt_traj[:, :, : self.step_gt + 1, :]

        # gt_valid: [n_batch, n_agent, step_gt]
        gt_valid = batch["agent/valid"].transpose(1, 2)
        gt_valid = gt_valid[:, :, : self.step_gt + 1]

        # agent_type: [n_batch, n_agent, 3] one_hot -> [n_batch, n_agent] [Vehicle=1, Pedestrian=2, Cyclist=3]
        agent_type = batch["agent/type"].float().argmax(dim=-1) + 1.0

        # [n_batch, step_start+1...step_end, n_agent, K, 2] -> downsample -> [n_batch, steps, n_agent, K, 2]
        pred_traj = pred_traj[:, 4 : self.track_future_samples : 5]
        if self.interactive_challenge:
            # [n_batch, 1, K, n_agent, steps, 2]
            pred_traj = pred_traj.permute((0, 3, 2, 1, 4)).unsqueeze(1)
            if pred_score is None:
                pred_score = torch.ones_like(pred_traj[:, :, :, 0, 0, 0]).softmax(-1)
            else:
                # [n_batch, 1, K]
                pred_score = pred_score.sum(dim=1, keepdim=True)
        else:
            # [n_batch, n_agent, K, 1, steps, 2]
            pred_traj = pred_traj.permute((0, 2, 3, 1, 4)).unsqueeze(3)
            if pred_score is None:
                pred_score = torch.ones_like(pred_traj[:, :, :, 0, 0, 0]).softmax(-1)

        n_batch = gt_traj.shape[0]
        n_agent = gt_traj.shape[1]
        n_gt_step = gt_traj.shape[2]
        n_pred_step = pred_traj.shape[-2]
        n_K = pred_traj.shape[2]
        device = pred_traj.device

        prediction_trajectory = torch.zeros(
            [n_batch, self.m_joint, n_K, self.n_pred, n_pred_step, 2], dtype=torch.float32, device=device
        )
        prediction_score = torch.zeros([n_batch, self.m_joint, n_K], dtype=torch.float32, device=device)
        ground_truth_trajectory = torch.zeros([n_batch, n_agent, n_gt_step, 7], dtype=torch.float32, device=device)
        ground_truth_is_valid = torch.zeros([n_batch, n_agent, n_gt_step], dtype=torch.bool, device=device)
        prediction_ground_truth_indices_mask = torch.zeros(
            [n_batch, self.m_joint, self.n_pred], dtype=torch.bool, device=device
        )
        object_type = torch.zeros([n_batch, n_agent], dtype=torch.float32, device=device)

        for i in range(n_batch):
            # reorder and reduce ground_truth_trajectory and ground_truth_is_valid, first pred_agent then other_agent
            n_pred_agent = mask_pred[i].sum()
            n_other_agent = mask_other[i].sum()

            if self.interactive_challenge:
                # pred_traj: [n_batch, 1, K, n_agent, steps, 2]
                prediction_trajectory[i, :, :, :n_pred_agent] = pred_traj[i, :, :, mask_pred[i]]
                prediction_score[i] = pred_score[i]
                prediction_ground_truth_indices_mask[i, :, :n_pred_agent] = True
            else:
                # pred_traj: [n_batch, n_agent, K, 1, steps, 2]
                prediction_trajectory[i, :n_pred_agent] = pred_traj[i, mask_pred[i]]
                prediction_score[i, :n_pred_agent] = pred_score[i][mask_pred[i]]
                prediction_ground_truth_indices_mask[i, :n_pred_agent] = True

            ground_truth_trajectory[i, :n_pred_agent] = gt_traj[i][mask_pred[i]]
            ground_truth_is_valid[i, :n_pred_agent] = gt_valid[i][mask_pred[i]]
            ground_truth_trajectory[i, n_pred_agent : n_pred_agent + n_other_agent] = gt_traj[i][mask_other[i]]
            ground_truth_is_valid[i, n_pred_agent : n_pred_agent + n_other_agent] = gt_valid[i][mask_other[i]]
            object_type[i, :n_pred_agent] = agent_type[i][mask_pred[i]]
            object_type[i, n_pred_agent : n_pred_agent + n_other_agent] = agent_type[i][mask_other[i]]

        self.prediction_trajectory_gpu.append(prediction_trajectory)
        self.prediction_score_gpu.append(prediction_score)
        self.ground_truth_trajectory_gpu.append(ground_truth_trajectory)
        self.ground_truth_is_valid_gpu.append(ground_truth_is_valid)
        self.prediction_ground_truth_indices_mask_gpu.append(prediction_ground_truth_indices_mask)
        self.object_type_gpu.append(object_type)

    def compute(self) -> Dict[str, Tensor]:
        out_dict = {
            "prediction_trajectory": self.prediction_trajectory_gpu,
            "prediction_score": self.prediction_score_gpu,
            "ground_truth_trajectory": self.ground_truth_trajectory_gpu,
            "ground_truth_is_valid": self.ground_truth_is_valid_gpu,
            "prediction_ground_truth_indices_mask": self.prediction_ground_truth_indices_mask_gpu,
            "object_type": self.object_type_gpu,
        }
        return out_dict

    def aggregate_on_cpu(self, gpu_dict_sync: Dict[str, Tensor]) -> None:
        for k, v in gpu_dict_sync.items():
            if type(v) is list:
                assert len(v) == 1
                v = v[0]
            if v.numel() == 1:
                v = v.unsqueeze(0)
                if k == "prediction_ground_truth_indices_mask":
                    v = v[:, None, None]
            self.ops_inputs_cpu[k].append(v.cpu())

    def compute_waymo_motion_metrics(self) -> Dict[str, Tensor]:
        tensorflow.config.set_visible_devices([], "GPU")
        ops_inputs = {}
        for k in self.ops_inputs_cpu.keys():
            ops_inputs[k] = torch.cat(self.ops_inputs_cpu[k], dim=0)
            self.ops_inputs_cpu[k] = []

        if self.interactive_challenge:
            # [1, 1, 2]
            indices = torch.arange(self.n_pred, dtype=torch.int64)[None, None, :]
        else:
            # [1, 8, 1]
            indices = torch.arange(self.m_joint, dtype=torch.int64)[None, :, None]
        # [n_batch, self.m_joint, self.n_pred]
        ops_inputs["prediction_ground_truth_indices"] = indices.expand(ops_inputs["object_type"].shape[0], -1, -1)

        out_dict = {}
        metric_values = py_metrics_ops.motion_metrics(
            config=self.metrics_config,
            prediction_trajectory=ops_inputs["prediction_trajectory"],
            prediction_score=ops_inputs["prediction_score"],
            ground_truth_trajectory=ops_inputs["ground_truth_trajectory"],
            ground_truth_is_valid=ops_inputs["ground_truth_is_valid"],
            prediction_ground_truth_indices_mask=ops_inputs["prediction_ground_truth_indices_mask"],
            object_type=ops_inputs["object_type"],
            prediction_ground_truth_indices=ops_inputs["prediction_ground_truth_indices"],
        )

        for m_type in self.metrics_type:  # e.g. min_ade
            values = np.array(getattr(metric_values, m_type))
            sum_VEHICLE = 0.0
            sum_PEDESTRIAN = 0.0
            sum_CYCLIST = 0.0
            counter_VEHICLE = 0.0
            counter_PEDESTRIAN = 0.0
            counter_CYCLIST = 0.0
            for i, m_name in enumerate(self.metrics_names):  # e.g. TYPE_CYCLIST_15
                out_dict[f"waymo_metrics/{self.prefix}_{m_type}_{m_name}"] = values[i]
                if "VEHICLE" in m_name:
                    sum_VEHICLE += values[i]
                    counter_VEHICLE += 1
                elif "PEDESTRIAN" in m_name:
                    sum_PEDESTRIAN += values[i]
                    counter_PEDESTRIAN += 1
                elif "CYCLIST" in m_name:
                    sum_CYCLIST += values[i]
                    counter_CYCLIST += 1
            out_dict[f"{self.prefix}/{m_type}"] = values.mean()
            out_dict[f"{self.prefix}/veh/{m_type}"] = sum_VEHICLE / counter_VEHICLE
            out_dict[f"{self.prefix}/ped/{m_type}"] = sum_PEDESTRIAN / counter_PEDESTRIAN
            out_dict[f"{self.prefix}/cyc/{m_type}"] = sum_CYCLIST / counter_CYCLIST
        return out_dict

    @staticmethod
    def _waymo_metrics_config_names(
        track_history_samples: int, track_future_samples: int
    ) -> Tuple[motion_metrics_pb2.MotionMetricsConfig, List[str]]:
        config = motion_metrics_pb2.MotionMetricsConfig()
        config_text = f"""
            track_steps_per_second: 10
            prediction_steps_per_second: 2
            track_history_samples: {track_history_samples}
            track_future_samples: {track_future_samples}
            speed_lower_bound: 1.4
            speed_upper_bound: 11.0
            speed_scale_lower: 0.5
            speed_scale_upper: 1.0
            max_predictions: 6
            """
        config_text += """
            step_configurations {
            measurement_step: 5
            lateral_miss_threshold: 1.0
            longitudinal_miss_threshold: 2.0
            }"""
        config_text += """
            step_configurations {
            measurement_step: 9
            lateral_miss_threshold: 1.8
            longitudinal_miss_threshold: 3.6
            }"""
        config_text += """
            step_configurations {
            measurement_step: 15
            lateral_miss_threshold: 3.0
            longitudinal_miss_threshold: 6.0
            }"""
        text_format.Parse(config_text, config)
        metric_names = get_breakdown_names_from_motion_config(config)
        return config, metric_names

src/models/metrics/womd.py
zhejz/TrafficBots
# Licensed under the CC BY-NC 4.0 license (https://creativecommons.org/licenses/by-nc/4.0/)
from typing import Optional, Tuple
from torch import Tensor, nn
import torch
from .distributions import DiagGaussian, MyDist
from .mlp import MLP


class ActionHead(nn.Module):
    def __init__(
        self,
        hidden_dim: int,
        action_dim: int,
        use_layernorm: bool,
        log_std: Optional[float] = None,
        branch_type: bool = False,
    ) -> None:
        super().__init__()
        self.branch_type = branch_type
        self.out_dim = action_dim

        if self.branch_type:
            self.mlp_mean = nn.ModuleList(
                [
                    MLP([hidden_dim, hidden_dim, self.out_dim], end_layer_activation=False, use_layernorm=use_layernorm)
                    for _ in range(3)
                ]
            )
            if log_std is None:
                self.log_std = None
                self.mlp_log_std = nn.ModuleList(
                    [
                        MLP(
                            [hidden_dim, hidden_dim, self.out_dim],
                            end_layer_activation=False,
                            use_layernorm=use_layernorm,
                        )
                        for _ in range(3)
                    ]
                )
            else:
                self.log_std = nn.ParameterList(
                    [nn.Parameter(log_std * torch.ones(self.out_dim), requires_grad=True) for _ in range(3)]
                )

        else:
            self.mlp_mean = MLP(
                [hidden_dim, hidden_dim, self.out_dim], end_layer_activation=False, use_layernorm=use_layernorm
            )
            if log_std is None:
                self.log_std = None
                self.mlp_log_std = MLP(
                    [hidden_dim, hidden_dim, self.out_dim], end_layer_activation=False, use_layernorm=use_layernorm
                )
            else:
                self.log_std = nn.Parameter(log_std * torch.ones(self.out_dim), requires_grad=True)

    def forward(self, x: Tensor, valid: Tensor, agent_type: Tensor) -> Tuple[MyDist]:
        """
        Args:
            x: [n_batch, n_agent, hidden_dim]
            valid: [n_batch, n_agent], bool
            agent_type: [n_batch, n_agent, 3], bool one_hot [Vehicle=0, Pedestrian=1, Cyclist=2]

        Returns:
            action_dist.mean: [n_batch, n_agent, action_dim]
            action_dist.covariance_matrix: [n_batch, n_agent, action_dim, action_dim]
        """
        n_batch, n_agent, n_type = agent_type.shape
        if self.branch_type:
            mask_type = agent_type & valid.unsqueeze(-1)  # [n_batch, n_agent, 3]
            mean = 0
            for i in range(n_type):
                mean += self.mlp_mean[i](x, mask_type[:, :, i])

            log_std = 0
            if self.log_std is None:
                for i in range(n_type):
                    log_std += self.mlp_log_std[i](x, mask_type[:, :, i])
            else:
                for i in range(n_type):
                    # [n_batch, n_agent, self.out_dim]
                    log_std += (
                        self.log_std[i][None, None, :]
                        .expand(n_batch, n_agent, -1)
                        .masked_fill(~mask_type[:, :, [i]], 0)
                    )
        else:
            mean = self.mlp_mean(x, valid)
            if self.log_std is None:
                # [n_batch, n_agent, self.out_dim]
                log_std = self.mlp_log_std(x, valid)
            else:
                # [self.out_dim] -> [n_batch, n_agent, self.out_dim]
                log_std = self.log_std[None, None, :].expand(n_batch, n_agent, -1)

        action_dist = DiagGaussian(mean, log_std)
        return action_dist

src/models/modules/action_head.py
zhejz/TrafficBots
# Licensed under the CC BY-NC 4.0 license (https://creativecommons.org/licenses/by-nc/4.0/)
from typing import Optional
from torch import Tensor, nn
from omegaconf import DictConfig
import torch
from .mlp import MLP


class AddLatentGoal(nn.Module):
    def __init__(
        self,
        hidden_dim: int,
        in_dim: int,
        dummy: bool,
        mode: str,
        n_layer_mlp_in: int,
        n_layer_mlp_out: int,
        mlp_in_cfg: DictConfig,
        mlp_out_cfg: DictConfig,
        res_cat: bool = False,
        res_add: bool = False,
    ) -> None:
        super().__init__()
        assert mode in ["add", "mul", "cat"]
        self.mode = mode
        self.dummy = dummy
        self.res_cat = res_cat
        self.res_add = res_add

        if not self.dummy:
            self.mlp_in = MLP([in_dim] + [hidden_dim] * n_layer_mlp_in, **mlp_in_cfg)

            if self.mode == "cat":
                out_dim = hidden_dim * 2
            else:
                out_dim = hidden_dim

            self.mlp_out = MLP([out_dim] + [hidden_dim] * n_layer_mlp_out, **mlp_out_cfg)

            if self.res_cat:
                self.mlp_res_cat = MLP([hidden_dim * 2 + in_dim] + [hidden_dim] * n_layer_mlp_out, **mlp_out_cfg)

    def forward(self, x: Tensor, x_valid: Tensor, z: Optional[Tensor], z_valid: Optional[Tensor]) -> Tensor:
        """
        Args:
            x: [n_batch, n_agent, hidden_dim]
            x_valid: [n_batch, n_agent]
            z: [n_batch, n_agent, hidden_dim], latent or goal
            z_valid: [n_batch, n_agent]

        Returns:
            h: [n_batch, n_agent, hidden_dim], x combined with z
        """
        if self.dummy:
            h = x
        else:
            z = self.mlp_in(z, z_valid)

            if self.mode == "add":
                h = x + z
            elif self.mode == "mul":
                h = x * z
            else:
                h = torch.cat([x, z], dim=-1)

            h = self.mlp_out(h)

            if self.res_cat:
                h = self.mlp_res_cat(torch.cat([x, h, z], dim=-1))

            h = h.masked_fill(~z_valid.unsqueeze(-1), 0)
            if self.res_add:  # h+x if z_valid else x
                h = h + x
            else:  # h if z_valid else x
                h = h + x.masked_fill(z_valid.unsqueeze(-1), 0)

        return h.masked_fill(~x_valid.unsqueeze(-1), 0)

src/models/modules/add_latent_goal.py
zhejz/TrafficBots
# Licensed under the CC BY-NC 4.0 license (https://creativecommons.org/licenses/by-nc/4.0/)
from typing import Optional, Tuple
import torch
from torch import Tensor, nn
from omegaconf import DictConfig
from .transformer import TransformerBlock


class MultiAgentTF(nn.Module):
    def __init__(
        self,
        hidden_dim: int,
        n_layer: int,
        attn_to_map_aware_feature: bool,
        mask_self_agent: bool,
        detach_tgt: bool,
        tf_cfg: DictConfig,
    ) -> None:
        super().__init__()
        self.hidden_dim = hidden_dim
        self.mask_self_agent = mask_self_agent
        self.attn_mask = None
        self.detach_tgt = detach_tgt
        self.attn_to_map_aware_feature = attn_to_map_aware_feature
        self.transformer = TransformerBlock(n_layer=n_layer, **tf_cfg)

    def forward(
        self, as_feature_map_aware: Tensor, as_feature: Tensor, as_valid: Tensor, need_weights: bool = False
    ) -> Tuple[Tensor, Optional[Tensor]]:
        """
        Args:
            as_feature_map_aware: [n_batch, (n_step), n_agent, hidden_dim]
            as_feature: [n_batch, (n_step), n_agent, hidden_dim]
            as_valid: [n_batch, (n_step), n_agent]

        Returns:
            x: [n_batch, (n_step), n_agent, out_dim]
            attn_weights: None or [n_batch, n_agent, n_agent]
        """
        valid_dim = as_valid.dim()
        if valid_dim == 3:
            n_batch, n_step, n_agent = as_valid.shape
            # [n_batch*n_step, n_agent, hidden_dim]
            as_feature_map_aware = as_feature_map_aware.flatten(start_dim=0, end_dim=1)
            as_feature = as_feature.flatten(start_dim=0, end_dim=1)
            # [n_batch*n_step, n_agent]
            as_valid = as_valid.flatten(start_dim=0, end_dim=1)
        elif valid_dim == 2:
            n_batch, n_agent = as_valid.shape

        x = as_feature_map_aware
        tgt_x = as_feature_map_aware if self.attn_to_map_aware_feature else as_feature
        if self.detach_tgt:
            tgt_x = tgt_x.detach()

        padding = ~as_valid
        if self.mask_self_agent:
            if self.attn_mask is None:
                self.attn_mask = torch.eye(n_agent, device=as_valid.device, dtype=torch.bool)

            invalid_batch = as_valid.sum(-1) == 1  # only one agent valid
            if invalid_batch.any():
                valid_batch = ~invalid_batch
                x_reduced = x[valid_batch]
                tgt_x_reduced = tgt_x[valid_batch]
                padding_reduced = padding[valid_batch]
                x_reduced, attn_weights_reduced = self.transformer(
                    src=x_reduced,
                    src_padding_mask=padding_reduced,
                    tgt=tgt_x_reduced,
                    tgt_padding_mask=padding_reduced,
                    need_weights=need_weights,
                    attn_mask=self.attn_mask,
                )

                x = x.masked_fill(valid_batch[:, None, None], 0.0)
                x[valid_batch] = x_reduced

                if need_weights:
                    n_batch, n_agent = as_valid.shape
                    attn_weights = torch.zeros([n_batch, n_agent, n_agent], device=x.device, dtype=x.dtype)
                    attn_weights[valid_batch] = attn_weights_reduced
                else:
                    attn_weights = None
            else:
                x, attn_weights = self.transformer(
                    src=x,  # [n_batch(*n_step), n_agent, hidden_dim]
                    src_padding_mask=padding,  # [n_batch(*n_step), n_agent]
                    tgt=tgt_x,
                    tgt_padding_mask=padding,
                    need_weights=need_weights,
                    attn_mask=self.attn_mask,
                )
        else:
            x, attn_weights = self.transformer(
                src=x,  # [n_batch(*n_step), n_agent, hidden_dim]
                src_padding_mask=padding,  # [n_batch(*n_step), n_agent]
                tgt=tgt_x,
                tgt_padding_mask=padding,
                need_weights=need_weights,
                attn_mask=None,
            )

        if valid_dim == 3:
            # [n_step, n_batch * n_agent, hidden_dim] -> [n_batch, n_step, n_agent, hidden_dim]
            x = x.view([n_batch, n_step, n_agent, self.hidden_dim])
        return x, attn_weights

src/models/modules/agent_interaction.py
zhejz/TrafficBots
# Licensed under the CC BY-NC 4.0 license (https://creativecommons.org/licenses/by-nc/4.0/)
from typing import Optional, Tuple, Any
import torch
from torch import Tensor, nn


class TemporalAggregate(nn.Module):
    def __init__(self, mode: str) -> None:
        """
        mode in {"max", "last", "max_valid", "last_valid", "mean_valid"}
        """
        super().__init__()
        self.mode = mode

    def forward(self, x: Tensor, valid: Tensor) -> Tuple[Tensor, Tensor]:
        """
        Args:
            x: [n_batch, n_step, n_agent, hidden_dim], invalid x is padded with 0
            valid: [n_batch, n_step, n_agent]
        Returns:
            x_aggregated: [n_batch, n_agent, hidden_dim]
            valid_aggregated: [n_batch, n_agent]
        """
        if self.mode == "max":
            x_aggregated = x.amax(1)
        elif self.mode == "last":
            x_aggregated = x[:, -1]
        elif self.mode == "max_valid":
            x_aggregated = x.masked_fill(~valid.unsqueeze(-1), -1e3).amax(1)
        elif self.mode == "last_valid":
            n_batch, n_step, n_agent = valid.shape
            idx_last_valid = n_step - 1 - torch.max(valid.flip(1), dim=1)[1]
            x_aggregated = x[torch.arange(n_batch).unsqueeze(1), idx_last_valid, torch.arange(n_agent).unsqueeze(0)]
        elif self.mode == "mean_valid":
            valid_sum = valid.sum(1) + torch.finfo(x.dtype).eps
            x_aggregated = x.sum(1) / valid_sum.unsqueeze(-1)

        valid_aggregated = torch.any(valid, axis=1)
        return x_aggregated.masked_fill(~valid_aggregated.unsqueeze(-1), 0), valid_aggregated


class MultiAgentDummy(nn.Module):
    def __init__(self, **kwargs) -> None:
        super().__init__()

    def forward(self, x: Tensor, valid: Tensor, h: Optional[Any] = None) -> Tuple[Tensor, Tensor]:
        return x, h


class MultiAgentGRUCell(nn.Module):
    def __init__(self, hidden_dim: int, num_layers: int, dropout: float = 0.1) -> None:
        super().__init__()
        self.hidden_dim = hidden_dim
        self.num_layers = num_layers
        self.dropout = nn.Dropout(dropout)

        self.layers = nn.ModuleList()
        for i in range(num_layers):
            self.layers.append(nn.GRUCell(hidden_dim, hidden_dim))

    def forward(self, x: Tensor, valid: Tensor, h: Optional[Any] = None) -> Tuple[Tensor, Tensor]:
        """
        Args:
            x: [n_batch, (n_step), n_agent, hidden_dim]
            valid: [n_batch, (n_step), n_agent]
            h: [num_layers, n_batch * n_agent, hidden_dim] or None
        Returns:
            x_1: [n_batch, (n_step), n_agent, hidden_dim]
            h_1: [num_layers, n_batch * n_agent, hidden_dim]
        """
        n_batch = valid.shape[0]
        n_agent = valid.shape[-1]
        if h is None:
            h = torch.zeros((self.num_layers, n_batch * n_agent, self.hidden_dim), device=x.device)

        if valid.dim() == 3:
            x_1 = []
            for k in range(valid.shape[1]):

                input = x[:, k].flatten(start_dim=0, end_dim=1)  # [n_batch * n_agent, hidden_dim]
                h_1 = []
                for i, layer in enumerate(self.layers):
                    h_1_i = layer(input, h[i])
                    input = h_1_i
                    if i + 1 != self.num_layers:
                        input = self.dropout(input)
                    h_1.append(h_1_i)

                invalid = ~valid[:, k].flatten(start_dim=0, end_dim=1).unsqueeze(-1)  # [n_batch * n_agent, 1]
                h = torch.stack(h_1, dim=0)  # [n_layers, n_batch * n_agent, hidden_dim]
                h.masked_fill_(invalid.unsqueeze(0), 0.0)
                x_1.append(input.masked_fill(invalid, 0.0).view(n_batch, n_agent, self.hidden_dim))
            x_1 = torch.stack(x_1, dim=1)  # [n_batch, n_step,  n_agent, hidden_dim]
            h_1 = None
        elif valid.dim() == 2:
            input = x.flatten(start_dim=0, end_dim=1)  # [n_batch * n_agent, hidden_dim]
            h_1 = []
            for i, layer in enumerate(self.layers):
                h_1_i = layer(input, h[i])
                input = h_1_i
                if i + 1 != self.num_layers:
                    input = self.dropout(input)
                h_1.append(h_1_i)
            invalid = ~valid.flatten(start_dim=0, end_dim=1).unsqueeze(-1)  # [n_batch * n_agent, 1]
            h_1 = torch.stack(h_1, dim=0)  # [n_layers, n_batch * n_agent, hidden_dim]
            h_1.masked_fill_(invalid.unsqueeze(0), 0.0)
            x_1 = input.masked_fill(invalid, 0.0).view(n_batch, n_agent, self.hidden_dim)
        return x_1, h_1


class MultiAgentGRULoop(nn.Module):
    def __init__(self, hidden_dim: int, num_layers: int, dropout: float = 0.1) -> None:
        super().__init__()
        self.hidden_dim = hidden_dim
        self.num_layers = num_layers
        self.rnn = nn.GRU(hidden_dim, hidden_dim, num_layers, dropout=dropout)

    def forward(self, x: Tensor, valid: Tensor, h: Optional[Any] = None) -> Tuple[Tensor, Tensor]:
        """
        Args:
            x: [n_batch, (n_step), n_agent, hidden_dim]
            valid: [n_batch, (n_step), n_agent]
            h: [num_layers, n_batch * n_agent, hidden_dim] or None
        Returns:
            x_1: [n_batch, (n_step), n_agent, hidden_dim]
            h_1: [num_layers, n_batch * n_agent, hidden_dim]
        """
        n_batch = valid.shape[0]
        n_agent = valid.shape[-1]
        if h is None:
            h = torch.zeros((self.num_layers, n_batch * n_agent, self.hidden_dim), device=x.device)

        if valid.dim() == 3:
            n_step = valid.shape[1]
            x_1 = []
            # [n_step, n_batch * n_agent, hidden_dim]
            x = x.transpose(0, 1).flatten(start_dim=1, end_dim=2)
            # [n_step, n_batch * n_agent, 1]
            invalid = ~valid.transpose(0, 1).flatten(start_dim=1, end_dim=2).unsqueeze(-1)
            for k in range(n_step):
                x_out, h = self.rnn(x[[k]], h)
                h = h.masked_fill(invalid[[k]], 0.0)
                x_1.append(x_out)  # [1, n_batch * n_agent, hidden_dim]
            x_1 = torch.cat(x_1, dim=0)  # [n_step, n_batch*n_agent, hidden_dim]
            x_1 = x_1.masked_fill(invalid, 0.0).view(n_step, n_batch, n_agent, self.hidden_dim).transpose(0, 1)
            h_1 = None
        elif valid.dim() == 2:
            input = x.flatten(start_dim=0, end_dim=1).unsqueeze(0)
            input, h = self.rnn(input, h)
            invalid = ~valid.flatten(start_dim=0, end_dim=1).unsqueeze(-1)  # [n_batch * n_agent, 1]
            h_1 = h.masked_fill(invalid.unsqueeze(0), 0.0)
            x_1 = input[0].masked_fill(invalid, 0.0).view(n_batch, n_agent, self.hidden_dim)
        return x_1, h_1


class MultiAgentGRU(nn.Module):
    # for TrafficSim where interaction_fist=False, temp_aggr_mode="last", forward not masked
    def __init__(self, hidden_dim: int, num_layers: int, dropout: Optional[float]) -> None:
        super().__init__()
        self.hidden_dim = hidden_dim
        self.num_layers = num_layers
        self.rnn = nn.GRU(hidden_dim, hidden_dim, num_layers, dropout=dropout)

    def forward(self, x: Tensor, valid: Tensor, h: Optional[Any] = None) -> Tuple[Tensor, Tensor]:
        """
        Args:
            x: [n_batch, (n_step), n_agent, hidden_dim]
            valid: [n_batch, (n_step), n_agent]
            h: [num_layers, n_batch * n_agent, hidden_dim] or None
        Returns:
            x: [n_batch, (n_step), n_agent, out_dim]
            h: [num_layers, n_batch * n_agent, hidden_dim]
        """
        if valid.dim() == 3:
            n_batch, n_step, n_agent = valid.shape
        elif valid.dim() == 2:
            n_batch, n_agent = valid.shape
            n_step = 1
            x = x.unsqueeze(dim=1)

        if h is None:
            h = torch.zeros((self.num_layers, n_batch * n_agent, self.hidden_dim), device=x.device)

        x, h = self.rnn((x.transpose(0, 1)).flatten(start_dim=1, end_dim=2), h)

        # [n_step, n_batch * n_agent, hidden_dim] -> [n_batch, n_step, n_agent, hidden_dim]
        x = x.view([n_step, n_batch, n_agent, self.hidden_dim]).transpose(0, 1)

        if valid.dim() == 2:
            x = x.squeeze(axis=1)
            # x = x.masked_fill(~valid.unsqueeze(-1), 0.0)

        return x, h

src/models/modules/agent_temporal.py
zhejz/TrafficBots
# Licensed under the CC BY-NC 4.0 license (https://creativecommons.org/licenses/by-nc/4.0/)
from typing import Optional, Tuple
import math
import torch
from torch import Tensor, nn
from torch.nn import functional as F

# refer to torch.nn.MultiheadAttention
class Attention(nn.Module):
    def __init__(self, d_model: int, n_head: int, dropout_p: float = 0.0, bias: bool = True) -> None:
        """
        Always batch first. Always src and tgt have the same d_model.
        """
        super().__init__()

        self.d_model = d_model
        self.n_head = n_head
        self.d_head = d_model // n_head

        assert self.d_head * n_head == d_model, "d_model must be divisible by n_head"

        self.in_proj_weight = nn.Parameter(torch.empty((3 * d_model, d_model)))
        self.out_proj_weight = nn.Parameter(torch.empty((d_model, d_model)))
        if bias:
            self.in_proj_bias = nn.Parameter(torch.empty(3 * d_model))
            self.out_proj_bias = nn.Parameter(torch.empty(d_model))
        else:
            self.register_parameter("in_proj_bias", None)
            self.register_parameter("out_proj_bias", None)

        self.dropout = nn.Dropout(p=dropout_p, inplace=False) if dropout_p > 0 else None

        self._reset_parameters()

    def _reset_parameters(self):
        nn.init.xavier_uniform_(self.in_proj_weight)
        nn.init.xavier_uniform_(self.out_proj_weight)
        if self.in_proj_bias is not None:
            nn.init.constant_(self.in_proj_bias, 0.0)
        if self.out_proj_bias is not None:
            nn.init.constant_(self.out_proj_bias, 0.0)

        # TODO: maybe more sophisticated init, could also use nn.init.xavier_normal_(self.bias_k)
        # nn.init.xavier_uniform_(self.in_proj_weight)
        # nn.init.kaiming_uniform_(self.out_proj_weight, a=math.sqrt(5))
        # fan_in, _ = nn.init._calculate_fan_in_and_fan_out(self.out_proj_weight)
        # bound = 1 / math.sqrt(fan_in) if fan_in > 0 else 0
        # nn.init.uniform_(self.out_proj_bias, -bound, bound)
        # if self.in_proj_bias is not None:
        #     nn.init.constant_(self.in_proj_bias, 0.0)
        #     nn.init.constant_(self.out_proj_bias, 0.0)

    def forward(
        self,
        src: Tensor,
        tgt: Optional[Tensor] = None,
        tgt_padding_mask: Optional[Tensor] = None,
        attn_mask: Optional[Tensor] = None,
        need_weights=False,
    ) -> Tuple[Tensor, Optional[Tensor]]:
        """
        Args:
            src: [n_batch, n_src, d_model]
            tgt: [n_batch, (n_src), n_tgt, d_model], None for self attention, (n_src) if using knn.
            tgt_padding_mask: [n_batch, (n_src), n_tgt], bool, if True, tgt is invalid, (n_src) if using knn.
            attn_mask: [n_batch, n_src, n_tgt], bool, if True, attn is disabled for that pair of src/tgt.

        Returns:
            out: [n_batch, n_src, d_model]
            attn_weights: [n_batch, n_src, n_tgt] if need_weights else None
        """
        n_batch, n_src, _ = src.shape
        if tgt is None:
            n_tgt = n_src
            # self-attention
            qkv = F.linear(src, self.in_proj_weight, self.in_proj_bias)
            q, k, v = qkv.chunk(3, dim=-1)
        else:
            n_tgt = tgt.shape[-2]
            # encoder-decoder attention
            w_src, w_tgt = self.in_proj_weight.split([self.d_model, self.d_model * 2])
            b_src, b_tgt = None, None
            if self.in_proj_bias is not None:
                b_src, b_tgt = self.in_proj_bias.split([self.d_model, self.d_model * 2])
            q = F.linear(src, w_src, b_src)
            kv = F.linear(tgt, w_tgt, b_tgt)
            k, v = kv.chunk(2, dim=-1)
        # q: [n_batch, n_src, d_model], k,v: [n_batch, (n_src), n_tgt, d_model]

        attn_invalid_mask = None  # [n_batch, n_src, n_tgt]
        if tgt_padding_mask is not None:  # [n_batch, n_tgt], bool
            attn_invalid_mask = tgt_padding_mask
            if attn_invalid_mask.dim() == 2:
                attn_invalid_mask = attn_invalid_mask.unsqueeze(1).expand(-1, n_src, -1)
        if attn_mask is not None:  # [n_batch, n_src, n_tgt], bool
            if attn_invalid_mask is None:
                attn_invalid_mask = attn_mask
            else:
                attn_invalid_mask = attn_invalid_mask | attn_mask

        mask_no_tgt_valid = None  # [n_batch, n_src]
        if attn_invalid_mask is not None:
            mask_no_tgt_valid = attn_invalid_mask.all(-1)
            if mask_no_tgt_valid.any():
                attn_invalid_mask = attn_invalid_mask & (~mask_no_tgt_valid.unsqueeze(-1))  # to avoid softmax nan
            else:
                mask_no_tgt_valid = None

        # get attn: [n_batch, n_head, n_src, n_tgt]
        if k.dim() == 3:
            # ! normal attention; q: [n_batch, n_src, d_model], k,v: [n_batch, n_tgt, d_model]
            q = q.view(n_batch, n_src, self.n_head, self.d_head).transpose(1, 2).contiguous()
            k = k.view(n_batch, n_tgt, self.n_head, self.d_head).transpose(1, 2).contiguous()
            v = v.view(n_batch, n_tgt, self.n_head, self.d_head).transpose(1, 2).contiguous()
            attn = torch.matmul(q, k.transpose(-2, -1))  # [n_batch, n_head, n_src, n_tgt]
            # q: [n_batch, n_head, n_src, d_head], k,v: [n_batch, n_head, n_tgt, d_head]
        else:
            # ! KNN attention; q: [n_batch, n_src, d_model], k,v: [n_batch, n_src, n_tgt, d_model]
            # k,v: [n_batch, n_src, n_tgt, d_model] -> [n_batch, n_head, n_src, n_tgt_knn, d_head]
            k = k.view(n_batch, n_src, n_tgt, self.n_head, self.d_head).movedim(3, 1)
            v = v.view(n_batch, n_src, n_tgt, self.n_head, self.d_head).movedim(3, 1)
            # [n_batch, n_src, d_model] -> [n_batch, n_head, n_src, 1, d_head]
            q = q.view(n_batch, n_src, self.n_head, self.d_head).transpose(1, 2).unsqueeze(3)
            attn = torch.sum(q * k, dim=-1)  # [n_batch, n_head, n_src, n_tgt_knn]

        if attn_invalid_mask is not None:
            # attn_invalid_mask: [n_batch, n_src, n_tgt], attn: [n_batch, n_head, n_src, n_tgt]
            attn = attn.masked_fill(attn_invalid_mask.unsqueeze(1), float("-inf"))

        attn = torch.softmax(attn / math.sqrt(self.d_head), dim=-1)
        if self.dropout is not None:
            attn = self.dropout(attn)

        # attn: [n_batch, n_head, n_src, n_tgt]
        if v.dim() == 4:
            out = torch.matmul(attn, v)  # v, [n_batch, n_head, n_tgt, d_head]
        else:
            out = torch.sum(v * attn.unsqueeze(-1), dim=3)  # v: [n_batch, n_head, n_src, n_tgt, d_head]

        # out: [n_batch, n_head, n_src, d_head]
        out = out.transpose(1, 2).flatten(2, 3)  # [n_batch, n_src, d_model]
        out = F.linear(out, self.out_proj_weight, self.out_proj_bias)

        if mask_no_tgt_valid is not None:
            # mask_no_tgt_valid: [n_batch, n_src], out: [n_batch, n_src, d_model]
            out = out.masked_fill(mask_no_tgt_valid.unsqueeze(-1), 0)

        if need_weights:
            attn_weights = attn.mean(1)  # [n_batch, n_src, n_tgt]
            if mask_no_tgt_valid is not None:
                attn_weights = attn_weights.masked_fill(mask_no_tgt_valid.unsqueeze(-1), 0)
            return out, attn_weights
        else:
            return out, None

src/models/modules/attention.py
zhejz/TrafficBots
# Licensed under the CC BY-NC 4.0 license (https://creativecommons.org/licenses/by-nc/4.0/)
from typing import Optional, Union
from torch import Tensor
import torch
from torch.distributions import Independent, Normal, MultivariateNormal, OneHotCategoricalStraightThrough, Categorical
from torch.nn import functional as F


class MyDist:
    def __init__(self, *args, **kwargs) -> None:
        self.distribution = None

    def log_prob(self, sample: Tensor) -> Tensor:
        """
        log_prob: [n_batch, n_agent]
        """
        return self.distribution.log_prob(sample)

    def sample(self, deterministic: Union[bool, Tensor]) -> Tensor:
        """
        Args:
            deterministic: bool, or Tensor for sampling relevant agents and determistic other agents.
        Returns:
            sample: [n_batch, n_agent, out_dim]
        """
        if type(deterministic) is Tensor:
            det_sample = self.distribution.mean
            rnd_sample = self.distribution.rsample()
            sample = det_sample.masked_fill(~deterministic.unsqueeze(-1), 0) + rnd_sample.masked_fill(
                deterministic.unsqueeze(-1), 0
            )
        else:
            if deterministic:
                sample = self.distribution.mean
            else:
                sample = self.distribution.rsample()
        return sample


class DiagGaussian(MyDist):
    def __init__(self, mean: Tensor, log_std: Tensor, valid: Optional[Tensor] = None) -> None:
        """
        mean: [n_batch, n_agent, (k_pred), out_dim]
        covariance_matrix: [n_batch, n_agent, (k_pred), out_dim, out_dim]
        """
        super().__init__()
        self.mean = mean
        self.valid = valid
        self.distribution = Independent(Normal(self.mean, log_std.exp()), 1)
        self.stddev = self.distribution.stddev
        self.covariance_matrix = torch.diag_embed(self.distribution.variance)

    def repeat_interleave_(self, repeats: int, dim: int) -> None:
        self.mean = self.mean.repeat_interleave(repeats, dim)
        self.stddev = self.stddev.repeat_interleave(repeats, dim)
        self.distribution = Independent(Normal(self.mean, self.stddev), 1)
        self.covariance_matrix = torch.diag_embed(self.distribution.variance)
        if self.valid is not None:
            self.valid = self.valid.repeat_interleave(repeats, dim)


class Gaussian(MyDist):
    def __init__(self, mean: Tensor, tril: Tensor) -> None:
        """
        mean: [n_batch, n_agent, (k_pred), out_dim]
        tril: [out_dim, out_dim], lower-triangular factor of covariance, with positive-valued diagonal, cov=LL^T
        covariance_matrix: [n_batch, n_agent, (k_pred), out_dim, out_dim]
        """
        super().__init__()
        self.mean = mean
        self.tril = tril
        self.distribution = MultivariateNormal(self.mean, scale_tril=self.tril)
        self.mean = self.distribution.mean
        self.covariance_matrix = self.distribution.covariance_matrix

    def repeat_interleave_(self, repeats: int, dim: int) -> None:
        self.mean = self.mean.repeat_interleave(repeats, dim)
        self.tril = self.tril.repeat_interleave(repeats, dim)
        self.distribution = MultivariateNormal(self.mean, scale_tril=self.tril)
        self.mean = self.distribution.mean
        self.covariance_matrix = self.distribution.covariance_matrix


class DummyLatent(MyDist):
    def __init__(self, x, valid) -> None:
        super().__init__()
        self._logp = torch.zeros_like(x[..., 0])
        self._sample = torch.zeros_like(x)
        self.valid = valid

    def log_prob(self, *args, **kwargs) -> Tensor:
        return self._logp

    def sample(self, *args, **kwargs) -> Tensor:
        return self._sample

    def repeat_interleave_(self, repeats: int, dim: int) -> None:
        self._logp = self._logp.repeat_interleave(repeats, dim)
        self._sample = self._sample.repeat_interleave(repeats, dim)


class MultiCategorical(MyDist):
    def __init__(self, probs: Tensor, valid: Optional[Tensor] = None):
        """
        probs: [n_batch, n_agent, n_cat, n_class]
        """
        super().__init__()
        self.probs = probs
        self.distribution = Independent(OneHotCategoricalStraightThrough(probs=self.probs), 1)
        self.n_cat = self.probs.shape[-2]
        self.n_class = self.probs.shape[-1]
        self._dtype = self.probs.dtype
        self.valid = valid

    def log_prob(self, sample: Tensor) -> Tensor:
        # [n_batch, n_agent]
        return self.distribution.log_prob(sample.view(*sample.shape[:-1], self.n_cat, self.n_class))

    def sample(self, deterministic: Union[bool, Tensor]) -> Tensor:
        """
        Args:
            deterministic: bool, or Tensor for sampling relevant agents and determistic other agents.
        Returns:
            sample: [n_batch, n_agent, out_dim]
        """
        # [n_batch, n_agent, n_cat, n_class]
        if type(deterministic) is Tensor:
            det_sample = (
                F.one_hot(self.distribution.base_dist.probs.argmax(-1), num_classes=self.n_class)
                .type(self._dtype)
                .flatten(start_dim=-2, end_dim=-1)
            )
            rnd_sample = self.distribution.rsample().flatten(start_dim=-2, end_dim=-1)
            sample = det_sample.masked_fill(~deterministic.unsqueeze(-1), 0) + rnd_sample.masked_fill(
                deterministic.unsqueeze(-1), 0
            )
        else:
            if deterministic:
                sample = (
                    F.one_hot(self.distribution.base_dist.probs.argmax(-1), num_classes=self.n_class)
                    .type(self._dtype)
                    .flatten(start_dim=-2, end_dim=-1)
                )
            else:
                sample = self.distribution.rsample().flatten(start_dim=-2, end_dim=-1)
        return sample

    def repeat_interleave_(self, repeats: int, dim: int) -> None:
        self.probs = self.probs.repeat_interleave(repeats, dim)
        self.distribution = Independent(OneHotCategoricalStraightThrough(probs=self.probs), 1)
        self.n_cat = self.probs.shape[-2]
        self.n_class = self.probs.shape[-1]
        self._dtype = self.probs.dtype
        if self.valid is not None:
            self.valid = self.valid.repeat_interleave(repeats, dim)


class DestCategorical(MyDist):
    def __init__(self, probs: Optional[Tensor] = None, logits: Optional[Tensor] = None, valid: Optional[Tensor] = None):
        """
        probs: [n_batch, n_agent, n_pl] >= 0, sum up to 1.
        """
        super().__init__()
        if probs is None:
            assert logits is not None
            self.distribution = Categorical(logits=logits)
            self.probs = self.distribution.probs
        else:
            assert probs is not None
            self.distribution = Categorical(probs=probs)
            self.probs = self.distribution.probs

        self.valid = valid

    def log_prob(self, sample: Tensor) -> Tensor:
        # [n_batch, n_agent]
        return self.distribution.log_prob(sample)

    def sample(self, deterministic: Union[bool, Tensor]) -> Tensor:
        """
        Args:
            deterministic: bool, or Tensor for sampling relevant agents and determistic other agents.
        Returns:
            sample: [n_batch, n_agent, out_dim]
        """
        if type(deterministic) is Tensor:
            det_sample = self.distribution.probs.argmax(-1)
            rnd_sample = self.distribution.sample()
            sample = det_sample.masked_fill(~deterministic, 0) + rnd_sample.masked_fill(deterministic, 0)
        else:
            if deterministic:
                sample = self.distribution.probs.argmax(-1)
            else:
                sample = self.distribution.sample()
        return sample

    def repeat_interleave_(self, repeats: int, dim: int) -> None:
        self.probs = self.probs.repeat_interleave(repeats, dim)
        self.distribution = Categorical(probs=self.probs)
        if self.valid is not None:
            self.valid = self.valid.repeat_interleave(repeats, dim)

src/models/modules/distributions.py
zhejz/TrafficBots
# Licensed under the CC BY-NC 4.0 license (https://creativecommons.org/licenses/by-nc/4.0/)
from typing import Optional
import torch
from torch import Tensor, nn
from .mlp import MLP


class InputPeEncoder(nn.Module):
    def __init__(
        self,
        hidden_dim: int,
        attr_dim: int,
        pe_dim: int,
        n_layer: int,
        mlp_dropout_p: Optional[float] = 0.1,
        mlp_use_layernorm: bool = True,
        pe_mode: str = "add",
    ) -> None:
        super().__init__()
        self.pe_mode = pe_mode

        if self.pe_mode == "input":
            mlp_in_dim = attr_dim + pe_dim
            mlp_out_dim = hidden_dim
        elif self.pe_mode == "cat":
            mlp_in_dim = attr_dim
            mlp_out_dim = hidden_dim - pe_dim
            assert mlp_out_dim >= 32, f"Make sure pe_dim is smaller than {hidden_dim-32}!"
        elif self.pe_mode == "add":
            mlp_in_dim = attr_dim
            mlp_out_dim = hidden_dim
            assert pe_dim == hidden_dim, f"Make sure pe_dim equals to hidden_dim={hidden_dim}!"

        self.mlp = MLP(
            [mlp_in_dim] + [mlp_out_dim] * n_layer,
            dropout_p=mlp_dropout_p,
            use_layernorm=mlp_use_layernorm,
            end_layer_activation=False,
        )

    def forward(self, valid: Tensor, attr: Tensor, pe: Tensor) -> Tensor:
        """
        Args:
            valid: [...], bool
            attr: [..., attr_dim], for input to MLP
            pe: [..., hidden_dim], for input/cat/add to MLP
        Returns:
            feature: [..., hidden_dim] float32
        """
        if self.pe_mode == "input":
            x = self.mlp(torch.cat([attr, pe], dim=-1))
        elif self.pe_mode == "cat":
            x = self.mlp(attr)
            x = torch.cat([x, pe], dim=-1)
        elif self.pe_mode == "add":
            x = self.mlp(attr)
            x = x + pe

        feature = x.masked_fill(~valid.unsqueeze(-1), 0)

        return feature

src/models/modules/input_pe_encoder.py
zhejz/TrafficBots
# Licensed under the CC BY-NC 4.0 license (https://creativecommons.org/licenses/by-nc/4.0/)
from typing import Tuple, List, Optional
import torch
from torch import Tensor, nn
from omegaconf import DictConfig
from .mlp import MLP
from .transformer import TransformerBlock
from .input_pe_encoder import InputPeEncoder


class MapEncoder(nn.Module):
    def __init__(
        self,
        hidden_dim: int,
        attr_dim: int,
        pe_dim: int,
        input_pe_encoder: DictConfig,
        tf_cfg: DictConfig,
        densetnt_vectornet: bool = False,
        pool_mode: str = "max",  # max, mean, first
        n_layer: int = 3,
        mlp_dropout_p: Optional[float] = 0.1,
        mlp_use_layernorm: bool = True,
    ) -> None:
        super().__init__()
        self.densetnt_vectornet = densetnt_vectornet
        self.pool_mode = pool_mode

        self.input_pe_encoder = InputPeEncoder(
            hidden_dim=hidden_dim, attr_dim=attr_dim, pe_dim=pe_dim, **input_pe_encoder
        )

        if self.densetnt_vectornet:
            self.transformer_densetnt = TransformerBlock(n_layer=n_layer, **tf_cfg)
        else:
            mlp_layers: List[nn.Module] = []
            for _ in range(n_layer - 1):
                mlp_layers.append(
                    MLP([hidden_dim, hidden_dim // 2], dropout_p=mlp_dropout_p, use_layernorm=mlp_use_layernorm)
                )

            if tf_cfg.norm_first:
                end_layer_activation = False
            else:
                end_layer_activation = True
            mlp_layers.append(
                MLP(
                    [hidden_dim, hidden_dim // 2],
                    dropout_p=mlp_dropout_p,
                    use_layernorm=mlp_use_layernorm,
                    end_layer_activation=end_layer_activation,
                )
            )
            self.mlp_layers = nn.ModuleList(mlp_layers)

        self.transformer_self_attn = TransformerBlock(n_layer=1, **tf_cfg)

    def forward(self, map_valid: Tensor, map_attr: Tensor, map_pe: Tensor) -> Tuple[Tensor, Tensor]:
        """
        c.f. VectorNet and SceneTransformer
        Aggregate polyline-level feature. n_pl polylines, n_node nodes per polyline.
        Args:
            map_valid: [n_scene, n_pl, n_pl_node], bool
            map_attr: [n_scene, n_pl, n_pl_node, map_attr_dim]
            map_pe: [n_scene, n_pl, n_pl_node, hidden_dim], for input/cat/add to MLP

        Returns:
            map_feature: float32, [n_scene, n_pl, hidden_dim]
            map_valid: bool, [n_scene, n_pl]
        """
        n_scene, n_pl, n_node = map_valid.shape
        pl_feature = self.input_pe_encoder(map_valid, map_attr, map_pe)

        if self.densetnt_vectornet:
            # map_valid: bool, [n_scene, n_pl, n_node]
            # transformer
            pl_feature = pl_feature.flatten(0, 1)
            map_valid = map_valid.flatten(0, 1)
            pl_feature, _ = self.transformer_densetnt(
                src=pl_feature,  # [n_scene*n_pl, n_node, hidden_dim]
                src_padding_mask=~map_valid,  # [n_scene*n_pl, n_node]
                tgt=pl_feature,  # [n_scene*n_pl, n_node, hidden_dim]
                tgt_padding_mask=~map_valid,  # [n_scene*n_pl, n_node]
                need_weights=False,
            )
            hidden_dim = pl_feature.shape[-1]
            pl_feature = pl_feature.view(n_scene, n_pl, n_node, hidden_dim)
            map_valid = map_valid.view(n_scene, n_pl, n_node)
        else:
            for mlp in self.mlp_layers:
                feature_encoded = mlp(pl_feature, map_valid, float("-inf"))
                feature_pooled = feature_encoded.amax(dim=2, keepdim=True)
                pl_feature = torch.cat((feature_encoded, feature_pooled.expand(-1, -1, n_node, -1)), dim=-1)

        if self.pool_mode == "max":
            pl_feature = pl_feature.masked_fill(~map_valid.unsqueeze(-1), float("-inf"))
            pl_feature = pl_feature.amax(dim=2, keepdim=False)  # [batch_size, n_pl, :]
        elif self.pool_mode == "first":
            pl_feature = pl_feature[:, :, 0]
        elif self.pool_mode == "mean":
            pl_feature = pl_feature.masked_fill(~map_valid.unsqueeze(-1), 0)
            pl_feature = pl_feature.sum(dim=2, keepdim=False)  # [batch_size, n_pl, :]
            pl_feature = pl_feature / (map_valid.sum(dim=-1, keepdim=True) + torch.finfo(pl_feature.dtype).eps)

        pl_valid = map_valid.any(-1)  # [batch_size, n_pl]
        pl_feature = pl_feature.masked_fill(~pl_valid.unsqueeze(-1), 0)

        pl_feature, _ = self.transformer_self_attn(
            src=pl_feature,  # [n_scene, n_pl, hidden_dim]
            src_padding_mask=~pl_valid,  # [n_scene, n_pl]
            tgt=pl_feature,  # [n_scene, n_pl, hidden_dim]
            tgt_padding_mask=~pl_valid,  # [n_scene, c_step]
            need_weights=False,
        )
        return pl_feature, pl_valid

src/models/modules/map_encoder.py
zhejz/TrafficBots
# Licensed under the CC BY-NC 4.0 license (https://creativecommons.org/licenses/by-nc/4.0/)
from typing import List, Tuple, Union, Optional
from torch import Tensor, nn


def _get_activation(activation: str, inplace: bool) -> nn.Module:
    if activation == "relu":
        return nn.ReLU(inplace=inplace)
    elif activation == "gelu":
        return nn.GELU()
    elif activation == "leaky_relu":
        return nn.LeakyReLU(inplace=inplace)
    elif activation == "elu":
        return nn.ELU(inplace=inplace)
    elif activation == "rrelu":
        return nn.RReLU(inplace=inplace)
    raise RuntimeError("activation {} not implemented".format(activation))


class MLP(nn.Module):
    def __init__(
        self,
        fc_dims: Union[List, Tuple],
        dropout_p: Optional[float] = None,
        use_layernorm: bool = False,
        activation: str = "relu",
        end_layer_activation: bool = True,
        init_weight_norm: bool = False,
        init_bias: Optional[float] = None,
        use_batchnorm: bool = False,
    ) -> None:
        super(MLP, self).__init__()
        assert len(fc_dims) >= 2
        assert not (use_layernorm and use_batchnorm)
        layers: List[nn.Module] = []
        for i in range(0, len(fc_dims) - 1):

            fc = nn.Linear(fc_dims[i], fc_dims[i + 1])

            if init_weight_norm:
                fc.weight.data *= 1.0 / fc.weight.norm(dim=1, p=2, keepdim=True)
            if init_bias is not None and i == len(fc_dims) - 2:
                fc.bias.data *= 0
                fc.bias.data += init_bias

            layers.append(fc)

            if i < len(fc_dims) - 2:
                if use_layernorm:
                    layers.append(nn.LayerNorm(fc_dims[i + 1]))
                elif use_batchnorm:
                    layers.append(nn.BatchNorm1d(fc_dims[i + 1]))
                if dropout_p is not None:
                    layers.append(nn.Dropout(p=dropout_p))
                layers.append(_get_activation(activation, inplace=True))
            if i == len(fc_dims) - 2:
                if end_layer_activation:
                    if use_layernorm:
                        layers.append(nn.LayerNorm(fc_dims[i + 1]))
                    elif use_batchnorm:
                        layers.append(nn.BatchNorm1d(fc_dims[i + 1]))
                    if dropout_p is not None:
                        layers.append(nn.Dropout(p=dropout_p))
                    self.end_layer_activation = _get_activation(activation, inplace=True)
                else:
                    self.end_layer_activation = None

        self.input_dim = fc_dims[0]
        self.output_dim = fc_dims[-1]
        self.fc_layers = nn.Sequential(*layers)

    def forward(self, x: Tensor, valid_mask: Optional[Tensor] = None, fill_invalid: float = 0.0) -> Tensor:
        """
        Args:
            x: [..., input_dim]
            valid_mask: [...]
        Returns:
            x: [..., output_dim]
        """
        x = self.fc_layers(x.flatten(0, -2)).view(*x.shape[:-1], self.output_dim)
        if valid_mask is not None:
            x.masked_fill_(~valid_mask.unsqueeze(-1), fill_invalid)
        if self.end_layer_activation is not None:
            self.end_layer_activation(x)
        return x

src/models/modules/mlp.py
zhejz/TrafficBots
# Licensed under the CC BY-NC 4.0 license (https://creativecommons.org/licenses/by-nc/4.0/)
from typing import Optional, Tuple
from torch import Tensor, nn
from torch.nn import functional as F
from .attention import Attention


def _get_activation_fn(activation):
    if activation == "relu":
        return F.relu
    elif activation == "gelu":
        return F.gelu
    elif activation == "elu":
        return F.elu
    raise RuntimeError("activation should be relu/gelu/elu, not {}".format(activation))


class TransformerBlock(nn.Module):
    __constants__ = ["norm"]

    def __init__(
        self,
        d_model: int,
        n_head: int = 2,
        d_feedforward: int = 256,
        dropout_p: float = 0.1,
        activation: str = "relu",
        n_layer: int = 1,
        norm_first: bool = True,
        decoder_self_attn: bool = False,
        bias: bool = True,
        out_layernorm: bool = False,
    ) -> None:
        super(TransformerBlock, self).__init__()
        self.layers = nn.ModuleList(
            [
                TransformerCrossAttention(
                    d_model=d_model,
                    n_head=n_head,
                    d_feedforward=d_feedforward,
                    dropout_p=dropout_p,
                    activation=activation,
                    norm_first=norm_first,
                    decoder_self_attn=decoder_self_attn,
                    bias=bias,
                )
                for _ in range(n_layer)
            ]
        )

        self.out_layernorm = nn.LayerNorm(d_model) if out_layernorm else None

    def forward(
        self,
        src: Tensor,
        src_padding_mask: Optional[Tensor] = None,
        tgt: Optional[Tensor] = None,
        tgt_padding_mask: Optional[Tensor] = None,
        decoder_tgt: Optional[Tensor] = None,
        decoder_tgt_padding_mask: Optional[Tensor] = None,
        attn_mask: Optional[Tensor] = None,
        need_weights: bool = False,
    ) -> Tuple[Tensor, Optional[Tensor]]:
        """
        Args:
            src: [n_batch, n_src, d_model]
            src_padding_mask: [n_batch, n_src], bool, if True, src is invalid.
            tgt: [n_batch, (n_src), n_tgt, d_model], None for self attention, (n_src) if using knn.
            tgt_padding_mask: [n_batch, (n_src), n_tgt], bool, if True, tgt is invalid, (n_src) if using knn.
            decoder_tgt: [n_batch, (n_src), n_tgt_decoder, d_model], (n_src) if using decoder_knn.
            decoder_tgt_padding_mask: [n_batch, (n_src), n_tgt_decoder], (n_src) if using decoder_knn.
            attn_mask: [n_batch, n_src, n_tgt], bool, if True, attn is disabled for that pair of src/tgt.

        Returns:
            src: [n_batch, n_src, d_model]
            attn_weights: [n_batch, n_src, n_tgt] if need_weights else None

        Remarks:
            absoulte_pe should be already added to src/tgt.
        """
        attn_weights = None
        for mod in self.layers:
            src, attn_weights = mod(
                src=src,
                src_padding_mask=src_padding_mask,
                tgt=tgt,
                tgt_padding_mask=tgt_padding_mask,
                decoder_tgt=decoder_tgt,
                decoder_tgt_padding_mask=decoder_tgt_padding_mask,
                attn_mask=attn_mask,
                need_weights=need_weights,
            )
        if self.out_layernorm is not None:
            src = self.out_layernorm(src)
        return src, attn_weights


class TransformerCrossAttention(nn.Module):
    def __init__(
        self,
        d_model: int,
        n_head: int,
        d_feedforward: int,
        dropout_p: float,
        activation: str,
        norm_first: bool,
        decoder_self_attn: bool,
        bias: bool,
    ) -> None:
        super().__init__()
        self.norm_first = norm_first
        self.d_feedforward = d_feedforward
        self.decoder_self_attn = decoder_self_attn
        inplace = False

        self.dropout = nn.Dropout(p=dropout_p, inplace=inplace) if dropout_p > 0 else None
        self.activation = _get_activation_fn(activation)
        self.norm1 = nn.LayerNorm(d_model)

        if self.decoder_self_attn:
            self.attn_src = Attention(d_model=d_model, n_head=n_head, dropout_p=dropout_p, bias=bias)
            self.norm_src = nn.LayerNorm(d_model)
            self.dropout_src = nn.Dropout(p=dropout_p, inplace=inplace) if dropout_p > 0 else None

        if self.norm_first:
            self.norm_tgt = nn.LayerNorm(d_model)

        self.attn = Attention(d_model=d_model, n_head=n_head, dropout_p=dropout_p, bias=bias)
        if self.d_feedforward > 0:
            self.linear1 = nn.Linear(d_model, d_feedforward)
            self.linear2 = nn.Linear(d_feedforward, d_model)
            self.norm2 = nn.LayerNorm(d_model)
            self.dropout1 = nn.Dropout(p=dropout_p, inplace=inplace) if dropout_p > 0 else None
            self.dropout2 = nn.Dropout(p=dropout_p, inplace=inplace) if dropout_p > 0 else None

    def forward(
        self,
        src: Tensor,
        src_padding_mask: Optional[Tensor] = None,
        tgt: Optional[Tensor] = None,
        tgt_padding_mask: Optional[Tensor] = None,
        decoder_tgt: Optional[Tensor] = None,
        decoder_tgt_padding_mask: Optional[Tensor] = None,
        attn_mask: Optional[Tensor] = None,
        need_weights: bool = False,
    ) -> Tuple[Tensor, Optional[Tensor]]:
        """
        Args:
            src: [n_batch, n_src, d_model]
            src_padding_mask: [n_batch, n_src], bool, if True, src is invalid.
            tgt: [n_batch, (n_src), n_tgt, d_model], None for self attention, (n_src) if using knn.
            tgt_padding_mask: [n_batch, (n_src), n_tgt], bool, if True, tgt is invalid, (n_src) if using knn.
            decoder_tgt: [n_batch, n_src, n_tgt_decoder, d_model], when use decoder_knn
            decoder_tgt_padding_mask: [n_batch, n_src, n_tgt_decoder], when use decoder_knn
            attn_mask: [n_batch, n_src, n_tgt], bool, if True, attn is disabled for that pair of src/tgt.

        Returns:
            out: [n_batch, n_src, d_model]
            attn_weights: [n_batch, n_src, n_tgt] if need_weights else None
        """
        if self.decoder_self_attn:
            # transformer decoder
            if self.norm_first:
                _s = self.norm_src(src)
                if decoder_tgt is None:
                    _s = self.attn_src(_s, tgt_padding_mask=src_padding_mask)[0]
                else:
                    decoder_tgt = self.norm_src(decoder_tgt)
                    _s = self.attn_src(_s, decoder_tgt, tgt_padding_mask=decoder_tgt_padding_mask)[0]

                if self.dropout_src is None:
                    src = src + _s
                else:
                    src = src + self.dropout_src(_s)
            else:
                if decoder_tgt is None:
                    _s = self.attn_src(src, tgt_padding_mask=src_padding_mask)[0]
                else:
                    _s = self.attn_src(src, decoder_tgt, tgt_padding_mask=decoder_tgt_padding_mask)[0]

                if self.dropout_src is None:
                    src = self.norm_src(src + _s)
                else:
                    src = self.norm_src(src + self.dropout_src(_s))

        if tgt is None:
            tgt_padding_mask = src_padding_mask

        if self.norm_first:
            src2 = self.norm1(src)
            if tgt is not None:
                tgt = self.norm_tgt(tgt)
        else:
            src2 = src

        # [n_batch, n_src, d_model]
        src2, attn_weights = self.attn(
            src=src2, tgt=tgt, tgt_padding_mask=tgt_padding_mask, attn_mask=attn_mask, need_weights=need_weights
        )

        if self.d_feedforward > 0:
            if self.dropout1 is None:
                src = src + src2
            else:
                src = src + self.dropout1(src2)

            if self.norm_first:
                src2 = self.norm2(src)
            else:
                src = self.norm1(src)
                src2 = src

            src2 = self.activation(self.linear1(src2))
            if self.dropout is None:
                src2 = self.linear2(src2)
            else:
                src2 = self.linear2(self.dropout(src2))

            if self.dropout2 is None:
                src = src + src2
            else:
                src = src + self.dropout2(src2)

            if not self.norm_first:
                src = self.norm2(src)
        else:
            # densetnt vectornet
            src2 = self.activation(src2)
            if self.dropout is None:
                src = src + src2
            else:
                src = src + self.dropout(src2)
            if not self.norm_first:
                src = self.norm1(src)

        if src_padding_mask is not None:
            src.masked_fill_(src_padding_mask.unsqueeze(-1), 0.0)
            if need_weights:
                attn_weights.masked_fill_(src_padding_mask.unsqueeze(-1), 0.0)
        return src, attn_weights

src/models/modules/transformer.py
zhejz/TrafficBots
# Licensed under the CC BY-NC 4.0 license (https://creativecommons.org/licenses/by-nc/4.0/)
from typing import Optional, Tuple, Union, Dict
from omegaconf import DictConfig
import hydra
from torch import Tensor, nn
from .modules.distributions import MyDist
from .modules.mlp import MLP
from .modules.transformer import TransformerBlock
from .modules.map_encoder import MapEncoder
from .modules.input_pe_encoder import InputPeEncoder
from .modules.agent_interaction import MultiAgentTF
from .modules.agent_temporal import TemporalAggregate
from .modules.add_latent_goal import AddLatentGoal
from .goal_manager import GoalManager
from .latent_encoder import LatentEncoder


class TrafficBots(nn.Module):
    def __init__(
        self,
        hidden_dim: int,
        map_attr_dim: int,
        tl_attr_dim: int,
        agent_attr_dim: int,
        map_pe_dim: int,
        tl_pe_dim: int,
        agent_pe_dim: int,
        map_encoder: DictConfig,
        input_pe_encoder: DictConfig,
        goal_manager: DictConfig,
        latent_encoder: DictConfig,
        tf_cfg: DictConfig,
        n_layer_tf_as2pl: int,
        n_layer_tf_as2tl: int,
        n_step_hist: int,
        n_pl_node: int,
        temporal_aggregate: DictConfig,
        agent_temporal: DictConfig,
        agent_interaction: DictConfig,
        add_latent: DictConfig,
        add_goal: DictConfig,
        interaction_first: bool,
        add_goal_latent_first: bool,
        resample_latent: bool,
        n_layer_final_mlp: int,
        final_mlp: DictConfig,
    ):
        super().__init__()
        self.resample_latent = resample_latent
        self.interaction_first = interaction_first
        self.add_goal_latent_first = add_goal_latent_first

        # if self.add_learned_pe:
        #     self.pe_target = nn.Parameter(torch.zeros([1, n_step_hist, hidden_dim]), requires_grad=True)
        #     self.pe_other = nn.Parameter(torch.zeros([1, 1, n_step_hist, hidden_dim]), requires_grad=True)
        #     self.pe_map = nn.Parameter(torch.zeros([1, 1, n_pl_node, hidden_dim]), requires_grad=True)
        #     if self.use_current_tl:
        #         self.pe_tl = nn.Parameter(torch.zeros([1, 1, 1, hidden_dim]), requires_grad=True)
        #     else:
        #         self.pe_tl = nn.Parameter(torch.zeros([1, n_step_hist, 1, hidden_dim]), requires_grad=True)

        # input encoder
        self.map_encoder = MapEncoder(
            hidden_dim=hidden_dim,
            attr_dim=map_attr_dim,
            pe_dim=map_pe_dim,
            input_pe_encoder=input_pe_encoder,
            tf_cfg=tf_cfg,
            **map_encoder,
        )
        self.tl_encoder = InputPeEncoder(
            hidden_dim=hidden_dim, attr_dim=tl_attr_dim, pe_dim=tl_pe_dim, **input_pe_encoder
        )
        self.agent_encoder = InputPeEncoder(
            hidden_dim=hidden_dim, attr_dim=agent_attr_dim, pe_dim=agent_pe_dim, **input_pe_encoder
        )
        # cross-attn transformers, could be shared between policy and latent_encoder
        self.transformer_as2pl = TransformerBlock(n_layer=n_layer_tf_as2pl, **tf_cfg)
        self.transformer_as2tl = TransformerBlock(n_layer=n_layer_tf_as2tl, **tf_cfg)
        # goal
        self.goal_manager = GoalManager(tf_cfg=tf_cfg, **goal_manager)
        # latent
        self.latent_encoder = LatentEncoder(
            tf_cfg=tf_cfg,
            interaction_first=interaction_first,
            transformer_as2pl=self.transformer_as2pl,
            transformer_as2tl=self.transformer_as2tl,
            agent_temporal=agent_temporal,
            agent_interaction=agent_interaction,
            temporal_aggregate=temporal_aggregate,
            **latent_encoder,
        )
        # policy
        self.agent_temporal = hydra.utils.instantiate(agent_temporal, hidden_dim=hidden_dim)
        self.agent_interaction = MultiAgentTF(hidden_dim=hidden_dim, tf_cfg=tf_cfg, **agent_interaction)
        self.temporal_aggregate = TemporalAggregate(**temporal_aggregate)
        self.add_goal = AddLatentGoal(
            hidden_dim=hidden_dim, in_dim=self.goal_manager.out_dim, dummy=self.goal_manager.dummy, **add_goal
        )
        self.add_latent = AddLatentGoal(
            hidden_dim=hidden_dim, in_dim=self.latent_encoder.out_dim, dummy=self.latent_encoder.dummy, **add_latent
        )

        if n_layer_final_mlp > 0:
            self.final_mlp = MLP([hidden_dim] * (n_layer_final_mlp + 1), **final_mlp)
        else:
            self.final_mlp = None

    def encode_input_features(
        self,
        agent_valid: Tensor,
        agent_attr: Tensor,
        agent_pe: Tensor,
        agent_pos: Tensor,
        map_valid: Tensor,
        map_attr: Tensor,
        map_pe: Tensor,
        map_pos: Tensor,
        tl_valid: Tensor,
        tl_attr: Tensor,
        tl_pe: Tensor,
        tl_pos: Tensor,
    ) -> Dict[str, Tensor]:
        """
        Args:
            agent_valid: [n_scene, n_step(_hist), n_agent], bool
            agent_attr: [n_scene, n_step(_hist), n_agent, agent_attr_dim]
            agent_pe: [n_scene, n_step(_hist), n_agent, hidden_dim]
            agent_pos: [n_scene, n_step(_hist), n_agent, 2]
            map_valid: [n_scene, n_pl, n_pl_node], bool
            map_attr: [n_scene, n_pl, n_pl_node, map_attr_dim]
            map_pe: [n_scene, n_pl, n_pl_node, hidden_dim]
            map_pos: [n_scene, n_pl, 2]
            tl_valid: [n_scene, n_step(_hist), n_tl], bool
            tl_attr: [n_scene, n_step(_hist), n_tl, tl_attr_dim]
            tl_pe: [n_scene, n_step(_hist), n_tl, hidden_dim]
            tl_pos: [n_scene, n_step(_hist), n_tl, 2]

        Returns: for each agent a latent distribution that considers temporal relation and interaction between agents.
            agent_feature: [n_scene, n_step_hist, n_agent, hidden_dim]
            agent_feature_valid: [n_scene, n_step_hist, n_agent]
            map_feature: [n_scene, n_pl, hidden_dim]
            map_feature_valid: [n_scene, n_pl]
            tl_feature: [n_scene, n_step_hist, n_tl, hidden_dim]
            tl_feature_valid: [n_scene, n_step_hist, n_tl]
        """
        feature_dict = {"agent_feature_valid": agent_valid, "tl_feature_valid": tl_valid}
        feature_dict["map_feature"], feature_dict["map_feature_valid"] = self.map_encoder(map_valid, map_attr, map_pe)
        feature_dict["agent_feature"] = self.agent_encoder(agent_valid, agent_attr, agent_pe)
        feature_dict["tl_feature"] = self.tl_encoder(tl_valid, tl_attr, tl_pe)
        return feature_dict

    def init(self, latent: MyDist, deterministic: Union[bool, Tensor]) -> None:
        """
        deterministic: bool or [n_batch, n_agent] for sampling relevant agents and determistic other agents.
        """
        self.latent = latent
        self.deterministic = deterministic
        self.hidden = None
        self.latent_sample = None
        self.latent_logp = None

    def forward(
        self,
        agent_valid: Tensor,
        agent_feature: Tensor,
        map_valid: Tensor,
        map_feature: Tensor,
        tl_valid: Tensor,
        tl_feature: Tensor,
        goal_valid: Optional[Tensor],
        goal_feature: Optional[Tensor],
        need_weights: bool = False,
    ) -> Tuple[
        Tensor, Tensor, Optional[MyDist], Optional[MyDist], Optional[Tensor], Optional[Tensor], Optional[Tensor]
    ]:
        """
        Args:
            agent_valid: [n_batch, n_agent], bool
            agent_feature: [n_batch, n_agent, hidden_dim]
            map_valid: [n_batch, n_pl], bool
            map_feature: [n_batch, n_pl, hidden_dim]
            tl_valid: [n_batch, n_tl], bool
            tl_feature: [n_batch, n_tl, hidden_dim]
            goal_valid: [n_batch, n_agent] or None
            goal_feature: [n_batch, n_agent, hidden_dim] or None

        Returns: for each agent a latent distribution that considers temporal relation and interaction between agents.
            policy_feature: [n_batch, n_agent, hidden_dim]
            latent_logp: [n_batch, n_agent]
            attn_pl: [n_batch, n_agent, n_pl]
            attn_tl: [n_batch, n_agent, n_tl]
            attn_agent: [n_batch, n_agent, n_agent]
        """
        # ! sample latent
        if self.resample_latent or (self.latent_sample is None):
            # sample latent at each time step
            self.latent_sample = self.latent.sample(self.deterministic)
            self.latent_logp = self.latent.log_prob(self.latent_sample.detach())

        # [n_batch, n_step, n_agent, hidden_dim]
        policy_feature = agent_feature

        # ! attention to map polyline
        policy_feature, attn_pl = self.transformer_as2pl(
            src=policy_feature,  # [n_batch, n_agent, hidden_dim]
            src_padding_mask=~agent_valid,  # [n_batch, n_agent]
            tgt=map_feature,  # [n_batch, n_pl, hidden_dim]
            tgt_padding_mask=~map_valid,  # [n_batch, n_pl]
            need_weights=need_weights,
        )
        # ! attention to traffic light
        policy_feature, attn_tl = self.transformer_as2tl(
            src=policy_feature,  # [n_batch, n_agent, hidden_dim]
            src_padding_mask=~agent_valid,  # [n_batch, n_agent]
            tgt=tl_feature,  # [n_batch, n_tl, hidden_dim]
            tgt_padding_mask=~tl_valid,  # [n_batch, n_tl]
            need_weights=need_weights,
        )

        if self.add_goal_latent_first:
            # ! add goal and latent before interaction
            policy_feature = self.add_goal(policy_feature, agent_valid, goal_feature, goal_valid)
            policy_feature = self.add_latent(policy_feature, agent_valid, self.latent_sample, agent_valid)

        # ! interaction and temporal
        if self.interaction_first:
            policy_feature, attn_agent = self.agent_interaction(
                policy_feature, agent_feature, agent_valid, need_weights=need_weights
            )
            policy_feature, self.hidden = self.agent_temporal(policy_feature, agent_valid, self.hidden)
        else:
            policy_feature, self.hidden = self.agent_temporal(policy_feature, agent_valid, self.hidden)
            policy_feature, attn_agent = self.agent_interaction(
                policy_feature, agent_feature, agent_valid, need_weights=need_weights
            )

        if not self.add_goal_latent_first:
            # ! add goal and latent before interaction
            policy_feature = self.add_goal(policy_feature, agent_valid, goal_feature, goal_valid)
            policy_feature = self.add_latent(policy_feature, agent_valid, self.latent_sample, agent_valid)

        # ! final mlp
        if self.final_mlp is not None:
            policy_feature = self.final_mlp(policy_feature, agent_valid)

        return policy_feature, self.latent_logp, attn_pl, attn_tl, attn_agent

src/models/traffic_bots.py
zhejz/TrafficBots
import sys

sys.path.append(".")

from argparse import ArgumentParser
from tqdm import tqdm
import h5py
import numpy as np
from pathlib import Path
from waymo_open_dataset.protos import scenario_pb2
import src.utils.pack_h5 as pack_utils
import tensorflow as tf

# Licensed under the CC BY-NC 4.0 license (https://creativecommons.org/licenses/by-nc/4.0/)
tf.config.set_visible_devices([], "GPU")

# "map/type"
# FREEWAY = 0
# SURFACE_STREET = 1
# STOP_SIGN = 2
# BIKE_LANE = 3
# TYPE_ROAD_EDGE_BOUNDARY = 4
# TYPE_ROAD_EDGE_MEDIAN = 5
# BROKEN = 6
# SOLID_SINGLE = 7
# DOUBLE = 8
# SPEED_BUMP = 9
# CROSSWALK = 10
N_PL_TYPE = 11
DIM_VEH_LANES = [0, 1, 2]
DIM_CYC_LANES = [3]
DIM_PED_LANES = [4]

# "tl_lane/state", "tl_stop/state"
# LANE_STATE_UNKNOWN = 0;
# LANE_STATE_STOP = 1;
# LANE_STATE_CAUTION = 2;
# LANE_STATE_GO = 3;
# LANE_STATE_FLASHING = 4;
N_TL_STATE = 5

N_PL_MAX = 3000
N_TL_MAX = 40
N_AGENT_MAX = 1300

N_PL = 1024
N_TL = 100
N_AGENT = 64
N_AGENT_NO_SIM = 256

THRESH_MAP = 500  # ! 120
THRESH_AGENT = 120

N_STEP = 91
STEP_CURRENT = 10


def collate_agent_features(tracks, sdc_track_index, track_index_predict, object_id_interest):
    agent_id = []
    agent_type = []
    agent_states = []
    agent_role = []
    for i, _track in enumerate(tracks):
        agent_id.append(_track.id)
        agent_type.append(_track.object_type - 1)  # [TYPE_VEHICLE=1, TYPE_PEDESTRIAN=2, TYPE_CYCLIST=3] -> [0,1,2]
        step_states = []
        for s in _track.states:
            step_states.append(
                [
                    s.center_x,
                    s.center_y,
                    s.center_z,
                    s.length,
                    s.width,
                    s.height,
                    s.heading,
                    s.velocity_x,
                    s.velocity_y,
                    s.valid,
                ]
            )
            # This angle is normalized to [-pi, pi). The velocity vector in m/s
        agent_states.append(step_states)

        agent_role.append([False, False, False])
        if i in track_index_predict:
            agent_role[-1][2] = True
        if _track.id in object_id_interest:
            agent_role[-1][1] = True
        if i == sdc_track_index:
            agent_role[-1][0] = True

    return agent_id, agent_type, agent_states, agent_role


def collate_tl_features(tl_features):
    tl_lane_state = []
    tl_lane_id = []
    tl_stop_point = []
    for _step_tl in tl_features:
        step_tl_lane_state = []
        step_tl_lane_id = []
        step_tl_stop_point = []
        for _tl in _step_tl.lane_states:
            if _tl.state == 0:  # LANE_STATE_UNKNOWN = 0;
                tl_state = 0  # LANE_STATE_UNKNOWN = 0;
            elif _tl.state in [1, 4]:  # LANE_STATE_ARROW_STOP = 1; LANE_STATE_STOP = 4;
                tl_state = 1  # LANE_STATE_STOP = 1;
            elif _tl.state in [2, 5]:  # LANE_STATE_ARROW_CAUTION = 2; LANE_STATE_CAUTION = 5;
                tl_state = 2  # LANE_STATE_CAUTION = 2;
            elif _tl.state in [3, 6]:  # LANE_STATE_ARROW_GO = 3; LANE_STATE_GO = 6;
                tl_state = 3  # LANE_STATE_GO = 3;
            elif _tl.state in [7, 8]:  # LANE_STATE_FLASHING_STOP = 7; LANE_STATE_FLASHING_CAUTION = 8;
                tl_state = 4  # LANE_STATE_FLASHING = 4;
            else:
                assert ValueError

            step_tl_lane_state.append(tl_state)
            step_tl_lane_id.append(_tl.lane)
            step_tl_stop_point.append([_tl.stop_point.x, _tl.stop_point.y, _tl.stop_point.z])

        tl_lane_state.append(step_tl_lane_state)
        tl_lane_id.append(step_tl_lane_id)
        tl_stop_point.append(step_tl_stop_point)
    return tl_lane_state, tl_lane_id, tl_stop_point


def collate_map_features(map_features):
    mf_id = []
    mf_xyz = []
    mf_type = []
    mf_edge = []
    for mf in map_features:
        feature_data_type = mf.WhichOneof("feature_data")
        if feature_data_type is None:  # pip install waymo-open-dataset-tf-2-6-0==1.4.9, not updated, should be driveway
            continue
        feature = getattr(mf, feature_data_type)
        if feature_data_type == "lane":
            if feature.type == 0:  # UNDEFINED
                mf_type.append(1)
            elif feature.type == 1:  # FREEWAY
                mf_type.append(0)
            elif feature.type == 2:  # SURFACE_STREET
                mf_type.append(1)
            elif feature.type == 3:  # BIKE_LANE
                mf_type.append(3)
            mf_id.append(mf.id)
            mf_xyz.append([[p.x, p.y, p.z] for p in feature.polyline][::2])
            if len(feature.exit_lanes) > 0:
                for _id_exit in feature.exit_lanes:
                    mf_edge.append([mf.id, _id_exit])
            else:
                mf_edge.append([mf.id, -1])
        elif feature_data_type == "stop_sign":
            for l_id in feature.lane:
                # override FREEWAY/SURFACE_STREET with stop sign lane
                # BIKE_LANE remains unchanged
                idx_lane = mf_id.index(l_id)
                if mf_type[idx_lane] < 2:
                    mf_type[idx_lane] = 2
        elif feature_data_type == "road_edge":
            assert feature.type > 0  # no UNKNOWN = 0
            mf_id.append(mf.id)
            mf_type.append(feature.type + 3)  # [1, 2] -> [4, 5]
            mf_xyz.append([[p.x, p.y, p.z] for p in feature.polyline][::2])
        elif feature_data_type == "road_line":
            assert feature.type > 0  # no UNKNOWN = 0
            # BROKEN_SINGLE_WHITE = 1
            # SOLID_SINGLE_WHITE = 2
            # SOLID_DOUBLE_WHITE = 3
            # BROKEN_SINGLE_YELLOW = 4
            # BROKEN_DOUBLE_YELLOW = 5
            # SOLID_SINGLE_YELLOW = 6
            # SOLID_DOUBLE_YELLOW = 7
            # PASSING_DOUBLE_YELLOW = 8
            if feature.type in [1, 4, 5]:
                feature_type_new = 6  # BROKEN
            elif feature.type in [2, 6]:
                feature_type_new = 7  # SOLID_SINGLE
            else:
                feature_type_new = 8  # DOUBLE
            mf_id.append(mf.id)
            mf_type.append(feature_type_new)
            mf_xyz.append([[p.x, p.y, p.z] for p in feature.polyline][::2])
        elif feature_data_type in ["speed_bump", "driveway", "crosswalk"]:
            xyz = np.array([[p.x, p.y, p.z] for p in feature.polygon])
            polygon_idx = np.linspace(0, xyz.shape[0], 4, endpoint=False, dtype=int)
            pl_polygon = pack_utils.get_polylines_from_polygon(xyz[polygon_idx])
            mf_xyz.extend(pl_polygon)
            mf_id.extend([mf.id] * len(pl_polygon))
            pl_type = 9 if feature_data_type in ["speed_bump", "driveway"] else 10
            mf_type.extend([pl_type] * len(pl_polygon))
        else:
            raise ValueError

    return mf_id, mf_xyz, mf_type, mf_edge


def main():
    parser = ArgumentParser(allow_abbrev=True)
    parser.add_argument("--data-dir", default="/cluster/scratch/zhejzhan/womd_scenario_v_1_2_0")
    parser.add_argument("--dataset", default="training")
    parser.add_argument("--out-dir", default="/cluster/scratch/zhejzhan/h5_womd_sim_agent")
    parser.add_argument("--rand-pos", default=50.0, type=float, help="Meter. Set to -1 to disable.")
    parser.add_argument("--rand-yaw", default=3.14, type=float, help="Radian. Set to -1 to disable.")
    parser.add_argument("--dest-no-pred", action="store_true")
    args = parser.parse_args()

    dataset_size = {
        "training": 486995,  # 487002
        "validation": 44097,
        "training_20s": 70541,
        "validation_interactive": 43479,
        "testing": 44920,
        "testing_interactive": 44154,
    }

    if "training" in args.dataset:
        pack_all = True  # ["agent/valid"]
        pack_history = False  # ["history/agent/valid"]
    elif "validation" in args.dataset:
        pack_all = True
        pack_history = True
    elif "testing" in args.dataset:
        pack_all = False
        pack_history = True

    out_path = Path(args.out_dir)
    out_path.mkdir(exist_ok=True)
    out_h5_path = out_path / (args.dataset + ".h5")

    data_path = Path(args.data_dir) / args.dataset
    dataset = tf.data.TFRecordDataset(sorted([p.as_posix() for p in data_path.glob("*")]), compression_type="")
    n_pl_max, n_tl_max, n_agent_max, n_agent_sim, n_agent_no_sim, data_len = 0, 0, 0, 0, 0, 0
    with h5py.File(out_h5_path, "w") as hf:
        for i, data in tqdm(enumerate(dataset), total=dataset_size[args.dataset]):
            # if i < 26:
            #     continue
            # if i > 26:
            #     break

            scenario = scenario_pb2.Scenario()
            scenario.ParseFromString(data.numpy())

            mf_id, mf_xyz, mf_type, mf_edge = collate_map_features(scenario.map_features)
            tl_lane_state, tl_lane_id, tl_stop_point = collate_tl_features(scenario.dynamic_map_states)
            agent_id, agent_type, agent_states, agent_role = collate_agent_features(
                scenario.tracks,
                sdc_track_index=scenario.sdc_track_index,
                track_index_predict=[i.track_index for i in scenario.tracks_to_predict],
                object_id_interest=[i for i in scenario.objects_of_interest],
            )

            episode = {}
            n_pl = pack_utils.pack_episode_map(
                episode=episode, mf_id=mf_id, mf_xyz=mf_xyz, mf_type=mf_type, mf_edge=mf_edge, n_pl_max=N_PL_MAX
            )
            n_tl = pack_utils.pack_episode_traffic_lights(
                episode=episode,
                tl_lane_state=tl_lane_state,
                tl_lane_id=tl_lane_id,
                tl_stop_point=tl_stop_point,
                pack_all=pack_all,
                pack_history=pack_history,
                n_tl_max=N_TL_MAX,
                step_current=STEP_CURRENT,
            )
            n_agent = pack_utils.pack_episode_agents(
                episode=episode,
                agent_id=agent_id,
                agent_type=agent_type,
                agent_states=agent_states,
                agent_role=agent_role,
                pack_all=pack_all,
                pack_history=pack_history,
                n_agent_max=N_AGENT_MAX,
                step_current=STEP_CURRENT,
            )
            scenario_center, scenario_yaw = pack_utils.center_at_sdc(episode, args.rand_pos, args.rand_yaw)
            n_pl_max = max(n_pl_max, n_pl)
            n_tl_max = max(n_tl_max, n_tl)
            n_agent_max = max(n_agent_max, n_agent)

            episode_reduced = {}
            pack_utils.filter_episode_map(episode, N_PL, THRESH_MAP, thresh_z=5)
            episode_with_map = episode["map/valid"].any(1).sum() > 0
            pack_utils.repack_episode_map(episode, episode_reduced, N_PL, N_PL_TYPE)

            pack_utils.filter_episode_traffic_lights(episode)
            pack_utils.repack_episode_traffic_lights(episode, episode_reduced, N_TL, N_TL_STATE)

            if "training" in args.dataset:
                mask_sim, mask_no_sim = pack_utils.filter_episode_agents(
                    episode=episode,
                    episode_reduced=episode_reduced,
                    n_agent=N_AGENT,
                    prefix="",
                    dim_veh_lanes=DIM_VEH_LANES,
                    dist_thresh_agent=THRESH_AGENT,
                    step_current=STEP_CURRENT,
                )
                pack_utils.repack_episode_agents(
                    episode=episode,
                    episode_reduced=episode_reduced,
                    mask_sim=mask_sim,
                    n_agent=N_AGENT,
                    prefix="",
                    dim_veh_lanes=DIM_VEH_LANES,
                    dim_cyc_lanes=DIM_CYC_LANES,
                    dim_ped_lanes=DIM_PED_LANES,
                    dest_no_pred=args.dest_no_pred,
                )
            elif "validation" in args.dataset:
                mask_sim, mask_no_sim = pack_utils.filter_episode_agents(
                    episode=episode,
                    episode_reduced=episode_reduced,
                    n_agent=N_AGENT,
                    prefix="history/",
                    dim_veh_lanes=DIM_VEH_LANES,
                    dist_thresh_agent=THRESH_AGENT,
                    step_current=STEP_CURRENT,
                )
                pack_utils.repack_episode_agents(
                    episode=episode,
                    episode_reduced=episode_reduced,
                    mask_sim=mask_sim,
                    n_agent=N_AGENT,
                    prefix="",
                    dim_veh_lanes=DIM_VEH_LANES,
                    dim_cyc_lanes=DIM_CYC_LANES,
                    dim_ped_lanes=DIM_PED_LANES,
                    dest_no_pred=args.dest_no_pred,
                )
                pack_utils.repack_episode_agents(episode, episode_reduced, mask_sim, N_AGENT, "history/")
                pack_utils.repack_episode_agents_no_sim(episode, episode_reduced, mask_no_sim, N_AGENT_NO_SIM, "")
                pack_utils.repack_episode_agents_no_sim(
                    episode, episode_reduced, mask_no_sim, N_AGENT_NO_SIM, "history/"
                )
            elif "testing" in args.dataset:
                if episode_with_map:
                    mask_sim, mask_no_sim = pack_utils.filter_episode_agents(
                        episode=episode,
                        episode_reduced=episode_reduced,
                        n_agent=N_AGENT,
                        prefix="history/",
                        dim_veh_lanes=DIM_VEH_LANES,
                        dist_thresh_agent=THRESH_AGENT,
                        step_current=STEP_CURRENT,
                    )
                else:
                    mask_valid = episode["history/agent/valid"].any(0)
                    mask_sim = episode["history/agent/role"].any(-1)
                    for _valid_idx in np.where(mask_valid)[0]:
                        mask_sim[_valid_idx] = True
                        if mask_sim.sum() >= N_AGENT:
                            break
                    mask_no_sim = mask_valid & (~mask_sim)

                pack_utils.repack_episode_agents(episode, episode_reduced, mask_sim, N_AGENT, "history/")
                pack_utils.repack_episode_agents_no_sim(
                    episode, episode_reduced, mask_no_sim, N_AGENT_NO_SIM, "history/"
                )
            n_agent_sim = max(n_agent_sim, mask_sim.sum())
            n_agent_no_sim = max(n_agent_no_sim, mask_no_sim.sum())

            if episode_with_map:
                episode_reduced["map/boundary"] = pack_utils.get_map_boundary(
                    episode_reduced["map/valid"], episode_reduced["map/pos"]
                )
            else:
                # only in waymo test split.
                assert args.dataset == "testing"
                episode_reduced["map/boundary"] = pack_utils.get_map_boundary(
                    episode["history/agent/valid"], episode["history/agent/pos"]
                )
                print(f"scenario {i} has no map! map boundary is: {episode_reduced['map/boundary']}")

            hf_episode = hf.create_group(str(i))
            hf_episode.attrs["scenario_id"] = scenario.scenario_id
            hf_episode.attrs["scenario_center"] = scenario_center
            hf_episode.attrs["scenario_yaw"] = scenario_yaw
            hf_episode.attrs["with_map"] = episode_with_map

            for k, v in episode_reduced.items():
                hf_episode.create_dataset(k, data=v, compression="gzip", compression_opts=4, shuffle=True)

            data_len += 1

        print(f"data_len: {data_len}, dataset_size: {dataset_size[args.dataset]}")
        print(f"n_pl_max: {n_pl_max}")
        print(f"n_agent_max: {n_agent_max}, n_agent_sim: {n_agent_sim}, n_agent_no_sim: {n_agent_no_sim}")
        hf.attrs["data_len"] = data_len


if __name__ == "__main__":
    main()

src/pack_h5_womd.py
zhejz/TrafficBots
# Licensed under the CC BY-NC 4.0 license (https://creativecommons.org/licenses/by-nc/4.0/)
from typing import Dict, List, Optional, Tuple, Union
import hydra
from pytorch_lightning import LightningModule
import torch
from omegaconf import DictConfig
from torch import Tensor, nn
import wandb
from pathlib import Path
from collections import OrderedDict

from models.metrics.training import TrainingMetrics
from models.metrics.logging import ErrorMetrics, TrafficRuleMetrics
from models.modules.distributions import MyDist
from models.modules.action_head import ActionHead
from utils.traffic_rule_checker import TrafficRuleChecker
from utils.teacher_forcing import TeacherForcing
from utils.buffer import RolloutBuffer
from utils.rewards import DifferentiableReward
from utils.dynamics import Dynamics
from utils.vis_waymo import VisWaymo
from data_modules.waymo_post_processing import WaymoPostProcessing
from models.metrics.womd import WOMDMetrics
from utils.submission import SubWOMD


class WaymoMotion(LightningModule):
    def __init__(
        self,
        time_step_current: int,
        time_step_gt: int,
        time_step_end: int,
        time_step_sim_start: int,
        hidden_dim: int,
        data_size: DictConfig,
        pre_processing: DictConfig,
        step_detach_hidden: int,
        model: DictConfig,
        p_training_rollout_prior: float,
        detach_state_policy: bool,
        training_deterministic_action: bool,
        differentiable_reward: DictConfig,
        p_drop_hidden: float,
        n_video_batch: int,
        n_joint_future: int,
        waymo_post_processing: DictConfig,
        dynamics: DictConfig,
        action_head: DictConfig,
        teacher_forcing_training: DictConfig,
        teacher_forcing_reactive_replay: DictConfig,
        teacher_forcing_joint_future_pred: DictConfig,
        training_metrics: DictConfig,
        traffic_rule_checker: DictConfig,
        optimizer: DictConfig,
        lr_scheduler: Optional[DictConfig],
        lr_goal: float,
        sub_womd_reactive_replay: DictConfig,
        sub_womd_joint_future_pred: DictConfig,
        interactive_challenge: bool = False,
        wb_artifact: Optional[str] = None,
    ) -> None:
        super().__init__()
        self.save_hyperparameters()

        # pre_processing
        self.pre_processing = []
        pre_proc_kwargs = {}
        for k, v in pre_processing.items():
            _pre_proc = hydra.utils.instantiate(v, time_step_current=time_step_current, data_size=data_size)
            self.pre_processing.append((k, _pre_proc))
            pre_proc_kwargs |= _pre_proc.model_kwargs
        self.pre_processing = nn.Sequential(OrderedDict(self.pre_processing))

        # model
        self.model = hydra.utils.instantiate(model, **pre_proc_kwargs, _recursive_=False)
        self.dynamics = Dynamics(**dynamics)
        self.action_head = ActionHead(hidden_dim=hidden_dim, action_dim=self.dynamics.action_dim, **action_head)

        # diffbar rewards
        self.diffbar_reward = DifferentiableReward(**differentiable_reward)

        # training
        self.teacher_forcing_training = TeacherForcing(**teacher_forcing_training)
        self.train_metrics_train = TrainingMetrics("training", **training_metrics)

        # reactive_replay: scene reconstruct given a complete episode, same setup as training
        self.teacher_forcing_reactive_replay = TeacherForcing(**teacher_forcing_reactive_replay)
        self.train_metrics_reactive_replay = TrainingMetrics("reactive_replay", **training_metrics)
        self.err_metrics_reactive_replay = ErrorMetrics("reactive_replay")
        self.rule_metrics_reactive_replay = TrafficRuleMetrics("reactive_replay")
        self.waymo_post_processing = WaymoPostProcessing(**waymo_post_processing)
        self.womd_metrics_reactive_replay = WOMDMetrics(
            "reactive_replay", time_step_end, time_step_current, interactive_challenge
        )

        # joint_future_pred: no spawn, tl_state from current_step, prior latent, goal/dest
        self.teacher_forcing_joint_future_pred = TeacherForcing(**teacher_forcing_joint_future_pred)
        self.err_metrics_joint_future_pred = ErrorMetrics("joint_future_pred")
        self.rule_metrics_joint_future_pred = TrafficRuleMetrics("joint_future_pred")
        self.womd_metrics_joint_future_pred = WOMDMetrics(
            "joint_future_pred", time_step_end, time_step_current, interactive_challenge
        )

        # save submission files
        self.sub_womd_reactive_replay = SubWOMD(wb_artifact=wb_artifact, **sub_womd_reactive_replay)
        self.sub_womd_joint_future_pred = SubWOMD(wb_artifact=wb_artifact, **sub_womd_joint_future_pred)

    def forward(
        self,
        map_feature: Tensor,
        map_valid: Tensor,
        tl_feature: Tensor,
        tl_valid: Tensor,
        goal_feature: Tensor,
        goal_valid: Tensor,
        action_override: Optional[Tensor] = None,
        mask_action_override: Optional[Tensor] = None,
        state_override: Optional[Dict[str, Tensor]] = None,
        mask_state_override: Optional[Tensor] = None,
        deterministic_action: bool = True,
        require_train_dict: bool = True,
        require_vis_dict: bool = False,
    ) -> Tuple[Tensor, Tensor, Dict[str, Tensor], Dict[str, Tensor]]:
        """
        Args:
            map_feature, map_valid: [n_batch, n_pl, (hidden_dim)]
            tl_feature, tl_valid: [n_batch, n_tl, (hidden_dim)]
            action_override: [n_batch, n_agent, k_pred, 2]
            mask_action_override: [n_batch, n_agent] bool
            state_override["as_pred"]: [n_batch, n_agent, k_pred, 4]
            state_override["acc"], state_override["yaw_rate]: [n_batch, n_agent, 1]
            mask_state_override: [n_batch, n_agent] bool, can be used to spawn new agents

        Return:
            state: [n_batch, n_agent, 4], x,y,theta,v
            valid: [n_batch, n_agent], bool
        """
        agent_valid = self.dynamics.agent_valid
        agent_state = self.dynamics.agent_state
        agent_attr, agent_pe = self.pre_processing.input.get_agent_attr_and_pe(
            agent_pos=agent_state[:, :, :2],
            agent_yaw_bbox=agent_state[:, :, 2:3],
            agent_spd=agent_state[:, :, 3:4],
            agent_vel=self.dynamics.vel,
            agent_yaw_rate=self.dynamics.yaw_rate,
            agent_acc=self.dynamics.acc,
            agent_size=self.dynamics.agent_size,
            agent_type=self.dynamics.agent_type,
        )
        if self.hparams.detach_state_policy:
            agent_state = agent_state.detach()
            agent_attr = agent_attr.detach()
            agent_pe = agent_pe.detach()

        agent_feature = self.model.agent_encoder(agent_valid, agent_attr, agent_pe)

        # some policies (e.g. TrafficSim) will infer a new pair of latent_post and latent_prior at each step.
        policy_feature, latent_logp, attn_pl, attn_tl, attn_agent = self.model(
            agent_valid=agent_valid,
            agent_feature=agent_feature,
            map_valid=map_valid,
            map_feature=map_feature,
            tl_valid=tl_valid,
            tl_feature=tl_feature,
            goal_valid=goal_valid,
            goal_feature=goal_feature,
            need_weights=require_vis_dict,
        )

        # action_dist: Gaussian or DiagGaussian
        action_dist = self.action_head(policy_feature, agent_valid, self.dynamics.agent_type)
        pred_state, pred_valid, action, action_log_prob = self.dynamics.update(
            action_dist=action_dist,
            action_override=action_override,
            mask_action_override=mask_action_override,
            deterministic=deterministic_action,
        )
        self.dynamics.override_states(state_override=state_override, mask_state_override=mask_state_override)

        if require_train_dict:
            train_dict = {
                # for logging
                "latent_log_prob": latent_logp,  # [n_batch, n_agent]
                "action_log_prob": action_log_prob,  # [n_batch, n_agent]
                # for train_metrics
                "pred_valid": pred_valid,  # old agent_valid, not updated with mask_override
                "pred_state": pred_state,
            }
        else:
            train_dict = {}
        if require_vis_dict:
            if goal_valid is None:
                goal_valid = torch.zeros_like(pred_valid)
            vis_dict = {
                "action": action.cpu(),
                "goal_valid": goal_valid.cpu(),
                "attn_weights_to_pl": attn_pl.cpu(),
                "attn_weights_to_tl": attn_tl.cpu(),
                "attn_weights_to_agent": attn_agent.cpu(),
            }
        else:
            vis_dict = {}
        return self.dynamics.agent_state, self.dynamics.agent_valid, train_dict, vis_dict

    def rollout(
        self,
        features: Dict[str, Tensor],
        latent: MyDist,
        goal: Tensor,
        goal_valid: Tensor,
        mask_teacher_forcing: Tensor,
        rule_checker: TrafficRuleChecker,
        deterministic_latent: Union[bool, Tensor],
        deterministic_action: bool,
        step_end: int,
        step_start: int,
        require_vis_dict: bool = False,
        gt_sdc: Optional[Dict[str, Tensor]] = None,
    ):
        """
        Args:
            features: step_current as step_gt during validation
                "map_valid": [n_batch, n_pl] bool
                "map_feature": [n_batch, n_pl, hidden_dim]
                "tl_valid": [n_batch, step_tl_gt+1, n_tl] bool
                "tl_feature": [n_batch, step_tl_gt+1, n_tl, hidden_dim]
                "agent_valid": [n_batch, step_gt+1, n_agent] bool
                "agent_state": [n_batch, step_gt+1, n_agent, 4], x,y,theta,spd
                "agent_type": [n_batch, n_agent, 3] one_hot bool [veh=0, ped=1, cyc=2]
                "agent_size": [n_batch, n_agent, 3] float32 [length, width, height]
                "vel": [n_batch, step_gt+1, n_agent, 2]
                "acc": [n_batch, step_gt+1, n_agent, 1]
                "yaw_rate": [n_batch, step_gt+1, n_agent, 1]
            goal: [n_batch, n_agent, 4] float32[x, y, theta, v] or [n_batch, n_agent] int64: index to map n_pl
            goal_valid: [n_batch, n_agent] bool
            latent: [n_batch, n_agent, hidden_dim]
            mask_teacher_forcing: [n_batch, step_gt+1, n_agent] bool
            deterministic_latent: bool or [n_batch, n_agent] for sampling relevant agents and determistic other agents.
            gt_sdc: None or {"agent_state", "vel", "acc", "yaw_rate"} [n_batch, n_step, 4]

        Returns:
            rollout_buffer [1, ..., step_end]
        """
        assert mask_teacher_forcing.shape[1] == features["agent_state"].shape[1]
        # init buffer and policy
        rollout_buffer = RolloutBuffer(
            step_start=step_start, step_end=step_end, step_current=self.hparams.time_step_current
        )
        self.model.init(latent, deterministic_latent)
        # init with first frame
        self.dynamics.init(
            agent_valid=features["agent_valid"][:, 0],
            agent_state=features["agent_state"][:, 0],
            agent_type=features["agent_type"],
            agent_size=features["agent_size"],
            vel=features["vel"][:, 0],
            acc=features["acc"][:, 0],
            yaw_rate=features["yaw_rate"][:, 0],
        )
        # encode goal
        if self.model.goal_manager.dummy:
            goal_feature = None
        else:
            goal_feature = self.model.goal_manager.get_goal_feature(
                goal=goal, as_state=self.dynamics.agent_state, map_feature=features["map_feature"]
            )

        # rollout
        for _step in range(step_start, step_end + 1):
            # prepare state_override: [n_batch, n_agent]
            if _step >= features["agent_valid"].shape[1]:
                mask_state_override = torch.zeros_like(mask_teacher_forcing[:, 0])
            else:
                mask_state_override = mask_teacher_forcing[:, _step]
            state_override = None
            if mask_state_override.any():
                state_override = {k: features[k][:, _step] for k in self.dynamics.state_keys}
            # what-if motion prediction
            if gt_sdc is not None:
                mask_state_override[:, 0] = True
                if state_override is None:
                    state_override = {k: features[k][:, 0] for k in self.dynamics.state_keys}
                for k in gt_sdc.keys():
                    state_override[k][:, 0] = gt_sdc[k][:, _step]

            # predict _step, given _step-1, using last observered tl_state
            step_tl = min(_step - 1, features["tl_valid"].shape[1] - 1)  # todo(maybe): could also use tl_state=UNKNOWN

            if self.model.goal_manager.update_goal:
                goal_feature = self.model.goal_manager.get_goal_feature(
                    goal=goal, as_state=self.dynamics.agent_state, map_feature=features["map_feature"]
                )

            state_new, valid_new, train_dict, vis_dict = self.forward(
                map_feature=features["map_feature"],
                map_valid=features["map_valid"],
                tl_feature=features["tl_feature"][:, step_tl],
                tl_valid=features["tl_valid"][:, step_tl],
                goal_feature=goal_feature,
                goal_valid=goal_valid,
                state_override=state_override,
                mask_state_override=mask_state_override,
                deterministic_action=deterministic_action,
                require_train_dict=True,
                require_vis_dict=require_vis_dict,
            )

            # check violations, kill agents if outside map, update self.dynamics.agent_valid
            _gt_valid = None if _step >= features["agent_valid"].shape[1] else features["agent_valid"][:, _step]
            _gt_state = None if _step >= features["agent_valid"].shape[1] else features["agent_state"][:, _step]
            violations = rule_checker.check(_step, valid_new, state_new)
            self.dynamics.kill(violations, _gt_valid)
            goal_valid = self.model.goal_manager.disable_goal_reached(
                goal_valid=goal_valid,
                agent_valid=self.dynamics.agent_valid,
                dest_reached=violations["dest_reached"],
                goal_reached=violations["goal_reached"],
            )

            # diffbar_reward
            if self.train_metrics_train.use_diffbar_reward:
                diffbar_reward, diffbar_reward_valid = self.diffbar_reward.get(
                    agent_valid=train_dict["pred_valid"],
                    agent_state=train_dict["pred_state"],
                    gt_valid=_gt_valid,
                    gt_state=_gt_state,
                    agent_size=features["agent_size"],
                )
            else:
                diffbar_reward, diffbar_reward_valid = None, None

            # finish iteration by adding to buffer
            rollout_buffer.add(
                valid=train_dict["pred_valid"],
                pred=train_dict["pred_state"],
                override_mask=mask_state_override,
                violation=violations,
                diffbar_reward=diffbar_reward,
                diffbar_reward_valid=diffbar_reward_valid,
                latent_log_prob=train_dict["latent_log_prob"],
                action_log_prob=train_dict["action_log_prob"],
                vis_dict=vis_dict,
            )

            # stop gradient or dropout during training for hidden
            if self.training and (self.model.hidden is not None):
                if _step <= self.hparams.step_detach_hidden:
                    self.model.hidden = self.model.hidden.detach()
                if self.hparams.p_drop_hidden > 0:
                    if torch.rand(1) < self.hparams.p_drop_hidden:
                        self.model.hidden = torch.zeros_like(self.model.hidden)

        rollout_buffer.finish()
        return rollout_buffer

    def training_step(self, batch: Dict[str, Tensor], batch_idx: int) -> Dict:
        with torch.no_grad():
            batch = self.pre_processing(batch)
            # for k, v in batch.items():
            #     if not v.is_contiguous():
            #         print(k)
            input_dict = {k.split("input/")[-1]: v for k, v in batch.items() if "input/" in k}
            latent_post_dict = {k.split("latent_post/")[-1]: v for k, v in batch.items() if "latent_post/" in k}
            latent_prior_dict = {k.split("latent_prior/")[-1]: v for k, v in batch.items() if "latent_prior/" in k}

        input_feature_dict = self.model.encode_input_features(**input_dict)
        latent_post_feature_dict = self.model.encode_input_features(**latent_post_dict)
        latent_prior_feature_dict = self.model.encode_input_features(**latent_prior_dict)

        # ! goal and destination
        goal_gt, goal_valid = self.model.goal_manager.get_gt_goal(
            agent_valid=input_dict["agent_valid"], gt_dest=batch["gt/dest"], gt_goal=batch["gt/goal"]
        )
        goal_pred = self.model.goal_manager.pred_goal(
            agent_type=batch["ref/agent_type"],
            map_type=batch["ref/map_type"],
            agent_state=batch["ref/agent_state"],
            **input_feature_dict,
        )

        # ! latents
        latent_post = self.model.latent_encoder(posterior=True, **latent_post_feature_dict)
        latent_prior = self.model.latent_encoder(**latent_prior_feature_dict)
        if torch.rand(1) < self.hparams.p_training_rollout_prior:
            latent = latent_prior
        else:
            latent = latent_post

        # ! rollout
        rollout_buffer = self.reactive_replay(
            batch=batch,
            input_feature_dict=input_feature_dict,
            mask_teacher_forcing=self.teacher_forcing_training.get(batch["gt/valid"], self.current_epoch),
            latent=latent,
            goal=goal_gt,
            goal_valid=goal_valid,
            deterministic_latent=False,
            deterministic_action=self.hparams.training_deterministic_action,
            require_vis_dict=False,
        )

        # ! metrics
        metrics_dict = self.train_metrics_train(
            pred_valid=rollout_buffer.valid,
            diffbar_rewards_valid=rollout_buffer.diffbar_rewards_valid,
            diffbar_rewards=rollout_buffer.diffbar_rewards,
            override_masks=rollout_buffer.override_masks,
            agent_role=batch["ref/agent_role"],
            goal_valid=goal_valid,
            goal_pred=goal_pred,
            goal_gt=goal_gt,
            latent_post=latent_post,
            latent_prior=latent_prior,
        )
        for k in metrics_dict.keys():
            self.log(k, metrics_dict[k], on_step=True)
        self.train_metrics_train.reset()
        return metrics_dict[f"{self.train_metrics_train.prefix}/loss"]

    def reactive_replay(
        self,
        batch: Dict[str, Tensor],
        input_feature_dict: Dict[str, Tensor],
        mask_teacher_forcing: Tensor,
        latent: Optional[MyDist],
        goal: Optional[Tensor],
        goal_valid: Optional[Tensor],
        deterministic_latent: bool,
        deterministic_action: bool,
        require_vis_dict: bool,
    ) -> RolloutBuffer:
        # init traffic rule checker
        rule_checker = TrafficRuleChecker(
            batch["map/boundary"],
            batch["map/valid"],
            batch["map/type"],
            batch["map/pos"],
            batch["map/dir"],
            batch["tl_stop/valid"],
            batch["tl_stop/pos"],
            batch["tl_stop/state"],
            batch["agent/type"],
            batch["agent/size"],
            batch["agent/goal"],
            batch["agent/dest"],
            **self.hparams.traffic_rule_checker,
        )
        # prepare features
        features = {
            "map_valid": input_feature_dict["map_feature_valid"],
            "map_feature": input_feature_dict["map_feature"],
            "tl_valid": input_feature_dict["tl_feature_valid"],
            "tl_feature": input_feature_dict["tl_feature"],
            "agent_type": batch["sc/agent_type"],
            "agent_size": batch["sc/agent_size"],
            # gt states for overriding
            "agent_valid": batch["agent/valid"],
            "vel": batch["agent/vel"],
            "acc": batch["agent/acc"],
            "yaw_rate": batch["agent/yaw_rate"],
            "agent_state": torch.cat([batch["agent/pos"], batch["agent/yaw_bbox"], batch["agent/spd"]], dim=-1),
        }
        rollout_buffer = self.rollout(
            features,
            latent=latent,
            goal=goal,
            goal_valid=goal_valid,
            mask_teacher_forcing=mask_teacher_forcing,
            rule_checker=rule_checker,
            step_start=self.hparams.time_step_sim_start,
            step_end=self.hparams.time_step_end,
            deterministic_latent=deterministic_latent,
            deterministic_action=deterministic_action,
            require_vis_dict=require_vis_dict,
        )
        return rollout_buffer

    def joint_future_pred(
        self,
        batch: Dict[str, Tensor],
        input_feature_dict: Dict[str, Tensor],
        latent: Optional[MyDist],
        goal: Optional[MyDist],
        goal_valid: Optional[Tensor],
        require_vis_dict: bool,
    ) -> Tuple[RolloutBuffer, Tensor, Tensor]:
        k_futures = self.hparams.n_joint_future
        # deterministic tensor for goal and latent: K=0 is deterministic
        deterministic = torch.zeros_like(batch["history/agent/valid"][:, 0])  # [n_batch, n_agent]
        deterministic = deterministic.repeat_interleave(k_futures, 0)  # [n_batch*k_futures, n_agent]
        deterministic[::k_futures] = True

        latent.repeat_interleave_(k_futures, 0)

        rule_checker_agent_dest = None
        rule_checker_agent_goal = None
        goal_sample = None
        if goal is not None:
            goal.repeat_interleave_(k_futures, 0)
            goal_sample = goal.sample(deterministic)
            goal_log_probs = goal.log_prob(goal_sample)
            goal_valid = goal_valid.repeat_interleave(k_futures, 0)

            # traffic rule checker
            if self.model.goal_manager.goal_attr_mode == "dest":
                rule_checker_agent_dest = goal_sample
            elif self.model.goal_manager.goal_attr_mode == "goal_xy":
                rule_checker_agent_goal = goal_sample

        if rule_checker_agent_dest is None and "agent/dest" in batch:
            rule_checker_agent_dest = batch["agent/dest"].repeat_interleave(k_futures, 0)

        if rule_checker_agent_goal is None and "agent/goal" in batch:
            rule_checker_agent_goal = batch["agent/goal"].repeat_interleave(k_futures, 0)

        rule_checker = TrafficRuleChecker(
            batch["map/boundary"].repeat_interleave(k_futures, 0),
            batch["map/valid"].repeat_interleave(k_futures, 0),
            batch["map/type"].repeat_interleave(k_futures, 0),
            batch["map/pos"].repeat_interleave(k_futures, 0),
            batch["map/dir"].repeat_interleave(k_futures, 0),
            batch["history/tl_stop/valid"].repeat_interleave(k_futures, 0),
            batch["history/tl_stop/pos"].repeat_interleave(k_futures, 0),
            batch["history/tl_stop/state"].repeat_interleave(k_futures, 0),
            batch["history/agent/type"].repeat_interleave(k_futures, 0),
            batch["history/agent/size"].repeat_interleave(k_futures, 0),
            rule_checker_agent_goal,  # ground truth goal
            rule_checker_agent_dest,  # pred dest
            **self.hparams.traffic_rule_checker,
        )

        # prepare features
        features = {
            "map_valid": input_feature_dict["map_feature_valid"],
            "map_feature": input_feature_dict["map_feature"],
            "tl_valid": input_feature_dict["tl_feature_valid"],
            "tl_feature": input_feature_dict["tl_feature"],
            "agent_type": batch["sc/agent_type"],
            "agent_size": batch["sc/agent_size"],
            # gt states for overriding
            "agent_valid": batch["agent/valid"],
            "vel": batch["agent/vel"],
            "acc": batch["agent/acc"],
            "yaw_rate": batch["agent/yaw_rate"],
            "agent_state": torch.cat([batch["agent/pos"], batch["agent/yaw_bbox"], batch["agent/spd"]], dim=-1),
        }
        for k in features.keys():
            features[k] = features[k].repeat_interleave(k_futures, 0)
        mask_teacher_forcing = self.teacher_forcing_joint_future_pred.get(features["agent_valid"], self.current_epoch)
        rollout_buffer = self.rollout(
            features,
            latent=latent,
            goal=goal_sample,
            goal_valid=goal_valid,
            mask_teacher_forcing=mask_teacher_forcing,
            rule_checker=rule_checker,
            step_start=self.hparams.time_step_sim_start,
            step_end=self.hparams.time_step_end,
            deterministic_latent=deterministic,
            deterministic_action=True,
            require_vis_dict=require_vis_dict,
        )
        rollout_buffer.flatten_repeat(self.hparams.n_joint_future)
        n_batch, n_agent, n_repeat, _ = rollout_buffer.valid.shape
        goal_log_probs = 0 if goal is None else goal_log_probs.view(n_batch, n_repeat, n_agent).transpose(1, 2)
        if self.model.goal_manager.goal_attr_mode == "dest":
            goal_sample = goal_sample.view(n_batch, n_repeat, n_agent).transpose(1, 2)
        elif self.model.goal_manager.goal_attr_mode == "goal_xy":
            _dim = goal_sample.shape[-1]
            goal_sample = goal_sample.view(n_batch, n_repeat, n_agent, _dim).transpose(1, 2)

        return rollout_buffer, goal_sample, goal_log_probs

    def validation_step(self, batch: Dict[str, Tensor], batch_idx: int) -> Dict:
        require_vis_dict = batch_idx < self.hparams.n_video_batch
        batch = self.pre_processing(batch)
        input_dict = {k.split("input/")[-1]: v for k, v in batch.items() if "input/" in k}
        latent_post_dict = {k.split("latent_post/")[-1]: v for k, v in batch.items() if "latent_post/" in k}
        latent_prior_dict = {k.split("latent_prior/")[-1]: v for k, v in batch.items() if "latent_prior/" in k}

        input_feature_dict = self.model.encode_input_features(**input_dict)
        latent_post_feature_dict = self.model.encode_input_features(**latent_post_dict)
        latent_prior_feature_dict = self.model.encode_input_features(**latent_prior_dict)

        # ! goal and destination
        goal_gt, goal_valid = self.model.goal_manager.get_gt_goal(
            agent_valid=input_dict["agent_valid"], gt_dest=batch["gt/dest"], gt_goal=batch["gt/goal"]
        )
        goal_pred = self.model.goal_manager.pred_goal(
            agent_type=batch["ref/agent_type"],
            map_type=batch["ref/map_type"],
            agent_state=batch["ref/agent_state"],
            **input_feature_dict,
        )

        # ! latents
        latent_post = self.model.latent_encoder(posterior=True, **latent_post_feature_dict)
        latent_prior = self.model.latent_encoder(**latent_prior_feature_dict)

        # ! reactive_replay: scene reconstruct given a complete episode, same setup as training
        rollout_buffer = self.reactive_replay(
            batch=batch,
            input_feature_dict=input_feature_dict,
            mask_teacher_forcing=self.teacher_forcing_reactive_replay.get(batch["gt/valid"], self.current_epoch),
            latent=latent_post,
            goal=goal_gt,
            goal_valid=goal_valid,
            deterministic_latent=True,
            deterministic_action=True,
            require_vis_dict=require_vis_dict,
        )
        rollout_buffer.flatten_repeat(1)
        self.err_metrics_reactive_replay(
            pred_valid=rollout_buffer.valid,
            pred_states=rollout_buffer.preds,
            gt_valid=batch["gt/valid"][:, self.hparams.time_step_sim_start :].transpose(1, 2),
            gt_states=batch["gt/state"][:, self.hparams.time_step_sim_start :].transpose(1, 2),
            override_masks=rollout_buffer.override_masks,
            agent_role=batch["ref/agent_role"],
        )
        self.rule_metrics_reactive_replay(
            valid=rollout_buffer.valid,
            override_masks=rollout_buffer.override_masks,
            outside_map=rollout_buffer.violations["outside_map"],
            collided=rollout_buffer.violations["collided"],
            run_road_edge=rollout_buffer.violations["run_road_edge"],
            run_red_light=rollout_buffer.violations["run_red_light"],
            passive=rollout_buffer.violations["passive"],
            goal_reached=rollout_buffer.violations["goal_reached"],
            dest_reached=rollout_buffer.violations["dest_reached"],
            agent_type=batch["ref/agent_type"],
        )
        self.train_metrics_reactive_replay(
            pred_valid=rollout_buffer.valid.squeeze(2),
            diffbar_rewards_valid=rollout_buffer.diffbar_rewards_valid.squeeze(2),
            diffbar_rewards=rollout_buffer.diffbar_rewards.squeeze(2),
            override_masks=rollout_buffer.override_masks.squeeze(2),
            agent_role=batch["ref/agent_role"],
            goal_valid=goal_valid,
            goal_pred=goal_pred,
            goal_gt=goal_gt,
            latent_post=latent_post,
            latent_prior=latent_prior,
        )

        # waymo metrics
        pred_dict = self.waymo_post_processing(
            valid=rollout_buffer.valid[:, :, 0].any(-1),
            scores=torch.ones_like(rollout_buffer.preds[:, :, :, 0, 0]),  # [n_batch, n_agent, 1],
            trajs=rollout_buffer.preds[:, :, :, rollout_buffer.step_future_start :],
            agent_type=batch["ref/agent_type"],
        )
        waymo_ops_inputs = self.womd_metrics_reactive_replay(batch, pred_dict["waymo_trajs"], pred_dict["waymo_scores"])
        self.womd_metrics_reactive_replay.aggregate_on_cpu(waymo_ops_inputs)
        self.womd_metrics_reactive_replay.reset()

        self.sub_womd_reactive_replay.add_to_submissions(
            waymo_trajs=pred_dict["waymo_trajs"],  # after nms
            waymo_scores=pred_dict["waymo_scores"],  # after nms
            mask_pred=batch["history/agent/role"][..., 2],
            object_id=batch["history/agent/object_id"],
            scenario_center=batch["scenario_center"],
            scenario_yaw=batch["scenario_yaw"],
            scenario_id=batch["scenario_id"],
        )
        if self.global_rank == 0 and require_vis_dict:
            _pred_goal, _pred_dest = None, None
            if self.model.goal_manager.goal_attr_mode == "dest":
                _pred_dest = goal_pred.sample(True).unsqueeze(2)
            elif self.model.goal_manager.goal_attr_mode == "goal_xy":
                _pred_goal = goal_pred.sample(True).unsqueeze(2)
            self.log_val_video(
                prefix="reactive_replay",
                batch_idx=batch_idx,
                batch=batch,
                buf=rollout_buffer,
                pred_goal=_pred_goal,
                pred_dest=_pred_dest,
                attn_video=False,
            )

        # ! joint_future_pred
        rollout_buffer, goal_sample, goal_log_probs = self.joint_future_pred(
            batch=batch,
            input_feature_dict=input_feature_dict,
            latent=latent_prior,
            goal=goal_pred,
            goal_valid=goal_valid,
            require_vis_dict=require_vis_dict,
        )
        self.err_metrics_joint_future_pred(
            pred_valid=rollout_buffer.valid,
            pred_states=rollout_buffer.preds,
            gt_valid=batch["gt/valid"][:, self.hparams.time_step_sim_start :].transpose(1, 2),
            gt_states=batch["gt/state"][:, self.hparams.time_step_sim_start :].transpose(1, 2),
            override_masks=rollout_buffer.override_masks,
            agent_role=batch["ref/agent_role"],
        )
        self.rule_metrics_joint_future_pred(
            valid=rollout_buffer.valid,
            override_masks=rollout_buffer.override_masks,
            outside_map=rollout_buffer.violations["outside_map"],
            collided=rollout_buffer.violations["collided"],
            run_road_edge=rollout_buffer.violations["run_road_edge"],
            run_red_light=rollout_buffer.violations["run_red_light"],
            passive=rollout_buffer.violations["passive"],
            goal_reached=rollout_buffer.violations["goal_reached"],
            dest_reached=rollout_buffer.violations["dest_reached"],
            agent_type=batch["ref/agent_type"],
        )

        pred_dict = self.waymo_post_processing(
            valid=rollout_buffer.valid[:, :, 0].any(-1),  # [n_scene, n_agent]
            scores=torch.exp(rollout_buffer.latent_log_probs[..., 0] + goal_log_probs),  # [n_batch, n_agent, n_repeat]
            trajs=rollout_buffer.preds[:, :, :, rollout_buffer.step_future_start :],
            agent_type=batch["ref/agent_type"],
        )
        waymo_ops_inputs = self.womd_metrics_joint_future_pred(
            batch, pred_dict["waymo_trajs"], pred_dict["waymo_scores"]
        )
        self.womd_metrics_joint_future_pred.aggregate_on_cpu(waymo_ops_inputs)
        self.womd_metrics_joint_future_pred.reset()

        self.sub_womd_joint_future_pred.add_to_submissions(
            waymo_trajs=pred_dict["waymo_trajs"],  # after nms
            waymo_scores=pred_dict["waymo_scores"],  # after nms
            mask_pred=batch["history/agent/role"][..., 2],
            object_id=batch["history/agent/object_id"],
            scenario_center=batch["scenario_center"],
            scenario_yaw=batch["scenario_yaw"],
            scenario_id=batch["scenario_id"],
        )

        if self.global_rank == 0 and require_vis_dict:
            _pred_goal, _pred_dest = None, None
            if self.model.goal_manager.goal_attr_mode == "dest":
                _pred_dest = goal_sample
            elif self.model.goal_manager.goal_attr_mode == "goal_xy":
                _pred_goal = goal_sample
            # log K=6 preds only for the first batch
            k_to_log = self.hparams.n_joint_future if batch_idx == 0 else 1
            self.log_val_video(
                prefix="joint_future_pred",
                batch_idx=batch_idx,
                batch=batch,
                buf=rollout_buffer,
                pred_scores=pred_dict["waymo_scores"],
                pred_goal=_pred_goal,
                pred_dest=_pred_dest,
                attn_video=False,
                k_to_log=k_to_log,
                as_goal_prior=goal_pred,
            )

    def validation_epoch_end(self, outputs):
        self.log("epoch", self.current_epoch, on_epoch=True)
        epoch_err_metrics_reactive_replay = self.err_metrics_reactive_replay.compute()
        for k, v in epoch_err_metrics_reactive_replay.items():
            self.log(k, v, on_epoch=True)
        self.err_metrics_reactive_replay.reset()
        epoch_rule_metrics_reactive_replay = self.rule_metrics_reactive_replay.compute()
        for k, v in epoch_rule_metrics_reactive_replay.items():
            self.log(k, v, on_epoch=True)
        self.rule_metrics_reactive_replay.reset()
        epoch_train_metrics_reactive_replay = self.train_metrics_reactive_replay.compute()
        for k, v in epoch_train_metrics_reactive_replay.items():
            self.log(k, v, on_epoch=True)
        self.train_metrics_reactive_replay.reset()
        epoch_err_metrics_joint_future_pred = self.err_metrics_joint_future_pred.compute()
        for k, v in epoch_err_metrics_joint_future_pred.items():
            self.log(k, v, on_epoch=True)
        self.err_metrics_joint_future_pred.reset()
        epoch_rule_metrics_joint_future_pred = self.rule_metrics_joint_future_pred.compute()
        for k, v in epoch_rule_metrics_joint_future_pred.items():
            self.log(k, v, on_epoch=True)
        self.rule_metrics_joint_future_pred.reset()

        epoch_womd_metrics_reactive_replay = self.womd_metrics_reactive_replay.compute_waymo_motion_metrics()
        for k, v in epoch_womd_metrics_reactive_replay.items():
            self.log(k, v, on_epoch=True)
        epoch_womd_metrics_joint_future_pred = self.womd_metrics_joint_future_pred.compute_waymo_motion_metrics()
        for k, v in epoch_womd_metrics_joint_future_pred.items():
            self.log(k, v, on_epoch=True)

        self.log(
            "val/loss",
            -epoch_womd_metrics_joint_future_pred[
                f"{self.womd_metrics_joint_future_pred.prefix}/mean_average_precision"
            ],
        )

        if self.global_rank == 0:
            self.sub_womd_reactive_replay.save_sub_files(self.logger[0])
            self.sub_womd_joint_future_pred.save_sub_files(self.logger[0])

    def log_val_video(
        self,
        prefix: str,
        batch_idx: int,
        batch: Dict[str, Tensor],
        buf,
        pred_scores: Optional[Tensor] = None,
        pred_goal: Optional[Tensor] = None,
        pred_dest: Optional[Tensor] = None,
        attn_video: bool = False,
        vis_eps_idx: List[int] = [],
        k_to_log: int = 1,
        as_goal_prior=None,
    ) -> Tuple[List[str], List[str]]:
        n_batch = batch["agent/valid"].shape[0]
        if len(vis_eps_idx) == 0:
            vis_eps_idx = range(n_batch)
        video_paths = []
        image_paths = []
        for idx in vis_eps_idx:
            video_dir = f"video_{batch_idx}-{idx}"
            _path = Path(video_dir)
            _path.mkdir(exist_ok=True, parents=True)
            episode_keys = [
                "agent/valid",
                "agent/pos",
                "agent/yaw_bbox",
                "agent/spd",
                "agent/role",
                "agent/size",
                "agent/goal",
                "map/valid",
                "map/type",
                "map/pos",
                "map/boundary",
                "tl_lane/valid",
                "tl_lane/state",
                "tl_lane/idx",
                "tl_stop/valid",
                "tl_stop/state",
                "tl_stop/pos",
                "tl_stop/dir",
                "episode_idx",
            ]
            episode = {}
            for k in episode_keys:
                episode[k] = batch[k][idx].cpu().numpy()
            # [n_agent, 2]
            episode["agent/dest"] = batch["agent/dest"][idx].cpu().numpy()

            for kf in range(k_to_log):
                # i = idx * k_futures + kf
                # [n_batch, n_step, n_agent, k_pred, 4] x,y,yaw,v -> [n_step, n_agent, 2]
                prediction = {
                    # [n_agent, n_step]
                    "agent/valid": buf.valid[idx, :, kf, buf.step_future_start :],
                    # [n_agent, n_step, 2]
                    "agent/pos": buf.preds[idx, :, kf, buf.step_future_start :, :2],
                    # n_agent, [n_ste,1]
                    "agent/yaw_bbox": buf.preds[idx, :, kf, buf.step_future_start :, [2]],
                    # [n_agent, n_step]
                    "speed": buf.preds[idx, :, kf, buf.step_future_start :, 3],
                    "rew_dr": buf.diffbar_rewards[idx, :, kf, buf.step_future_start :],
                    # [n_agent, n_step]
                    "lat_P": buf.latent_log_probs[idx, :, kf, buf.step_future_start :].float().exp(),
                    "act_P": buf.action_log_probs[idx, :, kf, buf.step_future_start :].float().exp(),
                }
                for k in prediction.keys():
                    prediction[k] = prediction[k].transpose(0, 1).cpu().numpy()
                prediction["step_current"] = self.hparams.time_step_current
                prediction["step_gt"] = self.hparams.time_step_gt
                prediction["step_end"] = self.hparams.time_step_end
                for k in buf.vis_dicts.keys():  # [n_step, n_agent, :]
                    prediction[k] = buf.vis_dicts[k][idx, :, kf, buf.step_future_start :].transpose(0, 1).numpy()
                for k in buf.violations.keys():  # [n_step, n_agent]
                    prediction[k] = buf.violations[k][idx, :, kf, buf.step_future_start :].transpose(0, 1).cpu().numpy()

                if pred_goal is not None:
                    prediction["agent/goal"] = pred_goal[idx, :, kf].cpu().numpy()
                if pred_dest is not None:
                    prediction["agent/dest"] = pred_dest[idx, :, kf].cpu().numpy()
                if pred_scores is not None:
                    prediction["score"] = pred_scores[idx, :, kf].cpu().numpy()

                vis_waymo = VisWaymo(
                    episode["map/valid"], episode["map/type"], episode["map/pos"], episode["map/boundary"]
                )
                video_paths_pred = vis_waymo.save_prediction_videos(f"{video_dir}/{prefix}_K{kf}", episode, prediction)
                video_paths += video_paths_pred
                if attn_video:
                    video_paths_attn = vis_waymo.save_attn_videos(f"{video_dir}/{prefix}_K{kf}", episode, prediction)
                    video_paths += video_paths_attn

                if (as_goal_prior is not None) and (kf == 0):
                    dest_img_paths = vis_waymo.get_dest_prob_image(
                        f"{video_dir}/dest_im", episode, prediction, as_goal_prior.probs[idx].cpu().numpy()
                    )
                    image_paths += dest_img_paths

        if self.logger is not None:
            for v_p in video_paths:
                self.logger[0].experiment.log({v_p: wandb.Video(v_p)}, commit=False)
            for i_p in image_paths:
                self.logger[0].experiment.log({i_p: wandb.Image(i_p)}, commit=False)
        return video_paths, image_paths

    def test_step(self, batch: Dict[str, Tensor], batch_idx: int) -> Dict:
        require_vis_dict = False
        batch = self.pre_processing(batch)
        input_dict = {k.split("input/")[-1]: v for k, v in batch.items() if "input/" in k}
        latent_prior_dict = {k.split("latent_prior/")[-1]: v for k, v in batch.items() if "latent_prior/" in k}
        input_feature_dict = self.model.encode_input_features(**input_dict)
        latent_prior_feature_dict = self.model.encode_input_features(**latent_prior_dict)

        # ! goal and destination
        goal_valid = input_dict["agent_valid"].any(1)
        goal_pred = self.model.goal_manager.pred_goal(
            agent_type=batch["ref/agent_type"],
            map_type=batch["ref/map_type"],
            agent_state=batch["ref/agent_state"],
            **input_feature_dict,
        )

        # ! latents
        latent_prior = self.model.latent_encoder(**latent_prior_feature_dict)

        # ! joint_future_pred
        for k in ["valid", "vel", "acc", "yaw_rate", "pos", "yaw_bbox", "spd", "size"]:
            batch[f"agent/{k}"] = batch[f"history/agent/{k}"]

        rollout_buffer, goal_sample, goal_log_probs = self.joint_future_pred(
            batch=batch,
            input_feature_dict=input_feature_dict,
            latent=latent_prior,
            goal=goal_pred,
            goal_valid=goal_valid,
            require_vis_dict=require_vis_dict,
        )
        pred_dict = self.waymo_post_processing(
            valid=rollout_buffer.valid[:, :, 0].any(-1),  # [n_scene, n_agent]
            scores=torch.exp(rollout_buffer.latent_log_probs[..., 0] + goal_log_probs),  # [n_batch, n_agent, n_repeat]
            trajs=rollout_buffer.preds[:, :, :, rollout_buffer.step_future_start :],
            agent_type=batch["ref/agent_type"],
        )

        self.sub_womd_joint_future_pred.add_to_submissions(
            waymo_trajs=pred_dict["waymo_trajs"],  # after nms
            waymo_scores=pred_dict["waymo_scores"],  # after nms
            mask_pred=batch["history/agent/role"][..., 2],
            object_id=batch["history/agent/object_id"],
            scenario_center=batch["scenario_center"],
            scenario_yaw=batch["scenario_yaw"],
            scenario_id=batch["scenario_id"],
        )

    def test_epoch_end(self, outputs):
        if self.global_rank == 0:
            self.sub_womd_joint_future_pred.save_sub_files(self.logger[0])

    def configure_optimizers(self):
        params = []
        params_goal = []
        for k, v in self.named_parameters():
            if "goal_predictor" in k:
                params_goal.append(v)
            else:
                params.append(v)
        optimizer = hydra.utils.instantiate(self.hparams.optimizer, params=params)
        if len(params_goal) > 0:
            optimizer.add_param_group({"params": params_goal, "lr": self.hparams.lr_goal})
        scheduler = {
            "scheduler": hydra.utils.instantiate(self.hparams.lr_scheduler, optimizer=optimizer),
            "monitor": "val/loss",
            "interval": "epoch",
            "frequency": self.trainer.check_val_every_n_epoch,
            "strict": True,
        }
        return [optimizer], [scheduler]

    def log_grad_norm(self, grad_norm_dict: Dict[str, float]) -> None:
        self.log_dict(grad_norm_dict, on_step=True, on_epoch=False, prog_bar=False, logger=True)

src/pl_modules/waymo_motion.py
zhejz/TrafficBots
# Licensed under the CC BY-NC 4.0 license (https://creativecommons.org/licenses/by-nc/4.0/)
import hydra
import torch
from omegaconf import DictConfig
from typing import List
from pytorch_lightning import seed_everything, LightningDataModule, LightningModule, Trainer, Callback
from pytorch_lightning.loggers import LightningLoggerBase
import os


def download_checkpoint(loggers, wb_ckpt) -> None:
    if os.environ.get("LOCAL_RANK", 0) == 0:
        artifact = loggers[0].experiment.use_artifact(wb_ckpt, type="model")
        artifact_dir = artifact.download("ckpt")


@hydra.main(config_path="../configs/", config_name="run.yaml")
def main(config: DictConfig) -> None:

    seed_everything(config.seed, workers=True)
    datamodule: LightningDataModule = hydra.utils.instantiate(config.datamodule)

    callbacks: List[Callback] = []
    if "callbacks" in config:
        for _, cb_conf in config.callbacks.items():
            callbacks.append(hydra.utils.instantiate(cb_conf))

    loggers: List[LightningLoggerBase] = []
    if "loggers" in config:
        for _, lg_conf in config.loggers.items():
            loggers.append(hydra.utils.instantiate(lg_conf))

    if config.resume.checkpoint is None:
        model: LightningModule = hydra.utils.instantiate(
            config.model, data_size=datamodule.tensor_size_train, _recursive_=False
        )
    else:
        download_checkpoint(loggers, config.resume.checkpoint)
        ckpt_path = "ckpt/model.ckpt"
        modelClass = hydra.utils.get_class(config.model._target_)

        model = modelClass.load_from_checkpoint(
            ckpt_path, wb_artifact=config.resume.checkpoint, **config.resume.model_overrides
        )

        if config.resume.resume_trainer and config.action == "fit":
            config.trainer.resume_from_checkpoint = ckpt_path

    # from pytorch_lightning.plugins import DDPPlugin
    # strategy = DDPPlugin(gradient_as_bucket_view=True)
    strategy = None
    if torch.cuda.device_count() > 1:
        strategy = "ddp"
    trainer: Trainer = hydra.utils.instantiate(
        config.trainer, strategy=strategy, callbacks=callbacks, logger=loggers, _convert_="partial"
    )
    if config.action == "fit":
        trainer.fit(model=model, datamodule=datamodule)
    elif config.action == "validate":
        trainer.validate(model=model, datamodule=datamodule)
    elif config.action == "test":
        trainer.test(model=model, datamodule=datamodule)
    else:
        raise NotImplementedError


if __name__ == "__main__":
    main()

src/run.py
zhejz/TrafficBots
# Licensed under the CC BY-NC 4.0 license (https://creativecommons.org/licenses/by-nc/4.0/)
from typing import Dict, List, Optional, Tuple
from torch import Tensor
import torch
from models.modules.distributions import MyDist


class RolloutBuffer:
    def __init__(self, step_start: int, step_end: int, step_current: int) -> None:
        """
        saves prediction, step [step_start,...,step_end]
        """
        self.step_start = step_start  # first step 0 in the rollout buffer corresponds to step_start in absolute time
        self.step_end = step_end  # last step in the rollout buffer corresponds to step_end in absolute time
        # this step in the rollout buffer corresponds to step_current+1 in absolute time
        self.step_future_start = step_current + 1 - step_start

        self.override_masks: List[Tensor] = []  # n_step * [n_batch, n_agent], bool

        # for logging
        self.valid: List[Tensor] = []  # n_step * [n_batch, n_agent], bool
        self.preds: List[Tensor] = []  # n_step * [n_batch, n_agent, 4], grad, x,y,yaw,v

        # for diffbar reward
        self.diffbar_rewards: List[Tensor] = []  # n_step * [n_batch, n_agent], grad
        self.diffbar_rewards_valid: List[Tensor] = []  # n_step * [n_batch, n_agent], bool

        # for VAE KL loss, len = n_step (traffic_sim) or 1(others that do not update latent)
        self.latents: List[Tuple[MyDist, MyDist]] = []
        self.latent_log_probs: List[Tensor] = []
        self.action_log_probs: List[Tensor] = []

        # for simulation metrics
        self.violations = {}  # Dict: n_step * [n_batch, n_agent], no_grad

        #  for visualizing video
        self.vis_dicts = {}

    def add(
        self,
        valid: Tensor,
        pred: Tensor,
        override_mask: Tensor,
        violation: Dict[str, Tensor],
        diffbar_reward: Optional[Tensor],
        diffbar_reward_valid: Optional[Tensor],
        latent_log_prob: Tensor,
        action_log_prob: Tensor,
        vis_dict: Dict[str, Tensor],
    ) -> None:
        self.valid.append(valid)  # [n_batch, n_agent]
        self.preds.append(pred)  # [n_batch, n_agent, 4]
        self.override_masks.append(override_mask)  # [n_batch, n_agent]

        if len(self.violations) == 0:
            self.violations = {k: [] for k in violation.keys()}
        for k, v in violation.items():
            self.violations[k].append(v)

        if diffbar_reward is not None:
            self.diffbar_rewards.append(diffbar_reward)  # [n_batch, n_agent]
            self.diffbar_rewards_valid.append(diffbar_reward_valid)

        self.latent_log_probs.append(latent_log_prob)  # [n_batch, n_agent]
        self.action_log_probs.append(action_log_prob)  # [n_batch, n_agent]

        if len(self.vis_dicts) == 0:
            self.vis_dicts = {k: [] for k in vis_dict.keys()}
        for k, v in vis_dict.items():
            self.vis_dicts[k].append(v)

    def finish(self) -> None:
        self.valid = torch.stack(self.valid, dim=2)  # [n_batch, n_agent, n_step]
        self.preds = torch.stack(self.preds, dim=2)  # [n_batch, n_agent, n_step, 4]
        self.override_masks = torch.stack(self.override_masks, dim=2)  # [n_batch, n_agent, n_step]

        # [n_batch, n_agent, n_step], no_grad
        for k in self.violations.keys():
            self.violations[k] = torch.stack(self.violations[k], dim=2)

        # [n_batch, n_agent, n_step]
        if len(self.diffbar_rewards) > 0:  # grad
            self.diffbar_rewards = torch.stack(self.diffbar_rewards, dim=2)
            self.diffbar_rewards_valid = torch.stack(self.diffbar_rewards_valid, dim=2)

        self.latent_log_probs = torch.stack(self.latent_log_probs, dim=2)  # [n_batch, n_agent, n_step]
        self.action_log_probs = torch.stack(self.action_log_probs, dim=2)  # [n_batch, n_agent, n_step]

        for k in self.vis_dicts.keys():
            self.vis_dicts[k] = torch.stack(self.vis_dicts[k], dim=2)

    def flatten_repeat(self, n_repeat: int) -> None:
        n_batch, n_agent, n_step, d_preds = self.preds.shape
        n_batch = n_batch // n_repeat

        # [n_batch, n_agent, n_repeat, n_step]
        self.valid = self.valid.view(n_batch, n_repeat, n_agent, n_step).transpose(1, 2)
        self.override_masks = self.override_masks.view(n_batch, n_repeat, n_agent, n_step).transpose(1, 2)

        # [n_batch, n_agent, n_repeat, n_step, 4]
        self.preds = self.preds.view(n_batch, n_repeat, n_agent, n_step, d_preds).transpose(1, 2)

        # [n_batch, n_agent, n_repeat, n_step], no_grad
        for k in self.violations.keys():
            self.violations[k] = self.violations[k].view(n_batch, n_repeat, n_agent, n_step).transpose(1, 2)

        # [n_batch, n_agent, n_repeat, n_step]
        if len(self.diffbar_rewards) > 0:  # grad
            self.diffbar_rewards = self.diffbar_rewards.view(n_batch, n_repeat, n_agent, n_step).transpose(1, 2)
            self.diffbar_rewards_valid = self.diffbar_rewards_valid.view(n_batch, n_repeat, n_agent, n_step).transpose(
                1, 2
            )

        # [n_batch, n_agent, n_repeat, n_step]
        self.latent_log_probs = self.latent_log_probs.view(n_batch, n_repeat, n_agent, n_step).transpose(1, 2)
        self.action_log_probs = self.action_log_probs.view(n_batch, n_repeat, n_agent, n_step).transpose(1, 2)

        for k in self.vis_dicts.keys():
            if "valid" in k:
                self.vis_dicts[k] = self.vis_dicts[k].view(n_batch, n_repeat, n_agent, n_step).transpose(1, 2)
            else:
                _dim = self.vis_dicts[k].shape[-1]
                self.vis_dicts[k] = self.vis_dicts[k].view(n_batch, n_repeat, n_agent, n_step, _dim).transpose(1, 2)

src/utils/buffer.py
zhejz/TrafficBots
# Licensed under the CC BY-NC 4.0 license (https://creativecommons.org/licenses/by-nc/4.0/)
import torch
from torch import Tensor
from typing import Optional, Dict, Tuple
from omegaconf import DictConfig
import hydra
from models.modules.distributions import MyDist
from utils.transform_utils import cast_rad


class Dynamics:
    def __init__(
        self, veh: DictConfig, ped: DictConfig, cyc: DictConfig, dt: float = 0.1, use_veh_dynamics_for_all: bool = False
    ) -> None:
        self.dt = dt
        self.action_dim = 2
        self.use_veh_dynamics_for_all = use_veh_dynamics_for_all
        self.state_keys = ["agent_state", "vel", "acc", "yaw_rate"]

        if self.use_veh_dynamics_for_all:
            self.agent_dynamics = hydra.utils.instantiate(veh, dt=dt)
        else:
            self.agent_dynamics = (
                hydra.utils.instantiate(veh, dt=dt),
                hydra.utils.instantiate(ped, dt=dt),
                hydra.utils.instantiate(cyc, dt=dt),
            )

    def init(
        self,
        agent_valid: Tensor,
        agent_state: Tensor,
        agent_size: Tensor,
        agent_type: Tensor,
        vel: Tensor,
        acc: Tensor,
        yaw_rate: Tensor,
    ) -> None:
        # constant
        self.agent_size = agent_size
        self.agent_type = agent_type
        # updating
        self.agent_valid = agent_valid  # [n_batch, n_agent], bool
        self.agent_killed = torch.zeros_like(agent_valid)  # [n_batch, n_agent], bool
        self.agent_state = agent_state  # [n_batch, n_agent, 4], x,y,yaw,v
        self.vel = vel  # [n_batch, n_agent, 2]
        self.acc = acc  # [n_batch, n_agent, 1]
        self.yaw_rate = yaw_rate  # [n_batch, n_agent, 1]

    def update(
        self,
        action_dist: MyDist,
        action_override: Optional[Tensor] = None,
        mask_action_override: Optional[Tensor] = None,
        deterministic: bool = True,
    ) -> Tuple[Tensor, Tensor]:
        """
        Args:
            action_dist: Gaussian or DiagGaussian, unbounded value range.
            action_override: [n_batch, n_agent, 2] in physical metrics.
            mask_action_override: [n_batch, n_agent] bool, for gym API.
        Updating attrs:
            self.agent_state: [n_batch, n_agent, 4], x,y,yaw,spd
            self.vel: [n_batch, n_agent, 2]
            self.acc: [n_batch, n_agent, 1]
            self.yaw_rate: [n_batch, n_agent, 1]

        Returns:
            agent_state: [n_batch, n_agent, 4], x,y,yaw,spd
            agent_valid: [n_batch, n_agent], bool
            action: [n_batch, n_agent, 2]
            action_log_prob: [n_batch, n_agent]
        """
        mask_type = ~self.agent_type  # [n_batch, n_agent, 3] reversed agent_type for masked_fill
        agent_invalid = ~self.agent_valid.unsqueeze(-1)
        # [n_batch, n_agent, 2], unbounded
        action_unbounded = action_dist.sample(deterministic)

        # [n_batch, n_agent]
        action_log_prob = action_dist.log_prob(action_unbounded.detach()).masked_fill(agent_invalid.squeeze(-1), 0)

        # action: unbounded value to value with physical metrics.
        if self.use_veh_dynamics_for_all:
            action = self.agent_dynamics.process_action(action_unbounded)
        else:
            action = 0
            # mask_type: [n_batch, n_agent, 3]
            for i in range(3):
                action += self.agent_dynamics[i].process_action(action_unbounded).masked_fill(mask_type[:, :, [i]], 0)
        action = action.masked_fill(agent_invalid, 0)

        # override actions: used for gym, for robot-vehicle. not used for now
        if action_override is not None:
            mask_action_override = mask_action_override & self.agent_valid  # [n_batch, n_agent], bool
            action = action.masked_fill(mask_action_override.unsqueeze(-1), 0)
            action[mask_action_override] = action_override[mask_action_override]  # physical metrics

        # dynamics update states using actions
        if self.use_veh_dynamics_for_all:
            state, vel, acc, yaw_rate = self.agent_dynamics.update(self.agent_state, action)
        else:
            state = 0  # [n_batch, n_agent, 4], x,y,yaw,spd
            vel = 0  # [n_batch, n_agent, 2]
            acc = 0  # [n_batch, n_agent, 1]
            yaw_rate = 0  # [n_batch, n_agent, 1]
            # mask_type: [n_batch, n_agent, 3]
            for i in range(3):
                _state, _vel, _acc, _yaw_rate = self.agent_dynamics[i].update(self.agent_state, action)
                state += _state.masked_fill(mask_type[:, :, [i]], 0)
                vel += _vel.masked_fill(mask_type[:, :, [i]], 0)
                acc += _acc.masked_fill(mask_type[:, :, [i]], 0)
                yaw_rate += _yaw_rate.masked_fill(mask_type[:, :, [i]], 0)

        # mask out invalid
        self.agent_state = state.masked_fill(agent_invalid, 0)
        self.agent_vel = vel.masked_fill(agent_invalid, 0)
        self.agent_acc = acc.masked_fill(agent_invalid, 0)
        self.agent_yaw_rate = yaw_rate.masked_fill(agent_invalid, 0)
        return self.agent_state, self.agent_valid, action, action_log_prob

    def override_states(
        self, state_override: Optional[Dict[str, Tensor]] = None, mask_state_override: Optional[Tensor] = None
    ) -> None:
        """For teacher forcing (existing agent) and spawn (new agent).
        Args:
            mask_state_override: [n_batch, n_agent] bool
            state_override["agent_state"]: [n_batch, n_agent, 4]
            state_override["vel]: [n_batch, n_agent, 2]
            state_override["acc"]: [n_batch, n_agent, 1]
            state_override["yaw_rate]: [n_batch, n_agent, 1]
        """
        if mask_state_override is not None:
            mask_state_override = mask_state_override & (~self.agent_killed)
            if mask_state_override.any():
                self.agent_valid = self.agent_valid | mask_state_override
                mask_state_override = mask_state_override.unsqueeze(-1)
                mask_old_state = ~mask_state_override
                if "agent_state" in state_override:
                    self.agent_state = self.agent_state.masked_fill(mask_state_override, 0)
                    self.agent_state = self.agent_state + state_override["agent_state"].masked_fill(mask_old_state, 0)
                if "vel" in state_override:
                    self.vel = self.vel.masked_fill(mask_state_override, 0)
                    self.vel = self.vel + state_override["vel"].masked_fill(mask_old_state, 0)
                if "acc" in state_override:
                    self.acc = self.acc.masked_fill(mask_state_override, 0)
                    self.acc = self.acc + state_override["acc"].masked_fill(mask_old_state, 0)
                if "yaw_rate" in state_override:
                    self.yaw_rate = self.yaw_rate.masked_fill(mask_state_override, 0)
                    self.yaw_rate = self.yaw_rate + state_override["yaw_rate"].masked_fill(mask_old_state, 0)

    @torch.no_grad()
    def kill(self, traffic_rule_violations: Dict[str, Tensor], gt_valid: Optional[Tensor] = None) -> None:
        """
        Args:
            traffic_rule_violations: at t
            gt_valid: [n_batch, n_agent] at t

        Update: will take affect at next update(), at t
            self.as_valid: [n_batch, n_agent], bool
        """
        mask_kill = traffic_rule_violations["outside_map_this_step"]
        if gt_valid is not None:  # do not kill agent that has gt_valid, such that the value can be computed at t+1.
            mask_kill = mask_kill & (~gt_valid)
        if mask_kill.any():
            self.agent_killed = self.agent_killed | mask_kill
            mask_survive = ~mask_kill
            self.agent_valid = self.agent_valid & mask_survive


class MultiPathPP:
    def __init__(self, dt: float, max_acc: float = 4, max_yaw_rate: float = 1, disable_neg_spd: bool = False) -> None:
        """
        max_acc:
        max_yaw_rate: veh=1rad/s, cyc=2.5rad/s, ped=5rad/s
        max_yaw_rate: veh=1.2rad/s, cyc=3rad/s, ped=6rad/s
        delta_theta per step (0.1sec), veh: 5m, 0.1rad, cyc: 2m, 0.25rad, ped:1m, 0.5rad
        """
        self.dt = dt
        self._max_acc = max_acc
        self._max_yaw_rate = max_yaw_rate
        self.disable_neg_spd = disable_neg_spd

    def init(self, *args, **kwargs) -> None:
        """Set prarameters for vehicle dynamics, called before each rollout."""
        pass

    def process_action(self, action: Tensor) -> Tensor:
        """
        Args:
            action: [n_batch, n_agent, 2] unbounded sample from Gaussian
        Returns:
            action: [n_batch, n_agent, 2], acc (m/s2), theta_rad (rad/s)
        """
        action = torch.tanh(action)
        action = torch.stack([action[..., 0] * self._max_acc, action[..., 1] * self._max_yaw_rate], dim=-1)
        return action

    def update(self, state: Tensor, action: Tensor) -> Tuple[Tensor, Tensor, Tensor, Tensor]:
        """
        Args:
            state: [n_batch, n_agent, 4] x,y,theta,spd
            action: [n_batch, n_agent, 2] acc(m/s2), yaw_rate(rad/s)

        Returns:
            as_pred: [n_batch, n_agent, 4], x,y,theta,spd
            vel: [n_batch, n_agent, 2]
            acc: [n_batch, n_agent, 1]
            yaw_rate [n_batch, n_agent, 1]
        """
        acc = action[:, :, 0]
        yaw_rate = action[:, :, 1]

        # [n_batch, n_agent]
        v_tilde = state[:, :, 3] + 0.5 * self.dt * acc
        theta_tilde = state[:, :, 2] + 0.5 * self.dt * yaw_rate
        cos_theta_tilde = torch.cos(theta_tilde)
        sin_theta_tilde = torch.sin(theta_tilde)

        # [n_batch, n_agent, 4]
        delta_state = torch.stack([v_tilde * cos_theta_tilde, v_tilde * sin_theta_tilde, yaw_rate, acc], dim=-1)
        as_pred = state + self.dt * delta_state
        if self.disable_neg_spd:
            as_pred[..., -1] = torch.relu(state[..., -1])

        vel = (as_pred[:, :, :2] - state[:, :, :2]) / self.dt
        acc = acc.unsqueeze(-1)
        yaw_rate = yaw_rate.unsqueeze(-1)
        return as_pred, vel, acc, yaw_rate


class StateIntegrator:
    def __init__(self, dt: float, max_v: float = 3) -> None:
        self.dt = dt
        self._max_v = max_v  # ped=3m/s

    def init(self, *args, **kwargs) -> None:
        pass

    def process_action(self, action: Tensor) -> Tensor:
        """
        Args:
            action: [n_batch, n_agent, 2] unbounded sample from Gaussian
        Returns:
            action: [n_batch, n_agent, 2], vx (m/s), vy (m/s)
        """
        action = torch.tanh(action) * self._max_v
        return action

    def update(self, state: Tensor, action: Tensor) -> Tuple[Tensor, Tensor, Tensor, Tensor]:
        """
        Args:
            state: [n_batch, n_agent, 4] x,y,theta,spd
            action: [n_batch, n_agent, 2] vx,vy in m/s

        Returns:
            as_pred: [n_batch, n_agent, 4], x,y,theta,spd
            vel: [n_batch, n_agent, 2]
            acc: [n_batch, n_agent, 1]
            yaw_rate [n_batch, n_agent, 1]
        """
        mask_theta_v = torch.zeros_like(state, dtype=torch.bool)
        mask_theta_v[:, :, 2:] = True

        # [n_batch, n_agent]
        vx = action[:, :, 0]
        vy = action[:, :, 1]
        theta = torch.atan2(vy, vx).detach()
        spd = torch.norm(action, dim=-1).detach()

        # [n_batch, n_agent, 4]
        delta_state = torch.stack([vx * self.dt, vy * self.dt, theta, spd], dim=-1)
        as_pred = state.masked_fill(mask_theta_v, 0) + delta_state

        acc = (spd - state[:, :, 3]) / self.dt
        acc = acc.unsqueeze(-1).detach()
        yaw_rate = cast_rad(theta - state[:, :, 2]) / self.dt
        yaw_rate = yaw_rate.unsqueeze(-1).detach()
        return as_pred, action, acc, yaw_rate

src/utils/dynamics.py
zhejz/TrafficBots
# Licensed under the CC BY-NC 4.0 license (https://creativecommons.org/licenses/by-nc/4.0/)
from typing import Dict, List, Tuple
import numpy as np
from scipy.interpolate import interp1d
from . import transform_utils

# "agent/cmd"
# STATIONARY = 0;
# STRAIGHT = 1;
# STRAIGHT_LEFT = 2;
# STRAIGHT_RIGHT = 3;
# LEFT_U_TURN = 4;
# LEFT_TURN = 5;
# RIGHT_U_TURN = 6;
# RIGHT_TURN = 7;
N_AGENT_CMD = 8


def get_polylines_from_polygon(polygon: np.ndarray) -> List[List[List]]:
    # polygon: [4, 3]
    l1 = np.linalg.norm(polygon[1, :2] - polygon[0, :2])
    l2 = np.linalg.norm(polygon[2, :2] - polygon[1, :2])

    def _pl_interp_start_end(start: np.ndarray, end: np.ndarray) -> List[List]:
        length = np.linalg.norm(start - end)
        unit_vec = (end - start) / length
        pl = []
        for i in range(int(length) + 1):  # 4.5 -> 5 [0,1,2,3,4]
            x, y, z = start + unit_vec * i
            pl.append([x, y, z])
        pl.append([end[0], end[1], end[2]])
        return pl

    # if l1 > l2:
    #     pl1 = _pl_interp_start_end((polygon[0] + polygon[3]) * 0.5, (polygon[1] + polygon[2]) * 0.5)
    # else:
    #     pl1 = _pl_interp_start_end((polygon[0] + polygon[1]) * 0.5, (polygon[2] + polygon[3]) * 0.5)
    # return [pl1, pl1[::-1]]

    if l1 > l2:
        pl1 = _pl_interp_start_end(polygon[0], polygon[1])
        pl2 = _pl_interp_start_end(polygon[2], polygon[3])
    else:
        pl1 = _pl_interp_start_end(polygon[0], polygon[3])
        pl2 = _pl_interp_start_end(polygon[2], polygon[1])
    return [pl1, pl1[::-1], pl2, pl2[::-1]]


def get_map_boundary(map_valid, map_pos):
    """
    Args:
        map_valid: [n_pl, 20],  # bool
        map_pos: [n_pl, 20, 3],  # float32
    Returns:
        map_boundary: [4]
    """
    pos = map_pos[map_valid]
    xmin = pos[:, 0].min()
    ymin = pos[:, 1].min()
    xmax = pos[:, 0].max()
    ymax = pos[:, 1].max()
    return np.array([xmin, xmax, ymin, ymax])


def classify_track(
    valid: np.ndarray,
    pos: np.ndarray,
    yaw: np.ndarray,
    spd: np.ndarray,
    kMaxSpeedForStationary: float = 2.0,  # (m/s)
    kMaxDisplacementForStationary: float = 5.0,  # (m)
    kMaxLateralDisplacementForStraight: float = 5.0,  # (m)
    kMinLongitudinalDisplacementForUTurn: float = -5.0,  # (m)
    kMaxAbsHeadingDiffForStraight: float = 0.5236,  # M_PI / 6.0
) -> int:
    """github.com/waymo-research/waymo-open-dataset/blob/master/waymo_open_dataset/metrics/motion_metrics_utils.cc
    Args:
        valid: [n_step], bool
        pos: [n_step, 2], x,y
        yaw: [n_step], float32
        spd: [n_step], float32
    Returns:
        traj_type: int in range(N_AGENT_CMD)
            # STATIONARY = 0;
            # STRAIGHT = 1;
            # STRAIGHT_LEFT = 2;
            # STRAIGHT_RIGHT = 3;
            # LEFT_U_TURN = 4;
            # LEFT_TURN = 5;
            # RIGHT_U_TURN = 6;
            # RIGHT_TURN = 7;
    """
    i0 = valid.argmax()
    i1 = len(valid) - 1 - np.flip(valid).argmax()

    x, y = pos[i1] - pos[i0]
    final_displacement = np.sqrt(x ** 2 + y ** 2)

    _c = np.cos(-yaw[i0])
    _s = np.sin(-yaw[i0])
    dx = x * _c - y * _s
    dy = x * _s + y * _c

    heading_diff = yaw[i1] - yaw[i0]
    max_speed = max(spd[i0], spd[i1])

    if max_speed < kMaxSpeedForStationary and final_displacement < kMaxDisplacementForStationary:
        return 0  # TrajectoryType::STATIONARY;

    if np.abs(heading_diff) < kMaxAbsHeadingDiffForStraight:
        if np.abs(dy) < kMaxLateralDisplacementForStraight:
            return 1  # TrajectoryType::STRAIGHT;
        if dy > 0:
            return 2  # TrajectoryType::STRAIGHT_LEFT
        else:
            return 3  # TrajectoryType::STRAIGHT_RIGHT

    if heading_diff < -kMaxAbsHeadingDiffForStraight and dy < 0:
        if dx < kMinLongitudinalDisplacementForUTurn:
            return 6  # TrajectoryType::RIGHT_U_TURN
        else:
            return 7  # TrajectoryType::RIGHT_TURN

    if dx < kMinLongitudinalDisplacementForUTurn:
        return 4  # TrajectoryType::LEFT_U_TURN;

    return 5  # TrajectoryType::LEFT_TURN;


def pack_episode_map(
    episode: Dict[str, np.ndarray],
    mf_id: List[int],
    mf_xyz: List[List[List[float]]],
    mf_type: List[int],
    mf_edge: List[List[int]],
    n_pl_max: int,
    n_nodes: int = 20,
) -> int:
    """
    Args:
        mf_id: [polyline]
        mf_xyz: [polyline, points, xyz]
        mf_type: [polyline]
        mf_edge: [edge, 2]
    """
    episode["map/valid"] = np.zeros([n_pl_max, n_nodes], dtype=bool)
    episode["map/id"] = np.zeros([n_pl_max], dtype=np.int64) - 1
    episode["map/type"] = np.zeros([n_pl_max], dtype=np.int64)
    episode["map/pos"] = np.zeros([n_pl_max, n_nodes, 3], dtype=np.float32)
    episode["map/dir"] = np.zeros([n_pl_max, n_nodes, 3], dtype=np.float32)
    episode["map/edge"] = np.array(mf_edge)

    pl_counter = 0
    for i_pl in range(len(mf_id)):
        pl_pos = np.array(mf_xyz[i_pl])
        pl_dir = np.diff(pl_pos, axis=0)
        polyline_len = pl_dir.shape[0]
        polyline_cuts = np.linspace(0, polyline_len, polyline_len // n_nodes + 1, dtype=int, endpoint=False)
        num_cuts = len(polyline_cuts)
        for idx_cut in range(num_cuts):
            idx_start = polyline_cuts[idx_cut]
            if idx_cut + 1 == num_cuts:
                # last slice
                idx_end = polyline_len
            else:
                idx_end = polyline_cuts[idx_cut + 1]

            episode["map/valid"][pl_counter, : idx_end - idx_start] = True
            episode["map/pos"][pl_counter, : idx_end - idx_start] = pl_pos[idx_start:idx_end]
            episode["map/dir"][pl_counter, : idx_end - idx_start] = pl_dir[idx_start:idx_end]
            episode["map/type"][pl_counter] = mf_type[i_pl]
            episode["map/id"][pl_counter] = mf_id[i_pl]
            pl_counter += 1
    return pl_counter


def pack_episode_agents(
    episode: Dict[str, np.ndarray],
    agent_id: List[int],
    agent_type: List[int],
    agent_states: List[List[List[float]]],
    agent_role: List[List[bool]],
    pack_all: bool,
    pack_history: bool,
    n_agent_max: int,
    step_current: int,
    n_agent_type: int = 3,
) -> None:
    """
    Args:
        agent_id: [agents]
        agent_type: [agents]
        agent_states: [agents, step, 10]; x,y,z,l,w,h,heading,vx,vy,valid
        agent_role: [agents, :], bool, [sdc=0]
    """
    n_step = len(agent_states[0])
    data_agent_valid = np.zeros([n_step, n_agent_max], dtype=bool)
    data_agent_pos = np.zeros([n_step, n_agent_max, 3], dtype=np.float32)
    data_agent_vel = np.zeros([n_step, n_agent_max, 2], dtype=np.float32)
    data_agent_spd = np.zeros([n_step, n_agent_max, 1], dtype=np.float32)
    data_agent_yaw_bbox = np.zeros([n_step, n_agent_max, 1], dtype=np.float32)

    data_agent_type = np.zeros([n_agent_max, n_agent_type], dtype=bool)
    data_agent_cmd = np.zeros([n_agent_max, N_AGENT_CMD], dtype=bool)
    data_agent_role = np.zeros([n_agent_max, len(agent_role[0])], dtype=bool)
    data_agent_size = np.zeros([n_agent_max, 3], dtype=np.float32)
    data_agent_goal = np.zeros([n_agent_max, 4], dtype=np.float32)
    data_agent_object_id = np.zeros([n_agent_max], dtype=np.int64) - 1

    for i in range(len(agent_id)):
        data_agent_type[i, agent_type[i]] = True
        data_agent_object_id[i] = agent_id[i]
        data_agent_role[i] = agent_role[i]

        length = 0.0
        width = 0.0
        height = 0.0
        count = 0.0
        for k in range(n_step):
            if agent_states[i][k][9]:
                data_agent_pos[k, i, 0] = agent_states[i][k][0]
                data_agent_pos[k, i, 1] = agent_states[i][k][1]
                data_agent_pos[k, i, 2] = agent_states[i][k][2]
                length += agent_states[i][k][3]
                width += agent_states[i][k][4]
                height += agent_states[i][k][5]
                data_agent_yaw_bbox[k, i, 0] = agent_states[i][k][6]
                data_agent_vel[k, i, 0] = agent_states[i][k][7]
                data_agent_vel[k, i, 1] = agent_states[i][k][8]
                data_agent_valid[k, i] = agent_states[i][k][9]

                count += data_agent_valid[k, i]

                spd = np.linalg.norm(data_agent_vel[k, i], axis=-1)

                spd_sign = np.sign(
                    np.cos(data_agent_yaw_bbox[k, i, 0]) * data_agent_vel[k, i, 0]
                    + np.sin(data_agent_yaw_bbox[k, i, 0]) * data_agent_vel[k, i, 1]
                )
                data_agent_spd[k, i, 0] = spd * spd_sign

                # set goal as the last valid state [x,y,theta,v]
                data_agent_goal[i, 0] = data_agent_pos[k, i, 0]
                data_agent_goal[i, 1] = data_agent_pos[k, i, 1]
                data_agent_goal[i, 2] = data_agent_yaw_bbox[k, i, 0]
                data_agent_goal[i, 3] = data_agent_spd[k, i, 0]

        cmd = classify_track(
            data_agent_valid[step_current:, i],
            data_agent_pos[step_current:, i, :2],
            data_agent_yaw_bbox[step_current:, i, 0],
            data_agent_spd[step_current:, i, 0],
        )
        data_agent_cmd[i, cmd] = True

        data_agent_size[i, 0] = length / count if count > 0 else 0.0
        data_agent_size[i, 1] = width / count if count > 0 else 0.0
        data_agent_size[i, 2] = height / count if count > 0 else 0.0

    # swap sdc to be the first agent
    sdc_track_index = np.where(data_agent_role[:, 0])[0][0]
    data_agent_valid[:, [0, sdc_track_index]] = data_agent_valid[:, [sdc_track_index, 0]]
    data_agent_pos[:, [0, sdc_track_index]] = data_agent_pos[:, [sdc_track_index, 0]]
    data_agent_vel[:, [0, sdc_track_index]] = data_agent_vel[:, [sdc_track_index, 0]]
    data_agent_spd[:, [0, sdc_track_index]] = data_agent_spd[:, [sdc_track_index, 0]]
    data_agent_yaw_bbox[:, [0, sdc_track_index]] = data_agent_yaw_bbox[:, [sdc_track_index, 0]]
    data_agent_object_id[[0, sdc_track_index]] = data_agent_object_id[[sdc_track_index, 0]]
    data_agent_type[[0, sdc_track_index]] = data_agent_type[[sdc_track_index, 0]]
    data_agent_role[[0, sdc_track_index]] = data_agent_role[[sdc_track_index, 0]]
    data_agent_size[[0, sdc_track_index]] = data_agent_size[[sdc_track_index, 0]]
    data_agent_goal[[0, sdc_track_index]] = data_agent_goal[[sdc_track_index, 0]]

    if pack_all:
        episode["agent/valid"] = data_agent_valid.copy()
        episode["agent/pos"] = data_agent_pos.copy()
        episode["agent/vel"] = data_agent_vel.copy()
        episode["agent/spd"] = data_agent_spd.copy()
        episode["agent/yaw_bbox"] = data_agent_yaw_bbox.copy()

        episode["agent/object_id"] = data_agent_object_id.copy()
        episode["agent/type"] = data_agent_type.copy()
        episode["agent/role"] = data_agent_role.copy()
        episode["agent/size"] = data_agent_size.copy()
        episode["agent/cmd"] = data_agent_cmd.copy()
        episode["agent/goal"] = data_agent_goal.copy()
    if pack_history:
        episode["history/agent/valid"] = data_agent_valid[: step_current + 1].copy()
        episode["history/agent/pos"] = data_agent_pos[: step_current + 1].copy()
        episode["history/agent/vel"] = data_agent_vel[: step_current + 1].copy()
        episode["history/agent/spd"] = data_agent_spd[: step_current + 1].copy()
        episode["history/agent/yaw_bbox"] = data_agent_yaw_bbox[: step_current + 1].copy()

        episode["history/agent/object_id"] = data_agent_object_id.copy()
        episode["history/agent/type"] = data_agent_type.copy()
        episode["history/agent/role"] = data_agent_role.copy()
        episode["history/agent/size"] = data_agent_size.copy()
        # no goal, no cmd
        invalid = ~(episode["history/agent/valid"].any(0))
        episode["history/agent/object_id"][invalid] = -1
        episode["history/agent/type"][invalid] = False
        episode["history/agent/size"][invalid] = 0
    return len(agent_id)


def pack_episode_traffic_lights(
    episode: Dict[str, np.ndarray],
    tl_lane_state: List[List[float]],
    tl_lane_id: List[List[int]],
    tl_stop_point: List[List[List[float]]],
    pack_all: bool,
    pack_history: bool,
    n_tl_max: int,
    step_current: int,
) -> int:
    """
    Args:
        tl_lane_state: [step, tl_lane]
        tl_lane_id: [step, tl_lane]
        tl_stop_point: [step, tl_lane, 3], x,y,z
    """
    n_step = len(tl_lane_state)
    data_tl_lane_valid = np.zeros([n_step, n_tl_max], dtype=bool)
    data_tl_lane_state = np.zeros([n_step, n_tl_max], dtype=np.int64)
    data_tl_lane_id = np.zeros([n_step, n_tl_max], dtype=np.int64) - 1
    data_tl_stop_pos = np.zeros([n_step, n_tl_max, 3], dtype=np.float32)

    for _step in range(n_step):
        n_tl_lane = len(tl_lane_state[_step])
        if n_tl_lane > 0:
            data_tl_lane_valid[_step, :n_tl_lane] = True
            data_tl_lane_state[_step, :n_tl_lane] = np.array(tl_lane_state[_step])
            data_tl_lane_id[_step, :n_tl_lane] = np.array(tl_lane_id[_step])
            data_tl_stop_pos[_step, :n_tl_lane] = np.array(tl_stop_point[_step])

    if pack_all:
        episode["tl_lane/valid"] = data_tl_lane_valid.copy()
        episode["tl_lane/state"] = data_tl_lane_state.copy()
        episode["tl_lane/id"] = data_tl_lane_id.copy()
        episode["tl_stop/pos"] = data_tl_stop_pos.copy()
    if pack_history:
        episode["history/tl_lane/valid"] = data_tl_lane_valid[: step_current + 1].copy()
        episode["history/tl_lane/state"] = data_tl_lane_state[: step_current + 1].copy()
        episode["history/tl_lane/id"] = data_tl_lane_id[: step_current + 1].copy()
        episode["history/tl_stop/pos"] = data_tl_stop_pos[: step_current + 1].copy()
    return data_tl_lane_valid.sum(1).max()


def center_at_sdc(
    episode: Dict[str, np.ndarray], rand_pos: float = -1, rand_yaw: float = -1
) -> Tuple[np.ndarray, float]:
    """episode
    # agent state
    "agent/valid": [n_step, N_AGENT_MAX],  # bool,
    "agent/pos": [n_step, N_AGENT_MAX, 3],  # float32
    "agent/vel": [n_step, N_AGENT_MAX, 2],  # float32, v_x, v_y
    "agent/yaw_bbox": [n_step, N_AGENT_MAX, 1],  # float32, yaw of the bbox heading
    "agent/goal": [N_AGENT_MAX, 4],  # float32: [x, y, theta, v]
    # map
    "map/valid": [N_PL_MAX, n_pl_nodes],  # bool
    "map/pos": [N_PL_MAX, n_pl_nodes, 3],  # float32
    "map/dir": [N_PL_MAX, n_pl_nodes, 3],  # float32
    # traffic light
    "tl_lane/valid: [step, n_tl], bool
    "tl_stop/pos: [step, n_tl, 3], x,y,z
    """
    prefix = []
    if "agent/pos" in episode:
        prefix.append("")
    if "history/agent/valid" in episode:
        prefix.append("history/")

    sdc_center = episode[prefix[0] + "agent/pos"][0, 0, :2].copy()
    sdc_yaw = episode[prefix[0] + "agent/yaw_bbox"][0, 0, 0].copy()

    if rand_pos > 0:
        sdc_center[0] += np.random.uniform(-rand_pos, rand_pos)
        sdc_center[1] += np.random.uniform(-rand_pos, rand_pos)
    if rand_yaw > 0:
        sdc_yaw += np.random.uniform(-rand_yaw, rand_yaw)

    to_sdc_se3 = transform_utils.get_transformation_matrix(sdc_center, sdc_yaw)  # for points
    to_sdc_yaw = transform_utils.get_yaw_from_se2(to_sdc_se3)  # for vector
    to_sdc_so2 = transform_utils.get_so2_from_se2(to_sdc_se3)  # for angle

    # map
    episode["map/pos"][..., :2][episode["map/valid"]] = transform_utils.transform_points(
        episode["map/pos"][..., :2][episode["map/valid"]], to_sdc_se3
    )
    episode["map/dir"][..., :2][episode["map/valid"]] = transform_utils.transform_points(
        episode["map/dir"][..., :2][episode["map/valid"]], to_sdc_so2
    )

    for pf in prefix:
        # agent: pos, vel, yaw_bbox
        episode[pf + "agent/pos"][..., :2][episode[pf + "agent/valid"]] = transform_utils.transform_points(
            episode[pf + "agent/pos"][..., :2][episode[pf + "agent/valid"]], to_sdc_se3
        )
        episode[pf + "agent/vel"][episode[pf + "agent/valid"]] = transform_utils.transform_points(
            episode[pf + "agent/vel"][episode[pf + "agent/valid"]], to_sdc_so2
        )
        episode[pf + "agent/yaw_bbox"][episode[pf + "agent/valid"]] += to_sdc_yaw
        # traffic light: [step, tl, 3]
        key_tl = pf + "tl_stop/pos"
        if key_tl in episode:
            episode[key_tl][..., :2][episode[pf + "tl_lane/valid"]] = transform_utils.transform_points(
                episode[key_tl][..., :2][episode[pf + "tl_lane/valid"]], to_sdc_se3
            )
        if pf == "":
            # goal: x, y, theta
            goal_valid = episode["agent/valid"].any(axis=0)
            episode["agent/goal"][..., :2][goal_valid] = transform_utils.transform_points(
                episode["agent/goal"][..., :2][goal_valid], to_sdc_se3
            )
            episode["agent/goal"][..., 2][goal_valid] += to_sdc_yaw

    return sdc_center, sdc_yaw


def filter_episode_traffic_lights(episode: Dict[str, np.ndarray]) -> None:
    """Filter traffic light based on map valid
    Args: episode
        # map
        "map/valid": [N_PL_MAX, 20],  # bool
        "map/id": [N_PL_MAX],  # int, with -1
        "map/type": [N_PL_MAX],  # int, >= 0
        "map/pos": [N_PL_MAX, 20, 3]
        "map/dir": [N_PL_MAX, 20, 3]
        # traffic light
        "tl_lane/valid": [n_step, N_TL_MAX],  # bool
        "tl_lane/state": [n_step, N_TL_MAX],  # >= 0
        "tl_lane/id": [n_step, N_TL_MAX],  # with -1
        "tl_stop/pos": [n_step, N_TL_MAX, 3]
    """
    prefix = []
    if "tl_lane/valid" in episode:
        prefix.append("")
    if "history/tl_lane/valid" in episode:
        prefix.append("history/")
    for pf in prefix:
        n_step, n_tl_max = episode[pf + "tl_lane/valid"].shape
        for tl_step in range(n_step):
            for tl_idx in range(n_tl_max):
                if episode[pf + "tl_lane/valid"][tl_step, tl_idx]:
                    tl_lane_map_id = episode["map/id"] == episode[pf + "tl_lane/id"][tl_step, tl_idx]
                    if episode["map/valid"][tl_lane_map_id].sum() == 0:
                        episode[pf + "tl_lane/valid"][tl_step, tl_idx] = False


def filter_episode_map(episode: Dict[str, np.ndarray], n_pl: int, thresh_map: float, thresh_z: float = 3):
    """
    Args: episode
        "agent/valid": [n_step, N_AGENT_MAX], bool,
        "agent/pos": [n_step, N_AGENT_MAX, 3], float32
        "agent/role": [N_AGENT_MAX, :], bool, one hot
        "map/valid": [N_PL_MAX, n_pl_nodes], bool
        "map/id": [N_PL_MAX], int, with -1
        "map/type": [N_PL_MAX], int, >= 0
        "map/pos": [N_PL_MAX, n_pl_nodes, 3]
        "map/dir": [N_PL_MAX, n_pl_nodes, 3]
    """
    # ! filter "map/" based on distance to relevant agents, based on the history only
    if "agent/valid" in episode:
        relevant_agent = episode["agent/role"].any(-1)
        agent_valid_relevant = episode["agent/valid"][:11, relevant_agent]
        agent_pos_relevant = episode["agent/pos"][:11, relevant_agent]
    elif "history/agent/valid" in episode:
        relevant_agent = episode["history/agent/role"].any(-1)
        agent_valid_relevant = episode["history/agent/valid"][:, relevant_agent]
        agent_pos_relevant = episode["history/agent/pos"][:, relevant_agent]
    agent_pos_relevant = agent_pos_relevant[agent_valid_relevant]  # [N, 3]

    xmin = agent_pos_relevant[:, 0].min()
    xmax = agent_pos_relevant[:, 0].max()
    ymin = agent_pos_relevant[:, 1].min()
    ymax = agent_pos_relevant[:, 1].max()
    x_thresh = max(xmax - xmin, thresh_map)
    y_thresh = max(ymax - ymin, thresh_map)

    old_map_valid = episode["map/valid"].copy()
    episode["map/valid"] = episode["map/valid"] & (episode["map/pos"][..., 0] > xmin - x_thresh)
    episode["map/valid"] = episode["map/valid"] & (episode["map/pos"][..., 0] < xmax + x_thresh)
    episode["map/valid"] = episode["map/valid"] & (episode["map/pos"][..., 1] > ymin - y_thresh)
    episode["map/valid"] = episode["map/valid"] & (episode["map/pos"][..., 1] < ymax + y_thresh)
    if thresh_z > 0:
        zmin = agent_pos_relevant[:, 2].min()
        zmax = agent_pos_relevant[:, 2].max()
        z_thresh = max(zmax - zmin, thresh_z)
        episode["map/valid"] = episode["map/valid"] & (episode["map/pos"][..., 2] > zmin - z_thresh)
        episode["map/valid"] = episode["map/valid"] & (episode["map/pos"][..., 2] < zmax + z_thresh)

    # assert episode["map/valid"].sum() > 0
    if episode["map/valid"].any(1).sum() < 10:
        # in waymo training 83954, agent and map z axis is off, filter using zmin/zmax will remove all polylines.
        print("something is wrong, check this episode out!")
        episode["map/valid"] = old_map_valid

    # ! filter "map/" that has too little valid entries
    mask_map_short = episode["map/valid"].sum(1) <= 3
    episode["map/valid"][mask_map_short] = False

    # ! still too many polylines: filter polylines based on distance to relevant agents
    dist_thresh_map = thresh_map
    while episode["map/valid"].any(1).sum() > n_pl:
        mask_map_remain = episode["map/valid"].any(1)
        for i in range(episode["map/valid"].shape[0]):
            if mask_map_remain[i]:
                pl_pos = episode["map/pos"][i, :, :][episode["map/valid"][i, :]]
                close_to_agent = (
                    min(
                        np.linalg.norm(agent_pos_relevant - pl_pos[0], axis=1).min(),
                        np.linalg.norm(agent_pos_relevant - pl_pos[-1], axis=1).min(),
                    )
                    < dist_thresh_map
                )
                if not close_to_agent:
                    episode["map/valid"][i, :] = False
                if episode["map/valid"].any(1).sum() == n_pl:
                    break
        dist_thresh_map = dist_thresh_map * 0.5


def repack_episode_traffic_lights(
    episode: Dict[str, np.ndarray], episode_reduced: Dict[str, np.ndarray], n_tl: int, n_tl_state: int
) -> None:
    """
    Args: episode
        # map
        "map/valid": [N_PL_MAX, 20],  # bool
        "map/id": [N_PL_MAX],  # int, with -1
        "map/type": [N_PL_MAX],  # int, >= 0
        "map/pos": [N_PL_MAX, 20, 3]
        "map/dir": [N_PL_MAX, 20, 3]
        # traffic light
        "tl_lane/valid": [n_step, N_TL_MAX],  # bool
        "tl_lane/state": [n_step, N_TL_MAX],  # >= 0
        "tl_lane/id": [n_step, N_TL_MAX],  # with -1
        "tl_stop/pos": [n_step, N_TL_MAX, 3], # x,y,z
    Returns:
        "tl_lane/valid": [n_step, N_TL_MAX],  # bool
        "tl_lane/state": [n_step, N_TL_MAX],  # >= 0
        "tl_lane/id": [n_step, N_TL_MAX],  # with -1
        "tl_stop/valid": [n_step, N_TL_MAX],  # bool
        "tl_stop/pos": [n_step, N_TL_MAX, 3], # x,y,z
        "tl_stop/dir": [n_step, N_TL_MAX, 3], # x,y,z
        "tl_stop/state": [n_step, N_TL_MAX],  # >= 0
    """
    prefix = []
    if "tl_lane/valid" in episode:
        prefix.append("")
    if "history/tl_lane/valid" in episode:
        prefix.append("history/")
    for pf in prefix:
        n_step, n_tl_max = episode[pf + "tl_lane/valid"].shape
        # tl_lane
        episode_reduced[pf + "tl_lane/valid"] = np.zeros([n_step, n_tl], dtype=bool)  # bool
        episode_reduced[pf + "tl_lane/state"] = np.zeros([n_step, n_tl], dtype=np.int64)  # will be one_hot
        episode_reduced[pf + "tl_lane/idx"] = np.zeros([n_step, n_tl], dtype=np.int64) - 1  # int, -1 means not valid
        # tl_stop
        episode_reduced[pf + "tl_stop/valid"] = np.zeros([n_step, n_tl_max], dtype=bool)  # bool
        episode_reduced[pf + "tl_stop/state"] = np.zeros([n_step, n_tl_max], dtype=np.int64)  # will be one_hot
        episode_reduced[pf + "tl_stop/pos"] = np.zeros([n_step, n_tl_max, 2], dtype=np.float32)  # x,y
        episode_reduced[pf + "tl_stop/dir"] = np.zeros([n_step, n_tl_max, 2], dtype=np.float32)  # x,y
        for i in range(n_step):
            counter_tl = 0
            counter_tl_stop = 0
            for j in range(n_tl_max):
                if episode[pf + "tl_lane/valid"][i, j]:
                    lane_idx = np.where(episode_reduced["map/id"] == episode[pf + "tl_lane/id"][i, j])[0]
                    n_lanes = lane_idx.shape[0]
                    # assert counter_tl + n_lanes <= N_TL, print("counter_tl, n_lanes:", counter_tl, n_lanes)
                    episode_reduced[pf + "tl_lane/valid"][i, counter_tl : counter_tl + n_lanes] = True
                    episode_reduced[pf + "tl_lane/state"][i, counter_tl : counter_tl + n_lanes] = episode[
                        pf + "tl_lane/state"
                    ][i, j]
                    episode_reduced[pf + "tl_lane/idx"][i, counter_tl : counter_tl + n_lanes] = lane_idx
                    counter_tl += n_lanes
                    # tl_stop
                    episode_reduced[pf + "tl_stop/valid"][i, counter_tl_stop] = True
                    episode_reduced[pf + "tl_stop/state"][i, counter_tl_stop] = episode[pf + "tl_lane/state"][i, j]
                    episode_reduced[pf + "tl_stop/pos"][i, counter_tl_stop] = episode[pf + "tl_stop/pos"][i, j, :2]
                    episode_reduced[pf + "tl_stop/dir"][i, counter_tl_stop] = episode_reduced["map/dir"][
                        lane_idx[0], 0, :2
                    ]
                    counter_tl_stop += 1

        # one_hot "tl_lane/state" and "tl_stop/state": [N_AGENT, N_AGENT_TYPE], bool
        episode_reduced[pf + "tl_lane/state"] = np.eye(n_tl_state, dtype=bool)[episode_reduced[pf + "tl_lane/state"]]
        episode_reduced[pf + "tl_stop/state"] = np.eye(n_tl_state, dtype=bool)[episode_reduced[pf + "tl_stop/state"]]

        episode_reduced[pf + "tl_lane/state"] = (
            episode_reduced[pf + "tl_lane/state"] & episode_reduced[pf + "tl_lane/valid"][:, :, None]
        )
        episode_reduced[pf + "tl_stop/state"] = (
            episode_reduced[pf + "tl_stop/state"] & episode_reduced[pf + "tl_stop/valid"][:, :, None]
        )


def repack_episode_map(
    episode: Dict[str, np.ndarray], episode_reduced: Dict[str, np.ndarray], n_pl: int, n_pl_type: int
) -> int:
    """
    Args: episode
        # map
        "map/valid": [N_PL_MAX, 20],  # bool
        "map/id": [N_PL_MAX],  # int, with -1
        "map/type": [N_PL_MAX],  # int, >= 0
        "map/pos": [N_PL_MAX, 20, 3]
        "map/dir": [N_PL_MAX, 20, 3]
    """
    n_pl_nodes = episode["map/valid"].shape[1]
    episode_reduced["map/valid"] = np.zeros([n_pl, n_pl_nodes], dtype=bool)  # bool
    episode_reduced["map/type"] = np.zeros([n_pl], dtype=np.int64)  # will be one_hot
    episode_reduced["map/pos"] = np.zeros([n_pl, n_pl_nodes, 2], dtype=np.float32)  # x,y
    episode_reduced["map/dir"] = np.zeros([n_pl, n_pl_nodes, 2], dtype=np.float32)  # x,y
    episode_reduced["map/id"] = np.zeros([n_pl], dtype=np.int64) - 1

    map_valid_mask = episode["map/valid"].any(1)
    n_pl_valid = map_valid_mask.sum()
    episode_reduced["map/valid"][:n_pl_valid] = episode["map/valid"][map_valid_mask]
    episode_reduced["map/type"][:n_pl_valid] = episode["map/type"][map_valid_mask]
    episode_reduced["map/pos"][:n_pl_valid] = episode["map/pos"][map_valid_mask, :, :2]
    episode_reduced["map/dir"][:n_pl_valid] = episode["map/dir"][map_valid_mask, :, :2]
    episode_reduced["map/id"][:n_pl_valid] = episode["map/id"][map_valid_mask]
    # one_hot "map/type": [N_PL, N_PL_TYPE], bool
    episode_reduced["map/type"] = np.eye(n_pl_type, dtype=bool)[episode_reduced["map/type"]]
    episode_reduced["map/type"] = episode_reduced["map/type"] & episode_reduced["map/valid"].any(-1, keepdims=True)


def repack_episode_agents_no_sim(
    episode: Dict[str, np.ndarray],
    episode_reduced: Dict[str, np.ndarray],
    mask_no_sim: np.ndarray,
    n_agent_no_sim: int,
    prefix: str,
) -> None:
    n_step = episode[prefix + "agent/valid"].shape[0]
    episode_reduced[prefix + "agent_no_sim/valid"] = np.zeros([n_step, n_agent_no_sim], dtype=bool)
    episode_reduced[prefix + "agent_no_sim/pos"] = np.zeros([n_step, n_agent_no_sim, 2], dtype=np.float32)
    episode_reduced[prefix + "agent_no_sim/z"] = np.zeros([n_step, n_agent_no_sim, 1], dtype=np.float32)
    episode_reduced[prefix + "agent_no_sim/vel"] = np.zeros([n_step, n_agent_no_sim, 2], dtype=np.float32)
    episode_reduced[prefix + "agent_no_sim/spd"] = np.zeros([n_step, n_agent_no_sim, 1], dtype=np.float32)
    episode_reduced[prefix + "agent_no_sim/yaw_bbox"] = np.zeros([n_step, n_agent_no_sim, 1], dtype=np.float32)
    episode_reduced[prefix + "agent_no_sim/object_id"] = np.zeros([n_agent_no_sim], dtype=np.int64) - 1
    episode_reduced[prefix + "agent_no_sim/type"] = np.zeros([n_agent_no_sim, 3], dtype=bool)
    episode_reduced[prefix + "agent_no_sim/size"] = np.zeros([n_agent_no_sim, 3], dtype=np.float32)
    # no role, no cmd, no goal
    for i, idx in enumerate(np.where(mask_no_sim)[0]):
        episode_reduced[prefix + "agent_no_sim/valid"][:, i] = episode[prefix + "agent/valid"][:, idx]
        episode_reduced[prefix + "agent_no_sim/pos"][:, i] = episode[prefix + "agent/pos"][:, idx, :2]
        episode_reduced[prefix + "agent_no_sim/z"][:, i] = episode[prefix + "agent/pos"][:, idx, 2:3]
        episode_reduced[prefix + "agent_no_sim/vel"][:, i] = episode[prefix + "agent/vel"][:, idx]
        episode_reduced[prefix + "agent_no_sim/spd"][:, i] = episode[prefix + "agent/spd"][:, idx]
        episode_reduced[prefix + "agent_no_sim/yaw_bbox"][:, i] = episode[prefix + "agent/yaw_bbox"][:, idx]
        episode_reduced[prefix + "agent_no_sim/object_id"][i] = episode[prefix + "agent/object_id"][idx]
        episode_reduced[prefix + "agent_no_sim/type"][i] = episode[prefix + "agent/type"][idx]
        episode_reduced[prefix + "agent_no_sim/size"][i] = episode[prefix + "agent/size"][idx]


def repack_episode_agents(
    episode: Dict[str, np.ndarray],
    episode_reduced: Dict[str, np.ndarray],
    mask_sim: np.ndarray,
    n_agent: int,
    prefix: str,
    dim_veh_lanes: List[int] = [],
    dim_cyc_lanes: List[int] = [],
    dim_ped_lanes: List[int] = [],
    dest_no_pred: bool = False,
) -> None:
    """fill episode_reduced["history/agent/"], episode_reduced["agent/"]
    Args: episode
        # agent state
        "agent/valid": [n_step, N_AGENT_MAX], bool,
        "agent/pos": [n_step, N_AGENT_MAX, 3], float32
        "agent/vel": [n_step, N_AGENT_MAX, 2], float32, v_x, v_y
        "agent/spd": [n_step, N_AGENT_MAX, 1], float32
        "agent/yaw_bbox": [n_step, N_AGENT_MAX, 1], float32, yaw of the bbox heading
        # agent attribute
        "agent/type": [N_AGENT_MAX, 3], bool, one hot [Vehicle=0, Pedestrian=1, Cyclist=2]
        "agent/cmd": [N_AGENT_MAX, N_AGENT_CMD], bool, one hot
        "agent/role": [N_AGENT_MAX, :], bool, one hot
        "agent/size": [N_AGENT_MAX, 3], float32: [length, width, height]
        "agent/goal": [N_AGENT_MAX, 4], float32: [x, y, theta, v]
        # map
        "map/valid": [N_PL_MAX, n_pl_nodes], bool
        "map/id": [N_PL_MAX], int, with -1
        "map/type": [N_PL_MAX], int, >= 0
        "map/pos": [N_PL_MAX, n_pl_nodes, 3]
        "map/dir": [N_PL_MAX, n_pl_nodes, 3]
    Returns: episode_reduced
    """
    n_step = episode[prefix + "agent/valid"].shape[0]
    # agent state
    episode_reduced[prefix + "agent/valid"] = np.zeros([n_step, n_agent], dtype=bool)  # bool,
    episode_reduced[prefix + "agent/pos"] = np.zeros([n_step, n_agent, 2], dtype=np.float32)  # x,y
    episode_reduced[prefix + "agent/z"] = np.zeros([n_step, n_agent, 1], dtype=np.float32)  # z
    episode_reduced[prefix + "agent/vel"] = np.zeros([n_step, n_agent, 2], dtype=np.float32)  # v_x, v_y, in m/s
    episode_reduced[prefix + "agent/spd"] = np.zeros([n_step, n_agent, 1], dtype=np.float32)  # m/s, signed
    # m/s2, acc[t] = (spd[t]-spd[t-1])/dt
    episode_reduced[prefix + "agent/acc"] = np.zeros([n_step, n_agent, 1], dtype=np.float32)
    episode_reduced[prefix + "agent/yaw_bbox"] = np.zeros([n_step, n_agent, 1], dtype=np.float32)  # [-pi, pi]
    # yaw_rate[t] = (yaw[t]-yaw[t-1])/dt
    episode_reduced[prefix + "agent/yaw_rate"] = np.zeros([n_step, n_agent, 1], dtype=np.float32)
    # agent attribute
    episode_reduced[prefix + "agent/object_id"] = np.zeros([n_agent], dtype=np.int64) - 1
    episode_reduced[prefix + "agent/type"] = np.zeros([n_agent, 3], dtype=bool)
    # one hot [sdc=0, interest=1, predict=2]
    episode_reduced[prefix + "agent/role"] = np.zeros([n_agent, episode[prefix + "agent/role"].shape[-1]], dtype=bool)
    episode_reduced[prefix + "agent/size"] = np.zeros([n_agent, 3], dtype=np.float32)  # float32 [length, width, height]

    if prefix == "":
        episode_reduced["agent/cmd"] = np.zeros([n_agent, N_AGENT_CMD], dtype=bool)
        episode_reduced["agent/goal"] = np.zeros([n_agent, 4], dtype=np.float32)  # float32 [x, y, theta, v]
        episode_reduced["agent/dest"] = np.zeros([n_agent], dtype=np.int64)  # index to "map/valid" in [0, N_PL]
        # ! map info for finding dest
        n_pl, n_pl_node = episode_reduced["map/valid"].shape
        mask_veh_lane = (episode_reduced["map/type"][:, dim_veh_lanes].any(axis=-1, keepdims=True)) & episode_reduced[
            "map/valid"
        ]
        pos_veh_lane = episode_reduced["map/pos"][mask_veh_lane]  # [?, 2]
        dir_veh_lane = episode_reduced["map/dir"][mask_veh_lane]  # [?, 2]
        dir_veh_lane = dir_veh_lane / np.linalg.norm(dir_veh_lane, axis=-1, keepdims=True)
        map_id_veh_lane = episode_reduced["map/id"][:, None].repeat(n_pl_node, 1)[mask_veh_lane]
        pl_idx_veh_lane = np.arange(n_pl)[:, None].repeat(n_pl_node, 1)[mask_veh_lane]
        # cyc_lane
        mask_cyc_lane = (episode_reduced["map/type"][:, dim_cyc_lanes].any(axis=-1, keepdims=True)) & episode_reduced[
            "map/valid"
        ]
        pos_cyc_lane = episode_reduced["map/pos"][mask_cyc_lane]  # [?, 2]
        dir_cyc_lane = episode_reduced["map/dir"][mask_cyc_lane]  # [?, 2]
        dir_cyc_lane = dir_cyc_lane / np.linalg.norm(dir_cyc_lane, axis=-1, keepdims=True)
        pl_idx_cyc_lane = np.arange(n_pl)[:, None].repeat(n_pl_node, 1)[mask_cyc_lane]
        # road_edge
        mask_road_edge = (episode_reduced["map/type"][:, dim_ped_lanes].any(axis=-1, keepdims=True)) & episode_reduced[
            "map/valid"
        ]
        pos_road_edge = episode_reduced["map/pos"][mask_road_edge]  # [?, 2]
        pl_idx_road_edge = np.arange(n_pl)[:, None].repeat(n_pl_node, 1)[mask_road_edge]

    for i, idx in enumerate(np.where(mask_sim)[0]):
        valid = episode[prefix + "agent/valid"][:, idx]
        if valid.sum() > 1:
            valid_steps = np.where(valid)[0]
            step_start = valid_steps[0]
            step_end = valid_steps[-1]

            f_pos = interp1d(valid_steps, episode[prefix + "agent/pos"][valid, idx, :2], axis=0)
            f_z = interp1d(valid_steps, episode[prefix + "agent/pos"][valid, idx, 2:3], axis=0)
            f_vel = interp1d(valid_steps, episode[prefix + "agent/vel"][valid, idx], axis=0)
            f_spd = interp1d(valid_steps, episode[prefix + "agent/spd"][valid, idx], axis=0)
            f_yaw_bbox = interp1d(
                valid_steps, np.unwrap(episode[prefix + "agent/yaw_bbox"][valid, idx], axis=0), axis=0
            )

            x = np.arange(step_start, step_end + 1)
            x_spd = f_spd(x)
            x_yaw = f_yaw_bbox(x)
            episode_reduced[prefix + "agent/valid"][step_start : step_end + 1, i] = True
            episode_reduced[prefix + "agent/pos"][step_start : step_end + 1, i] = f_pos(x)
            episode_reduced[prefix + "agent/z"][step_start : step_end + 1, i] = f_z(x)
            episode_reduced[prefix + "agent/vel"][step_start : step_end + 1, i] = f_vel(x)
            episode_reduced[prefix + "agent/spd"][step_start : step_end + 1, i] = x_spd
            episode_reduced[prefix + "agent/yaw_bbox"][step_start : step_end + 1, i] = x_yaw

            x_acc = np.diff(x_spd, axis=0) / 0.1
            episode_reduced[prefix + "agent/acc"][step_start + 1 : step_end + 1, i] = x_acc
            x_yaw_rate = np.diff(x_yaw, axis=0) / 0.1
            episode_reduced[prefix + "agent/yaw_rate"][step_start + 1 : step_end + 1, i] = x_yaw_rate

        else:
            valid_step = np.where(valid)[0][0]
            episode_reduced[prefix + "agent/valid"][valid_step, i] = True
            episode_reduced[prefix + "agent/pos"][valid_step, i] = episode[prefix + "agent/pos"][valid_step, idx, :2]
            episode_reduced[prefix + "agent/z"][valid_step, i] = episode[prefix + "agent/pos"][valid_step, idx, 2:3]
            episode_reduced[prefix + "agent/vel"][valid_step, i] = episode[prefix + "agent/vel"][valid_step, idx]
            episode_reduced[prefix + "agent/spd"][valid_step, i] = episode[prefix + "agent/spd"][valid_step, idx]
            episode_reduced[prefix + "agent/yaw_bbox"][valid_step, i] = episode[prefix + "agent/yaw_bbox"][
                valid_step, idx
            ]

        episode_reduced[prefix + "agent/object_id"][i] = episode[prefix + "agent/object_id"][idx]
        episode_reduced[prefix + "agent/type"][i] = episode[prefix + "agent/type"][idx]
        episode_reduced[prefix + "agent/role"][i] = episode[prefix + "agent/role"][idx]
        episode_reduced[prefix + "agent/size"][i] = episode[prefix + "agent/size"][idx]

        if prefix == "":
            episode_reduced["agent/goal"][i] = episode["agent/goal"][idx]
            episode_reduced["agent/cmd"][i] = episode["agent/cmd"][idx]
            episode_reduced["agent/dest"][i] = find_dest(
                episode_reduced["agent/type"][i],
                episode_reduced["agent/goal"][i],
                episode["map/edge"],
                pos_veh_lane,
                dir_veh_lane,
                map_id_veh_lane,
                pl_idx_veh_lane,
                pos_cyc_lane,
                dir_cyc_lane,
                pl_idx_cyc_lane,
                pos_road_edge,
                pl_idx_road_edge,
                dest_no_pred,
            )


def find_dest(
    agent_type,  # one_hot [3]
    agent_goal,  # [x, y, theta, v]
    map_edge,  # [?, 2] id0 -> id1, or id0 -> -1
    pos_veh_lane,  # [?, 2]
    dir_veh_lane,  # [?, 2]
    map_id_veh_lane,  # [?]
    pl_idx_veh_lane,  # [?]
    pos_cyc_lane,  # [?, 2]
    dir_cyc_lane,  # [?, 2]
    pl_idx_cyc_lane,  # [?, 2]
    pos_road_edge,  # [?, 2]
    pl_idx_road_edge,  # [?]
    no_pred,
):
    goal_yaw = agent_goal[2]
    goal_heading = np.array([np.cos(goal_yaw), np.sin(goal_yaw)])
    goal_pos = agent_goal[:2]
    if no_pred:
        extended_goal_pos = goal_pos
    else:
        extended_goal_pos = agent_goal[:2] + goal_heading * agent_goal[3] * 5  # 5 seconds with constant speed
    if agent_type[0]:  # veh
        dist_pos = np.linalg.norm((pos_veh_lane - goal_pos), axis=1)
        dist_rot = np.dot(dir_veh_lane, goal_heading)
        candidate_lanes = (dist_pos < 3) & (dist_rot > 0)
        if candidate_lanes.any():  # associate to a lane, extend with map topology
            if no_pred:
                idx_dest = pl_idx_veh_lane[candidate_lanes][np.argmin(dist_pos[candidate_lanes])]
            else:
                dest_map_id = map_id_veh_lane[candidate_lanes][np.argmin(dist_pos[candidate_lanes])]
                next_map_id = dest_map_id
                counter = 0
                continue_extend = True
                while continue_extend:
                    next_edges = np.where(map_edge[:, 0] == next_map_id)[0]
                    dest_map_id, next_map_id = map_edge[np.random.choice(next_edges)]
                    counter += 1
                    # if (
                    #     (next_map_id not in map_id_veh_lane)
                    #     or ((len(next_edges) > 1) and (counter > 3))
                    #     or (counter > 6)
                    # ):
                    if (
                        (next_map_id not in map_id_veh_lane)
                        or ((len(next_edges) > 1) and (counter > 1))
                        or (counter > 3)
                    ):
                        continue_extend = False
                idx_dest = pl_idx_veh_lane[np.where(map_id_veh_lane == dest_map_id)[0][-1]]
        else:  # not associate to a lane, use road edge
            idx_dest = pl_idx_road_edge[np.linalg.norm((pos_road_edge - extended_goal_pos), axis=1).argmin()]
    elif agent_type[1]:  # ped
        idx_dest = pl_idx_road_edge[np.linalg.norm((pos_road_edge - extended_goal_pos), axis=1).argmin()]
    elif agent_type[2]:  # cyc
        dist_pos = np.linalg.norm((pos_cyc_lane - extended_goal_pos), axis=1)
        dist_rot = np.dot(dir_cyc_lane, goal_heading)
        candidate_lanes = (dist_pos < 3) & (dist_rot > 0)
        if candidate_lanes.any():  # associate to a bike lane, extend with constant vel and find bike lane
            idx_dest = pl_idx_cyc_lane[candidate_lanes][np.argmin(dist_pos[candidate_lanes])]
        else:  # not associate to a lane, use road edge
            idx_dest = pl_idx_road_edge[np.linalg.norm((pos_road_edge - extended_goal_pos), axis=1).argmin()]
    return idx_dest


def filter_episode_agents(
    episode: Dict[str, np.ndarray],
    episode_reduced: Dict[str, np.ndarray],
    n_agent: int,
    prefix: str,
    dim_veh_lanes: List[int],
    dist_thresh_agent: float,
    step_current: int,
) -> Tuple[float, np.ndarray, np.ndarray]:
    """
    Args: episode
        # agent state
        "agent/valid": [n_step, N_AGENT_MAX], bool,
        "agent/pos": [n_step, N_AGENT_MAX, 3], float32
        "agent/vel": [n_step, N_AGENT_MAX, 2], float32, v_x, v_y
        "agent/spd": [n_step, N_AGENT_MAX, 1], float32
        "agent/yaw_bbox": [n_step, N_AGENT_MAX, 1], float32, yaw of the bbox heading
        # agent attribute
        "agent/type": [N_AGENT_MAX, 3], bool, one hot [Vehicle=0, Pedestrian=1, Cyclist=2]
        "agent/cmd": [N_AGENT_MAX, N_AGENT_CMD], bool, one hot
        "agent/role": [N_AGENT_MAX, :], bool, one hot
        "agent/size": [N_AGENT_MAX, 3], float32: [length, width, height]
        "agent/goal": [N_AGENT_MAX, 4], float32: [x, y, theta, v]
        # map
        "map/valid": [N_PL_MAX, 20], bool
        "map/id": [N_PL_MAX], int, with -1
        "map/type": [N_PL_MAX], int, >= 0
        "map/pos": [N_PL_MAX, 20, 3]
        "map/dir": [N_PL_MAX, 20, 3]
    """
    n_agent_max = episode[prefix + "agent/valid"].shape[1]
    agent_valid = episode[prefix + "agent/valid"].copy()
    relevant_agent = episode[prefix + "agent/role"].any(-1)
    agent_valid_relevant = agent_valid[:, relevant_agent]
    agent_pos_relevant = episode[prefix + "agent/pos"][:, relevant_agent]
    agent_pos_relevant = agent_pos_relevant[agent_valid_relevant][:, :2]  # [N, 2]
    thresh_spd = 2 if prefix == "" else 0.5

    # ! filter agents not seen in the history.
    mask_agent_history_not_seen = (~relevant_agent) & ~(agent_valid[: step_current + 1].any(axis=0))
    agent_valid = agent_valid & ~mask_agent_history_not_seen[None, :]

    # if prefix == "":
    #     mask_agent_short = (agent_valid.sum(axis=0) <= 20) & ~(relevant_agent)
    #     agent_valid[:, mask_agent_short] = False

    # ! filter agents that have small displacement and large distance to relevant agents or map polylines
    mask_agent_still = (
        (episode[prefix + "agent/spd"][..., 0].sum(axis=0) * 0.1 < thresh_spd)
        & (~relevant_agent)
        & (agent_valid.any(axis=0))
    )
    lane_pos = episode_reduced["map/pos"][episode_reduced["map/valid"]]
    for i in range(n_agent_max):
        # if mask_agent_still[i]:
        if mask_agent_still[i] and (agent_valid.any(axis=0).sum() > n_agent):
            agent_poses = episode[prefix + "agent/pos"][:, i, :2][agent_valid[:, i]]
            start_pos = agent_poses[0]
            end_pos = agent_poses[-1]
            far_to_relevant_agent = (np.linalg.norm(agent_pos_relevant - start_pos, axis=1).min() > 20) & (
                np.linalg.norm(agent_pos_relevant - end_pos, axis=1).min() > 20
            )
            far_to_lane = (np.linalg.norm(lane_pos - start_pos, axis=1).min() > 20) & (
                np.linalg.norm(lane_pos - end_pos, axis=1).min() > 20
            )
            if far_to_relevant_agent & far_to_lane:
                agent_valid[:, i] = False

    # ! filter parking vehicles that have large distance to relevant agents and lanes, Testing only
    mask_veh_lane = (episode_reduced["map/type"][:, dim_veh_lanes].any(axis=-1, keepdims=True)) & episode_reduced[
        "map/valid"
    ]
    pos_veh_lane = episode_reduced["map/pos"][mask_veh_lane]  # [?, 2]
    dir_veh_lane = episode_reduced["map/dir"][mask_veh_lane]  # [?, 2]
    dir_veh_lane = dir_veh_lane / np.linalg.norm(dir_veh_lane, axis=-1, keepdims=True)

    mask_vehicle_still = (
        (episode[prefix + "agent/spd"][..., 0].sum(axis=0) * 0.1 < thresh_spd)
        & (~relevant_agent)
        & (agent_valid.any(axis=0))
        & (episode[prefix + "agent/type"][:, 0])
    )
    for i in range(n_agent_max):
        if mask_vehicle_still[i] and (agent_valid.any(axis=0).sum() > n_agent):
            agent_pos = episode[prefix + "agent/pos"][:, i, :2][agent_valid[:, i]][-1]
            agent_yaw = episode[prefix + "agent/yaw_bbox"][:, i, 0][agent_valid[:, i]][-1]
            agent_heading = np.array([np.cos(agent_yaw), np.sin(agent_yaw)])

            dist_pos = np.linalg.norm((pos_veh_lane - agent_pos), axis=1)
            dist_rot = np.dot(dir_veh_lane, agent_heading)
            candidate_lanes = (dist_pos < 3) & (dist_rot > 0)
            not_associate_to_lane = ~(candidate_lanes.any())

            far_to_relevant_agent = (np.linalg.norm(agent_pos_relevant - start_pos, axis=1).min() > 10) & (
                np.linalg.norm(agent_pos_relevant - end_pos, axis=1).min() > 10
            )
            if far_to_relevant_agent & not_associate_to_lane:
                agent_valid[:, i] = False

    # ! filter vehicles that have small displacement and large yaw change. Training only.
    if prefix == "" and (agent_valid.any(axis=0).sum() > n_agent):
        yaw_diff = np.abs(transform_utils.cast_rad(np.diff(episode["agent/yaw_bbox"][..., 0], axis=0))) * (
            agent_valid[:-1] & agent_valid[1:]
        )
        max_yaw_diff = yaw_diff.max(axis=0)
        mask_large_yaw_diff_veh = ((episode["agent/spd"][..., 0].sum(axis=0) * 0.1 < 6) & (max_yaw_diff > 0.5)) | (
            max_yaw_diff > 1.5
        )
        mask_large_yaw_diff_veh = mask_large_yaw_diff_veh & (episode["agent/type"][:, 0])
        mask_large_yaw_diff_ped_cyc = ((episode["agent/spd"][..., 0].sum(axis=0) * 0.1 < 1) & (max_yaw_diff > 0.5)) | (
            max_yaw_diff > 1.5
        )
        mask_large_yaw_diff_ped_cyc = mask_large_yaw_diff_ped_cyc & (episode["agent/type"][:, 1:].any(-1))
        mask_agent_large_yaw_change = (
            (mask_large_yaw_diff_veh | mask_large_yaw_diff_ped_cyc) & (~relevant_agent) & (agent_valid.any(axis=0))
        )
        agent_valid[:, mask_agent_large_yaw_change] = False

    # ! still too many agents: filter agents based on distance to relevant agents
    while agent_valid.any(axis=0).sum() > n_agent:
        mask_agent_remain = ~(relevant_agent) & (agent_valid.any(axis=0))
        for i in range(n_agent_max):
            if mask_agent_remain[i]:
                agent_poses = episode[prefix + "agent/pos"][:, i, :2][agent_valid[:, i]]
                close_to_relevant_agent = (
                    min(
                        np.linalg.norm(agent_pos_relevant - agent_poses[0], axis=1).min(),
                        np.linalg.norm(agent_pos_relevant - agent_poses[-1], axis=1).min(),
                    )
                    < dist_thresh_agent
                )
                if not close_to_relevant_agent:
                    agent_valid[:, i] = False
                if agent_valid.any(axis=0).sum() == n_agent:
                    break
        dist_thresh_agent = dist_thresh_agent * 0.5
        # if dist_thresh_agent < 1.0:
        #     raise RuntimeError(f"n_agent={agent_valid.any(axis=0).sum()}, n_agent_max={n_agent_max}")

    mask_sim = agent_valid.any(0)
    mask_no_sim = episode[prefix + "agent/valid"].any(0) & (~mask_sim)
    return mask_sim, mask_no_sim

src/utils/pack_h5.py
zhejz/TrafficBots
# Licensed under the CC BY-NC 4.0 license (https://creativecommons.org/licenses/by-nc/4.0/)
import torch
from torch import Tensor, nn


class PositionalEmbedding(nn.Module):
    def __init__(self, dim: int, theta: float = 10000):
        super().__init__()
        assert dim % 2 == 0
        self.dim = dim
        freqs = 1.0 / (theta ** (torch.arange(0, dim, 2)[: (dim // 2)].float() / dim))
        # [dim]
        freqs = freqs.repeat_interleave(2, 0)
        self.register_buffer("freqs", freqs)

    def forward(self, x: Tensor):
        """
        Args:
            x: [...]
        Returns:
            pos_enc: [..., dim]
        """
        # [..., dim]
        pos_enc = x.unsqueeze(-1) * self.freqs.view([1] * x.dim() + [-1])
        pos_enc = torch.cat([torch.cos(pos_enc[..., ::2]), torch.sin(pos_enc[..., 1::2])], dim=-1)
        # pos_enc[..., ::2] = torch.cos(pos_enc[..., ::2])
        # pos_enc[..., 1::2] = torch.sin(pos_enc[..., 1::2])
        return pos_enc


class PositionalEmbeddingRad(nn.Module):
    def __init__(self, dim: int):
        """
        if dim=2, then just [cos(theta), sin(theta)]
        """
        super().__init__()
        assert dim % 2 == 0
        self.dim = dim
        # [dim]: [1,1,2,2,4,4,8,8]
        # freqs = 2 ** (torch.arange(0, dim // 2).float())
        # [dim]: [1,1,2,2,3,3,4,4]
        freqs = torch.arange(0, dim // 2) + 1.0
        freqs = freqs.repeat_interleave(2, 0)
        self.register_buffer("freqs", freqs)

    def forward(self, x: Tensor):
        """
        Args:
            x: [...], in rad
        Returns:
            pos_enc: [..., dim]
        """
        # [..., dim]
        pos_enc = x.unsqueeze(-1) * self.freqs.view([1] * x.dim() + [-1])
        pos_enc = torch.cat([torch.cos(pos_enc[..., ::2]), torch.sin(pos_enc[..., 1::2])], dim=-1)
        # pos_enc[..., ::2] = torch.cos(pos_enc[..., ::2])
        # pos_enc[..., 1::2] = torch.sin(pos_enc[..., 1::2])
        return pos_enc

src/utils/pos_emb.py
zhejz/TrafficBots
# Licensed under the CC BY-NC 4.0 license (https://creativecommons.org/licenses/by-nc/4.0/)
import torch
from torch import Tensor, nn
from utils.pos_emb import PositionalEmbedding, PositionalEmbeddingRad


class PosePE(nn.Module):
    def __init__(self, mode: str, pe_dim: int = 256, theta_xy: float = 1e3, theta_cs: float = 1e1):
        super().__init__()
        self.mode = mode
        if self.mode == "xy_dir":
            self.out_dim = 4
        elif self.mode == "mpa_pl":
            self.out_dim = 7
        elif self.mode == "pe_xy_unit_dir":
            self.out_dim = pe_dim + 2
            self.pe_xy = PositionalEmbedding(dim=pe_dim // 2, theta=theta_xy)
        elif self.mode == "pe_xy_dir":
            self.out_dim = pe_dim
            self.pe_xy = PositionalEmbedding(dim=pe_dim // 4, theta=theta_xy)
            self.pe_dir = PositionalEmbedding(dim=pe_dim // 4, theta=theta_cs)
        elif self.mode == "pe_xy_yaw":
            self.out_dim = pe_dim
            self.pe_xy = PositionalEmbedding(dim=pe_dim // 4, theta=theta_xy)
            self.pe_yaw = PositionalEmbeddingRad(dim=pe_dim // 2)
        else:
            raise NotImplementedError

    def forward(self, xy: Tensor, dir: Tensor):
        """
        Args: input either dir or yaw.
            xy: [..., 2]
            dir: cos/sin [..., 2] or yaw [..., 1]

        Returns:
            pos_out: [..., self.out_dim]
        """
        if self.mode == "xy_dir":
            if dir.shape[-1] == 1:
                dir = torch.cat([dir.cos(), dir.sin()], dim=-1)
            pos_out = torch.cat([xy, dir], dim=-1)
        elif self.mode == "mpa_pl":
            if dir.shape[-1] == 1:
                dir = torch.cat([dir.cos(), dir.sin()], dim=-1)
            pos_out = self.encode_polyline(xy, dir)
        elif self.mode == "pe_xy_unit_dir":
            if dir.shape[-1] == 1:
                dir = torch.cat([dir.cos(), dir.sin()], dim=-1)
            pos_out = torch.cat([self.pe_xy(xy[..., 0]), self.pe_xy(xy[..., 1]), dir[..., [0]], dir[..., [1]]], dim=-1)
        elif self.mode == "pe_xy_dir":
            if dir.shape[-1] == 1:
                dir = torch.cat([dir.cos(), dir.sin()], dim=-1)
            pos_out = torch.cat(
                [self.pe_xy(xy[..., 0]), self.pe_xy(xy[..., 1]), self.pe_dir(dir[..., 0]), self.pe_dir(dir[..., 1])],
                dim=-1,
            )
        elif self.mode == "pe_xy_yaw":
            if dir.shape[-1] == 1:
                dir = dir.squeeze(-1)
            else:
                dir = torch.atan2(dir[..., 1], dir[..., 0])
            pos_out = torch.cat([self.pe_xy(xy[..., 0]), self.pe_xy(xy[..., 1]), self.pe_yaw(dir)], dim=-1)
        return pos_out

    @staticmethod
    def encode_polyline(pos: Tensor, dir: Tensor) -> Tensor:
        """
        Args: pos and dir with respect to the agent
            pos: [..., 2]
            dir: [..., 2]

        Returns:
            pl_feature: [..., 7]
        """
        eps = torch.finfo(pos.dtype).eps
        # [n_scene, n_target, n_map, n_pl_node, 2]
        segments_start = pos
        segment_vec = dir
        # [n_scene, n_target, n_map, n_pl_node]
        segment_proj = (-segments_start * segment_vec).sum(-1) / ((segment_vec * segment_vec).sum(-1) + eps)
        # [n_scene, n_target, n_map, n_pl_node, 2]
        closest_points = segments_start + torch.clamp(segment_proj, min=0, max=1).unsqueeze(-1) * segment_vec
        # [n_scene, n_target, n_map, n_pl_node, 1]
        r_norm = torch.norm(closest_points, dim=-1, keepdim=True)
        segment_vec_norm = torch.norm(segment_vec, dim=-1, keepdim=True)
        pl_feature = torch.cat(
            [
                r_norm,  # 1
                closest_points / (r_norm + eps),  # 2
                segment_vec / (segment_vec_norm + eps),  # 2
                segment_vec_norm,  # 1
                torch.norm(segments_start + segment_vec - closest_points, dim=-1, keepdim=True),  # 1
            ],
            dim=-1,
        )
        return pl_feature

src/utils/pose_pe.py
zhejz/TrafficBots
# Licensed under the CC BY-NC 4.0 license (https://creativecommons.org/licenses/by-nc/4.0/)
from typing import Tuple
from torch import Tensor
import torch
from omegaconf import DictConfig
from models.metrics.loss import AngularError


class DifferentiableReward:
    def __init__(
        self,
        l_pos: DictConfig,
        l_rot: DictConfig,
        l_spd: DictConfig,
        w_collision: float,
        use_il_loss: bool,
        reduce_collsion_with_max: bool,
    ):
        # traffic_rule
        self.w_collision = w_collision
        self.reduce_collsion_with_max = reduce_collsion_with_max

        # imitation
        self.use_il_loss = use_il_loss
        if self.use_il_loss:
            self.il_l_pos = getattr(torch.nn, l_pos.criterion)(reduction="none")
            self.il_w_pos = l_pos.weight
            self.il_l_rot = AngularError(l_rot.criterion, l_rot.angular_type)
            self.il_w_rot = l_rot.weight
            self.il_l_spd = getattr(torch.nn, l_spd.criterion)(reduction="none")
            self.il_w_spd = l_spd.weight

    def get(
        self, agent_valid: Tensor, agent_state: Tensor, gt_valid: Tensor, gt_state: Tensor, agent_size: Tensor
    ) -> Tuple[Tensor, Tensor]:
        """
        Args:
            agent_valid: [n_batch, n_agent] bool
            agent_state: [n_batch, n_agent, 4]: x,y,theta,spd
            gt_valid: [n_batch, n_agent] bool
            gt_state: [n_batch, n_agent, 4] x,y,theta,spd
            agent_size: [n_batch, n_agent, 3], length, width, height
        Returns:
            reward: [n_batch, n_agent]
            reward_valid : [n_batch, n_agent]
        """
        reward = torch.zeros_like(agent_state[:, :, 0])
        reward_valid = agent_valid
        if self.w_collision > 0:
            n_batch, n_agent = agent_valid.shape
            # [n_batch, n_agent, 2]
            agent_xy = agent_state[..., :2]
            # [n_batch, n_agent]
            agent_yaw = agent_state[..., 2]
            # [n_batch, n_agent, 2]
            agent_heading = torch.stack([torch.cos(agent_yaw), torch.sin(agent_yaw)], axis=-1)

            # [n_batch, n_agent]
            agent_w = agent_size[:, :, :2].amin(-1)
            agent_l = agent_size[:, :, :2].amax(-1)
            agent_d = (agent_l - agent_w) / 4.0
            # [n_batch, n_agent, 2]
            agent_d = agent_d.unsqueeze(-1).expand(-1, -1, 2)

            # [n_batch, n_agent, 5, 2] centroid of 5 circles
            centroids = agent_xy.unsqueeze(2).expand(-1, -1, 5, -1) + torch.stack(
                [
                    -2 * agent_heading * agent_d,
                    -1 * agent_heading * agent_d,
                    0 * agent_heading * agent_d,
                    1 * agent_heading * agent_d,
                    2 * agent_heading * agent_d,
                ],
                dim=2,
            )

            # [n_batch, n_agent, 5, 2] -> [n_batch, n_agent, n_agent, 5, 2]
            centroids_0 = centroids.unsqueeze(2).expand(-1, -1, n_agent, -1, -1)
            centroids_1 = centroids_0.transpose(1, 2)
            # [n_batch, n_agent] -> [n_batch, n_agent, n_agent]
            agent_r = agent_w.unsqueeze(-1).expand(-1, -1, n_agent) / 2.0 + torch.finfo(agent_state.dtype).eps
            agent_r_sum = agent_r.transpose(1, 2) + agent_r

            distances = torch.zeros([n_batch, n_agent, n_agent, 5, 5], device=agent_valid.device)

            for i in range(5):
                for j in range(5):
                    # [n_batch, n_agent, n_agent, 2]
                    diff = centroids_0[:, :, :, i, :] - centroids_1[:, :, :, j, :]
                    # [n_batch, n_agent, n_agent]
                    _dist = torch.norm(diff, dim=-1) + torch.finfo(agent_state.dtype).eps
                    distances[:, :, :, i, j] = _dist

            # [n_batch, n_agent, n_agent]
            distances = torch.min(distances.flatten(start_dim=3, end_dim=4), dim=-1)[0]
            # relaxed collision: 1 for fully overlapped, 0 for not overlapped
            collision = torch.clamp(1 - distances / agent_r_sum, min=0)

            # [n_batch, n_agent, n_agent]
            ego_mask = torch.eye(n_agent, device=agent_valid.device, dtype=torch.bool)[None, :, :].expand(
                n_batch, -1, -1
            )
            ego_mask = ego_mask | (~agent_valid[:, :, None])
            ego_mask = ego_mask | (~agent_valid[:, None, :])
            collision.masked_fill_(ego_mask, 0.0)

            if self.reduce_collsion_with_max:
                # [n_batch, n_agent, n_agent] -> [n_batch, n_agent]: reduce dim 2
                collision = collision.amax(2)
            else:
                # [n_batch, n_agent, n_agent] -> [n_batch, n_agent, n_agent]
                collision = torch.clamp(collision, max=1)
                # [n_batch, n_agent]: reduce n_agent, as_valid: [n_batch, n_agent]
                collision = collision.sum(-1) / (agent_valid.sum(-1, keepdim=True))
            reward = reward - self.w_collision * collision.masked_fill(~agent_valid, 0.0)

        if self.use_il_loss and (gt_valid is not None):
            # [n_batch, n_agent, 1]
            il_invalid = ~(agent_valid & gt_valid).unsqueeze(-1)
            # [n_batch, n_agent, 4]
            gt_state = gt_state.masked_fill(il_invalid, 0)
            agent_state = agent_state.masked_fill(il_invalid, 0)

            error_pos = self.il_l_pos(gt_state[..., :2], agent_state[..., :2]).sum(-1)
            error_rot = self.il_l_rot.compute(gt_state[..., 2], agent_state[..., 2])
            error_spd = self.il_l_spd(gt_state[..., 3], agent_state[..., 3])
            il_loss = self.il_w_pos * error_pos + self.il_w_rot * error_rot + self.il_w_spd * error_spd
            reward = reward - il_loss
            reward_valid = agent_valid & gt_valid

        return reward.masked_fill(~reward_valid, 0.0), reward_valid

src/utils/rewards.py
zhejz/TrafficBots
# Licensed under the CC BY-NC 4.0 license (https://creativecommons.org/licenses/by-nc/4.0/)
from typing import List
from omegaconf import ListConfig
from pathlib import Path
import tarfile
import os
from torch import Tensor
from waymo_open_dataset.protos import motion_submission_pb2
from pytorch_lightning.loggers import WandbLogger
from .transform_utils import torch_pos2global, torch_rad2rot

# ! single GPU only


class SubWOMD:
    def __init__(
        self,
        k_futures: int,
        wb_artifact: str,
        interactive_challenge: bool,
        activate: bool,
        method_name: str,
        authors: ListConfig[str],
        affiliation: str,
        description: str,
        method_link: str,
    ) -> None:
        self.activate = activate
        self.method_name = method_name
        if activate:
            self.submissions = {}
            for _k in range(1, k_futures + 1):
                self.submissions[_k] = motion_submission_pb2.MotionChallengeSubmission()
                self.submissions[_k].account_name = "zhejun.zhang94@gmail.com"
                self.submissions[_k].unique_method_name = f"{method_name}_K{_k}"
                self.submissions[_k].authors.extend(list(authors))
                self.submissions[_k].affiliation = affiliation
                self.submissions[_k].description = f"{description}, wb_model: {wb_artifact}"
                self.submissions[_k].method_link = method_link
                if interactive_challenge:
                    self.submissions[_k].submission_type = 2
                else:
                    self.submissions[_k].submission_type = 1

    def add_to_submissions(
        self,
        waymo_trajs: Tensor,
        waymo_scores: Tensor,
        mask_pred: Tensor,
        object_id: Tensor,
        scenario_center: Tensor,
        scenario_yaw: Tensor,
        scenario_id: List[str],
    ) -> None:
        """
        Args:
            waymo_trajs: [n_batch, step_start+1...step_end, n_agent, K, 2]
            waymo_scores: [n_batch, n_agent, K]
            mask_pred: [n_batch, n_agent] bool
            object_id: [n_batch, n_agent] int
            scenario_center: [n_batch, 2]
            scenario_yaw: [n_batch]
            scenario_id: list of str
        """
        if not self.activate:
            return

        waymo_trajs = waymo_trajs[:, 4::5].permute(0, 2, 3, 1, 4)  # [n_batch, n_agent, K, n_step, 2]
        waymo_trajs = torch_pos2global(
            waymo_trajs.flatten(1, 3), scenario_center.unsqueeze(1), torch_rad2rot(scenario_yaw)
        ).view(waymo_trajs.shape)

        waymo_trajs = waymo_trajs.cpu().numpy()
        waymo_scores = waymo_scores.cpu().numpy()
        mask_pred = mask_pred.cpu().numpy()
        object_id = object_id.cpu().numpy()

        for i_batch in range(waymo_trajs.shape[0]):
            agent_pos = waymo_trajs[i_batch, mask_pred[i_batch]]  # [n_agent_pred, K, n_step, 2]
            agent_id = object_id[i_batch, mask_pred[i_batch]]  # [n_agent_pred]
            agent_score = waymo_scores[i_batch, mask_pred[i_batch]]  # [n_agent_pred, K]

            for n_K, submission in self.submissions.items():
                scenario_prediction = motion_submission_pb2.ChallengeScenarioPredictions()
                scenario_prediction.scenario_id = scenario_id[i_batch]

                if submission.submission_type == 1:
                    # single prediction
                    for i_track in range(agent_pos.shape[0]):
                        prediction = motion_submission_pb2.SingleObjectPrediction()
                        prediction.object_id = agent_id[i_track]
                        for _k in range(n_K):
                            scored_trajectory = motion_submission_pb2.ScoredTrajectory()
                            scored_trajectory.confidence = agent_score[i_track, _k]
                            scored_trajectory.trajectory.center_x.extend(agent_pos[i_track, _k, :, 0])
                            scored_trajectory.trajectory.center_y.extend(agent_pos[i_track, _k, :, 1])
                            prediction.trajectories.append(scored_trajectory)
                        scenario_prediction.single_predictions.predictions.append(prediction)
                else:
                    # joint prediction
                    for _k in range(n_K):
                        scored_joint_trajectory = motion_submission_pb2.ScoredJointTrajectory()
                        scored_joint_trajectory.confidence = agent_score[:, _k].sum(0)
                        for i_track in range(agent_pos.shape[0]):
                            object_trajectory = motion_submission_pb2.ObjectTrajectory()
                            object_trajectory.object_id = agent_id[i_track]
                            object_trajectory.trajectory.center_x.extend(agent_pos[i_track, _k, :, 0])
                            object_trajectory.trajectory.center_y.extend(agent_pos[i_track, _k, :, 1])
                            scored_joint_trajectory.trajectories.append(object_trajectory)
                        scenario_prediction.joint_prediction.joint_trajectories.append(scored_joint_trajectory)

                submission.scenario_predictions.append(scenario_prediction)

    def save_sub_files(self, logger: WandbLogger) -> List[str]:
        if not self.activate:
            return []

        print(f"saving womd submission files to {os.getcwd()}")
        file_paths = []
        for k, submission in self.submissions.items():
            submission_dir = Path(f"womd_{self.method_name}_K{k}")
            submission_dir.mkdir(exist_ok=True)
            f = open(submission_dir / f"womd_{self.method_name}_K{k}.bin", "wb")
            f.write(submission.SerializeToString())
            f.close()
            tar_file_name = submission_dir.as_posix() + ".tar.gz"
            with tarfile.open(tar_file_name, "w:gz") as tar:
                tar.add(submission_dir, arcname=submission_dir.name)
            if isinstance(logger, WandbLogger):
                logger.experiment.save(tar_file_name)
            else:
                file_paths.append(tar_file_name)
        return file_paths

src/utils/submission.py
zhejz/TrafficBots
# Licensed under the CC BY-NC 4.0 license (https://creativecommons.org/licenses/by-nc/4.0/)
from torch import Tensor
import torch


class TeacherForcing:
    def __init__(
        self,
        step_spawn_agent: int = 10,
        step_warm_start: int = 10,
        step_horizon: int = 0,
        step_horizon_decrease_per_epoch: int = 0,
        prob_forcing_agent: float = 0,
        prob_forcing_agent_decrease_per_epoch: float = 0,
    ):
        """
        Args:
            step_spawn_agent: spawn agents up to this time step
            step_warm_start: teacher forcing all agents up to this time step.
            step_horizon: from step 0 to step_horizon all agents will be teacher-forced
            step_horizon_decrease_per_epoch: decrease `step_horizon` every epoch. e.g. 10
            prob_forcing_agent: some agents will always be teacher-forced. e.g. 0.5
            prob_forcing_agent_decrease_per_epoch: decrease `prob_forcing_agent` every epoch. e.g. 0.1
        """
        self.step_spawn_agent = step_spawn_agent
        self.step_warm_start = step_warm_start
        self.step_horizon = step_horizon
        self.step_horizon_decrease_per_epoch = step_horizon_decrease_per_epoch
        self.prob_forcing_agent = prob_forcing_agent
        self.prob_forcing_agent_decrease_per_epoch = prob_forcing_agent_decrease_per_epoch

    @torch.no_grad()
    def get(self, as_valid: Tensor, current_epoch: int = 0, gt_sdc: bool = False) -> Tensor:
        """
        Args:
            as_valid: [n_batch, step_gt+1, n_agent] bool
            current_epoch: current training epoch

        Returns:
            mask_teacher_forcing: [n_batch, step_gt+1, n_agent]
        """
        mask_teacher_forcing = torch.zeros_like(as_valid)

        # always spawn at step 0
        mask_teacher_forcing[:, 0] |= as_valid[:, 0]
        if self.step_spawn_agent > 0:
            # spawn when valid change from False to True, because traj is interpolated.
            mask_spawn_agent = (~as_valid[:, :-1]) & as_valid[:, 1:]
            mask_spawn_agent[:, self.step_spawn_agent :] = False
            mask_teacher_forcing[:, 1:] |= mask_spawn_agent

        # warm start
        if self.step_warm_start >= 0:
            mask_teacher_forcing[:, : self.step_warm_start + 1] |= as_valid[:, : self.step_warm_start + 1]

        # horizon schedule
        step_horizon = self.step_horizon - self.step_horizon_decrease_per_epoch * current_epoch
        if step_horizon > 0:
            mask_teacher_forcing[:, :step_horizon] |= as_valid[:, :step_horizon]

        # agent schedule
        prob_forcing_agent = self.prob_forcing_agent - self.prob_forcing_agent_decrease_per_epoch * current_epoch
        if prob_forcing_agent > 0:
            # [n_batch, n_agent]
            mask_forcing_agent = torch.bernoulli(torch.ones_like(as_valid[:, 0]) * prob_forcing_agent).bool()
            mask_teacher_forcing |= mask_forcing_agent.unsqueeze(1) & as_valid

        # what-if motion prediction
        if gt_sdc:
            assert as_valid[:, :, 0].all()
            # mask_teacher_forcing[:, :, 0] = True
            mask_teacher_forcing[:, :, 0] |= as_valid[:, :, 0]

        return mask_teacher_forcing

src/utils/teacher_forcing.py
zhejz/TrafficBots
# Licensed under the CC BY-NC 4.0 license (https://creativecommons.org/licenses/by-nc/4.0/)
from typing import Dict, Tuple, Optional
from torch import Tensor
import numpy as np
import torch
from utils.transform_utils import cast_rad


class TrafficRuleChecker:
    def __init__(
        self,
        map_boundary: Tensor,
        map_valid: Tensor,
        map_type: Tensor,
        map_pos: Tensor,
        map_dir: Tensor,
        tl_stop_valid: Tensor,
        tl_stop_pos: Tensor,
        tl_stop_state: Tensor,
        agent_type: Tensor,
        agent_size: Tensor,
        agent_goal: Optional[Tensor],
        agent_dest: Optional[Tensor],
        enable_check_collided: bool,
        enable_check_run_road_edge: bool,
        enable_check_run_red_light: bool,
        enable_check_passive: bool,
        collision_size_scale: float = 1.1,
    ) -> None:
        self.agent_size = agent_size[..., :2] * collision_size_scale
        self.map_boundary = map_boundary
        self.tl_stop_valid = tl_stop_valid
        self.tl_stop_pos = tl_stop_pos
        self.tl_stop_state = tl_stop_state
        self.agent_goal = agent_goal
        self.agent_dest = agent_dest

        self.enable_check_collided = enable_check_collided
        self.enable_check_run_road_edge = enable_check_run_road_edge
        self.enable_check_run_red_light = enable_check_run_red_light
        self.enable_check_passive = enable_check_passive

        mask = agent_type[:, :, 0]
        n_batch, n_agent = mask.shape
        self.outside_map = torch.zeros_like(mask)
        self.collided = torch.zeros_like(mask)
        self.run_red_light = torch.zeros_like(mask)
        self.goal_reached = torch.zeros_like(mask)
        self.dest_reached = torch.zeros_like(mask)
        self.run_road_edge = torch.zeros_like(mask)
        self.passive = torch.zeros_like(mask)
        self.passive_counter = torch.zeros_like(mask, dtype=torch.float32)

        # for self._check_collided
        # [n_batch, n_agent, n_agent]: no self collision
        self.agent_ego_mask = torch.eye(n_agent, dtype=torch.bool, device=mask.device)[None, :, :].expand(
            n_batch, -1, -1
        )
        ped_cyc_mask = agent_type[:, :, 1]
        collision_ped_cyc_mask = ped_cyc_mask.unsqueeze(1) & ped_cyc_mask.unsqueeze(2)
        self.collision_invalid_mask = self.agent_ego_mask | collision_ped_cyc_mask

        # for self._check_run_road_edge
        self.road_edge, self.road_edge_valid = self._get_road_edge(map_valid, map_type, map_pos, map_dir)

        # for self._check_run_red_light
        # [n_batch, n_agent, 1]
        self.run_red_light_agent_length = agent_size[:, :, [0]] * 0.5 * 0.6
        self.run_red_light_agent_width = agent_size[:, :, [1]] * 0.5 * 1.8
        self.veh_mask = agent_type[:, :, 0]  # [n_batch, n_agent]

        # for self._check_passive
        self.lane_center, self.lane_center_valid = self._get_lane_center(map_valid, map_type, map_pos)

        # for self._check_goal_reached
        # [n_batch, n_agent]
        if self.agent_dest is not None:
            self.goal_thresh_pos = agent_size[:, :, 0] * 8
            self.goal_thresh_rot = np.deg2rad(15)

        # for self._check_dest_reached
        if self.agent_dest is not None:
            # [n_batch, 1]
            batch_idx = torch.arange(map_valid.shape[0]).unsqueeze(1)
            # [n_batch, n_agent, 20]
            self.dest_valid = map_valid[batch_idx, agent_dest]
            # [n_batch, n_agent, n_type]: one_hot bool, LANE<=3, TYPE_ROAD_EDGE_BOUNDARY = 4
            self.dest_type = map_type[batch_idx, agent_dest]
            # [n_batch, n_agent, 20, 2]
            self.dest_pos = map_pos[batch_idx, agent_dest]
            # [n_batch, n_agent, 20, 2]
            self.dest_dir = map_dir[batch_idx, agent_dest]
            self.dest_dir = self.dest_dir / torch.norm(self.dest_dir, dim=-1, keepdim=True)
            # dest_thresh_rot only valid for lane dests
            self.dest_thresh_rot = np.deg2rad(30)
            # [n_batch, n_agent]: thresh_lane=50, thresh_edge=10
            self.dest_thresh_pos = torch.ones_like(agent_size[:, :, 0]) * 50
            self.dest_thresh_pos = self.dest_thresh_pos * (1 - self.dest_type[:, :, 4] * 0.8)

    @staticmethod
    def _check_outside_map(as_valid: Tensor, as_state: Tensor, map_boundary: Tensor) -> Tensor:
        """
        Args:
            as_valid: [n_batch, n_agent] bool
            as_state: [n_batch, n_agent, 4] x,y,theta,v
            map_boundary: [n_batch, 4], xmin,xmax,ymin,ymax
        Returns:
            outside_map_this_step: [n_batch, n_agent], bool
        """
        # [n_batch, n_agent]
        x = as_state[:, :, 0]
        y = as_state[:, :, 1]
        # [n_batch, 1]
        xmin = map_boundary[:, [0]]
        xmax = map_boundary[:, [1]]
        ymin = map_boundary[:, [2]]
        ymax = map_boundary[:, [3]]
        outside_map_this_step = ((x > xmax) | (x < xmin) | (y > ymax) | (y < ymin)) & as_valid
        return outside_map_this_step

    @staticmethod
    def _check_collided(as_valid: Tensor, agent_bbox: Tensor, collision_invalid_mask: Tensor) -> Tensor:
        """
        Args:
            as_valid: [n_batch, n_agent] bool
            agent_bbox: [n_batch, n_agent, 4, 2], 4 corners, (x,y)
            collision_invalid_mask: [n_batch, n_agent, n_agent]: no self collision
        Returns:
            collided_this_step: [n_batch, n_agent], bool
        """
        bbox_next = agent_bbox.roll(-1, dims=2)

        # [n_batch, n_agent, 4, 3]
        bbox_line = torch.cat(  # ax+by+c=0
            [
                bbox_next[..., [1]] - agent_bbox[..., [1]],  # a
                agent_bbox[..., [0]] - bbox_next[..., [0]],  # b
                bbox_next[..., [0]] * agent_bbox[..., [1]] - bbox_next[..., [1]] * agent_bbox[..., [0]],
            ],  # c
            axis=-1,
        )
        bbox_point = torch.cat([agent_bbox, torch.ones_like(agent_bbox[..., [0]])], axis=-1)

        # [n_batch, n_agent, n_agent, 4, 4, 3]
        n_agent = agent_bbox.shape[1]
        bbox_line = bbox_line[:, :, None, :, None, :].expand(-1, -1, n_agent, -1, 4, -1)
        bbox_point = bbox_point[:, None, :, None, :, :].expand(-1, n_agent, -1, 4, -1, -1)

        # [n_batch, n_agent, n_agent, 4, 4]
        is_outside = torch.sum(bbox_line * bbox_point, axis=-1) > 0

        # [n_batch, n_agent, n_agent]
        no_collision = torch.any(torch.all(is_outside, axis=-1), axis=-1)
        no_collision = no_collision | no_collision.transpose(1, 2)

        # [n_batch, n_agent, n_agent]: no collision for invalid agent
        invalid_mask = ~(as_valid[:, :, None] & as_valid[:, None, :])
        no_collision = no_collision | collision_invalid_mask | invalid_mask
        collided_this_step = ~(no_collision.all(-1))
        return collided_this_step

    @staticmethod
    def _check_run_road_edge(
        as_valid: Tensor, agent_bbox: Tensor, veh_mask: Tensor, road_edge: Tensor, road_edge_valid: Tensor
    ) -> Tensor:
        """Check vehicles only.
            as_valid: [n_batch, n_agent] bool
            agent_bbox: [n_batch, n_agent, 4, 2], 4 corners, (x,y)
            veh_mask: [n_batch, n_agent] bool
            road_edge: [n_batch, n_pl*20, 2, 2], (start/end), (x,y)
            road_edge_valid: [n_batch, n_pl*20], bool
        Returns:
            run_road_edge_this_step: [n_batch, n_agent], bool
        """
        bbox_next = agent_bbox.roll(-1, dims=2)  # [n_batch, n_agent, 4, 2]
        # [n_batch, n_agent, 1, 4, 2, 2]
        bbox_line = torch.stack([agent_bbox, bbox_next], dim=-2).unsqueeze(2)

        # [n_batch, n_pl*20, 2, 2] -> [n_batch, 1, n_pl*20, 1, 2, 2]
        road_edge_line = road_edge[:, None, :, None, :, :]

        # [n_batch, n_agent, n_pl*20, 4, 2]
        A = bbox_line[:, :, :, :, 0]
        B = bbox_line[:, :, :, :, 1]
        C = road_edge_line[:, :, :, :, 0]
        D = road_edge_line[:, :, :, :, 1]

        # [n_batch, n_agent, n_pl*20, 4]
        run_road_edge_this_step = (ccw(A, C, D) != ccw(B, C, D)) & (ccw(A, B, C) != ccw(A, B, D))

        # [n_batch, n_agent, n_pl*20]
        run_road_edge_this_step = run_road_edge_this_step.any(-1) & road_edge_valid.unsqueeze(1)

        # [n_batch, n_agent]
        run_road_edge_this_step = run_road_edge_this_step.any(-1) & as_valid & veh_mask
        return run_road_edge_this_step

    @staticmethod
    def _check_run_red_light(
        as_valid: Tensor,
        as_state: Tensor,
        tl_valid: Tensor,
        tl_pos: Tensor,
        tl_state: Tensor,
        run_red_light_agent_length: Tensor,
        run_red_light_agent_width: Tensor,
        veh_mask: Tensor,
    ) -> Tensor:
        """Check up to step_gt
        Args:
            as_valid: [n_batch, n_agent] bool
            as_state: [n_batch, n_agent, 4] x,y,theta,v
            tl_valid: [n_batch, n_tl]
            tl_pos: [n_batch, n_tl, 2] xy of stop point
            tl_state: [n_batch, n_tl, n_tl_state=5] bool one_hot
            run_red_light_agent_length: [n_batch, n_agent, 1]
            run_red_light_agent_width: [n_batch, n_agent, 1]
            veh_mask: [n_batch, n_agent]
            # LANE_STATE_UNKNOWN = 0;
            # LANE_STATE_STOP = 1;
            # LANE_STATE_CAUTION = 2;
            # LANE_STATE_GO = 3;
            # LANE_STATE_FLASHING = 4;
        Returns:
            run_red_light_this_step: [n_batch, n_agent], bool
        """
        # [n_batch, n_agent]
        heading_cos = torch.cos(as_state[..., 2])
        heading_sin = torch.sin(as_state[..., 2])
        # [n_batch, n_agent, 1, 2]
        agent_heading_f = torch.stack([heading_cos, heading_sin], axis=-1).unsqueeze(2)
        agent_heading_r = torch.stack([heading_sin, -heading_cos], axis=-1).unsqueeze(2)

        # [n_batch, n_agent, 1, 2]
        agent_xy_0 = as_state[..., :2].unsqueeze(2)
        agent_xy_1 = agent_xy_0 + 0.1 * as_state[..., [3]].unsqueeze(2) * agent_heading_f

        # [n_batch, n_tl, 2] -> [n_batch, 1, n_tl, 2]
        tl_pos = tl_pos.unsqueeze(1)
        # [n_batch, n_agent, n_tl]
        inside_0 = torch.logical_and(
            torch.abs(torch.sum((tl_pos - agent_xy_0) * agent_heading_f, dim=-1)) < run_red_light_agent_length,
            torch.abs(torch.sum((tl_pos - agent_xy_0) * agent_heading_r, dim=-1)) < run_red_light_agent_width,
        )
        inside_1 = torch.logical_and(
            torch.abs(torch.sum((tl_pos - agent_xy_1) * agent_heading_f, dim=-1)) < run_red_light_agent_length,
            torch.abs(torch.sum((tl_pos - agent_xy_1) * agent_heading_r, dim=-1)) < run_red_light_agent_width,
        )

        # [n_batch, n_agent, 1]
        mask_valid_agent = (as_valid & veh_mask).unsqueeze(2)
        # [n_batch, 1, n_tl]
        mask_valid_tl = (tl_valid & tl_state[:, :, 1]).unsqueeze(1)
        # [n_batch, n_agent, n_tl]
        run_red_light_this_step = inside_0 & (~inside_1) & mask_valid_agent & mask_valid_tl
        # [n_batch, n_agent]
        run_red_light_this_step = run_red_light_this_step.any(-1)
        return run_red_light_this_step

    @staticmethod
    def _check_passive(
        as_valid: Tensor,
        as_state: Tensor,
        passive_counter: Tensor,
        tl_valid: Tensor,
        tl_pos: Tensor,
        tl_state: Tensor,
        lane_center: Tensor,
        lane_center_valid: Tensor,
        veh_mask: Tensor,
        agent_ego_mask: Tensor,
    ) -> Tuple[Tensor, Tensor]:
        """Check vehicles only.
            as_valid: [n_batch, n_agent] bool
            as_state: [n_batch, n_agent, 4] x,y,theta,v
            passive_counter: [n_batch, n_agent], float32
            tl_valid: [n_batch, n_tl]
            tl_pos: [n_batch, n_tl, 2] xy of stop point
            tl_state: [n_batch, n_tl, n_tl_state=5] bool one_hot
            lane_center: [n_batch, n_pl*20, 2]
            lane_center_valid: [n_batch, n_pl*20]
            veh_mask: [n_batch, n_agent] bool
            agent_ego_mask: [n_batch, n_agent, n_agent] bool, eye
        Returns:
            passive_this_step: [n_batch, n_agent], bool
            passive_counter: [n_batch, n_agent], float32
        """
        # [n_batch, n_agent, 1, 2] - [n_batch, 1, n_pl*20, 2] = [n_batch, n_agent, n_pl*20, 2]
        close_to_lane = torch.norm(as_state[:, :, :2].unsqueeze(2) - lane_center.unsqueeze(1), dim=-1)
        # [n_batch, n_agent, n_pl*20] bool
        close_to_lane = close_to_lane < 2  # m
        # [n_batch, n_agent]: [n_batch, n_agent, n_pl*20] & [n_batch, 1, n_pl*20]
        close_to_lane = (close_to_lane & lane_center_valid.unsqueeze(1)).any(-1)

        # [n_batch, n_agent]
        low_speed = as_state[:, :, 3] < 5  # m/s

        # [n_batch, n_agent, 1, 2]
        agent_heading_f = torch.stack([torch.cos(as_state[..., 2]), torch.sin(as_state[..., 2])], axis=-1).unsqueeze(2)

        # LANE_STATE_UNKNOWN = 0;
        # LANE_STATE_STOP = 1;
        # LANE_STATE_CAUTION = 2;
        # LANE_STATE_GO = 3;
        # LANE_STATE_FLASHING = 4;
        # check red (flashing,yellow) light ahead
        # [n_batch, 1, n_tl]
        mask_valid_tl = (tl_valid & tl_state[:, :, [0, 1, 2, 4]].any(-1)).unsqueeze(1)
        # [n_batch, n_agent, n_tl, 2]
        tl_vec = tl_pos.unsqueeze(1) - as_state[:, :, :2].unsqueeze(2)
        tl_vec_norm = torch.norm(tl_vec, dim=-1)
        # [n_batch, n_agent, n_tl]
        tl_is_close = tl_vec_norm < 10  # meter
        tl_is_ahead = ((agent_heading_f * tl_vec).sum(-1) / tl_vec_norm) > 0.95  # np.cos(np.deg2rad(18)) = 0.95
        # [n_batch, n_agent]
        red_tl_ahead = (tl_is_close & tl_is_ahead & mask_valid_tl).any(-1)

        # check other agents
        # [n_batch, n_agent, n_agent, 2]
        agent_vec = as_state[:, :, :2].unsqueeze(1) - as_state[:, :, :2].unsqueeze(2)
        agent_vec_norm = torch.norm(agent_vec, dim=-1)
        # [n_batch, n_agent, n_agent]
        agent_is_close = agent_vec_norm < 10  # meter
        agent_is_ahead = ((agent_heading_f * agent_vec).sum(-1) / agent_vec_norm) > 0.95
        # [n_batch, n_agent]
        agent_ahead = (
            agent_is_close & agent_is_ahead & as_valid.unsqueeze(1) & as_valid.unsqueeze(2) & (~agent_ego_mask)
        ).any(-1)

        passive_this_step = as_valid & veh_mask & close_to_lane & low_speed & (~red_tl_ahead) & (~agent_ahead)

        # accumulate, set to zero if not passive
        passive_counter = (passive_counter + passive_this_step) * passive_this_step
        passive_this_step = passive_counter > 20
        return passive_this_step, passive_counter

    @staticmethod
    def _check_goal_reached(
        as_valid: Tensor,
        as_state: Tensor,
        goal: Tensor,
        goal_reached: Tensor,
        goal_thresh_pos: Tensor,
        goal_thresh_rot: float,
    ) -> Tensor:
        """
        Args:
            as_valid: [n_batch, n_agent] bool
            as_state: [n_batch, n_agent, 4] x,y,theta,v
            goal: [n_batch, n_agent, 4], x,y,theta,v
            goal_reached: [n_batch, n_agent], bool
            goal_thresh_pos: [n_batch, n_agent] agent_length * 8
            goal_thresh_rot: float, 15 rad
        Returns:
            goal_reached_this_step: [n_batch, n_agent], bool
        """
        # [n_batch, n_agent]
        pos_reached = torch.norm(as_state[..., :2] - goal[..., :2], dim=-1) < goal_thresh_pos
        rot_reached = torch.abs(cast_rad(as_state[..., 2] - goal[..., 2])) < goal_thresh_rot
        goal_reached_this_step = pos_reached & rot_reached & as_valid & (~goal_reached)
        return goal_reached_this_step

    @staticmethod
    def _check_dest_reached(
        as_valid: Tensor,
        as_state: Tensor,
        dest_valid: Tensor,
        dest_type: Tensor,
        dest_pos: Tensor,
        dest_dir: Tensor,
        dest_reached: Tensor,
        dest_thresh_pos: Tensor,
        dest_thresh_rot: float,
    ) -> Tensor:
        """
        Args:
            as_valid: [n_batch, n_agent] bool
            as_state: [n_batch, n_agent, 4] x,y,theta,v
            dest_valid: [n_batch, n_agent, 20]
            dest_type: [n_batch, n_agent, n_pl_type] one_hot bool, LANE<=3, ROAD_EDGE = 4
            dest_pos: [n_batch, n_agent, 20, 2]
            dest_dir: [n_batch, n_agent, 20, 2], unit_vec
            dest_reached: [n_batch, n_agent], bool
            dest_thresh_pos: [n_batch, n_agent]
            dest_thresh_rot: float, in rad
        Returns:
            dest_reached_this_step: [n_batch, n_agent], bool
        """
        # [n_batch, n_agent, 20]
        dist_to_dest = torch.norm(as_state[..., :2].unsqueeze(2) - dest_pos, dim=-1).masked_fill(~dest_valid, 1e4)
        # [n_batch, n_agent]
        pos_reached = (dist_to_dest < dest_thresh_pos.unsqueeze(-1)).any(-1)

        heading_cos = torch.cos(as_state[..., 2])  # [n_batch, n_agent]
        heading_sin = torch.sin(as_state[..., 2])  # [n_batch, n_agent]
        agent_heading_f = torch.stack([heading_cos, heading_sin], axis=-1)  # [n_batch, n_agent, 2]

        # [n_batch, n_agent, 1, 2] * [n_batch, n_agent, 20, 2]
        # [n_batch, n_agent, 20]
        rot_diff_to_dest = (agent_heading_f.unsqueeze(2) * dest_dir).sum(-1).masked_fill(~dest_valid, 0)
        # [n_batch, n_agent]
        rot_reached = (rot_diff_to_dest > np.cos(dest_thresh_rot)).any(-1)

        # [n_batch, n_agent]: one_hot bool, LANE<=3, TYPE_ROAD_EDGE_BOUNDARY = 4
        mask_lane = dest_type[:, :, :4].any(-1)
        mask_edge = dest_type[:, :, 4]
        dest_reached_this_step = (
            (~dest_reached) & as_valid & ((mask_lane & pos_reached & rot_reached) | (mask_edge & pos_reached))
        )
        return dest_reached_this_step

    @torch.no_grad()
    def check(self, step: int, as_valid: Tensor, as_state: Tensor) -> Dict[str, Tensor]:
        """
        Args:
            as_valid: [n_batch, n_agent] bool
            as_state: [n_batch, n_agent, 4] x,y,theta,v
        Returns: for this current step
            violations: Dict {str -> Tensor [n_batch, n_agent]}
        """
        agent_bbox = self._get_agent_bbox(as_state, self.agent_size)

        outside_map_this_step = self._check_outside_map(as_valid, as_state, self.map_boundary)
        self.outside_map = self.outside_map | outside_map_this_step

        if self.enable_check_collided:
            collided_this_step = self._check_collided(as_valid, agent_bbox, self.collision_invalid_mask)
            self.collided = self.collided | collided_this_step
        else:
            collided_this_step = self.collided

        if self.enable_check_run_road_edge:
            run_road_edge_this_step = self._check_run_road_edge(
                as_valid, agent_bbox, self.veh_mask, self.road_edge, self.road_edge_valid
            )
            self.run_road_edge = self.run_road_edge | run_road_edge_this_step
        else:
            run_road_edge_this_step = self.run_road_edge

        # step can be larger than ground truth tl_step
        if self.enable_check_run_red_light:
            tl_step = min(step, self.tl_stop_valid.shape[1] - 1)
            run_red_light_this_step = self._check_run_red_light(
                as_valid,
                as_state,
                self.tl_stop_valid[:, tl_step],
                self.tl_stop_pos[:, tl_step],
                self.tl_stop_state[:, tl_step],
                self.run_red_light_agent_length,
                self.run_red_light_agent_width,
                self.veh_mask,
            )
            self.run_red_light = self.run_red_light | run_red_light_this_step
        else:
            run_red_light_this_step = self.run_red_light

        if self.enable_check_passive:
            passive_this_step, self.passive_counter = self._check_passive(
                as_valid,
                as_state,
                self.passive_counter,
                self.tl_stop_valid[:, tl_step],
                self.tl_stop_pos[:, tl_step],
                self.tl_stop_state[:, tl_step],
                self.lane_center,
                self.lane_center_valid,
                self.veh_mask,
                self.agent_ego_mask,
            )
            self.passive = self.passive | passive_this_step
        else:
            passive_this_step = self.passive

        if self.agent_goal is None:
            goal_reached_this_step = torch.zeros_like(self.goal_reached)
        else:
            goal_reached_this_step = self._check_goal_reached(
                as_valid, as_state, self.agent_goal, self.goal_reached, self.goal_thresh_pos, self.goal_thresh_rot
            )
        self.goal_reached = self.goal_reached | goal_reached_this_step

        if self.agent_dest is None:
            dest_reached_this_step = torch.zeros_like(self.dest_reached)
        else:
            dest_reached_this_step = self._check_dest_reached(
                as_valid,
                as_state,
                self.dest_valid,
                self.dest_type,
                self.dest_pos,
                self.dest_dir,
                self.dest_reached,
                self.dest_thresh_pos,
                self.dest_thresh_rot,
            )
        self.dest_reached = self.dest_reached | dest_reached_this_step

        # [n_batch, n_agent], bool
        violations = {
            # "agent_to_kill": outside_map_this_step,
            "outside_map": self.outside_map,
            "outside_map_this_step": outside_map_this_step,
            "collided": self.collided,  # no collision ped2ped
            "collided_this_step": collided_this_step,
            "run_road_edge": self.run_road_edge,  # only for vehicles
            "run_road_edge_this_step": run_road_edge_this_step,
            "run_red_light": self.run_red_light,  # only for vehicles
            "run_red_light_this_step": run_red_light_this_step,
            "passive": self.passive,  # only for vehicles
            "passive_this_step": passive_this_step,
            "goal_reached": self.goal_reached,
            "goal_reached_this_step": goal_reached_this_step,
            "dest_reached": self.dest_reached,
            "dest_reached_this_step": dest_reached_this_step,
        }
        return violations

    @staticmethod
    def _get_agent_bbox(agent_states: Tensor, agent_size: Tensor) -> Tensor:
        """
        Args:
            agent_states: [n_batch, n_agent, 4] x,y,theta,v
            agent_size: [n_batch, n_agent, 2], length, width
        Returns:
            agent_bbox: [n_batch, n_agent, 4, 2]
        """
        heading_cos = torch.cos(agent_states[..., 2])  # [n_batch, n_agent]
        heading_sin = torch.sin(agent_states[..., 2])  # [n_batch, n_agent]
        agent_heading_f = torch.stack([heading_cos, heading_sin], axis=-1)  # [n_batch, n_agent, 2]
        agent_heading_r = torch.stack([heading_sin, -heading_cos], axis=-1)  # [n_batch, n_agent, 2]
        offset_forward = 0.5 * agent_size[..., [0]].expand(-1, -1, 2) * agent_heading_f  # [n_batch, n_agent, 2]
        offset_right = 0.5 * agent_size[..., [1]].expand(-1, -1, 2) * agent_heading_r  # [n_batch, n_agent, 2]
        vertex_offset = torch.stack(
            [
                -offset_forward + offset_right,
                offset_forward + offset_right,
                offset_forward - offset_right,
                -offset_forward - offset_right,
            ],
            dim=2,
        )
        agent_bbox = agent_states[:, :, None, :2].expand(-1, -1, 4, -1) + vertex_offset
        return agent_bbox

    # FREEWAY = 0
    # SURFACE_STREET = 1
    # STOP_SIGN = 2
    # BIKE_LANE = 3
    # TYPE_ROAD_EDGE_BOUNDARY = 4
    # TYPE_ROAD_EDGE_MEDIAN = 5
    # SOLID_SINGLE = 6
    # SOLID_DOUBLE = 7
    # PASSING_DOUBLE_YELLOW = 8
    # SPEED_BUMP = 9
    # CROSSWALK = 10

    @staticmethod
    def _get_road_edge(map_valid: Tensor, map_type: Tensor, map_pos: Tensor, map_dir: Tensor) -> Tuple[Tensor, Tensor]:
        """
        Args:
            map_valid: [n_batch, n_pl, 20]
            map_type: [n_batch, n_pl, 11] one_hot bool
            map_pos: [n_batch, n_pl, 20, 2]
            map_dir: [n_batch, n_pl, 20, 2]
        Returns:
            road_edge: [n_batch, n_pl*20, 2, 2], (start/end), (x,y)
            road_edge_valid: [n_batch, n_pl*20], bool
        """
        # [n_batch, n_pl, 20]
        road_edge_valid = map_valid & map_type[:, :, [4, 5, 7]].any(dim=-1, keepdim=True)
        road_edge_valid = road_edge_valid.flatten(1, 2)  # [n_batch, n_pl*20]
        # [n_batch, n_pl, 20, 2] -> [n_batch, n_pl, 20, 2, 2] -> [n_batch, n_pl*20, 2, 2]
        road_edge = torch.stack([map_pos, map_pos + map_dir], dim=-2).flatten(1, 2)
        return road_edge, road_edge_valid

    @staticmethod
    def _get_lane_center(map_valid: Tensor, map_type: Tensor, map_pos: Tensor) -> Tuple[Tensor, Tensor]:
        """
        Args:
            map_valid: [n_batch, n_pl, 20]
            map_type: [n_batch, n_pl]
            map_pos: [n_batch, n_pl, 20, 2]
        Returns:
            lane_center: [n_batch, n_pl*20, 2]
            lane_center_valid: [n_batch, n_pl*20]
        """
        lane_center_valid = map_type[:, :, :3].any(dim=-1, keepdim=True)  # [n_batch, n_pl, 1]
        lane_center_valid = map_valid & lane_center_valid  # [n_batch, n_pl, 20]
        lane_center_valid = lane_center_valid.flatten(1, 2)  # [n_batch, n_pl*20]
        lane_center = map_pos.flatten(1, 2)
        return lane_center, lane_center_valid


def ccw(A: Tensor, B: Tensor, C: Tensor) -> Tensor:
    return (C[..., 1] - A[..., 1]) * (B[..., 0] - A[..., 0]) > (B[..., 1] - A[..., 1]) * (C[..., 0] - A[..., 0])

src/utils/traffic_rule_checker.py
zhejz/TrafficBots
# Licensed under the CC BY-NC 4.0 license (https://creativecommons.org/licenses/by-nc/4.0/)
import numpy as np
import torch
from torch import Tensor
import transforms3d
from typing import Union


def cast_rad(angle: Union[float, np.ndarray, Tensor]) -> Union[float, np.ndarray, Tensor]:
    """Cast angle such that they are always in the [-pi, pi) range."""
    return (angle + np.pi) % (2 * np.pi) - np.pi


def _rotation33_as_yaw(rotation: np.ndarray) -> float:
    """Compute the yaw component of given 3x3 rotation matrix.

    Args:
        rotation (np.ndarray): 3x3 rotation matrix (np.float64 dtype recommended)

    Returns:
        float: yaw rotation in radians
    """
    return transforms3d.euler.mat2euler(rotation)[2]


def _yaw_as_rotation33(yaw: float) -> np.ndarray:
    """Create a 3x3 rotation matrix from given yaw.
    The rotation is counter-clockwise and it is equivalent to:
    [cos(yaw), -sin(yaw), 0.0],
    [sin(yaw), cos(yaw), 0.0],
    [0.0, 0.0, 1.0],

    Args:
        yaw (float): yaw rotation in radians

    Returns:
        np.ndarray: 3x3 rotation matrix
    """
    return transforms3d.euler.euler2mat(0, 0, yaw)


def get_so2_from_se2(transform_se3: np.ndarray) -> np.ndarray:
    """Gets rotation component in SO(2) from transformation in SE(2).

    Args:
        transform_se3: se2 transformation.

    Returns:
        rotation component in so2
    """
    rotation = np.eye(3, dtype=np.float64)
    rotation[:2, :2] = transform_se3[:2, :2]
    return rotation


def get_yaw_from_se2(transform_se3: np.ndarray) -> float:
    """Gets yaw from transformation in SE(2).

    Args:
        transform_se3: se2 transformation.

    Returns:
        yaw
    """
    return _rotation33_as_yaw(get_so2_from_se2(transform_se3))


def transform_points(points: np.ndarray, transf_matrix: np.ndarray) -> np.ndarray:
    """Transform points using transformation matrix.
    Note this function assumes points.shape[1] == matrix.shape[1] - 1, which means that the last
    row in the matrix does not influence the final result.
    For 2D points only the first 2x3 part of the matrix will be used.

    Args:
        points (np.ndarray): Input points (Nx2) or (Nx3).
        transf_matrix (np.ndarray): np.float64, 3x3 or 4x4 transformation matrix for 2D and 3D input respectively

    Returns:
        np.ndarray: array of shape (N,2) for 2D input points, or (N,3) points for 3D input points
    """
    assert len(points.shape) == len(transf_matrix.shape) == 2, (
        f"dimensions mismatch, both points ({points.shape}) and "
        f"transf_matrix ({transf_matrix.shape}) needs to be 2D numpy ndarrays."
    )
    assert (
        transf_matrix.shape[0] == transf_matrix.shape[1]
    ), f"transf_matrix ({transf_matrix.shape}) should be a square matrix."

    if points.shape[1] not in [2, 3]:
        raise AssertionError(f"Points input should be (N, 2) or (N, 3) shape, received {points.shape}")

    assert points.shape[1] == transf_matrix.shape[1] - 1, "points dim should be one less than matrix dim"

    points_transformed = (points @ transf_matrix.T[:-1, :-1]) + transf_matrix[:-1, -1]

    return points_transformed.astype(points.dtype)


def get_transformation_matrix(agent_translation_m: np.ndarray, agent_yaw: float) -> np.ndarray:
    """Get transformation matrix from world to vehicle frame

    Args:
        agent_translation_m (np.ndarray): (x, y) position of the vehicle in world frame
        agent_yaw (float): rotation of the vehicle in the world frame

    Returns:
        (np.ndarray) transformation matrix from world to vehicle
    """

    # Translate world to ego by applying the negative ego translation.
    world_to_agent_in_2d = np.eye(3, dtype=np.float64)
    world_to_agent_in_2d[0:2, 2] = -agent_translation_m[0:2]

    # Rotate counter-clockwise by negative yaw to align world such that ego faces right.
    world_to_agent_in_2d = _yaw_as_rotation33(-agent_yaw) @ world_to_agent_in_2d

    return world_to_agent_in_2d


# transformation for torch
def torch_rad2rot(rad: Tensor) -> Tensor:
    """
    Args:
        rad: [n_batch] or [n_scene, n_agent] or etc.

    Returns:
        rot_mat: [{rad.shape}, 2, 2]
    """
    _cos = torch.cos(rad)
    _sin = torch.sin(rad)
    return torch.stack([torch.stack([_cos, -_sin], dim=-1), torch.stack([_sin, _cos], dim=-1)], dim=-2)


def torch_sincos2rot(in_sin: Tensor, in_cos: Tensor) -> Tensor:
    """
    Args:
        in_sin: [n_batch] or [n_scene, n_agent] or etc.
        in_cos: [n_batch] or [n_scene, n_agent] or etc.

    Returns:
        rot_mat: [{in_sin.shape}, 2, 2]
    """
    return torch.stack([torch.stack([in_cos, -in_sin], dim=-1), torch.stack([in_sin, in_cos], dim=-1)], dim=-2)


def torch_pos2local(in_pos: Tensor, local_pos: Tensor, local_rot: Tensor) -> Tensor:
    """Transform M position to the local coordinates.

    Args:
        in_pos: [..., M, 2]
        local_pos: [..., 1, 2]
        local_rot: [..., 2, 2]

    Returns:
        out_pos: [..., M, 2]
    """
    return torch.matmul(in_pos - local_pos, local_rot)


def torch_pos2global(in_pos: Tensor, local_pos: Tensor, local_rot: Tensor) -> Tensor:
    """Reverse torch_pos2local

    Args:
        in_pos: [..., M, 2]
        local_pos: [..., 1, 2]
        local_rot: [..., 2, 2]

    Returns:
        out_pos: [..., M, 2]
    """
    return torch.matmul(in_pos, local_rot.transpose(-1, -2)) + local_pos


def torch_dir2local(in_dir: Tensor, local_rot: Tensor) -> Tensor:
    """Transform M dir to the local coordinates.

    Args:
        in_dir: [..., M, 2]
        local_rot: [..., 2, 2]

    Returns:
        out_dir: [..., M, 2]
    """
    return torch.matmul(in_dir, local_rot)


def torch_dir2global(in_dir: Tensor, local_rot: Tensor) -> Tensor:
    """Reverse torch_dir2local

    Args:
        in_dir: [..., M, 2]
        local_rot: [..., 2, 2]

    Returns:
        out_dir: [..., M, 2]
    """
    return torch.matmul(in_dir, local_rot.transpose(-1, -2))


def torch_rad2local(in_rad: Tensor, local_rad: Tensor, cast: bool = True) -> Tensor:
    """Transform M rad angles to the local coordinates.

    Args:
        in_rad: [..., M]
        local_rad: [...]

    Returns:
        out_rad: [..., M]
    """
    out_rad = in_rad - local_rad.unsqueeze(-1)
    if cast:
        out_rad = cast_rad(out_rad)
    return out_rad


def torch_rad2global(in_rad: Tensor, local_rad: Tensor) -> Tensor:
    """Reverse torch_rad2local

    Args:
        in_rad: [..., M]
        local_rad: [...]

    Returns:
        out_rad: [..., M]
    """
    return cast_rad(in_rad + local_rad.unsqueeze(-1))

src/utils/transform_utils.py
zhejz/TrafficBots
# Licensed under the CC BY-NC 4.0 license (https://creativecommons.org/licenses/by-nc/4.0/)
import numpy as np
import cv2
from typing import Dict, List, Optional, Tuple
from gym.wrappers.monitoring.video_recorder import ImageEncoder

COLOR_BLACK = (0, 0, 0)
COLOR_WHITE = (255, 255, 255)

COLOR_RED = (255, 0, 0)
COLOR_GREEN = (0, 255, 0)
COLOR_BLUE = (0, 0, 255)
COLOR_CYAN = (0, 255, 255)
COLOR_MAGENTA = (255, 0, 255)
COLOR_MAGENTA_2 = (255, 140, 255)
COLOR_YELLOW = (255, 255, 0)
COLOR_YELLOW_2 = (160, 160, 0)
COLOR_VIOLET = (170, 0, 255)

COLOR_BUTTER_0 = (252, 233, 79)
COLOR_BUTTER_1 = (237, 212, 0)
COLOR_BUTTER_2 = (196, 160, 0)
COLOR_ORANGE_0 = (252, 175, 62)
COLOR_ORANGE_1 = (245, 121, 0)
COLOR_ORANGE_2 = (209, 92, 0)
COLOR_CHOCOLATE_0 = (233, 185, 110)
COLOR_CHOCOLATE_1 = (193, 125, 17)
COLOR_CHOCOLATE_2 = (143, 89, 2)
COLOR_CHAMELEON_0 = (138, 226, 52)
COLOR_CHAMELEON_1 = (115, 210, 22)
COLOR_CHAMELEON_2 = (78, 154, 6)
COLOR_SKY_BLUE_0 = (114, 159, 207)
COLOR_SKY_BLUE_1 = (52, 101, 164)
COLOR_SKY_BLUE_2 = (32, 74, 135)
COLOR_PLUM_0 = (173, 127, 168)
COLOR_PLUM_1 = (117, 80, 123)
COLOR_PLUM_2 = (92, 53, 102)
COLOR_SCARLET_RED_0 = (239, 41, 41)
COLOR_SCARLET_RED_1 = (204, 0, 0)
COLOR_SCARLET_RED_2 = (164, 0, 0)
COLOR_ALUMINIUM_0 = (238, 238, 236)
COLOR_ALUMINIUM_1 = (211, 215, 207)
COLOR_ALUMINIUM_2 = (186, 189, 182)
COLOR_ALUMINIUM_3 = (136, 138, 133)
COLOR_ALUMINIUM_4 = (85, 87, 83)
COLOR_ALUMINIUM_4_5 = (66, 62, 64)
COLOR_ALUMINIUM_5 = (46, 52, 54)


class VisWaymo:
    def __init__(
        self,
        map_valid: np.ndarray,
        map_type: np.ndarray,
        map_pos: np.ndarray,
        map_boundary: np.ndarray,
        px_per_m: float = 10.0,
        video_size: int = 960,
    ) -> None:
        # centered around ego vehicle first step, x=0, y=0, theta=0
        self.px_per_m = px_per_m
        self.video_size = video_size
        self.px_agent2bottom = video_size // 2

        # waymo
        self.lane_style = [
            (COLOR_WHITE, 6),  # FREEWAY = 0
            (COLOR_ALUMINIUM_4_5, 6),  # SURFACE_STREET = 1
            (COLOR_ORANGE_2, 6),  # STOP_SIGN = 2
            (COLOR_CHOCOLATE_2, 6),  # BIKE_LANE = 3
            (COLOR_SKY_BLUE_2, 4),  # TYPE_ROAD_EDGE_BOUNDARY = 4
            (COLOR_PLUM_2, 4),  # TYPE_ROAD_EDGE_MEDIAN = 5
            (COLOR_BUTTER_0, 2),  # BROKEN = 6
            (COLOR_MAGENTA, 2),  # SOLID_SINGLE = 7
            (COLOR_SCARLET_RED_2, 2),  # DOUBLE = 8
            (COLOR_CHAMELEON_2, 4),  # SPEED_BUMP = 9
            (COLOR_SKY_BLUE_0, 4),  # CROSSWALK = 10
        ]

        self.tl_style = [
            COLOR_ALUMINIUM_1,  # STATE_UNKNOWN = 0;
            COLOR_RED,  # STOP = 1;
            COLOR_YELLOW,  # CAUTION = 2;
            COLOR_GREEN,  # GO = 3;
            COLOR_VIOLET,  # FLASHING = 4;
        ]
        # sdc=0, interest=1, predict=2
        self.agent_role_style = [COLOR_CYAN, COLOR_CHAMELEON_2, COLOR_MAGENTA]

        self.agent_cmd_txt = [
            "STATIONARY",  # STATIONARY = 0;
            "STRAIGHT",  # STRAIGHT = 1;
            "STRAIGHT_LEFT",  # STRAIGHT_LEFT = 2;
            "STRAIGHT_RIGHT",  # STRAIGHT_RIGHT = 3;
            "LEFT_U_TURN",  # LEFT_U_TURN = 4;
            "LEFT_TURN",  # LEFT_TURN = 5;
            "RIGHT_U_TURN",  # RIGHT_U_TURN = 6;
            "RIGHT_TURN",  # RIGHT_TURN = 7;
        ]

        raster_map, self.top_left_px = self._register_map(map_boundary, self.px_per_m)
        self.raster_map = self._draw_map(raster_map, map_valid, map_type, map_pos)

    @staticmethod
    def _register_map(map_boundary: np.ndarray, px_per_m: float, edge_px: int = 100) -> Tuple[np.ndarray, np.ndarray]:
        """
        Args:
            map_boundary: [4], xmin, xmax, ymin, ymax
            px_per_m: float

        Returns:
            raster_map: empty image
            top_left_px
        """
        # y axis is inverted in pixel coordinate
        xmin, xmax, ymax, ymin = (map_boundary * px_per_m).astype(np.int64)
        ymax *= -1
        ymin *= -1
        xmin -= edge_px
        ymin -= edge_px
        xmax += edge_px
        ymax += edge_px

        raster_map = np.zeros([ymax - ymin, xmax - xmin, 3], dtype=np.uint8)
        top_left_px = np.array([xmin, ymin], dtype=np.float32)
        return raster_map, top_left_px

    def _draw_map(
        self,
        raster_map: np.ndarray,
        map_valid: np.ndarray,
        map_type: np.ndarray,
        map_pos: np.ndarray,
        attn_weights_to_pl: Optional[np.ndarray] = None,
    ) -> np.ndarray:
        """
        Args: numpy arrays
            map_valid: [n_pl, 20],  # bool
            map_type: [n_pl, 11],  # bool one_hot
            map_pos: [n_pl, 20, 2],  # float32
            attn_weights_to_pl: [n_pl], sum up to 1

        Returns:
            raster_map
        """
        mask_valid = map_valid.any(axis=1)
        if attn_weights_to_pl is None:
            attn_weights_to_pl = np.zeros(map_valid.shape[0]) - 1

        for type_to_draw in range(len(self.lane_style)):
            for i in np.where((map_type[:, type_to_draw]) & mask_valid)[0]:
                color, thickness = self.lane_style[type_to_draw]
                if attn_weights_to_pl[i] > 0:
                    color = tuple(np.array(color) * attn_weights_to_pl[i])
                cv2.polylines(
                    raster_map,
                    [self._to_pixel(map_pos[i][map_valid[i]])],
                    isClosed=False,
                    color=color,
                    thickness=thickness,
                    lineType=cv2.LINE_AA,
                )

        # for i in range(mask_valid.shape[0]):
        #     if mask_valid[i]:
        #         cv2.arrowedLine(
        #             raster_map,
        #             self._to_pixel(map_pos[i, 0]),
        #             self._to_pixel(map_pos[i, 0] + (map_pos[i, 2] - map_pos[i, 0]) * 1),
        #             color=COLOR_CYAN,
        #             thickness=4,
        #             line_type=cv2.LINE_AA,
        #             tipLength=0.5,
        #         )
        return raster_map

    def save_prediction_videos(
        self,
        video_base_name: str,
        episode: Dict[str, np.ndarray],
        prediction: Optional[Dict[str, np.ndarray]],
        save_agent_view: bool = True,
    ) -> List[str]:
        """
        Args:
            episode["agent/valid"]: np.zeros([n_step, N_AGENT], dtype=bool),  # bool,
            episode["agent/pos"]: np.zeros([n_step, N_AGENT, 2], dtype=np.float32),  # x,y
            episode["agent/yaw_bbox"]: np.zeros([n_step, N_AGENT, 1], dtype=np.float32),  # [-pi, pi]
            episode["agent/role"]: np.zeros([N_AGENT, 3], dtype=bool) one_hot [sdc=0, interest=1, predict=2]
            episode["agent/size"]: np.zeros([N_AGENT, 3], dtype=np.float32),  # float32: [length, width, height]
            episode["map/valid"]: np.zeros([N_PL, 20], dtype=bool),  # bool
            episode["map/pos"]: np.zeros([N_PL, 20, 2], dtype=np.float32),  # x,y
            episode["tl_lane/valid"]: np.zeros([n_step, N_TL], dtype=bool),  # bool
            episode["tl_lane/state"]: np.zeros([n_step, N_TL, N_TL_STATE], dtype=bool),  # one_hot
            episode["tl_lane/idx"]: np.zeros([n_step, N_TL], dtype=np.int64) - 1,  # int, -1 means not valid
            episode["tl_stop/valid"]: np.zeros([n_step, N_TL_STOP], dtype=bool),  # bool
            episode["tl_stop/state"]: np.zeros([n_step, N_TL_STOP, N_TL_STATE], dtype=bool),  # one_hot
            episode["tl_stop/pos"]: np.zeros([n_step, N_TL_STOP, 2], dtype=np.float32)  # x,y
            episode["tl_stop/dir"]: np.zeros([n_step, N_TL_STOP, 2], dtype=np.float32)  # x,y
            prediction["step_current"] <= prediction["step_gt"] <= prediction["step_end"]
            prediction["agent/valid"]: [step_current+1...step_end, N_AGENT]
            prediction["agent/pos"]: [step_current+1...step_end, N_AGENT, 2]
            prediction["agent/yaw_bbox"]: [step_current+1...step_end, N_AGENT, 1]
            prediction["speed"]: [step_current+1...step_end, N_AGENT], m/s
            prediction["collided"]: [step_current+1...step_end, N_AGENT], bool
            prediction["collided_this_step"]: [step_current+1...step_end, N_AGENT], bool
            prediction["target_reached"]: [step_current+1...step_end, N_AGENT], bool
            prediction["target_reached_this_step"]: [step_current+1...step_end, N_AGENT], bool
            prediction["route_completed"]: [step_current+1...step_end, N_AGENT], [0,1] percentage
            prediction["route_completed_this_step"]: [step_current+1...step_end, N_AGENT], [0,1] percentage
            prediction["run_red_light"]: [step_current+1...step_end, N_AGENT], bool
            prediction["run_red_light_this_step"]: [step_current+1...step_end, N_AGENT], bool
        """
        buffer_video = {f"{video_base_name}-gt.mp4": [[], None]}  # [List[im], agent_id]
        if prediction is None:
            step_end = episode["agent/valid"].shape[0] - 1
            step_gt = step_end
            step_current = step_end
        else:
            step_current = prediction["step_current"]
            step_end = prediction["step_end"]
            step_gt = prediction["step_gt"]
            buffer_video[f"{video_base_name}-pd.mp4"] = [[], None]
            buffer_video[f"{video_base_name}-mix.mp4"] = [[], None]
            if save_agent_view:
                buffer_video[f"{video_base_name}-sdc.mp4"] = [[], np.where(episode["agent/role"][:, 0])[0][0]]
                for i in np.where(episode["agent/role"][:, 1])[0]:
                    buffer_video[f"{video_base_name}-int_{i}.mp4"] = [[], i]
                for i in np.where(episode["agent/role"][:, 2])[0]:
                    buffer_video[f"{video_base_name}-pre_{i}.mp4"] = [[], i]
                n_others_to_vis = 5
                idx_others = np.where(prediction["agent/valid"].any(0) & ~(episode["agent/role"].any(-1)))[0]
                for i in idx_others[:n_others_to_vis]:
                    buffer_video[f"{video_base_name}-other_{i}.mp4"] = [[], i]

        for t in range(step_end + 1):
            step_image = self.raster_map.copy()
            # draw traffic lights
            if "tl_lane/valid" in episode:
                t_tl = min(t, step_gt)
                for i in range(episode["tl_lane/valid"].shape[1]):
                    if episode["tl_lane/valid"][t_tl, i]:
                        lane_idx = episode["tl_lane/idx"][t_tl, i]
                        tl_state = episode["tl_lane/state"][t_tl, i].argmax()
                        pos = self._to_pixel(episode["map/pos"][lane_idx][episode["map/valid"][lane_idx]])
                        cv2.polylines(
                            step_image,
                            [pos],
                            isClosed=False,
                            color=self.tl_style[tl_state],
                            thickness=8,
                            lineType=cv2.LINE_AA,
                        )
                        if tl_state >= 1 and tl_state <= 3:
                            cv2.drawMarker(
                                step_image,
                                pos[-1],
                                color=self.tl_style[tl_state],
                                markerType=cv2.MARKER_TILTED_CROSS,
                                markerSize=10,
                                thickness=6,
                            )
            # draw traffic lights stop points
            if "tl_stop/valid" in episode:
                for i in range(episode["tl_stop/valid"].shape[1]):
                    if episode["tl_stop/valid"][t_tl, i]:
                        tl_state = episode["tl_stop/state"][t_tl, i].argmax()
                        stop_point = self._to_pixel(episode["tl_stop/pos"][t_tl, i])
                        stop_point_end = self._to_pixel(
                            episode["tl_stop/pos"][t_tl, i] + 5 * episode["tl_stop/dir"][t_tl, i]
                        )
                        cv2.arrowedLine(
                            step_image,
                            stop_point,
                            stop_point_end,
                            color=self.tl_style[tl_state],
                            thickness=4,
                            line_type=cv2.LINE_AA,
                            tipLength=0.3,
                        )

            # draw agents: prediction["step_current"] <= prediction["step_gt"] <= prediction["step_end"]
            step_image_gt = step_image.copy()
            raster_blend_gt = np.zeros_like(step_image)
            if t <= step_gt:
                bbox_gt = self._get_agent_bbox(
                    episode["agent/valid"][t],
                    episode["agent/pos"][t],
                    episode["agent/yaw_bbox"][t],
                    episode["agent/size"],
                )
                bbox_gt = self._to_pixel(bbox_gt)
                agent_role = episode["agent/role"][episode["agent/valid"][t]]
                heading_start = self._to_pixel(episode["agent/pos"][t][episode["agent/valid"][t]])
                heading_end = self._to_pixel(
                    episode["agent/pos"][t][episode["agent/valid"][t]]
                    + 1.5
                    * np.stack(
                        [
                            np.cos(episode["agent/yaw_bbox"][t, :, 0][episode["agent/valid"][t]]),
                            np.sin(episode["agent/yaw_bbox"][t, :, 0][episode["agent/valid"][t]]),
                        ],
                        axis=-1,
                    )
                )
                for i in range(agent_role.shape[0]):
                    if not agent_role[i].any():
                        color = COLOR_ALUMINIUM_0
                    else:
                        color = self.agent_role_style[np.where(agent_role[i])[0].min()]
                    cv2.fillConvexPoly(step_image_gt, bbox_gt[i], color=color)
                    cv2.fillConvexPoly(raster_blend_gt, bbox_gt[i], color=color)
                    cv2.arrowedLine(
                        step_image_gt,
                        heading_start[i],
                        heading_end[i],
                        color=COLOR_BLACK,
                        thickness=4,
                        line_type=cv2.LINE_AA,
                        tipLength=0.6,
                    )
            buffer_video[f"{video_base_name}-gt.mp4"][0].append(step_image_gt)

            if prediction is not None:
                if t > prediction["step_current"]:
                    step_image_pd = step_image.copy()
                    t_pred = t - prediction["step_current"] - 1
                    bbox_pred = self._get_agent_bbox(
                        prediction["agent/valid"][t_pred],
                        prediction["agent/pos"][t_pred],
                        prediction["agent/yaw_bbox"][t_pred],
                        episode["agent/size"],
                    )
                    bbox_pred = self._to_pixel(bbox_pred)
                    heading_start = self._to_pixel(prediction["agent/pos"][t_pred][prediction["agent/valid"][t_pred]])
                    heading_end = self._to_pixel(
                        prediction["agent/pos"][t_pred][prediction["agent/valid"][t_pred]]
                        + 1.5
                        * np.stack(
                            [
                                np.cos(prediction["agent/yaw_bbox"][t_pred, :, 0][prediction["agent/valid"][t_pred]]),
                                np.sin(prediction["agent/yaw_bbox"][t_pred, :, 0][prediction["agent/valid"][t_pred]]),
                            ],
                            axis=-1,
                        )
                    )
                    agent_role = episode["agent/role"][prediction["agent/valid"][t_pred]]
                    for i in range(agent_role.shape[0]):
                        if not agent_role[i].any():
                            color = COLOR_ALUMINIUM_0
                        else:
                            color = self.agent_role_style[np.where(agent_role[i])[0].min()]
                        cv2.fillConvexPoly(step_image_pd, bbox_pred[i], color=color)
                        cv2.arrowedLine(
                            step_image_pd,
                            heading_start[i],
                            heading_end[i],
                            color=COLOR_BLACK,
                            thickness=4,
                            line_type=cv2.LINE_AA,
                            tipLength=0.6,
                        )
                    # step_image_mix = step_image.copy()
                    # cv2.addWeighted(raster_blend_gt, 0.6, step_image_pd, 1, 0, step_image_mix)
                    step_image_mix = cv2.addWeighted(raster_blend_gt, 0.6, step_image_pd, 1, 0)
                else:
                    step_image_pd = step_image_gt.copy()
                    step_image_mix = step_image_gt.copy()
                buffer_video[f"{video_base_name}-pd.mp4"][0].append(step_image_pd)
                buffer_video[f"{video_base_name}-mix.mp4"][0].append(step_image_mix)

            if save_agent_view:
                for k, v in buffer_video.items():
                    agent_idx = v[1]
                    if agent_idx is not None:
                        text_valid = True
                        if t <= step_current:
                            pred_started = False
                            valid_t = t
                            if not episode["agent/valid"][valid_t, agent_idx]:
                                # get the first valid step
                                text_valid = False
                                valid_t = np.where(episode["agent/valid"][:, agent_idx])[0][0]
                            ev_loc = self._to_pixel(episode["agent/pos"][valid_t, agent_idx])
                            ev_rot = episode["agent/yaw_bbox"][valid_t, agent_idx, 0]
                        else:
                            pred_started = True
                            valid_t = t - step_current - 1
                            if not prediction["agent/valid"][valid_t, agent_idx]:
                                # get the closest valid step
                                text_valid = False
                                valid_steps = np.where(prediction["agent/valid"][:, agent_idx])[0]
                                valid_t_idx = np.abs(valid_t - valid_steps).argmin()
                                valid_t = valid_steps[valid_t_idx]
                            ev_loc = self._to_pixel(prediction["agent/pos"][valid_t, agent_idx])
                            ev_rot = prediction["agent/yaw_bbox"][valid_t, agent_idx, 0]

                        agent_view = step_image_mix.copy()
                        if not episode["agent/role"][agent_idx].any():
                            color = COLOR_ALUMINIUM_0
                        else:
                            color = self.agent_role_style[np.where(episode["agent/role"][agent_idx])[0].min()]
                        # draw dest
                        cv2.arrowedLine(
                            agent_view,
                            ev_loc,
                            self._to_pixel(episode["map/pos"][episode["agent/dest"][agent_idx]][0, :2]),
                            color=COLOR_BUTTER_0,
                            thickness=4,
                            line_type=cv2.LINE_AA,
                            tipLength=0.05,
                        )
                        # draw goal
                        cv2.arrowedLine(
                            agent_view,
                            ev_loc,
                            self._to_pixel(episode["agent/goal"][agent_idx, :2]),
                            color=COLOR_MAGENTA,
                            thickness=4,
                            line_type=cv2.LINE_AA,
                            tipLength=0.05,
                        )
                        if "agent/dest" in prediction:
                            cv2.arrowedLine(
                                agent_view,
                                ev_loc,
                                self._to_pixel(episode["map/pos"][prediction["agent/dest"][agent_idx]][0, :2]),
                                color=COLOR_GREEN,
                                thickness=2,
                                line_type=cv2.LINE_AA,
                                tipLength=0.1,
                            )
                        if "agent/goal" in prediction:
                            cv2.arrowedLine(
                                agent_view,
                                ev_loc,
                                self._to_pixel(prediction["agent/goal"][agent_idx, :2]),
                                color=COLOR_GREEN,
                                thickness=2,
                                line_type=cv2.LINE_AA,
                                tipLength=0.1,
                            )
                        trans = self._get_warp_transform(ev_loc, ev_rot)
                        agent_view = cv2.warpAffine(agent_view, trans, (self.video_size, self.video_size))
                        agent_view = self._add_txt(
                            agent_view, episode, prediction, valid_t, agent_idx, pred_started, text_valid
                        )
                        v[0].append(agent_view)

        for k, v in buffer_video.items():
            encoder = ImageEncoder(k, v[0][0].shape, 20, 20)
            for im in v[0]:
                encoder.capture_frame(im)
            encoder.close()
            encoder = None

        return list(buffer_video.keys())

    def save_attn_videos(
        self, video_base_name: str, episode: Dict[str, np.ndarray], prediction: Optional[Dict[str, np.ndarray]]
    ) -> List[str]:
        # todo: deprecated, to be fixed
        step_current = prediction["step_current"]
        step_end = prediction["step_end"]
        step_gt = prediction["step_gt"]

        # normalize attention weights
        # [n_pl, 20]
        valid_pl = episode["map/valid"].any(1)
        # [n_step, n_agent, n_pl]
        attn_min = prediction["attn_weights_to_pl"][:, :, valid_pl].min(axis=(0, 2), keepdims=True)
        attn_max = prediction["attn_weights_to_pl"][:, :, valid_pl].max(axis=(0, 2), keepdims=True)
        attn_weights_to_pl = (prediction["attn_weights_to_pl"] - attn_min) / (attn_max - attn_min + 1e-5)
        # attn_weights_to_pl = prediction["attn_weights_to_pl"] / attn_max

        # [step_current+1...step_end, n_tl]
        n_tl = episode["tl_stop/valid"].shape[-1]
        n_step = prediction["agent/valid"].shape[0]
        valid_tl = np.zeros([n_step, n_tl], dtype=bool)
        valid_tl[: step_gt - step_current] = episode["tl_stop/valid"][step_current + 1 : step_gt + 1]
        valid_tl[step_gt - step_current :] = episode["tl_stop/valid"][step_gt]
        # [n_step, n_agent, n_tl]
        if valid_tl.sum() > 0:
            attn_min = prediction["attn_weights_to_tl"].transpose(0, 2, 1)[valid_tl].min(0)[None, :, None]
            attn_max = prediction["attn_weights_to_tl"].transpose(0, 2, 1)[valid_tl].max(0)[None, :, None]
            attn_weights_to_tl = (prediction["attn_weights_to_tl"] - attn_min) / (attn_max - attn_min + 1e-5)
            # attn_weights_to_tl = prediction["attn_weights_to_tl"] / attn_max
        else:
            # no normalization needed, no valid tl at all.
            attn_weights_to_tl = prediction["attn_weights_to_tl"]

        # [n_step, n_agent, n_agent]
        # attn_min = prediction["attn_weights_to_agent"].transpose(
        # 0, 2, 1)[prediction["agent/valid"]].min(0)[None, :, None]
        # attn_max = prediction["attn_weights_to_agent"].transpose(
        #     0, 2, 1)[prediction["agent/valid"]].max(0)[None, :, None]
        # attn_weights_to_agent = (prediction["attn_weights_to_agent"] - attn_min) / (attn_max - attn_min + 1e-5)
        attn_max = prediction["attn_weights_to_agent"].max(axis=(0, 2), keepdims=True)
        attn_weights_to_agent = prediction["attn_weights_to_agent"] / (attn_max + 1e-5)

        # video_names = {0: f"{video_base_name}-sdc-attn.mp4"} # in case sdc is always the first agent
        video_names = {np.where(episode["agent/role"][:, 0])[0][0]: f"{video_base_name}-sdc-attn.mp4"}

        for i in np.where(episode["agent/role"][:, 1])[0]:
            video_names[i] = f"{video_base_name}-int_{i}-attn.mp4"
        for i in np.where(episode["agent/role"][:, 2])[0]:
            video_names[i] = f"{video_base_name}-pre_{i}-attn.mp4"
        n_others_to_vis = 5
        idx_others = np.where(prediction["agent/valid"].any(0) & ~(episode["agent/role"].any(-1)))[0]
        for i in idx_others[:n_others_to_vis]:
            video_names[i] = f"{video_base_name}-other_{i}-attn.mp4"

        for idx_agent, v_name in video_names.items():
            im_list = []
            for t in range(step_end + 1):
                t_pred = t - step_current - 1
                if t <= step_current:
                    step_image = self.raster_map.copy()
                else:
                    step_image = self._draw_map(
                        raster_map=np.zeros_like(self.raster_map),
                        map_valid=episode["map/valid"],
                        map_type=episode["map/type"],
                        map_pos=episode["map/pos"],
                        attn_weights_to_pl=attn_weights_to_pl[t_pred, idx_agent],
                    )

                # draw agents
                if t <= step_current:
                    agent_bbox = self._get_agent_bbox(
                        episode["agent/valid"][t],
                        episode["agent/pos"][t],
                        episode["agent/yaw_bbox"][t],
                        episode["agent/size"],
                    )
                    agent_bbox = self._to_pixel(agent_bbox)
                    agent_role = np.zeros_like(episode["agent/role"][:, 0])
                    agent_role[idx_agent] = True
                    agent_role = agent_role[episode["agent/valid"][t]]
                    for i in range(agent_role.shape[0]):
                        color = COLOR_CHAMELEON_2 if agent_role[i] else COLOR_ALUMINIUM_0
                        cv2.fillConvexPoly(step_image, agent_bbox[i], color=color)
                else:
                    agent_bbox = self._get_agent_bbox(
                        prediction["agent/valid"][t_pred],
                        prediction["agent/pos"][t_pred],
                        prediction["agent/yaw_bbox"][t_pred],
                        episode["agent/size"],
                    )
                    agent_bbox = self._to_pixel(agent_bbox)
                    agent_role = np.zeros_like(episode["agent/role"][:, 0])
                    agent_role[idx_agent] = True
                    agent_role = agent_role[prediction["agent/valid"][t_pred]]
                    attn_agent = attn_weights_to_agent[t_pred, idx_agent][prediction["agent/valid"][t_pred]]
                    for i in range(agent_role.shape[0]):
                        if agent_role[i]:
                            color = COLOR_CHAMELEON_2
                        else:
                            color = tuple(np.array(COLOR_ALUMINIUM_0) * attn_agent[i])
                        cv2.fillConvexPoly(step_image, agent_bbox[i], color=color)

                # draw traffic lights
                t_tl = min(t, step_gt)
                valid_tl = episode["tl_stop/valid"][t_tl]
                if t <= step_current:
                    for idx_tl in np.where(valid_tl)[0]:
                        tl_state = episode["tl_stop/state"][t_tl, idx_tl].argmax()
                        stop_point = self._to_pixel(episode["tl_stop/pos"][t_tl, idx_tl])
                        cv2.circle(step_image, stop_point, 10, color=self.tl_style[tl_state], thickness=-1)
                else:
                    for idx_tl in np.where(valid_tl)[0]:
                        tl_state = episode["tl_stop/state"][t_tl, idx_tl].argmax()
                        stop_point = self._to_pixel(episode["tl_stop/pos"][t_tl, idx_tl])
                        attn_tl = attn_weights_to_tl[t_pred, idx_agent, idx_tl]
                        color = tuple(np.array(self.tl_style[tl_state]) * attn_tl)
                        cv2.circle(step_image, stop_point, 10, color=color, thickness=-1)

                im_list.append(step_image)

            encoder = ImageEncoder(v_name, im_list[0].shape, 20, 20)
            for im in im_list:
                encoder.capture_frame(im)
            encoder.close()
            encoder = None

        return [v_name for _, v_name in video_names.items()]

    @staticmethod
    def _add_txt(
        im: np.ndarray,
        episode: Dict[str, np.ndarray],
        prediction: Optional[Dict[str, np.ndarray]],
        t: int,
        idx: int,
        pred_started: bool,
        text_valid: bool,
        line_width: int = 30,
    ) -> np.ndarray:
        h, w, _ = im.shape
        agent_view = np.zeros([h, w + 200, 3], dtype=im.dtype)
        agent_view[:h, :w] = im
        if (prediction is not None) and pred_started:
            txt_list = [
                f'id:{int(episode["episode_idx"])}',
                f"valid:{int(text_valid)}",
                f'goal_valid:{int(prediction["goal_valid"][t, idx])}',
                f'out:{int(prediction["outside_map_this_step"][t, idx])}/{int(prediction["outside_map"][t, idx])}',
                f'col:{int(prediction["collided_this_step"][t, idx])}/{int(prediction["collided"][t, idx])}',
                f'red:{int(prediction["run_red_light_this_step"][t, idx])}/{int(prediction["run_red_light"][t, idx])}',
                f'edge:{int(prediction["run_road_edge_this_step"][t, idx])}/{int(prediction["run_road_edge"][t, idx])}',
                f'passive:{int(prediction["passive_this_step"][t, idx])}/{int(prediction["passive"][t, idx])}',
                f'r_goal:{int(prediction["goal_reached_this_step"][t, idx])}/{int(prediction["goal_reached"][t, idx])}',
                f'r_dest:{int(prediction["dest_reached_this_step"][t, idx])}/{int(prediction["dest_reached"][t, idx])}',
                f'x:{prediction["agent/pos"][t, idx, 0]:.2f}',
                f'y:{prediction["agent/pos"][t, idx, 1]:.2f}',
                f'spd:{prediction["speed"][t, idx]:.2f}',
                f'yaw:{prediction["agent/yaw_bbox"][t, idx, 0]:.2f}',
                f'acc:{prediction["action"][t, idx, 0]:.2f}',
                f'steer:{prediction["action"][t, idx, 1]:.2f}',
                f'act_P:{prediction["act_P"][t, idx]:.2f}',
                f'lat_P:{prediction["lat_P"][t, idx]:.2f}',
            ]
            if "score" in prediction:
                txt_list.append(f'score:{prediction["score"][idx]:.2f}')
            for k in [
                "rew_pg",
                "val_pg",
                "pb_pg",
                "adv_pg",
                "ret_pg",
                "rew_dr",
                "val_dr",
                "val_tgt",
                "ret_dr",
                "ret_tgt",
            ]:
                if k in prediction:
                    txt_list.append(f"{k}:{prediction[k][t, idx]:.2f}")
        else:
            txt_list = [
                f'id:{int(episode["episode_idx"])}',
                f"valid:{int(text_valid)}",
                f'g_x:{episode["agent/goal"][idx, 0]:.2f}',
                f'g_y:{episode["agent/goal"][idx, 1]:.2f}',
                f'x:{episode["agent/pos"][t, idx, 0]:.2f}',
                f'y:{episode["agent/pos"][t, idx, 1]:.2f}',
                f'yaw:{episode["agent/yaw_bbox"][t, idx, 0]:.2f}',
                f'spd:{episode["agent/spd"][t, idx, 0]:.2f}',
                f'role:{list(np.where(episode["agent/role"][idx])[0])}',
                f'size_x:{episode["agent/size"][idx, 0]:.2f}',
                f'size_y:{episode["agent/size"][idx, 1]:.2f}',
                f'size_z:{episode["agent/size"][idx, 2]:.2f}',
            ]

        for i, txt in enumerate(txt_list):
            agent_view = cv2.putText(
                agent_view, txt, (w, line_width * (i + 1)), cv2.FONT_HERSHEY_SIMPLEX, 0.8, (255, 255, 255), 1
            )
        return agent_view

    def _to_pixel(self, pos: np.ndarray) -> np.ndarray:
        pos = pos * self.px_per_m
        pos[..., 0] = pos[..., 0] - self.top_left_px[0]
        pos[..., 1] = -pos[..., 1] - self.top_left_px[1]
        return np.round(pos).astype(np.int32)

    def _get_warp_transform(self, loc: np.ndarray, yaw: float) -> np.ndarray:
        """
        loc: xy in pixel
        yaw: in rad
        """

        forward_vec = np.array([np.cos(yaw), -np.sin(yaw)])
        right_vec = np.array([np.sin(yaw), np.cos(yaw)])

        bottom_left = loc - self.px_agent2bottom * forward_vec - (0.5 * self.video_size) * right_vec
        top_left = loc + (self.video_size - self.px_agent2bottom) * forward_vec - (0.5 * self.video_size) * right_vec
        top_right = loc + (self.video_size - self.px_agent2bottom) * forward_vec + (0.5 * self.video_size) * right_vec

        src_pts = np.stack((bottom_left, top_left, top_right), axis=0).astype(np.float32)
        dst_pts = np.array([[0, self.video_size - 1], [0, 0], [self.video_size - 1, 0]], dtype=np.float32)
        return cv2.getAffineTransform(src_pts, dst_pts)

    @staticmethod
    def _normalize_attention(attn: np.ndarray, valid: np.ndarray) -> np.ndarray:
        # attn = attn / attn.max()
        attn_min = attn[valid].min()
        attn_max = attn[valid].max()
        attn = (attn - attn_min) / (attn_max - attn_min + 1e-5)
        return attn

    @staticmethod
    def _get_agent_bbox(
        agent_valid: np.ndarray, agent_pos: np.ndarray, agent_yaw: np.ndarray, agent_size: np.ndarray
    ) -> np.ndarray:
        yaw = agent_yaw[agent_valid]  # n, 1
        cos_yaw = np.cos(yaw)
        sin_yaw = np.sin(yaw)
        v_forward = np.concatenate([cos_yaw, sin_yaw], axis=-1)  # n,2
        v_right = np.concatenate([sin_yaw, -cos_yaw], axis=-1)

        offset_forward = 0.5 * agent_size[agent_valid, 0:1] * v_forward  # [n, 2]
        offset_right = 0.5 * agent_size[agent_valid, 1:2] * v_right  # [n, 2]

        vertex_offset = np.stack(
            [
                -offset_forward + offset_right,
                offset_forward + offset_right,
                offset_forward - offset_right,
                -offset_forward - offset_right,
            ],
            axis=1,
        )  # n,4,2

        agent_pos = agent_pos[agent_valid]
        bbox = agent_pos[:, None, :].repeat(4, 1) + vertex_offset  # n,4,2
        return bbox

    def get_dest_prob_image(
        self,
        im_base_name: str,
        episode: Dict[str, np.ndarray],
        prediction: Dict[str, np.ndarray],
        dest_prob: np.ndarray,
    ) -> Dict[str, np.ndarray]:
        """
        Args:
            episode["agent/valid"]: np.zeros([n_step, N_AGENT], dtype=bool),  # bool,
            episode["agent/pos"]: np.zeros([n_step, N_AGENT, 2], dtype=np.float32),  # x,y
            episode["agent/yaw_bbox"]: np.zeros([n_step, N_AGENT, 1], dtype=np.float32),  # [-pi, pi]
            episode["agent/role"]: np.zeros([N_AGENT, 3], dtype=bool) one_hot [sdc=0, interest=1, predict=2]
            episode["agent/size"]: np.zeros([N_AGENT, 3], dtype=np.float32),  # float32: [length, width, height]
            dest_prob: [n_agent, n_pl] float prob
        """
        list_im_idx = {}
        list_im_idx[f"{im_base_name}-sdc.jpg"] = 0
        for i in np.where(episode["agent/role"][:, 1])[0]:
            list_im_idx[f"{im_base_name}-int_{i}.jpg"] = i
        for i in np.where(episode["agent/role"][:, 2])[0]:
            list_im_idx[f"{im_base_name}-pre_{i}.jpg"] = i
        n_others_to_vis = 5
        idx_others = np.where(prediction["agent/valid"].any(0) & ~(episode["agent/role"].any(-1)))[0]
        for i in idx_others[:n_others_to_vis]:
            list_im_idx[f"{im_base_name}-other_{i}.jpg"] = i

        for im_path, i in list_im_idx.items():
            t = episode["agent/valid"][:, i].argmax()
            dest_valid = dest_prob[i] > 1e-4
            p_normalized = dest_prob[i, dest_valid].copy()
            p_max = p_normalized.max()
            p_min = p_normalized.min()
            p_normalized = (p_normalized - p_min) / (p_max - p_min + 1e-4) * 3.0
            m_type = np.zeros_like(episode["map/type"][dest_valid])
            m_type[:, 1] = True
            for k in p_normalized.argsort()[-6:]:
                m_type[k, 1] = False
                m_type[k, 3] = True
            im = self._draw_map(
                raster_map=np.zeros_like(self.raster_map),
                map_valid=episode["map/valid"][dest_valid],
                map_type=m_type,
                map_pos=episode["map/pos"][dest_valid],
                attn_weights_to_pl=p_normalized,
            )

            draw_gt_dest = self._to_pixel(episode["map/pos"][episode["agent/dest"][i]])
            # [n_pl, 20]
            draw_gt_dest_valid = episode["map/valid"][episode["agent/dest"][i]]
            cv2.polylines(
                im,
                [draw_gt_dest[draw_gt_dest_valid]],
                isClosed=False,
                color=COLOR_MAGENTA,
                thickness=2,
                lineType=cv2.LINE_AA,
            )

            bbox_gt = self._get_agent_bbox(
                episode["agent/valid"][t, i],
                episode["agent/pos"][t, i],
                episode["agent/yaw_bbox"][t, i],
                episode["agent/size"][i],
            )
            bbox_gt = self._to_pixel(bbox_gt)
            heading_start = self._to_pixel(episode["agent/pos"][t, i])
            heading_end = self._to_pixel(
                episode["agent/pos"][t, i]
                + 1.5
                * np.stack(
                    [np.cos(episode["agent/yaw_bbox"][t, i, 0]), np.sin(episode["agent/yaw_bbox"][t, i, 0])], axis=-1
                )
            )
            cv2.fillConvexPoly(im, bbox_gt, color=COLOR_RED)
            cv2.arrowedLine(
                im, heading_start, heading_end, color=COLOR_BLACK, thickness=4, line_type=cv2.LINE_AA, tipLength=0.6
            )
            cv2.imwrite(im_path, im[..., ::-1])
        return list(list_im_idx.keys())

src/utils/vis_waymo.py
zhejz/TrafficBots