{"cells":[{"metadata":{},"cell_type":"markdown","source":"# Overview\n\nThere's a lot of recent literature that tries to solve trajectory prediction problem using RNN/Transformer and considering multi-agents and their graph based relationships.\n\nTo name a few:\n\n* https://paperswithcode.com/paper/smart-simultaneous-multi-agent-recurrent\n* https://paperswithcode.com/paper/forecasting-trajectory-and-behavior-of-road\n* https://paperswithcode.com/paper/trajectron-multi-agent-generative-trajectory"},{"metadata":{},"cell_type":"markdown","source":"There are several things I want to handle before diving into these types of methods:\n\n1. decouple rasterization (and # of frames) from raw data of num of history frames to use\n2. co-locate all agents within the same frames in a same `torch` dataset format\n\nGiven the already wonderfully built `l5kit` library and the complexity of `AgentDataset` I'll try to reuse it by delegation.\n\nThis notebook uses `pytorch-lightning` for data-loading but you can still do it in the plain old vanilla pytorch way."},{"metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true},"cell_type":"code","source":"!pip install -q pip==20.2.3\n!pip uninstall -y typing\n!pip install -q l5kit==1.1 pytorch-lightning==0.10.0","execution_count":null,"outputs":[]},{"metadata":{"_uuid":"d629ff2d2480ee46fbb7e2d37f6b5fab8052498a","_cell_guid":"79c7e3d0-c299-4dcb-8224-4455121ee9b0","trusted":true},"cell_type":"code","source":"import bisect\nimport os\nfrom copy import deepcopy\nfrom operator import itemgetter\nfrom typing import Any, Dict, List, Optional, Tuple\n\nimport numpy as np\nimport pytorch_lightning as pl\nfrom l5kit.data import ChunkedDataset, LocalDataManager\nfrom l5kit.dataset import AgentDataset\nfrom l5kit.rasterization import StubRasterizer, build_rasterizer\nfrom torch.utils.data import DataLoader, Dataset, Subset\n\nis_kaggle = os.path.isdir(\"/kaggle\")\n\n\ndata_root = (\n    \"/kaggle/input/lyft-motion-prediction-autonomous-vehicles\"\n    if is_kaggle\n    else \"lyft-motion-prediction-autonomous-vehicles\"\n)\n\n\nCONFIG_DATA = {\n    \"format_version\": 4,\n    \"model_params\": {\n        \"model_architecture\": \"resnet34\",\n        \"history_num_frames\": 10,\n        \"history_step_size\": 1,\n        \"history_delta_time\": 0.1,\n        \"future_num_frames\": 50,\n        \"future_step_size\": 1,\n        \"future_delta_time\": 0.1,\n    },\n    \"raster_params\": {\n        \"raster_size\": [256, 256],\n        \"pixel_size\": [0.5, 0.5],\n        \"ego_center\": [0.25, 0.5],\n        \"map_type\": \"py_semantic\",\n        \"satellite_map_key\": \"aerial_map/aerial_map.png\",\n        \"semantic_map_key\": \"semantic_map/semantic_map.pb\",\n        \"dataset_meta_key\": \"meta.json\",\n        \"filter_agents_threshold\": 0.5,\n        \"disable_traffic_light_faces\": False,\n    },\n    \"train_dataloader\": {\n        \"key\": \"scenes/sample.zarr\",\n        \"batch_size\": 24,\n        \"shuffle\": True,\n        \"num_workers\": 0,\n    },\n    \"val_dataloader\": {\n        \"key\": \"scenes/validate.zarr\",\n        \"batch_size\": 24,\n        \"shuffle\": False,\n        \"num_workers\": 4,\n    },\n    \"test_dataloader\": {\n        \"key\": \"scenes/test.zarr\",\n        \"batch_size\": 24,\n        \"shuffle\": False,\n        \"num_workers\": 4,\n    },\n    \"train_params\": {\n        \"max_num_steps\": 400,\n        \"eval_every_n_steps\": 50,\n    },\n}","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"class MultiAgentDataset(Dataset):\n    def __init__(\n        self,\n        rast_only_agent_dataset: AgentDataset,\n        history_agent_dataset: AgentDataset,\n        num_neighbors: int = 10,\n    ):\n        super().__init__()\n        self.rast_only_agent_dataset = rast_only_agent_dataset\n        self.history_agent_dataset = history_agent_dataset\n        self.num_neighbors = num_neighbors\n\n    def __len__(self) -> int:\n        return len(self.rast_only_agent_dataset)\n\n    def get_others_dict(\n        self, index: int, ego_dict: Dict[str, Any]\n    ) -> Tuple[List[Dict[str, Any]], int]:\n        agent_index = self.rast_only_agent_dataset.agents_indices[index]\n        frame_index = bisect.bisect_right(\n            self.rast_only_agent_dataset.cumulative_sizes_agents, agent_index\n        )\n        frame_indices = self.rast_only_agent_dataset.get_frame_indices(frame_index)\n        assert len(frame_indices) >= 1, frame_indices\n        frame_indices = frame_indices[frame_indices != index]\n\n        others_dict = []\n        # The centroid of the AV in the current frame in world reference system. Unit is meters\n        for idx, agent in zip(  # type: ignore\n            frame_indices,\n            Subset(self.history_agent_dataset, frame_indices),\n        ):\n            agent[\"dataset_idx\"] = idx\n            agent[\"dist_to_ego\"] = np.linalg.norm(\n                agent[\"centroid\"] - ego_dict[\"centroid\"], ord=2\n            )\n            # TODO in future we can convert history positions via agent + ego transformation matrix\n            # TODO and get the normalized version\n            del agent[\"image\"]\n            others_dict.append(agent)\n\n        others_dict = sorted(others_dict, key=itemgetter(\"dist_to_ego\"))\n        others_dict = others_dict[: self.num_neighbors]\n        others_len = len(others_dict)\n\n        # have to pad because torch has no ragged tensor\n        # https://github.com/pytorch/pytorch/issues/25032\n        length_to_pad = self.num_neighbors - others_len\n        pad_item = deepcopy(ego_dict)\n        pad_item[\"dataset_idx\"] = index\n        pad_item[\"dist_to_ego\"] = np.nan  # set to nan so you don't by chance use this\n        del pad_item[\"image\"]\n        return (others_dict + [pad_item] * length_to_pad, others_len)\n\n    def __getitem__(self, index: int) -> Dict[str, Any]:\n        rast_dict = self.rast_only_agent_dataset[index]\n        ego_dict = self.history_agent_dataset[index]\n        others_dict, others_len = self.get_others_dict(index, ego_dict)\n        ego_dict[\"image\"] = rast_dict[\"image\"]\n        return {\n            \"ego_dict\": ego_dict,\n            \"others_dict\": others_dict,\n            \"others_len\": others_len,\n        }","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"class LyftAgentDataModule(pl.LightningDataModule):\n    def __init__(self, cfg: Dict = CONFIG_DATA, data_root: str = data_root):\n        super().__init__()\n        self.cfg = cfg\n        self.dm = LocalDataManager(data_root)\n        self.rast = build_rasterizer(self.cfg, self.dm)\n\n    def chunked_dataset(self, key: str):\n        dl_cfg = self.cfg[key]\n        dataset_path = self.dm.require(dl_cfg[\"key\"])\n        zarr_dataset = ChunkedDataset(dataset_path)\n        zarr_dataset.open()\n        return zarr_dataset\n\n    def get_dataloader_by_key(\n        self, key: str, mask: Optional[np.ndarray] = None\n    ) -> DataLoader:\n        dl_cfg = self.cfg[key]\n        zarr_dataset = self.chunked_dataset(key)\n        agent_dataset = AgentDataset(\n            self.cfg, zarr_dataset, self.rast, agents_mask=mask\n        )\n        return DataLoader(\n            agent_dataset,\n            shuffle=dl_cfg[\"shuffle\"],\n            batch_size=dl_cfg[\"batch_size\"],\n            num_workers=dl_cfg[\"num_workers\"],\n            pin_memory=True,\n        )\n\n    def train_dataloader(self):\n        key = \"train_dataloader\"\n        return self.get_dataloader_by_key(key)\n\n    def val_dataloader(self):\n        key = \"val_dataloader\"\n        return self.get_dataloader_by_key(key)\n\n    def test_dataloader(self):\n        key = \"test_dataloader\"\n        test_mask = np.load(f\"{data_root}/scenes/mask.npz\")[\"arr_0\"]\n        return self.get_dataloader_by_key(key, mask=test_mask)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"class MultiAgentDataModule(LyftAgentDataModule):\n    def __init__(self, cfg: Dict = CONFIG_DATA, data_root: str = data_root) -> None:\n        super().__init__(cfg=cfg, data_root=data_root)\n        stub_cfg = deepcopy(self.cfg)\n\n        # this can be removed once l5kit supporting passing None as rasterizer\n        # https://github.com/lyft/l5kit/pull/176\n        stub_cfg[\"raster_params\"][\"map_type\"] = \"stub_debug\"\n        self.stub_rast = build_rasterizer(stub_cfg, self.dm)\n        assert isinstance(self.stub_rast, StubRasterizer)\n\n    def get_dataloader_by_key(\n        self, key: str, mask: Optional[np.ndarray] = None\n    ) -> DataLoader:\n        dl_cfg = self.cfg[key]\n        zarr_dataset = self.chunked_dataset(key)\n        # for the rast only agent dataset, we'll disable rasterization for history frames, so the\n        # channel size will only be 3 + 2 (for current frame)\n        no_history_cfg = deepcopy(self.cfg)\n        no_history_cfg[\"model_params\"][\"history_num_frames\"] = 0\n        rast_only_agent_dataset = AgentDataset(\n            no_history_cfg, zarr_dataset, self.rast, agents_mask=mask\n        )\n        history_agent_dataset = AgentDataset(\n            self.cfg, zarr_dataset, self.stub_rast, agents_mask=mask\n        )\n        return DataLoader(\n            MultiAgentDataset(\n                rast_only_agent_dataset=rast_only_agent_dataset,\n                history_agent_dataset=history_agent_dataset,\n            ),\n            shuffle=dl_cfg[\"shuffle\"],\n            batch_size=dl_cfg[\"batch_size\"],\n            num_workers=dl_cfg[\"num_workers\"],\n            pin_memory=True,\n        )","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"datamodule = MultiAgentDataModule()","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"from pprint import pprint\nfor item in datamodule.train_dataloader():\n    pprint(item.keys())\n    print('ego_dict keys')\n    pprint(item['ego_dict'].keys())\n    pprint(len(item['others_dict']))\n    pprint(item['others_dict'][0].keys())\n    pprint(item['others_len'])\n    break","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"","execution_count":null,"outputs":[]}],"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"pygments_lexer":"ipython3","nbconvert_exporter":"python","version":"3.6.4","file_extension":".py","codemirror_mode":{"name":"ipython","version":3},"name":"python","mimetype":"text/x-python"}},"nbformat":4,"nbformat_minor":4}